{"id":50510387,"url":"https://github.com/hhh675597/revisiting_opd","last_synced_at":"2026-06-19T14:00:33.655Z","repository":{"id":345525094,"uuid":"1132059038","full_name":"hhh675597/revisiting_opd","owner":"hhh675597","description":null,"archived":false,"fork":false,"pushed_at":"2026-04-11T02:27:38.000Z","size":38465,"stargazers_count":25,"open_issues_count":0,"forks_count":1,"subscribers_count":0,"default_branch":"main","last_synced_at":"2026-04-11T04:24:37.713Z","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":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/hhh675597.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":".github/CODEOWNERS","security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null,"notice":"Notice.txt","maintainers":null,"copyright":null,"agents":null,"dco":null,"cla":null}},"created_at":"2026-01-11T08:33:57.000Z","updated_at":"2026-04-11T02:27:43.000Z","dependencies_parsed_at":null,"dependency_job_id":null,"html_url":"https://github.com/hhh675597/revisiting_opd","commit_stats":null,"previous_names":["hhh675597/revisiting_opd"],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/hhh675597/revisiting_opd","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/hhh675597%2Frevisiting_opd","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/hhh675597%2Frevisiting_opd/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/hhh675597%2Frevisiting_opd/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/hhh675597%2Frevisiting_opd/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/hhh675597","download_url":"https://codeload.github.com/hhh675597/revisiting_opd/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/hhh675597%2Frevisiting_opd/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":286080680,"owners_count":34534278,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2026-05-26T15:22:16.424Z","status":"online","status_checked_at":"2026-06-19T02:00:06.005Z","response_time":61,"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-06-02T20:00:26.251Z","updated_at":"2026-06-19T14:00:33.648Z","avatar_url":"https://github.com/hhh675597.png","language":"Python","funding_links":[],"categories":["📚 Surveys, Foundations \u0026 Position Papers"],"sub_categories":[],"readme":"# Revisiting On-Policy Distillation: Empirical Failure Modes and Simple Fixes\n\n- 📝 Blog: [Revisiting On-Policy Distillation: Empirical Failure Modes and Simple Fixes](https://notion.so/yuqianfu/Revisiting-On-Policy-Distillation-Empirical-Failure-Modes-and-Simple-Fixes-31dd5cc40dd181f89eead3de7181df1d)\n- 📄 Paper: [arXiv:2603.25562](https://arxiv.org/abs/2603.25562)\n- 💻 Code: [hhh675597/revisiting_opd](https://github.com/hhh675597/revisiting_opd)\n\n## Overview\n\nOn-policy distillation (OPD) trains a student on its own rollouts using teacher feedback. In practice, sampled-token OPD can be brittle because the learning signal is concentrated on a single sampled token, teacher guidance becomes less reliable on student-generated prefixes, and tokenizer or special-token mismatch can further distort the comparison. This repository includes our simple fixes, including a teacher top-K truncated reverse-KL objective, top-p rollouts, and special-token masking, together with the training recipes used in our reasoning and agentic multi-task experiments.\n\n## Table of Contents\n\n- [Overview](#overview)\n- [Installation](#installation)\n- [Data Preparation](#data-preparation)\n  - [Math Dataset](#math-dataset)\n  - [Multi-Task Dataset (ALFWorld + Math)](#multi-task-dataset-alfworld--math)\n- [Single-Task Training](#single-task-training)\n  - [Math](#math)\n  - [ALFWorld](#alfworld)\n- [Multi-Task Training](#multi-task-training)\n- [Configuring the Training Objective](#configuring-the-training-objective)\n  - [Original OPD](#original-opd)\n  - [Teacher-TopK (Ours)](#teacher-topk-ours)\n  - [GRPO Baseline (Multi-Task)](#grpo-baseline-multi-task)\n- [Evaluation](#evaluation)\n- [Model Merging](#model-merging)\n- [Key Configuration Reference](#key-configuration-reference)\n\n---\n\n## Installation\n\n```bash\nconda create -n verl-agent python==3.10 -y\nconda activate verl-agent\n\npip3 install vllm==0.8.5\npip3 install flash-attn==2.7.0.post2 --no-build-isolation --no-cache-dir\npip install -e .\n```\n\nFor ALFWorld environments, also install:\n\n```bash\npip3 install gymnasium==0.29.1\npip3 install stable-baselines3==2.6.0\npip install alfworld\nalfworld-download -f\n```\n\n---\n\n## Data Preparation\n\n### Math Dataset\n\nPrepare the math evaluation set (AIME 2024):\n\n```bash\npython examples/data_preprocess/verl_agent_math.py \\\n    --local_dir ./data/math_opd\n```\n\nThis saves `test.parquet` to the specified directory. Training data should be a parquet file in the same format (with `prompt`, `reward_model`, `env_kwargs` columns).\n\n### Multi-Task Dataset (ALFWorld + Math)\n\nFor multi-task training, combine ALFWorld placeholder samples and math problem samples into a single dataset:\n\n```bash\npython examples/data_preprocess/prepare_multitask_data.py \\\n    --alfworld_num 128 \\\n    --math_data_path /path/to/math/train.parquet \\\n    --output_dir ./data/multitask_data \\\n    --batching_mode sequential\n```\n\n- `--alfworld_num`: number of ALFWorld placeholder samples (ALFWorld loads its own game files internally, so these are just triggers).\n- `--math_data_path`: path to the math training parquet.\n- `--batching_mode`: `sequential` (recommended) groups samples by task type so each training step handles one task at a time.\n\n---\n\n## Single-Task Training\n\n### Math\n\nTrain a student model on math using Teacher-TopK distillation with a math teacher:\n\n```bash\nbash examples/opd/opd_math_qwen2.5-7b_it.sh\n```\n\nKey variables to configure in the script:\n\n```bash\nSTUDENT_MODEL=\"/path/to/Qwen2.5-7B-Instruct\"\nMATH_TEACHER=\"/path/to/OpenThinker3-7B\"           # teacher / reference model\nTRAIN_DATA=\"/path/to/math_opd/train.parquet\"\nVAL_DATA=\"/path/to/math_opd/test.parquet\"\n```\n\nThe entry point for single-task is `verl.trainer.main_ppo_multitask` (also used for multi-task). Single-task math uses `env.env_name=math`.\n\n### ALFWorld\n\nTrain a student on ALFWorld with an ALFWorld-specific teacher:\n\n```bash\nbash examples/opd/opd_alfworld_qwen2.5_7b_it.sh\n```\n\nKey variables:\n\n```bash\nSTUDENT_MODEL=\"/path/to/Qwen2.5-7B-Instruct\"\nALFWORLD_TEACHER=\"/path/to/alfworld-teacher-gigpo-qwen2.5-7b\"  # reference model\n```\n\nThe single-task ALFWorld example also uses `verl.trainer.main_ppo_multitask`. Single-task behavior comes from `env.env_name=alfworld/AlfredTWEnv`.\n\n---\n\n## Multi-Task Training\n\nMulti-task training alternates between tasks in each batch (sequential batching). The trainer dynamically swaps the reference model per task.\n\n```bash\nbash examples/opd/opd_multitask_qwen2.5-7b_it.sh\n```\n\nThe multi-task config block in the `.sh` file defines per-task settings:\n\n```bash\n+multitask.enable=True\n+multitask.batching_mode=sequential\n\n# Task 0: ALFWorld\n+multitask.tasks.task0.name=alfworld\n+multitask.tasks.task0.env_name=alfworld/AlfredTWEnv\n+multitask.tasks.task0.ref_model_path=${ALFWORLD_TEACHER}\n+multitask.tasks.task0.eval_dataset=eval_in_distribution\n+multitask.tasks.task0.max_response_length=512\n\n# Task 1: Math\n+multitask.tasks.task1.name=math\n+multitask.tasks.task1.env_name=math\n+multitask.tasks.task1.ref_model_path=${MATH_TEACHER}\n+multitask.tasks.task1.max_response_length=16384\n```\n\n- Each task specifies its own `ref_model_path` (teacher), `env_name`, and `max_response_length`.\n- `env.env_name=multitask` tells the trainer to dispatch to the appropriate environment per sample.\n- `data.batching_mode=sequential` ensures each step trains on one task at a time (e.g., step 1: math, step 2: alfworld, step 3: math, ...).\n\n---\n\n## Configuring the Training Objective\n\n### Original OPD\n\nThe original on-policy distillation uses token-level KL divergence between the student and teacher as the advantage signal. It is controlled by `algorithm.adv_estimator=opd`.\n\nExample script: `examples/opd/opd_original_math_qwen2.5_7b_it.sh`\n\nKey flags for original OPD:\n\n```bash\npython3 -m verl.trainer.main_ppo_multitask \\\n    algorithm.adv_estimator=opd \\\n    actor_rollout_ref.actor.kl_loss_type=k1 \\\n    +actor_rollout_ref.actor.opd_mask_special_tokens=False \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    algorithm.use_kl_in_reward=True \\\n    actor_rollout_ref.ref.model.path=${TEACHER_MODEL} \\\n    ...\n```\n\n| Flag | Value | Meaning |\n|------|-------|---------|\n| `algorithm.adv_estimator` | `opd` | Use OPD advantage estimator (KL-based advantages) |\n| `actor_rollout_ref.actor.kl_loss_type` | `k1` | Token-level KL approximation: `k1` = KL₁ = `exp(ref - student) - (ref - student) - 1` |\n| `actor_rollout_ref.actor.use_kl_loss` | `False` | No separate KL regularization loss term (KL is already in the advantage) |\n| `algorithm.use_kl_in_reward` | `True` | Fold KL penalty into the reward signal |\n| `+actor_rollout_ref.actor.opd_mask_special_tokens` | `True/False` | Whether to mask special tokens (e.g., `\u003cthink\u003e`, `\u003c/think\u003e`) from the KL computation |\n\n### Teacher-TopK (Ours)\n\nOur Teacher-TopK objective replaces the OPD advantage with a memory-efficient full reverse-KL loss computed over only the teacher's top-K token positions. Since the loss is applied directly (not via advantages), `algorithm.adv_estimator` is set to `placeholder`.\n\nExample script: `examples/opd/opd_math_qwen2.5-7b_it.sh`\n\nKey flags for Teacher-TopK:\n\n```bash\npython3 -m verl.trainer.main_ppo_multitask \\\n    algorithm.adv_estimator=placeholder \\\n    actor_rollout_ref.actor.kl_loss_type=full_reverse \\\n    +actor_rollout_ref.actor.kl_topk_tokens=32 \\\n    +actor_rollout_ref.actor.kl_topk_source=ref \\\n    +actor_rollout_ref.actor.norm_to_one_for_kl=True \\\n    +actor_rollout_ref.actor.clip_log_ratio=False \\\n    +actor_rollout_ref.actor.opd_mask_special_tokens=False \\\n    +actor_rollout_ref.actor.entropy_top_ratio=0.2 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=1 \\\n    algorithm.use_kl_in_reward=False \\\n    actor_rollout_ref.ref.model.path=${TEACHER_MODEL} \\\n    ...\n```\n\n| Flag | Value | Meaning |\n|------|-------|---------|\n| `algorithm.adv_estimator` | `placeholder` | No RL advantage computation; advantages are zeroed out |\n| `actor_rollout_ref.actor.kl_loss_type` | `full_reverse` | Full reverse KL divergence (KL(teacher \\|\\| student)) computed over top-K logits |\n| `+actor_rollout_ref.actor.kl_topk_tokens` | `32` | Number of top-K token positions to compute KL over |\n| `+actor_rollout_ref.actor.kl_topk_source` | `ref` | Whose top-K indices to use: `ref` (teacher) or `actor` (student) |\n| `+actor_rollout_ref.actor.norm_to_one_for_kl` | `True` | Normalize the top-K probability slice to sum to 1 before KL computation |\n| `+actor_rollout_ref.actor.clip_log_ratio` | `False` | Whether to clip log-probability ratios in KL |\n| `+actor_rollout_ref.actor.entropy_top_ratio` | `0.2` | Only train on the top 20% highest-entropy tokens per sample (entropy-based masking) |\n| `actor_rollout_ref.actor.use_kl_loss` | `True` | Enable the KL loss term in the actor update |\n| `actor_rollout_ref.actor.kl_loss_coef` | `1` | Weight of the KL loss |\n| `algorithm.use_kl_in_reward` | `False` | KL is applied as a direct loss, not folded into rewards |\n\n### GRPO Baseline (Multi-Task)\n\nFor comparison, you can also run multi-task training with GRPO (reward-based advantage, no distillation):\n\n```bash\nbash examples/opd/opd_grpo_multitask_qwen2.5-7b_it.sh\n```\n\nKey differences from OPD:\n\n```bash\nalgorithm.adv_estimator=grpo\nactor_rollout_ref.ref.model.path=${STUDENT_MODEL}   # ref = student itself (no teacher)\nactor_rollout_ref.actor.use_kl_loss=True\nactor_rollout_ref.actor.kl_loss_coef=1e-3            # small KL regularization\n```\n\n---\n\n## Evaluation\n\nUse the evaluation script to run validation-only mode on a trained checkpoint:\n\n```bash\nbash evaluate_multi.sh\n```\n\nConfigure the evaluation by setting:\n\n```bash\nSTUDENT_MODEL=\"/path/to/checkpoint/actor_merged\"  # merged checkpoint path\ntrainer.val_before_train=True\ntrainer.val_only=True\n```\n\nThe evaluation script supports the same multi-task config block, so you can evaluate on both ALFWorld and math simultaneously.\n\n---\n\n## Model Merging\n\nAfter training with FSDP, merge sharded checkpoints into a single model directory:\n\n```bash\npython scripts/model_merger.py merge \\\n    --backend fsdp \\\n    --local_dir /path/to/ckpts/global_step_XXX/actor \\\n    --target_dir /path/to/ckpts/global_step_XXX/actor_merged\n```\n\nThe merged model can then be loaded as `STUDENT_MODEL` for evaluation or further training.\n\n---\n\n## Key Configuration Reference\n\n### Common Settings\n\n| Config | Description | Typical Value |\n|--------|-------------|---------------|\n| `data.train_batch_size` | Number of prompts per training step | `16` |\n| `data.val_batch_size` | Number of prompts per validation step | `128` |\n| `env.rollout.n` | Group size (rollouts per prompt) | `8` |\n| `data.max_prompt_length` | Maximum prompt token length | `2048` |\n| `data.max_response_length` | Maximum response token length | `16384` (math) / `512` (alfworld) |\n| `data.batching_mode` | How to batch multi-task data | `sequential` |\n| `actor_rollout_ref.actor.optim.lr` | Learning rate | `1e-6` to `2e-6` |\n| `actor_rollout_ref.rollout.top_p` | Nucleus sampling top-p during rollout | `0.9` |\n| `trainer.n_gpus_per_node` | GPUs per node | `8` |\n| `trainer.save_freq` | Checkpoint save frequency (steps) | `80` |\n| `trainer.test_freq` | Validation frequency (steps) | `80` |\n\n### OPD-Specific\n\n| Config | Description |\n|--------|-------------|\n| `algorithm.adv_estimator` | `opd` for original OPD, `placeholder` for Teacher-TopK |\n| `algorithm.use_kl_in_reward` | `True` for original OPD, `False` for Teacher-TopK |\n| `actor_rollout_ref.actor.kl_loss_type` | `k1` (original OPD) or `full_reverse` (Teacher-TopK) |\n| `+actor_rollout_ref.actor.kl_topk_tokens` | Top-K token count for Teacher-TopK (e.g., `32`) |\n| `+actor_rollout_ref.actor.kl_topk_source` | Whose top-K indices to use: `ref` (teacher) or `actor` (student) |\n| `+actor_rollout_ref.actor.entropy_top_ratio` | Fraction of highest-entropy tokens to train on (e.g., `0.2`) |\n| `+actor_rollout_ref.actor.norm_to_one_for_kl` | Normalize top-K probabilities before KL |\n| `+actor_rollout_ref.actor.opd_mask_special_tokens` | Mask special tokens from KL computation |\n\n### Multi-Task Specific\n\n| Config | Description |\n|--------|-------------|\n| `+multitask.enable` | Enable multi-task mode |\n| `+multitask.batching_mode` | `sequential` (one task per step) |\n| `+multitask.tasks.taskN.name` | Task name (e.g., `alfworld`, `math`) |\n| `+multitask.tasks.taskN.env_name` | Environment name for this task |\n| `+multitask.tasks.taskN.ref_model_path` | Path to task-specific teacher model |\n| `+multitask.tasks.taskN.max_response_length` | Max response length for this task |\n| `+multitask.tasks.taskN.eval_dataset` | Evaluation dataset name (optional) |\n| `env.env_name` | Set to `multitask` for multi-task training |\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fhhh675597%2Frevisiting_opd","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fhhh675597%2Frevisiting_opd","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fhhh675597%2Frevisiting_opd/lists"}