# 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 = """\
"""
ARM_XML = """\
"""
OSCILLATOR_XML = """\
"""
# ---------------------------------------------------------------------------
# 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