Copybara-Service
3ec09f7296
PiperOrigin-RevId: 868229512 Change-Id: I790bc08fc8b0745583a2f92d9ee2c5a19ba558ea
225 lines
7.4 KiB
Python
225 lines
7.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.
|
|
# ==============================================================================
|
|
"""Shared fixtures for mujoco.sysid tests."""
|
|
|
|
import mujoco
|
|
from mujoco.sysid._src import parameter
|
|
from mujoco.sysid._src import timeseries
|
|
from mujoco.sysid._src.model_modifier import _infer_inertial
|
|
import numpy as np
|
|
import pytest
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Inline model XML strings — no external file dependencies
|
|
# ---------------------------------------------------------------------------
|
|
|
|
BOX_XML = """\
|
|
<mujoco model="box">
|
|
<compiler angle="radian"/>
|
|
<option integrator="implicitfast" timestep="0.001" cone="elliptic">
|
|
<flag contact="enable"/>
|
|
</option>
|
|
<asset>
|
|
<texture type="2d" name="groundplane" builtin="checker" width="300" height="300"/>
|
|
<material name="groundplane" texture="groundplane"/>
|
|
</asset>
|
|
<worldbody>
|
|
<body name="floor">
|
|
<geom name="floor" type="box" size="100 100 0.01" pos="0 0 -0.01"/>
|
|
</body>
|
|
<body name="box" pos="0 0 0.102">
|
|
<freejoint name="box_free"/>
|
|
<inertial pos="0 0 0" mass="5" diaginertia="0.0333 0.0333 0.0333"/>
|
|
<geom name="box" type="box" size="0.1 0.1 0.1"/>
|
|
</body>
|
|
</worldbody>
|
|
<actuator>
|
|
<motor name="push_x" joint="box_free" gear="1 0 0 0 0 0"
|
|
ctrllimited="true" ctrlrange="-20 20"/>
|
|
</actuator>
|
|
<contact>
|
|
<pair name="box_floor" geom1="box" geom2="floor"
|
|
friction="1.6 0.005 0.0001" solref="0.01 1.0"/>
|
|
</contact>
|
|
</mujoco>
|
|
"""
|
|
|
|
ARM_XML = """\
|
|
<mujoco model="test_arm">
|
|
<compiler angle="radian" autolimits="true"/>
|
|
<option integrator="implicitfast">
|
|
<flag contact="disable"/>
|
|
</option>
|
|
<asset>
|
|
<texture name="tex1" type="2d" builtin="checker" width="64" height="64"/>
|
|
<material name="mat1" texture="tex1"/>
|
|
</asset>
|
|
<worldbody>
|
|
<body name="link1" pos="0 0 0.1">
|
|
<inertial pos="0 0 0.05" mass="1.0" diaginertia="0.01 0.01 0.005"/>
|
|
<joint name="joint1" type="hinge" axis="0 0 1" range="-3.14 3.14"/>
|
|
<geom type="capsule" fromto="0 0 0 0 0 0.1" size="0.04"/>
|
|
<body name="link2" pos="0 0 0.1">
|
|
<inertial pos="0 0 0.05" mass="0.8" diaginertia="0.008 0.008 0.004"/>
|
|
<joint name="joint2" type="hinge" axis="0 1 0" range="-3.14 3.14"/>
|
|
<geom type="capsule" fromto="0 0 0 0 0 0.1" size="0.035"/>
|
|
<body name="link3" pos="0 0 0.1">
|
|
<inertial pos="0 0 0.05" mass="0.6" diaginertia="0.006 0.006 0.003"/>
|
|
<joint name="joint3" type="hinge" axis="0 1 0" range="-3.14 3.14"/>
|
|
<geom type="capsule" fromto="0 0 0 0 0 0.1" size="0.03"/>
|
|
<body name="link4" pos="0 0 0.1">
|
|
<inertial pos="0 0 0.04" mass="0.4" diaginertia="0.004 0.004 0.002"/>
|
|
<joint name="joint4" type="hinge" axis="0 0 1" range="-3.14 3.14"/>
|
|
<geom type="capsule" fromto="0 0 0 0 0 0.08" size="0.025"/>
|
|
<body name="link5" pos="0 0 0.08">
|
|
<inertial pos="0 0 0.03" mass="0.2" diaginertia="0.002 0.002 0.001"/>
|
|
<joint name="joint5" type="hinge" axis="0 1 0" range="-3.14 3.14"/>
|
|
<geom type="capsule" fromto="0 0 0 0 0 0.06" size="0.02"/>
|
|
</body>
|
|
</body>
|
|
</body>
|
|
</body>
|
|
</body>
|
|
</worldbody>
|
|
<actuator>
|
|
<position name="actuator1" joint="joint1" kp="400" kv="40"/>
|
|
<position name="actuator2" joint="joint2" kp="400" kv="40"/>
|
|
<position name="actuator3" joint="joint3" kp="400" kv="40"/>
|
|
<position name="actuator4" joint="joint4" kp="400" kv="40"/>
|
|
<position name="actuator5" joint="joint5" kp="200" kv="20"/>
|
|
</actuator>
|
|
<sensor>
|
|
<jointpos name="joint1_pos" joint="joint1"/>
|
|
<jointpos name="joint2_pos" joint="joint2"/>
|
|
<jointpos name="joint3_pos" joint="joint3"/>
|
|
<jointpos name="joint4_pos" joint="joint4"/>
|
|
<jointpos name="joint5_pos" joint="joint5"/>
|
|
<jointvel name="joint1_vel" joint="joint1"/>
|
|
<jointvel name="joint2_vel" joint="joint2"/>
|
|
<jointvel name="joint3_vel" joint="joint3"/>
|
|
<jointvel name="joint4_vel" joint="joint4"/>
|
|
<jointvel name="joint5_vel" joint="joint5"/>
|
|
<jointactuatorfrc name="joint1_torque" joint="joint1"/>
|
|
<jointactuatorfrc name="joint2_torque" joint="joint2"/>
|
|
<jointactuatorfrc name="joint3_torque" joint="joint3"/>
|
|
<jointactuatorfrc name="joint4_torque" joint="joint4"/>
|
|
<jointactuatorfrc name="joint5_torque" joint="joint5"/>
|
|
</sensor>
|
|
</mujoco>
|
|
"""
|
|
|
|
OSCILLATOR_XML = """\
|
|
<mujoco>
|
|
<compiler autolimits="true"/>
|
|
<option>
|
|
<flag gravity="disable" contact="disable" limit="disable"/>
|
|
</option>
|
|
<worldbody>
|
|
<body name="mass">
|
|
<joint name="j1" axis="1 0 0" type="slide" range="0 0.5"
|
|
stiffness="10" damping="1"/>
|
|
<geom type="box" size=".1 .1 .1" mass="0.1"/>
|
|
</body>
|
|
</worldbody>
|
|
</mujoco>
|
|
"""
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Fixtures
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.fixture
|
|
def box_spec() -> mujoco.MjSpec:
|
|
return mujoco.MjSpec.from_string(BOX_XML)
|
|
|
|
|
|
@pytest.fixture
|
|
def box_model(box_spec) -> mujoco.MjModel:
|
|
return box_spec.compile()
|
|
|
|
|
|
@pytest.fixture
|
|
def arm_spec() -> mujoco.MjSpec:
|
|
"""Minimal 5-joint arm with sensors, actuators, textures/materials."""
|
|
return mujoco.MjSpec.from_string(ARM_XML)
|
|
|
|
|
|
@pytest.fixture
|
|
def arm_model(arm_spec) -> mujoco.MjSpec:
|
|
return arm_spec.compile()
|
|
|
|
|
|
@pytest.fixture
|
|
def oscillator_spec() -> mujoco.MjSpec:
|
|
"""Single-body oscillator with implicit (geom-based) inertia."""
|
|
return mujoco.MjSpec.from_string(OSCILLATOR_XML)
|
|
|
|
|
|
@pytest.fixture
|
|
def simple_timeseries() -> timeseries.TimeSeries:
|
|
"""A TimeSeries with 5 data points and 2 columns: y = [x^2, 2*x^2]."""
|
|
times = np.array([0.0, 1.0, 2.0, 3.0, 4.0])
|
|
data = np.array([
|
|
[0.0, 0.0],
|
|
[1.0, 2.0],
|
|
[4.0, 8.0],
|
|
[9.0, 18.0],
|
|
[16.0, 32.0],
|
|
])
|
|
return timeseries.TimeSeries(times=times, data=data)
|
|
|
|
|
|
@pytest.fixture
|
|
def box_params(box_spec) -> parameter.ParameterDict:
|
|
"""ParameterDict for box model with modifier callbacks."""
|
|
del box_spec
|
|
pdict = parameter.ParameterDict()
|
|
pdict.add(
|
|
parameter.Parameter(
|
|
"box_mass",
|
|
[5],
|
|
min_value=[4.5],
|
|
max_value=[5.5],
|
|
modifier=lambda s, p: setattr(
|
|
_infer_inertial(s, "box"), "mass", p.value[0]
|
|
),
|
|
)
|
|
)
|
|
pdict.add(
|
|
parameter.Parameter(
|
|
"solref1",
|
|
[0.01],
|
|
min_value=[0.002],
|
|
max_value=[0.02],
|
|
modifier=lambda s, p: s.pair("box_floor").solref.__setitem__(
|
|
0, p.value[0]
|
|
),
|
|
)
|
|
)
|
|
pdict.add(
|
|
parameter.Parameter(
|
|
"friction2",
|
|
[0.005],
|
|
min_value=[0],
|
|
max_value=[0.01],
|
|
modifier=lambda s, p: s.pair("box_floor").friction.__setitem__(
|
|
1, p.value[0]
|
|
),
|
|
)
|
|
)
|
|
return pdict
|