PaperlayerView exact revision on arXiv

Source-linked AI summary

A Distributional Perspective on Reinforcement Learning

Marc G. Bellemare, Will Dabney, Rémi Munos

arXiv:1707.06887v1cs.LGcs.AIstat.ML

TL;DR

Reinforcement learning commonly models expected return, but the paper asks whether the full value distribution should instead be central. It analyzes distributional Bellman operators, introduces an approximate distribution-learning algorithm, and finds strong Atari performance alongside theoretical stability differences between evaluation and control.

  • Problem

    Value distributions had mainly been studied for specific purposes such as risk-sensitive behavior, leaving their broader role in reinforcement learning open.

  • Method

    The paper analyzes distributional Bellman operators and uses Bellman updates with projection to learn approximate value distributions in a categorical DQN architecture.

  • Results

    C51 significantly outperforms DQN, Double DQN, Dueling, and Prioritized Replay on the evaluated Atari games, surpassing the then-current state of the art by a large margin in several games.

  • Takeaways & Limitations

    Learning value distributions is reported as a powerful notion that surpasses most previous Atari 2600 gains without further algorithmic adjustments.

  • Takeaways & Limitations

    The practical occurrence of the theoretical control-setting instabilities remains unresolved, and the paper leaves more stable stochastic-policy behavior as an open question.

Abstract

from arXiv · show

In this paper we argue for the fundamental importance of the value distribution: the distribution of the random return received by a reinforcement learning agent. This is in contrast to the common approach to reinforcement learning which models the expectation of this return, or value. Although there is an established body of literature studying the value distribution, thus far it has always been used for a specific purpose such as implementing risk-aware behaviour. We begin with theoretical results in both the policy evaluation and control settings, exposing a significant distributional instability in the latter. We then use the distributional perspective to design a new algorithm which applies Bellman's equation to the learning of approximate value distributions. We evaluate our algorithm using the suite of games from the Arcade Learning Environment. We obtain both state-of-the-art results and anecdotal evidence demonstrating the importance of the value distribution in approximate reinforcement learning. Finally, we combine theoretical and empirical evidence to highlight the ways in which the value distribution impacts learning in the approximate setting.

1. Introduction

The paper argues that reinforcement learning should model the full random return distribution, not only its expectation, and develops theoretical and algorithmic consequences of this perspective.

  • Motivation: The paper shifts attention from the value Q, the expected return, to the random return Z and its value distribution.The value distribution is characterized recursively through the reward, next state-action, and next random return.
  • Motivation: Earlier work used value distributions mainly for parametric uncertainty, risk-sensitive algorithms, or theoretical analysis, whereas this paper argues for a central role in reinforcement learning.
  • Theory: For a fixed policy, the distributional Bellman operator is a contraction in the maximal Wasserstein metric, but not in total variation, Kullback-Leibler divergence, or Kolmogorov distance.
  • Theory: In control, the distributional optimality operator is not a contraction in any metric over distributions, despite contraction in expected value.The paper presents this instability as evidence favoring algorithms that model effects of nonstationary policies.
  • Algorithmic motivation: The distributional Bellman operator preserves multimodality, and the authors argue that learning approximate distributions can improve learning behavior and mitigate nonstationary-policy effects.
  • Empirical motivation: Modelling value distributions in a DQN agent produced considerably increased performance across Atari 2600 games and state-of-the-art performance on some games.

2. Setting

The paper formulates reinforcement learning as a discounted Markov decision process in which Bellman operators describe expected returns and the behavior of standard learning algorithms.

  • Markov Decision Process: The environment is modeled as a time-homogeneous Markov Decision Process (X, A, R, P, γ), with states, actions, transition kernel, discount factor, and reward function.
  • Returns and values: The return Zπ is the sum of discounted rewards along an agent’s trajectory, while Qπ gives its expected return after taking an action and then following policy π.
  • Returns and values: At each step, the state is sampled from the transition kernel and the action is sampled from the policy, starting from the specified initial state-action pair.
  • Bellman equations: Bellman’s equation describes the value function, and reinforcement learning commonly seeks policies that maximize the return.
  • Bellman equations: The optimality equation has a unique fixed point Q*, the optimal value function, associated with the set of optimal policies.
  • Bellman operators: The Bellman and optimality operators describe expected behavior of algorithms such as SARSA and Q-Learning, and repeated application converges exponentially to Qπ or Q*.

3. The Distributional Bellman Operators

