From 88df3607b7f3a040468e1bffd585d63da99e9490 Mon Sep 17 00:00:00 2001 From: Aidan Gao Date: Mon, 28 Sep 2026 14:57:26 -0400 Subject: [PATCH 01/20] eva: joint-space mode for the ABC YAM action space; robot image history ABC's station and its simulator (amazon-far/abc abc_sim) act in 14-D absolute joint positions + gripper, [left j1..j6, left grip, right j1..j6, right grip]. ABC zarrs already carry {left,right}.{obs,cmd}_{joints,gripper}, so a joint-space policy can train on real and sim ABC episodes with one schema and be stepped in the sim without IK. - Eva transform/keymap mode "joints": observations.state.joint_positions (14) and actions_joints (100, 14), the 45 raw 30 Hz steps resampled to the model horizon like the cartesian modes; no revert. - Embodiment.get_keymap(image_history_gap_s=...): the "_hist" front camera Human already had, so a robot domain can share an RDT trunk with trunk.image_history: 2. - layout_metrics(..., ac_key): actions_joints* get joint metrics (abs error in degrees, gripper MSE). A 14-D chunk otherwise reads as cartesian ypr. Co-Authored-By: Claude Opus 5.5 (1M context) --- egomimic/eval/action_metrics.py | 36 ++++++++- egomimic/eval/eval_hpt.py | 4 +- egomimic/eval/eval_pi.py | 4 +- egomimic/rldb/embodiment/embodiment.py | 21 ++++- egomimic/rldb/embodiment/eva.py | 84 +++++++++++++++++++- tests/unit/test_eva_joint_mode.py | 106 +++++++++++++++++++++++++ 6 files changed, 248 insertions(+), 7 deletions(-) create mode 100644 tests/unit/test_eva_joint_mode.py diff --git a/egomimic/eval/action_metrics.py b/egomimic/eval/action_metrics.py index 9c1d099562..f70346ee26 100644 --- a/egomimic/eval/action_metrics.py +++ b/egomimic/eval/action_metrics.py @@ -202,8 +202,40 @@ def _kp_l2(p, g): return metrics -def layout_metrics(pred: torch.Tensor, gt: torch.Tensor, prefix: str) -> dict: - """Dispatch on the last-dim width: cartesian, keypoint, or ``{}``.""" +# Bimanual joint layout (ABC YAM / abc_sim): [left j1..j6, left grip, +# right j1..j6, right grip]. +JOINT_LAYOUT_14 = {"joints": (0, 1, 2, 3, 4, 5, 7, 8, 9, 10, 11, 12), "grip": (6, 13)} + + +def joint_metrics(pred: torch.Tensor, gt: torch.Tensor, prefix: str) -> dict: + """Metrics for a bimanual joint-space chunk ``(B, T, 14)``; ``{}`` for + other widths. Joint errors are in degrees (mean absolute), the gripper in + its [0, 1] command units.""" + pred_cpu, gt_cpu = pred.detach().cpu().float(), gt.detach().cpu().float() + if pred_cpu.shape[-1] != 14: + return {} + j = list(JOINT_LAYOUT_14["joints"]) + g = list(JOINT_LAYOUT_14["grip"]) + err = (pred_cpu[..., j] - gt_cpu[..., j]).abs() * _RAD2DEG + return { + f"{prefix}_joint_abs_err_deg_avg": err.mean(), + f"{prefix}_joint_abs_err_deg_final": err[:, -1].mean(), + f"{prefix}_joint_abs_err_deg_max_joint": err.mean(dim=(0, 1)).max(), + f"{prefix}_grip_paired_mse_avg": _paired_mse(pred_cpu[..., g], gt_cpu[..., g]), + f"{prefix}_grip_final_mse_avg": _paired_mse( + pred_cpu[:, -1, g], gt_cpu[:, -1, g] + ), + } + + +def layout_metrics( + pred: torch.Tensor, gt: torch.Tensor, prefix: str, ac_key: str | None = None +) -> dict: + """Dispatch on the action key, then the last-dim width: joints, cartesian, + keypoint, or ``{}``. A 14-D joint chunk has a cartesian width (xyz+ypr+grip + per arm), so joint actions must be named: ``ac_key`` ``actions_joints*``.""" + if ac_key is not None and ac_key.startswith("actions_joints"): + return joint_metrics(pred, gt, prefix) metrics = cartesian_metrics(pred, gt, prefix) if metrics: return metrics diff --git a/egomimic/eval/eval_hpt.py b/egomimic/eval/eval_hpt.py index 8c5679d00f..e3c92e395c 100644 --- a/egomimic/eval/eval_hpt.py +++ b/egomimic/eval/eval_hpt.py @@ -59,7 +59,9 @@ def compute_metrics_and_viz(self, batch, do_viz=True): metrics[f"Valid/{pred_key}_frechet_gauss_min"] = fd.min().item() metrics[f"Valid/{pred_key}_frechet_gauss_max"] = fd.max().item() metrics.update( - layout_metrics(preds[pred_key], _batch[ac_key], f"Valid/{pred_key}") + layout_metrics( + preds[pred_key], _batch[ac_key], f"Valid/{pred_key}", ac_key + ) ) if embodiment_name in algo.auxiliary_ac_keys: diff --git a/egomimic/eval/eval_pi.py b/egomimic/eval/eval_pi.py index 6bd1aa4ae4..65c6f04af5 100644 --- a/egomimic/eval/eval_pi.py +++ b/egomimic/eval/eval_pi.py @@ -58,7 +58,9 @@ def compute_metrics_and_viz(self, batch, do_viz=True): if pred_key in preds: metrics.update( - layout_metrics(preds[pred_key], _batch[ac_key], f"Valid/{pred_key}") + layout_metrics( + preds[pred_key], _batch[ac_key], f"Valid/{pred_key}", ac_key + ) ) if do_viz: diff --git a/egomimic/rldb/embodiment/embodiment.py b/egomimic/rldb/embodiment/embodiment.py index cd4a2e30a2..7a5c0602be 100644 --- a/egomimic/rldb/embodiment/embodiment.py +++ b/egomimic/rldb/embodiment/embodiment.py @@ -224,9 +224,26 @@ def viz( ) @classmethod - def get_keymap(cls, keymap_mode: str, norm_mode: bool = False, annotation_key=None): - """Returns a dictionary mapping from the raw keys in the dataset to the canonical keys used by the model.""" + def get_keymap( + cls, + keymap_mode: str, + norm_mode: bool = False, + annotation_key=None, + image_history_gap_s: float | None = None, + ): + """Returns a dictionary mapping from the raw keys in the dataset to the canonical keys used by the model. + + ``image_history_gap_s`` adds ``_hist``, the front frame that + many seconds earlier (RDT's ``trunk.image_history: 2``); the same knob + as ``Human.get_keymap``, so a robot domain can share an RDT trunk with + a human one.""" key_map = cls._get_keymap(keymap_mode) + if image_history_gap_s is not None: + front = key_map[cls.VIZ_IMAGE_KEY] + key_map[f"{cls.VIZ_IMAGE_KEY}{IMAGE_HISTORY_SUFFIX}"] = { + **front, + "lag_s": float(image_history_gap_s), + } if annotation_key is not None and not norm_mode: key_map[annotation_key] = { "key_type": "annotation_keys", diff --git a/egomimic/rldb/embodiment/eva.py b/egomimic/rldb/embodiment/eva.py index b3edfa4f48..b811f39494 100644 --- a/egomimic/rldb/embodiment/eva.py +++ b/egomimic/rldb/embodiment/eva.py @@ -32,6 +32,16 @@ ) +# Joint-space mode: the ABC YAM station / abc_sim action space. +JOINT_MODE = "joints" +JOINT_ACTION_KEY = "actions_joints" +JOINT_STATE_KEY = "observations.state.joint_positions" +# Raw 30 Hz steps read per sample (1.5 s, the cartesian keymap's horizon), +# resampled to the model's chunk length like the cartesian modes. +JOINT_RAW_HORIZON = 45 +JOINT_ORDER = ("left.{}_joints", "left.{}_gripper", "right.{}_joints", "right.{}_gripper") + + class Eva(Embodiment): INTRINSICS = ARIA_INTRINSICS EXTRINSICS = { @@ -62,6 +72,7 @@ def get_transform_list( "cartesian_wristframe_ypr", "cartesian_wristframe_6d", "cartesian_wristframe_quat", + "joints", ], allow_legacy_rotation: bool = False, ) -> list[Transform]: @@ -97,6 +108,13 @@ def get_transform_list( ] if mode == "cartesian_wristframe_quat": return _build_eva_bimanual_eef_frame_transform_list(is_quat=True) + if mode == JOINT_MODE: + # Absolute joint positions + gripper, 14-D per step, in the + # [left j1..j6, left grip, right j1..j6, right grip] order of the + # ABC YAM station and its simulator (abc_sim's 14-D state/action). + # Needs the "joints" keymap (obs_/cmd_joints are not in the + # cartesian one). No frame, no rotation: the sim steps these as-is. + return _build_eva_bimanual_joint_transform_list() raise ValueError(f"Unsupported transform_list mode '{mode}' for Eva") @staticmethod @@ -105,7 +123,8 @@ def get_revert_transform_list(mode: str) -> list[Transform] | None: trained with ``get_transform_list(mode)`` back to camera-frame xyz+ypr+gripper (7/arm), given the batch's proprio ``observations.state.ee_pose``. ``None``: already in that layout.""" - if mode == "cartesian": + if mode in ("cartesian", JOINT_MODE): + # joints: predictions are already the commanded joint targets. return None if mode == "cartesian_6d": return _build_eva_cartesian_revert_6d_transform_list() @@ -120,6 +139,8 @@ def _get_keymap(cls, keymap_mode: str): # One camera naming for every algo. Pi renames onto openpi's slots # itself (egomimic.models.preprocess_pi_obs.PI_CAMERA_SLOTS). keymap_mode = _strip_pi_keymap_mode(cls, keymap_mode) + if keymap_mode == JOINT_MODE: + return cls._get_joint_keymap() front_key = cls.VIZ_IMAGE_KEY right_wrist_key = "observations.images.right_wrist_img" left_wrist_key = "observations.images.left_wrist_img" @@ -177,6 +198,36 @@ def _get_keymap(cls, keymap_mode: str): return key_map + @classmethod + def _get_joint_keymap(cls): + """Cameras + per-arm joint positions and gripper, observed (proprio) + and commanded (45-step action chunk). ABC zarrs (real and the + converted sim_224 episodes) carry ``{left,right}.{obs,cmd}_joints`` + (6) and ``_gripper`` (1, in [0, 1]).""" + key_map = { + cls.VIZ_IMAGE_KEY: {"key_type": "camera_keys", "zarr_key": "images.front_1"}, + "observations.images.right_wrist_img": { + "key_type": "camera_keys", + "zarr_key": "images.right_wrist", + }, + "observations.images.left_wrist_img": { + "key_type": "camera_keys", + "zarr_key": "images.left_wrist", + }, + } + for arm in ("left", "right"): + for part in ("joints", "gripper"): + key_map[f"{arm}.obs_{part}"] = { + "key_type": "proprio_keys", + "zarr_key": f"{arm}.obs_{part}", + } + key_map[f"{arm}.cmd_{part}"] = { + "key_type": "action_keys", + "zarr_key": f"{arm}.cmd_{part}", + "horizon": JOINT_RAW_HORIZON, + } + return key_map + @classmethod def dinov3_keymap(cls): """ @@ -615,3 +666,34 @@ def _build_eva_bimanual_transform_list( ] ) return transform_list + + +def _build_eva_bimanual_joint_transform_list( + *, + actions_key: str = JOINT_ACTION_KEY, + obs_key: str = JOINT_STATE_KEY, + chunk_length: int = 100, + stride: int = 1, +) -> list[Transform]: + """Joint-space pipeline: resample each commanded joint / gripper chunk to + ``chunk_length`` (as the cartesian modes do), then concatenate both arms in + ABC order into ``actions_joints`` (T, 14) and the observed joints into + ``observations.state.joint_positions`` (14,).""" + cmd_keys = [k.format("cmd") for k in JOINT_ORDER] + obs_keys = [k.format("obs") for k in JOINT_ORDER] + transform_list: list[Transform] = [ + InterpolateLinear( + new_chunk_length=chunk_length, + action_key=key, + output_action_key=key, + stride=stride, + ) + for key in cmd_keys + ] + transform_list += [ + InterpolatePadMask(new_chunk_length=chunk_length, stride=stride), + ConcatKeys(key_list=cmd_keys, new_key_name=actions_key, delete_old_keys=True), + ConcatKeys(key_list=obs_keys, new_key_name=obs_key, delete_old_keys=True), + NumpyToTensor(keys=[actions_key, obs_key]), + ] + return transform_list diff --git a/tests/unit/test_eva_joint_mode.py b/tests/unit/test_eva_joint_mode.py new file mode 100644 index 0000000000..4bb72469cf --- /dev/null +++ b/tests/unit/test_eva_joint_mode.py @@ -0,0 +1,106 @@ +"""Joint-space Eva mode (ABC YAM / abc_sim action space): the keymap reads +``{left,right}.{obs,cmd}_{joints,gripper}``, the transform emits a 14-D state +and a (100, 14) action chunk in ABC order, the robot keymap supports the same +``_hist`` camera as Human, and joint chunks get joint metrics, not the +cartesian ones their 14-D width would otherwise trigger.""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch + +from egomimic.eval.action_metrics import layout_metrics +from egomimic.rldb.embodiment.eva import ( + JOINT_ACTION_KEY, + JOINT_STATE_KEY, + Eva, +) +from egomimic.rldb.zarr.zarr_dataset_multi import LocalEpisodeResolver, MultiDataset +from egomimic.rldb.zarr.zarr_writer import ZarrWriter + +T = 120 +CAM = Eva.VIZ_IMAGE_KEY + + +def _write_episode(path, fps=30): + rng = np.random.default_rng(0) + t = np.arange(T, dtype=np.float32)[:, None] + numeric = {} + # Each channel carries a distinct offset so ordering is checkable. + for a, arm in enumerate(("left", "right")): + base = 10.0 * a + numeric[f"{arm}.obs_joints"] = base + np.arange(6, dtype=np.float32) + 0 * t + numeric[f"{arm}.cmd_joints"] = base + np.arange(6, dtype=np.float32) + 0.01 * t + numeric[f"{arm}.obs_gripper"] = np.full((T, 1), 0.25 + a * 0.5, np.float32) + numeric[f"{arm}.cmd_gripper"] = np.full((T, 1), 0.3 + a * 0.5, np.float32) + images = { + k: rng.integers(0, 255, (T, 24, 32, 3), np.uint8) + for k in ("images.front_1", "images.left_wrist", "images.right_wrist") + } + ZarrWriter.create_and_write( + path, + numeric_data=numeric, + image_data=images, + embodiment="eva_bimanual", + fps=fps, + task_name="synthetic", + intrinsics={"front_1": np.eye(3, 4)}, + ) + + +def _leaf(tmp_path, **keymap_kwargs): + _write_episode(tmp_path / "ep.zarr") + key_map = Eva.get_keymap("joints", **keymap_kwargs) + resolver = LocalEpisodeResolver( + tmp_path, key_map=key_map, transform_list=Eva.get_transform_list("joints") + ) + ds = MultiDataset._from_resolver(resolver, mode="total") + return next(iter(ds.datasets.values())) + + +def test_joint_sample_shapes_and_abc_order(tmp_path): + sample = _leaf(tmp_path)[10] + state = sample[JOINT_STATE_KEY].float() + actions = sample[JOINT_ACTION_KEY].float() + assert state.shape == (14,) and actions.shape == (100, 14) + expected = torch.tensor( + [0, 1, 2, 3, 4, 5, 0.25, 10, 11, 12, 13, 14, 15, 0.75], dtype=torch.float32 + ) + assert torch.allclose(state, expected) + # Chunk start = the commanded pose at idx 10 (+0.01 * 10 on the joints). + assert torch.allclose(actions[0, :6], expected[:6] + 0.1, atol=1e-5) + assert torch.allclose(actions[:, 6], torch.full((100,), 0.3)) + assert torch.allclose(actions[:, 13], torch.full((100,), 0.8)) + # 45 raw steps resampled onto 100: the last step is idx 10 + 44. + assert torch.allclose(actions[-1, 0], torch.tensor(0.54), atol=1e-5) + assert sample["fps"].item() == 30 + assert "left.cmd_joints" not in sample and "left.obs_joints" not in sample + + +def test_joint_keymap_supports_the_hist_camera(tmp_path): + leaf = _leaf(tmp_path, image_history_gap_s=0.1) + assert torch.equal(leaf[20][f"{CAM}_hist"], leaf[17][CAM]) + + +def test_joint_mode_has_no_revert(): + assert Eva.get_revert_transform_list("joints") is None + + +def test_cartesian_keymap_is_unchanged_by_the_joint_mode(): + km = Eva.get_keymap("cartesian") + assert "left.cmd_ee_pose" in km and "left.cmd_joints" not in km + + +def test_joint_actions_get_joint_metrics_not_cartesian(): + gt = torch.zeros(2, 100, 14) + pred = gt.clone() + pred[..., 0] = np.deg2rad(2.0) # left j1 off by 2 degrees + pred[..., 6] = 0.5 # left gripper off by 0.5 + m = layout_metrics(pred, gt, "Valid/x", ac_key=JOINT_ACTION_KEY) + assert not any("xyz" in k for k in m) + assert m["Valid/x_joint_abs_err_deg_avg"].item() == pytest.approx(2.0 / 12, rel=1e-4) + assert m["Valid/x_joint_abs_err_deg_max_joint"].item() == pytest.approx(2.0, rel=1e-4) + assert m["Valid/x_grip_paired_mse_avg"].item() == pytest.approx(0.125, rel=1e-4) + # Without the key name the width alone reads as cartesian ypr. + assert any("xyz" in k for k in layout_metrics(pred, gt, "Valid/x")) From f4e8b919af4274885c92dc2a5e0790f2cf0b633a Mon Sep 17 00:00:00 2001 From: Aidan Gao Date: Mon, 28 Sep 2026 15:01:37 -0400 Subject: [PATCH 02/20] train: init_weights_ckpt, a weights-only fine-tune init ckpt_path is a full Lightning resume (optimizer, schedule, epoch) and fails on keys a new embodiment adds, so a human-pretrained RDT could not be fine-tuned on robot data from this stack. The loader that did this on aidan/abc-stationery-pi never reached main. init_weights_ckpt= copies every name+shape-matched tensor and leaves the rest (a new embodiment's stems / head, a different action width) at their fresh init. It raises when under init_weights_min_fraction (0.5) of the model's elements matched, so an unrelated checkpoint cannot silently "fine-tune" from random init, and it is skipped whenever ckpt_path is set, so a Slurm requeue resumes the run's own weights. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../hydra_configs/train_zarr_cartesian.yaml | 6 ++ egomimic/trainHydra.py | 30 ++++++++- egomimic/utils/checkpoint_utils.py | 62 +++++++++++++++++ tests/unit/test_init_weights_ckpt.py | 67 +++++++++++++++++++ 4 files changed, 164 insertions(+), 1 deletion(-) create mode 100644 tests/unit/test_init_weights_ckpt.py diff --git a/egomimic/hydra_configs/train_zarr_cartesian.yaml b/egomimic/hydra_configs/train_zarr_cartesian.yaml index 7ccfe6bc95..1fd05f0557 100644 --- a/egomimic/hydra_configs/train_zarr_cartesian.yaml +++ b/egomimic/hydra_configs/train_zarr_cartesian.yaml @@ -13,6 +13,12 @@ defaults: name: test description: test ckpt_path: null +# Weights-only fine-tune init from a Lightning checkpoint (fresh optimizer, +# schedule and epoch; name+shape-matched tensors only, so new embodiment +# stems / heads start fresh). Ignored when ckpt_path resumes a run. +init_weights_ckpt: null +# Fail if less than this fraction of the model's elements came from it. +init_weights_min_fraction: 0.5 mode: train hydra: diff --git a/egomimic/trainHydra.py b/egomimic/trainHydra.py index cfaff0e529..23e704391a 100644 --- a/egomimic/trainHydra.py +++ b/egomimic/trainHydra.py @@ -28,7 +28,10 @@ PinError, pinned_episode_subset, ) -from egomimic.utils.checkpoint_utils import load_checkpoint_weights +from egomimic.utils.checkpoint_utils import ( + init_weights_from_checkpoint, + load_checkpoint_weights, +) from egomimic.utils.compile_cache import set_per_job_compile_cache_dir from egomimic.utils.env import load_env from egomimic.utils.gpu_orphans import arm_rank_pdeathsig, reap_orphans @@ -149,6 +152,30 @@ def _weights_from_checkpoint(cfg: DictConfig) -> bool: return bool(ckpt_path) and os.path.isfile(ckpt_path) and not cfg.get("pretrained") +def _apply_init_weights(cfg: DictConfig, model) -> Optional[dict]: + """``init_weights_ckpt``: weights-only fine-tune init (see + ``init_weights_from_checkpoint``). Skipped when ``ckpt_path`` is set -- a + resume (including a Slurm requeue, folded into ``ckpt_path`` by + ``_prepare_checkpoint_resume``) restores this run's own weights, which + must not be replaced by the base checkpoint's.""" + init_ckpt = cfg.get("init_weights_ckpt") + if not init_ckpt: + return None + if cfg.get("ckpt_path"): + log.info( + f"init_weights_ckpt={init_ckpt} ignored: resuming from " + f"ckpt_path={cfg.ckpt_path}" + ) + return None + if not os.path.isfile(init_ckpt): + raise FileNotFoundError(f"init_weights_ckpt {init_ckpt} does not exist") + return init_weights_from_checkpoint( + model, + init_ckpt, + min_fraction=float(cfg.get("init_weights_min_fraction", 0.5)), + ) + + def _log_dataset_frame_counts( train_datasets: dict, valid_datasets: dict, @@ -788,6 +815,7 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: os.makedirs(os.path.join(trainer.default_root_dir, "videos"), exist_ok=True) if mode == "train": + _apply_init_weights(cfg, model) if cfg.get("evaluator") is not None: eval_obj: Eval = hydra.utils.instantiate(cfg.evaluator) # data.valid_prefix (e.g. seen_op_valid in the opsplit configs) diff --git a/egomimic/utils/checkpoint_utils.py b/egomimic/utils/checkpoint_utils.py index 9d70bbe739..06025723b0 100644 --- a/egomimic/utils/checkpoint_utils.py +++ b/egomimic/utils/checkpoint_utils.py @@ -40,3 +40,65 @@ def load_checkpoint_weights( "on random weights. Is this checkpoint from the same model config?" ) log.info("Loaded weights from %s", ckpt_path) + + +def init_weights_from_checkpoint( + model: torch.nn.Module, + ckpt_path: str | os.PathLike, + *, + min_fraction: float = 0.5, +) -> dict: + """Fine-tune initialization: copy every tensor of a Lightning checkpoint's + ``state_dict`` whose name and shape match ``model``; leave the rest at their + fresh init. Unlike ``ckpt_path`` (a full resume: optimizer, scheduler, + epoch) nothing but weights is read, so a new embodiment's stems / head and + a different action width start fresh while the shared trunk transfers. + + Raises when fewer than ``min_fraction`` of the model's parameters (by + element count) were loaded: a wrong or unrelated checkpoint would otherwise + "fine-tune" from random init without a word. Returns a report dict. + """ + checkpoint = torch.load( + ckpt_path, map_location="cpu", weights_only=False, mmap=True + ) + ckpt_sd = checkpoint["state_dict"] + own_sd = model.state_dict() + loaded = { + k: v for k, v in ckpt_sd.items() if k in own_sd and own_sd[k].shape == v.shape + } + shape_mismatch = sorted( + k for k, v in ckpt_sd.items() if k in own_sd and own_sd[k].shape != v.shape + ) + unused = sorted(k for k in ckpt_sd if k not in own_sd) + fresh = sorted(k for k in own_sd if k not in loaded) + model.load_state_dict(loaded, strict=False) + + total = sum(v.numel() for v in own_sd.values()) + covered = sum(own_sd[k].numel() for k in loaded) + fraction = covered / max(total, 1) + report = { + "loaded": len(loaded), + "shape_mismatch": shape_mismatch, + "unused": unused, + "fresh": fresh, + "fraction": fraction, + } + log.info( + "init_weights_ckpt %s: loaded %d tensors (%.1f%% of the model's elements); " + "%d shape-mismatched, %d unused, %d left at fresh init, e.g. %s", + ckpt_path, + len(loaded), + 100 * fraction, + len(shape_mismatch), + len(unused), + len(fresh), + fresh[:5], + ) + if fraction < min_fraction: + raise RuntimeError( + f"init_weights_ckpt {ckpt_path}: only {100 * fraction:.1f}% of the " + f"model's weights matched (need {100 * min_fraction:.0f}%); is it " + f"from the same model family? e.g. fresh {fresh[:5]}, " + f"mismatched {shape_mismatch[:5]}" + ) + return report diff --git a/tests/unit/test_init_weights_ckpt.py b/tests/unit/test_init_weights_ckpt.py new file mode 100644 index 0000000000..496cd5c8e6 --- /dev/null +++ b/tests/unit/test_init_weights_ckpt.py @@ -0,0 +1,67 @@ +"""Weights-only fine-tune init: matched tensors load, a new head stays fresh, +an unrelated checkpoint fails loudly, and a resume ignores it.""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn as nn +from omegaconf import OmegaConf + +import egomimic.trainHydra as train_hydra +from egomimic.utils.checkpoint_utils import init_weights_from_checkpoint + + +class _Net(nn.Module): + def __init__(self, head_out: int): + super().__init__() + self.trunk = nn.Linear(64, 64) + self.head = nn.Linear(64, head_out) + + +def _save(model: nn.Module, path) -> str: + torch.save({"state_dict": model.state_dict(), "optimizer_states": [1]}, path) + return str(path) + + +def test_matched_tensors_load_and_the_new_head_stays_fresh(tmp_path): + base = _Net(head_out=18) + ckpt = _save(base, tmp_path / "base.ckpt") + target = _Net(head_out=14) + fresh_head = target.head.weight.detach().clone() + report = init_weights_from_checkpoint(target, ckpt) + assert torch.equal(target.trunk.weight, base.trunk.weight) + assert torch.equal(target.head.weight, fresh_head) + assert report["shape_mismatch"] == ["head.bias", "head.weight"] + assert report["fraction"] > 0.5 + + +def test_an_unrelated_checkpoint_raises(tmp_path): + other = nn.Sequential(nn.Linear(3, 3)) + ckpt = _save(other, tmp_path / "other.ckpt") + with pytest.raises(RuntimeError, match="matched"): + init_weights_from_checkpoint(_Net(14), ckpt) + + +def test_resume_ignores_init_weights(tmp_path): + ckpt = _save(_Net(18), tmp_path / "base.ckpt") + target = _Net(14) + before = target.trunk.weight.detach().clone() + cfg = OmegaConf.create( + {"init_weights_ckpt": ckpt, "ckpt_path": str(tmp_path / "last.ckpt")} + ) + assert train_hydra._apply_init_weights(cfg, target) is None + assert torch.equal(target.trunk.weight, before) + + +def test_missing_init_checkpoint_raises(tmp_path): + cfg = OmegaConf.create( + {"init_weights_ckpt": str(tmp_path / "nope.ckpt"), "ckpt_path": None} + ) + with pytest.raises(FileNotFoundError): + train_hydra._apply_init_weights(cfg, _Net(14)) + + +def test_train_config_declares_the_knob(compose_resolve): + cfg = compose_resolve("train_zarr_cartesian", []) + assert cfg.init_weights_ckpt is None and cfg.init_weights_min_fraction == 0.5 From a63771bc2b6dbd401f448b7f560048e23b6f5809 Mon Sep 17 00:00:00 2001 From: Aidan Gao Date: Mon, 28 Sep 2026 15:04:41 -0400 Subject: [PATCH 03/20] abc_sim: convert ABC sim_224 teleop episodes to EgoVerse zarrs amazon-far/abc publishes 24 simulator tasks of teleop data (sim_224, 412 h) as states_actions.bin + a stacked 3-camera mp4 per episode. The converter writes each episode as an eva_bimanual zarr with the real ABC zarrs' keys ({left,right}.{obs,cmd}_{joints,gripper}, images.{front_1,left_wrist, right_wrist}, annotations), so Eva's "joints" keymap reads real and sim ABC alike. - episodes are selected by task_name (prepare.py extracts every task into one {train,val}_sim folder) and keep ABC's own split as metadata `split`, plus lab=abc_sim and the source episode id; - prompts follow ABC's rule (prompt_timeline > instruction > task name) with the "sim " prefix ABC trains and evaluates sim episodes with; - intrinsics are the sim's fovy-58 pinhole at the stored resolution; - idempotent (done marker + tmp dir + rename), parallel, one bad episode is reported in conversion_report.json instead of stopping the task. Co-Authored-By: Claude Opus 5.5 (1M context) --- egomimic/scripts/abc_sim/__init__.py | 0 .../scripts/abc_sim/convert_sim_to_zarr.py | 273 ++++++++++++++++++ tests/unit/test_abc_sim_convert.py | 114 ++++++++ 3 files changed, 387 insertions(+) create mode 100644 egomimic/scripts/abc_sim/__init__.py create mode 100644 egomimic/scripts/abc_sim/convert_sim_to_zarr.py create mode 100644 tests/unit/test_abc_sim_convert.py diff --git a/egomimic/scripts/abc_sim/__init__.py b/egomimic/scripts/abc_sim/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/egomimic/scripts/abc_sim/convert_sim_to_zarr.py b/egomimic/scripts/abc_sim/convert_sim_to_zarr.py new file mode 100644 index 0000000000..a09eb4b2d2 --- /dev/null +++ b/egomimic/scripts/abc_sim/convert_sim_to_zarr.py @@ -0,0 +1,273 @@ +"""Convert ABC sim_224 teleop episodes (amazon-far/abc) to EgoVerse zarrs. + +ABC's ``prepare.py --sim-data `` downloads a task tar and extracts every +episode into ``$ABC_CACHE/{train,val}_sim//``: ``states_actions.bin`` +((T, 28) float64: 14 recorded + 14 commanded dofs, [left j1..j6, left grip, +right j1..j6, right grip]), ``combined_camera-images-rgb.mp4`` (the cameras +stacked vertically, 224x168 each, 30 fps) and ``episode_metadata.json``. All +tasks share those two folders, so episodes are selected by their +``task_name``. + +Each episode becomes one ``eva_bimanual`` zarr with the same keys as the real +ABC zarrs on the Phoenix mirror (``{left,right}.{obs,cmd}_{joints,gripper}``, +``images.{front_1,left_wrist,right_wrist}``, ``annotations``), so +``Eva.get_keymap("joints")`` reads real and sim ABC data alike. Camera map: +top -> front_1, left -> left_wrist, right -> right_wrist. Extra metadata: +``lab: abc_sim``, ``split`` (ABC's own train/val), ``abc_episode``, ``prompt``. + +Prompts follow ABC's training rule (``abc_minimal.dataloader``): the episode's +``prompt_timeline`` if it has one, else its ``instruction`` when that differs +from the task name, else the task name with ``_`` -> `` ``; ``sim `` is +prefixed when missing, as ABC's mixture prompts sim episodes and as +``abc_sim`` prompts the env at eval time. + + python -m egomimic.scripts.abc_sim.convert_sim_to_zarr \ + --src $ABC_CACHE --task sim_put_the_plastic_bottles_in_the_bin \ + --out /storage/project/r-dxu345-0/agao81/abc_sim/zarr/put_bottles +""" + +from __future__ import annotations + +import argparse +import json +import logging +import math +import os +import shutil +import sys +from concurrent.futures import ProcessPoolExecutor, as_completed +from pathlib import Path + +import numpy as np + +log = logging.getLogger("abc_sim.convert") + +STATE_DIM = 14 +DATA_FPS = 30 +DEFAULT_CAMERAS = ("top", "left", "right") +CAMERA_TO_ZARR = { + "top": "images.front_1", + "left": "images.left_wrist", + "right": "images.right_wrist", +} +# abc_sim scenes: every policy camera is a fovy=58 deg pinhole. +CAMERA_FOVY_DEG = 58.0 +SIM_PROMPT_PREFIX = "sim " +DONE_MARKER = ".abc_sim_done" + + +def task_name_to_prompt(task_name: str) -> str: + """abc_minimal.dit.task_name_to_prompt.""" + return " ".join(task_name.replace("-", " ").replace("_", " ").split()) + + +def with_sim_prefix(prompt: str) -> str: + prompt = prompt.strip() + return prompt if prompt.startswith(SIM_PROMPT_PREFIX) else SIM_PROMPT_PREFIX + prompt + + +def prompt_spans(meta: dict, num_steps: int) -> list[tuple[str, int, int]]: + """``(text, start, end)`` spans covering ``[0, num_steps)``, end exclusive.""" + task_name = str(meta.get("task_name") or "") + timeline = meta.get("prompt_timeline") + if timeline: + entries = sorted((int(e["frame"]), str(e["prompt"])) for e in timeline) + else: + instruction = meta.get("instruction") + if ( + isinstance(instruction, str) + and instruction.strip() + and task_name_to_prompt(instruction) != task_name_to_prompt(task_name) + ): + text = instruction.strip() + else: + text = task_name_to_prompt(task_name) + entries = [(0, text)] + spans = [] + for i, (start, text) in enumerate(entries): + end = entries[i + 1][0] if i + 1 < len(entries) else num_steps + start = max(0, min(start, num_steps)) + if i == 0: + start = 0 # the first directive holds from the episode start + if end > start: + spans.append((with_sim_prefix(text), start, end)) + return spans + + +def pinhole_intrinsics(height: int, width: int, fovy_deg: float = CAMERA_FOVY_DEG): + """3x4 K of a MuJoCo camera (square pixels, principal point centred).""" + f = 0.5 * height / math.tan(math.radians(fovy_deg) / 2.0) + return np.array( + [[f, 0.0, width / 2.0, 0.0], [0.0, f, height / 2.0, 0.0], [0.0, 0.0, 1.0, 0.0]], + dtype=np.float64, + ) + + +def load_states_actions(ep_dir: Path) -> tuple[np.ndarray, np.ndarray]: + raw = np.fromfile(ep_dir / "states_actions.bin", dtype=np.float64) + if raw.size == 0 or raw.size % (2 * STATE_DIM): + raise ValueError(f"states_actions.bin holds {raw.size} float64s") + table = raw.reshape(-1, 2 * STATE_DIM) + return table[:, :STATE_DIM], table[:, STATE_DIM:] + + +def decode_cameras(ep_dir: Path, cameras: tuple[str, ...]) -> dict[str, np.ndarray]: + """``{camera: (T, H, W, 3) uint8}`` from the vertically stacked mp4.""" + import av + + frames = [] + with av.open(str(ep_dir / "combined_camera-images-rgb.mp4")) as container: + for frame in container.decode(video=0): + frames.append(frame.to_ndarray(format="rgb24")) + if not frames: + raise ValueError("video has no frames") + stack = np.stack(frames) # (T, n * H, W, 3) + h = stack.shape[1] // len(cameras) + if h * len(cameras) != stack.shape[1]: + raise ValueError( + f"video height {stack.shape[1]} is not {len(cameras)} stacked cameras" + ) + out = {cam: stack[:, i * h : (i + 1) * h] for i, cam in enumerate(cameras)} + if "top" not in out and "top_left" in out and "top_right" in out: + # abc_minimal.dataloader: stereo episodes alias one eye to "top". + import hashlib + + digest = hashlib.sha1(ep_dir.name.encode("utf-8")).digest()[0] + out["top"] = out["top_left" if digest % 2 == 0 else "top_right"] + return out + + +def split_arms(x: np.ndarray, kind: str) -> dict[str, np.ndarray]: + x = np.asarray(x, dtype=np.float32) + return { + f"left.{kind}_joints": x[:, 0:6], + f"left.{kind}_gripper": x[:, 6:7], + f"right.{kind}_joints": x[:, 7:13], + f"right.{kind}_gripper": x[:, 13:14], + } + + +def episode_hash(split: str, ep_dir: Path) -> str: + return f"abcsim_{split}_{ep_dir.name}" + + +def convert_episode(ep_dir: Path, split: str, out_dir: Path, overwrite: bool = False) -> dict: + """Write one zarr; returns a report row. Idempotent via a done marker.""" + from egomimic.rldb.zarr.zarr_writer import ZarrWriter + + name = episode_hash(split, ep_dir) + dst = out_dir / f"{name}.zarr" + if (dst / DONE_MARKER).exists() and not overwrite: + return {"episode": name, "status": "skipped"} + meta = json.loads((ep_dir / "episode_metadata.json").read_text()) + states, actions = load_states_actions(ep_dir) + num_steps = len(states) + cameras = tuple(meta.get("cameras") or DEFAULT_CAMERAS) + images = decode_cameras(ep_dir, cameras) + n_frames = min(len(v) for v in images.values()) + if n_frames != num_steps: + raise ValueError(f"{n_frames} video frames for {num_steps} state rows") + missing = [c for c in DEFAULT_CAMERAS if c not in images] + if missing: + raise ValueError(f"cameras missing: {missing} (have {sorted(images)})") + h, w = images["top"].shape[1:3] + + tmp = out_dir / f".{name}.zarr.tmp" + shutil.rmtree(tmp, ignore_errors=True) + spans = prompt_spans(meta, num_steps) + ZarrWriter.create_and_write( + tmp, + numeric_data={**split_arms(states, "obs"), **split_arms(actions, "cmd")}, + image_data={CAMERA_TO_ZARR[c]: images[c] for c in DEFAULT_CAMERAS}, + embodiment="eva_bimanual", + fps=DATA_FPS, + task_name=str(meta.get("task_name") or ""), + task_description=spans[0][0] if spans else "", + annotations=spans, + intrinsics={"front_1": pinhole_intrinsics(h, w)}, + metadata_override={ + "lab": "abc_sim", + "split": split, + "abc_episode": ep_dir.name, + "prompt": spans[0][0] if spans else "", + "source": "abc sim_224", + }, + ) + (tmp / DONE_MARKER).write_text("ok\n") + shutil.rmtree(dst, ignore_errors=True) + os.replace(tmp, dst) + return {"episode": name, "status": "written", "frames": num_steps} + + +def find_episodes(src: Path, task: str, splits: list[str]) -> list[tuple[Path, str]]: + found = [] + for split in splits: + root = src / f"{split}_sim" + if not root.is_dir(): + log.warning("no %s", root) + continue + for bin_path in sorted(root.glob("*/states_actions.bin")): + ep_dir = bin_path.parent + meta_path = ep_dir / "episode_metadata.json" + if not meta_path.exists(): + continue + if json.loads(meta_path.read_text()).get("task_name") == task: + found.append((ep_dir, split)) + return found + + +def main(argv=None) -> int: + p = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + p.add_argument("--src", type=Path, required=True, help="ABC cache dir ($ABC_CACHE)") + p.add_argument("--task", required=True, help="dataset task_name, e.g. sim_pouring_beads") + p.add_argument("--out", type=Path, required=True) + p.add_argument("--splits", nargs="+", default=["train", "val"]) + p.add_argument("--limit", type=int, default=None, help="per split, for smokes") + p.add_argument("--workers", type=int, default=max(1, (os.cpu_count() or 2) // 2)) + p.add_argument("--overwrite", action="store_true") + args = p.parse_args(argv) + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s") + + episodes = find_episodes(args.src, args.task, args.splits) + if args.limit is not None: + by_split: dict[str, list] = {} + for ep in episodes: + by_split.setdefault(ep[1], []).append(ep) + episodes = [ep for eps in by_split.values() for ep in eps[: args.limit]] + if not episodes: + log.error("no %s episodes under %s/{%s}_sim", args.task, args.src, ",".join(args.splits)) + return 1 + args.out.mkdir(parents=True, exist_ok=True) + log.info("converting %d episodes of %s -> %s", len(episodes), args.task, args.out) + + rows, failures = [], [] + with ProcessPoolExecutor(max_workers=args.workers) as pool: + futures = { + pool.submit(convert_episode, ep, split, args.out, args.overwrite): ep + for ep, split in episodes + } + for i, fut in enumerate(as_completed(futures), 1): + ep = futures[fut] + try: + rows.append(fut.result()) + except Exception as e: # one bad episode must not sink the task + failures.append({"episode": str(ep), "error": f"{type(e).__name__}: {e}"}) + log.warning("FAILED %s: %s", ep, e) + if i % 100 == 0 or i == len(futures): + log.info("%d/%d done (%d failed)", i, len(futures), len(failures)) + + report = { + "task": args.task, + "src": str(args.src), + "written": sum(r["status"] == "written" for r in rows), + "skipped": sum(r["status"] == "skipped" for r in rows), + "failed": failures, + "frames_written": int(sum(r.get("frames", 0) for r in rows)), + } + (args.out / "conversion_report.json").write_text(json.dumps(report, indent=2)) + log.info("report: %s", {k: v for k, v in report.items() if k != "failed"}) + return 0 if not failures else 2 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/unit/test_abc_sim_convert.py b/tests/unit/test_abc_sim_convert.py new file mode 100644 index 0000000000..f32d6a47ab --- /dev/null +++ b/tests/unit/test_abc_sim_convert.py @@ -0,0 +1,114 @@ +"""ABC sim_224 episode -> EgoVerse zarr: states / actions land in the joint +keymap's keys in ABC order, the three stacked cameras are split onto +front/left/right, prompts follow ABC's sim rule, and the joint-mode dataset +reads the converted episode with ABC's own train/val split filterable.""" + +from __future__ import annotations + +import json + +import numpy as np +import pytest +import torch + +av = pytest.importorskip("av") + +from egomimic.rldb.embodiment.eva import JOINT_ACTION_KEY, JOINT_STATE_KEY, Eva +from egomimic.rldb.filters import DatasetFilter +from egomimic.rldb.zarr.zarr_dataset_multi import LocalEpisodeResolver, MultiDataset +from egomimic.scripts.abc_sim import convert_sim_to_zarr as conv + +T, H, W = 60, 16, 24 +TASK = "sim_put_the_plastic_bottles_in_the_bin" + + +def _write_abc_episode(root, split, name, meta_extra=None): + ep = root / f"{split}_sim" / name + ep.mkdir(parents=True) + t = np.arange(T, dtype=np.float64)[:, None] + states = np.tile(np.arange(14, dtype=np.float64), (T, 1)) + 0 * t + actions = states + 0.01 * t + np.concatenate([states, actions], axis=1).astype(np.float64).tofile( + ep / "states_actions.bin" + ) + # Three cameras stacked vertically, each a flat colour so the split is checkable. + with av.open(str(ep / "combined_camera-images-rgb.mp4"), "w") as c: + s = c.add_stream("libx264", rate=30) + s.width, s.height, s.pix_fmt = W, 3 * H, "yuv420p" + s.options = {"crf": "0"} + for _ in range(T): + img = np.zeros((3 * H, W, 3), np.uint8) + img[:H] = (200, 30, 30) + img[H : 2 * H] = (30, 200, 30) + img[2 * H :] = (30, 30, 200) + for pkt in s.encode(av.VideoFrame.from_ndarray(img, format="rgb24")): + c.mux(pkt) + for pkt in s.encode(): + c.mux(pkt) + meta = {"task_name": TASK, "num_steps": T, "cameras": ["top", "left", "right"]} + meta.update(meta_extra or {}) + (ep / "episode_metadata.json").write_text(json.dumps(meta)) + return ep + + +def test_prompt_rule_matches_abc(): + assert conv.prompt_spans({"task_name": TASK}, 10) == [ + ("sim put the plastic bottles in the bin", 0, 10) + ] + spans = conv.prompt_spans( + { + "task_name": "multi_drawer_search", + "prompt_timeline": [ + {"frame": 0, "prompt": "find the cup"}, + {"frame": 4, "prompt": "find the ball"}, + ], + }, + 10, + ) + assert spans == [("sim find the cup", 0, 4), ("sim find the ball", 4, 10)] + spans = conv.prompt_spans( + {"task_name": "put_relative", "instruction": "put the cube left of the mug"}, 5 + ) + assert spans == [("sim put the cube left of the mug", 0, 5)] + + +def test_intrinsics_are_the_sim_pinhole(): + K = conv.pinhole_intrinsics(168, 224) + assert K.shape == (3, 4) and K[0, 2] == 112 and K[1, 2] == 84 + assert K[1, 1] == pytest.approx(84 / np.tan(np.radians(29))) + + +def test_convert_and_read_back(tmp_path): + src, out = tmp_path / "cache", tmp_path / "zarr" + _write_abc_episode(src, "train", "episode_0001") + _write_abc_episode(src, "val", "episode_0002") + _write_abc_episode(src, "train", "episode_0003", {"task_name": "sim_pouring_beads"}) + assert conv.main(["--src", str(src), "--task", TASK, "--out", str(out), "--workers", "1"]) == 0 + report = json.loads((out / "conversion_report.json").read_text()) + assert report["written"] == 2 and not report["failed"] + # Idempotent: a rerun skips both. + conv.main(["--src", str(src), "--task", TASK, "--out", str(out), "--workers", "1"]) + assert json.loads((out / "conversion_report.json").read_text())["skipped"] == 2 + + resolver = LocalEpisodeResolver( + out, + key_map=Eva.get_keymap("joints", annotation_key="annotations"), + transform_list=Eva.get_transform_list("joints"), + ) + filters = DatasetFilter(filter_lambdas=["lambda row: row.get('split') == 'train'"]) + ds = MultiDataset._from_resolver(resolver, filters=filters, mode="total") + assert len(ds.datasets) == 1 + leaf = next(iter(ds.datasets.values())) + sample = leaf[5] + assert torch.allclose(sample[JOINT_STATE_KEY].float(), torch.arange(14.0)) + assert torch.allclose( + sample[JOINT_ACTION_KEY][0].float(), torch.arange(14.0) + 0.05, atol=1e-5 + ) + assert sample["annotations"] == ["sim put the plastic bottles in the bin"] + front = sample[Eva.VIZ_IMAGE_KEY] + assert front.shape == (3, H, W) + # top camera is red-dominant, left wrist green, right wrist blue + assert front[0].mean() > front[1].mean() + assert sample["observations.images.left_wrist_img"][1].mean() > 0.5 + assert sample["observations.images.right_wrist_img"][2].mean() > 0.5 + assert sample["fps"].item() == 30 From 68ac5df070f3c85e30a78b83ff0e17d5b7452281 Mon Sep 17 00:00:00 2001 From: Aidan Gao Date: Tue, 29 Sep 2026 00:59:20 -0400 Subject: [PATCH 04/20] abc: RDT-1B joint-space recipes, the fold ladder, a Lambda pack launcher The model and data side of the ABC sim pipeline (egomimic/scripts/abc_sim/ README.md): - model/rdt_1b_abc_joints_dinov3_vitb: the 1B DiT with an eva_bimanual head and state stem in the ABC station's 14-D joint space beside the mecka human head (shared trunk, tower and text stem). - data/abc_sim_joints (converted sim_224 episodes, ABC's own train/val split), abc_real_joints (lab=abc on the mirror; the 2026-09-01 relabel makes lab+task+path the stable filter; front frames resized to one size because ABC mixes 640x480 and 1280x720), abc_sim_joints_cotrain_mecka (single-task human+robot cotrain, one batch of each per step). - train_zarr_abc_sim_rdt_ft / _real_rdt_ft / _sim_rdt_cotrain recipes and evaluator/eval_hpt_abc_joints: joint metrics for the robot, the wrist-frame 6D overlay for the human; an embodiment without a viz entry gets metrics only. - train_zarr_fold_ladder_rdt1b: the RDT recipe at 1B on one rung of the mecka freeform-fold operator ladder (k = 1/2/4), 1000 x 100 steps, a checkpoint every 100 epochs -- the bases the ABC fine-tunes start from. - Eva keymap "joints_front" (front camera only) and get_transform_list image_hw (ResizeImages) for mixed-resolution sources. - EGOVERSE_EPISODE_TABLE: the S3 resolver reads a pickled episode-table snapshot where there is no SQL (Lambda). - scripts/abc_sim/lambda_pack.sbatch: several small-DDP runs on one exclusive loaner node, datasets staged to node-local /ephemeral first. Co-Authored-By: Claude Fable 5.1 --- egomimic/eval/eval_hpt.py | 4 + .../hydra_configs/data/abc_real_joints.yaml | 36 ++++++++ .../hydra_configs/data/abc_sim_joints.yaml | 56 ++++++++++++ .../data/abc_sim_joints_cotrain_mecka.yaml | 58 +++++++++++++ .../evaluator/eval_hpt_abc_joints.yaml | 17 ++++ .../viz/abc_joints_human_cartesian.yaml | 10 +++ .../model/rdt_1b_abc_joints_dinov3_vitb.yaml | 61 +++++++++++++ .../train_zarr_abc_real_rdt_ft.yaml | 10 +++ .../train_zarr_abc_sim_rdt_cotrain.yaml | 15 ++++ .../train_zarr_abc_sim_rdt_ft.yaml | 39 +++++++++ .../train_zarr_fold_ladder_rdt1b.yaml | 42 +++++++++ egomimic/rldb/embodiment/eva.py | 26 +++++- egomimic/rldb/zarr/action_chunk_transforms.py | 37 ++++++++ egomimic/rldb/zarr/zarr_dataset_multi.py | 5 ++ .../scripts/abc_sim/convert_sim_to_zarr.py | 10 ++- egomimic/scripts/abc_sim/lambda_pack.sbatch | 61 +++++++++++++ egomimic/scripts/abc_sim/tasks.py | 66 ++++++++++++++ tests/fixtures/synthetic_episodes.py | 10 +++ tests/unit/test_abc_joint_configs.py | 87 +++++++++++++++++++ tests/unit/test_episode_table_file.py | 16 ++++ 20 files changed, 661 insertions(+), 5 deletions(-) create mode 100644 egomimic/hydra_configs/data/abc_real_joints.yaml create mode 100644 egomimic/hydra_configs/data/abc_sim_joints.yaml create mode 100644 egomimic/hydra_configs/data/abc_sim_joints_cotrain_mecka.yaml create mode 100644 egomimic/hydra_configs/evaluator/eval_hpt_abc_joints.yaml create mode 100644 egomimic/hydra_configs/evaluator/viz/abc_joints_human_cartesian.yaml create mode 100644 egomimic/hydra_configs/model/rdt_1b_abc_joints_dinov3_vitb.yaml create mode 100644 egomimic/hydra_configs/train_zarr_abc_real_rdt_ft.yaml create mode 100644 egomimic/hydra_configs/train_zarr_abc_sim_rdt_cotrain.yaml create mode 100644 egomimic/hydra_configs/train_zarr_abc_sim_rdt_ft.yaml create mode 100644 egomimic/hydra_configs/train_zarr_fold_ladder_rdt1b.yaml create mode 100644 egomimic/scripts/abc_sim/lambda_pack.sbatch create mode 100644 egomimic/scripts/abc_sim/tasks.py create mode 100644 tests/unit/test_abc_joint_configs.py create mode 100644 tests/unit/test_episode_table_file.py diff --git a/egomimic/eval/eval_hpt.py b/egomimic/eval/eval_hpt.py index e3c92e395c..9638d852f2 100644 --- a/egomimic/eval/eval_hpt.py +++ b/egomimic/eval/eval_hpt.py @@ -105,6 +105,10 @@ def compute_metrics_and_viz(self, batch, do_viz=True): if not do_viz: continue + if self.viz_func is not None and embodiment_name not in self.viz_func: + # No overlay for this embodiment (e.g. joint-space actions, + # which have no image projection): metrics only. + continue transform_list = self.transform_lists.get(embodiment_name) main_pred_key = f"{embodiment_name}_{ac_key}" diff --git a/egomimic/hydra_configs/data/abc_real_joints.yaml b/egomimic/hydra_configs/data/abc_real_joints.yaml new file mode 100644 index 0000000000..323c4075f1 --- /dev/null +++ b/egomimic/hydra_configs/data/abc_real_joints.yaml @@ -0,0 +1,36 @@ +# ABC_REAL_JOINTS: one REAL ABC task (lab=abc on the Phoenix mirror) in the +# same joint space as abc_sim_joints -- the "real" side of the sim-real check. +# The SQL rows were relabelled eva_bimanual -> yam_bimanual on 2026-09-01 but +# the mirror zarrs still say eva_bimanual, so the filter keys on lab + task + +# the /abc/ path, never the embodiment. The paper's three sim-real tasks exist +# here: "put the plastic bottles in the bin" (100 eps, the default below), "load +# the mixed dishes into the dish rack" (88), "turn the mug right side up" (65); +# another task restates the train filter_lambdas list with its name. +defaults: + - abc_sim_joints + - _self_ + +train_datasets: + eva_bimanual: + resolver: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver + folder_path: ${paths.dataset_dir} + transform_list: + # ABC mixes 640x480 and 1280x720 front frames; one size so they collate + image_hw: [240, 320] + filters: + _target_: egomimic.rldb.filters.DatasetFilter + filter_lambdas: + - "lambda row: str(row.get('lab', '')).lower() == 'abc'" + - "lambda row: str(row.get('task', '')) == 'put the plastic bottles in the bin'" + - "lambda row: '/abc/' in str(row.get('zarr_processed_path', ''))" + - "lambda row: not bool(row.get('is_deleted', False))" + mode: train + valid_ratio: 0.1 + +valid_datasets: + eva_bimanual: + filters: ${data.train_datasets.eva_bimanual.filters} + mode: valid + valid_ratio: ${data.train_datasets.eva_bimanual.valid_ratio} + diff --git a/egomimic/hydra_configs/data/abc_sim_joints.yaml b/egomimic/hydra_configs/data/abc_sim_joints.yaml new file mode 100644 index 0000000000..c213d33f31 --- /dev/null +++ b/egomimic/hydra_configs/data/abc_sim_joints.yaml @@ -0,0 +1,56 @@ +# ABC_SIM_JOINTS: one abc_sim task's teleop episodes (amazon-far/abc sim_224), +# converted by egomimic/scripts/abc_sim/convert_sim_to_zarr.py into +# $ABC_SIM_ZARR_ROOT// (default task put_bottles; another task is +# data.train_datasets.eva_bimanual.resolver.folder_path=/turn_mug -- the +# valid set interpolates the resolver). JOINT space (Eva mode "joints"): the +# 14-D ABC state and a (100, 14) actions_joints chunk from 45 raw 30 Hz steps. +# Train / valid follow ABC's own split (zarr metadata `split`). +# +# The front camera (sim `top`) is the only one decoded (keymap joints_front); +# image_history_gap_s pairs with the RDT trunk's image_history: 2. +_target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper + +train_datasets: + eva_bimanual: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + resolver: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.LocalEpisodeResolver + folder_path: ${oc.env:ABC_SIM_ZARR_ROOT,/workspace/users/agao81/datasets/abc_sim/zarr}/put_bottles + key_map: + _target_: egomimic.rldb.embodiment.eva.Eva.get_keymap + keymap_mode: joints_front + annotation_key: annotations + image_history_gap_s: 0.1 + transform_list: + _target_: egomimic.rldb.embodiment.eva.Eva.get_transform_list + mode: joints + filters: + _target_: egomimic.rldb.filters.DatasetFilter + filter_lambdas: + - "lambda row: str(row.get('lab', '')) == 'abc_sim'" + - "lambda row: str(row.get('split', '')) == 'train'" + mode: total + +valid_datasets: + eva_bimanual: + _target_: ${data.train_datasets.eva_bimanual._target_} + resolver: ${data.train_datasets.eva_bimanual.resolver} + filters: + _target_: egomimic.rldb.filters.DatasetFilter + filter_lambdas: + - "lambda row: str(row.get('lab', '')) == 'abc_sim'" + - "lambda row: str(row.get('split', '')) == 'val'" + mode: total + + + +train_dataloader_params: + eva_bimanual: + batch_size: 32 + num_workers: 8 + +valid_dataloader_params: + eva_bimanual: + batch_size: 32 + num_workers: 4 + shuffle: false diff --git a/egomimic/hydra_configs/data/abc_sim_joints_cotrain_mecka.yaml b/egomimic/hydra_configs/data/abc_sim_joints_cotrain_mecka.yaml new file mode 100644 index 0000000000..7bbff1593a --- /dev/null +++ b/egomimic/hydra_configs/data/abc_sim_joints_cotrain_mecka.yaml @@ -0,0 +1,58 @@ +# ABC_SIM_JOINTS_COTRAIN_MECKA: single-task cotrain -- one abc_sim task (the +# robot, joint space) + one mecka task (human, wrist-frame 6D cartesian), +# one batch of each per step (CombinedLoader max_size_cycle), so the mix is +# set by the two batch sizes. Pairs with model=rdt_1b_abc_joints_dinov3_vitb. +# +# Default pairing put_bottles <- mecka cup_on_saucer (the largest mecka task); +# another human task restates the human filter_lambdas list (dishrack <- +# dishwashing, put_bottles <- packaging_items are the nearest by motion). +defaults: + - abc_sim_joints + - _self_ + +train_datasets: + human_bimanual: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + resolver: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver + folder_path: ${paths.dataset_dir} + key_map: + _target_: egomimic.rldb.embodiment.human.Human.get_keymap + keymap_mode: cartesian + annotation_key: annotations + image_history_gap_s: 0.1 + transform_list: + _target_: egomimic.rldb.embodiment.human.Human.get_transform_list + mode: cartesian_wristframe_6d + stride: 1 + fix_left_wrist_convention: true + pad_proprio_gripper: true + filters: + _target_: egomimic.rldb.filters.DatasetFilter + filter_lambdas: + - "lambda row: str(row.get('lab', '')).lower() == 'mecka'" + - "lambda row: str(row.get('task', '')) == 'cup_on_saucer'" + exclude_hashes: mecka_bad_egomotion + mode: train + valid_ratio: 0.05 + +valid_datasets: + human_bimanual: + _target_: ${data.train_datasets.human_bimanual._target_} + resolver: ${data.train_datasets.human_bimanual.resolver} + filters: ${data.train_datasets.human_bimanual.filters} + mode: valid + valid_ratio: ${data.train_datasets.human_bimanual.valid_ratio} + + + +train_dataloader_params: + human_bimanual: + batch_size: 32 + num_workers: 8 + +valid_dataloader_params: + human_bimanual: + batch_size: 32 + num_workers: 4 + shuffle: false diff --git a/egomimic/hydra_configs/evaluator/eval_hpt_abc_joints.yaml b/egomimic/hydra_configs/evaluator/eval_hpt_abc_joints.yaml new file mode 100644 index 0000000000..c51a298640 --- /dev/null +++ b/egomimic/hydra_configs/evaluator/eval_hpt_abc_joints.yaml @@ -0,0 +1,17 @@ +# HPT/RDT evaluator for the ABC joint-space runs: eva_bimanual (ABC station, +# actions_joints) gets the joint metrics (Valid/eva_bimanual_actions_joints_ +# joint_abs_err_deg_*, _grip_*) and no video; human_bimanual keeps the +# wrist-frame 6D revert + overlay of eval_hpt_wristframe_6d. The closed-loop +# number for these runs is the abc_sim rollout, not this. +defaults: + - viz@viz_func: abc_joints_human_cartesian + - _self_ + +_target_: egomimic.eval.eval_hpt.HPTEvalVideo + +viz_every_n_epochs: 100 +viz_max_batches: 10 + +transform_lists: + human_bimanual: + _target_: egomimic.rldb.embodiment.human._build_human_cartesian_revert_6d_wristframe_transform_list diff --git a/egomimic/hydra_configs/evaluator/viz/abc_joints_human_cartesian.yaml b/egomimic/hydra_configs/evaluator/viz/abc_joints_human_cartesian.yaml new file mode 100644 index 0000000000..11a6188b52 --- /dev/null +++ b/egomimic/hydra_configs/evaluator/viz/abc_joints_human_cartesian.yaml @@ -0,0 +1,10 @@ +# Overlay viz for the ABC joint-space cotrain: human only. The robot predicts +# joint angles, which have no image projection without forward kinematics, so +# eva_bimanual is left out and HPTEvalVideo logs its metrics without a video. +human_bimanual: + _target_: egomimic.rldb.embodiment.human.Human.viz_gt_preds + _partial_: true + image_key: observations.images.front_img_1 + action_key: actions_cartesian + annotation_key: sampled_prompt + mode: traj diff --git a/egomimic/hydra_configs/model/rdt_1b_abc_joints_dinov3_vitb.yaml b/egomimic/hydra_configs/model/rdt_1b_abc_joints_dinov3_vitb.yaml new file mode 100644 index 0000000000..9c956d5db2 --- /dev/null +++ b/egomimic/hydra_configs/model/rdt_1b_abc_joints_dinov3_vitb.yaml @@ -0,0 +1,61 @@ +# RDT_1B_ABC_JOINTS_DINOV3_VITB: RDT-1B (width 2048 x 28 x 32 heads) with a +# second embodiment, the ABC YAM station in JOINT space, beside the mecka +# human head -- the fine-tune / single-task cotrain model for the ABC sim +# rollout pipeline (egomimic/scripts/abc_sim/README.md). +# +# eva_bimanual here is the ABC station (real ABC zarrs and the converted +# sim_224 episodes both register as eva_bimanual): 14-D proprio +# (observations.state.joint_positions) and a (100, 14) actions_joints chunk, +# [left j1..j6, left grip, right j1..j6, right grip], exactly abc_sim's +# action space, so a rollout needs no IK. The DiT trunk, the DINOv3 front +# camera tower and the Qwen text stem are shared with human_bimanual; the +# robot gets its own state adaptor and denoiser head (RDT's per-embodiment +# adaptors at native width). +# +# Fine-tuning from a human-only RDT-1B (model=rdt_1b_bc_mecka_6d_dinov3_vitb): +# init_weights_ckpt= loads the trunk, towers and the human stems; +# the eva stem / head start fresh. ABC-only fine-tunes set +# model.robomimic_model.domains=[eva_bimanual] (the train configs do). +# +# Known gaps, deliberately left for later: the front camera only (the human +# pretrain has one camera; ABC's own policies use top + both wrists), and +# ABC's 4:3 frames go through the tower at the human 16:9 size [224, 400]. +defaults: + - rdt_1b_bc_mecka_6d_dinov3_vitb + - _self_ + +robomimic_model: + domains: ["human_bimanual", "eva_bimanual"] + dims: + eva_bimanual: {proprio: 14, action: 14} + ac_keys: + eva_bimanual: "actions_joints" + + head_specs: + eva_bimanual: + _target_: egomimic.models.fm_policy.FMPolicy + action_horizon: 100 + num_inference_steps: 10 + pooling: null + time_dist: "beta" + infer_ac_dims: + eva_bimanual: ${model.robomimic_model.dims.eva_bimanual.action} + model: + _target_: egomimic.models.rdt_nets.RDTDenoiser + act_dim: ${model.robomimic_model.dims.eva_bimanual.action} + act_seq: 100 + hidden_dim: ${model.robomimic_model.width} + n_state_tokens: 1 + time_scale: 1000.0 + + stem_specs: + eva_bimanual: + state_joint_positions: + _target_: egomimic.models.hpt_nets.MLPPolicyStem + input_dim: ${model.robomimic_model.dims.eva_bimanual.proprio} + output_dim: ${model.robomimic_model.width} + widths: + - ${model.robomimic_model.width} + - ${model.robomimic_model.width} + history_len: 1 + history_dropout: 0.2 diff --git a/egomimic/hydra_configs/train_zarr_abc_real_rdt_ft.yaml b/egomimic/hydra_configs/train_zarr_abc_real_rdt_ft.yaml new file mode 100644 index 0000000000..f1f57c2e9a --- /dev/null +++ b/egomimic/hydra_configs/train_zarr_abc_real_rdt_ft.yaml @@ -0,0 +1,10 @@ +# ABC REAL fine-tune: the train_zarr_abc_sim_rdt_ft arm on the matching REAL +# ABC task (data=abc_real_joints, lab=abc), for the real-robot side of the +# sim-real rank check. Same model, init, steps and batch as the sim arm. +# +# python egomimic/trainHydra.py --config-name=train_zarr_abc_real_rdt_ft \ +# init_weights_ckpt= +defaults: + - train_zarr_abc_sim_rdt_ft + - override data: abc_real_joints + - _self_ diff --git a/egomimic/hydra_configs/train_zarr_abc_sim_rdt_cotrain.yaml b/egomimic/hydra_configs/train_zarr_abc_sim_rdt_cotrain.yaml new file mode 100644 index 0000000000..4adcc0b0b3 --- /dev/null +++ b/egomimic/hydra_configs/train_zarr_abc_sim_rdt_cotrain.yaml @@ -0,0 +1,15 @@ +# Single-task cotrain: one abc_sim task (robot, joint space) + one mecka task +# (human, wrist-frame 6D), RDT-1B with both heads, one batch of each per step. +# Ryan's suggested starting point for the ABC tests (2026-09-27): single-task +# cotrain before the all-mecka base. +# +# python egomimic/trainHydra.py --config-name=train_zarr_abc_sim_rdt_cotrain \ +# name=abc_sim_cotrain description=put_bottles_cup_on_saucer +defaults: + - train_zarr_abc_sim_rdt_ft + - override data: abc_sim_joints_cotrain_mecka + - _self_ + +model: + robomimic_model: + domains: ["human_bimanual", "eva_bimanual"] diff --git a/egomimic/hydra_configs/train_zarr_abc_sim_rdt_ft.yaml b/egomimic/hydra_configs/train_zarr_abc_sim_rdt_ft.yaml new file mode 100644 index 0000000000..4cbe44c720 --- /dev/null +++ b/egomimic/hydra_configs/train_zarr_abc_sim_rdt_ft.yaml @@ -0,0 +1,39 @@ +# ABC sim fine-tune: RDT-1B in the ABC station's joint space on ONE abc_sim +# task's teleop data, initialised (weights only) from a pretrained RDT-1B -- +# the per-base step of the "rank pretraining runs by ABC sim rollouts" +# pipeline (egomimic/scripts/abc_sim/README.md). Rolled out afterwards with +# egomimic/scripts/abc_sim/rollout_sim.sbatch. +# +# python egomimic/trainHydra.py --config-name=train_zarr_abc_sim_rdt_ft \ +# init_weights_ckpt= name=abc_sim_ft description=_put_bottles +# +# init_weights_ckpt=null trains the same arm from scratch (the no-pretrain +# floor every base is compared against). Every arm of a comparison must share +# data, steps and batch; only the init changes. Norm stats are computed at +# startup on the ABC data (eva only: the human base's stats play no part). +defaults: + - train_zarr_cartesian + - override model: rdt_1b_abc_joints_dinov3_vitb + - override data: abc_sim_joints + - override evaluator: eval_hpt_abc_joints + - _self_ + +model: + robomimic_model: + domains: ["eva_bimanual"] + +# the train_viz head would need its own evaluator wiring; the sim rollout is +# this pipeline's check +train_viz: false + +trainer: + gradient_clip_val: 1.0 + max_epochs: 300 + min_epochs: 300 + limit_train_batches: 100 + limit_val_batches: 20 + check_val_every_n_epoch: 25 + +launch_params: + gpus_per_node: 1 + nodes: 1 diff --git a/egomimic/hydra_configs/train_zarr_fold_ladder_rdt1b.yaml b/egomimic/hydra_configs/train_zarr_fold_ladder_rdt1b.yaml new file mode 100644 index 0000000000..b468a128b0 --- /dev/null +++ b/egomimic/hydra_configs/train_zarr_fold_ladder_rdt1b.yaml @@ -0,0 +1,42 @@ +# Freeform fold-clothes OPERATOR ladder, RDT-1B: train_zarr_mecka_flagship_6d_rdt +# (RDT recipe: 2-frame image history, grad clip 1.0, seen/unseen-op heads) at the +# 1B DiT on one rung of data/mecka_fold_freeform_opscale{1,2,4}_6d. Pick the +# rung on the command line; everything else is identical across the ladder: +# +# python egomimic/trainHydra.py --config-name=train_zarr_fold_ladder_rdt1b \ +# data=mecka_fold_freeform_opscale1_6d name=fold_ladder description=k1 +# +# Budget per rung (the RDT memory ablation's): 1000 epochs x 100 steps at a +# global batch of 64 (32 x 2 GPUs), val every 25 epochs, a full checkpoint every +# 100 epochs (10 kept: the step ladder the ABC-sim fine-tunes start from). Norm +# stats are computed at startup per rung (each rung is its own distribution) +# and cached under paths.cache_dir. Runs on Lambda via +# egomimic/scripts/abc_sim/lambda_pack.sbatch, three rungs packed on one node. +defaults: + - train_zarr_mecka_flagship_6d_rdt + - override model: rdt_1b_bc_mecka_6d_dinov3_vitb + - override data: mecka_fold_freeform_opscale4_6d + - _self_ + +launch_params: + gpus_per_node: 2 + nodes: 1 + +data: + train_dataloader_params: + human_bimanual: + batch_size: 32 + num_workers: 12 + valid_dataloader_params: + human_bimanual: + batch_size: 32 + +trainer: + max_epochs: 1000 + min_epochs: 1000 + check_val_every_n_epoch: 25 + limit_train_batches: 100 + +callbacks: + model_checkpoint: + every_n_epochs: 100 diff --git a/egomimic/rldb/embodiment/eva.py b/egomimic/rldb/embodiment/eva.py index b811f39494..7ed667ff6e 100644 --- a/egomimic/rldb/embodiment/eva.py +++ b/egomimic/rldb/embodiment/eva.py @@ -23,6 +23,7 @@ NumpyToTensor, PoseCoordinateFrameTransform, QuaternionPoseToYPR, + ResizeImages, SplitKeys, Transform, XYZWXYZ_to_XYZYPR, @@ -75,14 +76,31 @@ def get_transform_list( "joints", ], allow_legacy_rotation: bool = False, + image_hw: tuple[int, int] | None = None, ) -> list[Transform]: - """Transform pipeline for the requested mode. + """Transform pipeline for the requested mode, then (``image_hw``) every + camera resized to ``(H, W)`` -- for sources mixing resolutions. ``allow_legacy_rotation`` opts out of the continuous-rotation rule for the modes in ``LEGACY_ROTATION_MODES``; only the data/rollout boundary (viz, a checkpoint that predates the 6D conversion) may set it. """ _reject_legacy_rotation(cls, mode, allow_legacy_rotation) + transforms = cls._mode_transform_list(mode) + if image_hw is not None: + transforms = [ResizeImages(cls.CAMERA_KEYS_ALL, tuple(image_hw))] + transforms + return transforms + + # Every camera entry the keymaps can emit, incl. the image-history twin. + CAMERA_KEYS_ALL = [ + "observations.images.front_img_1", + "observations.images.front_img_1_hist", + "observations.images.left_wrist_img", + "observations.images.right_wrist_img", + ] + + @classmethod + def _mode_transform_list(cls, mode: str) -> list[Transform]: if mode == "cartesian": return _build_eva_bimanual_transform_list(is_quat=True) if mode == "cartesian_6d": @@ -141,6 +159,12 @@ def _get_keymap(cls, keymap_mode: str): keymap_mode = _strip_pi_keymap_mode(cls, keymap_mode) if keymap_mode == JOINT_MODE: return cls._get_joint_keymap() + if keymap_mode == f"{JOINT_MODE}_front": + # front camera only: skip decoding wrist frames the model never reads + key_map = cls._get_joint_keymap() + for key in ("left_wrist_img", "right_wrist_img"): + del key_map[f"observations.images.{key}"] + return key_map front_key = cls.VIZ_IMAGE_KEY right_wrist_key = "observations.images.right_wrist_img" left_wrist_key = "observations.images.left_wrist_img" diff --git a/egomimic/rldb/zarr/action_chunk_transforms.py b/egomimic/rldb/zarr/action_chunk_transforms.py index 9150df551b..ac8c263b6e 100644 --- a/egomimic/rldb/zarr/action_chunk_transforms.py +++ b/egomimic/rldb/zarr/action_chunk_transforms.py @@ -963,6 +963,43 @@ def transform(self, batch: dict) -> dict: # --------------------------------------------------------------------------- +class ResizeImages(Transform): + """Resize every camera frame in ``keys`` to ``(H, W)``. Frames are the + decoded ``(3, H, W)`` float arrays ZarrDataset emits (or ``(K, 3, H, W)`` + windows); absent keys are skipped. For sources that mix resolutions (real + ABC: 640x480 and 1280x720 front frames), which default_collate cannot + stack; a model's own stem still resizes afterwards.""" + + def __init__(self, keys: list[str], hw: tuple[int, int]): + self.keys = list(keys) + self.hw = (int(hw[0]), int(hw[1])) + + def _resize(self, chw: np.ndarray) -> np.ndarray: + import cv2 + + if tuple(chw.shape[-2:]) == self.hw: + return chw + hwc = np.moveaxis(np.asarray(chw, dtype=np.float32), 0, -1) + out = cv2.resize(hwc, (self.hw[1], self.hw[0]), interpolation=cv2.INTER_AREA) + return np.moveaxis(out, -1, 0) + + def transform(self, batch: dict) -> dict: + for key in self.keys: + if key not in batch: + continue + x = batch[key] + if isinstance(x, torch.Tensor): + x = x.numpy() + x = np.asarray(x) + if x.ndim == 3: + batch[key] = self._resize(x) + elif x.ndim == 4: + batch[key] = np.stack([self._resize(f) for f in x]) + else: + raise ValueError(f"ResizeImages: '{key}' has shape {x.shape}") + return batch + + class NumpyToTensor(Transform): def __init__(self, keys: list[str]): self.keys = keys diff --git a/egomimic/rldb/zarr/zarr_dataset_multi.py b/egomimic/rldb/zarr/zarr_dataset_multi.py index 85c23c8d0c..46ab431598 100644 --- a/egomimic/rldb/zarr/zarr_dataset_multi.py +++ b/egomimic/rldb/zarr/zarr_dataset_multi.py @@ -88,6 +88,11 @@ def episode_table_to_df(engine): def _episode_table() -> pd.DataFrame: + # EGOVERSE_EPISODE_TABLE=: clusters without + # SQL access (Lambda) resolve from a snapshot exported where SQL exists. + path = os.environ.get("EGOVERSE_EPISODE_TABLE") + if path: + return memoized(("episode_table", path), lambda: pd.read_pickle(path)) return memoized( ("episode_table",), lambda: episode_table_to_df(create_default_engine()) ) diff --git a/egomimic/scripts/abc_sim/convert_sim_to_zarr.py b/egomimic/scripts/abc_sim/convert_sim_to_zarr.py index a09eb4b2d2..68d7bb0879 100644 --- a/egomimic/scripts/abc_sim/convert_sim_to_zarr.py +++ b/egomimic/scripts/abc_sim/convert_sim_to_zarr.py @@ -1,7 +1,8 @@ """Convert ABC sim_224 teleop episodes (amazon-far/abc) to EgoVerse zarrs. ABC's ``prepare.py --sim-data `` downloads a task tar and extracts every -episode into ``$ABC_CACHE/{train,val}_sim//``: ``states_actions.bin`` +episode into ``$ABC_CACHE/{train,val}_sim//`` (a plain ``tar -x`` of +the same tar gives ``{train,val}//``; both are read): ``states_actions.bin`` ((T, 28) float64: 14 recorded + 14 commanded dofs, [left j1..j6, left grip, right j1..j6, right grip]), ``combined_camera-images-rgb.mp4`` (the cameras stacked vertically, 224x168 each, 30 fps) and ``episode_metadata.json``. All @@ -202,9 +203,10 @@ def convert_episode(ep_dir: Path, split: str, out_dir: Path, overwrite: bool = F def find_episodes(src: Path, task: str, splits: list[str]) -> list[tuple[Path, str]]: found = [] for split in splits: - root = src / f"{split}_sim" - if not root.is_dir(): - log.warning("no %s", root) + # prepare.py extracts into {split}_sim/; a plain `tar -x` leaves {split}/ + root = next((src / d for d in (f"{split}_sim", split) if (src / d).is_dir()), None) + if root is None: + log.warning("no %s_sim/ or %s/ under %s", split, split, src) continue for bin_path in sorted(root.glob("*/states_actions.bin")): ep_dir = bin_path.parent diff --git a/egomimic/scripts/abc_sim/lambda_pack.sbatch b/egomimic/scripts/abc_sim/lambda_pack.sbatch new file mode 100644 index 0000000000..06c7fe35e8 --- /dev/null +++ b/egomimic/scripts/abc_sim/lambda_pack.sbatch @@ -0,0 +1,61 @@ +#!/bin/bash +# Pack several small-DDP main-EgoVerse runs on ONE exclusive Lambda loaner node +# (8 x H100), one srun step per run inside the pytorch image, each on its own +# GPUs / DDP port, resuming from its last.ckpt on requeue. The dataset dir is +# copied to the node-local /ephemeral first: zarr reads straight off the NFS +# /workspace stalled nodes and left cancelled jobs undead (Ryan, 2026-09-26). +# +# RUNS="k1 k2 k4" ARGS_k1="--config-name=train_zarr_fold_ladder_rdt1b data=mecka_fold_freeform_opscale1_6d" ... \ +# sbatch --qos=loaner-normal egomimic/scripts/abc_sim/lambda_pack.sbatch +# GPUS_PER_RUN=2 must match launch_params.gpus_per_node of the configs +# DATA_SRC=... NFS dataset dir (default datasets/mecka_abc_ladder), staged to /ephemeral +# TABLE=... EGOVERSE_EPISODE_TABLE pickle (default: the one in DATA_SRC) +# RUN_ROOT=... default /workspace/users/$USER/runs_abc +# EXTRA="..." appended to every run (smokes: trainer.max_epochs=1 ...) +#SBATCH --job-name=abc_pack +#SBATCH --account=loaner +#SBATCH --partition=loaner +#SBATCH --nodes=1 +#SBATCH --gres=gpu:h100:8 +#SBATCH --cpus-per-task=224 +#SBATCH --exclusive +#SBATCH --time=2-00:00:00 +#SBATCH --requeue +#SBATCH --output=/workspace/users/%u/logs/%x-%j.out +set -euo pipefail +SRC=${SRC:-/workspace/users/$USER/EgoVerse} +IMAGE=${IMAGE:-/workspace/users/$USER/images/nvidia-pytorch-25.06-py3.sqsh} +RUN_ROOT=${RUN_ROOT:-/workspace/users/$USER/runs_abc} +DATA_SRC=${DATA_SRC:-/workspace/users/$USER/datasets/mecka_abc_ladder} +TABLE=${TABLE:-$(ls "$DATA_SRC"/episode_table_*.pkl | tail -1)} +DATA=/ephemeral/$USER/$(basename "$DATA_SRC") +G=${GPUS_PER_RUN:-2} +read -r -a runs <<< "${RUNS:?set RUNS to space-separated run names, each with an ARGS_}" +(( ${#runs[@]} * G <= 8 )) || { echo "${#runs[@]} runs x $G GPUs do not fit one node" >&2; exit 2; } +CPUS=${CPUS_PER_TASK:-$(( 192 / (${#runs[@]} * G) ))} +test -s "$IMAGE" && test -s "$SRC/.venv/pyvenv.cfg" && test -s "$TABLE" +echo "=== abc_pack job=$SLURM_JOB_ID node=$(hostname) restart=${SLURM_RESTART_COUNT:-0} commit=$(git -C "$SRC" rev-parse --short HEAD) runs=${runs[*]}" +mkdir -p "$DATA" && rsync -a --exclude _sync "$DATA_SRC/" "$DATA/" && echo "staged $(du -sh "$DATA" | cut -f1) to $DATA" + +i=0 +for name in "${runs[@]}"; do + var="ARGS_$name"; args=${!var:?set $var} + port=$(( ${PORT_BASE:-29500} + 10 * i )); i=$(( i + 1 )) + run="$RUN_ROOT/$name"; mkdir -p "$run" + resume=""; [ -s "$run/checkpoints/last.ckpt" ] && resume="ckpt_path=$run/checkpoints/last.ckpt" + echo "--- $name -> $run ${resume:+(resuming)}: $args" + srun --exact --ntasks=$G --ntasks-per-node=$G --gres=gpu:h100:$G --gpu-bind=none --cpus-per-task=$CPUS --mem=200G \ + --container-image="$IMAGE" --container-mounts=/workspace:/workspace:rw,/ephemeral:/ephemeral:rw --container-workdir="$SRC" \ + --output="$run/train-%j.log" \ + bash -c "set -euo pipefail + ulimit -n \$(ulimit -Hn) + export HOME=/workspace/users/$USER HF_HOME=/workspace/users/$USER/.cache/huggingface HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 + export EGOVERSE_DATASET_DIR=$DATA EGOVERSE_EPISODE_TABLE=$TABLE EGOVERSE_LOG_DIR=$RUN_ROOT EGOVERSE_CACHE_DIR=$RUN_ROOT/cache + export MASTER_PORT=$port PYTHONPATH=$SRC HYDRA_FULL_ERROR=1 PYTHONUNBUFFERED=1 WANDB_MODE=\${WANDB_MODE:-online} WANDB_PROJECT=\${WANDB_PROJECT:-egoverse} + export OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 + export TORCHINDUCTOR_CACHE_DIR=/tmp/ti_${SLURM_JOB_ID}_$name TRITON_CACHE_DIR=/tmp/tr_${SLURM_JOB_ID}_$name + echo rank=\${SLURM_PROCID:-0} nofile=\$(ulimit -n) commit=\$(git -C $SRC rev-parse --short HEAD) job=$SLURM_JOB_ID restart=${SLURM_RESTART_COUNT:-0} + exec $SRC/.venv/bin/python egomimic/trainHydra.py $args launch_params.gpus_per_node=$G \ + hydra.run.dir=$run paths.output_dir=$run norm_stats.save_cache_dir=$run $resume ${EXTRA:-}" & +done +wait diff --git a/egomimic/scripts/abc_sim/tasks.py b/egomimic/scripts/abc_sim/tasks.py new file mode 100644 index 0000000000..b00d3f04b5 --- /dev/null +++ b/egomimic/scripts/abc_sim/tasks.py @@ -0,0 +1,66 @@ +"""The ABC tasks the sim rollout pipeline runs, by short name. + +``dataset``: the sim_224 ``task_name`` (convert_sim_to_zarr --task); +``sim``: the abc_sim task spec (the rollout env); ``real``: the matching REAL +ABC task on the Phoenix mirror (data=abc_real_joints data.abc_task=...), or +None. The first three are the tasks the ABC paper measured sim-real +correlation on (Pearson r = 0.85 success / 0.91 progress, 12 checkpoints). +The short name is also the zarr folder name under ABC_SIM_ZARR_ROOT. +""" + +from __future__ import annotations + +TASKS: dict[str, dict] = { + "put_bottles": { + "dataset": "sim_put_the_plastic_bottles_in_the_bin", + "sim": "put_plastic_bottles_in_bin", + "real": "put the plastic bottles in the bin", + }, + "dishrack": { + "dataset": "sim_load_the_plates_into_the_dish_rack", + "sim": "load_plates_into_dish_rack", + # real ABC loads mixed dishes; the sim rack holds plates only + "real": "load the mixed dishes into the dish rack", + }, + "turn_mug": { + "dataset": "sim_turn_the_mug_right_side_up", + "sim": "turn_mug_right_side_up", + "real": "turn the mug right side up", + }, + "hang_mug": { + "dataset": "sim_hang_the_mug_on_the_mug_rack", + "sim": "hang_mug_on_mug_rack", + "real": None, + }, + "pour": { + "dataset": "sim_pouring_beads", + "sim": "pouring", + "real": None, + }, + "sweep": { + "dataset": "sim_sweep_away_paper_scraps_from_the_table", + "sim": "sweep_away_paper_scraps_from_table", + "real": "sweep away the paper scraps from the table", + }, + "inhand_transfer": { + "dataset": "sim_inhand_transfer_the_item_to_other_side", + "sim": "inhand_transfer_item_to_other_side", + "real": None, + }, +} + + +def task(short_name: str) -> dict: + try: + return TASKS[short_name] + except KeyError: + raise KeyError( + f"unknown ABC task '{short_name}'; known: {', '.join(sorted(TASKS))}" + ) from None + + +if __name__ == "__main__": + import sys + + # `python -m egomimic.scripts.abc_sim.tasks ` for shell use. + print(task(sys.argv[1])[sys.argv[2]]) diff --git a/tests/fixtures/synthetic_episodes.py b/tests/fixtures/synthetic_episodes.py index a63c8324ae..2183db4266 100644 --- a/tests/fixtures/synthetic_episodes.py +++ b/tests/fixtures/synthetic_episodes.py @@ -52,6 +52,11 @@ def _pose(T: int, rng: np.random.Generator) -> np.ndarray: return np.repeat(one, T, axis=0).astype(np.float64) +def _joints(T: int, rng: np.random.Generator) -> np.ndarray: + """One random 6-joint pose (radians) held constant for all T frames.""" + return np.repeat(rng.uniform(-1.5, 1.5, (1, 6)), T, axis=0).astype(np.float32) + + def _scalar(T: int, rng: np.random.Generator) -> np.ndarray: """One uniform(0,1) value held constant for all T frames, shape (T, 1).""" return np.full((T, 1), rng.uniform(0.0, 1.0)) @@ -93,6 +98,11 @@ def write_episode( "right.obs_gripper": _scalar(T, rng), "left.cmd_gripper": _scalar(T, rng), "right.cmd_gripper": _scalar(T, rng), + # ABC / YAM joint space (Eva mode "joints"): 6 joints per arm. + "left.obs_joints": _joints(T, rng), + "right.obs_joints": _joints(T, rng), + "left.cmd_joints": _joints(T, rng), + "right.cmd_joints": _joints(T, rng), } ) extrinsics = { diff --git a/tests/unit/test_abc_joint_configs.py b/tests/unit/test_abc_joint_configs.py new file mode 100644 index 0000000000..66a11ee50e --- /dev/null +++ b/tests/unit/test_abc_joint_configs.py @@ -0,0 +1,87 @@ +"""The ABC joint-space recipes compose, and the fine-tune recipe trains: two +optimizer steps of a tiny RDT on synthetic eva JOINT episodes through the real +trainHydra.train() path (Eva "joints" keymap + transform, the eva RDT head and +state stem, joint metrics wiring, image history on the robot camera).""" + +from __future__ import annotations + +import math + +import pytest +from fixtures.recipes import LOCAL_RESOLVER, common_overrides, cpu_trainer_overrides +from fixtures.train_harness import BatchKeySpy, compose_recipe, hermetic_env, write_fixtures +from fixtures.recipes import Recipe + +import egomimic.trainHydra as train_hydra +from egomimic.algo.hpt import HPT +from egomimic.rldb.embodiment.eva import JOINT_ACTION_KEY, JOINT_STATE_KEY, Eva + +EMB = "eva_bimanual" +RM = "model.robomimic_model" + + +def _tiny_rdt(emb: str) -> list[str]: + enc = f"{RM}.encoder_specs.front_img_1" + return [ + f"{RM}.width=32", + f"{RM}.trunk.depth=2", + f"{RM}.trunk.num_heads=4", + f"{RM}.head_specs.{emb}.num_inference_steps=2", + f"~{RM}.shared_stem_specs.annotation", + f"{RM}.shared_obs_keys=[front_img_1]", + f"{RM}.annotation_key=null", + f"{enc}.model_name=vit_small_patch16_dinov3", + f"{enc}.image_size=[32,48]", + f"+{enc}.pretrained=false", + f"+{enc}.tower_kwargs={{embed_dim: 32, depth: 1, num_heads: 2}}", + ] + + +@pytest.mark.parametrize( + "top,domains", + [ + ("train_zarr_abc_sim_rdt_ft", ["eva_bimanual"]), + ("train_zarr_abc_real_rdt_ft", ["eva_bimanual"]), + ("train_zarr_abc_sim_rdt_cotrain", ["human_bimanual", "eva_bimanual"]), + ], +) +def test_abc_recipes_compose(compose_resolve, top, domains): + cfg = compose_resolve(top, []) + rm = cfg.model.robomimic_model + assert list(rm.domains) == domains + assert rm.dims.eva_bimanual.action == 14 and rm.dims.eva_bimanual.proprio == 14 + assert rm.ac_keys.eva_bimanual == JOINT_ACTION_KEY + assert "state_joint_positions" in rm.stem_specs.eva_bimanual + assert rm.width == 2048 # the 1B DiT Ryan asked for + assert cfg.data.train_datasets.eva_bimanual.resolver.transform_list.mode == "joints" + assert cfg.data.train_datasets.eva_bimanual.resolver.key_map.image_history_gap_s == 0.1 + assert "eva_bimanual" not in cfg.evaluator.viz_func # joints have no overlay + assert cfg.init_weights_ckpt is None and cfg.train_viz is False + + +def test_abc_sim_finetune_train_step(tmp_path, monkeypatch): + hermetic_env(monkeypatch) + data, out, hashes = write_fixtures(tmp_path, "eva") + recipe = Recipe("train_zarr_abc_sim_rdt_ft", "abc_sim_joints", "rdt_1b_abc_joints_dinov3_vitb", EMB) + cfg = compose_recipe( + recipe, + common_overrides(EMB, data, out, batch_size=2, num_workers=0, episode_hashes=hashes) + + cpu_trainer_overrides(2) + + _tiny_rdt(EMB) + + [ + # the abc_sim_joints resolver is already local; point it at the fixtures + f"data.train_datasets.{EMB}.resolver.folder_path={data}", + f"data.valid_datasets.{EMB}.filters=null", + ], + out, + ) + assert cfg.data.train_datasets[EMB].resolver._target_ == LOCAL_RESOLVER + spy = BatchKeySpy(monkeypatch, HPT) + metrics, objects = train_hydra.train(cfg) + loss = float(metrics["Train/action_loss"]) + assert math.isfinite(loss) + assert objects["trainer"].global_step == 2 + seen = spy.keys[EMB] + expected = {JOINT_ACTION_KEY, JOINT_STATE_KEY, Eva.VIZ_IMAGE_KEY, f"{Eva.VIZ_IMAGE_KEY}_hist", "fps"} + assert expected <= seen, f"missing {expected - seen}; saw {sorted(seen)}" + assert "observations.images.left_wrist_img" not in seen # joints_front keymap diff --git a/tests/unit/test_episode_table_file.py b/tests/unit/test_episode_table_file.py new file mode 100644 index 0000000000..77d92d037d --- /dev/null +++ b/tests/unit/test_episode_table_file.py @@ -0,0 +1,16 @@ +"""EGOVERSE_EPISODE_TABLE: the S3 resolver reads the episode table from a +pickled snapshot instead of SQL, memoized per path.""" + +import pandas as pd + +from egomimic.rldb.zarr import zarr_dataset_multi as zdm + + +def test_episode_table_from_pickle(tmp_path, monkeypatch): + df = pd.DataFrame({"episode_hash": ["a"], "lab": ["abc_sim"]}) + path = tmp_path / "episodes.pkl" + df.to_pickle(path) + monkeypatch.setenv("EGOVERSE_EPISODE_TABLE", str(path)) + monkeypatch.setattr(zdm, "create_default_engine", lambda: (_ for _ in ()).throw(AssertionError("SQL touched"))) + assert list(zdm._episode_table().episode_hash) == ["a"] + assert list(zdm._episode_table().episode_hash) == ["a"] # second call: still no SQL From d054b7e928e0d63a567db2c6b16076bc7995fdf4 Mon Sep 17 00:00:00 2001 From: Aidan Gao Date: Tue, 29 Sep 2026 00:59:20 -0400 Subject: [PATCH 05/20] abc_sim: policy server, sim client, Lambda rollout launcher Closed-loop evaluation of a joint-space checkpoint in ABC's simulator. The sim's Python (py3.12, MuJoCo 3.8, torch 2.11) cannot share the training venv, so the policy is served over HTTP + JSON (stdlib both ends): - policy_server.py builds the model input exactly as ZarrDataset builds a training sample for the checkpoint's data config (same keymap, transform list, the _hist frame image_history_gap_s earlier, fps) and normalizes it with the checkpoint's norm stats before process_batch_for_training -- the dataset normalizes in __getitem__, which the robot rollout skipped. - sim_client.py runs in ABC's venv on abc_minimal's SimTaskEnv with ABC's evaluators and eval defaults (seeded worlds, 15 executed steps per inference, 236 chunks, early stop on success), no RTC; summary.json has ABC's fields. --policy hold is the no-motion floor. - rollout_lambda.sbatch runs both in one container step on one H100. tests/unit/test_policy_server.py: the server's sample equals the dataset's for the same frames, normalization happens before the model, and the HTTP round trip. Co-Authored-By: Claude Fable 5.1 --- egomimic/scripts/abc_sim/README.md | 29 +++ egomimic/scripts/abc_sim/policy_server.py | 222 ++++++++++++++++++ .../scripts/abc_sim/rollout_lambda.sbatch | 43 ++++ egomimic/scripts/abc_sim/sim_client.py | 174 ++++++++++++++ tests/unit/test_policy_server.py | 129 ++++++++++ 5 files changed, 597 insertions(+) create mode 100644 egomimic/scripts/abc_sim/README.md create mode 100644 egomimic/scripts/abc_sim/policy_server.py create mode 100644 egomimic/scripts/abc_sim/rollout_lambda.sbatch create mode 100644 egomimic/scripts/abc_sim/sim_client.py create mode 100644 tests/unit/test_policy_server.py diff --git a/egomimic/scripts/abc_sim/README.md b/egomimic/scripts/abc_sim/README.md new file mode 100644 index 0000000000..4d542ea58c --- /dev/null +++ b/egomimic/scripts/abc_sim/README.md @@ -0,0 +1,29 @@ +# ABC fine-tune + abc_sim rollout pipeline + +Rank pretrained EgoVerse policies by fine-tuning each on one ABC task and +rolling the fine-tune out in ABC's simulator (amazon-far/abc, `abc_sim`): +closed-loop success where offline validation metrics all look alike. The ABC +paper reports Pearson r = 0.85 (success) / 0.91 (progress) between its sim and +real success over 12 checkpoints on bottles, dishrack and mug flip. + +Action space: the ABC station's 14-D absolute joints + gripper (Eva mode +`joints`, `egomimic/rldb/embodiment/eva.py`), which is also abc_sim's action, +so a rollout needs no IK. Tasks: `tasks.py` (short name -> sim_224 dataset +task, abc_sim task, real ABC task). + +Steps (Lambda; paths in the launchers): + +1. Data: `prepare.py --sim-data ` or `curl` the sim_224 tar, then + `convert_sim_to_zarr.py --src --task --out + $ABC_SIM_ZARR_ROOT/` -> eva_bimanual zarrs with the real ABC keys. +2. Base: `train_zarr_fold_ladder_rdt1b` (RDT-1B on the mecka freeform-fold + operator ladder) via `lambda_pack.sbatch`; every 100 epochs a checkpoint. +3. Fine-tune: `train_zarr_abc_sim_rdt_ft init_weights_ckpt=` + (same launcher); `init_weights_ckpt=null` is the from-scratch floor; + `train_zarr_abc_real_rdt_ft` the real-data twin; `train_zarr_abc_sim_rdt_cotrain` + the single-task human+robot cotrain. +4. Rollout: `rollout_lambda.sbatch` = `policy_server.py` (training venv) + + `sim_client.py` (ABC venv, ABC's evaluators and eval defaults, no RTC) -> + `summary.json` in ABC's format. `POLICY=hold` is the no-motion floor. +5. `val_vs_sim.py manifest.json --metric ` -> Spearman / Pearson with + bootstrap CIs and the log-linear fit of the ladder. diff --git a/egomimic/scripts/abc_sim/policy_server.py b/egomimic/scripts/abc_sim/policy_server.py new file mode 100644 index 0000000000..09576ff26d --- /dev/null +++ b/egomimic/scripts/abc_sim/policy_server.py @@ -0,0 +1,222 @@ +"""Serve an EgoVerse joint-space checkpoint over HTTP for the abc_sim client. + +The simulator's Python (amazon-far/abc: py3.12, MuJoCo 3.8, torch 2.11) cannot +share the training venv, so the policy runs here, in the training venv, and +the sim talks to it over plain HTTP + JSON (stdlib on both ends). + + python -m egomimic.scripts.abc_sim.policy_server --ckpt --port 8765 + +POST /infer {"episode": str, "t": float seconds, "state": [14], "prompt": str, + "images": {"top": {"shape": [H, W, 3], "b64": }, ...}} + -> {"actions": [[14] x 100], "dt": 1/30} +GET /health -> {"ok": true, "ckpt": ..., "mode": "joints"} + +The inputs are built EXACTLY as ZarrDataset builds a training sample for the +checkpoint's data config -- the same Eva keymap and transform list, then the +checkpoint's own norm stats (the dataset normalizes samples in __getitem__; +``process_batch_for_training`` does not, which the old robot rollout missed). +The ``_hist`` camera is the frame ``image_history_gap_s`` earlier, kept per +episode; at an episode start it is the current frame, as in training. +""" + +from __future__ import annotations + +import argparse +import base64 +import json +import logging +import threading +from collections import deque +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +import numpy as np +import torch +from omegaconf import OmegaConf +from torch.utils.data import default_collate + +from egomimic.rldb.embodiment.embodiment import IMAGE_HISTORY_SUFFIX, get_embodiment_id +from egomimic.rldb.embodiment.eva import ( + JOINT_ACTION_KEY, + JOINT_MODE, + JOINT_RAW_HORIZON, + Eva, +) + +log = logging.getLogger("abc_sim.policy_server") + +EMB = "eva_bimanual" +FPS = 30.0 +# abc_sim camera -> the zarr key the converter wrote it to +CAMERA_ZARR = {"top": "images.front_1", "left": "images.left_wrist", "right": "images.right_wrist"} + + +def decode_image(spec: dict) -> np.ndarray: + """{"shape": [H, W, 3], "b64": raw uint8} -> (3, H, W) float32 in [0, 1], RGB.""" + h, w, c = spec["shape"] + raw = np.frombuffer(base64.b64decode(spec["b64"]), dtype=np.uint8).reshape(h, w, c) + return np.moveaxis(raw, -1, -3).astype(np.float32) / np.float32(255.0) + + +def split_state(state) -> dict[str, np.ndarray]: + """ABC 14-D state -> the zarr proprio keys of Eva's joint keymap.""" + s = np.asarray(state, dtype=np.float32).reshape(14) + return { + "left.obs_joints": s[0:6], + "left.obs_gripper": s[6:7], + "right.obs_joints": s[7:13], + "right.obs_gripper": s[13:14], + } + + +class JointPolicy: + """The checkpoint plus the data pipeline it was trained with.""" + + def __init__(self, ckpt_path: str, device: str = "cuda"): + from egomimic.pl_utils.pl_model import ModelWrapper + + self.ckpt_path = ckpt_path + self.device = torch.device(device if torch.cuda.is_available() else "cpu") + wrapper = ModelWrapper.load_from_checkpoint(ckpt_path, weights_only=False, map_location="cpu") + self.algo = wrapper.model.to(self.device).eval() + self.algo.device = self.device + cfg = wrapper._as_config(getattr(wrapper.hparams, "config_tree", None)) + self.configure(cfg) + + def configure(self, cfg) -> None: + """Keymap + transform list from the checkpoint's data config.""" + resolver = OmegaConf.select(cfg, f"data.train_datasets.{EMB}.resolver") + if resolver is None: + raise ValueError(f"checkpoint config has no data.train_datasets.{EMB}") + km = {k: v for k, v in OmegaConf.to_container(resolver.key_map, resolve=True).items() if k != "_target_"} + tl = {k: v for k, v in OmegaConf.to_container(resolver.transform_list, resolve=True).items() if k != "_target_"} + if tl.get("mode") != JOINT_MODE: + raise ValueError(f"policy_server serves Eva mode '{JOINT_MODE}', checkpoint has {tl.get('mode')!r}") + self.key_map = Eva.get_keymap(**{k: v for k, v in km.items() if k != "annotation_key"}) + self.annotation_key = km.get("annotation_key") + self.transforms = Eva.get_transform_list(**tl) + self.gap_s = km.get("image_history_gap_s") + self.cameras = { + k: v["zarr_key"] + for k, v in self.key_map.items() + if v.get("key_type") == "camera_keys" and not k.endswith(IMAGE_HISTORY_SUFFIX) + } + self.emb_id = get_embodiment_id(EMB) + self.history: dict[str, deque] = {} + self.lock = threading.Lock() + + def _past_frame(self, episode: str, t: float, frame: np.ndarray) -> np.ndarray: + """The front frame >= gap_s before t, else the oldest one this episode + (an episode start: the current frame), as ZarrDataset clamps to frame 0.""" + hist = self.history.setdefault(episode, deque(maxlen=int(FPS * 5))) + hist.append((t, frame)) + past = [f for (tt, f) in hist if tt <= t - self.gap_s] + return past[-1] if past else hist[0][1] + + def build_sample(self, episode: str, t: float, state, images: dict[str, np.ndarray]) -> dict: + """One un-normalized sample with the dataset's keys (the 45-step cmd + chunk is faked from the current pose: only its shape reaches the model).""" + raw = split_state(state) + raw.update({f"{arm}.cmd_{part}": np.tile(raw[f"{arm}.obs_{part}"], (JOINT_RAW_HORIZON, 1)) + for arm in ("left", "right") for part in ("joints", "gripper")}) + sample = {} + for key, spec in self.key_map.items(): + zk = spec["zarr_key"] + if spec.get("key_type") == "camera_keys": + cam = next(c for c, z in CAMERA_ZARR.items() if z == zk) + if cam not in images: + raise ValueError(f"camera '{cam}' missing from the request (have {sorted(images)})") + frame = images[cam] + if key.endswith(IMAGE_HISTORY_SUFFIX): + frame = self._past_frame(episode, t, images[cam]) + sample[key] = frame + elif spec.get("key_type") in ("proprio_keys", "action_keys"): + sample[key] = raw[zk] + for tf in self.transforms: + sample = tf.transform(sample) + for k, v in sample.items(): + if isinstance(v, np.ndarray): + sample[k] = torch.from_numpy(v).to(torch.float32) + sample["fps"] = torch.tensor(FPS, dtype=torch.float32) + return sample + + @torch.no_grad() + def infer(self, episode: str, t: float, state, images: dict[str, np.ndarray], prompt: str) -> np.ndarray: + with self.lock: + sample = self.build_sample(episode, t, state, images) + sample = self.algo.norm_stats.normalize(sample, self.emb_id) + batch = default_collate([{k: v for k, v in sample.items() if isinstance(v, torch.Tensor)}]) + if self.annotation_key: + batch[self.annotation_key] = [[prompt]] + processed = self.algo.process_batch_for_training({EMB: batch}) + preds = self.algo.forward_eval(processed) + return preds[f"{EMB}_{JOINT_ACTION_KEY}"][0].float().cpu().numpy() + + def reset(self, episode: str) -> None: + with self.lock: + self.history.pop(episode, None) + + +def make_handler(policy: JointPolicy): + class Handler(BaseHTTPRequestHandler): + def _send(self, code: int, payload: dict) -> None: + body = json.dumps(payload).encode() + self.send_response(code) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def do_GET(self): + if self.path == "/health": + self._send(200, {"ok": True, "ckpt": policy.ckpt_path, "mode": JOINT_MODE}) + else: + self._send(404, {"error": self.path}) + + def do_POST(self): + req = json.loads(self.rfile.read(int(self.headers.get("Content-Length", 0)))) + try: + if self.path == "/reset": + policy.reset(req["episode"]) + self._send(200, {"ok": True}) + elif self.path == "/infer": + images = {k: decode_image(v) for k, v in req["images"].items()} + actions = policy.infer(req["episode"], float(req["t"]), req["state"], images, req.get("prompt", "")) + self._send(200, {"actions": actions.tolist(), "dt": 1.0 / FPS}) + else: + self._send(404, {"error": self.path}) + except Exception as e: # the client must see the reason, not a hang + log.exception("request failed") + self._send(500, {"error": f"{type(e).__name__}: {e}"}) + + def log_message(self, *a): # quiet + pass + + return Handler + + +def serve(policy: JointPolicy, host: str, port: int) -> ThreadingHTTPServer: + server = ThreadingHTTPServer((host, port), make_handler(policy)) + threading.Thread(target=server.serve_forever, daemon=True).start() + return server + + +def main(argv=None) -> int: + p = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + p.add_argument("--ckpt", required=True) + p.add_argument("--host", default="127.0.0.1") + p.add_argument("--port", type=int, default=8765) + p.add_argument("--device", default="cuda") + a = p.parse_args(argv) + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s") + policy = JointPolicy(a.ckpt, a.device) + server = serve(policy, a.host, a.port) + log.info("serving %s on http://%s:%d (cameras %s, gap %s s)", a.ckpt, a.host, a.port, sorted(policy.cameras), policy.gap_s) + try: + threading.Event().wait() + except KeyboardInterrupt: + server.shutdown() + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/egomimic/scripts/abc_sim/rollout_lambda.sbatch b/egomimic/scripts/abc_sim/rollout_lambda.sbatch new file mode 100644 index 0000000000..3a9b729561 --- /dev/null +++ b/egomimic/scripts/abc_sim/rollout_lambda.sbatch @@ -0,0 +1,43 @@ +#!/bin/bash +# Roll one checkpoint out in abc_sim on one Lambda H100: policy_server.py in the +# training venv (pytorch image) and sim_client.py in ABC's venv, both in one +# container step so they share the GPU (mjwarp renders on it too). +# +# CKPT= TASK=put_bottles OUT= sbatch --qos=loaner-normal \ +# egomimic/scripts/abc_sim/rollout_lambda.sbatch +# NUM_WORLDS=50 (default), POLICY=hold for the no-motion floor (no server) +#SBATCH --job-name=abc_rollout +#SBATCH --account=loaner +#SBATCH --partition=loaner +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --gres=gpu:h100:1 +#SBATCH --cpus-per-task=16 +#SBATCH --mem=120G +#SBATCH --time=06:00:00 +#SBATCH --output=/workspace/users/%u/logs/%x-%j.out +set -euo pipefail +SRC=${SRC:-/workspace/users/$USER/EgoVerse} +ABC=${ABC:-/workspace/users/$USER/abc} +IMAGE=${IMAGE:-/workspace/users/$USER/images/nvidia-pytorch-25.06-py3.sqsh} +TASK=${TASK:-put_bottles} +: "${OUT:?set OUT}" +POLICY=${POLICY:-served} +[ "$POLICY" = hold ] || : "${CKPT:?set CKPT}" +SIM_TASK=$($SRC/.venv/bin/python -m egomimic.scripts.abc_sim.tasks "$TASK" sim 2>/dev/null || python3 "$SRC/egomimic/scripts/abc_sim/tasks.py" "$TASK" sim) +PORT=$(( 8700 + SLURM_JOB_ID % 200 )) +mkdir -p "$OUT" +echo "=== abc_rollout job=$SLURM_JOB_ID node=$(hostname) task=$TASK/$SIM_TASK ckpt=${CKPT:-none} policy=$POLICY out=$OUT" +srun --ntasks=1 --container-image="$IMAGE" --container-mounts=/workspace:/workspace:rw --container-workdir="$SRC" \ + bash -c "set -euo pipefail + export HOME=/workspace/users/$USER HF_HOME=/workspace/users/$USER/.cache/huggingface HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 + export PYTHONPATH=$SRC PYTHONUNBUFFERED=1 OMP_NUM_THREADS=4 + if [ '$POLICY' = served ]; then + $SRC/.venv/bin/python -m egomimic.scripts.abc_sim.policy_server --ckpt '$CKPT' --port $PORT > '$OUT/server.log' 2>&1 & + for i in \$(seq 1 120); do curl -fsS http://127.0.0.1:$PORT/health > /dev/null 2>&1 && break; sleep 5; done + curl -fsS http://127.0.0.1:$PORT/health || { echo 'server never came up'; tail -50 '$OUT/server.log'; exit 1; } + fi + cd $ABC + MUJOCO_EGL_DEVICE_ID=0 .venv/bin/python $SRC/egomimic/scripts/abc_sim/sim_client.py --task '$SIM_TASK' \ + --server http://127.0.0.1:$PORT --policy '$POLICY' --out '$OUT' --num-worlds ${NUM_WORLDS:-50} ${CLIENT_EXTRA:-} + echo '{\"job\": '$SLURM_JOB_ID', \"ckpt\": \"${CKPT:-}\", \"task\": \"$TASK\", \"commit\": \"'\$(git -C $SRC rev-parse --short HEAD)'\"}' > '$OUT/run_info.json'" diff --git a/egomimic/scripts/abc_sim/sim_client.py b/egomimic/scripts/abc_sim/sim_client.py new file mode 100644 index 0000000000..8d65d386ac --- /dev/null +++ b/egomimic/scripts/abc_sim/sim_client.py @@ -0,0 +1,174 @@ +"""Roll a served EgoVerse policy out in abc_sim (amazon-far/abc) and score it +with ABC's own task evaluators. Runs in ABC's venv: stdlib + numpy + abc_minimal +only (no egomimic import; the policy lives behind policy_server.py). + + /bin/python egomimic/scripts/abc_sim/sim_client.py \ + --task put_plastic_bottles_in_bin --server http://127.0.0.1:8765 \ + --num-worlds 50 --out + +Protocol = ABC's SimEvalConfig defaults without RTC: worlds seeded seed+i, +one inference per 15 executed steps (execute_chunk_dim) out of the policy's +100-step chunk, 236 chunks max (~118 s at 30 Hz), early stop on success +(rollout_over), cameras top/left/right at 168x224. summary.json carries the +same fields ABC's build_summary writes (success_rate, num_success, +mean_reward, mean_max_progress, worlds[]) so the two are comparable side by +side. ``--policy hold`` needs no server: it repeats the current state (the +no-motion floor and a protocol smoke). +""" + +from __future__ import annotations + +import argparse +import base64 +import json +import time +import urllib.request +from pathlib import Path + +import numpy as np + +from abc_minimal.eval_policy import jsonable, rollout_over, video_frame +from abc_minimal.sim_env import SimTaskEnv, task_prompt + +CAMERAS = ("top", "left", "right") + + +class ServedPolicy: + def __init__(self, server: str): + self.server = server.rstrip("/") + with urllib.request.urlopen(f"{self.server}/health", timeout=30) as r: + self.health = json.load(r) + + def _post(self, path: str, payload: dict) -> dict: + req = urllib.request.Request( + f"{self.server}{path}", data=json.dumps(payload).encode(), + headers={"Content-Type": "application/json"}, method="POST", + ) + with urllib.request.urlopen(req, timeout=120) as r: + return json.load(r) + + def reset(self, episode: str) -> None: + self._post("/reset", {"episode": episode}) + + def infer(self, episode: str, t: float, obs: dict) -> np.ndarray: + images = {} + for cam, img in obs["images"].items(): + hwc = np.ascontiguousarray(np.moveaxis(np.asarray(img, dtype=np.uint8), 0, -1)) # (3,H,W)->(H,W,3) + images[cam] = {"shape": list(hwc.shape), "b64": base64.b64encode(hwc.tobytes()).decode()} + out = self._post("/infer", { + "episode": episode, "t": t, "state": np.asarray(obs["state"], dtype=np.float32).tolist(), + "prompt": obs["prompt"], "images": images, + }) + return np.asarray(out["actions"], dtype=np.float32) + + +class HoldPolicy: + health = {"ckpt": "hold", "mode": "hold"} + + def reset(self, episode: str) -> None: + pass + + def infer(self, episode: str, t: float, obs: dict) -> np.ndarray: + return np.tile(np.asarray(obs["state"], dtype=np.float32)[None], (100, 1)) + + +def rollout(env: SimTaskEnv, policy, *, seed: int, num_chunks: int, execute: int, video_path=None) -> dict: + episode = f"seed{seed}" + policy.reset(episode) + obs = env.reset(seed=seed) + video = None + if video_path is not None: + import imageio.v2 as imageio + + video = imageio.get_writer(str(video_path), fps=30, macro_block_size=1) + video.append_data(video_frame(obs["images"], CAMERAS)) + result = env.evaluate() + max_reward = float(result.get("reward", 0.0)) + steps, infer_s, t0 = 0, [], time.perf_counter() + try: + for _ in range(num_chunks): + t_inf = time.perf_counter() + actions = policy.infer(episode, steps / 30.0, obs) + infer_s.append(time.perf_counter() - t_inf) + for action in actions[:execute]: + env.step_one(action) + result = env.evaluate() + max_reward = max(max_reward, float(result.get("reward", 0.0))) + steps += 1 + if video is not None: + video.append_data(video_frame(env.render_cameras(), CAMERAS)) + if rollout_over(result): + break + if rollout_over(result): + break + obs = env.obs() + finally: + if video is not None: + video.close() + return { + "world_seed": seed, + "success": bool(result["ever_success"]), + "final_success": bool(result["success"]), + "reward": float(result["reward"]), + "max_reward": max_reward, + "steps": steps, + "wall_s": time.perf_counter() - t0, + "infer_s_mean": float(np.mean(infer_s)) if infer_s else None, + "randomization": env.randomization, + "final_task_eval": result, + "video_path": str(video_path) if video_path is not None else None, + } + + +def main(argv=None) -> int: + p = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + p.add_argument("--task", required=True, help="abc_sim task name / alias (e.g. put_plastic_bottles_in_bin)") + p.add_argument("--server", default="http://127.0.0.1:8765") + p.add_argument("--policy", choices=["served", "hold"], default="served") + p.add_argument("--out", type=Path, required=True) + p.add_argument("--num-worlds", type=int, default=50) + p.add_argument("--seed", type=int, default=20260511, help="ABC's SimEvalConfig default") + p.add_argument("--num-chunks", type=int, default=236) + p.add_argument("--execute", type=int, default=15, help="steps executed per inference (ABC's execute_chunk_dim)") + p.add_argument("--camera-backend", default="mjwarp") + p.add_argument("--gpu-id", type=int, default=None) + p.add_argument("--save-videos", type=int, default=2, help="record the first N worlds") + a = p.parse_args(argv) + a.out.mkdir(parents=True, exist_ok=True) + policy = HoldPolicy() if a.policy == "hold" else ServedPolicy(a.server) + env = SimTaskEnv(task=a.task, height=168, width=224, camera_keys=CAMERAS, prompt=task_prompt(a.task), + camera_backend=a.camera_backend, gpu_id=a.gpu_id) + worlds = [] + try: + for i in range(a.num_worlds): + video = a.out / f"world_{i:03d}.mp4" if i < a.save_videos else None + w = rollout(env, policy, seed=a.seed + i, num_chunks=a.num_chunks, execute=a.execute, video_path=video) + w["world_index"] = i + worlds.append(w) + print(f"world={i:03d} success={w['success']} max_progress={w['max_reward']:.3f} steps={w['steps']} " + f"infer={1000 * (w['infer_s_mean'] or 0):.0f}ms", flush=True) + (a.out / "worlds.jsonl").open("a").write(json.dumps(jsonable(w)) + "\n") + finally: + env.close() + succ = np.array([w["success"] for w in worlds], dtype=bool) + summary = { + "format": "egoverse_abc_sim_rollout/v1", + "policy": policy.health, + "task": env.spec.name, + "prompt": env.prompt, + "config": {k: (str(v) if isinstance(v, Path) else v) for k, v in vars(a).items()}, + "success_rate": float(succ.mean()) if len(succ) else None, + "num_success": int(succ.sum()), + "num_worlds": len(worlds), + "mean_reward": float(np.mean([w["reward"] for w in worlds])) if worlds else None, + "mean_max_progress": float(np.mean([w["max_reward"] for w in worlds])) if worlds else None, + "worlds": worlds, + } + (a.out / "summary.json").write_text(json.dumps(jsonable(summary), indent=2, sort_keys=True)) + print(f"summary: success_rate={summary['success_rate']} num_success={summary['num_success']}/{len(worlds)} " + f"mean_max_progress={summary['mean_max_progress']}", flush=True) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/unit/test_policy_server.py b/tests/unit/test_policy_server.py new file mode 100644 index 0000000000..f41a422aa6 --- /dev/null +++ b/tests/unit/test_policy_server.py @@ -0,0 +1,129 @@ +"""policy_server builds the model input exactly as ZarrDataset builds a +training sample (same keymap, transforms, history frame, fps), normalizes it +with the checkpoint's norm stats, and speaks the sim client's HTTP protocol.""" + +from __future__ import annotations + +import base64 +import json +import urllib.request + +import numpy as np +import torch +from fixtures.synthetic_episodes import write_episode +from omegaconf import OmegaConf + +from egomimic.rldb.embodiment.eva import JOINT_ACTION_KEY, JOINT_STATE_KEY, Eva +from egomimic.rldb.zarr.zarr_dataset_multi import LocalEpisodeResolver, MultiDataset +from egomimic.scripts.abc_sim import policy_server as ps + +CAM = Eva.VIZ_IMAGE_KEY +CFG = OmegaConf.create({"data": {"train_datasets": {"eva_bimanual": {"resolver": { + "key_map": {"_target_": "x", "keymap_mode": "joints_front", "annotation_key": "annotations", "image_history_gap_s": 0.1}, + "transform_list": {"_target_": "x", "mode": "joints"}, +}}}}}) + + +class _FakeNormStats: + """Scales the state by 2 so the test can see normalization happened.""" + + def normalize(self, data, emb_id): + out = dict(data) + out[JOINT_STATE_KEY] = data[JOINT_STATE_KEY] * 2 + return out + + +class _FakeAlgo: + """Records the batch it was given; predicts the (normalized) state tiled.""" + + norm_stats = _FakeNormStats() + annotation_key = "annotations" + + def process_batch_for_training(self, batch): + self.seen = batch["eva_bimanual"] + return batch + + def forward_eval(self, batch): + state = batch["eva_bimanual"][JOINT_STATE_KEY] + return {f"eva_bimanual_{JOINT_ACTION_KEY}": state[:, None, :].repeat(1, 100, 1)} + + +def _policy(): + pol = ps.JointPolicy.__new__(ps.JointPolicy) + pol.ckpt_path, pol.algo = "fake.ckpt", _FakeAlgo() + pol.configure(CFG) + return pol + + +def _episode_frames(root): + """Raw frames/states of the synthetic eva episode as the sim would send + them, read through the dataset's own reader and JPEG decode.""" + import simplejpeg + + from egomimic.rldb.zarr.zarr_dataset_multi import ZarrEpisode + + ep = ZarrEpisode(next(root.glob("*.zarr"))) + n = ep.metadata["total_frames"] + keys = ["images.front_1", "left.obs_joints", "left.obs_gripper", "right.obs_joints", "right.obs_gripper"] + got = ep.read_intervals({k: [(0, n)] for k in keys}) + rows = {k: got[k][0][1] for k in keys} + frames = [simplejpeg.decode_jpeg(rows["images.front_1"][i], colorspace="RGB") for i in range(n)] + state = np.concatenate([np.asarray(rows[k], dtype=np.float32) for k in keys[1:]], axis=1) + return frames, state + + +def test_server_sample_matches_the_dataset_sample(tmp_path): + write_episode(tmp_path, "eva", T=48) + frames, state = _episode_frames(tmp_path) + key_map = Eva.get_keymap("joints_front", image_history_gap_s=0.1) + resolver = LocalEpisodeResolver(tmp_path, key_map=key_map, transform_list=Eva.get_transform_list("joints")) + leaf = next(iter(MultiDataset._from_resolver(resolver, mode="total").datasets.values())) + pol = _policy() + for t in range(0, 12): # feed the frames in order so the history fills like an episode + sample = pol.build_sample("ep", t / 30.0, state[t], {"top": np.moveaxis(frames[t], -1, 0).astype(np.float32) / 255}) + ref = leaf[t] + assert torch.equal(sample[CAM], ref[CAM]), t + assert torch.equal(sample[f"{CAM}_hist"], ref[f"{CAM}_hist"]), t # 3 frames back, clamped at the start + assert torch.allclose(sample[JOINT_STATE_KEY], ref[JOINT_STATE_KEY].float()), t + assert sample[JOINT_ACTION_KEY].shape == ref[JOINT_ACTION_KEY].shape == (100, 14) + assert sample["fps"].item() == ref["fps"].item() == 30 + assert "observations.images.left_wrist_img" not in sample + + +def test_infer_normalizes_then_returns_unnormalized_chunk(tmp_path): + pol = _policy() + img = np.zeros((3, 16, 24), np.float32) + state = np.arange(14, dtype=np.float32) + out = pol.infer("ep", 0.0, state, {"top": img}, "sim put the bottles in the bin") + assert out.shape == (100, 14) + seen = pol.algo.seen + assert torch.allclose(seen[JOINT_STATE_KEY][0], torch.arange(14.0) * 2) # normalized before the model + assert seen["annotations"] == [["sim put the bottles in the bin"]] + assert seen["fps"].shape == (1,) + + +def test_http_protocol_round_trip(): + pol = _policy() + server = ps.serve(pol, "127.0.0.1", 0) + port = server.server_address[1] + try: + with urllib.request.urlopen(f"http://127.0.0.1:{port}/health") as r: + assert json.load(r)["mode"] == "joints" + hwc = np.random.default_rng(0).integers(0, 255, (16, 24, 3), np.uint8) + payload = {"episode": "e", "t": 0.0, "state": list(range(14)), "prompt": "sim x", + "images": {"top": {"shape": [16, 24, 3], "b64": base64.b64encode(hwc.tobytes()).decode()}}} + req = urllib.request.Request(f"http://127.0.0.1:{port}/infer", data=json.dumps(payload).encode(), + headers={"Content-Type": "application/json"}, method="POST") + with urllib.request.urlopen(req) as r: + out = json.load(r) + assert np.asarray(out["actions"]).shape == (100, 14) and out["dt"] == 1 / 30 + assert torch.allclose(pol.algo.seen[CAM][0], torch.from_numpy(np.moveaxis(hwc, -1, 0).astype(np.float32) / 255)) + bad = urllib.request.Request(f"http://127.0.0.1:{port}/infer", data=json.dumps({"episode": "e", "t": 0, "state": [0], "images": {}}).encode(), + headers={"Content-Type": "application/json"}, method="POST") + try: + urllib.request.urlopen(bad) + assert False, "expected 500" + except urllib.error.HTTPError as e: + assert e.code == 500 and "error" in json.load(e) + finally: + server.shutdown() From 31ccd9553134b5fd80b28fc386e331985c21cb12 Mon Sep 17 00:00:00 2001 From: Aidan Gao Date: Tue, 29 Sep 2026 00:59:21 -0400 Subject: [PATCH 06/20] abc_sim: val-metric vs sim-success report val_vs_sim.py: from a manifest of checkpoints (W&B run, epoch, ladder hours, rollout summary.json), Spearman / Pearson between the validation metric and success_rate / mean_max_progress with bootstrap CIs, and the log-linear fit of the metric (and success) against data hours across the ladder's final checkpoints. report.json + report.md. Co-Authored-By: Claude Fable 5.1 --- egomimic/scripts/abc_sim/val_vs_sim.py | 129 +++++++++++++++++++++++++ tests/unit/test_val_vs_sim.py | 28 ++++++ 2 files changed, 157 insertions(+) create mode 100644 egomimic/scripts/abc_sim/val_vs_sim.py create mode 100644 tests/unit/test_val_vs_sim.py diff --git a/egomimic/scripts/abc_sim/val_vs_sim.py b/egomimic/scripts/abc_sim/val_vs_sim.py new file mode 100644 index 0000000000..9af8947ffb --- /dev/null +++ b/egomimic/scripts/abc_sim/val_vs_sim.py @@ -0,0 +1,129 @@ +"""Does the offline validation metric predict closed-loop sim success? + +Input: a JSON manifest of checkpoints, each with the W&B run its validation +metrics live in, its epoch, its data hours (the ladder rung) and the +sim_client summary.json of its fine-tune's rollout: + + [{"name": "k4_ep400", "wandb": "rl2-group/egoverse/", "epoch": 400, + "hours": 1.096, "summary": ".../summary.json"}, ...] + + python -m egomimic.scripts.abc_sim.val_vs_sim manifest.json --metric \ + "seen_op_valid/Valid/human_bimanual_actions_cartesian_paired_mse_avg" --out report/ + +Writes report.json + report.md: per-checkpoint rows, Spearman and Pearson +between the val metric and success_rate / mean_max_progress with 2000-sample +bootstrap CIs, and the log-linear fit of the val metric against data hours +(one point per rung at its final epoch) with its slope and R^2. +""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path + +import numpy as np + + +def wandb_metric(run_path: str, metric: str, epoch: int) -> float: + """The metric's value at (the last logged step of) ``epoch``.""" + import wandb + + run = wandb.Api().run(run_path) + rows = [r for r in run.scan_history(keys=[metric, "epoch"]) if r.get(metric) is not None] + at = [r for r in rows if r.get("epoch") is not None and int(r["epoch"]) <= epoch] + if not at: + raise ValueError(f"{run_path}: no {metric} at epoch <= {epoch}") + return float(at[-1][metric]) + + +def _rank(x): + order = np.argsort(x) + r = np.empty(len(x)) + r[order] = np.arange(len(x)) + return r + + +def corr(x, y, kind: str) -> float: + x, y = np.asarray(x, float), np.asarray(y, float) + if kind == "spearman": + x, y = _rank(x), _rank(y) + if x.std() == 0 or y.std() == 0: + return float("nan") + return float(np.corrcoef(x, y)[0, 1]) + + +def bootstrap(x, y, kind: str, n: int = 2000, seed: int = 0) -> tuple[float, float]: + rng = np.random.default_rng(seed) + x, y = np.asarray(x, float), np.asarray(y, float) + vals = [] + for _ in range(n): + i = rng.integers(0, len(x), len(x)) + vals.append(corr(x[i], y[i], kind)) + vals = np.array([v for v in vals if not math.isnan(v)]) + return (float(np.percentile(vals, 2.5)), float(np.percentile(vals, 97.5))) if len(vals) else (float("nan"),) * 2 + + +def loglinear(hours, metric) -> dict: + """metric = a + b * log10(hours): slope, intercept, R^2.""" + x, y = np.log10(np.asarray(hours, float)), np.asarray(metric, float) + if len(x) < 2: + return {"slope": float("nan"), "intercept": float("nan"), "r2": float("nan")} + b, a = np.polyfit(x, y, 1) + pred = a + b * x + ss_res, ss_tot = float(((y - pred) ** 2).sum()), float(((y - y.mean()) ** 2).sum()) + return {"slope": float(b), "intercept": float(a), "r2": 1 - ss_res / ss_tot if ss_tot else float("nan")} + + +def main(argv=None) -> int: + p = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + p.add_argument("manifest", type=Path) + p.add_argument("--metric", required=True) + p.add_argument("--out", type=Path, required=True) + p.add_argument("--no-wandb", action="store_true", help="manifest rows already carry 'val'") + a = p.parse_args(argv) + rows = json.loads(a.manifest.read_text()) + for r in rows: + if "val" not in r or not a.no_wandb: + r["val"] = wandb_metric(r["wandb"], a.metric, int(r["epoch"])) + s = json.loads(Path(r["summary"]).read_text()) + r["success_rate"], r["mean_max_progress"], r["num_worlds"] = s["success_rate"], s.get("mean_max_progress"), s["num_worlds"] + val = [r["val"] for r in rows] + out = {"metric": a.metric, "rows": rows, "n": len(rows), "correlations": {}} + for target in ("success_rate", "mean_max_progress"): + y = [r[target] for r in rows] + if any(v is None for v in y): + continue + for kind in ("spearman", "pearson"): + lo, hi = bootstrap(val, y, kind) + out["correlations"][f"{kind}_{target}"] = {"r": corr(val, y, kind), "ci95": [lo, hi]} + # log-linear in data: the final epoch of each rung + final = {} + for r in rows: + if "hours" in r and (r["hours"] not in final or r["epoch"] > final[r["hours"]]["epoch"]): + final[r["hours"]] = r + if len(final) >= 2: + hs = sorted(final) + out["loglinear_val_vs_hours"] = {"hours": hs, "val": [final[h]["val"] for h in hs], **loglinear(hs, [final[h]["val"] for h in hs])} + out["loglinear_success_vs_hours"] = {"hours": hs, "success": [final[h]["success_rate"] for h in hs], **loglinear(hs, [final[h]["success_rate"] for h in hs])} + a.out.mkdir(parents=True, exist_ok=True) + (a.out / "report.json").write_text(json.dumps(out, indent=2)) + md = [f"# val metric vs sim success ({len(rows)} checkpoints)", "", f"metric: `{a.metric}`", "", + "| name | hours | epoch | val | success | max progress | worlds |", "|---|---|---|---|---|---|---|"] + for r in rows: + prog = "" if r["mean_max_progress"] is None else f"{r['mean_max_progress']:.3f}" + md.append(f"| {r['name']} | {r.get('hours', '')} | {r['epoch']} | {r['val']:.4g} | {r['success_rate']:.3f} | {prog} | {r['num_worlds']} |") + md += ["", "| correlation | r | 95% CI |", "|---|---|---|"] + md += [f"| {k} | {v['r']:.3f} | [{v['ci95'][0]:.3f}, {v['ci95'][1]:.3f}] |" for k, v in out["correlations"].items()] + for key in ("loglinear_val_vs_hours", "loglinear_success_vs_hours"): + if key in out: + f = out[key] + md += ["", f"{key}: slope {f['slope']:.4g} per decade, R^2 {f['r2']:.3f} over hours {f['hours']}"] + (a.out / "report.md").write_text("\n".join(md) + "\n") + print("\n".join(md)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/unit/test_val_vs_sim.py b/tests/unit/test_val_vs_sim.py new file mode 100644 index 0000000000..75144a474b --- /dev/null +++ b/tests/unit/test_val_vs_sim.py @@ -0,0 +1,28 @@ +"""val_vs_sim: correlations and the log-linear fit from a manifest with +pre-filled val metrics (no W&B).""" + +import json + +from egomimic.scripts.abc_sim import val_vs_sim as vs + + +def test_report_from_manifest(tmp_path): + rows = [] + for i, (h, ep, val, succ) in enumerate([(0.237, 1000, 0.9, 0.1), (0.559, 1000, 0.6, 0.4), (1.096, 1000, 0.3, 0.7), (1.096, 500, 0.5, 0.5)]): + s = tmp_path / f"s{i}.json" + s.write_text(json.dumps({"success_rate": succ, "mean_max_progress": succ + 0.1, "num_worlds": 50})) + rows.append({"name": f"r{i}", "epoch": ep, "hours": h, "val": val, "summary": str(s)}) + m = tmp_path / "m.json" + m.write_text(json.dumps(rows)) + assert vs.main([str(m), "--metric", "x", "--out", str(tmp_path / "out"), "--no-wandb"]) == 0 + rep = json.loads((tmp_path / "out" / "report.json").read_text()) + assert rep["correlations"]["spearman_success_rate"]["r"] == -1.0 # lower val mse, higher success + assert rep["loglinear_val_vs_hours"]["r2"] > 0.95 and rep["loglinear_val_vs_hours"]["slope"] < 0 + assert (tmp_path / "out" / "report.md").read_text().count("| r") >= 4 + + +def test_corr_helpers(): + assert vs.corr([1, 2, 3], [3, 2, 1], "spearman") == -1.0 + assert vs.corr([1, 1, 1], [1, 2, 3], "pearson") != vs.corr([1, 1, 1], [1, 2, 3], "pearson") # nan + lo, hi = vs.bootstrap([1, 2, 3, 4], [1, 2, 3, 4], "pearson", n=50) + assert lo <= 1.0 <= hi From 7f13aea194e173010c77e967ee520afbcfc79870 Mon Sep 17 00:00:00 2001 From: Aidan Gao Date: Tue, 29 Sep 2026 01:00:55 -0400 Subject: [PATCH 07/20] policy_server: the history lag in frames, as the dataset reads it ZarrDataset reads the _hist frame round(gap_s * fps) frames back; the server compared timestamps in seconds and, at 1/30 s resolution, could pick one frame too far back. Same frame arithmetic now; the equality test covers t = 0..11. Co-Authored-By: Claude Fable 5.1 --- egomimic/scripts/abc_sim/policy_server.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/egomimic/scripts/abc_sim/policy_server.py b/egomimic/scripts/abc_sim/policy_server.py index 09576ff26d..a2e03481f7 100644 --- a/egomimic/scripts/abc_sim/policy_server.py +++ b/egomimic/scripts/abc_sim/policy_server.py @@ -105,12 +105,14 @@ def configure(self, cfg) -> None: self.lock = threading.Lock() def _past_frame(self, episode: str, t: float, frame: np.ndarray) -> np.ndarray: - """The front frame >= gap_s before t, else the oldest one this episode - (an episode start: the current frame), as ZarrDataset clamps to frame 0.""" - hist = self.history.setdefault(episode, deque(maxlen=int(FPS * 5))) - hist.append((t, frame)) - past = [f for (tt, f) in hist if tt <= t - self.gap_s] - return past[-1] if past else hist[0][1] + """The front frame round(gap_s * fps) frames before t -- the dataset's + _seconds_to_frames lag, in frames not seconds -- clamped to the oldest + frame this episode (an episode start: the current frame).""" + lag = max(1, int(round(self.gap_s * FPS))) + hist = self.history.setdefault(episode, deque(maxlen=lag + 1)) + hist.append((int(round(t * FPS)), frame)) + idx = hist[-1][0] - lag + return next((f for (i, f) in hist if i == idx), hist[0][1]) def build_sample(self, episode: str, t: float, state, images: dict[str, np.ndarray]) -> dict: """One un-normalized sample with the dataset's keys (the 45-step cmd From 5c5382181618bed2c9ea08da8089dc8eea018ebd Mon Sep 17 00:00:00 2001 From: Aidan Gao Date: Tue, 29 Sep 2026 11:06:12 -0400 Subject: [PATCH 08/20] abc_sim: audit fixes (chunk timing, history frame, server load) + ponytail pass Fixes found by an audit of this branch against ABC's code: - policy_server could not start: it called .to()/.eval() on the algo, a plain object; the Lightning wrapper owns and moves the nets. - The served chunk played 2.2x too slow: the model's 100 steps span 45 raw 30 Hz frames (JOINT_RAW_HORIZON resampled to 100), and the client executed one per 30 Hz step. The server now resamples back onto the 45-frame grid, so the 15 executed steps are 0.5 s of the plan, as in ABC. - The _hist frame went 0.5-1.5 s stale: the server only ever saw re-plan frames. The client now renders the frame lag_frames (3) steps before each re-plan and sends it; the server is stateless. - Checkpoints stored only the model config, so a rollout could not rebuild the training inputs (robot/rollout.py's transform-mode lookup always fell back to the legacy ypr pipeline). The config tree now also carries each dataset's resolved keymap + transform list. - sim_client now drives ABC's own rollout_worlds + build_summary (the unplaceable-seed retry, chunk metrics, resolved physics, the summary fields) through an HTTP policy with ABC's infer() surface, instead of a hand copy of that loop. - Prompts: the converter keeps ABC's rule verbatim (no "sim " forced onto instruction/timeline prompts); the client conditions on the training prompt (the task name) by default, --prompt env for the randomized tasks. The two differed on 4 of the 7 candidate tasks. Episodes ABC's loader drops as unlabelled are skipped; a states/metadata length mismatch raises. Ponytail pass (less code, same behaviour): tasks.py deleted (every sim_224 task_name is already an abc_sim alias); val_vs_sim on scipy.stats; the cotrain data config inherits mecka_all_6d instead of copying it; the init_weights report dict and its min-fraction config knob removed (the guard stays); the unreachable stereo-camera alias and --limit removed; config headers cut to their role line (one pointed at a file that does not exist); config-echo tests replaced by an end-to-end test: train a tiny RDT, save the checkpoint, serve it with policy_server, infer. Co-Authored-By: Claude Opus 5.5 --- .../hydra_configs/data/abc_real_joints.yaml | 12 +- .../hydra_configs/data/abc_sim_joints.yaml | 14 +- .../data/abc_sim_joints_cotrain_mecka.yaml | 44 +--- .../model/rdt_1b_abc_joints_dinov3_vitb.yaml | 26 +-- .../train_zarr_abc_sim_rdt_ft.yaml | 13 +- .../hydra_configs/train_zarr_cartesian.yaml | 2 - .../train_zarr_fold_ladder_rdt1b.yaml | 15 +- egomimic/scripts/abc_sim/README.md | 14 +- .../scripts/abc_sim/convert_sim_to_zarr.py | 49 ++--- egomimic/scripts/abc_sim/policy_server.py | 191 ++++++----------- .../scripts/abc_sim/rollout_lambda.sbatch | 12 +- egomimic/scripts/abc_sim/sim_client.py | 200 +++++++----------- egomimic/scripts/abc_sim/tasks.py | 66 ------ egomimic/scripts/abc_sim/val_vs_sim.py | 40 ++-- egomimic/trainHydra.py | 27 ++- egomimic/utils/checkpoint_utils.py | 62 ++---- tests/unit/test_abc_joint_configs.py | 43 ++-- tests/unit/test_abc_sim_convert.py | 7 +- tests/unit/test_eva_joint_mode.py | 9 - tests/unit/test_init_weights_ckpt.py | 6 +- tests/unit/test_policy_server.py | 112 +++++----- tests/unit/test_val_vs_sim.py | 7 - 22 files changed, 319 insertions(+), 652 deletions(-) delete mode 100644 egomimic/scripts/abc_sim/tasks.py diff --git a/egomimic/hydra_configs/data/abc_real_joints.yaml b/egomimic/hydra_configs/data/abc_real_joints.yaml index 323c4075f1..533f277c65 100644 --- a/egomimic/hydra_configs/data/abc_real_joints.yaml +++ b/egomimic/hydra_configs/data/abc_real_joints.yaml @@ -1,11 +1,7 @@ -# ABC_REAL_JOINTS: one REAL ABC task (lab=abc on the Phoenix mirror) in the -# same joint space as abc_sim_joints -- the "real" side of the sim-real check. -# The SQL rows were relabelled eva_bimanual -> yam_bimanual on 2026-09-01 but -# the mirror zarrs still say eva_bimanual, so the filter keys on lab + task + -# the /abc/ path, never the embodiment. The paper's three sim-real tasks exist -# here: "put the plastic bottles in the bin" (100 eps, the default below), "load -# the mixed dishes into the dish rack" (88), "turn the mug right side up" (65); -# another task restates the train filter_lambdas list with its name. +# ABC_REAL_JOINTS: one REAL ABC task (lab=abc) in the same joint space. The SQL +# embodiment was relabelled to yam_bimanual on 2026-09-01 but the zarrs say +# eva_bimanual, so the filter keys on lab + task + the /abc/ path. Another task: +# restate the filter_lambdas with its name. defaults: - abc_sim_joints - _self_ diff --git a/egomimic/hydra_configs/data/abc_sim_joints.yaml b/egomimic/hydra_configs/data/abc_sim_joints.yaml index c213d33f31..e4bffbcacf 100644 --- a/egomimic/hydra_configs/data/abc_sim_joints.yaml +++ b/egomimic/hydra_configs/data/abc_sim_joints.yaml @@ -1,13 +1,7 @@ -# ABC_SIM_JOINTS: one abc_sim task's teleop episodes (amazon-far/abc sim_224), -# converted by egomimic/scripts/abc_sim/convert_sim_to_zarr.py into -# $ABC_SIM_ZARR_ROOT// (default task put_bottles; another task is -# data.train_datasets.eva_bimanual.resolver.folder_path=/turn_mug -- the -# valid set interpolates the resolver). JOINT space (Eva mode "joints"): the -# 14-D ABC state and a (100, 14) actions_joints chunk from 45 raw 30 Hz steps. -# Train / valid follow ABC's own split (zarr metadata `split`). -# -# The front camera (sim `top`) is the only one decoded (keymap joints_front); -# image_history_gap_s pairs with the RDT trunk's image_history: 2. +# ABC_SIM_JOINTS: one abc_sim task's sim_224 teleop episodes, converted by +# egomimic/scripts/abc_sim/convert_sim_to_zarr.py to $ABC_SIM_ZARR_ROOT/put_bottles +# (another task: override the resolver folder_path). Eva "joints" mode, front +# camera + its 0.1 s history frame; train/valid = ABC's own `split`. _target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper train_datasets: diff --git a/egomimic/hydra_configs/data/abc_sim_joints_cotrain_mecka.yaml b/egomimic/hydra_configs/data/abc_sim_joints_cotrain_mecka.yaml index 7bbff1593a..3e7e011a9f 100644 --- a/egomimic/hydra_configs/data/abc_sim_joints_cotrain_mecka.yaml +++ b/egomimic/hydra_configs/data/abc_sim_joints_cotrain_mecka.yaml @@ -1,58 +1,28 @@ -# ABC_SIM_JOINTS_COTRAIN_MECKA: single-task cotrain -- one abc_sim task (the -# robot, joint space) + one mecka task (human, wrist-frame 6D cartesian), -# one batch of each per step (CombinedLoader max_size_cycle), so the mix is -# set by the two batch sizes. Pairs with model=rdt_1b_abc_joints_dinov3_vitb. -# -# Default pairing put_bottles <- mecka cup_on_saucer (the largest mecka task); -# another human task restates the human filter_lambdas list (dishrack <- -# dishwashing, put_bottles <- packaging_items are the nearest by motion). +# ABC_SIM_JOINTS_COTRAIN_MECKA: single-task cotrain -- the abc_sim task (robot, +# joint space, abc_sim_joints) + one mecka task (human, mecka_all_6d's +# wrist-frame 6D pipeline), one batch of each per step (CombinedLoader +# max_size_cycle). Default human task cup_on_saucer (the largest mecka task); +# another one restates the human filter_lambdas. See egomimic/scripts/abc_sim/README.md. defaults: - abc_sim_joints + - mecka_all_6d - _self_ train_datasets: human_bimanual: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver resolver: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver - folder_path: ${paths.dataset_dir} key_map: - _target_: egomimic.rldb.embodiment.human.Human.get_keymap - keymap_mode: cartesian - annotation_key: annotations + # RDT's 2-frame image history, as on the robot side image_history_gap_s: 0.1 - transform_list: - _target_: egomimic.rldb.embodiment.human.Human.get_transform_list - mode: cartesian_wristframe_6d - stride: 1 - fix_left_wrist_convention: true - pad_proprio_gripper: true filters: - _target_: egomimic.rldb.filters.DatasetFilter filter_lambdas: - "lambda row: str(row.get('lab', '')).lower() == 'mecka'" - "lambda row: str(row.get('task', '')) == 'cup_on_saucer'" - exclude_hashes: mecka_bad_egomotion - mode: train - valid_ratio: 0.05 - -valid_datasets: - human_bimanual: - _target_: ${data.train_datasets.human_bimanual._target_} - resolver: ${data.train_datasets.human_bimanual.resolver} - filters: ${data.train_datasets.human_bimanual.filters} - mode: valid - valid_ratio: ${data.train_datasets.human_bimanual.valid_ratio} - - train_dataloader_params: human_bimanual: batch_size: 32 - num_workers: 8 valid_dataloader_params: human_bimanual: batch_size: 32 - num_workers: 4 - shuffle: false diff --git a/egomimic/hydra_configs/model/rdt_1b_abc_joints_dinov3_vitb.yaml b/egomimic/hydra_configs/model/rdt_1b_abc_joints_dinov3_vitb.yaml index 9c956d5db2..eb789c0483 100644 --- a/egomimic/hydra_configs/model/rdt_1b_abc_joints_dinov3_vitb.yaml +++ b/egomimic/hydra_configs/model/rdt_1b_abc_joints_dinov3_vitb.yaml @@ -1,25 +1,7 @@ -# RDT_1B_ABC_JOINTS_DINOV3_VITB: RDT-1B (width 2048 x 28 x 32 heads) with a -# second embodiment, the ABC YAM station in JOINT space, beside the mecka -# human head -- the fine-tune / single-task cotrain model for the ABC sim -# rollout pipeline (egomimic/scripts/abc_sim/README.md). -# -# eva_bimanual here is the ABC station (real ABC zarrs and the converted -# sim_224 episodes both register as eva_bimanual): 14-D proprio -# (observations.state.joint_positions) and a (100, 14) actions_joints chunk, -# [left j1..j6, left grip, right j1..j6, right grip], exactly abc_sim's -# action space, so a rollout needs no IK. The DiT trunk, the DINOv3 front -# camera tower and the Qwen text stem are shared with human_bimanual; the -# robot gets its own state adaptor and denoiser head (RDT's per-embodiment -# adaptors at native width). -# -# Fine-tuning from a human-only RDT-1B (model=rdt_1b_bc_mecka_6d_dinov3_vitb): -# init_weights_ckpt= loads the trunk, towers and the human stems; -# the eva stem / head start fresh. ABC-only fine-tunes set -# model.robomimic_model.domains=[eva_bimanual] (the train configs do). -# -# Known gaps, deliberately left for later: the front camera only (the human -# pretrain has one camera; ABC's own policies use top + both wrists), and -# ABC's 4:3 frames go through the tower at the human 16:9 size [224, 400]. +# RDT_1B_ABC_JOINTS_DINOV3_VITB: RDT-1B with a second head, eva_bimanual = the +# ABC YAM station in JOINT space (14-D state, (100, 14) actions_joints, abc_sim's +# action space), beside the mecka human head; trunk, front tower and text stem +# shared. Front camera only, at the human tower size [224, 400]. defaults: - rdt_1b_bc_mecka_6d_dinov3_vitb - _self_ diff --git a/egomimic/hydra_configs/train_zarr_abc_sim_rdt_ft.yaml b/egomimic/hydra_configs/train_zarr_abc_sim_rdt_ft.yaml index 4cbe44c720..ab10c17693 100644 --- a/egomimic/hydra_configs/train_zarr_abc_sim_rdt_ft.yaml +++ b/egomimic/hydra_configs/train_zarr_abc_sim_rdt_ft.yaml @@ -1,16 +1,9 @@ -# ABC sim fine-tune: RDT-1B in the ABC station's joint space on ONE abc_sim -# task's teleop data, initialised (weights only) from a pretrained RDT-1B -- -# the per-base step of the "rank pretraining runs by ABC sim rollouts" -# pipeline (egomimic/scripts/abc_sim/README.md). Rolled out afterwards with -# egomimic/scripts/abc_sim/rollout_sim.sbatch. +# ABC sim fine-tune: RDT-1B, joint space, one abc_sim task, weights-only init +# from a base; init_weights_ckpt=null is the from-scratch floor. Pipeline: +# egomimic/scripts/abc_sim/README.md. # # python egomimic/trainHydra.py --config-name=train_zarr_abc_sim_rdt_ft \ # init_weights_ckpt= name=abc_sim_ft description=_put_bottles -# -# init_weights_ckpt=null trains the same arm from scratch (the no-pretrain -# floor every base is compared against). Every arm of a comparison must share -# data, steps and batch; only the init changes. Norm stats are computed at -# startup on the ABC data (eva only: the human base's stats play no part). defaults: - train_zarr_cartesian - override model: rdt_1b_abc_joints_dinov3_vitb diff --git a/egomimic/hydra_configs/train_zarr_cartesian.yaml b/egomimic/hydra_configs/train_zarr_cartesian.yaml index 1fd05f0557..1fd804f923 100644 --- a/egomimic/hydra_configs/train_zarr_cartesian.yaml +++ b/egomimic/hydra_configs/train_zarr_cartesian.yaml @@ -17,8 +17,6 @@ ckpt_path: null # schedule and epoch; name+shape-matched tensors only, so new embodiment # stems / heads start fresh). Ignored when ckpt_path resumes a run. init_weights_ckpt: null -# Fail if less than this fraction of the model's elements came from it. -init_weights_min_fraction: 0.5 mode: train hydra: diff --git a/egomimic/hydra_configs/train_zarr_fold_ladder_rdt1b.yaml b/egomimic/hydra_configs/train_zarr_fold_ladder_rdt1b.yaml index b468a128b0..8aa017c173 100644 --- a/egomimic/hydra_configs/train_zarr_fold_ladder_rdt1b.yaml +++ b/egomimic/hydra_configs/train_zarr_fold_ladder_rdt1b.yaml @@ -1,17 +1,10 @@ -# Freeform fold-clothes OPERATOR ladder, RDT-1B: train_zarr_mecka_flagship_6d_rdt -# (RDT recipe: 2-frame image history, grad clip 1.0, seen/unseen-op heads) at the -# 1B DiT on one rung of data/mecka_fold_freeform_opscale{1,2,4}_6d. Pick the -# rung on the command line; everything else is identical across the ladder: +# Freeform fold-clothes OPERATOR ladder, RDT-1B: the RDT recipe at 1B on one +# rung of data/mecka_fold_freeform_opscale{1,2,4}_6d; 1000 x 100 steps at global +# batch 64 (32 x 2 GPUs), val every 25 epochs, a checkpoint every 100 (the bases +# the ABC-sim fine-tunes start from). Norm stats per rung, at startup. # # python egomimic/trainHydra.py --config-name=train_zarr_fold_ladder_rdt1b \ # data=mecka_fold_freeform_opscale1_6d name=fold_ladder description=k1 -# -# Budget per rung (the RDT memory ablation's): 1000 epochs x 100 steps at a -# global batch of 64 (32 x 2 GPUs), val every 25 epochs, a full checkpoint every -# 100 epochs (10 kept: the step ladder the ABC-sim fine-tunes start from). Norm -# stats are computed at startup per rung (each rung is its own distribution) -# and cached under paths.cache_dir. Runs on Lambda via -# egomimic/scripts/abc_sim/lambda_pack.sbatch, three rungs packed on one node. defaults: - train_zarr_mecka_flagship_6d_rdt - override model: rdt_1b_bc_mecka_6d_dinov3_vitb diff --git a/egomimic/scripts/abc_sim/README.md b/egomimic/scripts/abc_sim/README.md index 4d542ea58c..8c2fa59d7f 100644 --- a/egomimic/scripts/abc_sim/README.md +++ b/egomimic/scripts/abc_sim/README.md @@ -8,14 +8,14 @@ real success over 12 checkpoints on bottles, dishrack and mug flip. Action space: the ABC station's 14-D absolute joints + gripper (Eva mode `joints`, `egomimic/rldb/embodiment/eva.py`), which is also abc_sim's action, -so a rollout needs no IK. Tasks: `tasks.py` (short name -> sim_224 dataset -task, abc_sim task, real ABC task). +so a rollout needs no IK. A task is named by its sim_224 `task_name` (e.g. +`sim_put_the_plastic_bottles_in_the_bin`), which is also an abc_sim alias. Steps (Lambda; paths in the launchers): 1. Data: `prepare.py --sim-data ` or `curl` the sim_224 tar, then `convert_sim_to_zarr.py --src --task --out - $ABC_SIM_ZARR_ROOT/` -> eva_bimanual zarrs with the real ABC keys. + $ABC_SIM_ZARR_ROOT/put_bottles` -> eva_bimanual zarrs with the real ABC keys. 2. Base: `train_zarr_fold_ladder_rdt1b` (RDT-1B on the mecka freeform-fold operator ladder) via `lambda_pack.sbatch`; every 100 epochs a checkpoint. 3. Fine-tune: `train_zarr_abc_sim_rdt_ft init_weights_ckpt=` @@ -23,7 +23,11 @@ Steps (Lambda; paths in the launchers): `train_zarr_abc_real_rdt_ft` the real-data twin; `train_zarr_abc_sim_rdt_cotrain` the single-task human+robot cotrain. 4. Rollout: `rollout_lambda.sbatch` = `policy_server.py` (training venv) + - `sim_client.py` (ABC venv, ABC's evaluators and eval defaults, no RTC) -> - `summary.json` in ABC's format. `POLICY=hold` is the no-motion floor. + `sim_client.py` (ABC venv, ABC's own `rollout_worlds` + `build_summary`, + no RTC) -> ABC's `summary.json`. The server resamples the model's 100-step + chunk back onto the 45 raw 30 Hz frames it spans, so the client's 15 + executed steps are 0.5 s of the plan, as in ABC; the client renders the + history frame `lag` steps before each re-plan. `POLICY=hold` is the + no-motion floor. 5. `val_vs_sim.py manifest.json --metric ` -> Spearman / Pearson with bootstrap CIs and the log-linear fit of the ladder. diff --git a/egomimic/scripts/abc_sim/convert_sim_to_zarr.py b/egomimic/scripts/abc_sim/convert_sim_to_zarr.py index 68d7bb0879..24af7b97e1 100644 --- a/egomimic/scripts/abc_sim/convert_sim_to_zarr.py +++ b/egomimic/scripts/abc_sim/convert_sim_to_zarr.py @@ -16,11 +16,11 @@ top -> front_1, left -> left_wrist, right -> right_wrist. Extra metadata: ``lab: abc_sim``, ``split`` (ABC's own train/val), ``abc_episode``, ``prompt``. -Prompts follow ABC's training rule (``abc_minimal.dataloader``): the episode's -``prompt_timeline`` if it has one, else its ``instruction`` when that differs -from the task name, else the task name with ``_`` -> `` ``; ``sim `` is -prefixed when missing, as ABC's mixture prompts sim episodes and as -``abc_sim`` prompts the env at eval time. +Prompts follow ABC's training rule (``abc_minimal.dataloader``) verbatim: the +episode's ``prompt_timeline`` if it has one, else its ``instruction`` when that +differs from the task name, else the task name with ``_`` -> `` `` (the sim +tasks' names start with ``sim_``). Episodes ABC's loader drops as unlabelled +are skipped. python -m egomimic.scripts.abc_sim.convert_sim_to_zarr \ --src $ABC_CACHE --task sim_put_the_plastic_bottles_in_the_bin \ @@ -53,7 +53,6 @@ } # abc_sim scenes: every policy camera is a fovy=58 deg pinhole. CAMERA_FOVY_DEG = 58.0 -SIM_PROMPT_PREFIX = "sim " DONE_MARKER = ".abc_sim_done" @@ -62,11 +61,6 @@ def task_name_to_prompt(task_name: str) -> str: return " ".join(task_name.replace("-", " ").replace("_", " ").split()) -def with_sim_prefix(prompt: str) -> str: - prompt = prompt.strip() - return prompt if prompt.startswith(SIM_PROMPT_PREFIX) else SIM_PROMPT_PREFIX + prompt - - def prompt_spans(meta: dict, num_steps: int) -> list[tuple[str, int, int]]: """``(text, start, end)`` spans covering ``[0, num_steps)``, end exclusive.""" task_name = str(meta.get("task_name") or "") @@ -91,7 +85,7 @@ def prompt_spans(meta: dict, num_steps: int) -> list[tuple[str, int, int]]: if i == 0: start = 0 # the first directive holds from the episode start if end > start: - spans.append((with_sim_prefix(text), start, end)) + spans.append((text, start, end)) return spans @@ -116,26 +110,12 @@ def decode_cameras(ep_dir: Path, cameras: tuple[str, ...]) -> dict[str, np.ndarr """``{camera: (T, H, W, 3) uint8}`` from the vertically stacked mp4.""" import av - frames = [] with av.open(str(ep_dir / "combined_camera-images-rgb.mp4")) as container: - for frame in container.decode(video=0): - frames.append(frame.to_ndarray(format="rgb24")) - if not frames: - raise ValueError("video has no frames") - stack = np.stack(frames) # (T, n * H, W, 3) + stack = np.stack([f.to_ndarray(format="rgb24") for f in container.decode(video=0)]) h = stack.shape[1] // len(cameras) if h * len(cameras) != stack.shape[1]: - raise ValueError( - f"video height {stack.shape[1]} is not {len(cameras)} stacked cameras" - ) - out = {cam: stack[:, i * h : (i + 1) * h] for i, cam in enumerate(cameras)} - if "top" not in out and "top_left" in out and "top_right" in out: - # abc_minimal.dataloader: stereo episodes alias one eye to "top". - import hashlib - - digest = hashlib.sha1(ep_dir.name.encode("utf-8")).digest()[0] - out["top"] = out["top_left" if digest % 2 == 0 else "top_right"] - return out + raise ValueError(f"video height {stack.shape[1]} is not {len(cameras)} stacked cameras") + return {cam: stack[:, i * h : (i + 1) * h] for i, cam in enumerate(cameras)} def split_arms(x: np.ndarray, kind: str) -> dict[str, np.ndarray]: @@ -161,8 +141,12 @@ def convert_episode(ep_dir: Path, split: str, out_dir: Path, overwrite: bool = F if (dst / DONE_MARKER).exists() and not overwrite: return {"episode": name, "status": "skipped"} meta = json.loads((ep_dir / "episode_metadata.json").read_text()) + if not meta.get("prompt_timeline") and (meta.get("prompt_source") or {}).get("status") == "excluded": + return {"episode": name, "status": "skipped_unlabelled"} # abc_minimal.dataloader drops these states, actions = load_states_actions(ep_dir) num_steps = len(states) + if meta.get("num_steps") is not None and int(meta["num_steps"]) != num_steps: + raise ValueError(f"states_actions.bin has {num_steps} rows, metadata says {meta['num_steps']}") cameras = tuple(meta.get("cameras") or DEFAULT_CAMERAS) images = decode_cameras(ep_dir, cameras) n_frames = min(len(v) for v in images.values()) @@ -224,18 +208,12 @@ def main(argv=None) -> int: p.add_argument("--task", required=True, help="dataset task_name, e.g. sim_pouring_beads") p.add_argument("--out", type=Path, required=True) p.add_argument("--splits", nargs="+", default=["train", "val"]) - p.add_argument("--limit", type=int, default=None, help="per split, for smokes") p.add_argument("--workers", type=int, default=max(1, (os.cpu_count() or 2) // 2)) p.add_argument("--overwrite", action="store_true") args = p.parse_args(argv) logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s") episodes = find_episodes(args.src, args.task, args.splits) - if args.limit is not None: - by_split: dict[str, list] = {} - for ep in episodes: - by_split.setdefault(ep[1], []).append(ep) - episodes = [ep for eps in by_split.values() for ep in eps[: args.limit]] if not episodes: log.error("no %s episodes under %s/{%s}_sim", args.task, args.src, ",".join(args.splits)) return 1 @@ -263,6 +241,7 @@ def main(argv=None) -> int: "src": str(args.src), "written": sum(r["status"] == "written" for r in rows), "skipped": sum(r["status"] == "skipped" for r in rows), + "skipped_unlabelled": sum(r["status"] == "skipped_unlabelled" for r in rows), "failed": failures, "frames_written": int(sum(r.get("frames", 0) for r in rows)), } diff --git a/egomimic/scripts/abc_sim/policy_server.py b/egomimic/scripts/abc_sim/policy_server.py index a2e03481f7..d4c847fbd5 100644 --- a/egomimic/scripts/abc_sim/policy_server.py +++ b/egomimic/scripts/abc_sim/policy_server.py @@ -1,22 +1,25 @@ """Serve an EgoVerse joint-space checkpoint over HTTP for the abc_sim client. The simulator's Python (amazon-far/abc: py3.12, MuJoCo 3.8, torch 2.11) cannot -share the training venv, so the policy runs here, in the training venv, and -the sim talks to it over plain HTTP + JSON (stdlib on both ends). +share the training venv, so the policy runs here and the sim talks to it over +plain HTTP + JSON (stdlib on both ends). python -m egomimic.scripts.abc_sim.policy_server --ckpt --port 8765 -POST /infer {"episode": str, "t": float seconds, "state": [14], "prompt": str, - "images": {"top": {"shape": [H, W, 3], "b64": }, ...}} - -> {"actions": [[14] x 100], "dt": 1/30} -GET /health -> {"ok": true, "ckpt": ..., "mode": "joints"} - -The inputs are built EXACTLY as ZarrDataset builds a training sample for the -checkpoint's data config -- the same Eva keymap and transform list, then the -checkpoint's own norm stats (the dataset normalizes samples in __getitem__; -``process_batch_for_training`` does not, which the old robot rollout missed). -The ``_hist`` camera is the frame ``image_history_gap_s`` earlier, kept per -episode; at an episode start it is the current frame, as in training. +POST /infer {"state": [14], "prompt": str, + "images": {"top": {"shape": [H, W, 3], "b64": }, + "top_hist": {...}}} # optional: the frame lag_frames back + -> {"actions": [[14] x 45], "dt": 1/30} +GET /health -> {"ok": true, "ckpt", "cameras", "lag_frames", "horizon"} + +The input is built exactly as ZarrDataset builds a training sample for the +checkpoint's data config (keymap, transform list, fps) and normalized with +the checkpoint's norm stats: the dataset normalizes in __getitem__, and +process_batch_for_training does not. ``top_hist`` is the frame +``image_history_gap_s`` earlier; without it the current frame stands in, which +is what an episode start looks like in training. The model's 100-step chunk +spans 45 raw 30 Hz frames, so it is resampled back onto that grid: action k +is the command k frames ahead. """ from __future__ import annotations @@ -25,9 +28,7 @@ import base64 import json import logging -import threading -from collections import deque -from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from http.server import BaseHTTPRequestHandler, HTTPServer import numpy as np import torch @@ -35,39 +36,23 @@ from torch.utils.data import default_collate from egomimic.rldb.embodiment.embodiment import IMAGE_HISTORY_SUFFIX, get_embodiment_id -from egomimic.rldb.embodiment.eva import ( - JOINT_ACTION_KEY, - JOINT_MODE, - JOINT_RAW_HORIZON, - Eva, -) +from egomimic.rldb.embodiment.eva import JOINT_ACTION_KEY, JOINT_MODE, JOINT_RAW_HORIZON, Eva +from egomimic.scripts.abc_sim.convert_sim_to_zarr import CAMERA_TO_ZARR, split_arms +from egomimic.utils.pose_utils import _interpolate_linear log = logging.getLogger("abc_sim.policy_server") EMB = "eva_bimanual" FPS = 30.0 -# abc_sim camera -> the zarr key the converter wrote it to -CAMERA_ZARR = {"top": "images.front_1", "left": "images.left_wrist", "right": "images.right_wrist"} +ZARR_TO_CAMERA = {z: c for c, z in CAMERA_TO_ZARR.items()} def decode_image(spec: dict) -> np.ndarray: - """{"shape": [H, W, 3], "b64": raw uint8} -> (3, H, W) float32 in [0, 1], RGB.""" - h, w, c = spec["shape"] - raw = np.frombuffer(base64.b64decode(spec["b64"]), dtype=np.uint8).reshape(h, w, c) + """{"shape": [H, W, 3], "b64": raw uint8} -> (3, H, W) float32 in [0, 1].""" + raw = np.frombuffer(base64.b64decode(spec["b64"]), dtype=np.uint8).reshape(spec["shape"]) return np.moveaxis(raw, -1, -3).astype(np.float32) / np.float32(255.0) -def split_state(state) -> dict[str, np.ndarray]: - """ABC 14-D state -> the zarr proprio keys of Eva's joint keymap.""" - s = np.asarray(state, dtype=np.float32).reshape(14) - return { - "left.obs_joints": s[0:6], - "left.obs_gripper": s[6:7], - "right.obs_joints": s[7:13], - "right.obs_gripper": s[13:14], - } - - class JointPolicy: """The checkpoint plus the data pipeline it was trained with.""" @@ -75,12 +60,13 @@ def __init__(self, ckpt_path: str, device: str = "cuda"): from egomimic.pl_utils.pl_model import ModelWrapper self.ckpt_path = ckpt_path - self.device = torch.device(device if torch.cuda.is_available() else "cpu") - wrapper = ModelWrapper.load_from_checkpoint(ckpt_path, weights_only=False, map_location="cpu") - self.algo = wrapper.model.to(self.device).eval() - self.algo.device = self.device - cfg = wrapper._as_config(getattr(wrapper.hparams, "config_tree", None)) - self.configure(cfg) + device = torch.device(device if torch.cuda.is_available() else "cpu") + # The algo is a plain object; the Lightning wrapper owns (and moves) its nets. + self.wrapper = ModelWrapper.load_from_checkpoint(ckpt_path, weights_only=False, map_location="cpu") + self.wrapper.to(device).eval() + self.algo = self.wrapper.model + self.algo.device = device + self.configure(self.wrapper._as_config(self.wrapper.hparams.config_tree)) def configure(self, cfg) -> None: """Keymap + transform list from the checkpoint's data config.""" @@ -91,71 +77,44 @@ def configure(self, cfg) -> None: tl = {k: v for k, v in OmegaConf.to_container(resolver.transform_list, resolve=True).items() if k != "_target_"} if tl.get("mode") != JOINT_MODE: raise ValueError(f"policy_server serves Eva mode '{JOINT_MODE}', checkpoint has {tl.get('mode')!r}") - self.key_map = Eva.get_keymap(**{k: v for k, v in km.items() if k != "annotation_key"}) - self.annotation_key = km.get("annotation_key") + self.annotation_key = km.pop("annotation_key", None) + self.key_map = Eva.get_keymap(**km) self.transforms = Eva.get_transform_list(**tl) - self.gap_s = km.get("image_history_gap_s") - self.cameras = { - k: v["zarr_key"] - for k, v in self.key_map.items() - if v.get("key_type") == "camera_keys" and not k.endswith(IMAGE_HISTORY_SUFFIX) - } + self.lag_frames = max(1, round(km["image_history_gap_s"] * FPS)) if km.get("image_history_gap_s") else 0 + self.cameras = sorted({ZARR_TO_CAMERA[v["zarr_key"]] for v in self.key_map.values() if v.get("key_type") == "camera_keys"}) self.emb_id = get_embodiment_id(EMB) - self.history: dict[str, deque] = {} - self.lock = threading.Lock() - - def _past_frame(self, episode: str, t: float, frame: np.ndarray) -> np.ndarray: - """The front frame round(gap_s * fps) frames before t -- the dataset's - _seconds_to_frames lag, in frames not seconds -- clamped to the oldest - frame this episode (an episode start: the current frame).""" - lag = max(1, int(round(self.gap_s * FPS))) - hist = self.history.setdefault(episode, deque(maxlen=lag + 1)) - hist.append((int(round(t * FPS)), frame)) - idx = hist[-1][0] - lag - return next((f for (i, f) in hist if i == idx), hist[0][1]) - - def build_sample(self, episode: str, t: float, state, images: dict[str, np.ndarray]) -> dict: - """One un-normalized sample with the dataset's keys (the 45-step cmd - chunk is faked from the current pose: only its shape reaches the model).""" - raw = split_state(state) - raw.update({f"{arm}.cmd_{part}": np.tile(raw[f"{arm}.obs_{part}"], (JOINT_RAW_HORIZON, 1)) - for arm in ("left", "right") for part in ("joints", "gripper")}) + + def build_sample(self, state, images: dict[str, np.ndarray]) -> dict: + """One un-normalized sample with the dataset's keys. The 45-step cmd + chunk is the current pose tiled: only its shape reaches the model.""" + raw = {k: v[0] for k, v in split_arms(np.asarray(state, np.float32).reshape(1, 14), "obs").items()} + raw.update({k.replace("obs_", "cmd_"): np.tile(v, (JOINT_RAW_HORIZON, 1)) for k, v in list(raw.items())}) sample = {} for key, spec in self.key_map.items(): - zk = spec["zarr_key"] if spec.get("key_type") == "camera_keys": - cam = next(c for c, z in CAMERA_ZARR.items() if z == zk) + cam = ZARR_TO_CAMERA[spec["zarr_key"]] if cam not in images: raise ValueError(f"camera '{cam}' missing from the request (have {sorted(images)})") - frame = images[cam] - if key.endswith(IMAGE_HISTORY_SUFFIX): - frame = self._past_frame(episode, t, images[cam]) - sample[key] = frame + past = images.get(cam + IMAGE_HISTORY_SUFFIX, images[cam]) + sample[key] = past if key.endswith(IMAGE_HISTORY_SUFFIX) else images[cam] elif spec.get("key_type") in ("proprio_keys", "action_keys"): - sample[key] = raw[zk] + sample[key] = raw[spec["zarr_key"]] for tf in self.transforms: sample = tf.transform(sample) - for k, v in sample.items(): - if isinstance(v, np.ndarray): - sample[k] = torch.from_numpy(v).to(torch.float32) - sample["fps"] = torch.tensor(FPS, dtype=torch.float32) + sample = {k: torch.as_tensor(v, dtype=torch.float32) for k, v in sample.items() if isinstance(v, (np.ndarray, torch.Tensor))} + sample["fps"] = torch.tensor(FPS) return sample @torch.no_grad() - def infer(self, episode: str, t: float, state, images: dict[str, np.ndarray], prompt: str) -> np.ndarray: - with self.lock: - sample = self.build_sample(episode, t, state, images) - sample = self.algo.norm_stats.normalize(sample, self.emb_id) - batch = default_collate([{k: v for k, v in sample.items() if isinstance(v, torch.Tensor)}]) - if self.annotation_key: - batch[self.annotation_key] = [[prompt]] - processed = self.algo.process_batch_for_training({EMB: batch}) - preds = self.algo.forward_eval(processed) - return preds[f"{EMB}_{JOINT_ACTION_KEY}"][0].float().cpu().numpy() - - def reset(self, episode: str) -> None: - with self.lock: - self.history.pop(episode, None) + def infer(self, state, images: dict[str, np.ndarray], prompt: str) -> np.ndarray: + """(45, 14) absolute joint + gripper commands, one per 30 Hz step.""" + sample = self.algo.norm_stats.normalize(self.build_sample(state, images), self.emb_id) + batch = default_collate([sample]) + if self.annotation_key: + batch[self.annotation_key] = [[prompt]] + preds = self.algo.forward_eval(self.algo.process_batch_for_training({EMB: batch})) + chunk = preds[f"{EMB}_{JOINT_ACTION_KEY}"][0].float().cpu().numpy() + return _interpolate_linear(chunk, JOINT_RAW_HORIZON) def make_handler(policy: JointPolicy): @@ -169,39 +128,29 @@ def _send(self, code: int, payload: dict) -> None: self.wfile.write(body) def do_GET(self): - if self.path == "/health": - self._send(200, {"ok": True, "ckpt": policy.ckpt_path, "mode": JOINT_MODE}) - else: - self._send(404, {"error": self.path}) + if self.path != "/health": + return self._send(404, {"error": self.path}) + self._send(200, {"ok": True, "ckpt": policy.ckpt_path, "cameras": policy.cameras, + "lag_frames": policy.lag_frames, "horizon": JOINT_RAW_HORIZON}) def do_POST(self): - req = json.loads(self.rfile.read(int(self.headers.get("Content-Length", 0)))) + if self.path != "/infer": + return self._send(404, {"error": self.path}) try: - if self.path == "/reset": - policy.reset(req["episode"]) - self._send(200, {"ok": True}) - elif self.path == "/infer": - images = {k: decode_image(v) for k, v in req["images"].items()} - actions = policy.infer(req["episode"], float(req["t"]), req["state"], images, req.get("prompt", "")) - self._send(200, {"actions": actions.tolist(), "dt": 1.0 / FPS}) - else: - self._send(404, {"error": self.path}) + req = json.loads(self.rfile.read(int(self.headers.get("Content-Length", 0)))) + images = {k: decode_image(v) for k, v in req["images"].items()} + actions = policy.infer(req["state"], images, req.get("prompt", "")) + self._send(200, {"actions": actions.tolist(), "dt": 1.0 / FPS}) except Exception as e: # the client must see the reason, not a hang log.exception("request failed") self._send(500, {"error": f"{type(e).__name__}: {e}"}) - def log_message(self, *a): # quiet + def log_message(self, *a): pass return Handler -def serve(policy: JointPolicy, host: str, port: int) -> ThreadingHTTPServer: - server = ThreadingHTTPServer((host, port), make_handler(policy)) - threading.Thread(target=server.serve_forever, daemon=True).start() - return server - - def main(argv=None) -> int: p = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) p.add_argument("--ckpt", required=True) @@ -211,12 +160,8 @@ def main(argv=None) -> int: a = p.parse_args(argv) logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s") policy = JointPolicy(a.ckpt, a.device) - server = serve(policy, a.host, a.port) - log.info("serving %s on http://%s:%d (cameras %s, gap %s s)", a.ckpt, a.host, a.port, sorted(policy.cameras), policy.gap_s) - try: - threading.Event().wait() - except KeyboardInterrupt: - server.shutdown() + log.info("serving %s on http://%s:%d (cameras %s, lag %d frames)", a.ckpt, a.host, a.port, policy.cameras, policy.lag_frames) + HTTPServer((a.host, a.port), make_handler(policy)).serve_forever() return 0 diff --git a/egomimic/scripts/abc_sim/rollout_lambda.sbatch b/egomimic/scripts/abc_sim/rollout_lambda.sbatch index 3a9b729561..029a522ee0 100644 --- a/egomimic/scripts/abc_sim/rollout_lambda.sbatch +++ b/egomimic/scripts/abc_sim/rollout_lambda.sbatch @@ -3,9 +3,10 @@ # training venv (pytorch image) and sim_client.py in ABC's venv, both in one # container step so they share the GPU (mjwarp renders on it too). # -# CKPT= TASK=put_bottles OUT= sbatch --qos=loaner-normal \ +# CKPT= TASK=sim_put_the_plastic_bottles_in_the_bin OUT= sbatch --qos=loaner-normal \ # egomimic/scripts/abc_sim/rollout_lambda.sbatch -# NUM_WORLDS=50 (default), POLICY=hold for the no-motion floor (no server) +# TASK: any abc_sim name or alias (the sim_224 task_name); NUM_WORLDS=50; +# POLICY=hold for the no-motion floor (no server); CLIENT_EXTRA e.g. --save-video #SBATCH --job-name=abc_rollout #SBATCH --account=loaner #SBATCH --partition=loaner @@ -20,14 +21,13 @@ set -euo pipefail SRC=${SRC:-/workspace/users/$USER/EgoVerse} ABC=${ABC:-/workspace/users/$USER/abc} IMAGE=${IMAGE:-/workspace/users/$USER/images/nvidia-pytorch-25.06-py3.sqsh} -TASK=${TASK:-put_bottles} +TASK=${TASK:-sim_put_the_plastic_bottles_in_the_bin} : "${OUT:?set OUT}" POLICY=${POLICY:-served} [ "$POLICY" = hold ] || : "${CKPT:?set CKPT}" -SIM_TASK=$($SRC/.venv/bin/python -m egomimic.scripts.abc_sim.tasks "$TASK" sim 2>/dev/null || python3 "$SRC/egomimic/scripts/abc_sim/tasks.py" "$TASK" sim) PORT=$(( 8700 + SLURM_JOB_ID % 200 )) mkdir -p "$OUT" -echo "=== abc_rollout job=$SLURM_JOB_ID node=$(hostname) task=$TASK/$SIM_TASK ckpt=${CKPT:-none} policy=$POLICY out=$OUT" +echo "=== abc_rollout job=$SLURM_JOB_ID node=$(hostname) task=$TASK ckpt=${CKPT:-none} policy=$POLICY out=$OUT" srun --ntasks=1 --container-image="$IMAGE" --container-mounts=/workspace:/workspace:rw --container-workdir="$SRC" \ bash -c "set -euo pipefail export HOME=/workspace/users/$USER HF_HOME=/workspace/users/$USER/.cache/huggingface HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 @@ -38,6 +38,6 @@ srun --ntasks=1 --container-image="$IMAGE" --container-mounts=/workspace:/worksp curl -fsS http://127.0.0.1:$PORT/health || { echo 'server never came up'; tail -50 '$OUT/server.log'; exit 1; } fi cd $ABC - MUJOCO_EGL_DEVICE_ID=0 .venv/bin/python $SRC/egomimic/scripts/abc_sim/sim_client.py --task '$SIM_TASK' \ + MUJOCO_EGL_DEVICE_ID=0 .venv/bin/python $SRC/egomimic/scripts/abc_sim/sim_client.py --task '$TASK' \ --server http://127.0.0.1:$PORT --policy '$POLICY' --out '$OUT' --num-worlds ${NUM_WORLDS:-50} ${CLIENT_EXTRA:-} echo '{\"job\": '$SLURM_JOB_ID', \"ckpt\": \"${CKPT:-}\", \"task\": \"$TASK\", \"commit\": \"'\$(git -C $SRC rev-parse --short HEAD)'\"}' > '$OUT/run_info.json'" diff --git a/egomimic/scripts/abc_sim/sim_client.py b/egomimic/scripts/abc_sim/sim_client.py index 8d65d386ac..312f7115c5 100644 --- a/egomimic/scripts/abc_sim/sim_client.py +++ b/egomimic/scripts/abc_sim/sim_client.py @@ -1,19 +1,18 @@ -"""Roll a served EgoVerse policy out in abc_sim (amazon-far/abc) and score it -with ABC's own task evaluators. Runs in ABC's venv: stdlib + numpy + abc_minimal -only (no egomimic import; the policy lives behind policy_server.py). +"""Roll a served EgoVerse policy out in abc_sim (amazon-far/abc) through ABC's +own eval loop and summary (abc_minimal.eval_policy.rollout_worlds / +build_summary). Runs in ABC's venv; the policy lives behind policy_server.py. /bin/python egomimic/scripts/abc_sim/sim_client.py \ - --task put_plastic_bottles_in_bin --server http://127.0.0.1:8765 \ + --task sim_put_the_plastic_bottles_in_the_bin --server http://127.0.0.1:8765 \ --num-worlds 50 --out -Protocol = ABC's SimEvalConfig defaults without RTC: worlds seeded seed+i, -one inference per 15 executed steps (execute_chunk_dim) out of the policy's -100-step chunk, 236 chunks max (~118 s at 30 Hz), early stop on success -(rollout_over), cameras top/left/right at 168x224. summary.json carries the -same fields ABC's build_summary writes (success_rate, num_success, -mean_reward, mean_max_progress, worlds[]) so the two are comparable side by -side. ``--policy hold`` needs no server: it repeats the current state (the -no-motion floor and a protocol smoke). +ABC's SimEvalConfig defaults (seeded worlds seed+i with its unplaceable-seed +retry, 15 executed steps per inference, 236 chunks, early stop on success, +top/left/right at 168x224) with RTC off. ``--task`` is any abc_sim name or +alias -- the sim_224 task_name works. The prompt defaults to the one the +converter trained on (the task name with _ -> space); ``--prompt env`` keeps +the env's own (the prompt-randomized tasks). ``--policy hold`` needs no +server: it repeats the current state, the no-motion floor. """ from __future__ import annotations @@ -21,152 +20,99 @@ import argparse import base64 import json -import time import urllib.request +from dataclasses import replace from pathlib import Path import numpy as np -from abc_minimal.eval_policy import jsonable, rollout_over, video_frame +from abc_minimal.config import SimEvalConfig +from abc_minimal.dit import task_name_to_prompt +from abc_minimal.eval_policy import build_summary, resolved_physics, rollout_worlds from abc_minimal.sim_env import SimTaskEnv, task_prompt CAMERAS = ("top", "left", "right") +HORIZON = 45 # the served chunk: 1.5 s of 30 Hz commands + + +class HistEnv(SimTaskEnv): + """Renders the frame ``lag`` steps before each re-plan (ABC's loop only + renders at the re-plan itself), for the policy's 2-frame image history.""" + + def __init__(self, *args, lag: int, execute: int, **kwargs): + super().__init__(*args, **kwargs) + self.capture_at, self.hist, self._since_obs = execute - lag, None, 0 + + def reset(self, seed, options=None): + self.hist, self._since_obs = None, 0 + return super().reset(seed, options) + + def obs(self): + self._since_obs = 0 + return super().obs() + + def step_one(self, action): + super().step_one(action) + self._since_obs += 1 + if self.capture_at > 0 and self._since_obs == self.capture_at: + self.hist = self.render_cameras() class ServedPolicy: - def __init__(self, server: str): - self.server = server.rstrip("/") + """abc_minimal's SimPolicy.infer surface over HTTP (noise and prefix unused).""" + + def __init__(self, server: str, env: HistEnv | None = None): + self.server, self.env = server.rstrip("/"), env with urllib.request.urlopen(f"{self.server}/health", timeout=30) as r: self.health = json.load(r) - def _post(self, path: str, payload: dict) -> dict: - req = urllib.request.Request( - f"{self.server}{path}", data=json.dumps(payload).encode(), - headers={"Content-Type": "application/json"}, method="POST", - ) + def infer(self, obs, noise=None, action_prefix=None, prefix_length=0) -> np.ndarray: + images = {c: obs["images"][c] for c in self.health["cameras"]} + if self.env is not None and self.env.hist is not None: + images.update({f"{c}_hist": self.env.hist[c] for c in self.health["cameras"]}) + payload = {"state": np.asarray(obs["state"], np.float32).tolist(), "prompt": obs["prompt"], "images": { + k: {"shape": list(hwc.shape), "b64": base64.b64encode(hwc.tobytes()).decode()} + for k, hwc in ((k, np.ascontiguousarray(np.moveaxis(np.asarray(v, np.uint8), 0, -1))) for k, v in images.items())}} + req = urllib.request.Request(f"{self.server}/infer", data=json.dumps(payload).encode(), + headers={"Content-Type": "application/json"}, method="POST") with urllib.request.urlopen(req, timeout=120) as r: - return json.load(r) - - def reset(self, episode: str) -> None: - self._post("/reset", {"episode": episode}) - - def infer(self, episode: str, t: float, obs: dict) -> np.ndarray: - images = {} - for cam, img in obs["images"].items(): - hwc = np.ascontiguousarray(np.moveaxis(np.asarray(img, dtype=np.uint8), 0, -1)) # (3,H,W)->(H,W,3) - images[cam] = {"shape": list(hwc.shape), "b64": base64.b64encode(hwc.tobytes()).decode()} - out = self._post("/infer", { - "episode": episode, "t": t, "state": np.asarray(obs["state"], dtype=np.float32).tolist(), - "prompt": obs["prompt"], "images": images, - }) - return np.asarray(out["actions"], dtype=np.float32) + return np.asarray(json.load(r)["actions"], dtype=np.float32) class HoldPolicy: - health = {"ckpt": "hold", "mode": "hold"} - - def reset(self, episode: str) -> None: - pass - - def infer(self, episode: str, t: float, obs: dict) -> np.ndarray: - return np.tile(np.asarray(obs["state"], dtype=np.float32)[None], (100, 1)) - - -def rollout(env: SimTaskEnv, policy, *, seed: int, num_chunks: int, execute: int, video_path=None) -> dict: - episode = f"seed{seed}" - policy.reset(episode) - obs = env.reset(seed=seed) - video = None - if video_path is not None: - import imageio.v2 as imageio - - video = imageio.get_writer(str(video_path), fps=30, macro_block_size=1) - video.append_data(video_frame(obs["images"], CAMERAS)) - result = env.evaluate() - max_reward = float(result.get("reward", 0.0)) - steps, infer_s, t0 = 0, [], time.perf_counter() - try: - for _ in range(num_chunks): - t_inf = time.perf_counter() - actions = policy.infer(episode, steps / 30.0, obs) - infer_s.append(time.perf_counter() - t_inf) - for action in actions[:execute]: - env.step_one(action) - result = env.evaluate() - max_reward = max(max_reward, float(result.get("reward", 0.0))) - steps += 1 - if video is not None: - video.append_data(video_frame(env.render_cameras(), CAMERAS)) - if rollout_over(result): - break - if rollout_over(result): - break - obs = env.obs() - finally: - if video is not None: - video.close() - return { - "world_seed": seed, - "success": bool(result["ever_success"]), - "final_success": bool(result["success"]), - "reward": float(result["reward"]), - "max_reward": max_reward, - "steps": steps, - "wall_s": time.perf_counter() - t0, - "infer_s_mean": float(np.mean(infer_s)) if infer_s else None, - "randomization": env.randomization, - "final_task_eval": result, - "video_path": str(video_path) if video_path is not None else None, - } + health = {"ckpt": "hold", "cameras": [], "lag_frames": 0} + + def infer(self, obs, noise=None, action_prefix=None, prefix_length=0) -> np.ndarray: + return np.tile(np.asarray(obs["state"], np.float32)[None], (HORIZON, 1)) def main(argv=None) -> int: p = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) - p.add_argument("--task", required=True, help="abc_sim task name / alias (e.g. put_plastic_bottles_in_bin)") + p.add_argument("--task", required=True) p.add_argument("--server", default="http://127.0.0.1:8765") p.add_argument("--policy", choices=["served", "hold"], default="served") + p.add_argument("--prompt", default=None, help="default: the training prompt; 'env' = the env's own") p.add_argument("--out", type=Path, required=True) p.add_argument("--num-worlds", type=int, default=50) - p.add_argument("--seed", type=int, default=20260511, help="ABC's SimEvalConfig default") - p.add_argument("--num-chunks", type=int, default=236) - p.add_argument("--execute", type=int, default=15, help="steps executed per inference (ABC's execute_chunk_dim)") + p.add_argument("--save-video", action="store_true") p.add_argument("--camera-backend", default="mjwarp") p.add_argument("--gpu-id", type=int, default=None) - p.add_argument("--save-videos", type=int, default=2, help="record the first N worlds") a = p.parse_args(argv) a.out.mkdir(parents=True, exist_ok=True) + policy = HoldPolicy() if a.policy == "hold" else ServedPolicy(a.server) - env = SimTaskEnv(task=a.task, height=168, width=224, camera_keys=CAMERAS, prompt=task_prompt(a.task), - camera_backend=a.camera_backend, gpu_id=a.gpu_id) - worlds = [] - try: - for i in range(a.num_worlds): - video = a.out / f"world_{i:03d}.mp4" if i < a.save_videos else None - w = rollout(env, policy, seed=a.seed + i, num_chunks=a.num_chunks, execute=a.execute, video_path=video) - w["world_index"] = i - worlds.append(w) - print(f"world={i:03d} success={w['success']} max_progress={w['max_reward']:.3f} steps={w['steps']} " - f"infer={1000 * (w['infer_s_mean'] or 0):.0f}ms", flush=True) - (a.out / "worlds.jsonl").open("a").write(json.dumps(jsonable(w)) + "\n") - finally: - env.close() - succ = np.array([w["success"] for w in worlds], dtype=bool) - summary = { - "format": "egoverse_abc_sim_rollout/v1", - "policy": policy.health, - "task": env.spec.name, - "prompt": env.prompt, - "config": {k: (str(v) if isinstance(v, Path) else v) for k, v in vars(a).items()}, - "success_rate": float(succ.mean()) if len(succ) else None, - "num_success": int(succ.sum()), - "num_worlds": len(worlds), - "mean_reward": float(np.mean([w["reward"] for w in worlds])) if worlds else None, - "mean_max_progress": float(np.mean([w["max_reward"] for w in worlds])) if worlds else None, - "worlds": worlds, - } - (a.out / "summary.json").write_text(json.dumps(jsonable(summary), indent=2, sort_keys=True)) - print(f"summary: success_rate={summary['success_rate']} num_success={summary['num_success']}/{len(worlds)} " - f"mean_max_progress={summary['mean_max_progress']}", flush=True) + prompt = task_prompt(a.task) if a.prompt == "env" else (a.prompt or task_name_to_prompt(a.task)) + cfg = SimEvalConfig(checkpoint=str(policy.health["ckpt"]), task=a.task, num_worlds=a.num_worlds, + rtc=False, fast_inference=False, prefix_length=0, save_video=a.save_video, + camera_backend=a.camera_backend, gpu_id=a.gpu_id, output_dir=str(a.out), prompt=prompt) + model_config = replace(cfg.model, chunk_length=HORIZON, action_dim=14, camera_keys=CAMERAS) + env = HistEnv(task=a.task, height=cfg.camera_height, width=cfg.camera_width, camera_keys=CAMERAS, + prompt=prompt, camera_backend=cfg.camera_backend, gpu_id=cfg.gpu_id, + lag=int(policy.health["lag_frames"]), execute=cfg.execute_chunk_dim) + policy.env = env + physics = resolved_physics(env) + worlds = rollout_worlds(cfg, policy, env, 0, None, a.out, model_config) + build_summary(config=cfg, ckpt_path=Path(cfg.checkpoint), device="served", worlds=worlds, out_dir=a.out, physics=physics) return 0 diff --git a/egomimic/scripts/abc_sim/tasks.py b/egomimic/scripts/abc_sim/tasks.py deleted file mode 100644 index b00d3f04b5..0000000000 --- a/egomimic/scripts/abc_sim/tasks.py +++ /dev/null @@ -1,66 +0,0 @@ -"""The ABC tasks the sim rollout pipeline runs, by short name. - -``dataset``: the sim_224 ``task_name`` (convert_sim_to_zarr --task); -``sim``: the abc_sim task spec (the rollout env); ``real``: the matching REAL -ABC task on the Phoenix mirror (data=abc_real_joints data.abc_task=...), or -None. The first three are the tasks the ABC paper measured sim-real -correlation on (Pearson r = 0.85 success / 0.91 progress, 12 checkpoints). -The short name is also the zarr folder name under ABC_SIM_ZARR_ROOT. -""" - -from __future__ import annotations - -TASKS: dict[str, dict] = { - "put_bottles": { - "dataset": "sim_put_the_plastic_bottles_in_the_bin", - "sim": "put_plastic_bottles_in_bin", - "real": "put the plastic bottles in the bin", - }, - "dishrack": { - "dataset": "sim_load_the_plates_into_the_dish_rack", - "sim": "load_plates_into_dish_rack", - # real ABC loads mixed dishes; the sim rack holds plates only - "real": "load the mixed dishes into the dish rack", - }, - "turn_mug": { - "dataset": "sim_turn_the_mug_right_side_up", - "sim": "turn_mug_right_side_up", - "real": "turn the mug right side up", - }, - "hang_mug": { - "dataset": "sim_hang_the_mug_on_the_mug_rack", - "sim": "hang_mug_on_mug_rack", - "real": None, - }, - "pour": { - "dataset": "sim_pouring_beads", - "sim": "pouring", - "real": None, - }, - "sweep": { - "dataset": "sim_sweep_away_paper_scraps_from_the_table", - "sim": "sweep_away_paper_scraps_from_table", - "real": "sweep away the paper scraps from the table", - }, - "inhand_transfer": { - "dataset": "sim_inhand_transfer_the_item_to_other_side", - "sim": "inhand_transfer_item_to_other_side", - "real": None, - }, -} - - -def task(short_name: str) -> dict: - try: - return TASKS[short_name] - except KeyError: - raise KeyError( - f"unknown ABC task '{short_name}'; known: {', '.join(sorted(TASKS))}" - ) from None - - -if __name__ == "__main__": - import sys - - # `python -m egomimic.scripts.abc_sim.tasks ` for shell use. - print(task(sys.argv[1])[sys.argv[2]]) diff --git a/egomimic/scripts/abc_sim/val_vs_sim.py b/egomimic/scripts/abc_sim/val_vs_sim.py index 9af8947ffb..823ba6a3ae 100644 --- a/egomimic/scripts/abc_sim/val_vs_sim.py +++ b/egomimic/scripts/abc_sim/val_vs_sim.py @@ -20,10 +20,14 @@ import argparse import json -import math +import warnings from pathlib import Path import numpy as np +from scipy import stats + +# constant inputs inside a bootstrap resample give nan by design +warnings.filterwarnings("ignore", category=stats.ConstantInputWarning) def wandb_metric(run_path: str, metric: str, epoch: int) -> float: @@ -38,42 +42,22 @@ def wandb_metric(run_path: str, metric: str, epoch: int) -> float: return float(at[-1][metric]) -def _rank(x): - order = np.argsort(x) - r = np.empty(len(x)) - r[order] = np.arange(len(x)) - return r - - def corr(x, y, kind: str) -> float: - x, y = np.asarray(x, float), np.asarray(y, float) - if kind == "spearman": - x, y = _rank(x), _rank(y) - if x.std() == 0 or y.std() == 0: - return float("nan") - return float(np.corrcoef(x, y)[0, 1]) + return float((stats.spearmanr if kind == "spearman" else stats.pearsonr)(x, y)[0]) def bootstrap(x, y, kind: str, n: int = 2000, seed: int = 0) -> tuple[float, float]: + """95% percentile interval of ``corr`` over resampled checkpoints.""" rng = np.random.default_rng(seed) x, y = np.asarray(x, float), np.asarray(y, float) - vals = [] - for _ in range(n): - i = rng.integers(0, len(x), len(x)) - vals.append(corr(x[i], y[i], kind)) - vals = np.array([v for v in vals if not math.isnan(v)]) - return (float(np.percentile(vals, 2.5)), float(np.percentile(vals, 97.5))) if len(vals) else (float("nan"),) * 2 + vals = [corr(x[i], y[i], kind) for i in (rng.integers(0, len(x), len(x)) for _ in range(n))] + return float(np.nanpercentile(vals, 2.5)), float(np.nanpercentile(vals, 97.5)) def loglinear(hours, metric) -> dict: - """metric = a + b * log10(hours): slope, intercept, R^2.""" - x, y = np.log10(np.asarray(hours, float)), np.asarray(metric, float) - if len(x) < 2: - return {"slope": float("nan"), "intercept": float("nan"), "r2": float("nan")} - b, a = np.polyfit(x, y, 1) - pred = a + b * x - ss_res, ss_tot = float(((y - pred) ** 2).sum()), float(((y - y.mean()) ** 2).sum()) - return {"slope": float(b), "intercept": float(a), "r2": 1 - ss_res / ss_tot if ss_tot else float("nan")} + """metric = intercept + slope * log10(hours).""" + f = stats.linregress(np.log10(hours), metric) + return {"slope": float(f.slope), "intercept": float(f.intercept), "r2": float(f.rvalue**2)} def main(argv=None) -> int: diff --git a/egomimic/trainHydra.py b/egomimic/trainHydra.py index 23e704391a..7d6cb1d02a 100644 --- a/egomimic/trainHydra.py +++ b/egomimic/trainHydra.py @@ -59,7 +59,23 @@ def _build_model_config_tree(cfg: DictConfig) -> DictConfig: and "norm_stats" in model_cfg.robomimic_model ): model_cfg.robomimic_model.norm_stats = None - tree = OmegaConf.create({"model": model_cfg}) + # Each dataset's keymap + transform list, resolved: what a rollout needs to + # rebuild the training inputs (robot/rollout.py, scripts/abc_sim/policy_server.py + # read data.train_datasets..resolver.*). Nothing else of `data`. + pipelines = { + name: { + "resolver": { + k: OmegaConf.to_container(ds.resolver[k], resolve=True) + for k in ("key_map", "transform_list") + if ds.resolver.get(k) is not None + } + } + for name, ds in (cfg.get("data") or {}).get("train_datasets", {}).items() + if ds is not None and ds.get("resolver") is not None + } + tree = OmegaConf.create( + {"model": model_cfg, "data": {"train_datasets": pipelines}} + ) has_weights = OmegaConf.select(tree, _PI_WEIGHT_KEY, default=None) is not None if has_weights and _weights_from_checkpoint(cfg): log.info( @@ -152,7 +168,7 @@ def _weights_from_checkpoint(cfg: DictConfig) -> bool: return bool(ckpt_path) and os.path.isfile(ckpt_path) and not cfg.get("pretrained") -def _apply_init_weights(cfg: DictConfig, model) -> Optional[dict]: +def _apply_init_weights(cfg: DictConfig, model) -> Optional[str]: """``init_weights_ckpt``: weights-only fine-tune init (see ``init_weights_from_checkpoint``). Skipped when ``ckpt_path`` is set -- a resume (including a Slurm requeue, folded into ``ckpt_path`` by @@ -169,11 +185,8 @@ def _apply_init_weights(cfg: DictConfig, model) -> Optional[dict]: return None if not os.path.isfile(init_ckpt): raise FileNotFoundError(f"init_weights_ckpt {init_ckpt} does not exist") - return init_weights_from_checkpoint( - model, - init_ckpt, - min_fraction=float(cfg.get("init_weights_min_fraction", 0.5)), - ) + init_weights_from_checkpoint(model, init_ckpt) + return init_ckpt def _log_dataset_frame_counts( diff --git a/egomimic/utils/checkpoint_utils.py b/egomimic/utils/checkpoint_utils.py index 06025723b0..2142048dbc 100644 --- a/egomimic/utils/checkpoint_utils.py +++ b/egomimic/utils/checkpoint_utils.py @@ -42,63 +42,31 @@ def load_checkpoint_weights( log.info("Loaded weights from %s", ckpt_path) -def init_weights_from_checkpoint( - model: torch.nn.Module, - ckpt_path: str | os.PathLike, - *, - min_fraction: float = 0.5, -) -> dict: +# A fine-tune init that matched less than this share of the model is not one. +MIN_INIT_FRACTION = 0.5 + + +def init_weights_from_checkpoint(model: torch.nn.Module, ckpt_path: str | os.PathLike) -> None: """Fine-tune initialization: copy every tensor of a Lightning checkpoint's ``state_dict`` whose name and shape match ``model``; leave the rest at their fresh init. Unlike ``ckpt_path`` (a full resume: optimizer, scheduler, epoch) nothing but weights is read, so a new embodiment's stems / head and a different action width start fresh while the shared trunk transfers. - - Raises when fewer than ``min_fraction`` of the model's parameters (by - element count) were loaded: a wrong or unrelated checkpoint would otherwise - "fine-tune" from random init without a word. Returns a report dict. - """ - checkpoint = torch.load( - ckpt_path, map_location="cpu", weights_only=False, mmap=True - ) - ckpt_sd = checkpoint["state_dict"] + Raises when under MIN_INIT_FRACTION of the model's elements matched: an + unrelated checkpoint would otherwise "fine-tune" from random init.""" + ckpt_sd = torch.load(ckpt_path, map_location="cpu", weights_only=False, mmap=True)["state_dict"] own_sd = model.state_dict() - loaded = { - k: v for k, v in ckpt_sd.items() if k in own_sd and own_sd[k].shape == v.shape - } - shape_mismatch = sorted( - k for k, v in ckpt_sd.items() if k in own_sd and own_sd[k].shape != v.shape - ) - unused = sorted(k for k in ckpt_sd if k not in own_sd) + loaded = {k: v for k, v in ckpt_sd.items() if k in own_sd and own_sd[k].shape == v.shape} fresh = sorted(k for k in own_sd if k not in loaded) model.load_state_dict(loaded, strict=False) - - total = sum(v.numel() for v in own_sd.values()) - covered = sum(own_sd[k].numel() for k in loaded) - fraction = covered / max(total, 1) - report = { - "loaded": len(loaded), - "shape_mismatch": shape_mismatch, - "unused": unused, - "fresh": fresh, - "fraction": fraction, - } + fraction = sum(own_sd[k].numel() for k in loaded) / max(sum(v.numel() for v in own_sd.values()), 1) log.info( "init_weights_ckpt %s: loaded %d tensors (%.1f%% of the model's elements); " - "%d shape-mismatched, %d unused, %d left at fresh init, e.g. %s", - ckpt_path, - len(loaded), - 100 * fraction, - len(shape_mismatch), - len(unused), - len(fresh), - fresh[:5], + "%d checkpoint tensors unused; %d left at fresh init, e.g. %s", + ckpt_path, len(loaded), 100 * fraction, len(ckpt_sd) - len(loaded), len(fresh), fresh[:5], ) - if fraction < min_fraction: + if fraction < MIN_INIT_FRACTION: raise RuntimeError( - f"init_weights_ckpt {ckpt_path}: only {100 * fraction:.1f}% of the " - f"model's weights matched (need {100 * min_fraction:.0f}%); is it " - f"from the same model family? e.g. fresh {fresh[:5]}, " - f"mismatched {shape_mismatch[:5]}" + f"init_weights_ckpt {ckpt_path}: only {100 * fraction:.1f}% of the model's " + f"weights matched; is it from the same model family? e.g. fresh {fresh[:5]}" ) - return report diff --git a/tests/unit/test_abc_joint_configs.py b/tests/unit/test_abc_joint_configs.py index 66a11ee50e..3f97d4eb24 100644 --- a/tests/unit/test_abc_joint_configs.py +++ b/tests/unit/test_abc_joint_configs.py @@ -1,20 +1,22 @@ -"""The ABC joint-space recipes compose, and the fine-tune recipe trains: two -optimizer steps of a tiny RDT on synthetic eva JOINT episodes through the real -trainHydra.train() path (Eva "joints" keymap + transform, the eva RDT head and -state stem, joint metrics wiring, image history on the robot camera).""" +"""The ABC sim fine-tune recipe end to end on CPU: two optimizer steps of a +tiny RDT on synthetic eva JOINT episodes through the real trainHydra.train() +path, then the saved checkpoint served by policy_server (load, rebuild the +training inputs from the checkpoint's own config, infer).""" from __future__ import annotations import math -import pytest from fixtures.recipes import LOCAL_RESOLVER, common_overrides, cpu_trainer_overrides from fixtures.train_harness import BatchKeySpy, compose_recipe, hermetic_env, write_fixtures from fixtures.recipes import Recipe +import numpy as np + import egomimic.trainHydra as train_hydra from egomimic.algo.hpt import HPT from egomimic.rldb.embodiment.eva import JOINT_ACTION_KEY, JOINT_STATE_KEY, Eva +from egomimic.scripts.abc_sim.policy_server import JointPolicy EMB = "eva_bimanual" RM = "model.robomimic_model" @@ -37,28 +39,6 @@ def _tiny_rdt(emb: str) -> list[str]: ] -@pytest.mark.parametrize( - "top,domains", - [ - ("train_zarr_abc_sim_rdt_ft", ["eva_bimanual"]), - ("train_zarr_abc_real_rdt_ft", ["eva_bimanual"]), - ("train_zarr_abc_sim_rdt_cotrain", ["human_bimanual", "eva_bimanual"]), - ], -) -def test_abc_recipes_compose(compose_resolve, top, domains): - cfg = compose_resolve(top, []) - rm = cfg.model.robomimic_model - assert list(rm.domains) == domains - assert rm.dims.eva_bimanual.action == 14 and rm.dims.eva_bimanual.proprio == 14 - assert rm.ac_keys.eva_bimanual == JOINT_ACTION_KEY - assert "state_joint_positions" in rm.stem_specs.eva_bimanual - assert rm.width == 2048 # the 1B DiT Ryan asked for - assert cfg.data.train_datasets.eva_bimanual.resolver.transform_list.mode == "joints" - assert cfg.data.train_datasets.eva_bimanual.resolver.key_map.image_history_gap_s == 0.1 - assert "eva_bimanual" not in cfg.evaluator.viz_func # joints have no overlay - assert cfg.init_weights_ckpt is None and cfg.train_viz is False - - def test_abc_sim_finetune_train_step(tmp_path, monkeypatch): hermetic_env(monkeypatch) data, out, hashes = write_fixtures(tmp_path, "eva") @@ -85,3 +65,12 @@ def test_abc_sim_finetune_train_step(tmp_path, monkeypatch): expected = {JOINT_ACTION_KEY, JOINT_STATE_KEY, Eva.VIZ_IMAGE_KEY, f"{Eva.VIZ_IMAGE_KEY}_hist", "fps"} assert expected <= seen, f"missing {expected - seen}; saw {sorted(seen)}" assert "observations.images.left_wrist_img" not in seen # joints_front keymap + + # The checkpoint alone rebuilds the training pipeline and serves it. + ckpt = out / "served.ckpt" + objects["trainer"].save_checkpoint(ckpt) + policy = JointPolicy(str(ckpt), device="cpu") + assert policy.cameras == ["top"] and policy.lag_frames == 3 + img = np.random.default_rng(0).random((3, 64, 64), dtype=np.float32) + actions = policy.infer(np.zeros(14, np.float32), {"top": img, "top_hist": img}, "pick up the red cube") + assert actions.shape == (45, 14) and np.isfinite(actions).all() diff --git a/tests/unit/test_abc_sim_convert.py b/tests/unit/test_abc_sim_convert.py index f32d6a47ab..e37a87c287 100644 --- a/tests/unit/test_abc_sim_convert.py +++ b/tests/unit/test_abc_sim_convert.py @@ -65,11 +65,11 @@ def test_prompt_rule_matches_abc(): }, 10, ) - assert spans == [("sim find the cup", 0, 4), ("sim find the ball", 4, 10)] + assert spans == [("find the cup", 0, 4), ("find the ball", 4, 10)] spans = conv.prompt_spans( {"task_name": "put_relative", "instruction": "put the cube left of the mug"}, 5 ) - assert spans == [("sim put the cube left of the mug", 0, 5)] + assert spans == [("put the cube left of the mug", 0, 5)] def test_intrinsics_are_the_sim_pinhole(): @@ -83,9 +83,10 @@ def test_convert_and_read_back(tmp_path): _write_abc_episode(src, "train", "episode_0001") _write_abc_episode(src, "val", "episode_0002") _write_abc_episode(src, "train", "episode_0003", {"task_name": "sim_pouring_beads"}) + _write_abc_episode(src, "train", "episode_0004", {"prompt_source": {"status": "excluded"}}) assert conv.main(["--src", str(src), "--task", TASK, "--out", str(out), "--workers", "1"]) == 0 report = json.loads((out / "conversion_report.json").read_text()) - assert report["written"] == 2 and not report["failed"] + assert report["written"] == 2 and report["skipped_unlabelled"] == 1 and not report["failed"] # Idempotent: a rerun skips both. conv.main(["--src", str(src), "--task", TASK, "--out", str(out), "--workers", "1"]) assert json.loads((out / "conversion_report.json").read_text())["skipped"] == 2 diff --git a/tests/unit/test_eva_joint_mode.py b/tests/unit/test_eva_joint_mode.py index 4bb72469cf..181ca07443 100644 --- a/tests/unit/test_eva_joint_mode.py +++ b/tests/unit/test_eva_joint_mode.py @@ -83,15 +83,6 @@ def test_joint_keymap_supports_the_hist_camera(tmp_path): assert torch.equal(leaf[20][f"{CAM}_hist"], leaf[17][CAM]) -def test_joint_mode_has_no_revert(): - assert Eva.get_revert_transform_list("joints") is None - - -def test_cartesian_keymap_is_unchanged_by_the_joint_mode(): - km = Eva.get_keymap("cartesian") - assert "left.cmd_ee_pose" in km and "left.cmd_joints" not in km - - def test_joint_actions_get_joint_metrics_not_cartesian(): gt = torch.zeros(2, 100, 14) pred = gt.clone() diff --git a/tests/unit/test_init_weights_ckpt.py b/tests/unit/test_init_weights_ckpt.py index 496cd5c8e6..af8cadea06 100644 --- a/tests/unit/test_init_weights_ckpt.py +++ b/tests/unit/test_init_weights_ckpt.py @@ -29,11 +29,9 @@ def test_matched_tensors_load_and_the_new_head_stays_fresh(tmp_path): ckpt = _save(base, tmp_path / "base.ckpt") target = _Net(head_out=14) fresh_head = target.head.weight.detach().clone() - report = init_weights_from_checkpoint(target, ckpt) + init_weights_from_checkpoint(target, ckpt) assert torch.equal(target.trunk.weight, base.trunk.weight) assert torch.equal(target.head.weight, fresh_head) - assert report["shape_mismatch"] == ["head.bias", "head.weight"] - assert report["fraction"] > 0.5 def test_an_unrelated_checkpoint_raises(tmp_path): @@ -64,4 +62,4 @@ def test_missing_init_checkpoint_raises(tmp_path): def test_train_config_declares_the_knob(compose_resolve): cfg = compose_resolve("train_zarr_cartesian", []) - assert cfg.init_weights_ckpt is None and cfg.init_weights_min_fraction == 0.5 + assert cfg.init_weights_ckpt is None diff --git a/tests/unit/test_policy_server.py b/tests/unit/test_policy_server.py index f41a422aa6..0aac8ae81f 100644 --- a/tests/unit/test_policy_server.py +++ b/tests/unit/test_policy_server.py @@ -1,26 +1,31 @@ """policy_server builds the model input exactly as ZarrDataset builds a -training sample (same keymap, transforms, history frame, fps), normalizes it -with the checkpoint's norm stats, and speaks the sim client's HTTP protocol.""" +training sample (keymap, transforms, the 0.1 s history frame, fps), normalizes +it with the checkpoint's norm stats, returns the chunk on the raw 30 Hz grid, +and speaks the sim client's HTTP protocol.""" from __future__ import annotations import base64 import json +import threading +import urllib.error import urllib.request +from http.server import HTTPServer import numpy as np +import simplejpeg import torch from fixtures.synthetic_episodes import write_episode from omegaconf import OmegaConf from egomimic.rldb.embodiment.eva import JOINT_ACTION_KEY, JOINT_STATE_KEY, Eva -from egomimic.rldb.zarr.zarr_dataset_multi import LocalEpisodeResolver, MultiDataset +from egomimic.rldb.zarr.zarr_dataset_multi import LocalEpisodeResolver, MultiDataset, ZarrEpisode from egomimic.scripts.abc_sim import policy_server as ps CAM = Eva.VIZ_IMAGE_KEY CFG = OmegaConf.create({"data": {"train_datasets": {"eva_bimanual": {"resolver": { - "key_map": {"_target_": "x", "keymap_mode": "joints_front", "annotation_key": "annotations", "image_history_gap_s": 0.1}, - "transform_list": {"_target_": "x", "mode": "joints"}, + "key_map": {"keymap_mode": "joints_front", "annotation_key": "annotations", "image_history_gap_s": 0.1}, + "transform_list": {"mode": "joints"}, }}}}}) @@ -28,24 +33,21 @@ class _FakeNormStats: """Scales the state by 2 so the test can see normalization happened.""" def normalize(self, data, emb_id): - out = dict(data) - out[JOINT_STATE_KEY] = data[JOINT_STATE_KEY] * 2 - return out + return {**data, JOINT_STATE_KEY: data[JOINT_STATE_KEY] * 2} class _FakeAlgo: - """Records the batch it was given; predicts the (normalized) state tiled.""" + """Records the batch; predicts a ramp in time so the resampling is visible.""" norm_stats = _FakeNormStats() - annotation_key = "annotations" def process_batch_for_training(self, batch): self.seen = batch["eva_bimanual"] return batch def forward_eval(self, batch): - state = batch["eva_bimanual"][JOINT_STATE_KEY] - return {f"eva_bimanual_{JOINT_ACTION_KEY}": state[:, None, :].repeat(1, 100, 1)} + ramp = torch.linspace(0, 99, 100)[None, :, None].repeat(1, 1, 14) + return {f"eva_bimanual_{JOINT_ACTION_KEY}": ramp} def _policy(): @@ -55,74 +57,68 @@ def _policy(): return pol -def _episode_frames(root): - """Raw frames/states of the synthetic eva episode as the sim would send - them, read through the dataset's own reader and JPEG decode.""" - import simplejpeg - - from egomimic.rldb.zarr.zarr_dataset_multi import ZarrEpisode - - ep = ZarrEpisode(next(root.glob("*.zarr"))) - n = ep.metadata["total_frames"] - keys = ["images.front_1", "left.obs_joints", "left.obs_gripper", "right.obs_joints", "right.obs_gripper"] - got = ep.read_intervals({k: [(0, n)] for k in keys}) - rows = {k: got[k][0][1] for k in keys} - frames = [simplejpeg.decode_jpeg(rows["images.front_1"][i], colorspace="RGB") for i in range(n)] - state = np.concatenate([np.asarray(rows[k], dtype=np.float32) for k in keys[1:]], axis=1) - return frames, state +def _chw(hwc: np.ndarray) -> np.ndarray: + return np.moveaxis(hwc, -1, 0).astype(np.float32) / np.float32(255.0) def test_server_sample_matches_the_dataset_sample(tmp_path): write_episode(tmp_path, "eva", T=48) - frames, state = _episode_frames(tmp_path) + ep = ZarrEpisode(next(tmp_path.glob("*.zarr"))) + keys = ["images.front_1", "left.obs_joints", "left.obs_gripper", "right.obs_joints", "right.obs_gripper"] + rows = {k: g[0][1] for k, g in ep.read_intervals({k: [(0, 48)] for k in keys}).items()} + frames = [simplejpeg.decode_jpeg(rows["images.front_1"][i], colorspace="RGB") for i in range(48)] + state = np.concatenate([np.asarray(rows[k], np.float32) for k in keys[1:]], axis=1) key_map = Eva.get_keymap("joints_front", image_history_gap_s=0.1) resolver = LocalEpisodeResolver(tmp_path, key_map=key_map, transform_list=Eva.get_transform_list("joints")) leaf = next(iter(MultiDataset._from_resolver(resolver, mode="total").datasets.values())) pol = _policy() - for t in range(0, 12): # feed the frames in order so the history fills like an episode - sample = pol.build_sample("ep", t / 30.0, state[t], {"top": np.moveaxis(frames[t], -1, 0).astype(np.float32) / 255}) - ref = leaf[t] + assert pol.cameras == ["top"] and pol.lag_frames == 3 + for t in (0, 1, 3, 10): + images = {"top": _chw(frames[t])} + if t > 0: # the frame lag_frames back, clamped to the episode start; at t=0 none is sent + images["top_hist"] = _chw(frames[max(0, t - pol.lag_frames)]) + sample, ref = pol.build_sample(state[t], images), leaf[t] assert torch.equal(sample[CAM], ref[CAM]), t - assert torch.equal(sample[f"{CAM}_hist"], ref[f"{CAM}_hist"]), t # 3 frames back, clamped at the start + assert torch.equal(sample[f"{CAM}_hist"], ref[f"{CAM}_hist"]), t assert torch.allclose(sample[JOINT_STATE_KEY], ref[JOINT_STATE_KEY].float()), t assert sample[JOINT_ACTION_KEY].shape == ref[JOINT_ACTION_KEY].shape == (100, 14) assert sample["fps"].item() == ref["fps"].item() == 30 - assert "observations.images.left_wrist_img" not in sample -def test_infer_normalizes_then_returns_unnormalized_chunk(tmp_path): +def test_infer_normalizes_and_returns_the_raw_30hz_grid(): pol = _policy() - img = np.zeros((3, 16, 24), np.float32) - state = np.arange(14, dtype=np.float32) - out = pol.infer("ep", 0.0, state, {"top": img}, "sim put the bottles in the bin") - assert out.shape == (100, 14) - seen = pol.algo.seen - assert torch.allclose(seen[JOINT_STATE_KEY][0], torch.arange(14.0) * 2) # normalized before the model - assert seen["annotations"] == [["sim put the bottles in the bin"]] - assert seen["fps"].shape == (1,) + out = pol.infer(np.arange(14, dtype=np.float32), {"top": np.zeros((3, 16, 24), np.float32)}, "sim x") + # 100 model steps span 45 raw frames: raw frame k sits at model step 99k/44 + assert out.shape == (45, 14) + assert np.allclose(out[:, 0], np.arange(45) * 99 / 44, atol=1e-4) + assert torch.allclose(pol.algo.seen[JOINT_STATE_KEY][0], torch.arange(14.0) * 2) + assert pol.algo.seen["annotations"] == [["sim x"]] def test_http_protocol_round_trip(): pol = _policy() - server = ps.serve(pol, "127.0.0.1", 0) - port = server.server_address[1] - try: - with urllib.request.urlopen(f"http://127.0.0.1:{port}/health") as r: - assert json.load(r)["mode"] == "joints" - hwc = np.random.default_rng(0).integers(0, 255, (16, 24, 3), np.uint8) - payload = {"episode": "e", "t": 0.0, "state": list(range(14)), "prompt": "sim x", - "images": {"top": {"shape": [16, 24, 3], "b64": base64.b64encode(hwc.tobytes()).decode()}}} - req = urllib.request.Request(f"http://127.0.0.1:{port}/infer", data=json.dumps(payload).encode(), + server = HTTPServer(("127.0.0.1", 0), ps.make_handler(pol)) + threading.Thread(target=server.serve_forever, daemon=True).start() + url = f"http://127.0.0.1:{server.server_address[1]}" + + def post(payload): + req = urllib.request.Request(f"{url}/infer", data=json.dumps(payload).encode(), headers={"Content-Type": "application/json"}, method="POST") with urllib.request.urlopen(req) as r: - out = json.load(r) - assert np.asarray(out["actions"]).shape == (100, 14) and out["dt"] == 1 / 30 - assert torch.allclose(pol.algo.seen[CAM][0], torch.from_numpy(np.moveaxis(hwc, -1, 0).astype(np.float32) / 255)) - bad = urllib.request.Request(f"http://127.0.0.1:{port}/infer", data=json.dumps({"episode": "e", "t": 0, "state": [0], "images": {}}).encode(), - headers={"Content-Type": "application/json"}, method="POST") + return json.load(r) + + try: + with urllib.request.urlopen(f"{url}/health") as r: + health = json.load(r) + assert health["cameras"] == ["top"] and health["lag_frames"] == 3 and health["horizon"] == 45 + hwc = np.random.default_rng(0).integers(0, 255, (16, 24, 3), np.uint8) + out = post({"state": list(range(14)), "prompt": "sim x", + "images": {"top": {"shape": [16, 24, 3], "b64": base64.b64encode(hwc.tobytes()).decode()}}}) + assert np.asarray(out["actions"]).shape == (45, 14) and out["dt"] == 1 / 30 + assert torch.allclose(pol.algo.seen[CAM][0], torch.from_numpy(_chw(hwc))) try: - urllib.request.urlopen(bad) - assert False, "expected 500" + post({"state": [0], "images": {}}) + raise AssertionError("expected 500") except urllib.error.HTTPError as e: assert e.code == 500 and "error" in json.load(e) finally: diff --git a/tests/unit/test_val_vs_sim.py b/tests/unit/test_val_vs_sim.py index 75144a474b..f02d6ae231 100644 --- a/tests/unit/test_val_vs_sim.py +++ b/tests/unit/test_val_vs_sim.py @@ -19,10 +19,3 @@ def test_report_from_manifest(tmp_path): assert rep["correlations"]["spearman_success_rate"]["r"] == -1.0 # lower val mse, higher success assert rep["loglinear_val_vs_hours"]["r2"] > 0.95 and rep["loglinear_val_vs_hours"]["slope"] < 0 assert (tmp_path / "out" / "report.md").read_text().count("| r") >= 4 - - -def test_corr_helpers(): - assert vs.corr([1, 2, 3], [3, 2, 1], "spearman") == -1.0 - assert vs.corr([1, 1, 1], [1, 2, 3], "pearson") != vs.corr([1, 1, 1], [1, 2, 3], "pearson") # nan - lo, hi = vs.bootstrap([1, 2, 3, 4], [1, 2, 3, 4], "pearson", n=50) - assert lo <= 1.0 <= hi From 94086047d7492d82026fdeda76cb7cc6374ce7c0 Mon Sep 17 00:00:00 2001 From: Aidan Gao Date: Tue, 29 Sep 2026 11:12:21 -0400 Subject: [PATCH 09/20] lambda_pack: stage data from a node step; run onto a running job's spare GPUs The /ephemeral staging ran in the batch shell, so the documented "bash this file with SLURM_JOB_ID=" use copied onto the gateway, which has no /ephemeral. It is now a step. Run steps --overlap so they can share a running allocation's (mostly idle) CPUs; harmless in an exclusive one. Co-Authored-By: Claude Opus 5.5 --- egomimic/scripts/abc_sim/lambda_pack.sbatch | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/egomimic/scripts/abc_sim/lambda_pack.sbatch b/egomimic/scripts/abc_sim/lambda_pack.sbatch index 06c7fe35e8..bc3d9d0443 100644 --- a/egomimic/scripts/abc_sim/lambda_pack.sbatch +++ b/egomimic/scripts/abc_sim/lambda_pack.sbatch @@ -12,6 +12,8 @@ # TABLE=... EGOVERSE_EPISODE_TABLE pickle (default: the one in DATA_SRC) # RUN_ROOT=... default /workspace/users/$USER/runs_abc # EXTRA="..." appended to every run (smokes: trainer.max_epochs=1 ...) +# Onto a RUNNING job's spare GPUs: SLURM_JOB_ID= bash egomimic/scripts/abc_sim/lambda_pack.sbatch +# (steps --overlap its CPUs; they die with that job, so queue an afterany follow-up to resume). #SBATCH --job-name=abc_pack #SBATCH --account=loaner #SBATCH --partition=loaner @@ -35,7 +37,8 @@ read -r -a runs <<< "${RUNS:?set RUNS to space-separated run names, each with an CPUS=${CPUS_PER_TASK:-$(( 192 / (${#runs[@]} * G) ))} test -s "$IMAGE" && test -s "$SRC/.venv/pyvenv.cfg" && test -s "$TABLE" echo "=== abc_pack job=$SLURM_JOB_ID node=$(hostname) restart=${SLURM_RESTART_COUNT:-0} commit=$(git -C "$SRC" rev-parse --short HEAD) runs=${runs[*]}" -mkdir -p "$DATA" && rsync -a --exclude _sync "$DATA_SRC/" "$DATA/" && echo "staged $(du -sh "$DATA" | cut -f1) to $DATA" +# a step, so the copy lands on the compute node's /ephemeral also when this file is run with bash +srun --overlap --gres=none --ntasks=1 bash -c "mkdir -p $DATA && rsync -a --exclude _sync $DATA_SRC/ $DATA/ && echo staged \$(du -sh $DATA | cut -f1) to $DATA" i=0 for name in "${runs[@]}"; do @@ -44,7 +47,7 @@ for name in "${runs[@]}"; do run="$RUN_ROOT/$name"; mkdir -p "$run" resume=""; [ -s "$run/checkpoints/last.ckpt" ] && resume="ckpt_path=$run/checkpoints/last.ckpt" echo "--- $name -> $run ${resume:+(resuming)}: $args" - srun --exact --ntasks=$G --ntasks-per-node=$G --gres=gpu:h100:$G --gpu-bind=none --cpus-per-task=$CPUS --mem=200G \ + srun --exact --overlap --ntasks=$G --ntasks-per-node=$G --gres=gpu:h100:$G --gpu-bind=none --cpus-per-task=$CPUS --mem=200G \ --container-image="$IMAGE" --container-mounts=/workspace:/workspace:rw,/ephemeral:/ephemeral:rw --container-workdir="$SRC" \ --output="$run/train-%j.log" \ bash -c "set -euo pipefail From 4d0bae8ec6acb4f0d6575c387ad0f3c396b7905f Mon Sep 17 00:00:00 2001 From: Aidan Gao Date: Tue, 29 Sep 2026 11:16:07 -0400 Subject: [PATCH 10/20] lambda_pack: stage to node-local /tmp (/ephemeral is root-only on loaner nodes) Co-Authored-By: Claude Opus 5.5 --- egomimic/scripts/abc_sim/lambda_pack.sbatch | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/egomimic/scripts/abc_sim/lambda_pack.sbatch b/egomimic/scripts/abc_sim/lambda_pack.sbatch index bc3d9d0443..dbf47721cf 100644 --- a/egomimic/scripts/abc_sim/lambda_pack.sbatch +++ b/egomimic/scripts/abc_sim/lambda_pack.sbatch @@ -2,13 +2,14 @@ # Pack several small-DDP main-EgoVerse runs on ONE exclusive Lambda loaner node # (8 x H100), one srun step per run inside the pytorch image, each on its own # GPUs / DDP port, resuming from its last.ckpt on requeue. The dataset dir is -# copied to the node-local /ephemeral first: zarr reads straight off the NFS -# /workspace stalled nodes and left cancelled jobs undead (Ryan, 2026-09-26). +# copied to node-local disk first (/tmp; the /ephemeral array is root-only on +# the loaner nodes): zarr reads straight off the NFS /workspace stalled nodes +# and left cancelled jobs undead (Ryan, 2026-09-26). # # RUNS="k1 k2 k4" ARGS_k1="--config-name=train_zarr_fold_ladder_rdt1b data=mecka_fold_freeform_opscale1_6d" ... \ # sbatch --qos=loaner-normal egomimic/scripts/abc_sim/lambda_pack.sbatch # GPUS_PER_RUN=2 must match launch_params.gpus_per_node of the configs -# DATA_SRC=... NFS dataset dir (default datasets/mecka_abc_ladder), staged to /ephemeral +# DATA_SRC=... NFS dataset dir (default datasets/fold_ladder), staged to node-local /tmp # TABLE=... EGOVERSE_EPISODE_TABLE pickle (default: the one in DATA_SRC) # RUN_ROOT=... default /workspace/users/$USER/runs_abc # EXTRA="..." appended to every run (smokes: trainer.max_epochs=1 ...) @@ -28,16 +29,16 @@ set -euo pipefail SRC=${SRC:-/workspace/users/$USER/EgoVerse} IMAGE=${IMAGE:-/workspace/users/$USER/images/nvidia-pytorch-25.06-py3.sqsh} RUN_ROOT=${RUN_ROOT:-/workspace/users/$USER/runs_abc} -DATA_SRC=${DATA_SRC:-/workspace/users/$USER/datasets/mecka_abc_ladder} +DATA_SRC=${DATA_SRC:-/workspace/users/$USER/datasets/fold_ladder} TABLE=${TABLE:-$(ls "$DATA_SRC"/episode_table_*.pkl | tail -1)} -DATA=/ephemeral/$USER/$(basename "$DATA_SRC") +DATA=/tmp/$USER/$(basename "$DATA_SRC") G=${GPUS_PER_RUN:-2} read -r -a runs <<< "${RUNS:?set RUNS to space-separated run names, each with an ARGS_}" (( ${#runs[@]} * G <= 8 )) || { echo "${#runs[@]} runs x $G GPUs do not fit one node" >&2; exit 2; } CPUS=${CPUS_PER_TASK:-$(( 192 / (${#runs[@]} * G) ))} test -s "$IMAGE" && test -s "$SRC/.venv/pyvenv.cfg" && test -s "$TABLE" echo "=== abc_pack job=$SLURM_JOB_ID node=$(hostname) restart=${SLURM_RESTART_COUNT:-0} commit=$(git -C "$SRC" rev-parse --short HEAD) runs=${runs[*]}" -# a step, so the copy lands on the compute node's /ephemeral also when this file is run with bash +# a step, so the copy lands on the compute node also when this file is run with bash srun --overlap --gres=none --ntasks=1 bash -c "mkdir -p $DATA && rsync -a --exclude _sync $DATA_SRC/ $DATA/ && echo staged \$(du -sh $DATA | cut -f1) to $DATA" i=0 @@ -48,7 +49,7 @@ for name in "${runs[@]}"; do resume=""; [ -s "$run/checkpoints/last.ckpt" ] && resume="ckpt_path=$run/checkpoints/last.ckpt" echo "--- $name -> $run ${resume:+(resuming)}: $args" srun --exact --overlap --ntasks=$G --ntasks-per-node=$G --gres=gpu:h100:$G --gpu-bind=none --cpus-per-task=$CPUS --mem=200G \ - --container-image="$IMAGE" --container-mounts=/workspace:/workspace:rw,/ephemeral:/ephemeral:rw --container-workdir="$SRC" \ + --container-image="$IMAGE" --container-mounts=/workspace:/workspace:rw,$DATA:$DATA:ro --container-workdir="$SRC" \ --output="$run/train-%j.log" \ bash -c "set -euo pipefail ulimit -n \$(ulimit -Hn) From 9497082f9fc064339d451a88af799f8e4fa9dfff Mon Sep 17 00:00:00 2001 From: Aidan Gao Date: Tue, 29 Sep 2026 11:22:18 -0400 Subject: [PATCH 11/20] eval: the val subsample K is per-rank; the val loaders are not sharded _metric_frames_per_episode scaled K by devices * num_nodes on the premise that a DistributedSampler strides the val split across ranks. The val heads come back inside CombinedLoaders, which Lightning does not wrap (every rank runs every val batch; measured 2026-09-16, EvalVideo._video_fps), so on W GPUs each rank scored the first 1/W of a W-times-denser subsample, in episode-hash order: on the freeform unseen-operator head at 2 GPUs, half of the 35 episodes. K is now the single-rank value on every rank. Co-Authored-By: Claude Opus 5.5 --- egomimic/trainHydra.py | 44 +++++++-------------------------- tests/unit/test_eval_loaders.py | 43 +++++++------------------------- 2 files changed, 18 insertions(+), 69 deletions(-) diff --git a/egomimic/trainHydra.py b/egomimic/trainHydra.py index 7d6cb1d02a..ff24b5352f 100644 --- a/egomimic/trainHydra.py +++ b/egomimic/trainHydra.py @@ -346,32 +346,6 @@ def _unseen_op_valid_datasets(cfg: DictConfig, instantiate) -> dict: } -def _trainer_world_size(cfg: DictConfig) -> int: - """``devices * num_nodes`` as the config declares it, 1 if it cannot say. - - DistributedSampler strides rather than chunks, so W ranks between them walk - the WHOLE split; the per-rank ``limit_val_batches`` window is therefore W - times wider than it looks from one rank. Eval mode runs on one device - (``train`` forces ``devices=1`` after the loaders are built), so it is 1. - """ - legacy = "train" if cfg.get("train") else "eval" if cfg.get("eval") else None - if (cfg.get("mode") or legacy) == "eval": # _resolve_mode, minus its raise - return 1 - devices = cfg.get("trainer", {}).get("devices", 1) - if isinstance(devices, (list, ListConfig)): - n = len(devices) - elif isinstance(devices, int) and devices > 0: - n = devices - else: # "auto", -1: resolved by lightning at runtime, unknown here - log.warning( - f"trainer.devices={devices!r} is not a count; treating the metric " - "loaders as single-rank, which under-scores the val split." - ) - n = 1 - nodes = cfg.get("trainer", {}).get("num_nodes", 1) - return n * (int(nodes) if isinstance(nodes, int) and nodes > 0 else 1) - - def _metric_frames_per_episode( cfg: DictConfig, head: str, @@ -380,14 +354,14 @@ def _metric_frames_per_episode( ) -> int | None: """Frames per episode to keep on this val head, or None for no subsampling. - ``data.metric_frames_per_episode[head]`` is written for ONE rank -- it is - ``floor(limit_val_batches * batch_size / n_episodes)`` -- so it is scaled by - the world size here: the ranks stride through the subsampled set together - and each still reads at most ``limit_val_batches`` batches. Without this a - 4-GPU run scores a quarter of the frames the same config scores on 1 GPU, - which on the seen-val head is fewer than it scored before subsampling - existed at all. EvenStrideDataset keeps a whole episode when K exceeds its - length, so a split that fits entirely is not subsampled. + ``data.metric_frames_per_episode[head]`` is ``floor(limit_val_batches * + batch_size / n_episodes)``, the SAME on every rank: the val heads come back + inside CombinedLoaders, which Lightning does not wrap in a + DistributedSampler (every rank runs every val batch, measured 2026-09-16; + see EvalVideo._video_fps), so a K scaled by the world size only made each + rank score the first 1/W of a W-times-denser subsample. EvenStrideDataset + keeps a whole episode when K exceeds its length, so a split that fits + entirely is not subsampled. ``auto`` computes that formula from the resolved split (``n_episodes``) and the head's loader ``batch_size``, for splits defined by live SQL filters @@ -425,7 +399,7 @@ def _metric_frames_per_episode( f"data.metric_frames_per_episode.{head} must be a positive int " f"or 'auto', got {k}" ) - return k * _trainer_world_size(cfg) + return k def _subsample_val_datasets( diff --git a/tests/unit/test_eval_loaders.py b/tests/unit/test_eval_loaders.py index 1d8b92f22f..2760afa852 100644 --- a/tests/unit/test_eval_loaders.py +++ b/tests/unit/test_eval_loaders.py @@ -315,39 +315,15 @@ def test_every_head_can_be_subsampled_independently(): th._metric_frames_per_episode(_cfg(metric={"valid": 0}), "valid") -def test_k_scales_with_the_world_size(): - """K in the config is the single-rank value. DistributedSampler strides, - so W ranks between them read W times the frames in the same number of - batches each; without the scaling a 4-GPU run scores a quarter of what the - same config scores on one GPU.""" +def test_k_is_the_same_on_every_rank(): + """The val loaders are not sharded (no DistributedSampler inside the + CombinedLoaders; every rank runs every val batch), so K is per-rank and + must not grow with the world size: a W-scaled K had each rank score only + the first 1/W of the subsample.""" cfg = _cfg(metric={"valid": 7}) - for devices, nodes, expected in ( - (1, 1, 7), - (4, 1, 28), - (8, 2, 112), - ([0, 1], 1, 14), - ("auto", 1, 7), # unknown at config time: stay single-rank - (-1, 1, 7), - ): + for devices, nodes in ((1, 1), (4, 1), (8, 2), ([0, 1], 1)): cfg.trainer = {"devices": devices, "num_nodes": nodes} - assert th._metric_frames_per_episode(cfg, "valid") == expected, devices - # a split shorter than K is kept whole, so scaling cannot truncate - cfg.trainer = {"devices": 4, "num_nodes": 1} - split = {HUMAN: _split({"ep0": 20})} - assert len(th._subsample_val_datasets(cfg, "valid", split)[HUMAN]) == 20 - - -def test_eval_mode_is_single_rank(): - """Eval runs are forced onto one device after the loaders are built, so - their val split must be subsampled for one rank, not the training W.""" - cfg = _cfg(metric={"valid": 7}) - cfg.trainer = {"devices": 4, "num_nodes": 2} - assert th._metric_frames_per_episode(cfg, "valid") == 56 - cfg.mode = "eval" - assert th._metric_frames_per_episode(cfg, "valid") == 7 - del cfg["mode"] - cfg.eval = True - assert th._metric_frames_per_episode(cfg, "valid") == 7 + assert th._metric_frames_per_episode(cfg, "valid") == 7, devices @pytest.mark.parametrize("trainer_cfg", ["ddp", "ddp_pi"]) @@ -362,18 +338,17 @@ def test_ddp_devices_are_per_node(trainer_cfg, compose_resolve): ], ) assert cfg.trainer.devices == 8 and cfg.trainer.num_nodes == 2 - assert th._trainer_world_size(cfg) == 16 def test_auto_k_is_derived_from_the_resolved_split(): """`auto` = floor(limit_val_batches * batch_size / episodes), per dataset, - from the head's own loader batch size, then scaled by the world size.""" + from the head's own loader batch size, the same on every rank.""" cfg = _cfg(metric={"valid": "auto"}) cfg.trainer = {"limit_val_batches": 4, "devices": 2} split = {HUMAN: _split({"ep0": 100, "ep1": 100, "ep2": 100})} params = {HUMAN: {"batch_size": 5}} wrapped = th._subsample_val_datasets(cfg, "valid", split, params)[HUMAN] - assert wrapped.frames_per_episode == (4 * 5 // 3) * 2 + assert wrapped.frames_per_episode == 4 * 5 // 3 for limit in (1.0, None, True): cfg.trainer.limit_val_batches = limit with pytest.raises(ValueError, match="auto"): From d6bc6ff22b047cd87153194f76861212e4bbf1b2 Mon Sep 17 00:00:00 2001 From: Aidan Gao Date: Tue, 29 Sep 2026 11:22:18 -0400 Subject: [PATCH 12/20] fold ladder: the true replay stride (29/99 frames per chunk step) The replay's stride inference only fits keypoint layouts and falls back to 1 on the 12-D cartesian revert: a 3.4x-slow GT/pred replay that drifts off the hand. The mecka cartesian chunk is 30 raw frames resampled to 100. Co-Authored-By: Claude Opus 5.5 --- .../hydra_configs/train_zarr_fold_ladder_rdt1b.yaml | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/egomimic/hydra_configs/train_zarr_fold_ladder_rdt1b.yaml b/egomimic/hydra_configs/train_zarr_fold_ladder_rdt1b.yaml index 8aa017c173..3577c0d002 100644 --- a/egomimic/hydra_configs/train_zarr_fold_ladder_rdt1b.yaml +++ b/egomimic/hydra_configs/train_zarr_fold_ladder_rdt1b.yaml @@ -33,3 +33,13 @@ trainer: callbacks: model_checkpoint: every_n_epochs: 100 + +# The GT/pred replay needs the true frames per chunk step: the mecka cartesian +# chunk is 30 raw frames resampled to 100 steps (29/99), and the replay's own +# stride inference only fits keypoint layouts (it falls back to 1: a 3.4x-slow +# replay that drifts off the hand). The train_viz head composes its own copy. +evaluator: + action_stride: 0.29292929292929293 +train_viz_evaluator: + base: + action_stride: 0.29292929292929293 From d3e8f20d90726526d9bb5716d0715f1ca89a926b Mon Sep 17 00:00:00 2001 From: Aidan Gao Date: Tue, 29 Sep 2026 13:43:58 -0400 Subject: [PATCH 13/20] rollout_lambda: runnable as a one-GPU step inside a held job The srun step now asks for exactly one GPU (--exact --overlap), so the script can be run with bash and SLURM_JOB_ID= next to other steps (fine-tunes, other rollouts); PORT is overridable for that case. Co-Authored-By: Claude Opus 5.5 --- egomimic/scripts/abc_sim/rollout_lambda.sbatch | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/egomimic/scripts/abc_sim/rollout_lambda.sbatch b/egomimic/scripts/abc_sim/rollout_lambda.sbatch index 029a522ee0..65e04cc7d4 100644 --- a/egomimic/scripts/abc_sim/rollout_lambda.sbatch +++ b/egomimic/scripts/abc_sim/rollout_lambda.sbatch @@ -7,6 +7,7 @@ # egomimic/scripts/abc_sim/rollout_lambda.sbatch # TASK: any abc_sim name or alias (the sim_224 task_name); NUM_WORLDS=50; # POLICY=hold for the no-motion floor (no server); CLIENT_EXTRA e.g. --save-video +# Onto a running job: SLURM_JOB_ID= bash egomimic/scripts/abc_sim/rollout_lambda.sbatch #SBATCH --job-name=abc_rollout #SBATCH --account=loaner #SBATCH --partition=loaner @@ -25,10 +26,10 @@ TASK=${TASK:-sim_put_the_plastic_bottles_in_the_bin} : "${OUT:?set OUT}" POLICY=${POLICY:-served} [ "$POLICY" = hold ] || : "${CKPT:?set CKPT}" -PORT=$(( 8700 + SLURM_JOB_ID % 200 )) +PORT=${PORT:-$(( 8700 + SLURM_JOB_ID % 200 ))} # set PORT per step when several share a job mkdir -p "$OUT" echo "=== abc_rollout job=$SLURM_JOB_ID node=$(hostname) task=$TASK ckpt=${CKPT:-none} policy=$POLICY out=$OUT" -srun --ntasks=1 --container-image="$IMAGE" --container-mounts=/workspace:/workspace:rw --container-workdir="$SRC" \ +srun --exact --overlap --ntasks=1 --gres=gpu:h100:1 --cpus-per-task=${CPUS_PER_TASK:-16} --mem=120G --container-image="$IMAGE" --container-mounts=/workspace:/workspace:rw --container-workdir="$SRC" \ bash -c "set -euo pipefail export HOME=/workspace/users/$USER HF_HOME=/workspace/users/$USER/.cache/huggingface HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 export PYTHONPATH=$SRC PYTHONUNBUFFERED=1 OMP_NUM_THREADS=4 From b5f3b7ffc796922ffc07dc246df829fda19241ec Mon Sep 17 00:00:00 2001 From: Aidan Gao Date: Tue, 29 Sep 2026 14:23:27 -0400 Subject: [PATCH 14/20] abc_sim: ponytail-review cuts (-46 lines) Flags and fields nothing sets or reads: converter --splits/--overwrite and three unread metadata keys; policy_server --host/--device, /health ok+horizon and the response dt; sim_client --camera-backend/--gpu-id (ABC's defaults) and an unused constructor arg; val_vs_sim --no-wandb (a row carrying val is the signal). One-caller helpers inlined (episode_hash, the joint layout dict), unused kwargs dropped from the joint transform builder, the init-weights file pre-check dropped (torch.load raises the same error), and a config-echo test removed. Co-Authored-By: Claude Opus 5.5 --- egomimic/eval/action_metrics.py | 9 ++--- egomimic/rldb/embodiment/eva.py | 33 +++++-------------- .../scripts/abc_sim/convert_sim_to_zarr.py | 28 +++++----------- egomimic/scripts/abc_sim/policy_server.py | 17 ++++------ egomimic/scripts/abc_sim/sim_client.py | 8 ++--- egomimic/scripts/abc_sim/val_vs_sim.py | 3 +- egomimic/trainHydra.py | 9 ++--- tests/unit/test_init_weights_ckpt.py | 5 --- tests/unit/test_policy_server.py | 4 +-- tests/unit/test_val_vs_sim.py | 2 +- 10 files changed, 36 insertions(+), 82 deletions(-) diff --git a/egomimic/eval/action_metrics.py b/egomimic/eval/action_metrics.py index f70346ee26..497a79e94b 100644 --- a/egomimic/eval/action_metrics.py +++ b/egomimic/eval/action_metrics.py @@ -202,11 +202,6 @@ def _kp_l2(p, g): return metrics -# Bimanual joint layout (ABC YAM / abc_sim): [left j1..j6, left grip, -# right j1..j6, right grip]. -JOINT_LAYOUT_14 = {"joints": (0, 1, 2, 3, 4, 5, 7, 8, 9, 10, 11, 12), "grip": (6, 13)} - - def joint_metrics(pred: torch.Tensor, gt: torch.Tensor, prefix: str) -> dict: """Metrics for a bimanual joint-space chunk ``(B, T, 14)``; ``{}`` for other widths. Joint errors are in degrees (mean absolute), the gripper in @@ -214,8 +209,8 @@ def joint_metrics(pred: torch.Tensor, gt: torch.Tensor, prefix: str) -> dict: pred_cpu, gt_cpu = pred.detach().cpu().float(), gt.detach().cpu().float() if pred_cpu.shape[-1] != 14: return {} - j = list(JOINT_LAYOUT_14["joints"]) - g = list(JOINT_LAYOUT_14["grip"]) + # [left j1..j6, left grip, right j1..j6, right grip] (ABC YAM / abc_sim) + j, g = [0, 1, 2, 3, 4, 5, 7, 8, 9, 10, 11, 12], [6, 13] err = (pred_cpu[..., j] - gt_cpu[..., j]).abs() * _RAD2DEG return { f"{prefix}_joint_abs_err_deg_avg": err.mean(), diff --git a/egomimic/rldb/embodiment/eva.py b/egomimic/rldb/embodiment/eva.py index 7ed667ff6e..00e4cad995 100644 --- a/egomimic/rldb/embodiment/eva.py +++ b/egomimic/rldb/embodiment/eva.py @@ -692,32 +692,17 @@ def _build_eva_bimanual_transform_list( return transform_list -def _build_eva_bimanual_joint_transform_list( - *, - actions_key: str = JOINT_ACTION_KEY, - obs_key: str = JOINT_STATE_KEY, - chunk_length: int = 100, - stride: int = 1, -) -> list[Transform]: +def _build_eva_bimanual_joint_transform_list() -> list[Transform]: """Joint-space pipeline: resample each commanded joint / gripper chunk to - ``chunk_length`` (as the cartesian modes do), then concatenate both arms in - ABC order into ``actions_joints`` (T, 14) and the observed joints into + 100 steps (as the cartesian modes do), then concatenate both arms in ABC + order into ``actions_joints`` (T, 14) and the observed joints into ``observations.state.joint_positions`` (14,).""" cmd_keys = [k.format("cmd") for k in JOINT_ORDER] obs_keys = [k.format("obs") for k in JOINT_ORDER] - transform_list: list[Transform] = [ - InterpolateLinear( - new_chunk_length=chunk_length, - action_key=key, - output_action_key=key, - stride=stride, - ) - for key in cmd_keys - ] - transform_list += [ - InterpolatePadMask(new_chunk_length=chunk_length, stride=stride), - ConcatKeys(key_list=cmd_keys, new_key_name=actions_key, delete_old_keys=True), - ConcatKeys(key_list=obs_keys, new_key_name=obs_key, delete_old_keys=True), - NumpyToTensor(keys=[actions_key, obs_key]), + return [ + *(InterpolateLinear(new_chunk_length=100, action_key=k, output_action_key=k) for k in cmd_keys), + InterpolatePadMask(new_chunk_length=100), + ConcatKeys(key_list=cmd_keys, new_key_name=JOINT_ACTION_KEY, delete_old_keys=True), + ConcatKeys(key_list=obs_keys, new_key_name=JOINT_STATE_KEY, delete_old_keys=True), + NumpyToTensor(keys=[JOINT_ACTION_KEY, JOINT_STATE_KEY]), ] - return transform_list diff --git a/egomimic/scripts/abc_sim/convert_sim_to_zarr.py b/egomimic/scripts/abc_sim/convert_sim_to_zarr.py index 24af7b97e1..c0cab4d70b 100644 --- a/egomimic/scripts/abc_sim/convert_sim_to_zarr.py +++ b/egomimic/scripts/abc_sim/convert_sim_to_zarr.py @@ -14,7 +14,7 @@ ``images.{front_1,left_wrist,right_wrist}``, ``annotations``), so ``Eva.get_keymap("joints")`` reads real and sim ABC data alike. Camera map: top -> front_1, left -> left_wrist, right -> right_wrist. Extra metadata: -``lab: abc_sim``, ``split`` (ABC's own train/val), ``abc_episode``, ``prompt``. +``lab: abc_sim`` and ``split`` (ABC's own train/val). Prompts follow ABC's training rule (``abc_minimal.dataloader``) verbatim: the episode's ``prompt_timeline`` if it has one, else its ``instruction`` when that @@ -128,17 +128,13 @@ def split_arms(x: np.ndarray, kind: str) -> dict[str, np.ndarray]: } -def episode_hash(split: str, ep_dir: Path) -> str: - return f"abcsim_{split}_{ep_dir.name}" - - -def convert_episode(ep_dir: Path, split: str, out_dir: Path, overwrite: bool = False) -> dict: +def convert_episode(ep_dir: Path, split: str, out_dir: Path) -> dict: """Write one zarr; returns a report row. Idempotent via a done marker.""" from egomimic.rldb.zarr.zarr_writer import ZarrWriter - name = episode_hash(split, ep_dir) + name = f"abcsim_{split}_{ep_dir.name}" dst = out_dir / f"{name}.zarr" - if (dst / DONE_MARKER).exists() and not overwrite: + if (dst / DONE_MARKER).exists(): return {"episode": name, "status": "skipped"} meta = json.loads((ep_dir / "episode_metadata.json").read_text()) if not meta.get("prompt_timeline") and (meta.get("prompt_source") or {}).get("status") == "excluded": @@ -170,13 +166,7 @@ def convert_episode(ep_dir: Path, split: str, out_dir: Path, overwrite: bool = F task_description=spans[0][0] if spans else "", annotations=spans, intrinsics={"front_1": pinhole_intrinsics(h, w)}, - metadata_override={ - "lab": "abc_sim", - "split": split, - "abc_episode": ep_dir.name, - "prompt": spans[0][0] if spans else "", - "source": "abc sim_224", - }, + metadata_override={"lab": "abc_sim", "split": split}, ) (tmp / DONE_MARKER).write_text("ok\n") shutil.rmtree(dst, ignore_errors=True) @@ -207,15 +197,13 @@ def main(argv=None) -> int: p.add_argument("--src", type=Path, required=True, help="ABC cache dir ($ABC_CACHE)") p.add_argument("--task", required=True, help="dataset task_name, e.g. sim_pouring_beads") p.add_argument("--out", type=Path, required=True) - p.add_argument("--splits", nargs="+", default=["train", "val"]) p.add_argument("--workers", type=int, default=max(1, (os.cpu_count() or 2) // 2)) - p.add_argument("--overwrite", action="store_true") args = p.parse_args(argv) logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s") - episodes = find_episodes(args.src, args.task, args.splits) + episodes = find_episodes(args.src, args.task, ["train", "val"]) if not episodes: - log.error("no %s episodes under %s/{%s}_sim", args.task, args.src, ",".join(args.splits)) + log.error("no %s episodes under %s", args.task, args.src) return 1 args.out.mkdir(parents=True, exist_ok=True) log.info("converting %d episodes of %s -> %s", len(episodes), args.task, args.out) @@ -223,7 +211,7 @@ def main(argv=None) -> int: rows, failures = [], [] with ProcessPoolExecutor(max_workers=args.workers) as pool: futures = { - pool.submit(convert_episode, ep, split, args.out, args.overwrite): ep + pool.submit(convert_episode, ep, split, args.out): ep for ep, split in episodes } for i, fut in enumerate(as_completed(futures), 1): diff --git a/egomimic/scripts/abc_sim/policy_server.py b/egomimic/scripts/abc_sim/policy_server.py index d4c847fbd5..d36690b33d 100644 --- a/egomimic/scripts/abc_sim/policy_server.py +++ b/egomimic/scripts/abc_sim/policy_server.py @@ -9,8 +9,8 @@ POST /infer {"state": [14], "prompt": str, "images": {"top": {"shape": [H, W, 3], "b64": }, "top_hist": {...}}} # optional: the frame lag_frames back - -> {"actions": [[14] x 45], "dt": 1/30} -GET /health -> {"ok": true, "ckpt", "cameras", "lag_frames", "horizon"} + -> {"actions": [[14] x 45]} +GET /health -> {"ckpt", "cameras", "lag_frames"} The input is built exactly as ZarrDataset builds a training sample for the checkpoint's data config (keymap, transform list, fps) and normalized with @@ -130,8 +130,7 @@ def _send(self, code: int, payload: dict) -> None: def do_GET(self): if self.path != "/health": return self._send(404, {"error": self.path}) - self._send(200, {"ok": True, "ckpt": policy.ckpt_path, "cameras": policy.cameras, - "lag_frames": policy.lag_frames, "horizon": JOINT_RAW_HORIZON}) + self._send(200, {"ckpt": policy.ckpt_path, "cameras": policy.cameras, "lag_frames": policy.lag_frames}) def do_POST(self): if self.path != "/infer": @@ -140,7 +139,7 @@ def do_POST(self): req = json.loads(self.rfile.read(int(self.headers.get("Content-Length", 0)))) images = {k: decode_image(v) for k, v in req["images"].items()} actions = policy.infer(req["state"], images, req.get("prompt", "")) - self._send(200, {"actions": actions.tolist(), "dt": 1.0 / FPS}) + self._send(200, {"actions": actions.tolist()}) except Exception as e: # the client must see the reason, not a hang log.exception("request failed") self._send(500, {"error": f"{type(e).__name__}: {e}"}) @@ -154,14 +153,12 @@ def log_message(self, *a): def main(argv=None) -> int: p = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) p.add_argument("--ckpt", required=True) - p.add_argument("--host", default="127.0.0.1") p.add_argument("--port", type=int, default=8765) - p.add_argument("--device", default="cuda") a = p.parse_args(argv) logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s") - policy = JointPolicy(a.ckpt, a.device) - log.info("serving %s on http://%s:%d (cameras %s, lag %d frames)", a.ckpt, a.host, a.port, policy.cameras, policy.lag_frames) - HTTPServer((a.host, a.port), make_handler(policy)).serve_forever() + policy = JointPolicy(a.ckpt) + log.info("serving %s on port %d (cameras %s, lag %d frames)", a.ckpt, a.port, policy.cameras, policy.lag_frames) + HTTPServer(("127.0.0.1", a.port), make_handler(policy)).serve_forever() return 0 diff --git a/egomimic/scripts/abc_sim/sim_client.py b/egomimic/scripts/abc_sim/sim_client.py index 312f7115c5..fa61136de1 100644 --- a/egomimic/scripts/abc_sim/sim_client.py +++ b/egomimic/scripts/abc_sim/sim_client.py @@ -61,8 +61,8 @@ def step_one(self, action): class ServedPolicy: """abc_minimal's SimPolicy.infer surface over HTTP (noise and prefix unused).""" - def __init__(self, server: str, env: HistEnv | None = None): - self.server, self.env = server.rstrip("/"), env + def __init__(self, server: str): + self.server, self.env = server.rstrip("/"), None with urllib.request.urlopen(f"{self.server}/health", timeout=30) as r: self.health = json.load(r) @@ -95,8 +95,6 @@ def main(argv=None) -> int: p.add_argument("--out", type=Path, required=True) p.add_argument("--num-worlds", type=int, default=50) p.add_argument("--save-video", action="store_true") - p.add_argument("--camera-backend", default="mjwarp") - p.add_argument("--gpu-id", type=int, default=None) a = p.parse_args(argv) a.out.mkdir(parents=True, exist_ok=True) @@ -104,7 +102,7 @@ def main(argv=None) -> int: prompt = task_prompt(a.task) if a.prompt == "env" else (a.prompt or task_name_to_prompt(a.task)) cfg = SimEvalConfig(checkpoint=str(policy.health["ckpt"]), task=a.task, num_worlds=a.num_worlds, rtc=False, fast_inference=False, prefix_length=0, save_video=a.save_video, - camera_backend=a.camera_backend, gpu_id=a.gpu_id, output_dir=str(a.out), prompt=prompt) + output_dir=str(a.out), prompt=prompt) model_config = replace(cfg.model, chunk_length=HORIZON, action_dim=14, camera_keys=CAMERAS) env = HistEnv(task=a.task, height=cfg.camera_height, width=cfg.camera_width, camera_keys=CAMERAS, prompt=prompt, camera_backend=cfg.camera_backend, gpu_id=cfg.gpu_id, diff --git a/egomimic/scripts/abc_sim/val_vs_sim.py b/egomimic/scripts/abc_sim/val_vs_sim.py index 823ba6a3ae..62e25b2b6f 100644 --- a/egomimic/scripts/abc_sim/val_vs_sim.py +++ b/egomimic/scripts/abc_sim/val_vs_sim.py @@ -65,11 +65,10 @@ def main(argv=None) -> int: p.add_argument("manifest", type=Path) p.add_argument("--metric", required=True) p.add_argument("--out", type=Path, required=True) - p.add_argument("--no-wandb", action="store_true", help="manifest rows already carry 'val'") a = p.parse_args(argv) rows = json.loads(a.manifest.read_text()) for r in rows: - if "val" not in r or not a.no_wandb: + if "val" not in r: # a row may carry its val metric already r["val"] = wandb_metric(r["wandb"], a.metric, int(r["epoch"])) s = json.loads(Path(r["summary"]).read_text()) r["success_rate"], r["mean_max_progress"], r["num_worlds"] = s["success_rate"], s.get("mean_max_progress"), s["num_worlds"] diff --git a/egomimic/trainHydra.py b/egomimic/trainHydra.py index ff24b5352f..789a1a3792 100644 --- a/egomimic/trainHydra.py +++ b/egomimic/trainHydra.py @@ -168,7 +168,7 @@ def _weights_from_checkpoint(cfg: DictConfig) -> bool: return bool(ckpt_path) and os.path.isfile(ckpt_path) and not cfg.get("pretrained") -def _apply_init_weights(cfg: DictConfig, model) -> Optional[str]: +def _apply_init_weights(cfg: DictConfig, model) -> None: """``init_weights_ckpt``: weights-only fine-tune init (see ``init_weights_from_checkpoint``). Skipped when ``ckpt_path`` is set -- a resume (including a Slurm requeue, folded into ``ckpt_path`` by @@ -176,17 +176,14 @@ def _apply_init_weights(cfg: DictConfig, model) -> Optional[str]: must not be replaced by the base checkpoint's.""" init_ckpt = cfg.get("init_weights_ckpt") if not init_ckpt: - return None + return if cfg.get("ckpt_path"): log.info( f"init_weights_ckpt={init_ckpt} ignored: resuming from " f"ckpt_path={cfg.ckpt_path}" ) - return None - if not os.path.isfile(init_ckpt): - raise FileNotFoundError(f"init_weights_ckpt {init_ckpt} does not exist") + return init_weights_from_checkpoint(model, init_ckpt) - return init_ckpt def _log_dataset_frame_counts( diff --git a/tests/unit/test_init_weights_ckpt.py b/tests/unit/test_init_weights_ckpt.py index af8cadea06..a7b4e78188 100644 --- a/tests/unit/test_init_weights_ckpt.py +++ b/tests/unit/test_init_weights_ckpt.py @@ -58,8 +58,3 @@ def test_missing_init_checkpoint_raises(tmp_path): ) with pytest.raises(FileNotFoundError): train_hydra._apply_init_weights(cfg, _Net(14)) - - -def test_train_config_declares_the_knob(compose_resolve): - cfg = compose_resolve("train_zarr_cartesian", []) - assert cfg.init_weights_ckpt is None diff --git a/tests/unit/test_policy_server.py b/tests/unit/test_policy_server.py index 0aac8ae81f..8999a71889 100644 --- a/tests/unit/test_policy_server.py +++ b/tests/unit/test_policy_server.py @@ -110,11 +110,11 @@ def post(payload): try: with urllib.request.urlopen(f"{url}/health") as r: health = json.load(r) - assert health["cameras"] == ["top"] and health["lag_frames"] == 3 and health["horizon"] == 45 + assert health["cameras"] == ["top"] and health["lag_frames"] == 3 hwc = np.random.default_rng(0).integers(0, 255, (16, 24, 3), np.uint8) out = post({"state": list(range(14)), "prompt": "sim x", "images": {"top": {"shape": [16, 24, 3], "b64": base64.b64encode(hwc.tobytes()).decode()}}}) - assert np.asarray(out["actions"]).shape == (45, 14) and out["dt"] == 1 / 30 + assert np.asarray(out["actions"]).shape == (45, 14) assert torch.allclose(pol.algo.seen[CAM][0], torch.from_numpy(_chw(hwc))) try: post({"state": [0], "images": {}}) diff --git a/tests/unit/test_val_vs_sim.py b/tests/unit/test_val_vs_sim.py index f02d6ae231..7a25745d92 100644 --- a/tests/unit/test_val_vs_sim.py +++ b/tests/unit/test_val_vs_sim.py @@ -14,7 +14,7 @@ def test_report_from_manifest(tmp_path): rows.append({"name": f"r{i}", "epoch": ep, "hours": h, "val": val, "summary": str(s)}) m = tmp_path / "m.json" m.write_text(json.dumps(rows)) - assert vs.main([str(m), "--metric", "x", "--out", str(tmp_path / "out"), "--no-wandb"]) == 0 + assert vs.main([str(m), "--metric", "x", "--out", str(tmp_path / "out")]) == 0 rep = json.loads((tmp_path / "out" / "report.json").read_text()) assert rep["correlations"]["spearman_success_rate"]["r"] == -1.0 # lower val mse, higher success assert rep["loglinear_val_vs_hours"]["r2"] > 0.95 and rep["loglinear_val_vs_hours"]["slope"] < 0 From 57eeaf8f648f80dd07798e9f1de27e528433a1e9 Mon Sep 17 00:00:00 2001 From: Aidan Gao Date: Tue, 29 Sep 2026 19:54:35 -0400 Subject: [PATCH 15/20] abc_sim: rollout_lambda.sbatch takes POLICY=hold without CKPT, absolute OUT Found running the sim-validation hold floor on Lambda: under set -u the srun line expanded the unset CKPT and the step died, and a relative OUT resolved inside $ABC because the client cd's there before writing. Co-Authored-By: Claude Opus 5.5 --- egomimic/scripts/abc_sim/rollout_lambda.sbatch | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/egomimic/scripts/abc_sim/rollout_lambda.sbatch b/egomimic/scripts/abc_sim/rollout_lambda.sbatch index 65e04cc7d4..7ca8f6495c 100644 --- a/egomimic/scripts/abc_sim/rollout_lambda.sbatch +++ b/egomimic/scripts/abc_sim/rollout_lambda.sbatch @@ -26,8 +26,9 @@ TASK=${TASK:-sim_put_the_plastic_bottles_in_the_bin} : "${OUT:?set OUT}" POLICY=${POLICY:-served} [ "$POLICY" = hold ] || : "${CKPT:?set CKPT}" +CKPT=${CKPT:-} # hold needs none; set -u would trip on it below PORT=${PORT:-$(( 8700 + SLURM_JOB_ID % 200 ))} # set PORT per step when several share a job -mkdir -p "$OUT" +mkdir -p "$OUT" && OUT=$(cd "$OUT" && pwd) # absolute: the client runs from $ABC echo "=== abc_rollout job=$SLURM_JOB_ID node=$(hostname) task=$TASK ckpt=${CKPT:-none} policy=$POLICY out=$OUT" srun --exact --overlap --ntasks=1 --gres=gpu:h100:1 --cpus-per-task=${CPUS_PER_TASK:-16} --mem=120G --container-image="$IMAGE" --container-mounts=/workspace:/workspace:rw --container-workdir="$SRC" \ bash -c "set -euo pipefail From 2a512b6613afda3504ff28fca876ca60a94c2e22 Mon Sep 17 00:00:00 2001 From: Aidan Gao Date: Tue, 29 Sep 2026 21:43:18 -0400 Subject: [PATCH 16/20] abc_sim: rollout_lambda.sbatch stops its policy server when the step ends On the Lambda loaner a finished step does not reap processes it backgrounded (no proctrack/cgroup), so every rollout left its policy_server running with ~8 GB on the GPU and the port taken. Kill the shell's jobs on EXIT. Co-Authored-By: Claude Opus 5.5 --- egomimic/scripts/abc_sim/rollout_lambda.sbatch | 1 + 1 file changed, 1 insertion(+) diff --git a/egomimic/scripts/abc_sim/rollout_lambda.sbatch b/egomimic/scripts/abc_sim/rollout_lambda.sbatch index 7ca8f6495c..eb2783caa8 100644 --- a/egomimic/scripts/abc_sim/rollout_lambda.sbatch +++ b/egomimic/scripts/abc_sim/rollout_lambda.sbatch @@ -34,6 +34,7 @@ srun --exact --overlap --ntasks=1 --gres=gpu:h100:1 --cpus-per-task=${CPUS_PER_T bash -c "set -euo pipefail export HOME=/workspace/users/$USER HF_HOME=/workspace/users/$USER/.cache/huggingface HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 export PYTHONPATH=$SRC PYTHONUNBUFFERED=1 OMP_NUM_THREADS=4 + trap 'kill \$(jobs -p) 2>/dev/null' EXIT # the step's cleanup does not reap the backgrounded server on the loaner if [ '$POLICY' = served ]; then $SRC/.venv/bin/python -m egomimic.scripts.abc_sim.policy_server --ckpt '$CKPT' --port $PORT > '$OUT/server.log' 2>&1 & for i in \$(seq 1 120); do curl -fsS http://127.0.0.1:$PORT/health > /dev/null 2>&1 && break; sleep 5; done From 9588f1d120e3034203021aa609bba59d9e8efbfa Mon Sep 17 00:00:00 2001 From: Aidan Gao Date: Thu, 1 Oct 2026 22:40:53 -0400 Subject: [PATCH 17/20] abc: real-bottles RDT recipe and YAM robot labels read as Eva - train_zarr_abc_real_rdt: the base RDT (width 1024, the architecture of Ryan's rdt_mem_all checkpoints) on the 100 real ABC bottles episodes, 4 GPUs, 100k steps. The prompt is fixed in the model (default_prompt, no annotation track) and there is no image history, so a client that sends the current frames and the 14-D joint state can drive it: the abc_sim client, or an ABC-replica station. - model/rdt_abc_joints_dinov3_vitb: the ABC joints head on the base RDT; rdt_1b_abc_joints now inherits it at width 2048 (resolved configs unchanged). - embodiment: YAM_* collapses to EVA_*. The Lambda (R2) copies of the real ABC zarrs say yam_bimanual since the 2026-09-01 SQL relabel. This reverses test_pi05_norm_rot6d's "robot names are never aliased" assertion. - test: the real recipe trains two steps and its checkpoint serves without a history frame, using the model's own prompt. Co-Authored-By: Claude Opus 5.5 --- .../model/rdt_1b_abc_joints_dinov3_vitb.yaml | 45 +++---------------- .../model/rdt_abc_joints_dinov3_vitb.yaml | 44 ++++++++++++++++++ .../train_zarr_abc_real_rdt.yaml | 38 ++++++++++++++++ egomimic/rldb/embodiment/embodiment.py | 9 +++- tests/unit/test_abc_joint_configs.py | 44 ++++++++++++++++++ tests/unit/test_pi05_norm_rot6d.py | 4 +- 6 files changed, 141 insertions(+), 43 deletions(-) create mode 100644 egomimic/hydra_configs/model/rdt_abc_joints_dinov3_vitb.yaml create mode 100644 egomimic/hydra_configs/train_zarr_abc_real_rdt.yaml diff --git a/egomimic/hydra_configs/model/rdt_1b_abc_joints_dinov3_vitb.yaml b/egomimic/hydra_configs/model/rdt_1b_abc_joints_dinov3_vitb.yaml index eb789c0483..3b14436e2e 100644 --- a/egomimic/hydra_configs/model/rdt_1b_abc_joints_dinov3_vitb.yaml +++ b/egomimic/hydra_configs/model/rdt_1b_abc_joints_dinov3_vitb.yaml @@ -1,43 +1,10 @@ -# RDT_1B_ABC_JOINTS_DINOV3_VITB: RDT-1B with a second head, eva_bimanual = the -# ABC YAM station in JOINT space (14-D state, (100, 14) actions_joints, abc_sim's -# action space), beside the mecka human head; trunk, front tower and text stem -# shared. Front camera only, at the human tower size [224, 400]. +# RDT_1B_ABC_JOINTS_DINOV3_VITB: rdt_abc_joints_dinov3_vitb at RDT-1B width +# (rdt_1b_bc_mecka_6d_dinov3_vitb's width and heads). defaults: - - rdt_1b_bc_mecka_6d_dinov3_vitb + - rdt_abc_joints_dinov3_vitb - _self_ robomimic_model: - domains: ["human_bimanual", "eva_bimanual"] - dims: - eva_bimanual: {proprio: 14, action: 14} - ac_keys: - eva_bimanual: "actions_joints" - - head_specs: - eva_bimanual: - _target_: egomimic.models.fm_policy.FMPolicy - action_horizon: 100 - num_inference_steps: 10 - pooling: null - time_dist: "beta" - infer_ac_dims: - eva_bimanual: ${model.robomimic_model.dims.eva_bimanual.action} - model: - _target_: egomimic.models.rdt_nets.RDTDenoiser - act_dim: ${model.robomimic_model.dims.eva_bimanual.action} - act_seq: 100 - hidden_dim: ${model.robomimic_model.width} - n_state_tokens: 1 - time_scale: 1000.0 - - stem_specs: - eva_bimanual: - state_joint_positions: - _target_: egomimic.models.hpt_nets.MLPPolicyStem - input_dim: ${model.robomimic_model.dims.eva_bimanual.proprio} - output_dim: ${model.robomimic_model.width} - widths: - - ${model.robomimic_model.width} - - ${model.robomimic_model.width} - history_len: 1 - history_dropout: 0.2 + width: 2048 + trunk: + num_heads: 32 diff --git a/egomimic/hydra_configs/model/rdt_abc_joints_dinov3_vitb.yaml b/egomimic/hydra_configs/model/rdt_abc_joints_dinov3_vitb.yaml new file mode 100644 index 0000000000..49fe4cd2f1 --- /dev/null +++ b/egomimic/hydra_configs/model/rdt_abc_joints_dinov3_vitb.yaml @@ -0,0 +1,44 @@ +# RDT_ABC_JOINTS_DINOV3_VITB: the base RDT (width 1024, depth 28: the architecture of +# Ryan's all-mecka rdt_mem_all checkpoints) with a second head, eva_bimanual = the ABC +# YAM station in JOINT space (14-D state, (100, 14) actions_joints, abc_sim's action +# space), beside the mecka human head; trunk, front tower and text stem shared. Front +# camera only, at the human tower size [224, 400]. rdt_1b_abc_joints is this at width 2048. +defaults: + - rdt_bc_mecka_6d_dinov3_vitb + - _self_ + +robomimic_model: + domains: ["human_bimanual", "eva_bimanual"] + dims: + eva_bimanual: {proprio: 14, action: 14} + ac_keys: + eva_bimanual: "actions_joints" + + head_specs: + eva_bimanual: + _target_: egomimic.models.fm_policy.FMPolicy + action_horizon: 100 + num_inference_steps: 10 + pooling: null + time_dist: "beta" + infer_ac_dims: + eva_bimanual: ${model.robomimic_model.dims.eva_bimanual.action} + model: + _target_: egomimic.models.rdt_nets.RDTDenoiser + act_dim: ${model.robomimic_model.dims.eva_bimanual.action} + act_seq: 100 + hidden_dim: ${model.robomimic_model.width} + n_state_tokens: 1 + time_scale: 1000.0 + + stem_specs: + eva_bimanual: + state_joint_positions: + _target_: egomimic.models.hpt_nets.MLPPolicyStem + input_dim: ${model.robomimic_model.dims.eva_bimanual.proprio} + output_dim: ${model.robomimic_model.width} + widths: + - ${model.robomimic_model.width} + - ${model.robomimic_model.width} + history_len: 1 + history_dropout: 0.2 diff --git a/egomimic/hydra_configs/train_zarr_abc_real_rdt.yaml b/egomimic/hydra_configs/train_zarr_abc_real_rdt.yaml new file mode 100644 index 0000000000..c1c3b3df86 --- /dev/null +++ b/egomimic/hydra_configs/train_zarr_abc_real_rdt.yaml @@ -0,0 +1,38 @@ +# ABC REAL bottles: the base RDT (width 1024) trained on the real ABC episodes only +# (data=abc_real_joints, 90/10 episode split) and evaluated in abc_sim. Built to run +# unchanged on an ABC-replica YAM station: the prompt is fixed in the model +# (default_prompt; no annotation track is read) and there is no image history, so a +# client that sends the current camera frames + the 14-D joint state is enough +# (scripts/abc_sim/sim_client.py, ABC's deploy client through policy_server). +# +# from scratch: python egomimic/trainHydra.py --config-name=train_zarr_abc_real_rdt +# from Ryan's all-mecka base: ... init_weights_ckpt= +defaults: + - train_zarr_abc_real_rdt_ft + - override model: rdt_abc_joints_dinov3_vitb + - _self_ + +model: + robomimic_model: + default_prompt: "put the plastic bottles in the bin" + trunk: + image_history: 1 + +data: + train_datasets: + eva_bimanual: + resolver: + key_map: + annotation_key: null + image_history_gap_s: null + +# 1000 x 100 steps at 32 per GPU x 4 GPUs; a checkpoint every 50 epochs (5k steps) +trainer: + max_epochs: 1000 + min_epochs: 1000 + +launch_params: + gpus_per_node: 4 + +norm_stats: + max_samples: 200000 diff --git a/egomimic/rldb/embodiment/embodiment.py b/egomimic/rldb/embodiment/embodiment.py index 7a5c0602be..581f74a34e 100644 --- a/egomimic/rldb/embodiment/embodiment.py +++ b/egomimic/rldb/embodiment/embodiment.py @@ -63,17 +63,22 @@ def get_embodiment(index): # embodiment metadata (e.g. MECKA_BIMANUAL, SCALE_LEFT_ARM). Locally all human # demonstration data is ONE embodiment (see the EMBODIMENT docstring; the # source lives only in the SQL `lab` field), so those names collapse to -# HUMAN_*. Robot names (EVA_*) are never aliased. +# HUMAN_*. The one robot is Eva: YAM_* (the ABC station's label since the +# 2026-09-01 SQL relabel, also in the zarrs re-exported to R2) collapses to EVA_*. HUMAN_VENDOR_PREFIXES = ("MECKA", "SCALE", "ARIA", "LIGHTWHEEL") +ROBOT_VENDOR_PREFIXES = ("YAM",) def canonical_embodiment_name(embodiment_name: str) -> str: """Upper-case EMBODIMENT member name, with legacy vendor prefixes - (``MECKA_BIMANUAL`` ...) collapsed onto ``HUMAN_*``. Does not validate.""" + (``MECKA_BIMANUAL`` ...) collapsed onto ``HUMAN_*`` and ``YAM_*`` onto + ``EVA_*``. Does not validate.""" name = embodiment_name.upper() vendor, _, suffix = name.partition("_") if vendor in HUMAN_VENDOR_PREFIXES and suffix: return f"HUMAN_{suffix}" + if vendor in ROBOT_VENDOR_PREFIXES and suffix: + return f"EVA_{suffix}" return name diff --git a/tests/unit/test_abc_joint_configs.py b/tests/unit/test_abc_joint_configs.py index 3f97d4eb24..3ea2ea53e8 100644 --- a/tests/unit/test_abc_joint_configs.py +++ b/tests/unit/test_abc_joint_configs.py @@ -39,6 +39,50 @@ def _tiny_rdt(emb: str) -> list[str]: ] +def test_yam_robot_label_is_eva(): + from egomimic.rldb.embodiment.embodiment import get_embodiment_id + + assert get_embodiment_id("yam_bimanual") == get_embodiment_id("eva_bimanual") + + +def test_abc_real_recipe_serves_with_a_fixed_prompt_and_no_history(tmp_path, monkeypatch): + hermetic_env(monkeypatch) + data, out, hashes = write_fixtures(tmp_path, "eva") + recipe = Recipe("train_zarr_abc_real_rdt", "abc_real_joints", "rdt_abc_joints_dinov3_vitb", EMB) + tiny = [o for o in _tiny_rdt(EMB) if not o.startswith(f"{RM}.annotation_key")] + cfg = compose_recipe( + recipe, + common_overrides(EMB, data, out, batch_size=2, num_workers=0, episode_hashes=hashes) + + cpu_trainer_overrides(2) + + tiny + + [ + # abc_real_joints resolves through the SQL table; read the fixtures locally + f"data.train_datasets.{EMB}.resolver._target_={LOCAL_RESOLVER}", + f"data.train_datasets.{EMB}.resolver.folder_path={data}", + f"data.train_datasets.{EMB}.mode=total", + f"data.valid_datasets.{EMB}.filters=null", + f"data.valid_datasets.{EMB}.mode=total", + ], + out, + ) + assert cfg.model.robomimic_model.width == 32 and cfg.launch_params.gpus_per_node == 4 + spy = BatchKeySpy(monkeypatch, HPT) + metrics, objects = train_hydra.train(cfg) + assert math.isfinite(float(metrics["Train/action_loss"])) + seen = spy.keys[EMB] + assert {JOINT_ACTION_KEY, JOINT_STATE_KEY, Eva.VIZ_IMAGE_KEY, "fps"} <= seen + assert f"{Eva.VIZ_IMAGE_KEY}_hist" not in seen and "annotations" not in seen + + ckpt = out / "served.ckpt" + objects["trainer"].save_checkpoint(ckpt) + policy = JointPolicy(str(ckpt), device="cpu") + assert policy.cameras == ["top"] and policy.lag_frames == 0 and policy.annotation_key is None + assert policy.algo.default_prompt == "put the plastic bottles in the bin" + img = np.random.default_rng(0).random((3, 64, 64), dtype=np.float32) + actions = policy.infer(np.zeros(14, np.float32), {"top": img}, "any prompt; the model's own is used") + assert actions.shape == (45, 14) and np.isfinite(actions).all() + + def test_abc_sim_finetune_train_step(tmp_path, monkeypatch): hermetic_env(monkeypatch) data, out, hashes = write_fixtures(tmp_path, "eva") diff --git a/tests/unit/test_pi05_norm_rot6d.py b/tests/unit/test_pi05_norm_rot6d.py index 0b75298f2b..d2969c48fc 100644 --- a/tests/unit/test_pi05_norm_rot6d.py +++ b/tests/unit/test_pi05_norm_rot6d.py @@ -454,8 +454,8 @@ def test_vendor_embodiment_names_collapse_to_human(): assert get_embodiment_id("human_bimanual") == EMBODIMENT.HUMAN_BIMANUAL.value assert get_embodiment_id("eva_bimanual") == EMBODIMENT.EVA_BIMANUAL.value assert not is_legacy_vendor_embodiment("human_bimanual") - with pytest.raises(KeyError): - get_embodiment_id("yam_bimanual") # robot names are never aliased + # the one robot is Eva: YAM_* (the ABC station after the 2026-09-01 relabel) is Eva + assert get_embodiment_id("yam_bimanual") == EMBODIMENT.EVA_BIMANUAL.value def test_base_converter_rejects_norm_6d_encoding(): From 7c9330a33c1b287623706bdd189897af79d10960 Mon Sep 17 00:00:00 2001 From: Aidan Gao Date: Thu, 1 Oct 2026 22:45:05 -0400 Subject: [PATCH 18/20] lambda_pack: request 225G per GPU The loaner partition's DefMemPerGPU=256000 MB times 8 GPUs exceeds a node's 2,000,000 MB, so a full-node request is rejected ("Requested node configuration is not available"); Ryan's jobs request 225G per GPU. Co-Authored-By: Claude Opus 5.5 --- egomimic/scripts/abc_sim/lambda_pack.sbatch | 1 + 1 file changed, 1 insertion(+) diff --git a/egomimic/scripts/abc_sim/lambda_pack.sbatch b/egomimic/scripts/abc_sim/lambda_pack.sbatch index dbf47721cf..0403c33aa1 100644 --- a/egomimic/scripts/abc_sim/lambda_pack.sbatch +++ b/egomimic/scripts/abc_sim/lambda_pack.sbatch @@ -20,6 +20,7 @@ #SBATCH --partition=loaner #SBATCH --nodes=1 #SBATCH --gres=gpu:h100:8 +#SBATCH --mem-per-gpu=225G # the partition's DefMemPerGPU (256000 MB) x 8 exceeds a node #SBATCH --cpus-per-task=224 #SBATCH --exclusive #SBATCH --time=2-00:00:00 From 9a5496e11441f4f17245c8505cca0ec92a5f410a Mon Sep 17 00:00:00 2001 From: Aidan Gao Date: Sat, 3 Oct 2026 02:52:32 -0400 Subject: [PATCH 19/20] abc: dish-rack human + sim cotrain recipe abc_sim load_plates_into_dish_rack (robot, joints) + the Mecka episodes whose description has dishes and a rack (human, wrist-frame 6D; 232 episodes, 5.4 h), one batch of each per step, on the base RDT from scratch with the real-bottles recipe's deployable settings (fixed robot prompt, no image history), one full node. The dish rack is the one task in abc_sim, the real ABC data and Mecka. Co-Authored-By: Claude Opus 5.5 --- ...ain_zarr_abc_sim_dishrack_rdt_cotrain.yaml | 40 +++++++++++++++++++ 1 file changed, 40 insertions(+) create mode 100644 egomimic/hydra_configs/train_zarr_abc_sim_dishrack_rdt_cotrain.yaml diff --git a/egomimic/hydra_configs/train_zarr_abc_sim_dishrack_rdt_cotrain.yaml b/egomimic/hydra_configs/train_zarr_abc_sim_dishrack_rdt_cotrain.yaml new file mode 100644 index 0000000000..9d8aa50b6f --- /dev/null +++ b/egomimic/hydra_configs/train_zarr_abc_sim_dishrack_rdt_cotrain.yaml @@ -0,0 +1,40 @@ +# Dish-rack human + sim cotrain: abc_sim load_plates_into_dish_rack (robot, joint space, +# converted to $ABC_SIM_ZARR_ROOT/dishrack) + the Mecka episodes whose description has +# dishes and a rack (human, wrist-frame 6D; 232 episodes / 5.4 h on 2026-10-03), one batch +# of each per step. Base RDT (width 1024) from scratch with the real-bottles recipe's +# deployable settings: the robot prompt is fixed in the model, no image history. +# +# python egomimic/trainHydra.py --config-name=train_zarr_abc_sim_dishrack_rdt_cotrain +defaults: + - train_zarr_abc_real_rdt + - override data: abc_sim_joints_cotrain_mecka + - _self_ + +model: + robomimic_model: + domains: ["human_bimanual", "eva_bimanual"] + default_prompt: "load the plates into the dish rack" + +data: + train_datasets: + eva_bimanual: + resolver: + folder_path: ${oc.env:ABC_SIM_ZARR_ROOT,/workspace/users/agao81/datasets/abc_sim/zarr}/dishrack + human_bimanual: + resolver: + key_map: + image_history_gap_s: null # trunk.image_history is 1 + filters: + filter_lambdas: + - "lambda row: str(row.get('lab', '')).lower() == 'mecka'" + # the SQL selection of 2026-10-03: a rack, and dishes, in the episode description + - "lambda row: bool(__import__('re').search(r'\\b(?:dish|drying|plate|draining)[ -]?rack|\\brack\\b', str(row.get('task_description') or ''), 2))" + - "lambda row: bool(__import__('re').search(r'dish|plate|bowl|cup|mug', str(row.get('task_description') or ''), 2))" + +# 500 x 100 steps on one full node: 32 robot + 32 human samples per GPU per step +trainer: + max_epochs: 500 + min_epochs: 500 + +launch_params: + gpus_per_node: 8 From 2248a75d3388ffbdfedef20ac46db9b51c10e698 Mon Sep 17 00:00:00 2001 From: Aidan Gao Date: Sat, 3 Oct 2026 17:32:09 -0400 Subject: [PATCH 20/20] abc: dish-rack sim-only recipe, the cotrain's matched control abc_sim load_plates_into_dish_rack alone on the base RDT from scratch, with the same deployable settings and dish-rack prompt as train_zarr_abc_sim_dishrack_rdt_cotrain. Co-Authored-By: Claude Opus 5.5 --- .../train_zarr_abc_sim_dishrack_rdt.yaml | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) create mode 100644 egomimic/hydra_configs/train_zarr_abc_sim_dishrack_rdt.yaml diff --git a/egomimic/hydra_configs/train_zarr_abc_sim_dishrack_rdt.yaml b/egomimic/hydra_configs/train_zarr_abc_sim_dishrack_rdt.yaml new file mode 100644 index 0000000000..15e7263adb --- /dev/null +++ b/egomimic/hydra_configs/train_zarr_abc_sim_dishrack_rdt.yaml @@ -0,0 +1,20 @@ +# Dish-rack sim only: abc_sim load_plates_into_dish_rack (robot, joint space, converted to +# $ABC_SIM_ZARR_ROOT/dishrack), base RDT (width 1024) from scratch with the real-bottles +# recipe's deployable settings (robot prompt fixed in the model, no image history). The +# matched control for train_zarr_abc_sim_dishrack_rdt_cotrain. +# +# python egomimic/trainHydra.py --config-name=train_zarr_abc_sim_dishrack_rdt +defaults: + - train_zarr_abc_real_rdt + - override data: abc_sim_joints + - _self_ + +model: + robomimic_model: + default_prompt: "load the plates into the dish rack" + +data: + train_datasets: + eva_bimanual: + resolver: + folder_path: ${oc.env:ABC_SIM_ZARR_ROOT,/workspace/users/agao81/datasets/abc_sim/zarr}/dishrack