{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "# 15 Pi0 端到端：权限、长训、严格输入与失败桶\n\n            Pi0 这一版不是“完全复现历史 hard8”的主线，而是进阶诊断模型。Notebook 重点展示 gated 权限、长训配置、S8500 checkpoint、median reducer、严格评估和失败桶分析。\n"
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": "TOPIC_ROOT = $TOPIC_ROOT\nPROJECT_ROOT = $PROJECT_ROOT\nDATA_ROOT = $DATA_ROOT\nMODEL_ROOT = $MODEL_ROOT\nOUTPUT_ROOT = $OUTPUT_ROOT\n"
    }
   ],
   "source": [
    "from pathlib import Path\n",
    "import json\n",
    "import os\n",
    "import shlex\n",
    "import shutil\n",
    "import subprocess\n",
    "import sys\n",
    "\n",
    "try:\n",
    "    from IPython.display import HTML, Markdown, display\n",
    "except Exception:\n",
    "    class Markdown(str):\n",
    "        pass\n",
    "\n",
    "    class HTML(str):\n",
    "        pass\n",
    "\n",
    "    def display(obj):\n",
    "        print(obj)\n",
    "\n",
    "\n",
    "def 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",
    "\n",
    "TOPIC_ROOT = find_topic_root()\n",
    "ASSET_DIR = TOPIC_ROOT / \"assets\"\n",
    "PROJECT_ROOT = Path(os.environ.get(\"PROJECT_ROOT\", \"/path/to/04mujoco复现ACT、Pi0、SmolVLA\"))\n",
    "DATA_ROOT = Path(os.environ.get(\"DATA_ROOT\", \"/path/to/datasets/every_embodied\"))\n",
    "MODEL_ROOT = Path(os.environ.get(\"MODEL_ROOT\", \"/path/to/model/checkpoints\"))\n",
    "OUTPUT_ROOT = Path(os.environ.get(\"OUTPUT_ROOT\", TOPIC_ROOT / \"outputs\"))\n",
    "\n",
    "# The AMD teaching workflow should be runnable from local datasets/checkpoints.\n",
    "# Avoid surprising network calls during class or when AUP/Radeon Cloud cannot\n",
    "# reach Hugging Face.\n",
    "os.environ.setdefault(\"HF_HUB_OFFLINE\", os.environ.get(\"NOTEBOOK_HF_OFFLINE\", \"1\"))\n",
    "os.environ.setdefault(\"TRANSFORMERS_OFFLINE\", os.environ.get(\"NOTEBOOK_HF_OFFLINE\", \"1\"))\n",
    "os.environ.setdefault(\"HF_DATASETS_OFFLINE\", os.environ.get(\"NOTEBOOK_HF_OFFLINE\", \"1\"))\n",
    "os.environ.setdefault(\"HF_HOME\", str(Path(os.environ.get(\"CACHE_ROOT\", OUTPUT_ROOT / \"cache\")) / \"huggingface\"))\n",
    "os.environ.setdefault(\"HF_DATASETS_CACHE\", str(Path(os.environ[\"HF_HOME\"]) / \"datasets\"))\n",
    "\n",
    "def public_path(path):\n",
    "    path = Path(path)\n",
    "    replacements = [\n",
    "        (TOPIC_ROOT, \"$TOPIC_ROOT\"),\n",
    "        (PROJECT_ROOT, \"$PROJECT_ROOT\"),\n",
    "        (DATA_ROOT, \"$DATA_ROOT\"),\n",
    "        (MODEL_ROOT, \"$MODEL_ROOT\"),\n",
    "        (OUTPUT_ROOT, \"$OUTPUT_ROOT\"),\n",
    "    ]\n",
    "    value = str(path)\n",
    "    for root, label in sorted(replacements, key=lambda item: len(str(item[0])), reverse=True):\n",
    "        root_value = str(root)\n",
    "        if root_value and value.startswith(root_value):\n",
    "            return label + value[len(root_value):]\n",
    "    return value\n",
    "\n",
    "\n",
    "print(\"TOPIC_ROOT = $TOPIC_ROOT\")\n",
    "print(\"PROJECT_ROOT =\", public_path(PROJECT_ROOT))\n",
    "print(\"DATA_ROOT =\", public_path(DATA_ROOT))\n",
    "print(\"MODEL_ROOT =\", public_path(MODEL_ROOT))\n",
    "print(\"OUTPUT_ROOT =\", public_path(OUTPUT_ROOT))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def md_table(headers, rows):\n",
    "    lines = [\"| \" + \" | \".join(headers) + \" |\", \"| \" + \" | \".join([\"---\"] * len(headers)) + \" |\"]\n",
    "    for row in rows:\n",
    "        lines.append(\"| \" + \" | \".join(public_path(x) if isinstance(x, (str, Path)) else str(x) for x in row) + \" |\")\n",
    "    display(Markdown(\"\\n\".join(lines)))\n",
    "\n",
    "\n",
    "def show_json(path, max_chars=5000):\n",
    "    path = Path(path)\n",
    "    if not path.exists():\n",
    "        print(\"文件不存在：\", public_path(path))\n",
    "        return None\n",
    "    data = json.loads(path.read_text(encoding=\"utf-8\"))\n",
    "    text = json.dumps(data, ensure_ascii=False, indent=2)\n",
    "    print(text[:max_chars] + (\"\\n...\" if len(text) > max_chars else \"\"))\n",
    "    return data\n",
    "\n",
    "\n",
    "def show_video(filename, title):\n",
    "    path = ASSET_DIR / filename\n",
    "    display(Markdown(f\"**{title}**\"))\n",
    "    if path.exists():\n",
    "        rel = f\"../assets/{filename}\"\n",
    "        display(HTML(f\"<video controls muted preload='metadata' width='960'><source src='{rel}' type='video/mp4'></video>\"))\n",
    "    else:\n",
    "        print(\"缺少视频素材：\", public_path(path))\n",
    "\n",
    "\n",
    "def show_image(filename, title, width=960):\n",
    "    path = ASSET_DIR / filename\n",
    "    display(Markdown(f\"**{title}**\"))\n",
    "    if path.exists():\n",
    "        rel = f\"../assets/{filename}\"\n",
    "        display(HTML(f\"<img src='{rel}' width='{width}'>\"))\n",
    "    else:\n",
    "        print(\"缺少图片素材：\", public_path(path))\n",
    "\n",
    "\n",
    "def run_cmd_preview(command, cwd=None):\n",
    "    shown = [public_path(x) if isinstance(x, (str, Path)) else x for x in command]\n",
    "    print(\"$\", shlex.join([str(x) for x in shown]))\n",
    "    if cwd:\n",
    "        print(\"cwd =\", public_path(cwd))\n",
    "\n",
    "\n",
    "def tail_log(log_path, lines=40):\n",
    "    path = Path(log_path)\n",
    "    if not path.exists():\n",
    "        print(\"日志不存在：\", public_path(path))\n",
    "        return\n",
    "    content = path.read_text(encoding=\"utf-8\", errors=\"replace\").splitlines()\n",
    "    print(\"\\n\".join(content[-lines:]))\n",
    "\n",
    "\n",
    "def env_flag(name, default=False):\n",
    "    value = os.environ.get(name)\n",
    "    if value is None:\n",
    "        return default\n",
    "    return value.strip().lower() in {\"1\", \"true\", \"yes\", \"y\", \"on\"}\n",
    "\n",
    "\n",
    "RUN_SMOKE = env_flag(\"RUN_SMOKE\")\n",
    "RUN_LONG_TRAIN = env_flag(\"RUN_LONG_TRAIN\", True)\n",
    "RUN_EVAL = env_flag(\"RUN_EVAL\")\n",
    "EVAL_SCRIPT = Path(os.environ.get(\"EVAL_SCRIPT\", PROJECT_ROOT / \"eval_policy_success.py\"))\n",
    "\n",
    "\n",
    "_XVFB_PROCESS = None\n",
    "\n",
    "\n",
    "def ensure_xvfb_display():\n",
    "    \"\"\"Start a lightweight virtual display for headless MuJoCo evaluation.\"\"\"\n",
    "    global _XVFB_PROCESS\n",
    "    if os.environ.get(\"DISPLAY\"):\n",
    "        print(\"DISPLAY =\", os.environ[\"DISPLAY\"])\n",
    "        return None\n",
    "    xvfb_bin = shutil.which(\"Xvfb\")\n",
    "    if not xvfb_bin:\n",
    "        print(\"没有发现 Xvfb；如遇 GLFW DISPLAY 报错，请先安装 xvfb。\")\n",
    "        return None\n",
    "    display_id = os.environ.get(\"NOTEBOOK_XVFB_DISPLAY\", \":99\")\n",
    "    _XVFB_PROCESS = subprocess.Popen(\n",
    "        [xvfb_bin, display_id, \"-screen\", \"0\", \"1280x1024x24\"],\n",
    "        stdout=subprocess.DEVNULL,\n",
    "        stderr=subprocess.DEVNULL,\n",
    "    )\n",
    "    os.environ[\"DISPLAY\"] = display_id\n",
    "    print(\"已启动 Notebook 内部 Xvfb：DISPLAY =\", display_id)\n",
    "    return _XVFB_PROCESS\n",
    "\n",
    "\n",
    "def ensure_project_layout():\n",
    "    required = [PROJECT_ROOT / \"asset\" / \"example_scene_y2.xml\", PROJECT_ROOT / \"mujoco_env\"]\n",
    "    missing = [path for path in required if not path.exists()]\n",
    "    if missing:\n",
    "        print(\"当前 PROJECT_ROOT 还不是可运行工程，缺少：\")\n",
    "        for path in missing:\n",
    "            print(\" -\", public_path(path))\n",
    "        print(\"请先设置 PROJECT_ROOT，再运行训练或评估单元。\")\n",
    "        return False\n",
    "    return True\n",
    "\n",
    "\n",
    "def write_json_yaml(path, payload):\n",
    "    path = Path(path)\n",
    "    path.parent.mkdir(parents=True, exist_ok=True)\n",
    "    try:\n",
    "        import yaml\n",
    "        text = yaml.safe_dump(payload, allow_unicode=True, sort_keys=False)\n",
    "    except Exception:\n",
    "        text = json.dumps(payload, ensure_ascii=False, indent=2) + \"\\n\"\n",
    "    path.write_text(text, encoding=\"utf-8\")\n",
    "    print(\"写出配置：\", public_path(path))\n",
    "    return path\n",
    "\n",
    "\n",
    "def make_lerobot_train_config(policy_type, dataset_repo_id, dataset_root, output_dir, steps, batch_size, chunk_size, n_action_steps, seed=42):\n",
    "    save_freq = int(os.environ.get(f\"{policy_type.upper()}_SAVE_FREQ\", os.environ.get(\"SAVE_FREQ\", str(steps))))\n",
    "    return {\n",
    "        \"dataset\": {\n",
    "            \"repo_id\": dataset_repo_id,\n",
    "            \"root\": str(dataset_root),\n",
    "            \"use_imagenet_stats\": True,\n",
    "        },\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",
    "        \"output_dir\": str(output_dir),\n",
    "        \"job_name\": Path(output_dir).name,\n",
    "        \"batch_size\": int(batch_size),\n",
    "        \"steps\": int(steps),\n",
    "        \"save_freq\": max(1, save_freq),\n",
    "        \"log_freq\": 20,\n",
    "        \"num_workers\": 4,\n",
    "        \"seed\": int(seed),\n",
    "        \"resume\": False,\n",
    "        \"eval_freq\": -1,\n",
    "        \"save_checkpoint\": True,\n",
    "        \"use_policy_training_preset\": True,\n",
    "        \"wandb\": {\"enable\": False, \"disable_artifact\": True},\n",
    "    }\n",
    "\n",
    "\n",
    "def train_lerobot_config_in_notebook(config_path, enabled=False, progress_name=\"train\"):\n",
    "    \"\"\"Run LeRobot offline training directly inside the notebook kernel.\n",
    "\n",
    "    The notebook cell owns dataset creation, policy creation, optimizer steps,\n",
    "    checkpoint saving, tqdm progress, and metric JSONL writing.\n",
    "    \"\"\"\n",
    "    config_path = Path(config_path)\n",
    "    print(\"config =\", public_path(config_path))\n",
    "    if not enabled:\n",
    "        print(\"未启动。设置 RUN_SMOKE=1 或 RUN_LONG_TRAIN=1 后，本单元会直接在 Notebook 内训练。\")\n",
    "        return None\n",
    "    if not ensure_project_layout():\n",
    "        return None\n",
    "\n",
    "    import time\n",
    "    from contextlib import nullcontext\n",
    "\n",
    "    import draccus\n",
    "    import torch\n",
    "    from torch.amp import GradScaler\n",
    "\n",
    "    from lerobot.common.datasets.factory import make_dataset\n",
    "    from lerobot.common.datasets.sampler import EpisodeAwareSampler\n",
    "    from lerobot.common.datasets.utils import cycle\n",
    "    from lerobot.common.optim.factory import make_optimizer_and_scheduler\n",
    "    from lerobot.common.policies.factory import make_policy\n",
    "    from lerobot.common.policies.utils import get_device_from_parameters\n",
    "    from lerobot.common.utils.random_utils import set_seed\n",
    "    from lerobot.common.utils.train_utils import get_step_checkpoint_dir, save_checkpoint, update_last_checkpoint\n",
    "    from lerobot.common.utils.utils import get_safe_torch_device\n",
    "    from lerobot.configs.train import TrainPipelineConfig\n",
    "\n",
    "    class _CompactProgress:\n",
    "        def __init__(self, iterable, desc, total=None):\n",
    "            self.iterable = iterable\n",
    "            self.desc = desc\n",
    "            self.total = int(total if total is not None else len(iterable))\n",
    "            self.every = max(1, int(os.environ.get(\"NOTEBOOK_PROGRESS_EVERY\", \"100\")))\n",
    "            self.postfix = \"\"\n",
    "\n",
    "        def set_postfix(self, **kwargs):\n",
    "            self.postfix = \", \".join(f\"{key}={value}\" for key, value in kwargs.items())\n",
    "\n",
    "        def __iter__(self):\n",
    "            for index, value in enumerate(self.iterable, start=1):\n",
    "                yield value\n",
    "                if index == 1 or index == self.total or index % self.every == 0:\n",
    "                    suffix = f\" | {self.postfix}\" if self.postfix else \"\"\n",
    "                    print(f\"{self.desc}: {index}/{self.total}{suffix}\", flush=True)\n",
    "\n",
    "    def notebook_progress(iterable, desc, total=None):\n",
    "        return _CompactProgress(iterable, desc, total=total)\n",
    "\n",
    "    cfg = draccus.parse(TrainPipelineConfig, config_path=config_path, args=[])\n",
    "    cfg.validate()\n",
    "    if cfg.seed is not None:\n",
    "        set_seed(cfg.seed)\n",
    "\n",
    "    device = get_safe_torch_device(cfg.policy.device, log=True)\n",
    "    torch.backends.cudnn.benchmark = True\n",
    "    torch.backends.cuda.matmul.allow_tf32 = True\n",
    "\n",
    "    print(\"Creating dataset...\")\n",
    "    dataset = make_dataset(cfg)\n",
    "    print(\"Creating policy...\")\n",
    "    pretrained_override = os.environ.get(f\"{cfg.policy.type.upper()}_PRETRAINED_PATH_OVERRIDE\") or os.environ.get(\"POLICY_PRETRAINED_PATH_OVERRIDE\")\n",
    "    if pretrained_override and not cfg.resume:\n",
    "        cfg.policy.pretrained_path = str(Path(pretrained_override))\n",
    "        print(\"pretrained override =\", public_path(cfg.policy.pretrained_path))\n",
    "    elif cfg.policy.type == \"pi0\" and not cfg.resume:\n",
    "        cfg.policy.pretrained_path = \"lerobot/pi0\"\n",
    "    elif cfg.policy.type == \"smolvla\" and not cfg.resume:\n",
    "        smolvla_base_candidates = [\n",
    "            os.environ.get(\"SMOLVLA_BASE_PATH\"),\n",
    "            os.environ.get(\"SMOLVLA_PRETRAINED_BASE_PATH\"),\n",
    "            str(MODEL_ROOT / \"smolvla_base\" / \"pretrained_model\"),\n",
    "            str(MODEL_ROOT / \"lerobot_smolvla_base_legacy\"),\n",
    "            str(MODEL_ROOT / \"lerobot_smolvla_base\"),\n",
    "        ]\n",
    "        local_smolvla_base = next((Path(p) for p in smolvla_base_candidates if p and Path(p).exists()), None)\n",
    "        if local_smolvla_base is not None:\n",
    "            cfg.policy.pretrained_path = str(local_smolvla_base)\n",
    "            print(\"local smolvla base =\", public_path(cfg.policy.pretrained_path))\n",
    "        else:\n",
    "            cfg.policy.pretrained_path = \"lerobot/smolvla_base\"\n",
    "    policy = make_policy(cfg=cfg.policy, ds_meta=dataset.meta)\n",
    "\n",
    "    # Compatibility for newer Transformers: PaliGemmaForConditionalGeneration may expose\n",
    "    # language_model as GemmaModel directly, while this LeRobot Pi0 code expects\n",
    "    # language_model.model.  Use a non-Module proxy so checkpoints/state_dict stay clean.\n",
    "    if cfg.policy.type == \"pi0\":\n",
    "        try:\n",
    "            lm = policy.model.paligemma_with_expert.paligemma.language_model\n",
    "            if not hasattr(lm, \"model\"):\n",
    "                class _LanguageModelCoreProxy:\n",
    "                    def __init__(self, core):\n",
    "                        self._core = core\n",
    "\n",
    "                    def __getattr__(self, name):\n",
    "                        return getattr(self._core, name)\n",
    "\n",
    "                object.__setattr__(lm, \"model\", _LanguageModelCoreProxy(lm))\n",
    "                print(\"patched Pi0 PaliGemma language_model.model compatibility proxy\")\n",
    "        except Exception as exc:\n",
    "            print(f\"Pi0 PaliGemma compatibility patch skipped: {exc}\")\n",
    "\n",
    "    policy.to(device)\n",
    "    policy.train()\n",
    "\n",
    "    optimizer, lr_scheduler = make_optimizer_and_scheduler(cfg, policy)\n",
    "    grad_scaler = GradScaler(device.type, enabled=cfg.policy.use_amp)\n",
    "\n",
    "    def _dataset_column_values(name):\n",
    "        hf_dataset = getattr(dataset, \"hf_dataset\", None)\n",
    "        if hf_dataset is None or name not in getattr(hf_dataset, \"column_names\", []):\n",
    "            return None\n",
    "        values = hf_dataset[name]\n",
    "        try:\n",
    "            return list(values)\n",
    "        except TypeError:\n",
    "            return [values[i] for i in range(len(values))]\n",
    "\n",
    "    def _task_name_map():\n",
    "        meta = getattr(dataset, \"meta\", None)\n",
    "        tasks = getattr(meta, \"tasks\", None)\n",
    "        if tasks is None:\n",
    "            return {}\n",
    "        if isinstance(tasks, dict):\n",
    "            return {int(k): str(v) for k, v in tasks.items()}\n",
    "        try:\n",
    "            return {int(k): str(v) for k, v in dict(tasks).items()}\n",
    "        except Exception:\n",
    "            return {}\n",
    "\n",
    "    def _make_weighted_sampler(generator):\n",
    "        mode = os.environ.get(\"NOTEBOOK_FRAME_WEIGHT_MODE\", \"\").strip().lower()\n",
    "        if not mode or mode in {\"0\", \"none\", \"off\", \"false\"}:\n",
    "            return None, {\"mode\": \"none\"}\n",
    "        weights = torch.ones(len(dataset), dtype=torch.double)\n",
    "        info = {\"mode\": mode, \"num_frames\": len(dataset)}\n",
    "\n",
    "        if \"blue\" in mode:\n",
    "            blue_weight = float(os.environ.get(\"NOTEBOOK_BLUE_WEIGHT\", \"2.0\"))\n",
    "            mask = [False] * len(dataset)\n",
    "            task_indices = _dataset_column_values(\"task_index\")\n",
    "            task_names = _task_name_map()\n",
    "            if task_indices is not None and task_names:\n",
    "                for idx, task_index in enumerate(task_indices):\n",
    "                    task_text = task_names.get(int(task_index), \"\").lower()\n",
    "                    mask[idx] = (\"blue\" in task_text) or (\"蓝\" in task_text)\n",
    "            else:\n",
    "                for column in [\"task\", \"language_instruction\", \"instruction\"]:\n",
    "                    values = _dataset_column_values(column)\n",
    "                    if values is None:\n",
    "                        continue\n",
    "                    for idx, value in enumerate(values):\n",
    "                        text = str(value).lower()\n",
    "                        mask[idx] = (\"blue\" in text) or (\"蓝\" in text)\n",
    "                    break\n",
    "            blue_count = int(sum(mask))\n",
    "            if blue_count == 0:\n",
    "                print(\"警告：NOTEBOOK_FRAME_WEIGHT_MODE=blue 但没有识别到 blue/蓝 指令帧，采样退回均匀。\")\n",
    "            else:\n",
    "                for idx, is_blue in enumerate(mask):\n",
    "                    if is_blue:\n",
    "                        weights[idx] *= blue_weight\n",
    "            info.update({\"blue_weight\": blue_weight, \"blue_frames\": blue_count})\n",
    "\n",
    "        weight_file = os.environ.get(\"NOTEBOOK_FRAME_WEIGHT_JSON\")\n",
    "        if weight_file:\n",
    "            payload = json.loads(Path(weight_file).read_text(encoding=\"utf-8\"))\n",
    "            for key, value in payload.items():\n",
    "                weights[int(key)] *= float(value)\n",
    "            info.update({\"weight_json\": public_path(weight_file), \"json_entries\": len(payload)})\n",
    "\n",
    "        if float(weights.sum()) <= 0:\n",
    "            raise ValueError(\"采样权重总和为 0。\")\n",
    "        sampler = torch.utils.data.WeightedRandomSampler(\n",
    "            weights=weights,\n",
    "            num_samples=len(weights),\n",
    "            replacement=True,\n",
    "            generator=generator,\n",
    "        )\n",
    "        info.update(\n",
    "            {\n",
    "                \"weight_min\": float(weights.min()),\n",
    "                \"weight_max\": float(weights.max()),\n",
    "                \"weight_mean\": float(weights.mean()),\n",
    "            }\n",
    "        )\n",
    "        return sampler, info\n",
    "\n",
    "    generator = torch.Generator()\n",
    "    if cfg.seed is not None:\n",
    "        generator.manual_seed(int(cfg.seed))\n",
    "\n",
    "    weighted_sampler, sampler_info = _make_weighted_sampler(generator)\n",
    "    if weighted_sampler is not None:\n",
    "        shuffle = False\n",
    "        sampler = weighted_sampler\n",
    "        print(\"Notebook weighted sampler =\", json.dumps(sampler_info, ensure_ascii=False))\n",
    "    elif hasattr(cfg.policy, \"drop_n_last_frames\"):\n",
    "        shuffle = False\n",
    "        sampler = EpisodeAwareSampler(\n",
    "            dataset.episode_data_index,\n",
    "            drop_n_last_frames=cfg.policy.drop_n_last_frames,\n",
    "            shuffle=True,\n",
    "        )\n",
    "    else:\n",
    "        shuffle = True\n",
    "        sampler = None\n",
    "\n",
    "    dataloader = torch.utils.data.DataLoader(\n",
    "        dataset,\n",
    "        num_workers=cfg.num_workers,\n",
    "        batch_size=cfg.batch_size,\n",
    "        shuffle=shuffle,\n",
    "        sampler=sampler,\n",
    "        generator=generator if sampler is None else None,\n",
    "        pin_memory=device.type != \"cpu\",\n",
    "        drop_last=False,\n",
    "    )\n",
    "    dl_iter = cycle(dataloader)\n",
    "\n",
    "    output_dir = Path(cfg.output_dir)\n",
    "    output_dir.mkdir(parents=True, exist_ok=True)\n",
    "    metrics_path = output_dir / \"notebook_train_metrics.jsonl\"\n",
    "    num_learnable = sum(p.numel() for p in policy.parameters() if p.requires_grad)\n",
    "    num_total = sum(p.numel() for p in policy.parameters())\n",
    "    print(f\"output_dir = {public_path(output_dir)}\")\n",
    "    print(f\"steps = {cfg.steps}, batch_size = {cfg.batch_size}, frames = {dataset.num_frames}, episodes = {dataset.num_episodes}\")\n",
    "    print(f\"learnable_params = {num_learnable:,}, total_params = {num_total:,}\")\n",
    "\n",
    "    last_metrics = None\n",
    "    progress = notebook_progress(range(1, cfg.steps + 1), desc=progress_name, total=cfg.steps)\n",
    "    start_all = time.perf_counter()\n",
    "    for step in progress:\n",
    "        load_start = time.perf_counter()\n",
    "        batch = next(dl_iter)\n",
    "        data_s = time.perf_counter() - load_start\n",
    "        for key, value in batch.items():\n",
    "            if isinstance(value, torch.Tensor):\n",
    "                batch[key] = value.to(device, non_blocking=True)\n",
    "\n",
    "        update_start = time.perf_counter()\n",
    "        device_from_policy = get_device_from_parameters(policy)\n",
    "        with torch.autocast(device_type=device_from_policy.type) if cfg.policy.use_amp else nullcontext():\n",
    "            loss, output_dict = policy.forward(batch)\n",
    "        grad_scaler.scale(loss).backward()\n",
    "        grad_scaler.unscale_(optimizer)\n",
    "        grad_norm = torch.nn.utils.clip_grad_norm_(\n",
    "            policy.parameters(),\n",
    "            cfg.optimizer.grad_clip_norm,\n",
    "            error_if_nonfinite=False,\n",
    "        )\n",
    "        grad_scaler.step(optimizer)\n",
    "        grad_scaler.update()\n",
    "        optimizer.zero_grad()\n",
    "        if lr_scheduler is not None:\n",
    "            lr_scheduler.step()\n",
    "        if hasattr(policy, \"update\"):\n",
    "            policy.update()\n",
    "        update_s = time.perf_counter() - update_start\n",
    "\n",
    "        is_log_step = cfg.log_freq > 0 and (step % cfg.log_freq == 0 or step == 1 or step == cfg.steps)\n",
    "        is_saving_step = cfg.save_checkpoint and (step % cfg.save_freq == 0 or step == cfg.steps)\n",
    "        if is_log_step:\n",
    "            last_metrics = {\n",
    "                \"step\": step,\n",
    "                \"loss\": float(loss.detach().cpu()),\n",
    "                \"grad_norm\": float(grad_norm.detach().cpu()) if hasattr(grad_norm, \"detach\") else float(grad_norm),\n",
    "                \"lr\": float(optimizer.param_groups[0][\"lr\"]),\n",
    "                \"update_s\": float(update_s),\n",
    "                \"data_s\": float(data_s),\n",
    "                \"elapsed_s\": float(time.perf_counter() - start_all),\n",
    "            }\n",
    "            with metrics_path.open(\"a\", encoding=\"utf-8\") as f:\n",
    "                f.write(json.dumps(last_metrics, ensure_ascii=False) + \"\\n\")\n",
    "            progress.set_postfix(\n",
    "                loss=f\"{last_metrics['loss']:.4f}\",\n",
    "                lr=f\"{last_metrics['lr']:.1e}\",\n",
    "                updt_s=f\"{last_metrics['update_s']:.3f}\",\n",
    "            )\n",
    "        if is_saving_step:\n",
    "            checkpoint_dir = get_step_checkpoint_dir(cfg.output_dir, cfg.steps, step)\n",
    "            print(f\"\\nSaving checkpoint at step {step}: {public_path(checkpoint_dir)}\")\n",
    "            save_checkpoint(checkpoint_dir, step, cfg, policy, optimizer, lr_scheduler)\n",
    "            update_last_checkpoint(checkpoint_dir)\n",
    "\n",
    "    print(\"训练完成。metrics =\", public_path(metrics_path))\n",
    "    if last_metrics is not None:\n",
    "        print(json.dumps(last_metrics, ensure_ascii=False, indent=2))\n",
    "    return {\"output_dir\": output_dir, \"metrics_path\": metrics_path, \"last_metrics\": last_metrics}\n",
    "\n",
    "\n",
    "def load_eval_module():\n",
    "    import importlib.util\n",
    "\n",
    "    if not EVAL_SCRIPT.exists():\n",
    "        raise FileNotFoundError(f\"评估脚本不存在：{public_path(EVAL_SCRIPT)}\")\n",
    "    spec = importlib.util.spec_from_file_location(\"notebook_eval_policy_success\", EVAL_SCRIPT)\n",
    "    module = importlib.util.module_from_spec(spec)\n",
    "    assert spec.loader is not None\n",
    "    spec.loader.exec_module(module)\n",
    "    return module\n",
    "\n",
    "\n",
    "def run_eval_policy_in_notebook(\n",
    "    policy_name,\n",
    "    policy_path,\n",
    "    result_path,\n",
    "    episodes,\n",
    "    seed_start,\n",
    "    render=False,\n",
    "    enabled=False,\n",
    "    repo_id=None,\n",
    "    dataset_root=None,\n",
    "):\n",
    "    print(\"policy =\", policy_name)\n",
    "    print(\"policy_path =\", public_path(policy_path))\n",
    "    print(\"result =\", public_path(result_path))\n",
    "    if not enabled:\n",
    "        print(\"未启动。设置 RUN_EVAL=1 后，本单元会在 Notebook 内直接加载策略并闭环评估。\")\n",
    "        return None\n",
    "    if not ensure_project_layout():\n",
    "        return None\n",
    "\n",
    "    import argparse\n",
    "    from contextlib import contextmanager\n",
    "\n",
    "    @contextmanager\n",
    "    def pushd(path):\n",
    "        old = Path.cwd()\n",
    "        os.chdir(path)\n",
    "        try:\n",
    "            yield\n",
    "        finally:\n",
    "            os.chdir(old)\n",
    "\n",
    "    ensure_xvfb_display()\n",
    "    module = load_eval_module()\n",
    "    result_path = Path(result_path)\n",
    "    result_path.parent.mkdir(parents=True, exist_ok=True)\n",
    "    if result_path.exists():\n",
    "        result_path.unlink()\n",
    "\n",
    "    args = argparse.Namespace(\n",
    "        policy=policy_name,\n",
    "        episodes=int(episodes),\n",
    "        seed_start=int(seed_start),\n",
    "        max_action_steps=int(os.environ.get(\"EVAL_MAX_ACTION_STEPS\", \"400\")),\n",
    "        hz=float(os.environ.get(\"EVAL_HZ\", \"20\")),\n",
    "        render=bool(render),\n",
    "        output_jsonl=result_path,\n",
    "        device=os.environ.get(\"EVAL_DEVICE\", \"cuda\"),\n",
    "        reset_policy_each_action=env_flag(\"EVAL_RESET_POLICY_EACH_ACTION\", False),\n",
    "        act_n_action_steps=None,\n",
    "        act_force_dataset_gripper=False,\n",
    "        act_clamp_timestamp=False,\n",
    "        act_policy_path=Path(policy_path),\n",
    "        act_repo_id=repo_id or \"datawhale_eai_pnp\",\n",
    "        act_dataset_root=Path(dataset_root or \"./demo_data\"),\n",
    "        act_episode_timestamp_offsets=\"\",\n",
    "        act_episode_source_flags=\"\",\n",
    "        physical_success=env_flag(\"EVAL_PHYSICAL_SUCCESS\", True),\n",
    "        physical_min_lift=float(os.environ.get(\"EVAL_PHYSICAL_MIN_LIFT\", \"0.06\")),\n",
    "        physical_min_lift_steps=int(os.environ.get(\"EVAL_PHYSICAL_MIN_LIFT_STEPS\", \"3\")),\n",
    "        physical_final_upright_cos=float(os.environ.get(\"EVAL_PHYSICAL_FINAL_UPRIGHT_COS\", \"0.85\")),\n",
    "        smolvla_policy_path=Path(policy_path),\n",
    "        pi0_policy_path=Path(policy_path),\n",
    "        pi0_repo_id=repo_id or os.environ.get(\"PI0_DATASET_REPO_ID\", \"datawhale_eai_pnp_language\"),\n",
    "        pi0_dataset_root=Path(dataset_root or os.environ.get(\"PI0_DATASET_ROOT\", \"./demo_data_language\")),\n",
    "    )\n",
    "\n",
    "    with pushd(PROJECT_ROOT):\n",
    "        if policy_name == \"act\":\n",
    "            policy = module.make_act_policy(\n",
    "                args.device,\n",
    "                args.act_policy_path,\n",
    "                args.act_repo_id,\n",
    "                args.act_dataset_root,\n",
    "                n_action_steps=args.act_n_action_steps,\n",
    "                episode_timestamp_offsets=args.act_episode_timestamp_offsets,\n",
    "                episode_source_flags=args.act_episode_source_flags,\n",
    "            )\n",
    "            rollout = module.rollout_act\n",
    "        elif policy_name == \"smolvla\":\n",
    "            policy = module.make_smolvla_policy(args.device, args.smolvla_policy_path)\n",
    "            rollout = module.rollout_language_policy\n",
    "        elif policy_name == \"pi0\":\n",
    "            policy = module.make_pi0_policy(args.device, args.pi0_policy_path, args.pi0_repo_id, args.pi0_dataset_root)\n",
    "            rollout = module.rollout_language_policy\n",
    "        else:\n",
    "            raise ValueError(policy_name)\n",
    "\n",
    "        rows = []\n",
    "        for offset in notebook_progress(range(args.episodes), desc=f\"{policy_name} eval\", total=args.episodes):\n",
    "            seed = args.seed_start + offset\n",
    "            row = rollout(args, policy, seed)\n",
    "            rows.append(row)\n",
    "            with result_path.open(\"a\", encoding=\"utf-8\") as f:\n",
    "                f.write(json.dumps(row, ensure_ascii=False) + \"\\n\")\n",
    "            print(json.dumps(row, ensure_ascii=False))\n",
    "    summarize_jsonl(result_path)\n",
    "    return rows\n",
    "\n",
    "\n",
    "def list_checkpoints(run_dir):\n",
    "    run_dir = Path(run_dir)\n",
    "    candidates = []\n",
    "    for pattern in [\"checkpoints/*/pretrained_model\", \"checkpoint*/pretrained_model\", \"*/pretrained_model\", \"pretrained_model\"]:\n",
    "        candidates.extend(run_dir.glob(pattern))\n",
    "    unique = sorted(set(candidates))\n",
    "    if not unique:\n",
    "        print(\"尚未发现 checkpoint：\", public_path(run_dir))\n",
    "        return []\n",
    "    for path in unique:\n",
    "        print(\" -\", public_path(path))\n",
    "    return unique\n",
    "\n",
    "\n",
    "def resolve_eval_policy(default_path, trained_run_dir=None, env_name=None):\n",
    "    if env_name and os.environ.get(env_name):\n",
    "        path = Path(os.environ[env_name])\n",
    "        print(\"评估使用环境变量指定权重：\", public_path(path))\n",
    "        return path\n",
    "    if trained_run_dir is not None and env_flag(\"EVAL_USE_LONG_TRAIN\"):\n",
    "        checkpoints = list_checkpoints(trained_run_dir)\n",
    "        if checkpoints:\n",
    "            path = checkpoints[-1]\n",
    "            print(\"评估使用本次长训最新 checkpoint：\", public_path(path))\n",
    "            return path\n",
    "        print(\"未找到本次长训 checkpoint，回退到保护权重。\")\n",
    "    path = Path(default_path)\n",
    "    print(\"评估使用保护/预训练权重：\", public_path(path))\n",
    "    return path\n",
    "\n",
    "\n",
    "def summarize_jsonl(path):\n",
    "    path = Path(path)\n",
    "    if not path.exists():\n",
    "        print(\"结果 JSONL 尚不存在：\", public_path(path))\n",
    "        return None\n",
    "    rows = [json.loads(line) for line in path.read_text(encoding=\"utf-8\").splitlines() if line.strip()]\n",
    "    total = len(rows)\n",
    "    legacy = sum(bool(row.get(\"success\") or row.get(\"legacy_success\")) for row in rows)\n",
    "    if rows and all(\"physical_success\" in row for row in rows):\n",
    "        physical_count = sum(bool(row.get(\"physical_success\")) for row in rows)\n",
    "        physical_text = str(physical_count) + \"/\" + str(total)\n",
    "    else:\n",
    "        physical_text = \"未记录\"\n",
    "    md_table(\n",
    "        [\"结果文件\", \"episodes\", \"legacy_success\", \"physical_success\"],\n",
    "        [(public_path(path), total, f\"{legacy}/{total}\", physical_text)],\n",
    "    )\n",
    "    return rows\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 运行控制台：先在这里选择本次实验\n\n下面这个单元格是三个模型通用的开关。默认全部关闭，避免打开 Notebook 后误启动数小时训练。\n\n建议顺序：\n\n1. 第一次把 `RUN_SMOKE` 设为 `True`，确认环境、数据、模型和一次反向传播都正常。\n2. 确认无误后把 `RUN_LONG_TRAIN` 设为 `True`，执行后面的真实训练单元格，Notebook 会显示实时 `tqdm`、loss、耗时和 checkpoint。\n3. 训练完成后把 `RUN_EVAL` 设为 `True`，执行严格评估单元格，结果会汇总为 `x/y` 成功率，并在有渲染条件时保存视频。\n4. 如果要复现正式保护结果，再把 `RUN_PROTECTED_TRAIN` 设为 `True`，执行对应 protected recipe 单元格。它和普通教学长训是两条明确标注的训练谱系。\n\n训练循环和评估循环都在 Notebook Python kernel 内执行，不是 `cat` 静态日志，也不是把训练交给外部 shell 脚本。\n"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# RUN_CONTROL_CELL\n",
    "# 只修改下面四个布尔值，然后按顺序执行后面的单元格。\n",
    "RUN_SMOKE = False\n",
    "RUN_LONG_TRAIN = True\n",
    "RUN_EVAL = False\n",
    "RUN_PROTECTED_TRAIN = False\n",
    "\n",
    "print({\n",
    "    'RUN_SMOKE': RUN_SMOKE,\n",
    "    'RUN_LONG_TRAIN': RUN_LONG_TRAIN,\n",
    "    'RUN_EVAL': RUN_EVAL,\n",
    "    'RUN_PROTECTED_TRAIN': RUN_PROTECTED_TRAIN,\n",
    "})\n",
    "print('长训通常需要几十分钟到数小时；评估会逐 episode 闭环运行，时间更长。')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## Checkpoint 1：当前恢复边界\n"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 结果对照\n",
    "\n",
    "| 面板 | 当前结果 | 结论 |\n",
    "| --- | --- | --- |\n",
    "| unseen10 | 9/10 | 超过早期 7/10 |\n",
    "| full14 | 12/14 | 红 6/7，蓝 6/7 |\n",
    "| hard8 | 6/8 | 未恢复历史 8/8，不能写成完全复现 |\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## Checkpoint 2：gated 权限和缓存检查\n\n            预计耗时：几秒。Pi0 需要确认模型权限、Hugging Face 缓存和训练数据路径，Notebook 只检查状态，不打印任何 token。\n"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## Checkpoint 1.5：Pi0 保护权重的训练谱系\n\n            Pi0 的 `12/14` 不是从一个普通 `PI0_STEPS=8500` 单元自然保证出来的；它来自 clean40 success-only 数据、蓝杯补权重和 S7500→S8500 门禁选择。Notebook 默认轻量训练用于理解流程，protected recipe 用于复现报告口径。\n"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "### 结果口径对齐：本轮小面板、保护评估与 hard 面板\n\n            Pi0 的课堂小面板和保护权重门禁不是同一个结果。小面板只说明本轮普通长训 checkpoint 没有救回来；保护结果来自 clean40 success-only + blue2x 的 S8500 训练谱系。\n\n            | 口径 | 成功率 | 评估范围 | 教学解释 |\n            | --- | --- | --- | --- |\n            | 本轮 Notebook 小面板 | `0/4` | post-long eval seed3000-3003 | 说明普通本轮 checkpoint 仍失败，不能写成成功 |\n            | 正式保护评估 | `12/14` | full14 strict physical success | 当前 Pi0 发布候选，红 `6/7`、蓝 `6/7` |\n            | unseen 面板 | `9/10` | unseen seeds 3010-3019 | 比早期 unseen `7/10` 更好 |\n            | hard 面板 | `6/8` | hard8 | 尚未恢复历史 hard `8/8`，只能写成部分恢复 |\n"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# PROTECTED_RECIPE_CELL\n",
    "rows = [\n",
    "    (\"clean40-only S7500\", \"pi0_clean40_successonly\", \"成功轨迹过滤；不混失败轨迹\", \"7500 steps\", \"canary 2/4\", \"父权重，红杯开始恢复\"),\n",
    "    (\"blue2x S8500\", \"clean40 + 蓝杯加权/补采样\", \"蓝杯权重提高，但保留红杯保护门禁\", \"+1000 steps\", \"canary 3/4 -> full14\", \"保护结果 12/14，unseen 9/10，hard8 6/8\"),\n",
    "    (\"blue2x S9500\", \"同上\", \"继续加步数\", \"+2000 steps\", \"回落到 2/4\", \"说明不能盲目长训，S8500 才是选择点\"),\n",
    "]\n",
    "md_table([\"阶段\", \"数据\", \"采样/筛选策略\", \"训练步数\", \"门禁\", \"结论\"], rows)\n",
    "\n",
    "protected_env = {\n",
    "    \"TEACHING_RECIPE\": \"protected\",\n",
    "    \"PI0_TRAIN_DATA_ROOT\": str(DATA_ROOT / \"pi0_clean40_successonly\"),\n",
    "    \"PI0_STEPS\": \"8500\",\n",
    "    \"PI0_BATCH_SIZE\": \"2\",\n",
    "    \"PI0_EVAL_EPISODES\": \"14\",\n",
    "    \"PI0_EVAL_SEED_START\": \"3000\",\n",
    "    \"PI0_POLICY_PATH\": str(MODEL_ROOT / \"pi0_clean40_successonly_blue2x_from_s7500_s2000_v1\" / \"checkpoints\" / \"001000\" / \"pretrained_model\"),\n",
    "}\n",
    "print(json.dumps(protected_env, ensure_ascii=False, indent=2))\n",
    "print(\"The formal recipe uses clean40 S7500 followed by the blue2x continuation. Evaluate the model produced by this workflow.\")\n",
    "print(\"The fixed protocol results are listed above; new runs write their own per-episode evaluation files.\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "### 可执行 protected 训练：clean40 S7500 + blue2x S8500\n\n            设置 `RUN_PROTECTED_TRAIN=1` 后，这一格会先用 clean40 success-only 数据训练 S7500，再从 S7500 初始化做 blue2x 续训到 S8500。这样 Notebook 训练路径和 `12/14` 保护结果的因果链一致。\n\n            如果课堂或云端时长不够，也可以提供 `PI0_PARENT_CHECKPOINT_PATH=/path/to/s7500/pretrained_model`，Notebook 会跳过 S7500 父权重长训，直接从该父权重原生执行 blue2x 续训。这不是换方法，而是复用已经训练好的中间保护节点。\n"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# PROTECTED_TRAIN_CELL\n",
    "protected_train_enabled = globals().get(\"RUN_PROTECTED_TRAIN\", env_flag(\"RUN_PROTECTED_TRAIN\", False))\n",
    "if not protected_train_enabled:\n",
    "    print(\"未启动。设置 RUN_PROTECTED_TRAIN=1 后，本单元会原生训练 Pi0 protected recipe。\")\n",
    "else:\n",
    "    DATASET_REPO_ID = globals().get(\"DATASET_REPO_ID\", \"datawhale_eai_pnp_pi0_clean_oracle_y060z000_3100_3139_g8_rebuild_v1\")\n",
    "    TRAIN_DATA_ROOT = globals().get(\"TRAIN_DATA_ROOT\", Path(os.environ.get(\"PI0_TRAIN_DATA_ROOT\", DATA_ROOT / \"pi0_clean40_successonly\")))\n",
    "    CONFIG_DIR = OUTPUT_ROOT / \"configs\"\n",
    "    RUN_ROOT = OUTPUT_ROOT / \"runs\" / \"pi0_protected_recipe\"\n",
    "    CLEAN_OUTPUT = RUN_ROOT / \"clean40_s7500\"\n",
    "    BLUE_OUTPUT = RUN_ROOT / \"blue2x_s8500\"\n",
    "    clean_config = make_lerobot_train_config(\n",
    "        \"pi0\", DATASET_REPO_ID, TRAIN_DATA_ROOT, CLEAN_OUTPUT,\n",
    "        steps=int(os.environ.get(\"PI0_CLEAN_STEPS\", \"7500\")),\n",
    "        batch_size=int(os.environ.get(\"PI0_BATCH_SIZE\", \"2\")),\n",
    "        chunk_size=50,\n",
    "        n_action_steps=50,\n",
    "    )\n",
    "    blue_config = make_lerobot_train_config(\n",
    "        \"pi0\", DATASET_REPO_ID, TRAIN_DATA_ROOT, BLUE_OUTPUT,\n",
    "        steps=int(os.environ.get(\"PI0_BLUE_STEPS\", \"1000\")),\n",
    "        batch_size=int(os.environ.get(\"PI0_BATCH_SIZE\", \"2\")),\n",
    "        chunk_size=50,\n",
    "        n_action_steps=50,\n",
    "    )\n",
    "    clean_path = write_json_yaml(CONFIG_DIR / \"pi0_protected_clean40_s7500.yaml\", clean_config)\n",
    "    blue_path = write_json_yaml(CONFIG_DIR / \"pi0_protected_blue2x_s8500.yaml\", blue_config)\n",
    "    parent_checkpoint = os.environ.get(\"PI0_PARENT_CHECKPOINT_PATH\") or os.environ.get(\"PI0_S7500_CHECKPOINT_PATH\")\n",
    "    if parent_checkpoint:\n",
    "        clean_ckpt = Path(parent_checkpoint)\n",
    "        if not clean_ckpt.exists():\n",
    "            raise FileNotFoundError(f\"PI0_PARENT_CHECKPOINT_PATH does not exist: {clean_ckpt}\")\n",
    "        print(\"使用已提供的 S7500 父权重，跳过 clean40 长训：\", public_path(clean_ckpt))\n",
    "    else:\n",
    "        train_lerobot_config_in_notebook(clean_path, enabled=True, progress_name=\"Pi0 protected clean40\")\n",
    "        clean_ckpt = list_checkpoints(CLEAN_OUTPUT)[-1]\n",
    "    old_override = os.environ.get(\"PI0_PRETRAINED_PATH_OVERRIDE\")\n",
    "    old_mode = os.environ.get(\"NOTEBOOK_FRAME_WEIGHT_MODE\")\n",
    "    old_blue = os.environ.get(\"NOTEBOOK_BLUE_WEIGHT\")\n",
    "    os.environ[\"PI0_PRETRAINED_PATH_OVERRIDE\"] = str(clean_ckpt)\n",
    "    os.environ[\"NOTEBOOK_FRAME_WEIGHT_MODE\"] = \"blue\"\n",
    "    os.environ[\"NOTEBOOK_BLUE_WEIGHT\"] = os.environ.get(\"PI0_BLUE_WEIGHT\", \"2.0\")\n",
    "    try:\n",
    "        train_lerobot_config_in_notebook(blue_path, enabled=True, progress_name=\"Pi0 protected blue2x\")\n",
    "    finally:\n",
    "        if old_override is None:\n",
    "            os.environ.pop(\"PI0_PRETRAINED_PATH_OVERRIDE\", None)\n",
    "        else:\n",
    "            os.environ[\"PI0_PRETRAINED_PATH_OVERRIDE\"] = old_override\n",
    "        if old_mode is None:\n",
    "            os.environ.pop(\"NOTEBOOK_FRAME_WEIGHT_MODE\", None)\n",
    "        else:\n",
    "            os.environ[\"NOTEBOOK_FRAME_WEIGHT_MODE\"] = old_mode\n",
    "        if old_blue is None:\n",
    "            os.environ.pop(\"NOTEBOOK_BLUE_WEIGHT\", None)\n",
    "        else:\n",
    "            os.environ[\"NOTEBOOK_BLUE_WEIGHT\"] = old_blue\n",
    "    print(\"protected candidate checkpoints:\")\n",
    "    list_checkpoints(BLUE_OUTPUT)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": "HF_TOKEN 已注入： False\n"
    },
    {
     "output_type": "display_data",
     "data": {
      "text/markdown": "| 项目 | 说明 |\n| --- | --- |\n| HF_TOKEN | 只检查是否存在，不在 Notebook 打印 token |\n| HF_HOME | 已设置 |\n| 模型缓存 | 建议放到 $CACHE_ROOT/huggingface 或云平台持久盘 |"
     },
     "metadata": {}
    }
   ],
   "source": [
    "print(\"HF_TOKEN 已注入：\", bool(os.environ.get(\"HF_TOKEN\")))\n",
    "rows = [\n",
    "    (\"HF_TOKEN\", \"只检查是否存在，不在 Notebook 打印 token\"),\n",
    "    (\"HF_HOME\", \"已设置\" if os.environ.get(\"HF_HOME\") else \"未设置\"),\n",
    "    (\"模型缓存\", \"建议放到 $CACHE_ROOT/huggingface 或云平台持久盘\"),\n",
    "]\n",
    "md_table([\"项目\", \"说明\"], rows)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## Checkpoint 3：生成配置并真实启动训练/续训\n\n            预计耗时：smoke 约 1-5 分钟；`PI0_STEPS=8500` 属于长训，可能需要数小时。  \n            这一步是真实训练入口：开 `RUN_SMOKE=1` 做环境验证，开 `RUN_LONG_TRAIN=1` 后会在 Notebook kernel 内直接跑 Pi0 训练循环，并显示 tqdm 进度。\n"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATASET_REPO_ID = \"datawhale_eai_pnp_pi0_clean_oracle_y060z000_3100_3139_g8_rebuild_v1\"\n",
    "TRAIN_DATA_ROOT = Path(os.environ.get(\"PI0_TRAIN_DATA_ROOT\", DATA_ROOT / \"pi0_clean40_successonly\"))\n",
    "PI0_POLICY_PATH = Path(os.environ.get(\"PI0_POLICY_PATH\", MODEL_ROOT / \"pi0_clean40_successonly_blue2x_from_s7500_s2000_v1\" / \"checkpoints\" / \"001000\" / \"pretrained_model\"))\n",
    "CONFIG_DIR = OUTPUT_ROOT / \"configs\"\n",
    "LOG_DIR = OUTPUT_ROOT / \"logs\"\n",
    "RUN_ROOT = OUTPUT_ROOT / \"runs\" / \"pi0_s8500_repro\"\n",
    "SMOKE_OUTPUT = RUN_ROOT / \"smoke\"\n",
    "LONG_OUTPUT = RUN_ROOT / \"clean40_blue2x_s8500\"\n",
    "\n",
    "smoke_config = make_lerobot_train_config(\n",
    "    \"pi0\", DATASET_REPO_ID, TRAIN_DATA_ROOT, SMOKE_OUTPUT,\n",
    "    steps=2, batch_size=1, chunk_size=50, n_action_steps=50,\n",
    ")\n",
    "long_config = make_lerobot_train_config(\n",
    "    \"pi0\", DATASET_REPO_ID, TRAIN_DATA_ROOT, LONG_OUTPUT,\n",
    "    steps=int(os.environ.get(\"PI0_STEPS\", \"8500\")),\n",
    "    batch_size=int(os.environ.get(\"PI0_BATCH_SIZE\", \"2\")),\n",
    "    chunk_size=50,\n",
    "    n_action_steps=50,\n",
    ")\n",
    "smoke_config_path = write_json_yaml(CONFIG_DIR / \"pi0_smoke.yaml\", smoke_config)\n",
    "long_config_path = write_json_yaml(CONFIG_DIR / \"pi0_clean40_blue2x_s8500.yaml\", long_config)\n",
    "\n",
    "train_lerobot_config_in_notebook(smoke_config_path, enabled=RUN_SMOKE, progress_name=\"Pi0 smoke\")\n",
    "train_lerobot_config_in_notebook(long_config_path, enabled=RUN_LONG_TRAIN, progress_name=\"Pi0 long train\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## Checkpoint 4：实时查看训练日志和 checkpoint\n\n            预计耗时：几秒。训练中可以反复执行这一格，观察 loss、step、保存点和是否有 OOM/NaN。\n"
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": "smoke metrics:\n日志不存在： $OUTPUT_ROOT/runs/pi0_s8500_repro/smoke/notebook_train_metrics.jsonl\n\nlong train metrics:\n日志不存在： $OUTPUT_ROOT/runs/pi0_s8500_repro/clean40_blue2x_s8500/notebook_train_metrics.jsonl\n\ncheckpoints:\n尚未发现 checkpoint： $OUTPUT_ROOT/runs/pi0_s8500_repro/clean40_blue2x_s8500\n"
    }
   ],
   "source": [
    "print(\"smoke metrics:\")\n",
    "tail_log(SMOKE_OUTPUT / \"notebook_train_metrics.jsonl\", lines=20)\n",
    "print(\"\\nlong train metrics:\")\n",
    "tail_log(LONG_OUTPUT / \"notebook_train_metrics.jsonl\", lines=40)\n",
    "print(\"\\ncheckpoints:\")\n",
    "list_checkpoints(LONG_OUTPUT)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## Checkpoint 5：已完成长训的实测对照\n\n            这是此前已在 AMD 设备上完成的结果对照，帮助学习者知道 Pi0 当前恢复到了哪里；它不替代本次 Notebook 的真实运行输出。\n"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 结果对照\n",
    "\n",
    "| 阶段 | 门禁 | 结论 |\n",
    "| --- | --- | --- |\n",
    "| S7500 | canary2/4 | 继续训练的父 checkpoint |\n",
    "| S8500 | canary3/4 | 进入 full14 扩展评估 |\n",
    "| full14 median | 12/14 | 作为当前保护版本 |\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## Checkpoint 6：Notebook 内严格闭环评估\n\n            预计耗时：14 个 episode 常见为几十分钟；可以用 `PI0_EVAL_EPISODES=4` 做快速门禁。  \n            本单元会在 Notebook kernel 内直接加载 Pi0，并使用 checkpoint 对应的 8 维 `observation.state` 协议闭环评估。\n"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## Pi0 正式保护协议评估\n\n本评估采用与课程权重一致的闭环协议：\n\n- `eef_abs`：Pi0 输出 TCP 目标点，环境执行受限 TCP 增量；\n- 每次观测采样 4 个 action chunks，使用 `median` 保留一致轨迹；\n- 每次执行 5 步，再从动作队列继续执行；\n- 状态为 `6DoF + gripper + timestamp` 共 8 维；\n- 14 个固定 seed 输出逐回合结果和视频。\n",
   "execution_count": null
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# PI0_NATIVE_FORMAL_EVAL_CELL\n",
    "import gc\n",
    "import random\n",
    "import sys\n",
    "import time\n",
    "from collections import deque\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "from PIL import Image\n",
    "from torchvision import transforms\n",
    "\n",
    "_PI0_NATIVE_TO_TENSOR = transforms.ToTensor()\n",
    "\n",
    "def _pi0_native_to_tensor_image(image):\n",
    "    return _PI0_NATIVE_TO_TENSOR(Image.fromarray(image).resize((256, 256)))\n",
    "\n",
    "\n",
    "def _pi0_native_make_policy(policy_path, repo_id, dataset_root, device):\n",
    "    from lerobot.common.datasets.lerobot_dataset import LeRobotDatasetMetadata\n",
    "    from lerobot.common.policies.pi0.modeling_pi0 import PI0Policy\n",
    "\n",
    "    metadata = LeRobotDatasetMetadata(repo_id, root=dataset_root)\n",
    "    policy = PI0Policy.from_pretrained(policy_path, dataset_stats=metadata.stats)\n",
    "    policy.to(device)\n",
    "    policy.eval()\n",
    "    return policy\n",
    "\n",
    "\n",
    "def _pi0_native_expected_state_dim(policy):\n",
    "    feature = policy.config.input_features[\"observation.state\"]\n",
    "    return int(tuple(feature.shape)[0])\n",
    "\n",
    "\n",
    "def _pi0_native_gripper(env):\n",
    "    try:\n",
    "        return float(np.clip(float(env.env.get_qpos_joint(\"rh_r1\")[0]), 0.0, 1.0))\n",
    "    except Exception:\n",
    "        return 0.0\n",
    "\n",
    "\n",
    "def _pi0_native_state(env, policy, action_step, hz=20.0):\n",
    "    base = np.asarray(env.get_joint_state()[:6], dtype=np.float32).reshape(-1)\n",
    "    expected = _pi0_native_expected_state_dim(policy)\n",
    "    if expected == 6:\n",
    "        return base\n",
    "    if expected == 8:\n",
    "        # Match the protected evaluator exactly: timestamp is the action\n",
    "        # schedule time, not MuJoCo wall/simulation time after the 25 physics steps.\n",
    "        timestamp = float(action_step) / float(hz)\n",
    "        return np.concatenate(\n",
    "            [base, np.asarray([_pi0_native_gripper(env), timestamp], dtype=np.float32)]\n",
    "        )\n",
    "    raise ValueError(\n",
    "        f\"Pi0 checkpoint expects observation.state={expected}D; \"\n",
    "        \"formal raw VLA protocol only supplies 6D or 8D state.\"\n",
    "    )\n",
    "\n",
    "\n",
    "def _pi0_native_debug(env, tracker):\n",
    "    target_body = env.obj_target\n",
    "    target_pos = np.asarray(env.env.get_p_body(target_body), dtype=np.float32)\n",
    "    plate_pos = np.asarray(env.env.get_p_body(\"body_obj_plate_11\"), dtype=np.float32)\n",
    "    tcp_pos = np.asarray(env.env.get_p_body(\"tcp_link\"), dtype=np.float32)\n",
    "    try:\n",
    "        upright_cos = float(env.env.get_R_body(target_body)[2, 2])\n",
    "    except Exception:\n",
    "        upright_cos = float(\"nan\")\n",
    "    terminal = bool(env.check_success())\n",
    "    lifted = int(tracker[\"lifted_steps\"]) >= 3\n",
    "    upright = bool(np.isfinite(upright_cos) and upright_cos >= 0.7)\n",
    "    return {\n",
    "        \"xy_dist\": float(np.linalg.norm(target_pos[:2] - plate_pos[:2])),\n",
    "        \"target_pos\": [round(float(x), 4) for x in target_pos],\n",
    "        \"plate_pos\": [round(float(x), 4) for x in plate_pos],\n",
    "        \"tcp_pos\": [round(float(x), 4) for x in tcp_pos],\n",
    "        \"gripper\": _pi0_native_gripper(env),\n",
    "        \"success\": terminal,\n",
    "        \"physical_success\": bool(terminal and lifted and upright),\n",
    "        \"physical_lifted_enough\": bool(lifted),\n",
    "        \"physical_final_upright\": bool(upright),\n",
    "        \"physical_min_lift\": 0.03,\n",
    "        \"physical_min_lift_steps\": 3,\n",
    "        \"physical_final_upright_cos_threshold\": 0.7,\n",
    "        \"target_body\": target_body,\n",
    "        \"initial_target_pos\": tracker[\"initial_target_pos\"],\n",
    "        \"initial_plate_pos\": tracker[\"initial_plate_pos\"],\n",
    "        \"final_target_pos\": [round(float(x), 4) for x in target_pos],\n",
    "        \"max_target_lift\": float(tracker[\"max_target_lift\"]),\n",
    "        \"lifted_steps\": int(tracker[\"lifted_steps\"]),\n",
    "        \"final_target_upright_cos\": upright_cos,\n",
    "    }\n",
    "\n",
    "\n",
    "@torch.inference_mode()\n",
    "def _pi0_native_sample_chunk(policy, batch, samples=4, exec_chunk_steps=5):\n",
    "    # This follows the same Pi0 sampling path as the protected 12/14 run,\n",
    "    # while keeping the implementation visible in the Notebook.\n",
    "    from lerobot.common.constants import OBS_STATE\n",
    "\n",
    "    normalized = policy.normalize_inputs(batch)\n",
    "    images, image_masks = policy.prepare_images(normalized)\n",
    "    state = policy.prepare_state(normalized)\n",
    "    language_tokens, language_masks = policy.prepare_language(normalized)\n",
    "    chunks = []\n",
    "    for _ in range(int(samples)):\n",
    "        actions = policy.model.sample_actions(\n",
    "            images, image_masks, language_tokens, language_masks, state, noise=None\n",
    "        )\n",
    "        action_dim = policy.config.action_feature.shape[0]\n",
    "        actions = actions[:, :, :action_dim]\n",
    "        actions = policy.unnormalize_outputs({\"action\": actions})[\"action\"]\n",
    "        if getattr(policy.config, \"adapt_to_pi_aloha\", False):\n",
    "            actions = policy._pi_aloha_encode_actions(actions)\n",
    "        chunks.append(actions)\n",
    "    selected = torch.stack(chunks, dim=0).median(dim=0).values\n",
    "    # Keep a coherent short suffix from one predicted chunk before replanning.\n",
    "    return deque(selected[0, : int(exec_chunk_steps)].unbind(0))\n",
    "\n",
    "\n",
    "def _pi0_native_configure_env(env):\n",
    "    # The historical protected evaluator sets this reference pose once per episode.\n",
    "    # Recomputing env.p0 on every action changes the MuJoCo eef_pose controller.\n",
    "    from mujoco_env.transforms import rpy2r\n",
    "\n",
    "    env.action_type = \"eef_pose\"\n",
    "    env.p0, _ = env.env.get_pR_body(body_name=\"tcp_link\")\n",
    "    env.R0 = rpy2r(np.deg2rad([90.0, 0.0, 90.0]))\n",
    "\n",
    "\n",
    "def _pi0_native_eef_abs_to_env(action, env, max_step=0.004):\n",
    "    current_tcp = np.asarray(env.env.get_p_body(\"tcp_link\")[:3], dtype=np.float32)\n",
    "    target_tcp = np.asarray(action[:3], dtype=np.float32)\n",
    "    env_action = np.zeros(7, dtype=np.float32)\n",
    "    env_action[:3] = np.clip(target_tcp - current_tcp, -float(max_step), float(max_step))\n",
    "    env_action[6] = float(np.clip(action[6], 0.0, 1.0))\n",
    "    return env_action\n",
    "\n",
    "\n",
    "def _pi0_native_close_env(env):\n",
    "    try:\n",
    "        env.env.close_viewer()\n",
    "    except Exception:\n",
    "        pass\n",
    "    gc.collect()\n",
    "\n",
    "\n",
    "def run_pi0_eval_native_in_notebook(\n",
    "    policy_path,\n",
    "    result_path,\n",
    "    seeds,\n",
    "    repo_id,\n",
    "    dataset_root,\n",
    "    enabled=False,\n",
    "    render=False,\n",
    "):\n",
    "    print(\"Pi0 formal checkpoint =\", public_path(policy_path))\n",
    "    print(\"Pi0 formal seeds =\", list(seeds))\n",
    "    print(\"Pi0 protocol = eef_abs / samples=4 / reducer=median / execute=5 / state=8D / physics=25\")\n",
    "    if not enabled:\n",
    "        print(\"未启动。设置 RUN_EVAL=1 后运行 Pi0 正式保护协议。\")\n",
    "        return None\n",
    "    if not ensure_project_layout():\n",
    "        return None\n",
    "\n",
    "    from contextlib import contextmanager\n",
    "\n",
    "    project_path = str(Path(PROJECT_ROOT).resolve())\n",
    "    if project_path not in sys.path:\n",
    "        sys.path.insert(0, project_path)\n",
    "    from mujoco_env.y_env2 import SimpleEnv2\n",
    "\n",
    "    @contextmanager\n",
    "    def pushd(path):\n",
    "        old = Path.cwd()\n",
    "        os.chdir(path)\n",
    "        try:\n",
    "            yield\n",
    "        finally:\n",
    "            os.chdir(old)\n",
    "\n",
    "    ensure_xvfb_display()\n",
    "    result_path = Path(result_path)\n",
    "    result_path.parent.mkdir(parents=True, exist_ok=True)\n",
    "    existing_rows = []\n",
    "    completed_seeds = set()\n",
    "    if result_path.exists():\n",
    "        # Resume after an interrupted long Notebook run without duplicating seeds.\n",
    "        for line in result_path.read_text(encoding=\"utf-8\").splitlines():\n",
    "            line = line.strip()\n",
    "            if not line:\n",
    "                continue\n",
    "            row = json.loads(line)\n",
    "            existing_rows.append(row)\n",
    "            if row.get(\"seed\") is not None:\n",
    "                completed_seeds.add(int(row[\"seed\"]))\n",
    "        print(\"Pi0 resume rows =\", len(existing_rows), \"completed seeds =\", sorted(completed_seeds))\n",
    "    run_seeds = [int(seed) for seed in seeds if int(seed) not in completed_seeds]\n",
    "    if not run_seeds:\n",
    "        print(\"Pi0 formal result already complete; no seed will be repeated.\")\n",
    "\n",
    "    device = os.environ.get(\"EVAL_DEVICE\", \"cuda\")\n",
    "    with pushd(PROJECT_ROOT):\n",
    "        policy = _pi0_native_make_policy(Path(policy_path), repo_id, Path(dataset_root), device)\n",
    "        expected_dim = _pi0_native_expected_state_dim(policy)\n",
    "        print(\"Pi0 observation.state dimension =\", expected_dim)\n",
    "        if expected_dim not in (6, 8):\n",
    "            raise ValueError(f\"Unsupported Pi0 state dimension: {expected_dim}\")\n",
    "\n",
    "        rows = list(existing_rows)\n",
    "        for seed in notebook_progress(run_seeds, desc=\"pi0 formal eval\", total=len(run_seeds)):\n",
    "            env = SimpleEnv2(\"./asset/example_scene_y2.xml\", action_type=\"joint_angle\")\n",
    "            try:\n",
    "                try:\n",
    "                    env.env.reset(step=False)\n",
    "                except TypeError:\n",
    "                    env.env.reset()\n",
    "                env.reset(seed=int(seed))\n",
    "                random.seed(int(seed))\n",
    "                np.random.seed(int(seed))\n",
    "                torch.manual_seed(int(seed))\n",
    "                if torch.cuda.is_available():\n",
    "                    torch.cuda.manual_seed_all(int(seed))\n",
    "                _pi0_native_configure_env(env)\n",
    "                policy.reset()\n",
    "                tracker = {\n",
    "                    \"initial_target_pos\": [float(x) for x in env.env.get_p_body(env.obj_target)],\n",
    "                    \"initial_plate_pos\": [float(x) for x in env.env.get_p_body(\"body_obj_plate_11\")],\n",
    "                    \"initial_target_z\": float(env.env.get_p_body(env.obj_target)[2]),\n",
    "                    \"max_target_lift\": 0.0,\n",
    "                    \"lifted_steps\": 0,\n",
    "                }\n",
    "                action_queue = deque()\n",
    "                action_steps = 0\n",
    "                sim_steps = 0\n",
    "                success_ever = False\n",
    "                physical_success_ever = False\n",
    "                first_success_step = None\n",
    "                first_physical_success_step = None\n",
    "                start = time.time()\n",
    "                while action_steps < int(os.environ.get(\"PI0_EVAL_MAX_ACTION_STEPS\", \"700\")) and env.env.is_viewer_alive():\n",
    "                    for _ in range(25):\n",
    "                        env.step_env()\n",
    "                        sim_steps += 1\n",
    "                        target_z = float(env.env.get_p_body(env.obj_target)[2])\n",
    "                        lift = target_z - tracker[\"initial_target_z\"]\n",
    "                        tracker[\"max_target_lift\"] = max(tracker[\"max_target_lift\"], lift)\n",
    "                        if lift >= 0.03:\n",
    "                            tracker[\"lifted_steps\"] += 1\n",
    "                    debug = _pi0_native_debug(env, tracker)\n",
    "                    if debug[\"success\"]:\n",
    "                        success_ever = True\n",
    "                        first_success_step = action_steps if first_success_step is None else first_success_step\n",
    "                    if debug[\"physical_success\"]:\n",
    "                        physical_success_ever = True\n",
    "                        first_physical_success_step = action_steps if first_physical_success_step is None else first_physical_success_step\n",
    "                        break\n",
    "\n",
    "                    state = _pi0_native_state(env, policy, action_steps)\n",
    "                    image, wrist_image = env.grab_image()\n",
    "                    batch = {\n",
    "                        \"observation.state\": torch.tensor([state], dtype=torch.float32, device=device),\n",
    "                        \"observation.image\": _pi0_native_to_tensor_image(image).unsqueeze(0).to(device),\n",
    "                        \"observation.wrist_image\": _pi0_native_to_tensor_image(wrist_image).unsqueeze(0).to(device),\n",
    "                        \"task\": [env.instruction],\n",
    "                    }\n",
    "                    if not action_queue:\n",
    "                        action_queue = _pi0_native_sample_chunk(policy, batch, samples=4, exec_chunk_steps=5)\n",
    "                    action = action_queue.popleft().detach().float().cpu().numpy()\n",
    "                    env.step(_pi0_native_eef_abs_to_env(action, env))\n",
    "                    action_steps += 1\n",
    "                    if render:\n",
    "                        env.render()\n",
    "\n",
    "                    debug = _pi0_native_debug(env, tracker)\n",
    "                    if debug[\"success\"]:\n",
    "                        success_ever = True\n",
    "                        first_success_step = action_steps if first_success_step is None else first_success_step\n",
    "                    if debug[\"physical_success\"]:\n",
    "                        physical_success_ever = True\n",
    "                        first_physical_success_step = action_steps if first_physical_success_step is None else first_physical_success_step\n",
    "                        break\n",
    "\n",
    "                debug = _pi0_native_debug(env, tracker)\n",
    "                row = {\n",
    "                    \"policy\": \"pi0\",\n",
    "                    \"seed\": int(seed),\n",
    "                    \"success\": bool(success_ever or debug[\"success\"]),\n",
    "                    \"physical_success\": bool(physical_success_ever or debug[\"physical_success\"]),\n",
    "                    \"action_steps\": int(action_steps),\n",
    "                    \"sim_steps\": int(sim_steps),\n",
    "                    \"elapsed_s\": round(time.time() - start, 3),\n",
    "                    \"instruction\": getattr(env, \"instruction\", None),\n",
    "                    \"debug\": debug,\n",
    "                }\n",
    "                rows.append(row)\n",
    "                with result_path.open(\"a\", encoding=\"utf-8\") as f:\n",
    "                    f.write(json.dumps(row, ensure_ascii=False) + \"\\n\")\n",
    "                print(json.dumps(row, ensure_ascii=False))\n",
    "            finally:\n",
    "                _pi0_native_close_env(env)\n",
    "                if torch.cuda.is_available():\n",
    "                    torch.cuda.empty_cache()\n",
    "\n",
    "    summarize_jsonl(result_path)\n",
    "    success_count = sum(bool(row[\"success\"]) for row in rows)\n",
    "    physical_count = sum(bool(row[\"physical_success\"]) for row in rows)\n",
    "    print(f\"Pi0 formal result: success={success_count}/{len(rows)}, physical_success={physical_count}/{len(rows)}\")\n",
    "    return rows\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": "评估使用环境变量指定权重： $MODEL_ROOT/pi0_clean40_successonly_blue2x_from_s7500_s2000_v1/checkpoints/001000/pretrained_model\nPi0 formal checkpoint = $MODEL_ROOT/pi0_clean40_successonly_blue2x_from_s7500_s2000_v1/checkpoints/001000/pretrained_model\nPi0 formal seeds = [3007, 3002, 3004, 3006, 3010, 3011, 3012, 3013, 3014, 3015, 3016, 3017, 3018, 3019]\nPi0 protocol = eef_abs / samples=4 / reducer=median / execute=5 / state=8D / physics=25\nDISPLAY = :0\nPi0 resume rows = 11 completed seeds = [3002, 3004, 3006, 3007, 3010, 3011, 3012, 3013, 3014, 3015, 3016]\nLoading weights from local directory\nPi0 observation.state dimension = 8\n\n-----------------------------------------------------------------------------\nname:[Tabletop] dt:[0.002] HZ:[500]\n n_qpos:[31] n_qvel:[28] n_qacc:[28] n_ctrl:[10]\n integrator:[IMPLICITFAST]\n\nn_body:[23]\n [0/23] [world] mass:[0.00]kg\n [1/23] [front_object_table] mass:[1.00]kg\n [2/23] [camera] mass:[0.00]kg\n [3/23] [camera2] mass:[0.00]kg\n [4/23] [camera3] mass:[0.00]kg\n [5/23] [link1] mass:[2.06]kg\n [6/23] [link2] mass:[3.68]kg\n [7/23] [link3] mass:[2.39]kg\n [8/23] [link4] mass:[1.40]kg\n [9/23] [link5] mass:[1.40]kg\n [10/23] [link6] mass:[0.65]kg\n [11/23] [camera_center] mass:[0.00]kg\n [12/23] [tcp_link] mass:[0.32]kg\n [13/23] [rh_p12_rn_r1] mass:[0.07]kg\n [14/23] [rh_p12_rn_r2] mass:[0.02]kg\n [15/23] [rh_p12_rn_l1] mass:[0.07]kg\n [16/23] [rh_p12_rn_l2] mass:[0.02]kg\n [17/23] [body_obj_mug_5] mass:[0.00]kg\n [18/23] [object_mug_5] mass:[0.08]kg\n [19/23] [body_obj_plate_11] mass:[0.00]kg\n [20/23] [object_plate_11] mass:[0.10]kg\n [21/23] [body_obj_mug_6] mass:[0.00]kg\n [22/23] [object_mug_6] mass:[0.08]kg\nbody_total_mass:[13.35]kg\n\nn_geom:[116]\ngeom_names:['floor', None, 'front_object_table', None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None]\n\nn_mesh:[112]\nmesh_names:['base_unit', 'link1', 'link2', 'link3', 'link4', 'link5', 'link6', 'flange', 'base', 'r1', 'r2', 'l1', 'l2', 'mug_5_normalized_0_vis', 'mug_5_normalized_collision_22._coll', 'mug_5_normalized_collision_23._coll', 'mug_5_normalized_collision_21._coll', 'mug_5_normalized_collision_20._coll', 'mug_5_normalized_collision_24._coll', 'mug_5_normalized_collision_30._coll', 'mug_5_normalized_collision_18._coll', 'mug_5_normalized_collision_19._coll', 'mug_5_normalized_collision_31._coll', 'mug_5_normalized_collision_25._coll', 'mug_5_normalized_collision_27._coll', 'mug_5_normalized_collision_26._coll', 'mug_5_normalized_collision_9._coll', 'mug_5_normalized_collision_8._coll', 'mug_5_normalized_collision_6._coll', 'mug_5_normalized_collision_7._coll', 'mug_5_normalized_collision_5._coll', 'mug_5_normalized_collision_4._coll', 'mug_5_normalized_collision_0._coll', 'mug_5_normalized_collision_1._coll', 'mug_5_normalized_collision_3._coll', 'mug_5_normalized_collision_2._coll', 'mug_5_normalized_collision_17._coll', 'mug_5_normalized_collision_16._coll', 'mug_5_normalized_collision_28._coll', 'mug_5_normalized_collision_14._coll', 'mug_5_normalized_collision_15._coll', 'mug_5_normalized_collision_29._coll', 'mug_5_normalized_collision_11._coll', 'mug_5_normalized_collision_10._coll', 'mug_5_normalized_collision_12._coll', 'mug_5_normalized_collision_13._coll', 'plate_11_normalized_0_vis', 'plate_11_normalized_collision_22._coll', 'plate_11_normalized_collision_23._coll', 'plate_11_normalized_collision_21._coll', 'plate_11_normalized_collision_20._coll', 'plate_11_normalized_collision_24._coll', 'plate_11_normalized_collision_30._coll', 'plate_11_normalized_collision_18._coll', 'plate_11_normalized_collision_19._coll', 'plate_11_normalized_collision_31._coll', 'plate_11_normalized_collision_25._coll', 'plate_11_normalized_collision_27._coll', 'plate_11_normalized_collision_26._coll', 'plate_11_normalized_collision_9._coll', 'plate_11_normalized_collision_8._coll', 'plate_11_normalized_collision_6._coll', 'plate_11_normalized_collision_7._coll', 'plate_11_normalized_collision_5._coll', 'plate_11_normalized_collision_4._coll', 'plate_11_normalized_collision_0._coll', 'plate_11_normalized_collision_1._coll', 'plate_11_normalized_collision_3._coll', 'plate_11_normalized_collision_2._coll', 'plate_11_normalized_collision_17._coll', 'plate_11_normalized_collision_16._coll', 'plate_11_normalized_collision_28._coll', 'plate_11_normalized_collision_14._coll', 'plate_11_normalized_collision_15._coll', 'plate_11_normalized_collision_29._coll', 'plate_11_normalized_collision_11._coll', 'plate_11_normalized_collision_10._coll', 'plate_11_normalized_collision_12._coll', 'plate_11_normalized_collision_13._coll', 'mug_6_normalized_0_vis', 'mug_6_normalized_collision_22._coll', 'mug_6_normalized_collision_23._coll', 'mug_6_normalized_collision_21._coll', 'mug_6_normalized_collision_20._coll', 'mug_6_normalized_collision_24._coll', 'mug_6_normalized_collision_30._coll', 'mug_6_normalized_collision_18._coll', 'mug_6_normalized_collision_19._coll', 'mug_6_normalized_collision_31._coll', 'mug_6_normalized_collision_25._coll', 'mug_6_normalized_collision_27._coll', 'mug_6_normalized_collision_26._coll', 'mug_6_normalized_collision_9._coll', 'mug_6_normalized_collision_8._coll', 'mug_6_normalized_collision_6._coll', 'mug_6_normalized_collision_7._coll', 'mug_6_normalized_collision_5._coll', 'mug_6_normalized_collision_4._coll', 'mug_6_normalized_collision_0._coll', 'mug_6_normalized_collision_1._coll', 'mug_6_normalized_collision_3._coll', 'mug_6_normalized_collision_2._coll', 'mug_6_normalized_collision_17._coll', 'mug_6_normalized_collision_16._coll', 'mug_6_normalized_collision_28._coll', 'mug_6_normalized_collision_14._coll', 'mug_6_normalized_collision_15._coll', 'mug_6_normalized_collision_29._coll', 'mug_6_normalized_collision_11._coll', 'mug_6_normalized_collision_10._coll', 'mug_6_normalized_collision_12._coll', 'mug_6_normalized_collision_13._coll']\n\nn_joint:[13]\n [0/13] [joint1] axis:[0. 0. 1.]\n [1/13] [joint2] axis:[0. 1. 0.]\n [2/13] [joint3] axis:[0. 1. 0.]\n [3/13] [joint4] axis:[0. 1. 0.]\n [4/13] [joint5] axis:[0. 0. 1.]\n [5/13] [joint6] axis:[0. 1. 0.]\n [6/13] [rh_r1] axis:[1. 0. 0.]\n [7/13] [rh_r2] axis:[-1.  0.  0.]\n [8/13] [rh_l1] axis:[-1.  0.  0.]\n [9/13] [rh_l2] axis:[1. 0. 0.]\n [10/13] [None] axis:[0. 0. 1.]\n [11/13] [None] axis:[0. 0. 1.]\n [12/13] [None] axis:[0. 0. 1.]\n\nn_dof:[28] (=number of rows of Jacobian)\n [0/28] [None] attached joint:[joint1] body:[link1]\n [1/28] [None] attached joint:[joint2] body:[link2]\n [2/28] [None] attached joint:[joint3] body:[link3]\n [3/28] [None] attached joint:[joint4] body:[link4]\n [4/28] [None] attached joint:[joint5] body:[link5]\n [5/28] [None] attached joint:[joint6] body:[link6]\n [6/28] [None] attached joint:[rh_r1] body:[rh_p12_rn_r1]\n [7/28] [None] attached joint:[rh_r2] body:[rh_p12_rn_r2]\n [8/28] [None] attached joint:[rh_l1] body:[rh_p12_rn_l1]\n [9/28] [None] attached joint:[rh_l2] body:[rh_p12_rn_l2]\n [10/28] [None] attached joint:[None] body:[body_obj_mug_5]\n [11/28] [None] attached joint:[None] body:[body_obj_mug_5]\n [12/28] [None] attached joint:[None] body:[body_obj_mug_5]\n [13/28] [None] attached joint:[None] body:[body_obj_mug_5]\n [14/28] [None] attached joint:[None] body:[body_obj_mug_5]\n [15/28] [None] attached joint:[None] body:[body_obj_mug_5]\n [16/28] [None] attached joint:[None] body:[body_obj_plate_11]\n [17/28] [None] attached joint:[None] body:[body_obj_plate_11]\n [18/28] [None] attached joint:[None] body:[body_obj_plate_11]\n [19/28] [None] attached joint:[None] body:[body_obj_plate_11]\n [20/28] [None] attached joint:[None] body:[body_obj_plate_11]\n [21/28] [None] attached joint:[None] body:[body_obj_plate_11]\n [22/28] [None] attached joint:[None] body:[body_obj_mug_6]\n [23/28] [None] attached joint:[None] body:[body_obj_mug_6]\n [24/28] [None] attached joint:[None] body:[body_obj_mug_6]\n [25/28] [None] attached joint:[None] body:[body_obj_mug_6]\n [26/28] [None] attached joint:[None] body:[body_obj_mug_6]\n [27/28] [None] attached joint:[None] body:[body_obj_mug_6]\n\nFree joint information. n_free_joint:[3]\n [0/3] [None] body_name_attached:[body_obj_mug_5]\n [1/3] [None] body_name_attached:[body_obj_plate_11]\n [2/3] [None] body_name_attached:[body_obj_mug_6]\n\nRevolute joint information. n_rev_joint:[10]\n [0/10] [joint1] range:[-6.283]~[6.283]\n [1/10] [joint2] range:[-6.283]~[6.283]\n [2/10] [joint3] range:[-6.283]~[6.283]\n [3/10] [joint4] range:[-6.283]~[6.283]\n [4/10] [joint5] range:[-6.283]~[6.283]\n [5/10] [joint6] range:[-6.283]~[6.283]\n [6/10] [rh_r1] range:[0.000]~[1.100]\n [7/10] [rh_r2] range:[0.000]~[1.000]\n [8/10] [rh_l1] range:[0.000]~[1.100]\n [9/10] [rh_l2] range:[0.000]~[1.000]\n\nPrismatic joint information. n_pri_joint:[0]\n\nControl information. n_ctrl:[10]\n [0/10] [actuator_joint1] range:[-6.283]~[6.283] gear:[1.00] type:[JOINT]\n [1/10] [actuator_joint2] range:[-6.283]~[6.283] gear:[1.00] type:[JOINT]\n [2/10] [actuator_joint3] range:[-6.283]~[6.283] gear:[1.00] type:[JOINT]\n [3/10] [actuator_joint4] range:[-6.283]~[6.283] gear:[1.00] type:[JOINT]\n [4/10] [actuator_joint5] range:[-6.283]~[6.283] gear:[1.00] type:[JOINT]\n [5/10] [actuator_joint6] range:[-6.283]~[6.283] gear:[1.00] type:[JOINT]\n [6/10] [actuator_rh_r1] range:[0.000]~[1.100] gear:[1.00] type:[JOINT]\n [7/10] [actuator_rh_r2] range:[0.000]~[1.000] gear:[1.00] type:[JOINT]\n [8/10] [actuator_rh_l1] range:[0.000]~[1.100] gear:[1.00] type:[JOINT]\n [9/10] [actuator_rh_l2] range:[0.000]~[1.000] gear:[1.00] type:[JOINT]\n\nCamera information. n_cam:[4]\n [0/4] [agentview] fov:[60.0]\n [1/4] [topview] fov:[90.0]\n [2/4] [sideview] fov:[90.0]\n [3/4] [egocentric] fov:[90.0]\n\nn_sensor:[0]\nsensor_names:[]\nn_site:[9]\nsite_names:['bottom_site_mug_5', 'top_site_mug_5', 'horizontal_radius_site_mug_5', 'bottom_site_plate_11', 'top_site_plate_11', 'horizontal_radius_site_plate_11', 'bottom_site_mug_6', 'top_site_mug_6', 'horizontal_radius_site_mug_6']\n-----------------------------------------------------------------------------\nenv:[Tabletop] reset\nenv:[Tabletop] reset\nenv:[Tabletop] initalize viewer\nDONE INITIALIZATION\nenv:[Tabletop] reset\nDONE INITIALIZATION\n{\"policy\": \"pi0\", \"seed\": 3017, \"success\": true, \"physical_success\": true, \"action_steps\": 568, \"sim_steps\": 14225, \"elapsed_s\": 223.248, \"instruction\": \"Place the red mug on the plate.\", \"debug\": {\"xy_dist\": 0.032662779092788696, \"target_pos\": [0.314, -0.2797, 0.8527], \"plate_pos\": [0.3001, -0.2501, 0.8081], \"tcp_pos\": [0.3072, -0.2305, 0.903], \"gripper\": 0.06617078288012511, \"success\": true, \"physical_success\": true, \"physical_lifted_enough\": true, \"physical_final_upright\": true, \"physical_min_lift\": 0.03, \"physical_min_lift_steps\": 3, \"physical_final_upright_cos_threshold\": 0.7, \"target_body\": \"body_obj_mug_5\", \"initial_target_pos\": [0.3233920671807391, 0.015395458511440578, 0.84465294940151], \"initial_plate_pos\": [0.2999889147642407, -0.24998448478945673, 0.8078853128532457], \"final_target_pos\": [0.314, -0.2797, 0.8527], \"max_target_lift\": 0.10417330161349636, \"lifted_steps\": 5636, \"final_target_upright_cos\": 0.9999837934798487}}\n\n-----------------------------------------------------------------------------\nname:[Tabletop] dt:[0.002] HZ:[500]\n n_qpos:[31] n_qvel:[28] n_qacc:[28] n_ctrl:[10]\n integrator:[IMPLICITFAST]\n\nn_body:[23]\n [0/23] [world] mass:[0.00]kg\n [1/23] [front_object_table] mass:[1.00]kg\n [2/23] [camera] mass:[0.00]kg\n [3/23] [camera2] mass:[0.00]kg\n [4/23] [camera3] mass:[0.00]kg\n [5/23] [link1] mass:[2.06]kg\n [6/23] [link2] mass:[3.68]kg\n [7/23] [link3] mass:[2.39]kg\n [8/23] [link4] mass:[1.40]kg\n [9/23] [link5] mass:[1.40]kg\n [10/23] [link6] mass:[0.65]kg\n [11/23] [camera_center] mass:[0.00]kg\n [12/23] [tcp_link] mass:[0.32]kg\n [13/23] [rh_p12_rn_r1] mass:[0.07]kg\n [14/23] [rh_p12_rn_r2] mass:[0.02]kg\n [15/23] [rh_p12_rn_l1] mass:[0.07]kg\n [16/23] [rh_p12_rn_l2] mass:[0.02]kg\n [17/23] [body_obj_mug_5] mass:[0.00]kg\n [18/23] [object_mug_5] mass:[0.08]kg\n [19/23] [body_obj_plate_11] mass:[0.00]kg\n [20/23] [object_plate_11] mass:[0.10]kg\n [21/23] [body_obj_mug_6] mass:[0.00]kg\n [22/23] [object_mug_6] mass:[0.08]kg\nbody_total_mass:[13.35]kg\n\nn_geom:[116]\ngeom_names:['floor', None, 'front_object_table', None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None]\n\nn_mesh:[112]\nmesh_names:['base_unit', 'link1', 'link2', 'link3', 'link4', 'link5', 'link6', 'flange', 'base', 'r1', 'r2', 'l1', 'l2', 'mug_5_normalized_0_vis', 'mug_5_normalized_collision_22._coll', 'mug_5_normalized_collision_23._coll', 'mug_5_normalized_collision_21._coll', 'mug_5_normalized_collision_20._coll', 'mug_5_normalized_collision_24._coll', 'mug_5_normalized_collision_30._coll', 'mug_5_normalized_collision_18._coll', 'mug_5_normalized_collision_19._coll', 'mug_5_normalized_collision_31._coll', 'mug_5_normalized_collision_25._coll', 'mug_5_normalized_collision_27._coll', 'mug_5_normalized_collision_26._coll', 'mug_5_normalized_collision_9._coll', 'mug_5_normalized_collision_8._coll', 'mug_5_normalized_collision_6._coll', 'mug_5_normalized_collision_7._coll', 'mug_5_normalized_collision_5._coll', 'mug_5_normalized_collision_4._coll', 'mug_5_normalized_collision_0._coll', 'mug_5_normalized_collision_1._coll', 'mug_5_normalized_collision_3._coll', 'mug_5_normalized_collision_2._coll', 'mug_5_normalized_collision_17._coll', 'mug_5_normalized_collision_16._coll', 'mug_5_normalized_collision_28._coll', 'mug_5_normalized_collision_14._coll', 'mug_5_normalized_collision_15._coll', 'mug_5_normalized_collision_29._coll', 'mug_5_normalized_collision_11._coll', 'mug_5_normalized_collision_10._coll', 'mug_5_normalized_collision_12._coll', 'mug_5_normalized_collision_13._coll', 'plate_11_normalized_0_vis', 'plate_11_normalized_collision_22._coll', 'plate_11_normalized_collision_23._coll', 'plate_11_normalized_collision_21._coll', 'plate_11_normalized_collision_20._coll', 'plate_11_normalized_collision_24._coll', 'plate_11_normalized_collision_30._coll', 'plate_11_normalized_collision_18._coll', 'plate_11_normalized_collision_19._coll', 'plate_11_normalized_collision_31._coll', 'plate_11_normalized_collision_25._coll', 'plate_11_normalized_collision_27._coll', 'plate_11_normalized_collision_26._coll', 'plate_11_normalized_collision_9._coll', 'plate_11_normalized_collision_8._coll', 'plate_11_normalized_collision_6._coll', 'plate_11_normalized_collision_7._coll', 'plate_11_normalized_collision_5._coll', 'plate_11_normalized_collision_4._coll', 'plate_11_normalized_collision_0._coll', 'plate_11_normalized_collision_1._coll', 'plate_11_normalized_collision_3._coll', 'plate_11_normalized_collision_2._coll', 'plate_11_normalized_collision_17._coll', 'plate_11_normalized_collision_16._coll', 'plate_11_normalized_collision_28._coll', 'plate_11_normalized_collision_14._coll', 'plate_11_normalized_collision_15._coll', 'plate_11_normalized_collision_29._coll', 'plate_11_normalized_collision_11._coll', 'plate_11_normalized_collision_10._coll', 'plate_11_normalized_collision_12._coll', 'plate_11_normalized_collision_13._coll', 'mug_6_normalized_0_vis', 'mug_6_normalized_collision_22._coll', 'mug_6_normalized_collision_23._coll', 'mug_6_normalized_collision_21._coll', 'mug_6_normalized_collision_20._coll', 'mug_6_normalized_collision_24._coll', 'mug_6_normalized_collision_30._coll', 'mug_6_normalized_collision_18._coll', 'mug_6_normalized_collision_19._coll', 'mug_6_normalized_collision_31._coll', 'mug_6_normalized_collision_25._coll', 'mug_6_normalized_collision_27._coll', 'mug_6_normalized_collision_26._coll', 'mug_6_normalized_collision_9._coll', 'mug_6_normalized_collision_8._coll', 'mug_6_normalized_collision_6._coll', 'mug_6_normalized_collision_7._coll', 'mug_6_normalized_collision_5._coll', 'mug_6_normalized_collision_4._coll', 'mug_6_normalized_collision_0._coll', 'mug_6_normalized_collision_1._coll', 'mug_6_normalized_collision_3._coll', 'mug_6_normalized_collision_2._coll', 'mug_6_normalized_collision_17._coll', 'mug_6_normalized_collision_16._coll', 'mug_6_normalized_collision_28._coll', 'mug_6_normalized_collision_14._coll', 'mug_6_normalized_collision_15._coll', 'mug_6_normalized_collision_29._coll', 'mug_6_normalized_collision_11._coll', 'mug_6_normalized_collision_10._coll', 'mug_6_normalized_collision_12._coll', 'mug_6_normalized_collision_13._coll']\n\nn_joint:[13]\n [0/13] [joint1] axis:[0. 0. 1.]\n [1/13] [joint2] axis:[0. 1. 0.]\n [2/13] [joint3] axis:[0. 1. 0.]\n [3/13] [joint4] axis:[0. 1. 0.]\n [4/13] [joint5] axis:[0. 0. 1.]\n [5/13] [joint6] axis:[0. 1. 0.]\n [6/13] [rh_r1] axis:[1. 0. 0.]\n [7/13] [rh_r2] axis:[-1.  0.  0.]\n [8/13] [rh_l1] axis:[-1.  0.  0.]\n [9/13] [rh_l2] axis:[1. 0. 0.]\n [10/13] [None] axis:[0. 0. 1.]\n [11/13] [None] axis:[0. 0. 1.]\n [12/13] [None] axis:[0. 0. 1.]\n\nn_dof:[28] (=number of rows of Jacobian)\n [0/28] [None] attached joint:[joint1] body:[link1]\n [1/28] [None] attached joint:[joint2] body:[link2]\n [2/28] [None] attached joint:[joint3] body:[link3]\n [3/28] [None] attached joint:[joint4] body:[link4]\n [4/28] [None] attached joint:[joint5] body:[link5]\n [5/28] [None] attached joint:[joint6] body:[link6]\n [6/28] [None] attached joint:[rh_r1] body:[rh_p12_rn_r1]\n [7/28] [None] attached joint:[rh_r2] body:[rh_p12_rn_r2]\n [8/28] [None] attached joint:[rh_l1] body:[rh_p12_rn_l1]\n [9/28] [None] attached joint:[rh_l2] body:[rh_p12_rn_l2]\n [10/28] [None] attached joint:[None] body:[body_obj_mug_5]\n [11/28] [None] attached joint:[None] body:[body_obj_mug_5]\n [12/28] [None] attached joint:[None] body:[body_obj_mug_5]\n [13/28] [None] attached joint:[None] body:[body_obj_mug_5]\n [14/28] [None] attached joint:[None] body:[body_obj_mug_5]\n [15/28] [None] attached joint:[None] body:[body_obj_mug_5]\n [16/28] [None] attached joint:[None] body:[body_obj_plate_11]\n [17/28] [None] attached joint:[None] body:[body_obj_plate_11]\n [18/28] [None] attached joint:[None] body:[body_obj_plate_11]\n [19/28] [None] attached joint:[None] body:[body_obj_plate_11]\n [20/28] [None] attached joint:[None] body:[body_obj_plate_11]\n [21/28] [None] attached joint:[None] body:[body_obj_plate_11]\n [22/28] [None] attached joint:[None] body:[body_obj_mug_6]\n [23/28] [None] attached joint:[None] body:[body_obj_mug_6]\n [24/28] [None] attached joint:[None] body:[body_obj_mug_6]\n [25/28] [None] attached joint:[None] body:[body_obj_mug_6]\n [26/28] [None] attached joint:[None] body:[body_obj_mug_6]\n [27/28] [None] attached joint:[None] body:[body_obj_mug_6]\n\nFree joint information. n_free_joint:[3]\n [0/3] [None] body_name_attached:[body_obj_mug_5]\n [1/3] [None] body_name_attached:[body_obj_plate_11]\n [2/3] [None] body_name_attached:[body_obj_mug_6]\n\nRevolute joint information. n_rev_joint:[10]\n [0/10] [joint1] range:[-6.283]~[6.283]\n [1/10] [joint2] range:[-6.283]~[6.283]\n [2/10] [joint3] range:[-6.283]~[6.283]\n [3/10] [joint4] range:[-6.283]~[6.283]\n [4/10] [joint5] range:[-6.283]~[6.283]\n [5/10] [joint6] range:[-6.283]~[6.283]\n [6/10] [rh_r1] range:[0.000]~[1.100]\n [7/10] [rh_r2] range:[0.000]~[1.000]\n [8/10] [rh_l1] range:[0.000]~[1.100]\n [9/10] [rh_l2] range:[0.000]~[1.000]\n\nPrismatic joint information. n_pri_joint:[0]\n\nControl information. n_ctrl:[10]\n [0/10] [actuator_joint1] range:[-6.283]~[6.283] gear:[1.00] type:[JOINT]\n [1/10] [actuator_joint2] range:[-6.283]~[6.283] gear:[1.00] type:[JOINT]\n [2/10] [actuator_joint3] range:[-6.283]~[6.283] gear:[1.00] type:[JOINT]\n [3/10] [actuator_joint4] range:[-6.283]~[6.283] gear:[1.00] type:[JOINT]\n [4/10] [actuator_joint5] range:[-6.283]~[6.283] gear:[1.00] type:[JOINT]\n [5/10] [actuator_joint6] range:[-6.283]~[6.283] gear:[1.00] type:[JOINT]\n [6/10] [actuator_rh_r1] range:[0.000]~[1.100] gear:[1.00] type:[JOINT]\n [7/10] [actuator_rh_r2] range:[0.000]~[1.000] gear:[1.00] type:[JOINT]\n [8/10] [actuator_rh_l1] range:[0.000]~[1.100] gear:[1.00] type:[JOINT]\n [9/10] [actuator_rh_l2] range:[0.000]~[1.000] gear:[1.00] type:[JOINT]\n\nCamera information. n_cam:[4]\n [0/4] [agentview] fov:[60.0]\n [1/4] [topview] fov:[90.0]\n [2/4] [sideview] fov:[90.0]\n [3/4] [egocentric] fov:[90.0]\n\nn_sensor:[0]\nsensor_names:[]\nn_site:[9]\nsite_names:['bottom_site_mug_5', 'top_site_mug_5', 'horizontal_radius_site_mug_5', 'bottom_site_plate_11', 'top_site_plate_11', 'horizontal_radius_site_plate_11', 'bottom_site_mug_6', 'top_site_mug_6', 'horizontal_radius_site_mug_6']\n-----------------------------------------------------------------------------\nenv:[Tabletop] reset\nenv:[Tabletop] reset\nenv:[Tabletop] initalize viewer\nDONE INITIALIZATION\nenv:[Tabletop] reset\nDONE INITIALIZATION\n{\"policy\": \"pi0\", \"seed\": 3018, \"success\": true, \"physical_success\": true, \"action_steps\": 472, \"sim_steps\": 11825, \"elapsed_s\": 187.51, \"instruction\": \"Place the red mug on the plate.\", \"debug\": {\"xy_dist\": 0.019316906109452248, \"target_pos\": [0.299, -0.2693, 0.8555], \"plate_pos\": [0.3001, -0.25, 0.808], \"tcp_pos\": [0.3072, -0.2275, 0.9074], \"gripper\": 0.08093565056281846, \"success\": true, \"physical_success\": true, \"physical_lifted_enough\": true, \"physical_final_upright\": true, \"physical_min_lift\": 0.03, \"physical_min_lift_steps\": 3, \"physical_final_upright_cos_threshold\": 0.7, \"target_body\": \"body_obj_mug_5\", \"initial_target_pos\": [0.3218908213064042, 0.0052485008631459745, 0.84465294940151], \"initial_plate_pos\": [0.2999889147642407, -0.2499844847894567, 0.8078853128532457], \"final_target_pos\": [0.299, -0.2693, 0.8555], \"max_target_lift\": 0.10731077156529367, \"lifted_steps\": 3527, \"final_target_upright_cos\": 0.9748101428264967}}\n\n-----------------------------------------------------------------------------\nname:[Tabletop] dt:[0.002] HZ:[500]\n n_qpos:[31] n_qvel:[28] n_qacc:[28] n_ctrl:[10]\n integrator:[IMPLICITFAST]\n\nn_body:[23]\n [0/23] [world] mass:[0.00]kg\n [1/23] [front_object_table] mass:[1.00]kg\n [2/23] [camera] mass:[0.00]kg\n [3/23] [camera2] mass:[0.00]kg\n [4/23] [camera3] mass:[0.00]kg\n [5/23] [link1] mass:[2.06]kg\n [6/23] [link2] mass:[3.68]kg\n [7/23] [link3] mass:[2.39]kg\n [8/23] [link4] mass:[1.40]kg\n [9/23] [link5] mass:[1.40]kg\n [10/23] [link6] mass:[0.65]kg\n [11/23] [camera_center] mass:[0.00]kg\n [12/23] [tcp_link] mass:[0.32]kg\n [13/23] [rh_p12_rn_r1] mass:[0.07]kg\n [14/23] [rh_p12_rn_r2] mass:[0.02]kg\n [15/23] [rh_p12_rn_l1] mass:[0.07]kg\n [16/23] [rh_p12_rn_l2] mass:[0.02]kg\n [17/23] [body_obj_mug_5] mass:[0.00]kg\n [18/23] [object_mug_5] mass:[0.08]kg\n [19/23] [body_obj_plate_11] mass:[0.00]kg\n [20/23] [object_plate_11] mass:[0.10]kg\n [21/23] [body_obj_mug_6] mass:[0.00]kg\n [22/23] [object_mug_6] mass:[0.08]kg\nbody_total_mass:[13.35]kg\n\nn_geom:[116]\ngeom_names:['floor', None, 'front_object_table', None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None]\n\nn_mesh:[112]\nmesh_names:['base_unit', 'link1', 'link2', 'link3', 'link4', 'link5', 'link6', 'flange', 'base', 'r1', 'r2', 'l1', 'l2', 'mug_5_normalized_0_vis', 'mug_5_normalized_collision_22._coll', 'mug_5_normalized_collision_23._coll', 'mug_5_normalized_collision_21._coll', 'mug_5_normalized_collision_20._coll', 'mug_5_normalized_collision_24._coll', 'mug_5_normalized_collision_30._coll', 'mug_5_normalized_collision_18._coll', 'mug_5_normalized_collision_19._coll', 'mug_5_normalized_collision_31._coll', 'mug_5_normalized_collision_25._coll', 'mug_5_normalized_collision_27._coll', 'mug_5_normalized_collision_26._coll', 'mug_5_normalized_collision_9._coll', 'mug_5_normalized_collision_8._coll', 'mug_5_normalized_collision_6._coll', 'mug_5_normalized_collision_7._coll', 'mug_5_normalized_collision_5._coll', 'mug_5_normalized_collision_4._coll', 'mug_5_normalized_collision_0._coll', 'mug_5_normalized_collision_1._coll', 'mug_5_normalized_collision_3._coll', 'mug_5_normalized_collision_2._coll', 'mug_5_normalized_collision_17._coll', 'mug_5_normalized_collision_16._coll', 'mug_5_normalized_collision_28._coll', 'mug_5_normalized_collision_14._coll', 'mug_5_normalized_collision_15._coll', 'mug_5_normalized_collision_29._coll', 'mug_5_normalized_collision_11._coll', 'mug_5_normalized_collision_10._coll', 'mug_5_normalized_collision_12._coll', 'mug_5_normalized_collision_13._coll', 'plate_11_normalized_0_vis', 'plate_11_normalized_collision_22._coll', 'plate_11_normalized_collision_23._coll', 'plate_11_normalized_collision_21._coll', 'plate_11_normalized_collision_20._coll', 'plate_11_normalized_collision_24._coll', 'plate_11_normalized_collision_30._coll', 'plate_11_normalized_collision_18._coll', 'plate_11_normalized_collision_19._coll', 'plate_11_normalized_collision_31._coll', 'plate_11_normalized_collision_25._coll', 'plate_11_normalized_collision_27._coll', 'plate_11_normalized_collision_26._coll', 'plate_11_normalized_collision_9._coll', 'plate_11_normalized_collision_8._coll', 'plate_11_normalized_collision_6._coll', 'plate_11_normalized_collision_7._coll', 'plate_11_normalized_collision_5._coll', 'plate_11_normalized_collision_4._coll', 'plate_11_normalized_collision_0._coll', 'plate_11_normalized_collision_1._coll', 'plate_11_normalized_collision_3._coll', 'plate_11_normalized_collision_2._coll', 'plate_11_normalized_collision_17._coll', 'plate_11_normalized_collision_16._coll', 'plate_11_normalized_collision_28._coll', 'plate_11_normalized_collision_14._coll', 'plate_11_normalized_collision_15._coll', 'plate_11_normalized_collision_29._coll', 'plate_11_normalized_collision_11._coll', 'plate_11_normalized_collision_10._coll', 'plate_11_normalized_collision_12._coll', 'plate_11_normalized_collision_13._coll', 'mug_6_normalized_0_vis', 'mug_6_normalized_collision_22._coll', 'mug_6_normalized_collision_23._coll', 'mug_6_normalized_collision_21._coll', 'mug_6_normalized_collision_20._coll', 'mug_6_normalized_collision_24._coll', 'mug_6_normalized_collision_30._coll', 'mug_6_normalized_collision_18._coll', 'mug_6_normalized_collision_19._coll', 'mug_6_normalized_collision_31._coll', 'mug_6_normalized_collision_25._coll', 'mug_6_normalized_collision_27._coll', 'mug_6_normalized_collision_26._coll', 'mug_6_normalized_collision_9._coll', 'mug_6_normalized_collision_8._coll', 'mug_6_normalized_collision_6._coll', 'mug_6_normalized_collision_7._coll', 'mug_6_normalized_collision_5._coll', 'mug_6_normalized_collision_4._coll', 'mug_6_normalized_collision_0._coll', 'mug_6_normalized_collision_1._coll', 'mug_6_normalized_collision_3._coll', 'mug_6_normalized_collision_2._coll', 'mug_6_normalized_collision_17._coll', 'mug_6_normalized_collision_16._coll', 'mug_6_normalized_collision_28._coll', 'mug_6_normalized_collision_14._coll', 'mug_6_normalized_collision_15._coll', 'mug_6_normalized_collision_29._coll', 'mug_6_normalized_collision_11._coll', 'mug_6_normalized_collision_10._coll', 'mug_6_normalized_collision_12._coll', 'mug_6_normalized_collision_13._coll']\n\nn_joint:[13]\n [0/13] [joint1] axis:[0. 0. 1.]\n [1/13] [joint2] axis:[0. 1. 0.]\n [2/13] [joint3] axis:[0. 1. 0.]\n [3/13] [joint4] axis:[0. 1. 0.]\n [4/13] [joint5] axis:[0. 0. 1.]\n [5/13] [joint6] axis:[0. 1. 0.]\n [6/13] [rh_r1] axis:[1. 0. 0.]\n [7/13] [rh_r2] axis:[-1.  0.  0.]\n [8/13] [rh_l1] axis:[-1.  0.  0.]\n [9/13] [rh_l2] axis:[1. 0. 0.]\n [10/13] [None] axis:[0. 0. 1.]\n [11/13] [None] axis:[0. 0. 1.]\n [12/13] [None] axis:[0. 0. 1.]\n\nn_dof:[28] (=number of rows of Jacobian)\n [0/28] [None] attached joint:[joint1] body:[link1]\n [1/28] [None] attached joint:[joint2] body:[link2]\n [2/28] [None] attached joint:[joint3] body:[link3]\n [3/28] [None] attached joint:[joint4] body:[link4]\n [4/28] [None] attached joint:[joint5] body:[link5]\n [5/28] [None] attached joint:[joint6] body:[link6]\n [6/28] [None] attached joint:[rh_r1] body:[rh_p12_rn_r1]\n [7/28] [None] attached joint:[rh_r2] body:[rh_p12_rn_r2]\n [8/28] [None] attached joint:[rh_l1] body:[rh_p12_rn_l1]\n [9/28] [None] attached joint:[rh_l2] body:[rh_p12_rn_l2]\n [10/28] [None] attached joint:[None] body:[body_obj_mug_5]\n [11/28] [None] attached joint:[None] body:[body_obj_mug_5]\n [12/28] [None] attached joint:[None] body:[body_obj_mug_5]\n [13/28] [None] attached joint:[None] body:[body_obj_mug_5]\n [14/28] [None] attached joint:[None] body:[body_obj_mug_5]\n [15/28] [None] attached joint:[None] body:[body_obj_mug_5]\n [16/28] [None] attached joint:[None] body:[body_obj_plate_11]\n [17/28] [None] attached joint:[None] body:[body_obj_plate_11]\n [18/28] [None] attached joint:[None] body:[body_obj_plate_11]\n [19/28] [None] attached joint:[None] body:[body_obj_plate_11]\n [20/28] [None] attached joint:[None] body:[body_obj_plate_11]\n [21/28] [None] attached joint:[None] body:[body_obj_plate_11]\n [22/28] [None] attached joint:[None] body:[body_obj_mug_6]\n [23/28] [None] attached joint:[None] body:[body_obj_mug_6]\n [24/28] [None] attached joint:[None] body:[body_obj_mug_6]\n [25/28] [None] attached joint:[None] body:[body_obj_mug_6]\n [26/28] [None] attached joint:[None] body:[body_obj_mug_6]\n [27/28] [None] attached joint:[None] body:[body_obj_mug_6]\n\nFree joint information. n_free_joint:[3]\n [0/3] [None] body_name_attached:[body_obj_mug_5]\n [1/3] [None] body_name_attached:[body_obj_plate_11]\n [2/3] [None] body_name_attached:[body_obj_mug_6]\n\nRevolute joint information. n_rev_joint:[10]\n [0/10] [joint1] range:[-6.283]~[6.283]\n [1/10] [joint2] range:[-6.283]~[6.283]\n [2/10] [joint3] range:[-6.283]~[6.283]\n [3/10] [joint4] range:[-6.283]~[6.283]\n [4/10] [joint5] range:[-6.283]~[6.283]\n [5/10] [joint6] range:[-6.283]~[6.283]\n [6/10] [rh_r1] range:[0.000]~[1.100]\n [7/10] [rh_r2] range:[0.000]~[1.000]\n [8/10] [rh_l1] range:[0.000]~[1.100]\n [9/10] [rh_l2] range:[0.000]~[1.000]\n\nPrismatic joint information. n_pri_joint:[0]\n\nControl information. n_ctrl:[10]\n [0/10] [actuator_joint1] range:[-6.283]~[6.283] gear:[1.00] type:[JOINT]\n [1/10] [actuator_joint2] range:[-6.283]~[6.283] gear:[1.00] type:[JOINT]\n [2/10] [actuator_joint3] range:[-6.283]~[6.283] gear:[1.00] type:[JOINT]\n [3/10] [actuator_joint4] range:[-6.283]~[6.283] gear:[1.00] type:[JOINT]\n [4/10] [actuator_joint5] range:[-6.283]~[6.283] gear:[1.00] type:[JOINT]\n [5/10] [actuator_joint6] range:[-6.283]~[6.283] gear:[1.00] type:[JOINT]\n [6/10] [actuator_rh_r1] range:[0.000]~[1.100] gear:[1.00] type:[JOINT]\n [7/10] [actuator_rh_r2] range:[0.000]~[1.000] gear:[1.00] type:[JOINT]\n [8/10] [actuator_rh_l1] range:[0.000]~[1.100] gear:[1.00] type:[JOINT]\n [9/10] [actuator_rh_l2] range:[0.000]~[1.000] gear:[1.00] type:[JOINT]\n\nCamera information. n_cam:[4]\n [0/4] [agentview] fov:[60.0]\n [1/4] [topview] fov:[90.0]\n [2/4] [sideview] fov:[90.0]\n [3/4] [egocentric] fov:[90.0]\n\nn_sensor:[0]\nsensor_names:[]\nn_site:[9]\nsite_names:['bottom_site_mug_5', 'top_site_mug_5', 'horizontal_radius_site_mug_5', 'bottom_site_plate_11', 'top_site_plate_11', 'horizontal_radius_site_plate_11', 'bottom_site_mug_6', 'top_site_mug_6', 'horizontal_radius_site_mug_6']\n-----------------------------------------------------------------------------\nenv:[Tabletop] reset\nenv:[Tabletop] reset\nenv:[Tabletop] initalize viewer\nDONE INITIALIZATION\nenv:[Tabletop] reset\nDONE INITIALIZATION\n{\"policy\": \"pi0\", \"seed\": 3019, \"success\": true, \"physical_success\": true, \"action_steps\": 402, \"sim_steps\": 10075, \"elapsed_s\": 159.624, \"instruction\": \"Place the red mug on the plate.\", \"debug\": {\"xy_dist\": 0.02516002207994461, \"target_pos\": [0.3053, -0.2748, 0.8538], \"plate_pos\": [0.3001, -0.2501, 0.8084], \"tcp_pos\": [0.3087, -0.2286, 0.9082], \"gripper\": 0.08860601938282972, \"success\": true, \"physical_success\": true, \"physical_lifted_enough\": true, \"physical_final_upright\": true, \"physical_min_lift\": 0.03, \"physical_min_lift_steps\": 3, \"physical_final_upright_cos_threshold\": 0.7, \"target_body\": \"body_obj_mug_5\", \"initial_target_pos\": [0.3247087527800675, 0.012232299571383902, 0.8446529494015099], \"initial_plate_pos\": [0.2999889147642407, -0.2499844847894567, 0.8078853128532457], \"final_target_pos\": [0.3053, -0.2748, 0.8538], \"max_target_lift\": 0.11369674420924203, \"lifted_steps\": 3975, \"final_target_upright_cos\": 0.9969652570495078}}\nPi0 formal result: success=12/14, physical_success=12/14\n"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": "WARNING:root:Loading safetensors weights on CPU first, then moving the policy to cuda. This avoids ROCm crashes seen when safetensors loads large Pi0 weights directly onto the GPU.\n\rpi0 formal eval:   0%|          | 0/3 [00:00<?, ?it/s]15_pi0_end_to_end.ipynb:cell-20:263: UserWarning: Creating a tensor from a list of numpy.ndarrays is extremely slow. Please consider converting the list to a single numpy.ndarray with numpy.array() before converting to a tensor. (Triggered internally at /__w/TheRock/TheRock/external-builds/pytorch/pytorch/torch/csrc/utils/tensor_new.cpp:253.)\n$REMOTE_HOME/miniconda3/envs/lerobot/lib/python3.10/site-packages/transformers/integrations/sdpa_attention.py:66: UserWarning: Flash Efficient attention on Current AMD GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at /__w/TheRock/TheRock/external-builds/pytorch/pytorch/aten/src/ATen/native/transformers/hip/sdp_utils.cpp:323.)\n  attn_output = torch.nn.functional.scaled_dot_product_attention(\n$REMOTE_HOME/miniconda3/envs/lerobot/lib/python3.10/site-packages/transformers/integrations/sdpa_attention.py:66: UserWarning: Mem Efficient attention on Current AMD GPU is still experimental. Enable it with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. (Triggered internally at /__w/TheRock/TheRock/external-builds/pytorch/pytorch/aten/src/ATen/native/transformers/hip/sdp_utils.cpp:383.)\n  attn_output = torch.nn.functional.scaled_dot_product_attention(\n\rpi0 formal eval:  33%|###3      | 1/3 [03:49<07:39, 229.99s/it]\rpi0 formal eval:  67%|######6   | 2/3 [07:03<03:28, 208.80s/it]\rpi0 formal eval: 100%|##########| 3/3 [09:50<00:00, 189.39s/it]\rpi0 formal eval: 100%|##########| 3/3 [09:50<00:00, 196.75s/it]\n"
    },
    {
     "output_type": "display_data",
     "data": {
      "text/markdown": "| 结果文件 | episodes | legacy_success | physical_success |\n| --- | --- | --- | --- |\n| $OUTPUT_ROOT/pi0_formal_protected_seed3007_14ep.jsonl | 14 | 12/14 | 12/14 |"
     },
     "metadata": {}
    }
   ],
   "source": [
    "# PI0_FORMAL_EVAL_CELL\n",
    "seed_text = os.environ.get(\n",
    "    \"PI0_EVAL_SEEDS\",\n",
    "    \"3007,3002,3004,3006,3010,3011,3012,3013,3014,3015,3016,3017,3018,3019\",\n",
    ")\n",
    "eval_seeds = [int(x.strip()) for x in seed_text.split(\",\") if x.strip()]\n",
    "eval_policy = resolve_eval_policy(PI0_POLICY_PATH, LONG_OUTPUT, \"PI0_EVAL_POLICY_PATH\")\n",
    "result_path = OUTPUT_ROOT / \"pi0_formal_protected_seed3007_14ep.jsonl\"\n",
    "run_pi0_eval_native_in_notebook(\n",
    "    eval_policy,\n",
    "    result_path,\n",
    "    eval_seeds,\n",
    "    repo_id=DATASET_REPO_ID,\n",
    "    dataset_root=TRAIN_DATA_ROOT,\n",
    "    render=env_flag(\"RENDER_EVAL\"),\n",
    "    enabled=RUN_EVAL,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## Checkpoint 7：失败桶和视频复核\n"
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": "{\n  \"episodes\": 14,\n  \"physical_success_count\": 12,\n  \"by_color\": {\n    \"blue\": \"6/7\",\n    \"red\": \"6/7\"\n  },\n  \"failures\": [\n    \"blue seed3007: lifted but upright unstable\",\n    \"red seed3012: no enough lift/contact\"\n  ],\n  \"do_not_claim\": \"hard8 8/8 has not been recovered\"\n}\n"
    },
    {
     "output_type": "display_data",
     "data": {
      "text/markdown": "**Pi0 raw 与诊断 finisher 对照视频**"
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/html": "<video controls muted preload='metadata' width='960'><source src='../assets/pi0_ep2_raw_vs_finisher_side_by_side.mp4' type='video/mp4'></video>"
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/markdown": "**Pi0 诊断图：raw、scaffold 与严格判定**"
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/html": "<img src='../assets/pi0_raw_vs_finisher_diagnostic.png' width='960'>"
     },
     "metadata": {}
    }
   ],
   "source": [
    "summary = {\n",
    "    \"episodes\": 14,\n",
    "    \"physical_success_count\": 12,\n",
    "    \"by_color\": {\"blue\": \"6/7\", \"red\": \"6/7\"},\n",
    "    \"failures\": [\n",
    "        \"blue seed3007: lifted but upright unstable\",\n",
    "        \"red seed3012: no enough lift/contact\",\n",
    "    ],\n",
    "    \"do_not_claim\": \"hard8 8/8 has not been recovered\",\n",
    "}\n",
    "print(json.dumps(summary, ensure_ascii=False, indent=2))\n",
    "show_video(\"pi0_ep2_raw_vs_finisher_side_by_side.mp4\", \"Pi0 raw 与诊断 finisher 对照视频\")\n",
    "show_image(\"pi0_raw_vs_finisher_diagnostic.png\", \"Pi0 诊断图：raw、scaffold 与严格判定\")\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": 11,
   "errors": 0,
   "record_markers": [],
   "prepare_markers": [],
   "long_tasks_enabled": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
