From 8b726f36a8ce39c45cda0985ece56bff950f35b7 Mon Sep 17 00:00:00 2001 From: TATP-233 Date: Tue, 18 Aug 2026 16:40:23 +0800 Subject: [PATCH] feat(env): add manager pd gain reset event Refs #1092 [skip ci] --- src/unilab/base/entity.py | 96 +++++++++++ src/unilab/base/reset_state.py | 153 ++++++++++++++++++ src/unilab/envs/mdp/__init__.py | 2 + src/unilab/envs/mdp/events.py | 147 ++++++++++++++++- tests/base/test_reset_state.py | 72 ++++++++- .../locomotion/go2/test_manager_based_cfg.py | 54 +++++++ tests/envs/mdp/test_events.py | 146 ++++++++++++++++- 7 files changed, 660 insertions(+), 10 deletions(-) diff --git a/src/unilab/base/entity.py b/src/unilab/base/entity.py index 51487acf4..6e51998b8 100644 --- a/src/unilab/base/entity.py +++ b/src/unilab/base/entity.py @@ -474,6 +474,7 @@ def __init__( actuator_ids = self._resolve_enumerated_ids( "actuator", self._actuator_names, backend.get_actuator_names ) + self._actuator_ids = actuator_ids self._validate_joint_state(backend, joint_pos_ids, joint_vel_ids) self._validate_body_state(backend, root_body_ids, body_ids) @@ -1025,6 +1026,68 @@ def write_root_state_to_sim( term_name=f"{self.name}.write_root_state_to_sim", ) + def bind_actuator_gain_write( + self, + actuator_ids: np.ndarray | Sequence[int] | slice | None = None, + *, + term_name: str, + ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Bind selected actuator columns and immutable gain defaults on the cold path.""" + if self._reset_state is None: + raise self._capability_error( + "reset actuator-gain write", + "EntityScene was materialized without an env-owned reset transaction", + ) + if self._actuator_ids is None: + raise self._capability_error( + "reset actuator-gain write", + "actuator_names were not declared in EntityCfg", + ) + local_ids = self._normalize_local_actuator_ids( + actuator_ids, + capability="reset actuator-gain write", + ) + if local_ids.size == 0: + raise ValueError( + f"Entity '{self.name}' reset actuator-gain write selected no actuators" + ) + backend_ids = self._actuator_ids[local_ids] + _, default_kp, default_kd = self._reset_state.bind_actuator_gain_write( + backend_ids, + term_name=f"{term_name}:{self.name}", + ) + bound_local_ids = np.array(local_ids, copy=True) + bound_local_ids.setflags(write=False) + return bound_local_ids, default_kp, default_kd + + def write_actuator_gains_to_sim( + self, + kp: np.ndarray, + kd: np.ndarray, + actuator_ids: np.ndarray | Sequence[int] | slice | None = None, + env_ids: np.ndarray | slice | None = None, + *, + term_name: str = "pd_gains", + ) -> None: + """Stage entity-local actuator gains in the active reset transaction.""" + if self._reset_state is None or self._actuator_ids is None: + raise self._capability_error( + "reset actuator-gain write", + "actuator metadata or the env-owned reset transaction was not materialized", + ) + local_ids = self._normalize_local_actuator_ids( + actuator_ids, + capability="reset actuator-gain write", + ) + resolved_env_ids = self._normalize_reset_env_ids(env_ids) + self._reset_state.write_actuator_gains( + resolved_env_ids, + self._actuator_ids[local_ids], + kp, + kd, + term_name=f"{term_name}:{self.name}", + ) + def write_root_link_pose_to_sim( self, root_pose: np.ndarray, @@ -1160,6 +1223,39 @@ def _normalize_local_joint_ids( ) return ids + def _normalize_local_actuator_ids( + self, + actuator_ids: np.ndarray | Sequence[int] | slice | None, + *, + capability: str, + ) -> np.ndarray: + if actuator_ids is None: + ids = np.arange(self.num_actuators, dtype=np.intp) + elif isinstance(actuator_ids, slice): + ids = np.arange(self.num_actuators, dtype=np.intp)[actuator_ids] + else: + raw = np.asarray(actuator_ids) + if ( + raw.ndim != 1 + or not np.issubdtype(raw.dtype, np.integer) + or np.issubdtype(raw.dtype, np.bool_) + ): + raise TypeError( + f"Entity '{self.name}' {capability} actuator_ids must be a 1-D " + "integer array or slice" + ) + ids = np.asarray(raw, dtype=np.intp) + if np.any(ids < 0) or np.any(ids >= self.num_actuators): + raise IndexError( + f"Entity '{self.name}' {capability} actuator_ids out of range for " + f"{self.num_actuators} actuators: {ids.tolist()}" + ) + if np.unique(ids).size != ids.size: + raise ValueError( + f"Entity '{self.name}' {capability} actuator_ids contain duplicates: {ids.tolist()}" + ) + return ids + def _normalize_reset_env_ids(self, env_ids: np.ndarray | slice | None) -> np.ndarray: if env_ids is None: return np.arange(self._backend.num_envs, dtype=np.int32) diff --git a/src/unilab/base/reset_state.py b/src/unilab/base/reset_state.py index 6c94f2027..9c48d4520 100644 --- a/src/unilab/base/reset_state.py +++ b/src/unilab/base/reset_state.py @@ -13,6 +13,7 @@ import numpy as np from unilab.base.backend.base import BackendRootStateLayout, SimBackend +from unilab.dr.types import RESET_TERM_KD, RESET_TERM_KP, ResetRandomizationPayload from unilab.utils.rotation import np_quat_apply_inverse @@ -35,6 +36,11 @@ def __init__( self._default_qvel: np.ndarray | None = None self._qpos: np.ndarray | None = None self._qvel: np.ndarray | None = None + self._default_kp: np.ndarray | None = None + self._default_kd: np.ndarray | None = None + self._kp: np.ndarray | None = None + self._kd: np.ndarray | None = None + self._gain_dirty_mask = np.zeros(self._num_envs, dtype=np.bool_) self._requesting_terms: set[str] = set() @property @@ -62,9 +68,83 @@ def begin(self, env_ids: np.ndarray) -> None: self._active_mask.fill(False) self._active_mask[ids] = True self._dirty_mask.fill(False) + self._gain_dirty_mask.fill(False) self._requesting_terms.clear() self._active = True + def bind_actuator_gain_write( + self, + actuator_ids: np.ndarray, + *, + term_name: str, + ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Resolve gain mutation capability and immutable defaults on the cold path.""" + columns = self._validate_columns( + actuator_ids, + width=self._backend.num_actuators, + capability="actuator IDs", + term_name=term_name, + ) + self._materialize_default_actuator_gains(term_name) + assert self._default_kp is not None + assert self._default_kd is not None + selected_kp = np.array(self._default_kp[columns], copy=True) + selected_kd = np.array(self._default_kd[columns], copy=True) + selected_kp.setflags(write=False) + selected_kd.setflags(write=False) + bound_columns = np.array(columns, copy=True) + bound_columns.setflags(write=False) + return bound_columns, selected_kp, selected_kd + + def write_actuator_gains( + self, + env_ids: np.ndarray, + actuator_ids: np.ndarray, + kp: np.ndarray, + kd: np.ndarray, + *, + term_name: str, + ) -> None: + """Stage selected per-environment actuator gains in the reset transaction.""" + ids = self._prepare_state_write( + env_ids, + capability="actuator-gain", + term_name=term_name, + ) + columns = self._validate_columns( + actuator_ids, + width=self._backend.num_actuators, + capability="actuator IDs", + term_name=term_name, + ) + self._materialize_default_actuator_gains(term_name) + gains_shape = (ids.size, columns.size) + kp_values = self._validate_values( + kp, + shape=gains_shape, + capability="actuator kp", + term_name=term_name, + ) + kd_values = self._validate_values( + kd, + shape=gains_shape, + capability="actuator kd", + term_name=term_name, + ) + assert self._default_kp is not None + assert self._default_kd is not None + assert self._kp is not None + assert self._kd is not None + uninitialized = ids[~self._gain_dirty_mask[ids]] + if uninitialized.size: + self._kp[uninitialized] = self._default_kp + self._kd[uninitialized] = self._default_kd + if ids.size and columns.size: + self._kp[ids[:, None], columns[None, :]] = kp_values + self._kd[ids[:, None], columns[None, :]] = kd_values + self._gain_dirty_mask[ids] = True + self._dirty_mask[ids] = True + def reset_to_default(self, env_ids: np.ndarray, *, term_name: str) -> None: """Stage backend default qpos/qvel for a subset of the active reset.""" self._require_active() @@ -243,11 +323,13 @@ def commit(self) -> dict | None: return None assert self._qpos is not None assert self._qvel is not None + randomization = self._build_randomization_payload(dirty_ids) try: return self._backend.set_state( dirty_ids, self._qpos[dirty_ids], self._qvel[dirty_ids], + randomization=randomization, ) except (AttributeError, NotImplementedError) as exc: terms = ", ".join(sorted(self._requesting_terms)) @@ -284,6 +366,61 @@ def _materialize_default_state(self, term_name: str) -> None: self._qpos = np.empty((self._num_envs, default_qpos.size), dtype=default_qpos.dtype) self._qvel = np.empty((self._num_envs, default_qvel.size), dtype=default_qvel.dtype) + def _materialize_default_actuator_gains(self, term_name: str) -> None: + if self._default_kp is not None: + return + try: + capabilities = self._backend.get_dr_capabilities() + except (AttributeError, NotImplementedError) as exc: + raise self._capability_error(term_name, "actuator gain randomization", exc) from exc + required = frozenset((RESET_TERM_KP, RESET_TERM_KD)) + unsupported = capabilities.get_unsupported_reset_terms(required) + if unsupported: + detail = ", ".join(sorted(unsupported)) + raise self._capability_error( + term_name, + "actuator gain randomization", + NotImplementedError(f"unsupported reset payload fields: {detail}"), + ) + try: + kp, kd = self._backend.get_actuator_gains() + except (AttributeError, NotImplementedError) as exc: + raise self._capability_error(term_name, "default actuator gains", exc) from exc + default_kp = self._validate_gain_vector(kp, "default actuator kp", term_name) + default_kd = self._validate_gain_vector(kd, "default actuator kd", term_name) + self._default_kp = default_kp + self._default_kd = default_kd + self._kp = np.empty( + (self._num_envs, self._backend.num_actuators), + dtype=default_kp.dtype, + ) + self._kd = np.empty( + (self._num_envs, self._backend.num_actuators), + dtype=default_kd.dtype, + ) + + def _build_randomization_payload( + self, + dirty_ids: np.ndarray, + ) -> ResetRandomizationPayload | None: + gain_ids = np.flatnonzero(self._gain_dirty_mask).astype(np.int32, copy=False) + if gain_ids.size == 0: + return None + missing = dirty_ids[~self._gain_dirty_mask[dirty_ids]] + if missing.size: + terms = ", ".join(sorted(self._requesting_terms)) + raise RuntimeError( + "EventManager reset actuator-gain payload cannot represent sparse rows in " + f"one SimBackend.set_state call for term(s) [{terms}] on backend " + f"'{self._backend.backend_type}'; missing env IDs {missing.tolist()}" + ) + assert self._kp is not None + assert self._kd is not None + return ResetRandomizationPayload( + kp=np.array(self._kp[dirty_ids], copy=True), + kd=np.array(self._kd[dirty_ids], copy=True), + ) + def _prepare_state_write( self, env_ids: np.ndarray, @@ -378,6 +515,21 @@ def _validate_state_vector( result.setflags(write=False) return result + def _validate_gain_vector( + self, + value: np.ndarray, + capability: str, + term_name: str, + ) -> np.ndarray: + result = self._validate_state_vector(value, capability, term_name) + expected = (self._backend.num_actuators,) + if result.shape != expected: + raise ValueError( + f"EventManager term '{term_name}' capability '{capability}' on backend " + f"'{self._backend.backend_type}' returned shape {result.shape}; expected {expected}" + ) + return result + def _validate_ids(self, env_ids: np.ndarray, *, capability: str) -> np.ndarray: if not isinstance(env_ids, np.ndarray): raise TypeError( @@ -484,6 +636,7 @@ def _finish(self) -> None: self._active = False self._active_mask.fill(False) self._dirty_mask.fill(False) + self._gain_dirty_mask.fill(False) self._requesting_terms.clear() diff --git a/src/unilab/envs/mdp/__init__.py b/src/unilab/envs/mdp/__init__.py index 3b7fa400f..7e61bd0be 100644 --- a/src/unilab/envs/mdp/__init__.py +++ b/src/unilab/envs/mdp/__init__.py @@ -4,6 +4,7 @@ from unilab.envs.mdp.actions import JointPositionActionCfg as JointPositionActionCfg from unilab.envs.mdp.commands import UniformVelocityCommand as UniformVelocityCommand from unilab.envs.mdp.commands import UniformVelocityCommandCfg as UniformVelocityCommandCfg +from unilab.envs.mdp.events import pd_gains as pd_gains from unilab.envs.mdp.events import reset_root_state_uniform as reset_root_state_uniform from unilab.envs.mdp.events import reset_scene_to_default as reset_scene_to_default from unilab.envs.mdp.events import resolve_env_ids as resolve_env_ids @@ -53,6 +54,7 @@ "joint_vel_rel", "joint_vel_l2", "last_action", + "pd_gains", "is_alive", "is_terminated", "projected_gravity", diff --git a/src/unilab/envs/mdp/events.py b/src/unilab/envs/mdp/events.py index 81a9125c2..d86a382e4 100644 --- a/src/unilab/envs/mdp/events.py +++ b/src/unilab/envs/mdp/events.py @@ -6,10 +6,12 @@ from __future__ import annotations -from typing import TYPE_CHECKING, cast +from typing import TYPE_CHECKING, Any, Literal, cast import numpy as np +from unilab.managers.event_manager import EventTermCfg +from unilab.managers.manager_base import ManagerTermBase from unilab.managers.scene_entity_config import SceneEntityCfg from unilab.utils.rotation import np_quat_from_euler_xyz, np_quat_mul @@ -20,6 +22,53 @@ _DEFAULT_ASSET_CFG = SceneEntityCfg("robot") _SE3_KEYS = ("x", "y", "z", "roll", "pitch", "yaw") +_PD_GAIN_PARAM_NAMES = frozenset(("kp_range", "kd_range", "asset_cfg", "distribution", "operation")) + + +def _gain_range( + value: Any, + *, + name: str, + distribution: Literal["uniform", "log_uniform"], +) -> tuple[float, float]: + try: + bounds = np.asarray(value, dtype=np.float64) + except (TypeError, ValueError) as exc: + raise TypeError(f"pd_gains {name} must be a numeric (min, max) pair") from exc + if bounds.shape != (2,): + raise ValueError(f"pd_gains {name} must have shape (2,), got {bounds.shape}") + if not np.isfinite(bounds).all(): + raise ValueError(f"pd_gains {name} must contain only finite values") + lower, upper = float(bounds[0]), float(bounds[1]) + if lower > upper: + raise ValueError(f"pd_gains {name} minimum {lower} exceeds maximum {upper}") + if distribution == "log_uniform" and lower <= 0.0: + raise ValueError(f"pd_gains {name} must be positive for log_uniform sampling") + return lower, upper + + +def _gain_choice( + value: Any, + *, + name: str, + choices: tuple[str, ...], +) -> str: + if not isinstance(value, str): + raise TypeError(f"pd_gains {name} must be a string, got {type(value).__name__}") + if value not in choices: + raise ValueError(f"pd_gains {name} must be one of {choices}, got {value!r}") + return value + + +def _sample_gain_range( + rng: np.random.Generator, + bounds: tuple[float, float], + shape: tuple[int, int], + distribution: Literal["uniform", "log_uniform"], +) -> np.ndarray: + if distribution == "uniform": + return rng.uniform(bounds[0], bounds[1], size=shape) + return np.exp(rng.uniform(np.log(bounds[0]), np.log(bounds[1]), size=shape)) def _sample_se3_range( @@ -62,6 +111,100 @@ def resolve_env_ids(env: ManagerBasedRlEnv, env_ids: np.ndarray | None) -> np.nd return env_ids +class PdGains(ManagerTermBase): + """Pinned-mjlab-compatible PD gain randomization on UniLab reset payloads.""" + + def __init__(self, cfg: EventTermCfg, env: ManagerBasedRlEnv): + super().__init__(env) + if cfg.mode != "reset": + raise NotImplementedError( + "EventManager term 'pd_gains' only supports mode='reset' on the UniLab " + "set_state transaction; startup/interval/step model-field mutation is unavailable" + ) + if cfg.min_step_count_between_reset != 0: + raise NotImplementedError( + "EventManager term 'pd_gains' requires min_step_count_between_reset=0 " + "because sparse per-field reset rows cannot be represented by the current " + "SimBackend.set_state payload" + ) + unknown = sorted(set(cfg.params) - _PD_GAIN_PARAM_NAMES) + if unknown: + raise ValueError(f"EventManager term 'pd_gains' has unknown parameters {unknown}") + missing = [name for name in ("kp_range", "kd_range") if name not in cfg.params] + if missing: + raise ValueError(f"EventManager term 'pd_gains' is missing parameters {missing}") + + distribution = cast( + Literal["uniform", "log_uniform"], + _gain_choice( + cfg.params.get("distribution", "uniform"), + name="distribution", + choices=("uniform", "log_uniform"), + ), + ) + self._operation = cast( + Literal["scale", "abs"], + _gain_choice( + cfg.params.get("operation", "scale"), + name="operation", + choices=("scale", "abs"), + ), + ) + self._distribution = distribution + self._kp_range = _gain_range( + cfg.params["kp_range"], + name="kp_range", + distribution=distribution, + ) + self._kd_range = _gain_range( + cfg.params["kd_range"], + name="kd_range", + distribution=distribution, + ) + asset_cfg = cfg.params.get("asset_cfg", _DEFAULT_ASSET_CFG) + if not isinstance(asset_cfg, SceneEntityCfg): + raise TypeError( + "EventManager term 'pd_gains' asset_cfg must be SceneEntityCfg, got " + f"{type(asset_cfg).__name__}" + ) + self._entity = cast("Entity", env.scene[asset_cfg.name]) + self._actuator_ids, self._default_kp, self._default_kd = ( + self._entity.bind_actuator_gain_write( + asset_cfg.actuator_ids, + term_name="pd_gains", + ) + ) + + def __call__( + self, + env: ManagerBasedRlEnv, + env_ids: np.ndarray | None, + kp_range: tuple[float, float], + kd_range: tuple[float, float], + asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG, + distribution: Literal["uniform", "log_uniform"] = "uniform", + operation: Literal["scale", "abs"] = "scale", + ) -> None: + del kp_range, kd_range, asset_cfg, distribution, operation + ids = resolve_env_ids(env, env_ids) + shape = (len(ids), len(self._actuator_ids)) + kp = _sample_gain_range(env.rng, self._kp_range, shape, self._distribution) + kd = _sample_gain_range(env.rng, self._kd_range, shape, self._distribution) + if self._operation == "scale": + kp *= self._default_kp[None, :] + kd *= self._default_kd[None, :] + self._entity.write_actuator_gains_to_sim( + kp, + kd, + actuator_ids=self._actuator_ids, + env_ids=ids, + term_name="pd_gains", + ) + + +pd_gains = PdGains + + def reset_scene_to_default(env: ManagerBasedRlEnv, env_ids: np.ndarray | None) -> None: """Reset all materialized scene entities to backend default qpos/qvel.""" ids = resolve_env_ids(env, env_ids) @@ -106,4 +249,4 @@ def reset_root_state_uniform( asset.write_root_state_to_sim(root_states, env_ids=ids) -__all__ = ["reset_root_state_uniform", "reset_scene_to_default", "resolve_env_ids"] +__all__ = ["pd_gains", "reset_root_state_uniform", "reset_scene_to_default", "resolve_env_ids"] diff --git a/tests/base/test_reset_state.py b/tests/base/test_reset_state.py index 5242240ac..85939703a 100644 --- a/tests/base/test_reset_state.py +++ b/tests/base/test_reset_state.py @@ -9,10 +9,17 @@ from unilab.base.backend.base import BackendRootStateLayout, SimBackend from unilab.base.reset_state import ResetStateTransaction +from unilab.dr.types import ( + RESET_TERM_KD, + RESET_TERM_KP, + DomainRandomizationCapabilities, + ResetRandomizationPayload, +) class _Backend: backend_type = "fake" + num_actuators = 3 def __init__( self, @@ -28,6 +35,9 @@ def __init__( self.default_qpos_calls = 0 self.init_qvel_calls = 0 self.set_state_calls: list[tuple[np.ndarray, np.ndarray, np.ndarray]] = [] + self.randomization_calls: list[ResetRandomizationPayload | None] = [] + self.default_kp = np.array([10.0, 20.0, 30.0]) + self.default_kd = np.array([1.0, 2.0, 3.0]) def get_default_qpos(self): self.default_qpos_calls += 1 @@ -37,6 +47,14 @@ def get_init_qvel(self): self.init_qvel_calls += 1 return self.qvel + def get_dr_capabilities(self) -> DomainRandomizationCapabilities: + return DomainRandomizationCapabilities( + supported_reset_terms=frozenset((RESET_TERM_KP, RESET_TERM_KD)) + ) + + def get_actuator_gains(self) -> tuple[np.ndarray, np.ndarray]: + return self.default_kp.copy(), self.default_kd.copy() + def set_state( self, env_ids: np.ndarray, @@ -44,10 +62,10 @@ def set_state( qvel: np.ndarray, randomization=None, ) -> dict: - assert randomization is None if self.fail_set_state: raise NotImplementedError("reset upload disabled") self.set_state_calls.append((env_ids.copy(), qpos.copy(), qvel.copy())) + self.randomization_calls.append(randomization) return {"timing": {"set_state_ms": 1.0}} @@ -110,6 +128,58 @@ def test_exception_aborts_without_backend_mutation_and_next_reset_is_clean() -> np.testing.assert_array_equal(backend.set_state_calls[0][0], [1]) +def test_actuator_gains_compose_with_state_in_one_reset_commit() -> None: + backend = _Backend() + transaction = _transaction(backend) + columns, default_kp, default_kd = transaction.bind_actuator_gain_write( + np.array([2, 0], dtype=np.int32), + term_name="pd_gains:robot", + ) + np.testing.assert_array_equal(columns, [2, 0]) + np.testing.assert_array_equal(default_kp, [30.0, 10.0]) + np.testing.assert_array_equal(default_kd, [3.0, 1.0]) + + with transaction.scoped(np.array([0, 2], dtype=np.int32)): + transaction.reset_to_default(np.array([0, 2], dtype=np.int32), term_name="default") + transaction.write_actuator_gains( + np.array([2, 0], dtype=np.int32), + columns, + np.array([[5.0, 6.0], [7.0, 8.0]]), + np.array([[0.5, 0.6], [0.7, 0.8]]), + term_name="pd_gains:robot", + ) + + payload = backend.randomization_calls[0] + assert payload is not None + np.testing.assert_array_equal(payload.kp, [[8.0, 20.0, 7.0], [6.0, 20.0, 5.0]]) + np.testing.assert_array_equal(payload.kd, [[0.8, 2.0, 0.7], [0.6, 2.0, 0.5]]) + + +def test_actuator_gain_sparse_rows_abort_without_backend_mutation() -> None: + backend = _Backend() + transaction = _transaction(backend) + columns, _, _ = transaction.bind_actuator_gain_write( + np.array([0], dtype=np.int32), + term_name="pd_gains:robot", + ) + + with pytest.raises(RuntimeError, match=r"cannot represent sparse rows.*missing env IDs \[1\]"): + with transaction.scoped(np.array([0, 1], dtype=np.int32)): + transaction.reset_to_default( + np.array([0, 1], dtype=np.int32), + term_name="default", + ) + transaction.write_actuator_gains( + np.array([0], dtype=np.int32), + columns, + np.array([[11.0]]), + np.array([[1.1]]), + term_name="pd_gains:robot", + ) + + assert backend.set_state_calls == [] + + def test_joint_writes_initialize_defaults_and_compose_by_column() -> None: backend = _Backend() transaction = _transaction(backend) diff --git a/tests/envs/locomotion/go2/test_manager_based_cfg.py b/tests/envs/locomotion/go2/test_manager_based_cfg.py index 307912177..b70abbbea 100644 --- a/tests/envs/locomotion/go2/test_manager_based_cfg.py +++ b/tests/envs/locomotion/go2/test_manager_based_cfg.py @@ -16,6 +16,7 @@ make_go2_joystick_flat_manager_cfg, ) from unilab.envs.mdp import JointPositionAction, JointPositionActionCfg +from unilab.managers import EventTermCfg _JOINT_NAMES = ( "FL_hip_joint", @@ -294,3 +295,56 @@ def test_go2_manager_factory_executes_on_real_mujoco() -> None: assert state.truncated.dtype == np.bool_ finally: env.close() + + +def _read_runtime_actuator_gains(backend_type: str, backend) -> tuple[np.ndarray, np.ndarray]: + if backend_type == "mujoco": + assert backend._pool is not None + kp = np.stack([backend._pool.get_field(index, "kp") for index in range(backend.num_envs)]) + kd = np.stack([backend._pool.get_field(index, "kd") for index in range(backend.num_envs)]) + return kp, kd + assert backend_type == "motrix" + actuators = sorted(backend._position_actuators, key=lambda actuator: int(actuator.index)) + kp = np.column_stack( + [np.asarray(actuator.get_kp_override(backend._data)).reshape(-1) for actuator in actuators] + ) + kd = np.column_stack( + [np.asarray(actuator.get_kd_override(backend._data)).reshape(-1) for actuator in actuators] + ) + return kp, kd + + +@pytest.mark.parametrize("backend_type", ["mujoco", "motrix"]) +def test_go2_manager_pd_gains_mutates_real_backend_on_reset(backend_type: str) -> None: + cfg = make_go2_joystick_flat_manager_cfg() + cfg.events["pd_gains"] = EventTermCfg( + func=mdp.pd_gains, + mode="reset", + params={ + "kp_range": (31.5, 38.5), + "kd_range": (0.45, 0.55), + "operation": "abs", + }, + ) + assert cfg.scene is not None + backend = create_backend( + backend_type, + cfg.scene, + 2, + cfg.sim_dt, + base_name="base", + add_body_sensors=True, + **env_backend_kwargs(cfg), + ) + env = ManagerBasedRlEnv(cfg, backend, 2) + try: + env.reset(seed=29) + kp, kd = _read_runtime_actuator_gains(backend_type, backend) + assert kp.shape == (2, 12) + assert kd.shape == (2, 12) + assert np.all((kp >= 31.5) & (kp <= 38.5)) + assert np.all((kd >= 0.45) & (kd <= 0.55)) + assert np.unique(np.round(kp, 6)).size > 1 + assert np.unique(np.round(kd, 6)).size > 1 + finally: + env.close() diff --git a/tests/envs/mdp/test_events.py b/tests/envs/mdp/test_events.py index eb020e994..e73ace400 100644 --- a/tests/envs/mdp/test_events.py +++ b/tests/envs/mdp/test_events.py @@ -13,7 +13,14 @@ from unilab.base.backend.base import BackendRootStateLayout, SimBackend from unilab.base.entity import EntityCfg, EntityScene from unilab.base.reset_state import ResetStateTransaction +from unilab.dr.types import ( + RESET_TERM_KD, + RESET_TERM_KP, + DomainRandomizationCapabilities, + ResetRandomizationPayload, +) from unilab.envs import mdp +from unilab.managers import EventManager, EventTermCfg, SceneEntityCfg from unilab.managers._types import ManagerBasedRlEnv @@ -146,13 +153,20 @@ def test_uniform_root_state_none_ids_targets_all_environments() -> None: class _Backend: backend_type = "fake" num_envs = 3 - num_actuators = 0 + num_actuators = 3 - def __init__(self, *, root_layout_supported: bool = True) -> None: + def __init__( + self, + *, + root_layout_supported: bool = True, + gain_supported: bool = True, + ) -> None: self.root_layout_supported = root_layout_supported + self.gain_supported = gain_supported self.default_qpos = np.asarray([0.0, 0.0, 0.5, 1.0, 0.0, 0.0, 0.0]) self.init_qvel = np.zeros(6) self.set_state_calls: list[tuple[np.ndarray, np.ndarray, np.ndarray]] = [] + self.randomization_calls: list[ResetRandomizationPayload | None] = [] self.body_pos = np.zeros((self.num_envs, 1, 3)) self.body_quat = np.zeros((self.num_envs, 1, 4)) self.body_quat[:, :, 0] = 1.0 @@ -182,6 +196,19 @@ def get_dof_pos(self) -> np.ndarray: def get_dof_vel(self) -> np.ndarray: return np.empty((self.num_envs, 0)) + def get_actuator_names(self) -> tuple[str, ...]: + return ("a0", "a1", "a2") + + def get_actuator_ctrl_range(self) -> np.ndarray: + return np.tile([-1.0, 1.0], (self.num_actuators, 1)) + + def get_dr_capabilities(self) -> DomainRandomizationCapabilities: + terms = frozenset((RESET_TERM_KP, RESET_TERM_KD)) if self.gain_supported else frozenset() + return DomainRandomizationCapabilities(supported_reset_terms=terms) + + def get_actuator_gains(self) -> tuple[np.ndarray, np.ndarray]: + return np.array([10.0, 20.0, 30.0]), np.array([1.0, 2.0, 3.0]) + def get_body_pos_w(self, ids: np.ndarray) -> np.ndarray: return self.body_pos[:, ids] @@ -207,17 +234,20 @@ def set_state( qvel: np.ndarray, randomization=None, ) -> None: - assert randomization is None self.set_state_calls.append((env_ids.copy(), qpos.copy(), qvel.copy())) + self.randomization_calls.append(randomization) def _transaction_env( - *, root_layout_supported: bool = True + *, root_layout_supported: bool = True, gain_supported: bool = True, rng_seed: int = 5 ) -> tuple[ManagerBasedRlEnv, _Backend, ResetStateTransaction]: - backend = _Backend(root_layout_supported=root_layout_supported) + backend = _Backend( + root_layout_supported=root_layout_supported, + gain_supported=gain_supported, + ) transaction = ResetStateTransaction(cast(SimBackend, backend)) scene = EntityScene( - {"robot": EntityCfg(root_body_name="base")}, + {"robot": EntityCfg(root_body_name="base", actuator_names=("a0", "a1", "a2"))}, cast(SimBackend, backend), reset_state=transaction, ) @@ -225,7 +255,7 @@ def _transaction_env( ManagerBasedRlEnv, SimpleNamespace( num_envs=backend.num_envs, - rng=np.random.default_rng(5), + rng=np.random.default_rng(rng_seed), scene=scene, ), ) @@ -255,6 +285,108 @@ def test_uniform_root_state_composes_in_one_reset_transaction_commit() -> None: np.testing.assert_allclose(qvel[1], backend.init_qvel + [0.5, 0, 0, 0, 0, 0]) +def test_pd_gains_event_uses_selector_scale_and_exactly_once_reset_payload() -> None: + env, backend, transaction = _transaction_env(rng_seed=11) + manager = EventManager( + { + "randomize_pd": EventTermCfg( + func=mdp.pd_gains, + mode="reset", + params={ + "kp_range": (2.0, 2.0), + "kd_range": (3.0, 3.0), + "asset_cfg": SceneEntityCfg( + "robot", + actuator_names=["a2", "a0"], + preserve_order=True, + ), + }, + ) + }, + env, + ) + ids = np.array([0, 2], dtype=np.int32) + + with transaction.scoped(ids): + manager.apply(mode="reset", env_ids=ids, global_env_step_count=0) + assert backend.set_state_calls == [] + + assert len(backend.set_state_calls) == 1 + payload = backend.randomization_calls[0] + assert payload is not None + np.testing.assert_array_equal(payload.kp, [[20.0, 20.0, 60.0]] * 2) + np.testing.assert_array_equal(payload.kd, [[3.0, 2.0, 9.0]] * 2) + + +def test_pd_gains_event_supports_log_uniform_absolute_sampling() -> None: + env, backend, transaction = _transaction_env(rng_seed=11) + manager = EventManager( + { + "gain": EventTermCfg( + func=mdp.pd_gains, + mode="reset", + params={ + "kp_range": (0.25, 4.0), + "kd_range": (0.5, 2.0), + "distribution": "log_uniform", + "operation": "abs", + }, + ) + }, + env, + ) + ids = np.arange(3, dtype=np.int32) + + with transaction.scoped(ids): + manager.apply(mode="reset", env_ids=ids, global_env_step_count=0) + + payload = backend.randomization_calls[0] + assert payload is not None and payload.kp is not None and payload.kd is not None + assert np.all((payload.kp >= 0.25) & (payload.kp <= 4.0)) + assert np.all((payload.kd >= 0.5) & (payload.kd <= 2.0)) + assert np.unique(payload.kp[0]).size > 1 + + +@pytest.mark.parametrize( + ("cfg_kwargs", "match"), + [ + ({"mode": "startup"}, "only supports mode='reset'"), + ({"min_step_count_between_reset": 2}, "min_step_count_between_reset=0"), + ({"params": {"kp_range": (2.0, 1.0), "kd_range": (1.0, 1.0)}}, "minimum"), + ], +) +def test_pd_gains_invalid_config_fails_during_manager_construction( + cfg_kwargs: dict[str, Any], + match: str, +) -> None: + env, backend, _ = _transaction_env(rng_seed=11) + values: dict[str, Any] = { + "mode": "reset", + "params": {"kp_range": (1.0, 1.0), "kd_range": (1.0, 1.0)}, + } + values.update(cfg_kwargs) + + with pytest.raises((ValueError, NotImplementedError), match=match): + EventManager({"gain": EventTermCfg(func=mdp.pd_gains, **values)}, env) + assert backend.set_state_calls == [] + + +def test_pd_gains_missing_backend_capability_fails_during_manager_construction() -> None: + env, backend, _ = _transaction_env(gain_supported=False, rng_seed=11) + cfg = EventTermCfg( + func=mdp.pd_gains, + mode="reset", + params={"kp_range": (1.0, 1.0), "kd_range": (1.0, 1.0)}, + ) + + with pytest.raises( + NotImplementedError, + match="pd_gains:robot.*actuator gain randomization.*backend 'fake'", + ): + EventManager({"gain": cfg}, env) + assert backend.set_state_calls == [] + + def test_uniform_root_state_fixed_or_mocap_capability_fails_closed() -> None: env, backend, transaction = _transaction_env(root_layout_supported=False)