World Models Explained: How AI Learns to Make Decisions in Its Dreams

What Are World Models?

Recently, an open-source project called worldmodels_demo sparked widespread attention on Reddit. Published by developer Gananath, this demo project has refocused attention on a topic in AI research that is both longstanding and cutting-edge — World Models.

The core idea of World Models originates from the classic 2018 paper World Models by David Ha and Jürgen Schmidhuber. David Ha is a research scientist at Google Brain (who later joined Stability AI as Head of Research), known for his creative deep learning research. Jürgen Schmidhuber, often called one of the "fathers of modern AI," is a co-inventor of LSTM and has long conducted research on recurrent neural networks and meta-learning at Switzerland's IDSIA laboratory. In this paper, they chose OpenAI's racing game environment CarRacing-v0 and VizDoom as experimental platforms, systematically demonstrating for the first time that an agent can complete policy training entirely within its own learned model of the environment. With its elegant architecture design and intuitive visualizations, the paper became a milestone at the intersection of reinforcement learning and generative models.

The fundamental idea is this: rather than having an agent learn through massive trial-and-error directly in a complex real environment, it's better to first have it build an "internal model" of how the world works, and then train and plan within this compressed, abstract imaginary space.

In other words, an agent can "learn in its dreams." Once it has a sufficiently good internal representation of environmental dynamics, it can even detach from the real environment entirely and complete extensive decision-making training within its own constructed simulated world. This approach shows enormous potential in both sample efficiency and generalization ability. A core bottleneck of traditional reinforcement learning is extremely low sample efficiency — for example, in DeepMind's Atari game experiments, the DQN algorithm typically requires tens of millions of frames of game interaction to learn a single game, equivalent to hundreds of hours of human gameplay. In real physical environments (such as robotic manipulation), each trial-and-error attempt means real time consumption and hardware wear, further amplifying this problem. By replacing real environment interactions with massive "virtual trial-and-error" in internal simulations, World Models can theoretically reduce the required number of real samples by one to two orders of magnitude.

Three Core Components: VAE, RNN, and Controller

The classic World Models architecture typically consists of three modules working in concert — and this is exactly the core structure reproduced by worldmodels_demo.

VAE: Compressed Perception

The first component is the Variational Autoencoder (VAE), responsible for visual perception. Each frame of high-dimensional image input from the environment (e.g., game screens) is encoded and compressed by the VAE into a low-dimensional latent vector.

Proposed by Diederik Kingma and Max Welling in 2013, the VAE is one of the foundational architectures in generative modeling. Unlike standard autoencoders, a VAE doesn't map inputs to a deterministic encoding point, but instead maps them to the mean and variance parameters of a probability distribution (typically Gaussian). During training, samples are drawn from this distribution using the "reparameterization trick," allowing backpropagation to pass through the stochastic sampling step. Its loss function consists of two parts — reconstruction error and KL divergence. The former ensures that the decoded output reconstructs the input as faithfully as possible, while the latter constrains the latent space distribution to approximate a standard normal distribution, giving the latent space continuity and interpolability. In World Models, the VAE typically compresses 64×64 pixel game frames into 32- or 64-dimensional latent vectors, achieving compression ratios of several hundred times.

This step is crucial — it condenses redundant, massive pixel information into compact abstract representations, enabling subsequent reasoning and planning to operate efficiently in a space far smaller than the original.

MDN-RNN: Predicting the Future

The second component is the Recurrent Neural Network, typically implemented as a Mixture Density Network combined with an RNN (MDN-RNN), playing the dual role of "memory" and "imagination."

The MDN-RNN combines two key technologies. RNNs (particularly the LSTM variant) excel at processing time-series data, capturing long-range dependencies by passing information through hidden states across time steps. Meanwhile, the Mixture Density Network (MDN), proposed by Christopher Bishop in 1994, doesn't output a single deterministic prediction like ordinary neural networks. Instead, it outputs the parameters of a mixture model composed of multiple Gaussian distributions (including the mean, variance, and mixing weight of each component). This design enables the model to express multimodal uncertainty about future states — for example, at an intersection, a vehicle might turn left or right, and these two possibilities need to be modeled by different Gaussian components. In World Models, the MDN-RNN typically uses an LSTM with 256 hidden units and 5 Gaussian mixture components to predict the distribution of the next frame's latent vector.

