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
364 lines
12 KiB
Python
364 lines
12 KiB
Python
"""Tests for delayed actuators."""
|
|
|
|
import mujoco
|
|
import pytest
|
|
import torch
|
|
from conftest import get_test_device, load_fixture_xml
|
|
|
|
from mjlab.actuator import (
|
|
BuiltinPositionActuatorCfg,
|
|
IdealPdActuatorCfg,
|
|
)
|
|
from mjlab.entity import Entity, EntityArticulationInfoCfg, EntityCfg
|
|
from mjlab.sim.sim import Simulation, SimulationCfg
|
|
|
|
ROBOT_XML = load_fixture_xml("floating_base_articulated")
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def device():
|
|
return get_test_device()
|
|
|
|
|
|
def create_entity_with_delayed_builtin(delay_min_lag=0, delay_max_lag=3):
|
|
cfg = EntityCfg(
|
|
spec_fn=lambda: mujoco.MjSpec.from_string(ROBOT_XML),
|
|
articulation=EntityArticulationInfoCfg(
|
|
actuators=(
|
|
BuiltinPositionActuatorCfg(
|
|
target_names_expr=("joint.*",),
|
|
effort_limit=100.0,
|
|
stiffness=80.0,
|
|
damping=10.0,
|
|
delay_min_lag=delay_min_lag,
|
|
delay_max_lag=delay_max_lag,
|
|
),
|
|
)
|
|
),
|
|
)
|
|
return Entity(cfg)
|
|
|
|
|
|
def create_entity_with_delayed_ideal(delay_min_lag=0, delay_max_lag=3):
|
|
cfg = EntityCfg(
|
|
spec_fn=lambda: mujoco.MjSpec.from_string(ROBOT_XML),
|
|
articulation=EntityArticulationInfoCfg(
|
|
actuators=(
|
|
IdealPdActuatorCfg(
|
|
target_names_expr=("joint.*",),
|
|
effort_limit=100.0,
|
|
stiffness=80.0,
|
|
damping=10.0,
|
|
delay_min_lag=delay_min_lag,
|
|
delay_max_lag=delay_max_lag,
|
|
),
|
|
)
|
|
),
|
|
)
|
|
return Entity(cfg)
|
|
|
|
|
|
def initialize_entity(entity, device, num_envs=1):
|
|
model = entity.compile()
|
|
sim_cfg = SimulationCfg()
|
|
sim = Simulation(num_envs=num_envs, cfg=sim_cfg, model=model, device=device)
|
|
entity.initialize(model, sim.model, sim.data, device)
|
|
return entity, sim
|
|
|
|
|
|
def test_delayed_builtin_applies_constant_delay(device):
|
|
"""Test that delayed builtin actuator delays position targets."""
|
|
entity = create_entity_with_delayed_builtin(delay_min_lag=2, delay_max_lag=2)
|
|
entity, sim = initialize_entity(entity, device)
|
|
|
|
# Set position targets for 3 steps.
|
|
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),
|
|
]
|
|
|
|
joint_vel = torch.zeros(1, 2, device=device)
|
|
|
|
for target in targets:
|
|
entity.set_joint_position_target(target)
|
|
entity.set_joint_velocity_target(joint_vel)
|
|
entity.set_joint_effort_target(torch.zeros(1, 2, device=device))
|
|
entity.write_data_to_sim()
|
|
|
|
# After 3 steps with lag=2, the output should be the target from step 0.
|
|
ctrl = sim.data.ctrl[0]
|
|
# With constant lag=2, after 3 appends, we expect target from step 0.
|
|
assert torch.allclose(ctrl, targets[0][0], atol=1e-5)
|
|
|
|
|
|
def test_delayed_ideal_applies_delay(device):
|
|
"""Test that delayed ideal actuator delays 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)
|
|
|
|
# Set position targets for 3 steps.
|
|
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 target in targets:
|
|
entity.set_joint_position_target(target)
|
|
entity.set_joint_velocity_target(joint_vel)
|
|
entity.set_joint_effort_target(torch.zeros(1, 2, device=device))
|
|
entity.write_data_to_sim()
|
|
sim.forward() # Compute actuator forces
|
|
|
|
# The computed torque should use the delayed target from step 0.
|
|
joint_v_adr = entity.indexing.joint_v_adr
|
|
qfrc = sim.data.qfrc_actuator[0, joint_v_adr]
|
|
|
|
# Expected torque: kp * (delayed_target - joint_pos) + kd * (0 - joint_vel) + 0
|
|
# = 80.0 * targets[0] + 0 = 80.0 * [0.1, 0.2]
|
|
expected_torque = 80.0 * targets[0][0]
|
|
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)
|
|
entity, _ = initialize_entity(entity, device, num_envs=2)
|
|
|
|
# Set some targets to fill the buffer.
|
|
entity.set_joint_position_target(torch.ones(2, 2, device=device) * 0.5)
|
|
entity.set_joint_velocity_target(torch.zeros(2, 2, device=device))
|
|
entity.set_joint_effort_target(torch.zeros(2, 2, device=device))
|
|
entity.write_data_to_sim()
|
|
|
|
# Reset env 0.
|
|
entity.reset(torch.tensor([0], device=device))
|
|
|
|
# Check that delay buffer was reset for env 0.
|
|
actuator = entity.actuators[0]
|
|
assert actuator.has_delay
|
|
assert actuator._delay_buffer is not None
|
|
assert actuator._delay_buffer.current_lags[0] == 0
|
|
|
|
|
|
def test_delayed_actuator_set_lags(device):
|
|
"""Test that set_lags sets lag values on all delay buffers."""
|
|
entity = create_entity_with_delayed_builtin(delay_min_lag=0, delay_max_lag=5)
|
|
entity, _ = initialize_entity(entity, device, num_envs=4)
|
|
|
|
actuator = entity.actuators[0]
|
|
assert actuator.has_delay
|
|
|
|
# Set lags for all environments.
|
|
lags = torch.tensor([1, 2, 3, 4], device=device)
|
|
actuator.set_lags(lags)
|
|
|
|
# Check that lags were set.
|
|
buffer = actuator._delay_buffer
|
|
assert buffer is not None
|
|
assert torch.equal(buffer.current_lags, lags)
|
|
|
|
|
|
def test_delayed_actuator_set_lags_subset(device):
|
|
"""Test that set_lags can set lag values for a subset of environments."""
|
|
entity = create_entity_with_delayed_builtin(delay_min_lag=0, delay_max_lag=5)
|
|
entity, _ = initialize_entity(entity, device, num_envs=4)
|
|
|
|
actuator = entity.actuators[0]
|
|
assert actuator.has_delay
|
|
|
|
# Set lags for envs 1 and 3 only.
|
|
env_ids = torch.tensor([1, 3], device=device)
|
|
lags = torch.tensor([4, 5], device=device)
|
|
actuator.set_lags(lags, env_ids)
|
|
|
|
# Check that only specified envs were updated.
|
|
buffer = actuator._delay_buffer
|
|
assert buffer is not None
|
|
assert buffer.current_lags[0] == 0 # Unchanged (initial value)
|
|
assert buffer.current_lags[1] == 4
|
|
assert buffer.current_lags[2] == 0 # Unchanged
|
|
assert buffer.current_lags[3] == 5
|
|
|
|
|
|
def test_delayed_actuator_set_lags_clamps_to_range(device):
|
|
"""Test that set_lags clamps values to the configured lag range."""
|
|
entity = create_entity_with_delayed_builtin(delay_min_lag=1, delay_max_lag=3)
|
|
entity, _ = initialize_entity(entity, device, num_envs=2)
|
|
|
|
actuator = entity.actuators[0]
|
|
assert actuator.has_delay
|
|
|
|
# Try to set lags outside the valid range.
|
|
lags = torch.tensor([0, 10], device=device) # 0 < min_lag, 10 > max_lag
|
|
actuator.set_lags(lags)
|
|
|
|
# Lags should be clamped to [1, 3].
|
|
buffer = actuator._delay_buffer
|
|
assert buffer is not None
|
|
assert buffer.current_lags[0] == 1 # Clamped from 0
|
|
assert buffer.current_lags[1] == 3 # Clamped from 10
|
|
|
|
|
|
def test_delayed_actuator_set_lags_affects_delay(device):
|
|
"""Test that setting lags actually changes the delay behavior."""
|
|
# Use hold_prob=1.0 to prevent automatic lag resampling.
|
|
cfg = EntityCfg(
|
|
spec_fn=lambda: mujoco.MjSpec.from_string(ROBOT_XML),
|
|
articulation=EntityArticulationInfoCfg(
|
|
actuators=(
|
|
BuiltinPositionActuatorCfg(
|
|
target_names_expr=("joint.*",),
|
|
effort_limit=100.0,
|
|
stiffness=80.0,
|
|
damping=10.0,
|
|
delay_min_lag=0,
|
|
delay_max_lag=5,
|
|
delay_hold_prob=1.0, # Prevent automatic resampling
|
|
),
|
|
)
|
|
),
|
|
)
|
|
entity = Entity(cfg)
|
|
entity, sim = initialize_entity(entity, device, num_envs=1)
|
|
|
|
actuator = entity.actuators[0]
|
|
assert actuator.has_delay
|
|
|
|
# Set lag to 1.
|
|
actuator.set_lags(torch.tensor([1], device=device))
|
|
|
|
# Fill the buffer with known targets.
|
|
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 target in targets:
|
|
entity.set_joint_position_target(target)
|
|
entity.set_joint_velocity_target(torch.zeros(1, 2, device=device))
|
|
entity.set_joint_effort_target(torch.zeros(1, 2, device=device))
|
|
entity.write_data_to_sim()
|
|
|
|
# With lag=1, after 3 steps, ctrl should use target from step 1 (index 1).
|
|
ctrl = sim.data.ctrl[0]
|
|
assert torch.allclose(ctrl, targets[1][0], atol=1e-5)
|
|
|
|
|
|
def test_delayed_actuator_set_lags_overwritten_without_hold_prob(device):
|
|
"""Test that set_lags gets overwritten when delay_hold_prob < 1.0."""
|
|
# Use min_lag=max_lag=2 so resampling always produces 2.
|
|
cfg = EntityCfg(
|
|
spec_fn=lambda: mujoco.MjSpec.from_string(ROBOT_XML),
|
|
articulation=EntityArticulationInfoCfg(
|
|
actuators=(
|
|
BuiltinPositionActuatorCfg(
|
|
target_names_expr=("joint.*",),
|
|
effort_limit=100.0,
|
|
stiffness=80.0,
|
|
damping=10.0,
|
|
delay_min_lag=2,
|
|
delay_max_lag=2,
|
|
delay_hold_prob=0.0, # Always resample
|
|
),
|
|
)
|
|
),
|
|
)
|
|
entity = Entity(cfg)
|
|
entity, sim = initialize_entity(entity, device, num_envs=1)
|
|
|
|
actuator = entity.actuators[0]
|
|
assert actuator.has_delay
|
|
buffer = actuator._delay_buffer
|
|
assert buffer is not None
|
|
|
|
# Set lag to 2 (the only valid value, so set_lags won't clamp it).
|
|
actuator.set_lags(torch.tensor([2], device=device))
|
|
assert buffer.current_lags[0] == 2
|
|
|
|
# Now manually set _current_lags to 0 to simulate "we want 0".
|
|
# This bypasses clamping to test the resampling behavior.
|
|
buffer._current_lags[0] = 0
|
|
assert buffer.current_lags[0] == 0
|
|
|
|
# After compute, with hold_prob=0.0, it resamples to [min_lag, max_lag] = 2.
|
|
entity.set_joint_position_target(torch.zeros(1, 2, device=device))
|
|
entity.set_joint_velocity_target(torch.zeros(1, 2, device=device))
|
|
entity.set_joint_effort_target(torch.zeros(1, 2, device=device))
|
|
entity.write_data_to_sim()
|
|
|
|
# Lag should have been resampled back to 2.
|
|
assert buffer.current_lags[0] == 2
|