Category: RL Alias: World Models

Model-Based Reinforcement Learning

A paradigm where agents learn a predictive model of environment dynamics to plan, simulate, and optimize policies entirely within "imagined" rollouts.

In One Sentence

Instead of learning purely by trial-and-error in the real world, a Model-Based RL agent constructs an internal simulator (a "world model") of the environment's transitions and rewards, using it to "dream" up future scenarios and optimize its behavior without physical risk or sample-heavy environment interaction.

Intuition

Imagine learning to drive a car. A Model-Free agent learns by driving off a cliff thousands of times, slowly adjusting its policy parameters because "falling off cliff = bad reward."

A Model-Based agent, however, builds an internal mental model of physics. It understands that "if I turn the wheel sharply at high speed near a cliff, the car will fly off." Because it can simulate this outcome in its head, it never has to actually drive off the cliff to realize it is a bad idea.

✕ Model-Free RL

Interacts directly with the real world. Updates policies based on direct, experienced feedback. Extremely sample-inefficient but structurally simple.

✓ Model-Based RL

Learns a transition function $T(s'|s, a)$ and reward function $R(s, a)$. Plans actions or trains policies inside this "hallucinated" world model. Highly sample-efficient.

Formal Definition

A standard Markov Decision Process (MDP) is defined by the tuple $\mathcal{M} = (\mathcal{S}, \mathcal{A}, \mathcal{P}, \mathcal{R}, \gamma)$. Model-Based RL seeks to approximate the true transition dynamics $\mathcal{P}$ and reward function $\mathcal{R}$ with learned parameterized models:

Transition Model $$\hat{T}_\phi(s_{t+1} \mid s_t, a_t)$$

Reward Model $$\hat{R}_\psi(r_t \mid s_t, a_t)$$

Policy Optimization in Imagination $$\max_\theta \mathbb{E}_{\tau \sim \hat{P}_\phi, \pi_\theta} \left[ \sum_{t=0}^{H} \gamma^t \hat{R}_\psi(s_t, a_t) \right]$$

Where $H$ is the imagination horizon. Modern variants (like Dreamer) map high-dimensional observations $x_t$ into a compact latent state space $z_t$ using an encoder $q_\theta(z_t \mid x_t, z_{t-1}, a_{t-1})$, performing all transitions in latent space to prevent modeling irrelevant pixel-level noise.

When to Use & When to Avoid

Use When:

✓ Sample cost is high: Real-world robotics, clinical trials, or physical systems where taking an action is slow or dangerous.
✓ Planning is essential: Environments requiring deep search and lookahead (e.g., Chess, Go) where a perfect simulator is unavailable.
✓ Multi-task transfer: You want to learn a single model of the world and swap reward functions on the fly for different tasks.

Avoid When:

✕ Compounding errors dominate: High-dimensional, chaotic environments where small prediction errors accumulate, leading to "model exploitation."
✕ Simulations are cheap: If you have a fast, exact simulator (e.g., simple grid-worlds, Atari on massive GPU clusters), model-free RL is often simpler and more robust.

Milestones & History

2018

Unsupervised "Dreams" (World Models)

Ha & Schmidhuber demonstrate training a compact policy entirely inside the "hallucinated dream" generated by an unsupervised world model (VAE + RNN), proving policies can transfer seamlessly back to physical environments.

2019

Value-Equivalent Planning (MuZero)

Schrittwieser et al. introduce MuZero, which bypasses full reconstruction. By predicting only policy, value, and reward, it masters Atari, Go, Chess, and Shogi without knowing the underlying rules or dynamics.

2023

Generalization Across Domains (DreamerV3)

Hafner et al. present DreamerV3, showing that a single model-based configuration can learn across 150+ diverse tasks, achieving the milestone of collecting diamonds in Minecraft from scratch without human data.

Interactive Sandbox: The Compounding Error Problem

Adjust the sliders below to see how Model Accuracy and Imagination Horizon affect the agent's simulated trajectory. Notice how small errors compound over time, causing the "dreamed" trajectory to drift away from reality.

Model Accuracy: 90%

Imagination Horizon: 15 steps

True Trajectory Imagined "Dream" Trajectory

Takeaway: When planning deep into the future, even a 95% accurate model can hallucinate unfeasible states. This is why modern algorithms use short horizons or regular grounding updates from the real environment.

Pseudocode: Latent Imagination Loop

world_model_training.py Python

import torch
import torch.nn as nn

# Conceptual World Model training loop
def train_world_model(agent, env, episodes=1000):
    dataset = ReplayBuffer()
    
    for ep in range(episodes):
        obs = env.reset()
        while not done:
            # 1. Select action using policy optimized in imagination
            action = agent.policy(obs)
            next_obs, reward, done, _ = env.step(action)
            dataset.add(obs, action, reward, next_obs)
            obs = next_obs
            
        # 2. Update World Model (Transition & Reward models)
        batch = dataset.sample(batch_size=64)
        loss = agent.world_model.compute_loss(
            batch.obs, batch.actions, batch.rewards, batch.next_obs
        )
        agent.world_model.optimizer.step(loss)
        
        # 3. "Dream" & Optimize Policy inside Imagination
        imagined_states, imagined_rewards = agent.world_model.imagine(
            start_states=batch.obs, horizon=15
        )
        policy_loss = agent.policy.update(imagined_states, imagined_rewards)
        agent.policy.optimizer.step(policy_loss)

Common Variants

Dyna-Style RL

Uses the model to generate simulated transitions to augment real data, mixing both to update model-free Q-values or policies (e.g., Dyna-Q).

Latent World Models

Learns transitions in a compact, latent space rather than raw pixel space, ignoring visual noise (e.g., Dreamer).

Value-Equivalent Planning

Only models aspects of the environment critical to planning (rewards, values, action policies) rather than full state transitions (e.g., MuZero).

Foundational Components in Our Corpus

Foundational Component 5neujr5h96

Auto-Encoding Variational Bayes (VAE)

The Variational Autoencoder (VAE) is the cornerstone architecture used in early world models (such as Ha & Schmidhuber, 2018) to compress raw, high-dimensional visual observations into compact, low-dimensional latent states $z_t$ before passing them to the temporal transition network.

Grounding References

arXiv:1803.10122
World Models
Ha & Schmidhuber (2018). Explores training compact policies entirely inside hallucinated dreams generated by unsupervised latent models.
arXiv:1911.08265
Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model
Schrittwieser et al. (2019). Introduces MuZero, combining tree-based search with learned value-equivalent dynamics.
arXiv:2301.04104
Mastering Diverse Domains through World Models
Hafner et al. (2023). Introduces DreamerV3, demonstrating robust, single-configuration model-based RL across over 150 tasks.