PaperlayerView exact revision on arXiv

Source-linked AI summary

UniW(A)M: A Unified Mixture-of-Experts Framework for World and World-Action Models

Jingbo Xu, Jianjie Fang, Zhuohang Li, Zihao Song, Zhaolu Wang, Xvyuan Liu, Xin Zhang, Haisheng Su, Yinzhou Tang, Chen Gao, Wei Wu, Yong Li

arXiv:2610.18697v1cs.ROcs.AIcs.CVcs.LG

TL;DR

World models predict future environments and world-action models generate executable actions, but they are typically trained separately despite learning from embodied trajectories. UniW(A)M jointly trains action-conditioned world modeling and four WAM interfaces with reusable experts, then uses PAIR to select interfaces by execution chunk. It reports strong results across WorldArena and real-world manipulation, with PAIR also improving OOD success.

  • Problem

    World models and WAMs learn from embodied trajectories but are typically trained in isolation, leaving limited reuse of dynamics knowledge between prediction and control.

    Source: §1 · p. 1

  • Method

    UniW(A)M jointly trains action-conditioned world modeling and four WAM interfaces with shared and reusable experts, while PAIR selects a WAM interface for each chunk with the backbone frozen.

    Source: §1 · p. 2, §3.3 · p. 5

  • Results

    Across three WorldArena tasks UniW(A)M leads action-conditioned world modeling; on three real-world manipulation tasks, it achieves 88.3% average success versus 76.7% for the strongest baseline.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The results support joint learning and adaptive deployment of complementary world and action modeling capabilities.

    Source: §5 · p. 9

Abstract

from arXiv · show

World models (WMs) and world-action models (WAMs) focus on future world prediction and executable action generation, respectively. Although these models learn shared physical dynamics and interaction semantics from embodied trajectories, they are typically trained in isolation, which restricts knowledge reuse between prediction and control. The central challenge is therefore to develop a unified learning framework that shares these dynamics while preserving the distinct modeling capabilities required by each paradigm. We propose UniW(A)M, a mixture-of-experts world-(action) model built upon diffusion Transformers that unifies action-conditioned world modeling and diverse WAM variants. Shared experts capture common dynamics, while atomic capability experts are reused across interfaces so that each interface keeps its own conditional pathway. At deployment, a plug-in adaptive interface router (PAIR) learns from a small subset of target-task chunk scenarios to select a suitable WAM interface for each execution chunk, balancing task success and execution efficiency while keeping the backbone frozen. We evaluate UniW(A)M on LIBERO, RoboTwin 2.0, and three real-world manipulation tasks, where it achieves 88.3% average real-world success versus 76.7% for the strongest baseline. On WorldArena, it leads action-conditioned world modeling on all three tasks. AC-WM co-training contributes to action generation on every action interface and PAIR raises OOD success from 63.3% to 76.7%.

1 INTRODUCTION

World models and world-action models offer complementary prediction and action roles, while unified learning must reuse shared dynamics without erasing distinct conditional pathways. UniW(A)M jointly trains these roles and adapts WAM interface choice by execution chunk.

1 INTRODUCTION

  • Existing unified models demonstrate shared modeling and modular specialization, but organizing specialized computation around the conditional dependencies of jointly trained video–action interfaces remains a design issue.World models predict environment evolution, whereas WAMs use future-world representations to support action generation.Source: §1 · p. 1
  • PAIR learns from a small subset of target-task chunk scenarios to select among four WAM interfaces at each execution chunk while keeping the backbone frozen.Its choices span fast pathways that omit future-video prediction and interfaces that explicitly predict video as interaction demands change.Source: §1 · pp. 2–3
  • UniW(A)M achieves 88.3% average real-world success versus 76.7% for the strongest baseline and leads action-conditioned world modeling on all three WorldArena tasks.AC-WM co-training contributes to every action interface, and PAIR raises OOD success from 63.3% to 76.7%.Source: §1 · p. 2
  • Shared dynamics experts and reusable atomic capability experts support joint training while video–action masks preserve each interface’s conditional pathway.CATR routes token groups to eligible experts for roles including action-free video prediction, direct action generation, action-conditioned video prediction, and video-conditioned action generation.Source: §1 · p. 2

