Discrete Compositional Generation via General Soft Operators and Robust Reinforcement Learning
Authors: Marco Jiralerspong
Esther Derman
Danilo Vucetic
Nikolay Malkin
Bilun Sun
Tianyu Zhang
Pierre-Luc Bacon
Gauthier Gidel
Venue: ICLR 2026 Poster (2026)
Keywords: Reinforcement Learning
GFlowNets
Robust RL
Regularized RL
Generative Models
Scientific Discovery
Primary area: reinforcement learning
TL;DR: We propose a new general soft operator for discrete compositional generation tasks with experimental promise in various biological sequence design tasks and connect regularized RL/robust RL in these settings.

Abstract: A major bottleneck in scientific discovery involves narrowing an exponentially large set of objects, such as proteins or molecules, to a small set of promising candidates with desirable properties. While this process can rely on expert knowledge, recent methods leverage reinforcement learning (RL) guided by a proxy reward function to enable this filtering. By employing various forms of entropy regularization, these methods aim to learn samplers that generate diverse candidates that are highly rated by the proxy function. In this work, we make two main contributions. First, we show that these methods are liable to generate overly diverse, suboptimal candidates in large search spaces. To address this issue, we introduce a novel unified operator that combines several regularized RL operators into a general framework that better targets peakier sampling distributions. Secondly, we offer a novel, robust RL perspective of this filtering process. The regularization can be interpreted as robustness to a compositional form of uncertainty in the proxy function (i.e., the true evaluation of a candidate differs from the proxy's evaluation). Our analysis leads us to a novel, easy-to-use algorithm we name trajectory general mellowmax (TGM): we show it identifies higher quality, diverse candidates than baselines in both synthetic and real-world tasks.

Published as a conference paper at ICLR 2026
DISCRETE COMPOSITIONAL GENERATION VIA GEN- ERAL SOFT OPERATORS AND ROBUST REINFORCE- MENT LEARNING
Marco Jiralerspong" 1.2, Esther Derman Danilo Vucetic Nikolay Malkin Bilun Sun Tianyu Zhang*.2 Pierre-Luc Bacon Gauthier G.d
Université de Montréal
Mila University of Edinburgh
ACIFAR
Independent
ABSTRACT
A major bottleneck in scientific discovery involves narrowing an exponentially large set of objects, such as proteins molecules, a small set of promising candidates with desirable properties. While this process can rely on expert knowledge, recent methods leverage reinforcement learning (RL) guided by a proxy reward function to enable this filtering. By employing various forms of entropy regularization, these methods aim to learn samplers that generate diverse candidates that are highly rated by the proxy function. In this work, we make two main contributions. First, we show that these methods are liable to generate overly diverse, suboptimal candidates in large search spaces. To address this issue, introduce novel unified operator that combines several regularized RL operators into a general framework that better targets peakier sampling distributions. Second, offer a new, robust RL perspective of this filtering process, The regularization can be interpreted robustness function (i.e., the true evaluation of a candidate differs from the proxy's 's evaluation). Our analysis leads to a novel, easy-to-use algorithm we name trajectory general mellowmax (TGM): we show it identifies higher quality, diverse candidates than baselines both synthetic and real-world tasks. Code: https::///ie thu com/maa ht/.//aroooo////// rattttmm
compositional form of uncertainty in the proxy
INTRODUCTION
The task of scientific discovery centers on discovering novel objects with desirable properties such antimicrobial resistance or binding affinity. These objects often discrete and structured such that they can be constructed through a sequence of compositional steps, c.g., proteins as sequences of amino acids, molecules composed of various smaller fragments (Angermueller al. 2019 Wang jet al. 2023). Due to this compositionality, the number of possible objects grows exponentially with object size and becomes too large to search exhaustively. Additionally, experimentally evaluating potential objects can be too technical or too expensive. Nonetheless, in many cases, researchers have access to proxy reward model that quantitatively approximates the satisfaction level of the property of interest (Davis et al. 2018 Du et al. 2022).
large body of work leverages this proxy reward function to narrow the set of objects to small subset of promising candidates to test in a lab (Currin et al.) 2015 Gubernatis & Lookman 2018 Wu et al. 202T A14Science et al. 2023). An alternative approach that has demonstrated growing popularity and effectiveness is to pose this problem as an RL task (De Cao & Kipi 2018: Darvariu et al. [021 Bou al. 2024). what we will denominate discrete compositional processes (DCP), object construction viewed as Markov decision process (MDP) where terminal states correspond to complete objects and yield a reward given by the evaluation of that object by the proxy reward function. Then, the learned policy can be used as generative model that samples promising candidates.
Among these approaches, generative flow networks (GFNs) have shown considerable success (Bengio al. 2021 Jain et al. 2023 Zhou al. 2024 Cipcigan et al. 2024. GFNs learn an amortized policy and, at optimality, sample objects with probability proportional to the exponential of their reward
"Correspondence ma rco. aal quebecc

---

Published as a conference paper at ICLR 2026
(Eq. Connections with muximuum-entropy RL have been established, showing that this objective underpins the diversity of generated samples (Tappinn et al. 2024 Mohammadpour et al. 2024a Deleu et al., 2024). While this exact sampling goal is useful in certain applications, is problematic in exponentially large sets: the probability of sampling from a small set of objects with high rewards is dominated by the probability of sampling from an exponentially large set of objects with low rewards.
In this work, we leverage the connection between GFNs and regularized RL to use soft RL operators for scientific discovery. By combining these approaches. we derive an improved operator for the task and offer new perspective on regularization. Our contributions can be summarized as follows.
Our contributions: (1) We introduce general mellowmax, an interpolating operator that allows for a more flexible tradcoff' between quality and diversity than previous operators; (2) We derive a trajectory constraint for this operator, which we refer to as trajectory general mellowmax (TGM). TGM enables us to propose GFN-like algorithm for scientific discovery tasks that empirically finds diverse modes with higher reward than soft mellowmax, GFNs, and other RL baselines; (3) We provide a robust RL interpretation of these baselines in DCPs. Specifically, we show that the reward uncertainty sets induced by TGM are more interpretable than those associated with entropy regularization or GFNs. The analytical techniques we employ pave the way to more general trajectory constraints through an interpretable equivalence between regularization and reward robust sets.
2 RELATED WORK
The goal of sampling proportional to reward targeted by GFNs was historically tackled using Markov chain Monte Carlo (MCMC) (Gilks et al. 1995). These methods construct chains through a random walk: from set of initial points, new points are proposed and either accepted or rejected based on their function value (here, their proxy reward). While theoretically appealing, these methods struggle in high dimensions, in particular when comes to discovering separate modes (Bengio et al. 2021).
On the other hand, GFNs have demonstrated success in protein design (Jain et al. 2022). drug discovery (Shen et al. 2023), and material design (AI4Science et al. 2023). The standard way to tackle the excessive smoothness of GFN policies is through temperature conditioning (Zhang al. 2023
Kim et al. 2023 Zhou et al. 2024): the GFN target can be posed as learning p(x)) xx @ where 3 is an inverse temperature parameter. Temperature conditioning involves learning policy conditionally on B. with 3 varying throughout training. Another successful approach involves simultaneously learning a Q-function and sampling using a mixture of the GFN policy and the argmax over Q-values (Lau et al. 2024). Both methods complement our approach, as they can be used in addition to our framework. Mohammadpour et al. (2024b) have also proposed a dynamic regularization coefficient based on the number of actions, for maximum entropy RL. This change can also be used to address excessive smoothness in GFNs or incorporated through a varying wa in our Eq. 7
As for the RL perspective, many soft RL methods fit in the regularized framework of Geist al. (2019). Other forms of regularization and their corresponding operators include the mellowmax operator (Asadi & Littman 2017). which takes the logmeanexp Q-values instead of the logsumexp. and its extension, the soft mellowmax operator (Gan et al. 1021). Jiralerspong et al. (0024) also combine GFN perspective with RL by developing an adversarial version of the mellowmax operator called AFlowNets. None of these works use their operators for DCPs, whereas we generalize them, integrate them with trajectory constraints, and successfully apply these constraints to real DCP tasks.
Orthogonally, one can take generative modeling (with focus on property optimization) approach to the problem of scientific discovery. Of particular interest are methods that take a pretrained diffusion model and seek to finetune it to produce samples with a given property (in our case, the output of the proxy reward function) while maintaining high likelihood under the base model. [eehara et al. (2024a b) do so for continuous diffusion models while Wang et al. (2024) uses the gumbel softmax trick to differentiate with respect to a discrete diffusion model. Interestingly, these works also use an entropy term to ensure diversity of generated samples. While this approach is also applicable to similar problems, it additionally requires a pretrained diffusion model and can be computationally expensive (for example due to the need to approximately solve SDE at every step of training).
2

---

Published as a conference paper at ICLR 2026
Termination action
Complete objoct
B
A
(D)
xTx** (s)
(A)
(C)
12
YYY
YV..T
Figure 1: (Left) Illustrated issue with sampling proportional to reward. The high reward of the optimal sequence (green) drowned out by all the rewards in the combinatorial explosion of lower-reward longer subsequences. (Right) A DCP for protein sequence generation task. (A) Starting from an empty string, amino acids are added sequentially until termination. (B) Then, the full sequence is evaluated by a proxy reward function (xx) whose value is given reward for the termination action. (C) Over the course of protein generation. uncertainty accumulates through the per-step uncertainty terms 6,. (D) The true reward depends on (xx) and the accumulated uncertainty.
BACKGROUND AND MOTIVATION
MDPs. An MDP is a tuple (S, y, 7, P.r) where S and A are finite state and action spaces respectively, y E (0,1) 1) is a discount factor, P
- As a transition kernel and r S x + R a reward function. Define IT as the set of policy mappings π : S + A The goal is to find E IT that maximizes the value function of7 (s) := 1E 1_ 7'r(s,, ar) | 50 = s), Vs E S. The dependence of the value function the reward function is made explicit through the subscript r, clarified in the next paragraph. For any policy π 11, the expected reward is r (s) := Eae (a)rrr(s, a) and the expected transition PT (8, 8') = Enca T. (a) P(s, a, s'). The evaluation Bellman operator given by T := r" + 7P for all e E R is known to be contracting, admitting ef" as a fixed point. Similarly, the optimal value function U. (s) == max of is the fixed point of the optimal Bellman operator T (s) == max 7. (s),, Vu E R 8 E S. An optimal policy can be derived from iterative Bellman updates, which form the building block of RL (Puterman, 2014).
Robust MDPs. The true reward or transition function of an MDP is rarely known in practice. It may be estimated from trajectory data, but small error the MDP model can alter policy performance (Mannor et al. 2004). The robust MDP setting addresses this issue by assuming that (P,r) r) is unknown, lying
given uncertainty set. In our DCP setting where transitions are fully determined by the performed actions, we reasonably assume that the reward model r E R is the only uncertain element. Then, we aim to maximize performance for the worst-case model, namely:
U (s) == min (s), E PER
To solve this max- max-min problem, one can resort to robust Bellman operators 7 == mID T and T u= == max+en Ti v,, Vu E TRS. Indeed, both are contracting and admit the robust value vi and the optimal robust value vi as fixed points, respectively ([yengar 2005 Wiesemann et al. 2013).
Regularized MDPs provide general framework for regularization in RL, recovering celebrated algo- rithms such as soft Q-learning (Haarnoja et al. 2017). A regularized MDP is an MDP (8, A, P,r) combined with a family Ω := (S2,)sE of convex functions Ω2 A + R. At each state E S, Q defines a policy regularizer f2, (T.). (T.)).orr π. E A The regularized Bellman operator is given by:
( = (s)) (TT),
Vu R E
and its greedy equivalent by (s) max,, (y)) (7
Geist et al. 2019 Derman et al. (202T) have established an equivalence between policy-regularized and robust Bellman operators, thus highlighting formal motivation for regularized RL. The statement below a direct reformulation of
(Derman et al. 202T, Thm. 3.1).
Theorem 3.1 (Derman et al. (202T)). Assume that the reward function r is uncertain and satisfies funartnnnction rs E R2 = (s.-)) + Ry, where R2 I--R, Ri is closed and convex for all S. Then, for any E II, the robust value function Vi is the fixed point of the regularized Bellman operator T with D2 (TT. == max, maax,ER,, (-Rs, vs). other words, it holds that viz (s) = T ozz
(Tx), Vs E S.
Discrete compositional processes (DCP) special instances of standard MDPs with deterministic transitions. They can fully be described by a tuple (9, D), where 9 a directed acyclic graph with a single source state 8n and set of terminal states X' (see Fig. [ Each node in the graph corresponds
3

---

Published as a conference paper at ICLR 2026
Max validatioc
Unifoem
U7TR64
AMP1200
GFP22097)
12
10
15
Pereeat
10
31
10 12 01.0
0.5
1.0
1.5
2.5
3.0
35
Figure 2: Illustration of empirical distribution of proxy rewards for 3 domains compared to rewards of objects sampled by TGM policy at the end of training (blue). The uniform histogram (grey) was generated by sampling and evaluating million uniformly randomly drawn samples for each task, showing that the vast majority of possible objects are assigned a non-negligible reward by the proxy reward function. Max score refers to highest proxy reward given to an object in the validation set (sec App. F.4 for a description the validation set),
to partial object, being the empty set. An edge (s object s to get s' or a termination action if s' E X. Each completed object a X has an associated reward (x) given by the proxy reward model + X + R.
represents either adding part to the
In standard RL, the goal is to maximize the discounted cumulative reward. Instead, for scientific discovery,
interested in finding the most promising set of & candidates that are sufficiently distinct from one another (Jain et al. 2022). In practice, an imperfect but computationally efficient proxy score O is used to estimate the usefulness of each candidate, so naively maximizing usefulness can be problematic. Instead, given some metric d between completed objects, practitioners approxi- mately aim find a diverse and novel (w.r.t. the training set) set of candidates {aı (.I,...IR,, with high score (Jain et al. 2022). This goal can be formalized as maximizing the average mode reward (for simplicity, we only formalize the diversity constraint):
k
((xe)
subject to d(I;, 25) > dmins
Vi A
( CX K
3.1 LIMITATIONS OF GELOWNETS IN SCIENTIFIC DISCOVERY
Under a GFN perspective, the exponential of the proxy reward function can be viewed an unnormalized probability mass. Then, the goal is to learn per-state policy т. whose sequential application samples objects proportionally to this function. Formally, if we denote by T(x) trajectories in G ending at terminal state E X and アz(=) the probability of picking all edges from trajectory T (z)) under policy æ, the objective is to match the following distribution:
p(x)
Σ (()))T())
Tr(m) x c (0())
Vx E.
where 3 is the inverse temperature. We argue this objective suboptimal for solving Eq.
Motivating example. The main issue with Eq. 2 is that it yields a distribution assigning higher probability to large number of suboptimal reward objects rather than to a small number of high- reward objects. To demonstrate this, consider DCP (illustrated in Fig. I where the goal is to generate sequences of maximum length from a vocabulary of tokens of size B. Suppose a sampler is at an optimal sequence s* of length n < d. The sampler's policy then needs to decide whether to terminate and return x=" or to add more tokens and harm the usefulness of the sequence. doing so, weighs <<²) against the reward of all sequences starting with s. There are 52 2( such sequences. Suppose the sampler perfectly matches (2) and these sequences have rewards lower bounded by
0. Then, (ax () < (0(=") /r-BB- Ba-, and the probability of returning the optimal sequence at decreases exponentially in n. The problem occurs in many applications, as almost all objects
have non-negligible reward (see Fig. 2). When these empirical distributions are representative, GFNs
may conservatively estimate the probability of sampling a promising candidate.
ALTERNATIVE OPERATORS
To address this limitation, propose taking regularized RL perspective and modifying the operator used by GFNs. To simplify the discussion, we examine the case where G in the DCP is a tree. such settings, there is direct equivalence between the GFN flow operator and the soft Bellman operator
4

---

Published as a conference paper at ICLR 2026
GFN
Soft Bellman
MM
SMM
GM (ours)
Accumulation
log(k)
log(k)
(1=q) log(R))
~
Dilution
0
0
log(A:)
log(A))
9 glog({)
log(k)
(og(As
(lg
log(k)
max
(1 (1-9) log (k)
log(()
w
qartw
Table I: Worst-case accumulation, dilution and | for different operators when gotw 1. is the number of actions at a given state, w is the regularization coefficient, a is the softmax inverse temperature and q is the GM interpolation factor. For the same values of a, W, general mellowmax with q E (0, 1) trades off accumulation and dilution to get lower worst-case AA]] than the other operators. See App. for proofs.
with temperature = action a from s) and the vector Q-values at state by Qs, the soft Bellman operator is given by:
1. Denoting Qs(a) == us +r(s, a) (with being the state reached by taking
T SB U
logsumexp(w@Q.)
Recursively applying this operator in DCP, going from the leaves of G the root, causes value accumulate. At optimality, the value of a state is the logsumexp the rewards of all objects that can be reached from that state. The accumulation of the value of large number of suboptimal objects can overwhelm one high-reward object.
Instead of taking the sum, we could take the logmeunexp of the Q-values. Doing so yields the mel- lowmax operator (Asadi & Littman 2017). While the mellowmax operator solves the accumulation issue. is liable the opposite issue, dilution. Due to this averaging, the high reward of an object deep in the tree can get diluted by lower reward objects as we approach the root (since we divide by the number of children at each state).
T.M U.
log aECh(s)
Σ
Qa()
Ch(())
Luckily, we can alleviate the dilution issue by weighting the elements of the sum proportionally to the softmax (controlled by a) of the Q-values weights (instead
their value. In particular, of an equal weighting) and thus obtain the soft mellowmax operator (Gan et al. 2021):
SMM U2
(essoffmax(Q..... Q.).
Both accumulation and dilution consist deviations between the application of the operator and the max Q-value: A = Te, maXxgCh((s) ((s, a.. Formally, we define them as follows:
Acc = max (A 0) Dil (s) = -
(6)
Respectively, they represent the increase or decrease in the value of a state relative to its maximum Q-value. Ultimately, for scientific discovery applications, the ideal operator/ policy would assign probability to all actions while being resistant to accumulation and dilution.
We propose a new operator that interpolates between the accumulation of the soft Bellman op- erator and the dilution of the soft mellowmax operator with a hyperparameter q. Values of in (0, 1) have a smaller amount of dilution and accumulation instead of a large amount of either. Consequently, the resulting operator trades off both issues and has better worst-case bounds on | = max Aco (s), Dil" (s)). illustrated in Tab.
4.1 GENERAL MELLOWMAX
The operators mentioned above can be viewed through the lens of regularized MDPs, On one hand, the soft Bellman operator corresponds to regularizing by the Shannon entropy. On the other hand, given a distribution d, E AA,. the mellowmax/soft mellowmax operators result from regularization by (T, d,) where d, is the uniform distribution softmax((Q.)) respectively.
5

---

Published as a conference paper at ICLR 2026
We define the general mellowmax (GM) regularizer by interpolating between these regularizers. Then, we can recover these operators as special cases and also trade off their respective effects. We denote the general regularizer by S2g and the case dę softmax(aQu) by Da.
1 [qKL(T d,) (1 -1-1))))- ()(-(())) Vs E,T, E
2 (T&)
[0,1]. (7)
Notably, unlike other regularizers, 2t depends on the Q-values of the current state. Consequently, the associated operator is not a convex conjugate and deviates from the framework of (Geist et al. 2019). Nonetheless, we can still derive the associated operator which we denote Q&* using similar techniques (see App. E]:
2 (Qx) 1 (og)sttamax((())11 c.R.).
(8)
Using the regularized interpretation, trades off between maximizing the policy's 's entropy and having a policy that is close to the softmax of the Q-values. a allows us to put more or less weight on actions with high Q-values while controls the weight of regularization. In particular, setting GFN operators and policy. On the other hand, by 0 recover the mellowmax
allows
us to recover the standard entropy-regularin setting q 1, recover the soft mellowmax operator (with operator). The interpolated operator captures three important soft RI. operators and allows us to smoothly control the tension between accumulation/diiution using q and
We also observe that 22 is equivalent to a KL between a policy and tilted softmax, as stated below.
Proposition 4.1. For all q E [0, 1]. 22, (T.)
(T,, d.a)) log (gg(Zdd))) where Zg(d,)
EaEA da(a) and dea = d.//Z((d) the q-tilted softmax distribution.
4.2 TRAJECTORY GENERAL MELLOWMAX
Instead of training on transitions, possible to leverage trajectory-level constraints that relate policy/value functions over multiple steps to recover training algorithms that use subtrajectories as data units. In soft RI., this connection was first formalized in Nachum et al. (2017) through path consistency learning (PCL) for maximum entropy RL. Path consistency objectives have since been derived for Tsallis entropy (Chow et al. 2018) as well as the general class a-divergences (Brekelmans et al. 2022).
Trajectory-level constraints have been shown to be crucial to the performance of methods in DCP tasks. Trajectory balance (TB) (Malkin et al.) 2022) and subtrajectory balance (Madan et al. 2023)
are consistently used in GFNs, as DCP tasks only have terminal rewards. By connecting the policy at early states to this terminal reward, these constraints noticeably better at propagating signal and thus improving credit assignment.
Theorem 4.1 (Trajectory GM). For any DCP (G. ②) let (n" Q*) be the unique optimal value/policy functions for 12g Then, for given Q-function Q", the two following statements are equivalent:
softmax(/aq + w|Q?) arg max(Ts, 23) 1 109 (ms), for all states E EG.
Ne
if and only if
20 +
log softmaxx([(qq w] (OO,,,,a, q log softmax(a Ha.] r(x),
(10)
i=0
for all trajectories
ony x in the DCP.
We derive a novel, equivalent trajectory constraint for the general mellowmax operator. In particular, Thm 4.1 makes explicit that satisfying the trajectory constraint Eq. on all trajectories yields Qo such that softmax + w]Qo) is the optimal policy for the regularized problem. See App. E.3 for the proof.

---

Published as a conference paper at ICLR 2026
Trajectory general mellowmax (TGM). We now instantiate the interpolated regularizer practical algorithm for use in DCPs. Specifically, we make three design decisions.
(1) Using the general mellowmax operator with dg softmax(aQ2). Doing so allows to TGM both on-policy and off-policy and does not require a separate d,. Instead, the policy seeks to maximize reward while being close to softmax of the learned Q-values.
(2) Aiming to satisfy a novel trajectory constraint derived in App. E.3 Doing so recovers the benefits of trajectory constraints for DCPs.
(3) Training a single network Q" of Q-values through the VarGrad objective of Richter et al. (2020);
Zhang et al. (2023). It can be shown that minimizing the trajectory constraint Eq. [O equivalent to minimizing the following loss (where oe denotes the softmax with inverse temperature t):
1
Qo (7) = Var
Σ log Tqa+tal (a:) q log oa aalQ (a:))
Br(xx)
From Qº, the optimal policy is easily computed. The resulting algorithm is easy to implement, efficient (particularly for transformers), and equivalent to GFN training when q = 0,
5 DISCRETE COMPOSITIONAL GENERATION VIA ROBUST RL
In this section, we offer robust RL interpretation of regularized operators for DCPs. As described in Thm. [.I regularized MDPs are known to equivalent to robust MDPs with uncertain reward. the context of scientific discovery, since the reward is the evaluation of objects by proxy function, can conceptually consider the existence of a hidden true reward r* that we would like to maximize The filtering process can be interpreted as attempt to robustly maximize the proxy reward Φ accounting for the difference between and r". Thus, the goal is to solve:
max min Ea-p (xx + 8(x))
(11)
PEAx 8ER
where g| is the uncertainty set. The question then remains of how this uncertainty should be modeled in G. Traditionally, the reward in a DCP is identically 0 everywhere except for terminating actions. However, assuming that uncertainty only exists at the last action misses the compositional nature of the task. Instead, we decompose d(x)) into a sum of perturbations occurring at each step
of the generation process. More precisely, given a trajectory 50 d0
On
8n = H. we have
6(x) Σ 8,(a;), where d, E R...-
(12)
In this formulation, the uncertainty on is split between all actions taken to construct the object, instead of only existing at the final action, illustrated in Fig. [ From robust RL perspective, this uncertainty set structure corresponds to the common assumption of state-rectangularity (Wiesemann) al. 2013 Gadot et al. 2024).. The following section makes this model of uncertainty more explicit and analyzes the uncertainty sets of different operators. In particular, we show that the reward uncertainty set entailed by GFNs inadequate, whereas the set induced by the soft mellowmax operator provides more meaningful notion of uncertainty.
5.1 LENCHEL-ROBUST FORMULATION OF REGULARIZED MDPs
In this subsection, we use the robust MDP notations introduced in $ The following result provides an explicit mapping between reward-robust MDPs, as defined in 12) and regularized MDPs, through equivalent value functions.
Theorem 5.1 (Fenchel-Rooust MDP). For any state s S, let Ω2 AA + R be proper convex regularization function. Denote by D its convex conjugate which is known to be proper and convex. For any S, define the reward set
R2 ro(s..) fre E R.A. (.(-r.) 50)
Then, denoting the robust operator for R2 by TR, and the regularized operator for Pi by Tm:
T = R Vx E II E R
This reward would correspond the effective properties of interest measured in lab.
7

---

Published as a conference paper at ICLR 2026
Reg. Bellman
SMM
TGM
3
0.5
eDd
((x)
[,(1]
R,
R,
Figure 3: (Left) Regardless w. uncertainty set never contains 4 (x). As result, the soft Bellman/GFN operator is only robust to increases in reward. (Middle) For the soft mellowmax operator, for different values of with [.[1] > [2]. the uncertainty set contains & xx). Thus, the operator is robust to decreases in reward of one object (but not both the same time). When
(mellowmax), there is symmetry in this tradcoff, while increasing a skews such that the object with higher d, only admits small decrease in reward. (Right) uncertainty sets for GM interpolate between the two effects. While the uncertainty set not contain ((x), it contains points corresponding to decreases in reward.
q < does
Based on this theorem, solving Eq. Thm. 5.1 uses convex conjugacy in the context of dynamic programming. Similar results can be found in previous works (Eysenbach & Levine, 2021 Husain et al. [202]] Brekelmans et al. [022] Derman et al. (022). but as we carefully detail in App. A.5 they do not directly apply to our setting,
is equivalent to solving the associated regularized MDP.
5.2 ROBUST SETS INDUCED BY COMMON REGULARIZERS
Given this perspective, we now analyze the uncertainty sets corresponding to the regularizers we considered above. Since each is a special case of GM, we first give the uncertainty set of GM:
R, = ro + (r. E RAA E.EA do(a)ee (a) <1}-
(13)
A proof is in App. A.4 As opposed to Derman et al. (2021), the fact that our uncertainty sets are state-rectangular makes them independent of the executed policy, which results a formulation that fits with the standard robust MDP setting (Wiesemann et al. 2013).
Fig. 3 shows the uncertainty sets induced by three regularizers: negative Shannon entropy, soft mellowmax, and GM. For this illustration, consider single-step generation process with possible objects and associated proxy rewards [1, 1]. Similarly to (Brekelmans et al. 2022). ). we find that the entropy regularized Bellman operator (left) fails to capture a meaningful notion of robustness. Indeed, the uncertainty set never contains the proxy ((x), although higher values of w do bring the
set closer to the proxy reward. This problem is significantly exacerbated in DCP with multiple steps: at each layer, the uncertainty set becomes further away from the proxy reward (see App. D
EXPERIMENTS
Our experiments aim answer three main questions:
[Q1] How does the parameter q affect the peakiness of sampling in small and synthetic environments?
[Q2] Does TGM find better candidates than standard methods in large biological design DCP tasks?
[Q3] How robust is TGM to changes in hyperparameter settings?
6.1 IMPACT OF q IN SMALL AND SYNTHETIC ENVIRONMENTS [Q1]
TF-Bind-8. Originally proposed in (Trabucco et al. 2022). TF- -F-Bdnd-8 generates DNA sequences
of length 8 to find those with high binding activity to human transcription factors. The reward is the experimental binding activity of the sequence from (Barrera et al. 2016). As there are only 4 such sequences, the search space is small enough to compute the optimal value for each operator. We compare the optimal sampling densitics of each algorithm in Fig. for 4, a = 2 and w 2.
Bit sequence. Proposed in Malkin et al. 2022 the bit sequence task is significantly larger with 2220 possible sequences. The goal is to generate bit sequences length n by adding à bits at a time. M modes are selected remi-randomly and the reward is given by r(z) - minyeM d(x;, v)/n where is the Levenshtein edit distance (Levenshtein et al. 1966). Loosely, the reward of a sequence is the negative normalized edit distance the nearest mode.
8

---

Published as a conference paper at ICLR 2026
7
Unifume
OFN
TON:
2-0
q-0,25
q-0.5
q=0.76
GPN
Modos Pound
Min Dixt
22
3.2
1.0
1004
200k
100%
2004
Figure 4: (Left) Comparison the optimal sampling distribution of GFN and variants of TGM for TF-Bind-8 rewards. For the same 3 = 4, TGM concentrates significandyy more mass on the upper quantiles of the reward distribution. (Middle) Number of modes found by TGM and GFN. The increased peakiness of the TGM sampling does not harm its ability to find different modes. (Right) Average (over modes) of the distance of the closest sample found for each mode. On average, increasing q allows TGM to find closer samples to the true modes.
TGM (ours)
Valid.
SAC
PPO
GFN
q=
q-
0.25
- 0.5
- 0.75
q 1.0
Max
UTR
3.82+t.03
3.46+.02
1.12+.00
4.19+.011
1.15+.01
4.25+.00
4.27+.01
4.17+.01
4.26
AMP 5.69
9.79:1.02
10.044.04
10.13
9.93#.03
9.821.02
10.43=.07
10.29:1.05
9. 96
GFP 0.66+.02 1.90+.03 2.44++09
2.95+.08
2.664.12 2.34 2.05+.11
3.63
Table 2: Average mode reward of TGM compared to baselines on the biological sequence design tasks, We report the average and standard error 15 seeds. In all domains, variants of TGM either match outperform GFN/PPO/SAC with a particularly pronounced difference GFP, the largest domain. Valid. max refers to the object with the highest proxy reward in the validation set of the dataset used to train the proxy function.
As we have access to the ground- truth modes in this case, evaluate methods according to the following metrics: (1) the number of modes found, where a mode is deemed found if sample is generated within distance 8 = 28 of that mode; (2) for each mode, we track the distance to the closest sample generated during training. We report the average of this number over the M modes for the best performing hyperparameters (see App. F33 for further details). While the best GFN run matches 0, the increased peakiness of non-zero q-values yields noticeable benefits. The best-performing runs of q > 0 find samples that are much closer to the modes while also discovering more of them.
6.2 BIOLOGICAL SEQUENCE DESIGN [Q2]
We now focus on tasks based on woal-world domains. where proxy rewards are trained on actual datasets. The tasks are based on the setup of Malkin et al. (2022) and use the datasets from Trabucco al. (022). For each, transformer T is first trained on the training set (either for classification or regression). The transformer output is normalized before being used as a proxy reward (see App.
5' Untranslated Region Sequence (UTR). 5' UTR is mRNA region that regulates transcription of the main coding sequence. The goal of the task is to find a UTR sequence of length 50 that maximizes predicted gene expression level. We take the dataset of 280 000 filtered sequences from Trabucco et al. (2022) and associated ribosome loads. We train transformer as a regressor to predict the ribosome load from sequence. The vocabulary consists of nucleotides.
Antimicrobial Peptide (AMP). Antimicrobial peptides are short sequences of amino acids that have effects on microbes (bacteria, viruses, etc.). The goal is to discover novel peptides that are likely to have antimicrobial properties. The DBAASP database has collected known peptides with and without antimicrobial activity (Pirtskhalava al., 202I). We use the dataset (sourced from (Tabucco et al. )2022) of 9222 non-AMP sequences and 6438 AMP sequences and train a binary classifier (predicting whether a sequence has antimicrobial properties or not). We use the normalized logit of the classifier proxy reward. The vocabulary for this task consists of the 20 amino acids, plus a token corresponding to sequence termination. We mask the tokens such that the minimum sequence length is 14 and the maximum sequence length is 60. This is the only task where allow variable-length sequences (as it is the only one whose dataset contains sequences of variable length).
Green Fluorescent Protein (GFP). The green fluorescent protein is a length-233 protein whose fluorescence has numerous applications in biology. The goal of this task is to discover other proteins with high predicted fluorescence. We take the processed dataset from Trabucco al. (2022). which

---

Published as a conference paper at ICLR 2026
AMP 120A5)
CFP 22
Now moode NOW2IN
New mode
Ang mode
6.0
gred
gedd
PPo
PPP
Figure 5: Spread of final average mode rewards various algorithms from grid sweep over learning rates, and w. TGM on average performs better in AMP and GFP and similarly in UTR.
contains 56086 variations of the original sequence and their measured fluorescence. We train regressor to predict the measured fluorescence, and the vocabulary consists of the 20 amino acids.
Evaluation. Each method is trained for 100 000 generated samples during which regularly evaluate the learned network. By varying the temperature coefficient of the softmax representing the policy, we move along quality/diversity curve for each sampler. To approximately evaluate Eq. [T] we generate a set of samples for a range of temperature coefficients. We then aggregate the samples and aim to approximately determine the best k = 100 distinct samples (i.e., such that the minimum distance is 5) found by the policy. While this problem is NP-complete Karp (2009) (it can be seen as an instance of finding maximal independent set), greedy approach appears work well in practice. For each method, a sweep is performed over hyperparameters. The best performing setting (based on final average mode reward) run with 15 different seeds. See App.F..jp for further details.
Results. As illustrated in Tab. 2 for all three tasks, all variants of TGM either roughly match or exceed the performance of GFN, soft actor-critic (SAC) (Haarnoja et al.) 2018), and proximal policy oplimization (PPO) (Schulman et al. [017]. Similarly to what was found in (Malkin et al. 2022 Madan al., 2023), SAC and PPO perform relatively poorly. We hypothesize that the credit assignment problem associated with terminal rewards is a significant issue for these methods. In AMP, where generating shorter sequences is possible, PPO is able to achieve similar results to TGM/GEN The difference between TGM and GFN is most pronounced in the largest environment GFP, where TGM finds modes with significantly higher reward than the mean in Fig. 2 We hope this result shows the potential scalability of TGM. Interestingly, the effect of q more varied in these environments, showing the benefits of interpolation in the GM operator. The best performing variants balance dilution and accumulation by using q (0, 1).
6.3 HYPERPARAMETE ROBUSTNESS [Q3]
The best runs for Tab. were selected from grid sweep over learning rates {0.00001, 0.0001, 0.001}, B E {4, 16, 64, 256} and inverse temperature parameters the entropy coefficient to 1 l/w. For GFN, we set the sampling temperature of the policy to 1/w. To ensure an equal number of runs per method, fix a and only vary w for TGM. We expect that additionally varying would further improve results.. We plot the performance distribution of this sweep in Fig.
{1,4, 16). For SAC and PPO, we set
TGM variants seem relatively robust to these hyperparameter variations. The average SAC run is relatively stable but performs poorly. PPO is very sensitive to hyperparameter settings, with the best runs achieving strong performance in AMP. The mean average mode reward of TGM is higher than GFN across q and environments. Remarkably. TGM with q = 1 in AMP has little variability, with almost all runs performing well despite the significantly different hyperparameters. For GFP, there is significantly more variability, with the best performances coming from outliers for each method.
7 CONCLUSION
In this paper, motivated by the inadequacy of sampling proportional reward, we generalize various soft RL operators and propose GM as an interpolated operator. From this operator, we develop novel algorithm (TGM) and show it consistently able to outperform GFNs on a variety of DCP tasks. Finally, we adapt the reward-robust RL framework to DCP tasks where rewards are given for entire trajectories to offer a new perspective on regularization in these problems. Ultimately, we believe TGM has the potential to be a more effective framework for finding promising candidates in scientific discovery applications. Testing TGM on language generation tasks with verifiable rewards and extending the method to general DAG structured DCPs remain exciting directions for future work.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENTS
The authors would like to thank Glen Berseth, Muqeeth Mohammed and Sobhan Mohammadpour for fruitful discussions and feedback as well as Moksh Jain for providing helpful advice for setting up the biological sequence/bit sequence experiments. This research was enabled in part by compute resources provided by Mila. G. Gidel is CIFAR Al Chair, he supported by a Discovery Grant from the Natural Science and Engineering Research Council (NSERC) of Canada. M. Jiralespong is supported by an IVADO fund under the Canada First Research Excellence Fund grant to develop robust, reasoning, and responsible artificial intelligence. E. Derman is supported by Samsung. M. Jiralerspong acknowledges the support of the Natural Sciences and Engineering Research Council of Canada (NSERC) [funding reference no. 600916].
ETHICS STATEMENT
The purpose of this work is algorithmic nature, and we do not aim to produce actual applicable models. Nonetheless, the ideas in this work could potentially be used to improve the training of models used protein or other biological sequence design applications. While there are numerous beneficial applications for these types of models, they could also be used to design sequences that could eventually be harmful humans. This harm could come from bad actors or even accidentally by over-optimizing
single metrie (e-g, antimicrobial resistance) and not considering potential harmful effects on the body. The extension this work to multi-objective settings is a fruitful direction for future work.
In addition, the contribution of this work also has applications in language (another sequence design task). As such, it has the potential to improve the performance of large language models (LLMs) both useful and nefarious tasks. Given the recent improvements in the abilities of these models, this potential effect especially relevant.
REPRODUCIBILITY STATEMENT
To promote transparency and reproducibility, we share our code implementation as well as checkpoints of the proxy reward functions use. These allow for the reproduction of our results and fair future comparisons with new methods. We describe in ample detail the hyperparameters we use and design decisions made in App. as well Sec. 6 See the supplementary material for the code/checkpoints.
LLM USAGE STATEMENT
Claude Opus 4.6 was used for proofreading, spotting typos and improving wording/clarity for the camera-ready version of the paper.
REFERENCES
Mila AI4Science, Alex Hernandez-Garcia, Alexandre Duval, Alexandra Volokhova, Yoshua Bengio, Divya Sharma, Pierre Lue Carrier, Yasmine Benabed, Michal Koziarski, and Victor Schmidt. Crystal-GFN: sampling crystals with desirable properties and constraints. arXiv preprint 2310.04925, 2023.
Christof Angermueller, David Dohan, David Belanger, Ramya Deshpande, Kevin Murphy, and Lucy Colwell. Model-based reinforcement learning for biological sequence design. In International Conference Learning Representations (ICLR), 2019.
Kavosh Asadi and Michael Littman. An alternative softmax operator for reinforcement learning. In International Conference on Machine Learning (ICML), 2017.
Luis A Barrera, Anastasia Vedenko, Jesse V Kurland, Julia M Rogers, Stephen Gisselbrecht, Elizabeth Rossin, Jaie Woodard, Luca Mariani, Kian Hong Kock, Sachi Inukai, al. Survey
11

---

Published as a conference paper at ICLR 2026
of variation in human transcription factors reveals prevalent DNA binding changes. Science, 351 (1280)11450-4544 2016.
Emmanuel Bengio, Moksh Jain, Maksym Korablyov, Doina Precup, and Yoshua Bengio. Flow network based generative models for non-iterative diverse candidate generation. Neural Information Processing Systems (NeurIPS), 2021.
Dimitri Bertsekas. Convex optimization theory, volume 1, Athena Scientific, 2009.
Albert Bou, Morgan Thomas, Scbastian Dittert, Carles Navarro, Maciej Majewski. Ye Wang, Shivam Patel, Gary Tresadem, Mazen Ahmad, Vincent Moens, et al. ACEGEN: Reinforcement learning of generative chemical agents for drug discovery. Journal of Chemical Information and Modeling, 64 (15):5900-5911, 2024.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman- and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL com. jammm/jaxmm///a
Rob Brekelmans, Tim Genewein, Jordi Grau-Moya, Grégoire Delétang, Markus Kunesch, Shane Legg. and Pedro Ortega. Your policy regularizer is secretly adversary. arXiv preprint 2203. 12592, 2022.
Yinlam Chow, Ofir Nachum, and Mohammad Ghavamzadeh. Path consistency learning in Tsallis entropy regularized MDPs. In International Conference on Machine Learning (ICML), 2018.
Petros Christodoulou. Soft actor-critic for discrete action settings. arXiv 2019. arXiv preprint 1910.07207, 2019.
Flaviu Cipeigan, Jonathan Booth, Rodrigo Neumann Barros Ferreira, Carine Ribeiro Dos Santos, and Mathias Steiner. Discovery of novel reticular materials for carbon dioxide capture using GFlowNets. Digital Discovery, 3(3):449-455, 2024.
Andrew Currin, Neil Swainston, Philip J Day, and Douglas B Kell. Synthetic biology for the directed evolution of protein biocatalysts: navigating sequence space intelligently. Chemical Society Reviews, 44(5): (4455)1172-11) 11111122-2233 1239, 2015.
Victor-Alexandru Darvariu, Stephen Hailes, and Mirco Musolesi. Goal-directed graph construction using reinforcement learning. Proceedings of the Royal Society A, 27722225))222210668 2021.
Sarah E Davis, Selen Cremaschi, and Mario R Eden. Efficient surrogate model development: impact of sample size and underlying model dimensions. In Computer aided chemical engineering, volume 44, pp. 979- 984. Elsevier, 2018.
Nicola De Cao and Thomas Kipf. MolGAN: An implicit generative model for small molecular graphs. arXiv preprint 1805. / 1973, 2018
Tristan Deleu, Padideh Nouri, Nikolay Malkin, Doina Precup, and Yoshua Bengio. Discrete proba- bilistic inference as control in multi-path environments. Conference on Uncertainty in Artificial Intelligence (UAI), 2024.
Esther Derman, Matthieu Geist, and Shie Mannor. Twice regularized MDPs and the equivalence between robustness and regularization. Neural Information Processing Systems (NeurIPS). 2021.
Yuanqi Du, Tianfan Fu, Jimeng Sun, and Shengchao Liu. Molgensurvey: systematic survey in machine learning models for molecule design. arXiv preprint 2203. 14500, 2022.
Benjamin Eysenbach and Sergey Levine. Maximum entropy RL (provably) solves some robust RL problems. arXiv preprint 2103. 0103,06257, 2021.
Uri Gadot, Esther Derman, Navdeep Kumar, Maxence Mohamed Elfatihi, Kfir Levy, and Shie Mannor. Solving noctrectan reward-robust MDPs via frequency regularization. AAAI Conference Artificial Intelligence, 2024.
12

---

Published as a conference paper at ICLR 2026
Yaozhong Gan, Zhe Zhang, and Xiaoyang Tan. Stabilizing q learning via soft mellowmax operator. In AAAI Conference on Artificial Intelligence, 2021.
Matthieu Geist, Bruno Scherrer, and Olivier Pietquin. A theory of regularized markov decision processes. In International Conference on Machine Learning (ICML), 2019.
Walter R Gilks, Sylvia Richardson, and David Spiegelhalter. Markov chain Monte Carlo in practice. CRC press, 1995.
JE Gubernatis and TJPRM Lookman. Machine learning in materials design and discovery: Examples from the present and suggestions for the future. Physical Review Materials, 2(12): 120301, 2018.
Tuomas Haarooja, Haoran Tang, Pieter Abbeel, and Sergey Levine. Reinforcement learning with deep energy-based policies. In International Conference on Machine Learning (ICML), 2017.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In International Conference on Machine Learning (ICML), 2018.
Shengyi Huang, Rousslan Fermand Julien Dossa, Chang Ye, Jeff Braga, Dipam Chakraborty, Kinal Mehta, and João G.M. Araújo. CleanRL:: High-quality single-file implementations of deep reinforcement learning algorithms. Journal of Machine Learning Research, 23(274):1-18, 2022.
URL httt::ht::///000 r..
httml
Hisham Husain, Kamil Ciosek, and Ryota Tomioka. Regularized policies are reward robust. In International Conference on Artificial Intelligence and Statistics (AISTATS), 2021.
Garud N Iyengar. Robust dynamic programming. Mathematics of Operations Research, 30(2): 257-280, 2005.
Moksh Jain, Emmanuel Bengio, Alex Hernandez- Garciia, Jarrid Rector- Rector-Brooks, Bonaventure FP Dossou, Chanakya Ajit Ekbote, Jie Fu, Tianyu Zhang, Michael Kilgour, Dinghuai Zhang, et al. Biological sequence design with GFlowNets. In International Conference on Machine Learning (ICML), 2022.
Moksh Jain, Tristan Deleu, Jason Hartford, Cheng-Hao Liu, Alex Hernandez-Garcin and Yoshua Bengio. GFlowNets for Al-driven scientific discovery. Digital Discovery, 2(3):557-577, 2023.
Marco Jiralerspong, Bilun Sun, Danilo Vucetic, Tianyu Zhang, Yoshua Bengio, Gauthier Gidel, and Nikolay Malkin. Expected flow networks in stochastic environments and two-player zero-sum games. In International Conference on Learning Representations (ICLR), 2024.
Richard M Karp. Reducibility among combinatorial problems. In 50 Years of Integer Programming 1958-2008: from the Early Years to the State-of-the- Artt pp. 219 241. Springer, 2009.
Minsu Kim, Joohwan Ko, Taeyoung Yun, Dinghuai Zhang, Ling Pan, Woochang Kim, Jinkyoo Park, Emmanuel Bengio, and Yoshua Bengio. Learning scale logits for temperature-condiiiona GFlowNets. arXiv preprini 2310.02823, 2023.
Diederik Kingma, Adam: method for stochastic optimization. arXiv preprint 1412.6980, 2014.
Elaine Lau, Stephen Lu, Ling Pan, Doina Precup, and Emmanuel Bengio. Qgfn: Controllable greediness with action values. Advances in neural information processing systems, 37:81645- 81676, 2024.
Vladimir Levenshtein et al. Binary codes capable of correcting deletions, insertions, and reversals. In Soviet physics doklady, volume 10, pp. 707-710. Soviet Union, 1966.
Kanika Madan, Jarrid Rector- Rector-Brooks, Maksym Korablyov, Emmanuel Bengio, Moksh Jain, An- drei Cristian Nica, Tom Bosc, Yoshua Bengio, and Nikolay Malkin. Learning GFlowNets from partial episodes for improved convergence and stability. In International Conference on Machine Learning (ICMI), 2023.
13

---

Published as a conference paper at ICLR 2026
Nikolay Malkin, Moksh Jain, Emmanuel Bengio, Chen Sun, and Yoshua Bengio. Trajectory balance: Improved credit assignment in GFlowNets. Neural Information Processing Systems (NeurIPS), 2022.
Shie Mannor, Duncan Simester, Peng Sun, and John N Tsitsiklis. Bias and variance in value function estimation. In International Conference on Machine Learning (ICML), 2004.
Sobhan Mohammadpour, Emmanuel Bengio, Emma Frejinger, and Pierre-Luc Bacon. Maximum entropy GFlowNets with soft q-learning. In International Conference on Artificial Intelligence and Statistics, 2024a.
Sobhan Mohammadpour, Emma Frejinger, and Pierre-Luc Bacon. Decoupling regularization from the action space. International Conference on Learning Representations (ICLR), 2024b.
Ofir Nachum, Mohammad Norouzi, Kelvin Xu, and Dale Schuurmans. Bridging the gap between value and policy based reinforcement learning. Neural Information Processing Systems (NeurIPS), 2017.
Malak Pirtskhalava, Anthony Amstrong, Maia Grigolava, Mindia Chubinidze, Evgenia Alim- barashvili, Boris Vishnepolsky, Andrei Gabrielian, Alex Rosenthal, Darrell E Hurt, and Michael Tartakovsky. DBAASP v3: database of antimicrobialleytttoiii activity and structure peptides as resource for development of new therapeutics. Nucleic acids research, (9(D1):2288D222) 2021.
Martin L Puterman. Markov decision processes: Discrete stochastic dynamic programming. John Wiley & Sons, 2014.
Lorenz Richter, Ayman Boustati, Nikolas Nüsken, Francisco Ruiz, and Omer Deniz Akyildiz. Vargrad: A low-variance gradient estimator for variational inference. Neural Information Processing Systems (NeurIPS), 2020.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint 1707.06347, 2017.
Tony Shen, Mohit Pandey, and Martin Ester. Tacogfn: Target conditioned GFlowNet for drug design. In NeurlPS 2023 Generative Al and Biology (GenBio) Workshop. 2023.
Daniil Tiapkin, Nikita Morozov, Alexey Naumov, and Dmitry P Vetrov. Generative flow networks as entropy-regularized rl. In International Conference on Artificial Intelligence and Statistics (AISTATS), 2024.
Brandon Trabucco, Xinyang Geng, Aviral Kumar, and Sergey Levine. Design- bench: Benchmarks for data-driven offline model-based optimization. In International Conference Machine Learning (ICML), 2022.
Masatoshi Uehara, Yulai Zhao, Kevin Black, Ehsan Hajiramezanali, Gabriele Scalia, Nathaniel Lee Diamant, Alex M Tseng, Tommaso Biancalani, and Sergey Levine. Fine-tuning of continuous-time diffusion models as entropy-regularizze control. arXiv preprint arXiv:2402. 15194, 2024a.
Masatoshi Uchara, Yulai Zhao, Kevin Black, Ehsan Hajiramezanali, Gabriele Scalia, Nathaniel Lee Diamant, Alex M Tseng, Sergey Levine, and Tommaso Biancalani. Feedback efficient online fine-tuning of diffusion models. arXiv preprint arXiv:2402. 1633,9,9 2024b.
Nino Vieillard, Olivier Pietquin, and Matthieu Geist. Munchausen reinforcement learning. Advances in Neural Information Processing Systems, 01:4255-4466 2020.
Chenyu Wang, Masatoshi Uehara, Yichun He, Amy Wang, Tommaso Biancalani, Avantika Lal, Tommi Jaakkola, Sergey Levine, Hanchen Wang, and Aviv Regev. Fine-tuning discrete diffusion models via reward optimization with applications to dna and protein design. arXiv preprint orXiv:2410.13643 2024.
Hanchen Wang. Tianfan Fu, Yuanqi Du, Wenhao Gao, Kexin Huang, Ziming Liu, Payal Chandak, Shengchao Liu, Peter Van Katwyk, Andreea Deac, et al. Scientific discovery in the age of artificial intelligence. Nature, 620(7972):47-60.. 2023.
14

---

Published as a conference paper at ICLR 2026
Michael L. Waskom. seaborn: statistical data visualization. Journal of Open Source Software, 6(60): 3021, 2021. doi: 10.21105/joss.3002.. URL htt_tpppp//abb.ocgy.002111////// 03021
Wolfram Wiesemann, Daniel Kuhn, and Berç Rustem. Robust Markov decision processes. Mathe- matics of Operations Research, 38(1):153-183) 2013.
Zachary Wu, Kadina E. Johnston, Frances H Arnold, and Kevin Yang. Protein sequence design with deep generative models. Current opinion in chemical biology, 65:18-27, 2021.
David W Zhang, Corrado Rainone, Markus Peschl, and Roberto Bondesan. Robust scheduling with GFlowNets. International Conference on Learning Representations (ICLR), 2023.
Mingyang Zhou, Zichao Yan, Elliot Layne, Nikolay Malkin, Dinghuai Zhang, Moksh Jain, Mathieu Blanchette, and Yoshua Bengio. PhyloGFN: Phylogenetic inference with generative flow networks. International Conference on Learning Representations (ICLR), 2024.
15

---

Published as a conference paper at ICLR 2026
APPENDIX
A KOBUSTNESS-REGGLAAIIZZO DUALITY
A.1 CONVEX ANALYSIS
Before proving our results, we briefly recall notions of convex analysis below Bertsekas 2009).
Definition A.1 (Convex conjugate). Let a function f RRn R with domain dom (f S R." The convex conjugate of f is defined as
(
sup zEdom(J)
(x.1) f(z).
Definition A.2 (Infimal convolution). Given two functions f.g g 1R + R. the infimal convolution of and g defined as:
(z) = inf (f(x=2) - +g()}}
=EB"
Property A.1 (Operations on conjugate transforms). Let two functions fI, f2 R" R and positive real number w > 0.
(i) Defining g(x) (fr( the convex conjugate of g satisfies g²(w) (
(ii) For k + 0 and h(x)) = fi (Rar), the convex conjugate of satisfies h" (v) = (T(
(iil) The convex conjugate of the sum fi + f2 the infimal convolution of their conjugates f2 namely:
[fi f2l" (
Vy E"..
A.2 APPLICATION TO POLICY DIVERGENCE
For Shannon entropy and KL divergence, the convex conjugate can be derived in closed form and is known to be a logsumexp function, so we omit the proof Geist et 2001 Derman et al. (202T).
Proposition A.1 (Shannon conjugate). Define the negative Shannon entropy -Il IR" + R as
[-1](z) = >i=1 I; log(r:) over the simplex domain An. is convex function, and its convex conjugate is the logsumexp:
12
(1) = log
Σ
e
Vy E".
Proposition A.2 (KL conjugate). Let d E IR" be such that d > O and define the KL-divergence
function KLa(x)= = Σm= 2; log
x1
Vr E An It is a convex function, and its convex conjugate
di
the weighted logsumexp:
WLSE, == log
Σ de i=1
Vy E RD.
We are now interested in deriving the Fenchel conjugate of a convex combination of Shannon entropy and KL divergence. Its explicit form described below.
Proposition A.3 (Shannon-Ke conjugate). For any q E [o, 1). the convex combination (1- qKL satisfies:
(1-9)(-H)) + {xL) =
Σ ar log
Fi Vx E R
(d.)a
and admits as convex conjugate the function:
(1 - (--)) + "KLLa = -log
Σ (eve
Vy E R...
16

---

Published as a conference paper at ICLR 2026
Proof. The first statement comes from elementary algebra:
[11 - (1-
()(-))()))))))))) H) + (xL = (1 1--)) q) Σ Ia log(z;) + +9
Σ log
3j
di
==1
1=1
n
= (1 q + q Σ
24 log(z;)
q 2
Is log(d:)
1m1
im1
n
Ty log(
Σ
24 log(()))))
T3 log
(L))
(d.)9
For the second statement, we simply apply Prop. A.2 to obtain that
n
(-- + qKLa)] (y) WLSEn, (y)
log
Σ
(d)))ee))
more involved proof would combine notions of infimal convolution with the respective conjugates of the two divergences. We omit it for brevity.
A.3 PROOF OF THM. 5.1
We provide a slightly general proof below, with an arbitrary eg- set instead of just O. Theorem A.1 (Fenchel-Robee Regularized MDP). For any E 8, define the reward set
R2 = ro(s, :) fr. E RA fā(-rs) S {t,): and the regularization function (T, == fa(Tx) es: Vs E S. Then, for any policy a II, T = T892 Yu E TRS. and Vi = v.
Proof. First, establish the support function of a set R. = tr's E :5 (.)) at any state E S. By definition, for any policy E AA, we have
max EER,
max T. (M.ERAff))))))))..
(T,)
(T,) + Cg.
Next, we compute the robust Bellman operator associated with the uncertainty set R R.R rT((s u()
Tk)(s
min r.ER,
min r (s) + 7 r.Ero(e.)+
(s) + min (T,, r') + 7 pp v(s)
ER
= (s) max (Ta, - + 7pp c(s)
T v()) max (T.: -
E,R,
remains to compute
max (TT, -r/) = max r.) S. t. ft --r)) S
-max(T.,F.) F.) s. f (F.) S E
P,
Employing the above expression of the support function enables us to write:
( Fa) s. t. (F.) S cg = (T.)) +
so that T = T uu Ve E R.S. The unique fixed point each operator is vgz and 2
respectively, which leads the conclusion.
17

