wehub-resource-sync
185 lines
8.1 KiB
Python
185 lines
8.1 KiB
Python
"""Example of a custom Ray Tune experiment wrapping an RLlib Algorithm.
|
|
|
|
You should only use such a customized workflow if the following conditions apply:
|
|
- You know exactly what you are doing :)
|
|
- Configuring an existing RLlib Algorithm (e.g. PPO) via its AlgorithmConfig
|
|
is not sufficient and doesn't allow you to shape the Algorithm into behaving the way
|
|
you'd like. Note that for complex, custom evaluation procedures there are many
|
|
AlgorithmConfig options one can use (for more details, see:
|
|
https://github.com/ray-project/ray/blob/master/rllib/examples/evaluation/custom_evaluation.py). # noqa
|
|
- Subclassing an RLlib Algorithm class and overriding the new class' `training_step`
|
|
method is not sufficient and doesn't allow you to define the algorithm's execution
|
|
logic the way you'd like. See an example here on how to customize the algorithm's
|
|
`training_step()` method:
|
|
https://github.com/ray-project/ray/blob/master/rllib/examples/algorithm/custom_training_step_on_and_off_policy_combined.py # noqa
|
|
|
|
|
|
How to run this script
|
|
----------------------
|
|
`python [script file name].py`
|
|
|
|
|
|
Results to expect
|
|
-----------------
|
|
You should see the following output (at the end of the experiment) in your console:
|
|
|
|
╭───────────────────────────────────────────────────────────────────────────────────────
|
|
│ Trial name status iter total time (s) ts
|
|
├───────────────────────────────────────────────────────────────────────────────────────
|
|
│ my_experiment_CartPole-v1_77083_00000 TERMINATED 10 36.7799 60000
|
|
╰───────────────────────────────────────────────────────────────────────────────────────
|
|
╭───────────────────────────────────────────────────────╮
|
|
│ reward episode_len_mean episodes_this_iter │
|
|
├───────────────────────────────────────────────────────┤
|
|
│ 254.821 254.821 12 │
|
|
╰───────────────────────────────────────────────────────╯
|
|
evaluation episode returns=[500.0, 500.0, 500.0]
|
|
|
|
Note that evaluation results (on the CartPole-v1 env) should be close to perfect
|
|
(episode return of ~500.0) as we are acting greedily inside the evaluation procedure.
|
|
"""
|
|
from typing import Dict
|
|
|
|
import numpy as np
|
|
|
|
from ray import tune
|
|
from ray.rllib.algorithms.ppo import PPOConfig
|
|
from ray.rllib.utils.framework import try_import_torch
|
|
from ray.rllib.utils.metrics import NUM_ENV_STEPS_SAMPLED_LIFETIME
|
|
|
|
torch, _ = try_import_torch()
|
|
|
|
|
|
def my_experiment(config: Dict):
|
|
|
|
# Extract the number of iterations to run from the config.
|
|
train_iterations = config.pop("train-iterations", 2)
|
|
eval_episodes_to_do = config.pop("eval-episodes", 1)
|
|
|
|
config = (
|
|
PPOConfig()
|
|
.update_from_dict(config)
|
|
.api_stack(enable_rl_module_and_learner=True)
|
|
.environment("CartPole-v1")
|
|
)
|
|
|
|
# Train for n iterations with high LR.
|
|
config.training(lr=0.001)
|
|
algo_high_lr = config.build()
|
|
for _ in range(train_iterations):
|
|
train_results = algo_high_lr.train()
|
|
# Add the phase to the result dict.
|
|
train_results["phase"] = 1
|
|
tune.report(train_results)
|
|
phase_high_lr_time = train_results[NUM_ENV_STEPS_SAMPLED_LIFETIME]
|
|
checkpoint_training_high_lr = algo_high_lr.save()
|
|
algo_high_lr.stop()
|
|
|
|
# Train for n iterations with low LR.
|
|
config.training(lr=0.00001)
|
|
algo_low_lr = config.build()
|
|
# Load state from the high-lr algo into this one.
|
|
algo_low_lr.restore(checkpoint_training_high_lr)
|
|
for _ in range(train_iterations):
|
|
train_results = algo_low_lr.train()
|
|
# Add the phase to the result dict.
|
|
train_results["phase"] = 2
|
|
# keep time moving forward
|
|
train_results[NUM_ENV_STEPS_SAMPLED_LIFETIME] += phase_high_lr_time
|
|
tune.report(train_results)
|
|
|
|
checkpoint_training_low_lr = algo_low_lr.save()
|
|
algo_low_lr.stop()
|
|
|
|
# After training, run a manual evaluation procedure.
|
|
|
|
# Set the number of EnvRunners for collecting training data to 0 (local
|
|
# worker only).
|
|
config.env_runners(num_env_runners=0)
|
|
|
|
eval_algo = config.build()
|
|
# Load state from the low-lr algo into this one.
|
|
eval_algo.restore(checkpoint_training_low_lr)
|
|
# The algo's local worker (SingleAgentEnvRunner) that holds a
|
|
# gym.vector.Env object and an RLModule for computing actions.
|
|
local_env_runner = eval_algo.env_runner
|
|
# Extract the gymnasium env object from the created algo (its local
|
|
# SingleAgentEnvRunner worker). Note that the env in this single-agent
|
|
# case is a gymnasium vector env and that we get its first sub-env here.
|
|
env = local_env_runner.env.unwrapped.envs[0]
|
|
|
|
# The local worker (SingleAgentEnvRunner)
|
|
rl_module = local_env_runner.module
|
|
|
|
# Run a very simple env loop and add up rewards over a single episode.
|
|
obs, infos = env.reset()
|
|
episode_returns = []
|
|
episode_lengths = []
|
|
sum_rewards = length = 0
|
|
num_episodes = 0
|
|
while num_episodes < eval_episodes_to_do:
|
|
# Call the RLModule's `forward_inference()` method to compute an
|
|
# action.
|
|
rl_module_out = rl_module.forward_inference(
|
|
{
|
|
"obs": torch.from_numpy(np.expand_dims(obs, 0)), # <- add B=1
|
|
}
|
|
)
|
|
action_logits = rl_module_out["action_dist_inputs"][0] # <- remove B=1
|
|
action = np.argmax(action_logits.detach().cpu().numpy()) # act greedily
|
|
|
|
# Step the env.
|
|
obs, reward, terminated, truncated, info = env.step(action)
|
|
|
|
# Acculumate stats and reset the env, if necessary.
|
|
sum_rewards += reward
|
|
length += 1
|
|
if terminated or truncated:
|
|
num_episodes += 1
|
|
episode_returns.append(sum_rewards)
|
|
episode_lengths.append(length)
|
|
sum_rewards = length = 0
|
|
obs, infos = env.reset()
|
|
|
|
# Compile evaluation results.
|
|
eval_results = {
|
|
"eval_returns": episode_returns,
|
|
"eval_episode_lengths": episode_lengths,
|
|
}
|
|
# Combine the most recent training results with the just collected
|
|
# evaluation results.
|
|
results = {**train_results, **eval_results}
|
|
# Report everything.
|
|
tune.report(results)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
base_config = PPOConfig().environment("CartPole-v1").env_runners(num_env_runners=0)
|
|
# Convert to a plain dict for Tune. Note that this is usually not needed, you can
|
|
# pass into the below Tune Tuner any instantiated RLlib AlgorithmConfig object.
|
|
# However, for demonstration purposes, we show here how you can add other, arbitrary
|
|
# keys to the plain config dict and then pass these keys to your custom experiment
|
|
# function.
|
|
config_dict = base_config.to_dict()
|
|
|
|
# Set a Special flag signalling `my_experiment` how many training steps to
|
|
# perform on each: the high learning rate and low learning rate.
|
|
config_dict["train-iterations"] = 5
|
|
# Set a Special flag signalling `my_experiment` how many episodes to evaluate for.
|
|
config_dict["eval-episodes"] = 3
|
|
|
|
training_function = tune.with_resources(
|
|
my_experiment,
|
|
resources=base_config.algo_class.default_resource_request(base_config),
|
|
)
|
|
|
|
tuner = tune.Tuner(
|
|
training_function,
|
|
# Pass in your config dict.
|
|
param_space=config_dict,
|
|
)
|
|
results = tuner.fit()
|
|
best_results = results.get_best_result()
|
|
|
|
print(f"evaluation episode returns={best_results.metrics['eval_returns']}")
|