Source-linked AI summary

OA-WAM: Object-Addressable World Action Model for Robust Robot Manipulation

Yushan Liu, Peibo Sun, Shoujie Li, Yifan Xie, Lingfeng Zhang, Xintao Chao, Shiyuan Dong, Fang Chen, Xiao-Ping Zhang, Wenbo Ding

arXiv:2605.06481v1cs.RO

TL;DR

Existing WAMs represent predicted worlds holistically, making it difficult to address the language-referred object when scene context shifts. OA-WAM uses persistent object-addressable slots with joint world and action prediction, and reports strong benchmark performance plus substantially better causal target binding than holistic baselines.

  • Problem

    Holistic world representations entangle target identity with scene context, leaving action decoders without a stable interface for selecting language-named objects under perturbations.

  • Method

    OA-WAM decomposes frames into robot and object slots with frozen identity addresses and changing content, routes cross-slot keys through addresses, and jointly predicts future slot states and 16-step actions.

  • Results

    OA-WAM matches or exceeds strong baselines on LIBERO and SimplerEnv, leads the relevant geometric LIBERO-Plus axes, and achieves 0.87 swap-binding cosine versus at most 0.09 for holistic baselines.

  • Takeaways & Limitations

    Addressable object states provide a supported interface for world-action models that remains stable when scenes are geometrically rearranged while target identity is preserved.

  • Takeaways & Limitations

    Validation is simulator-only, and the approach inherits upstream slot-extraction errors; the frozen tokenizer struggles with small, reflective, transparent, occluded, or motion-blurred objects.

Abstract

from arXiv · show

World Action Models (WAMs) enhance Vision-Language-Action policies by jointly predicting scene evolution and robot actions, but existing methods usually represent the predicted world as holistic images, video tokens, or global latents. These representations are difficult for an action decoder to address when an instruction refers to a particular object, especially under scene shifts where object identity is entangled with context. We propose OA-WAM, an Object-Addressable World Action Model for robust robot manipulation. OA-WAM decomposes each frame into N+1 slot states, with one robot slot and N object slots. Each slot contains a persistent address vector and a time-varying content vector, and is fused with text, image, proprioception, and past-action tokens in a block-causal sequence. A world head predicts next-frame slot states, while a flow-matching action head decodes a 16-step continuous action chunk in the same forward pass. Addressability is enforced by routing cross-slot attention through address-only keys and resetting the address slice at every transformer layer, separating which object to act on from what that object currently is without adding extra tokens. OA-WAM matches strong VLA and WAM baselines on LIBERO (97.8%) and SimplerEnv (79.3%), reaches state-of-the-art performance on the most relevant LIBERO-Plus geometric axes, and remains competitive on the seven-axis aggregate. A causal slot-intervention test yields a swap-binding cosine of 0.87, versus at most 0.09 for holistic baselines. These results suggest that addressable object states provide an effective interface for robust world-action modeling under scene perturbations.

1 Introduction

OA-WAM addresses the fragility of holistic world-action models by representing scenes as object-addressable slots that separate persistent identity from changing content. It combines this representation with joint world prediction and action generation, achieving strong standard-benchmark and geometric-robustness results.

  • Object-level manipulation requires identifying language-named targets and reasoning about their relations to distractors before generating closed-loop actions.
  • Existing WAMs use holistic observations or global latents that capture scene evolution but provide no stable interface for selecting the instructed object under scene shifts.
  • OA-WAM decomposes each frame into N+1 slots with frozen identity addresses and time-varying content, then fuses them with multimodal inputs in a block-causal world-action sequence.
  • Addressability is enforced by restricting cross-slot keys to address subvectors and resetting the address stream at every transformer layer, separating object routing from content updates.
  • 97.8 LIBERO average and 79.3 SimplerEnv average lead published VLA and WAM baselines, while LIBERO-Plus results lead on geometric axes aligned with the hypothesis.
  • The authors frame poor robustness as missing object addressability and report that ablations and slot interventions support the addressable interface as the source of the robustness pattern.

2 Related Works

Prior work spans holistic world-action models, scalable VLA policies, and object-centric representations. OA-WAM connects these lines by adding persistent identity addresses and address-only routing to object slots.

  • Holistic WAMs couple action policies with future-world objectives through latent imagination, video planning, interleaved streams, JEPA-style models, and memory-based reasoning.
  • VLA research improves closed-loop manipulation through larger vision-language backbones, robot datasets, flow-matching or optimized action heads, and visuomotor designs.
  • Object-centric methods decompose scenes into object-level latent states, including generative decomposition, Slot-Attention, model-based-RL, and language-conditioned approaches.
  • OA-WAM closes the gap by partitioning each slot into a frozen identity address and time-varying content, routing cross-slot attention only through the address subvector at every layer.

