Skip to content

Index

utils

Color

Bases: StrEnum

Contains color aliases and converters to go between the various color formats. The color values are defined by the standard Tailwind CSS color theme.

MuJoCo uses a normalized rgba (all color channels clamped between 0 and 1). This enum defines with hex values, but they can be easily transformed with the provided methods (for example: Color.WHITE.rgba).

rgba property

Python
rgba: Vec4

Returns the normalized RGBA array for MuJoCo.

with_alpha

Python
with_alpha(alpha: float) -> Vec4

Returns the color with a custom transparency level.

Source code in src/mujoco_mojo/utils/color.py
Python
def with_alpha(self, alpha: float) -> Vec4:
    """Returns the color with a custom transparency level."""
    return Color.hex_to_rgba(self.value, alpha=alpha)

hex_to_rgba classmethod

Python
hex_to_rgba(hex_str: str, alpha: float = 1.0) -> Vec4

Converts '#RRGGBB' or 'RRGGBB' to normalized [0, 1] RGBA.

Source code in src/mujoco_mojo/utils/color.py
Python
@classmethod
def hex_to_rgba(cls, hex_str: str, alpha: float = 1.0) -> Vec4:
    """Converts '#RRGGBB' or 'RRGGBB' to normalized [0, 1] RGBA."""
    hex_str = hex_str.lstrip("#")
    # Convert hex to integers using bit-shifting
    lv = len(hex_str)
    if lv == 6:
        rgb = tuple(int(hex_str[i : i + 2], 16) for i in (0, 2, 4))
    else:
        raise ValueError(f"Invalid hex color: {hex_str}. Expected 6 characters.")

    return np.array([*(np.array(rgb) / 255.0), alpha])

rgba255_to_rgba classmethod

Python
rgba255_to_rgba(rgba255: Vec4) -> Vec4

Converts [0-255, 0-255, 0-255, 0-1] to normalized [0, 1] RGBA.

Source code in src/mujoco_mojo/utils/color.py
Python
@classmethod
def rgba255_to_rgba(cls, rgba255: Vec4) -> Vec4:
    """Converts [0-255, 0-255, 0-255, 0-1] to normalized [0, 1] RGBA."""
    res = np.array(rgba255, dtype=float)
    res[:3] /= 255.0
    return res

rgba_to_hex classmethod

Python
rgba_to_hex(rgba: Vec4) -> str

Converts normalized RGBA to '#RRGGBB' (alpha is discarded).

Source code in src/mujoco_mojo/utils/color.py
Python
@classmethod
def rgba_to_hex(cls, rgba: Vec4) -> str:
    """Converts normalized RGBA to '#RRGGBB' (alpha is discarded)."""
    rgba = np.asarray(rgba, dtype=float)
    rgb255 = (rgba[:3] * 255).astype(int)
    return "#{:02x}{:02x}{:02x}".format(*rgb255)

rgba_to_rgba255 classmethod

Python
rgba_to_rgba255(rgba: Vec4) -> Vec4

Converts normalized RGBA to [0-255, 0-255, 0-255, 0-1].

Source code in src/mujoco_mojo/utils/color.py
Python
@classmethod
def rgba_to_rgba255(cls, rgba: Vec4) -> Vec4:
    """Converts normalized RGBA to [0-255, 0-255, 0-255, 0-1]."""
    res = np.array(rgba, dtype=float)
    res[:3] *= 255.0
    return res

Interpolator

Bases: MojoBaseModel

Utility to handle 1D lookup tables for forcing functions.

x instance-attribute

Python
x: VecN

Interpolation abscissa values.

y instance-attribute

Python
y: VecN

Interpolation ordinate values

kind class-attribute instance-attribute

Python
kind: InterpOptions = 'linear'

What type of interpolation should be used.

from_series classmethod

Python
from_series(
    series: Series, kind: InterpOptions = "linear"
) -> Self

Create an interpolator directly from a Pandas Series.

Source code in src/mujoco_mojo/utils/interp.py
Python
@classmethod
def from_series(cls, series: pd.Series, kind: InterpOptions = "linear") -> Self:
    """Create an interpolator directly from a Pandas Series."""
    return cls(x=np.asarray(series.index), y=np.asarray(series.values), kind=kind)

Proximity

Bases: MojoBaseModel

Provide high-precision triangle-level distance queries.

geom_1 instance-attribute

Python
geom_1: Proximityable

First geometry to perform proximity calculations for.

geom_2 instance-attribute

Python
geom_2: Proximityable

Second geometry to perform proximity calculations for.

dist_max instance-attribute

Python
dist_max: float

The 'cutoff' distance. If objects are further than this (as estimated by a sphere to sphere test), the sphere to sphere estimate will be returned and exit early.

algorithm class-attribute instance-attribute

Python
algorithm: ProximityType = CONVEX_HULL

What algorithm should be used for the narrowphase test.

visualize class-attribute instance-attribute

Python
visualize: bool = True

Wheter or not to visualize this proximity in the MuJoCo viewer.

get_sphere_to_sphere_proximity

Python
get_sphere_to_sphere_proximity(
    state: MjState,
) -> tuple[float, Vec3, Vec3, bool]

Calculates the shortest distance between two geometries using their bounding spheres.

Parameters:

Name Type Description Default
state MjState

The paired MuJoCo model and data instance.

required

Returns:

Type Description
tuple[float, Vec3, Vec3, bool]

tuple[float, Vec3, Vec3, bool]: Unsigned (>= 0) minimum distance from geom_1 to geom_2, world location of minimum distance on geom_1, world location of minimum distance on geom_2, and if the estimated distance exceeds dist_max.

Source code in src/mujoco_mojo/utils/proximity.py
Python
def get_sphere_to_sphere_proximity(
    self,
    state: MjState,
) -> tuple[float, Vec3, Vec3, bool]:
    """
    Calculates the shortest distance between two geometries using their bounding spheres.

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

    Returns:
        tuple[float, Vec3, Vec3, bool]: Unsigned (`>= 0`) minimum distance from geom_1 to geom_2, world location of minimum distance on geom_1, world location of minimum distance on geom_2, and if the estimated distance exceeds dist_max.

    """
    # get world orientations and origins
    origin_geom_1 = self.geom_1.rt_pos(state)
    mat_geom_1 = self.geom_1.rt_xmat(state)

    origin_geom_2 = self.geom_2.rt_pos(state)
    mat_geom_2 = self.geom_2.rt_xmat(state)

    if np.isnan(self.geom_1._rad):
        self.geom_1._rad, self.geom_1._local_centroid = self.geom_1.vertex_max_norm(
            state.model
        )

    if np.isnan(self.geom_2._rad):
        self.geom_2._rad, self.geom_2._local_centroid = self.geom_2.vertex_max_norm(
            state.model
        )

    # shift centers to pre-calculated centroids
    pos_geom_1 = origin_geom_1 + (mat_geom_1 @ self.geom_1._local_centroid)
    pos_geom_2 = origin_geom_2 + (mat_geom_2 @ self.geom_2._local_centroid)

    rad_geom_1 = self.geom_1._rad
    rad_geom_2 = self.geom_2._rad

    vec_geom_1_to_geom_2 = pos_geom_2 - pos_geom_1
    d_centers = float(np.linalg.norm(vec_geom_1_to_geom_2))
    dist = d_centers - (rad_geom_1 + rad_geom_2)

    dist = max(0.0, dist)  # clip to zero
    exceeds_dist_max = dist > self.dist_max

    if d_centers > 1e-9:
        unit_vec = vec_geom_1_to_geom_2 / d_centers
        p1 = pos_geom_1 + (unit_vec * rad_geom_1)
        p2 = pos_geom_2 - (unit_vec * rad_geom_2)
    else:
        p1 = pos_geom_1
        p2 = pos_geom_2

    self.update_last(p1, p2, state)
    return dist, p1, p2, exceeds_dist_max

get_convex_hull_proximity

Python
get_convex_hull_proximity(
    state: MjState,
) -> tuple[float, Vec3, Vec3, ProximityType]

Calculates the shortest distance between two geometries using their convex hull.

Parameters:

Name Type Description Default
state MjState

The paired MuJoCo model and data instance.

required

Returns:

Type Description
tuple[float, Vec3, Vec3, ProximityType]

tuple[float, Vec3, Vec3, ProximityType]: Unsigned (>= 0) minimum distance from geom_1 to geom_2, world location of minimum distance on geom_1, world location of minimum distance on geom_2, and which phase the exit occurred in.

Source code in src/mujoco_mojo/utils/proximity.py
Python
def get_convex_hull_proximity(
    self,
    state: MjState,
) -> tuple[float, Vec3, Vec3, ProximityType]:
    """
    Calculates the shortest distance between two geometries using their convex hull.

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

    Returns:
        tuple[float, Vec3, Vec3, ProximityType]: Unsigned (`>= 0`) minimum distance from geom_1 to geom_2, world location of minimum distance on geom_1, world location of minimum distance on geom_2, and which phase the exit occurred in.

    """
    # ========== BROADPHASE ==========
    if self.geom_1._proximity_configured_for != ProximityType.CONVEX_HULL:
        self.geom_1.bake_proximity(state.model, ProximityType.CONVEX_HULL)

    if self.geom_2._proximity_configured_for != ProximityType.CONVEX_HULL:
        self.geom_2.bake_proximity(state.model, ProximityType.CONVEX_HULL)

    d_est, p1, p2, skip = self.get_sphere_to_sphere_proximity(state)

    if skip:
        return d_est, p1, p2, ProximityType.SPHERE_TO_SPHERE

    # ========== NARROWPHASE ==========
    # temp buffer for MuJoCo's 6-element output [x1,y1,z1, x2,y2,z2]
    mj_fromto = np.zeros(6)
    min_dist = mujoco.mj_geomDistance(
        m=state.model,
        d=state.data,
        geom1=self.geom_1.get_id(state.model),
        geom2=self.geom_2.get_id(state.model),
        distmax=self.dist_max,
        fromto=mj_fromto,
    )

    min_dist = max(0.0, min_dist)  # clip from below to zero

    p1 = mj_fromto[:3].copy()
    p2 = mj_fromto[3:6].copy()
    self.update_last(p1, p2, state)
    return min_dist, p1, p2, ProximityType.CONVEX_HULL

