Efficient World Action Model Inference with Adaptive Intermediate States

Zhinnan Liu; Haozhi Han; Ruge Zhang; Teng Ma; Tao Ma; Zheng Liu; Yifeng Chen; Yunquan Zhang

cs.ROSep 28, 2026

Efficient World Action Model Inference with Adaptive Intermediate States

Authors: Zhinnan Liu, Haozhi Han, Ruge Zhang, Teng Ma, Tao Ma, Zheng Liu, Yifeng Chen, Yunquan Zhang, +3 more

Organizations: Xiamen University · Institute for AI Industry Research, Tsinghua University · School of Computer Science, Peking University · Institute of Computing Technology, Chinese Academy of Sciences · Alibaba Group

Abstract

World Action Models (WAMs) enable future-aware control by jointly modeling actions and environment dynamics. However, iterative diffusion or flow inference incurs substantial denoising latency. Prior inference state offers a natural opportunity for acceleration, yet changing planning contexts, observations, and intermediate representations can quickly render retained state stale. Preserving useful computation therefore requires adapting inference state rather than reusing it as-is. To this end, we present $\mathrm{WAM}{\scriptstyle\mathrm{ACHINE}}$ , a training-free framework that accelerates WAM inference by preserving and adapting inference state for efficient and accurate continuation as the control loop evolves. Across closed-loop replans, Trajectory Remapping remaps replan state from the preceding replan to initialize the next replan, reducing redundant trajectory generation. Across denoising steps, Observation Rebinding performs anticipatory inference during action execution and rebinds retained denoising state to the real observation for continuation when consistency checks pass, reducing latency exposed to the control loop. Across Transformer layers, Residual Rescaling selectively rescales retained layer state and refreshes it through full computation of the middle layers when probe checks fail, reducing repeated Transformer computation. Evaluations of three representative WAM architectures on LIBERO and RoboTwin 2.0 show that $\mathrm{WAM}{\scriptstyle\mathrm{ACHINE}}$ achieves 1.47-3.05 $\times$ speedups in observation-to-action latency and 2.23-3.27 $\times$ speedups in GPU inference time per replan, while preserving 96.69-99.54% of native WAM task success.

Figures & tables

Figure 1: The main idea of WAMachine . WAMachine adapts state across replans, denoising steps, and Transformer layers to reduce computation and observation-to-action latency. $K$ denotes native denoising steps; $k_{1}$ and $k_{2}$ denote steps before and after real observation arrival, respectively; $k_{1}+k_{2}<K$ describes the illustrated accepted warm-start path.

Figure 2: Evidence for state continuity. (a) Consecutive replans show high cosine similarity and small relative $L_{2}$ distance between trajectory latents at matching early stages (top); remapped initialization enables three-step inference with low action RMSE and 96–100% success on LIBERO (bottom). (b) Most anticipatory prefixes are ready within the action execution window (below $y=x$ ); color shows action RMSE after observation rebinding relative to full inference under the real condition. (c) Relative $L_{2}$ changes between adjacent layer hidden states are small in the middle Action-DiT layers (top); residual rescaling and state refresh control relative output $L_{2}$ error compared with repeatedly reusing the first-step residual $R_{1}$ (bottom). Analysis protocols appear in Appendix E .

Figure 3: The framework of WAMachine . Left: Trajectory Remapping remaps replan state from the preceding replan to initialize the next replan. Top right: Observation Rebinding advances an anticipatory prefix during action execution, then rebinds retained denoising state to the real observation for continuation if consistency checks pass; otherwise, inference restarts from the remapped initialization under the real condition. Bottom: Residual Rescaling uses a shallow probe to rescale retained residuals and skip the middle layers; if consistency checks fail, it performs full computation of these layers and refreshes the retained layer state.

Cosmos Policy / LIBERO
	GPU inference time per replan		O2A latency		Subset SR
Method	Time (ms) $\downarrow$	Speedup $\uparrow$	Time (ms) $\downarrow$	Speedup $\uparrow$	(%) $\uparrow$
Native	284.29	1.00 $\times$	636.77	1.00 $\times$	98.0
RTI-DP	59.14	4.81 $\times$	399.15	1.60 $\times$	68.5
RTC	1024.53	0.28 $\times$	2203.99	0.29 $\times$	98.0
VLA-Cache	461.48	0.62 $\times$	813.54	0.78 $\times$	98.0