3 Method

OA-WAM represents each frame with object-level slots whose persistent addresses separate object identity from time-varying content, then jointly predicts future slot states and action chunks in one block-causal sequence. Its attention constraint routes slot selection through address information while preserving content flow for prediction and action generation.

  • Object-slot tokenization: Each frame is decomposed into N+1 slots, with per-slot address and content representations used in an object-level policy input.Addresses are computed once from language labels and initial visual features, while content is refreshed frame by frame.
  • Unified sequence: The unified sequence combines text, image-VQ, proprioception, past actions, and slot streams under block-causal temporal processing.The sequence uses a shared Chameleon-style multimodal trunk and fixed-capacity masked slot positions.
  • Object-addressable attention: Cross-slot keys read only the address subvector, while within-frame slots attend bidirectionally and cross-frame attention remains block-causal.The key restriction is implemented by masking coordinates beyond the first 32 before the pretrained key projection.
  • Object-addressable attention: Address slices are reset after every transformer block, preventing residual updates from changing the identity information used for slot routing.Together with address-only keys, this makes routing depend on frozen addresses rather than time-varying content or scene context, conditional on correct slot extraction.
  • Prediction heads: The world head predicts per-slot next-frame content and pose, while the action head uses the [ACT_Q] state to decode a 16-step continuous action chunk.The action head predicts a velocity field with flow matching, and inference uses four-step forward Euler integration.
  • Training objective: Training combines action, world, image-VQ, composition, and role losses, with fixed weights and scheduled composition and role regularization.The composition loss enforces invariance to distractor permutation and insertion, while the role loss aligns action attention with language-extracted labels when available.

4 Experiments

OA-WAM is evaluated on standard manipulation, robustness, and mechanism-diagnostic benchmarks. It preserves in-distribution performance, leads on geometric LIBERO-Plus perturbations, and shows that object-addressability drives the geometric robustness gains.

  • Standard benchmarks: 97.8% LIBERO and 79.3% SimplerEnv average success show OA-WAM preserves strong in-distribution manipulation performance.These results are reported across standard benchmark suites and SimplerEnv visual-matching tasks.
  • Robustness benchmark: 84.3 Geo Avg and +4.8% over π0.5 establish a new SOTA on the geometric LIBERO-Plus axes.The geometric average covers Camera, Robot init, and Layout perturbations.
  • Robustness benchmark: OA-WAM remains on par on appearance and language axes but trails by −17.1% on Sensor Noise.The reported Sensor Noise deficit is attributed to photometric distortion corrupting per-object content during slot extraction.
  • Ablation study: V0→V2 removes 20.0% on LP camera and 7.7% on LP avg when the addressability components are disabled.The variants share perception, tokenization, trunk weights, and training data, isolating the OA constraint as the main source of the geometric-axis effect.
  • Ablation study: Disabling the key mask reduces LP camera by 13.3% and LP robot by 18.2%, while LIBERO changes by only 1.5%.This separates geometric out-of-distribution degradation from comparatively small in-distribution drift.
  • Mechanism diagnostics: An address swap yields OA-WAM swap binding of 0.87, whereas holistic baselines remain at or below 0.09.The intervention measures cosine alignment between the resulting end-effector trajectory and displacement toward the swapped target.

5 Conclusion

OA-WAM makes world-action modeling object-addressable by routing cross-slot attention through frozen per-object identity addresses. Its benchmark and intervention results support stable target binding under geometric scene rearrangements, while simulator-only validation and perception limitations bound the current evidence.

  • Conclusion: OA-WAM uses a slot-aware 7B trunk with frozen per-object identity addresses to make world-action modeling object-addressable.The inference pipeline caches addresses at episode start and processes slot-aware representations before separate world and action readouts.
  • Conclusion: 0.87 swap-binding cosine versus ≤0.09 for holistic baselines verifies address-based target binding under causal slot intervention.The intervention swaps the language-bound target address with another in-scene slot while holding other inputs fixed.
  • Conclusion: The model leads on LIBERO and SimplerEnv and achieves a +4.8% geometric LIBERO-Plus Geo Avg over π0.5.The conclusion connects these results with OA isolation that degrades geometric axes while leaving LIBERO essentially unchanged.
  • Limitations: The reported robustness evidence is simulator-only and does not yet prove real-robot deployment.This is the stated scope boundary of the validation.
  • Limitations: Perception costs ∼95 ms/frame versus ∼5.6 ms for the trunk and head, and the tokenizer fails on several difficult object appearances.The listed cases include small reflective, transparent, occluded, or motion-blurred objects.

