Fix unit tests, Refactor, Add pusht env, (TODO pusht replay buffer, image preprocessing)

lerobot/common/envs/simxarm.py

@@ -167,18 +167,21 @@ class SimxarmEnv(EnvBase):
             device=self.device,
         )
 
-        self.done_spec = DiscreteTensorSpec(
-            2,
-            shape=(1,),
-            dtype=torch.bool,
-            device=self.device,
-        )
-
-        self.success_spec = DiscreteTensorSpec(
-            2,
-            shape=(1,),
-            dtype=torch.bool,
-            device=self.device,
+        self.done_spec = CompositeSpec(
+            {
+                "done": DiscreteTensorSpec(
+                    2,
+                    shape=(1,),
+                    dtype=torch.bool,
+                    device=self.device,
+                ),
+                "success": DiscreteTensorSpec(
+                    2,
+                    shape=(1,),
+                    dtype=torch.bool,
+                    device=self.device,
+                ),
+            }
         )
 
     def _set_seed(self, seed: Optional[int]):