diff --git a/eval.sh b/eval.sh index 151a292d0..7cc7551e4 100644 --- a/eval.sh +++ b/eval.sh @@ -3,4 +3,4 @@ lerobot-eval \ --env.type=libero \ --env.task=libero_spatial \ --eval.batch_size=1 \ - --eval.n_episodes=1 \ No newline at end of file + --eval.n_episodes=1 diff --git a/src/lerobot/envs/libero.py b/src/lerobot/envs/libero.py index a44f18881..c67dbcce0 100644 --- a/src/lerobot/envs/libero.py +++ b/src/lerobot/envs/libero.py @@ -180,15 +180,25 @@ class LiberoEnv(gym.Env): "eef": spaces.Dict( { "pos": spaces.Box(low=-np.inf, high=np.inf, shape=(3,), dtype=np.float64), - "quat": spaces.Box(low=-np.inf, high=np.inf, shape=(4,), dtype=np.float64), - "mat": spaces.Box(low=-np.inf, high=np.inf, shape=(3, 3), dtype=np.float64), - "axisangle": spaces.Box(low=-np.inf, high=np.inf, shape=(3,), dtype=np.float64), + "quat": spaces.Box( + low=-np.inf, high=np.inf, shape=(4,), dtype=np.float64 + ), + "mat": spaces.Box( + low=-np.inf, high=np.inf, shape=(3, 3), dtype=np.float64 + ), + "axisangle": spaces.Box( + low=-np.inf, high=np.inf, shape=(3,), dtype=np.float64 + ), } ), "gripper": spaces.Dict( { - "qpos": spaces.Box(low=-np.inf, high=np.inf, shape=(2,), dtype=np.float64), - "qvel": spaces.Box(low=-np.inf, high=np.inf, shape=(2,), dtype=np.float64), + "qpos": spaces.Box( + low=-np.inf, high=np.inf, shape=(2,), dtype=np.float64 + ), + "qvel": spaces.Box( + low=-np.inf, high=np.inf, shape=(2,), dtype=np.float64 + ), } ), "joints": spaces.Dict( @@ -202,7 +212,6 @@ class LiberoEnv(gym.Env): } ) - self.action_space = spaces.Box( low=ACTION_LOW, high=ACTION_HIGH, shape=(ACTION_DIM,), dtype=np.float32 ) @@ -232,8 +241,8 @@ class LiberoEnv(gym.Env): for camera_name in self.camera_name: image = raw_obs[camera_name] images[self.camera_name_mapping[camera_name]] = image - - eef_pos = raw_obs.get("robot0_eef_pos") + + eef_pos = raw_obs.get("robot0_eef_pos") eef_quat = raw_obs.get("robot0_eef_quat") # rotation matrix from controller @@ -241,34 +250,33 @@ class LiberoEnv(gym.Env): 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") + 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) + "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,) + "qpos": gripper_qpos, # (2,) + "qvel": gripper_qvel, # (2,) }, "joints": { - "pos": joint_pos, # (7,) - "vel": joint_vel, # (7,) + "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 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')." diff --git a/src/lerobot/envs/utils.py b/src/lerobot/envs/utils.py index e9eb17611..8eb9bf501 100644 --- a/src/lerobot/envs/utils.py +++ b/src/lerobot/envs/utils.py @@ -83,7 +83,7 @@ def preprocess_observation(observations: dict[str, np.ndarray]) -> dict[str, Ten if agent_pos.dim() == 1: agent_pos = agent_pos.unsqueeze(0) return_observations[OBS_STATE] = agent_pos - + if "robot_state" in observations: # simply copy nested dict as-is return_observations[f"{OBS_STR}.robot_state"] = { diff --git a/src/lerobot/processor/observation_processor.py b/src/lerobot/processor/observation_processor.py index cf314b781..ed776f37c 100644 --- a/src/lerobot/processor/observation_processor.py +++ b/src/lerobot/processor/observation_processor.py @@ -25,11 +25,6 @@ from lerobot.utils.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_ST from .pipeline import ObservationProcessorStep, ProcessorStepRegistry -try: - from robosuite.utils.transform_utils import quat2axisangle -except ImportError: - quat2axisangle = None - @dataclass @ProcessorStepRegistry.register(name="observation_processor") @@ -234,12 +229,6 @@ class LiberoProcessorStep(ObservationProcessorStep): """ Processes both image and robot_state observations from LIBERO. """ - if quat2axisangle is None: - raise ImportError( - "robosuite is required for LiberoProcessorStep. " - "Install it with: pip install robosuite" - ) - processed_obs = observation.copy() # Process robot_state into a flat state vector @@ -252,8 +241,8 @@ class LiberoProcessorStep(ObservationProcessorStep): gripper_qpos = robot_state["gripper"]["qpos"] # (2,) # Convert quaternion to axis-angle - eef_axisangle = quat2axisangle(eef_quat.squeeze(0)) # (3,) - eef_axisangle = eef_axisangle[np.newaxis, :] # (1, 3) + eef_axisangle = self._quat2axisangle(eef_quat.squeeze(0)) # (3,) + eef_axisangle = eef_axisangle[np.newaxis, :] # (1, 3) # Concatenate into a single state vector state = np.concatenate((eef_pos, eef_axisangle, gripper_qpos), axis=1) @@ -274,7 +263,33 @@ class LiberoProcessorStep(ObservationProcessorStep): """ new_features: dict[PipelineFeatureType, dict[str, PolicyFeature]] = {ft: {} for ft in features} return new_features - + def observation(self, observation): return self._process_observation(observation) + def _quat2axisangle(self, quat): + """ + # Copied from robosuite.utils.transform_utils.quat2axisangle + Converts quaternion to axis-angle format. + Returns a unit vector direction scaled by its angle in radians. + + Args: + quat (np.array): (x,y,z,w) vec4 float angles + + Returns: + np.array: (ax,ay,az) axis-angle exponential coordinates + """ + import math + + # clip quaternion + if quat[3] > 1.0: + quat[3] = 1.0 + elif quat[3] < -1.0: + quat[3] = -1.0 + + den = np.sqrt(1.0 - quat[3] * quat[3]) + if math.isclose(den, 0.0): + # This is (close to) a zero degree rotation, immediately return + return np.zeros(3) + + return (quat[:3] * 2.0 * math.acos(quat[3])) / den diff --git a/src/lerobot/scripts/lerobot_eval.py b/src/lerobot/scripts/lerobot_eval.py index 76e6b1242..0d66fa1aa 100644 --- a/src/lerobot/scripts/lerobot_eval.py +++ b/src/lerobot/scripts/lerobot_eval.py @@ -165,7 +165,6 @@ def rollout( # Infer "task" from attributes of environments. # TODO: works with SyncVectorEnv but not AsyncVectorEnv observation = add_envs_task(env, observation) - breakpoint() observation = preprocessor(observation) with torch.inference_mode(): action = policy.select_action(observation) diff --git a/test2.py b/test2.py deleted file mode 100644 index 9684333e0..000000000 --- a/test2.py +++ /dev/null @@ -1,22 +0,0 @@ -from lerobot.envs.factory import make_env, make_env_config -from lerobot.envs.utils import add_envs_task, preprocess_observation -from lerobot.processor.pipeline import PolicyProcessorPipeline -from lerobot.processor.observation_processor import LiberoProcessorStep -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) - -observation = preprocess_observation(obs1) -observation = add_envs_task(env, observation) - -libero_preprocessor = PolicyProcessorPipeline( - steps=[ - LiberoProcessorStep(), - ] -) -observation = libero_preprocessor(observation) diff --git a/test3.py b/test3.py deleted file mode 100644 index 83b28ee2a..000000000 --- a/test3.py +++ /dev/null @@ -1,20 +0,0 @@ -from lerobot.processor.observation_processor import VanillaObservationProcessorStep -from lerobot.processor.converters import create_transition -from lerobot.processor import TransitionKey -from lerobot.utils.constants import OBS_IMAGE -import numpy as np - -processor = VanillaObservationProcessorStep() - -# Create a mock image (H, W, C) format, uint8 -image = np.random.randint(0, 256, size=(64, 64, 3), dtype=np.uint8) - -observation = {"pixels": image} -transition = create_transition(observation=observation) -breakpoint() -result = processor(transition) -processed_obs = result[TransitionKey.OBSERVATION] - -# Check that the image was processed correctly -assert OBS_IMAGE in processed_obs -processed_img = processed_obs[OBS_IMAGE] diff --git a/test_libero.py b/test_libero.py deleted file mode 100644 index 75155b696..000000000 --- a/test_libero.py +++ /dev/null @@ -1,148 +0,0 @@ -#!/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 diff --git a/tests/processor/test_libero_processor.py b/tests/processor/test_libero_processor.py new file mode 100644 index 000000000..7717490b1 --- /dev/null +++ b/tests/processor/test_libero_processor.py @@ -0,0 +1,71 @@ +#!/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 numpy as np +from lerobot.envs.utils import preprocess_observation +from lerobot.processor.pipeline import PolicyProcessorPipeline +from lerobot.processor.observation_processor import LiberoProcessorStep +import torch +seed = 42 +np.random.seed(seed) + +obs1 = { + "pixels": { + "image": (np.random.rand(1, 256, 256, 3) * 255).astype(np.uint8), + "image2": (np.random.rand(1, 256, 256, 3) * 255).astype(np.uint8), + }, + "robot_state": { + "eef": { + "pos": np.random.randn(1, 3), + "quat": np.random.randn(1, 4), + "mat": np.random.randn(1, 3, 3), + "axisangle": np.random.randn(1, 3), + }, + "gripper": { + "qpos": np.random.randn(1, 2), + "qvel": np.random.randn(1, 2), + }, + "joints": { + "pos": np.random.randn(1, 7), + "vel": np.random.randn(1, 7), + } + } +} + +observation = preprocess_observation(obs1) + +libero_preprocessor = PolicyProcessorPipeline( + steps=[ + LiberoProcessorStep(), + ] +) +processed_obs = libero_preprocessor(observation) +assert "observation.state" in processed_obs +state = processed_obs["observation.state"] +assert isinstance(state, torch.Tensor) +assert state.dtype == torch.float32 + +assert state.shape[0] == 1 +assert state.shape[1] == 8 + +assert "observation.images.image" in processed_obs +assert "observation.images.image2" in processed_obs + +assert isinstance(processed_obs["observation.images.image"], torch.Tensor) +assert isinstance(processed_obs["observation.images.image2"], torch.Tensor) + +assert processed_obs["observation.images.image"].shape == (1, 3, 256, 256) +assert processed_obs["observation.images.image2"].shape == (1, 3, 256, 256)