Can World and Action Modeling Benefit Each Other under Unified Causal Training?

Paper: Can World and Action Modeling Benefit Each Other under Unified Causal Training?

Authors: Yue Su, Bokai Lin, Xinkai Wang, Chubin Zhang, Xiangnan Wu, Kuan Zhang, Weiyang Jin, Jiwen Yu, Yiming Li, Cewu Lu, Lixin Yang, and Xihui Liu

Affiliations: MMLab@HKU, MVIG@SJTU, THU, SII, UCAS, and IAAI Team

Paper: https://arxiv.org/abs/2610.21095

Code: The public implementation is available at https://github.com/Interactive-Active-AI/OneWM. The project page is https://selen-suyue.github.io/OneWMNet. This analysis was checked against commit 121e2e77f52200e040a5296a3bce605ce317de44.

One-Sentence Summary

OneWM places action-conditioned future prediction and action generation in one temporally causal video-action Mixture-of-Transformers, and its three-way clean/noisy action schedule provides controlled evidence that world modeling improves policy learning while action learning improves world prediction.

Why the Question Matters

A manipulation policy must answer two intimately related questions. Given a proposed action chunk, what will happen in the scene? Given the current visual context and a language instruction, which action chunk should the robot execute? The first question is forward dynamics or action-conditioned world modeling. The second is policy learning or action generation. Both require an understanding of robot-object geometry, contact, object persistence, and the way motion changes a scene, so it is tempting to train both objectives in one video-action foundation model.

Sharing a backbone, however, does not guarantee that two objectives will help each other. The result depends on how actions enter the video branch and how observed visual features enter the action branch. If action tokens are always noisy when the video model reads them, the video branch never receives an explicit lesson in the visual consequence of a specified control sequence. If action tokens are completely masked from video prediction, the problem is avoided by removing the very pathway that could carry forward-dynamics information into the shared representation.

Many unified systems favor one of these endpoints. They train action generation through denoising, then use the video branch as a predictive prior, but the video tokens receive either no action information or only the current noisy estimate. Such a video loss can teach visual continuation that is plausible under the instruction, yet it is not the same as learning what this exact action chunk does to this exact object configuration. The distinction becomes important when the policy must adapt to shifted camera poses, layouts, backgrounds, or sensor noise.

OneWM turns this design distinction into a controlled experiment. A video expert and an action expert form a shared causal MoT. Actions can be entirely clean, entirely noisy, or split into a clean prefix and noisy suffix. The clean portion acts as a specified condition for future-video prediction; the suffix acts as a denoising target for action generation. By mixing the three cases in every training batch, the model receives both kinds of supervision while retaining one set of temporal visibility rules.

The paper asks two questions that are deliberately symmetric. First, does learning forward dynamics under specified clean actions improve action learning? Second, does learning to generate actions improve action-conditioned world modeling when inference uses the same clean action chunk? Answering both questions requires architecture-matched baselines, not just comparisons with unrelated policies. OneWM therefore compares against a masked-action world-action model, a noisy-action world-action model, and a world model that uses the same MoT but has no action loss.

Technical Preliminaries

The observed visual prefix is kept clean, while future video latents are targets. Let the context be $\mathcal{C}=(\mathbf{z}^{\mathrm{ctx}},\ell)$, the future latents be $\mathbf{z}^{\mathrm{fut}}$, and an action chunk be $\mathbf{a}$. World prediction models $p_\theta(\mathbf{z}^{\mathrm{fut}}\mid\mathcal{C},\mathbf{a})$. Action generation models $p_\theta(\mathbf{a}\mid\mathcal{C})$. The two distributions share representations but condition on different variables: the first knows the controls and predicts their visual consequences; the second knows only the state and instruction and must infer controls.

