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-11 14:49:43 +00:00 · 2026-04-08 17:48:58 +02:00
parent 4eecbad32b
commit 5de7aa5a4f
7 changed files with 361 additions and 217 deletions
@@ -0,0 +1,143 @@
+"""Tests for the benchmark dispatch refactor (create_envs / get_env_processors on EnvConfig)."""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field
+
+import gymnasium as gym
+import pytest
+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.factory import make_env, make_env_config, make_env_pre_post_processors
+
+logger = logging.getLogger(__name__)
+
+
+def test_registry_all_types():
+    """make_env_config should resolve every registered EnvConfig subclass via the registry."""
+    known = list(EnvConfig.get_known_choices().keys())
+    assert len(known) >= 6
+    for t in known:
+        cfg = make_env_config(t)
+        if not isinstance(cfg, EnvConfig):
+            continue
+        assert cfg.type == t
+
+
+def test_unknown_type():
+    with pytest.raises(ValueError, match="not registered"):
+        make_env_config("nonexistent")
+
+
+def test_identity_processors():
+    """Base class get_env_processors() returns identity pipelines."""
+    cfg = make_env_config("aloha")
+    pre, post = cfg.get_env_processors()
+    assert len(pre.steps) == 0 and len(post.steps) == 0
+
+
+def test_delegation():
+    """make_env() should call cfg.create_envs(), not use if/elif dispatch."""
+    sentinel = {"delegated": {0: "marker"}}
+    fake = type(
+        "Fake",
+        (),
+        {
+            "hub_path": None,
+            "create_envs": lambda self, n_envs, use_async_envs=False: sentinel,
+        },
+    )()
+    result = make_env(fake, n_envs=1)
+    assert result is sentinel
+
+
+def test_processors_delegation():
+    """make_env_pre_post_processors delegates to cfg.get_env_processors()."""
+    cfg = make_env_config("aloha")
+    pre, post = make_env_pre_post_processors(cfg, policy_cfg=None)
+    assert len(pre.steps) == 0
+
+
+def test_base_create_envs():
+    """Base class create_envs() should build a single-task VectorEnv via gym.make()."""
+    gym_id = "_dispatch_test/CartPole-v99"
+    if gym_id not in gym_registry:
+        register(id=gym_id, entry_point="gymnasium.envs.classic_control:CartPoleEnv")
+
+    @EnvConfig.register_subclass("_dispatch_base_test")
+    @dataclass
+    class _Env(EnvConfig):
+        task: str = "CartPole-v99"
+        fps: int = 10
+        features: dict[str, PolicyFeature] = field(default_factory=dict)
+
+        @property
+        def package_name(self):
+            return "_dispatch_test"
+
+        @property
+        def gym_id(self):
+            return gym_id
+
+        @property
+        def gym_kwargs(self):
+            return {}
+
+    try:
+        envs = _Env().create_envs(n_envs=2)
+        assert "_dispatch_base_test" in envs
+        env = envs["_dispatch_base_test"][0]
+        assert isinstance(env, gym.vector.SyncVectorEnv)
+        assert env.num_envs == 2
+        env.close()
+    finally:
+        if gym_id in gym_registry:
+            del gym_registry[gym_id]
+
+
+def test_custom_create_envs_override():
+    """A custom EnvConfig subclass can override create_envs()."""
+    mock_vec = gym.vector.SyncVectorEnv([lambda: gym.make("CartPole-v1")])
+
+    @EnvConfig.register_subclass("_dispatch_custom_test")
+    @dataclass
+    class _Env(EnvConfig):
+        task: str = "x"
+        features: dict[str, PolicyFeature] = field(default_factory=dict)
+
+        @property
+        def gym_kwargs(self):
+            return {}
+
+        def create_envs(self, n_envs, use_async_envs=False):
+            return {"custom_suite": {0: mock_vec}}
+
+    try:
+        result = make_env(_Env(), n_envs=1)
+        assert "custom_suite" in result
+    finally:
+        mock_vec.close()
+
+
+def test_custom_get_env_processors_override():
+    """A custom EnvConfig subclass can override get_env_processors()."""
+    from lerobot.processor.pipeline import DataProcessorPipeline
+
+    @EnvConfig.register_subclass("_dispatch_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 = _Env().get_env_processors()
+    assert isinstance(pre, DataProcessorPipeline)