PBT Replay with RLlib

Hi everyone,

I’m having some problems to replay a PBT training. I’ve used the following code for the tuning:

tuner = tune.Tuner(
    "PPO",
    tune_config=tune.TuneConfig(
        metric="episode_reward_mean",
        mode="max",
        scheduler=pbt,
        num_samples=100,
    ),
    param_space=config_PPO,
    run_config=air.RunConfig(
        stop={"training_iteration": 200},
    )
)
results = tuner.fit()

I’m using a custom environment that I introduced directly as a class in PPOConfig with an env_config. For the PopulationBasedTrainingReplay a trainable class is needed, and using only “PPO” as in the above code is not an option. There is some examples on how to PBT Replay a generic Pytorch model, but I haven’t found any related to RLlib algorithms. Is there any example on how to define a trainable class for this same purpose? I would really appreciate any recommendation.