Policy Gradient Guidance Enables Test Time Control
This work addresses the need for controllable online reinforcement learning, offering an incremental extension of existing guidance techniques to standard on-policy methods.
The authors tackled the problem of enabling test-time control in reinforcement learning by introducing Policy Gradient Guidance (PGG), which adapts classifier-free guidance from diffusion models to policy gradient methods, resulting in improved stability, sample efficiency, and controllability in discrete and continuous control benchmarks.
We introduce Policy Gradient Guidance (PGG), a simple extension of classifier-free guidance from diffusion models to classical policy gradient methods. PGG augments the policy gradient with an unconditional branch and interpolates conditional and unconditional branches, yielding a test-time control knob that modulates behavior without retraining. We provide a theoretical derivation showing that the additional normalization term vanishes under advantage estimation, leading to a clean guided policy gradient update. Empirically, we evaluate PGG on discrete and continuous control benchmarks. We find that conditioning dropout-central to diffusion guidance-offers gains in simple discrete tasks and low sample regimes, but dropout destabilizes continuous control. Training with modestly larger guidance ($γ>1$) consistently improves stability, sample efficiency, and controllability. Our results show that guidance, previously confined to diffusion policies, can be adapted to standard on-policy methods, opening new directions for controllable online reinforcement learning.