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
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
Keywords: GFlowNets, amortized samplers, hierarchical planning, abstractions, macro-actions

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.

Decision: accept (poster)
Reviewer scores: 6, 8, 8, 6, 6

Reviews and discussion:
<thread>
creation date:
1730601571309

thread_type:
official_review

<thread_object>
creator: reviewer_juwr

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

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.

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

2. The abstract lacks a quantitative presentation of the results.

3. Can the proposed method be applied to offline-training RL tasks? Can the authors provide some discussion on offline training?

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

questions:
Please check the weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730641517712

thread_type:
official_review

<thread_object>
creator: reviewer_ulj8

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

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.

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

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

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

**References**

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

questions:
- How does your approach specifically differ from the method proposed by Zheng et al. as well as other action abstraction methods?

- The conclusions seem case-specific. Can you provide a more generalized summary of the key takeaways from your research?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730649923926

thread_type:
official_review

<thread_object>
creator: reviewer_a2lx

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

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.

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

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?

(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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730650996634

thread_type:
official_review

<thread_object>
creator: reviewer_qhvg

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

strengths:
- The idea of expanding the action space dynamically online with temporally extended action sequence is interesting and novel.
- The experiments are comprehensive and thorough with insightful analyses and visualizations that demonstrate the effectiveness of the proposed algorithm.

weaknesses:
*Unfounded claim*
- “the abstracted high-level actions are interpretable, …” — there is no evidence presented in the paper that illustrates the high-level actions are interpretable. 

*Comparison to prior chunking mechanisms is limited*
- 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. 
- 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].

[1] Kim, Taesup, Sungjin Ahn, and Yoshua Bengio. "Variational temporal abstraction." Advances in Neural Information Processing Systems 32 (2019).

[2] Srinivas, Aravind, et al. "Option discovery in hierarchical reinforcement learning using spatio-temporal clustering." arXiv preprint arXiv:1605.05359 (2016).

questions:
N/A
</content>
</thread_object>

</thread>
<thread>
creation date:
1730710967518

thread_type:
official_review

<thread_object>
creator: reviewer_3r38

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

strengths:
- The proposed ACTIONPIECE compatible with both RL and GFlowNets sampler.
- 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.

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

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

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

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

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

questions:
See Weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732122747396

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Response to all reviewers + new experiments

comment:
We thank the reviewers for their efforts in reviewing our paper. The comments have helped us to improve the work. We answer each reviewer separately, but use this comment to clarify some recurring points that the reviewers raised and new experiments we have added. Additionally, we highlighted the parts that we added to the manuscript in blue.

**Novelty and contribution**

Our aim is to study the effect of discovering action abstractions in the context of training an amortized sampler. As a consequence, we focus on environments that clearly have repeating substructures, rather than complex RL environments with entangled generalization factors. We believe that our work is the first to study action abstractions in the context of amortized sampling, and it gives actionable insights on how to improve the training of policies for amortized inference over compositional structures.

Our approach draws inspiration from action abstraction methods in RL. Conversely, we hope and expect that our work can inspire methods applicable to RL and control.

**New experiments**

- **Comparison to other macro-action frameworks**: We trained an Option-Critic [Bacon et al., The Option-Critic architecture] agent on all environments and summarized our findings in the following table:

Number of modes discovered during training in the `L14_RNA1` environment:
| **Sampler**         | **Atomic**     | **ActionPiece-Increment** | **ActionPiece-Replace** |
|----------------------|----------------|----------------------------|--------------------------|
| GFlowNet             | 43.24 ± 2.21  | 48.67 ± 1.25                    | 41.00 ± 3.27                  |
| MaxEnt-GFlowNet      | 43.24 ± 2.21  | 48.33 ± 0.47                    | 42.33 ± 2.87                   |
| ShortParse-GFlowNet  | 43.24 ± 2.21  | 48.27 ± 1.32                    | 36.67 ± 3.77                  |
| A2C                  | 37.67 ± 0.94  | 24.00 ± 0.82                    | 18.67 ± 3.77                  |
| SAC                  | 18.33 ± 2.05  | 0.67 ± 0.47                      | 4.33 ± 2.62                   |
| Random Sampler       | 0.00 ± 0.00   | 14.67 ± 2.62                     | 17.67 ± 4.50                  |
| Option-Critic       | 16.00 ± 2.16   | ---                  | ---                  |