2 RELATED WORK

Prior work unifies world prediction and action generation through expert mixtures, shared Transformer backbones, and mixture-of-experts designs, while leaving mutual reinforcement among modes comparatively unexplored.

Action-conditioned world models

World action models

  • WAM interfaces vary in generation targets and conditional structure, but are typically trained and deployed in isolation, limiting sharing of specialized action-generation capabilities.Joint-WAM predicts video and actions concurrently, visual-to-action interfaces predict visual futures before actions, and action-centric interfaces prioritize actions (Ye et al., 2026b; Li et al., 2026b; Ye et al., 2026a).Source: §2 · p. 3

Unified world-action modeling

3 METHOD

UniW(A)M represents action-conditioned world modeling and four WAM variants as interfaces over shared trajectory data, combining distinct conditional pathways, role-aligned experts, and chunk-level interface selection.

3.1 UNIFIED MODELING OF HETEROGENEOUS VIDEO-ACTION MASK RELATIONS

  • Five interfaces differ in whether future actions are conditions or targets and in future-token visibility, composing distinct predictive roles across video and action.Joint-WAM denoises video and actions jointly; IDM-WAM predicts video then actions; Action-Centric-WAM uses video prediction as auxiliary supervision; Action-Only-WAM omits future-video prediction at deployment; AC-WM predicts video given external actions.Source: §3.1 · p. 4

3.2 ARCHITECTURE OF UNIW(A)M

  • A shared diffusion Transformer uses interface-specific attention masks and shared and atomic experts, with CATR restricting eligible specialists by modality and predictive role.Tokens within a target group share learned mixture weights while retaining individual expert outputs.Source: §3.2 · pp. 4–5, Figure 2 · p. 5

3.3 OUT-OF-DISTRIBUTION DEPLOYMENT WITH PLUG-IN ADAPTIVE ROUTING

  • PAIR trains a classifier on frozen contextual features against MAE-aware targets, favoring a faster interface when its MAE is within 1% of the minimum.At deployment it selects the highest-probability interface for each action chunk, then updates the observation history before the next selection.Source: §3.3 · p. 6
  • PAIR is invoked at chunk transitions; routing averages approximately 63 ms in the stated benchmark, including shared-expert feature computation but excluding observation and language encoding.The router ensemble prediction itself takes approximately 2 ms.Source: §3.3 · p. 6

4 EXPERIMENTS

Across LIBERO, RoboTwin 2.0, WorldArena, and three real-world tasks, UniW(A)M is evaluated for action generation and future-world prediction. Results cover generalization, chunk-level routing, and ablations of co-training and expert routing.

4.1 EXPERIMENTAL SETUP

  • Evaluation covers LIBERO, RoboTwin 2.0, WorldArena, and cloth folding, box folding, and tabletop cleaning; OOD shifts retain task goals while changing object appearance or scene conditions.The reported shifts include garment colors and sizes, box types, and cluttered backgrounds.Source: §4.1 · p. 6

4.2 OVERALL EVALUATION

  • 97.5% average success on LIBERO (99.2% on Spatial) accompanies 84.38% C2C and 66.60% C2R success on RoboTwin, where UniW(A)M leads C2R and ranks second in C2C.RoboTwin results use clean-only training; C2C evaluates clean scenes and C2R randomized scenes.Source: §4.2 · pp. 6–7, Figure 4 · p. 7
  • 88.3% average real-world success for IDM-WAM exceeds WorldScape Policy 2.0’s 76.7% across cloth folding, box folding, and tabletop cleaning.IDM-WAM reaches 90%, 85%, and 90% on those tasks, respectively.Source: Table 2 · p. 7, §4.2 · p. 7
  • UniW(A)M achieves the highest aggregate EWMScore-P on all three WorldArena tasks, leading the compared methods in action-conditioned world modeling.The reported evaluation covers action-driven and interaction metrics.Source: §4.2 · p. 7, Table 3 · p. 8

