diff --git a/src/lerobot/envs/configs.py b/src/lerobot/envs/configs.py index 84c40472f..2e4de8c51 100644 --- a/src/lerobot/envs/configs.py +++ b/src/lerobot/envs/configs.py @@ -24,7 +24,12 @@ import gymnasium as gym from gymnasium.envs.registration import registry as gym_registry from lerobot.configs import FeatureType, PolicyFeature -from lerobot.processor import IsaaclabArenaProcessorStep, LiberoProcessorStep, PolicyProcessorPipeline +from lerobot.processor import ( + IsaaclabArenaProcessorStep, + LiberoActionProcessorStep, + LiberoProcessorStep, + PolicyProcessorPipeline, +) from lerobot.robots import RobotConfig from lerobot.teleoperators.config import TeleoperatorConfig from lerobot.utils.constants import ( @@ -123,7 +128,7 @@ class EnvConfig(draccus.ChoiceRegistry, abc.ABC): vec = env_cls([_make_one for _ in range(n_envs)], **extra_kwargs) return {self.type: {0: vec}} - def get_env_processors(self): + def get_env_processors(self, policy_cfg: Any | None = None): """Return (preprocessor, postprocessor) for this env. Default: identity.""" return PolicyProcessorPipeline(steps=[]), PolicyProcessorPipeline(steps=[]) @@ -436,10 +441,13 @@ class LiberoEnv(EnvConfig): is_libero_plus=self.is_libero_plus, ) - def get_env_processors(self): + def get_env_processors(self, policy_cfg: Any | None = None): + max_state_dim = getattr(policy_cfg, "max_state_dim", None) if getattr(policy_cfg, "type", None) == "evo1" else None + action_feature = self.features.get(ACTION) + action_dim = int(action_feature.shape[0]) if action_feature is not None else 7 return ( - PolicyProcessorPipeline(steps=[LiberoProcessorStep()]), - PolicyProcessorPipeline(steps=[]), + PolicyProcessorPipeline(steps=[LiberoProcessorStep(max_state_dim=max_state_dim)]), + PolicyProcessorPipeline(steps=[LiberoActionProcessorStep(action_dim=action_dim)]), ) @@ -705,7 +713,7 @@ class IsaaclabArenaEnv(HubEnvConfig): def gym_kwargs(self) -> dict: return {} - def get_env_processors(self): + def get_env_processors(self, policy_cfg: Any | None = None): state_keys = tuple(k.strip() for k in (self.state_keys or "").split(",") if k.strip()) camera_keys = tuple(k.strip() for k in (self.camera_keys or "").split(",") if k.strip()) if not state_keys and not camera_keys: diff --git a/src/lerobot/envs/factory.py b/src/lerobot/envs/factory.py index 317cf2e6f..22f43f4cd 100644 --- a/src/lerobot/envs/factory.py +++ b/src/lerobot/envs/factory.py @@ -15,6 +15,7 @@ # limitations under the License. from __future__ import annotations +import inspect from typing import Any import gymnasium as gym @@ -52,7 +53,14 @@ def make_env_pre_post_processors( return make_xvla_libero_pre_post_processors() - return env_cfg.get_env_processors() + get_processors = env_cfg.get_env_processors + signature = inspect.signature(get_processors) + supports_policy_cfg = "policy_cfg" in signature.parameters or any( + param.kind is inspect.Parameter.VAR_KEYWORD for param in signature.parameters.values() + ) + if supports_policy_cfg: + return get_processors(policy_cfg=policy_cfg) + return get_processors() def make_env( diff --git a/src/lerobot/processor/__init__.py b/src/lerobot/processor/__init__.py index 3688a4b8c..3329729fb 100644 --- a/src/lerobot/processor/__init__.py +++ b/src/lerobot/processor/__init__.py @@ -40,7 +40,7 @@ from .converters import ( ) from .delta_action_processor import MapDeltaActionToRobotActionStep, MapTensorToDeltaActionDictStep from .device_processor import DeviceProcessorStep -from .env_processor import IsaaclabArenaProcessorStep, LiberoProcessorStep +from .env_processor import IsaaclabArenaProcessorStep, LiberoActionProcessorStep, LiberoProcessorStep from .factory import ( make_default_processors, make_default_robot_action_processor, @@ -149,6 +149,7 @@ __all__ = [ "RewardProcessorStep", "DataProcessorPipeline", "IsaaclabArenaProcessorStep", + "LiberoActionProcessorStep", "LiberoProcessorStep", "TimeLimitProcessorStep", "AddBatchDimensionProcessorStep", diff --git a/src/lerobot/processor/env_processor.py b/src/lerobot/processor/env_processor.py index 75cbb79de..c8287b1aa 100644 --- a/src/lerobot/processor/env_processor.py +++ b/src/lerobot/processor/env_processor.py @@ -18,9 +18,9 @@ from dataclasses import dataclass import torch from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature -from lerobot.utils.constants import OBS_IMAGES, OBS_PREFIX, OBS_STATE, OBS_STR +from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_PREFIX, OBS_STATE, OBS_STR -from .pipeline import ObservationProcessorStep, ProcessorStepRegistry +from .pipeline import ActionProcessorStep, ObservationProcessorStep, ProcessorStepRegistry @dataclass @@ -46,6 +46,8 @@ class LiberoProcessorStep(ObservationProcessorStep): - This accounts for the HuggingFaceVLA/libero camera orientation convention. """ + max_state_dim: int | None = None + def _process_observation(self, observation): """ Processes both image and robot_state observations from LIBERO. @@ -78,6 +80,15 @@ class LiberoProcessorStep(ObservationProcessorStep): state = state.float() if state.dim() == 1: state = state.unsqueeze(0) + if self.max_state_dim is not None: + if state.shape[-1] > self.max_state_dim: + raise ValueError( + f"LIBERO state has {state.shape[-1]} dims, which is larger than " + f"configured max_state_dim={self.max_state_dim}." + ) + if state.shape[-1] < self.max_state_dim: + pad_width = self.max_state_dim - state.shape[-1] + state = torch.nn.functional.pad(state, (0, pad_width)) processed_obs[OBS_STATE] = state return processed_obs @@ -101,7 +112,7 @@ class LiberoProcessorStep(ObservationProcessorStep): # add our new flattened state state_feats[OBS_STATE] = PolicyFeature( type=FeatureType.STATE, - shape=(8,), # [eef_pos(3), axis_angle(3), gripper(2)] + shape=(self.max_state_dim or 8,), # [eef_pos(3), axis_angle(3), gripper(2)] plus padding ) new_features[FeatureType.STATE] = state_feats @@ -111,6 +122,9 @@ class LiberoProcessorStep(ObservationProcessorStep): def observation(self, observation): return self._process_observation(observation) + def get_config(self) -> dict: + return {"max_state_dim": self.max_state_dim} + def _quat2axisangle(self, quat: torch.Tensor) -> torch.Tensor: """ Convert batched quaternions to axis-angle format. @@ -153,6 +167,32 @@ class LiberoProcessorStep(ObservationProcessorStep): return result +@dataclass +@ProcessorStepRegistry.register(name="libero_action_processor") +class LiberoActionProcessorStep(ActionProcessorStep): + """Slices padded policy actions back to the executable LIBERO action space.""" + + action_dim: int = 7 + + def action(self, action): + if action.shape[-1] < self.action_dim: + raise ValueError( + f"LIBERO action has {action.shape[-1]} dims, which is smaller than action_dim={self.action_dim}." + ) + return action[..., : self.action_dim] + + def transform_features( + self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]] + ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]: + new_features = {ft: feats.copy() for ft, feats in features.items()} + action_feats = new_features.setdefault(FeatureType.ACTION, {}) + action_feats[ACTION] = PolicyFeature(type=FeatureType.ACTION, shape=(self.action_dim,)) + return new_features + + def get_config(self) -> dict: + return {"action_dim": self.action_dim} + + @dataclass @ProcessorStepRegistry.register(name="isaaclab_arena_processor") class IsaaclabArenaProcessorStep(ObservationProcessorStep): diff --git a/tests/envs/test_dispatch.py b/tests/envs/test_dispatch.py index 5bd2827f3..e26f34690 100644 --- a/tests/envs/test_dispatch.py +++ b/tests/envs/test_dispatch.py @@ -7,11 +7,14 @@ from dataclasses import dataclass, field import gymnasium as gym import pytest +import torch from gymnasium.envs.registration import register, registry as gym_registry from lerobot.configs.types import PolicyFeature -from lerobot.envs.configs import EnvConfig +from lerobot.envs.configs import EnvConfig, LiberoEnv from lerobot.envs.factory import make_env, make_env_config, make_env_pre_post_processors +from lerobot.processor import LiberoActionProcessorStep, LiberoProcessorStep +from lerobot.utils.constants import OBS_PREFIX, OBS_STATE logger = logging.getLogger(__name__) @@ -61,6 +64,80 @@ def test_processors_delegation(): assert len(pre.steps) == 0 +def test_processors_delegation_supports_legacy_override_signature(): + """External EnvConfig subclasses with the old get_env_processors() signature keep working.""" + from lerobot.processor.pipeline import DataProcessorPipeline + + @EnvConfig.register_subclass("_dispatch_legacy_proc_test") + @dataclass + class _Env(EnvConfig): + task: str = "x" + features: dict[str, PolicyFeature] = field(default_factory=dict) + + @property + def gym_kwargs(self): + return {} + + def get_env_processors(self): + return DataProcessorPipeline(steps=[]), DataProcessorPipeline(steps=[]) + + pre, post = make_env_pre_post_processors(_Env(), policy_cfg=object()) + assert isinstance(pre, DataProcessorPipeline) + assert isinstance(post, DataProcessorPipeline) + + +def test_libero_evo1_processors_use_padded_state_and_env_action_dim(): + """EVO1 uses padded LIBERO state features while env actions stay executable.""" + + class _Evo1Config: + type = "evo1" + max_state_dim = 24 + + cfg = LiberoEnv() + pre, post = make_env_pre_post_processors(cfg, policy_cfg=_Evo1Config()) + assert isinstance(pre.steps[0], LiberoProcessorStep) + assert pre.steps[0].max_state_dim == 24 + assert isinstance(post.steps[0], LiberoActionProcessorStep) + assert post.steps[0].action_dim == cfg.features["action"].shape[0] == 7 + + class _OtherConfig: + type = "other" + + pre_other, _ = make_env_pre_post_processors(cfg, policy_cfg=_OtherConfig()) + assert pre_other.steps[0].max_state_dim is None + + +def test_libero_processor_pads_state_to_max_dim(): + step = LiberoProcessorStep(max_state_dim=24) + observation = { + OBS_PREFIX + + "robot_state": { + "eef": { + "pos": torch.tensor([[1.0, 2.0, 3.0]]), + "quat": torch.tensor([[0.0, 0.0, 0.0, 1.0]]), + }, + "gripper": {"qpos": torch.tensor([[4.0, 5.0]])}, + } + } + + state = step.observation(observation)[OBS_STATE] + assert state.shape == (1, 24) + assert torch.allclose(state[:, :8], torch.tensor([[1.0, 2.0, 3.0, 0.0, 0.0, 0.0, 4.0, 5.0]])) + assert torch.count_nonzero(state[:, 8:]).item() == 0 + + +def test_libero_action_processor_slices_padded_action(): + step = LiberoActionProcessorStep(action_dim=7) + action = torch.arange(2 * 3 * 24, dtype=torch.float32).reshape(2, 3, 24) + + sliced = step.action(action) + assert sliced.shape == (2, 3, 7) + assert torch.equal(sliced, action[..., :7]) + + with pytest.raises(ValueError, match="smaller than action_dim=7"): + step.action(torch.zeros(2, 6)) + + def test_base_create_envs(): """Base class create_envs() should build a single-task VectorEnv via gym.make().""" gym_id = "_dispatch_test/CartPole-v99" @@ -136,7 +213,7 @@ def test_custom_get_env_processors_override(): def gym_kwargs(self): return {} - def get_env_processors(self): + def get_env_processors(self, policy_cfg=None): return DataProcessorPipeline(steps=[]), DataProcessorPipeline(steps=[]) pre, post = _Env().get_env_processors()