Introduction

Reinforcement Learning (RL) is a powerful approach for solving sequential decision-making problems, where an agent learns to act in an environment to maximize cumulative rewards. While many RL algorithms fall into the model-freecategory, there's another exciting branch of RL: Model-Based Reinforcement Learning (MBRL). In this post I'll introduce MBRL and walk through a SimPLe-style agent on Atari's Breakout — one that learns a model of the game from pixels and then trains its policy inside that model rather than on the real thing.

What is Model-Based RL?

In model-based RL, instead of learning the optimal policy purely by interacting with the environment, we first learn a model of the environment’s dynamics. This model predicts the next state and reward, given the current state and action. Once the agent has this learned model, it can plan by simulating future states, enabling more efficient decision-making compared to model-free approaches.

Model-based RL typically involves three key steps:

  1. Model Learning: Learn the dynamics of the environment, i.e., how states transition from one to another.
  2. Planning: Use the learned model to simulate the future and find the best actions.
  3. Policy Improvement: Refine the policy based on the outcomes from the planning phase.

Breakout puts a hard edge on all three, because the state the model has to predict is the screen itself. The rest of this post follows that through: what the world model predicts, the two Atari details the whole thing hinges on, and what the run actually scored.

Why Model-Based RL?

The advantage of model-based RL lies in its ability to learn policies more efficiently. Since the agent can simulate the environment instead of relying solely on real interactions, it can dramatically reduce the number of interactions required to learn a good policy. This is particularly useful in environments where real-world interactions are costly or limited (e.g., robotics).

Why Breakout Is Hard for a World Model

CartPole has four numbers of state. Breakout has a screen. A model-based agent here has to predict pixels, and a model that is even slightly wrong compounds its own error every step it dreams forward — which is why naive pixel regression fails: minimising squared error over an uncertain future produces a blurry average of the possibilities, and a blurry ball is no ball at all.

The approach used here is SimPLe (Simulated Policy Learning), from Kaiser et al. (2019), on ALE/Breakout-v5. Its loop is four steps:

  1. Collect real transitions — uniform random on the first pass, epsilon-greedy around the current policy afterwards.
  2. Train an action-conditional world model that predicts the next frame, the reward and whether the episode ended.
  3. Train a PPO actor-critic entirely inside that model — no real frames are touched during policy learning.
  4. Repeat — the improved policy collects better data, the world model is fine-tuned on everything so far, and the policy is retrained inside the updated model.

The point of all this is sample efficiency. Real Atari frames are the expensive resource; imagined ones are nearly free.

The World Model

The model takes the 4-frame stack and predicts the next single frame, plus a reward and a done signal. The interesting choice is that all three heads are categorical.

