Flow Equivariant World Models: Memory for Partially Observed Dynamic Environments

🍞 Hook: Imagine you’re playing hide-and-seek in a big room with friends running around. You can only see what’s in front of you, and you keep turning. Still, you somehow remember where everyone might be, even when you’re not looking right at them.

🥬 The Concept (World Modeling):

• What it is: World modeling is when an AI predicts how the world will change over time, including how its own actions change what it sees.
• How it works:
  1. Watch several frames of a scene.
  2. Remember important stuff (what and where things are, and how they move).
  3. Use actions (like turn left) to predict the next frames.
• Why it matters: Without a good world model, an AI forgets what left the screen, gets confused when it turns, and makes up wrong details (hallucinates).

🍞 Anchor: A robot vacuum remembering a chair it just passed so it doesn’t bump it when it turns around.

🍞 Hook: You know how peeking through a keyhole shows only part of the room? You still guess where the couch is, even if you can’t see it now.

🥬 The Concept (Partially Observed Dynamic Environments):

• What it is: Situations where you can’t see everything and things keep moving.
• How it works:
  1. Only a small view is visible.
  2. Objects can move while unseen.
  3. The agent also moves, changing what’s visible.
• Why it matters: If the model can’t remember and update hidden parts, it will be inconsistent when you look back.

🍞 Anchor: A self-driving car must track a pedestrian who temporarily goes behind a truck.

🍞 Hook: Think of a notebook where you jot notes so you won’t forget later.

🥬 The Concept (Memory-Augmented Models):

• What it is: Models with an extra memory that stores what was seen before.
• How it works:
  1. Save summaries of past frames into memory.
  2. Update the memory with each new frame.
  3. Read from memory to make better predictions.
• Why it matters: Without memory, once a frame falls out of the attention window, the model forgets it and may hallucinate.

🍞 Anchor: A video model that remembers where a ball rolled, even after the camera turned away.

🍞 Hook: Imagine rules like “turning your head right makes the world slide left” that always hold true, no matter the scene.

🥬 The Concept (Equivariance):

• What it is: A model property where if the input transforms in a known way, the output transforms in the same known way.
• How it works:
  1. Identify a symmetry (like shift or rotate).
  2. Build the model so outputs follow the same rule as inputs.
  3. Share weights across symmetries to learn faster.
• Why it matters: Without equivariance, the model must relearn the same pattern for every position or angle.

🍞 Anchor: A face detector that works the same if the photo is shifted a little.

🍞 Hook: Picture smooth motions like sliding or rotating—no jumps—just steady flows.

🥬 The Concept (Lie Groups, simply):

• What it is: A math way to describe smooth moves (like shifts and rotations) that can be combined.
• How it works:
  1. Describe tiny moves (a “velocity”).
  2. Add them up over time to get a big move (a “flow”).
  3. Combine moves with precise rules.
• Why it matters: Without these rules, the model can’t consistently track motion over time.

🍞 Anchor: Walking three steps forward then turning 90° right always lands you predictably if you repeat it.

The world before: Modern video diffusion transformers make beautiful short videos, but they use sliding attention windows. As videos get long, older frames are dropped, and the model forgets. When the camera turns back, they often hallucinate new scenes. Memory add-ons helped a bit—but usually stored viewpoint-specific pictures instead of a world-centered state, so they still couldn’t update what moved off-screen.

The problem: How can an AI keep a single, stable, world-centered memory that stays consistent while the agent moves and while objects move—even when they’re out of view?

Failed attempts: (1) Longer attention windows—too costly and still forgetful when truncated. (2) Memory banks that retrieve past images—view-dependent and weak at unseen changes. (3) State-space add-ons—compress history but don’t enforce motion-consistent updates.

The gap: A memory that moves exactly with the agent and updates internal object motion using precise symmetry rules was missing.

Real stakes: Robots, AR glasses, drones, and cars constantly face partial views. A reliable, symmetry-aware memory can reduce mistakes, improve safety, and make predictions stable far beyond the training horizon.

🍞 Hook: Imagine wearing a VR headset with a sticky-note map floating in front of you. When you turn, the map turns with you, and the notes about moving objects slide smoothly—even if you’re not looking directly at them.

🥬 The Concept (Flow Equivariance):

• What it is: Treat both self-motion and object motion as smooth time flows and design the model so its memory changes in lockstep with those flows.
• How it works:
  1. Represent motion as a one-parameter flow (like integrating a velocity over time).
  2. Keep a latent world map (memory) that is transformed by known self-actions.
  3. Use “velocity channels” so the memory also flows with external object motions.
• Why it matters: Without it, the model relearns the same motion over and over and loses track when things go out of view.

🍞 Anchor: When you pivot back to where you started, the map looks the same—no surprises—because the memory followed the same turns you did.

