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

Hanyang Yu; Haitao Lin; Jingbo Zhang; Wenyao Zhang; Chenghao Gu; Heng Li; Ping Tan

Source-linked AI summary

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

Hanyang Yu, Haitao Lin, Jingbo Zhang, Wenyao Zhang, Chenghao Gu, Heng Li, Ping Tan

arXiv:2606.13515v1Jun 11, 2026cs.CVcs.LGcs.RO

TL;DR

WAMs face weak semantic grounding in RGB prediction and spatial ambiguity in text conditioning. MaskWAM jointly predicts RGB, masks, and actions while using optional first-frame mask prompts, achieving strong performance across simulated and real-world tasks. Its main scope boundaries are reliance on mask supervision and segmentation-derived deployment prompts, plus deferred large-scale pretraining.

Problem
WAMs lack explicit semantic grounding in RGB predictions, while text conditioning can be imprecise for identifying targets among visually similar distractors.
Method
MaskWAM jointly models future RGB frames, task-relevant masks, and action chunks, using optional first-frame visual prompts within a unified Mixture of Transformers architecture.
Results
MaskWAM achieves 98.4% on LIBERO, 92.2% on RoboTwin, and outperforms the strongest baseline by 33.2% in real-world tasks.
Takeaways & Limitations
Future mask prediction provides object-centric semantic supervision, while first-frame mask prompts provide precise spatial anchoring for ambiguous targets.
Takeaways & Limitations
MaskWAM relies on mask supervision during training and segmentation-derived prompts during deployment, and large-scale RGB-mask-action pretraining is deferred.

Abstract

from arXiv · show

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.

1 Introduction

World-action models use future video prediction for robotic control, but RGB forecasts provide weak semantic grounding and language can leave target location ambiguous. MaskWAM unifies future mask prediction with visual prompting to improve object-focused, precise manipulation and generalization.

WAMs model future observations alongside actions, using predictive representations of physical dynamics and task-relevant temporal structure for downstream control.
RGB prediction treats regions similarly, allowing target objects to remain entangled with background clutter and motivating explicit semantic supervision.
Text conditioning lacks precision for complex spatial relationships and visually similar distractors, creating referential ambiguity in cluttered scenes.
MaskWAM unifies future mask prediction and first-frame visual prompting to provide semantic supervision and a precise spatial anchor.
MaskWAM integrates masks as explicit inputs and outputs while jointly modeling future RGB frames, task-relevant masks, and action chunks.

2 Related Work

Prior robotic world models predict future observations, while visual-prompting methods provide static spatial cues or intermediate representations. MaskWAM addresses the representation gap by jointly adding task-relevant future mask prediction to RGB futures and actions.

Video world models predict future observations alongside actions, offering richer supervision and a mechanism for reasoning about consequences before acting.
Most WAMs supervise visual futures mainly through RGB reconstruction, leaving a representation gap between high-dimensional pixels and control-relevant spatial structure.
MaskWAM adds future mask prediction for task-relevant regions and jointly trains it with RGB futures and actions.
Visual-prompting approaches use cues such as clicks, bounding boxes, or masks, while other policies predict intermediate representations including keypoints, trajectories, and object flows.

3 Method

MaskWAM jointly represents RGB, masks, and actions in a unified latent and transformer architecture. Optional first-frame masks provide spatial anchors, while joint flow matching trains visual and action predictions with separate noise schedules.

Overview: Given RGB, proprioception, language, and an optional first-frame target mask, MaskWAM predicts action chunks, future RGB frames, and future masks.
Overview: A Mixture of Transformers action expert lets video and action tokens interact through joint attention while extending world-action modeling beyond RGB-only prediction.
Future Masks as Predictive Targets: RGB and rendered mask frames share a causal 3D VAE encoder, producing aligned latent streams that preserve visual priors and encode spatial guidance.
Future Masks as Predictive Targets: RGB and mask latents are concatenated, and the pretrained patch embedding expands from C to 2C channels with added mask channels initialized to zero.
Unified Prompting: At deployment, an optional first-frame mask anchors the target, while mask dropout with p = 0.5 trains one policy for language-clear and language-ambiguous tasks.
Training Objective: The unified flow-matching objective sums RGB, mask, and action losses, synchronizing visual noise through τv while independently sampling action noise τa.

