Research map / Reinforcement learning / Policy gradient and actor-critic
Adaptive teachers for amortized samplers
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes a Teacher GFlowNet that adaptively samples high-loss regions of a Student amortized sampler (trained with trajectory balance) to improve mode coverage in off-policy training.
Problem
Exploration and mode coverage are poor when training amortized samplers (GFlowNets) with off-policy RL, since existing replay and perturbation methods only revisit modes near the current policy's samples.
Approach
An auxiliary Teacher GFlowNet is trained with a reward derived from the Student's trajectory-balance loss (weighted toward undersampled high-reward regions and mixed with the log reward), and its samples drive the Student's training, amortizing an ideal prioritized replay over the whole sample space.
Key results
- On the deceptive grid (d=2, H=256), Teacher discovers 2452.6 modes with L1 0.94e-5 versus 2165.2 modes for PRT and 2055.2 for PER.
- On Manywell (d=32), Teacher reaches ELBO 163.484 and EUBO 165.800, outperforming PER (161.537 and 210.440).
- On 25GMM, Teacher achieves W2^2 0.86 versus 1.87 for PER and 4.64 for on-policy TB.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Train the Student with an adaptive Teacher | TB sampler · Deceptive grid world (d=2, H=256) | # modes: 2165.2 -> 2452.6 | Helped |
| Train the Student with an adaptive Teacher | TB sampler · Deceptive grid world (d=4, H=32) | # modes: 120.4 -> 246.6 | Helped |
| Train the Student with an adaptive Teacher | TB sampler · 25GMM | EUBO: 1.833{ $\pm$2.366} -> 0.115{ $\pm$0.009} | Helped |
| Train the Student with an adaptive Teacher | TB sampler · Manywell | EUBO: 210.440{ $\pm$6.888} -> 165.800{ $\pm$0.045} | Helped |
| Use N_MC=3 instead of N_MC=1 to estimate the Teacher reward (ablation) | Teacher · Deceptive grid world (d=2, H=256) | # modes: $2452.6 \pm \text{ 21.7} -> $2489.6 \pm \text{ 15.8} | No clear effect |
| Set C=0 instead of C=19 (ablation) | Teacher · Deceptive grid world (d=2, H=256) | L1 error: $0.94 \pm \text{ 0.03} -> $1.06 \pm \text{ 0.08} | Hurt |
| Set C=0 instead of C=19 (ablation) | Teacher · Deceptive grid world | # modes | Mixed |
| Set α=0.5 instead of α=0.0 (ablation) | Teacher · Deceptive grid world (d=4, H=32) | # modes: $246.6 \pm \text{ 14.7} -> $85.6 \pm \text{ 8.5} | Hurt |
| Set α=0.5 instead of α=0.0 (ablation) | Teacher · 25GMM | ELBO: -0.144{ $\pm$0.001} -> -0.137{ $\pm$0.004} | Helped |
Surprising findings
- Reward mixing with alpha=0.5 degraded mode discovery on the high-dimensional deceptive grid (85.6 vs 246.6 modes at d=4, H=32), so pure high-loss targeting worked better there.
Limitations
- Added training complexity from training a second (teacher) policy network.
- In larger search spaces the teacher may collapse onto a subset of modes.
- Gains are task-dependent; unimodal or low-exploration tasks may not need a teacher.
Experimental setup
Datasets: QM9 sEH TFbind8 L14-RNA1 25GMM Manywell deceptive grid world
Benchmarks: deceptive grid world 25GMM Manywell QM9 sEH TFbind8 L14-RNA1
Metrics: number of modes discovered L1 distance ELBO ELBO-IS EUBO 2-Wasserstein distance
Baselines: on-policy TB epsilon-exploration GAFN PRT PER SMC GGNS PIS
Builds on: trajectory balance prioritized experience replay detailed balance GFlowNet local search GAFN
Closest papers by meaning
- Improved off-policy training of diffusion samplers (NeurIPS 2024 · similarity 0.86)
- Action abstractions for amortized sampling (ICLR 2025 · similarity 0.86)
- ALINE: Joint Amortization for Bayesian Inference and Active Data Acquisition (NeurIPS 2025 · similarity 0.86)
- Learning to Sample From Diffusion Models via Inverse Reinforcement Learning (NeurIPS 2026 · similarity 0.85)
- Efficient Adaptive Data Acquisition via Pretrained Belief Representations (NeurIPS 2026 · similarity 0.85)
- Amortized Active Causal Induction with Deep Reinforcement Learning (NeurIPS 2024 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.