Technical Summary: Emergence of Exploration in Policy Gradient Reinforcement Learning via Retrying

1. Problem Statement

In Reinforcement Learning (RL), exploration is typically driven by adding explicit bonuses (e.g., entropy, curiosity, or count-based bonuses) to the reward signal or by employing posterior sampling methods. The authors posit a distinct mechanism: exploration should emerge naturally from the agent's need to maximize rewards under uncertainty through retrying.

The core intuition is that in environments where an agent can retry actions or encounter similar states, a greedy policy is suboptimal if it ignores the potential for higher returns in subsequent attempts. Without the opportunity to retry, the rational choice is the action currently believed to yield the highest reward. However, under uncertainty, a policy that hedges against uncertainty by exploring alternative actions can achieve a higher expected maximum return over multiple trials. Existing methods often require complex auxiliary models to estimate uncertainty or construct bonuses; this paper asks if exploration can be induced purely by optimizing a reward-based objective that formalizes the value of retrying.

2. Methodology

2.1 The ReMax Objective

The authors propose ReMax (Retry Maximum), an objective that evaluates a policy $\pi$ by the expected maximum return over $M$ independent samples (trials), rather than the expected return of a single trial.

In a bandit setting with action values $\mu$ and a distribution over these values $\Pi$ (representing epistemic uncertainty), the ReMax objective is defined as:
$$J^M_{\text{ReMax}}(\pi) := \mathbb{E}_{\mu \sim \Pi} \left[ \mathbb{E}_{A^{[M]} \sim \pi} \left[ \max_{m \in [M]} \mu_{A_m} \mid \mu \right] \right]$$
where $A^{[M]}$ represents $M$ independent draws from the policy $\pi$.

• Mechanism: By maximizing the expected maximum of $M$ draws, the objective incentivizes the policy to assign probability mass to actions that might yield high rewards, even if they are not currently the estimated mean maximum. This induces stochasticity (exploration) as an emergent property to hedge against uncertainty.
• Deterministic Case: Even with fixed rewards, ReMax reshapes the objective geometry. Increasing the retry parameter $M$ flattens the gradient near the optimum, slowing convergence and sustaining exploration, whereas smaller $M$ accelerates updates.

2.2 Extension to RL (MDPs)

Extending ReMax to Markov Decision Processes (MDPs) presents challenges:

1. Uncertainty Unit: Uncertainty in RL encompasses the entire environment (rewards and transitions).
2. Feasibility of Retries: Physically resetting an environment to retry multiple actions from the same state is often infeasible.

To address this, the authors define the ReMax objective using a Q-function $Q^\pi_M(s, a)$, which represents the expected episodic return. The objective is formulated as the expected maximum over $M$ draws of Q-values at a state $s$:
$$J^M_{\text{ReMax}}(\pi, s, Q) := \mathbb{E}_{Q \sim \mathcal{Q}} \left[ \mathbb{E}_{A^{[M]} \sim \pi} \left[ \max_{m \in [M]} Q(s, A_m) \right] \right]$$
Here, the distribution $\mathcal{Q}$ captures uncertainty over Q-values (either via explicit modeling like ensembles or implicitly via the non-stationarity of deep critics).

2.3 Policy Gradient Formulation

To optimize ReMax efficiently, the authors derive a new policy gradient (PG) estimator that is computable from single-trajectory returns.

• Naïve Approach Failure: A direct application of the policy gradient theorem requires observing returns for all $M$ sampled actions simultaneously, which is infeasible in episodic RL.
• Expected Improvement (EI) Reformulation: The authors introduce a baseline based on the "Expected Improvement" (borrowed from Bayesian optimization). They decompose the max operator into an action-specific term and a baseline representing the maximum of the other $M-1$ draws ($W_{-m}$).
• Closed-Form Solution: By sorting the Q-values, they derive a closed-form expression for the expected improvement term. This allows the gradient to be estimated using a single trajectory return $R(s, a)$ and the current Q-estimates.
• Continuous Parameter $m$: The integer retry count $M$ is generalized to a continuous real parameter $m > 0$. This enables fine-grained control over the exploration-exploitation trade-off. The term $(1 - C_j)^{M-1}$ in the closed-form solution is replaced by $(1 - C_j)^{m-1}$.

2.4 RePPO Algorithm

The authors instantiate ReMax within the Proximal Policy Optimization (PPO) framework, creating ReMax PPO (RePPO).

