cs.ROOct 19, 2026

Training Controllable World Models via Action Influence Maximization

Authors: Samuel M. Bateman, Tenny Yin, Chongyi Zheng, Lei Huang, Brian Wang, Benjamin Eysenbach, Jaime Fernández Fisac, Dhruv Shah

Organizations: Princeton University · Columbia University

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

Action-conditioned video generative models trained on successful, goal-directed demonstrations tend to ignore the commanded action, rendering a plausible expert-like future regardless of what is asked of them. We trace this failure to training data that lack sufficient action influence. We analyze this metric and prove that it upper bounds the action influence that any model trained on that data can express. Expert demonstrations, whose narrow action distributions let next states be predicted without attending to the action, set a low ceiling that no amount of model capacity alone can raise. Raising it requires collecting different data. We introduce Action Influence Maximization (AIM), an algorithm for actively collecting training data for improving dataset action influence. We compare the performance and data-efficiency of AIM against exploration heuristics and intrinsic-motivation methods on simulated navigation and manipulation tasks. In pixel-based manipulation, AIM is the only method that consistently improves controllability: 500 actively collected episodes added to 3k demonstrations match a world model trained on 300k expert episodes on end-effector controllability, an 85×85\times reduction in data.

Figures & tables

Figure 1: Action-Conditioned Video Generative Models trained on near-expert data suffer from expert success bias and low controllability. The center panel shows a world model trained on expert data ignoring input actions and hallucinating a state closer to success. The right panel shows a world model trained with AIM data which remains faithful to the commanded action.
Figure 2: AIM is an algorithm that uses an information maximization objective to collect high-utility experimental transitions to improve controllability of learner policies similar to the prior demonstrations.
Figure 3: Action information and fidelity. Explained action information correlates with state-extracted simulation fidelity in Four Rooms, across dataset scales ( r=0.979r=0.979 ).
Figure 4: State-Action Coverage. Near-expert policies concentrate coverage around successful task execution, while random policies provide broader but less structured coverage. AIM targets action-informative regions useful for improving world model controllability.
Figure 5: Improving Controllability via Active Data Collection. Across discrete-action navigation (FourRooms: left) and continuous-action manipulation (Push-T: center, right), AIM improves controllability with just a few hundred episodes. While Random is ideal when uniform state sampling is possible (FourRooms), AIM approaches this performance with more data; in more realistic scenarios (Push-T), AIM is the only method that improves controllability. Most notably, AIM with just 500 new episodes matches the performance of a world model trained with 300K expert episodes.
ExperimentBC Recall (H=33, 95% CI)Preference ELOBC Failure Recall (H=33, 95% CI)Preference ELOPush Away Recall (H=33, 95% CI)Preference ELO
 EETEEJoint EETEEJoint EETEEJoint 
PPO-300K.275 ±\pm .013.400 ±\pm .024.224 ±\pm .012843 ±\pm 155.554 ±\pm .019.491 ±\pm .028.432 ±\pm .0261040 ±\pm 185.332 ±\pm .023.282 ±\pm .026.198 ±\pm .018816 ±\pm 128
PPO-3K.267 ±\pm .011.322 ±\pm .022.203 ±\pm .011732 ±\pm 155.517 ±\pm .020.424 ±\pm .030.378 ±\pm .028824 ±\pm 171.297 ±\pm .023.190 ±\pm .017.147 ±\pm .015889 ±\pm 132
PPO-3K+AIM500.485 ±\pm .014.403 ±\pm .023.344 ±\pm .0181190 ±\pm 167.564 ±\pm .018.454 ±\pm .031.411 ±\pm .0291049 ±\pm 185.375 ±\pm .016.201 ±\pm .020.165 ±\pm .0161029 ±\pm 134
PPO-3K+Recovery500.272 ±\pm .012.378 ±\pm .026.215 ±\pm .012979 ±\pm 155.545 ±\pm .022.455 ±\pm .030.406 ±\pm .027989 ±\pm 200.314 ±\pm .023.216 ±\pm .020.169 ±\pm .016735 ±\pm 128
PPO-3K+ICM500.222 ±\pm .011.393 ±\pm .025.186 ±\pm .011868 ±\pm 145.506 ±\pm .022.441 ±\pm .029.383 ±\pm .027934 ±\pm 186.267 ±\pm .020.242 ±\pm .023.158 ±\pm .016924 ±\pm 141
Table 1: Push-T collection comparisons and AIM ablations. The row groups compare expert data scale and 500-episode additions to PPO-3K (top), reuse of the same AIM500 with Recovery-3K (middle), and AIM component ablations (bottom). BC-Failure interleaves BC and Push-Away behavior to introduce anomalous contacts. Joint recall is the per-transition minimum of end-effector and Tee recall, averaged over evaluation transitions. Higher recall and preference scores are better.
Figure 6: Data Scaling Trends for World Models: We compare the simulation fidelity of world models trained on near-expert data (different types of PPO policies), recovery data, and random data on FourRooms [ 76 ] and Push-T [ 83 ] environments. Four Rooms Sim Realism estimates predicted states through maximum likelihood and enumerates the state-action space exhaustively for full evaluation reference distribution coverage. For Push-T, such exhaustive coverage is infeasible, so instead we evaluate the pixel recall of the ground truth sim end effector mask on H=10 horizon rollouts with on-policy Behavior Cloning and off-policy "Push Away" expert rollouts.
Figure 8
Expert Policy NameSteps TrainedEntropyEval Return
POMDP PPO300M0.0540.657
Early Stop PPO50M0.2520.587
State Oracle PPO300M0.0220.839
Table 2: Four Rooms Policy Statistics. All models are trained with PPO [ 78 ] in Stoix [ 84 ] .
 Softmax BCDiffusion BCSoftmax IQL+AWRDiffusion IQL+AWR
