32a241c28f
Some checks failed
nightly / Test against latest dependencies (py3.10) (push) Has been cancelled
nightly / Test against latest dependencies (py3.13) (push) Has been cancelled
tests / tests (3.13, locked) (push) Has been cancelled
tests / tests (3.13, unlocked) (push) Has been cancelled
tests / pyright (3.10) (push) Has been cancelled
tests / lint-format (push) Has been cancelled
tests / tests (3.10, locked) (push) Has been cancelled
tests / tests (3.11, locked) (push) Has been cancelled
tests / tests (3.12, locked) (push) Has been cancelled
tests / pyright (3.11) (push) Has been cancelled
tests / pyright (3.12) (push) Has been cancelled
tests / pyright (3.13) (push) Has been cancelled
tests / ty-check (3.10) (push) Has been cancelled
tests / ty-check (3.11) (push) Has been cancelled
tests / ty-check (3.12) (push) Has been cancelled
tests / ty-check (3.13) (push) Has been cancelled
tests / stubs (push) Has been cancelled
tests / smoke-test (push) Has been cancelled
Docker / check_paths (push) Has been cancelled
docs / build (push) Has been cancelled
Docker / build (push) Has been cancelled
Upstream: https://github.com/michaelgillett/mjlab Upstream-Commit: c19f713c415a699a79d71cd96aa13c3104a05047 Upstream-Branch: main
2470 lines
75 KiB
Python
2470 lines
75 KiB
Python
"""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)
|