Distilled Reinforcement Learning for LLM Post-Training

Standard reinforcement learning relies on coarse-grained outcome rewards, while on-policy distillation usually encourages the student to imitate the teacher distribution unconditionally. Distilled RL instead uses the teacher to redistribute the policy-gradient signal at the token level, providing selective and fine-grained guidance while preserving reward-driven optimization.

Overview

Distilled RL consists of three components:

1. Reverse importance sampling, which measures the teacher's relative preference for each student-generated token.
2. Negative sample reset, which disables teacher reweighting on negative-advantage trajectories.
3. Sequence-level geometric normalization, which removes sequence-level scale bias while preserving relative token preferences.

Method

Given a prompt $q$ and a response $o_i$ sampled from the old student policy, the standard policy ratio is

$$r_{i,t}(\theta )=\frac{{\pi }_{\theta }(o_{i,t}\mid q,o_{i,1:t-1})}{{\pi }_{\mathrm{o}\mathrm{l}\mathrm{d}}(o_{i,t}\mid q,o_{i,1:t-1})}r\_\{i,t\}(\backslash theta)\; =\; \backslash frac\{\; \backslash pi\_\{\backslash theta\}(o\_\{i,t\}\; \backslash mid\; q,\; o\_\{i,1:t-1\})\; \}\{\; \backslash pi\_\{\backslash mathrm\{old\}\}(o\_\{i,t\}\; \backslash mid\; q,\; o\_\{i,1:t-1\})\; \}$$

The response-level advantage is estimated using group-normalized rewards:

$$A_i=\frac{R_i-\mathrm{m}\mathrm{e}\mathrm{a}\mathrm{n}(\{R_j{\}}_{j=1}^G)}{\mathrm{s}\mathrm{t}\mathrm{d}(\{R_j{\}}_{j=1}^G)}\mathrm{.}A\_i\; =\; \backslash frac\{\; R\_i\; -\; \backslash mathrm\{mean\}(\backslash \{R\_j\backslash \}\_\{j=1\}^\{G\})\; \}\{\; \backslash mathrm\{std\}(\backslash \{R\_j\backslash \}\_\{j=1\}^\{G\})\; \}.$$

Reverse Importance Sampling

We measure the teacher's relative preference for each student-generated token using

$${\rho }_{i,t}=\frac{{\pi }_{\mathrm{t}\mathrm{e}\mathrm{a}\mathrm{c}\mathrm{h}\mathrm{e}\mathrm{r}}(o_{i,t}\mid q,o_{i,1:t-1})}{{\pi }_{{\theta }_{\mathrm{o}\mathrm{l}\mathrm{d}}}(o_{i,t}\mid q,o_{i,1:t-1})}\mathrm{.}\backslash rho\_\{i,t\}\; =\; \backslash frac\{\; \backslash pi\_\{\backslash mathrm\{teacher\}\}(o\_\{i,t\}\; \backslash mid\; q,\; o\_\{i,1:t-1\})\; \}\{\; \backslash pi\_\{\backslash theta\_\{\backslash mathrm\{old\}\}\}(o\_\{i,t\}\; \backslash mid\; q,\; o\_\{i,1:t-1\})\; \}.$$

To prevent extreme teacher–student likelihood ratios, we apply symmetric clipping:

$${\bar{\rho }}_{i,t}=\mathrm{c}\mathrm{l}\mathrm{i}\mathrm{p}({\rho }_{i,t},{ϵ}_{\rho }^{-1},{ϵ}_{\rho })\mathrm{.}\backslash bar\{\backslash rho\}\_\{i,t\}\; =\; \backslash mathrm\{clip\}\; \backslash left(\; \backslash rho\_\{i,t\},\; \backslash epsilon\_\{\backslash rho\}^\{-1\},\; \backslash epsilon\_\{\backslash rho\}\; \backslash right).$$

Sequence-Level Geometric Normalization

The clipped ratios are normalized within each response:

$${\tilde{\rho }}_{i,t}=\frac{{\bar{\rho }}_{i,t}}{\exp (\frac{1}{\mathrm{\mid }{o}_i\mathrm{\mid }}\sum _{s=1}^{\mathrm{\mid }{o}_i\mathrm{\mid }}\log {\bar{\rho }}_{i,s})}\mathrm{.}\backslash widetilde\{\backslash rho\}\_\{i,t\}\; =\; \backslash frac\{\; \backslash bar\{\backslash rho\}\_\{i,t\}\; \}\{\; \backslash exp\; \backslash left(\; \backslash frac\{1\}\{|o\_i|\}\; \backslash sum\_\{s=1\}^\{|o\_i|\}\; \backslash log\; \backslash bar\{\backslash rho\}\_\{i,s\}\; \backslash right)\; \}.$$

