Convergence Theorems for Entropy-Regularized and Distributional Reinforcement Learning

AI summary

Overview

Research area: Reinforcement learning theory, at the intersection of entropy-regularized RL (ERL) and distributional RL (DRL).

Technical level: Advanced. The paper works with discounted MDPs over Borel state and action spaces, probability kernels, generalized product measures and disintegration, Wasserstein metrics on spaces of kernels, and contraction arguments on non-separable function spaces.

Scope (one sentence): The paper develops a "temperature decoupling gambit" for vanishing-temperature entropy-regularized RL that guarantees convergence of policies and their return distributions to an interpretable, diversity-preserving reference-optimal limit, and uses this to define and analyze distributional ERL (DERL) with convergent iterates.

What This Paper Is About

In a generic MDP many policies are optimal and all of them share the same action-value function, so standard RL says nothing about which states a learned policy visits, which actions it plays, or what other statistics (such as variance) its returns have. Entropy regularization fixes this ambiguity for any positive temperature τ by making exactly one policy optimal, but outside tabular MDPs it is unclear whether the τ-optimal policy and its derived objects converge as τ → 0. The paper's goal is a policy-optimization framework that guarantees convergence to a particular optimal policy — one that quantifiably preserves state-wise action diversity — and that also yields an algorithm for estimating the associated return distribution to arbitrary accuracy.

Key Contributions

1. The temperature decoupling gambit. A scheme that specifies a second, amplified temperature σ = σ(τ) with σ/τ → 0 as τ → 0, and constructs policies π^{τ,σ} := 𝒢_τ q_σ^⋆ — Boltzmann–Gibbs policies at temperature τ built from action-values estimated at temperature σ. These policies are deliberately suboptimal for the τ-ERL objective, but they are the ones that converge.

2. A reference-optimality framework. The paper introduces the Bellman reference-optimality operator ℬ^⋆_ref, whose unique fixed point q^⋆_ref generally satisfies q^⋆_ref ≤ q^⋆, and defines a policy π as reference-optimal if q^π ≥ q^⋆_ref. It also defines the optimality-filtered reference policy π^{ref,⋆}, which reweights the reference policy π^ref onto the set of actions attaining the essential supremum of q^⋆ under π^ref.

3. Convergence theorems for policies and return distributions. Theorem 3.9 shows π^{τ,σ}_x → π^{ref,⋆}x for all states x, in total variation when the action space is discrete and weakly when it is continuous. Theorem 3.10 shows the return distribution functions ζ^{τ,σ} of the decoupled policies converge to ζ^{π^{ref,⋆}} in the metric d{p;p′,ω} for any p, p′ ∈ [1, ∞) and ω ∈ 𝒫(𝖷 × 𝖠).

4. Distributional ERL (DERL). A new branch of DRL built on a soft distributional Bellman operator 𝒯_τ^π and a soft distributional optimality operator 𝒯_τ^⋆. The paper proves 𝒯_τ^π is a γ-contraction in d̄_p under an integrability condition, proves the algebraic identity 𝒬𝒯_τ^⋆ = ℬ_τ^⋆𝒬, and proves in Theorem 4.5 that iterates of 𝒯_τ^⋆ converge — in contrast to iterates of the standard distributional optimality operator 𝒯^⋆, which are known not to converge in the control setting.

Main Findings

• The temperature decoupling gambit changes the limit. Even in tabular MDPs, where τ-optimal policies already converge (Theorem 2.5), the limiting policy produced by the gambit is distinct from the limit found otherwise. The gambit's limit preserves, quantifiably, more state-wise action diversity.

• The limit is the optimality-filtered reference policy. Under Assumption 3.4 — which holds when the action space is discrete and π^ref_x is uniformly lower bounded — the decoupled policies π^{τ,σ}_x converge to π^{ref,⋆}_x as τ → 0. When π^ref is the uniform policy, π^{ref,⋆}_x is the uniform distribution over the optimal actions in state x, i.e., the policy that samples all optimal actions uniformly.

• Action-values converge monotonically to a reference-optimal fixed point. Theorem 3.2 shows q^⋆_τ → q^⋆_ref monotonically as τ → 0. This implies optimal policies generally cannot be recovered by plain vanishing-temperature limits in ERL.

