Flow Equivariant World Models: Structured Memory for Dynamic Environments

Embodied systems experience the world as 'a symphony of flows': a combination of many continuous streams of sensory input coupled to self-motion, interwoven with the dynamics of external objects. These sensory streams and the underlying dynamics of the world obey smooth, time-parameterized symmetries which existing world models ignore. Without a memory that respects this structure, partial observability presents a major obstacle to existing methods: each observation reveals only a fraction of the world, while unobserved regions continue to evolve. In this work, we introduce Flow Equivariant World Modeling, a framework that leverages time-parameterized symmetries within a latent memory for stable and accurate dynamics prediction over long horizons. The latent memory shifts and transforms equivariantly with self-motion and inferred external object motion, keeping information about out-of-view regions aligned as time progresses. We demonstrate the advantage of this framework over state-of-the-art diffusion, memory-augmented, and recurrent world model architectures on 2D and 3D partially observed video world modeling benchmarks. More broadly, our results suggest that predictive representations become more powerful when they are organized in line with the temporal and dynamical structure of the world they model. Project page: https://flowequivariantworldmodels.github.io/

Imagine walking across the street, and noticing the cars coming from each of the four directions. For humans, it is natural to simulate the movement of the cars through time to avoid being hit, even when looking in a different direction. But for machine learning models, it is not so easy. In this work, we ask: how do intelligent embodied agents predict the dynamics of other objects in their environment, especially when they themselves are also moving? How does the agent remember what it saw previously, and simulate or imagine what will happen in the future, even if it the outcome is not playing out in their field of view? To study this problem, we propose a new type of machine learning model that maintains a spatial memory of the environment it is in. This spatial memory is then used to make predictions of what will happen in the future. The key novelty is that the spatial memory updates through time using a mathematical framework called Flow Equivariance that mirrors the structure of the world itself. Just as a human might simulate the future position of a moving object after they have observed and remembered how fast it is moving, the machine learning model can use Flow Equivariance to achieve the same goal. The movement of the agent itself is also built into the framework, allowing for seamless prediction of both external and internal dynamics. We demonstrate that machine learning models built within this framework can accurately predict the dynamics of the world around them, while existing methods that focus on static memory fail. For example, in an environment with moving blocks, after making a few seconds of observations, the models we train can accurately predict how that environment evolves long into the future. We see our work as an exciting step forward in building machine models that can predict and understand the 3D world's complexity as we humans do.