fix: restore go1 reset domain randomization path

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

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()