Video tokens and action tokens are organized into temporally aligned blocks. The observed visual context is $\mathcal{V}_0$, future latent frame $i$ is $\mathcal{V}_i$, and the action segment aligned with that frame is $\mathcal{A}_i$. A future video query may read visual and action blocks up to its own time. An action query may read only the observed visual context and its aligned action prefix. This asymmetry prevents future observations from leaking backward while allowing controls to explain future scene changes.

$$ \begin{aligned} \mathcal{S}_v(i)&=\mathcal{V}_{\leq i}\cup\mathcal{A}_{\leq i},\\ \mathcal{S}_a(i)&=\mathcal{V}_0\cup\mathcal{A}_{\leq i}. \end{aligned} $$

The video expert is initialized from Wan2.2-TI2V-5B and reuses its video VAE and text encoder. The action expert has modality-specific attention projections and feed-forward layers. Both experts receive the instruction through separate text cross-attention modules. This preserves language control over visual prediction and action generation without forcing the two token spaces to share one undifferentiated representation.

Method in Detail

1. Causal joint attention across two experts

MoT gives each expert its own query, key, and value projections, concatenates their keys and values, applies one structural visibility mask, and splits the resulting attended features back to each expert. The query determines which expert will consume the result, while the mask determines whether that query can see visual tokens, action tokens, or both. A single attention softmax therefore mixes the two modalities without discarding modality-specific projections.

$$ U_m=\operatorname{softmax}_{\mathcal{S}_m} \left(\frac{Q_mK^\top}{\sqrt d}\right)V,\qquad K=\begin{bmatrix}K_v\\K_a\end{bmatrix},\quad V=\begin{bmatrix}V_v\\V_a\end{bmatrix}. $$

This implementation is not merely a conceptual diagram. onewm/modules/mot.py performs the per-expert projections and joint forward pass, while onewm/wam/mot_wam.py::_build_mot_attention_mask encodes the temporal visibility relations. Tokens within the allowed future frame and action segment can interact freely; dependencies across blocks follow time. The result is a bidirectional representation inside an aligned block but a causal information path between blocks.

2. A clean prefix unifies conditioning and denoising

The central control variable is the length $r$ of a clean action prefix. When $r=K$, every action token is clean and serves as a condition for future video. When $r=0$, every action token is noisy and serves as a generation target. When $0<r<K$, the clean prefix specifies the beginning of the control sequence and the remaining noisy suffix must be inferred. The same sequence can therefore occupy both roles without introducing a second action encoder.

$$ a_k^{(r,\tau)}= \begin{cases} a_k,&k\leq r,\\ \widetilde a_k(\tau),&k>r. \end{cases} $$

Each action token also receives a noise-level embedding. Clean-prefix tokens use a zero level; all noisy-suffix tokens use the same sampled level. This marker lets the action expert distinguish a fixed condition from a prediction target even though both occur in one sequence. It also gives the video expert a consistent way to interpret partially observed controls.

Training examples are drawn from all-noisy, partial-clean, and all-clean regimes at probabilities $60\%$, $25\%$, and $15\%$. For partial-clean samples, the implementation can select an action boundary aligned with a video block and jitter it by up to two steps, using weights $[0.1,0.15,0.5,0.15,0.1]$, or sample a legal boundary uniformly. The aligned-boundary branch has probability $0.6$. The code is in onewm/training/flow.py::add_partial_clean_action_noise.

The alignment bias matters because a video frame and its action segment must refer to the same physical interval. If the boundary is always arbitrary, the video branch may learn broad scene dynamics while failing to use precise action timing. Small jitter prevents the model from depending on one deterministic cut point. Partial-clean training thus sits between the two easier endpoints: it supplies substantial control information while still requiring the action expert to infer a suffix from that prefix.

3. Flow matching with valid-target losses

Both experts use flow matching. For a clean target $\mathbf{x}$, the input at noise level $\tau$ is $\mathbf{x}(\tau)=(1-\tau)\mathbf{x}+\tau\boldsymbol{\epsilon}$, and the target velocity is $\boldsymbol{\epsilon}-\mathbf{x}$. Video and action noise levels are sampled independently. The observed context and clean action prefix remain conditions; only designated future latents and the noisy action suffix are reconstruction targets.

$$ \mathbf{x}(\tau)=(1-\tau)\mathbf{x}+\tau\boldsymbol{\epsilon}, \qquad \mathbf{u}^{*}=\boldsymbol{\epsilon}-\mathbf{x}. $$

The video loss compares the predicted velocity with the future-latent velocity. The action loss compares the predicted suffix velocity with the corresponding action suffix and therefore averages only over indices greater than $r$. An all-clean example contributes no action loss. The public configuration uses action dimension 7, chunk length 32, and equal video and action loss weights.

$$ \begin{aligned} \mathcal{L}_v&=\mathbb{E}\left[ w_v(\tau_v)\left\|\widehat{\mathbf{u}}_v- (\boldsymbol{\epsilon}_v-\mathbf{z}^{\mathrm{fut}})\right\|_{\mathrm{mean}}^2 \right],\\ \mathcal{L}_a&=\mathbb{E}_{r<K}\left[ w_a(\tau_a)\left\|\widehat{\mathbf{u}}_{a,>r}- (\boldsymbol{\epsilon}_{a,>r}-\mathbf{a}_{>r})\right\|_{\mathrm{mean}}^2 \right],\\ \mathcal{L}&=\lambda_v\mathcal{L}_v+\lambda_a\mathcal{L}_a. \end{aligned} $$

The expectation for the action term is explicitly restricted to $r<K$. Without that restriction, a prefix-only construction would contain no target suffix and the loss normalization would be ill-defined. The mean-squared norm is computed over valid target elements, so long and short suffixes do not silently change the effective scale. The video loss remains present in every regime because future observations are always predicted.

4. Two inference modes and an imagined closed loop

World prediction uses the $r=K$ endpoint. The specified action chunk stays fixed, future video latents are initialized from Gaussian noise, and the observed visual context and instruction remain unchanged. Action generation uses $r=0$. The action chunk is denoised while observed video representations are reused. In deployment, the policy does not synthesize future video for each control block; it generates an action chunk from the current observations and replans from the next observation.

The two modes can be composed into Endless Dream. At round $j$, the model first generates actions from context $\mathcal{C}_j$, then uses those exact actions as clean conditions for future prediction, then feeds only the final decoded predicted frame into the next round. The language instruction remains fixed and no new environment observation is obtained. This is a stringent test of whether the two learned modes can form a coherent internal simulator-policy loop.

$$ \begin{aligned} \widehat{\mathbf{a}}_j&\sim p_\theta(\mathbf{a}\mid\mathcal{C}_j),\\ \widehat{\mathbf{z}}_{j}^{\mathrm{fut}}&\sim p_\theta(\mathbf{z}^{\mathrm{fut}}\mid\mathcal{C}_j,\widehat{\mathbf{a}}_j),\\ \widehat{\mathbf{o}}_{j+1}&=\operatorname{last} \left(D(\widehat{\mathbf{z}}_{j}^{\mathrm{fut}})\right). \end{aligned} $$

On LIBERO, the loop runs for 16 rounds with 32 actions and eight frames per round. It produces 512 actions and a modeled horizon of 25.6 seconds. These are imagined controls and observations, not simulator-executed commands. The experiment demonstrates composability and qualitative task progression through drawer opening, pushing, and transport. It does not report task success after executing the imagined plan, so planning reliability remains a separate question.

Figure 1: The video and action experts exchange information through temporally causal joint attention. The lower panels show all-clean, partial-clean, and all-noisy action construction.

5. The method as an executable training graph