Number of modes discovered during training in the `bit_sequence` environment for sequence length 128:

| **Sampler**         | **Atomic**     | **ActionPiece-Increment** | **ActionPiece-Replace** |
|---------------------|----------------|---------------------------|-------------------------|
| GFlowNet                 | 32.81 (2.0)    | 35.0 (0.27)              | 21.25 (6.23)           |
| MaxEnt-GFlowNet          |     32.81 (2.0)      | 39.52 (2.47)              | 35.0 (2.94)             |
| ShortParse-GFlowNet      |  32.81 (2.0)   | 36.09 (1.64)             | 12.33 (1.25)           |
| A2C                 | 9.67 (7.76)    | 0.0 (0.0)                | 26.33 (18.66)          |
| SAC                 | 0.0 (0.0)      | 0.0 (0.0)                | 0.0 (0.0)              |
| Random Sampler      | 0.0 (0.0)      | 0.0 (0.0)                | 4.0 (0.82)             |
| Option-Critic       | 0.67 (0.94)    |---|---|

Number of modes discovered during training for the ``FractalGrid`` for size 257:

| **Sampler**         | **Atomic**     | **ActionPiece-Increment** | **ActionPiece-Replace** |
|---------------------|----------------|---------------------------|-------------------------|
| GFN                 | 1.0 ± 0.0      | 16.0 ± 0.0                | 16.0 ± 0.0              |
| A2C                 | 1.0 ± 0.0      | 3.67 ± 2.36               | 16.0 ± 0.0              |
| SAC                 | 1.0 ± 0.0      | 2.33 ± 1.89               | 8.33 ± 0.94             |
| Random Sampler      | 1.0 ± 0.0      | 10.33 ± 0.47              | 15.67 ± 0.47            |
| Option-Critic       | 1.0 ± 0.0      |---|---|

Note that Option-Critic should be compared with the ActionPiece-Increment/Replace columns of the other samplers since both learn macro-actions. We can see that the Option-Critic sampler struggles.

- **Comparison to WordPiece**: We studied action abstraction when replacing BPE with WordPiece and updated Table 3 from the appendix in the manuscript. We also provide a screenshot of the updated table [here](https://ibb.co/Xt30qLs). We observe that WordPiece results in better performance than BPE with all the samplers for both ActionPiece-Increment and ActionPiece-Replace. Thank you to Reviewer A2Lx for this insightful suggestion!

We appreciate greatly these comments, which believe have helped us strengthen the paper.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732752420019

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Final thank you for all reviewers.

comment:
We would like to take a moment to thank all of the reviewers for their helpful comments which we believe have dramatically improved our paper. In this latest revision which we just uploaded, we have introduced the results from doubling the number of seeds for the core set of experiments, improving the communication of our results, and improving the visualization of Figure 4. All meaningful changes to the text are highlighted in blue. and we look forward to addressing any further comments or concerns you have through the remainder of the discussion period.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposes ACTIONPIECE, a novel approach for incorporating action abstraction discovery into policy optimization for amortized sampling, using tokenizers to chunk action sequences from high-reward trajectories. While reviewers appreciated the paper's methodological rigor and comprehensive empirical validation across different environments, particularly the effectiveness demonstrated with GFlowNets, there were some concerns about the theoretical foundations and generalizability of the approach. The authors provided thorough responses, including additional experiments with Option-Critic baselines and offline datasets, which helped clarify the method's advantages in the amortized sampling setting. The revisions also improved clarity regarding the novelty relative to recent work like PRISE and strengthened the presentation of key findings. Notably, the results consistently show that GFlowNet-based methods make better use of the learned action abstractions compared to other sampling approaches. 

Given the paper's strong empirical contributions, clear exposition, and the authors' comprehensive response to reviewer concerns, I recommend acceptance.