{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 大模型 RL 的玩具演示：用一个极小的“策略模型”演示 GRPO 与 DPO 的损失函数。\n",
    "\n",
    "设定：共 20 个“问题”，每个问题有 8 个候选“回答”（相当于 8 个 token），其中恰有 1 个正确。\n",
    "可验证奖励：回答正确得 1，否则得 0。\n",
    "\n",
    "配套教程：[RLHF 与大模型中的强化学习（PPO / DPO / GRPO）](https://dqn.cn/advanced/rlhf-llm)（DQN.CN 强化学习中文教程）\n",
    "\n",
    "源码：`examples/advanced/rlhf_toy.py`，本 notebook 由 `scripts/build-notebooks.mjs` 自动生成，代码与 .py 文件一致。代码默认在 CPU 上运行，固定了随机种子；部分脚本支持设置环境变量 `DQN_QUICK=1` 只跑很短的版本（见下一个单元）。"
   ],
   "id": "advanced-rlhf_toy-py-0"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 环境准备：在 Colab 等缺少依赖的环境中自动安装（本地已装好依赖时什么也不做）\n",
    "import importlib.util, os, subprocess, sys\n",
    "sys.argv = sys.argv[:1]   # notebook 内核会带上 -f 等参数，清掉以免脚本把它们当成命令行参数\n",
    "# os.environ[\"DQN_QUICK\"] = \"1\"   # 取消注释：支持该开关的脚本只跑很短的版本，用于快速检查"
   ],
   "id": "advanced-rlhf_toy-py-1"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "\n",
    "torch.manual_seed(0)\n",
    "N_PROMPTS, N_ANS = 20, 8\n",
    "correct = torch.randint(0, N_ANS, (N_PROMPTS,))"
   ],
   "id": "advanced-rlhf_toy-py-2"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def new_policy():\n",
    "    # 每个问题一组 logits，相当于一个查表式的“语言模型”；初始化带少量随机性\n",
    "    return (0.1 * torch.randn(N_PROMPTS, N_ANS)).requires_grad_()"
   ],
   "id": "advanced-rlhf_toy-py-3"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def accuracy(logits):\n",
    "    return F.softmax(logits, -1)[torch.arange(N_PROMPTS), correct].mean().item()"
   ],
   "id": "advanced-rlhf_toy-py-4"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ======================= GRPO =======================\n",
    "G, CLIP, BETA, LR, STEPS = 8, 0.2, 0.04, 0.5, 60\n",
    "logits = new_policy()\n",
    "ref = logits.detach().clone()                      # 参考模型（冻结）\n",
    "opt = torch.optim.Adam([logits], lr=LR)\n",
    "print(f\"[GRPO] 初始：正确回答的平均概率 {accuracy(logits):.3f}\")\n",
    "for step in range(STEPS):\n",
    "    with torch.no_grad():\n",
    "        old_logp_all = F.log_softmax(logits, -1)\n",
    "        samples = torch.multinomial(old_logp_all.exp(), G, replacement=True)   # 每个问题采样 G 个回答\n",
    "        rewards = (samples == correct[:, None]).float()\n",
    "        # 组内相对优势：A_i = (r_i - mean(r)) / (std(r) + eps)，不需要价值网络\n",
    "        adv = (rewards - rewards.mean(1, keepdim=True)) / (rewards.std(1, keepdim=True) + 1e-4)\n",
    "        old_logp = old_logp_all.gather(1, samples)\n",
    "    for _ in range(2):                             # 同一批样本上做少量优化步\n",
    "        logp_all = F.log_softmax(logits, -1)\n",
    "        logp = logp_all.gather(1, samples)\n",
    "        ratio = (logp - old_logp).exp()\n",
    "        pg = torch.min(ratio * adv, ratio.clamp(1 - CLIP, 1 + CLIP) * adv)\n",
    "        # KL(π || π_ref) 的无偏估计（k3）：exp(ref - logp) - (ref - logp) - 1\n",
    "        ref_logp = F.log_softmax(ref, -1).gather(1, samples)\n",
    "        kl = (ref_logp - logp).exp() - (ref_logp - logp) - 1\n",
    "        loss = -(pg - BETA * kl).mean()\n",
    "        opt.zero_grad(); loss.backward(); opt.step()\n",
    "    if (step + 1) % 20 == 0:\n",
    "        print(f\"[GRPO] step {step + 1:3d}：正确回答的平均概率 {accuracy(logits):.3f}，本批平均奖励 {rewards.mean():.3f}\")\n",
    "\n",
    "# ======================= DPO =======================\n",
    "BETA_DPO, LR, STEPS = 0.1, 0.5, 200\n",
    "logits = new_policy()\n",
    "ref = logits.detach().clone()\n",
    "opt = torch.optim.Adam([logits], lr=LR)\n",
    "# 偏好数据：对每个问题，(正确回答 y_w) ≻ (随机错误回答 y_l)\n",
    "idx = torch.arange(N_PROMPTS)\n",
    "wrong = (correct + torch.randint(1, N_ANS, (N_PROMPTS,))) % N_ANS\n",
    "print(f\"[DPO ] 初始：正确回答的平均概率 {accuracy(logits):.3f}\")\n",
    "for step in range(STEPS):\n",
    "    logp, ref_logp = F.log_softmax(logits, -1), F.log_softmax(ref, -1)\n",
    "    # 隐式奖励差：β[(log π(y_w) - log π_ref(y_w)) - (log π(y_l) - log π_ref(y_l))]\n",
    "    margin = BETA_DPO * ((logp[idx, correct] - ref_logp[idx, correct]) - (logp[idx, wrong] - ref_logp[idx, wrong]))\n",
    "    loss = -F.logsigmoid(margin).mean()\n",
    "    opt.zero_grad(); loss.backward(); opt.step()\n",
    "    if (step + 1) % 50 == 0:\n",
    "        p = F.softmax(logits, -1)\n",
    "        print(f\"[DPO ] step {step + 1:3d}：loss {loss.item():.4f}，正确回答平均概率 {accuracy(logits):.3f}，\"\n",
    "              f\"被拒绝回答平均概率 {p[idx, wrong].mean():.3f}\")"
   ],
   "id": "advanced-rlhf_toy-py-5"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  },
  "colab": {
   "provenance": []
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}