scale configs for working training

egomimic/hydra_configs/data/scale_zarr_test.yaml

@@ -0,0 +1,71 @@
+_target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper
+train_datasets:
+  scale_bimanual:
+    _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver
+    resolver:
+      _target_: egomimic.rldb.zarr.zarr_dataset_multi.LocalEpisodeResolver
+      folder_path: /coc/flash7/scratch/egoverseDebugDatasets/scale/2026-02-24-01-49-24-166324
+      key_map:
+        observations.images.front_img_1: #batch key
+          key_type: camera_keys # key type
+          zarr_key: images.front_1
+        right.action_ee_pose:
+          key_type: action_keys
+          zarr_key: right.obs_ee_pose
+          horizon: 30
+        left.action_ee_pose:
+          key_type: action_keys
+          zarr_key: left.obs_ee_pose
+          horizon: 30
+        right.obs_ee_pose:
+          key_type: proprio_keys
+          zarr_key: right.obs_ee_pose
+        left.obs_ee_pose:
+          key_type: proprio_keys
+          zarr_key: left.obs_ee_pose
+        obs_head_pose:
+          zarr_key: obs_head_pose
+          key_type: proprio_keys
+      transform_list:
+        _target_: egomimic.rldb.zarr.action_chunk_transforms.build_aria_bimanual_transform_list
+        stride: 1
+    mode: train
+valid_datasets:
+  scale_bimanual:
+    _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver
+    resolver:
+      _target_: egomimic.rldb.zarr.zarr_dataset_multi.LocalEpisodeResolver
+      folder_path: /coc/flash7/scratch/egoverseDebugDatasets/scale/2026-02-24-01-49-24-166324
+      key_map:
+        observations.images.front_img_1: #batch key
+          key_type: camera_keys # key type
+          zarr_key: images.front_1
+        right.action_ee_pose:
+          key_type: action_keys
+          zarr_key: right.obs_ee_pose
+          horizon: 30
+        left.action_ee_pose:
+          key_type: action_keys
+          zarr_key: left.obs_ee_pose
+          horizon: 30
+        right.obs_ee_pose:
+          key_type: proprio_keys
+          zarr_key: right.obs_ee_pose
+        left.obs_ee_pose:
+          key_type: proprio_keys
+          zarr_key: left.obs_ee_pose
+        obs_head_pose:
+          zarr_key: obs_head_pose
+          key_type: proprio_keys
+      transform_list:
+        _target_: egomimic.rldb.zarr.action_chunk_transforms.build_aria_bimanual_transform_list
+        stride: 1
+    mode: valid
+train_dataloader_params:
+  scale_bimanual:
+    batch_size: 32
+    num_workers: 10
+valid_dataloader_params:
+  scale_bimanual:
+    batch_size: 32
+    num_workers: 10

egomimic/hydra_configs/model/hpt_bc_flow_scale.yaml

@@ -53,7 +53,7 @@ robomimic_model:
     scale_bimanual:
       state_ee_pose:
         _target_: egomimic.models.hpt_nets.MLPPolicyStem
-        input_dim: 14
+        input_dim: 12
         output_dim: 256
         widths: [256]
         specs:
@@ -73,13 +73,13 @@ robomimic_model:
       pooling: null
       time_dist: "beta"
       infer_ac_dims:
-        scale_bimanual: 14
+        scale_bimanual: 12
       model:
         _target_: egomimic.models.denoising_nets.CrossTransformer
         nblocks: 6
         cond_dim: 256
         hidden_dim: 128
-        act_dim: 14
+        act_dim: 12
         act_seq: 100
         n_heads: 4
         dropout: 0.1

egomimic/hydra_configs/train_zarr.yaml

@@ -4,7 +4,7 @@ defaults:
   - trainer: debug
   - debug: null
   - logger: debug
-  - data: eva_bc_zarr.yaml
+  - data: eva_bc_zarr
   - callbacks: checkpoints
   - override hydra/launcher: submitit
   - _self_
@@ -115,4 +115,4 @@ data_schematic: # Dynamically fill in these shapes from the dataset
     mecka_bimanual:
       front_img_1
     scale_bimanual:
-      front_img_1
+      front_img_1

egomimic/scripts/zarr_data_viz.ipynb

@@ -328,8 +328,8 @@
     "# EPISODE_PATH = Path(\"/coc/flash7/scratch/egoverseDebugDatasets/scale/697a9070da7b91acaf3f2d88_episode_000000.zarr\") # Scale\n",
     "# intrinsics_key = \"scale\"\n",
     "\n",
-    "EPISODE_PATH = Path(\"/coc/flash7/scratch/egoverseDebugDatasets/1767671007927.zarr/\") # Aria\n",
-    "intrinsics_key = \"base\"\n",
+    "EPISODE_PATH = Path(\"/coc/flash7/scratch/egoverseDebugDatasets/scale/2026-02-24-01-49-24-166324/697a9070da7b91acaf3f2d88_episode_000002.zarr\") # Aria\n",
+    "intrinsics_key = \"scale\"\n",
     "\n",
     "\n",
     "key_map = {\n",
