Technical Summary: Retry Policy Gradients in Continuous Action Spaces

Problem Statement

Exploration remains a central challenge in reinforcement learning (RL), particularly in continuous control. Standard approaches typically rely on stochastic policies combined with explicit entropy regularization (e.g., Soft Actor-Critic, SAC) or noise injection. While effective, these methods introduce an auxiliary entropy bonus that biases the optimal policy toward higher entropy, potentially preventing convergence to the true deterministic optimum.

Recently, "retry-based" objectives (e.g., pass@K, max@K) have shown promise in discrete action spaces and Large Language Model (LLM) post-training. These objectives optimize the maximum return obtained from $M$ sampled trajectories rather than the expected return. Prior work (ReMax) demonstrated that in discrete spaces, this objective naturally encourages exploration by adapting to return uncertainty, even without explicit entropy bonuses. However, the theoretical properties of ReMax in continuous action spaces remained unexplored, and existing gradient estimators for discrete spaces do not extend to continuous domains.

Methodology

Theoretical Framework

The authors extend the ReMax objective to continuous action spaces using a pathwise derivative estimator.

• Objective: Given a Q-function $Q(s, a)$ and a retry budget $M$, the objective is defined as:
  $$J^M(\pi, Q) := \mathbb{E}_{s \sim p_0} \left[ \mathbb{E}_{a_{1:M} \sim \pi(\cdot|s)} \left[ \max_{m=1:M} Q(s, a_m) \right] \right]$$
  Unlike entropy-regularized objectives, this formulation preserves the deterministic optimal policy ($\pi^*$) as the global optimum, avoiding the bias toward stochasticity inherent in entropy bonuses.

• Gradient Estimation: Since exact computation is infeasible in continuous spaces, the authors employ a Reparameterization (RP) gradient estimator. Actions are sampled via $a_i = \mu + \sigma \epsilon_i$ (where $\epsilon_i \sim \mathcal{N}(0, I)$). The objective is approximated by sampling $B$ actions, sorting their Q-values, and computing a weighted sum of the top $M$ values. This allows for backpropagation through the Q-network and the policy parameters.

Theoretical Analysis of Gradient Dynamics

The paper provides a detailed analysis of how the ReMax gradient reshapes the optimization landscape for a Gaussian policy $\mathcal{N}(\mu, \sigma^2)$ with a quadratic reward $r(a) = -a^2$. Two key mechanisms are identified when $M > 1$:

1. Entropy Increase (Directional Effect): When the policy mean is far from the optimum and the scale $\sigma$ is small, the gradient direction naturally points toward increasing $\sigma$. This occurs because the "max" operator over $M$ samples creates a positive expected upward fluctuation in reward (or downward in cost) that is only accessible by increasing the policy's scale. This promotes stochastic exploration without an explicit bonus.
2. Gradient Damping (Magnitude Effect): As the policy approaches the optimum, the gradient norm diminishes. Selecting the best of $M$ samples reduces the deviation from the optimum, effectively "forgiving" suboptimal samples and slowing convergence. This damping effect helps sustain higher entropy for longer during training.

The authors also analyze the interaction with the Adam optimizer. They show that Adam's adaptive normalization (specifically the stabilization parameter $\epsilon$) can counteract the gradient damping effect. A small $\epsilon$ (default $10^{-8}$) normalizes gradients, potentially accelerating convergence and reducing the entropy-sustaining effect, while a larger $\epsilon$ restores the damping behavior.

Algorithm: ReMAC

The authors instantiate these findings in ReMax Actor-Critic (ReMAC), an off-policy actor-critic algorithm.

• Critic: Trained via standard temporal-difference regression with clipped double Q-learning (using two critics to reduce overestimation).
• Actor: Optimized using the pathwise derivative of the ReMax objective.
• Implementation: ReMAC is presented as a minimal modification to SAC, removing the entropy bonus from the critic target and temperature tuning, and replacing the actor loss with the ReMax loss.

Key Results

Theoretical Findings

• Vector Field Analysis: Visualizations of the gradient vector field for 1D Gaussian policies confirm that for $M > 1$, gradients push the policy toward higher entropy when far from the optimum, whereas standard RL ($M=1$) always pushes toward lower entropy.
• Convergence Properties: ReMax converges to the deterministic optimum ($\mu=0, \sigma=0$), whereas entropy regularization converges to a stochastic optimum ($\mu=0, \sigma>0$).
• Adam Interaction: Theoretical and empirical analysis of a toy problem demonstrates that the damping effect is sensitive to Adam's $\epsilon$ parameter.

Empirical Evaluation

Experiments were conducted on six continuous control tasks from the Brax suite (Ant, HalfCheetah, Hopper, Reacher, Swimmer, Walker2d) over 3 million steps.

• Performance: ReMAC with $M > 1$ achieved performance comparable to SAC and generally outperformed PPO. In several environments (Ant, Reacher, Swimmer, Walker2d), larger retry budgets ($M=4, 8$) yielded higher average returns than $M=1$.
• Entropy Dynamics: Consistent with theoretical predictions, ReMAC with $M > 1$ maintained higher policy entropy than $M=1$ (standard RL) throughout training. Entropy generally increased with $M$, particularly in the Ant and Swimmer environments.
• Comparison to SAC: While ReMAC achieved higher entropy than standard RL ($M=1$), its entropy remained lower than SAC, which is expected given SAC's explicit entropy maximization.

Significance and Claims

The paper positions its contributions as follows:

1. Theoretical Extension: It is the first work to extend ReMax to continuous action spaces and provide a rigorous theoretical explanation for its exploration-promoting properties via pathwise derivatives.
2. Intrinsic Exploration: It demonstrates that retry-based objectives can induce stochastic exploration through the geometry of the gradient landscape (direction and magnitude) rather than through explicit entropy bonuses. This avoids the bias of shifting the optimal policy away from the deterministic solution.
3. Practical Viability: The introduction of ReMAC shows that these theoretical properties translate to practical deep RL settings, achieving SAC-level performance with minimal architectural changes.
4. Optimization Insight: The work highlights the critical role of optimizer hyperparameters (specifically Adam's $\epsilon$) in modulating the exploration-exploitation trade-off in retry-based objectives.

The authors remain modest regarding the scope, noting that their theoretical analysis assumes smoothness and strong convexity, which may not hold in all practical deep RL environments. They frame ReMAC as a proof-of-concept for stochastic exploration and suggest that future work could integrate ReMax with posterior sampling over Q-functions to enable "deep exploration" in sparse-reward settings.