Action abstractions for amortized sampling

We thank the reviewer for the helpful comments on our work.

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.

We appreciate the reviewer's concern for theoretical results. However, we decided to limit the scope of our present work to empirical analysis since theoretical analysis on amortized sampling and abstractions is itself limited in prior work.

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

This is an interesting question, though not quite in the scope of problems we aim to solve with the proposed metohd. We direct you to the response to Reviewer ULj8 regarding [1], which proposes using a two-stage quantized and encoding and BPE procedure to infer macro-actions from an offline continuous control dataset. In essence, the procedure in [1] can be seen as an offline variant of our approach that does not update the token library during training. Combinations of the two approaches could be interesting to explore.

That being said, we will be conducting additional experiments to show results from training using chunks abstracted from an offline dataset for the RNA Binding task.

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

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.

We are not sure to understand the question correctly:

• Note that the tokenizer is not fixed, as the library of chunks is learned online during training and thus evolves over time. BPE is used as the chunking mechanism, i.e., to decide which new tokens to add to the library.
• If it was meant that the use of BPE, rather than other compression algorithms, may be limiting -- this is a valid concern, and we performed additional experiments to study the effect of different tokenizers. We kindly direct you to the response to Reviewer A2Lx regarding WordPiece and a uniform tokenizer.

We would like to thank the reviewer for their helpful comments and would welcome further questions during the discussion period!

We thank the reviewer for the helpful comments on our work. See below for our comments:

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.

Thank you for drawing our attention to this recent work. While the use of BPE is similar between the methods, the algorithms and settings are quite different:

In [1], the domain of continuous control is studied. The proposed algorithm proceeds in two stages. First, a quantized encoding function is learned, which maps a pair of (a latent encoding of) an observation and a continuous-valued action to a discrete token. In the second stage, BPE is applied to the sequences of encoded actions to obtain macro-actions. This is very different from our work, which operates in action spaces and learns the library online, i.e., incrementally during training. As a concequence of this, the settings to which these approaches are applicable are very different.

Our work shares with [1] the motivation of extracting high-level, discrete macro-actions and the use of BPE on action sequences to do so. However, there are a number of important differences in the problem setting and approach:

• While [1] studies continuous control, we study amortized sampling in discrete domains, where the set of discrete actions is already available (obviating the need to learn a quantized encoder).
• On the other hand, while in [1] a dataset of expert trajectories is available, we consider an online setting, where the sampling agent must explore to discover high-reward trajectories.
• While [1] has a two-stage quantization and compression procedure, we perform discovery of macro-actions online. That is, we iteratively discover new macro-actions in a dynamic dataset of high-reward trajetories discovered so far by the sampler. The growing library of useful action abstractions and the improving amortized sampler reinforce each other in a virtuous loop.

It would be interesting to combine ideas from [1] and from our work. For example, a control agent could be trained in a loop such as the one we propose, in which an agent generates sequences of actions from a dynamically growing library that is initialized with the original quantized actions from [1]'s first stage and iteratively updated by performing BPE on discovered high-reward sequences.

Thank you again for pointing out this work: we now mention its relevance after introducing the algorithm in Section 4.

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.

Thank you for raising issues around the presentation of our results. On consistent finding in this study is that the utility of macro actions (in the form of ActionPiece) is dependent on the sampling method used. In our results, the diversity-seeking GFlowNet algorithms made the best use of the chunks learned by ActionPiece: they facilitated the discovery of new modes, learned libraries which provided the shortest parse of the path to the modes in the data, and produced chunks that transferred best to new models or tasks. We have added a section to the conclusions to this effect.

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. [...] The conclusions seem case-specific. Can you provide a more generalized summary of the key takeaways from your research?

Each section of the paper ends with a short description of the results found. Motivated by your comment, we have expanded the conclusions section of the paper to provide a high-level summary of the impact of this work and the main findings.

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

Thank you again for your comments and suggestions, which have helped us improve the paper. We hope that these responses have addressed your concerns and would welcome any further feedback.

We thank the reviewer for their helpful comments on our work.

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.

Our work is different from existing literature in the following aspects:

• We study the effect of action abstractions on the training of amortized samplers, which, to the best of our knowledge, has not been explored before.
• We investigate the discovery of macro-actions online during training and dynamically modify the action space. This means the action space is no longer static but evolves throughout the training process. Known frameworks for temporal abstraction, such as Options, focus on discovering policies while keeping the action space constant. Furthermore, these frameworks are closed-loop, whereas ours operates in an open-loop manner.
• We deliberately focus on environments where high-reward samples exhibit clear, repeating substructures, as opposed to complex environments (e.g., Atari) with entangled generalization factors that would make it challenging to derive clear and actionable conclusions.

