🤖 AI

Flow Equivariant World Models: Memory for Partially Observed Dynamic Environments

This paper introduces Flow Equivariant World Models, a framework that leverages time-parameterized symmetries within a latent memory to enable stable, long-horizon prediction in partially observed dynamic environments by ensuring the memory shifts and transforms equivariantly with self-motion and external object dynamics.

Original authors: Hansen Jin Lillemark, Benhao Huang, Fangneng Zhan, Yilun Du, Thomas Anderson Keller

Published 2026-06-01

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Hansen Jin Lillemark, Benhao Huang, Fangneng Zhan, Yilun Du, Thomas Anderson Keller

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are walking through a busy city square. You are looking at a street performer juggling. Suddenly, you turn your head to look at a friend who just arrived. While your eyes are on your friend, the street performer keeps juggling, and a bus drives by in the background. When you turn back to the square, you expect the performer to still be juggling and the bus to have moved further down the street. You don't expect the performer to vanish or the bus to suddenly teleport back to where it was.

This is exactly the problem that Flow Equivariant World Models (FloWM) tries to solve for artificial intelligence.

The Problem: The "Amnesia" of AI

Current AI systems that try to predict the future (called "World Models") are like people with very short attention spans. They can watch a video for a few seconds and guess what happens next. But if you ask them to predict what happens after they "look away" (which happens when an AI agent moves its camera or turns its body), they get confused.

Because they only see a small slice of the world at any moment, they forget what's happening outside their view. When they look back, they often "hallucinate" (make things up). They might think the street performer disappeared, or they might invent a brand new bus that never existed, just to fill in the gap. They lose track of the "flow" of time and motion.

The Solution: A Moving Map

The authors of this paper propose a new way for AI to remember the world. Instead of just storing a static picture of what it saw last, FloWM builds a structured, moving memory.

Think of it like this:

Old AI: Imagine you are drawing a picture of the city on a piece of paper. Every time you turn your head, you have to erase the old drawing and start a new one on a fresh piece of paper. You lose the context of what was behind you.
FloWM: Imagine you are drawing on a giant, transparent sheet of glass that is floating in front of you. As you walk and turn, the glass moves with you. If a bus drives by on the left, you draw it on the glass. When you turn your head, the bus stays painted on the glass, moving along with the flow of the world. When you turn back, the bus is right there, exactly where physics says it should be.

This "glass sheet" is the Flow Equivariant Memory. It respects the rules of motion (symmetries). It understands that if you move left, the world effectively moves right relative to you, and it updates its internal map accordingly without losing the data.

How It Works (The "Velocity Channels")

The paper introduces a clever trick called "Velocity Channels."

Imagine the AI's memory isn't just one big picture, but a stack of transparent overlays.

One overlay tracks things that are standing still.
Another tracks things moving slowly to the right.
Another tracks things moving fast to the left.

When the AI sees something moving, it knows exactly which "overlay" to update. Even if the AI turns its head (self-motion), the system shifts all these overlays together, keeping the relative positions correct. This allows the AI to predict what is happening behind it or off-screen with high accuracy, because the memory is mathematically designed to flow along with time and movement.

What They Tested

The researchers tested this idea in two main scenarios:

2D Digits: A simple world where numbers (like 1, 2, 3) float around on a screen. The AI had to predict where they would be after moving its "eyes" around.
3D Blocks: A more complex 3D room with colored blocks bouncing off walls. The AI had to predict the blocks' paths even when they went behind walls or out of view.

The Results

The results were clear:

Old AI models quickly lost track. After a while, they started guessing wrong, making up new objects, or forgetting old ones. Their predictions got blurry and chaotic.
FloWM stayed accurate for a very long time. It could predict the movement of objects 150 to 200 steps into the future, even though it was only trained on short sequences. It didn't hallucinate; it kept the "symphony of flows" playing correctly.

Why It Matters

The paper argues that for an AI to truly understand a dynamic world (like a robot navigating a house or a car driving down a street), it needs a memory that moves with the world, not just a memory that stores static snapshots. By organizing memory to respect the natural laws of motion and time, the AI becomes much better at predicting the future, learning faster, and making fewer mistakes.

In short, FloWM gives AI a "moving map" that never loses its place, allowing it to see the whole picture even when it's only looking at a tiny part.

Technical Summary: Flow Equivariant World Models

1. Problem Statement

Embodied agents operate in dynamic environments where they possess only a restricted, first-person field of view. This partial observability creates a significant challenge for existing world models: while the agent observes a fraction of the environment, unobserved regions continue to evolve. Current state-of-the-art methods, including large-scale latent diffusion transformers and recurrent models, often fail to maintain consistency over long horizons in these settings.

Without a memory structure that respects the underlying symmetries of the world, these models suffer from:

Information Loss: As the agent moves, relevant past information falls out of the context window (e.g., sliding windows in transformers) and is discarded.
Hallucination: When the agent turns back to a previously observed region, the model often fails to recall the object's state, instead generating new, hallucinated objects or fading existing ones.
Lack of Structural Priors: Existing models treat observations as static sequences or view-dependent data, ignoring the fact that both self-motion and external object motion obey smooth, time-parameterized symmetries (flows).

2. Methodology: Flow Equivariant World Models (FloWM)

The authors propose FloWM, a framework that integrates flow equivariance into the core of world model memory. The central hypothesis is that predictive representations become more powerful when organized in alignment with the temporal and dynamical structure of the world.

Core Concept: Flow Equivariance

The framework extends static group equivariance to time-parameterized sequence transformations, or "flows."