4 Experiments

MaskWAM achieves strong benchmark performance across LIBERO, RoboTwin, and real-world language-clear and language-ambiguous tasks. Ablations and attention maps associate mask supervision with more task-relevant visual grounding.

Benchmarks and Baselines: MaskWAM is evaluated against leading VLA and WAM baselines on LIBERO, and against π0 and FastWAM on randomized RoboTwin 2.0 tasks.
LIBERO: 98.4% average success on LIBERO establishes a state-of-the-art result and improves over the RGB-only variant’s 97.3%.
LIBERO: LIBERO attention maps show mask supervision shifts attention from spurious backgrounds toward task-relevant regions.
RoboTwin 2.0: 92.2% average success across six randomized RoboTwin 2.0 tasks exceeds π0 by 19.4% and FastWAM by 4.5%.
Real-world Tasks: Real-world evaluation separates eight tasks into language-clear settings and language-ambiguous settings requiring explicit spatial prompting.

RGB-only WAM

MaskWAM improves RGB-only world-action modeling by jointly predicting masks and RGB, providing semantic grounding for language-clear tasks and explicit spatial cues for ambiguous targets. Across simulated and real-world evaluations, mask prediction and conditioning improve robustness to clutter, novel objects, and appearance changes.

Language-clear tasks: Joint future mask prediction improves language-clear real-robot performance, raising Task 1 from 86% to 91% and Task 3 from 76% to 81%.The reported gains indicate that semantic grounding stabilizes rather than compromises text-conditioned control.
Language-ambiguous tasks: 92.9% average success in language-ambiguous in-distribution tasks demonstrates strong performance with explicit spatial prompting.The evaluation spans one in-distribution setting and three zero-shot generalization axes.
Zero-shot generalization: MaskWAM retains 90.4% success with distractors, 74.6% on novel instances, and 81.7% under altered lighting.These results exceed the π0-mask baselines of 52.9% and 44.6% where reported, supporting generalization across clutter, object variation, and appearance shifts.
Ablations: Mask prompting outperforms coordinate text, achieving 84.9% versus 18.2% success in the language-ambiguous ablation.Coordinate prompting produces dispersed attention, whereas mask conditioning supplies dense spatial-semantic priors.

5 Conclusion

MaskWAM augments RGB prediction with future mask prediction in a unified latent architecture, encouraging actions to focus on task-critical regions. Optional first-frame mask prompts support target disambiguation, while the combined design improves generalization across distractors, object variations, and lighting changes.

Architecture: MaskWAM jointly models future RGB frames, task-relevant masks, and action chunks within one unified world-action architecture.The model integrates masks as both explicit inputs and predicted outputs.
Conclusion: Future mask prediction encourages policy actions to ground in task-critical regions rather than relying only on RGB reconstruction.The formulation addresses distractors, object variations, and lighting changes.
Conclusion: Optional first-frame mask prompts provide precise target disambiguation for manipulation.The prompt acts as a spatial cue alongside the learned object-centric prediction objective.

6 Limitations and Future Work

MaskWAM’s deployment and training remain constrained by reliance on segmentation-derived prompts and mask supervision, while larger-scale RGB-mask-action pretraining is deferred.

Limitations: MaskWAM relies on mask supervision during training and segmentation-derived prompts during deployment.Reliable mask extraction in cluttered real-world environments remains non-trivial.
Future Work: Large-scale RGB-mask-action pretraining is deferred to future work because of computational constraints.The authors identify this pretraining as potential future support for mask-aware visual dynamics and downstream policy robustness.
Evaluation Platform: Real-world evaluation uses a Dual-arm Xtrainer robotic platform with head-mounted and hand-mounted depth cameras.The platform combines global scene understanding with localized, high-resolution visual feedback.

A.2 Task settings and evaluation in Real-world Tasks

The real-world evaluation covers eight challenging manipulation tasks, divided into language-clear and language-ambiguous settings, using success rate as the primary metric.

Task Settings: Eight representative real-world manipulation tasks are grouped into language-clear and language-ambiguous settings.The tasks include selecting, stacking, picking, placing, and grasping objects under varying ambiguity.
Data Collection: Models receive an average of 100 demonstrations per task.This collection supports the evaluation of the designed real-world manipulation tasks.
Evaluation Protocol: Language-clear tasks use 100 trials per task, while language-ambiguous tasks use 60 trials per task.The evaluation reports execution success rate as the primary performance metric.

