diff --git a/src/lerobot/envs/configs.py b/src/lerobot/envs/configs.py index dc526114d..029da95ac 100644 --- a/src/lerobot/envs/configs.py +++ b/src/lerobot/envs/configs.py @@ -261,13 +261,44 @@ class LiberoEnv(EnvConfig): type=FeatureType.VISUAL, shape=(self.observation_height, self.observation_width, 3) ) elif self.obs_type == "pixels_agent_pos": - self.features["agent_pos"] = PolicyFeature(type=FeatureType.STATE, shape=(8,)) self.features["pixels/agentview_image"] = PolicyFeature( type=FeatureType.VISUAL, shape=(self.observation_height, self.observation_width, 3) ) self.features["pixels/robot0_eye_in_hand_image"] = PolicyFeature( type=FeatureType.VISUAL, shape=(self.observation_height, self.observation_width, 3) ) + self.features["robot_state/eef/pos"] = PolicyFeature( + type=FeatureType.STATE, + shape=(3,), + ) + self.features["robot_state/eef/quat"] = PolicyFeature( + type=FeatureType.STATE, + shape=(4,), + ) + self.features["robot_state/eef/mat"] = PolicyFeature( + type=FeatureType.STATE, + shape=(3, 3), + ) + self.features["robot_state/eef/axisangle"] = PolicyFeature( + type=FeatureType.STATE, + shape=(3,), + ) + self.features["robot_state/gripper/qpos"] = PolicyFeature( + type=FeatureType.STATE, + shape=(2,), + ) + self.features["robot_state/gripper/qvel"] = PolicyFeature( + type=FeatureType.STATE, + shape=(2,), + ) + self.features["robot_state/joints/pos"] = PolicyFeature( + type=FeatureType.STATE, + shape=(7,), + ) + self.features["robot_state/joints/vel"] = PolicyFeature( + type=FeatureType.STATE, + shape=(7,), + ) else: raise ValueError(f"Unsupported obs_type: {self.obs_type}") diff --git a/src/lerobot/envs/libero.py b/src/lerobot/envs/libero.py index 94b08e991..5e692501a 100644 --- a/src/lerobot/envs/libero.py +++ b/src/lerobot/envs/libero.py @@ -212,23 +212,44 @@ class LiberoEnv(gym.Env): images = {} for camera_name in self.camera_name: image = raw_obs[camera_name] - image = image[::-1, ::-1] # rotate 180 degrees images[self.camera_name_mapping[camera_name]] = image - state = np.concatenate( - ( - raw_obs["robot0_eef_pos"], - quat2axisangle(raw_obs["robot0_eef_quat"]), - raw_obs["robot0_gripper_qpos"], - ) - ) - agent_pos = state + + eef_pos = raw_obs.get("robot0_eef_pos") + eef_quat = raw_obs.get("robot0_eef_quat") + + # rotation matrix from controller + eef_mat = self._env.robots[0].controller.ee_ori_mat if eef_pos is not None else None + eef_axisangle = quat2axisangle(eef_quat) if eef_quat is not None else None + gripper_qpos = raw_obs.get("robot0_gripper_qpos") + gripper_qvel = raw_obs.get("robot0_gripper_qvel") + joint_pos = raw_obs.get("robot0_joint_pos") + joint_vel = raw_obs.get("robot0_joint_vel") + obs = { + "pixels": images, + + "robot_state": { + "eef": { + "pos": eef_pos, # (3,) + "quat": eef_quat, # (4,) + "mat": eef_mat, # (3, 3) + "axisangle": eef_axisangle, # (3) + }, + "gripper": { + "qpos": gripper_qpos, # (2,) + "qvel": gripper_qvel, # (2,) + }, + "joints": { + "pos": joint_pos, # (7,) + "vel": joint_vel, # (7,) + }, + }, + } if self.obs_type == "pixels": return {"pixels": images.copy()} + if self.obs_type == "pixels_agent_pos": - return { - "pixels": images.copy(), - "agent_pos": agent_pos, - } + return obs + raise NotImplementedError( f"The observation type '{self.obs_type}' is not supported in LiberoEnv. " "Please switch to an image-based obs_type (e.g. 'pixels', 'pixels_agent_pos')." @@ -355,12 +376,11 @@ def create_libero_envs( print(f"Restricting to task_ids={task_ids_filter}") out: dict[str, dict[int, Any]] = defaultdict(dict) - for suite_name in suite_names: suite = _get_suite(suite_name) total = len(suite.tasks) selected = _select_task_ids(total, task_ids_filter) - + selected = [0] if not selected: raise ValueError(f"No tasks selected for suite '{suite_name}' (available: {total}).") diff --git a/test2.py b/test2.py new file mode 100644 index 000000000..9b2cf68a3 --- /dev/null +++ b/test2.py @@ -0,0 +1,9 @@ +from lerobot.envs.factory import make_env, make_env_config +config = make_env_config("libero", task="libero_spatial") +envs_dict = make_env(config) +env = envs_dict["libero_spatial"][0] + +seed = 42 + +# First rollout +obs1, info1 = env.reset(seed=seed) diff --git a/test_libero.py b/test_libero.py new file mode 100644 index 000000000..75155b696 --- /dev/null +++ b/test_libero.py @@ -0,0 +1,148 @@ +#!/usr/bin/env python + +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os + +import numpy as np +import pytest + +from lerobot.envs.factory import make_env, make_env_config + +# Set MuJoCo rendering backend before importing environment +os.environ["MUJOCO_GL"] = "egl" + + +def assert_observations_equal(obs1, obs2, path="", atol=1e-8): + """ + Recursively compare two observations and assert they are equal. + + Args: + obs1: First observation (dict or numpy array) + obs2: Second observation (dict or numpy array) + path: Current path in nested structure (for error messages) + atol: Absolute tolerance for numpy array comparisons + """ + if isinstance(obs1, dict) and isinstance(obs2, dict): + assert obs1.keys() == obs2.keys(), f"Keys differ at {path}: {obs1.keys()} != {obs2.keys()}" + for key in obs1: + assert_observations_equal(obs1[key], obs2[key], path=f"{path}.{key}" if path else key, atol=atol) + elif isinstance(obs1, np.ndarray) and isinstance(obs2, np.ndarray): + assert obs1.shape == obs2.shape, f"Shape mismatch at {path}: {obs1.shape} != {obs2.shape}" + assert obs1.dtype == obs2.dtype, f"Dtype mismatch at {path}: {obs1.dtype} != {obs2.dtype}" + assert np.allclose(obs1, obs2, atol=atol), ( + f"Array values differ at {path}: max abs diff = {np.abs(obs1 - obs2).max()}" + ) + else: + assert type(obs1) is type(obs2), f"Type mismatch at {path}: {type(obs1)} != {type(obs2)}" + assert obs1 == obs2, f"Values differ at {path}: {obs1} != {obs2}" + + +def test_libero_env_creation(): + """Test that the libero environment can be created successfully.""" + config = make_env_config("libero", task="libero_spatial") + envs_dict = make_env(config) + + assert "libero_spatial" in envs_dict + assert 0 in envs_dict["libero_spatial"] + + env = envs_dict["libero_spatial"][0] + assert env is not None + + # Test basic reset + observation, info = env.reset(seed=42) + assert observation is not None + assert info is not None + + env.close() + + +def test_libero_reset_determinism(): + """Test that resetting with the same seed produces identical observations.""" + config = make_env_config("libero", task="libero_spatial") + envs_dict = make_env(config) + env = envs_dict["libero_spatial"][0] + + # Reset multiple times with the same seed + obs1, info1 = env.reset(seed=42) + obs2, info2 = env.reset(seed=42) + obs3, info3 = env.reset(seed=42) + + # All observations should be identical + assert_observations_equal(obs1, obs2) + assert_observations_equal(obs1, obs3) + assert_observations_equal(obs2, obs3) + + env.close() + + +def test_libero_step_determinism(): + """Test that step() is deterministic when resetting with the same seed.""" + config = make_env_config("libero", task="libero_spatial") + envs_dict = make_env(config) + env = envs_dict["libero_spatial"][0] + + seed = 42 + + # First rollout + obs1, info1 = env.reset(seed=seed) + action = env.action_space.sample() + obs_after_step1, reward1, terminated1, truncated1, info_step1 = env.step(action) + + # Second rollout with identical seed and action + obs2, info2 = env.reset(seed=seed) + obs_after_step2, reward2, terminated2, truncated2, info_step2 = env.step(action) + + # Initial observations should be identical + assert_observations_equal(obs1, obs2) + + # Post-step observations should be identical + assert_observations_equal(obs_after_step1, obs_after_step2) + + # Rewards and termination flags should be identical + assert np.allclose(reward1, reward2), f"Rewards differ: {reward1} != {reward2}" + assert np.array_equal(terminated1, terminated2), ( + f"Terminated flags differ: {terminated1} != {terminated2}" + ) + assert np.array_equal(truncated1, truncated2), f"Truncated flags differ: {truncated1} != {truncated2}" + + env.close() + + +@pytest.mark.parametrize("task", ["libero_spatial", "libero_object", "libero_goal", "libero_10"]) +def test_libero_tasks(task): + """Test that different libero tasks can be created and used.""" + config = make_env_config("libero", task=task) + envs_dict = make_env(config) + + assert task in envs_dict + assert 0 in envs_dict[task] + + env = envs_dict[task][0] + observation, info = env.reset(seed=42) + + assert observation is not None + assert info is not None + + # Take a step + action = env.action_space.sample() + obs, reward, terminated, truncated, info = env.step(action) + + assert obs is not None + assert reward is not None + assert isinstance(terminated, (bool, np.ndarray)) + assert isinstance(truncated, (bool, np.ndarray)) + + env.close() \ No newline at end of file