Polychromic Objectives for Reinforcement Learning
Authors: Jubayer Ibn Hamid
Ifdita Hasan Orney
Ellen Xu
Chelsea Finn
Dorsa Sadigh
Venue: ICLR 2026 Poster (2026)
Keywords: Reinforcement Learning
Exploration
Primary area: reinforcement learning

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.

Published as a conference paper at ICLR 2026
POLYCHROMIC OBJECTIVES FOR REINFORCEMENT LEARNING
Jubayer Ihn Hamid", Ifdita Hasan Orney", Ellen Xu, Chelsea Finn, Dorsa Sadigh Stanford University
ABSTRACT
Reinforcement learning fine-tuning (RLFT) is dominant paradigm for improv- ing pretrained policies for downstream tasks. These pretrained policies, trained on large datasets, produce generations with a broad range of promising but un- refined behaviors. Often, critical failure mode of RLFT arises when policies lose this diversity and collapse into 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 scal- ing. To address this, we introduce an objective for policy gradient methods that explicitly enforces the exploration and refinement of diverse generations, which we call 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, Mini- grid, 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 diverse repertoire of strategies.
INTRODUCTION
Reinforcement learning fine-tuning (RLFT) is widely used to enhance the performance of pretrained models across diverse downstream domains. For instance, RLFT has been applied to steer large language models (LLMs) toward instruction following and complex reasoning (DeepSeek-A) et al., 2025: OpenAI et al., 2024; Qwen, 2025). A common thread across these settings is the availability of expressive generative models (i.c., pretrained distributions), trained on large and diverse datasets, that already exhibit broad repertoire of strategies. RLFT then refines these distributions by reinforcing the strategies that yield higher reliability and performance.
However, exploration during RLFT remains central challenge. Prior work (Cui et al., 2025; Zhao et al., 2025) has documented entropy collapse: instead of expanding their repertoire, fine-tuned policies concentrate probability mass on narrow of high-reward behaviors already present in the pretrained distribution, effectively sacrificing entropy and diversity. This limits exploration and prevents the discovery of alternative strategies that could expand the base model's capabilities. Empirically, this effect is captured by the pass @n metric, which measures the probability that at least one out of n independently sampled rollouts succeeds. When large, RL-fine-tuned models often underperform their pretrained counterparts because the latter retain greater diversity (Yue et al., 2025; Wu et al., 2025). Such diversity is practically important; it supports generalization to new tasks (Kumar et al., 2020) and amplifies test-time compute scaling (Snell et al., 2024).
The goal of this paper to study how to induce policies to explore and refine a diverse repertoire of generations through RLFT. Our key insight is that algorithms should optimize objectives that explicitly encourage exploration and refinement of the diverse generations already embedded in the pretrained distribution. Standard regularization techniques, such entropy bonuses, often induce local or token-level variation but fail to promote semantic or trajectory-level exploration and can be overshadowed by the RL objective. In contrast, we propose a unified formulation that directly
"qqual contribution. Correspondence [jubayer, ifdil 101 stsanford.edu.

---

Published as a conference paper at ICLR 2026
optimizes for a diverse set of successful behaviors, encouraging the policy to generate broad, varied trajectories rather than collapsing onto a few high-reward ones.
To this end, propose the framework of set reinforcement learning (set RL), where the objective is defined over a set of trajectories sampled independently and evaluated by a multi-sample objec- tive (Tang et al., 2025). Unlike standard RL, which maximizes the likelihood of a single optimal trajectory, set RL maximizes the likelihood of an optimal set of trajectories sampled independently under a set-level objective. Within this framework, combine reward and diversity by scoring sets highly only if they contain both successful and diverse trajectories. Optimizing policy with respect to this objective is a principled upproach towards encouraging the policy to explore and search for a diverse set of generations that also maximize reward. We then instantiate one such objective and show how proximal policy optimization (PPO) (Schulman et al., 2017b) can he adapted to optimize it effectively, yielding a practical algorithm we call polychromic PPO. We evaluate our method on BabyAl (Chevalier-Boiss 2019), Minigrid (Chevalier-Boissee 2025). Our results show that polychromic PPO achieves higher rewards and success rates, generates diverse trajectories that substantially improve pass @n pass@n coverage, and generalizes more robustly to perturbations in the initial state.
introduce polychromic objectives which
al., et al., 2023), and Algorithmic Creativity (Nagarajan et al.,
PRELIMINARIES
We consider a Markov decision process (MDP) defined by state space S, action space A, transition dynamics distribution p(s St, 0z), reward function : S x A + R, initial state distribution Pn and discount factor Y E (0, 1). In reinforcement learning (RL), the goal is to learn policy that maximizes the value ((mp) = Er-ro (
Er-r, [R(T)] where R(T) is the
(discounted) sum of rewards in trajectory т. One widely used RL algorithm is proximal policy optimization (PPO) (Schulman et al., 2017b) which, iteratively, collects rollouts under behavior policy F and updates policy To by constraining the divergence between the two policies from growing too large: letting r, no (a st)/mp(a; st), PPO optimizes an empirical estimate of the following:
Eennd ()arrep(aa))))) min (r=Å(s,, at), clip(r,, 1 e, e,1 +)) ]
where (s) the stationary state-visitation distribution under policy zz.
3 REINFORCING EXPLORATION DURING RLFT
We aim to address the problem of entropy collapse during RLFT through method that explicitly induces exploration by encouraging the generation of diverse trajectories. In $3.1, we introduce variant of RI. that allows for objectives beyond reward maximization and observe its various properties. This framework will provide us with way to optimize objectives that are beyond return maximization, such as objectives that also encourage exploration. In $3.2, we specify the objective used within this framework for that purpose, which we call a polychromic objective. In $3.3, we propose our practical algorithm for optimizing the objective.
3.1 SET REINFORCEMENT LEARNING
We introduce a variant of the standard RL setup in which, given an objective function, we optimize over a sel of trajectories. We call this framework set reinforcement learning (set RI.), where the goal is to solve:
max E EnaxEraatttlloo [f TT. TI. .T.].,
(2)
Here f(s.. Ti....,.7.. Tw)) is some objective function over trajectories Tim -.I.... Tn sam- pled independently from the policy. This is in contrast to standard RL where the problem, maxg E )o)) [. uses an objective function, R(T). defined over a single trajectory. The generality of set RL. makes it a powerful tool for objectives beyond sum of rewards. Set RL objec- tives can be optimized using policy gradient methods by noting that
VoE (- If TIm1)) = Er (-110)) [f((o: T2.n) - f(so)) ΣΣ Vo log (o(a [o(")] (3)
2

---

Published as a conference paper at ICLR 2026
where f(s) is variance-reduction baseline. In this paper, use the baseline f(su) (ry (s0: Fim)]. Note that the advantage term f Timn) f(0)) is shared across all trajectories in the set T1mn- This is key defining feature of the set RL framework; the log-probability gradient of all trajectories in a set must be multiplied by the same factor. This contrasts with Tang et al. (2025), which also use n-sample objectives but employs srajeii baselines (e.g., leave-one-out) leading to advantages of the form [00, TIsn) - [(00, Ti.--1, T1 so the update for trajectory depends on a baseline computed from the remaining trajectories, yielding individ- ualized credit assignment. In our case, uniform baseline provides a common update signal to all trajectories in the set. enabling optimization with respect to the quality of the entire ser of trajecto- ries. In other words, by definition, set RL does not distinguish between trajectories within a set, but instead optimizes the policy by comparing across sets whole. Note that the set RL framework can still be used to optimize the standard RL. objective by choosing (*u Ti=)) = Σi= R(T.)- However, the framework allows for a broader class of objectives, such inference-time objectives (Tang et al., 2025).
1:+
Notice that set RL is distinct from the framework of multi-objective reinforcement learning (Roijers et al., 2013). In multi-objectivg RL, we aim to maximize the same objective as in standard RL, except the rewards are vector-valued. As such, the objective still defined over single sampled trajectory, as opposed to being defined over a set of trajectories. While it is possible to consider vector- rewards in the set RL framework as well, in this paper we focus on optimizing scalar valued objectives of the form (so. Ti:n) E R.
Before we move on to our proposed algorithm, it is helpful to construct notion of value functions in the framework of set RL i.e., the expected sum of rewards as specified by the objective. We begin with a simplified (but impractical) setting. Suppose that at every state 8 encountered during an on-policy rollout, stemming out of every state that branch even further along timesteps. Then, given this set of actions, d1:.. taken from the state s, the policy gets the set reward ((s,)in)) with respect the objective function f. Although such a setup is impractical for long-horizon tasks, analyzing it will help us better understand what set RL algorithms should aim to achieve.
sample a set of actions, din. TT(- s). which lead to trajectories
Under this assumption, the data collection process naturally generates a state-visitation tree. Beginning at the root state 8p, each visited state s branches into n children, one for each sampled action. At depth t, the tree therefore contains n' states, denoted by s, s (n²))
For each state
.!!)) the corresponding set of n sampled actions is written (a.)(() Given this tree-structured rollout and assuming an infinit-hhorizon discounted return, we define the value functions associated with set RL as follows:
Definition 3.1 Given a policy n generating a state-visitation tree and a set objective f f S x An + R. the set value function V2 f) and the set Q-function Q4 (s, aimn f) are defined as
n!
vi(s;) = Er
*S(( (ac) =
(4)
n²
Q
aimi (s, E,
r f(s (a)1%)
sa=a
(5)
We use the notation V and Q# to distinguish it from the value function and Q-function in standard RL. Here, we assume that the discount factor y E (0, 4) to ensure values remain bounded the range is smaller since at we add the expected sum of rewards from n actions stemming out of each state. Intuitively, (s) is the expected discounted return of the entire state tree rooted at s, where the reward at each node is given by the set objective. In contrast, Qi (s, a)m) evaluates the expected return of the tree that begins at 8 with the specific action set aim. Note that, although we assumed sets at the action-level instead of at the trajectory-level, this setting is equivalent to the trajectory-
level set RL as in eq. (2) under the objective function F(s, TInn)) == ΣT_ En 7 f( (a)) in the finite-horizon setting (in the infinite horizon setting, we would have take the limit T - oc).
These definitions help us better understand the objective of set RL. In standard reinforcement learning, we want to learn the policy such that the expected return from trajectory is maximized.
3

---

Published as a conference paper at ICLR 2026
In set RL, we want learn the policy that maximizes the expected return from generated using our policy is maximized. Given these definitions, we have the following result which is an extension of the performance difference lemma (Kakade & Langford, 2002) to the set RL setting (see $D.1 for proof):
Lemma 3.2 Given any two policies no and Ta and a fixed initial state 80, under any set objective function f.
1 E snnde ( (a1ns)o)(11))) AA (8, armin)
Va do (80; () - V (50:
Similar to the standard reinforcement learning setting, this result says that if we update the policy n to Th such that, at all states visited by тe. the advantage A (s, Aimi f) = Q (s, 41: f) V (s) of a set of actions taken by our new policy Ta is positive, then we will get a policy that has strictly higher performance (as measured by value). This suggests that many of the principles underlying methods like PPO can be extended to the set RL paradigm well for policy improvement.
Having seen the generality of set RL, we now construct an objective function to be used within this framework that will allow us to induce exploration.
3.2 A PRACTICAL POLYCHROMIC OBJECTIVE
Our central construction is the notion of polychromic objectives that are aimed towards training an agent to explore and learn diverse set of behaviors. Intuitively, these are set objective functions Spoly S T + IR that jointly capture (1) the success of set of trajectories in terms of reward and (2) the degree to which the set exhibits exploration or diversity. While we later generalize this construction in $5, in this section we focus on the specific objective used in our algorithm and experiments:
Spoty (8, Tim)
Σ R(T,)d(s, Tim),
(6)
where R(T;) is the discounted sum of rewards in trajectory т; and d(s,, T1:n) is function that quantifies the diversity of trajectories within the set. We require that both R(T,) and d(s, Tim)) are normalized between 0 and 1. Because the set-RI. gradient uses a shared advantage for all trajectories in a set, this objective increases the likelihood of successful behaviors and diverse exploratory trajectories. Unlike prior approaches, the shared advantage term amplifies exploratory trajectories that do not (yet) yield high rewards, pushing the policy to discover diverse strategies.
Various diversity metrics have been studied and incorporated in reward functions in prior works; examples include the Vendi Score (Friedman & Dieng, 2023) and glassifier-uideed diversity (Zhang et al., 2025; Li et al., 2025). Our algorithm is designed to be agnostic to the choice of metric: given any diversity function, we evaluate the success and diversity of a set of trajectories and optimize the policy to maximize both in sets.
3.3 POLYCHROMIC PPO
In this section, present an algorithu for optimizing eq. (6) by modifying PPO, which is motivated by the extension of the performance difference lemma to the set RL framework (as shown in Lemma 3.2). We choose to modify PPO since learning fine-tuning (Ouyang et al., 2022; Stiennon et al., 2020) known for stability and greater sample-efficiency (Achiam, 2018). However, the modifications we propose to instantiate our algorithm can be used to modify other algorithms too, such as REINFORCE, in order to optimize eq. (6). Our approach differs from standard PPO in two key respects: the method using which we sample on- on-policy rollouts and the advantage function used in the update.
widely used algorithm for reinforcement
direct implementation of the definition of set advantage functions as in Lemma 3.2 would require sampling n actions from every visited state, leading to exponential data requirements. To avoid this, we instead rely on vine sampling (Schulman et al., 2017a; Kazemnejad et al., 2025) for on on-policy data collection. In vine sampling, after collecting an initial set of rollouts, {si...., sp} of the states visited, called rollout states. At each rollout state s,, we generate N additional rollouts (called vines) starting from 8; This procedure ensures that we obtain multiple
select subset
4

---

Published as a conference paper at ICLR 2026
states with independently sampled trajectory sets stemming out of them. The particular scheme we use is closely related to the vine sampling scheme in TRPO (Schulman et al., 2017a): details deferred to IB.1.1. Our algorithm, however, is compatible with any vime-sampling method that guarantees sufficient vine coverage. Note that, since vine sampling requires the ability to reset the environment, algorithm is only applicable environments where such resets are possible.
Given access to sets of trajectories from each rollout state, we can estimate the polychromic advantage. At rollout state 8; from which we generated N > n trajectories, estimate the polychromic advantage as
A3 at: [poty)
n Σ R(T.)d(BR Tim) (86i Spoty)
where at E T₁ for some i E I,...,n. Since PPO requires an advantage defined for individual actions, we assign to each action ar the advantage of the set 71mm that contains it. In other words, given a set of trajectories Timm all actions taken from & in this set receive the same update signal as desired in set RL. We use the following Monte Carlo estimate of the value baseline: V fpoly)
denotes the M independently sampled sets of trajectories starting from 8t. This unbiased estimate was sufficient for our experiments, but one can trade off variance further by using biased estimates which we leave to future work.
1 E fpoty (st, (4) ], where Timn1 i E (I,.... .M
M
For non-rollout states, the update remains the same as standard PPO, using generalized advantage estimation (GAE) (Schulman et al., 2018). As is often used in practical implementations of PPO, we also include per-state KI. penalty DKL ( s) ITO s)) at every state visited, which we found helpful for stability. The pseudocode is presented in algorithm 1, with modifications relative PPO highlighted; extended pseudocode and implementation details are given in appendix B.
Algorithm Polychromic PPO
1: for iteration 1, do
2: Collect trajectories under πæ; rollout vines T1:M from rollout states
3: if's s, rollout state then
4: Form sets, -1.-...9h IN of n trajectories from St =1= Spoly (8t, 95) (se, ax) gi
5: Set A(se, at) (8s. gi) E
6: else
7:
Compute A(se, a.) via GAE
8: end if
9: Update Te for K epochs on minibatches B by maximizing the PPO objective in eq. (1)
10: Set Rg f- TT
11:
end for
EXPERIMENTAL EVALUATION
Our experiments aim to answer two questions: (1) How does polychromic PPO, a set RL algorithm that explicitly encourages diverse trajectory generation, affect performance? More specifically, does improving diversity come at significant cost in accuracy and success rate? (2) Does polychromic PPO encourage the policy to explore and learn diverse behaviors? In particular, learning to solve a task through diverse generations should, ideally, increase the pass@n performance i.e., the probability of succeeding least once when given multiple attempts. (3) Does encouraging the policy to explore and maximize the diversity of generated trajectories help the policy to be more robust to perturbations in the state-viste visitavisittion distribution? To address these questions, we Snevaller-Boiivert et al., 2023), BabyAl (Chevalier-Boiiss and Algorithmic Creativity (Nagarajan al., 2025). We evaluate on these environments since they allow for diverse set of solutions We provide full environment and implementation details in $B, and briefly describe the environments below.
evaluate on Minigrid
et al., 2019),
Minigrid and BabyAl are grid-world platforms with multiple rooms populated with keys, balls, and distractors. The agent receives natural-language goals (e.g., open a red door and then go the ball your left after placing the grey ball next door). We pretrain expert demonstrations (Chevalier-Borts Boisvert et al., 2023; 2019); we then fine-tune and evaluate on 50 fixed configurations (each configuration specifies grid layout and mission pair). In the
policy
5

