Tasks and presets

Tasks and presets#

A preset is a pure Python function apply_*(cfg) -> None that mutates an existing ManagerBasedRlEnvCfg in place. A task is a registered name for one specific stack: robot entity + preset(s) + optional wiring.

Not sure which task to run first? See Which task should I use? on the Quickstart.

Presets never add new MDP term implementations — they reweight rewards, swap RSI flags, remove observation terms, or adjust termination thresholds on slots already defined in make_base_wbc_env_cfg(). Tasks are not separate MDP forks.

How they connect#

make_base_wbc_env_cfg()          # shared MDP template (env/wbc_env_cfg.py)
         ↓
<robot>_base_cfg()               # robot entity: MJCF, body names, sensors
         ↓
apply_<preset>(cfg, ...)         # paper / deploy recipe (presets/*.py)
         ↓
[optional] apply_se_actor(cfg)   # composable obs swap
         ↓
WbcTaskConfig(                   # CLI registration (tasks/config.py)
  task_id="Wbc-…",
  build_env_cfg=<builder above>,
)

Preset = recipe (reusable across robots). Task = one named builder + log dir. The same preset can back multiple tasks (e.g. apply_wbc on g1 and h2).

Available presets (core)#

Function	Module	Typical use
`apply_wbc`	`presets/wbc.py`	Default WBC + deploy extras (Table S4 on all keybodies)
`apply_zest`	`presets/zest.py`	ZEST paper repro (Table S4 on EE bodies only, no EE-z termination)
`apply_binary_failure`	`presets/binary_failure.py`	BeyondMimic-style `binary_failure` RSI
`apply_se_actor`	`presets/se_actor.py`	Actor obs swap (compose on top of wbc/zest)

Preset → task map (in-tree example)#

Task id	Preset stack	Builder (`robots/g1/tasks.py`)
`Wbc-G1`	`apply_wbc`	`g1_wbc_env_cfg`
`Wbc-G1-SE`	`apply_wbc` → `apply_se_actor` + IMU wiring	`g1_wbc_se_env_cfg`
`Wbc-G1-Zest`	`apply_zest`	`g1_wbc_zest_env_cfg`
`Wbc-G1-Zest-SE`	`apply_zest` → `apply_se_actor` + IMU wiring	`g1_wbc_zest_se_env_cfg`
`Wbc-G1-BinaryFailure`	`apply_binary_failure`	`g1_wbc_binary_failure_env_cfg`

Extension tasks (e.g. Wbc-H2) use the same pattern — see Extension packages.

Philosophy#

Principle	Meaning
Shared MDP	One `env/`; presets reweight cfg
Tasks, not forks	Paper differences = `presets/*.py` + task builders
Robot entity separate	Each platform registered via The robot entity
Neutral code	RSI strategies named generically; cite papers from Research & citations

Quick start (example)#

uv run wbc-mjlab-data-to-npz --robot g1 --dataset samples --batch-size 8
uv run wbc-mjlab-list-envs
uv run wbc-mjlab-train --task Wbc-G1 --dataset samples

Paper ↔ task map (in-tree example)#

Paper	Link	Example task ids
ZEST	arXiv:2602.00401	`Wbc-G1-Zest`, `Wbc-G1`, SE variants
BeyondMimic	arXiv:2508.08241	`Wbc-G1-BinaryFailure`
SONIC	arXiv:2511.07820	Planned — Roadmap

Preset comparison (in-tree example)#

	`Wbc-G1`	`Wbc-G1-Zest`	`Wbc-G1-BinaryFailure`
Preset	`apply_wbc`	`apply_zest`	`apply_binary_failure`
Body rewards	all keybodies	EE bodies only	all keybodies
RSI	reward-aligned `similarity_ema`	reward-aligned `similarity_ema`	`binary_failure`
`ee_body_pos` term	yes	removed	yes (base)

Example: `apply_zest` (preset code)#

presets/zest.py — mutates rewards, RSI, observations, and terminations on the cfg object passed in. Robot-specific body tuples come from arguments (supplied by the task builder):