get_vertex_to_face_proximity

Python
get_vertex_to_face_proximity(
    state: MjState,
) -> tuple[float, Vec3, Vec3, ProximityType]

Calculates the vertex to face distance using a multi-phase Bounding Volume Hierarchy (BVH) query.

Phases
  1. Broad Phase: Sphere-Sphere check (object level).
  2. Mid Phase: BVH Traversal (eliminating triangle groups). No exit here.
  3. Narrow Phase: Point-to-Face proximity.

Parameters:

Name Type Description Default
state MjState

The paired MuJoCo model and data instance.

required

Returns:

Type Description
tuple[float, Vec3, Vec3, ProximityType]

tuple[float, Vec3, Vec3, ProximityType]: Unsigned (>= 0) minimum distance from geom_1 to geom_2, world location of minimum distance on geom_1, world location of minimum distance on geom_2, and which phase the exit occurred in.

Source code in src/mujoco_mojo/utils/proximity.py
Python
def get_vertex_to_face_proximity(
    self,
    state: MjState,
) -> tuple[float, Vec3, Vec3, ProximityType]:
    """
    Calculates the vertex to face distance using a multi-phase Bounding Volume Hierarchy (BVH) query.

    Phases:
        1. Broad Phase: Sphere-Sphere check (object level).
        2. Mid Phase: BVH Traversal (eliminating triangle groups). No exit here.
        3. Narrow Phase: Point-to-Face proximity.

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

    Returns:
        tuple[float, Vec3, Vec3, ProximityType]: Unsigned (`>= 0`) minimum distance from geom_1 to geom_2, world location of minimum distance on geom_1, world location of minimum distance on geom_2, and which phase the exit occurred in.

    """
    if self.geom_1._proximity_configured_for != ProximityType.VERTEX_TO_FACE:
        self.geom_1.bake_proximity(state.model, ProximityType.VERTEX_TO_FACE)

    if self.geom_2._proximity_configured_for != ProximityType.VERTEX_TO_FACE:
        self.geom_2.bake_proximity(state.model, ProximityType.VERTEX_TO_FACE)

    assert self.geom_1._baked_query and self.geom_2._baked_query
    assert (
        self.geom_2._local_verts is not None
        and self.geom_2._local_verts is not None
    )

    # ========== BROADPHASE: Sphere-Sphere check ==========
    # find center to center to center distance and return early if broad phase
    d_est, p1, p2, skip = self.get_sphere_to_sphere_proximity(state)
    if skip:
        return d_est, p1, p2, ProximityType.SPHERE_TO_SPHERE

    # ========== COORDINATE TRANSFORMATION ==========
    pos_geom_1 = self.geom_1.rt_pos(state)
    pos_geom_2 = self.geom_2.rt_pos(state)

    mat_geom_1 = self.geom_1.rt_xmat(state)  # already Mat3 (3x3)
    mat_geom_2 = self.geom_2.rt_xmat(state)
    rel_pos = pos_geom_2 - pos_geom_1

    # ========== NARROWPHASE A: Geom_1-Surface vs. Geom_2-Vertices ==========
    # trimesh uses a BVH internall here (Mid-phase) to find closest triangles
    # combine transforms from geom_1 to geom_2: V_local_geom_1 = R_geom_1.T @ (R_geom_2 @ V_local_geom_2 + p_geom_2 - p_geom_1)
    geom_2_v_in_geom_1 = (
        self.geom_2._local_verts @ mat_geom_2.T + rel_pos
    ) @ mat_geom_1
    pts_on_geom_1, dist_a, _ = self.geom_1._baked_query.on_surface(
        geom_2_v_in_geom_1
    )
    idx_a = np.argmin(dist_a)
    min_a = dist_a[idx_a]

    # ========== NARROWPHASE B: Geom_1-Vertices vs. Geom_2-Surface  ==========
    # transform geom_1 vertices into geom_2's local frame
    geom_1_v_in_geom_2 = (
        self.geom_1._local_verts @ mat_geom_1.T - rel_pos
    ) @ mat_geom_2
    pts_on_geom_2, dist_b, _ = self.geom_2._baked_query.on_surface(
        geom_1_v_in_geom_2
    )
    idx_b = np.argmin(dist_b)
    min_b = dist_b[idx_b]

    # ========== CLEANUP ==========
    # find global min
    if min_a < min_b:
        min_dist = float(min_a)
        p1 = (pts_on_geom_1[idx_a] @ mat_geom_1.T) + pos_geom_1
        p2 = (geom_2_v_in_geom_1[idx_a] @ mat_geom_1.T) + pos_geom_1

        self.update_last(p1, p2, state)
        return min_dist, p1, p2, ProximityType.VERTEX_TO_FACE
    else:
        min_dist = float(min_b)

        # pt_on_geom_2 was calculated in geom_2's local frame
        p2 = (pts_on_geom_2[idx_b] @ mat_geom_2.T) + pos_geom_2
        p1 = (geom_1_v_in_geom_2[idx_b] @ mat_geom_2.T) + pos_geom_2

        self.update_last(p1, p2, state)
        return min_dist, p1, p2, ProximityType.VERTEX_TO_FACE

get_face_to_face_proximity

Python
get_face_to_face_proximity(
    state: MjState,
) -> tuple[float, Vec3, Vec3, ProximityType]

Calculates the face to face distance using a multi-phase Bounding Volume Hierarchy (BVH) query.

This is more accurate than the vertex to face method, but comes at higher computational cost.

Phases
  1. Broad Phase: Sphere-Sphere check (object level).
  2. Mid Phase: BVH Traversal (eliminating triangle groups). No exit here.
  3. Narrow Phase: Face-to-Face proximity.

Parameters:

Name Type Description Default
state MjState

The paired MuJoCo model and data instance.

required

Returns:

Type Description
tuple[float, Vec3, Vec3, ProximityType]

tuple[float, Vec3, Vec3, ProximityType]: Unsigned (>= 0) minimum distance from geom_1 to geom_2, world location of minimum distance on geom_1, world location of minimum distance on geom_2, and which phase the exit occurred in.

Source code in src/mujoco_mojo/utils/proximity.py
Python
def get_face_to_face_proximity(
    self,
    state: MjState,
) -> tuple[float, Vec3, Vec3, ProximityType]:
    """
    Calculates the face to face distance using a multi-phase Bounding Volume Hierarchy (BVH) query.

    This is more accurate than the vertex to face method, but comes at higher computational cost.

    Phases:
        1. Broad Phase: Sphere-Sphere check (object level).
        2. Mid Phase: BVH Traversal (eliminating triangle groups). No exit here.
        3. Narrow Phase: Face-to-Face proximity.

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

    Returns:
        tuple[float, Vec3, Vec3, ProximityType]: Unsigned (`>= 0`) minimum distance from geom_1 to geom_2, world location of minimum distance on geom_1, world location of minimum distance on geom_2, and which phase the exit occurred in.

    """
    if self.geom_1._proximity_configured_for != ProximityType.FACE_TO_FACE:
        self.geom_1.bake_proximity(state.model, ProximityType.FACE_TO_FACE)

    if self.geom_2._proximity_configured_for != ProximityType.FACE_TO_FACE:
        self.geom_2.bake_proximity(state.model, ProximityType.FACE_TO_FACE)

    assert self.geom_1._baked_manager and self.geom_2._baked_manager

    # ========== BROADPHASE: Sphere-Sphere check ==========

    # find center to center to center distance and return early if broad phase
    d_est, p1, p2, skip = self.get_sphere_to_sphere_proximity(state)
    if skip:
        return d_est, p1, p2, ProximityType.SPHERE_TO_SPHERE

    # ========== NARROWPHASE ==========
    # set the other transformation relative to geom_1's local frame
    t_geom_1 = np.eye(4)
    t_geom_1[:3, :3] = self.geom_1.rt_xmat(state)
    t_geom_1[:3, 3] = self.geom_1.rt_pos(state)

    t_geom_2 = np.eye(4)
    t_geom_2[:3, :3] = self.geom_2.rt_xmat(state)
    t_geom_2[:3, 3] = self.geom_2.rt_pos(state)

    self.geom_1._baked_manager.set_transform(self.geom_1.name, t_geom_1)
    self.geom_2._baked_manager.set_transform(self.geom_2.name, t_geom_2)

    # CollisionManager returns distance and the two closest points
    result = self.geom_1._baked_manager.min_distance_other(
        self.geom_2._baked_manager, return_data=True
    )
    min_dist = float(result[0])  # pyright: ignore[reportIndexIssue]
    data = result[1]  # pyright: ignore[reportIndexIssue]

    assert data
    p1 = data.point(self.geom_1.name)  # pyright: ignore[reportAttributeAccessIssue]
    p2 = data.point(self.geom_2.name)  # pyright: ignore[reportAttributeAccessIssue]

    self.update_last(p1, p2, state)
    return min_dist, p1, p2, ProximityType.FACE_TO_FACE

