{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/41_opd_loss.ipynb)\n",
    "\n",
    "# 🔴 Hard: OPD Loss\n",
    "\n",
    "Implement the **On-Policy Distillation (OPD)** loss used to distill one or more teacher policies into a student policy on trajectories sampled from the student.\n",
    "\n",
    "For each token position, OPD minimizes a weighted reverse KL from the student distribution to each teacher distribution:\n",
    "\n",
    "$$\\mathcal{L}_{\\text{OPD}} = \\sum_j w_j\\,D_{\\text{KL}}\\big(\\pi_\\theta(\\cdot \\mid x)\\;||\\;\\pi_{E_j}(\\cdot \\mid x)\\big)$$\n",
    "\n",
    "where\n",
    "\n",
    "$$D_{\\text{KL}}(p||q) = \\sum_v p(v)\\,[\\log p(v) - \\log q(v)].$$\n",
    "\n",
    "### Signature\n",
    "```python\n",
    "from torch import Tensor\n",
    "\n",
    "def opd_loss(student_logits: Tensor,\n",
    "             teacher_logits: Tensor,\n",
    "             teacher_weights: Tensor | None = None,\n",
    "             mask: Tensor | None = None,\n",
    "             temperature: float = 1.0) -> Tensor:\n",
    "    \"\"\"OPD reverse-KL distillation loss.\n",
    "\n",
    "    student_logits: (..., V) logits from the student policy\n",
    "    teacher_logits: (..., V) for one teacher, or (T, ..., V) for T teachers\n",
    "    teacher_weights: optional (T,) weights, normalized internally\n",
    "    mask: optional (...) token mask, where 1 = include and 0 = ignore\n",
    "    temperature: softmax temperature used for distillation\n",
    "    returns: scalar loss (Tensor)\n",
    "    \"\"\"\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
    "try:\n",
    "    import google.colab\n",
    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
    "except ImportError:\n",
    "    pass\n"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch import Tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ✏️ YOUR IMPLEMENTATION HERE\n",
    "\n",
    "def opd_loss(student_logits: Tensor,\n",
    "             teacher_logits: Tensor,\n",
    "             teacher_weights: Tensor | None = None,\n",
    "             mask: Tensor | None = None,\n",
    "             temperature: float = 1.0) -> Tensor:\n",
    "    pass  # reverse KL: sum_v p_student * (log p_student - log p_teacher)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🧪 Debug\n",
    "student_logits = torch.tensor([[[2.0, 0.0, -1.0], [0.5, 1.0, -0.5]]])\n",
    "teacher_logits = torch.tensor([[[1.0, 1.5, -0.5], [0.0, 2.0, -1.0]]])\n",
    "print('Loss:', opd_loss(student_logits, teacher_logits).item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ✅ SUBMIT\n",
    "from torch_judge import check\n",
    "check('opd_loss')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}