Jetson-PI

Towards Onboard Real-Time Robot Control via Foresight-Aligned Asynchronous Inference

Foresight-aligned asynchronous inference for responsive, real-time Vision-Language-Action control on edge robots.

Paper · Demo · Results · Training · Evaluation · Edge Runtime · Citation

---

Overview

This repository is the official implementation of:

Jetson-PI: Towards Onboard Real-Time Robot Control via Foresight-Aligned Asynchronous Inference
Zebin Yang, Qi Wang, Yunhe Wang, Xiurui Guo, Bo Yu, Shaoshan Liu, Jiafeng Xu, Hao Dong, and Meng Li.
arXiv:2607.12659, 2026.

Vision-Language-Action (VLA) models have achieved impressive performance on diverse embodied tasks, yet deploying them on low-power onboard devices such as NVIDIA Jetson Orin remains challenging due to high inference latency and limited compute. Asynchronous inference can partially mask this latency, but it introduces prediction–execution misalignment and long reaction time. Jetson-PI addresses both through Foresight-Aligned Asynchronous Correction (FAAC): we train a lightweight future correction module that predicts future environment representation conditioned on committed actions, enabling the action expert to directly predict actions from the future time step; we further introduce confidence-based scheduling optimization that adaptively balances VLM and action expert invocations.

This release open-sources LIBERO training and evaluation code built on π₀.₅. The accelerated llama.cpp-based onboard inference engine is available in PKU-SEC-Lab/Jetson-PI-Edge.

Real-world Demo

▶ Full video (mp4)

Results on LIBERO (π₀.₅)

Success rate comparison across four LIBERO sub-datasets. We report SR of Foresight-Aligned Asynchronous Correction alone (Ours) and with Confidence-based Scheduling Optimization (+Sched). Inference time of the action expert is estimated as $\Delta_{ae}=\lceil\Delta/3\rceil$, a common ratio on both high-end GPUs (e.g. RTX 4090) and onboard GPUs (e.g. Orin). We train one model for all $\Delta$ values.

	SPATIAL				OBJECT				GOAL				LIBERO-10
Sync.	97.3				99.6				96.7				93.5
Δ	VLASH	RTC	Ours	+Sched	VLASH	RTC	Ours	+Sched	VLASH	RTC	Ours	+Sched	VLASH	RTC	Ours	+Sched
1	98.8	97.1	97.7	98.6	99.2	98.5	98.7	98.8	96.7	96.5	97.0	97.5	94.4	92.3	93.4	93.1
2	97.5	95.4	97.1	97.6	99.2	98.5	98.5	99.3	97.0	96.0	96.2	97.3	94.6	92.1	92.7	92.9
3	94.4	94.5	97.1	97.2	98.8	96.4	98.2	98.4	93.3	94.3	96.7	96.4	91.9	89.6	91.3	91.9
4	92.5	92.5	96.7	96.9	96.9	97.7	97.8	98.9	93.3	93.3	95.9	96.9	89.6	85.8	91.8	92.3
5	84.3	91.2	95.9	96.7	94.3	96.9	96.9	98.6	89.9	93.9	96.0	96.1	80.3	83.6	91.5	92.0
6	74.4	91.7	97.0	97.5	88.5	97.4	98.3	98.4	81.3	94.0	96.7	96.9	78.5	84.6	92.2	92.9
7	51.3	91.7	97.0	97.3	81.4	97.1	98.0	98.4	76.1	92.6	97.1	97.2	77.0	85.3	92.6	92.6
8	46.7	90.9	97.2	97.5	65.1	93.8	98.3	98.8	70.8	93.4	96.8	96.8	65.5	83.2	92.3	92.7
9	30.1	88.8	97.0	97.3	51.7	93.1	97.5	97.6	59.7	92.5	96.3	96.5	59.6	81.0	92.0	92.2
Avg	74.4	92.6	97.0	97.4	86.1	96.6	98.0	98.6	84.2	94.1	96.5	96.8	81.3	86.4	92.2	92.5

---

Requirements

