Polychromic Objectives for Reinforcement Learning (2509.25424v1)

Abstract: Reinforcement learning fine-tuning (RLFT) is a dominant paradigm for improving pretrained policies for downstream tasks. These pretrained policies, trained on large datasets, produce generations with a broad range of promising but unrefined behaviors. Often, a critical failure mode of RLFT arises when policies lose this diversity and collapse into a handful of easily exploitable outputs. This convergence hinders exploration, which is essential for expanding the capabilities of the pretrained policy and for amplifying the benefits of test-time compute scaling. To address this, we introduce an objective for policy gradient methods that explicitly enforces the exploration and refinement of diverse generations, which we call a polychromic objective. We then show how proximal policy optimization (PPO) can be adapted to optimize this objective. Our method (1) employs vine sampling to collect on-policy rollouts and (2) modifies the advantage function to reflect the advantage under our new objective. Experiments on BabyAI, Minigrid, and Algorithmic Creativity show that our method improves success rates by reliably solving a larger set of environment configurations and generalizes better under large perturbations. Moreover, when given multiple attempts in pass@$k$ experiments, the policy achieves substantially higher coverage, demonstrating its ability to maintain and exploit a diverse repertoire of strategies.

• The paper's main contribution is the introduction of a polychromic objective that counteracts entropy collapse by enforcing diversity in reward optimization.
• It proposes a set RL framework and a polychromic PPO algorithm that use vine sampling and set-level advantage estimation to boost pass@k performance and robustness.
• Empirical results across BabyAI, Minigrid, and Algorithmic Creativity tasks show enhanced exploration, generalization, and multi-strategy behavior in challenging environments.

Polychromic Objectives for Reinforcement Learning: A Technical Analysis

Motivation and Problem Statement

Reinforcement learning fine-tuning (RLFT) of pretrained policies is a standard approach for adapting expressive generative models to downstream tasks. However, a persistent failure mode is entropy collapse: RLFT often drives the policy to concentrate probability mass on a narrow set of high-reward behaviors, sacrificing the diversity present in the pretrained distribution. This collapse impedes exploration, limits generalization, and undermines the benefits of test-time compute scaling, as evidenced by pass@$k$ metrics where RL-fine-tuned models can underperform their pretrained counterparts for large $k$.

The paper introduces a new objective—termed the polychromic objective—for policy gradient methods, designed to explicitly enforce exploration and refinement of diverse generations. The approach is instantiated via a modification of Proximal Policy Optimization (PPO), yielding the polychromic PPO algorithm. The method is evaluated on BabyAI, Minigrid, and Algorithmic Creativity environments, demonstrating improved success rates, diversity, and robustness to state perturbations.

Set Reinforcement Learning and the Polychromic Objective

The core technical contribution is the set reinforcement learning (set RL) framework, which generalizes the standard RL objective from single trajectories to sets of trajectories. The set RL objective is:

$$\max_\theta \mathbb{E}_{\tau_{1:n} \sim \pi_\theta(\cdot \mid s_0)}\left[ f(s_0, \tau_1,\cdots, \tau_n)\right]$$

where $f$ is a set-level objective function. This formulation enables direct optimization for properties such as diversity and coverage, which are not naturally captured by single-trajectory reward maximization.

The polychromic objective is a specific instantiation:

$$f_\mathrm{poly}(s, \tau_{1:n}) := \frac{1}{n} \sum_{i=1}^n R(\tau_i) \cdot d(s, \tau_{1:n})$$

where $R(\tau_i)$ is the (normalized) return of trajectory $\tau_i$, and $d(s, \tau_{1:n})$ is a diversity metric over the set. The diversity function is agnostic to the environment and can be instantiated as, e.g., the fraction of semantically distinct trajectories.

This objective ensures that the policy is rewarded only when it produces sets of trajectories that are both successful and diverse, thus directly counteracting entropy collapse.

Algorithmic Instantiation: Polychromic PPO

The polychromic PPO algorithm adapts PPO to optimize the set RL objective. The key modifications are:

• Vine Sampling: Instead of sampling $n$ actions at every state (which is computationally infeasible), the algorithm uses vine sampling to collect multiple rollouts from selected "rollout states." This enables estimation of set-level advantages without exponential sample complexity.
• Set-Level Advantage Estimation: At each rollout state, the advantage for all actions in a set is computed as the difference between the set's polychromic objective value and a baseline (Monte Carlo estimate over multiple sets). All actions in the set receive the same update signal, aligning with the set RL framework.
• KL Penalty: A per-state KL penalty is included for stability, anchoring the updated policy to the behavior policy.

