Feature Request: ONNX Export Support for Cloud VLM Backbones

[jian-dong]

Hi.
First, thank you for your excellent work and for open-sourcing this repository. Your tooling for VLA training and the existing ONNX export support for the action decoder is very valuable for deployment workflows.

I noticed that the current ONNX export script focuses on the action decoder component. Could you share whether there are any plans to extend ONNX export support to include Cloud VLM backbone models as well? This would greatly simplify inference deployment in backend systems using TensorRT-LLM for VLMs alongside TensorRT for the action decoder.

Specifically:

• Is there an intended roadmap for Cloud VLM → ONNX export?
• Would you recommend any interim approach for exporting the VLM backbone to ONNX for TensorRT-LLM-based deployment?

Thank you for your time and contributions

[shuheng-liu]

Hi @jian-dong

Yes, we will support exporting the whole model to ONNX in the near future. This helps people deploy their policy on real robots in TensorRT format.

Before that happens, we recommend running torch.onnx.export() manually on the whole policy. You may need some custom monkey patches to make ONNX happy. Let me know if you encounter any issues while doing that.