Adds onnx export to cleanrl example #161
Merged
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Adds support for exporting the trained model to onnx using the command-line argument
--onnx_export_path
.It works with a single observation space, and here is how the output onnx looks in Netron, compared to a file exported using the sb3 example script:
I assume the difference in flattening is due to the CleanRL export using a single observation space.
Quick test using:
python clean_rl_example.py --total-timesteps=50_000 --onnx_export_path=model_clean_rl.onnx
cleanrl_onnx_inference.mp4
The additional code for exporting is implemented in the same file, however it's still using a separate class similar to the SB3 export, so that the original agent class is not modified. The new class is added to implement the
forward
method, which in this case just returns the mean action and state_ins to keep compatibility with our onnx inference code.Note: Even though the implementation seems to work after a quick test, as I'm not very familiar with Torch/Numpy/Neural nets, please check the implementation in case I overlooked something.