wip - still need to verify full training run

2024-03-11 18:45:21 +00:00
parent 304355c917
commit 87fcc536f9
3 changed files with 9 additions and 7 deletions
--- a/lerobot/common/envs/pusht/pusht_image_env.py
+++ b/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