Task
EnvironmentFourRooms [ 75 , 76 ]Push-T [ 83 ]FourRoomsPush-T
Action space6 discrete6-D cont.6 discrete6-D cont.
Architecture
Model dim256256256256
Encoder layers612612
Table 3: Policy hyperparameters. All models share a frozen, per-domain–finetuned SD-VAE encoder [ 20 ] (bf16 activations, fp32 params), 144×144144{\times}144 inputs encoded to 18×18×418{\times}18{\times}4 latents ( 8181 tokens/camera), 2D sin–cos position embeddings [ 85 , 86 ] , and pre-LN transformer blocks with QK-norm [ 87 ] and SwiGLU FFNs [ 88 ] . IQL [ 71 ] variants add a cross-attention decoder on the shared BC backbone. Only the VAE is frozen across tasks.
HyperparameterValue
Initializationsd-vae-ft-mse [ 20 ]
Finetuned scopedecoder + post-quant conv (49.5M); encoder frozen
Image resolution144×144×3→[−1,1]144\times 144\times 3\to[-1,1]
Latent grid18×18×418\times 18\times 4 ( 8×8\times downsample)
Latent scaling0.182150.18215
LossMSE + LPIPS (VGG) [ 91 ]
Table 4: SD-VAE [ 20 ] decoder-only finetuning hyperparameters. One recipe produces the per-domain tokenizers shared by all policy and world models; only augmentation and frame budget differ across domains.
Figure 7: Latent Diffusion Forcing Model Architecture: We train DreamerV4-like [ 26 ] latent diffusion forcing [ 18 ] action-conditioned causal autoregressive video world models with VideoRoPE [ 98 ] positional encodings and Stable Diffusion VAE [ 20 ] latent representations. These represent a strong world model architectural baseline which is held fixed in architecture and training recipe across generated datasets. These models generate human interpretable videos which help to identify degradations in the model, allowing us to focus on the impact of the data the model is trained on rather than specific architectural details.
 Latent Diffusion WM [ 77 ]Diffusion Forcing WM [ 26 , 18 ]
