How to save RLlib model as Onnx

RLlib
1

After training the model with tune, how to save it as ONNX model ?

2

Hey @geekyneuro , we are actually working on an example script for this right now.

In general, you can get to the model after tune training like so:

results = tune.run(..., checkpoint_at_end=True)

best_checkpoint = results.get_best_checkpoint(
    results.trials[0], mode="max")
print(f".. best checkpoint was: {best_checkpoint}")

new_trainer = PPOTrainer(config=config)
new_trainer.restore(best_checkpoint)
policy = new_trainer.get_policy([policy ID or empty for "default_policy (single agent)"])
model = policy.model

# For torch, `model` is now a torch.nn.Module; for tf, try `model.base_model` to get the tf.keras.Model.
3

Thanks for sharing the sample code. Once the access to model object has been established, the dummy input is required to save it as onnx. How to figure out the dummy input ? Would the dummy input be same as the output of step function from the env ? Or would it be in some pre-processed format ?

Sincerely

Armando

4

I know this thread is probably too old, however for anyone looking for such answer >>

dummy input should be in size of observation space.

here is an example.

observation_size = model.observation_space.shape
dummy_input = torch.randn(1, *observation_size)