B Six-path tokenization and data preprocessing

The preprocessing pipeline converts demonstrations into six token paths, including language, images, object slots, proprioception, and past actions, then assembles them into a unified trunk-ready sequence. Object slots combine language-grounded masks, visual features, pose, labels, and shape descriptors, while cached features and explicit masks support efficient, permutation-aware processing.

  • Preprocessing: Images are resized to 256×256, encoded by frozen perception modules, and cached offline without RGB augmentation.Cached features are reused across stages and seeds, while Stage II uses slot-level distractor permutation and insertion augmentation.
  • Language: Instructions are passed unchanged to noun-phrase parsing and BPE tokenization, with average LIBERO instruction length around 30 BPE tokens.The noun-phrase output is used only to prompt SAM 3, whereas BPE tokens enter the trunk embedding stream.
  • Sequence construction: Training examples use four historical observations, past-action windows, and 16-step action targets sampled around a random pivot frame.Sequences are explicitly padded and masked when the pivot is near the beginning or the target chunk exceeds the demonstration.
  • Token paths: Six token paths represent text, noun-phrase prompts, image codes, object slots, proprioception, and past actions.Noun phrases guide SAM 3 but are excluded from the trunk; the other streams are consumed by the shared Chameleon sequence.
  • Object slots: Object slots combine SAM 3 masks, DINOv3 features, pose, noun-phrase labels, and mask-shape descriptors into raw 540-dimensional vectors.Up to 16 object slots are supported, with the robot fixed as slot 0 and padding slots masked downstream.
  • Slot representation: The slot adapter maps raw slot features into the 4096-dimensional trunk space, while separate networks produce address and content vectors.The address network uses label and visual features, and the content network processes the 540-dimensional raw slot vector.

D Prediction-head architectures

OA-WAM uses separate world and action heads over the shared trunk representation. The world head predicts per-object content and pose, while a flow-matching head predicts a 16-step action chunk through iterative velocity integration; slot-aware processing preserves permutation structure and supports distractor robustness.

  • World head: The world head uses independent MLPs to predict 256-dimensional content and 9-dimensional pose outputs from per-slot hidden states.Robot slots are excluded from world supervision, and padded slots are masked.
  • Action head: The action head conditions an 8-block residual flow MLP on the trunk action-query state, a noisy 16×7 action chunk, and a timestep embedding.It outputs a 16×7 velocity field for continuous action-chunk generation.
  • Auxiliary prediction: The auxiliary image-VQ head reuses the pretrained language-model head and applies weighted next-frame VQ supervision without adding parameters.Image-VQ tokens receive weight 0.04 so the auxiliary loss does not dominate action-learning gradients.
  • Action decoding: Inference uses 4-step forward Euler integration, with approximately 10 ms required per action chunk on one A100.Increasing from 4 to 8 Euler steps changed LIBERO success by less than 0.3% in the reported ablation.
  • Robustness objectives: Distractor permutation and insertion losses encourage invariant outputs under reordered or added non-target slots, while a temporary role loss aligns action attention with language-derived labels.The composition weight warms up during training, and the role-loss weight is removed after the first half.

G Three-stage training and inference latency

OA-WAM is trained in three stages: foundation pretraining, slot/world-head training, and LIBERO action finetuning with LoRA. The staged design freezes most of the trunk after pretraining, limits Stage II data to standard LIBERO demonstrations, and achieves parity between LoRA and full finetuning across reported axes.

  • Stage 0: Stage 0 pretrains the full Chameleon-7B trunk on approximately 2.5T mixed web and robot tokens.The mixture retains web image-text data for multimodal retention and robot data for slot-level world prediction.
  • Compute: Reproducing Stage 0 requires approximately 166k A100-hours, whereas downstream users can use the released checkpoint and scripts with approximately 60 GPU-days on 8×A100.The reported Stage 0 run used 384 A100-80GB GPUs for about 18 days.
  • Stage I: Stage I freezes the Stage 0 trunk and trains the slot adapter, world head, and address-reset hook with approximately 23.8M trainable parameters.Training uses LIBERO, DROID, and a subset of Open X-Embodiment with cached six-path features.
  • Stage II: Stage II adds rank-32 LoRA and the flow-matching action head, reaching approximately 127M trainable parameters on standard LIBERO demonstrations.The full objective includes world, auxiliary VQ, composition, and role losses, with composition weighting ramped during training.
  • Finetuning comparison: LIBERO and LIBERO-Plus results differ by at most ±1% between LoRA and full finetuning on every reported axis.The LoRA route is therefore used as the default despite the full-trunk alternative requiring substantially longer training.