Setting
DomainFourRooms [ 75 , 76 ]Push-T [ 83 ]
Predictionsingle-transition next-framesequence-level + AR rollout
Architecture
Model dim512512
Layers12 enc. + 12 dec.12 (9 spatial / 3 temporal) [ 26 ]
Table 5: World model hyperparameters. Both share a frozen, per-domain–finetuned SD-VAE (bf16 activations, fp32 params), 144×144144{\times}144 inputs encoded to 18×18×418{\times}18{\times}4 latents, pre-LN transformer blocks with QK-norm and SwiGLU FFNs, and a flow-matching velocity-MSE objective (EMA-normalized, α=0.99\alpha=0.99 ). Neither uses a target network; only the VAE is frozen. For human evaluation and ablations in Section 4.3 , we instead train the Push-T models at 256×256256{\times}256 inputs to make them easier to evaluate.
MethodIntrinsic rewardKey hyperparameterslearning rates
DIAYN [ 54 ]skill discriminability5050 discrete skills; disc. over T-block pose3×10−43\times 10^{-4}
DADS [ 44 ]skill–dynamics MIz∈[−1,1]8z\in[-1,1]^{8} , 44 mixtures, K=100K{=}100 prior samples3×10−43\times 10^{-4}
ICM [ 50 ]forward-model errorη=50\eta=50 , β=0.3\beta=0.3 , feature dim 2562563×10−43\times 10^{-4}
METRA [ 55 ]temporal-metric skills44 continuous skills, reward scale 1010 , dual reg.1×10−41\times 10^{-4}
Table 6: Hyperparameters for the exploration baselines in Push-T.
Figure 8: Exhaustive Evaluation of World Model Action Sensitivity in Four Rooms: This figure compares the evaluation results of one step action sensitivity of two data scales of a Four Rooms world model trained on expert data via the above “Dozen plots". Each quarter slice in the figure represents a state (grid position and direction of the agent), and within that slice it is split into 3 triangles representing each commanded Four Rooms action: forward, left, right. The color, as denoted in the top right, represents the action corresponding to the actual observed transition by the learned world model. One can see that, in low data regimes, the world model collapses to a near deterministic action policy where the model follows the most likely policy action regardless of action conditioning. However, the same model with more of the same data distribution becomes a near perfect action sensitive world model.
Figure 9: AIM-over-Random Fixed-Start Visitation Ratios: In this figure, we fix the initial state of data collection rollouts to the pose of the green triangle in all figures. Then, we compare state visitations in Four Rooms between rollouts generated by either successive AIM iterations or from a random policy as a heatmap of the clipped log ratio of their empirical state occupancy. More red represents more frequent AIM visitation, and more blue represents more frequent Random visitation. Pseudocounts with ϵ=10−6\epsilon=10^{-6} were added for normalization of the log ratios. One can see AIM provides significantly better medium and long horizon state visitation compared to pure random, which leaves the initial top left room less frequently. Additionally, each iteration shows AIM providing coverage of largely different subsets of the state space, showing the value of AIM iteration and its ability to avoid revisiting previously explored regions.
Figure 10: AIM Four Rooms Extended Results: We provide the full side-by-side metric results from batch online active data collection in Four Rooms. Of note is the similar quality of RMSE as an ordered quality metric compared to the state-focused Sim Fidelity in this simple environment, which is in contrast to conventional wisdom around evaluating generative models. Additionally, we see that measured world model action influence grows proportionately with the evaluated sim realism, a metric computed by enumerating all actions, rolling them out in both the world model and the simulator, and evaluating their state extractive simulation consistency, weighted according to the action influence evaluation measure.
Figure 11: The Need for Trust Regions for Density-Based Intrinsic Motivation in Push-T: In this experiment, we initialize the robot arm in the workspace in a grid and collect 4 transitions at each spot, one for each cardinal direction, then evaluate these transitions with the learned AIM reward model. We then plot the max AIM reward on the above heatmap. We see that the learned AIM intrinsic motivation reward provides significant reward for regions outside of the nominal action workspace, and unbounded policy optimization would likely encourage unfocused state coverage similar to prior work [ 54 , 44 , 55 ] . In contrast, the AWR [ 97 , 72 ] weighted reward focuses in a narrow region around pre-existing behavior, favoring task relevant information before broader exploration.
Figure 12: AIM Behavior Cloning Push-T Extended Results: While End-Effector controllability is significantly improved by AIM, we see all model variants struggle with improving Tee pixel recall, a measure of simulator fidelity. Nearly all exploration data collection policies we evaluate see noticeable improvement in tee recall, but even the top line 300k expert demonstration model struggles with performant tee manipulation. We hypothesize that this is a combination of poor coverage of tee positions in the dataset, as well as a lack of broad pre-training datasets in our experiments to improve extrapolation beyond the existing dataset. Additionally, we note that AIM shows significant improvement in full image pixel recall, outperforming the competing methods, and that pixel RMSE remains a strong metric of model fidelity in deterministic, nearly fully observable MDP settings despite the prevailing focus towards distributional evaluation metrics for generative world models [ 32 , 33 , 35 ] .
Figure 13: AIM Anti-Push-T Extended Results: We see similar results to that of Figure 12 , though dampened due to the significantly more off-policy nature of the “Push Away" Push-T task which results in actions which are significantly more counter factual to those which exist in the pretraining data and thus harder to reproduce. This results in the worst data efficiency scaling from all of the active collection methods.
More in this topic

Action-Conditioned World Models

Action-conditioned world models predict how environments change under actions to support planning, control, and simulation.

See all 345 papers

Explore similar work

CardsList
Apr 13, 2026

AIM: Intent-Aware Unified world action Modeling with Spatial Value Maps

