refactor(envs): move benchmark dispatch into EnvConfig subclasses (#3272)

* docs(benchmarks): add benchmark integration guide and standardize benchmark docs Add a comprehensive guide for adding new benchmarks to LeRobot, and refactor the existing LIBERO and Meta-World docs to follow the new standardized template. * refactor(envs): move dispatch logic from factory into EnvConfig subclasses Replace hardcoded if/elif chains in factory.py with create_envs() and get_env_processors() methods on EnvConfig. New benchmarks now only need to register a config subclass — no factory.py edits required. Net -23 lines: factory.py shrinks from ~200 to ~70 lines of logic. * docs(benchmarks): clean up adding-benchmarks guide for clarity Rewrite for simpler language, better structure, and easier navigation. Move quick-reference table to the top, fold eval explanation into architecture section, condense the doc template to a bulleted outline. * fix link * fix task count * fix(tests): fix 3 failing dispatch tests - test_registry_all_types: skip non-EnvConfig stubs (e.g. TestPluginConfig) - test_processors_delegation: use None instead of abstract PreTrainedConfig - test_custom_get_env_processors_override: use DataProcessorPipeline for isinstance check (PolicyProcessorPipeline is a subscripted generic) * fix: enable SmolVLA eval on LIBERO with custom camera mappings - Thread camera_name_mapping from LiberoEnv config through to gym envs - Sync features_map with camera_name_mapping in LiberoEnv.__post_init__ - Fix render() to use first available camera instead of hardcoded "image" - Handle non-dict final_info in rollout by falling back to info["is_success"] - Add use_peft legacy field to SmolVLAConfig for checkpoint compat - Add defaults to GR00TN15Config init=False fields for transformers 5.3 Made-with: Cursor * fix: use direct AutoresetMode import for gymnasium compat Made-with: Cursor * fix: handle gymnasium < 1.0 without AutoresetMode Made-with: Cursor * refactor: revert policy changes, keep env-only camera mapping fixes - Revert GR00T N1.5 default_factory/default changes (transformers compat) - Revert SmolVLA use_peft legacy field - Apply ruff formatting fixes - camera_name_mapping stays entirely in env/eval layer (no policy changes) Made-with: Cursor * Update docs/source/env_processor.mdx Co-authored-by: Khalil Meftah <khalil.meftah@huggingface.co> Signed-off-by: Pepijn <138571049+pkooij@users.noreply.github.com> * Update docs/source/env_processor.mdx Co-authored-by: Khalil Meftah <khalil.meftah@huggingface.co> Signed-off-by: Pepijn <138571049+pkooij@users.noreply.github.com> * Update docs/source/env_processor.mdx Co-authored-by: Khalil Meftah <khalil.meftah@huggingface.co> Signed-off-by: Pepijn <138571049+pkooij@users.noreply.github.com> * fix(eval): raise RuntimeError for unsupported final_info format (Gymnasium < 1.0) Made-with: Cursor * style: fix markdown code fences in env_processor.mdx Made-with: Cursor * docs: remove duplicate code blocks in env_processor.mdx Made-with: Cursor * style: revert quadruple backticks to triple (prettier compat) * docs(env_processor): add EnvConfig subclass step and policy_cfg examples - Add missing '### 2. Update Your EnvConfig Subclass' section with get_env_processors() snippet - Update factory usage example to show policy_cfg parameter and keyword-argument style for both SmolVLA and ACT cases * docs(env_processor): rename step 2 and fix policy_cfg examples - Rename '### 2. Update the Factory' → '### 2. Update Your EnvConfig Subclass' - Update factory usage examples to use keyword-argument style with policy_cfg parameter for both SmolVLA and ACT cases --------- Signed-off-by: Pepijn <138571049+pkooij@users.noreply.github.com> Co-authored-by: Khalil Meftah <khalil.meftah@huggingface.co>
2026-05-16 00:59:46 +00:00 · 2026-04-08 17:48:58 +02:00
parent 4eecbad32b
commit 5de7aa5a4f
7 changed files with 361 additions and 217 deletions
@@ -12,11 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.

+from __future__ import annotations
+
 import abc
+import importlib
 from dataclasses import dataclass, field, fields
 from typing import Any

 import draccus
+import gymnasium as gym
+from gymnasium.envs.registration import registry as gym_registry

 from lerobot.configs.types import FeatureType, PolicyFeature
 from lerobot.robots import RobotConfig
@@ -67,6 +72,49 @@ class EnvConfig(draccus.ChoiceRegistry, abc.ABC):
    def gym_kwargs(self) -> dict:
        raise NotImplementedError()

+    def create_envs(
+        self,
+        n_envs: int,
+        use_async_envs: bool = False,
+    ) -> dict[str, dict[int, gym.vector.VectorEnv]]:
+        """Create {suite: {task_id: VectorEnv}}.
+
+        Default: single-task env via gym.make(). Multi-task benchmarks override.
+        """
+        env_cls = gym.vector.AsyncVectorEnv if use_async_envs else gym.vector.SyncVectorEnv
+
+        if self.gym_id not in gym_registry:
+            print(f"gym id '{self.gym_id}' not found, attempting to import '{self.package_name}'...")
+            try:
+                importlib.import_module(self.package_name)
+            except ModuleNotFoundError as e:
+                raise ModuleNotFoundError(
+                    f"Package '{self.package_name}' required for env '{self.type}' not found. "
+                    f"Please install it or check PYTHONPATH."
+                ) from e
+
+            if self.gym_id not in gym_registry:
+                raise gym.error.NameNotFound(
+                    f"Environment '{self.gym_id}' not registered even after importing '{self.package_name}'."
+                )
+
+        def _make_one():
+            return gym.make(self.gym_id, disable_env_checker=self.disable_env_checker, **self.gym_kwargs)
+
+        try:
+            from gymnasium.vector import AutoresetMode
+
+            vec = env_cls([_make_one for _ in range(n_envs)], autoreset_mode=AutoresetMode.SAME_STEP)
+        except ImportError:
+            vec = env_cls([_make_one for _ in range(n_envs)])
+        return {self.type: {0: vec}}
+
+    def get_env_processors(self):
+        """Return (preprocessor, postprocessor) for this env. Default: identity."""
+        from lerobot.processor.pipeline import PolicyProcessorPipeline
+
+        return PolicyProcessorPipeline(steps=[]), PolicyProcessorPipeline(steps=[])
+

@dataclass
 class HubEnvConfig(EnvConfig):
@@ -338,6 +386,12 @@ class LiberoEnv(EnvConfig):
        else:
            raise ValueError(f"Unsupported obs_type: {self.obs_type}")

+        if self.camera_name_mapping is not None:
+            mapped_agentview = self.camera_name_mapping.get("agentview_image", "image")
+            mapped_eye_in_hand = self.camera_name_mapping.get("robot0_eye_in_hand_image", "image2")
+            self.features_map[LIBERO_KEY_PIXELS_AGENTVIEW] = f"{OBS_IMAGES}.{mapped_agentview}"
+            self.features_map[LIBERO_KEY_PIXELS_EYE_IN_HAND] = f"{OBS_IMAGES}.{mapped_eye_in_hand}"
+
    @property
    def gym_kwargs(self) -> dict:
        kwargs: dict[str, Any] = {"obs_type": self.obs_type, "render_mode": self.render_mode}
@@ -345,6 +399,33 @@ class LiberoEnv(EnvConfig):
            kwargs["task_ids"] = self.task_ids
        return kwargs

+    def create_envs(self, n_envs: int, use_async_envs: bool = False):
+        from lerobot.envs.libero import create_libero_envs
+
+        if self.task is None:
+            raise ValueError("LiberoEnv requires a task to be specified")
+        env_cls = gym.vector.AsyncVectorEnv if use_async_envs else gym.vector.SyncVectorEnv
+        return create_libero_envs(
+            task=self.task,
+            n_envs=n_envs,
+            camera_name=self.camera_name,
+            init_states=self.init_states,
+            gym_kwargs=self.gym_kwargs,
+            env_cls=env_cls,
+            control_mode=self.control_mode,
+            episode_length=self.episode_length,
+            camera_name_mapping=self.camera_name_mapping,
+        )
+
+    def get_env_processors(self):
+        from lerobot.processor.env_processor import LiberoProcessorStep
+        from lerobot.processor.pipeline import PolicyProcessorPipeline
+
+        return (
+            PolicyProcessorPipeline(steps=[LiberoProcessorStep()]),
+            PolicyProcessorPipeline(steps=[]),
+        )
+

@EnvConfig.register_subclass("metaworld")
@dataclass
@@ -387,6 +468,19 @@ class MetaworldEnv(EnvConfig):
            "render_mode": self.render_mode,
        }

+    def create_envs(self, n_envs: int, use_async_envs: bool = False):
+        from lerobot.envs.metaworld import create_metaworld_envs
+
+        if self.task is None:
+            raise ValueError("MetaWorld requires a task to be specified")
+        env_cls = gym.vector.AsyncVectorEnv if use_async_envs else gym.vector.SyncVectorEnv
+        return create_metaworld_envs(
+            task=self.task,
+            n_envs=n_envs,
+            gym_kwargs=self.gym_kwargs,
+            env_cls=env_cls,
+        )
+

@EnvConfig.register_subclass("isaaclab_arena")
@dataclass
@@ -454,3 +548,18 @@ class IsaaclabArenaEnv(HubEnvConfig):
    @property
    def gym_kwargs(self) -> dict:
        return {}
+
+    def get_env_processors(self):
+        from lerobot.processor.env_processor import IsaaclabArenaProcessorStep
+        from lerobot.processor.pipeline import PolicyProcessorPipeline
+
+        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:
+            raise ValueError("At least one of state_keys or camera_keys must be specified.")
+        return (
+            PolicyProcessorPipeline(
+                steps=[IsaaclabArenaProcessorStep(state_keys=state_keys, camera_keys=camera_keys)]
+            ),
+            PolicyProcessorPipeline(steps=[]),
+        )
@@ -13,90 +13,46 @@
 # 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 importlib
