English

Policy Gradient Guidance Enables Test Time Control

Machine Learning 2025-10-03 v1

Abstract

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\gamma>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.

Keywords

Cite

@article{arxiv.2510.02148,
  title  = {Policy Gradient Guidance Enables Test Time Control},
  author = {Jianing Qi and Hao Tang and Zhigang Zhu},
  journal= {arXiv preprint arXiv:2510.02148},
  year   = {2025}
}
R2 v1 2026-07-01T06:13:31.696Z