In Line 159, the paper appears to assume a deterministic state transition, where $𝑠'=𝑠+𝑎$. 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.

In this paper, we explore the use of action abstractions specifically for amortized sampling -- generation of samples from an unnormalized target density via a hierarchical generative process (i.e., by sequential decision-making). In such settings, the MDP is assumed to be deterministic (as is, in fact, done in all existing applications of GFlowNets for amortized sampling).

Generalization to sampling settings with stochastic environments -- while it would not constitute amortized inference in the commonly understood sense -- would be a direction for future work. One interesting case is action abstractions in adversarial games, where generalizations of GFlowNet objectives have been developed [Vucetic et al., "Expected flow networks...", ICLR'24]. We note that macro-action discovery in multi-agent or adversarial settings is largely unexplored in traditional RL domains as well.

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.

We have added the following to section 3 of the manuscript (Preliminaries): Amortized sampling is 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, enabling rapid sampling after training is complete.

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.

GFlowNets enable sampling from an unnormalized distribution, and our focus is specifically on their role in studying the impact of abstractions on generative modeling. We selected environments frequently used in GFlowNet research, as they provide a controlled setting for analyzing how abstractions influence generative modeling.

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 purpose of the action encoder is simply to provide embeddings for macro-actions. There is no separate loss to train the action encoder: for both GFlowNets and RL baselines, the policy logits are computed using Equation (4), so when the policy loss is backpropagated, it changes the parameters of the action encoder as well.

We thank the reviewer for the helpful comments on our work.

(W1) Minor typos, e.g.,’the the’ at line 319.

Thank you for pointing out the typo! It was fixed.

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

Thank you for this observation. The large number of experiments conducted prevents us from running more seeds for all experiments during the rebuttal period, but we will be sure to do this for the final version.

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

This is an interesting question, although it brings about a new layer of complexity in the schedule of retokenizations. One possibility would be to use ActionPiece-Increment every $x$ steps and then use ActionPiece-Replace on the learned library every $y$ steps, where $y>x$. Another is to use ActionPiece-Increment, but also to periodically remove macro-actions whose frequency of use falls below an adaptive threshold. Did you have in mind a specific way of combining both mechanisms?

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

In the appendix, we ablated the role of BPE against a uniform tokenizer. This tokenizer builds the library online by choosing two tokens at random from the current library and concatenating them to form a new token. We additionally BPE with WordPiece. Both have been included in the updated manuscript; a direct link to the results is here.

We see that BPE and WordPiece significantly outperform the uniform tokenizer, which indicates that tokenizing action sequences is the main driver of the performance. Note that the uniform tokenizer still outperforms the baseline of not introducing any new tokens during training (Atomic), which further reinforces the claim that action abstraction are helpful for training amortized samplers.

WordPiece performs strongly relative to BPE and will be evaluated on all tasks in the final version.

Thank you for the insightful suggestion, and thank you again for your other comments and questions. We are happy to provide further information if needed.

We thank the reviewer for helpful comments on our work.

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

In Section 6.2, we answer the interpretabilty question indirectly by showing, using various probes, that the learned macro-actions capture some underlying structure in the target distribution. While interpretabiilty is difficult to measure, in Appendix C we show the libraries of learned chunks for all tasks, which exhibit clear patterns consistent with the structure of the environment (see, e.g., discussion in Appendix C.1). We are happy to point to these results more prominently in the main text.

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

We thank the reviewer for pointing out specific examples in the literature:

• In [1], the authors learn a generative model that decomposes a sequence of observations into non-overlapping subsequences. They assume the existence of latent temporal abstractions that generate these subsequences and consider the number and length of each subsequence as latent variables, rather than treating them as mere hyperparameters. To achieve this, they introduce a "Binary Subsequence Indicator" variable to switch between temporal abstractions, akin to the termination probability in Options. However, their approach relies on the availability of a dataset and the overhead cost of training the generative model. In contrast, our method directly extracts the most common subsequences online at a negligible computational cost.
• In [2], authors cluster states of the MDP to extract abstract states (macrostates) and they learn options between these macrostates. While their method is similar in that it is conducted online, our approach operates directly in the action space, making it applicable to any environment. Their method, on the other hand, depends on the quality of state embeddings. Additionally, our approach incurs negligible computational cost, whereas theirs requires training the newly mined options.

In general, our method differs from existing literature by focusing on abstractions in the context of amortized sampling. We target environments where high-reward states exhibit clear repeating substructures. Moreover, we extract abstractions online using BPE, which incurs negligible cost and facilitates quick adaptation during training, thanks to our action encoder that can generate embeddings for any sequence of actions making up a macro-action, even one not previously seen.