Learning from Language Feedback via Variational Policy Distillation

2026-05-14

AI Summary

This work addresses the exploration bottleneck in reinforcement learning caused by sparse rewards and the difficulty of leveraging language feedback in complex reasoning tasks. The authors propose Variational Policy Distillation (VPD), a framework that formulates language-feedback-driven learning as a variational expectation-maximization (EM) problem. VPD jointly optimizes teacher and student policies to dynamically generate and internalize token-level dense supervision signals. Moving beyond conventional passive distillation, the method introduces an adaptive trust-region update mechanism and integrates on-policy rollouts to enable continuous evolution of the teacher policy as the student improves. Experiments demonstrate that VPD significantly outperforms standard RLVR and existing self-distillation approaches on scientific reasoning and code generation tasks, while also revealing performance limits of feedback-driven distillation in mathematical reasoning and cold-start scenarios.

Abstract

Reinforcement learning from verifiable rewards (RLVR) suffers from sparse outcome signals, creating severe exploration bottlenecks on complex reasoning tasks. Recent on-policy self-distillation methods attempt to address this by utilizing language feedback to generate dense, token-level supervision. However, these approaches rely on a fixed, passive teacher to interpret the feedback. As the student policy improves, the teacher's zero-shot assessment capabilities plateau, ultimately halting further learning. To overcome this, we propose Variational Policy Distillation (VPD), a framework that formalizes learning from language feedback as a Variational Expectation-Maximization (EM) problem. VPD co-evolves both policies: in the E-step, the teacher is actively refined on trajectory outcomes via an adaptive trust-region update, translating textual feedback into a dynamically improved target token distribution. In the M-step, the student internalizes this dense distributional guidance on its own on-policy rollouts. By continuously improving the teacher's ability to extract actionable signals from textual critique, VPD overcomes the limitations of passive distillation. Evaluated across diverse sources of diagnostic feedback on scientific reasoning and code generation tasks, VPD consistently outperforms both standard RLVR and existing self-distillation baselines. Finally, by stress-testing our framework on rigid mathematical reasoning and cold-start regimes, we illuminate the fundamental bounds of feedback-driven self-distillation compared to pure environment-driven RL.

Problem

• reinforcement learning
• language feedback
• sparse rewards
• policy distillation
• exploration bottleneck

Innovation

• Variational Policy Distillation
• Language Feedback
• Reinforcement Learning from Verifiable Rewards
• Self-Distillation
• Variational EM