From c5635b7d948a66a1cc87355f9a04665b10a96455 Mon Sep 17 00:00:00 2001 From: Simon Alibert Date: Mon, 25 Mar 2024 18:50:47 +0100 Subject: [PATCH] Minor fixes for #47 --- lerobot/common/datasets/abstract.py | 3 +-- lerobot/common/envs/aloha/env.py | 2 +- lerobot/common/envs/pusht/env.py | 2 +- lerobot/common/envs/simxarm/env.py | 10 +++------- lerobot/common/envs/simxarm/simxarm/__init__.py | 10 +++++----- .../envs/simxarm/simxarm/{task => tasks}/__init__.py | 0 .../simxarm/simxarm/{task => tasks}/assets/lift.xml | 0 .../simxarm/{task => tasks}/assets/mesh/base_link.stl | 0 .../{task => tasks}/assets/mesh/block_inner.stl | 0 .../{task => tasks}/assets/mesh/block_inner2.stl | 0 .../{task => tasks}/assets/mesh/block_outer.stl | 0 .../{task => tasks}/assets/mesh/left_finger.stl | 0 .../{task => tasks}/assets/mesh/left_inner_knuckle.stl | 0 .../{task => tasks}/assets/mesh/left_outer_knuckle.stl | 0 .../simxarm/{task => tasks}/assets/mesh/link1.stl | 0 .../simxarm/{task => tasks}/assets/mesh/link2.stl | 0 .../simxarm/{task => tasks}/assets/mesh/link3.stl | 0 .../simxarm/{task => tasks}/assets/mesh/link4.stl | 0 .../simxarm/{task => tasks}/assets/mesh/link5.stl | 0 .../simxarm/{task => tasks}/assets/mesh/link6.stl | 0 .../simxarm/{task => tasks}/assets/mesh/link7.stl | 0 .../simxarm/{task => tasks}/assets/mesh/link_base.stl | 0 .../{task => tasks}/assets/mesh/right_finger.stl | 0 .../assets/mesh/right_inner_knuckle.stl | 0 .../assets/mesh/right_outer_knuckle.stl | 0 .../simxarm/{task => tasks}/assets/peg_in_box.xml | 0 .../simxarm/simxarm/{task => tasks}/assets/push.xml | 0 .../simxarm/simxarm/{task => tasks}/assets/reach.xml | 0 .../simxarm/simxarm/{task => tasks}/assets/shared.xml | 0 .../simxarm/simxarm/{task => tasks}/assets/xarm.xml | 0 .../envs/simxarm/simxarm/{task => tasks}/base.py | 2 +- .../envs/simxarm/simxarm/{task => tasks}/lift.py | 0 .../envs/simxarm/simxarm/{task => tasks}/mocap.py | 0 .../envs/simxarm/simxarm/{task => tasks}/peg_in_box.py | 0 .../envs/simxarm/simxarm/{task => tasks}/push.py | 0 .../envs/simxarm/simxarm/{task => tasks}/reach.py | 0 lerobot/scripts/train.py | 2 +- 37 files changed, 13 insertions(+), 18 deletions(-) rename lerobot/common/envs/simxarm/simxarm/{task => tasks}/__init__.py (100%) rename lerobot/common/envs/simxarm/simxarm/{task => tasks}/assets/lift.xml (100%) rename lerobot/common/envs/simxarm/simxarm/{task => tasks}/assets/mesh/base_link.stl (100%) rename lerobot/common/envs/simxarm/simxarm/{task => tasks}/assets/mesh/block_inner.stl (100%) rename lerobot/common/envs/simxarm/simxarm/{task => tasks}/assets/mesh/block_inner2.stl (100%) rename lerobot/common/envs/simxarm/simxarm/{task => tasks}/assets/mesh/block_outer.stl (100%) rename lerobot/common/envs/simxarm/simxarm/{task => tasks}/assets/mesh/left_finger.stl (100%) rename lerobot/common/envs/simxarm/simxarm/{task => tasks}/assets/mesh/left_inner_knuckle.stl (100%) rename lerobot/common/envs/simxarm/simxarm/{task => tasks}/assets/mesh/left_outer_knuckle.stl (100%) rename lerobot/common/envs/simxarm/simxarm/{task => tasks}/assets/mesh/link1.stl (100%) rename lerobot/common/envs/simxarm/simxarm/{task => tasks}/assets/mesh/link2.stl (100%) rename lerobot/common/envs/simxarm/simxarm/{task => tasks}/assets/mesh/link3.stl (100%) rename lerobot/common/envs/simxarm/simxarm/{task => tasks}/assets/mesh/link4.stl (100%) rename lerobot/common/envs/simxarm/simxarm/{task => tasks}/assets/mesh/link5.stl (100%) rename lerobot/common/envs/simxarm/simxarm/{task => tasks}/assets/mesh/link6.stl (100%) rename lerobot/common/envs/simxarm/simxarm/{task => tasks}/assets/mesh/link7.stl (100%) rename lerobot/common/envs/simxarm/simxarm/{task => tasks}/assets/mesh/link_base.stl (100%) rename lerobot/common/envs/simxarm/simxarm/{task => tasks}/assets/mesh/right_finger.stl (100%) rename lerobot/common/envs/simxarm/simxarm/{task => tasks}/assets/mesh/right_inner_knuckle.stl (100%) rename lerobot/common/envs/simxarm/simxarm/{task => tasks}/assets/mesh/right_outer_knuckle.stl (100%) rename lerobot/common/envs/simxarm/simxarm/{task => tasks}/assets/peg_in_box.xml (100%) rename lerobot/common/envs/simxarm/simxarm/{task => tasks}/assets/push.xml (100%) rename lerobot/common/envs/simxarm/simxarm/{task => tasks}/assets/reach.xml (100%) rename lerobot/common/envs/simxarm/simxarm/{task => tasks}/assets/shared.xml (100%) rename lerobot/common/envs/simxarm/simxarm/{task => tasks}/assets/xarm.xml (100%) rename lerobot/common/envs/simxarm/simxarm/{task => tasks}/base.py (98%) rename lerobot/common/envs/simxarm/simxarm/{task => tasks}/lift.py (100%) rename lerobot/common/envs/simxarm/simxarm/{task => tasks}/mocap.py (100%) rename lerobot/common/envs/simxarm/simxarm/{task => tasks}/peg_in_box.py (100%) rename lerobot/common/envs/simxarm/simxarm/{task => tasks}/push.py (100%) rename lerobot/common/envs/simxarm/simxarm/{task => tasks}/reach.py (100%) diff --git a/lerobot/common/datasets/abstract.py b/lerobot/common/datasets/abstract.py index 2909507fd..8295ed480 100644 --- a/lerobot/common/datasets/abstract.py +++ b/lerobot/common/datasets/abstract.py @@ -21,7 +21,7 @@ class AbstractExperienceReplay(TensorDictReplayBuffer): def __init__( self, dataset_id: str, - version: str | None, + version: str | None = None, batch_size: int = None, *, shuffle: bool = True, @@ -32,7 +32,6 @@ class AbstractExperienceReplay(TensorDictReplayBuffer): collate_fn: Callable = None, writer: Writer = None, transform: "torchrl.envs.Transform" = None, - # storage = None, ): self.dataset_id = dataset_id self.version = version diff --git a/lerobot/common/envs/aloha/env.py b/lerobot/common/envs/aloha/env.py index 99f12cf0c..d38d7f018 100644 --- a/lerobot/common/envs/aloha/env.py +++ b/lerobot/common/envs/aloha/env.py @@ -63,7 +63,7 @@ class AlohaEnv(AbstractEnv): def _make_env(self): if not _has_gym: - raise ImportError("Cannot import gym.") + raise ImportError("Cannot import gymnasium.") if not self.from_pixels: raise NotImplementedError() diff --git a/lerobot/common/envs/pusht/env.py b/lerobot/common/envs/pusht/env.py index 6c9d211d0..ca39bf4d5 100644 --- a/lerobot/common/envs/pusht/env.py +++ b/lerobot/common/envs/pusht/env.py @@ -50,7 +50,7 @@ class PushtEnv(AbstractEnv): def _make_env(self): if not _has_gym: - raise ImportError("Cannot import gym.") + raise ImportError("Cannot import gymnasium.") # TODO(rcadene) (PushTEnv is similar to PushTImageEnv, but without the image rendering, it's faster to iterate on) # from lerobot.common.envs.pusht.pusht_env import PushTEnv diff --git a/lerobot/common/envs/simxarm/env.py b/lerobot/common/envs/simxarm/env.py index fc66d0130..f3c358d7a 100644 --- a/lerobot/common/envs/simxarm/env.py +++ b/lerobot/common/envs/simxarm/env.py @@ -49,10 +49,8 @@ class SimxarmEnv(AbstractEnv): ) def _make_env(self): - # if not _has_simxarm: - # raise ImportError("Cannot import simxarm.") if not _has_gym: - raise ImportError("Cannot import gym.") + raise ImportError("Cannot import gymnasium.") import gymnasium @@ -231,8 +229,6 @@ class SimxarmEnv(AbstractEnv): def _set_seed(self, seed: Optional[int]): set_global_seed(seed) - # self._env.seed(seed) - # self._env.action_space.seed(seed) - # self.set_seed(seed) - logging.warning("simxarm env is not seeded") self._seed = seed + # TODO(aliberts): change self._reset so that it takes in a seed value + logging.warning("simxarm env is not properly seeded") diff --git a/lerobot/common/envs/simxarm/simxarm/__init__.py b/lerobot/common/envs/simxarm/simxarm/__init__.py index 03928569f..903d60421 100644 --- a/lerobot/common/envs/simxarm/simxarm/__init__.py +++ b/lerobot/common/envs/simxarm/simxarm/__init__.py @@ -4,11 +4,11 @@ import gymnasium as gym import numpy as np from gymnasium.wrappers import TimeLimit -from lerobot.common.envs.simxarm.simxarm.task.base import Base as Base -from lerobot.common.envs.simxarm.simxarm.task.lift import Lift -from lerobot.common.envs.simxarm.simxarm.task.peg_in_box import PegInBox -from lerobot.common.envs.simxarm.simxarm.task.push import Push -from lerobot.common.envs.simxarm.simxarm.task.reach import Reach +from lerobot.common.envs.simxarm.simxarm.tasks.base import Base as Base +from lerobot.common.envs.simxarm.simxarm.tasks.lift import Lift +from lerobot.common.envs.simxarm.simxarm.tasks.peg_in_box import PegInBox +from lerobot.common.envs.simxarm.simxarm.tasks.push import Push +from lerobot.common.envs.simxarm.simxarm.tasks.reach import Reach TASKS = OrderedDict( ( diff --git a/lerobot/common/envs/simxarm/simxarm/task/__init__.py b/lerobot/common/envs/simxarm/simxarm/tasks/__init__.py similarity index 100% rename from lerobot/common/envs/simxarm/simxarm/task/__init__.py rename to lerobot/common/envs/simxarm/simxarm/tasks/__init__.py diff --git a/lerobot/common/envs/simxarm/simxarm/task/assets/lift.xml b/lerobot/common/envs/simxarm/simxarm/tasks/assets/lift.xml similarity index 100% rename from lerobot/common/envs/simxarm/simxarm/task/assets/lift.xml rename to lerobot/common/envs/simxarm/simxarm/tasks/assets/lift.xml diff --git a/lerobot/common/envs/simxarm/simxarm/task/assets/mesh/base_link.stl b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/base_link.stl similarity index 100% rename from lerobot/common/envs/simxarm/simxarm/task/assets/mesh/base_link.stl rename to lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/base_link.stl diff --git a/lerobot/common/envs/simxarm/simxarm/task/assets/mesh/block_inner.stl b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/block_inner.stl similarity index 100% rename from lerobot/common/envs/simxarm/simxarm/task/assets/mesh/block_inner.stl rename to lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/block_inner.stl diff --git a/lerobot/common/envs/simxarm/simxarm/task/assets/mesh/block_inner2.stl b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/block_inner2.stl similarity index 100% rename from lerobot/common/envs/simxarm/simxarm/task/assets/mesh/block_inner2.stl rename to lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/block_inner2.stl diff --git a/lerobot/common/envs/simxarm/simxarm/task/assets/mesh/block_outer.stl b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/block_outer.stl similarity index 100% rename from lerobot/common/envs/simxarm/simxarm/task/assets/mesh/block_outer.stl rename to lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/block_outer.stl diff --git a/lerobot/common/envs/simxarm/simxarm/task/assets/mesh/left_finger.stl b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/left_finger.stl similarity index 100% rename from lerobot/common/envs/simxarm/simxarm/task/assets/mesh/left_finger.stl rename to lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/left_finger.stl diff --git a/lerobot/common/envs/simxarm/simxarm/task/assets/mesh/left_inner_knuckle.stl b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/left_inner_knuckle.stl similarity index 100% rename from lerobot/common/envs/simxarm/simxarm/task/assets/mesh/left_inner_knuckle.stl rename to lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/left_inner_knuckle.stl diff --git a/lerobot/common/envs/simxarm/simxarm/task/assets/mesh/left_outer_knuckle.stl b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/left_outer_knuckle.stl similarity index 100% rename from lerobot/common/envs/simxarm/simxarm/task/assets/mesh/left_outer_knuckle.stl rename to lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/left_outer_knuckle.stl diff --git a/lerobot/common/envs/simxarm/simxarm/task/assets/mesh/link1.stl b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link1.stl similarity index 100% rename from lerobot/common/envs/simxarm/simxarm/task/assets/mesh/link1.stl rename to lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link1.stl diff --git a/lerobot/common/envs/simxarm/simxarm/task/assets/mesh/link2.stl b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link2.stl similarity index 100% rename from lerobot/common/envs/simxarm/simxarm/task/assets/mesh/link2.stl rename to lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link2.stl diff --git a/lerobot/common/envs/simxarm/simxarm/task/assets/mesh/link3.stl b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link3.stl similarity index 100% rename from lerobot/common/envs/simxarm/simxarm/task/assets/mesh/link3.stl rename to lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link3.stl diff --git a/lerobot/common/envs/simxarm/simxarm/task/assets/mesh/link4.stl b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link4.stl similarity index 100% rename from lerobot/common/envs/simxarm/simxarm/task/assets/mesh/link4.stl rename to lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link4.stl diff --git a/lerobot/common/envs/simxarm/simxarm/task/assets/mesh/link5.stl b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link5.stl similarity index 100% rename from lerobot/common/envs/simxarm/simxarm/task/assets/mesh/link5.stl rename to lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link5.stl diff --git a/lerobot/common/envs/simxarm/simxarm/task/assets/mesh/link6.stl b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link6.stl similarity index 100% rename from lerobot/common/envs/simxarm/simxarm/task/assets/mesh/link6.stl rename to lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link6.stl diff --git a/lerobot/common/envs/simxarm/simxarm/task/assets/mesh/link7.stl b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link7.stl similarity index 100% rename from lerobot/common/envs/simxarm/simxarm/task/assets/mesh/link7.stl rename to lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link7.stl diff --git a/lerobot/common/envs/simxarm/simxarm/task/assets/mesh/link_base.stl b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link_base.stl similarity index 100% rename from lerobot/common/envs/simxarm/simxarm/task/assets/mesh/link_base.stl rename to lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link_base.stl diff --git a/lerobot/common/envs/simxarm/simxarm/task/assets/mesh/right_finger.stl b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/right_finger.stl similarity index 100% rename from lerobot/common/envs/simxarm/simxarm/task/assets/mesh/right_finger.stl rename to lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/right_finger.stl diff --git a/lerobot/common/envs/simxarm/simxarm/task/assets/mesh/right_inner_knuckle.stl b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/right_inner_knuckle.stl similarity index 100% rename from lerobot/common/envs/simxarm/simxarm/task/assets/mesh/right_inner_knuckle.stl rename to lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/right_inner_knuckle.stl diff --git a/lerobot/common/envs/simxarm/simxarm/task/assets/mesh/right_outer_knuckle.stl b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/right_outer_knuckle.stl similarity index 100% rename from lerobot/common/envs/simxarm/simxarm/task/assets/mesh/right_outer_knuckle.stl rename to lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/right_outer_knuckle.stl diff --git a/lerobot/common/envs/simxarm/simxarm/task/assets/peg_in_box.xml b/lerobot/common/envs/simxarm/simxarm/tasks/assets/peg_in_box.xml similarity index 100% rename from lerobot/common/envs/simxarm/simxarm/task/assets/peg_in_box.xml rename to lerobot/common/envs/simxarm/simxarm/tasks/assets/peg_in_box.xml diff --git a/lerobot/common/envs/simxarm/simxarm/task/assets/push.xml b/lerobot/common/envs/simxarm/simxarm/tasks/assets/push.xml similarity index 100% rename from lerobot/common/envs/simxarm/simxarm/task/assets/push.xml rename to lerobot/common/envs/simxarm/simxarm/tasks/assets/push.xml diff --git a/lerobot/common/envs/simxarm/simxarm/task/assets/reach.xml b/lerobot/common/envs/simxarm/simxarm/tasks/assets/reach.xml similarity index 100% rename from lerobot/common/envs/simxarm/simxarm/task/assets/reach.xml rename to lerobot/common/envs/simxarm/simxarm/tasks/assets/reach.xml diff --git a/lerobot/common/envs/simxarm/simxarm/task/assets/shared.xml b/lerobot/common/envs/simxarm/simxarm/tasks/assets/shared.xml similarity index 100% rename from lerobot/common/envs/simxarm/simxarm/task/assets/shared.xml rename to lerobot/common/envs/simxarm/simxarm/tasks/assets/shared.xml diff --git a/lerobot/common/envs/simxarm/simxarm/task/assets/xarm.xml b/lerobot/common/envs/simxarm/simxarm/tasks/assets/xarm.xml similarity index 100% rename from lerobot/common/envs/simxarm/simxarm/task/assets/xarm.xml rename to lerobot/common/envs/simxarm/simxarm/tasks/assets/xarm.xml diff --git a/lerobot/common/envs/simxarm/simxarm/task/base.py b/lerobot/common/envs/simxarm/simxarm/tasks/base.py similarity index 98% rename from lerobot/common/envs/simxarm/simxarm/task/base.py rename to lerobot/common/envs/simxarm/simxarm/tasks/base.py index 589026010..b937b2901 100644 --- a/lerobot/common/envs/simxarm/simxarm/task/base.py +++ b/lerobot/common/envs/simxarm/simxarm/tasks/base.py @@ -4,7 +4,7 @@ import mujoco import numpy as np from gymnasium_robotics.envs import robot_env -from lerobot.common.envs.simxarm.simxarm.task import mocap +from lerobot.common.envs.simxarm.simxarm.tasks import mocap class Base(robot_env.MujocoRobotEnv): diff --git a/lerobot/common/envs/simxarm/simxarm/task/lift.py b/lerobot/common/envs/simxarm/simxarm/tasks/lift.py similarity index 100% rename from lerobot/common/envs/simxarm/simxarm/task/lift.py rename to lerobot/common/envs/simxarm/simxarm/tasks/lift.py diff --git a/lerobot/common/envs/simxarm/simxarm/task/mocap.py b/lerobot/common/envs/simxarm/simxarm/tasks/mocap.py similarity index 100% rename from lerobot/common/envs/simxarm/simxarm/task/mocap.py rename to lerobot/common/envs/simxarm/simxarm/tasks/mocap.py diff --git a/lerobot/common/envs/simxarm/simxarm/task/peg_in_box.py b/lerobot/common/envs/simxarm/simxarm/tasks/peg_in_box.py similarity index 100% rename from lerobot/common/envs/simxarm/simxarm/task/peg_in_box.py rename to lerobot/common/envs/simxarm/simxarm/tasks/peg_in_box.py diff --git a/lerobot/common/envs/simxarm/simxarm/task/push.py b/lerobot/common/envs/simxarm/simxarm/tasks/push.py similarity index 100% rename from lerobot/common/envs/simxarm/simxarm/task/push.py rename to lerobot/common/envs/simxarm/simxarm/tasks/push.py diff --git a/lerobot/common/envs/simxarm/simxarm/task/reach.py b/lerobot/common/envs/simxarm/simxarm/tasks/reach.py similarity index 100% rename from lerobot/common/envs/simxarm/simxarm/task/reach.py rename to lerobot/common/envs/simxarm/simxarm/tasks/reach.py diff --git a/lerobot/scripts/train.py b/lerobot/scripts/train.py index 3a45ddc9a..18c3715bf 100644 --- a/lerobot/scripts/train.py +++ b/lerobot/scripts/train.py @@ -231,7 +231,7 @@ def train(cfg: dict, out_dir=None, job_name=None): num_parallel_env = rollout.batch_size[0] if num_parallel_env != 1: - # TODO(rcadene): when num_parallel_env > 1, episode needs to be incremented and we need to add tests + # TODO(rcadene): when num_parallel_env > 1, rollout["episode"] needs to be properly set and we need to add tests raise NotImplementedError() num_max_steps = rollout.batch_size[1]