diff --git a/src/unilab/envs/locomotion/go1/joystick.py b/src/unilab/envs/locomotion/go1/joystick.py index 9fbb3521e..457fbbcf7 100644 --- a/src/unilab/envs/locomotion/go1/joystick.py +++ b/src/unilab/envs/locomotion/go1/joystick.py @@ -268,40 +268,3 @@ def _reward_swing_feet_z(self, ctx: RewardContext) -> np.ndarray: height_error = np.square(self.feet_pos[:, :, 2] - target_height) swing_rew = np.exp(-height_error / 0.01) * is_swing return np.asarray(np.sum(swing_rew, axis=1) / len(self._cfg.sensor.feet_pos)) - - def reset(self, env_indices: np.ndarray): - num_reset = len(env_indices) - qpos = np.tile(self._init_qpos, (num_reset, 1)) - qvel = np.tile(self._init_qvel, (num_reset, 1)) - - # Domain Randomization - dxy = np.random.uniform(-0.5, 0.5, (num_reset, 2)) - qpos[:, 0:2] += dxy - yaw = np.random.uniform(-np.pi, np.pi, (num_reset,)) - quat_yaw = np_yaw_to_quat(yaw) - qpos[:, 3:7] = np_quat_mul(qpos[:, 3:7], quat_yaw) - qvel[:, 0:6] = np.random.uniform(-0.5, 0.5, (num_reset, 6)) - - self._backend.set_state(env_indices, qpos, qvel) - - commands = np.random.uniform( - low=self._cfg.commands.vel_limit[0], - high=self._cfg.commands.vel_limit[1], - size=(num_reset, 3), - ) - - info = { - "commands": commands, - "current_actions": np.zeros((num_reset, self._num_action), dtype=get_global_dtype()), - "last_actions": np.zeros((num_reset, self._num_action), dtype=get_global_dtype()), - } - - linvel = self.get_local_linvel()[env_indices] - gyro = self.get_gyro()[env_indices] - gravity = self._backend.get_sensor_data("upvector")[env_indices] - dof_pos = self.get_dof_pos()[env_indices] - dof_vel = self.get_dof_vel()[env_indices] - obs = self._compute_obs( - info, linvel, gyro, gravity, dof_pos, dof_vel, self.feet_phase[env_indices] - ) - return obs, info diff --git a/tests/envs/test_go1_domain_randomization.py b/tests/envs/test_go1_domain_randomization.py new file mode 100644 index 000000000..9612d8bf5 --- /dev/null +++ b/tests/envs/test_go1_domain_randomization.py @@ -0,0 +1,62 @@ +from __future__ import annotations + +from typing import Any + +import numpy as np +import pytest + +pytest.importorskip("mujoco", reason="mujoco not installed") + +try: + from mujoco.batch_env import BatchEnvPool as _ # noqa: F401 +except Exception: + pytest.skip( + "mujoco.batch_env not available (platform/libstdc++ issue)", allow_module_level=True + ) + +from unilab.utils.algo_utils import ensure_registries + + +@pytest.mark.slow +def test_go1_mujoco_reset_applies_base_mass_and_com_domain_randomization( + default_go1_reward_config: dict[str, Any], +) -> None: + ensure_registries() + + from unilab.base import registry + + env = registry.make( + "Go1JoystickFlatTerrain", + num_envs=4, + sim_backend="mujoco", + env_cfg_override={"reward_config": default_go1_reward_config}, + ) + env_obj: Any = env + try: + env_obj.init_state() + backend: Any = env_obj._backend + base_body_id = int(backend._base_body_id) + body_mass = np.stack( + [backend._pool.get_field(i, "body_mass") for i in range(env_obj.num_envs)] + ) + body_ipos = np.stack( + [backend._pool.get_field(i, "body_ipos") for i in range(env_obj.num_envs)] + ) + body_ipos = body_ipos.reshape(env_obj.num_envs, -1, 3) + + base_mass = float(backend._base_body_mass[base_body_id]) + base_ipos = np.asarray(backend._base_body_ipos[base_body_id]) + randomized_base_mass = body_mass[:, base_body_id] + randomized_base_ipos = body_ipos[:, base_body_id] + delta_x = randomized_base_ipos[:, 0] - base_ipos[0] + delta_yz = randomized_base_ipos[:, 1:] - base_ipos[1:] + + assert np.unique(np.round(randomized_base_mass, 6)).size > 1 + assert np.unique(np.round(delta_x, 6)).size > 1 + assert np.all(randomized_base_mass >= base_mass - 1.5) + assert np.all(randomized_base_mass <= base_mass + 1.5) + assert np.all(delta_x >= -0.05) + assert np.all(delta_x <= 0.05) + np.testing.assert_allclose(delta_yz, 0.0) + finally: + env_obj.close()