A brief history of learning in imagination

Learning in imagination, a model-based reinforcement learning approach that trains policies solely on data generated by an action-conditioned world model, addresses sample efficiency by generating data for free, as exemplified by the Dyna paper (Sutton, 1991). The method contrasts with data-hungry deep learning models like AlexNet, which relied on the ImageNet dataset, and clarifies that world models are not merely video generators but model environment dynamics and rewards.

If there’s anything to learn from the ImageNet/AlexNet revolution, it’s that neural networks need tons and tons of data to work. Without the ImageNet dataset, the AlexNet model never would have been possible, and deep learning would forever be condemned to a fun fact section in an artificial intelligence textbook. This makes sample efficiency such an interesting problem. How do you learn well with limited data? Data is expensive to collect, in both the time and money sense. Coming up with a satisfying answer to the question would thus be a critical unlock to anyone building a production ML system. Learning in imagination is one of the classes of methods largely motivated by sample efficiency. It is a model-based reinforcement learning approach, which means that it has two main components†: Action-conditioned world model – to simulate the environment. Policy – the algorithm that interacts with the environment. The characteristic feature of learning in imagination is that the policy is trained solely on data generated by the world model. In other words, the policy learns in the imagination of the world model. From this, the obvious benefit of learning in imagination is that you don't need to collect huge amounts of data to train your agent; data is generated basically for free by your world model†. Analogies, metaphors, and allegories Imagine you’re a volleyball player. You’re on the court with 5 teammates facing off against another team of 6 on the other side of the net. During each rally, a complex series of stimuli and responses takes place almost instantly. The other team’s setter has set the ball to their outside hitter. Your blockers are getting ready to block a straight hit, so you in the back row are getting ready to dig a cross hit. The hitter hits that cross, you watch the ball rushing towards your right, and you crouch down, angle your arms, and position your platform so that the ball bounces back smoothly towards your setter for the next play. All of the things you just did, all of your actions , are done in response to some observation of your environment. In RL speak, you as the agent acted according to your policy , which maps states to probabilities of selecting each possible action†. Whenever you take an action which includes doing nothing , the environment reacts accordingly. For example, if you hadn’t positioned yourself for the cross hit and instead hid behind your blockers, the ball would’ve fallen straight to the ground and the other team would’ve gained a point. In this case, the environment gives you the next state of “ball on ground” and some feedback that tells you that you did the wrong thing. Again in RL speak, this feedback is called the reward , and its value for this case can be set to a negative value like -1. Modeling the environment thus means learning a mapping from states and actions to probabilities of next states and rewards. This mapping is your world model †. It’s worth doing a double take on the definition of a world model. In popular culture at least, popular in tech circles , the term “world model” is used to refer to something akin to a video generator, something that is purely visual and is optimized for photorealism or aesthetics. This definition isn’t quite right. A world model is actually much richer than that. It tells the agent how well it’s doing, and it’s also action-conditioned rather than only state-conditioned – in other words, the output changes based on the actions you take, not just the current state of the environment. In the Dyna https://doi.org/10.1145/122344.122377 paper Sutton, 1991 , this is equivalently called an action model . And it’s not really about pixel-perfect reconstruction, rather it’s about modeling the physics/dynamics in high fidelity; the exact pattern on the volleyball is less important than the trajectory of the ball after a serve or a hit. In recent times, we train world models on image/video reconstruction, since we assume that dynamics are apparent in visuals, so one can pardon the confusion between aesthetics and details for the layperson. However, this distinction is definitely something to be aware of and understand. The common misconception of what a world model is may also leave some confused about how a reward, which is an abstract idea, can be outputted by a world model, which is mistakenly assumed to produce just video. Don’t worry, this isn’t black magic All of this is just math and code. As a shorthand, you can think of a world model as a function that returns a 2-tuple, where one value is the pixels or latent representation of the next frame, and the other value is a scalar reward. So really, the two return values are separate. Okay, so how is the policy connected to the world model? An RL problem is a sequence of interactions between the agent and the environment. The classic agent-environment interaction diagram from the introduction to RL book by Sutton and Barto https://web.stanford.edu/class/psych209/Readings/SuttonBartoIPRLBook2ndEd.pdf illustrates this very concisely. Now imagine a futuristic world where you can practice volleyball in a simulator, kind of like how Lightning McQueen attempted to practice using the racing simulator in Cars 3. To build a simulator, you would need to collect actual footage of the circuit, or the volleyball court, so that the simulator can model the real environment with accurate physics, reward signals, etc. However, when the policy is learning – when you’re practicing – the observations consumed by the policy all come from the simulator, not the real environment. This is the setup for learning in imagination: you, the agent, learn from the simulated environment rather than the real one. So, if I were to create a diagram for learning in imagination, it would look something like this: Once again, the policy learns in the imagination of the world model. Data from the real environment is only used to train the world model itself, and the policy never sees this real data during training. In most existing work, the policy and the world model are being trained at the same time, circle of life style†. To do this, the agent has to send actions to the real environment so that we can collect the data necessary to train the world model. Keep in mind that this is different from the agent consuming data from the real environment, which is what learning in imagination avoids. A short detour first… Learning in imagination is a type of model-based RL, which implies that there are model-based RL methods that are not learning in imagination. And indeed, there are: , where there is no learned model the model is hard-coded . AlphaZero and MuZero , where real observations are used to train the model, which in turn predicts the policy. EfficientZero , where there is no learned policy, rather planning with Cross-Entropy Method CEM is used. PlaNet Knowing what is and isn’t something I find to be very useful. Reading up on the aforementioned works – which are quite well-known in the field, by the way – is left as an exercise for the reader. Things to optimize for in learning in imagination When your policy is trained on the generations of your world model, your policy can only be as good as your world model. If your simulated game of volleyball has the ball floating upwards instead of falling downwards, or if you get positively rewarded whenever the other team scores, you’ll never learn how to play the game the right way. This brings us to thing to optimize 1, the correctness of the world model . Earlier on, I talked about confusing optimizing details with optimizing aesthetics. The details of the states of the environment are important to model correctly, including the dynamics and the rewards. In Minecraft for example, if your world model hallucinates a diamond ore overground, your policy cannot actually learn to find actual diamonds in caves correctly; however, it is less important that the blue of the imagined diamond is the same as the blue of the real diamond. As you can imagine, imperfections in world model generations can yield policies that are suboptimal, or even policies that learn to exploit these imperfections to “reward hack” i.e. being a smart ss , both of which mean that the agent behaves sub-optimally in real settings. Furthermore, while the world model can interpolate between training data points to output counterfactual scenarios, it cannot interpolate between data points it has not yet seen from the real environment. Thing to optimize 2 is thus the scope of the world model . This particular thing is a challenge for environments with rare events, like encountering a woodland mansion in Minecraft. And this expectedly has downstream impacts on the policy; if we never discover a woodland mansion in the real environment, our world model can never be trained to simulate it, which means that the policy can never know how to navigate a woodland mansion nor how to fight evokers and vexes, let alone discover the lapis lazuli secret. This is an agent exploration and data problem. Speaking of data, thing to optimize 3 brings us back to sample efficiency . Your world model can generate unlimited data for your policy, but your world model doesn’t get the same treatment from the real environment†. Because of this, sample efficiency doesn’t come for free with learning in imagination. You have to first improve the sample efficiency of your world model training before you can claim that your learning in imagination method is indeed sample efficient. A brief academic history of learning in imagination Thanks for sticking with me throughout that long prologue. Now that that’s out of the way, it’s time for a short history of how we got here, where learning in imagination comes from, and some of the foundational papers / people that created it. . For easy reference, here are the papers I include in this article: Dyna https://doi.org/10.1145/122344.122377 Sutton, 1991 World Models https://arxiv.org/abs/1803.10122 Ha & Schmidhuber, 2018 SimPLe https://arxiv.org/abs/1903.00374 Kaiser et al., 2020 IRIS https://arxiv.org/abs/2209.00588 Micheli et al., 2023 DIAMOND https://arxiv.org/abs/2405.12399 Alonso et al., 2024 Dreamer 4 https://arxiv.org/abs/2509.24527 Hafner et al., 2025 There are obviously more papers, like the entire Dreamer series for example, but as before, reading these works is left as an exercise for the reader. Pioneers If I were to pick two papers to form the old guard of learning in imagination, they would be Dyna https://doi.org/10.1145/122344.122377 Sutton, 1991 and World Models https://arxiv.org/abs/1803.10122 Ha & Schmidhuber, 2018 . Let’s talk about the elder statesman first. Dyna https://doi.org/10.1145/122344.122377 predates the deep learning revolution brought about by ImageNet and AlexNet, so it doesn’t talk about image reconstruction as an objective for the world model or equivalently, action model . However, the idea for training a policy on “hypothetical world states” can be traced back to this paper. The classical components of learning in imagination – dynamics learning, behavior learning, and environment interaction – are introduced here. It’s really quite fitting that learning in imagination was proposed by the Father of RL himself, Rich Sutton. World Models https://arxiv.org/abs/1803.10122 the paper also proposes training a policy on the outputs of a world model, but since the ImageNet/AlexNet moment had passed, the world model is a neural network trained to do image reconstruction, bringing Dyna’s ideas into more complex environments. The architecture in World Models is quite familiar. You have a variational autoencoder VAE that encodes frames into latent representations, and a recurrent neural network RNN that predicts latent representations of future frames. These two components form the world model. Then, you have the action-taking controller model that takes as input these latent representations. This is a very simple policy network, consisting of only one linear layer that maps the vectors from the VAE and RNN directly to action at each time step. For our purposes, the architecture isn’t really the interesting part. What’s interesting is what the authors realized they could do given the setup. As mentioned earlier, the RNN is essentially tasked to predict the future. Given this, the authors thought maybe, just maybe, that they can get rid of the VAE altogether; in essence, doing away with the real environment observations entirely. Instead, they would sample a future latent representation from the RNN and send it straight to the policy. At first they tried this for inference, where an already-trained policy plays in the imagination of the world model. It worked. They then took this idea one step further and actually trained a policy purely in hallucination†. It also worked With these two pioneering papers – incredibly, more than 25 years apart from each other – we have the stage set for modern learning in imagination methods. World model-ers The next group of papers – SimPLe https://arxiv.org/abs/1903.00374 Kaiser et al., 2020 , IRIS https://arxiv.org/abs/2209.00588 Micheli et al., 2023 , and DIAMOND https://arxiv.org/abs/2405.12399 Alonso et al., 2024 – focus on advancing world modeling techniques done in conjunction with policy learning to advance learning in imagination. The emphasis on sample efficiency takes center stage in this era of papers, by virtue of their evaluation method if nothing else. A big common thread in these papers is the lopsided focus on innovating the world modeling half more than the policy learning half. Case in point, SimPLe simply haha uses proximal policy optimization PPO https://arxiv.org/abs/1707.06347 , which is one of, if not the most famous RL algorithm out there. You can tell that the authors of these papers are much more opinionated regarding the world modeling technique but much less so for the policy learning approach when you read lines like “many actor-critic methods could be employed to train the policy and the value network in imagination.” That quote is lifted verbatim from the IRIS paper, which along with DIAMOND uses the REINFORCE https://www.geeksforgeeks.org/machine-learning/reinforce-algorithm/ objective for policy learning, again a very standard RL objective. The astute would notice that across the three papers, the RL algorithm used – PPO and REINFORCE – is model-free. However, the overall method, given the presence of a world model, is considered model-based. I forget where I heard this advice from, but someone once told me that the easier your method is to implement and reproduce, the better and more citable it is. I think that’s good advice, and interestingly SimPLe, IRIS, and DIAMOND are in a way optimizing for it. In both the IRIS and DIAMOND papers, the authors write that since large amounts of trajectories can be generated in imagination, the RL objective used can be standard straightforward methods e.g. PPO, REINFORCE . In other words, the authors embrace the lack of necessity to reinvent the wheel, which is afforded by the “data is free” property of learning in imagination, which is in turn a design decision by the authors. How cool is that? Let’s now prove a point by walking through the different world models that SimPLe, IRIS, and DIAMOND built. SimPLe https://arxiv.org/abs/1903.00374 , the oldest of the trio, uses a convolutional encoder and decoder shaped like a U-Net https://en.wikipedia.org/wiki/U-Net , with subtle implementation differences. It makes sense that the U-Net shape is employed for this reconstruction task, since U-Net was originally created for image segmentation , where the final output is essentially another image with segmentation masks; contrast this with image classification , where the output is just class scores/probabilities, which have much fewer dimensions than a whole image. The SimPLe paper is also the paper to introduce the now-widely used Atari 100k benchmark. This benchmark evaluates model performance on Atari games under the condition of using only 100k interactions with the real environment, which translates to approximately only two hours of play time. While SimPLe can achieve a predetermined threshold score using fewer steps compared to older methods, these older methods push past SimPLe’s scores when trained on more data e.g. on 200k interactions . The chart below shows how SimPLe’s performance changes as the number of training samples it sees increases. The x-axis shows the number of samples SimPLe sees, and the y-axis shows the number of samples that PPO needs to match the performance of SimPLe. The number of samples that correspond to the point of intersection between the line and the gray diagonal y=x is the number of samples at which SimPLe’s performance equals PPO’s performance. Ideally, if you were SimPLe’s authors, you’d like to see that the line always lies above or at the y=x diagonal, since that indicates that there can’t be a matched sample count that leads to PPO outperforming SimPLe, at least for the number of samples within the tested region up to 1M samples . However, the chart shows a plateau, which means that while SimPLe is great in low data regimes, it is outperformed in high data regimes. The SimPLe authors hypothesized that building a better world model would enable future learning in imagination methods to also be competitive in higher data regimes. I believe this is a reasonable hypothesis, because the difference between SimPLe and PPO is basically only the fact that SimPLe uses the generations of a learned world model for policy training, which uses the PPO algorithm. Those generations are obviously bottlenecked by the quality of the world model. Enter IRIS https://arxiv.org/abs/2209.00588 . This is a special paper to me because this is the paper that introduced me to learning in imagination. Thanks to LLMs, transformers are all the rage nowadays. IRIS uses a transformer as the backbone of its world model, though it precedes LLMs entering the public zeitgeist. The world model in IRIS consists of a discrete autoencoder and an autoregressive transformer. The autoencoder is trained to reconstruct a 2D image frame . Like in the World Models paper, we intercept the encoder and decoder to get the latent tokens of the frame, which are then fed into the transformer. Given previous latent tokens and previous actions taken by the policy, the transformer autoregressively predicts future latent tokens, the reward, and a boolean value for episode termination. As can be seen from the diagram above, the encoder E is used only for the first frame. This is because the policy takes as input the reconstructed frames from the world model rather than the real frames from the environment once the world model has been seeded with an initial frame. On the Atari 100k benchmark, IRIS achieves new high scores and superhuman performance in 10 out of 26 games, beating SimPLe and model-free methods like CURL https://arxiv.org/pdf/2004.04136 and SPR https://arxiv.org/pdf/2007.05929 , and even model-based methods with lookahead search like MuZero https://www.nature.com/articles/s41586-020-03051-4.epdf?sharing token=kTk-xTZpQOF8Ym8nTQK6EdRgN0jAjWel9jnR3ZoTv0PMSWGj38iNIyNOw ooNp2BvzZ4nIcedo7GEXD7UmLqb0M V fop31mMY9VBBLNmGbm0K9jETKkZnJ9SgJ8Rwhp3ySvLuTcUr888puIYbngQ0fiMf45ZGDAQ7fUI66-u7Y%3D . While EfficientZero https://arxiv.org/pdf/2111.00210 is still generally stronger, it being a lookahead search method is computationally expensive, requiring parallel CPU and GPU threads while IRIS only needs a single GPU / single CPU implementation for inference. And what about performance in high data regimes? In the appendix, the authors report results for 10M samples, and they have great news to share. The performance drastically improved for most games, proving that IRIS can be scaled up to high data regimes without plateauing early. There is still room for improvement, however†. After all, ideally these methods should achieve superhuman performance on all 26 games even when given limited data DIAMOND https://arxiv.org/abs/2405.12399 claims that one way to do this is to better model visual details. The central claim in this paper is that improved visual details in world models can lead to improved agent performance. The authors argue that discrete latent variables, which up to that point had been the prevailing environment representation in world models, are lossy and do not preserve the minute but necessary visual elements integral for good policy learning. Instead, they turned to diffusion , which had increasingly been becoming the go-to approach for high-quality image generation. The DIAMOND paper thus trains a U-Net under a novel diffusion objective that makes diffusion-based world modeling actually viable. The visualization shows us two things. First, the imagination procedure is autoregressive in frames and actions, where a stack of T 1 past frames and actions inform the predicted next frame. Second, the diffusion component does not predict rewards; the world model has this diffusion component as depicted, as well as a separate reward and termination model. Again on the Atari 100k benchmark, DIAMOND achieves even better high scores and superhuman performance on 11 games, with numerical scores surpassing those of IRIS and DreamerV3 https://arxiv.org/abs/2301.04104 another learning in imagination method in most of the 26 games. In fact, it achieves a new best mean human normalized score among agents trained entirely within a world model. Perhaps modeling pixelated Atari games is not the most flashy showcase out there. The authors challenged themselves to model a more complex game to really prove their point about better modeling visual details. They showed that DIAMOND’s world model can simulate CS:GO In this case, only the world model is trained and there is no notion of policy, since the dataset used only contains frames. It’s of course not perfect e.g. hallucinating double jumps but the fact that it can model 3D environments with gameplay logic is impressive. This provides early evidence that learning in imagination can graduate from Atari games to more complex, higher fidelity environments. We’ll get back to this idea of a neural game engine in a bit. Latent space-er The final paper on our list, Dreamer 4 https://arxiv.org/abs/2509.24527 Hafner et al., 2025 , has a different approach from the previous three papers we discussed. Dreamer 4 https://arxiv.org/abs/2509.24527 is unique in that it is an offline RL algorithm, which means that it absolutely does not interact with the real environment even for training the world model, rather it uses a pre-collected dataset of observations. Furthermore, we’ve frequently seen the world model being “separated” from the policy, in that the policy receives decoded frames in pixel space from the world model. Conversely, Dreamer 4 has the model and policy share the same backbone, so the policy receives latents from the world model. An oversimplification of the architecture looks something like this†: The authors also included a diagram of their architecture in the paper, mainly focused on the frame tokenizer and interactive dynamics model world model . So how exactly does Dreamer 4 fuse the world model and the agent? At a high level, training was done in three phases: World model pretraining . The tokenizer is trained to reconstruct videos in the offline dataset. At the same time, the world model is trained on tokenized videos under a shortcut forcing objective, which is introduced in this paper. Both the tokenizer and world model are 2D transformers, equipped with time and space dimensions. The world model is therefore trained to do action-conditioned video prediction, predicting clean representations of frames from corrupted representations and actions and discrete signal levels and step sizes . Agent fine-tuning . The policy and reward heads, which are just small multilayer perceptrons MLPs , are trained to predict actions and rewards given task embeddings in the input stream. The world model is still fine-tuned using the video prediction loss, and the data is still from the offline dataset. Imagination training . The policy and value heads are trained with RL – the policy head via the PMPO https://arxiv.org/abs/2410.04166v1 objective and the value head via temporal difference learning – on imagined rollouts in order to maximize the learned reward model. In this phase, the world model is frozen. These imagined rollouts are generated by sampling latent representations from the world model and actions from the policy head. Since the policy head is already trained after agent fine-tuning, the model is already able to act after that phase. So why do we need imagination training? A useful way to think about this is to contrast imitation learning against reinforcement learning. The agent fine-tuning phase uses actions, observations, and rewards from the offline dataset, which means that the agent is learning to imitate existing behavior. This is because the training process dictates that the agent minimizes the loss against these pre-collected data points, so its peak performance is one that clones the existing behavior already exhibited in the dataset. In contrast, the imagination training phase uses reinforcement learning, which enables the agent to explore and exploit its simulated environment on its own. Therefore, the agent can learn behaviors outside of the offline dataset, such that its peak performance can potentially exceed that that is displayed in the dataset. Despite offline RL being a challenging paradigm, Dreamer 4 manages to perform well under that constraint. In fact, it is the first agent trained exclusively on offline data to successfully collect diamonds in Minecraft, using 100 times less data than OpenAI’s VPT https://arxiv.org/pdf/2206.11795 offline agent I’d like to think that this is because, in a way, imagination training enables an online training experience, since the policy and value heads leverage rollouts from an environment, albeit one that is simulated, that can also include counterfactual scenarios. It’s a neat cheat code. If you and I think alike, I think that you think that Dreamer 4’s architecture is quite reminiscent of the one in World Models. After all, the architecture in World Models has a world model component that passes latents directly to a small policy/controller head. In general, the Dreamer family v1, 2, 3, and 4 follows this “latent space” paradigm. Closing thoughts I hope I’ve convinced you that learning in imagination is cool, and I also hope that you’re excited to learn more about it on your own. There’s no denying that the prospect of infinite data to train your agents is appealing. I also promised you that we’ll go back to that idea of simulating a game engine. Very recently, General Intuition and Kyutai, in collaboration with Epic Games, published a multiplayer world model that simulates Rocket League called MIRA https://arxiv.org/pdf/2607.05352 . What’s exciting about this particular model is that it treats all agents as their own action streams, rather than treating only “Player 1” as an agent and other players as part of the environment. Plus, it runs at 20 FPS†, which makes it responsive enough to feel real-time I’ve tried it out†, it’s very impressive . You would’ve noticed that all of the learning in imagination methods I’ve shared in this article are single-player. Since architectures are designed with the data in mind, this is in part due to data constraints; most environments and datasets out there are for single agents only. These methods enable solving partially observable Markov decision processes POMDP which by construction are single-agent. However, there may be scenarios where you’d want to train multiple agents simultaneously, either cooperatively or competitively, like in a game of volleyball. To me, progress on multiplayer action-conditioned world models like MIRA means progress towards solving partially observable stochastic games POSG , the general multiagent form of POMDPs. POSGs are difficult to solve, even intractable in large settings since the number of actions and observations grow exponentially with the number of agents, but multiplayer world models enable modeling these multiagent environments to solve at least approximations of POSGs via learning in imagination. We’d then go into the realm of game theory, which is its own fun can of worms. Next up In the spirit of good science and engineering, I’ll be implementing a learning in imagination method myself, and share with you the nitty gritty details in my next post. Stay tuned :