Extension API

Extension API#

Public entry points for external robot packages. Register a robot and its tasks once at import time; train / play / data CLIs then discover them like in-tree G1.

How-to guides: Extensions. Reference extension: wbc-mjlab-extension-h2.

Minimal registration#

Call once from your package entry point (mjlab.tasks):

from wbc_mjlab.extension import WbcRobotSpec, register_wbc_extension
from wbc_mjlab.motion.robot_assets import RobotMotionSpec
from wbc_mjlab.tasks.config import WbcTaskConfig

register_wbc_extension(
  WbcRobotSpec(
    robot_id="my_robot",
    make_env_cfg=make_my_robot_wbc_env_cfg,
    make_rl_cfg=my_robot_wbc_rl_cfg,
    project_root=PROJECT_ROOT,
    motion_spec=RobotMotionSpec(
      scene_cfg_fn=lambda: make_my_robot_wbc_env_cfg().scene,
      # foot_body_names=..., foot_sole_z=..., mjcf_path=...
    ),
  ),
  WbcTaskConfig(
    task_id="Wbc-MyRobot",
    robot_id="my_robot",
    description="Default WBC stack",
    experiment_name="wbc_my_robot",
    build_env_cfg=make_my_robot_wbc_env_cfg,
  ),
)

Then install the extension editable alongside wbc-mjlab (see Installation) and confirm wbc-mjlab-list-envs lists Wbc-MyRobot. Optional symmetry_config on WbcRobotSpec enables wbc-mjlab-data-to-npz --mirror (see Motion data and The robot entity). Full walkthrough: Example: H2 extension.

Registration#

Public API for external WBC robot + task extensions.

class wbc_mjlab.extension.WbcRobotSpec[source]#

Bases: object

Robot wiring registered into wbc-mjlab from an external package.

robot_id: str#

Canonical robot id (e.g. "h2").

make_env_cfg: Callable[[...], Any]#

Builder returning a ManagerBasedRlEnvCfg.

make_rl_cfg: Callable[[], Any]#

Builder returning the RL runner config.

motion_spec: RobotMotionSpec | None = None#

Optional FK / debias metadata for wbc-mjlab-data-to-npz.

symmetry_config: RobotSymmetryConfig | None = None#

Optional RobotSymmetryConfig for --mirror.

aliases: tuple[str, ...] = ()#

Alternate ids accepted by CLIs.

project_root: Path | str | None = None#

Extension package root (for data/<robot_id>/ resolution).

mjcf_path: Path | str | None = None#

Optional MJCF path when not set on motion_spec.

wbc_mjlab.extension.register_robot(spec: WbcRobotSpec) → None[source]#

Register robot id, env/RL builders, and optional motion-conversion metadata.

Prefer register_wbc_extension() when also registering tasks.

wbc_mjlab.extension.register_wbc_extension(robot: WbcRobotSpec, tasks: WbcTaskConfig | tuple[WbcTaskConfig, ...]) → None[source]#

Register a robot and its WBC task table in one call.

Call at extension import time (e.g. from the package __init__ or the mjlab.tasks entry point) so wbc-mjlab-list-envs / train / play discover the robot.

Task config#

class wbc_mjlab.tasks.config.WbcTaskConfig[source]#

Bases: object

One registered mjlab task: metadata + env builder.

task_id: str#

CLI / registry id (e.g. "Wbc-G1", "Wbc-H2-Zest").

robot_id: str#

Robot this task binds to (must be registered).

description: str#

Short human-readable summary for listings.

experiment_name: str#

Log / experiment directory name.

build_env_cfg: Callable[[], Any]#

Zero-arg builder returning the env config.

Motion conversion metadata#

Used by wbc-mjlab-data-to-npz --robot <id> when the extension provides FK / debias metadata.

class wbc_mjlab.motion.robot_assets.RobotMotionSpec[source]#

Bases: object

Scene + optional foot metadata for motion NPZ conversion.

scene_cfg_fn: Callable[[], mjlab.scene.SceneCfg]#

Returns the robot scene used for FK during conversion.

actuated_joint_names: tuple[str, ...] | None = None#

Optional actuated DoF subset.

foot_body_names: tuple[str, ...] | None = None#

Feet used for ground-contact z debias.

foot_sole_z: float | None = None#

Sole offset (m) for z debias.

mjcf_path: Path | None = None#

Optional path to the robot MJCF.