Robocon/RC_WheelLeg
4
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[train] 更新新MJCF与第一版完整训练框架

Browse Source
This commit is contained in:
eastzio
2026-07-27 12:31:21 +08:00
parent dfd646cf47
commit b08956aec7
116 changed files with 7649 additions and 6760 deletions
Download Patch File Download Diff File Expand all files Collapse all files
05_software/train/rc_mjlab/mjlab/tests/conftest.py
+12 -4
View File
@@ -106,17 +106,25 @@ def initialize_entity(entity: Entity, device: str, num_envs: int = 1):
def make_scene_and_sim(
device: str,
xml: str,
xml: str | dict[str, str],
sensors: tuple,
num_envs: int = 1,
sim_cfg: SimulationCfg | None = None,
) -> tuple[Scene, Simulation]:
"""Create a scene and simulation from inline XML with sensors wired up."""
entity_cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(xml))
"""Create a scene and simulation from inline XML with sensors wired up.
``xml`` may be a single XML string (registered as the ``robot`` entity) or a
mapping of entity name to XML string for multi-entity scenes.
"""
xml_by_entity = {"robot": xml} if isinstance(xml, str) else xml
entities = {
name: EntityCfg(spec_fn=lambda s=s: mujoco.MjSpec.from_string(s))
for name, s in xml_by_entity.items()
}
scene_cfg = SceneCfg(
num_envs=num_envs,
env_spacing=5.0,
entities={"robot": entity_cfg},
entities=entities,
sensors=sensors,
)
scene = Scene(scene_cfg, device)
05_software/train/rc_mjlab/mjlab/tests/test_builtin_dcmotor_actuator.py
+653
View File
@@ -0,0 +1,653 @@
"""Tests for BuiltinDcMotorActuator.
Covers wiring of MuJoCo's native ``<dcmotor>`` element through mjlab: the
three input modes (voltage / position / velocity), torque saturation,
config validation, and DR integration.
"""
import math
from unittest.mock import Mock
import mujoco
import pytest
import torch
from conftest import (
create_entity_with_actuator,
get_test_device,
initialize_entity,
load_fixture_xml,
)
from mjlab.actuator import (
BuiltinDcMotorActuator,
BuiltinDcMotorActuatorCfg,
DcMotorDatasheetParams,
DcMotorInputMode,
DcMotorPhysicalParams,
)
from mjlab.actuator.actuator import TransmissionType
from mjlab.entity import Entity, EntityArticulationInfoCfg, 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
ROBOT_XML = load_fixture_xml("floating_base_articulated")
# Motor characterization used throughout (resolves to K=0.24, R=2.88).
V_NOM, TAU_STALL, OMEGA_NL = 24.0, 2.0, 100.0
K = V_NOM / OMEGA_NL
R = K * V_NOM / TAU_STALL
@pytest.fixture(scope="module")
def device():
return get_test_device()
DATASHEET = DcMotorDatasheetParams(
nominal_voltage=V_NOM, stall_torque=TAU_STALL, no_load_speed=OMEGA_NL
)
def _make_cfg(
*,
mode: DcMotorInputMode = DcMotorInputMode.POSITION,
stiffness=5.0,
damping=0.5,
voltage_limit=24.0,
**extra,
) -> BuiltinDcMotorActuatorCfg:
"""Build a cfg with sensible PID defaults. ``extra`` forwards any other
BuiltinDcMotorActuatorCfg kwarg (effort_limit, integral_gain, thermal,
delay_*, etc.)."""
return BuiltinDcMotorActuatorCfg(
target_names_expr=("joint.*",),
mode=mode,
motor_params=DATASHEET,
stiffness=stiffness,
damping=damping,
voltage_limit=voltage_limit,
**extra,
)
def _make_entity(**kwargs) -> Entity:
return create_entity_with_actuator(ROBOT_XML, _make_cfg(**kwargs))
def _make_initialized(device, **kwargs):
"""Build entity from cfg kwargs and initialize it through the sim."""
return initialize_entity(_make_entity(**kwargs), device)
def _drive(
entity: Entity,
sim,
device: str,
*,
pos_target=None,
vel_target=None,
effort_target=None,
q0=None,
qd0=None,
) -> None:
zero = torch.zeros(1, 2, device=device)
entity.write_joint_state_to_sim(
position=q0 if q0 is not None else zero,
velocity=qd0 if qd0 is not None else zero,
)
entity.set_joint_position_target(pos_target if pos_target is not None else zero)
entity.set_joint_velocity_target(vel_target if vel_target is not None else zero)
entity.set_joint_effort_target(effort_target if effort_target is not None else zero)
entity.write_data_to_sim()
sim.forward()
# Wiring sanity.
def test_kr_packed_into_gainprm(device):
"""The XML compiler derives K and R from the nominal triplet."""
_, sim = initialize_entity(_make_entity(effort_limit=1.5), device)
m = sim.mj_model
for i in range(2):
assert m.actuator_gainprm[i, 0] == pytest.approx(R, abs=1e-6)
assert m.actuator_gainprm[i, 1] == pytest.approx(K, abs=1e-6)
assert m.actuator_gainprm[i, 4] == pytest.approx(5.0) # kp
assert m.actuator_gainprm[i, 6] == pytest.approx(0.5) # kd
assert m.actuator_gainprm[i, 7] == pytest.approx(24.0) # Vmax
assert m.actuator_gainprm[i, 8] == pytest.approx(1.0) # input_mode=position
assert m.actuator_gaintype[i] == mujoco.mjtGain.mjGAIN_DCMOTOR
assert m.actuator_biastype[i] == mujoco.mjtBias.mjBIAS_DCMOTOR
# No activation state: ki=0, no inductance, no thermal/lugre/slew.
assert m.actuator_actnum[i] == 0
def test_motor_const_path(device):
"""Physical params pack K = sqrt(Kt*Ke) and R verbatim."""
cfg = BuiltinDcMotorActuatorCfg(
target_names_expr=("joint.*",),
mode=DcMotorInputMode.VOLTAGE,
motor_params=DcMotorPhysicalParams(kt=0.1, ke=0.05, resistance=2.0),
)
_, sim = initialize_entity(create_entity_with_actuator(ROBOT_XML, cfg), device)
m = sim.mj_model
for i in range(2):
assert m.actuator_gainprm[i, 0] == pytest.approx(2.0, abs=1e-6)
assert m.actuator_gainprm[i, 1] == pytest.approx((0.1 * 0.05) ** 0.5, abs=1e-6)
# Stateless motor physics.
def test_voltage_mode_steady_state(device):
"""At rest, ctrl = V -> tau = K * V / R."""
cfg = BuiltinDcMotorActuatorCfg(
target_names_expr=("joint.*",),
mode=DcMotorInputMode.VOLTAGE,
motor_params=DATASHEET,
)
entity, sim = initialize_entity(create_entity_with_actuator(ROBOT_XML, cfg), device)
V = torch.tensor([[10.0, -5.0]], device=device)
_drive(entity, sim, device, effort_target=V)
v_adr = entity.indexing.joint_v_adr
expected = K * V[0] / R
assert torch.allclose(sim.data.qfrc_actuator[0, v_adr], expected, atol=1e-4)
def test_voltage_mode_voltage_limit_zero_is_noop(device):
"""Docstring promises ``voltage_limit=0`` disables clamping. Verify against
MuJoCo's ``dcmotor_voltage`` (which only clamps when ``Vmax > 0``)."""
cfg = BuiltinDcMotorActuatorCfg(
target_names_expr=("joint.*",),
mode=DcMotorInputMode.VOLTAGE,
motor_params=DATASHEET,
voltage_limit=0.0,
)
entity, sim = initialize_entity(create_entity_with_actuator(ROBOT_XML, cfg), device)
V = torch.tensor([[1000.0, 0.0]], device=device) # absurdly high voltage.
_drive(entity, sim, device, effort_target=V)
v_adr = entity.indexing.joint_v_adr
expected = K * V[0] / R
assert torch.allclose(sim.data.qfrc_actuator[0, v_adr], expected, atol=1e-2)
def test_back_emf_reduces_torque_at_velocity(device):
"""Same V, joint moving at omega: tau = K * (V - K * omega) / R."""
cfg = BuiltinDcMotorActuatorCfg(
target_names_expr=("joint.*",),
mode=DcMotorInputMode.VOLTAGE,
motor_params=DATASHEET,
)
entity, sim = initialize_entity(create_entity_with_actuator(ROBOT_XML, cfg), device)
V = torch.tensor([[10.0, 0.0]], device=device)
omega0 = torch.tensor([[2.0, 0.0]], device=device)
_drive(entity, sim, device, effort_target=V, qd0=omega0)
v_adr = entity.indexing.joint_v_adr
expected = K * (V[0] - K * omega0[0]) / R
assert torch.allclose(sim.data.qfrc_actuator[0, v_adr], expected, atol=1e-4)
def test_position_mode_pid_at_rest(device):
"""kd=0, no Vmax clamp: tau = K * kp * (target - q) / R."""
# voltage_limit must be >0 (cfg invariant), pick it big enough not to clamp.
entity, sim = initialize_entity(
_make_entity(damping=0.0, voltage_limit=1000.0), device
)
pos = torch.tensor([[0.1, -0.05]], device=device)
_drive(entity, sim, device, pos_target=pos)
v_adr = entity.indexing.joint_v_adr
expected = K * 5.0 * pos[0] / R
assert torch.allclose(sim.data.qfrc_actuator[0, v_adr], expected, atol=1e-4)
def test_position_mode_voltage_clamp(device):
"""Huge position error -> PID voltage saturates at Vmax."""
entity, sim = initialize_entity(
_make_entity(stiffness=100.0, damping=0.0, voltage_limit=2.0),
device,
)
# kp * err = 100 * 0.5 = 50 V, well above Vmax=2.
pos = torch.tensor([[0.5, 0.0]], device=device)
_drive(entity, sim, device, pos_target=pos)
v_adr = entity.indexing.joint_v_adr
qfrc = sim.data.qfrc_actuator[0, v_adr]
expected_first = K * 2.0 / R # tau at clamped V.
assert qfrc[0].item() == pytest.approx(expected_first, abs=1e-4)
assert qfrc[1].item() == pytest.approx(0.0, abs=1e-4)
def test_velocity_mode_pid(device):
"""P-only velocity tracking: tau = K * kp * (target - qdot) / R."""
entity, sim = initialize_entity(
_make_entity(mode=DcMotorInputMode.VELOCITY, damping=0.0, voltage_limit=1000.0),
device,
)
qd0 = torch.tensor([[1.0, 0.0]], device=device)
vel_target = torch.tensor([[3.0, 0.0]], device=device)
_drive(entity, sim, device, vel_target=vel_target, qd0=qd0)
v_adr = entity.indexing.joint_v_adr
# back-EMF subtracts K*omega; this is folded into the dcmotor bias.
# voltage = kp*(target - qdot); tau = K*(voltage - K*omega)/R.
voltage = 5.0 * (vel_target[0] - qd0[0])
expected = K * (voltage - K * qd0[0]) / R
assert torch.allclose(sim.data.qfrc_actuator[0, v_adr], expected, atol=1e-4)
def test_effort_limit_clamps_torque(device):
"""forcerange clamps the algebraic torque output."""
entity, sim = initialize_entity(
_make_entity(stiffness=100.0, damping=0.0, voltage_limit=1000.0, effort_limit=0.1),
device,
)
m = sim.mj_model
for i in range(2):
assert m.actuator_forcelimited[i] == 1
assert m.actuator_forcerange[i, 0] == pytest.approx(-0.1)
assert m.actuator_forcerange[i, 1] == pytest.approx(0.1)
# Unclamped tau would be K * 100 * 0.5 / R ~= K*50/R, well above 0.1.
pos = torch.tensor([[0.5, 0.0]], device=device)
_drive(entity, sim, device, pos_target=pos)
v_adr = entity.indexing.joint_v_adr
qfrc = sim.data.qfrc_actuator[0, v_adr]
assert qfrc[0].item() == pytest.approx(0.1, abs=1e-4)
assert qfrc[1].item() == pytest.approx(0.0, abs=1e-4)
# Cogging.
def test_cogging_packed_into_biasprm(device):
"""``cogging=(A, Np, phi)`` packs into ``biasprm[0:3]``."""
cfg = BuiltinDcMotorActuatorCfg(
target_names_expr=("joint.*",),
mode=DcMotorInputMode.VOLTAGE,
motor_params=DATASHEET,
cogging=(0.5, 4.0, 0.1),
)
_, sim = initialize_entity(create_entity_with_actuator(ROBOT_XML, cfg), device)
m = sim.mj_model
for i in range(2):
assert m.actuator_biasprm[i, 0] == pytest.approx(0.5)
assert m.actuator_biasprm[i, 1] == pytest.approx(4.0)
assert m.actuator_biasprm[i, 2] == pytest.approx(0.1)
def test_cogging_contributes_torque(device):
"""At ctrl=0 (no electromagnetic torque), qfrc_actuator equals the cogging
term ``A * sin(Np * q + phi)`` evaluated at the joint angle."""
A, Np, phi = 0.5, 4.0, 0.1
cfg = BuiltinDcMotorActuatorCfg(
target_names_expr=("joint.*",),
mode=DcMotorInputMode.VOLTAGE,
motor_params=DATASHEET,
cogging=(A, Np, phi),
)
entity, sim = initialize_entity(create_entity_with_actuator(ROBOT_XML, cfg), device)
q0, q1 = 0.3, -0.2
_drive(
entity,
sim,
device,
q0=torch.tensor([[q0, q1]], device=device),
effort_target=torch.zeros(1, 2, device=device),
)
v_adr = entity.indexing.joint_v_adr
qfrc = sim.data.qfrc_actuator[0, v_adr]
assert qfrc[0].item() == pytest.approx(A * math.sin(Np * q0 + phi), abs=1e-5)
assert qfrc[1].item() == pytest.approx(A * math.sin(Np * q1 + phi), abs=1e-5)
def test_cogging_bypasses_effort_limit(device):
"""Cogging is added *after* the forcerange clamp (MuJoCo's intentional
model: ``effort_limit`` bounds electromagnetic torque, cogging is
mechanical). Total torque can exceed ``effort_limit`` by up to the
cogging amplitude."""
A, Np, phi = 0.5, 0.0, math.pi / 2 # sin(pi/2)=1, so cogging = A at any q.
cfg = BuiltinDcMotorActuatorCfg(
target_names_expr=("joint.*",),
mode=DcMotorInputMode.VOLTAGE,
motor_params=DATASHEET,
cogging=(A, Np, phi),
effort_limit=0.05, # An order of magnitude below A.
)
entity, sim = initialize_entity(create_entity_with_actuator(ROBOT_XML, cfg), device)
# Pick a voltage large enough that the electromagnetic torque alone
# would saturate forcerange at +/- 0.05.
V = torch.tensor([[100.0, 0.0]], device=device)
_drive(entity, sim, device, effort_target=V)
v_adr = entity.indexing.joint_v_adr
qfrc = sim.data.qfrc_actuator[0, v_adr]
# joint1: electromagnetic clamped to +0.05, plus cogging A=0.5.
assert qfrc[0].item() == pytest.approx(0.05 + A, abs=1e-5)
# joint2: zero voltage, electromagnetic=0, only cogging.
assert qfrc[1].item() == pytest.approx(A, abs=1e-5)
# Optional stateful extensions (integral, slew, inductance, thermal, LuGre).
# Each behavior check compares against a baseline with the feature disabled
# so that removing the wiring in edit_spec causes the comparison to fail.
def _step_n(entity, sim, device, n: int, *, pos_target=None, eff_target=None):
zero = torch.zeros(1, 2, device=device)
entity.write_joint_state_to_sim(position=zero, velocity=zero)
for _ in range(n):
entity.set_joint_position_target(pos_target if pos_target is not None else zero)
entity.set_joint_velocity_target(zero)
entity.set_joint_effort_target(eff_target if eff_target is not None else zero)
entity.write_data_to_sim()
sim.step()
def _qfrc(entity, sim) -> torch.Tensor:
return sim.data.qfrc_actuator[0, entity.indexing.joint_v_adr].clone()
def test_integral_gain_ramps_torque(device):
"""Integrator in position mode ramps torque over time even with ``kp``
and ``kd`` near zero."""
# stiffness must be > 0 (validation); choose tiny so ki dominates.
base = dict(
mode=DcMotorInputMode.POSITION, stiffness=1e-4, damping=0.0, voltage_limit=24.0
)
ent_off, sim_off = _make_initialized(device, **base, integral_gain=0.0)
ent_on, sim_on = _make_initialized(device, **base, integral_gain=10.0)
target = torch.tensor([[0.5, 0.0]], device=device)
for sim, ent in ((sim_off, ent_off), (sim_on, ent_on)):
_step_n(ent, sim, device, n=20, pos_target=target)
assert _qfrc(ent_on, sim_on)[0].abs() > 100.0 * _qfrc(ent_off, sim_off)[0].abs()
def test_slew_rate_limits_voltage(device):
"""``slew_rate`` rate-limits ``ctrl``: after one step, effective voltage
is far below the requested input."""
base = dict(
mode=DcMotorInputMode.VOLTAGE, stiffness=0.0, damping=0.0, voltage_limit=0.0
)
ent_off, sim_off = _make_initialized(device, **base, slew_rate=0.0)
ent_on, sim_on = _make_initialized(device, **base, slew_rate=10.0)
V = torch.tensor([[100.0, 0.0]], device=device)
for sim, ent in ((sim_off, ent_off), (sim_on, ent_on)):
_step_n(ent, sim, device, n=1, eff_target=V)
assert _qfrc(ent_off, sim_off)[0] > 100.0 * _qfrc(ent_on, sim_on)[0]
def test_inductance_lags_current(device):
"""Large ``inductance`` (te >> dt) suppresses early-step torque."""
base = dict(
mode=DcMotorInputMode.VOLTAGE, stiffness=0.0, damping=0.0, voltage_limit=0.0
)
ent_off, sim_off = _make_initialized(device, **base, inductance=0.0)
ent_on, sim_on = _make_initialized(device, **base, inductance=1.0)
V = torch.tensor([[10.0, 0.0]], device=device)
for sim, ent in ((sim_off, ent_off), (sim_on, ent_on)):
_step_n(ent, sim, device, n=2, eff_target=V)
assert _qfrc(ent_off, sim_off)[0].abs() > 10.0 * _qfrc(ent_on, sim_on)[0].abs()
def test_thermal_decays_torque(device):
"""I^2R heating raises T, which raises effective resistance and decays
torque over time."""
# Params chosen for visible effect in a handful of steps without going
# numerically unstable: small C (fast heating) and modest alpha.
base = dict(
mode=DcMotorInputMode.VOLTAGE, stiffness=0.0, damping=0.0, voltage_limit=0.0
)
ent_off, sim_off = _make_initialized(device, **base)
ent_on, sim_on = _make_initialized(
device, **base, thermal=(1.0, 0.1, 0.0, 0.01, 0.0, 0.0)
)
V = torch.tensor([[100.0, 0.0]], device=device)
for sim, ent in ((sim_off, ent_off), (sim_on, ent_on)):
_step_n(ent, sim, device, n=5, eff_target=V)
assert _qfrc(ent_on, sim_on)[0].abs() < 0.5 * _qfrc(ent_off, sim_off)[0].abs()
def test_lugre_subtracts_friction(device):
"""LuGre friction subtracts a velocity-dependent force after the
``effort_limit`` clamp (mechanical, like cogging)."""
# Static comparison at v>0, ctrl=0; avoids feedback between LuGre slowing
# the joint and back-EMF easing off under sim.step().
# no LuGre: qfrc = -K^2 * v / R (back-EMF only)
# w/ LuGre: qfrc = -K^2 * v / R - sigma1*v - ...
base = dict(
mode=DcMotorInputMode.VOLTAGE, stiffness=0.0, damping=0.0, voltage_limit=0.0
)
ent_off, sim_off = _make_initialized(device, **base)
ent_on, sim_on = _make_initialized(
device, **base, lugre=(1e4, 100.0, 0.1, 0.15, 0.01)
)
zero = torch.zeros(1, 2, device=device)
v0 = torch.tensor([[1.0, 0.0]], device=device)
for sim, ent in ((sim_off, ent_off), (sim_on, ent_on)):
ent.write_joint_state_to_sim(position=zero, velocity=v0)
ent.set_joint_position_target(zero)
ent.set_joint_velocity_target(zero)
ent.set_joint_effort_target(zero)
ent.write_data_to_sim()
sim.forward()
assert abs(_qfrc(ent_on, sim_on)[0]) > 100.0 * abs(_qfrc(ent_off, sim_off)[0])
# Config validation.
def test_pid_mode_requires_gains():
with pytest.raises(ValueError, match="stiffness"):
BuiltinDcMotorActuatorCfg(
target_names_expr=("j",),
mode=DcMotorInputMode.POSITION,
motor_params=DATASHEET,
voltage_limit=1.0,
)
with pytest.raises(ValueError, match="voltage_limit"):
BuiltinDcMotorActuatorCfg(
target_names_expr=("j",),
mode=DcMotorInputMode.POSITION,
motor_params=DATASHEET,
stiffness=1.0,
)
def test_voltage_mode_rejects_pid_gains():
with pytest.raises(ValueError, match="VOLTAGE"):
BuiltinDcMotorActuatorCfg(
target_names_expr=("j",),
mode=DcMotorInputMode.VOLTAGE,
motor_params=DATASHEET,
stiffness=1.0,
)
def test_site_rejected():
with pytest.raises(ValueError, match="SITE"):
BuiltinDcMotorActuatorCfg(
target_names_expr=("j",),
motor_params=DATASHEET,
stiffness=1.0,
voltage_limit=1.0,
transmission_type=TransmissionType.SITE,
)
# Joint-level passthrough.
def test_armature_applied(device):
_, sim = initialize_entity(_make_entity(armature=0.7), device)
m = sim.mj_model
for jname in ("joint1", "joint2"):
dof_id = m.jnt_dofadr[m.joint(jname).id]
assert m.dof_armature[dof_id] == pytest.approx(0.7)
# Domain randomization.
def _scene_env(
device,
num_envs=2,
mode: DcMotorInputMode = DcMotorInputMode.POSITION,
):
def spec_fn():
spec = mujoco.MjSpec.from_string(ROBOT_XML)
for a in list(spec.actuators):
spec.delete(a)
return spec
entity_cfg = EntityCfg(
spec_fn=spec_fn,
articulation=EntityArticulationInfoCfg(
actuators=(
BuiltinDcMotorActuatorCfg(
target_names_expr=("joint.*",),
mode=mode,
motor_params=DATASHEET,
stiffness=5.0 if mode != DcMotorInputMode.VOLTAGE else 0.0,
damping=0.5 if mode != DcMotorInputMode.VOLTAGE else 0.0,
voltage_limit=24.0 if mode != DcMotorInputMode.VOLTAGE else 0.0,
effort_limit=50.0,
),
)
),
)
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)
env = Mock()
env.num_envs = num_envs
env.device = device
env.scene = {"robot": scene["robot"]}
env.sim = sim
return env
@pytest.mark.parametrize(
"operation, kp_in, kd_in, kp_expected, kd_expected",
[
# scale: multiplies the configured defaults (kp=5.0, kd=0.5).
("scale", 2.0, 3.0, 2.0 * 5.0, 3.0 * 0.5),
# abs: writes the value directly.
("abs", 10.0, 2.0, 10.0, 2.0),
],
)
def test_dr_pd_gains_position_mode(
device, operation, kp_in, kd_in, kp_expected, kd_expected
):
env = _scene_env(device)
robot = env.scene["robot"]
act = robot.actuators[0]
assert isinstance(act, BuiltinDcMotorActuator)
ctrl_ids = act.global_ctrl_ids
env.sim.expand_model_fields(("actuator_gainprm", "actuator_biasprm"))
dr.pd_gains(
env,
env_ids=torch.tensor([0], device=device),
kp_range=(kp_in, kp_in),
kd_range=(kd_in, kd_in),
asset_cfg=SceneEntityCfg("robot"),
operation=operation,
)
m = env.sim.model
n = len(ctrl_ids)
assert torch.allclose(
m.actuator_gainprm[0, ctrl_ids, 4], torch.full((n,), kp_expected, device=device)
)
assert torch.allclose(
m.actuator_gainprm[0, ctrl_ids, 6], torch.full((n,), kd_expected, device=device)
)
# Other env untouched (cfg defaults).
assert torch.allclose(m.actuator_gainprm[1, ctrl_ids, 4], torch.tensor(5.0))
assert torch.allclose(m.actuator_gainprm[1, ctrl_ids, 6], torch.tensor(0.5))
def test_dr_pd_gains_voltage_mode_rejected(device):
env = _scene_env(device, mode=DcMotorInputMode.VOLTAGE)
env.sim.expand_model_fields(("actuator_gainprm", "actuator_biasprm"))
with pytest.raises(ValueError, match="VOLTAGE"):
dr.pd_gains(
env,
env_ids=torch.tensor([0], device=device),
kp_range=(1.0, 1.0),
kd_range=(1.0, 1.0),
asset_cfg=SceneEntityCfg("robot"),
)
def test_dr_effort_limits_writes_forcerange(device):
env = _scene_env(device)
robot = env.scene["robot"]
act = robot.actuators[0]
ctrl_ids = act.global_ctrl_ids
env.sim.expand_model_fields(
("actuator_forcerange", "jnt_actfrcrange", "tendon_actfrcrange")
)
dr.effort_limits(
env,
env_ids=torch.tensor([0], device=device),
effort_limit_range=(123.0, 123.0),
asset_cfg=SceneEntityCfg("robot"),
operation="abs",
)
m = env.sim.model
n = len(ctrl_ids)
assert torch.allclose(
m.actuator_forcerange[0, ctrl_ids, 0],
torch.full((n,), -123.0, device=device),
)
assert torch.allclose(
m.actuator_forcerange[0, ctrl_ids, 1],
torch.full((n,), 123.0, device=device),
)
# Env 1 keeps the configured default of 50.
assert torch.allclose(m.actuator_forcerange[1, ctrl_ids, 1], torch.tensor(50.0))
# Delay.
def test_delay_position_mode(device):
"""A 2-step lag should make position-mode torque reference step-0 target."""
entity, sim = initialize_entity(
_make_entity(
stiffness=10.0,
damping=0.0,
voltage_limit=1000.0,
delay_min_lag=2,
delay_max_lag=2,
),
device,
)
zero = torch.zeros(1, 2, device=device)
entity.write_joint_state_to_sim(position=zero, velocity=zero)
targets = [
torch.tensor([[0.1, 0.0]], device=device),
torch.tensor([[0.3, 0.0]], device=device),
torch.tensor([[0.5, 0.0]], device=device),
]
for p in targets:
entity.set_joint_position_target(p)
entity.set_joint_velocity_target(zero)
entity.set_joint_effort_target(zero)
entity.write_data_to_sim()
sim.forward()
v_adr = entity.indexing.joint_v_adr
# With lag=2 and three writes, the effective target is targets[0].
expected = K * 10.0 * targets[0][0] / R
assert torch.allclose(sim.data.qfrc_actuator[0, v_adr], expected, atol=1e-4)
05_software/train/rc_mjlab/mjlab/tests/test_builtin_pd_actuator.py
+468
View File
@@ -0,0 +1,468 @@
"""Tests for BuiltinPdActuator.
Covers the unique surface of the actuator: paired <position>/<velocity>
elements per target, joint/tendon-level actfrcrange sum-clamp, DR for both
gains and effort limits, delay synchronization, and the ordering invariant
that DR depends on.
"""
from unittest.mock import Mock
import mujoco
import pytest
import torch
from conftest import (
create_entity_with_actuator,
get_test_device,
initialize_entity,
load_fixture_xml,
)
from mjlab.actuator import BuiltinPdActuator, BuiltinPdActuatorCfg
from mjlab.actuator.actuator import TransmissionType
from mjlab.entity import Entity, EntityArticulationInfoCfg, 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
ROBOT_XML = load_fixture_xml("floating_base_articulated")
KP = 100.0
KD = 10.0
@pytest.fixture(scope="module")
def device():
return get_test_device()
def _make_entity(
*,
effort_limit: float | None = 50.0,
armature: float | None = None,
delay_max_lag: int = 0,
delay_min_lag: int = 0,
delay_hold_prob: float = 0.0,
) -> Entity:
cfg = BuiltinPdActuatorCfg(
target_names_expr=("joint.*",),
stiffness=KP,
damping=KD,
effort_limit=effort_limit,
armature=armature,
delay_min_lag=delay_min_lag,
delay_max_lag=delay_max_lag,
delay_hold_prob=delay_hold_prob,
)
return create_entity_with_actuator(ROBOT_XML, cfg)
def _at_rest_with_targets(
entity: Entity,
sim,
device: str,
pos_target: torch.Tensor,
vel_target: torch.Tensor,
) -> None:
entity.write_joint_state_to_sim(
position=torch.zeros(1, 2, device=device),
velocity=torch.zeros(1, 2, device=device),
)
entity.set_joint_position_target(pos_target)
entity.set_joint_velocity_target(vel_target)
entity.set_joint_effort_target(torch.zeros(1, 2, device=device))
entity.write_data_to_sim()
sim.forward()
# ---------------------------------------------------------------------------
# Structural invariants
# ---------------------------------------------------------------------------
def test_two_ctrls_per_target_with_pos_then_vel_layout(device):
"""Each target gets one <position> + one <velocity>, in halves."""
entity, sim = initialize_entity(_make_entity(), device)
act = entity.actuators[0]
assert isinstance(act, BuiltinPdActuator)
n = act.num_targets
assert n == len(act.target_names) == 2
assert len(act.ctrl_ids) == 2 * n
assert len(act.global_ctrl_ids) == 2 * n
names = [sim.mj_model.actuator(i).name for i in act.global_ctrl_ids.tolist()]
assert names[:n] == [f"{name}_pd_pos" for name in act.target_names]
assert names[n:] == [f"{name}_pd_vel" for name in act.target_names]
def test_site_transmission_rejected():
with pytest.raises(ValueError, match="SITE"):
BuiltinPdActuatorCfg(
target_names_expr=("x",),
stiffness=1.0,
damping=1.0,
transmission_type=TransmissionType.SITE,
)
def test_armature_applied_once(device):
"""Joint armature must come from the position element only; double-applying
would silently double dof_armature."""
_, sim = initialize_entity(_make_entity(armature=0.7), device)
m = sim.mj_model
for jname in ("joint1", "joint2"):
dof_id = m.jnt_dofadr[m.joint(jname).id]
assert m.dof_armature[dof_id] == pytest.approx(0.7)
# ---------------------------------------------------------------------------
# Force computation
# ---------------------------------------------------------------------------
def test_position_only(device):
"""Zero vel target: qfrc = kp * pos_target."""
entity, sim = initialize_entity(_make_entity(effort_limit=None), device)
pos = torch.tensor([[0.1, -0.05]], device=device)
_at_rest_with_targets(entity, sim, device, pos, torch.zeros(1, 2, device=device))
v_adr = entity.indexing.joint_v_adr
assert torch.allclose(sim.data.qfrc_actuator[0, v_adr], KP * pos[0], atol=1e-4)
def test_velocity_only(device):
"""Zero pos target, joint at rest: qfrc = kd * vel_target."""
entity, sim = initialize_entity(_make_entity(effort_limit=None), device)
vel = torch.tensor([[0.3, -0.2]], device=device)
_at_rest_with_targets(entity, sim, device, torch.zeros(1, 2, device=device), vel)
v_adr = entity.indexing.joint_v_adr
assert torch.allclose(sim.data.qfrc_actuator[0, v_adr], KD * vel[0], atol=1e-4)
def test_pd_superposition(device):
"""Both targets nonzero: qfrc = kp * pos_target + kd * vel_target."""
entity, sim = initialize_entity(_make_entity(effort_limit=None), device)
pos = torch.tensor([[0.1, -0.05]], device=device)
vel = torch.tensor([[0.2, -0.1]], device=device)
_at_rest_with_targets(entity, sim, device, pos, vel)
v_adr = entity.indexing.joint_v_adr
expected = KP * pos[0] + KD * vel[0]
assert torch.allclose(sim.data.qfrc_actuator[0, v_adr], expected, atol=1e-4)
def test_actfrcrange_sum_clamp(device):
"""A pos error big enough to make kp*err exceed effort_limit must be
clamped at the joint, not allowed to ride through the unbounded element."""
entity, sim = initialize_entity(_make_entity(effort_limit=5.0), device)
# kp * 10.0 = 1000, well over the 5.0 clamp.
pos = torch.tensor([[10.0, 0.0]], device=device)
_at_rest_with_targets(entity, sim, device, pos, torch.zeros(1, 2, device=device))
v_adr = entity.indexing.joint_v_adr
qfrc = sim.data.qfrc_actuator[0, v_adr]
assert qfrc[0].item() == pytest.approx(5.0, abs=1e-4)
assert qfrc[1].item() == pytest.approx(0.0, abs=1e-4)
def test_effort_limit_none_leaves_joint_unlimited(device):
"""effort_limit=None: jnt_actfrclimited stays 0 on the targeted joints."""
_, sim = initialize_entity(_make_entity(effort_limit=None), device)
m = sim.mj_model
for jname in ("joint1", "joint2"):
jid = m.joint(jname).id
assert m.jnt_actfrclimited[jid] == 0
def test_actuator_forcerange_not_set(device):
"""We deliberately leave per-element forcerange unset; the limit lives on
the joint. Inspection of actuator_force[i] thus shows the unclamped value."""
entity, sim = initialize_entity(_make_entity(effort_limit=5.0), device)
m = sim.mj_model
for ctrl_id in entity.actuators[0].global_ctrl_ids.tolist():
assert m.actuator_forcelimited[ctrl_id] == 0
# ---------------------------------------------------------------------------
# Delay synchronization
# ---------------------------------------------------------------------------
def test_delay_syncs_pos_and_vel(device):
"""The shared delay buffer must lag pos and vel together."""
entity, sim = initialize_entity(
_make_entity(effort_limit=None, delay_min_lag=2, delay_max_lag=2),
device,
)
pos_targets = [
torch.tensor([[0.1, 0.0]], device=device),
torch.tensor([[0.3, 0.0]], device=device),
torch.tensor([[0.5, 0.0]], device=device),
]
vel_targets = [
torch.tensor([[1.0, 0.0]], device=device),
torch.tensor([[2.0, 0.0]], device=device),
torch.tensor([[3.0, 0.0]], device=device),
]
entity.write_joint_state_to_sim(
position=torch.zeros(1, 2, device=device),
velocity=torch.zeros(1, 2, device=device),
)
for p, v in zip(pos_targets, vel_targets, strict=True):
entity.set_joint_position_target(p)
entity.set_joint_velocity_target(v)
entity.set_joint_effort_target(torch.zeros(1, 2, device=device))
entity.write_data_to_sim()
sim.forward()
v_adr = entity.indexing.joint_v_adr
# With lag=2, both halves should reference step-0 values.
expected = KP * pos_targets[0][0] + KD * vel_targets[0][0]
assert torch.allclose(sim.data.qfrc_actuator[0, v_adr], expected, atol=1e-4)
def test_reset_clears_delay_buffer(device):
entity, _ = initialize_entity(_make_entity(delay_min_lag=1, delay_max_lag=3), device)
act = entity.actuators[0]
assert act._delay_buffer is not None
entity.set_joint_position_target(torch.full((1, 2), 0.5, device=device))
entity.set_joint_velocity_target(torch.zeros(1, 2, device=device))
entity.write_data_to_sim()
entity.reset(torch.tensor([0], device=device))
assert act._delay_buffer.current_lags[0] == 0
# ---------------------------------------------------------------------------
# Domain randomization
# ---------------------------------------------------------------------------
def _scene_env(device, transmission=TransmissionType.JOINT, num_envs=2):
"""Build a real scene/sim with one BuiltinPd-driven entity for DR tests."""
if transmission == TransmissionType.JOINT:
xml = ROBOT_XML
targets = ("joint.*",)
else:
xml = load_fixture_xml("tendon_finger")
# tendon_finger ships with motor/position/velocity actuators; we need a
# bare spec so BuiltinPd can attach to the tendon without name clashes.
targets = ("finger_tendon",)
def spec_fn():
spec = mujoco.MjSpec.from_string(xml)
# Strip any pre-existing actuators so BuiltinPd's added elements own ctrl.
for a in list(spec.actuators):
spec.delete(a)
return spec
entity_cfg = EntityCfg(
spec_fn=spec_fn,
articulation=EntityArticulationInfoCfg(
actuators=(
BuiltinPdActuatorCfg(
target_names_expr=targets,
stiffness=KP,
damping=KD,
effort_limit=50.0,
transmission_type=transmission,
),
)
),
)
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)
env = Mock()
env.num_envs = num_envs
env.device = device
env.scene = {"robot": scene["robot"]}
env.sim = sim
return env
def test_dr_pd_gains_scales_halves_independently(device):
env = _scene_env(device)
robot = env.scene["robot"]
act = robot.actuators[0]
assert isinstance(act, BuiltinPdActuator)
n = act.num_targets
pos_ids = act.global_ctrl_ids[:n]
vel_ids = act.global_ctrl_ids[n:]
# Expand fields so DR can write per-env.
env.sim.expand_model_fields(("actuator_gainprm", "actuator_biasprm"))
dr.pd_gains(
env,
env_ids=torch.tensor([0], device=device),
kp_range=(2.0, 2.0),
kd_range=(3.0, 3.0),
asset_cfg=SceneEntityCfg("robot"),
operation="scale",
)
m = env.sim.model
# Position half: gainprm[0] and biasprm[1] both scaled by kp=2, biasprm[2]
# must stay zero (no kd injection).
assert torch.allclose(
m.actuator_gainprm[0, pos_ids, 0],
torch.full((n,), 2.0 * KP, device=device),
)
assert torch.allclose(
m.actuator_biasprm[0, pos_ids, 1],
torch.full((n,), -2.0 * KP, device=device),
)
assert torch.allclose(
m.actuator_biasprm[0, pos_ids, 2], torch.zeros(n, device=device)
)
# Velocity half: gainprm[0] and biasprm[2] both scaled by kd=3, biasprm[1]
# stays zero (no kp injection).
assert torch.allclose(
m.actuator_gainprm[0, vel_ids, 0],
torch.full((n,), 3.0 * KD, device=device),
)
assert torch.allclose(
m.actuator_biasprm[0, vel_ids, 2],
torch.full((n,), -3.0 * KD, device=device),
)
assert torch.allclose(
m.actuator_biasprm[0, vel_ids, 1], torch.zeros(n, device=device)
)
# The other env must be untouched.
assert torch.allclose(m.actuator_gainprm[1, pos_ids, 0], torch.tensor(KP))
assert torch.allclose(m.actuator_gainprm[1, vel_ids, 0], torch.tensor(KD))
def test_dr_pd_gains_abs_writes_correct_columns(device):
env = _scene_env(device)
robot = env.scene["robot"]
act = robot.actuators[0]
assert isinstance(act, BuiltinPdActuator)
n = act.num_targets
pos_ids = act.global_ctrl_ids[:n]
vel_ids = act.global_ctrl_ids[n:]
env.sim.expand_model_fields(("actuator_gainprm", "actuator_biasprm"))
dr.pd_gains(
env,
env_ids=torch.tensor([0, 1], device=device),
kp_range=(200.0, 200.0),
kd_range=(25.0, 25.0),
asset_cfg=SceneEntityCfg("robot"),
operation="abs",
)
m = env.sim.model
assert torch.allclose(
m.actuator_gainprm[:, pos_ids, 0],
torch.full((env.num_envs, n), 200.0, device=device),
)
assert torch.allclose(
m.actuator_biasprm[:, pos_ids, 1],
torch.full((env.num_envs, n), -200.0, device=device),
)
assert torch.allclose(
m.actuator_biasprm[:, pos_ids, 2],
torch.zeros(env.num_envs, n, device=device),
)
assert torch.allclose(
m.actuator_gainprm[:, vel_ids, 0],
torch.full((env.num_envs, n), 25.0, device=device),
)
assert torch.allclose(
m.actuator_biasprm[:, vel_ids, 2],
torch.full((env.num_envs, n), -25.0, device=device),
)
def test_dr_effort_limits_writes_jnt_actfrcrange(device):
env = _scene_env(device, transmission=TransmissionType.JOINT)
robot = env.scene["robot"]
act = robot.actuators[0]
assert isinstance(act, BuiltinPdActuator)
env.sim.expand_model_fields(
("actuator_forcerange", "jnt_actfrcrange", "tendon_actfrcrange")
)
joint_ids = robot.indexing.joint_ids[act.target_ids]
pre_forcerange = env.sim.model.actuator_forcerange.clone()
dr.effort_limits(
env,
env_ids=torch.tensor([0], device=device),
effort_limit_range=(123.0, 123.0),
asset_cfg=SceneEntityCfg("robot"),
operation="abs",
)
m = env.sim.model
# The joint sum-clamp was rewritten on env 0 only.
assert torch.allclose(
m.jnt_actfrcrange[0, joint_ids],
torch.tensor([[-123.0, 123.0]] * len(joint_ids), device=device),
)
assert torch.allclose(
m.jnt_actfrcrange[1, joint_ids],
torch.tensor([[-50.0, 50.0]] * len(joint_ids), device=device),
)
# Per-element actuator_forcerange must be untouched for BuiltinPd: that
# field belongs to the existing single-element actuator semantic.
assert torch.allclose(m.actuator_forcerange, pre_forcerange)
def test_dr_effort_limits_scale_multiplies_default(device):
"""``scale`` multiplies the configured ``effort_limit`` (50.0) by the sample."""
env = _scene_env(device, transmission=TransmissionType.JOINT)
robot = env.scene["robot"]
act = robot.actuators[0]
assert isinstance(act, BuiltinPdActuator)
env.sim.expand_model_fields(
("actuator_forcerange", "jnt_actfrcrange", "tendon_actfrcrange")
)
joint_ids = robot.indexing.joint_ids[act.target_ids]
dr.effort_limits(
env,
env_ids=torch.tensor([0], device=device),
effort_limit_range=(2.0, 2.0),
asset_cfg=SceneEntityCfg("robot"),
operation="scale",
)
m = env.sim.model
# Default is [-50, 50], scaled by 2 -> [-100, 100].
assert torch.allclose(
m.jnt_actfrcrange[0, joint_ids],
torch.tensor([[-100.0, 100.0]] * len(joint_ids), device=device),
)
def test_dr_effort_limits_writes_tendon_actfrcrange(device):
env = _scene_env(device, transmission=TransmissionType.TENDON)
robot = env.scene["robot"]
act = robot.actuators[0]
assert isinstance(act, BuiltinPdActuator)
env.sim.expand_model_fields(
("actuator_forcerange", "jnt_actfrcrange", "tendon_actfrcrange")
)
tendon_ids = robot.indexing.tendon_ids[act.target_ids]
dr.effort_limits(
env,
env_ids=torch.tensor([0], device=device),
effort_limit_range=(77.0, 77.0),
asset_cfg=SceneEntityCfg("robot"),
operation="abs",
)
m = env.sim.model
assert torch.allclose(
m.tendon_actfrcrange[0, tendon_ids],
torch.tensor([[-77.0, 77.0]] * len(tendon_ids), device=device),
)
assert torch.allclose(
m.tendon_actfrcrange[1, tendon_ids],
torch.tensor([[-50.0, 50.0]] * len(tendon_ids), device=device),
)
05_software/train/rc_mjlab/mjlab/tests/test_contact_sensor.py
+34
View File
@@ -1072,3 +1072,37 @@ def test_history_captures_impact_forces(device):
assert torch.all(max_force_seen > steady_state_force * 1.5), (
f"Peak force {max_force_seen} should be significantly above mg={steady_state_force}"
)
def test_global_frame_maxforce_rotation(device):
"""A box at rest on a plane has its contact normals all vertical."""
cfg = ContactSensorCfg(
name="box_contact",
primary=ContactMatch(mode="geom", pattern="box_geom", entity="box"),
fields=("found", "force", "normal", "tangent"),
reduce="maxforce",
global_frame=True,
)
scene, sim = create_scene_with_sensor(FALLING_BOX_XML, "box", cfg, device)
root_state = torch.zeros((2, 13), device=sim.device)
root_state[:, 2] = 0.11
root_state[:, 3] = 1.0
scene["box"].write_root_state_to_sim(root_state)
for _ in range(150):
sim.step()
scene.update(dt=sim.cfg.mujoco.timestep)
sensor_force = scene["box_contact"].data.force[:, 0, :]
# On a flat plane the contact normal is vertical, so a correctly rotated
# global-frame force should have its magnitude entirely on the z axis.
assert torch.all(sensor_force[:, 0].abs() < 0.05), (
f"sensor_force x-component should be ~0, got {sensor_force[:, 0].tolist()}"
)
assert torch.all(sensor_force[:, 1].abs() < 0.05), (
f"sensor_force y-component should be ~0, got {sensor_force[:, 1].tolist()}"
)
assert torch.all(sensor_force[:, 2].abs() > 1.0), (
f"sensor_force z-component should be non-trivial, got {sensor_force[:, 2].tolist()}"
)
05_software/train/rc_mjlab/mjlab/tests/test_delayed_actuator.py
+66
View File
@@ -125,6 +125,72 @@ def test_delayed_ideal_applies_delay(device):
assert torch.allclose(qfrc, expected_torque, atol=1e-4)
def test_delayed_ideal_delays_velocity(device):
"""Velocity targets share the same delay as position targets.
Regression test: the velocity reference used to bypass the delay buffer, so
the damping term consumed the latest target instead of the delayed one.
"""
entity = create_entity_with_delayed_ideal(delay_min_lag=2, delay_max_lag=2)
entity, sim = initialize_entity(entity, device)
joint_pos = torch.zeros(1, 2, device=device)
joint_vel = torch.zeros(1, 2, device=device)
entity.write_joint_state_to_sim(joint_pos, joint_vel)
# Only the velocity target varies; position and effort stay zero.
vel_targets = [
torch.tensor([[0.1, 0.2]], device=device),
torch.tensor([[0.3, 0.4]], device=device),
torch.tensor([[0.5, 0.6]], device=device),
]
for vel_target in vel_targets:
entity.set_joint_position_target(joint_pos)
entity.set_joint_velocity_target(vel_target)
entity.set_joint_effort_target(torch.zeros(1, 2, device=device))
entity.write_data_to_sim()
sim.forward()
joint_v_adr = entity.indexing.joint_v_adr
qfrc = sim.data.qfrc_actuator[0, joint_v_adr]
# With lag=2, the damping term uses the velocity target from step 0:
# kd * (delayed_vel_target - 0) = 10.0 * [0.1, 0.2].
expected_torque = 10.0 * vel_targets[0][0]
assert torch.allclose(qfrc, expected_torque, atol=1e-4)
def test_delayed_ideal_delays_effort(device):
"""Feedforward effort targets share the same delay as position targets."""
entity = create_entity_with_delayed_ideal(delay_min_lag=2, delay_max_lag=2)
entity, sim = initialize_entity(entity, device)
joint_pos = torch.zeros(1, 2, device=device)
joint_vel = torch.zeros(1, 2, device=device)
entity.write_joint_state_to_sim(joint_pos, joint_vel)
effort_targets = [
torch.tensor([[1.0, 2.0]], device=device),
torch.tensor([[3.0, 4.0]], device=device),
torch.tensor([[5.0, 6.0]], device=device),
]
for effort_target in effort_targets:
entity.set_joint_position_target(joint_pos)
entity.set_joint_velocity_target(joint_vel)
entity.set_joint_effort_target(effort_target)
entity.write_data_to_sim()
sim.forward()
joint_v_adr = entity.indexing.joint_v_adr
qfrc = sim.data.qfrc_actuator[0, joint_v_adr]
# With lag=2, the feedforward term uses the effort target from step 0.
expected_torque = effort_targets[0][0]
assert torch.allclose(qfrc, expected_torque, atol=1e-4)
def test_delayed_actuator_reset(device):
"""Test that reset clears the delay buffer."""
entity = create_entity_with_delayed_builtin(delay_min_lag=1, delay_max_lag=3)
05_software/train/rc_mjlab/mjlab/tests/test_entity.py
+21
View File
@@ -243,6 +243,27 @@ def test_unnamed_freejoint_gets_default_name():
assert "floating_base_joint" in entity.all_joint_names
def test_multiple_freejoints_raises():
"""An entity with more than one freejoint is rejected at construction."""
xml = """
<mujoco>
<worldbody>
<body name="object_a" pos="0 0 1">
<freejoint/>
<geom type="box" size="0.1 0.1 0.1" mass="0.1"/>
</body>
<body name="object_b" pos="1 0 1">
<freejoint/>
<geom type="box" size="0.1 0.1 0.1" mass="0.1"/>
</body>
</worldbody>
</mujoco>
"""
cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(xml))
with pytest.raises(ValueError, match="2 freejoints"):
Entity(cfg)
def test_find_methods():
"""Test find methods with exact and regex matches."""
entity = create_floating_articulated_entity()
05_software/train/rc_mjlab/mjlab/tests/test_events.py
+105 -6
View File
@@ -125,7 +125,9 @@ def test_dr_fields_registered_in_event_manager(device):
assert "actuator_gainprm" in manager.domain_randomization_fields
assert "actuator_biasprm" in manager.domain_randomization_fields
assert "actuator_forcerange" in manager.domain_randomization_fields
assert len(manager.domain_randomization_fields) == 5
assert "jnt_actfrcrange" in manager.domain_randomization_fields
assert "tendon_actfrcrange" in manager.domain_randomization_fields
assert len(manager.domain_randomization_fields) == 7
def test_recompute_level_ordering():
@@ -418,6 +420,62 @@ def test_effort_limits_scale_no_accumulation(device):
assert abs(actual_upper - 200.0) < 1e-5
def test_pd_gains_accepts_operation_object(device):
"""dr.scale / dr.abs Operation objects produce the same result as strings."""
env_str, ideal_str = _make_pd_env(device)
env_obj, ideal_obj = _make_pd_env(device)
ids = torch.tensor([0], device=device)
kwargs = dict(
kp_range=(1.5, 1.5), kd_range=(2.0, 2.0), asset_cfg=SceneEntityCfg("robot")
)
torch.manual_seed(0)
dr.pd_gains(env_str, ids, operation="scale", **kwargs)
torch.manual_seed(0)
dr.pd_gains(env_obj, ids, operation=dr.scale, **kwargs)
assert torch.allclose(
env_str.sim.model.actuator_gainprm[0], env_obj.sim.model.actuator_gainprm[0]
)
assert torch.allclose(ideal_str.stiffness, ideal_obj.stiffness)
def test_effort_limits_accepts_operation_object(device):
"""dr.abs Operation object produces the same result as the string."""
env_str, ideal_str = _make_effort_env(device)
env_obj, ideal_obj = _make_effort_env(device)
ids = torch.tensor([0], device=device)
kwargs = dict(effort_limit_range=(150.0, 150.0), asset_cfg=SceneEntityCfg("robot"))
dr.effort_limits(env_str, ids, operation="abs", **kwargs)
dr.effort_limits(env_obj, ids, operation=dr.abs, **kwargs)
assert torch.allclose(
env_str.sim.model.actuator_forcerange[0], env_obj.sim.model.actuator_forcerange[0]
)
assert torch.allclose(ideal_str.force_limit, ideal_obj.force_limit)
def test_pd_gains_rejects_unsupported_operation(device):
"""Operations other than scale/abs raise ValueError."""
env, _ = _make_pd_env(device)
ids = torch.tensor([0], device=device)
with pytest.raises(ValueError, match="only supports 'scale' and 'abs'"):
dr.pd_gains(env, ids, kp_range=(1.0, 1.0), kd_range=(1.0, 1.0), operation=dr.add)
def test_effort_limits_rejects_unsupported_operation(device):
"""Operations other than scale/abs raise ValueError."""
env, _ = _make_effort_env(device)
ids = torch.tensor([0], device=device)
with pytest.raises(ValueError, match="only supports 'scale' and 'abs'"):
dr.effort_limits(env, ids, effort_limit_range=(1.0, 1.0), operation=dr.add)
# ===========================================================================
# Section 3: Other events
# ===========================================================================
@@ -503,7 +561,9 @@ def test_step_mode_fires_every_call(device):
assert call_count[0] == 5
def _make_impulse_env(device, num_envs=2, num_bodies=1, body_ids=None):
def _make_impulse_env(
device, num_envs=2, num_bodies=1, body_ids=None, cooldown_s=(0.0, 0.0)
):
"""Create a mock env for apply_body_impulse tests."""
if body_ids is None:
body_ids = [0]
@@ -523,7 +583,7 @@ def _make_impulse_env(device, num_envs=2, num_bodies=1, body_ids=None):
asset_cfg = SceneEntityCfg("robot", body_ids=body_ids)
term_cfg = Mock()
term_cfg.params = {"asset_cfg": asset_cfg}
term_cfg.params = {"asset_cfg": asset_cfg, "cooldown_s": cooldown_s}
impulse = events.apply_body_impulse(cfg=term_cfg, env=env)
return env, mock_entity, asset_cfg, impulse
@@ -531,11 +591,13 @@ def _make_impulse_env(device, num_envs=2, num_bodies=1, body_ids=None):
def test_apply_body_impulse_basic(device):
"""Impulse is applied and cleared after duration expires."""
env, mock_entity, asset_cfg, impulse = _make_impulse_env(
device, num_envs=2, num_bodies=3, body_ids=[1]
device, num_envs=2, num_bodies=3, body_ids=[1], cooldown_s=(10.0, 10.0)
)
# First call: cooldown_s starts at 0 and gets decremented by dt,
# so it becomes <= 0 and triggers.
# Skip the initial cooldown so the first call triggers immediately;
# the trigger/sustain/expire cycle is what's under test here.
impulse._interval_time_left[:] = 0.0
impulse(
env,
None,
@@ -643,6 +705,43 @@ def test_apply_body_impulse_reset_clears(device):
assert env_ids_arg[0].item() == 0
def test_apply_body_impulse_initial_cooldown(device):
"""The first call after init/reset enters cooldown, not an immediate impulse.
Regression test for #973.
"""
env, mock_entity, asset_cfg, impulse = _make_impulse_env(
device, num_envs=1, num_bodies=1, body_ids=[0], cooldown_s=(0.05, 0.05)
)
def step():
impulse(
env,
None,
force_range=(10.0, 10.0),
torque_range=(0.0, 0.0),
duration_s=(1.0, 1.0),
cooldown_s=(0.05, 0.05), # ~2.5 steps at dt=0.02
asset_cfg=asset_cfg,
)
# First two steps consume the sampled cooldown; impulse must not fire yet.
step()
assert not impulse._active.any()
step()
assert not impulse._active.any()
# Third step crosses the cooldown boundary and triggers.
step()
assert impulse._active.all()
# Reset re-enters cooldown: next step should not immediately re-trigger.
impulse.reset(env_ids=torch.tensor([0], device=device))
assert not impulse._active.any()
step()
assert not impulse._active.any()
# ===========================================================================
# Section 5: Recomputation integration
# ===========================================================================
05_software/train/rc_mjlab/mjlab/tests/test_gpu_selection.py
+13
View File
@@ -113,3 +113,16 @@ def test_select_gpus_cpu_mode_empty_cuda_visible_devices():
selected, num = select_gpus([0])
assert selected is None
assert num == 0
def test_select_gpus_mig_uuids():
"""Handles MIG GPU UUIDs in CUDA_VISIBLE_DEVICES."""
os.environ["CUDA_VISIBLE_DEVICES"] = "MIG-GPU-abc-123,MIG-GPU-def-456"
selected, num = select_gpus("all")
assert selected == ["MIG-GPU-abc-123", "MIG-GPU-def-456"]
assert num == 2
selected, num = select_gpus([0])
assert selected == ["MIG-GPU-abc-123"]
assert num == 1
05_software/train/rc_mjlab/mjlab/tests/test_mesh_variants.py
-961
View File
@@ -1,961 +0,0 @@
"""Tests for per-world mesh variant support."""
from __future__ import annotations
from typing import Any, cast
import mujoco
import numpy as np
import pytest
import torch
from mjlab.entity import EntityCfg, VariantCfg, VariantEntityCfg
from mjlab.sim.mesh_variants import allocate_worlds, build_mesh_variant_model
from mjlab.viewer.model_sync import (
disable_model_sameframe_shortcuts,
sync_model_fields,
)
# Helpers: variant specs with visual + collision mesh geoms.
def _sphere_2col_spec() -> mujoco.MjSpec:
"""Sphere: 1 visual + 2 collision geoms."""
spec = mujoco.MjSpec()
mv = spec.add_mesh()
mv.name = "visual"
mv.make_sphere(subdivision=3)
for i in range(2):
mc = spec.add_mesh()
mc.name = f"col_{i}"
mc.make_sphere(subdivision=1)
body = spec.worldbody.add_body()
body.name = "prop"
body.add_freejoint()
gv = body.add_geom()
gv.name = "visual"
gv.type = mujoco.mjtGeom.mjGEOM_MESH
gv.meshname = "visual"
gv.contype = 0
gv.conaffinity = 0
for i in range(2):
gc = body.add_geom()
gc.name = f"col_{i}"
gc.type = mujoco.mjtGeom.mjGEOM_MESH
gc.meshname = f"col_{i}"
return spec
def _cone_4col_spec() -> mujoco.MjSpec:
"""Cone: 1 visual + 4 collision geoms (more than sphere)."""
spec = mujoco.MjSpec()
mv = spec.add_mesh()
mv.name = "visual"
mv.make_cone(nedge=8, radius=0.05)
for i in range(4):
mc = spec.add_mesh()
mc.name = f"col_{i}"
mc.make_sphere(subdivision=1)
body = spec.worldbody.add_body()
body.name = "prop"
body.add_freejoint()
gv = body.add_geom()
gv.name = "visual"
gv.type = mujoco.mjtGeom.mjGEOM_MESH
gv.meshname = "visual"
gv.contype = 0
gv.conaffinity = 0
for i in range(4):
gc = body.add_geom()
gc.name = f"col_{i}"
gc.type = mujoco.mjtGeom.mjGEOM_MESH
gc.meshname = f"col_{i}"
return spec
def _simple_sphere_spec() -> mujoco.MjSpec:
"""Single-geom sphere for simple tests."""
spec = mujoco.MjSpec()
m = spec.add_mesh()
m.name = "sphere"
m.make_sphere(subdivision=2)
body = spec.worldbody.add_body()
body.name = "prop"
body.add_freejoint()
g = body.add_geom()
g.name = "visual"
g.type = mujoco.mjtGeom.mjGEOM_MESH
g.meshname = "sphere"
return spec
def _simple_cone_spec() -> mujoco.MjSpec:
"""Single-geom cone for simple tests."""
spec = mujoco.MjSpec()
m = spec.add_mesh()
m.name = "cone"
m.make_cone(nedge=8, radius=0.05)
body = spec.worldbody.add_body()
body.name = "prop"
body.add_freejoint()
g = body.add_geom()
g.name = "visual"
g.type = mujoco.mjtGeom.mjGEOM_MESH
g.meshname = "cone"
return spec
def _hinge_spec() -> mujoco.MjSpec:
"""Object with a hinge joint (incompatible with freejoint variants)."""
spec = mujoco.MjSpec()
m = spec.add_mesh()
m.name = "box"
m.make_sphere(subdivision=1)
body = spec.worldbody.add_body()
body.name = "prop"
j = body.add_joint()
j.name = "hinge"
j.type = mujoco.mjtJoint.mjJNT_HINGE
g = body.add_geom()
g.name = "visual"
g.type = mujoco.mjtGeom.mjGEOM_MESH
g.meshname = "box"
return spec
def _build_scene_with_variants(
variant_a_fn, variant_b_fn, *, weight_a=0.5, weight_b=0.5
):
"""Build a scene spec + variant_info from two variant spec_fns."""
cfg = VariantEntityCfg(
variants={
"a": VariantCfg(spec_fn=variant_a_fn, weight=weight_a),
"b": VariantCfg(spec_fn=variant_b_fn, weight=weight_b),
},
)
entity = cfg.build()
assert entity.variant_metadata is not None
scene_spec = mujoco.MjSpec()
frame = scene_spec.worldbody.add_frame()
scene_spec.attach(entity.spec, prefix="object/", frame=frame)
return scene_spec, [("object/", entity.variant_metadata)]
# allocate_worlds.
def test_allocate_worlds_proportional():
result = allocate_worlds((0.6, 0.4), 10)
assert len(result) == 10
assert result.count(0) == 6
assert result.count(1) == 4
def test_allocate_worlds_uniform():
result = allocate_worlds((1.0, 1.0), 8)
assert result.count(0) == 4
assert result.count(1) == 4
def test_allocate_worlds_single_variant():
result = allocate_worlds((1.0,), 5)
assert result == [0, 0, 0, 0, 0]
def test_allocate_worlds_zero_weight_skips_variant():
"""A zero-weight variant gets zero worlds; the rest split nworld."""
result = allocate_worlds((1.0, 0.0, 1.0), 10)
assert len(result) == 10
assert result.count(1) == 0
assert result.count(0) == 5
assert result.count(2) == 5
def test_allocate_worlds_rejects_negative_weight():
with pytest.raises(ValueError, match="non-negative"):
allocate_worlds((1.0, -0.1), 10)
def test_allocate_worlds_rejects_all_zero():
with pytest.raises(ValueError, match="positive sum"):
allocate_worlds((0.0, 0.0), 10)
def test_allocate_worlds_largest_remainder_sums_to_nworld():
"""Largest-remainder rounding must always allocate exactly nworld worlds."""
for nworld in (3, 7, 100, 1000):
result = allocate_worlds((1.0, 1.0, 1.0), nworld)
assert len(result) == nworld
# Difference between any two variant counts is at most 1 (uniform).
counts = [result.count(i) for i in range(3)]
assert max(counts) - min(counts) <= 1
# Entity merging.
def test_entity_builds_with_variants():
cfg = VariantEntityCfg(
variants={
"sphere": VariantCfg(spec_fn=_simple_sphere_spec, weight=0.5),
"cone": VariantCfg(spec_fn=_simple_cone_spec, weight=0.5),
},
)
entity = cfg.build()
meta = entity.variant_metadata
assert meta is not None
assert meta.variant_names == ("sphere", "cone")
assert meta.num_mesh_geoms == 1
mesh_names = [m.name for m in entity.spec.meshes]
assert any("sphere" in n for n in mesh_names)
assert any("cone" in n for n in mesh_names)
def test_multi_geom_body_padding():
"""Sphere (3 geoms) + cone (5 geoms) -> body padded to 5 mesh geoms."""
cfg = VariantEntityCfg(
variants={
"sphere": VariantCfg(spec_fn=_sphere_2col_spec, weight=0.5),
"cone": VariantCfg(spec_fn=_cone_4col_spec, weight=0.5),
},
)
entity = cfg.build()
meta = entity.variant_metadata
assert meta is not None
assert meta.num_mesh_geoms == 5 # max(3, 5)
# Sphere: 3 real + 2 padding (None).
assert sum(1 for n in meta.variant_mesh_names[0] if n is None) == 2
# Cone: 5 real, no padding.
assert all(n is not None for n in meta.variant_mesh_names[1])
# Validation.
def test_mismatched_joint_structure_raises():
cfg = VariantEntityCfg(
variants={
"sphere": VariantCfg(spec_fn=_simple_sphere_spec, weight=0.5),
"hinge": VariantCfg(spec_fn=_hinge_spec, weight=0.5),
},
)
with pytest.raises(ValueError, match="joint"):
cfg.build()
def test_single_variant_builds():
"""A single variant degenerates cleanly; useful for templated variant sets."""
cfg = VariantEntityCfg(
variants={"only": VariantCfg(spec_fn=_simple_sphere_spec)},
)
entity = cfg.build()
assert entity.variant_metadata is not None
assert entity.variant_metadata.variant_names == ("only",)
def test_empty_variants_raises():
cfg = VariantEntityCfg(variants={})
with pytest.raises(ValueError, match="at least one"):
cfg.build()
def _fixed_base_sphere_spec() -> mujoco.MjSpec:
"""Fixed-base sphere variant (no free joint): currently unsupported."""
spec = mujoco.MjSpec()
m = spec.add_mesh(name="sphere")
m.make_sphere(subdivision=2)
body = spec.worldbody.add_body(name="prop")
body.add_geom(type=mujoco.mjtGeom.mjGEOM_MESH, meshname="sphere")
return spec
def test_fixed_base_variants_rejected():
"""Variants must be floating-base; fixed-base raises with a clear message."""
cfg = VariantEntityCfg(
variants={
"a": VariantCfg(spec_fn=_fixed_base_sphere_spec, weight=0.5),
"b": VariantCfg(spec_fn=_fixed_base_sphere_spec, weight=0.5),
},
)
with pytest.raises(ValueError, match="floating-base"):
cfg.build()
def test_setting_spec_fn_on_variant_cfg_raises():
"""VariantEntityCfg.spec_fn is unused; setting it should fail loudly."""
with pytest.raises(ValueError, match="spec_fn cannot be set"):
VariantEntityCfg(
variants={"only": VariantCfg(spec_fn=_simple_sphere_spec)},
spec_fn=_simple_sphere_spec,
)
def test_no_variants_unchanged():
cfg = EntityCfg(spec_fn=_simple_sphere_spec)
entity = cfg.build()
assert entity.variant_metadata is None
# build_mesh_variant_model: dataid and dependent fields.
def test_dataid_assigned_per_world():
"""Each world's geom_dataid points to its variant's meshes."""
scene_spec, vi = _build_scene_with_variants(_simple_sphere_spec, _simple_cone_spec)
result = build_mesh_variant_model(scene_spec, 4, vi)
dataid = result.wp_model.geom_dataid.numpy()
assert dataid.shape == (4, result.mj_model.ngeom)
assert dataid.ndim == 2
w2v = result.world_to_variant["object/"]
assert w2v[0] == 0 # variant a (sphere)
assert w2v[2] == 1 # variant b (cone)
# Sphere and cone worlds must have different dataid values.
assert not np.array_equal(dataid[0], dataid[2])
def test_padding_slots_get_disabled():
"""Shorter variant's padding geom slots have dataid == -1."""
scene_spec, vi = _build_scene_with_variants(_sphere_2col_spec, _cone_4col_spec)
result = build_mesh_variant_model(scene_spec, 4, vi)
dataid = result.wp_model.geom_dataid.numpy()
w2v = result.world_to_variant["object/"]
# Find a sphere world (variant 0, 3 mesh geoms -> 2 padding slots).
sphere_world = int(np.where(w2v == 0)[0][0])
# Find mesh geom columns (skip non-mesh geoms like worldbody).
mesh_geom_ids = [
gid
for gid in range(result.mj_model.ngeom)
if result.mj_model.geom_type[gid] == mujoco.mjtGeom.mjGEOM_MESH
]
sphere_dataid = dataid[sphere_world, mesh_geom_ids]
# Last 2 mesh geom slots should be -1 (disabled padding).
assert sphere_dataid[-1] == -1
assert sphere_dataid[-2] == -1
# Padding slots must still be collision-enabled in the template/warp model.
# Short variants are disabled by per-world dataid=-1; long variants need the
# same slots enabled so their extra hulls can collide.
assert np.all(result.mj_model.geom_contype[mesh_geom_ids[-2:]] == 1)
assert np.all(result.mj_model.geom_conaffinity[mesh_geom_ids[-2:]] == 1)
assert np.all(result.wp_model.geom_contype.numpy()[mesh_geom_ids[-2:]] == 1)
assert np.all(result.wp_model.geom_conaffinity.numpy()[mesh_geom_ids[-2:]] == 1)
# First 3 should be valid (>= 0).
assert all(d >= 0 for d in sphere_dataid[:3])
def test_dependent_fields_match_individual_compilation():
"""Per-world body_mass matches independently compiled variant models."""
scene_spec, vi = _build_scene_with_variants(_simple_sphere_spec, _simple_cone_spec)
result = build_mesh_variant_model(scene_spec, 4, vi)
# Compile each variant independently for reference values.
sphere_model = _simple_sphere_spec().compile()
cone_model = _simple_cone_spec().compile()
body_mass = result.wp_model.body_mass.numpy()
w2v = result.world_to_variant["object/"]
sphere_w = int(np.where(w2v == 0)[0][0])
cone_w = int(np.where(w2v == 1)[0][0])
# The object body is the last body in the scene.
obj_body = result.mj_model.nbody - 1
# Mass should match individually compiled models.
np.testing.assert_allclose(
body_mass[sphere_w, obj_body],
sphere_model.body_mass[-1],
atol=1e-4,
)
np.testing.assert_allclose(
body_mass[cone_w, obj_body],
cone_model.body_mass[-1],
atol=1e-4,
)
# Sphere and cone should have different masses.
assert not np.isclose(body_mass[sphere_w, obj_body], body_mass[cone_w, obj_body])
def test_select_default_values_uses_per_world_variant_defaults():
"""Per-world defaults are indexed by env first, then by entity."""
from mjlab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg
from mjlab.envs.mdp.dr._core import _select_default_values
from mjlab.scene import SceneCfg
from mjlab.terrains import TerrainEntityCfg
def _explicit_variant(
mesh_name: str,
mass: float,
inertia: tuple[float, float, float],
*,
cone: bool = False,
) -> mujoco.MjSpec:
spec = mujoco.MjSpec()
mesh = spec.add_mesh()
mesh.name = mesh_name
if cone:
mesh.make_cone(nedge=8, radius=0.05)
else:
mesh.make_sphere(subdivision=1)
body = spec.worldbody.add_body(name="prop")
body.add_freejoint()
body.explicitinertial = 1
body.mass = mass
body.ipos[:] = (0.0, 0.0, 0.0)
body.inertia[:] = inertia
body.iquat[:] = (1.0, 0.0, 0.0, 0.0)
body.add_geom(
name="visual",
type=mujoco.mjtGeom.mjGEOM_MESH,
meshname=mesh_name,
contype=0,
conaffinity=0,
mass=0.0,
)
return spec
object_cfg = VariantEntityCfg(
variants={
"sphere": VariantCfg(
lambda: _explicit_variant("sphere", 0.2, (1e-4, 2e-4, 3e-4)),
weight=0.5,
),
"cone": VariantCfg(
lambda: _explicit_variant("cone", 0.7, (4e-4, 5e-4, 6e-4), cone=True),
weight=0.5,
),
},
init_state=EntityCfg.InitialStateCfg(pos=(0.0, 0.0, 0.2)),
)
env_cfg = ManagerBasedRlEnvCfg(
decimation=1,
scene=SceneCfg(
terrain=TerrainEntityCfg(terrain_type="plane"),
num_envs=4,
env_spacing=1.0,
entities={"object": object_cfg},
),
)
env = ManagerBasedRlEnv(cfg=env_cfg, device="cpu")
try:
obj_body = int(env.scene["object"].indexing.root_body_id)
env_ids = torch.arange(env.num_envs, device=env.device)
body_ids = torch.tensor([obj_body], device=env.device)
for field in ("body_mass", "body_ipos", "body_inertia", "body_iquat"):
selected = _select_default_values(env, field, env_ids, body_ids)
torch.testing.assert_close(
selected[:, 0],
getattr(env.sim.model, field)[:, obj_body],
)
finally:
env.close()
def test_viser_builds_per_world_mesh_handles_for_variants():
"""Viser dynamic meshes must not collapse all worlds onto env0's mesh."""
from contextlib import nullcontext
from mjlab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg
from mjlab.scene import SceneCfg
from mjlab.terrains import TerrainEntityCfg
from mjlab.viewer.viser.scene import MjlabViserScene, _PerWorldMeshGroup
class _Handle:
def __init__(self, **kwargs):
self.visible = kwargs.get("visible", True)
self.batched_positions = kwargs.get("batched_positions", np.zeros((0, 3)))
self.batched_wxyzs = kwargs.get("batched_wxyzs", np.zeros((0, 4)))
self.batched_scales = kwargs.get("batched_scales")
self.batched_colors = kwargs.get("batched_colors")
self.batched_opacities = kwargs.get("batched_opacities")
self.position = kwargs.get("position", np.zeros(3))
self.wxyz = kwargs.get("wxyz", np.array([1.0, 0.0, 0.0, 0.0]))
def remove(self) -> None:
pass
class _Scene:
def __init__(self):
self.batched: list[tuple[tuple, dict, _Handle]] = []
def configure_environment_map(self, **_kwargs) -> None:
pass
def add_frame(self, *_args, **kwargs) -> _Handle:
return _Handle(**kwargs)
def add_grid(self, *_args, **kwargs) -> _Handle:
return _Handle(**kwargs)
def add_mesh_trimesh(self, *_args, **kwargs) -> _Handle:
return _Handle(**kwargs)
def add_batched_meshes_trimesh(self, *args, **kwargs) -> _Handle:
handle = _Handle(**kwargs)
self.batched.append((args, kwargs, handle))
return handle
def add_batched_meshes_simple(self, *args, **kwargs) -> _Handle:
handle = _Handle(**kwargs)
self.batched.append((args, kwargs, handle))
return handle
class _Server:
def __init__(self):
self.scene = _Scene()
def atomic(self):
return nullcontext()
def flush(self) -> None:
pass
env_cfg = ManagerBasedRlEnvCfg(
decimation=1,
scene=SceneCfg(
terrain=TerrainEntityCfg(terrain_type="plane"),
num_envs=4,
env_spacing=1.0,
entities={
"object": VariantEntityCfg(
variants={
"sphere": VariantCfg(_simple_sphere_spec, weight=0.5),
"cone": VariantCfg(_simple_cone_spec, weight=0.5),
},
init_state=EntityCfg.InitialStateCfg(pos=(0.0, 0.0, 0.2)),
)
},
),
)
env = ManagerBasedRlEnv(cfg=env_cfg, device="cpu")
try:
env.sim.expand_model_fields(("geom_rgba",))
env.sim.model.geom_rgba[:, :, :3] = torch.linspace(
0.2,
0.9,
env.num_envs,
device=env.device,
)[:, None, None]
server = _Server()
scene = MjlabViserScene(
cast(Any, server),
env.sim.mj_model,
env.num_envs,
sim_model=env.sim.model,
expanded_fields=env.sim.expanded_fields,
)
groups = [mg for mg in scene._mesh_groups if isinstance(mg, _PerWorldMeshGroup)]
assert groups
assert sum(len(mg.env_ids) for mg in groups) >= env.num_envs
body_xpos = env.sim.data.xpos.cpu().numpy()
body_xmat = env.sim.data.xmat.cpu().numpy()
mocap_pos = (
env.sim.data.mocap_pos.cpu().numpy() if env.sim.mj_model.nmocap > 0 else None
)
mocap_quat = (
env.sim.data.mocap_quat.cpu().numpy() if env.sim.mj_model.nmocap > 0 else None
)
scene.show_only_selected = True
scene.update_from_arrays(body_xpos, body_xmat, mocap_pos, mocap_quat, env_idx=0)
scene.update_from_arrays(body_xpos, body_xmat, mocap_pos, mocap_quat, env_idx=1)
assert any(mg.handle.visible for mg in groups)
handle_count = len(server.scene.batched)
env.sim.model.geom_rgba[:, :, :3] = torch.linspace(
0.9,
0.2,
env.num_envs,
device=env.device,
)[:, None, None]
scene.update_from_arrays(body_xpos, body_xmat, mocap_pos, mocap_quat, env_idx=0)
assert len(server.scene.batched) > handle_count
finally:
env.close()
def test_viser_convex_hulls_are_per_variant():
"""Convex-hull handles must differ across variants, not all show env0's hull."""
from contextlib import nullcontext
from mjlab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg
from mjlab.scene import SceneCfg
from mjlab.terrains import TerrainEntityCfg
from mjlab.viewer.viser.scene import MjlabViserScene, _PerWorldHullGroup
class _Handle:
def __init__(self, **kwargs):
self.visible = kwargs.get("visible", True)
self.batched_positions = kwargs.get("batched_positions", np.zeros((0, 3)))
self.batched_wxyzs = kwargs.get("batched_wxyzs", np.zeros((0, 4)))
self.batched_scales = kwargs.get("batched_scales")
self.batched_colors = kwargs.get("batched_colors")
self.batched_opacities = kwargs.get("batched_opacities")
self.position = kwargs.get("position", np.zeros(3))
self.wxyz = kwargs.get("wxyz", np.array([1.0, 0.0, 0.0, 0.0]))
self.vertices = kwargs.get("vertices")
self.faces = kwargs.get("faces")
def remove(self) -> None:
pass
class _Scene:
def __init__(self):
self.batched: list[tuple[tuple, dict, _Handle]] = []
def configure_environment_map(self, **_kwargs) -> None:
pass
def add_frame(self, *_args, **kwargs) -> _Handle:
return _Handle(**kwargs)
def add_grid(self, *_args, **kwargs) -> _Handle:
return _Handle(**kwargs)
def add_mesh_trimesh(self, *_args, **kwargs) -> _Handle:
return _Handle(**kwargs)
def add_batched_meshes_trimesh(self, *args, **kwargs) -> _Handle:
handle = _Handle(**kwargs)
self.batched.append((args, kwargs, handle))
return handle
def add_batched_meshes_simple(self, path, vertices, faces, **kwargs) -> _Handle:
# Capture the mesh identity so the test can compare hull shapes.
kwargs = dict(kwargs)
kwargs["vertices"] = np.asarray(vertices)
kwargs["faces"] = np.asarray(faces)
handle = _Handle(**kwargs)
self.batched.append(((path,), kwargs, handle))
return handle
class _Server:
def __init__(self):
self.scene = _Scene()
def atomic(self):
return nullcontext()
def flush(self) -> None:
pass
# Sphere and cone produce visibly different convex hulls.
env_cfg = ManagerBasedRlEnvCfg(
decimation=1,
scene=SceneCfg(
terrain=TerrainEntityCfg(terrain_type="plane"),
num_envs=4,
env_spacing=1.0,
entities={
"object": VariantEntityCfg(
variants={
"sphere": VariantCfg(_simple_sphere_spec, weight=0.5),
"cone": VariantCfg(_simple_cone_spec, weight=0.5),
},
init_state=EntityCfg.InitialStateCfg(pos=(0.0, 0.0, 0.2)),
)
},
),
)
env = ManagerBasedRlEnv(cfg=env_cfg, device="cpu")
try:
server = _Server()
scene = MjlabViserScene(
cast(Any, server),
env.sim.mj_model,
env.num_envs,
sim_model=env.sim.model,
expanded_fields=env.sim.expanded_fields,
)
groups: list[_PerWorldHullGroup] = list(scene._hull_per_world_groups)
# Two distinct variants -> at least two hull handles on the same body.
assert len(groups) >= 2, f"expected >=2 hull variants, got {len(groups)}"
all_envs = np.concatenate([g.env_ids for g in groups])
assert sorted(all_envs.tolist()) == list(range(env.num_envs))
# Hulls must be shape-distinct, not all copies of env0's hull.
shapes = {(g.handle.vertices.shape, g.handle.faces.shape) for g in groups}
assert len(shapes) >= 2, (
f"hull variants collapsed to one shape: {shapes} "
"(all envs would share env0's hull)"
)
body_xpos = env.sim.data.xpos.cpu().numpy()
body_xmat = env.sim.data.xmat.cpu().numpy()
scene.show_convex_hull = True
scene.show_only_selected = True
for target_env in range(env.num_envs):
scene.update_from_arrays(body_xpos, body_xmat, env_idx=target_env)
visible_groups = [g for g in groups if g.handle.visible]
assert len(visible_groups) == 1
assert target_env in visible_groups[0].env_ids
assert visible_groups[0].handle.batched_positions.shape[0] == 1
scene.show_only_selected = False
scene.update_from_arrays(body_xpos, body_xmat, env_idx=0)
assert all(g.handle.visible for g in groups)
finally:
env.close()
# DR consistency on variant scenes.
def _explicit_mass_variant(
mesh_name: str,
mass: float,
*,
cone: bool = False,
) -> mujoco.MjSpec:
"""Build a single-geom freejoint variant with an explicit body mass."""
spec = mujoco.MjSpec()
mesh = spec.add_mesh()
mesh.name = mesh_name
if cone:
mesh.make_cone(nedge=8, radius=0.05)
else:
mesh.make_sphere(subdivision=1)
body = spec.worldbody.add_body(name="prop")
body.add_freejoint()
body.explicitinertial = 1
body.mass = mass
body.ipos[:] = (0.0, 0.0, 0.0)
body.inertia[:] = (1e-4, 1e-4, 1e-4)
body.iquat[:] = (1.0, 0.0, 0.0, 0.0)
body.add_geom(
name="visual",
type=mujoco.mjtGeom.mjGEOM_MESH,
meshname=mesh_name,
contype=0,
conaffinity=0,
mass=0.0,
)
return spec
def test_dr_body_mass_scale_preserves_variant_baseline():
"""``dr.body_mass`` scale must use each variant's own baseline.
This is the load-bearing claim of ``_per_world_default_fields``: scaling
body_mass on a variant scene by a per-env factor must produce
``variant_default[env] * scale[env]``, not ``template_default * scale[env]``.
"""
from mjlab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg
from mjlab.envs.mdp import dr
from mjlab.managers.event_manager import EventTermCfg
from mjlab.managers.scene_entity_config import SceneEntityCfg
from mjlab.scene import SceneCfg
from mjlab.terrains import TerrainEntityCfg
light_mass = 0.1
heavy_mass = 1.0
scale = 2.0
object_cfg = VariantEntityCfg(
variants={
"light": VariantCfg(
lambda: _explicit_mass_variant("light", light_mass), weight=0.5
),
"heavy": VariantCfg(
lambda: _explicit_mass_variant("heavy", heavy_mass, cone=True), weight=0.5
),
},
init_state=EntityCfg.InitialStateCfg(pos=(0.0, 0.0, 0.2)),
)
env_cfg = ManagerBasedRlEnvCfg(
decimation=1,
scene=SceneCfg(
terrain=TerrainEntityCfg(terrain_type="plane"),
num_envs=4,
env_spacing=1.0,
entities={"object": object_cfg},
),
events={
"scale_mass": EventTermCfg(
func=dr.body_mass,
mode="startup",
params={
"asset_cfg": SceneEntityCfg("object", body_names=("prop",)),
"operation": "scale",
"ranges": (scale, scale), # deterministic factor
},
),
},
)
with pytest.warns(UserWarning, match="dr.body_mass only randomizes mass"):
env = ManagerBasedRlEnv(cfg=env_cfg, device="cpu")
try:
obj_body = int(env.scene["object"].indexing.root_body_id)
w2v = env.sim.world_to_variant["object"]
actual = env.sim.model.body_mass[:, obj_body].cpu()
variant_baseline = torch.tensor([light_mass, heavy_mass], dtype=actual.dtype)
expected = variant_baseline[w2v.cpu()] * scale
torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-5)
# Sanity: at least one env per variant, otherwise the test is vacuous.
assert (w2v == 0).any() and (w2v == 1).any()
finally:
env.close()
# Full env lifecycle.
def test_env_step_with_variants():
"""Build a full ManagerBasedRlEnv with variants; step without crashing."""
from mjlab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg
from mjlab.envs.mdp.events import reset_root_state_uniform
from mjlab.managers.event_manager import EventTermCfg
from mjlab.managers.scene_entity_config import SceneEntityCfg
from mjlab.scene import SceneCfg
from mjlab.terrains import TerrainEntityCfg
object_cfg = VariantEntityCfg(
variants={
"sphere": VariantCfg(_simple_sphere_spec, weight=0.5),
"cone": VariantCfg(_simple_cone_spec, weight=0.5),
},
init_state=EntityCfg.InitialStateCfg(pos=(0.0, 0.0, 0.2)),
)
env_cfg = ManagerBasedRlEnvCfg(
decimation=2,
scene=SceneCfg(
terrain=TerrainEntityCfg(terrain_type="plane"),
num_envs=4,
env_spacing=1.0,
entities={"object": object_cfg},
),
events={
"reset": EventTermCfg(
func=reset_root_state_uniform,
mode="reset",
params={
"pose_range": {},
"velocity_range": {},
"asset_cfg": SceneEntityCfg("object"),
},
),
},
)
env = ManagerBasedRlEnv(cfg=env_cfg, device="cpu")
obs, _ = env.reset()
actions = torch.zeros(env.num_envs, 0)
for _ in range(10):
obs, rew, term, trunc, info = env.step(actions)
# No NaN in positions.
qpos = env.sim.data.qpos[:].cpu().numpy()
assert np.all(np.isfinite(qpos))
env.close()
# Viewer: sameframe shortcut fix.
def _viewer_regression_sphere_spec() -> mujoco.MjSpec:
spec = mujoco.MjSpec()
m = spec.add_mesh()
m.name = "sphere"
m.make_sphere(subdivision=3)
m.scale[:] = (0.05, 0.05, 0.05)
body = spec.worldbody.add_body()
body.name = "prop"
body.add_freejoint()
g = body.add_geom()
g.name = "visual"
g.type = mujoco.mjtGeom.mjGEOM_MESH
g.meshname = "sphere"
return spec
def _viewer_regression_cone_spec() -> mujoco.MjSpec:
spec = mujoco.MjSpec()
m = spec.add_mesh()
m.name = "cone"
m.make_cone(nedge=16, radius=0.04)
m.scale[:] = (0.05, 0.05, 0.05)
body = spec.worldbody.add_body()
body.name = "prop"
body.add_freejoint()
g = body.add_geom()
g.name = "visual"
g.type = mujoco.mjtGeom.mjGEOM_MESH
g.meshname = "cone"
return spec
def test_sameframe_fix_makes_host_forward_match_variant():
"""Clearing sameframe shortcuts aligns host mj_forward with variant."""
base_model = _viewer_regression_sphere_spec().compile()
cone_model = _viewer_regression_cone_spec().compile()
# Sync cone's kinematic fields onto sphere's model (like viewer does).
for field in (
"geom_size",
"geom_pos",
"geom_quat",
"body_mass",
"body_inertia",
"body_ipos",
"body_iquat",
):
getattr(base_model, field)[:] = getattr(cone_model, field)
base_data = mujoco.MjData(base_model)
base_data.qpos[:] = cone_model.qpos0
base_data.qpos[2] = 0.05
mujoco.mj_forward(base_model, base_data)
cone_data = mujoco.MjData(cone_model)
cone_data.qpos[:] = cone_model.qpos0
cone_data.qpos[2] = 0.05
mujoco.mj_forward(cone_model, cone_data)
# Before fix: positions differ due to stale sameframe flags.
assert not np.allclose(base_data.geom_xpos, cone_data.geom_xpos)
# After fix: clearing sameframe makes them match.
disable_model_sameframe_shortcuts(base_model)
mujoco.mj_forward(base_model, base_data)
np.testing.assert_allclose(base_data.geom_xpos, cone_data.geom_xpos, atol=1e-6)
def test_sync_model_fields_copies_only_requested_env_fields():
"""Viewer model sync copies explicit fields and leaves others unchanged."""
model = _simple_sphere_spec().compile()
class _SimModel:
geom_rgba = torch.tensor(
[
[[0.1, 0.2, 0.3, 0.4]],
[[0.5, 0.6, 0.7, 0.8]],
],
dtype=torch.float32,
)
geom_pos = torch.tensor(
[
[[1.0, 2.0, 3.0]],
[[4.0, 5.0, 6.0]],
],
dtype=torch.float32,
)
original_geom_pos = model.geom_pos.copy()
sync_model_fields(model, _SimModel(), {"geom_rgba"}, env_idx=1)
np.testing.assert_allclose(model.geom_rgba, [[0.5, 0.6, 0.7, 0.8]])
np.testing.assert_allclose(model.geom_pos, original_geom_pos)
05_software/train/rc_mjlab/mjlab/tests/test_projected_gravity_sensor.py
+213
View File
@@ -0,0 +1,213 @@
"""Tests for sensor-based projected gravity (framezaxis up-vector sensor).
The shipped robots expose a ``framezaxis`` sensor that outputs the world Z-axis in the
IMU site frame; negating it gives projected gravity. These tests check the sensor (and
the ``projected_gravity_from_sensor`` observation that wraps it) against an independent
ground-truth computation, and verify that -- unlike the entity-data
``projected_gravity_b`` -- it tracks the IMU site orientation, which is what makes IMU
mounting domain randomization observable.
"""
from __future__ import annotations
import math
from typing import TYPE_CHECKING, cast
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.envs.mdp.observations import projected_gravity_from_sensor
from mjlab.managers.scene_entity_config import SceneEntityCfg
from mjlab.scene import Scene, SceneCfg
from mjlab.sim.sim import Simulation, SimulationCfg
if TYPE_CHECKING:
from mjlab.envs import ManagerBasedRlEnv
# Gravity points along world -Z; projected gravity is this expressed in a body frame.
_GRAVITY_DIR_W = (0.0, 0.0, -1.0)
def _quat_to_mat(q: tuple[float, float, float, float]) -> torch.Tensor:
"""Rotation matrix from a (w, x, y, z) quaternion. Independent of MuJoCo/mjlab."""
w, x, y, z = q
return torch.tensor(
[
[1 - 2 * (y * y + z * z), 2 * (x * y - w * z), 2 * (x * z + w * y)],
[2 * (x * y + w * z), 1 - 2 * (x * x + z * z), 2 * (y * z - w * x)],
[2 * (x * z - w * y), 2 * (y * z + w * x), 1 - 2 * (x * x + y * y)],
],
dtype=torch.float64,
)
def _expected_projected_gravity(q: tuple[float, float, float, float]) -> torch.Tensor:
"""Ground-truth projected gravity for a body with world orientation ``q``.
proj = R(q)^T @ g_world, computed from an explicit rotation matrix so it does not
share a code path with the sensor or with ``projected_gravity_b``.
"""
g_w = torch.tensor(_GRAVITY_DIR_W, dtype=torch.float64)
return _quat_to_mat(q).T @ g_w
class Env:
"""Minimal env stub for driving observation and dr functions in tests."""
def __init__(self, scene, sim, device):
self.scene = scene
self.sim = sim
self.num_envs = scene.num_envs
self.device = device
def _make_env(scene, sim, device) -> ManagerBasedRlEnv:
"""Build the env stub, typed as the real env for the functions under test."""
return cast("ManagerBasedRlEnv", Env(scene, sim, device))
@pytest.fixture(scope="module")
def device():
return get_test_device()
def _robot_xml(site_euler: str = "0 0 0") -> str:
"""Free-floating box with an IMU site and the framezaxis up-vector sensor."""
return f"""
<mujoco>
<worldbody>
<body name="base" pos="0 0 1">
<freejoint name="free_joint"/>
<geom name="base_geom" type="box" size="0.2 0.2 0.1" mass="5.0"/>
<site name="imu" pos="0.05 0 0" euler="{site_euler}"/>
</body>
</worldbody>
<sensor>
<framezaxis name="imu_upvector" objtype="body" objname="world"
reftype="site" refname="imu"/>
</sensor>
</mujoco>
"""
def _build(xml: str, device: str, num_envs: int = 2):
entity_cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(xml))
scene = Scene(
SceneCfg(num_envs=num_envs, env_spacing=3.0, entities={"robot": entity_cfg}),
device,
)
model = scene.compile()
sim = Simulation(
num_envs=num_envs, cfg=SimulationCfg(njmax=20), model=model, device=device
)
scene.initialize(sim.mj_model, sim.model, sim.data)
return scene, sim
def _set_root_quat(robot, q: tuple[float, float, float, float], device: str) -> None:
root_state = robot.data.default_root_state.clone()
root_state[:, 3:7] = torch.tensor(q, device=device, dtype=root_state.dtype)
robot.write_root_state_to_sim(root_state)
def test_sensor_matches_ground_truth_when_site_aligned(device):
"""Sensor and entity both equal hand-computed projected gravity for a tilted base."""
scene, sim = _build(_robot_xml(), device)
robot = scene["robot"]
# Compose a 0.6 rad roll with a 0.3 rad pitch into a single root quaternion.
ax = (math.cos(0.3), math.sin(0.3), 0.0, 0.0)
ay = (math.cos(0.15), 0.0, math.sin(0.15), 0.0)
q = (
ax[0] * ay[0] - ax[1] * ay[1] - ax[2] * ay[2] - ax[3] * ay[3],
ax[0] * ay[1] + ax[1] * ay[0] + ax[2] * ay[3] - ax[3] * ay[2],
ax[0] * ay[2] - ax[1] * ay[3] + ax[2] * ay[0] + ax[3] * ay[1],
ax[0] * ay[3] + ax[1] * ay[2] - ax[2] * ay[1] + ax[3] * ay[0],
)
_set_root_quat(robot, q, device)
sim.forward()
expected = _expected_projected_gravity(q).to(device=device, dtype=torch.float32)
# Guard against a vacuous pass: the tilt must actually move gravity off straight-down.
straight_down = torch.tensor(_GRAVITY_DIR_W, device=device)
assert (expected - straight_down).abs().max() > 0.3
sensor_grav = -scene["robot/imu_upvector"].data
entity_grav = robot.data.projected_gravity_b
torch.testing.assert_close(sensor_grav[0], expected, atol=1e-5, rtol=0)
torch.testing.assert_close(entity_grav[0], expected, atol=1e-5, rtol=0)
def test_observation_fn_tracks_site_orientation(device):
"""The observation fn reflects IMU site tilt; the entity-data version does not.
With the base upright but the IMU site rolled 30 deg about x, projected gravity in the
site frame is (0, -sin30, -cos30). The entity-data version stays straight-down because
it uses the root body orientation and is blind to the site.
"""
scene_rot, sim_rot = _build(_robot_xml(site_euler="30 0 0"), device)
scene_flat, sim_flat = _build(_robot_xml(site_euler="0 0 0"), device)
sim_rot.forward()
sim_flat.forward()
# Drive through the actual shipped observation function, not the raw sensor.
env_rot = _make_env(scene_rot, sim_rot, device)
env_flat = _make_env(scene_flat, sim_flat, device)
grav_rot = projected_gravity_from_sensor(env_rot, "robot/imu_upvector")
grav_flat = projected_gravity_from_sensor(env_flat, "robot/imu_upvector")
expected_rot = torch.tensor(
[0.0, -math.sin(math.radians(30)), -math.cos(math.radians(30))], device=device
)
straight_down = torch.tensor(_GRAVITY_DIR_W, device=device)
torch.testing.assert_close(grav_rot[0], expected_rot, atol=1e-5, rtol=0)
torch.testing.assert_close(grav_flat[0], straight_down, atol=1e-5, rtol=0)
# The entity-data version is unchanged by the site rotation (so it cannot be used to
# observe IMU mounting randomization), confirming why the sensor path is needed.
entity_rot = scene_rot["robot"].data.projected_gravity_b
torch.testing.assert_close(entity_rot[0], straight_down, atol=1e-5, rtol=0)
@pytest.mark.filterwarnings(
"ignore:Use of index_put_ on expanded tensors is deprecated:UserWarning"
)
def test_site_quat_randomization_changes_sensor(device):
"""The full DR path: running ``dr.site_quat`` perturbs the gravity observation.
This is what the G1 example configs rely on -- randomizing the IMU site orientation
must show up in the sensor-based projected gravity, per-environment.
"""
scene, sim = _build(_robot_xml(), device, num_envs=4)
sim.expand_model_fields(("site_quat",))
env = _make_env(scene, sim, device)
sim.forward()
straight_down = torch.tensor(_GRAVITY_DIR_W, device=device)
before = projected_gravity_from_sensor(env, "robot/imu_upvector").clone()
# Upright base + identity site quat => straight-down gravity in every env.
torch.testing.assert_close(before, straight_down.expand_as(before), atol=1e-5, rtol=0)
torch.manual_seed(0)
dr.site_quat(
env,
env_ids=None,
roll_range=(-0.3, 0.3),
pitch_range=(-0.3, 0.3),
yaw_range=(-0.3, 0.3),
asset_cfg=SceneEntityCfg("robot", site_names=("imu",)),
)
sim.forward()
after = projected_gravity_from_sensor(env, "robot/imu_upvector")
# Randomization moved the reading off straight-down and made it env-dependent.
assert (after - straight_down).abs().max() > 0.05
assert not torch.allclose(after, before, atol=1e-3)
assert torch.unique(after, dim=0).shape[0] >= 2
# The perturbation is a rotation, so gravity stays a unit vector.
norms = torch.linalg.norm(after, dim=-1)
torch.testing.assert_close(norms, torch.ones_like(norms), atol=1e-5, rtol=0)
05_software/train/rc_mjlab/mjlab/tests/test_random.py
-30
View File
@@ -1,30 +0,0 @@
"""Tests for mjlab.utils.random."""
import subprocess
import sys
import textwrap
def test_seed_rng_cpu_device_does_not_initialize_warp_cuda() -> None:
"""seed_rng(device="cpu") must not initialize Warp's CUDA runtime.
Runs in a subprocess so that Warp is guaranteed uninitialized before the
call.
"""
script = textwrap.dedent("""
import warp as wp
from mjlab.utils.random import seed_rng
assert wp._src.context.runtime is None, "Warp must not be initialized yet"
seed_rng(42, device="cpu")
rt = wp._src.context.runtime
if rt is not None:
cuda = [d for d in wp.get_devices() if "cuda" in str(d)]
assert not cuda, f"seed_rng(device='cpu') initialized CUDA devices {cuda}"
""")
result = subprocess.run(
[sys.executable, "-c", script], capture_output=True, text=True
)
assert result.returncode == 0, (
f"subprocess failed:\nstdout={result.stdout}\nstderr={result.stderr}"
)
05_software/train/rc_mjlab/mjlab/tests/test_raycast_sensor.py
+12 -4
View File
@@ -936,7 +936,7 @@ def test_multi_frame_body_exclusion(device):
should skip body_b's own geom but HIT body_a's platform. Frame A's
rays should skip body_a and hit the floor.
"""
xml = """
body_a_xml = """
<mujoco>
<worldbody>
<geom name="floor" type="plane" size="10 10 0.1" pos="0 0 0"/>
@@ -945,6 +945,12 @@ def test_multi_frame_body_exclusion(device):
<geom name="geom_a" type="box" size="2 2 0.1" mass="5.0"/>
<site name="site_a" pos="0 0 0"/>
</body>
</worldbody>
</mujoco>
"""
body_b_xml = """
<mujoco>
<worldbody>
<body name="body_b" pos="0 0 3">
<freejoint name="free_b"/>
<geom name="geom_b" type="box" size="0.5 0.5 0.5" mass="5.0"/>
@@ -957,15 +963,17 @@ def test_multi_frame_body_exclusion(device):
cfg = RayCastSensorCfg(
name="multi",
frame=(
ObjRef(type="site", name="site_a", entity="robot"),
ObjRef(type="site", name="site_b", entity="robot"),
ObjRef(type="site", name="site_a", entity="body_a"),
ObjRef(type="site", name="site_b", entity="body_b"),
),
pattern=GridPatternCfg(size=(0.0, 0.0), resolution=0.1),
max_distance=10.0,
exclude_parent_body=True,
)
scene, sim = make_scene_and_sim(device, xml, (cfg,))
scene, sim = make_scene_and_sim(
device, {"body_a": body_a_xml, "body_b": body_b_xml}, (cfg,)
)
sim.step()
sim.sense()
05_software/train/rc_mjlab/mjlab/tests/test_runner.py
+83
View File
@@ -4,6 +4,7 @@ import ast
import tempfile
from dataclasses import asdict
from pathlib import Path
from unittest.mock import MagicMock, patch
import mujoco
import onnx
@@ -521,3 +522,85 @@ def test_onnx_motion_model_clamps_out_of_bounds_time_step():
_, joint_pos, *_ = model(x, time_step)
torch.testing.assert_close(joint_pos, motion.joint_pos[num_steps - 1 : num_steps])
def _make_tracking_runner_shell(registry_name, logger_type, upload_model=True):
"""Build a MotionTrackingOnPolicyRunner with all heavy parts mocked out."""
from mjlab.tasks.tracking.rl.runner import MotionTrackingOnPolicyRunner
runner = MotionTrackingOnPolicyRunner.__new__(MotionTrackingOnPolicyRunner)
runner.registry_name = registry_name
runner.cfg = {"upload_model": upload_model}
runner.logger = MagicMock()
runner.logger.logger_type = logger_type
mock_motion_term = MagicMock()
mock_motion_term.cfg.anchor_body_name = "pelvis"
mock_motion_term.cfg.body_names = ["body1"]
runner.env = MagicMock()
runner.env.unwrapped.command_manager.get_term.return_value = mock_motion_term
return runner
@pytest.mark.parametrize("logger_type", ["wandb", "WandbLogWriter"])
def test_tracking_runner_registers_artifact_for_wandb_logger_types(
logger_type, monkeypatch, tmp_path
):
"""use_artifact is called for both legacy 'wandb' and current 'WandbLogWriter' logger types.
Regression test: rsl-rl-lib 5.4 renamed the WandB logger type from 'wandb'
to 'WandbLogWriter'. If only 'wandb' is checked, use_artifact is silently
skipped and the nightly report fails with 'No motion artifact found in the run.'
"""
from mjlab.rl.runner import MjlabOnPolicyRunner
from mjlab.tasks.tracking.rl import runner as runner_mod
runner = _make_tracking_runner_shell("org/motions/motion:latest", logger_type)
monkeypatch.setattr(MjlabOnPolicyRunner, "save", lambda *a, **kw: None)
monkeypatch.setattr(runner_mod, "get_base_metadata", lambda *a: {})
monkeypatch.setattr(runner_mod, "attach_metadata_to_onnx", lambda *a: None)
monkeypatch.setattr(
runner.env.unwrapped.__class__,
"export_policy_to_onnx",
lambda *a, **kw: None,
raising=False,
)
checkpoint = tmp_path / "run-dir" / "model_100.pt"
checkpoint.parent.mkdir()
checkpoint.touch()
mock_run = MagicMock()
mock_run.name = "test-run"
with patch.object(runner_mod, "wandb") as mock_wandb:
mock_wandb.run = mock_run
runner.export_policy_to_onnx = MagicMock()
runner.save(str(checkpoint))
mock_run.use_artifact.assert_called_once_with("org/motions/motion:latest")
def test_tracking_runner_does_not_register_artifact_for_tensorboard(
monkeypatch, tmp_path
):
"""use_artifact is NOT called when using the tensorboard logger."""
from mjlab.rl.runner import MjlabOnPolicyRunner
from mjlab.tasks.tracking.rl import runner as runner_mod
runner = _make_tracking_runner_shell("org/motions/motion:latest", "tensorboard")
monkeypatch.setattr(MjlabOnPolicyRunner, "save", lambda *a, **kw: None)
monkeypatch.setattr(runner_mod, "get_base_metadata", lambda *a: {})
monkeypatch.setattr(runner_mod, "attach_metadata_to_onnx", lambda *a: None)
checkpoint = tmp_path / "run-dir" / "model_100.pt"
checkpoint.parent.mkdir()
checkpoint.touch()
with patch.object(runner_mod, "wandb") as mock_wandb:
runner.export_policy_to_onnx = MagicMock()
runner.save(str(checkpoint))
mock_wandb.run.use_artifact.assert_not_called()
05_software/train/rc_mjlab/mjlab/tests/test_terrains.py
+59 -7
View File
@@ -2,8 +2,15 @@
import mujoco
import numpy as np
import pytest
from mjlab.terrains.primitive_terrains import BoxSteppingStonesTerrainCfg
from mjlab.terrains.config import ALL_TERRAIN_PRESETS
from mjlab.terrains.primitive_terrains import (
_MIN_BORDER_HEIGHT,
BoxInvertedPyramidStairsTerrainCfg,
BoxPyramidStairsTerrainCfg,
BoxSteppingStonesTerrainCfg,
)
_CFG = BoxSteppingStonesTerrainCfg(
proportion=1.0,
@@ -37,12 +44,10 @@ def _generate_stones(
if geom is None:
continue
pos, size = geom.pos, geom.size
# Skip platform, floor, and border geoms.
is_platform = (
np.isclose(pos[0], center)
and np.isclose(pos[1], center)
and np.isclose(size[0], cfg.platform_width / 2, atol=1e-4)
)
# Skip platform, floor, and border geoms. The platform is the geom centered
# exactly at the patch center (its size is grid-snapped, not the configured
# width, so it is identified by position alone).
is_platform = np.isclose(pos[0], center) and np.isclose(pos[1], center)
is_full_span = np.isclose(size[0], cfg.size[0] / 2) or np.isclose(
size[1], cfg.size[1] / 2
)
@@ -74,3 +79,50 @@ def test_stone_size_decreases_with_difficulty():
sizes[difficulty] = np.mean([hx + hy for _, _, hx, hy in stones])
assert sizes[0.0] > sizes[1.0]
@pytest.mark.parametrize(
"cfg_cls", [BoxPyramidStairsTerrainCfg, BoxInvertedPyramidStairsTerrainCfg]
)
def test_pyramid_stairs_border_present_at_zero_difficulty(cfg_cls):
"""At difficulty 0 the step height collapses to 0, but the flat border frame
must still be generated as solid, non-degenerate geometry (regression for the
empty-boundary bug, issue #1033)."""
cfg = cfg_cls(
size=(8.0, 8.0),
step_height_range=(0.0, 0.2),
step_width=0.3,
platform_width=3.0,
border_width=1.0,
)
spec = mujoco.MjSpec()
spec.worldbody.add_body(name="terrain")
output = cfg.function(difficulty=0.0, spec=spec, rng=np.random.default_rng(0))
# The border frame sits below z=0 (top flush at ground level); inner step
# boxes are centered at z=0. Identify the frame by its downward offset.
border_geoms = [
g.geom for g in output.geometries if g.geom is not None and g.geom.pos[2] < -1e-4
]
assert len(border_geoms) == 4, "Expected four border frame boxes."
for geom in border_geoms:
# Each frame box must be solid, not a degenerate zero-height geom, and its
# top must be flush with the ground plane at z=0.
assert geom.size[2] >= _MIN_BORDER_HEIGHT / 2 - 1e-9
assert np.isclose(geom.pos[2] + geom.size[2], 0.0, atol=1e-6)
@pytest.mark.parametrize("preset_name", sorted(ALL_TERRAIN_PRESETS))
@pytest.mark.parametrize("difficulty", [0.0, 1.0])
def test_preset_compiles_across_difficulty(preset_name, difficulty):
"""Every terrain preset must generate compilable MuJoCo geometry across the
full difficulty range. Difficulty 0 is exercised explicitly because curriculum
row 0 lands there deterministically, which previously produced degenerate
geometry (zero-height hfields, NaN colors, missing borders)."""
cfg = ALL_TERRAIN_PRESETS[preset_name](size=(8.0, 8.0))
spec = mujoco.MjSpec()
spec.worldbody.add_body(name="terrain")
cfg.function(difficulty=difficulty, spec=spec, rng=np.random.default_rng(0))
# Compiling validates geom/hfield sizes and rgba values (catches NaNs and
# non-positive sizes that MuJoCo rejects).
spec.compile()
05_software/train/rc_mjlab/mjlab/tests/test_tracking_metrics.py
+58 -34
View File
@@ -1,5 +1,6 @@
"""Tests for motion tracking evaluation metrics."""
import math
from unittest.mock import Mock
import pytest
@@ -34,11 +35,11 @@ def mock_command():
def test_mpkpe_zero_when_positions_match(mock_command):
"""Test MPKPE is zero when positions are identical."""
"""Test MPKPE is zero when global positions are identical."""
num_bodies = len(mock_command.cfg.body_names)
positions = torch.rand(mock_command.num_envs, num_bodies, 3)
mock_command.body_pos_relative_w = positions.clone()
mock_command.body_pos_w = positions.clone()
mock_command.robot_body_pos_w = positions.clone()
mpkpe = compute_mpkpe(mock_command)
@@ -48,10 +49,10 @@ def test_mpkpe_zero_when_positions_match(mock_command):
def test_mpkpe_correct_error(mock_command):
"""Test MPKPE computes correct mean error."""
"""Test MPKPE computes the correct mean global error."""
num_bodies = len(mock_command.cfg.body_names)
mock_command.body_pos_relative_w = torch.zeros(mock_command.num_envs, num_bodies, 3)
mock_command.body_pos_w = torch.zeros(mock_command.num_envs, num_bodies, 3)
mock_command.robot_body_pos_w = torch.zeros(mock_command.num_envs, num_bodies, 3)
mock_command.robot_body_pos_w[:, :, 0] = 1.0 # 1 unit offset in x
@@ -60,58 +61,71 @@ def test_mpkpe_correct_error(mock_command):
assert torch.allclose(mpkpe, torch.ones(mock_command.num_envs), atol=1e-6)
def test_r_mpkpe_invariant_to_global_translation(mock_command):
"""Test R-MPKPE is invariant to global translation."""
def test_mpkpe_uses_global_reference(mock_command):
"""MPKPE must read the global reference, not the drift-cancelled one.
Pins issue #1006: setting body_pos_relative_w to match the robot exactly
would yield zero error if it were (incorrectly) used; the metric must
instead follow body_pos_w.
"""
num_bodies = len(mock_command.cfg.body_names)
robot_pos = torch.rand(mock_command.num_envs, num_bodies, 3)
mock_command.robot_body_pos_w = robot_pos.clone()
mock_command.body_pos_relative_w = robot_pos.clone() # zero error if misused
mock_command.body_pos_w = robot_pos.clone()
mock_command.body_pos_w[:, :, 0] += 1.0 # 1 unit of global drift
mock_command.anchor_pos_w = torch.zeros(mock_command.num_envs, 3)
mock_command.body_pos_w = torch.rand(mock_command.num_envs, num_bodies, 3)
mock_command.robot_anchor_pos_w = torch.zeros(mock_command.num_envs, 3)
mock_command.robot_body_pos_w = mock_command.body_pos_w.clone()
mpkpe = compute_mpkpe(mock_command)
r_mpkpe_1 = compute_root_relative_mpkpe(mock_command)
# Translate everything by large offset.
offset = torch.tensor([100.0, 200.0, 300.0])
mock_command.anchor_pos_w = offset.expand(mock_command.num_envs, 3).clone()
mock_command.body_pos_w = mock_command.body_pos_w + offset
mock_command.robot_anchor_pos_w = offset.expand(mock_command.num_envs, 3).clone()
mock_command.robot_body_pos_w = mock_command.robot_body_pos_w + offset
r_mpkpe_2 = compute_root_relative_mpkpe(mock_command)
assert torch.allclose(r_mpkpe_1, r_mpkpe_2, atol=1e-5)
assert torch.allclose(mpkpe, torch.ones(mock_command.num_envs), atol=1e-6)
def test_r_mpkpe_detects_relative_error(mock_command):
"""Test R-MPKPE detects errors in relative positions."""
def test_r_mpkpe_zero_when_relative_positions_match(mock_command):
"""R-MPKPE is zero when re-anchored positions are identical."""
num_bodies = len(mock_command.cfg.body_names)
positions = torch.rand(mock_command.num_envs, num_bodies, 3)
mock_command.anchor_pos_w = torch.zeros(mock_command.num_envs, 3)
mock_command.body_pos_w = torch.zeros(mock_command.num_envs, num_bodies, 3)
mock_command.body_pos_w[:, :, 0] = 1.0 # Bodies 1 unit from anchor
mock_command.body_pos_relative_w = positions.clone()
mock_command.robot_body_pos_w = positions.clone()
mock_command.robot_anchor_pos_w = torch.zeros(mock_command.num_envs, 3)
mock_command.robot_body_pos_w = torch.zeros(mock_command.num_envs, num_bodies, 3)
mock_command.robot_body_pos_w[:, :, 0] = 2.0 # Bodies 2 units from anchor
r_mpkpe = compute_root_relative_mpkpe(mock_command)
assert r_mpkpe.shape == (mock_command.num_envs,)
assert torch.allclose(r_mpkpe, torch.zeros(mock_command.num_envs), atol=1e-6)
def test_r_mpkpe_uses_relative_reference(mock_command):
"""R-MPKPE reads the re-anchored reference, not the global one.
Setting body_pos_w to match the robot exactly would yield zero error if
it were (incorrectly) used; the metric must instead follow
body_pos_relative_w.
"""
num_bodies = len(mock_command.cfg.body_names)
robot_pos = torch.rand(mock_command.num_envs, num_bodies, 3)
mock_command.robot_body_pos_w = robot_pos.clone()
mock_command.body_pos_w = robot_pos.clone() # zero error if misused
mock_command.body_pos_relative_w = robot_pos.clone()
mock_command.body_pos_relative_w[:, :, 0] += 1.0 # 1 unit of local pose error
r_mpkpe = compute_root_relative_mpkpe(mock_command)
assert torch.allclose(r_mpkpe, torch.ones(mock_command.num_envs), atol=1e-6)
def test_joint_velocity_error(mock_command):
"""Test joint velocity error computes correct L2 norm."""
def test_joint_velocity_error_rms(mock_command):
"""Joint velocity error is the per-joint RMS of the velocity error."""
num_joints = 3
mock_command.joint_vel = torch.zeros(mock_command.num_envs, num_joints)
mock_command.robot_joint_vel = torch.zeros(mock_command.num_envs, num_joints)
mock_command.robot_joint_vel[:, 0] = 3.0
mock_command.robot_joint_vel[:, 1] = 4.0 # Error [3, 4, 0] has norm 5
mock_command.robot_joint_vel[:, 1] = 4.0 # Error [3, 4, 0]
error = compute_joint_velocity_error(mock_command)
assert torch.allclose(error, torch.ones(mock_command.num_envs) * 5.0, atol=1e-6)
expected = math.sqrt((3.0**2 + 4.0**2 + 0.0**2) / num_joints)
assert torch.allclose(error, torch.ones(mock_command.num_envs) * expected, atol=1e-6)
def test_ee_position_error_only_uses_specified_bodies(mock_command):
@@ -153,3 +167,13 @@ def test_ee_orientation_error_detects_rotation(mock_command):
# Error should be approximately pi/2 radians.
expected = torch.ones(mock_command.num_envs) * (3.14159 / 2)
assert torch.allclose(error, expected, atol=0.01)
def test_ee_metrics_raise_on_unknown_body(mock_command):
"""Unknown end-effector names raise instead of silently scoring zero."""
num_bodies = len(mock_command.cfg.body_names)
mock_command.body_pos_relative_w = torch.zeros(mock_command.num_envs, num_bodies, 3)
mock_command.robot_body_pos_w = torch.zeros(mock_command.num_envs, num_bodies, 3)
with pytest.raises(ValueError, match="not tracked"):
compute_ee_position_error(mock_command, ("nonexistent_body",))
05_software/train/rc_mjlab/mjlab/tests/test_variants.py
+1772
View File
File diff suppressed because it is too large Load Diff
05_software/train/rc_mjlab/mjlab/tests/test_xml_actuator.py
+2 -1
View File
@@ -4,7 +4,7 @@ import mujoco
import pytest
from conftest import get_test_device
from mjlab.actuator import XmlActuatorCfg
from mjlab.actuator import XmlActuator, XmlActuatorCfg
from mjlab.entity import Entity, EntityArticulationInfoCfg, EntityCfg
from mjlab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg, mdp
from mjlab.managers.observation_manager import ObservationGroupCfg, ObservationTermCfg
@@ -160,6 +160,7 @@ def test_xml_actuator_explicit_command_field_bypasses_detection():
entity.compile()
actuator = entity._actuators[0]
assert isinstance(actuator, XmlActuator)
assert actuator.command_field == "effort"
assert actuator._target_names == ["joint1"]