Can World and Action Modeling Benefit Each Other under Unified Causal Training?
Yue Su, Bokai Lin, Xinkai Wang, Chubin Zhang, Xiangnan Wu, Kuan Zhang, Weiyang Jin, Jiwen Yu, Yiming Li, Cewu Lu, Lixin Yang, Xihui Liu
OneWM unifies world and action modeling to show that predicting visual consequences improves control, and vice versa.
Can training a single model to simultaneously predict future video frames and future actions improve performance on both tasks compared to training them separately?
World models predict future scenes from actions, while action models generate controls from scenes; these tasks are often treated as separate objectives, limiting their mutual reinforcement. OneWM unifies these tasks in a single causal transformer architecture by mixing training regimes where actions serve as either clean conditions for video prediction or noisy targets for action generation. This unified causal training improves policy generalization across diverse scene perturbations and enhances the fidelity of action-conditioned world predictions.
Paper Primer
The paper addresses the disconnect between forward-dynamics learning (predicting how actions change a scene) and action generation (choosing actions to achieve a goal). By treating actions as both conditions for video prediction and targets for policy learning, OneWM forces the model to learn a shared representation of robot–object dynamics.
OneWM is a causal Mixture-of-Transformers (MoT) framework: it uses separate video and action experts that exchange information through temporally causal joint attention. The core move is a unified conditioning strategy where the model is trained on a mixture of "all-clean" actions (for forward dynamics), "all-noisy" actions (for generation), and "partial-clean" prefixes (to bridge the two).
Learning action-conditioned forward dynamics improves policy performance and generalization.
OneWM outperforms action-masked and noisy-action-conditioned baselines on LIBERO-Plus across background, camera, language, and sensor-noise perturbations. Substantial gains in success rates under distribution shift, particularly under sensor noise.
Action learning improves action-conditioned world modeling.
OneWM consistently achieves the best PSNR, SSIM, and LPIPS metrics compared to a world-model-only baseline (WM) on LIBERO, LIBERO-Plus, and RoboCasa365. Consistent ranking as the top performer across all three benchmarks on paired prediction metrics.
Why does this approach use a mixture of training regimes instead of just one?
The mixture allows the model to learn both specified-action dynamics (all-clean) and action generation (all-noisy) within the same architecture. This coupling ensures that the visual expert learns how controls change the scene, while the action expert learns the scene structure needed to select those controls.
What is the scope of this framework — does it require specific environment types?
The framework is evaluated on both stationary and mobile manipulation tasks across simulated (LIBERO, RoboCasa365) and real-world environments (block stacking, insertion, breakfast arrangement), suggesting it is robust to diverse scene layouts and articulated fixtures.
Researchers can now treat world prediction and action generation as reciprocal tasks rather than separate objectives, using unified causal training to strengthen the representations underlying both.
Motivation and Overview
Frames the gap: world and action modeling lack unified causal training.
Existing work shows that video prediction can aid policy learning, yet it is unclear how world and action modeling can reinforce each other when trained together. This section isolates that gap and poses two concrete questions about reciprocal benefits.
OneWM treats actions and future video frames as two sides of the same coin, training them together so each side provides a learning signal for the other.
How does OneWM differ from prior joint video‑action models that simply share a backbone?
Prior models typically mask video tokens from attending to action tokens or treat actions as noisy denoising targets. OneWM explicitly conditions world prediction on clean actions (all‑clean setting) and also trains an action‑generation branch, so the same forward‑dynamics objective directly informs policy learning rather than being an auxiliary signal.
The core hypothesis is that world and action modeling are mutually beneficial when trained under a unified causal objective.
The OneWM Architecture
We detail the causal joint attention mechanism that couples video and action experts via a Mixture‑of‑Tokens backbone.
Predicting future video frames while simultaneously generating the corresponding actions is challenging because the two tasks must remain temporally aligned yet mutually informative. We therefore formulate a joint problem: given an observed visual prefix $\\text{V}_0$, a language instruction $\\ell$, and an action chunk $a=(a_1,\\dots,a_K)$, we predict future video latents $z_{\\text{fut}}$ and the action chunk itself.
It lets video and action experts attend to each other’s tokens, but only within a causal window that respects the temporal order of frames and actions.
Visible keys for the video query: $S_v(2)=\\{V_0, V_1, V_2, A_1, A_2\\}$.
Compute raw scores $Q_v K_v^{\\top}$ for each visible key (5 scores).
Apply softmax over these 5 scores, yielding attention weights that sum to 1.
Weighted sum of the corresponding value vectors produces $U_v$ for block 2.
The video query can incorporate both past video and action information, but it never sees future tokens $V_{>2}$ or $A_{>2}$, preserving causality.
MoT treats tokens from different modalities as a single mixed sequence, applying a shared transformer stack while respecting modality‑specific masks.
**Figure 1.** **Action conditioning for world and action modeling.** Rows denote queries and columns denote keys/values. $V_c$, $V_{1,2}$, and $A_{1,2}$ represent visual context, future video, and aligned actions, respectively. Outlined blocks compare world prediction conditioned on clean ground-truth actions (solid orange) with WAM prediction using noisy actions (hatching) or masked action keys (light gray). Percentages indicate the fraction of clean ground-truth actions. OneWM (blue) mixes all-clean, partial-clean, and all-noisy examples to examine whether world and action modeling benefit each other under unified causal training.
**Figure 2.** **Overview of OneWM.** A VAE encodes frames into latents; the observed prefix remains clean, and the video expert denoises future latents. Action chunks are embedded for the action expert. Causal joint attention couples the two experts, each receiving instruction $\ell$ through its own text cross-attention (Sec. 3.2). The lower panels illustrate all-clean, clean-prefix/noisy-suffix, and all-noisy action construction, unifying world prediction and action generation (Sec. 3.3).
We vary the length $r$ of a clean action prefix so that the same backbone can be trained on fully clean, partially clean, or fully noisy action sequences.
Input sequence becomes $[a_1,\\; a_2(0.5),\\; a_3(0.5)]$.
The noise‑level embedding is $[0,\\;0.5,\\;0.5]$, signalling which tokens are conditions.
During training, the loss is computed only on the noisy suffix $a_{>r}=\\{a_2, a_3\\}$ for the action expert.
The video expert still receives the clean $a_1$ as part of its conditioning set.
Partial‑clean conditioning forces the model to learn to use both clean and noisy action information, bridging the gap between pure conditioning and pure denoising.
We train video and action experts together using flow‑matching losses that predict the velocity from noisy inputs back to clean targets.
At inference time we run two settings. For world prediction we feed clean actions ($r=K$) and sample video latents from noise to $\\tau=0$, decoding them with the VAE. For action generation we start from noisy actions ($r=0$) and denoise them while re‑using the observed video context, enabling closed‑loop replanning.
Experimental Setup
We benchmark OneWM across simulated and real tasks, reporting success rates for each method.
We evaluate how world and action modeling affect each other under unified causal training. The study uses three simulated benchmarks—LIBERO, LIBERO‑Plus, and RoboCasa365—and three real‑world tasks: block stacking, block insertion, and breakfast arrangement.
**Figure 3.** Evaluation tasks and environments. We evaluate OneWM on LIBERO, LIBERO-Plus, and RoboCasa365 in simulation, and on three real-world tasks: block stacking, block insertion, and breakfast arrangement.
Across the simulated benchmarks, Fast‑WAM reaches a perfect 100 % overall success on LIBERO, while OneWM attains 96.8 % on the same metric. On RoboCasa365, OneWM’s composite‑unseen success drops to 6.1 %, far below the baselines. LIBERO‑Plus results are incomplete in the source, but the reported numbers show OneWM comparable to other variants on most factors.
Impact on Action Learning
OneWM beats prior policies across simulation and real‑world tasks, delivering higher success rates.
OneWM achieves 66.7% overall success on the three manipulation tasks, outperforming `WAM_mask` by 17.8%.
Figure 4 reports success rates of 66.7% for OneWM and 48.9% for `WAM_mask`, a 17.8‑point gap.
**Table 1.** Policy performance in simulation. Success rates (%) on LIBERO, LIBERO-Plus, and RoboCasa365. Bold highlights the best result among internal variants. Bkgd., Lang., and Init. denote background, language, and robot initialization, respectively; S/U and Comp. denote seen/unseen and composite tasks in RoboCasa365.
**Figure.** Success rates on three manipulation tasks; Overall aggregated performance.
World Modeling is the ability to predict future visual observations given a sequence of robot actions, effectively learning how actions transform the scene.
OneWM consistently outperforms isolated policy baselines.
Impact on World Prediction
Action learning boosts world-model predictions, with OneWM leading across PSNR, SSIM, and LPIPS.
Recall that OneWM treats actions and future video frames as dual conditions in a unified causal training framework.
Action learning improves conditional world modeling, with OneWM achieving up to 13.68 PSNR gain over the next‑best method.
Table 2 shows OneWM’s PSNR advantage across LIBERO, LIBERO‑Plus, and RoboCasa365.
Action learning augments a world model with an auxiliary task that predicts the robot’s next action, forcing the latent representation to capture control‑relevant scene structure.
How does Action Learning differ from simply adding an extra loss to a video‑prediction model?
In Action Learning the model must generate the *exact* action sequence, not just any auxiliary signal. This forces the latent space to encode the precise robot‑object relationships needed for control, whereas a generic auxiliary loss could be satisfied by irrelevant features.
**Figure 5.** Main-camera world predictions. Six cases from LIBERO and LIBERO-Plus compare GT, OneWM, and the indicated world-model baseline. Orange boxes mark the regions enlarged at $t + 32$ in each panel's final column.
Ablation: Conditioning Mixtures
Evaluating how different conditioning mixtures impact policy success and world‑prediction quality.
We vary the sampling probabilities of all‑noisy, partial‑clean, and all‑clean examples and measure both policy success and world‑prediction quality.
The mixture controls how many training examples contain pure noise, partially cleaned data, or fully cleaned actions, thereby balancing supervision signals for policy learning and world prediction.
How does the Conditioning Mixture differ from simply adding more clean data?
Adding clean data alone would bias the model toward reconstruction, but the mixture also retains noisy examples that teach the model to denoise and generate actions from stochastic inputs. The balance is what yields both good policy generalization and accurate world prediction.
**Table 3.** **Conditioning mixtures affect both tasks.** Columns specify all-noisy / partial-clean / all-clean percentages; 60/25/15 is the default. SR is success rate (%). Bold marks the best value within each benchmark.
Switching from the default 60/25/15 mixture to 50/35/15 raises LIBERO success rate by +1.0 % while leaving world‑prediction metrics essentially unchanged.
SR improves from 94.95 % to 95.95 % on LIBERO; PSNR, SSIM, LPIPS, FID, and FVD vary by less than 0.2 % across the two mixtures.
Increasing the clean‑conditioning share to 30/45/25 reduces LIBERO‑Plus success rate by -21.26 % and degrades policy generalization under distribution shift.
SR falls from 59.27 % (default) to 38.01 % on LIBERO‑Plus; other metrics remain stable, indicating the drop is driven by policy performance rather than prediction quality.
Training Protocols and Details
Training, evaluation, and qualitative details for OneWM experiments.
Simulation training for LIBERO and LIBERO‑Plus follows the Fast‑WAM recipe, jointly training on demonstrations from the Spatial, Object, Goal, and Long suites for 20 epochs with a global batch size of 128. RoboCasa365 is trained on its own demonstration dataset for 90 k optimizer steps with a batch size of 1024. Real‑world tasks (block stacking, insertion, breakfast arrangement) use 50 demonstrations each, trained for 20 epochs with batch size 96.
All internal WAM and WM ablations share this dataset‑specific protocol, differing only in action‑conditioning probabilities (60 % all‑noisy, 25 % partial‑clean, 15 % all‑clean). During deployment, policies generate action chunks from Gaussian noise using current visual observations and language instructions, reusing observed video representations for action denoising without generating future video. For world prediction, the action chunk remains clean while the model denoises future video latents.
Policy evaluation measures closed‑loop task success according to each benchmark’s criterion; LIBERO reports results for four suites and their average, while LIBERO‑Plus aggregates seven perturbation categories. RoboCasa365 reports atomic‑seen, composite‑seen, and composite‑unseen tasks, and real‑world tasks are evaluated over 15 trials with success rates as percentages.
Video evaluation constructs 100‑episode sets from expert rollouts for LIBERO/PLUS (using $\pi$0.5) and RoboCasa365 (using Xiaomi‑Robotics‑1). Multi‑view predictions include external‑camera and wrist‑camera views; prediction chunks span 32 action steps with stride 4 for LIBERO/PLUS and 16 steps with stride 2 for RoboCasa365, yielding eight future frames per chunk.
PSNR is computed as $-10 \log_{10}(MSE + 10^{-8})$ dB, where $MSE$ averages over RGB channels and pixels; higher values indicate better reconstruction. SSIM uses an $11 \times 11$ Gaussian window ($\sigma$ = 1.5) with $K_1=0.01$, $K_2=0.03$; higher is better. LPIPS employs a VGG16 backbone; lower is better. FID compares 2048‑dimensional Inception‑v3 features, and FVD uses 400‑dimensional I3D‑R50 logits; both lower values indicate closer distributional similarity.
Qualitative analysis shows that OneWM’s main‑camera predictions more closely follow object poses (e.g., moka pot, gripper height) than WMAdaLN or WM baselines, and wrist‑camera predictions preserve scene geometry (e.g., drawer position, bottle orientation) more faithfully. The discussion highlights that PSNR/SSIM/LPIPS assess per‑frame fidelity, whereas FID/FVD reflect pooled feature statistics, so lower distributional distance does not guarantee better action‑conditioned geometry. RoboCasa365 examples further illustrate OneWM’s superior tracking of forearm and gripper motion.
In the “Endless Dream” experiment on LIBERO, the loop from Eq. 6 is instantiated with the jointly trained LIBERO model and the default conditioning mixture. Each rollout starts from a main‑camera image, a wrist‑camera image, and a task instruction, runs 10 denoising steps with a fixed seed of 42, and alternates action prediction with action‑conditioned world prediction for a fixed budget of 16 rounds, 32 actions, and eight predicted frames per round.