diff --git a/lerobot/common/envs/factory.py b/lerobot/common/envs/factory.py index acf089f6..984b866a 100644 --- a/lerobot/common/envs/factory.py +++ b/lerobot/common/envs/factory.py @@ -9,6 +9,7 @@ def make_env(cfg, transform=None): "image_size": cfg.env.image_size, # TODO(rcadene): do we want a specific eval_env_seed? "seed": cfg.seed, + "num_prev_obs": cfg.n_obs_steps - 1, } if cfg.env.name == "simxarm": diff --git a/lerobot/common/envs/pusht.py b/lerobot/common/envs/pusht.py index 62aa8d1b..7f9a3f63 100644 --- a/lerobot/common/envs/pusht.py +++ b/lerobot/common/envs/pusht.py @@ -2,6 +2,7 @@ import importlib from collections import deque from typing import Optional +import einops import torch from tensordict import TensorDict from torchrl.data.tensor_specs import ( @@ -28,7 +29,7 @@ class PushtEnv(EnvBase): image_size=None, seed=1337, device="cpu", - num_prev_obs=1, + num_prev_obs=0, num_prev_action=0, ): super().__init__(device=device, batch_size=[]) @@ -65,7 +66,8 @@ class PushtEnv(EnvBase): self._prev_obs_image_queue = deque(maxlen=self.num_prev_obs) self._prev_obs_state_queue = deque(maxlen=self.num_prev_obs) if self.num_prev_action > 0: - self._prev_action_queue = deque(maxlen=self.num_prev_action) + raise NotImplementedError() + # self._prev_action_queue = deque(maxlen=self.num_prev_action) def render(self, mode="rgb_array", width=384, height=384): if width != height: @@ -133,7 +135,7 @@ class PushtEnv(EnvBase): sum_reward = 0 if action.ndim == 1: - action = action.repeat(self.frame_skip, 1) + action = einops.repeat(action, "c -> t c", t=self.frame_skip) else: if self.frame_skip > 1: raise NotImplementedError() @@ -172,7 +174,7 @@ class PushtEnv(EnvBase): if self.from_pixels: image_shape = (3, self.image_size, self.image_size) if self.num_prev_obs > 0: - image_shape = (self.num_prev_obs, *image_shape) + image_shape = (self.num_prev_obs + 1, *image_shape) obs["image"] = BoundedTensorSpec( low=0, @@ -184,12 +186,12 @@ class PushtEnv(EnvBase): if not self.pixels_only: state_shape = self._env.observation_space["agent_pos"].shape if self.num_prev_obs > 0: - state_shape = (self.num_prev_obs, *state_shape) + state_shape = (self.num_prev_obs + 1, *state_shape) obs["state"] = BoundedTensorSpec( low=0, high=512, - shape=self._env.observation_space["agent_pos"].shape, + shape=state_shape, dtype=torch.float32, device=self.device, ) @@ -197,11 +199,11 @@ class PushtEnv(EnvBase): # TODO(rcadene): add observation_space achieved_goal and desired_goal? state_shape = self._env.observation_space["observation"].shape if self.num_prev_obs > 0: - state_shape = (self.num_prev_obs, *state_shape) + state_shape = (self.num_prev_obs + 1, *state_shape) obs["state"] = UnboundedContinuousTensorSpec( # TODO: - shape=self._env.observation_space["observation"].shape, + shape=state_shape, dtype=torch.float32, device=self.device, ) diff --git a/tests/test_datasets.py b/tests/test_datasets.py index b61873fb..a332a075 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -6,12 +6,17 @@ from .utils import init_config @pytest.mark.parametrize( - "env_name", + "env_name,dataset_id", [ - "simxarm", - "pusht", + # TODO(rcadene): simxarm is depreciated for now + # ("simxarm", "lift"), + ("pusht", "pusht"), + ("aloha", "sim_insertion_human"), + ("aloha", "sim_insertion_scripted"), + ("aloha", "sim_transfer_cube_human"), + ("aloha", "sim_transfer_cube_scripted"), ], ) -def test_factory(env_name): - cfg = init_config(overrides=[f"env={env_name}"]) +def test_factory(env_name, dataset_id): + cfg = init_config(overrides=[f"env={env_name}", f"env.task={dataset_id}"]) offline_buffer = make_offline_buffer(cfg) diff --git a/tests/test_envs.py b/tests/test_envs.py index d9fa4854..3382125a 100644 --- a/tests/test_envs.py +++ b/tests/test_envs.py @@ -36,6 +36,7 @@ def print_spec_rollout(env): print("data from rollout:", simple_rollout(100)) +@pytest.mark.skip(reason="Simxarm is deprecated") @pytest.mark.parametrize( "task,from_pixels,pixels_only", [ @@ -80,7 +81,7 @@ def test_pusht(from_pixels, pixels_only): @pytest.mark.parametrize( "env_name", [ - "simxarm", + # "simxarm", "pusht", ], )