MaskWAM: Unifying Mask Prompting and Prediction for World-Action Models

Abstract: World Action Models (WAMs) present a promising paradigm for robotic control via video prediction. However, current WAMs suffer from fundamental spatial bottlenecks: standard text inputs introduce referential ambiguity in cluttered scenes, while unstructured RGB predictions lack semantic grounding and remain biased by task-irrelevant backgrounds. To overcome these limitations, we introduce MaskWAM, an object-centric world-action model. By jointly integrating masks as both explicit inputs and predictions via a unified Mixture of Transformers (MoT), MaskWAM unlocks robust policy generalization. This design provides two key benefits: (1) predicting future masks yields object-centric semantic supervision that suppresses visual noise, significantly enhancing even standard text-conditioned WAMs; and (2) coupling this predictive supervision with first-frame visual prompts, such as target object masks, establishes a precise spatial anchor that substantially reduces language ambiguity. Crucially, as WAMs are inherently vision-driven architectures, direct mask conditioning yields substantially stronger guidance than text alone, establishing a precise and robust paradigm for manipulating unseen objects. Evaluations on LIBERO, RoboTwin, and real-world tasks demonstrate that MaskWAM significantly outperforms baselines in both language-clear and language-ambiguous tasks.

• The paper unifies first-frame mask prompting with future RGB-mask prediction in a flow-matching World Action Model, enabling object-focused action generation without costly full video denoising.
• The paper achieves 98.4% average success on LIBERO, 92.2% on RoboTwin 2.0, and 84.3% on language-clear real-world tasks, outperforming strong WAM and VLA baselines.
• The paper shows that prediction is essential for grounding: mask prompting without future mask prediction achieves only 21.6% success versus 84.9% for the full model, while mask supervision also improves robustness to distractors, novel objects, and lighting changes.

MaskWAM addresses two spatial bottlenecks that limit current World Action Models (WAMs): text conditioning introduces referential ambiguity in cluttered scenes, and RGB-only future prediction provides weakly structured supervision that treats all pixels equally, leaving task-relevant objects entangled with background clutter. The paper proposes an object-centric WAM that uses masks both as explicit first-frame visual prompts and as auxiliary prediction targets, trained jointly within a single flow-matching diffusion framework (2606.13515).

Motivation and positioning

WAMs couple action generation with future video prediction, treating observation forecasting as a proxy task that improves representations relative to direct VLA regression. The authors argue that most existing WAM formulations supervise the visual future only through RGB reconstruction, creating a representational gap between high-dimensional pixels and the low-dimensional spatial structure governing control. Prior work on grounding either conditions policies on static visual cues (points, boxes, masks) or predicts intermediate representations (keypoints, trajectories, masks) as isolated targets; MaskWAM unifies both paradigms by combining first-frame mask conditioning with future mask prediction inside one model.

Architecture

MaskWAM is built on Wan 2.2 with a frozen causal 3D Video VAE and a frozen T5 text encoder. Masks are rendered as three-channel images with a fixed color palette (background assigned a separate color), encoded by the same VAE as RGB, and channel-concatenated with RGB latents into a $2C$-channel stream. The pretrained patch embedding is expanded from $C$ to $2C$ channels, with new mask channels zero-initialized to preserve pretrained behavior; the output head is similarly expanded to predict both RGB and mask latent velocities. A Mixture of Transformers (MoT) comprises a visual branch denoising the joint RGB-mask stream and a lightweight action expert that attends jointly over visual context to denoise action chunks via flow matching.

Training uses a unified flow-matching objective summing video, mask, and action losses, with decoupled noise schedules: a shared visual timestep $\tau_v$ keeps RGB and mask streams spatiotemporally aligned, while an independent action timestep $\tau_a$ forces the action expert to condition on visual contexts at varying noise levels. This decoupling enables partial-denoising inference: at deployment, only one denoising step is performed on the joint stream, and the action expert generates actions from these partially denoised latents, avoiding full video generation cost. To support both language-clear and language-ambiguous regimes with a single policy, the first-frame mask is dropped to a zero tensor with probability 0.5 during training. At test time, an initial mask anchor is produced once per episode by SAM-3 from a phrase, click, box, or coarse mask—no real-time tracking or repeated prompting is required.

Benchmark results

