Flow Equivariant World Models: Memory for Partially Observed Dynamic Environments

TL;DR

Flow Equivariant World Models (FloWM) address long-horizon prediction in partially observed, dynamic environments by unifying self-motion and external object motion as time-parameterized flows on Lie groups. The framework enforces flow equivariance through a generalized recurrence with multiple velocity channels and a co-moving reference frame, and it provides two instantiations: Simple Recurrent FloWM and Transformer-Based FloWM. Empirical results on 2D MNIST World and 3D Dynamic Block World show FloWM achieves stable, out-of-view dynamics predictions far beyond training horizons and outperforms diffusion-based baselines, with velocity channels and self-motion equivariance driving data efficiency. Limitations include non-exact 3D analytic equivariance and continuous velocity extensions, with future work aiming to broaden action groups, improve efficiency, and integrate with downstream embodied tasks. Overall, FloWM offers a principled, symmetry-guided approach to building memory-rich, long-horizon embodied agents.

Abstract

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 streams obey smooth, time-parameterized symmetries, which combine through a precisely structured algebra; yet most neural network world models ignore this structure and instead repeatedly re-learn the same transformations from data. In this work, we introduce 'Flow Equivariant World Models', a framework in which both self-motion and external object motion are unified as one-parameter Lie group 'flows'. We leverage this unification to implement group equivariance with respect to these transformations, thereby providing a stable latent world representation over hundreds of timesteps. On both 2D and 3D partially observed video world modeling benchmarks, we demonstrate that Flow Equivariant World Models significantly outperform comparable state-of-the-art diffusion-based and memory-augmented world modeling architectures -- particularly when there are predictable world dynamics outside the agent's current field of view. We show that flow equivariance is particularly beneficial for long rollouts, generalizing far beyond the training horizon. By structuring world model representations with respect to internal and external motion, flow equivariance charts a scalable route to data efficient, symmetry-guided, embodied intelligence. Project link: https://flowequivariantworldmodels.github.io.

Figures (10)

• Figure 1: Partially observable dynamic world modeling. The agent observes dynamics, turns away, then turns back to the original viewpoint. Flow Equivariant World Models (FloWM) can successfully integrate dynamics through time in a stable manner, while existing work hallucinates.
• Figure 2: Comparing World Modeling Frameworks. a) Standard autoregressive video diffusion evicts frames beyond the sliding window. b) Information dependencies between past observations and generated frames cause inconsistency without memory. c) Existing memory solutions are view-dependent, and thus cannot predict dynamic scenes consistently. d) FloWM remembers past observations in the spatial latent memory, and continually updates them via internal dynamics.
• Figure 3: Visualization of the Simple Recurrent FloWM on MNIST World. FloWM Recurrence relation. Velocity channels are plotted as rows, with the 'read-in' and 'read-out' part of the hidden state in blue.
• Figure 4: Transformer-Based FloWM. a) Image observation $f_t$ at time $t$ and $\mathrm{FoV}$ selected map latents $h_t$ are passed through ViT encoder $\mathrm{E}_\theta$. Latent map $h_t$ is fully learned, visualized as a map here for clarity. b) Write to $h_t$ at the $\mathrm{FoV}$ locations, then transform latent map according to known action $a_{t}$ and internal flow $\psi_1(\nu)$, producing $h_{t+1}$. c) Decode using cross attention over $\mathrm{FoV}$ of $h_{t+1}$ with a ViT decoder $\mathrm{D}_\theta$ to predict next image $\hat{f}_{t+1}$.
• Figure 5: Dynamic MNIST World Prediction Rolloutsa) Timesteps 0 to 49 are given as observations. Models are trained to predict up to $t=69$. Note that FloWM does not diverge even at timestep 199, while baselines slowly degrade in image quality or lose track of the digits. b) MSE over different length rollouts show length generalization. c) Learning efficiency of the FloWM.

Theorems & Definitions (2)

• Theorem : The Generalized Flow Equivariant Recurrence Relation of Eqn. \ref{['eqn:general_floweq_app']} is Flow Equivariant
• proof