Index

def get_visuals(self, state: MjState) -> list[ArrowConfig]:
    """Returns a list of arrow configurations for the renderer."""
    return []

def resolve_ids(self, state: MjState) -> None:
    """Caches the actuator ID from the compiled MuJoCo model."""
    self._aid = self.actuator.get_id(state.model)

def apply_load(self, state: MjState) -> None:
    """Evaluates `control_func` and writes the result into `mjData.ctrl[self._aid]`."""
    if not self.active:
        self._last_ctrl = 0.0
        return

    self._last_ctrl = float(self.control_func(self.user_data, state))
    state.data.ctrl[self._aid] = self._last_ctrl

def request(
    self,
    signal_manager: SignalManager | None = None,
    metadata: dict[str, dict[str, Any]] | None = None,
) -> None:
    """
    Registers the applied control value for logging under `Loads/<name>:ctrl`.

    `ctrl`'s units depend on the driven actuator's transmission and `gear`/`dyntype` (the same ambiguity as `ActuatorBase.request()`'s `ctrl` channel), so no built-in metadata default is applied. Supply `metadata={"ctrl": {...}}` yourself if you know it.

    Args:
        signal_manager (SignalManager): Manager to register the sampler with. If omitted, the `SignalManager` of the active `RuntimeManager` `with` block is used. If that `RuntimeManager` has no `SignalManager` configured, this is a no-op.
        metadata: Metadata for the `ctrl` channel.

    """
    from mujoco_mojo.runtime.signal_manager import resolve_signal_manager

    signal_manager = resolve_signal_manager(signal_manager)
    if signal_manager is None:
        return

    def sample(state: MjState) -> None:
        signal_manager.post(
            value=self._last_ctrl if self.active else 0.0,
            category=SignalCategory.LOADS,
            subgroups=(self.name,),
            attr="ctrl",
            metadata=merge_signal_metadata(None, "ctrl", metadata),
        )

    signal_manager.register_sampler(sample)

@classmethod
def constant(
    cls,
    name: str,
    actuator: ActuatorBase,
    value: float | NamedValue[float],
) -> Self:
    """
    Drives the actuator with a fixed set point.

    Args:
        name (str): Load name used for telemetry column labeling.
        actuator (ActuatorBase): The MJCF actuator to drive.
        value (float | NamedValue[float]): Control value written to `mjData.ctrl` every timestep. Accepts `NamedValue[float]` for runtime mutation.

    """

    def func(ud: UserData | None, state: MjState) -> float:
        return value.value if isinstance(value, NamedValue) else value

    return cls(name=name, actuator=actuator, control_func=func)

def get_visuals(self, state: MjState) -> list[ArrowConfig]:
    """Returns a list of arrow configurations for the renderer."""
    if not self.active:
        return []

    visuals: list[ArrowConfig] = []
    action_pos = self.action_site.rt_pos(state)

    if self._last_f[3] > 1e-4 and self._vis.action_force:
        visuals.append(
            ArrowConfig(
                pos=action_pos,
                vec=self._last_f[:3],
                color=Color[self._vis.action_force].rgba,
                is_torque=False,
                length_scale=self.force_length_scale
                if self.force_length_scale is not None
                else self._vis.force_length_scale,
                width_scale=self.force_width_scale
                if self.force_width_scale is not None
                else self._vis.force_width_scale,
            )
        )

    if self._last_t[3] > 1e-4 and self._vis.torque:
        visuals.append(
            ArrowConfig(
                pos=action_pos,
                vec=self._last_t[:3],
                color=Color[self._vis.torque].rgba,
                is_torque=True,
                length_scale=self.torque_length_scale
                if self.torque_length_scale is not None
                else self._vis.torque_length_scale,
                width_scale=self.torque_width_scale
                if self.torque_width_scale is not None
                else self._vis.torque_width_scale,
            )
        )

    return visuals

@abstractmethod
def calculate(self, state: MjState) -> tuple[np.ndarray, np.ndarray]:
    """
    Calculate the force for the timestep.

    Args:
        state: The paired MuJoCo model and data instance.

    Returns:
        tuple[np.ndarray, np.ndarray]: The force and toque vector output.

    """