On LIBERO, MaskWAM achieves a state-of-the-art 98.4% average success rate, exceeding $\pi_{0.5}$ (96.8%), Motus (97.7%), and FastWAM (97.6%). Notably, the auxiliary mask prediction objective alone lifts the authors' RGB-only variant from 97.3% to 98.4% without any test-time visual prompt, indicating that mask supervision acts as a semantic regularizer even under standard text conditioning. Attention-map analysis shows the RGB-only model frequently highlights spurious background regions, whereas mask supervision yields consistently task-focused attention.

On RoboTwin 2.0 across six randomized tasks, MaskWAM reaches 92.2%, outperforming $\pi_0$ by 19.4 points and FastWAM by 4.5 points. The ablation pattern is informative: Mask-only (88.8%) outperforms RGB-only (87.3%), but joint prediction of both modalities maximizes performance at 92.2%, supporting the claim that paired prediction overcomes RGB's vulnerability to visual noise and masks' lack of textural context. The Mask-only result conceptually aligns with the concurrent Mask World Model work, while the full joint representation performs best.

Real-world evaluation

Eight real-world tasks on a dual-arm Xtrainer platform (RealSense D455 eye-on-base, D405 eye-on-hand) split into four language-clear tasks (~100 demonstrations each, 100 trials each) and four language-ambiguous tasks (60 trials per setting). On language-clear tasks MaskWAM averages 84.3%, versus 79.0% for FastWAM, 72.3% for $\pi_{0.5}$, and 55.8% for $\pi_0$; mask prediction improves the authors' own RGB-only variant on every task (e.g., 86%→91% on Task 1).

The language-ambiguous results carry the paper's strongest claims. In-distribution, MaskWAM attains 92.9%. A striking inversion appears among prompting baselines: the text-driven $\pi_0$-coord (41.7%) outperforms the vision-driven FastWAM-coord (26.3%), which the authors interpret as evidence that spatial disambiguation is better injected natively through the visual modality than through coordinate-augmented text. Under zero-shot generalization, MaskWAM reaches 90.4% with unseen distractors (vs. 52.9% for $C$0-mask), 74.6% on novel object instances (vs. 44.6%), and 81.7% under lighting changes. Across all settings, coordinate-text baselines collapse on precision-demanding tasks (FastWAM-coord scores 0–1.7% on Tasks 7–8).

Ablations

Three ablations isolate the design's components. First, joint RGB-mask prediction outperforms either modality alone on LIBERO (+2.0% Spatial, +1.2% Goal over RGB-only). Second, and most consequential: providing a first-frame mask prompt without future mask prediction drops success to 21.6%, versus 84.9% for the full model. This demonstrates that a visual prompt alone does not guarantee effective grounding—the predictive objective is what forces the policy to actually use the prompt. Third, replacing mask prompts with textual coordinates yields only 18.2%, with attention maps showing dispersed, background-distracted attention for coordinate prompts versus sharply localized, temporally consistent focus for mask conditioning.

Appendix experiments document negative results that reinforce the design rationale: direct mask downsampling, a separately learned 3D CNN mask encoder, ControlNet-style residual injection, and zero-initialized gated fusion for VLAs all underperform encoding rendered masks with the pretrained RGB VAE inside the main latent stream. The common failure mode is that masks treated as lightweight side signals remain untied to the future-prediction objective and are consequently underused. Robustness tests show MaskWAM tolerates moderate corruption of the first-frame prompt (erosion, dilation, shift, region dropout) with limited degradation, though severe mislocalization or heavy dropout clearly hurts—an important caveat given that prompts come from online SAM-3 segmentation.

Limitations

The paper concedes two limitations explicitly. MaskWAM depends on mask supervision during training and segmentation-derived prompts at deployment; automating reliable mask extraction in cluttered real-world environments remains non-trivial, and the annotation pipeline still requires human verification (91% of episodes needed no correction). Second, due to computational constraints, large-scale RGB-mask-action pretraining is deferred, leaving open whether mask-aware visual dynamics would improve further with foundation-scale training. An additional open question is how the approach scales beyond the eight-task real-world suite and whether the 50% mask-dropout schedule remains optimal at larger data scale.

Conclusion

MaskWAM shows that augmenting WAMs with future mask prediction serves a dual role: it regularizes attention toward task-relevant regions even without test-time prompts, and it makes first-frame mask prompting an effective spatial anchor that text-based alternatives cannot match. The central empirical finding—that mask prompting without mask prediction fails (21.6% vs. 84.9%)—establishes that conditioning and prediction must be coupled rather than treated independently.