Evaluation runtime

Runtime is dominated by perception and the slot-aware 7B trunk during closed-loop control. With cached or amortized components, OA-WAM processes each 16-step action chunk in approximately 233 ms, corresponding to about 4.3 Hz inside a 20 Hz simulator.

  • Evaluation cost: A full evaluation pass takes 1.7 hours for LIBERO, 11.4 hours for LIBERO-Plus, and 1.6 hours for SimplerEnv on one A100.Frozen-perception precomputation for the combined benchmark cache runs offline on 8 A100s.
  • Control loop: Closed-loop control processes one action chunk every 16 simulator steps, yielding an effective rate of approximately 4.3 Hz in a 20 Hz simulator.The total per-chunk latency is approximately 233 ms.
  • Latency breakdown: The per-chunk breakdown is approximately 138 ms for perception, 80 ms for the slot-aware trunk, 10 ms for the action head, and 5 ms for sequence construction.The perception stack includes SAM 3, DINOv3, and Qwen3-VL processing.

H Evaluation protocol

The evaluation covers standard LIBERO and SimplerEnv tasks, LIBERO-Plus perturbation axes, and mean success rates aggregated over repeated seeds and published or reproduced baselines.

  • LIBERO evaluates Spatial, Object, Goal, and Long suites using 100 episodes and three seeds per suite, with success defined by task completion.
  • LIBERO-Plus uses an official perturbation generator without training on perturbation factors and reports success separately for seven dimensions plus their average.
  • SimplerEnv evaluates four WidowX visual-matching tasks with 25 episodes per cell, three seeds, and an average across the four Bridge tasks.
  • Success rates are averaged over seeds, while perturbation-strength curves use the same seeds across methods.
  • Reported comparison rows combine published baseline results with reproduced π0.5 runs and leave unavailable LIBERO results blank.

Supporting ablations: A3 (world head) and A4 (distractor consistency)

Supporting ablations test the world-prediction head and distractor-consistency loss as separate contributors to robustness, with effects concentrated on specific geometric or layout metrics.

  • Supporting ablations: A3 and A4 support separate claims about the joint world-prediction objective and distractor-consistency regularization, alongside the A1 factorial analysis.
  • A3 – world head: Disabling future prediction removes the multi-step slot world loss and reduces LP camera by 7.1% and LIBERO by 2.2%, while LP avg remains statistically flat.
  • A4 – distractor consistency: Removing the distractor-consistency loss lowers LP layout by 4.3% and increases permutation KL and insertion drift by roughly 5×.

V3 of A1 – per-layer address-reset hook factorization

The V3 factorization isolates the per-layer address-reset hook within the A1 two-factor design, while the broader evaluation connects the intervention to geometric robustness and slot-binding diagnostics.

  • V3 factorization: V3 retains address-only key projection but removes the per-layer reset hook, allowing the address subvector to drift after layer 0.
  • V3 factorization: V3 lies between full OA-WAM and the mask-off variant on every metric, showing that both the hook and key mask contribute to robustness.
  • V3 factorization: The combined removal causes a 3.9% super-additive residual interaction on LP avg beyond the individual losses of 0.7% and 3.1%.
  • Architectural attribution: The slot-only control improves to full OA-WAM by +20.0% on LP-camera and +7.7% on LP-avg, while standard LIBERO changes by only +2.4%.
  • Architectural attribution: The causal swap-binding metric rises monotonically from 0.06 in the slot-only control to 0.87 when the mask and hook are restored.
  • Scope and diagnostics: The interpretation is bounded by upstream slot extraction, test-time robustness boundaries, and failure sources including perception errors, latency, blocked targets, and ambiguous instructions.
  • Evaluation coverage: The benchmark gallery spans four LIBERO suites, four SimplerEnv WidowX tasks, and seven LIBERO-Plus perturbation axes.
Loading 2605.06481v1…