Zero-shot World Models via Search in Memory

TL;DR

The paper proposes zero-shot World Models that predict environment dynamics by retrieving latent transitions from a memory of VAE-encoded trajectories, avoiding gradient-based training. It introduces Rollout and Replay memory structures, with L2 and KL-based retrieval, to estimate next states and extend predictions over horizons up to $h$, while optionally conditioning on future actions. Across SuperTuxKart, Minecraft, and Atari, the memory-based approach achieves competitive latent reconstruction and image-prediction quality compared with PlaNet, with strong performance on long-horizon dynamics in several visually diverse environments. The work highlights trade-offs between memory structure, retrieval metric, action conditioning, and latent-space coverage, arguing that these zero-shot models are a lightweight and effective alternative in smaller-scale tasks, though their efficacy scales with memory coverage of latent trajectories.

Abstract

World Models have vastly permeated the field of Reinforcement Learning. Their ability to model the transition dynamics of an environment have greatly improved sample efficiency in online RL. Among them, the most notorious example is Dreamer, a model that learns to act in a diverse set of image-based environments. In this paper, we leverage similarity search and stochastic representations to approximate a world model without a training procedure. We establish a comparison with PlaNet, a well-established world model of the Dreamer family. We evaluate the models on the quality of latent reconstruction and on the perceived similarity of the reconstructed image, on both next-step and long horizon dynamics prediction. The results of our study demonstrate that a search-based world model is comparable to a training based one in both cases. Notably, our model show stronger performance in long-horizon prediction with respect to the baseline on a range of visually different environments.

Figures (35)

• Figure 1: A simple schematic of the different data structures and sampling procedures used in our work. A replay buffer stores latents without a specific order; retrieving from it consists of computing a suitable distance of all stored data points (grey) w.r.t. the query (yellow), and retrieve the $k$-closest ones (blue). In a rollout buffer, latents are stored according to their temporal order; in this case, the search procedure only selects the $1$-most similar latent from each trajectory.
• Figure 2: Long-horizon reconstruction of $20$ consecutive frames in SuperTuxKart. The first row shows the real sampled sequence of frames. Each row from second to last corresponds to a model, respectively, a SSM baseline, a search-based world model with each trajectory encoded as a separate rollout, an L2-search-based world model with no constraint on trajectories, and a KL-search-based world model without temporal constraints. To predict the next frame, each model can only use the predicted context from the previous timestep.
• Figure 3: Numerical comparison of the models on the SuperTuxKart reconstruction benchmark for (A) one-step and (B) long-horizon predictions. Values are computed by averaging over five tracks, selecting $20$ random images from a disjoint set of test trajectories. For each value, we report mean and variance at each timestep, up to the horizon fixed at $t=20$.
• Figure 4: Cumulative benchmark for (A) one-step and (B) long-horizon predictions in Minecraft. The values are averaged over two tasks, "Treechop-v0" and "Navigate-v0". For each task, we select $20$ random transitions from the test dataset and report the evolution of KL divergence, L1 distance and structural similarity over time.
• Figure 5: Error measures over two tasks of the Atari benchmark, separated in (A) one-step and (B) long-horizon. In our experiments, we use "Seaquest" and "Space Invaders" to explore the limitations of our models in reconstructing small details. For each task, we average over $20$ random transitions of the test set and reconstruct for $t=20$ timesteps.