A Subtraction Trick Lets Robot World Models Ask 'What If'

A pre-print paper posted this week describes a method that lets world models, the systems robots use to predict what will happen next, answer a question they previously could not: what would have happened if the robot had acted differently in that same moment. The paper, "Counterfactual World Models: Natura Non Facit Saltus," comes from Randall Balestriero of Brown University's AMI Labs, researcher Kerda Varaku, and Yann LeCun of NYU's AMI Labs, who won the Turing Award for his work on deep learning.

World models predict the consequences of an action, then train on the gap between prediction and outcome. Joint Embedding Predictive Architecture models, which LeCun and colleagues introduced in 2023, do this by taking an action as input and forecasting the resulting state. The forecast is an average over many possible outcomes, so it can't say what would have happened under a hidden disturbance, like a current pulling a boat off course, that the model never observed directly.

Balestriero, in an email to The Register, described the fix as a single subtraction instead of a new model. He used the paper's own example: a boat steers left and ends up two meters further downstream than the world model predicted. Those two meters are the signature of a current the model never saw. To predict where the boat would have ended up steering right instead, the method takes the world model's forecast for steering right and adds the same two meters, the gap between what happened and what the model expected for the action actually taken.

"Physics enters to justify this rule, not as an equation the model fits," Balestriero said. Over a short time step, the effect of an action and the effect of a disturbance each grow in proportion to the duration, while their interaction grows only with the square of that duration. Over short intervals the two effects simply add. That is the content of Gottfried Leibniz's axiom natura non facit saltus, nature makes no leaps, which gives the paper its title. The authors prove that the share of the disturbance's effect this rule gets wrong shrinks with the square of the time step, and that no method working from the same ordinary observational data can do better in the worst case.

The team's first attempt at the problem added a learned component meant to infer the hidden disturbance directly. That approach proved fragile. "Learning the disturbance from ordinary data is not just hard, it is impossible in general," Balestriero said. "The hidden noise can be relabeled differently for each action without changing a single prediction, so the data cannot tell you which labeling is the right one." The missing piece has to come from outside the data, he said, from the physical fact that forces add and their interaction builds up slowly over time. Accepting that makes a separate disturbance model unnecessary.

The method works on any world model that predicts the average next state for a given action; it doesn't require a JEPA architecture specifically. Balestriero said JEPA models are especially well suited to the adaptation. In the team's experiments, they used off-the-shelf encoders, DINOv2 and MAE, without fine-tuning them, and trained only a small predictor on top.

Balestriero pointed to three uses for counterfactual world models. The first is letting a robot retry a failed action after an unforeseen event, such as an unexpected load or a weak motor. In simulated tasks, choosing the retry action this way reached the goal in 35 to 69 percent of events, against 14 to 23 percent for a planner that knows the true simulator but ignores the observed outcome. Two more attempts on the same event raised success to 51 to 97 percent. The other two uses he described are adapting to a persistent fault, like a miscalibrated motor, and helping engineers trace why a given action failed.

"World models are trained to predict, and they are getting good at it," Balestriero said. "But a robot, a car or a drone acting in the real world is constantly surprised by things it could not see coming. After a surprise, the useful question is no longer what was expected, but what would have happened had it acted differently in that same situation." Control engineers have applied corrections like this for decades to hand-built models, he noted. What the paper adds is a proof of why the same correction holds for learned world models that read the world through images, and where the method breaks down.

The paper was shared with The Register ahead of its posting to arXiv, expected this week. It has not yet been peer reviewed, and the reported success rates come from simulated tasks, not physical robots, so the next test of the claim is whether the subtraction rule holds up when the disturbance is measured by a camera instead of a simulator's ground truth.