Action abstractions for amortized sampling
Authors: Oussama Boussif
Lena Nehale Ezzine
Joseph D Viviano
Michał Koziarski
Moksh Jain
Nikolay Malkin
Emmanuel Bengio
Rim Assouel
Yoshua Bengio
Venue: ICLR 2025 Poster (2025)
Keywords: GFlowNets
amortized samplers
hierarchical planning
abstractions
macro-actions
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: Augmenting existing samplers with an action abstraction step using tokenizers on high-reward action sequences improves credit assignment.

Abstract: As trajectories sampled by policies used by reinforcement learning (RL) and generative flow networks (GFlowNets) grow longer, credit assignment and exploration become more challenging, and the long planning horizon hinders mode discovery and generalization.
The challenge is particularly pronounced in entropy-seeking RL methods, such as generative flow networks, where the agent must learn to sample from a structured distribution and discover multiple high-reward states, each of which take many steps to reach.
To tackle this challenge, we propose an approach to incorporate the discovery of action abstractions, or high-level actions, into the policy optimization process.
Our approach involves iteratively extracting action subsequences commonly used across many high-reward trajectories and `chunking' them into a single action that is added to the action space.
In empirical evaluation on synthetic and real-world environments, our approach demonstrates improved sample efficiency performance in discovering diverse high-reward objects, especially on harder exploration problems.
We also observe that the abstracted high-order actions are potentially interpretable, capturing the latent structure of the reward landscape of the action space.
This work provides a cognitively motivated approach to action abstraction in RL and is the first demonstration of hierarchical planning in amortized sequential sampling.

Published as a conference paper at ICLR 2025
ACTION ABSTRACTIONS FOR AMORTIZED SAMPLING
Oussama Boussifl.2 Léna Néhale Ezzine
Joseph Viviano
Michal Koziarski
Moksh Jain 1.2
Nikolay Malkin
Emmanuel Bengio*
Rim Assouel!2
Yoshua Bengio 1.2,6
Mila - Québec Al Institute Université de Montréal
3 University of Toronto CCIARR Al Chair
4 Valence Labs, Recursion
University of Edinburgh
ABSTRACT
As trajectories sampled by policies used by reinforcement learning (RL) and gen- erative flow networks (GFlowNets) grow longer, credit assignment and explo- ration become more challenging, and the long planning horizon hinders mode dis- covery and generalization. The challenge is particularly pronounced in entropy- seeking RL. methods, such as generative flow networks, where the agent must learn to sample from a structured distribution and discover multiple high-reward states, each of which take many steps to reach. To tackle this challenge, we propose an approach to incorporate the discovery of action abstractions, or high-level actions, into the policy optimization process. Our approach involves iteratively extracting action subsequences commonly used across many high-reward trajectories and 'chunking' them into a single action that is added to the action space, In empirical evaluation on synthetic and real-world environments, our approach demonstrates improved sample efficiency performance discovering diverse high-reward ob- jects, especially on harder exploration problems. We also observe that the ab- stracted high-order actions are potentially interpretable, capturing the latent struc- ture of the reward landscape of the action space. This work provides a cogni- tively motivated approach to action abstraction in RL and is the first demonstra- tion of hierarchical planning in amortized sequential sampling. Code is available at htttt::://// / http//gttttttttt thub. com Cooom///Nrgg//uuukk
INTRODUCTION
Reinforcement learning (RL) relies stochastic policy To to generate potentially long trajectories T of actions a to obtain a reward r. Standard RL methods take learning steps on the policy parameters
0 using loss function that reinforces actions which maximize reward (Sutton & Barto 2018). In
RL methods, such as generative flow networks (GFlowNets: [engio let al. 2021: 2023) special case of engropy-regulariied RL (Ziebart et al. (2008); see Tiapkin et al. 2023); Deleu et al. (2024)) - the loss function may instead encourage the policy to sample terminal states with probability proportional to their reward. In both cases, longer trajectories make training more difficult due to the problem of credit assignment, i.e., the propagation of a learning signal over long time horizons (Sutton & Barto 2018).
the case of diversity-seekinn
The difficulty of credit assignment grows with trajectory length: it has been shown that temporal difference learning methods require exponential (in the trajectory length) number of updates to correct learning bias, while Monte Carlo methods see the number of states affected by delayed re- wards grow exponentially with the number of delay steps (Arjona-Medina et al., 2019). Previous work addressed the challenge of credit assignment through long trajectories by propagating a learn- ing signal to specific moments in the trajectory, skipping over intermediate actions and effectively reducing the trajectory length (Ke et al., 2018) Liu et al., 2019; Hung et al., 2019: Sun et al.) 2023).
Recent work has shown that trajectory compression improves credit assignment in RL. For instance, Ramesh et al. (2024) show that it is beneficial to explicitly shorten trajectories by dropping the in- termediate states which are highly predictable given their predecessors, related to the principle of history compression (Schmidhuber, 1992). Other recent work used a similar compressibility crite- rion directly as a reward term: by encouraging the policy to implicitly learn trajectories that could be easily compressed or predicted (e.g., that contain repeating elements), the RL agents achieve better
*Correspondence to: oussama boussi f@mi 1a quebec

---

Published as a conference paper at ICLR 2025
Policy network
Sampled trajectories
Terminal states
Inner loop: Sampler training
policy update
R
12)
rollouts
02
41
82
44
R
otspollcy or exploraiory
ds
filter and compress
augment action space
Outer loop:
Abstraction
Figure I: Chunking procedure. Starting with policy, generate trajectories consisting of action sequences. The trajectories are then filtered to retain only the high-reward samples, which are passed a tokemizer. The tokenizer identifies frequently occurring "chunks" which are added the action space. The process is repeated till convergence.
sample efficiency (Saanum al. 2023). These works suggest that the lower description length of action sequences confers some benefit in terms of the credit assignment challenge.
Analogies exist between compression in artificial and natural learning systems. A long line of research in psychology investigates "chhnkknn",, phenomenon whereby humans tend to record smaller units of information as part of larger high-order units. Chunking believed reduce the load on working memory when reasoning about complex problems which would take too many bits encode using only lower-level units (Thalmann et al. 2019 Gobet et al., 2001 Miller 1956 Johnson 1970).
In RL, options and macro-actions generalize primitive actions and allow for temporally extended sequences of actions, temporal abstractions, that facilitate more efficient exploration and enable hierarchical learning. Macro-actions have inconsistently been found to speed-up learning in RI. and enhance credit assignment (Randlov 1998 [Sutonn et al.) 1999).
Drawing on the above work, we propose to abstract high-order "chunks" online from the trajectories sampled by GelowNet RL policy, producing action space that grows during training. In contrast to (Saanum et al. 2023). compressibility is not enforced through objectives rather, we harness the virtuous eycle of emergent compressibility of the policy's trajectories and the accelerated learning that results from shorter action sequences. In addition, we find that these high-order actions accurately recover the latent structure of the underlying distribution.
Our contributions can be summarized as follows:
We present ACTIONPIECF, method for extracting high-order actions, "chunks," from sampled trajectories using tokenization techniques. ACTIONPIECE includes two variants: ACTiONPIECE- INCREMENT and SERTONNIECE-RELLACE Both approaches compatible with any sampler at no significant cost.
We introduce ShortParse, new backward policy for the GFlowNet for sampling short trajectories.
We evaluate approach multiple environments and determine that accelerates mode dis- covery, improves density estimation, and reduces description length the samples.
We demonstrate that learned chunks represent the latent structure of the target distribution and are transferable different algorithms and new tasks.
2 RELATED WORK
Macro-actions are a sequence of actions assembled from an atomic action set (Randlov) 1998), form of temporal abstraction, and have been the focus of a long line of work. The idea stems from early proposals of automatic induction of macros for programs using various rules heuristics [lbal 1989 Laird et al. 1986); see Sutton et al.) (1999) on the related framework options and a review of macro-actions. In our context, macro-action is policy representing series of atomic
2

---

Published as a conference paper at ICLR 2025
actions to be executed at once, which in some cases has been shown to reduce the time required to discover the optimal policy and action-value function the RL context (MeGovern & Sutton [998] Laird et al. 1986). Previous work has proposed composing macro-actions from n-grams to speed up search during planning (Dulac et al. 2013) and employing them to construct form of a hierarchy of policies with different temporal resolutions (Precup & Sutton 1997| Dayan & Hinton 1992). Such methods have produced mixed results, either helping or hindering performance depending on the appropriateness of the macro-actions to the task (McGovern et al.) 1997, Jong et al.) 2008). In cases where macro-actions were helpful, their utility was explained by improved credit assignment and exploration. Previous work has also used evolutionary algorithms for discovering macro-actions, which was an effective approach on common Atari games (Chang al. 2022). Others proposed to learn macro-actions as either actions that are repeated number of times, or commonly-occuriin sub- sebbeeeqeeces of a fixed length (Durugkar et al. 2016) while training a deep RL policy. In this work, speed of convergence is the most reliable improvement from the introduction of macro-actions, but variance of returns is also decreased, and in some specific settings returns improved. The settings where macro-action are useful appear to contain more shared structure between the learned macros and the latent structure of the environment and reward landscape.
Work on library learning in program induction (e.g., Ellis et al. (2018); Nye et al. (2020) Chaudhuri al.] (202T) has often sought to incrementally modify the Nomain--peccin
(e.g.,Ellis
language of interest so as to reduce the complexity of useful programs, following the minimum description length principle (Rissanen 1978). DreamCoder and related approaches (Ellis et al. 2021 Claire et al. 2023 Ellis
ett 2018 Bowees et al.] 2023) accomplish this by introducing new subroutines during a distinct abstraction phase during training. These algorithms leverage a library of solutions to given problems which serve as prior (see also Tian et al. (2020) Liang et al. (2010)) Many other works in this space only implicitly learn library in the form of policy which emits program when conditioned on a program (e.g.. (Ellis et al. Dev9):Devlin et al. (2017)). Similar works exist in the related domain of proof synthesis (e.g.. Vyskočil et al. 2010 Zhou et al. (024)).
Credit assignment is fundamental problem in sequential decision- maaiiin, known to arise from few properties of the Markov decision process (MDP) in which the learner acts: its depth (the typical number of steps between the initiation of trajectory and reward); its density (the number of nonzero rewards in the typical trajectory) and its breadth (the number of possible routes from the initial state to a reward; for a review, see Pignatelli et al. (2024). In RL, Hindsight Credit Assignment et al., (209)) introduces a new family of algorithms where credit is assigned to past decisions based on the likelihood of them having led to the observed outcome. In GFlow the trajectory balance objectives (Malkin et al.) 2022: Madan et al.] 2023). which consist of loss depending on a whole action sequence rather than a single transition, and the newardd-lokingg flow parametrization (Pan et al. 2023), which is effective when the energy can be decomposed additively into per-time-step components, are both motivated by the credit assignment challenge.
[HCA:[Harutyunn
PRELIMINARIES
We begin by briefly introducing the concept of amortized sampling and summarize the learning problem solved by GFlowNets and the objective used, then introduce the two RL algorithms on which we also evaluate our action abstraction approach.
Amortized sampling refers to the process of sampling from a functional approximation of the target distribution, where the computational cost of iterative sampling is shifted to the model's optimization process. This allows for efficient and rapid sampling once training is complete.
Let G (S, be a directed acyelic graph. Nodes s S are called states and edges (s +> s²s E A are actions, which define the relations "s is a parent of y" and "s' is a child of s". We assume that there is unique state so with no parents (initial state) and denote the set of states with no children (terminal states) by X. One interprets X as the set of complete "objects" that one can sample and the actions as constructive steps that incrementally build object.
A key assumption needed to formulate action abstraction is that there is a correspondence between the actions available at different states. This a common assumption in RL in MDPs. the action space is typically the same at all states, but some actions may be "masked" or "invalid" at some states- but not a typical GFlowNet assumption. Here, we assume that there is a global set of actions A and that at each nonterminal state s subset A, S A of these actions is available. We fix a bijection between A and the set of children of s; thus, applying action a E A, at state transitions to the child state corresponding to a, denoted s + a. (For example, when constructing
3

---

Published as a conference paper at ICLR 2025
sequences over some alphabet by adding one symbol time to the end, A may be the alphabet, with A, = A for all s.)
A (forward) policy Pr is collection of distributions over A, (identified by the aforementioned bijection with the set of children of s) for every non-terminal state s. A forward policy induces a dis- tribution over complete trajectories. paths in the directed graph from so a terminal state Sn X:
PF(T (50
sn)
1
(a st).
(I)
where " 5/41 indicates that + S1+1- The marginal likelihood of sampling terminal state x E X is then
(x) = Σ PF(T), T-XX
where the sum is over all complete trajectories ending in x. (Note that this sum may be intractable to compute exactly if the number of trajectories leading to x is large.)
Trajectory balance, GFlowNet training aims to approximate PF (for example, parameterized as a neural network ( (.)) so as to make PF proportional to a given reward function R : X - Rt. Various objectives for achieving this exist; here, we use the trajectory balance (TB) objective (Malkin et al. 2022). TB requires introducing two auxiliary objects: a learned scalar Za (which estimates the normalizing constant Z = Exex R(x)) and a (fixed learned) backward policy Pa. which is a collection of distributions Ps( s) over the parents of every non-initial state s. Similarly to equation PB induces a distribution over complete trajectories ending in x:
PB(T (so
= xx) 12)
[ Pn(s: I (a+1). 2=0
The TB objective aims to make (T( proportional to P8(T I x) R(x) for every trajectory t ending in x, a constraint that implies that (x( R(x). The training loss for trajectory r enforces this proportionality:
L(;: = [log[ Zo Pf - (x))
(
(2)
If ((T: = 0 for all r, then the desired proportionality is satisfied and PT x R(x). The training procedure takes gradient steps with respect to 0 to minimize L(T: over trajectories sampled from some behaviour policy ((T). As we aim to simultancously minimize equation too for all trajecto- ries, the behaviour policy can be any full-support distribution, such as PF itself (on-policy training) or distribution chosen to promote exploration (off-policy training).
be translated to RL terminology in straightfoward manner: one defines an MDP with states S, action space A, deterministic transition function ass,a) s and reward r(s - s') log R(s') if e X and () otherwise. (The log is nec- essary for the equivalence of GFlowNet and RL algorithms certain cases to hold; see Tiapkin et al. (202)); Deleu et al. (2024b.) We now recall two RL objectives, restricting to this sparse-reward setting with discount factor
RL algorithms. The above setting
The Advantage Actor-Critic (A2C) algorithm (Degris et al. 2012) combines policy-based and value-based methods to improve training stability and elficiency. Two networks are learned jointly: a policy network Pj. which plays the role of the actor, and state-action value (or critic) network e to measure the 'goodness' of actions taken by the actor. The two objects are learned using the following update:
0 - --0-200 V log (f sx) An (r(s,, a,) + Q (S1+1. a++1) (w,, (s, where do and Ay are learning rates. To further stabilize training, baseline usually used, and Q(s,, a,) is replaced with @(s,. - B [Q(s,, a,)] reducing the variance of the gradient with respect w while keeping it unbiased.
The Soft Actor-Criti (SAC) algorithm (Haarnoja et al. 2018) is an off- off-policy algorithm that is trained maximize trade-off between the expected return and the entropy of the policy Pr, the
4

---

Published as a conference paper at ICLR 2025
latter encouraging more exploration. Using the notation introduced above, SAC's objective is:
arg max By-PE Pr
(3)
1=0
where a is an entropy regularization parameter. To satisfy this objective, we train a state-value function 2w and policy network Pe to satisfy:
Qu (s, a,) r(s;, + (8ot. (So..))) or log P (.1ss.i)) (Bellman equation)
exp (x
arg min Det. Pf
st) st))l
(q can be any off-policy distribution).
(s))
DISCOVERING ACTION ABSTRACTIONS FOR AMORTIZED SAMPLERS
Training amortized samplers in discrete compositional spaces becomes increasingly difficult as the number of steps in trajectories grows. Consider sampling graph one node time while connect- ing them to previous nodes. For example, a graph with и takes ( sampling steps. Sampling long strings autoregressively presents similar challenge: sampling one character at time makes the process difficult due to the vast search space of full sequences, which are predominantly low- likelihood samples. In the case of large language models, sampling subsequences (tokens) which commonly occur in the text corpora instead of individual characters reduces the trajectory length re- quired to encode high-likelihood sentences. This has generally been found to decrease both training and inference costs (Wul 2016 Kudo 2018), and can be viewed a strategy for trading off reducing the depth of the MDP in exchange for increased breadth, which potentially makes the task easier to learn (Pignatelli et al. 2024).
Drawing inspiration from this observation, we propose a general approach for sampling discrete objects that infers actions that effectively shorten the trajectory length while producing samples from the target distribution. In particular, with an action abstraction discovery step, a form of "chunking" as originally described by cognitive psychologists (Miller 1956), which is applicable any sampler, and consists of three major steps (see Algorithm|] for the overall approach):
approach consists of augmenting existing samplers
Generating action corpus: We first generate set of N action sequences from the sampler, optionally also drawing portion of these from a replay buffer of previously-drawn high- reward trajectories.
2. Chunking: We apply a tokenization algorithm common in NLP, byte pair encoding (Bge: 1994). the N action sequences obtain tokens ("chunks"" be added the action space. Augmenting the action space: Finally, we add the new abstracted actions to the action space. Whenever the abstracted action is chosen, its constituent actions executed in orde!")
Algorithm Training policies with chunking
1:
Initialize я (action space), B (replay buffer), M (maximum number of iterations), (pro- portion sample from the replay buffer if it exists) and N (number of trajectories), and k, (chunking frequency).
2: for 21, M) do
3: loss = tr, B,
Train the amortized sampler
4: if n%ok = 0 then
Sample LpNJ from to get action sequences LeN1
5:
Sample N LpN] from the sampler to get action sequences {af ( N-LeN]
6:
i=l
7:
S {al) i=1
N-INN U (at !
8:
C chunking(S)
Generate novel chunks
9:
A AUC
Update the action space
10: end if
end for
'This differs from the options framework where the agent is allowed only partially execute the option. If abstracted action representing sequence of actions is not possible execute (i.e., would lead to the application some of action not in A,), the abstracted action masked by the policy.
5

---

Published as a conference paper at ICLR 2025
This approach is similar to the offline method proposed by Zheng et al. which learns a static set of macro-actions from a dataset the continuous control setting. In contrast, our work introduces a method where the library of macro-actions is constructed online and updated iteratively within the context of amortized sampling.
Chunking mechanisms. We consider the following two approaches, comparing them the ATOMIC sampler (without action space compression):
We apply the tokenizer on the action corpus and add the most fre- quent token found to the action space, which grows by one element each time this is performed. This approach gradually builds up the action space but is susceptible to repetitions in chunks or low quality set of action sequences used for checking.
ACTIINNIEEE-INCREMEEEE
ACTIONPIECE-REPPP
This method instead takes the M most frequent tokens found by the tokenizer from the action corpus (to be appended to the initial action space). Each time abstraction is performed, the entire action space might be replaced by a new one (aside from the atomic to- kens). This effectively reduces the chance of storing redundant chunks and allows us to potentially discover many useful chunks in one step.
Chunking frequency. ACTIONPIECE require choosing "when to chunk" during training. A simple approach (used in most experiments) to simply chunk every k iterations. The alternative is to chunk based on some criterion being satisfied. In the case of GFlowNets, the TB loss can be used as a proxy for the "fit" of the sampler to the current action space (see Appendix B). and we therefore evaluated waiting for the loss to drop below certain threshold before a round of chunking.
5 EXPERIMENTAL SETUP
We benchmark SAC, A2C and GelowNet along with random sampler baseline that samples uni- formly at random a valid action at each step. For the GFlowNet, we use trajectory balance (Malkin
al. 2022 and include three different choices for the backward policy:
Uniform Ps: This backward policy chooses backward parents uniformly at random.
MaxEnt Ps: This backward policy chooses backward trajectories uniformly at random. Note that
MaxEnt P3 and Uniform PB are equivalent for Tree MDPs since there is only a single backward action at each step (Tiapkin al. 2023 Mohammadpour et al.) 2024)..
PhortParse pR: We introduce a new fixed backward policy aimed at sampling the most compact backward trajectory in terms of the number of trajectory steps (see suusecttionnn B.6).
A policy parametrization for nonstationary action spaces. Dealing with varying action space as a result of chunking requires a compatible parameterization of the policy to handle the evolving action space. Instead of a final layer that outputs logits for fixed number of actions, the policy
instead takes as input the current state and outputs an action embedding qt E Rd where d is the
!:91 where |.9] is the number of actions in the action space Я. We use an action encoder fa(a;) Rd that computes an embedding for the actions a;- We parametrize fo as an LSTM (Hochreiter & Schmidhuber 1997) that takes the sequence of atomic actions making up a chunk as input (possibly of length if the chunk is a single atomic action). Let (*A E R|A| denote the matrix of embeddings of all actions in the action space. The logits given by the policy network are:
action embedding dimension (similar to Chandak et al. 2019). Now let A = {at)]
A7) ERIA
E
(4)
vd
Tasks. To understand the impact of chunking, we consider diverse set of environments consisting of three synthetic tasks and practical task: bit sequence generation (Malkin et al.) (022). Fractal- Grid (adapted from the standard hypergrid; [Bengio al. (022). a graph generation environment,
and the practical task of RNA sequence generation (I.14 4 RNA Sinai et al., 2020). We provide
further details about each environment in Appendix B
RESULTS
6.1 DOES CHUNKING IMPROVE GENERATIVE MODELING?
Does chunking accelerate mode discovery? Figure 2 shows the number of modes discovered throughout training as a function of the number of visited states. Note that report the cumula- tive number of modes, i.e., counting all modes sampled for given number of visited states, and the difficulty of the task (in terms of the size of the MDP) increases from left right. Across all

---

Published as a conference paper at ICLR 2025
PractelCrid: 25y2
RVA Rinding: RNA1
Sequenca- sequenon lengt h 128
Sumbwr of nonion
20
10
ABC
1.0 Number
2.0
0.0
Number 1.0
Number
le6
(a) FractalGrid
(b) RNA Binding
(c) Bit sequence
Cumulative number of modes discovered during training. Chunking helps across all cnviron- ments, especially in FractalGrid where all samplers get stuck in the first mode but chunking unlocks exploratory abilitics to fetch faraway modes. Shaded area represents the standard deviation across 6 seeds.
Figure 2:
tasks, ACTIONPIECE outperforms the ATOMIC counterpart for the GelowNet sampler in terms of the speed of mode discovery. Results for other algorithms are mixed. ACTIONPIECE helps A2C and SAC in FractalGrid, but hurts performance for RNA binding. Moreover ACTIONPIECE- REPLACE helps A2C in BitSequence, while ACTIONPIECE-TNE INNREEEENT hurts performance. The greedy be- havior of both SAC and A2C provides possible explanation: when mode is discovered, it gets sampled more often resulting in the addition of a large chunk (e.g.. Figure 10[2). This chunk will result in high-reward samples and consequently chunks added beyond this will contain parts of this chunk hurting exploration. This is further supported by the the Option-Critte performance. This issue is mitigated in GFlowNets due to the exploratory behavior imparted by the learning objec- tive. The random sampler with ACTIONPIeCE achieves strong performance only on simple tasks, with a degradation in performance for harder problems, suggesting random sampling with action abstraction produces meaningful chunks only in cases where the problem is sufficienlly simple.
Does chunking improve density estimation? Next, we study the impact of chunking on den- sity estimation in the context of GFlowNets, which learn a policy to sample from a tar- get distribution. Figure 3a shows the tra- jectory balance loss for all chunking mech- anisms over the 2D grid in the FractalGrid environment. This serves as a visualization for how well each state is correctly learned by the GFlowNet. Chunking, and particularly ACTIONPIECE-INCRRMNE
Table I:
Comparison of chunking mechanisms based on various metries for the graph environment with nodes. We use three seeds with the standard deviation indicated in parentheses. We sample 10,000 graphs for computing the ELBO Gap and draw 40 backward tra- jectories for each graph to compute the JSD and LI dis- tance.
Mechanism
ELBO Gap
JSD
LI Distance
ATOMIC
0.72a039
0.38
1.44
ACTION INCREMENT
0.25
0.36
1.39
ACTIONPIECE CEEECEE
0.35amu1
0.40
positively rein-
forces the inherent exploration of GFlowNets (as seen in the previous section), helping enable the sampler to model the entire state space and sample accurately from the target distribution. This further supported by the LI distance, JSD, and ELBO Gap (see |subsection B.7) during training illustrated in Figure 3b for both grid size 652 and 129? and all chunking mechanisms. Across all sizes of the task, chunking accelerates the training of the sampler. FFiure 3c and [3]] show Spear- man's 's rank correlation coefficient between the target and the learned distribution for the L1 __NNA1 and Bit Sequence tasks respectively. The correlation is computed for a range of reward thresholds spanning all samples to only those with a reward larger than 0.93 to evaluate how well the sampler captures the high reward regions. For I.1 4_RNA1, ACTIONPIECR- CETTONPPECE--INCREEENN performs par with ATOMIC, and even slightly better for high-reward objects, whereas ACTIONPIECE-REDD lags behind. In the bit sequence task, however, chunking results in marginally worse performance. improves density estimation in the graph
Finally, Table I] shows that ACTIONPIECE-INCREEN environment, while ACTIONPIECE- REPLACE is also competitive. With some exceptions, these sults point towards chunking improving density estimation.
6.2 Do CHUNKS CAPTURE THE STRUCTURE IN THE UNDERLYING DISTRIBUTION?
In this section, we analyze the discovered chunks and study whether they capture some structure in the underlying distribution. For instance, RNA molecules are typically composed of building blocks (or chunks in our terminology) called codons that consist of sequences of three nucleotides. We use the RNA Binding (L1 4 4_RNA1) task for all the analyses in this section as we have access a dataset of high-reward objects.
Do chunks represent latent structure of the distribution? Here, we dive into the structural relationship between chunks and the objects generated using the chunks. We consider a Chunk Occurrence metric where we compute the average number of times a chunk occurs in objects in the dataset, and a Chunk Coverage metric where we compute the number of objects in the dataset that contain chunk, normalized over the total number of objects.
7

---

Published as a conference paper at ICLR 2025
Atomic
ActionPiece
Innaaleeeeee-eeenttttt Increment
ActionPiece PettonPlece-Replace
2
(a) TB loss for all chunking mechanisms in FractalGrid for grid size of 2572. All the losses clipped to maximum value of 10 for ease of visualization.
100
657
10-
Oip
10
Sloer 129
10nd
0.0
0.5
3.0
7.5
2.0
0.0
0.5
1.3
1.5
0.0
0.5
1.0
1.5
lantrvage
Number Mitaaeed stades 02.0%
statns (X00)
Number visbec states (220))
(b) The evolution of sampling metrics for all chunking mechanisms during training. The top is for Frac- talGrid of size 652 whereas the bottom row corresponds to that of size 129?. On the right, we plot the reward landscape for grid of size 652 for ease visualization.
Spearman Correiation
Spearman Cocrelation
1.0
0.9
0.9
0.8
0.8
0.7
0.7
0.6
Atomic
0.6
Atomic
ActionFiess-hcrreeeer
AddimPeoo-lnceemeee
0.5
Actionitace- Replace
0.5
Actionitace- Heplace
0.4 Reward
Reward Threshokd
(c) Spearman correlation between logreward and the learned log likelihood for the L14 RMA1 task.
(d) Spearman correlation hetween the logreward and the learned log likelihood for the bit sequence task.
Figure 3: Effect of chunking density estimation. Analysis of the effect of different chunking mechanisms density estimation in the FractalGrid. bit sequence, and RNA binding tasks.
Table 2: Chunk Occurrence and Coverage for different samplers and chunking mechanisms for the L14 RNA1 environment across seeds.
Chunk Occurrence
Chunk Coverage
Sampler
Mechanism
Mean
Median
SD
Mean
Median
GFlow/Net
ACTIONPIECE- INCREMENT
PINDNNPPECE-ICREEENT
1.11
1.00
0.98
0.69
0.73
0.25
ACTIONPIECE ECTIONPIECERREPLACE
0.28
0.00
0.55
0.24
0.10
0.28
ACTIONPIECE: INCREMENT
1.11
1.00
0.98
0.69
0.73
0.25
MaxEnt-GFlowNet
ACTIONPIECE-REIT REPLACE
0.26
0.00
0.47
0.24
0.18
0.23
ACTIONPIECE- INNREMEETT
1.10
1.00
0.98
0.68
0.73
0.26
CaonPPare-GGFloWNe
ACTIONPIECE
0.34
0.00
0..7
0.27
ACTIONPIECE: INCREMENT
0.36
0.00
0.23
0.35
A2C
CERIONPIECE-RPLACE
0.11
0.00
0.35
0.10
0.00
0.24
ACCIOOPPEEE--ICCRMENT
INCREMENT
0.43
0.00
0.78
0.29
0.02
0.36
SAC
ACTIONPIECE-REE REPPLEEEEEEREEEEEEEE
0.15
0.00
0.51
0.10
0.00
0.22
ACTIONPIECE PETIONNPIEE-INNRRMENT
INCREMENT
0.46
0.00
0.79
0.31
0.12
0.36
Random Sampler
ACTIONPIECE REPLACE
0.36
0.00
0.27
0.34
The dataset we consider includes all RNA sequences with a reward of 0.85 or higher. From Ta-
ble 2, it is clear that OCTIINPIECE-INCREEEE
values for GelowWet-based samplers (GFlowNet, consistently results in higher chunk occurrence and ShortParse-GFlow In contrast, A2C, SAC, and the random sampler exhibit moderate to low chunk occurrences un- der OETIONPIECE-INCREEII
while OCTIONPIFCE-REPI.CC yields near-zero median values
8

---

Published as a conference paper at ICLR 2025
ActionPaece Increment
Action?Pece
Avaraag Leopit
BPE the high roward sumples
GTN
Mac -
ShortPPmee-GF
ASC
SAC
Bandom Stampler
(a) Shortest parse of modes from L.4 IRAA reward distribution using 1,1 4.44 A.RNA1 learned chunks.
Avarage hare
30-
ActionPiece Reglace
ActioniPece- Increment
BIS the high roward sample
Slampler
(b) Shortest parse of modes from L14 RMN2 reward distribution using L1 RNA1 learned chunks.
Figure 4: Shortest parse modes using learned library. These plots show the average length of the shortest parses for high-reward samples across different models MahorrParse-GFlowWM MaxEnt-GFlowNee GFlowNet, SAC, A2C, and Random Sampler) when employing different chunking strategics, across 6 seeds (plotted sepa- rately)
across all samplers. A similar trend is observed when analyzing chunk coverage, where GelowNett based samplers again dominate in terms learned chunk representation within the high-reward landscape. Moreover, speculate because the ACTIONPIECE REPLACE REPLLEE strategy updates the library too rapidly, making hard for the policy to model it accurately. These findings suggest that GFlowNet-base samplers augmented with the PEETIINNPEEE-NNCEMMEE effectively capture latent structure of the distribu- tion, particularly around the modes.
CETTONPIECE-INCREMEE
also outperforms ACTIONPIECE-REPLA
Do chunks reduce description length/complexity of samples? To evaluate the compressive abil- ity of the chunks, we compute the shortest parse of the (known) modes of the distribution with the learned library of chunks. This defined as the average length of a trajectory required to sample mode. Note that the the longest parse of modes would be a trajectory using only atomic actions (letters string- environments). We also compute a minimum attainable shortest parse of the modes by performing BPE on the modes directly, which serves as lower bound.
Figure 4 shows that GFlowNet-based. approaches tend to have the lowest short parses of the modes, indicating greater diversity in the learned chunks. In contrast, RL-based methods tend to have higher average parse lengths, as their learned chunks focus on smaller region of the state space, suggesting that limited exploration chunking mechanism also affects different samplers in opposite ways. For RL-based methods, ACTIONPIECE- REPLACE learns library that doesn't 't compress samples very well as evident by the average length the shortest-parse as opposed to GFN-based methods that do indeed find libraries that capture the latent structure of the distribution and result in shorter average length compared to ACCTINPIECE-INCREMME
negatively impact the quality of the learned chunks. The
Are chunks transferable to new models and new tasks? In this section, we study the down- stream impact discovered chunks. We evaluate the downstream performance along two axes: generalization across different but related target distributions, and across different samplers. The aim this experiment to identify whether the chunks themselves generalize and which samplers provide the best chunks in terms of downstream performance.
Figure 5 shows that initializing samplers with learned library of chunks improves the number of modes found compared not performing any chunking. We see this specifically when samplers use learned library of chunks from L1 4_RNA1 to sample from 1.14_NAA and L14_RNA3, show- ing that the chunks do generalize structurally similar reward distributions. Zooming in the performances of each sampler, we see that GFlowNet-induuced libraries provide on-average better results in terms of mode-discovery than the libraries coming from other samplers. This can be

---

Published as a conference paper at ICLR 2025
CPN
16.0
05:30
27
31
QFN
40
00
GFN
1:67
10
60.29
06.00
50
SAC
8.33
9.33 11. 633 3,67
6.0
RAC
5.9
5.0
1000
200
1.33
4.0
SAC
2.67
2.67
0.0
0.0
0.0
0.67
40
Randon
26.0
13.38
10.33
Ranlon
17:07
24.33
10.33
0.67
Randon
3.0
0.33
3.33
OPF
Maxtro GPF
Maxian dPP
Sbartaammn 179
12PM
Kandon
Slort
Actindttos Hapiaot
GFN
2.33
26.33
GEN-
47:07
50.57
A2C
10.0
8.33
2
400
7,33
13:43
A2C
9.9
10.9
0.0
10
3.67
15:00
A2C
23 29.0
30.0
67
16.67
4.33
27.6
SAC
4.33
8.0
4,87
433 0.33 7:33
840-
1.0
1.87 4.23
0.57
5:33
8AC-
1.87 1.39 2.5 0:33
0.67
0.0
10
10.67
0.33
11:67
7.67
8:67
1.39
Bandom
0.9
2.33
CPW
Matha-oP
Shortremm- Ops
Nendom
OPW
Madia
Stontran--rp OPS
Sac
Nasium
CPW
Nodill 4PN
Stonthame-tr
Kasdom
Figure 5: Donwstream evaluation of discovered chunks. Each column represents a different environ- ment (1.14 RNA1, 1.1 RNA2, 1.1 RNA3) and each row presents a chunking mechanism (ACTIONPIeCE- In each heatmap, we show the number of modes discovered by
INCREMENT and OCTIONPIECE-REPLACC samplers on the y-axis trained on chunks found by samplers the x-axis. The color intensity represents the number of modes, with darker shades indicating higher numbers. Average of three random seeds.
plained by the fact that RL-based methods tend learn chunks specific prioritized subspace of the state-space, which fail as a result to generalize to the entire space. While these remarks are true for libraries learned using ACTIONPIECE- INCREMENT, the conclusions do not transfer well to the ACTIONPIECE PIETE-OPECE--ELLAAAC mechanism which achieves poor performance. We speculate that this mechanism is too abrupt when updating the library, since it updates all chunks at once, whereas ACTIONPIECE- INCREMENT adopts a more curriculum-lie strategy of adding only one chunk at a time, allowing the samplers enough time to adapt to the newly added chunk.
7 CONCLUSION
Abstraction and dynamic concept learning are critical in human cognition and should be modeled in artificial learning systems. In this paper, we investigated dynamic action space learning in the context of amortized samplers. We proposed ACTIONPIECE, an extensible approach which lever- ages tokenizers to chunk action sequences, which can be viewed as strategy for trading reduced depth of the MDP in exchange for increased breadth (Pignatelli et al. 2024). We demonstrated its effectiveness in improving the performance of samplers on variety of domains with respect to mode discovery and capturing the latent structure of the environment. Our empirical results also indicate that the selected chunks for samplers like GelowNets capture underlying structure in the target distribution and can generalize other target distributions. This approach has multiple pos- sible utilities: the learned chunks capture the underlying structure of the reward distribution (and therefore potentially interpretable), can facilitate transfer more complex problems where in- dividual trajectories would grow too long if constructed only of atomic actions, and under the right conditions, appears to facilitate exploration itself.
ACTIONPIECE builds on a long line of research into macro-actions (e.g.. Durugkar et al. (016);; Iba (1989): Chang al. (2022): Dulac et al.] (2013); Jong et al. (2008)) but previous results have shown they have mixed utility in general. They have not until now been evaluated in the context of diversity-seeking RL approaches such as GelowNets. Our results suggest that the use of online library construction works best in conjunction with methods that explicitly maximize diversity: they leam libraries that allow for the shorter parse lengths and more useful for discovering new modes when transferred to new tasks. This points towards an intriguing interaction between the trajectory sampling method and the ability to learn robust, general purpose macro-actions. Our results suggest that future studies on action abstraction for amortized samplers should explore the interaction between the method for macro-action construction and exploration.
Limitations and future work. In this work, we perform chunking with the BPE tokenizer for fixed target distribution. Future work can study learning task-conditioned libraries, as well as formulating the problem of learning abstractions as modelling the joint distribution P(X.L), where f denotes library of abstractions that is, modeling a distribution over possible concept libraries. Future work can also explore the application of the ideas of abstractions discussed in the paper in the context of code generation, where the right abstractions can naturally make complex tasks easier (Ellis et al. 2021 Stengel-Fskin et al. 2024).
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENTS
The authors acknowledge funding from CIFAR, NSERC, IVADO, NRC and Samsung. The search was enabled by computational resources provided by the Digital Research Alliance of Canada htttps://allaanees
quebec). and NVIDIA. ca). Mila (https: /mi lt..beee.......e..c/emmm/
REPRODUCIBILITY STATEMENT
We discuss all the details and hyperparameters used to reproduce the results in the paper in Ap- pendix B We also include the code to reproduce the experiments along with the submission:
REFERENCES
Jose A Arjona-Meding Medina, Michael Gillhofer, Michael Widrich, Thomas Unterthiner, Johannes Brand- stetter, and Sepp Hochreiter. Rudder: Return decomposition for delayed rewards. Neural Infor- mation Processing Systems (NeurIPS), 2019.
Emmanuel Bengio, Moksh Jain, Maksym Korablyov, Doina Precup, and Yoshua Bengio. Flow net- work based generative models for non-iterative diverse candidate generation. Neural Information Processing Systems (NeurIPS), 2021.
Yoshua Bengio, Salem Lahlou, Tristan Deleu, Edward J Hu, Mo Tiwari, and Emmanuel Bengio. GFlowNet foundations. Journal of Machine Learning Research, 24(210):1-55, 2023.
Matthew Bowers, Theo X Olausson, Lionel Wong, Gabriel Grand, Joshua Tenenbaum, Kevin Ellis, and Armando Solar-Lezama. Top-down synthesis for library learning. Proceedings of the ACM on Programming Languages, 7(POPL):1 122-1 2023.
Shaked Brody, Uri Alon, and Eran Yahav. How attentive are graph attention networks? International Conference Learning Representations (ICLR), 2021.
Yash Chandak, Georgios Theocharous, James Kostas, Scott Jordan, and Philip Thomas. Learning action representations for reinforcement learning. International Conference Machine Learning (ICML), 2019.
Yi-Hsiang Chang, Kuan- Yu Chang, Henry Kuo, and Chun-Yi Yi Lee. Reusability and transferability of macro actions for reinforcement learning. ACM Transactions Evolutionary Learning and Optimization, 2(1):1-16, 2022.
Swarat Chaudhuri, Kevin Ellis, Oleksandr Polozov, Rishabh Singh, Armando Solar-Lezama Yisong Yue, al. Neurosymbolic programming. Foundations and Trends@ in Programming Languages, 7(3):158-22) 243, 2021.
Glanois Claire, Shyam Sudhakaran, Elias Najarro, and Sebastian Risi. Open-ended library learning in unsupervised program synthesis. In ALIFE 2023: Ghost in the Machine: Proceedings of the 2023 Artificial Life Conference. MIT Press, 2023.
Peter Dayan and Geoffrey Hinton. Feudal reinforcement learning. Neural Information Processing Systems (NIPS), 1992.
Thomas Degris, Patrick M Pilarski, and Richard S Sutton. Model-free reinforcement learning with continuous action in practice. In 2012 American control conference (ACC), pp. 2177-2182. IEEE, 2012.
Tristan Deleu, Padideh Nouri, Nikolay Malkin, Doina Precup, and Yoshua Bengio. Discrete prob- abilistic inference as control in multi-path environments. Uncertainty in Artificial Intelligence (UAI), 2024.
Jacob Devlin, Rudy Bunel, Rishabh Singh, Matthew Hausknecht, and Pushmeet Kohli. Neural program meta-induccton. Neural Information Processing Systems (NIPS), 2017.
Adrien Dulac, Damien Pellier, Humbert Fiorino, and David Janiszek. Learning useful macro-actions for planning with n-grams. In 2013 IEEE 25th International Conference on Tools with Artificial Intelligence, pp. 803-810. IEFE, 2013.

---

Published as a conference paper at ICLR 2025
Ishan P Durugkar, Clemens Rosenbaum, Stefan Dembbach, and Sridhar Mahadevan. Deep reinforce- ment learning with macro-actions. arXiv preprini arXiv: 1606,04615, 2016.
Kevin Ellis, Lucas Morales, Mathias Sablé-Meyer, Armando Solar-Lezama, and Josh Tenenbaum. Learning libraries of subroutines for neurally-guides bayesian program induction. Neural Infor- mation Processing Systems (NIPS), 2018.
Kevin Ellis, Maxwell Nye, Yewen Pu, Felix Sosa, Josh Tenenbaum, and Armando Solar-Lezama. Write, execute, assess: Program synthesis with a REPL. Neural Information Processing Systems (NeurIPS), 2019.
Kevin Ellis, Catherine Wong, Maxwell Nye, Mathias Sablé-Meyer, Lucas Morales, Luke Hewitt, Luc Cary, Armando Solar-Lezama, and Joshua B Tenenbaum. Dreamcoder: Bootstrapping in- ductive program synthesis with wake-sleep library learning. In Proceedings of the 42nd acm sigplan international conference programming language design and implementation, pp. 835- 850, 2021.
Philip Gage. A new algorithm for data compression. C Users J., 12(2):23-38, feb 1994. ISSN 0898-9788.
Fernand Gobet, Peter CR Lane, Steve Croker, Peter CH Cheng, Gary Jones, lain Oliver, and Julian M Pine. Chunking mechanisms in human learning. Trends in cognitive sciences, 5(6):236-245) 243,, 2001.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. International Conference on Machine Learning (ICML), 2018.
Anna Harutyunyan, Will Dabney, Thomas Mesnard, Mohammad Gheshlaghi Azar, Bilal Piot, Nico- las Heess, Hado P van Hasselt, Gregory Wayne, Satinder Singh, Doina Precup, et al. Hindsight credit assignment. Neural Information Processing Systems (NeurIPS), 2019.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Comput., 98)): 1735-1780, nov 1997. ISSN 0899-7667. doi:
URL https:
[/doi.org//il 10. 11w./neco...ncoooooooooooooooo
neco. 1997. 7.9.7...... 8713
Chia-Chun Hung Timothy Lillicrap, Josh Abramson, Yan Wu, Mehdi Mirza, Federico Carnevale, Arun Ahuja, and Greg Wayne. Optimizing agent behavior over long time scales by transporting value. Nature communications, 10(1):5223, 2019.
Glenn A. Iba. A heuristic approach to the discovery of macro-operators. Mach. Learn., 3(4): 285-317, mar 1989. ISSN 0885-6125. doi: 10.1023/41022269377366. URL https: (adoi.org/In 1023/1: wtt.:////...22211111011
Neal Johnson. The role chunking and organization the process of recall. In Psychology of learning and motivation, volume 4, pp. 171-247. Elsevier, 1970.
Nicholas K Jong, Todd Hester, and Peter Stone. The utility of temporal abstraction in reinforcement learning. In AAMAS (1), pp. 299-306, 2008
Nan Rosemary Ke, Anirudh Goyal, Olexa Bilaniuk, Jonathan Binas, Michael C Mozer, Chris Pal, and Yoshua Bengio. Sparse attentive backtracking: Temporal credit assignment through remind- ing. Neural Information Processing Systems (NIPS), 2018
Kudo. SentencePicce: A simple and language independent subword tokenizer and detokenizer for neural text processing. arXiv preprint arXiv:1808.062. 2018.
Salem Lahlou, Tristan Deleu, Pablo Lemos, Dinghuai Zhang, Alexandra Volokhova, Alex Hernández- Gareia. Léna Néhale Ezzine, Yoshua Bengio, and Nikolay Malkin. A theory of con- tinuous generative flow networks. International Conference on Machine Learning (ICML), 2023.
John Laird, Paul S Rosenbloom, and Allen Newell. Chunking in soar: The anatomy of a general learning mechanism. Machine learning, 1:11 1986.
12

---

Published as a conference paper at ICLR 2025
Perey Liang, Michael Jordan, and Dan Klein. Learning programs: A hierarchical Bayesian ap- proach. International Conference on Machine Learning (ICML), 2010.
Yang Liu, Yunan Luo, Yuanyi Zhong. Xi Chen, Qiang Liu, and Jian Peng. Sequence modeling of temporal credit assignment for episodic reinforcement learning, arXiv preprint arXiv: 1a0e::995...444 13420, 2019.
Kanika Madan, Jarrid Rector- Brooos, Maksym Korablyov, Emmanuel Bengio, Moksh Jain, Andrei Nica, Tom Bosc, Yoshua Bengio, and Nikolay Malkin. Learning GFlowNets from partial episodes for improved convergence and stability, International Conference Machine Learning (ICML), 2023.
Nikolay Malkin, Moksh Jain, Emmanuel Bengio, Chen Sun, and Yoshua Bengio. Trajectory balance: Improved credit assignment in gflownets. Neural Information Processing Systems (NeurIPS), 2022.
Amy McGovern and Richard S Sutton. Macro-actions in reinforcement learning: An empirical analysis. Computer Science Department Faculty Publication Series, pp. 15, 1998.
Amy McGovemn Richard S Sutton, and Andrew Fagg. Roles of macro-actions accelerating reinforcement learning. In Gruce Hopper celebration of women in computing, volume 1317, pp- 15, 1997.
George A Miller The magical number seven, plus or minus two: Some limits on our capacity for processing information. Psychological review, 63(2):81, 1956.
Sobhan Mohammadpour, Emmanuel Bengio, Emma Frejinger, and Pierre-Lue Bacon. Maximum entropy GFlowNets with soft Q-learning. Artificial Intelligence and Statistics (AISTATS), 2024.
Vinod Nair and Geoffrey E. Hinton. Rectified linear units improve restricted Boltzmann machines. International Conference on Machine Learning (ICML). 2010.
Maxwell Nye, Armando Solar-Lezama, Josh Tenenbaum, and Brenden M Lake. Learning compo- sitional rules via neural program synthesis. Neural Information Processing Systems (NeurIPS), 2020.
Ling Pan, Nikolay Malkin, Dinghuai Zhang, and Yoshua Bengio. Better training of GFlowNets with local credit and incomplete trajectories. International Conference on Machine Learning (ICML), 2023.
Eduardo Pignatelli, Johan Ferret, Matthieu Geist, Thomas Mesnard, Hado van Hasselt, and Laura Toni. A survey of temporal credit assignment in deep reinforcement learning. Transactions Machine Learning Research, 2024.
Doina Precup and Richard S Sutton. Multi-time models for temporally abstract planning. Neural Information Processing Systems (NIPS), 1997.
Aditya A Ramesh, Kenny Young, Louis Kirsch, and Jürgen Schmidhuber. Sequence compression speeds up credit assignment in reinforcement learning. International Conference on Machine Learning (ICML), 2024.
Jette Randlov.
Learning macro-actions reinforcement learning. Neural Information Processing Systems (NIPS), 1998.
Jorma Rissanen. Modeling by shortest data description. Automatica, 14(5):465-47) 1978.
Tankred Saanum, Noémi Éltetō, Peter Dayan, Marcel Binz, and Eric Schulz. Reinforcement learning with simple sequence priors. Neural Information Processing Systems (NeurIPS, 2023.
Jürgen Schmidhuber. Learning complex, extended sequences using the principle of history com- pression. Neural Computation, 4(2):234 242, 1992.
Mike Schuster and Kaisuke Nakajima. Japanese and korean voice search. 2012 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 5149-5152, 2012. doi: 10. D0.1199CCASSPS22222220 URL
leee. org/ document/
6289079
13

---

Published as a conference paper at ICLR 2025
Marcin Sendera, Minsu Kim, Sarthak Mittal, Pablo Lemos, Luca Scimeca, Jarrid Rector-Brooks Alexandre Adam, Yoshua Bengio, and Nikolay Malkin. Improved off off-policy training of diffusion samplers. Neural Information Processing Systems (NeurIPS), 2024.
Rico Sennrich, B. Haddow, and Alexandra Birch. Neural machine translation of rare words with subword units, Association for Computational Linguistics (ACL), 2015.
Sam Sinai, Richard Wang, Alexander Whatley, Stewart Slocum, Elina Locane, and Eric D. Kelsic. AdaLead: A simple and robust adaptive greedy search algorithm for sequence design. arXiv preprint vrXiv:2010.021411 2020.
Elias Srengel-Esking Archiki Prasad, and Mohit Bansal. ReGAL: Refactoring programs to discover generalizable abstractions. International Conference on Machine Learning (ICML), 2024.
Chen Sun, Wannan Yang, Thomas Jiralerspong, Dane Malenfant, Benjamin Alsbury-Nealy, Yoshua Bengio, and Blake Richards. Contrastive retrospection: honing on critical steps for rapid learning and generalization in rl. Neural Information Processing Systems (NeurIPS), 2023.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Richard S Sutton, Doina Precup, and Satinder Singh. Between mdps and semi-mdps: A frame- work for temporal abstraction in reinforcement learning. Artificial intelligence, 112(1-2):181-
211. 1999.
Mirko Thalmann, Alessandra S Souza, and Klaus Oberauer. How does chunking help working memory? Journal of Experimental Psychology: Learning, Memory, and Cognition, 45(1):37, 2019.
Lucas Tian, Kevin Ellis, Marta Kryven, and Josh Tenenbaum. Learning abstract structure for draw- ing by efficient motor program induction. Neural Information Processing Systems (NeurIPS), 2020.
Daniil Tiapkin, Nikita Morozov, Alexey Naumov, and Dmitry Vetrov. Generative flow networks as entropy-regularize RL. Artificial Intelligence and Statistics (AISTATS), 2023.
Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pictro Liò, and Yoshua Bengio. Graph attention networks. International Conference Learning Representations
(ICLR), 2017.
Jiff Vyskočil, David Stanovský, and Josef Urban. Automated proof compression by invention of new definitions. In Edmund M. Clarke and Andrei Voronkov (eds.). Logic for Programming, Artificial Intelligence, and Reasoning, pp. 447-462. Springer, 2010.
Yonghui Wu. Google's neural machine translation system: Bridging the gap between human and 2016.
machine translation. arXiv preprint
Jingjing Xu, Xu Sun, Zhiyuan Zhang, Guangxiang Zhao, and Junyang Lin. Understanding and improving layer normalization. Neural Information Processing Systems (NeurIPS), 2019.
Linting Xue, Aditya Barua, Noah Constant, Rami Al-Rfou, Sharan Narang, Mihir Kale, Adam Roberts, and Colin Raffel. ByT5: Towards a token-free future with pre-trained byte-to-byte models. Transactions of the Association for Computational Linguistics (TACI.), 2021.
Qinsheng Zhang and Yongxin Chen. Path integral sampler: A stochastic control approach for sam- pling. International Conference on Learning Representations (ICLR), 2022.
Ruijie Zheng, Ching-An Cheng. Hal Daumé III. Furong Huang. and Andrey Kolobov. Prise: Llm- style sequence compression for learning temporal action abstractions in control. In Forty-first International Conference Machine Learning.
Jin Peng Zhou, Yuhuai Wu, Qiyang Li, and Roger Grosse. REFACTOR: Learning to extract theo- rems from proofs. International Conference on Learning Representations (ICLR), 2024.
Brian Ziebart, Andrew L Maas, J Andrew Bagnell, Anind K Dey, et al. Maximum entropy inverse reinforcement learning. Association for the Advancement of Artificial Intelligence (AAAI). 2008.
14

---

Published as a conference paper at ICLR 2025
Appendix
Table of Contents
Additional related works
16
B Experimental details
16
Orad
17
RNA Binding
B.3
Bit Sequence
B.4
Graph
18
B.S
Replay Bulfer
19
8.6
ShortParse
20
Sampling quality metrics
21
C
carned libraries of chunks
22
FradtalGrid
................................
22
2
RNA Binding
................................
22
Bit Sequence
................................
22
Graph
................................
24
D
The role of tokenization strategy
24
Sampler effectiveness in sampling high-reward objects
24
15

---

Published as a conference paper at ICLR 2025
A ADDITIONAL RELATED WORKS
Tokenization. Our work builds the principle of tokenization, which is pre-processing step for natural language tasks and a critical piece for language models (LMs). tokenizer produces vocabulary that be used by the L.M in order to correctlv process text in the language of the training datase Byte Pair Encoding (BPE) Gage 1994, Sennrich et al. 2015 starts with an initial vocabulary of tokens and successively merges the most frequent pairs of adjacent tokens get new composite onc. WordPiece (Schuster Nakajima 2012) works similarly to BPE where a pair only merged if the ratio of the likelihood of its joint to that the product of its marginals is the highest. For BPE and WordPiece both, the procedure repeated until a maximum vocabulary size is reached.
B EXPERIMENTAL DETAILS
FractalGrid reward landscape
def create_grid Ri, R2, side_length):
grid RD on.ones (side depth int (np. Log (eighe_ side Lede__enghh)
np. (2)
def fill r2 (x, y, current size, current ourrent_depth) if current_depth =- or eurrent_size < 1: return
Ensure don't go out of bounds max_y nyn(y, side_length max_x x min (x + current_size 1, side_length
$ Fill bottam-left cell
grid [maxy - 1, + 1] RC
RI + R2
grid max y
x] R0
grid [max_y, + 1] R0
grid [max_ y, x] R0
if surrent_deptin < depth:
$ Pill bottom-right cell grid max_y grid[nax_ 1, max_x] RD
max 1] R0
+ R2
grid max y, max 11 RD
24
mnaxy,, max_x] RC R1
25
26
$ Fill top-left cell
gridjy current_size + 2, + 1] RO R1 + R2
gridly current_size 2, x] R0
grid.y current_size 1, 11 R0
gridly current_siza
1, x] R0
31
# Calculate the size and position of the next level next size = current size // 2
next_x next_size next_y y neze_size
37
$ Recursive call for the next level fi (next_x, next_y, sizt_size, current_dapth - 1)
38
# Start from the bottom-left corner side length, depth)
fill (0, side length grid torch. fron_numpy (gridl: copy float ()
return grid
Listing 1: Code for implementing the FractalGrid reward landscape.
froce approaches
Xue
2021 do exist and only make of bytes/characters.
16

---

Published as a conference paper at ICLR 2025
In this section, we provide implementation details along with necessary the hyperparameters for reproducing the experiments presented in the paper. All runs use a single GPU with runs taking up to 36 hours maximum.
B.1 GRACTALGRID
Our implementation is inspired by hypergrid that was introduced in Bengio et al. (002)) In this study, we consider only the 2D case and present results for three settings with increasing grid size: 652, 1292 and 257². The main difference lies in the reward landscape which we have altered. We show the code used to compute the reward landscape (see Code
Environment. The starting state is so (0,0) R² and the atomic action space consists of the following actions: UP, RIGHT and <EXIT>. The sampler can choose to exit the trajectory at any stage. We chose Ro 0.1, R 0.5 and R₂ 2. We run the samplers for a total of 31250 iterations and a batch size of 64, adding up to a total of 2 million visited states during training.
Architecture. The forward policy in SAC, A2C. Option-Critie and GFlowNet is implemented using a feedforward network with 3 layers and a hidden dimension of 128. The last layer produces an action embedding of dimension 128 as well and layer norm applied just before (Xu et al. 2019). For the activation function, we use ReLU (Nair & Hinton 2010). For A2C, Option-Critic and GFlowNet, the forward policy has learning rate of 10 whereas for SAC, it is x 10-4. For GelowNet, we use an initial value for the learnable log-partition value of 90 and a learning rate of 10-3. The critic in A2C and Option-Critic is also parametrized by a feedforward network with 3 layers and a hidden dimension of 128 where a layer norm is applied before the last layer. The critic in A2C and Option- Critic has learning rate of 10 whereas in SAC, it has learning rate of 3 x 10"4. The termination probability network is parametrized using a feedforward network with layers and a hidden dimension of 128 with a learning rate of 10-4.
Samplers. Both GFlowNet and SAC are off- algorithms. For both, we use a prioritized replay buffer with diversity constraints (see section on replay buffer). Our diversity criterion the hamming distance between final states and the threshold is set to 1. The replay buffer capacity is set 1000 due to the small size of the state space. During training. 55% of samples come from the replay buffer and the rest from the sampler. We don't use e-greedy exploration for SAC, but for GelowNet we linearly decay from 0.5 down 0.1. SAC uses entropy coefficient of 0.2 and A2C and Option-Critic use an entropy coefficient of 0.5. For GFlowNet, we use a reward exponent of 1. We use 10 options for the Option-Critie and coefficient of 0.01 for regularizing the termination probability.
Chunking. We perform chunking every 1250 iterations to get 25 chunks at the end of training.
B.2 RNA BINDING
We consider sampling sequences of length 14 and sequences of length 50. The reward is scaled between 10- 10 and 1.
The starting state is the empty string and the atomic action space consists of the following actions: A, C, G, U which append the nucleobases to the current string and <EOS> for End-Of-String. We run the samplers for a total of 25000 iterations and a batch size of 64, adding up to total of 1.6 million visited states during training.
Environment
The forward policy in SAC, A2C, Option-Critic and GFlowNet is implemented us- ing an LSTM (Hochreiter & Schmidhuber 1997) with layers and a hidden dimension of 128 followed by a feedforward network with 2 layers. The last layer produces an action embedding of dimension 128 as and a layer norm is applied just before (Xu et al.. 2019). For the activation function, we use ReLU (Nair & Hinton 2010). For A2C, Option-Critic and GelowNet, the forward policy has a learning rate of 10 whereas for SAC, it's 3 10-4, For GFlowNet, we use an initial value for the learnable log-partition value of 11 for length 14 and 22 for length 50, the learning rate is set to 10-3. The critic in A2C and Option-Crite Critic is also parametrized by a two layer LSTM followed by a two layer feedforward network of hidden dimension of 128 where a layer norm is ap- plied before the last layer. The critic in A2C and Option- Critic has learning rate of 10-4 whereas in SAC, it has a learning rate of x 10-4, The termination probability network is parametrized LSTM (Hochreiter & Schmidhuber 1997) with 2 layers and a hidden dimension of 128 followed by a feedforward network with 2 layers and a learning rate of 10-4.
Architecture
Samplers. Both GelowNet and SAC are off-policy algorithms. For both, we use a prioritized replay buffer with diversity constraints (see section on replay buffer). Our diversity criterion is the
17

---

Published as a conference paper at ICLR 2025
hamming distance between final states and the threshold is set to 3 for length 14 and 10 for length 50. The replay buffer capacity is set to 10000. During training, 55% of samples come from the replay buffer and the rest from the sampler. We don't use e-greedy exploration for SAC, but for GFlowNet we linearly decay e from 0.1 down to 0.01. SAC uses an entropy coefficient of 0.1 for length 14 and 0.025 for length 50. A2C and Option-Critt use an entropy coefficient of 0.05 for length 14 and 0. 005 for length 50. For GFlowNet, we use reward exponent of 10 for length and 75 for length 50. We use 10 options for the Option-Critte and a coefficient of 0.01 for regularizing the termination probability.
Chunking. We perform chunking every 1000 iterations which means that at the end of the training, we have 25 chunks. For GFlowNettbassed samplers we use a different approach where we apply chunking once the TB loss gets below certain threshold. For length 14 the initial loss threshold is set to 1 whereas for length 50 it is set to 5. Once chunking is applied, the loss threshold is multiplied by 0.75.
B.3 BIT SEQUENCE
We modify the original bit-sequence task introduced in Malkin et al. (022 to have different reward function. We first initialize a list of "words": [00000000, 11111111, 11110000,
00001111, 00111100]. The reward function is then designed to favor strings that contain maximum number of these words (including repeated ones). As such, the reward function is:
R(s) number of words in s
(5)
N/8
Where N is the sequence length and 8 is the size of the words. N/8 represents the maximum number of words a string of length N. The numerator is computed using dynamic programming. We consider two settings: Sampling sequences of length 64 and those of length 128.
Environment. The starting state is the empty string and the atomic action space consists of the following actions: O, 1 and <EOS> for End-Of-String. We run the samplers for a total of 25000 iterations and a batch size of 64, adding up to total of 1.6 million visited states during training.
Architecture. The forward policy in SAC, A2C, Option-Critia and GFlowNet is implemented using an LSTM (Hochreiter & Schmichuber 1997) with 2 layers and a hidden dimension of 128 followed by a feedforward network with 2 layers. The last layer produces an action embedding of dimension 128 as and a layer norm is applied just before (Xu et al.] 2019). For the activation function, we use ReLU (Nair & Hinton| 2010). For A2C, Option-Critic and GFlowNet, the forward policy has a learning rate of 10 whereas for SAC, it is 3 x 10-4. For GelowNet, we use an initial value for the learnable part-patiiionn value of 40 for length 64 and 30 for length 128, the learning rate is set to 10-3. The critic in A2C is also parametrized by a two layer LSTM followed by a two layer feedforward network of hidden dimension of 128 where a layer norm is applied before the last layer. The critic in A2C has a learning rate of 10-4 whereas in SAC, it has a learning rate of 3 x 10-4, The termination probability network is parametrized LSTM (Hochreiter & Schmidhuber) 1997) with 2 layers and a hidden dimension of 128 followed by a feedforward network with 2 layers and learning rate of -0-4.
Samplers. Both GFlowNet and SAC are off-policy algorithms. For both, we use a prioritized replay buffer with diversity constraints (see section on replay buffer). Our diversity criterion the hamming distance between final states and the threshold is set to 8 for length 64 and 16 for length 128. The replay buffer capacity set 10000. During training. 55% of samples come from the replay buffer and the rest from the sampler. We don't use e-greedy exploration for SAC, but for GFlowNet we linearly decay E from 0.1 down to 0.01. SAC uses an entropy coefficient of 0.01 for length 64 and 0.005 for length 128. A2C and Option-Critic use an entropy coefficient of 0.05 for length 64 and 0.01 for length 128. For GFlowNet, we use reward exponent of 100 for length 64 and 200 for length 128. We use 10 options for the Option- Critic and a coefficient of 0.01 for regularizing the termination probability.
Chunking We perform chunking every 1000 iterations to get 25 chunks at the end of training.
B.4 GRAPH
In this paper, we departed from the usual way graph environments are set-up. Indeed, for chunking to make sense and essentially work for a graph environment, we need to be careful about the design of the MDP. For a modular structure to emerge in the action sequence, we need the action space to reflect that. In previous work, actions for adding edges would be to select any two nodes and connect them. However, connecting node 3 and 4 although similar in structure to connecting node 5
18

---

Published as a conference paper at ICLR 2025
and 6, would be seen as different for the tokenizer. Let N denote the maximum allowed number of nodes in the graph, we build the action space and its constraints as follows:
1. For an empty graph:
Only the ADD-NODE action is permitted.
2. For non-empty graph:
(a) If the graph contains exactly one node:
Choose between <EOG> (end of graph generation) or ADD-NODE
(b) If the graph contains two or more nodes:
i. When the last added node is connected:
Choose from <EOG>, ADD-NODE, or connect the last node to a previous node.
Edge addition actions are formatted as ADD-EDGE-1 where E (-1, .... -(k
j j+1 1)).
Here, k is the order of the last added node, and j is the order of the farthest connected node the last one.
ii. When the last added node is not connected:
You must connect it to one of the previous nodes. Allowed actions are ADD-EDGE- where i E (-1, -2... -k} where k is the order of the last node.
This structure ensures that edge additions always relative the last added node, creating consistent and potentially interpretable action space. For instance, ADD-EDGE- (-1) always means connecting the last node to the previous one, regardless of their absolute node numbers. We speculate that an objective of future work could be to explicitly produce chunks with interpretable properties, for example, using predefined rules which favors the retention of particular chunks, or using scores from human feedback. By defining actions in this relative manner, we create an action space that is more amenable to finding recurring structure in it. The reward function is defined as the total number of cycles divided by the maximum number of cycles that can be found in a graph with N nodes where N is the maximum number of nodes in the graph:
R(s) number of cycles in graph
(6)
N(N-1) N+1
Environment. The starting state is the empty graph and the atomic action space consists of the following actions: ADD-NOOS ADD- BOD-EDGE-B k where an edge is added between the last node and the k-th node before it. k takes values in the set (-1,-2..... -N -N + 1} where N is the maximum
number of nodes. The last action is <EOG> for End-Of- Of-Grafh........ We run the samplers for a total of 25000 iterations and batch size of 64, adding up to a total of 1.6 million visited states during training.
Architecture. The forward policy in SAC, A2C and GFlowNet has two Graph Attention Layers Veličković et al., 2017] [Brody et al. 2021) and a hidden dimension of 128 followed by feedfor- ward network with layers. The last layer produces a laver norm applied just before (Xu et al.) 2019). For the activation function, we use ReLU (Nair & Hinton 2010). For A2C and GFlowNet, the forward policy has a learning rate of 10-4 whereas for SAC, it 3 x 10-4, For GFlowNet, we use an initial value for the learnable log- dog-partition value of 10 for N = 7 and 30 for length N = 10, the learning rate is set to 10-3, The critic in A2C is also parametrized by a two layer LSTM followed by a two layer feedforward network of hidden dimen- sion of 128 where a layer norm is applied before the last layer. The critic in A2C has learning rate of 10-4 whereas SAC, it has a learning rate of x 10-4."
action embedding of dimension 128 and
Samplers. Both GelowNet and SAC are off-policy algorithms. For both, we use prioritized replay buffer with diversity constraints (see section on replay buffer). Our diversity criterion is the hamming distance between the graphs adjacency matrices and the threshold is set to 3 for N 7 and 5 for N 10. The replay buffer capacity is set to 10000. During training, 55% of samples come from the replay buffer and the rest from the sampler. We don't use e-greedy exploration for SAC, but for GFlowNet we linearly decay e from 0. down to 0.01. SAC uses an entropy coefficient of 0.1. A2C uses an entropy coefficient of 0.05. For GFlowNet, we use a reward exponent of 1.
Chunking. We perform chunking every 1000 iterations to get 25 chunks at the end of training.
B.5 REPLAY BUFFER
In this section, we provide additional details about the replay buffer used throughout the paper (see Algorithm 2). Our replay buffer maintains balance between having high-reward samples and
19

---

Published as a conference paper at ICLR 2025
diverse ones. When new batch of elements is added, only keep the ones that have reward higher than the minimal reward in the buffer. This ensures that the minimum reward the buffer never decreases, meaning that we don't hinder the quality of the samples in the buffer.
Once the high-reward states of the batch are kept, we iterate through each state s:
a distance metric. Given cutoff distance de that controls how much diversity want, we only add to the buffer B if d(s, B) de. the state is not diverse enough, we look for the most similar state 85 = arg miny dss,s') to it in the buffer and if state has higher reward than Sb, we replace by s. Finally, we sort the buffer by the reward and truncate to capacity.
Let d(s, B) miny: d(s,s") where
Algorithm 2 Adding trajectories to the Prioritized Replay Buffer
Require:
New experiences E, buffer B, cutoff distance de if B| capacity then
2: Add E to B
Sort B by log-reward
4: else
5:
e'
min(rb),
6:
for each e E' do if min(distance((,. b)) de, Vb B then
7:
8:
Add to B else if Fb for b arg minddistance(e, b)) then Replace with
9:
10:
11:
end if end for
12:
13:
Sort B by log-reward
14:
Truncate B to capacity
15:
end if
B.6 SHORTPARSE
In this section, describe in detail how ShortParse Pa is computed. We only ShortParse for string-based environment for its tractability, so the focus will be on DAGs that generate strings. Let's introduce some notation first:
Notation:
A: the alphabet of characters.
V: the set of tokens. We will not assume A C V unless stated, in which case the inclusion of in V forms the set of "atomic tokens".
X: the set of terminal sequences.
S X: state space, the set of sequences that can be reached on the way generating
sequence X.
T: the set of trajectories (sequences of tokens that when concatenated give some X).
T, (s S): the set sequences of tokens that when concatenated give s. Note = Uxex Tx.
V): the length of |r| (r T): the number of tokens in trajectory т. Note that in general for E Ts, |T| s |s].
|s| (s e S or
i: (s e S. 0 < [s]): the initial substring of of length i.
Let N(s) = 1Tg1. For given s, we can compute ((s.)) for = 0...., |s| by the following recurrence: ((sa))
N(s,)= e V: ends in
Σ
N(s.-1a)
(i > 0).
This computation is linear in the sequence length and can be done incrementally. as s is being con- structed by apposition of tokens. That is, maintains along with array containing (s..) for 0 & S |s]. When a token is appended to s, we append to this array the newly computed N(sa+)) for j l.....l...
20

---

Published as a conference paper at ICLR 2025
It is also possible to compute weighted version. Let
Na(s) Σ 1rl
so N(s) No(s). This be computed by the recurrence:
(x(5:)) = Nx(5.) ea
Σ
N(5
(i 0).
ends
family length-dependent Markovian backward policies For any d E R, we define back- ward policy pd as conditional distribution over trajectories, by
P edirl edl (xEX.TE Tg).
N2(x)
Proposition: This policy is Markovian, and its stepwise factorization is given by
e'Na(s)
P (s st) = Na(st) (s (SES,
Proof:
First, we have check that this stepwise policy is well-defined, that is, sums parents of any state st. This can be seen from the recurrence for N₂ above.
Now, suppose x the concatenation of tokens a trajectory T 1112...4, Then the expression for
Pn(T xx) given by the product of stepwise transitions is
e (0) eaNa(ti) (a(1)) Na(tit2)
e e'Na(t1.... (n-1) nd
Na(1 (n) Na(x)*
which exactly matches the trajectory- wise definition of P1 above. D
We note the following:
The stepwise factorization of pd can be useful, for example, if using forward-looking objective. As a special case, with = we have the maximum-entropp backward policy as studied in Mohammadpour al, (2024), which is uniform over trajectories conditioned on the terminal state notably different from the backward policy that is uniform each transition. that paper, the numbers of trajectories to every state had to be learned, but in our setting, computation is efficient (see above) and no learning necessary.
For A O, shorter trajectories are preferred, and as d -oo, approach policy that is peaky on optimal tokenizations. Conversely, if + +oo and s V, we approach a policy that only tokenizes into atomic tokens.
For ShortParse, we use -5 in all the paper.
B.7 SAMPLING QUALITY METRICS
In Infsubsection 6.1]
we computed the quality of the learned distribution using three metrics. We detail
them here:
JSD. This is the Jensen-Shannon divergence, it measures the similarity between two distributions and is symmetric and takes values in [O, 1J. Given two distributions and q, it defined as follows:
JSD(pllg) Dx. (pllM) - (xl.
Where D is the Kullback-Leibler divergence and M = ())))).
L1 distance. This computes the LI norm between two distributions and and is double the total variation distance:
||p - --h1 Σ 1(x))
(9)
21

---

Published as a conference paper at ICLR 2025
Table 3: Number of modes discovered during training in the L14_RNA1 environment. The ATOMIC case serves baseline for comparing BPF, WordPiece and Uniform tokenizers for hoth mechanisms. Greener shades indi- cate greater improvement over the ATOMIC baseline, while redder shades represent a decline in performance.
ATOMIC
ACTIONPIECE INCREMENT
ACTION
Sampler
BPE
Uniform
WordPiece
BPE
Uniform
WordPlece
GelowNet
47.12
49.00=1.15
46,33+2.9
40.00=14:
45 54g500
42 87
49 49.677904
Fettoe
47.12a411
48 48.83
37.33
49.67.0.7
45.00
27.33
48.33a1.30
ShortParse GilowNeet
47.12.4.1
38.97.1.4
39.00.2.12
49.67_094
43.01.699
22.00.8
50.33.644
36.67-130
23.50-0m4
23.67+4.88
26.67-0.n
18.00-1:1
25.67-094
19.67-2.9
SAC
19.50=149
1.17:1334
0:00s000
1.67e1255
3,50+2.2
4.33=205
1.33
Option-Critic
16.83=195
Random Sampler
0.17
18.67.525
1.00E4:
28.67-125
18.17
0.33.0.47
19.67a1.0
ELBO Gap.
Following past work (Lahlou et al. 2023 Zhang & Chen 2022 Sendera et al. 2024).
compute
variational lower bound on the log-pat pe-paatition function log 2,, R(x):
log Σ (xx)
R(x)
log E ())
Pr(T)
R(x)P
E B log
PF(T)
Where PB and PF are the backward and forward policy respectively. Thus log 2 1 --=(..-x)-Pp log R(x)))(T)x) that represents the lower bound estimate, is computed using Py(T))
K 10000 trajectories in this paper for all tasks. The ELBO gap is hence defined as:
ELBO Gap = log log 2)
P 1x)
Σ
= log Z
log
PF(T)
T=(...X)-PP
Note that only compute the above quantity for environments where the ground-truth partition function is tractable to compute.
C LEARNED LIBRARIES OF CHUNKS
In this section, we show the chunks learned by each sampler for both chunking mechanisms.
C.1 FRACTALGRID
Figure 7 shows subset of chunks found using the ACTIONPPECT-INCREMEEN INCREMENT mechanism. All samplers effectively learn diverse "paths" and these range from straight lines to a zigzag-like shape. These zigzag shapes allow the sampler to get from the initial state to the first mode in the middle of the grid. Note that straight UP and RIGHT chunks can also lead the modes by composing them together. This presents a straight-forward solution where the sampler can choose to first get the coordinates right by just using the straight line chunks towards the right and then switch to using straight line chunks going
C.2 RNA BINDING
Figure 8 shows the frequency of usage of chunks for different samplers for both chunking mecha- nisms. A big difference between both chunking mechanisms that seems be for all samplers, that the size of learned chunks is bigger which makes the learning not much more efficient than using the original ATOMIC baseline. Indeed in Figure 8b] lot of these chunks are rarely used compared to EETTOPPIEE-INCRRMEEN
When it comes to samplers however, regardless of the chunking mechanism, GFN-related approaches seem learn shorter chunks average compared to the other samplers which in turn compresses the trajectories enough to cover the whole state-space as opposed to RL-based methods.
C.3 BIT SEQUENCE
Figure 9 shows the frequency of of learned chunks for the bit sequence task for GFN-related approaches using the ACTION PIRCE-INCREMENN mechanism. We can see that most of the library comprises of short chunks with the exception of GFN and ShortParse-GFY that seem to have "mem-
22

---

Published as a conference paper at ICLR 2025
L14 RNAT
150 RNA1
1.0
A2C
0.9
Random Sampler
0.9
ActionPacee Ruplace
0.7
Atomac
0.7
Diverstty
(a) Average reward and diversity for the top 100 samples for the RNA Binding environment.
Suquaenco
GPN
A2C
0.5
Randem
0.8
0.8
ActsonPicce Repiace
Atrmac
0.6
Increment
0.5
(b) Average reward and diversity for the top 100 samples for the Bit Sequence environment.
Maximum number of nodes:
1.0
GFN
A2C
0.8
Random Sampler
SAC
Reward
ActionPiece RecionPiece-Replace
Atomic
iictonPiec-Inccemeen
2
3
Diversity
(c) Average reward and diversity for the top 100 samples the Graph environment.
Figure 6: Quality of high-reward samples. Average reward and diversity for the top 100 samples from 10000 generated objects from all samplers and all chunking mechanisms in the RNA Binding, Bit Sequence and Graph environments.
23

---

Published as a conference paper at ICLR 2025
GPN
A2C
SAC
Randem Sampler
Figure 7: Chunks visualization. Visualization of subset of chunks learned by the samplers for the FractalGrid environment of size 652 with the ACTIONPIECE-INCCREME mechanism.
orized" one of the modes as part of their learned library. For A2C and SAC however, they seem to have learned chunks that generate single mode that contains only This can be seen in (Fiiuu 10 where A2C doesn't It use the action at all and SAC sampling mostly big part of the same mode. The Random Sampler given as a reference keeps a balance in exploration although it also seems to have added one of the modes to its library. For the ACTIONPIECE-RE trend seems to be reversed, where GFN approaches discover shorter chunks than they had under the ACTIONPIECE- INCREMENT mechanism (see Figure 1T) and A2C being much more diverse under ActiONPieCE- REPLACE. REPPLLE Although SAC learned diverse library of chunks, it stuck to one mode for its last round of chunking as demonstrated Figure 12
mechanism, the
C.4 GRAPH
In this section, we won't show the frequency of all used graph chunks since it is harder to visualize. Instead, we will show some examples of learned subgraphs. We can see from Figure that all sam- plers are able to learn structures that maximize the number of cycles in the generated graph, which is the reward function can see that A2C already "overfits" the reward by seeking chunks that maximize the reward without regards to the diversity whereas this is not the case for SAC, Random Sampler and GFN.
used. All samplers effectively learn meaningful fragments. However,
D THE ROLE OF TOKENIZATION STRATEGY
In this section, we ablate the role of the tokenization strategy by comparing Byte Pair Encoding (BPE), WordPiece, and a uniform tokenizer as baseline. The uniform tokenizer does not require dataset; it builds new chunks by randomly sampling two tokens from the library and concatenating them. As shown in Table 3 the uniform tokenizer performs worse than both BPE and WordPiece in terms of the number of modes discovered during training for the L14_NNA1 task. Its performance deteriorates further with the ACTIONPIECE- -EPLECA mechanism compared to ACTIONPIECE- INCREMENT.
On the other hand, WordPiece outperforms BPF, showing particularly strong performance with the ACTIONPIECE- REPLACE mechanism. We conclude that the choice of tokenization algorithm crucial, and extra care should be taken when selecting an appropriate tokenizer.
E SAMPLER EFFECTIVENES! IN SAMPLING HIGH-REWARD OBJECTS
In this section, we look the top high-reward samples for all samplers as well their diversity. Starting with RNA Binding, we can see from Figure 6a that RI-based methods see their average reward of their distribution tail increase with chunking while the diversity of the samples decreasing.
24

---

Published as a conference paper at ICLR 2025
CNN
Mastii-GEN
Gart
23-
50
30
2000
10000000
33001>
20
20
20
70
000
50
50
500
2000
Sumpler
62
52
(a) Frequency of use of chunks for the ACTIONPIECE-INCREMEEE mechanism.
CN
Masiin-GUN
Shortiarsw-GGi
Randan
62
92
(b) Frequency of use chunks for the ACTIONPIECE REPLACE mechanism.
-EPPLLC
Figure 8: Frequency of use of chunks. Frequency of usage of chunks in the 1.1 1.1__NNA1 environment for both chunking mechanisms.
This highlights a major flaw of chunking on RL-based methods: They create a negative feedback- loop that harms diversity, which in turn gives chunks concentrated around small region of the search space which itself further harms exploration. While this is true for RL methods, it is clearly not the case for GFlowNet since can from the figure that the diversity of the top-100 samples doesn't decrease with chunking. For bit sequence, the conclusion is similar. Indeed, one can see from Figure 6b) that chunking harms RL-methods diversity and increases the average reward. For GFlowNet, the diversity actually increased with chunking. This may be attributed to the structure of the reward distribution where modes (corresponding
reward of 1) are themselves diverse. For the Graph environment (see Figure 6c), the previous observations stands as well, although for GFlowNet, it seems that diversity hurt as well in this case.
25

---

Published as a conference paper at ICLR 2025
0.25
GFN
0.25
Maxtnt GPN
0.25
ShortPParse GPN
0.20
Erequaacy
0.10
0.05
0.05
0.05
0.00
0.00
0.00
Figure 9: Frequency of use chunks the CERIONPPRCC--NCREMEEN mechanism for GFlowNet in bit sequence for length 128.
26

---

Published as a conference paper at ICLR 2025
0.25
A2C
0.25
SAC
0.25
Random Sampler
0.20
0.20
0.10
0.05
0.05
0.05
0.00
0.00
0.00
Figure 10: Frequency use chunks for CETTOONPPECE-INCREEME mechanism for RI. methods bit sequence for length 128.
27

---

Published as a conference paper at ICLR 2025
0.25
GPN
0.25
Maxtnt GPN
0.25
StertParse GPN
Eaquuncy
0.10
0.10
0.05
0.05
0.05
0.00
0.00
0.00
Figure 11: Frequency of use of chunks for ACTIONPIECE-REPL..C mechanism for GelowNet in bit sequence for length 128.
28

---

Published as a conference paper at ICLR 2025
0.25
A2C
0.25
SAC
0.25
Random Sampler
0.20
0.20
Erequaacy
0.10
0.05
0.05
0.05
0.00
0.00
0.00
100001 11 ROGANMMOT COoo
Figure 12: Frequency use of chunks for the ACTIONPIECE-REPM mechanism for methods bit sequence for length 128.
29

---

Published as a conference paper at ICLR 2025
GFN
A2C
SAC
Random Sampler
Figure 13: Visualization of subset chunks learned by the samplers the Graph environment with maxi- mum number of nodes of 7. Note for chunks where nodes connected to numbers, these represent the nodes in the graph relative to the last node, that they would be added to.
30


Reviews:
[{"review_id": "AjQjRflBn9", "reviewer": "Reviewer_3r38", "summary": "This paper introduces a approach to incorporate action abstractions into the policy optimization process of reinforcement learning and generative flow networks. The method, termed ACTIONPIECE, aims to improve sample efficiency and mode discovery by iteratively extracting and chunking action subsequences from high-reward trajectories into a growing action space. The experimental results demonstrate the effectiveness of the proposed method.", "questions": "See Weaknesses.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The proposed ACTIONPIECE compatible with both RL and GFlowNets sampler.\n- Empirical evaluation showing improved sample efficiency and mode discovery in different environments.", "weaknesses": "- As mentioned in the related works section, the discovery of macro-actions has been extensively studied. The authors should provide a more detailed discussion highlighting how the proposed method differs from existing methods.\n\n- In Line 159, the paper appears to assume a deterministic state transition, where $s'=s+a$. This assumption may be too strong and not applicable to real-world environments where state transitions involve a degree of randomness. The reviewers are concerned about the generalizability of the proposed method to stochastic environments. The authors should address whether and how the method can accommodate stochastic state transitions, which are common in many practical applications of RL and GFlowNets.\n\n- The reviewers suggest that the authors should provide a more comprehensive introduction to the concept of \"amortized samplers\" in the Preliminaries section. It is essential for readers who are unfamiliar with this concept.\n\n- The proposed Algorithm in Section 4 seems to be applicable to both GFlowNets and RL methods with discrete action spaces, as evidenced by the experiments conducted in the paper. However, the authors have chosen to focus heavily on GFlowNets as the primary background, which may not be immediately clear to readers. The reviewers recommend that the authors clarify why GFlowNets were chosen as the main framework and how the proposed method specifically leverages or addresses challenges unique to GFlowNets. \n\n- In Line 304, the paper mentions an action encoder but does not elaborate on how it is trained. The authors should provide details on the training process of the action encoder. \n\n- The experiments presented in the paper is simple. The reviewers question why the authors did not test their method in more complex environments, such as Atari games, which are known for their high dimensionality and complexity. Additionally, the paper lacks a comparison with existing methods for discovering macro-actions.", "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": "vvcNQCuG3F", "reviewer": "Reviewer_QHVg", "summary": "The paper proposes a new technique to iteratively extract action subsequences from high-reward trajectories as temporally extended skills that are then added to the action space for more efficient exploration. Applying this technique on top of generative flow networks (GFlowNets) yields improved sample efficiency in hard exploration problems. The authors experiment on three synthetic tasks and an RNA sequence generation task and show that the proposed method improves GFlowNets and outperform prior methods (e.g., RL methods like A2C and SAC) on these tasks.", "questions": "N/A", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The idea of expanding the action space dynamically online with temporally extended action sequence is interesting and novel.\n- The experiments are comprehensive and thorough with insightful analyses and visualizations that demonstrate the effectiveness of the proposed algorithm.", "weaknesses": "*Unfounded claim*\n- “the abstracted high-level actions are interpretable, …” — there is no evidence presented in the paper that illustrates the high-level actions are interpretable. \n\n*Comparison to prior chunking mechanisms is limited*\n- The authors considered two new chunking mechanisms, \"ActionPiece-Increment\" and \"ActionPiece-Replace\". Both of them use heuristics to expand action space with temporally extended action sequences. \n- It is unclear how these mechanisms compare to prior chunking mechanisms used in the options/unsupervised skill discovery/hierarchical RL literature (e.g., with variational formualtion [1], with clustering [2].\n\n[1] Kim, Taesup, Sungjin Ahn, and Yoshua Bengio. \"Variational temporal abstraction.\" Advances in Neural Information Processing Systems 32 (2019).\n\n[2] Srinivas, Aravind, et al. \"Option discovery in hierarchical reinforcement learning using spatio-temporal clustering.\" arXiv preprint arXiv:1605.05359 (2016).", "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": "epwA0nKjbe", "reviewer": "Reviewer_A2Lx", "summary": "This paper addresses the long-planning horizon problem in credit assignment by providing a trunking method ActionPiece. This approach aims to extract high-order actions from sampled trajectories and can be plugged into any sampler. In experiments with three classic algorithms and various environments, the capability of the proposed method in mode discovery, density estimation improvement, and sample length reduction is demonstrated. Abundant discussions and information are further provided.", "questions": "(Q1) As the two proposed chunking mechanisms show contrasting capabilities in multiple aspects, I am curious about the possibility of adaptively combining them. Do you have any related explorations?\n\n(Q2) How does the proposed method perform with other tokenization techniques? Does the chunking method particularly fit some of the tokenization techniques such as BPE used in the paper?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "Overall, this paper is well presented and provides an articulate method. I particularly appreciate the environment selections of a real-world orientation and informative way of discussion.", "weaknesses": "(W1) Minor typos, e.g.,’the the’ at line 319.\n\n(W2) It seems that all experiments are averaged from only three seeds per line 348 and 507, which is not enough to demonstrate statistical significance in some settings.", "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": "i49ElQGVxg", "reviewer": "Reviewer_ULj8", "summary": "The paper investigates the action abstraction for amortized sampling with reinforcement learning and generative flow networks. It proposes to iteratively extract high-reward action subsequences to be expanded as new actions, trading increased breadth for decreased depth of the Markov decision process.  In each iteration, it uses the byte pair encoding for generating high frequency subsequences from the high-reward sequences. Two variants are considered: one incrementally adds new subsequences, while the other periodically replaces old action abstractions.", "questions": "- How does your approach specifically differ from the method proposed by Zheng et al. as well as other action abstraction methods?\n\n- The conclusions seem case-specific. Can you provide a more generalized summary of the key takeaways from your research?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 1, "contribution": 1, "strengths": "**Quality**: The paper is of good methodological quality. It conducts a careful and meticulous empirical investigation of the use of action abstraction in amortized sampling.\n\n**Significance**: The exploration of action abstraction in this manner could potentially lead to more efficient sampling techniques, although its impact remains to be fully assessed.", "weaknesses": "**Novelty**: The use of byte pair encoding for learning action abstraction is not new, as it is already explored in the work by Zheng et al. [1]. To enhance the novelty, the authors should differentiate their approach more clearly or build upon the existing frameworks with significant innovations.\n\n**Significance**: The proposed algorithm exhibits varied performance across different tasks and samplers, lacking a consistent indication of its promise as a robust approach. This raises questions about its generalizability and efficacy in broader applications.\n\n**Clarity**: The conclusions are not effectively communicated. Despite thorough reading, the paper does not present clear, actionable takeaways or insights, as results appear highly case-specific. Improving the clarity of conclusions and providing more generalized insights would increase the paper's impact.\n\n**References**\n\n[1] Zheng, Ruijie, Ching-An Cheng, Hal Daumé III, Furong Huang, and Andrey Kolobov. \"PRISE: LLM-Style Sequence Compression for Learning Temporal Action Abstractions in Control.\" In Forty-first International Conference on Machine Learning.", "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": "FTiJAtzJs5", "reviewer": "Reviewer_juwr", "summary": "This paper proposes a method named ACTIONPIECE that incorporates action abstraction discovery into the policy optimization process. It leverages tokenizers to chunk action sequences, which can be viewed as a strategy for trading reduced depth of the MDP in exchange for increased breadth.", "questions": "Please check the weaknesses.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The proposed method is a novel approach to improve the performance of samplers on model discovery and capturing the latent structure of the environment.\n\n2. The method has been validated in three different scenarios, demonstrating its broad applicability and robustness.", "weaknesses": "1. The paper lacks rigorous theoretical analysis or proofs to support the observed improvements. While empirical results are strong, a more detailed theoretical discussion would strengthen the claims about the performance of the proposed method.\n\n2. The abstract lacks a quantitative presentation of the results.\n\n3. Can the proposed method be applied to offline-training RL tasks? Can the authors provide some discussion on offline training?\n\n4. As mentioned by the authors, the use of a fixed BPE tokenizer for chunking fixed target distributions in experiments may limit the method's generality and flexibility.", "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)