Chapter 6: Model-Based Reinforcement Learning

6.1 Introduction to Model-Based RL

Model-based RL learns a model of the environment's dynamics and uses it for planning. This contrasts with model-free methods that learn values or policies directly from experience without an explicit model.

Why Use Models?

• Sample Efficiency: Learn from simulated experience (no real environment interaction needed)
• Planning: Can reason about future consequences without trial and error
• Transfer: Models can generalize to new tasks
• Safety: Test policies in simulation before deployment
• Data Generation: Generate synthetic training data

6.2 Dyna Architecture

Dyna combines model-free and model-based learning by integrating planning, acting, and model-learning:

Algorithm: Dyna-Q

Initialize Q(s,a) and Model(s,a)

Loop:
    // (a) Acting: interact with environment
    s = current state
    a = ε-greedy(Q, s)
    Execute a, observe r, s'

    // (b) Model-free learning: direct RL
    Q(s,a) ← Q(s,a) + α[r + γ max Q(s',a') - Q(s,a)]

    // (c) Model learning: update transition/reward model
    Model(s,a) ← (s', r)

    // (d) Planning: learn from simulated experience
    Repeat n times:
        s_sim = random previously observed state
        a_sim = random previously taken action
        s'_sim, r_sim = Model(s_sim, a_sim)
        Q(s_sim, a_sim) ← Q(s_sim, a_sim) + α[r_sim + γ max Q(s'_sim, a') - Q(s_sim, a_sim)]

Dyna-Q Implementation

import numpy as np
from collections import defaultdict

class DynaQ:
    def __init__(self, n_states, n_actions, planning_steps=50, alpha=0.1, gamma=0.95, epsilon=0.1):
        self.Q = np.zeros((n_states, n_actions))
        self.model = {}  // Model(s,a) -> (s', r)
        self.planning_steps = planning_steps
        self.alpha = alpha
        self.gamma = gamma
        self.epsilon = epsilon

    def select_action(self, state):
        if np.random.random() < self.epsilon:
            return np.random.randint(self.Q.shape[1])
        return np.argmax(self.Q[state])

    def update(self, state, action, reward, next_state):
        // Direct RL (model-free)
        td_target = reward + self.gamma * np.max(self.Q[next_state])
        td_error = td_target - self.Q[state, action]
        self.Q[state, action] += self.alpha * td_error

        // Model learning
        self.model[(state, action)] = (next_state, reward)

        // Planning (model-based)
        for _ in range(self.planning_steps):
            // Sample random previously experienced (s,a)
            (s, a), (s_next, r) = random.choice(list(self.model.items()))

            // Simulate Q-learning update
            td_target = r + self.gamma * np.max(self.Q[s_next])
            td_error = td_target - self.Q[s, a]
            self.Q[s, a] += self.alpha * td_error

# Usage
env = gym.make('FrozenLake-v1')
agent = DynaQ(n_states=env.observation_space.n,
              n_actions=env.action_space.n,
              planning_steps=50)

state = env.reset()
for step in range(1000):
    action = agent.select_action(state)
    next_state, reward, done, _ = env.step(action)
    agent.update(state, action, reward, next_state)
    state = next_state if not done else env.reset()

6.3 Learning Dynamics Models

Modern model-based RL uses neural networks to learn complex dynamics:

Deterministic Dynamics Model

import torch
import torch.nn as nn

class DynamicsModel(nn.Module):
    def __init__(self, state_dim, action_dim, hidden_dim=256):
        super(DynamicsModel, self).__init__()

        self.network = nn.Sequential(
            nn.Linear(state_dim + action_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, state_dim + 1)  // Predict Δs and reward
        )

    def forward(self, state, action):
        """
        Predict next state and reward.

        Args:
            state: current state
            action: action taken

        Returns:
            next_state: predicted next state
            reward: predicted reward
        """
        x = torch.cat([state, action], dim=-1)
        output = self.network(x)

        delta_state = output[:, :-1]  // State change
        reward = output[:, -1:]       // Reward

        next_state = state + delta_state  // Residual prediction
        return next_state, reward

# Training the model
def train_dynamics_model(model, replay_buffer, batch_size=256, epochs=10):
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

    for epoch in range(epochs):
        states, actions, rewards, next_states, _ = replay_buffer.sample(batch_size)

        // Predict
        pred_next_states, pred_rewards = model(states, actions)

        // Loss
        state_loss = nn.MSELoss()(pred_next_states, next_states)
        reward_loss = nn.MSELoss()(pred_rewards, rewards)
        loss = state_loss + reward_loss

        // Update
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if (epoch + 1) % 10 == 0:
            print(f"Epoch {epoch+1}, Loss: {loss.item():.4f}")

Probabilistic Dynamics Model

For stochastic environments, predict distributions:

class ProbabilisticDynamicsModel(nn.Module):
    def __init__(self, state_dim, action_dim, hidden_dim=256):
        super(ProbabilisticDynamicsModel, self).__init__()

        self.shared = nn.Sequential(
            nn.Linear(state_dim + action_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU()
        )

        // Mean and log variance of next state
        self.mean_head = nn.Linear(hidden_dim, state_dim)
        self.logvar_head = nn.Linear(hidden_dim, state_dim)

        // Reward prediction
        self.reward_head = nn.Linear(hidden_dim, 1)

    def forward(self, state, action):
        features = self.shared(torch.cat([state, action], dim=-1))

        mean = self.mean_head(features) + state  // Residual
        logvar = self.logvar_head(features)
        reward = self.reward_head(features)

        return mean, logvar, reward

    def sample(self, state, action):
        """Sample next state from learned distribution."""
        mean, logvar, reward = self.forward(state, action)
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        next_state = mean + eps * std
        return next_state, reward

6.4 Model Predictive Control (MPC)

MPC uses the model to plan actions by optimizing over action sequences:

def mpc_planning(model, state, horizon=10, num_candidates=1000):
    """
    Model Predictive Control: plan best action sequence.

    Args:
        model: learned dynamics model
        state: current state
        horizon: planning horizon
        num_candidates: number of random action sequences to try

    Returns:
        best_action: first action of best sequence
    """
    best_return = -float('inf')
    best_action = None

    for _ in range(num_candidates):
        // Sample random action sequence
        action_sequence = [sample_random_action() for _ in range(horizon)]

        // Simulate trajectory
        total_return = 0
        sim_state = state.clone()

        for t, action in enumerate(action_sequence):
            // Predict next state and reward using model
            sim_state, reward = model(sim_state, action)
            total_return += (gamma ** t) * reward

        // Keep best sequence
        if total_return > best_return:
            best_return = total_return
            best_action = action_sequence[0]

    return best_action

// Usage
state = torch.FloatTensor(env.reset())
action = mpc_planning(dynamics_model, state, horizon=10)
next_state, reward, done, _ = env.step(action.numpy())

Cross-Entropy Method (CEM) for MPC

def cem_planning(model, state, horizon=10, population=500, elite_frac=0.1, iterations=5):
    """
    Cross-Entropy Method for action sequence optimization.
    """
    action_dim = model.action_dim

    // Initialize: mean=0, std=1 for each action in sequence
    mean = torch.zeros(horizon, action_dim)
    std = torch.ones(horizon, action_dim)

    for iteration in range(iterations):
        // Sample action sequences from current distribution
        action_sequences = []
        for _ in range(population):
            actions = torch.normal(mean, std)
            action_sequences.append(actions)

        // Evaluate each sequence
        returns = []
        for actions in action_sequences:
            total_return = simulate_trajectory(model, state, actions)
            returns.append(total_return)

        // Select elite sequences
        elite_count = int(population * elite_frac)
        elite_indices = torch.topk(torch.tensor(returns), elite_count).indices
        elite_sequences = [action_sequences[i] for i in elite_indices]

        // Update distribution
        elite_sequences = torch.stack(elite_sequences)
        mean = elite_sequences.mean(dim=0)
        std = elite_sequences.std(dim=0)

    // Return first action of best sequence
    return mean[0]

6.5 Monte Carlo Tree Search (MCTS)

MCTS builds a search tree by repeatedly simulating and expanding promising paths:

class MCTSNode:
    def __init__(self, state, parent=None):
        self.state = state
        self.parent = parent
        self.children = {}
        self.visits = 0
        self.value = 0

    def is_fully_expanded(self, action_space):
        return len(self.children) == len(action_space)

    def best_child(self, c=1.41):
        """UCB1 formula for child selection."""
        choices = []
        for action, child in self.children.items():
            if child.visits == 0:
                ucb = float('inf')
            else:
                exploit = child.value / child.visits
                explore = c * np.sqrt(np.log(self.visits) / child.visits)
                ucb = exploit + explore
            choices.append((ucb, action, child))

        return max(choices, key=lambda x: x[0])[2]

def mcts(root_state, model, simulations=1000):
    """
    Monte Carlo Tree Search using learned model.

    Steps:
    1. Selection: traverse tree using UCB
    2. Expansion: add new node
    3. Simulation: rollout using model
    4. Backpropagation: update values
    """
    root = MCTSNode(root_state)

    for _ in range(simulations):
        node = root
        state = root_state.clone()

        // 1. Selection
        while node.is_fully_expanded(action_space) and node.children:
            node = node.best_child()
            state, _ = model(state, node.action)

        // 2. Expansion
        if not node.is_fully_expanded(action_space):
            action = random.choice([a for a in action_space if a not in node.children])
            next_state, reward = model(state, action)
            child = MCTSNode(next_state, parent=node)
            child.action = action
            node.children[action] = child
            node = child
            state = next_state

        // 3. Simulation (rollout)
        total_reward = 0
        depth = 0
        while depth < max_depth:
            action = random.choice(action_space)
            state, reward = model(state, action)
            total_reward += (gamma ** depth) * reward
            depth += 1

        // 4. Backpropagation
        while node is not None:
            node.visits += 1
            node.value += total_reward
            node = node.parent

    // Return best action
    return max(root.children.items(), key=lambda x: x[1].visits)[0]

6.6 World Models

World Models learn compact latent representations of environments:

class WorldModel:
    def __init__(self, obs_dim, action_dim, latent_dim=32):
        // Vision model: encode observations
        self.vae = VAE(obs_dim, latent_dim)

        // Memory model: predict next latent state
        self.rnn = MDN_RNN(latent_dim, action_dim, hidden_dim=256)

        // Controller: simple linear policy
        self.controller = Controller(latent_dim + hidden_dim, action_dim)

    def encode(self, observation):
        """Compress observation to latent code."""
        return self.vae.encode(observation)

    def predict_next(self, latent_state, action, hidden):
        """Predict next latent state given current state and action."""
        return self.rnn(latent_state, action, hidden)

    def act(self, latent_state, hidden):
        """Select action based on latent state."""
        return self.controller(torch.cat([latent_state, hidden]))

    def dream(self, initial_latent, steps=100):
        """
        Simulate future trajectory in latent space.
        Train controller entirely in imagination!
        """
        latent = initial_latent
        hidden = self.rnn.init_hidden()
        trajectory = []

        for _ in range(steps):
            action = self.act(latent, hidden)
            next_latent, reward, hidden = self.predict_next(latent, action, hidden)

            trajectory.append((latent, action, reward, next_latent))
            latent = next_latent

        return trajectory

6.7 Model-Based vs Model-Free

6.8 Hybrid Approaches

Model-Based Value Expansion (MVE)

Use model to improve value estimates in model-free algorithms:

// Standard TD target
V_target = r + γ V(s')

// MVE: use model to expand k-step into future
V_target = r + γ r' + γ² r'' + ... + γᵏ V(s_k)

where s', r' = Model(s, a)
      s'', r'' = Model(s', a')
      ...

AlphaZero Architecture

class AlphaZero:
    """
    Combines deep learning, MCTS, and self-play.
    """
    def __init__(self):
        // Single neural network outputs (policy, value)
        self.network = ResidualNetwork()

    def search(self, state):
        """MCTS with network for policy and value."""
        root = MCTSNode(state)

        for _ in range(num_simulations):
            node = root

            // Selection
            while node.is_expanded():
                node = node.select_child()

            // Expansion and evaluation
            policy, value = self.network(node.state)
            node.expand(policy)

            // Backpropagation
            node.backup(value)

        // Return improved policy
        return root.get_action_probs()

    def self_play(self):
        """Generate training data through self-play."""
        game = Game()
        trajectory = []

        while not game.done():
            state = game.state()
            policy = self.search(state)
            action = sample(policy)

            game.step(action)
            trajectory.append((state, policy))

        // Game result as value target
        result = game.result()

        return [(s, p, result) for s, p in trajectory]

6.9 Practical Considerations

When to Use Model-Based RL?

• Limited Data: When real interactions are expensive
• Sim-to-Real: When good simulators exist
• Safety Critical: Test in simulation first
• Interpretability: Model provides understanding
• Transfer Learning: Reuse models across tasks

Challenges

• Model Error: Inaccurate models hurt performance
• Compounding Errors: Errors accumulate over long horizons
• Complexity: Harder to implement than model-free
• Computational Cost: Planning is expensive