diff --git a/docs/sphinx/source/en/2-user_guide/5-domain_randomization/2-writing_providers.md b/docs/sphinx/source/en/2-user_guide/5-domain_randomization/2-writing_providers.md index 061ab04e6..00799b7db 100644 --- a/docs/sphinx/source/en/2-user_guide/5-domain_randomization/2-writing_providers.md +++ b/docs/sphinx/source/en/2-user_guide/5-domain_randomization/2-writing_providers.md @@ -19,8 +19,35 @@ Current provider examples define one or more of these plan methods: - Return a reset plan with state updates and a reset randomization payload. - Return an interval plan for push or body-force perturbations. -The shared types live in `src/unilab/dr/types.py`, and the manager lives in -`src/unilab/dr/manager.py`. +Interval plans are built from `IntervalTermOp` descriptors (term name, NumPy +payload, optional `body_ids`; see {doc}`../../4-developer_guide/2-contracts/4-dr_contract`): + +```python +from unilab.dr import INTERVAL_TERM_BODY_FORCE, IntervalRandomizationPlan, IntervalTermOp + + +def build_interval_randomization_plan(self, env, step_counter): + ... + return IntervalRandomizationPlan( + ops=( + IntervalTermOp( + INTERVAL_TERM_BODY_FORCE, + force, # shape (num_envs, len(body_ids), 3) + body_ids=body_ids, + ), + ), + ) +``` + +Migration note: returning interval plans via the legacy fields +(`push_perturbation_limit`, `body_ids`, `body_force`, ...) is deprecated. Such +plans are still adapted 1:1 through `IntervalRandomizationPlan.iter_ops()`, +but new providers should populate `ops`; the legacy fields will be removed in +the next unisim-core major release. + +The shared types live in `unisim.dr.types` (interval term descriptors in +`unisim.dr.interval`), re-exported from `src/unilab/dr/__init__.py`, and the +manager lives in `src/unilab/dr/manager.py`. ## Rules diff --git a/docs/sphinx/source/en/4-developer_guide/2-contracts/4-dr_contract.md b/docs/sphinx/source/en/4-developer_guide/2-contracts/4-dr_contract.md index 1104d2e9b..30d6b4999 100644 --- a/docs/sphinx/source/en/4-developer_guide/2-contracts/4-dr_contract.md +++ b/docs/sphinx/source/en/4-developer_guide/2-contracts/4-dr_contract.md @@ -26,22 +26,57 @@ A task that uses DR should define: 4. Interval behavior through `IntervalRandomizationPlan` when needed. 5. Env construction that calls `self._init_domain_randomization(...)`. -Shared types live in `src/unilab/dr/types.py`, and manager behavior lives in -`src/unilab/dr/manager.py`. +Shared types live in `unisim.dr.types` (interval term descriptors in +`unisim.dr.interval`); both are re-exported from `src/unilab/dr/__init__.py`. +Manager behavior lives in `src/unilab/dr/manager.py`. ## Backend Capability Boundary Backend support is explicit. A reset or interval item only counts as a unified DR item when three pieces exist together: -1. `ResetRandomizationPayload` or `IntervalRandomizationPlan` has an explicit - field. +1. `ResetRandomizationPayload` has an explicit field, or + `IntervalRandomizationPlan.ops` carries an `IntervalTermOp` for the term. 2. The backend declares and implements the capability. -3. The task config/provider samples and dispatches that field. +3. The task config/provider samples and dispatches that field or op. MuJoCo and Motrix differences stay in backend capability declarations, backend implementations, and owner YAMLs. +## Interval Term Descriptors + +Interval plans are term-descriptor based: `IntervalRandomizationPlan.ops` +carries a tuple of `IntervalTermOp` entries (term name, NumPy payload, +optional `body_ids`) from `unisim.dr.interval`, re-exported through +`unilab.dr`. + +- Builtin term names are the `INTERVAL_TERM_*` constants; their payload + contracts are pinned by `INTERVAL_TERM_SPECS` (`push`: payload shape `(3,)`, + no `body_ids`; the four body terms: payload shape + `(num_envs, len(body_ids), 3)` with required `body_ids`). + `IntervalTermOp.validate()` enforces these contracts for builtin terms; + unknown custom terms pass validation through untouched. +- Capability ownership stays with the backend: + `DomainRandomizationCapabilities.supported_interval_terms` is the + authoritative declaration, queried via `supports_interval_term` / + `get_unsupported_interval_terms`. +- `DomainRandomizationManager.apply_interval_randomization_if_due` is generic: + it contains no term names and no per-term branches, so a backend-owned + custom term needs no manager change. Terms missing from the capability set + fail closed with `NotImplementedError` naming the backend type and the + terms; on the backend side, `SimBackend.apply_interval_randomization` + routes each op through its handler table and fails closed with the backend + class and term name when no handler exists. +- Ops and plans must stay pickle-safe (protocol 4) across spawn-based + collector processes: stdlib + NumPy frozen dataclasses only. +- The legacy plan fields (`push_perturbation_limit`, `body_ids`, + `body_linear_velocity_delta`, `body_angular_velocity_delta`, `body_force`, + `body_torque`) and the legacy `supports_interval_*` capability bools are + deprecated: `IntervalRandomizationPlan.iter_ops()` still adapts set legacy + fields into ops 1:1, and the bools remain as capability fallbacks. New + providers should populate `ops`; the legacy fields will be removed in the + next unisim-core major release. + ## MuJoCo BatchEnvPool Snapshot Current MuJoCo reset randomization uses `BatchEnvPool.reset(..., @@ -49,7 +84,7 @@ randomization=...)` with a fixed field whitelist. Indexed reads and writes are available through `get_field_indexed(...)` and `set_field_indexed(...)`. This interface lives in the `mujoco-uni-runtime` package (`mujoco_uni.batch_env`), not in this repository; the reset-term constants that map onto it are in -`src/unilab/dr/types.py`. +`unisim.dr.types`. The supported reset fields and their per-env block shapes are below. The leading dimension is always `len(env_ids)`; the trailing block size is the field's full @@ -76,7 +111,7 @@ Two caveats: - `geom_size` is not in `SUPPORTED_FIELDS`. Geometry size is expressed through init-lifecycle model materialization (see `GeomSizeOverride` / - `ModelVariantSpec` in `src/unilab/dr/types.py`), not reset randomization. + `ModelVariantSpec` in `unisim.dr.types`), not reset randomization. - `gravity` reset randomization requires a `mujoco-uni-runtime` build that ships it. This repository depends on the official `mujoco` package (`>=3.5`, with the default version pinned by `uv.lock`) @@ -97,7 +132,8 @@ payloads. ## Evidence In Repo -- DR types: `src/unilab/dr/types.py` +- DR types: `unisim.dr.types` and `unisim.dr.interval`, re-exported by + `src/unilab/dr/__init__.py` - DR manager: `src/unilab/dr/manager.py` - Backend interface: `unisim.backend.base` - Example providers: `src/unilab/tasks/locomotion/common/dr_provider.py`, diff --git a/docs/sphinx/source/zh_CN/2-user_guide/5-domain_randomization/2-writing_providers.md b/docs/sphinx/source/zh_CN/2-user_guide/5-domain_randomization/2-writing_providers.md index 901600ee3..a32798504 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/5-domain_randomization/2-writing_providers.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/5-domain_randomization/2-writing_providers.md @@ -17,7 +17,33 @@ provider;它们在 owner YAML 中通过 Hydra `events:` manager term 声明随 - 返回带有状态更新和 reset 随机化 payload 的 reset plan。 - 返回用于 push 或 body-force 扰动的 interval plan。 -共享类型位于 `src/unilab/dr/types.py`,manager 位于 +Interval plan 由 `IntervalTermOp` 描述符构建(term 名称、NumPy payload、 +可选的 `body_ids`;见 {doc}`../../4-developer_guide/2-contracts/4-dr_contract`): + +```python +from unilab.dr import INTERVAL_TERM_BODY_FORCE, IntervalRandomizationPlan, IntervalTermOp + + +def build_interval_randomization_plan(self, env, step_counter): + ... + return IntervalRandomizationPlan( + ops=( + IntervalTermOp( + INTERVAL_TERM_BODY_FORCE, + force, # 形状 (num_envs, len(body_ids), 3) + body_ids=body_ids, + ), + ), + ) +``` + +迁移说明:通过旧版字段(`push_perturbation_limit`、`body_ids`、 +`body_force` 等)返回 interval plan 已废弃。这类 plan 仍会经 +`IntervalRandomizationPlan.iter_ops()` 1:1 适配,但新 provider 应填充 +`ops`;旧字段将在下一个 unisim-core major release 中移除。 + +共享类型位于 `unisim.dr.types`(interval term 描述符位于 +`unisim.dr.interval`),由 `src/unilab/dr/__init__.py` 再导出,manager 位于 `src/unilab/dr/manager.py`。 ## 规则 diff --git a/docs/sphinx/source/zh_CN/4-developer_guide/2-contracts/4-dr_contract.md b/docs/sphinx/source/zh_CN/4-developer_guide/2-contracts/4-dr_contract.md index f84fdfe55..6572c4f8b 100644 --- a/docs/sphinx/source/zh_CN/4-developer_guide/2-contracts/4-dr_contract.md +++ b/docs/sphinx/source/zh_CN/4-developer_guide/2-contracts/4-dr_contract.md @@ -24,27 +24,59 @@ Domain randomization 是一个 env-owner 的 provider 契约,加上 backend 4. 必要时通过 `IntervalRandomizationPlan` 实现的 interval 行为。 5. 在 env 构造中调用 `self._init_domain_randomization(...)`。 -共享类型位于 `src/unilab/dr/types.py`,manager 行为位于 -`src/unilab/dr/manager.py`。 +共享类型位于 `unisim.dr.types`(interval term 描述符位于 +`unisim.dr.interval`),两者都由 `src/unilab/dr/__init__.py` 再导出。 +manager 行为位于 `src/unilab/dr/manager.py`。 ## Backend 能力边界 Backend 支持是显式的。只有当以下三个部分同时存在时,一个 reset 或 interval 条目才算作统一的 DR 条目: -1. `ResetRandomizationPayload` 或 `IntervalRandomizationPlan` 中有明确的字段。 +1. `ResetRandomizationPayload` 中有明确的字段,或 + `IntervalRandomizationPlan.ops` 携带该 term 的 `IntervalTermOp`。 2. backend 声明并实现了该能力。 -3. 任务 config/provider 对该字段进行采样并分发。 +3. 任务 config/provider 对该字段或 op 进行采样并分发。 MuJoCo 与 Motrix 的差异保留在 backend 能力声明、backend 实现与 owner YAML 中。 +## Interval Term 描述符 + +Interval plan 基于 term 描述符:`IntervalRandomizationPlan.ops` 携带一个 +`IntervalTermOp` 元组(term 名称、NumPy payload、可选的 `body_ids`),定义在 +`unisim.dr.interval`,并经 `unilab.dr` 再导出。 + +- 内置 term 名称即 `INTERVAL_TERM_*` 常量;其 payload 契约由 + `INTERVAL_TERM_SPECS` 固定(`push`:payload 形状 `(3,)`,不接受 + `body_ids`;四个 body term:payload 形状 `(num_envs, len(body_ids), 3)`, + 且必须携带 `body_ids`)。`IntervalTermOp.validate()` 对内置 term 强制 + 这些契约;未知的自定义 term 原样通过校验。 +- 能力所有权留在 backend: + `DomainRandomizationCapabilities.supported_interval_terms` 是权威声明, + 通过 `supports_interval_term` / `get_unsupported_interval_terms` 查询。 +- `DomainRandomizationManager.apply_interval_randomization_if_due` 是通用的: + 不包含任何 term 名称或按 term 分支,因此 backend 拥有的自定义 term 无需 + 修改 manager。不在能力集合中的 term 会 fail-closed,抛出带有 backend 类型 + 与 term 名称的 `NotImplementedError`;在 backend 一侧, + `SimBackend.apply_interval_randomization` 把每个 op 路由到 handler 表, + 缺少 handler 时 fail-closed,抛出带有 backend 类名与 term 名称的 + `NotImplementedError`。 +- Op 与 plan 必须保持 pickle 安全(protocol 4),以跨 spawn 方式的 collector + 子进程传递:只允许 stdlib + NumPy 的 frozen dataclass。 +- 旧版 plan 字段(`push_perturbation_limit`、`body_ids`、 + `body_linear_velocity_delta`、`body_angular_velocity_delta`、`body_force`、 + `body_torque`)与旧版 `supports_interval_*` 能力布尔位已废弃: + `IntervalRandomizationPlan.iter_ops()` 仍会把已设置的旧字段 1:1 适配为 + op,布尔位仍作为能力回退。新 provider 应填充 `ops`;旧字段将在下一个 + unisim-core major release 中移除。 + ## MuJoCo BatchEnvPool 快照 当前 MuJoCo 的 reset 随机化使用 `BatchEnvPool.reset(..., randomization=...)`, 并带有固定的字段白名单。带索引的读写可通过 `get_field_indexed(...)` 与 `set_field_indexed(...)` 实现。该接口位于 `mujoco-uni-runtime` 包 (`mujoco_uni.batch_env`),不在本仓库中;映射到它的 reset-term 常量定义在 -`src/unilab/dr/types.py`。 +`unisim.dr.types`。 支持的 reset 字段及其每 env 整块形状如下。首维始终是 `len(env_ids)`;尾部 整块大小是该字段在单个 `mjModel` 里的完整 flat 宽度。 @@ -68,7 +100,7 @@ refresh 行为由 backend 固定:`body_mass`、`body_ipos`、`body_iquat`、 两点注意: - `geom_size` 不在 `SUPPORTED_FIELDS` 里。几何尺寸通过 init-lifecycle 的模型 - materialization 表达(见 `src/unilab/dr/types.py` 中的 `GeomSizeOverride` / + materialization 表达(见 `unisim.dr.types` 中的 `GeomSizeOverride` / `ModelVariantSpec`),不走 reset 随机化。 - `gravity` 的 reset 随机化需要包含它的 `mujoco-uni-runtime` 构建。本仓库依赖 官方 `mujoco` 包(`>=3.5`,默认版本由 `uv.lock` 钉住)加 `mujoco-uni-runtime`,其 `SUPPORTED_FIELDS` @@ -87,7 +119,8 @@ actuator 的机制泄漏到共享 payload 里。 ## 仓库中的证据 -- DR 类型:`src/unilab/dr/types.py` +- DR 类型:`unisim.dr.types` 与 `unisim.dr.interval`,由 + `src/unilab/dr/__init__.py` 再导出 - DR manager:`src/unilab/dr/manager.py` - Backend 接口:`unisim.backend.base` - 示例 provider:`src/unilab/tasks/locomotion/common/dr_provider.py`、 diff --git a/pyproject.rocm.toml b/pyproject.rocm.toml index 24cdbad3e..5d446c4dd 100644 --- a/pyproject.rocm.toml +++ b/pyproject.rocm.toml @@ -25,7 +25,7 @@ dependencies = [ "numpy", # Physics implementations are provided by the independently released # unisim-core package from the production PyPI index. - "unisim-core>=0.1.14", + "unisim-core>=1.1.0", # RL algorithms and async runtimes live in the independently released # uni-rl package (distribution name ``unilab-rl``); see pyproject.toml. "unilab-rl==0.2.0", diff --git a/pyproject.toml b/pyproject.toml index a4f377499..09b06a80f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,7 +43,7 @@ dependencies = [ "numpy", # Physics implementations are provided by the independently released # unisim-core package from the production PyPI index. - "unisim-core>=0.1.14", + "unisim-core>=1.1.0", # RL algorithms and async runtimes (PPO/APPO/SAC/TD3/HORA runners, # collectors, IPC, logging) live in the independently released uni-rl # package (distribution name ``unilab-rl``), consumed via the injected diff --git a/src/unilab/dr/__init__.py b/src/unilab/dr/__init__.py index c349c1466..82cd4befa 100644 --- a/src/unilab/dr/__init__.py +++ b/src/unilab/dr/__init__.py @@ -4,10 +4,16 @@ """ from unisim.dr.types import ( + INTERVAL_TERM_BODY_ANGULAR_VELOCITY_DELTA, + INTERVAL_TERM_BODY_FORCE, + INTERVAL_TERM_BODY_LINEAR_VELOCITY_DELTA, + INTERVAL_TERM_BODY_TORQUE, + INTERVAL_TERM_PUSH, DomainRandomizationCapabilities, GeomSizeOverride, InitRandomizationPlan, IntervalRandomizationPlan, + IntervalTermOp, ModelVariantSpec, ResetPlan, ResetRandomizationPayload, @@ -17,12 +23,18 @@ from .provider import DomainRandomizationProvider __all__ = [ + "INTERVAL_TERM_BODY_ANGULAR_VELOCITY_DELTA", + "INTERVAL_TERM_BODY_FORCE", + "INTERVAL_TERM_BODY_LINEAR_VELOCITY_DELTA", + "INTERVAL_TERM_BODY_TORQUE", + "INTERVAL_TERM_PUSH", "DomainRandomizationCapabilities", "DomainRandomizationManager", "DomainRandomizationProvider", "GeomSizeOverride", "InitRandomizationPlan", "IntervalRandomizationPlan", + "IntervalTermOp", "ModelVariantSpec", "ResetPlan", "ResetRandomizationPayload", diff --git a/src/unilab/dr/dr_utils.py b/src/unilab/dr/dr_utils.py index dac016af0..58da49503 100644 --- a/src/unilab/dr/dr_utils.py +++ b/src/unilab/dr/dr_utils.py @@ -4,8 +4,10 @@ import numpy as np from unisim.dr.types import ( + INTERVAL_TERM_PUSH, DomainRandomizationCapabilities, IntervalRandomizationPlan, + IntervalTermOp, ResetRandomizationPayload, ) @@ -197,14 +199,16 @@ def build_interval_push_plan(env: Any, step_counter: int) -> IntervalRandomizati return None if step_counter % domain_rand.push_interval != 0: return None - return IntervalRandomizationPlan(push_perturbation_limit=domain_rand.max_force) + return IntervalRandomizationPlan( + ops=(IntervalTermOp(INTERVAL_TERM_PUSH, np.asarray(domain_rand.max_force)),) + ) def validate_interval_push_support(env: Any, capabilities: DomainRandomizationCapabilities) -> None: domain_rand = getattr(env.cfg, "domain_rand", None) if domain_rand is None or not getattr(domain_rand, "push_robots", False): return - if not capabilities.supports_interval_push: + if not capabilities.supports_interval_term(INTERVAL_TERM_PUSH): raise NotImplementedError( f"{env._backend.backend_type} backend does not support interval push" ) diff --git a/src/unilab/dr/manager.py b/src/unilab/dr/manager.py index fa2936ffe..94b41f11a 100644 --- a/src/unilab/dr/manager.py +++ b/src/unilab/dr/manager.py @@ -92,34 +92,13 @@ def apply_interval_randomization_if_due(self, step_counter: int) -> None: plan = self._provider.build_interval_randomization_plan(self._env, step_counter) if plan is None or plan.is_empty(): return - if ( - plan.push_perturbation_limit is not None - and not self._capabilities.supports_interval_push - ): - raise NotImplementedError( - f"{self._env._backend.backend_type} backend does not support interval push" - ) - if ( - plan.body_linear_velocity_delta is not None - and not self._capabilities.supports_interval_body_velocity_delta - ): - raise NotImplementedError( - f"{self._env._backend.backend_type} backend does not support interval body velocity perturbation" - ) - if ( - plan.body_angular_velocity_delta is not None - and not self._capabilities.supports_interval_body_angular_velocity_delta - ): - raise NotImplementedError( - f"{self._env._backend.backend_type} backend does not support interval body angular velocity perturbation" - ) - if plan.body_force is not None and not self._capabilities.supports_interval_body_force: - raise NotImplementedError( - f"{self._env._backend.backend_type} backend does not support interval body force perturbation" - ) - if plan.body_torque is not None and not self._capabilities.supports_interval_body_torque: + unsupported = self._capabilities.get_unsupported_interval_terms( + op.term for op in plan.iter_ops() + ) + if unsupported: raise NotImplementedError( - f"{self._env._backend.backend_type} backend does not support interval body torque perturbation" + f"{self._env._backend.backend_type} backend does not support " + f"interval terms: {', '.join(sorted(unsupported))}" ) self._env._backend.apply_interval_randomization(plan) diff --git a/src/unilab/dr/provider.py b/src/unilab/dr/provider.py index 2b6352be9..ea20cba23 100644 --- a/src/unilab/dr/provider.py +++ b/src/unilab/dr/provider.py @@ -33,4 +33,14 @@ def build_reset_observation( def build_interval_randomization_plan( self, env: Any, step_counter: int ) -> IntervalRandomizationPlan | None: + """Build the interval randomization plan for the upcoming step. + + Populate the plan's ``ops`` tuple with ``IntervalTermOp`` entries. + Returning plans via the legacy fields (``push_perturbation_limit``, + ``body_ids``, ``body_linear_velocity_delta``, + ``body_angular_velocity_delta``, ``body_force``, ``body_torque``) is + deprecated: they are still adapted 1:1 through + ``IntervalRandomizationPlan.iter_ops()``, but they will be removed in + the next unisim-core major release. + """ return None diff --git a/src/unilab/tasks/manipulation/sharpa_inhand/rotation.py b/src/unilab/tasks/manipulation/sharpa_inhand/rotation.py index 236be6a24..e940b6351 100644 --- a/src/unilab/tasks/manipulation/sharpa_inhand/rotation.py +++ b/src/unilab/tasks/manipulation/sharpa_inhand/rotation.py @@ -20,11 +20,13 @@ from unilab.base.backend_factory import create_backend, env_backend_kwargs from unilab.base.np_env import NpEnvState from unilab.dr import ( + INTERVAL_TERM_BODY_FORCE, DomainRandomizationCapabilities, DomainRandomizationProvider, GeomSizeOverride, InitRandomizationPlan, IntervalRandomizationPlan, + IntervalTermOp, ModelVariantSpec, ResetPlan, ) @@ -126,7 +128,7 @@ def validate(self, env: Any, capabilities: DomainRandomizationCapabilities) -> N if ( domain_rand is not None and domain_rand.force_scale > 0.0 - and not capabilities.supports_interval_body_force + and not capabilities.supports_interval_term(INTERVAL_TERM_BODY_FORCE) ): raise NotImplementedError( f"{env._backend.backend_type} backend does not support interval body force perturbation" @@ -436,8 +438,13 @@ def build_interval_randomization_plan( ) return IntervalRandomizationPlan( - body_ids=np.asarray([env._object_body_id], dtype=np.int32), - body_force=env._random_object_force[:, None, :].copy(), + ops=( + IntervalTermOp( + INTERVAL_TERM_BODY_FORCE, + env._random_object_force[:, None, :].copy(), + body_ids=np.asarray([env._object_body_id], dtype=np.int32), + ), + ), ) diff --git a/tests/base/test_genesis_backend.py b/tests/base/test_genesis_backend.py index 561f48cd5..b7fe99fa7 100644 --- a/tests/base/test_genesis_backend.py +++ b/tests/base/test_genesis_backend.py @@ -468,9 +468,13 @@ def test_interval_randomization_and_body_force(fake_genesis, tiny_model_file: st assert solver.external_forces[0][1] == [1] assert solver.external_forces[1][1] == [3] - with pytest.raises(NotImplementedError, match="interval push"): - backend.apply_interval_randomization(IntervalRandomizationPlan(push_perturbation_limit=5.0)) - with pytest.raises(NotImplementedError, match="body-velocity"): + with pytest.raises(NotImplementedError, match="does not support interval term 'push'"): + backend.apply_interval_randomization( + IntervalRandomizationPlan(push_perturbation_limit=np.ones(3)) + ) + with pytest.raises( + NotImplementedError, match="does not support interval term 'body_linear_velocity_delta'" + ): backend.apply_interval_randomization( IntervalRandomizationPlan( body_ids=np.array([1], dtype=np.int32), diff --git a/tests/base/test_isaacgym_backend.py b/tests/base/test_isaacgym_backend.py index d8e6e9ee6..352f280b9 100644 --- a/tests/base/test_isaacgym_backend.py +++ b/tests/base/test_isaacgym_backend.py @@ -26,6 +26,7 @@ resolve_isaacgym_runtime, ) from unisim.backend.isaacgym.sensors import scan_scene_metadata +from unisim.dr.types import IntervalTermOp from unilab.base.backend_factory import create_backend from unilab.base.scene import SceneCfg @@ -678,7 +679,10 @@ class _Plan: def is_empty(self) -> bool: return False - with pytest.raises(NotImplementedError, match="interval randomization"): + def iter_ops(self) -> tuple: + return (IntervalTermOp("custom_unsupported", np.zeros(3)),) + + with pytest.raises(NotImplementedError, match="does not support interval term"): backend.apply_interval_randomization(_Plan()) class _Payload: diff --git a/tests/base/test_sim_backend_smoke.py b/tests/base/test_sim_backend_smoke.py index 1540703cf..40eeab175 100644 --- a/tests/base/test_sim_backend_smoke.py +++ b/tests/base/test_sim_backend_smoke.py @@ -229,6 +229,13 @@ def test_mujoco_interval_root_velocity_kick_rejects_invalid_contracts(): ) ) with pytest.raises(ValueError, match=r"shape .* expected \(2, 1, 3\)"): + bkd.apply_interval_randomization( + IntervalRandomizationPlan( + body_ids=body_ids, + body_linear_velocity_delta=np.zeros((NUM_ENVS + 1, 1, 3), dtype=np.float64), + ) + ) + with pytest.raises(ValueError, match="payload must have ndim 3"): bkd.apply_interval_randomization( IntervalRandomizationPlan( body_ids=body_ids, diff --git a/tests/dr/test_manager.py b/tests/dr/test_manager.py index f35fba927..018164304 100644 --- a/tests/dr/test_manager.py +++ b/tests/dr/test_manager.py @@ -1,6 +1,7 @@ from __future__ import annotations import logging +import pickle from dataclasses import dataclass from types import SimpleNamespace from typing import Any @@ -17,9 +18,13 @@ ) from unilab.dr import ( + INTERVAL_TERM_BODY_FORCE, + INTERVAL_TERM_PUSH, DomainRandomizationCapabilities, DomainRandomizationManager, DomainRandomizationProvider, + IntervalRandomizationPlan, + IntervalTermOp, ResetPlan, ResetRandomizationPayload, ) @@ -122,6 +127,7 @@ class _FakeBackend: def __post_init__(self) -> None: self.last_randomization: ResetRandomizationPayload | None = None + self.interval_plans: list[IntervalRandomizationPlan] = [] def get_dr_capabilities(self) -> DomainRandomizationCapabilities: return self.capabilities @@ -135,6 +141,9 @@ def set_state( ) -> None: self.last_randomization = randomization + def apply_interval_randomization(self, plan: IntervalRandomizationPlan) -> None: + self.interval_plans.append(plan) + @dataclass class _FakeTimedBackend: @@ -289,3 +298,170 @@ def test_manager_tolerates_missing_or_malformed_backend_timing(): # Malformed value dropped, well-formed one merged. assert "set_state_mask_ms" not in timings assert timings["set_state_data_slice_ms"] == pytest.approx(0.5) + + +class _FakeIntervalProvider(_FakeProvider): + def __init__(self, plan: IntervalRandomizationPlan | None) -> None: + self._plan = plan + + def build_interval_randomization_plan( + self, env: Any, step_counter: int + ) -> IntervalRandomizationPlan | None: + return self._plan + + +def _interval_manager( + plan: IntervalRandomizationPlan | None, + capabilities: DomainRandomizationCapabilities, +) -> tuple[DomainRandomizationManager, _FakeBackend]: + backend = _FakeBackend(capabilities=capabilities) + env = SimpleNamespace(_backend=backend) + manager = DomainRandomizationManager(env, _FakeIntervalProvider(plan)) + return manager, backend + + +def test_manager_dispatches_custom_interval_term_without_manager_change(): + """A backend-owned custom term flows through the generic manager dispatch: + declaring it in ``supported_interval_terms`` is enough, no manager edit.""" + plan = IntervalRandomizationPlan( + ops=(IntervalTermOp("custom_shake", np.zeros((2, 3), dtype=np.float64)),) + ) + capabilities = DomainRandomizationCapabilities( + supported_interval_terms=frozenset({"custom_shake"}) + ) + manager, backend = _interval_manager(plan, capabilities) + + manager.apply_interval_randomization_if_due(step_counter=10) + + assert backend.interval_plans == [plan] + + +def test_manager_rejects_interval_term_missing_from_capabilities(): + plan = IntervalRandomizationPlan( + ops=(IntervalTermOp("custom_shake", np.zeros((2, 3), dtype=np.float64)),) + ) + manager, backend = _interval_manager(plan, DomainRandomizationCapabilities()) + + with pytest.raises(NotImplementedError) as excinfo: + manager.apply_interval_randomization_if_due(step_counter=10) + + assert "custom_shake" in str(excinfo.value) + assert backend.backend_type in str(excinfo.value) + assert backend.interval_plans == [] + + +def test_manager_dispatches_legacy_fields_plan_via_capability_bools(): + """Legacy-field plans are still adapted through ``iter_ops()`` and checked + against the deprecated legacy capability bools.""" + plan = IntervalRandomizationPlan( + push_perturbation_limit=np.asarray([10.0, 10.0, 5.0]), + body_ids=np.asarray([3], dtype=np.int32), + body_force=np.zeros((4, 1, 3), dtype=np.float64), + ) + capabilities = DomainRandomizationCapabilities( + supports_interval_push=True, + supports_interval_body_force=True, + ) + manager, backend = _interval_manager(plan, capabilities) + + manager.apply_interval_randomization_if_due(step_counter=10) + + assert backend.interval_plans == [plan] + + +def test_manager_skips_none_and_empty_interval_plans(): + capabilities = DomainRandomizationCapabilities( + supported_interval_terms=frozenset({INTERVAL_TERM_PUSH}) + ) + none_manager, none_backend = _interval_manager(None, capabilities) + none_manager.apply_interval_randomization_if_due(step_counter=10) + assert none_backend.interval_plans == [] + + empty_manager, empty_backend = _interval_manager(IntervalRandomizationPlan(), capabilities) + empty_manager.apply_interval_randomization_if_due(step_counter=10) + assert empty_backend.interval_plans == [] + + +def test_manager_dispatches_mixed_legacy_and_ops_plan(): + plan = IntervalRandomizationPlan( + push_perturbation_limit=np.asarray([10.0, 10.0, 5.0]), + ops=(IntervalTermOp("custom_shake", np.zeros((2, 3), dtype=np.float64)),), + ) + capabilities = DomainRandomizationCapabilities( + supports_interval_push=True, + supported_interval_terms=frozenset({"custom_shake"}), + ) + manager, backend = _interval_manager(plan, capabilities) + + manager.apply_interval_randomization_if_due(step_counter=10) + + assert backend.interval_plans == [plan] + + +def test_manager_detects_unsupported_terms_from_both_representations(): + capabilities = DomainRandomizationCapabilities( + supports_interval_push=True, + supported_interval_terms=frozenset({"custom_shake"}), + ) + # Legacy-derived term missing from capabilities. + legacy_plan = IntervalRandomizationPlan( + body_ids=np.asarray([0], dtype=np.int32), + body_torque=np.zeros((2, 1, 3), dtype=np.float64), + ) + manager, backend = _interval_manager(legacy_plan, capabilities) + with pytest.raises(NotImplementedError, match="body_torque"): + manager.apply_interval_randomization_if_due(step_counter=10) + assert backend.interval_plans == [] + + # Explicit op term missing from capabilities. + ops_plan = IntervalRandomizationPlan( + push_perturbation_limit=np.asarray([1.0, 1.0, 1.0]), + ops=(IntervalTermOp("custom_twist", np.zeros((2, 3), dtype=np.float64)),), + ) + manager, backend = _interval_manager(ops_plan, capabilities) + with pytest.raises(NotImplementedError, match="custom_twist"): + manager.apply_interval_randomization_if_due(step_counter=10) + assert backend.interval_plans == [] + + +def test_manager_dispatches_multi_op_plan_in_one_backend_call(): + plan = IntervalRandomizationPlan( + ops=( + IntervalTermOp(INTERVAL_TERM_PUSH, np.asarray([10.0, 10.0, 5.0])), + IntervalTermOp( + INTERVAL_TERM_BODY_FORCE, + np.zeros((4, 1, 3), dtype=np.float64), + body_ids=np.asarray([3], dtype=np.int32), + ), + ) + ) + capabilities = DomainRandomizationCapabilities( + supported_interval_terms=frozenset({INTERVAL_TERM_PUSH, INTERVAL_TERM_BODY_FORCE}) + ) + manager, backend = _interval_manager(plan, capabilities) + + manager.apply_interval_randomization_if_due(step_counter=10) + + assert backend.interval_plans == [plan] + + +def test_interval_plan_with_ops_pickle_round_trip(): + plan = IntervalRandomizationPlan( + ops=( + IntervalTermOp(INTERVAL_TERM_PUSH, np.asarray([10.0, 10.0, 5.0])), + IntervalTermOp( + "custom_shake", + np.ones((2, 3), dtype=np.float64), + body_ids=np.asarray([1, 2], dtype=np.int32), + ), + ) + ) + + restored = pickle.loads(pickle.dumps(plan, protocol=4)) + + assert [op.term for op in restored.ops] == [INTERVAL_TERM_PUSH, "custom_shake"] + np.testing.assert_array_equal(restored.ops[0].payload, plan.ops[0].payload) + np.testing.assert_array_equal(restored.ops[1].payload, plan.ops[1].payload) + assert restored.ops[0].body_ids is None + assert restored.ops[1].body_ids is not None + np.testing.assert_array_equal(restored.ops[1].body_ids, plan.ops[1].body_ids) diff --git a/uv.lock b/uv.lock index 13f18a6b1..24116fdbe 100644 --- a/uv.lock +++ b/uv.lock @@ -5054,7 +5054,7 @@ requires-dist = [ { name = "trimesh", marker = "extra == 'viser'", specifier = ">=3.21.7" }, { name = "typing-extensions" }, { name = "unilab-rl", specifier = "==1.0.0" }, - { name = "unisim-core", specifier = ">=0.1.14" }, + { name = "unisim-core", specifier = ">=1.1.0" }, { name = "viser", marker = "extra == 'viser'", specifier = ">=1.0.26" }, { name = "wandb" }, { name = "warp-lang", marker = "extra == 'mjwarp'", specifier = ">=1.14.0" }, @@ -5097,13 +5097,13 @@ wheels = [ [[package]] name = "unisim-core" -version = "0.1.14" +version = "1.1.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/6e/c9/f3a8b96037e41d7af6cf7d676650eff0f5078133937c02617e26f171cfbf/unisim_core-0.1.14.tar.gz", hash = "sha256:6290c2d1b440921ebd8a51bfd7160677df4897a0d26614d005ad1d7df95df23b", size = 187141, upload-time = "2026-09-02T14:59:15.011Z" } +sdist = { url = "https://files.pythonhosted.org/packages/df/0f/b0a3eb98522ebc6ba761124ac9dfb703fa0416413b203d8eae90bfabaeef/unisim_core-1.1.0.tar.gz", hash = "sha256:d01844cb189de027f3518139984e8ef1a10322021a35fa36482d95fc352dd98d", size = 193709, upload-time = "2026-09-05T09:32:54.183Z" } [[package]] name = "urllib3" diff --git a/uv.rocm.lock b/uv.rocm.lock index 53c3a5f9d..b05e1190d 100644 --- a/uv.rocm.lock +++ b/uv.rocm.lock @@ -3792,7 +3792,7 @@ requires-dist = [ { name = "triton-rocm", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'", specifier = "==3.6.0", index = "https://download.pytorch.org/whl/rocm7.2" }, { name = "typing-extensions" }, { name = "unilab-rl", specifier = "==0.2.0" }, - { name = "unisim-core", specifier = ">=0.1.14" }, + { name = "unisim-core", specifier = ">=1.1.0" }, { name = "viser", marker = "extra == 'viser'", specifier = ">=1.0.26" }, { name = "wandb" }, { name = "wheel", marker = "extra == 'mujoco'" }, @@ -3833,13 +3833,13 @@ wheels = [ [[package]] name = "unisim-core" -version = "0.1.14" +version = "1.1.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/6e/c9/f3a8b96037e41d7af6cf7d676650eff0f5078133937c02617e26f171cfbf/unisim_core-0.1.14.tar.gz", hash = "sha256:6290c2d1b440921ebd8a51bfd7160677df4897a0d26614d005ad1d7df95df23b", size = 187141, upload-time = "2026-09-02T14:59:15.011Z" } +sdist = { url = "https://files.pythonhosted.org/packages/df/0f/b0a3eb98522ebc6ba761124ac9dfb703fa0416413b203d8eae90bfabaeef/unisim_core-1.1.0.tar.gz", hash = "sha256:d01844cb189de027f3518139984e8ef1a10322021a35fa36482d95fc352dd98d", size = 193709, upload-time = "2026-09-05T09:32:54.183Z" } [[package]] name = "urllib3"