Learning the Supports for Categorical Critic in Reinforcement Learning

2026-07-02

AI Summary

This work addresses the limitations of traditional distributional reinforcement learning, which relies on fixed, predefined support intervals and struggles to handle the non-stationarity and stochasticity of return distributions. The authors formulate value function estimation as a classification task and propose a method that jointly learns adaptive support boundaries and categorical representations. By dynamically adjusting the support bounds using a Gaussian-smoothed target distribution and introducing a novel objective function, they theoretically establish a tighter upper bound on the mean squared Bellman error. Integrating Gaussian histogram loss, the distributional Bellman operator, and an actor-critic framework, the approach enables end-to-end training. Empirical results demonstrate that the method matches or exceeds the performance of HL-Gauss on most continuous control tasks while achieving stable, self-adaptive support intervals.

Abstract

Value functions are an essential component in actor-critic based deep reinforcement learning (RL). Conventionally, these functions are trained as a regression task by minimising the mean squared error (MSE) relative to bootstrapped target values. Meanwhile, in distributional RL, a distribution of returns is modelled based on the distributional Bellman operator. This work investigates the Gaussian Histogram Loss (HL-Gauss), a recent approach that reframes value estimation as classification by encoding each scalar Bellman target as a Gaussian-smoothed categorical target. Despite its potential, applying histogram-based losses to RL presents inherent challenges, most notably the requirement to pre-define a fixed support interval, which is often complicated by the non-stationary and stochastic nature of target values typically found in RL tasks. In this work, we propose an approach that dynamically learns the lower and upper bounds of the support instead of assigning them beforehand. We derive an objective that jointly learns these bounds whilst learning the categorical representation of the scalar values, and we show that this objective forms an upper bound on the mean-squared Bellman error. Our theoretical analysis further shows that this bound is tighter than that of non-learned supports of HL-Gauss. Empirically, the proposed objective enables stable adaptation of the support interval and matches HL-Gauss-based actor-critic algorithms on most continuous-control tasks whilst improving on a subset, without requiring a pre-specified support interval.

Problem

• distributional reinforcement learning
• categorical critic
• support interval
• value estimation
• non-stationary targets

Innovation

• distributional reinforcement learning
• categorical critic
• learnable support
• Gaussian Histogram Loss
• adaptive bounds