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

1. Problem Statement

Existing Vision-Language-Action (VLA) and World Action Models (WAMs) have achieved high success rates on standard manipulation benchmarks by scaling backbones and data. However, these models often fail under modest scene perturbations (e.g., changes in object layout, camera viewpoint, lighting, or robot initial state). The authors diagnose this fragility as a lack of object addressability.

In holistic WAMs, the predicted world state is represented as full-frame images, video-token streams, or global latents. While these representations capture scene evolution, they entangle the identity of a target object with its surrounding context within the token stream. When the scene geometry shifts, the action decoder struggles to isolate the specific object named in the instruction, leading to "drift" where the policy acts on the wrong object or fails to generalize. Current object-centric approaches often rely on dynamic slot discovery (e.g., Slot Attention) which is sensitive to occlusion and clutter, or they encode identity, appearance, and pose into a single vector, failing to decouple "which object to act on" from "what that object currently is."

2. Methodology: OA-WAM

The authors propose OA-WAM (Object-Addressable World Action Model), which reformulates world prediction as per-object slot-state prediction. The architecture enforces a strict separation between object identity and time-varying content at the tensor level.

2.1 Object-Addressable Tokenization

Each frame is decomposed into $N+1$ slots (1 robot slot, $N$ object slots). Each slot vector $s^t_k$ is constructed by concatenating:

• Identity Address ($addr_k$): A frozen 32-dimensional vector computed once per episode from the language label and the initial DINOv3 feature. It is never updated during the episode.
• Time-Varying Content ($cnt^t_k$): A 256-dimensional vector refreshed every step from SAM 3 masks and DINOv3 features, representing the object's current state (appearance, pose).
• Auxiliary Embeddings: Frame index ($\pi^t$) and role labels ($\rho_k$).

These slots are projected into a 4096-dimensional hidden space via a learnable slot adapter and fused with text, image-VQ, proprioception, and past-action tokens into a block-causal sequence processed by a 7B Chameleon-style transformer trunk.

2.2 Object-Addressable Attention (The OA Constraint)

The core architectural innovation is a parameter-free constraint applied at every transformer layer to ensure the action decoder attends to the correct object based solely on its identity, not its current appearance or context. This is achieved via two operations:

1. Address-Only Key Projection: For slot-typed positions, the key projection ($K$) reads only the first 32 dimensions (the $addr$ subvector) of the residual stream. The query ($Q$) and value ($V$) projections use the full hidden state. This forces cross-slot attention to be routed based on identity alone.
2. Per-Layer Address Reset: A forward hook overwrites the first 32 dimensions of the residual stream at slot positions with the cached $addr_k$ after every transformer block. This prevents the identity address from drifting or being corrupted by content updates in the residual stream.

2.3 Prediction Heads

The model performs a single forward pass to generate:

• World Head: Regresses the next-frame state for each slot, predicting both the content ($\hat{c}^{t+1}$) and pose ($\hat{p}^{t+1}$). This provides auxiliary supervision for temporal consistency.
• Action Head: A flow-matching MLP that decodes a 16-step continuous action chunk from a learnable action query token ($[ACT\_Q]$).
• Auxiliary Image-VQ Head: Reuses the trunk's language head to predict next-frame VQ tokens, aiding in pretraining without new parameters.

2.4 Training

The model is trained in three stages:

1. Stage 0: Pretraining the 7B trunk on a mixed corpus (web image-text + robot data) with the OA constraint enforced.
2. Stage I: Aligning the slot adapter and world head on frozen trunk weights.
3. Stage II: Full system finetuning using LoRA on the trunk and training the action head. The loss combines action flow-matching, world prediction MSE, and auxiliary regularizers for distractor permutation invariance and role alignment.

3. Key Contributions

1. Formulation of Object Addressability: The paper identifies the lack of stable, queryable object interfaces as the root cause of robustness failures in existing WAMs under geometric perturbations.
2. Architectural Decoupling: OA-WAM introduces a tensor-level mechanism (frozen address + address-only routing) that structurally decouples "which object to act on" from "what that object is," without introducing extra tokens or relying on dynamic slot discovery.
3. State-of-the-Art Robustness: The model achieves new SOTA results on the geometric axes of the LIBERO-Plus benchmark, which specifically tests robustness to layout, camera, and robot initial state changes.

4. Experimental Results

The authors evaluate OA-WAM on three benchmarks: LIBERO, SimplerEnv, and LIBERO-Plus.

• Standard Benchmarks: OA-WAM matches or exceeds SOTA VLA and WAM baselines on in-distribution tasks.
  • LIBERO: 97.8% average success (vs. 97.2% for VLA-JEPA).
  • SimplerEnv (WidowX): 79.3% average success (vs. 76.0% for CoWVLA).
• Robustness (LIBERO-Plus): OA-WAM sets a new SOTA on the geometric axes most aligned with the hypothesis (Camera, Robot Init, Layout).
  • Geo-Avg (Camera/Robot/Layout): 84.3%, a +4.8% improvement over $\pi0.5$.
  • Specific Axes: 80.5% on Camera (+4.7% over Cosmos-Policy) and 89.6% on Robot Init.
  • Sensor Noise: The model lags on sensor noise (-17.1% vs. prior best), which the authors attribute to failures in the upstream perception stack (slot extraction) rather than the policy itself.
• Mechanism Validation:
  • Ablation: Removing the OA constraint (mask or reset hook) causes a drastic drop in geometric robustness (e.g., -13.3% on Camera) while leaving in-distribution performance largely unchanged, confirming the bias is OOD-specific.
  • Causal Intervention: When the address of the target slot is swapped with a distractor's address at inference, OA-WAM's trajectory shifts to the new target with a cosine alignment of 0.87. In contrast, all holistic baselines show alignment $\le 0.09$, proving OA-WAM's actions are grounded in the explicit address subspace.

5. Significance and Claims

The paper claims that explicitly organizing a world model as addressable object states is an effective path toward robust robot manipulation under scene perturbations. By anchoring target selection to a frozen, language-grounded identity vector, OA-WAM makes robustness a structural property of the architecture rather than a result of training data diversity.

The authors emphasize that this approach provides an "intervenable object-level interface" where the policy can reliably query task-relevant objects regardless of how the scene is rearranged. They acknowledge limitations, noting that the approach is currently simulator-only and that performance degrades when the frozen perception stack fails (e.g., under heavy sensor noise or for transparent/occluded objects), but argue that the architectural contribution successfully isolates the routing mechanism from the perception and policy stages.