PaperlayerView exact revision on arXiv

Source-linked AI summary

Action abstractions for amortized sampling

Oussama Boussif, Léna Néhale Ezzine, Joseph D Viviano, Michał Koziarski, Moksh Jain, Esmeralda S. Whitammer, Emmanuel Bengio, Rim Assouel, Yoshua Bengio

arXiv:2410.15184v2cs.LGcs.AI

TL;DR

Long trajectories make credit assignment harder, while diversity-seeking samplers must discover multiple high-reward states. The paper adds tokenized action chunks to samplers’ action spaces and reports improved mode discovery, density estimation, and recovery of target-distribution structure.

  • Problem

    Longer trajectories make credit assignment more difficult, and diversity-seeking RL must sample terminal states in proportion to reward.

    Source: §1 · p. 1

  • Method

    ACTIONPIECE iteratively tokenizes sampled action sequences into chunks and adds the resulting abstracted actions to the sampler’s action space.

    Source: §1 · p. 2, §4 · p. 5

  • Results

    The authors report accelerated mode discovery, improved density estimation, and chunks that represent target-distribution structure and transfer across algorithms and tasks.

    Source: §1 · p. 2, §7 · p. 10

  • Takeaways & Limitations

    The authors conclude that learned chunks can capture underlying structure in the target distribution and improve sampler performance across varied domains.

    Source: §7 · p. 10

    Chunking uses BPE for a fixed target distribution; task-conditioned libraries and distributions over possible abstraction libraries remain future work.

    Source: §7 · p. 10

Abstract

from arXiv · show

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 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.

1 INTRODUCTION

Long trajectories make credit assignment harder in RL, while entropy-seeking methods such as GFlowNets must also discover multiple high-reward states. The paper introduces ACTIONPIECE and ShortParse, and reports gains in mode discovery, density estimation, and representation of target-distribution structure.

1 INTRODUCTION

  • Long trajectories intensify credit-assignment difficulty, while diversity-seeking RL must sample terminal states in proportion to reward.Prior work shortens trajectories or encourages compressible sequences to address this challenge.Source: §1 · p. 1
  • ACTIONPIECE extracts high-order actions, or chunks, from sampled trajectories using tokenization; its two variants are compatible with any sampler at no significant cost.The chunking procedure filters trajectories for high-reward samples, identifies frequent chunks, adds them to the action space, and repeats.Source: §1 · p. 2, Figure 1 · p. 2
  • ShortParse introduces a backward policy for GFlowNets to sample short trajectories.Source: §1 · p. 2
  • The paper reports that its approach accelerates mode discovery, improves density estimation, and reduces sample description length.Source: §1 · p. 2
  • Learned chunks are reported to represent target-distribution structure and transfer to different algorithms and new tasks.Source: §1 · p. 2

2 RELATED WORK

Related work covers macro-actions and policy libraries as ways to structure sequential decisions, alongside methods aimed at credit assignment. Prior macro-action results are mixed, with usefulness associated with shared structure between macros and the environment’s reward landscape.

2 RELATED WORK

  • Macro-actions assemble atomic actions into sequences executed at once; prior results are mixed, and usefulness has been associated with shared structure between macros and the environment’s reward landscape.Some settings report faster convergence, lower return variance, or improved returns.Source: §2 · pp. 2–3
  • Program-induction library-learning methods incrementally modify a domain-specific language to reduce useful-program complexity, often using an abstraction phase to introduce subroutines.Source: §2 · p. 3
  • Credit-assignment methods include assigning credit to past decisions based on outcome likelihood and GFlowNet objectives that use whole action sequences or additive per-step energy components.Source: §2 · p. 3

3 PRELIMINARIES

The preliminaries define amortized sampling over a directed acyclic graph and describe how GFlowNets and RL algorithms formulate learning over its actions. They introduce trajectory balance for GFlowNets and actor-critic methods for comparison.

3 PRELIMINARIES

  • Amortized sampling shifts the cost of iterative sampling to model optimization, enabling rapid sampling after training; the objects are terminal states in a directed acyclic graph.Actions incrementally build objects, with a global action set and state-specific available subsets.Source: §3 · p. 3
  • Trajectory balance trains a forward policy so each trajectory’s probability matches its backward-policy probability times terminal reward, implying terminal-state probabilities proportional to reward.The objective uses a learned normalizing-constant estimate and a backward policy.Source: §3 · p. 4
  • The GFlowNet setting translates to a deterministic-transition MDP with log reward at terminal states; the paper also recalls A2C and SAC actor-critic learning.The supplied update passage describes a critic update and a variance-reducing baseline; SAC trades off expected return and policy entropy.Source: §3 · pp. 4–5

