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 reward manager functionality."""
|
|
|
|
from unittest.mock import Mock
|
|
|
|
import mujoco
|
|
import pytest
|
|
import torch
|
|
from conftest import get_test_device
|
|
|
|
from mjlab.actuator import BuiltinPositionActuatorCfg
|
|
from mjlab.entity import Entity, EntityArticulationInfoCfg, EntityCfg
|
|
from mjlab.envs.mdp.rewards import electrical_power_cost, joint_torques_l2
|
|
from mjlab.managers.reward_manager import RewardManager, RewardTermCfg
|
|
from mjlab.managers.scene_entity_config import SceneEntityCfg
|
|
from mjlab.sim.sim import Simulation, SimulationCfg
|
|
|
|
PARTIALLY_ACTUATED_ROBOT_XML = """
|
|
<mujoco>
|
|
<worldbody>
|
|
<body name="base" pos="0 0 0.5">
|
|
<geom name="base_geom" type="cylinder" size="0.1 0.05" mass="1.0"/>
|
|
<body name="link1" pos="0 0 0.1">
|
|
<joint name="actuated_joint1" type="hinge" axis="0 0 1" range="-3.14 3.14"/>
|
|
<geom name="link1_geom" type="box" size="0.05 0.05 0.2" mass="0.5"/>
|
|
<body name="link2" pos="0 0 0.4">
|
|
<joint name="passive_joint" type="hinge" axis="0 1 0" range="-1.57 1.57"/>
|
|
<geom name="link2_geom" type="box" size="0.05 0.05 0.15" mass="0.3"/>
|
|
<body name="link3" pos="0 0 0.3">
|
|
<joint name="actuated_joint2" type="hinge" axis="1 0 0" range="-1.57 1.57"/>
|
|
<geom name="link3_geom" type="box" size="0.05 0.05 0.1" mass="0.2"/>
|
|
</body>
|
|
</body>
|
|
</body>
|
|
</body>
|
|
</worldbody>
|
|
</mujoco>
|
|
"""
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def device():
|
|
"""Test device fixture."""
|
|
return get_test_device()
|
|
|
|
|
|
class SimpleTestReward:
|
|
"""A simple class-based reward for testing that tracks state."""
|
|
|
|
def __init__(self, cfg: RewardTermCfg, env):
|
|
self.num_envs = env.num_envs
|
|
self.device = env.device
|
|
self.current_air_time = torch.zeros((self.num_envs, 1), device=self.device)
|
|
|
|
def __call__(self, env, **kwargs):
|
|
self.current_air_time += 0.01
|
|
return torch.ones(env.num_envs, device=env.device)
|
|
|
|
def reset(self, env_ids: torch.Tensor | None = None, env=None):
|
|
if env_ids is not None and len(env_ids) > 0:
|
|
self.current_air_time[env_ids] = 0
|
|
|
|
|
|
class StatelessReward:
|
|
"""A stateless class-based reward without reset method."""
|
|
|
|
def __init__(self, cfg: RewardTermCfg, env):
|
|
pass
|
|
|
|
def __call__(self, env, **kwargs):
|
|
return torch.ones(env.num_envs)
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_env():
|
|
"""Create a mock environment for testing."""
|
|
env = Mock()
|
|
env.num_envs = 4
|
|
env.device = "cpu"
|
|
env.step_dt = 0.01
|
|
env.max_episode_length_s = 10.0
|
|
robot = Mock()
|
|
env.scene = {"robot": robot}
|
|
env.command_manager.get_command = Mock(
|
|
return_value=torch.tensor([[1.0, 0.0, 0.0]] * 4)
|
|
)
|
|
return env
|
|
|
|
|
|
@pytest.fixture
|
|
def class_reward_config():
|
|
"""Config with a class-based reward."""
|
|
return {
|
|
"term": RewardTermCfg(
|
|
func=SimpleTestReward,
|
|
weight=1.0,
|
|
params={},
|
|
)
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def function_reward_config():
|
|
"""Config with a function-based reward."""
|
|
return {
|
|
"term": RewardTermCfg(
|
|
func=lambda env: torch.ones(env.num_envs),
|
|
weight=1.0,
|
|
params={},
|
|
)
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def stateless_reward_config():
|
|
"""Config with a stateless class-based reward."""
|
|
return {
|
|
"term": RewardTermCfg(
|
|
func=StatelessReward,
|
|
weight=1.0,
|
|
params={},
|
|
)
|
|
}
|
|
|
|
|
|
def test_class_based_reward_reset(mock_env, class_reward_config):
|
|
"""Test that class-based reward terms are tracked and have reset called."""
|
|
manager = RewardManager(class_reward_config, mock_env)
|
|
term = manager._class_term_cfgs[0].func
|
|
|
|
for _ in range(10):
|
|
manager.compute(dt=0.01)
|
|
assert (term.current_air_time > 0).all()
|
|
|
|
manager.reset(env_ids=torch.tensor([0, 2]))
|
|
|
|
# Check that only specified envs were reset.
|
|
assert term.current_air_time[0, 0] == 0
|
|
assert term.current_air_time[1, 0] > 0
|
|
assert term.current_air_time[2, 0] == 0
|
|
assert term.current_air_time[3, 0] > 0
|
|
|
|
|
|
def test_function_based_reward_not_tracked(mock_env, function_reward_config):
|
|
"""Test that function-based reward terms are not tracked as class terms."""
|
|
manager = RewardManager(function_reward_config, mock_env)
|
|
assert len(manager._class_term_cfgs) == 0
|
|
|
|
|
|
def test_stateless_class_reward_no_reset(mock_env, stateless_reward_config):
|
|
"""Test that stateless class-based rewards without reset don't break reset."""
|
|
manager = RewardManager(stateless_reward_config, mock_env)
|
|
|
|
# Stateless rewards without reset method should not be tracked.
|
|
assert len(manager._class_term_cfgs) == 0
|
|
|
|
# Reset should work without errors.
|
|
manager.reset(env_ids=torch.tensor([0, 2]))
|
|
|
|
|
|
def test_electrical_power_cost_partially_actuated(device):
|
|
"""Test electrical_power_cost on robots with passive joints.
|
|
|
|
This test verifies that:
|
|
1. The reward correctly identifies and uses only actuated joints (ignoring passive joints)
|
|
2. Power cost is computed as sum of max(0, force*velocity) for each actuator,
|
|
meaning only positive work is penalized (not regenerative braking)
|
|
"""
|
|
entity_cfg = EntityCfg(
|
|
spec_fn=lambda: mujoco.MjSpec.from_string(PARTIALLY_ACTUATED_ROBOT_XML),
|
|
articulation=EntityArticulationInfoCfg(
|
|
actuators=(
|
|
BuiltinPositionActuatorCfg(
|
|
target_names_expr=("actuated_joint1", "actuated_joint2"),
|
|
effort_limit=10.0,
|
|
stiffness=100.0,
|
|
damping=10.0,
|
|
),
|
|
)
|
|
),
|
|
)
|
|
entity = Entity(entity_cfg)
|
|
model = entity.compile()
|
|
|
|
num_envs = 2
|
|
sim_cfg = SimulationCfg()
|
|
sim = Simulation(num_envs=num_envs, cfg=sim_cfg, model=model, device=device)
|
|
entity.initialize(model, sim.model, sim.data, device)
|
|
|
|
env = Mock()
|
|
env.num_envs = num_envs
|
|
env.device = device
|
|
env.scene = {"robot": entity}
|
|
|
|
asset_cfg = SceneEntityCfg(name="robot", joint_names=("actuated.*",))
|
|
|
|
reward_cfg = RewardTermCfg(
|
|
func=electrical_power_cost, weight=1.0, params={"asset_cfg": asset_cfg}
|
|
)
|
|
|
|
reward = electrical_power_cost(reward_cfg, env)
|
|
|
|
assert len(reward._joint_ids) == 2
|
|
|
|
# Test case 1: All forces and velocities aligned (all positive work).
|
|
# actuated_joint1 (dof 0), actuated_joint2 (dof 2).
|
|
# Note: dof 1 is the passive joint, not used in power calculation.
|
|
sim.data.qfrc_actuator[:] = 0.0
|
|
sim.data.qfrc_actuator[:, 0] = torch.tensor([2.0, 1.0], device=device)
|
|
sim.data.qfrc_actuator[:, 2] = torch.tensor([3.0, 4.0], device=device)
|
|
sim.data.qvel[:] = 0.0
|
|
sim.data.qvel[:, 0] = torch.tensor([1.0, 2.0], device=device)
|
|
sim.data.qvel[:, 2] = torch.tensor([2.0, 1.0], device=device)
|
|
|
|
power_cost = reward(env, asset_cfg)
|
|
|
|
# Expected: env0 = 2.0*1.0 + 3.0*2.0 = 8.0, env1 = 1.0*2.0 + 4.0*1.0 = 6.0.
|
|
expected = torch.tensor([8.0, 6.0], device=device)
|
|
assert torch.allclose(power_cost, expected)
|
|
|
|
# Test case 2: Some negative forces (regenerative braking, not penalized).
|
|
sim.data.qfrc_actuator[:] = 0.0
|
|
sim.data.qfrc_actuator[:, 0] = torch.tensor([-2.0, 1.0], device=device)
|
|
sim.data.qfrc_actuator[:, 2] = torch.tensor([3.0, -4.0], device=device)
|
|
|
|
power_cost = reward(env, asset_cfg)
|
|
|
|
# Expected: env0 = max(0,-2.0*1.0) + max(0,3.0*2.0) = 0 + 6.0 = 6.0.
|
|
# env1 = max(0,1.0*2.0) + max(0,-4.0*1.0) = 2.0 + 0 = 2.0.
|
|
expected = torch.tensor([6.0, 2.0], device=device)
|
|
assert torch.allclose(power_cost, expected)
|
|
|
|
|
|
def test_reward_manager_handles_nan_values(mock_env):
|
|
"""Test that RewardManager converts NaN/Inf reward values to zero."""
|
|
|
|
def nan_reward(env):
|
|
"""Reward function that returns NaN for some environments."""
|
|
r = torch.ones(env.num_envs, device=env.device)
|
|
r[1] = float("nan")
|
|
r[3] = float("inf")
|
|
return r
|
|
|
|
cfg = {"nan_term": RewardTermCfg(func=nan_reward, weight=1.0, params={})}
|
|
manager = RewardManager(cfg, mock_env)
|
|
|
|
rewards = manager.compute(dt=0.01)
|
|
|
|
# NaN and Inf should be converted to 0.
|
|
assert not torch.isnan(rewards).any(), "Reward buffer contains NaN"
|
|
assert not torch.isinf(rewards).any(), "Reward buffer contains Inf"
|
|
assert rewards[0] == pytest.approx(0.01)
|
|
assert rewards[1] == 0.0
|
|
assert rewards[2] == pytest.approx(0.01)
|
|
assert rewards[3] == 0.0
|
|
|
|
|
|
def test_reward_manager_handles_neginf_values(mock_env):
|
|
"""Test that RewardManager converts negative infinity to zero."""
|
|
|
|
def neginf_reward(env):
|
|
r = torch.ones(env.num_envs, device=env.device)
|
|
r[2] = float("-inf")
|
|
return r
|
|
|
|
cfg = {"neginf_term": RewardTermCfg(func=neginf_reward, weight=1.0, params={})}
|
|
manager = RewardManager(cfg, mock_env)
|
|
|
|
rewards = manager.compute(dt=0.01)
|
|
|
|
assert not torch.isinf(rewards).any()
|
|
assert rewards[2] == 0.0
|
|
|
|
|
|
def test_reward_scaling_enabled(mock_env):
|
|
"""Test that rewards are scaled by dt when scale_by_dt=True (default)."""
|
|
|
|
def constant_reward(env):
|
|
return torch.ones(env.num_envs, device=env.device)
|
|
|
|
cfg = {"term": RewardTermCfg(func=constant_reward, weight=2.0, params={})}
|
|
manager = RewardManager(cfg, mock_env, scale_by_dt=True)
|
|
|
|
dt = 0.02
|
|
rewards = manager.compute(dt=dt)
|
|
|
|
# With scaling: reward = raw_value (1.0) * weight (2.0) * dt (0.02) = 0.04
|
|
expected = 1.0 * 2.0 * dt
|
|
assert torch.allclose(rewards, torch.full((4,), expected))
|
|
|
|
# _step_reward should be unscaled (raw_value * weight)
|
|
step_reward = manager._step_reward[:, 0]
|
|
expected_step = 1.0 * 2.0
|
|
assert torch.allclose(step_reward, torch.full((4,), expected_step))
|
|
|
|
|
|
def test_reward_scaling_disabled(mock_env):
|
|
"""Test that rewards are not scaled by dt when scale_by_dt=False."""
|
|
|
|
def constant_reward(env):
|
|
return torch.ones(env.num_envs, device=env.device)
|
|
|
|
cfg = {"term": RewardTermCfg(func=constant_reward, weight=2.0, params={})}
|
|
manager = RewardManager(cfg, mock_env, scale_by_dt=False)
|
|
|
|
dt = 0.02
|
|
rewards = manager.compute(dt=dt)
|
|
|
|
# Without scaling: reward = raw_value (1.0) * weight (2.0) = 2.0
|
|
expected = 1.0 * 2.0
|
|
assert torch.allclose(rewards, torch.full((4,), expected))
|
|
|
|
# _step_reward should still be unscaled (same as reward when not scaling)
|
|
step_reward = manager._step_reward[:, 0]
|
|
assert torch.allclose(step_reward, torch.full((4,), expected))
|
|
|
|
|
|
def test_reward_scaling_default_is_enabled(mock_env):
|
|
"""Test that scale_by_dt defaults to True for backward compatibility."""
|
|
|
|
def constant_reward(env):
|
|
return torch.ones(env.num_envs, device=env.device)
|
|
|
|
cfg = {"term": RewardTermCfg(func=constant_reward, weight=1.0, params={})}
|
|
# Don't pass scale_by_dt - should default to True
|
|
manager = RewardManager(cfg, mock_env)
|
|
|
|
dt = 0.01
|
|
rewards = manager.compute(dt=dt)
|
|
|
|
# Default (scaling enabled): reward = 1.0 * 1.0 * 0.01 = 0.01
|
|
assert torch.allclose(rewards, torch.full((4,), 0.01))
|
|
|
|
|
|
def test_reward_shape_validation_rejects_bad_compute_output(mock_env):
|
|
"""Reward terms must return one scalar per environment."""
|
|
cfg = {"bad": RewardTermCfg(func=lambda env: torch.ones(env.num_envs, 1), weight=1.0)}
|
|
manager = RewardManager(cfg, mock_env)
|
|
with pytest.raises(ValueError, match="RewardManager term 'bad'.*expected \\(4,\\)"):
|
|
manager.compute(dt=0.01)
|
|
|
|
|
|
def test_joint_torques_l2_with_actuator_ids(mock_env):
|
|
"""Test that joint_torques_l2 only penalizes specified actuators."""
|
|
mock_env.scene["robot"].data.actuator_force = torch.tensor([[1.0, 2.0, 3.0, 4.0]] * 4)
|
|
|
|
asset_cfg = SceneEntityCfg(name="robot", actuator_ids=[0, 2])
|
|
result = joint_torques_l2(mock_env, asset_cfg)
|
|
|
|
# Only actuators 0 and 2: 1^2 + 3^2 = 10.0
|
|
expected = torch.full((4,), 10.0)
|
|
assert torch.allclose(result, expected)
|
|
|
|
|
|
def test_joint_torques_l2_all_actuators(mock_env):
|
|
"""Test that joint_torques_l2 uses all actuators by default."""
|
|
mock_env.scene["robot"].data.actuator_force = torch.tensor([[1.0, 2.0, 3.0]] * 4)
|
|
|
|
asset_cfg = SceneEntityCfg(name="robot")
|
|
result = joint_torques_l2(mock_env, asset_cfg)
|
|
|
|
# All actuators: 1^2 + 2^2 + 3^2 = 14.0
|
|
expected = torch.full((4,), 14.0)
|
|
assert torch.allclose(result, expected)
|