121 lines
4.0 KiB
Python
121 lines
4.0 KiB
Python
from types import SimpleNamespace
|
|
|
|
import numpy as np
|
|
|
|
from motrix_envs.locomotion.go1.dreamwaq import (
|
|
DreamWaQCfg,
|
|
DreamWaQTask,
|
|
_advance_streak,
|
|
_sanitize_dof_pos,
|
|
)
|
|
from motrix_envs.locomotion.go1.walk_np import Go1WalkTask
|
|
|
|
|
|
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
|
|
|
|
|
|
def test_advance_streak_requires_consecutive_failures():
|
|
streak = np.array([0, 1, 3], dtype=np.int32)
|
|
|
|
streak = _advance_streak(streak, np.array([True, True, False]))
|
|
np.testing.assert_array_equal(streak, [1, 2, 0])
|
|
|
|
streak = _advance_streak(streak, np.array([False, True, True]))
|
|
np.testing.assert_array_equal(streak, [0, 3, 1])
|
|
|
|
|
|
def test_apply_action_advances_three_frame_action_history():
|
|
task = Go1WalkTask.__new__(Go1WalkTask)
|
|
task.get_dof_vel = lambda data: np.ones((1, 12), dtype=np.float32)
|
|
task._compute_torques = lambda actions, data: actions
|
|
state = SimpleNamespace(
|
|
data=SimpleNamespace(actuator_ctrls=None),
|
|
info={
|
|
"current_actions": np.full((1, 12), 2.0, dtype=np.float32),
|
|
"last_actions": np.full((1, 12), 1.0, dtype=np.float32),
|
|
"last_last_actions": np.zeros((1, 12), dtype=np.float32),
|
|
},
|
|
)
|
|
actions = np.full((1, 12), 3.0, dtype=np.float32)
|
|
|
|
task.apply_action(actions, state)
|
|
actions.fill(9.0)
|
|
|
|
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_commands_are_zero_or_symmetric_without_low_speeds(monkeypatch):
|
|
task = DreamWaQTask.__new__(DreamWaQTask)
|
|
task._cfg = SimpleNamespace(
|
|
commands=SimpleNamespace(
|
|
forward_speed_range=(0.5, 1.0), stand_probability=0.2
|
|
)
|
|
)
|
|
monkeypatch.setattr(
|
|
np.random,
|
|
"uniform",
|
|
lambda low, high, size: np.array([0.5, 0.75, 1.0], dtype=np.float32),
|
|
)
|
|
monkeypatch.setattr(
|
|
np.random,
|
|
"choice",
|
|
lambda values, size: np.array([-1.0, -1.0, 1.0], dtype=np.float32),
|
|
)
|
|
monkeypatch.setattr(
|
|
np.random,
|
|
"random",
|
|
lambda size: np.array([0.1, 0.2, 0.9], dtype=np.float32),
|
|
)
|
|
|
|
commands = task.resample_commands(3)
|
|
|
|
np.testing.assert_array_equal(commands[:, 0], [0.0, -0.75, 1.0])
|
|
np.testing.assert_array_equal(commands[:, 1:], 0.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
|