4 DISCOVERING ACTION ABSTRACTIONS FOR AMORTIZED SAMPLERS

The method augments amortized samplers with discovered action abstractions to shorten trajectories, using tokenization to turn action sequences into reusable chunks. It builds and updates the action space during training, with incremental and replacement variants.

4 DISCOVERING ACTION ABSTRACTIONS FOR AMORTIZED SAMPLERS

  • Action abstraction aims to shorten trajectories in discrete-object sampling, where longer paths make training difficult and can trade reduced MDP depth for increased breadth.Source: §4 · p. 5
  • The three-step procedure samples action sequences, applies byte pair encoding to obtain chunks, then adds those abstracted actions to the action space.When chosen, a chunk’s constituent actions execute in order.Source: §4 · p. 5
  • ACTIONPIECE-INCREMENT adds the most frequent token per chunking round, while ACTIONPIECE-REPLACE appends the M most frequent tokens and can replace the learned action set while retaining atomic tokens.Incremental updates risk repetitions or low-quality sequences; replacement reduces redundant chunks and can discover multiple chunks per round.Source: §4 · p. 6
  • Chunking can occur every k iterations or, for GFlowNets, after trajectory-balance loss falls below a threshold.A selected abstracted action is masked when its sequence includes an action unavailable at the current state; unlike options, chunks are executed fully.Source: §4 · pp. 5–6

5 EXPERIMENTAL SETUP

The experiments benchmark SAC, A2C, GFlowNet, and random sampling, with several GFlowNet backward policies. The policy supports an evolving action space by embedding actions, and evaluation spans synthetic and RNA-generation tasks.

5 EXPERIMENTAL SETUP

  • The benchmark includes SAC, A2C, GFlowNet, and a random sampler, with trajectory balance and Uniform, MaxEnt, or ShortParse backward policies for GFlowNet.Uniform and MaxEnt backward policies are equivalent for Tree MDPs.Source: §5 · p. 6
  • To handle the evolving action space, the policy produces an action embedding and scores each available action using an LSTM-encoded action representation.The LSTM encodes each action as its sequence of atomic actions, including length-one atomic actions.Source: §5 · p. 6
  • Evaluation covers bit-sequence generation, FractalGrid, graph generation, and the practical L14 RNA1 sequence-generation task.The three synthetic tasks are bit sequence, FractalGrid, and graph generation.Source: §5 · p. 6

6 RESULTS

Across tasks, ACTIONPIECE speeds GFlowNet mode discovery, while density-estimation results are favorable with exceptions. Analyses of RNA chunks find evidence of latent structure, shorter parses for GFlowNet-based methods, and transfer to related distributions.

6.1 DOES CHUNKING IMPROVE GENERATIVE MODELING?

  • ACTIONPIECE speeds GFlowNet mode discovery across tasks; Figure 2 reports cumulative discovered modes and especially highlights gains in FractalGrid, where atomic samplers get stuck at the first mode.Results for other algorithms are mixed: chunking helps A2C and SAC in FractalGrid but hurts them on RNA binding.Source: §6.1 · p. 7, Figure 2 · p. 7
  • Chunking generally improves GFlowNet density estimation, although BitSequence is marginally worse and ACTIONPIECE-REPLACE lags on L14 RNA1.Chunking accelerates training across the tested FractalGrid sizes; ACTIONPIECE-INCREMENT performs on par with ATOMIC on L14 RNA1 and slightly better for high-reward objects.Source: §6.1 · p. 7, Figure 3 · p. 8

6.2 DO CHUNKS CAPTURE THE STRUCTURE IN THE UNDERLYING DISTRIBUTION?

  • For high-reward L14 RNA1 sequences, ACTIONPIECE-INCREMENT yields higher chunk occurrence for GFlowNet-based samplers, and chunk coverage shows a similar pattern.A2C, SAC, and random sampling have moderate-to-low occurrence under INCREMENT, while REPLACE yields near-zero median occurrence across samplers; REPLACE also trails INCREMENT.Source: §6.2 · pp. 8–9, Table 2 · p. 8
  • GFlowNet-based approaches tend to produce the shortest parses of modes, while RL-based methods have longer average parses; Figure 4 compares models and chunking strategies.For RL-based methods, REPLACE compresses samples less effectively; GFlowNet-based methods produce shorter average parses with INCREMENT than with REPLACE.Source: §6.2 · p. 9, Figure 4 · p. 9
  • Libraries learned on L14 RNA1 improve mode discovery on related L14 RNA2 and RNA3 distributions compared with no chunking; GFlowNet-induced libraries perform better on average than libraries from other samplers.Figure 5 compares environments by column and chunking mechanisms by row, with heatmaps showing modes discovered by samplers using libraries found by other samplers.Source: §6.2 · p. 9, Figure 5 · p. 10

