[Paper Review] Distributional Reinforcement Learning with Quantile Regression

Will Dabney, Mark Rowland+2 more|arXiv (Cornell University)|Oct 27, 2017

Sports Analytics and PerformanceEconomics, Econometrics and Finance150 citations

TL;DR

This paper shows how to learn value distributions in reinforcement learning end-to-end under the Wasserstein metric by using quantile regression, introduces qr-dqn, and achieves state-of-the-art Atari results.

ABSTRACT

In reinforcement learning an agent interacts with the environment by taking actions and observing the next state and reward. When sampled probabilistically, these state transitions, rewards, and actions can all induce randomness in the observed long-term return. Traditionally, reinforcement learning algorithms average over this randomness to estimate the value function. In this paper, we build on recent work advocating a distributional approach to reinforcement learning in which the distribution over returns is modeled explicitly instead of only estimating the mean. That is, we examine methods of learning the value distribution instead of the value function. We give results that close a number of gaps between the theoretical and algorithmic results given by Bellemare, Dabney, and Munos (2017). First, we extend existing results to the approximate distribution setting. Second, we present a novel distributional reinforcement learning algorithm consistent with our theoretical formulation. Finally, we evaluate this new algorithm on the Atari 2600 games, observing that it significantly outperforms many of the recent improvements on DQN, including the related distributional algorithm C51.

Motivation & Objective

• Motivate modeling the full distribution of returns rather than only the mean in reinforcement learning.
• Bridge the theory-practice gap by enabling end-to-end optimization under the Wasserstein metric.
• Develop a practical algorithm (qr-dqn) that learns value distributions with quantile regression.
• Demonstrate superior performance on Atari 2600 benchmarks compared to prior distributional methods.

Proposed method

• Replace fixed-location, uniform-probability distributions (as in c51) with fixed-quantile locations and uniform weights, effectively estimating quantiles of the return distribution.
• Use quantile regression to minimize the Wasserstein-1 distance between the target and the predicted distribution, enabling unbiased stochastic gradient updates.
• Prove contraction properties for the combined quantile projection and distributional Bellman operator under the Wasserstein metric.
• Derive and implement a quantile regression TD (qrtd) for policy evaluation and a QR-DQN (qr-dqn) for control, with a quantile-Huber loss as an option.
• Adapt the DQN architecture to output N quantiles per action and train using the quantile regression loss instead of standard TD loss.
• Empirically validate on gridworld-like tasks and 57 Atari 2600 games, comparing to c51 and DQN variants.

Experimental results

Research questions

• RQ1 Can a distributional reinforcement learning algorithm operate end-to-end under the Wasserstein metric using quantile regression?
• RQ2 Does a quantile-based distributional representation, without projection steps, improve stability and performance over existing methods like c51?
• RQ3 Does qr-dqn achieve state-of-the-art performance on Atari 2600 benchmarks and how does it compare to prior distributional methods?
• RQ4 What are the theoretical contraction properties when combining quantile projection with the distributional Bellman operator?
• RQ5 How does quantile regression (with or without Huber smoothing) influence learning dynamics and final performance in Distributional RL?

Key findings

• A quantile-based distributional RL algorithm with uniform weights over fixed quantile locations converges under the Wasserstein distance to the distributional fixed point.
• The combined operator with quantile projection is a contraction in the infinity-Wasserstein metric, ensuring convergence.
• qr-dqn outperforms prior methods (including c51) on Atari 2600 benchmarks, achieving higher mean and median human-normalized scores.
• Using a quantile-Huber loss yields additional performance gains over the strict quantile loss.
• Empirical results show qrtd accurately minimizes the 1-Wasserstein distance to ground-truth distributions in a windy gridworld setting.