World Models: Learning Environment Models with PyTorch

World Models represents one of the most seminal breakthroughs in Reinforcement Learning (RL) and Artificial Intelligence. Originally proposed by David Ha and Jürgen Schmidhuber in 2018, this architecture enables AI agents to learn an internal spatial and temporal representation of their environment, allowing them to train and make decisions inside their own “hallucinated dream” environments before acting in the real world.

In this article, we explore the core concepts of World Models based on the clean PyTorch implementation from nik-55/world-models ↗.

3 Core Components of World Models#

The architecture decouples perception, memory, and control into three distinct modules:

Visual Component (V - VAE) ---> Memory Component (M - MDN-RNN) ---> Controller (C)

1. Vision Model (V - Variational Autoencoder): Compresses high-dimensional visual observations (e.g., 64x64x3 RGB images) into a low-dimensional latent vector (z).
2. Memory Model (M - MDN-RNN): Predicts future latent states (z_{t+1}) based on past observations and actions taken by the agent.
3. Controller (C): A simple linear model responsible for selecting action (a_t) based solely on current latent state (z_t) and RNN hidden state (h_t).

[!NOTE] Decoupling perception and memory keeps the Controller (C) extremely lightweight, enabling efficient policy optimization using Evolution Strategies (ES) or CMA-ES without needing full end-to-end backpropagation.

PyTorch Code Structure (nik-55/world-models)#

The nik-55/world-models ↗ repository recreates the classic CarRacing-v0 environment from OpenAI Gym with modular PyTorch code:

• models/vae.py: ConvVAE nén khung hình thành vector tiềm ẩn (z \in \mathbb{R}^{32}).
• models/mdnrnn.py: Combines LSTM with Mixture Density Networks (MDN) to output probability distributions of future states.
• models/controller.py: Linear controller determining steering, gas, and brake.
• trainvae.py, trainmdn.py, traincontroller.py: Separated training scripts.

import torch
import torch.nn as nn
import torch.nn.functional as F

class VAE(nn.Module):
    def __init__(self, img_channels=3, latent_size=32):
        super(VAE, self).__init__()
        self.latent_size = latent_size
        
        # Encoder
        self.encoder = nn.Sequential(
            nn.Conv2d(img_channels, 32, 4, stride=2),
            nn.ReLU(),
            nn.Conv2d(32, 64, 4, stride=2),
            nn.ReLU(),
            nn.Conv2d(64, 128, 4, stride=2),
            nn.ReLU(),
            nn.Conv2d(128, 256, 4, stride=2),
            nn.ReLU()
        )
        self.fc_mu = nn.Linear(256 * 2 * 2, latent_size)
        self.fc_logvar = nn.Linear(256 * 2 * 2, latent_size)
        
        # Decoder
        self.fc_decoder = nn.Linear(latent_size, 1024)
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(1024, 128, 5, stride=2),
            nn.ReLU(),
            nn.ConvTranspose2d(128, 64, 5, stride=2),
            nn.ReLU(),
            nn.ConvTranspose2d(64, 32, 6, stride=2),
            nn.ReLU(),
            nn.ConvTranspose2d(32, img_channels, 6, stride=2),
            nn.Sigmoid()
        )

    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std

3-Stage Training Pipeline#

1. Random Data Rollouts: Run a random policy agent to collect thousands of image observations across CarRacing episodes.
2. Train VAE & MDN-RNN:
   • VAE is trained via Reconstruction + KL Divergence loss.
   • MDN-RNN is trained to predict (z_{t+1}) modeled by a Gaussian Mixture Model (GMM).
3. Dream Training (Controller Optimization):
   • Controller is trained purely inside the MDN-RNN hallucinated environment without invoking the real Gym environment.

[!TIP] Training policies inside a simulated dream environment speeds up iterations significantly compared to interacting directly with physical or simulated environments.

References#

• Ha, D., & Schmidhuber, J. (2018). Recurrent World Models Facilitate Policy Evolution. NeurIPS.
• GitHub Repository: nik-55/world-models ↗