Copybara-Service
3ec09f7296
PiperOrigin-RevId: 868229512 Change-Id: I790bc08fc8b0745583a2f92d9ee2c5a19ba558ea
132 lines
4.4 KiB
Python
132 lines
4.4 KiB
Python
# 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.
|
|
# ==============================================================================
|
|
"""Tests for the model_modifier module."""
|
|
|
|
import mujoco
|
|
from mujoco.sysid._src import model_modifier
|
|
import numpy as np
|
|
|
|
|
|
def test_apply_pgain(arm_spec):
|
|
"""Setting a P gain correctly configures the underlying actuator parameters."""
|
|
actuator_name = "actuator5"
|
|
pgain_value = 74
|
|
|
|
modified_spec = model_modifier.apply_pgain(
|
|
arm_spec, actuator_name, pgain_value
|
|
)
|
|
model = modified_spec.compile()
|
|
|
|
assert model.actuator(actuator_name).gainprm[0] == pgain_value
|
|
assert model.actuator(actuator_name).biasprm[1] == -pgain_value
|
|
|
|
|
|
def test_apply_dgain(arm_spec):
|
|
"""Setting a D gain correctly configures the underlying actuator parameters."""
|
|
actuator_name = "actuator5"
|
|
dgain_value = 1.2
|
|
|
|
modified_spec = model_modifier.apply_dgain(
|
|
arm_spec, actuator_name, dgain_value
|
|
)
|
|
model = modified_spec.compile()
|
|
|
|
assert model.actuator(actuator_name).biasprm[2] == -dgain_value
|
|
|
|
|
|
def test_apply_pdgain(arm_spec):
|
|
"""Setting P and D gains together from a single array configures both correctly."""
|
|
actuator_name = "actuator5"
|
|
pdgain_value = np.array([74, 1.2])
|
|
|
|
modified_spec = model_modifier.apply_pdgain(
|
|
arm_spec, actuator_name, pdgain_value
|
|
)
|
|
model = modified_spec.compile()
|
|
|
|
assert model.actuator(actuator_name).gainprm[0] == pdgain_value[0]
|
|
assert model.actuator(actuator_name).biasprm[1] == -pdgain_value[0]
|
|
assert model.actuator(actuator_name).biasprm[2] == -pdgain_value[1]
|
|
|
|
|
|
def test_apply_body_mass_explicit(arm_spec):
|
|
"""Bodies with inertia defined in XML: changing mass proportionally scales inertia."""
|
|
body_name = "link1"
|
|
model = arm_spec.compile()
|
|
original_mass = model.body(body_name).mass[0]
|
|
original_inertia = model.body(body_name).inertia
|
|
del model
|
|
|
|
scale = 3.3
|
|
new_mass = scale * original_mass
|
|
|
|
modified_spec = model_modifier.apply_body_mass_ipos(
|
|
arm_spec, body_name, mass=new_mass, rot_inertia_scale=True
|
|
)
|
|
model = modified_spec.compile()
|
|
|
|
assert model.body(body_name).mass == new_mass
|
|
np.testing.assert_allclose(
|
|
model.body(body_name).inertia, original_inertia * scale
|
|
)
|
|
|
|
|
|
def test_apply_body_mass_implicit(oscillator_spec):
|
|
"""Bodies with inertia inferred from geoms: changing mass proportionally scales inertia."""
|
|
body_name = "mass"
|
|
model = oscillator_spec.compile()
|
|
original_mass = model.body(body_name).mass[0]
|
|
original_inertia = model.body(body_name).inertia
|
|
del model
|
|
|
|
scale = 0.077
|
|
new_mass = scale * original_mass
|
|
|
|
modified_spec = model_modifier.apply_body_mass_ipos(
|
|
oscillator_spec, body_name, mass=new_mass, rot_inertia_scale=True
|
|
)
|
|
model = modified_spec.compile()
|
|
|
|
assert model.body(body_name).mass == new_mass
|
|
np.testing.assert_allclose(
|
|
model.body(body_name).inertia, original_inertia * scale
|
|
)
|
|
|
|
|
|
def test_remove_visuals(arm_spec):
|
|
"""Stripping visuals removes all textures and materials for faster compilation."""
|
|
cleaned_spec = model_modifier.remove_visuals(arm_spec)
|
|
assert not cleaned_spec.textures
|
|
assert not cleaned_spec.materials
|
|
|
|
|
|
def test_apply_param_modifiers(box_spec, box_params):
|
|
"""The full modifier pipeline applies parameter callbacks and produces an updated model."""
|
|
spec = box_spec.copy()
|
|
original_model = spec.compile()
|
|
original_mass = original_model.body("box").mass[0]
|
|
|
|
# Change box_mass parameter.
|
|
box_params["box_mass"].update_from_vector(np.array([5.3]))
|
|
|
|
modified_model = model_modifier.apply_param_modifiers(box_params, spec)
|
|
assert modified_model.body("box").mass[0] != original_mass
|
|
np.testing.assert_allclose(modified_model.body("box").mass[0], 5.3, atol=1e-6)
|
|
|
|
# Also verify apply_param_modifiers_spec returns MjSpec.
|
|
spec2 = box_spec.copy()
|
|
result = model_modifier.apply_param_modifiers_spec(box_params, spec2)
|
|
assert isinstance(result, mujoco.MjSpec)
|