7 CONCLUSION

ACTIONPIECE learns reusable action chunks for amortized samplers, trading reduced MDP depth for increased breadth. The paper reports improved mode discovery and structure capture, while identifying fixed-target BPE chunking as a boundary for future work.

7 CONCLUSION

  • ACTIONPIECE uses tokenizers to chunk action sequences, trading reduced MDP depth for increased breadth in amortized sampling.The conclusion frames the method as extensible dynamic action-space learning.Source: §7 · p. 10
  • Across varied domains, the approach improves sampler performance in mode discovery and captures latent environment structure.Learned chunks can generalize to other target distributions and transfer to other tasks.Source: §7 · p. 10
  • Online library construction works best with methods that explicitly maximize diversity, according to the reported results.The authors connect this interaction to shorter parse lengths and discovering new modes when chunks transfer to new tasks.Source: §7 · p. 10
  • The experiments use BPE chunking for a fixed target distribution; task-conditioned libraries and modeling a distribution over libraries remain future directions.The authors also suggest applying abstraction to code generation.Source: §7 · p. 10

A ADDITIONAL RELATED WORKS

  • Tokenizers produce vocabularies for processing text; BPE merges frequent adjacent token pairs, while WordPiece merges pairs by a joint-to-marginal likelihood ratio.Both procedures repeat until reaching a maximum vocabulary size.Source: A ADDITIONAL RELATED WORKS · p. 16

B EXPERIMENTAL DETAILS

The experimental details describe the environments, sampler configurations, replay-buffer procedure, and ShortParse backward policy used in the study. Settings vary across domains, including reward definitions, sequence or graph sizes, and chunking schedules.

B EXPERIMENTAL DETAILS

  • Experiments use one GPU per run, with runs taking up to 36 hours; the appendix also specifies per-environment architectures and sampler hyperparameters.The implementation details include tokenization background and environment code.Source: B EXPERIMENTAL DETAILS · p. 16

B.1 FRACTALGRID

  • FractalGrid uses a 2D hypergrid-inspired environment with UP, RIGHT, and EXIT actions, three increasing grid sizes, and chunking every 1,250 iterations.Runs use 31,250 iterations with batches of 64, totaling two million visited states; 25 chunks are added by training's end.Source: §B.1 · p. 17

B.2 RNA BINDING

  • RNA Binding samples strings of length 14 or 50 from A, C, G, U, and EOS, with chunking every 1,000 iterations for most samplers.GFlowNet-based samplers instead trigger chunking at a trajectory-balance loss threshold, then multiply that threshold by 0.75 after each application.Source: §B.2 · pp. 17–18

B.3 BIT SEQUENCE

  • The modified bit-sequence reward favors strings containing a maximum number of five specified 8-bit words, including repeated instances, and experiments consider sequences of lengths 64 and 128.Chunking is performed every 1,000 iterations to produce 25 chunks at the end of training.Source: §B.3 · p. 18

B.4 GRAPH

  • The graph action design adds nodes sequentially and connects edges relative to the last-added node, making action patterns consistent across absolute node numbers.The reward is the number of cycles divided by the maximum possible for a graph with N nodes; chunking occurs every 1,000 iterations.Source: §B.4 · pp. 18–19

B.5 REPLAY BUFFER

  • The replay buffer retains batch samples whose rewards exceed its minimum, then admits diverse states or replaces their nearest buffer neighbor when the new state has higher reward.It sorts retained items by reward and truncates to capacity.Source: §B.5 · pp. 19–20

B.6 SHORTPARSE

  • ShortParse is a length-dependent Markovian backward policy for string-generating DAGs, with its trajectory counts computable incrementally in time linear in sequence length.For ShortParse, the authors use lambda = -5; negative lambda favors shorter trajectories.Source: §B.6 · pp. 20–21

B.7 SAMPLING QUALITY METRICS

The appendix defines metrics for comparing learned distributions and estimating the partition function, and describes a figure assessing high-reward sample quality. It limits ELBO-gap computation to settings with tractable ground-truth partition functions.

