Learning the Supports for Categorical Critic in Reinforcement Learning

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.

• The paper introduces DySEL, a method that adaptively learns critic support intervals to balance resolution and truncation biases.
• It formulates support learning as a min-max Lagrangian optimization that enforces a mass coverage constraint while minimizing support width.
• Empirical results on DeepMind Control Suite show improved IQM returns, especially in high-variance tasks, validating the dynamic support adaptation approach.

Learning Dynamic Supports for Categorical Critics in Reinforcement Learning

Introduction

Value estimation lies at the core of actor-critic reinforcement learning, where accurate critic targets are paramount for stable and efficient policy improvement. While conventional TD methods pursue scalar regression using mean-squared error (MSE), distributional RL methods model the return distribution explicitly, achieving superior performance in diverse domains. Recent developments in classification-based value learning, such as the Gaussian Histogram Loss (HL-Gauss), convert value regression to a multi-class classification problem, yielding favorable gradient properties and empirical gains. However, a persistent issue for categorical approaches is the necessity to pre-specify the discrete support interval for predicted values, which is infeasible for non-stationary, stochastic returns inherent in RL. The paper "Learning the Supports for Categorical Critic in Reinforcement Learning" (2607.01880) introduces Dynamic Support Endpoint Learning (DySEL), a principled method for adaptively learning the critic support interval, theoretically motivated by its role as an upper bound term in the MSE Bellman error.

Theoretical Foundations and Formulation

A central insight is that the HL-Gauss loss for a categorical critic tightly upper-bounds the mean-squared Bellman error, with the tightness of this bound explicitly controlled by the width of the support interval. The theoretical analysis exposes a tension:

• Narrower supports provide higher bin resolution (reducing quantization bias) and tighter bounds on MSE but risk truncation bias if the support does not capture all target mass.
• Wider supports eliminate truncation but disperse resolution, leading to loose bounds and degraded gradient signals.

This motivates a constrained optimization objective: minimize support width (for tight error bounds), subject to the constraint that the interval covers most of the Gaussian-smoothed target mass. The authors map this into a min-max Lagrangian formulation, where the critic and support endpoints minimize a compound objective (support width + cross-entropy) penalized by a Lagrangian enforcing mass coverage:

$$\mathcal{L}(\theta, \phi, \lambda) = \alpha \max(|\nu_\text{min}|, |\nu_\text{max}|) + \frac{1}{\alpha} \mathcal{L}_{\text{CE}} + \lambda \left( (1 - Z) - \epsilon \right)$$

Here, $Z$ is the fraction of the target distribution within the support, and $\lambda$ is learned adversarially to enforce the coverage constraint.

Practical Algorithm: Dynamic Support Endpoint Learning (DySEL)

The DySEL algorithm parameterizes the lower and upper support endpoints, which are adapted during learning alongside the critic. A Lagrangian min-max game, solved via gradient-based optimization, ensures the supports remain as narrow as possible without truncating significant probability mass. The only method-specific hyperparameter is $\alpha$, which trades off coverage against support width.

Empirical Evaluation and Numerical Results

The experimental evaluation targets 11 DeepMind Control Suite tasks via a TD3 backbone. DySEL is benchmarked against TD3 and TD3+HL-Gauss (fixed support $[-100, 100]$), using 128 bins for all categorical critics. Robust IQM (inter-quantile mean) returns are reported across 10 seeds.

Figure 1: The IQM returns for each task in DM Control of our proposed DySEL versus baselines. Our proposed method TD3+DySEL remains competitive with TD3+HLG, whilst significantly outperforming in some humanoid tasks, showing the empirical benefit of learnt supports.

Across most tasks, DySEL matches or outperforms fixed-support HL-Gauss, and exhibits a marked advantage on high-variance environments such as humanoid-run, humanoid-stand, and humanoid-walk. This validates the hypothesis that dynamic supports enable better adaptation to non-stationary value targets, maintaining high effective resolution without truncation.

Figure 2: The evolution of support interval of DySEL during training. We can observe that some support intervals remain stable whilst others gradually increase/decrease as training progresses.

Support bounds show task-dependent evolution: some stabilize early (e.g., for low-variance tasks), others expand adaptively as agent performance and the return distribution shift, highlighting the necessity for online support adaptation.

Ablation experiments demonstrate that both the width penalty and the mass constraint are critical for stability. Omitting the mass constraint leads to degenerate support shrinkage and catastrophic truncation bias. Removing the width penalty causes the supports to diverge, collapsing most probability mass into few bins and destroying learning signal.

Figure 3: IQM return and its support interval evolution whilst removing components. No-mass means we remove the mass enforcing adversarial penalty. No-width means we remove the absolute width penalty. In both cases, we can show that both elements are critical for stability. For support interval plots on the right, y-axis is in sym-log scale, because the support interval goes to exceedingly large values and diverges without the width penalty.

Robustness Analyses

Support interval adaptation exhibits low sensitivity to initialization: even narrow initial intervals ($[-1,1]$) are robustly expanded when needed, although this can transiently destabilize value targets on high-variance tasks.

Figure 4: IQM and bounds evolution of different initialisation. For all tasks, the supports evolution are similar, but those initialised with smaller initialising supports have some instability.

The key hyperparameter, $\alpha$, consistently controls the support adaptation scheme: small $\alpha$ results in narrow, stable intervals, while larger $\alpha$ enables sustained expansion for tasks demanding high bin resolution over a wider target range.

Figure 5: IQM and bounds evolution of different alpha.

Discussion and Implications

The work provides a theoretically principled and operationally robust mechanism for adaptively learning support intervals in categorical critics. DySEL eliminates the need for laborious task-specific support tuning, directly addressing a longstanding limitation in practical distributional RL. Dynamic support adaptation is essential for regimes with evolving or initially unknown return distributions, such as continuous control and long-horizon planning.

While quantile-based distributional methods do not require explicit supports, they can suffer from quantile crossing and lack the inductive bias provided by the categorical approach. DySEL bridges this gap for categorical critics, making HL-Gauss and related losses significantly more deployable.

Intrinsic limitations remain: the min-max Lagrangian is sensitive to dual learning rates, and the learned supports are empirically symmetric, potentially leaving performance on highly skewed distributions suboptimal. Extensions to asymmetric supports and adaptive Gaussian widths may further enhance flexibility and applicability.

Conclusion

DySEL introduces a theoretically justified, practically effective approach for learning dynamic support intervals in categorical critics. By reframing support selection as a constrained optimization and casting it as a min-max game with automated mass coverage enforcement, DySEL achieves tight error bounds and high-resolution representation throughout the learning process. Empirical gains in high-variance and non-stationary environments underscore its value for modern RL. Future research directions include generalizing DySEL to asymmetric supports and integrating adaptive smoothing for further improvements in distributional value estimation.