Reset invalid DreamWaQ physics states safely

motrix_envs/tests/test_dreamwaq_state_safety.py

@@ -0,0 +1,19 @@
+import numpy as np
+
+from motrix_envs.locomotion.go1.dreamwaq import _sanitize_dof_pos
+
+
+def test_sanitize_dof_pos_handles_zero_quaternion_with_nonfinite_joints():
+    dof_pos = np.zeros((2, 19), dtype=np.float32)
+    dof_pos[0, 3:7] = 0.0
+    dof_pos[0, 7] = np.inf
+    dof_pos[1, 3:7] = [0.0, 0.0, 0.5, 0.5]
+    dof_pos[1, 8] = np.nan
+
+    clean = _sanitize_dof_pos(dof_pos)
+
+    assert np.isfinite(clean).all()
+    np.testing.assert_allclose(clean[0, 3:7], [0.0, 0.0, 0.0, 1.0])
+    np.testing.assert_allclose(np.linalg.norm(clean[:, 3:7], axis=1), 1.0)
+    assert clean[0, 7] == 0.0
+    assert clean[1, 8] == 0.0