Proximal Policy Optimization for Amortized Discrete Sampling
Abstract
This paper explores policy gradient algorithms for training stochastic policies to sample from structured discrete probability distributions under the Generative Flow Network (GFlowNet) framework. Building on extensive theoretical connections between GFlowNets and entropy-regularized reinforcement learning, we derive equivalents of standard policy gradient algorithms for training GFlowNets, as well as experimentally explore their various methodological aspects, including baseline training and advantage estimation. Most importantly, our work is the first to derive and successfully apply proximal policy optimization to GFlowNets, showing its improved convergence speed and data efficiency compared to standard GFlowNet training objectives on benchmarks ranging from synthetic energies to molecular graph generation.
Related
- Weak Convexity and Proximal Bundle Methods for Nonsmooth Policy Optimization in Robust Control
- dFlowGRPO: Rate-Aware Policy Optimization for Discrete Flow Models
- ProximalFM: Amortized Proximal Causal Inference under Hidden Confounding
- Amortized Score-Hamiltonian Policy Iteration: A Grid-Free Scheme for Relaxed Stochastic Control Problems
- Amortized ratio-estimation importance sampling and localized simulation-based calibration for intractable likelihoods
BibTeX
@misc{zykovamyzina2026proximal,
title = {{Proximal Policy Optimization for Amortized Discrete Sampling}},
author = {Anna Zykova-Myzina and Timofei Gritsaev and Daniil Tiapkin and Nikita Morozov},
howpublished = {arXiv},
year = {2026},
month = nov,
doi = {10.48550/arxiv.2606.15793},
eprint = {2606.15793},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2606.15793},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - Proximal Policy Optimization for Amortized Discrete Sampling
AU - Anna Zykova-Myzina
AU - Timofei Gritsaev
AU - Daniil Tiapkin
AU - Nikita Morozov
PY - 2026
DA - 2026/11/02
PB - arXiv
DO - 10.48550/arxiv.2606.15793
UR - https://arxiv.org/abs/2606.15793
AB - This paper explores policy gradient algorithms for training stochastic policies to sample from structured discrete probability distributions under the Generative Flow Network (GFlowNet) framework. Building on extensive theoretical connections between GFlowNets and entropy-regularized reinforcement learning, we derive equivalents of standard policy gradient algorithms for training GFlowNets, as well as experimentally explore their various methodological aspects, including baseline training and advantage estimation. Most importantly, our work is the first to derive and successfully apply proximal policy optimization to GFlowNets, showing its improved convergence speed and data efficiency compared to standard GFlowNet training objectives on benchmarks ranging from synthetic energies to molecular graph generation.
ER -