more changes

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}")
 

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}).")
 

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)

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()