@@ -350,6 +350,47 @@
     ")\n"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1b4c03ca",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Full MultiDataset via LocalEpisodeResolver (mirrors test_multi_zarr.yaml)\n",
+    "from egomimic.rldb.zarr.zarr_dataset_multi import LocalEpisodeResolver, MultiDataset\n",
+    "from egomimic.rldb.zarr.action_chunk_transforms import build_aria_bimanual_transform_list\n",
+    "\n",
+    "SCALE_FOLDER = Path(\"/coc/flash7/scratch/egoverseDebugDatasets/scale/2026-02-24-01-49-24-166324\")\n",
+    "\n",
+    "key_map = {\n",
+    "    \"images.front_1\": {\"zarr_key\": \"images.front_1\"},\n",
+    "    \"right.obs_ee_pose\": {\"zarr_key\": \"right.obs_ee_pose\"},\n",
+    "    \"left.obs_ee_pose\": {\"zarr_key\": \"left.obs_ee_pose\"},\n",
+    "    \"right.action_ee_pose\": {\"zarr_key\": \"right.obs_ee_pose\", \"horizon\": 30},\n",
+    "    \"left.action_ee_pose\": {\"zarr_key\": \"left.obs_ee_pose\", \"horizon\": 30},\n",
+    "    \"obs_head_pose\": {\"zarr_key\": \"obs_head_pose\"},\n",
+    "}\n",
+    "\n",
+    "transform_list = build_aria_bimanual_transform_list(\n",
+    "    stride=1,\n",
+    ")\n",
+    "\n",
+    "resolver = LocalEpisodeResolver(\n",
+    "    folder_path=SCALE_FOLDER,\n",
+    "    key_map=key_map,\n",
+    "    transform_list=transform_list,\n",
+    ")\n",
+    "\n",
+    "multi_ds = MultiDataset._from_resolver(resolver, mode=\"total\")\n",
+    "print(f\"MultiDataset total frames: {len(multi_ds)}\")\n",
+    "print(f\"Underlying episodes: {list(multi_ds.datasets.keys())}\")\n",
+    "\n",
+    "loader = torch.utils.data.DataLoader(multi_ds, batch_size=1, shuffle=False)\n",
+    "batch = next(iter(loader))\n",
+    "print(\"Batch keys:\", list(batch.keys()))"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -379,7 +420,61 @@
    "outputs": [],
    "source": [
     "batch = next(iter(loader))\n",
-    "nds(batch)"
+    "nds(batch)\n",
+    "print(\"Batch keys:\", list(batch.keys()))\n",
+    "print(batch[\"right.action_ee_pose\"][0, 0])\n",
+    "print(batch[\"left.action_ee_pose\"][0, 0])\n",
+    "print(batch[\"obs_head_pose\"][0])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0e94799c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "left_hand_pose = []\n",
+    "right_hand_pose = []\n",
+    "head_pose = []\n",
+    "for i, batch in enumerate(loader):\n",
+    "    left_hand_pose.append(batch[\"left.action_ee_pose\"][0, 0])\n",
+    "    right_hand_pose.append(batch[\"right.action_ee_pose\"][0, 0])\n",
+    "    head_pose.append(batch[\"obs_head_pose\"][0])\n",
+    "    \n",
+    "    if i > 400:\n",
+    "        break\n",
+    "left_hand_pose = np.array(left_hand_pose)\n",
+    "right_hand_pose = np.array(right_hand_pose)\n",
+    "head_pose = np.array(head_pose)\n",
+    "\n",
+    "# chunk the pose to actions(N, 100, 3)\n",
+    "left_hand_pose_actions = []\n",
+    "right_hand_pose_actions = []\n",
+    "head_pose_actions = []\n",
+    "for i in range(left_hand_pose.shape[0] - 100):\n",
+    "    action_left_hand = left_hand_pose[i:i+100, :]\n",
+    "    action_right_hand = right_hand_pose[i:i+100, :]\n",
+    "    action_head = head_pose[i:i+100, :]\n",
+    "    left_hand_pose_actions.append(action_left_hand)\n",
+    "    right_hand_pose_actions.append(action_right_hand)\n",
+    "    head_pose_actions.append(action_head)\n",
+    "left_hand_pose_actions = np.array(left_hand_pose_actions)\n",
+    "right_hand_pose_actions = np.array(right_hand_pose_actions)\n",
+    "head_pose_actions = np.array(head_pose_actions)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6517f061",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from egomimic.utils.egomimicUtils import render_3d_traj_frames\n",
+    "\n",
+    "frames = render_3d_traj_frames([left_hand_pose_actions, right_hand_pose_actions, head_pose_actions], labels=[\"left hand\", \"right hand\", \"head\"], stride=10)\n",
+    "mpy.show_video(frames, fps=30)\n"
    ]
   },
   {