{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 多臂老虎机：比较 ε-贪心（固定 / 衰减 ε）、UCB1 与汤普森采样的累积懊悔（regret）。\n",
    "\n",
    "环境：K=10 臂伯努利老虎机，每次独立实验重新随机抽取各臂的中奖概率 p_k ~ U(0, 1)。\n",
    "累积懊悔（伪懊悔）：R(T) = Σ_t (p* - p_{a_t})，越小越好。所有实验向量化，一次跑 RUNS 个独立老虎机。\n",
    "\n",
    "配套教程：[多臂老虎机](https://dqn.cn/basics/bandit)（DQN.CN 强化学习中文教程）\n",
    "\n",
    "源码：`examples/basics/bandit.py`，本 notebook 由 `scripts/build-notebooks.mjs` 自动生成，代码与 .py 文件一致。代码默认在 CPU 上运行，固定了随机种子；部分脚本支持设置环境变量 `DQN_QUICK=1` 只跑很短的版本（见下一个单元）。"
   ],
   "id": "basics-bandit-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": "basics-bandit-py-1"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "QUICK = os.environ.get(\"DQN_QUICK\") == \"1\"\n",
    "K, T, RUNS = 10, (1000 if QUICK else 5000), (50 if QUICK else 200)\n",
    "SEED = 0"
   ],
   "id": "basics-bandit-py-2"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def simulate(choose, rng, p):\n",
    "    \"\"\"choose(t, counts, values, wins, rng) -> 每个实验本步要拉的臂（形状 [RUNS]）。返回累积懊悔曲线 [RUNS, T]。\"\"\"\n",
    "    runs = p.shape[0]\n",
    "    counts = np.zeros((runs, K)); values = np.zeros((runs, K)); wins = np.zeros((runs, K))\n",
    "    best, regret = p.max(1), np.zeros((runs, T))\n",
    "    idx = np.arange(runs)\n",
    "    for t in range(T):\n",
    "        a = choose(t, counts, values, wins, rng)\n",
    "        r = (rng.random(runs) < p[idx, a]).astype(float)       # 伯努利奖励\n",
    "        counts[idx, a] += 1; wins[idx, a] += r\n",
    "        values[idx, a] += (r - values[idx, a]) / counts[idx, a] # 增量式均值更新\n",
    "        regret[:, t] = best - p[idx, a]\n",
    "    return regret.cumsum(1)"
   ],
   "id": "basics-bandit-py-3"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def argmax_random_tie(x, rng):\n",
    "    return (x + 1e-9 * rng.random(x.shape)).argmax(1)           # 随机打破平局"
   ],
   "id": "basics-bandit-py-4"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def eps_greedy(eps):\n",
    "    def choose(t, counts, values, wins, rng):\n",
    "        explore = rng.random(len(values)) < eps\n",
    "        return np.where(explore, rng.integers(K, size=len(values)), argmax_random_tie(values, rng))\n",
    "    return choose"
   ],
   "id": "basics-bandit-py-5"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def decaying_eps(c=5.0):\n",
    "    \"\"\"ε_t = min(1, c·K / (t+1))：前期大量探索，之后探索率按 1/t 衰减。\"\"\"\n",
    "    def choose(t, counts, values, wins, rng):\n",
    "        return eps_greedy(min(1.0, c * K / (t + 1)))(t, counts, values, wins, rng)\n",
    "    return choose"
   ],
   "id": "basics-bandit-py-6"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ucb1(t, counts, values, wins, rng):\n",
    "    if t < K:                                                   # 先把每个臂各拉一次\n",
    "        return np.full(len(values), t)\n",
    "    bonus = np.sqrt(2 * np.log(t + 1) / counts)\n",
    "    return argmax_random_tie(values + bonus, rng)"
   ],
   "id": "basics-bandit-py-7"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def thompson(t, counts, values, wins, rng):\n",
    "    samples = rng.beta(1 + wins, 1 + counts - wins)             # Beta(1,1) 先验的后验采样\n",
    "    return samples.argmax(1)"
   ],
   "id": "basics-bandit-py-8"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if __name__ == \"__main__\":\n",
    "    methods = {\n",
    "        \"ε-贪心 ε=0.01\": eps_greedy(0.01),\n",
    "        \"ε-贪心 ε=0.1\": eps_greedy(0.1),\n",
    "        \"衰减 ε（c=5）\": decaying_eps(5.0),\n",
    "        \"UCB1\": ucb1,\n",
    "        \"汤普森采样\": thompson,\n",
    "    }\n",
    "    p = np.random.default_rng(SEED).random((RUNS, K))           # 所有方法面对同一批老虎机\n",
    "    checkpoints = [100, 1000, T]\n",
    "    print(f\"K={K} 臂伯努利老虎机，{RUNS} 次独立实验，T={T} 步；表中为平均累积懊悔（± 标准误）\")\n",
    "    for name, fn in methods.items():\n",
    "        reg = simulate(fn, np.random.default_rng(SEED + 1), p)\n",
    "        cells = \"  \".join(f\"T={c}: {reg[:, c - 1].mean():6.1f}±{reg[:, c - 1].std() / np.sqrt(RUNS):4.1f}\" for c in checkpoints)\n",
    "        print(f\"{name:14s} {cells}\")"
   ],
   "id": "basics-bandit-py-9"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  },
  "colab": {
   "provenance": []
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}