{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# SMOLVLA：保护训练与评估\n",
"\n",
"本 Notebook 只执行 protected recipe。训练、checkpoint 保存、耗时记录、14 条固定 seed 闭环评估、视频和动作序列图都在 Notebook kernel 内完成。评估严格绑定本次训练生成的 `TRAINED_POLICY_PATH`，不会回退到历史或预训练权重。\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"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' / 'protected' / 'smolvla'))\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\"\"))\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\"
\"))\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",
"\n",
"\n",
"class _NotebookCompactProgress:\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\", \"1\")))\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 _NotebookCompactProgress(iterable, desc, total=total)\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 datetime import datetime, timezone\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",
" run_started_at = datetime.now(timezone.utc).isoformat()\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",
" if cfg.policy.type == \"smolvla\":\n",
" local_vlm_model = os.environ.get(\"SMOLVLA_VLM_MODEL_PATH\")\n",
" if local_vlm_model and Path(local_vlm_model).exists():\n",
" cfg.policy.vlm_model_name = str(Path(local_vlm_model))\n",
" cfg.policy.load_vlm_weights = False\n",
" print(\"local smolvlm processor/config =\", public_path(cfg.policy.vlm_model_name))\n",
"\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",
" run_finished_at = datetime.now(timezone.utc).isoformat()\n",
" duration_s = float(time.perf_counter() - start_all)\n",
" checkpoint_paths = sorted(str(path) for path in output_dir.glob(\"checkpoints/*/pretrained_model\"))\n",
" run_summary = {\n",
" \"recipe\": str(cfg.policy.type),\n",
" \"native_notebook_kernel\": True,\n",
" \"started_at_utc\": run_started_at,\n",
" \"finished_at_utc\": run_finished_at,\n",
" \"duration_s\": duration_s,\n",
" \"steps\": int(cfg.steps),\n",
" \"batch_size\": int(cfg.batch_size),\n",
" \"output_dir\": str(output_dir),\n",
" \"metrics_path\": str(metrics_path),\n",
" \"checkpoint_paths\": checkpoint_paths,\n",
" \"last_metrics\": last_metrics,\n",
" }\n",
" summary_path = output_dir / \"training_run_summary.json\"\n",
" summary_path.write_text(json.dumps(run_summary, ensure_ascii=False, indent=2), encoding=\"utf-8\")\n",
" print(\"训练完成。metrics =\", public_path(metrics_path))\n",
" print(\"训练耗时 =\", f\"{duration_s / 60.0:.2f} 分钟\")\n",
" print(\"训练摘要 =\", public_path(summary_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, \"summary_path\": summary_path, \"last_metrics\": last_metrics, \"duration_s\": duration_s}\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",
" video_dir=None,\n",
" trace_dir=None,\n",
" seed_list=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",
" video_dir=Path(video_dir) if video_dir else None,\n",
" trace_dir=Path(trace_dir) if trace_dir else None,\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",
" selected_cases = set()\n",
" selected_cases = set()\n",
" selected_cases = set()\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 = (list(seed_list)[offset] if seed_list is not None else args.seed_start + offset)\n",
" row = rollout(args, policy, seed)\n",
" row = module.finalize_case_media(args, row, selected_cases)\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 require_trained_policy(run_dir, label=\"本轮训练\"):\n",
" checkpoints = list_checkpoints(run_dir)\n",
" if not checkpoints:\n",
" raise RuntimeError(\n",
" f\"{label}没有发现 checkpoint：{public_path(run_dir)}。\"\n",
" \"请先在本 Notebook 执行训练单元格，不能回退到历史或预训练权重。\"\n",
" )\n",
" path = checkpoints[-1]\n",
" if not path.exists():\n",
" raise FileNotFoundError(f\"训练 checkpoint 不存在：{public_path(path)}\")\n",
" print(f\"评估绑定到{label} checkpoint：\", 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": [
"## 1. 运行控制\n",
"\n",
"本 Notebook 默认执行保护训练和评估。首次运行可先将 `RUN_PROTECTED_TRAIN=False` 做环境检查；正式复现时保持 `True`。训练进度每 100 步刷新一次。\n"
]
},
{
"cell_type": "code",
"metadata": {},
"source": [
"# RUN_CONTROL_CELL\n",
"RUN_SMOKE = False\n",
"RUN_LONG_TRAIN = False\n",
"RUN_EVAL = True\n",
"RUN_PROTECTED_TRAIN = True\n",
"os.environ[\"NOTEBOOK_PROGRESS_EVERY\"] = os.environ.get(\"NOTEBOOK_PROGRESS_EVERY\", \"100\")\n",
"print({\"RUN_SMOKE\": RUN_SMOKE, \"RUN_LONG_TRAIN\": RUN_LONG_TRAIN, \"RUN_EVAL\": RUN_EVAL, \"RUN_PROTECTED_TRAIN\": RUN_PROTECTED_TRAIN})\n"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 2. 保护训练配方\n",
"\n",
"保护初始化权重只用于训练初始化；最终评估只能使用本单元新生成的最终 checkpoint。\n"
]
},
{
"cell_type": "code",
"metadata": {},
"source": [
"DATASET_REPO_ID = os.environ.get(\"DATASET_REPO_ID\", \"datawhale_eai_pnp_language\")\n",
"TRAIN_DATA_ROOT = Path(os.environ.get(\"TRAIN_DATA_ROOT\", DATA_ROOT / \"demo_data_language\"))\n",
"CONFIG_DIR = OUTPUT_ROOT / \"configs\"\n",
"print(\"DATASET_REPO_ID =\", DATASET_REPO_ID)\n",
"print(\"TRAIN_DATA_ROOT =\", public_path(TRAIN_DATA_ROOT))\n"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### 3. 执行 protected parent + weighted-blue 训练\n",
"\n",
"SmolVLA 先训练 parent 5000 steps，再训练 blue 加权续训 1000 steps。两阶段都会写入自己的训练时间摘要。\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 后，本单元会原生训练 SmolVLA protected recipe。\")\n",
"else:\n",
" DATASET_REPO_ID = globals().get(\"DATASET_REPO_ID\", os.environ.get(\"DATASET_REPO_ID\", \"datawhale_eai_pnp_language\"))\n",
" TRAIN_DATA_ROOT = globals().get(\"TRAIN_DATA_ROOT\", Path(os.environ.get(\"TRAIN_DATA_ROOT\", DATA_ROOT / \"demo_data_language\")))\n",
" CONFIG_DIR = OUTPUT_ROOT / \"configs\"\n",
" RUN_ROOT = OUTPUT_ROOT / \"runs\" / \"smolvla_protected_recipe\"\n",
" PARENT_OUTPUT = RUN_ROOT / \"parent_5000\"\n",
" WEIGHTED_OUTPUT = RUN_ROOT / \"weighted_blue2_step1000\"\n",
" parent_config = make_lerobot_train_config(\n",
" \"smolvla\", DATASET_REPO_ID, TRAIN_DATA_ROOT, PARENT_OUTPUT,\n",
" steps=int(os.environ.get(\"SMOLVLA_PARENT_STEPS\", \"5000\")),\n",
" batch_size=int(os.environ.get(\"SMOLVLA_BATCH_SIZE\", \"4\")),\n",
" chunk_size=50,\n",
" n_action_steps=50,\n",
" )\n",
" weighted_config = make_lerobot_train_config(\n",
" \"smolvla\", DATASET_REPO_ID, TRAIN_DATA_ROOT, WEIGHTED_OUTPUT,\n",
" steps=int(os.environ.get(\"SMOLVLA_WEIGHTED_STEPS\", \"1000\")),\n",
" batch_size=int(os.environ.get(\"SMOLVLA_BATCH_SIZE\", \"4\")),\n",
" chunk_size=50,\n",
" n_action_steps=50,\n",
" )\n",
" parent_path = write_json_yaml(CONFIG_DIR / \"smolvla_protected_parent_5000.yaml\", parent_config)\n",
" weighted_path = write_json_yaml(CONFIG_DIR / \"smolvla_protected_weighted_blue2.yaml\", weighted_config)\n",
" train_lerobot_config_in_notebook(parent_path, enabled=True, progress_name=\"SmolVLA protected parent\")\n",
" parent_ckpt = list_checkpoints(PARENT_OUTPUT)[-1]\n",
" old_override = os.environ.get(\"SMOLVLA_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[\"SMOLVLA_PRETRAINED_PATH_OVERRIDE\"] = str(parent_ckpt)\n",
" os.environ[\"NOTEBOOK_FRAME_WEIGHT_MODE\"] = \"blue\"\n",
" os.environ[\"NOTEBOOK_BLUE_WEIGHT\"] = os.environ.get(\"SMOLVLA_BLUE_WEIGHT\", \"2.0\")\n",
" try:\n",
" train_lerobot_config_in_notebook(weighted_path, enabled=True, progress_name=\"SmolVLA protected weighted-blue\")\n",
" finally:\n",
" if old_override is None:\n",
" os.environ.pop(\"SMOLVLA_PRETRAINED_PATH_OVERRIDE\", None)\n",
" else:\n",
" os.environ[\"SMOLVLA_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",
" protected_candidates = list_checkpoints(WEIGHTED_OUTPUT)\n",
" if not protected_candidates:\n",
" raise RuntimeError(\"SmolVLA protected 续训结束后没有 checkpoint，禁止进入评估。\")\n",
" TRAINED_POLICY_PATH = protected_candidates[-1]\n",
" TRAINING_SUMMARY_PATHS = [\n",
" PARENT_OUTPUT / \"training_run_summary.json\",\n",
" WEIGHTED_OUTPUT / \"training_run_summary.json\",\n",
" ]\n",
" print(\"本 Notebook protected checkpoint =\", public_path(TRAINED_POLICY_PATH))\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 4. 训练耗时与 checkpoint\n",
"\n",
"下面只读取本次训练写入的 `training_run_summary.json`，显示开始时间、结束时间、总耗时、步数和输出目录。\n"
]
},
{
"cell_type": "code",
"metadata": {},
"source": [
"summary_paths = [Path(path) for path in globals().get(\"TRAINING_SUMMARY_PATHS\", [])]\n",
"if not summary_paths:\n",
" print(\"尚未执行 protected 训练，暂无训练耗时。\")\n",
"else:\n",
" summaries = []\n",
" for summary_path in summary_paths:\n",
" if summary_path.exists():\n",
" summaries.append(json.loads(summary_path.read_text(encoding=\"utf-8\")))\n",
" else:\n",
" print(\"尚未生成：\", public_path(summary_path))\n",
" if summaries:\n",
" md_table(\n",
" [\"recipe\", \"steps\", \"duration_min\", \"started_at_utc\", \"finished_at_utc\", \"output_dir\"],\n",
" [(item.get(\"recipe\"), item.get(\"steps\"), f\"{float(item.get('duration_s', 0.0))/60.0:.2f}\", item.get(\"started_at_utc\"), item.get(\"finished_at_utc\"), item.get(\"output_dir\")) for item in summaries],\n",
" )\n"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 5. 14 条闭环评估\n",
"\n",
"固定 seed 逐条运行，保存结果 JSONL、成功/失败视频和动作 trace。\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# FULL_14_EVAL_WITH_MEDIA\n",
"FULL_EVAL_EPISODES = 14\n",
"FULL_EVAL_SEEDS = list(range(14))\n",
"FULL_EVAL_ROOT = OUTPUT_ROOT / \"full_eval_14\"\n",
"FULL_EVAL_RESULT = FULL_EVAL_ROOT / \"result.jsonl\"\n",
"if not globals().get(\"RUN_EVAL\", False):\n",
" print(\"未启动评估。请将 RUN_EVAL=True 后执行本单元。\")\n",
"else:\n",
" if globals().get(\"TRAINED_POLICY_PATH\") is None:\n",
" raise RuntimeError(\"本 Notebook 没有当前训练 checkpoint；请先执行本 Notebook 的训练单元格。\")\n",
" FULL_EVAL_POLICY = Path(TRAINED_POLICY_PATH)\n",
" if not FULL_EVAL_POLICY.exists():\n",
" raise FileNotFoundError(f\"当前训练 checkpoint 不存在：{public_path(FULL_EVAL_POLICY)}\")\n",
" print(\"评估模型（来自本 Notebook 当前训练）：\", public_path(FULL_EVAL_POLICY))\n",
" FULL_EVAL_ROWS = run_eval_policy_in_notebook(\n",
" \"smolvla\", FULL_EVAL_POLICY, FULL_EVAL_RESULT,\n",
" episodes=FULL_EVAL_EPISODES, seed_start=0, seed_list=FULL_EVAL_SEEDS,\n",
" render=True, enabled=True, repo_id=DATASET_REPO_ID, dataset_root=TRAIN_DATA_ROOT,\n",
" video_dir=FULL_EVAL_ROOT / \"videos\", trace_dir=FULL_EVAL_ROOT / \"traces\",\n",
" )\n",
"FULL_EVAL_ROWS = run_eval_policy_in_notebook(\n",
" \"smolvla\",\n",
" FULL_EVAL_POLICY,\n",
" FULL_EVAL_RESULT,\n",
" episodes=FULL_EVAL_EPISODES,\n",
" seed_start=0,\n",
" seed_list=FULL_EVAL_SEEDS,\n",
" render=True,\n",
" enabled=RUN_EVAL,\n",
" repo_id=DATASET_REPO_ID,\n",
" dataset_root=TRAIN_DATA_ROOT,\n",
" video_dir=FULL_EVAL_ROOT / \"videos\",\n",
" trace_dir=FULL_EVAL_ROOT / \"traces\",\n",
")\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# FULL_14_MEDIA_REVIEW\n",
"from IPython.display import Video\n",
"import matplotlib.pyplot as plt\n",
"\n",
"def display_full_eval_media(result_path, title=\"smolvla - 14 episode evaluation\"):\n",
" result_path = Path(result_path)\n",
" if not result_path.exists():\n",
" print(\"Evaluation has not run yet:\", public_path(result_path))\n",
" return\n",
" rows = [json.loads(line) for line in result_path.read_text(encoding=\"utf-8\").splitlines() if line.strip()]\n",
" success = sum(bool(row.get(\"physical_success\", row.get(\"success\"))) for row in rows)\n",
" print(f\"{title}: physical_success={success}/{len(rows)}\")\n",
" md_table([[\"metric\", \"value\"]][0], [[\"physical success\", f\"{success}/{len(rows)}\"]])\n",
" for kind in (\"success\", \"failure\"):\n",
" row = next((item for item in rows if bool(item.get(\"physical_success\", item.get(\"success\"))) == (kind == \"success\") and item.get(\"video_path\")), None)\n",
" if row is None:\n",
" print(f\"No {kind} episode was found.\")\n",
" continue\n",
" print(f\"{kind} case: seed={row['seed']}\")\n",
" video_path = Path(row[\"video_path\"])\n",
" if video_path.exists():\n",
" display(Video(filename=str(video_path), embed=False, width=960))\n",
" trace_path = Path(row.get(\"trace_path\", \"\"))\n",
" if not trace_path.exists():\n",
" continue\n",
" trace = json.loads(trace_path.read_text(encoding=\"utf-8\"))\n",
" actions = np.asarray([item.get(\"action\", []) for item in trace], dtype=np.float32)\n",
" if actions.ndim != 2 or not len(actions):\n",
" continue\n",
" fig, ax = plt.subplots(figsize=(12, 4))\n",
" channels = min(actions.shape[1], 7)\n",
" for channel in range(channels):\n",
" label = \"gripper\" if channel == 6 else f\"action[{channel}]\"\n",
" ax.plot(actions[:, channel], label=label, linewidth=1.2)\n",
" ax.set_title(f\"{title} - {kind} sequence, seed={row['seed']}\")\n",
" ax.set_xlabel(\"policy step\")\n",
" ax.set_ylabel(\"normalized action / command\")\n",
" ax.grid(alpha=0.25)\n",
" ax.legend(ncol=4, fontsize=8)\n",
" plt.tight_layout()\n",
" display(fig)\n",
" plt.close(fig)\n",
"\n",
"display_full_eval_media(FULL_EVAL_RESULT)\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 6. 复现检查\n",
"\n",
"评估结果必须来自本 Notebook 当前训练生成的 `TRAINED_POLICY_PATH`。没有当前 checkpoint 时，评估单元会直接报错。\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"
}
},
"nbformat": 4,
"nbformat_minor": 5
}