diff --git a/egomimic/eval/action_metrics.py b/egomimic/eval/action_metrics.py index 9c1d099562..497a79e94b 100644 --- a/egomimic/eval/action_metrics.py +++ b/egomimic/eval/action_metrics.py @@ -202,8 +202,35 @@ def _kp_l2(p, g): return metrics -def layout_metrics(pred: torch.Tensor, gt: torch.Tensor, prefix: str) -> dict: - """Dispatch on the last-dim width: cartesian, keypoint, or ``{}``.""" +def joint_metrics(pred: torch.Tensor, gt: torch.Tensor, prefix: str) -> dict: + """Metrics for a bimanual joint-space chunk ``(B, T, 14)``; ``{}`` for + other widths. Joint errors are in degrees (mean absolute), the gripper in + its [0, 1] command units.""" + pred_cpu, gt_cpu = pred.detach().cpu().float(), gt.detach().cpu().float() + if pred_cpu.shape[-1] != 14: + return {} + # [left j1..j6, left grip, right j1..j6, right grip] (ABC YAM / abc_sim) + j, g = [0, 1, 2, 3, 4, 5, 7, 8, 9, 10, 11, 12], [6, 13] + err = (pred_cpu[..., j] - gt_cpu[..., j]).abs() * _RAD2DEG + return { + f"{prefix}_joint_abs_err_deg_avg": err.mean(), + f"{prefix}_joint_abs_err_deg_final": err[:, -1].mean(), + f"{prefix}_joint_abs_err_deg_max_joint": err.mean(dim=(0, 1)).max(), + f"{prefix}_grip_paired_mse_avg": _paired_mse(pred_cpu[..., g], gt_cpu[..., g]), + f"{prefix}_grip_final_mse_avg": _paired_mse( + pred_cpu[:, -1, g], gt_cpu[:, -1, g] + ), + } + + +def layout_metrics( + pred: torch.Tensor, gt: torch.Tensor, prefix: str, ac_key: str | None = None +) -> dict: + """Dispatch on the action key, then the last-dim width: joints, cartesian, + keypoint, or ``{}``. A 14-D joint chunk has a cartesian width (xyz+ypr+grip + per arm), so joint actions must be named: ``ac_key`` ``actions_joints*``.""" + if ac_key is not None and ac_key.startswith("actions_joints"): + return joint_metrics(pred, gt, prefix) metrics = cartesian_metrics(pred, gt, prefix) if metrics: return metrics diff --git a/egomimic/eval/eval_hpt.py b/egomimic/eval/eval_hpt.py index 8c5679d00f..9638d852f2 100644 --- a/egomimic/eval/eval_hpt.py +++ b/egomimic/eval/eval_hpt.py @@ -59,7 +59,9 @@ def compute_metrics_and_viz(self, batch, do_viz=True): metrics[f"Valid/{pred_key}_frechet_gauss_min"] = fd.min().item() metrics[f"Valid/{pred_key}_frechet_gauss_max"] = fd.max().item() metrics.update( - layout_metrics(preds[pred_key], _batch[ac_key], f"Valid/{pred_key}") + layout_metrics( + preds[pred_key], _batch[ac_key], f"Valid/{pred_key}", ac_key + ) ) if embodiment_name in algo.auxiliary_ac_keys: @@ -103,6 +105,10 @@ def compute_metrics_and_viz(self, batch, do_viz=True): if not do_viz: continue + if self.viz_func is not None and embodiment_name not in self.viz_func: + # No overlay for this embodiment (e.g. joint-space actions, + # which have no image projection): metrics only. + continue transform_list = self.transform_lists.get(embodiment_name) main_pred_key = f"{embodiment_name}_{ac_key}" diff --git a/egomimic/eval/eval_pi.py b/egomimic/eval/eval_pi.py index 6bd1aa4ae4..65c6f04af5 100644 --- a/egomimic/eval/eval_pi.py +++ b/egomimic/eval/eval_pi.py @@ -58,7 +58,9 @@ def compute_metrics_and_viz(self, batch, do_viz=True): if pred_key in preds: metrics.update( - layout_metrics(preds[pred_key], _batch[ac_key], f"Valid/{pred_key}") + layout_metrics( + preds[pred_key], _batch[ac_key], f"Valid/{pred_key}", ac_key + ) ) if do_viz: diff --git a/egomimic/hydra_configs/data/abc_real_joints.yaml b/egomimic/hydra_configs/data/abc_real_joints.yaml new file mode 100644 index 0000000000..533f277c65 --- /dev/null +++ b/egomimic/hydra_configs/data/abc_real_joints.yaml @@ -0,0 +1,32 @@ +# ABC_REAL_JOINTS: one REAL ABC task (lab=abc) in the same joint space. The SQL +# embodiment was relabelled to yam_bimanual on 2026-09-01 but the zarrs say +# eva_bimanual, so the filter keys on lab + task + the /abc/ path. Another task: +# restate the filter_lambdas with its name. +defaults: + - abc_sim_joints + - _self_ + +train_datasets: + eva_bimanual: + resolver: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver + folder_path: ${paths.dataset_dir} + transform_list: + # ABC mixes 640x480 and 1280x720 front frames; one size so they collate + image_hw: [240, 320] + filters: + _target_: egomimic.rldb.filters.DatasetFilter + filter_lambdas: + - "lambda row: str(row.get('lab', '')).lower() == 'abc'" + - "lambda row: str(row.get('task', '')) == 'put the plastic bottles in the bin'" + - "lambda row: '/abc/' in str(row.get('zarr_processed_path', ''))" + - "lambda row: not bool(row.get('is_deleted', False))" + mode: train + valid_ratio: 0.1 + +valid_datasets: + eva_bimanual: + filters: ${data.train_datasets.eva_bimanual.filters} + mode: valid + valid_ratio: ${data.train_datasets.eva_bimanual.valid_ratio} + diff --git a/egomimic/hydra_configs/data/abc_sim_joints.yaml b/egomimic/hydra_configs/data/abc_sim_joints.yaml new file mode 100644 index 0000000000..e4bffbcacf --- /dev/null +++ b/egomimic/hydra_configs/data/abc_sim_joints.yaml @@ -0,0 +1,50 @@ +# ABC_SIM_JOINTS: one abc_sim task's sim_224 teleop episodes, converted by +# egomimic/scripts/abc_sim/convert_sim_to_zarr.py to $ABC_SIM_ZARR_ROOT/put_bottles +# (another task: override the resolver folder_path). Eva "joints" mode, front +# camera + its 0.1 s history frame; train/valid = ABC's own `split`. +_target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper + +train_datasets: + eva_bimanual: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + resolver: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.LocalEpisodeResolver + folder_path: ${oc.env:ABC_SIM_ZARR_ROOT,/workspace/users/agao81/datasets/abc_sim/zarr}/put_bottles + key_map: + _target_: egomimic.rldb.embodiment.eva.Eva.get_keymap + keymap_mode: joints_front + annotation_key: annotations + image_history_gap_s: 0.1 + transform_list: + _target_: egomimic.rldb.embodiment.eva.Eva.get_transform_list + mode: joints + filters: + _target_: egomimic.rldb.filters.DatasetFilter + filter_lambdas: + - "lambda row: str(row.get('lab', '')) == 'abc_sim'" + - "lambda row: str(row.get('split', '')) == 'train'" + mode: total + +valid_datasets: + eva_bimanual: + _target_: ${data.train_datasets.eva_bimanual._target_} + resolver: ${data.train_datasets.eva_bimanual.resolver} + filters: + _target_: egomimic.rldb.filters.DatasetFilter + filter_lambdas: + - "lambda row: str(row.get('lab', '')) == 'abc_sim'" + - "lambda row: str(row.get('split', '')) == 'val'" + mode: total + + + +train_dataloader_params: + eva_bimanual: + batch_size: 32 + num_workers: 8 + +valid_dataloader_params: + eva_bimanual: + batch_size: 32 + num_workers: 4 + shuffle: false diff --git a/egomimic/hydra_configs/data/abc_sim_joints_cotrain_mecka.yaml b/egomimic/hydra_configs/data/abc_sim_joints_cotrain_mecka.yaml new file mode 100644 index 0000000000..3e7e011a9f --- /dev/null +++ b/egomimic/hydra_configs/data/abc_sim_joints_cotrain_mecka.yaml @@ -0,0 +1,28 @@ +# ABC_SIM_JOINTS_COTRAIN_MECKA: single-task cotrain -- the abc_sim task (robot, +# joint space, abc_sim_joints) + one mecka task (human, mecka_all_6d's +# wrist-frame 6D pipeline), one batch of each per step (CombinedLoader +# max_size_cycle). Default human task cup_on_saucer (the largest mecka task); +# another one restates the human filter_lambdas. See egomimic/scripts/abc_sim/README.md. +defaults: + - abc_sim_joints + - mecka_all_6d + - _self_ + +train_datasets: + human_bimanual: + resolver: + key_map: + # RDT's 2-frame image history, as on the robot side + image_history_gap_s: 0.1 + filters: + filter_lambdas: + - "lambda row: str(row.get('lab', '')).lower() == 'mecka'" + - "lambda row: str(row.get('task', '')) == 'cup_on_saucer'" + +train_dataloader_params: + human_bimanual: + batch_size: 32 + +valid_dataloader_params: + human_bimanual: + batch_size: 32 diff --git a/egomimic/hydra_configs/evaluator/eval_hpt_abc_joints.yaml b/egomimic/hydra_configs/evaluator/eval_hpt_abc_joints.yaml new file mode 100644 index 0000000000..c51a298640 --- /dev/null +++ b/egomimic/hydra_configs/evaluator/eval_hpt_abc_joints.yaml @@ -0,0 +1,17 @@ +# HPT/RDT evaluator for the ABC joint-space runs: eva_bimanual (ABC station, +# actions_joints) gets the joint metrics (Valid/eva_bimanual_actions_joints_ +# joint_abs_err_deg_*, _grip_*) and no video; human_bimanual keeps the +# wrist-frame 6D revert + overlay of eval_hpt_wristframe_6d. The closed-loop +# number for these runs is the abc_sim rollout, not this. +defaults: + - viz@viz_func: abc_joints_human_cartesian + - _self_ + +_target_: egomimic.eval.eval_hpt.HPTEvalVideo + +viz_every_n_epochs: 100 +viz_max_batches: 10 + +transform_lists: + human_bimanual: + _target_: egomimic.rldb.embodiment.human._build_human_cartesian_revert_6d_wristframe_transform_list diff --git a/egomimic/hydra_configs/evaluator/viz/abc_joints_human_cartesian.yaml b/egomimic/hydra_configs/evaluator/viz/abc_joints_human_cartesian.yaml new file mode 100644 index 0000000000..11a6188b52 --- /dev/null +++ b/egomimic/hydra_configs/evaluator/viz/abc_joints_human_cartesian.yaml @@ -0,0 +1,10 @@ +# Overlay viz for the ABC joint-space cotrain: human only. The robot predicts +# joint angles, which have no image projection without forward kinematics, so +# eva_bimanual is left out and HPTEvalVideo logs its metrics without a video. +human_bimanual: + _target_: egomimic.rldb.embodiment.human.Human.viz_gt_preds + _partial_: true + image_key: observations.images.front_img_1 + action_key: actions_cartesian + annotation_key: sampled_prompt + mode: traj diff --git a/egomimic/hydra_configs/model/rdt_1b_abc_joints_dinov3_vitb.yaml b/egomimic/hydra_configs/model/rdt_1b_abc_joints_dinov3_vitb.yaml new file mode 100644 index 0000000000..3b14436e2e --- /dev/null +++ b/egomimic/hydra_configs/model/rdt_1b_abc_joints_dinov3_vitb.yaml @@ -0,0 +1,10 @@ +# RDT_1B_ABC_JOINTS_DINOV3_VITB: rdt_abc_joints_dinov3_vitb at RDT-1B width +# (rdt_1b_bc_mecka_6d_dinov3_vitb's width and heads). +defaults: + - rdt_abc_joints_dinov3_vitb + - _self_ + +robomimic_model: + width: 2048 + trunk: + num_heads: 32 diff --git a/egomimic/hydra_configs/model/rdt_abc_joints_dinov3_vitb.yaml b/egomimic/hydra_configs/model/rdt_abc_joints_dinov3_vitb.yaml new file mode 100644 index 0000000000..49fe4cd2f1 --- /dev/null +++ b/egomimic/hydra_configs/model/rdt_abc_joints_dinov3_vitb.yaml @@ -0,0 +1,44 @@ +# RDT_ABC_JOINTS_DINOV3_VITB: the base RDT (width 1024, depth 28: the architecture of +# Ryan's all-mecka rdt_mem_all checkpoints) with a second head, eva_bimanual = the ABC +# YAM station in JOINT space (14-D state, (100, 14) actions_joints, abc_sim's action +# space), beside the mecka human head; trunk, front tower and text stem shared. Front +# camera only, at the human tower size [224, 400]. rdt_1b_abc_joints is this at width 2048. +defaults: + - rdt_bc_mecka_6d_dinov3_vitb + - _self_ + +robomimic_model: + domains: ["human_bimanual", "eva_bimanual"] + dims: + eva_bimanual: {proprio: 14, action: 14} + ac_keys: + eva_bimanual: "actions_joints" + + head_specs: + eva_bimanual: + _target_: egomimic.models.fm_policy.FMPolicy + action_horizon: 100 + num_inference_steps: 10 + pooling: null + time_dist: "beta" + infer_ac_dims: + eva_bimanual: ${model.robomimic_model.dims.eva_bimanual.action} + model: + _target_: egomimic.models.rdt_nets.RDTDenoiser + act_dim: ${model.robomimic_model.dims.eva_bimanual.action} + act_seq: 100 + hidden_dim: ${model.robomimic_model.width} + n_state_tokens: 1 + time_scale: 1000.0 + + stem_specs: + eva_bimanual: + state_joint_positions: + _target_: egomimic.models.hpt_nets.MLPPolicyStem + input_dim: ${model.robomimic_model.dims.eva_bimanual.proprio} + output_dim: ${model.robomimic_model.width} + widths: + - ${model.robomimic_model.width} + - ${model.robomimic_model.width} + history_len: 1 + history_dropout: 0.2 diff --git a/egomimic/hydra_configs/train_zarr_abc_real_rdt.yaml b/egomimic/hydra_configs/train_zarr_abc_real_rdt.yaml new file mode 100644 index 0000000000..c1c3b3df86 --- /dev/null +++ b/egomimic/hydra_configs/train_zarr_abc_real_rdt.yaml @@ -0,0 +1,38 @@ +# ABC REAL bottles: the base RDT (width 1024) trained on the real ABC episodes only +# (data=abc_real_joints, 90/10 episode split) and evaluated in abc_sim. Built to run +# unchanged on an ABC-replica YAM station: the prompt is fixed in the model +# (default_prompt; no annotation track is read) and there is no image history, so a +# client that sends the current camera frames + the 14-D joint state is enough +# (scripts/abc_sim/sim_client.py, ABC's deploy client through policy_server). +# +# from scratch: python egomimic/trainHydra.py --config-name=train_zarr_abc_real_rdt +# from Ryan's all-mecka base: ... init_weights_ckpt= +defaults: + - train_zarr_abc_real_rdt_ft + - override model: rdt_abc_joints_dinov3_vitb + - _self_ + +model: + robomimic_model: + default_prompt: "put the plastic bottles in the bin" + trunk: + image_history: 1 + +data: + train_datasets: + eva_bimanual: + resolver: + key_map: + annotation_key: null + image_history_gap_s: null + +# 1000 x 100 steps at 32 per GPU x 4 GPUs; a checkpoint every 50 epochs (5k steps) +trainer: + max_epochs: 1000 + min_epochs: 1000 + +launch_params: + gpus_per_node: 4 + +norm_stats: + max_samples: 200000 diff --git a/egomimic/hydra_configs/train_zarr_abc_real_rdt_ft.yaml b/egomimic/hydra_configs/train_zarr_abc_real_rdt_ft.yaml new file mode 100644 index 0000000000..f1f57c2e9a --- /dev/null +++ b/egomimic/hydra_configs/train_zarr_abc_real_rdt_ft.yaml @@ -0,0 +1,10 @@ +# ABC REAL fine-tune: the train_zarr_abc_sim_rdt_ft arm on the matching REAL +# ABC task (data=abc_real_joints, lab=abc), for the real-robot side of the +# sim-real rank check. Same model, init, steps and batch as the sim arm. +# +# python egomimic/trainHydra.py --config-name=train_zarr_abc_real_rdt_ft \ +# init_weights_ckpt= +defaults: + - train_zarr_abc_sim_rdt_ft + - override data: abc_real_joints + - _self_ diff --git a/egomimic/hydra_configs/train_zarr_abc_sim_dishrack_rdt.yaml b/egomimic/hydra_configs/train_zarr_abc_sim_dishrack_rdt.yaml new file mode 100644 index 0000000000..15e7263adb --- /dev/null +++ b/egomimic/hydra_configs/train_zarr_abc_sim_dishrack_rdt.yaml @@ -0,0 +1,20 @@ +# Dish-rack sim only: abc_sim load_plates_into_dish_rack (robot, joint space, converted to +# $ABC_SIM_ZARR_ROOT/dishrack), base RDT (width 1024) from scratch with the real-bottles +# recipe's deployable settings (robot prompt fixed in the model, no image history). The +# matched control for train_zarr_abc_sim_dishrack_rdt_cotrain. +# +# python egomimic/trainHydra.py --config-name=train_zarr_abc_sim_dishrack_rdt +defaults: + - train_zarr_abc_real_rdt + - override data: abc_sim_joints + - _self_ + +model: + robomimic_model: + default_prompt: "load the plates into the dish rack" + +data: + train_datasets: + eva_bimanual: + resolver: + folder_path: ${oc.env:ABC_SIM_ZARR_ROOT,/workspace/users/agao81/datasets/abc_sim/zarr}/dishrack diff --git a/egomimic/hydra_configs/train_zarr_abc_sim_dishrack_rdt_cotrain.yaml b/egomimic/hydra_configs/train_zarr_abc_sim_dishrack_rdt_cotrain.yaml new file mode 100644 index 0000000000..9d8aa50b6f --- /dev/null +++ b/egomimic/hydra_configs/train_zarr_abc_sim_dishrack_rdt_cotrain.yaml @@ -0,0 +1,40 @@ +# Dish-rack human + sim cotrain: abc_sim load_plates_into_dish_rack (robot, joint space, +# converted to $ABC_SIM_ZARR_ROOT/dishrack) + the Mecka episodes whose description has +# dishes and a rack (human, wrist-frame 6D; 232 episodes / 5.4 h on 2026-10-03), one batch +# of each per step. Base RDT (width 1024) from scratch with the real-bottles recipe's +# deployable settings: the robot prompt is fixed in the model, no image history. +# +# python egomimic/trainHydra.py --config-name=train_zarr_abc_sim_dishrack_rdt_cotrain +defaults: + - train_zarr_abc_real_rdt + - override data: abc_sim_joints_cotrain_mecka + - _self_ + +model: + robomimic_model: + domains: ["human_bimanual", "eva_bimanual"] + default_prompt: "load the plates into the dish rack" + +data: + train_datasets: + eva_bimanual: + resolver: + folder_path: ${oc.env:ABC_SIM_ZARR_ROOT,/workspace/users/agao81/datasets/abc_sim/zarr}/dishrack + human_bimanual: + resolver: + key_map: + image_history_gap_s: null # trunk.image_history is 1 + filters: + filter_lambdas: + - "lambda row: str(row.get('lab', '')).lower() == 'mecka'" + # the SQL selection of 2026-10-03: a rack, and dishes, in the episode description + - "lambda row: bool(__import__('re').search(r'\\b(?:dish|drying|plate|draining)[ -]?rack|\\brack\\b', str(row.get('task_description') or ''), 2))" + - "lambda row: bool(__import__('re').search(r'dish|plate|bowl|cup|mug', str(row.get('task_description') or ''), 2))" + +# 500 x 100 steps on one full node: 32 robot + 32 human samples per GPU per step +trainer: + max_epochs: 500 + min_epochs: 500 + +launch_params: + gpus_per_node: 8 diff --git a/egomimic/hydra_configs/train_zarr_abc_sim_rdt_cotrain.yaml b/egomimic/hydra_configs/train_zarr_abc_sim_rdt_cotrain.yaml new file mode 100644 index 0000000000..4adcc0b0b3 --- /dev/null +++ b/egomimic/hydra_configs/train_zarr_abc_sim_rdt_cotrain.yaml @@ -0,0 +1,15 @@ +# Single-task cotrain: one abc_sim task (robot, joint space) + one mecka task +# (human, wrist-frame 6D), RDT-1B with both heads, one batch of each per step. +# Ryan's suggested starting point for the ABC tests (2026-09-27): single-task +# cotrain before the all-mecka base. +# +# python egomimic/trainHydra.py --config-name=train_zarr_abc_sim_rdt_cotrain \ +# name=abc_sim_cotrain description=put_bottles_cup_on_saucer +defaults: + - train_zarr_abc_sim_rdt_ft + - override data: abc_sim_joints_cotrain_mecka + - _self_ + +model: + robomimic_model: + domains: ["human_bimanual", "eva_bimanual"] diff --git a/egomimic/hydra_configs/train_zarr_abc_sim_rdt_ft.yaml b/egomimic/hydra_configs/train_zarr_abc_sim_rdt_ft.yaml new file mode 100644 index 0000000000..ab10c17693 --- /dev/null +++ b/egomimic/hydra_configs/train_zarr_abc_sim_rdt_ft.yaml @@ -0,0 +1,32 @@ +# ABC sim fine-tune: RDT-1B, joint space, one abc_sim task, weights-only init +# from a base; init_weights_ckpt=null is the from-scratch floor. Pipeline: +# egomimic/scripts/abc_sim/README.md. +# +# python egomimic/trainHydra.py --config-name=train_zarr_abc_sim_rdt_ft \ +# init_weights_ckpt= name=abc_sim_ft description=_put_bottles +defaults: + - train_zarr_cartesian + - override model: rdt_1b_abc_joints_dinov3_vitb + - override data: abc_sim_joints + - override evaluator: eval_hpt_abc_joints + - _self_ + +model: + robomimic_model: + domains: ["eva_bimanual"] + +# the train_viz head would need its own evaluator wiring; the sim rollout is +# this pipeline's check +train_viz: false + +trainer: + gradient_clip_val: 1.0 + max_epochs: 300 + min_epochs: 300 + limit_train_batches: 100 + limit_val_batches: 20 + check_val_every_n_epoch: 25 + +launch_params: + gpus_per_node: 1 + nodes: 1 diff --git a/egomimic/hydra_configs/train_zarr_cartesian.yaml b/egomimic/hydra_configs/train_zarr_cartesian.yaml index 7ccfe6bc95..1fd804f923 100644 --- a/egomimic/hydra_configs/train_zarr_cartesian.yaml +++ b/egomimic/hydra_configs/train_zarr_cartesian.yaml @@ -13,6 +13,10 @@ defaults: name: test description: test ckpt_path: null +# Weights-only fine-tune init from a Lightning checkpoint (fresh optimizer, +# schedule and epoch; name+shape-matched tensors only, so new embodiment +# stems / heads start fresh). Ignored when ckpt_path resumes a run. +init_weights_ckpt: null mode: train hydra: diff --git a/egomimic/hydra_configs/train_zarr_fold_ladder_rdt1b.yaml b/egomimic/hydra_configs/train_zarr_fold_ladder_rdt1b.yaml new file mode 100644 index 0000000000..3577c0d002 --- /dev/null +++ b/egomimic/hydra_configs/train_zarr_fold_ladder_rdt1b.yaml @@ -0,0 +1,45 @@ +# Freeform fold-clothes OPERATOR ladder, RDT-1B: the RDT recipe at 1B on one +# rung of data/mecka_fold_freeform_opscale{1,2,4}_6d; 1000 x 100 steps at global +# batch 64 (32 x 2 GPUs), val every 25 epochs, a checkpoint every 100 (the bases +# the ABC-sim fine-tunes start from). Norm stats per rung, at startup. +# +# python egomimic/trainHydra.py --config-name=train_zarr_fold_ladder_rdt1b \ +# data=mecka_fold_freeform_opscale1_6d name=fold_ladder description=k1 +defaults: + - train_zarr_mecka_flagship_6d_rdt + - override model: rdt_1b_bc_mecka_6d_dinov3_vitb + - override data: mecka_fold_freeform_opscale4_6d + - _self_ + +launch_params: + gpus_per_node: 2 + nodes: 1 + +data: + train_dataloader_params: + human_bimanual: + batch_size: 32 + num_workers: 12 + valid_dataloader_params: + human_bimanual: + batch_size: 32 + +trainer: + max_epochs: 1000 + min_epochs: 1000 + check_val_every_n_epoch: 25 + limit_train_batches: 100 + +callbacks: + model_checkpoint: + every_n_epochs: 100 + +# The GT/pred replay needs the true frames per chunk step: the mecka cartesian +# chunk is 30 raw frames resampled to 100 steps (29/99), and the replay's own +# stride inference only fits keypoint layouts (it falls back to 1: a 3.4x-slow +# replay that drifts off the hand). The train_viz head composes its own copy. +evaluator: + action_stride: 0.29292929292929293 +train_viz_evaluator: + base: + action_stride: 0.29292929292929293 diff --git a/egomimic/rldb/embodiment/embodiment.py b/egomimic/rldb/embodiment/embodiment.py index cd4a2e30a2..581f74a34e 100644 --- a/egomimic/rldb/embodiment/embodiment.py +++ b/egomimic/rldb/embodiment/embodiment.py @@ -63,17 +63,22 @@ def get_embodiment(index): # embodiment metadata (e.g. MECKA_BIMANUAL, SCALE_LEFT_ARM). Locally all human # demonstration data is ONE embodiment (see the EMBODIMENT docstring; the # source lives only in the SQL `lab` field), so those names collapse to -# HUMAN_*. Robot names (EVA_*) are never aliased. +# HUMAN_*. The one robot is Eva: YAM_* (the ABC station's label since the +# 2026-09-01 SQL relabel, also in the zarrs re-exported to R2) collapses to EVA_*. HUMAN_VENDOR_PREFIXES = ("MECKA", "SCALE", "ARIA", "LIGHTWHEEL") +ROBOT_VENDOR_PREFIXES = ("YAM",) def canonical_embodiment_name(embodiment_name: str) -> str: """Upper-case EMBODIMENT member name, with legacy vendor prefixes - (``MECKA_BIMANUAL`` ...) collapsed onto ``HUMAN_*``. Does not validate.""" + (``MECKA_BIMANUAL`` ...) collapsed onto ``HUMAN_*`` and ``YAM_*`` onto + ``EVA_*``. Does not validate.""" name = embodiment_name.upper() vendor, _, suffix = name.partition("_") if vendor in HUMAN_VENDOR_PREFIXES and suffix: return f"HUMAN_{suffix}" + if vendor in ROBOT_VENDOR_PREFIXES and suffix: + return f"EVA_{suffix}" return name @@ -224,9 +229,26 @@ def viz( ) @classmethod - def get_keymap(cls, keymap_mode: str, norm_mode: bool = False, annotation_key=None): - """Returns a dictionary mapping from the raw keys in the dataset to the canonical keys used by the model.""" + def get_keymap( + cls, + keymap_mode: str, + norm_mode: bool = False, + annotation_key=None, + image_history_gap_s: float | None = None, + ): + """Returns a dictionary mapping from the raw keys in the dataset to the canonical keys used by the model. + + ``image_history_gap_s`` adds ``_hist``, the front frame that + many seconds earlier (RDT's ``trunk.image_history: 2``); the same knob + as ``Human.get_keymap``, so a robot domain can share an RDT trunk with + a human one.""" key_map = cls._get_keymap(keymap_mode) + if image_history_gap_s is not None: + front = key_map[cls.VIZ_IMAGE_KEY] + key_map[f"{cls.VIZ_IMAGE_KEY}{IMAGE_HISTORY_SUFFIX}"] = { + **front, + "lag_s": float(image_history_gap_s), + } if annotation_key is not None and not norm_mode: key_map[annotation_key] = { "key_type": "annotation_keys", diff --git a/egomimic/rldb/embodiment/eva.py b/egomimic/rldb/embodiment/eva.py index b3edfa4f48..00e4cad995 100644 --- a/egomimic/rldb/embodiment/eva.py +++ b/egomimic/rldb/embodiment/eva.py @@ -23,6 +23,7 @@ NumpyToTensor, PoseCoordinateFrameTransform, QuaternionPoseToYPR, + ResizeImages, SplitKeys, Transform, XYZWXYZ_to_XYZYPR, @@ -32,6 +33,16 @@ ) +# Joint-space mode: the ABC YAM station / abc_sim action space. +JOINT_MODE = "joints" +JOINT_ACTION_KEY = "actions_joints" +JOINT_STATE_KEY = "observations.state.joint_positions" +# Raw 30 Hz steps read per sample (1.5 s, the cartesian keymap's horizon), +# resampled to the model's chunk length like the cartesian modes. +JOINT_RAW_HORIZON = 45 +JOINT_ORDER = ("left.{}_joints", "left.{}_gripper", "right.{}_joints", "right.{}_gripper") + + class Eva(Embodiment): INTRINSICS = ARIA_INTRINSICS EXTRINSICS = { @@ -62,16 +73,34 @@ def get_transform_list( "cartesian_wristframe_ypr", "cartesian_wristframe_6d", "cartesian_wristframe_quat", + "joints", ], allow_legacy_rotation: bool = False, + image_hw: tuple[int, int] | None = None, ) -> list[Transform]: - """Transform pipeline for the requested mode. + """Transform pipeline for the requested mode, then (``image_hw``) every + camera resized to ``(H, W)`` -- for sources mixing resolutions. ``allow_legacy_rotation`` opts out of the continuous-rotation rule for the modes in ``LEGACY_ROTATION_MODES``; only the data/rollout boundary (viz, a checkpoint that predates the 6D conversion) may set it. """ _reject_legacy_rotation(cls, mode, allow_legacy_rotation) + transforms = cls._mode_transform_list(mode) + if image_hw is not None: + transforms = [ResizeImages(cls.CAMERA_KEYS_ALL, tuple(image_hw))] + transforms + return transforms + + # Every camera entry the keymaps can emit, incl. the image-history twin. + CAMERA_KEYS_ALL = [ + "observations.images.front_img_1", + "observations.images.front_img_1_hist", + "observations.images.left_wrist_img", + "observations.images.right_wrist_img", + ] + + @classmethod + def _mode_transform_list(cls, mode: str) -> list[Transform]: if mode == "cartesian": return _build_eva_bimanual_transform_list(is_quat=True) if mode == "cartesian_6d": @@ -97,6 +126,13 @@ def get_transform_list( ] if mode == "cartesian_wristframe_quat": return _build_eva_bimanual_eef_frame_transform_list(is_quat=True) + if mode == JOINT_MODE: + # Absolute joint positions + gripper, 14-D per step, in the + # [left j1..j6, left grip, right j1..j6, right grip] order of the + # ABC YAM station and its simulator (abc_sim's 14-D state/action). + # Needs the "joints" keymap (obs_/cmd_joints are not in the + # cartesian one). No frame, no rotation: the sim steps these as-is. + return _build_eva_bimanual_joint_transform_list() raise ValueError(f"Unsupported transform_list mode '{mode}' for Eva") @staticmethod @@ -105,7 +141,8 @@ def get_revert_transform_list(mode: str) -> list[Transform] | None: trained with ``get_transform_list(mode)`` back to camera-frame xyz+ypr+gripper (7/arm), given the batch's proprio ``observations.state.ee_pose``. ``None``: already in that layout.""" - if mode == "cartesian": + if mode in ("cartesian", JOINT_MODE): + # joints: predictions are already the commanded joint targets. return None if mode == "cartesian_6d": return _build_eva_cartesian_revert_6d_transform_list() @@ -120,6 +157,14 @@ def _get_keymap(cls, keymap_mode: str): # One camera naming for every algo. Pi renames onto openpi's slots # itself (egomimic.models.preprocess_pi_obs.PI_CAMERA_SLOTS). keymap_mode = _strip_pi_keymap_mode(cls, keymap_mode) + if keymap_mode == JOINT_MODE: + return cls._get_joint_keymap() + if keymap_mode == f"{JOINT_MODE}_front": + # front camera only: skip decoding wrist frames the model never reads + key_map = cls._get_joint_keymap() + for key in ("left_wrist_img", "right_wrist_img"): + del key_map[f"observations.images.{key}"] + return key_map front_key = cls.VIZ_IMAGE_KEY right_wrist_key = "observations.images.right_wrist_img" left_wrist_key = "observations.images.left_wrist_img" @@ -177,6 +222,36 @@ def _get_keymap(cls, keymap_mode: str): return key_map + @classmethod + def _get_joint_keymap(cls): + """Cameras + per-arm joint positions and gripper, observed (proprio) + and commanded (45-step action chunk). ABC zarrs (real and the + converted sim_224 episodes) carry ``{left,right}.{obs,cmd}_joints`` + (6) and ``_gripper`` (1, in [0, 1]).""" + key_map = { + cls.VIZ_IMAGE_KEY: {"key_type": "camera_keys", "zarr_key": "images.front_1"}, + "observations.images.right_wrist_img": { + "key_type": "camera_keys", + "zarr_key": "images.right_wrist", + }, + "observations.images.left_wrist_img": { + "key_type": "camera_keys", + "zarr_key": "images.left_wrist", + }, + } + for arm in ("left", "right"): + for part in ("joints", "gripper"): + key_map[f"{arm}.obs_{part}"] = { + "key_type": "proprio_keys", + "zarr_key": f"{arm}.obs_{part}", + } + key_map[f"{arm}.cmd_{part}"] = { + "key_type": "action_keys", + "zarr_key": f"{arm}.cmd_{part}", + "horizon": JOINT_RAW_HORIZON, + } + return key_map + @classmethod def dinov3_keymap(cls): """ @@ -615,3 +690,19 @@ def _build_eva_bimanual_transform_list( ] ) return transform_list + + +def _build_eva_bimanual_joint_transform_list() -> list[Transform]: + """Joint-space pipeline: resample each commanded joint / gripper chunk to + 100 steps (as the cartesian modes do), then concatenate both arms in ABC + order into ``actions_joints`` (T, 14) and the observed joints into + ``observations.state.joint_positions`` (14,).""" + cmd_keys = [k.format("cmd") for k in JOINT_ORDER] + obs_keys = [k.format("obs") for k in JOINT_ORDER] + return [ + *(InterpolateLinear(new_chunk_length=100, action_key=k, output_action_key=k) for k in cmd_keys), + InterpolatePadMask(new_chunk_length=100), + ConcatKeys(key_list=cmd_keys, new_key_name=JOINT_ACTION_KEY, delete_old_keys=True), + ConcatKeys(key_list=obs_keys, new_key_name=JOINT_STATE_KEY, delete_old_keys=True), + NumpyToTensor(keys=[JOINT_ACTION_KEY, JOINT_STATE_KEY]), + ] diff --git a/egomimic/rldb/zarr/action_chunk_transforms.py b/egomimic/rldb/zarr/action_chunk_transforms.py index 9150df551b..ac8c263b6e 100644 --- a/egomimic/rldb/zarr/action_chunk_transforms.py +++ b/egomimic/rldb/zarr/action_chunk_transforms.py @@ -963,6 +963,43 @@ def transform(self, batch: dict) -> dict: # --------------------------------------------------------------------------- +class ResizeImages(Transform): + """Resize every camera frame in ``keys`` to ``(H, W)``. Frames are the + decoded ``(3, H, W)`` float arrays ZarrDataset emits (or ``(K, 3, H, W)`` + windows); absent keys are skipped. For sources that mix resolutions (real + ABC: 640x480 and 1280x720 front frames), which default_collate cannot + stack; a model's own stem still resizes afterwards.""" + + def __init__(self, keys: list[str], hw: tuple[int, int]): + self.keys = list(keys) + self.hw = (int(hw[0]), int(hw[1])) + + def _resize(self, chw: np.ndarray) -> np.ndarray: + import cv2 + + if tuple(chw.shape[-2:]) == self.hw: + return chw + hwc = np.moveaxis(np.asarray(chw, dtype=np.float32), 0, -1) + out = cv2.resize(hwc, (self.hw[1], self.hw[0]), interpolation=cv2.INTER_AREA) + return np.moveaxis(out, -1, 0) + + def transform(self, batch: dict) -> dict: + for key in self.keys: + if key not in batch: + continue + x = batch[key] + if isinstance(x, torch.Tensor): + x = x.numpy() + x = np.asarray(x) + if x.ndim == 3: + batch[key] = self._resize(x) + elif x.ndim == 4: + batch[key] = np.stack([self._resize(f) for f in x]) + else: + raise ValueError(f"ResizeImages: '{key}' has shape {x.shape}") + return batch + + class NumpyToTensor(Transform): def __init__(self, keys: list[str]): self.keys = keys diff --git a/egomimic/rldb/zarr/zarr_dataset_multi.py b/egomimic/rldb/zarr/zarr_dataset_multi.py index 85c23c8d0c..46ab431598 100644 --- a/egomimic/rldb/zarr/zarr_dataset_multi.py +++ b/egomimic/rldb/zarr/zarr_dataset_multi.py @@ -88,6 +88,11 @@ def episode_table_to_df(engine): def _episode_table() -> pd.DataFrame: + # EGOVERSE_EPISODE_TABLE=: clusters without + # SQL access (Lambda) resolve from a snapshot exported where SQL exists. + path = os.environ.get("EGOVERSE_EPISODE_TABLE") + if path: + return memoized(("episode_table", path), lambda: pd.read_pickle(path)) return memoized( ("episode_table",), lambda: episode_table_to_df(create_default_engine()) ) diff --git a/egomimic/scripts/abc_sim/README.md b/egomimic/scripts/abc_sim/README.md new file mode 100644 index 0000000000..8c2fa59d7f --- /dev/null +++ b/egomimic/scripts/abc_sim/README.md @@ -0,0 +1,33 @@ +# ABC fine-tune + abc_sim rollout pipeline + +Rank pretrained EgoVerse policies by fine-tuning each on one ABC task and +rolling the fine-tune out in ABC's simulator (amazon-far/abc, `abc_sim`): +closed-loop success where offline validation metrics all look alike. The ABC +paper reports Pearson r = 0.85 (success) / 0.91 (progress) between its sim and +real success over 12 checkpoints on bottles, dishrack and mug flip. + +Action space: the ABC station's 14-D absolute joints + gripper (Eva mode +`joints`, `egomimic/rldb/embodiment/eva.py`), which is also abc_sim's action, +so a rollout needs no IK. A task is named by its sim_224 `task_name` (e.g. +`sim_put_the_plastic_bottles_in_the_bin`), which is also an abc_sim alias. + +Steps (Lambda; paths in the launchers): + +1. Data: `prepare.py --sim-data ` or `curl` the sim_224 tar, then + `convert_sim_to_zarr.py --src --task --out + $ABC_SIM_ZARR_ROOT/put_bottles` -> eva_bimanual zarrs with the real ABC keys. +2. Base: `train_zarr_fold_ladder_rdt1b` (RDT-1B on the mecka freeform-fold + operator ladder) via `lambda_pack.sbatch`; every 100 epochs a checkpoint. +3. Fine-tune: `train_zarr_abc_sim_rdt_ft init_weights_ckpt=` + (same launcher); `init_weights_ckpt=null` is the from-scratch floor; + `train_zarr_abc_real_rdt_ft` the real-data twin; `train_zarr_abc_sim_rdt_cotrain` + the single-task human+robot cotrain. +4. Rollout: `rollout_lambda.sbatch` = `policy_server.py` (training venv) + + `sim_client.py` (ABC venv, ABC's own `rollout_worlds` + `build_summary`, + no RTC) -> ABC's `summary.json`. The server resamples the model's 100-step + chunk back onto the 45 raw 30 Hz frames it spans, so the client's 15 + executed steps are 0.5 s of the plan, as in ABC; the client renders the + history frame `lag` steps before each re-plan. `POLICY=hold` is the + no-motion floor. +5. `val_vs_sim.py manifest.json --metric ` -> Spearman / Pearson with + bootstrap CIs and the log-linear fit of the ladder. diff --git a/egomimic/scripts/abc_sim/__init__.py b/egomimic/scripts/abc_sim/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/egomimic/scripts/abc_sim/convert_sim_to_zarr.py b/egomimic/scripts/abc_sim/convert_sim_to_zarr.py new file mode 100644 index 0000000000..c0cab4d70b --- /dev/null +++ b/egomimic/scripts/abc_sim/convert_sim_to_zarr.py @@ -0,0 +1,242 @@ +"""Convert ABC sim_224 teleop episodes (amazon-far/abc) to EgoVerse zarrs. + +ABC's ``prepare.py --sim-data `` downloads a task tar and extracts every +episode into ``$ABC_CACHE/{train,val}_sim//`` (a plain ``tar -x`` of +the same tar gives ``{train,val}//``; both are read): ``states_actions.bin`` +((T, 28) float64: 14 recorded + 14 commanded dofs, [left j1..j6, left grip, +right j1..j6, right grip]), ``combined_camera-images-rgb.mp4`` (the cameras +stacked vertically, 224x168 each, 30 fps) and ``episode_metadata.json``. All +tasks share those two folders, so episodes are selected by their +``task_name``. + +Each episode becomes one ``eva_bimanual`` zarr with the same keys as the real +ABC zarrs on the Phoenix mirror (``{left,right}.{obs,cmd}_{joints,gripper}``, +``images.{front_1,left_wrist,right_wrist}``, ``annotations``), so +``Eva.get_keymap("joints")`` reads real and sim ABC data alike. Camera map: +top -> front_1, left -> left_wrist, right -> right_wrist. Extra metadata: +``lab: abc_sim`` and ``split`` (ABC's own train/val). + +Prompts follow ABC's training rule (``abc_minimal.dataloader``) verbatim: the +episode's ``prompt_timeline`` if it has one, else its ``instruction`` when that +differs from the task name, else the task name with ``_`` -> `` `` (the sim +tasks' names start with ``sim_``). Episodes ABC's loader drops as unlabelled +are skipped. + + python -m egomimic.scripts.abc_sim.convert_sim_to_zarr \ + --src $ABC_CACHE --task sim_put_the_plastic_bottles_in_the_bin \ + --out /storage/project/r-dxu345-0/agao81/abc_sim/zarr/put_bottles +""" + +from __future__ import annotations + +import argparse +import json +import logging +import math +import os +import shutil +import sys +from concurrent.futures import ProcessPoolExecutor, as_completed +from pathlib import Path + +import numpy as np + +log = logging.getLogger("abc_sim.convert") + +STATE_DIM = 14 +DATA_FPS = 30 +DEFAULT_CAMERAS = ("top", "left", "right") +CAMERA_TO_ZARR = { + "top": "images.front_1", + "left": "images.left_wrist", + "right": "images.right_wrist", +} +# abc_sim scenes: every policy camera is a fovy=58 deg pinhole. +CAMERA_FOVY_DEG = 58.0 +DONE_MARKER = ".abc_sim_done" + + +def task_name_to_prompt(task_name: str) -> str: + """abc_minimal.dit.task_name_to_prompt.""" + return " ".join(task_name.replace("-", " ").replace("_", " ").split()) + + +def prompt_spans(meta: dict, num_steps: int) -> list[tuple[str, int, int]]: + """``(text, start, end)`` spans covering ``[0, num_steps)``, end exclusive.""" + task_name = str(meta.get("task_name") or "") + timeline = meta.get("prompt_timeline") + if timeline: + entries = sorted((int(e["frame"]), str(e["prompt"])) for e in timeline) + else: + instruction = meta.get("instruction") + if ( + isinstance(instruction, str) + and instruction.strip() + and task_name_to_prompt(instruction) != task_name_to_prompt(task_name) + ): + text = instruction.strip() + else: + text = task_name_to_prompt(task_name) + entries = [(0, text)] + spans = [] + for i, (start, text) in enumerate(entries): + end = entries[i + 1][0] if i + 1 < len(entries) else num_steps + start = max(0, min(start, num_steps)) + if i == 0: + start = 0 # the first directive holds from the episode start + if end > start: + spans.append((text, start, end)) + return spans + + +def pinhole_intrinsics(height: int, width: int, fovy_deg: float = CAMERA_FOVY_DEG): + """3x4 K of a MuJoCo camera (square pixels, principal point centred).""" + f = 0.5 * height / math.tan(math.radians(fovy_deg) / 2.0) + return np.array( + [[f, 0.0, width / 2.0, 0.0], [0.0, f, height / 2.0, 0.0], [0.0, 0.0, 1.0, 0.0]], + dtype=np.float64, + ) + + +def load_states_actions(ep_dir: Path) -> tuple[np.ndarray, np.ndarray]: + raw = np.fromfile(ep_dir / "states_actions.bin", dtype=np.float64) + if raw.size == 0 or raw.size % (2 * STATE_DIM): + raise ValueError(f"states_actions.bin holds {raw.size} float64s") + table = raw.reshape(-1, 2 * STATE_DIM) + return table[:, :STATE_DIM], table[:, STATE_DIM:] + + +def decode_cameras(ep_dir: Path, cameras: tuple[str, ...]) -> dict[str, np.ndarray]: + """``{camera: (T, H, W, 3) uint8}`` from the vertically stacked mp4.""" + import av + + with av.open(str(ep_dir / "combined_camera-images-rgb.mp4")) as container: + stack = np.stack([f.to_ndarray(format="rgb24") for f in container.decode(video=0)]) + h = stack.shape[1] // len(cameras) + if h * len(cameras) != stack.shape[1]: + raise ValueError(f"video height {stack.shape[1]} is not {len(cameras)} stacked cameras") + return {cam: stack[:, i * h : (i + 1) * h] for i, cam in enumerate(cameras)} + + +def split_arms(x: np.ndarray, kind: str) -> dict[str, np.ndarray]: + x = np.asarray(x, dtype=np.float32) + return { + f"left.{kind}_joints": x[:, 0:6], + f"left.{kind}_gripper": x[:, 6:7], + f"right.{kind}_joints": x[:, 7:13], + f"right.{kind}_gripper": x[:, 13:14], + } + + +def convert_episode(ep_dir: Path, split: str, out_dir: Path) -> dict: + """Write one zarr; returns a report row. Idempotent via a done marker.""" + from egomimic.rldb.zarr.zarr_writer import ZarrWriter + + name = f"abcsim_{split}_{ep_dir.name}" + dst = out_dir / f"{name}.zarr" + if (dst / DONE_MARKER).exists(): + return {"episode": name, "status": "skipped"} + meta = json.loads((ep_dir / "episode_metadata.json").read_text()) + if not meta.get("prompt_timeline") and (meta.get("prompt_source") or {}).get("status") == "excluded": + return {"episode": name, "status": "skipped_unlabelled"} # abc_minimal.dataloader drops these + states, actions = load_states_actions(ep_dir) + num_steps = len(states) + if meta.get("num_steps") is not None and int(meta["num_steps"]) != num_steps: + raise ValueError(f"states_actions.bin has {num_steps} rows, metadata says {meta['num_steps']}") + cameras = tuple(meta.get("cameras") or DEFAULT_CAMERAS) + images = decode_cameras(ep_dir, cameras) + n_frames = min(len(v) for v in images.values()) + if n_frames != num_steps: + raise ValueError(f"{n_frames} video frames for {num_steps} state rows") + missing = [c for c in DEFAULT_CAMERAS if c not in images] + if missing: + raise ValueError(f"cameras missing: {missing} (have {sorted(images)})") + h, w = images["top"].shape[1:3] + + tmp = out_dir / f".{name}.zarr.tmp" + shutil.rmtree(tmp, ignore_errors=True) + spans = prompt_spans(meta, num_steps) + ZarrWriter.create_and_write( + tmp, + numeric_data={**split_arms(states, "obs"), **split_arms(actions, "cmd")}, + image_data={CAMERA_TO_ZARR[c]: images[c] for c in DEFAULT_CAMERAS}, + embodiment="eva_bimanual", + fps=DATA_FPS, + task_name=str(meta.get("task_name") or ""), + task_description=spans[0][0] if spans else "", + annotations=spans, + intrinsics={"front_1": pinhole_intrinsics(h, w)}, + metadata_override={"lab": "abc_sim", "split": split}, + ) + (tmp / DONE_MARKER).write_text("ok\n") + shutil.rmtree(dst, ignore_errors=True) + os.replace(tmp, dst) + return {"episode": name, "status": "written", "frames": num_steps} + + +def find_episodes(src: Path, task: str, splits: list[str]) -> list[tuple[Path, str]]: + found = [] + for split in splits: + # prepare.py extracts into {split}_sim/; a plain `tar -x` leaves {split}/ + root = next((src / d for d in (f"{split}_sim", split) if (src / d).is_dir()), None) + if root is None: + log.warning("no %s_sim/ or %s/ under %s", split, split, src) + continue + for bin_path in sorted(root.glob("*/states_actions.bin")): + ep_dir = bin_path.parent + meta_path = ep_dir / "episode_metadata.json" + if not meta_path.exists(): + continue + if json.loads(meta_path.read_text()).get("task_name") == task: + found.append((ep_dir, split)) + return found + + +def main(argv=None) -> int: + p = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + p.add_argument("--src", type=Path, required=True, help="ABC cache dir ($ABC_CACHE)") + p.add_argument("--task", required=True, help="dataset task_name, e.g. sim_pouring_beads") + p.add_argument("--out", type=Path, required=True) + p.add_argument("--workers", type=int, default=max(1, (os.cpu_count() or 2) // 2)) + args = p.parse_args(argv) + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s") + + episodes = find_episodes(args.src, args.task, ["train", "val"]) + if not episodes: + log.error("no %s episodes under %s", args.task, args.src) + return 1 + args.out.mkdir(parents=True, exist_ok=True) + log.info("converting %d episodes of %s -> %s", len(episodes), args.task, args.out) + + rows, failures = [], [] + with ProcessPoolExecutor(max_workers=args.workers) as pool: + futures = { + pool.submit(convert_episode, ep, split, args.out): ep + for ep, split in episodes + } + for i, fut in enumerate(as_completed(futures), 1): + ep = futures[fut] + try: + rows.append(fut.result()) + except Exception as e: # one bad episode must not sink the task + failures.append({"episode": str(ep), "error": f"{type(e).__name__}: {e}"}) + log.warning("FAILED %s: %s", ep, e) + if i % 100 == 0 or i == len(futures): + log.info("%d/%d done (%d failed)", i, len(futures), len(failures)) + + report = { + "task": args.task, + "src": str(args.src), + "written": sum(r["status"] == "written" for r in rows), + "skipped": sum(r["status"] == "skipped" for r in rows), + "skipped_unlabelled": sum(r["status"] == "skipped_unlabelled" for r in rows), + "failed": failures, + "frames_written": int(sum(r.get("frames", 0) for r in rows)), + } + (args.out / "conversion_report.json").write_text(json.dumps(report, indent=2)) + log.info("report: %s", {k: v for k, v in report.items() if k != "failed"}) + return 0 if not failures else 2 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/egomimic/scripts/abc_sim/lambda_pack.sbatch b/egomimic/scripts/abc_sim/lambda_pack.sbatch new file mode 100644 index 0000000000..0403c33aa1 --- /dev/null +++ b/egomimic/scripts/abc_sim/lambda_pack.sbatch @@ -0,0 +1,66 @@ +#!/bin/bash +# Pack several small-DDP main-EgoVerse runs on ONE exclusive Lambda loaner node +# (8 x H100), one srun step per run inside the pytorch image, each on its own +# GPUs / DDP port, resuming from its last.ckpt on requeue. The dataset dir is +# copied to node-local disk first (/tmp; the /ephemeral array is root-only on +# the loaner nodes): zarr reads straight off the NFS /workspace stalled nodes +# and left cancelled jobs undead (Ryan, 2026-09-26). +# +# RUNS="k1 k2 k4" ARGS_k1="--config-name=train_zarr_fold_ladder_rdt1b data=mecka_fold_freeform_opscale1_6d" ... \ +# sbatch --qos=loaner-normal egomimic/scripts/abc_sim/lambda_pack.sbatch +# GPUS_PER_RUN=2 must match launch_params.gpus_per_node of the configs +# DATA_SRC=... NFS dataset dir (default datasets/fold_ladder), staged to node-local /tmp +# TABLE=... EGOVERSE_EPISODE_TABLE pickle (default: the one in DATA_SRC) +# RUN_ROOT=... default /workspace/users/$USER/runs_abc +# EXTRA="..." appended to every run (smokes: trainer.max_epochs=1 ...) +# Onto a RUNNING job's spare GPUs: SLURM_JOB_ID= bash egomimic/scripts/abc_sim/lambda_pack.sbatch +# (steps --overlap its CPUs; they die with that job, so queue an afterany follow-up to resume). +#SBATCH --job-name=abc_pack +#SBATCH --account=loaner +#SBATCH --partition=loaner +#SBATCH --nodes=1 +#SBATCH --gres=gpu:h100:8 +#SBATCH --mem-per-gpu=225G # the partition's DefMemPerGPU (256000 MB) x 8 exceeds a node +#SBATCH --cpus-per-task=224 +#SBATCH --exclusive +#SBATCH --time=2-00:00:00 +#SBATCH --requeue +#SBATCH --output=/workspace/users/%u/logs/%x-%j.out +set -euo pipefail +SRC=${SRC:-/workspace/users/$USER/EgoVerse} +IMAGE=${IMAGE:-/workspace/users/$USER/images/nvidia-pytorch-25.06-py3.sqsh} +RUN_ROOT=${RUN_ROOT:-/workspace/users/$USER/runs_abc} +DATA_SRC=${DATA_SRC:-/workspace/users/$USER/datasets/fold_ladder} +TABLE=${TABLE:-$(ls "$DATA_SRC"/episode_table_*.pkl | tail -1)} +DATA=/tmp/$USER/$(basename "$DATA_SRC") +G=${GPUS_PER_RUN:-2} +read -r -a runs <<< "${RUNS:?set RUNS to space-separated run names, each with an ARGS_}" +(( ${#runs[@]} * G <= 8 )) || { echo "${#runs[@]} runs x $G GPUs do not fit one node" >&2; exit 2; } +CPUS=${CPUS_PER_TASK:-$(( 192 / (${#runs[@]} * G) ))} +test -s "$IMAGE" && test -s "$SRC/.venv/pyvenv.cfg" && test -s "$TABLE" +echo "=== abc_pack job=$SLURM_JOB_ID node=$(hostname) restart=${SLURM_RESTART_COUNT:-0} commit=$(git -C "$SRC" rev-parse --short HEAD) runs=${runs[*]}" +# a step, so the copy lands on the compute node also when this file is run with bash +srun --overlap --gres=none --ntasks=1 bash -c "mkdir -p $DATA && rsync -a --exclude _sync $DATA_SRC/ $DATA/ && echo staged \$(du -sh $DATA | cut -f1) to $DATA" + +i=0 +for name in "${runs[@]}"; do + var="ARGS_$name"; args=${!var:?set $var} + port=$(( ${PORT_BASE:-29500} + 10 * i )); i=$(( i + 1 )) + run="$RUN_ROOT/$name"; mkdir -p "$run" + resume=""; [ -s "$run/checkpoints/last.ckpt" ] && resume="ckpt_path=$run/checkpoints/last.ckpt" + echo "--- $name -> $run ${resume:+(resuming)}: $args" + srun --exact --overlap --ntasks=$G --ntasks-per-node=$G --gres=gpu:h100:$G --gpu-bind=none --cpus-per-task=$CPUS --mem=200G \ + --container-image="$IMAGE" --container-mounts=/workspace:/workspace:rw,$DATA:$DATA:ro --container-workdir="$SRC" \ + --output="$run/train-%j.log" \ + bash -c "set -euo pipefail + ulimit -n \$(ulimit -Hn) + export HOME=/workspace/users/$USER HF_HOME=/workspace/users/$USER/.cache/huggingface HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 + export EGOVERSE_DATASET_DIR=$DATA EGOVERSE_EPISODE_TABLE=$TABLE EGOVERSE_LOG_DIR=$RUN_ROOT EGOVERSE_CACHE_DIR=$RUN_ROOT/cache + export MASTER_PORT=$port PYTHONPATH=$SRC HYDRA_FULL_ERROR=1 PYTHONUNBUFFERED=1 WANDB_MODE=\${WANDB_MODE:-online} WANDB_PROJECT=\${WANDB_PROJECT:-egoverse} + export OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 + export TORCHINDUCTOR_CACHE_DIR=/tmp/ti_${SLURM_JOB_ID}_$name TRITON_CACHE_DIR=/tmp/tr_${SLURM_JOB_ID}_$name + echo rank=\${SLURM_PROCID:-0} nofile=\$(ulimit -n) commit=\$(git -C $SRC rev-parse --short HEAD) job=$SLURM_JOB_ID restart=${SLURM_RESTART_COUNT:-0} + exec $SRC/.venv/bin/python egomimic/trainHydra.py $args launch_params.gpus_per_node=$G \ + hydra.run.dir=$run paths.output_dir=$run norm_stats.save_cache_dir=$run $resume ${EXTRA:-}" & +done +wait diff --git a/egomimic/scripts/abc_sim/policy_server.py b/egomimic/scripts/abc_sim/policy_server.py new file mode 100644 index 0000000000..d36690b33d --- /dev/null +++ b/egomimic/scripts/abc_sim/policy_server.py @@ -0,0 +1,166 @@ +"""Serve an EgoVerse joint-space checkpoint over HTTP for the abc_sim client. + +The simulator's Python (amazon-far/abc: py3.12, MuJoCo 3.8, torch 2.11) cannot +share the training venv, so the policy runs here and the sim talks to it over +plain HTTP + JSON (stdlib on both ends). + + python -m egomimic.scripts.abc_sim.policy_server --ckpt --port 8765 + +POST /infer {"state": [14], "prompt": str, + "images": {"top": {"shape": [H, W, 3], "b64": }, + "top_hist": {...}}} # optional: the frame lag_frames back + -> {"actions": [[14] x 45]} +GET /health -> {"ckpt", "cameras", "lag_frames"} + +The input is built exactly as ZarrDataset builds a training sample for the +checkpoint's data config (keymap, transform list, fps) and normalized with +the checkpoint's norm stats: the dataset normalizes in __getitem__, and +process_batch_for_training does not. ``top_hist`` is the frame +``image_history_gap_s`` earlier; without it the current frame stands in, which +is what an episode start looks like in training. The model's 100-step chunk +spans 45 raw 30 Hz frames, so it is resampled back onto that grid: action k +is the command k frames ahead. +""" + +from __future__ import annotations + +import argparse +import base64 +import json +import logging +from http.server import BaseHTTPRequestHandler, HTTPServer + +import numpy as np +import torch +from omegaconf import OmegaConf +from torch.utils.data import default_collate + +from egomimic.rldb.embodiment.embodiment import IMAGE_HISTORY_SUFFIX, get_embodiment_id +from egomimic.rldb.embodiment.eva import JOINT_ACTION_KEY, JOINT_MODE, JOINT_RAW_HORIZON, Eva +from egomimic.scripts.abc_sim.convert_sim_to_zarr import CAMERA_TO_ZARR, split_arms +from egomimic.utils.pose_utils import _interpolate_linear + +log = logging.getLogger("abc_sim.policy_server") + +EMB = "eva_bimanual" +FPS = 30.0 +ZARR_TO_CAMERA = {z: c for c, z in CAMERA_TO_ZARR.items()} + + +def decode_image(spec: dict) -> np.ndarray: + """{"shape": [H, W, 3], "b64": raw uint8} -> (3, H, W) float32 in [0, 1].""" + raw = np.frombuffer(base64.b64decode(spec["b64"]), dtype=np.uint8).reshape(spec["shape"]) + return np.moveaxis(raw, -1, -3).astype(np.float32) / np.float32(255.0) + + +class JointPolicy: + """The checkpoint plus the data pipeline it was trained with.""" + + def __init__(self, ckpt_path: str, device: str = "cuda"): + from egomimic.pl_utils.pl_model import ModelWrapper + + self.ckpt_path = ckpt_path + device = torch.device(device if torch.cuda.is_available() else "cpu") + # The algo is a plain object; the Lightning wrapper owns (and moves) its nets. + self.wrapper = ModelWrapper.load_from_checkpoint(ckpt_path, weights_only=False, map_location="cpu") + self.wrapper.to(device).eval() + self.algo = self.wrapper.model + self.algo.device = device + self.configure(self.wrapper._as_config(self.wrapper.hparams.config_tree)) + + def configure(self, cfg) -> None: + """Keymap + transform list from the checkpoint's data config.""" + resolver = OmegaConf.select(cfg, f"data.train_datasets.{EMB}.resolver") + if resolver is None: + raise ValueError(f"checkpoint config has no data.train_datasets.{EMB}") + km = {k: v for k, v in OmegaConf.to_container(resolver.key_map, resolve=True).items() if k != "_target_"} + tl = {k: v for k, v in OmegaConf.to_container(resolver.transform_list, resolve=True).items() if k != "_target_"} + if tl.get("mode") != JOINT_MODE: + raise ValueError(f"policy_server serves Eva mode '{JOINT_MODE}', checkpoint has {tl.get('mode')!r}") + self.annotation_key = km.pop("annotation_key", None) + self.key_map = Eva.get_keymap(**km) + self.transforms = Eva.get_transform_list(**tl) + self.lag_frames = max(1, round(km["image_history_gap_s"] * FPS)) if km.get("image_history_gap_s") else 0 + self.cameras = sorted({ZARR_TO_CAMERA[v["zarr_key"]] for v in self.key_map.values() if v.get("key_type") == "camera_keys"}) + self.emb_id = get_embodiment_id(EMB) + + def build_sample(self, state, images: dict[str, np.ndarray]) -> dict: + """One un-normalized sample with the dataset's keys. The 45-step cmd + chunk is the current pose tiled: only its shape reaches the model.""" + raw = {k: v[0] for k, v in split_arms(np.asarray(state, np.float32).reshape(1, 14), "obs").items()} + raw.update({k.replace("obs_", "cmd_"): np.tile(v, (JOINT_RAW_HORIZON, 1)) for k, v in list(raw.items())}) + sample = {} + for key, spec in self.key_map.items(): + if spec.get("key_type") == "camera_keys": + cam = ZARR_TO_CAMERA[spec["zarr_key"]] + if cam not in images: + raise ValueError(f"camera '{cam}' missing from the request (have {sorted(images)})") + past = images.get(cam + IMAGE_HISTORY_SUFFIX, images[cam]) + sample[key] = past if key.endswith(IMAGE_HISTORY_SUFFIX) else images[cam] + elif spec.get("key_type") in ("proprio_keys", "action_keys"): + sample[key] = raw[spec["zarr_key"]] + for tf in self.transforms: + sample = tf.transform(sample) + sample = {k: torch.as_tensor(v, dtype=torch.float32) for k, v in sample.items() if isinstance(v, (np.ndarray, torch.Tensor))} + sample["fps"] = torch.tensor(FPS) + return sample + + @torch.no_grad() + def infer(self, state, images: dict[str, np.ndarray], prompt: str) -> np.ndarray: + """(45, 14) absolute joint + gripper commands, one per 30 Hz step.""" + sample = self.algo.norm_stats.normalize(self.build_sample(state, images), self.emb_id) + batch = default_collate([sample]) + if self.annotation_key: + batch[self.annotation_key] = [[prompt]] + preds = self.algo.forward_eval(self.algo.process_batch_for_training({EMB: batch})) + chunk = preds[f"{EMB}_{JOINT_ACTION_KEY}"][0].float().cpu().numpy() + return _interpolate_linear(chunk, JOINT_RAW_HORIZON) + + +def make_handler(policy: JointPolicy): + class Handler(BaseHTTPRequestHandler): + def _send(self, code: int, payload: dict) -> None: + body = json.dumps(payload).encode() + self.send_response(code) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def do_GET(self): + if self.path != "/health": + return self._send(404, {"error": self.path}) + self._send(200, {"ckpt": policy.ckpt_path, "cameras": policy.cameras, "lag_frames": policy.lag_frames}) + + def do_POST(self): + if self.path != "/infer": + return self._send(404, {"error": self.path}) + try: + req = json.loads(self.rfile.read(int(self.headers.get("Content-Length", 0)))) + images = {k: decode_image(v) for k, v in req["images"].items()} + actions = policy.infer(req["state"], images, req.get("prompt", "")) + self._send(200, {"actions": actions.tolist()}) + except Exception as e: # the client must see the reason, not a hang + log.exception("request failed") + self._send(500, {"error": f"{type(e).__name__}: {e}"}) + + def log_message(self, *a): + pass + + return Handler + + +def main(argv=None) -> int: + p = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + p.add_argument("--ckpt", required=True) + p.add_argument("--port", type=int, default=8765) + a = p.parse_args(argv) + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s") + policy = JointPolicy(a.ckpt) + log.info("serving %s on port %d (cameras %s, lag %d frames)", a.ckpt, a.port, policy.cameras, policy.lag_frames) + HTTPServer(("127.0.0.1", a.port), make_handler(policy)).serve_forever() + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/egomimic/scripts/abc_sim/rollout_lambda.sbatch b/egomimic/scripts/abc_sim/rollout_lambda.sbatch new file mode 100644 index 0000000000..eb2783caa8 --- /dev/null +++ b/egomimic/scripts/abc_sim/rollout_lambda.sbatch @@ -0,0 +1,46 @@ +#!/bin/bash +# Roll one checkpoint out in abc_sim on one Lambda H100: policy_server.py in the +# training venv (pytorch image) and sim_client.py in ABC's venv, both in one +# container step so they share the GPU (mjwarp renders on it too). +# +# CKPT= TASK=sim_put_the_plastic_bottles_in_the_bin OUT= sbatch --qos=loaner-normal \ +# egomimic/scripts/abc_sim/rollout_lambda.sbatch +# TASK: any abc_sim name or alias (the sim_224 task_name); NUM_WORLDS=50; +# POLICY=hold for the no-motion floor (no server); CLIENT_EXTRA e.g. --save-video +# Onto a running job: SLURM_JOB_ID= bash egomimic/scripts/abc_sim/rollout_lambda.sbatch +#SBATCH --job-name=abc_rollout +#SBATCH --account=loaner +#SBATCH --partition=loaner +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --gres=gpu:h100:1 +#SBATCH --cpus-per-task=16 +#SBATCH --mem=120G +#SBATCH --time=06:00:00 +#SBATCH --output=/workspace/users/%u/logs/%x-%j.out +set -euo pipefail +SRC=${SRC:-/workspace/users/$USER/EgoVerse} +ABC=${ABC:-/workspace/users/$USER/abc} +IMAGE=${IMAGE:-/workspace/users/$USER/images/nvidia-pytorch-25.06-py3.sqsh} +TASK=${TASK:-sim_put_the_plastic_bottles_in_the_bin} +: "${OUT:?set OUT}" +POLICY=${POLICY:-served} +[ "$POLICY" = hold ] || : "${CKPT:?set CKPT}" +CKPT=${CKPT:-} # hold needs none; set -u would trip on it below +PORT=${PORT:-$(( 8700 + SLURM_JOB_ID % 200 ))} # set PORT per step when several share a job +mkdir -p "$OUT" && OUT=$(cd "$OUT" && pwd) # absolute: the client runs from $ABC +echo "=== abc_rollout job=$SLURM_JOB_ID node=$(hostname) task=$TASK ckpt=${CKPT:-none} policy=$POLICY out=$OUT" +srun --exact --overlap --ntasks=1 --gres=gpu:h100:1 --cpus-per-task=${CPUS_PER_TASK:-16} --mem=120G --container-image="$IMAGE" --container-mounts=/workspace:/workspace:rw --container-workdir="$SRC" \ + bash -c "set -euo pipefail + export HOME=/workspace/users/$USER HF_HOME=/workspace/users/$USER/.cache/huggingface HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 + export PYTHONPATH=$SRC PYTHONUNBUFFERED=1 OMP_NUM_THREADS=4 + trap 'kill \$(jobs -p) 2>/dev/null' EXIT # the step's cleanup does not reap the backgrounded server on the loaner + if [ '$POLICY' = served ]; then + $SRC/.venv/bin/python -m egomimic.scripts.abc_sim.policy_server --ckpt '$CKPT' --port $PORT > '$OUT/server.log' 2>&1 & + for i in \$(seq 1 120); do curl -fsS http://127.0.0.1:$PORT/health > /dev/null 2>&1 && break; sleep 5; done + curl -fsS http://127.0.0.1:$PORT/health || { echo 'server never came up'; tail -50 '$OUT/server.log'; exit 1; } + fi + cd $ABC + MUJOCO_EGL_DEVICE_ID=0 .venv/bin/python $SRC/egomimic/scripts/abc_sim/sim_client.py --task '$TASK' \ + --server http://127.0.0.1:$PORT --policy '$POLICY' --out '$OUT' --num-worlds ${NUM_WORLDS:-50} ${CLIENT_EXTRA:-} + echo '{\"job\": '$SLURM_JOB_ID', \"ckpt\": \"${CKPT:-}\", \"task\": \"$TASK\", \"commit\": \"'\$(git -C $SRC rev-parse --short HEAD)'\"}' > '$OUT/run_info.json'" diff --git a/egomimic/scripts/abc_sim/sim_client.py b/egomimic/scripts/abc_sim/sim_client.py new file mode 100644 index 0000000000..fa61136de1 --- /dev/null +++ b/egomimic/scripts/abc_sim/sim_client.py @@ -0,0 +1,118 @@ +"""Roll a served EgoVerse policy out in abc_sim (amazon-far/abc) through ABC's +own eval loop and summary (abc_minimal.eval_policy.rollout_worlds / +build_summary). Runs in ABC's venv; the policy lives behind policy_server.py. + + /bin/python egomimic/scripts/abc_sim/sim_client.py \ + --task sim_put_the_plastic_bottles_in_the_bin --server http://127.0.0.1:8765 \ + --num-worlds 50 --out + +ABC's SimEvalConfig defaults (seeded worlds seed+i with its unplaceable-seed +retry, 15 executed steps per inference, 236 chunks, early stop on success, +top/left/right at 168x224) with RTC off. ``--task`` is any abc_sim name or +alias -- the sim_224 task_name works. The prompt defaults to the one the +converter trained on (the task name with _ -> space); ``--prompt env`` keeps +the env's own (the prompt-randomized tasks). ``--policy hold`` needs no +server: it repeats the current state, the no-motion floor. +""" + +from __future__ import annotations + +import argparse +import base64 +import json +import urllib.request +from dataclasses import replace +from pathlib import Path + +import numpy as np + +from abc_minimal.config import SimEvalConfig +from abc_minimal.dit import task_name_to_prompt +from abc_minimal.eval_policy import build_summary, resolved_physics, rollout_worlds +from abc_minimal.sim_env import SimTaskEnv, task_prompt + +CAMERAS = ("top", "left", "right") +HORIZON = 45 # the served chunk: 1.5 s of 30 Hz commands + + +class HistEnv(SimTaskEnv): + """Renders the frame ``lag`` steps before each re-plan (ABC's loop only + renders at the re-plan itself), for the policy's 2-frame image history.""" + + def __init__(self, *args, lag: int, execute: int, **kwargs): + super().__init__(*args, **kwargs) + self.capture_at, self.hist, self._since_obs = execute - lag, None, 0 + + def reset(self, seed, options=None): + self.hist, self._since_obs = None, 0 + return super().reset(seed, options) + + def obs(self): + self._since_obs = 0 + return super().obs() + + def step_one(self, action): + super().step_one(action) + self._since_obs += 1 + if self.capture_at > 0 and self._since_obs == self.capture_at: + self.hist = self.render_cameras() + + +class ServedPolicy: + """abc_minimal's SimPolicy.infer surface over HTTP (noise and prefix unused).""" + + def __init__(self, server: str): + self.server, self.env = server.rstrip("/"), None + with urllib.request.urlopen(f"{self.server}/health", timeout=30) as r: + self.health = json.load(r) + + def infer(self, obs, noise=None, action_prefix=None, prefix_length=0) -> np.ndarray: + images = {c: obs["images"][c] for c in self.health["cameras"]} + if self.env is not None and self.env.hist is not None: + images.update({f"{c}_hist": self.env.hist[c] for c in self.health["cameras"]}) + payload = {"state": np.asarray(obs["state"], np.float32).tolist(), "prompt": obs["prompt"], "images": { + k: {"shape": list(hwc.shape), "b64": base64.b64encode(hwc.tobytes()).decode()} + for k, hwc in ((k, np.ascontiguousarray(np.moveaxis(np.asarray(v, np.uint8), 0, -1))) for k, v in images.items())}} + req = urllib.request.Request(f"{self.server}/infer", data=json.dumps(payload).encode(), + headers={"Content-Type": "application/json"}, method="POST") + with urllib.request.urlopen(req, timeout=120) as r: + return np.asarray(json.load(r)["actions"], dtype=np.float32) + + +class HoldPolicy: + health = {"ckpt": "hold", "cameras": [], "lag_frames": 0} + + def infer(self, obs, noise=None, action_prefix=None, prefix_length=0) -> np.ndarray: + return np.tile(np.asarray(obs["state"], np.float32)[None], (HORIZON, 1)) + + +def main(argv=None) -> int: + p = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + p.add_argument("--task", required=True) + p.add_argument("--server", default="http://127.0.0.1:8765") + p.add_argument("--policy", choices=["served", "hold"], default="served") + p.add_argument("--prompt", default=None, help="default: the training prompt; 'env' = the env's own") + p.add_argument("--out", type=Path, required=True) + p.add_argument("--num-worlds", type=int, default=50) + p.add_argument("--save-video", action="store_true") + a = p.parse_args(argv) + a.out.mkdir(parents=True, exist_ok=True) + + policy = HoldPolicy() if a.policy == "hold" else ServedPolicy(a.server) + prompt = task_prompt(a.task) if a.prompt == "env" else (a.prompt or task_name_to_prompt(a.task)) + cfg = SimEvalConfig(checkpoint=str(policy.health["ckpt"]), task=a.task, num_worlds=a.num_worlds, + rtc=False, fast_inference=False, prefix_length=0, save_video=a.save_video, + output_dir=str(a.out), prompt=prompt) + model_config = replace(cfg.model, chunk_length=HORIZON, action_dim=14, camera_keys=CAMERAS) + env = HistEnv(task=a.task, height=cfg.camera_height, width=cfg.camera_width, camera_keys=CAMERAS, + prompt=prompt, camera_backend=cfg.camera_backend, gpu_id=cfg.gpu_id, + lag=int(policy.health["lag_frames"]), execute=cfg.execute_chunk_dim) + policy.env = env + physics = resolved_physics(env) + worlds = rollout_worlds(cfg, policy, env, 0, None, a.out, model_config) + build_summary(config=cfg, ckpt_path=Path(cfg.checkpoint), device="served", worlds=worlds, out_dir=a.out, physics=physics) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/egomimic/scripts/abc_sim/val_vs_sim.py b/egomimic/scripts/abc_sim/val_vs_sim.py new file mode 100644 index 0000000000..62e25b2b6f --- /dev/null +++ b/egomimic/scripts/abc_sim/val_vs_sim.py @@ -0,0 +1,112 @@ +"""Does the offline validation metric predict closed-loop sim success? + +Input: a JSON manifest of checkpoints, each with the W&B run its validation +metrics live in, its epoch, its data hours (the ladder rung) and the +sim_client summary.json of its fine-tune's rollout: + + [{"name": "k4_ep400", "wandb": "rl2-group/egoverse/", "epoch": 400, + "hours": 1.096, "summary": ".../summary.json"}, ...] + + python -m egomimic.scripts.abc_sim.val_vs_sim manifest.json --metric \ + "seen_op_valid/Valid/human_bimanual_actions_cartesian_paired_mse_avg" --out report/ + +Writes report.json + report.md: per-checkpoint rows, Spearman and Pearson +between the val metric and success_rate / mean_max_progress with 2000-sample +bootstrap CIs, and the log-linear fit of the val metric against data hours +(one point per rung at its final epoch) with its slope and R^2. +""" + +from __future__ import annotations + +import argparse +import json +import warnings +from pathlib import Path + +import numpy as np +from scipy import stats + +# constant inputs inside a bootstrap resample give nan by design +warnings.filterwarnings("ignore", category=stats.ConstantInputWarning) + + +def wandb_metric(run_path: str, metric: str, epoch: int) -> float: + """The metric's value at (the last logged step of) ``epoch``.""" + import wandb + + run = wandb.Api().run(run_path) + rows = [r for r in run.scan_history(keys=[metric, "epoch"]) if r.get(metric) is not None] + at = [r for r in rows if r.get("epoch") is not None and int(r["epoch"]) <= epoch] + if not at: + raise ValueError(f"{run_path}: no {metric} at epoch <= {epoch}") + return float(at[-1][metric]) + + +def corr(x, y, kind: str) -> float: + return float((stats.spearmanr if kind == "spearman" else stats.pearsonr)(x, y)[0]) + + +def bootstrap(x, y, kind: str, n: int = 2000, seed: int = 0) -> tuple[float, float]: + """95% percentile interval of ``corr`` over resampled checkpoints.""" + rng = np.random.default_rng(seed) + x, y = np.asarray(x, float), np.asarray(y, float) + vals = [corr(x[i], y[i], kind) for i in (rng.integers(0, len(x), len(x)) for _ in range(n))] + return float(np.nanpercentile(vals, 2.5)), float(np.nanpercentile(vals, 97.5)) + + +def loglinear(hours, metric) -> dict: + """metric = intercept + slope * log10(hours).""" + f = stats.linregress(np.log10(hours), metric) + return {"slope": float(f.slope), "intercept": float(f.intercept), "r2": float(f.rvalue**2)} + + +def main(argv=None) -> int: + p = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + p.add_argument("manifest", type=Path) + p.add_argument("--metric", required=True) + p.add_argument("--out", type=Path, required=True) + a = p.parse_args(argv) + rows = json.loads(a.manifest.read_text()) + for r in rows: + if "val" not in r: # a row may carry its val metric already + r["val"] = wandb_metric(r["wandb"], a.metric, int(r["epoch"])) + s = json.loads(Path(r["summary"]).read_text()) + r["success_rate"], r["mean_max_progress"], r["num_worlds"] = s["success_rate"], s.get("mean_max_progress"), s["num_worlds"] + val = [r["val"] for r in rows] + out = {"metric": a.metric, "rows": rows, "n": len(rows), "correlations": {}} + for target in ("success_rate", "mean_max_progress"): + y = [r[target] for r in rows] + if any(v is None for v in y): + continue + for kind in ("spearman", "pearson"): + lo, hi = bootstrap(val, y, kind) + out["correlations"][f"{kind}_{target}"] = {"r": corr(val, y, kind), "ci95": [lo, hi]} + # log-linear in data: the final epoch of each rung + final = {} + for r in rows: + if "hours" in r and (r["hours"] not in final or r["epoch"] > final[r["hours"]]["epoch"]): + final[r["hours"]] = r + if len(final) >= 2: + hs = sorted(final) + out["loglinear_val_vs_hours"] = {"hours": hs, "val": [final[h]["val"] for h in hs], **loglinear(hs, [final[h]["val"] for h in hs])} + out["loglinear_success_vs_hours"] = {"hours": hs, "success": [final[h]["success_rate"] for h in hs], **loglinear(hs, [final[h]["success_rate"] for h in hs])} + a.out.mkdir(parents=True, exist_ok=True) + (a.out / "report.json").write_text(json.dumps(out, indent=2)) + md = [f"# val metric vs sim success ({len(rows)} checkpoints)", "", f"metric: `{a.metric}`", "", + "| name | hours | epoch | val | success | max progress | worlds |", "|---|---|---|---|---|---|---|"] + for r in rows: + prog = "" if r["mean_max_progress"] is None else f"{r['mean_max_progress']:.3f}" + md.append(f"| {r['name']} | {r.get('hours', '')} | {r['epoch']} | {r['val']:.4g} | {r['success_rate']:.3f} | {prog} | {r['num_worlds']} |") + md += ["", "| correlation | r | 95% CI |", "|---|---|---|"] + md += [f"| {k} | {v['r']:.3f} | [{v['ci95'][0]:.3f}, {v['ci95'][1]:.3f}] |" for k, v in out["correlations"].items()] + for key in ("loglinear_val_vs_hours", "loglinear_success_vs_hours"): + if key in out: + f = out[key] + md += ["", f"{key}: slope {f['slope']:.4g} per decade, R^2 {f['r2']:.3f} over hours {f['hours']}"] + (a.out / "report.md").write_text("\n".join(md) + "\n") + print("\n".join(md)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/egomimic/trainHydra.py b/egomimic/trainHydra.py index cfaff0e529..789a1a3792 100644 --- a/egomimic/trainHydra.py +++ b/egomimic/trainHydra.py @@ -28,7 +28,10 @@ PinError, pinned_episode_subset, ) -from egomimic.utils.checkpoint_utils import load_checkpoint_weights +from egomimic.utils.checkpoint_utils import ( + init_weights_from_checkpoint, + load_checkpoint_weights, +) from egomimic.utils.compile_cache import set_per_job_compile_cache_dir from egomimic.utils.env import load_env from egomimic.utils.gpu_orphans import arm_rank_pdeathsig, reap_orphans @@ -56,7 +59,23 @@ def _build_model_config_tree(cfg: DictConfig) -> DictConfig: and "norm_stats" in model_cfg.robomimic_model ): model_cfg.robomimic_model.norm_stats = None - tree = OmegaConf.create({"model": model_cfg}) + # Each dataset's keymap + transform list, resolved: what a rollout needs to + # rebuild the training inputs (robot/rollout.py, scripts/abc_sim/policy_server.py + # read data.train_datasets..resolver.*). Nothing else of `data`. + pipelines = { + name: { + "resolver": { + k: OmegaConf.to_container(ds.resolver[k], resolve=True) + for k in ("key_map", "transform_list") + if ds.resolver.get(k) is not None + } + } + for name, ds in (cfg.get("data") or {}).get("train_datasets", {}).items() + if ds is not None and ds.get("resolver") is not None + } + tree = OmegaConf.create( + {"model": model_cfg, "data": {"train_datasets": pipelines}} + ) has_weights = OmegaConf.select(tree, _PI_WEIGHT_KEY, default=None) is not None if has_weights and _weights_from_checkpoint(cfg): log.info( @@ -149,6 +168,24 @@ def _weights_from_checkpoint(cfg: DictConfig) -> bool: return bool(ckpt_path) and os.path.isfile(ckpt_path) and not cfg.get("pretrained") +def _apply_init_weights(cfg: DictConfig, model) -> None: + """``init_weights_ckpt``: weights-only fine-tune init (see + ``init_weights_from_checkpoint``). Skipped when ``ckpt_path`` is set -- a + resume (including a Slurm requeue, folded into ``ckpt_path`` by + ``_prepare_checkpoint_resume``) restores this run's own weights, which + must not be replaced by the base checkpoint's.""" + init_ckpt = cfg.get("init_weights_ckpt") + if not init_ckpt: + return + if cfg.get("ckpt_path"): + log.info( + f"init_weights_ckpt={init_ckpt} ignored: resuming from " + f"ckpt_path={cfg.ckpt_path}" + ) + return + init_weights_from_checkpoint(model, init_ckpt) + + def _log_dataset_frame_counts( train_datasets: dict, valid_datasets: dict, @@ -306,32 +343,6 @@ def _unseen_op_valid_datasets(cfg: DictConfig, instantiate) -> dict: } -def _trainer_world_size(cfg: DictConfig) -> int: - """``devices * num_nodes`` as the config declares it, 1 if it cannot say. - - DistributedSampler strides rather than chunks, so W ranks between them walk - the WHOLE split; the per-rank ``limit_val_batches`` window is therefore W - times wider than it looks from one rank. Eval mode runs on one device - (``train`` forces ``devices=1`` after the loaders are built), so it is 1. - """ - legacy = "train" if cfg.get("train") else "eval" if cfg.get("eval") else None - if (cfg.get("mode") or legacy) == "eval": # _resolve_mode, minus its raise - return 1 - devices = cfg.get("trainer", {}).get("devices", 1) - if isinstance(devices, (list, ListConfig)): - n = len(devices) - elif isinstance(devices, int) and devices > 0: - n = devices - else: # "auto", -1: resolved by lightning at runtime, unknown here - log.warning( - f"trainer.devices={devices!r} is not a count; treating the metric " - "loaders as single-rank, which under-scores the val split." - ) - n = 1 - nodes = cfg.get("trainer", {}).get("num_nodes", 1) - return n * (int(nodes) if isinstance(nodes, int) and nodes > 0 else 1) - - def _metric_frames_per_episode( cfg: DictConfig, head: str, @@ -340,14 +351,14 @@ def _metric_frames_per_episode( ) -> int | None: """Frames per episode to keep on this val head, or None for no subsampling. - ``data.metric_frames_per_episode[head]`` is written for ONE rank -- it is - ``floor(limit_val_batches * batch_size / n_episodes)`` -- so it is scaled by - the world size here: the ranks stride through the subsampled set together - and each still reads at most ``limit_val_batches`` batches. Without this a - 4-GPU run scores a quarter of the frames the same config scores on 1 GPU, - which on the seen-val head is fewer than it scored before subsampling - existed at all. EvenStrideDataset keeps a whole episode when K exceeds its - length, so a split that fits entirely is not subsampled. + ``data.metric_frames_per_episode[head]`` is ``floor(limit_val_batches * + batch_size / n_episodes)``, the SAME on every rank: the val heads come back + inside CombinedLoaders, which Lightning does not wrap in a + DistributedSampler (every rank runs every val batch, measured 2026-09-16; + see EvalVideo._video_fps), so a K scaled by the world size only made each + rank score the first 1/W of a W-times-denser subsample. EvenStrideDataset + keeps a whole episode when K exceeds its length, so a split that fits + entirely is not subsampled. ``auto`` computes that formula from the resolved split (``n_episodes``) and the head's loader ``batch_size``, for splits defined by live SQL filters @@ -385,7 +396,7 @@ def _metric_frames_per_episode( f"data.metric_frames_per_episode.{head} must be a positive int " f"or 'auto', got {k}" ) - return k * _trainer_world_size(cfg) + return k def _subsample_val_datasets( @@ -788,6 +799,7 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: os.makedirs(os.path.join(trainer.default_root_dir, "videos"), exist_ok=True) if mode == "train": + _apply_init_weights(cfg, model) if cfg.get("evaluator") is not None: eval_obj: Eval = hydra.utils.instantiate(cfg.evaluator) # data.valid_prefix (e.g. seen_op_valid in the opsplit configs) diff --git a/egomimic/utils/checkpoint_utils.py b/egomimic/utils/checkpoint_utils.py index 9d70bbe739..2142048dbc 100644 --- a/egomimic/utils/checkpoint_utils.py +++ b/egomimic/utils/checkpoint_utils.py @@ -40,3 +40,33 @@ def load_checkpoint_weights( "on random weights. Is this checkpoint from the same model config?" ) log.info("Loaded weights from %s", ckpt_path) + + +# A fine-tune init that matched less than this share of the model is not one. +MIN_INIT_FRACTION = 0.5 + + +def init_weights_from_checkpoint(model: torch.nn.Module, ckpt_path: str | os.PathLike) -> None: + """Fine-tune initialization: copy every tensor of a Lightning checkpoint's + ``state_dict`` whose name and shape match ``model``; leave the rest at their + fresh init. Unlike ``ckpt_path`` (a full resume: optimizer, scheduler, + epoch) nothing but weights is read, so a new embodiment's stems / head and + a different action width start fresh while the shared trunk transfers. + Raises when under MIN_INIT_FRACTION of the model's elements matched: an + unrelated checkpoint would otherwise "fine-tune" from random init.""" + ckpt_sd = torch.load(ckpt_path, map_location="cpu", weights_only=False, mmap=True)["state_dict"] + own_sd = model.state_dict() + loaded = {k: v for k, v in ckpt_sd.items() if k in own_sd and own_sd[k].shape == v.shape} + fresh = sorted(k for k in own_sd if k not in loaded) + model.load_state_dict(loaded, strict=False) + fraction = sum(own_sd[k].numel() for k in loaded) / max(sum(v.numel() for v in own_sd.values()), 1) + log.info( + "init_weights_ckpt %s: loaded %d tensors (%.1f%% of the model's elements); " + "%d checkpoint tensors unused; %d left at fresh init, e.g. %s", + ckpt_path, len(loaded), 100 * fraction, len(ckpt_sd) - len(loaded), len(fresh), fresh[:5], + ) + if fraction < MIN_INIT_FRACTION: + raise RuntimeError( + f"init_weights_ckpt {ckpt_path}: only {100 * fraction:.1f}% of the model's " + f"weights matched; is it from the same model family? e.g. fresh {fresh[:5]}" + ) diff --git a/tests/fixtures/synthetic_episodes.py b/tests/fixtures/synthetic_episodes.py index a63c8324ae..2183db4266 100644 --- a/tests/fixtures/synthetic_episodes.py +++ b/tests/fixtures/synthetic_episodes.py @@ -52,6 +52,11 @@ def _pose(T: int, rng: np.random.Generator) -> np.ndarray: return np.repeat(one, T, axis=0).astype(np.float64) +def _joints(T: int, rng: np.random.Generator) -> np.ndarray: + """One random 6-joint pose (radians) held constant for all T frames.""" + return np.repeat(rng.uniform(-1.5, 1.5, (1, 6)), T, axis=0).astype(np.float32) + + def _scalar(T: int, rng: np.random.Generator) -> np.ndarray: """One uniform(0,1) value held constant for all T frames, shape (T, 1).""" return np.full((T, 1), rng.uniform(0.0, 1.0)) @@ -93,6 +98,11 @@ def write_episode( "right.obs_gripper": _scalar(T, rng), "left.cmd_gripper": _scalar(T, rng), "right.cmd_gripper": _scalar(T, rng), + # ABC / YAM joint space (Eva mode "joints"): 6 joints per arm. + "left.obs_joints": _joints(T, rng), + "right.obs_joints": _joints(T, rng), + "left.cmd_joints": _joints(T, rng), + "right.cmd_joints": _joints(T, rng), } ) extrinsics = { diff --git a/tests/unit/test_abc_joint_configs.py b/tests/unit/test_abc_joint_configs.py new file mode 100644 index 0000000000..3ea2ea53e8 --- /dev/null +++ b/tests/unit/test_abc_joint_configs.py @@ -0,0 +1,120 @@ +"""The ABC sim fine-tune recipe end to end on CPU: two optimizer steps of a +tiny RDT on synthetic eva JOINT episodes through the real trainHydra.train() +path, then the saved checkpoint served by policy_server (load, rebuild the +training inputs from the checkpoint's own config, infer).""" + +from __future__ import annotations + +import math + +from fixtures.recipes import LOCAL_RESOLVER, common_overrides, cpu_trainer_overrides +from fixtures.train_harness import BatchKeySpy, compose_recipe, hermetic_env, write_fixtures +from fixtures.recipes import Recipe + +import numpy as np + +import egomimic.trainHydra as train_hydra +from egomimic.algo.hpt import HPT +from egomimic.rldb.embodiment.eva import JOINT_ACTION_KEY, JOINT_STATE_KEY, Eva +from egomimic.scripts.abc_sim.policy_server import JointPolicy + +EMB = "eva_bimanual" +RM = "model.robomimic_model" + + +def _tiny_rdt(emb: str) -> list[str]: + enc = f"{RM}.encoder_specs.front_img_1" + return [ + f"{RM}.width=32", + f"{RM}.trunk.depth=2", + f"{RM}.trunk.num_heads=4", + f"{RM}.head_specs.{emb}.num_inference_steps=2", + f"~{RM}.shared_stem_specs.annotation", + f"{RM}.shared_obs_keys=[front_img_1]", + f"{RM}.annotation_key=null", + f"{enc}.model_name=vit_small_patch16_dinov3", + f"{enc}.image_size=[32,48]", + f"+{enc}.pretrained=false", + f"+{enc}.tower_kwargs={{embed_dim: 32, depth: 1, num_heads: 2}}", + ] + + +def test_yam_robot_label_is_eva(): + from egomimic.rldb.embodiment.embodiment import get_embodiment_id + + assert get_embodiment_id("yam_bimanual") == get_embodiment_id("eva_bimanual") + + +def test_abc_real_recipe_serves_with_a_fixed_prompt_and_no_history(tmp_path, monkeypatch): + hermetic_env(monkeypatch) + data, out, hashes = write_fixtures(tmp_path, "eva") + recipe = Recipe("train_zarr_abc_real_rdt", "abc_real_joints", "rdt_abc_joints_dinov3_vitb", EMB) + tiny = [o for o in _tiny_rdt(EMB) if not o.startswith(f"{RM}.annotation_key")] + cfg = compose_recipe( + recipe, + common_overrides(EMB, data, out, batch_size=2, num_workers=0, episode_hashes=hashes) + + cpu_trainer_overrides(2) + + tiny + + [ + # abc_real_joints resolves through the SQL table; read the fixtures locally + f"data.train_datasets.{EMB}.resolver._target_={LOCAL_RESOLVER}", + f"data.train_datasets.{EMB}.resolver.folder_path={data}", + f"data.train_datasets.{EMB}.mode=total", + f"data.valid_datasets.{EMB}.filters=null", + f"data.valid_datasets.{EMB}.mode=total", + ], + out, + ) + assert cfg.model.robomimic_model.width == 32 and cfg.launch_params.gpus_per_node == 4 + spy = BatchKeySpy(monkeypatch, HPT) + metrics, objects = train_hydra.train(cfg) + assert math.isfinite(float(metrics["Train/action_loss"])) + seen = spy.keys[EMB] + assert {JOINT_ACTION_KEY, JOINT_STATE_KEY, Eva.VIZ_IMAGE_KEY, "fps"} <= seen + assert f"{Eva.VIZ_IMAGE_KEY}_hist" not in seen and "annotations" not in seen + + ckpt = out / "served.ckpt" + objects["trainer"].save_checkpoint(ckpt) + policy = JointPolicy(str(ckpt), device="cpu") + assert policy.cameras == ["top"] and policy.lag_frames == 0 and policy.annotation_key is None + assert policy.algo.default_prompt == "put the plastic bottles in the bin" + img = np.random.default_rng(0).random((3, 64, 64), dtype=np.float32) + actions = policy.infer(np.zeros(14, np.float32), {"top": img}, "any prompt; the model's own is used") + assert actions.shape == (45, 14) and np.isfinite(actions).all() + + +def test_abc_sim_finetune_train_step(tmp_path, monkeypatch): + hermetic_env(monkeypatch) + data, out, hashes = write_fixtures(tmp_path, "eva") + recipe = Recipe("train_zarr_abc_sim_rdt_ft", "abc_sim_joints", "rdt_1b_abc_joints_dinov3_vitb", EMB) + cfg = compose_recipe( + recipe, + common_overrides(EMB, data, out, batch_size=2, num_workers=0, episode_hashes=hashes) + + cpu_trainer_overrides(2) + + _tiny_rdt(EMB) + + [ + # the abc_sim_joints resolver is already local; point it at the fixtures + f"data.train_datasets.{EMB}.resolver.folder_path={data}", + f"data.valid_datasets.{EMB}.filters=null", + ], + out, + ) + assert cfg.data.train_datasets[EMB].resolver._target_ == LOCAL_RESOLVER + spy = BatchKeySpy(monkeypatch, HPT) + metrics, objects = train_hydra.train(cfg) + loss = float(metrics["Train/action_loss"]) + assert math.isfinite(loss) + assert objects["trainer"].global_step == 2 + seen = spy.keys[EMB] + expected = {JOINT_ACTION_KEY, JOINT_STATE_KEY, Eva.VIZ_IMAGE_KEY, f"{Eva.VIZ_IMAGE_KEY}_hist", "fps"} + assert expected <= seen, f"missing {expected - seen}; saw {sorted(seen)}" + assert "observations.images.left_wrist_img" not in seen # joints_front keymap + + # The checkpoint alone rebuilds the training pipeline and serves it. + ckpt = out / "served.ckpt" + objects["trainer"].save_checkpoint(ckpt) + policy = JointPolicy(str(ckpt), device="cpu") + assert policy.cameras == ["top"] and policy.lag_frames == 3 + img = np.random.default_rng(0).random((3, 64, 64), dtype=np.float32) + actions = policy.infer(np.zeros(14, np.float32), {"top": img, "top_hist": img}, "pick up the red cube") + assert actions.shape == (45, 14) and np.isfinite(actions).all() diff --git a/tests/unit/test_abc_sim_convert.py b/tests/unit/test_abc_sim_convert.py new file mode 100644 index 0000000000..e37a87c287 --- /dev/null +++ b/tests/unit/test_abc_sim_convert.py @@ -0,0 +1,115 @@ +"""ABC sim_224 episode -> EgoVerse zarr: states / actions land in the joint +keymap's keys in ABC order, the three stacked cameras are split onto +front/left/right, prompts follow ABC's sim rule, and the joint-mode dataset +reads the converted episode with ABC's own train/val split filterable.""" + +from __future__ import annotations + +import json + +import numpy as np +import pytest +import torch + +av = pytest.importorskip("av") + +from egomimic.rldb.embodiment.eva import JOINT_ACTION_KEY, JOINT_STATE_KEY, Eva +from egomimic.rldb.filters import DatasetFilter +from egomimic.rldb.zarr.zarr_dataset_multi import LocalEpisodeResolver, MultiDataset +from egomimic.scripts.abc_sim import convert_sim_to_zarr as conv + +T, H, W = 60, 16, 24 +TASK = "sim_put_the_plastic_bottles_in_the_bin" + + +def _write_abc_episode(root, split, name, meta_extra=None): + ep = root / f"{split}_sim" / name + ep.mkdir(parents=True) + t = np.arange(T, dtype=np.float64)[:, None] + states = np.tile(np.arange(14, dtype=np.float64), (T, 1)) + 0 * t + actions = states + 0.01 * t + np.concatenate([states, actions], axis=1).astype(np.float64).tofile( + ep / "states_actions.bin" + ) + # Three cameras stacked vertically, each a flat colour so the split is checkable. + with av.open(str(ep / "combined_camera-images-rgb.mp4"), "w") as c: + s = c.add_stream("libx264", rate=30) + s.width, s.height, s.pix_fmt = W, 3 * H, "yuv420p" + s.options = {"crf": "0"} + for _ in range(T): + img = np.zeros((3 * H, W, 3), np.uint8) + img[:H] = (200, 30, 30) + img[H : 2 * H] = (30, 200, 30) + img[2 * H :] = (30, 30, 200) + for pkt in s.encode(av.VideoFrame.from_ndarray(img, format="rgb24")): + c.mux(pkt) + for pkt in s.encode(): + c.mux(pkt) + meta = {"task_name": TASK, "num_steps": T, "cameras": ["top", "left", "right"]} + meta.update(meta_extra or {}) + (ep / "episode_metadata.json").write_text(json.dumps(meta)) + return ep + + +def test_prompt_rule_matches_abc(): + assert conv.prompt_spans({"task_name": TASK}, 10) == [ + ("sim put the plastic bottles in the bin", 0, 10) + ] + spans = conv.prompt_spans( + { + "task_name": "multi_drawer_search", + "prompt_timeline": [ + {"frame": 0, "prompt": "find the cup"}, + {"frame": 4, "prompt": "find the ball"}, + ], + }, + 10, + ) + assert spans == [("find the cup", 0, 4), ("find the ball", 4, 10)] + spans = conv.prompt_spans( + {"task_name": "put_relative", "instruction": "put the cube left of the mug"}, 5 + ) + assert spans == [("put the cube left of the mug", 0, 5)] + + +def test_intrinsics_are_the_sim_pinhole(): + K = conv.pinhole_intrinsics(168, 224) + assert K.shape == (3, 4) and K[0, 2] == 112 and K[1, 2] == 84 + assert K[1, 1] == pytest.approx(84 / np.tan(np.radians(29))) + + +def test_convert_and_read_back(tmp_path): + src, out = tmp_path / "cache", tmp_path / "zarr" + _write_abc_episode(src, "train", "episode_0001") + _write_abc_episode(src, "val", "episode_0002") + _write_abc_episode(src, "train", "episode_0003", {"task_name": "sim_pouring_beads"}) + _write_abc_episode(src, "train", "episode_0004", {"prompt_source": {"status": "excluded"}}) + assert conv.main(["--src", str(src), "--task", TASK, "--out", str(out), "--workers", "1"]) == 0 + report = json.loads((out / "conversion_report.json").read_text()) + assert report["written"] == 2 and report["skipped_unlabelled"] == 1 and not report["failed"] + # Idempotent: a rerun skips both. + conv.main(["--src", str(src), "--task", TASK, "--out", str(out), "--workers", "1"]) + assert json.loads((out / "conversion_report.json").read_text())["skipped"] == 2 + + resolver = LocalEpisodeResolver( + out, + key_map=Eva.get_keymap("joints", annotation_key="annotations"), + transform_list=Eva.get_transform_list("joints"), + ) + filters = DatasetFilter(filter_lambdas=["lambda row: row.get('split') == 'train'"]) + ds = MultiDataset._from_resolver(resolver, filters=filters, mode="total") + assert len(ds.datasets) == 1 + leaf = next(iter(ds.datasets.values())) + sample = leaf[5] + assert torch.allclose(sample[JOINT_STATE_KEY].float(), torch.arange(14.0)) + assert torch.allclose( + sample[JOINT_ACTION_KEY][0].float(), torch.arange(14.0) + 0.05, atol=1e-5 + ) + assert sample["annotations"] == ["sim put the plastic bottles in the bin"] + front = sample[Eva.VIZ_IMAGE_KEY] + assert front.shape == (3, H, W) + # top camera is red-dominant, left wrist green, right wrist blue + assert front[0].mean() > front[1].mean() + assert sample["observations.images.left_wrist_img"][1].mean() > 0.5 + assert sample["observations.images.right_wrist_img"][2].mean() > 0.5 + assert sample["fps"].item() == 30 diff --git a/tests/unit/test_episode_table_file.py b/tests/unit/test_episode_table_file.py new file mode 100644 index 0000000000..77d92d037d --- /dev/null +++ b/tests/unit/test_episode_table_file.py @@ -0,0 +1,16 @@ +"""EGOVERSE_EPISODE_TABLE: the S3 resolver reads the episode table from a +pickled snapshot instead of SQL, memoized per path.""" + +import pandas as pd + +from egomimic.rldb.zarr import zarr_dataset_multi as zdm + + +def test_episode_table_from_pickle(tmp_path, monkeypatch): + df = pd.DataFrame({"episode_hash": ["a"], "lab": ["abc_sim"]}) + path = tmp_path / "episodes.pkl" + df.to_pickle(path) + monkeypatch.setenv("EGOVERSE_EPISODE_TABLE", str(path)) + monkeypatch.setattr(zdm, "create_default_engine", lambda: (_ for _ in ()).throw(AssertionError("SQL touched"))) + assert list(zdm._episode_table().episode_hash) == ["a"] + assert list(zdm._episode_table().episode_hash) == ["a"] # second call: still no SQL diff --git a/tests/unit/test_eva_joint_mode.py b/tests/unit/test_eva_joint_mode.py new file mode 100644 index 0000000000..181ca07443 --- /dev/null +++ b/tests/unit/test_eva_joint_mode.py @@ -0,0 +1,97 @@ +"""Joint-space Eva mode (ABC YAM / abc_sim action space): the keymap reads +``{left,right}.{obs,cmd}_{joints,gripper}``, the transform emits a 14-D state +and a (100, 14) action chunk in ABC order, the robot keymap supports the same +``_hist`` camera as Human, and joint chunks get joint metrics, not the +cartesian ones their 14-D width would otherwise trigger.""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch + +from egomimic.eval.action_metrics import layout_metrics +from egomimic.rldb.embodiment.eva import ( + JOINT_ACTION_KEY, + JOINT_STATE_KEY, + Eva, +) +from egomimic.rldb.zarr.zarr_dataset_multi import LocalEpisodeResolver, MultiDataset +from egomimic.rldb.zarr.zarr_writer import ZarrWriter + +T = 120 +CAM = Eva.VIZ_IMAGE_KEY + + +def _write_episode(path, fps=30): + rng = np.random.default_rng(0) + t = np.arange(T, dtype=np.float32)[:, None] + numeric = {} + # Each channel carries a distinct offset so ordering is checkable. + for a, arm in enumerate(("left", "right")): + base = 10.0 * a + numeric[f"{arm}.obs_joints"] = base + np.arange(6, dtype=np.float32) + 0 * t + numeric[f"{arm}.cmd_joints"] = base + np.arange(6, dtype=np.float32) + 0.01 * t + numeric[f"{arm}.obs_gripper"] = np.full((T, 1), 0.25 + a * 0.5, np.float32) + numeric[f"{arm}.cmd_gripper"] = np.full((T, 1), 0.3 + a * 0.5, np.float32) + images = { + k: rng.integers(0, 255, (T, 24, 32, 3), np.uint8) + for k in ("images.front_1", "images.left_wrist", "images.right_wrist") + } + ZarrWriter.create_and_write( + path, + numeric_data=numeric, + image_data=images, + embodiment="eva_bimanual", + fps=fps, + task_name="synthetic", + intrinsics={"front_1": np.eye(3, 4)}, + ) + + +def _leaf(tmp_path, **keymap_kwargs): + _write_episode(tmp_path / "ep.zarr") + key_map = Eva.get_keymap("joints", **keymap_kwargs) + resolver = LocalEpisodeResolver( + tmp_path, key_map=key_map, transform_list=Eva.get_transform_list("joints") + ) + ds = MultiDataset._from_resolver(resolver, mode="total") + return next(iter(ds.datasets.values())) + + +def test_joint_sample_shapes_and_abc_order(tmp_path): + sample = _leaf(tmp_path)[10] + state = sample[JOINT_STATE_KEY].float() + actions = sample[JOINT_ACTION_KEY].float() + assert state.shape == (14,) and actions.shape == (100, 14) + expected = torch.tensor( + [0, 1, 2, 3, 4, 5, 0.25, 10, 11, 12, 13, 14, 15, 0.75], dtype=torch.float32 + ) + assert torch.allclose(state, expected) + # Chunk start = the commanded pose at idx 10 (+0.01 * 10 on the joints). + assert torch.allclose(actions[0, :6], expected[:6] + 0.1, atol=1e-5) + assert torch.allclose(actions[:, 6], torch.full((100,), 0.3)) + assert torch.allclose(actions[:, 13], torch.full((100,), 0.8)) + # 45 raw steps resampled onto 100: the last step is idx 10 + 44. + assert torch.allclose(actions[-1, 0], torch.tensor(0.54), atol=1e-5) + assert sample["fps"].item() == 30 + assert "left.cmd_joints" not in sample and "left.obs_joints" not in sample + + +def test_joint_keymap_supports_the_hist_camera(tmp_path): + leaf = _leaf(tmp_path, image_history_gap_s=0.1) + assert torch.equal(leaf[20][f"{CAM}_hist"], leaf[17][CAM]) + + +def test_joint_actions_get_joint_metrics_not_cartesian(): + gt = torch.zeros(2, 100, 14) + pred = gt.clone() + pred[..., 0] = np.deg2rad(2.0) # left j1 off by 2 degrees + pred[..., 6] = 0.5 # left gripper off by 0.5 + m = layout_metrics(pred, gt, "Valid/x", ac_key=JOINT_ACTION_KEY) + assert not any("xyz" in k for k in m) + assert m["Valid/x_joint_abs_err_deg_avg"].item() == pytest.approx(2.0 / 12, rel=1e-4) + assert m["Valid/x_joint_abs_err_deg_max_joint"].item() == pytest.approx(2.0, rel=1e-4) + assert m["Valid/x_grip_paired_mse_avg"].item() == pytest.approx(0.125, rel=1e-4) + # Without the key name the width alone reads as cartesian ypr. + assert any("xyz" in k for k in layout_metrics(pred, gt, "Valid/x")) diff --git a/tests/unit/test_eval_loaders.py b/tests/unit/test_eval_loaders.py index 1d8b92f22f..2760afa852 100644 --- a/tests/unit/test_eval_loaders.py +++ b/tests/unit/test_eval_loaders.py @@ -315,39 +315,15 @@ def test_every_head_can_be_subsampled_independently(): th._metric_frames_per_episode(_cfg(metric={"valid": 0}), "valid") -def test_k_scales_with_the_world_size(): - """K in the config is the single-rank value. DistributedSampler strides, - so W ranks between them read W times the frames in the same number of - batches each; without the scaling a 4-GPU run scores a quarter of what the - same config scores on one GPU.""" +def test_k_is_the_same_on_every_rank(): + """The val loaders are not sharded (no DistributedSampler inside the + CombinedLoaders; every rank runs every val batch), so K is per-rank and + must not grow with the world size: a W-scaled K had each rank score only + the first 1/W of the subsample.""" cfg = _cfg(metric={"valid": 7}) - for devices, nodes, expected in ( - (1, 1, 7), - (4, 1, 28), - (8, 2, 112), - ([0, 1], 1, 14), - ("auto", 1, 7), # unknown at config time: stay single-rank - (-1, 1, 7), - ): + for devices, nodes in ((1, 1), (4, 1), (8, 2), ([0, 1], 1)): cfg.trainer = {"devices": devices, "num_nodes": nodes} - assert th._metric_frames_per_episode(cfg, "valid") == expected, devices - # a split shorter than K is kept whole, so scaling cannot truncate - cfg.trainer = {"devices": 4, "num_nodes": 1} - split = {HUMAN: _split({"ep0": 20})} - assert len(th._subsample_val_datasets(cfg, "valid", split)[HUMAN]) == 20 - - -def test_eval_mode_is_single_rank(): - """Eval runs are forced onto one device after the loaders are built, so - their val split must be subsampled for one rank, not the training W.""" - cfg = _cfg(metric={"valid": 7}) - cfg.trainer = {"devices": 4, "num_nodes": 2} - assert th._metric_frames_per_episode(cfg, "valid") == 56 - cfg.mode = "eval" - assert th._metric_frames_per_episode(cfg, "valid") == 7 - del cfg["mode"] - cfg.eval = True - assert th._metric_frames_per_episode(cfg, "valid") == 7 + assert th._metric_frames_per_episode(cfg, "valid") == 7, devices @pytest.mark.parametrize("trainer_cfg", ["ddp", "ddp_pi"]) @@ -362,18 +338,17 @@ def test_ddp_devices_are_per_node(trainer_cfg, compose_resolve): ], ) assert cfg.trainer.devices == 8 and cfg.trainer.num_nodes == 2 - assert th._trainer_world_size(cfg) == 16 def test_auto_k_is_derived_from_the_resolved_split(): """`auto` = floor(limit_val_batches * batch_size / episodes), per dataset, - from the head's own loader batch size, then scaled by the world size.""" + from the head's own loader batch size, the same on every rank.""" cfg = _cfg(metric={"valid": "auto"}) cfg.trainer = {"limit_val_batches": 4, "devices": 2} split = {HUMAN: _split({"ep0": 100, "ep1": 100, "ep2": 100})} params = {HUMAN: {"batch_size": 5}} wrapped = th._subsample_val_datasets(cfg, "valid", split, params)[HUMAN] - assert wrapped.frames_per_episode == (4 * 5 // 3) * 2 + assert wrapped.frames_per_episode == 4 * 5 // 3 for limit in (1.0, None, True): cfg.trainer.limit_val_batches = limit with pytest.raises(ValueError, match="auto"): diff --git a/tests/unit/test_init_weights_ckpt.py b/tests/unit/test_init_weights_ckpt.py new file mode 100644 index 0000000000..a7b4e78188 --- /dev/null +++ b/tests/unit/test_init_weights_ckpt.py @@ -0,0 +1,60 @@ +"""Weights-only fine-tune init: matched tensors load, a new head stays fresh, +an unrelated checkpoint fails loudly, and a resume ignores it.""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn as nn +from omegaconf import OmegaConf + +import egomimic.trainHydra as train_hydra +from egomimic.utils.checkpoint_utils import init_weights_from_checkpoint + + +class _Net(nn.Module): + def __init__(self, head_out: int): + super().__init__() + self.trunk = nn.Linear(64, 64) + self.head = nn.Linear(64, head_out) + + +def _save(model: nn.Module, path) -> str: + torch.save({"state_dict": model.state_dict(), "optimizer_states": [1]}, path) + return str(path) + + +def test_matched_tensors_load_and_the_new_head_stays_fresh(tmp_path): + base = _Net(head_out=18) + ckpt = _save(base, tmp_path / "base.ckpt") + target = _Net(head_out=14) + fresh_head = target.head.weight.detach().clone() + init_weights_from_checkpoint(target, ckpt) + assert torch.equal(target.trunk.weight, base.trunk.weight) + assert torch.equal(target.head.weight, fresh_head) + + +def test_an_unrelated_checkpoint_raises(tmp_path): + other = nn.Sequential(nn.Linear(3, 3)) + ckpt = _save(other, tmp_path / "other.ckpt") + with pytest.raises(RuntimeError, match="matched"): + init_weights_from_checkpoint(_Net(14), ckpt) + + +def test_resume_ignores_init_weights(tmp_path): + ckpt = _save(_Net(18), tmp_path / "base.ckpt") + target = _Net(14) + before = target.trunk.weight.detach().clone() + cfg = OmegaConf.create( + {"init_weights_ckpt": ckpt, "ckpt_path": str(tmp_path / "last.ckpt")} + ) + assert train_hydra._apply_init_weights(cfg, target) is None + assert torch.equal(target.trunk.weight, before) + + +def test_missing_init_checkpoint_raises(tmp_path): + cfg = OmegaConf.create( + {"init_weights_ckpt": str(tmp_path / "nope.ckpt"), "ckpt_path": None} + ) + with pytest.raises(FileNotFoundError): + train_hydra._apply_init_weights(cfg, _Net(14)) diff --git a/tests/unit/test_pi05_norm_rot6d.py b/tests/unit/test_pi05_norm_rot6d.py index 0b75298f2b..d2969c48fc 100644 --- a/tests/unit/test_pi05_norm_rot6d.py +++ b/tests/unit/test_pi05_norm_rot6d.py @@ -454,8 +454,8 @@ def test_vendor_embodiment_names_collapse_to_human(): assert get_embodiment_id("human_bimanual") == EMBODIMENT.HUMAN_BIMANUAL.value assert get_embodiment_id("eva_bimanual") == EMBODIMENT.EVA_BIMANUAL.value assert not is_legacy_vendor_embodiment("human_bimanual") - with pytest.raises(KeyError): - get_embodiment_id("yam_bimanual") # robot names are never aliased + # the one robot is Eva: YAM_* (the ABC station after the 2026-09-01 relabel) is Eva + assert get_embodiment_id("yam_bimanual") == EMBODIMENT.EVA_BIMANUAL.value def test_base_converter_rejects_norm_6d_encoding(): diff --git a/tests/unit/test_policy_server.py b/tests/unit/test_policy_server.py new file mode 100644 index 0000000000..8999a71889 --- /dev/null +++ b/tests/unit/test_policy_server.py @@ -0,0 +1,125 @@ +"""policy_server builds the model input exactly as ZarrDataset builds a +training sample (keymap, transforms, the 0.1 s history frame, fps), normalizes +it with the checkpoint's norm stats, returns the chunk on the raw 30 Hz grid, +and speaks the sim client's HTTP protocol.""" + +from __future__ import annotations + +import base64 +import json +import threading +import urllib.error +import urllib.request +from http.server import HTTPServer + +import numpy as np +import simplejpeg +import torch +from fixtures.synthetic_episodes import write_episode +from omegaconf import OmegaConf + +from egomimic.rldb.embodiment.eva import JOINT_ACTION_KEY, JOINT_STATE_KEY, Eva +from egomimic.rldb.zarr.zarr_dataset_multi import LocalEpisodeResolver, MultiDataset, ZarrEpisode +from egomimic.scripts.abc_sim import policy_server as ps + +CAM = Eva.VIZ_IMAGE_KEY +CFG = OmegaConf.create({"data": {"train_datasets": {"eva_bimanual": {"resolver": { + "key_map": {"keymap_mode": "joints_front", "annotation_key": "annotations", "image_history_gap_s": 0.1}, + "transform_list": {"mode": "joints"}, +}}}}}) + + +class _FakeNormStats: + """Scales the state by 2 so the test can see normalization happened.""" + + def normalize(self, data, emb_id): + return {**data, JOINT_STATE_KEY: data[JOINT_STATE_KEY] * 2} + + +class _FakeAlgo: + """Records the batch; predicts a ramp in time so the resampling is visible.""" + + norm_stats = _FakeNormStats() + + def process_batch_for_training(self, batch): + self.seen = batch["eva_bimanual"] + return batch + + def forward_eval(self, batch): + ramp = torch.linspace(0, 99, 100)[None, :, None].repeat(1, 1, 14) + return {f"eva_bimanual_{JOINT_ACTION_KEY}": ramp} + + +def _policy(): + pol = ps.JointPolicy.__new__(ps.JointPolicy) + pol.ckpt_path, pol.algo = "fake.ckpt", _FakeAlgo() + pol.configure(CFG) + return pol + + +def _chw(hwc: np.ndarray) -> np.ndarray: + return np.moveaxis(hwc, -1, 0).astype(np.float32) / np.float32(255.0) + + +def test_server_sample_matches_the_dataset_sample(tmp_path): + write_episode(tmp_path, "eva", T=48) + ep = ZarrEpisode(next(tmp_path.glob("*.zarr"))) + keys = ["images.front_1", "left.obs_joints", "left.obs_gripper", "right.obs_joints", "right.obs_gripper"] + rows = {k: g[0][1] for k, g in ep.read_intervals({k: [(0, 48)] for k in keys}).items()} + frames = [simplejpeg.decode_jpeg(rows["images.front_1"][i], colorspace="RGB") for i in range(48)] + state = np.concatenate([np.asarray(rows[k], np.float32) for k in keys[1:]], axis=1) + key_map = Eva.get_keymap("joints_front", image_history_gap_s=0.1) + resolver = LocalEpisodeResolver(tmp_path, key_map=key_map, transform_list=Eva.get_transform_list("joints")) + leaf = next(iter(MultiDataset._from_resolver(resolver, mode="total").datasets.values())) + pol = _policy() + assert pol.cameras == ["top"] and pol.lag_frames == 3 + for t in (0, 1, 3, 10): + images = {"top": _chw(frames[t])} + if t > 0: # the frame lag_frames back, clamped to the episode start; at t=0 none is sent + images["top_hist"] = _chw(frames[max(0, t - pol.lag_frames)]) + sample, ref = pol.build_sample(state[t], images), leaf[t] + assert torch.equal(sample[CAM], ref[CAM]), t + assert torch.equal(sample[f"{CAM}_hist"], ref[f"{CAM}_hist"]), t + assert torch.allclose(sample[JOINT_STATE_KEY], ref[JOINT_STATE_KEY].float()), t + assert sample[JOINT_ACTION_KEY].shape == ref[JOINT_ACTION_KEY].shape == (100, 14) + assert sample["fps"].item() == ref["fps"].item() == 30 + + +def test_infer_normalizes_and_returns_the_raw_30hz_grid(): + pol = _policy() + out = pol.infer(np.arange(14, dtype=np.float32), {"top": np.zeros((3, 16, 24), np.float32)}, "sim x") + # 100 model steps span 45 raw frames: raw frame k sits at model step 99k/44 + assert out.shape == (45, 14) + assert np.allclose(out[:, 0], np.arange(45) * 99 / 44, atol=1e-4) + assert torch.allclose(pol.algo.seen[JOINT_STATE_KEY][0], torch.arange(14.0) * 2) + assert pol.algo.seen["annotations"] == [["sim x"]] + + +def test_http_protocol_round_trip(): + pol = _policy() + server = HTTPServer(("127.0.0.1", 0), ps.make_handler(pol)) + threading.Thread(target=server.serve_forever, daemon=True).start() + url = f"http://127.0.0.1:{server.server_address[1]}" + + def post(payload): + req = urllib.request.Request(f"{url}/infer", data=json.dumps(payload).encode(), + headers={"Content-Type": "application/json"}, method="POST") + with urllib.request.urlopen(req) as r: + return json.load(r) + + try: + with urllib.request.urlopen(f"{url}/health") as r: + health = json.load(r) + assert health["cameras"] == ["top"] and health["lag_frames"] == 3 + hwc = np.random.default_rng(0).integers(0, 255, (16, 24, 3), np.uint8) + out = post({"state": list(range(14)), "prompt": "sim x", + "images": {"top": {"shape": [16, 24, 3], "b64": base64.b64encode(hwc.tobytes()).decode()}}}) + assert np.asarray(out["actions"]).shape == (45, 14) + assert torch.allclose(pol.algo.seen[CAM][0], torch.from_numpy(_chw(hwc))) + try: + post({"state": [0], "images": {}}) + raise AssertionError("expected 500") + except urllib.error.HTTPError as e: + assert e.code == 500 and "error" in json.load(e) + finally: + server.shutdown() diff --git a/tests/unit/test_val_vs_sim.py b/tests/unit/test_val_vs_sim.py new file mode 100644 index 0000000000..7a25745d92 --- /dev/null +++ b/tests/unit/test_val_vs_sim.py @@ -0,0 +1,21 @@ +"""val_vs_sim: correlations and the log-linear fit from a manifest with +pre-filled val metrics (no W&B).""" + +import json + +from egomimic.scripts.abc_sim import val_vs_sim as vs + + +def test_report_from_manifest(tmp_path): + rows = [] + for i, (h, ep, val, succ) in enumerate([(0.237, 1000, 0.9, 0.1), (0.559, 1000, 0.6, 0.4), (1.096, 1000, 0.3, 0.7), (1.096, 500, 0.5, 0.5)]): + s = tmp_path / f"s{i}.json" + s.write_text(json.dumps({"success_rate": succ, "mean_max_progress": succ + 0.1, "num_worlds": 50})) + rows.append({"name": f"r{i}", "epoch": ep, "hours": h, "val": val, "summary": str(s)}) + m = tmp_path / "m.json" + m.write_text(json.dumps(rows)) + assert vs.main([str(m), "--metric", "x", "--out", str(tmp_path / "out")]) == 0 + rep = json.loads((tmp_path / "out" / "report.json").read_text()) + assert rep["correlations"]["spearman_success_rate"]["r"] == -1.0 # lower val mse, higher success + assert rep["loglinear_val_vs_hours"]["r2"] > 0.95 and rep["loglinear_val_vs_hours"]["slope"] < 0 + assert (tmp_path / "out" / "report.md").read_text().count("| r") >= 4