Title: Zero-shot World Models via Search in Memory
Venue: NeurIPS 2025
Accepted: yes
Keywords: world models

reinforcement learning

imitation learning

similarity search
Abstract: World Models have vastly permeated the field of Reinforcement Learning. Their ability to model the transition dynamics of an environment have led to tremendous improvements in sample efficiency for 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 shows stronger performance in long-horizon prediction with respect to the baseline on a range of visually different environments.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces a zero-shot world model that leverages similarity search and stochastic representations (via VAEs) to predict environment dynamics without explicit training. The approach encodes trajectories into a latent space, retrieves similar transitions, and probabilistically estimates future states. It is compared to PlaNet, a well-established model-based RL method, showing competitive performance in next-step and long-horizon (e.g., 20-step) predictions across diverse environments (SuperTuxKart, Minecraft, Atari). Key contributions include: (1) demonstrating the feasibility of memory-based world models, (2) benchmarking against learning-based models, and (3) identifying task-specific applicability and limitations.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths** :\n\n1. **Novelty**: Proposes a training-free framework for world modeling, diverging from traditional data-hungry approaches. This offers a lightweight alternative for resource-constrained scenarios.\n2. **Empirical Validation**: Demonstrates competitive performance with PlaNet in long-horizon predictions (e.g., 20 steps) on visually complex environments, such as SuperTuxKart and Atari.\n3. **Simplicity**: The method is computationally lightweight (deploying in seconds) and works with small trajectory datasets (e.g., five trajectories).\n\n**Weaknesses** :\n\n1. **Theoretical Gaps**: Lacks rigorous analysis of why similarity search approximates actual dynamics. The assumption of Markovian transitions is unproven, and error propagation mechanisms are unclear.\n2. **Data Dependency**: Performance heavily relies on the quality and coverage of the trajectory. Results on Minecraft show significant degradation, indicating poor generalization in open-ended environments.\n3. **Action Selection Limitation**: Assumes future actions are known, decoupling dynamics prediction from policy optimization. This restricts applicability to planning tasks where actions are externally provided.\n4. **Reproducibility Concerns**: Implementation details (e.g., retrieval efficiency, hyperparameters) are sparse, and code and data availability are ambiguously stated.\n\nQUESTIONS:\n1. **Theoretical Justification**: What guarantees exist that retrieved local transitions approximate global dynamics? Is the Markov assumption empirically validated?\n2. **Generalization Limits**: Why does performance degrade in Minecraft? Are trajectory coverage statistics (e.g., diversity metrics) correlated with errors?\n3. **Action-Conditioning Dependency**: How would the model handle unknown future actions? Is there a mechanism to jointly optimize actions and dynamics?\n4. **Efficiency Claims**: How does retrieval time scale with the trajectory count? A direct comparison with PlaNet’s training/inference costs is needed to validate \"lightweight\" claims.\n\nLIMITATIONS:\nAuthors acknowledge data dependency, hallucinations, and variance in retrieval. However, they could quantify trajectory diversity requirements and propose mitigation strategies (e.g., causal priors or external knowledge integration).", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 5, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-21"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this paper the authors propose to use a world model zero shot via latent prediction with similarity search. They base their world model implementation on the PlaNet state space model [1],  from the family of Dreamer models [2].  The key innovation here is that rather than to learn a latent dynamics model as is done in [1], the authors here propose to design an algorithm that does this zero shot via similarity search [3] (a search method based upon a similarity metrics like L2, SSIM, KL etc.) and stochastic representations.\n\nThe stated main contributions of this work are:\n\n 1. Demonstrate feasibility of  memory-based world models.\n 2. Explore capabilities of said model and compare it to a training based baseline.\n 3. Determined tasks for which this approach works well, and where it is infeasible.\n\nThe authors define their goal with the following problem statement:\n\n\"given a dataset D of state-action pairs (x_t, a_t) ∈ D, where xt is an RGB image representing the state of a system at timestep t ∈ [0, T], can we predict the transition dynamics P(x_t+1|x_t, a_t) of an environment without learning them?\"\n\nThe authors observe that they can use similarity search over latent state space at t=t+1 to model z_t+1 effectively defining a latent state space transition model.  Thus given a set of trajectories each made up from successive transitions the set of latent transitions, Z, can be inferred over which search operations can be carried out to effect zero shot modeling of the state space model.  To test their method three model types are tested: *Rollout*, *Replay-L2*, *Replay-KL*.  *Rollout* stores independent latent trajectories which, when given a reference latent, can be used to estimate the next step latent distribution. *Replay* is simpler in that transitions independent of overall temporal dependence are retained and L2 and KL similarity variants are tested.\n\nBaseline is the PlaNet state space model (SSM) trained on trajectories from their three model game domains (SuperTuxkart, Atari, Minecraft) which is compared to the Zero Shot model variants. The authors report KL, L1, and SSIM scores on the reconstructed trajectories along with the visual samples demonstrating top performance for *Rollout* & *Replay-KL*. The authors also investigate the effect of data volume and the effect action conditioning has on the models.\n\n[1] Hafner, D., Lillicrap, T., Fischer, I., Villegas, R., Ha, D., Lee, H., and Davidson, J. (2018). Learning latent dynamics for planning from pixels. arXiv preprint arXiv:1811.04551.\n\n[2]  Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. arXiv preprint arXiv:2010.02193, 2020.\n\n[3] https://en.wikipedia.org/wiki/Similarity_search\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n\n[Quality & Clarity] The overall structure of the paper is well organized and makes clear the primary goal of the work, directly stated at the beginning of section 3, and the related and relevant past work, the methods, the experiments and results. The writing is also for the most part clear and concise, and the description of the models and experiments are for the most part reasonable descriptions although I think there could be some more descriptive artifacts (more on this below). \n\n[Originality] The authors build on past work done on Dreamer state-space models in a novel way. Using World Models zero shot by utilising a memory and similarity search is the main novelty of the paper and a good direction of inquiry that the authors appear to make work.  It is usually very common to use pre-trained language models zero shot so it's great to see this approach taken with a state-space model which can be expensive or tricky to optimise through training. \n\n[Significance] World Models are a growing area of focus that are likely to be increasingly important as agent based applications become more ubiquitous. \n\n**Weaknesses**\n\n[Clarity] In section three I think a figure detailing how the rollout and replay algorithms function in detail would be helpful.  It could be a single figure that could help tie together the concepts throughout the entire section, since these ideas really drive the rest of your work conceptually. It may even be worth including an algorithmic element, but in my opinion at least one or the other would make things a lot clearer.  Further, I think a bit more detail about how the buffer datasets are composed could be helpful, for instance, some concrete indication of where the samples come from (this might be more appropriate for the experiments section). \n\nI think it would be helpful to include some comparison of the computational burden of each approach at inference time, that is, baseline / rollout / replay approaches. It of course would be fair to consider any training resources for the baseline method but I think such considerations may be useful for these types of methods.   \n\nIn section 3.1 it is stated that *\"Then, we extract the next state and estimate the next latent distribution as previously discussed. \"*.  As this is one of the best performing methods it'd be helpful to be specific about what was \"previously discussed\".\n\nA bit more description of the SSM baseline training setup would be helpful if possible - and potentially some architectural details, although I suppose this can be found in the PlaNet paper, but it may be worth at least including in the appendix.\n\nThere is a typo in the second sentence of the abstract: \"Improve\" should be \"\"improved\".\n\nQUESTIONS:\nIs the computational complexity of inference for the method discussed somewhere?\n\nIn section 5.2, the ablation study, it sounds like the rollout method performance is sensitive to the volume of data available while Replay-KL is more robust. Is there a sense, for any of the experimental domains, what this sensitivity looks like?  Can we infer from figure 6? Would this overall make Replay-KL the preferred method? \n\nRegarding the third point of your stated contributions in the intro: \"Determined tasks for which this approach works well, and where it is infeasible.\", I didn't catch whether this is explicitly stated somewhere?  Could you clarify?\n\nLIMITATIONS:\nNot included by the authors directly, they should address this. They could probably discuss some things such as the dependence of their world models on good representations and exhibited some similarities to the chosen zero-shot domain, presumably drastically different world dynamics could lead to serious hallucinations.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-22"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces a method of approximating a world model without training a world model leveraging similarity search and stochastic representations. This method does similarity search in a latent memory buffer constructed using a pretrained variational autoencoder (VAE). The authors introduce three different memory structures and compare these against PlaNet, using reconstruction quality in latent and image space as their evaluation metric. The paper reports competitive or better results than PlaNet, evaluated on SuperTuxKart, Minecraft and Atari Seaquest and Space Invaders.\n\nThe aim of this paper is to introduce an alternative formulation of world models that does not require massive amounts of data and training.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n- The proposal of representing a world model using only memory search seems novel and appealing because of the huge requirement of resources of existing world models.\n- The paper is well written and builds well to form an understandable explanation of the proposed method. The paper also has a detailed description of the experiments and results.\n- The authors propose three different memory structures and evaluate all memory structures in their experiments.\n\n**Weaknesses**\n- the third contribution stated in the introduction is \"third, we determine a range of tasks for which such models are applicable, and clearly state situations where they are unfeasible.\" However, the paper does not clearly establish where the method fails or systematic analysis of feasibility, instead only pointing out that their method struggles with Minecraft because of insufficient data but this is far from clearly establishing where the method fails.\n-  The chosen benchmarks are favourable for similarity search and I feel they do not effectively stress test similarity search vs a learnt world model. For example, the authors design a memory buffer that can account for actions that result in huge visual differences in future states (the most obvious use with the proposed method) but do not evaluate on environments where this occurs. \n- The entire evaluation of this work is based on the assumption that high-fidelity image reconstruction is a good proxy for a good world model. However, world models are used for task learning so a model that filtered out task irrelevant features would be a better world model then a perfect pixel-for-pixel reconstruction. I would argue that using only image reconstruction as the evaluation metric is not sufficient to determine if this method is a valid alternative to existing world models.\n- Following from the previous point, the authors compare against PlaNet by comparing the quality of image reconstruction between their method and PlaNet (\"benchmark specific aspects of the compared models, namely the ability to reconstruct in a seemingly infinite observation space, and the ability to focus on small details.\"). However, the PlaNet design is optimized for policy learning and control and so claiming competitve or superior results over PlaNet seems disingenuous without including policy-learning experiments.\n- The authors point out that they rely on fixed actions and claim this is okay because they are only focused on the quality of dynamics prediction over time. However, the absence of exploring how their method handles autonomous decision making limits the paper's relevance. Without evaluating this method in a decision making process I don't think the authors can claim this work is a valid alternative to existing world models.\n\nQUESTIONS:\nI don't have any specific questions. Overall I feel that in the current form, the work does not provide enough evidence to reach the self-stated goal of offering a feasible alternative to existing world models. The work needs additional evaluation that effectively stress-tests similarity search and shows the feasibility of this method in a decision making process.\n\nLIMITATIONS:\nThe authors don't note their assumption of image reconstruction as a sufficient proxy for world models.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 1, "ORIGINALITY": 3, "DATE": "2025-07-01"}