---

Published as a conference paper at ICLR 2026
Pretrained policy
REINFORCE w/ UCB
PPO w/
Poly-PPO
Poly-PPO
Environment
REINFORCE
UCB
(ours)
w/ UCB (ours)
Goto
(0,246, 34.2)
(0.533, 73.0)
(0.538. 73.4)
(0,406, 46.2)
(0.428. 47.4)
(0.575, 80.2)
(0.561, 76.2)
Pickup
(0,141, 21.4)
(0,259, 39.8)
(0.991,56.0)
(0.283, 33.4)
(0.243. 27.8)
(0.452, 63.2)
(0.486, 65.6)
Synthseq
(0.157, 20.2)
(0.325, 45.4)
(0,361, 47.8)
(0.277, 32.2)
(0.224, 26.2)
(0.341, 47.0)
(0.317, 43.2)
Bosslevel
10.212. 20.6)
(0.266, 33.4)
(0.286, 36.4)
(0.336, 38.8)
(0.310. 35.8)
(0.378, 45.2)
(0.379, 46.8)
Four Rooms
(0.469, 70.4)
(0.639, 89.6)
(0.672, 92.6)
(0.618, 89.2)
(0.502, 78.6)
(0.666, 92.4)
(0.667, 93.2)
Table I: Performance on BabyAl and MiniGrid tasks. Each entry is reported as (2, y), where æ is the average episodic reward and y is the success rate (in %). Results are averaged over 100 rollouts across 50 configurations and 3 random seeds.
triangle discovery task in Algorithmic Creativity, an agent must output sequences of triangles from undirected graphs that are not revealed to the agent in context; the agent must learn the graph from recalling data has already seen and through further interactions during RLFT. We pretrain on triangles and edges drawn from 10 graphs, and then fine-tune on 3 graphs.
We compare polychromic PPO (Poly-PPO) to REINFORCE with baseline (Williams, 1992) and standard PPO (Schulman et al., 2017b). We also compare with a UCB-style regularization (Azar et al., 2017) where we add λUCB min(1, as, a)- to every advantage A(s, a). Here, N(s, a) is the number of times action a was sampled from state and λUCB is a hyperparameter. For Poly- PPO, we define the diversity funetion d(ss, Ti:n) ) to be the fraction of distinet trajectories in Tims in Minigrid/BabyAl, two trajectories are called distinct if they visit different sets of rooms and, in Algorithmic Creativity, two trajectories are distinct if they visit different sets of nodes. In both cases, d = 0 if all trajectories visit the same set of rooms or nodes.
All of these are long-horizon, sparse reward settings. Since the pretrained policy struggles solve several of these long-horizon tasks, the policy must strategically explore during RLFT and avoid collapsing onto behaviors that solve only a subset while failing to generalize to the rest.
4.1 How DOES POLYCHROMIC PPO AFFECT PERFORMANCE?
The results summarizing performance, in terms of average reward and success rate, in Minigrid and Al are provided in table 1. We find that Poly-PPO consistently matches or outperforms the best baseline in reward and success. Adding the UCB bonus helps the baselines, REINFORCE and PPO, improve performance in some environments. We find that UCB is complementary to Poly-PPO as well the bonus enables the agent to achieve higher performance in Pickup and Bosslevel.
The results on Triangle Discovery are shown in fig. 2, validity is the number of valid triangles constructed. Diversity is the number of unique valid triangles and creativity metric is defined as the percentage of generations that unique valid triangles not present in pretraining data (Nagarajan et al., 2025). PPO substantially increases validity compared to the pretrained policy, but lower creativity and diversity. On the other hand, Poly-PPO achieves slightly lower validity than vanilla PPO, but the gap is modest and highlights the tradeoff it strikes between success and exploration which we discuss in the next subsection.
4.2 DOES POLYCHROMIC PPO ENCOURAGE DIVERSE GENERATIONS?
Success rates do not adequately represent the coverage of tasks that the policy can solve. Since success rates are averaged across all configurations, a policy that overfits to a subset may achieve high reward there while failing elsewhere. To probe this, we examine pass@k k curves for each configuration, we provide the policy & attempts and find the fractal of configurations the policy can solve, which is called the pass rate. As k increases, methods that generate diverse trajectories should, ideally, achieve greater coverage/pass rate.
@ass@k k results on the BabyAl environments are shown in fig. 1. We first discuss all methods without the UCB bonus. We observe that REINFORCE does not improve pass rate sufficiently fast as the number of attempts increases: despite higher success rates overall, its coverage is lower than the pretrained policy at large k. PPO, on the other hand, starts off from much lower pass rate than other methods at small l but the pass rate increases as grows; however, is still lower than the pretrained policy and significantly lower than Poly-PPO. This indicates that these bascline methods suffer from an inherent trade-off between diversity and accuracy in the generations. In comparison, Poly-PPO achieves substantially higher pass rate than all baselines. It also achieves

