Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Adds onnx export to cleanrl example #161

Merged
merged 3 commits into from
Jan 18, 2024
Merged

Conversation

Ivan-267
Copy link
Collaborator

@Ivan-267 Ivan-267 commented Dec 17, 2023

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:

image

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.

@Ivan-267 Ivan-267 added the enhancement New feature or request label Dec 18, 2023
Copy link
Owner

@edbeeching edbeeching left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, thanks. I think you will need to merge/rebase on main to make the tests green. Then merge if they pass.

Adds basic information about how to use the script, removes the code from the doc as the python file is linked.
@Ivan-267 Ivan-267 force-pushed the add_clean_rl_onnx_export branch from 152e437 to 034ce4f Compare January 18, 2024 05:58
@Ivan-267 Ivan-267 merged commit cfaaf15 into main Jan 18, 2024
12 checks passed
@Ivan-267 Ivan-267 deleted the add_clean_rl_onnx_export branch January 18, 2024 06:23
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants