Fix EVO1 LIBERO rollout processors

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:

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(

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",

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

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