added k3 KL estimator

docs/en/advanced/miles_server_args.md

@@ -218,6 +218,7 @@ Arguments for reinforcement learning algorithms and loss calculation.
 | Argument | Description | Default | Options | Source |
 | :--- | :--- | :--- | :--- | :--- |
 | `--advantage-estimator` | Advantage estimator to use. | `"grpo"` | `grpo`, `gspo`, `ppo`, `reinforce_plus_plus`, `reinforce_plus_plus_baseline`, `on_policy_distillation` | Miles Native |
+| `--opd-reward-type` | Token-level reward for on-policy distillation. `logr` is `log(p_teacher/p_student)`; `k3` is the negative k3 KL estimator `1 + log(r) - r`. | `"logr"` | `logr`, `k3` | Miles Native |
 | `--loss-type` | Type of loss function to use. | `"policy_loss"` | `policy_loss`, `sft_loss`, `custom_loss` | Miles Native |
 | `--custom-loss-function-path` | Path to a custom loss calculation function (requires `--loss-type custom_loss`). [Ref](../get_started/customization.md#9-custom-loss-function---custom-loss-function-path) | `None` | Type: str | Miles Native |
 | `--critic-lr` | Learning rate for the Critic. Defaults to `--lr`. | `None` | Type: float | Miles Native |

examples/on_policy_distillation/README.md

@@ -35,6 +35,11 @@ PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \
 bash examples/on_policy_distillation/run-qwen3-8B-opd.sh
 ```
 
+The default token-level reward is `logr = log(p_teacher / p_student)`. Pass
+`--opd-reward-type k3` to instead use the negative k3 KL estimator
+`1 + log(r) - r`, where `r = p_teacher / p_student`. It has the same expected
+negative forward-KL objective while using the k3 control variate.
+
 
 # Preliminary Results
 Using Qwen3-8B-Base model sfted on part of the [OpenThoughts3-1.2M](https://huggingface.co/datasets/open-thoughts/OpenThoughts3-1.2M) dataset, we performed on-policy distillation with a Qwen3-32B teacher on the remaining data. Evaluation on Math500 shows:
@@ -56,4 +61,4 @@ The teacher runs on an independent SGLang server that miles treats as a reward m
 # References
 1. https://thinkingmachines.ai/blog/on-policy-distillation/
 2. https://arxiv.org/abs/2306.13649
-3. https://arxiv.org/abs/2306.08543
+3. https://arxiv.org/abs/2306.08543

miles/backends/training_utils/loss.py

@@ -11,6 +11,7 @@
     calculate_log_probs_and_entropy,
     compute_approx_kl,
     compute_gspo_kl,
+    compute_opd_reward,
     compute_opsm_mask,
     compute_policy_loss,
     get_advantages_and_returns_batch,
@@ -380,7 +381,7 @@ def compute_advantages_and_returns(args: Namespace, rollout_data: RolloutBatch)
             for t_log_prob, response_length in zip(teacher_log_probs, response_lengths, strict=False)
         ]
         advantages = [
-            teacher_log_prob - student_log_prob
+            compute_opd_reward(student_log_prob, teacher_log_prob, getattr(args, "opd_reward_type", "logr"))
             for teacher_log_prob, student_log_prob in zip(teacher_log_probs, student_log_probs, strict=False)
         ]
         returns = advantages

miles/utils/arguments.py

@@ -909,6 +909,17 @@ def add_algo_arguments(parser):
                 ],
                 default="grpo",
             )
+            parser.add_argument(
+                "--opd-reward-type",
+                type=str,
+                choices=["logr", "k3"],
+                default="logr",
+                help=(
+                    "Token-level reward for on_policy_distillation. logr uses log(p_teacher/p_student); "
+                    "k3 uses its lower-variance negative-KL counterpart "
+                    "1 + log(p_teacher/p_student) - p_teacher/p_student."
+                ),
+            )
             parser.add_argument(
                 "--disable-compute-advantages-and-returns",
                 action="store_false",

miles/utils/ppo_utils.py

@@ -51,6 +51,28 @@ def compute_approx_kl(
     return kl
 
 
+def compute_opd_reward(
+    student_log_prob: torch.Tensor,
+    teacher_log_prob: torch.Tensor,
+    reward_type: str,
+) -> torch.Tensor:
+    """Compute the token-level on-policy distillation reward.
+
+    Let ``r = p_teacher / p_student``. ``logr`` is the original OPD reward
+    ``log(r)``. ``k3`` uses the negative of Schulman's non-negative k3 KL
+    estimator, ``-(r - 1 - log(r))``, because this value is used as a reward
+    and OPD should minimize ``KL(p_student || p_teacher)``.
+    """
+    log_r = teacher_log_prob.float() - student_log_prob.float()
+
+    if reward_type == "logr":
+        return log_r
+    if reward_type == "k3":
+        # 1 + log(r) - r, written with expm1 for accuracy near r == 1.
+        return log_r - torch.expm1(log_r)
+    raise ValueError(f"Unknown OPD reward type: {reward_type}")
+
+
 def compute_opsm_mask(
     args: Namespace,
     full_log_probs: list[torch.Tensor],

tests/fast/utils/test_arguments.py

@@ -12,6 +12,16 @@
 REQUIRED_ARGS = ["--rollout-batch-size", "64"]
 
 
+@pytest.mark.parametrize("reward_type", ["logr", "k3"])
+def test_opd_reward_type_is_parsed(reward_type: str) -> None:
+    with patch.object(sys, "argv", ["test", "--opd-reward-type", reward_type] + REQUIRED_ARGS):
+        parser = argparse.ArgumentParser()
+        get_miles_extra_args_provider()(parser)
+        args, _ = parser.parse_known_args()
+
+    assert args.opd_reward_type == reward_type
+
+
 def make_class_with_add_arguments():
     class MyFn:
         @classmethod

tests/fast/utils/test_ppo_utils.py

@@ -0,0 +1,38 @@
+import pytest
+import torch
+
+from miles.utils.ppo_utils import compute_opd_reward
+
+
+def test_compute_opd_reward_logr() -> None:
+    student = torch.tensor([-2.0, -1.0])
+    teacher = torch.tensor([-1.5, -1.25])
+
+    actual = compute_opd_reward(student, teacher, "logr")
+
+    torch.testing.assert_close(actual, torch.tensor([0.5, -0.25]))
+
+
+def test_compute_opd_reward_k3_is_negative_kl_estimate() -> None:
+    student = torch.tensor([-2.0, -1.0])
+    teacher = torch.tensor([-1.5, -1.25])
+    log_r = teacher - student
+    expected = 1 + log_r - torch.exp(log_r)
+
+    actual = compute_opd_reward(student, teacher, "k3")
+
+    torch.testing.assert_close(actual, expected)
+    assert torch.all(actual <= 0)
+
+
+def test_compute_opd_reward_k3_is_zero_when_distributions_match() -> None:
+    log_probs = torch.tensor([-10.0, -1.0, 0.0])
+
+    actual = compute_opd_reward(log_probs, log_probs, "k3")
+
+    torch.testing.assert_close(actual, torch.zeros_like(log_probs))
+
+
+def test_compute_opd_reward_rejects_unknown_type() -> None:
+    with pytest.raises(ValueError, match="Unknown OPD reward type"):
+        compute_opd_reward(torch.tensor([0.0]), torch.tensor([0.0]), "unknown")