"""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 = """ """ ROBOT_XML_TWO_JOINTS = """ """ @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)