Reproduce “Stabilizing Policy Gradients for Sample-Efficient Reinforcement Learning in LLM Reasoning” on OpenResearch

About this paper

Luckeciano C. Melo, Alessandro Abate, Yarin Gal

Researchers at the University of Oxford developed Curvature-Aware Policy Optimization (CAPO) to stabilize policy gradient methods in reinforcement learning for large language model reasoning, preventing catastrophic policy collapse. CAPO demonstrated up to a 30-fold improvement in sample efficiency on mathematical reasoning tasks by using a novel, tractable second-order geometric model for data selection.

Key insights

• Second-order geometric information, even when approximated by modeling only the last layer of an LLM, effectively captures critical aspects of the optimization landscape for policy gradients, providing indicators for instability.
• Monitoring policy shifts (via FIM) and objective shifts (via Hessian) allows for anticipating and preventing abrupt, undesirable updates that lead to catastrophic policy collapse.
• Integrating curvature awareness into a data selection mechanism provides a sample-efficient and stable learning regime with minimal computational overhead, offering a principled alternative to heuristic stability improvements.

Summary by alphaXiv.