添加新 benchmark

本指南将带你为 LeRobot 添加一个新的 simulation benchmark。请按顺序执行这些步骤，并以现有 benchmark 作为模板。

LeRobot 中的 benchmark 是一组 Gymnasium 环境，它们在标准 gym.Env 接口背后封装第三方 simulation 器（如 LIBERO 或 Meta-World）。随后 lerobot-eval CLI 以统一方式在所有 benchmark 上运行评估。

现有 benchmark 一览

在深入之前，以下是已集成的内容：

使用 src/lerobot/envs/libero.py 和 src/lerobot/envs/metaworld.py 作为参考实现。

各部分如何协同

数据流

在评估期间，数据经过四个阶段：

1. gym.Env  ──→  raw observations (numpy dicts)

2. Preprocessing  ──→  standard LeRobot keys + task description
   (preprocess_observation in envs/utils.py, env.call("task_description"))

3. Processors  ──→  env-specific then policy-specific transforms
   (env_preprocessor, policy_preprocessor)

4. Policy  ──→  select_action()  ──→  action tensor
   then reverse: policy_postprocessor → env_postprocessor → numpy action → env.step()

大多数 benchmark 只需关注阶段 1（以正确的格式生成 observation），以及可选的阶段 3（如果需要特定于环境的变换）。

环境结构

make_env() 返回向量化环境的嵌套字典：

dict[str, dict[int, gym.vector.VectorEnv]]
#    ^suite       ^task_id

单任务环境（例如 PushT）形如 {"pusht": {0: vec_env}}。 多任务 benchmark（例如 LIBERO）形如 {"libero_spatial": {0: vec0, 1: vec1, ...}, ...}。

评估如何运行

所有 benchmark 都由 lerobot-eval 以相同方式评估：

1. make_env() 构建嵌套的 {suite: {task_id: VectorEnv}} 字典。
2. eval_policy_all() 遍历每个套件和任务。
3. 对每个任务，它通过 rollout() 运行 n_episodes 次 rollout。
4. 结果按层级聚合：episode、任务、套件、总体。
5. 指标包括 pc_success（成功率）、avg_sum_reward 和 avg_max_reward。

关键点：你的环境必须在每次 step() 调用时返回 info["is_success"]。评估循环正是据此判断任务是否完成。

你的环境必须提供什么

LeRobot 不强制严格的 observation 模式。相反，它依赖所有 benchmark 都遵循的一组约定。

环境属性

你的 gym.Env 必须设置以下属性：

成功上报

你的 step() 和 reset() 必须在 info 字典中包含 "is_success"：

info = {"is_success": True}   # or False
return observation, reward, terminated, truncated, info

observation

最简单的方法是将 simulation 器的输出映射到 preprocess_observation() 已理解的标准键。在你的 gym.Env 中执行此操作（例如在 _format_raw_obs() 辅助函数中）：

如果你的 simulation 器使用不同的键名，你有两个选择：

1. 推荐：在你的 gym.Env 封装内将它们重命名为标准键。
2. 替代方案：编写一个环境处理器，在 preprocess_observation() 运行后转换 observation（见下文步骤 4）。

action

action 是 gym.spaces.Box 中的连续 numpy 数组。维度取决于你的 benchmark（LIBERO 为 7，Meta-World 为 4，等等）。policy 通过其 input_features / output_features 配置适应不同的 action 维度。

特征声明

每个 EnvConfig 子类声明两个字典，告知 policy 应预期什么：

• features — 将特征名称映射到 PolicyFeature(type, shape)（例如 action 维度、图像形状）。
• features_map — 将原始 observation 键映射到 LeRobot 约定键（例如将 "agent_pos" 映射到 "observation.state"）。

分步说明

至少需要两个文件：一个 **gym.Env 封装**和一个带有 `create_envs()` 覆盖的 **EnvConfig 子类**。其余都是可选的或属于文档。无需修改 `factory.py`。

检查清单

1. gym.Env 封装（ src/lerobot/envs/<benchmark>.py ）

创建一个封装第三方 simulation 器的 gym.Env 子类：