4.3 ANALYSIS OF OUT-OF-DISTRIBUTION SCENARIOS

  • PAIR raises average OOD success from 63.3% to 76.7% using routing-only adaptation with the backbone frozen; it favors action-centric interfaces during approach and Joint-WAM or IDM-WAM during contact and fine manipulation.Adaptation uses 2,000 target-task chunk scenarios per task.Source: §4.3 · p. 8

4.4 PARETO ANALYSIS OF SUCCESS RATE AND INFERENCE LATENCY

  • 154.3 ms and 225.8 ms per configured step for Action-Centric-WAM and Action-Only-WAM compare with 428.1 ms and 532.7 ms for Joint-WAM and IDM-WAM; UniW(A)M + PAIR reaches the upper-left Pareto frontier.The routed system reports 76.7% success and 308.9 s average completion time over successful trials.Source: §4.4 · p. 8, Figure 5 · p. 9

4.5 ABLATION STUDIES

  • Removing AC-WM co-training lowers success across all four interfaces; CATR raises Action-Centric-WAM and Action-Only-WAM from 10.0% to 81.7% and 80.0%, respectively.AC-WM removal reduces average success by 15.0 points for Joint-WAM and IDM-WAM, and by 31.7 and 43.3 points for Action-Centric-WAM and Action-Only-WAM. With matched training steps, learning-rate settings, and data volume, the dense DiT baseline peaks at 16.7% versus UniW(A)M’s 88.3%.Source: §4.5 · p. 9

5 CONCLUSION

The conclusion presents UniW(A)M as a shared framework for action-conditioned world modeling and four WAM interfaces, with AC-WM co-training, CATR, and PAIR supporting joint learning and adaptive deployment. The appendix formalizes interface-specific conditional relationships and distinguishes those factorizations from the way predictions are realized or sampled.

5 CONCLUSION

  • The conclusion links video–action masks and capability-aligned expert routing to unified world modeling and four WAM interfaces, and reports benefits from AC-WM co-training, CATR, and frozen-backbone PAIR.The stated results support joint learning and adaptive deployment of complementary world and action modeling capabilities.Source: §5 · p. 9

A.1 CONDITIONAL FORMULATIONS OF THE FIVE INTERFACES

  • C_t denotes each interface’s observed context, including visible observation history, language, and robot-state inputs where used; AC-WM omits state tokens.The conditional factors express modeled relationships, while attention masks and denoising schedules specify their realization.Source: §A.1 · p. 13
  • Joint-WAM factorizes future video and action into a video marginal and an action conditional, but mutual token visibility during denoising allows video representations to use evolving action-token information.The factor labels identify predictive roles, not distributions independently parameterized by individual FFN experts.Source: §A.1 · p. 13
  • For IDM-WAM, video is generated from observed context before actions are inferred from that rollout, and a one-way mask blocks dependence on future action tokens.The factor order follows this generation procedure.Source: §A.1 · p. 13
  • One interface adds action-conditioned video training but samples only actions at inference, leaving its WM factor inactive.Action-Centric-WAM defines actions from observed context and models future observations conditioned on those actions.Source: §A.1 · p. 13
  • Action-Only-WAM shares observed context without cross-conditioning between future modalities, whereas AC-WM conditions future-video prediction on supplied actions and has no action-generation factor.For Action-Only-WAM, the video factor provides auxiliary training supervision; at deployment, future video is not generated.Source: §A.1 · p. 13

A.2 TRAINING AND INFERENCE VISIBILITY MASKS

Visibility masks define which modalities each interface may use during training and inference, while shared and capability-specific experts implement those pathways. The appendix also specifies how trajectory windows and routed token groups preserve those dependencies.

A.2 TRAINING AND INFERENCE VISIBILITY MASKS

  • Joint-WAM couples future video and action, IDM-WAM generates video before conditioning action generation on it, AC-WM predicts video from supplied actions, and the two fast interfaces omit future-video prediction at inference.Action-Centric-WAM retains observed context for action generation; Action-Only-WAM exposes neither future modality to the other during training.Source: Figure 6 · p. 14, Figure 7 · p. 14, Figure 10 · p. 15

