Technical Summary: Revisiting Mixture Policies in Entropy-Regularized Actor-Critic

Problem Statement

In continuous action reinforcement learning (RL), particularly within entropy-regularized actor-critic frameworks like Soft Actor-Critic (SAC), the standard policy parameterization is the unimodal Gaussian distribution (or its bounded "squashed" variant). While mixture policies (e.g., Gaussian Mixture Models) theoretically offer greater flexibility to represent multimodal optimal policies and facilitate exploration, they have been largely absent from state-of-the-art algorithms.

The primary barrier to adopting mixture policies is the lack of a low-variance reparameterization (RP) gradient estimator. Standard RP techniques, which are crucial for the stability and performance of SAC, rely on the ability to express actions as deterministic functions of noise and parameters ($a = f_\theta(\epsilon)$). For mixture policies, the discrete nature of the component selection (the mixing weights) prevents direct reparameterization, forcing reliance on high-variance likelihood-ratio (LR) estimators or ad-hoc workarounds that compromise flexibility. This paper investigates whether the theoretical benefits of mixture policies can be realized in practice by addressing this gradient estimation gap.

Methodology

1. Theoretical Analysis of Mixture Policies

The authors first establish theoretical guarantees regarding the robustness of mixture policies under entropy regularization:

• Stationary Points: They prove that the set of stationary points for a mixture policy is a superset of that for a base unimodal policy. Consequently, the optimal stationary point of a mixture policy is guaranteed to be at least as good as that of the base policy.
• Existence under High Entropy: In scenarios with strong entropy regularization (large entropy scale $\alpha$), unimodal Gaussian policies may fail to possess any stationary points (divergence). In contrast, Gaussian mixture (GM) policies are shown to retain stationary points for larger values of $\alpha$, offering greater robustness.
• Multimodality: The analysis suggests that mixture policies are better suited for environments with multimodal reward landscapes, as they can maintain high probability mass across multiple modes without collapsing to a single mode, a common failure mode for unimodal policies under high entropy.

2. The Marginalized Reparameterization (MRP) Estimator

To enable the practical training of mixture policies, the paper proposes the Marginalized Reparameterization (MRP) estimator.

• Derivation: The authors extend the reparameterization policy gradient theorem to mixture policies. Instead of reparameterizing the discrete mixing variable (which is impossible without bias), they marginalize over the discrete component selection variable $K$.
• Mechanism: The MRP estimator computes the gradient by summing the contributions of all mixture components weighted by their mixing probabilities, while reparameterizing only the continuous component policies.
  $$\hat{\nabla}_\theta \hat{J}(\pi^m_\theta) = \nabla_\theta \sum_{k=1}^N \pi^w_\theta(k|S_t) \left( Q_w(S_t, f_\theta(\epsilon_t; S_t, k)) - \alpha \log \pi^m_\theta(f_\theta(\epsilon_t; S_t, k)|S_t) \right)$$
• Variance Properties: The authors prove that under specific smoothness assumptions on the reward function, the trace of the covariance matrix of the MRP estimator is lower than that of the standard LR estimator. This theoretical variance reduction is critical for stability in high-dimensional continuous control.

3. Experimental Setup

The authors evaluate the proposed method using SAC as the base algorithm across a broad suite of benchmarks:

• Synthetic Bandits: Designed to mimic multimodal critic functions to test peak-finding capabilities.
• Continuous Control Benchmarks: 57 environments including OpenAI Gym MuJoCo (7 tasks), DeepMind Control Suite (10 tasks), MetaWorld (30 tasks), MyoSuite (10 tasks), and classic control environments (Pendulum, Acrobot, MountainCar).
• Baselines: Comparisons include standard Squashed Gaussian (SG) with RP, Squashed Gaussian Mixture (SGM) with LR, Half-Reparameterization (HalfRP), and Gumbel-Softmax RP (GumbelRP).

Key Results

1. Superiority of MRP over Alternative Estimators

In experiments comparing gradient estimators for mixture policies, the MRP estimator consistently outperforms the LR, HalfRP, and GumbelRP estimators.

• Stability: The LR estimator frequently diverges in high-dimensional environments (e.g., Humanoid, HumanoidStandup), whereas MRP maintains stability.
• Variance: Empirical measurements confirm that MRP exhibits significantly lower gradient variance than LR, even when the LR estimator is equipped with a baseline.
• Bias: Unlike the GumbelRP estimator, which introduces bias via the Gumbel-Softmax trick, MRP remains unbiased.

2. Performance of Mixture Policies

• General Benchmarks: Across the 57 continuous control environments, mixture policies (SGM-MRP) perform on par with or slightly better than base unimodal policies (SG-RP). In MetaWorld, mixture policies frequently achieve significantly higher success rates.
• Multimodal Bandits: In synthetic bandit tasks with multimodal reward functions, mixture policies are more likely to find the maximal peaks compared to base policies, demonstrating superior exploration capabilities.
• Unshaped Rewards: The most significant empirical advantage of mixture policies is observed in environments with unshaped (sparse) rewards (e.g., standard MountainCar, Acrobot). In these settings, where the critic is less smooth and exhibits multiple peaks, mixture policies substantially outperform base policies. The authors attribute this to the mixture policy's ability to maintain multiple action modes simultaneously, facilitating more efficient exploration in the absence of dense reward signals.
• Entropy Robustness: Consistent with theoretical predictions, mixture policies maintain performance at higher entropy scales where base policies often fail to converge or collapse to suboptimal modes.

Significance and Claims

The paper claims to elevate mixture policies from a "theoretical curiosity" to a "practical tool" for entropy-regularized actor-critic algorithms. Its primary contributions are:

1. Theoretical Justification: It provides the first proof that mixture policies offer improved solution quality and robustness to entropy regularization compared to unimodal policies, specifically regarding the existence of stationary points under high entropy.
2. Algorithmic Innovation: It introduces the Marginalized Reparameterization (MRP) estimator, the first unbiased RP-based gradient estimator for mixture policies that does not compromise flexibility. The paper proves this estimator has lower variance than the standard LR approach.
3. Empirical Validation: Through extensive experiments, the paper demonstrates that MRP enables mixture policies to match or exceed the performance of standard Gaussian policies, with clear advantages in environments characterized by multimodal reward landscapes and unshaped rewards.

The authors conclude that while mixture policies introduce modest computational overhead, the MRP estimator makes them a viable and often superior alternative to unimodal policies, particularly in challenging exploration scenarios. They do not claim mixture policies are universally superior for all tasks but highlight their specific utility in overcoming the limitations of unimodal representations in entropy-regularized settings.