---

Published as a conference paper at ICLR 2026
Goto
Pickup
Synthseq
Bosslevel
113
Noo Nu1.
105
45%
304
(Passig))
k(PassgH)
[Passigi
((oassgK)
1000
1th
10%
10%
309
Pretrained
REINFORCE REINFORCE w/ UCB
PPO
PPO UCB
Poly-PPO
Poly-PPO w/ UCB
Figure 1: Pass on BabyAl tasks. Top: methods without UCB, Bottom: methods with UCB. Each curve is pass rate vs. number of attempts.
equal or higher pass rate than the pretrained policy at almost all values of k. Another indicator for the higher diversity in generations is that the pass rate for Poly-PPO continues to rise until about = 80, whereas the baselines saturate much earlier (around k = 20). The effect is pronounced, especially, in Bosslevel where Poly-PPO achieves around 15% higher pass rate as & grows to 160 whereas baselines see modest increase.
We next examine the effect of adding the UCB bonus. For REINFORCE, the bonus improves pass rate at small k in Pickup and Synthseq, but the gain saturates around = 10 and vanishes for larger k; has no effect in Bosslevel or Goto. For PPO, the bonus reduces pass at small k, and although it improves performance at larger k. the gap to the pretrained policy and Poly-PPO remains, By contrast, combining UCB with Poly-PPO yields equal or higher coverage across most le (except small le in Synth- seq), showing that Poly-PPO preserves and refines pre- trained diversity rather than collapsing onto narrow be- haviors.
1.201
laldit
Diversity
Ereativity
Normalzzed Score
0.75
0.50
122
0.25
0.50
Pretrained
PPO
REINH DRCE
PPly
Figure 2: Results on Algorithmic Creativ- ity. Bars show normalized values for each metric, with raw values above each bar.
In the Triangle Discovery task, We find that polychromic PPO achieves substantially higher diversity and creativity. In particular, Poly-PPO outperforms all baselines, includ- ing the pretrained policy, both diversity and creativity metries, shown in fig. 2. Although it achieves slightly lower validity than PPO (significantly larger than REINFORCE though), Poly-PPO encourages broad exploration and the discovery of novel solutions. This trend further reflected in the pass @ass@/ & evaluation (see fig. 3). Specifically, validity pass & measures whether at least one of the A attempts forms valid triangle; diff@ k quantifies the number of unique triangles obtained in attempts; and creativity pass @ & assesses whether at least one of the & attempts is creative. We find that Poly-PPO outperforms baselines in both creativity and diversity metries, while attaining greater validity performance than the pretrained policy. Notably, even though REINFORCE achieves high diversity, it comes at the significant cost in validity @ where even below the pretrained policy.
We find that Poly-PPO achieves substantially higher diversity and creativity in the triangle discov- cry task. As shown in fig. 2, it outperforms all baselines, including the pretrained policy, on both metrics, while maintaining competitive validity above the pretrained policy. This pattern holds in the pass@1 evaluation (see fig. 3). Overall, Poly-PPO surpasses baselines in creativity and diversity while attaining higher validity than the pretrained policy, REINFORCE and nearly the same PPO. Notably, although REINFORCE achieves high diversity, it does so at a steep cost: its validity @ falls even below that of the pretrained policy.

---

Published as a conference paper at ICLR 2026
100m
100%
110
100
60%
60%
70.0
50m
50%
30.0
00.0-
30m
30%
20.5
10.0
05
16 (Validity Pass@lk)
10
60
100
10 20
30
160
0.0-
10
20
80
160
k (Creativity Pass
k
(Diffka))
Pretrained
REINFORCE
PPO
Poly-PPO
Figure 3: @ass@k results on Algorithmic Creativity. For validity pass@k l: and creativity pass@ k the agent gets a pass if at least one of the l attempts valid and creative triangle, respectively. In diff@ evaluation, we evaluate the number of generations that were unique given k: attempts.
Pretrained policy
REINFORCE Baseline
REINFORCE w/ UCB
PPO UCB
Polychromic PPO
Poly-PPO
Rnvironment
PPO
w/ UCB
30.2
41.3
37.1
18.4
Pickup
15.2
22.0
20.5
12.5
8.87
33.4
28.0
Synthscq
20.0
19.3
26.2
16.6
11.5
30.6
32.1
Bosslevel
23.8
22.5
27.6
26.6
28.2
34.3
32,8
Four Rooms
65.0
82.7
81.5
15.3
14.2
88.7
87.2
Table Average pass rate (%) one attempt on BabyAI and Minigrid tasks under large initial-state perturbations.
4.3 DOES POLYCHROMIC PPO GENERALIZE TO STATE PERTURBATIONS?
We evaluate generalization under perturbed initial states in Minigrid and BabyAl. For each grid-mission configuration, we first find all the rooms visited by the pretrained policy under high- temperature sampling over 100 rollouts. Then, we select 10 states randomly from inside each room as our initial states. Effectively, this changes the initial state to a completely different room in such a manner that the task remains solvable from the new initial state. Note that this randomization substantially changes the task for the agent; as shown in fig. 5, with respect to the new initial state, a successful trajectory would require very different strategies. From the new start state, we evaluate using pass@) for all states in each layout. As shown in table 2, consistently, Poly-PPO generalizes more reliably than baselines under these perturbations.
ENTROPY ANALYSIS
In this section, analyze how the entropy of policy evolves when trained optimize the poly- chromic objective in eq. (6). Our guiding question is: Under the polychromic objective, on which actions a policy most likely collapse probability mass? In our analysis, we restrict our atten- tion to the bandit setting i.c., time horizon H = 1, with binary rewards. We assume discrete action space, and that the diversity function d(s,, a1:n) equals the fraction of actions 41 that are distinct (d = O if the set is a singleton). We assume that our policy has softmax parameterization. Before turning to the polychromic objective itself, we extend the entropy analysis Cui et al. (2025) to the set RL setting. In standard reinforcement learning with policy gradient methods, Cui et al. (2025) showed that, at state s, after one update to the policy, the change in the entropy of the policy at can be approximated
H( H( ( s) H(mm (s xi @Cov(log m (a s), A(s,a)
where is the learning rate. This result formalizes the understanding that policy collapses its entropy onto high-reward actions when there is strong covariance between the policy's probability mass on an action and the advantage of the action.
Given any set objective f S x A" - TR, we ask: how does the entropy of the policy change ++11 when learning under this set-RL framework? The following proposition characterizes the first-order change (proof in $5:1):
after one step update (from a to n
8

---

Published as a conference paper at ICLR 2026
Proposition 5.1 Consider the set-RL setup state After update the policy, the change in entropy, A H (Tg k+1 H (Tt ()), is given by
~ CCova
log n (as 8), Cov (f(s,
11a
1aa =
where both covariances are taken with respect to Th(- 8) and a is the learning rate.
This result provides lens for understanding when and where entropy collapse occurs. Intuitively, suppose that for some reference set din there is a strong covariance between (i) the overlap of aim with sampled sets and (ii) the value the objective f. As the policy concentrates more probability mass on such sets, the entropy decreases. Conversely, when the covariance is strongly negative, the policy reallocates probability mass to sets with higher value under f. which increases entropy. Thus, the central question becomes: which sets d1xx are most prone to entropy collapse under the polychromic objective? Our analysis proceeds by introducing and studying the following key object:
Definition 5.2
The scaffold value of a set of actions, a1m under a policy and a set- Re-RRL objective
f: S x A -> RR. is defined to be
n
Σ Ay(arm; T) = Cov (f) I(a1m) s.j=1
1d.=
where [()).) the maximum size of the intersection of aim with any other set a'
The scaffold represents, for every action set 01mx measure the policy's 's propensity collapse its entropy around that set. We illustrate this further through the following lemma which shows us how the scaffold value of an action set affects the change the probability of policy sampling the set(proof in 15.1)
Lemma 5.3 Consider any set of actions 61 E A The change in the log probability of sampling this set of actions after one policy update using set RL. can be written as the following first-order approximation:
log log=/+1(ass |s) = log (aa () aay((11n) aC((t)
where C(th) is function independent of 01:n-
With this apparatus in hand, we next show that the polychromic objective rules out collapse onto homogeneous sets of actions (proof in $D.4):
Proposition 5.4 where r(s,a) = 1, there exists e (0, 1) such that Spory (a) < O when (a 8) > e. Furthermore,
Consider the polychromic objective in eq. (6). For any homogenous sel din (a)
the scaffold values of these homogenous sets satisfy the bound Asprr (a)
p(1
This result shows that once successful action accumulates sufficient probability mass, the poly- chromic objective automatically prevents further entropy collapse onto sets that only contain this action and, when homogeneous set decreases. Before that threshold, the scaffold value of homogeneous sets this action may be positive, but it tightly bounded with the bound decreasing further as n grows. This is desirable since it suggests that our policy learns to generate this action without incessantly col- lapsing probability mass on and memorizing such homogeneous sets. Next, we analyze the scaffold values of heterogenous sets with successful actions (proof in $D.5):
larger sets in the set RL framework, the maximum scaffold value of
Proposition 5.5 Suppose aim is heterogeneous where each a, is unique with probability p E (0, 4). Suppose exactly q of the actions satisfy r(s, a.) L. and that any other action a' $ 01:n with To(aa 8) > 0 yieldsr a') = 0. Then, the scaffold value of 0mm satisfies A (av.)
(1-(1-p)
Note that the set in this proposition includes unsuccessful actions as well that contribute diversity. As such, there are, likely, several such heterogeneous sets with positive scaffold values that attract more

---

Published as a conference paper at ICLR 2026
probability mass than homogeneous sets. Moreover, lower bound guarantee increases the number of successful actions in the set increases. The polychromic objective therefore channels entropy collapse toward those sets that balance success and exploration, rather than permitting collapse onto homogeneous behaviors. Together, these results motivate our construction of polychromic objectives in general. Broadly, such objectives reward both (i) the returns achieved by set of trajectories and (ii) the diversity of the trajectories.
Definition 5.6 A polychromic objective is a function P:Sx S + IR that factors as Tim) = ((8 (s, Tim))pl (8, Ti:a), where RB) and so() are scalar-valued functions satisfying:
1. Cov Tov (p( (s, Tim): Ei=1 R(T))) > O.
2. Cov Covn (i,)) 1r= = =}} < 0 for any T, and
3. (( -) and (s, share the same range, i.e.,
inf u = inf inllg////// (s,)) ) o sup up (s.
where the infinum and supremum are taken over all sets T1.m-
In other words, polychromic objective factors into reward component with positive covariance to return and a diversity component with negative covariance to homogeneity. Their product enforces that both success and diversity are indispensable.
RELATED WORK
RL fine-tuning often suffers from entropy collapse, where policies concentrate on a few high-reward behaviors and lose coverage of the pretrained distribution (Cui et al., 2025; Wu et., 2025; Yue al., 2025). Entropy bonuses (Haarnoja et al., 2017; 2018; Schulman et al., 2017b; Seo et al., 2021; Islam et al., 2019) mitigate this locally but may struggle to induce semantic or trajectory-level exploration. Ecoffet et al. (2021) aims to explore more from promising states and trains a policy to robustly handle stochasticity. Other approaches, more similar our work, explicitly reward diversity, e.g., wiversit objectives (Li et al., 2025), UCB bonuses (He al., 2025; Lanchantin et al., 2025), batch-level exploration rewards (Song et al., 2025) or use covariance controls (Cui et al., 2025). The crucial distinction work makes is the set RL framework that enables learning set of behaviors wherein some trajectories receive positive gradient updates for exploration despite not contributing the maximum rewards attained in the set.
7 CONCLUSION
We framed the problem of learning diverse set of successful behaviors in terms of set reinforce- ment learning and proposed optimizing the polychromic objective, which evaluates sets of actions using both reward and diversity. We derived polychromic-PPO, variant of PPO that incorporates vine sampling and a modified advantage estimator to optimize this objective. There are several lim- itations of our approach. Firstly. our approach requires the ability to reset the environment any set to enable vine sampling. Furthermore, ensuring sufficient vine coverage can be computationally demanding. Future work could develop efficient, low-variance gradient estimators, or adopt curriculum and annealing schemes that balance exploration early in training with exploitation later.
ACKNOWLEDGMENTS
We would like to thank Amber Xie, Joey Hejna, Suvir Mirchandani, Hung- Chieh Fang and Andy Tang for feedback on an earlier version of the paper. We thank Ayush Chakravarthy for helpful comments on the experimental evaluations. We thank Yoonho Lee for several valuable comments on formalizing diversity. We thank Hamidur Rahman and Pervin Akhter for helpful guidance on the for- malization of the problem and feedback on the project in general. We thank Zabed Hasan and Fazilat Afreen for their insightful feedback experimental evaluations. This work is supported by an NSF CAREER award, NSF award # 1941722, ONR YIP. DARPA YFA award # W and Schmidt Sciences.
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Joshua Achiam. Spinning Up in Deep Reinforcement Learning. 2018.
Mohammad Gheshlaghi Azar, lan Oshand, and Rémi Munos. Minimax regret bounds for reinforce- ment learning, 2017. URL https:t/////// //axxiv..vvv/////a org/ wrggabs//// abs /133. 05449.
Maxime Shevalier-Boisveer Dzmitry Bahdanau, Salem Lahlou, Lucas Willems, Chitwan Saharia, Thien Huu Nguyen, and Yoshua Bengio. Babyai: A platform to study the sample efficiency of grounded language learning, 2019. URL https: /tttpp:://rxiiii /wwpiviivvv./////////////////ii/ /otg/or///////aa//.eeaae attp/b/////// 08272.
Maxime Chevalier- -hevaliie-Boiiveert, Bolun Dai, Mark Towers, Rodrigo de Lazcano, Lucas Willems, Salem Lahlou, Suman Pal, Pablo Samuel Castro, and Jordan Terry. Minigrid and miniworld: Modu- lar and customizable reinforcement learning environments for goal-oriented tasks, 2023. URL https: //rpxiv// http//aaxxv...g//...a/../ wttp//rrg/////////o 13831.
Ganqu Cui, Yuchen Zhang, Jiacheng Chen, Lifan Yuan, Zhi Wang. Yuxin Zuo, Haozhan Li, Yuchen Fan, Huayu Chen, Weize Chen, Zhiyuan Liu, Hao Peng, Lei Bai, Wanli Ouyang, Yu Cheng, Bowen Zhou, and Ning Ding. The entropy mechanism of reinforcement learning for reasoning language models, 2025. URL https: /ttpaaaaiiii/iii///ii//aaa/aawia wwwpexiwww....eeeeeeeeeeeeeeeeee httg//bb///// /ttpa/as/////../ /abp///5..221..
DeepSeck-AL, AI, Daya Guo, Dejian Yang. Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, Xiaokang Zhang, Xingkai Yu, Yu Wu, Z. F. Wu, Zhibin Gou, Zhihong Shao, Zhuoshu Li, Ziyi Gao, Aixin Liu, Bing Xue, Bingxuan Wang. Bochao Wu, Bei Feng, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, Damai Dai, Deli Chen, Dongjie Ji, Erhang Li, Fangyun Lin, Fucong Dai, Fuli Luo, Guangbo Hao, Guanting Chen, Guowei Li, H. Zhang. Han Bao, Hanwei Xu, Haocheng Wang. Honghui Ding, Huajian Xin, Huazuo Gao. Hui Qu, Hui Li, Jianzhong Guo, Jiashi Li, Jiawei Wang, Jingchang Chen. Jingyang Yuan, Junjie Qiu, Junlong Li, J. Cai, Jiaqi Ni, Jian Liang, Jin Chen, Kai Dong, Kai Hu, Kaige Gao, Kang Guan, Kexin Huang, Kuai Yu, Lean Wang. Lecong Zhang, Liang Zhao, Litong Wang. Liyue Zhang, Lei Xu, Leyi Xia, Mingehuan Zhang, Minghua Zhang, Minghui Tang, Meng Li, Miaojun Wang, Mingming Li, Ning Tian, Panpan Huang, Peng Zhang, Qiancheng Wang. Qinyu Chen, Qiushi Du, Ruiqi Ge, Ruisong Zhang, Ruizhe Pan, Runji Wang, R. J. Chen, R. L. Jin, Ruyi Chen, Shanghao Lu, Shangyan Zhou, Shanhuang Chen, Shengfeng Ye, Shiyu Wang. Shuiping Yu, Shunfeng Zhou, Shuting Pan, S. S. Li, Shuang Zhou, Shaoqing Wu, Shengfeng Ye, Tao Yun, Tian Pei, Tianyu Sun, T. Wang, Wangding Zeng, Wanjia Zhao, Wen Liu, Wenfeng Liang, Wenjun Gao, Wenqin Yu, Wentao Zhang, W. L. Xiao, Wei An, Xiaodong Liu, Xiaohan Wang, Xiaokang Chen, Xiaotao Nie, Xin Cheng, Xin Liu, Xin Xie, Xingchao Liu, Xinyu Yang, Xinyuan Li, Xuecheng Su. Xuheng Lin, X. Q Li. Xiangyue Jin, Xiaojin Shen, Xi- aosha Chen, Xiaowen Sun, Xiaoxiang Wang, Xinnan Song, Xinyi Zhou, Xianzu Wang, Xinxia Shan, Y. K. Li, Y. Q Wang. Y. X. Wei, Yang Zhang, Yanhong Xu, Yao Li, Yao Zhao, Yaofeng Sun, Yaohui Wang, Yi Yu, Yichao Zhang. Yifan Shi, Yiliang Xiong, Ying He, Yishi Piao, Yisong Wang, Yixuan Tan, Yiyang Ma, Yiyuan Liu, Yongqiang Guo, Yuan Ou, Yuduan Wang, Yue Gong, Yuheng Zou, Yujia He, Yunfan Xiong, Yuxiang Luo, Yuxiang You, Yuxuan Liu, Yuyang Zhou, Y. X. Zhu, Yanhong Xu, Yanping Huang, Yaohui Li, Yi Zheng, Yuchen Zhu, Yunxian Ma. Ying Tang, Yukun Zha, Yuting Yan, Z. Z. Ren, Zehui Ren, Zhangli Sha. Zhe Fu. Zhean Xu, Zhenda Xie, Zhengyan Zhang, Zhewen Hao, Zhicheng Ma, Zhigang Yan, Zhiyu Wu, Zihui Gu, Zijia Zhu, Zijun Liu, Zilin Li, Ziwei Xie, Ziyang Song, Zizheng Pan, Zhen Huang, Zhipeng Xu, Zhongyu Zhang, and Zhen Zhang. Deepseek-rl: Incentivizing reasoning capability in llms via reinforce- 12948.
ment learning, 2025. URL https: http://wwww..////...a/////o wwt./oggg////oo.ooo.oo
Adrien Ecoffet, Joost Huizinga, Joel Lehman, Kenneth O. Stanley, and Jeff Clune. Go-explore: new approach for hard-exploration problems, 2021. URL 19C1. 10995.
Dan Friedman and Adji Bousso Dieng. The vendi score: A diversity evaluation metric for machine learning, 2023. URL https: // wwwaaaivvvvoo.....o.o abs abs//b/l.cc... 02410.
Tuomas Haamojj, Haoran Tang. Pieter Abbeel, and Sergey Levine. Reinforcement learning with deep energy-based policies, 2017. URL https: //tp://///xiviii wwt.v/oog.g////oo/ 702. 081 65.
11

---

Published as a conference paper at ICLR 2026
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforeement learning with a stochastic actor, 2018. URL https:
//arxiv. wwg. abs 11001.111100 01290.
Andre He, Daniel Fried, and Sean Welleck. Rewarding the unlikely: Lifting grpo beyond distribution sharpening, 2025. URL https: /tap:s://arae /ttp/vaavvvvvvva///////../a/.av. org/ abs/2b06.66.. 02355.
Riashat Islam, Zafarali Ahmed, and Doina Precup. Marginalized state distribution entropy regular- htt.ab///9....... 05128.
ization in policy optimization, 2019. URL
Sham M. Kakade and John Langford. Approximately optimal approximate reinforcement learn- ing. In International Conference on Machine Learning, 2002. URL https://api.
semanticscholest htt.sschhoo.aa..oooooooooooaaa:. 311444999
Amirhossein Kazemnejad, Milad Aghajohari, Eva Portelance, Alessandro Sordoni, Siva Reddy, Aaron Courville, and Nicolas Le Roux. Vineppo: Refining credit assignment in rl training of 1lms, 2025. URL https: / arwaa.iiivv.... orgivv.orr//aaa absgg//bs//// abs//410. 01679.
Saurabh Kumar, Aviral Kumar, Sergey Levine, and Chelsea Finn. One solution is not all you need: Few-shot extrapolation via structured maxent rl, 2020. URL 2010. 14
org/ abs/
Jack Lanchantin, Angelica Chen, Shehzaad Dhuliawala, Ping Yu, Jason Weston, Sainbayar Sukhbaatar, and Ilia Kulikov. Diverse preference optimization, 2025. URL httttt:s///aattaae hag//bb/2501 18101.
Tianjian Li, Yiming Zhang, Ping Yu, Swarnadeep Saha, Daniel Khashabi, Jason Weston, Jack Lan- chantin, and Tianlu Wang. Jointly reinforcing diversity and quality in language model generations, 2025. URL https: : //tp//aavvv/////////////////aav/ org/ abs /2509. 02534.
Vaishnavh Nagarajan, Chen Henry Wu, Charles Ding, and Aditi Raghunathan. Roll the dice and look before you leap: Going beyond the creative limits of next-token prediction, 2025. URL ht https: /atxxvivvvv///////////.////aaav. org/ abs 2504. 15266.
OpenAI, :, Aaron Jaech, Adam Kalai, Adam Lerer, Adam Richardson, Ahmed El- El-Kishky, Aiden Low, Alee Helyar, Aleksander Madry, Alex Beutel, Alex Carney, Alex Iftimie, Alex Karpenko, Alex Tachard Passos, Alexander Neitz, Alexander Prokofiev, Alexander Wei, Allison Tam, Ally Bennett, Ananya Kumar, Andre Saraiva, Andrea Vallone, Andrew Duberstein, Andrew Kondrich, Andrey Mishchenko, Andy Applebaum, Angela Jiang, Ashvin Nair, Barret Zoph, Behrooz Ghor- bani, Ben Rossen, Benjamin Sokolowsky, Boaz Barak, Bob McGrew, Borys Minaiev, Botao Hao, Bowen Baker, Brandon Houghton, Brandon McKinzie, Brydon Eastman, Camillo Lugaresi, Cary Bassin, Cary Hudson, Chak Ming Li, Charles de Bourcy, Chelsea Voss, Chen Shen, Chong Zhang, Chris Koch, Chris Orsinger, Christopher Hesse, Claudia Fischer, Clive Chan, Dan Roberts, Daniel Kappler, Daniel Levy, Daniel Selsam, David Dohan, David Farhi, David Mely, David Robinson, Dimitris Tsipras. Doug Li, Dragos Oprica, Eben Freeman, Eddie Zhang, Edmund Wong, Eliz- abeth Prochl, Enoch Cheung, Eric Mitchell, Eric Wallace, Erik Ritter, Evan Mays, Fan Wang, Felipe Petroski Such, Filippo Raso, Florencia Leoni, Foivos Tsimpourlas. Francis Song, Fred von Lohmann, Freddie Sulit, Geoff Salmon, Giambattista Parascandolo, Gildas Chabot, Grace Zhao, Greg Brockman, Guillaume Leclere, Hadi Salman, Haiming Bao, Hao Sheng, Hart An- drin, Hessam Bagherinezhad, Hongyu Ren, Hunter Lightman, Hyung Won Chung, Ian Kivlichan, Ian O'Connell, Ian Osband, Ignasi Clavera Gilaberte, 1lge Akkaya, Ilya Kostrikov, llya Sutskever, Irina Kofman, Jakub Pachocki, James Lennon, Jason Wei, Jean Harb, Jerry Twore, Jiacheng Feng, Jiahui Yu, Jiayi Weng, Jie Tang. Jieqi Yu, Joaquin Quiñonero Candela, Joe Palermo, Joel Parish, Johannes Heidecke, John Hallman, John Rizzo, Jonathan Gordon, Jonathan Uesato, Jonathan Ward, Joost Huizinga, Julie Wang, Kai Chen, Kai Xiao, Karan Singhal, Karina Nguyen, Karl Cobbe, Katy Shi, Kayla Wood, Kendra Rimbach, Keren Gu-Lemberg, Kevin Liu, Kevin Lu, Kevin Stone, Kevin Yu, Lama Ahmad, Lauren Yang. Leo Liu, Leon Maksin, Leyton Ho, Liam Fedus, Lilian Weng, Linden Li, Lindsay McCallum, Lindsey Held, Lorenz Kuhn, Lukas Kon- draciuk, Lukasz Kaiser, Luke Metz, Madelaine Boyd, Maja Trebacz, Manas Joglekar, Mark Chen, Marko Tintor, Mason Meyer, Matt Jones, Matt Kaufer, Max Schwarzer, Meghan Shah, Mehmet Yatbaz, Melody Y. Guan, Mengyuan Xu, Mengyuan Yan, Mia Glaese Mianna Chen, Michael Lampe, Michael Malek, Michele Wang, Michelle Fradin, Mike McClay, Mikhail Paylov, Miles
12

---

Published as a conference paper at ICLR 2026
Wang, Mingxuan Wang. Mira Murati, Mo Bavarian, Mostafa Rohaninejad, Nat McAleese, Neil Chowdhury, Neil Chowdhury, Nick Ryder, Nikolas Tezak, Noam Brown, Ofir Nachum, Oleg Boiko, Oleg Murk. Olivia Watkins. Patrick Chao, Paul Ashbourne, Pavel Izmailov, Peter Zhokhov, Rachel Dias, Rahul Arora, Randall Lin, Rapha Gontijo Lopes, Raz Gaon, Reah Miyara, Reimar Leike, Renny Hwang. Rhythm Garg, Robin Brown, Roshan James, Rui Shu, Ryan Cheu, Ryan Greene, Saachi Jain, Sam Altman, Sam Toizer, Sam Toyer, Samuel Miserendino, Sandhini Agar- wal, Santiago Hernandez, Sasha Baker, Scott McKinney, Scottie Yan, Shengjia Zhao, Shengli Hu, Shibani Santurkar, Shraman Ray Chaudhuri, Shuyuan Zhang, Siyuan Fu, Spencer Papay, Steph Lin, Suchir Balaji, Suvansh Sanjeev, Szymon Sidor, Tal Broda, Aidan Clark, Tao Wang. Tay- lor Gordon, Ted Sanders, Tejal Patwardhan, Thibault Sottiaux, Thomas Degry, Thomas Dimson, Tianhao Zheng, Timur Garipov, Tom Stasi, Trapit Bansal, Trevor Creech, Troy Peterson, Tyna Eloundou, Valerie Qi. Vineet Kosaraju, Vinnie Monaco, Vitchyr Pong, Vlad Fomenko, Weiyi Zheng, Wenda Zhou, Wes McCabe, Wojciech Zaremba, Yann Dubois, Yinghai Lu, Yining Chen, Young Cha, Yu Bai, Yuchen He, Yuchen Zhang, Yunyun Wang, Zheng Shao, and Zhuohan Li. Openai ol system card, 2024. URL https::///.......t arxiv. wrg. abs 222224122..22 1.412.16772 16720.
Long Ouyang, Jeff Wu, Xu Jiang, Diogo Almeida, Carroll L. Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, John Schulman, Jacob Hilton, Fraser Kel- ton, Luke Miller, Maddie Simens, Amanda Askell, Peter Welinder, Paul Christiano, Jan Leike, and Ryan Lowe. Training language models to follow instructions with human feedback, 2022. URL https: http://axivvvoo/ org/ abs 2203. 02155.
Qwen. Qwq-32b: Embracing the power of reinforcement learning, March 2025. URL https: // //qwenlm. g thub. fo/bllllloo log/
D. M. Roijers, P. Vamplew, S. Whiteson, and R. Dazeley. A survey of multi-objective sequential Decision-making. Journal of Artificial Intelligence Research, 48:67-113, October 2013. ISSN 1076-9757. doi: 10.1613/jair.3987. URL htth::: /xx doi org/10. 111./....11111.1.1 /t//////////////////e/a.aea. jair.
John Schulman, Sergey Levine, Philipp Moritz, Michael I. Jordan, and Pieter Abbeel. Trust region policy optimization, 2017a. URL https: /ttp::s////xvvvv/. /wwpr/////////////ee///ee///saa/ abs 1502. 05477.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms, 2017b. URL
John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel. High- dimensional continuous control using generalized advantage estimation, 2018. URL https: wwtarxii.orggaaaa. wawgg/abs/155/ 02438.
Younggyo Seo, Lili Chen, Jinwoo Shin, Honglak Lee, Pieter Abbeel, and Kimin Lee. State entropy maximization with random encoders for efficient exploration, 2021. URL https: //arxi://axxvv. org/abs/ /2102 09430.
Charlie Snell, Jaehoon Lee, Kelvin Xu, and Aviral Kumar. Scaling Ilm test-time compute optimally can be more effective than scaling model parameters, 2024. URL https: /atpiaxxvvvv.oo...oaaa abs/ /208 03314.
Yuda Song, Julia Kempe, and Remi Munos. Outcome-based exploration for Ilm reasoning, 2025.
URL htt.s://///// /tatttps//aaxxv...// //tp///a.ovrr////ee/eeeeee..e.aa org. http//b//2555 06941.
Nisan Stiennon, Long Ouyang, Jeffrey Wu, Daniel Ziegler, Ryan Lowe, Chelsea Voss, Alec Radford, Dario Amodei, and Paul F Christiano. Learning to summarize with human feed- back. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 3008-3021. Curran Associates, Inc., 2020. URL /ttp:::://p/ccceee/ http://ppoooeedii...oeeereeeeeee 2ilee111888855222990884444433111
Yunhao Tang, Kunhao Zheng, Gabriel Synnaeve, and Rémi Munos. Optimizing language models for inference time objectives using reinforcement learning, 2025. URL https: :ps:// / //arxiv..or. abs/2503. 9595.
13

---

Published as a conference paper at ICLR 2026
Ronald J. Williams. Simple statistical gradient-followigg algorithms for connectionist reinforce-
ment learning. Mach. Learn., 8(3 (33-4):229-256) May 1992. ISSN 0885-6125. doi: 10.1007/ BF00992696. URL https:: /ttas://dol /ww.//oi..... wwtp/o.gg.ooooooooooooo... 1007 www.110000/0099255555 /F00992666.
Fang Wu, Weihao Xuan, Ximing Lu, Zaid Harchaoui, and Yejin Choi. The invisible leash: Why rlvr may not escape its origin, 2025. URL
herrgg/bb/25000 14843.
Omar G. Younis, Rodrigo Perez- Vicente, John U. Balis, Will Dudley, Alex Davey, and Jordan K Terry. Minari, September 2024. URL https: com/ FarramaFFuudatioon Foundation/ Minari.
Yang Yue, Zhiqi Chen, Rui Lu, Andrew Zhao, Zhaokai Wang, Yang Yue, Shiji Song, and Gao Huang. Does reinforcement learning really incentivize reasoning capacity in Ilms beyond the base model?, 2025. URL https: /estts://aaee /ttp///axvvvvvv. org/ btog/abs//20 2504. 13837.
Yiming Zhang, Harshita Diddee, Susan Holm, Hanchen Liu, Xinyue Liu, Vinay Samuel, Barry Wang, and Daphne Ippolito. Noveltybench: Evaluating language models for humanlike diversity, 2025. URL https: arxiv. www.iiii.g....oooooooooooo abs http//ass/22000000
Rosie Zhao, Alexandru Meterez, Sham Kakade, Cengiz Pehlevan, Samy Jelassi, and Eran Malach. Echo chamber: RI post-training amplifies behaviors learned in pretraining, 2025. URL https:
07912.
A LLM USE
LLMSs were used to help with polishing the paper and for coding assistance in low-level implemen- tation. We take full responsibility for all ideas, experiments, results, and conclusions presented in this paper.
B IMPLEMENTATION DETAILS
BabyAI and MiniGrid. For BabyAl tasks, the policy conditions on the grid image, the agent's 's direction embedding, and the mission text (encoded by GRU); the action space is the standard Al/MiiGriiir discrete set left, right, forward, pickup, drop, toggle, done. We provide example configurations for each task in fig. 4. We train CNN-GRB policy that outputs action logits. In MiniGrid- Fourrooms the action space is identical, but the observation excludes mission (as the goal specification is fixed), so we use a compact MLP that produces action logits conditioned on flattened image observation. During pretraining, strations for each task, except for Synthseq and BossLevel, where we found that the full dataset was required obtain reasonably strong base policy. We used the dataset from Younis et al. (2024). We pretrain by minimizing the cross-entropy loss with an entropy regularizer.
80/20 train-test split of expert demon-
During RLFT, fine-tune 50 configurations. For all tasks except BossLevel and Synthseq. these configurations are drawn from the pretraining test set; for BossLevel and Synthseq, we did not carve separate
If the agent reaches the goal at time (episode horizon H 1100) it receives the reward 1 - 0.5 : otherwise r = O. (Baby AI/MiniGrid defaults to 1 - 0.9 4: we found that lowering the time penalty improved diversity during RLFT, especially for REINFORCE, by reducing the disadvantage of longer but successful trajectories.)
Algorithmic Creativity (Triangle Discovery). In Triangle Discovery, an agent must output sequence of edge tokens that form valid triangle in an unobserved undirected graph. The input sequence comprises a graph index plus a prefix prompt and the agent's past outputs. We set the maximum sequence length to be 11. The action space is discrete vocabulary of size 1017. We pretrain a decoder-only Transformer with masked cross-entropy for 25 epochs. The pretraining dataset, from Nagarajan et al. (2025), contains 15,000 samples per graph, with the same being triangle with probability 1 and an edge with probability & edges. Each graph has 999 nodes. For RLFT, fine-tune on fixed graphs. The reward is sparse: +1 for a valid triangle, 0 otherwise.
14

---

Published as a conference paper at ICLR 2026
(a) GoTo: mission "go to purple box."
(b) Pickup: mission "pick up green ball."
Synthseq: mission "put the ball your right next to red ball pick purple ball after you open a grey door."
(d) Bosslevel: mission "put yellow key next blue key and pick up a ball after you pick yellow key."
Figure 4: Example Baby Al environments and their missions.
B.1 POLYCHROMIC PPO
We summarize implementation details for polychromic PPO beginning with the vine sampling scheme used for on-policy data collection, then additional stability techniques and full pseudocode in algorithm 2.
B.1.1 VINE SAMPLING
In this section, we describe the vine sampling scheme (Schulman al., 2017a) we used for poly- chromic PPO In vine sampling, we first generate number of trajectories. Then, we select a subset of the states visited, denoted by {s1...., ,5. this is called the rollout set. For each state 8; in the rollout set (we call this rollout state), we generate set of trajectories TI:N n.(. s;). To avoid notational overload and to make sample accounting explicit, we distinguish:
the set size used by the set-RI. objective (number of trajectories per set)
N: the number of rollouts collected from each rollout state/vine state
pt the number of rollout states selected along each seed rollout
.M: the trajectory budget i.e., maximum number of trajectories can collect.
now discuss the method we used to select the set of rollout states. Our goal is to identify multiple states from which we want generate vines so that the policy at a large number states. Suppose, we have a trajectory budget M, i.e., we are allowed to generate most M trajectories during the on-policy data collection. We first sample N rollouts independently where N > n. Now, for each these N trajectories, we identify p rollout states according some criterion. Some examples of rollout state criterion (1) Top p states with the
use the polychromic advantage to update
15

---

Published as a conference paper at ICLR 2026
highest entropy; sample trajectories from states where the policy is uncertain, (2) Top p states with the highest critic losses; sample more trajectories where our critic is wrong/biased. and (3) p equally spaced out states. Suppose the main trajectory is T timesteps long. We select the states at timestep T 27'
pT
p+11 p+1
p+1
In this paper, we used the third criterion. Note that, in this scheme. we generate, total, N W+N2(p> 1) trajectories. Since we allowed to sample at most M trajectories, that > n and N + (p- 1) I M. Across all environments, we set a trajectory budget M 136 for all methods. As such, and p 2 rollout states per trajectory. Note that at each of these rollout states, given we generate N rollouts, we can find (A) sets for our set RI algorithm. For our chosen hyperparameters, this was sufficiently large number of sets from each rollout state allowing the baseline.
must select N and p such
sets of size
4 for set RL, N 8 vines at each rollout state,
several sets compute
B.1.2 OTHER IMPLEMENTATION DETAILS AND PSEUDOCODE
Now, will discuss other implementation details for polychromic PPO We provide the pseu- docode for the complete algorithm in algorithm 2.
We found that adding the KL penalty from the behavior policy helpful for stability. In the absence of it, in some tasks, the model's performance collapses after a certain number of training epochs. This is likely because our method explicitly encourages exploration and the KL penalty provides an anchor that prevents the model from drifting too far.
In practical implementation, tages to the polychromic advantage to be useful. As shown in algorithm 2, we do not only set the advantage the rollout state 8, to be the polychromic advantage instead, we the advantages at states st, St+1," S1+W to be the polychromic advantage even though we do not generate vines from S1+1." ... St+W. This ensures that the exploratory behavior that the polychromic advantage encourages is induced at all states throughout the window. Otherwise, although the policy might be exploratory at 8t, it might revert to being purely exploitative all subsequent states due to the updates using the standard PPO objective; this would cause the policy to not explore. This issue be- comes pronounced environments where, despite being exploratory at 8t, the exploitative behavior in all subsequent states may override the exploration in sc. This is why this implementation trick was important in Baby AI and Minigrid while, in Algorithmic Creativity, we set W = 0 since once the policy visits diverse set of nodes from 8t, it is not easy merge paths.
found that adding window within which update all the advan-
Hyperparameter
Value
PPO epochs
2
Minibatch size
64
Discount (y)
1.0
GAE parameter (A)
0.95
Clipping parameter (e)
0.2
Actor learning rate
x 10-5
Critic learning rate
10-4
Value loss coefficient (c.)
0.5
KL coefficient (3KL.)
{0.005, 0.01, 0.05, 0.13
Max grad norm Temperature
0.5
1.0
Size of sets (n)
4
Num. vines at state (N)
8
Polychrome window (W)
5 for BabyAl/Miniirid, 0 for Algorithmic Creativity
Num. vine states (p)
2
Table 3: Polychromic PPO hyperparameters. All hyperparameters are fixed apart from the KL coefficient 3KL for which we provide the set we sweep over.
16

---

Published as a conference paper at ICLR 2026
Algorithm 2 Polychromic PPO
Require:
pretrained Policy иa, value function Va. discount y, GAE parameter X, clipping policy epochs K, value coef Co, KL target coef ßKI. (Poly-PPO specific:) number of sets N, size n,, number of vine states p, number of vines N,
window length W.
Initialize Te < TTg
2: for iteration = 1, do
3: Collect on-policy data using fractal sampling:
4: Initialize rol loutllotttttteeeeeeeeeeeeaeeaaaa +- (}
dictionary: state + list of trajectories
5: Roll out N trajectories using R3
6: for each trajectory do select vine states from for each vine state do
7:
8:
9:
Roll out trajectories from & using x Append the new trajectories to rollout_states out_states( end for end for
10:
11:
13: Compute advantages:
if roltout.states then 6 + T1 + 7V55511))))))) Vo(ss) A, GAB(dET,
15:
b generalized advantage estimation
17:
R₂ + A, V
18:
else if rollout_states then
19:
Create groups 91,91, ,9M of n trajectories from
20:
Compute set scores score( (g:) fpoly (8, Tim) for Tim gileq. (6))
21:
Compute baseline f(s) x
Σ i= score(g:).
22:
Define polychromic advantage of pairs (st, (.......................... (st+,a for each E gs!
Apoly (ser, ae) - score(g;) f(s)
polychromic advantage
23:
end if
24:
Normalize {AA}
25:
for epoch 1,...: K do
26:
for minibatch B do
27:
Compute ratios r0)
18))
[y(as(8))
28:
Policy loss:
(0)
Σ
min (r)))),, c
[
tEB
29:
Value loss: Lv(d) = [
EiEr (Vo(sr) - Ro)²
30:
KI. penalty: (0) [8]
Eseg KT (T( (T((180))
Total loss:
C(O. Φ) (0) + c. Lv(() + BKG. (0).
Take gradient step on 0, minimize C Update Ty t- T
33:
34:
end for
35:
end for
36:
end for
17

---

Published as a conference paper at ICLR 2026
Figure 5: Generalization under state perturbations in BabyAI BossLevel environment. The figure shows successful rollouts across perturbed initial states (blue circles), highlighting diverse strategies learned by the agent.
C GENERALIZATION EXPERIMENT
To evaluate the agent's ability to generalize to initial state perturbations, we designed the following experiment. First, for each environment seed, perform 100 rollouts with the pretrained policy to identify all reachable rooms. Then, for each unique room visited, we evaluate the target policy (fine-tuned using RL) by placing the agent 10 randomly selected starting positions within that room. Performance is measured using the pass @ metric, which calculates the success rate on the first attempt from these novel starting points. This randomization presents significant challenge, as successful trajectory from new initial state often requires substantially different strategies than those effective from the original start state as shown fig. 5.
18

---

Published as a conference paper at ICLR 2026
D PROOFS
D. PROOF OF LEMMA 3.2
Given the set reinforcement learning framework, we first consider the set up where, from each state s visited during on-policy rollouts, we can generate n actions, dim- As such, from each state, we generate tree where each node branches out into children. We denote all the nodes at depth of this tree by 86
,t, (n')
Furthermore, if from state s, the agent generates the of actions dim. the agent gets reward fpoly (8, (1m) = 1 Ei= r(s, ag)d((, (1n)) where r(s, qq) is the original reward function in the MDP and d(s,, a1:)) is some function measuring diversity.
For simplicity, we assume that the initial state is fixed. Note that we can construct the distribution of states visited by policy wy time in this tree follows:
P(00 {86 1.. (8 (nn)) To)
P(so + {8| (1:55)) J, }. (,),
= Σ ve((an) im 80)
Σ
P(s (1m) 18, (ao)im)
(ag)1.n
(tm)
Σ
no ao((ar-1)1m (at-1) (1="1)
(1)
(n (n'-1)
18--11 1.
18-1
(ar Dimi" (ay Di 2)
P(s (1mn') (Imm 1)
18-11 (a:-1) (1) ... (ar-1) ((n 2)
Here, each (a(mn in 1s) I1 mo(a; |)) Similarly, the following shorthand:
To =o((ar-1) (1) 1))1 (n' (n'-1)
II no (o((ar-1) (0)
(21)
Before we prove the lemma, we make the following observation: suppose the environment's state transition dynamics is deterministic, then the polychromic Q-function can be written using the poly- chromic value function as follows:
Q (s, a1m) = Spoly (8, a1sa) + Σ (2()))
where (1)
si (n) are the states reached from s after taking actions a1mm independently.
One can easily verify this by starting from the definition of the set Q-function and noting that:
Soly (x) (1) (at) so 8, (5))))(0)) a1
E
Spoly (31)
(a
Solly
(av) a1..
1=2
i=1
(1)
fpoly(s()) (1)
(a1)
E
polyy (o (at) = (ao) aim 81 (a1) (6)
1=1
nt-
fooly(s(",(a)))
+
E
Spoly (s)
(n(i-1)+s) (at)in (((((1))+)))))))
$
(a)) 1.
(a
1=2
19

---

Published as a conference paper at ICLR 2026
vol)
Given this and our constructions of the polychromic value functions, the proof of Lemma 3.2 follows the same procedure the standard performance difference lemma in Kakade & Langford (2002).
Lemma D.1 Given the state visitation generated by policy we and any normalized objective function f.
D
E
(
E
oondto (3) (f amm)
(8)
1 7m8
where
20
12
(8) (1 mm) Σ
? PPoo - t,,
Proof.
E
[(s (a)
1= 1=1
f( (1) (m))
1=0
86
Psso st To Bonn [r(']) (a))
0=0
D
(1
7m)
(1) P(so - 86 , (o,
[5('f) (a))1
7m
E Baddl (1)(()))))))
[fss, aim]
- yn
Lemma 3.2 (restated). Given any two policies my and n.
V
(s)
E amat (h (mo()1a) [At (s, a1:)
(9)
1
y
Proof. We first, show that
80 n'
v (s)
V (s)
= E
ΣD
Q1
(s mg
(1)
(at)
Vt. V2, (s( so = s
(10)
t=
i=1
We show this by the following simplification:
Vd (s)
V
(s)
n
E
Soly Spoly (0 ( (ac)(o))
s0
V (s)
1=0
i=1
n
Exp
Spoly (s (a))
80 = 8
V1 (s)
1=0
i
20

---

Published as a conference paper at ICLR 2026
(s (m(i-1))))) 1)
+ E
[5+1
80
x0
1=0
12
E m00
V2
(n(é
=0
f=1
7
Spaly (5t (a))
V ((44(-))))))) up
+
(s;+1
j=1
n2'
Emg
Va Ta
(o(
i=1
Now, the first term on the right hand side can be simplified (we suppress the condition that So in terms of notation):
E
(ppoty( (0 (a)) + +
vB ns
(5(+1
((f1-))+t)))
j=1
n
Spoty (5% (at) +
mg (+1+1 1)+1))))
(=0
==1
00
n'
E
Spoly (ae) + Σ
v ()((()))++)))
(sf+1
se,5, (as) (8)
(=0
j=1
QF (5 (a))
(as (0)
E
QF ng (5(
(1) (a,) (a))
Therefore,
Vd n/
(s)
Vi3 (s)
n'
Exg
2 (s (a))
V2 (s 80
i=1
Then, using eq. (8), the statement follows.
D.2 PROOF OF PROPOSITION 5.1
We follow the set-up in Cui et al. (2025). We parameterize the policy softmax distribution:
exp (Zsa)
no | 8) = Eer exp (Zsa')
where Zea is the output logit of action a from state s. We also have the following derivative:
i
log To 8) l{aa a} Te (s)
Ozsn'
First, we prove the following that shows how the output logit changes after one-step parameter update in the set RI. paradigm.
21

---

Published as a conference paper at ICLR 2026
Lemma D.2 Given the softmax policy no is updated using eq. (3) using step-size a, the change in the output logit zsa after one step update is given by
k+1
200
aB Oin
(110))
[(, (m
Σ = =a}
amnz (a 8) omn [f a1:n)]
Proof. This can be proven by elementary properties of expectation:
80 Zea
k+1
Ozt
E
an - (1|s) (f(s,a1m)) aim)
a. aBai. (11))
log (Pg((P(aim (s) amm)
t=
① (ag(m)((;))) ( (s)
a. min (11((11)
dzh
(1a = a) -xā(al b)) (1))
aB ain Ne (11))
Now prove the main result:
Proposition 5.1 (restated), Consider the set- RL setup at state s. After one update to the policy, the change entropy, A H (T² ( (s) - (Tā s). is given by
1fa; d3 )
[(s,dli))
A ~ -CCova
Σ log (a: s), Covar.n
Σ i===
where both covariances taken with respect to x-(- s).
Proof. We have the following first order approximation A (Cui et al., 2025):
~ Cov awai
Coo
(og (a 8), _+11 2k
Using Lemma D.2, we get that
A =anlEff(s, (oor (log TH (s), =ā s))
Σ 1(a =a] m1) ))
=
aCov =CCovoss((lo)
(x))((1)) log (a 1s) E Em
((,,
1=1
Note that
Covano [(ss (mm); Σ
1aa = a}
=E a a (
f(s, (1n)
Σ lfa
i=1
- natiaa |s)
8)
(( [f(s)
Using this, we get that
A zonEif(s, an.)) Covan (log áă (a 8), mtia s))
log mt (al 8), Banaa [( > 1a; a) 1
- -aCovawastie
1=1
a1.n Cor (log m (a s), m/ s))
22

---

Published as a conference paper at ICLR 2026
aCov a017) log 8), Covasin (sis Σ
n
1(as 1aa ((a==a) a)
n u (a 8) B) [f(s, amm))
(s(,, 11a al )
--aCov a~st(1))
log a5(( 8); Covasin
All that remains to show is
Covar.n
1 Σ
log Tt s), Cova
f(s,
Σ
lsa = a;}
1m1
i,
Cov (-1
log=zl0 8); Cova
[(s,a1sn)
Σ Mas a}
This can be seen using the linearity of covariance and the fact that each os in a1m. is sampled independently:
Covaion
1
log T (a; s), Covarn
f(s,a
Σ 1a. = =
Covar.
log nt (as I 8). Covaiin
A(
Σ 1la = oزم}
j=1
n f(s,
E1fa, = =a a1)
Cov (111)) log T/ 8), Covann
D.3 PROOF OF LEMMA 5.3
Lemma 5.3 (restated). Consider any set of actions a1m A". We can write the change in the log probability of sampling this set of actions after one policy update using set RL as:
log 0
+++ (arm 8) = log nT s) +
m aC(B
where C(Ok) is function independent of di:--
Proof. This follows from using first-order Taylor approximation of E log क 8) E n (a;
++t1 (a:
8) and then using Lemma D.2. Here, C(Ob)
Cov Jo ~xb(-11) (f(s, d. E T (a ss)
D.4 PROOF PROPOSITION 5.4
Proposition 5.4 (restated). Consider the polychromic objective in eq. (6). For any homogeneous set
41 = {a} where r(s, a) = 1, there exists E (0, 1) such that A Spasy (a) < O when To (a 8) > 6.
Furthermore, the scaffold values of these homogeneous sets satisfy the bound A Apply (a) S / /
Proof.
We first prove the first part of the proposition. Let = To s). We will use the shorthand A(a) = Aspas (a; no), f fpoly and f = Baiaaroo(11)) (s,aa1..)) Then,
A(a) Cov ~ (f(8, a(m); T(a)
1a = a;1)
1,f=1
= =Cov ai To (-p) d(m))
Σ 1[c = a})
1=1
23

---

Published as a conference paper at ICLR 2026
To(disn ()) a(.s) Dr Eai. 1
-
j=0
10 (
n
Toldisn s)(f(s, o1.m) - 1)(2 p)
j=U
1iin nfaHl=j
Imp]
p)
Σ
To(dim (f), -
la May
] p)
Σ
na(aio Is) cis)) -f)
n2
mp
[np
p)
Σ
Toldin [f) dimn))
i=0
lai.n
p)
Σ
no(aii [s) dim)))
[pp]
10 Mayl=s
s) = pp)
Te(a
[=a
Now, we can simplify using properties of the binomial distribution as follows:
Σ Σ
Te(di.n
1s)( _ -) p) =
-
[ n --) p)
-1-p)
((((1pp)')) -p)"~ = =0.
j=0 1a
Dajl=y
==
Therefore, the scaffold value becomes:
mpl
Maa= = Σ
j
p)
Σ
To(dion (f) cbwl))
j=0
+
Σ
j
--p)
p)
Σ
Toldii s)( ((
1= Lapl +1
1o n)
ap
j
p)
2j
Σ
Tola s)
n
n²
1=0
a n(al==
p)
n
+ 1
Σ
Te(a s)
pp
+1
[xp]
p)
2j
11
p² - p)--
n
n2
1=0
+
Σ
j
p)
n - j
pd (1 - -p)ー
[xp]
+1
j
2j n1)
p)
pd (1 p)) n-
n
n?
5
1=
Σ
n-1
j p)
n j + 1
()
pd (1 - mp
+1
24

---

Published as a conference paper at ICLR 2026
Here, in the second line, we used the fact when [aa } is of size j. then the smallest possible value of [8, d'..)) is 2 since at least 2 out of n elements are unique and at least j out of elements attain reward +1. On the other hand, we can bound it above by -=itl which happens if all elements get reward +1 and j + 1 elements are unique. In the last line, we used the fact that when fs, dín {a}) as the diversity O. Now, let c n-1, Then, when p 2 c, Inpl n 1. Therefore,
[pl _ j=0
j
2j 11 p)
A(a)
T n² 1 p² (1 (1-1p))-- 5O.
Now we prove the second part. First, define the following upper bound of the scaffold value of the homogeneous sets:
(
-p) Cu(X)) -K-DCcx(x) Ba(n)
A(a)= & Ex (i-Bia(a.
2x
< npl
where (x)
n-
E [[np]] + 1,n-1] 11
== n
Now, E X p] 0 and E[( - p) Var( x) nº Var(X) p(1-2) On the other hand, we can
bound (X) as follows: when x < [np]. C,(x) 2/2 < 2 < 2. On the other hand, when
j E [[mp] + 1, 1), we have (x) = n-z+1 - 1. Combining, we have that C < 1. Therefore, using the Cauchy-Schwand inequalit:
X -p -D) (--)Ca(x)
Bn(n) ExxBin(npp)
X
E
E [CC
p(1 p)
71
D.5 PROOF OF PROPOSITION 5.5
Proposition 5.5 (restated). Suppose dim heterogeneous where each a; unique with probability p (0, )) Suppose exactly q of the actions satisfy r(s, a,) and that any other action a' & a1m with (a s) > 0 yields r(s,a') = Then, the scaffold value of aim satisfies
Appal (aim)
gp"(1-2)
Proof. This can be proven using very similar techniques. Let P, the shorthand Spoly. Then,
Again,
n A(ai.)) Σ
Pp( D) pai. 1= (f())
j=0
Inp]
Po (p(00-p)-00 0+
-0+
j Py
2 n
-pp) Dal [f(s, a(.n)l
#=1
Σ
P3
p) B Basi.. tss (( + Pa -p) !
[=[np]+1
25

---

Published as a conference paper at ICLR 2026
n-1
Σ
P3( ]
Ba
n0 [[(s, |m + Pn - -p)
j= mp +1
2 Pn (1 p)
2
71
qp" p)
7
26


Reviews:
[{"review_id": "ai5IZhB9Na", "reviewer": "Reviewer_GrRC", "summary": "This paper investigates a key limitation of reinforcement learning fine-tuning (RLFT) for pretrained policies namely, the collapse of behavioral diversity during fine-tuning, which leads to reduced exploration and exploitable outputs. The authors propose a new polychromic objective for policy gradient methods that explicitly encourages both exploration and refinement of diverse generations. They adapt proximal policy optimization (PPO) to this setting by introducing vine sampling for on-policy data collection and a modified advantage function consistent with the new objective. Experiments on BabyAI, Minigrid, and Algorithmic Creativity benchmarks show that the method improves success rates, generalization under perturbations, and coverage in pass@n evaluations. Overall, the work provides a principled approach to preserving and leveraging diversity in RL fine-tuning, addressing a major failure mode of current RLFT pipelines.", "questions": "Is there anything that can be added besides the conclusion that can better inform the reader of the potential limitations above, or potential truncated sequences, as done by traditional monte carlo methods similar to n step bootstrapping?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "* I appreciate the time the authors spent to differentiate standard RL from set RL. The authors do a good job of establishing the setting before discussing the proposed polychromic algorithm. \n* The results suggest that the polychromic algorithm is that set scoring is indeed useful at maintaining diversity while keeping the main performance/validity equal to that of PPO. \n* This work carefully considers some of the flaws consistent in standard RL, entropy collapse and thoroughly discusses the possibilities of Polychromic PPO failure. \n* I do think the contribution is a interesting novel extension of PPO.", "weaknesses": "* Perhaps the largest issue that the authors did not discuss is the time complexity of the generation of vines for the use of the algorithm, which might be the main trade off readers would consider when determining if the added diversity is worth the complexity cost. It may be the case that the methods were not explicitly fairly compared with respect to this trade off. \n* Going a bit further on the initial point, would it be a fairer comparison to compare against other full Monte Carlo based algorithms that perform roll-out such as MCTS? Maybe I am mistaken, but is it reasonable not to include any runtime analysis against the proposed baselines?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "MpCM5eADbV", "reviewer": "Reviewer_5izw", "summary": "This paper addresses the problem in RLFT (Reinforcement Learning Fine-Tuning) where policies lose diversity and \"collapse\" onto a handful of behaviors. It proposes a new problem, set RL, and improves policy gradient algorithms like PPO. For instance, in the advantage function calculation, it redefines a shared advantage term based on a \"polychromic objective\" (which simultaneously evaluates reward and diversity); it uses a \"vine sampling\" strategy to collect the on-policy data needed to evaluate the performance of trajectory sets. Experiments demonstrate that this method can effectively improve the policy's success rate, generalization ability to perturbations, and pass@n coverage.", "questions": "1. I am concerned about the requirement for excessive sampling, which appears to be a significant overhead. Could the authors provide a computational complexity analysis for the algorithm?\n\n\n2. Standard practice for evaluating RL algorithms often involves 'deterministic evaluation', where the action with the highest probability (i.e., argmax) is selected. If the environment is static, collapsing to a single deterministic solution is actually acceptable and often desired. In this context, wouldn't forcing diversity inherently degrade the model's optimal performance? Did the paper include any experiments using this standard deterministic evaluation to compare with the stochastic results?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 4, "strengths": "The paper defines the \"Set Reinforcement Learning\" (set RL) framework , which generalizes the optimization of a single trajectory in traditional RL to the optimization of a set of trajectories . This framework is not only theoretically novel (for example, the authors derive a performance difference lemma applicable to set RL ), but it also provides a clear mathematical foundation that allows algorithms to directly optimize more complex objectives beyond standard rewards (such as the \"polychromic objective\" in this paper, which combines reward and diversity ). This holds significant inspiration and meaning for future work.", "weaknesses": "1. The definition of trajectory diversity is not general but must be engineered for each specific environment. For instance, it is defined as visiting different \"sets of rooms\" in BabyAI/Minigrid but different \"sets of nodes\" in Algorithmic Creativity.\n\n2. To avoid exponential sampling complexity in long-horizon tasks, the method instead relies on 'vine sampling' . However, this sampling strategy itself imposes a major constraint, as it requires the ability to reset the environment to arbitrary states. This is unrealistic in most real-world robotics tasks, thus severely limiting its practical applicability.\n\n3. The algorithm is \"computationally demanding\". This high workload stems from the multiple layers of sampling required, including the complex \"vine sampling\" procedure to gather sets of trajectories and the Monte Carlo sampling used to estimate the value baseline for the polychromic advantage. This makes the algorithm's workload substantially larger than that of standard PPO.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1kWY1tv0g0", "reviewer": "Reviewer_tLb8", "summary": "This paper studies how to induce diverse generations (trajectories) from a\npolicy trained with reinforcement learning (RL). The motivation is that current\npolicies trained with RL finetuning (RLFT) tend to collapse to a single mode of\nbehavior. One symptom of this is that increasing the number of trials afforded\nto an RL policy does not increase its performance (i.e., pass@n coverage is\nlow), because similar trajectories are generated.\n\nThus, this paper proposes to create RL policies that generate diverse,\nhigh-performing trajectories. This approach begins with defining set RL, an\nextension of RL that considers objectives over a set of trajectories created by\na policy rather than over a single trajectory. Within set RL, the paper then\nproposes a \"polychromic objective\" that motivates the trajectories generated by\na policy to be diverse and high-performing. Finally, the paper solves this\nobjective with an extension of PPO, referred to as \"polychromic PPO.\"\nExperiments are conducted on Minigrid, BabyAI, and Algorithmic Creativity,\nshowing that polychromic PPO can finetune policies to generate trajectories with\nhigher diversity, as evinced by higher pass@n scores.", "questions": "Below are minor comments and questions I found while reading the paper:\n\n1. Please use the correct citation format for in-text citations. They should be\n   author-year rather than numbers. Per the instructions in the ICLR template:\n\n   > Citations within the text should be based on the natbib package and include\n   > the authors' last names and year (with the et al. construct for more than\n   > two authors).\n\n1. In Eq. 2, is $s_0$ assumed to be fixed? $s_0$ is not really discussed when\n   Eq. 2 is introduced.\n1. Line 153-154: grammar mistake; maximize is used twice in the sentence?\n1. I found it a bit confusing that the paper mentions Minigrid and BabyAI as two\n   separate domains on which polychromic PPO is evaluated (line 64, 249). This\n   is confusing because the results only seem to be shown for BabyAI (Table 1)\n   and not Minigrid; furthermore, Section 4.2 only discusses BabyAI and not\n   Minigrid. Are Minigrid and BabyAI the same benchmark?\n1. Diversity in the experiments (Section 4.2) seems to be measured indirectly by looking at the pass@n ratings.\n   Is it possible to measure diversity directly by sampling multiple rollouts of\n   the policy and computing the diversity of those trajectories, e.g., using the\n   diversity functions described on lines 263-264?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. I found the description of set RL and the polychromic objective intuitive and\n   easy to follow. Since I am not as familiar with this field, I do not know how\n   novel this objective is. However, I personally have not seen anything like it\n   before. The closest I am aware of are works that seek to find diverse\n   policies, like Diversity is All You Need and quality diversity algorithms.\n   These works seek to find a set of diverse policies, while the current paper\n   trains a single policy that can generate diverse trajectories.\n1. I think the experimental evaluation is quite thorough, in that it seeks to\n   answer a number of important questions about polychromic PPO, and it directly\n   addresses the pass@n issue initially raised in the introduction.", "weaknesses": "1. The experiments are only run over 3 seeds, which is quite few. Furthermore,\n   there do not seem to be error bars, and statistical testing was not performed\n   to verify significant differences. Including these things would make the\n   experiments more precise, e.g., statistical significance could be discussed\n   instead of saying \"achieves slightly lower validity\" on lines 313-314. Below\n   I list a couple of other minor issues with the presentation of the\n   experiments:\n   1. Table 1: It is unclear what values are bolded and why.\n   1. Figure 2 could be cleaned up a bit. Namely:\n      - The text is small and hard to read.\n      - The y-axis bounds should bound the values in the graph, e.g., the graphs\n        in the leftmost column should have bounds from 30% to 100% rather than\n        30% to 90%.\n      - The rows and columns should be labeled, instead of having the names only\n        be in the caption.\n   1. Figure 1 appears after Figure 2 in the paper.\n   1. Line 307 is a comma splice.\n   1. Figure 3 also has really small text.\n1. The Entropy Analysis in Section 5 seems to provide valuable insights into how\n   polychromic objectives work, but I found it difficult to follow, in part\n   because I do not have background on what \"entropy collapse\" means. Perhaps\n   some background could be provided on why we care to analyze the entropy of\n   the policy?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DiRMNEHQhO", "reviewer": "Reviewer_Bmic", "summary": "The work introduces the notion of set RL, in which agents are maximizing rewards not with respect to individual trajectories, but with respect to a set of trajectories. This formulation is appealing as it naturally allows to directly optimize for reward maximization as well as diversity. Based of the notion of set RL, the authors propose polychromic objectives, which give a practical way of opimizing both for success as well as diversity of the set of trajectories. The work then further discusses how to adapt PPO to work with this particular polychromic objective before evaluating it on three different environments. The Poly-PPO implementation seems ot generally improve over pretrained policies and outperforming baselines. Finally, the work provides a more theoretical analysis on the effect of polychromic objectives on the entropy of a policy.", "questions": "* How is set RL distinct from multi-objective RL? Specifically, how is the notion of \"sets\" different from pareto-fronts?\n* Why did you focus only on PPO and no other RL algorithm?\n* Is the choice of diversity function a crucial one or is Poly-PPO fairly insensitive to this choice?\n* How did you set the hyperparameters of your method? Did you tune the baselines and your new algorithm?", "limitations": "", "rating": 2, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "The notion of set RL seems appealing and could inspire novel learning approaches that are distinct from existing classical RL algorithms. Further, the polychromic objective seems like a good logical consequence of set RL.", "weaknesses": "The work discusses various aspects fairly shallowly. In the beginning, the work reads like it is written just to introduce the idea of set RL. The work raises the expectation in the reader that a more comprehensive discussion of set RL and how it might differentiate itself from multi-objective RL. However, before such a discussion starts, the work pivots to discuss a particular aspect of set RL in the form of polychromic objectives. Before going into depth on polychromic objectives, a particular instantiation of PPO is discussed that makes use of this objective. While the work states in line 166 that the \"generality of set RL\" has been discussed, the overall work only discusses a very particular instantiation that seems to require a lot of adaptation from classical RL to work.\nOverall, a lot of design decisions seem to just fall out of the blue without being adequately discussed. I fail to see why, e.g., no other RL algorithm is being considered for the extension to set RL. Similarly, why is there no ablation on the choice of diversity function for the Poly-PPO implementation? The choice of REINFORCE as a baseline is never justified, nor is the choice of environments.\n\nTaken altogether, the work does not seem fit for publication in its current state and there are too many loose ends that need to be taken care of before I would consider increasing my score. I am happy to adjust my score if the authors show that I have misunderstood crucial aspects but vote for rejection as is.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)