mirror of
https://github.com/huggingface/lerobot.git
synced 2026-05-30 16:09:44 +00:00
more changes
This commit is contained in:
@@ -261,13 +261,44 @@ class LiberoEnv(EnvConfig):
|
|||||||
type=FeatureType.VISUAL, shape=(self.observation_height, self.observation_width, 3)
|
type=FeatureType.VISUAL, shape=(self.observation_height, self.observation_width, 3)
|
||||||
)
|
)
|
||||||
elif self.obs_type == "pixels_agent_pos":
|
elif self.obs_type == "pixels_agent_pos":
|
||||||
self.features["agent_pos"] = PolicyFeature(type=FeatureType.STATE, shape=(8,))
|
|
||||||
self.features["pixels/agentview_image"] = PolicyFeature(
|
self.features["pixels/agentview_image"] = PolicyFeature(
|
||||||
type=FeatureType.VISUAL, shape=(self.observation_height, self.observation_width, 3)
|
type=FeatureType.VISUAL, shape=(self.observation_height, self.observation_width, 3)
|
||||||
)
|
)
|
||||||
self.features["pixels/robot0_eye_in_hand_image"] = PolicyFeature(
|
self.features["pixels/robot0_eye_in_hand_image"] = PolicyFeature(
|
||||||
type=FeatureType.VISUAL, shape=(self.observation_height, self.observation_width, 3)
|
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:
|
else:
|
||||||
raise ValueError(f"Unsupported obs_type: {self.obs_type}")
|
raise ValueError(f"Unsupported obs_type: {self.obs_type}")
|
||||||
|
|
||||||
|
|||||||
+35
-15
@@ -212,23 +212,44 @@ class LiberoEnv(gym.Env):
|
|||||||
images = {}
|
images = {}
|
||||||
for camera_name in self.camera_name:
|
for camera_name in self.camera_name:
|
||||||
image = raw_obs[camera_name]
|
image = raw_obs[camera_name]
|
||||||
image = image[::-1, ::-1] # rotate 180 degrees
|
|
||||||
images[self.camera_name_mapping[camera_name]] = image
|
images[self.camera_name_mapping[camera_name]] = image
|
||||||
state = np.concatenate(
|
|
||||||
(
|
eef_pos = raw_obs.get("robot0_eef_pos")
|
||||||
raw_obs["robot0_eef_pos"],
|
eef_quat = raw_obs.get("robot0_eef_quat")
|
||||||
quat2axisangle(raw_obs["robot0_eef_quat"]),
|
|
||||||
raw_obs["robot0_gripper_qpos"],
|
# 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
|
||||||
agent_pos = state
|
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":
|
if self.obs_type == "pixels":
|
||||||
return {"pixels": images.copy()}
|
return {"pixels": images.copy()}
|
||||||
|
|
||||||
if self.obs_type == "pixels_agent_pos":
|
if self.obs_type == "pixels_agent_pos":
|
||||||
return {
|
return obs
|
||||||
"pixels": images.copy(),
|
|
||||||
"agent_pos": agent_pos,
|
|
||||||
}
|
|
||||||
raise NotImplementedError(
|
raise NotImplementedError(
|
||||||
f"The observation type '{self.obs_type}' is not supported in LiberoEnv. "
|
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')."
|
"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}")
|
print(f"Restricting to task_ids={task_ids_filter}")
|
||||||
|
|
||||||
out: dict[str, dict[int, Any]] = defaultdict(dict)
|
out: dict[str, dict[int, Any]] = defaultdict(dict)
|
||||||
|
|
||||||
for suite_name in suite_names:
|
for suite_name in suite_names:
|
||||||
suite = _get_suite(suite_name)
|
suite = _get_suite(suite_name)
|
||||||
total = len(suite.tasks)
|
total = len(suite.tasks)
|
||||||
selected = _select_task_ids(total, task_ids_filter)
|
selected = _select_task_ids(total, task_ids_filter)
|
||||||
|
selected = [0]
|
||||||
if not selected:
|
if not selected:
|
||||||
raise ValueError(f"No tasks selected for suite '{suite_name}' (available: {total}).")
|
raise ValueError(f"No tasks selected for suite '{suite_name}' (available: {total}).")
|
||||||
|
|
||||||
|
|||||||
@@ -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)
|
||||||
+148
@@ -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()
|
||||||
Reference in New Issue
Block a user