The normalized ratios satisfy

$${\left(\prod _{t=1}^{\mathrm{\mid }{o}_i\mathrm{\mid }}{\tilde{\rho }}_{i,t}\right)}^{1\mathrm{/}\mathrm{\mid }{o}_i\mathrm{\mid }}=1.\backslash left(\; \backslash prod\_\{t=1\}^\{|o\_i|\}\; \backslash widetilde\{\backslash rho\}\_\{i,t\}\; \backslash right)^\{1/|o\_i|\}\; =\; 1.$$

This normalization removes the sequence-level mean shift in log importance ratios while preserving the teacher's relative preferences across tokens.

Negative Sample Reset

Teacher guidance is applied only to positive-advantage responses:

For negative-advantage responses, the update reduces to the original RL objective.

Distilled RL Objective

For responses sampled from the old student policy, the final policy optimization objective is

$${\mathcal{J}}_{\mathrm{D}\mathrm{i}\mathrm{s}\mathrm{t}\mathrm{i}\mathrm{l}\mathrm{l}\mathrm{e}\mathrm{d}\mathrm{R}\mathrm{L}}(\theta )=\mathbb{E}[\frac{1}{G}\sum _{i=1}^G\frac{1}{\mathrm{\mid }{o}_i\mathrm{\mid }}\sum _{t=1}^{\mathrm{\mid }{o}_i\mathrm{\mid }}\min (r_{i,t}(\theta )w_{i,t}{A}_i,{\widehat{r}}_{i,t}(\theta )w_{i,t}{A}_i)],\backslash mathcal\{J\}\_\{\backslash mathrm\{DistilledRL\}\}(\backslash theta)\; =\; \backslash mathbb\{E\}\; \backslash left[\; \backslash frac\{1\}\{G\}\; \backslash sum\_\{i=1\}^\{G\}\; \backslash frac\{1\}\{|o\_i|\}\; \backslash sum\_\{t=1\}^\{|o\_i|\}\; \backslash min\; \backslash left(\; r\_\{i,t\}(\backslash theta)\; w\_\{i,t\}\; A\_i,\; \backslash hat\{r\}\_\{i,t\}(\backslash theta)\; w\_\{i,t\}\; A\_i\; \backslash right)\; \backslash right],$$

where the clipped policy ratio is

$${\widehat{r}}_{i,t}(\theta )=\mathrm{c}\mathrm{l}\mathrm{i}\mathrm{p}(r_{i,t}(\theta ),1-{ϵ}_{\mathrm{l}\mathrm{o}\mathrm{w}},1+{ϵ}_{\mathrm{h}\mathrm{i}\mathrm{g}\mathrm{h}})\mathrm{.}\backslash hat\{r\}\_\{i,t\}(\backslash theta)\; =\; \backslash mathrm\{clip\}\; \backslash left(\; r\_\{i,t\}(\backslash theta),\; 1-\backslash epsilon\_\{\backslash mathrm\{low\}\},\; 1+\backslash epsilon\_\{\backslash mathrm\{high\}\}\; \backslash right).$$

Unlike KL-based on-policy distillation, Distilled RL does not treat the teacher as an unconditional imitation target. Instead, the teacher selectively redistributes the reward-driven policy-gradient signal at the token level.

Main Results

We evaluate Distilled RL on three student models using Qwen3-8B-GRPO as the teacher. The table below reports the average Pass@1 over ten mathematical reasoning benchmarks.

Distilled RL consistently improves over standard RL, OPD, and their direct combination across different student scales and teacher–student settings.

Requirements

Software

Clone the repository:

git clone https://github.com/597358816/Distilled-RL.git
cd Distilled-RL

Install the required dependencies:

pip install torch==2.6.0 torchaudio==2.6.0 torchvision==0.21.0 vllm==0.8.3 transformers==4.51.2
pip install ray==2.48.0 tensordict==0.9.1 pydantic==2.11.7
pip install flash-attn
pip install -e .
pip install tensorboard
cd examples
bash XX.sh