class MyBenchmarkEnv(gym.Env):
    metadata = {"render_modes": ["rgb_array"], "render_fps": <fps>}

    def __init__(self, task_suite, task_id, ...):
        super().__init__()
        self.task = <task_name_string>
        self.task_description = <natural_language_instruction>
        self._max_episode_steps = <max_steps>
        self.observation_space = spaces.Dict({...})
        self.action_space = spaces.Box(low=..., high=..., shape=(...,), dtype=np.float32)

    def reset(self, seed=None, **kwargs):
        ...  # return (observation, info) — info must contain {"is_success": False}

    def step(self, action: np.ndarray):
        ...  # return (obs, reward, terminated, truncated, info) — info must contain {"is_success": <bool>}

    def render(self):
        ...  # return RGB image as numpy array

    def close(self):
        ...

基于 GPU 的 simulation 器（例如使用 EGL 渲染的 MuJoCo）：如果你的 simulation 器在 __init__ 期间分配 GPU/EGL 上下文，请将分配推迟到首次 reset()/step() 时调用的 _ensure_env() 辅助函数中。这可以避免在 AsyncVectorEnv 生成 worker 进程时继承过期的 GPU 句柄。参见 LiberoEnv._ensure_env() 了解此模式。

还要提供一个返回嵌套字典结构的工厂函数：

def create_mybenchmark_envs(
    task: str,
    n_envs: int,
    gym_kwargs: dict | None = None,
    env_cls: type | None = None,
) -> dict[str, dict[int, Any]]:
    """Create {suite_name: {task_id: VectorEnv}} for MyBenchmark."""
    ...

参见 create_libero_envs()（多套件、多任务）和 create_metaworld_envs()（按难度分组的任务）作为参考。

2. 配置（ src/lerobot/envs/configs.py ）

注册一个配置 dataclass，让用户可以用 --env.type=<name> 选择你的 benchmark。每个配置通过两个方法拥有其环境创建和处理器逻辑：

• create_envs(n_envs, use_async_envs) — 返回 {suite: {task_id: VectorEnv}}。基类默认对单任务环境使用 gym.make()。多任务 benchmark 会覆盖此方法。
• get_env_processors() — 返回 (preprocessor, postprocessor)。基类默认返回恒等（空操作）流水线。如果你的 benchmark 需要 observation/action 变换，请覆盖此方法。

@EnvConfig.register_subclass("<benchmark_name>")
@dataclass
class MyBenchmarkEnvConfig(EnvConfig):
    task: str = "<default_task>"
    fps: int = <fps>
    obs_type: str = "pixels_agent_pos"

    features: dict[str, PolicyFeature] = field(default_factory=lambda: {
        ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(<action_dim>,)),
    })
    features_map: dict[str, str] = field(default_factory=lambda: {
        ACTION: ACTION,
        "agent_pos": OBS_STATE,
        "pixels": OBS_IMAGE,
    })

    def __post_init__(self):
        ...  # populate features based on obs_type

    @property
    def gym_kwargs(self) -> dict:
        return {"obs_type": self.obs_type, "render_mode": self.render_mode}

    def create_envs(self, n_envs: int, use_async_envs: bool = True):
        """Override for multi-task benchmarks or custom env creation."""
        from lerobot.envs.<benchmark> import create_<benchmark>_envs
        return create_<benchmark>_envs(task=self.task, n_envs=n_envs, ...)

    def get_env_processors(self):
        """Override if your benchmark needs observation/action transforms."""
        from lerobot.processor import PolicyProcessorPipeline
        from lerobot.processor.env_processor import MyBenchmarkProcessorStep
        return (
            PolicyProcessorPipeline(steps=[MyBenchmarkProcessorStep()]),
            PolicyProcessorPipeline(steps=[]),
        )

要点：

• register_subclass 名称是用户在 CLI 上传递的名称（--env.type=<name>）。
• features 告知 policy 环境产生什么。
• features_map 将原始 observation 键映射到 LeRobot 约定键。
• 无需修改 factory.py — 工厂会自动委托给 cfg.create_envs() 和 cfg.get_env_processors()。

