Loading files
314965a3cThis file is pinned to the selected release. Download the original to open it in your design tool.
# Copyright 2026 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Model modifiers."""
from typing import Any
import mujoco
from mujoco.sysid._src.parameter import InertiaType
from mujoco.sysid._src.parameter import ModifierFn
from mujoco.sysid._src.parameter import Parameter
from mujoco.sysid._src.parameter import ParameterDict
import numpy as np
def remove_visuals(in_spec: mujoco.MjSpec) -> mujoco.MjSpec:
"""Remove visual elements from a Spec."""
spec = in_spec.copy()
all_geoms = spec.worldbody.find_all("geom")
for geom in all_geoms:
if geom.contype == 0 and geom.conaffinity == 0:
if geom.type == mujoco.mjtGeom.mjGEOM_MESH and geom.meshname:
meshname = geom.meshname
mesh = spec.mesh(meshname)
if mesh: # multiple geoms can ref same mesh.
spec.delete(mesh)
spec.delete(geom)
for mat in spec.materials:
spec.delete(mat)
for tex in spec.textures:
spec.delete(tex)
spec.compile() # TODO(b/0): is this compile necessary?
return spec
def _get_obj_or_raise(spec: mujoco.MjSpec, obj_type: str, obj_name: str) -> Any:
getter = getattr(spec, obj_type, None)
if not callable(getter):
raise AttributeError(f"MjSpec has no method '{obj_type}'")
obj = getter(obj_name)
if obj is None:
raise ValueError(f"{obj_type.capitalize()} '{obj_name}' not found in spec.")
return obj
def apply_param_modifiers_spec(
params: ParameterDict, spec: mujoco.MjSpec
) -> mujoco.MjSpec:
for key in params.keys():
param = params[key]
if not param.frozen:
param.apply_modifier(spec)
return spec
def apply_param_modifiers(
params: ParameterDict, spec: mujoco.MjSpec
) -> mujoco.MjModel:
return apply_param_modifiers_spec(params, spec).compile()
def _infer_inertial(spec: mujoco.MjSpec, body_name: str) -> mujoco.MjsBody:
"""Override spec inertia using inferred inertia from compiled model."""
body = _get_obj_or_raise(spec, "body", body_name)
assert isinstance(body, mujoco.MjsBody)
spec.compiler.inertiafromgeom = 2
model = spec.compile()
body.explicitinertial = True
body.fullinertia = np.full((6, 1), np.nan)
body.mass = model.body(body_name).mass[0]
body.inertia = model.body(body_name).inertia
body.ipos = model.body(body_name).ipos
body.iquat = model.body(body_name).iquat
return body
def is_position_actuator(actuator) -> bool:
"""Check if an actuator is a position actuator.
This function works on both model.actuator and spec.actuator objects.
Args:
actuator: An actuator object from model or spec.
Returns:
True if the actuator is a position actuator.
"""
return (
actuator.gaintype == mujoco.mjtGain.mjGAIN_FIXED
and actuator.biastype == mujoco.mjtBias.mjBIAS_AFFINE
and actuator.dyntype
in (mujoco.mjtDyn.mjDYN_NONE, mujoco.mjtDyn.mjDYN_FILTEREXACT)
and actuator.gainprm[0] == -actuator.biasprm[1]
)
def get_actuator_pd_gains(
model: mujoco.MjModel, actuator_name: str
) -> tuple[float, float]:
"""Return the (P, D) gains of a position actuator.
Args:
model: MuJoCo model.
actuator_name: Name of the actuator.
"""
actuator_id = mujoco.mj_name2id(
model, mujoco.mjtObj.mjOBJ_ACTUATOR.value, actuator_name
)
if actuator_id == -1:
raise ValueError(f"Actuator {actuator_name} not found in model.")
actuator = model.actuator(actuator_id)
if not is_position_actuator(actuator):
raise ValueError(f"Actuator {actuator_name} is not a position actuator.")
return -actuator.biasprm[1], -actuator.biasprm[2]
def apply_pgain(
spec: mujoco.MjSpec,
actuator_name: str,
value: float | np.ndarray,
) -> mujoco.MjSpec:
"""Set the proportional gain for a position actuator.
Args:
spec: MuJoCo model specification.
actuator_name: Name of the actuator.
value: Proportional gain value.
"""
# TODO(b/0): assert scalar
actuator = _get_obj_or_raise(spec, "actuator", actuator_name)
assert isinstance(actuator, mujoco.MjsActuator)
if not is_position_actuator(actuator):
raise ValueError(f"Actuator {actuator_name} is not a position actuator.")
actuator.gainprm[0] = value
actuator.biasprm[1] = -value
return spec
def apply_dgain(
spec: mujoco.MjSpec,
actuator_name: str,
value: float | np.ndarray,
) -> mujoco.MjSpec:
"""Set the derivative gain for a position actuator.
Args:
spec: MuJoCo model specification.
actuator_name: Name of the actuator.
value: Derivative gain value.
"""
# TODO(b/0): assert scalar
actuator = _get_obj_or_raise(spec, "actuator", actuator_name)
assert isinstance(actuator, mujoco.MjsActuator)
if not is_position_actuator(actuator):
raise ValueError(f"Actuator {actuator_name} is not a position actuator.")
actuator.biasprm[2] = -value
return spec
def apply_pdgain(
spec: mujoco.MjSpec,
actuator_name: str,
value: np.ndarray,
) -> mujoco.MjSpec:
"""Set both proportional and derivative gains for a position actuator.
Args:
spec: MuJoCo model specification.
actuator_name: Name of the actuator.
value: 2-element array ``[P_gain, D_gain]``.
"""
if value.size != 2:
raise ValueError(f"pdgain must be a 2-element array, got {value.size}.")
apply_pgain(spec, actuator_name, value[0])
apply_dgain(spec, actuator_name, value[1])
return spec
def apply_body_mass_ipos(
spec: mujoco.MjSpec,
body_name: str,
mass: np.ndarray | None = None,
ipos: np.ndarray | None = None,
rot_inertia_scale: bool = False,
) -> mujoco.MjSpec:
"""Apply mass and center-of-mass position to a body.
Args:
spec: MuJoCo model specification.
body_name: Name of the body.
mass: Optional new mass value.
ipos: Optional new center-of-mass position.
rot_inertia_scale: If True, scale rotational inertia proportionally to
mass change.
"""
# TODO(b/0): assert mass and ipos shapes
body = _infer_inertial(spec, body_name)
mass_original = body.mass
if mass is not None:
body.mass = mass
if rot_inertia_scale:
scale = mass / mass_original
body.inertia *= scale
if ipos is not None:
body.ipos = ipos
return spec
def scale_body_inertia(
spec: mujoco.MjSpec,
body_name: str,
value: np.ndarray,
) -> mujoco.MjSpec:
# TODO(b/0): assert scalar
body = _infer_inertial(spec, body_name)
body.inertia *= value
return spec
def pi_from_theta(theta: np.ndarray) -> np.ndarray:
"""Convert base parameters θ to inertial parameters π.
Args:
theta: 10-D array [alpha, d1, d2, d3, s12, s23, s13, t1, t2, t3] where:
alpha: Scale parameter (log of U[3,3])
[d1, d2, d3]: Log of diagonal elements
[s12, s23, s13]: Shear parameters from upper triangle
[t1, t2, t3]: Translation parameters from last column
Returns:
10-D array π = [m, hx, hy, hz, Ixx, Iyy, Izz, Ixy, Iyz, Ixz].
"""
alpha, d1, d2, d3, s12, s23, s13, t1, t2, t3 = theta
exp_alpha = np.exp(alpha)
exp_d1 = np.exp(d1)
exp_d2 = np.exp(d2)
exp_d3 = np.exp(d3)
U = np.zeros((4, 4))
U[0, 0] = exp_d1
U[0, 1] = s12
U[0, 2] = s13
U[0, 3] = t1
U[1, 1] = exp_d2
U[1, 2] = s23
U[1, 3] = t2
U[2, 2] = exp_d3
U[2, 3] = t3
U[3, 3] = 1
U *= exp_alpha
J = U @ U.T
sigma = J[:3, :3]
I_bar = np.trace(sigma) * np.eye(3) - sigma
h = J[:3, 3]
m = J[3, 3]
return np.concatenate(([m], h, I_bar.flatten()))
def pseudoinertia_from_pi(pi: np.ndarray) -> np.ndarray:
"""Converts inertial parameters π to a 4x4 pseudoinertia matrix J.
Args:
pi: A 10-D array [m, hx, hy, hz, Ixx, Iyy, Izz, Ixy, Iyz, Ixz].
Returns:
A 4x4 pseudoinertia matrix J.
"""
m = pi[0]
h = pi[1:4]
I_bar = pi[4:].reshape((3, 3))
Sigma = 0.5 * np.trace(I_bar) * np.eye(3) - I_bar
J = np.zeros((4, 4))
J[:3, :3] = Sigma
J[:3, 3] = h
J[3, :3] = h
J[3, 3] = m
return J
def cholesky_decompose_upper(J: np.ndarray) -> np.ndarray:
"""Perform an upper-triangular Cholesky decomposition of J.
The returned matrix U is such that J = U @ U.T.
Args:
J: A 4x4 positive-definite matrix.
Returns:
An upper-triangular matrix U.
"""
n = J.shape[0]
indices = np.arange(n - 1, -1, -1)
J_reversed = J[indices][:, indices]
L_prime = np.linalg.cholesky(J_reversed)
return L_prime[indices][:, indices]
def theta_from_pseudoinertia(J: np.ndarray) -> np.ndarray:
"""Extract the 10-D vector of base parameters θ from the pseudoinertia J.
Args:
J: A 4x4 pseudoinertia.
Returns:
A 10-D array θ = [alpha, d1, d2, d3, s12, s23, s13, t1, t2, t3] where:
alpha: Scale parameter (log of U[3,3])
[d1, d2, d3]: Log of diagonal elements
[s12, s23, s13]: Shear parameters from upper triangle
[t1, t2, t3]: Translation parameters from last column
"""
# U: A 4x4 upper-triangular matrix from Cholesky decomposition
U = cholesky_decompose_upper(J)
# Extract exp(α) from the bottom-right element of U.
exp_alpha = U[3, 3]
alpha = np.log(exp_alpha)
# Compute the d parameters from the diagonal entries (adjusted by alpha).
d1 = np.log(U[0, 0] / exp_alpha)
d2 = np.log(U[1, 1] / exp_alpha)
d3 = np.log(U[2, 2] / exp_alpha)
# Extract the shear parameters (off-diagonals in the upper triangle).
s12 = U[0, 1] / exp_alpha
s13 = U[0, 2] / exp_alpha
s23 = U[1, 2] / exp_alpha
# Extract the translation parameters (last column, except the bottom element).
t1 = U[0, 3] / exp_alpha
t2 = U[1, 3] / exp_alpha
t3 = U[2, 3] / exp_alpha
return np.array([alpha, d1, d2, d3, s12, s23, s13, t1, t2, t3])
def skew(v: np.ndarray) -> np.ndarray:
"""Skew-symmetric matrix from a length-3 vector."""
return np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]])
def inertia_to_fullinertia(q: np.ndarray, inertia: np.ndarray) -> np.ndarray:
"""Convert diagonal inertia with quaternion to full inertia matrix."""
xmat = np.empty(9)
mujoco.mju_quat2Mat(xmat, q)
R = xmat.reshape(3, 3)
return R @ np.diag(inertia) @ R.T
def pi_from_body(spec: mujoco.MjSpec, body_name: str) -> np.ndarray:
"""Extracts the 10-D vector of inertial parameters π from a MuJoCo body.
Args:
spec: MuJoCo model specification object.
body_name: Name of the body to extract parameters from.
Returns:
A 10-D numpy array π = [m, hx, hy, hz, Ixx, Iyy, Izz, Ixy, Iyz, Ixz] where:
m: Mass of the body
[hx, hy, hz]: First moment of mass (m * com, where com is center of mass)
[Ixx, Iyy, Izz, Ixy, Iyz, Ixz]: Rotational inertia about the origin of the
body-fixed reference frame.
"""
body = _infer_inertial(spec, body_name)
mass = body.mass
ipos = body.ipos
inertia = body.inertia
iquat = body.iquat
fullinertia = inertia_to_fullinertia(iquat, inertia)
# Transform inertial from ipos origin to body origin.
I_bar = fullinertia - (mass * skew(ipos) @ skew(ipos))
return np.concatenate([[mass], mass * ipos, I_bar.flatten()])
def theta_inertia_from_body(spec: mujoco.MjSpec, body_name: str) -> np.ndarray:
"""Extract base parameters θ from a body's inertial properties."""
pi = pi_from_body(spec, body_name)
J = pseudoinertia_from_pi(pi)
return theta_from_pseudoinertia(J)
def apply_body_theta_inertia(
spec: mujoco.MjSpec,
body_name: str,
theta: np.ndarray,
) -> mujoco.MjSpec:
"""Apply base-parameter inertia θ to a body in the spec.
Args:
spec: MuJoCo model specification.
body_name: Name of the body.
theta: 10-element array [alpha, d1, d2, d3, s12, s23, s13, t1, t2, t3].
"""
if theta.size != 10:
raise ValueError(f"theta must be a 10-element array, got {theta.size}.")
pi = pi_from_theta(theta)
body = _infer_inertial(spec, body_name)
body.mass = pi[0]
body.ipos = pi[1:4] / pi[0]
# This tells the compiler to ignore the diagonal inertia and instead
# calculate it from the full inertia.
body.inertia[:] = 0.0
body.iquat[:] = np.nan
I_bar = pi[4:].reshape((3, 3))
skew_ipos = skew(body.ipos)
fullinertia = I_bar + (body.mass * skew_ipos @ skew_ipos)
# MuJoCo's ordering is: M(1,1), M(2,2), M(3,3), M(1,2), M(1,3), M(2,3) which
# corresponds to Ixx, Iyy, Izz, Ixy, Ixz
body.fullinertia[0] = fullinertia[0, 0] # Ixx
body.fullinertia[1] = fullinertia[1, 1] # Iyy
body.fullinertia[2] = fullinertia[2, 2] # Izz
body.fullinertia[3] = fullinertia[0, 1] # Ixy
body.fullinertia[4] = fullinertia[0, 2] # Ixz
body.fullinertia[5] = fullinertia[1, 2] # Iyz
return spec
def apply_body_inertia(spec: mujoco.MjSpec, name: str, param: Parameter):
"""Apply inertia parameters to a body based on the parameter type.
Args:
spec: MuJoCo model specification.
name: Name of the body.
param: Parameter with an ``inertia_type`` attribute.
"""
if not hasattr(param, "inertia_type"):
raise ValueError(
f"Parameter {param.name} does not have inertia_type attribute."
)
if param.inertia_type == InertiaType.Mass:
apply_body_mass_ipos(
spec, name, mass=param.value, rot_inertia_scale=param.scale_rot_inertia
)
elif param.inertia_type == InertiaType.MassIpos:
apply_body_mass_ipos(
spec,
name,
mass=param.value[0],
ipos=param.value[1:4],
rot_inertia_scale=param.scale_rot_inertia,
)
elif param.inertia_type == InertiaType.Pseudo:
apply_body_theta_inertia(spec, name, param.value)
def body_inertia_param(
spec: mujoco.MjSpec,
model: mujoco.MjModel,
body_name: str,
inertia_type: InertiaType = InertiaType.MassIpos,
scale_rot_inertia: bool = False,
mass_bound_mult: np.ndarray | None = None,
ipos_bound_off: np.ndarray | None = None,
stretch_bound_mult: np.ndarray | None = None,
shear_bound_off: np.ndarray | None = None,
param_name: str | None = None,
modifier: ModifierFn | None = None,
) -> Parameter:
"""Creates Parameter objects for the inertia of a body in a simplified manner.
Args:
spec: MuJoCo model specification object.
model: The MuJoCo model.
body_name: Name of the body to create the parameter for.
inertia_type: The type of inertia parameterization to use.
scale_rot_inertia: Whether to scale the original inertia when mass changes,
ignored with pseudo inertia.
mass_bound_mult: Multiplicative bounds for the mass parameter.
ipos_bound_off: Additive bounds for the ipos parameter.
stretch_bound_mult: Multiplicative bounds for the stretch parameters in the
pseudo-inertia parameterization.
shear_bound_off: Additive bounds for the shear parameters in the
pseudo-inertia parameterization.
param_name: Optional name for the parameter. Defaults to
``"{body_name}_inertia"``.
modifier: Optional custom modifier callback. If None, the default
:func:`apply_body_inertia` modifier is registered on the Parameter.
Returns:
A Parameter configured for the body's inertia.
"""
if mass_bound_mult is None:
mass_bound_mult = np.array([0.1, 10.0])
if ipos_bound_off is None:
ipos_bound_off = np.array([-0.5, 0.5])
if stretch_bound_mult is None:
stretch_bound_mult = np.array([0.5, 2.0])
if shear_bound_off is None:
shear_bound_off = np.array([-0.5, 0.5])
body = model.body(body_name)
if param_name is None:
param_name = f"{body_name}_inertia"
if modifier is None:
def _default_modifier(spec, param):
return apply_body_inertia(spec, body_name, param)
modifier = _default_modifier
if inertia_type == InertiaType.Mass:
param = Parameter(
param_name,
body.mass,
body.mass * mass_bound_mult[0],
body.mass * mass_bound_mult[1],
modifier=modifier,
)
param.inertia_type = inertia_type
param.scale_rot_inertia = scale_rot_inertia
elif inertia_type == InertiaType.MassIpos:
massipos0 = np.concatenate((body.mass, body.ipos))
massipos_low = np.concatenate(
(body.mass * mass_bound_mult[0], body.ipos + ipos_bound_off[0])
)
massipos_high = np.concatenate(
(body.mass * mass_bound_mult[1], body.ipos + ipos_bound_off[1])
)
param = Parameter(
param_name, massipos0, massipos_low, massipos_high, modifier=modifier
)
param.inertia_type = inertia_type
param.scale_rot_inertia = scale_rot_inertia
elif inertia_type == InertiaType.Pseudo:
theta_i_0 = theta_inertia_from_body(spec, body_name)
# mass = exp(2*alpha)
alpha = theta_i_0[0]
mass = np.exp(2 * alpha)
mass_bounds = mass * mass_bound_mult
alpha_bounds = 0.5 * np.log(mass_bounds)
# d1, d2, d3, stretch = exp(2*d)
# stretches body along principal axes
d = theta_i_0[1 : 1 + 3]
stretch = np.exp(2 * d)
stretch_bounds = stretch[:, np.newaxis] * np.atleast_2d(stretch_bound_mult)
d_bounds = 0.5 * np.log(stretch_bounds)
# s12, s23, s13
# shear the body
s_bounds = theta_i_0[4 : 4 + 3, np.newaxis] + np.atleast_2d(shear_bound_off)
# t1, t2, t3
# center of mass
t_bounds = theta_i_0[7:10, np.newaxis] + np.atleast_2d(ipos_bound_off)
theta_bounds = np.vstack([
alpha_bounds,
d_bounds,
s_bounds,
t_bounds,
])
param = Parameter(
param_name,
theta_i_0,
theta_bounds[:, 0],
theta_bounds[:, 1],
modifier=modifier,
)
param.inertia_type = inertia_type
else:
raise ValueError(f"Unknown inertia_type: {inertia_type}")
return param