A.3 TRAJECTORY ORGANIZATION AND TEACHER FORCING

  • The shared trajectory window contains 33 RGB frames, 96 action vectors, four robot states, and language; teacher-forced clean history supplies context rather than extra prediction targets.The initial frame is conditioning-only, and target blocks access clean video prefixes without receiving a clean copy of their own future target.Source: §A.3 · pp. 15–16, Table 5 · p. 16

A.4 CAPABILITY-ALIGNED TOKEN ROUTING

  • Shared experts remain active across interfaces, while atomic experts follow each interface’s predicted modalities and conditioning pathway.Action-Centric-WAM deactivates WM at inference; Action-Only-WAM retains VGM for its initial conditioning frame, with future action tokens using shared and VLA experts.Source: §A.4 · p. 16, Figure 11 · p. 17
  • CATR learns mixture weights only among each token group’s eligible experts, while shared-only groups use a fixed weight of one.Weights are shared within a routed group, but tokens retain their individual expert outputs.Source: §A.4 · p. 17
  • CATR routes expert computation within an interface, whereas PAIR selects the action-generating interface at deployment.For Action-Only-WAM’s seed group, pooled frame-token states determine a shared/VGM mixture; Action-Centric-WAM uses shared alone.Source: §A.4 · p. 18

B.1 MAE-BASED TRAINING TARGETS

PAIR learns chunk-level interface choices from offline target-task scenarios using a frozen UniW(A)M backbone. Its selections vary across execution phases, favoring fast pathways during approach and video-aware interfaces during finer manipulation.

B.1 MAE-BASED TRAINING TARGETS

  • Each task contributes 2,000 sampled 24-step chunk scenarios, with all four frozen interfaces evaluated against the same recorded reference action chunk.Scenarios come from 500 OOD episodes per task; the 6,000 selected scenarios are approximately 0.0274% of the candidate-window pool, not an equivalent reduction in episode acquisition.Source: §B.1 · pp. 18–19
  • PAIR trains on normalized action errors, assigning soft targets while preferring the lowest-latency fast interface among near-best candidates when one qualifies.At deployment, the learned router selects without ground-truth actions or online MAE comparisons.Source: §B.1 · p. 19

B.2 ROUTER INPUTS AND CHUNK-LEVEL SELECTION

  • PAIR combines video, CLIP-derived semantic, robot-state, language, and chunk-position features; its first-chunk routing call takes 63 ms on average under the reported timing setup.The timing includes input preparation, shared-expert feature computation, classification, and selection, but excludes observation and language encoding and subsequent action generation.Source: §B.2 · pp. 19–20

B.3 INTERFACE SELECTION ACROSS EXECUTION PHASES

  • During Reach, the two fast interfaces account for 60%–75% of selections; Joint-WAM and IDM-WAM together account for approximately 60%–70% during Contact and 70%–90% during Core.The phase-conditioned frequencies are reported across three real-world tasks.Source: §B.3 · p. 20, Figure 12 · p. 20

C BENCHMARKS AND EVALUATION PROTOCOLS

The evaluation covers action generation on LIBERO, RoboTwin 2.0, and real-world tasks, alongside action-conditioned world prediction on WorldArena. Real-world uncertainty estimates are based on 20 trials per task and are reported with boundaries on their interpretation.

C BENCHMARKS AND EVALUATION PROTOCOLS

  • Action generation is evaluated on LIBERO, RoboTwin 2.0, and three real-world tasks, while WorldArena assesses action-conditioned future-world prediction across visual, motion, consistency, physical, 3D, and controllability metrics.Real-world tasks use 20 trials each; RoboTwin reports mean success over 50 tasks under both clean and randomized conditions.Source: C BENCHMARKS AND EVALUATION PROTOCOLS · p. 21

C.1 UNCERTAINTY IN REAL-WORLD SUCCESS RATES

  • Real-world per-task success rates are summarized with two-sided 95% Wilson intervals from 20 trials, assuming approximately independent Bernoulli outcomes with a common probability within each setting.Intervals quantify finite-trial uncertainty for the evaluated model and setting, not variation across training runs.Source: §C.1 · pp. 21–22, Table 8 · p. 22, Table 9 · p. 22
  • One additional success changes a task’s observed rate by five percentage points, so small point-estimate differences should not be treated as stable rankings.The reported intervals are marginal, not simultaneous bounds or tests of method differences.Source: §C.1 · p. 22