+from __future__ import annotations
+
 from typing import Any

 import gymnasium as gym
-from gymnasium.envs.registration import registry as gym_registry

-from lerobot.configs.policies import PreTrainedConfig
-from lerobot.envs.configs import AlohaEnv, EnvConfig, HubEnvConfig, IsaaclabArenaEnv, LiberoEnv, PushtEnv
+from lerobot.envs.configs import EnvConfig, HubEnvConfig
 from lerobot.envs.utils import _call_make_env, _download_hub_file, _import_hub_module, _normalize_hub_result
-from lerobot.policies.xvla.configuration_xvla import XVLAConfig
-from lerobot.processor import ProcessorStep
-from lerobot.processor.env_processor import IsaaclabArenaProcessorStep, LiberoProcessorStep
-from lerobot.processor.pipeline import PolicyProcessorPipeline


 def make_env_config(env_type: str, **kwargs) -> EnvConfig:
-    if env_type == "aloha":
-        return AlohaEnv(**kwargs)
-    elif env_type == "pusht":
-        return PushtEnv(**kwargs)
-    elif env_type == "libero":
-        return LiberoEnv(**kwargs)
-    else:
-        raise ValueError(f"Policy type '{env_type}' is not available.")
+    try:
+        cls = EnvConfig.get_choice_class(env_type)
+    except KeyError as err:
+        raise ValueError(
+            f"Environment type '{env_type}' is not registered. "
+            f"Available: {list(EnvConfig.get_known_choices().keys())}"
+        ) from err
+    return cls(**kwargs)


 def make_env_pre_post_processors(
    env_cfg: EnvConfig,
-    policy_cfg: PreTrainedConfig,
-) -> tuple[
-    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
-    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
-]:
+    policy_cfg: Any,
+) -> tuple[Any, Any]:
    """
    Create preprocessor and postprocessor pipelines for environment observations.

-    This function creates processor pipelines that transform raw environment
-    observations and actions. By default, it returns identity processors that do nothing.
-    For specific environments like LIBERO, it adds environment-specific processing steps.
-
-    Args:
-        env_cfg: The configuration of the environment.
-
-    Returns:
-        A tuple containing:
-            - preprocessor: Pipeline that processes environment observations
-            - postprocessor: Pipeline that processes environment outputs (currently identity)
+    Returns a tuple of (preprocessor, postprocessor). By default, delegates to
+    ``env_cfg.get_env_processors()``.  The XVLAConfig policy-specific override
+    stays here because it depends on the *policy* config, not the env config.
    """
