PaperlayerView exact revision on arXiv

Source-linked AI summary

WAMJET: A Harness for World Action Model Acceleration

Le Chen, Lixin Liu, Jan Schneider, Zeju Qiu, Simon Guist, Bernhard Schölkopf, Dieter Büchler

arXiv:2610.03797v3cs.CVcs.AIcs.RO

TL;DR

WAM inference is costly, and choosing and combining accelerations requires deployment-specific engineering while preserving policy behavior. WAMJET equips coding agents with reusable guidance and measurement and validation tools for iterative, bottleneck-driven optimization. It accelerates inference across WAMs and settings, with lossless speedups and further gains from approximate optimization.

  • Problem

    Selecting and combining acceleration techniques for each WAM and hardware platform requires repeated engineering, profiling, and measurement while preserving policy behavior.

    Source: §I · p. 1

  • Method

    WAMJET guides coding agents to profile a trained WAM, target bottlenecks, test accelerations, and validate candidates using measurement and analysis tools.

    Source: §I · pp. 1–2, A. Overview · p. 4

  • Results

    WAMJET accelerates inference across coding agents and WAMs; LingBot-VA achieves a 9.95× speedup with GPT-6-astra.

    Source: B. Lossless Acceleration · p. 6

  • Takeaways & Limitations

    Experiments across six models show gains from conditioning reuse, reduced execution overhead, and kernel optimization, with additional gains from continued search.

    Source: §VI · p. 8

    The authors identify broader paired task tests and controlled comparisons under matched search budgets as future work.

    Source: §VI · p. 8

Abstract

from arXiv · show

World Action Models (WAMs) leverage pretrained video foundation models for robot manipulation, but their large backbones and video-action co-prediction are expensive. Although existing acceleration techniques offer many ways to reduce this cost, selecting and composing them requires substantial engineering for each model and hardware platform. To tackle this bottleneck, we present WAMJET, an agentic harness that accelerates WAM inference by equipping coding agents with reusable optimization guidance and measurement and validation tools. WAMJET follows a bottleneck-driven workflow where the agent profiles inference, modifies targeted code, validates effects, and iteratively refines the acceleration stack as bottlenecks shift, while preserving action quality. Experiments span six WAMs, three coding agents, and two GPU architectures. WAMJET achieves up to 9.95x lossless speedup over upstream implementations. Approximation and hardware-aware optimization yield additional latency reductions, with comparable success rates. The results show that WAMJET can produce effective acceleration stacks for WAM deployment.

I. INTRODUCTION

WAMs draw on pretrained video models to predict future states and actions, but their costly inference and deployment-specific optimization make acceleration difficult. WAMJET gives coding agents reusable guidance and validation tools to build tailored acceleration stacks.

I. INTRODUCTION

  • Video-model scaling, repeated denoising, and data movement make WAM inference costly for real-time robot control.WAMs use video-model priors to predict future states alongside actions.Source: §I · p. 1
  • Choosing and composing accelerations requires repeated profiling because bottlenecks shift across models and hardware, while changes can affect robot actions.Quantization and feature reuse can introduce errors that tensor comparisons alone may not predict for manipulation.Source: §I · p. 1
  • WAMJET equips coding agents with reusable optimization guidance and tools for profiling, latency measurement, numerical comparison, and paired task analysis.The agent starts from a runnable policy, tests lossless or approximate accelerations, and benchmarks and validates candidates.Source: §I · pp. 1–2
  • WAMJET’s workflow iterates through bottleneck identification, targeted changes, and validation, retaining useful changes and adapting to the target hardware.The work evaluates the framework across WAMs and GPU architectures, including approximate inference and its limitations.Source: §I · p. 2

II. RELATED WORK

Prior work brings vision-language and video-generation models into robot manipulation, including approaches that connect future visual prediction with action generation. WAMJET instead accelerates the inference code of trained policies.

