Training controllable world models
via Action Influence Maximization

Samuel M. Bateman1 · Tenny Yin1 · Chongyi Zheng1 · Lei Huang1,2 · Brian Wang1
Benjamin Eysenbach1 · Jaime Fernández Fisac1 · Dhruv Shah1

1 Princeton University & 2 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 85× reduction in data.

1

Problem

World models lack useful controllability.

In the last couple years, world models (predictive models of the future based on one's actions) have gone from "showing signs of life on domain-specific problems" to "possibly simulates physics flawlessly."

In robotics, this is pretty interesting: a lot of cool work has pointed at the power of training policies in scalable simulation despite the sim2real gap. What if we could eliminate this gap entirely by data-driven simulation of the real world?

But generating plausible video is very different than predicting action consequences reliably. This is sometimes called the controllability of the world model, relating to the influence a policy's actions have on the outcomes in the world model. This work and contemporaneous work by N. Hansen and X. Wang have found that varying the training distribution of a world model changes the predictive errors of the model.

Try it: compare predictions on BC failures

Figure 1 · World model predictions with different training data

Simulator
3K Expert Demos
3K Recovery Demos + 500 AIM
The same recorded commands, shown in simulation and two world-model predictions. Training datasets: 3,000 expert demos versus 3,000 recovery demos plus 500 AIM trajectories. This selected example compares different starting datasets as well as the AIM addition; it is not a matched-prior ablation. The AIM trajectories were reused from PPO-prior collection.

Action influence versus fidelity

Four Rooms navigation
In Four Rooms, higher model-explained action influence accompanies higher simulator fidelity across the tested collection policies. These are the original measured points, separated from the scaling curves. Association here does not establish a universal test-time ceiling.

Since we want data-driven simulators for improving policies, we are particularly interested in controllability, separate from pure prediction error. In a discrete environment (Four Rooms) where this can be measured exhaustively, we find controllability/action influence is a pretty good proxy for fidelity anyways.

This leads our work to answer three key questions: what properties of data influence downstream world model controllability? And can we actively collect for these properties? Can we do it without training a world model in the first place? Our work suggests the answer to all three is yes.

2

Framing

What properties of data influence downstream world model controllability?

The deficiency of expert demonstration and non-active data collection has a pretty simple explanation: expert demonstrations in high-dimensional action spaces (e.g. human teleoperation) are efficient with their actions and tend to solve the same task in similar ways each time.

A downstream learner policy or task has no reason to solve their task with similar states or actions to the expert unless their task is identical, and this results in poor coverage of the learner sample distribution.

In high-dimensional search spaces with non-uniform initialization, trying to cover this deficiency through random behavior runs into the curse of dimensionality: interesting environments allow for a lot of behavior, most of which is useless, and in robotics, randomly stumbling on an interesting part of the environment is pretty unlikely.

Figure 2 · Training data and learner behavior need not match

Three panels compare world-model training trajectories, learner-policy samples, and behavior-focused AIM exploration, with train, test, and explore regions.
World-model training data, learner-policy samples, and behavior-focused AIM exploration, as illustrated in the supplied teaser. Full-resolution figure

Figure 3 · Different data reveal different behaviors

PushT collection behaviors: expert, noisy expert, recovery, random walk, random, and AIM. The animation illustrates the different motion coverage of these collection policies.

Some form of action collection is critical, but what determines utility? We show in our work that a learned model is bound in what action consequences it can explain by the mutual information of its dataset. If the model is good at extracting information from a dataset, then we can ignore it entirely and optimize the dataset directly, independent of any specific model.

3

Solution

Optimize the data, not the world model

With this insight, that dataset mutual information relates to the best controllability a model can produce, we do a bit of analysis to derive a reward function to directly optimize this quantity in an ideal contextual bandits setting:

rAIMk = log qk(a | s,s′) − log βk(a | s) + log ρeval(s) − log dk(s)

Excitingly, if our practical implementation is similar enough to the ideal contextual bandits problem this is derived from, optimizing this reward has higher mutual information than the previously collected data, and thus a higher upper bound on world model explainable action information. Full mathematical details are in the paper.

This reward popping out is pretty neat: if ρeval is uniform, then this looks like a classic mutual information reward and state coverage curiosity reward.

We estimate these learned reward functions with any prior collected data with flow matching likelihood estimation and self distillation. These allow us to train RL policies to focus collection of a batch of data on the highest information utility samples, which we can add back into our dataset for the next round of reward estimation and collection.

Walk through one AIM collection round

Click a line or a cycle station, or step through the procedure. The reward-fitting call opens Algorithm 2 and returns to the collector.

Summarize data
distribution Deploy on
robot

Repeat with the enlarged replay when more rounds remain.

    Symbolic steps with recorded AIM environment footage from different collections; no live fitting or collection. Algorithm 2 standardizes the reward terms, distinct from the contextual formula above. The final replay is returned for downstream world-model training.
    4

    Results

    The Promising Power of Small Data

    We test different data collection controllability with two downstream learner policies: a Push-T behavior cloning (BC) policy, and a Push-Away policy that moves the Tee away from the goal. Starting with roughly 3,000 expert trajectories, we give collectors five 100-episode rounds in Push-T, and compare AIM with expert, random, recovery, and intrinsic exploration-based collection. On the original short-horizon PushT BC test, 500 AIM trajectories match the 300,000-expert model on end-effector recall: roughly 85× fewer training episodes in total. On push-away, AIM approaches the performance of the 19k-expert. Both show significant data scaling efficiency improvements over alternate data collection efforts in high-dimensional manipulation:

    Figure 5 · Improving controllability via active data collection

    Original evaluation · PushT H10
    Figure 5: Behavior Cloning Policy and Push-Away Policy end-effector recall acquisition curves.
    Push-T BC and Push-Away. Horizontal references use 3k and 300k expert data. These original H10 curves are distinct from the H33 archive above; confidence bands are omitted here. Full Figure 5

    For a more indepth analysis and human A/B preference testing, we train larger, higher resolution models with longer horizons (H=33). We also introduced a BC failure policy, that alternates between BC and push-away to simulate a moderately weak policy somewhere between the two extremes above. Human preference consistently prefers the AIM data mixtures, which is consistent with the controllability and fidelity recall metrics as well. See our paper for full experimental details!

    Behavioral fidelity · Aggregate human preference

    PushT H33 · Table 1
    0Arena score1800
    Human comparisons judge how faithfully predictions reproduce the simulator's robot behavior. Higher Arena scores indicate stronger preference within the selected distribution; whiskers show the reported Table 1 intervals. Scores are relative, not ratio-valued or comparable across distributions. These aggregate results are separate from pixel recall and your sampled notebook. Source: Table 1

    Interactive Experimental Sandbox Try different data mixtures and compare policy rollout predictions.

    Example experiments:

    Choose a question to set up a comparison, or pick your own models below.

    GTRecorded Simulation
    Simulator

    The actual outcome of these actions. Both models start from the same recorded history.

    AModel A
    Prediction

    The model's predicted consequence

    BModel B
    Prediction

    Same actions. Different training data.

    01 / 33
    Action · mm +x+y Unavailable

    Robot recall is the percentage of ground-truth robot pixels reproduced within the evaluator's RGB-error threshold. Higher is better.

    Learner Policy Distribution:

    Your experiments so far

    View your experiments
    How this experiment works

    Run Rollouts samples precomputed results, not live training or inference. The two models receive the same recorded actions and start from the same ground-truth history. We sample a shared evaluation window and a recorded seed-1 prediction for each model. Playback loops; only a completed experiment adds a result to your notebook.

    Pixel recall measures the percentage of ground-truth robot or object pixels reproduced within a mean absolute RGB error of 0.098. Both takes the lower robot/object recall at each frame where both are valid. The large numbers show the current frame; the curves show the window; your notebook averages valid frames. This acts as a proxy for task relevant controllability and fidelity.

    Your experiments

    Robot recall

    Your sample

    Behavior cloning · mean recall over your sampled windows

    Run your first experiment to start comparing models.

    These summaries reflect your chosen models and sampled cases. Different sample sizes and cases can change their order.

    Experiment log

    0 experiments

    Each completed experiment is saved here, automatically.
    Your notebook stays in this browser.

    Citation

    @inproceedings{bateman2026aim,
  title = {Training Controllable World Models via Action Influence Maximization},
  author = {Bateman, Samuel M. and Yin, Tenny and Zheng, Chongyi and Huang, Lei
            and Wang, Brian and Eysenbach, Benjamin and Fernández Fisac, Jaime and Shah, Dhruv},
  booktitle = {Conference on Robot Learning},
  year = {2026}
}