Poly-EPO: Training Exploratory Reasoning Models
Ifdita Hasan Orney, Jubayer Ibn Hamid, Shreya S Ramanujam, Shirley Wu, Hengyuan Hu, Noah Goodman, Dorsa Sadigh, Chelsea Finn
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 language models (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@ coverage, preserves greater diversity in model generations, and effectively scales with test-time compute.
AI Impact Assessments
(1 models)Scientific Impact Assessment: Poly-EPO: Training Exploratory Reasoning Models
1. Core Contribution
Poly-EPO addresses a well-documented problem in RL post-training of language models: diversity collapse, where RL fine-tuning narrows the model's generation distribution onto a small set of high-reward behaviors. The paper makes three interleaved contributions:
1. A general recipe for set RL on LMs: The authors show how to adapt standard RL algorithms (GRPO, PPO) for set-level objectives by computing a "marginal set advantage" — the sum of set-level advantages over all sets containing a given generation. This avoids the exponential sample complexity of naive set RL implementations.
2. The Poly-EPO algorithm: Instantiates set RL with a "polychromic objective" — the product of average reward and diversity (measured via LM-judge clustering). The multiplicative structure means sets must simultaneously achieve high reward *and* high diversity.
3. Theoretical analysis of exploration-exploitation synergy: The advantage decomposition (Eq. 15) reveals a covariance term that explicitly rewards generations enabling joint high reward and high diversity — a structural property absent from additive reward-shaping approaches.
2. Methodological Rigor
Strengths in theory: The paper provides formal justification (Proposition 3.1) that the proposed estimator is unbiased up to a scaling factor absorbable into the learning rate, established through U-statistic arguments. Lemma 5.1 connects set RL to standard RL via logit-shift analysis, providing a clean justification for why standard RL algorithms can be used with the marginal set advantage.
The analysis comparing Poly-EPO's advantage (Eq. 15) against reward-shaped RL (Eq. 16) is insightful: it shows that in standard reward shaping, the exploration threshold for incorrect-but-diverse generations becomes harder to satisfy as model accuracy improves (requiring d(x,y) ≥ p/λ + E[d]), while set RL's shared credit assignment naturally sustains exploration incentives.
Concerns: The LM-judge clustering introduces a non-trivial dependency on judge quality that is not rigorously validated. The paper acknowledges this as a limitation but does not quantify clustering reliability or sensitivity. The choice to construct all sets (70 sets from 8 rollouts with set size 4) is computationally feasible here but the scaling behavior with larger N and n is not studied. The unbiasedness proof requires a scaling factor M that is not implemented in practice — the authors note this but don't assess the impact of this approximation.
3. Potential Impact
Practical impact: The method is compatible with existing RL training infrastructure (it only modifies advantage computation), making adoption straightforward. The pass@k improvements (up to 20%) are significant for settings where multiple attempts are available, which is increasingly the norm in reasoning applications paired with verifiers (e.g., Lean for theorem proving).
Broader implications: The framework of set RL for LMs is generalizable beyond exploration — any set-level objective (coverage, calibration, etc.) could be optimized using this recipe. This opens a design space for post-training objectives that reason about collections of outputs rather than individual ones.
Test-time compute scaling: The demonstration that Poly-EPO scales better with majority voting is timely given the industry's shift toward inference-time scaling. However, the majority voting results show modest improvements over baselines on some benchmarks (e.g., HMMT Feb), suggesting the gains are not universal.
4. Timeliness & Relevance
The paper arrives at a moment when diversity collapse in RL post-training is widely recognized as a bottleneck, with multiple concurrent works (SKM25, TFK+25, CHZ+25, etc.) proposing solutions. Poly-EPO differentiates itself by providing a principled framework (set RL) rather than ad-hoc reward bonuses. The paper directly addresses the λ-tuning problem that plagues additive exploration bonuses — a practical pain point in LLM training pipelines.
5. Strengths & Limitations
Key Strengths:
Notable Weaknesses:
6. Additional Observations
The connection between set RL and pass@k training objectives is clarifying: the authors show that under pass@n via set RL, incorrect generations always receive negative advantage, distinguishing their framework from leave-one-out approaches. The branching analysis (Fig. 3) provides compelling visual evidence that Poly-EPO induces fundamentally different generation behavior, with earlier divergence in reasoning trajectories.
Generated May 5, 2026
Comparison History (28)
Paper 2 likely has higher scientific impact due to its strong real-world applicability and breadth: a large-scale foundation model trained on nationwide claims enables immediate advances in disease prediction (1,000+ tasks), expenditure forecasting, and causal/RWE analyses, with external validation and demonstrated bias reduction in target trial emulation. The scale (43.8B events, up to 1.7B parameters) and rigorous retrospective/prospective evaluations support methodological robustness and timeliness for healthcare and regulatory use. Paper 1 is novel for RL-based exploratory reasoning, but its impact is narrower and more benchmark-centric.
Paper 1 addresses a core challenge in LLM post-training—improving exploration and reasoning diversity—with a concrete, implementable framework (Poly-EPO) backed by empirical results on reasoning benchmarks. This directly advances the rapidly growing field of LLM reasoning and test-time compute scaling, with immediate practical applications. Paper 2 presents an interesting theoretical impossibility result for AI governance, but its impact is limited by reliance on synthetic experiments, strong assumptions that may not map cleanly to real-world governance, and a more niche audience. Paper 1's methodology is more likely to be widely adopted and built upon.
Paper 1 addresses a core challenge in LLM post-training—improving exploration and reasoning diversity—with a practical algorithmic framework (Poly-EPO) demonstrating empirical gains across benchmarks. This directly impacts the rapidly growing field of reasoning model training and test-time compute scaling. Paper 2 presents an interesting theoretical impossibility result for AI governance, but its impact is limited by reliance on synthetic experiments, strong assumptions in its axiomatization, and the gap between formal proofs and real governance practice. Paper 1's methodological contributions are more immediately actionable and broadly applicable to the ML community.
Paper 1 addresses fundamental challenges in AI, specifically language model reasoning, exploration, and test-time compute scaling. Its methodological innovations in set reinforcement learning have broad applicability across the entire AI field. In contrast, Paper 2 presents an applied, domain-specific AI system for drilling operations. While practically valuable, its scientific impact is narrower and largely confined to the petroleum engineering sector.
Paper 2 introduces a novel and general training framework (Poly-EPO) for improving exploration in language model reasoning via set reinforcement learning, with broad applicability across reasoning tasks and test-time compute scaling. This addresses a fundamental challenge in LLM post-training with theoretical grounding and empirical results across multiple benchmarks. Paper 1, while technically competent, is a domain-specific engineering application (drilling operations) with limited generalizability beyond its niche. Paper 2's contributions to RL-based LM training methodology have far greater potential to influence the broader AI/ML research community.
Paper 2 has higher likely scientific impact due to broader cross-domain applicability: a general post-training framework (set RL) and Poly-EPO can influence many areas of LM training, reasoning, and test-time compute scaling. Its methodological contribution (adapting RL advantage for set objectives) is more foundational and reusable than a domain-specific multi-agent pipeline. Paper 1 is timely and impactful for digital health, but its advances are more applied and likely constrained by dataset/clinical validation requirements, limiting breadth and near-term generalization beyond wearables.
While Paper 1 is a timely survey that will likely attract many citations, Paper 2 presents a fundamental methodological advancement in post-training language models. By introducing a novel set reinforcement learning framework to enhance exploratory reasoning and scale test-time compute, Paper 2 directly addresses the current frontier of AI research, offering higher potential for transformative scientific impact and real-world performance gains in complex problem-solving.
Paper 2 has higher potential scientific impact due to its broad methodological advancements in foundational AI. While Paper 1 offers a highly valuable clinical application, Paper 2 addresses a fundamental bottleneck in large language models: improving reasoning, generalization, and test-time compute scaling through set reinforcement learning. Because advanced LLM reasoning is critical to nearly all downstream AI applications across various scientific disciplines, the Poly-EPO framework has a substantially wider breadth of impact and aligns with the most urgent, timely challenges in artificial intelligence research.
Poly-EPO introduces a fundamentally novel training paradigm (set RL with exploration-exploitation synergy) that addresses a core challenge in LM reasoning—scaling test-time compute through diverse exploration. Its contributions are methodologically clean, broadly applicable across reasoning tasks, and directly advance the rapidly growing field of reasoning model training. Paper 1 (TRUST) is ambitious but attempts to solve too many problems simultaneously (decentralized verification, consensus, privacy, attribution) with moderate empirical gains (72.4% accuracy). Poly-EPO's focus on improving how models reason, with clear generalization and scaling benefits, has broader foundational impact.
Paper 2 addresses a highly timely and impactful topic—training reasoning LMs with exploration strategies—at the intersection of RL and large language models. The Poly-EPO framework introduces a novel set RL approach that improves generalization and test-time compute scaling, with broad applicability across reasoning tasks. Paper 1, while technically solid in advancing exact cover counting with a new parallel algorithm and decision-ZDNNF representation, addresses a more niche combinatorial problem with narrower impact. The explosive interest in LM reasoning and RL-based training gives Paper 2 significantly greater potential for broad scientific influence.
Paper 2 likely has higher scientific impact: it introduces a general set-RL post-training framework and a concrete algorithm (Poly-EPO) that targets a broadly important capability—exploratory reasoning that scales with test-time compute. This is timely for reasoning LMs and could influence training methods across tasks (planning, search, code, theorem proving), with clear methodological novelty in adapting RL advantage computation to set objectives and emphasizing diversity/coverage (pass@k). Paper 1 is practical and useful for tool-calling agents, but is more architectural/metrics-focused and narrower in scope.
Paper 1 introduces a novel theoretical framework (set RL, Poly-EPO) for training language models with exploratory reasoning, addressing fundamental challenges in RL-based LM post-training. Its contributions—new optimization methods, improved generalization, and scalable test-time compute—are broadly applicable across reasoning tasks and advance core ML methodology. Paper 2, while practically useful, is primarily an application/system paper integrating existing components (stuttering classification, LLM agents) into a clinical workflow, with narrower scope and more incremental technical contributions.
Paper 1 targets a timely, high-impact area—post-training and inference-time scaling for language-model reasoning—introducing a set-RL framework and a concrete algorithm (Poly-EPO) aimed at improving diversity, generalization, and pass@k coverage, with broad applicability across reasoning tasks and LM training paradigms. Paper 2 offers a solid algorithmic advance (decision-ZDNNF and parallel exact-cover counting) with strong rigor, but its impact is likely narrower to knowledge compilation/combinatorial counting. Given current momentum and cross-field applicability, Paper 1 has higher estimated scientific impact.
Paper 1 proposes a more novel, paradigm-shifting framework: open-ended coevolution of both tasks and LLM populations, enabling continual capability discovery in a single run and yielding archives of specialized experts with strong coverage without direct benchmark optimization. This has broad potential applications (automated model discovery, efficient expert ensembles, continual learning) and could influence multiple fields (AutoML, multi-agent systems, evolutionary computation, synthetic data/task generation). Paper 2 is timely and methodologically solid but is more incremental within post-training/RL, with impact primarily on reasoning and test-time compute scaling.
Paper 2 is likely higher impact: it introduces a general set-RL post-training framework and a concrete algorithm (Poly-EPO) that targets a fundamental limitation in LM reasoning—insufficient exploratory diversity—while improving pass@k scaling with test-time compute. The methodological contribution (adapting RL via modified advantage for set objectives) is broadly reusable across tasks, models, and labs, and connects to core RL exploration theory, potentially influencing multiple fields (LLM alignment, reasoning, agents). Paper 1 is impactful for systems/cost orchestration, but is more engineering-centric and may be less foundational.
Paper 2 has higher potential scientific impact due to greater novelty and breadth: it introduces a general set-RL post-training recipe (adapting standard RL via advantage modification) and a concrete algorithm (Poly-EPO) targeting exploration–exploitation in reasoning. This could influence RLHF/post-training, test-time compute scaling, and diverse decoding across many domains. Paper 1 is highly practical and timely for cost-effective multi-model orchestration, but it is more systems/engineering-oriented and may be easier to reproduce as an incremental orchestration principle rather than a new learning framework. Both are relevant; Paper 2 likely generalizes more broadly.
Paper 1 is more likely to have higher scientific impact due to greater methodological novelty and broader applicability: it proposes a general set-RL post-training recipe plus a specific algorithm (Poly-EPO) that improves reasoning generalization, diversity, and test-time compute scaling—capabilities relevant across many LM tasks and fields. Paper 2 is timely and useful for GIScience workflows, but its contributions are more systems/engineering- and domain-specific, with evaluation largely via case studies, which may limit perceived rigor and breadth of downstream impact.
Poly-EPO tackles a highly timely and critical area in LLM research: test-time compute scaling and exploration. While HiPO offers a practical modification to DPO for segment-wise reasoning, Poly-EPO introduces a more foundational framework (set RL) that explicitly addresses generalization, response diversity, and test-time scaling. This broad applicability and strong alignment with current frontier developments in reasoning models give Paper 2 a significantly higher potential for broad scientific impact.
While Paper 1 offers a foundational improvement to deep learning optimization, Paper 2 tackles the highly critical and timely frontier of exploratory reasoning and test-time compute scaling in language models. By introducing set RL to synergize exploration and exploitation, Paper 2 directly addresses the current bottleneck in scaling LLM reasoning capabilities, promising profound and immediate impact across AI research and real-world deployment of reasoning agents.
Paper 2 likely has higher impact: it introduces a principled, training-free inference algorithm (APPS) that can be applied immediately to any base LLM, improving accuracy–compute trade-offs via particle-based approximation to sequence-level power sampling with future-value guidance. This is timely given strong interest in test-time compute and “reasoning without training,” and its methodological framing (SMC/particle methods + decoding) is broadly relevant across ML, probabilistic inference, and NLP. Paper 1 is innovative but depends on post-training with set RL, making deployment heavier and narrower in applicability.