Ticket Type
Feature Request / Improvement
Description
LeRobot covers the full pipeline from data collection to training to async inference, but there's currently no built-in way to export a trained policy (ACT, Diffusion Policy, etc.) to optimized formats like ONNX or TensorRT for edge deployment.
This is a common pain point for anyone deploying on resource-constrained hardware (Jetson Orin Nano, etc.) where PyTorch inference is too slow or memory-heavy. Right now the only option is D-Robotics' hardware-specific rdk_LeRobot_tools, but there's nothing general-purpose.
Proposed scope
A lerobot-export CLI command (or equivalent API) that:
- Loads a trained policy checkpoint (local path or Hub repo)
- Exports to ONNX with correct dynamic axes for batch size / sequence length
- Optionally optimizes with TensorRT (FP16/INT8) when available
- Validates correctness by comparing PyTorch vs exported model outputs on sample inputs
Example usage:
lerobot-export \
--policy.path=lerobot/act_pusht \
--format=onnx \
--output=act_pusht.onnx
# With TensorRT optimization
lerobot-export \
--policy.path=lerobot/act_pusht \
--format=tensorrt \
--precision=fp16 \
--output=act_pusht.engine
Why this matters
- The SO-101 + Jetson Orin Nano combo is increasingly popular and PyTorch inference can be a bottleneck
- ONNX export enables deployment across diverse hardware (Jetson, Intel NPUs, Qualcomm, etc.) without vendor lock-in
- TensorRT typically gives 2-5x inference speedup over vanilla PyTorch on NVIDIA GPUs
- This would complement the existing async inference stack by making the policy server itself faster
What I've looked into
torch.onnx.export works for ACT and Diffusion Policy architectures with some attention to dynamic axes- The main challenge is handling policy-specific preprocessing/postprocessing (action chunking, denoising loops for diffusion, etc.)
- VLA models (Pi0, SmolVLA) are harder due to size but ONNX export is still feasible; TensorRT would require more work
Happy to put together a PR if this is something the team would want upstream. Would also appreciate guidance on which policies to prioritize.
Ticket Type
Feature Request / Improvement
Description
LeRobot covers the full pipeline from data collection to training to async inference, but there's currently no built-in way to export a trained policy (ACT, Diffusion Policy, etc.) to optimized formats like ONNX or TensorRT for edge deployment.
This is a common pain point for anyone deploying on resource-constrained hardware (Jetson Orin Nano, etc.) where PyTorch inference is too slow or memory-heavy. Right now the only option is D-Robotics' hardware-specific rdk_LeRobot_tools, but there's nothing general-purpose.
Proposed scope
A
lerobot-exportCLI command (or equivalent API) that:Example usage:
lerobot-export \ --policy.path=lerobot/act_pusht \ --format=onnx \ --output=act_pusht.onnx # With TensorRT optimization lerobot-export \ --policy.path=lerobot/act_pusht \ --format=tensorrt \ --precision=fp16 \ --output=act_pusht.engineWhy this matters
What I've looked into
torch.onnx.exportworks for ACT and Diffusion Policy architectures with some attention to dynamic axesHappy to put together a PR if this is something the team would want upstream. Would also appreciate guidance on which policies to prioritize.