{"id":51489894,"url":"https://github.com/BWGZK-keke/BRTS","last_synced_at":"2026-07-26T06:00:37.788Z","repository":{"id":357505871,"uuid":"1234207967","full_name":"BWGZK-keke/BRTS","owner":"BWGZK-keke","description":"code for our paper Best of N teacher rollouts for on-policy distillation","archived":false,"fork":false,"pushed_at":"2026-05-13T03:02:43.000Z","size":2526,"stargazers_count":1,"open_issues_count":0,"forks_count":0,"subscribers_count":0,"default_branch":"main","last_synced_at":"2026-05-13T05:06:55.276Z","etag":null,"topics":[],"latest_commit_sha":null,"homepage":null,"language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":null,"status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/BWGZK-keke.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":null,"code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null,"notice":null,"maintainers":null,"copyright":null,"agents":null,"dco":null,"cla":null}},"created_at":"2026-05-09T22:12:52.000Z","updated_at":"2026-05-13T03:02:47.000Z","dependencies_parsed_at":null,"dependency_job_id":null,"html_url":"https://github.com/BWGZK-keke/BRTS","commit_stats":null,"previous_names":["bwgzk-keke/brts"],"tags_count":null,"template":false,"template_full_name":null,"purl":"pkg:github/BWGZK-keke/BRTS","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/BWGZK-keke%2FBRTS","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/BWGZK-keke%2FBRTS/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/BWGZK-keke%2FBRTS/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/BWGZK-keke%2FBRTS/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/BWGZK-keke","download_url":"https://codeload.github.com/BWGZK-keke/BRTS/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/BWGZK-keke%2FBRTS/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":286080680,"owners_count":35902828,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2026-07-20T02:08:10.276Z","status":"online","status_checked_at":"2026-07-26T02:00:06.503Z","response_time":89,"last_error":null,"robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":true,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":[],"created_at":"2026-07-07T11:00:22.090Z","updated_at":"2026-07-26T06:00:37.778Z","avatar_url":"https://github.com/BWGZK-keke.png","language":"Python","funding_links":[],"categories":["🔬 OPD with Larger External Teachers — White-Box"],"sub_categories":[],"readme":"\u003cdiv align=\"center\"\u003e\n\n# On-Policy Distillation with Best-of-N Teacher Rollout Selection\n\n\u003c/div\u003e\n\n\u003cdiv align=\"center\" style=\"font-family: Arial, sans-serif;\"\u003e\n  \u003cp\u003e\n    \u003ca href=\"#overview\" style=\"text-decoration: none; font-weight: bold;\"\u003e📖 Overview\u003c/a\u003e •\n    \u003ca href=\"#getting-started\" style=\"text-decoration: none; font-weight: bold;\"\u003e✨ Getting Started\u003c/a\u003e •\n    \u003ca href=\"#repository-layout\" style=\"text-decoration: none; font-weight: bold;\"\u003e📁 Repository Layout\u003c/a\u003e •\n    \u003ca href=\"#citation\" style=\"text-decoration: none; font-weight: bold;\"\u003e🎈 Citation\u003c/a\u003e\n  \u003c/p\u003e\n\u003c/div\u003e\n\n---\n\n## 📖Overview\n\nThis repository contains the official implementation of **BRTS** (**B**est-of-N **R**ollout **T**eacher **S**election), a lightweight extension to on-policy distillation (OPD) that improves how teacher trajectories are selected and used as supervision.\n\nStandard OPD supervises the student on its own sampled trajectories using a single stochastic teacher rollout per prompt. As a result, the supervision signal can be high-variance: the sampled teacher trajectory can be incorrect, uninformative, or poorly matched to the student's current reasoning behavior. BRTS addresses this with a **correctness- and alignment-aware** teacher-context branch:\n\n1. **Tier 1 — Multi-rollout sampling.** For each prompt, draw `N` unconditioned teacher rollouts. If at least one is correct, select the correct rollout with the highest top-K token-level overlap with the student.\n2. **Tier 2 — Ground-truth-guided recovery.** If all Tier-1 rollouts are wrong, re-prompt the teacher with the ground-truth answer silently injected as a validation signal, and retain the resulting rollout only if its extracted answer is correct.\n3. **Tier 3 — Fallback.** If neither tier yields a correct rollout, fall back to the Tier-1 rollout with the highest student overlap.\n\nThe selected trajectory is then used in an auxiliary teacher-context distillation loss alongside the standard student-context OPD loss:\n\n```\nL_total = L_student-context + λ · L_teacher-context\n```\n\nExperiments on AIME 2024, AIME 2025, and AMC 2023 show that BRTS improves over standard OPD, with the largest gains on harder benchmarks where reliable teacher trajectories are sparse.\n\n## ✨Getting Started\n\n### Environment Setup\n\nThis codebase is built on top of [verl](https://github.com/volcengine/verl). To prepare the environment:\n\n```bash\nconda create -n verl python==3.12\nconda activate verl\ncd verl/\nUSE_MEGATRON=0 bash scripts/install_vllm_sglang_mcore.sh\npip install math-verify\n```\n\nSet Hugging Face credentials in your environment if your models require authentication:\n\n```bash\nexport HF_TOKEN=\u003cyour_token\u003e\n# Optional: point HF_HOME at a shared cache directory\nexport HF_HOME=/path/to/hf_cache\n```\n\n### Training\n\n#### BRTS (Tier-1 only)\n\nMulti-rollout teacher sampling with overlap-based selection, no ground-truth-guided recovery:\n\n```bash\nbash forward.sh\n```\n\nThis launches training with the following BRTS settings:\n\n| Variable | Value | Description |\n|----------|-------|-------------|\n| `TEACHER_ROLLOUT_MULTI_FORWARD_SELECT_BY_OVERLAP` | `True` | Enable Best-of-N selection by top-K overlap |\n| `TEACHER_ROLLOUT_TIER1_ONLY` | `True` | Disable Tier-2 ground-truth-guided recovery |\n| `TEACHER_ROLLOUT_N_ROLLOUTS_HINT` | `0` | No hinted retries |\n\n#### BRTS (Tier-1 + Tier-2)\n\nAdds ground-truth-guided recovery for prompts where all Tier-1 rollouts fail:\n\n```bash\nbash forward_tier2.sh\n```\n\nThis sets `TEACHER_ROLLOUT_TIER1_ONLY=False` and `TEACHER_ROLLOUT_N_ROLLOUTS_HINT=1`.\n\n\u003cdetails\u003e\n\u003csummary\u003e\u003cb\u003eKey Parameters (override via env vars)\u003c/b\u003e\u003c/summary\u003e\n\n| Parameter | Default | Description |\n|-----------|---------|-------------|\n| **Models** |||\n| `ACTOR_MODEL_PATH` | `deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B` | Path or HF repo id of the student (policy) model |\n| `REWARD_MODEL_PATH` | `model/JustRL-DeepSeek-1.5B` (or actor model fallback) | Path of the teacher model that provides token-level distillation signals |\n| **Distillation** |||\n| `ADV_ESTIMATOR` | `token_reward_direct` | Required for OPD; do not change |\n| `LOG_PROB_TOP_K` | `16` | Number of top-K tokens retained per position; `0` falls back to sampled-token OPD |\n| `TOP_K_STRATEGY` | `only_stu` | Top-K candidate set source: `only_stu`, `only_tch`, `intersection`, `union`, `union-intersection` |\n| `REWARD_WEIGHT_MODE` | `student_p` | Token-reward weighting: `student_p`, `teacher_p`, `none` |\n| **Generation** |||\n| `N_RESPONSES` | `2` | Student rollouts per prompt |\n| `MAX_PROMPT_LENGTH` | `1024` | Max prompt length |\n| `MAX_RESP_LENGTH` | `7168` | Max student response length during training |\n| `MAX_VAL_RESP_LENGTH` | `31744` | Max response length during validation |\n| **BRTS — Teacher Rollout** |||\n| `TEACHER_ROLLOUT_ENABLE` | `True` | Enable teacher rollouts for distillation |\n| `N_TEACHER_ROLLOUTS` | `2` | Number of Tier-1 unconditioned teacher candidates |\n| `TEACHER_ROLLOUT_GEN_TEMPERATURE` | `0.7` | Teacher sampling temperature |\n| `TEACHER_ROLLOUT_TOP_P` | `0.95` | Teacher nucleus sampling |\n| `TEACHER_ROLLOUT_TIER1_ONLY` | `False` | Disable Tier-2 hinted recovery |\n| `TEACHER_ROLLOUT_N_ROLLOUTS_HINT` | `1` | Number of Tier-2 ground-truth-guided rollouts |\n| `TEACHER_ROLLOUT_FALLBACK` | `most_similar` | Tier-3 fallback: `most_similar` or `skip` |\n| `TEACHER_ROLLOUT_MULTI_FORWARD_SELECT_BY_OVERLAP` | `False` | Run teacher forward on both student and selected contexts and pick by overlap |\n| **BRTS — Auxiliary Teacher-Context KD** |||\n| `AUX_TEACHER_CTX_KD_ENABLE` | `False` | Enable auxiliary teacher-context distillation loss |\n| `AUX_TEACHER_CTX_KD_COEF` | `0.05` | Coefficient λ on the auxiliary teacher-context loss |\n| `AUX_TEACHER_CTX_KD_WEIGHT_MODE` | `student_p` | Weighting for the auxiliary loss |\n| **BRTS — Prompt Perturbation Ablation** |||\n| `TEACHER_PROMPT_DISTURB_ENABLE` | `False` | Append a perturbation string to one teacher rollout's prompt |\n| `TEACHER_PROMPT_DISTURB_ROLLOUT_INDEX` | `2` | Which rollout (1-indexed) to perturb |\n| `TEACHER_PROMPT_DISTURB_TEXT` | `\"Please reason step by step and rethink in detail before giving the final answer.\"` | Perturbation text |\n\n\u003c/details\u003e\n\n\u003e [!IMPORTANT]\n\u003e **Non-thinking Models:** When training a non-thinking model (e.g. `Qwen3-1.7B (Non-thinking)`), add `+data.apply_chat_template_kwargs.enable_thinking=False` to the training script.\n\n### Validation\n\nThe evaluation pipeline is adapted from [JustRL](https://github.com/RyanLiu112/JustRL).\n\n**Generation (optional):**\n\n```bash\ncd scripts/val/eval\npython gen_vllm.py\n```\n\nBefore running generation, set `MODEL_NAMES` in `gen_vllm.py` to the checkpoint(s) you want to evaluate, and set an appropriate `available_workers`.\n\n**Grading:**\n\n```bash\ncd scripts/val/eval\npython grade.py\n```\n\nThe grading script processes all JSONL files in the output directory and produces `grading_results.json`. Enable the LLM-based verifier with:\n\n```bash\npython grade.py --enable_model_verifier\n```\n\n*All main experiments were conducted on 8 × NVIDIA B200 GPUs.*\n\n## 📁Repository Layout\n\n```\n.\n├── forward.sh                 # BRTS Tier-1-only launcher\n├── forward_tier2.sh           # BRTS Tier-1 + Tier-2 launcher\n├── on_policy_distillation.sh  # Main OPD/BRTS training script\n├── grpo.sh                    # GRPO RL baseline\n├── scripts/\n│   ├── fetch_val_acc.py       # Swanlab log → accuracy table utility\n│   ├── infer/\n│   │   ├── vllm_rollout.py    # Teacher rollout generator\n│   │   └── dedup_deepmath.py  # Deduplicate DeepMath against DAPO-Math-17k\n│   └── val/eval/\n│       ├── gen_vllm.py        # Validation generation\n│       ├── grade.py           # Validation grading (with optional LLM verifier)\n│       └── utils.py\n└── verl/                      # verl framework with BRTS extensions\n    └── verl/trainer/ppo/\n        ├── teacher_selection.py     # ★ Tier-1/2/3 teacher rollout selection\n        ├── rollout_corr_helper.py   # ★ Correctness checking \u0026 overlap utilities\n        ├── ray_trainer.py           # PPO trainer (with BRTS hooks)\n        └── ...\n```\n\nThe core BRTS implementation lives in `verl/verl/trainer/ppo/teacher_selection.py`, which implements the 3-tier waterfall (Tier-1 unconditioned sampling → Tier-2 ground-truth-guided recovery → Tier-3 most-similar fallback).\n\n## 🎈Citation\n\nIf you find this work helpful, please cite us:\n\n```bibtex\n@misc{zhang2026onpolicydistillationbestofnteacher,\n      title={On-Policy Distillation with Best-of-N Teacher Rollout Selection}, \n      author={Ke Zhang and Yunjie Tian and DongDi Zhao and Yijiang Li and Yuanye Liu and Vishal M Patel and Di Fu},\n      year={2026},\n      eprint={2605.09725},\n      archivePrefix={arXiv},\n      primaryClass={cs.CV},\n      url={https://arxiv.org/abs/2605.09725}, \n}\n```\n\n## Acknowledgements\n\nThis codebase is built on top of [verl](https://github.com/volcengine/verl), and the evaluation pipeline is adapted from [JustRL](https://github.com/RyanLiu112/JustRL). We thank the authors of these projects for their excellent open-source work.\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2FBWGZK-keke%2FBRTS","html_url":"https://awesome.ecosyste.ms/projects/github.com%2FBWGZK-keke%2FBRTS","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2FBWGZK-keke%2FBRTS/lists"}