fix: downgrade unsupported DR terms and move G1 motrix override to Hydra

src/unilab/dr/dr_utils.py

@@ -50,16 +50,11 @@ def build_common_reset_randomization(env: Any, num_reset: int) -> ResetRandomiza
 
 def validate_common_reset_randomization(
     env: Any, capabilities: DomainRandomizationCapabilities
-) -> None:
+) -> frozenset[str]:
     payload = build_common_reset_randomization(env, num_reset=1)
     if payload is None:
-        return
-    unsupported = payload.requested_terms() - capabilities.supported_reset_terms
-    if unsupported:
-        terms = ", ".join(sorted(unsupported))
-        raise NotImplementedError(
-            f"{env._backend.backend_type} backend does not support reset randomization terms: {terms}"
-        )
+        return frozenset()
+    return capabilities.get_unsupported_reset_terms(payload.requested_terms())
 
 
 def build_interval_push_plan(env: Any, step_counter: int) -> IntervalRandomizationPlan | None:

src/unilab/dr/manager.py

@@ -1,35 +1,36 @@
 from __future__ import annotations
 
+import logging
 from typing import Any
 
 import numpy as np
 
 from .provider import DomainRandomizationProvider
 from .types import DomainRandomizationCapabilities
 
+logger = logging.getLogger(__name__)
+
 
 class DomainRandomizationManager:
     def __init__(self, env: Any, provider: DomainRandomizationProvider):
         self._env = env
         self._provider = provider
         self._capabilities: DomainRandomizationCapabilities = env._backend.get_dr_capabilities()
+        self._warned_reset_terms: frozenset[str] = frozenset()
         self._provider.validate(env, self._capabilities)
 
     def reset(self, env_ids: np.ndarray) -> tuple[dict[str, np.ndarray], dict]:
         plan = self._provider.build_reset_plan(self._env, env_ids)
         payload = plan.randomization
         if payload is not None:
-            unsupported = payload.requested_terms() - self._capabilities.supported_reset_terms
+            payload, unsupported = self._capabilities.filter_reset_payload(payload)
             if unsupported:
-                terms = ", ".join(sorted(unsupported))
-                raise NotImplementedError(
-                    f"{self._env._backend.backend_type} backend does not support reset randomization terms: {terms}"
-                )
+                self._log_unsupported_reset_terms(unsupported)
         self._env._backend.set_state(
             plan.env_ids,
             plan.qpos,
             plan.qvel,
-            randomization=plan.randomization,
+            randomization=payload,
         )
         obs = self._provider.build_reset_observation(self._env, plan.env_ids, plan.info_updates)
         return obs, plan.info_updates
@@ -46,3 +47,14 @@ def apply_interval_randomization_if_due(self, step_counter: int) -> None:
                 f"{self._env._backend.backend_type} backend does not support interval push"
             )
         self._env._backend.apply_interval_randomization(plan)
+
+    def _log_unsupported_reset_terms(self, unsupported: frozenset[str]) -> None:
+        new_terms = frozenset(term for term in unsupported if term not in self._warned_reset_terms)
+        if not new_terms:
+            return
+        self._warned_reset_terms |= new_terms
+        logger.warning(
+            "%s backend does not support reset randomization terms: %s; skipping them.",
+            self._env._backend.backend_type,
+            ", ".join(sorted(new_terms)),
+        )

src/unilab/dr/types.py

@@ -18,6 +18,37 @@ class DomainRandomizationCapabilities:
     supported_reset_terms: frozenset[str] = field(default_factory=frozenset)
     supports_interval_push: bool = False
 
+    def supports_reset_term(self, term: str) -> bool:
+        return term in self.supported_reset_terms
+
+    def get_unsupported_reset_terms(self, requested_terms: frozenset[str]) -> frozenset[str]:
+        return frozenset(term for term in requested_terms if not self.supports_reset_term(term))
+
+    def filter_reset_payload(
+        self, payload: ResetRandomizationPayload
+    ) -> tuple[ResetRandomizationPayload | None, frozenset[str]]:
+        unsupported = self.get_unsupported_reset_terms(payload.requested_terms())
+        if not unsupported:
+            return payload, frozenset()
+
+        filtered = ResetRandomizationPayload(
+            base_mass_delta=(
+                payload.base_mass_delta if self.supports_reset_term(RESET_TERM_BASE_MASS) else None
+            ),
+            base_com_offset=(
+                payload.base_com_offset if self.supports_reset_term(RESET_TERM_BASE_COM) else None
+            ),
+            body_iquat=(
+                payload.body_iquat if self.supports_reset_term(RESET_TERM_BODY_IQUAT) else None
+            ),
+            body_inertia=(
+                payload.body_inertia if self.supports_reset_term(RESET_TERM_BODY_INERTIA) else None
+            ),
+            kp=payload.kp if self.supports_reset_term(RESET_TERM_KP) else None,
+            kd=payload.kd if self.supports_reset_term(RESET_TERM_KD) else None,
+        )
+        return (None if filtered.is_empty() else filtered), unsupported
+
 
 @dataclass
 class ResetRandomizationPayload:

tests/dr/test_manager.py

@@ -0,0 +1,100 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from types import SimpleNamespace
+from typing import Any
+
+import numpy as np
+
+from unilab.dr import (
+    DomainRandomizationCapabilities,
+    DomainRandomizationManager,
+    DomainRandomizationProvider,
+    ResetPlan,
+    ResetRandomizationPayload,
+)
+from unilab.dr.types import RESET_TERM_BASE_MASS, RESET_TERM_KP
+
+
+def test_capabilities_filter_reset_payload_drops_unsupported_terms():
+    capabilities = DomainRandomizationCapabilities(
+        supported_reset_terms=frozenset({RESET_TERM_BASE_MASS})
+    )
+    payload = ResetRandomizationPayload(
+        base_mass_delta=np.array([0.25]),
+        kp=np.array([[12.0, 12.0]]),
+    )
+
+    filtered, unsupported = capabilities.filter_reset_payload(payload)
+
+    assert unsupported == frozenset({RESET_TERM_KP})
+    assert filtered is not None
+    assert filtered.base_mass_delta is not None
+    np.testing.assert_allclose(filtered.base_mass_delta, np.array([0.25]))
+    assert filtered.kp is None
+
+
+@dataclass
+class _FakeBackend:
+    capabilities: DomainRandomizationCapabilities
+    backend_type: str = "motrix"
+
+    def __post_init__(self) -> None:
+        self.last_randomization: ResetRandomizationPayload | None = None
+
+    def get_dr_capabilities(self) -> DomainRandomizationCapabilities:
+        return self.capabilities
+
+    def set_state(
+        self,
+        env_indices: np.ndarray,
+        qpos: np.ndarray,
+        qvel: np.ndarray,
+        randomization: ResetRandomizationPayload | None = None,
+    ) -> None:
+        self.last_randomization = randomization
+
+
+class _FakeProvider(DomainRandomizationProvider):
+    def validate(self, env: Any, capabilities: DomainRandomizationCapabilities) -> None:
+        return None
+
+    def build_reset_plan(self, env: Any, env_ids: np.ndarray) -> ResetPlan:
+        return ResetPlan(
+            env_ids=env_ids,
+            qpos=np.zeros((len(env_ids), 8), dtype=np.float32),
+            qvel=np.zeros((len(env_ids), 7), dtype=np.float32),
+            info_updates={"commands": np.zeros((len(env_ids), 3), dtype=np.float32)},
+            randomization=ResetRandomizationPayload(
+                base_mass_delta=np.full((len(env_ids),), 0.1, dtype=np.float32),
+                kp=np.full((len(env_ids), 2), 5.0, dtype=np.float32),
+            ),
+        )
+
+    def build_reset_observation(
+        self, env: Any, env_ids: np.ndarray, info_updates: dict[str, Any]
+    ) -> dict[str, np.ndarray]:
+        return {"obs": np.zeros((len(env_ids), 1), dtype=np.float32)}
+
+
+def test_manager_skips_unsupported_reset_terms_with_warning(caplog):
+    backend = _FakeBackend(
+        capabilities=DomainRandomizationCapabilities(
+            supported_reset_terms=frozenset({RESET_TERM_BASE_MASS})
+        )
+    )
+    env = SimpleNamespace(_backend=backend)
+    manager = DomainRandomizationManager(env, _FakeProvider())
+
+    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
+    np.testing.assert_allclose(backend.last_randomization.base_mass_delta, np.array([0.1, 0.1]))
+    assert backend.last_randomization.kp is None
+    assert (
+        "motrix backend does not support reset randomization terms: kp; skipping them."
+        in caplog.text
+    )