Skip to content

Load

load

Load

Bases: MojoBaseModel, ABC

Base class to build forcing functions off of.

name instance-attribute

Python
name: str

Name of the forcing function. Used in data output column naming.

active class-attribute instance-attribute

Python
active: bool = True

Whether or not this force should be active.

action_site instance-attribute

Python
action_site: AnySite

Site on which the forcing function acts.

rel_to_site class-attribute instance-attribute

Python
rel_to_site: AnySite | None = None

Frame of reference for the calculated force. If None, uses worldbody.

resolve_ids

Python
resolve_ids(mj_model: MjModel, mj_data: MjData)

Caches the integer IDs from the compiled MuJoCo model.

Source code in src/mujoco_mojo/runtime/load.py
Python
def resolve_ids(self, mj_model: mujoco.MjModel, mj_data: mujoco.MjData):
    """Caches the integer IDs from the compiled MuJoCo model."""
    self.action_site.get_id(mj_model)

    if self.rel_to_site:
        self.rel_to_site.get_id(mj_model)

calculate abstractmethod

Python
calculate(
    mj_model: MjModel, mj_data: MjData
) -> tuple[ndarray, ndarray]

Calculate the force for the timestep.

Parameters:

Name Type Description Default
mj_model MjModel

description

required
mj_data MjData

description

required

Returns:

Type Description
tuple[ndarray, ndarray]

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

Source code in src/mujoco_mojo/runtime/load.py
Python
@abstractmethod
def calculate(
    self, mj_model: mujoco.MjModel, mj_data: mujoco.MjData
) -> tuple[np.ndarray, np.ndarray]:
    """
    Calculate the force for the timestep.

    Args:
        mj_model (mujoco.MjModel): _description_
        mj_data (mujoco.MjData): _description_

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

    """

get_visuals

Python
get_visuals(
    mj_model: MjModel, mj_data: MjData
) -> list[ArrowConfig]

Returns a list of arrow configurations for the renderer.

Source code in src/mujoco_mojo/runtime/load.py
Python
def get_visuals(
    self, mj_model: mujoco.MjModel, mj_data: mujoco.MjData
) -> 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(mj_model, mj_data)

    # force arrow
    f_vec = self._last_f[:3]
    if np.linalg.norm(f_vec) > 1e-4:
        visuals.append(
            ArrowConfig(
                pos=action_pos,
                vec=f_vec,
                color=Color.EMERALD_500.rgba,
                is_torque=False,
            )
        )

    t_vec = self._last_t[:3]
    if np.linalg.norm(t_vec) > 1e-4:
        visuals.append(
            ArrowConfig(
                pos=action_pos,
                vec=t_vec,
                color=Color.AMBER_500.rgba,
                is_torque=True,
            )
        )

    return visuals

PointToPointForce

Bases: Load

Acts along the line-of-sight between two sites.

xtion_site instance-attribute

Python
xtion_site: AnySite

Site on which the forcing function will apply a reation force. Leave as None to use the worldbody.

This is called xtion to limit confusion between "reaction" and "relative".

magnitude_func instance-attribute

Python
magnitude_func: Callable[
    [float, float, float, MjModel, MjData], float
]

Func(distance, velocity, initial distance, MjModel, MjData) -> scalar_force. Can be a regular function, lambda, etc.

name instance-attribute

Python
name: str

Name of the forcing function. Used in data output column naming.

active class-attribute instance-attribute

Python
active: bool = True

Whether or not this force should be active.

action_site instance-attribute

Python
action_site: AnySite

Site on which the forcing function acts.

rel_to_site class-attribute instance-attribute

Python
rel_to_site: AnySite | None = None

Frame of reference for the calculated force. If None, uses worldbody.

resolve_ids

Python
resolve_ids(mj_model: MjModel, mj_data: MjData)

Caches the integer IDs from the compiled MuJoCo model.

Source code in src/mujoco_mojo/runtime/load.py
Python
def resolve_ids(self, mj_model: mujoco.MjModel, mj_data: mujoco.MjData):
    """Caches the integer IDs from the compiled MuJoCo model."""
    super().resolve_ids(mj_model, mj_data)
    self.xtion_site.get_id(mj_model)
    self._r0_mag = self.action_site.rt_dm(self.xtion_site, mj_model, mj_data)

get_visuals

Python
get_visuals(
    mj_model: MjModel, mj_data: MjData
) -> list[ArrowConfig]

Returns a list of arrow configurations for the renderer.

Source code in src/mujoco_mojo/runtime/load.py
Python
def get_visuals(
    self, mj_model: mujoco.MjModel, mj_data: mujoco.MjData
) -> list[ArrowConfig]:
    """Returns a list of arrow configurations for the renderer."""
    visuals = super().get_visuals(mj_model, mj_data)

    if not self.active:
        return []

    # Add the reaction force arrow at the xtion site
    xtion_pos = self.xtion_site.rt_pos(mj_model, mj_data)
    f_vec = self._last_f[:3]

    if np.linalg.norm(f_vec) > 1e-4:
        visuals.append(
            ArrowConfig(
                pos=xtion_pos,
                vec=-f_vec,  # opposite direction
                color=Color.ROSE_500.rgba,  # Red for Reaction
                is_torque=False,
            )
        )

    return visuals

