feat: support motrix kp/kd reset randomization

docs/users/zh_CN/06-domain-randomization.md

@@ -20,7 +20,7 @@
 2. 形式上基本都是结构化的：任务文件内定义 `domain_rand` 配置 dataclass、`DomainRandomizationProvider`、`ResetPlan`，`G1WalkFlat` 复用 `G1Joystick` 的 provider。
 3. 现在“统一”的主要是入口和执行流程，不是所有随机项本身。公共 helper [`build_common_reset_randomization()`](../../../src/unilab/dr/dr_utils.py) 目前生成 `base_mass_delta`、`base_com_offset`、`kp`、`kd`；公共 interval helper 目前只生成 push。
 4. [`ResetRandomizationPayload`](../../../src/unilab/dr/types.py) 已经能表达 `body_iquat`、`body_inertia`、`kp`、`kd`，且 [`MuJoCoBackend`](../../../src/unilab/base/backend/mujoco_backend.py) 已声明支持。是否真正使用这些项，仍取决于任务 provider 是否采样并下发。
-5. [`MotrixBackend`](../../../src/unilab/base/backend/motrix_backend.py) 当前只支持 `base_mass_delta`、`base_com_offset` 和 interval push。
+5. [`MotrixBackend`](../../../src/unilab/base/backend/motrix_backend.py) 当前支持 `base_mass_delta`、`base_com_offset`、`kp`、`kd` 和 interval push；并在初始化阶段要求模型 actuator 全部为 position actuator。
 6. `geom_size` 不属于 reset-lifecycle 字段；Sharpa-hand object geom scale 通过 init-lifecycle 的 model materialization 完成。
 
 ## 统一性评估表
@@ -92,7 +92,7 @@
 backend capability 当前是：
 
 - [`MuJoCoBackend`](../../../src/unilab/base/backend/mujoco_backend.py)：支持上面 6 个 reset term，且支持 interval push
-- [`MotrixBackend`](../../../src/unilab/base/backend/motrix_backend.py)：只支持 `base_mass_delta`、`base_com_offset`，且支持 interval push
+- [`MotrixBackend`](../../../src/unilab/base/backend/motrix_backend.py)：支持 `base_mass_delta`、`base_com_offset`、`kp`、`kd`，且支持 interval push；初始化阶段要求 actuator 全为 position
 
 但任务侧当前实际情况是：并不是所有 provider 都构造这些字段。backend contract 是能力边界，任务配置和 provider 是否下发 payload 才决定该任务是否实际启用对应 DR 项。
 

src/unilab/base/backend/motrix_backend.py

@@ -1,13 +1,15 @@
 import os
 import time
 from collections.abc import Sequence
