[Feat] Add on-policy distillation support by HwVanICI · Pull Request #964 · inclusionAI/AReaL

HwVanICI · 2026-03-02T23:02:57Z

Description

This PR adds support for on-policy knowledge distillation (KD) in AReaL, enabling a student model to learn from a teacher while sampling trajectories from its own policy, thereby reducing distribution mismatch. It also includes combined KD + RL training, allowing the student to imitate the teacher while simultaneously optimizing reward through reinforcement learning.

Related Issue

#960

Type of Change

Bug fix (non-breaking change that fixes an issue)
New feature (non-breaking change that adds functionality)
Breaking change (fix or feature that would cause existing functionality to not
work as expected)
Documentation update
Code refactoring (no functional changes)
Performance improvement
Test coverage improvement

Checklist

I have read the Contributing Guide
I have run formatting tools (pre-commit or manual)
I have run relevant unit tests and they pass
I have added tests for new functionality
I have updated documentation if needed
My branch is up to date with main
This PR introduces breaking changes (if yes, fill out details below)
If this PR changes documentation, I have built and previewed it locally with
jb build docs
No critical issues raised by AI reviewers (/gemini review)

Breaking Change Details (if applicable):

Additional Context

Need help? Check the Contributing Guide or ask in
GitHub Discussions!

gemini-code-assist · 2026-03-02T23:03:19Z

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances AReaL's capabilities by introducing on-policy knowledge distillation (KD) and a combined KD + Reinforcement Learning (KDRL) framework. This allows student models to learn from a teacher's policy while actively exploring with their own, reducing distribution mismatch and improving training stability. The changes involve new configuration options for teacher models, modifications to the PPO loss function to incorporate distillation, and updates to the RL trainer to manage the teacher model and its contributions to the learning process.

Highlights

On-Policy Knowledge Distillation (KD) Support: Introduced support for on-policy knowledge distillation, allowing a student model to learn from a teacher by sampling trajectories from its own policy, which helps reduce distribution mismatch.
Combined KD + Reinforcement Learning (KDRL): Enabled a combined training approach where the student model can simultaneously imitate the teacher and optimize for reward through reinforcement learning.
New Teacher Configuration: Added a TeacherConfig dataclass to cli_args.py for defining teacher model parameters, including rl_loss_weight and distill_loss_weight.
PPO Loss Function Enhancement: Modified the PPO actor's grpo_loss_fn to incorporate both pure reverse KL distillation and a weighted combination of RL and distillation losses.
RL Trainer Integration: Integrated the teacher model initialization and log-probability computation into the RLTrainer to facilitate on-policy distillation during the training loop.
Distillation Example and Documentation: Provided a new example configuration (gsm8k_grpo_distill.yaml) and a comprehensive README.md to demonstrate and explain the on-policy distillation feature.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

areal/api/cli_args.py
- Added TeacherConfig dataclass, inheriting from PPOActorConfig, with fields for allocation_mode, rl_loss_weight, and distill_loss_weight.
- Added an optional teacher: TeacherConfig field to PPOConfig.
areal/trainer/ppo/actor.py
- Modified grpo_loss_fn to conditionally calculate and apply a joint distillation KL loss.
- Implemented logic for pure reverse KL distillation when rl_loss_weight is zero.
- Implemented KDRL by combining RL loss with a forward KL distillation term when rl_loss_weight is non-zero.
- Added teacher_kl_loss and teacher_kl scalars to stats_tracker.
areal/trainer/rl_trainer.py
- Added self.teacher attribute, initialized to None.
- Included logic in __init__ to create and initialize the teacher model if config.teacher is provided.
- Modified the train method to compute teacher_logp and pass rl_loss_weight and distill_loss_weight to the rollout batch if a teacher is present.
- Added a new private method _create_teacher to handle the instantiation and setup of the teacher model based on its configuration.
examples/distillation/README.md
- Added a new markdown file explaining on-policy distillation, its idea, how to run an example, and results.
- Included YAML configuration snippets and a bash command for running the example.
- Provided a reference to the KDRL paper.
examples/distillation/gsm8k_grpo_distill.yaml
- Added a new YAML configuration file for running a GSM8K GRPO distillation experiment.
- Defined teacher configuration parameters, including allocation_mode, rl_loss_weight, distill_loss_weight, and model path.

Activity

