How to export model with tune.run?

Ray Libraries (Data, Train, Tune, Serve) Ray Tune
1

I am using tune.run and specified export_formats=["model", "checkpoint"], with a valid local_dir
PPOTrainer is from from ray.rllib.agents.ppo import PPOTrainer

analysis = tune.run(
        PPOTrainer, # uses PPO algorithm
        name="lets_play_tag",
        stop=TimeStopper(),
        checkpoint_freq=5,
        checkpoint_at_end=True,
        resume=args.resume_training,
        local_dir=local_dir,
        reuse_actors=True,
        max_failures=1,
        export_formats=["model", "checkpoint"],
        config=tune_config,
        scheduler=pbt,
    )

I only found checkpoint in the export directory with no model folder. There are no models in the checkpoint folder either.
How do I export the trained model? Thanks!

image
image1577×900 75.3 KB

I am aware that I can export a model when using the trainable like this:


from ray.rllib.agents.ppo import PPOTrainer
import os
import ray

from ray.rllib.agents.registry import get_agent_class
from ray.rllib.utils.framework import try_import_tf

tf1, tf, tfv = try_import_tf()

ray.init(num_cpus=10)

training = PPOTrainer(config={}, env="CartPole-v0")
for _ in range(3):
    training.train()

prefix = "model.ckpt"
model_dir = os.path.join("model_export_dir")
training.export_policy_checkpoint('model', filename_prefix=prefix)
training.export_policy_model(model_dir)

But I am not sure how to export the model when using tune.run to run the trainable

2

Hm, I’ll cc @sven1977 for RLLib related questions

3

Hi @Jiffer_PengPeng,

You can use a custom trainer function or trainable class which can call the export_policy_model function when you need to.

The other option is to use the callbacks, with something like:

class Callbacks(DefaultCallbacks    ):
        def on_train_result(self, trainer, result, **kwargs):
            trainer.export_policy_model(model_dir)

I couldn’t find any way to get tune.run to automatically call export_model on the trainer like you can do with checkpoint_freq