-from typing import TypeVar
+from typing import TypeVar, cast
 
 import numpy as np
 
 from unilab.dr.types import (
     RESET_TERM_BASE_COM,
     RESET_TERM_BASE_MASS,
+    RESET_TERM_KD,
+    RESET_TERM_KP,
     DomainRandomizationCapabilities,
     IntervalRandomizationPlan,
     ResetRandomizationPayload,
@@ -93,6 +95,27 @@ def __init__(
         self._body_floatingbase = self._body.floatingbase
         self._joint_dof_pos_indices = np.asarray(self._model.joint_dof_pos_indices, dtype=np.intp)
         self._joint_dof_vel_indices = np.asarray(self._model.joint_dof_vel_indices, dtype=np.intp)
+        position_actuators: list["mtx.PositionActuator"] = []
+        for actuator in self._model.actuators:
+            if actuator.typ == "position":
+                position_actuators.append(cast("mtx.PositionActuator", actuator))
+        self._position_actuators = position_actuators
+        if len(self._position_actuators) != int(self._model.num_actuators):
+            raise ValueError(
+                "Motrix backend requires all actuators to be position actuators when "
+                "domain-randomization kp/kd support is enabled."
+            )
+        self._default_actuator_kp = np.zeros((self.num_actuators,), dtype=np.float64)
+        self._default_actuator_kd = np.zeros((self.num_actuators,), dtype=np.float64)
+        for actuator in self._position_actuators:
+            idx = int(actuator.index)
+            # TODO: switch to motrixsim model-level actuator gain API once available.
+            self._default_actuator_kp[idx] = float(
+                np.asarray(actuator.get_kp_override(self._data), dtype=np.float64)[0]
+            )
+            self._default_actuator_kd[idx] = float(
+                np.asarray(actuator.get_kd_override(self._data), dtype=np.float64)[0]
+            )
         self._floating_base_quat_indices: tuple[np.ndarray, ...] = tuple(
             np.asarray(floating_base.dof_pos_indices[3:7], dtype=np.intp)
             for floating_base in getattr(self._model, "floating_bases", [])
@@ -225,7 +248,9 @@ def set_state(
 
     def get_dr_capabilities(self) -> DomainRandomizationCapabilities:
         return DomainRandomizationCapabilities(
-            supported_reset_terms=frozenset({RESET_TERM_BASE_MASS, RESET_TERM_BASE_COM}),
+            supported_reset_terms=frozenset(
+                {RESET_TERM_BASE_MASS, RESET_TERM_BASE_COM, RESET_TERM_KP, RESET_TERM_KD}
+            ),
             supports_interval_push=True,
         )
 
@@ -413,8 +438,8 @@ def _apply_reset_randomization(
     ) -> None:
         if randomization is None or randomization.is_empty():
             return
-        unsupported = randomization.requested_terms() - frozenset(
-            {RESET_TERM_BASE_MASS, RESET_TERM_BASE_COM}
+        unsupported = (
+            randomization.requested_terms() - self.get_dr_capabilities().supported_reset_terms
         )
         if unsupported:
             terms = ", ".join(sorted(unsupported))
@@ -432,3 +457,29 @@ def _apply_reset_randomization(
             base_com = self._default_base_com_override[env_ids].copy()
             randomized_com = base_com + randomization.base_com_offset
             self._body_link.set_center_of_mass_override(data_slice, randomized_com)
+
+        num_reset = len(env_ids)
+        if randomization.kp is not None:
+            kp = np.asarray(randomization.kp, dtype=np.float32)
+            expected_shape = (num_reset, self.num_actuators)
+            if kp.shape != expected_shape:
+                raise ValueError(f"kp must have shape {expected_shape}, got {kp.shape}")
+            self._set_position_actuator_kp_override(data_slice, kp)
+
+        if randomization.kd is not None:
+            kd = np.asarray(randomization.kd, dtype=np.float32)
+            expected_shape = (num_reset, self.num_actuators)
+            if kd.shape != expected_shape:
+                raise ValueError(f"kd must have shape {expected_shape}, got {kd.shape}")
+            self._set_position_actuator_kd_override(data_slice, kd)
+
+    def _set_position_actuator_kp_override(self, data_slice, kp: np.ndarray) -> None:
+        for actuator in self._position_actuators:
+            actuator.set_kp_override(data_slice, kp[:, int(actuator.index)])
+
+    def _set_position_actuator_kd_override(self, data_slice, kd: np.ndarray) -> None:
+        for actuator in self._position_actuators:
+            actuator.set_kd_override(data_slice, kd[:, int(actuator.index)])
+
+    def get_actuator_gains(self) -> tuple[np.ndarray, np.ndarray]:
+        return self._default_actuator_kp.copy(), self._default_actuator_kd.copy()

tests/base/test_sim_backend.py

@@ -16,6 +16,7 @@
 from unilab.dr import (
     GeomSizeOverride,
     InitRandomizationPlan,
+    IntervalRandomizationPlan,
     ModelVariantSpec,
     ResetRandomizationPayload,
 )
@@ -602,21 +603,116 @@ def test_set_state_randomization_only_affects_target_envs(self, _ctx):
         np.testing.assert_allclose(updated_mass[1], original_mass[1], atol=1e-6)
         np.testing.assert_allclose(updated_mass[0], original_mass[0] + delta[0], atol=1e-6)
 
+    def test_set_state_base_com_randomization_only_affects_target_envs(self, _ctx):
+        bkd, _ = _ctx
+        nq = bkd.get_dof_pos().shape[-1] + 7
+        nv = bkd.get_dof_vel().shape[-1] + 6
+        qpos = _identity_qpos_mujoco(nq)
+        qvel = np.zeros((1, nv))
+        original_com = np.asarray(bkd._body_link.get_center_of_mass_override(bkd.data)).copy()
+        delta = np.array([[0.01, 0.0, 0.0]], dtype=np.float64)
+
+        bkd.set_state(
+            np.array([0]),
+            qpos,
+            qvel,
+            randomization=ResetRandomizationPayload(base_com_offset=delta),
+        )
+
+        updated_com = np.asarray(bkd._body_link.get_center_of_mass_override(bkd.data))
+        np.testing.assert_allclose(updated_com[1], original_com[1], atol=1e-6)
+        np.testing.assert_allclose(updated_com[0], original_com[0] + delta[0], atol=1e-6)
+
+    def test_set_state_kp_kd_randomization_only_affects_target_envs(self, _ctx):
+        bkd, _ = _ctx
+        nq = bkd.get_dof_pos().shape[-1] + 7
+        nv = bkd.get_dof_vel().shape[-1] + 6
+        qpos = _identity_qpos_mujoco(nq)
+        qvel = np.zeros((1, nv))
+
+        base_kp, base_kd = bkd.get_actuator_gains()
+        randomized_kp = (base_kp + 0.5)[None, :]
+        randomized_kd = (base_kd + 0.1)[None, :]
+
+        check_indices = [0]
+        if bkd.num_actuators > 1:
+            check_indices.append(bkd.num_actuators - 1)
+        position_actuators = {int(actuator.index): actuator for actuator in bkd._position_actuators}
+        original_kp = {
+            idx: np.asarray(position_actuators[idx].get_kp_override(bkd.data)).copy()
+            for idx in check_indices
+        }
+        original_kd = {
+            idx: np.asarray(position_actuators[idx].get_kd_override(bkd.data)).copy()
+            for idx in check_indices
+        }
+
+        bkd.set_state(
+            np.array([0]),
+            qpos,
+            qvel,
+            randomization=ResetRandomizationPayload(kp=randomized_kp, kd=randomized_kd),
+        )
+
+        for idx in check_indices:
+            updated_kp = np.asarray(position_actuators[idx].get_kp_override(bkd.data))
+            updated_kd = np.asarray(position_actuators[idx].get_kd_override(bkd.data))
+            np.testing.assert_allclose(updated_kp[1], original_kp[idx][1], atol=1e-6)
+            np.testing.assert_allclose(updated_kd[1], original_kd[idx][1], atol=1e-6)
+            np.testing.assert_allclose(updated_kp[0], randomized_kp[0, idx], atol=1e-6)
+            np.testing.assert_allclose(updated_kd[0], randomized_kd[0, idx], atol=1e-6)
+
+    def test_set_state_kp_randomization_shape_validation(self, _ctx):
+        bkd, _ = _ctx
+        nq = bkd.get_dof_pos().shape[-1] + 7
+        nv = bkd.get_dof_vel().shape[-1] + 6
+        qpos = _identity_qpos_mujoco(nq)
+        qvel = np.zeros((1, nv))
+        invalid_cols = bkd.num_actuators + 1
+
+        with pytest.raises(ValueError, match="kp must have shape"):
+            bkd.set_state(
+                np.array([0]),
+                qpos,
+                qvel,
+                randomization=ResetRandomizationPayload(kp=np.zeros((1, invalid_cols))),
+            )
+
     def test_set_state_unsupported_randomization_raises(self, _ctx):
         bkd, _ = _ctx
         nq = bkd.get_dof_pos().shape[-1] + 7
         nv = bkd.get_dof_vel().shape[-1] + 6
         qpos = _identity_qpos_mujoco(nq)
         qvel = np.zeros((1, nv))
 
-        with pytest.raises(NotImplementedError, match="kp"):
+        with pytest.raises(NotImplementedError, match="body_inertia"):
             bkd.set_state(
                 np.array([0]),
                 qpos,
                 qvel,
-                randomization=ResetRandomizationPayload(kp=np.zeros((1, 1))),
+                randomization=ResetRandomizationPayload(body_inertia=np.zeros((1, 1, 3))),
             )
 
+    def test_get_dr_capabilities_include_expected_terms(self, bkd):
+        caps = bkd.get_dr_capabilities()
+        assert {"base_mass_delta", "base_com_offset", "kp", "kd"}.issubset(
+            caps.supported_reset_terms
+        )
+        assert caps.supports_interval_push
+
+    def test_apply_interval_randomization_calls_push_robots(self, bkd):
+        called: dict[str, np.ndarray] = {}
+
+        def _fake_push_robots(force_range):
+            called["force_range"] = np.asarray(force_range, dtype=np.float64)
+
+        bkd.push_robots = _fake_push_robots  # type: ignore[method-assign]
+        limit = np.array([0.7, 0.3, 0.2], dtype=np.float64)
+
+        bkd.apply_interval_randomization(IntervalRandomizationPlan(push_perturbation_limit=limit))
+
+        np.testing.assert_allclose(called["force_range"], limit)
+
     # base kinematics
 
     def test_get_base_pos_shape(self, bkd):

tests/dr/test_manager.py

@@ -100,3 +100,23 @@ def test_manager_skips_unsupported_reset_terms_with_warning(caplog):
         "motrix backend does not support reset randomization terms: kp; skipping them."
         in caplog.text
     )
+
+
+def test_manager_keeps_supported_reset_terms_without_warning(caplog):
+    backend = _FakeBackend(
+        capabilities=DomainRandomizationCapabilities(
+            supported_reset_terms=frozenset({RESET_TERM_BASE_MASS, RESET_TERM_KP})
+        )
+    )
+    env = SimpleNamespace(_backend=backend)
+    manager = DomainRandomizationManager(env, _FakeProvider())
+
+    with caplog.at_level(logging.WARNING):
+        obs, info = manager.reset(np.array([0, 1], dtype=np.int32))
+
+    assert obs["obs"].shape == (2, 1)
+    assert info["commands"].shape == (2, 3)
+    assert backend.last_randomization is not None
+    assert backend.last_randomization.base_mass_delta is not None
+    assert backend.last_randomization.kp is not None
+    assert "skipping them" not in caplog.text

tests/ipc/test_async_runner.py

@@ -211,6 +211,7 @@ def test_start_collector_spawns_process():
     r.close()
 
 
+@pytest.mark.slow
 def test_start_collector_does_not_merge_runner_runtime_fields():
     r = _make_runner(sim_backend="motrix")
     report_queue = _SPAWN_CTX.Queue()