It receives the current latent state and the agent's action, then predicts the next latent state. This is essentially learning the dynamics of the environment: given the current state and a particular action, how will the future most likely unfold? It's this module that gives the agent the ability to "rehearse in its dreams."

Controller: Decision Output

The third component is the Controller, which typically has a very simple structure with only a small number of parameters. Based on the current perception provided by the VAE and the future predictions provided by the RNN, it outputs concrete action decisions.

Since the heavy lifting of representation learning and dynamics modeling is handled by the first two modules, the controller can remain lightweight and is often optimized using evolutionary strategies such as CMA-ES. CMA-ES (Covariance Matrix Adaptation Evolution Strategy), proposed by Nikolaus Hansen in 2001, is a gradient-free optimization algorithm. Unlike gradient descent-based methods, CMA-ES searches the parameter space by maintaining a multivariate normal distribution: each generation samples a batch of candidate solutions from the current distribution, evaluates their fitness (i.e., cumulative reward), and then updates the distribution's mean and covariance matrix based on the best-performing candidates. This approach is particularly suited for optimization problems with relatively few parameters (typically hundreds to thousands), which perfectly matches the controller's lightweight design. Choosing evolutionary strategies over backpropagation for optimizing the controller has another advantage: it naturally supports parallelization, allowing hundreds of candidate policies to be evaluated simultaneously without requiring differentiation of the reward function.

Why This Open-Source Demo Deserves Attention

For developers looking to get started with World Models, a runnable minimal demo project has significant practical value. Compared to cutting-edge research reproductions that often require massive computational resources, worldmodels_demo substantially lowers the barrier to understanding, allowing users to intuitively observe the complete pipeline of VAE reconstruction, RNN imagination sequences, and controller decision-making.

The World Models concept has once again become a research hotspot in the past two years. From DeepMind's Dreamer series of algorithms to various "neural network-generated interactive worlds" research (such as Google's Genie and various game screen generation models), they all share the same underlying logic as World Models: teaching AI to build predictable internal representations of the world.

The Dreamer series represents the most successful modern extension of the World Models concept. Dreamer V1 (2020) was the first to demonstrate that "imagined backpropagation" through a learned world model in continuous control tasks could match or even surpass model-free reinforcement learning methods. Dreamer V2 (2021) introduced discrete latent representations and was the first model-based method to surpass human benchmarks across all 55 Atari games. Dreamer V3 (2023) achieved cross-domain generality, delivering excellent performance with a single set of hyperparameters across tasks as diverse as Atari pixel games and Minecraft diamond collection — even becoming the first reinforcement learning algorithm to collect diamonds from scratch in Minecraft. This series of work continuously validates the engineering viability of the World Models paradigm.

Google DeepMind's Genie, released in early 2024, represents an important breakthrough for World Models in the generative AI direction. Genie is an 11-billion-parameter foundation world model that learned to generate interactive 2D virtual worlds by watching large amounts of unlabeled internet gameplay videos. Users need only provide a single image or even a text description, and Genie can generate an environment explorable through action controls. Its key innovation lies in unsupervised discovery of a "latent action space" from video — even without any action labels in the training data, the model can automatically infer the basic actions that control environmental changes. This represents a paradigm shift for World Models from "task-specific" to "general foundation models," hinting at a future where a "world foundation model" analogous to GPT for language might emerge.

It's fair to say that World Models are evolving from an academic concept into critical infrastructure for general-purpose agents and embodied intelligence.

World Models and Their Relationship to the Current AI Wave

In the current era dominated by large language models, World Models represent a different yet complementary technical path. LLMs excel at processing patterns in language and symbolic worlds, while World Models are better suited for environments with continuous dynamics and physical interactions.

As AI applications extend into real or simulated physical scenarios such as robotics, autonomous driving, and game NPCs, models that can "predict how the world will evolve" become increasingly important. The industry widely believes that only agents equipped with reliable world models can truly achieve long-horizon planning and efficient sample utilization. Notably, these two paths are beginning to converge — researchers are exploring ways to combine the reasoning capabilities of large language models with the physical prediction capabilities of world models, building unified agent architectures that can both understand language instructions and predict physical consequences.

For developers, starting with open-source demos like worldmodels_demo and understanding the collaboration mechanism between VAE, MDN-RNN, and Controller is an excellent entry point into this field. Although it's just a demo project, it's small but complete — neatly outlining the full picture of the powerful World Models paradigm.

Key Takeaways