EgoMimic Changes

.gitignore

@@ -123,3 +123,5 @@ venv.bak/
 
 # private macros
 macros_private.py
+*.pyc
+act/detr/models/__pycache__

requirements.txt

@@ -10,4 +10,4 @@ imageio-ffmpeg
 matplotlib
 egl_probe>=1.0.1
 torch
-torchvision
+torchvision

robomimic/algo/algo.py

@@ -118,6 +118,7 @@ def __init__(
         self.global_config = global_config
 
         self.ac_dim = ac_dim
+        self.ac_key = global_config.train.ac_key
         self.device = device
         self.obs_key_shapes = obs_key_shapes
 
@@ -201,7 +202,7 @@ def process_batch_for_training(self, batch):
         """
         return batch
 
-    def postprocess_batch_for_training(self, batch, obs_normalization_stats):
+    def postprocess_batch_for_training(self, batch, normalization_stats, normalize_actions=True):
         """
         Does some operations (like channel swap, uint8 to float conversion, normalization)
         after @process_batch_for_training is called, in order to ensure these operations
@@ -222,7 +223,11 @@ def postprocess_batch_for_training(self, batch, obs_normalization_stats):
         """
 
         # ensure obs_normalization_stats are torch Tensors on proper device
-        obs_normalization_stats = TensorUtils.to_float(TensorUtils.to_device(TensorUtils.to_tensor(obs_normalization_stats), self.device))
+        normalization_stats = TensorUtils.to_float(
+            TensorUtils.to_device(
+                TensorUtils.to_tensor(normalization_stats), self.device
+            )
+        )
 
         # we will search the nested batch dictionary for the following special batch dict keys
         # and apply the processing function to their values (which correspond to observations)
@@ -236,14 +241,16 @@ def recurse_helper(d):
                 if k in obs_keys:
                     # found key - stop search and process observation
                     if d[k] is not None:
-                        d[k] = ObsUtils.process_obs_dict(d[k])
-                        if obs_normalization_stats is not None:
-                            d[k] = ObsUtils.normalize_obs(d[k], obs_normalization_stats=obs_normalization_stats)
+                        d[k] = ObsUtils.process_obs_dict(d[k], imagenet_normalize=self.global_config.train.imagenet_normalize_images)
                 elif isinstance(d[k], dict):
                     # search down into dictionary
                     recurse_helper(d[k])
 
         recurse_helper(batch)
+        if normalization_stats is not None:
+            batch = ObsUtils.normalize_batch(
+                batch, normalization_stats=normalization_stats, normalize_actions=normalize_actions
+            )
         return batch
 
     def train_on_batch(self, batch, epoch, validate=False):
@@ -502,8 +509,10 @@ def _prepare_observation(self, ob):
             # ensure obs_normalization_stats are torch Tensors on proper device
             obs_normalization_stats = TensorUtils.to_float(TensorUtils.to_device(TensorUtils.to_tensor(self.obs_normalization_stats), self.policy.device))
             # limit normalization to obs keys being used, in case environment includes extra keys
-            ob = { k : ob[k] for k in self.policy.global_config.all_obs_keys }
-            ob = ObsUtils.normalize_obs(ob, obs_normalization_stats=obs_normalization_stats)
+            ob = {k: ob[k] for k in self.policy.global_config.all_obs_keys}
+            ob = ObsUtils.normalize_batch(
+                ob, obs_normalization_stats=obs_normalization_stats
+            )
         return ob
 
     def __repr__(self):

robomimic/algo/bc.py

@@ -107,7 +107,8 @@ def process_batch_for_training(self, batch):
                 will be used for training 
         """
         input_batch = dict()
-        input_batch["obs"] = {k: batch["obs"][k][:, 0, :] for k in batch["obs"]}
+        #input_batch["obs"] = {k: batch["obs"][k][:, 0, :] for k in batch["obs"]}
+        input_batch["obs"] = {k: v[:, 0, :] if v.ndim != 1 else v for k, v in batch['obs'].items()}
         input_batch["goal_obs"] = batch.get("goal_obs", None) # goals may not be present
         input_batch["actions"] = batch["actions"][:, 0, :]
         # we move to device first before float conversion because image observation modalities will be uint8 -