Add Prod transform, Add test_factory

lerobot/common/envs/factory.py

@@ -2,6 +2,7 @@ from torchrl.envs.transforms import StepCounter, TransformedEnv
 
 from lerobot.common.envs.pusht import PushtEnv
 from lerobot.common.envs.simxarm import SimxarmEnv
+from lerobot.common.envs.transforms import Prod
 
 
 def make_env(cfg):
@@ -25,6 +26,10 @@ def make_env(cfg):
     # limit rollout to max_steps
     env = TransformedEnv(env, StepCounter(max_steps=cfg.episode_length))
 
+    if cfg.env == "pusht":
+        # to ensure pusht is in [0,255] like simxarm
+        env.append_transform(Prod(in_keys=[("observation", "image")], prod=255.0))
+
     return env
 
 

lerobot/common/envs/pusht.py

@@ -48,23 +48,16 @@ class PushtEnv(EnvBase):
         if not _has_gym:
             raise ImportError("Cannot import gym.")
 
-        import gym
         from diffusion_policy.env.pusht.pusht_env import PushTEnv
+
+        if not from_pixels:
+            raise NotImplementedError("Use PushTEnv, instead of PushTImageEnv")
         from diffusion_policy.env.pusht.pusht_image_env import PushTImageEnv
         from gym.wrappers import TimeLimit
 
         self._env = PushTImageEnv(render_size=self.image_size)
         self._env = TimeLimit(self._env, self.max_episode_length)
 
-        # MAX_NUM_ACTIONS = 4
-        # num_actions = len(TASKS[self.task]["action_space"])
-        # self._action_space = gym.spaces.Box(low=-1.0, high=1.0, shape=(num_actions,))
-        # self._action_padding = np.zeros(
-        #     (MAX_NUM_ACTIONS - num_actions), dtype=np.float32
-        # )
-        # if "w" not in TASKS[self.task]["action_space"]:
-        #     self._action_padding[-1] = 1.0
-
         self._make_spec()
         self.set_seed(seed)
 

lerobot/common/envs/transforms.py

@@ -0,0 +1,22 @@
+from typing import Sequence
+
+from tensordict.utils import NestedKey
+from torchrl.envs.transforms import ObservationTransform
+
+
+class Prod(ObservationTransform):
+
+    def __init__(self, in_keys: Sequence[NestedKey], prod: float):
+        super().__init__()
+        self.in_keys = in_keys
+        self.prod = prod
+
+    def _call(self, td):
+        for key in self.in_keys:
+            td[key] *= self.prod
+        return td
+
+    def transform_observation_spec(self, obs_spec):
+        for key in self.in_keys:
+            obs_spec[key].space.high *= self.prod
+        return obs_spec