diff --git a/egomimic/algo/rdt.py b/egomimic/algo/rdt.py index 74043f0deb..9615ec3c47 100644 --- a/egomimic/algo/rdt.py +++ b/egomimic/algo/rdt.py @@ -10,6 +10,7 @@ import torch import torch.nn as nn +import torch.nn.functional as F from egomimic.algo.hpt import HPT, HPTModel from egomimic.models.rdt_nets import ( @@ -48,6 +49,10 @@ class RDTModel(HPTModel): steps are masked out. ``short_dropout`` replaces a sample's whole short image stream by a learned null token (training only), so the policy cannot lean on the latest frames alone. + + ``lang_pad_to`` right-pads the language tokens (masked) to a multiple of it, + so torch.compile sees a fixed length instead of each batch's longest prompt; + masked keys get zero attention weight, so outputs do not change. """ def __init__( @@ -61,6 +66,7 @@ def __init__( image_history_dropout: float = 0.0, memory: dict | None = None, short_dropout: float = 0.0, + lang_pad_to: int | None = None, **kwargs, ): super().__init__( @@ -74,6 +80,7 @@ def __init__( self.image_history_dropout = float(image_history_dropout) self.stem_modality = {} self.short_dropout = float(short_dropout) + self.lang_pad_to = lang_pad_to self.long_memory = ( None if memory is None else RDTMemory(hidden_dim=embed_dim, **memory) ) @@ -137,7 +144,7 @@ def _image_tokens(self, name, modality, stem, x): raise ValueError( f"'{modality}' carries {T} frames but trunk.image_history=" f"{self.image_history}; set it to match the data config's " - "image_history_gap_s (2 frames when set, else 1)" + "key_map.image_history (when image_history_gap_s is set, else 1)" ) drop = self._drop(self.image_history_dropout, B, x.device) if T > 1 else None if drop is not None: @@ -222,6 +229,10 @@ def forward_features(self, domain, data): if lang: cond.lang = torch.cat(lang, dim=1) cond.lang_mask = torch.cat(lang_mask, dim=1) + if self.lang_pad_to: + pad = -cond.lang.shape[1] % self.lang_pad_to + cond.lang = F.pad(cond.lang, (0, 0, 0, pad)) + cond.lang_mask = F.pad(cond.lang_mask, (0, pad), value=False) if self.long_memory is not None: # after state dropout, so a fused memory loses its proprio with it cond.memory, cond.memory_mask = self.long_memory( @@ -251,7 +262,7 @@ def _robomimic_to_hpt_data( domain="", ): """HPT's layout plus ``fps``; per camera with a ``*_hist`` twin, the - (past, current) frame pair on the image input's time axis; and a + (past..., current) frames on the image input's time axis; and a ``*_mem`` window as ``memory`` / ``memory_mask``, un-augmented (its tower is frozen and the stem normalizes).""" past = { @@ -272,13 +283,16 @@ def _robomimic_to_hpt_data( ) for key, past_key in past.items(): short = key.rsplit(".", 1)[-1] - # (B * 2, 3, H, W), each sample's (past, current) adjacent: one - # jitter draw per sample covers both frames, and every op -- a - # contrast mean, a crop, a resize -- still sees one frame. - pair = torch.stack([batch[past_key], batch[key]], dim=1) - n = pair.shape[0] - pair = self._apply_image_augs(pair.flatten(0, 1), short, frames=2) - data[short] = pair.reshape(n, 2, *pair.shape[1:]).unsqueeze(2) + # (B * T, 3, H, W), each sample's frames adjacent: one jitter draw + # per sample covers all of them, and every op -- a contrast mean, + # a crop, a resize -- still sees one frame. + past_frames = batch[past_key] + if past_frames.dim() == 4: + past_frames = past_frames.unsqueeze(1) + clip = torch.cat([past_frames, batch[key].unsqueeze(1)], dim=1) + n, t = clip.shape[:2] + clip = self._apply_image_augs(clip.flatten(0, 1), short, frames=t) + data[short] = clip.reshape(n, t, *clip.shape[1:]).unsqueeze(2) for key in memory: data["memory"] = batch[key] data["memory_mask"] = batch[f"{key}{HISTORY_MASK_SUFFIX}"] > 0.5 diff --git a/egomimic/eval/eval_video.py b/egomimic/eval/eval_video.py index d524c0203c..c732656c25 100644 --- a/egomimic/eval/eval_video.py +++ b/egomimic/eval/eval_video.py @@ -35,7 +35,7 @@ def __init__( viz_every_n_epochs: int = 1, viz_max_batches: int | None = None, viz_mode: str = "replay", - replay_chunks: int = 5, + replay_chunks: int = 17, replay_trail: int = 5, action_stride: float | None = None, ): @@ -55,7 +55,8 @@ def __init__( # "replay" (egomimic/eval/replay_viz.py) plays each chunk as GT then # prediction; it needs consecutive frames, so it runs only on a pinned # video loader and caps it at ``replay_chunks`` chunks instead of - # ``viz_max_batches``. "overlay" draws each frame's whole GT and + # ``viz_max_batches`` (17: ~17 s of a 30-frame-chunk episode, 1020 video + # frames). "overlay" draws each frame's whole GT and # predicted chunk, and is what a head without a video loader gets. # ``action_stride`` is frames per chunk step, fractional when the chunk # is resampled (None: inferred from the GT keypoints, see diff --git a/egomimic/hydra_configs/data/mecka_all_kp_pinned_val.yaml b/egomimic/hydra_configs/data/mecka_all_kp_pinned_val.yaml new file mode 100644 index 0000000000..723ffe8573 --- /dev/null +++ b/egomimic/hydra_configs/data/mecka_all_kp_pinned_val.yaml @@ -0,0 +1,45 @@ +# All mecka episodes with the wrist-frame KEYPOINT pipeline, and mecka_20pct_6d's +# validation heads and video pins so runs are comparable to the 20 % ones: +# actions_keypoints (100, 144) per hand [wrist xyz+rot6d (9) | +# 21 keypoints (63)], wrist frame +# observations.state.keypoints (144,) per hand [wrist pose in HEAD +# frame | keypoints in wrist frame] +# train = every lab=mecka episode except the bad-egomotion blocklist, the 3 valid +# pins and all 10 episodes of the 3 unseen_op_valid operators (38,596 episodes, +# ~77M frames vs 4,078 / 10.5M for the 20 % subset, SQL 2026-09-23). +defaults: + - mecka_20pct_6d + - _self_ + +# SQL app.episodes.operator of the 3 unseen_op_valid pins. +held_out_operators: + - 69475d885c0479a0fb81408e + - 69a84eed47e12c9b339208fd + - 6974f58bb393af9134c82247 + +train_datasets: + human_bimanual: + resolver: + key_map: + keymap_mode: keypoints + transform_list: + mode: keypoints_wristframe_6d + pad_proprio_gripper: false + filters: + episode_hashes: null + filter_lambdas: + - "lambda row: str(row.get('lab', '')).lower() == 'mecka'" + - "lambda row: str(row.get('operator', '')) not in ${data.held_out_operators}" + - "lambda row: str(row.get('episode_hash', '')) not in ${data.valid_datasets.human_bimanual.filters.episode_hashes}" + +# The 20 % heads interpolate the train lambdas, which now exclude their own pins. +valid_datasets: + human_bimanual: + filters: + filter_lambdas: + - "lambda row: str(row.get('lab', '')).lower() == 'mecka'" + +unseen_op_valid_datasets: + human_bimanual: + filters: + filter_lambdas: ${data.valid_datasets.human_bimanual.filters.filter_lambdas} diff --git a/egomimic/hydra_configs/data/mecka_flagship_fold_cup_kp_holdout.yaml b/egomimic/hydra_configs/data/mecka_flagship_fold_cup_kp_holdout.yaml new file mode 100644 index 0000000000..6c6ff1c788 --- /dev/null +++ b/egomimic/hydra_configs/data/mecka_flagship_fold_cup_kp_holdout.yaml @@ -0,0 +1,21 @@ +# mecka_flagship_fold_kp_holdout plus the flagship cup_on_saucer episodes (5409 +# eps / 905k frames): 7016 eps / 1.22M frames. 3 held out, 1 fold (shared with +# the parent's valid) and 2 cup, each from a different operator. +defaults: + - mecka_flagship_fold_kp_holdout + - _self_ + +tasks: [fold_clothes, cup_on_saucer] + +held_out_hashes: + - 6930852199946db9586cf530 + - 692ed3f8f2e4e72c033b9970 + - 692eb00f1e78e176ba6cf9b2 + +metric_frames_per_episode: + train_viz: 1 + +video_episodes: + valid: + - 6930852199946db9586cf530 + - 692ed3f8f2e4e72c033b9970 diff --git a/egomimic/hydra_configs/data/mecka_flagship_fold_kp_holdout.yaml b/egomimic/hydra_configs/data/mecka_flagship_fold_kp_holdout.yaml new file mode 100644 index 0000000000..52db47184e --- /dev/null +++ b/egomimic/hydra_configs/data/mecka_flagship_fold_kp_holdout.yaml @@ -0,0 +1,88 @@ +# Mecka FLAGSHIP fold_clothes (SQL lab='mecka', task='fold_clothes', 1607 eps / +# 316k frames under processed_v3/mecka/flagship/) with the wrist-frame KEYPOINT +# pipeline of mecka_all_kp_pinned_val, and 3 hand-picked held-out episodes as the +# only validation: a fitting test, not an operator-generalization split. +# train = every other episode (mode total), minus the bad-egomotion blocklist +# valid = the `held_out_hashes` (3 operators, ~200 frames each, run in full) +# Children swap `tasks` / `held_out_hashes` (mecka_flagship_fold_cup_kp_holdout, +# mecka_flagship_fold_kp_overfit1). Norm stats are split-specific. +_target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper + +tasks: [fold_clothes] + +# 2026-09-25, fold_clothes rows with 40-60th percentile length, one per operator. +held_out_hashes: + - 6930852199946db9586cf530 + - 692e9f2cba0400c70689fdfd + - 692eefe40ba268eb16564ab9 + +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: keypoints + annotation_key: annotations + transform_list: + _target_: egomimic.rldb.embodiment.human.Human.get_transform_list + mode: keypoints_wristframe_6d + stride: 1 + fix_left_wrist_convention: true + pad_proprio_gripper: false + filters: + _target_: egomimic.rldb.filters.DatasetFilter + filter_lambdas: + - "lambda row: str(row.get('lab', '')).lower() == 'mecka'" + - "lambda row: str(row.get('task', '')).lower() in ${data.tasks}" + - "lambda row: '/mecka/flagship/' in str(row.get('zarr_processed_path', ''))" + - "lambda row: str(row.get('episode_hash', '')) not in ${data.held_out_hashes}" + exclude_hashes: mecka_bad_egomotion + mode: total + +valid_datasets: + human_bimanual: + _target_: ${data.train_datasets.human_bimanual._target_} + resolver: ${data.train_datasets.human_bimanual.resolver} + filters: + _target_: egomimic.rldb.filters.DatasetFilter + filter_lambdas: + - "lambda row: str(row.get('lab', '')).lower() == 'mecka'" + episode_hashes: ${data.held_out_hashes} + mode: total + +train_viz_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: total + +train_dataloader_params: + human_bimanual: + batch_size: 32 + num_workers: 12 + +valid_dataloader_params: + human_bimanual: + batch_size: 32 + num_workers: 8 + shuffle: false + +train_viz_dataloader_params: + human_bimanual: + batch_size: 32 + num_workers: 4 + shuffle: false + +# Single-rank K = floor(5120 / episodes); the 3 valid episodes run in full. +metric_frames_per_episode: + train_viz: 3 + +video_episodes: + valid: + - 6930852199946db9586cf530 + train_viz: + - 692fe87a3c1dfb13a7a99ee6 diff --git a/egomimic/hydra_configs/data/mecka_flagship_fold_kp_overfit1.yaml b/egomimic/hydra_configs/data/mecka_flagship_fold_kp_overfit1.yaml new file mode 100644 index 0000000000..0c243a62dc --- /dev/null +++ b/egomimic/hydra_configs/data/mecka_flagship_fold_kp_overfit1.yaml @@ -0,0 +1,22 @@ +# Overfit check: ONE flagship fold_clothes episode (191 frames), trained on and +# validated on itself (valid runs without dropout or augmentation, so it reads +# how well the episode was memorized). train_viz would duplicate valid; the +# train config turns it off. +defaults: + - mecka_flagship_fold_kp_holdout + - _self_ + +train_datasets: + human_bimanual: + filters: + episode_hashes: [692fe87a3c1dfb13a7a99ee6] + +valid_datasets: + human_bimanual: + filters: + episode_hashes: [692fe87a3c1dfb13a7a99ee6] + +video_episodes: + valid: + - 692fe87a3c1dfb13a7a99ee6 + train_viz: [] diff --git a/egomimic/hydra_configs/data/mecka_freeform_fold_kp_all_holdout.yaml b/egomimic/hydra_configs/data/mecka_freeform_fold_kp_all_holdout.yaml new file mode 100644 index 0000000000..56999e1e9b --- /dev/null +++ b/egomimic/hydra_configs/data/mecka_freeform_fold_kp_all_holdout.yaml @@ -0,0 +1,18 @@ +# mecka_freeform_fold_kp_top3_holdout on every freeform folding_clothes operator: +# the same seen_op_valid (top3's 6 held-out episodes) and unseen_op_valid, whose 5 +# operators (`held_out_operators`) are dropped from train. After the blocklist +# (2026-09-29): 423 of 462 episodes, 9.6 h; train is 171 operators, 407 +# episodes, 9.2 h. +defaults: + - mecka_freeform_fold_kp_top3_holdout + - _self_ + +train_datasets: + human_bimanual: + filters: + filter_lambdas: + - "lambda row: str(row.get('lab', '')).lower() == 'mecka'" + - "lambda row: str(row.get('task', '')).lower() in ${data.tasks}" + - "lambda row: '/mecka/freeform/' in str(row.get('zarr_processed_path', ''))" + - "lambda row: str(row.get('operator', '')) not in ${data.held_out_operators}" + - "lambda row: str(row.get('episode_hash', '')) not in ${data.held_out_hashes}" diff --git a/egomimic/hydra_configs/data/mecka_freeform_fold_kp_op2_holdout.yaml b/egomimic/hydra_configs/data/mecka_freeform_fold_kp_op2_holdout.yaml new file mode 100644 index 0000000000..866c003511 --- /dev/null +++ b/egomimic/hydra_configs/data/mecka_freeform_fold_kp_op2_holdout.yaml @@ -0,0 +1,100 @@ +# Mecka FREEFORM fold clothes from ONE operator, 68e0b875c33a1abcb8fc55b9 (#2 by +# total freeform hours: 21 eps / 0.43 h), on the keypoint pipeline of +# mecka_flagship_fold_kp_holdout. All episodes, annotated or not: the 6 without +# segments get the empty default prompt, RDT's no-instruction condition. +# train = 19 eps, 42.0k frames = 0.39 h (mode total) +# seen_op_valid = the 2 `held_out_hashes`, median-length episodes, 4.7k frames +# unseen_op_valid = the 2 episodes of each of 5 operators with exactly 2 +# annotated freeform fold episodes, so the all-operator +# run (mecka_freeform_fold_kp_all_holdout) loses no more +# than these 10 to hold them out +# None of the operator's episodes are on the bad-egomotion blocklist (2026-09-29). +# mecka_freeform_fold_kp_top3_holdout inherits this file. +defaults: + - mecka_flagship_fold_kp_holdout + - _self_ + +tasks: [folding_clothes] + +held_out_hashes: + - 69b49d838cd7957ebe478c88 + - 69b4b52e6d78ad6736570e97 + +valid_prefix: seen_op_valid + +# The unseen_op_valid operators; mecka_freeform_fold_kp_all_holdout keeps them out +# of train. Each has exactly the 2 freeform fold episodes pinned below. +held_out_operators: + - 69439c6ea2a8b5aee7660435 + - 699f2e580aa6425a245d2e6f + - 699a5a3b6fe91711a3b1c45b + - 69ac8ebed84e313181d82dd4 + - 6992b19ab9d55c9c2d48432c + +train_datasets: + human_bimanual: + filters: + filter_lambdas: + - "lambda row: str(row.get('lab', '')).lower() == 'mecka'" + - "lambda row: str(row.get('task', '')).lower() in ${data.tasks}" + - "lambda row: '/mecka/freeform/' in str(row.get('zarr_processed_path', ''))" + - "lambda row: str(row.get('operator', '')) == '68e0b875c33a1abcb8fc55b9'" + - "lambda row: str(row.get('episode_hash', '')) not in ${data.held_out_hashes}" + +# Val heads widen the norm-stat bounds from 0.5x to 10x each cell's quantile +# range. Episodes the stats never saw (other operators; with op2's +# one-operator stats, even its own held-out episodes) land just outside it, and +# swapping those frames out biased the metric and then aborted validation +# (406042). 10x still rejects garbage, but not a ~1 m egomotion jump, so these +# heads mask frames around pose glitches instead. train_viz is the stats' own +# split and keeps the defaults. +valid_datasets: + human_bimanual: + bounds_rel_slack: 10.0 + pose_glitch_mask: {step_m: 0.05, rot_deg: 10.0} + +unseen_op_valid_datasets: + human_bimanual: + bounds_rel_slack: 10.0 + pose_glitch_mask: {step_m: 0.05, rot_deg: 10.0} + _target_: ${data.train_datasets.human_bimanual._target_} + resolver: ${data.train_datasets.human_bimanual.resolver} + filters: + _target_: egomimic.rldb.filters.DatasetFilter + filter_lambdas: + - "lambda row: str(row.get('lab', '')).lower() == 'mecka'" + # a pin from any other operator is a PinError + - "lambda row: str(row.get('operator', '')) in ${data.held_out_operators}" + episode_hashes: + - 696ba3f375020a567fd375b4 # 69439c6ea2a8b5aee7660435 + - 696ba3b9a491d5abbd79528b + - 69b9674cfa0ca5629329edae # 699f2e580aa6425a245d2e6f + - 69b9996a9d8f2a52dc2e0406 + - 69b91de3bd0a1226e8720ad0 # 699a5a3b6fe91711a3b1c45b + - 69b91a05de59539246182a53 + - 69bad9296984eb7262b04a2a # 69ac8ebed84e313181d82dd4 + - 69ba0421ad41965cca81ccca + - 69b93da22e71fa6c931cbf9c # 6992b19ab9d55c9c2d48432c + - 69b92d11b3cc90fa8d9acd2c + mode: total + +unseen_op_valid_dataloader_params: + human_bimanual: + batch_size: 32 + num_workers: 8 + shuffle: false + +# limit_val_batches x 32 = 2560 frames per head, fewer than the val episodes +# hold, so every head subsamples evenly per episode. +metric_frames_per_episode: + valid: auto + unseen_op_valid: auto + train_viz: auto + +video_episodes: + valid: + - 69b49d838cd7957ebe478c88 + unseen_op_valid: + - 69b9674cfa0ca5629329edae # 699f2e580aa6425a245d2e6f + train_viz: + - 69b4bd2c47952ada7592905d diff --git a/egomimic/hydra_configs/data/mecka_freeform_fold_kp_top3_holdout.yaml b/egomimic/hydra_configs/data/mecka_freeform_fold_kp_top3_holdout.yaml new file mode 100644 index 0000000000..8b17dbfd10 --- /dev/null +++ b/egomimic/hydra_configs/data/mecka_freeform_fold_kp_top3_holdout.yaml @@ -0,0 +1,28 @@ +# mecka_freeform_fold_kp_op2_holdout on freeform folding_clothes operators #2-#4 +# by total hours (#1, 695d09fc83a9fdf2d84d9c11, has 15 of 21 episodes on the +# bad-egomotion blocklist), 2 median-length episodes of each held out; the same +# unseen_op_valid. After the blocklist (2026-09-29): +# 68e0b875c33a1abcb8fc55b9 21 of 21 -> 19 train, 2 val +# 695d0a7783a9fdf2d84d9db5 15 of 16 -> 13 train, 2 val +# 69624c9583a9fdf2d8601cf9 13 of 13 -> 11 train, 2 val +defaults: + - mecka_freeform_fold_kp_op2_holdout + - _self_ + +held_out_hashes: + - 696d119e1fe17643f95f9fb9 # 69624c9583a9fdf2d8601cf9 + - 696d0fe4a199224e8d66242c + - 69b49d838cd7957ebe478c88 # 68e0b875c33a1abcb8fc55b9 + - 69b4b52e6d78ad6736570e97 + - 69b4b33c5dac6770e6021e42 # 695d0a7783a9fdf2d84d9db5 + - 69b4b9abe8279eda42670225 + +train_datasets: + human_bimanual: + filters: + filter_lambdas: + - "lambda row: str(row.get('lab', '')).lower() == 'mecka'" + - "lambda row: str(row.get('task', '')).lower() in ${data.tasks}" + - "lambda row: '/mecka/freeform/' in str(row.get('zarr_processed_path', ''))" + - "lambda row: str(row.get('operator', '')) in ('68e0b875c33a1abcb8fc55b9', '695d0a7783a9fdf2d84d9db5', '69624c9583a9fdf2d8601cf9')" + - "lambda row: str(row.get('episode_hash', '')) not in ${data.held_out_hashes}" diff --git a/egomimic/hydra_configs/model/rdt_1b_bc_mecka_6d_dinov3_vitb.yaml b/egomimic/hydra_configs/model/rdt_1b_bc_mecka_6d_dinov3_vitb.yaml index d60059a7b1..a0201b2c5a 100644 --- a/egomimic/hydra_configs/model/rdt_1b_bc_mecka_6d_dinov3_vitb.yaml +++ b/egomimic/hydra_configs/model/rdt_1b_bc_mecka_6d_dinov3_vitb.yaml @@ -8,4 +8,5 @@ defaults: robomimic_model: width: 2048 trunk: + depth: 28 num_heads: 32 diff --git a/egomimic/hydra_configs/model/rdt_1b_bc_mecka_kp_dinov3_vitb.yaml b/egomimic/hydra_configs/model/rdt_1b_bc_mecka_kp_dinov3_vitb.yaml new file mode 100644 index 0000000000..2987863d9b --- /dev/null +++ b/egomimic/hydra_configs/model/rdt_1b_bc_mecka_kp_dinov3_vitb.yaml @@ -0,0 +1,11 @@ +# RDT_1B_BC_MECKA_KP_DINOV3_VITB: rdt_bc_mecka_kp_dinov3_vitb at the paper's +# RDT-1B shape (width 2048 x 28 blocks x 32 heads, ~1.2B DiT). +defaults: + - rdt_bc_mecka_kp_dinov3_vitb + - _self_ + +robomimic_model: + width: 2048 + trunk: + depth: 28 + num_heads: 32 diff --git a/egomimic/hydra_configs/model/rdt_bc_mecka_6d_dinov3_vitb.yaml b/egomimic/hydra_configs/model/rdt_bc_mecka_6d_dinov3_vitb.yaml index f0539898f9..a095de10d5 100644 --- a/egomimic/hydra_configs/model/rdt_bc_mecka_6d_dinov3_vitb.yaml +++ b/egomimic/hydra_configs/model/rdt_bc_mecka_6d_dinov3_vitb.yaml @@ -32,11 +32,12 @@ # - condition dropout per stream: state and language 0.1, no image dropout # with a single camera. # -# DiT: width 1024 x 28 blocks x 16 heads (head_dim 64, RDT-1B's), RDT's FFN -# ratio 1 -> 10.5M / block. -# Trainable 392.7M = DiT 296.6M + DINOv3 85.6M + adaptors; Qwen 596M frozen. -# H100, batch 32, bf16: 212 ms / step, 30 GiB (ViT-L: 365 ms, 49 GiB; the 1B -# DiT: 359 ms, 57 GiB); sampling 0.3 s per batch. +# DiT: width 1280 x 18 blocks x 20 heads (head_dim 64, RDT-1B's), RDT's FFN +# ratio 1 -> 16.4M / block, ~295M; width and depth both ~0.63x RDT-1B's 2048 x 28. +# Trainable ~391M = DiT ~295M + DINOv3 85.6M + adaptors; Qwen 596M frozen. +# Measured at the former 1024 x 28 shape -- H100, batch 32, bf16: 212 ms / step, +# 30 GiB (ViT-L: 365 ms, 49 GiB; the 1B DiT: 359 ms, 57 GiB); sampling 0.3 s +# per batch. _target_: egomimic.pl_utils.pl_model.ModelWrapper robomimic_model: _target_: egomimic.algo.rdt.RDT @@ -44,7 +45,7 @@ robomimic_model: dims: human_bimanual: {proprio: 20, action: 18} # DiT width; every stem / encoder output_dim and the head's hidden_dim. - width: 1024 + width: 1280 ac_keys: human_bimanual: "actions_cartesian" @@ -60,8 +61,8 @@ robomimic_model: trunk: embed_dim: ${model.robomimic_model.width} - depth: 28 - num_heads: 16 + depth: 18 + num_heads: 20 mlp_ratio: 1.0 # RDT's cond_mask_prob, per stream (train only). A dropped prompt becomes # "", the no-prompt deployment condition; `img` is per camera and never diff --git a/egomimic/hydra_configs/model/rdt_bc_mecka_kp_dinov3_vitb.yaml b/egomimic/hydra_configs/model/rdt_bc_mecka_kp_dinov3_vitb.yaml new file mode 100644 index 0000000000..a7dd284c70 --- /dev/null +++ b/egomimic/hydra_configs/model/rdt_bc_mecka_kp_dinov3_vitb.yaml @@ -0,0 +1,148 @@ +# RDT_BC_MECKA_KP_DINOV3_VITB: rdt_bc_mecka_6d_dinov3_vitb (see it for the +# architecture notes) on the 144-D wrist-frame hand-keypoint action with the +# 144-D keypoint state at t-3 .. t, and RDT-1B's image pair: the current frame +# and the one stored before it, so the data config must set +# key_map.proprio_history: 4, history_stride: 1 and image_history_gap_s: 0.033 +# (t-1 at mecka's 30 fps). Pairs with evaluator=eval_hpt_keypoints_wrist. +# +# ~1B total: DiT 1280 x 18 (~295M) + DINOv3 ViT-B (86M) + frozen Qwen (596M). +# Width and depth both ~0.63x RDT-1B's 2048 x 28 (head_dim 64), not width alone. +# Standalone rather than a child of the 6d config: the state stem is renamed, +# and an inherited, unused state_ee_pose stem would trip DDP. +_target_: egomimic.pl_utils.pl_model.ModelWrapper +robomimic_model: + _target_: egomimic.algo.rdt.RDT + camera_transforms: {} + dims: + human_bimanual: {proprio: 144, action: 144} + # DiT width; every stem / encoder output_dim and the head's hidden_dim. + width: 1280 + ac_keys: + human_bimanual: "actions_keypoints" + + diffusion: true + + annotation_key: "annotations" + annotation_sampling_mode: "random" + annotation_modality: "annotation" + default_prompt: "" + + use_pad_mask: false + + trunk: + embed_dim: ${model.robomimic_model.width} + depth: 18 + num_heads: 20 + mlp_ratio: 1.0 + # RDT's cond_mask_prob, per stream (train only). A dropped prompt becomes + # "", the no-prompt deployment condition; `img` is per camera and never + # blinds a sample, so it is inert with one camera. + cond_drop: {state: 0.1, lang: 0.1, img: 0.0} + # frames per camera (current last); 2 <-> key_map.image_history_gap_s set + image_history: 2 + image_history_dropout: 0.0 + # the prompt encoder's max_length: one language shape for torch.compile + lang_pad_to: 128 + + multitask: false + pretrained: false + pretrained_checkpoint: null + domains: ["human_bimanual"] + shared_obs_keys: ["front_img_1", "annotation"] + + head_specs: + human_bimanual: + _target_: egomimic.models.fm_policy.FMPolicy + action_horizon: 100 + num_inference_steps: 10 + pooling: null + time_dist: "beta" + infer_ac_dims: + human_bimanual: ${model.robomimic_model.dims.human_bimanual.action} + model: + _target_: egomimic.models.rdt_nets.RDTDenoiser + act_dim: ${model.robomimic_model.dims.human_bimanual.action} + act_seq: 100 + hidden_dim: ${model.robomimic_model.width} + # proprio stems x history_len + n_state_tokens: 4 + # FMPolicy's t in [0, 1] -> RDT's [0, 1000) timestep embedding + time_scale: 1000.0 + + shared_stem_specs: + # image adaptor (DINOv3Stem.proj + this MLP ~ RDT's mlp2x_gelu) + front_img_1: + _target_: egomimic.models.hpt_nets.MLPPolicyStem + input_dim: ${model.robomimic_model.width} + output_dim: ${model.robomimic_model.width} + widths: ["${model.robomimic_model.width}"] + # RDTModel adds the mlp2x_gelu language adaptor on top of these tokens. + annotation: + _target_: egomimic.models.hpt_nets.QwenPerTokenEncoder + model_name: "Qwen/Qwen3-Embedding-0.6B" + max_length: 128 + freeze: true + dtype: "float16" + cache_text_features: true + text_cache_max_entries: 8192 + output_dim: ${model.robomimic_model.width} + + stem_specs: + human_bimanual: + # state adaptor (RDT's mlp3x_gelu); one token per proprio step + state_keypoints: + _target_: egomimic.models.hpt_nets.MLPPolicyStem + input_dim: ${model.robomimic_model.dims.human_bimanual.proprio} + output_dim: ${model.robomimic_model.width} + widths: + - ${model.robomimic_model.width} + - ${model.robomimic_model.width} + history_len: 4 + history_dropout: 0.2 + + encoder_specs: + front_img_1: + _target_: egomimic.models.hpt_nets.DINOv3Stem + model_name: "vit_base_patch16_dinov3.lvd1689m" + output_dim: ${model.robomimic_model.width} + freeze_backbone: false + # [H, W]; 360x640 frames at their native aspect + image_size: [224, 400] + + train_image_augs: + _target_: torchvision.transforms.Compose + transforms: + - _target_: torchvision.transforms.ColorJitter + brightness: 0.1 + contrast: 0.1 + saturation: 0.1 + hue: 0.05 + eval_image_augs: null + +# DINOv3 tower at 1e-5 against 1e-4 everywhere else (ABC-DiT's 0.1). +backbone_lr_scale: 0.1 + +# RDT and ABC-DiT agree: AdamW 1e-4, wd 1e-2, constant LR. ABC's 1k-step +# warmup and (0.9, 0.95) betas; gradient clipping is in the train config. +optimizer: + _target_: torch.optim.AdamW + _partial_: true + _convert_: all + lr: 1e-4 + weight_decay: 0.01 + betas: [0.9, 0.95] + +# LinearLR holds end_factor after total_iters: warmup, then constant. +scheduler: + _target_: torch.optim.lr_scheduler.LinearLR + _partial_: true + start_factor: 1.0e-3 + end_factor: 1.0 + total_iters: 1000 +scheduler_interval: step + +# RDT's EMA (pl_utils/ema.py); when on, validation and eval use the EMA weights. +ema: + enabled: false + power: 0.75 + max_decay: 0.9999 diff --git a/egomimic/hydra_configs/train_zarr_cartesian.yaml b/egomimic/hydra_configs/train_zarr_cartesian.yaml index 7ccfe6bc95..72dad6a35e 100644 --- a/egomimic/hydra_configs/train_zarr_cartesian.yaml +++ b/egomimic/hydra_configs/train_zarr_cartesian.yaml @@ -56,6 +56,10 @@ norm_stats: num_workers: 4 save_cache_dir: ${hydra:runtime.output_dir} precomputed_norm_path: null + # history_stats.json from precompute_norm_stats.py --history-max-lag: per-lag stats for + # the head-frame wrist translation of a proprio history (K > 1). Unset, every + # history step uses the current-step stats. + history_stats_path: null # Content-keyed cache shared across runs: one entry per split, plus per-episode # samples under episodes/ that any other split of the same recipe reuses. cache_dir: ${paths.cache_dir}/norm_stats diff --git a/egomimic/hydra_configs/train_zarr_mecka_all_kp_rdt_mem_base.yaml b/egomimic/hydra_configs/train_zarr_mecka_all_kp_rdt_mem_base.yaml new file mode 100644 index 0000000000..9b67e915db --- /dev/null +++ b/egomimic/hydra_configs/train_zarr_mecka_all_kp_rdt_mem_base.yaml @@ -0,0 +1,43 @@ +# RDT memory ablation on ALL mecka with the 144-D keypoint action and state, +# BASELINE arm: the recipe of train_zarr_mecka_20pct_6d_rdt_mem_base (8xH100, +# 32 / GPU, 1000 x 100 steps) with RDT-1B's short pair: image and proprio at t and +# t-1 frame. +# _proprio + proprio history: 8 steps 5 frames apart (~1.2 s) +# _proprio_frames _proprio with 4 full front frames 10 apart (t-30 .. t) +# through the main encoder instead of the pair +# _short_long_encoder _proprio + 8 front frames on the proprio grid, one CLS +# token each from a second DINOv3, fused with proprio; +# pair dropped for 30 % +# All arms share one set of norm stats; K>1 wrist xyz uses per-step history stats. +defaults: + - train_zarr_mecka_20pct_6d_rdt_mem_base + - override evaluator@train_viz_evaluator: train_viz_hpt_keypoints_wrist + - override model: rdt_bc_mecka_kp_dinov3_vitb + - override data: mecka_all_kp_pinned_val + - override evaluator: eval_hpt_keypoints_wrist + - _self_ + +data: + train_datasets: + human_bimanual: + resolver: + key_map: + # RDT-1B's pair is the previous stored frame (hdf5_vla_dataset.py): + # one frame back at 30 fps, for both the image and the proprio. + image_history_gap_s: 0.033 + proprio_history: 2 + history_stride: 1 + +model: + robomimic_model: + stem_specs: + human_bimanual: + state_keypoints: + history_len: 2 + # [s_t, s_t] would read as a stationary hand; cond_drop.state + # already guards against leaning on proprio. + history_dropout: 0.0 + head_specs: + human_bimanual: + model: + n_state_tokens: 2 diff --git a/egomimic/hydra_configs/train_zarr_mecka_all_kp_rdt_mem_proprio.yaml b/egomimic/hydra_configs/train_zarr_mecka_all_kp_rdt_mem_proprio.yaml new file mode 100644 index 0000000000..8b028d2fa9 --- /dev/null +++ b/egomimic/hydra_configs/train_zarr_mecka_all_kp_rdt_mem_proprio.yaml @@ -0,0 +1,26 @@ +# RDT memory ablation (all mecka, keypoints), PROPRIO arm: the baseline plus 8 +# keypoint-state steps 5 frames apart (~1.2 s, current last), one token each; +# 0.25 % of samples fail the bounds check vs 0.05 % at K=1 (reject_probe.py). +# history_dropout stays 0 as in the base. See train_zarr_mecka_all_kp_rdt_mem_base.yaml. +defaults: + - train_zarr_mecka_all_kp_rdt_mem_base + - _self_ + +data: + train_datasets: + human_bimanual: + resolver: + key_map: + proprio_history: 8 + history_stride: 5 + +model: + robomimic_model: + stem_specs: + human_bimanual: + state_keypoints: + history_len: 8 + head_specs: + human_bimanual: + model: + n_state_tokens: 8 diff --git a/egomimic/hydra_configs/train_zarr_mecka_all_kp_rdt_mem_proprio_frames.yaml b/egomimic/hydra_configs/train_zarr_mecka_all_kp_rdt_mem_proprio_frames.yaml new file mode 100644 index 0000000000..351fb76879 --- /dev/null +++ b/egomimic/hydra_configs/train_zarr_mecka_all_kp_rdt_mem_proprio_frames.yaml @@ -0,0 +1,20 @@ +# RDT memory ablation (all mecka, keypoints), PROPRIO_FRAMES arm: the proprio arm with 4 +# full front frames 10 apart (t-30 .. t, 1 s) through the main encoder instead of +# RDT's pair; every image lands on every other proprio step. +# See train_zarr_mecka_all_kp_rdt_mem_base.yaml. +defaults: + - train_zarr_mecka_all_kp_rdt_mem_proprio + - _self_ + +data: + train_datasets: + human_bimanual: + resolver: + key_map: + image_history_gap_s: 0.333 + image_history: 4 + +model: + robomimic_model: + trunk: + image_history: 4 diff --git a/egomimic/hydra_configs/train_zarr_mecka_all_kp_rdt_mem_short_long_encoder.yaml b/egomimic/hydra_configs/train_zarr_mecka_all_kp_rdt_mem_short_long_encoder.yaml new file mode 100644 index 0000000000..00db387c13 --- /dev/null +++ b/egomimic/hydra_configs/train_zarr_mecka_all_kp_rdt_mem_short_long_encoder.yaml @@ -0,0 +1,44 @@ +# RDT memory ablation (all mecka, keypoints), SHORT_LONG_ENCODER arm: the proprio arm plus 8 +# front frames on the proprio's grid (5 frames apart, t-35 .. t), one CLS token +# each from a second fine-tuned DINOv3, fused with that step's proprio token; +# RDT's pair is dropped for 30 % of training samples. +# See train_zarr_mecka_all_kp_rdt_mem_base.yaml. +defaults: + - train_zarr_mecka_all_kp_rdt_mem_proprio + - _self_ + +data: + train_datasets: + human_bimanual: + resolver: + key_map: + # Seconds, not frames, so fusion can check it against the memory's grid; + # 5 frames at mecka's 30 fps (history stats raise on any other fps). + history_stride: 1 + history_stride_s: 0.16667 + image_memory: 8 + image_memory_stride_s: 0.16667 + +model: + robomimic_model: + trunk: + short_dropout: 0.3 + memory: + encoder: + _target_: egomimic.models.hpt_nets.DINOv3Stem + model_name: ${model.robomimic_model.encoder_specs.front_img_1.model_name} + output_dim: ${model.robomimic_model.width} + freeze_backbone: false + image_size: ${model.robomimic_model.encoder_specs.front_img_1.image_size} + pool: cls + grad_checkpointing: true + # Must equal key_map.image_memory / image_memory_stride_s above: the + # model is built from the model config alone, so it cannot interpolate + # them. A frame-count mismatch raises at the first batch. + frames: 8 + stride_s: 0.16667 + fuse_proprio: true + head_specs: + human_bimanual: + model: + n_memory_tokens: 8 diff --git a/egomimic/hydra_configs/train_zarr_mecka_flagship_fold_cup_kp_rdt.yaml b/egomimic/hydra_configs/train_zarr_mecka_flagship_fold_cup_kp_rdt.yaml new file mode 100644 index 0000000000..45433f2645 --- /dev/null +++ b/egomimic/hydra_configs/train_zarr_mecka_flagship_fold_cup_kp_rdt.yaml @@ -0,0 +1,5 @@ +# train_zarr_mecka_flagship_fold_kp_rdt on fold_clothes + cup_on_saucer, 3 held out. +defaults: + - train_zarr_mecka_flagship_fold_kp_rdt + - override data: mecka_flagship_fold_cup_kp_holdout + - _self_ diff --git a/egomimic/hydra_configs/train_zarr_mecka_flagship_fold_kp_rdt.yaml b/egomimic/hydra_configs/train_zarr_mecka_flagship_fold_kp_rdt.yaml new file mode 100644 index 0000000000..77e33eb926 --- /dev/null +++ b/egomimic/hydra_configs/train_zarr_mecka_flagship_fold_kp_rdt.yaml @@ -0,0 +1,28 @@ +# RDT (model/rdt_bc_mecka_kp_dinov3_vitb: DiT 1280 x 18, DINOv3 ViT-B, frozen +# Qwen; image pair t, t-1 frame; keypoint state t-3 .. t) fitting the mecka +# flagship data, in three steps: +# train_zarr_mecka_flagship_fold_kp_rdt_overfit1 one fold_clothes episode, 1 GPU +# train_zarr_mecka_flagship_fold_kp_rdt fold_clothes, 3 held out +# train_zarr_mecka_flagship_fold_cup_kp_rdt + cup_on_saucer, 3 held out +# Recipe of train_zarr_mecka_20pct_6d_rdt_mem_base: 8xH100, 32 / GPU, AdamW 1e-4, +# 1k-step warmup, grad clip 1.0, 1000 x 100 steps, val every 25 epochs. +defaults: + - train_zarr_mecka_20pct_6d_rdt_mem_base + - override evaluator@train_viz_evaluator: train_viz_hpt_keypoints_wrist + - override model: rdt_bc_mecka_kp_dinov3_vitb + - override data: mecka_flagship_fold_kp_holdout + - override evaluator: eval_hpt_keypoints_wrist + - _self_ + +data: + train_datasets: + human_bimanual: + resolver: + key_map: + # t-1 at mecka's 30 fps, RDT-1B's pair + image_history_gap_s: 0.033 + proprio_history: 4 + history_stride: 1 + +norm_stats: + max_samples: 200000 diff --git a/egomimic/hydra_configs/train_zarr_mecka_flagship_fold_kp_rdt_overfit1.yaml b/egomimic/hydra_configs/train_zarr_mecka_flagship_fold_kp_rdt_overfit1.yaml new file mode 100644 index 0000000000..079dee2150 --- /dev/null +++ b/egomimic/hydra_configs/train_zarr_mecka_flagship_fold_kp_rdt_overfit1.yaml @@ -0,0 +1,27 @@ +# train_zarr_mecka_flagship_fold_kp_rdt on a single episode, 1 GPU. Its 191 frames +# are 6 batches of 32, so an epoch is 6 steps: 840 epochs = 5040 steps (~840 +# passes), val every 84 epochs (~500 steps), video + checkpoint every 168. +defaults: + - train_zarr_mecka_flagship_fold_kp_rdt + - override data: mecka_flagship_fold_kp_overfit1 + - override hydra/launcher: submitit_loaner + - _self_ + +launch_params: + gpus_per_node: 1 + nodes: 1 + +train_viz: false + +trainer: + max_epochs: 840 + min_epochs: 840 + check_val_every_n_epoch: 84 + log_every_n_steps: 6 + +evaluator: + viz_every_n_epochs: 168 + +callbacks: + model_checkpoint: + every_n_epochs: 168 diff --git a/egomimic/hydra_configs/train_zarr_mecka_freeform_fold_kp_rdt_all.yaml b/egomimic/hydra_configs/train_zarr_mecka_freeform_fold_kp_rdt_all.yaml new file mode 100644 index 0000000000..9d277b7916 --- /dev/null +++ b/egomimic/hydra_configs/train_zarr_mecka_freeform_fold_kp_rdt_all.yaml @@ -0,0 +1,20 @@ +# train_zarr_mecka_freeform_fold_kp_rdt_op2 on every freeform fold operator but the +# 5 unseen ones (171 operators, 407 train episodes, 9.2 h), on a whole node +# (logs/freeform_fold/launch_all.sbatch). At 8 x 32 / step one pass is ~3.9k +# steps, so 1500 x 100 steps is ~38 passes, sized to fit the 48 h limit at ~1 s/step. +# LR: 1k warmup, 119k stable, 30k decay (set in launch_all.sbatch). +defaults: + - train_zarr_mecka_freeform_fold_kp_rdt_op2 + - override data: mecka_freeform_fold_kp_all_holdout + - _self_ + +launch_params: + gpus_per_node: 8 + +trainer: + max_epochs: 1500 + min_epochs: 1500 + +callbacks: + model_checkpoint: + every_n_epochs: 250 diff --git a/egomimic/hydra_configs/train_zarr_mecka_freeform_fold_kp_rdt_op2.yaml b/egomimic/hydra_configs/train_zarr_mecka_freeform_fold_kp_rdt_op2.yaml new file mode 100644 index 0000000000..eaf27b9f64 --- /dev/null +++ b/egomimic/hydra_configs/train_zarr_mecka_freeform_fold_kp_rdt_op2.yaml @@ -0,0 +1,29 @@ +# train_zarr_mecka_flagship_fold_kp_rdt at the RDT-1B shape, on +# one freeform fold operator's 19 episodes (0.39 h), 2 held out; 4 GPUs, so two +# runs share a node (logs/freeform_fold/launch_pair.sbatch). At 4 x 32 / step +# one pass is ~330 steps, so 600 x 100 steps is ~180 passes: the target is train +# error, which kept falling to ~80 passes (+ decay) on the flagship fold run. +# LR: 1k warmup, 47k stable, 12k linear decay to 0 (warmup_stable_decay); the +# launcher swaps it in, since a config merge cannot drop LinearLR's arguments. +defaults: + - train_zarr_mecka_flagship_fold_kp_rdt + - override model: rdt_1b_bc_mecka_kp_dinov3_vitb + - override data: mecka_freeform_fold_kp_op2_holdout + - override hydra/launcher: submitit_loaner + - _self_ + +launch_params: + gpus_per_node: 4 + nodes: 1 + +trainer: + max_epochs: 600 + min_epochs: 600 + check_val_every_n_epoch: 25 + +evaluator: + viz_every_n_epochs: 50 + +callbacks: + model_checkpoint: + every_n_epochs: 100 diff --git a/egomimic/hydra_configs/train_zarr_mecka_freeform_fold_kp_rdt_top3.yaml b/egomimic/hydra_configs/train_zarr_mecka_freeform_fold_kp_rdt_top3.yaml new file mode 100644 index 0000000000..ce6769b929 --- /dev/null +++ b/egomimic/hydra_configs/train_zarr_mecka_freeform_fold_kp_rdt_top3.yaml @@ -0,0 +1,7 @@ +# train_zarr_mecka_freeform_fold_kp_rdt_op2 on freeform fold operators #2-#4 +# (43 train episodes, 0.96 h). At 4 x 32 / step one pass is ~810 steps, so +# 600 x 100 steps is ~75 passes. +defaults: + - train_zarr_mecka_freeform_fold_kp_rdt_op2 + - override data: mecka_freeform_fold_kp_top3_holdout + - _self_ diff --git a/egomimic/pl_utils/pl_data_utils.py b/egomimic/pl_utils/pl_data_utils.py index d7fd4ae114..e82080d19f 100644 --- a/egomimic/pl_utils/pl_data_utils.py +++ b/egomimic/pl_utils/pl_data_utils.py @@ -9,10 +9,22 @@ logger = logging.getLogger(__name__) # Val heads that can carry a metric loader, in val_dataloader() order. -VAL_HEADS = ("valid", "train_viz", "unseen_op_valid") +VAL_HEADS = ("valid", "train_viz", "unseen_op_valid", "extra_valid") # Suffix of the video-only companion loader of a head (``valid`` -> # ``valid_video``). ModelWrapper.validation_step splits on it. VIDEO_SUFFIX = "_video" +# A hung worker blocks only its own rank; without a timeout the other ranks die +# 30 min later on the NCCL watchdog with no stack trace from the stuck one. +WORKER_TIMEOUT_S = 600 + + +def _with_worker_defaults(params: dict, *, persistent: bool) -> dict: + params = orphan_guarded(params) + if params.get("num_workers", 0) > 0: + params.setdefault("timeout", WORKER_TIMEOUT_S) + if persistent: + params.setdefault("persistent_workers", True) + return params def video_loader_name(head: str) -> str: @@ -81,11 +93,17 @@ def __init__( train_viz_dataloader_params: dict | None = None, unseen_op_valid_datasets: dict | None = None, unseen_op_valid_dataloader_params: dict | None = None, + extra_valid_datasets: dict | None = None, + extra_valid_dataloader_params: dict | None = None, + extra_valid_prefix: str | None = None, valid_prefix: str | None = None, held_out_operators: list | None = None, video_datasets: dict | None = None, metric_frames_per_episode: dict | None = None, video_episodes: dict | None = None, + tasks: list | None = None, + held_out_hashes: list | None = None, + train_operators: list | None = None, ): """ Args: @@ -105,6 +123,11 @@ def __init__( operators' held-out episodes and ``train_viz`` the train split. unseen_op_valid_dataloader_params: dict of per-dataset DataLoader kwargs for the unseen_op_valid loader. + extra_valid_datasets / extra_valid_dataloader_params: a fourth + val head, like unseen_op_valid, for a second seen split + (e.g. held-out episodes of several operators next to + ``valid``'s one). Logs under ``extra_valid_prefix`` + (config-only, default ``extra_valid``). valid_prefix: config-only. When set, trainHydra wraps the canonical evaluator so the valid loader logs ``/...`` and writes ``videos_/`` (``seen_op_valid`` in the opsplit configs). @@ -131,6 +154,10 @@ def __init__( video_episodes: config-only ``{head: [episode_hash, ...]}`` -- the pinned video episodes, one per operator per head. Provenance only; trainHydra is what reads it. + tasks, held_out_hashes, train_operators: config-only lists + interpolated by the filter lambdas of data/mecka_*_holdout.yaml and + similar split configs; accepted + like ``held_out_operators``, provenance only. Tokenization (sampling a prompt from per-sample annotation lists, splicing in embodiment / control-mode / proprio blocks, and running @@ -155,6 +182,11 @@ def __init__( k: v for k, v in (unseen_op_valid_datasets or {}).items() if v is not None } self.unseen_op_valid_dataloader_params = unseen_op_valid_dataloader_params or {} + self.extra_valid_datasets = { + k: v for k, v in (extra_valid_datasets or {}).items() if v is not None + } + self.extra_valid_dataloader_params = extra_valid_dataloader_params or {} + self.extra_valid_prefix = extra_valid_prefix self.valid_prefix = valid_prefix self.held_out_operators = list(held_out_operators or []) # {head: {dataset_name: dataset}}; heads with no pinned episodes are @@ -166,6 +198,9 @@ def __init__( self.video_datasets = {h: d for h, d in self.video_datasets.items() if d} self.metric_frames_per_episode = dict(metric_frames_per_episode or {}) self.video_episodes = dict(video_episodes or {}) + self.tasks = list(tasks or []) + self.held_out_hashes = list(held_out_hashes or []) + self.train_operators = list(train_operators or []) self.collate_fn = annotation_collate def train_dataloader(self): @@ -180,7 +215,9 @@ def train_dataloader(self): dataset, shuffle=True, collate_fn=self.collate_fn, - **orphan_guarded(dataset_params), + # Persistent: re-forking the workers every 100-step epoch is where + # a rank hung for good (job 405831). + **_with_worker_defaults(dataset_params, persistent=True), ) return CombinedLoader(iterables, "max_size_cycle") @@ -195,7 +232,7 @@ def _build_val_style_loader( raise ValueError( f"No dataloader params found for dataset {dataset_name}. Please add {dataset_name} into your data config {kind}_dataloader_params." ) - dataset_params = orphan_guarded(dataset_params) + dataset_params = _with_worker_defaults(dataset_params, persistent=False) shuffle = dataset_params.pop("shuffle", False) iterables[dataset_name] = DataLoader( dataset, @@ -247,6 +284,10 @@ def _metric_sources(self) -> dict: self.unseen_op_valid_datasets, self.unseen_op_valid_dataloader_params, ), + "extra_valid": ( + self.extra_valid_datasets, + self.extra_valid_dataloader_params, + ), } def val_loader_names(self) -> list[str]: diff --git a/egomimic/pl_utils/pl_model.py b/egomimic/pl_utils/pl_model.py index 2790dc7051..bf929c6c6e 100644 --- a/egomimic/pl_utils/pl_model.py +++ b/egomimic/pl_utils/pl_model.py @@ -85,6 +85,8 @@ def __init__( # Optional third evaluator for the unseen_op_valid loader (held-out # operators in the opsplit data configs). self.unseen_op_valid_evaluator = None + # Optional fourth evaluator for the extra_valid loader. + self.extra_valid_evaluator = None # Val loader names in dataloader_idx order, from # MultiDataModuleWrapper.val_loader_names(); None = [valid, train_viz]. self.val_loader_names = None @@ -290,6 +292,7 @@ def _val_heads(self) -> dict: "valid": self.evaluator, "train_viz": self.train_viz_evaluator, "unseen_op_valid": getattr(self, "unseen_op_valid_evaluator", None), + "extra_valid": getattr(self, "extra_valid_evaluator", None), } def on_validation_start(self): diff --git a/egomimic/rldb/embodiment/human.py b/egomimic/rldb/embodiment/human.py index ea5aa6de41..2b085be6f7 100644 --- a/egomimic/rldb/embodiment/human.py +++ b/egomimic/rldb/embodiment/human.py @@ -207,6 +207,7 @@ def get_keymap( history_stride: int = 1, history_stride_s: float | None = None, image_history_gap_s: float | None = None, + image_history: int = 2, image_memory: int = 0, image_memory_stride_s: float = 1.0, ): @@ -233,7 +234,9 @@ def get_keymap( one (frames = gap x the episode's fps; frame 0 at an episode start). The algo must consume it explicitly: HPT has no stem for it and skips augs and normalization on cameras outside its encoders, so under plain - HPT it is decoded and dropped. + HPT it is decoded and dropped. ``image_history`` (T) > 2 makes it + the T - 1 past frames that many seconds apart, ``(T - 1, 3, H, W)`` + oldest first, plus a ``_hist_mask``. ``history_stride_s`` replaces ``history_stride`` by a spacing in seconds, read in each episode's fps. ``image_memory`` (N) adds @@ -253,10 +256,16 @@ def get_keymap( ) 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 image_history < 2: + raise ValueError( + f"image_history must be >= 2 with image_history_gap_s, " + f"got {image_history}" + ) + past = {**front, "lag_s": float(image_history_gap_s)} + if image_history > 2: + past["history"] = int(image_history) - 1 + past["history_stride_s"] = float(image_history_gap_s) + key_map[f"{cls.VIZ_IMAGE_KEY}{IMAGE_HISTORY_SUFFIX}"] = past if image_memory > 0: if image_memory_stride_s <= 0: raise ValueError( diff --git a/egomimic/rldb/zarr/history_stats.py b/egomimic/rldb/zarr/history_stats.py new file mode 100644 index 0000000000..25df682339 --- /dev/null +++ b/egomimic/rldb/zarr/history_stats.py @@ -0,0 +1,277 @@ +"""History norm stats: per-lag stats for proprio observation history. + +A history step ``lag`` frames old is expressed in the CURRENT step's head +frame, so its wrist translation spreads wider than the current step's (the +past wrist carries all the head motion since). Current-step stats then clip or +reject old steps. The table holds stats per lag ``0..max_lag`` for the +channels that depend on it (the wrist / ee translation); every other channel +keeps the current-step stats. ``gather`` picks the rows for a given +``(K, stride)`` and returns ordinary ``(K, D)`` per-cell stats, so normalize, +the bounds check and checkpoints need nothing new. + +Keypoints are excluded: each history step's keypoints are in that step's own +wrist frame, so their distribution does not depend on the lag. Rotation +channels are excluded because rot6d is left unnormalized. +""" + +from __future__ import annotations + +import json +import logging +import os + +import numpy as np +from torch.utils.data import DataLoader, Dataset, Subset + +from egomimic.utils.pose_utils import ( + bimanual_cartesian_layout, + bimanual_keypoint_layout, +) + +logger = logging.getLogger(__name__) + +HISTORY_STATS_FILE = "history_stats.json" +HISTORY_MASK_KEY = "proprio_history_mask" + + +def lag_channels(zarr_key: str, width: int) -> list[int] | None: + """Channels of a proprio key whose distribution depends on the lag: the + head-frame wrist translation. ``None`` for keys without one.""" + if zarr_key == "observations.state.keypoints": + layout = bimanual_keypoint_layout(width) + return None if layout is None else sorted(layout["wrist_xyz"]) + if zarr_key == "observations.state.ee_pose": + layout = bimanual_cartesian_layout(width) + return None if layout is None else sorted(layout["xyz"]) + return None + + +def history_keymap(key_map: dict, max_lag: int) -> dict: + """The norm-mode keymap config reading a ``max_lag + 1`` step history at + stride 1, so one sample carries every lag.""" + km = dict(key_map) + km["proprio_history"] = int(max_lag) + 1 + km["history_stride"] = 1 + km["history_stride_s"] = None + return km + + +class _LagSlice(Dataset): + """Worker-side slice to the lag channels and the history mask, so a + ``(max_lag + 1, D)`` window does not cross the worker pipe whole.""" + + def __init__(self, dataset, zarr_keys: dict[str, str], channels: dict[str, list]): + self.dataset = dataset + self.zarr_keys = zarr_keys + self.channels = channels + + def __len__(self): + return len(self.dataset) + + def __getitem__(self, i): + data = self.dataset[i] + out = {HISTORY_MASK_KEY: np.asarray(data[HISTORY_MASK_KEY], dtype=np.float32)} + for name, zarr_key in self.zarr_keys.items(): + v = np.asarray(data[zarr_key], dtype=np.float32) + out[name] = v[..., self.channels[name]] + return out + + +def compute( + norm_stats, + lag_dataset, + embodiment_id: int, + max_lag: int, + n_anchors: int, + seed: int = 42, + batch_size: int = 64, + num_workers: int = 8, +) -> dict: + """Per-lag stats for every history proprio key of ``lag_dataset``. + + ``lag_dataset`` is a norm-mode MultiDataset built from ``history_keymap``; + ``norm_stats`` supplies the key names. Each lag only counts anchors whose + step at that lag lies inside the episode (history mask 1), so the + front-padding copies never shape the stats. + """ + from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset + + probe = lag_dataset[0] + if HISTORY_MASK_KEY not in probe: + raise ValueError( + "lag dataset has no proprio history; build it with history_keymap" + ) + L = int(max_lag) + 1 + zarr_keys, channels = {}, {} + for name in norm_stats.keys_of_type("proprio_keys", embodiment_id): + zk = norm_stats.keyname_to_zarr_key(name, embodiment_id) + v = probe.get(zk) + if v is None or np.ndim(v) != 2 or np.shape(v)[0] != L: + continue + ch = lag_channels(zk, np.shape(v)[-1]) + if ch: + zarr_keys[name], channels[name] = zk, ch + if not zarr_keys: + raise ValueError("no history proprio key with lag-dependent channels") + + n = min(int(n_anchors), len(lag_dataset)) + idx = np.random.default_rng(seed).choice(len(lag_dataset), size=n, replace=False) + loader = DataLoader( + Subset(_LagSlice(lag_dataset, zarr_keys, channels), idx.tolist()), + batch_size=batch_size, + num_workers=num_workers, + ) + vals = {name: [] for name in zarr_keys} + masks = [] + for batch in loader: + masks.append(batch[HISTORY_MASK_KEY].numpy()) + for name in zarr_keys: + vals[name].append(batch[name].numpy()) + mask = np.concatenate(masks) > 0 # (N, L), oldest first + + out = {} + for name in zarr_keys: + X = np.concatenate(vals[name]) # (N, L, C) + per_lag = {} + count = np.zeros(L, dtype=np.int64) + for lag in range(L): + j = L - 1 - lag + rows = X[mask[:, j], j] + rows = rows[np.isfinite(rows).all(axis=1)] + count[lag] = len(rows) + if len(rows) == 0: + raise ValueError( + f"{name}: no anchor has a real step at lag {lag}; lower max_lag" + ) + for stat, arr in MultiDataset._compute_stats_for_array(rows).items(): + per_lag.setdefault(stat, []).append(arr) + out[name] = { + "zarr_key": zarr_keys[name], + "channels": channels[name], + "count": count, + "stats": {s: np.stack(a).astype(np.float32) for s, a in per_lag.items()}, + } + logger.info( + f"[history_stats] {name}: {n} anchors, lag {max_lag} has {count[-1]} real steps" + ) + return out + + +def save(path: str, embodiment_id: int, table: dict, max_lag: int) -> str: + """Write/merge ``table`` for ``embodiment_id`` into ``path`` (a file or a + norm_stats dir).""" + if os.path.isdir(path) or not path.endswith(".json"): + os.makedirs(path, exist_ok=True) + path = os.path.join(path, HISTORY_STATS_FILE) + payload = {"max_lag": int(max_lag), "stats": {}} + if os.path.isfile(path): + with open(path) as f: + payload = json.load(f) + if payload["max_lag"] != int(max_lag): + raise ValueError(f"{path} has max_lag={payload['max_lag']}, not {max_lag}") + payload["stats"][str(embodiment_id)] = { + name: { + "zarr_key": e["zarr_key"], + "channels": list(map(int, e["channels"])), + "count": np.asarray(e["count"]).tolist(), + "stats": {s: np.asarray(a).tolist() for s, a in e["stats"].items()}, + } + for name, e in table.items() + } + with open(path, "w") as f: + json.dump(payload, f) + return path + + +def load(path: str, embodiment_id: int) -> tuple[dict, int]: + """``(table, max_lag)`` for ``embodiment_id`` from a file or norm_stats dir.""" + if os.path.isdir(path): + path = os.path.join(path, HISTORY_STATS_FILE) + with open(path) as f: + payload = json.load(f) + raw = payload["stats"].get(str(embodiment_id)) + if raw is None: + raise ValueError(f"{path} has no lag stats for embodiment id {embodiment_id}") + table = { + name: { + "zarr_key": e["zarr_key"], + "channels": e["channels"], + "count": np.asarray(e["count"]), + "stats": { + s: np.asarray(a, dtype=np.float32) for s, a in e["stats"].items() + }, + } + for name, e in raw.items() + } + return table, int(payload["max_lag"]) + + +def gather(base: dict, entry: dict, history: int, stride: int) -> dict: + """``(K, D)`` stats for a K-step history ``stride`` frames apart (oldest + first): the current-step ``(D,)`` stats tiled, with the lag channels of + step k taken from lag ``stride * (K - 1 - k)``.""" + lags = stride * (history - 1 - np.arange(history)) + max_lag = entry["stats"]["mean"].shape[0] - 1 + if lags[0] > max_lag: + raise ValueError( + f"history {history} x stride {stride} reaches lag {lags[0]}, past the " + f"table's max_lag {max_lag}" + ) + ch = np.asarray(entry["channels"]) + out = {} + for stat, arr in base.items(): + arr = np.asarray(arr, dtype=np.float32) + if arr.ndim != 1: + raise ValueError(f"expected current-step (D,) stats, got {arr.shape}") + tiled = np.tile(arr, (history, 1)) + if stat in entry["stats"]: + tiled[:, ch] = entry["stats"][stat][lags] + out[stat] = tiled + return out + + +def apply(norm_stats, embodiment_id: int, path: str, history: int, stride: int) -> list: + """Replace the current-step stats of each history key in ``norm_stats`` + with its gathered ``(K, D)`` stats. Keys already at ``(K, D)`` (stats loaded + from a run that gathered them) are left as they are. Returns the keys + changed.""" + table, _ = load(path, embodiment_id) + per_emb = norm_stats.norm_stats.get(embodiment_id, {}) + changed = [] + for name, entry in table.items(): + base = per_emb.get(name) + if base is None: + continue + if np.ndim(base["mean"]) == 2: + if np.shape(base["mean"])[0] != history: + raise ValueError( + f"{name}: stats are {np.shape(base['mean'])}, not a " + f"{history}-step history" + ) + continue + per_emb[name] = gather(base, entry, history, stride) + changed.append(name) + return changed + + +def history_stride_frames(key_map: dict, dataset) -> int: + """The history stride in frames: ``history_stride`` or + ``history_stride_s`` x fps, which must then be the same for every + episode (the stats have one row per lag in frames).""" + stride_s = key_map.get("history_stride_s") + if stride_s is None: + return int(key_map.get("history_stride", 1) or 1) + strides = set() + stack = [dataset] + while stack: + ds = stack.pop() + if hasattr(ds, "datasets"): + stack.extend(ds.datasets.values()) + elif hasattr(ds, "_seconds_to_frames"): + strides.add(ds._seconds_to_frames(stride_s)) + if len(strides) != 1: + raise ValueError( + f"history_stride_s={stride_s} is {sorted(strides)} frames across " + "episodes; per-lag stats need one stride in frames (set history_stride)" + ) + return strides.pop() diff --git a/egomimic/rldb/zarr/zarr_dataset_multi.py b/egomimic/rldb/zarr/zarr_dataset_multi.py index def8107c69..a1450ea5ef 100644 --- a/egomimic/rldb/zarr/zarr_dataset_multi.py +++ b/egomimic/rldb/zarr/zarr_dataset_multi.py @@ -22,6 +22,7 @@ import copy import fcntl +import hashlib import json import logging import math @@ -31,7 +32,6 @@ import subprocess import tempfile import time -from concurrent.futures import ThreadPoolExecutor from pathlib import Path from typing import TYPE_CHECKING, Any, Iterable, Mapping @@ -216,6 +216,29 @@ def get_fallback_idx( return random.choice(valid_candidates), attempts +def pose_glitch_frames( + head_pose: np.ndarray, step_m: float = 0.05, rot_deg: float = 10.0 +) -> np.ndarray: + """Frames on either side of a camera-pose jump: a one-frame move over + ``step_m`` or turn over ``rot_deg``. ``head_pose`` is (T, 7) xyz + wxyz.""" + from scipy.spatial.transform import Rotation + + step = np.linalg.norm(np.diff(head_pose[:, :3], axis=0), axis=1) + q = head_pose[:, 3:7] + ok = np.linalg.norm(q, axis=1) > 0 + rot = np.zeros(len(step)) + pair = ok[1:] & ok[:-1] + if pair.any(): + r = ( + Rotation.from_quat(q[1:][pair], scalar_first=True) + * Rotation.from_quat(q[:-1][pair], scalar_first=True).inv() + ) + rot[pair] = np.degrees(r.magnitude()) + jump = np.flatnonzero((step > step_m) | (rot > rot_deg)) + 1 + # Which side of a jump is wrong is unknown, so both are. + return np.unique(np.concatenate([jump - 1, jump])) + + def _bounds_check_channels(zarr_key: str, width: int) -> list[int] | None: """Channels of a known bimanual layout that quantile bounds apply to (translation, gripper, keypoints), or ``None`` to check the full vector. @@ -319,6 +342,13 @@ def _warn_embodiment_mismatch( # Never read from /workspace NFS: those reads have hung ranks unkillably and drained nodes. DEFAULT_STAGE_DIR = f"/ephemeral/loaner-jobs/{os.getuid()}/egoverse-datasets" _STAGE_MIN_FREE_FRAC = 0.1 +_LAUNCH_RECORD_TTL_S = 7 * 24 * 3600 +# A staged copy is re-checked against the source at least this often, which +# catches NFS data rewritten without its SQL updated_at changing. +_STAGED_TRUST_S = 7 * 24 * 3600 +_STAGE_CHUNK = 64 +# Longer than any job (48 h limit), so an episode this stale has no reader. +_STAGE_EVICT_AGE_S = 3 * 24 * 3600 class StagingError(RuntimeError): @@ -332,83 +362,280 @@ def _episode_dir(root: Path, name: str) -> Path | None: return None +def _launch_key() -> str | None: + """Identifies the ranks of one launch: an ``srun`` step, or else the process + group a launcher spawned its ranks in. None outside a Slurm job.""" + job = os.environ.get("SLURM_JOB_ID") + if job is None: + return None + step = os.environ.get("SLURM_STEP_ID") + step = f"s{step}" if step is not None else f"pg{os.getpgid(0)}" + return f"{job}.{step}.{os.environ.get('SLURM_RESTART_COUNT', '0')}" + + +def _launch_record(stage_root: Path, kind: str, source: Path) -> Path | None: + key = _launch_key() + if key is None: + return None + src = hashlib.sha256(str(source).encode()).hexdigest()[:16] + return stage_root / kind / f"{key}.{src}.json" + + +def _read_json(path: Path | None) -> dict: + if path is None: + return {} + try: + return json.loads(path.read_text()) + except (OSError, ValueError): + return {} + + +def _write_json(path: Path, data: dict) -> None: + path.parent.mkdir(exist_ok=True) + tmp = path.with_name(f".{path.name}.tmp") + tmp.write_text(json.dumps(data)) + tmp.rename(path) + + +def _prune(directory: Path, ttl_s: float) -> None: + cutoff = time.time() - ttl_s + for old in [*directory.glob("*.json"), *directory.glob(".*.tmp")]: + try: + if old.stat().st_mtime < cutoff: + old.unlink() + except OSError: + pass + + +def _tar_copy(src_root: Path, names: list[str], dest: Path) -> None: + """Copy ``src_root/`` trees into ``dest`` as one sequential tar stream.""" + # --dereference: a symlinked episode must land as data, not a link back to NFS. + # pax keeps sub-second mtimes, which the fingerprints and norm-stat keys compare. + create = ["tar", "-C", str(src_root), "--dereference", "--format=posix"] + with tempfile.TemporaryFile() as read_err, tempfile.TemporaryFile() as write_err: + reader = subprocess.Popen( + [*create, "--null", "-T", "-", "-cf", "-"], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=read_err, + ) + writer = subprocess.Popen( + ["tar", "-C", str(dest), "-xf", "-"], + stdin=reader.stdout, + stderr=write_err, + ) + reader.stdout.close() + try: + reader.stdin.write(b"".join(n.encode() + b"\0" for n in names)) + reader.stdin.close() + except BrokenPipeError: + pass + writer.wait() + reader.wait() + if reader.returncode or writer.returncode: + read_err.seek(0) + write_err.seek(0) + detail = (read_err.read()[-2000:] + write_err.read()[-2000:]).decode( + errors="replace" + ) + raise OSError(f"tar {src_root} -> {dest} failed: {detail.strip()}") + + +def staged_episode_dir(stage_root: Path, src_root: Path) -> Path: + """Where ``stage_episodes`` puts episodes from ``src_root``: one subdirectory + per source, so the same hash from two sources never shares a copy.""" + return stage_root / hashlib.sha256(str(src_root).encode()).hexdigest()[:16] + + +def _evict_unused(stage_root: Path, keep: set[Path], target_free: float) -> int: + """Delete staged episodes untouched for ``_STAGE_EVICT_AGE_S`` (older than + any job, so no running job reads them), least recently used first, until + ``target_free`` of the disk is free. Returns how many were removed.""" + cutoff = time.time() - _STAGE_EVICT_AGE_S + candidates = [] + for ns in stage_root.iterdir(): + if not ns.is_dir() or ns.name.startswith("."): + continue + for ep in ns.iterdir(): + if ep.name.startswith(".") or ep in keep: + continue + try: + mtime = ep.stat().st_mtime + except OSError: + continue + if mtime < cutoff: + candidates.append((mtime, ep)) + removed = 0 + for _, ep in sorted(candidates): + usage = shutil.disk_usage(stage_root) + if usage.free >= target_free * usage.total: + break + (ep.parent / ".staged" / f"{ep.name.removesuffix('.zarr')}.json").unlink( + missing_ok=True + ) + shutil.rmtree(ep, ignore_errors=True) + removed += 1 + return removed + + def stage_episodes( - src_root: Path, stage_root: Path, names: set[str], workers: int = 16 + src_root: Path, + stage_root: Path, + names: set[str], + versions: Mapping[str, str] | None = None, ) -> set[str]: - """Copy episodes from ``src_root`` into node-local ``stage_root`` and return - the names now loadable from it (those present in ``src_root``). Ranks on a - node serialize on a lock file, so the first one copies and the rest find - fresh copies. Raises ``StagingError`` if any present episode can't be staged.""" + """Copy episodes from ``src_root`` into ``staged_episode_dir(stage_root, + src_root)`` on node-local disk and return the names now loadable from it + (those present in ``src_root``). Raises ``StagingError`` if any present + episode can't be staged. + + ``/workspace`` NFS falls over under many small metadata ops, so the source + is touched as little as possible. Callers on a node serialize on a lock + file. An episode staged at the same ``versions[name]`` (e.g. the SQL + ``updated_at``), re-checked against the source within ``_STAGED_TRUST_S`` + and still intact locally, is used without reading the source. The rest are + checked against the source (~45 metadata RPCs each) one at a time, at most + once per launch, whose other ranks and heads reuse the result; stale ones + are copied in sequential tar streams of ``_STAGE_CHUNK`` episodes, each + committed before the next, and swapped in by rename so a killed restage + never leaves a half-deleted copy in place.""" from egomimic.rldb.zarr.norm_cache import episode_fingerprint + versions = versions or {} + ns = staged_episode_dir(stage_root, src_root) try: - stage_root.mkdir(parents=True, exist_ok=True) + ns.mkdir(parents=True, exist_ok=True) lock = open(stage_root / ".lock", "w") except OSError as e: raise StagingError( f"Stage dir {stage_root} is not writable ({e}); refusing to read " f"episodes from {src_root}. Set ${STAGE_DIR_ENV} to a node-local dir." ) from e + records = ns / ".staged" + incoming = stage_root / ".incoming" + trash = stage_root / ".trash" + + def trusted(name): + version = versions.get(name) + if version is None: + return False + rec = _read_json(records / f"{name}.json") + if ( + rec.get("version") != version + or time.time() - rec.get("checked_at", 0) > _STAGED_TRUST_S + ): + return False + dst = ns / rec.get("dir", "") + return dst.is_dir() and episode_fingerprint(dst) == rec.get("fp") + + def record(name, dst): + if versions.get(name) is None: + return + _write_json( + records / f"{name}.json", + { + "dir": dst.name, + "version": versions[name], + "fp": episode_fingerprint(dst), + "checked_at": time.time(), + }, + ) def check(name): src = _episode_dir(src_root, name) fp = None if src is None else episode_fingerprint(src) if fp is None: return None - dst = stage_root / src.name + dst = ns / src.name return name, src, dst, dst.is_dir() and fp == episode_fingerprint(dst) - def check_all(names): - with ThreadPoolExecutor(workers) as ex: - return [c for c in ex.map(check, sorted(names)) if c is not None] - - def copy(item): - name, src, dst, _ = item + def ensure_space(done, total, keep): + usage = shutil.disk_usage(stage_root) + if usage.free >= _STAGE_MIN_FREE_FRAC * usage.total: + return + evicted = _evict_unused(stage_root, keep, _STAGE_MIN_FREE_FRAC + 0.05) usage = shutil.disk_usage(stage_root) if usage.free < _STAGE_MIN_FREE_FRAC * usage.total: - return name, f"{stage_root} is below {_STAGE_MIN_FREE_FRAC:.0%} free" - tmp = stage_root / f".{src.name}.partial" - try: - shutil.rmtree(tmp, ignore_errors=True) - # copy2 keeps mtimes, so norm-stat cache keys match the source's. - shutil.copytree(src, tmp) - shutil.rmtree(dst, ignore_errors=True) - tmp.rename(dst) - except OSError as e: - shutil.rmtree(tmp, ignore_errors=True) - return name, str(e) - return name, None + raise StagingError( + f"{stage_root} is below {_STAGE_MIN_FREE_FRAC:.0%} free after " + f"{done}/{total} episodes (evicted {evicted} unused); refusing to " + f"read the rest from {src_root}." + ) - with lock: - checked = check_all(names) - staged = {c[0] for c in checked if c[3]} - stale = [c[0] for c in checked if not c[3]] - if stale: - fcntl.flock(lock, fcntl.LOCK_EX) - # Another rank may have copied these while this one waited. - todo = [c for c in check_all(stale) if not c[3]] - staged |= set(stale) - {c[0] for c in todo} - if todo: - logger.info("Staging %d episodes into %s", len(todo), stage_root) - t0 = time.time() - with ThreadPoolExecutor(workers) as ex: - results = list(ex.map(copy, todo)) - failed = [(n, err) for n, err in results if err is not None] - if failed: + def copy(todo, keep): + logger.info("Staging %d episodes into %s", len(todo), ns) + t0 = time.time() + for i in range(0, len(todo), _STAGE_CHUNK): + chunk = todo[i : i + _STAGE_CHUNK] + ensure_space(i, len(todo), keep) + try: + shutil.rmtree(incoming, ignore_errors=True) + incoming.mkdir() + _tar_copy(src_root, [c[1].name for c in chunk], incoming) + for name, src, dst, _ in chunk: + (records / f"{name}.json").unlink(missing_ok=True) + if dst.exists(): + trash.mkdir(exist_ok=True) + dst.rename(trash / f"{dst.name}.{time.time_ns()}") + (incoming / src.name).rename(dst) + record(name, dst) + except OSError as e: raise StagingError( - f"Failed to stage {len(failed)}/{len(todo)} episodes into " - f"{stage_root}; refusing to read them from {src_root}. " - f"First: {failed[0][0]}: {failed[0][1]}" - ) - staged |= {n for n, _ in results} + f"Failed to stage episodes {i}-{i + len(chunk)} of {len(todo)} " + f"into {ns}; refusing to read them from {src_root}: {e}" + ) from e + finally: + shutil.rmtree(incoming, ignore_errors=True) + shutil.rmtree(trash, ignore_errors=True) + logger.info( + "Staged %d episodes into %s in %.0fs", len(todo), ns, time.time() - t0 + ) + + with lock: + fcntl.flock(lock, fcntl.LOCK_EX) + # Left by a caller killed mid-copy; the lock makes them nobody's now. + shutil.rmtree(incoming, ignore_errors=True) + shutil.rmtree(trash, ignore_errors=True) + launch = _launch_record(stage_root, ".verified", src_root) + verified = _read_json(launch) + local_ok = {n: trusted(n) for n in sorted(names)} + unchecked = sorted( + n for n, ok in local_ok.items() if not ok and n not in verified + ) + if unchecked: + logger.info( + "Checking %d episodes against %s, one at a time", + len(unchecked), + src_root, + ) + checked = [c for c in map(check, unchecked) if c is not None] + todo = [c for c in checked if not c[3]] if todo: - logger.info( - "Staged %d episodes into %s in %.0fs", - len(todo), - stage_root, - time.time() - t0, - ) - logger.info("%d episodes load from %s", len(staged), stage_root) + copy(todo, keep={c[2] for c in checked}) + for name, _, dst, fresh in checked: + if fresh: + record(name, dst) + present = {c[0] for c in checked} + verified.update({n: n in present for n in unchecked}) + if launch is not None: + _write_json(launch, verified) + _prune(launch.parent, _LAUNCH_RECORD_TTL_S) + staged = { + n + for n in names + if local_ok[n] or (verified.get(n) and _episode_dir(ns, n) is not None) + } + for n in staged: + os.utime(_episode_dir(ns, n)) # recency for _evict_unused + logger.info( + "%d episodes load from %s (%d trusted from earlier staging, %d checked " + "against %s)", + len(staged), + ns, + sum(local_ok.values()), + len(unchecked), + src_root, + ) return staged @@ -510,15 +737,22 @@ def load(self, paths: list[tuple[str, str]]) -> dict[str, "ZarrDataset"]: """Datasets for ``paths``, first copied to a node-local stage directory (``$EGOVERSE_STAGE_DIR`` or ``DEFAULT_STAGE_DIR``) and read from there.""" valid = {h for _, h in paths} - stage_dir = os.environ.get(STAGE_DIR_ENV) or DEFAULT_STAGE_DIR - staged = stage_episodes(self.folder_path, Path(stage_dir), valid) + stage_dir = Path(os.environ.get(STAGE_DIR_ENV) or DEFAULT_STAGE_DIR) + staged = stage_episodes( + self.folder_path, stage_dir, valid, self._stage_versions(valid) + ) if not staged: return {} datasets = self._load_zarr_datasets( - search_path=Path(stage_dir), valid_folder_names=staged + search_path=staged_episode_dir(stage_dir, self.folder_path), + valid_folder_names=staged, ) return dict(sorted(datasets.items())) + def _stage_versions(self, names: set[str]) -> dict[str, str] | None: + """Per-episode version tokens that let staging trust an earlier copy.""" + return None + @classmethod def _episode_already_present(cls, local_dir: Path, episode_hash: str) -> bool: direct = local_dir / episode_hash @@ -552,6 +786,17 @@ def __init__( transform_list=transform_list, ) + def _stage_versions(self, names: set[str]) -> dict[str, str] | None: + df = _episode_table() + if "updated_at" not in df.columns: + return None + rows = df[df["episode_hash"].astype(str).isin(names)] + return { + str(h): str(t) + for h, t in zip(rows["episode_hash"], rows["updated_at"]) + if pd.notna(t) + } + # Deliberately keyed on "S3" rather than type(self).__name__: every # subclass shares this exact path resolution and differs only in load(). def _memo_fields(self) -> tuple: @@ -723,6 +968,40 @@ def _sync_s3_to_local( s3_paths: list[tuple[str, str]], local_dir: Path, numworkers: int = 10, + ): + """Sync ``s3_paths`` into ``local_dir``. Inside an ``srun`` step, once + per node per launch: callers serialize on a node-local lock, and the + node's other ranks and heads skip what an earlier caller of the launch + already synced instead of re-checking ``local_dir`` (NFS) episode by + episode. Elsewhere (e.g. a download script on the login node) it syncs + directly.""" + if not s3_paths: + return + stage_root = Path(os.environ.get(STAGE_DIR_ENV) or DEFAULT_STAGE_DIR) + record = _launch_record(stage_root, ".synced", local_dir) + try: + if record is None: + raise OSError("not inside an srun step") + stage_root.mkdir(parents=True, exist_ok=True) + lock = open(stage_root / ".sync.lock", "w") + except OSError: + cls._sync_s3_to_local_unlocked(bucket_name, s3_paths, local_dir, numworkers) + return + with lock: + fcntl.flock(lock, fcntl.LOCK_EX) + done = set(_read_json(record).get("hashes", [])) + pending = [(p, h) for p, h in s3_paths if h not in done] + cls._sync_s3_to_local_unlocked(bucket_name, pending, local_dir, numworkers) + _write_json(record, {"hashes": sorted(done | {h for _, h in s3_paths})}) + _prune(record.parent, _LAUNCH_RECORD_TTL_S) + + @classmethod + def _sync_s3_to_local_unlocked( + cls, + bucket_name: str, + s3_paths: list[tuple[str, str]], + local_dir: Path, + numworkers: int, ): if not s3_paths: return @@ -972,6 +1251,11 @@ def resolve( return kept +# ``{episode_hash: attrs}`` for a LocalEpisodeResolver folder; when present it +# replaces opening every episode to filter on its attrs. +LOCAL_EPISODE_INDEX = "_episode_index.json" + + class LocalEpisodeResolver(EpisodeResolver): """ Resolves episodes from local Zarr stores, filtering via local metadata. @@ -1002,18 +1286,26 @@ def _get_local_filtered_paths( rows: list[dict] = [] # one normalized attrs row per episode dir paths_by_hash: dict[str, str] = {} - for p in sorted(search_path.iterdir()): - if not p.is_dir(): - continue - episode_hash = p.name[:-5] if p.name.endswith(".zarr") else p.name - try: - store = zarr.open_group(str(p), mode="r") - metadata = dict(store.attrs) - except Exception as e: - logger.warning("Failed to read metadata for %s: %s", p, e) - continue - rows.append(_normalize_filter_row(metadata, episode_hash=episode_hash)) - paths_by_hash[episode_hash] = str(p) + index = search_path / LOCAL_EPISODE_INDEX + if index.is_file(): + # One read instead of opening every episode's attrs: on NFS that + # scan is thousands of small RPCs per rank per head. + for episode_hash, metadata in sorted(json.loads(index.read_text()).items()): + rows.append(_normalize_filter_row(metadata, episode_hash=episode_hash)) + paths_by_hash[episode_hash] = str(search_path / f"{episode_hash}.zarr") + else: + for p in sorted(search_path.iterdir()): + if not p.is_dir(): + continue + episode_hash = p.name[:-5] if p.name.endswith(".zarr") else p.name + try: + store = zarr.open_group(str(p), mode="r") + metadata = dict(store.attrs) + except Exception as e: + logger.warning("Failed to read metadata for %s: %s", p, e) + continue + rows.append(_normalize_filter_row(metadata, episode_hash=episode_hash)) + paths_by_hash[episode_hash] = str(p) if filters.episode_hashes: present = {r["episode_hash"]: r for r in rows} @@ -1070,6 +1362,19 @@ def _get_local_filtered_paths( def _memo_fields(self) -> tuple: return (str(self.folder_path), self.debug) + def _stage_versions(self, names: set[str]) -> dict[str, str] | None: + """A hash of each episode's index row, so a staged copy of an unchanged + episode is trusted without re-fingerprinting the NFS source.""" + index = Path(self.folder_path) / LOCAL_EPISODE_INDEX + if not index.is_file(): + return None + rows = json.loads(index.read_text()) + return { + n: hashlib.sha1(json.dumps(rows[n], sort_keys=True).encode()).hexdigest() + for n in names + if n in rows + } + def _compute_paths( self, filters: DatasetFilter, expected_embodiment: str | None ) -> list[tuple[str, str]]: @@ -1129,7 +1434,8 @@ class MultiDataset(torch.utils.data.Dataset): # Bounds-check slack per (timestep, channel) cell; see _check_bounds. BOUNDS_ABS_TOL = 1e-6 - BOUNDS_REL_SLACK = 0.5 + bounds_rel_slack = 0.5 + bounds_check = True def __init__( self, @@ -1141,6 +1447,8 @@ def __init__( valid_ratio: float = 0.2, norm_mode: str = "zscore", state: dict | None = None, + bounds_check: bool = True, + bounds_rel_slack: float | None = None, **kwargs, ): """ @@ -1151,8 +1459,18 @@ def __init__( valid_ratio: Train/valid split ratio. norm_mode: One of "zscore", "minmax", "quantile". state: If provided, populate stats fields from this dict (deploy mode). + bounds_check: reject samples outside the norm-stat quantile range + (NaN/Inf is rejected either way). + bounds_rel_slack: how far past each cell's quantile range, as a + fraction of it, a sample may go. Eval heads on other operators + widen it: their frames legitimately land just outside the train + range, and swapping them out biases the metric, while corrupt + values (fill constants, wrong frames) are orders of magnitude out. """ super().__init__() + self.bounds_check = bounds_check + if bounds_rel_slack is not None: + self.bounds_rel_slack = bounds_rel_slack # ---- Stats fields (always present, may be empty) ---- self.norm_mode = norm_mode @@ -1207,7 +1525,8 @@ def __init__( self._global_indices_by_dataset = {n: [] for n in self.datasets} for dataset_name, dataset in self.datasets.items(): - for local_idx in range(len(dataset)): + starts = getattr(dataset, "sample_indices", None) + for local_idx in range(len(dataset)) if starts is None else starts.tolist(): global_idx = len(self.index_map) self.index_map.append((dataset_name, local_idx)) self._global_indices_by_dataset[dataset_name].append(global_idx) @@ -1306,6 +1625,9 @@ def _check_bounds( logger.warning(prefix) return prefix + if not self.bounds_check: + continue + # Rotation channels are either Euler ypr (wraps at +-pi) or # continuous 6D columns; quantile bounds on them are meaningless # and reject otherwise-valid frames, so only the translation / @@ -1338,15 +1660,15 @@ def _check_bounds( # same way before comparing. Corrupt values # (fill constants, wrong-frame data) sit orders of magnitude # outside the range and are still caught. - tol = self.BOUNDS_ABS_TOL + self.BOUNDS_REL_SLACK * (q_high - q_low) + tol = self.BOUNDS_ABS_TOL + self.bounds_rel_slack * (q_high - q_low) below = arr_q < q_low - tol above = arr_q > q_high + tol # Cells normalize() treats as constant (wrist-frame t=0, the kp0 / # kp9 structural zeros) cannot blow up the loss whatever offset # they carry, so don't reject on them. Which stat decides that is # norm_mode's business -- std under zscore, the quantile range - # under quantile -- and every shipped config is zscore, so reading - # q99 - q1 here would exempt cells the normalizer still divides by. + # under quantile -- so reading q99 - q1 here regardless would + # exempt cells a zscore normalizer still divides by. constant = self._degenerate_cells(stats, arr.shape) if constant is not None: if check_idx is not None: @@ -1480,6 +1802,7 @@ def _from_resolver( skipped by the resolver, is a PinError rather than a silent drop.""" sync_from_s3 = kwargs.pop("sync_from_s3", False) filters = kwargs.pop("filters", None) + pose_glitch_mask = kwargs.pop("pose_glitch_mask", None) expected = None if dataset_name is not None: try: @@ -1505,6 +1828,25 @@ def _from_resolver( "were skipped by the resolver (see the log above)." ) + if pose_glitch_mask: + n_masked = n_total = n_eps = 0 + for ds in resolved.values(): + if isinstance(ds, ZarrDataset): + dropped_frames = ds.set_pose_glitch_mask(**dict(pose_glitch_mask)) + n_masked += dropped_frames + n_total += ds.total_frames + n_eps += dropped_frames > 0 + logger.info( + "pose_glitch_mask on '%s': %d of %d start frames dropped (%.1f%%) " + "in %d of %d episodes", + dataset_name or "", + n_masked, + n_total, + 100 * n_masked / max(1, n_total), + n_eps, + len(resolved), + ) + return cls(datasets=resolved, **kwargs) # ===================================================================== @@ -2095,6 +2437,9 @@ def __init__( ) torch.utils.data.Dataset.__init__(self) self.base = base + # MultiDataset's class defaults would otherwise shadow the base's values. + self.bounds_check = base.bounds_check + self.bounds_rel_slack = base.bounds_rel_slack self.frames_per_episode = frames_per_episode self.stride = stride self.datasets = base.datasets @@ -2148,6 +2493,17 @@ def __getattr__(self, item): return getattr(base, item) +def _unmasked(leaf): + """A shallow copy of ``leaf`` that starts at every frame: the replay video + assumes consecutive frames, so it plays through pose_glitch_mask gaps (the + metric loaders keep the mask).""" + if getattr(leaf, "sample_indices", None) is None: + return leaf + leaf = copy.copy(leaf) + leaf.sample_indices = None + return leaf + + def pinned_episode_subset( base: MultiDataset, episode_hashes, dataset_name: str | None = None ) -> MultiDataset: @@ -2182,9 +2538,11 @@ def pinned_episode_subset( f"'{dataset_name or ''}'" ) subset = MultiDataset( - datasets={h: base.datasets[h] for h in hashes}, + datasets={h: _unmasked(base.datasets[h]) for h in hashes}, mode="total", norm_mode=base.norm_mode, + bounds_check=base.bounds_check, + bounds_rel_slack=base.bounds_rel_slack, ) subset.set_norm_stats_from(base) return subset @@ -2216,8 +2574,55 @@ def __init__( self.key_map = key_map self.transform = transform_list + # Frames a sample may start at; None = all of them (set_pose_glitch_mask). + self.sample_indices: np.ndarray | None = None super().__init__() + def set_pose_glitch_mask( + self, + step_m: float = 0.05, + rot_deg: float = 10.0, + head_key: str = "obs_head_pose", + ) -> int: + """Drop the start frames whose non-image reads touch a camera-pose jump + (``pose_glitch_frames``): world-frame hand poses jump with the camera + there. Returns the number of start frames dropped.""" + T = self.total_frames + if head_key not in self.episode_reader._store or T < 2: + self.sample_indices = None + return 0 + bad = pose_glitch_frames( + np.asarray(self.episode_reader._store[head_key][:T]), step_m, rot_deg + ) + if bad.size == 0: + self.sample_indices = None + return 0 + lo, hi = self._read_span_offsets() + # Start frame t reads [t + lo, t + hi]; it is bad iff that span holds a + # bad frame b, i.e. t in [b - hi, b - lo]. + delta = np.zeros(T + 1, dtype=np.int64) + np.add.at(delta, np.clip(bad - hi, 0, T), 1) + np.add.at(delta, np.clip(bad - lo + 1, 0, T), -1) + masked = np.cumsum(delta[:T]) > 0 + self.sample_indices = np.flatnonzero(~masked) + return int(masked.sum()) + + def _read_span_offsets(self) -> tuple[int, int]: + """(lo, hi) frame offsets, relative to the start frame, of everything a + sample reads outside the image and annotation keys, measured mid-episode + so no window is clipped.""" + t0 = self.total_frames // 2 + frames = [t0] + for k, (_, window) in self._read_windows(t0).items(): + if self.key_map[k].get("key_type") == "camera_keys": + continue + if isinstance(window, np.ndarray): + frames.extend(int(i) for i in window) + else: + start, end = window + frames.extend([start, start if end is None else end - 1]) + return min(frames) - t0, max(frames) - t0 + def init_episode(self): """ inits the zarr episode and all the metadata associated, as well as total_frames for len @@ -2467,7 +2872,9 @@ def _next(reason: str, key: str = "") -> int: nonlocal attempts next_idx, attempts = get_fallback_idx( idx=idx, - candidates=range(self.total_frames), + candidates=range(self.total_frames) + if self.sample_indices is None + else self.sample_indices.tolist(), _attempts=attempts, max_attempts=self.total_frames, exhausted_error=( diff --git a/egomimic/scripts/precompute_norm_stats.py b/egomimic/scripts/precompute_norm_stats.py index 2c140cb9cc..c87bb35f53 100644 --- a/egomimic/scripts/precompute_norm_stats.py +++ b/egomimic/scripts/precompute_norm_stats.py @@ -42,7 +42,7 @@ import egomimic from egomimic.rldb.embodiment.embodiment import get_embodiment_id -from egomimic.rldb.zarr import episode_norm_samples, norm_cache +from egomimic.rldb.zarr import episode_norm_samples, history_stats, norm_cache from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset from egomimic.utils.env import load_env @@ -82,6 +82,26 @@ def main() -> None: action="store_true", help="only write the explicit file; skip the content-keyed cache entry", ) + ap.add_argument( + "--history-max-lag", + type=int, + default=0, + help="also write /norm_stats/history_stats.json: per-lag stats (lags " + "0..N frames) of the head-frame wrist translation, for " + "norm_stats.history_stats_path on a proprio-history run", + ) + ap.add_argument( + "--history-anchors", + type=int, + default=100000, + help="frames sampled for the history stats; lag L only counts those at " + "least L frames into their episode", + ) + ap.add_argument( + "--history-only", + action="store_true", + help="compute only the history stats (the current-step stats exist already)", + ) ap.add_argument( "overrides", nargs="*", @@ -89,6 +109,8 @@ def main() -> None: "paths.dataset_dir=/path/to/zarr/mirror)", ) args = ap.parse_args() + if args.history_only and args.history_max_lag <= 0: + ap.error("--history-only needs --history-max-lag") cfg_dir = os.path.join(os.path.dirname(egomimic.__file__), "hydra_configs") GlobalHydra.instance().clear() @@ -144,6 +166,35 @@ def main() -> None: # frame is exactly `[..., -1, :]` of the window, minus the K-fold # read. km["proprio_history"] = 1 + if args.history_max_lag > 0: + lag_inst = copy.deepcopy(inst) + lag_inst.resolver.key_map = history_stats.history_keymap( + km, args.history_max_lag + ) + lag_dataset = hydra.utils.instantiate(lag_inst, dataset_name=dataset_name) + t0 = time.perf_counter() + table = history_stats.compute( + norm_stats, + lag_dataset, + get_embodiment_id(dataset_name), + args.history_max_lag, + args.history_anchors, + seed=cfg.seed, + num_workers=args.num_workers, + ) + written = history_stats.save( + os.path.join(args.out, "norm_stats"), + get_embodiment_id(dataset_name), + table, + args.history_max_lag, + ) + print( + f"[precompute] {dataset_name}: history stats in " + f"{time.perf_counter() - t0:.1f}s -> {written}" + ) + if args.history_only: + continue + inst.resolver.key_map = km norm_dataset = hydra.utils.instantiate(inst, dataset_name=dataset_name) @@ -194,10 +245,17 @@ def main() -> None: ) print(f"[precompute] {dataset_name}: cache entry {written}") - norm_stats.cache_stats(save_cache_dir=args.out) out_dir = os.path.join(args.out, "norm_stats") + if args.history_only: + print( + f"\nDONE. Use this in training:\n norm_stats.history_stats_path={out_dir}" + ) + return + norm_stats.cache_stats(save_cache_dir=args.out) print("\nDONE. Use this in training:") print(f" norm_stats.precomputed_norm_path={out_dir}") + if args.history_max_lag > 0: + print(f" norm_stats.history_stats_path={out_dir}") if cache_dir: print( f"(or nothing: the content-keyed cache under {cache_dir} now holds " diff --git a/egomimic/trainHydra.py b/egomimic/trainHydra.py index ca4e688082..f05e56bc86 100644 --- a/egomimic/trainHydra.py +++ b/egomimic/trainHydra.py @@ -21,6 +21,7 @@ from egomimic.pl_utils.pl_model import ModelWrapper from egomimic.rldb.embodiment.embodiment import get_embodiment_id from egomimic.rldb.resolve_memo import resolve_once +from egomimic.rldb.zarr import history_stats from egomimic.rldb.zarr.utils import set_global_seed from egomimic.rldb.zarr.zarr_dataset_multi import ( EvenStrideDataset, @@ -173,6 +174,7 @@ def _log_dataset_frame_counts( train_datasets: dict, valid_datasets: dict, unseen_op_valid_datasets: dict | None = None, + extra_valid_datasets: dict | None = None, ) -> None: rows = [] for name, ds in train_datasets.items(): @@ -189,6 +191,8 @@ def _log_dataset_frame_counts( ) for name, ds in (unseen_op_valid_datasets or {}).items(): rows.append(("unseen_op_valid", name, len(ds))) + for name, ds in (extra_valid_datasets or {}).items(): + rows.append(("extra_valid", name, len(ds))) table = tabulate( rows, headers=["Split", "Dataset", "Frames"], @@ -309,11 +313,11 @@ def _build_train_viz_evaluator(cfg: DictConfig): return TrainVizEvalVideo(hydra.utils.instantiate(cfg.evaluator)) -def _unseen_op_valid_datasets(cfg: DictConfig, instantiate) -> dict: - """Datasets for the third (unseen_op_valid) val loader: a data config's own - ``unseen_op_valid_datasets``, in training runs with an evaluator only (eval - mode validates ``valid_datasets`` alone).""" - explicit = cfg.data.get("unseen_op_valid_datasets") +def _config_val_datasets(cfg: DictConfig, head: str, instantiate) -> dict: + """Datasets for a config-only val head (unseen_op_valid, extra_valid): the + data config's own ``_datasets``, in training runs with an evaluator + only (eval mode validates ``valid_datasets`` alone).""" + explicit = cfg.data.get(f"{head}_datasets") if ( explicit is None or _resolve_mode(cfg) != "train" @@ -510,14 +514,12 @@ def _require_capped_video_heads(model, datamodule) -> None: ) -def _build_unseen_op_valid_evaluator(cfg: DictConfig): +def _build_prefixed_evaluator(cfg: DictConfig, prefix: str): """The canonical evaluator (fresh instance, own frame buffers) wrapped to - log ``unseen_op_valid/`` and write ``videos_unseen_op_valid/``.""" + log ``/`` and write ``videos_/``.""" from egomimic.eval.eval_train_viz import TrainVizEvalVideo - return TrainVizEvalVideo( - hydra.utils.instantiate(cfg.evaluator), prefix="unseen_op_valid" - ) + return TrainVizEvalVideo(hydra.utils.instantiate(cfg.evaluator), prefix=prefix) @task_wrapper @@ -560,8 +562,11 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: train_viz_datasets, train_viz_params = _train_viz_datasets( cfg, train_datasets, instantiate=_instantiate_dataset ) - unseen_op_valid_datasets = _unseen_op_valid_datasets( - cfg, instantiate=_instantiate_dataset + unseen_op_valid_datasets = _config_val_datasets( + cfg, "unseen_op_valid", instantiate=_instantiate_dataset + ) + extra_valid_datasets = _config_val_datasets( + cfg, "extra_valid", instantiate=_instantiate_dataset ) # Split the val heads in two (lane A): a per-episode SUBSAMPLED metric @@ -575,6 +580,7 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: "valid": valid_datasets, "train_viz": train_viz_datasets, "unseen_op_valid": unseen_op_valid_datasets, + "extra_valid": extra_valid_datasets, }, ) valid_datasets = _subsample_val_datasets( @@ -597,6 +603,12 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: unseen_op_valid_datasets, cfg.data.get("unseen_op_valid_dataloader_params"), ) + extra_valid_datasets = _subsample_val_datasets( + cfg, + "extra_valid", + extra_valid_datasets, + cfg.data.get("extra_valid_dataloader_params"), + ) log.info(f"Instantiating datamodule <{cfg.data._target_}>") assert ( @@ -613,6 +625,7 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: # Always passed too: hydra would otherwise instantiate the raw # config nodes itself (eval mode included). datamodule_kwargs["unseen_op_valid_datasets"] = unseen_op_valid_datasets + datamodule_kwargs["extra_valid_datasets"] = extra_valid_datasets # Built above; hydra has nothing to instantiate for these. datamodule_kwargs["video_datasets"] = video_datasets datamodule: LightningDataModule = hydra.utils.instantiate( @@ -711,6 +724,25 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: norm_stats.norm_stats[emb], norm_stats._norm_run_metadata, ) + history_stats_path = OmegaConf.select( + cfg, "norm_stats.history_stats_path", default=None + ) + full_km = OmegaConf.to_container( + cfg.data.train_datasets[dataset_name].resolver.key_map, resolve=True + ) + history = int(full_km.get("proprio_history") or 1) + if history_stats_path and history > 1: + # After the content-keyed cache write, which holds current-step + # stats, and before save_cache_dir, so the run's own + # norm_stats.json carries the gathered (K, D) stats. + stride = history_stats.history_stride_frames(full_km, dataset) + changed = history_stats.apply( + norm_stats, emb, history_stats_path, history, stride + ) + log.info( + f"history stats for <{dataset_name}>: {changed} at " + f"K={history}, stride={stride} frames" + ) # Cache norm stats if save_cache_dir is set save_cache_dir = OmegaConf.select( cfg, "norm_stats.save_cache_dir", default=None @@ -730,6 +762,8 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: ds.set_norm_stats_from(norm_stats) for ds in getattr(datamodule, "unseen_op_valid_datasets", {}).values(): ds.set_norm_stats_from(norm_stats) + for ds in getattr(datamodule, "extra_valid_datasets", {}).values(): + ds.set_norm_stats_from(norm_stats) # The video subsets share leaves with their head's split but are separate # MultiDatasets, so they need their own wiring (EvenStrideDataset forwards # the call to the base it indexes into). @@ -751,6 +785,7 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: datamodule.train_datasets, datamodule.valid_datasets, datamodule.unseen_op_valid_datasets, + datamodule.extra_valid_datasets, ) log.info("Instantiating callbacks...") @@ -829,10 +864,17 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: train_viz_eval_obj.model = model.model model.train_viz_evaluator = train_viz_eval_obj if datamodule.unseen_op_valid_datasets: - unseen_eval_obj = _build_unseen_op_valid_evaluator(cfg) + unseen_eval_obj = _build_prefixed_evaluator(cfg, "unseen_op_valid") unseen_eval_obj.trainer = trainer unseen_eval_obj.model = model.model model.unseen_op_valid_evaluator = unseen_eval_obj + if datamodule.extra_valid_datasets: + extra_eval_obj = _build_prefixed_evaluator( + cfg, cfg.data.get("extra_valid_prefix") or "extra_valid" + ) + extra_eval_obj.trainer = trainer + extra_eval_obj.model = model.model + model.extra_valid_evaluator = extra_eval_obj _require_capped_video_heads(model, datamodule) model.val_loader_names = datamodule.val_loader_names() # Pre-fit baseline val. Skipped on requeues AND checkpoint resumes: diff --git a/egomimic/utils/git_info.py b/egomimic/utils/git_info.py index be2e11dd1a..2e3f75a839 100644 --- a/egomimic/utils/git_info.py +++ b/egomimic/utils/git_info.py @@ -9,6 +9,12 @@ # Any path inside the checkout; git walks up to the work tree root itself. _CWD = Path(__file__).resolve().parent +# stage-runtime.sh runs a node-local copy of the code, which has no .git; it +# records these facts about the checkout when it takes the copy. +_SNAPSHOT = _CWD.parents[1] / ".git_provenance" +# stage-runtime.sh runs the same commands when it records a snapshot. +DIFF_ARGS = ("diff", "HEAD", "--submodule=diff", "--no-ext-diff", "--no-color") +UNTRACKED_ARGS = ("ls-files", "--others", "--exclude-standard") # Bounds every git call (a large dirty checkout on loaded NFS can be slow). _TIMEOUT_S = 30.0 @@ -17,6 +23,17 @@ class GitError(RuntimeError): """git is missing, timed out, or exited non-zero.""" +def _snapshot(name: str) -> bytes | None: + """``name`` as recorded by stage-runtime.sh, or None when not running from a + staged copy. Raises GitError if git failed when the copy was taken.""" + if not _SNAPSHOT.is_dir(): + return None + err = _SNAPSHOT / f"{name}.error" + if err.exists(): + raise GitError(err.read_text(errors="replace").strip()) + return (_SNAPSHOT / name).read_bytes() + + def _git(*args: str) -> bytes: cmd = ["git", *args] try: @@ -32,7 +49,12 @@ def _git(*args: str) -> bytes: def git_sha() -> str | None: """HEAD commit, or None (no git, not inside a work tree, or git failed).""" try: - sha = _git("rev-parse", "HEAD").decode(errors="replace").strip() + raw = _snapshot("sha") + sha = ( + (raw if raw is not None else _git("rev-parse", "HEAD")) + .decode(errors="replace") + .strip() + ) except GitError: return None return sha if len(sha) == 40 else None @@ -42,10 +64,16 @@ def git_diff() -> bytes: """``git diff HEAD``: staged + unstaged tracked changes, including file diffs inside submodules (external/openpi is installed editable). ``b""`` when clean; raises GitError rather than passing a failure off as clean.""" - return _git("diff", "HEAD", "--submodule=diff", "--no-ext-diff", "--no-color") + raw = _snapshot("diff") + if raw is not None: + return raw + return _git(*DIFF_ARGS) def git_untracked() -> bytes: """Newline-separated paths of untracked, non-ignored files (names only; untracked files inside submodules are not listed). Raises GitError.""" - return _git("ls-files", "--others", "--exclude-standard") + raw = _snapshot("untracked") + if raw is not None: + return raw + return _git(*UNTRACKED_ARGS) diff --git a/egomimic/utils/scheduler_utils.py b/egomimic/utils/scheduler_utils.py index 5378909bf2..01f44743e0 100644 --- a/egomimic/utils/scheduler_utils.py +++ b/egomimic/utils/scheduler_utils.py @@ -52,3 +52,22 @@ def warmup_then_cosine( schedulers=[warmup, cosine], milestones=[warmup_epochs], ) + + +def warmup_stable_decay( + optimizer: torch.optim.Optimizer, + warmup: int, + stable: int, + decay: int, + warmup_start_factor: float = 1.0e-3, +) -> torch.optim.lr_scheduler.LambdaLR: + """Linear warmup, constant peak LR, then linear decay to 0.""" + + def factor(step: int) -> float: + if step < warmup: + return warmup_start_factor + (1 - warmup_start_factor) * step / warmup + if step < warmup + stable: + return 1.0 + return max(0.0, 1 - (step - warmup - stable) / decay) + + return torch.optim.lr_scheduler.LambdaLR(optimizer, factor) diff --git a/stage-runtime.sh b/stage-runtime.sh new file mode 100644 index 0000000000..9647176bd4 --- /dev/null +++ b/stage-runtime.sh @@ -0,0 +1,137 @@ +# Source (do not execute) from a Slurm batch script, after wt-env.sh: +# source ./stage-runtime.sh || exit 1 +# Copies the venv, its uv Python and this tree's code to node-local /ephemeral on +# every node of the job, then exports EGOVERSE_PYTHON and EGOVERSE_CODE: +# "$EGOVERSE_PYTHON" "$EGOVERSE_CODE/egomimic/trainHydra.py" ... +# Imports are thousands of small-file lookups per rank; on /workspace NFS they +# are what livelocks nodes. The venv is copied once per node and reused while +# unchanged. Each sourcing snapshots the code into one tarball on NFS (with its +# git provenance) that every node unpacks; a Slurm requeue reuses it, so it +# reruns the same code. All copies are single sequential tar streams. The cwd +# stays in the tree, so run dirs, logs and checkpoints still land on /workspace. +# +# The work runs in child processes (--snapshot, then --copy on every node), all +# reading the _RT_* paths computed once here: errexit is ignored in a function +# or subshell on the left of ||, which would let a failed copy be marked complete. + +_rt_copy() { + mkdir -p "$_RT_ROOT" + exec 9>"$_RT_ROOT/.lock" + flock 9 + # Under the lock, any .partial is left from a killed or failed copy. + rm -rf "$_RT_ROOT"/env-*.partial "$_RT_ROOT"/code-*.partial + if [ ! -e "$_RT_ENV/.complete" ]; then + rm -rf "$_RT_ENV" + mkdir -p "$_RT_ENV.partial" + tar -C / -cf - "${_RT_VENV#/}" "${_RT_PYHOME#/}" | tar -C "$_RT_ENV.partial" -xf - + local bin="$_RT_ENV.partial$_RT_VENV/bin" + ln -sfn "$_RT_ENV$_RT_PYTHON" "$bin/python" + sed -i "s#^home = .*#home = $_RT_ENV$_RT_PYHOME/bin#" "$_RT_ENV.partial$_RT_VENV/pyvenv.cfg" + # Console scripts (torchrun, wandb, ...) would otherwise start the NFS python. + { grep -lZ -m1 "^#!$_RT_VENV/bin/python" "$bin"/* 2>/dev/null || true; } \ + | xargs -0 -r sed -i "1s|^#!$_RT_VENV/bin/python|#!$_RT_ENV$_RT_VENV/bin/python|" + mv "$_RT_ENV.partial" "$_RT_ENV" + if ! "$_RT_ENV$_RT_VENV/bin/python" -c "import sys, torch; assert sys.prefix == '$_RT_ENV$_RT_VENV', sys.prefix"; then + rm -rf "$_RT_ENV" + echo "stage-runtime: the staged interpreter at $_RT_ENV does not run" >&2 + return 1 + fi + touch "$_RT_ENV/.complete" + fi + touch "$_RT_ENV" + if [ ! -d "$_RT_CODE" ]; then + mkdir -p "$_RT_CODE.partial" + tar -C "$_RT_CODE.partial" -xf "$_RT_TARBALL" + mv "$_RT_CODE.partial" "$_RT_CODE" + fi + # Jobs share nodes, so only drop copies no queued or running job of ours can use: + # code dirs are named code--..., and an env is touched by every job that starts on it. + local live d + if live="$(squeue -h -u "$(id -un)" -o %A)"; then + for d in "$_RT_ROOT"/code-*; do + [ -d "$d" ] || continue + grep -qx "$(basename "$d" | cut -d- -f2)" <<<"$live" || rm -rf "$d" + done + fi + ls -1dt "$_RT_ROOT"/env-* | tail -n +3 | while read -r d; do + [ -n "$(find "$d" -maxdepth 0 -mtime +2)" ] && rm -rf "$d" + done || true +} + +_rt_snapshot() { + local tmp prov + tmp="$(mktemp -d)" + prov="$tmp/.git_provenance" + mkdir "$prov" + # The same commands as egomimic/utils/git_info.py (DIFF_ARGS, UNTRACKED_ARGS). + git -C "$_RT_TREE" rev-parse HEAD >"$prov/sha" 2>"$prov/sha.error" && rm "$prov/sha.error" || rm -f "$prov/sha" + git -C "$_RT_TREE" diff HEAD --submodule=diff --no-ext-diff --no-color >"$prov/diff" 2>"$prov/diff.error" && rm "$prov/diff.error" || rm -f "$prov/diff" + git -C "$_RT_TREE" ls-files --others --exclude-standard >"$prov/untracked" 2>"$prov/untracked.error" && rm "$prov/untracked.error" || rm -f "$prov/untracked" + mkdir -p "$(dirname "$_RT_TARBALL")" + tar -cf "$_RT_TARBALL.partial" -C "$_RT_TREE" --exclude=__pycache__ egomimic external/openpi/src -C "$tmp" .git_provenance + mv "$_RT_TARBALL.partial" "$_RT_TARBALL" + rm -rf "$tmp" + # Generous: a requeued job may have waited in the queue a long time. + find "$(dirname "$_RT_TARBALL")" -maxdepth 1 -name 'code-*.tar*' -mtime +30 -delete 2>/dev/null || true +} + +case "${1:-}" in + --copy) set -euo pipefail; _rt_copy; exit ;; + --snapshot) set -euo pipefail; _rt_snapshot; exit ;; +esac +if [ "${BASH_SOURCE[0]}" = "$0" ]; then + echo "stage-runtime.sh must be sourced: source ./stage-runtime.sh" >&2 + exit 1 +fi + +_rt_main() { + if [ -z "${SLURM_JOB_ID:-}" ]; then + echo "stage-runtime.sh: not in a Slurm job; nothing staged" >&2 + return 1 + fi + local tree="${EGOVERSE_ROOT:-}" venv="${VIRTUAL_ENV:-}" cache="${EGOVERSE_CACHE_DIR:-}" + if [ -z "$tree" ] || [ -z "$venv" ] || [ -z "$cache" ]; then + echo "stage-runtime.sh: source wt-env.sh first" >&2 + return 1 + fi + local python pyhome key + python="$(readlink -f "$venv/bin/python")" || true + pyhome="$(dirname "$(dirname "$python")")" + # A dangling link would make pyhome "." and the copy tar the whole filesystem. + if [ ! -x "$python" ] || [ "${pyhome#/}" = "$pyhome" ] || [ "$pyhome" = / ]; then + echo "stage-runtime.sh: $venv/bin/python does not resolve to an interpreter" >&2 + return 1 + fi + # Directory mtimes change when packages are added or removed, not when a + # file inside one is edited in place. + key="$(set -o pipefail; stat -c '%n %Y' "$venv/pyvenv.cfg" "$venv"/lib/python3*/site-packages "$pyhome" | sha256sum | cut -c1-16)" || return 1 + export _RT_ROOT="/ephemeral/loaner-jobs/$(id -u)/runtime" + export _RT_TREE="$tree" _RT_VENV="$venv" _RT_PYTHON="$python" _RT_PYHOME="$pyhome" + export _RT_ENV="$_RT_ROOT/env-$key" + export _RT_TARBALL="$cache/runtime/code-$SLURM_JOB_ID-$(printf %s "$tree" | sha256sum | cut -c1-12).tar" + if [ "${SLURM_RESTART_COUNT:-0}" -gt 0 ] && [ -e "$_RT_TARBALL" ]; then + echo "stage-runtime: requeue, reusing $_RT_TARBALL" + else + bash "$tree/stage-runtime.sh" --snapshot || return 1 + fi + # A fresh snapshot unpacks to a fresh dir; a requeue finds its old one. + export _RT_CODE="$_RT_ROOT/$(basename "$_RT_TARBALL" .tar)-$(stat -c %Y "$_RT_TARBALL")" + srun --nodes="${SLURM_JOB_NUM_NODES:-1}" --ntasks-per-node=1 --ntasks="${SLURM_JOB_NUM_NODES:-1}" \ + -c 1 --mem=0 --gres=none --overlap \ + bash "$tree/stage-runtime.sh" --copy || return 1 + export EGOVERSE_PYTHON="$_RT_ENV$_RT_VENV/bin/python" + export EGOVERSE_CODE="$_RT_CODE" + export PYTHONPATH="$_RT_CODE:$_RT_CODE/external/openpi/src" + export PATH="$_RT_ENV$_RT_VENV/bin:$PATH" + # torch.compile/triton caches are many small files; keep them off NFS too. + export TRITON_CACHE_DIR="$_RT_ROOT/triton" + echo "stage-runtime: $_RT_ENV, $_RT_CODE" +} + +unset EGOVERSE_PYTHON EGOVERSE_CODE +_rt_main +_rt_status=$? +[ "$_rt_status" = 0 ] || echo "stage-runtime: staging failed" >&2 +unset _RT_ROOT _RT_TREE _RT_VENV _RT_PYTHON _RT_PYHOME _RT_ENV _RT_TARBALL _RT_CODE +unset -f _rt_main _rt_copy _rt_snapshot +return $_rt_status diff --git a/tests/unit/test_eval_loaders.py b/tests/unit/test_eval_loaders.py index 1d8b92f22f..4253fb5e3f 100644 --- a/tests/unit/test_eval_loaders.py +++ b/tests/unit/test_eval_loaders.py @@ -246,6 +246,18 @@ def test_config_without_the_keys_is_unchanged(compose_resolve): assert dm.val_loader_names() == ["valid", "train_viz", "unseen_op_valid"] assert dm.valid_datasets[HUMAN] is split + dm = MultiDataModuleWrapper( + train_datasets=datasets, + valid_datasets=datasets, + train_dataloader_params=params, + valid_dataloader_params=params, + extra_valid_datasets=datasets, + extra_valid_dataloader_params=params, + extra_valid_prefix="seen_ops_valid", + ) + assert dm.val_loader_names() == ["valid", "extra_valid"] + assert len(dm.val_dataloader()) == 2 + # ------------------------------------------------------------ 2. metric loaders def test_metric_loader_is_even_stride_over_every_episode(): @@ -402,6 +414,17 @@ def test_video_dataset_holds_only_the_pins_in_frame_order(): pinned_episode_subset(split, ["nope"], dataset_name=HUMAN) +def test_video_dataset_keeps_the_heads_bounds_settings(): + split = _split({"ep0": 30, "ep1": 20}) + assert pinned_episode_subset(split, ["ep0"]).bounds_check + split.bounds_check = False + assert not pinned_episode_subset(split, ["ep0"]).bounds_check + split.bounds_rel_slack = 10.0 + assert pinned_episode_subset(split, ["ep0"]).bounds_rel_slack == 10.0 + wrapped = EvenStrideDataset(split, frames_per_episode=5) + assert wrapped.bounds_rel_slack == 10.0 and not wrapped.bounds_check + + def test_video_datasets_are_built_per_head_from_the_resolved_split(): valid = _split({"a": 10, "b": 10}) unseen = _split({"c": 10, "d": 10}) diff --git a/tests/unit/test_git_provenance.py b/tests/unit/test_git_provenance.py index 9e7f8d0f3f..13890b55ab 100644 --- a/tests/unit/test_git_provenance.py +++ b/tests/unit/test_git_provenance.py @@ -34,6 +34,21 @@ def test_failed_git_diff_raises_instead_of_reading_as_clean(tmp_path, monkeypatc git_info.git_diff() +def test_staged_copy_reads_the_recorded_snapshot(tmp_path, monkeypatch): + """A node-local code copy has no .git; stage-runtime.sh records the facts.""" + snap = tmp_path / ".git_provenance" + snap.mkdir() + (snap / "sha").write_bytes(b"a" * 40 + b"\n") + (snap / "diff").write_bytes(b"diff --git a/x b/x\n\xff\n") + (snap / "untracked.error").write_text("git ls-files exited 128") + monkeypatch.setattr(git_info, "_SNAPSHOT", snap) + monkeypatch.setattr(git_info, "_CWD", tmp_path) # not a work tree + assert git_info.git_sha() == "a" * 40 + assert git_info.git_diff() == b"diff --git a/x b/x\n\xff\n" + with pytest.raises(GitError, match="exited 128"): + git_info.git_untracked() + + def _trainer(tmp_path, rank0=True): return SimpleNamespace(default_root_dir=str(tmp_path), is_global_zero=rank0) @@ -260,3 +275,9 @@ def test_callback_writes_on_validate_only_runs(tmp_path, monkeypatch): ) trainer.validate(_Tiny(), dataloaders=_loader()) assert (tmp_path / "git_sha.txt").read_text() == "d" * 40 + "\n" + + +def test_stage_runtime_records_what_git_info_would_run(): + script = (Path(git_info.__file__).parents[2] / "stage-runtime.sh").read_text() + for args in (("rev-parse", "HEAD"), git_info.DIFF_ARGS, git_info.UNTRACKED_ARGS): + assert f'git -C "$_RT_TREE" {" ".join(args)} ' in script diff --git a/tests/unit/test_history_stats.py b/tests/unit/test_history_stats.py new file mode 100644 index 0000000000..6f0b2e4b7a --- /dev/null +++ b/tests/unit/test_history_stats.py @@ -0,0 +1,204 @@ +"""History norm stats (egomimic/rldb/zarr/history_stats.py).""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch + +from egomimic.rldb.embodiment.embodiment import get_embodiment_id +from egomimic.rldb.embodiment.human import Human +from egomimic.rldb.zarr import history_stats +from egomimic.rldb.zarr.zarr_dataset_multi import LocalEpisodeResolver, MultiDataset +from tests.unit.test_proprio_history import _write_varying_episode + +STATE = "observations.state.ee_pose" +EMB = get_embodiment_id("human_bimanual") +STATS = ("mean", "std", "min", "max", "median", "quantile_1", "quantile_99") + + +def _base(width=20): + return {s: np.full(width, 100.0 + i, dtype=np.float32) for i, s in enumerate(STATS)} + + +def _entry(max_lag=20, channels=(0, 1, 2)): + lags = np.arange(max_lag + 1, dtype=np.float32) + return { + "zarr_key": STATE, + "channels": list(channels), + "count": np.full(max_lag + 1, 7), + "stats": { + s: np.tile(lags[:, None] + 1000 * i, (1, len(channels))) + for i, s in enumerate(STATS) + }, + } + + +def test_gather_tiles_the_base_and_takes_each_steps_lag_row(): + out = history_stats.gather(_base(), _entry(), history=3, stride=5) + assert out["mean"].shape == (3, 20) + # oldest first: lags 10, 5, 0 + np.testing.assert_array_equal(out["mean"][:, 0], [10.0, 5.0, 0.0]) + np.testing.assert_array_equal(out["quantile_99"][:, 2], [6010.0, 6005.0, 6000.0]) + # every other channel is the current-step stat at every step + np.testing.assert_array_equal(out["mean"][:, 3:], np.full((3, 17), 100.0)) + + +def test_gather_refuses_a_history_past_the_table(): + with pytest.raises(ValueError, match="max_lag 20"): + history_stats.gather(_base(), _entry(max_lag=20), history=5, stride=6) + + +def test_gather_keeps_stats_the_table_does_not_hold(): + base = {**_base(), "quantile_0_01": np.full(20, -5.0, dtype=np.float32)} + entry = _entry() + del entry["stats"]["min"] + out = history_stats.gather(base, entry, history=2, stride=1) + np.testing.assert_array_equal(out["min"], np.full((2, 20), 102.0)) + np.testing.assert_array_equal(out["quantile_0_01"], np.full((2, 20), -5.0)) + + +def test_save_load_round_trip(tmp_path): + path = history_stats.save(str(tmp_path / "norm_stats"), EMB, {"s": _entry()}, 20) + assert path.endswith(history_stats.HISTORY_STATS_FILE) + table, max_lag = history_stats.load(str(tmp_path / "norm_stats"), EMB) + assert max_lag == 20 + np.testing.assert_array_equal( + table["s"]["stats"]["mean"], _entry()["stats"]["mean"] + ) + assert table["s"]["channels"] == [0, 1, 2] + with pytest.raises(ValueError, match="max_lag"): + history_stats.save(str(tmp_path / "norm_stats"), EMB, {"s": _entry(10)}, 10) + + +class _Shell: + def __init__(self, stats): + self.norm_stats = {EMB: stats} + + +def test_apply_gathers_once_and_leaves_gathered_stats_alone(tmp_path): + history_stats.save(str(tmp_path), EMB, {"s": _entry()}, 20) + shell = _Shell({"s": _base(), "other": _base()}) + assert history_stats.apply(shell, EMB, str(tmp_path), history=3, stride=5) == ["s"] + assert shell.norm_stats[EMB]["s"]["mean"].shape == (3, 20) + assert shell.norm_stats[EMB]["other"]["mean"].shape == (20,) + # a run's saved norm_stats.json already holds (K, D): reloading it is a no-op + gathered = shell.norm_stats[EMB]["s"] + assert history_stats.apply(shell, EMB, str(tmp_path), history=3, stride=5) == [] + assert shell.norm_stats[EMB]["s"] is gathered + with pytest.raises(ValueError, match="4-step"): + history_stats.apply(shell, EMB, str(tmp_path), history=4, stride=5) + + +def test_lag_channels_are_the_wrist_translation_only(): + assert history_stats.lag_channels(STATE, 20) == [0, 1, 2, 10, 11, 12] + assert history_stats.lag_channels("observations.state.keypoints", 144) == [ + 0, + 1, + 2, + 72, + 73, + 74, + ] + assert history_stats.lag_channels("actions_keypoints", 144) is None + assert history_stats.lag_channels(STATE, 7) is None + + +# --------------------------------------------------------------------------- +# compute() on a real dataset + transform pipeline +# --------------------------------------------------------------------------- + +TRANSFORMS = dict( + mode="cartesian_wristframe_6d", + stride=1, + fix_left_wrist_convention=True, + pad_proprio_gripper=True, +) + + +def _md(tmp_path, **km): + key_map = Human.get_keymap(keymap_mode="cartesian", norm_mode=True, **km) + tl = Human.get_transform_list(**TRANSFORMS) + resolver = LocalEpisodeResolver(tmp_path, key_map=key_map, transform_list=tl) + return MultiDataset._from_resolver(resolver, mode="total") + + +@pytest.fixture +def episodes(tmp_path): + for seed in range(2): + _write_varying_episode(tmp_path, seed=seed) + return tmp_path + + +def test_compute_matches_a_direct_per_lag_computation(episodes): + max_lag = 6 + base = _md(episodes) + norm = MultiDataset(state={}, norm_mode="quantile") + norm.populate_from_datasets({"human_bimanual": base}) + lag_ds = _md(episodes, **history_stats.history_keymap({}, max_lag)) + table = history_stats.compute( + norm, lag_ds, EMB, max_lag, n_anchors=10**6, num_workers=0 + ) + ((name, entry),) = table.items() + assert entry["zarr_key"] == STATE + ch = entry["channels"] + assert ch == [0, 1, 2, 10, 11, 12] + + # every frame is an anchor; lag L is real for frames >= L of each episode + lens = [len(leaf) for leaf in base.datasets.values()] + for lag in (0, 3, max_lag): + assert entry["count"][lag] == sum(n - lag for n in lens) + + # lag 0 is the current-step value + cur = np.stack([np.asarray(base[i][STATE])[ch] for i in range(len(base))]) + np.testing.assert_allclose(entry["stats"]["mean"][0], cur.mean(0), atol=1e-5) + + # lag L: the past wrist in the anchor's head frame, from a K = L + 1 read + lag = 4 + hist = _md(episodes, proprio_history=lag + 1, history_stride=1) + rows = [] + for i in range(len(hist)): + s = hist[i] + if np.asarray(s["proprio_history_mask"])[0] > 0: + rows.append(np.asarray(s[STATE])[0, ch]) + rows = np.stack(rows) + np.testing.assert_allclose(entry["stats"]["mean"][lag], rows.mean(0), atol=1e-5) + np.testing.assert_allclose(entry["stats"]["max"][lag], rows.max(0), atol=1e-5) + # the head moves, so the past wrist is not distributed like the current one + assert not np.allclose(entry["stats"]["mean"][lag], entry["stats"]["mean"][0]) + + +def test_gathered_stats_normalize_a_history_sample(episodes, tmp_path_factory): + max_lag, K, stride = 6, 3, 2 + base = _md(episodes) + norm = MultiDataset(state={}, norm_mode="quantile") + norm.populate_from_datasets({"human_bimanual": base}) + norm.infer_norm_from_dataset(base, "human_bimanual", sample_frac=1.0, num_workers=0) + lag_ds = _md(episodes, **history_stats.history_keymap({}, max_lag)) + out = tmp_path_factory.mktemp("stats") + history_stats.save( + str(out), + EMB, + history_stats.compute( + norm, lag_ds, EMB, max_lag, n_anchors=10**6, num_workers=0 + ), + max_lag, + ) + name = norm.zarr_key_to_keyname(STATE, EMB) + assert history_stats.apply(norm, EMB, str(out), K, stride) == [name] + + hist = _md(episodes, proprio_history=K, history_stride=stride) + hist.set_norm_stats_from(norm) + sample = hist[len(hist) // 2] + assert tuple(sample[STATE].shape) == (K, 20) + assert torch.isfinite(torch.as_tensor(sample[STATE])).all() + + +def test_history_stride_frames(): + assert history_stats.history_stride_frames({"history_stride": 5}, None) == 5 + assert history_stats.history_stride_frames({}, None) == 1 + + +def test_history_stride_seconds_resolve_through_each_episodes_fps(episodes): + ds = _md(episodes) + assert history_stats.history_stride_frames({"history_stride_s": 0.5}, ds) == 15 diff --git a/tests/unit/test_image_history.py b/tests/unit/test_image_history.py index 3bdaa6da67..2f75b244c5 100644 --- a/tests/unit/test_image_history.py +++ b/tests/unit/test_image_history.py @@ -100,3 +100,26 @@ def spy(self, intervals): sample = leaf[idx] assert rows[-1]["images.front_1"] == 2 assert torch.equal(sample[PAST], leaf[idx - round(gap_s * 30)][CAM]) + + +def test_image_history_reads_the_past_frames_gap_apart_oldest_first(tmp_path): + leaf = _leaf(tmp_path, 30, image_history_gap_s=0.1, image_history=4) + idx = 20 + sample = leaf[idx] + assert sample[PAST].shape == (3, 3, 32, 32) + for i, lag in enumerate((9, 6, 3)): + assert torch.equal(sample[PAST][i], leaf[idx - lag][CAM]) + assert sample[f"{PAST}_mask"].tolist() == [1.0, 1.0, 1.0] + + +def test_image_history_pads_before_the_episode_start(tmp_path): + leaf = _leaf(tmp_path, 30, image_history_gap_s=0.1, image_history=4) + sample = leaf[4] + assert sample[f"{PAST}_mask"].tolist() == [0.0, 0.0, 1.0] + for i in range(3): + assert torch.equal(sample[PAST][i], leaf[1][CAM]) + + +def test_image_history_below_two_frames_is_rejected(): + with pytest.raises(ValueError, match="image_history"): + Human.get_keymap("cartesian", image_history_gap_s=0.1, image_history=1) diff --git a/tests/unit/test_infra_port_wiring.py b/tests/unit/test_infra_port_wiring.py index a150f562a7..e21b7b06ba 100644 --- a/tests/unit/test_infra_port_wiring.py +++ b/tests/unit/test_infra_port_wiring.py @@ -538,6 +538,13 @@ def check(v): assert check(1.0 + 0.6 * 2.0) is not None assert check(1e9) is not None + md.bounds_rel_slack = 10.0 + assert check(1.0 + 5.0 * 2.0) is None + assert check(1e9) is not None + md.bounds_check = False + assert check(1e9) is None + assert check(float("nan")) is not None + # ------------------------------------------------------ train_viz on by default def _viz_cfg(**top): diff --git a/tests/unit/test_local_episode_index.py b/tests/unit/test_local_episode_index.py new file mode 100644 index 0000000000..53aa59d621 --- /dev/null +++ b/tests/unit/test_local_episode_index.py @@ -0,0 +1,57 @@ +"""``LocalEpisodeResolver`` filters from ``_episode_index.json`` when the folder +has one, without opening the episodes.""" + +from __future__ import annotations + +import json +from unittest import mock + +from egomimic.rldb.filters import DatasetFilter +from egomimic.rldb.zarr.zarr_dataset_multi import ( + LOCAL_EPISODE_INDEX, + LocalEpisodeResolver, +) + + +def test_index_rows_drive_the_filter_and_no_episode_is_opened(tmp_path): + index = { + "a": {"embodiment": "human_bimanual", "task": "fold_tshirt", "user_id": "u1"}, + "b": {"embodiment": "human_bimanual", "task": "fold_tshirt", "user_id": "u2"}, + "c": {"embodiment": "human_bimanual", "task": "hang_clothes", "user_id": "u1"}, + } + (tmp_path / LOCAL_EPISODE_INDEX).write_text(json.dumps(index)) + filters = DatasetFilter( + filter_lambdas=[ + "lambda row: row.get('task') == 'fold_tshirt'", + "lambda row: row.get('user_id') == 'u1'", + ] + ) + with mock.patch("zarr.open_group", side_effect=AssertionError("opened")): + paths = LocalEpisodeResolver._get_local_filtered_paths( + tmp_path, filters, expected_embodiment="human_bimanual" + ) + assert paths == [(str(tmp_path / "a.zarr"), "a")] + + +def test_pins_are_checked_against_the_index(tmp_path): + (tmp_path / LOCAL_EPISODE_INDEX).write_text( + json.dumps({"a": {"embodiment": "human_bimanual"}}) + ) + paths = LocalEpisodeResolver._get_local_filtered_paths( + tmp_path, + DatasetFilter(episode_hashes=["a"]), + expected_embodiment="human_bimanual", + ) + assert paths == [(str(tmp_path / "a.zarr"), "a")] + + +def test_stage_versions_come_from_the_index_rows(tmp_path): + rows = {"a": {"total_frames": 10}, "b": {"total_frames": 20}} + (tmp_path / LOCAL_EPISODE_INDEX).write_text(json.dumps(rows)) + resolver = LocalEpisodeResolver(tmp_path) + v = resolver._stage_versions({"a", "b", "missing"}) + assert set(v) == {"a", "b"} and v["a"] != v["b"] + assert resolver._stage_versions({"a"}) == {"a": v["a"]} + rows["a"]["total_frames"] = 11 + (tmp_path / LOCAL_EPISODE_INDEX).write_text(json.dumps(rows)) + assert resolver._stage_versions({"a"})["a"] != v["a"] diff --git a/tests/unit/test_pose_glitch_mask.py b/tests/unit/test_pose_glitch_mask.py new file mode 100644 index 0000000000..662d3cb09c --- /dev/null +++ b/tests/unit/test_pose_glitch_mask.py @@ -0,0 +1,99 @@ +"""``pose_glitch_mask``: start frames whose non-image reads touch a camera-pose +jump are left out of ``MultiDataset``'s index, everything else is kept.""" + +from __future__ import annotations + +import numpy as np + +from egomimic.rldb.embodiment.human import Human +from egomimic.rldb.zarr.zarr_dataset_multi import ( + EvenStrideDataset, + LocalEpisodeResolver, + MultiDataset, + pose_glitch_frames, +) +from egomimic.rldb.zarr.zarr_writer import ZarrWriter + +T = 120 + + +def _pose(jumps=(), turn_at=None): + t = np.arange(T, dtype=np.float64)[:, None] + xyz = t * 0.001 + np.zeros((T, 3)) + for f in jumps: + xyz[f:, 0] += 0.2 + quat = np.tile([[1.0, 0.0, 0.0, 0.0]], (T, 1)) + if turn_at is not None: + half = np.radians(30) / 2 + quat[turn_at:] = [np.cos(half), 0.0, 0.0, np.sin(half)] + return np.concatenate([xyz, quat], axis=1) + + +def _dataset(tmp_path, head_pose, mask=True, **keymap_kwargs): + rng = np.random.default_rng(0) + ee = _pose() + ZarrWriter.create_and_write( + tmp_path / "ep.zarr", + numeric_data={ + "left.obs_ee_pose": ee, + "right.obs_ee_pose": ee, + "obs_head_pose": head_pose, + }, + image_data={"images.front_1": rng.integers(0, 255, (T, 16, 16, 3), np.uint8)}, + embodiment="human_bimanual", + fps=30, + task_name="synthetic", + intrinsics={"front_1": np.eye(3, 4)}, + ) + key_map = Human.get_keymap(keymap_mode="cartesian", **keymap_kwargs) + resolver = LocalEpisodeResolver(tmp_path, key_map=key_map, transform_list=[]) + return MultiDataset._from_resolver( + resolver, + mode="total", + pose_glitch_mask={"step_m": 0.05, "rot_deg": 10.0} if mask else None, + ) + + +def test_glitch_frames_are_both_sides_of_a_jump_or_turn(): + assert pose_glitch_frames(_pose(jumps=[40])).tolist() == [39, 40] + assert pose_glitch_frames(_pose(turn_at=70)).tolist() == [69, 70] + assert pose_glitch_frames(_pose()).size == 0 + + +def test_masked_starts_are_exactly_those_whose_reads_touch_a_glitch(tmp_path): + ds = _dataset(tmp_path, _pose(jumps=[60]), proprio_history=4, history_stride=2) + leaf = next(iter(ds.datasets.values())) + lo, hi = leaf._read_span_offsets() + assert lo == -6 and hi > 0 + expected = [t for t in range(T) if not (t + lo <= 60 and 59 <= t + hi)] + assert leaf.sample_indices.tolist() == expected + assert [i for _, i in ds.index_map] == expected + assert len(ds) == len(expected) < T + + +def test_clean_episode_and_mask_off_keep_every_frame(tmp_path): + clean = _dataset(tmp_path / "a", _pose()) + assert ( + len(clean) == T and next(iter(clean.datasets.values())).sample_indices is None + ) + off = _dataset(tmp_path / "b", _pose(jumps=[60]), mask=False) + assert len(off) == T + + +def test_even_stride_draws_only_kept_starts(tmp_path): + ds = _dataset(tmp_path, _pose(jumps=[30, 90])) + kept = set(next(iter(ds.datasets.values())).sample_indices.tolist()) + strided = EvenStrideDataset(ds, frames_per_episode=10) + for i in range(len(strided)): + _, local = ds.index_map[strided.indices[i]] + assert local in kept + + +def test_video_subset_plays_through_the_mask_but_metrics_keep_it(tmp_path): + from egomimic.rldb.zarr.zarr_dataset_multi import pinned_episode_subset + + ds = _dataset(tmp_path, _pose(jumps=[60])) + (name,) = ds.datasets + video = pinned_episode_subset(ds, [name]) + assert len(video) == T + assert len(ds) < T and ds.datasets[name].sample_indices is not None diff --git a/tests/unit/test_rdt.py b/tests/unit/test_rdt.py index 06bb103ce5..ec7ba0a4ee 100644 --- a/tests/unit/test_rdt.py +++ b/tests/unit/test_rdt.py @@ -160,6 +160,7 @@ def _model( frame_dropout=0.0, memory=None, short_dropout=0.0, + lang_pad_to=None, ): torch.manual_seed(0) model = RDTModel( @@ -171,6 +172,7 @@ def _model( image_history_dropout=frame_dropout, memory=memory, short_dropout=short_dropout, + lang_pad_to=lang_pad_to, ) shared = {"front_img_1": MLPPolicyStem(input_dim=D, output_dim=D, widths=[D])} if text: @@ -239,6 +241,21 @@ def test_model_streams_and_shapes(): assert cond.backbone is model.trunk["trunk"] +def test_lang_pad_to_fixes_the_length_without_changing_the_output(): + data = _data() + plain, padded = _model().eval(), _model(lang_pad_to=8).eval() + cond, _ = plain.forward_features(DOMAIN, data) + cond_pad, _ = padded.forward_features(DOMAIN, data) + assert cond_pad.lang.shape == (B, 8, D) + assert cond_pad.lang_mask.sum(1).tolist() == [1, 2, 3] + torch.manual_seed(3) + x, t = torch.randn(B, SEQ, ACT), torch.rand(B) + with torch.no_grad(): + out = plain.heads[DOMAIN].model(x, t, cond) + out_pad = padded.heads[DOMAIN].model(x, t, cond_pad) + torch.testing.assert_close(out_pad, out) + + def test_model_shares_the_dit_in_the_trunk_not_the_head(): names = [n for n, _ in _model().named_parameters()] assert any(n.startswith("trunk.trunk.blocks.") for n in names) @@ -397,6 +414,26 @@ def test_algo_keeps_past_and_current_in_order(): assert frames[:, 0].sum() == 0 and frames[:, 1].min() == 1 # current frame last +def test_algo_stacks_a_multi_frame_history_oldest_first(): + algo = _pairing_algo(nn.Identity()) + cam = "observations.images.front_img_1" + past = torch.arange(3.0).view(1, 3, 1, 1, 1).expand(B, 3, 3, 8, 8).contiguous() + batch = { + cam: torch.full((B, 3, 8, 8), 3.0), + f"{cam}_hist": past, + f"{cam}_hist_mask": torch.ones(B, 3), + "actions": torch.zeros(B, SEQ, ACT), + "pad_mask": torch.ones(B, SEQ, 1), + "embodiment": torch.tensor([3]), + "fps": torch.full((B,), 30.0), + } + frames = algo._robomimic_to_hpt_data( + batch, [cam, f"{cam}_hist"], [], [], "actions" + )["front_img_1"] + assert frames.shape == (B, 4, 1, 3, 8, 8) + assert frames.flatten(2).mean(-1)[0].tolist() == [0.0, 1.0, 2.0, 3.0] + + def _pair_batch(current, past): cam = "observations.images.front_img_1" n = current.shape[0] diff --git a/tests/unit/test_rdt_mem_configs.py b/tests/unit/test_rdt_mem_configs.py index cfabc58c12..87dbda3a83 100644 --- a/tests/unit/test_rdt_mem_configs.py +++ b/tests/unit/test_rdt_mem_configs.py @@ -9,13 +9,19 @@ from egomimic.trainHydra import _build_model_config_tree -ARMS = ("base", "proprio", "vision", "both") +ARM_NAMES = { + "20pct_6d": ("base", "proprio", "vision", "both"), + "all_kp": ("base", "proprio", "proprio_frames", "short_long_encoder"), +} +ARMS = [ + f"train_zarr_mecka_{variant}_rdt_mem_{arm}" + for variant in ("20pct_6d", "all_kp") + for arm in ARM_NAMES[variant] +] def _compose(arm, compose_resolve): - cfg = compose_resolve( - f"train_zarr_mecka_20pct_6d_rdt_mem_{arm}", [], keep_hydra=True - ) + cfg = compose_resolve(arm, [], keep_hydra=True) return OmegaConf.masked_copy(cfg, [k for k in cfg if k != "hydra"]) @@ -34,8 +40,14 @@ def test_model_sizes_match_the_keymap(arm, compose_resolve): model = cfg.model.robomimic_model denoiser = model.head_specs.human_bimanual.model history = km.get("proprio_history", 1) - assert model.stem_specs.human_bimanual.state_ee_pose.history_len == history + (state_stem,) = model.stem_specs.human_bimanual.values() + assert state_stem.history_len == history assert denoiser.n_state_tokens == history + frames = model.trunk.get("image_history", 1) + if km.get("image_history_gap_s") is None: + assert frames == 1 + else: + assert frames == km.get("image_history", 2) memory = model.trunk.get("memory") if memory is None: assert not km.get("image_memory") and not denoiser.get("n_memory_tokens") diff --git a/tests/unit/test_stage_episodes.py b/tests/unit/test_stage_episodes.py index 96500e0e34..aaa5d772e4 100644 --- a/tests/unit/test_stage_episodes.py +++ b/tests/unit/test_stage_episodes.py @@ -19,24 +19,52 @@ StagingError, stage_episodes, ) +from egomimic.rldb.zarr.zarr_dataset_multi import ( + staged_episode_dir as ns, +) VENDOR = next(v for v in sorted(VENDORS) if VENDORS[v].embodiment == "human_bimanual") +@pytest.fixture(autouse=True) +def _no_slurm(monkeypatch): + for var in ("SLURM_JOB_ID", "SLURM_STEP_ID", "SLURM_RESTART_COUNT"): + monkeypatch.delenv(var, raising=False) + + +def _launch(monkeypatch, job="7", step="0"): + monkeypatch.setenv("SLURM_JOB_ID", job) + monkeypatch.setenv("SLURM_STEP_ID", step) + + +def _count_src_checks(monkeypatch, src): + from egomimic.rldb.zarr import norm_cache + + calls = [] + real = norm_cache.episode_fingerprint + + def counting(path): + if Path(path).parent == src: + calls.append(Path(path).name) + return real(path) + + monkeypatch.setattr(norm_cache, "episode_fingerprint", counting) + return calls + + def _episodes(root, n=2): return {write_episode(root, VENDOR, seed=i).name[: -len(".zarr")] for i in range(n)} def _count_copies(monkeypatch): calls = [] - real = shutil.copytree + real = zdm._tar_copy - def counting(src, dst, *a, **k): - if Path(src).name.endswith(".zarr"): # copytree recurses into itself - calls.append(Path(src)) - return real(src, dst, *a, **k) + def counting(src_root, names, dest): + calls.extend(Path(src_root) / n for n in names) + return real(src_root, names, dest) - monkeypatch.setattr(zdm.shutil, "copytree", counting) + monkeypatch.setattr(zdm, "_tar_copy", counting) return calls @@ -46,10 +74,10 @@ def test_stages_and_keeps_fingerprints(tmp_path): staged = stage_episodes(src, stage, names | {"not_synced"}) assert staged == names for n in names: - assert episode_fingerprint(stage / f"{n}.zarr") == episode_fingerprint( + assert episode_fingerprint(ns(stage, src) / f"{n}.zarr") == episode_fingerprint( src / f"{n}.zarr" ) - assert not list(stage.glob(".*.partial")) + assert not (stage / ".incoming").exists() def test_fresh_copies_are_reused(tmp_path, monkeypatch): @@ -119,10 +147,10 @@ def test_failed_copy_raises(tmp_path, monkeypatch): def boom(*a, **k): raise OSError("disk on fire") - monkeypatch.setattr(zdm.shutil, "copytree", boom) + monkeypatch.setattr(zdm, "_tar_copy", boom) with pytest.raises(StagingError, match="disk on fire"): stage_episodes(src, stage, names) - assert not list(stage.glob(".*.partial")) + assert not list(stage.glob(".incoming*")) def test_resolver_reads_from_stage(tmp_path, monkeypatch): @@ -141,7 +169,7 @@ def test_resolver_reads_from_stage(tmp_path, monkeypatch): datasets = resolver.load([(None, n) for n in names]) assert sorted(datasets) == sorted(names) for ds in datasets.values(): - assert Path(ds.episode_path).parent == stage + assert Path(ds.episode_path).parent == ns(stage, src) staged = MultiDataset._from_resolver(resolver, mode="total") assert len(staged) == len(baseline) @@ -149,3 +177,279 @@ def test_resolver_reads_from_stage(tmp_path, monkeypatch): np.asarray(staged[0]["actions_cartesian"]), np.asarray(baseline[0]["actions_cartesian"]), ) + + +def test_launch_checks_the_source_once(tmp_path, monkeypatch): + """The node's other ranks and later heads of one launch reuse the first + caller's check instead of re-reading the source over NFS.""" + src, stage = tmp_path / "nfs", tmp_path / "local" + names = _episodes(src, n=3) + _launch(monkeypatch) + checks = _count_src_checks(monkeypatch, src) + assert stage_episodes(src, stage, names | {"not_synced"}) == names + assert len(checks) == 3 + subset = set(sorted(names)[:2]) + assert stage_episodes(src, stage, subset) == subset + assert stage_episodes(src, stage, names | {"not_synced"}) == names + assert len(checks) == 3 + + +def test_new_launch_rechecks_the_source(tmp_path, monkeypatch): + src, stage = tmp_path / "nfs", tmp_path / "local" + names = _episodes(src) + _launch(monkeypatch, step="0") + stage_episodes(src, stage, names) + changed = sorted(names)[0] + meta = src / f"{changed}.zarr" / "zarr.json" + st = meta.stat() + os.utime(meta, ns=(st.st_atime_ns, st.st_mtime_ns + 10**9)) + calls = _count_copies(monkeypatch) + assert stage_episodes(src, stage, names) == names + assert calls == [] + _launch(monkeypatch, step="1") + assert stage_episodes(src, stage, names) == names + assert [p.name for p in calls] == [f"{changed}.zarr"] + + +def test_launch_record_is_per_source(tmp_path, monkeypatch): + stage = tmp_path / "local" + names = _episodes(tmp_path / "a") + _launch(monkeypatch) + assert stage_episodes(tmp_path / "a", stage, names) == names + assert stage_episodes(tmp_path / "b", stage, names) == set() + + +def _bump_mtime(path): + st = path.stat() + os.utime(path, ns=(st.st_atime_ns, st.st_mtime_ns + 10**9)) + + +def test_versioned_copy_is_trusted_without_reading_the_source(tmp_path, monkeypatch): + src, stage = tmp_path / "nfs", tmp_path / "local" + names = _episodes(src) + versions = {n: "v1" for n in names} + _launch(monkeypatch, step="0") + stage_episodes(src, stage, names, versions) + _launch(monkeypatch, step="1") + checks = _count_src_checks(monkeypatch, src) + assert stage_episodes(src, stage, names, versions) == names + assert checks == [] + + +def test_new_version_rechecks_and_recopies(tmp_path, monkeypatch): + src, stage = tmp_path / "nfs", tmp_path / "local" + names = _episodes(src) + stage_episodes(src, stage, names, {n: "v1" for n in names}) + changed = sorted(names)[0] + _bump_mtime(src / f"{changed}.zarr" / "zarr.json") + calls = _count_copies(monkeypatch) + assert stage_episodes(src, stage, names, {n: "v1" for n in names}) == names + assert calls == [] + versions = {n: "v1" for n in names} | {changed: "v2"} + assert stage_episodes(src, stage, names, versions) == names + assert [p.name for p in calls] == [f"{changed}.zarr"] + + +def test_damaged_local_copy_is_restaged(tmp_path, monkeypatch): + src, stage = tmp_path / "nfs", tmp_path / "local" + names = _episodes(src) + versions = {n: "v1" for n in names} + stage_episodes(src, stage, names, versions) + damaged = sorted(names)[0] + _bump_mtime(ns(stage, src) / f"{damaged}.zarr" / "zarr.json") + calls = _count_copies(monkeypatch) + assert stage_episodes(src, stage, names, versions) == names + assert [p.name for p in calls] == [f"{damaged}.zarr"] + + +def test_s3_sync_runs_once_per_launch(tmp_path, monkeypatch): + monkeypatch.setenv(zdm.STAGE_DIR_ENV, str(tmp_path / "local")) + synced = [] + monkeypatch.setattr( + zdm.S3EpisodeResolver, + "_sync_s3_to_local_unlocked", + classmethod(lambda cls, b, paths, d, n: synced.append([h for _, h in paths])), + ) + paths = [("s3://b/p/a", "a"), ("s3://b/p/b", "b")] + + def sync(p): + zdm.S3EpisodeResolver._sync_s3_to_local("b", p, tmp_path / "nfs") + + _launch(monkeypatch) + sync(paths) + sync(paths) + sync(paths + [("s3://b/p/c", "c")]) + assert synced == [["a", "b"], [], ["c"]] + _launch(monkeypatch, step="1") + sync(paths) + assert synced[-1] == ["a", "b"] + + +def test_s3_sync_outside_a_step_needs_no_stage_dir(tmp_path, monkeypatch): + """The download script runs on the login node, which has no /ephemeral.""" + blocker = tmp_path / "file" + blocker.write_text("") + monkeypatch.setenv(zdm.STAGE_DIR_ENV, str(blocker / "stage")) + synced = [] + monkeypatch.setattr( + zdm.S3EpisodeResolver, + "_sync_s3_to_local_unlocked", + classmethod(lambda cls, b, paths, d, n: synced.append([h for _, h in paths])), + ) + zdm.S3EpisodeResolver._sync_s3_to_local("b", [("s3://b/p/a", "a")], tmp_path) + _launch(monkeypatch) + zdm.S3EpisodeResolver._sync_s3_to_local("b", [("s3://b/p/a", "a")], tmp_path) + assert synced == [["a"], ["a"]] + + +def test_symlinked_source_is_staged_as_data(tmp_path): + real, src, stage = tmp_path / "real", tmp_path / "nfs", tmp_path / "local" + names = _episodes(real, n=1) + src.mkdir() + for n in names: + (src / f"{n}.zarr").symlink_to(real / f"{n}.zarr") + assert stage_episodes(src, stage, names) == names + for n in names: + staged = ns(stage, src) / f"{n}.zarr" + assert staged.is_dir() and not staged.is_symlink() + assert not any(p.is_symlink() for p in staged.rglob("*")) + + +def test_trust_is_per_source(tmp_path, monkeypatch): + stage = tmp_path / "local" + names = _episodes(tmp_path / "a", n=1) + versions = {n: "v1" for n in names} + assert stage_episodes(tmp_path / "a", stage, names, versions) == names + assert stage_episodes(tmp_path / "missing", stage, names, versions) == set() + + +def test_trust_expires(tmp_path, monkeypatch): + src, stage = tmp_path / "nfs", tmp_path / "local" + names = _episodes(src, n=1) + versions = {n: "v1" for n in names} + stage_episodes(src, stage, names, versions) + monkeypatch.setattr(zdm, "_STAGED_TRUST_S", -1) + checks = _count_src_checks(monkeypatch, src) + assert stage_episodes(src, stage, names, versions) == names + assert len(checks) == 1 + + +def test_failed_chunk_keeps_earlier_chunks(tmp_path, monkeypatch): + src, stage = tmp_path / "nfs", tmp_path / "local" + names = _episodes(src, n=3) + monkeypatch.setattr(zdm, "_STAGE_CHUNK", 1) + real = zdm._tar_copy + calls = [] + + def flaky(src_root, chunk, dest): + calls.append(chunk) + if len(calls) == 2: + raise OSError("disk on fire") + return real(src_root, chunk, dest) + + monkeypatch.setattr(zdm, "_tar_copy", flaky) + versions = {n: "v1" for n in names} + with pytest.raises(StagingError, match="disk on fire"): + stage_episodes(src, stage, names, versions) + first = calls[0][0][: -len(".zarr")] + monkeypatch.setattr(zdm, "_tar_copy", real) + recopied = _count_copies(monkeypatch) + assert stage_episodes(src, stage, names, versions) == names + assert f"{first}.zarr" not in [p.name for p in recopied] + assert len(recopied) == 2 + + +def test_leftover_incoming_is_removed(tmp_path): + src, stage = tmp_path / "nfs", tmp_path / "local" + names = _episodes(src, n=1) + (stage / ".incoming" / "junk").mkdir(parents=True) + assert stage_episodes(src, stage, names) == names + assert not (stage / ".incoming").exists() + + +def test_tar_reports_errors_without_hanging(tmp_path): + """Thousands of per-file warnings must not fill a pipe and deadlock.""" + if os.geteuid() == 0: + pytest.skip("root reads mode-000 files") + src = tmp_path / "nfs" / "ep.zarr" + src.mkdir(parents=True) + for i in range(3000): + f = src / f"f{i:05d}_{'x' * 40}" + f.write_text("") + f.chmod(0) + dest = tmp_path / "out" + dest.mkdir() + with pytest.raises(OSError, match="Permission denied"): + zdm._tar_copy(tmp_path / "nfs", ["ep.zarr"], dest) + + +def test_same_hash_from_two_sources_keeps_both(tmp_path): + stage = tmp_path / "local" + names = _episodes(tmp_path / "a", n=1) + _episodes(tmp_path / "b", n=1) + versions = {n: "v1" for n in names} + assert stage_episodes(tmp_path / "a", stage, names, versions) == names + assert stage_episodes(tmp_path / "b", stage, names, versions) == names + calls = [] + for src in ("a", "b"): + assert stage_episodes(tmp_path / src, stage, names, versions) == names + calls += [ns(stage, tmp_path / src) / f"{n}.zarr" for n in names] + assert all(p.is_dir() for p in calls) + + +def test_verified_name_needs_a_local_copy(tmp_path, monkeypatch): + src, stage = tmp_path / "nfs", tmp_path / "local" + names = _episodes(src) + _launch(monkeypatch) + stage_episodes(src, stage, names) + gone = sorted(names)[0] + shutil.rmtree(ns(stage, src) / f"{gone}.zarr") + assert stage_episodes(src, stage, names) == names - {gone} + + +def test_restage_swaps_by_rename(tmp_path, monkeypatch): + """The old copy is never deleted in place, so a kill can't leave a + half-deleted episode that still fingerprints as intact.""" + src, stage = tmp_path / "nfs", tmp_path / "local" + names = _episodes(src, n=1) + stage_episodes(src, stage, names) + (name,) = names + _bump_mtime(src / f"{name}.zarr" / "zarr.json") + deleted = [] + real = shutil.rmtree + + def watching(path, *a, **k): + deleted.append(Path(path)) + return real(path, *a, **k) + + monkeypatch.setattr(zdm.shutil, "rmtree", watching) + assert stage_episodes(src, stage, names) == names + assert ns(stage, src) / f"{name}.zarr" not in deleted + assert not (stage / ".trash").exists() + + +def test_full_disk_evicts_stale_unused_episodes(tmp_path, monkeypatch): + src, stage = tmp_path / "nfs", tmp_path / "local" + old = _episodes(tmp_path / "old", n=1) + stage_episodes(tmp_path / "old", stage, old) + (old_name,) = old + old_dir = ns(stage, tmp_path / "old") / f"{old_name}.zarr" + os.utime(old_dir, (0, 0)) + new = {write_episode(src, VENDOR, seed=7).name[: -len(".zarr")]} + + def usage(_): + full = old_dir.exists() + return SimpleNamespace( + total=100, used=95 if full else 50, free=5 if full else 50 + ) + + monkeypatch.setattr(zdm.shutil, "disk_usage", usage) + assert stage_episodes(src, stage, new) == new + assert not old_dir.exists() + + +def test_launch_key_without_a_step_uses_the_process_group(monkeypatch): + monkeypatch.setenv("SLURM_JOB_ID", "7") + assert zdm._launch_key() == f"7.pg{os.getpgid(0)}.0" + monkeypatch.setenv("SLURM_STEP_ID", "3") + assert zdm._launch_key() == "7.s3.0"