No specific activity (comments, reviews, progress updates) was provided in the pull request context.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request introduces on-policy knowledge distillation, which is a valuable feature. The implementation looks mostly correct, but I've identified a few areas for improvement. There's a significant code duplication in rl_trainer.py when creating the teacher model, which also contains a bug that will cause failures in SPMD mode. Additionally, there's an inconsistency in the documentation and comments regarding the type of KL divergence used in the loss function in ppo/actor.py. Finally, the new README file contains a broken link to a reference paper. Addressing these points will improve the maintainability and correctness of the new feature.

areal/trainer/rl_trainer.py

areal/trainer/ppo/actor.py

examples/distillation/README.md

ZiyiTsang · 2026-03-03T11:06:27Z

Actually, I would recommend using a rollout engine combined with a custom workflow to address this issue, as no training is required for teacher...?
If FSDP/Megatron is used, it will create optimizers and gradients, which could end up consuming more memory instead....

HwVanICI · 2026-03-03T17:58:56Z

/gemini review

gemini-code-assist

Code Review

The pull request successfully adds support for on-policy knowledge distillation (KD) and combined KD + RL training. The new TeacherConfig dataclass is well-integrated into the PPOConfig, and the distillation loss logic is correctly implemented in areal/trainer/ppo/actor.py. The _create_teacher method in areal/trainer/rl_trainer.py handles the teacher model's initialization, and the provided README.md and example YAML configuration offer clear guidance for using the new feature. The overall implementation is clean and aligns with the stated objective of the PR.

_{Note: Security Review did not run due to the size of the PR.}

areal/trainer/ppo/actor.py

areal/trainer/rl_trainer.py

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

gemini-code-assist

Code Review

This pull request introduces on-policy knowledge distillation, which is a valuable addition. The implementation is mostly solid, but I've identified a few areas for improvement in areal/trainer/rl_trainer.py concerning code consistency and maintainability. Specifically, the creation and initialization of the teacher component could be refactored to align with how other components like actor and critic are handled, which would also resolve a bug where a backend is missing. Additionally, the new README.md for distillation contains a potentially confusing explanation of the KL divergence that could be clarified for better understanding.

areal/trainer/rl_trainer.py

examples/distillation/README.md

HwVanICI · 2026-03-03T19:41:32Z

/gemini review

gemini-code-assist

Code Review

This pull request introduces on-policy knowledge distillation, which is a great feature. The implementation adds a TeacherConfig, integrates a teacher model into the PPOTrainer, and implements the distillation loss.

My review has identified a few points for improvement:

There appears to be an issue in the pure knowledge distillation loss calculation where a tensor that should be detached is not, potentially leading to incorrect gradients.
There is some code duplication in the creation of the teacher engine, which could be refactored for better maintainability.
The new README.md has a slightly confusing explanation regarding loss minimization that could be clarified.

Overall, the changes are well-structured. Addressing these points will help improve the correctness and maintainability of the new feature.

areal/trainer/ppo/actor.py

gemini-code-assist · 2026-03-03T19:44:02Z

docs/en/algorithms/distillation.md

+where the reward $R_{i,t} = \log \pi_T(o_{i,t}) - \log \pi_\theta(o_{i,t})$. This encourages the student to increase the probability of tokens the teacher prefers and suppress those it deems unlikely.
+
+- Implementation Detail:
+During pure KD, we need to set `rl_loss_weight` to 0, so the implementation estimates the RKL gradient using importance sampling. The code calculates the reward as teacher_logp - logprobs ($R_{i,t}$) and applies a negative coefficient to the loss to perform minimization (check `areal/trainer/ppo/actor.py`).


The phrase "applies a negative coefficient to the loss to perform minimization" is slightly misleading. In gradient-based optimization, applying a negative coefficient to a loss term corresponds to maximizing the objective that the term represents (in this case, the expected reward). For better clarity, you could rephrase this to state that the negative coefficient is used to perform gradient ascent on the reward objective. For example: "...and applies a negative coefficient to the loss, effectively performing gradient ascent to maximize the reward."

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

garrett4wade

LGTM but somes nits

areal/api/cli_args.py

areal/trainer/ppo/actor.py

examples/distillation/gsm8k_grpo_distill.yaml

docs/en/algorithms/distillation.md

areal/trainer/ppo/actor.py

xssstory · 2026-03-04T06:37:30Z

areal/trainer/ppo/actor.py

+            teacher_kl = teacher_logp - logprobs.detach()
+            prob_ratio = torch.exp(logprobs - old_logp)
+
+            loss_per_token = prob_ratio * teacher_kl * loss_mask


Why is prob_ratio required here, yet it's ignored in line 460?

