Flow Equivariant World Models: Structured Memory for Dynamic Environments

Comparison between different world modeling frameworks: a) Standard diffusion forcing uses a fixed length sliding window to generate video autoregressively. Frames must be evicted if they exceed the window. b) When there are information dependencies between past observations and the generated frame, without a memory mechanism, DFoT is not able to generate consistently. c) Existing memory solutions are view-dependent, and cannot handle dynamic scenes, still resulting in inconsistent generation. d) In FloWM, past frames are remembered in the spatial latent memory and continually updated through FloWM's internal dynamics, resulting in consistent generation.

FloWM Recurrence relation in 3D: a) Information passes from the image observations to the hidden state through a ViT encoder. b) The new updates are combined with the existing hidden state, and the action and internal flows roll the hidden state to the next timestep. c) The updated hidden state is used to predict the next timestep observation.

Below we provide full rollouts on the Dynamic, Textured, and Static splits of the 3D Blockworld Dataset

We train simple probe models on the activation spaces of FloWM and baselines on Dynamic Block World to see whether the representation readiliy contains the position of each block at each timestep. a) FloWM can make practically perfect predictions, even while the agent is moving around and the blocks are moving. This demonstrates the accuracy of the learned latent space. b) The DFoT baseline makes sporadic predictions of each of the blocks, demonstrating its unprincipled representation space.

We find that the probe can accurately predict the correct block position ~96% of the time for FloWM, while the accuracy prediction of the probe is less than 1% for DFoT. Further, the L2 distance between predicted positions, which can be seen as a proxy for testing whether the FloWM model has learned to be flow equivariant, is significantly lower (0.22) for FloWM, while remaining high for DFoT (2.36, much better than an untrained model at 6.96, but significantly worse than FloWM). See the paper for more details!

We present new results on an additional baseline, an RSSM model based on the Dreamer V3 architecture. Two variants are trained, one with 50M parameters and one with 400M parameters, following the description and hyperparameters in the paper. They are trained for 300k steps, the same as the other baselines. In the table below, we can see that the RSSM model performs poorly, approximately matching the performance of the other baselines. Further, model scale improves the quantitative metrics slightly but not significantly, and perceptually the videos still look quite poor, reinforcing our hypothesis that structured memory is important for solving partially observed dynamic world modeling problems.

Figure and Table caption: we present results on downstream planning experiments in the Blockworld environment using the existing pretrained world models. The task is to maximize the amount of red pixels on the screen, analogous to being as close to the red block (in this experiment there is one of each color) as possible without being on top of it (reward is maximized at distance of 2 away). We use exhaustive search in an MPC framework with a horizon of 3 and a rollout of length 70 after feeding in 70 observed frames. In total the experiment is run across 8 episodes, requiring no additional trained parameters. In the figure below, we can see that FloWM quickly finds its way to the red block, while the baselines stay at far away. In the table below we can see that the final absolute distance is far lower for FloWM compared to the baselines. Example visualizations of the rollouts for each model are also available below. Though this is a preliminary simple planning experiment, we believe it demonstrates the downstream utility of being able to properly predict dynamics in partially observed settings.

a) FloWM Recurrence relation in 2D. Velocity channels are plotted as rows. At each timestep, the internal flow and action flow compose and act upon the latent memory map representation, which is used to predict the observation at the next timestep. b) Rollout error over time, for FloWM and ablations, and baselines. c) Training loss across batches, demonstrating how the full equivariance allows the model to learn significantly quicker.

Below we provide full rollouts on the Dynamic Partially observed, Static partially observed, and dynamic fully observed splits of the 2D MNIST World Dataset.

Below we report the compute and memory cost of rollouts with FloWM, DFoT, and DFoT-SSM models. We evaluate each model at batch size 8 on the Dynamic Blockworld Validation set on a single H200 GPU. We evaluate the each model under two different settings: one for one step of autoregressive generation (1 frame), and one for the inference setting we actually use (70 frames). We report the rollout time in seconds, the throughput, the resident allocated memory (GiB allocated before the generation begins), and the peak delta (measuring the additional memory allocated during rollout above the resident alloc baseline). We can see that FloWM remains substantially more efficient than the baselines as the sequence length increases, due to the recurrent nature of the model. Importantly, memory usage grows much slower with sequence length than the other baselines.

The below table reports the difference in memory allocation and rollout time depending on the map and FoV size. The default configuration uses a 33x33 spatial map and a 60 degree field of view. We can see that map size and FoV have an impact on the Peak Delta, and the throughput, but not the resident alloc; this originates through difference in the KV and Map update tokens. Through this we can understand the scaling and tradeoff of having a particular size map and FoV.