Update pipeline.py

pipeline.py

@@ -14,8 +14,17 @@ class ValueGuidedDiffuserPipeline(DiffusionPipeline):
         self.scheduler = scheduler
         self.env = env
         self.data = env.get_dataset()
-        self.means = dict((key, val.mean(axis=0)) for key, val in self.data.items())
-        self.stds = dict((key, val.std(axis=0)) for key, val in self.data.items())
+        for key in self.data.keys():
+            try:
+                self.means[key] = self.data[key].mean()
+            except AxisError:
+                pass
+        self.stds = dict()
+        for key in self.data.keys():
+            try:
+                self.stds[key] = self.data[key].std()
+            except AxisError:
+                pass
         self.device = self.unet.device
         self.state_dim = env.observation_space.shape[0]
         self.action_dim = env.action_space.shape[0]