D.1 TASK SEQUENCES ACROSS WAM INTERFACES

The appendix presents recorded sequences for the three real-world manipulation tasks across four WAM interfaces. Each selected frame combines scene-level and local interaction views.

D.1 TASK SEQUENCES ACROSS WAM INTERFACES

  • The recorded sequences cover Folding Clothes, Folding Box, and Tabletop Cleaning for Joint-WAM, IDM-WAM, Action-Centric-WAM, and Action-Only-WAM.Each interface group contains ten selected frames, read left to right across the first row and then the second.Source: §D.1 · p. 23, Figure 13 · p. 23, Figure 14 · p. 24

D.2 CASE STUDY: ROUTING AT CRITICAL MANIPULATION STEPS

The case study shows how local execution errors can disrupt cloth folding, box folding, and tabletop cleaning. PAIR usually selects Joint-WAM and IDM-WAM for the observed fine-manipulation steps, though its context-based preference does not guarantee an optimal action for every chunk.

D.2 CASE STUDY: ROUTING AT CRITICAL MANIPULATION STEPS

  • Small release, contact, and wrinkle-recovery errors can disrupt successive stages of cloth folding.Examples include failed garment releases, a pull-back without gripper contact, and lifting the garment to restart after a wrinkled intermediate result.Source: §D.2 · p. 25
  • Incorrect flap placement and uncoordinated pushes prevent completion of the corresponding box folds.The cases involve pushing either side flap into the intended position.Source: §D.2 · p. 25
  • In tabletop cleaning, inaccurate grasping, slippage during transfer, and failed basket placement are observed execution errors.Figure 16 contrasts successful operations in the upper row with corresponding failures below, across distinct cases read left to right.Source: §D.2 · p. 26, Figure 16 · p. 26
  • PAIR usually selects Joint-WAM and IDM-WAM for the observed fine-manipulation operations, suggesting chunk-level selection can reduce the impact of one interface’s local weaknesses.The router expresses a preference from current context; it does not guarantee an optimal action for every chunk.Source: §D.2 · p. 26

E SIMULATION BENCHMARK RESULTS

The simulation results cover action-generation success on LIBERO and RoboTwin 2.0, while the WorldArena tables report component scores across six evaluation dimensions. UniW(A)M performs strongly on LIBERO and on RoboTwin’s randomized-scene evaluation.

E SIMULATION BENCHMARK RESULTS

  • 97.5% average success on LIBERO, with above 95% in every suite and 99.2% on Spatial, 0.6 percentage points above π0.5 in average success.Table 10 reports success rates across four suites and their arithmetic mean.Source: E SIMULATION BENCHMARK RESULTS · p. 27, Table 10 · p. 27
  • 66.60% C2R success on RoboTwin 2.0 leads the compared methods and exceeds π0.5 by 20.60 percentage points; its C2C success is 84.38%, near WorldScape Policy 2.0’s 84.92%.Table 11 reports success rates averaged across 50 tasks separately for C2C and C2R, using clean-only training.Source: E SIMULATION BENCHMARK RESULTS · p. 27, Table 11 · p. 27

F.1 RESULTS BY EVALUATION METRIC

  • Tables 12 and 13 report WorldArena component scores across visual quality, motion quality, content consistency, physics adherence, 3D accuracy, and controllability for three tasks.The tables mark the highest score for each metric within each task, including ties; no component-score values are provided in these passages.Source: §F.1 · p. 27, Table 12 · p. 28, Table 13 · p. 28

F.2 QUALITATIVE COMPARISON OF FUTURE-WORLD PREDICTION

The qualitative comparisons examine whether AC-WM predictions retain object persistence, appearance consistency, and local interaction details across three manipulation tasks. The accompanying training and implementation passages describe uniform interface sampling, interface-specific objectives, and initialization and optimization settings.