flowchart LR
  O["Multi-view observation"] --> E["Video VAE encoder"]
  L["Language instruction"] --> TV["Video text cross-attention"]
  L --> TA["Action text cross-attention"]
  E --> V["Video expert"]
  A["Action chunk"] --> R{"Clean prefix r"}
  R -->|"r = K"| C["All clean"]
  R -->|"0 < r < K"| P["Clean prefix and noisy suffix"]
  R -->|"r = 0"| N["All noisy"]
  C --> J["Temporally causal joint attention"]
  P --> J
  N --> J
  V --> J
  TV --> V
  TA --> AEXP["Action expert"]
  J --> AEXP
  J --> VP["Future latent video"]
  VP --> LV["Video flow loss"]
  AEXP --> LA["Action suffix flow loss"]
  

Evidence from Policies

The experiments cover LIBERO, LIBERO-Plus, RoboCasa365, and three real-world manipulation tasks. LIBERO measures in-distribution manipulation across spatial, object, goal, and long-horizon suites. LIBERO-Plus perturbs background, camera, language, lighting, layout, robot initialization, and sensor signals. RoboCasa365 tests stationary and mobile manipulation in complex scenes. The real tasks are block stacking, block insertion, and breakfast arrangement, chosen to stress relative placement, precise alignment, and multi-object state changes. Figure 2 summarizes the settings.

Figure 2: LIBERO, LIBERO-Plus, and RoboCasa365 simulations plus block stacking, insertion, and breakfast arrangement in the real world.

To test whether world modeling helps action learning, all internal variants jointly train video and action prediction. WAM mask blocks future video tokens from reading actions. WAM noisy lets the video branch read noisy actions. OneWM mixes the three clean-action regimes. OneWM reaches $94.95\%$ on LIBERO versus $92.45\%$ for WAM mask and $92.60\%$ for WAM noisy. Its LIBERO-Plus average rises to $59.27\%$, compared with $53.37\%$ and $52.50\%$. On RoboCasa365 it reaches $36.68\%$, compared with $28.32\%$ for WAM mask.

Fast-WAM remains stronger on in-distribution LIBERO at $97.60\%$, and the paper reports that result directly. OneWM's main advantage appears on LIBERO-Plus, where it outperforms Fast-WAM in every perturbation category and obtains the best listed overall score. The largest contrasts occur under camera, layout, and sensor noise. This pattern supports a representation benefit from learning specified-action consequences, but it does not establish universal superiority over every policy architecture.

Real-world results point in the same direction with smaller samples. OneWM reaches $73.3\%$ overall and WAM mask reaches $60.0\%$. Per-task success is $66.7\%$ versus $60.0\%$ for stacking, $20.0\%$ versus $0\%$ for insertion, and $86.7\%$ versus $40.0\%$ for breakfast arrangement. The gains are largest when the robot must track object relations over several changes. Each model is tested only 15 times per task, so these percentages are useful directional evidence rather than precise population estimates.

Method	LIBERO	LIBERO-Plus	RoboCasa365
OneWM	94.95	59.27	36.68
WAM mask	92.45	53.37	28.32
WAM noisy	92.60	52.50	Not reported
Fast-WAM	97.60	49.70	Not reported
Pi0	94.40	53.60	14.80

Table 1: Simulation policy success rates in percent. OneWM is not the best LIBERO in-distribution model, but it performs best overall on the two more distribution-sensitive comparisons shown here.

Evidence from World Prediction

The reverse comparison isolates action learning. The world-model baseline WM uses the same MoT and receives clean actions at inference, but it has no action-prediction loss. OneWM improves PSNR, SSIM, and LPIPS over WM on all three benchmarks. PSNR is $22.95$ versus $22.33$ on LIBERO, $21.05$ versus $20.54$ on LIBERO-Plus, and $25.21$ versus $24.72$ on RoboCasa365. Because the architectures match, these differences are stronger evidence for the action objective than a comparison with an unrelated video generator.