A. Foundation Models for Robot Manipulation

  • Foundation models support manipulation through direct action prediction, language-conditioned skill selection, and executable programs for control.Examples include RT-2 and OpenVLA, π0 and Octo, and SayCan and Code as Policies.Source: A. Foundation Models for Robot… · p. 2
  • Video-prediction approaches use generated visual plans or jointly learned video-and-action representations to connect future visual states with robot actions.UniPi extracts actions from text-conditioned video plans; UVA, UWM, and DreamZero couple video and action learning in different ways.Source: A. Foundation Models for Robot… · p. 2
  • WAMJET accelerates the inference code of a trained policy and uses numerical and task evaluations to assess each optimization’s effects.Source: A. Foundation Models for Robot… · p. 2

B. Accelerating Video Generation and World Action Models

Video and WAM acceleration methods target different parts of inference, from sampling steps and redundant computation to attention, precision, and deployment overhead. WAMJET selects and combines techniques according to model and hardware bottlenecks.

B. Accelerating Video Generation and World Action Models

  • Diffusion acceleration reduces model evaluations, reuses or forecasts intermediate features, sparsifies attention, merges tokens, or enables low-bit inference.DPM-Solver and progressive distillation reduce sampling steps; TeaCache and TaylorSeer exploit denoising-trajectory redundancy.Source: B. Accelerating Video Generation and… · pp. 2–3
  • The workflow profiles inference, tries lossless acceleration before permitted approximation, validates candidates, and reassesses bottlenecks within the search budget.Source: Figure 2 · p. 3
  • FlashAttention accelerates exact dense attention through IO-aware tiling and reduced memory traffic, unlike approximate acceleration methods.Source: B. Accelerating Video Generation and… · p. 3
  • DreamZero combines deployment-time optimizations, while Fast-WAM and Light-WAM change model design and require retraining; C3ache is training-free.C3ache reuses action-expert residuals across consecutive inference chunks.Source: B. Accelerating Video Generation and… · p. 3
  • WAMJET selects and combines complementary acceleration techniques according to each model’s and hardware platform’s bottlenecks.Numerical equivalence, perceptual similarity, and task success are distinct evaluation criteria; video fidelity alone does not ensure manipulation performance.Source: B. Accelerating Video Generation and… · p. 3

C. Agentic Workflows and Harness Engineering

Agentic harnesses structure iterative tool use and feedback across software, GPU-kernel, and robotics tasks. The preliminaries describe WAMs as jointly predicting future states and actions, with costs that vary by architecture and deployment.

C. Agentic Workflows and Harness Engineering

  • Agentic workflows pair iterative tool use with external feedback, while performance-engineering systems use correctness checks and hardware measurements to improve GPU kernels.ReAct introduced a reasoning–action loop, SWE-agent an agent–computer interface, and harnesses structure the runtime around the agent.Source: C. Agentic Workflows and Harness… · p. 3
  • Robotics harnesses address different lifecycle stages, from RL training and benchmark integration to embodied control; WAMJET optimizes trained-WAM inference code.Source: C. Agentic Workflows and Harness… · p. 3
  • The Sol Video Inference Engine combines parallel optimization agents with human quality feedback; WAMJET adds recurrent-state and changing-observation checks plus paired task evaluation.Source: C. Agentic Workflows and Harness… · p. 3

III. PRELIMINARIES

  • A WAM combines a pretrained video-generation world-model prior with predictions of future states and robot actions.Source: §III · pp. 3–4
  • WAMs factor joint future-state and action prediction into future-state generation followed by action inference conditioned on the imagined future.Implementations may jointly denoise video and actions or generate video before a separate inverse-dynamics module; dominant costs vary by architecture and platform.Source: §III · p. 4

IV. WAMJET

WAMJET guides a coding agent through a bottleneck-driven search for faster implementations of a trained WAM policy. The harness records experiments and returns an implementation with measurements and validation evidence.

A. Overview

  • WAMJET combines optimization guidance with Python tools for timing, trace analysis, numerical comparisons, and paired task analysis.Table I summarizes the components by functional role.Source: A. Overview · p. 4, Table I · p. 4
  • The user specifies the policy code and checkpoint, inference configuration, quality requirements, and search budget.The agent returns the selected implementation with measurements and validation evidence for reproduction and assessment.Source: A. Overview · p. 4
  • The agent profiles inference, tests optimizations matched to bottlenecks, records accept/revert decisions, and reassesses bottlenecks after each change.Search continues until user requirements are met or the budget limit is reached.Source: A. Overview · p. 4

