eva: joint-space mode for the ABC YAM action space; robot image history

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

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}"

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:

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}
+

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

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

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

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

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

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

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=<rdt_mem_all .../epoch_epoch=899.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

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=<base.ckpt>
+defaults:
+  - train_zarr_abc_sim_rdt_ft
+  - override data: abc_real_joints
+  - _self_

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

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

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"]