-    # Preprocessor and Postprocessor steps are Identity for most environments
-    preprocessor_steps: list[ProcessorStep] = []
-    postprocessor_steps: list[ProcessorStep] = []
+    from lerobot.policies.xvla.configuration_xvla import XVLAConfig
+
    if isinstance(policy_cfg, XVLAConfig):
        from lerobot.policies.xvla.processor_xvla import make_xvla_libero_pre_post_processors

        return make_xvla_libero_pre_post_processors()

-    # For LIBERO environments, add the LiberoProcessorStep to preprocessor
-    if isinstance(env_cfg, LiberoEnv) or "libero" in env_cfg.type:
-        preprocessor_steps.append(LiberoProcessorStep())
-
-    # For Isaaclab Arena environments, add the IsaaclabArenaProcessorStep
-    if isinstance(env_cfg, IsaaclabArenaEnv) or "isaaclab_arena" in env_cfg.type:
-        # Parse comma-separated keys (handle None for state-based policies)
-        if env_cfg.state_keys:
-            state_keys = tuple(k.strip() for k in env_cfg.state_keys.split(",") if k.strip())
-        else:
-            state_keys = ()
-        if env_cfg.camera_keys:
-            camera_keys = tuple(k.strip() for k in env_cfg.camera_keys.split(",") if k.strip())
-        else:
-            camera_keys = ()
-        if not state_keys and not camera_keys:
-            raise ValueError("At least one of state_keys or camera_keys must be specified.")
-        preprocessor_steps.append(
-            IsaaclabArenaProcessorStep(
-                state_keys=state_keys,
-                camera_keys=camera_keys,
-            )
-        )
-
-    preprocessor = PolicyProcessorPipeline(steps=preprocessor_steps)
-    postprocessor = PolicyProcessorPipeline(steps=postprocessor_steps)
-
-    return preprocessor, postprocessor
+    return env_cfg.get_env_processors()


 def make_env(
@@ -163,57 +119,4 @@ def make_env(
    if n_envs < 1:
        raise ValueError("`n_envs` must be at least 1")

-    env_cls = gym.vector.AsyncVectorEnv if use_async_envs else gym.vector.SyncVectorEnv
-
-    if "libero" in cfg.type:
-        from lerobot.envs.libero import create_libero_envs
-
-        if cfg.task is None:
-            raise ValueError("LiberoEnv requires a task to be specified")
-
-        return create_libero_envs(
-            task=cfg.task,
-            n_envs=n_envs,
-            camera_name=cfg.camera_name,
-            init_states=cfg.init_states,
-            gym_kwargs=cfg.gym_kwargs,
-            env_cls=env_cls,
-            control_mode=cfg.control_mode,
-            episode_length=cfg.episode_length,
-        )
-    elif "metaworld" in cfg.type:
-        from lerobot.envs.metaworld import create_metaworld_envs
-
-        if cfg.task is None:
-            raise ValueError("MetaWorld requires a task to be specified")
-
-        return create_metaworld_envs(
-            task=cfg.task,
-            n_envs=n_envs,
-            gym_kwargs=cfg.gym_kwargs,
-            env_cls=env_cls,
-        )
-
-    if cfg.gym_id not in gym_registry:
-        print(f"gym id '{cfg.gym_id}' not found, attempting to import '{cfg.package_name}'...")
-        try:
-            importlib.import_module(cfg.package_name)
-        except ModuleNotFoundError as e:
-            raise ModuleNotFoundError(
-                f"Package '{cfg.package_name}' required for env '{cfg.type}' not found. "
-                f"Please install it or check PYTHONPATH."
-            ) from e
-
-        if cfg.gym_id not in gym_registry:
-            raise gym.error.NameNotFound(
-                f"Environment '{cfg.gym_id}' not registered even after importing '{cfg.package_name}'."
-            )
-
-    def _make_one():
-        return gym.make(cfg.gym_id, disable_env_checker=cfg.disable_env_checker, **(cfg.gym_kwargs or {}))
-
-    vec = env_cls([_make_one for _ in range(n_envs)], autoreset_mode=gym.vector.AutoresetMode.SAME_STEP)
-
-    # normalize to {suite: {task_id: vec_env}} for consistency
-    suite_name = cfg.type  # e.g., "pusht", "aloha"
-    return {suite_name: {0: vec}}
+    return cfg.create_envs(n_envs=n_envs, use_async_envs=use_async_envs)
@@ -223,7 +223,8 @@ class LiberoEnv(gym.Env):

    def render(self):
        raw_obs = self._env.env._get_observations()
-        image = self._format_raw_obs(raw_obs)["pixels"]["image"]
+        pixels = self._format_raw_obs(raw_obs)["pixels"]
+        image = next(iter(pixels.values()))
        image = image[::-1, ::-1]  # flip both H and W for visualization
        return image

@@ -339,12 +340,6 @@ class LiberoEnv(gym.Env):
        )
        observation = self._format_raw_obs(raw_obs)
        if terminated:
-            info["final_info"] = {
-                "task": self.task,
-                "task_id": self.task_id,
-                "done": bool(done),
-                "is_success": bool(is_success),
-            }
            self.reset()
        truncated = False
        return observation, reward, terminated, truncated, info
@@ -364,6 +359,7 @@ def _make_env_fns(
    init_states: bool,
    gym_kwargs: Mapping[str, Any],
    control_mode: str,
+    camera_name_mapping: dict[str, str] | None = None,
 ) -> list[Callable[[], LiberoEnv]]:
    """Build n_envs factory callables for a single (suite, task_id)."""

@@ -379,6 +375,7 @@ def _make_env_fns(
            episode_index=episode_index,
            n_envs=n_envs,
            control_mode=control_mode,
+            camera_name_mapping=camera_name_mapping,
            **local_kwargs,
        )

@@ -400,6 +397,7 @@ def create_libero_envs(
    env_cls: Callable[[Sequence[Callable[[], Any]]], Any] | None = None,
    control_mode: str = "relative",
    episode_length: int | None = None,
+    camera_name_mapping: dict[str, str] | None = None,
 ) -> dict[str, dict[int, Any]]:
    """
    Create vectorized LIBERO environments with a consistent return shape.
@@ -449,6 +447,7 @@ def create_libero_envs(
                init_states=init_states,
                gym_kwargs=gym_kwargs,
                control_mode=control_mode,
+                camera_name_mapping=camera_name_mapping,
            )
            out[suite_name][tid] = env_cls(fns)
            print(f"Built vec env | suite={suite_name} | task_id={tid} | n_envs={n_envs}")
@@ -201,6 +201,11 @@ def rollout(
                    "You're likely using an older version of gymnasium (< 1.0). Please upgrade."
                )
            successes = final_info["is_success"].tolist()
+        elif "is_success" in info:
+            is_success = info["is_success"]
+            successes = (
+                is_success.tolist() if hasattr(is_success, "tolist") else [bool(is_success)] * env.num_envs
+            )
        else:
            successes = [False] * env.num_envs