get_proximity

Python
get_proximity(
    state: MjState,
) -> tuple[float, Vec3, Vec3, ProximityType]

Calculates the shortest distance between two geometries using the specified proximity algorithm.

This is a general dispatcher method that routes to different proximity calculation algorithms based on the algorithm parameter. Each mode offers different tradeoffs between speed and precision:

Modes: - SPHERE_TO_SPHERE: Fastest. Uses bounding sphere distance only (broadphase). - CONVEX_HULL: Fast & accurate. Uses MuJoCo's convex hull-based distance (default). - VERTEX_TO_FACE: Accurate. Multi-phase BVH with vertex-to-surface queries. - FACE_TO_FACE: Most accurate but slowest. Full mesh-to-mesh distance calculation.

Phases (for non-sphere modes): 1. Broad Phase: Sphere-Sphere check (object level). 2. Narrow Phase: Algorithm-specific distance calculation.

Parameters:

Name Type Description Default
state MjState

The paired MuJoCo model and data instance.

required

Returns:

Type Description
float

tuple[float, ProximityType]: If fromto=False, returns the unsigned (>= 0) minimum distance and which algorithm produced the result.

Vec3

tuple[tuple[float, Vec3, Vec3], ProximityType]: If fromto=True, returns the minimum distance, world location of minimum distance on geom_1, world location of minimum distance on geom_2, and which algorithm produced the result.

