{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "# 10 pi_0 从权限门控到正式训练\n\n            pi_0 的模型加载、gated 权限和显存压力都更高。这里先把权限、数据、模型构造和 2-step 反向传播逐项过门，再启动正式训练。\n\n            这一节从 LeRobot 数据检查开始，依次完成 2-step smoke、正式训练、ROCm 资源观察和 checkpoint 定位。所有长任务默认关闭，先确认生成的配置，再显式打开运行开关。\n"
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": "TOPIC_ROOT = $NOTEBOOK_TOPIC_ROOT\nPROJECT_ROOT = $PROJECT_ROOT\nDATA_ROOT = $PROJECT_ROOT\nOUTPUT_ROOT = $OUTPUT_ROOT\nMODEL_ROOT = $MODEL_ROOT\n"
        }
      ],
      "source": "from pathlib import Path\nimport json\nimport os\nimport shutil\nimport subprocess\nimport sys\n\n\ndef find_topic_root():\n    cwd = Path.cwd().resolve()\n    for candidate in [cwd, *cwd.parents]:\n        if (candidate / \"assets\" / \"metrics_snapshot.json\").exists():\n            return candidate\n    raise RuntimeError(\"请从 AMD ROCm 专题目录或 notebooks 子目录启动 Jupyter。\")\n\n\nTOPIC_ROOT = find_topic_root()\nASSET_DIR = TOPIC_ROOT / \"assets\"\nNOTEBOOK_DIR = TOPIC_ROOT / \"notebooks\"\nPROJECT_ROOT = Path(os.environ.get(\"PROJECT_ROOT\", \"/path/to/every-embodied/mujoco_pnp\"))\nDATA_ROOT = Path(os.environ.get(\"DATA_ROOT\", \"/path/to/datasets/every_embodied\"))\nOUTPUT_ROOT = Path(os.environ.get(\"OUTPUT_ROOT\", TOPIC_ROOT / \"outputs\"))\nMODEL_ROOT = Path(os.environ.get(\"MODEL_ROOT\", PROJECT_ROOT / \"ckpt\"))\n\nprint(\"TOPIC_ROOT =\", TOPIC_ROOT)\nprint(\"PROJECT_ROOT =\", PROJECT_ROOT)\nprint(\"DATA_ROOT =\", DATA_ROOT)\nprint(\"OUTPUT_ROOT =\", OUTPUT_ROOT)\nprint(\"MODEL_ROOT =\", MODEL_ROOT)\n"
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {},
      "outputs": [],
      "source": "try:\n    from IPython.display import Image, Markdown, display\nexcept Exception:\n    class Markdown(str):\n        pass\n\n    def display(obj):\n        print(obj)\n\n    def Image(filename=None, width=None):\n        return f\"[image] {filename}\"\n\n\ndef show_asset(filename, width=960):\n    path = ASSET_DIR / filename\n    if path.exists():\n        display(Image(filename=str(path), width=width))\n    else:\n        print(f\"缺少素材：{path}\")\n\n\ndef md_table(headers, rows):\n    lines = [\"| \" + \" | \".join(headers) + \" |\", \"| \" + \" | \".join([\"---\"] * len(headers)) + \" |\"]\n    for row in rows:\n        lines.append(\"| \" + \" | \".join(str(x) for x in row) + \" |\")\n    display(Markdown(\"\\n\".join(lines)))\n"
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {},
      "outputs": [],
      "source": "import shlex\n\ntry:\n    import yaml\nexcept ImportError as exc:\n    raise RuntimeError(\"当前环境缺少 PyYAML，请先执行 pip install pyyaml。\") from exc\n\n\ndef require_project_layout():\n    required = [\n        PROJECT_ROOT / \"train_model.py\",\n        PROJECT_ROOT / \"env_config.py\",\n        PROJECT_ROOT / \"asset\" / \"example_scene_y2.xml\",\n        PROJECT_ROOT / \"mujoco_env\" / \"y_env2.py\",\n    ]\n    missing = [path for path in required if not path.exists()]\n    if missing:\n        raise FileNotFoundError(\n            \"PROJECT_ROOT 不是 04mujoco 教程目录，缺少：\\n\"\n            + \"\\n\".join(str(path) for path in missing)\n        )\n    return True\n\n\ndef dataset_report(dataset_root):\n    dataset_root = Path(dataset_root)\n    info_path = dataset_root / \"meta\" / \"info.json\"\n    tasks_path = dataset_root / \"meta\" / \"tasks.jsonl\"\n    if not info_path.exists():\n        raise FileNotFoundError(f\"找不到 LeRobot 元数据：{info_path}\")\n    info = json.loads(info_path.read_text(encoding=\"utf-8\"))\n    tasks = []\n    if tasks_path.exists():\n        tasks = [\n            json.loads(line)[\"task\"]\n            for line in tasks_path.read_text(encoding=\"utf-8\").splitlines()\n            if line.strip()\n        ]\n    features = info.get(\"features\", {})\n    rows = [\n        (\"repo_id\", info.get(\"repo_id\", \"\")),\n        (\"episodes\", info.get(\"total_episodes\", 0)),\n        (\"frames\", info.get(\"total_frames\", 0)),\n        (\"fps\", info.get(\"fps\", \"\")),\n        (\"state shape\", features.get(\"observation.state\", {}).get(\"shape\")),\n        (\"action shape\", features.get(\"action\", {}).get(\"shape\")),\n        (\"tasks\", \" / \".join(tasks)),\n    ]\n    md_table([\"数据项\", \"读取结果\"], rows)\n    return info, tasks\n\n\ndef make_training_config(\n    policy_type,\n    dataset_repo_id,\n    dataset_root,\n    output_dir,\n    steps,\n    batch_size,\n    chunk_size,\n    n_action_steps,\n    save_freq,\n    seed=42,\n):\n    return {\n        \"dataset\": {\"repo_id\": dataset_repo_id, \"root\": str(Path(dataset_root))},\n        \"policy\": {\n            \"type\": policy_type,\n            \"chunk_size\": int(chunk_size),\n            \"n_action_steps\": int(n_action_steps),\n            \"device\": \"cuda\",\n        },\n        \"save_checkpoint\": True,\n        \"output_dir\": str(Path(output_dir)),\n        \"batch_size\": int(batch_size),\n        \"job_name\": Path(output_dir).name,\n        \"resume\": False,\n        \"seed\": int(seed),\n        \"num_workers\": 4,\n        \"steps\": int(steps),\n        \"eval_freq\": 0,\n        \"log_freq\": max(1, min(50, int(steps))),\n        \"save_freq\": int(save_freq),\n        \"use_policy_training_preset\": True,\n        \"wandb\": {\n            \"enable\": False,\n            \"project\": f\"every_embodied_{policy_type}\",\n            \"entity\": None,\n            \"disable_artifact\": True,\n        },\n    }\n\n\ndef write_training_config(name, config):\n    config_dir = OUTPUT_ROOT / \"configs\"\n    config_dir.mkdir(parents=True, exist_ok=True)\n    path = config_dir / f\"{name}.yaml\"\n    path.write_text(\n        yaml.safe_dump(config, allow_unicode=True, sort_keys=False),\n        encoding=\"utf-8\",\n    )\n    print(\"已写出配置：\", path)\n    return path\n\n\ndef training_command(config_path):\n    return [\n        sys.executable,\n        str(PROJECT_ROOT / \"train_model.py\"),\n        \"--config_path\",\n        str(Path(config_path)),\n    ]\n\n\ndef run_training(config_path, enabled=False):\n    require_project_layout()\n    command = training_command(config_path)\n    print(\"$\", shlex.join(command))\n    if not enabled:\n        print(\"当前只预览命令。确认配置后，把对应 RUN_* 开关改为 True。\")\n        return None\n    env = os.environ.copy()\n    env[\"PYTHONPATH\"] = f\"{PROJECT_ROOT}:{env.get('PYTHONPATH', '')}\"\n    return subprocess.run(command, cwd=PROJECT_ROOT, env=env, check=True)\n\n\ndef show_rocm_resources():\n    command = [\"rocm-smi\", \"--showuse\", \"--showmemuse\", \"--showtemp\"]\n    print(\"$\", shlex.join(command))\n    if shutil.which(command[0]) is None:\n        print(\"未找到 rocm-smi。确认当前机器是否安装 ROCm。\")\n        return\n    result = subprocess.run(command, check=False, capture_output=True, text=True)\n    if result.stdout:\n        print(result.stdout.rstrip())\n    if result.stderr:\n        print(result.stderr.rstrip(), file=sys.stderr)\n"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "## Checkpoint 1：确认项目和训练数据\n"
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": "train_model.py = $PROJECT_ROOT/train_model.py\n训练数据 = $TRAIN_DATA_ROOT\n"
        }
      ],
      "source": "require_project_layout()\nprint(\"train_model.py =\", PROJECT_ROOT / \"train_model.py\")\nprint(\"训练数据 =\", Path(os.environ.get(\"TRAIN_DATA_ROOT\", DATA_ROOT / \"omy_pnp_language\")))\n"
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": "| 数据项 | 读取结果 |\n| --- | --- |\n| repo_id |  |\n| episodes | 20 |\n| frames | 2621 |\n| fps | 20 |\n| state shape | [6] |\n| action shape | [7] |\n| tasks | Place the blue mug on the plate. / Place the red mug on the plate. |\n"
        }
      ],
      "source": "candidate = Path(os.environ.get(\"TRAIN_DATA_ROOT\", DATA_ROOT / \"omy_pnp_language\"))\nif (candidate / \"meta\" / \"info.json\").exists():\n    dataset_report(candidate)\nelse:\n    print(\"训练数据还没准备好：\", candidate)\n    print(\"先完成 07_data_collection_and_audit.ipynb，或设置 TRAIN_DATA_ROOT 指向已有 LeRobot 数据集。\")\n"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "## Checkpoint 2：检查 Hugging Face gated 权限\n"
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": "HF_TOKEN 已注入当前进程： False\n请在终端私密设置 HF_TOKEN，再重新启动 Jupyter kernel。不要把 token 写进 Notebook。\n"
        }
      ],
      "source": "token_present = bool(os.environ.get(\"HF_TOKEN\"))\nprint(\"HF_TOKEN 已注入当前进程：\", token_present)\nif not token_present:\n    print(\"请在终端私密设置 HF_TOKEN，再重新启动 Jupyter kernel。不要把 token 写进 Notebook。\")\n"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "这里只检查环境变量是否存在，不打印 token。pi_0 训练还需要账号已经获得 PaliGemma 等 gated 权重的访问权限。\n"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "## Checkpoint 3：生成 smoke 与正式训练配置\n"
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": "已写出配置： $OUTPUT_ROOT/configs/pi0_smoke.yaml\n已写出配置： $OUTPUT_ROOT/configs/pi0_full.yaml\nsmoke output = $MODEL_ROOT/pi0_rocm_smoke\nfull output = $MODEL_ROOT/pi0_rocm_full\n"
        }
      ],
      "source": "MODEL_TYPE = 'pi0'\nDATASET_REPO_ID = os.environ.get(\"DATASET_REPO_ID\", \"datawhale_eai_pnp_language\")\nTRAIN_DATA_ROOT = Path(\n    os.environ.get(\"TRAIN_DATA_ROOT\", DATA_ROOT / \"omy_pnp_language\")\n).expanduser()\n\nRUN_SMOKE = False\nRUN_FULL_TRAIN = False\n\nSMOKE_OUTPUT = MODEL_ROOT / f\"{MODEL_TYPE}_rocm_smoke\"\nFULL_OUTPUT = MODEL_ROOT / f\"{MODEL_TYPE}_rocm_full\"\n\nsmoke_config = make_training_config(\n    policy_type=MODEL_TYPE,\n    dataset_repo_id=DATASET_REPO_ID,\n    dataset_root=TRAIN_DATA_ROOT,\n    output_dir=SMOKE_OUTPUT,\n    steps=2,\n    batch_size=min(4, 4),\n    chunk_size=5,\n    n_action_steps=5,\n    save_freq=2,\n)\nfull_config = make_training_config(\n    policy_type=MODEL_TYPE,\n    dataset_repo_id=DATASET_REPO_ID,\n    dataset_root=TRAIN_DATA_ROOT,\n    output_dir=FULL_OUTPUT,\n    steps=20000,\n    batch_size=4,\n    chunk_size=5,\n    n_action_steps=5,\n    save_freq=max(1, 20000 // 2),\n)\n\nsmoke_config_path = write_training_config(f\"{MODEL_TYPE}_smoke\", smoke_config)\nfull_config_path = write_training_config(f\"{MODEL_TYPE}_full\", full_config)\nprint(\"smoke output =\", SMOKE_OUTPUT)\nprint(\"full output =\", FULL_OUTPUT)\n"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "配置文件会写到 `$OUTPUT_ROOT/configs`，checkpoint 写到 `$MODEL_ROOT`。这两个目录应放在容量充足的磁盘，不要写进 Git 仓库。\n"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "## Checkpoint 4：先跑 2-step smoke\n"
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": "$ rocm-smi --showuse --showmemuse --showtemp\n\n\n============================ ROCm System Management Interface ============================\n====================================== Temperature =======================================\nGPU[0]\t\t: Temperature (Sensor edge) (C): 31.0\n==========================================================================================\n=================================== % time GPU is busy ===================================\nGPU[0]\t\t: GPU use (%): 0\n==========================================================================================\n=================================== Current Memory Use ===================================\nGPU[0]\t\t: GPU Memory Allocated (VRAM%): 9\nGPU[0]\t\t: Memory Activity: N/A\nGPU[0]\t\t: Not supported on the given system\n==========================================================================================\n================================== End of ROCm SMI Log ===================================\n$ $NOTEBOOK_PYTHON '$PROJECT_ROOT/train_model.py' --config_path '$OUTPUT_ROOT/configs/pi0_smoke.yaml'\n当前只预览命令。确认配置后，把对应 RUN_* 开关改为 True。\n"
        }
      ],
      "source": "show_rocm_resources()\nrun_training(smoke_config_path, enabled=RUN_SMOKE)\n"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "2-step smoke 只证明数据加载、模型构造、forward/backward、optimizer step 和 checkpoint 写出可用；它不证明模型收敛，更不能替代 MuJoCo closed-loop 成功率。\n"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "## Checkpoint 5：启动正式训练\n"
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": "$ $NOTEBOOK_PYTHON '$PROJECT_ROOT/train_model.py' --config_path '$OUTPUT_ROOT/configs/pi0_full.yaml'\n当前只预览命令。确认配置后，把对应 RUN_* 开关改为 True。\n"
        }
      ],
      "source": "run_training(full_config_path, enabled=RUN_FULL_TRAIN)\n"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "pi_0 的训练 loss 下降并不保证闭环抓取成功。小数据接触任务尤其容易出现 gripper 时序和误差累积问题。raw policy、learned auxiliary head 和带脚手架的 hybrid 结果必须分开报告。\n"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "常见恢复顺序：显存不足先减小 `batch_size`；出现 DataLoader worker / pickle 错误时把 YAML 中的 `num_workers` 改为 `0`；下载报 `401/403` 时检查账号权限和私密 `HF_TOKEN`；写 checkpoint 失败先检查 `$MODEL_ROOT` 所在磁盘。ROCm 的 MIOpen 数据库 warning 如果没有伴随训练退出可以记录后继续观察，出现 kernel 退出、NaN 或进程消失则必须停止长训练并回到 smoke。\n"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "## Checkpoint 6：定位 checkpoint 并检查资源\n"
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": "$ rocm-smi --showuse --showmemuse --showtemp\n\n\n============================ ROCm System Management Interface ============================\n====================================== Temperature =======================================\nGPU[0]\t\t: Temperature (Sensor edge) (C): 31.0\n==========================================================================================\n=================================== % time GPU is busy ===================================\nGPU[0]\t\t: GPU use (%): 0\n==========================================================================================\n=================================== Current Memory Use ===================================\nGPU[0]\t\t: GPU Memory Allocated (VRAM%): 9\nGPU[0]\t\t: Memory Activity: N/A\nGPU[0]\t\t: Not supported on the given system\n==========================================================================================\n================================== End of ROCm SMI Log ===================================\n正式训练尚未运行： $MODEL_ROOT/pi0_rocm_full\n"
        }
      ],
      "source": "show_rocm_resources()\nif (FULL_OUTPUT / \"checkpoints\").exists():\n    checkpoints = sorted((FULL_OUTPUT / \"checkpoints\").glob(\"*/pretrained_model\"))\n    print(\"可用 checkpoint：\")\n    for path in checkpoints:\n        print(\" -\", path)\nelse:\n    print(\"正式训练尚未运行：\", FULL_OUTPUT)\n"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "训练完成后不要只挑 loss 最低的节点。下一步进入 `11_mujoco_closed_loop_deploy.ipynb`，用固定 seed、严格 `physical_success` 和视频复核比较中间 checkpoint 与最终 checkpoint。\n"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "## 已完成训练的日志快照\n"
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": "               pi0 [####################################] 1500/1500\nstage: public raw baseline\nclosed-loop physical_success: 0/20\nnote: loss 从 step500 约 0.179 降到 step1500 约 0.046，但 raw closed-loop 没有成功\n"
        }
      ],
      "source": "progress = json.loads(\n    (ASSET_DIR / \"training_progress_snapshot.json\").read_text(encoding=\"utf-8\")\n)\nrecords = [row for row in progress[\"records\"] if row[\"model\"] == 'pi0']\nfor row in records:\n    ratio = row[\"current_step\"] / max(1, row[\"total_steps\"])\n    filled = round(36 * ratio)\n    bar = \"#\" * filled + \".\" * (36 - filled)\n    print(f'{row[\"model\"]:>18} [{bar}] {row[\"current_step\"]}/{row[\"total_steps\"]}')\n    print(\"stage:\", row[\"stage\"])\n    print(\"closed-loop physical_success:\", row[\"physical_success\"])\n    print(\"note:\", row[\"note\"])\n"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "![AMD ROCm 历史训练进度与闭环结果](../assets/training_progress_overview.png)\n\n                这张图来自已经完成的 AMD 实验日志，不是假装本次打开 Notebook 后重新跑完的训练。上面的 smoke/full 单元提供真实启动代码；运行长训练时，终端进度会继续写入输出目录，完成后再用 11 的闭环评估决定 checkpoint 是否可用。\n"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "## 在 Notebook 中跟踪后台训练日志\n"
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": "训练运行后可重复执行：tail_training_log(TRAIN_LOG)\n训练日志尚不存在： $MODEL_ROOT/pi0_rocm_full/train.log\n"
        }
      ],
      "source": "def tail_training_log(log_path, lines=25):\n    log_path = Path(log_path)\n    if not log_path.exists():\n        print(\"训练日志尚不存在：\", log_path)\n        return\n    content = log_path.read_text(encoding=\"utf-8\", errors=\"replace\").splitlines()\n    print(\"\\n\".join(content[-lines:]))\n\n\nTRAIN_LOG = Path(os.environ.get(\"TRAIN_LOG\", FULL_OUTPUT / \"train.log\"))\nprint(\"训练运行后可重复执行：tail_training_log(TRAIN_LOG)\")\ntail_training_log(TRAIN_LOG)\n"
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3"
    },
    "amd_rocm_tutorial_execution": {
      "executor": "code/execute_tutorial_notebooks.py",
      "python": "3.10.19",
      "code_cells_executed": 12,
      "errors": 0,
      "long_tasks_enabled": false
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}