The “Aha!” in one sentence: Unify self-motion and external motion as Lie-group flows, then make the memory equivariant to those flows so the world map stays stable and predictive over long horizons.

Three analogies:

1. Moving walkway: If you stand on a walkway and toss a ball, the walkway’s motion (self-motion) and the ball’s throw (external motion) both affect where the ball ends up; this model accounts for both exactly.
2. Sheet protector: Your notes (memory) are in a clear sheet that rotates and slides exactly as you do, while small arrows on the notes show how off-screen things continue to move.
3. Orchestra conductor: Your baton movement (self-motion) shifts the whole score, while each instrument’s melody (external motion) continues on its own line (velocity channel), all in sync.

Before vs. After:

• Before: Sliding windows forget; memory banks memorize views; long rollouts drift and hallucinate.
• After: A co-moving, flow-equivariant memory map keeps consistent state; objects continue moving correctly off-screen; long rollouts remain stable.

Why it works (intuition):

• Motions like shifts/rotations are symmetries with strict combination rules (Lie groups). If you design the memory to transform by these same rules, then turning right and then left brings you truly back. Velocity channels let the memory also keep evolving when objects move while unseen. This prevents drift and hallucinations.

Building blocks:

• 🍞 Hook: Like sorting mail into cubbies by street and then shifting the whole shelf when the city map rotates.
• 🥬 The Concept (Egocentric Latent Map Memory):
  • What it is: A 2D grid of latent tokens centered on the agent that rotates/shift precisely with the agent.
  • How it works: Read/update only the Field-of-View (FoV) region with the new image; transform the whole map by the known action.
  • Why it matters: Keeps a single, world-centric memory that aligns across turns.
  • 🍞 Anchor: Turning your head right makes the map rotate left around you, keeping front-as-front.
• 🍞 Hook: Imagine having different shelves for slow, medium, and fast movers.
• 🥬 The Concept (Velocity Channels):
  • What it is: Parallel layers of memory that each flow at a specific velocity.
  • How it works: At each step, shift each channel by its velocity; combine them when reading.
  • Why it matters: Lets the model carry forward off-screen dynamics faithfully.
  • 🍞 Anchor: A digit that left the view at speed 2 reappears exactly where it should after many steps.
• 🍞 Hook: When you walk forward, the scene slides backward.
• 🥬 The Concept (Self-Motion Equivariance):
  • What it is: A rule to transform memory by the exact inverse of the agent’s action.
  • How it works: Apply a known representation T_a^{-1} (e.g., rotate 90°, shift 1 cell) to the map before the next update.
  • Why it matters: Guarantees closure—coming back to a spot gives the same memory.
  • 🍞 Anchor: Walk a square and end up with a map identical to the start.
• Encoder/Update/Decoder:
  • Encoder reads the current image and FoV map tokens and fuses them.
  • Update writes new info into the FoV positions with a learned gate.
  • Decoder cross-attends to the updated FoV to predict the next frame.

Together, these pieces create a memory that behaves like physics: when inputs transform, the memory transforms accordingly, so predictions remain consistent far beyond the training horizon.

High-level pipeline: Input (image_t, action_t) → [Encode FoV + Map] → [Update FoV tokens] → [Apply internal flows (velocity channels)] → [Apply self-action transform (co-moving)] → [Decode next image] → Repeat.

Step-by-step with the Sandwich pattern on key steps:

1. Observation encoding

• 🍞 Hook: You look through a flashlight cone; you only see a wedge of the room.
• 🥬 The Concept: Encode the current image and the matching FoV region of the map.
  • What: A Vision Transformer (or small CNN in 2D) reads image patches plus the FoV tokens of the latent map.
  • How:
    1. Patchify the image into tokens and add image-position embeddings.
    2. Select the FoV tokens from the egocentric map and add map-position embeddings.
    3. Concatenate and process with a ViT to produce updated FoV features o_t.
  • Why: Without fusing image with the current map slice, the model can’t align new evidence to its memory.
• 🍞 Anchor: Seeing a red block near the right wall updates the right-side map tiles.

1. Map update (write to memory)

• 🍞 Hook: Like writing with a pencil where your light shines, not scribbling the whole notebook.
• 🥬 The Concept: Gated write into FoV map positions.
  • What: For each FoV cell (x,y), blend old token h_t(x,y) with new token o_t(x,y) via a learned gate α.
  • How: h_{t+1}(x,y) = (1−α) h_t(x,y) + α o_t(x,y) for FoV cells; keep others unchanged.
  • Why: Without gating, you’d either overwrite good memory (forget) or never update (stale info).
• 🍞 Anchor: If the block’s color is seen clearly, α goes up and the map trusts the new color.

1. Internal flows for external motion (velocity channels)

