Poly-EPO: Training Exploratory Reasoning Models

Abstract: Exploration is a cornerstone of learning from experience: it enables agents to find solutions to complex problems, generalize to novel ones, and scale performance with test-time compute. In this paper, we present a framework for post-training LMs that explicitly encourages optimistic exploration and promotes a synergy between exploration and exploitation. The central idea is to train the LM to generate sets of responses that are collectively accurate under the reward function and exploratory in their reasoning strategies. We first develop a general recipe for optimizing LMs with set reinforcement learning (set RL) under arbitrary objective functions, showing how standard RL algorithms can be adapted to this setting through a modification to the advantage computation. We then propose Polychromic Exploratory Policy Optimization (Poly-EPO), which instantiates this framework with an objective that explicitly synergizes exploration and exploitation. Across a range of reasoning benchmarks, we show that Poly-EPO improves generalization, as evidenced by higher pass@$k$ coverage, preserves greater diversity in model generations, and effectively scales with test-time compute.

• The paper's main contribution is a scalable set RL framework, Poly-EPO, that explicitly incentivizes both exploration and exploitation in language models.
• It employs a set-level reward structure with a polychromic objective to maintain diverse reasoning strategies, achieving up to 20% improvement in pass@k metrics.
• The approach demonstrates robust generalization and efficient test-time compute scaling by synergizing reward and diversity, validated across mathematical and synthetic reasoning tasks.

Poly-EPO: Training Exploratory Reasoning Models

Abstract and Motivation

The paper introduces Polychromic Exploratory Policy Optimization (Poly-EPO), a scalable set reinforcement learning (set RL) framework for post-training LMs that explicitly incentivizes both optimism under uncertainty and a synergy between exploration and exploitation. Poly-EPO operationalizes this via a set-level reward structure: for each prompt, the LM is optimized to generate a set of responses that are collectively accurate and diverse in reasoning strategies, measured via clustering with an LM-judge. The method addresses the collapse of diversity observed in standard RL fine-tuning of LMs, offering improved generalization and scalability with test-time compute.

Set RL Framework and Poly-EPO Algorithm

Set RL generalizes standard RL by assigning rewards to sets of generations rather than individual samples. For prompts $x$, sets of $n$ generations $y_{1:n} \sim \pi_\theta(\cdot~|~x)$ are scored via a symmetric set-level objective $f(x, y_{1:n})$. The paper details a tractable unbiased gradient estimator for optimizing LMs under set RL, leveraging advantage computation over all sampled sets for variance reduction and computational efficiency.

Poly-EPO instantiates this with the polychromic objective:

$$f_\text{poly}(x, y_{1:n}) = \frac{1}{n}\sum_{i=1}^n r(x, y_i) \cdot d(x, y_{1:n})$$

where $r$ is task reward and $d$ is a diversity function measuring the distinct reasoning-strategy clusters (via an LM-judge, e.g., Qwen-3-4B-Instruct) among generations. Notably, incorrect but exploratory generations can receive positive advantage if they contribute to strategic diversity, ensuring the model continues to attempt and refine novel strategies (optimism under uncertainty).

Analysis: Synergy of Exploration and Exploitation

The paper offers a decomposition of the marginal set advantage under Poly-EPO, showing explicit synergy:

• The advantage assigned to a generation depends on its contribution to both the mean reward and mean diversity in sets containing it, and
• On the covariance between reward and diversity, thereby encouraging generations that foster joint high performance and strategic exploration.

Comparisons with standard RL and reward-shaping approaches reveal that, as accuracy increases, Poly-EPO maintains strong incentive for exploratory behaviors, unlike reward-shaped RL where the exploration bonus typically vanishes unless manually tuned.

Empirical Evaluation

Mathematical Reasoning Benchmarks

Extensive experiments on the POLARIS-53K dataset, using Qwen-3-4B-Base/GRPO for baseline comparisons, demonstrate that Poly-EPO significantly enhances generalization, as evidenced by pass@$k$ metrics across tasks such as BeyondAIME, AIME, HMMT, and Minerva. Poly-EPO achieves up to 20% higher pass@$k$ coverage compared to strong RL baselines.

Figure 1: Pass@k evaluations demonstrate Poly-EPO's superior coverage growth with increasing attempts $k$, indicating effective diversity of sampled solutions.

Training Dynamics and Reasoning Diversity

Poly-EPO consistently expands the diversity of successful reasoning strategies, as measured by reasoning-strategy clusters among correct generations. The number of unique clusters rises steadily during training, contrasting sharply with a rapid collapse to few strategies under GRPO.

Figure 2: Unique reasoning-strategy clusters and prompt coverage during training, showing Poly-EPO's sustained discovery of diverse approaches and higher coverage.

Branching analysis on sampled generations reveals that Poly-EPO enables LMs to differentiate reasoning trajectories earlier and maintain broader exploration, unlike baselines that delay branching and strategic differentiation.

Scaling with Test-Time Compute

Poly-EPO's diversity translates directly to improved performance under majority voting. As the number of sampled generations increases, Poly-EPO retains higher entropy over candidate answers and achieves equal or greater gains in accuracy compared to baselines, suggesting increased test-time compute is efficiently converted into performance improvements due to the model's broader strategic repertoire.

Synthetic Infinite-Strategy Domains

The method was further evaluated in domains with unbounded correct solution strategies (multi-digit multiplication and polynomial solving). Poly-EPO outperforms both GRPO and base models, maintaining substantially greater diversity in both correct and incorrect responses and discovering orders of magnitude more successful strategies.

Figure 3: Poly-EPO maintains higher strategic diversity and uncovers more distinct correct strategies in synthetic tasks with infinite valid solutions.

Practical and Theoretical Implications

Poly-EPO represents a scalable solution for post-training LLMs to pursue a principled balance between exploration and exploitation, overcoming diversity collapse inherent to standard RL protocols. The set RL framework enables direct optimization of objectives aligned with downstream inference and reasoning diversity, opening avenues for robust test-time compute strategies (e.g., verification, aggregation, majority voting). The explicit synergy encoded in the advantage function obviates the need for manual hyperparameter tuning to regulate exploration, facilitating reliable optimization.

Theoretically, Poly-EPO advances the RL paradigm for LMs by formalizing the contribution of covariance between reward and diversity in strategic exploration, challenging reward-shaping methods where this interaction is lost. The unbiased, tractable gradient estimator paves the way for scalable set RL in large generative models.

Limitations and Future Directions

• Dependency on LM-judge quality: As with all LM-judge-based reward assignment methodologies, Poly-EPO is susceptible to reward hacking and clustering inaccuracies, especially as response lengths increase.
• Scaling laws: The interplay between set size, number of rollouts, and constructed sets remains analytically unresolved; further work should examine scaling behavior on larger models and tasks.
• Diversity measurement: Future research should explore alternative diversity metrics beyond semantic clustering, including non-judge-based approaches.

Potential extensions include application to instruction-following, coding, and proof domains, integration with long-context generative pipelines, and systematic study of set RL scaling effects.

Conclusion

Poly-EPO delivers a scalable set RL algorithm for post-training LMs that explicitly synergizes exploration and exploitation. Empirical results across mathematical and synthetic reasoning domains confirm improved generalization, robust test-time compute scaling, and sustained strategic diversity in generations. The framework advances practical RL for LMs, suggesting future developments in robust, exploratory LLM training and inference.