Technical Summary: Quotient-Categorical Representations for Bellman-Compatible Average-Reward Distributional Reinforcement Learning

1. Problem Statement

Distributional Reinforcement Learning (DRL) has achieved significant success in discounted Markov Decision Processes (MDPs), where the contraction geometry of the Bellman operator facilitates analysis and algorithm design. However, the average-reward setting presents a fundamental structural challenge for distributional modeling.

In average-reward MDPs, evaluation relies on two objects: the gain (long-run average reward) and the bias (relative value function). The bias satisfies the Poisson equation, which is the average-reward analogue of the Bellman equation. A critical property of the Poisson equation is that its solutions are unique only up to an additive constant. This translation invariance creates an obstruction for standard distributional modeling: a state-indexed family of probability laws on $\mathbb{R}$ cannot be directly defined with a fixed origin, because translating every state's bias distribution by the same constant should not alter the learned object.

Existing average-reward DRL work often focuses on refining the reward criterion itself or uses scalar representations. This paper addresses the gap in constructing a Bellman-compatible distributional analogue of the bias that respects the additive-constant symmetry, enabling distributional dynamic programming and temporal-difference (TD) learning in the average-reward regime.

2. Methodology

2.1 Quotient-Space Formulation

The authors propose modeling state-indexed bias laws modulo common translations.

• Quotient Space: Let $\mathcal{F}_S$ be the set of state-indexed collections of probability laws on $\mathbb{R}$. An equivalence relation $\approx$ is defined such that $\eta \approx \zeta$ if $\zeta$ is a common translation of $\eta$. The resulting space $\mathcal{F}_S / \approx$ preserves the additive-constant symmetry of the Poisson equation.
• Distributional Operator: The average-reward distributional operator $T$ is defined on $\mathcal{F}_S$. The paper proves that $T$ commutes with common translations, thereby inducing a well-defined operator on the quotient space $\mathcal{F}_S / \approx$.

2.2 Categorical Parameterization

To make the quotient space computationally tractable, the authors introduce a quotient-categorical representation:

• Representation: Each state's bias law is represented by a categorical distribution on a fixed grid $\Theta = \{\theta_1, \dots, \theta_d\}$.
• Shift Equivariance: A representative family is defined as $\eta_{p,c} = \sum p_k \delta_{\theta_k + c}$, where $c$ is an arbitrary shift. The coefficient vector $p$ uniquely characterizes the quotient class $[\eta_{p,c}]$, independent of $c$.
• Projection: A shift-equivariant categorical projection $\Pi_{\Theta_c}^C$ is defined. The paper establishes that projecting a translated law onto a translated grid yields the same coefficient vector as projecting the original law onto the original grid.

2.3 Projected Operator and Non-Expansiveness

The authors define a projected operator $G_g$ on the categorical coefficients, which corresponds to applying the distributional operator and projecting back onto the categorical grid.

• Metric: Convergence is analyzed using the coordinate Cramér metric ($\ell_{\Theta, C, \infty}$), which coincides with the standard Cramér metric on the representative laws.
• Key Property: Unlike discounted DRL operators which are contractive, the projected operator $G_g$ in the average-reward setting is non-expansive. The paper proves that $G_g$ is well-defined, non-expansive in the coordinate Cramér metric, and admits fixed points (via Brouwer's fixed-point theorem).

2.4 Stochastic Approximation (SA) Recursions

Since the operator is non-expansive rather than contractive, standard contraction arguments do not apply. The authors analyze sampled recursions as asynchronous relaxations of the projected operator:

1. Centered Recursion: Assuming the gain $\bar{r}^\pi$ is known, they define a one-state TD update (Stochastic Krasnosel'skiĭ–Mann iteration) using centered rewards $B = R - \bar{r}^\pi$. The mean-field map of this recursion is shown to be non-expansive and share fixed points with $G$.
2. Coupled Recursion (Unknown Gain): When the gain is unknown, the recursion is augmented with an online gain estimator $g_k$. The authors define a product metric $d_\lambda$ combining the categorical residual and the gain error. They prove that the coupled update map is non-expansive in this product metric and that its fixed points correspond to the correct quotient-categorical bias laws and the true gain.

3. Key Contributions

1. Quotient-Space Foundation: The paper establishes that the quotient space of state-indexed law families by common translation is the natural domain for a distributional analogue of the scalar bias.
2. Well-Defined Projected Operator: It defines an exact quotient-space operator and a quotient-categorical projection algorithm, proving the projected operator is well-defined, non-expansive, and possesses fixed points.
3. Convergence Guarantees:
   • For centered rewards (known gain), the paper proves almost sure (a.s.) convergence and finite-iteration residual bounds for the TD recursion under both i.i.d. and Markovian sampling.
   • For unknown gains, it introduces a coupled recursion with an online gain estimator, proving non-expansiveness and Markovian convergence of the joint scheme.
4. Structural Insight: The work highlights a structural contrast between ideal quotient distributions (where centering constants induce only common translations) and practical fixed-grid categorical representations (where projection can alter coefficients, necessitating correct centering).

4. Experimental Results

The paper validates its theoretical claims through two experiments:

• Five-State MDP: A tabular experiment comparing exact KM iterations, centered i.i.d. SKM, centered Markovian SKM, and coupled gain-estimated SKM.
  • Results confirm that exact and stochastic recursions converge to the same fixed points.
  • The coupled recursion successfully learns both the bias distribution and the gain.
  • An ablation study demonstrates that using an incorrect centering constant ($g=0$ instead of $\bar{r}^\pi$) leads to convergence to a degenerate solution, validating the necessity of correct centering in fixed-grid representations.
• Continuous-State Pendulum-v1: A proof-of-concept using neural function approximation.
  • The proposed gain-centered quotient-categorical critic significantly outperforms a naive raw-reward critic in terms of projected-operator residuals.
  • The coupled online-gain recursion successfully decomposes the error into categorical and gain components, consistent with theoretical bounds.

5. Significance and Claims

The paper claims to put average-reward distributional bias estimation on firm foundations. Its primary significance lies in:

• Resolving the additive-constant ambiguity of the scalar bias by introducing a quotient-categorical formulation.
• Providing the first Bellman-compatible distributional analogue of the bias that is suitable for dynamic programming and TD-learning arguments in the average-reward setting.
• Establishing convergence guarantees for non-expansive maps in the context of distributional RL, bridging the gap between theoretical fixed-point methods for non-expansive operators and practical TD algorithms.

The authors note that their objective is complementary to prior work refining the average-reward criterion itself; instead, they refine the bias as the object of learning. The theory is currently focused on finite-state, fixed-policy evaluation with fixed categorical supports, with richer function approximation and control settings identified as natural next steps.