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 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 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:
- AlphaZeroMastering Chess and Shogi by Self-Play with a General Reinforcement Learning AlgorithmThe game of chess is the most widely-studied domain in the history of artificial intelligence. The strongest programs are based on a combination of sophisticated search techniques, domain-specific adaptations, and handcrafted evaluation functions that have been refined by human experts over several decades. In contrast, the AlphaGo Zero program recently achieved superhuman performance in the game of Go, by tabula rasa reinforcement learning from games of self-play. In this paper, we generalise this approach into a single AlphaZero algorithm that can achieve, tabula rasa, superhuman performance in many challenging domains. Starting from random play, and given no domain knowledge except the game rules, AlphaZero achieved within 24 hours a superhuman level of play in the games of chess and shogi (Japanese chess) as well as Go, and convincingly defeated a world-champion program in each case.arXiv:1712.01815v1View paper, where there is no learned model (the model is hard-coded).
- MuZero and EfficientZeroMastering Atari Games with Limited DataReinforcement learning has achieved great success in many applications. However, sample efficiency remains a key challenge, with prominent methods requiring millions (or even billions) of environment steps to train. Recently, there has been significant progress in sample efficient image-based RL algorithms; however, consistent human-level performance on the Atari game benchmark remains an elusive goal. We propose a sample efficient model-based visual RL algorithm built on MuZero, which we name EfficientZero. Our method achieves 194.3% mean human performance and 109.0% median performance on the Atari 100k benchmark with only two hours of real-time game experience and outperforms the state SAC in some tasks on the DMControl 100k benchmark. This is the first time an algorithm achieves super-human performance on Atari games with such little data. EfficientZero's performance is also close to DQN's performance at 200 million frames while we consume 500 times less data. EfficientZero's low sample complexity and high performance can bring RL closer to real-world applicability. We implement our algorithm in an easy-to-understand manner and it is available at https://github.com/YeWR/EfficientZero. We hope it will accelerate the research of MCTS-based RL algorithms in the wider community.arXiv:2111.00210v2View paper, where real observations are used to train the model, which in turn predicts the policy.
- PlaNetLearning Latent Dynamics for Planning from PixelsPlanning has been very successful for control tasks with known environment dynamics. To leverage planning in unknown environments, the agent needs to learn the dynamics from interactions with the world. However, learning dynamics models that are accurate enough for planning has been a long-standing challenge, especially in image-based domains. We propose the Deep Planning Network (PlaNet), a purely model-based agent that learns the environment dynamics from images and chooses actions through fast online planning in latent space. To achieve high performance, the dynamics model must accurately predict the rewards ahead for multiple time steps. We approach this using a latent dynamics model with both deterministic and stochastic transition components. Moreover, we propose a multi-step variational inference objective that we name latent overshooting. Using only pixel observations, our agent solves continuous control tasks with contact dynamics, partial observability, and sparse rewards, which exceed the difficulty of tasks that were previously solved by planning with learned models. PlaNet uses substantially fewer episodes and reaches final performance close to and sometimes higher than strong model-free algorithms.arXiv:1811.04551v5View paper, where there is no learned policy, rather planning with Cross-Entropy Method (CEM) is used.
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 (Sutton, 1991)
- World ModelsWorld ModelsWe explore building generative neural network models of popular reinforcement learning environments. Our world model can be trained quickly in an unsupervised manner to learn a compressed spatial and temporal representation of the environment. By using features extracted from the world model as inputs to an agent, we can train a very compact and simple policy that can solve the required task. We can even train our agent entirely inside of its own hallucinated dream generated by its world model, and transfer this policy back into the actual environment. An interactive version of this paper is available at https://worldmodels.github.io/arXiv:1803.10122v4View paper (Ha & Schmidhuber, 2018)
- SimPLeModel-Based Reinforcement Learning for AtariModel-free reinforcement learning (RL) can be used to learn effective policies for complex tasks, such as Atari games, even from image observations. However, this typically requires very large amounts of interaction -- substantially more, in fact, than a human would need to learn the same games. How can people learn so quickly? Part of the answer may be that people can learn how the game works and predict which actions will lead to desirable outcomes. In this paper, we explore how video prediction models can similarly enable agents to solve Atari games with fewer interactions than model-free methods. We describe Simulated Policy Learning (SimPLe), a complete model-based deep RL algorithm based on video prediction models and present a comparison of several model architectures, including a novel architecture that yields the best results in our setting. Our experiments evaluate SimPLe on a range of Atari games in low data regime of 100k interactions between the agent and the environment, which corresponds to two hours of real-time play. In most games SimPLe outperforms state-of-the-art model-free algorithms, in some games by over an order of magnitude.arXiv:1903.00374v5View paper (Kaiser et al., 2020)
- IRISTransformers are Sample-Efficient World ModelsDeep reinforcement learning agents are notoriously sample inefficient, which considerably limits their application to real-world problems. Recently, many model-based methods have been designed to address this issue, with learning in the imagination of a world model being one of the most prominent approaches. However, while virtually unlimited interaction with a simulated environment sounds appealing, the world model has to be accurate over extended periods of time. Motivated by the success of Transformers in sequence modeling tasks, we introduce IRIS, a data-efficient agent that learns in a world model composed of a discrete autoencoder and an autoregressive Transformer. With the equivalent of only two hours of gameplay in the Atari 100k benchmark, IRIS achieves a mean human normalized score of 1.046, and outperforms humans on 10 out of 26 games, setting a new state of the art for methods without lookahead search. To foster future research on Transformers and world models for sample-efficient reinforcement learning, we release our code and models at https://github.com/eloialonso/iris.arXiv:2209.00588v2View paper (Micheli et al., 2023)
- DIAMONDDiffusion for World Modeling: Visual Details Matter in AtariWorld models constitute a promising approach for training reinforcement learning agents in a safe and sample-efficient manner. Recent world models predominantly operate on sequences of discrete latent variables to model environment dynamics. However, this compression into a compact discrete representation may ignore visual details that are important for reinforcement learning. Concurrently, diffusion models have become a dominant approach for image generation, challenging well-established methods modeling discrete latents. Motivated by this paradigm shift, we introduce DIAMOND (DIffusion As a Model Of eNvironment Dreams), a reinforcement learning agent trained in a diffusion world model. We analyze the key design choices that are required to make diffusion suitable for world modeling, and demonstrate how improved visual details can lead to improved agent performance. DIAMOND achieves a mean human normalized score of 1.46 on the competitive Atari 100k benchmark; a new best for agents trained entirely within a world model. We further demonstrate that DIAMOND's diffusion world model can stand alone as an interactive neural game engine by training on static Counter-Strike: Global Offensive gameplay. To foster future research on diffusion for world modeling, we release our code, agents, videos and playable world models at https://diamond-wm.github.io.arXiv:2405.12399v2View paper (Alonso et al., 2024)
- Dreamer 4Training Agents Inside of Scalable World ModelsWorld models learn general knowledge from videos and simulate experience for training behaviors in imagination, offering a path towards intelligent agents. However, previous world models have been unable to accurately predict object interactions in complex environments. We introduce Dreamer 4, a scalable agent that learns to solve control tasks by reinforcement learning inside of a fast and accurate world model. In the complex video game Minecraft, the world model accurately predicts object interactions and game mechanics, outperforming previous world models by a large margin. The world model achieves real-time interactive inference on a single GPU through a shortcut forcing objective and an efficient transformer architecture. Moreover, the world model learns general action conditioning from only a small amount of data, allowing it to extract the majority of its knowledge from diverse unlabeled videos. We propose the challenge of obtaining diamonds in Minecraft from only offline data, aligning with practical applications such as robotics where learning from environment interaction can be unsafe and slow. This task requires choosing sequences of over 20,000 mouse and keyboard actions from raw pixels. By learning behaviors in imagination, Dreamer 4 is the first agent to obtain diamonds in Minecraft purely from offline data, without environment interaction. Our work provides a scalable recipe for imagination training, marking a step towards intelligent agents.arXiv:2509.24527v1View paper (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 (Sutton, 1991) and World ModelsWorld ModelsWe explore building generative neural network models of popular reinforcement learning environments. Our world model can be trained quickly in an unsupervised manner to learn a compressed spatial and temporal representation of the environment. By using features extracted from the world model as inputs to an agent, we can train a very compact and simple policy that can solve the required task. We can even train our agent entirely inside of its own hallucinated dream generated by its world model, and transfer this policy back into the actual environment. An interactive version of this paper is available at https://worldmodels.github.io/arXiv:1803.10122v4View paper (Ha & Schmidhuber, 2018). Let’s talk about the elder statesman first.
Dyna 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 ModelsWorld ModelsWe explore building generative neural network models of popular reinforcement learning environments. Our world model can be trained quickly in an unsupervised manner to learn a compressed spatial and temporal representation of the environment. By using features extracted from the world model as inputs to an agent, we can train a very compact and simple policy that can solve the required task. We can even train our agent entirely inside of its own hallucinated dream generated by its world model, and transfer this policy back into the actual environment. An interactive version of this paper is available at https://worldmodels.github.io/arXiv:1803.10122v4View paper (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 – SimPLeModel-Based Reinforcement Learning for AtariModel-free reinforcement learning (RL) can be used to learn effective policies for complex tasks, such as Atari games, even from image observations. However, this typically requires very large amounts of interaction -- substantially more, in fact, than a human would need to learn the same games. How can people learn so quickly? Part of the answer may be that people can learn how the game works and predict which actions will lead to desirable outcomes. In this paper, we explore how video prediction models can similarly enable agents to solve Atari games with fewer interactions than model-free methods. We describe Simulated Policy Learning (SimPLe), a complete model-based deep RL algorithm based on video prediction models and present a comparison of several model architectures, including a novel architecture that yields the best results in our setting. Our experiments evaluate SimPLe on a range of Atari games in low data regime of 100k interactions between the agent and the environment, which corresponds to two hours of real-time play. In most games SimPLe outperforms state-of-the-art model-free algorithms, in some games by over an order of magnitude.arXiv:1903.00374v5View paper (Kaiser et al., 2020), IRISTransformers are Sample-Efficient World ModelsDeep reinforcement learning agents are notoriously sample inefficient, which considerably limits their application to real-world problems. Recently, many model-based methods have been designed to address this issue, with learning in the imagination of a world model being one of the most prominent approaches. However, while virtually unlimited interaction with a simulated environment sounds appealing, the world model has to be accurate over extended periods of time. Motivated by the success of Transformers in sequence modeling tasks, we introduce IRIS, a data-efficient agent that learns in a world model composed of a discrete autoencoder and an autoregressive Transformer. With the equivalent of only two hours of gameplay in the Atari 100k benchmark, IRIS achieves a mean human normalized score of 1.046, and outperforms humans on 10 out of 26 games, setting a new state of the art for methods without lookahead search. To foster future research on Transformers and world models for sample-efficient reinforcement learning, we release our code and models at https://github.com/eloialonso/iris.arXiv:2209.00588v2View paper (Micheli et al., 2023), and DIAMONDDiffusion for World Modeling: Visual Details Matter in AtariWorld models constitute a promising approach for training reinforcement learning agents in a safe and sample-efficient manner. Recent world models predominantly operate on sequences of discrete latent variables to model environment dynamics. However, this compression into a compact discrete representation may ignore visual details that are important for reinforcement learning. Concurrently, diffusion models have become a dominant approach for image generation, challenging well-established methods modeling discrete latents. Motivated by this paradigm shift, we introduce DIAMOND (DIffusion As a Model Of eNvironment Dreams), a reinforcement learning agent trained in a diffusion world model. We analyze the key design choices that are required to make diffusion suitable for world modeling, and demonstrate how improved visual details can lead to improved agent performance. DIAMOND achieves a mean human normalized score of 1.46 on the competitive Atari 100k benchmark; a new best for agents trained entirely within a world model. We further demonstrate that DIAMOND's diffusion world model can stand alone as an interactive neural game engine by training on static Counter-Strike: Global Offensive gameplay. To foster future research on diffusion for world modeling, we release our code, agents, videos and playable world models at https://diamond-wm.github.io.arXiv:2405.12399v2View paper (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)Proximal Policy Optimization AlgorithmsWe propose a new family of policy gradient methods for reinforcement learning, which alternate between sampling data through interaction with the environment, and optimizing a "surrogate" objective function using stochastic gradient ascent. Whereas standard policy gradient methods perform one gradient update per data sample, we propose a novel objective function that enables multiple epochs of minibatch updates. The new methods, which we call proximal policy optimization (PPO), have some of the benefits of trust region policy optimization (TRPO), but they are much simpler to implement, more general, and have better sample complexity (empirically). Our experiments test PPO on a collection of benchmark tasks, including simulated robotic locomotion and Atari game playing, and we show that PPO outperforms other online policy gradient methods, and overall strikes a favorable balance between sample complexity, simplicity, and wall-time.arXiv:1707.06347v2View paper, 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 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.
SimPLeModel-Based Reinforcement Learning for AtariModel-free reinforcement learning (RL) can be used to learn effective policies for complex tasks, such as Atari games, even from image observations. However, this typically requires very large amounts of interaction -- substantially more, in fact, than a human would need to learn the same games. How can people learn so quickly? Part of the answer may be that people can learn how the game works and predict which actions will lead to desirable outcomes. In this paper, we explore how video prediction models can similarly enable agents to solve Atari games with fewer interactions than model-free methods. We describe Simulated Policy Learning (SimPLe), a complete model-based deep RL algorithm based on video prediction models and present a comparison of several model architectures, including a novel architecture that yields the best results in our setting. Our experiments evaluate SimPLe on a range of Atari games in low data regime of 100k interactions between the agent and the environment, which corresponds to two hours of real-time play. In most games SimPLe outperforms state-of-the-art model-free algorithms, in some games by over an order of magnitude.arXiv:1903.00374v5View paper, the oldest of the trio, uses a convolutional encoder and decoder shaped like a 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 IRISTransformers are Sample-Efficient World ModelsDeep reinforcement learning agents are notoriously sample inefficient, which considerably limits their application to real-world problems. Recently, many model-based methods have been designed to address this issue, with learning in the imagination of a world model being one of the most prominent approaches. However, while virtually unlimited interaction with a simulated environment sounds appealing, the world model has to be accurate over extended periods of time. Motivated by the success of Transformers in sequence modeling tasks, we introduce IRIS, a data-efficient agent that learns in a world model composed of a discrete autoencoder and an autoregressive Transformer. With the equivalent of only two hours of gameplay in the Atari 100k benchmark, IRIS achieves a mean human normalized score of 1.046, and outperforms humans on 10 out of 26 games, setting a new state of the art for methods without lookahead search. To foster future research on Transformers and world models for sample-efficient reinforcement learning, we release our code and models at https://github.com/eloialonso/iris.arXiv:2209.00588v2View paper. 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 CURLCURL: Contrastive Unsupervised Representations for Reinforcement LearningWe present CURL: Contrastive Unsupervised Representations for Reinforcement Learning. CURL extracts high-level features from raw pixels using contrastive learning and performs off-policy control on top of the extracted features. CURL outperforms prior pixel-based methods, both model-based and model-free, on complex tasks in the DeepMind Control Suite and Atari Games showing 1.9x and 1.2x performance gains at the 100K environment and interaction steps benchmarks respectively. On the DeepMind Control Suite, CURL is the first image-based algorithm to nearly match the sample-efficiency of methods that use state-based features. Our code is open-sourced and available at https://github.com/MishaLaskin/curl.arXiv:2004.04136v4View paper and SPRData-Efficient Reinforcement Learning with Self-Predictive RepresentationsWhile deep reinforcement learning excels at solving tasks where large amounts of data can be collected through virtually unlimited interaction with the environment, learning from limited interaction remains a key challenge. We posit that an agent can learn more efficiently if we augment reward maximization with self-supervised objectives based on structure in its visual input and sequential interaction with the environment. Our method, Self-Predictive Representations(SPR), trains an agent to predict its own latent state representations multiple steps into the future. We compute target representations for future states using an encoder which is an exponential moving average of the agent's parameters and we make predictions using a learned transition model. On its own, this future prediction objective outperforms prior methods for sample-efficient deep RL from pixels. We further improve performance by adding data augmentation to the future prediction loss, which forces the agent's representations to be consistent across multiple views of an observation. Our full self-supervised objective, which combines future prediction and data augmentation, achieves a median human-normalized score of 0.415 on Atari in a setting limited to 100k steps of environment interaction, which represents a 55% relative improvement over the previous state-of-the-art. Notably, even in this limited data regime, SPR exceeds expert human scores on 7 out of 26 games. The code associated with this work is available at https://github.com/mila-iqia/sprarXiv:2007.05929v4View paper, and even model-based methods with lookahead search like MuZero. While EfficientZeroMastering Atari Games with Limited DataReinforcement learning has achieved great success in many applications. However, sample efficiency remains a key challenge, with prominent methods requiring millions (or even billions) of environment steps to train. Recently, there has been significant progress in sample efficient image-based RL algorithms; however, consistent human-level performance on the Atari game benchmark remains an elusive goal. We propose a sample efficient model-based visual RL algorithm built on MuZero, which we name EfficientZero. Our method achieves 194.3% mean human performance and 109.0% median performance on the Atari 100k benchmark with only two hours of real-time game experience and outperforms the state SAC in some tasks on the DMControl 100k benchmark. This is the first time an algorithm achieves super-human performance on Atari games with such little data. EfficientZero's performance is also close to DQN's performance at 200 million frames while we consume 500 times less data. EfficientZero's low sample complexity and high performance can bring RL closer to real-world applicability. We implement our algorithm in an easy-to-understand manner and it is available at https://github.com/YeWR/EfficientZero. We hope it will accelerate the research of MCTS-based RL algorithms in the wider community.arXiv:2111.00210v2View paper 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!
DIAMONDDiffusion for World Modeling: Visual Details Matter in AtariWorld models constitute a promising approach for training reinforcement learning agents in a safe and sample-efficient manner. Recent world models predominantly operate on sequences of discrete latent variables to model environment dynamics. However, this compression into a compact discrete representation may ignore visual details that are important for reinforcement learning. Concurrently, diffusion models have become a dominant approach for image generation, challenging well-established methods modeling discrete latents. Motivated by this paradigm shift, we introduce DIAMOND (DIffusion As a Model Of eNvironment Dreams), a reinforcement learning agent trained in a diffusion world model. We analyze the key design choices that are required to make diffusion suitable for world modeling, and demonstrate how improved visual details can lead to improved agent performance. DIAMOND achieves a mean human normalized score of 1.46 on the competitive Atari 100k benchmark; a new best for agents trained entirely within a world model. We further demonstrate that DIAMOND's diffusion world model can stand alone as an interactive neural game engine by training on static Counter-Strike: Global Offensive gameplay. To foster future research on diffusion for world modeling, we release our code, agents, videos and playable world models at https://diamond-wm.github.io.arXiv:2405.12399v2View paper 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 DreamerV3Mastering Diverse Domains through World ModelsDeveloping a general algorithm that learns to solve tasks across a wide range of applications has been a fundamental challenge in artificial intelligence. Although current reinforcement learning algorithms can be readily applied to tasks similar to what they have been developed for, configuring them for new application domains requires significant human expertise and experimentation. We present DreamerV3, a general algorithm that outperforms specialized methods across over 150 diverse tasks, with a single configuration. Dreamer learns a model of the environment and improves its behavior by imagining future scenarios. Robustness techniques based on normalization, balancing, and transformations enable stable learning across domains. Applied out of the box, Dreamer is the first algorithm to collect diamonds in Minecraft from scratch without human data or curricula. This achievement has been posed as a significant challenge in artificial intelligence that requires exploring farsighted strategies from pixels and sparse rewards in an open world. Our work allows solving challenging control problems without extensive experimentation, making reinforcement learning broadly applicable.arXiv:2301.04104v2View paper (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 4Training Agents Inside of Scalable World ModelsWorld models learn general knowledge from videos and simulate experience for training behaviors in imagination, offering a path towards intelligent agents. However, previous world models have been unable to accurately predict object interactions in complex environments. We introduce Dreamer 4, a scalable agent that learns to solve control tasks by reinforcement learning inside of a fast and accurate world model. In the complex video game Minecraft, the world model accurately predicts object interactions and game mechanics, outperforming previous world models by a large margin. The world model achieves real-time interactive inference on a single GPU through a shortcut forcing objective and an efficient transformer architecture. Moreover, the world model learns general action conditioning from only a small amount of data, allowing it to extract the majority of its knowledge from diverse unlabeled videos. We propose the challenge of obtaining diamonds in Minecraft from only offline data, aligning with practical applications such as robotics where learning from environment interaction can be unsafe and slow. This task requires choosing sequences of over 20,000 mouse and keyboard actions from raw pixels. By learning behaviors in imagination, Dreamer 4 is the first agent to obtain diamonds in Minecraft purely from offline data, without environment interaction. Our work provides a scalable recipe for imagination training, marking a step towards intelligent agents.arXiv:2509.24527v1View paper (Hafner et al., 2025), has a different approach from the previous three papers we discussed.
Dreamer 4Training Agents Inside of Scalable World ModelsWorld models learn general knowledge from videos and simulate experience for training behaviors in imagination, offering a path towards intelligent agents. However, previous world models have been unable to accurately predict object interactions in complex environments. We introduce Dreamer 4, a scalable agent that learns to solve control tasks by reinforcement learning inside of a fast and accurate world model. In the complex video game Minecraft, the world model accurately predicts object interactions and game mechanics, outperforming previous world models by a large margin. The world model achieves real-time interactive inference on a single GPU through a shortcut forcing objective and an efficient transformer architecture. Moreover, the world model learns general action conditioning from only a small amount of data, allowing it to extract the majority of its knowledge from diverse unlabeled videos. We propose the challenge of obtaining diamonds in Minecraft from only offline data, aligning with practical applications such as robotics where learning from environment interaction can be unsafe and slow. This task requires choosing sequences of over 20,000 mouse and keyboard actions from raw pixels. By learning behaviors in imagination, Dreamer 4 is the first agent to obtain diamonds in Minecraft purely from offline data, without environment interaction. Our work provides a scalable recipe for imagination training, marking a step towards intelligent agents.arXiv:2509.24527v1View paper 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 PMPOPreference Optimization as Probabilistic InferenceExisting preference optimization methods are mainly designed for directly learning from human feedback with the assumption that paired examples (preferred vs. dis-preferred) are available. In contrast, we propose a method that can leverage unpaired preferred or dis-preferred examples, and works even when only one type of feedback (positive or negative) is available. This flexibility allows us to apply it in scenarios with varying forms of feedback and models, including training generative language models based on human feedback as well as training policies for sequential decision-making problems, where learned (value) functions are available. Our approach builds upon the probabilistic framework introduced in (Dayan and Hinton, 1997), which proposes to use expectation-maximization (EM) to directly optimize the probability of preferred outcomes (as opposed to classic expected reward maximization). To obtain a practical algorithm, we identify and address a key limitation in current EM-based methods: when applied to preference optimization, they solely maximize the likelihood of preferred examples, while neglecting dis-preferred samples. We show how one can extend EM algorithms to explicitly incorporate dis-preferred outcomes, leading to a novel, theoretically grounded, preference optimization algorithm that offers an intuitive and versatile way to learn from both positive and negative feedback.arXiv:2410.04166v1View paper 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 VPTVideo PreTraining (VPT): Learning to Act by Watching Unlabeled Online VideosPretraining on noisy, internet-scale datasets has been heavily studied as a technique for training models with broad, general capabilities for text, images, and other modalities. However, for many sequential decision domains such as robotics, video games, and computer use, publicly available data does not contain the labels required to train behavioral priors in the same way. We extend the internet-scale pretraining paradigm to sequential decision domains through semi-supervised imitation learning wherein agents learn to act by watching online unlabeled videos. Specifically, we show that with a small amount of labeled data we can train an inverse dynamics model accurate enough to label a huge unlabeled source of online data -- here, online videos of people playing Minecraft -- from which we can then train a general behavioral prior. Despite using the native human interface (mouse and keyboard at 20Hz), we show that this behavioral prior has nontrivial zero-shot capabilities and that it can be fine-tuned, with both imitation learning and reinforcement learning, to hard-exploration tasks that are impossible to learn from scratch via reinforcement learning. For many tasks our models exhibit human-level performance, and we are the first to report computer agents that can craft diamond tools, which can take proficient humans upwards of 20 minutes (24,000 environment actions) of gameplay to accomplish.arXiv:2206.11795v1View paper 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 MIRAMultiplayer Interactive World Models with Representation AutoencodersWe introduce the first multiplayer world model for highly dynamic environments governed by complex physical interactions. Whereas single-player world models treat the other agents as part of the environment, ours conditions on the action streams of multiple agents, learning to attribute changes in the scene to the correct player and to stay coherent under arbitrary combinations of their actions. We study this problem in the game of Rocket League, where players compete and cooperate under fast, tightly coupled dynamics. Trained on 10,000 hours of gameplay collected with publicly available bots, our 5-billion-parameter latent diffusion model generates four-player matches in real time, producing 20 frames per second on a single Nvidia B200 GPU. Although trained only on short clips, its rollouts stay stable far beyond the training horizon: distributional quality holds steady out to five minutes, the longest horizon we measure, and in practice we observe rollouts continuing for hours with no sign of collapse. We systematically investigate the central design choices: the video codec, the generative objective, and the multiplayer conditioning scheme. In addition, we characterize how behavior changes with model and data scale, including the capabilities that emerge and the failure modes that persist. We further develop targeted evaluations that probe the model's physical understanding rather than visual appearance alone. To support continued research on multiplayer world models, we release our dataset, our full training and inference codebase, and a live demo.arXiv:2607.05352v2View paper. 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 :)
