Machine learning researchers have developed a technique that allows world models, used to train physical AI systems like robots how to navigate the real world, to make better predictions about how objects will behave by inferring counterfactual scenarios.
The work, described in a paper titled 'Counterfactual World Models: Natura Non Facit Saltus', comes from Randall Balestriero (Brown University, AMI Labs), Kerda Varaku, and Turing Award winner Yann LeCun (NYU, AMI Labs). The Latin in the title refers to Gottfried Leibniz's axiom 'nature makes no leaps,' which the authors say explains how they devised a method for zero-shot physics inference.
World models predict the consequences of actions, but they do so based on observed data. A Joint Embedding Predictive Architecture (JEPA) model, developed by LeCun and colleagues in 2023, can take an action as input and predict what the world would be like with that action. However, a JEPA world model produces an average of outcomes and cannot anticipate counterfactual scenarios, events that did not happen.
The new technique is not specific to JEPAs but can be applied to any world model that predicts the average next state under any action. Balestriero said the counterfactual world model experiments used off-the-shelf encoders (DINOv2 and MAE) without any fine-tuning and only a small predictor trained on top. Creating a counterfactual world model from that starting point would require only minor additions.
Counterfactual world models predict alternative scenarios more effectively by inferring possibilities from what did occur. 'In causal inference, answering 'What would have happened had I acted differently?' takes three steps,' Balestriero explained. 'You use what you observed to learn something about those hidden factors, you change the action, and you predict again with the hidden factors kept exactly as they were. Our counterfactual world model does all three steps with one subtraction.'
Balestriero cited an example from the paper: a boat crossing a river with a hidden current. 'The world model predicted where the boat would end up after steering left, and it actually ended up two meters further downstream,' he said. 'Those two meters are the signature of the current that occurred. To answer 'where would the boat be had I steered right?', we take the world model's prediction for steering right and add the same two meters.'
'We prove that the fraction of the disturbance's effect that this rule gets wrong shrinks with the square of the time step, and that no method using the same ordinary data can do better than this rate in the worst case,' Balestriero said.
Counterfactual world models have three likely uses, according to Balestriero. The first is having a robot retry an action after a surprise event. 'When an attempt fails because of something the robot could not see, such as an unexpected load, a current or a weak motor, the counterfactual world model lets it estimate which action would have worked in that same situation,' he explained. 'In our simulated tasks, choosing the action this way reaches 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. With two more attempts on the same event, success rises to 51 to 97 percent.'
He also sees these models being useful for adapting to persistent fault, like a miscalibrated motor, and helping people understand failures. 'World models are trained to predict, and they are getting good at it,' said Balestriero. '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.'
The researchers initially tried adding a learned component to the world model to infer hidden disturbances, but that approach was too fragile. 'Learning the disturbance from ordinary data is not just hard, it is impossible in general: 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,' Balestriero said. 'For physical systems, it comes from the fact that forces add and that their interaction builds up slowly in time.'
Balestriero said control engineers have been implementing corrections of this sort for decades, but that what was missing was why the same correction is justified for learned world models, what can and cannot be recovered without paired data, and where it breaks. The paper, provided to The Register prior to its release, is expected to be released on ArXiv this week.