• 🍞 Hook: Imagine drawers labeled “moves left 1”, “moves right 1”, etc.; each drawer slides a little every tick.
• 🥬 The Concept: Each velocity channel shifts by its preset motion.
  • What: For each channel ν (e.g., +1 x, 0 y), shift that layer by ψ(ν) to reflect objects’ motion.
  • How: Roll/shift the memory grid for that channel by the exact pixels of ν at each step.
  • Why: Without this, off-screen objects wouldn’t advance in time and would be misplaced later.
• 🍞 Anchor: A digit moving right at speed 2 keeps drifting right in its channel while unseen.

1. Self-motion transform (co-moving frame)

• 🍞 Hook: Turning your head right makes the whole world seem to slide left.
• 🥬 The Concept: Apply the inverse action to the entire map.
  • What: Use the known action representation T_{a_t}^{-1} (e.g., rotate map 90°, shift up by 1) to align the map to the new egocentric view.
  • How: After updating/following internal flows, transform all channels together by the inverse of the agent’s action.
  • Why: Without this, coming back to the same place wouldn’t restore the same memory, breaking consistency.
• 🍞 Anchor: Walk forward one cell; the map shifts backward one cell, keeping the agent at the center.

1. Readout (predict next view)

• 🍞 Hook: To see the next camera frame, you only need the part of the map in your headlamp.
• 🥬 The Concept: Decode the next image from the current FoV.
  • What: A decoder (ViT or CNN) uses the FoV of the updated map to render the next frame.
  • How: Cross-attend from a grid of output image tokens to the FoV map tokens; reconstruct pixels.
  • Why: Without tying the view to the map FoV, the model would drift from its own memory.
• 🍞 Anchor: The predicted camera image matches the part of the world the agent will see after its action.

Two concrete instantiations:

A) Simple Recurrent FloWM (2D MNIST World)

• Representation: h_t has 25 velocity channels for ν ∈ {−2…2}×{−2…2} and 64 hidden features per cell, over a full world-size grid.
• Partial observability: The input image is written only into a centered 32×32 window; the rest of the map remains but flows.
• Update rule: h_{t+1}(ν) = ψ(ν − a_t) ⋅ σ(W ⋆ h_t(ν) + pad(U ⋆ f_t)). Here ψ applies channel-wise shifts by ν − a_t; W and U are small 3×3 convs; σ is ReLU.
• Readout: Crop the same window, max-pool over ν, then decode to the next 32×32 image.
• Why it’s robust: Exact pixel shifts make both self- and object-motion precise; simple ops, strong stability.

B) Transformer-Based FloWM (3D Dynamic Block World)

• Representation: A 2D egocentric map of token embeddings (256-dim) with 5 velocity channels (−1, 0, +1 in x/y, no diagonals) over a 32×32 grid.
• Encoder: ViT over concat(image patches + FoV map tokens) to produce updated FoV tokens.
• Update: Gated blend into FoV positions; others unchanged.
• Internal flows: Shift each velocity channel by its ν.
• Self-motion transform: Rotate 90° or shift one cell using T_{a_t}^{−1} to match turn/forward.
• Decoder: ViT cross-attends FoV map tokens to predict the next 128×128 RGB frame.
• Practical note: The 3D encoder is not analytically equivariant; the recurrence encourages learning approximate action-equivariance during training.

What breaks without each step:

• No encoder fusion: New observations can’t be aligned to the map, causing stale or conflicting memory.
• No gating: Either overwriting good memory (forgetfulness) or under-updating (staleness).
• No velocity channels: Off-screen dynamics aren’t carried forward; long-horizon drift appears.
• No self-motion inverse: Returning to a spot won’t restore the same memory; inconsistency blooms.
• No egocentric map: Memory becomes view-dependent snapshots rather than a world-centric state.

The secret sauce:

• Symmetry first: By baking in the group rules (flows), the model doesn’t have to rediscover them from data.
• Co-moving memory: Transforming the entire map by the inverse action keeps alignment perfect.
• Velocity channels: A compact, data-efficient way to carry forward off-screen dynamics without expensive 3D geometry or huge retrieval banks.
• Result: Stable, long rollouts that generalize far beyond training length.

The tests: Can the model predict future frames accurately when the agent and objects both move, even when objects go out of view? We measure pixel error (MSE), clarity (PSNR), and structural similarity (SSIM) over many future steps.

The competition: Two strong diffusion-based world modeling baselines.

• DFoT: History-guided Diffusion Forcing with a CogVideoX-style transformer backbone (powerful short-horizon quality, sliding window for long videos).
• DFoT-SSM: DFoT plus a state-space memory module (aims for longer memory).

Benchmarks:

1. 2D Dynamic MNIST World (partial observability): Multiple digits move with constant velocities on a larger canvas. The agent’s small window moves via actions; edges wrap. Train on 50 context frames → predict 20; test generalization to 150.
2. 3D Dynamic Block World: In a small room, colored blocks move and bounce off walls. The agent can turn left/right (90°), move forward, or stay. Train on ~50–70 context → predict 70; test up to 210.