def request(
    self,
    signal_manager: SignalManager | None = None,
    channels: list[Literal["force", "torque"]]
    | dict[Literal["force", "torque"], dict[str, Any] | None] = ["force", "torque"],
):
    """
    Registers specific channels for logging.

    | Channel  | Description                        | Type |
    |:---------|:-----------------------------------|:-----|
    | `force`  | applied force in the world frame   | xyzm |
    | `torque` | applied torque in the world frame  | xyzm |

    Each channel is posted under `subgroups=(load_name, channel)`.

    * An `xyzm` is a cartesian vector, posted as 4 values (`x`, `y`, `z`, and its magnitude `m`).

    Each signal is tagged with built-in `dimension` metadata for its channel (`force` as force, `torque` as torque).

    If `signal_manager` is omitted, the `SignalManager` of the active `RuntimeManager` `with` block is used. If that `RuntimeManager` has no `SignalManager` configured, this is a no-op.

    Args:
        signal_manager: The signal manager to register the sampler with.
        channels: The load data channels to log. Pass a list to select channels, or a dict mapping channel name to metadata overrides (or `None`) to select channels and attach per-channel metadata in one step.

    """
    from mujoco_mojo.runtime.signal_manager import resolve_signal_manager

    signal_manager = resolve_signal_manager(signal_manager)
    if signal_manager is None:
        return

    if isinstance(channels, dict):
        _meta = cast("dict[str, dict[str, Any] | None]", channels)
        channels = list(channels.keys())
    else:
        _meta = {}

    channel_metadata = {
        "force": dim(Dimension.FORCE),
        "torque": torque_metadata(),
    }

    def sample(state: MjState):
        for channel in channels:
            source = self._last_f if channel == "force" else self._last_t
            meta = merge_signal_metadata(
                channel_metadata.get(channel), channel, _meta, unit_system=state.us
            )

            # iterate through x, y, z, and magnitude (pop. pop.)
            for i, attr in enumerate("xyzm"):
                signal_manager.post(
                    value=float(source[i]) if self.active else 0.0,
                    category=SignalCategory.LOADS,
                    # nest the force/torque under the function name
                    subgroups=(f"{self.name}", channel),
                    attr=attr,
                    metadata=meta,
                )

    signal_manager.register_sampler(sample)

def get_visuals(self, state: MjState) -> list[ArrowConfig]:
    """Returns a list of arrow configurations for the renderer."""
    if not self.active:
        return []

    visuals: list[ArrowConfig] = []
    action_pos = self.action_site.rt_pos(state)

    if self._last_f[3] > 1e-4 and self._vis.action_force:
        visuals.append(
            ArrowConfig(
                pos=action_pos,
                vec=self._last_f[:3],
                color=Color[self._vis.action_force].rgba,
                is_torque=False,
                length_scale=self.force_length_scale
                if self.force_length_scale is not None
                else self._vis.force_length_scale,
                width_scale=self.force_width_scale
                if self.force_width_scale is not None
                else self._vis.force_width_scale,
            )
        )

    if self._last_t[3] > 1e-4 and self._vis.torque:
        visuals.append(
            ArrowConfig(
                pos=action_pos,
                vec=self._last_t[:3],
                color=Color[self._vis.torque].rgba,
                is_torque=True,
                length_scale=self.torque_length_scale
                if self.torque_length_scale is not None
                else self._vis.torque_length_scale,
                width_scale=self.torque_width_scale
                if self.torque_width_scale is not None
                else self._vis.torque_width_scale,
            )
        )

    return visuals