comments: SUMMARY:
This paper leverages similarity search and stochastic representations to approximate a world model without a training procedure. Comparison experiments with PlaNet showed comparable quality of latent reconstruction and reconstructed images. The model shows stronger performance in long-horizon prediction with respect to the baseline on a range of visually different environments. Besides, the authors mentioned a possible limitation of this method: "Conversely, our method may suffer from hallucinations and lack of generalization when paired with insufficient data, as shown in the Minecraft benchmark."

STRENGTHS AND WEAKNESSES:
Strengths:
- The paper shows a new Zero-shot World Model without a training procedure. The reviewer is not very familiar with this field, but feels that this setting is still very meaningful.
- Comparison experiments with PlaNet showed comparable quality of latent reconstruction and reconstructed images.
- The paper is well-written and easy to follow.

Weaknesses:
- The images generated by the world model are very blurry and cannot be clearly seen by the human eye. I am not sure whether this generation quality can provide any useful information for the model, or how the action model can benefit from this world model.
- The experiments only show the results on generation metrics, KL divergence, L1 distance and SSIM. How the action model utilizes the world model has not been evaluated.
- (minor) I am quite curious about the comparison and discussion between this type of world model and the method of using diffusion to achieve video generation as a world model.

In all, if the author can thoroughly address my concern, I am very willing to raise my score.

QUESTIONS:
Please see Strengths and Weaknesses above.

LIMITATIONS:
yes
TITLE: Official Review
CLARITY: 3
IMPACT: 2
ORIGINALITY: 3
DATE: 2025-07-03