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