C51

Definition

The C51 algorithm is a Distributional Reinforcement Learning algorithm that approximates the distribution of the random return a using discretized distribution.

Architecture

The range of possible returns is divided into a fixed support of $N$-equally spaced bins $\{z_i{\}}_{i=0}^{N-1}$, where $z_i=R_{\text{min}}+i\Delta z$ and $\Delta z=\frac{R_{\text{max}}-R_{\text{min}}}{N-1}$. $R_{\text{min}}$ and $R_{\text{max}}$ are the minimum and maximum possible return values, respectively. In the C51 algorithm, the support determined once at the beginning and remains fixed throughout the entire learning process.

For each state-action pair $(s,a)$, the algorithm estimates the probabilities $\{p_i(s,a)\}$ that the return will fall into each $i$-th bin. The probabilities are estimated by a neural network, which takes the state $s$ as input and outputs a vector of probabilities $\{p_i(s,a)\}$ for each state-possible action pair $(s,a),\forall a\in \mathcal{A}(s)$. The Softmax Function ensures the properties of the probability of the output for each state-action pair. $p_i(s,a;\theta )=\frac{e^{{\theta }_i(s,a)}}{{\sum }_{j=0}^{N-1}{e}^{{\theta }_j(s,a)}}$ where ${\theta }_i(s,a)$ are the raw outputs (logits) of the network for the $i$-th atom and action $a$, and $\theta$ represents the network’s parameters.

The discrete distribution $Z_{\theta }$ over the fixed support $\{z_i{\}}_{i=0}^{N-1}$ is constructed as $Z_{\theta }(s,a)=\sum _{i=0}^{N-1}{p}_i(s,a){\delta }_{z_i}$ where ${\delta }_x$ is a Dirac’s delta function at $x$.

Projection

In a policy evaluation process, discrete distributions have a problem: $\mathcal{T}{Z}_{\theta }$ and $Z_{\theta }$ have disjoint supports. Though, Wasserstein Metric in the original setup is robust to this issue, but in practice, KL-Divergence is used instead due to the distributional instability and differentiability. Therefore, the problem still matters.

The disjoint support problem is solved by projection.

1. Compute the distributional Bellman update: ${\hat{Z}}_{\theta }\leftarrow r+\gamma Z_{\theta }(s^{\prime },a^{\ast })$ where $a^{\ast }=\underset{a}{\mathrm{argmax}}Q(s^{\prime },a)$
2. Distribute the computed probability mass ${\hat{Z}}_{\theta }$ to the neighboring bins in the support $\{z_i{\}}_{i=0}^{N-1}$ proportionally to the distance from each original point: $(\Phi \mathcal{T}{Z}_{\theta })(s,a)$, where $\Phi$ is the projection operator.

Loss Function

C51 uses the KL-Divergence as a loss function. $L_{s,a}(\theta )=D_{KL}(\Phi \hat{\mathcal{T}}{Z}_{\hat{\theta }}(s,a)||Z_{\theta }(s,a)))$ where $\theta$ and $\hat{\theta }$ are the parameters for behavior and target networks, respectively.

The loss function can be simplified to the Cross-Entropy Loss. $L(\theta )=-\sum _{i=0}^{N-1}(\Phi \hat{\mathcal{T}}Z(s,a;\hat{\theta }){)}_i\ln p_i(s,a;\theta )$

Algorithm

1. Initialize behavior network $Q$ and target network $\hat{Q}$ with random weights $\theta$, and the replay buffer $\mathcal{R}$ to max size $N$.
2. Repeat for each episode:
   1. Initialize sequence $s_0$.
   2. Repeat for each step of an episode until terminal, $t=0,1,\dots ,T-1$ :
      1. With probability $ϵ$, select a random action $a_t$ otherwise select $a_t=\underset{a}{\mathrm{argmax}}Q(s_t,a;\theta )$.
      2. Take the action $a_t$ and observe a reward $r_{t+1}$ and a next state $s_{t+1}$.
      3. Store transition $(s_t,a_t,r_{t+1},s_{t+1})$ in the replay buffer $\mathcal{R}$.
      4. Sample a random transition $(s_t,a_t,r_{t+1},s_{t+1})$ from $\mathcal{R}$.
      5. Compute $Q(s_{t+1},a)=\sum _{i=0}^{N-1}{z}_i{p}_i(s_{t+1},\alpha ;\hat{\theta })$ and select a greedy action $a^{\ast }=\underset{a}{\mathrm{argmax}}Q(s_{t+1},a)$.
      6. Perform Gradient Descent on the loss $L(\theta )=-\sum _{i=0}^{N-1}(\Phi \hat{\mathcal{T}}Z(s_{t+1},a^{\ast };\hat{\theta }){)}_i\ln p_i(s_t,a_t;\theta )$.
      7. Update the target network parameter $\hat{\theta }\leftarrow \theta$ every $C$ steps.
      8. Update $s_t\leftarrow s_{t+1}$.