ideal_spring classmethod

Python
ideal_spring(
    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

Standard linear spring-damper (works in both tension and compression).

Source code in src/mujoco_mojo/runtime/load.py
Python
@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:
    """Standard linear spring-damper (works in both tension and compression)."""

    def logic(
        d: float,
        v: float,
        r0: float,
        mj_model: mujoco.MjModel,
        mj_data: mujoco.MjData,
    ) -> float:
        return _ideal_force_logic(d, v, stiffness, damping, rest_length)

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

stroke_compression_spring classmethod

Python
stroke_compression_spring(
    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

Creates a spring-damper that only acts when the runtime length is between rest_length and (rest_legnth + stroke_length)

Source code in src/mujoco_mojo/runtime/load.py
Python
@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:
    """Creates a spring-damper that only acts when the runtime length is between rest_length and (rest_legnth + stroke_length)"""

    def logic(
        d: float,
        v: float,
        r0: float,
        mj_model: mujoco.MjModel,
        mj_data: mujoco.MjData,
    ) -> float:
        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 = d - r0

        if 0 <= delta_d <= d_f:
            f_mag = f_0 - (k * delta_d) - (c * v)
            return max(0.0, f_mag)
        return 0.0

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

compression_spring classmethod

Python
compression_spring(
    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.

Source code in src/mujoco_mojo/runtime/load.py
Python
@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(
        d: float,
        v: float,
        r0: float,
        mj_model: mujoco.MjModel,
        mj_data: mujoco.MjData,
    ) -> float:
        if d < rest_length:
            return _ideal_force_logic(d, v, stiffness, damping, rest_length)
        return 0.0

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

tension_spring classmethod

Python
tension_spring(
    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.

Source code in src/mujoco_mojo/runtime/load.py
Python
@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(
        d: float,
        v: float,
        r0: float,
        mj_model: mujoco.MjModel,
        mj_data: mujoco.MjData,
    ) -> float:
        if d > rest_length:
            return _ideal_force_logic(d, v, stiffness, damping, rest_length)
        return 0.0

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

BodyReactionForce

Bases: Load

xtion_body class-attribute instance-attribute

Python
xtion_body: Body | None = None

Body on which the load should be acted on. If None the world will be used.

name instance-attribute

Python
name: str

Name of the forcing function. Used in data output column naming.

active class-attribute instance-attribute

Python
active: bool = True

Whether or not this force should be active.

action_site instance-attribute

Python
action_site: AnySite

Site on which the forcing function acts.

rel_to_site class-attribute instance-attribute

Python
rel_to_site: AnySite | None = None

Frame of reference for the calculated force. If None, uses worldbody.

calculate abstractmethod

Python
calculate(
    mj_model: MjModel, mj_data: MjData
) -> tuple[ndarray, ndarray]

Calculate the force for the timestep.

Parameters:

Name Type Description Default
mj_model MjModel

description

required
mj_data MjData

description

required

Returns:

Type Description
tuple[ndarray, ndarray]

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

Source code in src/mujoco_mojo/runtime/load.py
Python
@abstractmethod
def calculate(
    self, mj_model: mujoco.MjModel, mj_data: mujoco.MjData
) -> tuple[np.ndarray, np.ndarray]:
    """
    Calculate the force for the timestep.

    Args:
        mj_model (mujoco.MjModel): _description_
        mj_data (mujoco.MjData): _description_

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

    """

get_visuals

Python
get_visuals(
    mj_model: MjModel, mj_data: MjData
) -> list[ArrowConfig]

Returns a list of arrow configurations for the renderer.

Source code in src/mujoco_mojo/runtime/load.py
Python
def get_visuals(
    self, mj_model: mujoco.MjModel, mj_data: mujoco.MjData
) -> 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(mj_model, mj_data)

    # force arrow
    f_vec = self._last_f[:3]
    if np.linalg.norm(f_vec) > 1e-4:
        visuals.append(
            ArrowConfig(
                pos=action_pos,
                vec=f_vec,
                color=Color.EMERALD_500.rgba,
                is_torque=False,
            )
        )

    t_vec = self._last_t[:3]
    if np.linalg.norm(t_vec) > 1e-4:
        visuals.append(
            ArrowConfig(
                pos=action_pos,
                vec=t_vec,
                color=Color.AMBER_500.rgba,
                is_torque=True,
            )
        )

    return visuals

ScalarForce

Bases: BodyReactionForce

Applies a scalar force along the local X-axis of the action_site.

scalar_func class-attribute instance-attribute

Python
scalar_func: Callable[
    [float, ndarray, MjModel, MjData], float
] = lambda t, unit_vec, m, d: 0.0

Func(time, action_site x axis unit vector, MjModel, MjData) -> scalar force value.

name instance-attribute

Python
name: str

Name of the forcing function. Used in data output column naming.

active class-attribute instance-attribute

Python
active: bool = True

Whether or not this force should be active.

action_site instance-attribute

Python
action_site: AnySite

Site on which the forcing function acts.

rel_to_site class-attribute instance-attribute

Python
rel_to_site: AnySite | None = None

Frame of reference for the calculated force. If None, uses worldbody.

xtion_body class-attribute instance-attribute

Python
xtion_body: Body | None = None

Body on which the load should be acted on. If None the world will be used.

get_visuals

Python
get_visuals(
    mj_model: MjModel, mj_data: MjData
) -> list[ArrowConfig]

Returns a list of arrow configurations for the renderer.

Source code in src/mujoco_mojo/runtime/load.py
Python
def get_visuals(
    self, mj_model: mujoco.MjModel, mj_data: mujoco.MjData
) -> 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(mj_model, mj_data)

    # force arrow
    f_vec = self._last_f[:3]
    if np.linalg.norm(f_vec) > 1e-4:
        visuals.append(
            ArrowConfig(
                pos=action_pos,
                vec=f_vec,
                color=Color.EMERALD_500.rgba,
                is_torque=False,
            )
        )

    t_vec = self._last_t[:3]
    if np.linalg.norm(t_vec) > 1e-4:
        visuals.append(
            ArrowConfig(
                pos=action_pos,
                vec=t_vec,
                color=Color.AMBER_500.rgba,
                is_torque=True,
            )
        )

    return visuals

ScalarTorque

Bases: BodyReactionForce

Applies a scalar torque along the local X-axis of the action_site.

scalar_func class-attribute instance-attribute

Python
scalar_func: Callable[
    [float, ndarray, MjModel, MjData], float
] = lambda t, unit_vec, m, d: 0.0

Func(time, action_site x-axis unit vector, MjModel, MjData) -> scalar torque value.

name instance-attribute

Python
name: str

Name of the forcing function. Used in data output column naming.

active class-attribute instance-attribute

Python
active: bool = True

Whether or not this force should be active.

action_site instance-attribute

Python
action_site: AnySite

Site on which the forcing function acts.

rel_to_site class-attribute instance-attribute

Python
rel_to_site: AnySite | None = None

Frame of reference for the calculated force. If None, uses worldbody.

xtion_body class-attribute instance-attribute

Python
xtion_body: Body | None = None

Body on which the load should be acted on. If None the world will be used.

get_visuals

Python
get_visuals(
    mj_model: MjModel, mj_data: MjData
) -> list[ArrowConfig]

Returns a list of arrow configurations for the renderer.

Source code in src/mujoco_mojo/runtime/load.py
Python
def get_visuals(
    self, mj_model: mujoco.MjModel, mj_data: mujoco.MjData
) -> 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(mj_model, mj_data)

    # force arrow
    f_vec = self._last_f[:3]
    if np.linalg.norm(f_vec) > 1e-4:
        visuals.append(
            ArrowConfig(
                pos=action_pos,
                vec=f_vec,
                color=Color.EMERALD_500.rgba,
                is_torque=False,
            )
        )

    t_vec = self._last_t[:3]
    if np.linalg.norm(t_vec) > 1e-4:
        visuals.append(
            ArrowConfig(
                pos=action_pos,
                vec=t_vec,
                color=Color.AMBER_500.rgba,
                is_torque=True,
            )
        )

    return visuals

GeneralLoad

Bases: VectorForce, VectorTorque

A 6-DOF force/torque applier.

name instance-attribute

Python
name: str

Name of the forcing function. Used in data output column naming.

active class-attribute instance-attribute

Python
active: bool = True

Whether or not this force should be active.

action_site instance-attribute

Python
action_site: AnySite

Site on which the forcing function acts.

rel_to_site class-attribute instance-attribute

Python
rel_to_site: AnySite | None = None

Frame of reference for the calculated force. If None, uses worldbody.

xtion_body class-attribute instance-attribute

Python
xtion_body: Body | None = None

Body on which the load should be acted on. If None the world will be used.

get_visuals

Python
get_visuals(
    mj_model: MjModel, mj_data: MjData
) -> list[ArrowConfig]

Returns a list of arrow configurations for the renderer.

Source code in src/mujoco_mojo/runtime/load.py
Python
def get_visuals(
    self, mj_model: mujoco.MjModel, mj_data: mujoco.MjData
) -> 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(mj_model, mj_data)

    # force arrow
    f_vec = self._last_f[:3]
    if np.linalg.norm(f_vec) > 1e-4:
        visuals.append(
            ArrowConfig(
                pos=action_pos,
                vec=f_vec,
                color=Color.EMERALD_500.rgba,
                is_torque=False,
            )
        )

    t_vec = self._last_t[:3]
    if np.linalg.norm(t_vec) > 1e-4:
        visuals.append(
            ArrowConfig(
                pos=action_pos,
                vec=t_vec,
                color=Color.AMBER_500.rgba,
                is_torque=True,
            )
        )

    return visuals