Flow Equivariant World Models: Memory for Partially Observed Dynamic Environments

AI summary

Overview

Research area: Machine learning / generative world modeling, specifically partially observed dynamic environment modeling, video world models, group equivariance, and recurrent latent memory.

Technical level: Advanced. The paper builds on Lie group theory, flow equivariance, diffusion transformers, and vision transformers; the framework is written in mathematical notation, though the core intuition (a memory that moves with the agent) is describable in plain language.

Scope: The paper introduces Flow Equivariant World Modeling and the FloWM architecture, evaluates it against diffusion and recurrent world model baselines on two newly proposed partially observed video benchmarks (2D MNIST World and 3D Dynamic Block World), and demonstrates gains in long-horizon consistency, length generalization, planning, and learned equivariant structure.

What This Paper Is About

Existing video world models represent the world only through what is currently visible, so when an agent turns away and comes back, the model has lost the unobserved regions and often hallucinates new content in their place. The authors argue that both an agent's self-motion and the motion of external objects are smooth, time-parameterized symmetries ("flows"), and that a world model's latent memory should shift and transform according to those same flows. Their goal is a world model whose memory behaves like a structured, continuously updated map of the environment, rather than a view-dependent sliding window of past frames.

Key Contributions

1. Flow Equivariant World Modeling framework. The authors generalize static group equivariance to time-parameterized flows and extend it to partially observed settings, where equivariance is defined with respect to an unobserved world state sequence and an observation function, rather than directly on observations.

2. A generalized flow equivariant recurrence relation. They abstract the recurrence into an arbitrary encoder E_θ[f_t; h_t] and update operation U_θ[h_t; o_t], flowing each velocity channel by ψ_1(ν), and state conditions (Equations 5 and 6, plus a "trivial lift" requirement) under which flow equivariance is retained. A formal proof is given in Appendix Section A.

3. Self-motion equivariance without extra velocity channels. By using the known action a_t and its latent group representation T_{a_t}, the hidden state is transformed into the co-moving reference frame of the agent. The authors note that a byproduct is a memory spatially structured in a manner homomorphic to the world, which they connect to theories of prefrontal cortex memory (Whittington et al., 2025).

4. Two concrete instantiations and two new benchmarks. A Simple Recurrent FloWM (convolutional, with velocity channels and self-motion equivariance) for 2D MNIST World, and a Transformer-Based FloWM with a ViT encoder and ViT decoder for 3D Dynamic Block World built in the Miniworld environment. The paper reports that code is available on the project page.

5. Empirical comparison against three baseline families. History-guided Diffusion Forcing (DFoT) with a CogVideoX-style transformer backbone, DFoT-SSM (diffusion forcing transformer plus a blockwise scan state space model), and an RSSM based on Dreamer V3 for the 3D dataset.

Main Findings

• Long-horizon consistency on MNIST World: Trained with 50 observation frames and 20 prediction frames, FloWM predictions remain consistent with ground truth for 150 timesteps past the observation window, well beyond its training prediction horizon of 20 timesteps. The ablation FloWM (no VC, no SME) fails, and FloWM (no VC) diverges over time.

• Baselines degrade or hallucinate on MNIST World: DFoT's predictions quickly diverge from ground truth and generate plausible digit-like artifacts; the paper reports that such hallucinations make DFoT's MSE error worse than a simple all-black baseline. DFoT-SSM's digits slowly fade to black as it loses track of digits through decaying SSM memory.

• Learning efficiency from structured priors: Models combining self-motion equivariance (SME) and velocity channels (VC) require orders of magnitude fewer training steps to converge compared with those without these priors.

• Long-horizon consistency on 3D Dynamic Block World: Given 70 frames of context, FloWM predictions remain consistent for as many as 210 frames of future prediction, while baselines diverge. Metrics are reported on rollouts of 70 and 210 prediction frames (the numeric MSE, PSNR and SSIM values themselves are not included in the provided content).

• Baseline failure modes on Block World: DFoT and DFoT-SSM frequently hallucinate new objects and forget old ones; the RSSM degrades to a blurry average of many overlapping blocks, even in early prediction frames.

• Robustness to visual complexity: The same performance gap holds on more visually complex variants of Block World, where textures and shapes are randomly assigned per video rather than just color. The authors conclude the bottleneck of existing world models is less about visual complexity and more about predicting consistent environment dynamics.

• Downstream planning: In a "Find the red block" task with reward maximizing red pixels (max reward at distance 2, one red block per episode), using exhaustive search in an MPC framework with search horizon 3 and rollout length 70 after 70 context frames, over 8 episodes and requiring no additional trained parameters, FloWM quickly finds its way to the red block while baselines often hallucinate a red block and follow that action.

• Latent structured memory is decodable: A small probe network trained to decode block locations from internal states shows FloWM decoding block position with 96% accuracy, while DFoT and DFoT-SSM probes decode less than 1% of timesteps correctly. Visualizations for two of six blocks (blue and green) show FloWM tracking ground-truth linear trajectories while DFoT decodes a jagged inconsistent path.

