{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SmolVLA 微调 · 右臂拿杯 · Colab\n",
    "\n",
    "**这份 notebook 是写好但没有跑过的** —— 写它的 session 连不上 huggingface.co（网络策略封的），没法自己验证。第一次跑时自己盯着每一格的输出，別盲信它一定能跑通。\n",
    "\n",
    "**前提**：数据集 `Suyang99/xlerobot-cup-grasp-20260820-0230` 在 HuggingFace Hub 上。如果是 private 的，在下面登录那一格自己粘贴 token（不要把 token 写进 notebook 保存下来，运行时输入就好）。\n",
    "\n",
    "**运行前先确认**：菜单 `代码执行程序` → `更改运行时类型` → 硬件加速器选 **T4 GPU**（免费档有）。没选对的话下面会用 CPU 跑，慢到没法用。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. 确认 GPU 在线"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!nvidia-smi\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. 装 lerobot（带 smolvla 依赖）"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q 'lerobot[smolvla]'\n",
    "\n",
    "import torch\n",
    "print('torch', torch.__version__, 'cuda available:', torch.cuda.is_available())\n",
    "assert torch.cuda.is_available(), '没拿到 GPU —— 回到上面把运行时类型改成 T4，别用 CPU 跑。'\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. 登录 HuggingFace（只有数据集是 private 时才需要）\n",
    "\n",
    "如果数据集是 public 的，这一格跳过不跑。是 private 的话跑下面这格，运行时会跑出一个输入框，把你自己的 HF token 粘进去。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import login\n",
    "login()  # 弹出输入框，粘贴你自己的 token，不会写进这个 notebook 文件\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. 挂载 Google Drive（保存 checkpoint用）\n",
    "\n",
    "**必做，不是可选。** Colab 运行时过一段时间会自动断开回收，到时候本地磁盘里的 checkpoint 会跟着消失。挂到 Drive 才能在断开之后接着跑。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from google.colab import drive\n",
    "drive.mount('/content/drive')\n",
    "\n",
    "import os\n",
    "OUT_DIR = '/content/drive/MyDrive/xlerobot-training/smolvla_cup_grasp_right'\n",
    "os.makedirs(OUT_DIR, exist_ok=True)\n",
    "print('checkpoints 会存在:', OUT_DIR)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. 跑训练\n",
    "\n",
    "参数跟 `03-software/DATA.md` 里的示例命令一致（`--policy.path=lerobot/smolvla_base`、`--policy.dtype=bfloat16`）。\n",
    "\n",
    "`batch_size` 我写了 4 作为起点 —— T4 有 16GB、比 Jetson 那 8GB 共享内存宽裕，但这**不是实测值**，是按 Jetson 那边 `train_smolvla.sh` 用的 batch=1 推算放大的。如果 OOM，把下面 `BATCH` 改小。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "REPO_ID = 'Suyang99/xlerobot-cup-grasp-20260820-0230'\n",
    "BATCH = 4          # 没实测过， OOM 就改小\n",
    "STEPS = 20000\n",
    "SAVE_FREQ = 2000\n",
    "LOG_FREQ = 50\n",
    "\n",
    "!python -m lerobot.scripts.lerobot_train \\\n",
    "  --dataset.repo_id={REPO_ID} \\\n",
    "  --policy.path=lerobot/smolvla_base \\\n",
    "  --policy.device=cuda \\\n",
    "  --policy.dtype=bfloat16 \\\n",
    "  --policy.push_to_hub=false \\\n",
    "  --batch_size={BATCH} \\\n",
    "  --steps={STEPS} \\\n",
    "  --save_freq={SAVE_FREQ} \\\n",
    "  --log_freq={LOG_FREQ} \\\n",
    "  --output_dir={OUT_DIR} \\\n",
    "  --job_name=smolvla_cup_grasp_right_colab\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. 跑完之后\n",
    "\n",
    "- checkpoint 在 Google Drive 里，不会跟着 Colab 运行时消失。\n",
    "- **这个模型和 Jetson 上那个 ACT 训练是同一份 50 集数据，两个独立的模型。** 那个未解决的注意事项仍然成立：不管哪个，训练完都必须用**训练时没出现过的新位置**实测，不能只看 loss 曲线就算数。详见 `03-software/DATA-PLAN.md` §6。"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "name": "train_smolvla_colab.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}