A.3 Baseline Details for Language-Ambiguous Tasks

The baselines compare language-based coordinates, mask visual prompts, and automated mask annotation procedures for language-ambiguous manipulation.

Coordinate-Text Baselines: Coordinate-text baselines convert the target-mask centroid into normalized pixel coordinates embedded in language instructions.These prompts provide coarse spatial control for π0 on relatively simple tasks with large objects.
Mask Visual Prompting: π0-mask receives both RGB and mask images, encodes them separately with SigLIP, and fuses the embeddings by direct addition.The fused representation is passed into the VLM transformer layers.
Baseline Comparison: π0-mask performs better than π0-coord but still falls short on high-precision tasks.This comparison separates the benefit of mask prompts from coordinate-augmented language prompts.
Baseline Comparison: MaskWAM performs best, followed by π0-mask, π0-coord, and FastWAM-coord.WAM-based models appear less sensitive to coordinate-augmented language prompts, whereas MaskWAM benefits more from mask-based visual prompting.
Annotation Pipeline: The annotation pipeline uses language parsing and SAM3 tracking for language-clear tasks, but human point prompts for language-ambiguous tasks.All first-round annotations are human-verified to avoid oracle leakage and improve mask quality.
Annotation Efficiency: 91% of episodes require no human correction, and annotating 50 episodes takes approximately 3 minutes.For language-ambiguous tasks, the initial point prompt adds 5–10 seconds per episode.

B.2 Robustness to Noisy First-frame Mask Prompts

MaskWAM is tested with corrupted first-frame masks while all other inference inputs and model weights remain fixed, showing tolerance to moderate corruption but degradation under severe target-region damage.

Experimental Setup: Only the first-frame prompt mask M0 is perturbed at inference time; observations, language, proprioception, and model weights remain unchanged.This isolates sensitivity to prompt-mask quality.
Evaluation Tasks: The robustness study evaluates Task 5 with separated objects and Task 8 with densely arranged, visually similar cosmetics.Each task and mask condition uses 20 trials with diverse initial object placements.
Mask Corruptions: The study applies erosion, dilation, spatial shift, and region dropout to the online SAM3 mask.These corruptions represent under-segmentation, over-segmentation, mislocalization, and partial mask loss.
Results: MaskWAM shows limited performance degradation under mild erosion, dilation, shift, and region dropout when overlap with the original target mask remains sufficient.The x-axis measures mask IoU with the unperturbed online SAM3 mask, while the y-axis reports real-world success rate.
Results: Severe spatial shifts or heavy region dropout produce clearer performance drops by mislocalizing the prompt or removing most of the target area.The first-frame mask therefore functions as both a target-identity signal and a spatial anchor.

C Ablations on Alternative Mask-Conditioning Designs

Ablations show that mask conditioning is most effective when masks are aligned with the pretrained visual latent space and tied to future prediction, rather than added as an auxiliary pathway. Simpler encoders, side-branch injection, and zero-initialized fusion can leave the mask signal underused.

Design conclusion: The design analysis concludes that masks should be aligned with the pretrained visual latent space and coupled with an explicit future prediction objective.This conclusion summarizes why the tested lightweight or auxiliary conditioning designs underperformed.
Mask representation: Encoding rendered mask frames with the pretrained RGB video VAE outperformed direct mask downsampling and a lightweight 3D CNN mask encoder.The ablation attributes the weaker alternatives to lacking semantic and structural alignment with the video latent space.
Mask injection: ControlNet-style mask injection performed poorly because its auxiliary mask branch was not explicitly tied to the future prediction objective.The policy could therefore underuse the mask, particularly in language-ambiguous tasks.
Gated fusion: Zero-initialized gated fusion preserved pretrained behavior initially but led the model to keep the mask pathway weak and largely ignore the mask embedding during training.Zero initialization alone did not ensure fine-grained spatial grounding.
Spatial grounding: MaskWAM’s visual prompting produces more focused attention on task-relevant regions than text prompting, which can scatter attention over task-irrelevant regions.The comparison is shown for first-frame mask prompting combined with future mask prediction.