Skip to content

mjlab v1.4.0

Latest

Choose a tag to compare

@kevinzakka kevinzakka released this 26 May 22:52
· 18 commits to main since this release
Immutable release. Only release title and notes can be modified.

mjlab 1.4.0 adds per-world mesh variants, a native implicit PD actuator, sensor-frame projected gravity, and stricter shape validation in reward, termination, and metrics managers. It also updates MuJoCo and mujoco-warp to 3.8 and fixes several issues around actuator delay, multi-node seeding, contact sensor frames, rough-terrain videos, and tracking metrics.

MuJoCo 3.8

mjlab now uses MuJoCo 3.8 and requires mujoco-warp 3.8.0.3 or newer within the 3.8 series. The multiccd enable flag was removed upstream because it is now always enabled. If your config includes "multiccd" in MujocoCfg.enableflags, remove it.

Per-World Mesh Variants

Batched simulations can now use different mesh assets for the same logical entity in different worlds. Define a VariantEntityCfg with named spec callables, then choose how worlds are assigned to variants:

VariantEntityCfg(
    variants={
        "cube": make_cube_spec,
        "sphere": make_sphere_spec,
        "bunny": make_bunny_spec,
    },
    assignment={"cube": 0.5, "sphere": 0.3, "bunny": 0.2},
)

assignment can be:

  • None, for a uniform assignment across variants
  • a dict[str, float], for weighted assignment
  • a Callable[[int], Sequence[int]], for a custom world-to-variant mapping

Variants must have the same kinematic structure: the same bodies, joints, and joint types. Mesh geoms may differ. Mesh-derived constants such as collision bounds, body inertials, subtree mass, and inverse weights are compiled per-variant and stored as per-world arrays in the Warp model. Domain randomization, the native viewer, the offscreen renderer, and the Viser viewer all use the assigned variant automatically.

Per-variant materials and textures are supported. Each variant can reference its own named material, which is prefixed and scattered through geom_matid alongside geom_dataid. Variant assignment is fixed at simulation init.

See the heterogeneous worlds docs for usage. Contributions by @XiangruiJiang and @omarrayyann.

levered_grasp.mp4

A YAM-adapted lift task trained with per-world mesh variants. The objects vary in shape and size, and one flat rectangular object leads the policy to use a pivot grasp strategy.

BuiltinPdActuator

BuiltinPdActuator provides native MuJoCo PD control with both position and velocity targets.

The existing actuator options covered different parts of this behavior: BuiltinPositionActuator used MuJoCo's native PD path, so implicit and implicitfast included the gain derivatives in the velocity update, but it only supported a zero velocity target. IdealPdActuator supported a non-zero velocity target, but computed torque in Python and applied it through a <motor>, which MuJoCo treats as an external force and hence is less stable.

BuiltinPdActuator is implemented as paired <position> and <velocity> actuators:

kp * (p_target - q) + kd * (v_target - qdot)

effort_limit is enforced as a clamp on the summed PD force through jnt_actfrcrange or tendon_actfrcrange. dr.pd_gains and dr.effort_limits both support the new actuator.

Sensor-Frame Projected Gravity

mdp.projected_gravity is computed from the root body orientation. That means it does not reflect IMU mounting randomization applied to the IMU site, even though gyro and linear acceleration observations from that site do.

mdp.projected_gravity_from_sensor reads from a framezaxis up-vector sensor attached to the IMU site and negates it. Use it when projected gravity should move with IMU site randomization such as dr.site_quat. Go1 and G1 now include an imu_upvector sensor for this observation.

Manager Shape Validation

RewardManager, TerminationManager, and MetricsManager now validate term outputs at compute time. Each term must return a tensor with shape (num_envs,). If a term returns the wrong shape, mjlab raises a ValueError that names the offending term instead of allowing a silent broadcast or a later training-time failure.

Fixes

  • Actuator command delay now applies to all targets. IdealPdActuator and DcMotorActuator use velocity and effort targets as well as position, and those targets were previously not delayed. Zero-reference setups are unaffected.
  • Multi-node training seeds now use global RANK instead of LOCAL_RANK, so seeds do not repeat across nodes. Contribution by @bd-pdomanico.
  • Reward entries written to extras["log"] are no longer dropped on reset. The log dict is cleared at the top of step() and reset(), before reward computation writes into it.
  • ContactSensor(global_frame=True) now rotates contact-frame values with columns ordered [normal, tangent, tangent2]. The previous order projected normal force onto a tangent axis. Contribution by @bd-pdomanico.
  • Rough-terrain offscreen videos no longer flicker when the terrain curriculum mutates env_origins on reset. The renderer now computes the neighboring robot set once and reuses it.
  • apply_body_impulse no longer fires on the first step or immediately after every reset. Cooldown is sampled lazily on first use.

Other Changes

  • ContactSensor.primary_names exposes resolved primary names in per-contact tensor order, so contact-data columns can be mapped back to primaries.
  • DebugVisualizer.add_box draws axis-oriented boxes in the native and Viser viewers.
  • train, play, and evaluate now accept --log-root for routing logs to a scratch disk or shared mount. Contribution by @louislelay.
  • CameraSensorData.segmentation now stores typed (object_id, object_type) pairs with shape [B, H, W, 2], matching mujoco-warp's output. Contribution by @tkelestemur.
  • RayCaster post-processing avoids CUDA syncs by replacing boolean-mask indexing with masked_fill_ and a clamped-distance formulation of hit_pos_w. Contribution by @bd-pdomanico.
  • Tracking evaluation now snapshots the reference before each step, raises on unknown end-effector body names, computes MPKPE against the global reference, re-anchors root-relative MPKPE to the robot root each step, and reports joint velocity error as per-joint RMS instead of an L2 norm.
  • Entity now raises during construction if its spec contains more than one freejoint.
  • Event helpers now share resolve_env_ids; push_by_setting_velocity and apply_external_force_torque accept env_ids=None for global-time interval terms.
  • TerrainEntityCfg no longer adds environment-origin, terrain-origin, and flat-patch debug sites unless debug_vis=True.
  • Task package load failures now print full tracebacks to stderr, including the entry point module path. Contribution by @saikishor.
  • rsl-rl-lib was bumped from 5.0.1 to 5.2.0 for torch.compile support in PPO and Distillation, plus optional std clamping in GaussianDistribution. No mjlab code changes are required.
  • Deprecated warp-lang symbols removed in newer releases were replaced. Contribution by @rdeits.

Breaking Changes

  • Remove "multiccd" from MujocoCfg.enableflags.
  • CameraSensorData.segmentation changed from legacy [B, H, W] geom IDs to [B, H, W, 2] typed (object_id, object_type) pairs.

Contributors

First-time contributors: @CodingCatMountain, @paLeziart, @bd-pdomanico, @tkelestemur.

Returning contributors: @rdeits, @louislelay, @omarrayyann, @saikishor, @XiangruiJiang.

Full Changelog: v1.3.0...v1.4.0