Studio presets, robots and plots

Studio presets, robots and plots#

The Studio’s files are plain YAML handled by two modules you can use from scripts too; a third turns clips and Play rollouts into plot channels. Guides: WBC Studio and Add a robot from model files.

Presets and modules#

YAML diffs on top of a registered task (~/.config/wbc-mjlab/presets, …/modules), registered as <base>--<name>.

wbc_mjlab.studio.presets.build(base: str, overrides: dict, *, play: bool = False, modules: list[dict] | list[str] = ())[source]#: Env + runner configs of base with modules (names, or resolved {name, overrides}) then overrides applied (fresh copies).

wbc_mjlab.studio.presets.apply(env_cfg, agent_cfg, overrides: dict, robot_id: str | None = None) → None[source]#: Apply preset overrides to an env config and a runner config in place (robot_id: where added observation terms come from).

wbc_mjlab.studio.presets.register(name: str, snapshot: dict | None = None) → str[source]#: Register preset name (or a run’s resolved snapshot of it) as the mjlab / WBC task <base>--<name>, re-registering when the content changed; returns the task id.

wbc_mjlab.studio.presets.resolve(name: str) → dict[source]#: Preset name with its modules inlined: all it takes to rebuild its configs (the run snapshot).

wbc_mjlab.studio.presets.save(name: str, base: str, overrides: dict, description: str = '', modules: list[str] | None = None) → dict[source]#

wbc_mjlab.studio.presets.save_module(name: str, overrides: dict, description: str = '', robot: str = '') → dict[source]#

wbc_mjlab.studio.presets.list_modules(robot: str | None = None) → list[dict][source]#

wbc_mjlab.studio.presets.describe(env_cfg, agent_cfg) → dict[source]#: Every editable knob of a task’s configs, with its current (default) value.

wbc_mjlab.studio.presets.library(robot_id: str, joints: int = 0, bodies: int = 0, env_cfg=None) → list[dict][source]#: The term library as rows for the UI (dim when known; bodies = keybody subset of the variant a preset gets when it adds the term to the actor; requires: a scene sensor env_cfg lacks, e.g. the height scan of the -Terrain-Scan tasks).

wbc_mjlab.studio.presets.observation_diff(a, b) → dict[source]#: Observation overrides turning env cfg a into b (terms removed / added).

wbc_mjlab.studio.presets.deployable(name: str) → bool[source]#: Actor terms the exporter (tracking_params.yaml) knows how to rebuild on the robot.

wbc_mjlab.user_config.config_dir() → Path[source]#

Robots without code#

Robots from a model file + a YAML spec, without code (the Studio’s Add robot).

Each robot is a folder in <config>/robots/<id>/ (wbc_mjlab.user_config.config_dir(), ~/.config/wbc-mjlab):

robot.yaml       the spec: model file, anchor / keybodies / end-effectors / feet, PD gains
                 per joint, initial pose (:class:`CustomRobotSpec`)
<model>.xml|urdf the uploaded model and its meshes, untouched
wbc_robot.xml    the model as wbc-mjlab uses it, written next to the source on save:
                 actuators, keyframes and world geoms removed (wbc-mjlab adds its own PD
                 actuators and terrain), a free joint on the root, and the sensors the WBC
                 tasks read (``imu_ang_vel`` / ``imu_lin_vel`` IMU, ``root_angmom``) when the
                 model has none

register_all() (run by the mjlab.tasks entry point, so in every CLI and the Studio) registers each robot like an extension: robot id <id>, tasks Wbc-<Id> (the default WBC stack), Wbc-<Id>-SE (state-estimation actor) and Wbc-<Id>-Terrain-Scan (curriculum terrain + height scan), data in data/<id>/.

class wbc_mjlab.robots.custom.CustomRobotSpec[source]#

Bases: object

id: str#

model: str#: Source model file, relative to the robot folder.

anchor_body: str#

keybodies: list[str]#: Bodies the motion command tracks (rewards, RSI, terminations).

ee_bodies: list[str]#: End-effectors (EE height termination, the EE task’s actor command).

feet_bodies: list[str]#

foot_sites: list[str]#

the reward is dropped).

Type:: Sites under the feet for the foot-slip reward (none

foot_geoms: list[str]#

not randomized).