3. Env 处理器（可选 — src/lerobot/processor/env_processor.py ）

仅当你的 benchmark 需要超出 preprocess_observation() 所处理的 observation 变换（例如图像翻转、坐标转换）时才需要。在此定义处理器步骤，并在你的配置中从 get_env_processors() 返回它（见步骤 2）：

@dataclass
@ProcessorStepRegistry.register(name="<benchmark>_processor")
class MyBenchmarkProcessorStep(ObservationProcessorStep):
    def _process_observation(self, observation):
        processed = observation.copy()
        # your transforms here
        return processed

    def transform_features(self, features):
        return features  # update if shapes change

    def observation(self, observation):
        return self._process_observation(observation)

参见 LiberoProcessorStep 获取完整示例（图像旋转、四元数转轴角）。

4. 依赖（ pyproject.toml ）

添加一个新的可选依赖组：

mybenchmark = ["my-benchmark-pkg==1.2.3", "lerobot[scipy-dep]"]

固定版本规则：

• 始终固定benchmark 包到精确版本以保证可复现性（例如 metaworld==3.0.0）。
• 在需要时添加平台标记（例如 ; sys_platform == 'linux'）。
• 固定脆弱的传递依赖（如果已知）（例如 Meta-World 的 gymnasium==1.1.0）。
• 记录约束在你的 benchmark 文档页面中。

用户使用以下命令安装：

pip install -e ".[mybenchmark]"

5. 文档（ docs/source/<benchmark>.mdx ）

按照下一节的模板编写面向用户的页面。参见 docs/source/libero.mdx 和 docs/source/metaworld.mdx 获取完整示例。

6. 目录（ docs/source/_toctree.yml ）

将你的 benchmark 添加到“Benchmarks”部分：

- sections:
    - local: libero
      title: LIBERO
    - local: metaworld
      title: Meta-World
    - local: envhub_isaaclab_arena
      title: NVIDIA IsaacLab Arena Environments
    - local: <your_benchmark>
      title: <Your Benchmark Name>
  title: "Benchmarks"

验证你的集成

完成上述步骤后，确认一切正常：

1. 安装 — pip install -e ".[mybenchmark]" 并验证依赖组能干净安装。
2. 环境创建冒烟测试 — 在 Python 中使用你的配置调用 make_env()，检查返回的字典具有预期的 {suite: {task_id: VectorEnv}} 形状，并且 reset() 返回的 observation 带有正确的键。
3. 运行完整评估 — lerobot-eval --env.type=<name> --env.task=<task> --eval.n_episodes=1 --policy.path=<any_compatible_policy> 以端到端地运行完整流水线。（batch_size 默认根据 CPU 核心数自动调整；传入 --eval.batch_size=1 以强制使用单个环境。）
4. 检查成功检测 — 验证任务实际完成时 info["is_success"] 会翻转为 True。评估循环正是据此计算成功率。

编写 benchmark 文档页面

每个 benchmark 的 .mdx 页面应包含：

• 标题和描述 — 用 1-2 段说明该 benchmark 测什么以及为什么重要。
• 链接 — 论文、GitHub 仓库、项目网站（如果有）。
• 概览图片或 GIF。
• 可用任务 — 任务套件表格，包含数量和简要描述。
• 安装 — pip install -e ".[<benchmark>]" 以及任何额外步骤（环境变量、系统包）。
• 评估 — 推荐的 lerobot-eval 命令，配合 n_episodes 以获得可复现的结果。batch_size 默认为自动；仅在需要时指定。如果适用，包含单任务和多任务示例。
• policy 输入和输出 — 带形状的 observation 键、action space 描述。
• 推荐的评估 episode 数 — 每个任务多少个 episode 是标准做法。
• 训练 — 示例 lerobot-train 命令。
• 复现已发表结果 — 链接到预训练模型、评估命令、结果表格（如果有）。

参见 docs/source/libero.mdx 和 docs/source/metaworld.mdx 获取完整示例。