The algorithm is compatible with any diversity metric and can be combined with UCB-style bonuses for further exploration incentives.

Empirical Results

Performance and Diversity

Polychromic PPO is evaluated on BabyAI, Minigrid, and Algorithmic Creativity tasks. The main findings are:

• Success Rate and Reward: Polychromic PPO matches or exceeds the best baselines in average reward and success rate across tasks. The addition of UCB bonuses is complementary, further improving performance in some environments.
• Pass@$k$ Coverage: Polychromic PPO achieves substantially higher pass@$k$ rates than both standard PPO and REINFORCE, especially as $k$ increases. This indicates that the policy maintains and exploits a diverse repertoire of strategies, in contrast to baselines that saturate quickly or even underperform the pretrained policy at large $k$.

Figure 1: Pass@k on BabyAI tasks. Top: methods without UCB. Bottom: methods with UCB. Columns show Goto, Pickup, Synthseq, and Bosslevel. Each curve is pass rate vs. number of attempts.

• Algorithmic Creativity: On the triangle discovery task, polychromic PPO achieves higher diversity and creativity metrics, with only a modest reduction in validity compared to PPO. The method outperforms all baselines in generating unique, valid, and creative solutions.

Figure 2: Pass@k results on Algorithmic Creativity. For validity pass@k and creativity pass@k, the agent gets a pass if at least one of the k attempts was a valid and creative triangle, respectively. In diff@k evaluation, we evaluate the number of generations that were unique given k attempts.

Generalization to State Perturbations

Polychromic PPO demonstrates superior generalization under large initial-state perturbations. When evaluated from novel starting positions in BabyAI environments, the method achieves higher pass@1 rates than all baselines, indicating that the learned policy is robust to distributional shifts and does not overfit to specific initial conditions.

Figure 3: Generalization under state perturbations in BabyAI BossLevel environment. The mission here is ``Pickup a key". The figure shows successful rollouts across perturbed initial states (blue circles), highlighting diverse strategies learned by the agent.

Qualitative Analysis

Visualization of rollouts in BabyAI environments shows that polychromic PPO learns to solve tasks via multiple distinct strategies, in contrast to baselines that converge to a small set of behaviors.

Figure 4: GoTo: mission “go to purple box.”

Theoretical Analysis: Entropy and Scaffold Value

The paper provides a detailed analysis of entropy dynamics under the polychromic objective. The key result is that the polychromic objective prevents entropy collapse onto homogeneous sets of actions, instead channeling probability mass toward heterogeneous sets that balance reward and diversity. The scaffold value is introduced as a measure of the policy's propensity to collapse onto a set; the analysis shows that homogeneous sets have negative scaffold value (discouraging collapse), while heterogeneous sets with multiple successful actions have positive scaffold value (attracting probability mass).

This theoretical framework generalizes the performance difference lemma to set RL and provides formal guarantees that the polychromic objective enforces a trade-off between exploitation and exploration at the set level.

Implementation Considerations

• Vine Sampling: Requires the ability to reset the environment to arbitrary states. For environments without this capability, approximate replay or deterministic transitions are necessary.
• Diversity Metric: The choice of diversity function is critical and may be nontrivial in high-dimensional or continuous control settings.
• Sample Complexity: While vine sampling mitigates exponential sample requirements, the method still incurs higher computational cost than standard PPO due to the need for multiple rollouts per state.
• Advantage Estimation: Monte Carlo estimates are used for the set-level baseline, which may introduce high variance; more efficient estimators or variance reduction techniques could further improve stability.

Implications and Future Directions

The polychromic objective and set RL framework provide a principled approach to maintaining and refining diversity during RLFT. This has significant implications for:

• Generalization: Policies trained with polychromic objectives are more robust to distributional shifts and can better exploit test-time compute via diverse sampling.
• Exploration-Exploitation Trade-off: The method offers a direct mechanism for balancing exploration and exploitation at the trajectory level, beyond local entropy regularization.
• Scalability: While demonstrated on grid-world and algorithmic tasks, extending the approach to high-dimensional, continuous, or real-world domains will require advances in diversity metrics and scalable sampling strategies.

Potential future work includes learning diversity functions, developing more efficient advantage estimators, and integrating curriculum or annealing schemes to balance exploration and exploitation dynamically.

Conclusion

The polychromic objective and its instantiation via polychromic PPO represent a substantive advance in reinforcement learning fine-tuning, addressing the critical issue of entropy collapse by directly optimizing for sets of diverse, successful behaviors. Theoretical analysis and empirical results demonstrate that this approach yields policies with superior coverage, robustness, and generalization, providing a foundation for further research in set-level RL objectives and diversity-aware policy optimization.