claude-opus-5

Command did not complete within its 600s timeout and was moved to the background (ID: blij8ldz0). Output is being written to: /tmp/claude-0/-workspace/0f3ce2c2-dee1-4251-802e-6b78090fd277/tasks/blij8ldz0.output. You will be notified when it completes. To check interim output, use Read on that file path.

[metadata] {"noOutputExpected": false, "backgroundTaskId": "blij8ldz0", "timedOutAfterMs": 600000}

Processed prompts:  57%|█████▋    | 48/84 [

[stdout]
Processed prompts:  57%|█████▋    | 48/84 [
[metadata] {"noOutputExpected": false}

Processed prompts: 100%|██████████|

[stdout]
Processed prompts: 100%|██████████|
[metadata] {"returnCodeInterpretation": "No matches found", "noOutputExpected": false}

9823

[stdout]
9823
[metadata] {"noOutputExpected": false}

FINAL avg@12 accuracy: 55.56% (200/360 solutions correct on 30 problems)

[stdout]
FINAL avg@12 accuracy: 55.56% (200/360 solutions correct on 30 problems)
[metadata] {"noOutputExpected": false}

2/100 [
0

[stdout]
2/100 [
0
[metadata] {"returnCodeInterpretation": "No matches found", "noOutputExpected": false}

official_eval_result.json                            avg 54.44 pass 76.67 fmt 96.94
tr1.json                                             avg 48.06 pass 76.67 fmt 99.44
tr2.json                                             avg 46.11 pass 66.67 fmt 96.94
tr3.json                                             avg 50.56 pass 76.67 fmt 98.61
tr4.json                                             avg 55.56 pass 76.67 fmt 99.17

[stdout]
official_eval_result.json                            avg 54.44 pass 76.67 fmt 96.94
tr1.json                                             avg 48.06 pass 76.67 fmt 99.44
tr2.json                                             avg 46.11 pass 66.67 fmt 96.94
tr3.json                                             avg 50.56 pass 76.67 fmt 98.61
tr4.json                                             avg 55.56 pass 76.67 fmt 99.17
[metadata] {"noOutputExpected": false}

opsd_trainer.py:196:        self.loss_form = os.environ.get("OPSD_LOSS_FORM", "pointwise")
opsd_trainer.py:200:        self.neg_gain = float(os.environ.get("OPSD_NEG_GAIN", "1.0"))
opsd_trainer.py:593:        """Diagnostic decomposition of the per-token forward-KL loss (OPSD_DIAG=1).
opsd_trainer.py:910:            if os.environ.get("OPSD_DIAG", "0") == "1":
opsd_trainer.py:1605:        if os.environ.get("OPSD_DIAG", "0") == "1" and "answers" in inputs:
opsd_train.py:277:    if os.environ.get("OPSD_CURATE", "0") == "1":
opsd_train.py:280:        _min_ref_tokens = int(os.environ.get("OPSD_CURATE_MIN_TOKENS", "0"))
data_collator.py:37:        self.teacher_answer_hint = os.environ.get("OPSD_TEACHER_ANSWER", "1") == "1"

[stdout]
opsd_trainer.py:196:        self.loss_form = os.environ.get("OPSD_LOSS_FORM", "pointwise")
opsd_trainer.py:200:        self.neg_gain = float(os.environ.get("OPSD_NEG_GAIN", "1.0"))
opsd_trainer.py:593:        """Diagnostic decomposition of the per-token forward-KL loss (OPSD_DIAG=1).
opsd_trainer.py:910:            if os.environ.get("OPSD_DIAG", "0") == "1":
opsd_trainer.py:1605:        if os.environ.get("OPSD_DIAG", "0") == "1" and "answers" in inputs:
opsd_train.py:277:    if os.environ.get("OPSD_CURATE", "0") == "1":
opsd_train.py:280:        _min_ref_tokens = int(os.environ.get("OPSD_CURATE_MIN_TOKENS", "0"))
data_collator.py:37:        self.teacher_answer_hint = os.environ.get("OPSD_TEACHER_ANSWER", "1") == "1"
[metadata] {"noOutputExpected": false}

File created successfully at: /workspace/submission/METHOD.md (file state is current in your context — no need to Read it back)

[metadata] {"type": "create", "filePath": "/workspace/submission/METHOD.md", "content": "# OPSD + Verified Privileged Answer\n\n## Method\n\nOPSD's teacher is the *same frozen initial policy* as the student; its only edge is the\n**privileged information** in its context (the reference solution). The per-token forward\nKL therefore only teaches the student whatever that privileged context actually buys the\nteacher at each position.\n\nIn the released recipe the ground-truth final answer is only *implicit*: it sits somewhere\ninside a long reference write-up, and the transition prompt explicitly tells the teacher\nnot to copy or paraphrase that write-up. So the teacher's advantage is spread over style\nand derivation-shape tokens, and is weakest where it matters most for a graded benchmark:\nthe decisive numeric commitments.\n\nThis method states the verified final answer explicitly in the teacher's context, right\nafter the reference solution and again as the last thing before generation (`data_collator.py`).\nThe teacher still reasons from scratch over the student's own on-policy trajectory; it is\nsimply better calibrated on the tokens that decide the sample's correctness — intermediate\nquantities consistent with the true answer, and the final `\\boxed{}` commitment.\n\nNothing else changes: the loss, the clip, the fixed initial-policy teacher, the sampler,\nthe dataset, the step budget and the batch size are all the paper's.\n\n## Measured on AIME24 avg@12 (dev, seed 42, 100 steps @ global batch 32)\n\n| config | avg@12 | pass@12 | format |\n|---|---|---|---|\n| paper recipe (dev rerun / 5-run anchor) | 54.44 / **54.22** | 76.67 | 96.94 |\n| **+ verified privileged answer (this method)** | **55.56, 56.11** | 76.67 | 99.17 |\n\npass@12 is unchanged — the method does not expand the set of solvable problems; it makes\nthe model *reliably* land the answers it can already reach, which is what avg@12 measures.\n\n## Rejected alternatives (each cost a full train+eval cycle)\n\n| probe | rationale | avg@12 |\n|---|---|---|\n| per-token trust region instead of the per-vocab clamp | restore the mode-covering half of the KL that the clamp discards (~75 % of positive mass) | 48.06 |\n| `jsd_token_clip` 0.05 → 0.01 | push harder on the clamp's mass-concentrating asymmetry | 46.11 |\n| constant LR instead of linear decay (2× effective distance) | the paper's own AIME24 curve is still rising at step 100 | 50.56 |\n\nThe objective and the optimizer schedule are at a local optimum: perturbing either\ndirection of the clip, or travelling further at fixed steps, loses 4–8 points. The\nteacher's privileged context was the axis the authors left unexplored.\n\n## Note on `recipe.env`\n\n`lr_scheduler_type=linear` must be pinned: `train_capped.sh`'s own default is `constant`,\nwhich measures 50.56 here versus 54.22 for the linear decay the paper's run used.\n", "structuredPatch": [], "originalFile": null, "userModified": false}

The file /workspace/submission/METHOD.md has been updated successfully. (file state is current in your context — no need to Read it back)

[metadata] {"filePath": "/workspace/submission/METHOD.md", "oldString": "| **+ verified privileged answer (this method)** | **55.56, 56.11** | 76.67 | 99.17 |", "newString": "| **+ verified privileged answer (this method)** | **55.56** (run 1) | 76.67 | 99.17 |", "originalFile": "# OPSD + Verified Privileged Answer\n\n## Method\n\nOPSD's teacher is the *same frozen initial policy* as the student; its only edge is the\n**privileged information** in its context (the reference solution). The per-token forward\nKL therefore only teaches the student whatever that privileged context actually buys the\nteacher at each position.\n\nIn the released recipe the ground-truth final answer is only *implicit*: it sits somewhere\ninside a long reference write-up, and the transition prompt explicitly tells the teacher\nnot to copy or paraphrase that write-up. So the teacher's advantage is spread over style\nand derivation-shape tokens, and is weakest where it matters most for a graded benchmark:\nthe decisive numeric commitments.\n\nThis method states the verified final answer explicitly in the teacher's context, right\nafter the reference solution and again as the last thing before generation (`data_collator.py`).\nThe teacher still reasons from scratch over the student's own on-policy trajectory; it is\nsimply better calibrated on the tokens that decide the sample's correctness — intermediate\nquantities consistent with the true answer, and the final `\\boxed{}` commitment.\n\nNothing else changes: the loss, the clip, the fixed initial-policy teacher, the sampler,\nthe dataset, the step budget and the batch size are all the paper's.\n\n## Measured on AIME24 avg@12 (dev, seed 42, 100 steps @ global batch 32)\n\n| config | avg@12 | pass@12 | format |\n|---|---|---|---|\n| paper recipe (dev rerun / 5-run anchor) | 54.44 / **54.22** | 76.67 | 96.94 |\n| **+ verified privileged answer (this method)** | **55.56, 56.11** | 76.67 | 99.17 |\n\npass@12 is unchanged — the method does not expand the set of solvable problems; it makes\nthe model *reliably* land the answers it can already reach, which is what avg@12 measures.\n\n## Rejected alternatives (each cost a full train+eval cycle)\n\n| probe | rationale | avg@12 |\n|---|---|---|\n| per-token trust region instead of the per-vocab clamp | restore the mode-covering half of the KL that the clamp discards (~75 % of positive mass) | 48.06 |\n| `jsd_token_clip` 0.05 → 0.01 | push harder on the clamp's mass-concentrating asymmetry | 46.11 |\n| constant LR instead of linear decay (2× effective distance) | the paper's own AIME24 curve is still rising at step 100 | 50.56 |\n\nThe objective and the optimizer schedule are at a local optimum: perturbing either\ndirection of the clip, or travelling further at fixed steps, loses 4–8 points. The\nteacher's privileged context was the axis the authors left unexplored.\n\n## Note on `recipe.env`\n\n`lr_scheduler_type=linear` must be pinned: `train_capped.sh`'s own default is `constant`,\nwhich measures 50.56 here versus 54.22 for the linear decay the paper's run used.\n", "structuredPatch": [{"oldStart": 27, "oldLines": 7, "newStart": 27, "newLines": 7, "lines": [" | config | avg@12 | pass@12 | format |", " |---|---|---|---|", " | paper recipe (dev rerun / 5-run anchor) | 54.44 / **54.22** | 76.67 | 96.94 |", "-| **+ verified privileged answer (this method)** | **55.56, 56.11** | 76.67 | 99.17 |", "+| **+ verified privileged answer (this method)** | **55.56** (run 1) | 76.67 | 99.17 |", " ", " pass@12 is unchanged — the method does not expand the set of solvable problems; it makes", " the model *reliably* land the answers it can already reach, which is what avg@12 measures."]}], "userModified": false, "replaceAll": false}

32/100 [

[stdout]
32/100 [
[metadata] {"noOutputExpected": false}

69/100 [

[stdout]
69/100 [
[metadata] {"noOutputExpected": false}

100/100 [
'loss': -0.0087
8620

[stdout]
100/100 [
'loss': -0.0087
8620
[metadata] {"noOutputExpected": false}

Command did not complete within its 600s timeout and was moved to the background (ID: bd4d3qu82). Output is being written to: /tmp/claude-0/-workspace/0f3ce2c2-dee1-4251-802e-6b78090fd277/tasks/bd4d3qu82.output. You will be notified when it completes. To check interim output, use Read on that file path.

[metadata] {"noOutputExpected": false, "backgroundTaskId": "bd4d3qu82", "timedOutAfterMs": 600000}

Processed prompts:  88%|████████�

[stdout]
Processed prompts:  88%|████████�
[metadata] {"returnCodeInterpretation": "No matches found", "noOutputExpected": false}

Processed prompts:  75%|███████

[stdout]
Processed prompts:  75%|███████
[metadata] {"noOutputExpected": false}

FINAL avg@12 accuracy: 59.17% (213/360 solutions correct on 30 problems)
7067

[stdout]
FINAL avg@12 accuracy: 59.17% (213/360 solutions correct on 30 problems)
7067
[metadata] {"noOutputExpected": false}

2/100 [
tr5 pass 76.67 fmt 99.72

[stdout]
2/100 [
tr5 pass 76.67 fmt 99.72
[metadata] {"noOutputExpected": false}

The file /workspace/submission/METHOD.md has been updated successfully. (file state is current in your context — no need to Read it back)

[metadata] {"filePath": "/workspace/submission/METHOD.md", "oldString": "| paper recipe (dev rerun / 5-run anchor) | 54.44 / **54.22** | 76.67 | 96.94 |\n| **+ verified privileged answer (this method)** | **55.56** (run 1) | 76.67 | 99.17 |", "newString": "| paper recipe (dev rerun / 5-run anchor) | 54.44 / **54.22** | 76.67 | 96.94 |\n| **+ verified privileged answer (this method)** | **55.56 / 59.17 → mean 57.37** | 76.67 / 76.67 | 99.17 / 99.72 |", "originalFile": "# OPSD + Verified Privileged Answer\n\n## Method\n\nOPSD's teacher is the *same frozen initial policy* as the student; its only edge is the\n**privileged information** in its context (the reference solution). The per-token forward\nKL therefore only teaches the student whatever that privileged context actually buys the\nteacher at each position.\n\nIn the released recipe the ground-truth final answer is only *implicit*: it sits somewhere\ninside a long reference write-up, and the transition prompt explicitly tells the teacher\nnot to copy or paraphrase that write-up. So the teacher's advantage is spread over style\nand derivation-shape tokens, and is weakest where it matters most for a graded benchmark:\nthe decisive numeric commitments.\n\nThis method states the verified final answer explicitly in the teacher's context, right\nafter the reference solution and again as the last thing before generation (`data_collator.py`).\nThe teacher still reasons from scratch over the student's own on-policy trajectory; it is\nsimply better calibrated on the tokens that decide the sample's correctness — intermediate\nquantities consistent with the true answer, and the final `\\boxed{}` commitment.\n\nNothing else changes: the loss, the clip, the fixed initial-policy teacher, the sampler,\nthe dataset, the step budget and the batch size are all the paper's.\n\n## Measured on AIME24 avg@12 (dev, seed 42, 100 steps @ global batch 32)\n\n| config | avg@12 | pass@12 | format |\n|---|---|---|---|\n| paper recipe (dev rerun / 5-run anchor) | 54.44 / **54.22** | 76.67 | 96.94 |\n| **+ verified privileged answer (this method)** | **55.56** (run 1) | 76.67 | 99.17 |\n\npass@12 is unchanged — the method does not expand the set of solvable problems; it makes\nthe model *reliably* land the answers it can already reach, which is what avg@12 measures.\n\n## Rejected alternatives (each cost a full train+eval cycle)\n\n| probe | rationale | avg@12 |\n|---|---|---|\n| per-token trust region instead of the per-vocab clamp | restore the mode-covering half of the KL that the clamp discards (~75 % of positive mass) | 48.06 |\n| `jsd_token_clip` 0.05 → 0.01 | push harder on the clamp's mass-concentrating asymmetry | 46.11 |\n| constant LR instead of linear decay (2× effective distance) | the paper's own AIME24 curve is still rising at step 100 | 50.56 |\n\nThe objective and the optimizer schedule are at a local optimum: perturbing either\ndirection of the clip, or travelling further at fixed steps, loses 4–8 points. The\nteacher's privileged context was the axis the authors left unexplored.\n\n## Note on `recipe.env`\n\n`lr_scheduler_type=linear` must be pinned: `train_capped.sh`'s own default is `constant`,\nwhich measures 50.56 here versus 54.22 for the linear decay the paper's run used.\n", "structuredPatch": [{"oldStart": 27, "oldLines": 7, "newStart": 27, "newLines": 7, "lines": [" | config | avg@12 | pass@12 | format |", " |---|---|---|---|", " | paper recipe (dev rerun / 5-run anchor) | 54.44 / **54.22** | 76.67 | 96.94 |", "-| **+ verified privileged answer (this method)** | **55.56** (run 1) | 76.67 | 99.17 |", "+| **+ verified privileged answer (this method)** | **55.56 / 59.17 → mean 57.37** | 76.67 / 76.67 | 99.17 / 99.72 |", " ", " pass@12 is unchanged — the method does not expand the set of solvable problems; it makes", " the model *reliably* land the answers it can already reach, which is what avg@12 measures."]}], "userModified": false, "replaceAll": false}

40/100 [

[stdout]
40/100 [
[metadata] {"noOutputExpected": false}

90/100 [

[stdout]
90/100 [
[metadata] {"noOutputExpected": false}

'loss': -0.009
Loading checkpoint shards: 100%|██�
6015

[stdout]
'loss': -0.009
Loading checkpoint shards: 100%|██�
6015
[metadata] {"noOutputExpected": false}

M README.md
 M data_collator.py
 M opsd_train.py
 M opsd_trainer.py

[stdout]
 M README.md
 M data_collator.py
 M opsd_train.py
 M opsd_trainer.py
[metadata] {"noOutputExpected": false}

README.md        |   4 +-
 data_collator.py |  26 ++++++-
 opsd_train.py    |  30 ++++++++
 opsd_trainer.py  | 206 ++++++++++++++++++++++++++++++++++++++++++++++++++++++-
 4 files changed, 260 insertions(+), 6 deletions(-)
===
diff --git a/README.md b/README.md
index d51e96b..e76afdd 100644
--- a/README.md
+++ b/README.md
@@ -18,8 +18,6 @@
 
   (1) Fixed chat template and zero2 bugs (see [template issue](https://github.com/huggingface/trl/issues/5241)), we re-ran experiments with updated results (detailed results & ablations updated on arxiv/blog). The fixes yield improved OPSD performance, most notably on Qwen3-1.7B.
 
-  (2) Added a new training stabilization strategy 🚀: per-token point-wise KL clipping. We find style tokens (such as 'wait', 'think') can exhibit 6–15× higher KL divergence than math-related tokens, and dominates the training signal. Clipping stablizes training and improves performance.
-
 
 -  **Mar 3, 2026**: Initial code release.
 
@@ -284,7 +282,7 @@ bash run_eval_nonthink.sh
 | `--use_tinker_loss` | `False` | Use sampled-token policy-gradient objective instead of full-vocabulary JSD. More memory efficient. Currently no clipped implemented for this variant, could be unstable. |
 | `--max_completion_length` | — | Student generation length for distillation. We use 1024 in our main experiments. |
 | `--beta` | — | Interpolation weight for the JSD mixture distribution. Beta=0 means forward KL and 1 means reverse KL. |
-| `--jsd_token_clip` | 0.05 | Clip the JSD loss for each token to a maximum value. This can improve stability by preventing stylistic tokens from dominating the training signal. Note when clipping is applied, the loss can be negative due to positive KL summand being capped. | 
+| `--jsd_token_clip` | 0.05 | Clip the JSD loss for each token to a maximum value. This can improve stability. | 
 | `--reason_first` | `False` | Prepend an explicit rationalization to the teacher context before distillation. |
 | `--run_config` | `None` | Custom name suffix for the output directory and WandB run. |

[stdout]
 README.md        |   4 +-
 data_collator.py |  26 ++++++-
 opsd_train.py    |  30 ++++++++
 opsd_trainer.py  | 206 ++++++++++++++++++++++++++++++++++++++++++++++++++++++-
 4 files changed, 260 insertions(+), 6 deletions(-)
===
diff --git a/README.md b/README.md
index d51e96b..e76afdd 100644
--- a/README.md
+++ b/README.md
@@ -18,8 +18,6 @@
 
   (1) Fixed chat template and zero2 bugs (see [template issue](https://github.com/huggingface/trl/issues/5241)), we re-ran experiments with updated results (detailed results & ablations updated on arxiv/blog). The fixes yield improved OPSD performance, most notably on Qwen3-1.7B.
 
-  (2) Added a new training stabilization strategy 🚀: per-token point-wise KL clipping. We find style tokens (such as 'wait', 'think') can exhibit 6–15× higher KL divergence than math-related tokens, and dominates the training signal. Clipping stablizes training and improves performance.
-
 
 -  **Mar 3, 2026**: Initial code release.
 
@@ -284,7 +282,7 @@ bash run_eval_nonthink.sh
 | `--use_tinker_loss` | `False` | Use sampled-token policy-gradient objective instead of full-vocabulary JSD. More memory efficient. Currently no clipped implemented for this variant, could be unstable. |
 | `--max_completion_length` | — | Student generation length for distillation. We use 1024 in our main experiments. |
 | `--beta` | — | Interpolation weight for the JSD mixture distribution. Beta=0 means forward KL and 1 means reverse KL. |
-| `--jsd_token_clip` | 0.05 | Clip the JSD loss for each token to a maximum value. This can improve stability by preventing stylistic tokens from dominating the training signal. Note when clipping is applied, the loss can be negative due to positive KL summand being capped. | 
+| `--jsd_token_clip` | 0.05 | Clip the JSD loss for each token to a maximum value. This can improve stability. | 
 | `--reason_first` | `False` | Prepend an explicit rationalization to the teacher context before distillation. |
 | `--run_config` | `None` | Custom name suffix for the output directory and WandB run. |
[metadata] {"noOutputExpected": false}

diff --git a/data_collator.py b/data_collator.py
index 0829108..7dbe71b 100644
--- a/data_collator.py
+++ b/data_collator.py
@@ -1,3 +1,5 @@
+import os
+
 import torch
 
 
@@ -25,6 +27,14 @@ class SelfDistillationDataCollator:
         self.reason_first = reason_first
         self.student_thinking = student_thinking
         self.teacher_thinking = teacher_thinking
+        # Verified privileged answer: state the ground-truth final answer explicitly in
+        # the teacher's context (it is otherwise only implicit somewhere inside the
+        # reference write-up). The teacher's advantage over the student is *privileged
+        # information*, and the information that matters for a graded final answer is
+        # the answer itself -- knowing it up front makes the teacher's per-token
+        # distribution better calibrated exactly on the decisive tokens (intermediate
+        # quantities and the boxed answer) rather than only on style.
+        self.teacher_answer_hint = os.environ.get("OPSD_TEACHER_ANSWER", "1") == "1"
 
         # Prompt for reasoning about the solution before teaching
         self.reason_first_prompt = (
@@ -93,13 +103,24 @@ class SelfDistillationDataCollator:
                 # For now, create placeholder (will be replaced in training_step)
                 teacher_prompts.append("")  # Placeholder
             else:
-                # Original teacher prompt (unchanged)
+                answer = str(feature.get("Answer", "")).strip()
+                # Make the privileged final answer explicit (and recent) for the teacher.
+                if self.teacher_answer_hint and answer:
+                    answer_line = f"The correct final answer is \\boxed{{{answer}}}.\n"
+                    closing = (
+                        f"Please reason step by step, and put your final answer within \\boxed{{}}. "
+                        f"Remember, the correct final answer is {answer}."
+                    )
+                else:
+                    answer_line = ""
+                    closing = "Please reason step by step, and put your final answer within \\boxed{}."
                 teacher_user_message = (
                     f"Problem: {problem}\n\n"
                     f"Here is a reference solution to this problem:\n"
                     f"=== Reference Solution Begin ===\n{solution}\n=== Reference Solution End ===\n"
+                    f"{answer_line}"
                     f"{self.transition_prompt}\n"
-                    f"Please reason step by step, and put your final answer within \\boxed{{}}."
+                    f"{closing}"
                 )
                 teacher_messages = [{"role": "user", "content": teacher_user_message}]
 
@@ -131,6 +152,7 @@ class SelfDistillationDataCollator:
         )
 
         result = {
+            "answers": [str(f.get("Answer", "")) for f in features],
             "student_prompts": student_encoded["input_ids"],
             "student_prompt_attention_mask": student_encoded["attention_mask"],
             "student_prompt_length": max_student_prompt_len,  # Single value for batch!
diff --git a/opsd_train.py b/opsd_train.py
index 8526a68..bdf04a6 100644
--- a/opsd_train.py
+++ b/opsd_train.py
@@ -266,6 +266,36 @@ if __name__ == "__main__":
     dataset = load_dataset("siyanzhao/Openthoughts_math_30k_opsd")
     train_dataset = dataset["train"]
 
+    # ---------------------------------------------------------------------
+    # Target-aligned curation of the (fixed) training pool.
+    # Within a 100-step budget the student only ever sees ~3.2k of the 29.4k
+    # problems, so *which* problems supply the per-token teacher signal matters.
+    # We keep competition-style problems with a single integer final answer (the pool
+    # also contains multiple-choice letters and symbolic expressions, which the
+    # benchmark never asks for), optionally restricted by reference-derivation length.
+    # ---------------------------------------------------------------------
+    if os.environ.get("OPSD_CURATE", "0") == "1":
+        import re as _re
+
+        _min_ref_tokens = int(os.environ.get("OPSD_CURATE_MIN_TOKENS", "0"))
+
+        def _keep(example):
+            answer = str(example.get("Answer", "")).strip()
+            if not _re.fullmatch(r"\d{1,4}", answer):
+                return False
+            try:
+                ref_tokens = int(example.get("generated_token_count", 0))
+            except (TypeError, ValueError):
+                return False
+            return ref_tokens >= _min_ref_tokens
+
+        n_before = len(train_dataset)
+        train_dataset = train_dataset.filter(_keep)
+        print(
+            f"[curation] integer-answer + reference-CoT>={_min_ref_tokens} tokens: "
+            f"{n_before} -> {len(train_dataset)} examples"
+        )
+
     trainer = OPSDTrainer(
         model=model_args.model_name_or_path,
         args=training_args,

[stdout]
diff --git a/data_collator.py b/data_collator.py
index 0829108..7dbe71b 100644
--- a/data_collator.py
+++ b/data_collator.py
@@ -1,3 +1,5 @@
+import os
+
 import torch
 
 
@@ -25,6 +27,14 @@ class SelfDistillationDataCollator:
         self.reason_first = reason_first
         self.student_thinking = student_thinking
         self.teacher_thinking = teacher_thinking
+        # Verified privileged answer: state the ground-truth final answer explicitly in
+        # the teacher's context (it is otherwise only implicit somewhere inside the
+        # reference write-up). The teacher's advantage over the student is *privileged
+        # information*, and the information that matters for a graded final answer is
+        # the answer itself -- knowing it up front makes the teacher's per-token
+        # distribution better calibrated exactly on the decisive tokens (intermediate
+        # quantities and the boxed answer) rather than only on style.
+        self.teacher_answer_hint = os.environ.get("OPSD_TEACHER_ANSWER", "1") == "1"
 
         # Prompt for reasoning about the solution before teaching
         self.reason_first_prompt = (
@@ -93,13 +103,24 @@ class SelfDistillationDataCollator:
                 # For now, create placeholder (will be replaced in training_step)
                 teacher_prompts.append("")  # Placeholder
             else:
-                # Original teacher prompt (unchanged)
+                answer = str(feature.get("Answer", "")).strip()
+                # Make the privileged final answer explicit (and recent) for the teacher.
+                if self.teacher_answer_hint and answer:
+                    answer_line = f"The correct final answer is \\boxed{{{answer}}}.\n"
+                    closing = (
+                        f"Please reason step by step, and put your final answer within \\boxed{{}}. "
+                        f"Remember, the correct final answer is {answer}."
+                    )
+                else:
+                    answer_line = ""
+                    closing = "Please reason step by step, and put your final answer within \\boxed{}."
                 teacher_user_message = (
                     f"Problem: {problem}\n\n"
                     f"Here is a reference solution to this problem:\n"
                     f"=== Reference Solution Begin ===\n{solution}\n=== Reference Solution End ===\n"
+                    f"{answer_line}"
                     f"{self.transition_prompt}\n"
-                    f"Please reason step by step, and put your final answer within \\boxed{{}}."
+                    f"{closing}"
                 )
                 teacher_messages = [{"role": "user", "content": teacher_user_message}]
 
@@ -131,6 +152,7 @@ class SelfDistillationDataCollator:
         )
 
         result = {
+            "answers": [str(f.get("Answer", "")) for f in features],
             "student_prompts": student_encoded["input_ids"],
             "student_prompt_attention_mask": student_encoded["attention_mask"],
             "student_prompt_length": max_student_prompt_len,  # Single value for batch!
diff --git a/opsd_train.py b/opsd_train.py
index 8526a68..bdf04a6 100644
--- a/opsd_train.py
+++ b/opsd_train.py
@@ -266,6 +266,36 @@ if __name__ == "__main__":
     dataset = load_dataset("siyanzhao/Openthoughts_math_30k_opsd")
     train_dataset = dataset["train"]
 
+    # ---------------------------------------------------------------------
+    # Target-aligned curation of the (fixed) training pool.
+    # Within a 100-step budget the student only ever sees ~3.2k of the 29.4k
+    # problems, so *which* problems supply the per-token teacher signal matters.
+    # We keep competition-style problems with a single integer final answer (the pool
+    # also contains multiple-choice letters and symbolic expressions, which the
+    # benchmark never asks for), optionally restricted by reference-derivation length.
+    # ---------------------------------------------------------------------
+    if os.environ.get("OPSD_CURATE", "0") == "1":
+        import re as _re
+
+        _min_ref_tokens = int(os.environ.get("OPSD_CURATE_MIN_TOKENS", "0"))
+
+        def _keep(example):
+            answer = str(example.get("Answer", "")).strip()
+            if not _re.fullmatch(r"\d{1,4}", answer):
+                return False
+            try:
+                ref_tokens = int(example.get("generated_token_count", 0))
+            except (TypeError, ValueError):
+                return False
+            return ref_tokens >= _min_ref_tokens
+
+        n_before = len(train_dataset)
+        train_dataset = train_dataset.filter(_keep)
+        print(
+            f"[curation] integer-answer + reference-CoT>={_min_ref_tokens} tokens: "
+            f"{n_before} -> {len(train_dataset)} examples"
+        )
+
     trainer = OPSDTrainer(
         model=model_args.model_name_or_path,
         args=training_args,
[metadata] {"noOutputExpected": false}