def request(
    self,
    signal_manager: SignalManager | None = None,
    channels: list[Literal["force", "torque"]]
    | dict[Literal["force", "torque"], dict[str, Any] | None] = ["force", "torque"],
):
    """
    Registers specific channels for logging.

    | Channel  | Description                        | Type |
    |:---------|:-----------------------------------|:-----|
    | `force`  | applied force in the world frame   | xyzm |
    | `torque` | applied torque in the world frame  | xyzm |

    Each channel is posted under `subgroups=(load_name, channel)`.

    * An `xyzm` is a cartesian vector, posted as 4 values (`x`, `y`, `z`, and its magnitude `m`).

    Each signal is tagged with built-in `dimension` metadata for its channel (`force` as force, `torque` as torque).

    If `signal_manager` is omitted, the `SignalManager` of the active `RuntimeManager` `with` block is used. If that `RuntimeManager` has no `SignalManager` configured, this is a no-op.

    Args:
        signal_manager: The signal manager to register the sampler with.
        channels: The load data channels to log. Pass a list to select channels, or a dict mapping channel name to metadata overrides (or `None`) to select channels and attach per-channel metadata in one step.

    """
    from mujoco_mojo.runtime.signal_manager import resolve_signal_manager

    signal_manager = resolve_signal_manager(signal_manager)
    if signal_manager is None:
        return

    if isinstance(channels, dict):
        _meta = cast("dict[str, dict[str, Any] | None]", channels)
        channels = list(channels.keys())
    else:
        _meta = {}

    channel_metadata = {
        "force": dim(Dimension.FORCE),
        "torque": torque_metadata(),
    }

    def sample(state: MjState):
        for channel in channels:
            source = self._last_f if channel == "force" else self._last_t
            meta = merge_signal_metadata(
                channel_metadata.get(channel), channel, _meta, unit_system=state.us
            )

            # iterate through x, y, z, and magnitude (pop. pop.)
            for i, attr in enumerate("xyzm"):
                signal_manager.post(
                    value=float(source[i]) if self.active else 0.0,
                    category=SignalCategory.LOADS,
                    # nest the force/torque under the function name
                    subgroups=(f"{self.name}", channel),
                    attr=attr,
                    metadata=meta,
                )

    signal_manager.register_sampler(sample)

def resolve_ids(self, state: MjState) -> None:
    """Caches the joint ID, DOF address, and DOF count from the compiled MuJoCo model."""
    self._jid = self.joint.get_id(state.model)
    jnt_type = state.model.jnt_type[self._jid]
    match jnt_type:
        case mujoco.mjtJoint.mjJNT_HINGE | mujoco.mjtJoint.mjJNT_SLIDE:
            self._nv = 1
        case mujoco.mjtJoint.mjJNT_BALL:
            self._nv = 3
        case _:
            msg = f"JointLoad '{self.name}': joint '{self.joint.name}' must be hinge, slide, or ball (got type {jnt_type})"
            logger.error(msg)
            raise ValueError(msg)
    self._dof_adr = int(state.model.jnt_dofadr[self._jid])
    logger.debug(
        f"resolved joint '{self.joint.name}' to dof_adr={self._dof_adr} nv={self._nv}",
    )

def get_visuals(self, state: MjState) -> list[ArrowConfig]:
    """Returns a list of arrow configurations for the renderer."""
    return []

def apply_load(self, state: MjState) -> None:
    """Evaluates `friction_func` at the current joint velocity vector and accumulates the result into `qfrc_applied`."""
    if not self.active or self._dof_adr < 0:
        self._last_force = np.zeros(3)
        return

    vel = np.array(state.data.qvel[self._dof_adr : self._dof_adr + self._nv])
    frc = self.friction_func(vel, state)
    state.data.qfrc_applied[self._dof_adr : self._dof_adr + self._nv] += frc
    self._last_force = self._to_world_frc(frc, state)

def request(
    self,
    signal_manager: SignalManager | None = None,
    metadata: dict[str, dict[str, Any]] | None = None,
) -> None:
    """
    Registers specific channels for logging.

    | Channel    | Description                                       | Type |
    |:-----------|:--------------------------------------------------|:-----|
    | `friction` | applied friction force/torque in the world frame  | xyzm |

    Each channel is posted under `subgroups=(load_name, channel)`.

    * An `xyzm` is a cartesian vector, posted as 4 values (`x`, `y`, `z`, and its magnitude `m`).

    `friction` is tagged with built-in `dimension` metadata resolved from the joint's type (force for slide, torque for hinge/ball).

    Args:
        signal_manager (SignalManager): Manager to register the sampler with. If omitted, the `SignalManager` of the active `RuntimeManager` `with` block is used. If that `RuntimeManager` has no `SignalManager` configured, this is a no-op.
        metadata: Metadata overriding or extending the built-in default for the `friction` channel.

    Raises:
        ValueError: If the joint has no name.

    """
    from mujoco_mojo.runtime.signal_manager import resolve_signal_manager

    signal_manager = resolve_signal_manager(signal_manager)
    if signal_manager is None:
        return

    if self.joint.name is None:
        msg = f"Cannot request telemetry for JointFriction '{self.name}': joint has no name."
        logger.error(msg)
        raise ValueError(msg)

    def sample(state: MjState) -> None:
        jnt_id = self.joint.get_id(state.model)
        jnt_type = int(state.model.jnt_type[jnt_id])
        meta = merge_signal_metadata(
            force_or_torque(jnt_type), "friction", metadata, unit_system=state.us
        )

        frc = self._last_force if self.active else np.zeros(3)
        for v, attr in zip(frc, ("x", "y", "z")):
            signal_manager.post(
                value=float(v),
                category=SignalCategory.LOADS,
                subgroups=(self.name, "friction"),
                attr=attr,
                metadata=meta,
            )
        signal_manager.post(
            value=float(np.linalg.norm(frc)),
            category=SignalCategory.LOADS,
            subgroups=(self.name, "friction"),
            attr="m",
            metadata=meta,
        )

    signal_manager.register_sampler(sample)

