Making VLA Truly Run on Robots: PKU Team Open-Sources Jetson-PI, Boosting Jetson Orin Control Frequency by 8.66x

Deploying VLA models on onboard devices faces dual challenges. On one hand, relying on external high-end GPUs significantly increases power consumption and reduces battery life, limiting the robot's range of motion; on the other hand, running directly on edge devices like Jetson Orin, π0.5 takes about 1.4 seconds per inference, resulting in a control frequency of only 0.7Hz, causing sluggish robot movements and noticeable pauses.

Naive asynchronous inference can hide some latency but introduces two new problems: perception-action misalignment (the environment changes during model inference) and excessive reaction time (still limited by full inference latency). Previous methods attempt to fuse old and new trajectories, but may still fail in complex tasks due to action errors.

Jetson-PI proposes Foresight-Aligned Asynchronous Correction, training a lightweight future correction module with about 40M parameters. Based on the current VLM environment representation and submitted actions, this module predicts the environment representation at a future time, allowing the action expert to generate actions starting from a future time point rather than based on outdated observations. The module only predicts the last layer representation of VLM, avoiding the high cost of generating full images.

The confidence-based scheduling mechanism lets the system decide when to skip VLM: the future correction module also outputs confidence; when above a threshold, directly invoke the action expert (faster inference); when below, re-invoke VLM to refresh the environment. In actual tests, confidence drops during key steps like grasping and placing, triggering VLM updates, while remaining high during smooth actions.

The research team developed Jetson-PI-Edge based on llama.cpp, with three optimizations tailored for VLA: computation graph reuse (build CUDA Graph once for fixed input sizes and reuse), persistent GPU intermediate cache (avoid data transfer between ViT, LLM, and action expert via CPU), and flow matching unrolling (merge multiple denoising steps into a unified computation graph).

After optimization, π0.5 single inference time on Jetson Orin drops to 412.9ms, and action expert inference from 536.8ms to 123.1ms. Combined with confidence scheduling, system reaction time further reduces to 165.1ms, achieving a control frequency of 6.06Hz, an 8.66x improvement over native PyTorch, without using quantization or pruning.

On four LIBERO sub-datasets, Jetson-PI's average success rate is 14.8 percentage points higher than VLASH and 3.9 points higher than RTC. The advantage grows with larger latency: when latency steps Δ=9, Jetson-PI's average success rate is 45.6 points higher than VLASH and 7.0 points higher than RTC, indicating that future environment prediction is crucial under long delays.

In cloth folding tasks on a PrimeBot X2-W robot, Jetson-PI achieves smoother and more continuous actions, avoiding pauses from synchronous inference and misalignment from naive asynchronous inference. The robot is equipped with one head camera and two wrist cameras, executing actions at 15Hz.