adding support for single arm training

egomimic/hydra_configs/data/eva_right_arm.yaml

@@ -0,0 +1,41 @@
+_target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper
+train_datasets:
+  eva_right_arm:
+    _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
+        arm_mode: right_arm
+      transform_list:
+        _target_: egomimic.rldb.embodiment.eva.Eva.get_transform_list
+        arm_mode: right_arm
+    filters:
+      robot_name: "eva_right_arm"
+    mode: total
+
+valid_datasets:
+  eva_right_arm:
+    _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
+        arm_mode: right_arm
+      transform_list:
+        _target_: egomimic.rldb.embodiment.eva.Eva.get_transform_list
+        arm_mode: right_arm
+    filters:
+      robot_name: "eva_right_arm"
+    mode: total
+
+train_dataloader_params:
+  eva_right_arm:
+    batch_size: 32
+    num_workers: 10
+valid_dataloader_params:
+  eva_right_arm:
+    batch_size: 32
+    num_workers: 10

egomimic/hydra_configs/model/hpt_bc_flow_eva_right_arm.yaml

@@ -0,0 +1,142 @@
+_target_: egomimic.pl_utils.pl_model.ModelWrapper
+robomimic_model:
+  _target_: egomimic.algo.hpt.HPT
+  data_schematic: _${data.dataset.data_schematic}
+  camera_transforms:
+    eva_right_arm:
+      _target_: egomimic.utils.egomimicUtils.CameraTransforms
+      intrinsics_key: "base" # change to base_half if using half res 
+      extrinsics_key: "x5Dec13_2"
+
+  diffusion: true 
+  6dof: true
+
+  ac_keys:
+    eva_right_arm: "actions_cartesian"
+
+  kinematics_solver:
+    _target_: egomimic.robot.eva.eva_kinematics.EvaMinkKinematicsSolver
+    model_path: /coc/flash7/rpunamiya6/Projects/EgoVerse/egomimic/resources/model_x5.xml
+
+  trunk:
+    embed_dim: 256
+    num_blocks: 16
+    num_heads: 8
+    token_postprocessing: "action_token"
+    observation_horizon: 1
+    action_horizon: 64
+    no_trunk: false
+    use_domain_embedding: true
+    drop_path: 0.1
+    weight_init_style: "pytorch"
+
+  multitask: false
+  pretrained: false
+  pretrained_checkpoint: "" # TODO
+  reverse_kl_samples: 8
+
+  domains: ["eva_right_arm"]
+  shared_obs_keys: ["front_img_1"]
+
+  shared_stem_specs:
+    front_img_1:
+      _target_: egomimic.models.hpt_nets.MLPPolicyStem
+      input_dim: 256
+      output_dim: 256
+      widths: [256]
+      specs:
+        random_horizon_masking: false
+        cross_attn:
+          crossattn_latent: 16
+          crossattn_heads: 8
+          crossattn_dim_head: 64
+          crossattn_modality_dropout: 0.1
+          modality_embed_dim: 256
+
+  stem_specs:
+    eva_right_arm:
+      right_wrist_img:
+        _target_: egomimic.models.hpt_nets.MLPPolicyStem
+        input_dim: 256
+        output_dim: 256
+        widths: [256]
+        specs:
+          random_horizon_masking: false
+          cross_attn:
+            crossattn_latent: 16
+            crossattn_heads: 8
+            crossattn_dim_head: 64
+            crossattn_modality_dropout: 0.1
+            modality_embed_dim: 256
+      state_ee_pose:
+        _target_: egomimic.models.hpt_nets.MLPPolicyStem
+        input_dim: 7  # single arm: xyz + ypr + gripper (bimanual would be 14)
+        output_dim: 256
+        widths: [256]
+        specs:
+          random_horizon_masking: false
+          cross_attn:
+            crossattn_latent: 16
+            crossattn_heads: 8
+            crossattn_dim_head: 64
+            crossattn_modality_dropout: 0.1
+            modality_embed_dim: 256
+
+  head_specs:
+    eva_right_arm: 
+      _target_: egomimic.models.fm_policy.FMPolicy
+      action_horizon: 100
+      num_inference_steps: 50
+      pooling: null
+      time_dist: "beta"
+      infer_ac_dims:
+        eva_right_arm: 7
+      model:
+        _target_: egomimic.models.denoising_nets.CrossTransformer
+        nblocks: 6
+        cond_dim: 256
+        hidden_dim: 128
+        act_dim: 7
+        act_seq: 100
+        n_heads: 4
+        dropout: 0.1
+        mlp_layers: 4
+        mlp_ratio: 4
+
+  encoder_specs:
+    front_img_1:
+      _target_: egomimic.models.hpt_nets.ResNet
+      output_dim: 256
+    right_wrist_img:
+      _target_: egomimic.models.hpt_nets.ResNet
+      output_dim: 256
+
+  train_image_augs:
+    _target_: torchvision.transforms.Compose
+    transforms:
+      - _target_: torchvision.transforms.ColorJitter
+        brightness: 0.1
+        contrast: 0.1
+        saturation: 0.1
+        hue: 0.05
+      - _target_: torchvision.transforms.Normalize
+        mean: [0.485, 0.456, 0.406]
+        std: [0.229, 0.224, 0.225]
+  eval_image_augs:
+    _target_: torchvision.transforms.Compose
+    transforms:
+      - _target_: torchvision.transforms.Normalize
+        mean: [0.485, 0.456, 0.406]
+        std: [0.229, 0.224, 0.225]
+
+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: 1400
+  eta_min: 1e-5