---

Published as a conference paper at ICLR 2026
Equation 13 comes from combining Thm. 5.1 with Prop. A.3 We can similarly obtain the uncertainty set resulting from KL and negative Shannon regularizations, based on their corresponding dual described in Sec. A.I This leads us the following table, summarizing the different regularizers used this paper along with their corresponding uncertainty sets.
A.4 PROOF OF PROP.
Proposition A.4. For all q E [0, (z, wa KL - where Zg(d,) =
EoEA dz(a) and de = (d)) is the q-tilted softmax distribution.
Proof. By definition,
wg = [qKL (T,, do) + (11-))(-(TT.))
(a((a)
(a) log
+ (1 -1-)) q) Σ T log(((()a))
d,(a)
a'd
aEA
- Σ
(a) (n) + Σ
Tx(a) log(T.(a)))
aEA
aEA
(a(a)
Σ Ta(a))og
d,(a)9
(a)
(a) log
Σ
Ta(a) Log(Za(d.))
d.(a)9/Zg(d.)
aEA
aEA
(a(a)
Σ
(a(a) log
log (g(d)),
d.(a)9/2g(d.))
acA
which yields the desired result.
Table 3: Summary table of policy regularizers.
Neg. Shannon
KL divergence
GSM
General convex
Regularizer
?.
H(TT)
(I,(F,,
1 (ssL( wP
f((Te) +
+(1 (--((((((TT)))
Conjugate
LSE)
WI.SF.A. (q:)
w WI.S. A.
f((g)) -
Reward Uncer-
(Tx (2 50) (rx ((rr)) S (r.:0((-ra) 50 ((rr.) e.)
tainty
A.5 COMPARISON WITH OTHER REGULARIZED RL WORKS
Eysenbach & Levine ) focuses Shannon entropy, while we encompass a broad class of regularizers. Husain et al. (2021) b: Brekelmans et al. (022) analyze robustness from the LP-dual perspective of RL. Although equivalent& optimum, this approach may not hold for any given policy. The notion of occupancy measure is also obscure the context of DCP where transitions are fully determined by the agent's decisions and the time horizon is finite. Differently, Thm. 5.1 establisiss
18

---

Published as a conference paper at ICLR 2026
robustness reeuuaaiiaignl equivalence for any policy via Bellman evaluation operators. It thus provides a principled identity between robust dynamic programming in the sense (Iyengar] (2005) and regularized MDPs in the sense of (Geist et al. 2019). Finally, the research motivation of Derman et al. (202T) being different, they proceed the opposite way from ours: they deduce a regularizer from generic uncertainty sets, whereas we deduce uncertainty sets from generic regularizers. This enables us to clarify the robustness properties caused by regularization.
19

---

Published as a conference paper at ICLR 2026
B ADDITIONAL RESULTS
B.I IMPACT OF 3 ON GFN PERFORMANCE
Another potential solution to the issue of sampling proportional (**)) consists of modifying the reward function. In particular, the reward exponent hyperparameter 3 could be used to arbitrarily increase the value of bigh-reward objects, hopefully overwhelming even an exponential amount of perform additional grid sweep for GFNs using
low-reward objects. To test this hypothesis, higher values of B. We sweep over the same learning rates {0.00001, 0.0001, 0.001 } and sampling
temperatures {1, 4, 16) but also over 3 E {512, 1024, 2048, 4096, 8192). We plot the final average mode reward of the best performing setting for each | in Fig. 6
Going past 256 worsens performance for UTR and AMP but does noticeably improve performance for GFP. However, for all values of 3. GFNs do not manage to reach the best-performing TGM setting. This discrepancy indicates TGM is exploring meaningfully different peaky distributions with better quality and diversity. Surprisingly, GFNs still perform decently even with very high 3 values (up to 8192). seems that gradient clipping and the VarGrad objective are enough to ensure some level of training stability at these high / values. These values yield equally high losses, reaching values the millions for 3 8192 on GFP.
UTR
Arg. modee reward
4.2
Avr. mode reward
2.5
4.0
8
2.0
3.8
3.6
3.4
Best TGM: 4.27
Bost TGM: 10.43
0.5
Bost TGM: 2.95
256
S12
1024
2048
4n86
8192
256
512
1024
2048
4086
0192
256
512
1024
2048
4096
0152
Figure 6: Average mode reward of best GFN setting for different values of /3 compared to the best performing setting TGM.
B.2 COMPUTATIONAL COMPLEXITY
TGM is generalization of TB but nearly identical in terms of computational complexity (essentially just requires an additional log-softmas, in the loss computation). In particular, for sequence tasks with transformers, by acting on trajectories, TGM/TB are significantly faster than methods that act on transitions. The loss computation only requires a single forward pass per batch of trajectories instead of forward pass per batch of transitions. We've included a table of training speed on an L40s GPU below. TGM/TB are roughly 10x faster than PPO/SAC when the latter splits generated trajectories into minibatches. TGM is only slightly more complicated to implement than GFNs, and simpler to implement than PPO (see the code submission for details).
Table 4: Wallclock training speed on an L40s for various methods on biological sequence tasks.
Algorithm
UTR (samples/s) 1
AMP (samples/s)
GFP (samples/s) t
SAC
115
92
5
PPO
191
144
8
TGM/TB
1680
1216
120
B.3 DIVERSITY AND NOVELTY EVALUATION
In addition to the aggregate metric of average mode reward, which seeks to approximate Eq. [ also evaluate the diversity and novelty of generated samples by each method. To do so, we use the metrics proposed in Jain al. (2022). For diversity, the metric evaluates the mean pairwise
20

---

Published as a conference paper at ICLR 2026
distance between generated samples. For novelty, the metric evaluates the mean minimum distance to a training set sample. We evaluate the modes generated by each method at the end of training and use the Levenshtein edit distance (Levenshtein et al., 1966).
Table 5: Diversity of TGM compared to baselines on the 3 biological sequence design tasks. Included standard error over 15 seeds.
GFN
TGM (ours)
0.25
9 0.5
9 0.75
1.0
UTR 18.38:00.06
18.75
18.29:10.12 0212 18.5840.09
18.62:+0.07
18. 18120-10.11 04.06=0.75
AMP 17.76=1.18
16.54
14.82=00.00
03.69+0.05
12.94=0.03
GFP 149.90+2.42
84.18+2.63
89.95+7.59
106.6847.69
96.36+6.58 107.83+55.6
Table 6: Novelty of TGM compared to baselines on the biological sequence design tasks. Included standard error over 15 seeds.
GFN
TGM (ours)
q
q 0.25
q 0.5
0.75
1.0
UTR 31.99:10.12
30.97+0.06
30.85=00006
32 32.37=0.10 10 32.09:0.07
30.94+0.07
AMP 36.21+4.14
35.87=41.05
27.65-+3.33
21.11#0.02
21.65
20.91+0.03
GFP 205.6510.26 212 412.4810.78 210.38:10998 209.86::0.60 211.03++0.6 210.93
In terms diversity, GFNs and TGM obtain similar values for UTR. For AMP, diversity is higher for GFNs but this is partially due to all the modes GFN generates being of length 60. Since the Levenshtein edit distance grow with sequence length. Conversely, the sequence lengths of modes generated by TGM are more variable and shorter, limiting the potential edit distance. For GFP, GFNs obtain noticeably higher diversity, though at the cost of significantly lower reward.
normalized for distance,
expected that the pairwise distance will
For novelty, all methods are relatively similar for UTR and GFP. There a significant difference for AMP, however this is likely due, once again, to the longer sequences being generated by GFNs. Overall, the increase in average mode reward does come at the cost of reduced diversity, though we argue this tradeoff is fair given the goals of scientific discovery (there is still enough diversity to have different interesting candidates).
B.4 ADDITIONAL BASELINES
We compare the performance of TGM with additional baselines: Munchausen DQN (Tiapkin ct al. 2024; Vieillard et al. 2020), QGFN (Lau et al. 2024) and lemperaturee-condiioneee GFNs (Zhang et al. 2023] Kim al. 2023). We note that QGFN and lemperaturre-condiiooeee GFNs complementary to our approach: the techniques they use can also be used with TGM.
Table 7: Mean mode reward of M-DQN, GFN and TGM as well as Qfeemperature-condttione variants. Included the standard error over 5 seeds.
M-DQN
GPN
QGEN
GFN + Temp Best TGM
QTGM
Best TGM Temp
UTR 3.37+0.12
4.12#0.00
4.08=0.04 4.1810.05
4.27+0.01
4.05
4.1810.01
8.96=0.11 9.97+0.12 AMP 2.97++0.82 10.04+0.01
10.43-0.07
9.115
9.83+-031
GFP 1.90+0.03
1.08=0.07 1.8910.04
2.9510.08
0.9910.11 3.10=0.15
B.5 MUNCHAUSEN DQN
We set a = 0.1,4 = 20 and separate the entropy coefficient from a as we observe that entropy coefficients tend to perform better (which would require < 0). Like other methods, run
a hyperparameter sweep: over learning rates E {le - 5, le 4, le - 3), 3 E (4, 16, 64, 256) and entropy coefficients (1, 1/4, 1/16) and select the best performing hyperparameters. We find that
21

---

Published as a conference paper at ICLR 2026
Munchausen DQN (M-DQN) struggles to learn a good sampler. Similarly to SAC, the non-trajectory based loss appears to lead to unstable training and poor performance. A trajectory based version would likely address some of these issues.
B.6 QGEN
We test the p-greedy QGFN setup (consisting of a mixture of forward policy with the argmax of the Q-values). We train a separate Q network in parallel (using the same architecture as the GFN) and set (in the n-step return) to be the maximum sequence length of the environment divided by We sweep over learning rates E {le 5, 4, le 3). 3 E {4, 16, 64, 256) and values of p E {0.25, 0.5, 0.75) while fixing the sampling temperature 1.
We also take the best performing hyperparameter setting from TGM and test with the p-greedy setup (using p 0.25). Doing so worsens performance in all 3 domains. While we have not tested them, it is possible that other variants QGFN or the annealing schedule mentioned Lau et al. (2024) could improve performance. A more careful combination is beyond the scope of this work, but offers potentially interesting future avenues for research.
B.7 CRAPERAATURE-CONDITIONN GFNS
We follow the method described in Kim et al. (2023) which expands on Zhang al. (2023). Specif- ically, during training, for each batch we uniformly randomly sample a temperature between the set minimum temperature and maximum temperature. This temperature is encoded by two-layer MLP before being added to each transformer layer. Then, as in Kim et al. (2023), the temperature decoded into scalar which is used as a temperature parameter, modulating the output logits. minimum temperatures {1/ L//256, 1/64, 1/16, 1/4] and maximum temperatures {1/128, 1/32, 1/8} as well as learning rates E {le 5, le - 4, le - 3}. For evaluation, we generate samples using the minimum temperature. Overall, this leads to similar or slightly better performance (on UTR) compared to regular GFN.
We then sweep
We also take the best performing hyperparameter setting from TGM and test it with temperature conditioning between I and 3- The change results significant performance increase for GFP showing that temperature conditioning is a general technique that can be applied to both GFNs and TGM.
3
22

---

Published as a conference paper at ICLR 2026
C ACCUMULATION AND DILUTION
We prove the bounds in Tab. I below. For convenience, we denote Q = maX Q(s, a),
k = [Ch(s)] and 3 == go go
Soft Bellman: For the soft Bellman operator, A A7 2 since the logsumexp is always greater
than equal to the max. Then, since the logsumexp is strictly increasing, worse-case accumulation occurs when all the Q-values are equal, in which case A t log ke Q log k
Mellowmax:
For the mellowmax operator, AT S 0 since the logmeanexp is bounded by the
MM
max input. In the worst-case, we have the following bound log Σ k (e s 1 log e" log(k)
Q:
Thus, A MM
= A AM
log(k)
C.0.1 GENERAL MELLOWMAX
The general mellowmax operator 2 can display accumulation dilution depending on the value of q. To bound | Ar we consider the cases A and ATCM separately, We show bounds for 0,w 0, q E (0, and k E N when 3 > 1.
TCM
We begin by noting that by factoring out Q and denoting 6; Q(s, a,) - Qú, we can rewrite AT as follows. Without loss of generality, we assume the first Q-value is the largest, i.e., Q = Q a1):
A TCM = (og()( e - Q- log(1 + Σ ((a)+)))) - - log(1 + Σ co)))
(14)
We can then make the change of variables = adó which is continuous and bijective on [0, + In particular, since 8, S we are only interested in yı E [0, 1]. To simplify our analysis, we consider I3 defined as follows
&
ga(u2 "e/V),...,,))) = log(1 + Σ v glog(1 + Σ 9:),
(15)
1=2
with ge
endi) A
TG
Accumulation bound: We begin by showing that 15) is maximized at y2 _... Vk = for
Ogg By
y/ E [0, 1]. To see this, we consider the partial derivatives By,
1+ + v. 9 -ー 34
which is
strictly positive for yi < (since 3 >> 1). As a result, the function is increasing in all the y on [0, 1]. implying the maximum is attained the boundary. Hence,
GM
1 (1 1) = (log(A) - q glog(k)) (1 q) log(k)
w
Dilution bound: The dilution bound is more complex. Using the above, we have that - -77 TCM -U2 S - (p.......))Euuuu ..., y.)).
Proposition C.1. Let a > O, q E (0, is > 0 and k E N with The function
ga(u..... yn) log(1 + - log(1 +
Σ vi)
im
attains unique minimum 00, 1 and at this minimum, 3/2
yk
Proof. It is elementary to compute
0ge Dy:
Ogs 8
1/
8-1
y
8-1
ay;
1 Σ 2 Vi
23

---

Published as a conference paper at ICLR 2026
Since 3 - 0, for (1/2..... Mke [0,
aga
age if and only if ye Vy. This implies all By;
By:
local minima of gs in [0, 1,,1|=1 are on the diagonal y/2 To show the minimizer is unique, we now consider
yk.
ho(y) == 9..................
...... .)
log(1 + ny - q logg(1 + ny),
where n := k - 1. By computation, (0) = and h = (1 q) log + n) 2 To show the minimizer in y E (0, we need to show that hz(r)) 0 for some y (0,1) Using I+r S log(1 + x) S x for x > -1, 1, we get that
hz(v) S ny
quly ny
my 1
q
1 + my
Taking y = (3) we get
8-1
q + ny
q
1 + ())
where the left term and the right term is > 3. satisfying hv(v) < O. Hence, since hz(0) = 0. ha(v) < 0 for some y E (0, 1) and hz( 2 0 for 2 1. by continuity, the minimum of hgg on [0, +00) exists, is attained at least once, and all minimizers are in (0, 1). Furthermore, to show uniqueness, consider
Bny + ny
( q) + By ql
(v(
9 + ny
(16)
(1 + n1) + +ny²)
We remark that the numerator of (16) is strictly increasing function of for y 2 0. Thus, h1(()) has at most one solution.
0
From Prop. C1
we have that
min (e Wh ) 10
(0( min hz(3) max -ho(y).
yE[0,1]
ve[0.1]
We aim to find upper bound for -ko(y). For 0. > 1, define
g(x, B, == h2(x)) glog(1 + cr) logg1 + cx²).
Proposition C.2. Let q (0, 3 and > O. Let x² and M (3, c) be the argmax and max of (,,3,e) (defined ubove) for a E [0, 11, respectively. Then, we have that
M(3.c) S q (0-11-11 (1-1) log(1 + cc..
Proof. By Prop. C.I x= (3, c) is well-defined and 0 < (B,c) < 1 for and 0. x" (B, c) not defined when 3 = or O because g(-, 1,... Nonetheless, M c) is continuous from the right = 0, where it has limit 0, and
where it has limit
We denote by Jx. ga, Jc the partial derivatives of g. The constraint that gı (=(x" (())))))))))))))))) 8,),8,c) = -0 simplifies
3x (
q
(17)
1 cz
1 1 cx (3
24

---

Published as a conference paper at ICLR 2026
Using this, we can simplify the derivative of M with respect to
(M(3)))
2x
(r,
90(x" (8, + go(x" (B,c),3,0
dc
dc
(3, c)
x'(3,c) c).
1t cx (8,c)
1 cx (3,c)8
(B,c)
(1-3)
(1-
(Using
1t cx (3,c)
8
9
(Using that x² =B,,c) < 1)
1+c
3
Integrating this inequality, for 2
1tt (1-3) dt 10 log(1 +c).
q (1
S M(8.0) $
Using Prop. C.2 we can plug in our parameters of interest by replacing 3 with gart and e with k get
qa w a qa + w
M(S c) .
log(k) S
qa + log(A)
Finally, getting back to - TTM have
-A TO
- hb(y)) M(3,k- 1)
log(k)
min 24
go(U2. 3k)
yE|0.1]
qa +
Notably, this bound allows recover worst-case bound on dilution of log(k) for the soft mellowmax operator (where 1). Combining bounds, we can solve for the that yields the same bound for accumulation and dilution for GG Doing so yields q a- 4474
a+w
20
25

---

Published as a conference paper at ICLR 2026
D MULTI-STEP UNCERTAINTY SETS
While Fig. 3 examines the case of uncertainty sets for a single step, we now seek to explore the values that can be taken by ó; throughout a trajectory and how this affects the final uncertainty set. For the entropic regularizer, if wis is the same for all s, then Ra, is the same at every state. However, the final uncertainty set R for an object x constructed through trajectory of length given by E; Rs (i.e. the Minkowski sum of the individual uncertainty sets).
We extend the binary generation example DCP consisting of sequence generation task with two tokens (i.e. binary tree) and length d. Then, using Prop Dr1 in Fig. 7 we illustrate how the sum of uncertainty sets changes as the depth of the tree increases for different operators.
Proposition D.1. For a convex regularizer if and associated robust set R = {6 RRA (*) = 0}. have that:
Σ R = (rE kf"
SO}.
(18)
1
Proof. We seek to show that, for given E R RRIAA, there exists 51, 62 ... 5y such that Σ 6, and 8 Ry if and only if kf" (t) S O.
[+]
Since f is convex, so is f". Thus, using Jensen's 's inequality, we get that:
6, k S Σ (ó.) 00
kf 3 = kf
where the last inequality stems from 8; E R2 = (5:) 0.
Taking 6; & for all 6; we have that Σ, 6, = Then, kf (t) <0 = Vi f(6.) and thus 8, R,.
Soft Bellman
SMM
TGM (q 0.5)
15
Depth
15
Depth
15
Depth
d= 1
d= 16
d= 16
d 2
8
d 8
5,(2]
4
5
5
d
5
d 16
d 1
1
10
15
10
15
10
15
6[1]
6
6:13
Figure 7: Illustration of sum of uncertainty sets along trajectory. (Left) The final uncertainty set shifts farther away from (0, 0) as we increase d. As such, for objects that are deeper in the tree, the operator is only robust to objects having much higher reward than their proxy evaluation. This behavior provides an alternative viewpoint for the preference of the soft Bellman operator for longer objects. (Middle) On the other hand, the SMM operator stays centered around (0, 0) and becomes increasingly robust to decreases rewards. (Right) The final uncertainty set of the GM operator stays close to (0, 0) and is robust to some decreases in reward.
26

---

Published as a conference paper at ICLR 2026
E TRAJECTORY GENERAL MELLOWMAX
E TGM HYPERPARAMETERS
We summarize the TGM hyperparameters (with reward being given by (0(x)) below. Recall that TGM solves the following regularized problem:
arg max Qa)
(q (ofl + (1-
4)(-((T)))]
whose corresponding operator is given by
2 (Q.) =
log log(softmax(aQ4)) euQ.).
Parameter
Interpretation
Range
Temperature the softmax the Q-values used by K1. regularizer
Regularization coefficient
Interpolation coefficient between entropy and KL regularization
[0 1]
Exponent of the reward function
R
Table 8: Summary of TGM hyperparameters and their interpretation.
E.2 PROOF OF OPTIMALITY
Notation: We denote by Ch(s) the set of possible actions from state o, (Qa)
En'ec erQateT is the softmax with inverse temperature т. For a DCP and a given state 8, we denote by 8 the state reached by taking action a (unique since transitions are deterministic). We denote by E the set of edges in G, each of which represents a transition (s, a, sy,). Finally, for convenience, we denote
o(Q,,a) 1
log gw (Q.)a( q glog (a)))
w
1
(Q.) = log(o, (Q.) R.)..
Proposition F.1. DCP (9, r), for any state 8,
0 max Q.) (T)) log
(19)
and
Faotw (Qa).
(20)
is the unique maximizer of the state-wise regularized problem.
Proof. Unlike the case of a generic n. 19 is not convex conjugate since Q. appears in both the dot product and Ω. Nonetheless, Ω is still convex in æ. Thus, similarly to (Nachum et al. 2017). consider the Lagrangian:
L( (T,, Q.) - (T.) + X(1 (1, m.)).
Since Q is convex, the overall function to optimize concave. Given that Slater's condition holds, the KKT conditions are both sufficient and necessary for optimality. Due to the lack of cycle in DCPs,
we have that V, Q. = 0. Hence, the KKT conditions yield the following system of equations:
Q₂ 1
(qllo n, log TQ. + + (1 g)ooo म. +1) -
(21)
(T.. 1)
(22)
27

---

Published as a conference paper at ICLR 2026
where we denote "q, Fa (R.). Solving for n, in the first, get that the optimal policy n. is given by:
T exp([Q. - Alo + qlogg= =q, 1)
which ensures that n 2 0. Replacing the above in 22 we have that:
x
log
Σ
(aqtw]Q.lul log
@aQu
1
(oo)) Qu 1).
oECh(s)
aECh(s)
We note the similarity of and the maximal value 19
Plugging X back in to ny we get
Qqutt (Q.).
20
Finally, plugging Th in (T.. Q.) obtain (9)). Since both the KL divergence and the negative entropy are strictly convex, (00 is the unique maximizer.
y
22 (T,) with some straightforward algebraic manipulations, we
Corollary E.1. For any DCP (9, ②). the GM operator has unique optimal value/policy functions (v",
Proof. We show this by induction on an inverse topological sorting 80, 81 .... 8n the states of G. For state with children, the optimal value is defined to be O and there is no policy. Then, assume each state 8,, í < N has a unique optimal value/policy. By the properties of the topological sorting, all the children of sn have unique optimal value/policy. Then, by Prop. E.1 (v xgw) is also uniquely defined.
Lemma E.1. For any Qs and a, we have that
g(Qs, a) = Qalaa-g" gg"
(23)
Proof. By simple algebra,
g(Q,,a) a)
1 # (log qw (0o++((..(4)) - glog o. (Qa)(a))
((qa + Q.(aa - (aQu(aa)
9
1 log [aa++) log
Σ
Qulal
(C)())
a' a'ECh(s)
qqa (a') Qulaaa
(wQ, (a))
log
w
w
Ex'ECn(O) eaQula")
('ECh(a)
[a 1 (log(oa(Q.)) ew))))))))))))))
Q, [a] g" (Q.)
Lemma E.2. For any state s in a DCP, the the following are equivalent:
aaqt (Q+)) arg max (m Q3) ~ 1
(mx),
(24)
n,
Ya Ch(s) v = + +Q a) r(s,
(25)
Proof.
28

---

Published as a conference paper at ICLR 2026
24
25
Since ni is unique, by 24 and 20 by the properties of the softmax, have
that Oaatwl (22)
gqatw(Q2 (Q²) and aa (Q2) Fa (Q.) Then, for any a E Ch(s):
0 g(Q
= = (log T
qart (?a)[a] - q log 0a (Q.)aa)
(1o)( (.) ) + log go +w [Qt)[a] - loo 0. (Q3)
Using
19
log EareCh(s) e(ga+w))1(a')
Eaveh(s)) eaQila'l + log aa tw 10lolo-lo000221
1 (loge (lgge(n+))) log ((14l)
w
Q [a]
oo + a(
Using the definition of Q. [a).
24
25
For any Ch(s), we have that:
" = V + s(Qu, a) r(s,a)
o. vf + Qa(a] - g" g²( r(s.a)
Using
9 (Q.) [a]a]] [a
Since the above equality holds for all actions and vi and 9 (Q.) do not depend on a, we have that
Qi - [a] must be equal for all actions. Thus, Va E Ch(s) : Qi[a] Q [a] which implies
[1 Q1(4]
(26)
for some action independent value Finally, we have
Tagtw (Q.) = o oaqtw (Qr -c) = (oq+
by the properties of the softmax.
E.3 PROOF OF THEOREM
We note that it is straightforward to derive the equivalent subtrajectory constraint.
Theorem E.1 (Trajectory GM). For any DCP (G, 2), let (u", x²) x²) be the unique optimal value/policy functions for the GM operator. Then, for a given Q-function Q?,
Tagtu (agtw(Q2)) arg max( Qa) (.23((.))
holds for all states s if and only if
(log aga ( glog a. (Q.))a -r(si,
10
a(Q" ms)
holds for full trajectories 8o do
8n an in the DCP.
Proof. Similarly to (Nahhum et al. 2017). ) we aim to show that optimality implies consistency and vice versa. A refresher of relevant notation be found at the start of App.
Optimality implies consistency
[0P] From Lemma E.2 for any given trajectory, we have that (25) holds for all transitions in the trajectory. Then, (10) follows straightforwardly by expanding that trajectory in recursive manner.
25
29

---

Published as a conference paper at ICLR 2026
Consistency implies optimality of G such that for any state s', any child s of s' has index j < i.
9
I0
Consider an inverse topological sorting s' s? S
Now, seek to show by induction over this sorting that the following holds for any state and subtrajectory 5o 7 51 ... 7 Sn + s ending in s:
66
(log ta (%%, [a.] - glogon (logoa(Q2)((a)) r(8i, 0z) v.-
(27)
For the base case st s must be leaf and belong to X. Then, by assumption Vga = and (7 holds by (10). Now, suppose that (27) holds for any s', < N. Then, by the definition of the sorting, for all E Ch(s (27) holds for 8, where sN S Sa-
Let sn + $1 + Sn + be a trajectory ending in sN. From the induction hypothesis, da E Ch(s):
to
g(0 a;) - r((s;, ((Q0,, - -r(s",00 = 1%
+ Σ E9 r((s, -00---g'000 00(xxx) = -ox - -Qa [a]
Using 23
Similarly to the proof of Lemma we have that the LHS independent of a must be equal to Qi [a] 22 [a] for all a. As such, we guaranteed
2º [a] = Q2 [a]
for some action independent value and all a E Ch(s Then, it is straightforward to compute g" (Q"xx) = pQ"x c) = g" (Q²x)) c. Plugging both back and canceling c, we get
g( a.) - = =g" ViN
00
and the induction is complete. Finally, it remains to show that ao is optimal for all s has no children, optimality holds trivially. If s has at least one child a, take a trajectory ending in From (27). we have that both:
so + $1... + Sn
Vn +
Σ g(0 ai) - r(888, as) = v.
1=0
Ya E Ch(s): vo + Σ gloa a;) r(s, a:) + g(Q.. a) as v=.-
By subtracting the first from the second, we get that, Va E Ch(s):
5(O r(s,a) = v/. v.
3 + 9Q a) r(s, 1
By Lemma
B.2 Uagtu (Q?) thus optimal s and we are done.
30

---

Published as a conference paper at ICLR 2026
F EXPERIMENTAL DETAILS
F.1 ALGORITHMS
We follow most the hyperparameter decisions (Malkin et al., 2022). For each method, we train a transformer with 3 layers, attention heads, and embedding dimension of 64. We set dropout to 0.1 and use causal masking. Each output head uses a fully- connected network with hidden layers of dimensions 256.
We use weight decay of 10 and use gradient clipping with a threshold of to help stabilize training at higher / values. When generating samples for training, a policy corresponding to mixture of the method's policy and a uniform policy (with weight 0.01). Each method is trained with Adam (Kingma 2014) with 10-5
TGM/GFN:
Since GFNs correspond to specific hyperparamettr setting for TGM, training details are essentially identical between the two methods. Both methods are trained online with a batch size of 16 trajectories and have single head that outputs logits. We sample trajectories by taking the tempered softmax of the logits, using the temperature eq + w from the optimal policy (20) For GFNs, affects the sampling policy but not the actual training objective this ensures the network is learning the actual GFN objective.
For SAC, we use three networks: two Q networks (to reduce overestimation bias) and a policy network, following the implementation of (Christodoulou 2019 Huang al. (022). We use replay buffer of size 100 000 and a batch size of 024 transitions. Generation/traiiin is adjusted to have replay ratio of 1.0. The entropy coefficient is fixed at 1/w, and we use target networks (for the Q networks), which get updated every 10 iterations.
SAC:
PPO: For PPO, we use two networks: a policy network and a value network, following the implementation of (Huang al. 2022). We follow the hyperparameter settings of (Huang al.) 2022) and use a generalized advantage estimator À of 0.95 and = to avoid discounting. The clipping epsilon 0.1 and the value coefficient is 0.5. During online training, 16 trajectories are generated, and the resulting transitions are split into 4 minibatches for training. The entropy coefficient set to 1 1/w.
F.2 SEQUENCES
For each sequence generation task, we add the following 3 tokens to the vocabulary: BOS, PAD, EOS. Sequences are padded to be of length equal to the maximum length for the task, with an added BOS, EOS. Sequences look like the following: [BOS, 21, , ,,, PAD..... PAD, EOS]. When generating sequences, they start with BOS, and tokens are added until terminating action is chosen, which point PAD/EOS are added automatically. The terminating action masked until the task's minimum length is reached.
F.3 SYNTHETIC TASKS
TF-Bind-8. Hyperparameters described in the paper. We sample 10000 sequences for each method using their respective optimal policy. Fig. i the kernel density estimate computed using these samples with the default parameters used by seaborn (Waskom 2021
Bit sequence. We use 120 and k 8 with MM 60 modes and force sequences to be of length 120 through action masking. The modes are sampled randomly and fixed across runs to ensure there is no variation from certain runs having harder-to tonfr--t--fin modes. Each method is trained for 200 000 samples.
F.4 BIOLOGICAL SEQUENCE DESIGN
For each task, once the proxy reward model is trained, we compute the mean and standard deviation of the logits the validation set. Then, use the following normalized value reward
31

---

Published as a conference paper at ICLR 2026
6(x)
Mval
- - (
(28)
vval
where ((x) is the logit output by the trained model. To train the proxy reward models, we follow the hyperparameters used by (Malkin al. 2022).
UTR: Using the described dataset, we train transformer with 4 layers, 8 attention heads, and an embedding dimension of 64. 80% of the data used for training and the rest for validation. We train for 250 epochs with early stopping (using patience of 15). The learning rate is set to 10-1, and we use a batch size 128 and weight decay of le-6. set both the minimum and maximum length to 50 since all sequences in the dataset of length 50.
We use the same hyperparameters for training the model as UTR. Since the dataset has sequences of variable length, set the minimum length to 14 and the maximum length to 60. Unlike the other tasks, the model is trained for binary classification, and we use the logit passed to the sigmoid as (xx).
AMP:
GFP: GFP is the largest task. We train transformer with 3 layers, 8 attention heads, and an embedding dimension 128. Since sequences are length 237, we set both the minimum length and maximum length to 237.
F.4.1 EVALUATION
For evaluation of the average mode reward, we sweep over inverse temperature modifiers t E (0.005, 0.01, 0.02, 0.05, 0.1, 0.2, 0.5, 2,2,5) (i.e. if the regular sampling policy was xx eT we sample using n' x e eTa²) and generate 512 samples for each. We then aggregate the samples and greedily find the top 100 samples such that they are at least at a distance dmin of each other. We use the Levenshtein edit distance metric (Levenshtein et al. 1966) and set dmin = 0.25 minLen++malLer (i.e. dauin = 13 for UTR, dunin 10 for AMP and dmin = 60 for GFP). The constraint can be roughly interpreted requiring least 25% of sequence be different for it to be considered distinct. In Tab. we report the max of this value achieved over the course training.
F.5 COMPUTE RESOURCES
Experiments were conducted on cluster with mix of GPUs. All experiments were run using jobs that requested 1xL40s, 24G of RAM, and 4 CPUs. The cluster allows for dozens of jobs to be run in parallel. Overall, the experiments (including the hyperparameter sweeps, baselines, and various seeds) used roughly:
Bit Sequence: 35 GPU hours.
UTR: 48 GPU hours.
AMP: 48 GPU hours.
GFP: 300 GPU hours.
for total of roughly x 20 GPU days. In particular, the entire codebase uses Jax (Bradbury et al.) 2018) and a parallelized environment allowing for fast generation. Since we use transformers, training generated trajectory for TGM/GFNs only requires single forward/backward pass. This significantly speeds up training (roughly 4x faster) compared to the PPO/SAC baselines, which operate individual transitions (though there potential engineering optimizations improve the speed of these baselines).
32


Reviews:
[{"review_id": "TYPyco7Hif", "reviewer": "Reviewer_9b9M", "summary": "The authors introduce the Trajectory General Mellowmax operator, motivated by the framework of regularized reinforcement learning operators, which tend to encourage more concentrated policy distributions, and seem to show improvements across datasets.", "questions": "1. It seems that several baselines ([A], [B], and [C]) are missing. Since these approaches also use a Q-function, comparing against them would make the evaluation fairer. In contrast, GFlowNets do not rely on a Q-function, which makes direct comparison more challenging, as their learning dynamics and compute requirements differ fundamentally. Moreover, in the middle panel of Figure 4, the GFlowNet performance appears to be increasing, does this suggest that GFlowNets could further improve with additional compute (e.g., equivalent to what is used for Q-function training)?\n\n2. In the intuition for the motivating example, doesn’t the argument implicitly depend on the quality of the reward model? For instance, if the rewards were accurate, then each sequence wouldn’t necessarily need to have its reward lower bounded by a constant r, since poorly performing samples could simply receive large negative rewards. Wouldn’t this issue therefore be mitigated by using a more accurate reward model?\n\n3. In general, GFlowNets aim to generate diverse and novel samples. With that in mind, could the paper report diversity and novelty metrics as discussed in [E] and [F]? Additionally, similar to Figure 5 in [D] and Figure 2 in [E], would it be possible to conduct synthetic experiments to evaluate whether TGM successfully recovers all modes?\n\n4. Could you clarify how to interpret Figure 2? What exactly should the reader take away from it? (I don't think I understand the caption)\n\n5. In lines 403–404, could you explain why $\\delta = 28?$\n\n---\n\n**References**\n\n[A] Mohammadpour, Sobhan, Emmanuel Bengio, Emma Frejinger, and Pierre-Luc Bacon. \"Maximum entropy gflownets with soft q-learning.\" In International Conference on Artificial Intelligence and Statistics, pp. 2593-2601. PMLR, 2024.\n\n[B] Lau, E., Lu, S., Pan, L., Precup, D., & Bengio, E. (2024). Qgfn: Controllable greediness with action values. Advances in neural information processing systems, 37, 81645-81676.\n\n[C] Tiapkin, D., Morozov, N., Naumov, A., & Vetrov, D. P. (2024, April). Generative flow networks as entropy-regularized rl. In International Conference on Artificial Intelligence and Statistics (pp. 4213-4221). PMLR.\n\n[D] Zhang, D., Malkin, N., Liu, Z., Volokhova, A., Courville, A., & Bengio, Y. (2022, June). Generative flow networks for discrete probabilistic modeling. In International Conference on Machine Learning (pp. 26412-26428). PMLR.\n\n[E] Jain, M., Bengio, E., Hernandez-Garcia, A., Rector-Brooks, J., Dossou, B. F., Ekbote, C. A., ... & Bengio, Y. (2022, June). Biological sequence design with gflownets. In International Conference on Machine Learning (pp. 9786-9801). PMLR.\n\n[F] Ekbote, C., Jain, M., Das, P., & Bengio, Y. (2022). Consistent training via energy-based gflownets for modeling discrete joint distributions. arXiv preprint arXiv:2211.00568.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 2, "strengths": "1. The paper addresses an important problem—developing a reinforcement learning method that can generate diverse responses while also producing high-scoring candidates useful for biological and chemical discovery.\n\n2. The paper is well written and generally easy to follow.", "weaknesses": "1. It seems that several important baseline comparisons are missing (see questions for details).\n\n2. I’m not sure it’s appropriate to compare GFN with TGM, since GFN involves training a Q-network and therefore requires significantly more computation. The compute budget should be made fair for a valid comparison (see questions for details).\n\n3. I also find the motivating example somewhat unclear or unconvincing (see questions for details).\n\n4. Some of the figures are hard to understand.", "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": "RVyi3Fu3yp", "reviewer": "Reviewer_5kkr", "summary": "The paper proposes Trajectory General Mellowmax, a novel algorithm that outperforms GFN in finding high-reward targets on discrete state spaces. The experiments conducted in both the synthetic environment and the biological sequence search demonstrate that the newly proposed methods are effective and relevant.", "questions": "See weakness.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The proposed methods are novel and theoretically grounded, with a clear motivation.\n- The numerical experiments are comprehensive, showing the relevance of TGM", "weaknesses": "- The writing is pretty challenging to follow.\n- The motivating examples of \"high reward path dominated by many low reward paths\" seem to be the issue of the designed target distribution rather than the problem of the learning algorithm. For a 0-1 reward, if we reduce the temperature of the distribution to be proportional to exp(r/gamma) where gamma -> 0, the valid target examples are only going to be concentrated on positive reward ones, and the cases presented in the motivating examples no longer seem to be valid.", "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": "gQk3BOQrGL", "reviewer": "Reviewer_QUN8", "summary": "This paper introduces trajectory general mellowmax, a reinforcement learning method for discrete compositional generation tasks like molecular or protein design. The authors highlight how methods that learn a policy whose density/distribution is proportional to the reward (such as GFlowNets) tend to produce diverse but generally low-quality samples (since the high reward, low density samples are drowned out by the low reward, high density samples). To address this, the authors propose the general mellowmax operator, which unifies and interpolates between existing soft RL operators to better balance diversity and quality. Experiments on synthetic and biological benchmarks show that the proposed approach consistently generates higher reward and still diverse candidates than prior methods.", "questions": "Are there ways to experimentally demonstrate the connection to robust RL? E.g., by introducing noise into the proxy reward to evaluate how TGM performs relative to GFlowNets?\n\nHave you compared TGM against established alternatives that address policy over smoothness / diversity such as temperature conditioned GFlowNets? I saw temperature conditioned GFNs mentioned in the paper, but not used as a baseline for comparison.\n\nFor the biological sequence design experiments in 6.2, it would be useful to understand how the diversity of the generated sequences varies across the baselines and TGM with different choices of q.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "This paper is original in unifying multiple soft RL operators through the proposed general mellowmax framework. The proposed trajectory general mellowmax is significant in bridging GFlowNets and robust RL under a common theoretical perspective. The paper contains rigorous mathematical derivations, a well motivated algorithmic design, and thorough experiments across both synthetic and real wold biological design tasks. The problem motivation, operator formulation, and empirical findings are clearly articulated, and the work is significant in meaningfully addressing a key limitation of excess diversity (at the expense of high reward designs) in GFlowNets / similar constructions.", "weaknesses": "The mathematical notation can feel a bit dense at times and I found myself getting occasionally lost on what the various key parameters are meant to govern. It would be helpful to have a better illustration or explanation on the role of alpha, q, and omega in TGM - this is attempted in Figure 3, but requires the reader to go to other portions of the paper to understand what the axes are describing. So in general, I feel that the presentation and clarity of the paper could be improved. I also feel like the experimental section would benefit from including a baseline that is also designed to address the mode concentration issues of GFlowNets, such as temperature conditioned GFNs.", "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": "aAIPEP3zF4", "reviewer": "Reviewer_iNss", "summary": "The paper addresses the problem of constructing/composing discrete objects that maximize an uncertain (learnt) proxy reward and are diverse. The starting point of the paper is Generative Flow Networks (GFNs) which learn to sample proportionally to some reward function. However the number of possible discrete objects is exponential to their length/number of their components. Thus sampling proportionally to the value of the reward function will tend to produce samples that are of low reward, simply because low reward samples will have collectively a larger probability of being selected than the few high reward samples. Thus the problem that the paper sets to address is given equation 1 of the paper:\n$$ \\max_{x_1, ..., x_k} \\sum_i \\exp{\\Phi(x_i)} \\text{ subject to } d(x_i, x_j) > \\delta, \\text{ eq. 1 } $$\nwhere $\\Phi(.)$ is the learnt reward function and $d(.,.)$ some distance measure meant to capture diversity in the set of selected structures. \n\nThe paper casts this problem as learning a policy that is robust to the worst-case reward model, i.e. the value function that they seek to maximize is: \n$$v_{\\mathcal R}^{\\pi} (s) := \\min_{ r \\in \\mathcal R} v^{\\pi}(s)$$\nUncertain rewards create reward uncertainty sets and a robust policy should be optimal under the worst case reward; however at least as eq. 1 above is stated the initial optimization problem that the paper sets to address looks rather like a standard optimization problem and not a robust optimization problem.  \n\nGiven that one should learn a robust policy to address 1 then the paper proceeds with the following contributions:\n* a regulariser that explicitly trades off between different regularisation functions associated with GFN, and ((soft)mellowmax), \n* an algorithm that uses the above regulariser and operates on the level of trajectories and samples from high reward regions\n* a characterisation of the reward uncertainty sets (created by the uncertain reward) that are induced by convex regularises\n\nThe paper discusses the relation between GFNs, soft Bellman, and the (soft)-mellow max operators: \n* There is a direct equivalence between GFNs and soft Bellman (entropy regularised MDP/policy $H(\\pi)$, both suffer from value accumulation: the value accumulation of a large number of suboptimal objects can overwhelm a single high reward object\n* The mellowmax operator solves accumulation but suffers from dilution; a high reward object can get diluted by lower reward objects.\n* The soft mellowmax addresses the dilution of mellowmax \n\nThe paper argues for controlling the balance of accumulation and dilution and does so by addressing a regulariser that allows to seamslessly transition between the three operators.  GFNs and soft Bellman are entropy regularised MDP/policy $H(\\pi)$); mellowmax/soft-mellowmax learn  KL regularised policies, $KL(\\pi,d_s)$, where $d_s$ are the uniform distribution and the $softmax(Q)$ distributions respectively. The paper introduces the general mellow max regulariser that interpolates between the two as:\n$$ q KL(\\pi,d_s) + (1-q) (-H(\\pi_s)) \\text{ eq. 6 } $$\ninterpolating between policy entropy maximization and getting a policy that is close to the softmax of the $Q$ values, capturing GFNs, and (soft)-mellowmax for different values of $q$ and temperature. \n\nThe paper then proceeds to introduce the trajectory general mellow max which builds upon eq 6 and operates over trajectories. The algorithm is evaluated over a set of synthetic and real world sequence generation benchmarks that show that TGM achieves high rewards against RL baselines, namely PPO and SAC, and it is on a par or better than GFN. The experiments show that TGM has a performance that is quite stable with respect to difference values of the interpolation coefficient .\n\n\n\nIn addition the paper discusses the equivalence of entropy regularised MDPs to robust MDPs with uncertain reward (robust RL). \n\nThe paper proceeds with a discussion of the uncertainty sets introduced by the different regularisers, i.e. the entropy regulariser used in GFNs, the softmellowmax, and the general mellowmax. In particular the setting of discovering interesting structures one would like to maximize the proxy reward $\\Phi$ while accounting for the difference, $\\delta$, between $\\Phi$ and the true uknown reward $r^*$, by learning a sampling distribution $p$ that satisfies:  \n$$ \\max_{p} \\min_{\\delta \\in R} E_{x \\sim p}[\\Phi(x) + \\delta(x)]  \\text{ eq. 10 } $$\nwhere $R$ is the reward uncertainty set. It attributes the failure of GFNs to maximise $\\Phi$ to the fact that the uncertainty set produced by its regulariser does not contain the proxy reward.", "questions": "* Robust optimization seeks to maximize an objective function in the worst case under uncertainties. The objective as given in equation (1) does not look like a robust optimization problem but rather takes $\\Phi$ at face value and does not state anything about a worst-case setting. So why robust optimization is appropriate for eq. (1)? \n$$ \\max_{x_1, ..., x_k} \\sum_i \\exp{\\Phi(x_i)} \\text{ subject to } d(x_i, x_j) > \\delta, \\text{ eq. 1 } $$\n\n* how should I understand equation 10: \n$$ \\max_{p} \\min_{\\delta \\in R} E_{x \\sim p}[\\Phi(x) + \\delta(x)], \\text{ eq. 10 } $$\nin the sense that $\\delta$ is not something we have access to, is this more a theoretical concept to study the properties of the different regularisers? Reading it as is stated above there are two things that can be adjusted/learnt the sampling distribution $p$ and the $\\delta$ function (or I guess equivalently the $\\Phi$), however $\\Phi$ is fixed. In addition the max/min problem reads as: find the distribution that maximizes $\\Phi$ under the lowest error, this does not sound like optimization in the worst case setting.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The introduces an interpolated regulariser which unifies three different operators and a respective algorithm that achieves the desiderata of oversampling high reward, demonstrated by the empirical evaluation\n* The paper characterises the uncertainty reward sets associated with the different regularisers and provides an explanation on why GFlowNets do not oversample high reward regions by relying on their equivalence entropy regularised policies: their reward uncertainty set does not include the proxy reward.", "weaknesses": "As I have said in a previous review of the same paper: The paper addresses the property optimization problem for (robust) sequence generation in an RL setting. I understand that the core of the paper is the robust property optimization with RL [...] however there are different approaches to the same problem that are [still] not discussed in the related work. In particular there is quite some work in generative models that seek to do exactly the same, i.e. property optimization. \n\nSee references provided in the previous review. I think it is still useful to discuss also such different modelling approaches to the same problem.", "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)