Quantile Geometry Regularization for Distributional Reinforcement Learning

TL;DR

This paper introduces RQIQN, a novel distributional RL method that enhances quantile estimation robustness using Wasserstein regularization, improving performance in risk-sensitive tasks.

Contribution

The paper proposes RQIQN, a Wasserstein distributionally robust enhancement for quantile regression in distributional RL, addressing distributional degeneration without altering the core value objective.

Findings

RQIQN outperforms existing quantile-based methods in risk-sensitive navigation.

RQIQN achieves superior results in Atari game benchmarks.

The method effectively regularizes quantile geometry, preventing distribution collapse.

Abstract

Quantile-based distributional reinforcement learning methods learn return distributions through sampled quantile regression, but their bootstrapped target quantiles may induce distorted or degenerate distribution estimates. We propose Robust Quantile-based Implicit Quantile Networks (RQIQN), a lightweight Wasserstein distributionally robust enhancement boosted from a quantile estimation perspective. We first reinterpret a snapshot of IQN loss as a collection of local empirical quantile estimation problems over sampled current fractions. We then robustify each local slot with a Wasserstein distributionally robust quantile estimation formulation, yielding a closed-form, fraction-dependent correction to the Bellman target. This correction directly addresses distributional degeneration: its median antisymmetry preserves the risk-neutral quantile average, while its monotonicity enlarges upper-lower quantile gaps and counteracts collapsed distributional spread. RQIQN thus regularizes quantile geometry without changing the underlying value objective or requiring additional sample set reconstruction. Finally, we empirically show that the proposed RQIQN outperforms other existing quantile-based distributional reinforcement learning algorithms in risk-sensitive navigation and Atari games.