microduck/mjlab
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
32a241c28f218be149318befefa570b0fcc839cf
mjlab/tests/test_domain_randomization.py
T

2470 lines
75 KiB
Python
Raw Normal View History

Import upstream snapshot c19f713c415a699a79d71cd96aa13c3104a05047
2026-08-28 15:42:17 +08:00
"""Tests for the dr.* model-field randomization engine."""
import math
import mujoco
import pytest
import torch
from conftest import get_test_device
from mjlab.entity import EntityCfg
from mjlab.envs.mdp import dr
from mjlab.managers.scene_entity_config import SceneEntityCfg
from mjlab.scene import Scene, SceneCfg
from mjlab.sim.sim import Simulation, SimulationCfg
from mjlab.utils.lab_api.math import quat_from_euler_xyz, quat_mul
pytestmark = pytest.mark.filterwarnings(
"ignore:Use of index_put_ on expanded tensors is deprecated:UserWarning"
)
# Shared helpers.
ROBOT_XML = """
<mujoco>
<worldbody>
<body name="base" pos="0 0 1">
<freejoint name="free_joint"/>
<geom name="base_geom" type="box" size="0.1 0.1 0.1" mass="1.0"
friction="0.5 0.01 0.005"/>
<site name="base_site" pos="0 0 0.1"/>
<camera name="front_cam" pos="0.5 0 0.3" xyaxes="0 -1 0 0 0 1" fovy="60"/>
<camera name="side_cam" pos="0 0.5 0.3" xyaxes="1 0 0 0 0 1" fovy="45"/>
<light name="top_light" pos="0 0 2" dir="0 0 -1"/>
<body name="foot1" pos="0.2 0 0">
<joint name="joint1" type="hinge" axis="0 0 1" range="0 1.57"/>
<geom name="foot1_geom" type="box" size="0.05 0.05 0.05" mass="0.1"
friction="0.5 0.01 0.005"/>
</body>
<body name="foot2" pos="-0.2 0 0">
<joint name="joint2" type="hinge" axis="0 0 1" range="0 1.57"/>
<geom name="foot2_geom" type="box" size="0.05 0.05 0.05" mass="0.1"
friction="0.5 0.01 0.005"/>
</body>
<body name="offset_body" pos="0 0.3 0">
<joint name="joint3" type="hinge" axis="1 0 0" range="-1.57 1.57"/>
<inertial pos="0.05 -0.02 0.03" mass="0.5"
fullinertia="0.001 0.002 0.0015 0.0001 -0.00005 0.00008"/>
<geom name="offset_geom" type="sphere" size="0.04"/>
</body>
</body>
</worldbody>
</mujoco>
"""
NUM_ENVS = 4
@pytest.fixture(scope="module")
def device():
return get_test_device()
class Env:
def __init__(self, scene, sim, device):
self.scene = scene
self.sim = sim
self.num_envs = scene.num_envs
self.device = device
def create_test_env(
device, num_envs=NUM_ENVS, expand_fields=("geom_friction", "dof_damping")
):
entity_cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(ROBOT_XML))
scene_cfg = SceneCfg(num_envs=num_envs, entities={"robot": entity_cfg})
scene = Scene(scene_cfg, device)
model = scene.compile()
sim = Simulation(num_envs=num_envs, cfg=SimulationCfg(), model=model, device=device)
scene.initialize(model, sim.model, sim.data)
if expand_fields:
sim.expand_model_fields(expand_fields)
return Env(scene, sim, device)
@pytest.fixture(scope="module")
def env(device):
return create_test_env(device)
# Operations: abs, scale, add.
def test_abs_operation(env):
"""Values set directly within range, diversity across envs."""
torch.manual_seed(42)
robot = env.scene["robot"]
dr.geom_friction(
env,
env_ids=None,
ranges=(0.3, 1.2),
operation="abs",
asset_cfg=SceneEntityCfg("robot", geom_names=(".*",)),
axes=[0],
)
friction = env.sim.model.geom_friction[:, robot.indexing.geom_ids, 0]
assert torch.all((friction >= 0.3) & (friction <= 1.2))
assert len(torch.unique(friction)) >= 2
def test_scale_operation(env):
"""Values = default * sample, within expected bounds."""
torch.manual_seed(42)
robot = env.scene["robot"]
geom_idx = robot.indexing.geom_ids[0]
default_val = env.sim.get_default_field("geom_friction")[geom_idx, 0].item()
dr.geom_friction(
env,
env_ids=None,
ranges=(0.5, 2.0),
operation="scale",
asset_cfg=SceneEntityCfg("robot", geom_ids=[0]),
axes=[0],
)
result = env.sim.model.geom_friction[:, geom_idx, 0]
assert torch.all(
(result >= default_val * 0.5 - 1e-5) & (result <= default_val * 2.0 + 1e-5)
)
def test_add_operation(env):
"""Values = default + sample, within expected bounds."""
torch.manual_seed(42)
robot = env.scene["robot"]
geom_idx = robot.indexing.geom_ids[0]
default_val = env.sim.get_default_field("geom_friction")[geom_idx, 0].item()
dr.geom_friction(
env,
env_ids=None,
ranges=(-0.1, 0.1),
operation="add",
asset_cfg=SceneEntityCfg("robot", geom_ids=[0]),
axes=[0],
)
result = env.sim.model.geom_friction[:, geom_idx, 0]
assert torch.all(
(result >= default_val - 0.1 - 1e-5) & (result <= default_val + 0.1 + 1e-5)
)
# No accumulation (scale and add share the same mechanism).
@pytest.mark.parametrize(
"operation, fixed_val, expected_fn",
[
("scale", (2.0, 2.0), lambda d: d * 2.0),
("add", (0.1, 0.1), lambda d: d + 0.1),
],
)
def test_no_accumulation(device, operation, fixed_val, expected_fn):
"""3x with fixed value produces same result as 1x (uses defaults)."""
env = create_test_env(device, num_envs=2)
robot = env.scene["robot"]
geom_idx = robot.indexing.geom_ids[0]
default_friction = env.sim.get_default_field("geom_friction")[geom_idx, 0].item()
for _ in range(3):
dr.geom_friction(
env,
env_ids=None,
ranges=fixed_val,
operation=operation,
asset_cfg=SceneEntityCfg("robot", geom_ids=[0]),
axes=[0],
)
final = env.sim.model.geom_friction[0, geom_idx, 0].item()
assert abs(final - expected_fn(default_friction)) < 1e-5
# Distributions.
def test_log_uniform_distribution(device):
"""Log-uniform produces values in [lo, hi], skewed toward lower."""
torch.manual_seed(42)
env = create_test_env(device, num_envs=64)
dr.geom_friction(
env,
env_ids=None,
ranges=(0.1, 10.0),
operation="abs",
distribution="log_uniform",
asset_cfg=SceneEntityCfg("robot", geom_ids=[0]),
axes=[0],
)
vals = env.sim.model.geom_friction[:, env.scene["robot"].indexing.geom_ids[0], 0]
assert torch.all((vals >= 0.1 - 1e-5) & (vals <= 10.0 + 1e-5))
# Geometric mean of log-uniform(0.1, 10) is 1.0, well below (0.1+10)/2=5.05.
assert vals.median() < (0.1 + 10.0) / 2
def test_gaussian_distribution(device):
"""Gaussian: mean/std interpretation, values cluster around mean."""
torch.manual_seed(42)
env = create_test_env(device, num_envs=128)
mean, std = 0.5, 0.05
dr.geom_friction(
env,
env_ids=None,
ranges=(mean, std),
operation="abs",
distribution="gaussian",
asset_cfg=SceneEntityCfg("robot", geom_ids=[0]),
axes=[0],
)
vals = env.sim.model.geom_friction[:, env.scene["robot"].indexing.geom_ids[0], 0]
assert abs(vals.mean().item() - mean) < 0.05
# Axes.
@pytest.mark.parametrize(
"axes, changed, unchanged",
[
(None, [0], [1, 2]), # Default axes for geom_friction = [0]
([1, 2], [1, 2], [0]), # Explicit axes
],
)
def test_axes_selectivity(device, axes, changed, unchanged):
"""Only the specified axes are modified; others stay at defaults."""
env = create_test_env(device, num_envs=2)
robot = env.scene["robot"]
geom_idx = robot.indexing.geom_ids[0]
default_friction = env.sim.get_default_field("geom_friction")[geom_idx].clone()
dr.geom_friction(
env,
env_ids=None,
ranges=(2.0, 2.0),
operation="scale",
asset_cfg=SceneEntityCfg("robot", geom_ids=[0]),
axes=axes,
)
final = env.sim.model.geom_friction[0, geom_idx]
for ax in changed:
assert abs(final[ax] - default_friction[ax] * 2.0) < 1e-5
for ax in unchanged:
assert abs(final[ax] - default_friction[ax]) < 1e-5
def test_dict_int_axis_ranges(device):
"""{0: (lo, hi), 1: (lo2, hi2)} per-axis ranges."""
torch.manual_seed(42)
env = create_test_env(device, num_envs=2)
robot = env.scene["robot"]
geom_idx = robot.indexing.geom_ids[0]
dr.geom_friction(
env,
env_ids=None,
ranges={0: (0.3, 0.4), 1: (0.001, 0.002)},
operation="abs",
asset_cfg=SceneEntityCfg("robot", geom_ids=[0]),
)
final = env.sim.model.geom_friction[0, geom_idx]
assert 0.3 - 1e-5 <= final[0].item() <= 0.4 + 1e-5
assert 0.001 - 1e-5 <= final[1].item() <= 0.002 + 1e-5
def test_invalid_axes_raises(env):
"""axes outside valid_axes raises ValueError."""
with pytest.raises(ValueError, match="Invalid axes"):
dr.geom_friction(
env,
env_ids=None,
ranges=(0.3, 1.2),
operation="abs",
asset_cfg=SceneEntityCfg("robot", geom_names=(".*",)),
axes=[3], # geom_friction valid_axes=[0,1,2]
)
# String-keyed ranges.
def test_string_keyed_ranges(env):
"""Per-component DR with regex patterns."""
torch.manual_seed(42)
dr.joint_damping(
env,
env_ids=None,
ranges={".*joint1": (0.5, 0.5), ".*joint2": (1.5, 1.5)},
operation="abs",
asset_cfg=SceneEntityCfg("robot", joint_names=(".*",)),
)
robot = env.scene["robot"]
dof_adr = robot.indexing.joint_v_adr
damping = env.sim.model.dof_damping[0, dof_adr]
assert abs(damping[0].item() - 0.5) < 1e-5
assert abs(damping[1].item() - 1.5) < 1e-5
def test_string_keyed_ranges_no_match_raises(env):
"""Pattern matching no names raises ValueError."""
with pytest.raises(ValueError, match="matched no"):
dr.joint_damping(
env,
env_ids=None,
ranges={"nonexistent_joint": (0.5, 1.5)},
operation="abs",
asset_cfg=SceneEntityCfg("robot", joint_names=(".*",)),
)
# Shared random.
def test_shared_random(env):
"""All entities within same env get same value, envs differ."""
torch.manual_seed(42)
robot = env.scene["robot"]
geom_ids = robot.indexing.geom_ids
dr.geom_friction(
env,
env_ids=None,
ranges=(0.3, 1.2),
operation="abs",
asset_cfg=SceneEntityCfg("robot", geom_names=(".*",)),
axes=[0],
shared_random=True,
)
friction = env.sim.model.geom_friction[:, geom_ids, 0]
for env_idx in range(env.num_envs):
env_friction = friction[env_idx]
assert torch.allclose(env_friction, env_friction[0].expand_as(env_friction))
env_frictions = friction[:, 0]
assert len(torch.unique(env_frictions)) > 1
assert torch.all((friction >= 0.3) & (friction <= 1.2))
# Edge cases.
def test_single_env_without_expand(device):
"""num_envs=1 works without expand_model_fields, no accumulation."""
env = create_test_env(device, num_envs=1, expand_fields=())
robot = env.scene["robot"]
geom_idx = robot.indexing.geom_ids[0]
original_friction = env.sim.model.geom_friction[0, geom_idx, 0].item()
for _ in range(2):
dr.geom_friction(
env,
env_ids=None,
ranges=(2.0, 2.0),
operation="scale",
asset_cfg=SceneEntityCfg("robot", geom_ids=[0]),
axes=[0],
)
final = env.sim.model.geom_friction[0, geom_idx, 0].item()
assert abs(final - original_friction * 2.0) < 1e-5
def test_partial_env_ids(env):
"""Randomizing subset of envs leaves others unchanged."""
torch.manual_seed(42)
robot = env.scene["robot"]
geom_idx = robot.indexing.geom_ids[0]
original = env.sim.model.geom_friction[:, geom_idx, 0].clone()
dr.geom_friction(
env,
env_ids=torch.tensor([0, 2], device=env.device),
ranges=(0.1, 0.2),
operation="abs",
asset_cfg=SceneEntityCfg("robot", geom_ids=[0]),
axes=[0],
)
result = env.sim.model.geom_friction[:, geom_idx, 0]
# Envs 0,2 should have changed.
assert torch.all((result[0] >= 0.1) & (result[0] <= 0.2))
assert torch.all((result[2] >= 0.1) & (result[2] <= 0.2))
# Envs 1,3 should be unchanged.
assert result[1] == original[1]
assert result[3] == original[3]
# CUDA graph.
@pytest.mark.skipif(
not torch.cuda.is_available(), reason="CUDA required for graph capture"
)
def test_expand_model_fields_recreates_cuda_graph(device):
"""Verify CUDA graph is recreated after expand_model_fields."""
entity_cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(ROBOT_XML))
scene_cfg = SceneCfg(num_envs=NUM_ENVS, entities={"robot": entity_cfg})
scene = Scene(scene_cfg, device)
model = scene.compile()
sim = Simulation(num_envs=NUM_ENVS, cfg=SimulationCfg(), model=model, device=device)
scene.initialize(model, sim.model, sim.data)
if not sim.use_cuda_graph:
pytest.skip("CUDA graph capture not enabled on this device")
original_step_graph = sim.step_graph
sim.expand_model_fields(("geom_friction",))
assert sim.step_graph is not original_step_graph, (
"CUDA graph was not recreated after expand_model_fields"
)
# Integration test.
@pytest.mark.slow
def test_g1_foot_friction_shared_across_geoms(device):
"""G1 velocity env has uniform foot friction across all collision geoms."""
import io
import warnings
from contextlib import redirect_stderr, redirect_stdout
from mjlab.envs.manager_based_rl_env import ManagerBasedRlEnv
from mjlab.tasks.velocity.config.g1.env_cfgs import unitree_g1_flat_env_cfg
cfg = unitree_g1_flat_env_cfg()
with warnings.catch_warnings():
warnings.simplefilter("ignore")
with redirect_stdout(io.StringIO()), redirect_stderr(io.StringIO()):
env = ManagerBasedRlEnv(cfg, device=device)
try:
robot = env.scene["robot"]
foot_geom_names = [
f"{side}_foot{i}_collision" for side in ("left", "right") for i in range(1, 8)
]
foot_geom_ids, _ = robot.find_geoms(foot_geom_names)
foot_geom_indices = robot.indexing.geom_ids[foot_geom_ids]
friction = env.sim.model.geom_friction[:, foot_geom_indices, 0]
for env_idx in range(env.num_envs):
env_friction = friction[env_idx]
assert torch.allclose(env_friction, env_friction[0].expand_as(env_friction))
if env.num_envs > 1:
env_frictions = friction[:, 0]
assert len(torch.unique(env_frictions)) > 1
finally:
env.close()
# Quaternion DR tests.
def _make_quat_env(device, num_envs=NUM_ENVS):
"""Create an env with body_quat, geom_quat, site_quat expanded."""
return create_test_env(
device,
num_envs=num_envs,
expand_fields=("body_quat", "geom_quat", "site_quat"),
)
@pytest.fixture(scope="module")
def quat_env(device):
return _make_quat_env(device)
_QUAT_CASES = [
pytest.param(
dr.body_quat,
"body_quat",
SceneEntityCfg("robot", body_names=(".*",)),
"body_ids",
id="body",
),
pytest.param(
dr.geom_quat,
"geom_quat",
SceneEntityCfg("robot", geom_names=(".*",)),
"geom_ids",
id="geom",
),
pytest.param(
dr.site_quat,
"site_quat",
SceneEntityCfg("robot", site_names=(".*",)),
"site_ids",
id="site",
),
]
@pytest.mark.parametrize("dr_func, field, asset_cfg, ids_attr", _QUAT_CASES)
def test_quat_is_unit_quaternion(quat_env, dr_func, field, asset_cfg, ids_attr):
"""All output quaternions must have unit norm (within 1e-6)."""
torch.manual_seed(0)
env = quat_env
robot = env.scene["robot"]
dr_func(
env,
env_ids=None,
roll_range=(-0.5, 0.5),
pitch_range=(-0.5, 0.5),
yaw_range=(-math.pi, math.pi),
asset_cfg=asset_cfg,
)
ids = getattr(robot.indexing, ids_attr)
quats = getattr(env.sim.model, field)[:, ids, :]
norms = quats.norm(dim=-1)
assert torch.all((norms - 1.0).abs() < 1e-6), (
f"Non-unit quaternion found; max deviation: {(norms - 1.0).abs().max()}"
)
@pytest.mark.parametrize("dr_func, field, asset_cfg, ids_attr", _QUAT_CASES)
def test_quat_zero_range_unchanged(quat_env, dr_func, field, asset_cfg, ids_attr):
"""All ranges (0, 0): result equals default quaternion."""
env = quat_env
robot = env.scene["robot"]
ids = getattr(robot.indexing, ids_attr)
default_quat = env.sim.get_default_field(field)[ids].clone()
dr_func(env, env_ids=None, asset_cfg=asset_cfg)
result = getattr(env.sim.model, field)[:, ids, :]
assert torch.allclose(result, default_quat.unsqueeze(0).expand_as(result), atol=1e-6)
@pytest.mark.parametrize("dr_func, field, asset_cfg, ids_attr", _QUAT_CASES)
def test_quat_composes_with_default(quat_env, dr_func, field, asset_cfg, ids_attr):
"""Fixed yaw=pi/4 matches manual quat_mul(q_yaw, q_default)."""
env = quat_env
robot = env.scene["robot"]
ids = getattr(robot.indexing, ids_attr)
n = len(ids)
yaw = math.pi / 4
dr_func(env, env_ids=None, yaw_range=(yaw, yaw), asset_cfg=asset_cfg)
result = getattr(env.sim.model, field)[:, ids, :]
q_default = env.sim.get_default_field(field)[ids]
zeros = torch.zeros(n, device=env.device)
q_yaw = quat_from_euler_xyz(zeros, zeros, zeros + yaw)
q_expected = quat_mul(q_yaw, q_default)
assert torch.allclose(result, q_expected.unsqueeze(0).expand_as(result), atol=1e-6)
def test_body_quat_only_specified_axes(quat_env):
"""Yaw-only perturbation: roll/pitch of result match default."""
torch.manual_seed(1)
env = quat_env
robot = env.scene["robot"]
body_cfg = SceneEntityCfg("robot", body_names=("base",))
body_cfg.resolve(env.scene)
body_ids = robot.indexing.body_ids[body_cfg.body_ids]
dr.body_quat(
env,
env_ids=None,
yaw_range=(-0.3, 0.3),
asset_cfg=SceneEntityCfg("robot", body_names=("base",)),
)
result = env.sim.model.body_quat[:, body_ids, :]
# For the default quat [1,0,0,0] composed with a yaw-only rotation, qx and qy should
# remain 0 (pure yaw -> no roll/pitch).
q_default = env.sim.get_default_field("body_quat")[body_ids]
if torch.allclose(q_default, torch.tensor([[1.0, 0.0, 0.0, 0.0]], device=env.device)):
assert torch.allclose(result[..., 1], torch.zeros_like(result[..., 1]), atol=1e-6)
assert torch.allclose(result[..., 2], torch.zeros_like(result[..., 2]), atol=1e-6)
# pseudo_inertia tests.
def _make_inertia_env(device, num_envs=NUM_ENVS):
"""Create an env with all inertia-related fields expanded."""
return create_test_env(
device,
num_envs=num_envs,
expand_fields=("body_mass", "body_ipos", "body_inertia", "body_iquat"),
)
@pytest.fixture(scope="module")
def inertia_env(device):
return _make_inertia_env(device)
def _matrix_from_quat(q: torch.Tensor) -> torch.Tensor:
"""Convert wxyz quaternion(s) to 3x3 rotation matrix.
Shape: (..., 4) → (..., 3, 3).
"""
w, x, y, z = q[..., 0], q[..., 1], q[..., 2], q[..., 3]
return torch.stack(
[
torch.stack(
[1 - 2 * (y * y + z * z), 2 * (x * y - z * w), 2 * (x * z + y * w)], dim=-1
),
torch.stack(
[2 * (x * y + z * w), 1 - 2 * (x * x + z * z), 2 * (y * z - x * w)], dim=-1
),
torch.stack(
[2 * (x * z - y * w), 2 * (y * z + x * w), 1 - 2 * (x * x + y * y)], dim=-1
),
],
dim=-2,
)
def _reconstruct_J(mass, ipos, inertia, iquat):
"""Reconstruct pseudo-inertia matrix from MuJoCo fields (test helper).
Correctly accounts for body_iquat (the principal-frame rotation) and the
parallel-axis theorem (body_ipos) so that the reconstructed J is exact regardless of
shear magnitude.
"""
I3 = torch.eye(3, device=mass.device, dtype=mass.dtype)
# MuJoCo body_iquat maps principal→body.
R = _matrix_from_quat(iquat) # (..., 3, 3)
# Inertia tensor at COM in body frame: I_com = R @ diag(inertia) @ R^T
I_com = R @ torch.diag_embed(inertia) @ R.mT # (..., 3, 3)
# Parallel-axis theorem: shift from COM to body origin.
c = ipos # (..., 3)
c_sq = (c * c).sum(dim=-1) # (...,)
c_outer = c.unsqueeze(-1) * c.unsqueeze(-2) # (..., 3, 3)
m = mass.unsqueeze(-1).unsqueeze(-1) # (..., 1, 1)
I_origin = I_com + m * (c_sq.unsqueeze(-1).unsqueeze(-1) * I3 - c_outer)
trace = I_origin.diagonal(dim1=-2, dim2=-1).sum(dim=-1)
sigma = 0.5 * trace.unsqueeze(-1).unsqueeze(-1) * I3 - I_origin
h = mass.unsqueeze(-1) * ipos
batch = mass.shape
J = torch.zeros(*batch, 4, 4, device=mass.device, dtype=mass.dtype)
J[..., :3, :3] = sigma
J[..., :3, 3] = h
J[..., 3, :3] = h
J[..., 3, 3] = mass
return J
def test_pseudo_inertia_physical_consistency(inertia_env):
"""After wide perturbation, J ≻ 0, mass > 0, triangle inequality holds."""
torch.manual_seed(10)
env = inertia_env
robot = env.scene["robot"]
body_cfg = SceneEntityCfg("robot", body_names=("base", "foot1", "foot2"))
body_cfg.resolve(env.scene)
dr.pseudo_inertia(
env,
env_ids=None,
alpha_range=(-0.5, 0.5),
d_range=(-0.3, 0.3),
s12_range=(-0.1, 0.1),
s13_range=(-0.1, 0.1),
s23_range=(-0.1, 0.1),
t_range=(-0.05, 0.05),
asset_cfg=body_cfg,
)
body_ids = robot.indexing.body_ids[body_cfg.body_ids]
mass = env.sim.model.body_mass[:, body_ids]
ipos = env.sim.model.body_ipos[:, body_ids, :]
inertia = env.sim.model.body_inertia[:, body_ids, :]
iquat = env.sim.model.body_iquat[:, body_ids, :]
# Mass must be positive.
assert torch.all(mass > 0), f"Non-positive mass: {mass.min()}"
# Principal moments must all be positive.
assert torch.all(inertia > 0), f"Non-positive principal moment: {inertia.min()}"
# body_iquat must be a unit quaternion.
iquat_norms = iquat.norm(dim=-1)
assert torch.all((iquat_norms - 1.0).abs() < 1e-5), (
f"Non-unit body_iquat; max deviation: {(iquat_norms - 1.0).abs().max()}"
)
# Triangle inequality on principal moments: D_i + D_j >= D_k.
d1, d2, d3 = inertia[..., 0], inertia[..., 1], inertia[..., 2]
assert torch.all(d1 + d2 >= d3 - 1e-6)
assert torch.all(d1 + d3 >= d2 - 1e-6)
assert torch.all(d2 + d3 >= d1 - 1e-6)
# Round-trip check: reconstruct J from all four stored fields (including body_iquat)
# and verify it is positive definite. This is exact because the implementation writes
# the eigendecomposition of J'.
J = _reconstruct_J(mass, ipos, inertia, iquat)
eigvals = torch.linalg.eigvalsh(J)
assert torch.all(eigvals > -1e-5), f"Non-positive J eigenvalue: {eigvals.min()}"
def test_pseudo_inertia_zero_perturbation_unchanged(inertia_env):
"""All ranges (0, 0): mass, ipos, inertia, iquat equal defaults (1e-5)."""
env = inertia_env
robot = env.scene["robot"]
body_cfg = SceneEntityCfg("robot", body_names=("base", "foot1", "foot2"))
body_cfg.resolve(env.scene)
body_ids = robot.indexing.body_ids[body_cfg.body_ids]
def_mass = env.sim.get_default_field("body_mass")[body_ids].clone()
def_ipos = env.sim.get_default_field("body_ipos")[body_ids].clone()
def_inertia = env.sim.get_default_field("body_inertia")[body_ids].clone()
def_iquat = env.sim.get_default_field("body_iquat")[body_ids].clone()
dr.pseudo_inertia(env, env_ids=None, asset_cfg=body_cfg)
mass = env.sim.model.body_mass[:, body_ids]
ipos = env.sim.model.body_ipos[:, body_ids, :]
inertia = env.sim.model.body_inertia[:, body_ids, :]
iquat = env.sim.model.body_iquat[:, body_ids, :]
assert torch.allclose(mass, def_mass.unsqueeze(0).expand_as(mass), atol=1e-5)
assert torch.allclose(ipos, def_ipos.unsqueeze(0).expand_as(ipos), atol=1e-5)
# For bodies with degenerate eigenvalues (2+ equal principal moments), the
# eigenvector basis is non-unique, so the (inertia, iquat) decomposition may differ
# while producing the same physical inertia tensor. Compare the full tensor
# R @ diag(inertia) @ R^T rather than individual components.
R_def = _matrix_from_quat(def_iquat)
I_def = R_def @ torch.diag_embed(def_inertia) @ R_def.mT
R_new = _matrix_from_quat(iquat)
I_new = R_new @ torch.diag_embed(inertia) @ R_new.mT
assert torch.allclose(I_new, I_def.unsqueeze(0).expand_as(I_new), atol=1e-5), (
"Inertia tensor changed unexpectedly under zero perturbation"
)
def test_pseudo_inertia_zero_perturbation_offset_body(inertia_env):
"""Zero perturbation preserves physics for a body with non-trivial ipos/iquat.
Since ``eigh`` returns eigenvalues in ascending order while MuJoCo may store them
differently, we compare the reconstructed pseudo-inertia matrix J (which is
representation-independent) rather than the raw principal moments and iquat
individually.
"""
env = inertia_env
robot = env.scene["robot"]
body_cfg = SceneEntityCfg("robot", body_names=("offset_body",))
body_cfg.resolve(env.scene)
body_ids = robot.indexing.body_ids[body_cfg.body_ids]
def_mass = env.sim.get_default_field("body_mass")[body_ids].clone()
def_ipos = env.sim.get_default_field("body_ipos")[body_ids].clone()
def_inertia = env.sim.get_default_field("body_inertia")[body_ids].clone()
def_iquat = env.sim.get_default_field("body_iquat")[body_ids].clone()
# Verify the test body actually has non-trivial ipos and iquat.
assert not torch.allclose(def_ipos, torch.zeros_like(def_ipos), atol=1e-6), (
"Test body should have non-zero ipos"
)
identity_q = torch.tensor([1.0, 0.0, 0.0, 0.0], device=env.device)
assert not (
torch.allclose(def_iquat, identity_q.expand_as(def_iquat), atol=1e-6)
or torch.allclose(def_iquat, -identity_q.expand_as(def_iquat), atol=1e-6)
), "Test body should have non-identity iquat"
dr.pseudo_inertia(env, env_ids=None, asset_cfg=body_cfg)
mass = env.sim.model.body_mass[:, body_ids]
ipos = env.sim.model.body_ipos[:, body_ids, :]
inertia = env.sim.model.body_inertia[:, body_ids, :]
iquat = env.sim.model.body_iquat[:, body_ids, :]
# Mass and ipos must be exactly preserved.
assert torch.allclose(mass, def_mass.unsqueeze(0).expand_as(mass), atol=1e-5)
assert torch.allclose(ipos, def_ipos.unsqueeze(0).expand_as(ipos), atol=1e-5)
# Inertia and iquat may have different eigenvalue ordering, so compare the
# reconstructed J matrix instead.
J_default = _reconstruct_J(def_mass, def_ipos, def_inertia, def_iquat)
J_result = _reconstruct_J(mass, ipos, inertia, iquat)
assert torch.allclose(
J_result, J_default.unsqueeze(0).expand_as(J_result), atol=1e-5
), (
f"Reconstructed J differs under zero perturbation.\n"
f"Max deviation: {(J_result - J_default.unsqueeze(0)).abs().max()}"
)
def test_pseudo_inertia_alpha_only_scales_mass(inertia_env):
"""alpha-only: mass and inertia scale by e^{2α}, ipos unchanged."""
env = inertia_env
robot = env.scene["robot"]
body_cfg = SceneEntityCfg("robot", body_names=("base", "foot1", "foot2"))
body_cfg.resolve(env.scene)
body_ids = robot.indexing.body_ids[body_cfg.body_ids]
a = 0.3
scale = math.exp(2 * a)
def_mass = env.sim.get_default_field("body_mass")[body_ids].clone()
def_ipos = env.sim.get_default_field("body_ipos")[body_ids].clone()
def_inertia = env.sim.get_default_field("body_inertia")[body_ids].clone()
dr.pseudo_inertia(env, env_ids=None, alpha_range=(a, a), asset_cfg=body_cfg)
mass = env.sim.model.body_mass[:, body_ids]
ipos = env.sim.model.body_ipos[:, body_ids, :]
inertia = env.sim.model.body_inertia[:, body_ids, :]
expected_mass = def_mass * scale
expected_inertia = def_inertia * scale
assert torch.allclose(mass, expected_mass.unsqueeze(0).expand_as(mass), atol=1e-5)
assert torch.allclose(
inertia, expected_inertia.unsqueeze(0).expand_as(inertia), atol=1e-5
)
assert torch.allclose(ipos, def_ipos.unsqueeze(0).expand_as(ipos), atol=1e-5)
def test_pseudo_inertia_t1_shifts_com(inertia_env):
"""t1-only: mass unchanged, ipos[0] shifts by exactly t1."""
env = inertia_env
robot = env.scene["robot"]
body_cfg = SceneEntityCfg("robot", body_names=("base", "foot1", "foot2"))
body_cfg.resolve(env.scene)
body_ids = robot.indexing.body_ids[body_cfg.body_ids]
t = 0.05
def_mass = env.sim.get_default_field("body_mass")[body_ids].clone()
def_ipos = env.sim.get_default_field("body_ipos")[body_ids].clone()
dr.pseudo_inertia(env, env_ids=None, t1_range=(t, t), asset_cfg=body_cfg)
mass = env.sim.model.body_mass[:, body_ids]
ipos = env.sim.model.body_ipos[:, body_ids, :]
assert torch.allclose(mass, def_mass.unsqueeze(0).expand_as(mass), atol=1e-5)
# ipos x-component shifts by t; y and z are unchanged.
assert torch.allclose(
ipos[..., 0], (def_ipos[..., 0] + t).unsqueeze(0).expand_as(ipos[..., 0]), atol=1e-5
)
assert torch.allclose(
ipos[..., 1:], def_ipos[..., 1:].unsqueeze(0).expand_as(ipos[..., 1:]), atol=1e-5
)
def test_pseudo_inertia_d_isotropic_vs_per_axis(device):
"""d_range=(r,r) gives same result as d1=d2=d3=(r,r)."""
r = 0.2
seed = 42
torch.manual_seed(seed)
env1 = _make_inertia_env(device)
body_cfg = SceneEntityCfg("robot", body_names=("base", "foot1", "foot2"))
body_cfg.resolve(env1.scene)
dr.pseudo_inertia(env1, env_ids=None, d_range=(r, r), asset_cfg=body_cfg)
torch.manual_seed(seed)
env2 = _make_inertia_env(device)
dr.pseudo_inertia(
env2,
env_ids=None,
d1_range=(r, r),
d2_range=(r, r),
d3_range=(r, r),
asset_cfg=body_cfg,
)
robot1 = env1.scene["robot"]
body_ids = robot1.indexing.body_ids[body_cfg.body_ids]
assert torch.allclose(
env1.sim.model.body_mass[:, body_ids],
env2.sim.model.body_mass[:, body_ids],
atol=1e-6,
)
assert torch.allclose(
env1.sim.model.body_inertia[:, body_ids, :],
env2.sim.model.body_inertia[:, body_ids, :],
atol=1e-6,
)
def test_pseudo_inertia_no_accumulation(device):
"""3x with alpha=0.1 gives same result as 1x (uses defaults each time)."""
env = _make_inertia_env(device)
robot = env.scene["robot"]
body_cfg = SceneEntityCfg("robot", body_names=("base",))
body_cfg.resolve(env.scene)
body_ids = robot.indexing.body_ids[body_cfg.body_ids]
for _ in range(3):
dr.pseudo_inertia(env, env_ids=None, alpha_range=(0.1, 0.1), asset_cfg=body_cfg)
mass_3x = env.sim.model.body_mass[:, body_ids].clone()
env2 = _make_inertia_env(device)
dr.pseudo_inertia(env2, env_ids=None, alpha_range=(0.1, 0.1), asset_cfg=body_cfg)
mass_1x = env2.sim.model.body_mass[:, body_ids]
assert torch.allclose(mass_3x, mass_1x, atol=1e-5)
def test_pseudo_inertia_partial_env_ids(device):
"""Randomizing only env 0 leaves env 1 unchanged."""
env = _make_inertia_env(device, num_envs=2)
robot = env.scene["robot"]
body_cfg = SceneEntityCfg("robot", body_names=("base",))
body_cfg.resolve(env.scene)
body_ids = robot.indexing.body_ids[body_cfg.body_ids]
mass_before = env.sim.model.body_mass[:, body_ids].clone()
dr.pseudo_inertia(
env,
env_ids=torch.tensor([0], device=device),
alpha_range=(0.5, 0.5),
asset_cfg=body_cfg,
)
mass_after = env.sim.model.body_mass[:, body_ids]
# Env 0 changed.
assert not torch.allclose(mass_after[0], mass_before[0], atol=1e-6)
# Env 1 unchanged.
assert torch.allclose(mass_after[1], mass_before[1], atol=1e-6)
# Camera / Light DR tests.
def _make_cam_light_env(device, num_envs=NUM_ENVS):
"""Create an env with camera and light fields expanded."""
return create_test_env(
device,
num_envs=num_envs,
expand_fields=(
"cam_fovy",
"cam_pos",
"cam_quat",
"cam_intrinsic",
"light_pos",
"light_dir",
"light_diffuse",
"light_specular",
"light_ambient",
"light_attenuation",
"light_cutoff",
"light_exponent",
),
)
@pytest.fixture(scope="module")
def cam_light_env(device):
return _make_cam_light_env(device)
def test_cam_fovy_abs(cam_light_env):
"""Set fovy to absolute range, check bounds."""
torch.manual_seed(42)
env = cam_light_env
robot = env.scene["robot"]
cam_cfg = SceneEntityCfg("robot", camera_names=(".*",))
cam_cfg.resolve(env.scene)
cam_ids = robot.indexing.cam_ids[cam_cfg.camera_ids]
dr.cam_fovy(
env,
env_ids=None,
ranges=(30.0, 90.0),
operation="abs",
asset_cfg=cam_cfg,
)
fovy = env.sim.model.cam_fovy[:, cam_ids]
assert torch.all((fovy >= 30.0 - 1e-3) & (fovy <= 90.0 + 1e-3))
def test_cam_pos_add(cam_light_env):
"""Add offset to camera positions, check default + offset."""
torch.manual_seed(42)
env = cam_light_env
robot = env.scene["robot"]
cam_cfg = SceneEntityCfg("robot", camera_names=("front_cam",))
cam_cfg.resolve(env.scene)
cam_ids = robot.indexing.cam_ids[cam_cfg.camera_ids]
default_pos = env.sim.get_default_field("cam_pos")[cam_ids].clone()
dr.cam_pos(
env,
env_ids=None,
ranges=(-0.1, 0.1),
operation="add",
asset_cfg=cam_cfg,
)
result = env.sim.model.cam_pos[:, cam_ids, :]
for ax in range(3):
lo = default_pos[..., ax] - 0.1 - 1e-5
hi = default_pos[..., ax] + 0.1 + 1e-5
assert torch.all((result[..., ax] >= lo) & (result[..., ax] <= hi))
def test_cam_quat_zero_unchanged(cam_light_env):
"""Zero RPY range preserves default quaternion."""
env = cam_light_env
robot = env.scene["robot"]
cam_cfg = SceneEntityCfg("robot", camera_names=(".*",))
cam_cfg.resolve(env.scene)
cam_ids = robot.indexing.cam_ids[cam_cfg.camera_ids]
default_quat = env.sim.get_default_field("cam_quat")[cam_ids].clone()
dr.cam_quat(env, env_ids=None, asset_cfg=cam_cfg)
result = env.sim.model.cam_quat[:, cam_ids, :]
assert torch.allclose(
result,
default_quat.unsqueeze(0).expand_as(result),
atol=1e-6,
)
def test_light_pos_abs(cam_light_env):
"""Set light position to absolute range, check bounds."""
torch.manual_seed(42)
env = cam_light_env
robot = env.scene["robot"]
light_cfg = SceneEntityCfg("robot", light_names=(".*",))
light_cfg.resolve(env.scene)
light_ids = robot.indexing.light_ids[light_cfg.light_ids]
dr.light_pos(
env,
env_ids=None,
ranges=(-5.0, 5.0),
operation="abs",
asset_cfg=light_cfg,
)
pos = env.sim.model.light_pos[:, light_ids, :]
assert torch.all((pos >= -5.0 - 1e-3) & (pos <= 5.0 + 1e-3))
def test_light_dir_add(cam_light_env):
"""Add offset to light direction, check default + offset."""
torch.manual_seed(42)
env = cam_light_env
robot = env.scene["robot"]
light_cfg = SceneEntityCfg("robot", light_names=(".*",))
light_cfg.resolve(env.scene)
light_ids = robot.indexing.light_ids[light_cfg.light_ids]
default_dir = env.sim.get_default_field("light_dir")[light_ids].clone()
dr.light_dir(
env,
env_ids=None,
ranges=(-0.5, 0.5),
operation="add",
asset_cfg=light_cfg,
)
result = env.sim.model.light_dir[:, light_ids, :]
for ax in range(3):
lo = default_dir[..., ax] - 0.5 - 1e-5
hi = default_dir[..., ax] + 0.5 + 1e-5
assert torch.all((result[..., ax] >= lo) & (result[..., ax] <= hi))
@pytest.mark.parametrize(
("func_name", "field", "ranges"),
[
("light_diffuse", "light_diffuse", (0.2, 0.8)),
("light_specular", "light_specular", (0.0, 1.0)),
("light_ambient", "light_ambient", (0.1, 0.5)),
("light_attenuation", "light_attenuation", (0.0, 1.0)),
],
)
def test_light_vec3_fields_abs(cam_light_env, func_name, field, ranges):
"""Vec3 light fields randomize per selected light."""
torch.manual_seed(42)
env = cam_light_env
robot = env.scene["robot"]
light_cfg = SceneEntityCfg("robot", light_names=(".*",))
light_cfg.resolve(env.scene)
light_ids = robot.indexing.light_ids[light_cfg.light_ids]
func = getattr(dr, func_name)
func(
env,
env_ids=None,
ranges=ranges,
operation="abs",
asset_cfg=light_cfg,
)
values = getattr(env.sim.model, field)[:, light_ids, :]
lower, upper = ranges
assert torch.all((values >= lower - 1e-5) & (values <= upper + 1e-5))
assert len(torch.unique(values[:, 0, 0])) >= 2
@pytest.mark.parametrize(
("func_name", "field", "ranges"),
[
("light_cutoff", "light_cutoff", (20.0, 60.0)),
("light_exponent", "light_exponent", (1.0, 20.0)),
],
)
def test_light_scalar_fields_abs(cam_light_env, func_name, field, ranges):
"""Scalar light fields randomize per selected light."""
torch.manual_seed(42)
env = cam_light_env
robot = env.scene["robot"]
light_cfg = SceneEntityCfg("robot", light_names=(".*",))
light_cfg.resolve(env.scene)
light_ids = robot.indexing.light_ids[light_cfg.light_ids]
func = getattr(dr, func_name)
func(
env,
env_ids=None,
ranges=ranges,
operation="abs",
asset_cfg=light_cfg,
)
values = getattr(env.sim.model, field)[:, light_ids]
lower, upper = ranges
assert torch.all((values >= lower - 1e-5) & (values <= upper + 1e-5))
assert len(torch.unique(values[:, 0])) >= 2
def test_camera_partial_env_ids(cam_light_env):
"""Subset of envs randomized, others unchanged."""
torch.manual_seed(42)
env = cam_light_env
robot = env.scene["robot"]
cam_cfg = SceneEntityCfg("robot", camera_names=(".*",))
cam_cfg.resolve(env.scene)
cam_ids = robot.indexing.cam_ids[cam_cfg.camera_ids]
original_fovy = env.sim.model.cam_fovy[:, cam_ids].clone()
dr.cam_fovy(
env,
env_ids=torch.tensor([0, 2], device=env.device),
ranges=(10.0, 20.0),
operation="abs",
asset_cfg=cam_cfg,
)
result = env.sim.model.cam_fovy[:, cam_ids]
# Envs 0, 2 should have changed.
assert torch.all((result[0] >= 10.0 - 1e-3) & (result[0] <= 20.0 + 1e-3))
assert torch.all((result[2] >= 10.0 - 1e-3) & (result[2] <= 20.0 + 1e-3))
# Envs 1, 3 should be unchanged.
assert torch.allclose(result[1], original_fovy[1])
assert torch.allclose(result[3], original_fovy[3])
# geom_size DR tests.
GEOM_SIZE_XML = """
<mujoco>
<worldbody>
<body name="base" pos="0 0 1">
<freejoint name="free_joint"/>
<geom name="box_geom" type="box" size="0.1 0.2 0.3" mass="1.0"/>
<geom name="sphere_geom" type="sphere" size="0.15" mass="0.5"/>
<geom name="capsule_geom" type="capsule" size="0.05 0.2" mass="0.3"/>
<geom name="cylinder_geom" type="cylinder" size="0.08 0.15" mass="0.4"/>
<geom name="ellipsoid_geom" type="ellipsoid" size="0.1 0.2 0.05"
mass="0.2"/>
</body>
</worldbody>
</mujoco>
"""
def _make_geom_size_env(device, num_envs=NUM_ENVS):
entity_cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(GEOM_SIZE_XML))
scene_cfg = SceneCfg(num_envs=num_envs, entities={"robot": entity_cfg})
scene = Scene(scene_cfg, device)
model = scene.compile()
sim = Simulation(num_envs=num_envs, cfg=SimulationCfg(), model=model, device=device)
scene.initialize(model, sim.model, sim.data)
sim.expand_model_fields(("geom_size", "geom_rbound", "geom_aabb"))
return Env(scene, sim, device)
@pytest.fixture(scope="module")
def geom_size_env(device):
return _make_geom_size_env(device)
def _expected_rbound(geom_type, s0, s1, s2):
"""Reference rbound computation matching MuJoCo GetRBound."""
if geom_type == "sphere":
return s0
elif geom_type == "capsule":
return s0 + s1
elif geom_type == "cylinder":
return math.sqrt(s0**2 + s1**2)
elif geom_type == "ellipsoid":
return max(s0, s1, s2)
elif geom_type == "box":
return math.sqrt(s0**2 + s1**2 + s2**2)
raise ValueError(geom_type)
def _expected_aabb_half(geom_type, s0, s1, s2):
"""Reference aabb half-size computation matching MuJoCo ComputeAABB."""
if geom_type == "sphere":
return (s0, s0, s0)
elif geom_type == "capsule":
return (s0, s0, s0 + s1)
elif geom_type == "cylinder":
return (s0, s0, s1)
elif geom_type == "ellipsoid":
return (s0, s1, s2)
elif geom_type == "box":
return (s0, s1, s2)
raise ValueError(geom_type)
def test_geom_size_scale_updates_bounds(geom_size_env):
"""Scaling geom_size updates rbound and aabb consistently."""
torch.manual_seed(42)
env = geom_size_env
robot = env.scene["robot"]
geom_cfg = SceneEntityCfg("robot", geom_names=(".*",))
geom_cfg.resolve(env.scene)
dr.geom_size(
env,
env_ids=None,
ranges=(0.8, 1.5),
operation="scale",
asset_cfg=geom_cfg,
)
geom_ids = robot.indexing.geom_ids[geom_cfg.geom_ids]
size = env.sim.model.geom_size[:, geom_ids] # (E, G, 3)
rbound = env.sim.model.geom_rbound[:, geom_ids] # (E, G)
aabb = env.sim.model.geom_aabb[:, geom_ids] # (E, G, 2, 3)
geom_names = list(robot.geom_names)
type_names = {
"box_geom": "box",
"sphere_geom": "sphere",
"capsule_geom": "capsule",
"cylinder_geom": "cylinder",
"ellipsoid_geom": "ellipsoid",
}
for g_local, gname in enumerate(geom_names):
tname = type_names[gname]
for e in range(env.num_envs):
s0 = size[e, g_local, 0].item()
s1 = size[e, g_local, 1].item()
s2 = size[e, g_local, 2].item()
expected_rb = _expected_rbound(tname, s0, s1, s2)
actual_rb = rbound[e, g_local].item()
assert abs(actual_rb - expected_rb) < 1e-5, (
f"{gname} env {e}: rbound {actual_rb} != {expected_rb}"
)
ex, ey, ez = _expected_aabb_half(tname, s0, s1, s2)
actual_half = aabb[e, g_local, 1] # half-size is index 1
assert abs(actual_half[0].item() - ex) < 1e-5
assert abs(actual_half[1].item() - ey) < 1e-5
assert abs(actual_half[2].item() - ez) < 1e-5
# Center should be zero for all primitives.
center = aabb[e, g_local, 0]
assert torch.allclose(center, torch.zeros(3, device=env.device), atol=1e-6)
def test_geom_size_no_accumulation(device):
"""Repeated scale operations don't accumulate."""
env = _make_geom_size_env(device, num_envs=2)
robot = env.scene["robot"]
geom_cfg = SceneEntityCfg("robot", geom_names=("box_geom",))
geom_cfg.resolve(env.scene)
geom_ids = robot.indexing.geom_ids[geom_cfg.geom_ids]
default_size = env.sim.get_default_field("geom_size")[geom_ids].clone()
for _ in range(3):
dr.geom_size(
env,
env_ids=None,
ranges=(2.0, 2.0),
operation="scale",
asset_cfg=geom_cfg,
)
result = env.sim.model.geom_size[0, geom_ids]
expected = default_size * 2.0
assert torch.allclose(result, expected, atol=1e-5)
def test_geom_size_separate_axes_dont_clobber(device):
"""Events targeting different axes of the same geom compose, not clobber."""
env = _make_geom_size_env(device, num_envs=2)
robot = env.scene["robot"]
geom_cfg = SceneEntityCfg("robot", geom_names=("box_geom",))
geom_cfg.resolve(env.scene)
geom_ids = robot.indexing.geom_ids[geom_cfg.geom_ids]
default_size = env.sim.get_default_field("geom_size")[geom_ids].clone()
# Scale axis 0 by 0.1, then axis 1 by 3.0 in a separate event.
dr.geom_size(
env,
env_ids=None,
ranges=(0.1, 0.1),
operation="scale",
axes=[0],
asset_cfg=geom_cfg,
)
dr.geom_size(
env,
env_ids=None,
ranges=(3.0, 3.0),
operation="scale",
axes=[1],
asset_cfg=geom_cfg,
)
result = env.sim.model.geom_size[0, geom_ids]
expected = default_size.clone()
expected[..., 0] *= 0.1
expected[..., 1] *= 3.0
assert torch.allclose(result, expected, atol=1e-5)
def test_geom_size_partial_env_ids(device):
"""Only specified envs are updated."""
env = _make_geom_size_env(device, num_envs=4)
robot = env.scene["robot"]
geom_cfg = SceneEntityCfg("robot", geom_names=("sphere_geom",))
geom_cfg.resolve(env.scene)
geom_ids = robot.indexing.geom_ids[geom_cfg.geom_ids]
original_rbound = env.sim.model.geom_rbound[:, geom_ids].clone()
dr.geom_size(
env,
env_ids=torch.tensor([0, 2], device=device),
ranges=(2.0, 2.0),
operation="scale",
asset_cfg=geom_cfg,
)
rbound = env.sim.model.geom_rbound[:, geom_ids]
# Envs 0, 2 should have changed.
assert not torch.allclose(rbound[0], original_rbound[0])
assert not torch.allclose(rbound[2], original_rbound[2])
# Envs 1, 3 should be unchanged.
assert torch.allclose(rbound[1], original_rbound[1])
assert torch.allclose(rbound[3], original_rbound[3])
def test_geom_size_raises_on_unsupported_type(device):
"""dr.geom_size raises ValueError when a non-primitive geom type is selected."""
PLANE_XML = """
<mujoco>
<worldbody>
<geom name="floor" type="plane" size="5 5 0.1"/>
<body name="base" pos="0 0 1">
<freejoint name="free_joint"/>
<geom name="box_geom" type="box" size="0.1 0.1 0.1" mass="1.0"/>
</body>
</worldbody>
</mujoco>
"""
entity_cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(PLANE_XML))
scene_cfg = SceneCfg(num_envs=2, entities={"robot": entity_cfg})
scene = Scene(scene_cfg, device)
model = scene.compile()
sim = Simulation(num_envs=2, cfg=SimulationCfg(), model=model, device=device)
scene.initialize(model, sim.model, sim.data)
sim.expand_model_fields(("geom_size", "geom_rbound", "geom_aabb"))
env = Env(scene, sim, device)
plane_cfg = SceneEntityCfg("robot", geom_names=("floor",))
plane_cfg.resolve(env.scene)
with pytest.raises(ValueError, match="unsupported types"):
dr.geom_size(env, env_ids=None, ranges=(1.0, 2.0), asset_cfg=plane_cfg)
# Tendon DR tests.
TENDON_XML = """
<mujoco>
<worldbody>
<body name="base" pos="0 0 1">
<freejoint name="free_joint"/>
<geom name="base_geom" type="box" size="0.1 0.1 0.1" mass="1.0"/>
<site name="anchor" pos="0 0 0.1"/>
<body name="child" pos="0.3 0 0">
<joint name="hinge1" type="hinge" axis="0 0 1"/>
<geom name="child_geom" type="sphere" size="0.05" mass="0.2"/>
<site name="hook" pos="0 0 -0.05"/>
</body>
</body>
</worldbody>
<tendon>
<spatial name="wire" limited="true" range="0 0.5" width="0.003"
damping="2.0" stiffness="10.0" frictionloss="0.5" armature="0.1">
<site site="anchor"/>
<site site="hook"/>
</spatial>
</tendon>
</mujoco>
"""
def _make_tendon_env(device, num_envs=NUM_ENVS):
entity_cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(TENDON_XML))
scene_cfg = SceneCfg(num_envs=num_envs, entities={"robot": entity_cfg})
scene = Scene(scene_cfg, device)
model = scene.compile()
sim = Simulation(num_envs=num_envs, cfg=SimulationCfg(), model=model, device=device)
scene.initialize(model, sim.model, sim.data)
sim.expand_model_fields(
(
"tendon_damping",
"tendon_stiffness",
"tendon_frictionloss",
"tendon_armature",
)
)
return Env(scene, sim, device)
@pytest.fixture(scope="module")
def tendon_env(device):
return _make_tendon_env(device)
@pytest.mark.parametrize(
"dr_func, field",
[
(dr.tendon_damping, "tendon_damping"),
(dr.tendon_stiffness, "tendon_stiffness"),
(dr.tendon_friction, "tendon_frictionloss"),
(dr.tendon_armature, "tendon_armature"),
],
)
def test_tendon_field_scale(tendon_env, dr_func, field):
"""Scale operation on tendon fields: result = default * sample."""
torch.manual_seed(42)
env = tendon_env
robot = env.scene["robot"]
tendon_cfg = SceneEntityCfg("robot", tendon_names=(".*",))
tendon_cfg.resolve(env.scene)
tendon_ids = robot.indexing.tendon_ids[tendon_cfg.tendon_ids]
default_val = env.sim.get_default_field(field)[tendon_ids].clone()
dr_func(
env,
env_ids=None,
ranges=(0.5, 2.0),
operation="scale",
asset_cfg=tendon_cfg,
)
result = getattr(env.sim.model, field)[:, tendon_ids]
lo = default_val * 0.5 - 1e-5
hi = default_val * 2.0 + 1e-5
assert torch.all((result >= lo) & (result <= hi))
def test_tendon_armature_no_accumulation(device):
"""Repeated scale does not accumulate."""
env = _make_tendon_env(device, num_envs=2)
robot = env.scene["robot"]
tendon_cfg = SceneEntityCfg("robot", tendon_names=(".*",))
tendon_cfg.resolve(env.scene)
tendon_ids = robot.indexing.tendon_ids[tendon_cfg.tendon_ids]
default_val = env.sim.get_default_field("tendon_armature")[tendon_ids]
for _ in range(3):
dr.tendon_armature(
env,
env_ids=None,
ranges=(2.0, 2.0),
operation="scale",
asset_cfg=tendon_cfg,
)
result = env.sim.model.tendon_armature[0, tendon_ids]
assert torch.allclose(result, default_val * 2.0, atol=1e-5)
# Extensible Operation / Distribution types.
def test_operation_instance_abs(device):
"""Passing an Operation instance works identically to the string."""
from mjlab.envs.mdp.dr._types import abs as abs_op
torch.manual_seed(42)
env1 = create_test_env(device)
robot = env1.scene["robot"]
dr.geom_friction(
env1,
env_ids=None,
ranges=(0.3, 1.2),
operation="abs",
asset_cfg=SceneEntityCfg("robot", geom_names=(".*",)),
axes=[0],
)
result_str = env1.sim.model.geom_friction[:, robot.indexing.geom_ids, 0].clone()
torch.manual_seed(42)
env2 = create_test_env(device)
dr.geom_friction(
env2,
env_ids=None,
ranges=(0.3, 1.2),
operation=abs_op,
asset_cfg=SceneEntityCfg("robot", geom_names=(".*",)),
axes=[0],
)
result_inst = env2.sim.model.geom_friction[:, robot.indexing.geom_ids, 0]
assert torch.allclose(result_str, result_inst)
def test_distribution_instance_uniform(device):
"""Passing a Distribution instance works identically to the string."""
from mjlab.envs.mdp.dr._types import uniform as uniform_dist
torch.manual_seed(42)
env1 = create_test_env(device)
robot = env1.scene["robot"]
dr.geom_friction(
env1,
env_ids=None,
ranges=(0.3, 1.2),
distribution="uniform",
operation="abs",
asset_cfg=SceneEntityCfg("robot", geom_names=(".*",)),
axes=[0],
)
result_str = env1.sim.model.geom_friction[:, robot.indexing.geom_ids, 0].clone()
torch.manual_seed(42)
env2 = create_test_env(device)
dr.geom_friction(
env2,
env_ids=None,
ranges=(0.3, 1.2),
distribution=uniform_dist,
operation="abs",
asset_cfg=SceneEntityCfg("robot", geom_names=(".*",)),
axes=[0],
)
result_inst = env2.sim.model.geom_friction[:, robot.indexing.geom_ids, 0]
assert torch.allclose(result_str, result_inst)
def test_custom_operation(env):
"""A user-defined Operation works end-to-end."""
from mjlab.envs.mdp.dr._types import Operation
clamp_op = Operation(
name="clamp",
initialize=torch.Tensor.clone,
combine=lambda _base, random: torch.clamp(random, min=0.4, max=0.9),
uses_defaults=False,
)
torch.manual_seed(42)
robot = env.scene["robot"]
dr.geom_friction(
env,
env_ids=None,
ranges=(0.1, 1.5),
operation=clamp_op,
asset_cfg=SceneEntityCfg("robot", geom_names=(".*",)),
axes=[0],
)
friction = env.sim.model.geom_friction[:, robot.indexing.geom_ids, 0]
assert torch.all((friction >= 0.4 - 1e-5) & (friction <= 0.9 + 1e-5))
def test_custom_distribution(env):
"""A user-defined Distribution works end-to-end."""
from mjlab.envs.mdp.dr._types import Distribution
# Distribution that always returns the midpoint.
midpoint_dist = Distribution(
name="midpoint",
sample=lambda lo, hi, shape, _device: ((lo + hi) / 2).expand(shape),
)
robot = env.scene["robot"]
dr.geom_friction(
env,
env_ids=None,
ranges=(0.2, 0.8),
distribution=midpoint_dist,
operation="abs",
asset_cfg=SceneEntityCfg("robot", geom_names=(".*",)),
axes=[0],
)
friction = env.sim.model.geom_friction[:, robot.indexing.geom_ids, 0]
assert torch.allclose(friction, torch.tensor(0.5, device=env.device), atol=1e-5)
def test_resolve_unknown_operation_raises(device):
"""Unknown operation string raises ValueError."""
from mjlab.envs.mdp.dr._types import resolve_operation
with pytest.raises(ValueError, match="Unknown operation"):
resolve_operation("nonexistent")
def test_resolve_unknown_distribution_raises(device):
"""Unknown distribution string raises ValueError."""
from mjlab.envs.mdp.dr._types import resolve_distribution
with pytest.raises(ValueError, match="Unknown distribution"):
resolve_distribution("nonexistent")
# mat_rgba tests.
MAT_XML = """
<mujoco>
<asset>
<material name="test_mat" rgba="1 1 1 1" emission="0.1"
specular="0.2" shininess="0.3" texrepeat="1 2"/>
</asset>
<worldbody>
<body name="base" pos="0 0 1">
<freejoint name="free_joint"/>
<geom name="base_geom" type="box" size="0.1 0.1 0.1" mass="1.0"
material="test_mat"/>
</body>
</worldbody>
</mujoco>
"""
def _make_mat_env(device, num_envs=NUM_ENVS):
entity_cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(MAT_XML))
scene_cfg = SceneCfg(num_envs=num_envs, entities={"robot": entity_cfg})
scene = Scene(scene_cfg, device)
model = scene.compile()
sim = Simulation(num_envs=num_envs, cfg=SimulationCfg(), model=model, device=device)
scene.initialize(model, sim.model, sim.data)
sim.expand_model_fields(
("mat_rgba", "mat_emission", "mat_specular", "mat_shininess", "mat_texrepeat")
)
return Env(scene, sim, device)
@pytest.fixture(scope="module")
def mat_env(device):
return _make_mat_env(device)
def test_mat_rgba_abs(mat_env):
"""Values within range, diversity across envs."""
torch.manual_seed(42)
env = mat_env
dr.mat_rgba(
env,
env_ids=None,
ranges=(0.2, 0.8),
asset_cfg=SceneEntityCfg("robot", material_names=("test_mat",)),
operation="abs",
)
mat_id = mujoco.mj_name2id(
env.sim.mj_model, mujoco.mjtObj.mjOBJ_MATERIAL, "robot/test_mat"
)
rgba = env.sim.model.mat_rgba[:, mat_id, :]
assert torch.all((rgba >= 0.2 - 1e-5) & (rgba <= 0.8 + 1e-5))
assert len(torch.unique(rgba[:, 0])) >= 2
@pytest.mark.parametrize(
("func_name", "field", "ranges"),
[
("mat_emission", "mat_emission", (0.1, 0.7)),
("mat_specular", "mat_specular", (0.0, 1.0)),
("mat_shininess", "mat_shininess", (0.2, 0.9)),
],
)
def test_scalar_material_fields_abs(mat_env, func_name, field, ranges):
"""Scalar material fields randomize per selected material."""
torch.manual_seed(42)
env = mat_env
func = getattr(dr, func_name)
func(
env,
env_ids=None,
ranges=ranges,
asset_cfg=SceneEntityCfg("robot", material_names=("test_mat",)),
operation="abs",
)
mat_id = mujoco.mj_name2id(
env.sim.mj_model, mujoco.mjtObj.mjOBJ_MATERIAL, "robot/test_mat"
)
values = getattr(env.sim.model, field)[:, mat_id]
lower, upper = ranges
assert torch.all((values >= lower - 1e-5) & (values <= upper + 1e-5))
assert len(torch.unique(values)) >= 2
def test_mat_texrepeat_abs(mat_env):
"""Texture repeat randomizes S/T axes for selected materials."""
torch.manual_seed(42)
env = mat_env
dr.mat_texrepeat(
env,
env_ids=None,
ranges={0: (1.0, 3.0), 1: (4.0, 6.0)},
asset_cfg=SceneEntityCfg("robot", material_names=("test_mat",)),
operation="abs",
)
mat_id = mujoco.mj_name2id(
env.sim.mj_model, mujoco.mjtObj.mjOBJ_MATERIAL, "robot/test_mat"
)
texrepeat = env.sim.model.mat_texrepeat[:, mat_id, :]
assert torch.all((texrepeat[:, 0] >= 1.0 - 1e-5) & (texrepeat[:, 0] <= 3.0 + 1e-5))
assert torch.all((texrepeat[:, 1] >= 4.0 - 1e-5) & (texrepeat[:, 1] <= 6.0 + 1e-5))
assert len(torch.unique(texrepeat[:, 0])) >= 2
@pytest.mark.parametrize(
("func_name", "field"),
[
("mat_emission", "mat_emission"),
("mat_specular", "mat_specular"),
("mat_shininess", "mat_shininess"),
],
)
def test_scalar_material_fields_scale_use_defaults(mat_env, func_name, field):
"""Scale operation uses compiled defaults for scalar material fields."""
env = mat_env
func = getattr(dr, func_name)
func(
env,
env_ids=None,
ranges=(2.0, 2.0),
asset_cfg=SceneEntityCfg("robot", material_names=("test_mat",)),
operation="scale",
)
mat_id = mujoco.mj_name2id(
env.sim.mj_model, mujoco.mjtObj.mjOBJ_MATERIAL, "robot/test_mat"
)
expected = env.sim.get_default_field(field)[mat_id] * 2.0
values = getattr(env.sim.model, field)[:, mat_id]
assert torch.allclose(values, expected.expand_as(values))
def test_mat_texrepeat_scale_uses_defaults(mat_env):
"""Scale operation uses compiled defaults and both S/T axes by default."""
env = mat_env
dr.mat_texrepeat(
env,
env_ids=None,
ranges=(2.0, 2.0),
asset_cfg=SceneEntityCfg("robot", material_names=("test_mat",)),
operation="scale",
)
mat_id = mujoco.mj_name2id(
env.sim.mj_model, mujoco.mjtObj.mjOBJ_MATERIAL, "robot/test_mat"
)
expected = env.sim.get_default_field("mat_texrepeat")[mat_id] * 2.0
texrepeat = env.sim.model.mat_texrepeat[:, mat_id, :]
assert torch.allclose(texrepeat, expected.expand_as(texrepeat))
def test_mat_rgba_invalid_name(mat_env):
"""ValueError for unknown material name."""
env = mat_env
with pytest.raises(ValueError, match="nonexistent_material"):
cfg = SceneEntityCfg("robot", material_names=("nonexistent_material",))
cfg.resolve(env.scene)
dr.mat_rgba(
env,
env_ids=None,
ranges=(0.2, 0.8),
asset_cfg=cfg,
)
# geom_matid DR tests.
GEOM_MATID_XML = """
<mujoco>
<asset>
<material name="mat_a" rgba="1 0 0 1"/>
<material name="mat_b" rgba="0 1 0 1"/>
<material name="mat_c" rgba="0 0 1 1"/>
</asset>
<worldbody>
<body name="base" pos="0 0 1">
<freejoint name="free_joint"/>
<geom name="geom1" type="box" size="0.1 0.1 0.1" mass="1.0" material="mat_a"/>
<geom name="geom2" type="sphere" size="0.05" mass="0.3" material="mat_b"/>
</body>
</worldbody>
</mujoco>
"""
def _make_matid_env(device, num_envs=NUM_ENVS):
entity_cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(GEOM_MATID_XML))
scene_cfg = SceneCfg(num_envs=num_envs, entities={"robot": entity_cfg})
scene = Scene(scene_cfg, device)
model = scene.compile()
sim = Simulation(num_envs=num_envs, cfg=SimulationCfg(), model=model, device=device)
scene.initialize(model, sim.model, sim.data)
sim.expand_model_fields(("geom_matid",))
return Env(scene, sim, device)
@pytest.fixture(scope="module")
def matid_env(device):
return _make_matid_env(device)
def test_geom_matid_draws_from_pool(matid_env):
"""Assigned matids come from the pool and actually change from the defaults."""
torch.manual_seed(42)
env = matid_env
robot = env.scene["robot"]
asset_cfg = SceneEntityCfg("robot", geom_names=(".*",), material_names=(".*",))
asset_cfg.resolve(env.scene)
geom_ids = robot.indexing.geom_ids[asset_cfg.geom_ids]
mat_ids = robot.indexing.mat_ids[asset_cfg.material_ids]
original = env.sim.model.geom_matid[:, geom_ids].clone()
dr.geom_matid(env, env_ids=None, asset_cfg=asset_cfg)
assigned = env.sim.model.geom_matid[:, geom_ids]
# Every assignment is a valid member of the selected pool.
assert torch.all(torch.isin(assigned, mat_ids))
# Randomization is not a no-op: values differ from the baked defaults, and at
# least one geom draws a material it did not start with (proving the full pool
# is sampled, not just the compile-time defaults).
assert not torch.equal(assigned, original)
assert torch.isin(assigned, torch.unique(original), invert=True).any()
def test_geom_matid_partial_env_ids(matid_env):
"""Randomizing subset of envs leaves others unchanged."""
torch.manual_seed(42)
env = matid_env
robot = env.scene["robot"]
asset_cfg = SceneEntityCfg("robot", geom_names=(".*",), material_names=(".*",))
asset_cfg.resolve(env.scene)
geom_ids = robot.indexing.geom_ids[asset_cfg.geom_ids]
original = env.sim.model.geom_matid[:, geom_ids].clone()
dr.geom_matid(
env, env_ids=torch.tensor([0, 2], device=env.device), asset_cfg=asset_cfg
)
result = env.sim.model.geom_matid[:, geom_ids]
assert torch.all(result[1] == original[1])
assert torch.all(result[3] == original[3])
def test_geom_matid_invalid_material_name(matid_env):
"""Unknown material name is rejected during config resolution."""
env = matid_env
with pytest.raises(ValueError, match="nonexistent_material"):
cfg = SceneEntityCfg(
"robot", geom_names=(".*",), material_names=("nonexistent_material",)
)
cfg.resolve(env.scene)
dr.geom_matid(env, env_ids=None, asset_cfg=cfg)
def test_geom_matid_empty_material_selection(matid_env):
"""geom_matid raises when the resolved material pool is empty."""
env = matid_env
cfg = SceneEntityCfg("robot", geom_names=(".*",), material_names=())
cfg.resolve(env.scene)
with pytest.raises(ValueError, match="No materials selected"):
dr.geom_matid(env, env_ids=None, asset_cfg=cfg)
def test_geom_matid_shared_random(matid_env):
"""All geoms within same env get same material, envs differ."""
torch.manual_seed(42)
env = matid_env
robot = env.scene["robot"]
asset_cfg = SceneEntityCfg("robot", geom_names=(".*",), material_names=(".*",))
asset_cfg.resolve(env.scene)
geom_ids = robot.indexing.geom_ids[asset_cfg.geom_ids]
dr.geom_matid(env, env_ids=None, asset_cfg=asset_cfg, shared_random=True)
assigned = env.sim.model.geom_matid[:, geom_ids]
for env_idx in range(env.num_envs):
env_matids = assigned[env_idx]
assert torch.all(env_matids == env_matids[0])
assert len(torch.unique(assigned[:, 0])) > 1
# mat_texid DR tests.
MAT_TEXID_XML = """
<mujoco>
<asset>
<texture name="tex_a" type="2d" builtin="flat" rgb1="1 0 0" width="4" height="4"/>
<texture name="tex_b" type="2d" builtin="flat" rgb1="0 1 0" width="4" height="4"/>
<texture name="tex_c" type="2d" builtin="flat" rgb1="0 0 1" width="4" height="4"/>
<texture name="tex_d" type="2d" builtin="flat" rgb1="1 1 0" width="4" height="4"/>
<material name="mat_a" texture="tex_a"/>
<material name="mat_b" texture="tex_b"/>
</asset>
<worldbody>
<body name="base" pos="0 0 1">
<freejoint name="free_joint"/>
<geom name="geom1" type="box" size="0.1 0.1 0.1" mass="1.0" material="mat_a"/>
<geom name="geom2" type="sphere" size="0.05" mass="0.3" material="mat_b"/>
</body>
</worldbody>
</mujoco>
"""
def _make_texid_env(device, num_envs=NUM_ENVS):
entity_cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(MAT_TEXID_XML))
scene_cfg = SceneCfg(num_envs=num_envs, entities={"robot": entity_cfg})
scene = Scene(scene_cfg, device)
model = scene.compile()
sim = Simulation(num_envs=num_envs, cfg=SimulationCfg(), model=model, device=device)
scene.initialize(model, sim.model, sim.data)
sim.expand_model_fields(("mat_texid",))
return Env(scene, sim, device)
@pytest.fixture
def texid_env(device):
# Function-scoped so tests compare against pristine mat_texid defaults.
return _make_texid_env(device)
def test_mat_texid_draws_from_pool(texid_env):
"""Assigned texids come from the pool and actually change from the defaults."""
torch.manual_seed(42)
env = texid_env
robot = env.scene["robot"]
asset_cfg = SceneEntityCfg("robot", material_names=(".*",), texture_names=(".*",))
asset_cfg.resolve(env.scene)
mat_ids = robot.indexing.mat_ids[asset_cfg.material_ids]
tex_ids = robot.indexing.tex_ids[asset_cfg.texture_ids]
role = mujoco.mjtTextureRole.mjTEXROLE_RGB.value
original = env.sim.model.mat_texid[:, mat_ids, role].clone()
dr.mat_texid(env, env_ids=None, asset_cfg=asset_cfg)
assigned = env.sim.model.mat_texid[:, mat_ids, role]
assert torch.all(torch.isin(assigned, tex_ids))
assert not torch.equal(assigned, original)
assert torch.isin(assigned, torch.unique(original), invert=True).any()
def test_mat_texid_respects_texture_subset(texid_env):
"""Textures outside the selection are never sampled."""
torch.manual_seed(42)
env = texid_env
robot = env.scene["robot"]
# The pool excludes tex_a and tex_b, which are the two materials' defaults.
asset_cfg = SceneEntityCfg(
"robot", material_names=(".*",), texture_names=("tex_c", "tex_d")
)
asset_cfg.resolve(env.scene)
mat_ids = robot.indexing.mat_ids[asset_cfg.material_ids]
tex_ids = robot.indexing.tex_ids[asset_cfg.texture_ids]
role = mujoco.mjtTextureRole.mjTEXROLE_RGB.value
excluded = torch.unique(env.sim.model.mat_texid[:, mat_ids, role])
dr.mat_texid(env, env_ids=None, asset_cfg=asset_cfg)
assigned = env.sim.model.mat_texid[:, mat_ids, role]
assert torch.all(torch.isin(assigned, tex_ids))
assert not torch.isin(assigned, excluded).any()
assert len(torch.unique(assigned)) >= 2
def test_mat_texid_partial_env_ids(texid_env):
"""Randomizing subset of envs leaves others unchanged."""
torch.manual_seed(42)
env = texid_env
robot = env.scene["robot"]
asset_cfg = SceneEntityCfg(
"robot", material_names=(".*",), texture_names=("tex_c", "tex_d")
)
asset_cfg.resolve(env.scene)
mat_ids = robot.indexing.mat_ids[asset_cfg.material_ids]
tex_ids = robot.indexing.tex_ids[asset_cfg.texture_ids]
role = mujoco.mjtTextureRole.mjTEXROLE_RGB.value
original = env.sim.model.mat_texid[:, mat_ids, role].clone()
dr.mat_texid(
env, env_ids=torch.tensor([0, 2], device=env.device), asset_cfg=asset_cfg
)
result = env.sim.model.mat_texid[:, mat_ids, role]
# Envs 0, 2 changed, the pool excludes their defaults.
assert torch.all(torch.isin(result[0], tex_ids))
assert torch.all(torch.isin(result[2], tex_ids))
assert not torch.equal(result[0], original[0])
assert not torch.equal(result[2], original[2])
# Envs 1, 3 unchanged.
assert torch.all(result[1] == original[1])
assert torch.all(result[3] == original[3])
def test_mat_texid_invalid_texture_name(texid_env):
"""Unknown texture name is rejected during config resolution."""
env = texid_env
with pytest.raises(ValueError, match="nonexistent_texture"):
cfg = SceneEntityCfg(
"robot", material_names=(".*",), texture_names=("nonexistent_texture",)
)
cfg.resolve(env.scene)
dr.mat_texid(env, env_ids=None, asset_cfg=cfg)
def test_mat_texid_empty_texture_selection(texid_env):
"""mat_texid raises when the resolved texture pool is empty."""
env = texid_env
cfg = SceneEntityCfg("robot", material_names=(".*",), texture_names=())
cfg.resolve(env.scene)
with pytest.raises(ValueError, match="No textures selected"):
dr.mat_texid(env, env_ids=None, asset_cfg=cfg)
def test_mat_texid_shared_random(texid_env):
"""All materials within same env get same texture, envs differ."""
torch.manual_seed(42)
env = texid_env
robot = env.scene["robot"]
asset_cfg = SceneEntityCfg("robot", material_names=(".*",), texture_names=(".*",))
asset_cfg.resolve(env.scene)
mat_ids = robot.indexing.mat_ids[asset_cfg.material_ids]
role = mujoco.mjtTextureRole.mjTEXROLE_RGB.value
dr.mat_texid(env, env_ids=None, asset_cfg=asset_cfg, shared_random=True)
assigned = env.sim.model.mat_texid[:, mat_ids, role]
for env_idx in range(env.num_envs):
env_texids = assigned[env_idx]
assert torch.all(env_texids == env_texids[0])
assert len(torch.unique(assigned[:, 0])) > 1
# pair_friction tests.
PAIR_XML = """
<mujoco>
<worldbody>
<body name="base" pos="0 0 1">
<freejoint name="free_joint"/>
<geom name="base_geom" type="box" size="0.1 0.1 0.1" mass="1.0"/>
<body name="foot" pos="0 0 -0.15">
<geom name="foot_geom" type="box" size="0.05 0.05 0.02" mass="0.1"/>
</body>
</body>
<geom name="floor" type="plane" size="2 2 0.1"/>
</worldbody>
<contact>
<pair name="base_floor" geom1="base_geom" geom2="floor"
friction="0.5 0.5 0.005 0.0001 0.0001"/>
<pair name="foot_floor" geom1="foot_geom" geom2="floor"
friction="0.3 0.3 0.003 0.0002 0.0002"/>
</contact>
</mujoco>
"""
@pytest.fixture(scope="module")
def pair_env(device):
entity_cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(PAIR_XML))
scene_cfg = SceneCfg(num_envs=NUM_ENVS, entities={"robot": entity_cfg})
scene = Scene(scene_cfg, device)
model = scene.compile()
sim = Simulation(num_envs=NUM_ENVS, cfg=SimulationCfg(), model=model, device=device)
scene.initialize(model, sim.model, sim.data)
sim.expand_model_fields(("pair_friction",))
return Env(scene, sim, device)
def test_pair_friction_abs(pair_env):
"""Values set directly within range, diverse across envs."""
torch.manual_seed(42)
env = pair_env
robot = env.scene["robot"]
dr.pair_friction(env, env_ids=None, ranges=(0.4, 1.5), operation="abs")
friction = env.sim.model.pair_friction[:, robot.indexing.pair_ids, 0]
assert torch.all((friction >= 0.4) & (friction <= 1.5))
assert len(torch.unique(friction)) >= 2
def test_pair_friction_scale(pair_env):
"""Values = default * sample; no accumulation after 3 calls."""
env = pair_env
robot = env.scene["robot"]
pair_idx = robot.indexing.pair_ids[0]
default_val = env.sim.get_default_field("pair_friction")[pair_idx, 0].item()
for _ in range(3):
dr.pair_friction(
env,
env_ids=None,
ranges=(2.0, 2.0),
operation="scale",
asset_cfg=SceneEntityCfg("robot", pair_ids=[0]),
axes=[0],
)
result = env.sim.model.pair_friction[:, pair_idx, 0]
assert torch.allclose(result, torch.full_like(result, default_val * 2.0), atol=1e-5)
def test_pair_friction_add(pair_env):
"""Values = default + sample; no accumulation after 3 calls."""
env = pair_env
robot = env.scene["robot"]
pair_idx = robot.indexing.pair_ids[0]
default_val = env.sim.get_default_field("pair_friction")[pair_idx, 0].item()
for _ in range(3):
dr.pair_friction(
env,
env_ids=None,
ranges=(0.1, 0.1),
operation="add",
asset_cfg=SceneEntityCfg("robot", pair_ids=[0]),
axes=[0],
)
result = env.sim.model.pair_friction[:, pair_idx, 0]
assert torch.allclose(result, torch.full_like(result, default_val + 0.1), atol=1e-5)
def test_pair_friction_axes_selectivity(pair_env):
"""Only the targeted axis changes; others stay at their original values."""
torch.manual_seed(7)
env = pair_env
robot = env.scene["robot"]
before = env.sim.model.pair_friction.clone()
dr.pair_friction(env, env_ids=None, ranges=(0.8, 1.2), operation="abs", axes=[0])
after = env.sim.model.pair_friction
pair_ids = robot.indexing.pair_ids
# Axis 0 changed.
assert not torch.allclose(after[:, pair_ids, 0], before[:, pair_ids, 0])
# Axes 1-4 unchanged.
for ax in range(1, 5):
assert torch.allclose(after[:, pair_ids, ax], before[:, pair_ids, ax])
def test_pair_friction_multi_axis(pair_env):
"""axes=[0, 1] randomizes both tangent components; axes 2-4 stay unchanged."""
torch.manual_seed(17)
env = pair_env
robot = env.scene["robot"]
before = env.sim.model.pair_friction.clone()
dr.pair_friction(env, env_ids=None, ranges=(0.4, 1.0), operation="abs", axes=[0, 1])
after = env.sim.model.pair_friction
pair_ids = robot.indexing.pair_ids
# Both tangent axes changed.
assert not torch.allclose(after[:, pair_ids, 0], before[:, pair_ids, 0])
assert not torch.allclose(after[:, pair_ids, 1], before[:, pair_ids, 1])
# Axes 2-4 unchanged.
for ax in range(2, 5):
assert torch.allclose(after[:, pair_ids, ax], before[:, pair_ids, ax])
# Axis 0 and axis 1 are sampled independently (not forced equal).
assert not torch.allclose(after[:, pair_ids, 0], after[:, pair_ids, 1])
def test_pair_friction_partial_env_ids(pair_env):
"""Randomizing a subset of envs leaves the others unchanged."""
torch.manual_seed(99)
env = pair_env
robot = env.scene["robot"]
before = env.sim.model.pair_friction.clone()
randomized_ids = torch.tensor([0, 1], device=env.device)
dr.pair_friction(
env, env_ids=randomized_ids, ranges=(2.0, 3.0), operation="abs", axes=[0]
)
after = env.sim.model.pair_friction
pair_ids = robot.indexing.pair_ids
unchanged_ids = torch.tensor([2, 3], device=env.device)
assert torch.allclose(
after[unchanged_ids[:, None], pair_ids[None, :], :],
before[unchanged_ids[:, None], pair_ids[None, :], :],
)
def test_pair_friction_shared_random(pair_env):
"""All pairs in the same env get the same value; envs differ."""
torch.manual_seed(5)
env = pair_env
robot = env.scene["robot"]
dr.pair_friction(
env, env_ids=None, ranges=(0.5, 2.0), operation="abs", axes=[0], shared_random=True
)
pair_ids = robot.indexing.pair_ids
friction = env.sim.model.pair_friction[:, pair_ids, 0]
# Within each env, all pairs have the same value.
for e in range(env.num_envs):
row = friction[e]
assert torch.allclose(row, row[0].expand_as(row), atol=1e-6)
# Across envs, values differ.
assert not torch.allclose(friction[0], friction[1])
def test_pair_friction_by_name(pair_env):
"""Targeting pair by name only modifies that pair."""
torch.manual_seed(11)
env = pair_env
robot = env.scene["robot"]
foot_cfg = SceneEntityCfg("robot", pair_names=("foot_floor",))
foot_cfg.resolve(env.scene)
base_cfg = SceneEntityCfg("robot", pair_names=("base_floor",))
base_cfg.resolve(env.scene)
foot_pair_id = robot.indexing.pair_ids[foot_cfg.pair_ids]
base_pair_id = robot.indexing.pair_ids[base_cfg.pair_ids]
before = env.sim.model.pair_friction.clone()
dr.pair_friction(
env, env_ids=None, ranges=(1.5, 2.5), operation="abs", asset_cfg=foot_cfg
)
after = env.sim.model.pair_friction
# foot_floor changed.
assert not torch.allclose(after[:, foot_pair_id, 0], before[:, foot_pair_id, 0])
# base_floor unchanged.
assert torch.allclose(after[:, base_pair_id, :], before[:, base_pair_id, :])
def test_pair_friction_string_ranges(pair_env):
"""Dict with pair name patterns randomizes the matched pairs correctly."""
torch.manual_seed(13)
env = pair_env
robot = env.scene["robot"]
dr.pair_friction(
env,
env_ids=None,
ranges={"base_floor": (1.0, 1.0), "foot_floor": (0.5, 0.5)},
operation="abs",
axes=[0],
)
base_cfg = SceneEntityCfg("robot", pair_names=("base_floor",))
base_cfg.resolve(env.scene)
foot_cfg = SceneEntityCfg("robot", pair_names=("foot_floor",))
foot_cfg.resolve(env.scene)
base_id = robot.indexing.pair_ids[base_cfg.pair_ids]
foot_id = robot.indexing.pair_ids[foot_cfg.pair_ids]
assert torch.allclose(
env.sim.model.pair_friction[:, base_id, 0],
torch.ones(env.num_envs, device=env.device),
atol=1e-5,
)
assert torch.allclose(
env.sim.model.pair_friction[:, foot_id, 0],
torch.full((env.num_envs,), 0.5, device=env.device),
atol=1e-5,
)
def test_pair_friction_isotropic_tangent(pair_env):
"""isotropic=True mirrors tangent2 = tangent1; spin and roll axes unchanged."""
torch.manual_seed(21)
env = pair_env
robot = env.scene["robot"]
pair_ids = robot.indexing.pair_ids
before = env.sim.model.pair_friction.clone()
dr.pair_friction(
env, env_ids=None, ranges=(0.4, 1.0), operation="abs", axes=[0], isotropic=True
)
after = env.sim.model.pair_friction
# Axis 0 was randomized into range.
assert torch.all((after[:, pair_ids, 0] >= 0.4) & (after[:, pair_ids, 0] <= 1.0))
# Axis 1 equals axis 0 exactly.
assert torch.allclose(after[:, pair_ids, 1], after[:, pair_ids, 0])
# Axes 2, 3, 4 untouched.
for ax in (2, 3, 4):
assert torch.allclose(after[:, pair_ids, ax], before[:, pair_ids, ax])
def test_pair_friction_isotropic_shared(pair_env):
"""isotropic=True + shared_random=True: all pairs same value, tangent2 == tangent1."""
torch.manual_seed(33)
env = pair_env
robot = env.scene["robot"]
pair_ids = robot.indexing.pair_ids
dr.pair_friction(
env,
env_ids=None,
ranges=(0.4, 1.0),
operation="abs",
axes=[0],
shared_random=True,
isotropic=True,
)
friction = env.sim.model.pair_friction[:, pair_ids, :]
# All pairs in each env share the same tangent1 value.
for e in range(env.num_envs):
assert torch.allclose(
friction[e, :, 0], friction[e, 0, 0].expand(len(pair_ids)), atol=1e-6
)
# Across envs, values differ.
assert not torch.allclose(friction[0, 0, 0], friction[1, 0, 0])
# tangent2 == tangent1 everywhere.
assert torch.allclose(friction[:, :, 1], friction[:, :, 0])
def test_pair_friction_isotropic_roll(pair_env):
"""isotropic=True with roll axes mirrors roll2 = roll1; tangent axes unchanged."""
torch.manual_seed(41)
env = pair_env
robot = env.scene["robot"]
pair_ids = robot.indexing.pair_ids
before = env.sim.model.pair_friction.clone()
dr.pair_friction(
env, env_ids=None, ranges=(0.001, 0.01), operation="abs", axes=[3], isotropic=True
)
after = env.sim.model.pair_friction
# Axis 3 was randomized.
assert not torch.allclose(after[:, pair_ids, 3], before[:, pair_ids, 3])
# Axis 4 equals axis 3 exactly.
assert torch.allclose(after[:, pair_ids, 4], after[:, pair_ids, 3])
# Tangent and spin axes untouched.
for ax in (0, 1, 2):
assert torch.allclose(after[:, pair_ids, ax], before[:, pair_ids, ax])
def test_pair_friction_isotropic_roll_dict_ranges(pair_env):
"""isotropic=True with dict-int ranges containing roll axis triggers roll2 = roll1."""
torch.manual_seed(53)
env = pair_env
robot = env.scene["robot"]
pair_ids = robot.indexing.pair_ids
before = env.sim.model.pair_friction.clone()
dr.pair_friction(
env,
env_ids=None,
ranges={3: (0.001, 0.01)},
operation="abs",
isotropic=True,
)
after = env.sim.model.pair_friction
# Axis 3 was randomized.
assert not torch.allclose(after[:, pair_ids, 3], before[:, pair_ids, 3])
# Axis 4 equals axis 3 exactly (dict-int ranges must trigger the roll copy).
assert torch.allclose(after[:, pair_ids, 4], after[:, pair_ids, 3])
# Tangent and spin axes untouched.
for ax in (0, 1, 2):
assert torch.allclose(after[:, pair_ids, ax], before[:, pair_ids, ax])
def test_pair_friction_invalid_name(pair_env):
"""ValueError for unknown pair name."""
env = pair_env
with pytest.raises(ValueError, match="nonexistent_pair"):
cfg = SceneEntityCfg("robot", pair_names=("nonexistent_pair",))
cfg.resolve(env.scene)
Reference in New Issue Copy Permalink