cs.ROOct 19, 2026
SaveTraining Controllable World Models via Action Influence Maximization
Organizations: Princeton University · Columbia University
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 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 ( ).
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.
| Experiment | BC Recall (H=33, 95% CI) | Preference ELO | BC Failure Recall (H=33, 95% CI) | Preference ELO | Push Away Recall (H=33, 95% CI) | Preference ELO | ||||||
| EE | TEE | Joint | EE | TEE | Joint | EE | TEE | Joint | ||||
| PPO-300K | .275 .013 | .400 .024 | .224 .012 | 843 155 | .554 .019 | .491 .028 | .432 .026 | 1040 185 | .332 .023 | .282 .026 | .198 .018 | 816 128 |
| PPO-3K | .267 .011 | .322 .022 | .203 .011 | 732 155 | .517 .020 | .424 .030 | .378 .028 | 824 171 | .297 .023 | .190 .017 | .147 .015 | 889 132 |
| PPO-3K+AIM500 | .485 .014 | .403 .023 | .344 .018 | 1190 167 | .564 .018 | .454 .031 | .411 .029 | 1049 185 | .375 .016 | .201 .020 | .165 .016 | 1029 134 |
| PPO-3K+Recovery500 | .272 .012 | .378 .026 | .215 .012 | 979 155 | .545 .022 | .455 .030 | .406 .027 | 989 200 | .314 .023 | .216 .020 | .169 .016 | 735 128 |
| PPO-3K+ICM500 | .222 .011 | .393 .025 | .186 .011 | 868 145 | .506 .022 | .441 .029 | .383 .027 | 934 186 | .267 .020 | .242 .023 | .158 .016 | 924 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.
| Expert Policy Name | Steps Trained | Entropy | Eval Return |
|---|---|---|---|
| POMDP PPO | 300M | 0.054 | 0.657 |
| Early Stop PPO | 50M | 0.252 | 0.587 |
| State Oracle PPO | 300M | 0.022 | 0.839 |
Table 2: Four Rooms Policy Statistics. All models are trained with PPO [ 78 ] in Stoix [ 84 ] .
| Softmax BC | Diffusion BC | Softmax IQL+AWR | Diffusion IQL+AWR | |
| Task | ||||
| Environment | FourRooms [ 75 , 76 ] | Push-T [ 83 ] | FourRooms | Push-T |
| Action space | 6 discrete | 6-D cont. | 6 discrete | 6-D cont. |
| Architecture | ||||
| Model dim | 256 | 256 | 256 | 256 |
| Encoder layers | 6 | 12 | 6 | 12 |
Table 3: Policy hyperparameters. All models share a frozen, per-domain–finetuned SD-VAE encoder [ 20 ] (bf16 activations, fp32 params), inputs encoded to latents ( 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.
| Hyperparameter | Value |
|---|---|
| Initialization | sd-vae-ft-mse [ 20 ] |
| Finetuned scope | decoder + post-quant conv (49.5M); encoder frozen |
| Image resolution | |
| Latent grid | ( downsample) |
| Latent scaling | |
| Loss | MSE + 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 | ||
| Domain | FourRooms [ 75 , 76 ] | Push-T [ 83 ] |
| Prediction | single-transition next-frame | sequence-level + AR rollout |
| Architecture | ||
| Model dim | 512 | 512 |
| Layers | 12 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), inputs encoded to latents, pre-LN transformer blocks with QK-norm and SwiGLU FFNs, and a flow-matching velocity-MSE objective (EMA-normalized, ). 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 inputs to make them easier to evaluate.
| Method | Intrinsic reward | Key hyperparameters | learning rates |
|---|---|---|---|
| DIAYN [ 54 ] | skill discriminability | discrete skills; disc. over T-block pose | |
| DADS [ 44 ] | skill–dynamics MI | , mixtures, prior samples | |
| ICM [ 50 ] | forward-model error | , , feature dim | |
| METRA [ 55 ] | temporal-metric skills | continuous skills, reward scale , dual reg. |
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 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