Zero-shot World Models via Search in Memory
Authors: Federico Malato
Ville Hautamaki
Venue: NeurIPS 2025 poster ()
Keywords: world models
reinforcement learning
imitation learning
similarity search
Primary area: reinforcement_learning
TL;DR: We compare the capabilities of modern neural-networks-based world models with a simple, search-based alternative for dynamics prediction.

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.

Zero-shot World Models via Search in Memory
Federico Malato School Computing
Ville Hautamäki
School of Computing
University of Eastern Finland Joensuu, 80101 federico www..co....atttuee...... fi
University of Eastern Finland
Joensuu, FI 80101
ville www.hhattaaaannna.aa @uefiie@uf..oo
Abstract
World Models have vastly permeated the field of Reinforcement Learning. Their ability model the transition dynamics of an environment have greatly improved sample efficiency in online RL. Among them, the most notorious example is Dreamer, model that learns In this paper, we leverage similarity search and stochastic representations to approx- imate world model without a training procedure. We establish a comparison with PlaNet, well-establishe 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 a training based one in both cases. Notably, our model show stronger performance in long-horizer prediction with respect to the baseline on range of visually different environments.
in diverse of image-based environments.
world model of the Dreamer family. We evaluate the
Introduction
World Models (WMs) ? have played fundamental role in the most recent advances in reinforcement learning (RL) 19 their ability predict future states of a process, along with the possibility to enhance planning have led to tremendous achievements in autonomous agents. WMs have demonstrated their impressive capabilities in range of tasks, from playing small scope (9] and open ended [10] video games, to autonomous driving [4, up to robotic control 220) Recently, WMs have shown also to be adaptable diverse domain [10]. enlarging the pool of their potential application
even further. Moreover, recent discussions hint a possible connection between world models and
Large Language Models (LLMs)
However, learning the dynamics an environment is not a trivial task: modeling temporal dynamics implies learning underlying causal structure of the environment (5] 16], which in turn implies strong generalization and adaptability skills. As such, WMs typically require massive amounts of data construct solid internal model. Moreover, they are subject to hallucinations and error accumulation in long-term prediction ( which limit the effectiveness of the prediction. To ease this limitation, more complex architectures are required computational resources.
3 which in turn require more data and more
In this paper, introduce alternative formulation of World Models derived from similarity search and probabilistic modeling, which we refer to as zero-shot World Models due to their independency from standard training procedure. The contributions our study three-fold: first, explore the theoretical feasibility of a memory-based world model; second, we explore the capabilities and well-known, learning-based model; third, we determine range of tasks for which such models are applicable, and clearly state situations where they are unfeasible. We remark that our aim is explore a valid alternative current world models, focusing on the task of dynamics prediction and reconstruction, rather than action selection.
compare
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

2
Related Work
Our study draws its main inspirations from PlaNet [8 and its evolution Dreamer we hase our study on previous work on similarity search, namely [13] 15 12 briefly revise and introduce the main concepts of each work.
Moreover, In this Section, we I
9 10
In [8, authors propose PlaNet, model-based RL agents that learns to plan from pixels. In their study, authors model the state of an environment as composed by a deterministic and a stochastic part. To successfully predict future states, they define a recurrent state-space model (SSM) composed of a vari- ational autoencoder (VAE) [T] an observation model P(o,|s,) transition model P(s (1-t-1,,;1)) 1). a reward model (r))(ss) and recurrent, deterministic state model h2 = f(h St-1, Qt-1). While leaving the general structure of the model substantially unaltered, evolutions of PlaNet introduce. respectively, a discrete underlying distribution of the latent space 9 and an improved loss for more stable predictions [10].
In I1 authors define locally weighted learning (LWL.), a framework to train a model for continuous control by using an ensemble of local models. In particular, the dataset projected into metric state space and divided into neighborhoods, hence producing subsets of closely related data. Then, a local model is trained on each specialized dataset. Finally, an agent selects actions by querying each local model and performing a weighted average over their answers.
In [15], authors apply LWL in the context of robotic manipulation. In particular, they demonstrate that separating representation and behavior learning improves robustness in robots. In their study, they pre-train an encoder on a dataset of images to extrapolate suitable latent dataset for their task; then, they apply LWL to predict actions from an ensemble of local agents.
Zero-shot Imitation Policy (ZIP) [3] illustrates how agent can be successfully controlled without learning even in open-ended tasks: using tasks from Minecraft [6] as their benchmark, they encode temporally extended latents in latent space using Video PreTraining (VPT) [2] and track the divergence of the current state from a retrieved reference state in latent space. For each timestep, they copy the action of the retrieved sequence. When the two states become too distant, and a new sequence of actions is followed.
search is repeated,
In 12). authors combine ideas from LWL. and ZIP to adapt a learning- hased agent online, leveraging only Bayesian statisties. In detail, authors pair a pre-trained imitation learning agent with a search- based policy. At each timestep, the search-based policy retrieves batch of latents similar to the current state, from which they build a probability distribution for the current state following the "suggestions" of an expert. Then, authors combine the imitation policy and the expert policy suggestions by building a posterior distribution over actions from the two.
Zero-shot World Models
We derive our approach by combining previous studies similarity search [3 15 12
with
probabilistic modeling, specifically Variational Autoencoders (VAEs) [[]] We state our problem as follows: given dataset D of state-action pairs (I,, (z), E D, where 2; is RGB image representing the state of a system at timestep E [o, T]. can we predict the transition dynamics PP P( ze, at) of an environment without learning them?
Initially, we train VAE to reconstruct images 2t from our dataset (Itt, ~ D. Importantly, VAEs operate in two consecutive steps: first, given an image It, stochastic, latent representation & ~ go( Ze (z)) is obtained. Then, an approximation of the initial image ž ~ po(Xx (2) is recovered by passing through the decoder. After training,
encode each 2, to obtain a latent dataset = {(,,) ~ qa Ze (xt), (z,, at) ~ D). Similar to previous work [TT] we choose qo Ze (x:) to be multivariate Normal distribution (u Ex(())
Following previous work in similarity search [ [5], given a latent 2, at timestep t, we can easily produce batch of K relevant samples {(21r ak,>) } 1= E [0, T7. Moreover, assuming that the retrieved latents are sufficiently close to each other, we could estimate go Ze| (e]]=)). From this, a natural question arises: what if, instead of approximating 4a Ze (a( directly, we moved one step forward in time from each of our retrieved latents? Namely, what would happen if we tried to approximate
qo (Z+1 (x++1) from ((šk,++1, k
2

---

Intuitively, if we could do that, then we could sample general representation of the expected next state in time, 2++1 (o(z+1|(z41)) We point out that since the pairs (Ir, 0t) represent states and actions of system, we could view them as trajectories T {( Fi, By exploiting this simple trick, we can reshape also the latent dataset 2 to store transitions (2t, Qz. 2++1) instead. Now, whenever we search for latents that are similar to 2t. we retrieve a batch of relevant transitions { 4k.r: 2k.r+ ) k=1- In summary, given reference latent ere,, and a latent dataset of transitions 2, we can obtain a new latent 4 i by leveraging similarity search and the stochastic representation of a VAE. That is, we can predict an approximation of the next state 2 ~ 1( from the current one.
Notably, the last statement is very similar to the transition dynamics P(s (sc, ae) of Markov Decision Process (MDP) [9]. Hence, we advance our theory this direction: can we approximate the transition dynamics by repeating the same procedure, but imposing a certain action? Following from our previous remarks, the answer to this question yes: given reference pair (2ref,t, azef.t), we can either perform an unconstrained search on 2zet,t obtain a batch {(šk,r ak,r +1)++))) and subsequently remove latents for which ag,7 * Gref,s, or directly restrict the search to Zarae, = {(Zr, 07, 27+1) E 2 a+ = aref,t)-
Until now, we have established that given pair (zt, az) from set of trajectories Z, we can approximately predict how the state will evolve one step in the future by approximating [P | az) in latent space through similarity search and probabilistic modeling. In summary, we have effectively recovered the functionality of a next-state predictor without learning. Still, an effective World Model would ideally extend its prediction to longer horizon h = At, At >> 1, as demonstrated by PlaNet, Dreamer, and similar models [ Therefore, how can we extend our idea to include long-horizon predictions?
Intuitively, we could apply the same "retrieve & reconstruct" procedure iteratively. More formally, given an initial pair (z,,ag ~ 2 and the next predicted state 2t+1 ~ P(=++1lze. at), we could search a new batch {āā,,r 4k.5; 11,r+1)] LK using ê.ạ1 and estimate (2 Ze+ G(+1). effectively evolving the state for a second step in the future. We provide some implementation details of this procedure in the supplementary material of this paper.
However, one problem arises: how do we select a&+1? ? From our latest search, we have retrieved a batch of actions {01 {01.++1..... . which, in general, will not be all equal. One solution to this includes estimating a probability distribution over actions P( 1& and sampling from it, similarly to [12] for the discrete case and [5] for continuous actions. Although this valid alternative, it implies moving the focus towards the problem of action selection, which is outside the scope of this paper. We solve the problem by making an additional assumption that highlights the effects of dynamics prediction. To avoid the fuzziness that comes from sampling future actions, we assume complete access to future actions. As such, given an initial latent 2,, we retrieve a sequence of deterministic actions {a...... J } up to an horizon h = t + At. We remark that, while this assumption would make our approach unfeasible for an action selection process, we are interested in assessing the quality of dynamics prediction over time.
Intuitively, searching constitutes a major factor in the success of our approach. As such, without completeness, we propose three different structures for searching. In the following Subsections, we highlight the unique traits of each and discuss their strengths and weaknesses. A visualization of these modalities for retrieval, along with visual representation of the different search procedures is provided in Figure
3.1 Rollout buffer
In a MDP, a trajectory can be loosely defined a succession of transitions (Iz, az I(+1). In- tuitively, consecutive transitions are temporally dependent on each other. To reflect this fact, we propose to use a structured storage method for the encoded dataset Z. However, differently from a MDP, our transitions are incomplete, refer to a transition as a tuple (2,, d,, 2 Ho, (x:))) where (x, and op. are the output vectors of the encoder of a VAE used sample é, ~ (N)( (x)(
lack information of the reward. As such, in our study
To reflect the temporal dependency across consecutive transitions and the substantial independence across trajectories, draw inspiration from the concept of rollout typically used in on-policy
3

---

Replay buffer
Rollout buffer
Latent
Query
Retrieved
Distance
Trajectory
Figure 1: A simple schematic of the different data structures and sampling procedures used in work. A replay buffer stores latents without specific order; retrieving from it consists of computing 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.
RL 19 17 18 In practice, we define rollout buffer, specifically designed to keep each trajectory temporally ordered and independent from the others.
Given N trajectories, a rollout buffer stores them in a way that generates N independent threads; each thread has specific order and, given starting point in time, it can be explored only in one direction. However, multiple threads can be explored independently. In a rollout buffer, retrieval is operated
follows: given reference latent 2ref,t: we retrieve the 1-most simi- lar latent from each trajectory T using L2 distance and combine them in batch of transi- tions { : 0r,1 1,1, 1 ( or, (x+)), .......... (Zt,n; Ut,n, It+1,n, l, (Lt.n (2ix) Tt. ])]. Then, we extract the next state and estimate the next latent distribution as previously discussed. Similarly, condi- tioning the prediction the action immediate by restricting each search to To, = {(2, a, z', o),
T|a a.). In the remainder of this paper, will refer to this method Rollout.
3.2 Replay buffer
While rollout buffer heuristically gives structure to the search space, we propose simpler alternative inspired by off-policy RI. (19] 14]. In detail, we convert each pair (zt, a.) E Z transition (=e: ae, St+1; ( 2t41- (s.)( and store them in a replay buffer, which impose no constraint
on ordering or temporal dependence. We propose two alternative formulations of this method, which differ in the metric used for searching. From this, other differences follow, which we detail in the next paragraphs.
L.2 distance When searching, we retrieve the k-most similar transitions by computing the L. distance between the reference latent eref,r and each encoded If conditioning, we restrict the search to Zo,- This way we extract a batch transitions, from which we estimate 2 We name this method Replay-L2.
KL divergence Taking advantage of the vectors (z( and ds(x), we propose to retrieve the most similar distribution directly, rather than estimating from a batch of latents. To do so, given transition (zt, at, 2t+1, (4,1,,,, (x) (x:),, we build reference distribution Neat (z), (x)) index of the closest
and compute the KL divergence with all the stored distributions, extracting one. Then, we retrieve the subsequent transition and estimate غ from it. Similarly to the other methods, action conditioning simply restricts the search space by masking irrelevant transitions according to their action. Perhaps unsurprisingly, refer to this method Replay-KL.
3.3 Comparing the methods
Intuitively, using replay buffer gives more freedom for retrieval. However, conditioning on the action might severely impact the process, as the average distance across retrieved samples will generally increase. Hence, expect some high variance in the retrieved batches, indicating uncertainty in

---

estimation. Nonetheless, each batch will produce valid point for sampling, regardless of its semantic meaningfulness.
Searching using KI. divergence solves some of the previous problems: first, since the encoded μ and obtained from actual trajectories, they correspond to meaningful regions of the latent space. Moreover, given that in this case we perform no batch estimation, we are guaranteed to sample a meaningful E. However, action conditioning might still pose threat: whenever we diverge too much from the reference, we might get consistent, but unrelated samples.
Finally, using rollout buffer constitutes a middle point between the other two alternatives: by imposing independency between trajectories and by sampling from each them, we intuitively reduce the variance in the latents batch. However, if there are no similar transitions, or if transition with the imposed action is temporally too distant from the reference, we might still observe hallucinations.
Experiments
We test our approach against PlaNet 8 We justify the choice of using a legacy method from the Dreamer family comparison by highlighting the good properties of this model. First, PlaNet a standard VAE modeled as Normal distribution; conversely, other models of the family, starting from D2eamerV2 [ model the latent space after discrete distribution, which may be harder to study. We remark that in our study we assume a stochastic representation, but impose no condition on a specific one. As such, using Normal representation comes with no loss of generality. Second, PlaNet includes less modules than its evolutions. As such, observing the isolated effect of the dynamics predictor is easier.
In each experiment we compare four models, namely the baseline PlaNet along with the three versions of our approach as detailed in Section 3 We test the models on range of visually different, image-based environments extracted from well-known benchmarks in RL. Specifically, we use five tracks from SuperTuuxart consistent visuals with very diverse features; two tasks from Minecraft [6], to benchmark tasks with a seemingly limitless, diverse observation space; and two tasks from Atari [14], study how WMs predict the dynamics of small details. All our models are trained on consumer hardware, consisting of a single RTX 4080 GPU.
racing game with complex visuals, test the performance of WMs on
For each experiment, we compare the models both in latent and image spaces, using KL divergence for the first, and L1 distance and structural similarity (SSIM) for the decoded images. In particular, we use the KL divergence to assess how well the dynamics of the tasks have been reconstructed; conversely,
MSE and SSIM determine the visual relevance of the reconstructed dynamics. We highlight that, despite our best efforts, a VAE-reconstructte image will inevitably lose quality with respect the original. Similarly, the KI divergence assesses the overall distance between pair of distributions, but does not consider positional differences. As such, we invite the reader to consider the measures as correlated, and consider the absolute values of one in light of the others.
We test our approach on both single-step and long-term predictions, using a horizon of = 20 as reference for the latter. For single-step experiments, we refresh the hidden state of each model at every timestep, using the actual observation. For long-term comparisons, we let the model evolve independently until the horizon is reached, using the generated latents intermediate representations. In both cases, we extract 20 random starting samples from a separate batch of unseen trajectories and average each measure over them.
Furthermore, since our approach uses a number of encoded trajectories for retrieval and prediction, we conduct ablations study two track of SuperTuxKart. We test the performance of our model for 5, 6, 7, 8, 9, 10, 15. 20 & 30 encoded trajectories. To establish a fair comparison, we train a SSM for each subset of data, using the same trajectories in both cases. Similarly to other experiments, we extract 20 random samples from a disjoint set of test trajectories and report the average for each measure.
Finally, we recall how our approach enables us to estimate P(*++ar/s)) with / E N, At > that is, how the state is action. Since we deem this difference interesting, we compare our models when acting with and without action conditioning.
as well, expected to evolve from the current one with restriction on the chosen
5

---

Tacpoy 1L Raploy Rallue
Timestep
Figure 2: Long-horizon horizon reconstruction of 20 consecutive frames in SuperTuxKart. The first row shows the real sampled sequence of frames. Each from second to last corresponds to model, respectively. a SSM baseline, a seurch-based world model with each trajectory encoded as separate rollout, 1.2-search-basee beaaed world model with constraint trajectories, and (L-ssarrch-baee world model without temporal constraints. To predict the next frame, each model can only use the predicted context from the previous timestep.
Baseline
Roliout
A
Replay-L2 L2
KL Divergence
Replay-KL
Distance
200
0.1
SSIM
0.4
16
B
KL Divergence
Distance
600
S5IM
400
0.4
200
0.2
16
Timestep
Figure 3: Numerical comparison of the models the SuperTuxKan reconstruction benchmark for (A) one-step and (B) long-horizon predictions. Values are computed by averaging over five tracks, selecting 20 random images from disjoint set of test trajectories. For each value, we report mean and variance at each timestep. up to the horizon fixed at 20.
5 Results & Discussion
In this Section, we present the results of study. The Section is organized in three subsections, each dedicated to specific experiment as described in Section 4 Due to page limit constraints, we report only small subset of the available results and present only a handful examples. However, report more details and examples in Appendices.
5.1 Prediction & reconstruction quality
Figure 2 shows example of long-horizon dynamics prediction using each model. In the one-step case, we see how re-initializing the hidden state of the model benefits their performance. We highlight how PlaNet hallucinates after 10 timesteps, while Rollout and Replay-KL. maintain consistency in their prediction. In particular, Replay-KL, shows an unmatched resemblance to the real sequence. In both cases, Replay-L_2 is affected by noticeable hallucinations, thus making it completely unreliable. In the supplementary material, we include more examples from different tracks, and examples of one-step predictions.

---

A
Replay-L2 Replay-KL
Divergence
Distance
SSIM
02
0.1
16
B
KL Divergencce
Distance
2000
0.2
SSIM
1000
0.1
16
16
Timestep
Figure 4: Cumulative benchmark for (A) one-step and (B) long-horizon predictions in Minecraft. The values are averaged over two tasks, "Treechop-vee and "Navigate-ve 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.
Rollout Replay-L2
A
KL Divergence
Distance
Replay-KL
SSIM
200
0.8
100
B
KL Divergence
400
Distance
S5IM
200
0.8
16
Timestep
Figure 5: Error measures 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, average over 20 random transitions of the test set and reconstruct for t = 20 timesteps.
The numerical comparison reported in Figure confirms our qualitative assessment: for one-step predictions (Figure 3)) the baseline, Rollout and Replay-KL are visually indistigguishable while in long-horizon regime Rollout and especially Replay-KL achieve remarkable performance over the baseline. notable how all our models are generally closer in distribution to the real representation than the baseline. We explain this mismatch by remarking that VAF-reconstructed image will necessarily carry some error. As such, comparison between latent reconstructions and decoded images not trivial.
Notably, Replay-L2 features small KL error, but significantly higher LI and lower SSIM values. We explain this fact as follows: searching in replay buffer while conditioning on action may lead to high variance the retrieved batch. Hence, estimating the posterior from it could lead an unstable region of the latent space. Therefore, sampling in that point may produce visually incoherent reconstructions.
the other hand, giving structure to the latent space through the separation in rollouts appear to regularize the reconstruction of 2++1 p(X++1lá searching directly for the most similar prior greatly reduces this instability. Hence, both cases we have more coherent predictions.
As for Replay-KL,
7

---

SSN
Pacplay IL, Replay Ialluu
Timestep
Figure 6: Long-horize reconstruction from Minecraft "Treechop" task. First row reports the real sequence; rows from second to last show the reconstruction of, respectively, baseline, Rollout, Replay- L2, and Replay-KL.
We report the reconstruction results for Minecraft and Atari respectively in Figures 4 We leverage these two environments to benchmark specific aspects of the compared models, namely the ability reconstruct in a seemingly infinite observation space, and the ability focus on small details. In both cases, we confirm the patterns discussed for Super SuperTuxKart Tark Rollout and Replay- KL are generally better than the SSM baseline, while Replay-L2 is worse than any other model in terms of actual reconstruction, while remaining overall competitive in latent space.
By comparing values in Figures 3 & conclude that, expected, Minecraft represents challenging environment for reconstruction: despite the similar patterns, we see how values are overall higher across all measures. In particular, Rollout and Replay-L2 behave similarly both in terms of dynamics prediction and reconstruction. However, the example shown in Figure 6 shows how both models produce a completely unreliable reconstruction.
Despite obtaining similar values in L1 and SSIM, the baseline and Replay-KI produce very dissimilar results:
hand, the SSM baseline produces blurry. incoherent images; the other hand, Replay-KL generates visually appealing sequences that recall the main elements of a scene, but gradually diverges. By isually analyzing the samples, we notice "hit-or-miss" pattern in both models, with Replay-KI. being generally sharper but more variable, and the baseline being much blurrier but more accurate in overall structure. In the supplementary material of this study report some examples of this fact and offer a more detailed analysis to support this claim.
5.2 Ablation study: number of trajectories
We report the results our ablation study on SuperTuxKart in Figure 7 and provide an additional ablation on Minecraft in supplementary material. As expected, the SSM baseline benefits from the greater availability of training data; on the other hand, increasing the dataset only marginally improves models, and only in some cases.
In detail, KL divergence (Figure (7) in the SSM baseline drops steadily from 339.48 169.46, amounting to a total decrease of 50.08%, in the one-step case. In long-horizon predictions, the decrease is less consistent (-26 26..%.).. Similarly, Rollout benefits from more data and reduces the average KI. divergence from 311 .98 to a mere $6.27 (-72.34%). In contrast, Replay-L.2 and Replay- KL are substantially unaffected by the increase in available data. We explain this fact by noting that the result of an unconstrained search changes only if there is no meaningful data. As for Rollout, point out that introducing a temporal constraint on the encoded trajectories severely impacts search, and adding more data seems relax this constraint. However, if we consider the absolute values for each model, we see how the best baseline model (30 trajectories) still reports higher average KL than Rollout and Replay- KLL with only five encoded trajectories.
LI distance and SSIM substantially confirm our previous discussion: decodings from the SSM baseline improve with more data, while Replay-KI. and Rollout remain substantially unaffected, and perform either comparably or better than baseline in all cases. In the Appendices we show some image reconstructions directly related to these results.

---

divergence
dbttacce
150
300
6:00
6.14
0.14
025
0:68
0.43
0.50
S5M
2709
Model
Replay (1.
UMS
Reploy-E. pana see SAL.C 196 2a0d paan 121-30 100.0 221.00 $5.04
81
01
6.6
240
Cm
cm
63
55M4
621.1 100.1
118.1 200.70 311.0
0.11
031
022
Nodel
Replay-1.
118 40.1 004-7 14L-7 ITANO
36.39
031
631
610
0:1
011
944
641
541
cm
on
10
15
20
10
25
10
Number of trajectorios
Figure 7: Results of the ablation study tracks from SuperTuxKart. Each model uses increasing number of trajectories for either training (SSM) or retrieval (Rollout, Replay-L2, Replay- KL). Top row shows results of one-step predictions, bottom row corresponds long-horizon dynamics reconstruction. Heatmaps analyze (from left to right) KL divergence, L1 distance, and structural similarity.
Table 1: Performance for planning in SuperTuxKart. Results are reported using mean reward & standard deviation, and a 95% confidence interval (parentheses).
Model
fortmagma
lighthouse
ressrbobnowooad
snowmountain
volcano island
0.0880 0.1125
3.0780 + 1,5162 (2.6427, 3,5133)
3.1600 + 2.9678
5.3220 3,7877
0.9280 0.6440 (0.7432, 1.1 1.1128)
PlaNet
(0.0557, 0. 1203)
(2.3080, 4.0120)
(4.2346, 6.4094)
5,8100 3.4761
2.6400 +: 2.5952
4.8220 +: 1.1188
5.5560 6.5701
1.5300 0.5849
Rollout
(4.8121, 6.8079)
(1.8950, 3,3850)
(4.5008. 5.1432)
(3.6698, 7.4422)
(1.3621, 1.6979)
4.3500
1,4506
5.1140 +: 9.5549
6.3360 +: 5.2006
4.9560 = 6.5035
3.3480 = 6.1291
Replay-L2
(3.9336, 46764)
(2.3709, 7.8571)
(4.8430, 7.8290)
(3.0890, 6.8230)
(1.5885, 5.1075)
4.7740 3.8107
0.3200 = 0.7558
4.7140 + 0.6515
15.6640 11.8248
4.4780 7.9652
Replay-KL. KL
(3.6800, 5.8680)
10.1030, 0.5370)
(4.5270, 4.9010)
(12.2693, 19.0587)
(2.1913, 6.7647)
5.3
Action selection
WMs are typically embedded in more complex architectures to inform the state with future infor- mation. For instance, our baseline model PlaNet (8) designed for planning, that is, selecting sequence of actions by evolving the environment dynamics from the current observation. Although not the primary scope of this paper, we deemed interesting to explore the planning capabilities of our methods.
We test each model the realized returns over 50 evaluation episodes using the SuperTux: Tux Kart benchmark, as we believe that it represents the typical use case for our approach. Each model is allowed to observe the first state of an episode. From this single observation, agents asked to evolve the dynamics according to their WM and plan the next 20 actions. Then, the planned actions are executed with no room for corrections; finally, the model is allowed observe a new state and plan the next sequence of actions. The process repeats until the end of the episode.
Actions are selected deterministically using the mode of a retrieved action distribution. We report the results in Table highlighting the best model for each environment. For each track, we report mean, standard deviation, and 95% confidence interval.
The results clearly show that our approaches perform generally better than the PlaNet baseline, thus suggesting that search-based approaches can plan. In particular, replay -aafer-based methods yield significant improvements
out of environments. However, from our experiments, no specific method clearly dominates on the others.

---

Polfout wrt.
Raplay
Divergence
Replay 12- n- art.
300
Reploy
Reploy
200
100
Divergence
1000
800
400
200
10 Timestep
9
21
14
15
16
19
Figure 8: Evolution of KL divergence w.r.f. action conditioning for (Top) one-step and (Bottom) long-horizon predictions. A reference model is reported using solid color; corresponding models without action conditioning are characterized by a pastel tone of the same color.
5.4 Action conditioning
Figure 8 shows the results of our study on the effects of action conditioning. We remark that the scope of this experiment is to assess the differences between predicting P( (12++1/2, 22, a,) and P( l2)). Since an SSM can not compute the latter, we only compare Rollout, Replay-1.2, and Replay-KI..
The results confirm the generally poor performance of Replay-1.2: both in one-step (Figure ( top) and longterm (Figure g bottom) predictions, KL divergence significantly decreases. This indicates reduction variance over the batch of retrieved samples, consistently with our speculation. However, by analyzing the decoded sequences related to this study (reported in Appendix B results are still too noisy be considered reliable.
Rollout marginally benefits from the absence of actions in one-step predictions, while apparently remains unaffected in the long-hone horizon regime. Analyzing the decoded sequences (in Appendix B reveals no substantial difference. We explain this fact in conjunction with the result obtained in Section 5.2 intuitively, predicting 1(2) rather than P(2++1|22 2t, az) enlarges the pool of retrievable images.
In Replay-KL, one-step predictions unaffected by action conditioning. Contrarily, long-tern reconstructions show a significant increase in KL divergence for conditioned retrieval. We explain this effect by highlighting that Replay-KL recovers only one example for each search and, in the long-term scenario, only retrieves the first latent. Forcing the model to take a specific action may rarely result in planning
very dissimilar sequence of latents.
However, by considering the one-step case, where the same retrieval produces virtually no effect, this search failure appears to be a very rare occurrence. Moreover, we believe that the best use for our method lies the "no action conditioning" regime. Additionally, we do not deem this error to serious limitation: intuitively, enlarging the data pool in memory is likely to mitigate this effect.
6 Conclusions
We have presented an alternative way to predict the transition dynamics on a number of image-based environments. Our proposal explores similarity search stochastic representations way to build WMs without training. Our results confirm that our proposed WM can act as a valid alternative in a range of scenarios.
Our methods outperform the baseline while also being more immediate in implementation and application. Despite showing acceptable performance in open-ended tasks such as Minecraft and being comparable to the PlaNet baseline, our results suggest that the performance zero-shot WMs depends on latent space coverage of the encoded trajectories, limiting their effectiveness in tasks with vast state spaces. However, they represent a valid, lightweight alternative to regular WMs in small-scoper tasks, such SuperTuxKart and Atari.
10

---

The authors wish to acknowledge CSC- - IT Center for Science, Finland, for computational resources. Additionally, Ville Hautamäki thanks the Jane and Aatos Erkko Foundation for partial funding.
Acknowledgments.
References
[1] Christopher G. Atkeson, Andrew W. Moore, and Stefan Schaal. Locally weighted learning, page 11-73. Kluwer Academic Publishers, USA, 1997.
[2] Bowen Baker, llge Akkaya, Peter Zhokhov, Joost Huizinga, Jie Tang, Adrien Ecoffet, Brandon Houghton. Raul Sampedro, and Jeff Clune. Video pretraining (vpt): Learning to act by watching unlabeled online videos, 2022.
[3] Jingtao Ding, Yunke Zhang, Yu Shang, Yuheng Zhang, Zefang Zong, Jie Feng. Yuan Yuan, Hongyuan Su, Nian Li. Nicholas Sukiennik, Fengli Xu, and Yong Li. Understanding world or predicting future? comprehensive survey of world models, 2024.
[4] Yanchen Guan, Haicheng Liao, Zhenning Li. Jia Hu, Runze Yuan, Yunjian Li, Guohui Zhang, and Chengzhong Xu. World models for autonomous driving: An initial survey, 2024.
[5] Tarun Gupta, Wenbo Gong, Chao Ma, Nick Pawlowski, Agrin Hilmkil, Meyer Scethon, Mare Rigter, Ade Famoti, Ashley Juan Llorens, Jianfeng Gao, Stefan Bauer, Danica Kragic, Bernhard Schölkopf, and Cheng Zhang. The essential role causality foundation world models for embodied ai, 2024.
[6] William H. Guss, Brandon Houghton, Nicholay Topin, Phillip Wang. Cayden Codel, Manuela Veloso, and Ruslan Salakhutdinov. Minerl: large-scale dataset minecraft demonstrations, 2019.
[7] David Ha and Jürgen Schmidhuber. World models, 2018
[8] Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels, 2019.
[9] Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models, 2022.
[10] Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models, 2024.
[11] Diederik P Kingma and Max Welling, Auto-encoding variational bayes, 2022.
[12] Federico Malato and Ville Hautamāki. Online adaptation for enhancing imitation learning policies. In 2024 IEEE Conference Games (CoG), pages 1-8, 2024.
[13] Federico Malato, Florian Leopold, Andrew Mclnik, and Ville Hautamäki. Zero-shot imitation policy via search in demonstration dataset. In ICASSP 2024 2024 IEEE International Conference Acoustics, Speech and Signal Processing (ICASSP), pages 7590-7594, 2024.
[14] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning, 2013.
[15] Jyothish Pari, Nur Muhammad Shafiullah, Sridhar Pandian Arunachalam, and Lerrel Pinto. The surprising effectiveness of representation learning for visual imitation, 2021.
[16] Jonathan Richens and Tom Everitt. Robust agents learn causal world models, 2024.
[17] John Schulman, Sergey Levine, Philipp Moritz, Michael I. Jordan, and Pieter Abbeel. Trust region policy optimization, 2017.
[18] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms, 2017.
[19] Richard S. Sutton and Andrew G. Barto, Reinforcement Learning: An Introduction. The MIT Press, second edition, 2018
[20] Philipp Wu, Alejandro Escontrela, Danijar Hafner, Ken Goldberg, and Pieter Abbeel, Daydreamer: World models for physical robot learning, 2022.

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Methods are described Sections and experiments and results cited the
abstract are discussed in Section
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: Limitations of our method are discussed in Sections 5 and and in Appendices using both text and figures.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: Method builds on top of VAEs theory. Innovative theory is explained step-by- step either in Section or Appendices.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Ye]]
Justification: Every method used in the paper extensively described. All the necessary information is disclosed in Sections ?? and 3
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: Code is provided to reproduce all trainings and experiments. Due to size, we include links the data used for training.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification: Experimental procedure is described in Section 4 code includes scripts with hyperparameters set as used during the study.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: Each experiment has been repeated and reports margins. Exception is made for heatmaps where we report only the mean due to readability of the results.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
12

---

Answer: [Yes]
Justification: Resources used are cited in Section
9. Code ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics httppp////errrrs..ccc/////ittttp
Answer: [Yes]
Justification: Our work adheres the NeurIPS Code of Ethics with no deviation from it.
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: To the best of our knowledge, our study represents primer. While consider its impact to be highly relevant for its field, such relevance does not extend to a broader context.
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: To the best of our knowledge, the limited scope of our models and data do not pose such risks.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: Ground work for the present study is discussed and properly cited multiple times in text, e.g. Sections 2 ??, and Sourced for external datasets are cited in Section 4
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does release any new asset, aside from a of collected human demonstrations for the enviroment "SuperTuxKart".
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: The study does not include research with human subjects, crowdsourcing experiments.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Question:
Answer: [NA]
Justification: The present study does not include studies where participants needed.
13

---

16. Declaration of LLM usage
Question:
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
Justification: Our study is disjoint from LLMs.
14

---

A Prediction & reconstruction benchmark results
In this Appendix we report the evaluations for each environment we tested. We report the results following the same template as Figure 3 [nd Additionally, we discuss each result individually, better highlight environment-relate strengths and weaknesses of each approach.
Baseline
Rollout
A
Divergencc
800
Distance
Replay L2
Replay
S5IM
0.1-
B
800
Divergence
Distance
SSIM
200
0
Timestep
Figure 9: Numerical evaluation in SuperTuuKart "fortmagma" track. Results reported in terms of KL divergence for latent dynamics prediction, and Ll distance & SSIM index for (A) one-step and (B) long-horizon predictions.
Results for the "fortmagma" (Tigure9) reflect the general behavior we have discussed in Section 5 all models achieve significantly lower error in KL divergence, meaning that they are closer the true dynamics in latent space. As for the reconstruction, which we remark could be affected also by an error in decoding the latent, Rollout and Replay-Ki show comparable reconstructions with the baseline, while Replay-L2 is generally affected by hallucinations or inconsistencies.
Baseliee
Rollout Replay.
A
0.6
KL Divergence
1000
Distance
Replay-KL
500
SSIM
-500
16
16
20
12
20
B
Divergence
1000
Distance
0.2
SSIM
500
0.1
12
16
20
8
12
16
20
12 16
20
Timestep
Figure 10: Numerical evaluation in SuperTuxKar "snes in terms of KL divergence for latent dynamics prediction, and L1 distance & SSIM index for (A) one-step and (B) long-horizon predictions.
"snes_rriinowwrreeeee rainbowroad" track. Results are reported
The results of our benchmark "snes ratesbbrbnowwa"" reported in Figure 10 show the importance of interpreting the quantitative and qualitative results jointly. The general behavior of the models is quite different from Figure 9 models Rollout and Replay-KL perform either similarly (KL divergence) or slightly worse (L1, SSIM) than the baseline for one-step predictions, while match or improve in long-term dynamics prediction.
We attribute this difference to the visual complexity the images from this specific environment, as shown in Appendix B the sequence of colors in the track, along with the "falls" that the player might
15

---

experience heavily contribute to the divergence from the real sequence. However, we point out how our models are generally more consistent reconstructing visually coherent sequences, even though they might be slightly more distant in terms of visual accuracy.
Baseline Replay-L2 Replay-Kt KL
A
Divergence
200
0.2
SSIM
0.4
100
0.1
0.2
0.0
12
16
20
16
20
12
16
20
B
Divergence
0.2
SSIM
0.4
0.1
0.2
0.0
16
12
Timestep
Figure 11: Numerical evaluation in SuperTux "volcano island" track. Results are reported KL divergence for latent dynamics prediction, and L1 distance SSIM index for (A) one-step and (B) long-horizon predictions.
The results from track "volcano island" (Figure [1] support the general results we have shown in the main text. Our agents Rollout ad Replay-KL are generally closer to the real sequence in terms of latent dynamics prediction, while either match or slightly improve the visual reconstructions. Notably, the gap in difference is more pronounced in long-horizon predictions, where Rollout and Replay-KL. diverge from the real sequence much more slowly than the baseline.
Baseline
Rollout
A
0.6
KL DiDeeegennce
Distance
Replay
300
SSIM
200
0.4
100
10
B
0.
0.6
KL Divergence
800
Distance
600
SSIM
0.1
0.4
200
12
16
20
8
16
20
12
16
20
Timestep
Figure 12: Numerical evaluation SuperTuxKart "snowmountain" track. Results are reported in terms of KL divergence for latent dynamics prediction, and L1 distance & SSIM index for (A) one-step and (B) long-horizon predictions.
In the "snowmountain" track, the overall trend is also confirmed: all models improve latent reconstructions, while either match or improve the baseline in decoded reconstructions. The only exception is represented by Replay-L2 that, despite achieving better latent prediction, systematically fails to decode it. As reported in main text, we track this effect to the latent space learned by the VAE: while estimating the generative distributions from batch of similar samples correctly points the model in the right direction, sampling from that distribution for the purpose decoding usually results in visually unpleasant images.
As can be seen from the examples differences in this track generally smaller and harder to assess from quantitative measures.
B "lighthouse" features a visually dark theme. As such,
16

---

Baseliee Replay-L2 Replay-KX KL
A
Divergence
600
Distance
SSIM
200
0.1
0.4
0-
0.0
16
12
16
12
16
20
B
Divergence
600
Distance
0.2
SSIM
0.1
0.4
200
0.0
0.2
12
Timestep
Figure 13: Numerical evaluation in SuperTuxKart "lighthouse" track. Results are reported in terms of KL divergence for latent dynamics prediction, and L1 distance SSIM index for (A) one-step and (B) long-horizon predictions.
However, we from Figure 13 how the general trend is confirmed, with Replay-KL and Rollout either matching or improving over the baseline, while Replay-L2 behaving better than the baseline in latent reconstructions (KI divergence), while failing to properly decode its predictions.
Baseline
Rollout
A
Replay-
Divergence
Distance
Replay-
1000
SSIM
500
B
KL Divergencc
Distance
3000
S5iM
2000
0.2
1000
Timestep
Figure 14: Numerical comparison for Minecraft "Treechop-vO" task. Results divided in (A) one-step and (B) long-horizon predictions.
The results from "Treechop-vO" (Figure 14 from the Minecraft benchmark further confirm findings. Intuitively, the state space of the task limitless, with a few "bottlenecks" occurring in the proximity of tree. Nonetheless, and on average match the baseline on visual reconstructions. We found this result surprising, as agents only store limited number of trajectories, while SSM should theoretically be able generalize better.
agents generally better at reconstructing latent dynamics,
Similarly, in "Navigate-vO" human expert is asked to explore the open world of the game, with no particular aim. Hence, it is nearly impossible to fully represent the full space without having access to infinite number of trajectories. The results in Figure 15 reflect this intrinsic complexity: this the only instance in our experiments where the latent reconstruction favors the baseline over our strongest model, Replay-Ko KL. However, notice how, despite similar higher KL divergence, models' decoded latents are generally closer to the real sequence.
Our experiments on "Space Invaders". reported in Figure 16 reflect the average trend shown in the main text. Notably, in this case the KI. divergence computed on our models is very low, meaning that the predicted dynamics is generally very close to the original sequence. Also, due the images
17

---

Baseline
A
Divergence
2000
0.4
Replay-L2 Raplay-KL
1500
SSIM
12
1000
500
12
16
16
20
12
B
0.4
Divergence
2000
SSIM
1000
500
14
0.0
16
12
Timestep
Figure 15: Numerical comparison for Minecraft "Navigate-vDl v0" task. Results are divided in (A)
one-step and (B) long-horizon predictions.
Baseir Rollout
A
1.0
Replay Replay
MMW
KL Divergence
400
Distance
.03
0.9
SS/M
200-
.01
B
KL Divergence
600
Distance
1.0
400
SSIM
0.7
8
12
16
12
16
20
Timestep
Figure 16: Numerical evaluation of "Space Invaders" from the Atari benchmark. Results report (A)
one-step and (B) long horizon predictions.
differing only on small details, the error on L1 distance is an order of magnitude lower w.r.t. the other benchmarks, and SSIM is close to perfect.
Basell Rollout
A
.03
1.0
Replay
KL Divergence
300
Distance
Replay
200
.02
SSIM
0.9
100
.01
0.8
0.7
0
12
16
20
36
20
B
1.0
KL Divergence
300
L1 Distance
0.9
SSIM
200
100
.0
0.8
0.7
8
12
16
12
16
20
Timestep
Figure 17: Numerical evaluation of "Seaquest" from the Atari benchmark. Results report (A) one-step and (B) long-horizon predictions.
18

---

SSM Replay-12 Repley-K.
SuperTinKaar
Minecraft
KL 1Ldeegoncc
1000
500
0.7
Figure 18:
comparison between latent space coverage ratio, that is, the fraction of latent space "sectors" covered by the encoded trajectories, and the KL. prediction error of each world model.
Table 2: Pearson correlation index of the tested models representing the correlation between encoded trajectories and KL prediction error.
Model
Pearson Correlation
PlaNet
-0.8269
Rollout
-0.6709
Replay-L2 -0.6711
Replay-KL -0.6366
Similarly, Figure I7 further confirms the validity of our approach, with Replay-KL performing either on par or better than the baseline. Moreover, like in "Space Invaders", the gap KL divergence is abyssal, while LI distance and SSIM are similar, but always favoring Replay-KL, especially in long-horizon reconstructions.
A.1 Coverage statistics
In Section we concluded that our methods are most expressive in smaller-scoped tasks, in which the set encoded trajectories are sufficient representations to page limitations, no empirical result was provided. In this Section, we test the generalization capabilities of our proposed approaches w.r.t. to the PlaNet baseline. latent state coverage ratio of the encoded trajectories with the KL prediction error of each model. Additionally, we compute the Pearson's 's correlation index for each tested environment and model.
the task dynamics. However, due
do so, we compare the
The results are shown in Figure 18 and Table 2
The experiment shows that, as denoted in Section higher coverage of the latent space corresponds to lower prediction error. This supports our claim that our models are better suited for smaller scoped tasks, such SuperTuxKart and Atari, where a significant fraction of the environment dynamics is implicitly reproduced in the encoded trajectories. However, interestingly the strongest correlation is observed in the PlaNet baseline, indicating that both learning and search-based world models benefit similarly from more data.
B Prediction & reconstruction qualitative evaluation
This Appendix reports a number of visual reconstructions of predicted dynamics. For each environ- ment, we report two one-step and two long-horizon predictions, namely positive example and failure case for our models. Given the number of models test, finding a totally positive/negative
19

---

example that include all models is hard. As such, in order to maintain the length of our Appendix reasonable, report the best overall positive/negative examples. Each subsection refers to an environment.
A
D
PORAL
10
15
Timestep
Figure 19: Examples from Super Karrt "fortmagma" track. (A) Long-horizon success; (B) long- horizon failure; (C) One-step success; (D) One-step failure.
we could not find failure case for Replay-KL. However, Rollout and especiallt Replay-L2 are subject to hallucinations both in one-step and long-term reconstruction.
In "fortmagma" (Figure 19
As stated in the main text, "snes rainbowroad" represents particularly hard case for this benchmark due to the visual complexity of the track and to the "falls" that player may experience. As such, in Figure 20we we see how all models can incur into hallucinations.
The track "volcano island" features a of visual variety, as it includes track sections on asphalt, rock, sand, and even water. Such visual diversity generally makes prediction harder. As expected,
"isladd"
20

---

A
D
Timestep
Figure 20: Examples from SuperTuxKar "snes_ raaessswbnoowaadd track. (A) Long-horizon success; (B) long-horizon failure; (C) One-step success; (D) One-step failure.
in Figure 21
we see how, to some extent, all models incur into hallucinations. An exception is represented by Replay-KL, KL, which exhibits strong performance throughout the whole track.
As visible in Figure 22 "snowmountain" all models except Replay-L2 are mostly capable of predicting both long-term and one-step situations. However, we see in Figure 22A that the baseline SSM model might diverge from the reference sequence.
As previously stated, "lighthouse" is peculiar case of this benchmark due to its intrinsic darkness. hard task. However, the tracks also features lightnings occurring at random times, which we use to determine failure cases. In particular, Figure 23 shows Replay-KI. KL predicting lightning when not existing. Also, in section D of the same Figure we how all models fail to reconstruct coherent background.
As such, spotting hallucinations
21

---

A
B
Timestep
Figure 21: Examples from SuperTux "volcano joolcano_islan"" track. (A) Long-horizon success; (B) long-horizon failure; (C) One-step success; (D) One-step failure.
Tasks from the MineRL benchmark feature an almost limitless state space. As such, expecting a perfect prediction from quite simple models such as ours is quite irrealistic. However, in Figure 24 it surprising to see Replay-KI KL achieving quite good reconstructions in A and C. Notably, in this task baseline model (SSM) produces mostly blurry reconstructions.
An even more extreme case of limitless state space is represented by "Navigate-vO", for which we report some examples in Figure 25) Interestingly, quite some trajectories are centered navigating the sea, hence easily leading all models to hallucinations if water is shown in the picture. For this task, finding "good prediction" was particularly hard.
22

---

A
D
Timestep
Figure 22: Examples from SuperTuxKar "snowmountain" track. (A) Long-horizon success; (B)
long-horizon failure; (C) One-step success; (D) One-step failure.
The most notable difference between the models in "Space Invaders" (Figure 26 the consistency with which they predict the alien ships. In general, Replay-KL was the most consistent one, even though it was not perfect.
Similarly, as shown in Figure 27 in "Seaquest" the most notable example of hallucination is repre- sented by the enemy boats that, in some models, flicker throughout the sequence. Like in "fortmagma", Replay-KI. was mostly correct and coherent in predictions, while both SSM, Rollout and Replay- L2 struggle to maintain consistency in their predictions.
23

---

A
B
Timestep
Figure 23: Examples from SuperTuxKart "lighthouse" track. (A) Long horizon success; (B) long-
horizon failure; (C) One-step success; (D) One-step failure.
C Ablation study additional results
In this Appendix we report some visual reconstructions coming from the ablation study presented in Section 5.2) Moreover, we report the results of an additional ablation study performed on Minecraft "Treechop-vO", using [10, 80) trajectories and step of 10. This additional ablation study is also used to compute the wallclock time needed each method to predict the latent dynamics. To fit all the models in one run, the test is run CPU, using Intel i7 12650HX.

---

A
B
TARDAA NO Raplayaa NA USE
Sentaa XLL Seleee 43 SOOR
MSC
Replay IL Replay
10
15
Timestep
Figure 24: Examples from MineRL "Treechop-vO" task. (A) Long-horizon success; (B) long-horizon
failure; (C) One-step success; (D) One-step failure.
C.1 Wallelock time
Figure 28 shows the average time required to recover a transition for each method in the one-step (top) and long-horizon (bottom) cases. As expected. the time required to perform the search scales linearly in the number of stored transitions. During the test, we occupied a maximum of 12GB of RAM, including VAE, SSMs, rollout and replay buffer.
Clearly, computing the KI divergence requires significantly more time than computing L.2 distance. However, we highlight that for long-term predictions we only require to compute the distances once every N steps, where N the horizon. Additionally, we inform that in our implementation used the torch. distribution butions package to ensure correctness, regardless of efficiency. The
25

---

NOS
D
WCS
MSC
D
Replay Replay
10
15
Timestep
Figure 25: Examples from MineRL "Navigate-vO" task. (A) Long-horizon success; (B) long-horizon
failure; (C) One-step success; (D) One-step failure.
package requires two Distribution objects to compute the KL distance between them exactly, hence making the computation quite expensive. However, if the type of distribution is set, efficiency can be improved by computing a closed-fone form solution and leveraging PyTorch parallel computation. remark that the test has been done in CPU, since all models could not fit in the VRAM of our GPU. However, when using only a single method, the computation can be parallelized in GPU to vastly improve speed.
Additionally,
Finally, we point out how encoding almost half a million transitions represents quite an extreme case needed only in open-ended tasks: in SuperTuxKart, for example, each dataset is composed of roughly 10k transitions; similarly, in Atari we store around 20k transitions per task. Therefore, we expect the time gap between methods to be quite limited in practice.
26

---

A
PT
S0H
B
Raplay NO Rapllyntt
poodaay VCL Ruplay 13
Replay ICL suplayy
10
15
Timestep
Figure 26: Examples from Atari "Space Invaders" games. (A) Long-horizon success; (B) long-
horizon failure; (C) One-step success; (D) One-step failure.
To account for this expected practical scenario, we computed the inference time (dynamics evolution & action selection) for each model, by using typical number of encoded trajectories for tasks in SuperTuxKart. As the inference time only depends on the number of encoded transitions, the results hold for any task using a comparable number of encodings of similar size. We show the results in Table g
In the "best-case scenario" of a small scoped task with a reasonable amount of encoded trajectories, PlaNet inference time is significantly higher than our methods. However, one might argue that in this case, performance could justify the gap if, for example, our approaches were to be comparable to random policy. We address this point in Appendix ?? using data from the same experiment used to collect the numbers in Table 3
27

---

DOCS
A
B
Repley Ruplayy
10
15
Timestep
Figure 27: Examples from Atari "Seaquest" game. (A) Long-horizon success; (B) long-horizon
failure; (C) One-step success; (D) One-step failure.
C.2 Visual results
In Figure 29 we report example of reconstruction for 10, 20, 40 and 80 encoded trajectories. The image suggests that encoding more images slightly improves visual resemblance our method, mostly thanks to the more complete coverage of the state space. However, reconstruction is not always perfect. Interestingly, despite improving its latent prediction skills as reported in Section 5 the baseline cannot produce convincing decoded sequences from its predictions.
28

---

1.0 SSM
Rollout
Replay-L2
Raplay-KL
0.5
0.0
1.0
0.75
60000
120000
180000
240000
300000
360000
420000
480000
encoded transitions
Figure 28: Average time needed for one-step (top) and long-horizon (bottom) latent dynamics prediction.
Table 3: Average inference time for each model in SuperTuxKart tasks. Results are consistent in all other environments, when using a similar number encoded trajectories.
Method Inference time (μ o) (S)
PlaNet
0.04448 + 0.00392
Rollout
0.00248 0.00021
L2
0.00084 + 0.00007
0.00188 0.00044
D Action conditioning visual examples
report some visual examples of the effects of action conditioning SuperTuxkee TuxKart. In each Figure, the first row shows the real sequence, while the three pairs of sequences show (top to bottom) Rollout, Replay-L2 and Replay-KL KL with (first row of a pair) and without (second row of a pair) action conditioning.
In Figures 30
33 and
for each track of Super
As reported in main text, conditioning on the action generally makes the prediction task harder, especially for Rollout and Replay-L.2, as both models need to sample batch on transitions to estimate the distribution. On the contrary, Replay-KL seems generally unaffected. We hypothesize that this the result matching the distribution rather than estimating it.
Moreover, in Figure 35 we report an example of retrieval failure for Replay-KL, as discussed in Section 5 is notable how retrieving a significantly different first latent leads to a completely different predicted sequence. Fortunately, such occurrences are very rare, as discussed in Section
E Algorithm & implementation details
In this Appendix, we report the architectures we have used for our study, along with some details and hyperparameters, to facilitate reproducibiliity of our results. working implementation our code wttprrroo__oottooooddddoeeeeeee
provided at hittt:::s... //tp::/ggtttttttt http:///////ttttooooooo/aaaaaaaa wtw.hhbb..oo//aaalaatteeeeoooooe
29

---

A
Repteyoks Parpleyta
B
Racstyya Sepley
Tan
D Faplay-ki XL Rapilys22 Redioo
10
15
Timestep
Figure 29: Visual examples of reconstructions from MineRL "Trrechop-vO" using (A) ten, (B) twenty,
(C) forty, and (D) eighty trajectories.
E.1 VAE
Table shows the hyperparameters used to train the VAEs. We train the models using the 3-VAE variant, which includes warm-up procedure for the KL divergence term. Specifically, the KL divergence part of the loss is masked in the first 10% the epochs, then it is gradually increased during the next 80% of training time. The architecture of the encoder and decoder follows [8]
For MineRL, we have increased the latent size due to higher amount of relevant information of an image. Additionally, we have decreased the number of epochs to 50, as empirically the model showed signs of convergence around that time. Finally, we have increased the learning rate from 5 x 1 to 3 10-1, as it empirically led to the best results.
30

---

No en Na Raa
R
KI no act
Timestep
Figure 30: Effects of action conditioning on dynamics prediction in "fortmagma" track. First row reports the real sequence; after that, each pair of rows reports our three methods, respectively with and without action conditioning.
Nea
No act Ratoa
no act.
Timestep
Figure 31: Effects of action conditioning on dynamics prediction Synes_ainbowroale track. First row reports the real sequence; after that, each pair of rows reports our three methods, respectively with and without action conditioning.
E.2 SSM
We report the hyperparameters used to train the SSM model in Table The training procedure follows [8) From the original paper, changed the values of the hyperparameters to obtain the best results, according to our empirical tests. Our PyTorch implementation uses abhayrav1/planet- torch as reference, even though the source code of the repo has been used only as a guide. We re-implemented the model as shown in the source code linked to this study.
com/
31

---

Nalbut
No OCT
n0 KC.
Timestep
Figure 32: Effects action conditioning on dynamics prediction "volcano island" track. First row reports the real sequence; after that, each pair of rows reports three methods, respectively with and without action conditioning.
hallus
No oct.
no BCT.
KL
XI art
10
15
Timestep
Figure 33; Effects of action conditioning dynamics prediction in "snowmountain" track. First row reports the real sequence; after that, each pair of rows reports our three methods, respectively with and without action conditioning.
Table 4: Hyperparameters for VAE models used in this study.
Name
SuperTuxKart
MineRI.
Atari
image size
64x64x3
64x64x3
64x64x3
latent size
128
512
128
learning rate
10-5
3 10-4
5 10-5
epochs batch size beta beta interval (epochs)
250
50
250
128
128
128
0.0 + 5 x 10
0.0 + 5 x 10
0.0 + 5 x 10 -
25 225
5 45
25 225
32

---

Nallou
No KCT.
122 no ICC
no aCC
Timestep
Figure 34: Effects of action conditioning on dynamics prediction "lighthouse" track. First row reports the real sequence; after that, each pair of rows reports our three methods, respectively with and without action conditioning.
A
RL no
10
15
Timestep
B
no aEE.
15
Timestep
Figure 35: A comparison between (A) a successful retrieval and (B) failed search by Replay-Ke KL. Both sequences are extracted from the "fortmagma" track.
Table 5: Hyperparameters used train the PlaNet baseline.
Name
Super TuxKat
MineRI
Atari
latent size
128
512
128
hidden size
256
256
256
learning rate
10
10
1 10
epochs
250
250
250
batch size
64
64
64
beta
0.1
0.1
0.1
33


Reviews:
[{"review_id": "JYrWm1RX1U", "reviewer": "Reviewer_4hop", "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.\"", "questions": "Please see Strengths and Weaknesses above.", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\n- 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.\n- Comparison experiments with PlaNet showed comparable quality of latent reconstruction and reconstructed images.\n- The paper is well-written and easy to follow.\n\nWeaknesses:\n- 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.\n- 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.\n- (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.\n\nIn all, if the author can thoroughly address my concern, I am very willing to raise my score.", "was_revised": true, "final_justification": "I have carefully read the author's rebuttal. The reviewer greatly appreciates the additional experiments and explanations regarding novelty provided by the author. I will raise my score.", "rebuttal": "We thank the Reviewer for their valuable suggestions and gladly answer the interesting points they raised. We address each question in a separate item, following the order specified by the Reviewer. At present time in the rebuttal process, it is not possible for us to share an updated version of the manuscript. However, changes to the manuscript following from the comments are marked for reference of the Reviewer. Please find our answers below.\n\n- We agree that the images are not extremely clear and may appear blurry; however, we highlight that the images shown in the manuscript are intended for visual evaluation purposes only. Additionally, most of the images are VAE reconstructions of the predicted dynamics, which are notoriously affected by a small reconstruction error causing the blur. Regarding the informative content of the images, we agree with the Reviewer that greater resolution could help in preserving details and hence add meaningful information for a model. However, we point out that our experimental protocol follows common practices found in relevant previous literature, such as using 84x84 [1], or 64x64 [2, 3] images for learning to act & plan with partial pixel observations.\n\n- We agree with the Reviewer that comparing the models for action selection and planning would constitute a valuable addition to our manuscript. Therefore, we have designed an additional experiment to test the action models. The experiment is designed as follows: each model is tested on the realized returns over 50 evaluation episodes. Each model is allowed to observe the first state of an episode. From this single observation, agents are asked to evolve the dynamics according to their world model and plan the next 20 actions. Then, the planned actions are executed with no room for corrections. Finally, the model is allowed to observe a new state and plan the next sequence of actions. The process is repeated until the end of the episode. We report the results in Table 1 (below) using both _mean_ $\\pm$ _std_ and [95% confidence interval]. Additionally, we have reported the same results in our manuscript in Appendix F, Table 5. We highlight that the test was run only on the SuperTuxKart environments due to time restrictions. The results show that zero-shot world models consistently outperform the PlaNet baseline, even though no specific approach shows consistently superior performance. These new results support our conclusion that zero-shot world models can be a valid alternative to learning-based ones.\n\n- We thank the Reviewer for providing such an interesting point of discussion. Despite the vast gap between the two ideas, we believe that a diffusion model would indeed achieve superior performance to our methods. However, it is worth noting that, with respect to our proposal _that requires no training at all_, a diffusion model would need a significant amount of data and time to learn the dynamics of each environment. With our study, we aim at showing that \"lightweight\" alternatives to computationally heavy models might exist. Therefore, despite recognizing the validity of the proposal advanced by the Reviewer, we believe that a comparison between the two approaches would not be fair due to the vast gap in computational resources needed.\n\n### Table 1: Mean episodic return in SuperTuxKart environments. Each run is composed by 50 evaluation episodes.\n\n| Environment        | Model       | μ ± σ           | [CI 95%]           |\n|-------------------|-------------|------------------|--------------------|\n| **fortmagma**      | PlaNet      | 0.088 ± 0.113     | [0.056, 0.120]     |\n|                   | **Rollout**     | **5.810 ± 3.476**   | **[4.812, 6.808]** |\n|                   | Replay-L2   | 4.350 ± 1.451     | [3.934, 4.766]     |\n|                   | Replay-KL   | 4.774 ± 3.811     | [3.680, 5.868]     |\n| **lighthouse**     | PlaNet      | 3.078 ± 1.516     | [2.643, 3.513]     |\n|                   | Rollout     | 2.640 ± 2.595     | [1.895, 3.385]     |\n|                   | **Replay-L2**   | **5.114 ± 9.555**   | **[2.371, 7.857]** |\n|                   | Replay-KL   | 0.320 ± 0.756     | [0.103, 0.537]     |\n| **snes_rainbowroad** | PlaNet      | 3.160 ± 2.968     | [2.308, 4.012]     |\n|                   | Rollout     | 4.822 ± 1.119     | [4.501, 5.143]     |\n|                   | **Replay-L2**   | **6.336 ± 5.201**   | **[4.843, 7.829]** |\n|                   | Replay-KL   | 4.714 ± 0.652     | [4.527, 4.901]     |\n| **snowmountain**   | PlaNet      | 5.322 ± 3.788     | [4.235, 6.409]     |\n|                   | Rollout     | 5.556 ± 6.570     | [3.670, 7.442]     |\n|                   | Replay-L2   | 4.956 ± 6.504     | [3.089, 6.823]     |\n|                   | **Replay-KL**   | **15.664 ± 11.825** | **[12.269, 19.059]** |\n| **volcano_island** | PlaNet      | 0.928 ± 0.644     | [0.743, 1.113]     |\n|                   | Rollout     | 1.530 ± 0.585     | [1.362, 1.698]     |\n|                   | Replay-L2   | 3.348 ± 6.129     | [1.589, 5.108]     |\n|                   | **Replay-KL**   | **4.478 ± 7.965**   | **[2.191, 6.765]** |\n\n## References:\n[1] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, Martin A. Riedmiller, _Playing Atari with Deep Reinforcement Learning._ CoRR abs/1312.5602 (2013)\n\n[2] Danijar Hafner, Timothy P. Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, James Davidson, _Learning Latent Dynamics for Planning from Pixels._ ICML 2019: 2555-2565\n\n[3] Danijar Hafner, Timothy P. Lillicrap, Mohammad Norouzi, Jimmy Ba, _Mastering Atari with Discrete World Models._ ICLR 2021"}, {"review_id": "ijbsOIzHer", "reviewer": "Reviewer_JymM", "summary": "This 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.", "questions": "I 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.", "limitations": "The authors don't note their assumption of image reconstruction as a sufficient proxy for world models.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 1, "clarity": 3, "significance": 1, "originality": 3, "strengths_and_weaknesses": "**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.", "was_revised": true, "final_justification": "Given the new results and other changes the authors have provided I think this is an interesting direction of work that has been well-motivated and effectively evaluated.", "rebuttal": "We thank the Reviewer for their insightful analysis and provide discussion points to the “Weaknesses” listed by them, following the same order. At present time in the rebuttal process, it is not possible for us to share an updated version of the manuscript. However, changes to the manuscript following from the comments are marked for reference of the Reviewer. Please find our answers below.\n\n-  We recognize that the text was previously unclear about the limitations of our method. To address this, we updated **Section 6 (lines 317–321)** to better articulate both the strengths and shortcomings. Below is the revised text for the Reviewer's reference:\n\n  **Original text:**\n  > *Our method introduces three main benefits over the baseline: first, it lower the prediction error on latent dynamics, and either matches or lowers the reconstruction error for decoded trajectories; second, it improves consistency and visual appeal in long-horizon prediction; third, requires seconds to be implemented. Conversely, our method may suffer from hallucinations and lack of generalization when paired with insufficient data, as shown in the Minecraft benchmark. Additionally, the results obtained with our models may significantly vary, depending on the strategy used for retrieval.*\n\n  **Revised text:**\n  > *Our methods outperform the baseline on latent prediction and reconstruction errors, while also being more immediate in implementation and application. Despite showing acceptable performance in open-ended tasks such as Minecraft and being comparable to the PlaNet baseline, our results suggest that the performance of zero-shot WMs depend on latent space coverage of the encoded trajectories, limiting their effectiveness in tasks with vast state spaces. However, they represent a valid, lightweight alternative to regular WMs in small-scoped tasks, such as SuperTuxKart and Atari.*\n\n\n- We thank the Reviewer for raising this point. We would like to highlight that tasks from Minecraft are substantially unpredictable due to the highly diverse observations that an agent could encounter. These tasks have been selected precisely to test the limits of similarity search on vast state spaces and have been tested coherently with the other selected benchmarks. Examples of this additional difficulty can be seen by comparing the values in Figures 2, 3, and 4 from our manuscript: in general, error values in Figure 3 (related to Minecraft tasks) are significantly higher than in the other two benchmarks, suggesting that, as expected, all models struggle more in predicting dynamics from Minecraft. However, we would appreciate it if the Reviewer could provide some examples of environments deemed to be suitable for this point.\n\n- We thank the Reviewer for providing this comment. To answer, we would like to highlight that our methods are evaluated both in the _visual domain_ through L1 and SSIM reconstruction error and in _latent space_ through KL divergence *(lines 203 - 206)*, which account for the *dynamics prediction error*. Moreover, we specify that L1 and SSIM are necessarily affected by decoding errors and, as such, should be considered in conjunction with KL and the visual examples *(lines 207 - 210)*.\n\n- We agree that testing the action selection performance would be beneficial to our manuscript. Therefore, we have added an experiments that compares the models during planning. The test is designed as follows: on the first timestep of an episode, each model observes the environment. Then, based solely on this observation, models are asked to evolve the dynamics using their world model and plan 20 actions accordingly. The actions are then executed as predicted, giving models no chance to change their plan. Finally,  We report the results in Table 1 (below) as _mean_ $\\pm$ _std_, and use a 95% confidence interval as additional metric. We add the same results in our manuscript in Appendix F (Table 5). We highlight that the test was run only on the SuperTuxKart environments due to time restrictions. The results show that action selected by our models are, on average, either comparable or better than the one chosen by the selected baseline. For reference of the Reviewer, these results were obtained by encoding 10 trajectories for each task (time taken: ~15 seconds) for our methods, as opposed to training a PlaNet baseline for 25 epochs (time taken: ~5 hours). All methods were tested using the same hardware. In Table 2 (below) we also report the average inference time for each model, as recorded during the same test.\n\n- We believe that the additional experiment highlithed in the previous point may be sufficient to cover this discussion point as well.\n\n### Table 1: Mean episodic return in SuperTuxKart environments. Each run is composed by 50 evaluation episodes.\n\n| Environment        | Model       | μ ± σ           | [CI 95%]           |\n|-------------------|-------------|------------------|--------------------|\n| **fortmagma**      | PlaNet      | 0.088 ± 0.113     | [0.056, 0.120]     |\n|                   | **Rollout**     | **5.810 ± 3.476**   | **[4.812, 6.808]** |\n|                   | Replay-L2   | 4.350 ± 1.451     | [3.934, 4.766]     |\n|                   | Replay-KL   | 4.774 ± 3.811     | [3.680, 5.868]     |\n| **lighthouse**     | PlaNet      | 3.078 ± 1.516     | [2.643, 3.513]     |\n|                   | Rollout     | 2.640 ± 2.595     | [1.895, 3.385]     |\n|                   | **Replay-L2**   | **5.114 ± 9.555**   | **[2.371, 7.857]** |\n|                   | Replay-KL   | 0.320 ± 0.756     | [0.103, 0.537]     |\n| **snes_rainbowroad** | PlaNet      | 3.160 ± 2.968     | [2.308, 4.012]     |\n|                   | Rollout     | 4.822 ± 1.119     | [4.501, 5.143]     |\n|                   | **Replay-L2**   | **6.336 ± 5.201**   | **[4.843, 7.829]** |\n|                   | Replay-KL   | 4.714 ± 0.652     | [4.527, 4.901]     |\n| **snowmountain**   | PlaNet      | 5.322 ± 3.788     | [4.235, 6.409]     |\n|                   | Rollout     | 5.556 ± 6.570     | [3.670, 7.442]     |\n|                   | Replay-L2   | 4.956 ± 6.504     | [3.089, 6.823]     |\n|                   | **Replay-KL**   | **15.664 ± 11.825** | **[12.269, 19.059]** |\n| **volcano_island** | PlaNet      | 0.928 ± 0.644     | [0.743, 1.113]     |\n|                   | Rollout     | 1.530 ± 0.585     | [1.362, 1.698]     |\n|                   | Replay-L2   | 3.348 ± 6.129     | [1.589, 5.108]     |\n|                   | **Replay-KL**   | **4.478 ± 7.965**   | **[2.191, 6.765]** |\n\n### Table 2: Average inference time in SuperTuxKart environments.\n\n| Method  | Inference time ($\\mu \\pm \\sigma$) (s)       |\n|---------|----------------------------------|\n| KL      | $0.00188 \\pm 0.00044$           |\n| L2      | $0.00084 \\pm 0.00007$           |\n| Rollout | $0.00248 \\pm 0.00021$           |\n| PlaNet  | $0.04448 \\pm 0.00392$           |"}, {"review_id": "O0lqsDZepu", "reviewer": "Reviewer_zx8U", "summary": "In 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", "questions": "Is 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?", "limitations": "Not 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.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**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\".", "was_revised": true, "final_justification": "The authors have addressed issues around writing and presentation clarity and also on computational complexity of the method effectively. They have shown that world models can be leveraged zero shot through search methods to model environment dynamics and that the approach can also help with action selection.  I am supportive overall of the method which utilises world models in a novel fashion to to good effect and I am therefore happy to see this work published.", "rebuttal": "We thank the Reviewer for their valuable comments. We will address each question in a separate item for clear readability, in the order provided by the Reviewer themselves. At present time in the rebuttal process, it is not possible for us to share an updated version of the manuscript. However, changes to the manuscript following from the comments are marked for reference of the Reviewer. Please find our answers below.\n\n- Computational complexity is discussed in Appendix C. We highlight that in the Appendix we evaluate how retrieval time scales with the number of encoded transitions (and, therefore, trajectories). However, we point out that this represents an unrealistic use case for our method, as according to our conclusions (Section 6, lines 318 - 321 & **Revised text** in item 2 of this answer), our method is best suited for small-scoped tasks, which will typically require a fraction of those transitions. To this end, we ran an additional experiment addressing action selection performance (added to Appendix F, Table 5 and shown in Table 1 below) on the tracks of the SuperTuxKart benchmark, in which we also computed the inference time for each agent (Table 2, below). We believe this test to be a much more interesting and likely scenario for our method. These results have been added also in Appendix C of our manuscript (Table 2). Furthermore, as a reference for the Reviewer, in this scenario the \"training time\" for our proposals consist of the time needed to encode trajectories, which is ~15 seconds; conversely, training the PlaNet baseline on the same hardware to obtain the results in Table 1 took ~5 hours for each environment.\n\n- We appreciate the comment and agree that a visual example showing the differences in reconstructing the dynamics w.r.t. the number of encoded trajectories would improve our manuscript. We have added a visualization of this fact in Appendix C of our updated version of the manuscript. To recover it while maintaining the visualization readable, we set the model to what we believe should be its standard use, namely long-term prediction without action conditioning. However, as correctly pointed out by the Reviewer, according to Figure 6 \"Replay-KL\" shows the most robust performance in predicting the evolution of the task dynamics over time. Relative to this, we highlight that the dynamics prediction does not necessarily correlate with better performance in planning, as shown in Table 1 (below). By testing the planning capabilities of our proposed methods, we show that while their performance is either on par or superior to the PlaNet baseline, no specific method dominates over the other across all tested tasks. \n\n- We agree with the Reviewer that the third point of our contribution could have been discussed more in detail. To this end, we have expanded our discussion on it in **Section 6, lines 318 - 321**. For reference of the Reviewer, please find our updated text below:\n\n  **Original text:**\n  > *Our method introduces three main benefits over the baseline: first, it lower the prediction error on latent dynamics, and either matches or lowers the reconstruction error for decoded trajectories; second, it improves consistency and visual appeal in long-horizon prediction; third, requires seconds to be implemented. Conversely, our method may suffer from hallucinations and lack of generalization when paired with insufficient data, as shown in the Minecraft benchmark. Additionally, the results obtained with our models may significantly vary, depending on the strategy used for retrieval.*\n\n  **Revised text:**\n  > *Our methods outperform the baseline on latent prediction and reconstruction errors, while also being more immediate in implementation and application. Despite showing acceptable performance in open-ended tasks such as Minecraft and being comparable to the PlaNet baseline, our results suggest that the performance of zero-shot WMs depend on latent space coverage of the encoded trajectories, limiting their effectiveness in tasks with vast state spaces. However, they represent a valid, lightweight alternative to regular WMs in small-scoped tasks, such as SuperTuxKart and Atari.*\n\nIn addition to the questions posed by the Reviewer, we would like to briefly address the points listed under the \"Weaknesses\" category. Also in this instance we will answer following the order proposed by the Reviewer.\n\n- A Figure highlighting the differences among our proposed models has not been included due to page limits. However, we agree with the Reviewer that it would greatly help in conveying our proposal. As such, we are working on a solution to include one either in main text, or in the Appendices, while abiding by the submission limits.\n\n- A discussion on computational burden at inference time has been included above.\n\n- We thank the Reviewer for pointing out this shortcoming, and agree that the sentence is cryptic w.r.t. the relevance of the method. The sentence implicitly refers to Section 3, implying that after retrieving the 1-most similar transition from each trajectory, we leverage the minibatch of next observations {$o_{t+1,i}$}$_{i=1}^{|\\mathcal{T}|}$ to estimate the latent distribution of the next sample, namely the next predicted step. For reference of the Reviewer, please find the updated text below:\n\n**Original text:**\n  > *Then, we extract the next state and estimate the next latent distribution as previously discussed.*\n\n  **Revised text:**\n  > *Then, we form the minibatch of next observations extracted by the current retrieval process. Finally, we estimate the next latent distribution by consideting each point in the minibatch as a realization of it, as discussed in Section 3.*\n\n- Some details regarding the training setup, e.g. hyperparameters, are provided in Appendix C, Tables 3 & 4. All the relevant architectural details are provided in the referenced papers. However, we agree with the Reviewer that providing the most relevant details in our Appendix could benefit to our manuscript, and are working on including them in Appendix C.\n\n- Typo fixed, we thank the Reviewer for spotting it.\n \n### Table 1: Mean episodic return in SuperTuxKart environments. Each run is composed by 50 evaluation episodes.\n\n| Environment        | Model       | μ ± σ           | [CI 95%]           |\n|-------------------|-------------|------------------|--------------------|\n| **fortmagma**      | PlaNet      | 0.088 ± 0.113     | [0.056, 0.120]     |\n|                   | **Rollout**     | **5.810 ± 3.476**   | **[4.812, 6.808]** |\n|                   | Replay-L2   | 4.350 ± 1.451     | [3.934, 4.766]     |\n|                   | Replay-KL   | 4.774 ± 3.811     | [3.680, 5.868]     |\n| **lighthouse**     | PlaNet      | 3.078 ± 1.516     | [2.643, 3.513]     |\n|                   | Rollout     | 2.640 ± 2.595     | [1.895, 3.385]     |\n|                   | **Replay-L2**   | **5.114 ± 9.555**   | **[2.371, 7.857]** |\n|                   | Replay-KL   | 0.320 ± 0.756     | [0.103, 0.537]     |\n| **snes_rainbowroad** | PlaNet      | 3.160 ± 2.968     | [2.308, 4.012]     |\n|                   | Rollout     | 4.822 ± 1.119     | [4.501, 5.143]     |\n|                   | **Replay-L2**   | **6.336 ± 5.201**   | **[4.843, 7.829]** |\n|                   | Replay-KL   | 4.714 ± 0.652     | [4.527, 4.901]     |\n| **snowmountain**   | PlaNet      | 5.322 ± 3.788     | [4.235, 6.409]     |\n|                   | Rollout     | 5.556 ± 6.570     | [3.670, 7.442]     |\n|                   | Replay-L2   | 4.956 ± 6.504     | [3.089, 6.823]     |\n|                   | **Replay-KL**   | **15.664 ± 11.825** | **[12.269, 19.059]** |\n| **volcano_island** | PlaNet      | 0.928 ± 0.644     | [0.743, 1.113]     |\n|                   | Rollout     | 1.530 ± 0.585     | [1.362, 1.698]     |\n|                   | Replay-L2   | 3.348 ± 6.129     | [1.589, 5.108]     |\n|                   | **Replay-KL**   | **4.478 ± 7.965**   | **[2.191, 6.765]** |\n\n### Table 2: Average inference time in SuperTuxKart environments.\n\n| Method  | Inference time ($\\mu \\pm \\sigma$) (s)       |\n|---------|----------------------------------|\n| KL      | $0.00188 \\pm 0.00044$           |\n| L2      | $0.00084 \\pm 0.00007$           |\n| Rollout | $0.00248 \\pm 0.00021$           |\n| PlaNet  | $0.04448 \\pm 0.00392$           |"}, {"review_id": "jq4Bm8u3CC", "reviewer": "Reviewer_AMmz", "summary": "This 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.", "questions": "1. **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.", "limitations": "Authors 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).", "rating": 4, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**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.", "was_revised": true, "final_justification": "Thank you to the authors for providing a thorough and constructive response to my initial review comments. After careful consideration of the additional empirical data and analysis you provided, I have decided to adjust my recommendation from \"Reject\" to \"Borderline Reject.\"\n\nWhile the new empirical analysis is valuable, there remains a lack of deeper theoretical explanation for why similarity search effectively approximates dynamics. For example, there is no analysis of how errors accumulate with prediction horizon length or under what conditions this method would fail.\n\nWhile comparison with PlaNet is reasonable, the lack of comparison with more recent world model approaches (such as DreamerV3) limits the generalizability of the results.", "rebuttal": "We thank the Reviewer for their insightful analysis of our manuscript and for the interesting questions raised. We provide the answers below, using an itemized list that follows the same order. At present time in the rebuttal process, it is not possible for us to share an updated version of the manuscript. However, changes to the manuscript following from the comments are marked for reference of the Reviewer. Please find our answers below.\n\n- Despite not having theoretical guarantees on the fact that retrieved local transitions approximate global dynamics, some reasonable assumptions used in our work could justify this fact. Firstly, the retrieval process implies that the error between an observation query $o$ and the corresponding retrieved observations $o'$ is minimized; assuming that such error is small enough, we could conclude that local transitions observed by a model approximate local transitions stored in the trajectories dataset $\\mathcal{T}$. In second instance, considering that the encoded trajectories are collected from human gameplay, they will by definition implicitly follow the global dynamics of the environment; this is a main difference between our model and Dreamer-like architectures such as PlaNet, where the global dynamics are explicitly tracked by the hidden state of the RSSM. Finally, since the task of predicting dynamics with zero-shot world models can be reduced to compounding local transitions (that follow the Markov property) to estimate a generative distribution in a continuous latent space, we can reasonably assume that the retrieval process in its entirety follows the Markov property. Regarding the empirical validation of the property, we have not validated it. \n\n- We thank the Reviewer for highlighting this important point of discussion. To address it, we have computed the latent space binning coverage ratio using our test trajectories, and compared it to the KL error obtained by each model. Additionally, we computed the Pearson correlation coefficient for each model. As correctly stated by the Reviewer, after comparing the coverage statistics for SuperTuxKart and Minecraft, we found that lower coverage corresponds to higher mean error in transition retrieval. We report our results in Tables 1 & 2 (below), reporting respectively the dynamics prediction error (KL divergence) compared to the state space coverage ratio for each environment and model, and the Pearson correlation coefficient for the two quantities, for each model. Additionally, we have added the same results in Appendix A of our manuscript in visual form. Notably, all tested models show a strong correlation between data diversity and generalization capabilities. However, the correlation is particularly strong in PlaNet, accordingly with the well-known fact that neural networks greatly benefit from a more diverse dataset. Overall, our proposed models are slightly less dependent on data diversity.\n\n- We thank the Reviewer for raising this valid point. To answer this point we have performed an additional experiment, which we have also added to Appendix F of our manuscript. Please find the results in the Table 3 below. To account for unknown actions, we have compared the planning capabilities of each model. In particular, we let each model observe one frame every 20 steps and plan a sequence of actions up to the planning horizon; the actions are selected only based on the predicted dynamics. The results show that all our models yield comparable or superior performance to the PlaNet baseline, while requiring drastically less time to be applied (~15 seconds as opposed to ~5 hours for each environment, using same hardware). Furthermore, regarding the joint optimization of action and dynamics, we highlight that the only optimization procedure in our method consists of an independent training for the VAE encoder, which we operate following [1]. Aside from this, _none of our proposed methods requires any optimization_. The relevant processing needed to prepare our model for predicting dynamics is discussed in **Section 3** of our manuscript. Similarly, action selection and planning is discussed in **Section 3 (lines 117 - 120)**.\n\n- We thank the Reviewer for addressing this concern. We report the retrieval time for each model in Appendix C of our manuscript, where we test the feature under extreme conditions (Minecraft, 50k -> 500k encoded transitions). However, as addressed in Section 6 of our manuscript, we believet that our approach would be best suited for tasks with a more limited scope. As such, we conduct additional tests on the inference time of each model, that is, the time required to map an observation into an action. We have added these results in Appendix C, and we show them in Table 4 below. Finally, we agree with the Reviewer that claiming our method to be “lightweight” should be limited to such smaller scoped tasks, and have therefore updated our manuscript accordingly (**Section 6, lines 318 - 321**). For reference of the Reviewer, please find the added text below:\n\n**Added text:**\n> *Despite showing acceptable performance in open-ended tasks such as Minecraft and being comparable to the PlaNet baseline, our results suggest that the performance of zero-shot WMs depend on latent space coverage of the encoded trajectories, limiting their effectiveness in tasks with vast state spaces. However, they represent a valid, lightweight alternative to regular WMs in small-scoped tasks, such as SuperTuxKart and Atari.*\n\n\n### Table 1: State space coverage ratio and dynamics prediction error (KL divergence) for the tested models on environments from SuperTuxKart and Minecraft.\n\n| Environment      | State Space Coverage Ratio |     PlaNet     | Rollout   | Replay-L2 | Replay-KL  |\n|------------------------|----------------------|-------------|-----------|-----------|------------|\n|TuxKart-lighthouse| 0.7675               |  571.3605   | 173.8899  | 192.4496  | 231.5324   |\n|TuxKart-fortmagma| 0.7300               |  519.0421   |  68.6900  |  86.9360  | 136.4911   |\n|TuxKart-rainbowroad| 0.8475               |  588.6913   | 146.5262  | 148.1211  | 521.2641   |\n|TuxKart-snowmountain| 0.6750               |  569.5750   |  67.0347  |  76.2962  | 133.6698   |\n|TuxKart-volcano_island| 0.5925               |  499.0704   |  85.8624  |  94.0208  | 217.8786   |\n|Minecraft-Treechop| 0.3950               | 2612.7190   | 394.5654  | 407.0560  | 1065.6871  |\n|Minecraft-Navigate| 0.6175               | 1156.5831   | 208.7204  | 196.5164  | 578.7786   |\n\n### Table 2: Pearson correlation index of each model across the tested environments.\n| Model     | Pearson Correlation Coefficient |\n|-----------|---------------------|\n| SSM       | -0.8269             |\n| Rollout   | -0.6709             |\n| Replay-L2 | -0.6711             |\n| Replay-KL | -0.6366             |\n\n### Table 3: Mean episodic return in SuperTuxKart environments. Each run is composed by 50 evaluation episodes.\n\n| Environment        | Model       | μ ± σ           | [CI 95%]           |\n|-------------------|-------------|------------------|--------------------|\n| **fortmagma**      | PlaNet      | 0.088 ± 0.113     | [0.056, 0.120]     |\n|                   | **Rollout**     | **5.810 ± 3.476**   | **[4.812, 6.808]** |\n|                   | Replay-L2   | 4.350 ± 1.451     | [3.934, 4.766]     |\n|                   | Replay-KL   | 4.774 ± 3.811     | [3.680, 5.868]     |\n| **lighthouse**     | PlaNet      | 3.078 ± 1.516     | [2.643, 3.513]     |\n|                   | Rollout     | 2.640 ± 2.595     | [1.895, 3.385]     |\n|                   | **Replay-L2**   | **5.114 ± 9.555**   | **[2.371, 7.857]** |\n|                   | Replay-KL   | 0.320 ± 0.756     | [0.103, 0.537]     |\n| **snes_rainbowroad** | PlaNet      | 3.160 ± 2.968     | [2.308, 4.012]     |\n|                   | Rollout     | 4.822 ± 1.119     | [4.501, 5.143]     |\n|                   | **Replay-L2**   | **6.336 ± 5.201**   | **[4.843, 7.829]** |\n|                   | Replay-KL   | 4.714 ± 0.652     | [4.527, 4.901]     |\n| **snowmountain**   | PlaNet      | 5.322 ± 3.788     | [4.235, 6.409]     |\n|                   | Rollout     | 5.556 ± 6.570     | [3.670, 7.442]     |\n|                   | Replay-L2   | 4.956 ± 6.504     | [3.089, 6.823]     |\n|                   | **Replay-KL**   | **15.664 ± 11.825** | **[12.269, 19.059]** |\n| **volcano_island** | PlaNet      | 0.928 ± 0.644     | [0.743, 1.113]     |\n|                   | Rollout     | 1.530 ± 0.585     | [1.362, 1.698]     |\n|                   | Replay-L2   | 3.348 ± 6.129     | [1.589, 5.108]     |\n|                   | **Replay-KL**   | **4.478 ± 7.965**   | **[2.191, 6.765]** |\n\n### Table 4: Average inference time in SuperTuxKart environments.\n\n| Method  | Inference time ($\\mu \\pm \\sigma$) (s)       |\n|---------|----------------------------------|\n| KL      | $0.00188 \\pm 0.00044$           |\n| L2      | $0.00084 \\pm 0.00007$           |\n| Rollout | $0.00248 \\pm 0.00021$           |\n| PlaNet  | $0.04448 \\pm 0.00392$           |\n\n## References:\n[1] Diederik P Kingma, Max Welling, _Auto-Encoding Variational Bayes_, ICLR 2014."}]

Author rebuttal:

Meta-review: This paper proposes a zero-shot world model that leverages similarity search over a latent memory buffer and stochastic representations via a pretrained VAE to predict environment dynamics without explicit training. The approach stores trajectories and retrieves similar transitions to estimate future states, enabling long-horizon predictions in diverse environments such as SuperTuxKart, Minecraft, and Atari. 

The paper is practical and well-written, offering a lightweight alternative to traditional data-hungry world models. Strengths include its simplicity, empirical validation across multiple domains, and a clear description of memory structures and experiments. Weaknesses include limited theoretical analysis, reliance on known future actions, sensitivity to dataset quality, and incomplete evaluation of downstream control tasks. Reviewers also suggested clarifying algorithmic steps, memory buffer composition, and minor implementation details.

During rebuttal, the authors addressed key concerns by providing additional planning experiments, clarifying limitations, and expanding methodological discussions. These updates demonstrate that the method is usable as a world model and highlight future research directions. Considering the strong response to reviewer concerns, I recommend accept (Consensus Reached), with the requirement that the final version includes the promised clarifications, supplementary experiments, and minor corrections.
Decision: Accept (poster)