wehub-resource-sync
265 lines
12 KiB
Python
265 lines
12 KiB
Python
"""Example showing how to restore an Algorithm from a checkpoint and resume training.
|
|
|
|
Use the setup shown in this script if your experiments tend to crash after some time,
|
|
and you would therefore like to make your setup more robust and fault-tolerant.
|
|
|
|
This example:
|
|
- runs a single- or multi-agent CartPole experiment (for multi-agent, we use
|
|
different learning rates) thereby checkpointing the state of the Algorithm every n
|
|
iterations.
|
|
- stops the experiment due to an expected crash in the algorithm's main process
|
|
after a certain number of iterations.
|
|
- just for testing purposes, restores the entire algorithm from the latest
|
|
checkpoint and checks, whether the state of the restored algo exactly match the
|
|
state of the crashed one.
|
|
- then continues training with the restored algorithm until the desired final
|
|
episode return is reached.
|
|
|
|
|
|
How to run this script
|
|
----------------------
|
|
`python [script file name].py --num-agents=[0 or 2]
|
|
--stop-reward-crash=[the episode return after which the algo should crash]
|
|
--stop-reward=[the final episode return to achieve after(!) restoration from the
|
|
checkpoint]
|
|
`
|
|
|
|
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
|
|
-----------------
|
|
First, you should see the initial tune.Tuner do it's thing:
|
|
|
|
Trial status: 1 RUNNING
|
|
Current time: 2024-06-03 12:03:39. Total running time: 30s
|
|
Logical resource usage: 3.0/12 CPUs, 0/0 GPUs
|
|
╭────────────────────────────────────────────────────────────────────────
|
|
│ Trial name status iter total time (s)
|
|
├────────────────────────────────────────────────────────────────────────
|
|
│ PPO_CartPole-v1_7b1eb_00000 RUNNING 6 15.362
|
|
╰────────────────────────────────────────────────────────────────────────
|
|
───────────────────────────────────────────────────────────────────────╮
|
|
..._sampled_lifetime ..._trained_lifetime ...episodes_lifetime │
|
|
───────────────────────────────────────────────────────────────────────┤
|
|
24000 24000 340 │
|
|
───────────────────────────────────────────────────────────────────────╯
|
|
...
|
|
|
|
then, you should see the experiment crashing as soon as the `--stop-reward-crash`
|
|
has been reached:
|
|
|
|
```RuntimeError: Intended crash after reaching trigger return.```
|
|
|
|
At some point, the experiment should resume exactly where it left off (using
|
|
the checkpoint and restored Tuner):
|
|
|
|
Trial status: 1 RUNNING
|
|
Current time: 2024-06-03 12:05:00. Total running time: 1min 0s
|
|
Logical resource usage: 3.0/12 CPUs, 0/0 GPUs
|
|
╭────────────────────────────────────────────────────────────────────────
|
|
│ Trial name status iter total time (s)
|
|
├────────────────────────────────────────────────────────────────────────
|
|
│ PPO_CartPole-v1_7b1eb_00000 RUNNING 27 66.1451
|
|
╰────────────────────────────────────────────────────────────────────────
|
|
───────────────────────────────────────────────────────────────────────╮
|
|
..._sampled_lifetime ..._trained_lifetime ...episodes_lifetime │
|
|
───────────────────────────────────────────────────────────────────────┤
|
|
108000 108000 531 │
|
|
───────────────────────────────────────────────────────────────────────╯
|
|
|
|
And if you are using the `--as-test` option, you should see a finel message:
|
|
|
|
```
|
|
`env_runners/episode_return_mean` of 500.0 reached! ok
|
|
```
|
|
"""
|
|
import re
|
|
import time
|
|
|
|
from ray import tune
|
|
from ray.air.integrations.wandb import WandbLoggerCallback
|
|
from ray.rllib.algorithms.algorithm_config import AlgorithmConfig
|
|
from ray.rllib.callbacks.callbacks import RLlibCallback
|
|
from ray.rllib.examples.envs.classes.multi_agent import MultiAgentCartPole
|
|
from ray.rllib.examples.utils import add_rllib_example_script_args
|
|
from ray.rllib.policy.policy import PolicySpec
|
|
from ray.rllib.utils.metrics import (
|
|
ENV_RUNNER_RESULTS,
|
|
EPISODE_RETURN_MEAN,
|
|
)
|
|
from ray.rllib.utils.test_utils import check_learning_achieved
|
|
from ray.tune.registry import get_trainable_cls, register_env
|
|
|
|
parser = add_rllib_example_script_args(
|
|
default_reward=500.0, default_timesteps=10000000, default_iters=2000
|
|
)
|
|
parser.add_argument(
|
|
"--stop-reward-crash",
|
|
type=float,
|
|
default=200.0,
|
|
help="Mean episode return after which the Algorithm should crash.",
|
|
)
|
|
# By default, set `args.checkpoint_freq` to 1 and `args.checkpoint_at_end` to True.
|
|
parser.set_defaults(
|
|
checkpoint_freq=1,
|
|
checkpoint_at_end=True,
|
|
)
|
|
|
|
|
|
class CrashAfterNIters(RLlibCallback):
|
|
"""Callback that makes the algo crash after a certain avg. return is reached."""
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
# We have to delay crashing by one iteration just so the checkpoint still
|
|
# gets created by Tune after(!) we have reached the trigger avg. return.
|
|
self._should_crash = False
|
|
|
|
def on_train_result(self, *, algorithm, metrics_logger, result, **kwargs):
|
|
# We had already reached the mean-return to crash, the last checkpoint written
|
|
# (the one from the previous iteration) should yield that exact avg. return.
|
|
if self._should_crash:
|
|
raise RuntimeError("Intended crash after reaching trigger return.")
|
|
# Reached crashing criterion, crash on next iteration.
|
|
elif result[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN] >= args.stop_reward_crash:
|
|
print(
|
|
"Reached trigger return of "
|
|
f"{result[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN]}"
|
|
)
|
|
self._should_crash = True
|
|
|
|
|
|
if __name__ == "__main__":
|
|
args = parser.parse_args()
|
|
|
|
register_env(
|
|
"ma_cart", lambda cfg: MultiAgentCartPole({"num_agents": args.num_agents})
|
|
)
|
|
|
|
# Simple generic config.
|
|
config = (
|
|
get_trainable_cls(args.algo)
|
|
.get_default_config()
|
|
.environment("CartPole-v1" if args.num_agents == 0 else "ma_cart")
|
|
.env_runners(create_env_on_local_worker=True)
|
|
.training(lr=0.0001)
|
|
.callbacks(CrashAfterNIters)
|
|
)
|
|
|
|
# Tune config.
|
|
# Need a WandB callback?
|
|
tune_callbacks = []
|
|
if args.wandb_key:
|
|
project = args.wandb_project or (
|
|
args.algo.lower() + "-" + re.sub("\\W+", "-", str(config.env).lower())
|
|
)
|
|
tune_callbacks.append(
|
|
WandbLoggerCallback(
|
|
api_key=args.wandb_key,
|
|
project=args.wandb_project,
|
|
upload_checkpoints=False,
|
|
**({"name": args.wandb_run_name} if args.wandb_run_name else {}),
|
|
)
|
|
)
|
|
|
|
# Setup multi-agent, if required.
|
|
if args.num_agents > 0:
|
|
config.multi_agent(
|
|
policies={
|
|
f"p{aid}": PolicySpec(
|
|
config=AlgorithmConfig.overrides(
|
|
lr=5e-5
|
|
* (aid + 1), # agent 1 has double the learning rate as 0.
|
|
)
|
|
)
|
|
for aid in range(args.num_agents)
|
|
},
|
|
policy_mapping_fn=lambda aid, *a, **kw: f"p{aid}",
|
|
)
|
|
|
|
# Define some stopping criterion. Note that this criterion is an avg episode return
|
|
# to be reached. The stop criterion does not consider the built-in crash we are
|
|
# triggering through our callback.
|
|
stop = {
|
|
f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": args.stop_reward,
|
|
}
|
|
|
|
# Run tune for some iterations and generate checkpoints.
|
|
tuner = tune.Tuner(
|
|
trainable=config.algo_class,
|
|
param_space=config,
|
|
run_config=tune.RunConfig(
|
|
callbacks=tune_callbacks,
|
|
checkpoint_config=tune.CheckpointConfig(
|
|
checkpoint_frequency=args.checkpoint_freq,
|
|
checkpoint_at_end=args.checkpoint_at_end,
|
|
),
|
|
stop=stop,
|
|
),
|
|
)
|
|
tuner_results = tuner.fit()
|
|
|
|
# Perform a very quick test to make sure our algo (upon restoration) did not lose
|
|
# its ability to perform well in the env.
|
|
# - Extract the best checkpoint.
|
|
metric = f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}"
|
|
best_result = tuner_results.get_best_result(metric=metric, mode="max")
|
|
assert (
|
|
best_result.metrics[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN]
|
|
>= args.stop_reward_crash
|
|
)
|
|
# - Change our config, such that the restored algo will have an env on the local
|
|
# EnvRunner (to perform evaluation) and won't crash anymore (remove the crashing
|
|
# callback).
|
|
config.callbacks(None)
|
|
# Rebuild the algorithm (just for testing purposes).
|
|
test_algo = config.build()
|
|
# Load algo's state from best checkpoint.
|
|
test_algo.restore(best_result.checkpoint)
|
|
# Perform some checks on the restored state.
|
|
assert test_algo.training_iteration > 0
|
|
# Evaluate on the restored algorithm.
|
|
test_eval_results = test_algo.evaluate()
|
|
assert (
|
|
test_eval_results[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN]
|
|
>= args.stop_reward_crash
|
|
), test_eval_results[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN]
|
|
# Train one iteration to make sure, the performance does not collapse (e.g. due
|
|
# to the optimizer weights not having been restored properly).
|
|
test_results = test_algo.train()
|
|
assert (
|
|
test_results[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN] >= args.stop_reward_crash
|
|
), test_results[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN]
|
|
# Stop the test algorithm again.
|
|
test_algo.stop()
|
|
|
|
# Create a new Tuner from the existing experiment path (which contains the tuner's
|
|
# own checkpoint file). Note that even the WandB logging will be continued without
|
|
# creating a new WandB run name.
|
|
restored_tuner = tune.Tuner.restore(
|
|
path=tuner_results.experiment_path,
|
|
trainable=config.algo_class,
|
|
param_space=config,
|
|
# Important to set this to True b/c the previous trial had failed (due to our
|
|
# `CrashAfterNIters` callback).
|
|
resume_errored=True,
|
|
)
|
|
# Continue the experiment exactly where we left off.
|
|
tuner_results = restored_tuner.fit()
|
|
|
|
# Not sure, whether this is really necessary, but we have observed the WandB
|
|
# logger sometimes not logging some of the last iterations. This sleep here might
|
|
# give it enough time to do so.
|
|
time.sleep(20)
|
|
|
|
if args.as_test:
|
|
check_learning_achieved(tuner_results, args.stop_reward, metric=metric)
|