Table 1: Inference efficiency and subset success rate (SR) on the same fixed 200-episode subset for each WAM. GPU inference time per replan and observation-to-action (O2A) latency are averaged over replans and reported in milliseconds; speedups are relative to Native. Each acceleration method is applied independently to Native, and WAMachine includes CUDA Graphs. All methods use one policy worker on one A100 80GB GPU; WAMachine uses an additional rendering GPU for Motus. Large-sample task success appears in Table 2 . Lower (

\downarrow

) or higher (

\uparrow

) is better.

LIBERO Method Spatial Object Goal Long Avg./Ret. Cosmos Policy Native 97.40 99.60 97.67 97.53 98.05/– WAMachine 97.00 99.13 96.47 95.47 97.02/98.95 Fast-WAM-IDM Native 99.27 99.67 98.47 97.00 98.60/– WAMachine 98.80 98.87 98.00 96.93 98.15/99.54

RoboTwin 2.0 Method Clean Randomized Avg./Ret. Motus Native 86.50 85.70 86.10/– WAMachine 84.30 82.20 83.25/96.69

Table 2: Large-sample closed-loop task success (%) for Native and WAMachine . Each method uses 6,000 LIBERO episodes for each of Cosmos Policy and Fast-WAM-IDM, and 1,000 Clean and 1,000 Randomized RoboTwin 2.0 episodes for Motus. Avg./Ret. reports average task success and the percentage of Native average task success retained.

Variant	TR	OR	RR	Graph	GPU (ms) $\downarrow$	O2A (ms) $\downarrow$	Subset SR (%) $\uparrow$
Native	✗	✗	✗	✗	415.19	470.96	97.5
Native + Graph	✗	✗	✗	✓	295.29	350.25	98.0
WAMachine	✓	✓	✓	✓	114.75	144.75	98.0
w/o TR	✗	✓	✓	✓	188.02	177.20	97.0
w/o OR	✓	✗	✓	✓	108.45	165.62	95.5
w/o RR	✓	✓	✗	✓	188.99	205.45	99.0

Table 3: Ablation of the three mechanisms and CUDA Graphs on LIBERO with Fast-WAM-IDM, using the same fixed 200-episode subset as Table 1 . TR, OR, and RR denote Trajectory Remapping, Observation Rebinding, and Residual Rescaling, respectively; Graph denotes CUDA Graphs. GPU and O2A report mean GPU inference time per replan and observation-to-action latency in milliseconds; subset SR reports task success (%) from the same runs. Lower (

\downarrow

) or higher (

\uparrow

) is better.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Model	Benchmark	Large sample	Fixed subset
Cosmos Policy	LIBERO	6,000	200
Fast-WAM-IDM	LIBERO	6,000	200
Motus	RoboTwin 2.0	$1{,}000+1{,}000$	200 Clean

Appendix

Table 4: Evaluation episodes per configuration. Fixed subsets supply both task success and timing; the Motus subset uses Clean scenes only.

Model	Trajectory	First replan	Accepted prefix	Refresh
Cosmos Policy	Joint video/action	5	$1+2$	3
Fast-WAM-IDM	Video	10	$2+4$	6
	Action	10	5	5
Motus	Joint video/action	10	$3+4$	7

Appendix

Table 5: Denoising steps by execution path. In the accepted prefix column,

a+b

denotes

a

anticipatory steps followed by

b

steps under the real condition; a single number denotes only steps under the real condition. Refresh counts start from the remapped initialization.

Branch	$L$	Probe	Middle layers	Tail	Evaluated
Cosmos Policy	28	$[2,4)$	$[4,24)$	$[24,28)$	8
Fast-WAM-IDM video	30	$[2,4)$	$[4,26)$	$[26,30)$	8
Fast-WAM-IDM action	30	$[2,4)$	$[4,26)$	$[26,30)$	8
Motus joint	30	$[2,7)$	$[7,26)$	$[26,30)$	11

Appendix

Table 6: RR layer boundaries and layer evaluations remaining after accepted reuse. Motus counts each coupled layer group once.

Parameter	Fast-WAM-IDM	Cosmos Policy	Motus
Minimum cosine similarity	$0.80$	$0.984$	$0.95$
Maximum relative fitting error	$0.60$	$0.176$	$0.40$

Appendix

Table 7: RR consistency thresholds used in Equation 8 .