Scoreboard with context:

• 2D Dynamic MNIST World (20-step train horizon vs. 150-step long horizon):
  • FloWM: MSE ≈ 0.0005 (20) and 0.0018 (150), SSIM ≈ 0.9900 and 0.9813. This is like getting an A+ and still an A when the exam becomes 7.5× longer than practice.
  • DFoT: MSE ≈ 0.1448 and 0.2111, SSIM ≈ 0.4045 and 0.2434 (drift and hallucinations).
  • DFoT-SSM: MSE ≈ 0.1277 and 0.1688, SSIM ≈ 0.4550 and 0.3146 (fades/forgets).
  • Ablations show: removing velocity channels or self-motion equivariance hurts badly, confirming both are important; no-VC still works somewhat short-term but drifts long-term.
• 3D Dynamic Block World (70-step vs. 210-step):
  • FloWM: MSE ≈ 0.000603 and 0.001539, SSIM ≈ 0.9673 and 0.9525—stable across 3× training horizon.
  • DFoT: MSE ≈ 0.011759 and 0.021684, SSIM ≈ 0.9377 and 0.8885—hallucinations and forgetting.
  • DFoT-SSM: MSE ≈ 0.022616 and 0.022570, SSIM ≈ 0.8877 and 0.8879—limited long-horizon accuracy.
• Textured 3D variant (harder visuals): FloWM still wins strongly (e.g., MSE ~0.000826 at 70), while baselines degrade faster.
• Static 3D variant: With no external motion, the no-velocity version of FloWM slightly edges out full FloWM (velocity channels unnecessary noise) but still far better than diffusion baselines.

Surprising findings:

• Length generalization: FloWM trained for short horizons stays accurate hundreds of steps later—a strong sign the symmetry-structured memory really tracks state rather than memorizing frames.
• Data efficiency: Combining self-motion equivariance with velocity channels speeds learning and reduces training steps to converge.
• Baselines’ behavior: Diffusion models, even with added memory, tend to memorize view-specific appearances; when views change or objects move off-screen, they either hallucinate new objects or fade existing ones.

Takeaway: By aligning memory to actions and flowing internal dynamics via velocity channels, the model predicts the right future even when the evidence was seen long ago and then left the view.

Limitations:

• Rigid/simple motions: Current setups target shifts and right-angle rotations (2D grid actions) and constant-velocity object motion. Articulated bodies or deforming objects aren’t yet covered.
• Deterministic focus: Experiments assume a single correct future given actions; stochastic futures (coin flips, human choices) aren’t modeled yet.
• Discrete velocity set: Velocity channels are discretized; real-world speeds are continuous. Still, partial/discrete equivariance often helps a lot in practice.
• Approximate 3D encoder: The ViT encoder learns action-equivariance rather than being analytically guaranteed, slowing early training.
• Fixed-size egocentric map: Scaling to large, open worlds will need variable-sized, multi-scale maps.

Required resources:

• GPUs suitable for ViT-based training (e.g., L40S/H100 class in the paper’s settings), video datasets with action logs, and standard deep learning tooling.

When NOT to use:

• Highly non-rigid, semantic actions (e.g., “open door” changes topology) without a plan to extend group structure.
• Fully observable static scenes where simple diffusion models suffice.
• Cases with no action information at all (though a learned ego-motion estimator could bridge this).

Open questions:

• How to extend flows to richer action groups (articulations, continuous rotations) and continuous velocity spectra.
• How to build analytically equivariant 3D encoders for faster learning.
• How to incorporate uncertainty and multi-modal futures with stochastic latent variables while preserving flow structure.
• How to scale the map (multi-resolution, hierarchical, or object-centric) without losing equivariance.
• How to combine with control/planning methods (e.g., JEPA/TDMPC2) for end-to-end embodied tasks.

Three-sentence summary: This paper unifies self-motion and object motion as time-parameterized flows and makes a world model’s memory equivariant to those flows. Doing so creates a stable, egocentric latent map with velocity channels that keeps track of off-screen dynamics and stays consistent across turns. The result is far better long-horizon prediction than diffusion baselines, with fewer hallucinations and strong length generalization.

Main achievement: A practical, symmetry-guided recurrent memory (FloWM) that preserves global consistency under partial observability by aligning exactly with action and external motion flows.

Future directions: Extend to richer action groups (articulations and continuous rotations), add stochasticity for multi-modal futures, design analytically equivariant 3D encoders, and scale maps to open-world settings; integrate with planning/control for real tasks.

Why remember this: It shows how building the laws of motion into the memory itself—rather than hoping a big network will rediscover them—yields stable, data-efficient, and much longer-horizon world understanding for embodied AI.