{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "# 07 从零采集 MuJoCo LeRobot 示教数据\n\n        这一节把 `5.language_env.ipynb` 的交互采集流程整理成可配置、可复核的版本。最终产物是一个包含双相机图像、6 维关节状态、7 维动作和红/蓝杯语言指令的 LeRobot 数据集。\n\n        采集不是纯 GPU 任务，NVIDIA 或 AMD 设备都可以完成。真正需要保持一致的是 MuJoCo 场景、LeRobot 数据格式、控制频率、state/action 定义和成功判定。远端 Jupyter 无法稳定接收 MuJoCo viewer 键盘事件时，可以在有桌面的机器采集，再把数据目录同步到 AMD 训练机。\n"
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": "TOPIC_ROOT = $NOTEBOOK_TOPIC_ROOT\nPROJECT_ROOT = $PROJECT_ROOT\nDATA_ROOT = $PROJECT_ROOT\nOUTPUT_ROOT = $OUTPUT_ROOT\nMODEL_ROOT = $MODEL_ROOT\n"
        }
      ],
      "source": "from pathlib import Path\nimport json\nimport os\nimport shutil\nimport subprocess\nimport sys\n\n\ndef find_topic_root():\n    cwd = Path.cwd().resolve()\n    for candidate in [cwd, *cwd.parents]:\n        if (candidate / \"assets\" / \"metrics_snapshot.json\").exists():\n            return candidate\n    raise RuntimeError(\"请从 AMD ROCm 专题目录或 notebooks 子目录启动 Jupyter。\")\n\n\nTOPIC_ROOT = find_topic_root()\nASSET_DIR = TOPIC_ROOT / \"assets\"\nNOTEBOOK_DIR = TOPIC_ROOT / \"notebooks\"\nPROJECT_ROOT = Path(os.environ.get(\"PROJECT_ROOT\", \"/path/to/every-embodied/mujoco_pnp\"))\nDATA_ROOT = Path(os.environ.get(\"DATA_ROOT\", \"/path/to/datasets/every_embodied\"))\nOUTPUT_ROOT = Path(os.environ.get(\"OUTPUT_ROOT\", TOPIC_ROOT / \"outputs\"))\nMODEL_ROOT = Path(os.environ.get(\"MODEL_ROOT\", PROJECT_ROOT / \"ckpt\"))\n\nprint(\"TOPIC_ROOT =\", TOPIC_ROOT)\nprint(\"PROJECT_ROOT =\", PROJECT_ROOT)\nprint(\"DATA_ROOT =\", DATA_ROOT)\nprint(\"OUTPUT_ROOT =\", OUTPUT_ROOT)\nprint(\"MODEL_ROOT =\", MODEL_ROOT)\n"
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {},
      "outputs": [],
      "source": "try:\n    from IPython.display import Image, Markdown, display\nexcept Exception:\n    class Markdown(str):\n        pass\n\n    def display(obj):\n        print(obj)\n\n    def Image(filename=None, width=None):\n        return f\"[image] {filename}\"\n\n\ndef show_asset(filename, width=960):\n    path = ASSET_DIR / filename\n    if path.exists():\n        display(Image(filename=str(path), width=width))\n    else:\n        print(f\"缺少素材：{path}\")\n\n\ndef md_table(headers, rows):\n    lines = [\"| \" + \" | \".join(headers) + \" |\", \"| \" + \" | \".join([\"---\"] * len(headers)) + \" |\"]\n    for row in rows:\n        lines.append(\"| \" + \" | \".join(str(x) for x in row) + \" |\")\n    display(Markdown(\"\\n\".join(lines)))\n"
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {},
      "outputs": [],
      "source": "import shlex\n\ntry:\n    import yaml\nexcept ImportError as exc:\n    raise RuntimeError(\"当前环境缺少 PyYAML，请先执行 pip install pyyaml。\") from exc\n\n\ndef require_project_layout():\n    required = [\n        PROJECT_ROOT / \"train_model.py\",\n        PROJECT_ROOT / \"env_config.py\",\n        PROJECT_ROOT / \"asset\" / \"example_scene_y2.xml\",\n        PROJECT_ROOT / \"mujoco_env\" / \"y_env2.py\",\n    ]\n    missing = [path for path in required if not path.exists()]\n    if missing:\n        raise FileNotFoundError(\n            \"PROJECT_ROOT 不是 04mujoco 教程目录，缺少：\\n\"\n            + \"\\n\".join(str(path) for path in missing)\n        )\n    return True\n\n\ndef dataset_report(dataset_root):\n    dataset_root = Path(dataset_root)\n    info_path = dataset_root / \"meta\" / \"info.json\"\n    tasks_path = dataset_root / \"meta\" / \"tasks.jsonl\"\n    if not info_path.exists():\n        raise FileNotFoundError(f\"找不到 LeRobot 元数据：{info_path}\")\n    info = json.loads(info_path.read_text(encoding=\"utf-8\"))\n    tasks = []\n    if tasks_path.exists():\n        tasks = [\n            json.loads(line)[\"task\"]\n            for line in tasks_path.read_text(encoding=\"utf-8\").splitlines()\n            if line.strip()\n        ]\n    features = info.get(\"features\", {})\n    rows = [\n        (\"repo_id\", info.get(\"repo_id\", \"\")),\n        (\"episodes\", info.get(\"total_episodes\", 0)),\n        (\"frames\", info.get(\"total_frames\", 0)),\n        (\"fps\", info.get(\"fps\", \"\")),\n        (\"state shape\", features.get(\"observation.state\", {}).get(\"shape\")),\n        (\"action shape\", features.get(\"action\", {}).get(\"shape\")),\n        (\"tasks\", \" / \".join(tasks)),\n    ]\n    md_table([\"数据项\", \"读取结果\"], rows)\n    return info, tasks\n\n\ndef make_training_config(\n    policy_type,\n    dataset_repo_id,\n    dataset_root,\n    output_dir,\n    steps,\n    batch_size,\n    chunk_size,\n    n_action_steps,\n    save_freq,\n    seed=42,\n):\n    return {\n        \"dataset\": {\"repo_id\": dataset_repo_id, \"root\": str(Path(dataset_root))},\n        \"policy\": {\n            \"type\": policy_type,\n            \"chunk_size\": int(chunk_size),\n            \"n_action_steps\": int(n_action_steps),\n            \"device\": \"cuda\",\n        },\n        \"save_checkpoint\": True,\n        \"output_dir\": str(Path(output_dir)),\n        \"batch_size\": int(batch_size),\n        \"job_name\": Path(output_dir).name,\n        \"resume\": False,\n        \"seed\": int(seed),\n        \"num_workers\": 4,\n        \"steps\": int(steps),\n        \"eval_freq\": 0,\n        \"log_freq\": max(1, min(50, int(steps))),\n        \"save_freq\": int(save_freq),\n        \"use_policy_training_preset\": True,\n        \"wandb\": {\n            \"enable\": False,\n            \"project\": f\"every_embodied_{policy_type}\",\n            \"entity\": None,\n            \"disable_artifact\": True,\n        },\n    }\n\n\ndef write_training_config(name, config):\n    config_dir = OUTPUT_ROOT / \"configs\"\n    config_dir.mkdir(parents=True, exist_ok=True)\n    path = config_dir / f\"{name}.yaml\"\n    path.write_text(\n        yaml.safe_dump(config, allow_unicode=True, sort_keys=False),\n        encoding=\"utf-8\",\n    )\n    print(\"已写出配置：\", path)\n    return path\n\n\ndef training_command(config_path):\n    return [\n        sys.executable,\n        str(PROJECT_ROOT / \"train_model.py\"),\n        \"--config_path\",\n        str(Path(config_path)),\n    ]\n\n\ndef run_training(config_path, enabled=False):\n    require_project_layout()\n    command = training_command(config_path)\n    print(\"$\", shlex.join(command))\n    if not enabled:\n        print(\"当前只预览命令。确认配置后，把对应 RUN_* 开关改为 True。\")\n        return None\n    env = os.environ.copy()\n    env[\"PYTHONPATH\"] = f\"{PROJECT_ROOT}:{env.get('PYTHONPATH', '')}\"\n    return subprocess.run(command, cwd=PROJECT_ROOT, env=env, check=True)\n\n\ndef show_rocm_resources():\n    command = [\"rocm-smi\", \"--showuse\", \"--showmemuse\", \"--showtemp\"]\n    print(\"$\", shlex.join(command))\n    if shutil.which(command[0]) is None:\n        print(\"未找到 rocm-smi。确认当前机器是否安装 ROCm。\")\n        return\n    result = subprocess.run(command, check=False, capture_output=True, text=True)\n    if result.stdout:\n        print(result.stdout.rstrip())\n    if result.stderr:\n        print(result.stderr.rstrip(), file=sys.stderr)\n"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "## Checkpoint 1：设置采集目录与安全开关\n"
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": "dataset repo_id = datawhale_eai_pnp_language_local\ncollection root = $TRAIN_DATA_ROOT\nnum demos = 20\ninstructions = ['Place the red mug on the plate.', 'Place the blue mug on the plate.']\nfour-view videos = $OUTPUT_ROOT/collection_four_view\n"
        }
      ],
      "source": "DATASET_REPO_ID = os.environ.get(\"DATASET_REPO_ID\", \"datawhale_eai_pnp_language_local\")\nCOLLECTION_ROOT = Path(\n    os.environ.get(\"COLLECTION_ROOT\", DATA_ROOT / \"omy_pnp_language\")\n).expanduser()\nNUM_DEMOS = int(os.environ.get(\"NUM_DEMOS\", \"20\"))\nSEED_START = int(os.environ.get(\"SEED_START\", \"0\"))\nINSTRUCTIONS = [\n    \"Place the red mug on the plate.\",\n    \"Place the blue mug on the plate.\",\n]\nRUN_INTERACTIVE_COLLECTION = False\nOVERWRITE_DATASET = False\nRECORD_FOUR_VIEW_VIDEO = True\nCOLLECTION_VIDEO_DIR = Path(\n    os.environ.get(\"COLLECTION_VIDEO_DIR\", OUTPUT_ROOT / \"collection_four_view\")\n).expanduser()\n\nprint(\"dataset repo_id =\", DATASET_REPO_ID)\nprint(\"collection root =\", COLLECTION_ROOT)\nprint(\"num demos =\", NUM_DEMOS)\nprint(\"instructions =\", INSTRUCTIONS)\nprint(\"four-view videos =\", COLLECTION_VIDEO_DIR)\n"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "`COLLECTION_ROOT` 应放在大容量数据盘。默认不开启采集，也不会删除已有目录；只有明确把 `OVERWRITE_DATASET=True` 后才允许覆盖。20 条可以跑通语言条件小实验，想做位置泛化时建议采 30–50 条高质量轨迹，并让红杯、蓝杯和初始位置都得到覆盖。\n"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "## Checkpoint 2：检查项目、显示和数据 schema\n"
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": "MuJoCo scene = $PROJECT_ROOT/asset/example_scene_y2.xml\nDISPLAY = None\n当前没有 DISPLAY。交互采集需要本地桌面、远程桌面或可用的 X11 会话。\n| 字段 | shape | 作用 |\n| --- | --- | --- |\n| observation.image | 256x256x3 | 固定相机 |\n| observation.wrist_image | 256x256x3 | 腕部相机 |\n| observation.state | 6 | 当前 6 关节状态 |\n| action | 7 | 下一步 6 关节目标 + 夹爪 |\n| obj_init | 9 | 三件物体初始 xyz，仅用于回放/审计 |\n"
        }
      ],
      "source": "require_project_layout()\nprint(\"MuJoCo scene =\", PROJECT_ROOT / \"asset\" / \"example_scene_y2.xml\")\nprint(\"DISPLAY =\", os.environ.get(\"DISPLAY\"))\nif not os.environ.get(\"DISPLAY\"):\n    print(\"当前没有 DISPLAY。交互采集需要本地桌面、远程桌面或可用的 X11 会话。\")\n\nFEATURES = {\n    \"observation.image\": {\n        \"dtype\": \"image\",\n        \"shape\": (256, 256, 3),\n        \"names\": [\"height\", \"width\", \"channels\"],\n    },\n    \"observation.wrist_image\": {\n        \"dtype\": \"image\",\n        \"shape\": (256, 256, 3),\n        \"names\": [\"height\", \"width\", \"channels\"],\n    },\n    \"observation.state\": {\n        \"dtype\": \"float32\",\n        \"shape\": (6,),\n        \"names\": [\"state\"],\n    },\n    \"action\": {\n        \"dtype\": \"float32\",\n        \"shape\": (7,),\n        \"names\": [\"action\"],\n    },\n    \"obj_init\": {\n        \"dtype\": \"float32\",\n        \"shape\": (9,),\n        \"names\": [\n            \"red_x\", \"red_y\", \"red_z\",\n            \"blue_x\", \"blue_y\", \"blue_z\",\n            \"plate_x\", \"plate_y\", \"plate_z\",\n        ],\n    },\n}\nmd_table(\n    [\"字段\", \"shape\", \"作用\"],\n    [\n        (\"observation.image\", \"256x256x3\", \"固定相机\"),\n        (\"observation.wrist_image\", \"256x256x3\", \"腕部相机\"),\n        (\"observation.state\", \"6\", \"当前 6 关节状态\"),\n        (\"action\", \"7\", \"下一步 6 关节目标 + 夹爪\"),\n        (\"obj_init\", \"9\", \"三件物体初始 xyz，仅用于回放/审计\"),\n    ],\n)\n"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "## Checkpoint 3：定义严格成功判定\n"
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {},
      "outputs": [],
      "source": "import numpy as np\n\n\ndef collection_snapshot(env, initial_target_z, initial_plate_pos, max_target_lift, max_lifted_run):\n    target_pos = np.asarray(env.env.get_p_body(env.obj_target), dtype=np.float64)\n    plate_pos = np.asarray(env.env.get_p_body(\"body_obj_plate_11\"), dtype=np.float64)\n    target_R = np.asarray(env.env.get_R_body(env.obj_target), dtype=np.float64)\n    legacy_success = bool(env.check_success())\n    upright_cos = float(target_R[2, 2])\n    z_gap = abs(float(target_pos[2] - plate_pos[2]))\n    plate_xy_displacement = float(\n        np.linalg.norm(plate_pos[:2] - np.asarray(initial_plate_pos)[:2])\n    )\n    placement_candidate = bool(\n        legacy_success\n        and max_target_lift >= 0.03\n        and max_lifted_run >= 3\n        and upright_cos >= 0.7\n        and z_gap < 0.08\n        and plate_xy_displacement < 0.05\n    )\n    return {\n        \"legacy_success\": legacy_success,\n        \"placement_candidate\": placement_candidate,\n        \"xy_dist\": float(np.linalg.norm(target_pos[:2] - plate_pos[:2])),\n        \"plate_z_gap\": z_gap,\n        \"plate_xy_displacement\": plate_xy_displacement,\n        \"max_target_lift\": float(max_target_lift),\n        \"max_lifted_run\": int(max_lifted_run),\n        \"upright_cos\": upright_cos,\n    }\n"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "旧 `check_success()` 只看终态几何关系，杯子被推到盘子附近也可能触发。这里额外要求杯子至少抬升 3 cm、连续保持 3 个控制步、终态保持直立，并且杯子高度与盘子相符。只有 `physical_success=True` 才写入 episode。\n"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "## Checkpoint 4：创建或加载 LeRobot 数据集\n"
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {},
      "outputs": [],
      "source": "def create_or_load_dataset():\n    from lerobot.common.datasets.lerobot_dataset import LeRobotDataset\n\n    COLLECTION_ROOT.parent.mkdir(parents=True, exist_ok=True)\n    if COLLECTION_ROOT.exists() and OVERWRITE_DATASET:\n        shutil.rmtree(COLLECTION_ROOT)\n    if COLLECTION_ROOT.exists():\n        print(\"继续写入已有数据集：\", COLLECTION_ROOT)\n        return LeRobotDataset(DATASET_REPO_ID, root=COLLECTION_ROOT)\n    print(\"创建新数据集：\", COLLECTION_ROOT)\n    return LeRobotDataset.create(\n        repo_id=DATASET_REPO_ID,\n        root=COLLECTION_ROOT,\n        robot_type=\"omy\",\n        fps=20,\n        features=FEATURES,\n        image_writer_threads=10,\n        image_writer_processes=0,\n    )\n"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "## Checkpoint 5：键盘采集完整循环\n"
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": "采集默认关闭。确认 DISPLAY、路径和覆盖开关后，将 RUN_INTERACTIVE_COLLECTION 改为 True。\n"
        }
      ],
      "source": "def collect_demonstrations():\n    import random\n    import numpy as np\n    from PIL import Image\n    from mujoco_env.y_env2 import SimpleEnv2\n\n    dataset = create_or_load_dataset()\n    xml_path = PROJECT_ROOT / \"asset\" / \"example_scene_y2.xml\"\n    COLLECTION_VIDEO_DIR.mkdir(parents=True, exist_ok=True)\n\n    try:\n        import cv2\n    except ImportError:\n        cv2 = None\n        if RECORD_FOUR_VIEW_VIDEO:\n            print(\"未安装 OpenCV，数据仍会采集，但跳过四视角 MP4。\")\n\n    video_writer = None\n    video_tmp_path = None\n\n    def open_episode_video(episode_id):\n        nonlocal video_writer, video_tmp_path\n        if not RECORD_FOUR_VIEW_VIDEO or cv2 is None:\n            return\n        video_tmp_path = COLLECTION_VIDEO_DIR / f\"episode_{episode_id:03d}.recording.mp4\"\n        video_writer = cv2.VideoWriter(\n            str(video_tmp_path),\n            cv2.VideoWriter_fourcc(*\"mp4v\"),\n            20.0,\n            (640, 480),\n        )\n        if not video_writer.isOpened():\n            raise RuntimeError(f\"无法创建四视角视频：{video_tmp_path}\")\n\n    def write_episode_video(env, agent_image, wrist_image):\n        if video_writer is None:\n            return\n        views = [\n            (\"Agent\", agent_image),\n            (\"Egocentric\", wrist_image),\n            (\"Top\", env.env.get_fixed_cam_rgb(cam_name=\"topview\")),\n            (\"Side\", env.env.get_fixed_cam_rgb(cam_name=\"sideview\")),\n        ]\n        panels = []\n        for label, image in views:\n            panel = cv2.resize(np.asarray(image), (320, 240), interpolation=cv2.INTER_AREA)\n            panel = cv2.cvtColor(panel, cv2.COLOR_RGB2BGR)\n            cv2.rectangle(panel, (0, 0), (150, 28), (20, 20, 20), thickness=-1)\n            cv2.putText(\n                panel, label, (8, 20), cv2.FONT_HERSHEY_SIMPLEX,\n                0.55, (255, 255, 255), 1, cv2.LINE_AA,\n            )\n            panels.append(panel)\n        video_writer.write(np.vstack([np.hstack(panels[:2]), np.hstack(panels[2:])]))\n\n    def close_episode_video(episode_id, keep):\n        nonlocal video_writer, video_tmp_path\n        if video_writer is not None:\n            video_writer.release()\n        if video_tmp_path is not None and video_tmp_path.exists():\n            if keep:\n                final_path = COLLECTION_VIDEO_DIR / f\"episode_{episode_id:03d}_success.mp4\"\n                video_tmp_path.replace(final_path)\n                print(\"saved four-view video:\", final_path)\n            else:\n                video_tmp_path.unlink()\n        video_writer = None\n        video_tmp_path = None\n\n    def reset_episode(env, episode_id):\n        seed = SEED_START + episode_id\n        np.random.seed(seed)\n        random.seed(seed)\n        env.reset(seed=seed)\n        instruction = INSTRUCTIONS[episode_id % len(INSTRUCTIONS)]\n        env.set_instruction(instruction)\n        initial_z = float(env.env.get_p_body(env.obj_target)[2])\n        initial_plate_pos = np.asarray(\n            env.env.get_p_body(\"body_obj_plate_11\"), dtype=np.float64\n        ).copy()\n        print(f\"episode={episode_id} seed={seed} task={instruction}\")\n        return initial_z, initial_plate_pos\n\n    np.random.seed(SEED_START)\n    random.seed(SEED_START)\n    env = SimpleEnv2(str(xml_path), seed=SEED_START, state_type=\"joint_angle\")\n    episode_id = 0\n    record_flag = False\n    initial_target_z, initial_plate_pos = reset_episode(env, episode_id)\n    max_target_lift = 0.0\n    lifted_run = 0\n    max_lifted_run = 0\n    stable_place_steps = 0\n\n    try:\n        while env.env.is_viewer_alive() and episode_id < NUM_DEMOS:\n            env.step_env()\n            if not env.env.loop_every(HZ=20):\n                continue\n\n            target_z = float(env.env.get_p_body(env.obj_target)[2])\n            lift = target_z - initial_target_z\n            max_target_lift = max(max_target_lift, lift)\n            lifted_run = lifted_run + 1 if lift >= 0.03 else 0\n            max_lifted_run = max(max_lifted_run, lifted_run)\n            status = collection_snapshot(\n                env,\n                initial_target_z,\n                initial_plate_pos,\n                max_target_lift,\n                max_lifted_run,\n            )\n            stable_place_steps = stable_place_steps + 1 if status[\"placement_candidate\"] else 0\n            status[\"stable_place_steps\"] = stable_place_steps\n            status[\"physical_success\"] = bool(stable_place_steps >= 5)\n\n            if record_flag and status[\"physical_success\"]:\n                dataset.save_episode()\n                close_episode_video(episode_id, keep=True)\n                print(\"saved:\", json.dumps(status, ensure_ascii=False))\n                episode_id += 1\n                record_flag = False\n                if episode_id >= NUM_DEMOS:\n                    break\n                initial_target_z, initial_plate_pos = reset_episode(env, episode_id)\n                max_target_lift = 0.0\n                lifted_run = 0\n                max_lifted_run = 0\n                stable_place_steps = 0\n                continue\n\n            teleop_delta, reset = env.teleop_robot()\n            if reset:\n                dataset.clear_episode_buffer()\n                close_episode_video(episode_id, keep=False)\n                record_flag = False\n                initial_target_z, initial_plate_pos = reset_episode(env, episode_id)\n                max_target_lift = 0.0\n                lifted_run = 0\n                max_lifted_run = 0\n                stable_place_steps = 0\n                continue\n\n            if not record_flag and np.any(np.abs(teleop_delta) > 1e-8):\n                record_flag = True\n                open_episode_video(episode_id)\n                print(\"Start recording\")\n\n            agent_image, wrist_image = env.grab_image()\n            if record_flag:\n                write_episode_video(env, agent_image, wrist_image)\n            state = env.get_joint_state()[:6].astype(np.float32)\n            env.step(teleop_delta)\n            target_action = env.q[:7].astype(np.float32)\n\n            if record_flag:\n                dataset.add_frame(\n                    {\n                        \"observation.image\": np.asarray(Image.fromarray(agent_image).resize((256, 256))),\n                        \"observation.wrist_image\": np.asarray(Image.fromarray(wrist_image).resize((256, 256))),\n                        \"observation.state\": state,\n                        \"action\": target_action,\n                        \"obj_init\": np.asarray(env.obj_init_pose, dtype=np.float32),\n                    },\n                    task=env.instruction,\n                )\n            env.render(teleop=True, idx=episode_id)\n    finally:\n        close_episode_video(episode_id, keep=False)\n        env.env.close_viewer()\n        shutil.rmtree(dataset.root / \"images\", ignore_errors=True)\n    print(f\"采集完成：{episode_id}/{NUM_DEMOS}\")\n    return dataset\n\n\nif RUN_INTERACTIVE_COLLECTION:\n    dataset = collect_demonstrations()\nelse:\n    print(\"采集默认关闭。确认 DISPLAY、路径和覆盖开关后，将 RUN_INTERACTIVE_COLLECTION 改为 True。\")\n"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "键位与上游教程一致：`W/A/S/D` 控制平面移动，`R/F` 控制高度，`Q/E` 与方向键控制姿态，空格切换夹爪，`Z` 丢弃当前失败回合。每次成功保存后会重新关闭记录开关，避免把 reset 过程混进下一条 episode。模型训练仍只使用固定相机和腕部相机；Top/Side 视角写入独立 MP4，专门用于人工复核，不偷偷扩充模型输入。\n"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "## Checkpoint 6：采集后审计\n"
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": "| 数据项 | 读取结果 |\n| --- | --- |\n| repo_id |  |\n| episodes | 20 |\n| frames | 2621 |\n| fps | 20 |\n| state shape | [6] |\n| action shape | [7] |\n| tasks | Place the blue mug on the plate. / Place the red mug on the plate. |\n| 指令 | episode 数 |\n| --- | --- |\n| Place the blue mug on the plate. | 10 |\n| Place the red mug on the plate. | 10 |\n"
        }
      ],
      "source": "info_path = COLLECTION_ROOT / \"meta\" / \"info.json\"\nif info_path.exists():\n    info, tasks = dataset_report(COLLECTION_ROOT)\n    episodes_path = COLLECTION_ROOT / \"meta\" / \"episodes.jsonl\"\n    episodes = [\n        json.loads(line)\n        for line in episodes_path.read_text(encoding=\"utf-8\").splitlines()\n        if line.strip()\n    ]\n    task_counts = {}\n    for episode in episodes:\n        task = episode.get(\"tasks\", [\"\"])[0]\n        task_counts[task] = task_counts.get(task, 0) + 1\n    md_table([\"指令\", \"episode 数\"], sorted(task_counts.items()))\nelse:\n    print(\"还没有可审计的数据：\", COLLECTION_ROOT)\n"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "数据表通过后，再打开上游 `6.visualize_data.ipynb` 随机回放若干 episode。至少检查图像、state/action shape、夹爪开闭时序、是否真的抬起，以及红蓝杯指令是否平衡。完成这些检查后再进入 08–10 的训练 Notebook。\n"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "## Checkpoint 7：查看已经跑过的数据集实测结果\n"
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": "| 数据项 | AMD 实测值 |\n| --- | --- |\n| repo_id | datawhale_eai_pnp_language |\n| episodes | 20 |\n| frames | 2621 |\n| fps | 20 |\n| state / action | [6] / [7] |\n| 模型输入相机 | observation.image / observation.wrist_image |\n| 复核视频视角 | Agent / Egocentric / Top / Side |\n平均每条轨迹帧数 = 131.1\n"
        }
      ],
      "source": "snapshot = json.loads(\n    (ASSET_DIR / \"collection_dataset_snapshot.json\").read_text(encoding=\"utf-8\")\n)\nrows = [\n    (\"repo_id\", snapshot[\"dataset_repo_id\"]),\n    (\"episodes\", snapshot[\"total_episodes\"]),\n    (\"frames\", snapshot[\"total_frames\"]),\n    (\"fps\", snapshot[\"fps\"]),\n    (\"state / action\", f'{snapshot[\"state_shape\"]} / {snapshot[\"action_shape\"]}'),\n    (\"模型输入相机\", \" / \".join(snapshot[\"stored_cameras\"])),\n    (\"复核视频视角\", \" / \".join(snapshot[\"recorder_views\"])),\n]\nmd_table([\"数据项\", \"AMD 实测值\"], rows)\nprint(\"平均每条轨迹帧数 =\", round(snapshot[\"total_frames\"] / snapshot[\"total_episodes\"], 1))\n"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "## Checkpoint 8：四视角采集/回放样例\n"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "下面的视频是在 AMD 设备上用与键盘采集相同的四视角 recorder 实际录制的严格成功策略回放。它用于先检查画面布局、接触、抬升、搬运和释放是否清楚；它不是人工键盘示教。真正采集时，将 `RUN_INTERACTIVE_COLLECTION=True`，每条成功示教会在 `COLLECTION_VIDEO_DIR` 生成同样布局的 MP4。\n\n            ![MuJoCo 四视角严格成功序列](../assets/pnp_four_view_strict_success_sequence.jpg)\n\n            <video controls width=\"960\" src=\"../assets/pnp_four_view_strict_success.mp4\"></video>\n"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "## Checkpoint 9：确认环境 seed 真的改变物体位置\n"
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": "seed source = $PROJECT_ROOT/mujoco_env/y_env2.py\n发现旧的固定 seed 写法： False\n发现修正后的 seed 写法： True\n"
        }
      ],
      "source": "source_path = PROJECT_ROOT / \"mujoco_env\" / \"y_env2.py\"\nsource = source_path.read_text(encoding=\"utf-8\")\nbad_pattern = \"np.random.seed(seed=0)\"\ngood_pattern = \"np.random.seed(seed=seed)\"\nprint(\"seed source =\", source_path)\nprint(\"发现旧的固定 seed 写法：\", bad_pattern in source)\nprint(\"发现修正后的 seed 写法：\", good_pattern in source)\nif bad_pattern in source:\n    raise RuntimeError(\n        \"当前 y_env2.py 会把所有环境固定成 seed 0。先修复它，再采集位置随机化数据。\"\n    )\n"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "这个检查很重要。旧实现曾把任意传入 seed 都改成 `0`，导致多 seed 评估只改变了策略采样，杯子和盘子位置并没有真正变化。修复后要重新采 30–50 条覆盖不同位置的轨迹；旧固定场景数据仍可用于链路 smoke，但不能作为空间泛化训练集。\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,
      "long_tasks_enabled": false
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}