Load

@abstractmethod
def resolve_ids(self, state: MjState) -> None:
    """Cache any integer IDs from the compiled MuJoCo model."""

@abstractmethod
def apply_load(self, state: MjState) -> None:
    """Write the load contribution into the MuJoCo force buffers."""

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

def resolve_ids(self, state: MjState):
    """Caches the integer IDs from the compiled MuJoCo model."""
    self._vis = MujocoMojoSettings().visualization
    self.action_site.get_id(state.model)

    if self.rel_to_site:
        self.rel_to_site.get_id(state.model)

@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 resolve_ids(self, state: MjState):
    """Caches the integer IDs from the compiled MuJoCo model."""
    super().resolve_ids(state)
    self.xtion_site.get_id(state.model)
    self._r0_mag = self.action_site.rt_dm(self.xtion_site, state)
    on_resolve = getattr(self.magnitude_func, "on_resolve", None)
    if callable(on_resolve):
        on_resolve(self._r0_mag)

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

    visuals = super().get_visuals(state)

    if self._last_f[3] > 1e-4 and self._vis.reaction_force:
        xtion_pos = self.xtion_site.rt_pos(state)
        visuals.append(
            ArrowConfig(
                pos=xtion_pos,
                vec=-self._last_f[:3],
                color=Color[self._vis.reaction_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,
            )
        )

    return visuals

@classmethod
def ideal_spring(
    cls,
    name: str,
    action_site: AnySite,
    xtion_site: AnySite,
    stiffness: float | NamedValue[float] = 0.0,
    damping: float | NamedValue[float] = 0.0,
    rest_length: float = 0.0,
) -> Self:
    """Bidirectional spring-damper active in both tension and compression. `F = -k*(d - d0) - c*v`."""

    def logic(ud: UserData | None, state: MjState) -> float:
        dr = action_site.rt_displacements(xtion_site, state)
        dist = float(np.linalg.norm(dr))
        unit_vec = dr / dist if dist > 1e-9 else np.zeros(3)
        vel = float(
            np.dot(
                action_site.rt_velocities(xtion_site, state)[3:6],
                unit_vec,
            )
        )
        return _ideal_force_logic(dist, vel, stiffness, damping, rest_length)

    return cls(
        name=name,
        action_site=action_site,
        xtion_site=xtion_site,
        magnitude_func=logic,
    )

@classmethod
def stroke_compression_spring(
    cls,
    name: str,
    action_site: AnySite,
    xtion_site: AnySite,
    stiffness: float | NamedValue[float] = 0.0,
    damping: float | NamedValue[float] = 0.0,
    preload: float | NamedValue[float] = 0.0,
    max_stroke: float | NamedValue[float] = 0.1,
) -> Self:
    """Compression spring with a finite stroke; only active while the sites are compressed between rest and (rest + max_stroke). Useful for preloaded gas springs and mechanical end-stops."""

    class _Logic:
        def __init__(self) -> None:
            self._r0: float = 0.0

        def on_resolve(self, r0: float) -> None:
            self._r0 = r0

        def __call__(self, ud: UserData | None, state: MjState) -> float:
            dr = action_site.rt_displacements(xtion_site, state)
            dist = float(np.linalg.norm(dr))
            unit_vec = dr / dist if dist > 1e-9 else np.zeros(3)
            vel = float(
                np.dot(
                    action_site.rt_velocities(xtion_site, state)[3:6],
                    unit_vec,
                )
            )

            k = stiffness.value if isinstance(stiffness, NamedValue) else stiffness
            c = damping.value if isinstance(damping, NamedValue) else damping
            f_0 = preload.value if isinstance(preload, NamedValue) else preload
            d_f = (
                max_stroke.value
                if isinstance(max_stroke, NamedValue)
                else max_stroke
            )

            delta_d = dist - self._r0
            if 0 <= delta_d <= d_f:
                f_mag = f_0 - (k * delta_d) - (c * vel)
                return max(0.0, f_mag)
            return 0.0

    logic = _Logic()

    return cls(
        name=name,
        action_site=action_site,
        xtion_site=xtion_site,
        magnitude_func=logic,
    )

@classmethod
def compression_spring(
    cls,
    name: str,
    action_site: AnySite,
    xtion_site: AnySite,
    stiffness: float | NamedValue[float] = 0.0,
    damping: float | NamedValue[float] = 0.0,
    rest_length: float = 0.0,
) -> Self:
    """
    Creates a spring-damper that only acts when compressed (dist < rest_length). Useful for bumpers, feet, push-off springs, or end-stops.
    """

    def logic(ud: UserData | None, state: MjState) -> float:
        dr = action_site.rt_displacements(xtion_site, state)
        dist = float(np.linalg.norm(dr))
        unit_vec = dr / dist if dist > 1e-9 else np.zeros(3)
        vel = float(
            np.dot(
                action_site.rt_velocities(xtion_site, state)[3:6],
                unit_vec,
            )
        )
        if dist < rest_length:
            return _ideal_force_logic(dist, vel, stiffness, damping, rest_length)
        return 0.0

    return cls(
        name=name,
        action_site=action_site,
        xtion_site=xtion_site,
        magnitude_func=logic,
    )

@classmethod
def tension_spring(
    cls,
    name: str,
    action_site: AnySite,
    xtion_site: AnySite,
    stiffness: float | NamedValue[float] = 0.0,
    damping: float | NamedValue[float] = 0.0,
    rest_length: float = 0.0,
) -> Self:
    """
    Creates a spring-damper that only acts when extended (dist > rest_length). Useful for cables, bungees, or tendons.
    """

    def logic(ud: UserData | None, state: MjState) -> float:
        dr = action_site.rt_displacements(xtion_site, state)
        dist = float(np.linalg.norm(dr))
        unit_vec = dr / dist if dist > 1e-9 else np.zeros(3)
        vel = float(
            np.dot(
                action_site.rt_velocities(xtion_site, state)[3:6],
                unit_vec,
            )
        )
        if dist > rest_length:
            return _ideal_force_logic(dist, vel, stiffness, damping, rest_length)
        return 0.0

    return cls(
        name=name,
        action_site=action_site,
        xtion_site=xtion_site,
        magnitude_func=logic,
    )

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

@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)