egomimic/hydra_configs/train_zarr.yaml

@@ -6,7 +6,7 @@ defaults:
   - logger: wandb
   - data: eva
   - callbacks: checkpoints
-  - override hydra/launcher: submitit
+  - override hydra/launcher: submitit_skynet
   - _self_
 
 name: test
@@ -62,6 +62,50 @@ data_schematic: # Dynamically fill in these shapes from the dataset
       embodiment:
         key_type: metadata_keys
         zarr_key: metadata.embodiment
+    eva_right_arm:
+      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
+      ee_pose:
+        key_type: proprio_keys
+        zarr_key: observations.state.ee_pose
+      joint_positions:
+        key_type: proprio_keys
+        zarr_key: observations.state.joint_positions
+      actions_joints:
+        key_type: action_keys
+        zarr_key: actions_joints
+      actions_cartesian:
+        key_type: action_keys
+        zarr_key: actions_cartesian
+      embodiment:
+        key_type: metadata_keys
+        zarr_key: metadata.embodiment
+    eva_left_arm:
+      front_img_1: #batch key
+        key_type: camera_keys # key type
+        zarr_key: observations.images.front_img_1 # dataset key
+      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
+      joint_positions:
+        key_type: proprio_keys
+        zarr_key: observations.state.joint_positions
+      actions_joints:
+        key_type: action_keys
+        zarr_key: actions_joints
+      actions_cartesian:
+        key_type: action_keys
+        zarr_key: actions_cartesian
+      embodiment:
+        key_type: metadata_keys
+        zarr_key: metadata.embodiment
     aria_bimanual:
       front_img_1:
         key_type: camera_keys
@@ -102,13 +146,11 @@ data_schematic: # Dynamically fill in these shapes from the dataset
         key_type: metadata_keys
         zarr_key: metadata.embodiment
   viz_img_key:
-    eva_bimanual:
-      front_img_1
-    aria_bimanual:
-      front_img_1
-    mecka_bimanual:
-      front_img_1
-    scale_bimanual:
-      front_img_1
+    eva_right_arm: front_img_1
+    eva_left_arm: front_img_1
+    eva_bimanual: front_img_1
+    aria_bimanual: front_img_1
+    mecka_bimanual: front_img_1
+    scale_bimanual: front_img_1
 
 seed: 42