B.7 SAMPLING QUALITY METRICS

  • Jensen–Shannon divergence measures similarity between distributions, is symmetric, and ranges from 0 to 1; L1 distance is twice total variation distance.The section introduces these among three metrics for evaluating learned-distribution quality.Source: §B.7 · p. 21
  • The ELBO gap compares a variational lower-bound estimate of log Z, computed from 10,000 trajectories, with the ground-truth value where tractable.The authors compute this quantity only in environments where the true partition function can be computed tractably.Source: §B.7 · pp. 21–22

C LEARNED LIBRARIES OF CHUNKS

  • Figure 6 compares average reward and diversity among the top 100 of 10,000 generated objects across samplers and chunking mechanisms in three environments.The environments are RNA Binding, Bit Sequence, and Graph; the section also introduces learned chunk examples.Source: Figure 6 · p. 22, C LEARNED LIBRARIES OF CHUNKS · p. 22

C.1 FRACTALGRID

In FractalGrid, learned chunks trace routes to a grid mode; in RNA Binding, chunk size and use vary by mechanism and sampler.

C.1 FRACTALGRID

  • Learned FractalGrid chunks form paths ranging from straight lines to zigzags, and the zigzags reach the first mode in the grid.Straight UP and RIGHT chunks can also reach it when composed in sequence.Source: §C.1 · pp. 22–23

C.2 RNA BINDING

  • Table 3 compares modes discovered during L14 RNA1 training across ATOMIC, BPE, WordPiece, and Uniform for both chunking mechanisms.Source: Table 3 · p. 23
  • In RNA Binding, ACTIONPIECE-REPLACE learns larger chunks that are often rarely used, while GFlowNet-related samplers learn shorter chunks than other samplers.Regardless of mechanism, GFlowNet-related approaches compress trajectories enough to cover the whole state space, unlike RL-based methods.Source: §C.2 · p. 23

C.3 BIT SEQUENCE

Bit-sequence chunk use varies by sampler and mechanism, while graph samplers learn cycle-maximizing fragments with differences in diversity.

C.3 BIT SEQUENCE

  • With ACTIONPIECE-INCREMENT, GFN and ShortParse-GFN appear to memorize a bit-sequence mode, while A2C and SAC learn chunks generating the all-1 mode.The accompanying discussion says A2C does not use action 0, while SAC sampling concentrates on much of the same mode.Source: §C.3 · p. 23
  • With ACTIONPIECE-REPLACE, GFN approaches discover shorter chunks and A2C is more diverse than with ACTIONPIECE-INCREMENT, while SAC sticks to one mode in its final chunking round.Source: §C.3 · p. 24

C.4 GRAPH

  • In graphs, all samplers learn fragments that maximize cycle count, but A2C overfits this reward without regard to diversity, unlike SAC, Random Sampler, and GFN.Source: §C.4 · p. 24

D THE ROLE OF TOKENIZATION STRATEGY

The tokenizer choice affects the number of modes discovered in L14 RNA1, with WordPiece performing especially well under ACTIONPIECE-REPLACE.

D THE ROLE OF TOKENIZATION STRATEGY

E SAMPLER EFFECTIVENESS IN SAMPLING HIGH-REWARD OBJECTS

Chunking is associated with higher reward but reduced diversity for RL methods; GFlowNet diversity varies by task, remaining stable or increasing in some and falling in graph sampling.

E SAMPLER EFFECTIVENESS IN SAMPLING HIGH-REWARD OBJECTS

  • For RL methods in RNA Binding, lower diversity creates a negative feedback loop: chunks concentrate around a small search-space region and further harm exploration.Source: E SAMPLER EFFECTIVENESS IN SAMPLING… · p. 25
  • Figure 10 depicts chunk-use frequency for ACTIONPIECE-INCREMENT in RL bit-sequence methods at length 128; Figures 11 and 12 depict ACTIONPIECE-REPLACE chunk use for GFlowNet and RL.Source: Figure 10 · p. 26, Figure 11 · p. 27, Figure 12 · p. 28
  • Across RNA Binding and bit sequence, chunking raises RL methods’ average or tail reward while reducing diversity; GFlowNet diversity does not fall in RNA Binding and increases in bit sequence.In graph sampling, diversity is also hurt for GFlowNet.Source: E SAMPLER EFFECTIVENESS IN SAMPLING… · pp. 25–26
  • Figure 13 visualizes learned graph chunks for samplers in the graph environment with a maximum of seven nodes.The caption explains that numbered node connections are relative to the last node.Source: Figure 13 · p. 29

Related papers

Loading 2410.15184v2…