B. Agent-Native Acceleration Workflow

The workflow establishes a comparable baseline, profiles costly inference stages, and tests targeted changes. It iterates within a time budget and validates the final candidate against numerical, state, and task requirements as applicable.

B. Agent-Native Acceleration Workflow

  • Baseline and candidate latency are measured after warmup with fixed checkpoints, inputs, random seeds, and inference settings.Repeated baseline outputs can reveal numerical variation; setup and cache states are controlled for comparable measurements.Source: B. Agent-Native Acceleration Workflow · p. 4
  • After each improvement, the agent examines remaining computation, chooses follow-up experiments, and reserves time to remeasure the baseline and best candidate.Final checks cover numerical outputs, stateful behavior, and task performance as required by the change.Source: B. Agent-Native Acceleration Workflow · p. 5
  • External memory stores validated environment fixes and model-specific findings for reuse, while harness evaluations hold guidance and resources fixed across settings.The evaluation settings receive the same search budget and access to prior information.Source: B. Agent-Native Acceleration Workflow · p. 5

C. Correctness and Candidate Acceptance

WAMJET separates lossless changes, validated approximations, and changed or unvalidated workloads, with acceptance checks matched to the change. Its acceleration search targets measured bottlenecks and evaluates effects on complete inference and task performance.

C. Correctness and Candidate Acceptance

  • Tier A preserves computation, precision, workload, and state behavior; Tier B approximates internal computation but fixes the denoising steps, schedule, action horizon, and required outputs.Tier B must pass paired task evaluation within a specified quality-loss bound; Tier C includes changed workloads or approximations without passing task evaluation.Source: C. Correctness and Candidate Acceptance · p. 5
  • The numerical drift checker may make acceleration too conservative, so it is usually turned off in practice.The original reference is retained to detect cumulative drift across successive changes.Source: C. Correctness and Candidate Acceptance · p. 5
  • Stateful changes are checked on changing input sequences and resets, while task performance is tested on matched episodes with identical initial states and random seeds.Episode counts and acceptable loss are set in advance; final evaluation uses episodes excluded from calibration and optimization.Source: C. Correctness and Candidate Acceptance · p. 5

D. Acceleration Techniques

  • Repeated conditioning can be cached when inputs remain unchanged, while static buffers and CUDA graphs can reduce allocation and dispatch overhead.Changes must preserve cache updates and ensure graph replay receives current inputs.Source: D. Acceleration Techniques · p. 5
  • For kernel- or transfer-bound inference, the agent can try packed projections, attention implementations, layouts, compiler specialization, or fusion to reduce intermediate memory traffic.Each change is assessed on the complete inference call, since a faster kernel may have little effect on total latency.Source: D. Acceleration Techniques · p. 5
  • When approximation is allowed, quantization and approximate feature or K/V reuse are measured for the actual workload, including conversion and cache-maintenance costs.Useful candidates are evaluated on robotic tasks, and earlier changes can make initially unhelpful techniques worth revisiting.Source: D. Acceleration Techniques · p. 5

V. EXPERIMENTS

The experiments examine lossless and approximate acceleration across WAMs, coding agents, and GPU architectures. The reported comparisons use warmed inference latency and recorded timing and state-update protocols.

A. Experimental Setup

  • The evaluation spans six WAMs, three coding agents, and two GPU architectures, with approximate acceleration tested on H100 and B200.A recently released state-of-the-art WAM on B200 is included as a robustness check for further acceleration of an already optimized pipeline.Source: A. Experimental Setup · pp. 5–6
  • All accelerations reported in Table II are lossless, with no statistically significant difference in policy performance.The study evaluates lossless acceleration across coding-agent and WAM combinations on H100.Source: A. Experimental Setup · pp. 5–6, Table II · p. 6
  • Inference latency is measured after warmup, excluding model loading, compilation, and CUDA-graph capture, under recorded timing and state-update protocols.Table III covers two WAMs, baseline and different acceleration levels, and two GPU architectures.Source: A. Experimental Setup · p. 6, Table III · p. 6