Cosmos Policy / LIBERO
	All 200 episodes		Pairwise common success
Method	Time (s) $\downarrow$	Speedup $\uparrow$	Time (s) $\downarrow$	Speedup $\uparrow$
Native	10.34	1.00 $\times$	–	–
RTI-DP ( $n=135$ )	95.85	$0.11\times$	63.79	$0.14\times$
RTC ( $n=193$ )	27.13	$0.38\times$	26.23	$0.38\times$
VLA-Cache ( $n=193$ )	12.51	$0.83\times$	12.08	$0.83\times$

Appendix

Table 8: Episode time and speedup on all 200 conditions and on pairwise common-success subsets.

n

gives the common-success count for each Native–method pair. Times are mean seconds per episode; each speedup uses Native time on the same subset. Bold and underline mark the best and second-best values within each model, respectively. Measurement boundaries are defined in Appendix D.1 .

	Residual Rescaling (RR)		Observation Rebinding (OR)
Model	Acceptance rate	Layer-skip fraction	Acceptance rate	Replan coverage
Cosmos Policy	81.55%	38.43%	83.75%	74.79%
Fast-WAM-IDM	79.20%	45.80%	83.47%	71.40%
Motus	49.34%	10.36%	54.34%	38.67%

Appendix

Table 9: RR acceptance rate, layer-skip fraction, OR acceptance rate, and replan coverage of the complete WAMachine on the 200-episode subsets. All values are percentages; metric definitions are given in Appendix D.2 .

Explore similar work

Latent Action as Intention Enables Efficient Future Imagination for World Action Models
Aug 25, 2026Xiang Li, Yupeng Zheng, Songen Gu +11Efficient World-Action Model Latent Actions
1CollegeAI, THU · 2AIR, THU · 4TARS Robotics +6
Efficient-WAM: A 1B-Parameter World-Action Model with Low-Cost Future Imagination
Jun 8, 2026Jiajun Li, Tiecheng Guo, Yifan Ye +9Efficient World-Action Model Faster-Wam
1The University of Hong Kong · 2Peking University · 3Muka Robotics +2
Faster-WAM: Efficient Inference-Time Future Conditioning for Robust World Action Models
Aug 5, 2026Weiheng Zhao, Haoyi Jiang, Xin Shi +5Efficient World-Action Model Recent World-Action Models
Huazhong University of Science and Technology · D-Robotics · Horizon Robotics +1

Aug 25, 2026, cs.ROJ/K move · Enter open · S save

Latent Action as Intention Enables Efficient Future Imagination for World Action Models

Xiang Li, Yupeng Zheng, Songen Gu, Huailiang Ma+10

1CollegeAI, THU · 2AIR, THU · 4TARS Robotics+6

Efficient World-Action Model Latent Actions Faster-Wam Imagination World Models Robocasa

Open paper AskPDF arXiv

World action models (WAMs) improve robot control by modeling how observations evolve, but generating future observations at test time incurs substantial latency. Fast-WAM removes this process for efficiency; however, our matched implementations show lower generalization for Fast-WAM than for future-aware alternatives, especially with scarce robot demonstrations and in out-of-distribution scenarios. To bridge this gap, we introduce LAWA, a WAM architecture that uses compact latent actions as an operational representation of future intentions, enabling efficient test-time future imagination without generating future observations. Specifically, a discrete tokenizer enhanced by action-free pre-training produces manipulation-centric codebook targets. LAWA jointly denoises a continuous latent state anchored to these targets with executable action chunks while omitting the future-video branch at inference. On RoboCasa, LAWA achieves state-of-the-art average success rates of 65.6% and 80.8% in the few-shot and full data settings, improving over the matched Fast-WAM baseline by 9.6 and 4.5 points, respectively. It also preserves the performance level of the matched Joint-WAM variant while requiring 42.9% lower inference latency. LAWA also demonstrates competitive zero-shot robustness on LIBERO-Plus and superior performance on real-world tasks. These results show that future imagination need not be discarded: retaining it with compact latent actions yields an effective trade-off among performance, generalization, and latency. Code and models will be released.

Abstract

Explore similar work

Latent Action as Intention Enables Efficient Future Imagination for World Action Models

Efficient-WAM: A 1B-Parameter World-Action Model with Low-Cost Future Imagination

Faster-WAM: Efficient Inference-Time Future Conditioning for Robust World Action Models