Flow Equivariant World Models: Memory for Partially Observed Dynamic Environments (2601.01075v1)

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.

• The paper introduces Flow Equivariant World Models (FloWM) that leverage Lie group flows to ensure memory consistency in partially observed environments.
• It presents both 2D and 3D instantiations with recurrent and transformer-based architectures that achieve superior long-horizon predictions compared to state-of-the-art methods.
• Extensive experiments on MNIST and simulated block worlds demonstrate significantly lower errors and improved spatial consistency, highlighting practical benefits for embodied AI.

Flow Equivariant World Models: Memory for Partially Observed Dynamic Environments

Introduction and Motivation

The paper presents Flow Equivariant World Models (FloWM), a principled framework for generative world modeling in partially observed dynamic environments, integrating both self-motion and external object motion via a unified mathematical treatment using Lie group flows. Unlike dominant video diffusion and transformer-based generative models, FloWM imposes strong inductive bias through explicit group equivariance, resulting in stable, persistent, and spatially-consistent latent world representations. The study addresses the core challenge in video world modeling of maintaining memory and consistent predictions over long horizons in the presence of partial observability and complex dynamics, a regime largely unmet by existing architectural approaches.

Theoretical Framework: Flow Equivariance

Central to FloWM is the generalization of group equivariance to time-parameterized flows, allowing both agent (internal) and object (external) motions to be modeled as Lie group flows. The model’s hidden state maintains a set of “velocity channels,” each corresponding to a basis in the Lie algebra of the underlying motion group. The recurrent computation—whether simple RNN-style or transformer-based—acts in the egocentric, co-moving reference frame of the agent while faithfully encoding world states in a manner equivariant to both self-motion and external transformations.

The formalism ensures that if the input sequence is transformed according to some flow, the output (both in terms of predicted observations and latent memory state) transforms accordingly. Comprehensive proofs, including a generalized recurrence relation and commutation properties, guarantee this equivariance.

Model Instantiations: 2D and 3D Architectures

Two principal instantiations are provided:

• Simple Recurrent FloWM (2D): Designed for environments with 2D translations. The hidden state is spatial, windowed for partial observability, and augmented with velocity channels representing external object dynamics. Self-motion actions are handled via known group transformations (e.g., translation/roll operations), ensuring that the representation is invariant when the agent returns to a previously visited state.
• Transformer-based FloWM (3D): Built for visually and spatially richer data, the hidden state is a top-down, tokenized map (ViT-based) structured by both agent and object transformation groups (including rotations and translations). The encoder jointly processes current field-of-view tokens and observations, the update writes selectively to the map, and group-structured updates guarantee the desired symmetry. Though not analytically equivariant for arbitrary 3D transformations, the recurrent and gating structure empirically encourages equivariant behavior.

Experimental Validation

FloWM is empirically validated against state-of-the-art video diffusion models and memory-augmented video transformers, specifically History-guided Diffusion Forcing Transformers (DFoT) and a hybrid DFoT-State Space Model (SSM) baseline.

2D MNIST World

• Task: Predict the evolution of MNIST digits moving at constant velocity in a partially observed 2D world, with random agent self-motion.
• Results: FloWM achieves orders-of-magnitude lower MSE and higher SSIM, and maintains error-free rollouts for up to 150 steps—far beyond the training rollout horizon—while all baselines quickly drift, hallucinate scene content, or degrade.
• Ablations confirm the necessity of both velocity channels (external motion equivariance) and self-motion equivariance for robust, long-horizon prediction and sample efficiency.

3D Dynamic Block World

• Task: Simulate long-term object dynamics and predict agent observations from egocentric images in a 3D simulated environment, including texture and viewpoint variability.
• Results: FloWM substantially outperforms SOTA video diffusion methods (both standard and SSM-augmented) in MSE/PSNR/SSIM for both training-length and extrapolated rollouts, exhibiting stable memory of out-of-view states and robust spatial consistency.
• Generalization: The advantage persists with greater visual complexity (textured environments) and in partial observability regimes.

Relation to Prior Work

Previous memory-augmented and consistency-seeking approaches lack a unified or principled treatment of group structure, limiting their effectiveness for handling dynamics outside the field of view or under long-term compositional action sequences. Retrieval-based or geometric memory architectures either scale poorly or are unsuited to dynamic, out-of-view prediction.

Neural Map [Parisotto & Salakhutdinov, 2017] and EgoMap [Beeching et al., 2020] provide early forms of group-structured latent maps, but FloWM generalizes and formalizes this notion, supporting arbitrary Lie groups and endowing the recurrent memory with flow-level equivariance. Unlike prior approximate and loss-conditioned equivariant approaches [Park et al., 2022; Ghaemi et al., 2025], FloWM constructs equivariance by architectural design.

Limitations

FloWM is currently constrained to relatively simple, rigid motions with known parameterizations. Extending to non-rigid object dynamics, rich semantic actions, and continuous symmetries remains future work. Additionally, the current transformer encoder for 3D settings does not guarantee analytic equivariance, which slows convergence. Scaling to open-world, large-scale mixed-reality environments will likely require hierarchical and scalable memory updates.

The approach is also inherently more compute-intensive during training, though inferential cost remains tractable and within an order of magnitude of comparable baselines.

Implications and Future Directions

FloWM demonstrates that symmetry-guided, equivariant memory is critical for persistent, long-horizon world modeling under partial observability. This architectural principle supports rapidly learning robust, stable, and generalizable latent representations, essential for planners and agents in embodied AI.

Key implications include:

• Enhanced memory and consistency: The ability to retain, recall, and update scene state for both observed and unobserved world regions enables robust planning, counterfactual reasoning, and accurate environment simulation.
• Sample and compute efficiency: Embedding group symmetry in the model reduces redundant learning and enables faster convergence.
• Scalability: The mathematical generality of the flow equivariant framework points toward extensions in more complex settings, including full 3D, articulated agents, and rich semantic interaction spaces.

Future research will likely focus on extending flow equivariance to continuous sets, richer action spaces, semantic dynamics, and integrating stochastic latent variables, merging with advances in non-generative predictive control world models.

Conclusion

Flow Equivariant World Models provide a unified, principled solution for memory and representation in partially observed, dynamic environments by leveraging Lie group flows for both self-motion and external dynamics. Empirical results establish clear advantages over dominant video diffusion models and memory-augmented variants, particularly in long-horizon, partially observed tasks. The theoretical and empirical framework presented opens new directions for symmetry-guided, data-efficient, and generalizable embodied world modeling in artificial intelligence.