Definition: A model is flow equivariant if, when the input sequence undergoes a flow (e.g., a shift due to time or motion), the output sequence transforms according to the same flow.
Global Equivariance: In partially observed settings, the model must maintain a latent memory of the entire world state (observed and unobserved). The latent representation must transform equivariantly with respect to the global world state, not just the current observation.
Co-moving Reference Frame: To achieve this, the recurrent memory is transported by the same flow as the input before incorporating new observations. This ensures features remain spatially aligned with moving objects rather than being left behind in a fixed frame.

Architectural Components

The framework is instantiated through a generalized flow equivariant recurrence relation:
$h_{t+1}(\nu) = \psi_1(\nu) \cdot U_\theta [h_t(\nu); E_\theta [f_t; h_t] (\nu)]$
Where:

$h_t(\nu)$ represents the hidden state organized into "velocity channels" ( $\nu$ ), corresponding to different vector fields of motion.
$E_\theta$ is an encoder mapping observations to the latent space.
$U_\theta$ is the update operation.
$\psi_1(\nu)$ represents the flow action over one timestep.

Self-Motion Equivariance:
The model treats agent self-motion as a relative flow. By transforming the hidden state according to the agent's action ( $T_{a_t}$ ), the memory effectively shifts in the agent's co-moving reference frame. This allows the model to build a structured "map" of the environment where new observations are integrated into correct relative positions based on the agent's actions.

Two Instantiations

Simple Recurrent FloWM (2D): Designed for 2D environments (e.g., MNIST World). It uses convolutional operations over hidden states and spatial inputs. It employs a fixed window to read/write to the hidden state while the rest of the state flows around the agent's field of view.
Transformer-Based FloWM (3D): Designed for complex 3D environments (e.g., Block World). It utilizes a Vision Transformer (ViT) encoder and decoder. The latent memory is a set of spatially organized token embeddings acting as a top-down abstract map.
- Note on Equivariance: While exact encoder equivariance is difficult for 3D-to-2D projections without costly depth unprojection, the authors demonstrate that the recurrent flow transformation ( $T^{-1}_{a_t} \psi_1(\nu)$ ) encourages the encoder to learn equivariance during training, achieving high fidelity without explicit geometric constraints.

3. Key Contributions

Formalization of Flow Equivariance: The paper formalizes how time-parameterized symmetries (flows) can be used to construct structured, persistent memory for partially observed dynamic world modeling.
Structured Memory Architecture: Introduction of a recurrent memory with "velocity channels" that shifts and transforms equivariantly with both self-motion and external object motion, keeping information about out-of-view regions aligned over time.
New Benchmarks: The authors introduce new benchmarks for partially observed world modeling:
- 2D MNIST World: Digits moving at random velocities with partial observation.
- 3D Dynamic Block World: Colored blocks bouncing in a 3D room with an agent navigating via discrete actions.
Empirical Validation: Demonstration that FloWM significantly outperforms state-of-the-art diffusion (DFoT, DFoT-SSM) and recurrent (RSSM) baselines in long-horizon prediction and length generalization.

4. Experimental Results

The authors evaluated FloWM against baselines including History-Guided Diffusion Forcing (DFoT), DFoT-SSM, and DreamerV3's RSSM.

Performance Metrics

Long-Horizon Consistency: On the 2D MNIST World, FloWM maintained consistent predictions for 150 timesteps beyond the training horizon (20 timesteps), whereas baselines diverged or hallucinated. On the 3D Block World, FloWM remained consistent for 210 frames, while baselines lost track of objects or hallucinated new ones.
Error Reduction: FloWM achieved significantly lower Mean Squared Error (MSE) and higher PSIM/SSIM scores compared to all baselines. For example, in 3D Block World, FloWM's MSE was ~0.0006 (70 frames) vs. 0.0117 for DFoT.
Learning Efficiency: Models combining Self-Motion Equivariance (SME) and Velocity Channels (VC) required orders of magnitude fewer training steps to converge compared to ablated versions.
Downstream Planning: In a "Find the red block" planning task, FloWM successfully navigated to the target. Baselines often hallucinated the target's location, leading to poor planning performance.
Probe Analysis: A probe network trained to decode block positions from the latent space showed 96% accuracy for FloWM versus <1% for DFoT, confirming that FloWM learns a structured, equivariant representation of object trajectories.

Ablation Studies

Removing Velocity Channels (VC) or Self-Motion Equivariance (SME) caused performance to degrade significantly, often matching or falling below baseline diffusion models, highlighting that both components are critical for handling partial observability and dynamics.

5. Significance and Claims

The paper claims that predictive representations are fundamentally more powerful when organized in line with the temporal and dynamical structure of the world.

Theoretical Groundwork: The work provides a theoretical framework (flow equivariance) that dictates the necessity of structured memory for world models in dynamic, partially observed environments.
Solving Partial Observability: Unlike previous methods that rely on extending context windows or static memory banks, FloWM inherently handles the evolution of unobserved regions by maintaining a memory that moves with the world state.
Generalization: The framework enables models to generalize to significantly longer sequences than seen during training, a critical capability for real-world embodied agents.
Modest Scope: The authors acknowledge that current evaluations are in controlled settings with rigid-body geometric actions. They position this work as a foundation for future extensions to open-world datasets, continuous action spaces, and more complex physical interactions.

In summary, FloWM demonstrates that by explicitly encoding the symmetries of motion into the memory architecture, world models can achieve stable, accurate, and hallucination-free predictions over long horizons in partially observed dynamic environments.