{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "# 12 pi0 strict-input 端到端诊断\n\n        这一节把 pi0 调试中最容易混淆的四件事拆开：raw policy、只用 VLA 合法输入的 learned head、固定场景策略采样稳定性、真正随机物体位置泛化。所有表格都来自 AMD 设备实测结果。\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": "markdown",
      "metadata": {},
      "source": "## Checkpoint 1：读取最终严格评估结果\n"
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": "| 策略 | 协议 | final strict |\n| --- | --- | --- |\n| raw pi0 | env seed 0；12 个 policy sampling seeds | 0/12 |\n| pi0 + visual/history GRU head | env seed 0；同一批 12 个 policy sampling seeds | 6/12 |\n| pi0 + visual/history GRU head | 修正 seed 后的随机环境 41–44 | 1/4 |\n"
        }
      ],
      "source": "results = json.loads(\n    (ASSET_DIR / \"pi0_strict_input_results.json\").read_text(encoding=\"utf-8\")\n)\nfixed = results[\"fixed_scene\"]\nrandomized = results[\"randomized_environment_gate\"]\nrows = [\n    (\n        \"raw pi0\",\n        \"env seed 0；12 个 policy sampling seeds\",\n        f'{fixed[\"raw_pi0\"][\"success\"]}/{fixed[\"raw_pi0\"][\"total\"]}',\n    ),\n    (\n        \"pi0 + visual/history GRU head\",\n        \"env seed 0；同一批 12 个 policy sampling seeds\",\n        f'{fixed[\"pi0_visual_history_gru\"][\"success\"]}/{fixed[\"pi0_visual_history_gru\"][\"total\"]}',\n    ),\n    (\n        \"pi0 + visual/history GRU head\",\n        \"修正 seed 后的随机环境 41–44\",\n        f'{randomized[\"pi0_visual_history_gru\"][\"success\"]}/{randomized[\"pi0_visual_history_gru\"][\"total\"]}',\n    ),\n]\nmd_table([\"策略\", \"协议\", \"final strict\"], rows)\n"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "![pi0 strict-input 对照](../assets/pi0_strict_input_progress.png)\n\n        固定场景从 `0/12` 提到 `6/12` 是实质进步，但不能跳过随机环境 `1/4`。这也是为什么教程把“固定场景拟合”和“位置泛化”分开写。\n"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "## Checkpoint 2：learned head 到底看了什么\n"
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": "| 类别 | 内容 |\n| --- | --- |\n| 允许 | agent/wrist 图像、语言、6D robot proprio、上一条策略自己执行过的 7D EEF/gripper 命令 |\n| 禁止 | target xyz、plate xyz、GT phase、oracle action、当前或未来 GT gripper 标签 |\n| 输出 | EEF keyframe/历史动作与 gripper close probability |\n| 部署 | gripper probability threshold=0.5，再转成二值执行命令 |\n"
        }
      ],
      "source": "rows = [\n    (\"允许\", \"agent/wrist 图像、语言、6D robot proprio、上一条策略自己执行过的 7D EEF/gripper 命令\"),\n    (\"禁止\", \"target xyz、plate xyz、GT phase、oracle action、当前或未来 GT gripper 标签\"),\n    (\"输出\", \"EEF keyframe/历史动作与 gripper close probability\"),\n    (\"部署\", \"gripper probability threshold=0.5，再转成二值执行命令\"),\n]\nmd_table([\"类别\", \"内容\"], rows)\n"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "这个 head 具备 VLA-compatible 的输入边界，但它是额外训练的视觉/历史控制头，因此报告时写 `pi0 + learned head`，不能写成 raw pi0。它学的是图像与机器人历史中的阶段线索，不依赖鲜艳标记或手工目标坐标。\n"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "## Checkpoint 3：visual/history head 的实际训练命令\n"
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": "$ $NOTEBOOK_PYTHON $NOTEBOOK_TOPIC_ROOT/code/pi0/train_pi0_visual_contact_head.py --source your_multiseed_fulltask_dataset '$PROJECT_ROOT/your_multiseed_fulltask_dataset' '$OUTPUT_ROOT/collection_results.jsonl' --policy-path '$MODEL_ROOT/pi0_base/pretrained_model' --stats-repo-id your_multiseed_fulltask_dataset --stats-dataset-root '$PROJECT_ROOT/your_multiseed_fulltask_dataset' --train-seeds 21-32 --val-seeds 33-36 --append-prev-action --feature-cache '$MODEL_ROOT/pi0_visual_history_head/visual_features.npz' --output '$MODEL_ROOT/pi0_visual_history_head/visual_keyframe_head.npz' --summary-json '$MODEL_ROOT/pi0_visual_history_head/visual_keyframe_head_summary.json'\n$ $NOTEBOOK_PYTHON $NOTEBOOK_TOPIC_ROOT/code/pi0/train_pi0_visual_gripper_gru.py --source your_multiseed_fulltask_dataset '$PROJECT_ROOT/your_multiseed_fulltask_dataset' '$OUTPUT_ROOT/collection_results.jsonl' --feature-cache '$MODEL_ROOT/pi0_visual_history_head/visual_features.npz' --direct-head '$MODEL_ROOT/pi0_visual_history_head/visual_keyframe_head.npz' --train-seeds 21-32 --val-seeds 33-36 --epochs 300 --transition-weight 20 --release-weight 40 --output '$MODEL_ROOT/pi0_visual_history_head/visual_keyframe_gru_head.npz' --summary-json '$MODEL_ROOT/pi0_visual_history_head/visual_keyframe_gru_summary.json'\n$ $NOTEBOOK_PYTHON $NOTEBOOK_TOPIC_ROOT/code/pi0/test_pi0_visual_gripper_gru_parity.py --head '$MODEL_ROOT/pi0_visual_history_head/visual_keyframe_gru_head.npz' --steps 32\n默认只预览。先替换多位置数据、collector JSONL 和 base policy 路径，再打开 RUN_AUXILIARY_TRAIN。\n"
        }
      ],
      "source": "import shlex\n\nHEAD_SOURCE_REPO = os.environ.get(\"HEAD_SOURCE_REPO\", \"your_multiseed_fulltask_dataset\")\nHEAD_SOURCE_ROOT = Path(os.environ.get(\"HEAD_SOURCE_ROOT\", DATA_ROOT / HEAD_SOURCE_REPO))\nHEAD_SOURCE_JSONL = Path(os.environ.get(\"HEAD_SOURCE_JSONL\", OUTPUT_ROOT / \"collection_results.jsonl\"))\nHEAD_POLICY_PATH = Path(os.environ.get(\"HEAD_POLICY_PATH\", MODEL_ROOT / \"pi0_base\" / \"pretrained_model\"))\nHEAD_DIR = Path(os.environ.get(\"HEAD_DIR\", MODEL_ROOT / \"pi0_visual_history_head\"))\nFEATURE_CACHE = HEAD_DIR / \"visual_features.npz\"\nDIRECT_HEAD = HEAD_DIR / \"visual_keyframe_head.npz\"\nGRU_HEAD = HEAD_DIR / \"visual_keyframe_gru_head.npz\"\nRUN_AUXILIARY_TRAIN = False\n\ncommands = [\n    [\n        sys.executable,\n        str(TOPIC_ROOT / \"code\" / \"pi0\" / \"train_pi0_visual_contact_head.py\"),\n        \"--source\", HEAD_SOURCE_REPO, str(HEAD_SOURCE_ROOT), str(HEAD_SOURCE_JSONL),\n        \"--policy-path\", str(HEAD_POLICY_PATH),\n        \"--stats-repo-id\", HEAD_SOURCE_REPO,\n        \"--stats-dataset-root\", str(HEAD_SOURCE_ROOT),\n        \"--train-seeds\", \"21-32\",\n        \"--val-seeds\", \"33-36\",\n        \"--append-prev-action\",\n        \"--feature-cache\", str(FEATURE_CACHE),\n        \"--output\", str(DIRECT_HEAD),\n        \"--summary-json\", str(HEAD_DIR / \"visual_keyframe_head_summary.json\"),\n    ],\n    [\n        sys.executable,\n        str(TOPIC_ROOT / \"code\" / \"pi0\" / \"train_pi0_visual_gripper_gru.py\"),\n        \"--source\", HEAD_SOURCE_REPO, str(HEAD_SOURCE_ROOT), str(HEAD_SOURCE_JSONL),\n        \"--feature-cache\", str(FEATURE_CACHE),\n        \"--direct-head\", str(DIRECT_HEAD),\n        \"--train-seeds\", \"21-32\",\n        \"--val-seeds\", \"33-36\",\n        \"--epochs\", \"300\",\n        \"--transition-weight\", \"20\",\n        \"--release-weight\", \"40\",\n        \"--output\", str(GRU_HEAD),\n        \"--summary-json\", str(HEAD_DIR / \"visual_keyframe_gru_summary.json\"),\n    ],\n    [\n        sys.executable,\n        str(TOPIC_ROOT / \"code\" / \"pi0\" / \"test_pi0_visual_gripper_gru_parity.py\"),\n        \"--head\", str(GRU_HEAD),\n        \"--steps\", \"32\",\n    ],\n]\n\nfor command in commands:\n    print(\"$\", shlex.join(command))\nif RUN_AUXILIARY_TRAIN:\n    HEAD_DIR.mkdir(parents=True, exist_ok=True)\n    for command in commands:\n        subprocess.run(command, cwd=PROJECT_ROOT, check=True)\nelse:\n    print(\"默认只预览。先替换多位置数据、collector JSONL 和 base policy 路径，再打开 RUN_AUXILIARY_TRAIN。\")\n"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "## Checkpoint 4：语言表述也属于部署协议\n"
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": "| instruction | learned head | strict | 现象 |\n| --- | --- | --- | --- |\n| Place the blue mug on the plate. | 0 active steps | 0/1 | 短 prompt 超出小头训练分布，prototype gate 一直拒绝接管 |\n| Pick up the blue mug and place it on the plate. | 269 active steps | 1/1 | 与训练指令一致；seed3 在 270 步通过 final strict |\n"
        }
      ],
      "source": "rows = [\n    (\n        \"Place the blue mug on the plate.\",\n        \"0 active steps\",\n        \"0/1\",\n        \"短 prompt 超出小头训练分布，prototype gate 一直拒绝接管\",\n    ),\n    (\n        \"Pick up the blue mug and place it on the plate.\",\n        \"269 active steps\",\n        \"1/1\",\n        \"与训练指令一致；seed3 在 270 步通过 final strict\",\n    ),\n]\nmd_table([\"instruction\", \"learned head\", \"strict\", \"现象\"], rows)\n"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "这不是说 VLA 永远不能接受同义指令，而是当前 4 条小数据训练出来的 auxiliary head 还没有语言改写泛化。部署时要先固定训练/评估 prompt；下一轮再加入同义指令增强，并单独测试 language paraphrase gate。\n"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "## Checkpoint 5：为什么 overall accuracy 会骗人\n"
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": "| 模型 | 验证指标 | held-out release | 解读 |\n| --- | --- | --- | --- |\n| 旧线性 gripper head | 99.66% | 0/4 | release 只有极少帧，总体准确率掩盖了全漏 |\n| GRU sequence head | 96.8% balanced | 5/5 | 按 close/release transition 单独选阈值 |\n"
        }
      ],
      "source": "rows = [\n    (\"旧线性 gripper head\", \"99.66%\", \"0/4\", \"release 只有极少帧，总体准确率掩盖了全漏\"),\n    (\"GRU sequence head\", \"96.8% balanced\", \"5/5\", \"按 close/release transition 单独选阈值\"),\n]\nmd_table([\"模型\", \"验证指标\", \"held-out release\", \"解读\"], rows)\n"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "## Checkpoint 6：用代码复现最终成功谓词\n"
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": "示例是否通过最终严格判定： True\n"
        }
      ],
      "source": "def final_strict_success(row):\n    p = results[\"strict_predicate\"]\n    return bool(\n        row[\"legacy_success\"]\n        and row[\"max_target_lift\"] >= p[\"min_lift_m\"]\n        and row[\"max_lifted_run\"] >= p[\"min_lift_steps\"]\n        and row[\"upright_cos\"] >= p[\"min_upright_cos\"]\n        and row[\"plate_z_gap\"] <= p[\"max_plate_z_gap_m\"]\n        and row[\"plate_xy_displacement\"] <= p[\"max_plate_xy_displacement_m\"]\n        and row[\"stable_place_steps\"] >= p[\"stable_place_steps\"]\n    )\n\n\nknown_success = {\n    \"legacy_success\": True,\n    \"max_target_lift\": 0.11,\n    \"max_lifted_run\": 8,\n    \"upright_cos\": 0.95,\n    \"plate_z_gap\": 0.0448,\n    \"plate_xy_displacement\": 0.0105,\n    \"stable_place_steps\": 5,\n}\nprint(\"示例是否通过最终严格判定：\", final_strict_success(known_success))\n"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "## Checkpoint 7：seed bug 对结论的影响\n"
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": "旧代码： if seed != None: np.random.seed(seed=0)\n修复后： if seed is not None: np.random.seed(seed=seed)\n影响： 旧 30-seed 面板只改变策略采样随机性，环境始终是 seed 0，不能作为空间泛化证据\n"
        }
      ],
      "source": "seed_bug = results[\"seed_bug\"]\nprint(\"旧代码：\", seed_bug[\"old_code\"])\nprint(\"修复后：\", seed_bug[\"fixed_code\"])\nprint(\"影响：\", seed_bug[\"impact\"])\n"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "## Checkpoint 8：四视角成功回放\n\n        ![四视角成功回放关键帧](../assets/pnp_four_view_strict_success_sequence.jpg)\n\n        <video controls width=\"960\" src=\"../assets/pnp_four_view_strict_success.mp4\"></video>\n"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "## Checkpoint 9：下一轮数据计划\n"
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": "| 阶段 | 动作 |\n| --- | --- |\n| 采集 | 修正 seed 后采 30–50 条多位置完整成功轨迹 |\n| 分层 | 按红/蓝杯、起点区域、contact miss、transport tip、release fail 分桶 |\n| 训练 | 先保护 fixed-scene 6/12，再训练 visual/history head；不要继续盲加 raw pi0 steps |\n| 评估 | fixed scene 与 random env 分开；随机环境至少 20–30 条 |\n| 替换门槛 | 新模型同协议超过保护基线，且视频无推杯/悬空误判 |\n"
        }
      ],
      "source": "plan = [\n    (\"采集\", \"修正 seed 后采 30–50 条多位置完整成功轨迹\"),\n    (\"分层\", \"按红/蓝杯、起点区域、contact miss、transport tip、release fail 分桶\"),\n    (\"训练\", \"先保护 fixed-scene 6/12，再训练 visual/history head；不要继续盲加 raw pi0 steps\"),\n    (\"评估\", \"fixed scene 与 random env 分开；随机环境至少 20–30 条\"),\n    (\"替换门槛\", \"新模型同协议超过保护基线，且视频无推杯/悬空误判\"),\n]\nmd_table([\"阶段\", \"动作\"], plan)\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": 10,
      "errors": 0,
      "long_tasks_enabled": false
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}