The paper replaces expected returns with value distributions and analyzes their Bellman operators. Policy evaluation is contractive under a Wasserstein metric, whereas control exhibits substantial distributional instability despite convergence toward optimal distributions.

  • Value distributions map state-action pairs to distributions over random returns rather than their expectations.
  • The distributional Bellman equation combines randomness from rewards, transitions, and the next-state value distribution.
  • Policy evaluation: For a fixed policy, the distributional Bellman operator is a γ-contraction in the maximal Wasserstein metric, yielding convergence to the unique value-distribution fixed point.
  • Policy evaluation: The contraction guarantee depends on the metric: the policy-evaluation operator is not contractive in total variation, Kullback-Leibler divergence, or Kolmogorov distance.
  • Control: Control iterates nevertheless converge weakly to the set of optimal value distributions, while the mean can converge exponentially faster than the distribution itself.
  • Control: In control, the optimality operator is not a contraction in any distributional metric, although it remains contractive in expected value.
  • Control: The control operator can lack a fixed point, fail to converge even when one exists, and increase distributional distance under updates.
  • Control: The analysis leaves open whether stochastic policies could produce more stable control behavior.

4. Approximate Distributional Learning

The proposed algorithm learns approximate value distributions using a discrete categorical representation and a projected sample Bellman update. This converts distributional learning into a tractable classification problem while addressing disjoint supports and sample-transition constraints.

  • The algorithm applies the distributional Bellman optimality operator while selecting a rich class of parametric approximating distributions.
  • Value distributions are represented by N discrete atoms between VMIN and VMAX, with probabilities produced by a parametric model.
  • The discrete representation is both highly expressive and computationally friendly.
  • Because the Bellman update and model distribution usually have disjoint supports, directly minimizing Wasserstein distance is inconvenient for sample-based learning.
  • Categorical algorithm: The method projects each sample Bellman update onto the fixed atom support by distributing probability to neighboring atoms, reducing the update to multiclass classification.
  • The categorical algorithm minimizes a cross-entropy loss, while the Bernoulli algorithm provides a one-parameter alternative when N = 2.

5. Evaluation on Atari 2600 Games

Categorical DQN models value distributions on Atari games, where more atoms generally improve performance and C51 achieves strong results across deterministic and stochastic settings.

  • Algorithm: The Categorical DQN architecture outputs atom probabilities instead of action-values while retaining DQN’s training regime and ε-greedy action selection over expected values.The study used VMAX = −VMIN = 10 and a target network.
  • Learned value distributions: The distributional update separates low-value losing events from higher-value outcomes, as illustrated by distinct SPACE INVADERS action distributions.Actions that release the laser too early assign substantial probability to the terminal value 0, unlike safer actions.
  • Learned value distributions: The learned distributions can remain multimodal and often resemble Gaussians despite the deterministic ALE, which the authors attribute to discount-induced diffusion.PONG distributions also show two modes reflecting intrinsic randomness in reward timing.
  • Varying the number of atoms: Using too few atoms can produce poor behavior, while increasing the atom count improved performance in the five training games.The 51-atom version outperformed DQN in all five games and reached state-of-the-art performance in SEAQUEST; N = 51 trained at roughly 75% of DQN’s speed.
  • State-of-the-Art Results: C51 significantly outperformed DQN, Double DQN, Dueling, and Prioritized Replay, with especially strong results on SEAQUEST and sparse-reward games.The authors suggest value distributions may better propagate rarely occurring events.
  • State-of-the-Art Results: Within 50 million frames, C51 outperformed fully trained DQN on 45 of 57 games, suggesting that 200 million frames may be unnecessary for ALE evaluation.The comparison was averaged over three seeds.
  • State-of-the-Art Results: In stochastic execution, C51 achieved mean and median score improvements of 126% and 21.5%, respectively, relative to the random and DQN agents.The ALE rejected the selected action with probability p = 0.25 in this setting.

6. Discussion

The discussion argues that explicitly learning value distributions improves approximate reinforcement learning by addressing instability, preserving richer return information, and providing useful inductive biases.

  • Learning value distributions helped surpass most previous Atari 2600 gains without further algorithmic adjustments.
  • Why does learning a distribution matter?: Distributional learning may reduce chattering by averaging effects of the unstable Bellman optimality operator under function approximation.
  • Why does learning a distribution matter?: Explicitly modelling return distributions provides more stable learning targets when reward timing is difficult to predict.
  • Why does learning a distribution matter?: The approach supplies auxiliary predictions about the probability of each return value, whose accuracy is coupled to agent performance.
  • Why does learning a distribution matter?: Bounded support acts as an inductive-bias hyperparameter by treating returns beyond VMIN and VMAX as equivalent.
  • Why does learning a distribution matter?: The authors attribute some empirical performance to the tractability of minimizing KL divergence between categorical distributions.
  • The discussion concludes that algorithms should account for distributions in their theoretical and practical design.

Erratum

The erratum corrects C51’s reported mean score because one Atlantis evaluation episode exceeded the standard 30-minute cap; the corrected result remains strong.

  • 701% is the corrected C51 mean score, replacing the incorrectly reported 1010% and remaining higher than comparable baselines.
  • 178% is C51’s unchanged median score after correcting the mean-score calculation.
  • Capping Atlantis episodes at 30 minutes changed its human-normalized score from 22824% to 5199%, materially affecting the outlier-sensitive mean.

