Skip to content

visualization

Traceable

Bases: Protocol

Anything with a world position that can be followed by a Tracer, e.g. a Body, Site, or Geom.

Returned by: Used by:

ArrowConfig dataclass

Python
ArrowConfig(
    pos: Vec3,
    vec: Vec3,
    color: Vec4,
    is_torque: bool,
    length_scale: float = 1.0,
    width_scale: float = 1.0,
)
Returned by: Used by:

length_scale class-attribute instance-attribute

Python
length_scale: float = 1.0

Extra multiplier applied on top of MuJoCo's native length scaling, for stretching or shrinking this arrow's length independently of the others.

width_scale class-attribute instance-attribute

Python
width_scale: float = 1.0

Extra multiplier applied on top of MuJoCo's native width scaling, for thickening or thinning this arrow's shaft independently of the others.

resolve_arrow_coords

Python
resolve_arrow_coords(
    mj_model: MjModel,
) -> tuple[Vec3, Vec3, float]

Calculates the 'from' and 'to' points and width for an arrow.

Source code in src/mujoco_mojo/visualization.py
Python
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
def resolve_arrow_coords(
    self, mj_model: mujoco.MjModel
) -> tuple[Vec3, Vec3, float]:
    """Calculates the 'from' and 'to' points and width for an arrow."""
    # calculate native scaling
    v_map = mj_model.vis.map
    v_scale = mj_model.vis.scale
    stat = mj_model.stat

    if self.is_torque:
        mag_scale = v_map.torque
        width = v_scale.jointwidth * stat.meansize
    else:
        mag_scale = v_map.force
        width = v_scale.forcewidth * stat.meansize

    # normalize length by mean body mass, then apply the caller's custom length scale
    scaled_vec = (
        self.vec * (mag_scale / max(stat.meanmass, 1e-6)) * self.length_scale
    )
    start = np.asarray(self.pos)
    end = start + scaled_vec

    # apply the caller's custom width scale, then cap it so a short vector doesn't render as
    # a disk (see _MAX_WIDTH_TO_LENGTH_RATIO)
    length = float(np.linalg.norm(scaled_vec))
    width = min(width * self.width_scale, length * self._MAX_WIDTH_TO_LENGTH_RATIO)

    return start, end, width