{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# ACT：保护训练与评估\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' / 'act'))\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.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", " 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 = iter(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", " try:\n", " batch = next(dl_iter)\n", " except StopIteration:\n", " dl_iter = iter(dataloader)\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(\"ACT_EVAL_MAX_ACTION_STEPS\", 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=(int(os.environ[\"ACT_EVAL_N_ACTION_STEPS\"]) if os.environ.get(\"ACT_EVAL_N_ACTION_STEPS\") else None),\n", " act_force_dataset_gripper=env_flag(\"ACT_EVAL_FORCE_DATASET_GRIPPER\", False),\n", " act_clamp_timestamp=env_flag(\"ACT_EVAL_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=os.environ.get(\"ACT_EVAL_EPISODE_TIMESTAMP_OFFSETS\", \"\"),\n", " act_episode_source_flags=os.environ.get(\"ACT_EVAL_EPISODE_SOURCE_FLAGS\", \"\"),\n", " physical_success=env_flag(\"ACT_EVAL_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": [ "ACT_RECIPE = os.environ.get(\"ACT_RECIPE\", \"repair15\").strip().lower()\n", "ACT_IS_REPAIR = ACT_RECIPE in {\"repair15\", \"stable61_to_dagger\", \"nomemleak\"}\n", "ACT_IS_DAGGER = ACT_RECIPE in {\"dagger\", \"dagger_best025\", \"protected\"} and not ACT_IS_REPAIR\n", "ACT_DATASET_DEFAULT = \"datawhale_eai_pnp_act_base72_plus_dagger3x3_rebuild_v1\" if ACT_IS_REPAIR else (\"datawhale_eai_pnp_act_dagger_best025_rebuild_v1\" if ACT_IS_DAGGER else \"datawhale_eai_pnp_language\")\n", "ACT_DATA_ROOT_DEFAULT = DATA_ROOT / (\"act_base72_plus_dagger3x3_rebuild_v1\" if ACT_IS_REPAIR else (\"act_dagger_meta\" if ACT_IS_DAGGER else \"demo_data_language\"))\n", "DATASET_REPO_ID = os.environ.get(\"ACT_DATASET_REPO_ID\", ACT_DATASET_DEFAULT)\n", "TRAIN_DATA_ROOT = Path(os.environ.get(\"ACT_TRAIN_DATA_ROOT\", ACT_DATA_ROOT_DEFAULT))\n", "CONFIG_DIR = OUTPUT_ROOT / \"configs\"\n", "print(\"ACT_RECIPE =\", ACT_RECIPE)\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 ACT 训练\n", "\n", "`repair15` 使用 Notebook 原生训练循环，最终评估只读取本次 `REPAIR_OUTPUT` 中产生的模型。\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# ACT_NATIVE_PROTECTED_TRAINING_CELL\n", "# This is the protected ACT recipe itself. It runs in the Notebook kernel:\n", "# dataset wrappers, weighted sampler, forward/backward, tqdm, checkpoints and metrics JSONL.\n", "def train_act_repair15_in_notebook(repo_id, data_root, output_dir, init_policy_path, steps=2500, batch_size=16):\n", " import copy\n", " import io\n", " import json\n", " import time\n", " from datetime import datetime, timezone\n", " from collections import defaultdict\n", " from pathlib import Path\n", "\n", " import numpy as np\n", " import torch\n", " import torch.nn.functional as F\n", " import torchvision.transforms as T\n", " from PIL import Image\n", " from safetensors.torch import load_file\n", "\n", " from lerobot.common.datasets.factory import resolve_delta_timestamps\n", " from lerobot.common.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata\n", " from lerobot.common.datasets.utils import dataset_to_policy_features\n", " from lerobot.configs.types import FeatureType, PolicyFeature\n", " import sys\n", " lerobot_src = os.environ.get(\"LEROBOT_SRC\", \"/home/aup/jiahang/lerobot_10b7_legacy\")\n", " if Path(lerobot_src).exists():\n", " sys.path.insert(0, lerobot_src)\n", " try:\n", " from lerobot.policies.act.configuration_act import ACTConfig\n", " from lerobot.policies.act.modeling_act import ACTPolicy\n", " except ModuleNotFoundError:\n", " from lerobot.common.policies.act.configuration_act import ACTConfig\n", " from lerobot.common.policies.act.modeling_act import ACTPolicy\n", "\n", " output_dir = Path(output_dir)\n", " data_root = Path(data_root)\n", " init_policy_path = Path(init_policy_path)\n", " output_dir.mkdir(parents=True, exist_ok=True)\n", " if not (init_policy_path / \"model.safetensors\").exists():\n", " raise FileNotFoundError(f\"protected init weight missing: {init_policy_path}\")\n", "\n", " episode_indexes = [0,1,2,3,4,5,6,8,9,10,13,14,15,16,17,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,40,41,42,43,44,45,47,48,50,51,52,53,54,55,56,57,58,59,60,61,62,63,65,67,68,69,71,72,73,74,75,76,77,78,79,80]\n", " timestamp_offsets = {episode: 2.0 for episode in range(49, 81)}\n", "\n", " def patch_image_transform():\n", " import lerobot.common.datasets.lerobot_dataset as ds_mod\n", " import lerobot.common.datasets.utils as ds_utils\n", " def hf_transform_to_torch(items_dict):\n", " to_tensor = T.ToTensor()\n", " for key in items_dict:\n", " first = items_dict[key][0]\n", " if isinstance(first, Image.Image):\n", " items_dict[key] = [to_tensor(img) for img in items_dict[key]]\n", " elif isinstance(first, dict) and (\"bytes\" in first or \"path\" in first):\n", " converted = []\n", " for item in items_dict[key]:\n", " if isinstance(item, dict) and item.get(\"bytes\") is not None:\n", " converted.append(to_tensor(Image.open(io.BytesIO(item[\"bytes\"])).convert(\"RGB\")))\n", " else:\n", " converted.append(item)\n", " items_dict[key] = converted\n", " elif first is not None:\n", " items_dict[key] = [item if isinstance(item, str) else torch.as_tensor(item) for item in items_dict[key]]\n", " return items_dict\n", " ds_utils.hf_transform_to_torch = hf_transform_to_torch\n", " ds_mod.hf_transform_to_torch = hf_transform_to_torch\n", "\n", " class EpisodeSubset(torch.utils.data.Dataset):\n", " def __init__(self, dataset, indexes):\n", " self.dataset = dataset\n", " self.indexes = []\n", " for episode in indexes:\n", " lo = int(dataset.episode_data_index[\"from\"][episode].item())\n", " hi = int(dataset.episode_data_index[\"to\"][episode].item())\n", " self.indexes.extend(range(lo, hi))\n", " if not self.indexes:\n", " raise ValueError(\"protected ACT episode subset is empty\")\n", " def __len__(self):\n", " return len(self.indexes)\n", " def __getitem__(self, index):\n", " return self.dataset[self.indexes[int(index)]]\n", "\n", " class TimestampOffset(torch.utils.data.Dataset):\n", " def __init__(self, dataset, offsets):\n", " self.dataset, self.offsets = dataset, offsets\n", " def __len__(self):\n", " return len(self.dataset)\n", " def __getitem__(self, index):\n", " item = dict(self.dataset[int(index)])\n", " episode = item.get(\"episode_index\", -1)\n", " if isinstance(episode, torch.Tensor):\n", " episode = int(episode.detach().cpu().reshape(-1)[0])\n", " if episode in self.offsets:\n", " timestamp = item[\"timestamp\"]\n", " if isinstance(timestamp, torch.Tensor):\n", " item[\"timestamp\"] = timestamp + timestamp.new_tensor(self.offsets[episode])\n", " else:\n", " item[\"timestamp\"] = np.asarray(timestamp, dtype=np.float32) + np.float32(self.offsets[episode])\n", " return item\n", "\n", " class TimestampAndObjectInit(torch.utils.data.Dataset):\n", " def __init__(self, dataset):\n", " self.dataset = dataset\n", " def __len__(self):\n", " return len(self.dataset)\n", " def __getitem__(self, index):\n", " item = dict(self.dataset[int(index)])\n", " state = item[\"observation.state\"]\n", " timestamp = torch.as_tensor(item[\"timestamp\"], dtype=state.dtype).reshape(-1)\n", " obj_init = torch.as_tensor(item[\"obj_init\"], dtype=state.dtype).reshape(-1)\n", " item[\"observation.state\"] = torch.cat([state, timestamp, obj_init])\n", " return item\n", "\n", " def append_stats(stats, extra_key):\n", " patched = copy.deepcopy(stats)\n", " state = patched[\"observation.state\"]\n", " extra = patched[extra_key]\n", " for key in (\"min\", \"max\", \"mean\", \"std\"):\n", " state[key] = np.concatenate([np.asarray(state[key]).reshape(-1), np.asarray(extra[key]).reshape(-1)])\n", " state[\"std\"] = np.maximum(state[\"std\"], 1e-6)\n", " return patched\n", "\n", " def stats_with_phase_and_object(stats, timestamp_values):\n", " patched = copy.deepcopy(stats)\n", " timestamp_values = np.asarray(timestamp_values, dtype=np.float32).reshape(-1)\n", " timestamp_stats = patched[\"timestamp\"]\n", " for key in (\"min\", \"max\", \"mean\", \"std\"):\n", " if key == \"min\":\n", " value = timestamp_values.min()\n", " elif key == \"max\":\n", " value = timestamp_values.max()\n", " elif key == \"mean\":\n", " value = timestamp_values.mean()\n", " elif key == \"std\":\n", " value = max(float(timestamp_values.std()), 1e-6)\n", " patched[\"timestamp\"][key] = np.asarray([value], dtype=np.float32)\n", " state = patched[\"observation.state\"]\n", " for key in (\"min\", \"max\", \"mean\", \"std\"):\n", " state[key] = np.concatenate([np.asarray(state[key]).reshape(-1), np.asarray(patched[\"timestamp\"][key]).reshape(-1)])\n", " return append_stats(patched, \"obj_init\")\n", "\n", " def episode_ids(dataset):\n", " values = []\n", " for index in range(len(dataset)):\n", " value = dataset[index].get(\"episode_index\", -1)\n", " if isinstance(value, torch.Tensor):\n", " value = int(value.detach().cpu().reshape(-1)[0])\n", " values.append(int(value))\n", " return values\n", "\n", " def first_action(dataset, index):\n", " value = dataset[index][\"action\"]\n", " value = value.detach().cpu().numpy() if isinstance(value, torch.Tensor) else np.asarray(value)\n", " return value[0] if value.ndim == 2 else value\n", "\n", " def build_protected_sampler(dataset):\n", " ids = episode_ids(dataset)\n", " actions = np.stack([first_action(dataset, i) for i in range(len(dataset))]).astype(np.float64)\n", " weights = np.ones(len(dataset), dtype=np.float64)\n", " ranges = defaultdict(list)\n", " for index, episode in enumerate(ids):\n", " ranges[episode].append(index)\n", " for episode, indexes in ranges.items():\n", " if episode in timestamp_offsets:\n", " weights[indexes] *= 0.25\n", " weights[indexes[:250]] *= 6.0\n", " gripper = actions[:, 6]\n", " weights[gripper > 0.5] *= 6.0\n", " transitions = np.flatnonzero(np.abs(np.diff(gripper, prepend=gripper[0])) > 0.5)\n", " for transition in transitions:\n", " lo, hi = max(int(transition) - 24, 0), min(int(transition) + 25, len(weights))\n", " weights[lo:hi] *= 25.0\n", " motion = np.zeros(len(actions), dtype=np.float64)\n", " if len(actions) > 1:\n", " motion[1:] = np.linalg.norm(actions[1:, :6] - actions[:-1, :6], axis=1)\n", " scale = np.percentile(motion, 95) if np.any(motion > 0) else 0.0\n", " if scale > 1e-9:\n", " weights *= 1.0 + 8.0 * np.clip(motion / scale, 0.0, 1.0)\n", " sampler = torch.utils.data.WeightedRandomSampler(torch.as_tensor(weights, dtype=torch.double), len(weights), replacement=True)\n", " info = {\"mode\":\"trajectory\", \"frames\":len(dataset), \"episodes\":len(ranges), \"weight_min\":float(weights.min()), \"weight_max\":float(weights.max()), \"transition_count\":int(len(transitions)), \"offset_episodes\":sorted(timestamp_offsets)}\n", " print(\"protected sampler =\", json.dumps(info, ensure_ascii=False))\n", " return sampler, info\n", "\n", " patch_image_transform()\n", " metadata = LeRobotDatasetMetadata(repo_id, root=data_root)\n", " features = dataset_to_policy_features(metadata.features)\n", " output_features = {key: feature for key, feature in features.items() if feature.type is FeatureType.ACTION}\n", " input_features = {key: feature for key, feature in features.items() if feature.type is not FeatureType.ACTION}\n", " input_features.pop(\"observation.wrist_image\", None)\n", " base_state = int(input_features[\"observation.state\"].shape[0])\n", " input_features[\"observation.state\"] = PolicyFeature(type=FeatureType.STATE, shape=(base_state + 7,))\n", " config = ACTConfig(\n", " input_features=input_features,\n", " output_features=output_features,\n", " chunk_size=20,\n", " n_action_steps=10,\n", " use_vae=False,\n", " dropout=0.0,\n", " kl_weight=10.0,\n", " optimizer_lr=1e-5,\n", " optimizer_weight_decay=1e-4,\n", " optimizer_lr_backbone=1e-5,\n", " device=\"cuda\",\n", " )\n", " delta_timestamps = resolve_delta_timestamps(config, metadata)\n", " dataset = LeRobotDataset(repo_id, delta_timestamps=delta_timestamps, root=data_root)\n", " dataset = EpisodeSubset(dataset, episode_indexes)\n", " dataset = TimestampOffset(dataset, timestamp_offsets)\n", " timestamp_values = [float(torch.as_tensor(dataset[i][\"timestamp\"]).reshape(-1)[0]) for i in range(len(dataset))]\n", " stats = stats_with_phase_and_object(metadata.stats, timestamp_values)\n", " dataset = TimestampAndObjectInit(dataset)\n", " device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n", " policy = ACTPolicy(config, dataset_stats=stats).to(device)\n", " state = load_file(str(init_policy_path / \"model.safetensors\"), device=\"cpu\")\n", " missing, unexpected = policy.load_state_dict(state, strict=False)\n", " if missing or unexpected:\n", " raise RuntimeError(f\"protected init mismatch: missing={missing}, unexpected={unexpected}\")\n", " policy.train()\n", " optimizer = torch.optim.Adam(policy.parameters(), lr=1e-5)\n", " sampler, sampler_info = build_protected_sampler(dataset)\n", " loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, sampler=sampler, num_workers=0, pin_memory=device.type == \"cuda\", drop_last=True)\n", " iterator = iter(loader)\n", " metrics_path = output_dir / \"notebook_native_train_metrics.jsonl\"\n", " manifest = {\n", " \"recipe\": \"stable61_to_dagger_low_lr_step1500\",\n", " \"native_notebook_kernel\": True,\n", " \"repo_id\": repo_id,\n", " \"data_root\": str(data_root),\n", " \"init_policy_path\": str(init_policy_path),\n", " \"steps\": int(steps),\n", " \"batch_size\": int(batch_size),\n", " \"chunk_size\": 20,\n", " \"n_action_steps\": 10,\n", " \"use_vae\": False,\n", " \"optimizer\": \"Adam(lr=1e-5)\",\n", " \"state_features\": \"proprio + timestamp + obj_init\",\n", " \"sampler\": sampler_info,\n", " }\n", " (output_dir / \"notebook_native_train_manifest.json\").write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding=\"utf-8\")\n", " print(json.dumps(manifest, ensure_ascii=False, indent=2))\n", " run_started_at = datetime.now(timezone.utc).isoformat()\n", " start = time.perf_counter()\n", " last = None\n", " for step in notebook_progress(range(1, int(steps) + 1), desc=\"ACT protected native Notebook train\", total=int(steps)):\n", " try:\n", " batch = next(iterator)\n", " except StopIteration:\n", " iterator = iter(loader)\n", " batch = next(iterator)\n", " raw_action = batch[\"action\"]\n", " batch = {key: (value.to(device, non_blocking=True) if isinstance(value, torch.Tensor) else value) for key, value in batch.items()}\n", " raw_action = raw_action.to(device, non_blocking=True)\n", " optimizer.zero_grad(set_to_none=True)\n", " normalized = policy.normalize_inputs(batch)\n", " if policy.config.image_features:\n", " normalized = dict(normalized)\n", " normalized[\"observation.images\"] = [normalized[key] for key in policy.config.image_features]\n", " normalized = policy.normalize_targets(normalized)\n", " actions_hat, _ = policy.model(normalized)\n", " valid = ~normalized[\"action_is_pad\"]\n", " l1 = (F.l1_loss(normalized[\"action\"], actions_hat, reduction=\"none\") * valid.unsqueeze(-1)).mean()\n", " threshold = torch.tensor(0.5, dtype=torch.float32, device=device)\n", " action_buffer = policy.normalize_targets.buffer_action\n", " threshold_norm = (threshold - action_buffer[\"mean\"][..., 6].to(device)) / (action_buffer[\"std\"][..., 6].to(device) + 1e-8)\n", " target_bin = (raw_action[..., 6] > 0.5).to(dtype=actions_hat.dtype)\n", " logits = (actions_hat[..., 6] - threshold_norm) * 4.0\n", " bce_each = F.binary_cross_entropy_with_logits(logits, target_bin, reduction=\"none\")\n", " valid_float = valid.to(dtype=bce_each.dtype)\n", " bce = (bce_each * valid_float).sum() / valid_float.sum().clamp_min(1.0)\n", " loss = l1 + 0.5 * bce\n", " loss.backward()\n", " optimizer.step()\n", " if step == 1 or step % 20 == 0 or step == int(steps):\n", " last = {\"step\":step, \"loss\":float(loss.detach().cpu()), \"l1_loss\":float(l1.detach().cpu()), \"gripper_bce_loss\":float(bce.detach().cpu()), \"elapsed_s\":float(time.perf_counter()-start)}\n", " with metrics_path.open(\"a\", encoding=\"utf-8\") as handle:\n", " handle.write(json.dumps(last, ensure_ascii=False) + \"\\n\")\n", " print(json.dumps(last, ensure_ascii=False))\n", " if step % 500 == 0 or step == int(steps):\n", " checkpoint = output_dir / f\"step_{step}\"\n", " policy.save_pretrained(checkpoint)\n", " print(\"saved native Notebook checkpoint:\", checkpoint)\n", " policy.save_pretrained(output_dir)\n", " run_finished_at = datetime.now(timezone.utc).isoformat()\n", " duration_s = float(time.perf_counter() - start)\n", " run_summary = {\n", " \"recipe\": \"act_protected_repair15\",\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(steps),\n", " \"batch_size\": int(batch_size),\n", " \"output_dir\": str(output_dir),\n", " \"metrics_path\": str(metrics_path),\n", " \"checkpoint_path\": str(output_dir),\n", " \"last_metrics\": last,\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(\"ACT 训练耗时 =\", f\"{duration_s / 60.0:.2f} 分钟\")\n", " print(\"ACT 训练摘要 =\", public_path(summary_path))\n", " return {\"output_dir\": output_dir, \"metrics_path\": metrics_path, \"manifest\": output_dir / \"notebook_native_train_manifest.json\", \"summary_path\": summary_path, \"last\": last}\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 后，本单元会原生训练 ACT protected recipe；设置 ACT_RECIPE=repair15 可按当前保护协议准备 chunk20/n10 配置。\")\n", "else:\n", " DATASET_REPO_ID = globals().get(\"DATASET_REPO_ID\", os.environ.get(\"ACT_DATASET_REPO_ID\", \"datawhale_eai_pnp_language\"))\n", " TRAIN_DATA_ROOT = globals().get(\"TRAIN_DATA_ROOT\", Path(os.environ.get(\"ACT_TRAIN_DATA_ROOT\", DATA_ROOT / \"demo_data_language\")))\n", " CONFIG_DIR = OUTPUT_ROOT / \"configs\"\n", " RUN_ROOT = OUTPUT_ROOT / \"runs\" / \"act_protected_current_recipe\"\n", " is_repair = os.environ.get(\"ACT_RECIPE\", \"stable61\").strip().lower() in {\"repair15\", \"stable61_to_dagger\", \"nomemleak\"}\n", " if is_repair:\n", " REPAIR_OUTPUT = RUN_ROOT / \"repair15_continuation\"\n", " initial_path = Path(os.environ.get(\"ACT_INITIAL_PROTECTED_PATH\", str(MODEL_ROOT / \"act_fallback_best_stable61_step2500\")))\n", " repair_data_root = Path(os.environ.get(\"ACT_REPAIR_TRAIN_DATA_ROOT\", TRAIN_DATA_ROOT))\n", " print(\"本单元在 Notebook kernel 内原生执行 repair15：不调用外部训练脚本。\")\n", " native_result = train_act_repair15_in_notebook(\n", " repo_id=os.environ.get(\"ACT_DATASET_REPO_ID\", DATASET_REPO_ID),\n", " data_root=repair_data_root,\n", " output_dir=REPAIR_OUTPUT,\n", " init_policy_path=initial_path,\n", " steps=int(os.environ.get(\"ACT_REPAIR_STEPS\", \"2500\")),\n", " batch_size=int(os.environ.get(\"ACT_BATCH_SIZE\", \"16\")),\n", " )\n", " ACT_REPAIR_NATIVE_OUTPUT = Path(native_result[\"output_dir\"])\n", " ACT_REPAIR_NATIVE_METRICS = Path(native_result[\"metrics_path\"])\n", " print(\"Notebook native training result:\", public_path(ACT_REPAIR_NATIVE_OUTPUT))\n", " if not (ACT_REPAIR_NATIVE_OUTPUT / \"model.safetensors\").exists():\n", " raise RuntimeError(\"ACT protected native 训练没有产生最终 model.safetensors。\")\n", " else:\n", " STABLE_OUTPUT = RUN_ROOT / \"stable61_full\"\n", " stable_config = make_lerobot_train_config(\n", " \"act\",\n", " DATASET_REPO_ID,\n", " Path(os.environ.get(\"ACT_STABLE61_TRAIN_DATA_ROOT\", TRAIN_DATA_ROOT)),\n", " STABLE_OUTPUT,\n", " steps=int(os.environ.get(\"ACT_STEPS\", \"5000\")),\n", " batch_size=int(os.environ.get(\"ACT_BATCH_SIZE\", \"8\")),\n", " chunk_size=50,\n", " n_action_steps=50,\n", " )\n", " stable_path = write_json_yaml(CONFIG_DIR / \"act_protected_current_stable61.yaml\", stable_config)\n", " train_lerobot_config_in_notebook(stable_path, enabled=True, progress_name=\"ACT protected-current stable61\")\n", " print(\"protected-current candidate checkpoints:\")\n", " stable_candidates = list_checkpoints(STABLE_OUTPUT)\n", " if not stable_candidates:\n", " raise RuntimeError(\"ACT protected stable61 训练没有产生 checkpoint。\")\n", " ACT_REPAIR_NATIVE_OUTPUT = stable_candidates[-1]\n", "\n", "if is_repair:\n", " TRAINED_POLICY_PATH = ACT_REPAIR_NATIVE_OUTPUT\n", " TRAINING_SUMMARY_PATHS = [REPAIR_OUTPUT / \"training_run_summary.json\"]\n", "else:\n", " TRAINED_POLICY_PATH = ACT_REPAIR_NATIVE_OUTPUT\n", " TRAINING_SUMMARY_PATHS = [STABLE_OUTPUT / \"training_run_summary.json\"]\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(1030, 1044))\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", " \"act\", FULL_EVAL_POLICY, FULL_EVAL_RESULT,\n", " episodes=FULL_EVAL_EPISODES, seed_start=1030, 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", " \"act\",\n", " FULL_EVAL_POLICY,\n", " FULL_EVAL_RESULT,\n", " episodes=FULL_EVAL_EPISODES,\n", " seed_start=1030,\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=\"act - 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 }