Extension packages

External robots ship as separate Python packages that call register_wbc_extension at import time. Core stays robot-agnostic; extensions add ids, builders, data roots, and tasks.

API overview

register_wbc_extension(robot, tasks) (extension.py) performs:

1. register_robot_id — --robot resolution and data/<id>/ paths

2. register_robot_builders — make_env_cfg / make_rl_cfg for train/play

3. register_robot_motion_spec — FK scene for wbc-mjlab-data-to-npz

4. register_project_root — extension-owned data/<robot>/ tree

5. register_wbc_tasks — --task ids

6. register_robot_symmetry_config (optional) — when symmetry_config is set on the spec

WbcRobotSpec fields

Field	Purpose
robot_id	Canonical id (lowercase), e.g. h2, mybot
aliases	Extra names for --robot (unitree_h2, …)
make_env_cfg	(play=False, task_id=…) -> ManagerBasedRlEnvCfg
make_rl_cfg	RL runner cfg (PPO hyperparameters, …)
motion_spec	RobotMotionSpec — scene fn, foot debias, MJCF path for FK
project_root	Root containing data/<robot_id>/
mjcf_path	Optional shortcut when motion_spec omits mjcf_path
symmetry_config	Optional RobotSymmetryConfig for --mirror on wbc-mjlab-data-to-npz

Minimal registration

from wbc_mjlab.extension import WbcRobotSpec, register_wbc_extension
from wbc_mjlab.motion.robot_assets import RobotMotionSpec
from wbc_mjlab.tasks.config import WbcTaskConfig
from my_robot_wbc.robots.mybot.symmetry import MYBOT_SYMMETRY_CONFIG

ROOT = Path(__file__).resolve().parents[2]

def mybot_wbc_env_cfg() -> ManagerBasedRlEnvCfg:
  cfg = mybot_base_cfg()
  apply_wbc(cfg, motion_body_names=MY_MOTION_BODIES, ...)
  return cfg

register_wbc_extension(
  WbcRobotSpec(
    robot_id="mybot",
    aliases=("vendor_mybot",),
    project_root=ROOT,
    make_env_cfg=make_mybot_wbc_env_cfg,
    make_rl_cfg=mybot_wbc_rl_cfg,
    motion_spec=RobotMotionSpec(
      scene_cfg_fn=lambda: mybot_base_cfg().scene,
      mjcf_path=MYBOT_XML,
    ),
    symmetry_config=MYBOT_SYMMETRY_CONFIG,
  ),
  WbcTaskConfig(
    task_id="Wbc-Mybot",
    robot_id="mybot",
    description="Default WBC stack (apply_wbc).",
    experiment_name="wbc_mybot",
    build_env_cfg=mybot_wbc_env_cfg,
  ),
)

Package layout (recommended)

my-robot-wbc/
  pyproject.toml          # mjlab.tasks entry point
  data/<robot_id>/        # datasets (versioned or git-lfs)
  src/my_robot_wbc/
    mjlab_entry.py        # register_wbc_extension(...)
    robots/<id>/
      constants.py        # body names, XML paths
      symmetry.py         # RobotSymmetryConfig for --mirror (optional)
      base.py             # <robot>_base_cfg()
      actuators.py
      tasks.py            # WbcTaskConfig table
      rl_cfg.py

Discovery via mjlab

Add to pyproject.toml:

[project.entry-points."mjlab.tasks"]
my_robot_wbc = "my_robot_wbc.mjlab_entry"

When the package is installed, mjlab imports the entry module and registration runs before wbc-mjlab-list-envs or train/play.

Typical task stack

Most extensions expose one default task with apply_wbc only. Paper-specific variants (Zest, binary failure, SE) can be added by stacking additional presets in tasks.py — same pattern as core robots/g1/tasks.py.

Data paths

With project_root set, CLI resolves:

<project_root>/data/<robot_id>/<dataset>/

Motion conversion uses --robot <robot_id> (entity + FK), not --task.

Reference implementation

See Example: H2 extension for a walkthrough of the published H2 extension layout (wbc-mjlab-extension-h2).

Related: The robot entity, Motion data, Usage.