APXinf-robo — Edge Inference Engine for VLA/WAM Models on Jetson Thor / Orin

TL;DR

APXinf-robo (RLinf / Infinigence) is an open-source Rust-native edge inference engine for VLA and World-Action models targeting Jetson AGX Thor / Thor-U / Orin and desktop GPUs (RTX 4090), shipping first with a highly-optimized π0.5 implementation across BF16, FP8 (Thor-only), and INT8 (Orin/Ada) with an OpenPI-compatible websocket server as a drop-in replacement for the standard π0.5 serving stack. Headline numbers: 41.16 ms / 24.3 Hz on Jetson Thor FP8 vs 72.45 ms / 13.8 Hz BF16 (1.76× speedup) for the two-view π0.5 workload, LIBERO-10 task success within noise of the published pi05_libero_base reference (92.8% Thor BF16, 92.2% Thor FP8, 92.0% Orin BF16 vs 92.4% reference across 500 episodes / 10 tasks × 50). The release bundles an agent-driven model-port-workflow skill for Claude Code / Codex that automates porting new VLAs (GR00T N1.7 named as the reference example) to the engine, positioning APXinf as both a deployment runtime and an agentically-portable substrate. Same lab as Harness VLA: Steering Frozen VLAs into Reliable Manipulation Primitives via Memory-Guided Agents Harness VLA — the RLinf group is now shipping the serving half of its own agentic-VLA thesis.

Key claims

Rust-native inference engine with no external ML dependencies: “implemented with system language Rust with no other external dependencies”, positioned as the model-plane counterpart to LLM-serving stacks like FasterTransformer / TensorRT-LLM / vLLM / SGLang but scoped to embodied AI [README §Header].
Jetson Thor FP8 is the fastest supported path: 41.16 ms P50 / 24.3 Hz on the two-view / 10-flow-step / H=10 / batch-1 π0.5 workload, vs 72.45 ms / 13.8 Hz BF16 on the same hardware — a 1.76× latency reduction from BF16 to FP8 [README §Performance].
Onestep action-generation pruning adds another ~1.6× on top: FP8 Thor drops to 26.32 ms / 38.0 Hz and BF16 RTX 4090 to 20.36 ms / 49.1 Hz with onestep pruning enabled [README §Performance, second table].
LIBERO-10 task success is preserved across precisions: 500-episode evaluation (10 tasks × 50) reports 92.8% Thor BF16, 92.2% Thor FP8, 92.0% Orin BF16 vs pi05_libero_base’s published 92.4% — the quantization + acceleration path does not degrade closed-loop policy quality within evaluation noise [README §Performance, LIBERO table].
OpenPI-compatible websocket serving protocol: an unmodified openpi-client connects to apxinf-robo serve without code changes, making APXinf a drop-in replacement for the reference π0.5 serving stack [README §Serve it with OpenPI compatible websocket server].
Three layered Python APIs (L1 / L2 / L3) factor deployment: L1 load_bare_model takes preprocessed inputs and returns normalized actions; L2 build_robot_policy / load_policy adds robot presets (franka_libero, unitree_g1) with per-preset image/state/action encoding; L3 websocket_server wraps L2 in the OpenPI wire protocol [README §Using APXinf-robo from Python].
INT8 W8A8 supports Orin (SM87) and Ada (SM89) with no calibration required beyond the checkpoint, while FP8 requires per-tensor activation scales calibrated from representative observations (fallback to <model-dir>/calibration.json) [README §Precisions].
Compute-capability-locked build: the CUDA kernels compile for exactly the visible GPU’s compute capability, so builds are hardware-specific — cross-compilation requires APXINF_CUDA_ARCH naming the target (sm_87 Orin, sm_101 Thor-U, sm_110 Thor) [README §Build APXinf-robo].
Agent-driven model-port workflow: ships a Claude Code / Codex skill at apxinf/skills/model-port-workflow that drives a full port sequence (survey → implement → validate → parity check) against documented porting guides in the sibling infinigence/ApxInf engine repo, with GR00T N1.7 named as a worked example [README §Port a new model with an agent].

Method

APXinf-robo is a thin Python + Rust package layered on top of an engine submodule at apxinf/ (the infinigence/ApxInf core), so the release is really two pieces: a general-purpose Rust CUDA inference engine (APXinf) and a robotics-specific harness that wraps it for VLA deployment. The engine vendors CUTLASS and FlashAttention sources — no external ML dependencies — and builds through maturin into a Python-visible extension (apxinf_py) plus a robotics Python package (apxinf_robo). The three-layer L1/L2/L3 API factoring is worth pulling out: L1 (bare model) takes already-resized [views, H, W, 3] uint8 frames, uint32 token IDs, and float32 noise and returns normalized-domain actions — the caller owns resize / tokenization / noise / unnormalize; L2 (policy) adds robot presets that supply image keys (observation/image, observation/wrist_image), state encoding (franka_libero uses 7-dim EEF delta over the standard state, unitree_g1 discretizes 16-dim state into the prompt and returns 32→16 delta joints), prompt keys, and action dims, per-call overridable; L3 (websocket server) wraps L2 in the OpenPI wire protocol and publishes the resolved wire contract in connect-time metadata so the client asserts rather than guesses. The FP8 path requires offline calibration — apxinf-robo capture-libero captures NPZ observations per preset, scripts/calibrate_pi05.py produces the per-tensor activation-scale profile, and apxinf-robo serve --precision fp8 --calibration <path> loads it at inference; INT8 W8A8 needs only the checkpoint. Benchmarking layers up cleanly: l1 measures model + Python/Rust call, l2 adds engine policy preprocessing/postprocessing (but not robot-preset code), and l3 measures client-to-server RTT including the served policy and transport. The LIBERO evaluation harness (apxinf-robo eval-libero) reproduces the published π0.5 protocol (all 10 LIBERO-10 tasks × 50 episodes × seed 7 = 500 episodes) and supports resumable runs via a JSONL ledger.