The distributional metrics are more nuanced. On LIBERO, OneWM has FID $12.45$, while WM has $11.89$; the action-trained model is therefore worse on this pooled image-feature statistic. The authors explain that PSNR, SSIM, and LPIPS evaluate each prediction against the reference produced under the same action sequence. FID and FVD compare marginal feature distributions and do not explicitly test action-outcome correspondence. Reassigning generated clips to different action sequences would not necessarily change those pooled statistics. Lower FID alone therefore does not prove stronger control following.

OneWM does improve FVD against the architecture-matched WM on all three benchmarks: $13.58$ versus $13.68$, $31.66$ versus $32.48$, and $17.74$ versus $18.55$. This combines with the consistent paired-metric gains. The paper also presents matched qualitative windows from external and wrist cameras, checking gripper position, object edges, bottle orientation, door motion, and viewpoint changes. These panels are useful because they expose action-conditioned geometry that a single FID value can obscure.

Figure 3 shows main-camera predictions and enlarges the same region at $t+32$. Figure 4 uses wrist views to inspect close-range geometry. Figure 5 extends the comparison to left, wrist, and right views in RoboCasa365. Figure 6 shows the all-imagined Endless Dream rollouts. In each comparison, the observed context, instruction, clean actions, and camera are matched across rows, so the relevant question is whether the predicted robot-object relationship advances in the right way.

Figure 3: Main-camera predictions on LIBERO and LIBERO-Plus, with matching regions enlarged at $t+32$.

Figure 4: Wrist-camera predictions under matched clean-action conditions, emphasizing local object position, shape, and pose.

Benchmark	Method	PSNR up	SSIM up	LPIPS down	FID down	FVD down
LIBERO	OneWM	22.95	0.8423	0.1458	12.45	13.58
LIBERO	WM	22.33	0.8325	0.1529	11.89	13.68
LIBERO-Plus	OneWM	21.05	0.8049	0.1935	17.01	31.66
LIBERO-Plus	WM	20.54	0.7937	0.2014	18.06	32.48
RoboCasa365	OneWM	25.21	0.8515	0.1015	9.16	17.74
RoboCasa365	WM	24.72	0.8433	0.1087	9.20	18.55

Table 2: Architecture-matched world-prediction comparison. Paired metrics and FVD favor OneWM, while LIBERO FID favors WM and prevents an overbroad claim of uniform metric improvement.

The Balancing Act Behind Mutual Benefit

The conditioning mixture is not a minor hyperparameter. With a $60/25/15$ split among all-noisy, partial-clean, and all-clean examples, OneWM reaches $59.27\%$ on LIBERO-Plus. A $50/35/15$ split improves LIBERO in-distribution success slightly to $95.95\%$ but lowers LIBERO-Plus to $56.98\%$. A $30/45/25$ split collapses LIBERO-Plus to $38.01\%$, even though several world-prediction metrics remain stable or improve slightly.

This dissociation is one of the paper's more useful findings. More clean-action supervision does not monotonically improve the policy. It strengthens reconstruction under familiar conditions, but deployment requires generating actions from noise. Preserving enough all-noisy training keeps that path well trained. The best configuration therefore balances the forward-dynamics signal against the action-generation distribution that is actually used at test time.

All noisy / partial / all clean	LIBERO success	LIBERO-Plus success	LIBERO PSNR	LIBERO-Plus PSNR
60 / 25 / 15	94.95	59.27	22.95	21.05
50 / 35 / 15	95.95	56.98	22.44	20.57
30 / 45 / 25	87.80	38.01	23.12	20.90

Table 3: Action-conditioning mixture ablation. Moderate rebalancing changes both tasks, while a large shift toward clean conditions mainly damages out-of-distribution policy performance.

Figure 5: Multi-view predictions in RoboCasa365 across left, wrist, and right cameras.

What Endless Dream Does and Does Not Show

Endless Dream alternates generated actions with clean-action video prediction. Because the last generated frame becomes the next observation, errors can accumulate without any environmental correction. The fact that 16 rounds remain visually coherent and task-directed is evidence that the two inference endpoints share compatible state representations. The blue outlines in Figure 6 mark only the initial observation; every later frame is produced by the model.

