47372443ff
Upstream: https://github.com/pollen-robotics/microduck_rl Upstream-Commit: d424a0c899f6b33cbd3daeb279913134349c0b63 Upstream-Branch: develop
138 lines
4.4 KiB
Python
138 lines
4.4 KiB
Python
"""The critic obs must survive a non-finite sensor reading.
|
|
|
|
Regression for the 2026-08-21 crash: rsl_rl's check_nan killed a
|
|
Velocity2-Rough-Backlash run with "observation group 'critic' contains NaN".
|
|
`nan_state` (robot_state_is_nan) only covered joint + root state, but the
|
|
critic also carries three SENSOR-derived terms (raycast heights, contact
|
|
air-time, contact forces). MuJoCo can return a non-finite contact force while
|
|
the integrated robot state is still clean, so the env was never reset and the
|
|
NaN reached the runner.
|
|
"""
|
|
|
|
import torch
|
|
|
|
from mjlab_microduck.tasks import mdp as microduck_mdp
|
|
|
|
|
|
class _SensorData:
|
|
def __init__(self, force=None, heights=None):
|
|
self.force = force
|
|
self.heights = heights
|
|
|
|
|
|
class _Sensor:
|
|
def __init__(self, data):
|
|
self.data = data
|
|
|
|
|
|
class _Scene:
|
|
def __init__(self, sensors, asset):
|
|
self.sensors = sensors
|
|
self._asset = asset
|
|
|
|
def __getitem__(self, key):
|
|
return self.sensors[key] if key in self.sensors else self._asset
|
|
|
|
|
|
class _AssetData:
|
|
def __init__(self, n):
|
|
self.joint_pos = torch.zeros(n, 4)
|
|
self.joint_vel = torch.zeros(n, 4)
|
|
self.root_link_pos_w = torch.zeros(n, 3)
|
|
self.root_link_quat_w = torch.zeros(n, 4)
|
|
self.root_link_lin_vel_w = torch.zeros(n, 3)
|
|
self.root_link_ang_vel_w = torch.zeros(n, 3)
|
|
|
|
|
|
class _Asset:
|
|
def __init__(self, data):
|
|
self.data = data
|
|
|
|
|
|
class _Env:
|
|
def __init__(self, n, force):
|
|
self.num_envs = n
|
|
self.device = "cpu"
|
|
asset = _Asset(_AssetData(n))
|
|
self.scene = _Scene({"feet": _Sensor(_SensorData(force=force))}, asset)
|
|
|
|
|
|
def _force(n, bad_env=None, value=float("nan")):
|
|
f = torch.ones(n, 2, 3)
|
|
if bad_env is not None:
|
|
f[bad_env, 0, 0] = value
|
|
return f
|
|
|
|
|
|
def test_state_only_check_misses_bad_contact_force():
|
|
# This is the gap that killed the run: robot state is clean, force is not.
|
|
env = _Env(3, _force(3, bad_env=1))
|
|
assert not microduck_mdp.robot_state_is_nan(env).any()
|
|
|
|
|
|
def test_termination_catches_nan_contact_force():
|
|
env = _Env(3, _force(3, bad_env=1))
|
|
out = microduck_mdp.robot_state_is_nan(env, sensor_names=("feet",))
|
|
assert out.tolist() == [False, True, False]
|
|
|
|
|
|
def test_termination_catches_inf_contact_force():
|
|
env = _Env(3, _force(3, bad_env=2, value=float("inf")))
|
|
out = microduck_mdp.robot_state_is_nan(env, sensor_names=("feet",))
|
|
assert out.tolist() == [False, False, True]
|
|
|
|
|
|
def test_termination_ignores_missing_sensor():
|
|
env = _Env(2, _force(2))
|
|
assert not microduck_mdp.robot_state_is_nan(env, sensor_names=("nope",)).any()
|
|
|
|
|
|
def test_finite_helper_sanitizes_nan_and_inf():
|
|
x = torch.tensor([[1.0, float("nan"), float("inf"), float("-inf")]])
|
|
out = microduck_mdp._finite(x)
|
|
assert torch.isfinite(out).all()
|
|
assert out[0, 0] == 1.0
|
|
|
|
|
|
def test_safe_obs_wrappers_are_wired_into_the_critic():
|
|
# Guards must actually be installed on the env cfg, not just exist.
|
|
from mjlab_microduck.tasks.microduck_velocity_env_cfg import (
|
|
make_microduck_velocity_env_cfg,
|
|
)
|
|
|
|
cfg = make_microduck_velocity_env_cfg(rough=True)
|
|
terms = cfg.observations["critic"].terms
|
|
for name in ("foot_contact_forces", "foot_height", "foot_air_time"):
|
|
assert terms[name].func.__name__.endswith("_safe"), (
|
|
f"critic/{name} lost its NaN guard"
|
|
)
|
|
|
|
|
|
def test_nan_state_termination_watches_the_contact_sensor():
|
|
from mjlab_microduck.tasks.microduck_velocity_env_cfg import (
|
|
make_microduck_velocity_env_cfg,
|
|
)
|
|
|
|
cfg = make_microduck_velocity_env_cfg(rough=True)
|
|
params = cfg.terminations["nan_state"].params
|
|
assert params.get("sensor_names"), "nan_state no longer watches contact forces"
|
|
|
|
|
|
def test_standup_env_is_also_guarded():
|
|
# The deployed standing policy trains on StandUp, which builds on mjlab's
|
|
# base env (NOT the microduck velocity env) and therefore does not inherit
|
|
# the guards wired there.
|
|
from mjlab_microduck.tasks.microduck_standup_env_cfg import (
|
|
make_microduck_standup_env_cfg,
|
|
)
|
|
|
|
cfg = make_microduck_standup_env_cfg()
|
|
terms = cfg.observations["critic"].terms
|
|
for name in ("foot_contact_forces", "foot_air_time"):
|
|
assert terms[name].func.__name__.endswith("_safe"), (
|
|
f"standup critic/{name} lost its NaN guard"
|
|
)
|
|
assert cfg.terminations["nan_state"].params.get("sensor_names"), (
|
|
"standup nan_state no longer watches contact forces"
|
|
)
|