class EnvModel(nn.Module):
    """Action-conditional world model with categorical output heads.

    - Next frame: per-pixel 256-way softmax (discretized intensity bins),
      which avoids the blurry-MSE failure mode of naive L2 pixel regression.
    - Reward: 3-way classifier over clipped reward classes {-1, 0, +1}.
    - Done: Bernoulli logit (binary cross-entropy).
    """

    def __init__(self, obs_shape, n_actions: int, cfg: Config):
        super().__init__()
        C, H, W = obs_shape
        hid = cfg.hidden_size
        self.action_embedding = nn.Embedding(n_actions, hid)

        self.encoder = nn.Sequential(
            nn.Conv2d(C, hid, kernel_size=4, stride=2, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(hid, hid, kernel_size=4, stride=2, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(hid, hid, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
        )

        # Per-pixel 256-way softmax head over one predicted frame.
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(hid, hid, kernel_size=4, stride=2, padding=1),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(hid, hid, kernel_size=4, stride=2, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(hid, cfg.pixel_bins, kernel_size=3, stride=1, padding=1),
        )

        self.reward_head = nn.Sequential(
            nn.Linear(hid, 64), nn.ReLU(inplace=True), nn.Linear(64, 3),
        )
        self.done_head = nn.Sequential(
            nn.Linear(hid, 64), nn.ReLU(inplace=True), nn.Linear(64, 1),
        )

Treating each pixel as a 256-way classification rather than a regression is the fix for the blurring problem: the model puts probability mass on discrete intensities instead of averaging them. Rewards get clipped to three classes, and done is a single Bernoulli logit. Only one frame is predicted at a time — in imagination the stack rolls forward by dropping the oldest channel and appending the prediction.

Two Atari Details That Decide Whether This Works

Breakout will not train without these, and neither is obvious.

Something has to press FIRE. Breakout does not launch the ball until the FIRE action is taken, and it needs it again after every lost life. A FireResetWrapper presses it automatically on reset and after each life loss; without it the agent sits watching an empty screen and learns nothing.

Life loss has to be a terminal during collection but not during evaluation. The config splits these deliberately:

@dataclass
class Config:
    env_name: str = "ALE/Breakout-v5"
    screen_size: int = 84
    frame_skip: int = 4
    frame_stack: int = 4
    terminal_on_life_loss: bool = True
    eval_terminal_on_life_loss: bool = False

    # SimPLe outer loop
    num_simple_iters: int = 3
    steps_per_iter: int = 4_000       # real env steps collected per iteration
    collect_epsilon: float = 0.1      # epsilon-greedy around policy (iter >= 1)
    first_iter_random: bool = True    # uniform random on iteration 0

    # World model
    hidden_size: int = 64
    world_model_epochs: int = 5
    pixel_bins: int = 256             # per-pixel 256-way softmax

    # Imagination-based policy training (PPO)
    imagination_horizon: int = 15
    imagination_iters: int = 200
    gamma: float = 0.99
    gae_lambda: float = 0.95
    ppo_clip: float = 0.2

Marking each life loss as terminal during collection gives the world model explicit per-life done labels, so imagination never has to model the death → FIRE → respawn sequence — a transition that is rare in the data and, left unmodelled, makes dreamed rollouts diverge the moment the first life is lost. Evaluation then uses full multi-life episodes, so the score stays comparable to the model-free baselines and to the numbers in the SimPLe paper.

The 4-frame stack matters for the same reason it does in any Atari agent: a single frame tells you where the ball is, not where it is going.

Results

All runs on a Google Colab L4.

EnvironmentModelAverage rewardTotal training steps
BreakoutNoFrameskip-v4SimPLenot recorded5,000,000
BreakoutNoFrameskip-v4DQN239.20 +/- 73.635,000,000
BreakoutNoFrameskip-v4PPO187.80 +/- 114.625,000,000
ALE/Breakout-v5DQN298.70 +/- 33.817,500,000
ALE/Breakout-v5PPO398.30 +/- 19.097,500,000

The model-free baselines are the honest headline: DQN reaches 239.20 and PPO 187.80 on BreakoutNoFrameskip-v4 at five million steps, and both go higher on ALE/Breakout-v5 given seven and a half million.

No comparable SimPLe score is recorded. That is worth stating plainly rather than papering over, because it is the usual outcome when you first stand up a model-based pipeline: the defaults in this notebook are deliberately small so it finishes in a few minutes, and small SimPLe runs produce a world model good enough to dream in but not good enough to train a policy that scores well on the real game. Reaching paper-scale numbers means the Config.full_training() preset — roughly 100,000 real environment steps over 15 SimPLe iterations with a deeper model and a 50-step imagination horizon, which is four to eight hours on an L4 and needs the high-RAM runtime for the replay buffer.

So the fair reading is that this is a working SimPLe implementation, not a competitive Breakout agent. The comparison it does support is about where the compute goes: the model-free runs spend millions of real frames, while SimPLe spends thousands of real frames and millions of imagined ones.

Next Steps

Where to take this next:

  • Scale the run: the Config.full_training() preset targets roughly 100,000 real environment steps over 15 SimPLe iterations, which is what a comparable score would need.
  • Uncertainty-Aware Models: Learn a distribution over possible outcomes (using Bayesian methods or ensembles) to account for model uncertainty.
  • Sample the world model: turning on imagination_sample draws frames and rewards from the categorical heads instead of taking the argmax, which stops PPO exploiting a mode-collapsed dream.
  • Compare on real frames: measure SimPLe against the DQN and PPO baselines at equal real environment steps rather than equal total steps — that is the comparison sample efficiency is actually about.

Conclusion

Model-based RL trades a hard problem for a different hard problem. Instead of needing millions of real frames, you need a model accurate enough that a policy trained inside it still works outside it — and on Atari that means predicting pixels well enough that a ball two frames from now is still a ball.

SimPLe shows the trade is workable: categorical pixel heads instead of squared error, per-life terminals so imagination never has to dream the respawn, and a policy that touches no real frames while it learns. What this run does not show is a competitive score, and the small defaults are why. The model-free baselines here reach 239.20 and 187.80 at five million real steps; matching that from a world model is a far longer run than a few minutes on a Colab GPU.

Additional Learning Materials

Code Repository & Models