diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst index a940a8871b..a633b39679 100644 --- a/docs/source/changelog.rst +++ b/docs/source/changelog.rst @@ -8,6 +8,8 @@ Upcoming version (not yet released) Added ^^^^^ +- Added ``lock_joints`` spec wrapper to permanently weld specified joints, + removing them from the compiled model and action space. - Added ``BuiltinDcMotorActuator``, a native MuJoCo ```` wrapper. Supports voltage / position / velocity input modes with back-EMF, configurable motor constants, and optional integral, slew, inductance, diff --git a/docs/source/entity/index.rst b/docs/source/entity/index.rst index 4e9b714ce7..902578b33f 100644 --- a/docs/source/entity/index.rst +++ b/docs/source/entity/index.rst @@ -132,6 +132,24 @@ Because ``spec_fn`` is an arbitrary callable, you can perform any before returning: add bodies, change joint limits, swap materials, or build the entire model programmatically without an XML file at all. +Locking joints +"""""""""""""" + +Use :func:`~mjlab.utils.spec.lock_joints` to permanently weld specific +joints. It wraps a ``spec_fn`` and deletes the listed joints before +compilation: + +.. code-block:: python + + from mjlab.utils.spec import lock_joints + + locked_spec_fn = lock_joints(get_spec, ["wrist_roll", "wrist_pitch"]) + + robot_cfg = EntityCfg(spec_fn=locked_spec_fn, ...) + +The deleted joints are frozen at their default position and removed from +the action space entirely. + ``init_state`` ^^^^^^^^^^^^^^ diff --git a/src/mjlab/utils/spec.py b/src/mjlab/utils/spec.py index 033e36230a..3512418d79 100644 --- a/src/mjlab/utils/spec.py +++ b/src/mjlab/utils/spec.py @@ -5,6 +5,7 @@ import shutil import xml.etree.ElementTree as ET import zipfile +from collections.abc import Sequence from pathlib import Path from typing import Callable @@ -200,6 +201,21 @@ def wrapper() -> mujoco.MjSpec: return wrapper +def lock_joints( + spec_fn: Callable[[], mujoco.MjSpec], + joint_names: Sequence[str], +) -> Callable[[], mujoco.MjSpec]: + """Wrap spec_fn to delete joints before compilation.""" + + def wrapper() -> mujoco.MjSpec: + spec = spec_fn() + for name in joint_names: + spec.delete(spec.joint(name)) + return spec + + return wrapper + + def get_non_free_joints(spec: mujoco.MjSpec) -> tuple[mujoco.MjsJoint, ...]: """Returns all joints except the free joint.""" joints: list[mujoco.MjsJoint] = [] diff --git a/tests/test_spec_utils.py b/tests/test_spec_utils.py index 83dc4d438a..57622f3f19 100644 --- a/tests/test_spec_utils.py +++ b/tests/test_spec_utils.py @@ -4,7 +4,11 @@ import numpy as np import pytest -from mjlab.utils.spec import create_position_actuator, create_velocity_actuator +from mjlab.utils.spec import ( + create_position_actuator, + create_velocity_actuator, + lock_joints, +) @pytest.fixture @@ -212,3 +216,57 @@ def test_velocity_actuator_forces_clipped_to_effort_limit( data.ctrl[0] = 100.0 mujoco.mj_forward(model, data) assert abs(data.actuator_force[0]) <= effort_limit + 1e-6 + + +# --------------------------------------------------------------------------- +# lock_joints +# --------------------------------------------------------------------------- + + +@pytest.fixture +def two_joint_spec(): + """Spec with two hinge joints on a serial chain.""" + spec = mujoco.MjSpec() + b1 = spec.worldbody.add_body(name="link1") + b1.add_joint(name="j1", type=mujoco.mjtJoint.mjJNT_HINGE, axis=[0, 0, 1]) + b1.add_geom(type=mujoco.mjtGeom.mjGEOM_BOX, size=[0.1, 0.1, 0.1], mass=1.0) + b2 = b1.add_body(name="link2") + b2.add_joint(name="j2", type=mujoco.mjtJoint.mjJNT_HINGE, axis=[0, 0, 1]) + b2.add_geom(type=mujoco.mjtGeom.mjGEOM_BOX, size=[0.1, 0.1, 0.1], mass=1.0) + return spec + + +def test_lock_joints_removes_specified_joints(two_joint_spec): + """Locking a joint removes it from the compiled model.""" + spec_fn = lock_joints(lambda: two_joint_spec, ["j1"]) + spec = spec_fn() + + joint_names = [j.name for j in spec.joints] + assert "j1" not in joint_names + assert "j2" in joint_names + + model = spec.compile() + assert model.njnt == 1 + + +def test_lock_joints_removes_multiple_joints(two_joint_spec): + """Locking all joints produces a fully welded model.""" + spec_fn = lock_joints(lambda: two_joint_spec, ["j1", "j2"]) + spec = spec_fn() + + assert len(list(spec.joints)) == 0 + model = spec.compile() + assert model.njnt == 0 + assert model.nq == 0 + + +def test_lock_joints_missing_joint_raises(): + """Requesting a non-existent joint raises an error.""" + spec = mujoco.MjSpec() + body = spec.worldbody.add_body(name="body") + body.add_joint(name="j1", type=mujoco.mjtJoint.mjJNT_HINGE, axis=[0, 0, 1]) + body.add_geom(type=mujoco.mjtGeom.mjGEOM_BOX, size=[0.1, 0.1, 0.1], mass=1.0) + + spec_fn = lock_joints(lambda: spec, ["nonexistent"]) + with pytest.raises(Exception): + spec_fn()