• Architecture: An on-policy actor-critic method using a Q-critic.
• Surrogate Objective: RePPO modifies the standard PPO surrogate by replacing the standard advantage with an EI-based advantage.
  $$A^+(t) = R^+(t) - b^+(s_t)$$
  where $R^+(t) = \text{EIm}(R^\lambda_t, \pi_\theta, Q_\phi(s_t, \cdot))$ and $b^+$ is the expected value of the EI-transformed Q-values.
• Q-Replacement Strategy: To prevent the critic from overestimating the improvement of the sampled action (due to underestimation of the Q-value), the authors replace the Q-value of the sampled action with the actual trajectory return when computing the EI term. This ensures the "improvement" is measured against the true return, not a potentially biased estimate.

3. Key Contributions

1. ReMax Objective: A novel RL objective that induces exploration via greedy reward maximization under uncertainty (retrying) without explicit exploration bonuses.
2. Theoretical Derivation: A new policy gradient formulation for ReMax that is estimable from single-trajectory returns, utilizing a closed-form expression for the expected maximum over $M$ draws.
3. Continuous Control: The generalization of the discrete retry count $M$ to a continuous parameter $m$, allowing for precise tuning of the exploration rate.
4. RePPO Algorithm: A practical deep RL algorithm (RePPO) that implements ReMax within a PPO framework, demonstrating that exploration can emerge from the objective function itself.

4. Experimental Results

4.1 Bandit Experiments

• Stochasticity: In a two-armed bandit, ReMax induces stochastic policies for $M \ge 2$, whereas standard RL ($M=1$) yields deterministic policies.
• Adaptivity: ReMax adapts exploration to the magnitude of reward uncertainty. As variance increases, the optimal ReMax policy increases exploration to capture rare high-reward outcomes, unlike entropy-regularized softmax which remains fixed based on the mean.
• Regret: In posterior bandit settings (Beta-Bernoulli and Gaussian-Gaussian), ReMax exhibits empirically sublinear cumulative regret, comparable to Thompson Sampling and UCB, and superior to a Softmax baseline.

4.2 MinAtar Benchmarks

Evaluated on four games (Breakout, Asterix, Freeway, Space Invaders):

• Performance: RePPO (with $m \in [1.2, 1.4]$) outperformed standard PPO variants (PPO-V, PPO-Q) and PPO with entropy bonuses or Random Network Distillation (RND) bonuses.
• Exploration: RePPO maintained higher policy entropy throughout training compared to PPO with an explicit entropy bonus, despite having no such bonus.
• Parameter Sensitivity: Performance peaked around $m=1.2$ to $1.4$. Larger $m$ slowed entropy decay (sustaining exploration), while smaller $m$ led to faster convergence but lower performance.
• Ablation: Removing the action-independent baseline or the Q-replacement strategy significantly degraded performance, confirming their necessity.

4.3 Craftax Benchmark

Evaluated on Craftax, a large-scale, open-ended RL environment:

• Performance: RePPO ($m=1.2$) achieved competitive performance (11.87% of max reward) compared to PPO variants with entropy and RND bonuses (11.66% and 11.68%), and outperformed PPO without bonuses.
• Scalability: RePPO maintained higher entropy without explicit bonuses, demonstrating that the mechanism scales to complex, long-horizon tasks.

5. Significance and Claims

The paper claims that ReMax offers a promising alternative to traditional exploration methods by:

• Eliminating the need for explicit bonuses: Exploration emerges naturally from the objective of maximizing the expected maximum return over retries.
• Simplifying algorithm design: It avoids the computational overhead and complexity of designing curiosity bonuses or maintaining explicit posterior models (though it can utilize them for uncertainty estimation).
• Providing fine-grained control: The continuous parameter $m$ allows for precise adjustment of the exploration-exploitation trade-off.

The authors note that while their deep RL experiments primarily demonstrated stochastic exploration (increased policy entropy), the theoretical framework supports deep exploration (temporally coherent information gathering) if combined with explicit uncertainty modeling (e.g., ensembles). They position ReMax as a general principle applicable to domains where retrying is valuable, including reasoning tasks with Large Language Models (LLMs), drawing parallels to pass@K objectives in that field.

The work is modest in its theoretical claims, noting that formal regret bounds for the deep RL setting are left for future work, and the current results are empirical demonstrations of effectiveness.