@classmethod
def coulomb_simple(
    cls,
    name: str,
    joint: Joint,
    magnitude: float | NamedValue[float],
) -> Self:
    """
    Constant-magnitude Coulomb (dry) friction opposing motion. Zero force at standstill. Good for brake pads, dry contacts, and cable friction.

    "Simple" because the magnitude is a fixed number you choose, not derived from any actual load on the joint. See `karnopp` for friction that responds to the joint's real bearing/pin reaction load.

    Args:
        name (str): Load name used for telemetry column labeling.
        joint (Joint): The MJCF joint to act on (slide, hinge, or ball).
        magnitude (float | NamedValue[float]): Friction force or torque magnitude. Accepts `NamedValue[float]` for runtime mutation.

    """

    def func(vel: np.ndarray, state: MjState) -> np.ndarray:
        speed = float(np.linalg.norm(vel))
        return (
            -float(magnitude) * vel / speed if speed > 1e-9 else np.zeros_like(vel)
        )

    return cls(name=name, joint=joint, friction_func=func)

@classmethod
def viscous_simple(
    cls,
    name: str,
    joint: Joint,
    damping: float | NamedValue[float],
) -> Self:
    """
    Velocity-proportional viscous damping. Force is continuous through zero with no discontinuity at standstill. Good for grease-lubricated joints and fluid drag.

    Args:
        name (str): Load name used for telemetry column labeling.
        joint (Joint): The MJCF joint to act on (slide, hinge, or ball).
        damping (float | NamedValue[float]): Damping coefficient (force per velocity, or torque per angular velocity). Accepts `NamedValue[float]` for runtime mutation.

    """

    def func(vel: np.ndarray, state: MjState) -> np.ndarray:
        return -float(damping) * vel

    return cls(name=name, joint=joint, friction_func=func)

@classmethod
def coulomb_viscous_simple(
    cls,
    name: str,
    joint: Joint,
    coulomb: float | NamedValue[float],
    viscous: float | NamedValue[float],
) -> Self:
    """
    Coulomb and viscous friction combined, both at fixed magnitudes you choose. Constant sliding friction plus a velocity-proportional drag term.

    "Simple" because neither term is derived from any actual load on the joint. See `karnopp` for friction that responds to the joint's real bearing/pin reaction load.

    Args:
        name (str): Load name used for telemetry column labeling.
        joint (Joint): The MJCF joint to act on (slide, hinge, or ball).
        coulomb (float | NamedValue[float]): Coulomb friction force or torque magnitude. Accepts `NamedValue[float]` for runtime mutation.
        viscous (float | NamedValue[float]): Viscous damping coefficient (force per velocity, or torque per angular velocity). Accepts `NamedValue[float]` for runtime mutation.

    """

    def func(vel: np.ndarray, state: MjState) -> np.ndarray:
        speed = float(np.linalg.norm(vel))
        coulomb_term = (
            -float(coulomb) * vel / speed if speed > 1e-9 else np.zeros_like(vel)
        )
        return coulomb_term - float(viscous) * vel

    return cls(name=name, joint=joint, friction_func=func)

