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
467 lines
15 KiB
Python
467 lines
15 KiB
Python
"""Tests for learned MLP actuator."""
|
|
|
|
import mujoco
|
|
import pytest
|
|
import torch
|
|
from conftest import get_test_device
|
|
|
|
from mjlab.actuator import LearnedMlpActuator, LearnedMlpActuatorCfg
|
|
from mjlab.entity import Entity, EntityArticulationInfoCfg, EntityCfg
|
|
from mjlab.sim.sim import Simulation, SimulationCfg
|
|
|
|
ROBOT_XML = """
|
|
<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="1.0"/>
|
|
<body name="link1" pos="0 0 0">
|
|
<joint name="joint1" type="hinge" axis="0 0 1" range="-3.14 3.14"/>
|
|
<geom name="link1_geom" type="box" size="0.1 0.1 0.1" mass="0.1"/>
|
|
</body>
|
|
</body>
|
|
</worldbody>
|
|
</mujoco>
|
|
"""
|
|
|
|
ROBOT_XML_TWO_JOINTS = """
|
|
<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="1.0"/>
|
|
<body name="link1" pos="0 0 0">
|
|
<joint name="joint1" type="hinge" axis="0 0 1" range="-3.14 3.14"/>
|
|
<geom name="link1_geom" type="box" size="0.1 0.1 0.1" mass="0.1"/>
|
|
<body name="link2" pos="0.2 0 0">
|
|
<joint name="joint2" type="hinge" axis="0 0 1" range="-3.14 3.14"/>
|
|
<geom name="link2_geom" type="box" size="0.1 0.1 0.1" mass="0.1"/>
|
|
</body>
|
|
</body>
|
|
</body>
|
|
</worldbody>
|
|
</mujoco>
|
|
"""
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def device():
|
|
return get_test_device()
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def identity_network_file(tmp_path_factory):
|
|
"""Create a simple identity network: output = sum of all inputs."""
|
|
|
|
class IdentityNet(torch.nn.Module):
|
|
def forward(self, x):
|
|
# Sum all inputs and return as single output per joint.
|
|
return x.sum(dim=1, keepdim=True)
|
|
|
|
net = IdentityNet()
|
|
tmp_dir = tmp_path_factory.mktemp("networks")
|
|
network_path = tmp_dir / "identity_net.pt"
|
|
torch.jit.script(net).save(str(network_path))
|
|
return str(network_path)
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def constant_network_file(tmp_path_factory):
|
|
"""Create a network that returns constant value regardless of input."""
|
|
|
|
class ConstantNet(torch.nn.Module):
|
|
def forward(self, x):
|
|
# Return constant 10.0 for each joint.
|
|
return torch.full((x.shape[0], 1), 10.0, device=x.device)
|
|
|
|
net = ConstantNet()
|
|
tmp_dir = tmp_path_factory.mktemp("networks")
|
|
network_path = tmp_dir / "constant_net.pt"
|
|
torch.jit.script(net).save(str(network_path))
|
|
return str(network_path)
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def subtract_network_file(tmp_path_factory):
|
|
"""Create a network that subtracts second half from first half of inputs.
|
|
|
|
For input [a, b, c, d], returns (a + b) - (c + d).
|
|
This is sensitive to input order, unlike sum-based networks.
|
|
"""
|
|
|
|
class SubtractNet(torch.nn.Module):
|
|
def forward(self, x):
|
|
# Split input in half and subtract second half from first half.
|
|
mid = x.shape[1] // 2
|
|
first_half = x[:, :mid].sum(dim=1, keepdim=True)
|
|
second_half = x[:, mid:].sum(dim=1, keepdim=True)
|
|
return first_half - second_half
|
|
|
|
net = SubtractNet()
|
|
tmp_dir = tmp_path_factory.mktemp("networks")
|
|
network_path = tmp_dir / "subtract_net.pt"
|
|
torch.jit.script(net).save(str(network_path))
|
|
return str(network_path)
|
|
|
|
|
|
def create_entity_with_actuator(actuator_cfg):
|
|
cfg = EntityCfg(
|
|
spec_fn=lambda: mujoco.MjSpec.from_string(ROBOT_XML),
|
|
articulation=EntityArticulationInfoCfg(actuators=(actuator_cfg,)),
|
|
)
|
|
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_learned_mlp_network_loads(device, identity_network_file):
|
|
"""Verify network loads from TorchScript file and initializes."""
|
|
entity = create_entity_with_actuator(
|
|
LearnedMlpActuatorCfg(
|
|
target_names_expr=("joint.*",),
|
|
network_file=identity_network_file,
|
|
pos_scale=1.0,
|
|
vel_scale=1.0,
|
|
torque_scale=1.0,
|
|
input_order="pos_vel",
|
|
history_length=1,
|
|
saturation_effort=100.0,
|
|
velocity_limit=30.0,
|
|
effort_limit=50.0,
|
|
)
|
|
)
|
|
|
|
entity, sim = initialize_entity(entity, device)
|
|
|
|
# Verify network was loaded.
|
|
actuator = entity.actuators[0]
|
|
assert isinstance(actuator, LearnedMlpActuator)
|
|
assert actuator.network is not None
|
|
assert actuator._pos_error_history is not None
|
|
assert actuator._vel_history is not None
|
|
|
|
|
|
def test_learned_mlp_history_indexing(device, identity_network_file):
|
|
"""Test that history_length uses consecutive timesteps."""
|
|
entity = create_entity_with_actuator(
|
|
LearnedMlpActuatorCfg(
|
|
target_names_expr=("joint.*",),
|
|
network_file=identity_network_file,
|
|
pos_scale=1.0,
|
|
vel_scale=0.0, # Zero out velocity contribution.
|
|
torque_scale=1.0,
|
|
input_order="pos_vel",
|
|
history_length=3, # Use current + 2 previous timesteps.
|
|
saturation_effort=100.0,
|
|
velocity_limit=30.0,
|
|
effort_limit=100.0,
|
|
)
|
|
)
|
|
|
|
entity, sim = initialize_entity(entity, device)
|
|
|
|
# Set up state at rest.
|
|
joint_pos = torch.tensor([[0.0]], device=device)
|
|
joint_vel = torch.tensor([[0.0]], device=device)
|
|
entity.write_joint_state_to_sim(joint_pos, joint_vel)
|
|
|
|
# Set position targets to create known position errors.
|
|
# Step 0: pos_error = 1.0
|
|
# Step 1: pos_error = 2.0
|
|
# Step 2: pos_error = 3.0
|
|
targets = [
|
|
torch.tensor([[1.0]], device=device),
|
|
torch.tensor([[2.0]], device=device),
|
|
torch.tensor([[3.0]], 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, 1, device=device))
|
|
entity.write_data_to_sim()
|
|
|
|
# After 3 steps with history_length=3:
|
|
# Network input: [current, -1 step, -2 steps] = [3.0, 2.0, 1.0]
|
|
# Identity network sums: 3.0 + 2.0 + 1.0 = 6.0
|
|
ctrl = sim.data.ctrl[0]
|
|
assert torch.allclose(ctrl, torch.tensor([6.0], device=device), atol=1e-4)
|
|
|
|
|
|
def test_learned_mlp_input_order_pos_vel(device, subtract_network_file):
|
|
"""Verify input_order='pos_vel' concatenates position then velocity."""
|
|
entity = create_entity_with_actuator(
|
|
LearnedMlpActuatorCfg(
|
|
target_names_expr=("joint.*",),
|
|
network_file=subtract_network_file,
|
|
pos_scale=1.0,
|
|
vel_scale=1.0,
|
|
torque_scale=1.0,
|
|
input_order="pos_vel",
|
|
history_length=1,
|
|
saturation_effort=100.0,
|
|
velocity_limit=30.0,
|
|
effort_limit=100.0,
|
|
)
|
|
)
|
|
|
|
entity, sim = initialize_entity(entity, device)
|
|
|
|
# Set state: pos_error = 2.0, vel = 3.0.
|
|
joint_pos = torch.tensor([[0.0]], device=device)
|
|
joint_vel = torch.tensor([[3.0]], device=device)
|
|
entity.write_joint_state_to_sim(joint_pos, joint_vel)
|
|
|
|
entity.set_joint_position_target(torch.tensor([[2.0]], device=device))
|
|
entity.set_joint_velocity_target(torch.zeros(1, 1, device=device))
|
|
entity.set_joint_effort_target(torch.zeros(1, 1, device=device))
|
|
entity.write_data_to_sim()
|
|
|
|
# Network input: [pos_error, vel] = [2.0, 3.0]
|
|
# Subtract network: 2.0 - 3.0 = -1.0
|
|
ctrl = sim.data.ctrl[0]
|
|
assert torch.allclose(ctrl, torch.tensor([-1.0], device=device), atol=1e-4)
|
|
|
|
|
|
def test_learned_mlp_input_order_vel_pos(device, subtract_network_file):
|
|
"""Verify input_order='vel_pos' concatenates velocity then position."""
|
|
entity = create_entity_with_actuator(
|
|
LearnedMlpActuatorCfg(
|
|
target_names_expr=("joint.*",),
|
|
network_file=subtract_network_file,
|
|
pos_scale=1.0,
|
|
vel_scale=1.0,
|
|
torque_scale=1.0,
|
|
input_order="vel_pos",
|
|
history_length=1,
|
|
saturation_effort=100.0,
|
|
velocity_limit=30.0,
|
|
effort_limit=100.0,
|
|
)
|
|
)
|
|
|
|
entity, sim = initialize_entity(entity, device)
|
|
|
|
# Set state: pos_error = 2.0, vel = 3.0.
|
|
joint_pos = torch.tensor([[0.0]], device=device)
|
|
joint_vel = torch.tensor([[3.0]], device=device)
|
|
entity.write_joint_state_to_sim(joint_pos, joint_vel)
|
|
|
|
entity.set_joint_position_target(torch.tensor([[2.0]], device=device))
|
|
entity.set_joint_velocity_target(torch.zeros(1, 1, device=device))
|
|
entity.set_joint_effort_target(torch.zeros(1, 1, device=device))
|
|
entity.write_data_to_sim()
|
|
|
|
# Network input: [vel, pos_error] = [3.0, 2.0]
|
|
# Subtract network: 3.0 - 2.0 = 1.0 (opposite of pos_vel!)
|
|
ctrl = sim.data.ctrl[0]
|
|
assert torch.allclose(ctrl, torch.tensor([1.0], device=device), atol=1e-4)
|
|
|
|
|
|
def test_learned_mlp_scaling_applied(device, identity_network_file):
|
|
"""Test pos_scale, vel_scale, torque_scale are applied correctly."""
|
|
entity = create_entity_with_actuator(
|
|
LearnedMlpActuatorCfg(
|
|
target_names_expr=("joint.*",),
|
|
network_file=identity_network_file,
|
|
pos_scale=2.0,
|
|
vel_scale=3.0,
|
|
torque_scale=4.0,
|
|
input_order="pos_vel",
|
|
history_length=1,
|
|
saturation_effort=100.0,
|
|
velocity_limit=30.0,
|
|
effort_limit=100.0,
|
|
)
|
|
)
|
|
|
|
entity, sim = initialize_entity(entity, device)
|
|
|
|
# Set state: pos_error = 1.0, vel = 1.0.
|
|
joint_pos = torch.tensor([[0.0]], device=device)
|
|
joint_vel = torch.tensor([[1.0]], device=device)
|
|
entity.write_joint_state_to_sim(joint_pos, joint_vel)
|
|
|
|
entity.set_joint_position_target(torch.tensor([[1.0]], device=device))
|
|
entity.set_joint_velocity_target(torch.zeros(1, 1, device=device))
|
|
entity.set_joint_effort_target(torch.zeros(1, 1, device=device))
|
|
entity.write_data_to_sim()
|
|
|
|
# Network input: [1.0 * 2.0, 1.0 * 3.0] = [2.0, 3.0]
|
|
# Identity network sums: 2.0 + 3.0 = 5.0
|
|
# Output scaled: 5.0 * 4.0 = 20.0
|
|
ctrl = sim.data.ctrl[0]
|
|
assert torch.allclose(ctrl, torch.tensor([20.0], device=device), atol=1e-4)
|
|
|
|
|
|
def test_learned_mlp_reset_clears_history(device, identity_network_file):
|
|
"""Test reset zeroes history buffers for specified environments."""
|
|
entity = create_entity_with_actuator(
|
|
LearnedMlpActuatorCfg(
|
|
target_names_expr=("joint.*",),
|
|
network_file=identity_network_file,
|
|
pos_scale=1.0,
|
|
vel_scale=1.0,
|
|
torque_scale=1.0,
|
|
input_order="pos_vel",
|
|
history_length=2,
|
|
saturation_effort=100.0,
|
|
velocity_limit=30.0,
|
|
effort_limit=100.0,
|
|
)
|
|
)
|
|
|
|
entity, sim = initialize_entity(entity, device, num_envs=2)
|
|
|
|
# Set state and targets to fill buffers.
|
|
joint_pos = torch.tensor([[0.0], [0.0]], device=device)
|
|
joint_vel = torch.tensor([[1.0], [1.0]], device=device)
|
|
entity.write_joint_state_to_sim(joint_pos, joint_vel)
|
|
|
|
entity.set_joint_position_target(torch.tensor([[5.0], [5.0]], device=device))
|
|
entity.set_joint_velocity_target(torch.zeros(2, 1, device=device))
|
|
entity.set_joint_effort_target(torch.zeros(2, 1, device=device))
|
|
entity.write_data_to_sim()
|
|
|
|
# Reset only env 0.
|
|
entity.reset(torch.tensor([0], device=device))
|
|
|
|
# Check history buffers.
|
|
actuator = entity.actuators[0]
|
|
assert isinstance(actuator, LearnedMlpActuator)
|
|
assert actuator._pos_error_history is not None
|
|
assert actuator._vel_history is not None
|
|
assert actuator._pos_error_history.current_length[0] == 0
|
|
assert actuator._pos_error_history.current_length[1] > 0
|
|
assert actuator._vel_history.current_length[0] == 0
|
|
assert actuator._vel_history.current_length[1] > 0
|
|
|
|
|
|
def test_learned_mlp_inherits_dc_motor_limits(device, constant_network_file):
|
|
"""Test that DC motor saturation_effort and velocity limits apply."""
|
|
saturation_effort = 20.0
|
|
velocity_limit = 30.0
|
|
|
|
entity = create_entity_with_actuator(
|
|
LearnedMlpActuatorCfg(
|
|
target_names_expr=("joint.*",),
|
|
network_file=constant_network_file, # Always outputs 10.0.
|
|
pos_scale=1.0,
|
|
vel_scale=1.0,
|
|
torque_scale=100.0, # Scale to 1000.0 to exceed limits.
|
|
input_order="pos_vel",
|
|
history_length=1,
|
|
saturation_effort=saturation_effort,
|
|
velocity_limit=velocity_limit,
|
|
effort_limit=saturation_effort, # Set to saturation to not constrain.
|
|
)
|
|
)
|
|
|
|
entity, sim = initialize_entity(entity, device)
|
|
|
|
# At zero velocity, should be clipped to saturation_effort.
|
|
joint_pos = torch.tensor([[0.0]], device=device)
|
|
joint_vel = torch.tensor([[0.0]], device=device)
|
|
entity.write_joint_state_to_sim(joint_pos, joint_vel)
|
|
|
|
entity.set_joint_position_target(torch.tensor([[1.0]], device=device))
|
|
entity.set_joint_velocity_target(torch.zeros(1, 1, device=device))
|
|
entity.set_joint_effort_target(torch.zeros(1, 1, device=device))
|
|
entity.write_data_to_sim()
|
|
|
|
# Network outputs 10.0, scaled by 100.0 = 1000.0.
|
|
# Should be clipped to saturation_effort = 20.0.
|
|
ctrl = sim.data.ctrl[0]
|
|
assert torch.allclose(ctrl, torch.tensor([saturation_effort], device=device))
|
|
|
|
|
|
def test_learned_mlp_dc_motor_zero_torque_at_max_velocity(
|
|
device, constant_network_file
|
|
):
|
|
"""Test that DC motor produces zero torque at max velocity."""
|
|
saturation_effort = 20.0
|
|
velocity_limit = 30.0
|
|
|
|
entity = create_entity_with_actuator(
|
|
LearnedMlpActuatorCfg(
|
|
target_names_expr=("joint.*",),
|
|
network_file=constant_network_file, # Always outputs 10.0.
|
|
pos_scale=1.0,
|
|
vel_scale=1.0,
|
|
torque_scale=100.0, # Scale to 1000.0.
|
|
input_order="pos_vel",
|
|
history_length=1,
|
|
saturation_effort=saturation_effort,
|
|
velocity_limit=velocity_limit,
|
|
effort_limit=saturation_effort, # Set to saturation to not constrain.
|
|
)
|
|
)
|
|
|
|
entity, sim = initialize_entity(entity, device)
|
|
|
|
# At max velocity, should produce zero torque.
|
|
joint_pos = torch.tensor([[0.0]], device=device)
|
|
joint_vel = torch.tensor([[velocity_limit]], device=device)
|
|
entity.write_joint_state_to_sim(joint_pos, joint_vel)
|
|
|
|
entity.set_joint_position_target(torch.tensor([[1.0]], device=device))
|
|
entity.set_joint_velocity_target(torch.zeros(1, 1, device=device))
|
|
entity.set_joint_effort_target(torch.zeros(1, 1, device=device))
|
|
entity.write_data_to_sim()
|
|
|
|
# At max velocity, DC motor should clip to zero.
|
|
ctrl = sim.data.ctrl[0]
|
|
assert torch.allclose(ctrl, torch.tensor([0.0], device=device), atol=1e-5)
|
|
|
|
|
|
def test_learned_mlp_multi_joint_reshaping(device, identity_network_file):
|
|
"""Test that multi-joint reshaping works correctly."""
|
|
cfg = EntityCfg(
|
|
spec_fn=lambda: mujoco.MjSpec.from_string(ROBOT_XML_TWO_JOINTS),
|
|
articulation=EntityArticulationInfoCfg(
|
|
actuators=(
|
|
LearnedMlpActuatorCfg(
|
|
target_names_expr=("joint.*",),
|
|
network_file=identity_network_file,
|
|
pos_scale=1.0,
|
|
vel_scale=1.0,
|
|
torque_scale=1.0,
|
|
input_order="pos_vel",
|
|
history_length=1,
|
|
saturation_effort=100.0,
|
|
velocity_limit=30.0,
|
|
effort_limit=100.0,
|
|
),
|
|
)
|
|
),
|
|
)
|
|
|
|
entity = Entity(cfg)
|
|
model = entity.compile()
|
|
sim_cfg = SimulationCfg()
|
|
sim = Simulation(num_envs=1, cfg=sim_cfg, model=model, device=device)
|
|
entity.initialize(model, sim.model, sim.data, device)
|
|
|
|
# Set different states for each joint.
|
|
# Joint 1: pos_error = 1.0, vel = 2.0
|
|
# Joint 2: pos_error = 3.0, vel = 4.0
|
|
joint_pos = torch.tensor([[0.0, 0.0]], device=device)
|
|
joint_vel = torch.tensor([[2.0, 4.0]], device=device)
|
|
entity.write_joint_state_to_sim(joint_pos, joint_vel)
|
|
|
|
entity.set_joint_position_target(torch.tensor([[1.0, 3.0]], 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()
|
|
|
|
# Network input for joint 1: [1.0, 2.0] → sum = 3.0
|
|
# Network input for joint 2: [3.0, 4.0] → sum = 7.0
|
|
ctrl = sim.data.ctrl[0]
|
|
assert torch.allclose(ctrl, torch.tensor([3.0, 7.0], device=device), atol=1e-4)
|