Don't use batch dimension in normalization

lerobot/common/datasets/abstract.py

@@ -49,9 +49,9 @@ class AbstractExperienceReplay(TensorDictReplayBuffer):
     @property
     def stats_patterns(self) -> dict:
         return {
-            ("observation", "state"): "b c -> 1 c",
-            ("observation", "image"): "b c h w -> 1 c 1 1",
-            ("action",): "b c -> 1 c",
+            ("observation", "state"): "b c -> c",
+            ("observation", "image"): "b c h w -> c 1 1",
+            ("action",): "b c -> c",
         }
 
     @property

lerobot/common/datasets/aloha.py

@@ -113,11 +113,11 @@ class AlohaExperienceReplay(AbstractExperienceReplay):
     @property
     def stats_patterns(self) -> dict:
         d = {
-            ("observation", "state"): "b c -> 1 c",
-            ("action",): "b c -> 1 c",
+            ("observation", "state"): "b c -> c",
+            ("action",): "b c -> c",
         }
         for cam in CAMERAS[self.dataset_id]:
-            d[("observation", "image", cam)] = "b c h w -> 1 c 1 1"
+            d[("observation", "image", cam)] = "b c h w -> c 1 1"
         return d
 
     @property