Merge pull request #3079 from aftersomemath:sysid-pr
PiperOrigin-RevId: 868229512 Change-Id: I790bc08fc8b0745583a2f92d9ee2c5a19ba558ea
This commit is contained in:
Copybara-Service
@@ -0,0 +1,224 @@
|
||||
# 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
|
||||
Reference in New Issue
Block a user