Diversity-Preserving Optimal Policy

• The paper introduces a decoupling mechanism for temperature parameters that ensures diversity among optimal actions.
• It employs entropy regularization with a soft Q-function to prevent policy collapse on a single optimal action.
• Theoretical guarantees include convergence proofs and distributional analyses, enhancing robustness and safety in RL applications.

A diversity-preserving optimal policy is a reinforcement learning (RL) solution that seeks not only to maximize return but also to retain and actively promote diversity among high-performing policies or actions. Instead of converging to a single (possibly over-specialized) optimal solution, such frameworks seek to ensure that the mixture of learned behaviors or policies covers the full set of optimal (or near-optimal) solutions, thereby providing improved exploration, robustness, and adaptability. A principled approach to this problem is articulated in "Convergence Theorems for Entropy-Regularized and Distributional Reinforcement Learning" (Jhaveri et al., 9 Oct 2025), where diversity is guaranteed by carefully controlling the entropy regularization and by employing a temperature decoupling mechanism that ensures uniform support over optimal actions.

1. Entropy Regularization and Its Limitations

Classical entropy-regularized reinforcement learning (ERL) augments the expected reward objective with an entropy or KL-divergence penalty, yielding a "soft" optimal policy for each temperature τ > 0:

$$\pi^{\tau,\star}(a|x) = \frac{\exp(q^\star_\tau(x,a)/\tau)}{\int \exp(q^\star_\tau(x,a')/\tau)\,d\pi_\text{ref}(a')}$$

with the soft Q-function $q^\star_\tau$ satisfying a soft BeLLMan equation. As τ → 0, this policy is theoretically expected to recover a deterministic optimal policy. However, the direct vanishing of τ often does not select a policy that preserves the diversity of optimal actions. Instead, in the standard ERL limit, the policy may collapse onto a single optimal action even when multiple actions have identical Q-values, erasing the diversity present in the optimal set.

2. Temperature Decoupling Gambit: Achieving Diversity Preservation

To resolve the loss of diversity in the τ → 0 limit, the temperature decoupling gambit is introduced. The approach is to decouple the temperature parameter used for computing the soft Q-function (denoted as σ) from that used in the Boltzmann action selection (still denoted as τ). By ensuring that σ(τ)/τ → 0 as τ → 0, one constructs a decoupled policy

$$\pi^{\tau,\sigma}(a|x) = \frac{\exp(q^\star_\sigma(x,a)/\tau)}{\int \exp(q^\star_\sigma(x,a')/\tau)\,d\pi_\text{ref}(a')}$$

where $q^\star_\sigma$ is the soft Q-function at temperature σ. In this formulation, σ is sent to zero more rapidly than τ. Under the assumption that the reference policy $\pi_\text{ref}$ has support over all optimal actions in each state, it can be shown that as τ → 0:

$\lim_{\tau\to 0} \pi^{\tau,\sigma}(a|x) = \pistarref_x \propto \pi_{\text{ref},x}\cdot \mathbf{1}_{\{a\,:\,q^\star(x,a)=\esssup_{\pi_{\text{ref},x}} q^\star(x,\cdot)\}}$

That is, the limiting policy assigns probability (according to $\pi_\text{ref}$) to all optimal actions, preserving their diversity instead of arbitrarily selecting one.

3. Theoretical Guarantees and Convergence Analysis

The main theoretical contributions include:

• Monotonic convergence of soft Q-functions $q^\star_\tau(x,a)$ to a reference-optimal Q-function as τ → 0.
• Proof that the decoupled policy converges (with respect to total variation for discrete action sets) to the diversity-preserving optimal policy $\pistarref$.
• The mean of the limiting return distributions equals the expected value for the optimal policy, but the distributional shape (variance, skewness) is also preserved, which is critical for safety and risk-sensitive applications.

The combination of vanishing-entropy regularization and temperature decoupling thus yields both optimality and maximal diversity among optimal actions.

4. Diversity-Preserving Policy in Distributional RL

Beyond expectation, understanding the full distribution of returns under the diversity-preserving optimal policy is essential, especially in risk-sensitive or safe RL. The paper defines a soft distributional BeLLMan operator:

$(\mathcal{T}^\pi_{\tau}\zeta)_{x,a} = \left(\bootfn{r(x,a)}\circ\proj{\mathbb{R} - \gamma\tau\,\mathrm{KL}(\pi_{x'}\,\Vert\,\pi_{\text{ref},x'})}\right)_\#(\zeta_{x',a'}\otimes P_{x,a}(dx'))$

and pairs it with corresponding soft policy improvement. By iteratively applying the operator, one obtains a sequence of estimated return distributions that converge (in Wasserstein metric) to the return distribution for the diversity-preserving policy $\pistarref$. The contraction property of the operator guarantees convergence and accuracy up to arbitrary precision.

5. Interpretability and Algorithmic Implementation

The approach is notable for its interpretability:

• The limiting diversity-preserving policy is analytically characterized: it is uniform (or reference-weighted) over the set of optimal actions.
• The convergence of associated return distributions provides full probabilistic information, not only mean values, for policy analysis and deployment.
• The decoupled-annealing schedule is straightforward to implement, requiring only adjustment of the schedules for σ and τ during training, with their ratio tending to zero.

Algorithmically, one can plug the decoupling mechanism into any entropy-regularized value or policy iteration loop, monitoring both convergence of value functions and the policy support over optimal actions.

6. Implications and Applications

Preserving diversity among optimal policies has several critical implications:

• Robustness and Exploration: Uniform coverage over optimal actions precludes over-specialization, mitigating the risk of catastrophic failure in environments with multiple equivalent high-reward strategies.
• Safety and Distributional Considerations: Knowing the return distribution (and not just the mean) is vital in safety-critical RL, and the proposed approach provides this through a provably convergent implementation.
• Fairness and Control: By varying $\pi_\text{ref}$, a practitioner can embed fairness or behavioral priors into the limiting diversity-preserving policy.
• Modularity: The method can be combined with other RL paradigms (e.g., constrained RL, distributional RL, safe RL) where control over the diversity of actions is valuable.

7. Summary of Key Equations and Theoretical Results

Concept	Formula
Boltzmann policy (ERL)	$$\pi^{\tau,\star}(a|x) = \frac{\exp(q^\star_\tau(x,a)/\tau)}{\int \exp(q^\star_\tau(x,a')/\tau)\,d\pi_\text{ref}(a')}$$
Decoupled temperature policy	$\pi^{\tau,\sigma}(a|x)= \boltzmann{\tau} q^\star_{\sigma}(x,a)$
Limiting diversity-preserving policy	$\pistarref_x \propto \pi_{\text{ref},x} \cdot \mathbf{1}_{\{a: q^\star(x,a) = \esssup_{\pi_{\text{ref},x}} q^\star(x,\cdot)\}}$
Soft distributional BeLLMan op.	$(\mathcal{T}^\pi_{\tau}\zeta)_{x,a} = \left(\bootfn{r(x,a)}\circ\proj{\mathbb{R} - \gamma\tau\,\mathrm{KL}(\pi_{x'}\,\Vert\,\pi_{\text{ref},x'})}\right)_\#(\zeta_{x',a'}\otimes P_{x,a}(dx'))$
Contraction and convergence	The operator is a contraction and iterates converge in Wasserstein metric to return distribution associated to $\pistarref$

These results provide both a theoretical and computational recipe for realizing interpretable, diversity-preserving optimal policies in RL via entropy regularization and controlled annealing of the temperature parameter (Jhaveri et al., 9 Oct 2025).