F.2 QUALITATIVE COMPARISON OF FUTURE-WORLD PREDICTION

  • The video comparisons use matched timestamps across eight frames per model, arranged in two rows of four and read left to right, then top to bottom.Red circles mark selected baseline prediction errors.Source: §F.2 · p. 28
  • UniW(A)M more consistently retains garment silhouette, dark fabric, and folds around the grippers than the compared predictions in the selected Folding Clothes frames.CtrlWorld loses garment regions and shows inconsistent colors, while Cosmos-Predict 2.5 shifts the garment toward a pale bluish appearance and distorts folds.Source: §F.2 · p. 28, Figure 17 · p. 29
  • In Folding Box, UniW(A)M retains more discernible cardboard edges and gripper–object boundaries; the comparison mainly concerns local detail, not task sequence.The figure comparison emphasizes cardboard structure and wrist-view contact details.Source: §F.2 · p. 29, Figure 18 · p. 30
  • In Tabletop Cleaning, UniW(A)M preserves more recognizable bottle lettering and label structure, while the baselines show lost detail, color drift, or shape distortion.The selected comparisons also note more stable basket structure and background content in UniW(A)M predictions.Source: §F.2 · p. 30, Figure 19 · p. 31

G.1 UNIFORM INTERFACE SAMPLING

  • Five interfaces are sampled uniformly at probability 0.2 each during joint training, without a manually tuned sampling schedule.The four-WAM-only ablation samples each remaining interface at 0.25, giving each 25% more expected updates when total optimization steps are fixed.Source: §G.1 · p. 31

G.2 INTERFACE-SPECIFIC TRAINING OBJECTIVES

  • All five interfaces use flow matching with conditioning determined by visibility masks; AC-WM predicts future video conditioned on clean ground-truth actions, while Action-Centric-WAM and Action-Only-WAM retain video loss during training.The latter two block different cross-modal pathways and neither predicts future video at deployment; Action-Centric-WAM has inactive WM, while Action-Only-WAM retains VGM for the initial conditioning frame.Source: §G.2 · pp. 31–32
  • Joint-WAM permits joint video–action denoising, whereas IDM-WAM uses interleaved teacher forcing and may noise video conditioning copies with probability 0.5.Both use the same loss composition, but IDM-WAM keeps action conditioning copies clean and supervises only noisy target copies.Source: §G.2 · p. 32

G.3 INITIALIZATION AND OPTIMIZATION SETTINGS

  • MoE sublayers copy the corresponding pretrained dense feed-forward network into five independently trainable experts, with equal initial router weights preserving the feed-forward output at initialization.The stated identity concerns the feed-forward replacement under identical inputs, not equivalence across interfaces with different attention masks.Source: §G.3 · p. 33

G.4 ARCHITECTURE, ROUTING, AND TRAINING-PARADIGM ABLATIONS

The ablations show that CATR supports substantially stronger performance across the four WAM interfaces than dense or sample-level routing configurations. Including AC-WM in training improves every interface–task pair, even though each WAM interface receives fewer expected updates.

G.4 ARCHITECTURE, ROUTING, AND TRAINING-PARADIGM ABLATIONS

  • The dense-versus-MoE comparison is not strictly parameter-matched, but sample-level routing and CATR use identical expert architectures and parameter counts.The latter comparison evaluates the combined token-group routing design rather than isolating pooling granularity alone.Source: §G.4 · p. 33
  • CATR raises the four interfaces to 80.0%–88.3% average success, while sample-level MoE routing leaves both action-only pathways at 10.0%.CATR reaches 81.7% and 80.0% for Action-Centric-WAM and Action-Only-WAM; the dense configuration reaches at most 16.7%. Table 15 compares the configurations.Source: §G.4 · p. 34, Table 15 · p. 34
  • Including AC-WM improves every interface–task pair; average success rises from 71.7%, 73.3%, 50.0%, and 36.7% to 86.7%, 88.3%, 81.7%, and 80.0% across the four WAM interfaces.The largest box-folding gains are 30% to 75% for Action-Centric-WAM and 10% to 70% for Action-Only-WAM, which omit future-video prediction at deployment.Source: §G.4 · p. 34

Related papers

Loading 2610.18697v1…