feat(grpo): add RLSD self-distilled advantage reweighting

[qychen2001]

PR type

• Bug Fix
• New Feature
• Document Updates
• More Models or Datasets Support

PR information

Introduce RLSD (Self-Distilled RLVR) token-level advantage reweighting for GRPO, enabled via --advantage_reweight rlsd. The per-sequence advantage is redistributed inside each trajectory using the teacher-vs-student logprob gap as a sign-aware, clipped multiplicative weight, without flipping the reward sign. The teacher reuses the OPSD self-distillation forward (current policy conditioned on a privileged teacher_prompt column), and a frozen local teacher via --teacher_model is also supported.

• add apply_rlsd_reweight in rl_core/advantage.py
• add RLSD args and lambda warmup/decay schedule to the GRPO trainer
• add _check_rlsd fail-fast validation in RLHFArguments
• add unit tests in tests/train/test_rlsd_reweight.py

src: https://github.com/iie-ycx/RLSD/tree/main
paper: https://arxiv.org/abs/2604.03128

[gemini-code-assist]

Code Review

This pull request implements RLSD (Self-Distilled RLVR) token-level advantage reweighting for GRPO. It introduces the validation checks, configuration arguments, the core reweighting algorithm, trainer integration, and comprehensive unit tests. The review feedback highlights a critical bug where NaN/inf values in padded positions can propagate and corrupt the training loss, and a medium issue where missing teacher logprobs can cause a silent fallback to standard GRPO instead of raising an error.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Critical Bug: If teacher_per_token_logps or policy_per_token_logps contains NaN or inf (e.g., -inf in padded positions), the subtraction teacher_per_token_logps - policy_per_token_logps will produce NaN or inf. Multiplying by mask (which is 0 at padded positions) still results in NaN because NaN * 0 = NaN in PyTorch. This NaN then propagates through torch.exp, torch.clamp, and the final multiplication, causing the entire returned advantage tensor (and consequently the training loss) to become NaN. To prevent this, use torch.where to mask the logprob difference before any exponential or clamp operations. Additionally, cast seq_neg to reweight.dtype instead of .float() to support mixed-precision/half-precision training seamlessly.

Suggested change

	mask = completion_mask
	base_bt = base_advantages.unsqueeze(1).expand_as(mask)
	delta = (teacher_per_token_logps.detach() - policy_per_token_logps.detach()) * mask
	sign_a = torch.sign(base_bt)
	weights = torch.exp(sign_a * delta) * mask
	clipped = torch.clamp(weights, min=1.0 - clip_range, max=1.0 + clip_range)
	reweight = (1.0 - lam) + lam * clipped
	if negative_only:
	seq_neg = (base_advantages < 0).float().unsqueeze(1)
	reweight = seq_neg * reweight + (1.0 - seq_neg)
	return base_bt * reweight.detach() * mask
	mask = completion_mask.to(teacher_per_token_logps.dtype)
	base_bt = base_advantages.unsqueeze(1).expand_as(mask)
	raw_delta = teacher_per_token_logps.detach() - policy_per_token_logps.detach()
	delta = torch.where(mask.bool(), raw_delta, torch.zeros_like(raw_delta))
	sign_a = torch.sign(base_bt)
	weights = torch.exp(sign_a * delta) * mask
	clipped = torch.clamp(weights, min=1.0 - clip_range, max=1.0 + clip_range)
	reweight = (1.0 - lam) + lam * clipped
	if negative_only:
	seq_neg = (base_advantages < 0).to(reweight.dtype).unsqueeze(1)
	reweight = seq_neg * reweight + (1.0 - seq_neg)
	return base_bt * reweight.detach() * mask

[gemini-code-assist]

Medium Issue: If the user explicitly configures --advantage_reweight rlsd but forgets to provide a teacher_prompt column in the dataset or fails to configure a teacher, grpo_batch.teacher_per_token_logps will be None. In this case, use_rlsd silently evaluates to False, and the trainer falls back to standard GRPO without any warning or error. This silent failure makes it very hard for users to debug why RLSD is not active. We should raise a clear ValueError if RLSD is requested but the required teacher logprobs are missing.

            if self.advantage_reweight == 'rlsd':
                if grpo_batch.teacher_per_token_logps is None:
                    raise ValueError(
                        "advantage_reweight='rlsd' is enabled, but teacher_per_token_logps is None. "
                        "Please ensure that your dataset contains a 'teacher_prompt' column and that "
                        "the teacher model or self-distillation is correctly configured."
                    )
                use_rlsd = True
            else:
                use_rlsd = False
            if use_rlsd: