| """Pinned official Diffusion Policy components, adapted to the SO101 data contract. |
| |
| The upstream encoder, ResNet-18 getter, crop randomizer, Mish U-Net and EMA are |
| vendored without edits. This wrapper deliberately keeps Issue 18's *current-t* |
| action convention, train-only min/max statistics and full-window dataset. It is |
| not an unmodified official task/runner or a pretrained real-robot policy. |
| """ |
| from __future__ import annotations |
|
|
| import copy |
| import hashlib |
| import importlib.metadata |
| import inspect |
| import json |
| import os |
| from pathlib import Path |
| import sys |
| import types |
|
|
| import numpy as np |
| import torch |
| from torch import Tensor, nn |
|
|
| HERE = Path(__file__).resolve().parent |
| REPO_ROOT = HERE.parents[1] |
| VENDOR = HERE / "vendor" |
| UPSTREAM_COMMIT = "5ba07ac6661db573af695b419a7947ecb704690f" |
| FORMAT = "so101_official_dp.visual_policy.v1" |
| NORMALIZER_KEYS = ("action_center", "action_span", "state_center", "state_span") |
| STATS_KEYS = ("action_min", "action_max", "state_min", "state_max") |
| DEFAULTS = { |
| "history": 2, "horizon": 16, "image_size": 96, |
| "crop_shape": None, "resize_shape": None, "random_crop": False, "imagenet_norm": True, |
| "cameras": ["front", "gripper"], "context_dim": 2060, |
| "execution_start": 0, "execute_steps": 8, "control_hz": 10, |
| "down_dims": [512, 1024, 2048], "diffusion_step_embed_dim": 128, |
| "kernel_size": 5, "n_groups": 8, "cond_predict_scale": True, |
| "num_train_timesteps": 100, "beta_schedule": "squaredcos_cap_v2", |
| "prediction_type": "epsilon", "clip_sample": True, |
| "variance_type": "fixed_small", "default_sampler": "ddpm", |
| "default_inference_steps": 100, |
| "total_updates": 100000, "batch_size": 64, "microbatch_size": 8, |
| "learning_rate": 1e-4, "adam_betas": [0.95, 0.999], |
| "weight_decay": 1e-6, "warmup_steps": 500, "gradient_clip": None, |
| "ema_power": 0.75, "ema_max_value": 0.9999, |
| "checkpoint_every": 10000, "log_every": 100, "num_workers": 0, |
| "validation_batch_size": 8, "validation_sampler": "ddim", |
| "validation_inference_steps": 16, "validation_noise_seed": 7041, |
| "no_real_robot": True, "test_only": False, |
| } |
|
|
|
|
| def canonical_config(config: dict) -> dict: |
| cfg = copy.deepcopy(DEFAULTS) |
| cfg.update(copy.deepcopy(config)) |
| required = {"history": 2, "horizon": 16, "image_size": 96, |
| "crop_shape": None, "resize_shape": None, "random_crop": False, "imagenet_norm": True, |
| "cameras": ["front", "gripper"], "context_dim": 2060, |
| "execution_start": 0, "execute_steps": 8, "control_hz": 10, |
| "no_real_robot": True, "num_train_timesteps": 100, |
| "beta_schedule": "squaredcos_cap_v2", "prediction_type": "epsilon", |
| "clip_sample": True, "variance_type": "fixed_small"} |
| for key, value in required.items(): |
| if cfg[key] != value: |
| raise ValueError(f"SO101 official-component contract requires {key}={value!r}") |
| if not cfg["test_only"]: |
| for key in ("down_dims", "diffusion_step_embed_dim", "kernel_size", "n_groups", |
| "cond_predict_scale"): |
| if cfg[key] != DEFAULTS[key]: |
| raise ValueError(f"Production architecture requires {key}={DEFAULTS[key]!r}") |
| if len(cfg["down_dims"]) != 3 or any(int(x) <= 0 or int(x) % cfg["n_groups"] for x in cfg["down_dims"]): |
| raise ValueError("Use three positive group-divisible U-Net channel widths") |
| for key in ("batch_size", "microbatch_size", "total_updates", "checkpoint_every", |
| "log_every", "validation_batch_size", "validation_inference_steps"): |
| if isinstance(cfg[key], bool) or int(cfg[key]) != cfg[key] or cfg[key] <= 0: |
| raise ValueError(f"{key} must be a positive integer") |
| if cfg["batch_size"] % cfg["microbatch_size"]: |
| raise ValueError("microbatch_size must divide effective batch_size") |
| if cfg["num_workers"] != 0: |
| raise ValueError("Exact-resume batch stream uses num_workers=0") |
| if cfg["warmup_steps"] < 0 or cfg["learning_rate"] <= 0: |
| raise ValueError("Invalid optimizer schedule") |
| if cfg["gradient_clip"] is not None and cfg["gradient_clip"] <= 0: |
| raise ValueError("gradient_clip must be positive or None") |
| return cfg |
|
|
|
|
| def canonical_stats(stats: dict) -> dict: |
| out = copy.deepcopy(stats) |
| for key in STATS_KEYS: |
| value = np.asarray(stats[key], dtype=np.float32) |
| if value.shape != (6,) or not np.isfinite(value).all(): |
| raise ValueError(f"{key} requires six finite values") |
| out[key] = value.tolist() |
| for prefix in ("action", "state"): |
| if np.any(np.asarray(out[prefix + "_max"]) < np.asarray(out[prefix + "_min"])): |
| raise ValueError(f"Invalid {prefix} min/max") |
| return out |
|
|
|
|
| def sha256(path: str | Path) -> str: |
| h = hashlib.sha256() |
| with Path(path).open("rb") as stream: |
| for part in iter(lambda: stream.read(1024 * 1024), b""): |
| h.update(part) |
| return h.hexdigest() |
|
|
|
|
| def _official_imports(): |
| |
| |
| package_path = VENDOR / "diffusion_policy" |
| existing = sys.modules.get("diffusion_policy") |
| if existing is None: |
| package = types.ModuleType("diffusion_policy") |
| package.__path__ = [str(package_path)] |
| package.__package__ = "diffusion_policy" |
| sys.modules["diffusion_policy"] = package |
| elif [Path(p).resolve() for p in getattr(existing, "__path__", [])] != [package_path.resolve()]: |
| raise ImportError("Another diffusion_policy package is loaded; use an isolated process") |
| from diffusion_policy.model.vision.multi_image_obs_encoder import MultiImageObsEncoder |
| from diffusion_policy.model.vision.model_getter import get_resnet |
| from diffusion_policy.model.diffusion.conditional_unet1d import ConditionalUnet1D |
| from diffusion_policy.model.diffusion.ema_model import EMAModel |
| for component in (MultiImageObsEncoder, get_resnet, ConditionalUnet1D, EMAModel): |
| if not Path(inspect.getfile(component)).resolve().is_relative_to(VENDOR.resolve()): |
| raise ImportError("An official component resolved outside the pinned vendor directory") |
| return MultiImageObsEncoder, get_resnet, ConditionalUnet1D, EMAModel |
|
|
|
|
| def official_ema(model: nn.Module, cfg: dict): |
| """Upstream EMA owns its input; deepcopy is REQUIRED, not done upstream.""" |
| *_, EMA = _official_imports() |
| return EMA(copy.deepcopy(model), update_after_step=0, inv_gamma=1.0, |
| power=float(cfg["ema_power"]), min_value=0.0, |
| max_value=float(cfg["ema_max_value"])) |
|
|
|
|
| class OfficialDP(nn.Module): |
| kind = "teacher" |
|
|
| def __init__(self, cfg: dict, stats: dict): |
| super().__init__() |
| self.cfg = canonical_config(cfg) |
| self.stats = canonical_stats(stats) |
| for prefix in ("action", "state"): |
| low = torch.tensor(self.stats[prefix + "_min"], dtype=torch.float32) |
| high = torch.tensor(self.stats[prefix + "_max"], dtype=torch.float32) |
| self.register_buffer(prefix + "_center", (high + low) / 2) |
| self.register_buffer(prefix + "_span", (high - low).clamp_min(1e-3)) |
| Encoder, get_resnet, Unet, _ = _official_imports() |
| shape_meta = {"obs": { |
| "front": {"shape": [3, 96, 96], "type": "rgb"}, |
| "gripper": {"shape": [3, 96, 96], "type": "rgb"}, |
| "q": {"shape": [6], "type": "low_dim"}}, |
| "action": {"shape": [6]}} |
| self.encoder = Encoder( |
| shape_meta=shape_meta, rgb_model=get_resnet("resnet18", weights=None), |
| resize_shape=None, crop_shape=None, random_crop=False, |
| use_group_norm=True, share_rgb_model=False, imagenet_norm=True) |
| self.denoiser = Unet( |
| input_dim=6, local_cond_dim=None, global_cond_dim=2060, |
| diffusion_step_embed_dim=self.cfg["diffusion_step_embed_dim"], |
| down_dims=self.cfg["down_dims"], kernel_size=self.cfg["kernel_size"], |
| n_groups=self.cfg["n_groups"], cond_predict_scale=self.cfg["cond_predict_scale"]) |
| from diffusers.schedulers.scheduling_ddpm import DDPMScheduler |
| from diffusers.schedulers.scheduling_ddim import DDIMScheduler |
| schedule = dict(num_train_timesteps=100, beta_start=0.0001, beta_end=0.02, |
| beta_schedule="squaredcos_cap_v2", clip_sample=True, |
| prediction_type="epsilon", timestep_spacing="leading", |
| steps_offset=0) |
| self.noise_scheduler = DDPMScheduler(**schedule, variance_type="fixed_small") |
| self.ddim_scheduler = DDIMScheduler(**schedule, set_alpha_to_one=True) |
|
|
| def normalize_action(self, action: Tensor) -> Tensor: |
| return 2 * (action - self.action_center) / self.action_span |
|
|
| def normalize_state(self, q: Tensor) -> Tensor: |
| return 2 * (q - self.state_center) / self.state_span |
|
|
| def decode(self, normalized_action: Tensor) -> Tensor: |
| """Linear inverse only; the evaluator logs both explicit clamp stages.""" |
| return normalized_action * self.action_span / 2 + self.action_center |
|
|
| def encode(self, images: Tensor, q: Tensor) -> Tensor: |
| if images.ndim != 6 or tuple(images.shape[1:]) != (2, 2, 3, 96, 96): |
| raise ValueError("Expected [B,previous/current,front/gripper,RGB,96,96]") |
| if images.dtype != torch.uint8 or tuple(q.shape) != (images.shape[0], 2, 6): |
| raise ValueError("Images must be uint8 and q must be [B,2,6] environment coordinates") |
| if not q.is_floating_point(): |
| raise ValueError("q must contain floating-point environment-normalized joint states") |
| b = images.shape[0] |
| rgb = images.float() / 255.0 |
| state = self.normalize_state(q) |
| obs = {"front": rgb[:, :, 0].reshape(b * 2, 3, 96, 96), |
| "gripper": rgb[:, :, 1].reshape(b * 2, 3, 96, 96), |
| "q": state.reshape(b * 2, 6)} |
| return self.encoder(obs).reshape(b, 2060) |
|
|
| def head_predict(self, context: Tensor, noise: Tensor | None = None, *, |
| sampler: str = "ddpm", inference_steps: int = 100, |
| generator: torch.Generator | None = None) -> Tensor: |
| """Return all 16 actions starting at current t, BEFORE deployment clamps. |
| |
| DDPM draws new Gaussian noise at every nonterminal reverse step. Supplying |
| only initial noise does not make DDPM deterministic: pass a seeded generator |
| as well. DDIM uses eta=0. Both use diffusers' official update equations. |
| """ |
| if context.ndim != 2 or context.shape[1] != 2060: |
| raise ValueError("Expected context [B,2060]") |
| if sampler not in ("ddpm", "ddim") or not 1 <= int(inference_steps) <= 100: |
| raise ValueError("Use ddpm/ddim with 1..100 inference steps") |
| shape = (context.shape[0], 16, 6) |
| if noise is None: |
| sample = torch.randn(shape, dtype=context.dtype, device=context.device, generator=generator) |
| else: |
| if tuple(noise.shape) != shape or noise.device != context.device or noise.dtype != context.dtype: |
| raise ValueError("Initial noise must match prediction shape/device/dtype") |
| sample = noise.clone() |
| scheduler = self.noise_scheduler if sampler == "ddpm" else self.ddim_scheduler |
| scheduler.set_timesteps(int(inference_steps), device=context.device) |
| finite = torch.isfinite(sample).all() & torch.isfinite(context).all() |
| for t in scheduler.timesteps: |
| epsilon = self.denoiser(sample, t, global_cond=context) |
| finite = finite & torch.isfinite(epsilon).all() |
| if sampler == "ddpm": |
| sample = scheduler.step(epsilon, t, sample, generator=generator).prev_sample |
| else: |
| sample = scheduler.step(epsilon, t, sample, eta=0.0, generator=generator).prev_sample |
| finite = finite & torch.isfinite(sample).all() |
| if not bool(finite): |
| raise FloatingPointError("Non-finite intermediate diffusion sample") |
| return sample |
|
|
| def predict(self, images: Tensor, q: Tensor, noise: Tensor | None = None, *, |
| sampler: str = "ddpm", inference_steps: int = 100, |
| generator: torch.Generator | None = None) -> Tensor: |
| raw = self.head_predict(self.encode(images, q), noise, sampler=sampler, |
| inference_steps=inference_steps, generator=generator) |
| return self.decode(raw.clamp(-1, 1)).clamp(-1, 1) |
|
|
| def training_loss(self, batch: dict, *, generator: torch.Generator | None = None) -> Tensor: |
| target = self.normalize_action(batch["action"]) |
| if target.shape[1:] != (16, 6) or batch["mask"].shape != target.shape[:2]: |
| raise ValueError("Training target must be a masked [B,16,6] current-time chunk") |
| noise = torch.randn(target.shape, device=target.device, dtype=target.dtype, generator=generator) |
| timestep = torch.randint(0, 100, (target.shape[0],), device=target.device, generator=generator) |
| noisy = self.noise_scheduler.add_noise(target, noise, timestep) |
| prediction = self.denoiser(noisy, timestep, global_cond=self.encode(batch["images"], batch["q"])) |
| mask = batch["mask"].unsqueeze(-1).expand_as(target) |
| count = mask.sum() |
| if count <= 0: |
| raise ValueError("Empty action mask") |
| return ((prediction - noise).square() * mask).sum() / count |
|
|
|
|
| def make_dataset(root: str | Path, cfg: dict, split: str): |
| """Use the ORIGINAL manifest configuration; new optimizer keys are unrelated. |
| |
| No new demonstrations, boundary padding, statistics fit or test examples are |
| introduced. Requested semantic fields must agree with the recorded dataset. |
| """ |
| if str(REPO_ROOT) not in sys.path: |
| sys.path.insert(0, str(REPO_ROOT)) |
| from experiments.so101_pbit_sim.data import SimDataset |
| manifest = json.loads((Path(root) / "manifest.json").read_text(encoding="utf-8")) |
| original = manifest["config"] |
| shared = ("task_id", "cameras", "camera_layout", "image_size", "horizon", "history", |
| "execution_start", "control_hz", "success_contract", "action", "no_real_robot", |
| "train_seeds", "validation_seeds", "development_seeds") |
| checked = canonical_config(cfg) |
| for key in shared: |
| if key in checked and checked[key] != original.get(key): |
| raise ValueError(f"Original dataset semantic contract differs at {key}") |
| if "test_seeds" in checked and checked["test_seeds"] != original["test_seeds"]: |
| if checked.get("historical_test_seeds") != original["test_seeds"]: |
| raise ValueError("Declare the complete historical test set before choosing a fresh holdout") |
| used = set().union(*(set(original[k]) for k in |
| ("train_seeds", "validation_seeds", "test_seeds", "development_seeds"))) |
| if used.intersection(checked["test_seeds"]) or len(set(checked["test_seeds"])) != len(checked["test_seeds"]): |
| raise ValueError("Fresh holdout overlaps a prior split or contains duplicate seeds") |
| data = SimDataset(root, original, split) |
| if data.stats.get("fit_seeds") != original["train_seeds"]: |
| raise ValueError("Normalization was not fit exclusively on training episodes") |
| return data |
|
|
|
|
| def source_record() -> dict: |
| manifest_path = VENDOR / "source_manifest.json" |
| manifest = json.loads(manifest_path.read_text(encoding="utf-8")) |
| if manifest.get("commit") != UPSTREAM_COMMIT: |
| raise ValueError("Unexpected upstream source revision") |
| for row in manifest["files"]: |
| path = VENDOR / row["path"] |
| if sha256(path) != row["sha256"]: |
| raise ValueError(f"Pinned vendor file changed: {row['path']}") |
| return {"upstream_commit": UPSTREAM_COMMIT, |
| "vendor_manifest_sha256": sha256(manifest_path), |
| "git_commit_declared_by_launcher": os.environ.get("SOURCE_SHA"), |
| "file_sha256": {p.relative_to(REPO_ROOT).as_posix(): sha256(p) |
| for p in (HERE / "model.py", HERE / "train.py") if p.is_file()}, |
| "versions": {name: importlib.metadata.version(name) |
| for name in ("torch", "torchvision", "diffusers", "einops")}} |
|
|
|
|
| def policy_payload(model: OfficialDP, *, step: int, manifest_sha: str, |
| extra: dict | None = None) -> dict: |
| return {"format": FORMAT, "kind": "teacher", "config": model.cfg, |
| "stats": model.stats, "state_dict": model.state_dict(), "step": int(step), |
| "dataset_manifest_sha256": manifest_sha, "source": source_record(), |
| "action_time_origin": "current_t", "execution_start": 0, |
| "clamp_contract": "data_clamp_then_linear_decode_then_environment_clamp", |
| "weights": "EMA", "extra": copy.deepcopy(extra or {})} |
|
|
|
|
| def atomic_torch_save(payload: dict, path: str | Path) -> dict: |
| path = Path(path) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| temporary = path.with_name(path.name + ".tmp") |
| torch.save(payload, temporary) |
| temporary.replace(path) |
| return {"path": path.name, "sha256": sha256(path), "bytes": path.stat().st_size, |
| "step": payload.get("step")} |
|
|
|
|
| def load_policy(path: str | Path, device: str | torch.device = "cpu") -> tuple[OfficialDP, dict]: |
| ck = torch.load(Path(path), map_location="cpu", weights_only=True) |
| if ck.get("format") != FORMAT or ck.get("action_time_origin") != "current_t": |
| raise ValueError("Not an official-component SO101 EMA inference checkpoint") |
| current_source = source_record() |
| for key in ("upstream_commit", "vendor_manifest_sha256"): |
| if ck.get("source", {}).get(key) != current_source[key]: |
| raise ValueError(f"Checkpoint official-component identity differs: {key}") |
| model = OfficialDP(ck["config"], ck["stats"]) |
| expected = {name: getattr(model, name).clone() for name in NORMALIZER_KEYS} |
| model.load_state_dict(ck["state_dict"], strict=True) |
| for name, value in expected.items(): |
| if not torch.equal(value, getattr(model, name)): |
| raise ValueError(f"Checkpoint normalizer contradicts train statistics: {name}") |
| for name, value in model.state_dict().items(): |
| if not torch.isfinite(value).all(): |
| raise ValueError(f"Non-finite checkpoint tensor: {name}") |
| return model.to(device).eval(), ck |
|
|
