pi zarr training configs

egomimic/algo/pi.py

@@ -86,19 +86,19 @@ def __init__(
             self.camera_keys[embodiment_id] = []
             self.proprio_keys[embodiment_id] = []
             self.lang_keys[embodiment_id] = []
-            for key in data_schematic.keys_of_type("action_keys"):
+            for key in data_schematic.keys_of_type("action_keys", embodiment_id):
                 if (
                     data_schematic.is_key_with_embodiment(key, embodiment_id)
                     and key == self.ac_keys[embodiment]
                 ):
                     self.ac_keys[embodiment_id] = key
-            for key in data_schematic.keys_of_type("camera_keys"):
+            for key in data_schematic.keys_of_type("camera_keys", embodiment_id):
                 if data_schematic.is_key_with_embodiment(key, embodiment_id):
                     self.camera_keys[embodiment_id].append(key)
-            for key in data_schematic.keys_of_type("proprio_keys"):
+            for key in data_schematic.keys_of_type("proprio_keys", embodiment_id):
                 if data_schematic.is_key_with_embodiment(key, embodiment_id):
                     self.proprio_keys[embodiment_id].append(key)
-            for key in data_schematic.keys_of_type("lang_keys"):
+            for key in data_schematic.keys_of_type("lang_keys", embodiment_id):
                 if data_schematic.is_key_with_embodiment(key, embodiment_id):
                     self.lang_keys[embodiment_id].append(key)
 
@@ -165,12 +165,11 @@ def process_batch_for_training(self, batch):
         """
         processed_batch = {}
 
-        for embodiment_id, _batch in batch.items():
+        for embodiment_name, _batch in batch.items():
+            embodiment_id = get_embodiment_id(embodiment_name)
             processed_batch[embodiment_id] = {}
             for key, value in _batch.items():
-                key_name = self.data_schematic.lerobot_key_to_keyname(
-                    key, embodiment_id
-                )
+                key_name = self.data_schematic.zarr_key_to_keyname(key, embodiment_id)
                 if key_name is not None:
                     processed_batch[embodiment_id][key_name] = value
 
@@ -196,6 +195,13 @@ def process_batch_for_training(self, batch):
             processed_batch[embodiment_id] = self.data_schematic.normalize_data(
                 processed_batch[embodiment_id], embodiment_id
             )
+            processed_batch[embodiment_id]["embodiment"] = torch.tensor(
+                [embodiment_id], device=device, dtype=torch.int64
+            )
+
+            for key, value in processed_batch[embodiment_id].items():
+                if isinstance(value, torch.Tensor) and value.dtype == torch.float64:
+                    processed_batch[embodiment_id][key] = value.float()
 
         if not processed_batch:
             raise ValueError(

egomimic/hydra_configs/hydra/launcher/submitit_pace.yaml

@@ -11,7 +11,7 @@ cpus_per_task: 8                                   # Number of CPUs per task (ma
 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/model/pi0.5.yaml → egomimic/hydra_configs/model/pi0.5_bc_aria.yaml

@@ -20,7 +20,7 @@ robomimic_model:
 
   config:      
     pytorch_training_precision: bfloat16
-    pytorch_weight_path: /storage/home/hcoda1/5/rpunamiya6/cedar-dx/rpunamiya6/Projects/EgoVerse/egomimic/algo/pi_checkpoints/pi05_base_pytorch
+    pytorch_weight_path: /storage/project/r-dxu345-0/rco3/EgoVerse/egomimic/algo/pi_checkpoints/pi05_base_pytorch
     model:
       pi05: true
       action_dim: 32

egomimic/hydra_configs/model/pi0.5_bc_eva.yaml

@@ -0,0 +1,56 @@
+_target_: egomimic.pl_utils.pl_model.ModelWrapper
+robomimic_model:
+  _target_: egomimic.algo.pi.PI
+  data_schematic: _${data.dataset.data_schematic}
+  camera_transforms:
+    _target_: egomimic.utils.egomimicUtils.CameraTransforms
+    intrinsics_key: "base" # change to base_half if using half res 
+    extrinsics_key: "x5Dec13_2"
+  ac_keys:
+    eva_bimanual: "actions_cartesian"
+  domains: ["eva_bimanual"]
+
+  action_converters:
+    rules:
+      EVA_BIMANUAL:
+        _target_: egomimic.utils.action_utils.RobotBimanualCartesianEuler
+    # optional fallback if no match is found
+    fallback:
+      _target_: egomimic.utils.action_utils.BaseActionConverter
+
+  config:      
+    pytorch_training_precision: bfloat16
+    pytorch_weight_path: /storage/project/r-dxu345-0/rco3/EgoVerse/egomimic/algo/pi_checkpoints/pi05_base_pytorch
+    model:
+      pi05: true
+      action_dim: 32
+      action_horizon: 100
+      max_token_len: 180
+
+  train_image_augs:
+    _target_: torchvision.transforms.Compose
+    transforms:
+      - _target_: torchvision.transforms.Resize
+        size: 224
+        interpolation: 3
+  eval_image_augs:
+    _target_: torchvision.transforms.Compose
+    transforms:
+      - _target_: torchvision.transforms.Resize
+        size: 224
+        interpolation: 3
+
+optimizer:
+  _target_: torch.optim.AdamW
+  _partial_: true
+  lr: 5e-4
+  betas: [0.9, 0.999]
+  eps: 1e-8
+  weight_decay: 0.0
+
+scheduler:
+  _target_: transformers.get_cosine_schedule_with_warmup
+  _partial_: true
+  num_warmup_steps: 1000
+  num_training_steps: 100000
+  num_cycles: 0.5

egomimic/hydra_configs/model/pi0.5_bc_mecka.yaml

@@ -0,0 +1,56 @@
+_target_: egomimic.pl_utils.pl_model.ModelWrapper
+robomimic_model:
+  _target_: egomimic.algo.pi.PI
+  data_schematic: _${data.dataset.data_schematic}
+  camera_transforms:
+    _target_: egomimic.utils.egomimicUtils.CameraTransforms
+    intrinsics_key: "mecka" # change to base_half if using half res 
+    extrinsics_key: "mecka"
+  ac_keys:
+    mecka_bimanual: "actions_cartesian"
+  domains: ["mecka_bimanual"]
+
+  action_converters:
+    rules:
+      MECKA_BIMANUAL:
+        _target_: egomimic.utils.action_utils.HumanBimanualCartesianEuler
+    # optional fallback if no match is found
+    fallback:
+      _target_: egomimic.utils.action_utils.BaseActionConverter
+
+  config:      
+    pytorch_training_precision: bfloat16
+    pytorch_weight_path: /storage/project/r-dxu345-0/rco3/EgoVerse/egomimic/algo/pi_checkpoints/pi05_base_pytorch
+    model:
+      pi05: true
+      action_dim: 32
+      action_horizon: 100
+      max_token_len: 180
+
+  train_image_augs:
+    _target_: torchvision.transforms.Compose
+    transforms:
+      - _target_: torchvision.transforms.Resize
+        size: 224
+        interpolation: 3
+  eval_image_augs:
+    _target_: torchvision.transforms.Compose
+    transforms:
+      - _target_: torchvision.transforms.Resize
+        size: 224
+        interpolation: 3
+
+optimizer:
+  _target_: torch.optim.AdamW
+  _partial_: true
+  lr: 5e-4
+  betas: [0.9, 0.999]
+  eps: 1e-8
+  weight_decay: 0.0
+
+scheduler:
+  _target_: transformers.get_cosine_schedule_with_warmup
+  _partial_: true
+  num_warmup_steps: 1000
+  num_training_steps: 100000
+  num_cycles: 0.5

egomimic/hydra_configs/model/pi0.5_bc_scale.yaml

@@ -0,0 +1,56 @@
+_target_: egomimic.pl_utils.pl_model.ModelWrapper
+robomimic_model:
+  _target_: egomimic.algo.pi.PI
+  data_schematic: _${data.dataset.data_schematic}
+  camera_transforms:
+    _target_: egomimic.utils.egomimicUtils.CameraTransforms
+    intrinsics_key: "scale" # change to base_half if using half res 
+    extrinsics_key: "scale"
+  ac_keys:
+    scale_bimanual: "actions_cartesian"
+  domains: ["scale_bimanual"]
+
+  action_converters:
+    rules:
+      SCALE_BIMANUAL:
+        _target_: egomimic.utils.action_utils.HumanBimanualCartesianEuler
+    # optional fallback if no match is found
+    fallback:
+      _target_: egomimic.utils.action_utils.BaseActionConverter
+
+  config:      
+    pytorch_training_precision: bfloat16
+    pytorch_weight_path: /storage/project/r-dxu345-0/rco3/EgoVerse/egomimic/algo/pi_checkpoints/pi05_base_pytorch
+    model:
+      pi05: true
+      action_dim: 32
+      action_horizon: 100
+      max_token_len: 180
+
+  train_image_augs:
+    _target_: torchvision.transforms.Compose
+    transforms:
+      - _target_: torchvision.transforms.Resize
+        size: 224
+        interpolation: 3
+  eval_image_augs:
+    _target_: torchvision.transforms.Compose
+    transforms:
+      - _target_: torchvision.transforms.Resize
+        size: 224
+        interpolation: 3
+
+optimizer:
+  _target_: torch.optim.AdamW
+  _partial_: true
+  lr: 5e-4
+  betas: [0.9, 0.999]
+  eps: 1e-8
+  weight_decay: 0.0
+
+scheduler:
+  _target_: transformers.get_cosine_schedule_with_warmup
+  _partial_: true
+  num_warmup_steps: 1000
+  num_training_steps: 100000
+  num_cycles: 0.5

egomimic/hydra_configs/model/pi0.5_cotrain_eva_aria.yaml

@@ -0,0 +1,64 @@
+_target_: egomimic.pl_utils.pl_model.ModelWrapper
+robomimic_model:
+  _target_: egomimic.algo.pi.PI
+  data_schematic: _${data.dataset.data_schematic}
+  camera_transforms:
+    eva_bimanual:
+      _target_: egomimic.utils.egomimicUtils.CameraTransforms
+      intrinsics_key: "base" # change to base_half if using half res 
+      extrinsics_key: "x5Dec13_2"
+    aria_bimanual:
+      _target_: egomimic.utils.egomimicUtils.CameraTransforms
+      intrinsics_key: "base" # change to base_half if using half res 
+      extrinsics_key: "ariaJun7"
+  ac_keys:
+    eva_bimanual: "actions_cartesian"
+    aria_bimanual: "actions_cartesian"
+  domains: ["eva_bimanual", "aria_bimanual"]
+
+  action_converters:
+    rules:
+      EVA_BIMANUAL:
+        _target_: egomimic.utils.action_utils.RobotBimanualCartesianEuler
+      ARIA_BIMANUAL:
+        _target_: egomimic.utils.action_utils.HumanBimanualCartesianEuler
+    # optional fallback if no match is found
+    fallback:
+      _target_: egomimic.utils.action_utils.BaseActionConverter
+
+  config:      
+    pytorch_training_precision: bfloat16
+    pytorch_weight_path: /storage/project/r-dxu345-0/rco3/EgoVerse/egomimic/algo/pi_checkpoints/pi05_base_pytorch
+    model:
+      pi05: true
+      action_dim: 32
+      action_horizon: 100
+      max_token_len: 180
+
+  train_image_augs:
+    _target_: torchvision.transforms.Compose
+    transforms:
+      - _target_: torchvision.transforms.Resize
+        size: 224
+        interpolation: 3
+  eval_image_augs:
+    _target_: torchvision.transforms.Compose
+    transforms:
+      - _target_: torchvision.transforms.Resize
+        size: 224
+        interpolation: 3
+
+optimizer:
+  _target_: torch.optim.AdamW
+  _partial_: true
+  lr: 5e-4
+  betas: [0.9, 0.999]
+  eps: 1e-8
+  weight_decay: 0.0
+
+scheduler:
+  _target_: transformers.get_cosine_schedule_with_warmup
+  _partial_: true
+  num_warmup_steps: 1000
+  num_training_steps: 100000
+  num_cycles: 0.5