Distributional Reinforcement Learning with Dual Expectile-Quantile Regression

TL;DR

This paper addresses the instability and distributional collapse observed when using asymmetric $L_2$ losses (Huber) in distributional RL by proposing a dual expectile-quantile learning framework. It jointly learns $Z_\theta(s,a,\tau)$ as expectile statistics and a mapper $m_\phi(s,a,\tau)$ to generate quantile fractions, enabling efficient learning with distributional guarantees. The authors prove that the dual operator converges to the distributional Bellman operator in the limit of infinite quantile/expectile fractions and demonstrate practical effectiveness with IEQN, which matches the performance of huber-based IQN baselines on Atari-5 while avoiding distributional collapse. Empirically, the method shows non-collapsing, well-spread distributions and strong performance on both a toy MDP and large-scale Atari benchmarks, highlighting its potential for robust, distribution-aware RL in stochastic environments.

Abstract

Distributional reinforcement learning (RL) has proven useful in multiple benchmarks as it enables approximating the full distribution of returns and extracts rich feedback from environment samples. The commonly used quantile regression approach to distributional RL -- based on asymmetric $L_1$ losses -- provides a flexible and effective way of learning arbitrary return distributions. In practice, it is often improved by using a more efficient, asymmetric hybrid $L_1$-$L_2$ Huber loss for quantile regression. However, by doing so, distributional estimation guarantees vanish, and we empirically observe that the estimated distribution rapidly collapses to its mean. Indeed, asymmetric $L_2$ losses, corresponding to expectile regression, cannot be readily used for distributional temporal difference. Motivated by the efficiency of $L_2$-based learning, we propose to jointly learn expectiles and quantiles of the return distribution in a way that allows efficient learning while keeping an estimate of the full distribution of returns. We prove that our proposed operator converges to the distributional Bellman operator in the limit of infinite estimated quantile and expectile fractions, and we benchmark a practical implementation on a toy example and at scale. On the Atari benchmark, our approach matches the performance of the Huber-based IQN-1 baseline after $200$M training frames but avoids distributional collapse and keeps estimates of the full distribution of returns.

Figures (4)

• Figure 1: (a) Approximating a bimodal distribution with quantile and expectile regression. Quantile regression approximates the inverse CDF, albeit with high variance, especially on extreme values (left, blue curves). Expectiles converge very quickly to the expectile function (left, red curves). When training a mapper to generate quantiles from expectiles, quantile estimation becomes much more efficient (right). (b) Distributional RL with function approximation in a chain MDP with 4 states, and a bimodal reward distribution at the last state. The expectile function collapses as the temporal difference error propagates to previous states (left, red curves) while the quantile function is a poor approximation of the inverse CDF (left, blue curves). Our dual method solves both problems (right).
• Figure 2: Interquartile mean of the human normalized score of distributional RL agents on the Atari-5 benchmark with 5 random seeds per environment. Shaded areas correspond to the 25-th and 75-th percentiles of a bootstrap distribution. A rolling average with window size of $20M$ frames is performed to enhance readability.
• Figure 3: Toy Markov decision process.
• Figure 4: Comparison of estimated variance against observed variance of unfolding the greedy policy.

Theorems & Definitions (9)

• Lemma 0
• Theorem 1
• Corollary 1
• Lemma 1
• proof
• Theorem 1
• proof
• Corollary 1
• proof