Item	Recommendation
OS	Ubuntu 22.04
GPU	NVIDIA GPU with ≥ 48 GB VRAM for full three-stage training (batch 16)
Python	3.11 for training (uv / JAX)
CUDA	CUDA 12.x (installed via project dependencies; no system CUDA required)

---

Environment Setup

1. Clone and submodules

git clone --recurse-submodules <this-repo-url>
cd <repo>
git submodule update --init --recursive

2. Training environment (JAX)

Install with uv:

export PYTHONNOUSERSITE=1
GIT_LFS_SKIP_SMUDGE=1 uv sync
GIT_LFS_SKIP_SMUDGE=1 uv pip install -e .

Pin versions are listed in training_requirements.txt (Python 3.11, JAX 0.5.3, etc.).

Apply the transformers patch required for π₀.₅ PyTorch/JAX compatibility:

cp -r ./src/openpi/models_pytorch/transformers_replace/* .venv/lib/python3.11/site-packages/transformers/

3. LIBERO eval client (simulation)

The policy server uses the training venv (PY_SERVER). The LIBERO client can use examples/libero/.venv:

uv venv --python 3.11 examples/libero/.venv
source examples/libero/.venv/bin/activate
uv pip sync examples/libero/requirements.txt third_party/libero/requirements.txt \
  --extra-index-url https://download.pytorch.org/whl/cu121 --index-strategy=unsafe-best-match
uv pip install -e packages/openpi-client
uv pip install -e third_party/libero

Eval-only dependencies are listed in test_requirements.txt.

4. Paths and checkpoints

Download checkpoints

Pretrained π₀.₅-LIBERO + future correction module (LIBERO-spatial, step 65000) are hosted on:

Hugging Face: diantoudefengshan/Jetson-PI-pi05

ModelScope: zebinyang/Jetson-PI-pi05

pip install modelscope
python -c "from modelscope import snapshot_download; snapshot_download('zebinyang/Jetson-PI-pi05', local_dir='./checkpoints/jetson-pi-pi05')"
export PI0_CHECKPOINT=./checkpoints/jetson-pi-pi05/pi05_libero
export WM=./checkpoints/jetson-pi-pi05/future_correction_module

The bundle contains two separate directories (pi05_libero/, future_correction_module/); do not merge their params/ trees.

Set these before training or evaluation:

Variable	Description
PI0_CHECKPOINT	π₀.₅-LIBERO weights (local dir with params/). Use pi05_libero/ from ModelScope, or download upstream: gs://openpi-assets/checkpoints/pi05_libero
WM	Eval only. Path to trained future correction module dir (must contain params/). Use future_correction_module/ from ModelScope
OPENPI_LIBERO_LOCAL_DATASET_DIR	LeRobot LIBERO dataset root (parquet + meta/tasks.jsonl)
PY	Python for training (JAX venv)
PY_SERVER	Python for serve_policy (must have JAX; often same as PY)

Example:

export PI0_CHECKPOINT=PATH/TO/CHECKPOINT/pi05_libero
export OPENPI_LIBERO_LOCAL_DATASET_DIR=PATH/TO/DATASET/libero
export PY=PATH/TO/PYTHON
export PY_SERVER=PATH/TO/PYTHON
export PYTHONNOUSERSITE=1

---

Training

Recipe (current default)

Three-stage schedule on π₀.₅-LIBERO:

Stage	Steps	What is trained
1	30,000	Action Expert + token reducer (L_act)
2	15,000	Future correction module (L_cond, no logvar head)
3	55,000	L_cond on future correction module (no reducer) + L_act on Pi0 AE + full LLM (μ detached)

Fixed handover H = 10, max_delta_t = 10, action_encoder = transformer_block.

Checkpoints: checkpoints/<EXP_NAME>/ (Orbax) and checkpoints/<EXP_NAME>/world_model_step_<N>/.

Launch

Use the dedicated launcher (configure paths first):

cd PATH/TO/REPO
export PI0_CHECKPOINT=PATH/TO/CHECKPOINT/pi05_libero
export OPENPI_LIBERO_LOCAL_DATASET_DIR=PATH/TO/DATASET/libero
export PY=PATH/TO/PYTHON
export CUDA_VISIBLE_DEVICES=0

bash scripts/train_wm_libero_spatial_four_stage.sh

Optional overrides:

export STAGE1_STEPS=30000
export STAGE2_STEPS=15000
export STAGE3_STEPS=55000
export BATCH_SIZE=16
export NUM_WORKERS=4
export EXP_NAME=my_wm_spatial_run
bash scripts/train_wm_libero_spatial_four_stage.sh

Logs: logs/<EXP_NAME>.log.

---

Evaluation

Evaluation runs serve_policy.py (future correction module + π₀.₅ action expert) and examples/libero/main.py (LIBERO sim). Set WM to the trained future correction module checkpoint directory.

Single run (FAAC + async action expert)

cd PATH/TO/REPO
export PI0_CHECKPOINT=PATH/TO/CHECKPOINT/pi05_libero
export PY_SERVER=PATH/TO/PYTHON
export WM=PATH/TO/future-correction-module
export CUDA_VISIBLE_DEVICES=0
export PORT=8000             # pick a free port

bash scripts/eval_wm_libero_spatial.sh

Alternatively, if the checkpoint lives under checkpoints/<EXP_NAME>/world_model_step_<N>/:

export EXP_NAME=<your_training_exp_name>
export STEP=<N>
bash scripts/eval_wm_libero_spatial.sh

Outputs under logs/<run_dir>/: serve.log, client.log, run_meta.txt, videos/.

Default: libero_spatial, 50 trials/task, H=10, K=9, overlap=1.

Adaptive multi-rollout (confidence-based scheduling)

export LIBERO_WM_EVAL_ADAPTIVE_KAPPA=1
export LIBERO_WM_EVAL_KAPPA_DELTA=0.4
bash scripts/eval_wm_libero_spatial.sh

Other LIBERO suites

export LIBERO_WM_EVAL_TASK_SUITE=libero_object   # or libero_goal, libero_10
bash scripts/eval_wm_libero_spatial.sh

K-sweep (advanced)

For sweeping trigger step K from 9 down to 1 with adaptive kappa:

export WM=PATH/TO/future-correction-module
export PI0_CHECKPOINT=PATH/TO/CHECKPOINT/pi05_libero
export PY_SERVER=PATH/TO/PYTHON
export LIBERO_WM_EVAL_KAPPA_DELTA=0.4
bash scripts/libero_wm_eval_spatial_k9to1_adaptive_kappa_low_replan_gpu2_kd0p4.sh

---

Troubleshooting

Issue	Fix
ModuleNotFoundError: No module named 'jax' in eval	Set PY_SERVER to the JAX training venv Python, not system Python
OOM during training	Lower BATCH_SIZE, set XLA_PYTHON_CLIENT_MEM_FRACTION=0.85, or NUM_WORKERS=0
Missing norm_stats in eval	Point --pi0-norm-checkpoint-dir / PI0_CHECKPOINT to a tree containing assets/physical-intelligence/libero/norm_stats.json
LIBERO EGL / display errors	Install xvfb; eval script falls back to MUJOCO_GL=egl if xvfb is missing
Checkpoint save killed (no traceback)	First Orbax save can spike host RAM; see docs/docker.md and reduce save frequency for smoke tests

---

License

See LICENSE and LICENSE_GEMMA.txt. LIBERO and upstream openpi components retain their respective licenses.

Citation

If Jetson-PI helps your research, please cite our paper:

@article{yang2026jetson,
  title={Jetson-PI: Towards Onboard Real-Time Robot Control via Foresight-Aligned Asynchronous Inference},
  author={Yang, Zebin and Wang, Qi and Wang, Yunhe and Guo, Xiurui and Yu, Bo and Liu, Shaoshan and Xu, Jiafeng and Dong, Hao and Li, Meng},
  journal={arXiv preprint arXiv:2607.12659},
  year={2026}
}

Acknowledgments

Jetson-PI builds on OpenPI and the π model family from Physical Intelligence, and uses LIBERO for simulation evaluation. The onboard inference engine is developed in Jetson-PI-Edge, based on llama.cpp, with integration support for FlashRT.