• Equivariance is learned, not enforced: Before training, FloWM's approximate equivariance error is 6.96 in L2 distance, meaning predictions are off by roughly 5 units in both spatial directions (square root of 5² + 5²). After training, the error is reduced to 0.22, an error of only 0.16 units in each dimension, compared to 2.36 for DFoT.

• Partial observability, dynamics, and self-motion individually are not enough: Additional results in Section F show DFoT can sometimes handle partial observability, object dynamics, and self-motion individually, but not in any combination.

Methodology in Plain English

The central idea is that a memory should move with the things it remembers. If an object slides to the right, the part of the internal memory representing that object should slide right by the same amount; if the agent itself turns, the whole memory should rotate or shift accordingly. This keeps stored information aligned with the world instead of being stranded in a fixed frame.

For the 2D experiments, the model is a convolutional recurrent network with a hidden state organized into stacked "velocity channels," each flowing according to its own vector field. Only a window of that hidden state smaller than the world is written to and read from at each step, mimicking a limited field of view; the rest keeps flowing around. The agent's action is folded in as a flow with a negative sign, so internal and action-induced flows combine into a single transformation. To predict an image, the windowed memory is max-pooled over velocity channels and passed through a decoder.

For the 3D experiments, the memory is a spatially organized set of ViT tokens forming an abstract top-down map. A ViT encoder takes the current image patches together with the memory tokens inside the agent's triangular wedge field of view, produces updated tokens, and those are written back into the correct positions of the map. Before each step, the map is transformed by the agent's action (2D translation and 90-degree rotation) and by the internal flow. Prediction uses a ViT decoder that cross-attends over the new field-of-view tokens.

One important design choice: the authors do not enforce exact equivariance of the ViT encoder by design, since mapping from a 3D first-person view to a top-down map would require costly explicit depth unprojection. They instead treat the encoder output as if it were doing the equivariant lift and let the inter-timestep transformation push the encoder to learn equivariance, citing prior work (Keller and Welling, 2021; Keurti et al., 2023). The probe-based equivariance error measurement is how they verify this empirically.

Why This Matters

Impact on research. The paper reframes memory in world models as a geometric object that should transform with the symmetries of the environment, rather than an attention window or an unstructured latent. It unifies and formalizes earlier structured-memory models such as Neural Map (Parisotto and Salakhutdinov, 2017), a variant of which can be seen as a special case of FloWM without velocity channels, and EgoMap (Beeching et al., 2020). It also argues the field's main weakness is dynamic consistency rather than visual fidelity. The broader claim is that predictive representations become more powerful when organized in line with the temporal and dynamical structure of the world they model.

Real-world applications (as motivated or implied by the paper):

• Robotics and embodied navigation, where an agent must keep track of objects, people, or hazards that leave its field of view and reappear later.
• Autonomous driving, where vehicles and obstacles behind or beside the ego vehicle continue to move and matter within seconds.
• Multi-agent coordination problems such as pack hunting, which the introduction uses as the motivating natural example of estimating a target's location and velocity while predicting other agents.
• Simulated training environments and model-based planning, where rollout consistency over long horizons determines whether a planner can be trusted.

Industry relevance. Long-horizon consistency is a direct bottleneck for using generative video world models as simulators for planning and reinforcement learning. The reported planning experiment, where baselines hallucinate a target and follow it, illustrates the practical cost of memory that is not structured. The paper also notes that existing memory solutions, such as memory banks indexed by camera position (WORLDMEM) and 3D voxel maps retrieved to condition diffusion, have limitations: they cannot handle dynamic environments, rely on self-attention to integrate self-motion, use max pooling instead of a recurrence relation, or are prohibitively expensive for large hidden state sizes.

Future Directions

• Exact versus learned equivariance. The ViT encoder's equivariance is only approximately correct (error reduced from 6.96 to 0.22 after training). Whether encoders can achieve exact equivariance without costly depth unprojection, or whether learned equivariance scales to harder visual domains, remains open.

• Richer action and world structure. The framework is instantiated for 2D translation, 90-degree rotation, and external object translation. Extending to more sophisticated action spaces such as 6-DoF motion, non-rigid or articulated objects, and deformable scenes is not demonstrated in this work.

• Scaling the structured memory map. The Transformer-Based FloWM maintains a top-down tokened map; the related-work discussion notes that 3D voxel approaches become prohibitively expensive for large hidden state sizes. How FloWM's map scales in resolution, memory footprint, and compute for large 3D environments is not reported.

• Realistic and real-world data. The paper reports that the same gap holds on more visually complex variants of Block World, but the benchmarks remain synthetic (MNIST digits, Miniworld rooms). Whether structured flow equivariant memory transfers to photorealistic video and real sensor data is not established here.

Target Audience

This paper is most useful to graduate researchers and practitioners working on video generation, model-based reinforcement learning, and world models, especially those focused on long-horizon prediction and partial observability. It also suits researchers interested in geometric deep learning and equivariance who want to see group-theoretic structure applied to sequential latent memory. Readers should be comfortable with group theory notation, latent variable models, and diffusion transformers; those without that background can still follow the conceptual framing and the empirical comparisons against DFoT, DFoT-SSM, and Dreamer V3-based RSSM baselines.

Authors’ 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 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/