Results

Two-view π0.5 workload (224×224 NHWC uint8, 10 flow steps, H=10, batch 1, CUDA Graph replay P50) [README §Performance]:
- Jetson AGX Thor BF16: 72.45 ms / 13.8 Hz
- Jetson AGX Thor FP8: 41.16 ms / 24.3 Hz (1.76× faster than BF16)
- Jetson AGX Orin BF16: 165.67 ms / 6.0 Hz
- RTX 4090 BF16: 31.38 ms / 31.9 Hz
- RTX 4090 INT8: 25.99 ms / 38.5 Hz
With onestep action-generation pruning [README §Performance, second table]:
- Thor FP8: 26.32 ms / 38.0 Hz (~1.56× on top of standard FP8)
- RTX 4090 BF16: 20.36 ms / 49.1 Hz (~1.54× on top of standard BF16)
- Orin BF16: 119.05 ms / 8.4 Hz
LIBERO-10 closed-loop task success, 500 episodes (10 tasks × 50, H=10, replan=5, seed 7) [README §Performance, LIBERO table]. Reference pi05_libero_base published number: 92.4%.
- Thor BF16: 464/500 = 92.8%
- Thor FP8: 461/500 = 92.2%
- Orin BF16: 460/500 = 92.0%
Deployment surface: OpenPI wire-protocol server binaries per model, GGUF-free / checkpoint-directory-native (model.safetensors + config.json + tokenizer/normalizer), plus a resumable LIBERO evaluator that outputs per-task JSONL + summary JSON [README §OpenPI-compatible serving, §LIBERO evaluation].
Validated toolkits: CUDA 13.0 on Thor, CUDA 12.8 on Thor-U and RTX 4090, CUDA 12.6/13.2 on Orin [README §NVIDIA build environment].

Why it’s interesting

APXinf-robo is the second inference-runtime-shaped release on the wiki after Embodied.cpp: A Portable Inference Runtime of Embodied AI Models on Heterogeneous Robots Embodied.cpp — but stakes out a different point on the substrate axis: where Embodied.cpp went the llama.cpp route (GGUF weights, portable C++, one backend abstraction spanning CPU / CUDA / NPU, Apache-2.0 with three built-in models), APXinf goes the TensorRT-LLM route (Rust + CUTLASS + FlashAttention, hardware-locked builds, compute-capability-specific kernels, Jetson-family-optimized precisions). The result is that APXinf’s FP8 Thor path (41.16 ms / 24.3 Hz standard, 26.32 ms / 38.0 Hz with onestep pruning) is a much sharper Jetson-specific number than Embodied.cpp reports, at the cost of Embodied.cpp’s portability. Both keep policy quality within noise of the PyTorch reference (Embodied.cpp reports 91.0% closed-loop π0.5 success on LIBERO; APXinf reports 92.0–92.8% across three precisions on the exact pi05_libero_base protocol) — so the emerging picture from these two releases is that quantized VLA serving on Jetson-class hardware preserves policy quality, the open question is what runtime substrate the field converges on.

The release is also a data point on RLinf’s own strategy: Harness VLA: Steering Frozen VLAs into Reliable Manipulation Primitives via Memory-Guided Agents Harness VLA — same lab — makes the argument that the right way to deploy π0.5 is as a retryable contact-rich primitive inside an LLM planner, with the VLA reserved for the local sub-phases it’s reliable at. APXinf now supplies the serving half of that thesis: if you want to invoke a frozen π0.5 dozens of times per rollout inside a planner loop, the per-call latency of 26–41 ms on Thor makes the harness-style pattern economically plausible in a way a 400 ms end-to-end PyTorch inference does not. Sibling in strategy to FluxVLA Engine — Standardized VLA Engineering Foundation FluxVLA Engine (LimX’s training-to-deployment authoring framework) but scoped strictly to the post-checkpoint deployment plane, with concrete Jetson benchmarks the FluxVLA page conspicuously lacked.

The agent-driven model-port-workflow is the more speculative bet: it treats “port a new model to APXinf” as a task Claude Code / Codex can execute against the documented infinigence/ApxInf porting guides, with GR00T N1.7 named as the reference target. Reads as a recipe-lever bet on the recent AI-for-AI-research thread (AIDE²: First Evidence of Recursive Self-Improvement, Redwood: A Frontier AI Accelerator Designed, Verified, and Deployed from Scratch in 2 Weeks by AI) applied to the narrower runtime-porting subproblem — an early instance of “agentically optimized CUDA kernels” as a shipping deliverable rather than a paper claim.