Machine learning researchers build a better world model

Machine learning researchers have created a method that enables world models to make more accurate predictions about how objects in the physical world will behave. World models are utilized to instruct physical AI systems, such as robots, on navigating the real world by predicting the outcomes of actions. However, these models rely on observed data, which limits their ability to make predictions about alternative scenarios without prior training for every possible outcome.

The researchers, Randall Balestriero from Brown University's AMI Labs, Kerda Varaku, and Turing Award winner Yann LeCun from NYU's AMI Labs, propose a method called Counterfactual World Models in a pre-print paper titled "Counterfactual World Models: Natura Non Facit Saltus." The Latin phrase refers to Gottfried Leibniz's axiom that "nature makes no leaps." Despite this axiom's validity until quantum mechanics, the researchers have devised a technique for zero-shot physics inference using counterfactual world models.

Counterfactual world models leverage the difference between the model's prediction and the actual world to identify errors and enhance future predictions, enabling debugging, reasoning, and inference. This approach can be applied to any world model that predicts the average next state for any given action, although JEPA (Joint Embedding Predictive Architecture) models developed by LeCun and colleagues are particularly well-suited for adaptation.

The counterfactual world model works by inferring possible scenarios from observed events, answering questions like "What would have happened if I acted differently?" in three steps: learning hidden factors from observed data, changing the action, and predicting again with those hidden factors kept constant. For instance, in a scenario where a boat crosses a river with an unknown current, the JEPA world model predicts the boat's position, and the counterfactual model can estimate the boat's position if it had steered in a different direction by adding the same displacement caused by the actual current.

Counterfactual world models have three primary applications: allowing robots to retry an action after an unexpected event, adapting to persistent faults such as miscalibrated motors, and helping users understand failures. The researchers found that using counterfactual world models enables robots to reach their goals in 35 to 69 percent of events, compared to 14 to 23 percent for planners that ignore the observed outcome. With two additional attempts, success rates increased to 51 to 97 percent.