B. Lossless Acceleration

WAMJET improves lossless inference latency across coding agents and WAMs, and additional approximate optimizations yield further speedups across GPU architectures. Its quantization strategies vary by model and GPU, while the reported evaluation checks whether approximate acceleration preserves task performance.

B. Lossless Acceleration

  • The H100 lossless study pairs three coding agents with four WAMs and compares upstream baselines, independent agent optimization, and WAMJET-guided optimization.The study states that lossless acceleration leaves policy performance unchanged.Source: B. Lossless Acceleration · p. 6
  • All WAMJET configurations speed up inference over shared baselines, with lower latency than runs without WAMJET in all 12 entries; LingBot-VA reaches 9.95× speedup with GPT-6-astra.A second WAMJET round can further reduce inference latency.Source: B. Lossless Acceleration · p. 6

C. Approximate Acceleration across GPU Architectures

  • With approximate optimization, overall speedups are 2.14× and 2.21× for DreamZero and 1.25× and 1.47× for Cosmos3 on H100 and B200, respectively.Lossless-only speedups are 1.56× and 1.94× for DreamZero, and 1.06× and 1.21× for Cosmos3, on H100 and B200, respectively.Source: C. Approximate Acceleration across GPU… · p. 6
  • The reported RoboLab evaluation indicates that approximate acceleration preserves action quality and task performance.The evaluation uses 20 tasks with 10 episodes per task; the supplied table caption identifies success rates for Cosmos3-Nano-Policy-DROID and DreamZero at baseline and different acceleration levels.Source: C. Approximate Acceleration across GPU… · p. 6, Table IV · p. 7

D. Further Optimization of a Strong WAM Baseline

WAMJET-guided optimization further accelerates OpenWAM on B200, despite its existing inference optimizations. Lossless and approximate approaches achieve similar speedups while preserving action quality.

D. Further Optimization of a Strong WAM Baseline

  • WAMJET reduces OpenWAM latency by 32.69 ms with lossless optimization and 31.70 ms with approximate optimization, achieving 2.67× and 2.75× speedups, respectively.OpenWAM’s baseline latency is 87.25 ms on B200; both approaches preserve action quality.Source: Table V · p. 7, D. Further Optimization of a Strong… · p. 7
  • The evaluation applies WAMJET-guided GPT-6-astra to OpenWAM on B200 and tests full LIBERO and full LIBERO-plus.OpenWAM already uses the PyTorch compiler and velocity cache.Source: D. Further Optimization of a Strong… · p. 7
  • Similar lossless and approximate speedups indicate that quantization adds limited benefit in this setting, with most gains coming from lossless methods.The authors conclude that the guided agent finds substantial opportunities even in an already highly optimized pipeline.Source: D. Further Optimization of a Strong… · p. 7

E. Optimization Progress over Search Time

The optimization progress analysis shows iterative exploration on DreamZero and OpenWAM, with model-specific final recipes despite some shared techniques. WAMJET also reduces startup times, and the conclusion reports improvements across six models with further gains from continued search.

E. Optimization Progress over Search Time

  • On B200, accepted optimization trajectories improve through iterative exploration, while DreamZero and OpenWAM finish with different recipes despite sharing techniques such as CUDA Graphs and operator fusion.Without WAMJET, coding agents tend to terminate early, leaving optimization opportunities unexplored.Source: E. Optimization Progress over Search… · pp. 7–8

F. Startup Cost Acceleration

VI. CONCLUSION

  • Across six models, WAMJET improves inference through conditioning reuse, reduced execution overhead, and kernel optimization, with additional gains from continued search.The conclusion notes that quantization can further reduce latency, while its effect on task performance requires evaluation.Source: §VI · p. 8

Related papers

Loading 2610.03797v3…