Latent Flow matching that works regardless of dimension dim

egomimic/algo/hpt.py

@@ -359,7 +359,7 @@ def stem_process(self, domain, data):
         feat_dict = {}
         for modality in self.modalities.get(domain, []) + self.shared_keys:
             if modality not in data:
-                continue
+                raise ValueError(f"Modality {modality} not found in data")
             if modality in self.shared_keys:
                 domain = "shared"
 
@@ -829,7 +829,6 @@ def __init__(
         self.domains = domains.copy()
         self.auxiliary_ac_keys = auxiliary_ac_keys.copy()
         self.shared_ac_key = kwargs.get("shared_ac_key", None)
-        self.is_6dof = kwargs.get("6dof", False)
         self.kinematics_solver = kwargs.get("kinematics_solver", None)
 
         model = HPTModel(**trunk)
@@ -1282,13 +1281,16 @@ def compute_losses(self, predictions, batch):
             embodiment_name = get_embodiment(embodiment_id).lower()
             bc_loss = predictions[f"{embodiment_name}_loss"]
             scaled_bc_loss = bc_weight * bc_loss
-            total_action_loss += scaled_bc_loss
+            total_action_loss = total_action_loss + scaled_bc_loss
             loss_dict[f"{embodiment_name}_loss"] = bc_loss  # for logging
 
         if self.ot:
             loss_dict["ot_loss"] = predictions["ot_loss"]
             loss_dict["avg_feature_distance"] = predictions["avg_feature_distance"]
-            total_action_loss += ot_weight * self.temperature * predictions["ot_loss"]
+            total_action_loss = (
+                total_action_loss
+                + ot_weight * self.temperature * predictions["ot_loss"]
+            )
 
         loss_dict["action_loss"] = total_action_loss / len(self.domains)
         return loss_dict
@@ -1372,7 +1374,6 @@ def _robomimic_to_hpt_data(
             if key in batch:
                 data[key] = batch[key]
 
-        data["is_6dof"] = self.is_6dof
         data["pad_mask"] = batch["pad_mask"]
         data["embodiment"] = batch["embodiment"]
 

egomimic/hydra_configs/data/aria_debug.yaml

@@ -0,0 +1,39 @@
+_target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper
+train_datasets:
+  aria_bimanual:
+    _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver
+    resolver:
+      _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver
+      folder_path: /coc/flash7/scratch/egoverseDebugDatasets/aria
+      key_map:
+        _target_: egomimic.rldb.embodiment.human.Aria.get_keymap
+        mode: keypoints
+      transform_list:
+        _target_: egomimic.rldb.embodiment.human.Aria.get_transform_list
+        mode: keypoints
+    filters:
+      episode_hash: "2025-11-27-23-44-43-234000"
+    mode: total
+valid_datasets:
+  aria_bimanual:
+    _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver
+    resolver:
+      _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver
+      folder_path: /coc/flash7/scratch/egoverseDebugDatasets/aria
+      key_map:
+        _target_: egomimic.rldb.embodiment.human.Aria.get_keymap
+        mode: keypoints
+      transform_list:
+        _target_: egomimic.rldb.embodiment.human.Aria.get_transform_list
+        mode: keypoints
+    filters:
+      episode_hash: "2025-11-27-23-44-43-234000"
+    mode: total
+train_dataloader_params:
+  aria_bimanual:
+    batch_size: 64
+    num_workers: 10
+valid_dataloader_params:
+  aria_bimanual:
+    batch_size: 64
+    num_workers: 10

egomimic/hydra_configs/data/eva_human_cotrain.yaml

@@ -4,7 +4,7 @@ train_datasets:
     _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver
     resolver:
       _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver
-      folder_path: /coc/flash7/scratch/egoverseDebugDatasets/egoverseS3DatasetTest/
+      folder_path: /storage/project/r-dxu345-0/paphiwetsa3/datasets/temp_train
       key_map:
         _target_: egomimic.rldb.embodiment.eva.Eva.get_keymap
       transform_list:
@@ -16,7 +16,7 @@ train_datasets:
     _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver
     resolver:
       _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver
-      folder_path: /coc/flash7/scratch/egoverseDebugDatasets/aria
+      folder_path: /storage/project/r-dxu345-0/paphiwetsa3/datasets/temp_train
       key_map:
         _target_: egomimic.rldb.embodiment.human.Aria.get_keymap
         mode: cartesian
@@ -31,7 +31,7 @@ valid_datasets:
     _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver
     resolver:
       _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver
-      folder_path: /coc/flash7/scratch/egoverseDebugDatasets/egoverseS3DatasetTest/
+      folder_path: /storage/project/r-dxu345-0/paphiwetsa3/datasets/temp_train
       key_map:
         _target_: egomimic.rldb.embodiment.eva.Eva.get_keymap
       transform_list:
@@ -43,7 +43,7 @@ valid_datasets:
     _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver
     resolver:
       _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver
-      folder_path: /coc/flash7/scratch/egoverseDebugDatasets/aria
+      folder_path: /storage/project/r-dxu345-0/paphiwetsa3/datasets/temp_train
       key_map:
         _target_: egomimic.rldb.embodiment.human.Aria.get_keymap
         mode: cartesian

egomimic/hydra_configs/data/eva_human_keypoints_cotrain.yaml

@@ -0,0 +1,73 @@
+_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.S3EpisodeResolver
+      folder_path: /storage/project/r-dxu345-0/paphiwetsa3/datasets/temp_train
+      key_map:
+        _target_: egomimic.rldb.embodiment.eva.Eva.get_keymap
+      transform_list:
+        _target_: egomimic.rldb.embodiment.eva.Eva.get_transform_list
+    filters:
+      robot_name: "eva_bimanual"
+      task: "fold_clothes"
+    mode: total
+  aria_bimanual:
+    _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver
+    resolver:
+      _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver
+      folder_path: /storage/project/r-dxu345-0/paphiwetsa3/datasets/temp_train
+      key_map:
+        _target_: egomimic.rldb.embodiment.human.Aria.get_keymap
+        mode: keypoints
+      transform_list:
+        _target_: egomimic.rldb.embodiment.human.Aria.get_transform_list
+        mode: keypoints
+    filters:
+      robot_name: "aria_bimanual"
+      task: "fold_clothes_indomain"
+    mode: total
+valid_datasets:
+  eva_bimanual:
+    _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver
+    resolver:
+      _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver
+      folder_path: /storage/project/r-dxu345-0/paphiwetsa3/datasets/temp_train
+      key_map:
+        _target_: egomimic.rldb.embodiment.eva.Eva.get_keymap
+      transform_list:
+        _target_: egomimic.rldb.embodiment.eva.Eva.get_transform_list
+    filters:
+      robot_name: "eva_bimanual"
+      task: "fold_clothes"
+    mode: total
+  aria_bimanual:
+    _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver
+    resolver:
+      _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver
+      folder_path: /storage/project/r-dxu345-0/paphiwetsa3/datasets/temp_train
+      key_map:
+        _target_: egomimic.rldb.embodiment.human.Aria.get_keymap
+        mode: keypoints
+      transform_list:
+        _target_: egomimic.rldb.embodiment.human.Aria.get_transform_list
+        mode: keypoints
+    filters:
+      robot_name: "aria_bimanual"
+      task: "fold_clothes_indomain"
+    mode: total
+train_dataloader_params:
+  eva_bimanual:
+    batch_size: 256
+    num_workers: 6
+  aria_bimanual:
+    batch_size: 256
+    num_workers: 6
+valid_dataloader_params:
+  eva_bimanual:
+    batch_size: 256
+    num_workers: 6
+  aria_bimanual:
+    batch_size: 256
+    num_workers: 6

egomimic/hydra_configs/data/eva_human_keypoints_cotrain_wrist.yaml

@@ -0,0 +1,77 @@
+_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.S3EpisodeResolver
+      folder_path: /storage/project/r-dxu345-0/paphiwetsa3/datasets/temp_train
+      key_map:
+        _target_: egomimic.rldb.embodiment.eva.Eva.get_keymap
+      transform_list:
+        _target_: egomimic.rldb.embodiment.eva.Eva.get_transform_list
+        mode: cartesian_wristframe_ypr
+    filters:
+      robot_name: "eva_bimanual"
+      task: "fold_clothes"
+    mode: total
+  aria_bimanual:
+    _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver
+    resolver:
+      _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver
+      folder_path: /storage/project/r-dxu345-0/paphiwetsa3/datasets/temp_train
+      key_map:
+        _target_: egomimic.rldb.embodiment.human.Aria.get_keymap
+        mode: keypoints
+      transform_list:
+        _target_: egomimic.rldb.embodiment.human.Aria.get_transform_list
+        mode: keypoints_wristframe_ypr
+    filters:
+      robot_name: "aria_bimanual"
+      task: "fold_clothes_indomain"
+      lab: "rl2"
+    mode: total
+valid_datasets:
+  eva_bimanual:
+    _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver
+    resolver:
+      _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver
+      folder_path: /storage/project/r-dxu345-0/paphiwetsa3/datasets/temp_train
+      key_map:
+        _target_: egomimic.rldb.embodiment.eva.Eva.get_keymap
+      transform_list:
+        _target_: egomimic.rldb.embodiment.eva.Eva.get_transform_list
+        mode: cartesian_wristframe_ypr
+    filters:
+      robot_name: "eva_bimanual"
+      task: "fold_clothes"
+    mode: total
+  aria_bimanual:
+    _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver
+    resolver:
+      _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver
+      folder_path: /storage/project/r-dxu345-0/paphiwetsa3/datasets/temp_train
+      key_map:
+        _target_: egomimic.rldb.embodiment.human.Aria.get_keymap
+        mode: keypoints
+      transform_list:
+        _target_: egomimic.rldb.embodiment.human.Aria.get_transform_list
+        mode: keypoints_wristframe_ypr
+    filters:
+      robot_name: "aria_bimanual"
+      task: "fold_clothes"
+      lab: "rl2"
+    mode: total
+train_dataloader_params:
+  eva_bimanual:
+    batch_size: 32
+    num_workers: 6
+  aria_bimanual:
+    batch_size: 64
+    num_workers: 6
+valid_dataloader_params:
+  eva_bimanual:
+    batch_size: 16
+    num_workers: 6
+  aria_bimanual:
+    batch_size: 16
+    num_workers: 6

egomimic/hydra_configs/hydra/launcher/submitit.yaml

@@ -4,15 +4,15 @@ defaults:
 _target_: hydra_plugins.hydra_submitit_launcher.submitit_launcher.SlurmLauncher
 
 # Slurm configuration
-name: ${hydra.job.name}                             # Default job name
-partition: "rl2-lab"                            # Slurm partition (e.g., 'gpu' or 'compute')
-account: "rl2-lab"                              # Slurm account (e.g., 'my_account')
-cpus_per_task: 12                                   # Number of CPUs per task
-nodes: ${launch_params.nodes}                       # Number of nodes
-tasks_per_node: ${launch_params.gpus_per_node}      # Use variable for tasks per node
+name: ${hydra.job.name} # Default job name
+partition: "hoffman-lab" # Slurm partition (e.g., 'gpu' or 'compute')
+account: "hoffman-lab" # Slurm account (e.g., 'my_account')
+cpus_per_task: 12 # Number of CPUs per task
+nodes: ${launch_params.nodes} # Number of nodes
+tasks_per_node: ${launch_params.gpus_per_node} # Use variable for tasks per node
 gres: "gpu:a40:${eval:'${launch_params.gpus_per_node} * ${launch_params.nodes}'}" # GPU type and count
-qos: "short"                                        # Slurm QoS
-timeout_min: 2880                                   # Timeout in minutes (48 hours)
-exclude: "protocol, puma"     # Nodes to exclude
+qos: "short" # Slurm QoS
+timeout_min: 2880 # Timeout in minutes (48 hours)
+exclude: "protocol, puma" # Nodes to exclude
 additional_parameters:
-  requeue: true
+  requeue: true

egomimic/hydra_configs/hydra/launcher/submitit_pace.yaml

@@ -4,14 +4,15 @@ defaults:
 _target_: hydra_plugins.hydra_submitit_launcher.submitit_launcher.SlurmLauncher
 
 # Slurm configuration
+
 name: ${hydra.job.name}                             # Default job name
 partition: "gpu-h200"                             # Slurm partition
 account: "gts-dxu345-rl2"                              # Slurm account
 cpus_per_task: 8                                   # Number of CPUs per task (max 4:1 CPU:GPU ratio)
 nodes: ${launch_params.nodes}                       # Number of nodes
 tasks_per_node: ${launch_params.gpus_per_node}      # Use variable for tasks per node
 gres: "gpu:h200:${eval:'${launch_params.gpus_per_node} * ${launch_params.nodes}'}" # GPU type and count (h100 for H100 GPUs)
-qos: "short"                                        # Slurm QoS
+qos: "inferno"                                        # Slurm QoS
 mem_per_gpu: 250G
 timeout_min: 2880                                   # Timeout in minutes (48 hours)
 # exclude: "protocol, puma"     # Nodes to exclude

egomimic/hydra_configs/logger/wandb.yaml

@@ -7,7 +7,7 @@ wandb:
   offline: False
   id: "${name}_${description}_${now:%Y-%m-%d_%H-%M-%S}" # pass correct id to resume experiment!
   anonymous: null # enable anonymous logging
-  project: "zarr_test"
+  project: "keypoints_cotrain"
   log_model: False # upload lightning ckpts
   prefix: "" # a string to put at the beginning of metric keys
   entity: "rl2-group" # set to name of your wandb team

egomimic/hydra_configs/model/hpt_cotrain_flow_aria_keypoints.yaml

@@ -0,0 +1,54 @@
+defaults:
+  - hpt_cotrain_keypoints_base
+
+robomimic_model:
+  ac_keys:
+    eva_bimanual: "actions_eva_cart_aria_keypoints"
+    aria_bimanual: "actions_eva_cart_aria_keypoints"
+  shared_ac_key: "actions_eva_cart_aria_keypoints"
+
+  6dof: true
+  diffusion: true
+
+  head_specs:
+    aria_bimanual: null
+    eva_bimanual: null
+    shared:
+      _target_: egomimic.models.fm_policy.FMPolicy
+      action_horizon: 100
+      num_inference_steps: 50
+      pooling: null
+      padding: "zero"
+      time_dist: "beta"
+      model:
+        _target_: egomimic.models.denoising_nets.CrossTransformer
+        nblocks: 8
+        cond_dim: 256
+        hidden_dim: 256
+        act_dim: 256
+        act_seq: 32
+        n_heads: 4
+        dropout: 0.1
+        mlp_layers: 4
+        mlp_ratio: 4
+      embodiment_specs:
+        aria_bimanual:
+          ac_dims: 140
+          encoder:
+            _target_: egomimic.models.codec.identity.Identity
+          decoder:
+            _target_: egomimic.models.codec.identity.Identity
+  stem_specs:
+    eva_bimanual: null
+
+optimizer:
+  _target_: torch.optim.AdamW
+  _partial_: true
+  lr: 3e-4
+  weight_decay: 0.0001
+
+scheduler:
+  _target_: torch.optim.lr_scheduler.CosineAnnealingLR
+  _partial_: true
+  T_max: 1800
+  eta_min: 1e-5