ray-project--ray/rllib/examples/offline_rl/cartpole_recording.py

"""Example showing how to record expert data from a trained policy.

This example:
    - demonstrates how you can train a single-agent expert PPO Policy (RLModule)
    and checkpoint it.
    - shows how you can then record expert data from the trained PPO Policy to
    disk during evaluation.

How to run this script
----------------------
`python [script file name].py --checkpoint-at-end`

For debugging, use the following additional command line options
`--no-tune --num-env-runners=0`
which should allow you to set breakpoints anywhere in the RLlib code and
have the execution stop there for inspection and debugging.

For logging to your WandB account, use:
`--wandb-key=[your WandB API key] --wandb-project=[some project name]
--wandb-run-name=[optional: WandB run name (within the defined project)]`

Results to expect
-----------------
In the console output you can see that the episode return of 350.0 is reached
before the timestep stop criteria is touched. Afterwards evaluation starts and
runs 10 iterations while recording the data. The number of recorded experiences
might differ from evaluation run to evaluation run because evaluation
`EnvRunner`s sample episodes while recording timesteps and episodes contain
usually different numbers of timesteps. Note, this is different when recording
episodes - in this case each row is one episode.

+-----------------------------+------------+----------------------+
| Trial name                  | status     | loc                  |
|                             |            |                      |
|-----------------------------+------------+----------------------+
| PPO_CartPole-v1_df83f_00000 | TERMINATED | 192.168.0.119:233661 |
+-----------------------------+------------+----------------------+
+--------+------------------+------------------------+------------------------+
|   iter |   total time (s) |   num_training_step_ca |   num_env_steps_sample |
|        |                  |      lls_per_iteration |             d_lifetime |
+--------+------------------+------------------------+------------------------|
|     21 |          25.9162 |                      1 |                  84000 |
+--------+------------------+------------------------+------------------------+

...

Number of experiences recorded: 26644
"""
from pathlib import Path

import ray
from ray.rllib.algorithms.ppo import PPOConfig
from ray.rllib.core import (
    COMPONENT_LEARNER,
    COMPONENT_LEARNER_GROUP,
    COMPONENT_RL_MODULE,
)
from ray.rllib.core.columns import Columns
from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
from ray.rllib.examples.utils import (
    add_rllib_example_script_args,
    run_rllib_example_script_experiment,
)
from ray.rllib.utils.metrics import (
    ENV_RUNNER_RESULTS,
    EPISODE_RETURN_MEAN,
    EVALUATION_RESULTS,
    NUM_ENV_STEPS_SAMPLED_LIFETIME,
)

parser = add_rllib_example_script_args(
    default_timesteps=200000,
    default_reward=350.0,
)
parser.set_defaults(
    checkpoint_at_end=True,
    max_concurrent_trials=1,
)
# Use `parser` to add your own custom command line options to this script
# and (if needed) use their values to set up `config` below.
args = parser.parse_args()

config = (
    PPOConfig()
    .env_runners(
        num_env_runners=5,
    )
    .environment("CartPole-v1")
    .rl_module(
        model_config=DefaultModelConfig(
            fcnet_hiddens=[32],
            fcnet_activation="linear",
            vf_share_layers=True,
        ),
    )
    .training(
        lr=0.0003,
        num_epochs=6,
        vf_loss_coeff=0.01,
    )
    .evaluation(
        evaluation_num_env_runners=1,
        evaluation_interval=1,
        evaluation_parallel_to_training=True,
        evaluation_config=PPOConfig.overrides(explore=False),
    )
)

stop = {
    f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": args.stop_timesteps,
    f"{EVALUATION_RESULTS}/{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": (
        args.stop_reward
    ),
}


if __name__ == "__main__":
    results = run_rllib_example_script_experiment(config, args, stop=stop)

    # Store the best checkpoint for recording.
    best_checkpoint = results.get_best_result(
        metric=f"{EVALUATION_RESULTS}/{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}",
        mode="max",
    ).checkpoint.path

    # Configure the algorithm for offline recording.
    config.offline_data(
        output="local:///tmp/cartpole/",
        # Store columnar (tabular) data.
        output_write_episodes=True,
        # Each file should hold 1,000 rows.
        output_max_rows_per_file=1000,
        output_write_remaining_data=True,
        # LZ4-compress columns 'obs', 'new_obs', and 'actions' to
        # save disk space and increase performance. Note, this means
        # that you have to use `input_compress_columns` in the same
        # way when using the data for training in `RLlib`.
        output_compress_columns=[Columns.OBS, Columns.ACTIONS],
    )
    config.env_runners(
        batch_mode="complete_episodes",
    )
    # Change the evaluation settings to sample exactly 50 episodes
    # per evaluation iteration and increase the number of evaluation
    # env-runners to 5.
    config.evaluation(
        evaluation_num_env_runners=5,
        evaluation_duration=50,
        evaluation_duration_unit="episodes",
        evaluation_interval=1,
        evaluation_parallel_to_training=False,
        evaluation_config=PPOConfig.overrides(explore=False),
    )

    # Build the algorithm for evaluation.
    algo = config.build()
    # Load the checkpoint stored above.
    rl_module_checkpoint = (
        Path(best_checkpoint)
        / COMPONENT_LEARNER_GROUP
        / COMPONENT_LEARNER
        / COMPONENT_RL_MODULE
    )
    algo.restore_from_path(
        rl_module_checkpoint,
        # Note, to load only a component of a checkpoint, specify the
        # `component` argument and provide the path to the component.
        component=(
            f"{COMPONENT_LEARNER_GROUP}"
            f"/{COMPONENT_LEARNER}"
            f"/{COMPONENT_RL_MODULE}"
        ),
    )

    # Evaluate over 10 iterations and record the data.
    for i in range(10):
        print(f"Iteration: {i + 1}:\n")
        res = algo.evaluate()
        print(res)

    # Stop the algorithm.
    algo.stop()

    # Check the number of rows in the dataset.
    ds = ray.data.read_parquet("local:///tmp/cartpole")
    print(f"Number of experiences recorded: {ds.count()}")