@classmethod
def stribeck_simple(
    cls,
    name: str,
    joint: Joint,
    coulomb: float | NamedValue[float],
    static: float | NamedValue[float],
    stribeck_velocity: float | NamedValue[float],
    viscous: float | NamedValue[float] = 0.0,
) -> Self:
    """
    Full Stribeck friction model, at fixed magnitudes you choose. Friction peaks at standstill, drops to the kinetic level as motion begins, then rises with speed. Best for brake and clutch models where stick-slip matters.

    F = -(coulomb + (static - coulomb) * exp(-|v| / stribeck_velocity) + viscous * |v|) * v / |v|

    Args:
        name (str): Load name used for telemetry column labeling.
        joint (Joint): The MJCF joint to act on (slide, hinge, or ball).
        coulomb (float | NamedValue[float]): Kinetic friction force or torque magnitude. Accepts `NamedValue[float]`.
        static (float | NamedValue[float]): Peak static friction force or torque at zero velocity. Must be >= `coulomb`. Accepts `NamedValue[float]`.
        stribeck_velocity (float | NamedValue[float]): Characteristic velocity at which friction transitions from static to kinetic. Smaller values give a sharper transition. Accepts `NamedValue[float]`.
        viscous (float | NamedValue[float], optional): Velocity-proportional damping (force per velocity, or torque per angular velocity). Defaults to 0. Accepts `NamedValue[float]`.

    """

    def func(vel: np.ndarray, state: MjState) -> np.ndarray:
        speed = float(np.linalg.norm(vel))
        if speed < 1e-9:
            return np.zeros_like(vel)
        fc, fs = float(coulomb), float(static)
        vs, fv = float(stribeck_velocity), float(viscous)
        mag = fc + (fs - fc) * np.exp(-speed / max(vs, 1e-9)) + fv * speed
        return -mag * vel / speed

    return cls(name=name, joint=joint, friction_func=func)

@classmethod
def karnopp(
    cls,
    name: str,
    joint: Joint,
    mu_kinetic: float | NamedValue[float],
    mu_static: float | NamedValue[float],
    velocity_threshold: float | NamedValue[float],
    viscous: float | NamedValue[float] = 0.0,
) -> Self:
    """
    Karnopp friction whose normal force is the joint's actual bearing/pin reaction load (`Joint.rt_bearing_load`), not a fixed magnitude. Responds to ordinary side loads (e.g. a body's weight pressing on a hinge pin, or a side load on a slider) not just a number you pick. Properly distinguishes "stuck" from "sliding" instead of picking a coefficient based on speed alone:

    - Sliding (`|v| >= velocity_threshold`): ordinary kinetic Coulomb friction, `F = -mu_kinetic * bearing_load * v/|v|`, plus an optional `viscous` term.
    - Stuck (`|v| < velocity_threshold`): rather than picking a direction from a near-zero (and possibly noisy) velocity, friction is set to exactly cancel whatever other smooth force is currently acting on the joint (`qfrc_smooth`: passive + actuator + applied + bias), clamped to the static limit `mu_static * bearing_load`. If the driving force exceeds that limit the joint breaks away and friction saturates at the limit, opposing the driving force.

    This avoids the chattering a pure velocity-direction switch can produce right at standstill, since the held force no longer depends on the sign of a near-zero, noisy velocity.

    The bearing load comes from `Joint.rt_cfrc_int`, which refreshes itself on demand via `state.ensure_rne_post_constraint()`, so no extra wiring is needed regardless of how the simulation is driven. Note that `qfrc_smooth` reflects the previous step's solve, since this step's values are not yet known when loads are applied; this mirrors the existing one-step lag already inherent in using `vel` to set this step's force.

    Args:
        name (str): Load name used for telemetry column labeling.
        joint (Joint): The MJCF joint to act on (slide, hinge, or ball).
        mu_kinetic (float | NamedValue[float]): Friction coefficient applied to the bearing load while sliding. Accepts `NamedValue[float]`.
        mu_static (float | NamedValue[float]): Friction coefficient applied to the bearing load that the joint can hold against before breaking away. Accepts `NamedValue[float]`.
        velocity_threshold (float | NamedValue[float]): Speed below which the joint is considered stuck rather than sliding. Accepts `NamedValue[float]`.
        viscous (float | NamedValue[float], optional): Velocity-proportional damping added on top of the kinetic term while sliding. Defaults to 0. Accepts `NamedValue[float]`.

    """

    def func(vel: np.ndarray, state: MjState) -> np.ndarray:
        speed = float(np.linalg.norm(vel))
        normal = joint.rt_bearing_load(state)

        if speed < float(velocity_threshold):
            driving = np.array(joint.rt_qfrc_smooth(state))
            drive_mag = float(np.linalg.norm(driving))
            if drive_mag < 1e-9:
                return np.zeros_like(vel)
            max_static = float(mu_static) * normal
            return -min(drive_mag, max_static) * driving / drive_mag

        if speed < 1e-9:
            return np.zeros_like(vel)
        coulomb_term = -float(mu_kinetic) * normal * vel / speed
        return coulomb_term - float(viscous) * vel

    return cls(name=name, joint=joint, friction_func=func)