garrett4wade

LGTM except for a nit

garrett4wade · 2026-03-05T12:34:12Z

areal/trainer/ppo/actor.py

+            stats_tracker.denominator(token=loss_mask)
+            stats_tracker.stat(
+                rkl_loss=-1 * rkl_weighted_term,
+                denominator="token",
+            )


nit: L479 has logged n_valid_tokens, so you should remove the denominator log and defer stat logging after L479.

garrett4wade

LGTM

* [feat] Add on-policy knowledge distillation support * Correct file extension * Delete comment on distill_loss_weight Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * explain RKL and Joint loss * Add reference to global README * teacher config refactoring * refactor: remove redundant stats_tracker denominator logging for KD loss --------- Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

HwVanICI added 2 commits March 2, 2026 13:55

[feat] Add on-policy knowledge distillation support

c97df15

Correct file extension

d980069

gemini-code-assist bot reviewed Mar 2, 2026

View reviewed changes

areal/trainer/rl_trainer.py Outdated Show resolved Hide resolved

areal/trainer/ppo/actor.py Show resolved Hide resolved

examples/distillation/README.md Outdated Show resolved Hide resolved

SPMD mode comment

c37d918

format problem solved

2df99b9

gemini-code-assist bot reviewed Mar 3, 2026

View reviewed changes

areal/trainer/ppo/actor.py Outdated Show resolved Hide resolved

areal/trainer/rl_trainer.py Outdated Show resolved Hide resolved

Delete comment on distill_loss_weight

eadd411

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

gemini-code-assist bot reviewed Mar 3, 2026

View reviewed changes

areal/trainer/rl_trainer.py Show resolved Hide resolved

areal/trainer/rl_trainer.py Show resolved Hide resolved

examples/distillation/README.md Outdated Show resolved Hide resolved

explain RKL and Joint loss

061138f

gemini-code-assist bot reviewed Mar 3, 2026

View reviewed changes

inclusionAI deleted a comment from gemini-code-assist bot Mar 3, 2026

HwVanICI and others added 2 commits March 3, 2026 12:22

detach logprobs during pure KD

b654749

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Merge branch 'main' into distillation

a65cc54

HwVanICI marked this pull request as ready for review March 3, 2026 20:25

garrett4wade reviewed Mar 4, 2026

View reviewed changes

areal/api/cli_args.py Outdated Show resolved Hide resolved

areal/trainer/ppo/actor.py Outdated Show resolved Hide resolved

examples/distillation/gsm8k_grpo_distill.yaml Outdated Show resolved Hide resolved

docs/en/algorithms/distillation.md Show resolved Hide resolved

xssstory reviewed Mar 4, 2026

View reviewed changes

areal/trainer/ppo/actor.py Outdated Show resolved Hide resolved

xssstory reviewed Mar 4, 2026

View reviewed changes

HwVanICI added 4 commits March 4, 2026 14:11

Merge branch 'main' into distillation

ac82f73

refactored

5e24f76

Add reference to global README

89fe9e5

teacher config refactoring

f5b9c9c

garrett4wade reviewed Mar 5, 2026

View reviewed changes

HwVanICI added 2 commits March 5, 2026 15:23

Merge branch 'main' into distillation

56f66bf

refactor: remove redundant stats_tracker denominator logging for KD loss

748671e

garrett4wade approved these changes Mar 6, 2026

View reviewed changes

garrett4wade merged commit d1cdac3 into inclusionAI:main Mar 6, 2026
5 checks passed

This was referenced Mar 6, 2026

[Feature] On-Policy Distillation Support #960

Closed

[Question] Supporting RL training of other model families #945

Open

Conversation

HwVanICI commented Mar 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Related Issue

Type of Change

Checklist

Additional Context

Uh oh!

gemini-code-assist bot commented Mar 2, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ZiyiTsang commented Mar 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

HwVanICI commented Mar 3, 2026

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

HwVanICI commented Mar 3, 2026

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

gemini-code-assist bot Mar 3, 2026

Choose a reason for hiding this comment

Uh oh!

garrett4wade left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

xssstory Mar 4, 2026

Choose a reason for hiding this comment

Uh oh!

garrett4wade left a comment

Choose a reason for hiding this comment

Uh oh!

garrett4wade Mar 5, 2026

Choose a reason for hiding this comment

Uh oh!

garrett4wade left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

HwVanICI commented Mar 2, 2026 •

edited

Loading

ZiyiTsang commented Mar 3, 2026 •

edited

Loading