ray restore checkpoint in rllib

user_5228890 0
Viewed 1281

Ray saves a bunch of checkpoints during a call of agent.train(). How do I know which one is the checkpoint with the best agent to load? Is there any function like tune-analysis-output.get_best_checkpoint(path, mode="max") to explore different loading possibilities over the checkpoints?

ray
2 Answers

As answered in https://discuss.ray.io/t/ray-restore-checkpoint-in-rllib/3186/2 you can use:

analysis = tune.Analysis(experiment_path)  # can also be the result of `tune.run()`

trial_logdir = analysis.get_best_logdir(metric="metric", mode="max")  # Can also just specify trial dir directly

checkpoints = analysis.get_trial_checkpoints_paths(trial_logdir)  # Returns tuples of (logdir, metric)
best_checkpoint = analysis.get_best_checkpoint(trial_logdir, metric="metric", mode="max")

See https://docs.ray.io/en/master/tune/api_docs/analysis.html#id1

user_152626410 
analysis = tune.run( 
    "A2C", 
    name = model_name,
    config = config, 
    ...
    checkpoint_freq = 5,
    checkpoint_at_end = True,
    restore = best_checkpoint 
    )



trial_logdir = analysis.get_best_logdir(metric="episode_reward_mean", mode="max")  
best_checkpoint = analysis.get_best_checkpoint(trial_logdir, metric="episode_reward_mean", mode="max")
user_71069150
Related