• Reference-optimality is a "skyline" for performance. When q^⋆_ref ≠ q^⋆, optimal actions form a measure-zero set, so even rich policy classes such as neural-network-parameterized Gaussian policies or diffusion policies will not sample optimal actions with probability 1. The paper's unique fixed point for the reference-optimality equation upper-bounds the RL performance of τ-optimal policies in general.

• A total-variation control on Boltzmann–Gibbs policy differences. Theorem 3.6 bounds ‖(𝒢_τ q)x − (𝒢_τ q′)x‖TV by min{ √(τ^{−1}‖q(x,·) − q′(x,·)‖{L^∞(π^ref_x)}), ½ sinh(4τ^{−1}‖q(x,·) − q′(x,·)‖{L^∞(π^ref_x)}) }, and by ((2e − 3)/4)·τ^{−1}‖q(x,·) − q′(x,·)‖{L^∞(π^ref_x)} when that sup-distance is less than τ/2. The convergence of the gambit follows from the derived inequality lim sup_x ‖(𝒢_τ q_σ^⋆)_x − (𝒢_τ q^⋆_ref)_x‖_TV ≲ −lim (σ/τ) log p_ref, which is where the requirement σ/τ → 0 originates.

• Tabular limits are the minimizer of the regularizer among optimal policies. Theorem 2.5: if the regularizer is not identically +∞ on the set of optimal policies (Assumption 2.4), the sequence (μ^⋆τ){τ>0} of τ-optimal occupancy measures has a unique setwise limit μ^⋆0, which minimizes ℛ over arg sup{𝒪(ν_0)} 𝒥_0.

• Joint-distribution convergence is not policy convergence. Remark 2.6 observes that convergence of a sequence of joint distributions does not imply convergence of the associated conditional distributions, which is why the tabular result is stated for occupancy measures and not for policies outside the tabular setting.

• Numerical demonstration. On a tristate MDP with two actions (a_1 in blue, a_2 in green), γ = 0.9, uniform reference policy, and σ = τ², soft Q-learning estimates of π^{τ,⋆} and π^{τ,σ} differ in state x_0: π̂^{τ,⋆}{x_0} collapses to δ{a_1} as τ → 0, while π̂^{τ,σ} converges to π^{ref,⋆}. The paper describes the two limits as emphasizing different notions of diversity: the limit of π^{τ,⋆} filters out optimal actions to play actions more uniformly on average with respect to long-term state occupancy, while the limit of π^{τ,σ} maximizes state-wise action diversity.

• Reported as not established elsewhere. The paper states that the only comparable well-posedness result it is aware of is for tabular MDPs and establishes convexity on 𝒪(ν_0) rather than on all of 𝒫(𝖷 × 𝖠); no benchmark datasets, dataset sizes, or numeric performance comparisons are reported in the provided content.

Methodology in Plain English

The authors start from a standard trick: penalize a policy by how far it diverges (in KL divergence) from a fixed reference policy π^ref. Adding that penalty at temperature τ makes the optimal policy unique and makes it a Boltzmann–Gibbs policy — a softmax-like reweighting of the reference policy by action-values.

The problem is that at temperature τ, the log-probabilities of the Boltzmann–Gibbs policy are amplified by 1/τ. So even though the action-values q^⋆_τ themselves converge nicely to a limit q^⋆_ref, the induced policy may not, because tiny action-value errors get blown up by the amplification.

The gambit is to decouple the two roles temperature plays. The policy is built at temperature τ, but the action-values fed into it are estimated at a smaller temperature σ(τ) with σ/τ → 0. This makes the policy "look like" a Boltzmann–Gibbs policy applied to a fixed potential q^⋆_ref, whose vanishing-temperature behaviour is well understood. The policy played at any finite τ is not optimal for the τ-regularized objective — the authors compare this deliberate sacrifice to a chess gambit — but the long-run limit is a well-defined, diverse, reference-optimal policy.

For the distributional half, the authors generalize distributional dynamic programming by adding the KL penalty into the reward, defining a soft distributional Bellman operator (for evaluation) and a soft distributional optimality operator (for control). They show the evaluation operator is a γ-contraction under a boundedness condition, and that control iterates converge — which standard distributional control iterates do not. A small tabular experiment with soft Q-learning illustrates the policy-level difference between the coupled and decoupled schemes.