Pretrained video generation models provide strong priors for robot control, but existing unified world action models still struggle to decode reliable actions without substantial robot-specific training. We attribute this limitation to a structural mismatch: while video models capture how scenes evolve, action generation requires explicit reasoning about where to interact and the underlying manipulation intent. We introduce AIM, an intent-aware unified world action model that bridges this gap via an explicit spatial interface. Instead of decoding actions directly from future visual representations, AIM predicts an aligned spatial value map that encodes task-relevant interaction structure, enabling a control-oriented abstraction of future dynamics. Built on a pretrained video generation model, AIM jointly models future observations and value maps within a shared mixture-of-transformers architecture. It employs intent-causal attention to route future information to the action branch exclusively through the value representation. We further propose a self-distillation reinforcement learning stage that freezes the video and value branches and optimizes only the action head using dense rewards derived from projected value-map responses together with sparse task-level signals. To support training and evaluation, we construct a simulation dataset of 30K manipulation trajectories with synchronized multi-view observations, actions, and value-map annotations. Experiments on RoboTwin 2.0 benchmark show that AIM achieves a 94.0% average success rate, significantly outperforming prior unified world action baselines. Notably, the improvement is more pronounced in long-horizon and contact-sensitive manipulation tasks, demonstrating the effectiveness of explicit spatial-intent modeling as a bridge between visual world modeling and robot control.

Liaoyuan Fan, Zetian Xu, Chen Cao +3

INFIFORCE Intelligent Technology Co., Ltd., Hangzhou, China · The University of Hong Kong, Hong Kong SAR, China · Shanghai Jiao Tong University, Shanghai, China

World Action ModelsRobotic ManipulationLong-Horizon Robotic ManipulationRobotic RL
AskPDF
Aug 10, 2026

FACT: Failure-Aware Causal Training for World-Action Models

Recent world-action models (WAMs) show that co-training policies with future prediction can provide physical priors for action generation. Building on the future-prediction ability of video models, many WAMs generate future videos and recover actions with inverse-dynamics models, or use these predicted videos as goal conditions for action generation. In both cases, the world model is trained mostly on successful demonstrations and has little reason to predict the consequences of bad actions. We introduce FACT, a causal World-Action Model that predicts future video and task progress conditioned on the executed action. This action-conditioned interface allows failure rollouts to supervise action consequences, turning bad actions into valid future targets rather than being discarded. Failure-aware training makes the progress predictor aware of both successful and failed action outcomes, which can optionally be used to score sampled action candidates at inference. Extensive experiments on simulation and real-world bimanual manipulation tasks show that FACT outperforms many existing baselines, improves as failure data are incorporated into training, and reduces success-biased future hallucination under bad actions. See more details at https://fact-wam.github.io/

Quanquan Peng, Yutong Liang, Rui Yan +2

University of California San Diego

Causal World ModelsAction-Conditioned World ModelsRobotic ManipulationWorld Action Models
AskPDF
Aug 5, 2026

Overcoming Statistical Bias in Action-Controllable World Models

Action-conditioned world models aim to predict how visual environments evolve under an agent's actions. Yet future frames are often highly predictable from visual inertia and recurring motion patterns alone. This creates a shortcut: models can fit the data by exploiting statistical biases without making their visible dynamics meaningfully depend on the action. As a result, different actions may produce similar futures, while motion may persist even under zero action. The key question is how to reduce reliance on statistical shortcuts from dominating action-conditioned prediction. We argue that action control requires more than injecting action features; it requires enforcing consistency under counterfactual changes to actions and observations. Based on this insight, we introduce CoCo, a Counterfactual Consistency framework to enhance action controllability through two complementary constraints. Multi-step counterfactual consistency constrains reference, inverse-action, and zero-action rollouts, while action-spatial counterfactual consistency enforces consistent predictions under mirrored scenes and transformed actions. Together, they reduce reliance on statistical shortcuts from substituting for action-dependent dynamics. We further introduce Action Response Consistency (ARC) and Drift Energy (DE) to assess action controllability, together with Mini-SSMB for same-state, multi-action counterfactual evaluation. On Mini-SSMB, our full model achieved ARC_inv of 0.412 and ARC_ref of 0.483, while reducing DE by 17.07% relative to the baseline. On VP2 visual planning, it achieves the highest average success rate among SOTA models, at 73.1%. Experiments on BAIR and RoboNet further show that these gains preserve video prediction quality and transfer across model settings.

Yuhong Shi, Zhenhao Chu, Jie Wei +3

Institute of Artificial Intelligence and Robotics, Xi’an Jiaotong University, Xi’an, China · China Mobile (Shanghai) Information and Communication Technology Co., Ltd., Shanghai, China · Zhongguancun Academy, Beijing, China

World ModelsAction-Conditioned World ModelsVideo PredictionCounterfactual Learning
AskPDF