The result is cached per-timestep (keyed on state.data.time), so calling this more than once during the same step (e.g. once from request()'s telemetry sampler and again from a user-defined runtime input/Load that reads the same Proximity instance) only pays for the underlying calculation once.

Source code in src/mujoco_mojo/utils/proximity.py
Python
def get_proximity(self, state: MjState) -> tuple[float, Vec3, Vec3, ProximityType]:
    """
    Calculates the shortest distance between two geometries using the specified proximity algorithm.

    This is a general dispatcher method that routes to different proximity calculation algorithms based on the `algorithm` parameter. Each mode offers different tradeoffs between speed and precision:

    **Modes:**
        - `SPHERE_TO_SPHERE`: Fastest. Uses bounding sphere distance only (broadphase).
        - `CONVEX_HULL`: Fast & accurate. Uses MuJoCo's convex hull-based distance (default).
        - `VERTEX_TO_FACE`: Accurate. Multi-phase BVH with vertex-to-surface queries.
        - `FACE_TO_FACE`: Most accurate but slowest. Full mesh-to-mesh distance calculation.

    **Phases (for non-sphere modes):**
        1. Broad Phase: Sphere-Sphere check (object level).
        2. Narrow Phase: Algorithm-specific distance calculation.

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

    Returns:
        tuple[float, ProximityType]: If fromto=False, returns the unsigned (`>= 0`) minimum distance and which algorithm produced the result.

        tuple[tuple[float, Vec3, Vec3], ProximityType]: If fromto=True, returns the minimum distance, world location of minimum distance on geom_1, world location of minimum distance on geom_2, and which algorithm produced the result.

    The result is cached per-timestep (keyed on `state.data.time`), so calling this more than once during the same step (e.g. once from `request()`'s telemetry sampler and again from a user-defined runtime input/Load that reads the same `Proximity` instance) only pays for the underlying calculation once.

    """
    is_cached = (
        self._last_prox_type is not None
        and not np.isnan(self._last_dist)
        and state.data.time == self._last_t
    )
    if is_cached:
        assert self._last_prox_type is not None  # narrows for the type checker
        return self._last_dist, self._last_p1, self._last_p2, self._last_prox_type

    match self.algorithm:
        case ProximityType.SPHERE_TO_SPHERE:
            d_est, p1, p2, _skip = self.get_sphere_to_sphere_proximity(state)
            result = (d_est, p1, p2, ProximityType.SPHERE_TO_SPHERE)
        case ProximityType.CONVEX_HULL:
            result = self.get_convex_hull_proximity(state)
        case ProximityType.VERTEX_TO_FACE:
            result = self.get_vertex_to_face_proximity(state)
        case ProximityType.FACE_TO_FACE:
            result = self.get_face_to_face_proximity(state)
        case _:
            msg = f"Method for {self.algorithm.name} not implemented."
            logger.error(msg)
            raise NotImplementedError(msg)

    # the configured algorithm's broadphase sphere-to-sphere check can exit early
    # (objects farther apart than dist_max), so the phase that actually produced
    # this result can flip between SPHERE_TO_SPHERE and the configured narrowphase
    # algorithm from one call to the next; surface that for performance debugging
    prox_type = result[3]
    if self._last_prox_type is not None and prox_type != self._last_prox_type:
        logger.debug(
            f"Proximity {self.pair_name} switched phase: "
            f"{self._last_prox_type.name} -> {prox_type.name} (t={state.data.time:.4f})"
        )
    self._last_prox_type = prox_type
    self._last_dist = result[0]
    # _last_t/_last_p1/_last_p2 were already set by update_last() inside whichever
    # method above produced `result`

    return result

request

Python
request(
    signal_manager: SignalManager | None = None,
    channels: list[Literal["dist", "fromto", "prox_type"]]
    | dict[
        Literal["dist", "fromto", "prox_type"],
        dict[str, Any] | None,
    ] = ["dist", "prox_type"],
)

Registers specific channels for logging.

Channel Description Type
dist minimum distance between the geom pair, per the proximity algorithm scalar
fromto world coordinates of the nearest point on each geom xyz
prox_type the ProximityType used to compute dist and fromto, as an integer scalar

Each channel is posted under subgroups=(pair_name,), where pair_name is f"{geom_1.name}_to_{geom_2.name}".

  • A scalar is posted as a single value with attr=channel.
  • An xyz is a cartesian vector without a magnitude component, posted as 3 values (x, y, z). fromto posts one xyz for each geom in the pair, under subgroups=(pair_name, "fromto", geom_name).

Only the computations required by the requested channels are performed each timestep.

Each signal is tagged with built-in dimension/unit metadata for its channel (dist/fromto as a length, prox_type as dimensionless).

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.

Parameters:

Name Type Description Default
signal_manager SignalManager | None

The signal manager to register the sampler with.

None
channels list[Literal['dist', 'fromto', 'prox_type']] | dict[Literal['dist', 'fromto', 'prox_type'], dict[str, Any] | None]

The proximity 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.

['dist', 'prox_type']
Source code in src/mujoco_mojo/utils/proximity.py
Python
def request(
    self,
    signal_manager: SignalManager | None = None,
    channels: list[Literal["dist", "fromto", "prox_type"]]
    | dict[Literal["dist", "fromto", "prox_type"], dict[str, Any] | None] = [
        "dist",
        "prox_type",
    ],
):
    """
    Registers specific channels for logging.

    | Channel     | Description                                                            | Type   |
    |:------------|:-----------------------------------------------------------------------|:-------|
    | `dist`      | minimum distance between the geom pair, per the proximity algorithm    | scalar |
    | `fromto`    | world coordinates of the nearest point on each geom                    | xyz    |
    | `prox_type` | the `ProximityType` used to compute `dist` and `fromto`, as an integer | scalar |

    Each channel is posted under `subgroups=(pair_name,)`, where `pair_name` is `f"{geom_1.name}_to_{geom_2.name}"`.

    * A `scalar` is posted as a single value with `attr=channel`.
    * An `xyz` is a cartesian vector without a magnitude component, posted as 3 values (`x`, `y`, `z`). `fromto` posts one `xyz` for each geom in the pair, under `subgroups=(pair_name, "fromto", geom_name)`.

    Only the computations required by the requested channels are performed each timestep.

    Each signal is tagged with built-in `dimension`/`unit` metadata for its channel (`dist`/`fromto` as a length, `prox_type` as dimensionless).

    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 proximity 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 = {}

    pair_name = self.pair_name
    _prox_attrs = {"dist", "fromto", "prox_type"}
    needs_proximity = bool(set(channels) & _prox_attrs)

    if needs_proximity:
        self._requested = True

    def sample(state: MjState):
        dist: float = np.nan
        p1: Vec3 = np.zeros(3)
        p2: Vec3 = np.zeros(3)
        prox_type = ProximityType.SPHERE_TO_SPHERE

        if needs_proximity:
            dist, p1, p2, prox_type = self.get_proximity(state)

        for channel in channels:
            meta = merge_signal_metadata(
                _REQUEST_CHANNEL_METADATA.get(channel),
                channel,
                _meta,
                unit_system=state.us,
            )

            match channel:
                case "dist":
                    signal_manager.post(
                        value=dist,
                        category=SignalCategory.PROXIMITIES,
                        subgroups=(pair_name,),
                        attr=channel,
                        metadata=meta,
                    )
                case "fromto":
                    for i, attr in enumerate("xyz"):
                        signal_manager.post(
                            value=float(p1[i]),
                            category=SignalCategory.PROXIMITIES,
                            subgroups=(pair_name, channel, str(self.geom_1.name)),
                            attr=attr,
                            metadata=meta,
                        )
                    for i, attr in enumerate("xyz"):
                        signal_manager.post(
                            value=float(p2[i]),
                            category=SignalCategory.PROXIMITIES,
                            subgroups=(pair_name, channel, str(self.geom_2.name)),
                            attr=attr,
                            metadata=meta,
                        )
                case "prox_type":
                    signal_manager.post(
                        value=float(prox_type.value),
                        category=SignalCategory.PROXIMITIES,
                        subgroups=(pair_name,),
                        attr=channel,
                        metadata=meta,
                    )
                case _:
                    continue

    signal_manager.register_sampler(sample)

MojoGenerator

Bases: Protocol

Definition of a function that generates a MojoModel model instance.

MojoObjective

Bases: Protocol

Definition of a function that scores a completed Mojo simulation.

MojoRunner dataclass

Python
MojoRunner(
    generator: MojoGenerator,
    generator_path: str | None = None,
    runtime: MojoRuntime | None = DEFAULT_RUNTIME,
    runtime_path: str | None = None,
    objective: MojoObjective | None = None,
    objective_path: str | None = None,
    seed: int | None = DEFAULT_SEED,
    workdir: Path = DEFAULT_WORKDIR,
    model_config_name: str
    | None = DEFAULT_MODEL_CONFIG_NAME,
    xml_name: str = DEFAULT_XML_NAME,
    config: MonteCarloConfig
    | OptimizerConfig = MonteCarloConfig(),
    gen_args: list[Any] = list(),
    gen_kwargs: dict[str, Any] = dict(),
    run_args: list[Any] = list(),
    run_kwargs: dict[str, Any] = dict(),
)

slurm_trial_id property

Python
slurm_trial_id: int | None

Returns the current SLURM task ID if running as part of an array job.

run

Python
run(
    global_overrides: NamedValueDict[
        NDArray
    ] = NamedValueDict[NDArray](),
    clean_workdir: bool = False,
    cleanup_delay: int = 10,
    execution_mode: ExecutionMode = LOCAL,
    trial_nums: list[int] | None = None,
) -> bool

Vectors a job to be either computed locally or to be orchestrated by SLURM.

Source code in src/mujoco_mojo/utils/runner.py
Python
def run(
    self,
    global_overrides: NamedValueDict[NDArray] = NamedValueDict[NDArray](),
    clean_workdir: bool = False,
    cleanup_delay: int = 10,
    execution_mode: ExecutionMode = ExecutionMode.LOCAL,
    trial_nums: list[int] | None = None,
) -> bool:
    """Vectors a job to be either computed locally or to be orchestrated by SLURM."""
    funcs = {
        "generator": self.generator,
        "runtime": self.runtime,
        "objective": self.objective,
    }
    non_none = {k: v for k, v in funcs.items() if v is not None}
    seen: dict[int, str] = {}
    for role, func in non_none.items():
        fid = id(func)
        if fid in seen:
            msg = f"'{getattr(func, '__qualname__', func)}' is assigned to both '{seen[fid]}' and '{role}'. Each function must be unique."
            logger.error(msg)
            raise ValueError(msg)
        seen[fid] = role

    if clean_workdir:
        if self.config.resume:
            msg = "clean_workdir and resume are mutually exclusive with one another. Use one or the other."
            logger.error(msg)
            raise ValueError(msg)

        self.force_remove_dir(countdown_from=cleanup_delay, path=self.workdir)

    self.workdir.mkdir(parents=True, exist_ok=True)
    if not (self.workdir / ".gitignore").exists():
        (self.workdir / ".gitignore").write_text("*", encoding="utf-8")

    # SLURM workers re-enter run() with execution_mode=LOCAL; these only need
    # to happen once, during orchestration, to avoid every worker racing to
    # write the same shared files
    if self.slurm_trial_id is None:
        write_dojo_script(self.workdir)
        self.capture_environment()

    match execution_mode:
        case ExecutionMode.LOCAL:
            return self.run_local(
                global_overrides=global_overrides,
                trial_nums=trial_nums,
            )
        case ExecutionMode.SLURM:
            return self.orchestrate_slurm(
                global_overrides=global_overrides,
                trial_nums=trial_nums,
            )
        case _:
            msg = f"No run command has been configured for execution mode {execution_mode}"
            logger.error(msg)
            raise NotImplementedError(msg)

orchestrate_slurm

Python
orchestrate_slurm(
    global_overrides: NamedValueDict[NDArray],
    trial_nums: list[int] | None = None,
) -> bool

Generates an sbatch script and submits the job array to SLURM for a given config.

Source code in src/mujoco_mojo/utils/runner.py
Python
def orchestrate_slurm(
    self,
    global_overrides: NamedValueDict[NDArray],
    trial_nums: list[int] | None = None,
) -> bool:
    """Generates an sbatch script and submits the job array to SLURM for a given config."""
    (self.workdir / "logs").mkdir(exist_ok=True)

    if isinstance(self.config, MonteCarloConfig):
        try:
            had_fails = self.orchestrate_slurm_monte_carlo(
                global_overrides=global_overrides,
                trial_nums=trial_nums,
            )
        except (BdbQuit, KeyboardInterrupt):
            print("\n")
            logger.error("Aborted SLURM orchestration!")
            had_fails = True
    else:
        msg = f"A SLURM configuration for {self.config.__class__.__name__} has not been implemented."
        logger.error(msg)
        raise NotImplementedError(msg)
    return had_fails

execute_single_trial

Python
execute_single_trial(
    trial_num: int,
    seed: int | None,
    overrides_payload: dict,
) -> tuple[MojoModel | None, TrialStatus, MjState | None]

Helper to package a Trial and run it.

The model and state are live in-process objects; callers that need to cross a ProcessPoolExecutor boundary (i.e. parallel MC) must use _execute_trial_subprocess instead, which drops the non-picklable objects before the result is serialized back.

Returns (mojo_model, trial_status, state).

Source code in src/mujoco_mojo/utils/runner.py
Python
def execute_single_trial(
    self, trial_num: int, seed: int | None, overrides_payload: dict
) -> tuple[MojoModel | None, TrialStatus, MjState | None]:
    """
    Helper to package a Trial and run it.

    The model and state are live in-process objects; callers that need to cross a ProcessPoolExecutor boundary (i.e. parallel MC) must use `_execute_trial_subprocess` instead, which drops the non-picklable objects before the result is serialized back.

    Returns `(mojo_model, trial_status, state)`.
    """
    overrides = NamedValueDict[NDArray].model_validate(overrides_payload)

    trial = Trial(
        trial_num=trial_num,
        base_dir=self.workdir,
        xml_name=self.xml_name,
        model_config_name=self.model_config_name,
        padding_style=self.config.padding_style,
    )

    mojo_model, trial_status, state = trial.run(
        generator=self.generator,
        runtime=self.runtime,
        seed=seed,
        overrides=overrides,
        gen_args=self.gen_args,
        gen_kwargs=self.gen_kwargs,
        run_args=self.run_args,
        run_kwargs=self.run_kwargs,
    )

    # write distribution tables from this process so they never need to
    # cross the pickle boundary; FileLock serializes concurrent workers,
    # and the inner exists() check ensures only the first one writes
    if mojo_model is not None and mojo_model.dists:
        stochas_dir = self.workdir / STOCHAS_DIR_NAME
        stochas_dir.mkdir(exist_ok=True)
        dists_json_path = stochas_dir / STOCHAS_DISTS_FNAME
        with FileLock(stochas_dir / ".dists.lock"):
            if not dists_json_path.exists():
                mojo_model.dists.to_tables(stochas_dir)
                tmp = stochas_dir / "dists.tmp.json"
                tmp.write_text(mojo_model.dists.model_dump_json(), encoding="utf-8")
                tmp.replace(dists_json_path)

    return mojo_model, trial_status, state

run_monte_carlo

Python
run_monte_carlo(
    global_overrides: NamedValueDict[
        NDArray
    ] = NamedValueDict[NDArray](),
    trial_nums: list[int] | None = None,
) -> bool

Orchestrates a Monte Carlo job.

Source code in src/mujoco_mojo/utils/runner.py
Python
def run_monte_carlo(
    self,
    global_overrides: NamedValueDict[NDArray] = NamedValueDict[NDArray](),
    trial_nums: list[int] | None = None,
) -> bool:
    """Orchestrates a Monte Carlo job."""
    if self.slurm_trial_id is not None:
        tn = self.slurm_trial_id
        logger.info(f"SLURM Worker detected. Executing Trial {tn}")
        _mojo_model, trial_status, _ = self.execute_single_trial(
            trial_num=tn,
            seed=self.seed,
            overrides_payload=global_overrides.model_dump(),
        )

        return trial_status.completion == Completion.ERROR

    job_trial_nums = trial_nums if trial_nums else self.config.trial_nums

    # initialize the status tracker
    status_tracker = JobStatus(
        workdir=self.workdir.resolve(),
        job_type=JobType.MONTE_CARLO,
        execution_mode=ExecutionMode.LOCAL,
        n_proc=self.config.n_proc,
        seed=self.seed,
        padding_style=self.config.padding_style,
        generator=MojoRunner.inspect_protocol(self.generator),
        runtime=MojoRunner.inspect_protocol(self.runtime),
        objective=MojoRunner.inspect_protocol(self.objective),
        gen_args_used=bool(self.gen_args),
        gen_kwargs_used=bool(self.gen_kwargs),
        run_args_used=bool(self.run_args),
        run_kwargs_used=bool(self.run_kwargs),
        trial_nums=job_trial_nums,
    )

    # decide which trials to execute
    if self.config.resume:
        self._renumber_trial_folders(
            self.workdir.resolve(), self.config.padding_style
        )
        status_tracker.refresh_from_disk(n_proc=self.config.n_proc)
    status_tracker.dump_to_path(self.workdir / JOB_STATUS_FNAME)

    to_run = status_tracker.pending_trial_nums

    if not to_run:
        logger.info("All trials were already completed. Nothing to do.")
        return bool(status_tracker.unsuccessful_trial_nums)

    if self.config.is_parallel:
        logger.info(
            f"Running {len(to_run)} trials with {self.config.n_proc} processors. {status_tracker.n_done}/{self.config.n_trial} ({status_tracker.progress:.2%}) trials completed."
        )
        # needed for logging on Windows
        with multiprocessing.Manager() as m:
            log_queue = m.Queue()

            parent_log_level = logging.getLogger().getEffectiveLevel()
            listener = QueueListener(log_queue, *logging.getLogger().handlers)
            listener.start()

            executor = ProcessPoolExecutor(
                max_workers=self.config.n_proc,
                initializer=worker_init,
                initargs=(log_queue, parent_log_level),
            )
            try:
                future_to_tn = {
                    executor.submit(
                        self._execute_trial_subprocess,
                        tn,
                        self.seed,
                        global_overrides.model_dump(),
                    ): tn
                    for tn in to_run
                }
                for f in as_completed(future_to_tn):
                    tn = future_to_tn[f]
                    try:
                        trial_status = f.result()
                        status_tracker.update_trial(status=trial_status)
                    except (BdbQuit, KeyboardInterrupt):
                        # user is quitting from breakpoint() or CTRL+C
                        raise
                    except Exception as e:
                        logger.exception(f"Trial {tn} failed: {e}")
                        status_tracker.update_trial(
                            status=TrialStatus(
                                trial_num=tn, completion=Completion.ERROR
                            )
                        )
                    status_tracker.generate_report()
            except (BdbQuit, KeyboardInterrupt):
                # allows killing the job with one CTRL+C
                logger.warning("Interrupt recieved. Stopping all trials.")
                executor.shutdown(wait=False, cancel_futures=True)
                raise
            finally:
                listener.stop()
                executor.shutdown(wait=True)
    else:
        for tn in to_run:
            try:
                _, trial_status, _ = self.execute_single_trial(
                    trial_num=tn,
                    seed=self.seed,
                    overrides_payload=global_overrides.model_dump(),
                )
                status_tracker.update_trial(
                    status=trial_status,
                )
            except (BdbQuit, KeyboardInterrupt):
                # user is quitting from breakpoint() or CTRL+C
                raise
            except Exception as e:
                logger.exception(f"A trial failed with error: {e}")
                status_tracker.update_trial(
                    status=TrialStatus(trial_num=tn, completion=Completion.ERROR)
                )
            status_tracker.generate_report()

    status_tracker.generate_report(alert_generation=True)
    return bool(status_tracker.unsuccessful_trial_nums)

get_slurm_array_string staticmethod

Python
get_slurm_array_string(ids: list[int]) -> str

Collapses [0, 1, 2, 5, 6] into '0-2,5-6' for SLURM.

Source code in src/mujoco_mojo/utils/runner.py
Python
@staticmethod
def get_slurm_array_string(ids: list[int]) -> str:
    """Collapses [0, 1, 2, 5, 6] into '0-2,5-6' for SLURM."""
    if not ids:
        return ""

    ranges = []
    # Identify groups of consecutive integers
    for _, group in itertools.groupby(
        enumerate(sorted(ids)), lambda x: x[1] - x[0]
    ):
        group = list(group)
        start = group[0][1]
        end = group[-1][1]
        if start == end:
            ranges.append(str(start))
        else:
            ranges.append(f"{start}-{end}")

    return ",".join(ranges)

get_slurm_partitions staticmethod

Python
get_slurm_partitions() -> tuple[list[str], str | None]

Queries sinfo for available partitions and identifies the default.

Source code in src/mujoco_mojo/utils/runner.py
Python
@staticmethod
def get_slurm_partitions() -> tuple[list[str], str | None]:
    """Queries sinfo for available partitions and identifies the default."""
    try:
        result = subprocess.run(
            ["sinfo", "-h", "--format=%P"],
            capture_output=True,
            text=True,
            timeout=5,
        )
        if result.returncode == 0:
            raw_partitions = [
                p.strip() for p in result.stdout.splitlines() if p.strip()
            ]

            default_partition = None
            clean_partitions = []

            for p in raw_partitions:
                if p.endswith("*"):
                    name = p.replace("*", "")
                    default_partition = name
                    clean_partitions.append(name)
                else:
                    clean_partitions.append(p)

            return sorted(list(set(clean_partitions))), default_partition
    except Exception:
        pass
    return [], None

normalize_to_mb staticmethod

Python
normalize_to_mb(mem_str: str) -> int

Converts SLURM memory strings (e.g., '1000', '1G', '1024M') to integer MB.

Source code in src/mujoco_mojo/utils/runner.py
Python
@staticmethod
def normalize_to_mb(mem_str: str) -> int:
    """Converts SLURM memory strings (e.g., '1000', '1G', '1024M') to integer MB."""
    # Split the number from the unit (e.g., '1024M' -> '1024', 'M')
    match = re.match(r"(\d+)([KMGTP]?)", mem_str.upper())
    if not match:
        return 0

    value, unit = match.groups()
    value = int(value)

    # SLURM units are powers of 1024
    multipliers = {
        "K": 1 / 1024,  # Kilobytes to MB
        "M": 1,  # Megabytes
        "G": 1024,  # Gigabytes to MB
        "T": 1024**2,  # Terabytes to MB
        "P": 1024**3,  # Petabytes to MB
    }

    return int(value * multipliers.get(unit or "M", 1))

format_bytes staticmethod

Python
format_bytes(mb_value: int) -> str

Scales MB back up to the most readable unit (G, T, etc.).

Source code in src/mujoco_mojo/utils/runner.py
Python
@staticmethod
def format_bytes(mb_value: int) -> str:
    """Scales MB back up to the most readable unit (G, T, etc.)."""
    units = ["M", "G", "T", "P"]
    value = float(mb_value)
    unit_index = 0

    # Keep dividing by 1024 as long as it's a clean multiple
    while value >= 1024 and unit_index < len(units) - 1:
        value /= 1024
        unit_index += 1

    # If it's a whole number (like 1.0G), show it as 1G.
    # Otherwise, show one decimal place (like 1.5G).
    if value.is_integer():
        return f"{int(value)}{units[unit_index]}"
    return f"{value:.1f}{units[unit_index]}"

get_slurm_node_mem_limit staticmethod

Python
get_slurm_node_mem_limit(partition_name: str) -> str

Finds the MINIMUM RealMemory limit, normalized to MB.

Source code in src/mujoco_mojo/utils/runner.py
Python
@staticmethod
def get_slurm_node_mem_limit(partition_name: str) -> str:
    """Finds the MINIMUM RealMemory limit, normalized to MB."""
    try:
        # 1. Get nodes from partition
        part_info = subprocess.run(
            ["scontrol", "show", "partition", partition_name, "-o"],
            capture_output=True,
            text=True,
        ).stdout
        node_match = re.search(r"\bNodes=(\S+)", part_info)
        if not node_match:
            return "<UNKNOWN>"

        # 2. Get node info (can handle ranges like c[1-2])
        node_info = subprocess.run(
            ["scontrol", "show", "node", node_match.group(1), "-o"],
            capture_output=True,
            text=True,
        ).stdout

        # 3. Find all RealMemory values (capturing optional suffixes)
        # Regex captures digits and any trailing letters (like G or M)
        mem_matches = re.findall(r"\bRealMemory=(\d+[KMGTP]?)", node_info)
        if mem_matches:
            # Normalize every match to MB and find the lowest one
            min_mb = min(MojoRunner.normalize_to_mb(m) for m in mem_matches)
            return MojoRunner.format_bytes(min_mb)

    except Exception:
        pass

    return "<UNKNOWN>"

get_slurm_cpu_limit staticmethod

Python
get_slurm_cpu_limit(partition_name: str) -> str

Finds the minimum of the max allowed CPUs (CPUTot) among nodes in a partition.

Source code in src/mujoco_mojo/utils/runner.py
Python
@staticmethod
def get_slurm_cpu_limit(partition_name: str) -> str:
    """Finds the minimum of the max allowed CPUs (CPUTot) among nodes in a partition."""
    try:
        # 1. Get the nodes in the partition
        part_cmd = ["scontrol", "show", "partition", partition_name, "-o"]
        part_info = subprocess.run(part_cmd, capture_output=True, text=True).stdout

        node_match = re.search(r"\bNodes=(\S+)", part_info)
        if not node_match:
            return "<UNKNOWN>"

        # 2. Get the node info
        node_cmd = ["scontrol", "show", "node", node_match.group(1), "-o"]
        nodes_info = subprocess.run(node_cmd, capture_output=True, text=True).stdout

        # 3. Find all CPUTot values (The total physical/logical CPUs on the node)
        cpu_values = re.findall(r"\bCPUTot=(\d+)", nodes_info)

        if cpu_values:
            # Find the minimum to ensure any node can handle the task
            return str(min(int(v) for v in cpu_values))

    except Exception:
        pass

    return "<UNKNOWN>"

slurm_time_to_seconds staticmethod

Python
slurm_time_to_seconds(time_str: str) -> int

Converts SLURM time (D-HH:MM:SS or HH:MM:SS) to total seconds.

Source code in src/mujoco_mojo/utils/runner.py
Python
@staticmethod
def slurm_time_to_seconds(time_str: str) -> int:
    """Converts SLURM time (D-HH:MM:SS or HH:MM:SS) to total seconds."""
    if time_str.upper() == "UNLIMITED":
        return -1

    # Format: Days-Hours:Minutes:Seconds
    days = 0
    if "-" in time_str:
        days_part, time_str = time_str.split("-")
        days = int(days_part)

    parts = list(map(int, time_str.split(":")))
    if len(parts) == 3:  # HH:MM:SS
        return days * 86400 + parts[0] * 3600 + parts[1] * 60 + parts[2]
    if len(parts) == 2:  # MM:SS
        return days * 86400 + parts[0] * 60 + parts[1]
    if len(parts) == 1:  # MM
        return days * 86400 + parts[0] * 60
    return 0

get_slurm_time_limit staticmethod

Python
get_slurm_time_limit(partition_name: str) -> str

Finds the MaxTime limit for a specific partition.

Source code in src/mujoco_mojo/utils/runner.py
Python
@staticmethod
def get_slurm_time_limit(partition_name: str) -> str:
    """Finds the MaxTime limit for a specific partition."""
    try:
        # Get partition info
        part_cmd = ["scontrol", "show", "partition", partition_name, "-o"]
        part_info = subprocess.run(part_cmd, capture_output=True, text=True).stdout

        # Look for MaxTime= followed by the time string (e.g., 1-00:00:00 or UNLIMITED)
        time_match = re.search(r"\bMaxTime=(\S+)", part_info)
        if time_match:
            return time_match.group(1)

    except Exception:
        pass

    return "<UNKNOWN>"

get_max_array_size staticmethod

Python
get_max_array_size() -> int

Queries the global SLURM configuration for MaxArraySize.

Source code in src/mujoco_mojo/utils/runner.py
Python
@staticmethod
def get_max_array_size() -> int:
    """Queries the global SLURM configuration for MaxArraySize."""
    try:
        result = subprocess.run(
            ["scontrol", "show", "config"], capture_output=True, text=True
        ).stdout
        match = re.search(r"MaxArraySize=(\d+)", result)
        return int(match.group(1)) if match else 1001
    except Exception:
        return 1001

orchestrate_slurm_monte_carlo

Python
orchestrate_slurm_monte_carlo(
    global_overrides: NamedValueDict[
        NDArray
    ] = NamedValueDict[NDArray](),
    trial_nums: list[int] | None = None,
) -> bool

Orchestrates a Monte Carlo SLURM submission.

Source code in src/mujoco_mojo/utils/runner.py
Python
    def orchestrate_slurm_monte_carlo(
        self,
        global_overrides: NamedValueDict[NDArray] = NamedValueDict[NDArray](),
        trial_nums: list[int] | None = None,
    ) -> bool:
        """Orchestrates a Monte Carlo SLURM submission."""
        from rich.console import Console
        from rich.prompt import Confirm, Prompt

        console = Console()

        project_root = Path.cwd().resolve()

        # persist overrides so workers can access them
        overrides_path = self.workdir.resolve() / "global_overrides.json"
        if len(global_overrides) > 0:
            logger.info(f"Persisting global overrides to {overrides_path}")
            overrides_path.write_text(global_overrides.model_dump_json())

        # initialize the status tracker
        job_trial_nums = trial_nums if trial_nums else self.config.trial_nums

        status_tracker = JobStatus(
            workdir=self.workdir.resolve(),
            job_type=JobType.MONTE_CARLO,
            execution_mode=ExecutionMode.SLURM,
            n_proc=self.config.n_proc,
            seed=self.seed,
            padding_style=self.config.padding_style,
            generator=MojoRunner.inspect_protocol(self.generator),
            runtime=MojoRunner.inspect_protocol(self.runtime),
            objective=MojoRunner.inspect_protocol(self.objective),
            gen_args_used=bool(self.gen_args),
            gen_kwargs_used=bool(self.gen_kwargs),
            run_args_used=bool(self.run_args),
            run_kwargs_used=bool(self.run_kwargs),
            trial_nums=job_trial_nums,
        )

        # decide which trials to execute
        if self.config.resume:
            self._renumber_trial_folders(
                self.workdir.resolve(), self.config.padding_style
            )
            status_tracker.refresh_from_disk(n_proc=self.config.n_proc)
        status_tracker.dump_to_path(self.workdir / JOB_STATUS_FNAME)

        to_run = status_tracker.pending_trial_nums

        if not to_run:
            logger.info("All trials were already completed. Nothing to do.")
            return False
        else:
            logger.info(f"{len(to_run)} trials were identified for running.")

        # reconstruct the CLI command for the worker
        gen_args_str = " ".join([f'--gen-arg "{a}"' for a in self.gen_args])
        gen_kwargs_str = " ".join(
            [f'--gen-kwarg "{k}={v}"' for k, v in self.gen_kwargs.items()]
        )
        run_args_str = " ".join([f'--run-arg "{a}"' for a in self.run_args])
        run_kwargs_str = " ".join(
            [f'--run-kwarg "{k}={v}"' for k, v in self.run_kwargs.items()]
        )

        runtime_flag = f'--runtime "{self.runtime_path}"' if self.runtime_path else ""
        seed_flag = f"--seed {self.seed}" if self.seed is not None else ""
        overrides_flag = (
            f'--overrides "{overrides_path}"' if len(global_overrides) > 0 else ""
        )

        # get the path to the mujoco-mojo CLI executable
        py_bin_dir = Path(sys.executable).parent.resolve()
        mojo_cmd = py_bin_dir / "mujoco-mojo"

        cmd = (
            f"{mojo_cmd} run monte-carlo "
            f'--generator "{self.generator_path}" '
            f"{runtime_flag} {seed_flag} {overrides_flag} "
            f'--workdir "{self.workdir.resolve()}" '
            f"{gen_args_str} {gen_kwargs_str} "
            f"{run_args_str} {run_kwargs_str} "
            f"--n-trials 0 "  # force to zero to prevent an expected warning
            f"--trial-num $SLURM_ARRAY_TASK_ID "  # execute its onw trial_num
            f"--execution-mode local "  # using local since slurm will just send us back to this method
            f"--n-proc 1"  # A worker only needs 1 process
        )

        # ask for sbatch settings with a bunch of console inputs with default values
        available_partitions, default_partition = self.get_slurm_partitions()

        console.print(
            "\n[bold cyan]MuJoCo Mojo Orchestrator: SLURM Resource Setup[/bold cyan]"
        )
        console.print(f"\t            [dim]Python:[/dim] {sys.executable}")
        console.print(f"\t              [dim]Root:[/dim] {project_root}")
        console.print(f"\t[dim]mujoco-mojo binary:[/dim] {mojo_cmd}\n")

        # Standard colors only for Rich compatibility
        job_name = Prompt.ask(
            "  [white]Job Name[/]", default=f"mojo-sim-{status_tracker.id}"
        )
        if available_partitions:
            # Use the actual SLURM default if we found one, otherwise the first in list
            initial_default = (
                default_partition if default_partition else available_partitions[0]
            )

            partition = Prompt.ask(
                "  [white]Partition[/]",
                choices=available_partitions,
                default=initial_default,
            )
        else:
            partition = Prompt.ask(
                "  [white]Partition[/] [dim](optional)[/]", default=""
            )

        # === get partition limits ===
        cpu_limit = self.get_slurm_cpu_limit(partition)
        mem_limit = self.get_slurm_node_mem_limit(partition)
        time_limit_hint = self.get_slurm_time_limit(partition)

        # === get cpus per task ===
        cpus_per_task = Prompt.ask(
            f"  [white]CPUs per task[/] [dim](Node Limit: {cpu_limit})[/]",
            default="1",
        )
        cpus_per_task = max([1, int(cpus_per_task)])
        if cpu_limit != "<UNKNOWN>" and cpus_per_task > int(cpu_limit):
            console.print(
                f"\n[bold red]WARNING:[/] Requested CPUs ({cpus_per_task}) exceeds "
                f"the physical node limit ({cpu_limit})."
            )
            if not Confirm.ask("Do you want to proceed anyway?", default=False):
                return True

        # === get memory ===
        mem_per_node = Prompt.ask(
            f"  [white]Memory per node[/] (e.g., 256M) [dim](Node Limit: {mem_limit})[/]",
            default=mem_limit,
        )
        if self.normalize_to_mb(mem_limit) > 0 and self.normalize_to_mb(
            mem_per_node
        ) > self.normalize_to_mb(mem_limit):
            console.print(
                f"\n[bold red]WARNING:[/] Requested memory ({mem_per_node}) exceeds "
                f"the partition node limit ({mem_limit})."
            )
            console.print("[red]This job will likely be rejected by SLURM.[/]\n")
            if not Confirm.ask("Do you want to proceed anyway?", default=False):
                return True

        # === get time ===
        time_limit = Prompt.ask(
            f"  [white]Time limit[/] (HH:MM:SS) [dim](Partition Limit: {time_limit_hint})[/]",
            default="01:00:00",
        )
        requested_seconds = self.slurm_time_to_seconds(time_limit)
        max_seconds = self.slurm_time_to_seconds(time_limit_hint)
        if requested_seconds > max_seconds and max_seconds != -1:  # -1 means infinite
            console.print(
                f"\n[bold red]WARNING:[/] Requested time ({time_limit}) exceeds "
                f"partition MaxTime ({time_limit_hint})."
            )
            if not Confirm.ask("Proceed anyway?", default=False):
                return True

        # === get concurrency throttle ===
        max_concurrent = Prompt.ask(
            "  [white]Max concurrent tasks[/] [dim](blank = no limit, let SLURM decide)[/]",
            default="",
        )
        if max_concurrent and not max_concurrent.isdigit():
            console.print(
                f"\n[bold red]WARNING:[/] '{max_concurrent}' is not a number. Ignoring."
            )
            max_concurrent = ""

        current_pythonpath = os.getenv("PYTHONPATH", "")
        if str(project_root) not in current_pythonpath:
            console.print(
                f"\n[yellow]Warning:[/] Project root [italic]{project_root}[/] is not in your PYTHONPATH."
            )
            if Confirm.ask(
                "Should Mojo automatically include it in the SLURM submission?",
                default=True,
            ):
                # We'll handle this in the sbatch content generation
                include_root_in_path = True
            else:
                include_root_in_path = False
        else:
            include_root_in_path = True

        partition_line = f"#SBATCH --partition={partition}" if partition else ""
        python_path_line = (
            f"export PYTHONPATH=$PYTHONPATH:{project_root}"
            if include_root_in_path
            else ""
        )

        # generate the .sh script
        max_array = self.get_max_array_size()
        if len(to_run) > max_array:
            console.print(
                f"\n[bold red]ERROR:[/] Trial count ({len(to_run)}) exceeds "
                f"SLURM MaxArraySize ({max_array})."
            )
            if not Confirm.ask("Proceed anyway?", default=False):
                return True
        array_range = self.get_slurm_array_string(to_run)
        if max_concurrent:
            array_range += f"%{max_concurrent}"
        script_path = (self.workdir / "mujoco_mojo_submit.sh").resolve()

        sbatch_content = f"""#!/bin/bash
#SBATCH --job-name={job_name}
#SBATCH --array={array_range}
#SBATCH --output={self.workdir.resolve()}/logs/trial_%a.log
#SBATCH --cpus-per-task={cpus_per_task}
#SBATCH --mem={mem_per_node}
#SBATCH --time={time_limit}
{partition_line}

# Move to the project root so imports work
cd {project_root}
{python_path_line}

# Avoid concurrent array tasks racing to write .pyc files into the shared venv
export PYTHONDONTWRITEBYTECODE=1

# Execute the worker command
{cmd}
"""

        script_path.write_text(sbatch_content, encoding="utf-8")
        logger.info(f"SLURM submission script written to {script_path}")

        # final submission
        if Confirm.ask(
            f"\n[cyan]Submit {len(to_run)} trials to SLURM now?[/]", default=True
        ):
            # automatic submission
            logger.info(f"Submitting {len(to_run)} trials...")
            result = subprocess.run(
                ["sbatch", str(script_path)], capture_output=True, text=True
            )
            if result.returncode == 0:
                job_id_msg = result.stdout.strip()
                # Extract just the numeric ID if possible (e.g. "Submitted batch job 2" -> "2")
                job_id = job_id_msg.split()[-1] if job_id_msg else "UNKNOWN"

                console.print(f"\n[bold green]Success![/] {job_id_msg}")

                # === Monitoring Dashboard ===
                console.print("\n[bold cyan]Monitoring Status:[/bold cyan]")
                console.print(
                    f"  - [white]Check status:[/]       [green]squeue -j {job_id}[/]"
                )
                console.print(
                    f"  - [white]Watch live:[/]         [green]watch -n 1 squeue -j {job_id}[/]"
                )
                console.print(
                    f"  - [white]View first log:[/]     [green]tail -f {self.workdir.resolve()}/logs/trial_0.log[/]"
                )
                console.print(
                    f"  - [white]Cancel all trials:[/]  [green]scancel {job_id}[/]"
                )

                console.print(
                    f"\n[dim]Logs are being written to: {self.workdir.resolve()}/logs/[/]"
                )
                return False
            else:
                logger.error(f"SLURM Submission Failed: {result.stderr}")
                return True
        else:
            # deffered submission
            console.print(
                f"\n[yellow]Orchestration complete.[/] Submit manually with:\n[bold green]sbatch {script_path}[/]"
            )
            return False

MojoRuntime

Bases: Protocol

Definition of a function that executes a generated MojoModel model.

MonteCarloConfig

Bases: BaseConfig

n_proc class-attribute instance-attribute

Python
n_proc: int = DEFAULT_N_PROC

Number of proccesses to allow.

This value is used to determine how many parallel jobs can be run. It is also used for the discovery of trial status. Using a value of 1 will result in the slowest runtime, but highest reliability.

Important

Be a good citizen. Use a reasonable number if you are working on a shared resource. You are a jerk if you use everything.

resume class-attribute instance-attribute

Python
resume: bool = DEFAULT_RESUME

Whether to resume a study if the study already exists.

padding_style property

Python
padding_style: str

This dynamically defines the padding style for trial numbers. This is helpful to ensure the filesystem consistently sorts the trials.

Examples: - Suppose you n_trials is 2000 (and the nominal trial_num is 0) - This method would return 04d - Trial number 0 maps to 0000 - Trial number 123 maps to 0123 - Trial number 1999 will still map to 1999

n_trial class-attribute instance-attribute

Python
n_trial: int = DEFAULT_MC_N_TRIAL

Number of trials to run.

You are able to resume a previous job and modify the number of runs desired by changing this value. A job already in progress will not be dynamically stopped though if you change this value at runtime.

OptimizerConfig

Bases: BaseConfig

n_proc class-attribute instance-attribute

Python
n_proc: int = DEFAULT_N_PROC

Number of proccesses to allow.

This value is used to determine how many parallel jobs can be run. It is also used for the discovery of trial status. Using a value of 1 will result in the slowest runtime, but highest reliability.

Important

Be a good citizen. Use a reasonable number if you are working on a shared resource. You are a jerk if you use everything.

resume class-attribute instance-attribute

Python
resume: bool = DEFAULT_RESUME

Whether to resume a study if the study already exists.

padding_style property

Python
padding_style: str

This dynamically defines the padding style for trial numbers. This is helpful to ensure the filesystem consistently sorts the trials.

Examples: - Suppose you n_trials is 2000 (and the nominal trial_num is 0) - This method would return 04d - Trial number 0 maps to 0000 - Trial number 123 maps to 0123 - Trial number 1999 will still map to 1999

n_trial class-attribute instance-attribute

Python
n_trial: int = DEFAULT_OP_N_TRIAL

Number of trials to run.

study_name class-attribute instance-attribute

Python
study_name: str = DEFAULT_OP_STUDY_NAME

Unique identifier for the Optuna study.

direction instance-attribute

Python
direction: Literal['minimize', 'maximize']

Whether we want to find the lowest or highest objective value.

timeout class-attribute instance-attribute

Python
timeout: float | None = DEFAULT_OP_TIMEOUT

Stop the study after this many seconds, regardless of trial count.

storage class-attribute instance-attribute

Python
storage: str | None = DEFAULT_OP_STORAGE

Database URL (e.g., 'sqlite:///study.db') for multi-node persistence.

sampler class-attribute instance-attribute

Python
sampler: SamplerOptions = DEFAULT_OP_SAMPLER

The search algorithm. TPE is generally best for noisy physics.

evals_per_trial class-attribute instance-attribute

Python
evals_per_trial: int = Field(
    default=DEFAULT_OP_EVALS_PER_TRIAL, ge=1
)

Number of times to run the sim with different seeds per trial. The average score is returned to Optuna.

Reduces 'lucky' trials in noisy physics.

refine_search_factor class-attribute instance-attribute

Python
refine_search_factor: float | None = Field(
    default=DEFAULT_OP_REFINE_SEARCH_FACTOR, gt=0, lt=1
)

If set (e.g., 0.5), and resume is True, the Runner will shrink the search space bounds by this factor around the current best trial to focus on local refinement.

Smaller values will more aggressively refine the search space.

prune_failed_trials class-attribute instance-attribute

Python
prune_failed_trials: bool = DEFAULT_OP_PRUNE_FAILED_TRIALS

Whether to immediately stop trials that violate physical constraints (e.g., Mujoco instability) to save compute time.

Trial dataclass

Python
Trial(
    trial_num: int,
    base_dir: Path,
    xml_name: str,
    model_config_name: str | None,
    padding_style: str,
)

Handles the lifecycle of a single simulation run.

The Trial object is responsible for the 'dirty work' of a Monte Carlo run: - Creating directories - Writing the MJCF XML - Saving the configuration snapshot - Triggering the physics runtime.

trial_num instance-attribute

Python
trial_num: int

Unique identifier for this trial iteration.

base_dir instance-attribute

Python
base_dir: Path

Root directory where all simulation trials are stored.

xml_name instance-attribute

Python
xml_name: str

Filename for the generated MJCF XML (e.g., 'model.xml').

model_config_name instance-attribute

Python
model_config_name: str | None

Filename for the serialized MojoModel configuration (e.g., 'config.json').

padding_style instance-attribute

Python
padding_style: str

Format specifier for directory naming (e.g., '04d').

trial_dir property

Python
trial_dir: Path

The absolute path to this trial's unique workspace.

Example

If base_dir is './sims' and trial_num is 7 with '03d' padding, this returns './sims/trial_007'.

xml_path property

Python
xml_path: Path

The full path to the MJCF XML file for this trial.

model_config_path property

Python
model_config_path: Path

The full path to the JSON configuration file for this trial.

named_value_path property

Python
named_value_path: Path

The full path to the JSON NamedValue file for this trial.

run

Python
run(
    generator: MojoGenerator,
    runtime: MojoRuntime | None,
    seed: int | None,
    overrides: NamedValueDict[NDArray],
    gen_args: list[Any],
    gen_kwargs: dict[str, Any],
    run_args: list[Any],
    run_kwargs: dict[str, Any],
) -> tuple[MojoModel | None, TrialStatus, MjState | None]

Executes the complete simulation pipeline for this trial.

This method coordinates three main phases: 1. Generation: Calls the user-provided generator to build a MojoModel model. 2. Persistence: Creates the workspace and writes the model/config to disk. 3. Execution: Triggers the physics runtime if one is provided.

Parameters:

Name Type Description Default
generator MojoGenerator

Function that returns a MojoModel instance.

required
runtime MojoRuntime | None

Optional function to run the simulation (MuJoCo).

required
seed int | None

Seed to use to define the trial.

required
overrides NamedValueDict[NDArray]

Key-value pairs that override random distributions.

required
gen_args list[Any]

Positional arguments for the generator.

required
gen_kwargs dict[str, Any]

Keyword arguments for the generator.

required
run_args list[Any]

Positional arguments for the runtime.

required
run_kwargs dict[str, Any]

Keyword arguments for the runtime.

required

Returns:

Type Description
tuple[MojoModel | None, TrialStatus, MjState | None]

The MojoModel object for the trial or None if there was a failure prior to generating the MojoModel and the status of the trial.

Source code in src/mujoco_mojo/utils/runner.py
Python
def run(
    self,
    generator: MojoGenerator,
    runtime: MojoRuntime | None,
    seed: int | None,
    overrides: NamedValueDict[NDArray],
    gen_args: list[Any],
    gen_kwargs: dict[str, Any],
    run_args: list[Any],
    run_kwargs: dict[str, Any],
) -> tuple[MojoModel | None, TrialStatus, MjState | None]:
    """
    Executes the complete simulation pipeline for this trial.

    This method coordinates three main phases:
    1.  **Generation**: Calls the user-provided generator to build a `MojoModel` model.
    2.  **Persistence**: Creates the workspace and writes the model/config to disk.
    3.  **Execution**: Triggers the physics runtime if one is provided.

    Args:
        generator: Function that returns a `MojoModel` instance.
        runtime: Optional function to run the simulation (MuJoCo).
        seed: Seed to use to define the trial.
        overrides: Key-value pairs that override random distributions.
        gen_args: Positional arguments for the generator.
        gen_kwargs: Keyword arguments for the generator.
        run_args: Positional arguments for the runtime.
        run_kwargs: Keyword arguments for the runtime.

    Returns:
        The `MojoModel` object for the trial or None if there was a failure prior to generating the MojoModel and the status of the trial.

    """
    status = TrialStatus(trial_num=self.trial_num)
    status._path = self.trial_dir / TRIAL_STATUS_FNAME

    # set up a per-trial log file capturing everything logged during this trial
    self.trial_dir.mkdir(parents=True, exist_ok=True)
    trial_log_handler = get_trial_log_handler(self.trial_dir / "mojo.log")
    root_logger = logging.getLogger()
    root_logger.addHandler(trial_log_handler)

    with status.record_step(step_name="pending"):
        pass

    from mujoco_mojo.runtime.requirements_manager import RequirementSatisfied
    from mujoco_mojo.runtime.runtime_manager import SimulationStopped

    result = None
    state = None
    runtime_manager: RuntimeManager | None = None
    try:
        # 1. Generate
        with status.record_step(step_name="generating"):
            logger.info(f"Generating trial_num={self.trial_num}")
            mojo_model = (
                MojoModel()
                .with_overrides(overrides=overrides)
                .with_seed(seed=seed)
                .with_trial_num(self.trial_num)
            )
            mojo_model = generator(mojo_model, overrides, *gen_args, **gen_kwargs)
            mojo_model._trial_dir = self.trial_dir

            # 2. Setup Workspace & Save Metadata
            logger.info(f"Saving trial_num={self.trial_num} to {self.trial_dir}")

            # bundle assets, this remaps DepPath attributes to point to the shared asset dir
            rel_to_xml = Path(
                os.path.relpath(self.shared_asset_dir, self.trial_dir)
            )
            mojo_model.mjcf.bundle_assets(
                target_dir=self.shared_asset_dir, rel_to_xml=rel_to_xml
            )

            # save XML (with modified DepPath)
            if runtime is None:
                mojo_model.mjcf.write_xml(self.xml_path)
            if self.model_config_name:
                mojo_model.dump_to_path(self.model_config_path)
            self.named_value_path.write_text(mojo_model.named.model_dump_json())

        with status.record_step(step_name="solving"):
            # 3. Execute (if runtime provided)
            if runtime is not None:
                logger.info(f"Executing trial_num={self.trial_num} runtime")
                import mujoco_mojo.runtime as rt

                runtime_manager = rt.RuntimeManager(
                    signal_manager=rt.SignalManager(
                        export_path=self.trial_dir
                        / rt.SignalManager.default_output_name(),
                        unit_system=mojo_model.us,
                    )
                )
                runtime_manager._mojo_model = mojo_model
                state = mojo_model.mjcf.prep_for_sim(
                    self.xml_path, unit_system=mojo_model.us
                )
                result = runtime(
                    mojo_model,
                    runtime_manager,
                    state,
                    *run_args,
                    **run_kwargs,
                )
                status.requirements = runtime_manager.requirement_results
            else:
                logger.info(
                    f"No runtime definition was provided for trial_num={self.trial_num} so MuJoCo will not be run."
                )
                result = mojo_model
                state = None

        # serialize again in case new named values were added during the run
        if self.model_config_name:
            mojo_model.dump_to_path(self.model_config_path)
        self.named_value_path.write_text(mojo_model.named.model_dump_json())

        # clear unpicklable collision managers before returing to avoid multiprocessing serialization erros
        mojo_model.clear_unpickleable_data()

        status.step = "done"
        if status.requirements and not all(r.passed for r in status.requirements):
            status.completion = Completion.FAILURE
        else:
            status.completion = Completion.SUCCESS

    except (BdbQuit, KeyboardInterrupt):
        logger.warning("Quit command detected. Exiting execution...")
        raise
    except RequirementSatisfied:
        # a live requirement ended the trial early as a success: complete
        # normally, with the outcome decided by the requirement results
        status.step = "done"
        if runtime_manager is not None:
            status.requirements = runtime_manager.requirement_results
        if status.requirements and not all(r.passed for r in status.requirements):
            status.completion = Completion.FAILURE
        else:
            status.completion = Completion.SUCCESS
    except SimulationStopped:
        status.step = "done"
        status.completion = Completion.TERMINATED
        if runtime_manager is not None:
            status.requirements = runtime_manager.requirement_results
    except Exception as e:
        status.step = "done"
        status.completion = Completion.ERROR
        logger.exception(
            f"Trial {self.trial_num} failed with the following error: {e}"
        )
    finally:
        status.dump_to_path(status._path)
        root_logger.removeHandler(trial_log_handler)
        trial_log_handler.close()

    return result, status, state

setup_logger

Python
setup_logger(
    level=INFO,
    terminal: bool = True,
    log_file: Path | str | None = Path("mojo.log"),
    max_bytes: int = 10 * 1024 * 1024,
    backup_count: int = 5,
) -> Logger

Shortcut to a sensible MuJoCo Mojo logger using Rich for the terminal.

Source code in src/mujoco_mojo/utils/log.py
Python
def setup_logger(
    level=logging.INFO,
    terminal: bool = True,
    log_file: Path | str | None = Path("mojo.log"),
    max_bytes: int = 10 * 1024 * 1024,  # 10MB
    backup_count: int = 5,
) -> logging.Logger:
    """Shortcut to a sensible MuJoCo Mojo logger using Rich for the terminal."""
    logger = logging.getLogger()
    logger.setLevel(level)

    # avoid double logging by preventing logs from being sent to the root logger
    logger.propagate = False

    if not terminal and log_file is None:
        return logger

    # clear existing handlers
    if logger.hasHandlers():
        logger.handlers.clear()

    # --- RICH HANDLER FOR TERMINAL ---
    if terminal:
        console = Console(theme=mojo_theme)
        rich_handler = RichHandler(
            console=console,
            show_time=True,
            log_time_format="[%Y-%m-%d %H:%M:%S]",
            show_path=True,
            enable_link_path=False,
            # highlighter=None,
            markup=True,  # allows color in the log message itself currently disabled in case an array is logged. Square brackets will break logging.
            rich_tracebacks=True,
        )
        rich_handler.addFilter(TerminalFilter())
        logger.addHandler(rich_handler)

    # --- PLAIN TEXT HANDLER ---
    if log_file is not None:
        log_file = Path(log_file)
        log_file.parent.mkdir(parents=True, exist_ok=True)

        file_format = (
            "%(asctime)s [%(levelname)s] %(pathname)s:%(lineno)d - %(message)s"
        )
        file_handler = RotatingFileHandler(
            log_file, maxBytes=max_bytes, backupCount=backup_count
        )
        file_handler.setFormatter(logging.Formatter(file_format))
        file_handler.addFilter(FileFilter())
        logger.addHandler(file_handler)

    return logger