Learning the Supports for Categorical Critic in Reinforcement Learning

Abstract

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 approximating the categorical return distribution by framing the dynamic support learning problem as a constrained optimisation problem. We call our algorithm DySEL (Dynamic Support Endpoint Learning).

Problem

We addresses a key limitation in Categorical Value Learning [1]: histogram-based value losses such as HL-Gauss [2] [3] require a pre-defined fixed support interval, even though the true range of RL target values is usually unknown, non-stationary, and task-dependent. If the support interval is too narrow, the target distribution is truncated and important return information can be lost; if it is too broad, the fixed number of bins must cover a larger range, reducing resolution and weakening the critic’s learning signal. We argue that this support-selection problem directly affects both truncation bias and quantisation bias, making fixed supports an strong design choice for actor-critic RL with categorical critics.

Proposed Method

Our proposed method DySEL dynamically learns the lower and upper bounds of the categorical critic’s support interval instead of manually fixing them. We derive a connection showing that the mean-squared Bellman error is upper-bounded by the HL-Gauss objective, and that this bound becomes tighter when the support interval is narrower. However, the interval must still be wide enough to cover the target distribution’s probability mass. To balance these competing pressures, we formulates the support learning as a constrained optimisation problem: minimise the support width and distribution-matching loss (cross-entropy) while enforcing sufficient probability-mass coverage. We further implemented this constraed optimisation problem as a min-max Lagrangian objective, where the critic and support parameters minimise the loss while a Lagrange multiplier encourages expansion when too much probability mass leaks outside the support.

Experiment results

We evaluate on DM Control Suite tasks using TD3 [4] as the base algorithm, comparing standard TD3, TD3 with fixed-support HL-Gauss, and TD3+DySEL (proposed method). Across most tasks, TD3+DySEL matches or outperforms the baselines, with especially large gains on humanoid tasks, where the dynamically learned support provides finer value resolution and helps the critic distinguish subtle return differences.

How the dynamic supports evolve

We further plot how the supports evolve during training. The results also show that the learned supports adapt differently across tasks: some converge quickly to stable intervals, while others gradually expand during training.

References

[1] Jesse Farebrother et al., Stop Regressing: Training Value Functions via Classification for Scalable Deep RL. In ICML, 2024
[2] Ehsan Imani et al., Investigating the Histogram Loss in Regression In JMLR, 2026
[3] Scott Fujimoto et al., Addressing Function Approximation Error in Actor-Critic Methods. In ICML, 2018
[4] Marc G. Bellemare et al., A Distributional Perspective on Reinforcement Learning. In ICML, 2017