From f57301f88d327d8d637d75db7c72f19d1b5b6a37 Mon Sep 17 00:00:00 2001
From: ElmoPA <elmoworld2005@gmail.com>
Date: Thu, 12 Feb 2026 00:21:40 -0500
Subject: [PATCH] "Debug training runs"

---
 .../hydra_configs/data/test_multi_zarr.yaml   | 60 ++++++-------------
 egomimic/rldb/zarr/zarr_dataset_multi.py      | 11 +++-
 2 files changed, 27 insertions(+), 44 deletions(-)

diff --git a/egomimic/hydra_configs/data/test_multi_zarr.yaml b/egomimic/hydra_configs/data/test_multi_zarr.yaml
index 8b1953bf..b98bc752 100644
--- a/egomimic/hydra_configs/data/test_multi_zarr.yaml
+++ b/egomimic/hydra_configs/data/test_multi_zarr.yaml
@@ -1,67 +1,45 @@
 _target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper
 train_datasets:
-  dataset1:
+  aria_bimanual:
     _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver
     resolver:
       _target_: egomimic.rldb.zarr.zarr_dataset_multi.LocalEpisodeResolver
-      folder_path: /nethome/paphiwetsa3/flash/datasets/test_zarr/
+      folder_path: /nethome/paphiwetsa3/flash/datasets/proc_zarr
       key_map:
         front_img_1: #batch key
           key_type: camera_keys # key type
-          zarr_key: observations.images.front_img_1 # dataset key
-        right_wrist_img:
-          key_type: camera_keys
-          zarr_key: observations.images.right_wrist_img
-        left_wrist_img:
-          key_type: camera_keys
-          zarr_key: observations.images.left_wrist_img
-        ee_pose:
-          key_type: proprio_keys
-          zarr_key: observations.state.ee_pose
-          horizon: 100
-        joint_positions:
-          key_type: proprio_keys
-          zarr_key: observations.state.joint_positions
-          horizon: 100
+          zarr_key: front_img_1
         actions_cartesian:
           key_type: action_keys
-          zarr_key: actions_cartesian
+          zarr_key: obs_ee_pose_6d
           horizon: 100
-    mode: total 
+        ee_pose:
+          key_type: proprio_keys
+          zarr_key: obs_ee_pose_6d
+    mode: total
 valid_datasets:
-  dataset1:
+  aria_bimanual:
     _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver
     resolver:
       _target_: egomimic.rldb.zarr.zarr_dataset_multi.LocalEpisodeResolver
-      folder_path: /nethome/paphiwetsa3/flash/datasets/test_zarr/
+      folder_path: /nethome/paphiwetsa3/flash/datasets/proc_zarr
       key_map:
         front_img_1: #batch key
           key_type: camera_keys # key type
-          zarr_key: observations.images.front_img_1 # dataset key
-        right_wrist_img:
-          key_type: camera_keys
-          zarr_key: observations.images.right_wrist_img
-        left_wrist_img:
-          key_type: camera_keys
-          zarr_key: observations.images.left_wrist_img
-        ee_pose:
-          key_type: proprio_keys
-          zarr_key: observations.state.ee_pose
-          horizon: 100
-        joint_positions:
-          key_type: proprio_keys
-          zarr_key: observations.state.joint_positions
-          horizon: 100
+          zarr_key: front_img_1
         actions_cartesian:
           key_type: action_keys
-          zarr_key: actions_cartesian
+          zarr_key: obs_ee_pose_6d
           horizon: 100
+        ee_pose:
+          key_type: proprio_keys
+          zarr_key: obs_ee_pose_6d
     mode: total
 train_dataloader_params:
-  dataset1:
-    batch_size: 2
+  aria_bimanual:
+    batch_size: 32
     num_workers: 10
 valid_dataloader_params:
-  dataset1:
-    batch_size: 2
+  aria_bimanual:
+    batch_size: 32
     num_workers: 10
diff --git a/egomimic/rldb/zarr/zarr_dataset_multi.py b/egomimic/rldb/zarr/zarr_dataset_multi.py
index 37aa09b6..fab19ef5 100644
--- a/egomimic/rldb/zarr/zarr_dataset_multi.py
+++ b/egomimic/rldb/zarr/zarr_dataset_multi.py
@@ -399,8 +399,8 @@ def resolve(
 
         filtered_paths = self._get_local_filtered_paths(self.folder_path, filters)
 
-        valid_hashes = {hashes for _, hashes in filtered_paths}
-        if not valid_hashes:
+        valid_folder_names = {folder_name for _, folder_name in filtered_paths}
+        if not valid_folder_names:
             raise ValueError(
                 "No valid collection names from local filtering: "
                 "filters matched no episodes in the local directory."
@@ -408,7 +408,7 @@ def resolve(
 
         datasets = self._load_zarr_datasets(
             search_path=self.folder_path,
-            valid_folder_names=valid_hashes
+            valid_folder_names=valid_folder_names
         )
 
         return datasets
@@ -473,6 +473,10 @@ def __len__(self) -> int:
     def __getitem__(self, idx):
         dataset_name, local_idx = self.index_map[idx]
         data = self.datasets[dataset_name][local_idx]
+
+        robot_name = self.datasets[dataset_name].embodiment
+        data["metadata.robot_name"] = robot_name
+        data["embodiment"] = robot_name
             
         return data
     
@@ -542,6 +546,7 @@ def init_episode(self):
         self.episode_reader = ZarrEpisode(self.episode_path)
         self.metadata = self.episode_reader.metadata
         self.total_frames = self.metadata["total_frames"]
+        self.embodiment = self.metadata["embodiment"]
         self.keys_dict = {k: (0, None) for k in self.episode_reader._collect_keys()}
 
         # Detect JPEG-encoded image keys from metadata