Action abstractions for amortized sampling
Oussama Boussif, Léna Néhale Ezzine, Joseph D. Viviano, Michal Koziarski, Moksh Jain, Nikolay Malkin, Emmanuel Bengio, Rim Assouel, Yoshua Bengio
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 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.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Cited by top-tier papers2
- Scalable and Cost-Efficient de Novo Template-Based Molecular GenerationPiotr Gainski, Oussama Boussif, Andrei Rekesh, Dmytro Shevchuk et al.NeurIPS 2025 · 6 citations
- Path-dependent Discrete Amortized InferenceTiago Silva, Esmeralda S. Whitammer, Salem LahlouICML 2026
Builds on15
- How Attentive are Graph Attention Networks?Shaked Brody, Uri Alon, Eran YahavICLR 2022 · 1,717 citations
- Flow Network based Generative Models for Non-Iterative Diverse Candidate GenerationEmmanuel Bengio, Moksh Jain, Maksym Korablyov, Doina Precup et al.NeurIPS 2021 · 565 citations
- Path Integral Sampler: A Stochastic Control Approach For SamplingQinsheng Zhang, Yongxin ChenICLR 2022 · 177 citations
- Learning GFlowNets From Partial Episodes For Improved Convergence And StabilityKanika Madan, Jarrid Rector-Brooks, Maksym Korablyov, Emmanuel Bengio et al.ICML 2023 · 138 citations
- Learning Compositional Rules via Neural Program SynthesisMaxwell I. Nye, Armando Solar-Lezama, Josh Tenenbaum, Brenden M. LakeNeurIPS 2020 · 120 citations
Related papers
- Better Training of GFlowNets with Local Credit and Incomplete TrajectoriesLing Pan, Nikolay Malkin, Dinghuai Zhang, Yoshua BengioICML 2023 · 100 citations
- Pre-Training and Fine-Tuning Generative Flow NetworksLing Pan, Moksh Jain, Kanika Madan, Yoshua BengioICLR 2024 · 24 citations
- Avoid What You Know: Divergent Trajectory Balance for GFlowNetsPedro Dall’Antonia, Tiago Silva, Daniel Csillag, Salem Lahlou et al.ICML 2026 · 2 citations
- Optimizing Backward Policies in GFlowNets via Trajectory Likelihood MaximizationTimofei Gritsaev, Nikita Morozov, Sergey Samsonov, Daniil TiapkinICLR 2025
- Flow Factorization for Efficient Generative Flow NetworksJiashun Liu, Chunhui Li, Cheng-Hao Liu, Dianbo Liu et al.AAAI 2025
