Distributional Reinforcement Learning

Standard reinforcement learning methods focus on estimating the expected value (mean) of future discounted returns:

$$Q^\pi(s, a) = \mathbb{E}\left[ \sum_{t=0}^\infty \gamma^t R(s_t, a_t) ;\Big|; s_0=s, a_0=a \right]$$

While optimizing expected return is theoretically sufficient for risk-neutral decision making in Markov Decision Processes (MDPs), averaging collapses critical information about environmental uncertainty, multi-modality, and tail risks.

Distributional Reinforcement Learning, pioneered by Bellemare, Dabney, and Munos (2017), fundamentally transforms this paradigm. Instead of approximating a single scalar $Q$-value, distributional RL models the complete probability distribution of returns, denoted as the random variable $Z^\pi(s, a)$:

$$Q^\pi(s, a) = \mathbb{E}\left[ Z^\pi(s, a) \right]$$

Why Model the Full Return Distribution?

Two Actions with Identical Expected Value E[Z] = 10:

Action A (Deterministic Safe Route):
Probability Density
    ▲
    │          │ (Single spike at return = 10)
    │          │
────┴──────────┼───────────────► Return
               10

Action B (Stochastic High-Risk Route):
Probability Density
    ▲
    │   ▲             ▲ (50% crash -> -100, 50% jackpot -> +120)
    │   │             │
────┴───┼─────────────┼────────► Return
      -100     10    +120

Under classical RL, Action A and Action B have identical $Q$-values ($Q = 10$), making them indistinguishable to the agent. In reality:

• Modeling the full distribution preserves multi-modal outcomes (e.g., bifurcating paths or winning vs. losing a game).
• It provides non-linear representations that stabilize neural network feature learning and eliminate gradient noise.
• It enables risk-sensitive decision making (such as avoiding catastrophic tail events in autonomous driving or financial trading).

The Distributional Bellman Equation

In traditional RL, the scalar Bellman operator acts on expected values:

$$\mathcal{T}^\pi Q(s, a) := R(s, a) + \gamma \mathbb{E}_{P, \pi} \left[ Q(S’, A’) \right]$$

In Distributional RL, the operator acts directly on probability distributions:

$$Z(s, a) \stackrel{D}{:=} R(s, a) + \gamma Z(S’, A’)$$

where $\stackrel{D}{:=}$ denotes equality in distribution. The random variable on the right-hand side is formed by shifting the distribution $Z(S’, A’)$ by discount factor $\gamma$ and adding stochastic reward $R(s, a)$.

Bellemare et al. proved that while the distributional Bellman operator is not a contraction in the Wasserstein metric under the $\max$ policy operator, it is a contraction in the Wasserstein-1 ($W_1$) metric for policy evaluation.

The Three Core Architectures

1. Categorical DQN (C51)

Categorical DQN discretizes the continuous return space into $N = 51$ fixed, equally-spaced atoms support vector:

$$z = {z_{\min} + i \cdot \Delta z}{i=0}^{N-1}, \quad \Delta z = \frac{z{\max} - z_{\min}}{N - 1}$$

The neural network outputs 51 logits per action, normalized via softmax to produce probabilities $\mathbf{p}(s, a) \in \Delta^{51}$.

                 Categorical C51 Pipeline
State s ──► [ CNN Backbone ] ──► Linear Heads ──► Softmax ──► 51 Discrete Probabilities
                                                                (Histogram over [Vmin, Vmax])

The Projection Step

When the Bellman operator shifts and scales the distribution ($r + \gamma z_j$), the resulting atoms no longer align with the fixed support grid $z$. C51 projects the probability mass of each shifted atom onto its nearest neighboring support bins:

$$\Phi \hat{\mathcal{T}} Z(s, a) = \sum_{j=0}^{N-1} \left[ 1 - \frac{|\text{clip}(\hat{\mathcal{T}} z_j, z_{\min}, z_{\max}) - z_i|}{\Delta z} \right]_0^1 p_j(s’, a^*)$$

The network is trained by minimizing the Kullback-Leibler (KL) divergence or cross-entropy between the projected target distribution and the predicted distribution.

2. Quantile Regression DQN (QR-DQN)

Limitation of C51: C51 requires manual tuning of bounds $[z_{\min}, z_{\max}]$ and cannot extrapolate outside this fixed range.

QR-DQN (Dabney et al., 2018) transposes the problem: instead of fixing locations and learning probabilities, QR-DQN fixes probabilities and learns location quantiles:

$$\tau_i = \frac{2i - 1}{2N} \quad \text{for } i \in {1, \dots, N}$$

The network outputs $N$ quantile values $\theta_i(s, a)$ representing the inverse cumulative distribution function (CDF) at fixed cumulative probability fractions $\tau_i$.

Quantile Huber Loss

The model optimizes the asymmetric Quantile Huber Loss:

$$\rho_\tau^\kappa(\delta) = |\tau - \mathbb{I}(\delta < 0)| \cdot \mathcal{L}_\kappa(\delta)$$

where $\mathcal{L}_\kappa(\delta)$ is the standard smooth Huber loss with threshold $\kappa$. This formulation eliminates manual support clipping and provides unbiased $W_1$-metric contraction.

3. Implicit Quantile Networks (IQN)

While QR-DQN models a discrete set of quantiles (e.g., $N=200$), Implicit Quantile Networks (IQN) (Dabney et al., 2018) learn an infinite, continuous representation of the quantile function:

State s       ──► [ State Feature Extractor ψ(s) ] ──┐
                                                     ├──► Hadamard Product (ψ ⊙ φ) ──► MLP ──► Quantile Output Q_τ(s, a)
Quantile τ~U  ──► [ Cosine Embedding Layer φ(τ)  ] ──┘

1. A scalar probability $\tau \sim U(0, 1)$ is sampled uniformly at random.
2. $\tau$ is mapped into an embedding vector using harmonic cosine basis functions: $$\phi_j(\tau) = \text{ReLU}\left( \sum_{i=0}^{n-1} \cos(\pi i \tau) w_{ij} + b_j \right)$$
3. The state features $\psi(s)$ and quantile embedding $\phi(\tau)$ are merged via an element-wise product to predict quantile returns $Z_\tau(s, a)$.

Risk-Sensitive Policy Selection

A unique advantage of distributional RL is the ability to implement arbitrary risk profiles without retraining the neural network:

• Risk-Neutral (Standard Expected Value): $$\pi(s) = \arg\max_a \frac{1}{K} \sum_{k=1}^K Z_{\tau_k}(s, a)$$

• Risk-Averse (Conditional Value at Risk - CVaR): Evaluate only the bottom 10% worst-case quantile outcomes: $$\pi(s) = \arg\max_a \frac{1}{\alpha} \int_0^\alpha Z_\tau(s, a) , d\tau \quad (\text{with } \alpha = 0.1)$$ Critical for industrial robots and financial portfolios where worst-case drawdowns must be minimized.

• Optimistic Exploration: Act based on upper quantiles ($\tau \in [0.8, 1.0]$), driving agents to explore environments with rare, high-upside rewards.