mjcf

Helpers for building MuJoCo models with the native mujoco.MjSpec API.

FlyGym composes its models with mujoco.MjSpec (MuJoCo's native model-editing API) rather than dm_control.mjcf (PyMJCF). MjSpec's element constructors take typed enum values where PyMJCF accepted strings, and it has no high-level actuator "shortcut" classes (position/velocity/...). This module provides the thin translation layer FlyGym needs: string-to-enum lookups and an add_actuator helper that reproduces MuJoCo's actuator shortcuts.

add_actuator(spec, kind, *, name, joint=None, body=None, tendon=None, site=None, forcelimited=None, forcerange=None, ctrllimited=None, ctrlrange=None, gear=None, kp=None, kv=None, gain=None, **kwargs)

Add an actuator to spec, reproducing MuJoCo's actuator shortcuts.

MjSpec only exposes the low-level "general" actuator (gain/bias/dyn/trn). This helper expands the familiar motor/position/velocity/adhesion/ general shortcuts into the equivalent low-level parameters, matching what the XML compiler (and PyMJCF) produced. The expansions were verified against dm_control.mjcf compiled output.

Exactly one transmission target (joint, body, tendon or site) must be given.

Parameters:

Name	Type	Description	Default
spec	MjSpec	The spec to add the actuator to.	required
kind	str	Actuator shortcut name (e.g. "position", "motor").	required
name	str	Actuator name.	required
joint	str | None	Transmission target joint name (give exactly one of joint/body/tendon/site).	None
body	str | None	Transmission target body name.	None
tendon	str | None	Transmission target tendon name.	None
site	str | None	Transmission target site name.	None
forcelimited	bool | None	Whether actuator force is clamped to forcerange.	None
forcerange	tuple[float, float] | None	Min/max actuator force.	None
ctrllimited	bool | None	Whether the control input is clamped to ctrlrange.	None
ctrlrange	tuple[float, float] | None	Min/max control input.	None
gear	float | None	Transmission gear ratio (gear[0]).	None
kp	float | None	Position/intvelocity gain.	None
kv	float | None	Position/velocity/intvelocity damping.	None
gain	float | None	Gain (gainprm[0]) for adhesion, motor, and general actuators.	None
**kwargs	Any	Extra low-level attributes set directly on the actuator (e.g. gainprm, biasprm for general).	{}

Returns:

Type	Description
MjsActuator	The created MjsActuator.

Source code in src/flygym/utils/mjcf.py

def add_actuator(
    spec: mj.MjSpec,
    kind: str,
    *,
    name: str,
    joint: str | None = None,
    body: str | None = None,
    tendon: str | None = None,
    site: str | None = None,
    forcelimited: bool | None = None,
    forcerange: tuple[float, float] | None = None,
    ctrllimited: bool | None = None,
    ctrlrange: tuple[float, float] | None = None,
    gear: float | None = None,
    kp: float | None = None,
    kv: float | None = None,
    gain: float | None = None,
    **kwargs: Any,
) -> mj.MjsActuator:
    """Add an actuator to ``spec``, reproducing MuJoCo's actuator shortcuts.

    MjSpec only exposes the low-level "general" actuator (gain/bias/dyn/trn). This
    helper expands the familiar ``motor``/``position``/``velocity``/``adhesion``/
    ``general`` shortcuts into the equivalent low-level parameters, matching what
    the XML compiler (and PyMJCF) produced. The expansions were verified against
    ``dm_control.mjcf`` compiled output.

    Exactly one transmission target (``joint``, ``body``, ``tendon`` or ``site``)
    must be given.

    Args:
        spec: The spec to add the actuator to.
        kind: Actuator shortcut name (e.g. ``"position"``, ``"motor"``).
        name: Actuator name.
        joint: Transmission target joint name (give exactly one of
            ``joint``/``body``/``tendon``/``site``).
        body: Transmission target body name.
        tendon: Transmission target tendon name.
        site: Transmission target site name.
        forcelimited: Whether actuator force is clamped to ``forcerange``.
        forcerange: Min/max actuator force.
        ctrllimited: Whether the control input is clamped to ``ctrlrange``.
        ctrlrange: Min/max control input.
        gear: Transmission gear ratio (``gear[0]``).
        kp: Position/intvelocity gain.
        kv: Position/velocity/intvelocity damping.
        gain: Gain (``gainprm[0]``) for adhesion, motor, and general actuators.
        **kwargs: Extra low-level attributes set directly on the actuator (e.g.
            ``gainprm``, ``biasprm`` for ``general``).

    Returns:
        The created ``MjsActuator``.
    """
    params: dict[str, Any] = {"name": name}

    # Transmission target.
    targets = {"joint": joint, "body": body, "tendon": tendon, "site": site}
    set_targets = {k: v for k, v in targets.items() if v is not None}
    if len(set_targets) != 1:
        raise ValueError(
            f"Exactly one transmission target required, got {list(set_targets)}."
        )
    trn, target = next(iter(set_targets.items()))
    trntype = {
        "joint": mj.mjtTrn.mjTRN_JOINT,
        "body": mj.mjtTrn.mjTRN_BODY,
        "tendon": mj.mjtTrn.mjTRN_TENDON,
        "site": mj.mjtTrn.mjTRN_SITE,
    }[trn]
    params["trntype"] = trntype
    params["target"] = target

    # Shortcut expansion (gain/bias/dyn). Verified against PyMJCF output. MjSpec
    # requires gainprm/biasprm to be length 10, so leading entries are padded.
    if kind == "motor" or kind == "general":
        gain_v = 1.0 if gain is None else gain
        params["gaintype"] = mj.mjtGain.mjGAIN_FIXED
        params["gainprm"] = _prm(gain_v)
        params["biastype"] = mj.mjtBias.mjBIAS_NONE
    elif kind == "position":
        kp_v = 1.0 if kp is None else kp
        kv_v = 0.0 if kv is None else kv
        params["gaintype"] = mj.mjtGain.mjGAIN_FIXED
        params["gainprm"] = _prm(kp_v)
        params["biastype"] = mj.mjtBias.mjBIAS_AFFINE
        params["biasprm"] = _prm(0.0, -kp_v, -kv_v)
    elif kind == "velocity":
        kv_v = 1.0 if kv is None else kv
        params["gaintype"] = mj.mjtGain.mjGAIN_FIXED
        params["gainprm"] = _prm(kv_v)
        params["biastype"] = mj.mjtBias.mjBIAS_AFFINE
        params["biasprm"] = _prm(0.0, 0.0, -kv_v)
    elif kind == "adhesion":
        gain_v = 1.0 if gain is None else gain
        params["gaintype"] = mj.mjtGain.mjGAIN_FIXED
        params["gainprm"] = _prm(gain_v)
        params["biastype"] = mj.mjtBias.mjBIAS_NONE
        if ctrlrange is None:
            ctrlrange = (0.0, 1.0)
        if ctrllimited is None:
            ctrllimited = True
    else:
        raise NotImplementedError(f"Actuator shortcut '{kind}' is not supported.")

    # Common attributes.
    if forcerange is not None:
        params["forcerange"] = list(forcerange)
    if forcelimited is not None:
        params["forcelimited"] = _limited_enum(forcelimited)
    if ctrlrange is not None:
        params["ctrlrange"] = list(ctrlrange)
    if ctrllimited is not None:
        params["ctrllimited"] = _limited_enum(ctrllimited)
    if gear is not None:
        params["gear"] = [gear, 0, 0, 0, 0, 0]

    params.update(kwargs)
    return spec.add_actuator(**params)

add_material(spec, *, texture=None, **params)

Add a material, optionally linking a texture by name in the RGB texture role (PyMJCF exposed this as a single material.texture attribute).

Source code in src/flygym/utils/mjcf.py

def add_material(
    spec: mj.MjSpec, *, texture: str | None = None, **params: Any
) -> mj.MjsMaterial:
    """Add a material, optionally linking a texture by name in the RGB texture role
    (PyMJCF exposed this as a single ``material.texture`` attribute)."""
    material = spec.add_material(**params)
    if texture is not None:
        material.textures[int(mj.mjtTextureRole.mjTEXROLE_RGB)] = texture
    return material

add_texture(spec, **params)

Add a texture, converting the type/builtin/mark string attributes (which PyMJCF accepted as strings) to MjSpec enums.

Source code in src/flygym/utils/mjcf.py

def add_texture(spec: mj.MjSpec, **params: Any) -> mj.MjsTexture:
    """Add a texture, converting the ``type``/``builtin``/``mark`` string attributes
    (which PyMJCF accepted as strings) to MjSpec enums."""
    if "type" in params:
        params["type"] = TEXTURE_TYPES[params["type"]]
    if "builtin" in params:
        params["builtin"] = _BUILTIN_TYPES[params["builtin"]]
    if "mark" in params:
        params["mark"] = _MARK_TYPES[params["mark"]]
    return spec.add_texture(**params)

set_mujoco_globals(spec, mujoco_globals_path)

Load a YAML file of global MuJoCo settings and apply them to a spec.

Handles the compiler/option/statistic/visual groups, including the cases that differ from a flat attribute set: compiler.angle maps to compiler.degree, option.flag maps to the enable/disable bitmasks, option.integrator and option.solver are string-named enums, statistic maps to spec.stat, and visual.global maps to spec.visual.global_ (global is reserved).

Parameters:

Name	Type	Description	Default
spec	MjSpec	The spec to update.	required
mujoco_globals_path	PathLike	Path to the YAML file of global parameter overrides.	required

Source code in src/flygym/utils/mjcf.py

def set_mujoco_globals(spec: mj.MjSpec, mujoco_globals_path: PathLike) -> None:
    """Load a YAML file of global MuJoCo settings and apply them to a spec.

    Handles the compiler/option/statistic/visual groups, including the cases that
    differ from a flat attribute set: ``compiler.angle`` maps to ``compiler.degree``,
    ``option.flag`` maps to the enable/disable bitmasks, ``option.integrator`` and
    ``option.solver`` are string-named enums, ``statistic`` maps to ``spec.stat``,
    and ``visual.global`` maps to ``spec.visual.global_`` (``global`` is reserved).

    Args:
        spec: The spec to update.
        mujoco_globals_path: Path to the YAML file of global parameter overrides.
    """
    with open(mujoco_globals_path) as f:
        cfg = yaml.safe_load(f)

    for group, params in cfg.items():
        if group == "compiler":
            _apply_compiler(spec.compiler, params)
        elif group == "option":
            _apply_option(spec.option, params)
        elif group == "statistic":
            _set_attrs(spec.stat, params)
        elif group == "visual":
            _apply_visual(spec.visual, params)
        elif group == "size":
            # The <size> directives (njmax, nconmax, nkey) are legacy: MuJoCo 3.x
            # allocates these dynamically, so they are no-ops and intentionally
            # ignored here.
            continue
        else:
            raise ValueError(f"Unsupported global settings group: '{group}'.")