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
305 lines
9.3 KiB
Python
305 lines
9.3 KiB
Python
"""Tests for encoder bias simulation.
|
|
|
|
We simulate this by storing `encoder_bias` per joint in EntityData:
|
|
- Observations: policy sees (true_position + bias)
|
|
- Actions: position commands are converted via (command - bias) before simulation
|
|
|
|
This ensures joint limits apply to the true physical position, not biased readings.
|
|
"""
|
|
|
|
from collections.abc import Callable
|
|
from functools import partial
|
|
|
|
import mujoco
|
|
import pytest
|
|
import torch
|
|
from conftest import get_test_device
|
|
|
|
from mjlab.actuator import BuiltinPositionActuatorCfg
|
|
from mjlab.entity import EntityArticulationInfoCfg, EntityCfg
|
|
from mjlab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg, mdp
|
|
from mjlab.managers.event_manager import EventTermCfg
|
|
from mjlab.managers.observation_manager import ObservationGroupCfg, ObservationTermCfg
|
|
from mjlab.managers.scene_entity_config import SceneEntityCfg
|
|
from mjlab.scene import SceneCfg
|
|
from mjlab.sim import MujocoCfg, SimulationCfg
|
|
from mjlab.terrains import TerrainEntityCfg
|
|
|
|
# =============================================================================
|
|
# Test fixtures and helpers
|
|
# =============================================================================
|
|
|
|
SLIDING_MASS_XML = """
|
|
<mujoco>
|
|
<option timestep="0.002"/>
|
|
<worldbody>
|
|
<body name="mass" pos="0 0 0">
|
|
<joint name="slide" type="slide" axis="1 0 0" range="-1 1" limited="true"/>
|
|
<geom name="mass_geom" type="sphere" size="0.1" mass="1.0"/>
|
|
</body>
|
|
</worldbody>
|
|
<sensor>
|
|
<jointpos name="slide_pos" joint="slide"/>
|
|
</sensor>
|
|
</mujoco>
|
|
"""
|
|
|
|
|
|
def _make_robot_cfg():
|
|
return EntityCfg(
|
|
spec_fn=lambda: mujoco.MjSpec.from_string(SLIDING_MASS_XML),
|
|
articulation=EntityArticulationInfoCfg(
|
|
actuators=(
|
|
BuiltinPositionActuatorCfg(
|
|
target_names_expr=(".*",), stiffness=1000.0, damping=100.0
|
|
),
|
|
)
|
|
),
|
|
)
|
|
|
|
|
|
def _make_env_cfg(
|
|
obs_func: Callable | None = None,
|
|
num_envs: int = 2,
|
|
events: dict | None = None,
|
|
) -> ManagerBasedRlEnvCfg:
|
|
if obs_func is None:
|
|
obs_func = partial(mdp.joint_pos_rel, biased=True)
|
|
return ManagerBasedRlEnvCfg(
|
|
scene=SceneCfg(
|
|
terrain=TerrainEntityCfg(terrain_type="plane"),
|
|
num_envs=num_envs,
|
|
extent=1.0,
|
|
entities={"robot": _make_robot_cfg()},
|
|
),
|
|
observations={
|
|
"actor": ObservationGroupCfg(
|
|
terms={"obs": ObservationTermCfg(func=obs_func)},
|
|
),
|
|
},
|
|
actions={
|
|
"joint_pos": mdp.JointPositionActionCfg(
|
|
entity_name="robot", actuator_names=(".*",), scale=1.0
|
|
)
|
|
},
|
|
events=events or {},
|
|
sim=SimulationCfg(mujoco=MujocoCfg(timestep=0.002, iterations=1)),
|
|
decimation=1,
|
|
episode_length_s=10.0,
|
|
)
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def device():
|
|
return get_test_device()
|
|
|
|
|
|
# =============================================================================
|
|
# EntityData.encoder_bias field tests
|
|
# =============================================================================
|
|
|
|
|
|
def test_encoder_bias_initialized_to_zero(device):
|
|
"""encoder_bias should be zeros with shape (num_envs, num_joints)."""
|
|
env = ManagerBasedRlEnv(cfg=_make_env_cfg(num_envs=4), device=device)
|
|
env.reset()
|
|
|
|
robot = env.scene["robot"]
|
|
assert robot.data.encoder_bias.shape == (4, 1)
|
|
assert (robot.data.encoder_bias == 0).all()
|
|
|
|
env.close()
|
|
|
|
|
|
def test_encoder_bias_can_be_set_per_env(device):
|
|
"""encoder_bias can be set differently per environment."""
|
|
env = ManagerBasedRlEnv(cfg=_make_env_cfg(), device=device)
|
|
env.reset()
|
|
|
|
robot = env.scene["robot"]
|
|
robot.data.encoder_bias[0, 0] = 0.1
|
|
robot.data.encoder_bias[1, 0] = -0.2
|
|
|
|
assert robot.data.encoder_bias[0, 0].item() == pytest.approx(0.1)
|
|
assert robot.data.encoder_bias[1, 0].item() == pytest.approx(-0.2)
|
|
|
|
env.close()
|
|
|
|
|
|
def test_joint_pos_biased_equals_joint_pos_plus_bias(device):
|
|
"""joint_pos_biased property should return joint_pos + encoder_bias."""
|
|
env = ManagerBasedRlEnv(cfg=_make_env_cfg(), device=device)
|
|
env.reset()
|
|
|
|
robot = env.scene["robot"]
|
|
robot.data.encoder_bias[:, 0] = 0.25
|
|
|
|
torch.testing.assert_close(
|
|
robot.data.joint_pos_biased, robot.data.joint_pos + robot.data.encoder_bias
|
|
)
|
|
|
|
env.close()
|
|
|
|
|
|
# =============================================================================
|
|
# Observation tests
|
|
# =============================================================================
|
|
|
|
|
|
def test_joint_pos_rel_includes_encoder_bias(device):
|
|
"""joint_pos_rel observation should include encoder_bias."""
|
|
env = ManagerBasedRlEnv(cfg=_make_env_cfg(), device=device)
|
|
env.reset()
|
|
|
|
robot = env.scene["robot"]
|
|
bias = 0.5
|
|
|
|
obs_before = env.observation_manager.compute()["actor"]
|
|
assert isinstance(obs_before, torch.Tensor)
|
|
obs_before = obs_before.clone()
|
|
robot.data.encoder_bias[:, 0] = bias
|
|
env.observation_manager._obs_buffer = None # Invalidate cache.
|
|
obs_after = env.observation_manager.compute()["actor"]
|
|
|
|
# Observation should increase by bias amount.
|
|
torch.testing.assert_close(obs_after, obs_before + bias, atol=1e-5, rtol=0)
|
|
|
|
env.close()
|
|
|
|
|
|
def test_joint_vel_rel_ignores_encoder_bias(device):
|
|
"""joint_vel_rel should NOT include encoder_bias (bias is constant, d/dt = 0)."""
|
|
env = ManagerBasedRlEnv(cfg=_make_env_cfg(obs_func=mdp.joint_vel_rel), device=device)
|
|
env.reset()
|
|
|
|
robot = env.scene["robot"]
|
|
|
|
obs_before = env.observation_manager.compute()["actor"]
|
|
assert isinstance(obs_before, torch.Tensor)
|
|
obs_before = obs_before.clone()
|
|
robot.data.encoder_bias[:, 0] = 0.5
|
|
env.observation_manager._obs_buffer = None
|
|
obs_after = env.observation_manager.compute()["actor"]
|
|
|
|
torch.testing.assert_close(obs_before, obs_after, atol=1e-6, rtol=0)
|
|
|
|
env.close()
|
|
|
|
|
|
# =============================================================================
|
|
# Action tests
|
|
# =============================================================================
|
|
|
|
|
|
def test_position_action_subtracts_encoder_bias(device):
|
|
"""JointPositionAction should subtract encoder_bias from commands."""
|
|
env = ManagerBasedRlEnv(cfg=_make_env_cfg(), device=device)
|
|
env.reset()
|
|
|
|
robot = env.scene["robot"]
|
|
robot.data.encoder_bias[0, 0] = 0.0
|
|
robot.data.encoder_bias[1, 0] = 0.3
|
|
|
|
# Same command to both envs.
|
|
env.step(torch.tensor([[0.5], [0.5]], device=device))
|
|
|
|
# Actual targets should differ by bias.
|
|
assert robot.data.joint_pos_target[0, 0].item() == pytest.approx(0.5, abs=1e-5)
|
|
assert robot.data.joint_pos_target[1, 0].item() == pytest.approx(0.2, abs=1e-5)
|
|
|
|
env.close()
|
|
|
|
|
|
# =============================================================================
|
|
# End-to-end consistency test
|
|
# =============================================================================
|
|
|
|
|
|
def test_bias_compensation_produces_identical_physical_behavior(device):
|
|
"""Envs with different biases should have identical physics when compensated.
|
|
|
|
If env0 has bias=0 and env1 has bias=0.3, and we want both to reach physical
|
|
position 0.4, we command:
|
|
- env0: 0.4 (in encoder frame)
|
|
- env1: 0.7 (in encoder frame, which becomes 0.4 after bias subtraction)
|
|
|
|
Both should reach the same physical position. Their observations should differ
|
|
by the bias amount.
|
|
"""
|
|
env = ManagerBasedRlEnv(cfg=_make_env_cfg(), device=device)
|
|
env.reset()
|
|
|
|
robot = env.scene["robot"]
|
|
bias_env0, bias_env1 = 0.0, 0.3
|
|
robot.data.encoder_bias[0, 0] = bias_env0
|
|
robot.data.encoder_bias[1, 0] = bias_env1
|
|
|
|
# Command same physical target via encoder frame.
|
|
target_physical = 0.4
|
|
action = torch.tensor(
|
|
[[target_physical + bias_env0], [target_physical + bias_env1]], device=device
|
|
)
|
|
|
|
for _ in range(100):
|
|
env.step(action)
|
|
|
|
# Physical positions should match.
|
|
pos_env0 = robot.data.joint_pos[0, 0].item()
|
|
pos_env1 = robot.data.joint_pos[1, 0].item()
|
|
assert pos_env0 == pytest.approx(pos_env1, abs=1e-4)
|
|
|
|
# Observations should differ by bias.
|
|
env.observation_manager._obs_buffer = None
|
|
obs = env.observation_manager.compute()["actor"]
|
|
assert isinstance(obs, torch.Tensor)
|
|
obs_diff = obs[1, 0].item() - obs[0, 0].item()
|
|
assert obs_diff == pytest.approx(bias_env1 - bias_env0, abs=1e-4)
|
|
|
|
env.close()
|
|
|
|
|
|
# =============================================================================
|
|
# Event randomization test
|
|
# =============================================================================
|
|
|
|
|
|
def test_randomize_encoder_bias_event(device):
|
|
"""randomize_encoder_bias should sample values within specified range."""
|
|
env_cfg = ManagerBasedRlEnvCfg(
|
|
scene=SceneCfg(
|
|
terrain=TerrainEntityCfg(terrain_type="plane"),
|
|
num_envs=100,
|
|
extent=10.0,
|
|
entities={"robot": _make_robot_cfg()},
|
|
),
|
|
observations={
|
|
"actor": ObservationGroupCfg(
|
|
terms={"obs": ObservationTermCfg(func=partial(mdp.joint_pos_rel, biased=True))},
|
|
),
|
|
},
|
|
actions={
|
|
"joint_pos": mdp.JointPositionActionCfg(
|
|
entity_name="robot", actuator_names=(".*",), scale=1.0
|
|
)
|
|
},
|
|
events={
|
|
"randomize_bias": EventTermCfg(
|
|
func=mdp.dr.encoder_bias,
|
|
mode="startup",
|
|
params={"bias_range": (-0.1, 0.1), "asset_cfg": SceneEntityCfg("robot")},
|
|
)
|
|
},
|
|
sim=SimulationCfg(mujoco=MujocoCfg(timestep=0.002, iterations=1)),
|
|
decimation=1,
|
|
episode_length_s=10.0,
|
|
)
|
|
|
|
env = ManagerBasedRlEnv(cfg=env_cfg, device=device)
|
|
env.reset()
|
|
|
|
biases = env.scene["robot"].data.encoder_bias[:, 0]
|
|
assert (biases >= -0.1).all() and (biases <= 0.1).all()
|
|
assert biases.std() > 0.01 # Not all identical.
|
|
|
|
env.close()
|