diff --git a/lerobot/common/envs/aloha/env.py b/lerobot/common/envs/aloha.py similarity index 81% rename from lerobot/common/envs/aloha/env.py rename to lerobot/common/envs/aloha.py index 8f907650..55219dcd 100644 --- a/lerobot/common/envs/aloha/env.py +++ b/lerobot/common/envs/aloha.py @@ -6,8 +6,6 @@ from typing import Optional import einops import numpy as np import torch -from dm_control import mujoco -from dm_control.rl import control from tensordict import TensorDict from torchrl.data.tensor_specs import ( BoundedTensorSpec, @@ -17,21 +15,10 @@ from torchrl.data.tensor_specs import ( ) from lerobot.common.envs.abstract import AbstractEnv -from lerobot.common.envs.aloha.constants import ( - ACTIONS, - ASSETS_DIR, - DT, - JOINTS, -) -from lerobot.common.envs.aloha.tasks.sim import BOX_POSE, InsertionTask, TransferCubeTask -from lerobot.common.envs.aloha.tasks.sim_end_effector import ( - InsertionEndEffectorTask, - TransferCubeEndEffectorTask, -) -from lerobot.common.envs.aloha.utils import sample_box_pose, sample_insertion_pose from lerobot.common.utils import set_global_seed _has_gym = importlib.util.find_spec("gymnasium") is not None +_has_aloha = importlib.util.find_spec("aloha") is not None class AlohaEnv(AbstractEnv): @@ -70,43 +57,20 @@ class AlohaEnv(AbstractEnv): if not self.from_pixels: raise NotImplementedError() - self._env = self._make_env_task(self.task) + if not _has_aloha: + raise ImportError( + "Cannot import aloha env. Please install it with `python -m pip install 'lerobot[aloha]'`" + ) + + from aloha.env import make_env_task + + self._env = make_env_task(self.task) def render(self, mode="rgb_array", width=640, height=480): # TODO(rcadene): render and visualizer several cameras (e.g. angle, front_close) image = self._env.physics.render(height=height, width=width, camera_id="top") return image - def _make_env_task(self, task_name): - # time limit is controlled by StepCounter in env factory - time_limit = float("inf") - - if "sim_transfer_cube" in task_name: - xml_path = ASSETS_DIR / "bimanual_viperx_transfer_cube.xml" - physics = mujoco.Physics.from_xml_path(str(xml_path)) - task = TransferCubeTask(random=False) - elif "sim_insertion" in task_name: - xml_path = ASSETS_DIR / "bimanual_viperx_insertion.xml" - physics = mujoco.Physics.from_xml_path(str(xml_path)) - task = InsertionTask(random=False) - elif "sim_end_effector_transfer_cube" in task_name: - raise NotImplementedError() - xml_path = ASSETS_DIR / "bimanual_viperx_end_effector_transfer_cube.xml" - physics = mujoco.Physics.from_xml_path(str(xml_path)) - task = TransferCubeEndEffectorTask(random=False) - elif "sim_end_effector_insertion" in task_name: - raise NotImplementedError() - xml_path = ASSETS_DIR / "bimanual_viperx_end_effector_insertion.xml" - physics = mujoco.Physics.from_xml_path(str(xml_path)) - task = InsertionEndEffectorTask(random=False) - else: - raise NotImplementedError(task_name) - - env = control.Environment( - physics, task, time_limit, control_timestep=DT, n_sub_steps=None, flat_observation=False - ) - return env - def _format_raw_obs(self, raw_obs): if self.from_pixels: image = torch.from_numpy(raw_obs["images"]["top"].copy()) @@ -124,6 +88,9 @@ class AlohaEnv(AbstractEnv): return obs def _reset(self, tensordict: Optional[TensorDict] = None): + from aloha.tasks.sim import BOX_POSE + from aloha.utils import sample_box_pose, sample_insertion_pose + if tensordict is not None and not AlohaEnv._reset_warning_issued: logging.warning(f"{self.__class__.__name__}._reset ignores the provided tensordict.") AlohaEnv._reset_warning_issued = True @@ -200,9 +167,14 @@ class AlohaEnv(AbstractEnv): return td def _make_spec(self): - obs = {} from omegaconf import OmegaConf + from aloha.constants import ( + ACTIONS, + JOINTS, + ) + + obs = {} if self.from_pixels: if isinstance(self.image_size, int): image_shape = (3, self.image_size, self.image_size) diff --git a/lerobot/common/envs/aloha/assets/bimanual_viperx_end_effector_insertion.xml b/lerobot/common/envs/aloha/assets/bimanual_viperx_end_effector_insertion.xml deleted file mode 100644 index 8002838c..00000000 --- a/lerobot/common/envs/aloha/assets/bimanual_viperx_end_effector_insertion.xml +++ /dev/null @@ -1,59 +0,0 @@ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - diff --git a/lerobot/common/envs/aloha/assets/bimanual_viperx_end_effector_transfer_cube.xml b/lerobot/common/envs/aloha/assets/bimanual_viperx_end_effector_transfer_cube.xml deleted file mode 100644 index 05249ad2..00000000 --- a/lerobot/common/envs/aloha/assets/bimanual_viperx_end_effector_transfer_cube.xml +++ /dev/null @@ -1,48 +0,0 @@ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - diff --git a/lerobot/common/envs/aloha/assets/bimanual_viperx_insertion.xml b/lerobot/common/envs/aloha/assets/bimanual_viperx_insertion.xml deleted file mode 100644 index 511f7947..00000000 --- a/lerobot/common/envs/aloha/assets/bimanual_viperx_insertion.xml +++ /dev/null @@ -1,53 +0,0 @@ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - diff --git a/lerobot/common/envs/aloha/assets/bimanual_viperx_transfer_cube.xml b/lerobot/common/envs/aloha/assets/bimanual_viperx_transfer_cube.xml deleted file mode 100644 index 2d85a47c..00000000 --- a/lerobot/common/envs/aloha/assets/bimanual_viperx_transfer_cube.xml +++ /dev/null @@ -1,42 +0,0 @@ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - diff --git a/lerobot/common/envs/aloha/assets/scene.xml b/lerobot/common/envs/aloha/assets/scene.xml deleted file mode 100644 index 0f61b8a5..00000000 --- a/lerobot/common/envs/aloha/assets/scene.xml +++ /dev/null @@ -1,38 +0,0 @@ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - diff --git a/lerobot/common/envs/aloha/assets/tabletop.stl b/lerobot/common/envs/aloha/assets/tabletop.stl deleted file mode 100644 index 1c17d3f0..00000000 --- a/lerobot/common/envs/aloha/assets/tabletop.stl +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:76a1571d1aa36520f2bd81c268991b99816c2a7819464d718e0fd9976fe30dce -size 684 diff --git a/lerobot/common/envs/aloha/assets/vx300s_10_custom_finger_left.stl b/lerobot/common/envs/aloha/assets/vx300s_10_custom_finger_left.stl deleted file mode 100644 index ef1f3f35..00000000 --- a/lerobot/common/envs/aloha/assets/vx300s_10_custom_finger_left.stl +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:df73ae5b9058e5d50a6409ac2ab687dade75053a86591bb5e23ab051dbf2d659 -size 83384 diff --git a/lerobot/common/envs/aloha/assets/vx300s_10_custom_finger_right.stl b/lerobot/common/envs/aloha/assets/vx300s_10_custom_finger_right.stl deleted file mode 100644 index 7eb8aefd..00000000 --- a/lerobot/common/envs/aloha/assets/vx300s_10_custom_finger_right.stl +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:56fb3cc1236d4193106038adf8e457c7252ae9e86c7cee6dabf0578c53666358 -size 83384 diff --git a/lerobot/common/envs/aloha/assets/vx300s_10_gripper_finger.stl b/lerobot/common/envs/aloha/assets/vx300s_10_gripper_finger.stl deleted file mode 100644 index 4c2b3a1f..00000000 --- a/lerobot/common/envs/aloha/assets/vx300s_10_gripper_finger.stl +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:a4baacd9a64df1be60ea5e98f50f3c660e1b7a1fe9684aace6004c5058c09483 -size 42884 diff --git a/lerobot/common/envs/aloha/assets/vx300s_11_ar_tag.stl b/lerobot/common/envs/aloha/assets/vx300s_11_ar_tag.stl deleted file mode 100644 index 8a30f7cc..00000000 --- a/lerobot/common/envs/aloha/assets/vx300s_11_ar_tag.stl +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:a18a1601074d29ed1d546ead70cd18fbb063f1db7b5b96b9f0365be714f3136a -size 3884 diff --git a/lerobot/common/envs/aloha/assets/vx300s_1_base.stl b/lerobot/common/envs/aloha/assets/vx300s_1_base.stl deleted file mode 100644 index 9198e625..00000000 --- a/lerobot/common/envs/aloha/assets/vx300s_1_base.stl +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:d100cafe656671ca8fde98fb6a4cf2d1b746995c51c61c25ad9ea2715635d146 -size 99984 diff --git a/lerobot/common/envs/aloha/assets/vx300s_2_shoulder.stl b/lerobot/common/envs/aloha/assets/vx300s_2_shoulder.stl deleted file mode 100644 index ab3d9570..00000000 --- a/lerobot/common/envs/aloha/assets/vx300s_2_shoulder.stl +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:139745a74055cb0b23430bb5bc032bf68cf7bea5e4975c8f4c04107ae005f7f0 -size 63884 diff --git a/lerobot/common/envs/aloha/assets/vx300s_3_upper_arm.stl b/lerobot/common/envs/aloha/assets/vx300s_3_upper_arm.stl deleted file mode 100644 index 3d6f663c..00000000 --- a/lerobot/common/envs/aloha/assets/vx300s_3_upper_arm.stl +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:900f236320dd3d500870c5fde763b2d47502d51e043a5c377875e70237108729 -size 102984 diff --git a/lerobot/common/envs/aloha/assets/vx300s_4_upper_forearm.stl b/lerobot/common/envs/aloha/assets/vx300s_4_upper_forearm.stl deleted file mode 100644 index 4eb249e7..00000000 --- a/lerobot/common/envs/aloha/assets/vx300s_4_upper_forearm.stl +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:4104fc54bbfb8a9b533029f1e7e3ade3d54d638372b3195daa0c98f57e0295b5 -size 49584 diff --git a/lerobot/common/envs/aloha/assets/vx300s_5_lower_forearm.stl b/lerobot/common/envs/aloha/assets/vx300s_5_lower_forearm.stl deleted file mode 100644 index 34c76221..00000000 --- a/lerobot/common/envs/aloha/assets/vx300s_5_lower_forearm.stl +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:66814e27fa728056416e25e02e89eb7d34c51d51c51e7c3df873829037ddc6b8 -size 99884 diff --git a/lerobot/common/envs/aloha/assets/vx300s_6_wrist.stl b/lerobot/common/envs/aloha/assets/vx300s_6_wrist.stl deleted file mode 100644 index 232fabf7..00000000 --- a/lerobot/common/envs/aloha/assets/vx300s_6_wrist.stl +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:90eb145c85627968c3776ae6de23ccff7e112c9dd713c46bc9acdfdaa859a048 -size 70784 diff --git a/lerobot/common/envs/aloha/assets/vx300s_7_gripper.stl b/lerobot/common/envs/aloha/assets/vx300s_7_gripper.stl deleted file mode 100644 index 946c3c86..00000000 --- a/lerobot/common/envs/aloha/assets/vx300s_7_gripper.stl +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:786c1077bfd226f14219581b11d5f19464ca95b17132e0bb7532503568f5af90 -size 450084 diff --git a/lerobot/common/envs/aloha/assets/vx300s_8_gripper_prop.stl b/lerobot/common/envs/aloha/assets/vx300s_8_gripper_prop.stl deleted file mode 100644 index 28d5bd76..00000000 --- a/lerobot/common/envs/aloha/assets/vx300s_8_gripper_prop.stl +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:d1275a93fe2157c83dbc095617fb7e672888bdd48ec070a35ef4ab9ebd9755b0 -size 31684 diff --git a/lerobot/common/envs/aloha/assets/vx300s_9_gripper_bar.stl b/lerobot/common/envs/aloha/assets/vx300s_9_gripper_bar.stl deleted file mode 100644 index 5201d5ea..00000000 --- a/lerobot/common/envs/aloha/assets/vx300s_9_gripper_bar.stl +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:a4de62c9a2ed2c78433010e4c05530a1254b1774a7651967f406120c9bf8973e -size 379484 diff --git a/lerobot/common/envs/aloha/assets/vx300s_dependencies.xml b/lerobot/common/envs/aloha/assets/vx300s_dependencies.xml deleted file mode 100644 index 93037ab7..00000000 --- a/lerobot/common/envs/aloha/assets/vx300s_dependencies.xml +++ /dev/null @@ -1,17 +0,0 @@ - - - - - - - - - - - - - - - - - diff --git a/lerobot/common/envs/aloha/assets/vx300s_left.xml b/lerobot/common/envs/aloha/assets/vx300s_left.xml deleted file mode 100644 index 3af6c235..00000000 --- a/lerobot/common/envs/aloha/assets/vx300s_left.xml +++ /dev/null @@ -1,59 +0,0 @@ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - diff --git a/lerobot/common/envs/aloha/assets/vx300s_right.xml b/lerobot/common/envs/aloha/assets/vx300s_right.xml deleted file mode 100644 index 495df478..00000000 --- a/lerobot/common/envs/aloha/assets/vx300s_right.xml +++ /dev/null @@ -1,59 +0,0 @@ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - diff --git a/lerobot/common/envs/aloha/constants.py b/lerobot/common/envs/aloha/constants.py deleted file mode 100644 index e582e5f3..00000000 --- a/lerobot/common/envs/aloha/constants.py +++ /dev/null @@ -1,163 +0,0 @@ -from pathlib import Path - -### Simulation envs fixed constants -DT = 0.02 # 0.02 ms -> 1/0.2 = 50 hz -FPS = 50 - - -JOINTS = [ - # absolute joint position - "left_arm_waist", - "left_arm_shoulder", - "left_arm_elbow", - "left_arm_forearm_roll", - "left_arm_wrist_angle", - "left_arm_wrist_rotate", - # normalized gripper position 0: close, 1: open - "left_arm_gripper", - # absolute joint position - "right_arm_waist", - "right_arm_shoulder", - "right_arm_elbow", - "right_arm_forearm_roll", - "right_arm_wrist_angle", - "right_arm_wrist_rotate", - # normalized gripper position 0: close, 1: open - "right_arm_gripper", -] - -ACTIONS = [ - # position and quaternion for end effector - "left_arm_waist", - "left_arm_shoulder", - "left_arm_elbow", - "left_arm_forearm_roll", - "left_arm_wrist_angle", - "left_arm_wrist_rotate", - # normalized gripper position (0: close, 1: open) - "left_arm_gripper", - "right_arm_waist", - "right_arm_shoulder", - "right_arm_elbow", - "right_arm_forearm_roll", - "right_arm_wrist_angle", - "right_arm_wrist_rotate", - # normalized gripper position (0: close, 1: open) - "right_arm_gripper", -] - - -START_ARM_POSE = [ - 0, - -0.96, - 1.16, - 0, - -0.3, - 0, - 0.02239, - -0.02239, - 0, - -0.96, - 1.16, - 0, - -0.3, - 0, - 0.02239, - -0.02239, -] - -ASSETS_DIR = Path(__file__).parent.resolve() / "assets" # note: absolute path - -# Left finger position limits (qpos[7]), right_finger = -1 * left_finger -MASTER_GRIPPER_POSITION_OPEN = 0.02417 -MASTER_GRIPPER_POSITION_CLOSE = 0.01244 -PUPPET_GRIPPER_POSITION_OPEN = 0.05800 -PUPPET_GRIPPER_POSITION_CLOSE = 0.01844 - -# Gripper joint limits (qpos[6]) -MASTER_GRIPPER_JOINT_OPEN = 0.3083 -MASTER_GRIPPER_JOINT_CLOSE = -0.6842 -PUPPET_GRIPPER_JOINT_OPEN = 1.4910 -PUPPET_GRIPPER_JOINT_CLOSE = -0.6213 - -MASTER_GRIPPER_JOINT_MID = (MASTER_GRIPPER_JOINT_OPEN + MASTER_GRIPPER_JOINT_CLOSE) / 2 - -############################ Helper functions ############################ - - -def normalize_master_gripper_position(x): - return (x - MASTER_GRIPPER_POSITION_CLOSE) / ( - MASTER_GRIPPER_POSITION_OPEN - MASTER_GRIPPER_POSITION_CLOSE - ) - - -def normalize_puppet_gripper_position(x): - return (x - PUPPET_GRIPPER_POSITION_CLOSE) / ( - PUPPET_GRIPPER_POSITION_OPEN - PUPPET_GRIPPER_POSITION_CLOSE - ) - - -def unnormalize_master_gripper_position(x): - return x * (MASTER_GRIPPER_POSITION_OPEN - MASTER_GRIPPER_POSITION_CLOSE) + MASTER_GRIPPER_POSITION_CLOSE - - -def unnormalize_puppet_gripper_position(x): - return x * (PUPPET_GRIPPER_POSITION_OPEN - PUPPET_GRIPPER_POSITION_CLOSE) + PUPPET_GRIPPER_POSITION_CLOSE - - -def convert_position_from_master_to_puppet(x): - return unnormalize_puppet_gripper_position(normalize_master_gripper_position(x)) - - -def normalizer_master_gripper_joint(x): - return (x - MASTER_GRIPPER_JOINT_CLOSE) / (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE) - - -def normalize_puppet_gripper_joint(x): - return (x - PUPPET_GRIPPER_JOINT_CLOSE) / (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE) - - -def unnormalize_master_gripper_joint(x): - return x * (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE) + MASTER_GRIPPER_JOINT_CLOSE - - -def unnormalize_puppet_gripper_joint(x): - return x * (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE) + PUPPET_GRIPPER_JOINT_CLOSE - - -def convert_join_from_master_to_puppet(x): - return unnormalize_puppet_gripper_joint(normalizer_master_gripper_joint(x)) - - -def normalize_master_gripper_velocity(x): - return x / (MASTER_GRIPPER_POSITION_OPEN - MASTER_GRIPPER_POSITION_CLOSE) - - -def normalize_puppet_gripper_velocity(x): - return x / (PUPPET_GRIPPER_POSITION_OPEN - PUPPET_GRIPPER_POSITION_CLOSE) - - -def convert_master_from_position_to_joint(x): - return ( - normalize_master_gripper_position(x) * (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE) - + MASTER_GRIPPER_JOINT_CLOSE - ) - - -def convert_master_from_joint_to_position(x): - return unnormalize_master_gripper_position( - (x - MASTER_GRIPPER_JOINT_CLOSE) / (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE) - ) - - -def convert_puppet_from_position_to_join(x): - return ( - normalize_puppet_gripper_position(x) * (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE) - + PUPPET_GRIPPER_JOINT_CLOSE - ) - - -def convert_puppet_from_joint_to_position(x): - return unnormalize_puppet_gripper_position( - (x - PUPPET_GRIPPER_JOINT_CLOSE) / (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE) - ) diff --git a/lerobot/common/envs/aloha/tasks/sim.py b/lerobot/common/envs/aloha/tasks/sim.py deleted file mode 100644 index ee1d0927..00000000 --- a/lerobot/common/envs/aloha/tasks/sim.py +++ /dev/null @@ -1,219 +0,0 @@ -import collections - -import numpy as np -from dm_control.suite import base - -from lerobot.common.envs.aloha.constants import ( - START_ARM_POSE, - normalize_puppet_gripper_position, - normalize_puppet_gripper_velocity, - unnormalize_puppet_gripper_position, -) - -BOX_POSE = [None] # to be changed from outside - -""" -Environment for simulated robot bi-manual manipulation, with joint position control -Action space: [left_arm_qpos (6), # absolute joint position - left_gripper_positions (1), # normalized gripper position (0: close, 1: open) - right_arm_qpos (6), # absolute joint position - right_gripper_positions (1),] # normalized gripper position (0: close, 1: open) - -Observation space: {"qpos": Concat[ left_arm_qpos (6), # absolute joint position - left_gripper_position (1), # normalized gripper position (0: close, 1: open) - right_arm_qpos (6), # absolute joint position - right_gripper_qpos (1)] # normalized gripper position (0: close, 1: open) - "qvel": Concat[ left_arm_qvel (6), # absolute joint velocity (rad) - left_gripper_velocity (1), # normalized gripper velocity (pos: opening, neg: closing) - right_arm_qvel (6), # absolute joint velocity (rad) - right_gripper_qvel (1)] # normalized gripper velocity (pos: opening, neg: closing) - "images": {"main": (480x640x3)} # h, w, c, dtype='uint8' -""" - - -class BimanualViperXTask(base.Task): - def __init__(self, random=None): - super().__init__(random=random) - - def before_step(self, action, physics): - left_arm_action = action[:6] - right_arm_action = action[7 : 7 + 6] - normalized_left_gripper_action = action[6] - normalized_right_gripper_action = action[7 + 6] - - left_gripper_action = unnormalize_puppet_gripper_position(normalized_left_gripper_action) - right_gripper_action = unnormalize_puppet_gripper_position(normalized_right_gripper_action) - - full_left_gripper_action = [left_gripper_action, -left_gripper_action] - full_right_gripper_action = [right_gripper_action, -right_gripper_action] - - env_action = np.concatenate( - [left_arm_action, full_left_gripper_action, right_arm_action, full_right_gripper_action] - ) - super().before_step(env_action, physics) - return - - def initialize_episode(self, physics): - """Sets the state of the environment at the start of each episode.""" - super().initialize_episode(physics) - - @staticmethod - def get_qpos(physics): - qpos_raw = physics.data.qpos.copy() - left_qpos_raw = qpos_raw[:8] - right_qpos_raw = qpos_raw[8:16] - left_arm_qpos = left_qpos_raw[:6] - right_arm_qpos = right_qpos_raw[:6] - left_gripper_qpos = [normalize_puppet_gripper_position(left_qpos_raw[6])] - right_gripper_qpos = [normalize_puppet_gripper_position(right_qpos_raw[6])] - return np.concatenate([left_arm_qpos, left_gripper_qpos, right_arm_qpos, right_gripper_qpos]) - - @staticmethod - def get_qvel(physics): - qvel_raw = physics.data.qvel.copy() - left_qvel_raw = qvel_raw[:8] - right_qvel_raw = qvel_raw[8:16] - left_arm_qvel = left_qvel_raw[:6] - right_arm_qvel = right_qvel_raw[:6] - left_gripper_qvel = [normalize_puppet_gripper_velocity(left_qvel_raw[6])] - right_gripper_qvel = [normalize_puppet_gripper_velocity(right_qvel_raw[6])] - return np.concatenate([left_arm_qvel, left_gripper_qvel, right_arm_qvel, right_gripper_qvel]) - - @staticmethod - def get_env_state(physics): - raise NotImplementedError - - def get_observation(self, physics): - obs = collections.OrderedDict() - obs["qpos"] = self.get_qpos(physics) - obs["qvel"] = self.get_qvel(physics) - obs["env_state"] = self.get_env_state(physics) - obs["images"] = {} - obs["images"]["top"] = physics.render(height=480, width=640, camera_id="top") - obs["images"]["angle"] = physics.render(height=480, width=640, camera_id="angle") - obs["images"]["vis"] = physics.render(height=480, width=640, camera_id="front_close") - - return obs - - def get_reward(self, physics): - # return whether left gripper is holding the box - raise NotImplementedError - - -class TransferCubeTask(BimanualViperXTask): - def __init__(self, random=None): - super().__init__(random=random) - self.max_reward = 4 - - def initialize_episode(self, physics): - """Sets the state of the environment at the start of each episode.""" - # TODO Notice: this function does not randomize the env configuration. Instead, set BOX_POSE from outside - # reset qpos, control and box position - with physics.reset_context(): - physics.named.data.qpos[:16] = START_ARM_POSE - np.copyto(physics.data.ctrl, START_ARM_POSE) - assert BOX_POSE[0] is not None - physics.named.data.qpos[-7:] = BOX_POSE[0] - # print(f"{BOX_POSE=}") - super().initialize_episode(physics) - - @staticmethod - def get_env_state(physics): - env_state = physics.data.qpos.copy()[16:] - return env_state - - def get_reward(self, physics): - # return whether left gripper is holding the box - all_contact_pairs = [] - for i_contact in range(physics.data.ncon): - id_geom_1 = physics.data.contact[i_contact].geom1 - id_geom_2 = physics.data.contact[i_contact].geom2 - name_geom_1 = physics.model.id2name(id_geom_1, "geom") - name_geom_2 = physics.model.id2name(id_geom_2, "geom") - contact_pair = (name_geom_1, name_geom_2) - all_contact_pairs.append(contact_pair) - - touch_left_gripper = ("red_box", "vx300s_left/10_left_gripper_finger") in all_contact_pairs - touch_right_gripper = ("red_box", "vx300s_right/10_right_gripper_finger") in all_contact_pairs - touch_table = ("red_box", "table") in all_contact_pairs - - reward = 0 - if touch_right_gripper: - reward = 1 - if touch_right_gripper and not touch_table: # lifted - reward = 2 - if touch_left_gripper: # attempted transfer - reward = 3 - if touch_left_gripper and not touch_table: # successful transfer - reward = 4 - return reward - - -class InsertionTask(BimanualViperXTask): - def __init__(self, random=None): - super().__init__(random=random) - self.max_reward = 4 - - def initialize_episode(self, physics): - """Sets the state of the environment at the start of each episode.""" - # TODO Notice: this function does not randomize the env configuration. Instead, set BOX_POSE from outside - # reset qpos, control and box position - with physics.reset_context(): - physics.named.data.qpos[:16] = START_ARM_POSE - np.copyto(physics.data.ctrl, START_ARM_POSE) - assert BOX_POSE[0] is not None - physics.named.data.qpos[-7 * 2 :] = BOX_POSE[0] # two objects - # print(f"{BOX_POSE=}") - super().initialize_episode(physics) - - @staticmethod - def get_env_state(physics): - env_state = physics.data.qpos.copy()[16:] - return env_state - - def get_reward(self, physics): - # return whether peg touches the pin - all_contact_pairs = [] - for i_contact in range(physics.data.ncon): - id_geom_1 = physics.data.contact[i_contact].geom1 - id_geom_2 = physics.data.contact[i_contact].geom2 - name_geom_1 = physics.model.id2name(id_geom_1, "geom") - name_geom_2 = physics.model.id2name(id_geom_2, "geom") - contact_pair = (name_geom_1, name_geom_2) - all_contact_pairs.append(contact_pair) - - touch_right_gripper = ("red_peg", "vx300s_right/10_right_gripper_finger") in all_contact_pairs - touch_left_gripper = ( - ("socket-1", "vx300s_left/10_left_gripper_finger") in all_contact_pairs - or ("socket-2", "vx300s_left/10_left_gripper_finger") in all_contact_pairs - or ("socket-3", "vx300s_left/10_left_gripper_finger") in all_contact_pairs - or ("socket-4", "vx300s_left/10_left_gripper_finger") in all_contact_pairs - ) - - peg_touch_table = ("red_peg", "table") in all_contact_pairs - socket_touch_table = ( - ("socket-1", "table") in all_contact_pairs - or ("socket-2", "table") in all_contact_pairs - or ("socket-3", "table") in all_contact_pairs - or ("socket-4", "table") in all_contact_pairs - ) - peg_touch_socket = ( - ("red_peg", "socket-1") in all_contact_pairs - or ("red_peg", "socket-2") in all_contact_pairs - or ("red_peg", "socket-3") in all_contact_pairs - or ("red_peg", "socket-4") in all_contact_pairs - ) - pin_touched = ("red_peg", "pin") in all_contact_pairs - - reward = 0 - if touch_left_gripper and touch_right_gripper: # touch both - reward = 1 - if ( - touch_left_gripper and touch_right_gripper and (not peg_touch_table) and (not socket_touch_table) - ): # grasp both - reward = 2 - if peg_touch_socket and (not peg_touch_table) and (not socket_touch_table): # peg and socket touching - reward = 3 - if pin_touched: # successful insertion - reward = 4 - return reward diff --git a/lerobot/common/envs/aloha/tasks/sim_end_effector.py b/lerobot/common/envs/aloha/tasks/sim_end_effector.py deleted file mode 100644 index d93c8330..00000000 --- a/lerobot/common/envs/aloha/tasks/sim_end_effector.py +++ /dev/null @@ -1,263 +0,0 @@ -import collections - -import numpy as np -from dm_control.suite import base - -from lerobot.common.envs.aloha.constants import ( - PUPPET_GRIPPER_POSITION_CLOSE, - START_ARM_POSE, - normalize_puppet_gripper_position, - normalize_puppet_gripper_velocity, - unnormalize_puppet_gripper_position, -) -from lerobot.common.envs.aloha.utils import sample_box_pose, sample_insertion_pose - -""" -Environment for simulated robot bi-manual manipulation, with end-effector control. -Action space: [left_arm_pose (7), # position and quaternion for end effector - left_gripper_positions (1), # normalized gripper position (0: close, 1: open) - right_arm_pose (7), # position and quaternion for end effector - right_gripper_positions (1),] # normalized gripper position (0: close, 1: open) - -Observation space: {"qpos": Concat[ left_arm_qpos (6), # absolute joint position - left_gripper_position (1), # normalized gripper position (0: close, 1: open) - right_arm_qpos (6), # absolute joint position - right_gripper_qpos (1)] # normalized gripper position (0: close, 1: open) - "qvel": Concat[ left_arm_qvel (6), # absolute joint velocity (rad) - left_gripper_velocity (1), # normalized gripper velocity (pos: opening, neg: closing) - right_arm_qvel (6), # absolute joint velocity (rad) - right_gripper_qvel (1)] # normalized gripper velocity (pos: opening, neg: closing) - "images": {"main": (480x640x3)} # h, w, c, dtype='uint8' -""" - - -class BimanualViperXEndEffectorTask(base.Task): - def __init__(self, random=None): - super().__init__(random=random) - - def before_step(self, action, physics): - a_len = len(action) // 2 - action_left = action[:a_len] - action_right = action[a_len:] - - # set mocap position and quat - # left - np.copyto(physics.data.mocap_pos[0], action_left[:3]) - np.copyto(physics.data.mocap_quat[0], action_left[3:7]) - # right - np.copyto(physics.data.mocap_pos[1], action_right[:3]) - np.copyto(physics.data.mocap_quat[1], action_right[3:7]) - - # set gripper - g_left_ctrl = unnormalize_puppet_gripper_position(action_left[7]) - g_right_ctrl = unnormalize_puppet_gripper_position(action_right[7]) - np.copyto(physics.data.ctrl, np.array([g_left_ctrl, -g_left_ctrl, g_right_ctrl, -g_right_ctrl])) - - def initialize_robots(self, physics): - # reset joint position - physics.named.data.qpos[:16] = START_ARM_POSE - - # reset mocap to align with end effector - # to obtain these numbers: - # (1) make an ee_sim env and reset to the same start_pose - # (2) get env._physics.named.data.xpos['vx300s_left/gripper_link'] - # get env._physics.named.data.xquat['vx300s_left/gripper_link'] - # repeat the same for right side - np.copyto(physics.data.mocap_pos[0], [-0.31718881, 0.5, 0.29525084]) - np.copyto(physics.data.mocap_quat[0], [1, 0, 0, 0]) - # right - np.copyto(physics.data.mocap_pos[1], np.array([0.31718881, 0.49999888, 0.29525084])) - np.copyto(physics.data.mocap_quat[1], [1, 0, 0, 0]) - - # reset gripper control - close_gripper_control = np.array( - [ - PUPPET_GRIPPER_POSITION_CLOSE, - -PUPPET_GRIPPER_POSITION_CLOSE, - PUPPET_GRIPPER_POSITION_CLOSE, - -PUPPET_GRIPPER_POSITION_CLOSE, - ] - ) - np.copyto(physics.data.ctrl, close_gripper_control) - - def initialize_episode(self, physics): - """Sets the state of the environment at the start of each episode.""" - super().initialize_episode(physics) - - @staticmethod - def get_qpos(physics): - qpos_raw = physics.data.qpos.copy() - left_qpos_raw = qpos_raw[:8] - right_qpos_raw = qpos_raw[8:16] - left_arm_qpos = left_qpos_raw[:6] - right_arm_qpos = right_qpos_raw[:6] - left_gripper_qpos = [normalize_puppet_gripper_position(left_qpos_raw[6])] - right_gripper_qpos = [normalize_puppet_gripper_position(right_qpos_raw[6])] - return np.concatenate([left_arm_qpos, left_gripper_qpos, right_arm_qpos, right_gripper_qpos]) - - @staticmethod - def get_qvel(physics): - qvel_raw = physics.data.qvel.copy() - left_qvel_raw = qvel_raw[:8] - right_qvel_raw = qvel_raw[8:16] - left_arm_qvel = left_qvel_raw[:6] - right_arm_qvel = right_qvel_raw[:6] - left_gripper_qvel = [normalize_puppet_gripper_velocity(left_qvel_raw[6])] - right_gripper_qvel = [normalize_puppet_gripper_velocity(right_qvel_raw[6])] - return np.concatenate([left_arm_qvel, left_gripper_qvel, right_arm_qvel, right_gripper_qvel]) - - @staticmethod - def get_env_state(physics): - raise NotImplementedError - - def get_observation(self, physics): - # note: it is important to do .copy() - obs = collections.OrderedDict() - obs["qpos"] = self.get_qpos(physics) - obs["qvel"] = self.get_qvel(physics) - obs["env_state"] = self.get_env_state(physics) - obs["images"] = {} - obs["images"]["top"] = physics.render(height=480, width=640, camera_id="top") - obs["images"]["angle"] = physics.render(height=480, width=640, camera_id="angle") - obs["images"]["vis"] = physics.render(height=480, width=640, camera_id="front_close") - # used in scripted policy to obtain starting pose - obs["mocap_pose_left"] = np.concatenate( - [physics.data.mocap_pos[0], physics.data.mocap_quat[0]] - ).copy() - obs["mocap_pose_right"] = np.concatenate( - [physics.data.mocap_pos[1], physics.data.mocap_quat[1]] - ).copy() - - # used when replaying joint trajectory - obs["gripper_ctrl"] = physics.data.ctrl.copy() - return obs - - def get_reward(self, physics): - raise NotImplementedError - - -class TransferCubeEndEffectorTask(BimanualViperXEndEffectorTask): - def __init__(self, random=None): - super().__init__(random=random) - self.max_reward = 4 - - def initialize_episode(self, physics): - """Sets the state of the environment at the start of each episode.""" - self.initialize_robots(physics) - # randomize box position - cube_pose = sample_box_pose() - box_start_idx = physics.model.name2id("red_box_joint", "joint") - np.copyto(physics.data.qpos[box_start_idx : box_start_idx + 7], cube_pose) - # print(f"randomized cube position to {cube_position}") - - super().initialize_episode(physics) - - @staticmethod - def get_env_state(physics): - env_state = physics.data.qpos.copy()[16:] - return env_state - - def get_reward(self, physics): - # return whether left gripper is holding the box - all_contact_pairs = [] - for i_contact in range(physics.data.ncon): - id_geom_1 = physics.data.contact[i_contact].geom1 - id_geom_2 = physics.data.contact[i_contact].geom2 - name_geom_1 = physics.model.id2name(id_geom_1, "geom") - name_geom_2 = physics.model.id2name(id_geom_2, "geom") - contact_pair = (name_geom_1, name_geom_2) - all_contact_pairs.append(contact_pair) - - touch_left_gripper = ("red_box", "vx300s_left/10_left_gripper_finger") in all_contact_pairs - touch_right_gripper = ("red_box", "vx300s_right/10_right_gripper_finger") in all_contact_pairs - touch_table = ("red_box", "table") in all_contact_pairs - - reward = 0 - if touch_right_gripper: - reward = 1 - if touch_right_gripper and not touch_table: # lifted - reward = 2 - if touch_left_gripper: # attempted transfer - reward = 3 - if touch_left_gripper and not touch_table: # successful transfer - reward = 4 - return reward - - -class InsertionEndEffectorTask(BimanualViperXEndEffectorTask): - def __init__(self, random=None): - super().__init__(random=random) - self.max_reward = 4 - - def initialize_episode(self, physics): - """Sets the state of the environment at the start of each episode.""" - self.initialize_robots(physics) - # randomize peg and socket position - peg_pose, socket_pose = sample_insertion_pose() - - def id2index(j_id): - return 16 + (j_id - 16) * 7 # first 16 is robot qpos, 7 is pose dim # hacky - - peg_start_id = physics.model.name2id("red_peg_joint", "joint") - peg_start_idx = id2index(peg_start_id) - np.copyto(physics.data.qpos[peg_start_idx : peg_start_idx + 7], peg_pose) - # print(f"randomized cube position to {cube_position}") - - socket_start_id = physics.model.name2id("blue_socket_joint", "joint") - socket_start_idx = id2index(socket_start_id) - np.copyto(physics.data.qpos[socket_start_idx : socket_start_idx + 7], socket_pose) - # print(f"randomized cube position to {cube_position}") - - super().initialize_episode(physics) - - @staticmethod - def get_env_state(physics): - env_state = physics.data.qpos.copy()[16:] - return env_state - - def get_reward(self, physics): - # return whether peg touches the pin - all_contact_pairs = [] - for i_contact in range(physics.data.ncon): - id_geom_1 = physics.data.contact[i_contact].geom1 - id_geom_2 = physics.data.contact[i_contact].geom2 - name_geom_1 = physics.model.id2name(id_geom_1, "geom") - name_geom_2 = physics.model.id2name(id_geom_2, "geom") - contact_pair = (name_geom_1, name_geom_2) - all_contact_pairs.append(contact_pair) - - touch_right_gripper = ("red_peg", "vx300s_right/10_right_gripper_finger") in all_contact_pairs - touch_left_gripper = ( - ("socket-1", "vx300s_left/10_left_gripper_finger") in all_contact_pairs - or ("socket-2", "vx300s_left/10_left_gripper_finger") in all_contact_pairs - or ("socket-3", "vx300s_left/10_left_gripper_finger") in all_contact_pairs - or ("socket-4", "vx300s_left/10_left_gripper_finger") in all_contact_pairs - ) - - peg_touch_table = ("red_peg", "table") in all_contact_pairs - socket_touch_table = ( - ("socket-1", "table") in all_contact_pairs - or ("socket-2", "table") in all_contact_pairs - or ("socket-3", "table") in all_contact_pairs - or ("socket-4", "table") in all_contact_pairs - ) - peg_touch_socket = ( - ("red_peg", "socket-1") in all_contact_pairs - or ("red_peg", "socket-2") in all_contact_pairs - or ("red_peg", "socket-3") in all_contact_pairs - or ("red_peg", "socket-4") in all_contact_pairs - ) - pin_touched = ("red_peg", "pin") in all_contact_pairs - - reward = 0 - if touch_left_gripper and touch_right_gripper: # touch both - reward = 1 - if ( - touch_left_gripper and touch_right_gripper and (not peg_touch_table) and (not socket_touch_table) - ): # grasp both - reward = 2 - if peg_touch_socket and (not peg_touch_table) and (not socket_touch_table): # peg and socket touching - reward = 3 - if pin_touched: # successful insertion - reward = 4 - return reward diff --git a/lerobot/common/envs/aloha/utils.py b/lerobot/common/envs/aloha/utils.py deleted file mode 100644 index 5ac8b955..00000000 --- a/lerobot/common/envs/aloha/utils.py +++ /dev/null @@ -1,39 +0,0 @@ -import numpy as np - - -def sample_box_pose(): - x_range = [0.0, 0.2] - y_range = [0.4, 0.6] - z_range = [0.05, 0.05] - - ranges = np.vstack([x_range, y_range, z_range]) - cube_position = np.random.uniform(ranges[:, 0], ranges[:, 1]) - - cube_quat = np.array([1, 0, 0, 0]) - return np.concatenate([cube_position, cube_quat]) - - -def sample_insertion_pose(): - # Peg - x_range = [0.1, 0.2] - y_range = [0.4, 0.6] - z_range = [0.05, 0.05] - - ranges = np.vstack([x_range, y_range, z_range]) - peg_position = np.random.uniform(ranges[:, 0], ranges[:, 1]) - - peg_quat = np.array([1, 0, 0, 0]) - peg_pose = np.concatenate([peg_position, peg_quat]) - - # Socket - x_range = [-0.2, -0.1] - y_range = [0.4, 0.6] - z_range = [0.05, 0.05] - - ranges = np.vstack([x_range, y_range, z_range]) - socket_position = np.random.uniform(ranges[:, 0], ranges[:, 1]) - - socket_quat = np.array([1, 0, 0, 0]) - socket_pose = np.concatenate([socket_position, socket_quat]) - - return peg_pose, socket_pose diff --git a/lerobot/common/envs/factory.py b/lerobot/common/envs/factory.py index 855e073b..0b3d0ae5 100644 --- a/lerobot/common/envs/factory.py +++ b/lerobot/common/envs/factory.py @@ -17,18 +17,18 @@ def make_env(cfg, transform=None): } if cfg.env.name == "simxarm": - from lerobot.common.envs.simxarm.env import SimxarmEnv + from lerobot.common.envs.xarm import SimxarmEnv kwargs["task"] = cfg.env.task clsfunc = SimxarmEnv elif cfg.env.name == "pusht": - from lerobot.common.envs.pusht.env import PushtEnv + from lerobot.common.envs.pusht import PushtEnv # assert kwargs["seed"] > 200, "Seed 0-200 are used for the demonstration dataset, so we don't want to seed the eval env with this range." clsfunc = PushtEnv elif cfg.env.name == "aloha": - from lerobot.common.envs.aloha.env import AlohaEnv + from lerobot.common.envs.aloha import AlohaEnv kwargs["task"] = cfg.env.task clsfunc = AlohaEnv diff --git a/lerobot/common/envs/pusht/env.py b/lerobot/common/envs/pusht.py similarity index 100% rename from lerobot/common/envs/pusht/env.py rename to lerobot/common/envs/pusht.py diff --git a/lerobot/common/envs/simxarm/env.py b/lerobot/common/envs/xarm.py similarity index 100% rename from lerobot/common/envs/simxarm/env.py rename to lerobot/common/envs/xarm.py diff --git a/poetry.lock b/poetry.lock index a1a59b69..ab3a0867 100644 --- a/poetry.lock +++ b/poetry.lock @@ -3096,6 +3096,20 @@ numpy = ">=1.14,<2" docs = ["matplotlib", "numpydoc (==1.1.*)", "sphinx", "sphinx-book-theme", "sphinx-remove-toctrees"] test = ["pytest", "pytest-cov"] +[[package]] +name = "sim-aloha" +version = "0.1.2" +description = "ALOHA environment for LeRobot" +optional = true +python-versions = "<4.0,>=3.10" +files = [ + {file = "sim_aloha-0.1.2-py3-none-any.whl", hash = "sha256:3b13c1ee474481f5d686b57bf1a9ed350e01ca4da7d65f7501446eb74b02653a"}, + {file = "sim_aloha-0.1.2.tar.gz", hash = "sha256:44d36bbdb13e98e0c74f4d2a682f38683f4f63951618da28175f89cbb1c6f324"}, +] + +[package.dependencies] +dm-control = "1.0.14" + [[package]] name = "sim-pusht" version = "0.1.0" @@ -3622,10 +3636,11 @@ docs = ["furo", "jaraco.packaging (>=9.3)", "jaraco.tidelift (>=1.4)", "rst.link testing = ["big-O", "jaraco.functools", "jaraco.itertools", "more-itertools", "pytest (>=6)", "pytest-checkdocs (>=2.4)", "pytest-cov", "pytest-enabler (>=2.2)", "pytest-ignore-flaky", "pytest-mypy", "pytest-ruff (>=0.2.1)"] [extras] +aloha = ["sim-aloha"] pusht = ["sim-pusht"] xarm = ["sim-xarm"] [metadata] lock-version = "2.0" python-versions = "^3.10" -content-hash = "0ac21898262a0d5410340f822b4aa3ae4580348ebb2bc58dbd351f46a0feedbd" +content-hash = "e42b69516ddf3626847e98266fb844dcbcf2cdfc3d48524326e212bc494eafc7" diff --git a/pyproject.toml b/pyproject.toml index bcf27b05..dd24f0ec 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -59,11 +59,13 @@ gymnasium = "^0.29.1" cmake = "^3.29.0.1" sim-pusht = { version = "^0.1.0", optional = true} sim-xarm = { version = "^0.1.0", optional = true} +sim-aloha = { version = "^0.1.2", optional = true} [tool.poetry.extras] pusht = ["sim-pusht"] xarm = ["sim-xarm"] +aloha = ["sim-aloha"] [tool.poetry.group.dev.dependencies] @@ -109,7 +111,7 @@ select = ["E4", "E7", "E9", "F", "I", "N", "B", "C4", "SIM"] ignore-init-module-imports = true [tool.ruff.lint.isort] -known-first-party = ["pusht", "xarm"] +known-first-party = ["pusht", "xarm", "aloha"] [tool.poetry-dynamic-versioning] enable = true diff --git a/tests/test_available.py b/tests/test_available.py index 9cc91efa..827f07ff 100644 --- a/tests/test_available.py +++ b/tests/test_available.py @@ -15,9 +15,9 @@ Note: import pytest import lerobot -from lerobot.common.envs.aloha.env import AlohaEnv -from lerobot.common.envs.pusht.env import PushtEnv -from lerobot.common.envs.simxarm.env import SimxarmEnv +from lerobot.common.envs.aloha import AlohaEnv +from lerobot.common.envs.pusht import PushtEnv +from lerobot.common.envs.xarm import SimxarmEnv from lerobot.common.datasets.simxarm import SimxarmDataset from lerobot.common.datasets.aloha import AlohaDataset diff --git a/tests/test_envs.py b/tests/test_envs.py index a91bb707..7c32b30f 100644 --- a/tests/test_envs.py +++ b/tests/test_envs.py @@ -5,9 +5,9 @@ from torchrl.envs.utils import check_env_specs, step_mdp from lerobot.common.datasets.factory import make_offline_buffer from lerobot.common.envs.factory import make_env -from lerobot.common.envs.pusht.env import PushtEnv -from lerobot.common.envs.simxarm.env import SimxarmEnv -from lerobot.common.envs.aloha.env import AlohaEnv +from lerobot.common.envs.pusht import PushtEnv +from lerobot.common.envs.xarm import SimxarmEnv +from lerobot.common.envs.aloha import AlohaEnv from .utils import DEVICE, init_config