def apply_zest(
  cfg: ManagerBasedRlEnvCfg,
  *,
  reward_body_names: tuple[str, ...],
  contact_body_names: tuple[str, ...],
  force_threshold: float = 2000.0,
) -> None:
  rw = cfg.rewards

  # Table S4 tracking kernels on selected bodies
  rw["motion_body_pos"].weight = 1.0
  rw["motion_body_pos"].params["sigma_per_keybody"] = 0.2
  rw["motion_body_pos"].params["body_names"] = reward_body_names
  rw["motion_body_ori"].params["sigma_per_keybody"] = 0.4
  rw["motion_body_ori"].params["body_names"] = reward_body_names
  # ... root + joint terms, kappa=0.25 ...

  # Disable extras not in ZEST Table S4
  rw["motion_body_lin_vel"].weight = 0.0
  rw["motion_joint_vel"].weight = 0.0
  rw["foot_slip"].weight = 0.0

  # Reward-aligned RSI
  motion_cmd = cfg.commands["motion"]
  motion_cmd.rsi = replace(
    motion_cmd.rsi,
    similarity_from_rewards=True,
    bin_width_s=4.0,
    similarity_norm_by_remaining_clip=True,
    min_bin_span_ratio=0.5,
    persist_failure_levels=True,
  )

  # Actor: drop ref_joint_vel
  cfg.observations["actor"].terms.pop("ref_joint_vel", None)

  # Terminations: ZEST style — no EE height cutoff
  cfg.terminations.pop("ee_body_pos", None)
  cfg.terminations["anchor_pos"].params["threshold"] = 0.35
  cfg.terminations["keybody_ground_contact_force"] = TerminationTermCfg(
    func=mdp.excessive_keybody_ground_contact_force,
    params={"body_names": contact_body_names, "force_threshold": force_threshold, ...},
  )

Full source: src/wbc_mjlab/presets/zest.py.

Example: task builder + registration#

def g1_wbc_zest_env_cfg() -> ManagerBasedRlEnvCfg:
  cfg = g1_base_cfg()
  apply_zest(
    cfg,
    reward_body_names=G1_ENDEFFECTOR_BODY_NAMES,
    contact_body_names=G1_MOTION_BODY_NAMES,
  )
  return cfg

G1_WBC_TASKS = (
  WbcTaskConfig(
    task_id="Wbc-G1-Zest",
    robot_id="g1",
    experiment_name="wbc_g1_zest",
    build_env_cfg=g1_wbc_zest_env_cfg,
  ),
  # ...
)

Training uses the task id, not the preset name:

uv run wbc-mjlab-train --task Wbc-G1-Zest --dataset samples

Composing presets#

def g1_wbc_zest_se_env_cfg() -> ManagerBasedRlEnvCfg:
  cfg = g1_wbc_zest_env_cfg()
  apply_se_actor(cfg)
  wire_g1_imu_sensors(cfg)
  return cfg

Order matters: base cfg → primary paper preset → optional obs/IMU tweaks.

When to add a preset vs a task#

Change	Where
New reward/RSI implementation	`env/mdp/` (shared; needs review)
New paper weights / RSI flags / term toggles	New `presets/<method>.py`
New `--task` id for a robot	Task builder + `WbcTaskConfig` in `robots/<id>/tasks.py` or extension
Same preset, different robot	New task builder with that robot’s body name tuples — preset unchanged

Task guides (in-tree)#

Task	Logs	Guide
`Wbc-G1`	`logs/rsl_rl/wbc_g1/`	Default WBC stack
`Wbc-G1-SE`	`logs/rsl_rl/wbc_g1_se/`	State-estimation observations
`Wbc-G1-Zest`	`logs/rsl_rl/wbc_g1_zest/`	ZEST reproduction
`Wbc-G1-Zest-SE`	`logs/rsl_rl/wbc_g1_zest_se/`	ZEST + state-estimation
`Wbc-G1-BinaryFailure`	`logs/rsl_rl/wbc_g1_binary/`	BeyondMimic-style RSI