State Entropy Regularization for Robust Reinforcement Learning

TL;DR

This work analyzes state entropy regularization as a robustness-promoting mechanism in reinforcement learning. It provides a rigorous characterization showing that state entropy exactly solves a reward-robust RL problem via an explicit uncertainty set and delivers nontrivial bounds under kernel uncertainty, with policy entropy offering different robustness properties. It also establishes fundamental limits: entropy regularization cannot fully solve kernel robustness and can harm risk-averse performance, with robustness benefits being sensitive to rollout budgets. Empirically, state entropy improves resilience to spatially structured perturbations in both discrete and continuous tasks, albeit requiring sufficient rollouts for reliable entropy estimation. Overall, the paper clarifies when and how state entropy regularization can enhance robustness and where its applicability is limited in practice.

Abstract

State entropy regularization has empirically shown better exploration and sample complexity in reinforcement learning (RL). However, its theoretical guarantees have not been studied. In this paper, we show that state entropy regularization improves robustness to structured and spatially correlated perturbations. These types of variation are common in transfer learning but often overlooked by standard robust RL methods, which typically focus on small, uncorrelated changes. We provide a comprehensive characterization of these robustness properties, including formal guarantees under reward and transition uncertainty, as well as settings where the method performs poorly. Much of our analysis contrasts state entropy with the widely used policy entropy regularization, highlighting their different benefits. Finally, from a practical standpoint, we illustrate that compared with policy entropy, the robustness advantages of state entropy are more sensitive to the number of rollouts used for policy evaluation.

Figures (21)

• Figure 1: State entropy regularization leads to diverse behavior and improved robustness. We plot 50 trajectories from standard RL, policy-entropy–regularized RL and state-entropy–regularized RL on a manipulation task (left). Policies are trained on the nominal environment and evaluated in a perturbed version with a red obstacle along the optimal trajectory. The state-entropy–regularized agent often bypasses the obstacle and solves the perturbed task.
• Figure 2: Uncertainty sets $\tilde{\mathcal{R}}^\pi_\epsilon(\alpha)$ induced by state entropy regularization in a two-state MDP with $r^\pi(s_1) = 1$, $r^\pi(s_2) = 0$, and $\epsilon = \frac{1}{2}$. The $x$- and $y$-axes correspond to the perturbed rewards $\tilde{r}^\pi(s_1)$ and $\tilde{r}^\pi(s_2)$, respectively. As $\alpha$ increases, the uncertainty set expands, interpolating between $\ell_\infty$- and $\ell_1$-type constraints in the limits $\alpha \to 0$ and $\alpha \to \infty$.
• Figure 3: An illustrative toy MDP. The grid is a $2\times6$ environment with the initial state marked with "I" and the goal state with "G". The red path shows the longest deterministic trajectory that visits the maximum number of states before reaching the goal. Reward is given only at the goal state.
• Figure 5:
• Figure 6:

Theorems & Definitions (14)

• Theorem 3.1
• Theorem 3.2: Limit uncertainty in $\alpha$
• Theorem 3.3
• Theorem 4.1
• Theorem 4.2
• proof
• proof
• proof
• proof
• proof