Type:: Named collision geoms of the feet for friction randomization (none

foot_sole_z: float = 0.0#: Sole below the foot bodies (m), for z debias in conversion.

init_height: float = 1.0#

init_joint_pos: dict[str, float]#

joints: dict[str, JointGains]#

action_scale: float = 0.25#: Action = scale × effort / stiffness per joint (as G1 / H2).

description: str = ''#

nconmax: int = 64#

njmax: int = 400#

property folder: Path#

property source: Path#

property prepared: Path#

property task_id: str#

class wbc_mjlab.robots.custom.JointGains[source]#

Bases: object

effort: float = 100.0#: Peak torque magnitude (N·m).

stiffness: float = 0.0#

damping: float = 0.0#

armature: float = 0.01#: Reflected rotor inertia (kg·m²).

wbc_mjlab.robots.custom.prepare(path: Path) → tuple[mujoco.MjSpec, dict][source]#: The model as wbc-mjlab uses it (see the module docstring) and what was changed.

wbc_mjlab.robots.custom.inspect_model(path: Path) → dict[source]#: What the Studio’s Add robot form needs: the tree, joints with suggested PD gains, sites and a suggested wiring (all computed on the prepared model).

wbc_mjlab.robots.custom.suggest(m: mujoco.MjModel, d: mujoco.MjData, root: int) → dict[source]#: Best-guess wiring from body names and the kinematic tree.

wbc_mjlab.robots.custom.settle_test(spec: CustomRobotSpec, seconds: float = 0.5) → dict[source]#: Drop the robot in its initial pose onto a floor and hold the pose with the spec’s PD gains (CPU MuJoCo, no policy) for seconds. Plain PD doesn’t balance a humanoid for long (G1 with its own gains sags after ~1 s), so this checks the setup, not standing: stable simulation, no start inside the floor, lands on its feet, holds the pose, torque limits not saturated. Returns {"ok", "checks": [{name, ok, text, warn_only?}], …} (ok: no failed check; a warn_only check only warns).

wbc_mjlab.robots.custom.save_spec(spec: CustomRobotSpec) → CustomRobotSpec[source]#: Validate, write wbc_robot.xml + robot.yaml and (re)register the robot.

wbc_mjlab.robots.custom.load_spec(robot_id: str) → CustomRobotSpec[source]#

wbc_mjlab.robots.custom.env_cfg(spec: CustomRobotSpec, *, se: bool = False, terrain: bool = False, play: bool = False)[source]#: Wbc-<Id> (se: -SE, terrain: -Terrain-Scan — procedural curriculum terrain, a height scan under the root body and terrain-aware terminations, as G1).

wbc_mjlab.robots.custom.register(spec: CustomRobotSpec) → None[source]#: Register (or re-register after an edit) one custom robot and its tasks.

wbc_mjlab.robots.custom.register_all() → list[str][source]#: Register every custom robot (a broken one is reported and skipped).

wbc_mjlab.robots.custom.task_id_for(robot_id: str) → str[source]#: h1 → Wbc-H1, h2x → Wbc-H2X, my_bot → Wbc-My-Bot.

Plot channels#

Plot channels for the Studio’s Plots window.

Two sources, one layout ({channels: [{id, name, group, unit, lower, upper}], data, ref} with data / ref (rows, channels) float32 values):

Clip (Data step) — clip_channels(): the open clip’s root position, root roll / pitch / yaw, every joint (with its limits) and the feet’s sole heights. A converted clip previewed against its raw source gets the source as ref (resampled to the clip’s frames), drawn dashed. Sent once per clip as JSON (values little-endian, base64).
Play — PlayRecorder: a rolling record of the session’s envs while the policy runs: the robot’s root, joints, PD targets, torques and actions, the reference it tracks as ref, per-keybody tracking errors and every reward term (weighted, as summed into the reward, and raw: before the weight). Row k pairs the robot after step k with the reference that step’s rewards compared it to (the command advances after the rewards, so that is the frame set at the end of step k - 1). Recording is one vectorized write per step for every recorded env, on the sim device; the client pulls new rows as binary (encode_rows()) with every stage frame.

wbc_mjlab.studio.plots.clip_channels(clip, model, *, feet: list[str], sole_z: float, compare=None) → dict[source]#: Channels of a clip (Data step); compare (the raw source of a converted clip) is resampled to the clip’s frames and sent as ref.

class wbc_mjlab.studio.plots.PlayRecorder[source]#

Bases: object

Rolling record of a Play session, several envs at once (each a track): the robot (data) next to the reference it tracks (ref, the ref_cols channels only), plus per-row meta columns (clip, frame in the clip, episode step) the client uses to mark episode and clip changes.

A session of up to MAX_TRACKS envs (memory permitting) records every env from the start, so switching the focused env shows its history at once; a larger one records the most recently focused envs (an env starts recording when first focused or fetched). One vectorized write per step for all tracks; nothing leaves the GPU until fetched.

META = ('clip', 'clip_frame', 'episode_step')#

property num_tracks: int#

follow(env: int) → None[source]#: Keep env (the focused one) recorded.

clear() → None[source]#: Drop every record (restart): the client starts over.

record() → None[source]#: One row for every track (call after each env step).

fetch(env: int = 0, since: int = 0, epoch: int | None = None) → dict[source]#: Rows since… of env’s record as (rows, columns) float32 arrays (data, ref for the ref_cols channels, meta); all of it, with the channel list, when epoch is not the record’s (another env, a restart: the client starts over). An env not recorded yet starts recording now.

wbc_mjlab.studio.plots.encode_rows(out: dict) → bytes[source]#: A PlayRecorder.fetch() result as [u32 header length][header JSON][pad to 4] then the float32 data, ref and meta rows (little-endian); the header lists their column counts (cols).

wbc_mjlab.studio.plots.effort_limits(session) → dict[str, tuple[float, float, str]][source]#: Joint name → (lower, upper, source) torque limits of a Play session’s robot: the tightest of the compiled model’s joint actuatorfrcrange / joint-actuator forcerange (× gear) — from the asset’s MJCF —, the actuators’ constant effort_limit and the peak of a torque-speed envelope (Unitree actuators: the real cap is lower at speed). Joints with none of them are left out.

wbc_mjlab.studio.plots.joint_group(name: str) → str[source]#: “Left leg”, “Right arm”, “Torso”, “Head” or “Other” from a joint’s name.

Studio presets, robots and plots

Contents

Studio presets, robots and plots#

Presets and modules#

Robots without code#

Plot channels#