Align DreamWaQ Go1 environment rewards and initialization

motrix_envs/tests/test_dreamwaq_state_safety.py

@@ -2,7 +2,7 @@ from types import SimpleNamespace
 
 import numpy as np
 
-from motrix_envs.locomotion.go1.dreamwaq import _sanitize_dof_pos
+from motrix_envs.locomotion.go1.dreamwaq import DreamWaQCfg, DreamWaQTask, _sanitize_dof_pos
 from motrix_envs.locomotion.go1.walk_np import Go1WalkTask
 
 
@@ -42,3 +42,35 @@ def test_apply_action_advances_three_frame_action_history():
     np.testing.assert_array_equal(state.info["last_last_actions"], 1.0)
     np.testing.assert_array_equal(state.info["last_actions"], 2.0)
     np.testing.assert_array_equal(state.info["current_actions"], 3.0)
+
+
+def test_dreamwaq_feet_air_time_matches_upstream_contact_filtering():
+    task = DreamWaQTask.__new__(DreamWaQTask)
+    task._cfg = SimpleNamespace(ctrl_dt=0.02)
+    task._num_envs = 1
+    info = {
+        "feet_air_time": np.array([[0.6, 0.0]], dtype=np.float32),
+        "contacts": np.array([[True, False]]),
+        "last_contacts": np.array([[False, False]]),
+    }
+
+    task.update_feet_air_time(info)
+    reward = task._reward_feet_air_time(
+        np.array([[1.0, 0.0, 0.0]], dtype=np.float32), info
+    )
+
+    np.testing.assert_allclose(reward, [0.12])
+    np.testing.assert_allclose(info["feet_air_time"], [[0.0, 0.02]])
+    np.testing.assert_array_equal(info["last_contacts"], info["contacts"])
+
+
+def test_dreamwaq_config_matches_upstream_go1_defaults():
+    cfg = DreamWaQCfg()
+
+    assert cfg.reward_config.only_positive_rewards
+    assert cfg.reward_config.scales["tracking_lin_vel"] == 1.0
+    assert cfg.reward_config.scales["tracking_ang_vel"] == 0.5
+    assert cfg.reward_config.scales["feet_air_time"] == 0.1
+    assert cfg.reward_config.scales["base_height"] == -10.0
+    assert cfg.init_state.default_joint_angles["FL_hip"] == 0.1
+    assert cfg.init_state.default_joint_angles["RR_hip"] == -0.1