Global Optimality for Constrained Exploration via Penalty Regularization

1. Problem Formulation

The paper addresses the challenge of Constrained Maximum-Entropy Exploration (CME) in Reinforcement Learning (RL).

• Goal: Maximize the entropy of the state-action occupancy measure ($H(\lambda_\pi)$) to encourage diverse exploration of the environment.
• Constraints: The exploration must satisfy general convex constraints on the occupancy measure, such as safety limits (cumulative cost), resource budgets, or imitation requirements (closeness to a reference policy via KL divergence).
• Mathematical Formulation:
  $$\max_{\theta \in \Theta} H(\lambda_{\pi_\theta}) \quad \text{s.t.} \quad R(\lambda_{\pi_\theta}) \leq 0$$
  where $\pi_\theta$ is a parameterized policy (e.g., softmax), and $\lambda_{\pi_\theta}$ is the induced state-action occupancy measure.
• Challenges:
  1. Non-Convexity: While the problem is convex in the occupancy measure space ($\lambda$), parameterizing the policy ($\theta$) renders the objective and constraints non-convex.
  2. Lack of Additive Structure: Entropy maximization lacks the additive Bellman structure, making standard Dynamic Programming or Actor-Critic methods inapplicable.
  3. Limitations of Prior Work: Existing methods (e.g., Ying et al., 2025) rely on primal-dual approaches that only guarantee ergodic averages or weak regret. These do not ensure that the final output policy (the last iterate) is feasible or near-optimal, which is critical for deployment.

2. Methodology: Policy Gradient Penalty (PGP)

The authors propose Policy Gradient Penalty (PGP), a single-loop, primal-only algorithm that operates directly in the policy parameter space.

Core Mechanism

Instead of using dual variables (Lagrange multipliers) which require complex dual updates, PGP transforms the constrained problem into an unconstrained one using quadratic penalty regularization:
$$\min_{\theta \in \Theta} P(\theta) := -H(\lambda_{\pi_\theta}) + \beta [R(\lambda_{\pi_\theta})]_+^2$$
where $[x]_+ = \max(0, x)$ and $\beta$ is the penalty parameter.

Key Algorithmic Innovations

1. Pseudo-Rewards: A major novelty is the construction of "pseudo-rewards" to compute gradients.

   • Standard stochastic penalty methods require separate estimators for the constraint value and its gradient, leading to bias.
   • PGP leverages the Policy Gradient Theorem to reinterpret the gradient of the penalized objective as a standard policy gradient with respect to a constructed pseudo-reward vector.
   • Specifically, the gradient of the penalty term $\nabla_\theta [R(\lambda)]_+^2$ is computed as $\nabla_\theta V^{\pi_\theta}(r_{pseudo})$, where $r_{pseudo} = 2[R(\lambda)]_+ \nabla_\lambda R(\lambda)$.
   • This allows the use of a single trajectory batch to estimate the full gradient of the penalized objective without auxiliary function-value estimators.

2. Single-Loop SGD: The algorithm performs Stochastic Gradient Descent (SGD) on the penalized objective using these pseudo-rewards, avoiding the nested loops or dual updates typical of primal-dual methods.

Theoretical Foundation

The method relies on the concept of Hidden Convexity:

• Although the problem is non-convex in $\theta$, the mapping from policy parameters to occupancy measures preserves a hidden convex structure.
• The authors prove that the quadratic penalty formulation preserves this hidden convexity and smoothness properties.
• Under the assumption of Strong Duality (which is weaker than Slater's condition and holds for hidden convex problems), they establish that the penalized problem's solution converges to the constrained solution.

3. Key Contributions

1. Algorithm (PGP): Introduction of a single-loop, primal-only policy gradient method for constrained entropy maximization that enforces constraints via quadratic penalties without dual variables.
2. Global Last-Iterate Convergence: Theoretical proof that the last iterate of the algorithm converges to an $(\epsilon, \epsilon)$-optimal solution (near-optimal objective and near-feasible constraints) with high probability. This is a significant improvement over prior ergodic/weak-regret guarantees.
3. Sample Complexity: Establishment of a non-asymptotic sample complexity of $\tilde{O}(\epsilon^{-6})$ for general policy parameterizations under strong duality. This is the first such guarantee for constrained maximum-entropy exploration in this setting.
4. Pseudo-Reward Construction: A novel technique to handle the chain rule in stochastic penalty methods by reinterpreting gradients as policy gradients, eliminating the need for multiple sampling steps per iteration.

4. Experimental Results

The authors validated PGP on both discrete and continuous control tasks:

• Grid World (FrozenLake):

  • Task: Maximize entropy while avoiding "hole" states (safety constraint).
  • Comparison: Outperformed the state-of-the-art Primal-Dual Policy Gradient (PDPG) method (Ying et al., 2025).
  • Observation: PDPG exhibited oscillations around the constraint boundary (weak regret behavior), whereas PGP converged to a stable policy that strictly satisfied constraints while maintaining high entropy.
  • Robustness: PGP showed robustness to the penalty parameter $\beta$ and step size $\eta$.

• Continuous Control (PointMass & SafeCartpole):

  • PointMass (Imitation): Explored the environment while staying close to a reference SAC policy. PGP successfully balanced exploration and imitation, with the policy deviating from the reference as the constraint budget ($R_{max}$) increased.
  • SafeCartpole: A challenging swing-up task where the cart position must remain within safety bounds. PGP learned a "bang-bang" control strategy to swing the pole up while strictly adhering to the position constraints, demonstrating scalability to high-dimensional continuous spaces.

5. Significance and Impact

• Theoretical Breakthrough: This work bridges the gap between convex optimization theory and practical RL by providing last-iterate convergence guarantees for non-convex, constrained exploration problems. Previous methods could not guarantee that the final deployed policy was safe or optimal.
• Practical Applicability: The single-loop nature of PGP makes it computationally efficient and easier to implement than primal-dual methods, which often suffer from tuning difficulties with dual variables.
• Generalizability: The framework extends beyond RL to general hidden-convex constrained optimization problems, offering a unified approach for safety-critical exploration, imitation learning, and resource-constrained tasks.
• Addressing Real-World Needs: By handling general convex constraints (safety, cost, imitation) with rigorous guarantees, the method addresses a critical bottleneck in deploying RL agents in real-world scenarios where safety and feasibility are non-negotiable.

In summary, the paper presents a rigorous, scalable, and theoretically grounded solution to constrained exploration, moving beyond heuristic approaches and weak convergence guarantees to provide a deployable algorithm with global optimality assurances.