diff --git a/src/unilab/envs/locomotion/common/__init__.py b/src/unilab/envs/locomotion/common/__init__.py index 70e4ce59a..4528273b8 100644 --- a/src/unilab/envs/locomotion/common/__init__.py +++ b/src/unilab/envs/locomotion/common/__init__.py @@ -1,6 +1,7 @@ from .base import ControlConfigBase, LocomotionBaseCfg, LocomotionBaseEnv, Sensor from .commands import Commands, sample_velocity_commands from .domain_rand import DomainRandConfig +from .dr_provider import LocomotionDRProvider from .rewards import RewardContext __all__ = [ @@ -9,6 +10,7 @@ "DomainRandConfig", "LocomotionBaseCfg", "LocomotionBaseEnv", + "LocomotionDRProvider", "RewardContext", "Sensor", "sample_velocity_commands", diff --git a/src/unilab/envs/locomotion/common/dr_provider.py b/src/unilab/envs/locomotion/common/dr_provider.py new file mode 100644 index 000000000..018d61f17 --- /dev/null +++ b/src/unilab/envs/locomotion/common/dr_provider.py @@ -0,0 +1,127 @@ +"""Shared DomainRandomizationProvider for locomotion environments. + +Implements the common reset/interval randomization logic shared by +G1, Go1, and Go2 joystick environments. Subclasses override hooks +to provide robot-specific behaviour. +""" + +from __future__ import annotations + +from typing import Any, cast + +import numpy as np + +from unilab.base.dtype_config import get_global_dtype +from unilab.dr import ( + DomainRandomizationCapabilities, + DomainRandomizationProvider, + IntervalRandomizationPlan, + ResetPlan, +) +from unilab.dr.dr_utils import ( + build_common_reset_randomization, + build_interval_push_plan, + validate_common_reset_randomization, + validate_interval_push_support, + zero_actions, +) +from unilab.utils.math_utils import np_quat_mul, np_yaw_to_quat + + +class LocomotionDRProvider(DomainRandomizationProvider): + """Base DR provider for locomotion joystick environments. + + Shared logic: + - ``validate``, ``build_interval_randomization_plan``, ``_sample_commands`` + - ``build_reset_plan`` (template with hooks) + - ``build_reset_observation`` (template with hook) + + Override these hooks in subclasses: + - ``_get_qvel_limit`` — default ``0.5`` + - ``_build_extra_info_updates`` — default empty dict + - ``_compute_reset_obs`` — must be implemented per robot + """ + + # ── shared methods ─────────────────────────────────────────────── + + def validate(self, env: Any, capabilities: DomainRandomizationCapabilities) -> None: + validate_common_reset_randomization(env, capabilities) + validate_interval_push_support(env, capabilities) + + def build_interval_randomization_plan( + self, env: Any, step_counter: int + ) -> IntervalRandomizationPlan | None: + return build_interval_push_plan(env, step_counter) + + def _sample_commands(self, env: Any, num_reset: int) -> np.ndarray: + low = np.asarray(env.cfg.commands.vel_limit[0], dtype=get_global_dtype()) + high = np.asarray(env.cfg.commands.vel_limit[1], dtype=get_global_dtype()) + return np.asarray( + np.random.uniform(low=low, high=high, size=(num_reset, 3)), dtype=get_global_dtype() + ) + + # ── template: build_reset_plan ─────────────────────────────────── + + def _get_qvel_limit(self, env: Any) -> float: + """Return the base qvel reset limit. Override for configurable limits.""" + return 0.5 + + def _build_extra_info_updates(self, env: Any, num_reset: int) -> dict[str, np.ndarray]: + """Return additional info_updates entries (e.g. gait_phase for G1).""" + return {} + + def build_reset_plan(self, env: Any, env_ids: np.ndarray) -> ResetPlan: + num_reset = len(env_ids) + qpos = np.tile(env._init_qpos, (num_reset, 1)) + qvel = np.tile(env._init_qvel, (num_reset, 1)) + qpos[:, 0:2] += np.random.uniform(-0.5, 0.5, (num_reset, 2)) + yaw = np.random.uniform(-np.pi, np.pi, (num_reset,)) + qpos[:, 3:7] = np_quat_mul(qpos[:, 3:7], np_yaw_to_quat(yaw)) + limit = self._get_qvel_limit(env) + qvel[:, 0:6] = np.asarray( + np.random.uniform(-limit, limit, size=(num_reset, 6)), dtype=get_global_dtype() + ) + info_updates: dict[str, Any] = { + "commands": self._sample_commands(env, num_reset), + "current_actions": zero_actions(num_reset, env._num_action), + "last_actions": zero_actions(num_reset, env._num_action), + } + info_updates.update(self._build_extra_info_updates(env, num_reset)) + return ResetPlan( + env_ids=env_ids, + qpos=qpos, + qvel=qvel, + info_updates=info_updates, + randomization=build_common_reset_randomization(env, num_reset), + ) + + # ── template: build_reset_observation ──────────────────────────── + + def _compute_reset_obs( + self, + env: Any, + env_ids: np.ndarray, + info_updates: dict[str, Any], + linvel: np.ndarray, + gyro: np.ndarray, + gravity: np.ndarray, + dof_pos: np.ndarray, + dof_vel: np.ndarray, + ) -> dict[str, np.ndarray]: + """Compute reset observations. Override per robot to pass correct args to _compute_obs.""" + raise NotImplementedError + + def build_reset_observation( + self, env: Any, env_ids: np.ndarray, info_updates: dict[str, Any] + ) -> dict[str, np.ndarray]: + linvel = env.get_local_linvel()[env_ids] + gyro = env.get_gyro()[env_ids] + gravity = env._backend.get_sensor_data("upvector")[env_ids] + dof_pos = env.get_dof_pos()[env_ids] + dof_vel = env.get_dof_vel()[env_ids] + return cast( + dict[str, np.ndarray], + self._compute_reset_obs( + env, env_ids, info_updates, linvel, gyro, gravity, dof_pos, dof_vel + ), + ) diff --git a/src/unilab/envs/locomotion/g1/joystick.py b/src/unilab/envs/locomotion/g1/joystick.py index 3b849435a..5726621c2 100644 --- a/src/unilab/envs/locomotion/g1/joystick.py +++ b/src/unilab/envs/locomotion/g1/joystick.py @@ -4,7 +4,7 @@ import math from dataclasses import dataclass, field -from typing import Any, cast +from typing import Any import numpy as np from etils import epath @@ -14,25 +14,12 @@ from unilab.base.backend import create_backend from unilab.base.dtype_config import get_global_dtype from unilab.base.np_env import NpEnvState -from unilab.dr import ( - DomainRandomizationCapabilities, - DomainRandomizationProvider, - IntervalRandomizationPlan, - ResetPlan, -) -from unilab.dr.dr_utils import ( - build_common_reset_randomization, - build_interval_push_plan, - validate_common_reset_randomization, - validate_interval_push_support, - zero_actions, -) from unilab.envs.locomotion.common import rewards from unilab.envs.locomotion.common.commands import Commands from unilab.envs.locomotion.common.domain_rand import DomainRandConfig +from unilab.envs.locomotion.common.dr_provider import LocomotionDRProvider from unilab.envs.locomotion.common.rewards import RewardContext from unilab.envs.locomotion.g1.base import G1BaseCfg, G1BaseEnv -from unilab.utils.math_utils import np_quat_mul, np_yaw_to_quat @dataclass @@ -154,22 +141,12 @@ def apply_backend_overrides(self, backend_type: str) -> None: self.reset_base_qvel_limit = float(reset_base_qvel_limit) -class G1JoystickDomainRandomizationProvider(DomainRandomizationProvider): - def validate(self, env: Any, capabilities: DomainRandomizationCapabilities) -> None: - validate_common_reset_randomization(env, capabilities) - validate_interval_push_support(env, capabilities) +class G1JoystickDomainRandomizationProvider(LocomotionDRProvider): + def _get_qvel_limit(self, env: Any) -> float: + return float(env.cfg.reset_base_qvel_limit) - def build_interval_randomization_plan( - self, env: Any, step_counter: int - ) -> IntervalRandomizationPlan | None: - return build_interval_push_plan(env, step_counter) - - def _sample_commands(self, env: Any, num_reset: int) -> np.ndarray: - low = np.asarray(env.cfg.commands.vel_limit[0], dtype=get_global_dtype()) - high = np.asarray(env.cfg.commands.vel_limit[1], dtype=get_global_dtype()) - return np.asarray( - np.random.uniform(low=low, high=high, size=(num_reset, 3)), dtype=get_global_dtype() - ) + def _build_extra_info_updates(self, env: Any, num_reset: int) -> dict[str, np.ndarray]: + return {"gait_phase": self._sample_gait_phase(env, num_reset)} def _sample_gait_phase(self, env: Any, num_reset: int) -> np.ndarray: mode = env.cfg.gait_phase_init_mode @@ -181,43 +158,18 @@ def _sample_gait_phase(self, env: Any, num_reset: int) -> np.ndarray: phase = np.random.uniform(0.0, 2.0 * np.pi, size=(num_reset,)) return np.asarray(np.column_stack([phase, phase + np.pi]), dtype=get_global_dtype()) - def build_reset_plan(self, env: Any, env_ids: np.ndarray) -> ResetPlan: - num_reset = len(env_ids) - qpos = np.tile(env._init_qpos, (num_reset, 1)) - qvel = np.tile(env._init_qvel, (num_reset, 1)) - qpos[:, 0:2] += np.random.uniform(-0.5, 0.5, (num_reset, 2)) - yaw = np.random.uniform(-np.pi, np.pi, (num_reset,)) - qpos[:, 3:7] = np_quat_mul(qpos[:, 3:7], np_yaw_to_quat(yaw)) - limit = float(env.cfg.reset_base_qvel_limit) - qvel[:, 0:6] = np.asarray( - np.random.uniform(-limit, limit, size=(num_reset, 6)), dtype=get_global_dtype() - ) - info_updates = { - "commands": self._sample_commands(env, num_reset), - "current_actions": zero_actions(num_reset, env._num_action), - "last_actions": zero_actions(num_reset, env._num_action), - "gait_phase": self._sample_gait_phase(env, num_reset), - } - return ResetPlan( - env_ids=env_ids, - qpos=qpos, - qvel=qvel, - info_updates=info_updates, - randomization=build_common_reset_randomization(env, num_reset), - ) - - def build_reset_observation( - self, env: Any, env_ids: np.ndarray, info_updates: dict[str, Any] + def _compute_reset_obs( + self, + env: Any, + env_ids: Any, + info_updates: Any, + linvel: Any, + gyro: Any, + gravity: Any, + dof_pos: Any, + dof_vel: Any, ) -> dict[str, np.ndarray]: - linvel = env.get_local_linvel()[env_ids] - gyro = env.get_gyro()[env_ids] - gravity = env._backend.get_sensor_data("upvector")[env_ids] - dof_pos = env.get_dof_pos()[env_ids] - dof_vel = env.get_dof_vel()[env_ids] - return cast( - dict[str, np.ndarray], - env._compute_obs(info_updates, linvel, gyro, gravity, dof_pos, dof_vel), - ) + return env._compute_obs(info_updates, linvel, gyro, gravity, dof_pos, dof_vel) # type: ignore[no-any-return] @registry.env("G1JoystickFlatTerrain", sim_backend="mujoco") diff --git a/src/unilab/envs/locomotion/go1/joystick.py b/src/unilab/envs/locomotion/go1/joystick.py index 9fbb3521e..4a95aacda 100644 --- a/src/unilab/envs/locomotion/go1/joystick.py +++ b/src/unilab/envs/locomotion/go1/joystick.py @@ -1,7 +1,7 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import Any, cast +from typing import Any import numpy as np from etils import epath @@ -11,22 +11,10 @@ from unilab.base.backend import create_backend from unilab.base.dtype_config import get_global_dtype from unilab.base.np_env import NpEnvState -from unilab.dr import ( - DomainRandomizationCapabilities, - DomainRandomizationProvider, - IntervalRandomizationPlan, - ResetPlan, -) -from unilab.dr.dr_utils import ( - build_common_reset_randomization, - build_interval_push_plan, - validate_common_reset_randomization, - validate_interval_push_support, - zero_actions, -) from unilab.envs.locomotion.common import rewards from unilab.envs.locomotion.common.commands import Commands from unilab.envs.locomotion.common.domain_rand import DomainRandConfig +from unilab.envs.locomotion.common.dr_provider import LocomotionDRProvider from unilab.envs.locomotion.common.rewards import RewardContext from unilab.envs.locomotion.go1.base import Go1BaseCfg, Go1BaseEnv from unilab.utils.math_utils import np_quat_mul, np_yaw_to_quat @@ -84,57 +72,20 @@ def apply_legacy_motrix_profile(self) -> None: self.commands.vel_limit = [list(v) for v in profile.command_vel_limit] -class Go1JoystickDomainRandomizationProvider(DomainRandomizationProvider): - def validate(self, env: Any, capabilities: DomainRandomizationCapabilities) -> None: - validate_common_reset_randomization(env, capabilities) - validate_interval_push_support(env, capabilities) - - def build_interval_randomization_plan( - self, env: Any, step_counter: int - ) -> IntervalRandomizationPlan | None: - return build_interval_push_plan(env, step_counter) - - def _sample_commands(self, env: Any, num_reset: int) -> np.ndarray: - low = np.asarray(env.cfg.commands.vel_limit[0], dtype=get_global_dtype()) - high = np.asarray(env.cfg.commands.vel_limit[1], dtype=get_global_dtype()) - return np.asarray( - np.random.uniform(low=low, high=high, size=(num_reset, 3)), dtype=get_global_dtype() - ) - - def build_reset_plan(self, env: Any, env_ids: np.ndarray) -> ResetPlan: - num_reset = len(env_ids) - qpos = np.tile(env._init_qpos, (num_reset, 1)) - qvel = np.tile(env._init_qvel, (num_reset, 1)) - qpos[:, 0:2] += np.random.uniform(-0.5, 0.5, (num_reset, 2)) - yaw = np.random.uniform(-np.pi, np.pi, (num_reset,)) - qpos[:, 3:7] = np_quat_mul(qpos[:, 3:7], np_yaw_to_quat(yaw)) - qvel[:, 0:6] = np.random.uniform(-0.5, 0.5, (num_reset, 6)) - info_updates = { - "commands": self._sample_commands(env, num_reset), - "current_actions": zero_actions(num_reset, env._num_action), - "last_actions": zero_actions(num_reset, env._num_action), - } - return ResetPlan( - env_ids=env_ids, - qpos=qpos, - qvel=qvel, - info_updates=info_updates, - randomization=build_common_reset_randomization(env, num_reset), - ) - - def build_reset_observation( - self, env: Any, env_ids: np.ndarray, info_updates: dict[str, Any] +class Go1JoystickDomainRandomizationProvider(LocomotionDRProvider): + def _compute_reset_obs( + self, + env: Any, + env_ids: Any, + info_updates: Any, + linvel: Any, + gyro: Any, + gravity: Any, + dof_pos: Any, + dof_vel: Any, ) -> dict[str, np.ndarray]: - linvel = env.get_local_linvel()[env_ids] - gyro = env.get_gyro()[env_ids] - gravity = env._backend.get_sensor_data("upvector")[env_ids] - dof_pos = env.get_dof_pos()[env_ids] - dof_vel = env.get_dof_vel()[env_ids] - return cast( - dict[str, np.ndarray], - env._compute_obs( - info_updates, linvel, gyro, gravity, dof_pos, dof_vel, env.feet_phase[env_ids] - ), + return env._compute_obs( # type: ignore[no-any-return] + info_updates, linvel, gyro, gravity, dof_pos, dof_vel, env.feet_phase[env_ids] ) diff --git a/src/unilab/envs/locomotion/go2/joystick.py b/src/unilab/envs/locomotion/go2/joystick.py index fdf05bbf4..c9a01b816 100644 --- a/src/unilab/envs/locomotion/go2/joystick.py +++ b/src/unilab/envs/locomotion/go2/joystick.py @@ -1,7 +1,7 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import Any, cast +from typing import Any import numpy as np from etils import epath @@ -11,25 +11,12 @@ from unilab.base.backend import create_backend from unilab.base.dtype_config import get_global_dtype from unilab.base.np_env import NpEnvState -from unilab.dr import ( - DomainRandomizationCapabilities, - DomainRandomizationProvider, - IntervalRandomizationPlan, - ResetPlan, -) -from unilab.dr.dr_utils import ( - build_common_reset_randomization, - build_interval_push_plan, - validate_common_reset_randomization, - validate_interval_push_support, - zero_actions, -) from unilab.envs.locomotion.common import rewards from unilab.envs.locomotion.common.commands import Commands from unilab.envs.locomotion.common.domain_rand import DomainRandConfig +from unilab.envs.locomotion.common.dr_provider import LocomotionDRProvider from unilab.envs.locomotion.common.rewards import RewardContext from unilab.envs.locomotion.go2.base import Go2BaseCfg, Go2BaseEnv -from unilab.utils.math_utils import np_quat_mul, np_yaw_to_quat @dataclass @@ -88,57 +75,20 @@ def apply_legacy_motrix_profile(self) -> None: self.commands.vel_limit = [list(v) for v in profile.command_vel_limit] -class Go2JoystickDomainRandomizationProvider(DomainRandomizationProvider): - def validate(self, env: Any, capabilities: DomainRandomizationCapabilities) -> None: - validate_common_reset_randomization(env, capabilities) - validate_interval_push_support(env, capabilities) - - def build_interval_randomization_plan( - self, env: Any, step_counter: int - ) -> IntervalRandomizationPlan | None: - return build_interval_push_plan(env, step_counter) - - def _sample_commands(self, env: Any, num_reset: int) -> np.ndarray: - low = np.asarray(env.cfg.commands.vel_limit[0], dtype=get_global_dtype()) - high = np.asarray(env.cfg.commands.vel_limit[1], dtype=get_global_dtype()) - return np.asarray( - np.random.uniform(low=low, high=high, size=(num_reset, 3)), dtype=get_global_dtype() - ) - - def build_reset_plan(self, env: Any, env_ids: np.ndarray) -> ResetPlan: - num_reset = len(env_ids) - qpos = np.tile(env._init_qpos, (num_reset, 1)) - qvel = np.tile(env._init_qvel, (num_reset, 1)) - qpos[:, 0:2] += np.random.uniform(-0.5, 0.5, (num_reset, 2)) - yaw = np.random.uniform(-np.pi, np.pi, (num_reset,)) - qpos[:, 3:7] = np_quat_mul(qpos[:, 3:7], np_yaw_to_quat(yaw)) - qvel[:, 0:6] = np.random.uniform(-0.5, 0.5, (num_reset, 6)) - info_updates = { - "commands": self._sample_commands(env, num_reset), - "current_actions": zero_actions(num_reset, env._num_action), - "last_actions": zero_actions(num_reset, env._num_action), - } - return ResetPlan( - env_ids=env_ids, - qpos=qpos, - qvel=qvel, - info_updates=info_updates, - randomization=build_common_reset_randomization(env, num_reset), - ) - - def build_reset_observation( - self, env: Any, env_ids: np.ndarray, info_updates: dict[str, Any] +class Go2JoystickDomainRandomizationProvider(LocomotionDRProvider): + def _compute_reset_obs( + self, + env: Any, + env_ids: Any, + info_updates: Any, + linvel: Any, + gyro: Any, + gravity: Any, + dof_pos: Any, + dof_vel: Any, ) -> dict[str, np.ndarray]: - linvel = env.get_local_linvel()[env_ids] - gyro = env.get_gyro()[env_ids] - gravity = env._backend.get_sensor_data("upvector")[env_ids] - dof_pos = env.get_dof_pos()[env_ids] - dof_vel = env.get_dof_vel()[env_ids] - return cast( - dict[str, np.ndarray], - env._compute_obs( - info_updates, linvel, gyro, gravity, dof_pos, dof_vel, env.feet_phase[env_ids] - ), + return env._compute_obs( # type: ignore[no-any-return] + info_updates, linvel, gyro, gravity, dof_pos, dof_vel, env.feet_phase[env_ids] )