[Question] Need help on exporting the final Policy and ONNX file after training using Robomimic

[cyborgsince90s]

Question

Hi,

I have trained a pick and place policy using Robomimic and have got the Final epoch outputs in "PTH" format.

Now similar to Reinforcement learning,

1. How do I get the Exported Policy file and ONNX file for physical deployment?
2. I have a locomotion policy trained and working, I want to combine the Locomotion and Pick and Place policy together - Please let me know the process for accomplishing this.

Thanks,

[RandomOakForest]

Thank you for posting this. I'll move to our Discussions for follow up. Here is a summary of guidelines that may be of help.

Exporting Your Robomimic Policy to ONNX

1. Load the PTH checkpoint and script model:
   Using PyTorch, load your final epoch .pth checkpoint. If the policy contains custom layers or wrappers, ensure the forward() is compatible with ONNX tracing.
2. Export to ONNX using PyTorch:
   Use the standard PyTorch ONNX exporter. Example:

import torch

# Assume your policy network class is PolicyNetwork and checkpoint path is known
model = PolicyNetwork(...)
model.load_state_dict(torch.load('path_to_your_final_policy.pth'))
model.eval()

# Create a dummy input that matches the input shape your policy expects
dummy_input = torch.randn(1, input_dim)  # use correct shape and device
torch.onnx.export(
    model,
    dummy_input,
    "exported_policy.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}},
    opset_version=13
)

This produces an ONNX file that you can use for deployment via inference infrastructure like Isaac Sim or real-robot middleware.¹²³
3. Isaac Lab to Isaac Sim Workflow:
If you used Isaac Lab for training, replay in Isaac Lab to generate an exported folder with the policy and supporting files, then migrate these to Isaac Sim as described in the deployment demo.⁴⁵

Combining the Locomotion and Pick-and-Place Policies

• Policy Composition:
  • Sequential Switching: If the locomotion and pick-and-place policies correspond to distinct phases, use a higher-level task manager (state machine or policy orchestrator) that selects and feeds observations to the correct policy at runtime.
  • Input/Output Coordination: Ensure that the observation provided for each policy matches the expected input shape and semantics; you may need to combine or split the robot's observation vector appropriately.
• Deployment in Isaac Sim:
  • Use the helper functions from the isaac_sim_policy_example extension to load and run multiple policies, as shown in the latest Isaac Sim reinforcement learning policies examples and demos.⁶
  • Reference the unitree or Spot policy deployment demos for structure; adapt to support two ONNX models and selection logic, or merge their outputs if parallel behaviors are desired.⁵⁶

Footnotes

1. https://developer.nvidia.com/blog/end-to-end-ai-for-pcs-transitioning-ai-models-with-onnx/ ↩

2. https://docs.pytorch.org/tutorials/beginner/onnx/export_simple_model_to_onnx_tutorial.html ↩

3. https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/examples/onnx/onnx_export.html ↩

4. https://www.youtube.com/watch?v=Cx9ss0v-DMs ↩

5. https://docs.isaacsim.omniverse.nvidia.com/4.5.0/isaac_lab_tutorials/tutorial_policy_deployment.html ↩ ↩²

6. https://docs.isaacsim.omniverse.nvidia.com/5.1.0/robot_simulation/ext_isaacsim_robot_policy_example.html ↩ ↩²