The experiment does not evaluate execution. No simulated controller receives the 512 generated actions, and no success-rate table is reported for the imagined sequence. Long-horizon contact and object-state errors may therefore be invisible in a visually plausible rollout. A stronger follow-up would close the loop with a simulator or real robot, use the learned dynamics to rank candidate controls, and measure whether imagined evaluation predicts actual success.

Figure 6: Endless Dream rollouts on LIBERO. Each round alternates action generation and clean-action world prediction; only the first observation comes from the environment.

Implementation Traceability

The public code maps cleanly to the paper's central claims. onewm/training/flow.py constructs the clean prefix and noisy suffix. onewm/wam/mot_wam.py::training_step combines the video and action flow objectives. onewm/modules/mot.py::MoT.forward performs joint attention, and the WAM mask builder enforces temporal visibility. deploy/libero/endless_dream.py::imagine alternates action generation with world prediction. The LIBERO configuration specifies seven-dimensional actions, 32-step chunks, and unit weights for both losses.

This level of correspondence matters for judging the contribution. It establishes that the clean action endpoint is a real training and inference path, not only an explanatory drawing. It also makes the controlled baselines credible: the key difference between OneWM and the same-architecture WM is the presence of the action objective, rather than a new video backbone or a different action-conditioned interface.

Limitations

The first limitation is evaluation data. The three benchmarks were designed for policy testing and do not provide dedicated validation sets for action-conditioned video prediction. The authors build video sets from 100 expert episodes collected with Pi0.5 or Xiaomi-Robotics-1. This standardizes references, but those episodes do not represent an independent held-out distribution and their policy quality affects the absolute metrics.

The second limitation is the scale of real-world evidence. The study uses 150 demonstrations across three tasks and only 15 trials per task and model. The insertion comparison, $20.0\%$ versus $0\%$, consists of a small number of successes and cannot support a precise effect-size estimate. More objects, embodiments, camera placements, and randomized initial states would be needed to establish robustness across real deployments.

The third limitation is the gap between imagination and planning. Endless Dream shows that actions and predicted observations can be composed for 16 rounds, but it does not execute those actions or compare imagined trajectories with real outcomes. The conclusion appropriately identifies planning and policy improvement as future work. Until that step is tested, mutual training benefit should not be confused with a validated model-based controller.

The fourth limitation is hyperparameter and compute coverage. Only three mixture ratios are reported, while model size, training budget, chunk length, and action head are fixed. The balance may shift with more data or longer horizons. The paper also does not provide a detailed latency and throughput study, which matters because a video foundation model with joint attention may be expensive for high-frequency robot control.

Finally, success and fidelity are not interchangeable. A model can improve task success because the action expert becomes stronger, while the video expert changes only modestly, or it can produce sharper video without capturing action-dependent geometry. The architecture-matched comparisons and qualitative panels reduce this ambiguity, but a systematic action-counterfactual benchmark would be a cleaner test of whether predicted futures respond correctly when the control sequence is changed.

Conclusion

OneWM makes a focused claim: actions can serve as both specified conditions and learning targets inside one temporally causal model. The experiments support reciprocal benefit under that construction. Clean-action forward dynamics improves policy performance and distribution-shift generalization. Action generation improves action-conditioned future-video prediction when both models use the same clean-action inference endpoint. The mixture study then shows that the benefit depends on preserving enough noisy-action training for the policy path.

The next research step is to connect this mutual training signal to decision quality. A learned world model becomes a true planner only when its imagined outcomes are accurate enough to rank candidate actions and when closed-loop observations can correct accumulating error. OneWM provides a credible architecture and a controlled starting point for that next experiment, while its reported limitations keep the current claim at the level of reciprocal representation learning.

Key Takeaway

Clean actions teach the video branch what control changes, while action generation forces the same representation to decide what must change next. OneWM's reciprocal benefit comes from applying both signals to one shared causal structure.