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