A. Related Work

The related-work discussion places distributional reinforcement learning within prior work on risk sensitivity, uncertainty, higher moments, probabilistic formulations, and return-distribution modelling.

  • Prior distributional Bellman research studied return distributions mainly for risk-sensitive reinforcement learning, using parametric and nonparametric approaches.
  • Return variance has been studied for policy evaluation, risk-sensitive actor-critic design, and variance-constrained optimal control.
  • Earlier work used Gaussian value-distribution approximations and probabilistic models to represent uncertainty about value functions or their parameters.
  • Foundational studies established distributional or higher-moment Bellman equations and moment-based optimality criteria.
  • Probabilistic reinforcement-learning frameworks embed returns in graphical models or optimize stationary distributions under transition constraints.
  • Compress and Control describes value functions by learning return distributions with density models.

B. Proofs

The proofs establish metric properties and contraction results for distributional operators, then analyze control-setting convergence, including convergence toward optimal value distributions despite a lack of contraction.

  • The partition lemma decomposes random variables across mutually exclusive events while preserving conditional structure.
  • The Wasserstein-based bound can be strict, and the corresponding gradients can differ.
  • The proof constructs a metric over value distributions and uses its triangle inequality and Wasserstein properties to bound distributional Bellman updates.
  • For policy evaluation, the distributional operator contracts under the selected metric, yielding exponential convergence of expected values to Q∗.
  • Convergence in the control setting: In control, iterates converge in a weak sense to the set of optimal value distributions, uniformly when the state space is finite.
  • Convergence in the control setting: The proof shows that greedy policies eventually select optimal actions at every state and that recursively defined solved-state sets eventually contain every state.
  • A fixed point of the distributional optimality operator does not by itself guarantee convergence of the iterates to that fixed point.
  • A nonstationary policy can produce a return distribution unavailable under any stationary mixture, motivating convergence toward nonstationary optimal value distributions.

C. Algorithmic Details

The categorical algorithm applies Bellman updates separately to distributional atoms, then projects each update onto neighboring support atoms. Training follows DQN closely but uses Adam and tuned VMAX and ϵadam hyperparameters.

  • The training regime closely follows DQN but uses Adam instead of RMSProp for gradient rescaling.
  • The final results use hyperparameters selected by evaluating VMAX and ϵadam over five training games.The considered VMAX values were {3, 10, 100}, and ϵadam values were {1/L, 0.1/L, 0.01/L, 0.001/L, 0.0001/L}, with L = 32.
  • The categorical algorithm applies the Bellman update to each atom separately and projects it into the two nearest atoms on the original support.Terminal transitions use γt = 0.

D. Comparison of Sampled Wasserstein Loss and Categorical Projection

CliffWalk experiments compare sampled Wasserstein optimization with categorical projection against ground-truth return distributions. The categorical method captures the true distribution more accurately, while sampled Wasserstein optimization can converge to inferior fixed points or local minima.

  • The distributional Bellman operator is a γ-contraction for a fixed policy, and CliffWalk experiments empirically test convergence under sampled Wasserstein loss and categorical projection.The policy follows the safe path but has a 10% chance of taking another action uniformly at random.
  • Ground-truth return distributions are estimated from 10000 Monte-Carlo rollouts per state and used to evaluate discrete approximations.The experiments compare supervised targets with sampled one-step distributional Bellman targets.
  • Some approximation error is unavoidable because a small probability of larger negative returns lies beyond the bounded support.The passage states that this effect is relatively negligible in the experiments.
  • The categorical algorithm minimizes the Wasserstein metric in both supervised and sampled Bellman settings, whereas stochastic-gradient Wasserstein optimization is flawed and can reach different local minima.
  • Categorical projection captures the true return variance more accurately than the sampled Wasserstein fixed point across five safe-path states.

E. Supplemental Videos and Results

Supplemental figures document C51 training videos, relative performance, raw evaluation scores, and selected Atari frames. The frame examples illustrate action-value distributions that encode fatal outcomes, uncertainty, and likely success or failure.

  • Figure 12 tracks the number of Atari games where training performance exceeds fully trained DQN and human baselines.Error bands show standard deviations, and averages are computed over games.
  • Supplemental videos show the C51 agent during training on various Atari 2600 games.
  • Figure 14 reports raw scores across all games after 30 no-op actions and compares them with reference values from Wang et al. (2016).
  • FREEWAY depicts the agent differentiating action-value distributions under pressure.
  • Q*BERT distributions predict unrecoverably fatal actions and show steep consequences for wrong actions.The figure also depicts a frame where the agent has made a huge mistake.
  • SEAQUEST shows bimodal distributions alongside situations that might or definitely will hit the fish.
  • SPACE INVADERS illustrates multimodal uncertainty, clearer action differences, uncertain survival, and certain success or demise.
Loading 1707.06887v1…