A Harness for World Action Model Acceleration

Hand-tuning every WAM does not scale

Accelerating WAM inference is not a one-time optimization. It requires a recipe that composes techniques for a specific model, hardware platform, and task, and the best recipe for one WAM may not transfer directly to another.

Manual optimization requires coordination across kernel experts, model authors, and robotics engineers, and can take days to weeks for each model-GPU pair. A coding agent alone may stop at generic fixes, such as prompt caching or CUDA graphs, leaving deeper optimization opportunities unexplored.

A bottleneck-driven loop

The agent first reduces startup cost, then profiles a single inference call and targets the bottleneck. Each candidate is validated, and each accepted change becomes the new baseline. The agent then re-profiles the updated implementation and moves on to the next bottleneck.

• Lossless techniques first: compilation, fusion, CUDA Graphs, reuse, kernel tuning.
• Approximate techniques only when permitted, such as quantization matched to the GPU.
• The loop continues within the available budget, rather than stopping at the first win.

Where WAM latency hides

WAMJET's guidance helps agents identify where time is spent. Pick a latency source to see a schematic CPU/GPU timeline before and after the optimization, along with the code change behind it.

Every agent goes further with WAMJET

Our first experiment gives three coding agents the same four WAMs on an H100, with and without WAMJET, and allows only lossless changes. R1 is one round of WAMJET optimization; R2 continues for a second round and is the reported result.

• 12 of 12 agent-WAM pairs end faster with WAMJET than the same agent without it.
• A second round helps in every pair: R2 beats R1 all 12 times.
• The largest gain is 9.95× on LingBot-VA with GPT-6-astra; the same agent reaches 5.51× without WAMJET.

Each kept change behind the GPT-6-astra results is in the loop replay above.

Quantization recipes follow the GPU

When approximation is permitted, the agent applies quantization on top of the lossless optimizations and selects formats natively supported by the GPU. We test DreamZero[1] and Cosmos3-Nano-Policy-DROID[5] on Hopper (H100) and Blackwell (B200).

Faster must not mean worse. We report the success rate on a RoboLab[6] subset of 20 tasks with 10 episodes per task. Task performance remains comparable to upstream.

Room left in already-optimized SOTA WAMs

OpenWAM

OpenWAM[7], released on September 7, 2026, reports state-of-the-art results on several robot benchmarks and already incorporates several inference optimizations, such as the PyTorch compiler and velocity cache. We let WAMJET-guided GPT-6-astra optimize it on a B200 and evaluate on full LIBERO[8] and full LIBERO-plus[9].

Long-WAM

Long-WAM[10], released on October 7, 2026, is another state-of-the-art WAM. It already incorporates NVFP4 quantization, CUDA Graphs, compilation, denoising-invariant reuse, and GEMM/VAE tuning. We use its V4/A4 configuration and let WAMJET-guided GPT-6-astra optimize it on a B200.

On OpenWAM, the guided agent identifies substantial acceleration opportunities in an already highly optimized inference pipeline, achieving a 2.67× speedup while preserving action quality. On Long-WAM, WAMJET reduces latency from 108.92 ms to 43.32 ms, a further 2.51× speedup over the highly optimized upstream implementation.

How optimization progresses

Here is how GPT-6-astra achieved its DreamZero, OpenWAM, and Long-WAM results on a B200. Each dot is a measured trial and the step line is the best configuration kept so far. Gray is the same agent without WAMJET.

• Without the harness, the agent stops early: after 1.38 h on DreamZero (1.11×), 0.62 h on OpenWAM (1.30×), and 2.38 h on Long-WAM (1.72×).
• With WAMJET it keeps improving for hours, reaching 2.21× on DreamZero (NVFP4 + FP8), 2.75× on OpenWAM (FP8), and 2.51× on Long-WAM.
• Shared techniques, different recipes: CUDA Graphs and kernel tuning recur across the runs, but the final optimization recipes differ between models.

Run WAMJET on your own WAM

To use WAMJET with your coding agent, start by installing its tools:

git clone https://github.com/liulixinkerry/WAMJET.git
cd WAMJET
python -m pip install -e '.[numerics]'

Then start your coding agent in the repository and launch a campaign with the example prompts, replacing the placeholders with your paths and campaign ID:

Lossless and approximate prompts → Usage guide: workflow, tools, validation →

Citation

@article{chen2026wamjet,
  title={WAMJET: A Harness for World Action Model Acceleration},
  author={Chen, Le and Liu, Lixin and Schneider, Jan and Qiu, Zeju and Guist, Simon and Sch{\"o}lkopf, Bernhard and B{\"u}chler, Dieter},
  journal={arXiv preprint arXiv:2610.03797},
  year={2026}
}

References

1. S. Ye et al. World action models are zero-shot policies. arXiv:2602.15922, 2026.
2. T. Yuan, Z. Dong, Y. Liu, and H. Zhao. Fast-WAM: Do world action models need test-time future imagination? arXiv:2603.16666, 2026.
3. L. Li et al. Causal world modeling for robot control. arXiv:2601.21998, 2026.
4. M. J. Kim et al. Cosmos Policy: Fine-tuning video models for visuomotor control and planning. arXiv:2601.16163, 2026.
5. N. Agarwal et al. Cosmos 3: Omnimodal world models for physical AI. arXiv:2606.02800, 2026.
6. X. Yang et al. RoboLab: A high-fidelity simulation benchmark for analysis of task generalist policies. Proceedings of Robotics: Science and Systems, Sydney, Australia, July 2026.
7. Y. Wang et al. OpenWAM: An open, modular exploration towards systematic world-action model pretraining. arXiv:2609.07398, 2026.
8. B. Liu et al. LIBERO: Benchmarking knowledge transfer for lifelong robot learning. Advances in Neural Information Processing Systems (NeurIPS), Datasets and Benchmarks Track, 2023.
9. S. Fei et al. LIBERO-Plus: In-depth Robustness analysis of vision-language-action models. arXiv:2510.13626, 2025.
10. W. Huang et al. Long-WAM: Scaling the Context of World-Action Models. arXiv:2610.10528, 2026.