Distributional Reinforcement Learning: Beyond Expected Returns

Module: Machine Learning | Difficulty: Advanced

Distributional Bellman Equation

C51 (Categorical DQN)

Quantile Regression DQN

where

Distributional vs Expected RL

import torch
import torch.nn as nn

class C51(nn.Module):
    def __init__(self, n_actions, n_atoms=51, v_min=-10, v_max=10):
        super().__init__()
        self.n_atoms = n_atoms
        self.v_min = v_min; self.v_max = v_max
        self.delta_z = (v_max - v_min) / (n_atoms - 1)
        self.z = torch.linspace(v_min, v_max, n_atoms)
        self.net = nn.Sequential(
            nn.Linear(84, 128), nn.ReLU(),
            nn.Linear(128, 128), nn.ReLU(),
            nn.Linear(128, n_actions * n_atoms))
    def forward(self, x):
        logits = self.net(x).view(-1, n_actions, self.n_atoms)
        return torch.softmax(logits, dim=-1)
    def q_values(self, probs):
        return (probs * self.z).sum(dim=-1)

Research Insight: Distributional RL learns the full return distribution, not just the expected value. This leads to better exploration and risk-sensitive behavior. The key insight is that the Bellman equation holds exactly for the full distribution.