Latent Flow matching that works regardless of dimension dim

egomimic/algo/hpt.py

@@ -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/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: /coc/flash7/scratch/egoverseDebugDatasets/egoverseS3DatasetTest/
+      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: /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:
+      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: /coc/flash7/scratch/egoverseDebugDatasets/egoverseS3DatasetTest/
+      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: /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:
+      robot_name: "aria_bimanual"
+      task: "fold_clothes_indomain"
+    mode: total
+train_dataloader_params:
+  eva_bimanual:
+    batch_size: 64
+    num_workers: 10
+  aria_bimanual:
+    batch_size: 64
+    num_workers: 10
+valid_dataloader_params:
+  eva_bimanual:
+    batch_size: 64
+    num_workers: 10
+  aria_bimanual:
+    batch_size: 64
+    num_workers: 10

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/model/hpt_cotrain_flow_shared_head_latent.yaml

@@ -0,0 +1,74 @@
+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: 6
+        cond_dim: 256
+        hidden_dim: 256
+        act_dim: 128
+        act_seq: 12
+        n_heads: 4
+        dropout: 0.1
+        mlp_layers: 4
+        mlp_ratio: 4
+      embodiment_specs:
+        eva_bimanual:
+          ac_dims: 14
+          encoder:
+            _target_: egomimic.models.codec.temporal_enc_dec.SmallTemporalEncoder
+            action_dim: 14
+            hidden_dim: 128
+            activation: "gelu"
+            use_layernorm: false
+          decoder:
+            _target_: egomimic.models.codec.temporal_enc_dec.SmallTemporalDecoder
+            action_dim: 14
+            hidden_dim: 128
+            activation: "gelu"
+            use_layernorm: true
+        aria_bimanual:
+          ac_dims: 140
+          encoder:
+            _target_: egomimic.models.codec.temporal_enc_dec.LargeTemporalEncoder
+            action_dim: 140
+            hidden_dim: 128
+            activation: "gelu"
+            use_layernorm: false
+          decoder:
+            _target_: egomimic.models.codec.temporal_enc_dec.LargeTemporalDecoder
+            action_dim: 140
+            hidden_dim: 128
+            activation: "gelu"
+            use_layernorm: true
+
+optimizer:
+  _target_: torch.optim.AdamW
+  _partial_: true
+  lr: 1e-4
+  weight_decay: 0.0001
+
+scheduler:
+  _target_: torch.optim.lr_scheduler.CosineAnnealingLR
+  _partial_: true
+  T_max: 1800
+  eta_min: 1e-5

egomimic/hydra_configs/model/hpt_cotrain_flow_shared_head_latent_large.yaml

@@ -0,0 +1,74 @@
+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:
+        eva_bimanual:
+          ac_dims: 14
+          encoder:
+            _target_: egomimic.models.codec.temporal_enc_dec.SmallTemporalEncoder_32_256
+            action_dim: 14
+            hidden_dim: 256
+            activation: "gelu"
+            use_layernorm: false
+          decoder:
+            _target_: egomimic.models.codec.temporal_enc_dec.SmallTemporalDecoder_32_256
+            action_dim: 14
+            hidden_dim: 256
+            activation: "gelu"
+            use_layernorm: true
+        aria_bimanual:
+          ac_dims: 140
+          encoder:
+            _target_: egomimic.models.codec.temporal_enc_dec.LargeTemporalEncoder_32_256
+            action_dim: 140
+            hidden_dim: 256
+            activation: "gelu"
+            use_layernorm: false
+          decoder:
+            _target_: egomimic.models.codec.temporal_enc_dec.LargeTemporalDecoder_32_256
+            action_dim: 140
+            hidden_dim: 256
+            activation: "gelu"
+            use_layernorm: true
+
+optimizer:
+  _target_: torch.optim.AdamW
+  _partial_: true
+  lr: 1e-4
+  weight_decay: 0.0001
+
+scheduler:
+  _target_: torch.optim.lr_scheduler.CosineAnnealingLR
+  _partial_: true
+  T_max: 1800
+  eta_min: 1e-5

egomimic/hydra_configs/model/hpt_cotrain_flow_shared_head_latent_mlp.yaml

@@ -0,0 +1,70 @@
+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: 100
+        n_heads: 4
+        dropout: 0.1
+        mlp_layers: 4
+        mlp_ratio: 4
+      embodiment_specs:
+        eva_bimanual:
+          ac_dims: 14
+          encoder:
+            _target_: egomimic.models.codec.mlp.MLPProjection
+            input_dim: 14
+            hidden_dim: 256
+            output_dim: 256
+          decoder:
+            _target_: egomimic.models.codec.mlp.MLPProjection
+            input_dim: 256
+            hidden_dim: 256
+            output_dim: 14
+        aria_bimanual:
+          ac_dims: 140
+          encoder:
+            _target_: egomimic.models.codec.mlp.MLPProjection
+            input_dim: 140
+            hidden_dim: 256
+            output_dim: 256
+          decoder:
+            _target_: egomimic.models.codec.mlp.MLPProjection
+            input_dim: 256
+            hidden_dim: 256
+            output_dim: 140
+
+optimizer:
+  _target_: torch.optim.AdamW
+  _partial_: true
+  lr: 1e-4
+  weight_decay: 0.0001
+
+scheduler:
+  _target_: torch.optim.lr_scheduler.CosineAnnealingLR
+  _partial_: true
+  T_max: 1800
+  eta_min: 1e-5