Import upstream snapshot d424a0c899f6b33cbd3daeb279913134349c0b63

Upstream: https://github.com/pollen-robotics/microduck_rl
Upstream-Commit: d424a0c899f6b33cbd3daeb279913134349c0b63
Upstream-Branch: develop

tests/test_obs_nan_guard.py

@@ -0,0 +1,137 @@
+"""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"
+    )