wip - still need to verify full training run

lerobot/common/envs/pusht/pusht_image_env.py

@@ -25,7 +25,7 @@ class PushTImageEnv(PushTEnv):
         img = super()._render_frame(mode="rgb_array")
 
         agent_pos = np.array(self.agent.position)
-        img_obs = np.moveaxis(img.astype(np.float32) / 255, -1, 0)
+        img_obs = np.moveaxis(img.astype(np.float32), -1, 0)
         obs = {"image": img_obs, "agent_pos": agent_pos}
 
         # draw action

lerobot/common/policies/diffusion/model/multi_image_obs_encoder.py

@@ -123,6 +123,8 @@ class MultiImageObsEncoder(ModuleAttrMixin):
                 if imagenet_norm:
                     # TODO(rcadene): move normalizer to dataset and env
                     this_normalizer = torchvision.transforms.Normalize(
+                        # Note: This matches the normalization in the original impl. for PushT Image. This may not be
+                        # the case for other tasks.
                         mean=[127.5, 127.5, 127.5],
                         std=[127.5, 127.5, 127.5],
                     )