Why This Matters

Impact on research. The paper supplies convergence guarantees where the field previously had none: for vanishing-temperature entropy-regularized policies outside tabular MDPs, and for distributional control, where existing distributional methods are known not to produce convergent iterates. It also reframes what the "right" target of regularization should be by introducing reference-optimality and a Bellman-like equation whose unique fixed point upper-bounds the RL performance of τ-optimal policies generally. The authors position reference-optimality as the realistic performance ceiling for continuous-control methods when q^⋆_ref ≠ q^⋆ and optimal actions form a measure-zero set.

Real-world applications (the paper motivates distributional RL by the need to understand the distribution over returns, not just its mean):

• Safety-critical applications, explicitly cited in the paper as a setting where understanding the distribution over returns is crucial, and where which states a policy visits and which actions it plays matters.
• Continuous-control settings that use rich policy classes such as neural-network-parameterized Gaussian policies or diffusion policies, which the paper argues cannot sample measure-zero optimal action sets.
• Settings where diversity of behaviour at optimality is desirable — the paper notes that the gambit's limit does not discriminate between optimal actions and maximizes state-wise action diversity, and the uniform-reference case samples all optimal actions with equal probability.
• Risk-aware decision making, which follows from the paper's emphasis on return distributions and statistics such as variance rather than expected value alone.

Industry relevance. The results give practitioners a principled way to choose among the many optimal policies of a given MDP — controlling which actions are played at optimality — and a recipe (temperature decoupling, plus DERL operators) for estimating the associated return distribution to arbitrary accuracy. The temperature schedule guidance is concrete: the paper notes that σ(τ) = τ³ makes 𝒢_τ q^⋆_σ behave more like 𝒢_τ q^⋆_ref than σ(τ) = τ² does.

Future Directions

1. Continuous action spaces under weaker assumptions. Assumption 3.4 is standard and easy to satisfy in the discrete case when π^ref_x is uniformly lower bounded, but the paper states it is more difficult to guarantee when the action space is continuous. Relaxing it would widen the applicability of Theorem 3.9.
2. Policy convergence outside the tabular setting. Remark 2.6 shows that occupancy-measure convergence does not imply policy convergence in continuous MDPs, and the authors say they are unaware of any result of that type. Whether the structure of 𝒪(ν_0) permits policy convergence remains open.
3. Computational guarantees for return-distribution approximation. Theorem 3.10 establishes convergence of return distribution functions but, as the authors state, does not yet provide an algorithm for approximating ζ^⋆; this motivates the DERL section. The statements of Theorems 4.6 and 4.7, which the paper lists among its main results for convergent iterative schemes, are cut off in the provided content and are not reported here.
4. Extension beyond Euclidean state and action spaces. A footnote states the authors expect many results to extend to Polish spaces; verifying this is left open.

Target Audience

Theoretical reinforcement learning researchers and graduate students comfortable with measure-theoretic probability, Markov decision processes, and contraction-based dynamic programming arguments. It is most directly useful to those working on entropy-regularized RL, KL-regularized policy optimization, distributional RL and distributional dynamic programming, and the theory of optimal-policy selection. Practitioners in safety-critical or risk-sensitive RL will find the framing (return distributions, interpretable limits, schedule conditions such as σ/τ → 0 and σ = τ²) relevant, though the paper's development is abstract rather than empirical.

Authors’ abstract

In the pursuit of finding an optimal policy, reinforcement learning (RL) methods generally ignore the properties of learned policies apart from their expected return. Thus, even when successful, it is difficult to characterize which policies will be learned and what they will do. In this work, we present a theoretical framework for policy optimization that guarantees convergence to a particular optimal policy, via vanishing entropy regularization and a temperature decoupling gambit. Our approach realizes an interpretable, diversity-preserving optimal policy as the regularization temperature vanishes and ensures the convergence of policy derived objects--value functions and return distributions. In a particular instance of our method, for example, the realized policy samples all optimal actions uniformly. Leveraging our temperature decoupling gambit, we present an algorithm that estimates, to arbitrary accuracy, the return distribution associated to its interpretable, diversity-preserving optimal policy.