Improving World Models: A Neuroscience-Inspired Approach to Latent Planning
Humans instinctively map out the physical consequences of their actions before taking them, seamlessly predicting that a dropped glass will shatter or a jumped gap might be too wide. In new research, CDS PhD student Ying Wang and her colleagues applied a concept from human neuroscience to help AI systems plan their physical movements with similar intuitive efficiency.
The paper, “Temporal Straightening for Latent Planning,” was co-authored by Wang, CDS founding director Yann LeCun, CDS Assistant Professor Mengye Ren, NYU postdoctoral researcher Oumayma Bounou, and colleagues, and has been accepted at ICML 2026.
The authors tackled a fundamental issue in how world models perform planning. World models are AI systems that learn the dynamics of their environment to simulate the outcomes of actions. Currently, many models use powerful pre-trained visual encoders to compress high-dimensional observations, like images of an environment, into more compact latent representations. While these encoders excel at capturing static visual semantics, they struggle to organize this data in a way that helps an agent plot a physical course from a starting point to a goal. In the latent space created by these encoders, a feasible trajectory often looks like a zigzag, meaning the latent distance between two points fails to reflect the actual geodesic distance the agent must travel.
To solve this, the team looked to human biology, drawing inspiration from a 2019 paper on perceptual straightening co-authored by CDS-associated Professor of Neural Science, Mathematics, Data Science, and Psychology Eero P. Simoncelli. That research demonstrated that human visual systems inherently transform complex, highly curved natural video sequences into straighter internal representations.
Wang and her collaborators incorporated this mechanism into their world models by applying a curvature regularizer during the training of the model’s encoder and predictor. They reduced the dimensionality of the features dramatically — compressing 384 dimensions down to just eight — while still preserving the structural information from the original visual input necessary for successful planning.
“By making feasible trajectories straighter in latent space, we make prediction and planning easier,” Wang said.
The researchers observed that reducing the curvature of these latent trajectories completely transformed the optimization landscape. It makes the landscape closer to convex, making gradient-based planners less likely to get trapped in poor local minima. This straightening effect made gradient descent significantly more stable and yielded higher success rates across a suite of goal-reaching tasks, bypassing the need for computationally heavy search-based planning methods.
They also found that the Joint-Embedding Predictive Architecture (JEPA) used in the model naturally encouraged this simplified geometry even without the explicit curvature penalty.
“Because we are training the predictor and encoder together, the pressure to predict implicitly forces the encoder to straighten,” Wang said.
By Stephen Thomas
