From 4315086194ef09b41373c4664da236df5f3cdd60 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Mon, 17 Aug 2026 18:24:16 +0800 Subject: [PATCH 001/169] docs: define NumPy manager API contract (#1043) (#1044) --- docs/sphinx/source/adr/ADR-0000-index.md | 1 + ...-community-manager-api-on-numpy-runtime.md | 245 ++++++++++++++++++ docs/sphinx/source/adr/README.md | 1 + docs/sphinx/source/conf.py | 3 + .../1-architecture/0-index.md | 7 + .../1-architecture/6-manager_based_api.md | 29 +++ .../1-architecture/0-index.md | 7 + .../1-architecture/6-manager_based_api.md | 23 ++ 8 files changed, 316 insertions(+) create mode 100644 docs/sphinx/source/adr/ADR-0006-community-manager-api-on-numpy-runtime.md create mode 100644 docs/sphinx/source/en/4-developer_guide/1-architecture/6-manager_based_api.md create mode 100644 docs/sphinx/source/zh_CN/4-developer_guide/1-architecture/6-manager_based_api.md diff --git a/docs/sphinx/source/adr/ADR-0000-index.md b/docs/sphinx/source/adr/ADR-0000-index.md index 007ef110a..772be3814 100644 --- a/docs/sphinx/source/adr/ADR-0000-index.md +++ b/docs/sphinx/source/adr/ADR-0000-index.md @@ -19,6 +19,7 @@ orphan: true | [ADR-0003 Task Owner And Config Compose Contract](ADR-0003-task-owner-and-config-compose-contract.md) | Config owner | Accepted | | [ADR-0004 Registry Bootstrap Contract](ADR-0004-registry-bootstrap-contract.md) | Registry bootstrap | Accepted | | [ADR-0005 Unified Obs Critic Env And IPC Contract](ADR-0005-unified-obs-critic-env-and-ipc-contract.md) | Observation / IPC | Accepted | +| [ADR-0006 Community Manager API On NumPy Runtime](ADR-0006-community-manager-api-on-numpy-runtime.md) | Manager API / NumPy runtime | Accepted | ## ADR Governance diff --git a/docs/sphinx/source/adr/ADR-0006-community-manager-api-on-numpy-runtime.md b/docs/sphinx/source/adr/ADR-0006-community-manager-api-on-numpy-runtime.md new file mode 100644 index 000000000..51291ede7 --- /dev/null +++ b/docs/sphinx/source/adr/ADR-0006-community-manager-api-on-numpy-runtime.md @@ -0,0 +1,245 @@ +--- +orphan: true +--- + +# ADR-0006 Community Manager API On NumPy Runtime + +语言: 简体中文 + +- Status: Accepted +- Date: 2026-08-17 +- Owners: Env / Config / Backend maintainers +- Supersedes: None +- Superseded by: None + +## Context + +UniLab 的 `NpEnv`、Hydra owner YAML、registry、`SimBackend` 与 heterogeneous +training runtime 已形成稳定 contract，但 task 的 observation、action、reward、 +termination、event、command 与 curriculum 仍主要由各 env 的私有方法组装。用户迁移 +Isaac Lab 或 mjlab task 时，需要重写 manager term、配置和 lifecycle。 + +本决策采用 mjlab v1.6.0 的 manager package 作为可逐文件审查的迁移基线： + +- repository: `mujocolab/mjlab` +- commit: `0fb8a681136be94ffc636a3dd423cabb97d91f10` +- source: `src/mjlab/managers/` 的 12 个 Python 文件 +- license: Apache-2.0；上游 `LICENSE` 声明 + `Copyright 2025, The mjlab Developers` + +该基线只定义 manager-facing API 与语义。它不把 mjlab 的 Torch、Warp、scene +composer、viewer、simulation 或 training runtime 带入 UniLab，也不恢复或参考 UniLab +历史上的 Manager-Based API 实现。 + +## Decision + +### 1. Source-aligned public surface + +`src/unilab/managers/` 按 pinned mjlab package 的模块职责和 exports 直接迁移。以下名称 +是 canonical public surface；Torch 类型替换为 NumPy 类型不构成改名： + +| Module | Canonical exports | +| --- | --- | +| `manager_base` | `ManagerBase`, `ManagerTermBase`, `ManagerTermBaseCfg` | +| `action_manager` | `ActionManager`, `ActionTerm`, `ActionTermCfg` | +| `observation_manager` | `ObservationManager`, `ObservationGroupCfg`, `ObservationTermCfg` | +| `reward_manager` | `RewardManager`, `RewardTermCfg` | +| `termination_manager` | `TerminationManager`, `TerminationTermCfg` | +| `event_manager` | `EventManager`, `EventMode`, `EventTermCfg` | +| `command_manager` | `CommandManager`, `CommandTerm`, `CommandTermCfg`, `NullCommandManager` | +| `curriculum_manager` | `CurriculumManager`, `CurriculumTermCfg`, `NullCurriculumManager` | +| `metrics_manager` | `MetricsManager`, `MetricsTermCfg`, `NullMetricsManager` | +| `recorder_manager` | `RecorderManager`, `RecorderTerm`, `RecorderTermCfg`, `NullRecorderManager` | +| `scene_entity_config` | `SceneEntityCfg` | + +Manager cfg 使用 plain dataclass instance；term 集合使用保持插入顺序的 typed `dict`。 +`func + params`、function/class term、class term 的 `(cfg, env)` 构造和局部 +`reset(env_ids)` 语义保持不变。显式空配置或 term 值为 `None` 表示用户选择禁用，允许 +使用 upstream Null manager/no-op 语义。 + +未来 env lifecycle 的 canonical 名称沿用迁移源的 `ManagerBasedRlEnv` 与 +`ManagerBasedRlEnvCfg`。如果为 Isaac Lab 拼写提供 `ManagerBasedRLEnv` / +`ManagerBasedRLEnvCfg`，它们必须是同一对象的无分支 alias，不能形成第二套实现。 +其他别名必须由实际 migration fixture 证明有价值，不能预先扩张 API。 + +### 2. NumPy runtime boundary + +Manager-facing tensor、buffer、term return、env IDs 和 entity view 使用 +`np.ndarray` 或 `slice`。Torch 的 `device`、`.to()`、`.cpu()` 与 Tensor-only API 不属于 +UniLab manager contract；manager package 不能 import Torch、runner、learner 或 IPC。 + +数值转换保持下列语义： + +- shape、dtype 和更新时序与上游一致；action history 与 observation history 不改变顺序； +- buffer 在 manager 构造或 reset owner 边界分配，step 热路径复用； +- 随机采样使用由 env 拥有并可复现的 NumPy generator，不依赖进程全局 RNG； +- shape 不匹配以及非有限 term 输出在最近 manager/term 边界直接报错；reward 不使用 + `nan_to_num` 把非法值静默变成零； +- observation 明确配置的 noise/delay/history/NaN policy 可以保留，但默认不能掩盖非法 + 输出。 + +### 3. UniLab env、config 与 IPC boundary + +Managers 只依赖一个 typed env context。P0 context 包含 `num_envs`、physics/control dt、 +episode counters、NumPy RNG、各 manager 属性，以及正式 scene/entity facade；不能要求 +`device` 或 backend 私有对象。 + +Manager 内可以使用社区常见的 `policy` / `actor` / `critic` observation group。env owner +必须显式把 actor-facing group 映射为 `NpEnvState.obs["obs"]`，并把可选 critic group +映射为 `NpEnvState.obs["critic"]`。runner、learner 与 IPC 不推断、不拼接 group。 +`reset() -> (obs_dict, info_dict)`、final observation 与 `obs_groups_spec` 保持现有 contract。 + +Callable 和 typed config instance 由 task-owned Python factory 声明；Hydra owner YAML 只按 +manager/term 名覆盖可序列化字段。现有 CLI、registry、algorithm YAML 和 +`task=/` owner compose 保持唯一配置入口。Scripts 不解释 term 业务规则。 + +### 4. Scene/entity owner boundary + +`SceneEntityCfg` 和 term 所需的最小 NumPy entity facade 属于 `src/unilab/base/` 公共 +contract；backend 负责通过 `SimBackend` materialize 名称、ID 和 state/control view，env +负责把 facade 组合进 manager context。该决策解决 #586 的 owner 问题，但不引入完整 +scene composer 或通用 asset hierarchy。 + +- entity name 以及 joint/body/geom/site/actuator selector 在 init/materialization/cache + 冷路径解析一次；热路径只持有已解析 `list[int]`、`np.ndarray` 或 `slice`； +- selector 保留上游 name/regex、`preserve_order`、names/IDs consistency check 和全选压缩为 + `slice(None)` 的语义； +- root/joint/body/site/geom/control 能力只能来自 `SimBackend` 已声明方法；不暴露 backend + model/data 私有对象； +- tendon/camera/light/material/texture/pair 等迁移表面可以存在，但 backend 未声明能力时在 + resolve/materialization 直接 `NotImplementedError`，不能返回空 ID 或跳过； +- 新 backend 能力必须作为独立 child 扩展 `SimBackend` 并补 conformance tests，不能在 + manager 或 env 中用 `getattr` / `hasattr` 探测私有实现。 + +### 5. Fail-closed capability rule + +用户显式禁用与实现缺失是两种不同状态。前者允许 Null manager；后者必须失败： + +| Failure | Required behavior | +| --- | --- | +| cfg/term 类型错误、签名或 shape 不匹配 | `TypeError` / `ValueError`，包含 manager 与 term | +| term 输出 NaN/Inf | `ValueError`，包含 manager、group/term 与非法值类别 | +| backend/entity capability 未实现 | `NotImplementedError`，包含 manager、term、capability 与 backend | +| selector name/ID 不存在或不一致 | `KeyError` / `ValueError`，包含 entity 与 selector | + +不得 warning 后 skip、返回零/旧值、自动换 backend、禁用 feature 或回退到旧 env。当前 +不新增公共 exception hierarchy；只有 consumer 证明需要 machine-readable 分类时再单独 +决策。 + +### 6. Performance and deletion policy + +优先级固定为：社区 Manager-Based API 语义与结构一致性，优先于改变公共设计的局部性能 +优化。在此约束下，生产级 NumPy 热路径不能引入明显可避免的重复解析、逐环境 Python +循环、数组复制或临时分配。优化必须由同配置、同硬件 benchmark 证明有足够收益，并优先 +保持在内部预解析、预分配和批量 NumPy 实现；低收益但增加专用 fast path、缓存协议或长期 +复杂度的方案不采用。 + +Production task 迁移后必须在同一 task-family child 删除被替代的旧 dispatch、重复 +reward/config helper 和 bridge。Umbrella 完成时只保留一套 manager lifecycle，不保留 +fallback 到旧单体 env 的永久兼容路径。 + +## Stable Contracts + +### Compatibility matrix + +状态只表示本 ADR 固定的迁移目标；实际 support claim 仍需要代码、注册、配置和测试证据。 + +| Surface | Target | Notes | +| --- | --- | --- | +| manager modules、class/config names、dict order | Compatible | 直接保留 pinned mjlab 1.6.0 表面 | +| function/class term、`params`、local reset | Compatible | class term 在冷路径实例化 | +| action split/apply/history、reward dt scaling、termination timeout split | Compatible | NumPy 实现保持时序 | +| observation groups、clip/scale/noise/delay/history | Adapted | 数值为 NumPy；group 在 env boundary 显式映射 | +| manager buffers、env IDs、RNG | Adapted | Torch→NumPy；无 device API | +| `ManagerBasedRlEnv` return | Adapted | 保留 `NpEnvState` 与 UniLab reset/final-observation contract | +| config container | Adapted | plain instances + Hydra owner YAML overlay，不引入第二套 runtime | +| `SceneEntityCfg` selectors | Adapted | 语义保留；只解析 `SimBackend` 已声明能力 | +| event/domain randomization | Adapted | 调度语义保留；mutation 走 backend DR/capability contract | +| Metrics/Recorder | Adapted | lifecycle hook 存在时启用；缺失时显式失败或显式空配置 | +| Torch device、Warp mutation、viewer glue | Unsupported | 不进入 manager core，不提供静默替代 | +| Omniverse/USD/mjlab Scene/Simulation | Unsupported | 不属于 UniLab runtime | + +### Mechanical migration example + +迁移前的 mjlab term： + +```python +import torch +from mjlab.managers import RewardTermCfg + +def joint_error(env) -> torch.Tensor: + return torch.square(env.joint_pos - env.target_joint_pos).sum(dim=1) + +term = RewardTermCfg(func=joint_error, weight=-1.0) +``` + +迁移后的 UniLab term 只改 import、数值类型和对应 NumPy 运算： + +```python +import numpy as np +from unilab.managers import RewardTermCfg + +def joint_error(env) -> np.ndarray: + return np.square(env.joint_pos - env.target_joint_pos).sum(axis=1) + +term = RewardTermCfg(func=joint_error, weight=-1.0) +``` + +如果迁移还要求重写 term 结构、增加 backend 分支或改 runner/IPC，说明 adapter boundary +不够薄，必须停止并拆出 owner child。 + +### Provenance and change accounting + +每个 source-derived Python 文件必须注明上游 repository、tag/commit、原始路径、 +Apache-2.0 和 UniLab 的修改类别。实现 PR 分别报告： + +1. source-derived：保留的上游结构/语义； +2. mechanical：import、typing、Torch→NumPy 和格式转换； +3. UniLab-specific glue：新 facade、contract adapter 或行为； +4. deleted：删除的上游不适用代码和 UniLab 旧实现。 + +不得通过重新分类隐藏 glue 超预算；不建立长期 upstream mirror 或自动 sync tooling。 + +## Alternatives Considered + +- 重新设计一套更适合 UniLab 的 managers，再提供兼容 facade。拒绝：会形成 + UniLab-only 方言和两套行为，增加用户迁移与长期维护成本。 +- 在 manager 热路径保留 Torch。拒绝：破坏 NumPy runtime、backend isolation 与 + heterogeneous CPU physics → accelerator learner 数据面。 +- 一次迁移完整 mjlab scene/simulation/entity runtime。拒绝：复制第二套 backend/scene + abstraction，并引入 Warp/MuJoCo/Viewer 假设。 +- 先设计 compiler、fused term protocol 或专用 fast path。拒绝：在 benchmark 证明瓶颈前 + 增加结构复杂度，并可能牺牲社区 term 语义。 +- 缺失能力 warning + skip 或回退旧 env。拒绝：配置表面与真实执行不一致，不能用于生产。 + +## Consequences + +- Manager port 的审查基线是 pinned upstream diff，而不是重新解释每个 manager 的职责。 +- NumPy、UniLab env/config contract 和显式 unsupported 是允许的偏离；其他偏离必须在 + compatibility matrix 中先记录。 +- Scene/entity 采用最小 base facade，#586 不再阻塞 manager port；真实 backend 能力仍按 + 独立 child 和 conformance evidence 接入。 +- 迁移初期允许 production 旧 task 与未接入的 manager package 同时存在，但 task 一旦迁移 + 就必须删除对应旧实现；umbrella 结束时不能保留双 lifecycle。 +- 性能 gate 关注明显低效与实测瓶颈，不以复杂度换取未经证明的小收益。 + +## Evidence In Repo + +- Env contract: `src/unilab/base/np_env.py` +- Backend contract: `src/unilab/base/backend/base.py` +- Scene config owner: `src/unilab/base/scene.py` +- Config schema and registry: `src/unilab/structured_configs.py`, + `src/unilab/base/registry.py`, `conf/` +- Observation/IPC contract: `docs/sphinx/source/adr/ADR-0005-unified-obs-critic-env-and-ipc-contract.md` +- Layer boundary: `docs/sphinx/source/adr/ADR-0001-runtime-model-and-layer-boundaries.md` +- Upstream checkout used for the decision: + `/home/user/ws/simulator/mjlab/src/mjlab/managers/` at `0fb8a681` + +## Related Documents + +- {doc}`ADR Index ` +- {doc}`Manager-Based API contract ` +- {doc}`RL Infrastructure 开发标准 ` +- [Roadmap #1042](https://github.com/unilabsim/UniLab/issues/1042) +- [Implementation issue #1043](https://github.com/unilabsim/UniLab/issues/1043) +- [Entity abstraction decision #586](https://github.com/unilabsim/UniLab/issues/586) diff --git a/docs/sphinx/source/adr/README.md b/docs/sphinx/source/adr/README.md index 007ef110a..772be3814 100644 --- a/docs/sphinx/source/adr/README.md +++ b/docs/sphinx/source/adr/README.md @@ -19,6 +19,7 @@ orphan: true | [ADR-0003 Task Owner And Config Compose Contract](ADR-0003-task-owner-and-config-compose-contract.md) | Config owner | Accepted | | [ADR-0004 Registry Bootstrap Contract](ADR-0004-registry-bootstrap-contract.md) | Registry bootstrap | Accepted | | [ADR-0005 Unified Obs Critic Env And IPC Contract](ADR-0005-unified-obs-critic-env-and-ipc-contract.md) | Observation / IPC | Accepted | +| [ADR-0006 Community Manager API On NumPy Runtime](ADR-0006-community-manager-api-on-numpy-runtime.md) | Manager API / NumPy runtime | Accepted | ## ADR Governance diff --git a/docs/sphinx/source/conf.py b/docs/sphinx/source/conf.py index b818430fa..65effaca0 100644 --- a/docs/sphinx/source/conf.py +++ b/docs/sphinx/source/conf.py @@ -276,6 +276,9 @@ # map is computed below. _LANGUAGE_PATH_FORWARD: dict[str, str] = { "en/1-getting_started/5-faq": "zh_CN/1-getting_started/5-faq", + "en/4-developer_guide/1-architecture/6-manager_based_api": ( + "zh_CN/4-developer_guide/1-architecture/6-manager_based_api" + ), } # Keyed by (current_pagename, target_language) → target_pagename. _LANGUAGE_PATH_MAP: dict[tuple[str, str], str] = {} diff --git a/docs/sphinx/source/en/4-developer_guide/1-architecture/0-index.md b/docs/sphinx/source/en/4-developer_guide/1-architecture/0-index.md index f31456d43..9e48c9faf 100644 --- a/docs/sphinx/source/en/4-developer_guide/1-architecture/0-index.md +++ b/docs/sphinx/source/en/4-developer_guide/1-architecture/0-index.md @@ -36,6 +36,12 @@ Scene fragments, assets, and cold-path materialization. Bootstrap imports and env/backend registration. ::: +:::{grid-item-card} Manager-Based API +:link: 6-manager_based_api +:link-type: doc +Community manager semantics, NumPy runtime, and fail-closed boundaries. +::: + :::: ```{toctree} @@ -46,4 +52,5 @@ Bootstrap imports and env/backend registration. 3-layer_boundaries 4-scene_composition 5-registry +6-manager_based_api ``` diff --git a/docs/sphinx/source/en/4-developer_guide/1-architecture/6-manager_based_api.md b/docs/sphinx/source/en/4-developer_guide/1-architecture/6-manager_based_api.md new file mode 100644 index 000000000..6090d9117 --- /dev/null +++ b/docs/sphinx/source/en/4-developer_guide/1-architecture/6-manager_based_api.md @@ -0,0 +1,29 @@ +# Manager-Based API + +UniLab uses a community-compatible manager API on its NumPy runtime. Manager modules, +term configs, function/class terms, lifecycle ordering, and reset semantics follow the +pinned mjlab 1.6.0 source. Numeric execution uses NumPy while preserving UniLab's +`NpEnvState`, Hydra owner YAML, `SimBackend`, registry, and IPC contracts. + +The normative compatibility matrix and mechanical migration example are in +{doc}`ADR-0006 `. + +## Invariants + +- Community manager semantics and a general structure take priority over local + optimizations that would create a UniLab-only term API. +- Manager buffers, term returns, environment IDs, and entity views use `np.ndarray` + or `slice`; manager core does not depend on Torch, Warp, runners, learners, or IPC. +- `SceneEntityCfg` resolves through a base-owned scene/entity facade on the cold path. + The facade uses only the public `SimBackend` contract, and hot paths reuse cached IDs + and views. +- Explicitly empty configuration may use a Null manager. A requested capability that + is unavailable fails at the nearest boundary; it is never skipped, zero-filled, or + routed back to a legacy environment. +- Hot paths avoid obvious repeated parsing, per-environment Python loops, copies, and + temporary allocations. Further optimization requires benchmark evidence and must + not add disproportionate structural complexity. + +Only surfaces backed by registration, configuration, and tests may be called +Compatible. NumPy/env/config adapters are Adapted; capabilities without a formal +backend contract are Unsupported and fail closed. diff --git a/docs/sphinx/source/zh_CN/4-developer_guide/1-architecture/0-index.md b/docs/sphinx/source/zh_CN/4-developer_guide/1-architecture/0-index.md index 531791167..ee5e0fd29 100644 --- a/docs/sphinx/source/zh_CN/4-developer_guide/1-architecture/0-index.md +++ b/docs/sphinx/source/zh_CN/4-developer_guide/1-architecture/0-index.md @@ -36,6 +36,12 @@ Runner 生命周期、worker/learner 拆分与数据流。 Bootstrap 导入与 env/backend 注册。 ::: +:::{grid-item-card} Manager-Based API +:link: 6-manager_based_api +:link-type: doc +社区 manager 语义、NumPy runtime 与 fail-closed 边界。 +::: + :::: ```{toctree} @@ -46,4 +52,5 @@ Bootstrap 导入与 env/backend 注册。 3-layer_boundaries 4-scene_composition 5-registry +6-manager_based_api ``` diff --git a/docs/sphinx/source/zh_CN/4-developer_guide/1-architecture/6-manager_based_api.md b/docs/sphinx/source/zh_CN/4-developer_guide/1-architecture/6-manager_based_api.md new file mode 100644 index 000000000..9cba1ead7 --- /dev/null +++ b/docs/sphinx/source/zh_CN/4-developer_guide/1-architecture/6-manager_based_api.md @@ -0,0 +1,23 @@ +# Manager-Based API + +UniLab 采用“社区兼容 API + UniLab NumPy runtime”：manager-facing 模块、term cfg、 +function/class term、生命周期和顺序语义以固定的 mjlab 1.6.0 source 为基线；数值实现使用 +NumPy，并保留现有 `NpEnvState`、Hydra owner YAML、`SimBackend`、registry 与 IPC contract。 + +完整决策、兼容矩阵和机械迁移示例见 +{doc}`/adr/ADR-0006-community-manager-api-on-numpy-runtime`。 + +## 不变量 + +- 公共 manager 结构优先保持社区语义；不为局部性能制造 UniLab-only term API。 +- manager buffer、term return、env ID 和 entity view 使用 `np.ndarray` / `slice`，core 不依赖 + Torch、Warp、runner、learner 或 IPC。 +- `SceneEntityCfg` 在冷路径通过 base scene/entity facade 解析；facade 只调用正式 + `SimBackend` contract，热路径复用缓存 ID/view。 +- 用户显式空配置可以使用 Null manager；配置请求但 runtime/backend 不支持的能力必须在 + 最近边界报错，不能 warning、skip、返回零或回退旧 env。 +- 热路径避免明显的重复解析、逐环境 Python 循环、复制和临时分配；进一步优化需要 + benchmark 证明收益，且不能增加不成比例的结构复杂度。 + +只有被注册、配置和测试覆盖的表面才能声明 Compatible。NumPy/env/config adapter 标为 +Adapted；缺少正式 backend contract 的能力标为 Unsupported 并 fail-closed。 From 1145e172ff7e150e2c622bbbd26914ee8906f934 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Mon, 17 Aug 2026 19:06:17 +0800 Subject: [PATCH 002/169] feat(managers): port manager core to numpy (#1046) --- pyproject.toml | 1 + src/unilab/managers/__init__.py | 40 ++ src/unilab/managers/_buffers/__init__.py | 7 + .../managers/_buffers/circular_buffer.py | 258 ++++++++ src/unilab/managers/_buffers/delay_buffer.py | 304 ++++++++++ src/unilab/managers/_noise/__init__.py | 15 + src/unilab/managers/_noise/noise_cfg.py | 143 +++++ src/unilab/managers/_noise/noise_model.py | 97 +++ src/unilab/managers/_types.py | 52 ++ src/unilab/managers/action_manager.py | 219 +++++++ src/unilab/managers/command_manager.py | 313 ++++++++++ src/unilab/managers/curriculum_manager.py | 166 +++++ src/unilab/managers/event_manager.py | 276 +++++++++ src/unilab/managers/manager_base.py | 159 +++++ src/unilab/managers/metrics_manager.py | 242 ++++++++ src/unilab/managers/observation_manager.py | 570 ++++++++++++++++++ src/unilab/managers/recorder_manager.py | 267 ++++++++ src/unilab/managers/reward_manager.py | 154 +++++ src/unilab/managers/scene_entity_config.py | 263 ++++++++ src/unilab/managers/termination_manager.py | 139 +++++ tests/managers/__init__.py | 1 + tests/managers/conftest.py | 58 ++ tests/managers/test_core_managers.py | 195 ++++++ .../test_event_command_metrics_recorder.py | 263 ++++++++ .../test_observation_buffers_noise.py | 205 +++++++ uv.lock | 14 + 26 files changed, 4421 insertions(+) create mode 100644 src/unilab/managers/__init__.py create mode 100644 src/unilab/managers/_buffers/__init__.py create mode 100644 src/unilab/managers/_buffers/circular_buffer.py create mode 100644 src/unilab/managers/_buffers/delay_buffer.py create mode 100644 src/unilab/managers/_noise/__init__.py create mode 100644 src/unilab/managers/_noise/noise_cfg.py create mode 100644 src/unilab/managers/_noise/noise_model.py create mode 100644 src/unilab/managers/_types.py create mode 100644 src/unilab/managers/action_manager.py create mode 100644 src/unilab/managers/command_manager.py create mode 100644 src/unilab/managers/curriculum_manager.py create mode 100644 src/unilab/managers/event_manager.py create mode 100644 src/unilab/managers/manager_base.py create mode 100644 src/unilab/managers/metrics_manager.py create mode 100644 src/unilab/managers/observation_manager.py create mode 100644 src/unilab/managers/recorder_manager.py create mode 100644 src/unilab/managers/reward_manager.py create mode 100644 src/unilab/managers/scene_entity_config.py create mode 100644 src/unilab/managers/termination_manager.py create mode 100644 tests/managers/__init__.py create mode 100644 tests/managers/conftest.py create mode 100644 tests/managers/test_core_managers.py create mode 100644 tests/managers/test_event_command_metrics_recorder.py create mode 100644 tests/managers/test_observation_buffers_noise.py diff --git a/pyproject.toml b/pyproject.toml index c9374d0fe..7b56b062f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,6 +12,7 @@ license-files = ["LICENSE"] requires-python = ">=3.10,<3.14" dependencies = [ "numpy", + "prettytable>=3.10", "torch==2.9.0 ; sys_platform == 'linux' and platform_machine == 'aarch64'", "torch==2.7.0 ; sys_platform != 'linux' or platform_machine != 'aarch64'", "gymnasium", diff --git a/src/unilab/managers/__init__.py b/src/unilab/managers/__init__.py new file mode 100644 index 000000000..be3224090 --- /dev/null +++ b/src/unilab/managers/__init__.py @@ -0,0 +1,40 @@ +# Derived from mujocolab/mjlab v1.6.0 (0fb8a681), src/mjlab/managers/__init__.py. +# Copyright 2025, The mjlab Developers. +# Modified by UniLab for NumPy and UniLab contracts; licensed under Apache-2.0. +"""Environment managers.""" + +from unilab.managers.action_manager import ActionManager as ActionManager +from unilab.managers.action_manager import ActionTerm as ActionTerm +from unilab.managers.action_manager import ActionTermCfg as ActionTermCfg +from unilab.managers.command_manager import CommandManager as CommandManager +from unilab.managers.command_manager import CommandTerm as CommandTerm +from unilab.managers.command_manager import CommandTermCfg as CommandTermCfg +from unilab.managers.command_manager import NullCommandManager as NullCommandManager +from unilab.managers.curriculum_manager import CurriculumManager as CurriculumManager +from unilab.managers.curriculum_manager import CurriculumTermCfg as CurriculumTermCfg +from unilab.managers.curriculum_manager import ( + NullCurriculumManager as NullCurriculumManager, +) +from unilab.managers.event_manager import EventManager as EventManager +from unilab.managers.event_manager import EventMode as EventMode +from unilab.managers.event_manager import EventTermCfg as EventTermCfg +from unilab.managers.manager_base import ManagerBase as ManagerBase +from unilab.managers.manager_base import ManagerTermBase as ManagerTermBase +from unilab.managers.manager_base import ManagerTermBaseCfg as ManagerTermBaseCfg +from unilab.managers.metrics_manager import MetricsManager as MetricsManager +from unilab.managers.metrics_manager import MetricsTermCfg as MetricsTermCfg +from unilab.managers.metrics_manager import NullMetricsManager as NullMetricsManager +from unilab.managers.observation_manager import ( + ObservationGroupCfg as ObservationGroupCfg, +) +from unilab.managers.observation_manager import ObservationManager as ObservationManager +from unilab.managers.observation_manager import ObservationTermCfg as ObservationTermCfg +from unilab.managers.recorder_manager import NullRecorderManager as NullRecorderManager +from unilab.managers.recorder_manager import RecorderManager as RecorderManager +from unilab.managers.recorder_manager import RecorderTerm as RecorderTerm +from unilab.managers.recorder_manager import RecorderTermCfg as RecorderTermCfg +from unilab.managers.reward_manager import RewardManager as RewardManager +from unilab.managers.reward_manager import RewardTermCfg as RewardTermCfg +from unilab.managers.scene_entity_config import SceneEntityCfg as SceneEntityCfg +from unilab.managers.termination_manager import TerminationManager as TerminationManager +from unilab.managers.termination_manager import TerminationTermCfg as TerminationTermCfg diff --git a/src/unilab/managers/_buffers/__init__.py b/src/unilab/managers/_buffers/__init__.py new file mode 100644 index 000000000..2599f33d2 --- /dev/null +++ b/src/unilab/managers/_buffers/__init__.py @@ -0,0 +1,7 @@ +# Derived from mujocolab/mjlab v1.6.0 (0fb8a681), src/mjlab/utils/buffers/__init__.py. +# Copyright 2025, The mjlab Developers. +# Modified by UniLab for NumPy and UniLab contracts; licensed under Apache-2.0. +"""Buffer utilities.""" + +from unilab.managers._buffers.circular_buffer import CircularBuffer as CircularBuffer +from unilab.managers._buffers.delay_buffer import DelayBuffer as DelayBuffer diff --git a/src/unilab/managers/_buffers/circular_buffer.py b/src/unilab/managers/_buffers/circular_buffer.py new file mode 100644 index 000000000..b0575dd50 --- /dev/null +++ b/src/unilab/managers/_buffers/circular_buffer.py @@ -0,0 +1,258 @@ +# Derived from mujocolab/mjlab v1.6.0 (0fb8a681), src/mjlab/utils/buffers/circular_buffer.py. +# Copyright 2025, The mjlab Developers. +# Modified by UniLab for NumPy and UniLab contracts; licensed under Apache-2.0. +"""Circular buffer for storing a history of batched tensor data. + +Understanding Dimensions +======================== + +Internal storage shape: (max_len, batch_size, ...) + ↑ ↑ + time environments + axis axis + +External view (buffer property): (batch_size, max_len, ...) + ↑ ↑ + environments time + axis axis + +Why different? + - Internal: time-first for clean API (self._buffer[pointer] = data) + - External: batch-first for typical use (iterate over environments) + +Backfill Behavior +================= + +When you first append to a new or reset batch row, that first value is copied +to ALL history slots for that row: + + buffer = CircularBuffer(max_len=3, batch_size=2) + buffer.append(np.array([[5.0], [10.0]])) + + # buffer.buffer contains (shape: 2, 3, 1): + # Batch 0: [5.0, 5.0, 5.0] <- all slots filled with first value + # Batch 1: [10.0, 10.0, 10.0] <- all slots filled with first value + +Why backfill? + You always have valid data. If training expects 3 frames of history, you + don't want garbage/zeros for the first 2 timesteps. + +Per-Batch Reset +=============== + +Reset affects specific batch rows. The circular pointer advances globally, but +reset rows get "first-append" treatment on their next write: + + buffer = CircularBuffer(max_len=3, batch_size=3) + + buffer.append(np.array([[1.0], [10.0], [100.0]])) # t0 + buffer.append(np.array([[2.0], [20.0], [200.0]])) # t1 + buffer.append(np.array([[3.0], [30.0], [300.0]])) # t2 + + # buffer.buffer: + # Batch 0: [1.0, 2.0, 3.0] + # Batch 1: [10.0, 20.0, 30.0] + # Batch 2: [100.0, 200.0, 300.0] + + buffer.reset(batch_ids=[1]) # Reset only batch 1 + + # buffer.buffer after reset: + # Batch 0: [1.0, 2.0, 3.0] <- unchanged + # Batch 1: [0.0, 0.0, 0.0] <- zeroed + # Batch 2: [100.0, 200.0, 300.0] <- unchanged + + # current_length: [3, 0, 3] <- batch 1 has 0 valid frames + + buffer.append(np.array([[4.0], [99.0], [400.0]])) # t3 + + # buffer.buffer after append: + # Batch 0: [2.0, 3.0, 4.0] <- oldest overwritten (normal) + # Batch 1: [99.0, 99.0, 99.0] <- BACKFILLED with 99.0 + # Batch 2: [200.0, 300.0, 400.0] <- oldest overwritten (normal) + +Key insight: + Reset only affects specific batch rows. The pointer keeps advancing for + everyone, but reset rows get backfilled on their next append. +""" + +from __future__ import annotations + +from collections.abc import Sequence + +import numpy as np + + +class CircularBuffer: + """Fixed-length circular buffer for batched tensor history. + + Stores history with shape (max_len, batch_size, ...) internally. + The `buffer` property returns chronologically ordered data (oldest to newest) + with shape (batch_size, max_len, ...). + + Storage and Retrieval + --------------------- + Internal storage (circular): + ┌─────┬─────┬─────┐ + │ 2 │ 3 │ 1 │ <- pointer at index 1 (newest = 3) + └─────┴─────┴─────┘ + idx:0 idx:1 idx:2 + + buffer property returns (chronological, oldest→newest): + ┌─────┬─────┬─────┐ + │ 1 │ 2 │ 3 │ + └─────┴─────┴─────┘ + + LIFO Retrieval via __getitem__ + ------------------------------- + Given buffer with [1, 2, 3] (oldest to newest): + buffer[lag=0] -> 3 (most recent) + buffer[lag=1] -> 2 (one step back) + buffer[lag=2] -> 1 (oldest) + + Per-Batch Reset + --------------- + After reset(batch_ids=[1]): + Batch 0: [1, 2, 3] (unchanged) + Batch 1: [0, 0, 0] (zeroed, current_length=0) + Batch 2: [1, 2, 3] (unchanged) + + Next append backfills reset rows with first value. + + Args: + max_len: Maximum number of historical frames to retain. + batch_size: Size of the batch dimension. + """ + + def __init__(self, max_len: int, batch_size: int) -> None: + if max_len < 1: + raise ValueError(f"Buffer size must be >= 1, got {max_len}") + + self._max_len = max_len + self._batch_size = batch_size + self._pointer: int = -1 + self._buffer: np.ndarray | None = None + self._all_indices = np.arange(batch_size) + self._num_pushes = np.zeros(batch_size, dtype=np.int64) + self._max_len_array = np.full(batch_size, max_len, dtype=np.int64) + + @property + def batch_size(self) -> int: + return self._batch_size + + @property + def max_length(self) -> int: + return self._max_len + + @property + def current_length(self) -> np.ndarray: + """Per-batch count of valid frames. Shape: (batch_size,).""" + return np.minimum(self._num_pushes, self._max_len_array) + + @property + def is_initialized(self) -> bool: + """Check if the buffer has been initialized with at least one append.""" + return self._buffer is not None + + @property + def buffer(self) -> np.ndarray: + """History in chronological order (oldest to newest). + + Returns: + Tensor of shape (batch_size, max_len, ...) where index 0 is oldest + and index -1 is newest. + """ + if self._buffer is None: + raise RuntimeError("Buffer not initialized. Call append() first.") + + start = (self._pointer + 1) % self._max_len + idx = (np.arange(self._max_len) + start) % self._max_len + buf = self._buffer[idx] # (max_len, batch, ...) + return np.swapaxes(buf, 0, 1) # (batch, max_len, ...) + + def reset(self, batch_ids: Sequence[int] | np.ndarray | slice | None = None) -> None: + """Zero out values and counters for specified batch rows. + + Args: + batch_ids: Batch indices to reset, or None to reset all. + """ + ids: Sequence[int] | np.ndarray | slice = slice(None) if batch_ids is None else batch_ids + self._num_pushes[ids] = 0 + if self._buffer is not None: + self._buffer[:, ids] = 0.0 + + def backfill(self, data: np.ndarray, batch_ids: np.ndarray) -> None: + """Fill the given rows' entire history with one frame, without advancing time. + + Unlike append, the global pointer does not move and other rows are + untouched. Used after a partial reset: the reset rows get their first + post-reset frame in every slot (the same backfill their next append would + apply) while the remaining rows keep their history intact. + + Args: + data: Tensor of shape (batch_size, ...); only rows at batch_ids are read. + batch_ids: Batch indices to backfill. + """ + if data.shape[0] != self._batch_size: + raise ValueError(f"Expected batch size {self._batch_size}, got {data.shape[0]}") + if self._buffer is None: + raise RuntimeError("Buffer not initialized. Call append() first.") + + data = np.asarray(data) + self._buffer[:, batch_ids] = np.expand_dims(data[batch_ids], axis=0) + self._num_pushes[batch_ids] = 1 + + def append(self, data: np.ndarray) -> None: + """Append a new frame for all batch elements. + + Args: + data: Tensor of shape (batch_size, ...). + """ + if data.shape[0] != self._batch_size: + raise ValueError(f"Expected batch size {self._batch_size}, got {data.shape[0]}") + + data = np.asarray(data) + + if self._buffer is None: + self._pointer = -1 + self._buffer = np.empty((self._max_len, *data.shape), dtype=data.dtype) + + self._pointer = (self._pointer + 1) % self._max_len + self._buffer[self._pointer] = data + + # Backfill only newly initialized rows. After warm-up this branch avoids + # scanning the full history buffer on every hot-path append. + is_first_push = self._num_pushes == 0 + if np.any(is_first_push): + first_ids = np.flatnonzero(is_first_push) + self._buffer[:, first_ids] = np.expand_dims(data[first_ids], axis=0) + + self._num_pushes += 1 + + def __getitem__(self, key: np.ndarray | int) -> np.ndarray: + """Retrieve lagged frames per batch (LIFO). + + Args: + key: Per-batch lags (Tensor) or shared lag (int). Shape (batch_size,) or scalar. + """ + if self._buffer is None: + raise RuntimeError("Buffer not initialized. Call append() first.") + + if isinstance(key, int): + key = np.full(self._batch_size, key, dtype=np.int64) + else: + key = np.asarray(key, dtype=np.int64) + if key.ndim == 0: + key = np.full(self._batch_size, key.item(), dtype=np.int64) + + if key.size != self._batch_size: + raise ValueError(f"Expected {self._batch_size} lags, got {key.size}") + + # Clamp to the oldest retained frame: without the max_len bound, a lag + # beyond the buffer length would wrap around to a newer frame once + # num_pushes exceeds max_len. + pushes = np.maximum(self._num_pushes, 1) + max_lag = np.minimum(pushes, self._max_len_array) - 1 + valid = np.maximum(np.minimum(key, max_lag), 0) + + idx = np.remainder(self._pointer - valid, self._max_len) + return self._buffer[idx, self._all_indices] diff --git a/src/unilab/managers/_buffers/delay_buffer.py b/src/unilab/managers/_buffers/delay_buffer.py new file mode 100644 index 000000000..da29bb29c --- /dev/null +++ b/src/unilab/managers/_buffers/delay_buffer.py @@ -0,0 +1,304 @@ +# Derived from mujocolab/mjlab v1.6.0 (0fb8a681), src/mjlab/utils/buffers/delay_buffer.py. +# Copyright 2025, The mjlab Developers. +# Modified by UniLab for NumPy and UniLab contracts; licensed under Apache-2.0. +"""Delay buffer for stochastically delayed observations.""" + +from __future__ import annotations + +from collections.abc import Sequence + +import numpy as np + +from unilab.managers._buffers import CircularBuffer + + +class DelayBuffer: + """Serve stochastically delayed observations from a rolling history. + + Wraps a CircularBuffer to simulate observation delays by returning frames from T-lag + timesteps ago, where lag is sampled from [min_lag, max_lag]. + + Core Behavior + ============= + + At each timestep: + 1. Append new observation to history + 2. Sample or hold lag value (0 = no delay, 3 = 3 timesteps old) + 3. Return observation from T-lag + + Example with lag=2: + t=0: append obs_0 → return obs_0 (not enough history) + t=1: append obs_1 → return obs_0 (clamped to available history) + t=2: append obs_2 → return obs_0 (lag=2, so T-2 = 0) + t=3: append obs_3 → return obs_1 (lag=2, so T-2 = 1) + + Lag Update Policy + ================= + + Lags can be refreshed every step or periodically: + + **Every-step updates (update_period=0)** + Each timestep may sample a new lag (subject to hold_prob). + + **Periodic updates (update_period=N)** + Lags refresh only every N steps per environment: + if (step_count + phase_offset) % N == 0: + sample new lag + else: + keep previous lag + + **Staggered updates (per_env_phase=True)** + Each environment gets a random phase_offset ∈ [0, N), causing + lag updates to occur on different timesteps: + Env 0: updates at t=0, N, 2N, ... + Env 1: updates at t=3, N+3, 2N+3, ... + Env 2: updates at t=7, N+7, 2N+7, ... + + **Hold probability (hold_prob=0.2)** + Even when an update would occur, keep previous lag with 20% chance. + Creates temporal correlation in delay patterns. + + Per-Environment vs Shared Lags + ============================== + + **per_env=True** (default) + Each environment has independent lag: + Batch 0: lag=1 → returns obs from t-1 + Batch 1: lag=3 → returns obs from t-3 + Batch 2: lag=0 → returns current obs + + **per_env=False** + All environments share one sampled lag: + All batches: lag=2 → all return obs from t-2 + + Reset Behavior + ============== + + reset(batch_ids=[1]) clears history for specified environments: + - Sets lag and step counter to zero + - Clears circular buffer for those rows + - Next append backfills their history with first new value + - Until that append, compute() returns zeros for reset rows + + Args: + min_lag (int, optional): Minimum lag (inclusive). Must be >= 0. + max_lag (int, optional): Maximum lag (inclusive). Must be >= `min_lag`. + batch_size (int, optional): Number of parallel environments (leading + dimension of inputs). + per_env (bool, optional): If True, sample a separate lag per environment; + otherwise sample one lag and share it across environments. + hold_prob (float, optional): Probability in `[0.0, 1.0]` to keep the previous + lag when an update would occur. Creates temporal correlation in delays. + update_period (int, optional): If > 0, refresh lags every N steps per + environment; if 0, consider updating every step. + per_env_phase (bool, optional): If True and `update_period > 0`, each + environment uses a different phase offset in `[0, update_period)`, causing + staggered refresh steps across the batch. + generator (np.random.Generator | None, optional): RNG for sampling lags. + + Examples: + Constant delay (lag = 2): + >>> buf = DelayBuffer(min_lag=2, max_lag=2, batch_size=4) + >>> buf.append(obs) # obs.shape == (4, ...) + >>> delayed = buf.compute() # delayed[t] = obs[t-2] + + Stochastic delay (uniform 0-3): + >>> buf = DelayBuffer(min_lag=0, max_lag=3, batch_size=4) + >>> buf.append(obs) + >>> delayed = buf.compute() # per-env lag sampled in {0,1,2,3} + + Periodic updates with staggering: + >>> buf = DelayBuffer( + ... min_lag=1, max_lag=5, batch_size=8, + ... update_period=10, # refresh every 10 steps + ... per_env_phase=True, # stagger across envs + ... hold_prob=0.2 # 20% chance to hold lag + ... ) + >>> # Env 0 refreshes at t=0,10,20,... + >>> # Env 1 refreshes at t=3,13,23,... (random offset) + >>> # But each refresh has 20% chance to keep previous lag + """ + + def __init__( + self, + min_lag: int = 0, + max_lag: int = 3, + batch_size: int = 1, + per_env: bool = True, + hold_prob: float = 0.0, + update_period: int = 0, + per_env_phase: bool = True, + generator: np.random.Generator | None = None, + ) -> None: + if min_lag < 0: + raise ValueError(f"min_lag must be >= 0, got {min_lag}") + if max_lag < min_lag: + raise ValueError(f"max_lag ({max_lag}) must be >= min_lag ({min_lag})") + if not 0.0 <= hold_prob <= 1.0: + raise ValueError(f"hold_prob must be in [0, 1], got {hold_prob}") + if update_period < 0: + raise ValueError(f"update_period must be >= 0, got {update_period}") + + self.min_lag = min_lag + self.max_lag = max_lag + self.batch_size = batch_size + self.per_env = per_env + self.hold_prob = hold_prob + self.update_period = update_period + self.per_env_phase = per_env_phase + self.generator = generator + + buffer_size = max_lag + 1 if max_lag > 0 else 1 + self._buffer = CircularBuffer(max_len=buffer_size, batch_size=batch_size) + self._current_lags = np.zeros(batch_size, dtype=np.int64) + self._step_count = np.zeros(batch_size, dtype=np.int64) + + if update_period > 0 and per_env_phase: + self._phase_offsets = self._require_generator().integers( + 0, update_period, size=batch_size, dtype=np.int64 + ) + else: + self._phase_offsets = np.zeros(batch_size, dtype=np.int64) + + @property + def is_initialized(self) -> bool: + """Check if buffer has been initialized with at least one append.""" + return self._buffer.is_initialized + + @property + def current_lags(self) -> np.ndarray: + """Current lag per environment. Shape: (batch_size,).""" + return self._current_lags + + def set_lags( + self, + lags: np.ndarray, + batch_ids: Sequence[int] | np.ndarray | slice | None = None, + ) -> None: + """Set lag values for specified environments. + + Args: + lags: Lag values to set. Shape: (num_batch_ids,) or scalar. + batch_ids: Batch indices to set, or None to set all. + """ + idx = slice(None) if batch_ids is None else batch_ids + self._current_lags[idx] = np.clip(lags, self.min_lag, self.max_lag) + + def reset(self, batch_ids: Sequence[int] | np.ndarray | slice | None = None) -> None: + """Reset specified environments to initial state. + + Args: + batch_ids: Batch indices to reset, or None to reset all. + """ + if isinstance(batch_ids, slice): + indices = range(*batch_ids.indices(self.batch_size)) + batch_ids = list(indices) + + self._buffer.reset(batch_ids=batch_ids) + idx = slice(None) if batch_ids is None else batch_ids + self._current_lags[idx] = 0 + self._step_count[idx] = 0 + if self.update_period > 0 and self.per_env_phase: + new_phases = self._require_generator().integers( + 0, self.update_period, size=self.batch_size, dtype=np.int64 + ) + self._phase_offsets[idx] = new_phases[idx] + + def append(self, data: np.ndarray) -> None: + """Append new observation to buffer. + + Args: + data: Observation tensor of shape (batch_size, ...). + """ + self._buffer.append(data) + + def backfill(self, data: np.ndarray, batch_ids: np.ndarray) -> None: + """Backfill the given rows with one frame, without advancing time. + + Used after a partial reset: the reset rows (whose lags and step counters + were just zeroed by reset) get their first post-reset frame while the + remaining rows keep their history, lags, and update schedule intact. + + Args: + data: Tensor of shape (batch_size, ...); only rows at batch_ids are read. + batch_ids: Batch indices to backfill. + """ + self._buffer.backfill(data, batch_ids) + + def compute(self) -> np.ndarray: + """Compute delayed observation for current step. + + Advances the lag update schedule, then returns the delayed observation. + + Returns: + Delayed observation with shape (batch_size, ...). + """ + if not self.is_initialized: + raise RuntimeError("Buffer not initialized. Call append() first.") + + self._update_lags() + return self.peek() + + def peek(self) -> np.ndarray: + """Return the delayed observation using current lags, without advancing. + + Unlike compute, this neither steps the update schedule nor resamples lags, + so it is safe to call outside the once-per-step cadence. + + Returns: + Delayed observation with shape (batch_size, ...). + """ + if not self.is_initialized: + raise RuntimeError("Buffer not initialized. Call append() first.") + + # Clamp lags to valid range [0, buffer_length - 1]. + # Buffer may not be full yet (e.g., only 2 frames but sampled lag=3). + valid_lags = np.minimum(self._current_lags, self._buffer.current_length - 1) + valid_lags = np.maximum(valid_lags, 0) + + return self._buffer[valid_lags] + + def _update_lags(self) -> None: + """Update current lags according to configured policy.""" + if self.update_period > 0: + phase_adjusted_count = (self._step_count + self._phase_offsets) % (self.update_period) + should_update = phase_adjusted_count == 0 + else: + should_update = np.ones(self.batch_size, dtype=np.bool_) + new_lags = self._sample_lags(should_update) + self._current_lags = np.where(should_update, new_lags, self._current_lags) + self._step_count += 1 + + def _sample_lags(self, mask: np.ndarray) -> np.ndarray: + """Sample new lags for specified environments. + + Args: + mask: Boolean mask of shape (batch_size,) indicating which envs to sample. + + Returns: + New lags with shape (batch_size,). + """ + if self.min_lag == self.max_lag: + candidate_lags = np.full(self.batch_size, self.min_lag, dtype=np.int64) + elif self.per_env: + candidate_lags = self._require_generator().integers( + self.min_lag, self.max_lag + 1, size=self.batch_size, dtype=np.int64 + ) + else: + shared_lag = self._require_generator().integers(self.min_lag, self.max_lag + 1) + candidate_lags = np.full(self.batch_size, shared_lag, dtype=np.int64) + + if self.hold_prob > 0.0: + should_sample = self._require_generator().random(self.batch_size) >= self.hold_prob + update_mask = mask & should_sample + else: + update_mask = mask + + return np.where(update_mask, candidate_lags, self._current_lags) + + def _require_generator(self) -> np.random.Generator: + if self.generator is None: + raise ValueError( + "DelayBuffer stochastic sampling requires an env-owned NumPy generator." + ) + return self.generator diff --git a/src/unilab/managers/_noise/__init__.py b/src/unilab/managers/_noise/__init__.py new file mode 100644 index 000000000..baa5838b5 --- /dev/null +++ b/src/unilab/managers/_noise/__init__.py @@ -0,0 +1,15 @@ +# Derived from mujocolab/mjlab v1.6.0 (0fb8a681), src/mjlab/utils/noise/__init__.py. +# Copyright 2025, The mjlab Developers. +# Modified by UniLab for NumPy and UniLab contracts; licensed under Apache-2.0. +from unilab.managers._noise.noise_cfg import ConstantNoiseCfg as ConstantNoiseCfg +from unilab.managers._noise.noise_cfg import GaussianNoiseCfg as GaussianNoiseCfg +from unilab.managers._noise.noise_cfg import NoiseCfg as NoiseCfg +from unilab.managers._noise.noise_cfg import NoiseModelCfg as NoiseModelCfg +from unilab.managers._noise.noise_cfg import ( + NoiseModelWithAdditiveBiasCfg as NoiseModelWithAdditiveBiasCfg, +) +from unilab.managers._noise.noise_cfg import UniformNoiseCfg as UniformNoiseCfg +from unilab.managers._noise.noise_model import NoiseModel as NoiseModel +from unilab.managers._noise.noise_model import ( + NoiseModelWithAdditiveBias as NoiseModelWithAdditiveBias, +) diff --git a/src/unilab/managers/_noise/noise_cfg.py b/src/unilab/managers/_noise/noise_cfg.py new file mode 100644 index 000000000..d5cf2d8ec --- /dev/null +++ b/src/unilab/managers/_noise/noise_cfg.py @@ -0,0 +1,143 @@ +# Derived from mujocolab/mjlab v1.6.0 (0fb8a681), src/mjlab/utils/noise/noise_cfg.py. +# Copyright 2025, The mjlab Developers. +# Modified by UniLab for NumPy and UniLab contracts; licensed under Apache-2.0. +from __future__ import annotations + +import abc +from dataclasses import dataclass +from typing import ClassVar, Literal + +import numpy as np +from typing_extensions import override + +from unilab.managers._noise import noise_model + +# Type alias for noise parameters: scalar or per-dimension values. +NoiseParam = float | tuple[float, ...] + + +@dataclass(kw_only=True) +class NoiseCfg(abc.ABC): + """Base configuration for a noise term.""" + + operation: Literal["add", "scale", "abs"] = "add" + + @staticmethod + def _as_array(value: NoiseParam, dtype: np.dtype) -> np.ndarray: + """Convert a scalar or per-component parameter without a device abstraction.""" + return np.asarray(value, dtype=dtype) + + @abc.abstractmethod + def apply(self, data: np.ndarray, *, rng: np.random.Generator | None = None) -> np.ndarray: + """Apply noise to the input data.""" + + +@dataclass +class ConstantNoiseCfg(NoiseCfg): + bias: NoiseParam = 0.0 + + @override + def apply(self, data: np.ndarray, *, rng: np.random.Generator | None = None) -> np.ndarray: + del rng + bias = self._as_array(self.bias, data.dtype) + + if self.operation == "add": + return data + bias + elif self.operation == "scale": + return data * bias + elif self.operation == "abs": + return np.zeros_like(data) + bias + else: + raise ValueError(f"Unsupported noise operation: {self.operation}") + + +@dataclass +class UniformNoiseCfg(NoiseCfg): + n_min: NoiseParam = -1.0 + n_max: NoiseParam = 1.0 + + def __post_init__(self): + if isinstance(self.n_min, float) and isinstance(self.n_max, float): + if self.n_min >= self.n_max: + raise ValueError(f"n_min ({self.n_min}) must be less than n_max ({self.n_max})") + + @override + def apply(self, data: np.ndarray, *, rng: np.random.Generator | None = None) -> np.ndarray: + if rng is None: + raise ValueError("UniformNoiseCfg requires an env-owned NumPy generator.") + n_min = self._as_array(self.n_min, data.dtype) + n_max = self._as_array(self.n_max, data.dtype) + + # Generate uniform noise in [0, 1) and scale to [n_min, n_max). + noise = rng.random(data.shape).astype(data.dtype, copy=False) + noise = noise * (n_max - n_min) + n_min + + if self.operation == "add": + return data + noise + elif self.operation == "scale": + return data * noise + elif self.operation == "abs": + return noise + else: + raise ValueError(f"Unsupported noise operation: {self.operation}") + + +@dataclass +class GaussianNoiseCfg(NoiseCfg): + mean: NoiseParam = 0.0 + std: NoiseParam = 1.0 + + def __post_init__(self): + if isinstance(self.std, float) and self.std <= 0: + raise ValueError(f"std ({self.std}) must be positive") + + @override + def apply(self, data: np.ndarray, *, rng: np.random.Generator | None = None) -> np.ndarray: + if rng is None: + raise ValueError("GaussianNoiseCfg requires an env-owned NumPy generator.") + mean = self._as_array(self.mean, data.dtype) + std = self._as_array(self.std, data.dtype) + + # Generate standard normal noise and scale. + noise = rng.standard_normal(data.shape).astype(data.dtype, copy=False) + noise = mean + std * noise + + if self.operation == "add": + return data + noise + elif self.operation == "scale": + return data * noise + elif self.operation == "abs": + return noise + else: + raise ValueError(f"Unsupported noise operation: {self.operation}") + + +## +# Noise models. +## + + +@dataclass(kw_only=True) +class NoiseModelCfg: + """Configuration for a noise model.""" + + noise_cfg: NoiseCfg + + class_type: ClassVar[type[noise_model.NoiseModel]] = noise_model.NoiseModel + + def __init_subclass__(cls, class_type: type[noise_model.NoiseModel]): + cls.class_type = class_type + + +@dataclass(kw_only=True) +class NoiseModelWithAdditiveBiasCfg( + NoiseModelCfg, class_type=noise_model.NoiseModelWithAdditiveBias +): + """Configuration for an additive Gaussian noise with bias model.""" + + bias_noise_cfg: NoiseCfg | None = None + sample_bias_per_component: bool = True + + def __post_init__(self): + if self.bias_noise_cfg is None: + raise ValueError("bias_noise_cfg must be specified for NoiseModelWithAdditiveBiasCfg") diff --git a/src/unilab/managers/_noise/noise_model.py b/src/unilab/managers/_noise/noise_model.py new file mode 100644 index 000000000..7cd8868a4 --- /dev/null +++ b/src/unilab/managers/_noise/noise_model.py @@ -0,0 +1,97 @@ +# Derived from mujocolab/mjlab v1.6.0 (0fb8a681), src/mjlab/utils/noise/noise_model.py. +# Copyright 2025, The mjlab Developers. +# Modified by UniLab for NumPy and UniLab contracts; licensed under Apache-2.0. +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np +from typing_extensions import override + +if TYPE_CHECKING: + from unilab.managers._noise import noise_cfg + + +class NoiseModel: + """Base class for noise models.""" + + def __init__( + self, + noise_model_cfg: noise_cfg.NoiseModelCfg, + num_envs: int, + rng: np.random.Generator, + ): + self._noise_model_cfg = noise_model_cfg + self._num_envs = num_envs + self._rng = rng + + # Validate configuration. + if not hasattr(noise_model_cfg, "noise_cfg") or noise_model_cfg.noise_cfg is None: + raise ValueError("NoiseModelCfg must have a valid noise_cfg") + + def reset(self, env_ids: np.ndarray | slice | None = None) -> None: + """Reset noise model state. Override in subclasses if needed.""" + + def __call__(self, data: np.ndarray) -> np.ndarray: + """Apply noise to input data.""" + assert self._noise_model_cfg.noise_cfg is not None + return self._noise_model_cfg.noise_cfg.apply(data, rng=self._rng) + + +class NoiseModelWithAdditiveBias(NoiseModel): + """Noise model with additional additive bias that is constant for the duration + of the entire episode.""" + + def __init__( + self, + noise_model_cfg: noise_cfg.NoiseModelWithAdditiveBiasCfg, + num_envs: int, + rng: np.random.Generator, + ): + super().__init__(noise_model_cfg, num_envs, rng) + + # Validate bias configuration. + if not hasattr(noise_model_cfg, "bias_noise_cfg") or noise_model_cfg.bias_noise_cfg is None: + raise ValueError("NoiseModelWithAdditiveBiasCfg must have a valid bias_noise_cfg") + + self._bias_noise_cfg = noise_model_cfg.bias_noise_cfg + self._sample_bias_per_component = noise_model_cfg.sample_bias_per_component + + # Shape is materialized from the first observation so scalar and + # higher-rank terms broadcast without a device-specific convention. + self._bias = np.zeros((num_envs, 1), dtype=np.float32) + self._bias_initialized = False + + @override + def reset(self, env_ids: np.ndarray | slice | None = None) -> None: + """Reset bias values for specified environments.""" + indices = slice(None) if env_ids is None else env_ids + # Sample new bias values. + self._bias[indices] = self._bias_noise_cfg.apply(self._bias[indices], rng=self._rng) + + def _initialize_bias_shape(self, data: np.ndarray) -> None: + """Initialize bias tensor shape based on data and configuration.""" + if not self._bias_initialized: + if data.ndim == 0 or data.shape[0] != self._num_envs: + raise ValueError( + f"NoiseModel expected leading dimension {self._num_envs}, " + f"received shape {data.shape}." + ) + if self._sample_bias_per_component: + bias_shape = data.shape + else: + bias_shape = (self._num_envs, *([1] * (data.ndim - 1))) + self._bias = np.zeros(bias_shape, dtype=data.dtype) + self._bias_initialized = True + self.reset() + elif self._bias.shape != data.shape and self._sample_bias_per_component: + raise ValueError( + f"NoiseModel observation shape changed from {self._bias.shape} to {data.shape}." + ) + + @override + def __call__(self, data: np.ndarray) -> np.ndarray: + """Apply noise and additive bias to input data.""" + self._initialize_bias_shape(data) + noisy_data = super().__call__(data) + return noisy_data + self._bias diff --git a/src/unilab/managers/_types.py b/src/unilab/managers/_types.py new file mode 100644 index 000000000..9ff2cbfeb --- /dev/null +++ b/src/unilab/managers/_types.py @@ -0,0 +1,52 @@ +"""Typing-only contracts used by the standalone manager package. + +The production environment and scene adapters implement these structural protocols in +later integration layers. Keeping them here prevents the manager core from importing +an environment, backend, runner, or IPC implementation. +""" + +from __future__ import annotations + +from typing import Any, Protocol + +import numpy as np + + +class ManagerEntity(Protocol): + """Cold-path entity metadata required by :class:`SceneEntityCfg`.""" + + joint_names: list[str] + body_names: list[str] + geom_names: list[str] + site_names: list[str] + actuator_names: list[str] + tendon_names: list[str] + camera_names: list[str] + light_names: list[str] + material_names: list[str] + texture_names: list[str] + pair_names: list[str] + + +class ManagerScene(Protocol): + """Minimal name-addressable scene surface consumed by managers.""" + + def __getitem__(self, name: str) -> ManagerEntity: ... + + +class ManagerBasedRlEnv(Protocol): + """Structural context visible to manager terms. + + Additional task-owned state is intentionally not enumerated: term callables may use + their concrete environment type, while the manager core depends only on this seam. + """ + + num_envs: int + rng: np.random.Generator + scene: ManagerScene + max_episode_length_s: float + + def __getattr__(self, name: str) -> Any: ... + + +DebugVisualizer = Any diff --git a/src/unilab/managers/action_manager.py b/src/unilab/managers/action_manager.py new file mode 100644 index 000000000..b966446ab --- /dev/null +++ b/src/unilab/managers/action_manager.py @@ -0,0 +1,219 @@ +# Derived from mujocolab/mjlab v1.6.0 (0fb8a681), src/mjlab/managers/action_manager.py. +# Copyright 2025, The mjlab Developers. +# Modified by UniLab for NumPy and UniLab contracts; licensed under Apache-2.0. +"""Action manager for processing actions sent to the environment.""" + +from __future__ import annotations + +import abc +from dataclasses import dataclass +from typing import TYPE_CHECKING, Sequence + +import numpy as np +from prettytable import PrettyTable + +from unilab.managers.manager_base import ManagerBase, ManagerTermBase + +if TYPE_CHECKING: + from unilab.managers._types import ManagerBasedRlEnv + + +@dataclass(kw_only=True) +class ActionTermCfg(abc.ABC): + """Configuration for an action term. + + Action terms process raw actions from the policy and apply them to entities + in the scene (e.g., setting joint positions, velocities, or efforts). + """ + + entity_name: str + """Name of the entity in the scene that this action term controls.""" + + clip: dict[str, tuple] | None = None + """Optional clipping bounds applied to processed actions (after scale + and offset). Dict maps actuator name regex patterns to (min, max) + tuples, resolved the same way as ``scale`` and ``offset``.""" + + @abc.abstractmethod + def build(self, env: ManagerBasedRlEnv) -> ActionTerm: + """Build the action term from this config.""" + raise NotImplementedError + + +class ActionTerm(ManagerTermBase): + """Base class for action terms. + + The action term is responsible for processing the raw actions sent to the environment + and applying them to the entity managed by the term. + """ + + def __init__(self, cfg: ActionTermCfg, env: ManagerBasedRlEnv): + self.cfg = cfg + super().__init__(env) + self._entity = self._env.scene[self.cfg.entity_name] + + @property + @abc.abstractmethod + def action_dim(self) -> int: + raise NotImplementedError + + @abc.abstractmethod + def process_actions(self, actions: np.ndarray) -> None: + raise NotImplementedError + + @abc.abstractmethod + def apply_actions(self) -> None: + raise NotImplementedError + + @property + @abc.abstractmethod + def raw_action(self) -> np.ndarray: + raise NotImplementedError + + +class ActionManager(ManagerBase): + """Manages action processing for the environment. + + The action manager aggregates multiple action terms, each controlling a different + entity or aspect of the simulation. It splits the policy's action tensor and + routes each slice to the appropriate action term. + """ + + def __init__(self, cfg: dict[str, ActionTermCfg | None], env: ManagerBasedRlEnv): + self.cfg = cfg + super().__init__(env=env) + + # Create buffers to store actions. + self._action = np.zeros((self.num_envs, self.total_action_dim), dtype=np.float32) + self._prev_action = np.zeros_like(self._action) + self._prev_prev_action = np.zeros_like(self._action) + + def __str__(self) -> str: + msg = f" contains {len(self._term_names)} active terms.\n" + table = PrettyTable() + table.title = f"Active Action Terms (shape: {self.total_action_dim})" + table.field_names = ["Index", "Name", "Dimension"] + table.align["Name"] = "l" + table.align["Dimension"] = "r" + for index, (name, term) in enumerate(self._terms.items()): + table.add_row([index, name, term.action_dim]) + msg += str(table.get_string()) + msg += "\n" + return msg + + # Properties. + + @property + def total_action_dim(self) -> int: + return sum(self.action_term_dim) + + @property + def action_term_dim(self) -> list[int]: + return [term.action_dim for term in self._terms.values()] + + @property + def action(self) -> np.ndarray: + """Raw policy output from the current step, before per-term + scale/offset. Shape: ``(num_envs, total_action_dim)``.""" + return self._action + + @property + def prev_action(self) -> np.ndarray: + """Raw policy output from the previous step, before per-term + scale/offset. Shape: ``(num_envs, total_action_dim)``.""" + return self._prev_action + + @property + def prev_prev_action(self) -> np.ndarray: + """Raw policy output from two steps ago, before per-term + scale/offset. Shape: ``(num_envs, total_action_dim)``.""" + return self._prev_prev_action + + @property + def active_terms(self) -> list[str]: + return self._term_names + + # Methods. + + def get_term(self, name: str) -> ActionTerm: + return self._terms[name] + + def reset(self, env_ids: np.ndarray | slice | None = None) -> dict[str, float]: + if env_ids is None: + env_ids = slice(None) + # Reset action history. + self._prev_action[env_ids] = 0.0 + self._prev_prev_action[env_ids] = 0.0 + self._action[env_ids] = 0.0 + # Reset action terms. + for term in self._terms.values(): + term.reset(env_ids=env_ids) + return {} + + def process_action(self, action: np.ndarray) -> None: + """Store the raw policy output and route slices to each action term. + + Called once per policy step. The raw action tensor is saved into the + history buffers (``action``, ``prev_action``, ``prev_prev_action``) *before* + any per-term scale/offset is applied. Each term then receives its slice and + independently applies its own affine transformation via + :meth:`ActionTerm.process_actions`. + """ + if not isinstance(action, np.ndarray): + raise TypeError(f"ActionManager expected np.ndarray, received {type(action).__name__}.") + expected_shape = (self.num_envs, self.total_action_dim) + if action.shape != expected_shape: + raise ValueError( + f"Invalid action shape, expected {expected_shape}, received {action.shape}." + ) + if not np.isfinite(action).all(): + raise ValueError("ActionManager received an action containing NaN or Inf.") + # Shift history: prev_prev ← prev ← current ← new. + self._prev_prev_action[:] = self._prev_action + self._prev_action[:] = self._action + self._action[:] = action + # Split the flat action vector and route each slice to its term. + idx = 0 + for term in self._terms.values(): + term_actions = self._action[:, idx : idx + term.action_dim] + term.process_actions(term_actions) + idx += term.action_dim + + def apply_action(self) -> None: + """Write processed actions to entity actuator targets. + + Called on every decimation substep (physics step), not just once per policy + step. Each term writes its most recently processed targets to the simulation. + """ + for term in self._terms.values(): + term.apply_actions() + + def get_active_iterable_terms(self, env_idx: int) -> Sequence[tuple[str, Sequence[float]]]: + terms = [] + idx = 0 + for name, term in self._terms.items(): + term_actions = self._action[env_idx, idx : idx + term.action_dim] + terms.append((name, term_actions.tolist())) + idx += term.action_dim + return terms + + def _prepare_terms(self) -> None: + self._term_names: list[str] = list() + self._terms: dict[str, ActionTerm] = dict() + + for term_name, term_cfg in self.cfg.items(): + if term_cfg is None: + print(f"term: {term_name} set to None, skipping...") + continue + term = term_cfg.build(self._env) + if not isinstance(term, ActionTerm): + raise TypeError( + f"ActionManager term '{term_name}' build() returned " + f"{type(term).__name__}, expected ActionTerm." + ) + if not isinstance(term.action_dim, int) or term.action_dim < 0: + raise ValueError( + f"ActionManager term '{term_name}' has invalid action_dim {term.action_dim}." + ) + self._term_names.append(term_name) + self._terms[term_name] = term diff --git a/src/unilab/managers/command_manager.py b/src/unilab/managers/command_manager.py new file mode 100644 index 000000000..2b20323ee --- /dev/null +++ b/src/unilab/managers/command_manager.py @@ -0,0 +1,313 @@ +# Derived from mujocolab/mjlab v1.6.0 (0fb8a681), src/mjlab/managers/command_manager.py. +# Copyright 2025, The mjlab Developers. +# Modified by UniLab for NumPy and UniLab contracts; licensed under Apache-2.0. +"""Command manager for generating and updating commands.""" + +from __future__ import annotations + +import abc +import inspect +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Sequence + +import numpy as np +from prettytable import PrettyTable + +from unilab.managers.manager_base import ManagerBase, ManagerTermBase + +if TYPE_CHECKING: + from unilab.managers._types import ManagerBasedRlEnv + + +@dataclass(kw_only=True) +class CommandTermCfg(abc.ABC): + """Configuration for a command generator term. + + Command terms generate goal commands for the agent (e.g., target velocity, + target position). Commands are automatically resampled at configurable + intervals and can track metrics for logging. + """ + + resampling_time_range: tuple[float, float] + """Time range in seconds for command resampling. When the timer expires, a new + command is sampled and the timer is reset to a value uniformly drawn from + ``[min, max]``. Set both values equal for fixed-interval resampling.""" + + debug_vis: bool = False + """Whether to enable debug visualization for this command term. When True, + the command term's ``_debug_vis_impl`` method is called each frame to render + visual aids (e.g., velocity arrows, target markers).""" + + @abc.abstractmethod + def build(self, env: ManagerBasedRlEnv) -> CommandTerm: + """Build the command term from this config.""" + raise NotImplementedError + + +class CommandTerm(ManagerTermBase): + """Base class for command terms.""" + + def __init__(self, cfg: CommandTermCfg, env: ManagerBasedRlEnv): + self.cfg = cfg + super().__init__(env) + lower, upper = cfg.resampling_time_range + if not np.isfinite((lower, upper)).all() or lower > upper: + raise ValueError( + f"CommandTerm '{self.name}' has invalid resampling_time_range " + f"{cfg.resampling_time_range}." + ) + self._resampling_time_range = (lower, upper) + self._check_update_command_signature() + self.metrics: dict[str, np.ndarray] = {} + self.time_left = np.zeros(self.num_envs, dtype=np.float32) + self.command_counter = np.zeros(self.num_envs, dtype=np.int64) + + @property + @abc.abstractmethod + def command(self): + raise NotImplementedError + + def reset(self, env_ids: np.ndarray | slice | None) -> dict[str, float]: + assert isinstance(env_ids, np.ndarray) + extras = {} + for metric_name, metric_value in self.metrics.items(): + metric_slice = metric_value[env_ids] + if not np.isfinite(metric_slice).all(): + raise ValueError( + f"CommandTerm '{self.name}' metric '{metric_name}' contains NaN or Inf." + ) + extras[metric_name] = float(np.mean(metric_slice)) + metric_value[env_ids] = 0.0 + self.command_counter[env_ids] = 0 + self._resample(env_ids) + return extras + + def compute(self, dt: float | np.ndarray, env_ids: np.ndarray | None = None) -> None: + """Advance the command state by dt. + + With env_ids=None (the per-step path) all envs are updated; with env_ids + (the reset path) timers and the command update are scoped to those envs. + Metrics are always refreshed. + + dt may be a scalar (all envs) or a per-env tensor (auto-reset path, + where freshly reset envs get zero to keep their timers full). A tensor + dt requires env_ids=None. + """ + if isinstance(dt, np.ndarray): + if env_ids is not None: + raise ValueError("Per-environment command dt requires env_ids=None.") + if dt.shape != (self.num_envs,): + raise ValueError( + f"CommandTerm '{self.name}' expected dt shape ({self.num_envs},), " + f"received {dt.shape}." + ) + dt_is_finite = np.isfinite(dt).all() if isinstance(dt, np.ndarray) else np.isfinite(dt) + if not dt_is_finite: + raise ValueError(f"CommandTerm '{self.name}' received non-finite dt.") + self._update_metrics() + self._validate_metrics() + if env_ids is None: + self.time_left -= dt + resample_env_ids = np.flatnonzero(self.time_left <= 0.0) + else: + assert not isinstance(dt, np.ndarray) + self.time_left[env_ids] -= dt + resample_env_ids = env_ids[self.time_left[env_ids] <= 0.0] + if len(resample_env_ids) > 0: + self._resample(resample_env_ids) + self._update_command(env_ids) + + def _validate_metrics(self) -> None: + for metric_name, metric_value in self.metrics.items(): + if not isinstance(metric_value, np.ndarray): + raise TypeError( + f"CommandTerm '{self.name}' metric '{metric_name}' returned " + f"{type(metric_value).__name__}, expected np.ndarray." + ) + if metric_value.ndim == 0 or metric_value.shape[0] != self.num_envs: + raise ValueError( + f"CommandTerm '{self.name}' metric '{metric_name}' returned shape " + f"{metric_value.shape}, expected leading dimension {self.num_envs}." + ) + if not np.isfinite(metric_value).all(): + raise ValueError( + f"CommandTerm '{self.name}' metric '{metric_name}' contains NaN or Inf." + ) + + def _check_update_command_signature(self) -> None: + """Fail fast with a migration hint for terms with the old signature.""" + try: + sig = inspect.signature(self._update_command) + except (TypeError, ValueError): + return + if len(sig.parameters) == 0: + raise TypeError( + f"{type(self).__name__}._update_command must accept env_ids: " + "_update_command(self, env_ids: np.ndarray | None). It receives " + "None on the per-step update and the reset env ids on reset(); " + "scope per-step state advances to env_ids." + ) + + def _resample(self, env_ids: np.ndarray) -> None: + if len(env_ids) != 0: + lower, upper = self._resampling_time_range + self.time_left[env_ids] = self._env.rng.uniform(lower, upper, len(env_ids)) + self._resample_command(env_ids) + self.command_counter[env_ids] += 1 + + @abc.abstractmethod + def _update_metrics(self) -> None: + """Update the metrics based on the current state.""" + raise NotImplementedError + + @abc.abstractmethod + def _resample_command(self, env_ids: np.ndarray) -> None: + """Resample the command for the specified environments.""" + raise NotImplementedError + + @abc.abstractmethod + def _update_command(self, env_ids: np.ndarray | None) -> None: + """Update the command based on the current state. + + env_ids is None on the per-step update (all envs) and the reset env ids on reset(). + Scope per-step state advances (e.g. a motion frame index) to env_ids; pure + functions of the current state may ignore it. + """ + raise NotImplementedError + + +class CommandManager(ManagerBase): + """Manages command generation for the environment. + + The command manager generates and updates goal commands for the agent (e.g., + target velocity, target position). Commands are resampled at configurable + intervals and can track metrics for logging. + """ + + _env: ManagerBasedRlEnv + + def __init__(self, cfg: dict[str, CommandTermCfg | None], env: ManagerBasedRlEnv): + self._terms: dict[str, CommandTerm] = dict() + + self.cfg = cfg + super().__init__(env) + + def __str__(self) -> str: + msg = f" contains {len(self._terms.values())} active terms.\n" + table = PrettyTable() + table.title = "Active Command Terms" + table.field_names = ["Index", "Name", "Type"] + table.align["Name"] = "l" + for index, (name, term) in enumerate(self._terms.items()): + table.add_row([index, name, term.__class__.__name__]) + msg += str(table.get_string()) + msg += "\n" + return msg + + # Properties. + + @property + def active_terms(self) -> list[str]: + return list(self._terms.keys()) + + def get_active_iterable_terms(self, env_idx: int) -> Sequence[tuple[str, Sequence[float]]]: + terms = [] + for name, term in self._terms.items(): + command = self._validate_command(name, term.command) + terms.append((name, command[env_idx].tolist())) + return terms + + def reset(self, env_ids: np.ndarray | slice | None) -> dict[str, float]: + if env_ids is None: + env_ids = np.arange(self.num_envs) + elif isinstance(env_ids, slice): + env_ids = np.arange(self.num_envs)[env_ids] + extras = {} + for name, term in self._terms.items(): + metrics = term.reset(env_ids=env_ids) + self._validate_command(name, term.command) + for metric_name, metric_value in metrics.items(): + extras[f"Metrics/{name}/{metric_name}"] = metric_value + return extras + + def compute(self, dt: float | np.ndarray, env_ids: np.ndarray | None = None) -> None: + for name, term in self._terms.items(): + term.compute(dt, env_ids) + self._validate_command(name, term.command) + + def get_command(self, name: str) -> np.ndarray: + return self._validate_command(name, self._terms[name].command) + + def get_term(self, name: str) -> CommandTerm: + return self._terms[name] + + def get_term_cfg(self, name: str) -> CommandTermCfg: + term_cfg = self.cfg[name] + if term_cfg is None: + raise KeyError(f"Command term '{name}' is disabled.") + return term_cfg + + def _prepare_terms(self) -> None: + for term_name, term_cfg in self.cfg.items(): + if term_cfg is None: + print(f"term: {term_name} set to None, skipping...") + continue + if term_cfg.debug_vis: + raise NotImplementedError( + f"CommandManager term '{term_name}' requested viewer debug visualization; " + "viewer glue is unsupported by the UniLab manager core." + ) + term = term_cfg.build(self._env) + if not isinstance(term, CommandTerm): + raise TypeError( + f"Returned object for the term {term_name} is not of type CommandType." + ) + self._terms[term_name] = term + + def _validate_command(self, name: str, command: np.ndarray) -> np.ndarray: + if not isinstance(command, np.ndarray): + raise TypeError( + f"CommandManager term '{name}' returned {type(command).__name__}, " + "expected np.ndarray." + ) + if command.ndim < 1 or command.shape[0] != self.num_envs: + raise ValueError( + f"CommandManager term '{name}' returned shape {command.shape}, " + f"expected leading dimension {self.num_envs}." + ) + if not np.isfinite(command).all(): + raise ValueError(f"CommandManager term '{name}' returned NaN or Inf.") + return command + + +class NullCommandManager: + """Placeholder for absent command manager that safely no-ops all operations.""" + + def __init__(self): + self.active_terms: list[str] = [] + self._terms: dict[str, Any] = {} + self.cfg = None + + def __str__(self) -> str: + return " (inactive)" + + def __repr__(self) -> str: + return "NullCommandManager()" + + def get_active_iterable_terms(self, env_idx: int) -> Sequence[tuple[str, Sequence[float]]]: + return [] + + def reset(self, env_ids: np.ndarray | None = None) -> dict[str, np.ndarray]: + return {} + + def compute(self, dt: float | np.ndarray, env_ids: np.ndarray | None = None) -> None: + pass + + def get_command(self, name: str) -> None: + return None + + def get_term(self, name: str) -> None: + return None + + def get_term_cfg(self, name: str) -> None: + return None diff --git a/src/unilab/managers/curriculum_manager.py b/src/unilab/managers/curriculum_manager.py new file mode 100644 index 000000000..e2814b86a --- /dev/null +++ b/src/unilab/managers/curriculum_manager.py @@ -0,0 +1,166 @@ +# Derived from mujocolab/mjlab v1.6.0 (0fb8a681), src/mjlab/managers/curriculum_manager.py. +# Copyright 2025, The mjlab Developers. +# Modified by UniLab for NumPy and UniLab contracts; licensed under Apache-2.0. +"""Curriculum manager for updating environment quantities subject to a training curriculum.""" + +from __future__ import annotations + +from copy import deepcopy +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Sequence + +import numpy as np +from prettytable import PrettyTable + +from unilab.managers.manager_base import ManagerBase, ManagerTermBaseCfg + +if TYPE_CHECKING: + from unilab.managers._types import ManagerBasedRlEnv + + +@dataclass(kw_only=True) +class CurriculumTermCfg(ManagerTermBaseCfg): + """Configuration for a curriculum term. + + Curriculum terms modify environment parameters during training to implement + curriculum learning strategies (e.g., gradually increasing task difficulty). + """ + + pass + + +class CurriculumManager(ManagerBase): + """Manages curriculum learning for the environment. + + The curriculum manager updates environment parameters during training based + on agent performance. Each term can modify different aspects of the task + difficulty (e.g., terrain complexity, command ranges). + """ + + _env: ManagerBasedRlEnv + + def __init__(self, cfg: dict[str, CurriculumTermCfg | None], env: ManagerBasedRlEnv): + self._term_names: list[str] = list() + self._term_cfgs: list[CurriculumTermCfg] = list() + self._class_term_cfgs: list[CurriculumTermCfg] = list() + + self.cfg = deepcopy(cfg) + super().__init__(env) + + self._curriculum_state: dict[str, Any] = {} + for term_name in self._term_names: + self._curriculum_state[term_name] = None + + def __str__(self) -> str: + msg = f" contains {len(self._term_names)} active terms.\n" + table = PrettyTable() + table.title = "Active Curriculum Terms" + table.field_names = ["Index", "Name"] + table.align["Name"] = "l" + for index, name in enumerate(self._term_names): + table.add_row([index, name]) + msg += str(table.get_string()) + msg += "\n" + return msg + + # Properties. + + @property + def active_terms(self) -> list[str]: + return self._term_names + + # Methods. + + def get_term_cfg(self, term_name: str) -> CurriculumTermCfg: + if term_name not in self._term_names: + raise ValueError(f"Term '{term_name}' not found in active terms.") + return self._term_cfgs[self._term_names.index(term_name)] + + def get_active_iterable_terms(self, env_idx: int) -> Sequence[tuple[str, Sequence[float]]]: + terms = [] + for term_name, term_state in self._curriculum_state.items(): + if term_state is not None: + data = [] + if isinstance(term_state, dict): + for _key, value in term_state.items(): + if isinstance(value, np.ndarray): + value = value.item() + data.append(value) + else: + if isinstance(term_state, np.ndarray): + term_state = term_state.item() + data.append(term_state) + terms.append((term_name, data)) + return terms + + def reset(self, env_ids: np.ndarray | slice | None = None) -> dict[str, float]: + extras = {} + for term_name, term_state in self._curriculum_state.items(): + if term_state is not None: + if isinstance(term_state, dict): + for key, value in term_state.items(): + if isinstance(value, np.ndarray): + value = value.item() + extras[f"Curriculum/{term_name}/{key}"] = value + else: + if isinstance(term_state, np.ndarray): + term_state = term_state.item() + extras[f"Curriculum/{term_name}"] = term_state + for term_cfg in self._class_term_cfgs: + term_cfg.func.reset(env_ids=env_ids) + return extras + + def compute(self, env_ids: np.ndarray | slice | None = None) -> None: + if env_ids is None: + env_ids = slice(None) + for name, term_cfg in zip(self._term_names, self._term_cfgs, strict=False): + state = term_cfg.func(self._env, env_ids, **term_cfg.params) + self._validate_state(name, state) + self._curriculum_state[name] = state + + def _validate_state(self, term_name: str, state: Any) -> None: + values = state.values() if isinstance(state, dict) else (state,) + for value in values: + if isinstance(value, np.ndarray): + finite = np.isfinite(value).all() + elif isinstance(value, (int, float, np.number)): + finite = bool(np.isfinite(value)) + else: + continue + if not finite: + raise ValueError(f"CurriculumManager term '{term_name}' returned NaN or Inf.") + + def _prepare_terms(self) -> None: + for term_name, term_cfg in self.cfg.items(): + if term_cfg is None: + print(f"term: {term_name} set to None, skipping...") + continue + self._resolve_common_term_cfg(term_name, term_cfg) + self._term_names.append(term_name) + self._term_cfgs.append(term_cfg) + if hasattr(term_cfg.func, "reset") and callable(term_cfg.func.reset): + self._class_term_cfgs.append(term_cfg) + + +class NullCurriculumManager: + """Placeholder for absent curriculum manager that safely no-ops all operations.""" + + def __init__(self): + self.active_terms: list[str] = [] + self._curriculum_state: dict[str, Any] = {} + self.cfg = None + + def __str__(self) -> str: + return " (inactive)" + + def __repr__(self) -> str: + return "NullCurriculumManager()" + + def get_active_iterable_terms(self, env_idx: int) -> Sequence[tuple[str, Sequence[float]]]: + return [] + + def reset(self, env_ids: np.ndarray | None = None) -> dict[str, float]: + return {} + + def compute(self, env_ids: np.ndarray | None = None) -> None: + pass diff --git a/src/unilab/managers/event_manager.py b/src/unilab/managers/event_manager.py new file mode 100644 index 000000000..fb98d7ea9 --- /dev/null +++ b/src/unilab/managers/event_manager.py @@ -0,0 +1,276 @@ +# Derived from mujocolab/mjlab v1.6.0 (0fb8a681), src/mjlab/managers/event_manager.py. +# Copyright 2025, The mjlab Developers. +# Modified by UniLab for NumPy and UniLab contracts; licensed under Apache-2.0. +"""Event manager for orchestrating operations based on different simulation events.""" + +from __future__ import annotations + +from copy import deepcopy +from dataclasses import dataclass +from typing import TYPE_CHECKING, Literal + +import numpy as np +from prettytable import PrettyTable + +from unilab.managers.manager_base import ManagerBase, ManagerTermBaseCfg + +if TYPE_CHECKING: + from unilab.managers._types import ManagerBasedRlEnv + +EventMode = Literal["startup", "reset", "interval", "step"] + + +@dataclass(kw_only=True) +class EventTermCfg(ManagerTermBaseCfg): + """Configuration for an event term. + + Event terms trigger operations at specific simulation events. They're commonly + used for domain randomization, state resets, and periodic perturbations. + + The four modes determine when the event fires: + + - ``"startup"``: Once when the environment initializes. Use for parameters that + should be randomized per-environment but stay constant within an episode (e.g., + domain randomization). + + - ``"reset"``: On every episode reset. Use for parameters that should vary between + episodes (e.g., initial robot pose, domain randomization). + + - ``"interval"``: Periodically during simulation, controlled by ``interval_range_s``. + Use for perturbations that should happen during episodes (e.g., pushing the robot, + external disturbances). + + - ``"step"``: Every environment step, unconditionally on all envs. Use for terms that + manage per-step state such as force lifetimes (e.g., ``apply_body_impulse``). + """ + + mode: EventMode + """When the event triggers: ``"startup"`` (once at init), ``"reset"`` (every + episode), ``"interval"`` (periodically during simulation), or ``"step"`` (every + environment step).""" + + interval_range_s: tuple[float, float] | None = None + """Time range in seconds for interval mode. The next trigger time is uniformly + sampled from ``[min, max]``. Required when ``mode="interval"``.""" + + is_global_time: bool = False + """Whether all environments share the same timer. If True, all envs trigger + simultaneously. If False (default), each env has an independent timer that + resets on episode reset. Only applies to ``mode="interval"``.""" + + min_step_count_between_reset: int = 0 + """Minimum environment steps between triggers. Prevents the event from firing + too frequently when episodes reset rapidly. Only applies to ``mode="reset"``. + Set to 0 (default) to trigger on every reset.""" + + +class EventManager(ManagerBase): + """Manages event-based operations for the environment. + + The event manager triggers operations at different simulation events: startup + (once at initialization), reset (on episode reset), or interval (periodically + during simulation). Common uses include domain randomization and state resets. + """ + + _env: ManagerBasedRlEnv + + def __init__(self, cfg: dict[str, EventTermCfg | None], env: ManagerBasedRlEnv): + self.cfg = deepcopy(cfg) + self._mode_term_names: dict[EventMode, list[str]] = dict() + self._mode_term_cfgs: dict[EventMode, list[EventTermCfg]] = dict() + self._mode_class_term_cfgs: dict[EventMode, list[EventTermCfg]] = dict() + + super().__init__(env=env) + + def __str__(self) -> str: + msg = f" contains {len(self._mode_term_names)} active terms.\n" + for mode in self._mode_term_names: + table = PrettyTable() + table.title = f"Active Event Terms in Mode: '{mode}'" + if mode == "interval": + table.field_names = ["Index", "Name", "Interval time range (s)"] + table.align["Name"] = "l" + for index, (name, cfg) in enumerate( + zip(self._mode_term_names[mode], self._mode_term_cfgs[mode], strict=False) + ): + table.add_row([index, name, cfg.interval_range_s]) + else: + table.field_names = ["Index", "Name"] + table.align["Name"] = "l" + for index, name in enumerate(self._mode_term_names[mode]): + table.add_row([index, name]) + msg += str(table.get_string()) + msg += "\n" + return msg + + # Properties. + + @property + def active_terms(self) -> dict[EventMode, list[str]]: + return self._mode_term_names + + @property + def available_modes(self) -> list[EventMode]: + return list(self._mode_term_names.keys()) + + # Methods. + + def get_term_cfg(self, term_name: str) -> EventTermCfg: + """Get the configuration of a specific event term by name.""" + for mode in self._mode_term_names: + if term_name in self._mode_term_names[mode]: + index = self._mode_term_names[mode].index(term_name) + return self._mode_term_cfgs[mode][index] + raise ValueError(f"Event term '{term_name}' not found in active terms.") + + def reset(self, env_ids: np.ndarray | None = None): + for mode_cfg in self._mode_class_term_cfgs.values(): + for term_cfg in mode_cfg: + term_cfg.func.reset(env_ids=env_ids) + if env_ids is None: + num_envs = self._env.num_envs + ids: np.ndarray | slice = slice(None) + else: + num_envs = len(env_ids) + ids = env_ids + # Iterate the full interval term list: _interval_term_time_left is parallel + # to _mode_term_cfgs["interval"], not the class-only subset. + if "interval" in self._mode_term_cfgs: + for index, term_cfg in enumerate(self._mode_term_cfgs["interval"]): + if not term_cfg.is_global_time: + assert term_cfg.interval_range_s is not None + lower, upper = term_cfg.interval_range_s + sampled_interval = self._env.rng.uniform(lower, upper, num_envs) + self._interval_term_time_left[index][ids] = sampled_interval + return {} + + def apply( + self, + mode: EventMode, + env_ids: np.ndarray | slice | None = None, + dt: float | None = None, + global_env_step_count: int | None = None, + ): + if mode not in ("startup", "reset", "interval", "step"): + raise ValueError(f"Unsupported event mode '{mode}'.") + if mode not in self._mode_term_cfgs: + return + if mode == "interval" and dt is None: + raise ValueError(f"Event mode '{mode}' requires the time-step of the environment.") + if mode == "interval" and env_ids is not None: + raise ValueError( + f"Event mode '{mode}' does not require environment indices. This is an undefined behavior" + " as the environment indices are computed based on the time left for each environment." + ) + if mode == "reset" and global_env_step_count is None: + raise ValueError( + f"Event mode '{mode}' requires the total number of environment steps to be provided." + ) + if mode == "step" and dt is None: + raise ValueError(f"Event mode '{mode}' requires the time-step of the environment.") + + for index, term_cfg in enumerate(self._mode_term_cfgs[mode]): + if mode == "interval": + time_left = self._interval_term_time_left[index] + assert dt is not None + time_left -= dt + if term_cfg.is_global_time: + if time_left < 1e-6: + assert term_cfg.interval_range_s is not None + lower, upper = term_cfg.interval_range_s + sampled_interval = self._env.rng.uniform(lower, upper, 1) + self._interval_term_time_left[index][:] = sampled_interval + term_cfg.func(self._env, None, **term_cfg.params) + else: + valid_env_ids = np.flatnonzero(time_left < 1e-6) + if len(valid_env_ids) > 0: + assert term_cfg.interval_range_s is not None + lower, upper = term_cfg.interval_range_s + sampled_time = self._env.rng.uniform(lower, upper, len(valid_env_ids)) + self._interval_term_time_left[index][valid_env_ids] = sampled_time + term_cfg.func(self._env, valid_env_ids, **term_cfg.params) + elif mode == "step": + term_cfg.func(self._env, None, **term_cfg.params) + elif mode == "reset": + assert global_env_step_count is not None + # Reset events require concrete indices: callers (e.g. ManagerBasedRlEnv) + # resolve None to all environments upstream. Enforce that here so a future + # caller passing None fails loudly instead of leaking a slice into event + # functions, which only understand None or a tensor. + if env_ids is None: + raise ValueError("Event mode 'reset' requires concrete environment indices.") + min_step_count = term_cfg.min_step_count_between_reset + if min_step_count == 0: + self._reset_term_last_triggered_step_id[index][env_ids] = global_env_step_count + self._reset_term_last_triggered_once[index][env_ids] = True + term_cfg.func(self._env, env_ids, **term_cfg.params) + else: + last_triggered_step = self._reset_term_last_triggered_step_id[index][env_ids] + triggered_at_least_once = self._reset_term_last_triggered_once[index][env_ids] + steps_since_triggered = global_env_step_count - last_triggered_step + valid_trigger = steps_since_triggered >= min_step_count + valid_trigger |= (last_triggered_step == 0) & ~triggered_at_least_once + if isinstance(env_ids, np.ndarray): + valid_env_ids = env_ids[valid_trigger] + else: + valid_env_ids = np.flatnonzero(valid_trigger) + if len(valid_env_ids) > 0: + self._reset_term_last_triggered_once[index][valid_env_ids] = True + self._reset_term_last_triggered_step_id[index][valid_env_ids] = ( + global_env_step_count + ) + term_cfg.func(self._env, valid_env_ids, **term_cfg.params) + else: + term_cfg.func(self._env, env_ids, **term_cfg.params) + + def _prepare_terms(self) -> None: + self._interval_term_time_left: list[np.ndarray] = list() + self._reset_term_last_triggered_step_id: list[np.ndarray] = list() + self._reset_term_last_triggered_once: list[np.ndarray] = list() + + for term_name, term_cfg in self.cfg.items(): + if term_cfg is None: + print(f"term: {term_name} set to None, skipping...") + continue + self._resolve_common_term_cfg(term_name, term_cfg) + if term_cfg.mode not in ("startup", "reset", "interval", "step"): + raise ValueError( + f"EventManager term '{term_name}' has unsupported mode '{term_cfg.mode}'." + ) + if term_cfg.mode not in self._mode_term_names: + self._mode_term_names[term_cfg.mode] = list() + self._mode_term_cfgs[term_cfg.mode] = list() + self._mode_class_term_cfgs[term_cfg.mode] = list() + self._mode_term_names[term_cfg.mode].append(term_name) + self._mode_term_cfgs[term_cfg.mode].append(term_cfg) + if hasattr(term_cfg.func, "reset") and callable(term_cfg.func.reset): + self._mode_class_term_cfgs[term_cfg.mode].append(term_cfg) + if term_cfg.mode == "interval": + if term_cfg.interval_range_s is None: + raise ValueError( + f"Event term '{term_name}' has mode 'interval' but 'interval_range_s' is not specified." + ) + lower, upper = term_cfg.interval_range_s + if not np.isfinite((lower, upper)).all() or lower > upper: + raise ValueError( + f"EventManager term '{term_name}' has invalid interval_range_s " + f"{term_cfg.interval_range_s}." + ) + if term_cfg.is_global_time: + time_left = self._env.rng.uniform(lower, upper, 1) + self._interval_term_time_left.append(time_left) + else: + time_left = self._env.rng.uniform(lower, upper, self.num_envs) + self._interval_term_time_left.append(time_left) + elif term_cfg.mode == "reset": + step_count = np.zeros(self.num_envs, dtype=np.int64) + self._reset_term_last_triggered_step_id.append(step_count) + no_trigger = np.zeros(self.num_envs, dtype=np.bool_) + self._reset_term_last_triggered_once.append(no_trigger) + + func = term_cfg.func + if hasattr(func, "model_fields"): + raise NotImplementedError( + f"EventManager term '{term_name}' requests direct model-field mutation; " + "this capability is unsupported by the standalone UniLab manager core." + ) diff --git a/src/unilab/managers/manager_base.py b/src/unilab/managers/manager_base.py new file mode 100644 index 000000000..bac7275d6 --- /dev/null +++ b/src/unilab/managers/manager_base.py @@ -0,0 +1,159 @@ +# Derived from mujocolab/mjlab v1.6.0 (0fb8a681), src/mjlab/managers/manager_base.py. +# Copyright 2025, The mjlab Developers. +# Modified by UniLab for NumPy and UniLab contracts; licensed under Apache-2.0. +from __future__ import annotations + +import abc +import inspect +from collections.abc import Sequence +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any + +import numpy as np + +from unilab.managers.scene_entity_config import SceneEntityCfg + +if TYPE_CHECKING: + from unilab.managers._types import ManagerBasedRlEnv + + +@dataclass +class ManagerTermBaseCfg: + """Base configuration for manager terms. + + This is the base config for terms in observation, reward, termination, curriculum, + and event managers. It provides a common interface for specifying a callable + and its parameters. + + The ``func`` field accepts either a function or a class: + + **Function-based terms** are simpler and suitable for stateless computations: + + .. code-block:: python + + RewardTermCfg(func=mdp.joint_torques_l2, weight=-0.01) + + **Class-based terms** are instantiated with ``(cfg, env)`` and useful when you need + to: + + - Cache computed values at initialization (e.g., resolve regex patterns to indices) + - Maintain state across calls + - Perform expensive setup once rather than every call + + .. code-block:: python + + class posture: + def __init__(self, cfg: RewardTermCfg, env: ManagerBasedRlEnv): + # Resolve std dict to tensor once at init + self.std = resolve_std_to_tensor(cfg.params["std"], env) + + def __call__(self, env, **kwargs) -> np.ndarray: + # Use cached self.std + return compute_posture_reward(env, self.std) + + RewardTermCfg(func=posture, params={"std": {".*knee.*": 0.3}}, weight=1.0) + + Class-based terms can optionally implement ``reset(env_ids)`` for per-episode state. + """ + + func: Any + """The callable that computes this term's value. Can be a function or a class. + Classes are auto-instantiated with ``(cfg=term_cfg, env=env)``.""" + + params: dict[str, Any] = field(default_factory=lambda: {}) + """Additional keyword arguments passed to func when called.""" + + +class ManagerTermBase: + def __init__(self, env: ManagerBasedRlEnv): + self._env = env + + # Properties. + + @property + def num_envs(self) -> int: + return self._env.num_envs + + @property + def name(self) -> str: + return self.__class__.__name__ + + # Methods. + + def reset(self, env_ids: np.ndarray | slice | None) -> Any: + """Resets the manager term.""" + del env_ids # Unused. + pass + + def __call__(self, *args, **kwargs) -> Any: + """Returns the value of the term required by the manager.""" + raise NotImplementedError + + +class ManagerBase(abc.ABC): + """Base class for all managers.""" + + def __init__(self, env: ManagerBasedRlEnv): + self._env = env + + self._prepare_terms() + + # Properties. + + @property + def num_envs(self) -> int: + return self._env.num_envs + + @property + @abc.abstractmethod + def active_terms(self) -> list[str] | dict[Any, list[str]]: + raise NotImplementedError + + # Methods. + + def reset(self, env_ids: np.ndarray) -> dict[str, Any]: + """Resets the manager and returns logging info for the current step.""" + del env_ids # Unused. + return {} + + def _check_term_shape(self, term_name: str, value: np.ndarray) -> None: + if not isinstance(value, np.ndarray): + manager_name = type(self).__name__ + raise TypeError( + f"{manager_name} term '{term_name}' returned {type(value).__name__}, " + "expected np.ndarray." + ) + if value.shape != (self.num_envs,): + manager_name = type(self).__name__ + raise ValueError( + f"{manager_name} term '{term_name}' returned shape {tuple(value.shape)}, " + f"expected ({self.num_envs},)." + ) + + def _check_term_finite(self, term_name: str, value: np.ndarray) -> None: + if np.isfinite(value).all(): + return + manager_name = type(self).__name__ + has_nan = np.isnan(value).any() + has_inf = np.isinf(value).any() + invalid_kind = "NaN/Inf" if has_nan and has_inf else "NaN" if has_nan else "Inf" + env_ids = np.flatnonzero(~np.isfinite(value)).tolist() + raise ValueError( + f"{manager_name} term '{term_name}' returned {invalid_kind} for " + f"environments {env_ids[:10]}." + ) + + def get_active_iterable_terms(self, env_idx: int) -> Sequence[tuple[str, Sequence[float]]]: + raise NotImplementedError + + @abc.abstractmethod + def _prepare_terms(self) -> None: + raise NotImplementedError + + def _resolve_common_term_cfg(self, term_name: str, term_cfg: ManagerTermBaseCfg) -> None: + del term_name # Unused. + for value in term_cfg.params.values(): + if isinstance(value, SceneEntityCfg): + value.resolve(self._env.scene) + if inspect.isclass(term_cfg.func): + term_cfg.func = term_cfg.func(cfg=term_cfg, env=self._env) diff --git a/src/unilab/managers/metrics_manager.py b/src/unilab/managers/metrics_manager.py new file mode 100644 index 000000000..961ae135c --- /dev/null +++ b/src/unilab/managers/metrics_manager.py @@ -0,0 +1,242 @@ +# Derived from mujocolab/mjlab v1.6.0 (0fb8a681), src/mjlab/managers/metrics_manager.py. +# Copyright 2025, The mjlab Developers. +# Modified by UniLab for NumPy and UniLab contracts; licensed under Apache-2.0. +"""Metrics manager for logging custom per-step metrics during training.""" + +from __future__ import annotations + +from copy import deepcopy +from dataclasses import dataclass +from typing import TYPE_CHECKING, Literal, Sequence + +import numpy as np +from prettytable import PrettyTable + +from unilab.managers.manager_base import ManagerBase, ManagerTermBaseCfg + +if TYPE_CHECKING: + from unilab.managers._types import ManagerBasedRlEnv + +REDUCE_OPTIONS = ("last", "max", "mean", "sum") + + +@dataclass(kw_only=True) +class MetricsTermCfg(ManagerTermBaseCfg): + """Configuration for a metrics term. + + Attributes: + per_substep: If True, evaluate this term once per physics substep inside + the decimation loop and report the per-step mean. Only the integrated + state (qpos, qvel, act) is current mid-loop; all derived quantities + (xpos, xquat, site_xpos, actuator_force, contacts, ...) are stale. + + reduce: How to aggregate per-step values into an episode metric. + - ``"mean"`` (default) reports ``sum / step_count``. + - ``"last"`` reports the value from the final step of the episode, + useful for binary success metrics that should not be averaged over + timesteps. + - ``"max"`` reports the highest value seen during the episode, useful + for peak metrics like maximum power or contact force. + - ``"sum"`` reports the accumulated total over the episode, useful for + cumulative quantities like episodic reward or total distance + traveled. + """ + + per_substep: bool = False + reduce: Literal["last", "max", "mean", "sum"] = "mean" + + +class MetricsManager(ManagerBase): + """Accumulates per-step metric values, reports episode averages. + + Unlike rewards, metrics have no weight, no dt scaling, and no + normalization by episode length. Episode values are true per-step + averages (sum / step_count), so a metric in [0,1] stays in [0,1] + in the logger. + """ + + _env: ManagerBasedRlEnv + + def __init__(self, cfg: dict[str, MetricsTermCfg | None], env: ManagerBasedRlEnv): + self._term_names: list[str] = list() + self._term_cfgs: list[MetricsTermCfg] = list() + self._class_term_cfgs: list[MetricsTermCfg] = list() + self._step_term_indices: list[int] = list() + self._substep_term_indices: list[int] = list() + + self.cfg = deepcopy(cfg) + super().__init__(env=env) + + self._episode_sums: dict[str, np.ndarray] = {} + self._episode_max: dict[str, np.ndarray] = {} + for idx, term_name in enumerate(self._term_names): + if self._term_cfgs[idx].reduce not in REDUCE_OPTIONS: + msg = ( + f"The reduce method '{self._term_cfgs[idx].reduce}' for metric '{term_name}' " + f"is unknown. Valid options are {REDUCE_OPTIONS}." + ) + raise ValueError(msg) + + self._episode_sums[term_name] = np.zeros(self.num_envs, dtype=np.float32) + + if self._term_cfgs[idx].reduce == "max": + self._episode_max[term_name] = np.full( + self.num_envs, float("-inf"), dtype=np.float32 + ) + # Pre-resolved tensor refs for substep terms to avoid dict lookups in + # the hot loop. + self._substep_accum: list[np.ndarray] = [] + self._substep_episode_sums: list[np.ndarray] = [] + self._substep_episode_max: list[np.ndarray | None] = [] + for idx in self._substep_term_indices: + name = self._term_names[idx] + buf = np.zeros(self.num_envs, dtype=np.float32) + self._substep_accum.append(buf) + self._substep_episode_sums.append(self._episode_sums[name]) + self._substep_episode_max.append(self._episode_max.get(name)) + self._substep_count: int = 0 + self._step_count = np.zeros(self.num_envs, dtype=np.int64) + self._step_values = np.zeros((self.num_envs, len(self._term_names)), dtype=np.float32) + + def __str__(self) -> str: + msg = f" contains {len(self._term_names)} active terms.\n" + table = PrettyTable() + table.title = "Active Metrics Terms" + table.field_names = ["Index", "Name"] + table.align["Name"] = "l" + for index, name in enumerate(self._term_names): + table.add_row([index, name]) + msg += str(table.get_string()) + msg += "\n" + return msg + + # Properties. + + @property + def active_terms(self) -> list[str]: + return self._term_names + + # Methods. + + def reset(self, env_ids: np.ndarray | slice | None = None) -> dict[str, float]: + if env_ids is None: + env_ids = slice(None) + extras = {} + counts = self._step_count[env_ids].astype(np.float32) + # Avoid division by zero for envs that haven't stepped. + safe_counts = np.maximum(counts, 1.0) + for idx, key in enumerate(self._episode_sums): + reduce = self._term_cfgs[idx].reduce + if reduce == "max": + extras["Episode_Metrics/" + key] = float(np.mean(self._episode_max[key][env_ids])) + self._episode_max[key][env_ids] = float("-inf") + + elif reduce == "last": + extras["Episode_Metrics/" + key] = float(np.mean(self._step_values[env_ids, idx])) + + elif reduce == "sum": + extras["Episode_Metrics/" + key] = float(np.mean(self._episode_sums[key][env_ids])) + + else: + extras["Episode_Metrics/" + key] = float( + np.mean(self._episode_sums[key][env_ids] / safe_counts) + ) + + self._episode_sums[key][env_ids] = 0.0 + self._step_count[env_ids] = 0 + + for buf in self._substep_accum: + buf[env_ids] = 0.0 + + for term_cfg in self._class_term_cfgs: + term_cfg.func.reset(env_ids=env_ids) + + return extras + + def compute_substep(self) -> None: + """Accumulate per-substep metric values inside the decimation loop. + + No-op when no ``per_substep`` terms are configured. + """ + if not self._substep_term_indices: + return + for i, idx in enumerate(self._substep_term_indices): + value = self._compute_term(idx) + self._substep_accum[i] += value + self._substep_count += 1 + + def compute(self) -> None: + self._step_count += 1 + if self._substep_term_indices and self._substep_count > 0: + for i, idx in enumerate(self._substep_term_indices): + avg = self._substep_accum[i] / self._substep_count + self._substep_episode_sums[i] += avg + self._step_values[:, idx] = avg + max_buf = self._substep_episode_max[i] + if max_buf is not None: + np.maximum(max_buf, avg, out=max_buf) + self._substep_accum[i].fill(0.0) + self._substep_count = 0 + for idx in self._step_term_indices: + name = self._term_names[idx] + value = self._compute_term(idx) + self._episode_sums[name] += value + self._step_values[:, idx] = value + if name in self._episode_max: + np.maximum(self._episode_max[name], value, out=self._episode_max[name]) + + def get_active_iterable_terms(self, env_idx: int) -> Sequence[tuple[str, Sequence[float]]]: + terms = [] + for idx, name in enumerate(self._term_names): + terms.append((name, [self._step_values[env_idx, idx].item()])) + return terms + + def _prepare_terms(self) -> None: + for term_name, term_cfg in self.cfg.items(): + if term_cfg is None: + print(f"term: {term_name} set to None, skipping...") + continue + self._resolve_common_term_cfg(term_name, term_cfg) + idx = len(self._term_names) + self._term_names.append(term_name) + self._term_cfgs.append(term_cfg) + if term_cfg.per_substep: + self._substep_term_indices.append(idx) + else: + self._step_term_indices.append(idx) + if hasattr(term_cfg.func, "reset") and callable(term_cfg.func.reset): + self._class_term_cfgs.append(term_cfg) + + def _compute_term(self, idx: int) -> np.ndarray: + name = self._term_names[idx] + term_cfg = self._term_cfgs[idx] + value = term_cfg.func(self._env, **term_cfg.params) + self._check_term_shape(name, value) + self._check_term_finite(name, value) + return value + + +class NullMetricsManager: + """Placeholder for absent metrics manager that safely no-ops all operations.""" + + def __init__(self): + self.active_terms: list[str] = [] + self.cfg = None + + def __str__(self) -> str: + return " (inactive)" + + def __repr__(self) -> str: + return "NullMetricsManager()" + + def get_active_iterable_terms(self, env_idx: int) -> Sequence[tuple[str, Sequence[float]]]: + return [] + + def reset(self, env_ids: np.ndarray | None = None) -> dict[str, float]: + return {} + + def compute_substep(self) -> None: + pass + + def compute(self) -> None: + pass diff --git a/src/unilab/managers/observation_manager.py b/src/unilab/managers/observation_manager.py new file mode 100644 index 000000000..b5c531ca3 --- /dev/null +++ b/src/unilab/managers/observation_manager.py @@ -0,0 +1,570 @@ +# Derived from mujocolab/mjlab v1.6.0 (0fb8a681), src/mjlab/managers/observation_manager.py. +# Copyright 2025, The mjlab Developers. +# Modified by UniLab for NumPy and UniLab contracts; licensed under Apache-2.0. +"""Observation manager for computing observations.""" + +from __future__ import annotations + +from copy import deepcopy +from dataclasses import dataclass +from typing import TYPE_CHECKING, Literal, Sequence + +import numpy as np +from prettytable import PrettyTable + +from unilab.managers._buffers import CircularBuffer, DelayBuffer +from unilab.managers._noise import noise_cfg, noise_model +from unilab.managers._noise.noise_cfg import NoiseCfg, NoiseModelCfg +from unilab.managers.manager_base import ManagerBase, ManagerTermBaseCfg + +if TYPE_CHECKING: + from unilab.managers._types import ManagerBasedRlEnv + + +@dataclass +class ObservationTermCfg(ManagerTermBaseCfg): + """Configuration for an observation term. + + Processing pipeline: compute → noise → clip → scale → delay → history. + Delay models sensor latency. History provides temporal context. Both are optional + and can be combined. + """ + + noise: NoiseCfg | NoiseModelCfg | None = None + """Noise model to apply to the observation.""" + + clip: tuple[float, float] | None = None + """Range (min, max) to clip the observation values.""" + + scale: tuple[float, ...] | float | np.ndarray | None = None + """Scaling factor(s) to multiply the observation by.""" + + delay_min_lag: int = 0 + """Minimum lag (in steps) for delayed observations. Lag sampled uniformly from + [min_lag, max_lag]. Convert to ms: lag * (1000 / control_hz).""" + + delay_max_lag: int = 0 + """Maximum lag (in steps) for delayed observations. Use min=max for constant delay.""" + + delay_per_env: bool = True + """If True, each environment samples its own lag. If False, all environments share + the same lag at each step.""" + + delay_hold_prob: float = 0.0 + """Probability of reusing the previous lag instead of resampling. Useful for + temporally correlated latency patterns.""" + + delay_update_period: int = 0 + """Resample lag every N steps (models multi-rate sensors). If 0, update every step.""" + + delay_per_env_phase: bool = True + """If True and update_period > 0, stagger update timing across envs to avoid + synchronized resampling.""" + + history_length: int = 0 + """Number of past observations to keep in history. 0 = no history.""" + + flatten_history_dim: bool = True + """Whether to flatten the history dimension into observation. + + When True and concatenate_terms=True, uses term-major ordering: + [A_t0, A_t1, ..., A_tH-1, B_t0, B_t1, ..., B_tH-1, ...] + See docs/source/observation.rst for details on ordering.""" + + +@dataclass +class ObservationGroupCfg: + """Configuration for an observation group. + + An observation group bundles multiple observation terms together. Groups are + typically used to separate observations for different purposes (e.g., "actor" + for the actor, "critic" for the value function). + """ + + terms: dict[str, ObservationTermCfg | None] + """Dictionary mapping term names to their configurations.""" + + concatenate_terms: bool = True + """Whether to concatenate all terms into a single tensor. If False, returns + a dict mapping term names to their individual tensors.""" + + concatenate_dim: int = -1 + """Dimension along which to concatenate terms. Default -1 (last dimension).""" + + enable_corruption: bool = False + """Whether to apply noise corruption to observations. Set to True during + training for domain randomization, False during evaluation.""" + + history_length: int | None = None + """Group-level history length override. If set, applies to all terms in + this group. If None, each term uses its own ``history_length`` setting.""" + + flatten_history_dim: bool = True + """Whether to flatten history into the observation dimension. If True, + observations have shape ``(num_envs, obs_dim * history_length)``. If False, + shape is ``(num_envs, history_length, obs_dim)``.""" + + nan_policy: Literal["disabled", "warn", "sanitize", "error"] = "error" + """NaN/Inf handling policy for observations in this group. + + - 'disabled': No checks (explicit opt-out) + - 'warn': Log warning with term name and env IDs, then sanitize (debugging) + - 'sanitize': Silent sanitization to 0.0 like reward manager (safe for production) + - 'error': Raise ValueError on NaN/Inf (strict development mode) + """ + + nan_check_per_term: bool = True + """If True, check each observation term individually to identify NaN source. + If False, check only the final concatenated output (faster but less informative). + Only applies when nan_policy != 'disabled'.""" + + +class ObservationManager(ManagerBase): + """Manages observation computation for the environment. + + The observation manager computes observations from multiple terms organized + into groups. Each term can have noise, clipping, scaling, delay, and history + applied. Groups can optionally concatenate their terms into a single tensor. + """ + + def __init__(self, cfg: dict[str, ObservationGroupCfg | None], env: ManagerBasedRlEnv): + self.cfg = deepcopy(cfg) + super().__init__(env=env) + + self._group_obs_dim: dict[str, tuple[int, ...] | list[tuple[int, ...]]] = dict() + + for group_name, group_term_dims in self._group_obs_term_dim.items(): + if self._group_obs_concatenate[group_name]: + term_dims = np.stack([np.asarray(dims) for dims in group_term_dims], axis=0) + if len(term_dims.shape) > 1: + if self._group_obs_concatenate_dim[group_name] >= 0: + dim = self._group_obs_concatenate_dim[group_name] - 1 + else: + dim = self._group_obs_concatenate_dim[group_name] + dim_sum = np.sum(term_dims[:, dim], axis=0) + term_dims[0, dim] = dim_sum + term_dims = term_dims[0] + else: + term_dims = np.sum(term_dims, axis=0) + self._group_obs_dim[group_name] = tuple(term_dims.tolist()) + else: + self._group_obs_dim[group_name] = group_term_dims + + self._obs_buffer: dict[str, np.ndarray | dict[str, np.ndarray]] | None = None + + def __str__(self) -> str: + msg = f" contains {len(self._group_obs_term_names)} groups.\n" + for group_name, group_dim in self._group_obs_dim.items(): + table = PrettyTable() + table.title = f"Active Observation Terms in Group: '{group_name}'" + if self._group_obs_concatenate[group_name]: + table.title += f" (shape: {group_dim})" # type: ignore + table.field_names = ["Index", "Name", "Shape"] + table.align["Name"] = "l" + obs_terms = zip( + self._group_obs_term_names[group_name], + self._group_obs_term_dim[group_name], + self._group_obs_term_cfgs[group_name], + strict=False, + ) + for index, (name, dims, term_cfg) in enumerate(obs_terms): + if term_cfg.history_length > 0 and term_cfg.flatten_history_dim: + # Flattened history: show (9,) ← 3×(3,) + original_size = int(np.prod(dims)) // term_cfg.history_length + original_shape = (original_size,) if len(dims) == 1 else dims[1:] + shape_str = f"{dims} ← {term_cfg.history_length}×{original_shape}" + else: + shape_str = str(tuple(dims)) + table.add_row([index, name, shape_str]) + msg += str(table.get_string()) + msg += "\n" + return msg + + def get_active_iterable_terms(self, env_idx: int) -> Sequence[tuple[str, Sequence[float]]]: + terms = [] + + if self._obs_buffer is None: + self.compute() + assert self._obs_buffer is not None + obs_buffer: dict[str, np.ndarray | dict[str, np.ndarray]] = self._obs_buffer + + for group_name, _ in self.group_obs_dim.items(): + if not self.group_obs_concatenate[group_name]: + buffers = obs_buffer[group_name] + assert isinstance(buffers, dict) + for name, term in buffers.items(): + terms.append((group_name + "-" + name, term[env_idx].tolist())) + continue + + idx = 0 + data = obs_buffer[group_name] + assert isinstance(data, np.ndarray) + for name, shape in zip( + self._group_obs_term_names[group_name], + self._group_obs_term_dim[group_name], + strict=False, + ): + data_length = np.prod(shape) + term = data[env_idx, idx : idx + data_length] + terms.append((group_name + "-" + name, term.tolist())) + idx += data_length + + return terms + + # Properties. + + @property + def active_terms(self) -> dict[str, list[str]]: + return self._group_obs_term_names + + @property + def group_obs_dim(self) -> dict[str, tuple[int, ...] | list[tuple[int, ...]]]: + return self._group_obs_dim + + @property + def group_obs_term_dim(self) -> dict[str, list[tuple[int, ...]]]: + return self._group_obs_term_dim + + @property + def group_obs_concatenate(self) -> dict[str, bool]: + return self._group_obs_concatenate + + # Methods. + + def get_term_cfg(self, group_name: str, term_name: str) -> ObservationTermCfg: + if group_name not in self._group_obs_term_names: + raise ValueError(f"Group '{group_name}' not found in active groups.") + if term_name not in self._group_obs_term_names[group_name]: + raise ValueError(f"Term '{term_name}' not found in group '{group_name}'.") + index = self._group_obs_term_names[group_name].index(term_name) + return self._group_obs_term_cfgs[group_name][index] + + def reset(self, env_ids: np.ndarray | slice | None = None) -> dict[str, float]: + # Invalidate cache since reset envs will have different observations. + self._obs_buffer = None + + for group_name, group_cfg in self._group_obs_class_term_cfgs.items(): + for term_cfg in group_cfg: + term_cfg.func.reset(env_ids=env_ids) + for term_name in self._group_obs_term_names[group_name]: + batch_ids = env_ids + if term_name in self._group_obs_term_delay_buffer[group_name]: + self._group_obs_term_delay_buffer[group_name][term_name].reset( + batch_ids=batch_ids + ) + if term_name in self._group_obs_term_history_buffer[group_name]: + self._group_obs_term_history_buffer[group_name][term_name].reset( + batch_ids=batch_ids + ) + for group_mods in self._group_obs_class_instances.values(): + for mod in group_mods.values(): + mod.reset(env_ids=env_ids) + return {} + + def _check_and_handle_nans(self, tensor: np.ndarray, context: str, policy: str) -> np.ndarray: + """Check for NaN/Inf and handle according to policy. + + Args: + tensor: Observation tensor to check. + context: Context string for error/warning messages (e.g., "actor/base_lin_vel"). + policy: NaN handling policy ("disabled", "warn", "sanitize", "error"). + + Returns: + The tensor, potentially sanitized depending on policy. + + Raises: + ValueError: If policy is "error" and NaN/Inf detected. + """ + if policy == "disabled": + return tensor + + has_nan = np.isnan(tensor).any() + has_inf = np.isinf(tensor).any() + + if not (has_nan or has_inf): + return tensor + + if policy == "error": + invalid = ~np.isfinite(tensor) + nan_mask = invalid.reshape(self.num_envs, -1).any(axis=1) + nan_env_ids = np.flatnonzero(nan_mask).tolist() + invalid_kind = ( + "NaN" + if has_nan and not has_inf + else "Inf" + if has_inf and not has_nan + else "NaN/Inf" + ) + raise ValueError( + f"{invalid_kind} detected in ObservationManager term '{context}' " + f"for environments: {nan_env_ids[:10]}" + ) + + if policy == "warn": + invalid = ~np.isfinite(tensor) + nan_mask = invalid.reshape(self.num_envs, -1).any(axis=1) + nan_env_ids = np.flatnonzero(nan_mask).tolist() + print( + f"[ObservationManager] NaN/Inf in '{context}' " + f"(envs: {nan_env_ids[:5]}). Sanitizing to 0." + ) + + # Sanitize (applies to both "warn" and "sanitize" policies). + return np.nan_to_num(tensor, nan=0.0, posinf=0.0, neginf=0.0) + + def compute( + self, + update_history: bool = False, + env_ids: np.ndarray | None = None, + ) -> dict[str, np.ndarray | dict[str, np.ndarray]]: + """Compute observations for all groups. + + With env_ids=None (the per-step path), history and delay buffers advance + for all envs. With env_ids (the reset path), only the reset envs' buffers + receive their post-reset frame (a backfill); other envs' buffers, delay + schedules, and lag draws are untouched, so a partial reset does not + advance their observation timelines. + """ + if env_ids is not None and not update_history: + raise ValueError("env_ids is only meaningful with update_history=True.") + # Return cached observations if not updating and cache exists. + # This prevents double-pushing to delay buffers when compute() is called + # multiple times per control step (e.g., in get_observations() after step()). + if not update_history and self._obs_buffer is not None: + return self._obs_buffer + + obs_buffer: dict[str, np.ndarray | dict[str, np.ndarray]] = dict() + for group_name in self._group_obs_term_names: + obs_buffer[group_name] = self.compute_group(group_name, update_history, env_ids) + self._obs_buffer = obs_buffer + return obs_buffer + + def compute_group( + self, + group_name: str, + update_history: bool = False, + env_ids: np.ndarray | None = None, + ) -> np.ndarray | dict[str, np.ndarray]: + group_cfg = self.cfg[group_name] + if group_cfg is None: + raise KeyError(f"Observation group '{group_name}' is disabled.") + group_term_names = self._group_obs_term_names[group_name] + group_obs: dict[str, np.ndarray] = {} + obs_terms = zip(group_term_names, self._group_obs_term_cfgs[group_name], strict=False) + for term_name, term_cfg in obs_terms: + obs = term_cfg.func(self._env, **term_cfg.params) + if not isinstance(obs, np.ndarray): + raise TypeError( + f"ObservationManager term '{group_name}/{term_name}' returned " + f"{type(obs).__name__}, expected np.ndarray." + ) + if obs.ndim < 2 or obs.shape[0] != self.num_envs: + raise ValueError( + f"ObservationManager term '{group_name}/{term_name}' returned shape " + f"{obs.shape}, expected (num_envs, ...) with num_envs={self.num_envs}." + ) + obs = obs.copy() + if isinstance(term_cfg.noise, noise_cfg.NoiseCfg): + obs = term_cfg.noise.apply(obs, rng=self._env.rng) + elif isinstance(term_cfg.noise, noise_cfg.NoiseModelCfg): + obs = self._group_obs_class_instances[group_name][term_name](obs) + if term_cfg.clip: + np.clip(obs, term_cfg.clip[0], term_cfg.clip[1], out=obs) + if term_cfg.scale is not None: + scale = term_cfg.scale + assert isinstance(scale, np.ndarray) + np.multiply(obs, scale, out=obs) + + # Check for NaN/Inf before delay/history buffers (per-term checking). + if group_cfg.nan_check_per_term and group_cfg.nan_policy != "disabled": + obs = self._check_and_handle_nans( + obs, context=f"{group_name}/{term_name}", policy=group_cfg.nan_policy + ) + + if term_cfg.delay_max_lag > 0: + delay_buffer = self._group_obs_term_delay_buffer[group_name][term_name] + if env_ids is None or not delay_buffer.is_initialized: + delay_buffer.append(obs) + obs = delay_buffer.compute() + else: + delay_buffer.backfill(obs, env_ids) + obs = delay_buffer.peek() + if term_cfg.history_length > 0: + circular_buffer = self._group_obs_term_history_buffer[group_name][term_name] + if env_ids is None or not circular_buffer.is_initialized: + if update_history or not circular_buffer.is_initialized: + circular_buffer.append(obs) + else: + circular_buffer.backfill(obs, env_ids) + + if term_cfg.flatten_history_dim: + group_obs[term_name] = circular_buffer.buffer.reshape(self._env.num_envs, -1) + else: + group_obs[term_name] = circular_buffer.buffer + else: + group_obs[term_name] = obs + + # Final NaN check for non-per-term checking. + if not group_cfg.nan_check_per_term and group_cfg.nan_policy != "disabled": + if self._group_obs_concatenate[group_name]: + # Will check after concatenation below. + pass + else: + for term_name in group_obs: + group_obs[term_name] = self._check_and_handle_nans( + group_obs[term_name], + context=f"{group_name}/{term_name}", + policy=group_cfg.nan_policy, + ) + + if self._group_obs_concatenate[group_name]: + result = np.concatenate( + list(group_obs.values()), axis=self._group_obs_concatenate_dim[group_name] + ) + # Final check for concatenated result (non-per-term checking). + if not group_cfg.nan_check_per_term and group_cfg.nan_policy != "disabled": + result = self._check_and_handle_nans( + result, context=group_name, policy=group_cfg.nan_policy + ) + return result + return group_obs + + def _prepare_terms(self) -> None: + self._group_obs_term_names: dict[str, list[str]] = dict() + self._group_obs_term_dim: dict[str, list[tuple[int, ...]]] = dict() + self._group_obs_term_cfgs: dict[str, list[ObservationTermCfg]] = dict() + self._group_obs_class_term_cfgs: dict[str, list[ObservationTermCfg]] = dict() + self._group_obs_concatenate: dict[str, bool] = dict() + self._group_obs_concatenate_dim: dict[str, int] = dict() + self._group_obs_class_instances: dict[str, dict[str, noise_model.NoiseModel]] = {} + self._group_obs_term_delay_buffer: dict[str, dict[str, DelayBuffer]] = dict() + self._group_obs_term_history_buffer: dict[str, dict[str, CircularBuffer]] = dict() + + for group_name, group_cfg in self.cfg.items(): + if group_cfg is None: + print(f"group: {group_name} set to None, skipping...") + continue + + if not any(t is not None for t in group_cfg.terms.values()): + print(f"group: {group_name} has no active terms, skipping...") + continue + + if group_cfg.nan_policy not in ("disabled", "warn", "sanitize", "error"): + raise ValueError( + f"Observation group '{group_name}' has unsupported NaN policy " + f"'{group_cfg.nan_policy}'." + ) + if group_cfg.history_length is not None and group_cfg.history_length < 0: + raise ValueError( + f"Observation group '{group_name}' has negative history_length " + f"{group_cfg.history_length}." + ) + + self._group_obs_term_names[group_name] = list() + self._group_obs_term_dim[group_name] = list() + self._group_obs_term_cfgs[group_name] = list() + self._group_obs_class_term_cfgs[group_name] = list() + self._group_obs_class_instances[group_name] = {} + group_entry_delay_buffer: dict[str, DelayBuffer] = dict() + group_entry_history_buffer: dict[str, CircularBuffer] = dict() + + self._group_obs_concatenate[group_name] = group_cfg.concatenate_terms + self._group_obs_concatenate_dim[group_name] = ( + group_cfg.concatenate_dim + 1 + if group_cfg.concatenate_dim >= 0 + else group_cfg.concatenate_dim + ) + + for term_name, term_cfg in group_cfg.terms.items(): + if term_cfg is None: + print(f"term: {term_name} set to None, skipping...") + continue + + if term_cfg.delay_min_lag < 0 or term_cfg.delay_max_lag < term_cfg.delay_min_lag: + raise ValueError( + f"ObservationManager term '{group_name}/{term_name}' has invalid " + f"delay range [{term_cfg.delay_min_lag}, {term_cfg.delay_max_lag}]." + ) + if term_cfg.history_length < 0: + raise ValueError( + f"ObservationManager term '{group_name}/{term_name}' has negative " + f"history_length {term_cfg.history_length}." + ) + if term_cfg.clip is not None and term_cfg.clip[0] > term_cfg.clip[1]: + raise ValueError( + f"ObservationManager term '{group_name}/{term_name}' has invalid " + f"clip range {term_cfg.clip}." + ) + + # NOTE: This deepcopy is important to avoid cross-group contamination of term + # configs. + term_cfg = deepcopy(term_cfg) + self._resolve_common_term_cfg(term_name, term_cfg) + + if not group_cfg.enable_corruption: + term_cfg.noise = None + if group_cfg.history_length is not None: + term_cfg.history_length = group_cfg.history_length + term_cfg.flatten_history_dim = group_cfg.flatten_history_dim + self._group_obs_term_names[group_name].append(term_name) + self._group_obs_term_cfgs[group_name].append(term_cfg) + if hasattr(term_cfg.func, "reset") and callable(term_cfg.func.reset): + self._group_obs_class_term_cfgs[group_name].append(term_cfg) + + initial_obs = term_cfg.func(self._env, **term_cfg.params) + if not isinstance(initial_obs, np.ndarray): + raise TypeError( + f"ObservationManager term '{group_name}/{term_name}' returned " + f"{type(initial_obs).__name__}, expected np.ndarray." + ) + if initial_obs.ndim < 2 or initial_obs.shape[0] != self.num_envs: + raise ValueError( + f"ObservationManager term '{group_name}/{term_name}' returned shape " + f"{initial_obs.shape}, expected (num_envs, ...) with " + f"num_envs={self.num_envs}." + ) + obs_dims = tuple(initial_obs.shape) + + if term_cfg.scale is not None: + term_cfg.scale = np.asarray(term_cfg.scale, dtype=np.float32).copy() + + if term_cfg.noise is not None and isinstance( + term_cfg.noise, noise_cfg.NoiseModelCfg + ): + noise_model_cls = term_cfg.noise.class_type + if not issubclass(noise_model_cls, noise_model.NoiseModel): + raise TypeError( + f"ObservationManager term '{group_name}/{term_name}' noise model " + f"{noise_model_cls} is not a NoiseModel subclass." + ) + self._group_obs_class_instances[group_name][term_name] = noise_model_cls( + term_cfg.noise, num_envs=self._env.num_envs, rng=self._env.rng + ) + + if term_cfg.delay_max_lag > 0: + group_entry_delay_buffer[term_name] = DelayBuffer( + min_lag=term_cfg.delay_min_lag, + max_lag=term_cfg.delay_max_lag, + batch_size=self._env.num_envs, + per_env=term_cfg.delay_per_env, + hold_prob=term_cfg.delay_hold_prob, + update_period=term_cfg.delay_update_period, + per_env_phase=term_cfg.delay_per_env_phase, + generator=self._env.rng, + ) + + if term_cfg.history_length > 0: + group_entry_history_buffer[term_name] = CircularBuffer( + max_len=term_cfg.history_length, + batch_size=self._env.num_envs, + ) + old_dims = list(obs_dims) + old_dims.insert(1, term_cfg.history_length) + obs_dims = tuple(old_dims) + if term_cfg.flatten_history_dim: + obs_dims = (obs_dims[0], int(np.prod(obs_dims[1:]))) + + self._group_obs_term_dim[group_name].append(obs_dims[1:]) + + self._group_obs_term_delay_buffer[group_name] = group_entry_delay_buffer + self._group_obs_term_history_buffer[group_name] = group_entry_history_buffer diff --git a/src/unilab/managers/recorder_manager.py b/src/unilab/managers/recorder_manager.py new file mode 100644 index 000000000..a2a379921 --- /dev/null +++ b/src/unilab/managers/recorder_manager.py @@ -0,0 +1,267 @@ +# Derived from mujocolab/mjlab v1.6.0 (0fb8a681), src/mjlab/managers/recorder_manager.py. +# Copyright 2025, The mjlab Developers. +# Modified by UniLab for NumPy and UniLab contracts; licensed under Apache-2.0. +"""Recorder manager for logging environment data during rollouts.""" + +from __future__ import annotations + +from copy import deepcopy +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import numpy as np +from prettytable import PrettyTable + +from unilab.managers.manager_base import ManagerBase, ManagerTermBase, ManagerTermBaseCfg + +if TYPE_CHECKING: + from unilab.managers._types import ManagerBasedRlEnv + + +@dataclass +class RecorderTermCfg(ManagerTermBaseCfg): + """Configuration for a recorder term. + + ``func`` must be a :class:`RecorderTerm` subclass. Function-based terms are not + supported because recorder terms are stateful (file handles, buffers, etc.). + """ + + +class RecorderTerm(ManagerTermBase): + """Base class for recorder terms. + + Override only the lifecycle methods you need. Each method is a no-op by default so + subclasses are not required to implement all of them. + + The environment is available as ``self._env``, giving access to ``self._env.obs_buf``, + ``self._env.action_manager.action``, and all other environment state. + + Example:: + + class CsvRecorder(RecorderTerm): + def __init__(self, cfg, env): + super().__init__(cfg, env) + self._file = open(cfg.params["path"], "w", newline="") + self._writer = csv.writer(self._file) + + def record_pre_reset(self, env_ids): + # Terminal transition: action is still intact here. + # It will be zeroed by _reset_idx immediately after this returns. + obs = self._env.obs_buf["actor"][env_ids] + act = self._env.action_manager.action[env_ids] + for o, a in zip(obs, act): + self._writer.writerow(o.tolist() + a.tolist()) + + def record_post_step(self): + # Skip envs that just reset: their terminal pair was written in record_pre_reset + # and their action is now zeroed. + mask = ~self._env.reset_buf + obs = self._env.obs_buf["actor"][mask] + act = self._env.action_manager.action[mask] + for o, a in zip(obs, act): + self._writer.writerow(o.tolist() + a.tolist()) + + def close(self): + self._file.close() + """ + + def __init__(self, cfg: RecorderTermCfg, env: ManagerBasedRlEnv): + super().__init__(env) # ManagerTermBase only accepts env + self.cfg = cfg + + def record_pre_reset(self, env_ids: np.ndarray) -> None: + """Called in ``env.step()`` before terminated environments are reset. + + **What is available:** + + - ``obs_buf`` contains the observation from the *end of the previous step* (the + input the agent used to choose the terminal action). It does **not** contain the + post-action terminal observation (the state reached after applying the action), + which is never computed for resetting environments. + - ``action_manager.action`` contains the action applied during this step. This is + the correct terminal action. It will be zeroed for these environments by + ``_reset_idx`` immediately after this hook returns, so capture it here if you + need it later. + - ``reward_buf`` contains the reward for this terminal step. + - ``reset_terminated`` and ``reset_time_outs`` reflect why each environment is + resetting. + + This is the right hook to record the terminal transition + ``(obs_t, action_t, reward_t, done=True)`` for each resetting environment. + + Args: + env_ids: Indices of environments that are about to be reset. + """ + del env_ids # Unused in base implementation. + + def record_post_reset(self, env_ids: np.ndarray) -> None: + """Called after a reset completes with fresh observations computed. + + Fires at the end of ``env.reset()`` (covering all environments on the initial call) + and within ``env.step()`` for each batch of environments that terminates, after + state has been overwritten and new observations computed. + + At this point ``obs_buf[env_ids]`` holds the initial observation of the new episode + and ``action_manager.action[env_ids]`` is zero (no action has been taken in the new + episode yet). + + Use this hook to initialize per-episode state or record the first observation of a + new episode. + + Args: + env_ids: Indices of environments that were reset. + """ + del env_ids # Unused in base implementation. + + def record_post_step(self) -> None: + """Called at the end of every ``env.step()`` with fresh observations. + + At this point ``obs_buf`` holds the new observation for every environment and + ``action_manager.action`` holds the action that was applied during this step. + **Exception:** for environments that reset during this step, + ``action_manager.action`` has been zeroed by ``_reset_idx`` and ``obs_buf`` holds + the initial observation of the new episode rather than the post-action terminal + observation. Use ``record_pre_reset`` to capture the terminal ``(obs, action)`` + pair for those environments. Resetting environments are identified by + ``self._env.reset_buf``. + """ + + def close(self) -> None: + """Called when the environment closes. + + Release file handles, flush write buffers, or finalize output here. + """ + + def __call__(self): + raise NotImplementedError( + "RecorderTerm is not invoked via __call__. Override the lifecycle methods instead." + ) + + +class RecorderManager(ManagerBase): + """Orchestrates recorder terms during environment rollouts. + + Holds a collection of :class:`RecorderTerm` instances and calls their lifecycle + methods at the appropriate points in the environment loop. The manager has no opinion + on how data is stored; each term handles its own I/O entirely. + + Register terms by adding them to the ``recorders`` dict on + :class:`ManagerBasedRlEnvCfg`. If the dict is empty, the environment + substitutes a :class:`NullRecorderManager` with zero overhead. + """ + + def __init__(self, cfg: dict[str, RecorderTermCfg | None], env: ManagerBasedRlEnv): + self._terms: dict[str, RecorderTerm] = {} + self.cfg = deepcopy(cfg) + super().__init__(env) # calls _prepare_terms() + + def __str__(self) -> str: + msg = f" contains {len(self._terms)} active terms.\n" + table = PrettyTable() + table.title = "Active Recorder Terms" + table.field_names = ["Index", "Name"] + for idx, name in enumerate(self._terms): + table.add_row([idx, name]) + return msg + str(table.get_string()) + + def __contains__(self, name: str) -> bool: + """Return True if a term named ``name`` is registered.""" + return name in self._terms + + @property + def active_terms(self) -> list[str]: + """List of active term names.""" + return list(self._terms.keys()) + + def get_term(self, name: str) -> RecorderTerm: + """Return the recorder term registered under ``name``. + + Use this to reach a recorder's public methods (e.g. to start/stop logging) + from outside the env loop without touching private state. + + Args: + name: Term name as registered in ``ManagerBasedRlEnvCfg.recorders``. + + Raises: + KeyError: If no term is registered under ``name``. + """ + try: + return self._terms[name] + except KeyError: + msg = f"No recorder term named '{name}'. Active terms: {self.active_terms}" + raise KeyError(msg) from None + + def record_pre_reset(self, env_ids: np.ndarray) -> None: + """Forward to each term's :meth:`RecorderTerm.record_pre_reset`.""" + for term in self._terms.values(): + term.record_pre_reset(env_ids) + + def record_post_reset(self, env_ids: np.ndarray) -> None: + """Forward to each term's :meth:`RecorderTerm.record_post_reset`.""" + for term in self._terms.values(): + term.record_post_reset(env_ids) + + def record_post_step(self) -> None: + """Forward to each term's :meth:`RecorderTerm.record_post_step`.""" + for term in self._terms.values(): + term.record_post_step() + + def close(self) -> None: + """Forward to each term's :meth:`RecorderTerm.close`.""" + for term in self._terms.values(): + term.close() + + def _prepare_terms(self) -> None: + for name, cfg in self.cfg.items(): + if cfg is None: + continue + self._resolve_common_term_cfg(name, cfg) + # _resolve_common_term_cfg instantiates class-based terms in-place. + if not isinstance(cfg.func, RecorderTerm): + raise TypeError( + f"Recorder term '{name}': func must be a RecorderTerm subclass," + f" got {type(cfg.func).__name__}. Function-based terms are not" + " supported." + ) + self._terms[name] = cfg.func + + +class NullRecorderManager: + """No-op fallback used when no recorder terms are configured. + + All methods are no-ops. This class is not a :class:`ManagerBase` subclass + so it carries zero overhead. + """ + + def __init__(self): + self.active_terms: list[str] = [] + self.cfg = None + + def __str__(self) -> str: + return " (inactive)" + + def __repr__(self) -> str: + return "NullRecorderManager()" + + def __contains__(self, name: str) -> bool: + """Always returns False since there are no terms.""" + del name + return False + + def get_term(self, name: str) -> RecorderTerm: + """Always raises KeyError since there are no terms.""" + del name + msg = "NullRecorderManager has no terms." + raise KeyError(msg) + + def record_pre_reset(self, env_ids: np.ndarray) -> None: + del env_ids + + def record_post_reset(self, env_ids: np.ndarray) -> None: + del env_ids + + def record_post_step(self) -> None: + pass + + def close(self) -> None: + pass diff --git a/src/unilab/managers/reward_manager.py b/src/unilab/managers/reward_manager.py new file mode 100644 index 000000000..67b2a5a1c --- /dev/null +++ b/src/unilab/managers/reward_manager.py @@ -0,0 +1,154 @@ +# Derived from mujocolab/mjlab v1.6.0 (0fb8a681), src/mjlab/managers/reward_manager.py. +# Copyright 2025, The mjlab Developers. +# Modified by UniLab for NumPy and UniLab contracts; licensed under Apache-2.0. +"""Reward manager for computing reward signals.""" + +from __future__ import annotations + +from copy import deepcopy +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +import numpy as np +from prettytable import PrettyTable + +from unilab.managers.manager_base import ManagerBase, ManagerTermBaseCfg + +if TYPE_CHECKING: + from unilab.managers._types import DebugVisualizer, ManagerBasedRlEnv + + +@dataclass(kw_only=True) +class RewardTermCfg(ManagerTermBaseCfg): + """Configuration for a reward term.""" + + func: Any + """The callable that computes this reward term's value.""" + + weight: float + """Weight multiplier for this reward term.""" + + +class RewardManager(ManagerBase): + """Manages reward computation by aggregating weighted reward terms. + + Reward Scaling Behavior: + By default, rewards are scaled by the environment step duration (dt). This + normalizes cumulative episodic rewards across different simulation frequencies. + The scaling can be disabled via the ``scale_by_dt`` parameter. + + When ``scale_by_dt=True`` (default): + - ``reward_buf`` (returned by ``compute()``) = raw_value * weight * dt + - ``_episode_sums`` (cumulative rewards) are scaled by dt + - ``Episode_Reward/*`` logged metrics are scaled by dt + + When ``scale_by_dt=False``: + - ``reward_buf`` = raw_value * weight (no dt scaling) + + Regardless of the scaling setting: + - ``_step_reward`` (via ``get_active_iterable_terms()``) always contains + the unscaled reward rate (raw_value * weight) + """ + + _env: ManagerBasedRlEnv + + def __init__( + self, + cfg: dict[str, RewardTermCfg | None], + env: ManagerBasedRlEnv, + *, + scale_by_dt: bool = True, + ): + self._term_names: list[str] = list() + self._term_cfgs: list[RewardTermCfg] = list() + self._class_term_cfgs: list[RewardTermCfg] = list() + self._scale_by_dt = scale_by_dt + + self.cfg = deepcopy(cfg) + super().__init__(env=env) + self._episode_sums = dict() + for term_name in self._term_names: + self._episode_sums[term_name] = np.zeros(self.num_envs, dtype=np.float32) + self._reward_buf = np.zeros(self.num_envs, dtype=np.float32) + self._step_reward = np.zeros((self.num_envs, len(self._term_names)), dtype=np.float32) + + def __str__(self) -> str: + msg = f" contains {len(self._term_names)} active terms.\n" + table = PrettyTable() + table.title = "Active Reward Terms" + table.field_names = ["Index", "Name", "Weight"] + table.align["Name"] = "l" + table.align["Weight"] = "r" + for index, (name, term_cfg) in enumerate( + zip(self._term_names, self._term_cfgs, strict=False) + ): + table.add_row([index, name, term_cfg.weight]) + msg += str(table.get_string()) + msg += "\n" + return msg + + # Properties. + + @property + def active_terms(self) -> list[str]: + return self._term_names + + # Methods. + + def reset(self, env_ids: np.ndarray | slice | None = None) -> dict[str, float]: + if env_ids is None: + env_ids = slice(None) + extras = {} + for key in self._episode_sums.keys(): + episodic_sum_avg = float(np.mean(self._episode_sums[key][env_ids])) + extras["Episode_Reward/" + key] = episodic_sum_avg / self._env.max_episode_length_s + self._episode_sums[key][env_ids] = 0.0 + for term_cfg in self._class_term_cfgs: + term_cfg.func.reset(env_ids=env_ids) + return extras + + def compute(self, dt: float) -> np.ndarray: + if not np.isfinite(dt) or (self._scale_by_dt and dt <= 0.0): + raise ValueError(f"RewardManager received invalid dt {dt}.") + self._reward_buf[:] = 0.0 + scale = dt if self._scale_by_dt else 1.0 + for term_idx, (name, term_cfg) in enumerate( + zip(self._term_names, self._term_cfgs, strict=False) + ): + if term_cfg.weight == 0.0: + self._step_reward[:, term_idx] = 0.0 + continue + value = term_cfg.func(self._env, **term_cfg.params) + self._check_term_shape(name, value) + self._check_term_finite(name, value) + value = value * term_cfg.weight * scale + self._reward_buf += value + self._episode_sums[name] += value + self._step_reward[:, term_idx] = value / scale + return self._reward_buf + + def get_active_iterable_terms(self, env_idx: int) -> list[tuple[str, list[float]]]: + terms = [] + for idx, name in enumerate(self._term_names): + terms.append((name, [self._step_reward[env_idx, idx].item()])) + return terms + + def get_term_cfg(self, term_name: str) -> RewardTermCfg: + if term_name not in self._term_names: + raise ValueError(f"Term '{term_name}' not found in active terms.") + return self._term_cfgs[self._term_names.index(term_name)] + + def _prepare_terms(self) -> None: + for term_name, term_cfg in self.cfg.items(): + if term_cfg is None: + print(f"term: {term_name} set to None, skipping...") + continue + if not np.isfinite(term_cfg.weight): + raise ValueError( + f"RewardManager term '{term_name}' has non-finite weight {term_cfg.weight}." + ) + self._resolve_common_term_cfg(term_name, term_cfg) + self._term_names.append(term_name) + self._term_cfgs.append(term_cfg) + if hasattr(term_cfg.func, "reset") and callable(term_cfg.func.reset): + self._class_term_cfgs.append(term_cfg) diff --git a/src/unilab/managers/scene_entity_config.py b/src/unilab/managers/scene_entity_config.py new file mode 100644 index 000000000..cb90c84d4 --- /dev/null +++ b/src/unilab/managers/scene_entity_config.py @@ -0,0 +1,263 @@ +# Derived from mujocolab/mjlab v1.6.0 (0fb8a681), src/mjlab/managers/scene_entity_config.py. +# Copyright 2025, The mjlab Developers. +# Modified by UniLab for NumPy and UniLab contracts; licensed under Apache-2.0. +"""Configuration for scene entities used by manager terms.""" + +from dataclasses import dataclass, field +from typing import NamedTuple + +from unilab.managers._types import ManagerEntity, ManagerScene + + +class _FieldConfig(NamedTuple): + """Configuration for a resolvable entity field.""" + + names_attr: str + ids_attr: str + find_method: str + num_attr: str + kind_label: str + + +_FIELD_CONFIGS = [ + _FieldConfig("joint_names", "joint_ids", "find_joints", "num_joints", "joint"), + _FieldConfig("body_names", "body_ids", "find_bodies", "num_bodies", "body"), + _FieldConfig("geom_names", "geom_ids", "find_geoms", "num_geoms", "geom"), + _FieldConfig("site_names", "site_ids", "find_sites", "num_sites", "site"), + _FieldConfig("actuator_names", "actuator_ids", "find_actuators", "num_actuators", "actuator"), + _FieldConfig("tendon_names", "tendon_ids", "find_tendons", "num_tendons", "tendon"), + _FieldConfig( + "camera_names", + "camera_ids", + "find_cameras", + "num_cameras", + "camera", + ), + _FieldConfig( + "light_names", + "light_ids", + "find_lights", + "num_lights", + "light", + ), + _FieldConfig( + "material_names", + "material_ids", + "find_materials", + "num_materials", + "material", + ), + _FieldConfig( + "texture_names", + "texture_ids", + "find_textures", + "num_textures", + "texture", + ), + _FieldConfig("pair_names", "pair_ids", "find_pairs", "num_pairs", "pair"), +] + + +@dataclass +class SceneEntityCfg: + """Configuration for a scene entity that is used by the manager's term. + + This configuration allows flexible specification of entity components either by name + or by ID. During resolution, it ensures consistency between names and IDs, and can + optimize to slice(None) when all components are selected. + """ + + name: str + """The name of the entity in the scene.""" + + joint_names: str | tuple[str, ...] | None = None + """Names of joints to include. Can be a single string or tuple.""" + + joint_ids: list[int] | slice = field(default_factory=lambda: slice(None)) + """IDs of joints to include. Can be a list or slice.""" + + body_names: str | tuple[str, ...] | None = None + """Names of bodies to include. Can be a single string or tuple.""" + + body_ids: list[int] | slice = field(default_factory=lambda: slice(None)) + """IDs of bodies to include. Can be a list or slice.""" + + geom_names: str | tuple[str, ...] | None = None + """Names of geometries to include. Can be a single string or tuple.""" + + geom_ids: list[int] | slice = field(default_factory=lambda: slice(None)) + """IDs of geometries to include. Can be a list or slice.""" + + site_names: str | tuple[str, ...] | None = None + """Names of sites to include. Can be a single string or tuple.""" + + site_ids: list[int] | slice = field(default_factory=lambda: slice(None)) + """IDs of sites to include. Can be a list or slice.""" + + actuator_names: str | list[str] | None = None + """Names of actuators to include. Can be a single string or list.""" + + actuator_ids: list[int] | slice = field(default_factory=lambda: slice(None)) + """IDs of actuators to include. Can be a list or slice.""" + + tendon_names: str | tuple[str, ...] | None = None + """Names of tendons to include. Can be a single string or tuple.""" + + tendon_ids: list[int] | slice = field(default_factory=lambda: slice(None)) + """IDs of tendons to include. Can be a list or slice.""" + + camera_names: str | tuple[str, ...] | None = None + """Names of cameras to include. Can be a single string or tuple.""" + + camera_ids: list[int] | slice = field(default_factory=lambda: slice(None)) + """IDs of cameras to include. Can be a list or slice.""" + + light_names: str | tuple[str, ...] | None = None + """Names of lights to include. Can be a single string or tuple.""" + + light_ids: list[int] | slice = field(default_factory=lambda: slice(None)) + """IDs of lights to include. Can be a list or slice.""" + + material_names: str | tuple[str, ...] | None = None + """Names of materials to include. Can be a single string or tuple.""" + + material_ids: list[int] | slice = field(default_factory=lambda: slice(None)) + """IDs of materials to include. Can be a list or slice.""" + + texture_names: str | tuple[str, ...] | None = None + """Names of textures to include. Can be a single string or tuple.""" + + texture_ids: list[int] | slice = field(default_factory=lambda: slice(None)) + """IDs of textures to include. Can be a list or slice.""" + + pair_names: str | tuple[str, ...] | None = None + """Names of contact pairs to include. Can be a single string or tuple.""" + + pair_ids: list[int] | slice = field(default_factory=lambda: slice(None)) + """IDs of contact pairs to include. Can be a list or slice.""" + + preserve_order: bool = False + """If True, maintains the order of components as specified.""" + + def resolve(self, scene: ManagerScene) -> None: + """Resolve names and IDs for all configured fields. + + This method ensures consistency between names and IDs for each field type. + It handles three cases: + 1. Both names and IDs provided: Validates they match + 2. Only names provided: Computes IDs (optimizes to slice(None) if all selected) + 3. Only IDs provided: Computes names + + Args: + scene: The scene containing the entity to resolve against. + + Raises: + ValueError: If provided names and IDs are inconsistent. + KeyError: If the entity name is not found in the scene. + """ + entity = scene[self.name] + + for config in _FIELD_CONFIGS: + self._resolve_field(entity, config) + + def _resolve_field(self, entity: ManagerEntity, config: _FieldConfig) -> None: + """Resolve a single field's names and IDs. + + Args: + entity: The entity to resolve against. + config: Field configuration specifying attribute names and methods. + """ + names = getattr(self, config.names_attr) + ids = getattr(self, config.ids_attr) + + # Early return if nothing to resolve. + if names is None and not isinstance(ids, list): + return + + # Get entity metadata. + entity_all_names = getattr(entity, config.names_attr) + entity_count = getattr(entity, config.num_attr) + find_method = getattr(entity, config.find_method) + + # Normalize single values to lists for uniform processing. + names = self._normalize_to_list(names) + if isinstance(ids, (int, list)): + ids = self._normalize_to_list(ids) + setattr(self, config.ids_attr, ids) + + # Handle three resolution cases. + if names is not None and isinstance(ids, list): + setattr(self, config.names_attr, names) + self._validate_consistency(names, ids, entity_all_names, find_method, config.kind_label) + elif names is not None: + self._resolve_names_to_ids( + names, + entity_all_names, + entity_count, + find_method, + config.names_attr, + config.ids_attr, + ) + elif isinstance(ids, list): + self._resolve_ids_to_names(ids, entity_all_names, config.names_attr) + + def _normalize_to_list(self, value: str | int | tuple | list | None) -> list | None: + """Convert single values to lists for uniform processing.""" + if value is None: + return None + if isinstance(value, (str, int)): + return [value] + if isinstance(value, list): + return value + return list(value) + + def _validate_consistency( + self, + names: list[str], + ids: list[int], + entity_all_names: list[str], + find_method, + kind_label: str, + ) -> None: + """Validate that provided names and IDs are consistent. + + Raises: + ValueError: If names and IDs don't match. + """ + found_ids, _ = find_method(names, preserve_order=self.preserve_order) + computed_names = [entity_all_names[i] for i in ids] + + if found_ids != ids or computed_names != names: + raise ValueError( + f"Inconsistent {kind_label} names and indices. " + f"Names {names} resolved to indices {found_ids}, " + f"but indices {ids} (mapping to names {computed_names}) were provided." + ) + + def _resolve_names_to_ids( + self, + names: list[str], + entity_all_names: list[str], + entity_count: int, + find_method, + names_attr: str, + ids_attr: str, + ) -> None: + """Resolve names to IDs, optimizing to slice(None) when all are selected.""" + found_ids, found_names = find_method(names, preserve_order=self.preserve_order) + + # Keep names and IDs in the same order. + setattr(self, names_attr, found_names) + + # Optimize to slice(None) if all components are selected in order. + if len(found_ids) == entity_count and found_names == list(entity_all_names): + setattr(self, ids_attr, slice(None)) + else: + setattr(self, ids_attr, found_ids) + + def _resolve_ids_to_names( + self, ids: list[int], entity_all_names: list[str], names_attr: str + ) -> None: + """Resolve IDs to their corresponding names.""" + resolved_names = [entity_all_names[i] for i in ids] + setattr(self, names_attr, resolved_names) diff --git a/src/unilab/managers/termination_manager.py b/src/unilab/managers/termination_manager.py new file mode 100644 index 000000000..999cac9bb --- /dev/null +++ b/src/unilab/managers/termination_manager.py @@ -0,0 +1,139 @@ +# Derived from mujocolab/mjlab v1.6.0 (0fb8a681), src/mjlab/managers/termination_manager.py. +# Copyright 2025, The mjlab Developers. +# Modified by UniLab for NumPy and UniLab contracts; licensed under Apache-2.0. +"""Termination manager for computing done signals.""" + +from __future__ import annotations + +from copy import deepcopy +from dataclasses import dataclass +from typing import TYPE_CHECKING, Sequence + +import numpy as np +from prettytable import PrettyTable + +from unilab.managers.manager_base import ManagerBase, ManagerTermBaseCfg + +if TYPE_CHECKING: + from unilab.managers._types import ManagerBasedRlEnv + + +@dataclass +class TerminationTermCfg(ManagerTermBaseCfg): + """Configuration for a termination term.""" + + time_out: bool = False + """Whether the term contributes towards episodic timeouts.""" + + +class TerminationManager(ManagerBase): + """Manages termination conditions for the environment. + + The termination manager aggregates multiple termination terms to compute + episode done signals. Terms can be either truncations (time-based) or + terminations (failure conditions). + """ + + _env: ManagerBasedRlEnv + + def __init__(self, cfg: dict[str, TerminationTermCfg | None], env: ManagerBasedRlEnv): + self._term_names: list[str] = list() + self._term_cfgs: list[TerminationTermCfg] = list() + self._class_term_cfgs: list[TerminationTermCfg] = list() + + self.cfg = deepcopy(cfg) + super().__init__(env) + + self._term_dones = dict() + for term_name in self._term_names: + self._term_dones[term_name] = np.zeros(self.num_envs, dtype=np.bool_) + self._truncated_buf = np.zeros(self.num_envs, dtype=np.bool_) + self._terminated_buf = np.zeros_like(self._truncated_buf) + + def __str__(self) -> str: + msg = f" contains {len(self._term_names)} active terms.\n" + table = PrettyTable() + table.title = "Active Termination Terms" + table.field_names = ["Index", "Name", "Time Out"] + table.align["Name"] = "l" + for index, (name, term_cfg) in enumerate( + zip(self._term_names, self._term_cfgs, strict=False) + ): + table.add_row([index, name, term_cfg.time_out]) + msg += str(table.get_string()) + msg += "\n" + return msg + + # Properties. + + @property + def active_terms(self) -> list[str]: + return self._term_names + + @property + def dones(self) -> np.ndarray: + return self._truncated_buf | self._terminated_buf + + @property + def time_outs(self) -> np.ndarray: + return self._truncated_buf + + @property + def terminated(self) -> np.ndarray: + return self._terminated_buf + + # Methods. + + def reset(self, env_ids: np.ndarray | slice | None = None) -> dict[str, int]: + if env_ids is None: + env_ids = slice(None) + extras = {} + for key in self._term_dones.keys(): + extras["Episode_Termination/" + key] = int( + np.count_nonzero(self._term_dones[key][env_ids]) + ) + for term_cfg in self._class_term_cfgs: + term_cfg.func.reset(env_ids=env_ids) + return extras + + def compute(self) -> np.ndarray: + self._truncated_buf[:] = False + self._terminated_buf[:] = False + for name, term_cfg in zip(self._term_names, self._term_cfgs, strict=False): + value = term_cfg.func(self._env, **term_cfg.params) + self._check_term_shape(name, value) + if value.dtype != np.bool_: + raise TypeError( + f"TerminationManager term '{name}' returned dtype {value.dtype}, expected bool." + ) + if term_cfg.time_out: + self._truncated_buf |= value + else: + self._terminated_buf |= value + self._term_dones[name][:] = value + return self._truncated_buf | self._terminated_buf + + def get_term(self, name: str) -> np.ndarray: + return self._term_dones[name] + + def get_term_cfg(self, term_name: str) -> TerminationTermCfg: + if term_name not in self._term_names: + raise ValueError(f"Term '{term_name}' not found in active terms.") + return self._term_cfgs[self._term_names.index(term_name)] + + def get_active_iterable_terms(self, env_idx: int) -> Sequence[tuple[str, Sequence[float]]]: + terms = [] + for key in self._term_dones.keys(): + terms.append((key, [float(self._term_dones[key][env_idx])])) + return terms + + def _prepare_terms(self) -> None: + for term_name, term_cfg in self.cfg.items(): + if term_cfg is None: + print(f"term: {term_name} set to None, skipping...") + continue + self._resolve_common_term_cfg(term_name, term_cfg) + self._term_names.append(term_name) + self._term_cfgs.append(term_cfg) + if hasattr(term_cfg.func, "reset") and callable(term_cfg.func.reset): + self._class_term_cfgs.append(term_cfg) diff --git a/tests/managers/__init__.py b/tests/managers/__init__.py new file mode 100644 index 000000000..3eb352d0c --- /dev/null +++ b/tests/managers/__init__.py @@ -0,0 +1 @@ +"""Tests for the standalone NumPy manager package.""" diff --git a/tests/managers/conftest.py b/tests/managers/conftest.py new file mode 100644 index 000000000..241c2b32c --- /dev/null +++ b/tests/managers/conftest.py @@ -0,0 +1,58 @@ +# Derived from mujocolab/mjlab v1.6.0 (0fb8a681), manager test fixtures. +# Modified by UniLab for NumPy and the standalone manager contract; Apache-2.0. + +from __future__ import annotations + +import re +from types import SimpleNamespace + +import numpy as np +import pytest + + +class FakeEntity: + def __init__(self) -> None: + self.joint_names = ["hip", "knee", "ankle"] + self.body_names = ["base", "foot"] + self.num_joints = len(self.joint_names) + self.num_bodies = len(self.body_names) + + @staticmethod + def _find( + all_names: list[str], patterns: list[str], preserve_order: bool + ) -> tuple[list[int], list[str]]: + if preserve_order: + found = [ + name for pattern in patterns for name in all_names if re.fullmatch(pattern, name) + ] + else: + found = [name for name in all_names if any(re.fullmatch(p, name) for p in patterns)] + return [all_names.index(name) for name in found], found + + def find_joints( + self, patterns: list[str], *, preserve_order: bool = False + ) -> tuple[list[int], list[str]]: + return self._find(self.joint_names, patterns, preserve_order) + + def find_bodies( + self, patterns: list[str], *, preserve_order: bool = False + ) -> tuple[list[int], list[str]]: + return self._find(self.body_names, patterns, preserve_order) + + +class FakeEnv: + def __init__(self, seed: int = 7, num_envs: int = 4) -> None: + self.num_envs = num_envs + self.rng = np.random.default_rng(seed) + self.scene = {"robot": FakeEntity()} + self.max_episode_length_s = 2.0 + self.value = np.arange(num_envs, dtype=np.float32) + self.obs = np.arange(num_envs * 2, dtype=np.float32).reshape(num_envs, 2) + self.calls: list[tuple[str, np.ndarray | None]] = [] + self.obs_buf: dict[str, np.ndarray] = {} + self.reset_buf = np.zeros(num_envs, dtype=np.bool_) + + +@pytest.fixture +def fake_env() -> FakeEnv: + return FakeEnv() diff --git a/tests/managers/test_core_managers.py b/tests/managers/test_core_managers.py new file mode 100644 index 000000000..18aa53f64 --- /dev/null +++ b/tests/managers/test_core_managers.py @@ -0,0 +1,195 @@ +# Derived from mujocolab/mjlab v1.6.0 (0fb8a681), manager tests. +# Modified by UniLab for NumPy and fail-closed term validation; Apache-2.0. + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np +import pytest + +from unilab.managers import ( + ActionManager, + ActionTerm, + ActionTermCfg, + CurriculumManager, + CurriculumTermCfg, + NullCurriculumManager, + RewardManager, + RewardTermCfg, + SceneEntityCfg, + TerminationManager, + TerminationTermCfg, +) + +from .conftest import FakeEnv + + +class DummyAction(ActionTerm): + def __init__(self, cfg: DummyActionCfg, env: FakeEnv): + super().__init__(cfg, env) + self._raw = np.zeros((env.num_envs, cfg.dim), dtype=np.float32) + self.applied = 0 + self.reset_ids: np.ndarray | slice | None = None + + @property + def action_dim(self) -> int: + return self._raw.shape[1] + + @property + def raw_action(self) -> np.ndarray: + return self._raw + + def process_actions(self, actions: np.ndarray) -> None: + self._raw[:] = actions + + def apply_actions(self) -> None: + self.applied += 1 + + def reset(self, env_ids: np.ndarray | slice | None) -> None: + self.reset_ids = env_ids + self._raw[env_ids] = 0.0 + + +@dataclass(kw_only=True) +class DummyActionCfg(ActionTermCfg): + dim: int + + def build(self, env: FakeEnv) -> DummyAction: + return DummyAction(self, env) + + +def test_action_split_history_apply_and_partial_reset(fake_env: FakeEnv) -> None: + manager = ActionManager( + { + "legs": DummyActionCfg(entity_name="robot", dim=2), + "disabled": None, + "arm": DummyActionCfg(entity_name="robot", dim=1), + }, + fake_env, + ) + assert manager.active_terms == ["legs", "arm"] + first = np.arange(12, dtype=np.float32).reshape(4, 3) + second = first + 20 + manager.process_action(first) + manager.process_action(second) + np.testing.assert_array_equal(manager.prev_action, first) + np.testing.assert_array_equal(manager.action, second) + np.testing.assert_array_equal(manager.get_term("legs").raw_action, second[:, :2]) + manager.apply_action() + assert manager.get_term("legs").applied == 1 + manager.reset(np.array([1, 3])) + np.testing.assert_array_equal(manager.action[[1, 3]], 0.0) + np.testing.assert_array_equal(manager.action[[0, 2]], second[[0, 2]]) + + +@pytest.mark.parametrize( + "action,match", + [ + (np.zeros((4, 2), dtype=np.float32), "Invalid action shape"), + (np.full((4, 3), np.nan, dtype=np.float32), "NaN or Inf"), + ], +) +def test_action_rejects_invalid_input(fake_env: FakeEnv, action: np.ndarray, match: str) -> None: + manager = ActionManager({"a": DummyActionCfg(entity_name="robot", dim=3)}, fake_env) + with pytest.raises(ValueError, match=match): + manager.process_action(action) + + +class StatefulReward: + def __init__(self, cfg: RewardTermCfg, env: FakeEnv): + self.reset_ids = None + + def __call__(self, env: FakeEnv) -> np.ndarray: + return env.value.copy() + + def reset(self, env_ids: np.ndarray | slice | None) -> None: + self.reset_ids = env_ids + + +def test_reward_dt_scaling_reset_and_config_immutability(fake_env: FakeEnv) -> None: + cfg = {"stateful": RewardTermCfg(func=StatefulReward, weight=2.0)} + manager = RewardManager(cfg, fake_env) + np.testing.assert_allclose(manager.compute(dt=0.25), fake_env.value * 0.5) + assert manager.get_active_iterable_terms(2) == [("stateful", [4.0])] + extras = manager.reset(np.array([1, 2])) + assert extras["Episode_Reward/stateful"] == pytest.approx(0.375) + assert cfg["stateful"].func is StatefulReward + assert isinstance(manager.get_term_cfg("stateful").func, StatefulReward) + + +@pytest.mark.parametrize("bad", [np.nan, np.inf, -np.inf]) +def test_reward_nonfinite_is_an_error(fake_env: FakeEnv, bad: float) -> None: + def reward(env: FakeEnv) -> np.ndarray: + value = np.ones(env.num_envs, dtype=np.float32) + value[2] = bad + return value + + manager = RewardManager({"bad_reward": RewardTermCfg(func=reward, weight=1.0)}, fake_env) + with pytest.raises(ValueError, match="RewardManager term 'bad_reward'"): + manager.compute(0.01) + + +def test_reward_and_termination_shape_validation(fake_env: FakeEnv) -> None: + reward = RewardManager( + {"bad": RewardTermCfg(func=lambda env: np.zeros((env.num_envs, 1)), weight=1.0)}, + fake_env, + ) + with pytest.raises(ValueError, match=r"expected \(4,\)"): + reward.compute(0.1) + + termination = TerminationManager( + {"bad": TerminationTermCfg(func=lambda env: np.zeros(env.num_envs))}, fake_env + ) + with pytest.raises(TypeError, match="expected bool"): + termination.compute() + + +def test_termination_splits_timeouts_and_failures(fake_env: FakeEnv) -> None: + timeout = np.array([True, False, False, True]) + failure = np.array([False, True, False, True]) + manager = TerminationManager( + { + "timeout": TerminationTermCfg(func=lambda env: timeout.copy(), time_out=True), + "failure": TerminationTermCfg(func=lambda env: failure.copy()), + }, + fake_env, + ) + np.testing.assert_array_equal(manager.compute(), timeout | failure) + np.testing.assert_array_equal(manager.time_outs, timeout) + np.testing.assert_array_equal(manager.terminated, failure) + assert manager.reset(np.array([0, 1])) == { + "Episode_Termination/timeout": 1, + "Episode_Termination/failure": 1, + } + + +def test_curriculum_and_null_semantics(fake_env: FakeEnv) -> None: + def update(env: FakeEnv, env_ids: np.ndarray | slice) -> dict[str, float]: + return {"difficulty": 3.0} + + manager = CurriculumManager({"terrain": CurriculumTermCfg(func=update)}, fake_env) + manager.compute(np.array([1, 2])) + assert manager.reset()["Curriculum/terrain/difficulty"] == 3.0 + null = NullCurriculumManager() + assert null.active_terms == [] + assert null.reset() == {} + + bad = CurriculumManager({"bad": CurriculumTermCfg(func=lambda env, env_ids: np.nan)}, fake_env) + with pytest.raises(ValueError, match="CurriculumManager term 'bad'"): + bad.compute() + + +def test_scene_entity_selector_resolution(fake_env: FakeEnv) -> None: + cfg = SceneEntityCfg(name="robot", joint_names=("ankle", "hip"), preserve_order=True) + cfg.resolve(fake_env.scene) + assert cfg.joint_names == ["ankle", "hip"] + assert cfg.joint_ids == [2, 0] + + all_joints = SceneEntityCfg(name="robot", joint_names=".*") + all_joints.resolve(fake_env.scene) + assert all_joints.joint_ids == slice(None) + + inconsistent = SceneEntityCfg(name="robot", joint_names="hip", joint_ids=[1]) + with pytest.raises(ValueError, match="Inconsistent joint"): + inconsistent.resolve(fake_env.scene) diff --git a/tests/managers/test_event_command_metrics_recorder.py b/tests/managers/test_event_command_metrics_recorder.py new file mode 100644 index 000000000..17f9e7830 --- /dev/null +++ b/tests/managers/test_event_command_metrics_recorder.py @@ -0,0 +1,263 @@ +# Derived from mujocolab/mjlab v1.6.0 (0fb8a681), manager tests. +# Modified by UniLab for NumPy scheduling and unsupported capability errors; Apache-2.0. + +from __future__ import annotations + +import ast +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +import pytest + +import unilab.managers as managers +from unilab.managers import ( + CommandManager, + CommandTerm, + CommandTermCfg, + EventManager, + EventTermCfg, + MetricsManager, + MetricsTermCfg, + NullCommandManager, + NullMetricsManager, + NullRecorderManager, + RecorderManager, + RecorderTerm, + RecorderTermCfg, +) + +from .conftest import FakeEnv + + +def _record(env: FakeEnv, env_ids: np.ndarray | None, *, label: str) -> None: + copied = None if env_ids is None else np.asarray(env_ids).copy() + env.calls.append((label, copied)) + + +def test_event_modes_interval_reset_throttle_and_order(fake_env: FakeEnv) -> None: + cfg = { + "startup": EventTermCfg(func=_record, params={"label": "startup"}, mode="startup"), + "reset": EventTermCfg( + func=_record, + params={"label": "reset"}, + mode="reset", + min_step_count_between_reset=3, + ), + "interval": EventTermCfg( + func=_record, + params={"label": "interval"}, + mode="interval", + interval_range_s=(0.1, 0.1), + ), + "step": EventTermCfg(func=_record, params={"label": "step"}, mode="step"), + } + manager = EventManager(cfg, fake_env) + assert list(manager.active_terms) == ["startup", "reset", "interval", "step"] + manager.apply("startup", env_ids=np.array([0, 2])) + manager.apply("step", dt=0.01) + manager.apply("interval", dt=0.1) + manager.apply("reset", env_ids=np.array([1, 3]), global_env_step_count=1) + manager.apply("reset", env_ids=np.array([1, 3]), global_env_step_count=2) + assert [label for label, _ in fake_env.calls] == ["startup", "step", "interval", "reset"] + np.testing.assert_array_equal(fake_env.calls[2][1], np.arange(fake_env.num_envs)) + + +def test_event_validation_and_model_mutation_failure(fake_env: FakeEnv) -> None: + with pytest.raises(ValueError, match="interval_range_s"): + EventManager({"bad": EventTermCfg(func=_record, mode="interval")}, fake_env) + + def model_mutation(env: FakeEnv, env_ids: np.ndarray | None) -> None: + pass + + model_mutation.model_fields = ("body_mass",) # type: ignore[attr-defined] + with pytest.raises(NotImplementedError, match="model-field mutation"): + EventManager({"unsupported": EventTermCfg(func=model_mutation, mode="startup")}, fake_env) + + empty = EventManager({}, fake_env) + empty.apply("interval") + + +def test_event_interval_rng_is_reproducible() -> None: + cfg = { + "interval": EventTermCfg( + func=_record, + params={"label": "interval"}, + mode="interval", + interval_range_s=(0.1, 2.0), + ) + } + left = EventManager(cfg, FakeEnv(seed=17)) + right = EventManager(cfg, FakeEnv(seed=17)) + np.testing.assert_array_equal(left._interval_term_time_left, right._interval_term_time_left) + + +class DummyCommand(CommandTerm): + def __init__(self, cfg: DummyCommandCfg, env: FakeEnv): + super().__init__(cfg, env) + self._command = np.zeros((env.num_envs, 1), dtype=np.float32) + self.metrics["error"] = np.arange(env.num_envs, dtype=np.float32) + + @property + def command(self) -> np.ndarray: + return self._command + + def _update_metrics(self) -> None: + self.metrics["error"] += 1.0 + + def _resample_command(self, env_ids: np.ndarray) -> None: + self._command[env_ids, 0] = self.command_counter[env_ids] + + def _update_command(self, env_ids: np.ndarray | None) -> None: + pass + + +@dataclass(kw_only=True) +class DummyCommandCfg(CommandTermCfg): + def build(self, env: FakeEnv) -> DummyCommand: + return DummyCommand(self, env) + + +def test_command_resample_metrics_validation_and_null(fake_env: FakeEnv) -> None: + manager = CommandManager({"goal": DummyCommandCfg(resampling_time_range=(0.5, 0.5))}, fake_env) + extras = manager.reset(np.array([1, 2])) + assert extras == {"Metrics/goal/error": 1.5} + np.testing.assert_array_equal(manager.get_command("goal")[[1, 2]], 0.0) + manager.compute(0.5) + np.testing.assert_array_equal(manager.get_term("goal").command[:, 0], [0, 1, 1, 0]) + assert manager.get_term("goal").command_counter.tolist() == [1, 2, 2, 1] + + manager.get_term("goal").command[0, 0] = np.nan + with pytest.raises(ValueError, match="NaN or Inf"): + manager.get_command("goal") + + null = NullCommandManager() + assert null.get_command("missing") is None + assert null.reset() == {} + + +def test_command_viewer_request_and_old_signature_fail_closed(fake_env: FakeEnv) -> None: + with pytest.raises(NotImplementedError, match="viewer"): + CommandManager( + {"goal": DummyCommandCfg(resampling_time_range=(1.0, 1.0), debug_vis=True)}, + fake_env, + ) + + class OldCommand(DummyCommand): + def _update_command(self) -> None: # type: ignore[override] + pass + + @dataclass(kw_only=True) + class OldCfg(CommandTermCfg): + def build(self, env: FakeEnv) -> OldCommand: + return OldCommand(self, env) + + with pytest.raises(TypeError, match="must accept env_ids"): + CommandManager({"old": OldCfg(resampling_time_range=(1.0, 1.0))}, fake_env) + + +def test_metrics_reductions_substeps_reset_and_finite_failure(fake_env: FakeEnv) -> None: + manager = MetricsManager( + { + "mean": MetricsTermCfg(func=lambda env: env.value.copy(), reduce="mean"), + "max": MetricsTermCfg(func=lambda env: env.value.copy(), reduce="max"), + "sum": MetricsTermCfg(func=lambda env: env.value.copy(), reduce="sum"), + "last": MetricsTermCfg(func=lambda env: env.value.copy(), reduce="last"), + "substep": MetricsTermCfg( + func=lambda env: env.value.copy(), per_substep=True, reduce="mean" + ), + }, + fake_env, + ) + manager.compute_substep() + fake_env.value += 2 + manager.compute_substep() + manager.compute() + extras = manager.reset(np.array([1, 2])) + assert extras["Episode_Metrics/mean"] == pytest.approx(3.5) + assert extras["Episode_Metrics/max"] == pytest.approx(3.5) + assert extras["Episode_Metrics/sum"] == pytest.approx(3.5) + assert extras["Episode_Metrics/last"] == pytest.approx(3.5) + assert extras["Episode_Metrics/substep"] == pytest.approx(2.5) + + bad = MetricsManager( + {"bad": MetricsTermCfg(func=lambda env: np.full(env.num_envs, np.inf))}, fake_env + ) + with pytest.raises(ValueError, match="MetricsManager term 'bad'"): + bad.compute() + assert NullMetricsManager().reset() == {} + + +class TraceRecorder(RecorderTerm): + def __init__(self, cfg: RecorderTermCfg, env: FakeEnv): + super().__init__(cfg, env) + self.events: list[tuple[str, list[int] | None]] = [] + + def record_pre_reset(self, env_ids: np.ndarray) -> None: + self.events.append(("pre", env_ids.tolist())) + + def record_post_reset(self, env_ids: np.ndarray) -> None: + self.events.append(("post_reset", env_ids.tolist())) + + def record_post_step(self) -> None: + self.events.append(("step", None)) + + def close(self) -> None: + self.events.append(("close", None)) + + +def test_recorder_lifecycle_and_null(fake_env: FakeEnv) -> None: + cfg = {"trace": RecorderTermCfg(func=TraceRecorder)} + manager = RecorderManager(cfg, fake_env) + ids = np.array([0, 3]) + manager.record_pre_reset(ids) + manager.record_post_reset(ids) + manager.record_post_step() + manager.close() + assert manager.get_term("trace").events == [ + ("pre", [0, 3]), + ("post_reset", [0, 3]), + ("step", None), + ("close", None), + ] + assert cfg["trace"].func is TraceRecorder + null = NullRecorderManager() + with pytest.raises(KeyError, match="has no terms"): + null.get_term("trace") + + +def test_public_exports_and_repository_import_boundary() -> None: + expected = { + "ManagerBase", + "ManagerTermBase", + "ManagerTermBaseCfg", + "ActionManager", + "ObservationManager", + "RewardManager", + "TerminationManager", + "EventManager", + "CommandManager", + "CurriculumManager", + "MetricsManager", + "RecorderManager", + "SceneEntityCfg", + } + assert expected <= set(vars(managers)) + + package_root = Path(managers.__file__).parent + forbidden_roots = {"torch", "mjlab"} + forbidden_unilab = {"unilab.ipc", "unilab.runners", "unilab.scripts", "unilab.base.backend"} + for path in package_root.rglob("*.py"): + tree = ast.parse(path.read_text()) + imports = [] + for node in ast.walk(tree): + if isinstance(node, ast.Import): + imports.extend(alias.name for alias in node.names) + elif isinstance(node, ast.ImportFrom) and node.module: + imports.append(node.module) + assert not ({name.split(".")[0] for name in imports} & forbidden_roots), path + assert not any( + name == prefix or name.startswith(prefix + ".") + for name in imports + for prefix in forbidden_unilab + ), path diff --git a/tests/managers/test_observation_buffers_noise.py b/tests/managers/test_observation_buffers_noise.py new file mode 100644 index 000000000..2b7e10974 --- /dev/null +++ b/tests/managers/test_observation_buffers_noise.py @@ -0,0 +1,205 @@ +# Derived from mujocolab/mjlab v1.6.0 (0fb8a681), observation/buffer/noise tests. +# Modified by UniLab for NumPy and env-owned RNG; Apache-2.0. + +from __future__ import annotations + +import numpy as np +import pytest + +from unilab.managers import ObservationGroupCfg, ObservationManager, ObservationTermCfg +from unilab.managers._buffers import CircularBuffer, DelayBuffer +from unilab.managers._noise import ( + ConstantNoiseCfg, + GaussianNoiseCfg, + NoiseModelWithAdditiveBiasCfg, + UniformNoiseCfg, +) + +from .conftest import FakeEnv + + +def test_circular_buffer_history_backfill_lag_and_partial_reset() -> None: + buffer = CircularBuffer(max_len=3, batch_size=2) + first = np.array([[1.0], [10.0]], dtype=np.float32) + buffer.append(first) + np.testing.assert_array_equal(buffer.buffer[:, :, 0], [[1, 1, 1], [10, 10, 10]]) + buffer.append(np.array([[2.0], [20.0]], dtype=np.float32)) + buffer.append(np.array([[3.0], [30.0]], dtype=np.float32)) + np.testing.assert_array_equal(buffer[np.array([0, 2])][:, 0], [3, 10]) + + buffer.reset([1]) + buffer.append(np.array([[4.0], [99.0]], dtype=np.float32)) + np.testing.assert_array_equal(buffer.buffer[0, :, 0], [2, 3, 4]) + np.testing.assert_array_equal(buffer.buffer[1, :, 0], [99, 99, 99]) + + +def test_circular_buffer_rejects_invalid_usage() -> None: + with pytest.raises(ValueError, match=">= 1"): + CircularBuffer(max_len=0, batch_size=2) + buffer = CircularBuffer(max_len=2, batch_size=2) + with pytest.raises(RuntimeError, match="not initialized"): + _ = buffer.buffer + with pytest.raises(ValueError, match="batch size"): + buffer.append(np.zeros((3, 1))) + + +def test_delay_buffer_constant_delay_and_partial_backfill() -> None: + buffer = DelayBuffer(min_lag=2, max_lag=2, batch_size=2) + outputs = [] + for value in (1.0, 2.0, 3.0, 4.0): + buffer.append(np.full((2, 1), value, dtype=np.float32)) + outputs.append(buffer.compute().copy()) + np.testing.assert_array_equal(np.stack(outputs)[:, 0, 0], [1, 1, 1, 2]) + + buffer.reset(np.array([1])) + buffer.backfill(np.array([[8.0], [9.0]], dtype=np.float32), np.array([1])) + np.testing.assert_array_equal(buffer.peek()[1], [9.0]) + + +def test_delay_rng_is_reproducible_and_required() -> None: + def draw(seed: int) -> list[np.ndarray]: + buffer = DelayBuffer( + min_lag=0, + max_lag=3, + batch_size=8, + generator=np.random.default_rng(seed), + ) + values = [] + for step in range(5): + buffer.append(np.full((8, 1), step, dtype=np.float32)) + buffer.compute() + values.append(buffer.current_lags.copy()) + return values + + for left, right in zip(draw(123), draw(123), strict=True): + np.testing.assert_array_equal(left, right) + + missing_rng = DelayBuffer(min_lag=0, max_lag=2, batch_size=2) + missing_rng.append(np.zeros((2, 1))) + with pytest.raises(ValueError, match="env-owned"): + missing_rng.compute() + + +def test_noise_configs_use_supplied_generator() -> None: + data = np.ones((4, 3), dtype=np.float32) + uniform = UniformNoiseCfg(n_min=-0.2, n_max=0.2) + first = uniform.apply(data, rng=np.random.default_rng(9)) + second = uniform.apply(data, rng=np.random.default_rng(9)) + np.testing.assert_array_equal(first, second) + assert first.dtype == np.float32 + with pytest.raises(ValueError, match="env-owned"): + uniform.apply(data) + + gaussian = GaussianNoiseCfg(mean=0.0, std=0.1) + assert gaussian.apply(data, rng=np.random.default_rng(2)).shape == data.shape + np.testing.assert_array_equal(ConstantNoiseCfg(bias=2.0, operation="abs").apply(data), 2.0) + + +def test_additive_bias_noise_supports_scalar_terms() -> None: + from unilab.managers._noise import NoiseModelWithAdditiveBias + + cfg = NoiseModelWithAdditiveBiasCfg( + noise_cfg=ConstantNoiseCfg(bias=0.0), + bias_noise_cfg=ConstantNoiseCfg(bias=0.5), + ) + model = NoiseModelWithAdditiveBias(cfg, num_envs=4, rng=np.random.default_rng(2)) + result = model(np.ones(4, dtype=np.float32)) + np.testing.assert_array_equal(result, 1.5) + assert result.shape == (4,) + + +def test_observation_groups_pipeline_order_and_history(fake_env: FakeEnv) -> None: + cfg = { + "policy": ObservationGroupCfg( + terms={ + "state": ObservationTermCfg( + func=lambda env: env.obs, + clip=(-1.0, 4.0), + scale=2.0, + history_length=2, + ), + "bias": ObservationTermCfg( + func=lambda env: np.ones((env.num_envs, 1), dtype=np.float32) + ), + } + ), + "dict_group": ObservationGroupCfg( + terms={"state": ObservationTermCfg(func=lambda env: env.obs)}, + concatenate_terms=False, + ), + } + manager = ObservationManager(cfg, fake_env) + first = manager.compute(update_history=True) + expected_first = np.clip(fake_env.obs, -1, 4) * 2 + np.testing.assert_array_equal(first["policy"][:, :4], np.tile(expected_first, (1, 2))) + assert list(first["dict_group"]) == ["state"] + assert manager.group_obs_dim["policy"] == (5,) + + fake_env.obs = fake_env.obs + 10 + second = manager.compute(update_history=True)["policy"] + expected_second = np.clip(fake_env.obs, -1, 4) * 2 + np.testing.assert_array_equal(second[:, :2], expected_first) + np.testing.assert_array_equal(second[:, 2:4], expected_second) + assert manager.get_active_iterable_terms(0)[0][0] == "policy-state" + + +def test_observation_noise_model_delay_and_seed_reproducibility() -> None: + cfg = { + "policy": ObservationGroupCfg( + terms={ + "state": ObservationTermCfg( + func=lambda env: env.obs, + noise=NoiseModelWithAdditiveBiasCfg( + noise_cfg=GaussianNoiseCfg(std=0.1), + bias_noise_cfg=UniformNoiseCfg(n_min=-0.2, n_max=0.2), + ), + delay_min_lag=1, + delay_max_lag=1, + ) + }, + enable_corruption=True, + ) + } + left = ObservationManager(cfg, FakeEnv(seed=11)).compute(update_history=True) + right = ObservationManager(cfg, FakeEnv(seed=11)).compute(update_history=True) + np.testing.assert_array_equal(left["policy"], right["policy"]) + + +@pytest.mark.parametrize("bad", [np.nan, np.inf]) +def test_observation_default_finite_policy_fails_closed(fake_env: FakeEnv, bad: float) -> None: + def invalid(env: FakeEnv) -> np.ndarray: + result = env.obs.copy() + result[1, 0] = bad + return result + + manager = ObservationManager( + {"policy": ObservationGroupCfg(terms={"bad": ObservationTermCfg(func=invalid)})}, + fake_env, + ) + with pytest.raises(ValueError, match="ObservationManager term 'policy/bad'"): + manager.compute() + + +def test_observation_explicit_sanitize_and_shape_error(fake_env: FakeEnv) -> None: + sanitize = ObservationManager( + { + "policy": ObservationGroupCfg( + terms={ + "bad": ObservationTermCfg(func=lambda env: np.full((env.num_envs, 1), np.nan)) + }, + nan_policy="sanitize", + ) + }, + fake_env, + ) + np.testing.assert_array_equal(sanitize.compute()["policy"], 0.0) + + with pytest.raises(ValueError, match="num_envs"): + ObservationManager( + { + "policy": ObservationGroupCfg( + terms={"bad": ObservationTermCfg(func=lambda env: np.zeros((2, 1)))} + ) + }, + fake_env, + ) diff --git a/uv.lock b/uv.lock index 7c4d63516..9be5d1e83 100644 --- a/uv.lock +++ b/uv.lock @@ -2733,6 +2733,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, ] +[[package]] +name = "prettytable" +version = "3.18.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "wcwidth" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/81/74/ba08d81e668ccfe8658d7520a307e63c19862c08eb4ccb26f356c5239a7a/prettytable-3.18.0.tar.gz", hash = "sha256:439217116152244369caf3d9f1caf2f9fe29b03bd79e88d2928c8e718c95d680", size = 76373, upload-time = "2026-06-22T16:07:50.174Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fe/be/2e6798ace5cc036f5d05d36b7b2fd85346f1a708c87060890b070d0ec607/prettytable-3.18.0-py3-none-any.whl", hash = "sha256:b3346e0e6f79180833aebaac088ae926340586cf6d7d991b9eb125b65f72313a", size = 37357, upload-time = "2026-06-22T16:07:48.595Z" }, +] + [[package]] name = "prompt-toolkit" version = "3.0.52" @@ -3921,6 +3933,7 @@ dependencies = [ { name = "onnxruntime", version = "1.19.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "onnxruntime", version = "1.24.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "packaging" }, + { name = "prettytable" }, { name = "rich" }, { name = "rsl-rl-lib" }, { name = "setuptools" }, @@ -3981,6 +3994,7 @@ requires-dist = [ { name = "onnxruntime", marker = "python_full_version < '3.11'", specifier = "<1.20" }, { name = "onnxruntime", marker = "python_full_version >= '3.11'", specifier = ">=1.20" }, { name = "packaging" }, + { name = "prettytable", specifier = ">=3.10" }, { name = "pybind11", marker = "extra == 'mujoco'", specifier = ">=2.12" }, { name = "rich" }, { name = "rsl-rl-lib", specifier = ">=5.0.0" }, From b7673d78c311988d0afceb6933dc433f082e5ddf Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Mon, 17 Aug 2026 19:34:00 +0800 Subject: [PATCH 003/169] feat(base): add numpy entity facade (#1048) --- src/unilab/base/entity.py | 635 +++++++++++++++++++++ src/unilab/base/scene.py | 3 + src/unilab/managers/_types.py | 121 +++- src/unilab/managers/manager_base.py | 11 +- src/unilab/managers/scene_entity_config.py | 30 +- tests/base/test_entity_facade.py | 351 ++++++++++++ 6 files changed, 1134 insertions(+), 17 deletions(-) create mode 100644 src/unilab/base/entity.py create mode 100644 tests/base/test_entity_facade.py diff --git a/src/unilab/base/entity.py b/src/unilab/base/entity.py new file mode 100644 index 000000000..e9c3e0842 --- /dev/null +++ b/src/unilab/base/entity.py @@ -0,0 +1,635 @@ +"""Base-owned NumPy scene/entity facade for manager terms. + +The facade deliberately describes partitions of an already materialized UniLab scene. +It is not a second scene composer: all name resolution and state reads go through the +public :class:`~unilab.base.backend.base.SimBackend` contract. +""" + +from __future__ import annotations + +import re +from collections.abc import Iterator, Mapping, Sequence +from dataclasses import dataclass +from types import MappingProxyType +from typing import TYPE_CHECKING, NoReturn + +import numpy as np + +from unilab.base.backend.base import SimBackend + +if TYPE_CHECKING: + from unilab.base.scene import SceneCfg + + +NamesCfg = tuple[str, ...] | list[str] | None + + +@dataclass(frozen=True) +class EntityCfg: + """Declare one logical entity inside an existing backend scene. + + Names are explicit because UniLab keeps scene composition in task-owned XML and + backend adapters. ``None`` means that the namespace is not exposed by this + entity; an empty sequence means that it is exposed but contains no elements. + """ + + root_body_name: str | None = None + joint_names: NamesCfg = None + body_names: NamesCfg = None + geom_names: NamesCfg = None + site_names: NamesCfg = None + actuator_names: NamesCfg = None + + +def _normalize_names(entity_name: str, kind: str, names: NamesCfg) -> tuple[str, ...] | None: + if names is None: + return None + if isinstance(names, str): + raise TypeError( + f"Entity '{entity_name}' {kind} names must be a sequence of strings, not a scalar" + ) + invalid = [value for value in names if not isinstance(value, str)] + if invalid: + raise TypeError(f"Entity '{entity_name}' {kind} names must be strings; got {invalid}") + normalized = tuple(names) + if any(not name for name in normalized): + raise ValueError(f"Entity '{entity_name}' {kind} names must be non-empty strings") + if len(set(normalized)) != len(normalized): + raise ValueError(f"Entity '{entity_name}' {kind} names must be unique: {normalized}") + return normalized + + +def _readonly_ids(values: np.ndarray | Sequence[int], *, expected: int, label: str) -> np.ndarray: + raw_ids = np.asarray(values) + if not np.issubdtype(raw_ids.dtype, np.integer) or np.issubdtype(raw_ids.dtype, np.bool_): + raise TypeError(f"{label} resolver must return integer IDs, got dtype {raw_ids.dtype}") + ids = np.asarray(raw_ids, dtype=np.int32) + if ids.shape != (expected,): + raise ValueError(f"{label} resolver returned shape {ids.shape}, expected ({expected},)") + if np.any(ids < 0): + raise ValueError(f"{label} resolver returned negative IDs: {ids.tolist()}") + if np.unique(ids).size != ids.size: + raise ValueError(f"{label} resolver returned duplicate IDs: {ids.tolist()}") + ids = np.array(ids, copy=True, dtype=np.int32) + ids.setflags(write=False) + return ids + + +def _as_column_index(ids: np.ndarray) -> slice | np.ndarray: + """Use a slice for contiguous columns and advanced indexing otherwise.""" + if ids.size: + start = int(ids[0]) + if np.array_equal(ids, np.arange(start, start + ids.size, dtype=ids.dtype)): + return slice(start, start + ids.size) + index = np.asarray(ids, dtype=np.intp).copy() + index.setflags(write=False) + return index + + +# Matching semantics derived from mujocolab/mjlab v1.6.0 (0fb8a681), +# src/mjlab/utils/lab_api/string.py. Copyright 2025, The mjlab Developers; +# adapted for the UniLab NumPy facade under Apache-2.0. +def _resolve_matching_names( + keys: str | Sequence[str], names: Sequence[str], preserve_order: bool +) -> tuple[list[int], list[str]]: + """Pinned mjlab-compatible full-regex matching over cached entity names.""" + patterns = (keys,) if isinstance(keys, str) else tuple(keys) + matches: list[tuple[int, int, str]] = [] + matched_by: list[str | None] = [None] * len(names) + per_pattern: list[list[str]] = [[] for _ in patterns] + + for name_index, candidate in enumerate(names): + for pattern_index, pattern in enumerate(patterns): + try: + matched = re.fullmatch(pattern, candidate) is not None + except re.error as exc: + raise ValueError(f"Invalid entity selector regex {pattern!r}: {exc}") from exc + if not matched: + continue + if matched_by[name_index] is not None: + raise ValueError( + f"Multiple matches for '{candidate}': " + f"'{matched_by[name_index]}' and '{pattern}'!" + ) + matched_by[name_index] = pattern + matches.append((pattern_index, name_index, candidate)) + per_pattern[pattern_index].append(candidate) + + if any(not values for values in per_pattern): + rendered = ", ".join( + f"{pattern!r}: {values}" for pattern, values in zip(patterns, per_pattern) + ) + raise ValueError( + "Not all entity selector regular expressions matched; " + f"matches=({rendered}), available={list(names)}" + ) + + if preserve_order: + matches.sort(key=lambda item: item[0]) + return [item[1] for item in matches], [item[2] for item in matches] + + +class EntityData: + """Hot-path NumPy state surface backed by cached backend IDs.""" + + def __init__( + self, + backend: SimBackend, + *, + root_body_ids: np.ndarray | None, + joint_pos_ids: np.ndarray | None, + joint_vel_ids: np.ndarray | None, + body_ids: np.ndarray | None, + actuator_ctrl_range: np.ndarray | None, + entity_name: str, + backend_type: str, + ) -> None: + self._backend = backend + self._entity_name = entity_name + self._backend_type = backend_type + self._root_body_ids = root_body_ids + self._joint_pos_index = None if joint_pos_ids is None else _as_column_index(joint_pos_ids) + self._joint_vel_index = None if joint_vel_ids is None else _as_column_index(joint_vel_ids) + self._body_ids = body_ids + self._actuator_ctrl_range = actuator_ctrl_range + + def _require(self, value, capability: str): + if value is None: + raise NotImplementedError( + f"Entity '{self._entity_name}' data capability '{capability}' is unavailable " + f"on backend '{self._backend_type}': it was not materialized" + ) + return value + + @property + def root_link_pos_w(self) -> np.ndarray: + ids = self._require(self._root_body_ids, "root body state") + return self._backend.get_body_pos_w(ids)[:, 0] + + @property + def root_link_quat_w(self) -> np.ndarray: + ids = self._require(self._root_body_ids, "root body state") + return self._backend.get_body_quat_w(ids)[:, 0] + + @property + def root_link_lin_vel_w(self) -> np.ndarray: + ids = self._require(self._root_body_ids, "root body state") + return self._backend.get_body_lin_vel_w(ids)[:, 0] + + @property + def root_link_ang_vel_w(self) -> np.ndarray: + ids = self._require(self._root_body_ids, "root body state") + return self._backend.get_body_ang_vel_w(ids)[:, 0] + + @property + def root_link_pose_w(self) -> np.ndarray: + return np.concatenate((self.root_link_pos_w, self.root_link_quat_w), axis=-1) + + @property + def root_link_vel_w(self) -> np.ndarray: + return np.concatenate((self.root_link_lin_vel_w, self.root_link_ang_vel_w), axis=-1) + + @property + def joint_pos(self) -> np.ndarray: + index = self._require(self._joint_pos_index, "joint position") + return self._backend.get_dof_pos()[:, index] + + @property + def joint_vel(self) -> np.ndarray: + index = self._require(self._joint_vel_index, "joint velocity") + return self._backend.get_dof_vel()[:, index] + + @property + def body_link_pos_w(self) -> np.ndarray: + ids = self._require(self._body_ids, "body state") + return self._backend.get_body_pos_w(ids) + + @property + def body_link_quat_w(self) -> np.ndarray: + ids = self._require(self._body_ids, "body state") + return self._backend.get_body_quat_w(ids) + + @property + def body_link_lin_vel_w(self) -> np.ndarray: + ids = self._require(self._body_ids, "body state") + return self._backend.get_body_lin_vel_w(ids) + + @property + def body_link_ang_vel_w(self) -> np.ndarray: + ids = self._require(self._body_ids, "body state") + return self._backend.get_body_ang_vel_w(ids) + + @property + def body_link_pose_w(self) -> np.ndarray: + return np.concatenate((self.body_link_pos_w, self.body_link_quat_w), axis=-1) + + @property + def body_link_vel_w(self) -> np.ndarray: + return np.concatenate((self.body_link_lin_vel_w, self.body_link_ang_vel_w), axis=-1) + + @property + def actuator_ctrl_range(self) -> np.ndarray: + return self._require(self._actuator_ctrl_range, "actuator control range") + + +class Entity: + """Logical entity with cached local-to-backend mappings.""" + + def __init__(self, name: str, cfg: EntityCfg, backend: SimBackend) -> None: + if not name: + raise ValueError("Entity name must be a non-empty string") + self.name = name + self._backend_type = backend.backend_type + + self._joint_names = _normalize_names(name, "joint", cfg.joint_names) + self._body_names = _normalize_names(name, "body", cfg.body_names) + self._geom_names = _normalize_names(name, "geom", cfg.geom_names) + self._site_names = _normalize_names(name, "site", cfg.site_names) + self._actuator_names = _normalize_names(name, "actuator", cfg.actuator_names) + + root_body_ids = None + if cfg.root_body_name is not None: + if not isinstance(cfg.root_body_name, str) or not cfg.root_body_name: + raise TypeError(f"Entity '{self.name}' root_body_name must be a non-empty string") + root_body_ids = self._resolve_ids( + "root body", + (cfg.root_body_name,), + backend.get_body_ids, + ) + + joint_pos_ids = joint_vel_ids = None + if self._joint_names is not None: + joint_pos_ids = self._resolve_ids( + "joint position index", + self._joint_names, + backend.get_joint_dof_pos_indices, + ) + joint_vel_ids = self._resolve_ids( + "joint velocity index", + self._joint_names, + backend.get_joint_dof_vel_indices, + ) + + body_ids = None + if self._body_names is not None: + body_ids = self._resolve_ids("body", self._body_names, backend.get_body_ids) + + self._geom_ids = None + if self._geom_names is not None: + self._geom_ids = self._resolve_enumerated_ids( + "geom", self._geom_names, backend.get_geom_names + ) + + self._site_ids = None + if self._site_names is not None: + self._site_ids = self._resolve_ids("site", self._site_names, backend.get_site_ids) + + actuator_ids = None + if self._actuator_names is not None: + actuator_ids = self._resolve_enumerated_ids( + "actuator", self._actuator_names, backend.get_actuator_names + ) + + self._validate_joint_state(backend, joint_pos_ids, joint_vel_ids) + self._validate_body_state(backend, root_body_ids, body_ids) + actuator_ctrl_range = self._materialize_actuator_ctrl_range(backend, actuator_ids) + + self.data = EntityData( + backend, + root_body_ids=root_body_ids, + joint_pos_ids=joint_pos_ids, + joint_vel_ids=joint_vel_ids, + body_ids=body_ids, + actuator_ctrl_range=actuator_ctrl_range, + entity_name=self.name, + backend_type=self._backend_type, + ) + + def _capability_error(self, capability: str, detail: str) -> NotImplementedError: + return NotImplementedError( + f"Entity '{self.name}' capability '{capability}' is unavailable on " + f"backend '{self._backend_type}': {detail}" + ) + + def _resolve_ids(self, capability: str, names: tuple[str, ...], resolver) -> np.ndarray: + try: + values = resolver(names) + except NotImplementedError as exc: + raise self._capability_error(capability, str(exc)) from exc + except (KeyError, ValueError) as exc: + raise ValueError( + f"Entity '{self.name}' could not resolve {capability} names {list(names)} " + f"on backend '{self._backend_type}': {exc}" + ) from exc + return _readonly_ids( + values, + expected=len(names), + label=f"Entity '{self.name}' {capability}", + ) + + def _resolve_enumerated_ids( + self, capability: str, names: tuple[str, ...], resolver + ) -> np.ndarray: + try: + all_names = tuple(resolver()) + except NotImplementedError as exc: + raise self._capability_error(capability, str(exc)) from exc + invalid = [value for value in all_names if not isinstance(value, str)] + if invalid: + raise TypeError( + f"Entity '{self.name}' {capability} name resolver on backend " + f"'{self._backend_type}' returned non-string names: {invalid}" + ) + nonempty_names = [value for value in all_names if value] + if len(set(nonempty_names)) != len(nonempty_names): + raise ValueError( + f"Entity '{self.name}' {capability} name resolver on backend " + f"'{self._backend_type}' returned duplicate names" + ) + ids_by_name = {value: index for index, value in enumerate(all_names) if value} + missing = [value for value in names if value not in ids_by_name] + if missing: + raise ValueError( + f"Entity '{self.name}' could not resolve {capability} names {missing} on " + f"backend '{self._backend_type}'; available={list(all_names)}" + ) + return _readonly_ids( + [ids_by_name[value] for value in names], + expected=len(names), + label=f"Entity '{self.name}' {capability}", + ) + + def _read_state(self, capability: str, getter, *args) -> np.ndarray: + try: + return np.asarray(getter(*args)) + except (AttributeError, NotImplementedError) as exc: + raise self._capability_error(capability, str(exc)) from exc + + def _validate_joint_state( + self, + backend: SimBackend, + pos_ids: np.ndarray | None, + vel_ids: np.ndarray | None, + ) -> None: + for capability, getter, ids in ( + ("joint position state", backend.get_dof_pos, pos_ids), + ("joint velocity state", backend.get_dof_vel, vel_ids), + ): + if ids is None: + continue + value = self._read_state(capability, getter) + if value.ndim != 2 or value.shape[0] != backend.num_envs: + raise ValueError( + f"Entity '{self.name}' capability '{capability}' on backend " + f"'{self._backend_type}' returned shape {value.shape}; expected " + f"({backend.num_envs}, num_dof)" + ) + if ids.size and int(np.max(ids)) >= value.shape[1]: + raise ValueError( + f"Entity '{self.name}' capability '{capability}' resolved index " + f"{int(np.max(ids))}, but backend '{self._backend_type}' returned " + f"only {value.shape[1]} columns" + ) + + def _validate_body_state( + self, + backend: SimBackend, + root_body_ids: np.ndarray | None, + body_ids: np.ndarray | None, + ) -> None: + arrays = [values for values in (root_body_ids, body_ids) if values is not None] + if not arrays: + return + validation_ids = np.unique(np.concatenate(arrays)).astype(np.int32, copy=False) + for capability, getter, width in ( + ("body position state", backend.get_body_pos_w, 3), + ("body quaternion state", backend.get_body_quat_w, 4), + ("body linear velocity state", backend.get_body_lin_vel_w, 3), + ("body angular velocity state", backend.get_body_ang_vel_w, 3), + ): + value = self._read_state(capability, getter, validation_ids) + expected = (backend.num_envs, len(validation_ids), width) + if value.shape != expected: + raise ValueError( + f"Entity '{self.name}' capability '{capability}' on backend " + f"'{self._backend_type}' returned shape {value.shape}; expected {expected}" + ) + + def _materialize_actuator_ctrl_range( + self, backend: SimBackend, actuator_ids: np.ndarray | None + ) -> np.ndarray | None: + if actuator_ids is None: + return None + ranges = self._read_state("actuator control range", backend.get_actuator_ctrl_range) + expected = (backend.num_actuators, 2) + if ranges.shape != expected: + raise ValueError( + f"Entity '{self.name}' capability 'actuator control range' on backend " + f"'{self._backend_type}' returned shape {ranges.shape}; expected {expected}" + ) + selected = np.array(ranges[_as_column_index(actuator_ids)], copy=True) + selected.setflags(write=False) + return selected + + def _require_names(self, kind: str, names: tuple[str, ...] | None) -> tuple[str, ...]: + if names is None: + raise self._capability_error(kind, "the namespace was not declared in EntityCfg") + return names + + def _unsupported_names(self, kind: str) -> NoReturn: + raise self._capability_error(kind, "SimBackend does not declare this namespace") + + @property + def joint_names(self) -> tuple[str, ...]: + return self._require_names("joint", self._joint_names) + + @property + def body_names(self) -> tuple[str, ...]: + return self._require_names("body", self._body_names) + + @property + def geom_names(self) -> tuple[str, ...]: + return self._require_names("geom", self._geom_names) + + @property + def site_names(self) -> tuple[str, ...]: + return self._require_names("site", self._site_names) + + @property + def actuator_names(self) -> tuple[str, ...]: + return self._require_names("actuator", self._actuator_names) + + @property + def tendon_names(self) -> tuple[str, ...]: + return self._unsupported_names("tendon") + + @property + def camera_names(self) -> tuple[str, ...]: + return self._unsupported_names("camera") + + @property + def light_names(self) -> tuple[str, ...]: + return self._unsupported_names("light") + + @property + def material_names(self) -> tuple[str, ...]: + return self._unsupported_names("material") + + @property + def texture_names(self) -> tuple[str, ...]: + return self._unsupported_names("texture") + + @property + def pair_names(self) -> tuple[str, ...]: + return self._unsupported_names("pair") + + @property + def num_joints(self) -> int: + return len(self.joint_names) + + @property + def num_bodies(self) -> int: + return len(self.body_names) + + @property + def num_geoms(self) -> int: + return len(self.geom_names) + + @property + def num_sites(self) -> int: + return len(self.site_names) + + @property + def num_actuators(self) -> int: + return len(self.actuator_names) + + @property + def num_tendons(self) -> int: + return len(self.tendon_names) + + @property + def num_cameras(self) -> int: + return len(self.camera_names) + + @property + def num_lights(self) -> int: + return len(self.light_names) + + @property + def num_materials(self) -> int: + return len(self.material_names) + + @property + def num_textures(self) -> int: + return len(self.texture_names) + + @property + def num_pairs(self) -> int: + return len(self.pair_names) + + def _find( + self, + kind: str, + names: tuple[str, ...] | None, + keys: str | Sequence[str], + preserve_order: bool, + ) -> tuple[list[int], list[str]]: + return _resolve_matching_names(keys, self._require_names(kind, names), preserve_order) + + def find_joints( + self, keys: str | Sequence[str], preserve_order: bool = False + ) -> tuple[list[int], list[str]]: + return self._find("joint", self._joint_names, keys, preserve_order) + + def find_bodies( + self, keys: str | Sequence[str], preserve_order: bool = False + ) -> tuple[list[int], list[str]]: + return self._find("body", self._body_names, keys, preserve_order) + + def find_geoms( + self, keys: str | Sequence[str], preserve_order: bool = False + ) -> tuple[list[int], list[str]]: + return self._find("geom", self._geom_names, keys, preserve_order) + + def find_sites( + self, keys: str | Sequence[str], preserve_order: bool = False + ) -> tuple[list[int], list[str]]: + return self._find("site", self._site_names, keys, preserve_order) + + def find_actuators( + self, keys: str | Sequence[str], preserve_order: bool = False + ) -> tuple[list[int], list[str]]: + return self._find("actuator", self._actuator_names, keys, preserve_order) + + def find_tendons( + self, keys: str | Sequence[str], preserve_order: bool = False + ) -> tuple[list[int], list[str]]: + del keys, preserve_order + return self._unsupported_names("tendon") + + def find_cameras( + self, keys: str | Sequence[str], preserve_order: bool = False + ) -> tuple[list[int], list[str]]: + del keys, preserve_order + return self._unsupported_names("camera") + + def find_lights( + self, keys: str | Sequence[str], preserve_order: bool = False + ) -> tuple[list[int], list[str]]: + del keys, preserve_order + return self._unsupported_names("light") + + def find_materials( + self, keys: str | Sequence[str], preserve_order: bool = False + ) -> tuple[list[int], list[str]]: + del keys, preserve_order + return self._unsupported_names("material") + + def find_textures( + self, keys: str | Sequence[str], preserve_order: bool = False + ) -> tuple[list[int], list[str]]: + del keys, preserve_order + return self._unsupported_names("texture") + + def find_pairs( + self, keys: str | Sequence[str], preserve_order: bool = False + ) -> tuple[list[int], list[str]]: + del keys, preserve_order + return self._unsupported_names("pair") + + +class EntityScene(Mapping[str, Entity]): + """Read-only name-addressable collection of backend-bound entities.""" + + def __init__(self, entities: Mapping[str, EntityCfg], backend: SimBackend) -> None: + materialized: dict[str, Entity] = {} + for name, cfg in entities.items(): + if not isinstance(name, str) or not name: + raise TypeError(f"Scene entity names must be non-empty strings; got {name!r}") + if not isinstance(cfg, EntityCfg): + raise TypeError( + f"Scene entity '{name}' must be EntityCfg, got {type(cfg).__name__}" + ) + materialized[name] = Entity(name, cfg, backend) + self._entities = MappingProxyType(materialized) + + @classmethod + def from_scene_cfg(cls, cfg: SceneCfg, backend: SimBackend) -> EntityScene: + return cls(cfg.entities, backend) + + def __getitem__(self, name: str) -> Entity: + try: + return self._entities[name] + except KeyError as exc: + raise KeyError( + f"Scene entity '{name}' not found; available={list(self._entities)}" + ) from exc + + def __iter__(self) -> Iterator[str]: + return iter(self._entities) + + def __len__(self) -> int: + return len(self._entities) + + +__all__ = ["Entity", "EntityCfg", "EntityData", "EntityScene"] diff --git a/src/unilab/base/scene.py b/src/unilab/base/scene.py index dca7bc0df..4063e45f1 100644 --- a/src/unilab/base/scene.py +++ b/src/unilab/base/scene.py @@ -3,6 +3,7 @@ from dataclasses import dataclass, field from pathlib import Path +from unilab.base.entity import EntityCfg from unilab.terrains.terrain_generator import TerrainGeneratorCfg @@ -38,6 +39,8 @@ class SceneCfg: model_file: str fragment_files: list[str] = field(default_factory=list) terrain: TerrainSceneCfg | None = None + entities: dict[str, EntityCfg] = field(default_factory=dict) + """Logical entity partitions materialized by the base-owned manager facade.""" # Optional render-only model override. When set, offline playback/video # export renders this XML instead of ``model_file`` while physics keeps # using ``model_file``. Used to give the renderer a visual twin of the diff --git a/src/unilab/managers/_types.py b/src/unilab/managers/_types.py index 9ff2cbfeb..4459c33cb 100644 --- a/src/unilab/managers/_types.py +++ b/src/unilab/managers/_types.py @@ -7,6 +7,7 @@ from __future__ import annotations +from collections.abc import Sequence from typing import Any, Protocol import numpy as np @@ -15,17 +16,115 @@ class ManagerEntity(Protocol): """Cold-path entity metadata required by :class:`SceneEntityCfg`.""" - joint_names: list[str] - body_names: list[str] - geom_names: list[str] - site_names: list[str] - actuator_names: list[str] - tendon_names: list[str] - camera_names: list[str] - light_names: list[str] - material_names: list[str] - texture_names: list[str] - pair_names: list[str] + @property + def joint_names(self) -> Sequence[str]: ... + + @property + def body_names(self) -> Sequence[str]: ... + + @property + def geom_names(self) -> Sequence[str]: ... + + @property + def site_names(self) -> Sequence[str]: ... + + @property + def actuator_names(self) -> Sequence[str]: ... + + @property + def tendon_names(self) -> Sequence[str]: ... + + @property + def camera_names(self) -> Sequence[str]: ... + + @property + def light_names(self) -> Sequence[str]: ... + + @property + def material_names(self) -> Sequence[str]: ... + + @property + def texture_names(self) -> Sequence[str]: ... + + @property + def pair_names(self) -> Sequence[str]: ... + + @property + def num_joints(self) -> int: ... + + @property + def num_bodies(self) -> int: ... + + @property + def num_geoms(self) -> int: ... + + @property + def num_sites(self) -> int: ... + + @property + def num_actuators(self) -> int: ... + + @property + def num_tendons(self) -> int: ... + + @property + def num_cameras(self) -> int: ... + + @property + def num_lights(self) -> int: ... + + @property + def num_materials(self) -> int: ... + + @property + def num_textures(self) -> int: ... + + @property + def num_pairs(self) -> int: ... + + def find_joints( + self, keys: str | Sequence[str], preserve_order: bool = False + ) -> tuple[list[int], list[str]]: ... + + def find_bodies( + self, keys: str | Sequence[str], preserve_order: bool = False + ) -> tuple[list[int], list[str]]: ... + + def find_geoms( + self, keys: str | Sequence[str], preserve_order: bool = False + ) -> tuple[list[int], list[str]]: ... + + def find_sites( + self, keys: str | Sequence[str], preserve_order: bool = False + ) -> tuple[list[int], list[str]]: ... + + def find_actuators( + self, keys: str | Sequence[str], preserve_order: bool = False + ) -> tuple[list[int], list[str]]: ... + + def find_tendons( + self, keys: str | Sequence[str], preserve_order: bool = False + ) -> tuple[list[int], list[str]]: ... + + def find_cameras( + self, keys: str | Sequence[str], preserve_order: bool = False + ) -> tuple[list[int], list[str]]: ... + + def find_lights( + self, keys: str | Sequence[str], preserve_order: bool = False + ) -> tuple[list[int], list[str]]: ... + + def find_materials( + self, keys: str | Sequence[str], preserve_order: bool = False + ) -> tuple[list[int], list[str]]: ... + + def find_textures( + self, keys: str | Sequence[str], preserve_order: bool = False + ) -> tuple[list[int], list[str]]: ... + + def find_pairs( + self, keys: str | Sequence[str], preserve_order: bool = False + ) -> tuple[list[int], list[str]]: ... class ManagerScene(Protocol): diff --git a/src/unilab/managers/manager_base.py b/src/unilab/managers/manager_base.py index bac7275d6..dcf9fee16 100644 --- a/src/unilab/managers/manager_base.py +++ b/src/unilab/managers/manager_base.py @@ -151,9 +151,14 @@ def _prepare_terms(self) -> None: raise NotImplementedError def _resolve_common_term_cfg(self, term_name: str, term_cfg: ManagerTermBaseCfg) -> None: - del term_name # Unused. - for value in term_cfg.params.values(): + for param_name, value in term_cfg.params.items(): if isinstance(value, SceneEntityCfg): - value.resolve(self._env.scene) + try: + value.resolve(self._env.scene) + except (KeyError, TypeError, ValueError, NotImplementedError) as exc: + message = ( + f"{type(self).__name__} term '{term_name}' parameter '{param_name}': {exc}" + ) + raise type(exc)(message) from exc if inspect.isclass(term_cfg.func): term_cfg.func = term_cfg.func(cfg=term_cfg, env=self._env) diff --git a/src/unilab/managers/scene_entity_config.py b/src/unilab/managers/scene_entity_config.py index cb90c84d4..e7035b199 100644 --- a/src/unilab/managers/scene_entity_config.py +++ b/src/unilab/managers/scene_entity_config.py @@ -158,7 +158,12 @@ def resolve(self, scene: ManagerScene) -> None: entity = scene[self.name] for config in _FIELD_CONFIGS: - self._resolve_field(entity, config) + try: + self._resolve_field(entity, config) + except (TypeError, ValueError, NotImplementedError) as exc: + raise type(exc)( + f"SceneEntityCfg entity '{self.name}' {config.kind_label} selector: {exc}" + ) from exc def _resolve_field(self, entity: ManagerEntity, config: _FieldConfig) -> None: """Resolve a single field's names and IDs. @@ -199,7 +204,7 @@ def _resolve_field(self, entity: ManagerEntity, config: _FieldConfig) -> None: config.ids_attr, ) elif isinstance(ids, list): - self._resolve_ids_to_names(ids, entity_all_names, config.names_attr) + self._resolve_ids_to_names(ids, entity_all_names, config.names_attr, config.kind_label) def _normalize_to_list(self, value: str | int | tuple | list | None) -> list | None: """Convert single values to lists for uniform processing.""" @@ -225,6 +230,7 @@ def _validate_consistency( ValueError: If names and IDs don't match. """ found_ids, _ = find_method(names, preserve_order=self.preserve_order) + self._validate_ids(ids, len(entity_all_names), kind_label) computed_names = [entity_all_names[i] for i in ids] if found_ids != ids or computed_names != names: @@ -256,8 +262,26 @@ def _resolve_names_to_ids( setattr(self, ids_attr, found_ids) def _resolve_ids_to_names( - self, ids: list[int], entity_all_names: list[str], names_attr: str + self, + ids: list[int], + entity_all_names: list[str], + names_attr: str, + kind_label: str, ) -> None: """Resolve IDs to their corresponding names.""" + self._validate_ids(ids, len(entity_all_names), kind_label) resolved_names = [entity_all_names[i] for i in ids] setattr(self, names_attr, resolved_names) + + @staticmethod + def _validate_ids(ids: list[int], count: int, kind_label: str) -> None: + invalid_types = [ + value for value in ids if isinstance(value, bool) or not isinstance(value, int) + ] + if invalid_types: + raise TypeError(f"{kind_label} IDs must be integers; got {invalid_types}") + out_of_range = [value for value in ids if value < 0 or value >= count] + if out_of_range: + raise ValueError( + f"{kind_label} IDs {out_of_range} are out of range for {count} available entries" + ) diff --git a/tests/base/test_entity_facade.py b/tests/base/test_entity_facade.py new file mode 100644 index 000000000..a254da4b6 --- /dev/null +++ b/tests/base/test_entity_facade.py @@ -0,0 +1,351 @@ +from __future__ import annotations + +import ast +import inspect +from collections import Counter +from types import SimpleNamespace +from typing import Any, cast + +import numpy as np +import pytest + +import unilab.base.entity as entity_module +from unilab.assets import ASSETS_ROOT_PATH +from unilab.base.backend.base import SimBackend +from unilab.base.entity import EntityCfg, EntityScene +from unilab.base.scene import SceneCfg +from unilab.managers import RewardManager, RewardTermCfg, SceneEntityCfg + + +class _StrictBackendProfile: + """Strict public-contract fake shared by backend capability profiles.""" + + num_envs = 3 + num_actuators = 5 + + def __init__(self, backend_type: str, *, unsupported: frozenset[str] = frozenset()) -> None: + self.backend_type = backend_type + self.unsupported = unsupported + self.calls: Counter[str] = Counter() + self.joint_ids = {"hip": 2, "knee": 0, "ankle": 4} + self.body_ids = {"base": 4, "foot": 7} + self.site_ids = {"imu": 3} + self.geom_names = ("floor", "base_collision", "foot_collision") + self.actuator_names = ("knee", "unused", "hip", "unused_2", "ankle") + self.dof_pos = np.arange(self.num_envs * 5, dtype=np.float32).reshape(self.num_envs, 5) + self.dof_vel = self.dof_pos + 100.0 + base = np.arange(self.num_envs * 10 * 3, dtype=np.float32) + self.body_pos = base.reshape(self.num_envs, 10, 3) + self.body_quat = np.zeros((self.num_envs, 10, 4), dtype=np.float32) + self.body_quat[..., 0] = 1.0 + self.body_lin_vel = self.body_pos + 200.0 + self.body_ang_vel = self.body_pos + 300.0 + + def _check(self, capability: str) -> None: + self.calls[capability] += 1 + if capability in self.unsupported: + raise NotImplementedError(f"{self.backend_type} lacks {capability}") + + def get_body_ids(self, names) -> np.ndarray: + self._check("body names") + return np.asarray([self.body_ids[name] for name in names], dtype=np.int32) + + def get_joint_dof_pos_indices(self, names) -> np.ndarray: + self._check("joint position names") + return np.asarray([self.joint_ids[name] for name in names], dtype=np.int32) + + def get_joint_dof_vel_indices(self, names) -> np.ndarray: + self._check("joint velocity names") + return np.asarray([self.joint_ids[name] for name in names], dtype=np.int32) + + def get_site_ids(self, names) -> np.ndarray: + self._check("site names") + return np.asarray([self.site_ids[name] for name in names], dtype=np.int32) + + def get_geom_names(self) -> tuple[str, ...]: + self._check("geom names") + return self.geom_names + + def get_actuator_names(self) -> tuple[str, ...]: + self._check("actuator names") + return self.actuator_names + + def get_actuator_ctrl_range(self) -> np.ndarray: + self.calls["actuator range"] += 1 + return np.arange(10, dtype=np.float32).reshape(5, 2) + + def get_dof_pos(self) -> np.ndarray: + self._check("joint position state") + return self.dof_pos + + def get_dof_vel(self) -> np.ndarray: + self._check("joint velocity state") + return self.dof_vel + + def get_body_pos_w(self, ids: np.ndarray) -> np.ndarray: + self._check("body position state") + return self.body_pos[:, ids] + + def get_body_quat_w(self, ids: np.ndarray) -> np.ndarray: + self._check("body quaternion state") + return self.body_quat[:, ids] + + def get_body_lin_vel_w(self, ids: np.ndarray) -> np.ndarray: + self._check("body linear velocity state") + return self.body_lin_vel[:, ids] + + def get_body_ang_vel_w(self, ids: np.ndarray) -> np.ndarray: + self._check("body angular velocity state") + return self.body_ang_vel[:, ids] + + +def _scene(backend_type: str = "mujoco") -> tuple[_StrictBackendProfile, EntityScene]: + backend = _StrictBackendProfile(backend_type) + cfg = SceneCfg( + model_file="unused.xml", + entities={ + "robot": EntityCfg( + root_body_name="base", + joint_names=("ankle", "hip"), + body_names=("foot", "base"), + geom_names=("foot_collision", "base_collision"), + site_names=("imu",), + actuator_names=("ankle", "hip"), + ) + }, + ) + return backend, EntityScene.from_scene_cfg(cfg, cast(SimBackend, backend)) + + +@pytest.mark.parametrize("backend_type", ["mujoco", "motrix", "drake"]) +def test_backend_profiles_materialize_identical_local_entity_contract(backend_type: str) -> None: + backend, scene = _scene(backend_type) + robot = scene["robot"] + + assert robot.joint_names == ("ankle", "hip") + assert robot.body_names == ("foot", "base") + np.testing.assert_array_equal(robot.data.joint_pos, backend.dof_pos[:, [4, 2]]) + np.testing.assert_array_equal(robot.data.joint_vel, backend.dof_vel[:, [4, 2]]) + np.testing.assert_array_equal(robot.data.body_link_pos_w, backend.body_pos[:, [7, 4]]) + np.testing.assert_array_equal(robot.data.root_link_pos_w, backend.body_pos[:, 4]) + np.testing.assert_array_equal( + robot.data.actuator_ctrl_range, + np.arange(10, dtype=np.float32).reshape(5, 2)[[4, 2]], + ) + + +def test_scene_entity_cfg_resolves_only_against_cached_names() -> None: + backend, scene = _scene() + cold_path_calls = backend.calls.copy() + + cfg = SceneEntityCfg( + "robot", + joint_names=("hip", "ankle"), + body_names=".*", + geom_names="foot_.*", + site_names="imu", + actuator_names=["hip", "ankle"], + preserve_order=True, + ) + cfg.resolve(scene) + + assert cfg.joint_ids == [1, 0] + assert cfg.body_ids == slice(None) + assert cfg.geom_ids == [0] + assert cfg.site_ids == slice(None) + assert cfg.actuator_ids == [1, 0] + assert backend.calls == cold_path_calls + + for _ in range(3): + scene["robot"].data.joint_pos + scene["robot"].data.root_link_quat_w + for key in ( + "body names", + "joint position names", + "joint velocity names", + "site names", + "geom names", + "actuator names", + ): + assert backend.calls[key] == cold_path_calls[key] + + +@pytest.mark.parametrize( + ("ids", "error_type", "message"), + [ + ([-1], ValueError, "EntityCfg entity 'robot' joint selector.*out of range"), + ([2], ValueError, "EntityCfg entity 'robot' joint selector.*out of range"), + ([True], TypeError, "EntityCfg entity 'robot' joint selector.*must be integers"), + (["0"], TypeError, "EntityCfg entity 'robot' joint selector.*must be integers"), + ], +) +def test_scene_entity_cfg_rejects_invalid_ids(ids, error_type, message: str) -> None: + _, scene = _scene() + cfg = SceneEntityCfg("robot", joint_ids=ids) + with pytest.raises(error_type, match=message): + cfg.resolve(scene) + + +def test_scene_entity_cfg_reports_entity_for_invalid_regex() -> None: + _, scene = _scene() + with pytest.raises( + ValueError, + match="SceneEntityCfg entity 'robot' joint selector.*Invalid entity selector regex", + ): + SceneEntityCfg("robot", joint_names="[").resolve(scene) + + +def test_missing_entity_namespace_and_backend_capability_fail_closed() -> None: + backend = _StrictBackendProfile("drake", unsupported=frozenset({"actuator names"})) + with pytest.raises( + NotImplementedError, + match="Entity 'robot'.*actuator.*backend 'drake'", + ): + EntityScene( + {"robot": EntityCfg(actuator_names=("hip",))}, + cast(SimBackend, backend), + ) + + _, scene = _scene("motrix") + with pytest.raises( + NotImplementedError, + match="Entity 'robot'.*tendon.*backend 'motrix'", + ): + SceneEntityCfg("robot", tendon_names=".*").resolve(scene) + + sparse = EntityScene( + {"robot": EntityCfg(joint_names=("hip",))}, + cast(SimBackend, _StrictBackendProfile("mujoco")), + ) + with pytest.raises(NotImplementedError, match="body.*not declared"): + SceneEntityCfg("robot", body_names=".*").resolve(sparse) + + state_missing = _StrictBackendProfile("mujoco", unsupported=frozenset({"body position state"})) + with pytest.raises( + NotImplementedError, + match="Entity 'robot'.*body position state.*backend 'mujoco'", + ): + EntityScene( + {"robot": EntityCfg(root_body_name="base")}, + cast(SimBackend, state_missing), + ) + + +def test_entity_declaration_rejects_coercion_and_duplicate_names() -> None: + backend = cast(SimBackend, _StrictBackendProfile("mujoco")) + with pytest.raises(TypeError, match="sequence of strings, not a scalar"): + EntityScene( + {"robot": EntityCfg(joint_names=cast(Any, "hip"))}, + backend, + ) + with pytest.raises(TypeError, match="joint names must be strings"): + EntityScene( + {"robot": EntityCfg(joint_names=cast(Any, ("hip", 1)))}, + backend, + ) + with pytest.raises(ValueError, match="joint names must be unique"): + EntityScene( + {"robot": EntityCfg(joint_names=("hip", "hip"))}, + backend, + ) + + +def test_manager_resolution_error_has_manager_term_entity_capability_and_backend() -> None: + _, scene = _scene("drake") + env = SimpleNamespace( + num_envs=3, + scene=scene, + rng=np.random.default_rng(7), + max_episode_length_s=2.0, + ) + + with pytest.raises( + NotImplementedError, + match=( + "RewardManager term 'unsupported' parameter 'asset_cfg'.*" + "Entity 'robot'.*tendon.*backend 'drake'" + ), + ): + RewardManager( + { + "unsupported": RewardTermCfg( + func=lambda env, asset_cfg: np.zeros(env.num_envs), + weight=1.0, + params={"asset_cfg": SceneEntityCfg("robot", tendon_names=".*")}, + ) + }, + cast(Any, env), + ) + + +def test_entity_facade_has_no_backend_model_or_asset_access() -> None: + tree = ast.parse(inspect.getsource(entity_module)) + accessed_attributes = {node.attr for node in ast.walk(tree) if isinstance(node, ast.Attribute)} + assert "model" not in accessed_attributes + assert "scene_model_file" not in accessed_attributes + + +def test_real_mujoco_entity_selector_and_numpy_state_smoke() -> None: + from unilab.base.backend.mujoco.backend import MuJoCoBackend + + joint_names = ( + "FL_hip_joint", + "FL_thigh_joint", + "FL_calf_joint", + "FR_hip_joint", + "FR_thigh_joint", + "FR_calf_joint", + "RL_hip_joint", + "RL_thigh_joint", + "RL_calf_joint", + "RR_hip_joint", + "RR_thigh_joint", + "RR_calf_joint", + ) + actuator_names = ( + "FR_hip", + "FR_thigh", + "FR_calf", + "FL_hip", + "FL_thigh", + "FL_calf", + "RR_hip", + "RR_thigh", + "RR_calf", + "RL_hip", + "RL_thigh", + "RL_calf", + ) + scene_cfg = SceneCfg( + model_file=str(ASSETS_ROOT_PATH / "robots" / "go2" / "scene_flat.xml"), + entities={ + "robot": EntityCfg( + root_body_name="base", + joint_names=joint_names, + body_names=("base",), + actuator_names=actuator_names, + ) + }, + ) + backend = MuJoCoBackend( + scene_cfg, + num_envs=2, + sim_dt=0.01, + base_name="base", + add_body_sensors=True, + ) + backend.materialize() + scene = EntityScene.from_scene_cfg(scene_cfg, backend) + + selector = SceneEntityCfg("robot", joint_names=".*_calf_joint") + selector.resolve(scene) + assert selector.joint_ids == [2, 5, 8, 11] + assert scene["robot"].data.joint_pos.shape == (2, 12) + assert scene["robot"].data.root_link_pose_w.shape == (2, 7) + + +def test_scene_cfg_entity_defaults_are_not_shared() -> None: + first = SceneCfg(model_file="first.xml") + second = SceneCfg(model_file="second.xml") + first.entities["robot"] = EntityCfg() + assert second.entities == {} From 211c48dd137e079fbf1e23b177d2d45b85ca63c6 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Mon, 17 Aug 2026 20:11:16 +0800 Subject: [PATCH 004/169] feat(env): add manager-based lifecycle (#1051) --- src/unilab/base/entity.py | 108 ++++- src/unilab/base/np_env.py | 30 +- src/unilab/envs/__init__.py | 14 +- src/unilab/envs/manager_based_rl_env.py | 556 ++++++++++++++++++++++++ src/unilab/managers/_types.py | 18 +- src/unilab/managers/action_manager.py | 14 +- tests/base/test_entity_facade.py | 36 ++ tests/envs/test_manager_based_rl_env.py | 449 +++++++++++++++++++ tests/managers/test_core_managers.py | 26 ++ 9 files changed, 1229 insertions(+), 22 deletions(-) create mode 100644 src/unilab/envs/manager_based_rl_env.py create mode 100644 tests/envs/test_manager_based_rl_env.py diff --git a/src/unilab/base/entity.py b/src/unilab/base/entity.py index e9c3e0842..771a5f2fe 100644 --- a/src/unilab/base/entity.py +++ b/src/unilab/base/entity.py @@ -140,7 +140,9 @@ def __init__( joint_pos_ids: np.ndarray | None, joint_vel_ids: np.ndarray | None, body_ids: np.ndarray | None, + actuator_ids: np.ndarray | None, actuator_ctrl_range: np.ndarray | None, + control_buffer: np.ndarray | None, entity_name: str, backend_type: str, ) -> None: @@ -151,7 +153,9 @@ def __init__( self._joint_pos_index = None if joint_pos_ids is None else _as_column_index(joint_pos_ids) self._joint_vel_index = None if joint_vel_ids is None else _as_column_index(joint_vel_ids) self._body_ids = body_ids + self._actuator_index = None if actuator_ids is None else _as_column_index(actuator_ids) self._actuator_ctrl_range = actuator_ctrl_range + self._control_buffer = control_buffer def _require(self, value, capability: str): if value is None: @@ -231,11 +235,81 @@ def body_link_vel_w(self) -> np.ndarray: def actuator_ctrl_range(self) -> np.ndarray: return self._require(self._actuator_ctrl_range, "actuator control range") + def write_ctrl( + self, + values: np.ndarray, + env_ids: np.ndarray | slice | None = None, + ) -> None: + """Write entity-local actuator controls into the env-owned control buffer. + + This is an in-memory scene write, analogous to the pinned manager runtime's + entity target buffers. Physics remains owned by ``NpEnv``/``SimBackend``; + this method never steps or calls a backend-private API. + """ + actuator_index = self._require(self._actuator_index, "actuator control write") + control = self._require(self._control_buffer, "actuator control write") + if not isinstance(values, np.ndarray): + raise TypeError( + f"Entity '{self._entity_name}' write_ctrl expected np.ndarray, " + f"received {type(values).__name__}" + ) + row_index: np.ndarray | slice + if env_ids is None: + row_index = slice(None) + row_count = control.shape[0] + elif isinstance(env_ids, slice): + row_index = env_ids + row_count = len(range(*env_ids.indices(control.shape[0]))) + else: + raw_ids = np.asarray(env_ids) + if ( + raw_ids.ndim != 1 + or not np.issubdtype(raw_ids.dtype, np.integer) + or np.issubdtype(raw_ids.dtype, np.bool_) + ): + raise TypeError( + f"Entity '{self._entity_name}' write_ctrl env_ids must be a 1-D " + f"integer array or slice, got shape={raw_ids.shape}, dtype={raw_ids.dtype}" + ) + row_index = np.asarray(raw_ids, dtype=np.intp) + if np.any(row_index < 0) or np.any(row_index >= control.shape[0]): + raise IndexError( + f"Entity '{self._entity_name}' write_ctrl env_ids out of range for " + f"{control.shape[0]} environments: {row_index.tolist()}" + ) + if np.unique(row_index).size != row_index.size: + raise ValueError( + f"Entity '{self._entity_name}' write_ctrl env_ids contain duplicates: " + f"{row_index.tolist()}" + ) + row_count = len(row_index) + + actuator_count = len(self._require(self._actuator_ctrl_range, "actuator control write")) + expected = (row_count, actuator_count) + if values.shape != expected: + raise ValueError( + f"Entity '{self._entity_name}' write_ctrl expected shape {expected}, " + f"received {values.shape}" + ) + if not np.isfinite(values).all(): + raise ValueError(f"Entity '{self._entity_name}' write_ctrl received NaN or Inf") + + if isinstance(row_index, slice) or isinstance(actuator_index, slice): + control[row_index, actuator_index] = values + else: + control[row_index[:, None], actuator_index[None, :]] = values + class Entity: """Logical entity with cached local-to-backend mappings.""" - def __init__(self, name: str, cfg: EntityCfg, backend: SimBackend) -> None: + def __init__( + self, + name: str, + cfg: EntityCfg, + backend: SimBackend, + control_buffer: np.ndarray | None = None, + ) -> None: if not name: raise ValueError("Entity name must be a non-empty string") self.name = name @@ -293,6 +367,18 @@ def __init__(self, name: str, cfg: EntityCfg, backend: SimBackend) -> None: self._validate_joint_state(backend, joint_pos_ids, joint_vel_ids) self._validate_body_state(backend, root_body_ids, body_ids) actuator_ctrl_range = self._materialize_actuator_ctrl_range(backend, actuator_ids) + if control_buffer is not None: + expected_control_shape = (backend.num_envs, backend.num_actuators) + if control_buffer.shape != expected_control_shape: + raise ValueError( + f"Entity '{self.name}' control buffer has shape {control_buffer.shape}; " + f"expected {expected_control_shape} on backend '{self._backend_type}'" + ) + if not np.issubdtype(control_buffer.dtype, np.floating): + raise TypeError( + f"Entity '{self.name}' control buffer must have floating dtype, " + f"got {control_buffer.dtype}" + ) self.data = EntityData( backend, @@ -300,7 +386,9 @@ def __init__(self, name: str, cfg: EntityCfg, backend: SimBackend) -> None: joint_pos_ids=joint_pos_ids, joint_vel_ids=joint_vel_ids, body_ids=body_ids, + actuator_ids=actuator_ids, actuator_ctrl_range=actuator_ctrl_range, + control_buffer=control_buffer, entity_name=self.name, backend_type=self._backend_type, ) @@ -601,7 +689,12 @@ def find_pairs( class EntityScene(Mapping[str, Entity]): """Read-only name-addressable collection of backend-bound entities.""" - def __init__(self, entities: Mapping[str, EntityCfg], backend: SimBackend) -> None: + def __init__( + self, + entities: Mapping[str, EntityCfg], + backend: SimBackend, + control_buffer: np.ndarray | None = None, + ) -> None: materialized: dict[str, Entity] = {} for name, cfg in entities.items(): if not isinstance(name, str) or not name: @@ -610,12 +703,17 @@ def __init__(self, entities: Mapping[str, EntityCfg], backend: SimBackend) -> No raise TypeError( f"Scene entity '{name}' must be EntityCfg, got {type(cfg).__name__}" ) - materialized[name] = Entity(name, cfg, backend) + materialized[name] = Entity(name, cfg, backend, control_buffer) self._entities = MappingProxyType(materialized) @classmethod - def from_scene_cfg(cls, cfg: SceneCfg, backend: SimBackend) -> EntityScene: - return cls(cfg.entities, backend) + def from_scene_cfg( + cls, + cfg: SceneCfg, + backend: SimBackend, + control_buffer: np.ndarray | None = None, + ) -> EntityScene: + return cls(cfg.entities, backend, control_buffer) def __getitem__(self, name: str) -> Entity: try: diff --git a/src/unilab/base/np_env.py b/src/unilab/base/np_env.py index 0e74f9640..b219b06f3 100644 --- a/src/unilab/base/np_env.py +++ b/src/unilab/base/np_env.py @@ -117,6 +117,7 @@ def __init__(self, cfg: EnvCfg, backend: SimBackend, num_envs: int): self._init_randomization_applied = False self._nan_guard: NanGuard | None = None self._autoreset = True + self._autoreset_reset_active = False self._nan_guard_model_file = self._resolve_nan_guard_model_file() @property @@ -156,12 +157,7 @@ def init_state(self) -> NpEnvState: reward = np.zeros((self._num_envs,), dtype=dtype) terminated = np.ones((self._num_envs,), dtype=bool) truncated = np.zeros((self._num_envs,), dtype=bool) - if self._cfg.max_episode_steps: - steps = np.random.randint( - 0, self._cfg.max_episode_steps, size=(self._num_envs,), dtype=np.uint32 - ) - else: - steps = np.zeros((self._num_envs,), dtype=np.uint32) + steps = self._initial_episode_steps() info: dict = {"steps": steps} self._state = NpEnvState(obs, reward, terminated, truncated, info) @@ -169,6 +165,22 @@ def init_state(self) -> NpEnvState: self._clear_step_final_observation() return self._state + def _initial_episode_steps(self) -> np.ndarray: + """Return initial per-env episode counters. + + Existing monolithic tasks keep their randomized initialization. A lifecycle + with different public semantics can override this cold-path hook without + duplicating :meth:`init_state` or the autoreset machinery. + """ + if self._cfg.max_episode_steps: + return np.random.randint( + 0, + self._cfg.max_episode_steps, + size=(self._num_envs,), + dtype=np.uint32, + ) + return np.zeros((self._num_envs,), dtype=np.uint32) + def step(self, actions: np.ndarray) -> NpEnvState: step_t0 = time.perf_counter() @@ -274,7 +286,11 @@ def _reset_done_envs(self) -> None: detail_timing["reset_done_terminal_obs_ms"] = (time.perf_counter() - t0) * 1000.0 t0 = time.perf_counter() - new_obs, info1 = self.reset(env_indices) + self._autoreset_reset_active = True + try: + new_obs, info1 = self.reset(env_indices) + finally: + self._autoreset_reset_active = False detail_timing["reset_done_reset_call_ms"] = (time.perf_counter() - t0) * 1000.0 if self._dr_manager is not None: detail_timing.update(self._dr_manager.last_reset_timing_ms) diff --git a/src/unilab/envs/__init__.py b/src/unilab/envs/__init__.py index 85a4be3f6..4542de1a5 100644 --- a/src/unilab/envs/__init__.py +++ b/src/unilab/envs/__init__.py @@ -1 +1,13 @@ -"""Environment""" +"""Environment public API.""" + +from unilab.envs.manager_based_rl_env import ManagerBasedRLEnv as ManagerBasedRLEnv +from unilab.envs.manager_based_rl_env import ManagerBasedRlEnv as ManagerBasedRlEnv +from unilab.envs.manager_based_rl_env import ManagerBasedRLEnvCfg as ManagerBasedRLEnvCfg +from unilab.envs.manager_based_rl_env import ManagerBasedRlEnvCfg as ManagerBasedRlEnvCfg + +__all__ = [ + "ManagerBasedRLEnv", + "ManagerBasedRLEnvCfg", + "ManagerBasedRlEnv", + "ManagerBasedRlEnvCfg", +] diff --git a/src/unilab/envs/manager_based_rl_env.py b/src/unilab/envs/manager_based_rl_env.py new file mode 100644 index 000000000..1b6773812 --- /dev/null +++ b/src/unilab/envs/manager_based_rl_env.py @@ -0,0 +1,556 @@ +# Derived from mujocolab/mjlab v1.6.0 (0fb8a681), +# src/mjlab/envs/manager_based_rl_env.py. +# Copyright 2025, The mjlab Developers. +# Modified by UniLab for the NumPy NpEnv/SimBackend contracts; Apache-2.0. +"""Community-compatible manager lifecycle on UniLab's NumPy runtime.""" + +from __future__ import annotations + +import math +import secrets +from dataclasses import dataclass, field +from typing import Any + +import gymnasium as gym +import numpy as np + +from unilab.base.backend import SimBackend +from unilab.base.base import EnvCfg +from unilab.base.entity import EntityScene +from unilab.base.np_env import NpEnv, NpEnvState +from unilab.base.scene import SceneCfg +from unilab.dtype_config import get_global_dtype +from unilab.managers import ( + ActionManager, + ActionTermCfg, + CommandManager, + CommandTermCfg, + CurriculumManager, + CurriculumTermCfg, + EventManager, + EventTermCfg, + MetricsManager, + MetricsTermCfg, + NullCommandManager, + NullCurriculumManager, + NullMetricsManager, + NullRecorderManager, + ObservationGroupCfg, + ObservationManager, + RecorderManager, + RecorderTermCfg, + RewardManager, + RewardTermCfg, + TerminationManager, + TerminationTermCfg, +) + + +@dataclass +class ManagerBasedRlEnvCfg(EnvCfg): + """Configuration for the manager-based NumPy environment. + + Serializable values remain ordinary dataclass fields so task-owned Hydra owner + configs can overlay them without introducing a second configuration runtime. + """ + + observations: dict[str, ObservationGroupCfg | None] = field(default_factory=dict) + actions: dict[str, ActionTermCfg | None] = field(default_factory=dict) + events: dict[str, EventTermCfg | None] = field(default_factory=dict) + rewards: dict[str, RewardTermCfg | None] = field(default_factory=dict) + terminations: dict[str, TerminationTermCfg | None] = field(default_factory=dict) + commands: dict[str, CommandTermCfg | None] = field(default_factory=dict) + curriculum: dict[str, CurriculumTermCfg | None] = field(default_factory=dict) + metrics: dict[str, MetricsTermCfg | None] = field(default_factory=dict) + recorders: dict[str, RecorderTermCfg | None] = field(default_factory=dict) + + seed: int | None = None + is_finite_horizon: bool = False + auto_reset: bool = True + scale_rewards_by_dt: bool = True + policy_observation_group: str = "policy" + critic_observation_group: str | None = None + + def validate(self) -> None: + for name, value in (("sim_dt", self.sim_dt), ("ctrl_dt", self.ctrl_dt)): + if isinstance(value, bool) or not isinstance(value, (int, float, np.number)): + raise TypeError(f"ManagerBasedRlEnvCfg {name} must be a real number") + if not np.isfinite(value) or value <= 0.0: + raise ValueError(f"ManagerBasedRlEnvCfg {name} must be finite and positive") + super().validate() + ratio = self.ctrl_dt / self.sim_dt + if not np.isclose(ratio, round(ratio), rtol=0.0, atol=1e-9): + raise ValueError( + "ManagerBasedRlEnvCfg ctrl_dt must be an integer multiple of sim_dt; " + f"received ctrl_dt={self.ctrl_dt}, sim_dt={self.sim_dt}" + ) + if self.max_episode_seconds is None: + raise ValueError("ManagerBasedRlEnvCfg max_episode_seconds must be finite and positive") + if isinstance(self.max_episode_seconds, bool) or not isinstance( + self.max_episode_seconds, (int, float, np.number) + ): + raise TypeError("ManagerBasedRlEnvCfg max_episode_seconds must be a real number") + if not np.isfinite(self.max_episode_seconds) or self.max_episode_seconds <= 0.0: + raise ValueError("ManagerBasedRlEnvCfg max_episode_seconds must be finite and positive") + if self.seed is not None and ( + isinstance(self.seed, bool) + or not isinstance(self.seed, (int, np.integer)) + or self.seed < 0 + ): + raise ValueError("ManagerBasedRlEnvCfg seed must be a non-negative integer or None") + if self.seed is not None: + self.seed = int(self.seed) + for name in ( + "observations", + "actions", + "events", + "rewards", + "terminations", + "commands", + "curriculum", + "metrics", + "recorders", + ): + if not isinstance(getattr(self, name), dict): + raise TypeError(f"ManagerBasedRlEnvCfg {name} must be a dict") + for name in ("is_finite_horizon", "auto_reset", "scale_rewards_by_dt"): + if not isinstance(getattr(self, name), bool): + raise TypeError(f"ManagerBasedRlEnvCfg {name} must be bool") + if not isinstance(self.policy_observation_group, str) or not self.policy_observation_group: + raise ValueError("policy_observation_group must be a non-empty string") + if self.critic_observation_group is not None: + if ( + not isinstance(self.critic_observation_group, str) + or not self.critic_observation_group + ): + raise ValueError("critic_observation_group must be a non-empty string or None") + if self.critic_observation_group == self.policy_observation_group: + raise ValueError("policy and critic observation groups must be different") + if not isinstance(self.scene, SceneCfg): + raise TypeError( + "ManagerBasedRlEnvCfg scene must be a SceneCfg instance, " + f"got {type(self.scene).__name__}" + ) + + +class ManagerBasedRlEnv(NpEnv): + """Manager-Based API adapter that reuses the single :class:`NpEnv` lifecycle.""" + + is_vector_env = True + _cfg: ManagerBasedRlEnvCfg + + def __init__(self, cfg: ManagerBasedRlEnvCfg, backend: SimBackend, num_envs: int): + if not isinstance(cfg, ManagerBasedRlEnvCfg): + raise TypeError( + f"ManagerBasedRlEnv expected ManagerBasedRlEnvCfg, received {type(cfg).__name__}" + ) + cfg.validate() + if isinstance(num_envs, bool) or not isinstance(num_envs, int) or num_envs <= 0: + raise ValueError( + f"ManagerBasedRlEnv num_envs must be a positive integer, got {num_envs!r}" + ) + if backend.num_envs != num_envs: + raise ValueError( + f"ManagerBasedRlEnv num_envs={num_envs} does not match backend " + f"'{backend.backend_type}' num_envs={backend.num_envs}" + ) + + super().__init__(cfg, backend, num_envs) + actual_seed = cfg.seed if cfg.seed is not None else secrets.randbits(63) + cfg.seed = actual_seed + self.rng = np.random.default_rng(actual_seed) + + self._control = np.zeros((num_envs, backend.num_actuators), dtype=get_global_dtype()) + assert cfg.scene is not None + self.scene = EntityScene.from_scene_cfg(cfg.scene, backend, self._control) + + self.common_step_counter = 0 + self._sim_step_counter = 0 + self.episode_length_buf = np.zeros(num_envs, dtype=np.int64) + self.reset_buf = np.zeros(num_envs, dtype=np.bool_) + self.reset_terminated = np.zeros(num_envs, dtype=np.bool_) + self.reset_time_outs = np.zeros(num_envs, dtype=np.bool_) + self.reward_buf = np.zeros(num_envs, dtype=get_global_dtype()) + self.obs_buf: dict[str, np.ndarray] = {} + self.extras: dict[str, Any] = {"log": {}} + self._command_dt = np.zeros(num_envs, dtype=get_global_dtype()) + self._no_truncation = np.zeros(num_envs, dtype=np.bool_) + self._manual_reset_pending = np.zeros(num_envs, dtype=np.bool_) + self._has_transition = False + self._uses_pre_step_control = False + + self._load_managers() + self._mapped_obs_dims = self._validate_observation_mapping() + self._validate_substep_capabilities() + self._configure_action_control() + self.set_autoreset(cfg.auto_reset) + + if "startup" in self.event_manager.available_modes: + self.event_manager.apply(mode="startup") + + @property + def physics_dt(self) -> float: + return self._cfg.sim_dt + + @property + def step_dt(self) -> float: + return self._cfg.ctrl_dt + + @property + def max_episode_length_s(self) -> float: + assert self._cfg.max_episode_seconds is not None + return self._cfg.max_episode_seconds + + @property + def max_episode_length(self) -> int: + return math.ceil(self.max_episode_length_s / self.step_dt) + + @property + def obs_groups_spec(self) -> dict[str, int]: + return dict(self._mapped_obs_dims) + + @property + def action_space(self) -> gym.Space: + return gym.spaces.Box( + low=-np.inf, + high=np.inf, + shape=(self.action_manager.total_action_dim,), + dtype=get_global_dtype(), + ) + + @property + def unwrapped(self) -> ManagerBasedRlEnv: + return self + + def _load_managers(self) -> None: + """Construct managers in the pinned community dependency order.""" + self.event_manager = EventManager(self._cfg.events, self) + self.command_manager = ( + CommandManager(self._cfg.commands, self) if self._cfg.commands else NullCommandManager() + ) + self.action_manager = ActionManager(self._cfg.actions, self) + self.observation_manager = ObservationManager(self._cfg.observations, self) + self.termination_manager = TerminationManager(self._cfg.terminations, self) + self.reward_manager = RewardManager( + self._cfg.rewards, + self, + scale_by_dt=self._cfg.scale_rewards_by_dt, + ) + self.curriculum_manager = ( + CurriculumManager(self._cfg.curriculum, self) + if self._cfg.curriculum + else NullCurriculumManager() + ) + self.metrics_manager = ( + MetricsManager(self._cfg.metrics, self) if self._cfg.metrics else NullMetricsManager() + ) + self.recorder_manager = ( + RecorderManager(self._cfg.recorders, self) + if self._cfg.recorders + else NullRecorderManager() + ) + + def _validate_observation_mapping(self) -> dict[str, int]: + mapping = {"obs": self._cfg.policy_observation_group} + if self._cfg.critic_observation_group is not None: + mapping["critic"] = self._cfg.critic_observation_group + dims: dict[str, int] = {} + for output_name, group_name in mapping.items(): + if group_name not in self.observation_manager.active_terms: + raise KeyError( + f"ManagerBasedRlEnv observation mapping '{output_name}' requests group " + f"'{group_name}', available={list(self.observation_manager.active_terms)}" + ) + if not self.observation_manager.group_obs_concatenate[group_name]: + raise ValueError( + f"ManagerBasedRlEnv observation group '{group_name}' mapped to " + f"NpEnvState.obs['{output_name}'] must concatenate terms" + ) + group_dim = self.observation_manager.group_obs_dim[group_name] + if not isinstance(group_dim, tuple) or len(group_dim) != 1: + raise ValueError( + f"ManagerBasedRlEnv observation group '{group_name}' mapped to " + f"NpEnvState.obs['{output_name}'] must be one-dimensional; got {group_dim}" + ) + dims[output_name] = int(group_dim[0]) + return dims + + def _validate_substep_capabilities(self) -> None: + per_substep_terms = [ + name + for name, term_cfg in self._cfg.metrics.items() + if term_cfg is not None and term_cfg.per_substep + ] + if self._cfg.sim_substeps > 1 and per_substep_terms: + raise NotImplementedError( + "MetricsManager capability 'post-physics per-substep metrics' is unavailable " + f"on backend '{self._backend.backend_type}' with sim_substeps=" + f"{self._cfg.sim_substeps}; terms={per_substep_terms}. SimBackend does not " + "declare a post-substep hook." + ) + + def _configure_action_control(self) -> None: + if self._cfg.sim_substeps <= 1 or not self.action_manager.active_terms: + return + try: + self._backend.set_pre_step_control(self._apply_manager_control) + except NotImplementedError as exc: + raise NotImplementedError( + "ActionManager capability 'apply actions on every physics substep' is " + f"unavailable on backend '{self._backend.backend_type}': {exc}" + ) from exc + self._uses_pre_step_control = True + + def _apply_manager_control( + self, + backend: SimBackend, + control: np.ndarray, + ) -> np.ndarray: + del backend, control + self._sim_step_counter += 1 + self.action_manager.apply_action() + return self._control + + def _initial_episode_steps(self) -> np.ndarray: + return np.zeros((self.num_envs,), dtype=np.uint32) + + def init_state(self) -> NpEnvState: + state = super().init_state() + self.obs_buf = state.obs + self.reward_buf = state.reward + self.extras = state.info + return state + + def step(self, actions: np.ndarray) -> NpEnvState: + if not self._autoreset and np.any(self._manual_reset_pending): + pending = np.flatnonzero(self._manual_reset_pending).tolist() + raise RuntimeError( + f"ManagerBasedRlEnv environments {pending} must be reset before step() " + "when auto_reset=False" + ) + state = super().step(actions) + if not self._autoreset: + self._manual_reset_pending |= state.terminated | state.truncated + self.recorder_manager.record_post_step() + return state + + def apply_action(self, actions: np.ndarray, state: NpEnvState) -> np.ndarray: + del state + self.action_manager.process_action(actions) + if not self._uses_pre_step_control: + self._sim_step_counter += self._cfg.sim_substeps + self.action_manager.apply_action() + return self._control + + def update_state(self, state: NpEnvState) -> NpEnvState: + log: dict[str, Any] = {} + state.info["log"] = log + self.extras = state.info + + np.add(state.info["steps"], 1, out=self.episode_length_buf) + self.common_step_counter = self.step_counter + 1 + self._sim_step_counter = self.common_step_counter * self._cfg.sim_substeps + + self.termination_manager.compute() + if self._cfg.is_finite_horizon: + np.logical_or( + self.termination_manager.terminated, + self.termination_manager.time_outs, + out=self.reset_terminated, + ) + self.reset_time_outs.fill(False) + else: + np.copyto(self.reset_terminated, self.termination_manager.terminated) + np.copyto(self.reset_time_outs, self.termination_manager.time_outs) + np.logical_or(self.reset_terminated, self.reset_time_outs, out=self.reset_buf) + + self.reward_buf = self.reward_manager.compute(dt=self.step_dt) + if self._cfg.sim_substeps == 1: + self.metrics_manager.compute_substep() + self.metrics_manager.compute() + + if "step" in self.event_manager.available_modes: + self.event_manager.apply(mode="step", dt=self.step_dt) + if "interval" in self.event_manager.available_modes: + self.event_manager.apply(mode="interval", dt=self.step_dt) + + self._command_dt.fill(self.step_dt) + self._command_dt[self.reset_buf] = 0.0 + self.command_manager.compute(dt=self._command_dt) + manager_obs = self.observation_manager.compute(update_history=True) + self.obs_buf = self._map_observations(manager_obs) + self._has_transition = True + + return state.replace( + obs=self.obs_buf, + reward=self.reward_buf, + terminated=self.reset_terminated, + truncated=self.reset_time_outs, + info=state.info, + ) + + def _compute_truncated(self, state: NpEnvState) -> np.ndarray: + del state + self._no_truncation.fill(False) + return self._no_truncation + + def reset( + self, + env_indices: np.ndarray | None = None, + *, + seed: int | None = None, + env_ids: np.ndarray | None = None, + options: dict[str, Any] | None = None, + ) -> tuple[dict[str, np.ndarray], dict[str, Any]]: + del options + ids = self._normalize_reset_ids(env_indices, env_ids) + if seed is not None: + self.seed(seed) + if self._state is None: + all_ids = np.arange(self.num_envs, dtype=np.int32) + if not np.array_equal(ids, all_ids): + raise RuntimeError( + "ManagerBasedRlEnv requires a full reset before the first partial reset" + ) + state = self.init_state() + return state.obs, {"log": state.info.get("log", {})} + + done_ids = ids[self.reset_buf[ids]] + if self._has_transition and len(done_ids) > 0: + self.recorder_manager.record_pre_reset(done_ids) + + log: dict[str, Any] = {} + self.curriculum_manager.compute(env_ids=ids) + if "reset" in self.event_manager.available_modes: + self.event_manager.apply( + mode="reset", + env_ids=ids, + global_env_step_count=self.step_counter, + ) + + for manager in ( + self.observation_manager, + self.action_manager, + self.reward_manager, + self.metrics_manager, + self.curriculum_manager, + self.command_manager, + self.event_manager, + self.termination_manager, + ): + log.update(manager.reset(ids)) + + self.episode_length_buf[ids] = 0 + self._control[ids] = 0.0 + self._manual_reset_pending[ids] = False + if self._state is not None: + self._state.info["steps"][ids] = 0 + + self.command_manager.compute(dt=0.0, env_ids=ids) + manager_obs = self.observation_manager.compute(update_history=True, env_ids=ids) + mapped_obs = self._map_observations(manager_obs) + reset_obs = {name: values[ids].copy() for name, values in mapped_obs.items()} + + if self._state is not None: + for name, values in reset_obs.items(): + self._state.obs[name][ids] = values + self._state.info["log"] = log + if not self._autoreset_reset_active: + self._state.terminated[ids] = False + self._state.truncated[ids] = False + self.reset_buf[ids] = False + self.reset_terminated[ids] = False + self.reset_time_outs[ids] = False + self.obs_buf = self._state.obs if self._state is not None else mapped_obs + self.extras = self._state.info if self._state is not None else {"log": log} + self.recorder_manager.record_post_reset(ids) + return reset_obs, {"log": log} + + def _normalize_reset_ids( + self, + env_indices: np.ndarray | None, + env_ids: np.ndarray | None, + ) -> np.ndarray: + if env_indices is not None and env_ids is not None: + raise ValueError("Pass either env_indices or env_ids, not both") + values = env_ids if env_ids is not None else env_indices + if values is None: + return np.arange(self.num_envs, dtype=np.int32) + raw = np.asarray(values) + if ( + raw.ndim != 1 + or not np.issubdtype(raw.dtype, np.integer) + or np.issubdtype(raw.dtype, np.bool_) + ): + raise TypeError( + "ManagerBasedRlEnv reset env IDs must be a 1-D integer np.ndarray; " + f"got shape={raw.shape}, dtype={raw.dtype}" + ) + ids = np.asarray(raw, dtype=np.int32) + if np.any(ids < 0) or np.any(ids >= self.num_envs): + raise IndexError( + f"ManagerBasedRlEnv reset env IDs out of range for {self.num_envs} envs: " + f"{ids.tolist()}" + ) + if np.unique(ids).size != ids.size: + raise ValueError(f"ManagerBasedRlEnv reset env IDs contain duplicates: {ids.tolist()}") + return ids + + def _map_observations( + self, + manager_obs: dict[str, np.ndarray | dict[str, np.ndarray]], + ) -> dict[str, np.ndarray]: + mapping = {"obs": self._cfg.policy_observation_group} + if self._cfg.critic_observation_group is not None: + mapping["critic"] = self._cfg.critic_observation_group + mapped: dict[str, np.ndarray] = {} + for output_name, group_name in mapping.items(): + value = manager_obs[group_name] + if not isinstance(value, np.ndarray): + raise TypeError( + f"ManagerBasedRlEnv observation group '{group_name}' returned " + f"{type(value).__name__}, expected np.ndarray" + ) + expected = (self.num_envs, self._mapped_obs_dims[output_name]) + if value.shape != expected: + raise ValueError( + f"ManagerBasedRlEnv observation group '{group_name}' returned shape " + f"{value.shape}, expected {expected} for NpEnvState.obs['{output_name}']" + ) + mapped[output_name] = value + return mapped + + def get_observations(self) -> dict[str, np.ndarray]: + if self._state is None: + return self.init_state().obs + self.obs_buf = self._state.obs + return self.obs_buf + + def seed(self, seed: int = -1) -> int: + if seed == -1: + seed = secrets.randbits(63) + if isinstance(seed, bool) or not isinstance(seed, int) or seed < 0: + raise ValueError(f"ManagerBasedRlEnv seed must be a non-negative integer, got {seed!r}") + replacement = np.random.default_rng(seed) + self.rng.bit_generator.state = replacement.bit_generator.state + self._cfg.seed = seed + return seed + + def close(self) -> None: + self.recorder_manager.close() + if self._uses_pre_step_control: + self._backend.set_pre_step_control(None) + self._uses_pre_step_control = False + super().close() + + +# Isaac Lab capitalization is a spelling-only alias. There is one implementation. +ManagerBasedRLEnv = ManagerBasedRlEnv +ManagerBasedRLEnvCfg = ManagerBasedRlEnvCfg + +__all__ = [ + "ManagerBasedRLEnv", + "ManagerBasedRLEnvCfg", + "ManagerBasedRlEnv", + "ManagerBasedRlEnvCfg", +] diff --git a/src/unilab/managers/_types.py b/src/unilab/managers/_types.py index 4459c33cb..37f89dd89 100644 --- a/src/unilab/managers/_types.py +++ b/src/unilab/managers/_types.py @@ -140,12 +140,20 @@ class ManagerBasedRlEnv(Protocol): their concrete environment type, while the manager core depends only on this seam. """ - num_envs: int - rng: np.random.Generator - scene: ManagerScene - max_episode_length_s: float + @property + def num_envs(self) -> int: ... + + @property + def rng(self) -> np.random.Generator: ... + + @property + def scene(self) -> ManagerScene: ... + + @property + def max_episode_length_s(self) -> float: ... - def __getattr__(self, name: str) -> Any: ... + # Concrete task terms may still type their own richer env subclass. The + # standalone manager core deliberately depends only on the properties above. DebugVisualizer = Any diff --git a/src/unilab/managers/action_manager.py b/src/unilab/managers/action_manager.py index b966446ab..9840edb31 100644 --- a/src/unilab/managers/action_manager.py +++ b/src/unilab/managers/action_manager.py @@ -174,9 +174,12 @@ def process_action(self, action: np.ndarray) -> None: self._action[:] = action # Split the flat action vector and route each slice to its term. idx = 0 - for term in self._terms.values(): + for name, term in self._terms.items(): term_actions = self._action[:, idx : idx + term.action_dim] - term.process_actions(term_actions) + try: + term.process_actions(term_actions) + except (TypeError, ValueError, NotImplementedError) as exc: + raise type(exc)(f"ActionManager term '{name}': {exc}") from exc idx += term.action_dim def apply_action(self) -> None: @@ -185,8 +188,11 @@ def apply_action(self) -> None: Called on every decimation substep (physics step), not just once per policy step. Each term writes its most recently processed targets to the simulation. """ - for term in self._terms.values(): - term.apply_actions() + for name, term in self._terms.items(): + try: + term.apply_actions() + except (TypeError, ValueError, NotImplementedError) as exc: + raise type(exc)(f"ActionManager term '{name}': {exc}") from exc def get_active_iterable_terms(self, env_idx: int) -> Sequence[tuple[str, Sequence[float]]]: terms = [] diff --git a/tests/base/test_entity_facade.py b/tests/base/test_entity_facade.py index a254da4b6..0342964e2 100644 --- a/tests/base/test_entity_facade.py +++ b/tests/base/test_entity_facade.py @@ -170,6 +170,42 @@ def test_scene_entity_cfg_resolves_only_against_cached_names() -> None: assert backend.calls[key] == cold_path_calls[key] +def test_entity_control_write_uses_cached_actuator_columns_and_fails_closed() -> None: + backend = _StrictBackendProfile("mujoco") + control = np.zeros((backend.num_envs, backend.num_actuators), dtype=np.float32) + scene = EntityScene( + {"robot": EntityCfg(actuator_names=("ankle", "hip"))}, + cast(SimBackend, backend), + control, + ) + data = scene["robot"].data + + data.write_ctrl(np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], dtype=np.float32)) + np.testing.assert_array_equal(control[:, 4], [1.0, 3.0, 5.0]) + np.testing.assert_array_equal(control[:, 2], [2.0, 4.0, 6.0]) + data.write_ctrl( + np.array([[9.0, 8.0], [7.0, 6.0]], dtype=np.float32), + env_ids=np.array([2, 0], dtype=np.int32), + ) + np.testing.assert_array_equal(control[:, 4], [7.0, 3.0, 9.0]) + np.testing.assert_array_equal(control[:, 2], [6.0, 4.0, 8.0]) + assert backend.calls["actuator names"] == 1 + + with pytest.raises(ValueError, match="expected shape"): + data.write_ctrl(np.zeros((backend.num_envs, 1), dtype=np.float32)) + with pytest.raises(ValueError, match="NaN or Inf"): + data.write_ctrl(np.full((backend.num_envs, 2), np.nan, dtype=np.float32)) + with pytest.raises(IndexError, match="out of range"): + data.write_ctrl( + np.zeros((1, 2), dtype=np.float32), + env_ids=np.array([backend.num_envs], dtype=np.int32), + ) + + _, read_only_scene = _scene() + with pytest.raises(NotImplementedError, match="actuator control write.*not materialized"): + read_only_scene["robot"].data.write_ctrl(np.zeros((backend.num_envs, 2), dtype=np.float32)) + + @pytest.mark.parametrize( ("ids", "error_type", "message"), [ diff --git a/tests/envs/test_manager_based_rl_env.py b/tests/envs/test_manager_based_rl_env.py new file mode 100644 index 000000000..448163b1a --- /dev/null +++ b/tests/envs/test_manager_based_rl_env.py @@ -0,0 +1,449 @@ +"""Focused contract tests for the NumPy Manager-Based environment lifecycle.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, cast + +import numpy as np +import pytest + +import unilab.envs.manager_based_rl_env as manager_env_module +from unilab.base.backend.base import SimBackend +from unilab.base.entity import EntityCfg +from unilab.base.scene import SceneCfg +from unilab.envs import ( + ManagerBasedRLEnv, + ManagerBasedRlEnv, + ManagerBasedRLEnvCfg, + ManagerBasedRlEnvCfg, +) +from unilab.managers import ( + ActionTerm, + ActionTermCfg, + CommandTerm, + CommandTermCfg, + CurriculumTermCfg, + EventTermCfg, + MetricsTermCfg, + ObservationGroupCfg, + ObservationTermCfg, + RecorderTerm, + RecorderTermCfg, + RewardTermCfg, + TerminationTermCfg, +) + + +class _FakeBackend: + backend_type = "fake" + + def __init__(self, num_envs: int, *, reject_pre_step: bool = False) -> None: + self.num_envs = num_envs + self.num_actuators = 1 + self.reject_pre_step = reject_pre_step + self.pre_step_control = None + self.applied_controls: list[np.ndarray] = [] + self.cleanup_calls = 0 + + def get_actuator_names(self) -> tuple[str, ...]: + return ("motor",) + + def get_actuator_ctrl_range(self) -> np.ndarray: + return np.array([[-2.0, 2.0]], dtype=np.float32) + + def get_dof_pos(self) -> np.ndarray: + return np.empty((self.num_envs, 0), dtype=np.float32) + + def get_dof_vel(self) -> np.ndarray: + return np.empty((self.num_envs, 0), dtype=np.float32) + + def set_pre_step_control(self, fn) -> None: + if fn is not None and self.reject_pre_step: + raise NotImplementedError("host callback disabled") + self.pre_step_control = fn + + def step(self, ctrl: np.ndarray, nsteps: int = 1) -> None: + native = ctrl + for _ in range(nsteps): + if self.pre_step_control is not None: + native = self.pre_step_control(self, ctrl) + self.applied_controls.append(native.copy()) + + def cleanup_scene_assets(self) -> None: + self.cleanup_calls += 1 + + +@dataclass(kw_only=True) +class _DriveCfg(ActionTermCfg): + gain: float = 1.0 + + def build(self, env) -> ActionTerm: + return _DriveAction(self, env) + + +class _DriveAction(ActionTerm): + def __init__(self, cfg: _DriveCfg, env) -> None: + super().__init__(cfg, env) + self._processed = np.zeros((self.num_envs, 1), dtype=np.float32) + + @property + def action_dim(self) -> int: + return 1 + + @property + def raw_action(self) -> np.ndarray: + return self._processed + + def process_actions(self, actions: np.ndarray) -> None: + self._processed[:] = actions * cast(_DriveCfg, self.cfg).gain + + def apply_actions(self) -> None: + self._env.trace.append("action_apply") + self._env.action_sim_steps.append(self._env._sim_step_counter) + self._entity.data.write_ctrl(self._processed) + + +@dataclass(kw_only=True) +class _CommandCfg(CommandTermCfg): + def build(self, env) -> CommandTerm: + return _Command(self, env) + + +class _Command(CommandTerm): + def __init__(self, cfg: _CommandCfg, env) -> None: + super().__init__(cfg, env) + self._command = np.zeros((self.num_envs, 1), dtype=np.float32) + + @property + def command(self) -> np.ndarray: + return self._command + + def _update_metrics(self) -> None: + return None + + def _resample_command(self, env_ids: np.ndarray) -> None: + self._command[env_ids, 0] = self._env.rng.uniform(size=len(env_ids)) + + def _update_command(self, env_ids: np.ndarray | None) -> None: + ids = None if env_ids is None else env_ids.copy() + self._env.command_update_ids.append(ids) + + +class _Recorder(RecorderTerm): + def record_pre_reset(self, env_ids: np.ndarray) -> None: + self._env.trace.append(("pre_reset", env_ids.tolist())) + + def record_post_reset(self, env_ids: np.ndarray) -> None: + self._env.trace.append(("post_reset", env_ids.tolist())) + + def record_post_step(self) -> None: + self._env.trace.append("post_step") + + def close(self) -> None: + self._env.trace.append("recorder_close") + + +class _TestEnv(ManagerBasedRlEnv): + def __init__(self, cfg: ManagerBasedRlEnvCfg, backend: SimBackend, num_envs: int) -> None: + self.trace: list[Any] = [] + self.command_update_ids: list[np.ndarray | None] = [] + self.action_sim_steps: list[int] = [] + super().__init__(cfg, backend, num_envs) + + +def _policy_obs(env: _TestEnv) -> np.ndarray: + return np.column_stack( + (env.episode_length_buf.astype(np.float32), env.action_manager.action[:, 0]) + ) + + +def _critic_obs(env: _TestEnv) -> np.ndarray: + return env.episode_length_buf[:, None].astype(np.float32) + + +def _reward(env: _TestEnv) -> np.ndarray: + return env.action_manager.action[:, 0].copy() + + +def _failure(env: _TestEnv) -> np.ndarray: + return env.action_manager.action[:, 0] > 0.8 + + +def _time_out(env: _TestEnv) -> np.ndarray: + return env.episode_length_buf >= 2 + + +def _metric(env: _TestEnv) -> np.ndarray: + return env.episode_length_buf.astype(np.float32) + + +def _curriculum(env: _TestEnv, env_ids: np.ndarray | slice) -> float: + del env + return float(len(env_ids)) if isinstance(env_ids, np.ndarray) else 0.0 + + +def _event(env: _TestEnv, env_ids: np.ndarray | None) -> None: + rendered_ids = None if env_ids is None else env_ids.tolist() + env.trace.append(("event", rendered_ids)) + + +def _make_cfg( + *, + sim_substeps: int = 2, + finite_horizon: bool = False, + auto_reset: bool = True, + metrics: dict[str, MetricsTermCfg | None] | None = None, + observations: dict[str, ObservationGroupCfg | None] | None = None, + include_optional_managers: bool = True, +) -> ManagerBasedRlEnvCfg: + if observations is None: + observations = { + "actor": ObservationGroupCfg(terms={"policy": ObservationTermCfg(func=_policy_obs)}), + "value": ObservationGroupCfg(terms={"critic": ObservationTermCfg(func=_critic_obs)}), + } + cfg = ManagerBasedRlEnvCfg( + scene=SceneCfg( + model_file="fake.xml", + entities={"robot": EntityCfg(actuator_names=("motor",))}, + ), + sim_dt=0.01, + ctrl_dt=0.01 * sim_substeps, + max_episode_seconds=1.0, + seed=7, + observations=observations, + actions={"drive": _DriveCfg(entity_name="robot")}, + rewards={"track": RewardTermCfg(func=_reward, weight=1.0)}, + terminations={ + "failure": TerminationTermCfg(func=_failure), + "time_out": TerminationTermCfg(func=_time_out, time_out=True), + }, + events={"reset": EventTermCfg(func=_event, mode="reset")}, + metrics={} if metrics is None else metrics, + policy_observation_group="actor", + critic_observation_group="value", + is_finite_horizon=finite_horizon, + auto_reset=auto_reset, + ) + if include_optional_managers: + cfg.commands = {"target": _CommandCfg(resampling_time_range=(1.0, 1.0))} + cfg.curriculum = {"difficulty": CurriculumTermCfg(func=_curriculum)} + cfg.recorders = {"trace": RecorderTermCfg(func=_Recorder)} + return cfg + + +def _make_env( + cfg: ManagerBasedRlEnvCfg | None = None, + *, + num_envs: int = 2, + reject_pre_step: bool = False, +) -> tuple[_TestEnv, _FakeBackend]: + backend = _FakeBackend(num_envs, reject_pre_step=reject_pre_step) + env = _TestEnv( + cfg or _make_cfg(), + cast(SimBackend, backend), + num_envs, + ) + return env, backend + + +def test_public_names_are_spelling_only_aliases() -> None: + assert ManagerBasedRLEnv is ManagerBasedRlEnv + assert ManagerBasedRLEnvCfg is ManagerBasedRlEnvCfg + + +@pytest.mark.parametrize( + ("mutate", "error", "match"), + [ + (lambda cfg: setattr(cfg, "sim_dt", True), TypeError, "sim_dt must be a real"), + ( + lambda cfg: setattr(cfg, "ctrl_dt", 0.015), + ValueError, + "integer multiple", + ), + ( + lambda cfg: setattr(cfg, "max_episode_seconds", None), + ValueError, + "max_episode_seconds", + ), + (lambda cfg: setattr(cfg, "scene", None), TypeError, "scene must be a SceneCfg"), + ], +) +def test_manager_based_config_rejects_invalid_contracts(mutate, error, match: str) -> None: + cfg = _make_cfg() + mutate(cfg) + with pytest.raises(error, match=match): + cfg.validate() + + +def test_manager_construction_uses_pinned_order(monkeypatch: pytest.MonkeyPatch) -> None: + order: list[str] = [] + names = ( + "EventManager", + "CommandManager", + "ActionManager", + "ObservationManager", + "TerminationManager", + "RewardManager", + "CurriculumManager", + "MetricsManager", + "RecorderManager", + ) + cfg = _make_cfg(metrics={"progress": MetricsTermCfg(func=_metric)}) + for name in names: + original = getattr(manager_env_module, name) + + def wrapped(*args, _name=name, _original=original, **kwargs): + order.append(_name) + return _original(*args, **kwargs) + + monkeypatch.setattr(manager_env_module, name, wrapped) + + _make_env(cfg) + assert order == list(names) + + +def test_np_env_owns_substeps_autoreset_and_final_observation() -> None: + env, backend = _make_env() + initial_obs, initial_info = env.reset() + assert env.state is not None + initial = env.state + + assert initial_obs["obs"].shape == (2, 2) + assert initial_obs["critic"].shape == (2, 1) + assert "log" in initial_info + np.testing.assert_array_equal(initial.info["steps"], [0, 0]) + assert env._dr_manager is None + + state = env.step(np.array([[0.25], [0.5]], dtype=np.float32)) + assert len(backend.applied_controls) == 2 + assert env.action_sim_steps == [1, 2] + np.testing.assert_allclose(backend.applied_controls[0][:, 0], [0.25, 0.5]) + np.testing.assert_allclose(backend.applied_controls[1][:, 0], [0.25, 0.5]) + np.testing.assert_allclose(state.reward, [0.005, 0.01]) + np.testing.assert_array_equal(state.terminated, [False, False]) + np.testing.assert_array_equal(state.truncated, [False, False]) + np.testing.assert_array_equal(state.info["steps"], [1, 1]) + + state = env.step(np.array([[0.25], [0.5]], dtype=np.float32)) + np.testing.assert_array_equal(state.truncated, [True, True]) + np.testing.assert_array_equal(state.terminated, [False, False]) + np.testing.assert_array_equal(state.info["steps"], [0, 0]) + np.testing.assert_array_equal(state.obs["obs"][:, 0], [0.0, 0.0]) + assert state.final_observation is not None + np.testing.assert_array_equal(state.final_observation["obs"][:, 0], [2.0, 2.0]) + assert state.info["_final_observation"].tolist() == [True, True] + assert state.info["log"]["Episode_Termination/time_out"] == 2 + + pre_index = env.trace.index(("pre_reset", [0, 1])) + post_reset_index = env.trace.index(("post_reset", [0, 1]), pre_index) + post_step_index = len(env.trace) - 1 + assert env.trace[post_step_index] == "post_step" + assert pre_index < post_reset_index < post_step_index + + +def test_partial_reset_preserves_other_env_counter_and_terminal_obs() -> None: + env, _ = _make_env() + env.init_state() + + state = env.step(np.array([[1.0], [0.0]], dtype=np.float32)) + + np.testing.assert_array_equal(state.terminated, [True, False]) + np.testing.assert_array_equal(state.truncated, [False, False]) + np.testing.assert_array_equal(state.info["steps"], [0, 1]) + np.testing.assert_array_equal(state.obs["obs"][:, 0], [0.0, 1.0]) + assert state.final_observation is not None + np.testing.assert_array_equal(state.final_observation["obs"][0], [1.0, 1.0]) + assert state.info["_final_observation"].tolist() == [True, False] + assert ("event", [0]) in env.trace + + +def test_finite_horizon_maps_time_out_to_terminated() -> None: + env, _ = _make_env(_make_cfg(finite_horizon=True)) + env.init_state() + env.step(np.zeros((2, 1), dtype=np.float32)) + state = env.step(np.zeros((2, 1), dtype=np.float32)) + np.testing.assert_array_equal(state.terminated, [True, True]) + np.testing.assert_array_equal(state.truncated, [False, False]) + + +def test_manual_reset_is_required_when_autoreset_is_disabled() -> None: + env, _ = _make_env(_make_cfg(auto_reset=False)) + env.init_state() + state = env.step(np.array([[1.0], [0.0]], dtype=np.float32)) + np.testing.assert_array_equal(state.info["steps"], [1, 1]) + with pytest.raises(RuntimeError, match="must be reset"): + env.step(np.zeros((2, 1), dtype=np.float32)) + + reset_obs, info = env.reset(env_ids=np.array([0], dtype=np.int32)) + assert reset_obs["obs"].shape == (1, 2) + assert "log" in info + assert not state.terminated[0] + env.step(np.zeros((2, 1), dtype=np.float32)) + + +def test_reset_seed_updates_the_shared_generator_in_place() -> None: + cfg = _make_cfg(include_optional_managers=False) + cfg.events = {} + env, _ = _make_env(cfg) + env.init_state() + generator = env.rng + + env.reset(seed=123, env_ids=np.array([0], dtype=np.int32)) + + assert env.rng is generator + expected = np.random.default_rng(123).random() + assert env.rng.random() == pytest.approx(expected) + + +def test_per_substep_metrics_fail_without_post_substep_backend_hook() -> None: + cfg = _make_cfg( + sim_substeps=2, + metrics={"energy": MetricsTermCfg(func=_metric, per_substep=True)}, + ) + with pytest.raises(NotImplementedError, match="post-physics per-substep metrics.*fake"): + _make_env(cfg) + + +def test_multisubstep_action_fails_when_backend_rejects_callback() -> None: + with pytest.raises(NotImplementedError, match="ActionManager.*every physics substep.*fake"): + _make_env(reject_pre_step=True) + + +@pytest.mark.parametrize( + ("mutate", "error", "match"), + [ + ( + lambda cfg: setattr(cfg, "policy_observation_group", "missing"), + KeyError, + "requests group 'missing'", + ), + ( + lambda cfg: setattr( + cfg.observations["actor"], + "concatenate_terms", + False, # type: ignore[union-attr] + ), + ValueError, + "must concatenate terms", + ), + ( + lambda cfg: setattr(cfg, "critic_observation_group", "actor"), + ValueError, + "must be different", + ), + ], +) +def test_observation_mapping_fails_closed(mutate, error, match: str) -> None: + cfg = _make_cfg() + mutate(cfg) + with pytest.raises(error, match=match): + _make_env(cfg) + + +def test_close_unhooks_callback_and_closes_owned_resources() -> None: + env, backend = _make_env() + env.close() + assert backend.pre_step_control is None + assert backend.cleanup_calls == 1 + assert env.trace[-1] == "recorder_close" diff --git a/tests/managers/test_core_managers.py b/tests/managers/test_core_managers.py index 18aa53f64..653362c16 100644 --- a/tests/managers/test_core_managers.py +++ b/tests/managers/test_core_managers.py @@ -96,6 +96,32 @@ def test_action_rejects_invalid_input(fake_env: FakeEnv, action: np.ndarray, mat manager.process_action(action) +class FailingAction(DummyAction): + def process_actions(self, actions: np.ndarray) -> None: + del actions + raise ValueError("invalid processed target") + + def apply_actions(self) -> None: + raise NotImplementedError("backend control write unavailable") + + +@dataclass(kw_only=True) +class FailingActionCfg(DummyActionCfg): + def build(self, env: FakeEnv) -> FailingAction: + return FailingAction(self, env) + + +def test_action_term_errors_include_manager_and_term_context(fake_env: FakeEnv) -> None: + manager = ActionManager( + {"broken": FailingActionCfg(entity_name="robot", dim=1)}, + fake_env, + ) + with pytest.raises(ValueError, match="ActionManager term 'broken'.*invalid processed"): + manager.process_action(np.zeros((fake_env.num_envs, 1), dtype=np.float32)) + with pytest.raises(NotImplementedError, match="ActionManager term 'broken'.*control write"): + manager.apply_action() + + class StatefulReward: def __init__(self, cfg: RewardTermCfg, env: FakeEnv): self.reset_ids = None From f8ebee0bd9f44bb2da6e87eb1f491c93a762c5ee Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Mon, 17 Aug 2026 20:30:59 +0800 Subject: [PATCH 005/169] feat(backend): expose actuation metadata (#1053) --- src/unilab/base/backend/base.py | 23 ++++++++++ src/unilab/base/backend/drake/backend.py | 49 +++++++++++++++++++++ src/unilab/base/backend/mjwarp/backend.py | 38 ++++++++++++++++ src/unilab/base/backend/motrix/backend.py | 48 ++++++++++++++++++++ src/unilab/base/backend/mujoco/backend.py | 38 ++++++++++++++++ tests/base/backend/test_drake_batch_pool.py | 30 +++++++++++++ tests/base/test_backend_conformance.py | 39 ++++++++++++++++ tests/base/test_mjwarp_identity.py | 27 ++++++++++++ 8 files changed, 292 insertions(+) diff --git a/src/unilab/base/backend/base.py b/src/unilab/base/backend/base.py index 3989c35d2..32b4a612f 100644 --- a/src/unilab/base/backend/base.py +++ b/src/unilab/base/backend/base.py @@ -137,6 +137,18 @@ def get_actuator_names(self) -> tuple[str, ...]: """Return actuator names in control-vector order on the cold path.""" raise NotImplementedError(f"{self.__class__.__name__} does not expose actuator names") + def get_actuator_joint_names(self) -> tuple[str, ...]: + """Return each actuator's target single-DoF joint in control-vector order. + + Backends must fail closed when an actuator does not target exactly one + hinge/slide joint. Manager action terms use this cold-path metadata to + map community joint selectors onto the backend control vector without + inspecting backend-private model objects. + """ + raise NotImplementedError( + f"{self.__class__.__name__} does not expose actuator target joints" + ) + def get_scene_model_file(self) -> str | None: """Return the materialized scene path for diagnostics, when available.""" return None @@ -172,6 +184,17 @@ def get_default_qpos(self) -> np.ndarray: """Return the backend/model default qpos through a stable contract.""" raise NotImplementedError(f"{self.__class__.__name__} does not expose default qpos") + def get_default_dof_pos(self) -> np.ndarray: + """Return default joint positions in the same column order as ``get_dof_pos``. + + The returned array is detached, one-dimensional, and excludes floating + root coordinates. Backends whose DoF view is actuator-indexed must use + that same actuator-target order here. + """ + raise NotImplementedError( + f"{self.__class__.__name__} does not expose default DoF positions" + ) + @abc.abstractmethod def get_init_qvel(self) -> np.ndarray: """Return a zero-initialized qvel vector compatible with ``set_state``. diff --git a/src/unilab/base/backend/drake/backend.py b/src/unilab/base/backend/drake/backend.py index b753eed46..dd29977e3 100644 --- a/src/unilab/base/backend/drake/backend.py +++ b/src/unilab/base/backend/drake/backend.py @@ -222,6 +222,10 @@ def __init__( self._actuator_damping = model_info.actuator_damping.copy() self._actuator_qpos_adr = model_info.actuator_qpos_adr.astype(np.intp, copy=True) self._actuator_qvel_adr = model_info.actuator_qvel_adr.astype(np.intp, copy=True) + raw_actuator_names = getattr(model_info, "actuator_names", None) + self._actuator_names = ( + None if raw_actuator_names is None else tuple(str(name) for name in raw_actuator_names) + ) self._sensor_names = tuple(model_info.sensor_names) self._sensor_adr = model_info.sensor_adr.copy() self._sensor_dim = model_info.sensor_dim.copy() @@ -253,6 +257,19 @@ def __init__( strict=True, ) } + joint_name_by_qpos_adr = { + int(adr): str(name) + for name, adr, dim in zip( + getattr(model_info, "joint_names", ()), + getattr(model_info, "joint_qpos_adr", ()), + getattr(model_info, "joint_qpos_dim", ()), + strict=True, + ) + if int(dim) == 1 + } + self._actuator_joint_names = tuple( + joint_name_by_qpos_adr.get(int(adr), "") for adr in self._actuator_qpos_adr + ) self._root_qpos_dim = ( int(np.min(self._actuator_qpos_adr)) if self._actuator_qpos_adr.size else 0 ) @@ -311,6 +328,35 @@ def num_dof_vel(self) -> int: def get_actuator_ctrl_range(self) -> np.ndarray: return self._ctrl_limits.copy() + def get_actuator_names(self) -> tuple[str, ...]: + names = self._actuator_names + if names is None: + raise NotImplementedError( + "backend 'drake' capability 'actuator names' is unavailable: " + "DrakeUni model_info does not expose actuator_names" + ) + if len(names) != self.num_actuators or any(not name for name in names): + raise NotImplementedError( + "backend 'drake' capability 'actuator names' requires one non-empty name " + f"per control column; received {names}" + ) + if len(set(names)) != len(names): + raise NotImplementedError( + "backend 'drake' capability 'actuator names' requires unique names; " + f"received {names}" + ) + return names + + def get_actuator_joint_names(self) -> tuple[str, ...]: + names = self._actuator_joint_names + if len(names) != self.num_actuators or any(not name for name in names): + raise NotImplementedError( + "backend 'drake' capability 'actuator target joint' requires every " + "actuator_qpos_adr to resolve to one named single-DoF joint; " + f"received {names}" + ) + return names + def get_scene_model_file(self) -> str | None: return self._scene_model_file @@ -325,6 +371,9 @@ def get_keyframe_qpos(self, name: str) -> np.ndarray: def get_default_qpos(self) -> np.ndarray: return self._home_qpos_mujoco.copy() + def get_default_dof_pos(self) -> np.ndarray: + return np.asarray(self._home_qpos_mujoco[self._actuator_qpos_adr], dtype=np.float64).copy() + def get_init_qvel(self) -> np.ndarray: return self._home_qvel_mujoco.copy() diff --git a/src/unilab/base/backend/mjwarp/backend.py b/src/unilab/base/backend/mjwarp/backend.py index de84edf78..04e473b15 100644 --- a/src/unilab/base/backend/mjwarp/backend.py +++ b/src/unilab/base/backend/mjwarp/backend.py @@ -292,6 +292,41 @@ def get_actuator_ctrl_range(self) -> np.ndarray: def get_actuator_names(self) -> tuple[str, ...]: return self._actuator_names + def get_actuator_joint_names(self) -> tuple[str, ...]: + supported_transmissions = { + int(self._mujoco.mjtTrn.mjTRN_JOINT), + int(self._mujoco.mjtTrn.mjTRN_JOINTINPARENT), + } + supported_joint_types = { + int(self._mujoco.mjtJoint.mjJNT_HINGE), + int(self._mujoco.mjtJoint.mjJNT_SLIDE), + } + names: list[str] = [] + for actuator_id, actuator_name in enumerate(self._actuator_names): + transmission = int(self._cpu_model.actuator_trntype[actuator_id]) + joint_id = int(self._cpu_model.actuator_trnid[actuator_id, 0]) + if transmission not in supported_transmissions or joint_id < 0: + raise NotImplementedError( + "backend 'mjwarp' capability 'actuator target joint' requires a " + f"joint transmission; actuator '{actuator_name}' uses " + f"transmission type {transmission}" + ) + if int(self._cpu_model.jnt_type[joint_id]) not in supported_joint_types: + raise NotImplementedError( + "backend 'mjwarp' capability 'actuator target joint' requires a " + f"single-DoF joint; actuator '{actuator_name}' targets joint id {joint_id}" + ) + joint_name = self._mujoco.mj_id2name( + self._cpu_model, self._mujoco.mjtObj.mjOBJ_JOINT, joint_id + ) + if not joint_name: + raise NotImplementedError( + "backend 'mjwarp' capability 'actuator target joint' requires named " + f"joints; actuator '{actuator_name}' targets unnamed joint id {joint_id}" + ) + names.append(str(joint_name)) + return tuple(names) + def get_scene_model_file(self) -> str | None: return self.scene_model_file @@ -305,6 +340,9 @@ def get_keyframe_qpos(self, name: str) -> np.ndarray: def get_default_qpos(self) -> np.ndarray: return np.asarray(self._cpu_model.qpos0, dtype=np.float32).copy() + def get_default_dof_pos(self) -> np.ndarray: + return np.asarray(self._cpu_model.qpos0[self._root_qpos_dim :], dtype=np.float32).copy() + def get_init_qvel(self) -> np.ndarray: return np.zeros((self._nv,), dtype=np.float32) diff --git a/src/unilab/base/backend/motrix/backend.py b/src/unilab/base/backend/motrix/backend.py index 2d6ae5090..2a25ff35f 100644 --- a/src/unilab/base/backend/motrix/backend.py +++ b/src/unilab/base/backend/motrix/backend.py @@ -383,6 +383,45 @@ def get_actuator_ctrl_range(self) -> np.ndarray: result: np.ndarray = arr.T.copy() return result + def get_actuator_names(self) -> tuple[str, ...]: + actuators = sorted(self._model.actuators, key=lambda actuator: int(actuator.index)) + names = tuple(str(actuator.name) for actuator in actuators) + if len(names) != self.num_actuators or any(not name for name in names): + raise NotImplementedError( + "backend 'motrix' capability 'actuator names' requires one non-empty name " + f"per control column; received {names}" + ) + if len(set(names)) != len(names): + raise NotImplementedError( + "backend 'motrix' capability 'actuator names' requires unique names; " + f"received {names}" + ) + return names + + def get_actuator_joint_names(self) -> tuple[str, ...]: + actuators = sorted(self._model.actuators, key=lambda actuator: int(actuator.index)) + names: list[str] = [] + for actuator in actuators: + if actuator.target_type != "joint": + raise NotImplementedError( + "backend 'motrix' capability 'actuator target joint' requires a joint " + f"transmission; actuator '{actuator.name}' targets '{actuator.target_type}'" + ) + joint = self._model.get_joint(actuator.target_name) + if joint is None or int(joint.num_dof_pos) != 1 or int(joint.num_dof_vel) != 1: + raise NotImplementedError( + "backend 'motrix' capability 'actuator target joint' requires a " + f"single-DoF joint; actuator '{actuator.name}' targets " + f"'{actuator.target_name}'" + ) + names.append(str(actuator.target_name)) + if len(names) != self.num_actuators: + raise NotImplementedError( + "backend 'motrix' capability 'actuator target joint' returned " + f"{len(names)} targets for {self.num_actuators} actuators" + ) + return tuple(names) + def get_terrain_spawn_data(self) -> BackendTerrainSpawnData | None: return self._terrain_spawn_data @@ -397,6 +436,15 @@ def get_default_qpos(self) -> np.ndarray: qpos = np.array(self._model.compute_init_dof_pos(), dtype=self._np_dtype) return self._motrix_qpos_to_mujoco(qpos) + def get_default_dof_pos(self) -> np.ndarray: + qpos = np.asarray(self._model.compute_init_dof_pos(), dtype=self._np_dtype) + indices = ( + self._actuator_joint_pos_indices + if self._actuator_joint_pos_indices is not None + else self._joint_dof_pos_indices + ) + return np.asarray(qpos[indices], dtype=self._np_dtype).copy() + def get_init_qvel(self) -> np.ndarray: return np.zeros((self._model.num_dof_vel,), dtype=self._np_dtype) diff --git a/src/unilab/base/backend/mujoco/backend.py b/src/unilab/base/backend/mujoco/backend.py index 38c66837c..0de0e3bee 100644 --- a/src/unilab/base/backend/mujoco/backend.py +++ b/src/unilab/base/backend/mujoco/backend.py @@ -692,6 +692,41 @@ def get_actuator_names(self) -> tuple[str, ...]: for actuator_id in range(int(self._model.nu)) ) + def get_actuator_joint_names(self) -> tuple[str, ...]: + supported_transmissions = { + int(mujoco.mjtTrn.mjTRN_JOINT), + int(mujoco.mjtTrn.mjTRN_JOINTINPARENT), + } + supported_joint_types = { + int(mujoco.mjtJoint.mjJNT_HINGE), + int(mujoco.mjtJoint.mjJNT_SLIDE), + } + names: list[str] = [] + for actuator_id in range(int(self._model.nu)): + transmission = int(self._model.actuator_trntype[actuator_id]) + joint_id = int(self._model.actuator_trnid[actuator_id, 0]) + if transmission not in supported_transmissions or joint_id < 0: + actuator_name = self.get_actuator_names()[actuator_id] + raise NotImplementedError( + "backend 'mujoco' capability 'actuator target joint' requires " + f"a joint transmission; actuator '{actuator_name}' uses " + f"transmission type {transmission}" + ) + if int(self._model.jnt_type[joint_id]) not in supported_joint_types: + actuator_name = self.get_actuator_names()[actuator_id] + raise NotImplementedError( + "backend 'mujoco' capability 'actuator target joint' requires " + f"a single-DoF joint; actuator '{actuator_name}' targets joint id {joint_id}" + ) + joint_name = mujoco.mj_id2name(self._model, mujoco.mjtObj.mjOBJ_JOINT, joint_id) + if not joint_name: + raise NotImplementedError( + "backend 'mujoco' capability 'actuator target joint' requires named joints; " + f"actuator id {actuator_id} targets unnamed joint id {joint_id}" + ) + names.append(joint_name) + return tuple(names) + def get_scene_model_file(self) -> str | None: return str(self.scene_model_file) if self.scene_model_file else None @@ -710,6 +745,9 @@ def get_keyframe_qpos(self, name: str) -> np.ndarray: def get_default_qpos(self) -> np.ndarray: return np.asarray(self._model.qpos0, dtype=np.float64).copy() + def get_default_dof_pos(self) -> np.ndarray: + return np.asarray(self._model.qpos0[self._root_qpos_dim :], dtype=self._np_dtype).copy() + def get_init_qvel(self) -> np.ndarray: return np.zeros((self.nv,), dtype=self._np_dtype) diff --git a/tests/base/backend/test_drake_batch_pool.py b/tests/base/backend/test_drake_batch_pool.py index 456a7ada0..aa39f80de 100644 --- a/tests/base/backend/test_drake_batch_pool.py +++ b/tests/base/backend/test_drake_batch_pool.py @@ -5,7 +5,9 @@ import subprocess import sys import textwrap +from types import SimpleNamespace +import numpy as np import pytest @@ -150,6 +152,34 @@ def test_drake_batch_thread_policy_matches_mujoco_auto(monkeypatch: pytest.Monke assert backend._resolve_batch_nthread(2, 8) == 2 +def test_drake_actuation_metadata_is_detached_and_fails_closed() -> None: + from unilab.base.backend.drake.backend import DrakeBackend + + backend = object.__new__(DrakeBackend) + backend._model = SimpleNamespace(nu=2) + backend._actuator_names = ("hip_motor", "knee_motor") + backend._actuator_joint_names = ("hip", "knee") + backend._actuator_qpos_adr = np.asarray([7, 8], dtype=np.intp) + backend._home_qpos_mujoco = np.arange(9, dtype=np.float64) + + assert backend.get_actuator_names() == ("hip_motor", "knee_motor") + assert backend.get_actuator_joint_names() == ("hip", "knee") + default = backend.get_default_dof_pos() + np.testing.assert_array_equal(default, np.asarray([7.0, 8.0])) + default[:] = -1.0 + np.testing.assert_array_equal(backend.get_default_dof_pos(), np.asarray([7.0, 8.0])) + + backend._actuator_names = None + with pytest.raises(NotImplementedError, match="DrakeUni model_info.*actuator_names"): + backend.get_actuator_names() + backend._actuator_names = ("duplicate", "duplicate") + with pytest.raises(NotImplementedError, match="unique names"): + backend.get_actuator_names() + backend._actuator_joint_names = ("hip", "") + with pytest.raises(NotImplementedError, match="single-DoF joint"): + backend.get_actuator_joint_names() + + def test_batch_backend_mode_rejects_existing_pydrake_module() -> None: output = _run_clean_python( """ diff --git a/tests/base/test_backend_conformance.py b/tests/base/test_backend_conformance.py index 6e8578e40..bee1f09e8 100644 --- a/tests/base/test_backend_conformance.py +++ b/tests/base/test_backend_conformance.py @@ -130,6 +130,45 @@ def sample_height(xy: np.ndarray) -> np.ndarray: ) +def test_actuation_metadata_defaults_fail_closed() -> None: + with pytest.raises(NotImplementedError, match="actuator target joints"): + SimBackend.get_actuator_joint_names(object()) # type: ignore[arg-type] + with pytest.raises(NotImplementedError, match="default DoF positions"): + SimBackend.get_default_dof_pos(object()) # type: ignore[arg-type] + + +@pytest.mark.parametrize("backend_type", _BACKEND_PARAMS) +def test_actuation_metadata_contract(backend_type: str) -> None: + _require_backend(backend_type) + + backend = create_backend( + backend_type, + SceneCfg(model_file=_G1_SCENE), + NUM_ENVS, + SIM_DT, + base_name="pelvis", + ) + backend.materialize() + + actuator_names = backend.get_actuator_names() + target_joint_names = backend.get_actuator_joint_names() + default_dof_pos = backend.get_default_dof_pos() + + assert len(actuator_names) == backend.num_actuators + assert len(set(actuator_names)) == len(actuator_names) + assert all(actuator_names) + assert len(target_joint_names) == backend.num_actuators + assert all(target_joint_names) + assert default_dof_pos.shape == backend.get_dof_pos().shape[1:] + assert np.issubdtype(default_dof_pos.dtype, np.floating) + assert np.isfinite(default_dof_pos).all() + np.testing.assert_allclose(default_dof_pos, backend.get_dof_pos()[0], atol=1e-6) + + detached = default_dof_pos.copy() + default_dof_pos[:] = np.nan + np.testing.assert_array_equal(backend.get_default_dof_pos(), detached) + + def test_terrain_spawn_consumers_do_not_probe_private_backend_capabilities() -> None: forbidden_names = {"terrain_origins", "terrain_surface_sampler", "sample_height"} offenders: list[str] = [] diff --git a/tests/base/test_mjwarp_identity.py b/tests/base/test_mjwarp_identity.py index 3fe150bc9..eed646cfa 100644 --- a/tests/base/test_mjwarp_identity.py +++ b/tests/base/test_mjwarp_identity.py @@ -9,6 +9,7 @@ from importlib.machinery import ModuleSpec from pathlib import Path +import numpy as np import pytest from unilab.base.backend.mjwarp import dependencies @@ -140,3 +141,29 @@ def test_mjwarp_getters_do_not_materialize_warp_arrays() -> None: ): offenders.append(getter.name) assert offenders == [] + + +def test_mjwarp_actuation_metadata_uses_cpu_model_only() -> None: + import mujoco + + from unilab.assets import ASSETS_ROOT_PATH + from unilab.base.backend.mjwarp.backend import MjwarpBackend + + model = mujoco.MjModel.from_xml_path( + str(ASSETS_ROOT_PATH / "robots" / "go2" / "scene_flat.xml") + ) + backend = object.__new__(MjwarpBackend) + backend._mujoco = mujoco + backend._cpu_model = model + backend._root_qpos_dim = 7 + backend._actuator_names = tuple( + mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_ACTUATOR, index) or f"#{index}" + for index in range(model.nu) + ) + + assert len(backend.get_actuator_joint_names()) == model.nu + assert backend.get_actuator_joint_names()[0] == "FR_hip_joint" + default = backend.get_default_dof_pos() + np.testing.assert_array_equal(default, model.qpos0[7:]) + default[:] = np.nan + assert np.isfinite(backend.get_default_dof_pos()).all() From 6ab34eaccb4d1741f6acd0295fe9a07f711fdf08 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Mon, 17 Aug 2026 20:53:01 +0800 Subject: [PATCH 006/169] feat(env): add NumPy joint position action (#1055) --- src/unilab/base/entity.py | 185 +++++++++++++- src/unilab/envs/mdp/__init__.py | 6 + src/unilab/envs/mdp/actions/__init__.py | 6 + src/unilab/envs/mdp/actions/actions.py | 223 +++++++++++++++++ tests/base/test_entity_facade.py | 83 +++++++ tests/envs/mdp/test_joint_position_action.py | 246 +++++++++++++++++++ 6 files changed, 746 insertions(+), 3 deletions(-) create mode 100644 src/unilab/envs/mdp/__init__.py create mode 100644 src/unilab/envs/mdp/actions/__init__.py create mode 100644 src/unilab/envs/mdp/actions/actions.py create mode 100644 tests/envs/mdp/test_joint_position_action.py diff --git a/src/unilab/base/entity.py b/src/unilab/base/entity.py index 771a5f2fe..678629048 100644 --- a/src/unilab/base/entity.py +++ b/src/unilab/base/entity.py @@ -139,6 +139,7 @@ def __init__( root_body_ids: np.ndarray | None, joint_pos_ids: np.ndarray | None, joint_vel_ids: np.ndarray | None, + default_joint_pos: np.ndarray | None, body_ids: np.ndarray | None, actuator_ids: np.ndarray | None, actuator_ctrl_range: np.ndarray | None, @@ -152,8 +153,14 @@ def __init__( self._root_body_ids = root_body_ids self._joint_pos_index = None if joint_pos_ids is None else _as_column_index(joint_pos_ids) self._joint_vel_index = None if joint_vel_ids is None else _as_column_index(joint_vel_ids) + self._default_joint_pos = default_joint_pos + self._encoder_bias = ( + None + if default_joint_pos is None + else np.zeros(default_joint_pos.shape, dtype=default_joint_pos.dtype) + ) self._body_ids = body_ids - self._actuator_index = None if actuator_ids is None else _as_column_index(actuator_ids) + self._actuator_ids = actuator_ids self._actuator_ctrl_range = actuator_ctrl_range self._control_buffer = control_buffer @@ -203,6 +210,16 @@ def joint_vel(self) -> np.ndarray: index = self._require(self._joint_vel_index, "joint velocity") return self._backend.get_dof_vel()[:, index] + @property + def default_joint_pos(self) -> np.ndarray: + """Read-only per-environment default joint positions.""" + return self._require(self._default_joint_pos, "default joint position") + + @property + def encoder_bias(self) -> np.ndarray: + """Mutable per-environment joint encoder bias used by position actions.""" + return self._require(self._encoder_bias, "joint encoder bias") + @property def body_link_pos_w(self) -> np.ndarray: ids = self._require(self._body_ids, "body state") @@ -239,6 +256,8 @@ def write_ctrl( self, values: np.ndarray, env_ids: np.ndarray | slice | None = None, + *, + actuator_ids: np.ndarray | Sequence[int] | slice | None = None, ) -> None: """Write entity-local actuator controls into the env-owned control buffer. @@ -246,7 +265,7 @@ def write_ctrl( entity target buffers. Physics remains owned by ``NpEnv``/``SimBackend``; this method never steps or calls a backend-private API. """ - actuator_index = self._require(self._actuator_index, "actuator control write") + entity_actuator_ids = self._require(self._actuator_ids, "actuator control write") control = self._require(self._control_buffer, "actuator control write") if not isinstance(values, np.ndarray): raise TypeError( @@ -284,7 +303,38 @@ def write_ctrl( ) row_count = len(row_index) - actuator_count = len(self._require(self._actuator_ctrl_range, "actuator control write")) + if actuator_ids is None: + selected_actuator_ids = entity_actuator_ids + elif isinstance(actuator_ids, slice): + selected_actuator_ids = entity_actuator_ids[actuator_ids] + else: + raw_actuator_ids = np.asarray(actuator_ids) + if ( + raw_actuator_ids.ndim != 1 + or not np.issubdtype(raw_actuator_ids.dtype, np.integer) + or np.issubdtype(raw_actuator_ids.dtype, np.bool_) + ): + raise TypeError( + f"Entity '{self._entity_name}' write_ctrl actuator_ids must be a 1-D " + "integer array or slice" + ) + local_actuator_ids = np.asarray(raw_actuator_ids, dtype=np.intp) + if np.any(local_actuator_ids < 0) or np.any( + local_actuator_ids >= len(entity_actuator_ids) + ): + raise IndexError( + f"Entity '{self._entity_name}' write_ctrl actuator_ids out of range for " + f"{len(entity_actuator_ids)} entity actuators: {local_actuator_ids.tolist()}" + ) + if np.unique(local_actuator_ids).size != local_actuator_ids.size: + raise ValueError( + f"Entity '{self._entity_name}' write_ctrl actuator_ids contain duplicates: " + f"{local_actuator_ids.tolist()}" + ) + selected_actuator_ids = entity_actuator_ids[local_actuator_ids] + + actuator_index = _as_column_index(np.asarray(selected_actuator_ids, dtype=np.int32)) + actuator_count = len(selected_actuator_ids) expected = (row_count, actuator_count) if values.shape != expected: raise ValueError( @@ -366,7 +416,12 @@ def __init__( self._validate_joint_state(backend, joint_pos_ids, joint_vel_ids) self._validate_body_state(backend, root_body_ids, body_ids) + default_joint_pos = self._materialize_default_joint_pos(backend, joint_pos_ids) actuator_ctrl_range = self._materialize_actuator_ctrl_range(backend, actuator_ids) + ( + self._actuator_target_joint_names, + self._joint_to_actuator_local, + ) = self._materialize_joint_actuator_mapping(backend, actuator_ids) if control_buffer is not None: expected_control_shape = (backend.num_envs, backend.num_actuators) if control_buffer.shape != expected_control_shape: @@ -385,6 +440,7 @@ def __init__( root_body_ids=root_body_ids, joint_pos_ids=joint_pos_ids, joint_vel_ids=joint_vel_ids, + default_joint_pos=default_joint_pos, body_ids=body_ids, actuator_ids=actuator_ids, actuator_ctrl_range=actuator_ctrl_range, @@ -519,6 +575,63 @@ def _materialize_actuator_ctrl_range( selected.setflags(write=False) return selected + def _materialize_default_joint_pos( + self, backend: SimBackend, joint_pos_ids: np.ndarray | None + ) -> np.ndarray | None: + if joint_pos_ids is None: + return None + defaults = self._read_state("default joint position", backend.get_default_dof_pos) + current = self._read_state("joint position state", backend.get_dof_pos) + if defaults.shape != current.shape[1:]: + raise ValueError( + f"Entity '{self.name}' capability 'default joint position' on backend " + f"'{self._backend_type}' returned shape {defaults.shape}; expected " + f"{current.shape[1:]} to match get_dof_pos()" + ) + selected = np.asarray(defaults[_as_column_index(joint_pos_ids)]) + materialized = np.broadcast_to(selected, (backend.num_envs, len(joint_pos_ids))).copy() + materialized.setflags(write=False) + return materialized + + def _materialize_joint_actuator_mapping( + self, backend: SimBackend, actuator_ids: np.ndarray | None + ) -> tuple[tuple[str, ...] | None, np.ndarray | None]: + if actuator_ids is None or self._joint_names is None: + return None, None + try: + all_target_names = tuple(backend.get_actuator_joint_names()) + except NotImplementedError as exc: + raise self._capability_error("actuator target joint", str(exc)) from exc + if len(all_target_names) != backend.num_actuators: + raise ValueError( + f"Entity '{self.name}' capability 'actuator target joint' on backend " + f"'{self._backend_type}' returned {len(all_target_names)} names for " + f"{backend.num_actuators} actuators" + ) + target_names = tuple(all_target_names[int(index)] for index in actuator_ids) + if any(not isinstance(name, str) or not name for name in target_names): + raise ValueError( + f"Entity '{self.name}' actuator target joint names must be non-empty strings; " + f"got {target_names}" + ) + if len(set(target_names)) != len(target_names): + raise ValueError( + f"Entity '{self.name}' actuator target joints must be unique for position " + f"control; got {target_names}" + ) + joint_index_by_name = {name: index for index, name in enumerate(self._joint_names)} + missing = [name for name in target_names if name not in joint_index_by_name] + if missing: + raise ValueError( + f"Entity '{self.name}' actuators target joints outside its declared joint " + f"partition on backend '{self._backend_type}': {missing}" + ) + joint_to_actuator = np.full(len(self._joint_names), -1, dtype=np.int32) + for actuator_local_id, joint_name in enumerate(target_names): + joint_to_actuator[joint_index_by_name[joint_name]] = actuator_local_id + joint_to_actuator.setflags(write=False) + return target_names, joint_to_actuator + def _require_names(self, kind: str, names: tuple[str, ...] | None) -> tuple[str, ...]: if names is None: raise self._capability_error(kind, "the namespace was not declared in EntityCfg") @@ -629,6 +742,72 @@ def find_joints( ) -> tuple[list[int], list[str]]: return self._find("joint", self._joint_names, keys, preserve_order) + def find_joints_by_actuator_names( + self, keys: str | Sequence[str] + ) -> tuple[list[int], list[str]]: + """Resolve actuator-target joint patterns in natural entity joint order.""" + target_names = self._actuator_target_joint_names + if target_names is None: + raise self._capability_error( + "actuator target joint", + "joint_names and actuator_names must both be declared in EntityCfg", + ) + target_set = set(target_names) + natural_ids = [index for index, name in enumerate(self.joint_names) if name in target_set] + natural_names = [self.joint_names[index] for index in natural_ids] + matched_ids, matched_names = _resolve_matching_names(keys, natural_names, False) + return [natural_ids[index] for index in matched_ids], matched_names + + def set_joint_position_target( + self, + target: np.ndarray, + joint_ids: np.ndarray | Sequence[int] | slice | None = None, + env_ids: np.ndarray | slice | None = None, + ) -> None: + """Map entity-local joint targets to the env-owned actuator control buffer.""" + joint_to_actuator = self._joint_to_actuator_local + if joint_to_actuator is None: + raise self._capability_error( + "joint position target", + "joint-to-actuator metadata was not materialized", + ) + if joint_ids is None: + local_joint_ids = np.arange(self.num_joints, dtype=np.intp) + elif isinstance(joint_ids, slice): + local_joint_ids = np.arange(self.num_joints, dtype=np.intp)[joint_ids] + else: + raw_joint_ids = np.asarray(joint_ids) + if ( + raw_joint_ids.ndim != 1 + or not np.issubdtype(raw_joint_ids.dtype, np.integer) + or np.issubdtype(raw_joint_ids.dtype, np.bool_) + ): + raise TypeError( + f"Entity '{self.name}' joint position target joint_ids must be a 1-D " + "integer array or slice" + ) + local_joint_ids = np.asarray(raw_joint_ids, dtype=np.intp) + if np.any(local_joint_ids < 0) or np.any(local_joint_ids >= self.num_joints): + raise IndexError( + f"Entity '{self.name}' joint position target joint_ids out of range for " + f"{self.num_joints} joints: {local_joint_ids.tolist()}" + ) + if np.unique(local_joint_ids).size != local_joint_ids.size: + raise ValueError( + f"Entity '{self.name}' joint position target joint_ids contain duplicates: " + f"{local_joint_ids.tolist()}" + ) + actuator_ids = joint_to_actuator[local_joint_ids] + if np.any(actuator_ids < 0): + passive_names = [ + self.joint_names[int(index)] for index in local_joint_ids[actuator_ids < 0] + ] + raise NotImplementedError( + f"Entity '{self.name}' capability 'joint position target' is unavailable " + f"for passive joints on backend '{self._backend_type}': {passive_names}" + ) + self.data.write_ctrl(target, env_ids, actuator_ids=actuator_ids) + def find_bodies( self, keys: str | Sequence[str], preserve_order: bool = False ) -> tuple[list[int], list[str]]: diff --git a/src/unilab/envs/mdp/__init__.py b/src/unilab/envs/mdp/__init__.py new file mode 100644 index 000000000..98de19312 --- /dev/null +++ b/src/unilab/envs/mdp/__init__.py @@ -0,0 +1,6 @@ +"""Community-style built-in MDP terms for UniLab's NumPy manager runtime.""" + +from unilab.envs.mdp.actions import JointPositionAction as JointPositionAction +from unilab.envs.mdp.actions import JointPositionActionCfg as JointPositionActionCfg + +__all__ = ["JointPositionAction", "JointPositionActionCfg"] diff --git a/src/unilab/envs/mdp/actions/__init__.py b/src/unilab/envs/mdp/actions/__init__.py new file mode 100644 index 000000000..f0fc2d98a --- /dev/null +++ b/src/unilab/envs/mdp/actions/__init__.py @@ -0,0 +1,6 @@ +"""Built-in action terms supported by the NumPy runtime.""" + +from unilab.envs.mdp.actions.actions import JointPositionAction as JointPositionAction +from unilab.envs.mdp.actions.actions import JointPositionActionCfg as JointPositionActionCfg + +__all__ = ["JointPositionAction", "JointPositionActionCfg"] diff --git a/src/unilab/envs/mdp/actions/actions.py b/src/unilab/envs/mdp/actions/actions.py new file mode 100644 index 000000000..4db452169 --- /dev/null +++ b/src/unilab/envs/mdp/actions/actions.py @@ -0,0 +1,223 @@ +# Derived from mujocolab/mjlab v1.6.0 (0fb8a681), +# src/mjlab/envs/mdp/actions/actions.py. +# Copyright 2025, The mjlab Developers. +# Modified by UniLab for NumPy and the SimBackend/entity contracts; Apache-2.0. +"""Actions that write joint-position targets to entity actuator controls.""" + +from __future__ import annotations + +import math +import re +from dataclasses import dataclass +from numbers import Real +from typing import TYPE_CHECKING, Any + +import numpy as np + +from unilab.dtype_config import get_global_dtype +from unilab.managers.action_manager import ActionTerm, ActionTermCfg + +if TYPE_CHECKING: + from unilab.base.entity import Entity + from unilab.managers._types import ManagerBasedRlEnv + + +def _real(value: Any, *, label: str) -> float: + if isinstance(value, bool) or not isinstance(value, Real): + raise TypeError(f"{label} must be a real number, got {type(value).__name__}") + result = float(value) + if not math.isfinite(result): + raise ValueError(f"{label} must be finite, got {result}") + return result + + +def _resolve_named_values( + values: dict[str, Any], names: list[str], *, label: str +) -> tuple[list[int], list[Any]]: + """Resolve regex-keyed values once, preserving target-name order.""" + if not isinstance(values, dict): + raise TypeError(f"{label} must be a dict") + patterns = list(values) + matched_by_pattern = [False] * len(patterns) + indices: list[int] = [] + resolved: list[Any] = [] + for index, name in enumerate(names): + matches: list[int] = [] + for pattern_index, pattern in enumerate(patterns): + try: + matches_pattern = re.fullmatch(pattern, name) is not None + except re.error as exc: + raise ValueError(f"{label} contains invalid regex {pattern!r}: {exc}") from exc + if matches_pattern: + matches.append(pattern_index) + if len(matches) > 1: + rendered = [patterns[pattern_index] for pattern_index in matches] + raise ValueError(f"{label} patterns {rendered} both match target '{name}'") + if matches: + pattern_index = matches[0] + matched_by_pattern[pattern_index] = True + indices.append(index) + resolved.append(values[patterns[pattern_index]]) + missing = [pattern for pattern, matched in zip(patterns, matched_by_pattern) if not matched] + if missing: + raise ValueError(f"{label} patterns {missing} match no targets; available={names}") + return indices, resolved + + +@dataclass(kw_only=True) +class BaseActionCfg(ActionTermCfg): + """Configuration shared by entity joint actions.""" + + actuator_names: tuple[str, ...] | list[str] + scale: float | dict[str, float] = 1.0 + offset: float | dict[str, float] = 0.0 + preserve_order: bool = False + + +class BaseAction(ActionTerm): + """Apply a cold-path-resolved affine transform to raw policy actions.""" + + cfg: BaseActionCfg + _entity: Entity + + def __init__(self, cfg: BaseActionCfg, env: ManagerBasedRlEnv): + super().__init__(cfg=cfg, env=env) + target_ids, target_names = self._entity.find_joints_by_actuator_names(cfg.actuator_names) + self._target_ids = np.asarray(target_ids, dtype=np.intp) + self._target_ids.setflags(write=False) + self._target_names = list(target_names) + self._action_dim = len(target_ids) + dtype = get_global_dtype() + self._raw_actions = np.zeros((self.num_envs, self.action_dim), dtype=dtype) + self._processed_actions = np.zeros_like(self._raw_actions) + self._scale = self._resolve_affine(cfg.scale, default=1.0, label="scale") + self._offset = self._resolve_affine(cfg.offset, default=0.0, label="offset") + self._clip = self._resolve_clip(cfg.clip) + + def _resolve_affine( + self, value: float | dict[str, float], *, default: float, label: str + ) -> float | np.ndarray: + if isinstance(value, dict): + result = np.full_like(self._raw_actions, default) + indices, resolved = _resolve_named_values( + value, self._target_names, label=f"{type(self).__name__} {label}" + ) + result[:, indices] = [ + _real(item, label=f"{type(self).__name__} {label}") for item in resolved + ] + return result + return _real(value, label=f"{type(self).__name__} {label}") + + def _resolve_clip(self, value: dict[str, tuple] | None) -> np.ndarray | None: + if value is None: + return None + result = np.empty((*self._raw_actions.shape, 2), dtype=self._raw_actions.dtype) + result[..., 0] = -np.inf + result[..., 1] = np.inf + indices, bounds = _resolve_named_values( + value, self._target_names, label=f"{type(self).__name__} clip" + ) + for index, raw_bounds in zip(indices, bounds, strict=True): + if not isinstance(raw_bounds, (tuple, list)) or len(raw_bounds) != 2: + raise TypeError( + f"{type(self).__name__} clip for '{self._target_names[index]}' " + "must be a (min, max) pair" + ) + lower = _real(raw_bounds[0], label=f"{type(self).__name__} clip lower") + upper = _real(raw_bounds[1], label=f"{type(self).__name__} clip upper") + if lower > upper: + raise ValueError( + f"{type(self).__name__} clip lower {lower} exceeds upper {upper} " + f"for '{self._target_names[index]}'" + ) + result[:, index, 0] = lower + result[:, index, 1] = upper + return result + + @property + def scale(self) -> float | np.ndarray: + return self._scale + + @property + def offset(self) -> float | np.ndarray: + return self._offset + + @property + def raw_action(self) -> np.ndarray: + return self._raw_actions + + @property + def processed_action(self) -> np.ndarray: + return self._processed_actions + + @property + def action_dim(self) -> int: + return self._action_dim + + @property + def target_ids(self) -> np.ndarray: + return self._target_ids + + @property + def target_names(self) -> list[str]: + return list(self._target_names) + + def process_actions(self, actions: np.ndarray) -> None: + if not isinstance(actions, np.ndarray): + raise TypeError( + f"{type(self).__name__} expected np.ndarray, got {type(actions).__name__}" + ) + if actions.shape != self._raw_actions.shape: + raise ValueError( + f"{type(self).__name__} expected action shape {self._raw_actions.shape}, " + f"got {actions.shape}" + ) + if not np.isfinite(actions).all(): + raise ValueError(f"{type(self).__name__} received NaN or Inf actions") + self._raw_actions[:] = actions + np.multiply(self._raw_actions, self._scale, out=self._processed_actions) + np.add(self._processed_actions, self._offset, out=self._processed_actions) + if self._clip is not None: + np.clip( + self._processed_actions, + self._clip[..., 0], + self._clip[..., 1], + out=self._processed_actions, + ) + + def reset(self, env_ids: np.ndarray | slice | None = None) -> None: + if env_ids is None: + env_ids = slice(None) + self._raw_actions[env_ids] = 0.0 + + +@dataclass(kw_only=True) +class JointPositionActionCfg(BaseActionCfg): + """Configuration for joint-position control.""" + + use_default_offset: bool = True + + def build(self, env: ManagerBasedRlEnv) -> JointPositionAction: + return JointPositionAction(self, env) + + +class JointPositionAction(BaseAction): + """Convert policy actions into entity joint-position targets.""" + + cfg: JointPositionActionCfg + + def __init__(self, cfg: JointPositionActionCfg, env: ManagerBasedRlEnv): + if not isinstance(cfg.use_default_offset, bool): + raise TypeError("JointPositionActionCfg use_default_offset must be bool") + super().__init__(cfg=cfg, env=env) + if cfg.use_default_offset: + self._offset = self._entity.data.default_joint_pos[:, self._target_ids].copy() + self._target = np.empty_like(self._processed_actions) + + def apply_actions(self) -> None: + encoder_bias = self._entity.data.encoder_bias[:, self._target_ids] + np.subtract(self._processed_actions, encoder_bias, out=self._target) + self._entity.set_joint_position_target(self._target, joint_ids=self._target_ids) + + +__all__ = ["JointPositionAction", "JointPositionActionCfg"] diff --git a/tests/base/test_entity_facade.py b/tests/base/test_entity_facade.py index 0342964e2..d30d5874b 100644 --- a/tests/base/test_entity_facade.py +++ b/tests/base/test_entity_facade.py @@ -32,6 +32,13 @@ def __init__(self, backend_type: str, *, unsupported: frozenset[str] = frozenset self.site_ids = {"imu": 3} self.geom_names = ("floor", "base_collision", "foot_collision") self.actuator_names = ("knee", "unused", "hip", "unused_2", "ankle") + self.actuator_joint_names = ( + "knee", + "unused_joint", + "hip", + "unused_2_joint", + "ankle", + ) self.dof_pos = np.arange(self.num_envs * 5, dtype=np.float32).reshape(self.num_envs, 5) self.dof_vel = self.dof_pos + 100.0 base = np.arange(self.num_envs * 10 * 3, dtype=np.float32) @@ -70,6 +77,10 @@ def get_actuator_names(self) -> tuple[str, ...]: self._check("actuator names") return self.actuator_names + def get_actuator_joint_names(self) -> tuple[str, ...]: + self._check("actuator target joints") + return self.actuator_joint_names + def get_actuator_ctrl_range(self) -> np.ndarray: self.calls["actuator range"] += 1 return np.arange(10, dtype=np.float32).reshape(5, 2) @@ -78,6 +89,10 @@ def get_dof_pos(self) -> np.ndarray: self._check("joint position state") return self.dof_pos + def get_default_dof_pos(self) -> np.ndarray: + self._check("default joint position") + return np.arange(5, dtype=np.float32) + 10.0 + def get_dof_vel(self) -> np.ndarray: self._check("joint velocity state") return self.dof_vel @@ -206,6 +221,74 @@ def test_entity_control_write_uses_cached_actuator_columns_and_fails_closed() -> read_only_scene["robot"].data.write_ctrl(np.zeros((backend.num_envs, 2), dtype=np.float32)) +def test_entity_joint_position_target_maps_natural_joint_order_to_control_order() -> None: + backend = _StrictBackendProfile("mujoco") + control = np.zeros((backend.num_envs, backend.num_actuators), dtype=np.float32) + scene = EntityScene( + { + "robot": EntityCfg( + joint_names=("ankle", "hip", "knee"), + actuator_names=("ankle", "hip"), + ) + }, + cast(SimBackend, backend), + control, + ) + robot = scene["robot"] + + ids, names = robot.find_joints_by_actuator_names(".*") + assert ids == [0, 1] + assert names == ["ankle", "hip"] + np.testing.assert_array_equal( + robot.data.default_joint_pos, + np.asarray([[14.0, 12.0, 10.0]] * backend.num_envs, dtype=np.float32), + ) + np.testing.assert_array_equal(robot.data.encoder_bias, 0.0) + + targets = np.asarray([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], dtype=np.float32) + robot.set_joint_position_target(targets, joint_ids=np.asarray(ids, dtype=np.int32)) + np.testing.assert_array_equal(control[:, 4], targets[:, 0]) + np.testing.assert_array_equal(control[:, 2], targets[:, 1]) + np.testing.assert_array_equal(control[:, [0, 1, 3]], 0.0) + + with pytest.raises(NotImplementedError, match="passive joints.*knee"): + robot.set_joint_position_target( + np.ones((backend.num_envs, 1), dtype=np.float32), + joint_ids=np.asarray([2], dtype=np.int32), + ) + + +@pytest.mark.parametrize( + ("actuator_joint_names", "joint_names", "message"), + [ + (("hip", "unused", "hip", "unused_2", "knee"), ("hip",), "must be unique"), + ( + ("knee", "unused_joint", "hip", "unused_2_joint", "ankle"), + ("hip",), + "outside its declared joint partition.*knee", + ), + ], +) +def test_entity_joint_actuator_mapping_rejects_ambiguous_or_missing_targets( + actuator_joint_names: tuple[str, ...], + joint_names: tuple[str, ...], + message: str, +) -> None: + backend = _StrictBackendProfile("mujoco") + backend.actuator_joint_names = actuator_joint_names + + with pytest.raises(ValueError, match=message): + EntityScene( + { + "robot": EntityCfg( + joint_names=joint_names, + actuator_names=("knee", "hip"), + ) + }, + cast(SimBackend, backend), + ) + + @pytest.mark.parametrize( ("ids", "error_type", "message"), [ diff --git a/tests/envs/mdp/test_joint_position_action.py b/tests/envs/mdp/test_joint_position_action.py new file mode 100644 index 000000000..dbba0a161 --- /dev/null +++ b/tests/envs/mdp/test_joint_position_action.py @@ -0,0 +1,246 @@ +"""Upstream-derived NumPy tests for the manager joint-position action.""" + +from __future__ import annotations + +import ast +from pathlib import Path +from types import SimpleNamespace +from typing import cast + +import numpy as np +import pytest + +from unilab.assets import ASSETS_ROOT_PATH +from unilab.base.backend import create_backend +from unilab.base.backend.base import SimBackend +from unilab.base.entity import EntityCfg, EntityScene +from unilab.base.scene import SceneCfg +from unilab.envs.mdp import JointPositionAction, JointPositionActionCfg +from unilab.envs.mdp.actions import ( + JointPositionAction as ExportedJointPositionAction, +) +from unilab.envs.mdp.actions import ( + JointPositionActionCfg as ExportedJointPositionActionCfg, +) +from unilab.managers._types import ManagerBasedRlEnv + + +class _Backend: + backend_type = "fake" + num_envs = 2 + num_actuators = 3 + + def __init__(self) -> None: + self.actuator_names = ("knee_motor", "hip_motor", "ankle_motor") + self.target_joint_names = ("knee", "hip", "ankle") + self.joint_index = {"hip": 0, "knee": 1, "ankle": 2} + self.dof_pos = np.zeros((self.num_envs, 3), dtype=np.float32) + + def get_actuator_names(self) -> tuple[str, ...]: + return self.actuator_names + + def get_actuator_joint_names(self) -> tuple[str, ...]: + return self.target_joint_names + + def get_actuator_ctrl_range(self) -> np.ndarray: + return np.tile(np.asarray([[-10.0, 10.0]], dtype=np.float32), (3, 1)) + + def get_joint_dof_pos_indices(self, names) -> np.ndarray: + return np.asarray([self.joint_index[name] for name in names], dtype=np.int32) + + def get_joint_dof_vel_indices(self, names) -> np.ndarray: + return self.get_joint_dof_pos_indices(names) + + def get_dof_pos(self) -> np.ndarray: + return self.dof_pos + + def get_dof_vel(self) -> np.ndarray: + return np.zeros_like(self.dof_pos) + + def get_default_dof_pos(self) -> np.ndarray: + return np.asarray([0.1, 0.2, 0.3], dtype=np.float32) + + +def _action( + **overrides, +) -> tuple[JointPositionAction, np.ndarray, EntityScene]: + backend = _Backend() + control = np.zeros((backend.num_envs, backend.num_actuators), dtype=np.float32) + scene = EntityScene( + { + "robot": EntityCfg( + joint_names=("hip", "knee", "ankle"), + actuator_names=backend.actuator_names, + ) + }, + cast(SimBackend, backend), + control, + ) + cfg_values = { + "entity_name": "robot", + "actuator_names": ("hip|knee",), + **overrides, + } + cfg = JointPositionActionCfg(**cfg_values) + env = cast(ManagerBasedRlEnv, SimpleNamespace(num_envs=backend.num_envs, scene=scene)) + return cfg.build(env), control, scene + + +def test_public_exports_are_canonical_objects() -> None: + assert JointPositionAction is ExportedJointPositionAction + assert JointPositionActionCfg is ExportedJointPositionActionCfg + + +def test_default_offset_encoder_bias_and_control_order() -> None: + action, control, scene = _action(scale=2.0) + raw = np.asarray([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32) + scene["robot"].data.encoder_bias[:, 0] = np.asarray([0.05, 0.1]) + + action.process_actions(raw) + action.apply_actions() + + assert action.target_names == ["hip", "knee"] + np.testing.assert_array_equal(action.target_ids, [0, 1]) + np.testing.assert_allclose(action.processed_action, raw * 2.0 + [0.1, 0.2]) + np.testing.assert_allclose(control[:, 1], action.processed_action[:, 0] - [0.05, 0.1]) + np.testing.assert_allclose(control[:, 0], action.processed_action[:, 1]) + np.testing.assert_array_equal(control[:, 2], 0.0) + + +def test_regex_scale_offset_clip_and_local_reset() -> None: + action, _, _ = _action( + scale={"hip": 2.0, "knee": 3.0}, + offset={"hip": 0.5, "knee": -0.5}, + clip={"hip": (-1.0, 1.0)}, + use_default_offset=False, + ) + raw = np.asarray([[2.0, 2.0], [-2.0, -2.0]], dtype=np.float32) + + action.process_actions(raw) + + np.testing.assert_allclose(action.processed_action, [[1.0, 5.5], [-1.0, -6.5]]) + np.testing.assert_array_equal(action.raw_action, raw) + action.reset(np.asarray([1], dtype=np.int32)) + np.testing.assert_array_equal(action.raw_action[0], raw[0]) + np.testing.assert_array_equal(action.raw_action[1], 0.0) + + +@pytest.mark.parametrize( + ("overrides", "error", "message"), + [ + ({"scale": {"missing": 1.0}}, ValueError, "match no targets"), + ( + {"scale": {"hip|knee": 1.0, ".*": 2.0}}, + ValueError, + "both match target", + ), + ({"clip": {"hip": (1.0, -1.0)}}, ValueError, "exceeds upper"), + ({"offset": float("nan")}, ValueError, "must be finite"), + ({"use_default_offset": 1}, TypeError, "must be bool"), + ], +) +def test_invalid_action_config_fails_at_construction(overrides, error, message) -> None: + with pytest.raises(error, match=message): + _action(**overrides) + + +def test_non_finite_and_wrong_shape_actions_fail_before_control_write() -> None: + action, control, _ = _action() + with pytest.raises(ValueError, match="expected action shape"): + action.process_actions(np.zeros((2, 1), dtype=np.float32)) + with pytest.raises(ValueError, match="NaN or Inf"): + action.process_actions(np.full((2, 2), np.nan, dtype=np.float32)) + np.testing.assert_array_equal(control, 0.0) + + +@pytest.mark.parametrize("backend_type", ["mujoco", "motrix"]) +def test_go2_joint_targets_are_mapped_to_backend_control_order(backend_type: str) -> None: + if backend_type == "motrix": + pytest.importorskip("motrixsim") + joint_names = ( + "FL_hip_joint", + "FL_thigh_joint", + "FL_calf_joint", + "FR_hip_joint", + "FR_thigh_joint", + "FR_calf_joint", + "RL_hip_joint", + "RL_thigh_joint", + "RL_calf_joint", + "RR_hip_joint", + "RR_thigh_joint", + "RR_calf_joint", + ) + scene_cfg = SceneCfg( + model_file=str(ASSETS_ROOT_PATH / "robots" / "go2" / "scene_flat.xml"), + entities={ + "robot": EntityCfg( + joint_names=joint_names, + actuator_names=( + "FR_hip", + "FR_thigh", + "FR_calf", + "FL_hip", + "FL_thigh", + "FL_calf", + "RR_hip", + "RR_thigh", + "RR_calf", + "RL_hip", + "RL_thigh", + "RL_calf", + ), + ) + }, + ) + backend = create_backend( + backend_type, + scene_cfg, + 2, + 0.01, + base_name="base", + ) + control = np.zeros((2, backend.num_actuators), dtype=np.float32) + scene = EntityScene.from_scene_cfg(scene_cfg, backend, control) + env = cast(ManagerBasedRlEnv, SimpleNamespace(num_envs=2, scene=scene)) + action = JointPositionActionCfg( + entity_name="robot", + actuator_names=(".*",), + scale=0.25, + offset={".*_hip_joint": 0.1, ".*_thigh_joint": 0.2, ".*_calf_joint": -0.3}, + use_default_offset=False, + ).build(env) + raw = np.arange(24, dtype=np.float32).reshape(2, 12) / 10.0 + + action.process_actions(raw) + action.apply_actions() + + target_index = {name: index for index, name in enumerate(joint_names)} + expected = np.column_stack( + [ + action.processed_action[:, target_index[name]] + for name in backend.get_actuator_joint_names() + ] + ) + np.testing.assert_allclose(control, expected) + + +def test_action_module_has_no_runtime_or_backend_private_dependencies() -> None: + path = ( + Path(__file__).resolve().parents[3] + / "src" + / "unilab" + / "envs" + / "mdp" + / "actions" + / "actions.py" + ) + tree = ast.parse(path.read_text(encoding="utf-8")) + forbidden = ("torch", "unilab.ipc", "unilab.algos", "unilab.training", "unilab.base.backend") + imports = [node.module or "" for node in ast.walk(tree) if isinstance(node, ast.ImportFrom)] + [ + alias.name + for node in ast.walk(tree) + if isinstance(node, ast.Import) + for alias in node.names + ] + assert not [name for name in imports if name.startswith(forbidden)] From 8ff448710fb1a30544ee6020ccc13f9cf913f626 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Mon, 17 Aug 2026 21:18:07 +0800 Subject: [PATCH 007/169] feat(env): add NumPy velocity command (#1057) --- src/unilab/base/entity.py | 39 +++ src/unilab/envs/mdp/__init__.py | 9 +- src/unilab/envs/mdp/commands/__init__.py | 10 + .../envs/mdp/commands/velocity_command.py | 226 ++++++++++++++ src/unilab/managers/_types.py | 6 + tests/base/test_entity_facade.py | 19 ++ .../envs/mdp/test_uniform_velocity_command.py | 294 ++++++++++++++++++ 7 files changed, 602 insertions(+), 1 deletion(-) create mode 100644 src/unilab/envs/mdp/commands/__init__.py create mode 100644 src/unilab/envs/mdp/commands/velocity_command.py create mode 100644 tests/envs/mdp/test_uniform_velocity_command.py diff --git a/src/unilab/base/entity.py b/src/unilab/base/entity.py index 678629048..dd116eb6f 100644 --- a/src/unilab/base/entity.py +++ b/src/unilab/base/entity.py @@ -16,6 +16,7 @@ import numpy as np from unilab.base.backend.base import SimBackend +from unilab.utils.rotation import np_yaw_from_quat if TYPE_CHECKING: from unilab.base.scene import SceneCfg @@ -192,6 +193,21 @@ def root_link_ang_vel_w(self) -> np.ndarray: ids = self._require(self._root_body_ids, "root body state") return self._backend.get_body_ang_vel_w(ids)[:, 0] + @property + def root_link_lin_vel_b(self) -> np.ndarray: + ids = self._require(self._root_body_ids, "root body state") + return self._backend.get_body_lin_vel_b(ids)[:, 0] + + @property + def root_link_ang_vel_b(self) -> np.ndarray: + ids = self._require(self._root_body_ids, "root body state") + return self._backend.get_body_ang_vel_b(ids)[:, 0] + + @property + def heading_w(self) -> np.ndarray: + """Root yaw in the world frame, derived from the backend quaternion view.""" + return np_yaw_from_quat(self.root_link_quat_w) + @property def root_link_pose_w(self) -> np.ndarray: return np.concatenate((self.root_link_pos_w, self.root_link_quat_w), axis=-1) @@ -558,6 +574,29 @@ def _validate_body_state( f"Entity '{self.name}' capability '{capability}' on backend " f"'{self._backend_type}' returned shape {value.shape}; expected {expected}" ) + if not np.isfinite(value).all(): + raise ValueError( + f"Entity '{self.name}' capability '{capability}' on backend " + f"'{self._backend_type}' returned NaN or Inf" + ) + if root_body_ids is None: + return + for capability, getter in ( + ("body-frame linear velocity state", backend.get_body_lin_vel_b), + ("body-frame angular velocity state", backend.get_body_ang_vel_b), + ): + value = self._read_state(capability, getter, root_body_ids) + expected = (backend.num_envs, len(root_body_ids), 3) + if value.shape != expected: + raise ValueError( + f"Entity '{self.name}' capability '{capability}' on backend " + f"'{self._backend_type}' returned shape {value.shape}; expected {expected}" + ) + if not np.isfinite(value).all(): + raise ValueError( + f"Entity '{self.name}' capability '{capability}' on backend " + f"'{self._backend_type}' returned NaN or Inf" + ) def _materialize_actuator_ctrl_range( self, backend: SimBackend, actuator_ids: np.ndarray | None diff --git a/src/unilab/envs/mdp/__init__.py b/src/unilab/envs/mdp/__init__.py index 98de19312..5062bd425 100644 --- a/src/unilab/envs/mdp/__init__.py +++ b/src/unilab/envs/mdp/__init__.py @@ -2,5 +2,12 @@ from unilab.envs.mdp.actions import JointPositionAction as JointPositionAction 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 -__all__ = ["JointPositionAction", "JointPositionActionCfg"] +__all__ = [ + "JointPositionAction", + "JointPositionActionCfg", + "UniformVelocityCommand", + "UniformVelocityCommandCfg", +] diff --git a/src/unilab/envs/mdp/commands/__init__.py b/src/unilab/envs/mdp/commands/__init__.py new file mode 100644 index 000000000..6647e1910 --- /dev/null +++ b/src/unilab/envs/mdp/commands/__init__.py @@ -0,0 +1,10 @@ +"""Built-in command terms supported by the NumPy runtime.""" + +from unilab.envs.mdp.commands.velocity_command import ( + UniformVelocityCommand as UniformVelocityCommand, +) +from unilab.envs.mdp.commands.velocity_command import ( + UniformVelocityCommandCfg as UniformVelocityCommandCfg, +) + +__all__ = ["UniformVelocityCommand", "UniformVelocityCommandCfg"] diff --git a/src/unilab/envs/mdp/commands/velocity_command.py b/src/unilab/envs/mdp/commands/velocity_command.py new file mode 100644 index 000000000..3932d77d3 --- /dev/null +++ b/src/unilab/envs/mdp/commands/velocity_command.py @@ -0,0 +1,226 @@ +# Derived from mujocolab/mjlab v1.6.0 (0fb8a681), +# src/mjlab/tasks/velocity/mdp/velocity_command.py. +# Copyright 2025, The mjlab Developers. +# Modified by UniLab for NumPy and the base-owned entity facade; Apache-2.0. +"""Uniform velocity commands for locomotion tasks.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass, field +from numbers import Real +from typing import TYPE_CHECKING, Any, cast + +import numpy as np + +from unilab.dtype_config import get_global_dtype +from unilab.managers.command_manager import CommandTerm, CommandTermCfg +from unilab.utils.rotation import np_wrap_to_pi + +if TYPE_CHECKING: + from unilab.base.entity import Entity + from unilab.managers._types import ManagerBasedRlEnv + + +def _real(value: Any, *, label: str) -> float: + if isinstance(value, bool) or not isinstance(value, Real): + raise TypeError(f"{label} must be a real number, got {type(value).__name__}") + result = float(value) + if not math.isfinite(result): + raise ValueError(f"{label} must be finite, got {result}") + return result + + +def _range_pair(value: Any, *, label: str) -> tuple[float, float]: + if not isinstance(value, (tuple, list)) or len(value) != 2: + raise TypeError(f"{label} must be a two-value range") + lower = _real(value[0], label=f"{label} lower") + upper = _real(value[1], label=f"{label} upper") + if lower > upper: + raise ValueError(f"{label} lower {lower} exceeds upper {upper}") + return lower, upper + + +def _ratio(value: Any, *, label: str) -> float: + result = _real(value, label=label) + if result < 0.0 or result > 1.0: + raise ValueError(f"{label} must be within [0, 1], got {result}") + return result + + +class UniformVelocityCommand(CommandTerm): + """Sample planar velocity commands and update frame-dependent components.""" + + cfg: UniformVelocityCommandCfg + + def __init__(self, cfg: UniformVelocityCommandCfg, env: ManagerBasedRlEnv): + self._validate_cfg(cfg) + super().__init__(cfg, env) + if cfg.init_velocity_prob > 0.0: + raise NotImplementedError( + "UniformVelocityCommand capability 'initial root velocity write' is " + "unavailable in the UniLab entity facade; set init_velocity_prob=0" + ) + + self.robot = cast("Entity", env.scene[cfg.entity_name]) + dtype = get_global_dtype() + self.vel_command_b = np.zeros((self.num_envs, 3), dtype=dtype) + self.vel_command_w = np.zeros_like(self.vel_command_b) + self.heading_target = np.zeros(self.num_envs, dtype=dtype) + self.heading_error = np.zeros(self.num_envs, dtype=dtype) + self.is_heading_env = np.zeros(self.num_envs, dtype=np.bool_) + self.is_standing_env = np.zeros(self.num_envs, dtype=np.bool_) + self.is_world_env = np.zeros(self.num_envs, dtype=np.bool_) + self.is_forward_env = np.zeros(self.num_envs, dtype=np.bool_) + self.metrics["error_vel_xy"] = np.zeros(self.num_envs, dtype=dtype) + self.metrics["error_vel_yaw"] = np.zeros(self.num_envs, dtype=dtype) + + @staticmethod + def _validate_cfg(cfg: UniformVelocityCommandCfg) -> None: + if not isinstance(cfg.entity_name, str) or not cfg.entity_name: + raise ValueError("UniformVelocityCommandCfg entity_name must be non-empty") + if not isinstance(cfg.heading_command, bool): + raise TypeError("UniformVelocityCommandCfg heading_command must be bool") + _real( + cfg.heading_control_stiffness, + label="UniformVelocityCommandCfg heading_control_stiffness", + ) + if cfg.heading_control_stiffness < 0.0: + raise ValueError("heading_control_stiffness must be non-negative") + for name in ( + "rel_standing_envs", + "rel_heading_envs", + "rel_world_envs", + "rel_forward_envs", + "init_velocity_prob", + ): + _ratio(getattr(cfg, name), label=f"UniformVelocityCommandCfg {name}") + if not isinstance(cfg.ranges, UniformVelocityCommandCfg.Ranges): + raise TypeError("UniformVelocityCommandCfg ranges must be a Ranges instance") + _range_pair(cfg.ranges.lin_vel_x, label="ranges.lin_vel_x") + _range_pair(cfg.ranges.lin_vel_y, label="ranges.lin_vel_y") + _range_pair(cfg.ranges.ang_vel_z, label="ranges.ang_vel_z") + if cfg.ranges.heading is not None: + _range_pair(cfg.ranges.heading, label="ranges.heading") + if cfg.heading_command and cfg.ranges.heading is None: + raise ValueError("heading_command=True but ranges.heading is None") + if cfg.ranges.heading is not None and not cfg.heading_command: + raise ValueError("ranges.heading is set but heading_command=False") + _, upper = _range_pair( + cfg.resampling_time_range, + label="UniformVelocityCommandCfg resampling_time_range", + ) + if upper <= 0.0: + raise ValueError("resampling_time_range upper bound must be positive") + + @property + def command(self) -> np.ndarray: + return self.vel_command_b + + def _update_metrics(self) -> None: + max_command_steps = self.cfg.resampling_time_range[1] / self._env.step_dt + self.metrics["error_vel_xy"] += ( + np.linalg.norm( + self.vel_command_b[:, :2] - self.robot.data.root_link_lin_vel_b[:, :2], + axis=-1, + ) + / max_command_steps + ) + self.metrics["error_vel_yaw"] += ( + np.abs(self.vel_command_b[:, 2] - self.robot.data.root_link_ang_vel_b[:, 2]) + / max_command_steps + ) + + def _resample_command(self, env_ids: np.ndarray) -> None: + count = len(env_ids) + rng = self._env.rng + ranges = self.cfg.ranges + self.vel_command_b[env_ids, 0] = rng.uniform(*ranges.lin_vel_x, size=count) + self.vel_command_b[env_ids, 1] = rng.uniform(*ranges.lin_vel_y, size=count) + self.vel_command_b[env_ids, 2] = rng.uniform(*ranges.ang_vel_z, size=count) + + if self.cfg.heading_command: + assert ranges.heading is not None + self.heading_target[env_ids] = rng.uniform(*ranges.heading, size=count) + self.is_heading_env[env_ids] = ( + rng.uniform(0.0, 1.0, size=count) <= self.cfg.rel_heading_envs + ) + self.is_standing_env[env_ids] = ( + rng.uniform(0.0, 1.0, size=count) <= self.cfg.rel_standing_envs + ) + self.is_world_env[env_ids] = rng.uniform(0.0, 1.0, size=count) <= self.cfg.rel_world_envs + self.vel_command_w[env_ids] = self.vel_command_b[env_ids] + self.is_forward_env[env_ids] = ( + rng.uniform(0.0, 1.0, size=count) <= self.cfg.rel_forward_envs + ) + forward_ids = env_ids[self.is_forward_env[env_ids]] + if len(forward_ids) > 0: + self.vel_command_b[forward_ids, 0] = np.maximum( + np.abs(self.vel_command_b[forward_ids, 0]), 0.3 + ) + self.vel_command_b[forward_ids, 1:] = 0.0 + + def _update_command(self, env_ids: np.ndarray | None = None) -> None: + del env_ids + if self.cfg.heading_command: + self.heading_error[:] = np_wrap_to_pi(self.heading_target - self.robot.data.heading_w) + heading_ids = np.flatnonzero(self.is_heading_env) + self.vel_command_b[heading_ids, 2] = np.clip( + self.cfg.heading_control_stiffness * self.heading_error[heading_ids], + self.cfg.ranges.ang_vel_z[0], + self.cfg.ranges.ang_vel_z[1], + ) + + world_ids = np.flatnonzero(self.is_world_env) + if len(world_ids) > 0: + heading = self.robot.data.heading_w[world_ids] + cos_heading = np.cos(heading) + sin_heading = np.sin(heading) + velocity_x_w = self.vel_command_w[world_ids, 0] + velocity_y_w = self.vel_command_w[world_ids, 1] + self.vel_command_b[world_ids, 0] = ( + cos_heading * velocity_x_w + sin_heading * velocity_y_w + ) + self.vel_command_b[world_ids, 1] = ( + -sin_heading * velocity_x_w + cos_heading * velocity_y_w + ) + + standing_ids = np.flatnonzero(self.is_standing_env) + self.vel_command_b[standing_ids] = 0.0 + self.vel_command_w[standing_ids] = 0.0 + + +@dataclass(kw_only=True) +class UniformVelocityCommandCfg(CommandTermCfg): + """Configuration for uniformly sampled planar velocity commands.""" + + entity_name: str + heading_command: bool = False + heading_control_stiffness: float = 1.0 + rel_standing_envs: float = 0.0 + rel_heading_envs: float = 1.0 + rel_world_envs: float = 0.0 + rel_forward_envs: float = 0.0 + init_velocity_prob: float = 0.0 + + @dataclass + class Ranges: + lin_vel_x: tuple[float, float] + lin_vel_y: tuple[float, float] + ang_vel_z: tuple[float, float] + heading: tuple[float, float] | None = None + + ranges: Ranges + + @dataclass + class VizCfg: + z_offset: float = 0.2 + scale: float = 0.5 + + viz: VizCfg = field(default_factory=VizCfg) + + def build(self, env: ManagerBasedRlEnv) -> UniformVelocityCommand: + return UniformVelocityCommand(self, env) + + +__all__ = ["UniformVelocityCommand", "UniformVelocityCommandCfg"] diff --git a/src/unilab/managers/_types.py b/src/unilab/managers/_types.py index 37f89dd89..6744cc825 100644 --- a/src/unilab/managers/_types.py +++ b/src/unilab/managers/_types.py @@ -146,6 +146,12 @@ def num_envs(self) -> int: ... @property def rng(self) -> np.random.Generator: ... + @property + def physics_dt(self) -> float: ... + + @property + def step_dt(self) -> float: ... + @property def scene(self) -> ManagerScene: ... diff --git a/tests/base/test_entity_facade.py b/tests/base/test_entity_facade.py index d30d5874b..d24d002f5 100644 --- a/tests/base/test_entity_facade.py +++ b/tests/base/test_entity_facade.py @@ -47,6 +47,8 @@ def __init__(self, backend_type: str, *, unsupported: frozenset[str] = frozenset self.body_quat[..., 0] = 1.0 self.body_lin_vel = self.body_pos + 200.0 self.body_ang_vel = self.body_pos + 300.0 + self.body_lin_vel_b = self.body_pos + 400.0 + self.body_ang_vel_b = self.body_pos + 500.0 def _check(self, capability: str) -> None: self.calls[capability] += 1 @@ -113,6 +115,14 @@ def get_body_ang_vel_w(self, ids: np.ndarray) -> np.ndarray: self._check("body angular velocity state") return self.body_ang_vel[:, ids] + def get_body_lin_vel_b(self, ids: np.ndarray) -> np.ndarray: + self._check("body-frame linear velocity state") + return self.body_lin_vel_b[:, ids] + + def get_body_ang_vel_b(self, ids: np.ndarray) -> np.ndarray: + self._check("body-frame angular velocity state") + return self.body_ang_vel_b[:, ids] + def _scene(backend_type: str = "mujoco") -> tuple[_StrictBackendProfile, EntityScene]: backend = _StrictBackendProfile(backend_type) @@ -143,6 +153,9 @@ def test_backend_profiles_materialize_identical_local_entity_contract(backend_ty np.testing.assert_array_equal(robot.data.joint_vel, backend.dof_vel[:, [4, 2]]) np.testing.assert_array_equal(robot.data.body_link_pos_w, backend.body_pos[:, [7, 4]]) np.testing.assert_array_equal(robot.data.root_link_pos_w, backend.body_pos[:, 4]) + np.testing.assert_array_equal(robot.data.root_link_lin_vel_b, backend.body_lin_vel_b[:, 4]) + np.testing.assert_array_equal(robot.data.root_link_ang_vel_b, backend.body_ang_vel_b[:, 4]) + np.testing.assert_array_equal(robot.data.heading_w, 0.0) np.testing.assert_array_equal( robot.data.actuator_ctrl_range, np.arange(10, dtype=np.float32).reshape(5, 2)[[4, 2]], @@ -461,6 +474,12 @@ def test_real_mujoco_entity_selector_and_numpy_state_smoke() -> None: assert selector.joint_ids == [2, 5, 8, 11] assert scene["robot"].data.joint_pos.shape == (2, 12) assert scene["robot"].data.root_link_pose_w.shape == (2, 7) + assert scene["robot"].data.root_link_lin_vel_b.shape == (2, 3) + assert scene["robot"].data.root_link_ang_vel_b.shape == (2, 3) + assert scene["robot"].data.heading_w.shape == (2,) + assert np.isfinite(scene["robot"].data.root_link_lin_vel_b).all() + assert np.isfinite(scene["robot"].data.root_link_ang_vel_b).all() + assert np.isfinite(scene["robot"].data.heading_w).all() def test_scene_cfg_entity_defaults_are_not_shared() -> None: diff --git a/tests/envs/mdp/test_uniform_velocity_command.py b/tests/envs/mdp/test_uniform_velocity_command.py new file mode 100644 index 000000000..f22d95a3a --- /dev/null +++ b/tests/envs/mdp/test_uniform_velocity_command.py @@ -0,0 +1,294 @@ +"""Upstream-derived NumPy tests for uniform velocity commands.""" + +from __future__ import annotations + +import ast +from pathlib import Path +from types import SimpleNamespace +from typing import Any, cast + +import numpy as np +import pytest + +from unilab.base.backend.base import SimBackend +from unilab.base.entity import EntityCfg, EntityScene +from unilab.envs.mdp import UniformVelocityCommand, UniformVelocityCommandCfg +from unilab.envs.mdp.commands import ( + UniformVelocityCommand as ExportedUniformVelocityCommand, +) +from unilab.envs.mdp.commands import ( + UniformVelocityCommandCfg as ExportedUniformVelocityCommandCfg, +) +from unilab.managers import CommandManager +from unilab.managers._types import ManagerBasedRlEnv +from unilab.utils.rotation import np_yaw_to_quat + + +class _Backend: + backend_type = "fake" + num_actuators = 0 + + def __init__(self, num_envs: int = 4, *, unsupported: str | None = None) -> None: + self.num_envs = num_envs + self.unsupported = unsupported + self.body_pos = np.zeros((num_envs, 1, 3), dtype=np.float32) + self.body_quat = np.zeros((num_envs, 1, 4), dtype=np.float32) + self.body_quat[..., 0] = 1.0 + self.body_lin_vel_w = np.zeros((num_envs, 1, 3), dtype=np.float32) + self.body_ang_vel_w = np.zeros((num_envs, 1, 3), dtype=np.float32) + self.body_lin_vel_b = np.zeros((num_envs, 1, 3), dtype=np.float32) + self.body_ang_vel_b = np.zeros((num_envs, 1, 3), dtype=np.float32) + + def get_body_ids(self, names) -> np.ndarray: + if tuple(names) != ("base",): + raise KeyError(names) + return np.asarray([0], dtype=np.int32) + + def get_dof_pos(self) -> np.ndarray: + return np.empty((self.num_envs, 0), dtype=np.float32) + + def get_dof_vel(self) -> np.ndarray: + return np.empty((self.num_envs, 0), dtype=np.float32) + + def get_body_pos_w(self, ids: np.ndarray) -> np.ndarray: + return self.body_pos[:, ids] + + def get_body_quat_w(self, ids: np.ndarray) -> np.ndarray: + return self.body_quat[:, ids] + + def get_body_lin_vel_w(self, ids: np.ndarray) -> np.ndarray: + return self.body_lin_vel_w[:, ids] + + def get_body_ang_vel_w(self, ids: np.ndarray) -> np.ndarray: + return self.body_ang_vel_w[:, ids] + + def get_body_lin_vel_b(self, ids: np.ndarray) -> np.ndarray: + if self.unsupported == "body-frame linear velocity": + raise NotImplementedError("fake lacks body-frame linear velocity") + return self.body_lin_vel_b[:, ids] + + def get_body_ang_vel_b(self, ids: np.ndarray) -> np.ndarray: + return self.body_ang_vel_b[:, ids] + + +def _env(seed: int = 7) -> tuple[ManagerBasedRlEnv, _Backend]: + backend = _Backend() + scene = EntityScene( + {"robot": EntityCfg(root_body_name="base")}, + cast(SimBackend, backend), + ) + env = cast( + ManagerBasedRlEnv, + SimpleNamespace( + num_envs=backend.num_envs, + rng=np.random.default_rng(seed), + scene=scene, + step_dt=0.02, + ), + ) + return env, backend + + +def _cfg(**overrides: Any) -> UniformVelocityCommandCfg: + values: dict[str, Any] = { + "entity_name": "robot", + "resampling_time_range": (1.0, 1.0), + "ranges": UniformVelocityCommandCfg.Ranges( + lin_vel_x=(-1.0, 1.0), + lin_vel_y=(-0.5, 0.5), + ang_vel_z=(-0.8, 0.8), + ), + **overrides, + } + return UniformVelocityCommandCfg(**values) + + +def _manager(env: ManagerBasedRlEnv, **overrides: Any) -> CommandManager: + return CommandManager({"twist": _cfg(**overrides)}, env) + + +def test_public_exports_and_cfg_build_are_canonical() -> None: + env, _ = _env() + term = _cfg().build(env) + + assert UniformVelocityCommand is ExportedUniformVelocityCommand + assert UniformVelocityCommandCfg is ExportedUniformVelocityCommandCfg + assert isinstance(term, UniformVelocityCommand) + assert term.command.shape == (env.num_envs, 3) + assert term.command.dtype == np.float32 + + +def test_resampling_is_seeded_and_partial_reset_is_local() -> None: + left_env, _ = _env(seed=19) + right_env, _ = _env(seed=19) + left = _manager(left_env, resampling_time_range=(0.5, 1.5)) + right = _manager(right_env, resampling_time_range=(0.5, 1.5)) + all_ids = np.arange(left_env.num_envs, dtype=np.int32) + + left.reset(all_ids) + right.reset(all_ids) + np.testing.assert_array_equal(left.get_command("twist"), right.get_command("twist")) + np.testing.assert_array_equal( + left.get_term("twist").time_left, + right.get_term("twist").time_left, + ) + + before_command = left.get_command("twist").copy() + before_counter = left.get_term("twist").command_counter.copy() + left.reset(np.asarray([1, 3], dtype=np.int32)) + np.testing.assert_array_equal(left.get_command("twist")[[0, 2]], before_command[[0, 2]]) + np.testing.assert_array_equal( + left.get_term("twist").command_counter[[0, 2]], before_counter[[0, 2]] + ) + np.testing.assert_array_equal(left.get_term("twist").command_counter[[1, 3]], 1) + + +def test_metrics_and_fixed_interval_resampling_follow_manager_schedule() -> None: + env, backend = _env() + backend.body_lin_vel_b[:, 0, :2] = [0.5, -0.25] + backend.body_ang_vel_b[:, 0, 2] = 0.1 + manager = _manager( + env, + resampling_time_range=(0.2, 0.2), + ranges=UniformVelocityCommandCfg.Ranges( + lin_vel_x=(1.0, 1.0), + lin_vel_y=(0.0, 0.0), + ang_vel_z=(0.2, 0.2), + ), + ) + all_ids = np.arange(env.num_envs, dtype=np.int32) + manager.reset(all_ids) + term = manager.get_term("twist") + + manager.compute(0.1) + np.testing.assert_allclose(term.metrics["error_vel_xy"], np.hypot(0.5, 0.25) / 10.0) + np.testing.assert_allclose(term.metrics["error_vel_yaw"], 0.1 / 10.0) + np.testing.assert_array_equal(term.command_counter, 1) + manager.compute(0.1) + np.testing.assert_array_equal(term.command_counter, 2) + np.testing.assert_allclose(term.time_left, 0.2) + + +def test_heading_world_forward_and_standing_modes() -> None: + all_ids = np.arange(4, dtype=np.int32) + + heading_env, _ = _env() + heading = _manager( + heading_env, + heading_command=True, + rel_heading_envs=1.0, + heading_control_stiffness=1.0, + ranges=UniformVelocityCommandCfg.Ranges( + lin_vel_x=(0.0, 0.0), + lin_vel_y=(0.0, 0.0), + ang_vel_z=(-0.5, 0.5), + heading=(np.pi / 2, np.pi / 2), + ), + ) + heading.reset(all_ids) + heading.compute(0.0, env_ids=all_ids) + np.testing.assert_allclose(heading.get_command("twist")[:, 2], 0.5) + + world_env, world_backend = _env() + world_backend.body_quat[:, 0] = np_yaw_to_quat(np.full(4, np.pi / 2)) + world = _manager( + world_env, + rel_world_envs=1.0, + ranges=UniformVelocityCommandCfg.Ranges( + lin_vel_x=(1.0, 1.0), + lin_vel_y=(0.0, 0.0), + ang_vel_z=(0.0, 0.0), + ), + ) + world.reset(all_ids) + world.compute(0.0, env_ids=all_ids) + np.testing.assert_allclose(world.get_command("twist")[:, :2], [[0.0, -1.0]] * 4, atol=1e-6) + + forward_env, _ = _env() + forward = _manager( + forward_env, + rel_forward_envs=1.0, + ranges=UniformVelocityCommandCfg.Ranges( + lin_vel_x=(-0.1, -0.1), + lin_vel_y=(0.4, 0.4), + ang_vel_z=(0.2, 0.2), + ), + ) + forward.reset(all_ids) + np.testing.assert_allclose(forward.get_command("twist"), [[0.3, 0.0, 0.0]] * 4) + + standing_env, _ = _env() + standing = _manager(standing_env, rel_standing_envs=1.0) + standing.reset(all_ids) + standing.compute(0.0, env_ids=all_ids) + np.testing.assert_array_equal(standing.get_command("twist"), 0.0) + + +@pytest.mark.parametrize( + ("overrides", "error", "message"), + [ + ({"rel_standing_envs": 1.1}, ValueError, "within.*0, 1"), + ({"heading_command": True}, ValueError, "heading_command=True"), + ( + { + "ranges": UniformVelocityCommandCfg.Ranges( + lin_vel_x=(1.0, -1.0), + lin_vel_y=(0.0, 0.0), + ang_vel_z=(0.0, 0.0), + ) + }, + ValueError, + "exceeds upper", + ), + ({"heading_control_stiffness": float("nan")}, ValueError, "must be finite"), + ({"resampling_time_range": (0.0, 0.0)}, ValueError, "must be positive"), + ({"init_velocity_prob": 0.1}, NotImplementedError, "root velocity write"), + ], +) +def test_invalid_or_unsupported_config_fails_at_construction( + overrides: dict[str, Any], error: type[Exception], message: str +) -> None: + env, _ = _env() + with pytest.raises(error, match=message): + _cfg(**overrides).build(env) + + +def test_root_body_frame_capability_and_finite_values_fail_closed() -> None: + unsupported = _Backend(unsupported="body-frame linear velocity") + with pytest.raises( + NotImplementedError, + match="body-frame linear velocity state.*backend 'fake'", + ): + EntityScene( + {"robot": EntityCfg(root_body_name="base")}, + cast(SimBackend, unsupported), + ) + + non_finite = _Backend() + non_finite.body_ang_vel_b[0, 0, 0] = np.nan + with pytest.raises(ValueError, match="body-frame angular velocity state.*NaN or Inf"): + EntityScene( + {"robot": EntityCfg(root_body_name="base")}, + cast(SimBackend, non_finite), + ) + + +def test_command_module_has_no_forbidden_runtime_dependencies() -> None: + path = ( + Path(__file__).resolve().parents[3] + / "src" + / "unilab" + / "envs" + / "mdp" + / "commands" + / "velocity_command.py" + ) + tree = ast.parse(path.read_text(encoding="utf-8")) + forbidden = ("torch", "unilab.ipc", "unilab.algos", "unilab.training", "unilab.base.backend") + imports = [node.module or "" for node in ast.walk(tree) if isinstance(node, ast.ImportFrom)] + [ + alias.name + for node in ast.walk(tree) + if isinstance(node, ast.Import) + for alias in node.names + ] + assert not [name for name in imports if name.startswith(forbidden)] From e843604435a59ba2358321c293e0b684ba70d3cb Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Mon, 17 Aug 2026 21:35:00 +0800 Subject: [PATCH 008/169] feat(env): add NumPy observation terms (#1060) --- src/unilab/base/entity.py | 52 ++++- src/unilab/envs/manager_based_rl_env.py | 9 + src/unilab/envs/mdp/__init__.py | 14 ++ src/unilab/envs/mdp/observations.py | 96 +++++++++ src/unilab/managers/_types.py | 22 +++ tests/base/test_entity_facade.py | 7 + tests/envs/mdp/test_observations.py | 247 ++++++++++++++++++++++++ 7 files changed, 446 insertions(+), 1 deletion(-) create mode 100644 src/unilab/envs/mdp/observations.py create mode 100644 tests/envs/mdp/test_observations.py diff --git a/src/unilab/base/entity.py b/src/unilab/base/entity.py index dd116eb6f..7c9017c4c 100644 --- a/src/unilab/base/entity.py +++ b/src/unilab/base/entity.py @@ -16,7 +16,7 @@ import numpy as np from unilab.base.backend.base import SimBackend -from unilab.utils.rotation import np_yaw_from_quat +from unilab.utils.rotation import np_quat_apply_inverse, np_yaw_from_quat if TYPE_CHECKING: from unilab.base.scene import SceneCfg @@ -141,6 +141,8 @@ def __init__( joint_pos_ids: np.ndarray | None, joint_vel_ids: np.ndarray | None, default_joint_pos: np.ndarray | None, + default_joint_vel: np.ndarray | None, + gravity_vec_w: np.ndarray | None, body_ids: np.ndarray | None, actuator_ids: np.ndarray | None, actuator_ctrl_range: np.ndarray | None, @@ -155,6 +157,8 @@ def __init__( self._joint_pos_index = None if joint_pos_ids is None else _as_column_index(joint_pos_ids) self._joint_vel_index = None if joint_vel_ids is None else _as_column_index(joint_vel_ids) self._default_joint_pos = default_joint_pos + self._default_joint_vel = default_joint_vel + self._gravity_vec_w = gravity_vec_w self._encoder_bias = ( None if default_joint_pos is None @@ -208,6 +212,12 @@ def heading_w(self) -> np.ndarray: """Root yaw in the world frame, derived from the backend quaternion view.""" return np_yaw_from_quat(self.root_link_quat_w) + @property + def projected_gravity_b(self) -> np.ndarray: + """Unit gravity vector projected into the root link frame.""" + gravity = self._require(self._gravity_vec_w, "projected gravity") + return np_quat_apply_inverse(self.root_link_quat_w, gravity) + @property def root_link_pose_w(self) -> np.ndarray: return np.concatenate((self.root_link_pos_w, self.root_link_quat_w), axis=-1) @@ -226,11 +236,21 @@ def joint_vel(self) -> np.ndarray: index = self._require(self._joint_vel_index, "joint velocity") return self._backend.get_dof_vel()[:, index] + @property + def joint_pos_biased(self) -> np.ndarray: + """Joint positions with the manager-owned encoder bias applied.""" + return self.joint_pos + self.encoder_bias + @property def default_joint_pos(self) -> np.ndarray: """Read-only per-environment default joint positions.""" return self._require(self._default_joint_pos, "default joint position") + @property + def default_joint_vel(self) -> np.ndarray: + """Read-only zero default velocities from the UniLab reset contract.""" + return self._require(self._default_joint_vel, "default joint velocity") + @property def encoder_bias(self) -> np.ndarray: """Mutable per-environment joint encoder bias used by position actions.""" @@ -433,6 +453,8 @@ def __init__( self._validate_joint_state(backend, joint_pos_ids, joint_vel_ids) self._validate_body_state(backend, root_body_ids, body_ids) default_joint_pos = self._materialize_default_joint_pos(backend, joint_pos_ids) + default_joint_vel = self._materialize_default_joint_vel(backend, joint_vel_ids) + gravity_vec_w = self._materialize_gravity_vector(backend, root_body_ids) actuator_ctrl_range = self._materialize_actuator_ctrl_range(backend, actuator_ids) ( self._actuator_target_joint_names, @@ -457,6 +479,8 @@ def __init__( joint_pos_ids=joint_pos_ids, joint_vel_ids=joint_vel_ids, default_joint_pos=default_joint_pos, + default_joint_vel=default_joint_vel, + gravity_vec_w=gravity_vec_w, body_ids=body_ids, actuator_ids=actuator_ids, actuator_ctrl_range=actuator_ctrl_range, @@ -632,6 +656,32 @@ def _materialize_default_joint_pos( materialized.setflags(write=False) return materialized + def _materialize_default_joint_vel( + self, backend: SimBackend, joint_vel_ids: np.ndarray | None + ) -> np.ndarray | None: + if joint_vel_ids is None: + return None + current = self._read_state("joint velocity state", backend.get_dof_vel) + materialized = np.zeros( + (backend.num_envs, len(joint_vel_ids)), + dtype=current.dtype, + ) + materialized.setflags(write=False) + return materialized + + def _materialize_gravity_vector( + self, backend: SimBackend, root_body_ids: np.ndarray | None + ) -> np.ndarray | None: + if root_body_ids is None: + return None + quat = self._read_state( + "root body quaternion state", backend.get_body_quat_w, root_body_ids + ) + gravity = np.zeros((backend.num_envs, 3), dtype=quat.dtype) + gravity[:, 2] = -1.0 + gravity.setflags(write=False) + return gravity + def _materialize_joint_actuator_mapping( self, backend: SimBackend, actuator_ids: np.ndarray | None ) -> tuple[tuple[str, ...] | None, np.ndarray | None]: diff --git a/src/unilab/envs/manager_based_rl_env.py b/src/unilab/envs/manager_based_rl_env.py index 1b6773812..e5d8c0f63 100644 --- a/src/unilab/envs/manager_based_rl_env.py +++ b/src/unilab/envs/manager_based_rl_env.py @@ -138,6 +138,15 @@ class ManagerBasedRlEnv(NpEnv): is_vector_env = True _cfg: ManagerBasedRlEnvCfg + event_manager: EventManager + command_manager: CommandManager | NullCommandManager + action_manager: ActionManager + observation_manager: ObservationManager + termination_manager: TerminationManager + reward_manager: RewardManager + curriculum_manager: CurriculumManager | NullCurriculumManager + metrics_manager: MetricsManager | NullMetricsManager + recorder_manager: RecorderManager | NullRecorderManager def __init__(self, cfg: ManagerBasedRlEnvCfg, backend: SimBackend, num_envs: int): if not isinstance(cfg, ManagerBasedRlEnvCfg): diff --git a/src/unilab/envs/mdp/__init__.py b/src/unilab/envs/mdp/__init__.py index 5062bd425..34037ff38 100644 --- a/src/unilab/envs/mdp/__init__.py +++ b/src/unilab/envs/mdp/__init__.py @@ -4,10 +4,24 @@ 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.observations import base_ang_vel as base_ang_vel +from unilab.envs.mdp.observations import base_lin_vel as base_lin_vel +from unilab.envs.mdp.observations import generated_commands as generated_commands +from unilab.envs.mdp.observations import joint_pos_rel as joint_pos_rel +from unilab.envs.mdp.observations import joint_vel_rel as joint_vel_rel +from unilab.envs.mdp.observations import last_action as last_action +from unilab.envs.mdp.observations import projected_gravity as projected_gravity __all__ = [ "JointPositionAction", "JointPositionActionCfg", "UniformVelocityCommand", "UniformVelocityCommandCfg", + "base_ang_vel", + "base_lin_vel", + "generated_commands", + "joint_pos_rel", + "joint_vel_rel", + "last_action", + "projected_gravity", ] diff --git a/src/unilab/envs/mdp/observations.py b/src/unilab/envs/mdp/observations.py new file mode 100644 index 000000000..f1de28afe --- /dev/null +++ b/src/unilab/envs/mdp/observations.py @@ -0,0 +1,96 @@ +# Derived from mujocolab/mjlab v1.6.0 (0fb8a681), +# src/mjlab/envs/mdp/observations.py. +# Copyright 2025, The mjlab Developers. +# Modified by UniLab for NumPy and the base-owned entity facade; Apache-2.0. +"""Community-style observation terms for the NumPy manager runtime.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, cast + +import numpy as np + +from unilab.managers.scene_entity_config import SceneEntityCfg + +if TYPE_CHECKING: + from unilab.base.entity import Entity + from unilab.managers._types import ManagerBasedRlEnv + + +_DEFAULT_ASSET_CFG = SceneEntityCfg("robot") + + +def base_lin_vel( + env: ManagerBasedRlEnv, + asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG, +) -> np.ndarray: + asset = cast("Entity", env.scene[asset_cfg.name]) + return asset.data.root_link_lin_vel_b + + +def base_ang_vel( + env: ManagerBasedRlEnv, + asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG, +) -> np.ndarray: + asset = cast("Entity", env.scene[asset_cfg.name]) + return asset.data.root_link_ang_vel_b + + +def projected_gravity( + env: ManagerBasedRlEnv, + asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG, +) -> np.ndarray: + asset = cast("Entity", env.scene[asset_cfg.name]) + return asset.data.projected_gravity_b + + +def joint_pos_rel( + env: ManagerBasedRlEnv, + biased: bool = False, + asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG, +) -> np.ndarray: + if not isinstance(biased, bool): + raise TypeError(f"joint_pos_rel biased must be bool, got {type(biased).__name__}") + asset = cast("Entity", env.scene[asset_cfg.name]) + joint_ids = asset_cfg.joint_ids + joint_pos = asset.data.joint_pos_biased if biased else asset.data.joint_pos + return joint_pos[:, joint_ids] - asset.data.default_joint_pos[:, joint_ids] + + +def joint_vel_rel( + env: ManagerBasedRlEnv, + asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG, +) -> np.ndarray: + asset = cast("Entity", env.scene[asset_cfg.name]) + joint_ids = asset_cfg.joint_ids + return asset.data.joint_vel[:, joint_ids] - asset.data.default_joint_vel[:, joint_ids] + + +def last_action(env: ManagerBasedRlEnv, action_name: str | None = None) -> np.ndarray: + if action_name is None: + return env.action_manager.action + try: + return env.action_manager.get_term(action_name).raw_action + except KeyError as exc: + raise KeyError(f"Action term '{action_name}' not found") from exc + + +def generated_commands(env: ManagerBasedRlEnv, command_name: str) -> np.ndarray: + try: + command = env.command_manager.get_command(command_name) + except KeyError as exc: + raise KeyError(f"Command term '{command_name}' not found") from exc + if command is None: + raise KeyError(f"Command term '{command_name}' not found") + return command + + +__all__ = [ + "base_ang_vel", + "base_lin_vel", + "generated_commands", + "joint_pos_rel", + "joint_vel_rel", + "last_action", + "projected_gravity", +] diff --git a/src/unilab/managers/_types.py b/src/unilab/managers/_types.py index 6744cc825..4e29c7e61 100644 --- a/src/unilab/managers/_types.py +++ b/src/unilab/managers/_types.py @@ -133,6 +133,22 @@ class ManagerScene(Protocol): def __getitem__(self, name: str) -> ManagerEntity: ... +class ManagerActionTerm(Protocol): + @property + def raw_action(self) -> np.ndarray: ... + + +class ManagerActionManager(Protocol): + @property + def action(self) -> np.ndarray: ... + + def get_term(self, name: str) -> ManagerActionTerm: ... + + +class ManagerCommandManager(Protocol): + def get_command(self, name: str) -> np.ndarray | None: ... + + class ManagerBasedRlEnv(Protocol): """Structural context visible to manager terms. @@ -155,6 +171,12 @@ def step_dt(self) -> float: ... @property def scene(self) -> ManagerScene: ... + @property + def action_manager(self) -> ManagerActionManager: ... + + @property + def command_manager(self) -> ManagerCommandManager: ... + @property def max_episode_length_s(self) -> float: ... diff --git a/tests/base/test_entity_facade.py b/tests/base/test_entity_facade.py index d24d002f5..ed048f1eb 100644 --- a/tests/base/test_entity_facade.py +++ b/tests/base/test_entity_facade.py @@ -151,11 +151,14 @@ def test_backend_profiles_materialize_identical_local_entity_contract(backend_ty assert robot.body_names == ("foot", "base") np.testing.assert_array_equal(robot.data.joint_pos, backend.dof_pos[:, [4, 2]]) np.testing.assert_array_equal(robot.data.joint_vel, backend.dof_vel[:, [4, 2]]) + np.testing.assert_array_equal(robot.data.default_joint_vel, 0.0) + np.testing.assert_array_equal(robot.data.joint_pos_biased, backend.dof_pos[:, [4, 2]]) np.testing.assert_array_equal(robot.data.body_link_pos_w, backend.body_pos[:, [7, 4]]) np.testing.assert_array_equal(robot.data.root_link_pos_w, backend.body_pos[:, 4]) np.testing.assert_array_equal(robot.data.root_link_lin_vel_b, backend.body_lin_vel_b[:, 4]) np.testing.assert_array_equal(robot.data.root_link_ang_vel_b, backend.body_ang_vel_b[:, 4]) np.testing.assert_array_equal(robot.data.heading_w, 0.0) + np.testing.assert_array_equal(robot.data.projected_gravity_b, [[0.0, 0.0, -1.0]] * 3) np.testing.assert_array_equal( robot.data.actuator_ctrl_range, np.arange(10, dtype=np.float32).reshape(5, 2)[[4, 2]], @@ -477,9 +480,13 @@ def test_real_mujoco_entity_selector_and_numpy_state_smoke() -> None: assert scene["robot"].data.root_link_lin_vel_b.shape == (2, 3) assert scene["robot"].data.root_link_ang_vel_b.shape == (2, 3) assert scene["robot"].data.heading_w.shape == (2,) + assert scene["robot"].data.projected_gravity_b.shape == (2, 3) + assert scene["robot"].data.default_joint_vel.shape == (2, 12) assert np.isfinite(scene["robot"].data.root_link_lin_vel_b).all() assert np.isfinite(scene["robot"].data.root_link_ang_vel_b).all() assert np.isfinite(scene["robot"].data.heading_w).all() + assert np.isfinite(scene["robot"].data.projected_gravity_b).all() + np.testing.assert_array_equal(scene["robot"].data.default_joint_vel, 0.0) def test_scene_cfg_entity_defaults_are_not_shared() -> None: diff --git a/tests/envs/mdp/test_observations.py b/tests/envs/mdp/test_observations.py new file mode 100644 index 000000000..b64106587 --- /dev/null +++ b/tests/envs/mdp/test_observations.py @@ -0,0 +1,247 @@ +"""Upstream-derived NumPy tests for basic manager observation terms.""" + +from __future__ import annotations + +import ast +from pathlib import Path +from types import SimpleNamespace +from typing import Any, cast + +import numpy as np +import pytest + +from unilab.base.backend.base import SimBackend +from unilab.base.entity import EntityCfg, EntityScene +from unilab.envs import mdp +from unilab.managers import ObservationGroupCfg, ObservationManager, ObservationTermCfg +from unilab.managers._types import ManagerBasedRlEnv +from unilab.managers.scene_entity_config import SceneEntityCfg + + +class _Backend: + backend_type = "fake" + num_envs = 2 + num_actuators = 0 + + def __init__(self) -> None: + self.joint_names = ("hip", "knee", "ankle") + self.dof_pos = np.asarray( + [[0.4, 0.1, -0.2], [0.0, 0.3, 0.7]], + dtype=np.float32, + ) + self.dof_vel = np.asarray( + [[1.0, 2.0, 3.0], [-1.0, -2.0, -3.0]], + dtype=np.float32, + ) + self.body_pos = np.zeros((2, 1, 3), dtype=np.float32) + self.body_quat = np.zeros((2, 1, 4), dtype=np.float32) + self.body_quat[0, 0, 0] = 1.0 + half = np.sqrt(0.5) + self.body_quat[1, 0] = [half, half, 0.0, 0.0] + self.body_lin_vel_w = np.zeros((2, 1, 3), dtype=np.float32) + self.body_ang_vel_w = np.zeros((2, 1, 3), dtype=np.float32) + self.body_lin_vel_b = np.asarray( + [[[0.5, 0.25, 0.0]], [[-0.5, -0.25, 0.1]]], dtype=np.float32 + ) + self.body_ang_vel_b = np.asarray( + [[[0.1, 0.2, 0.3]], [[-0.1, -0.2, -0.3]]], dtype=np.float32 + ) + + def get_joint_dof_pos_indices(self, names) -> np.ndarray: + return np.asarray([self.joint_names.index(name) for name in names], dtype=np.int32) + + def get_joint_dof_vel_indices(self, names) -> np.ndarray: + return self.get_joint_dof_pos_indices(names) + + def get_body_ids(self, names) -> np.ndarray: + if tuple(names) != ("base",): + raise KeyError(names) + return np.asarray([0], dtype=np.int32) + + def get_dof_pos(self) -> np.ndarray: + return self.dof_pos + + def get_dof_vel(self) -> np.ndarray: + return self.dof_vel + + def get_default_dof_pos(self) -> np.ndarray: + return np.asarray([0.1, 0.2, 0.3], dtype=np.float32) + + def get_body_pos_w(self, ids: np.ndarray) -> np.ndarray: + return self.body_pos[:, ids] + + def get_body_quat_w(self, ids: np.ndarray) -> np.ndarray: + return self.body_quat[:, ids] + + def get_body_lin_vel_w(self, ids: np.ndarray) -> np.ndarray: + return self.body_lin_vel_w[:, ids] + + def get_body_ang_vel_w(self, ids: np.ndarray) -> np.ndarray: + return self.body_ang_vel_w[:, ids] + + def get_body_lin_vel_b(self, ids: np.ndarray) -> np.ndarray: + return self.body_lin_vel_b[:, ids] + + def get_body_ang_vel_b(self, ids: np.ndarray) -> np.ndarray: + return self.body_ang_vel_b[:, ids] + + +class _ActionManager: + def __init__(self) -> None: + self.action = np.arange(6, dtype=np.float32).reshape(2, 3) + self._terms = { + "legs": SimpleNamespace(raw_action=self.action[:, [0, 2]]), + } + + def get_term(self, name: str): + return self._terms[name] + + +class _CommandManager: + def __init__(self) -> None: + self.command = np.asarray([[1.0, 0.0, 0.2], [0.5, -0.1, -0.2]], dtype=np.float32) + + def get_command(self, name: str) -> np.ndarray: + if name != "twist": + raise KeyError(name) + return self.command + + +def _env() -> tuple[ManagerBasedRlEnv, _Backend]: + backend = _Backend() + scene = EntityScene( + { + "robot": EntityCfg( + root_body_name="base", + joint_names=backend.joint_names, + ) + }, + cast(SimBackend, backend), + ) + env = cast( + ManagerBasedRlEnv, + SimpleNamespace( + num_envs=backend.num_envs, + scene=scene, + action_manager=_ActionManager(), + command_manager=_CommandManager(), + rng=np.random.default_rng(4), + ), + ) + return env, backend + + +def test_root_joint_action_and_command_terms_match_numpy_contract() -> None: + env, backend = _env() + robot = cast(Any, env.scene["robot"]) + robot.data.encoder_bias[:] = [[0.01, 0.02, 0.03], [-0.01, -0.02, -0.03]] + + np.testing.assert_array_equal(mdp.base_lin_vel(env), backend.body_lin_vel_b[:, 0]) + np.testing.assert_array_equal(mdp.base_ang_vel(env), backend.body_ang_vel_b[:, 0]) + np.testing.assert_allclose( + mdp.projected_gravity(env), + [[0.0, 0.0, -1.0], [0.0, -1.0, 0.0]], + atol=1e-6, + ) + np.testing.assert_allclose( + mdp.joint_pos_rel(env), + backend.dof_pos - [0.1, 0.2, 0.3], + ) + np.testing.assert_allclose( + mdp.joint_pos_rel(env, biased=True), + backend.dof_pos + robot.data.encoder_bias - [0.1, 0.2, 0.3], + ) + np.testing.assert_array_equal(mdp.joint_vel_rel(env), backend.dof_vel) + np.testing.assert_array_equal(mdp.last_action(env), env.action_manager.action) + np.testing.assert_array_equal( + mdp.last_action(env, "legs"), env.action_manager.get_term("legs").raw_action + ) + np.testing.assert_array_equal( + mdp.generated_commands(env, "twist"), env.command_manager.get_command("twist") + ) + + +def test_scene_entity_selector_is_resolved_once_by_observation_manager() -> None: + env, backend = _env() + selector = SceneEntityCfg("robot", joint_names=("ankle", "hip"), preserve_order=True) + manager = ObservationManager( + { + "policy": ObservationGroupCfg( + terms={ + "joint_pos": ObservationTermCfg( + func=mdp.joint_pos_rel, + params={"asset_cfg": selector}, + ), + "joint_vel": ObservationTermCfg( + func=mdp.joint_vel_rel, + params={"asset_cfg": selector}, + ), + "command": ObservationTermCfg( + func=mdp.generated_commands, + params={"command_name": "twist"}, + ), + } + ) + }, + env, + ) + + result = manager.compute()["policy"] + assert isinstance(result, np.ndarray) + assert result.shape == (2, 7) + np.testing.assert_allclose(result[:, :2], backend.dof_pos[:, [2, 0]] - [0.3, 0.1]) + np.testing.assert_array_equal(result[:, 2:4], backend.dof_vel[:, [2, 0]]) + assert selector.joint_ids == slice(None) + resolved = manager.get_term_cfg("policy", "joint_pos").params["asset_cfg"] + assert resolved.joint_ids == [2, 0] + + +@pytest.mark.parametrize( + ("call", "message"), + [ + (lambda env: mdp.last_action(env, "missing"), "Action term 'missing' not found"), + ( + lambda env: mdp.generated_commands(env, "missing"), + "Command term 'missing' not found", + ), + (lambda env: mdp.joint_pos_rel(env, biased=1), "biased must be bool"), + ], +) +def test_invalid_term_requests_fail_explicitly(call, message: str) -> None: + env, _ = _env() + with pytest.raises((KeyError, TypeError), match=message): + call(env) + + +def test_missing_entity_capability_fails_instead_of_returning_zeros() -> None: + backend = _Backend() + scene = EntityScene( + {"robot": EntityCfg(joint_names=backend.joint_names)}, + cast(SimBackend, backend), + ) + env = cast( + ManagerBasedRlEnv, + SimpleNamespace( + num_envs=2, + scene=scene, + action_manager=_ActionManager(), + command_manager=_CommandManager(), + ), + ) + with pytest.raises(NotImplementedError, match="projected gravity.*not materialized"): + mdp.projected_gravity(env) + + +def test_observation_module_has_no_forbidden_runtime_dependencies() -> None: + path = ( + Path(__file__).resolve().parents[3] / "src" / "unilab" / "envs" / "mdp" / "observations.py" + ) + tree = ast.parse(path.read_text(encoding="utf-8")) + forbidden = ("torch", "unilab.ipc", "unilab.algos", "unilab.training", "unilab.base.backend") + imports = [node.module or "" for node in ast.walk(tree) if isinstance(node, ast.ImportFrom)] + [ + alias.name + for node in ast.walk(tree) + if isinstance(node, ast.Import) + for alias in node.names + ] + assert not [name for name in imports if name.startswith(forbidden)] From b77960f27be9f8f40a265ec1a1fd5c22df8e89c8 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Mon, 17 Aug 2026 21:48:14 +0800 Subject: [PATCH 009/169] feat(env): add NumPy termination terms (#1062) --- src/unilab/envs/mdp/__init__.py | 8 ++ src/unilab/envs/mdp/terminations.py | 82 ++++++++++++++++++ src/unilab/managers/_types.py | 6 ++ tests/envs/mdp/test_terminations.py | 127 ++++++++++++++++++++++++++++ 4 files changed, 223 insertions(+) create mode 100644 src/unilab/envs/mdp/terminations.py create mode 100644 tests/envs/mdp/test_terminations.py diff --git a/src/unilab/envs/mdp/__init__.py b/src/unilab/envs/mdp/__init__.py index 34037ff38..3df60715f 100644 --- a/src/unilab/envs/mdp/__init__.py +++ b/src/unilab/envs/mdp/__init__.py @@ -11,6 +11,11 @@ from unilab.envs.mdp.observations import joint_vel_rel as joint_vel_rel from unilab.envs.mdp.observations import last_action as last_action from unilab.envs.mdp.observations import projected_gravity as projected_gravity +from unilab.envs.mdp.terminations import bad_orientation as bad_orientation +from unilab.envs.mdp.terminations import ( + root_height_below_minimum as root_height_below_minimum, +) +from unilab.envs.mdp.terminations import time_out as time_out __all__ = [ "JointPositionAction", @@ -19,9 +24,12 @@ "UniformVelocityCommandCfg", "base_ang_vel", "base_lin_vel", + "bad_orientation", "generated_commands", "joint_pos_rel", "joint_vel_rel", "last_action", "projected_gravity", + "root_height_below_minimum", + "time_out", ] diff --git a/src/unilab/envs/mdp/terminations.py b/src/unilab/envs/mdp/terminations.py new file mode 100644 index 000000000..1437e54cb --- /dev/null +++ b/src/unilab/envs/mdp/terminations.py @@ -0,0 +1,82 @@ +# Derived from mujocolab/mjlab v1.6.0 (0fb8a681), +# src/mjlab/envs/mdp/terminations.py. +# Copyright 2025, The mjlab Developers. +# Modified by UniLab for NumPy and the base-owned entity facade; Apache-2.0. +"""Community-style termination terms for the NumPy manager runtime.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, cast + +import numpy as np + +from unilab.managers.scene_entity_config import SceneEntityCfg + +if TYPE_CHECKING: + from unilab.base.entity import Entity + from unilab.managers._types import ManagerBasedRlEnv + + +_DEFAULT_ASSET_CFG = SceneEntityCfg("robot") + + +def _require_state(term_name: str, value: np.ndarray, num_envs: int) -> np.ndarray: + if not isinstance(value, np.ndarray): + raise TypeError( + f"Termination term '{term_name}' expected an np.ndarray entity state, " + f"got {type(value).__name__}" + ) + if value.shape != (num_envs, 3): + raise ValueError( + f"Termination term '{term_name}' received entity state shape {value.shape}, " + f"expected ({num_envs}, 3)" + ) + if not np.isfinite(value).all(): + env_ids = np.flatnonzero(~np.isfinite(value).all(axis=1)).tolist() + raise ValueError( + f"Termination term '{term_name}' received NaN or Inf entity state for " + f"environments {env_ids[:10]}" + ) + return value + + +def time_out(env: ManagerBasedRlEnv) -> np.ndarray: + """Terminate when the episode length reaches its maximum.""" + return env.episode_length_buf >= env.max_episode_length + + +def bad_orientation( + env: ManagerBasedRlEnv, + limit_angle: float, + asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG, +) -> np.ndarray: + """Terminate when the asset orientation exceeds ``limit_angle``.""" + if isinstance(limit_angle, bool) or not isinstance(limit_angle, (int, float, np.number)): + raise TypeError("bad_orientation limit_angle must be a real number") + if not np.isfinite(limit_angle): + raise ValueError("bad_orientation limit_angle must be finite") + asset = cast("Entity", env.scene[asset_cfg.name]) + projected_gravity = _require_state( + "bad_orientation", asset.data.projected_gravity_b, env.num_envs + ) + return np.abs(np.arccos(np.clip(-projected_gravity[:, 2], -1.0, 1.0))) > limit_angle + + +def root_height_below_minimum( + env: ManagerBasedRlEnv, + minimum_height: float, + asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG, +) -> np.ndarray: + """Terminate when the asset root height is below ``minimum_height``.""" + if isinstance(minimum_height, bool) or not isinstance(minimum_height, (int, float, np.number)): + raise TypeError("root_height_below_minimum minimum_height must be a real number") + if not np.isfinite(minimum_height): + raise ValueError("root_height_below_minimum minimum_height must be finite") + asset = cast("Entity", env.scene[asset_cfg.name]) + root_pos_w = _require_state( + "root_height_below_minimum", asset.data.root_link_pos_w, env.num_envs + ) + return root_pos_w[:, 2] < minimum_height + + +__all__ = ["bad_orientation", "root_height_below_minimum", "time_out"] diff --git a/src/unilab/managers/_types.py b/src/unilab/managers/_types.py index 4e29c7e61..26d126a82 100644 --- a/src/unilab/managers/_types.py +++ b/src/unilab/managers/_types.py @@ -177,6 +177,12 @@ def action_manager(self) -> ManagerActionManager: ... @property def command_manager(self) -> ManagerCommandManager: ... + @property + def episode_length_buf(self) -> np.ndarray: ... + + @property + def max_episode_length(self) -> int: ... + @property def max_episode_length_s(self) -> float: ... diff --git a/tests/envs/mdp/test_terminations.py b/tests/envs/mdp/test_terminations.py new file mode 100644 index 000000000..aafddcb18 --- /dev/null +++ b/tests/envs/mdp/test_terminations.py @@ -0,0 +1,127 @@ +"""Upstream-derived NumPy tests for basic manager termination terms.""" + +from __future__ import annotations + +import ast +from pathlib import Path +from types import SimpleNamespace +from typing import Any, cast + +import numpy as np +import pytest + +from unilab.envs import mdp +from unilab.managers import TerminationManager, TerminationTermCfg +from unilab.managers._types import ManagerBasedRlEnv +from unilab.managers.scene_entity_config import SceneEntityCfg + + +def _env() -> ManagerBasedRlEnv: + angles = np.asarray([0.0, 0.5, 1.0], dtype=np.float32) + gravity = np.zeros((3, 3), dtype=np.float32) + gravity[:, 2] = -np.cos(angles) + root_pos = np.zeros((3, 3), dtype=np.float32) + root_pos[:, 2] = [0.2, 0.4, 0.6] + entity = SimpleNamespace( + data=SimpleNamespace( + projected_gravity_b=gravity, + root_link_pos_w=root_pos, + ) + ) + return cast( + ManagerBasedRlEnv, + SimpleNamespace( + num_envs=3, + episode_length_buf=np.asarray([9, 10, 11], dtype=np.int64), + max_episode_length=10, + scene={"robot": entity}, + ), + ) + + +def test_basic_termination_terms_match_pinned_semantics() -> None: + env = _env() + + np.testing.assert_array_equal(mdp.time_out(env), [False, True, True]) + np.testing.assert_array_equal(mdp.bad_orientation(env, limit_angle=0.6), [False, False, True]) + np.testing.assert_array_equal( + mdp.root_height_below_minimum(env, minimum_height=0.4), + [True, False, False], + ) + + +def test_terms_integrate_with_termination_manager_and_selector_resolution() -> None: + env = _env() + selector = SceneEntityCfg("robot") + manager = TerminationManager( + { + "time_out": TerminationTermCfg(func=mdp.time_out, time_out=True), + "bad_orientation": TerminationTermCfg( + func=mdp.bad_orientation, + params={"limit_angle": 0.6, "asset_cfg": selector}, + ), + "low_height": TerminationTermCfg( + func=mdp.root_height_below_minimum, + params={"minimum_height": 0.4}, + ), + }, + env, + ) + + np.testing.assert_array_equal(manager.compute(), [True, True, True]) + np.testing.assert_array_equal(manager.time_outs, [False, True, True]) + np.testing.assert_array_equal(manager.terminated, [True, False, True]) + assert selector.name == "robot" + + +@pytest.mark.parametrize( + ("term", "parameter", "message"), + [ + (mdp.bad_orientation, np.nan, "limit_angle must be finite"), + (mdp.bad_orientation, True, "limit_angle must be a real number"), + ( + mdp.root_height_below_minimum, + np.inf, + "minimum_height must be finite", + ), + ], +) +def test_invalid_scalar_parameters_fail_explicitly(term, parameter, message: str) -> None: + env = _env() + keyword = "limit_angle" if term is mdp.bad_orientation else "minimum_height" + with pytest.raises((TypeError, ValueError), match=message): + term(env, **{keyword: parameter}) + + +@pytest.mark.parametrize("field", ["projected_gravity_b", "root_link_pos_w"]) +def test_invalid_entity_state_fails_instead_of_becoming_false(field: str) -> None: + env = _env() + entity = cast(Any, env.scene["robot"]) + setattr(entity.data, field, np.full((3, 3), np.nan, dtype=np.float32)) + + with pytest.raises(ValueError, match="NaN or Inf entity state"): + if field == "projected_gravity_b": + mdp.bad_orientation(env, limit_angle=0.5) + else: + mdp.root_height_below_minimum(env, minimum_height=0.5) + + +def test_missing_entity_is_not_silently_replaced() -> None: + env = _env() + with pytest.raises(KeyError, match="missing"): + mdp.bad_orientation(env, 0.5, SceneEntityCfg("missing")) + + +def test_termination_module_has_no_forbidden_runtime_dependencies() -> None: + path = ( + Path(__file__).resolve().parents[3] / "src" / "unilab" / "envs" / "mdp" / "terminations.py" + ) + tree = ast.parse(path.read_text(encoding="utf-8")) + forbidden = ("torch", "unilab.ipc", "unilab.algos", "unilab.training", "unilab.base.backend") + imports = [node.module or "" for node in ast.walk(tree) if isinstance(node, ast.ImportFrom)] + [ + alias.name + for node in ast.walk(tree) + if isinstance(node, ast.Import) + for alias in node.names + ] + assert not [name for name in imports if name.startswith(forbidden)] From 63dad7dd54953a6b4558258ff519633157890a51 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Mon, 17 Aug 2026 22:00:31 +0800 Subject: [PATCH 010/169] feat(env): add NumPy reward terms (#1065) --- src/unilab/envs/mdp/__init__.py | 20 +++ src/unilab/envs/mdp/rewards.py | 143 +++++++++++++++++++++ src/unilab/managers/_types.py | 14 +++ tests/envs/mdp/test_rewards.py | 216 ++++++++++++++++++++++++++++++++ 4 files changed, 393 insertions(+) create mode 100644 src/unilab/envs/mdp/rewards.py create mode 100644 tests/envs/mdp/test_rewards.py diff --git a/src/unilab/envs/mdp/__init__.py b/src/unilab/envs/mdp/__init__.py index 3df60715f..925e49264 100644 --- a/src/unilab/envs/mdp/__init__.py +++ b/src/unilab/envs/mdp/__init__.py @@ -11,6 +11,17 @@ from unilab.envs.mdp.observations import joint_vel_rel as joint_vel_rel from unilab.envs.mdp.observations import last_action as last_action from unilab.envs.mdp.observations import projected_gravity as projected_gravity +from unilab.envs.mdp.rewards import action_acc_l2 as action_acc_l2 +from unilab.envs.mdp.rewards import action_rate_l2 as action_rate_l2 +from unilab.envs.mdp.rewards import ( + body_angular_velocity_penalty as body_angular_velocity_penalty, +) +from unilab.envs.mdp.rewards import flat_orientation_l2 as flat_orientation_l2 +from unilab.envs.mdp.rewards import is_alive as is_alive +from unilab.envs.mdp.rewards import is_terminated as is_terminated +from unilab.envs.mdp.rewards import joint_vel_l2 as joint_vel_l2 +from unilab.envs.mdp.rewards import track_angular_velocity as track_angular_velocity +from unilab.envs.mdp.rewards import track_linear_velocity as track_linear_velocity from unilab.envs.mdp.terminations import bad_orientation as bad_orientation from unilab.envs.mdp.terminations import ( root_height_below_minimum as root_height_below_minimum, @@ -22,14 +33,23 @@ "JointPositionActionCfg", "UniformVelocityCommand", "UniformVelocityCommandCfg", + "action_acc_l2", + "action_rate_l2", "base_ang_vel", "base_lin_vel", "bad_orientation", + "body_angular_velocity_penalty", + "flat_orientation_l2", "generated_commands", "joint_pos_rel", "joint_vel_rel", + "joint_vel_l2", "last_action", + "is_alive", + "is_terminated", "projected_gravity", "root_height_below_minimum", "time_out", + "track_angular_velocity", + "track_linear_velocity", ] diff --git a/src/unilab/envs/mdp/rewards.py b/src/unilab/envs/mdp/rewards.py new file mode 100644 index 000000000..dab37abe9 --- /dev/null +++ b/src/unilab/envs/mdp/rewards.py @@ -0,0 +1,143 @@ +# Derived from mujocolab/mjlab v1.6.0 (0fb8a681), +# src/mjlab/envs/mdp/rewards.py and src/mjlab/tasks/velocity/mdp/rewards.py. +# Copyright 2025, The mjlab Developers. +# Modified by UniLab for NumPy and the base-owned entity facade; Apache-2.0. +"""Community-style reward terms for the NumPy manager runtime.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, cast + +import numpy as np + +from unilab.managers.scene_entity_config import SceneEntityCfg + +if TYPE_CHECKING: + from unilab.base.entity import Entity + from unilab.managers._types import ManagerBasedRlEnv + + +_DEFAULT_ASSET_CFG = SceneEntityCfg("robot") + + +def _positive_std(term_name: str, std: float) -> float: + if isinstance(std, bool) or not isinstance(std, (int, float, np.number)): + raise TypeError(f"{term_name} std must be a real number") + value = float(std) + if not np.isfinite(value) or value <= 0.0: + raise ValueError(f"{term_name} std must be finite and positive") + return value + + +def _command(env: ManagerBasedRlEnv, command_name: str) -> np.ndarray: + try: + command = env.command_manager.get_command(command_name) + except KeyError as exc: + raise KeyError(f"Command term '{command_name}' not found") from exc + if command is None: + raise KeyError(f"Command term '{command_name}' not found") + return command + + +def is_alive(env: ManagerBasedRlEnv) -> np.ndarray: + """Reward environments that have not reached a non-timeout termination.""" + return np.logical_not(env.termination_manager.terminated).astype(np.float32, copy=False) + + +def is_terminated(env: ManagerBasedRlEnv) -> np.ndarray: + """Return one for non-timeout terminations.""" + return env.termination_manager.terminated.astype(np.float32, copy=False) + + +def joint_vel_l2( + env: ManagerBasedRlEnv, + asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG, +) -> np.ndarray: + """Penalize selected joint velocities with an L2-squared kernel.""" + asset = cast("Entity", env.scene[asset_cfg.name]) + return np.sum(np.square(asset.data.joint_vel[:, asset_cfg.joint_ids]), axis=1) + + +def action_rate_l2(env: ManagerBasedRlEnv) -> np.ndarray: + """Penalize the first difference of raw policy actions.""" + delta = env.action_manager.action - env.action_manager.prev_action + return np.sum(np.square(delta), axis=1) + + +def action_acc_l2(env: ManagerBasedRlEnv) -> np.ndarray: + """Penalize the second difference of raw policy actions.""" + action_acc = ( + env.action_manager.action + - 2.0 * env.action_manager.prev_action + + env.action_manager.prev_prev_action + ) + return np.sum(np.square(action_acc), axis=1) + + +def flat_orientation_l2( + env: ManagerBasedRlEnv, + asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG, +) -> np.ndarray: + """Penalize non-flat base orientation.""" + asset = cast("Entity", env.scene[asset_cfg.name]) + return np.sum(np.square(asset.data.projected_gravity_b[:, :2]), axis=1) + + +def track_linear_velocity( + env: ManagerBasedRlEnv, + std: float, + command_name: str, + asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG, +) -> np.ndarray: + """Reward commanded base linear velocity, assuming commanded z is zero.""" + scale = _positive_std("track_linear_velocity", std) + asset = cast("Entity", env.scene[asset_cfg.name]) + command = _command(env, command_name) + actual = asset.data.root_link_lin_vel_b + xy_error = np.sum(np.square(command[:, :2] - actual[:, :2]), axis=1) + z_error = np.square(actual[:, 2]) + return np.exp(-(xy_error + z_error) / scale**2) + + +def track_angular_velocity( + env: ManagerBasedRlEnv, + std: float, + command_name: str, + asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG, +) -> np.ndarray: + """Reward commanded yaw rate while keeping roll/pitch rates near zero.""" + scale = _positive_std("track_angular_velocity", std) + asset = cast("Entity", env.scene[asset_cfg.name]) + command = _command(env, command_name) + actual = asset.data.root_link_ang_vel_b + z_error = np.square(command[:, 2] - actual[:, 2]) + xy_error = np.sum(np.square(actual[:, :2]), axis=1) + return np.exp(-(z_error + xy_error) / scale**2) + + +def body_angular_velocity_penalty( + env: ManagerBasedRlEnv, + asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG, +) -> np.ndarray: + """Penalize roll/pitch angular velocity of one selected body.""" + asset = cast("Entity", env.scene[asset_cfg.name]) + ang_vel = asset.data.body_link_ang_vel_w[:, asset_cfg.body_ids, :] + if ang_vel.shape != (env.num_envs, 1, 3): + raise ValueError( + "body_angular_velocity_penalty requires exactly one body; " + f"received state shape {ang_vel.shape}" + ) + return np.sum(np.square(ang_vel[:, 0, :2]), axis=1) + + +__all__ = [ + "action_acc_l2", + "action_rate_l2", + "body_angular_velocity_penalty", + "flat_orientation_l2", + "is_alive", + "is_terminated", + "joint_vel_l2", + "track_angular_velocity", + "track_linear_velocity", +] diff --git a/src/unilab/managers/_types.py b/src/unilab/managers/_types.py index 26d126a82..eace62476 100644 --- a/src/unilab/managers/_types.py +++ b/src/unilab/managers/_types.py @@ -142,6 +142,12 @@ class ManagerActionManager(Protocol): @property def action(self) -> np.ndarray: ... + @property + def prev_action(self) -> np.ndarray: ... + + @property + def prev_prev_action(self) -> np.ndarray: ... + def get_term(self, name: str) -> ManagerActionTerm: ... @@ -149,6 +155,11 @@ class ManagerCommandManager(Protocol): def get_command(self, name: str) -> np.ndarray | None: ... +class ManagerTerminationManager(Protocol): + @property + def terminated(self) -> np.ndarray: ... + + class ManagerBasedRlEnv(Protocol): """Structural context visible to manager terms. @@ -177,6 +188,9 @@ def action_manager(self) -> ManagerActionManager: ... @property def command_manager(self) -> ManagerCommandManager: ... + @property + def termination_manager(self) -> ManagerTerminationManager: ... + @property def episode_length_buf(self) -> np.ndarray: ... diff --git a/tests/envs/mdp/test_rewards.py b/tests/envs/mdp/test_rewards.py new file mode 100644 index 000000000..b7fefd4b3 --- /dev/null +++ b/tests/envs/mdp/test_rewards.py @@ -0,0 +1,216 @@ +"""Upstream-derived NumPy tests for sensor-free manager reward terms.""" + +from __future__ import annotations + +import ast +from pathlib import Path +from types import SimpleNamespace +from typing import Any, cast + +import numpy as np +import pytest + +from unilab.envs import mdp +from unilab.managers import RewardManager, RewardTermCfg +from unilab.managers._types import ManagerBasedRlEnv +from unilab.managers.scene_entity_config import SceneEntityCfg + + +class _Entity: + joint_names = ("hip", "knee", "ankle") + body_names = ("base", "torso") + num_joints = 3 + num_bodies = 2 + + def __init__(self) -> None: + self.data = SimpleNamespace( + joint_vel=np.asarray( + [[1.0, 2.0, 3.0], [-1.0, 0.5, 2.0], [0.0, -2.0, 1.0]], + dtype=np.float32, + ), + projected_gravity_b=np.asarray( + [[0.0, 0.0, -1.0], [0.3, 0.4, -0.866], [0.6, 0.0, -0.8]], + dtype=np.float32, + ), + root_link_lin_vel_b=np.asarray( + [[1.0, 0.0, 0.0], [0.0, 0.5, 0.25], [-0.5, 0.0, -0.2]], + dtype=np.float32, + ), + root_link_ang_vel_b=np.asarray( + [[0.0, 0.0, 0.5], [0.1, 0.2, -0.25], [-0.3, 0.0, 0.1]], + dtype=np.float32, + ), + body_link_ang_vel_w=np.asarray( + [ + [[0.1, 0.2, 0.3], [1.0, 2.0, 3.0]], + [[-0.1, 0.4, 0.2], [-1.0, 0.5, 0.0]], + [[0.0, -0.2, 0.7], [0.25, -0.75, 0.1]], + ], + dtype=np.float32, + ), + ) + + def find_joints(self, keys, preserve_order: bool = False): + patterns = (keys,) if isinstance(keys, str) else tuple(keys) + ids = [self.joint_names.index(name) for name in patterns] + return ids, [self.joint_names[index] for index in ids] + + +class _CommandManager: + def __init__(self) -> None: + self.command = np.asarray( + [[1.0, 0.0, 0.5], [0.25, 0.0, -0.5], [0.0, 0.0, 0.0]], dtype=np.float32 + ) + + def get_command(self, name: str) -> np.ndarray: + if name != "twist": + raise KeyError(name) + return self.command + + +def _env() -> ManagerBasedRlEnv: + action = np.asarray( + [[1.0, 2.0], [0.5, -0.5], [-1.0, 0.25]], + dtype=np.float32, + ) + return cast( + ManagerBasedRlEnv, + SimpleNamespace( + num_envs=3, + scene={"robot": _Entity()}, + action_manager=SimpleNamespace( + action=action, + prev_action=action - 0.25, + prev_prev_action=action - 0.75, + ), + command_manager=_CommandManager(), + termination_manager=SimpleNamespace( + terminated=np.asarray([False, True, False], dtype=np.bool_) + ), + max_episode_length_s=2.0, + ), + ) + + +def test_generic_reward_terms_match_pinned_numpy_semantics() -> None: + env = _env() + entity = cast(Any, env.scene["robot"]) + + np.testing.assert_array_equal(mdp.is_alive(env), [1.0, 0.0, 1.0]) + np.testing.assert_array_equal(mdp.is_terminated(env), [0.0, 1.0, 0.0]) + np.testing.assert_allclose( + mdp.joint_vel_l2(env, SceneEntityCfg("robot", joint_ids=[2, 0])), + np.sum(np.square(entity.data.joint_vel[:, [2, 0]]), axis=1), + ) + np.testing.assert_allclose(mdp.action_rate_l2(env), 2 * 0.25**2) + np.testing.assert_allclose(mdp.action_acc_l2(env), 2 * 0.25**2) + np.testing.assert_allclose( + mdp.flat_orientation_l2(env), + np.sum(np.square(entity.data.projected_gravity_b[:, :2]), axis=1), + ) + + +def test_velocity_tracking_terms_match_pinned_equations() -> None: + env = _env() + entity = cast(Any, env.scene["robot"]) + command = cast(np.ndarray, env.command_manager.get_command("twist")) + + linear_error = np.sum( + np.square(command[:, :2] - entity.data.root_link_lin_vel_b[:, :2]), axis=1 + ) + np.square(entity.data.root_link_lin_vel_b[:, 2]) + angular_error = np.square(command[:, 2] - entity.data.root_link_ang_vel_b[:, 2]) + np.sum( + np.square(entity.data.root_link_ang_vel_b[:, :2]), axis=1 + ) + np.testing.assert_allclose( + mdp.track_linear_velocity(env, std=0.5, command_name="twist"), + np.exp(-linear_error / 0.5**2), + ) + np.testing.assert_allclose( + mdp.track_angular_velocity(env, std=0.4, command_name="twist"), + np.exp(-angular_error / 0.4**2), + ) + + +def test_body_angular_velocity_requires_one_selected_body() -> None: + env = _env() + np.testing.assert_allclose( + mdp.body_angular_velocity_penalty(env, SceneEntityCfg("robot", body_ids=[1])), + [5.0, 1.25, 0.625], + ) + with pytest.raises(ValueError, match="requires exactly one body"): + mdp.body_angular_velocity_penalty(env) + + +def test_terms_integrate_with_reward_manager_and_cold_selector_resolution() -> None: + env = _env() + selector = SceneEntityCfg("robot", joint_names=("ankle", "hip"), preserve_order=True) + manager = RewardManager( + { + "joint_velocity": RewardTermCfg( + func=mdp.joint_vel_l2, + weight=-0.5, + params={"asset_cfg": selector}, + ), + "track_linear": RewardTermCfg( + func=mdp.track_linear_velocity, + weight=2.0, + params={"std": 0.5, "command_name": "twist"}, + ), + }, + env, + scale_by_dt=False, + ) + + result = manager.compute(dt=0.02) + entity = cast(Any, env.scene["robot"]) + expected = -0.5 * np.sum(np.square(entity.data.joint_vel[:, [2, 0]]), axis=1) + expected += 2.0 * mdp.track_linear_velocity(env, std=0.5, command_name="twist") + np.testing.assert_allclose(result, expected) + resolved = manager.get_term_cfg("joint_velocity").params["asset_cfg"] + assert resolved.joint_ids == [2, 0] + + +@pytest.mark.parametrize( + ("call", "message"), + [ + ( + lambda env: mdp.track_linear_velocity(env, std=0.0, command_name="twist"), + "std must be finite and positive", + ), + ( + lambda env: mdp.track_angular_velocity(env, std=np.nan, command_name="twist"), + "std must be finite and positive", + ), + ( + lambda env: mdp.track_linear_velocity(env, std=0.5, command_name="missing"), + "Command term 'missing' not found", + ), + ], +) +def test_invalid_reward_requests_fail_explicitly(call, message: str) -> None: + with pytest.raises((KeyError, ValueError), match=message): + call(_env()) + + +def test_reward_manager_reports_nonfinite_term_and_name() -> None: + env = _env() + cast(Any, env.scene["robot"]).data.joint_vel[1, 0] = np.nan + manager = RewardManager( + {"joint_velocity": RewardTermCfg(func=mdp.joint_vel_l2, weight=-1.0)}, + env, + ) + with pytest.raises(ValueError, match="RewardManager term 'joint_velocity'.*NaN"): + manager.compute(dt=0.02) + + +def test_reward_module_has_no_forbidden_runtime_dependencies() -> None: + path = Path(__file__).resolve().parents[3] / "src" / "unilab" / "envs" / "mdp" / "rewards.py" + tree = ast.parse(path.read_text(encoding="utf-8")) + forbidden = ("torch", "unilab.ipc", "unilab.algos", "unilab.training", "unilab.base.backend") + imports = [node.module or "" for node in ast.walk(tree) if isinstance(node, ast.ImportFrom)] + [ + alias.name + for node in ast.walk(tree) + if isinstance(node, ast.Import) + for alias in node.names + ] + assert not [name for name in imports if name.startswith(forbidden)] From f28d63708a65450fe421de10e739c418c9da367c Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 01:30:33 +0800 Subject: [PATCH 011/169] feat(env): add reset event state transaction (#1067) Refs #1066\n\nUmbrella: #1042 --- src/unilab/base/backend/base.py | 17 ++ src/unilab/base/backend/drake/backend.py | 22 ++ src/unilab/base/backend/mjwarp/backend.py | 8 + src/unilab/base/backend/motrix/backend.py | 8 + src/unilab/base/backend/mujoco/backend.py | 6 + src/unilab/base/entity.py | 165 +++++++++-- src/unilab/base/reset_state.py | 337 ++++++++++++++++++++++ src/unilab/envs/manager_based_rl_env.py | 20 +- src/unilab/envs/mdp/__init__.py | 4 + src/unilab/envs/mdp/events.py | 32 ++ src/unilab/managers/_types.py | 18 +- tests/base/test_backend_conformance.py | 23 ++ tests/base/test_entity_facade.py | 13 + tests/base/test_reset_state.py | 233 +++++++++++++++ tests/envs/test_manager_based_rl_env.py | 125 ++++++++ 15 files changed, 997 insertions(+), 34 deletions(-) create mode 100644 src/unilab/base/reset_state.py create mode 100644 src/unilab/envs/mdp/events.py create mode 100644 tests/base/test_reset_state.py diff --git a/src/unilab/base/backend/base.py b/src/unilab/base/backend/base.py index 32b4a612f..ba23ea3d7 100644 --- a/src/unilab/base/backend/base.py +++ b/src/unilab/base/backend/base.py @@ -774,6 +774,23 @@ def get_joint_dof_vel_indices(self, names: Sequence[str]) -> np.ndarray: f"{type(self).__name__} does not implement get_joint_dof_vel_indices" ) + def get_joint_state_qpos_indices(self, names: Sequence[str]) -> np.ndarray: + """Resolve single-DoF joints to full ``set_state`` qpos columns. + + Unlike :meth:`get_joint_dof_pos_indices`, these indices address the + complete qpos vector accepted by :meth:`set_state`, including any root + coordinates. Manager reset transactions resolve them on the cold path. + """ + raise NotImplementedError( + f"{type(self).__name__} does not implement get_joint_state_qpos_indices" + ) + + def get_joint_state_qvel_indices(self, names: Sequence[str]) -> np.ndarray: + """Resolve single-DoF joints to full ``set_state`` qvel columns.""" + raise NotImplementedError( + f"{type(self).__name__} does not implement get_joint_state_qvel_indices" + ) + def get_site_jacobian_w( self, site_id: int, diff --git a/src/unilab/base/backend/drake/backend.py b/src/unilab/base/backend/drake/backend.py index dd29977e3..b2591f4f5 100644 --- a/src/unilab/base/backend/drake/backend.py +++ b/src/unilab/base/backend/drake/backend.py @@ -431,6 +431,28 @@ def get_joint_dof_vel_indices(self, names: Sequence[str]) -> np.ndarray: raise ValueError(f"Drake model does not contain joint {key!r}") from exc return np.asarray(indices, dtype=np.int32) + def get_joint_state_qpos_indices(self, names: Sequence[str]) -> np.ndarray: + indices: list[int] = [] + for name in names: + key = str(name) + self._require_single_dof_joint(key) + try: + indices.append(self._joint_qpos_adr_by_name[key]) + except KeyError as exc: + raise ValueError(f"Drake model does not contain joint {key!r}") from exc + return np.asarray(indices, dtype=np.int32) + + def get_joint_state_qvel_indices(self, names: Sequence[str]) -> np.ndarray: + indices: list[int] = [] + for name in names: + key = str(name) + self._require_single_dof_joint(key) + try: + indices.append(self._joint_qvel_adr_by_name[key]) + except KeyError as exc: + raise ValueError(f"Drake model does not contain joint {key!r}") from exc + return np.asarray(indices, dtype=np.int32) + # Stepping and reset. def step(self, ctrl: np.ndarray, nsteps: int = 1) -> dict | None: # UniLab passes one actuator command per env. An optional pre-step hook diff --git a/src/unilab/base/backend/mjwarp/backend.py b/src/unilab/base/backend/mjwarp/backend.py index 04e473b15..8d2ebe77a 100644 --- a/src/unilab/base/backend/mjwarp/backend.py +++ b/src/unilab/base/backend/mjwarp/backend.py @@ -461,6 +461,14 @@ def get_joint_dof_vel_indices(self, names: Sequence[str]) -> np.ndarray: return self.get_joint_dof_indices(names) - self._root_qvel_dim + def get_joint_state_qpos_indices(self, names: Sequence[str]) -> np.ndarray: + """Resolve named joints to full reset qpos columns.""" + return self.get_joint_dof_pos_indices(names) + self._root_qpos_dim + + def get_joint_state_qvel_indices(self, names: Sequence[str]) -> np.ndarray: + """Resolve named joints to full reset qvel columns.""" + return self.get_joint_dof_vel_indices(names) + self._root_qvel_dim + def get_actuator_gains(self) -> tuple[np.ndarray, np.ndarray]: """Expose immutable model defaults; this does not advertise gain DR support.""" kp = np.asarray(self._cpu_model.actuator_gainprm[:, 0], dtype=np.float32).copy() diff --git a/src/unilab/base/backend/motrix/backend.py b/src/unilab/base/backend/motrix/backend.py index 2a25ff35f..41536d43c 100644 --- a/src/unilab/base/backend/motrix/backend.py +++ b/src/unilab/base/backend/motrix/backend.py @@ -487,6 +487,14 @@ def get_joint_dof_vel_indices(self, names: Sequence[str]) -> np.ndarray: indices.append(self._joint_dof_local_index(name, int(joint.dof_vel_index), pos=False)) return np.array(indices, dtype=np.int32) + def get_joint_state_qpos_indices(self, names: Sequence[str]) -> np.ndarray: + indices = [int(self._resolve_single_dof_joint(name).dof_pos_index) for name in names] + return np.asarray(indices, dtype=np.int32) + + def get_joint_state_qvel_indices(self, names: Sequence[str]) -> np.ndarray: + indices = [int(self._resolve_single_dof_joint(name).dof_vel_index) for name in names] + return np.asarray(indices, dtype=np.int32) + def get_site_jacobian_w( self, site_id: int, diff --git a/src/unilab/base/backend/mujoco/backend.py b/src/unilab/base/backend/mujoco/backend.py index 0de0e3bee..d7a86afea 100644 --- a/src/unilab/base/backend/mujoco/backend.py +++ b/src/unilab/base/backend/mujoco/backend.py @@ -872,6 +872,12 @@ def get_joint_dof_pos_indices(self, names: Sequence[str]) -> np.ndarray: def get_joint_dof_vel_indices(self, names: Sequence[str]) -> np.ndarray: return self.get_joint_dof_indices(names) - self._root_qvel_dim + def get_joint_state_qpos_indices(self, names: Sequence[str]) -> np.ndarray: + return self.get_joint_dof_pos_indices(names) + self._root_qpos_dim + + def get_joint_state_qvel_indices(self, names: Sequence[str]) -> np.ndarray: + return self.get_joint_dof_vel_indices(names) + self._root_qvel_dim + def get_joint_range(self) -> np.ndarray | None: jnt_range = self._model.jnt_range mask = self._model.jnt_type != int(mujoco.mjtJoint.mjJNT_FREE) diff --git a/src/unilab/base/entity.py b/src/unilab/base/entity.py index 7c9017c4c..3dd913ca1 100644 --- a/src/unilab/base/entity.py +++ b/src/unilab/base/entity.py @@ -19,6 +19,7 @@ from unilab.utils.rotation import np_quat_apply_inverse, np_yaw_from_quat if TYPE_CHECKING: + from unilab.base.reset_state import ResetStateTransaction from unilab.base.scene import SceneCfg @@ -395,11 +396,16 @@ def __init__( cfg: EntityCfg, backend: SimBackend, control_buffer: np.ndarray | None = None, + reset_state: ResetStateTransaction | None = None, ) -> None: if not name: raise ValueError("Entity name must be a non-empty string") self.name = name self._backend_type = backend.backend_type + self._backend = backend + self._reset_state = reset_state + self._reset_joint_qpos_ids: np.ndarray | None = None + self._reset_joint_qvel_ids: np.ndarray | None = None self._joint_names = _normalize_names(name, "joint", cfg.joint_names) self._body_names = _normalize_names(name, "body", cfg.body_names) @@ -860,42 +866,115 @@ def set_joint_position_target( "joint position target", "joint-to-actuator metadata was not materialized", ) + local_joint_ids = self._normalize_local_joint_ids( + joint_ids, + capability="joint position target", + ) + actuator_ids = joint_to_actuator[local_joint_ids] + if np.any(actuator_ids < 0): + passive_names = [ + self.joint_names[int(index)] for index in local_joint_ids[actuator_ids < 0] + ] + raise NotImplementedError( + f"Entity '{self.name}' capability 'joint position target' is unavailable " + f"for passive joints on backend '{self._backend_type}': {passive_names}" + ) + self.data.write_ctrl(target, env_ids, actuator_ids=actuator_ids) + + def write_joint_state_to_sim( + self, + position: np.ndarray, + velocity: np.ndarray, + joint_ids: np.ndarray | Sequence[int] | slice | None = None, + env_ids: np.ndarray | slice | None = None, + ) -> None: + """Stage community-style joint state writes in the active reset transaction.""" + if self._reset_state is None: + raise self._capability_error( + "reset joint-state write", + "EntityScene was materialized without an env-owned reset transaction", + ) + if self._joint_names is None: + raise self._capability_error( + "reset joint-state write", + "joint_names were not declared in EntityCfg", + ) + local_joint_ids = self._normalize_local_joint_ids( + joint_ids, + capability="reset joint-state write", + ) + resolved_env_ids = self._normalize_reset_env_ids(env_ids) + self._materialize_reset_joint_indices() + assert self._reset_joint_qpos_ids is not None + assert self._reset_joint_qvel_ids is not None + self._reset_state.write_joint_state( + resolved_env_ids, + self._reset_joint_qpos_ids[local_joint_ids], + self._reset_joint_qvel_ids[local_joint_ids], + position, + velocity, + term_name=f"{self.name}.write_joint_state_to_sim", + ) + + def _materialize_reset_joint_indices(self) -> None: + if self._reset_joint_qpos_ids is not None: + return + assert self._joint_names is not None + try: + qpos_ids = self._backend.get_joint_state_qpos_indices(self._joint_names) + qvel_ids = self._backend.get_joint_state_qvel_indices(self._joint_names) + except (AttributeError, NotImplementedError) as exc: + raise self._capability_error("reset joint-state layout", str(exc)) from exc + self._reset_joint_qpos_ids = _readonly_ids( + qpos_ids, + expected=len(self._joint_names), + label=f"Entity '{self.name}' reset qpos", + ) + self._reset_joint_qvel_ids = _readonly_ids( + qvel_ids, + expected=len(self._joint_names), + label=f"Entity '{self.name}' reset qvel", + ) + + def _normalize_local_joint_ids( + self, + joint_ids: np.ndarray | Sequence[int] | slice | None, + *, + capability: str, + ) -> np.ndarray: if joint_ids is None: - local_joint_ids = np.arange(self.num_joints, dtype=np.intp) + ids = np.arange(self.num_joints, dtype=np.intp) elif isinstance(joint_ids, slice): - local_joint_ids = np.arange(self.num_joints, dtype=np.intp)[joint_ids] + ids = np.arange(self.num_joints, dtype=np.intp)[joint_ids] else: - raw_joint_ids = np.asarray(joint_ids) + raw = np.asarray(joint_ids) if ( - raw_joint_ids.ndim != 1 - or not np.issubdtype(raw_joint_ids.dtype, np.integer) - or np.issubdtype(raw_joint_ids.dtype, np.bool_) + raw.ndim != 1 + or not np.issubdtype(raw.dtype, np.integer) + or np.issubdtype(raw.dtype, np.bool_) ): raise TypeError( - f"Entity '{self.name}' joint position target joint_ids must be a 1-D " - "integer array or slice" + f"Entity '{self.name}' {capability} joint_ids must be a 1-D integer " + "array or slice" ) - local_joint_ids = np.asarray(raw_joint_ids, dtype=np.intp) - if np.any(local_joint_ids < 0) or np.any(local_joint_ids >= self.num_joints): + ids = np.asarray(raw, dtype=np.intp) + if np.any(ids < 0) or np.any(ids >= self.num_joints): raise IndexError( - f"Entity '{self.name}' joint position target joint_ids out of range for " - f"{self.num_joints} joints: {local_joint_ids.tolist()}" + f"Entity '{self.name}' {capability} joint_ids out of range for " + f"{self.num_joints} joints: {ids.tolist()}" ) - if np.unique(local_joint_ids).size != local_joint_ids.size: + if np.unique(ids).size != ids.size: raise ValueError( - f"Entity '{self.name}' joint position target joint_ids contain duplicates: " - f"{local_joint_ids.tolist()}" - ) - actuator_ids = joint_to_actuator[local_joint_ids] - if np.any(actuator_ids < 0): - passive_names = [ - self.joint_names[int(index)] for index in local_joint_ids[actuator_ids < 0] - ] - raise NotImplementedError( - f"Entity '{self.name}' capability 'joint position target' is unavailable " - f"for passive joints on backend '{self._backend_type}': {passive_names}" + f"Entity '{self.name}' {capability} joint_ids contain duplicates: {ids.tolist()}" ) - self.data.write_ctrl(target, env_ids, actuator_ids=actuator_ids) + 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) + if isinstance(env_ids, slice): + return np.arange(self._backend.num_envs, dtype=np.int32)[env_ids] + return env_ids def find_bodies( self, keys: str | Sequence[str], preserve_order: bool = False @@ -962,6 +1041,8 @@ def __init__( entities: Mapping[str, EntityCfg], backend: SimBackend, control_buffer: np.ndarray | None = None, + *, + reset_state: ResetStateTransaction | None = None, ) -> None: materialized: dict[str, Entity] = {} for name, cfg in entities.items(): @@ -971,8 +1052,12 @@ def __init__( raise TypeError( f"Scene entity '{name}' must be EntityCfg, got {type(cfg).__name__}" ) - materialized[name] = Entity(name, cfg, backend, control_buffer) + materialized[name] = Entity(name, cfg, backend, control_buffer, reset_state) self._entities = MappingProxyType(materialized) + self._reset_state = reset_state + env_origins = np.zeros((backend.num_envs, 3), dtype=np.float32) + env_origins.setflags(write=False) + self._env_origins = env_origins @classmethod def from_scene_cfg( @@ -980,8 +1065,34 @@ def from_scene_cfg( cfg: SceneCfg, backend: SimBackend, control_buffer: np.ndarray | None = None, + *, + reset_state: ResetStateTransaction | None = None, ) -> EntityScene: - return cls(cfg.entities, backend, control_buffer) + return cls(cfg.entities, backend, control_buffer, reset_state=reset_state) + + @property + def entities(self) -> Mapping[str, Entity]: + """Pinned community-style read-only entity mapping.""" + return self._entities + + @property + def env_origins(self) -> np.ndarray: + """Read-only per-environment origins; flat UniLab scenes default to zero.""" + return self._env_origins + + def reset_to_default(self, env_ids: np.ndarray, *, term_name: str) -> None: + """Stage a full-scene default state in the active reset transaction.""" + if self._reset_state is None: + raise NotImplementedError( + f"EventManager term '{term_name}' reset-state capability is unavailable: " + "EntityScene was materialized without an env-owned reset transaction" + ) + if np.any(self._env_origins): + raise NotImplementedError( + f"EventManager term '{term_name}' cannot apply non-zero env_origins without " + "a formal backend root-state layout" + ) + self._reset_state.reset_to_default(env_ids, term_name=term_name) def __getitem__(self, name: str) -> Entity: try: diff --git a/src/unilab/base/reset_state.py b/src/unilab/base/reset_state.py new file mode 100644 index 000000000..86a1bdffc --- /dev/null +++ b/src/unilab/base/reset_state.py @@ -0,0 +1,337 @@ +"""Base-owned reset-state transaction for Manager-Based event terms. + +The transaction composes NumPy state writes in memory and hands the finished +batch to :meth:`SimBackend.set_state` exactly once. It deliberately knows +nothing about task configuration, IPC, runners, or backend-private state. +""" + +from __future__ import annotations + +from collections.abc import Iterator +from contextlib import contextmanager + +import numpy as np + +from unilab.base.backend.base import SimBackend + + +class ResetStateTransaction: + """Reusable, fail-closed transaction for reset-mode state mutation.""" + + def __init__(self, backend: SimBackend) -> None: + self._backend = backend + self._num_envs = backend.num_envs + self._active = False + self._active_mask = np.zeros(self._num_envs, dtype=np.bool_) + self._dirty_mask = np.zeros(self._num_envs, dtype=np.bool_) + self._default_qpos: np.ndarray | None = None + self._default_qvel: np.ndarray | None = None + self._qpos: np.ndarray | None = None + self._qvel: np.ndarray | None = None + self._requesting_terms: set[str] = set() + + @property + def active(self) -> bool: + """Whether a reset lifecycle currently owns the transaction.""" + return self._active + + @contextmanager + def scoped(self, env_ids: np.ndarray) -> Iterator[ResetStateTransaction]: + """Begin a reset transaction and commit it only after all terms succeed.""" + self.begin(env_ids) + try: + yield self + except BaseException: + self.abort() + raise + else: + self.commit() + + def begin(self, env_ids: np.ndarray) -> None: + """Open a transaction for the concrete reset environment IDs.""" + if self._active: + raise RuntimeError("ManagerBased reset-state transaction is already active") + ids = self._validate_ids(env_ids, capability="begin") + self._active_mask.fill(False) + self._active_mask[ids] = True + self._dirty_mask.fill(False) + self._requesting_terms.clear() + self._active = 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() + ids = self._validate_ids(env_ids, capability="reset_to_default") + outside = ids[~self._active_mask[ids]] + if outside.size: + raise ValueError( + "EventManager term " + f"'{term_name}' attempted reset-state mutation outside the active reset: " + f"{outside.tolist()}" + ) + if ids.size == 0: + return + self._requesting_terms.add(term_name) + self._materialize_default_state(term_name) + assert self._default_qpos is not None + assert self._default_qvel is not None + assert self._qpos is not None + assert self._qvel is not None + self._qpos[ids] = self._default_qpos + self._qvel[ids] = self._default_qvel + self._dirty_mask[ids] = True + + def write_joint_state( + self, + env_ids: np.ndarray, + qpos_indices: np.ndarray, + qvel_indices: np.ndarray, + position: np.ndarray, + velocity: np.ndarray, + *, + term_name: str, + ) -> None: + """Stage selected joint position and velocity columns in the reset batch.""" + self._require_active() + ids = self._validate_ids(env_ids, capability="write_joint_state") + outside = ids[~self._active_mask[ids]] + if outside.size: + raise ValueError( + f"EventManager term '{term_name}' attempted joint-state mutation outside " + f"the active reset: {outside.tolist()}" + ) + self._requesting_terms.add(term_name) + self._materialize_default_state(term_name) + assert self._default_qpos is not None + assert self._default_qvel is not None + assert self._qpos is not None + assert self._qvel is not None + + pos_columns = self._validate_columns( + qpos_indices, + width=self._default_qpos.size, + capability="qpos indices", + term_name=term_name, + ) + vel_columns = self._validate_columns( + qvel_indices, + width=self._default_qvel.size, + capability="qvel indices", + term_name=term_name, + ) + if pos_columns.size != vel_columns.size: + raise ValueError( + f"EventManager term '{term_name}' joint-state qpos/qvel index counts differ: " + f"{pos_columns.size} != {vel_columns.size}" + ) + positions = self._validate_values( + position, + shape=(ids.size, pos_columns.size), + capability="joint position", + term_name=term_name, + ) + velocities = self._validate_values( + velocity, + shape=(ids.size, vel_columns.size), + capability="joint velocity", + term_name=term_name, + ) + + uninitialized = ids[~self._dirty_mask[ids]] + if uninitialized.size: + self._qpos[uninitialized] = self._default_qpos + self._qvel[uninitialized] = self._default_qvel + if ids.size and pos_columns.size: + self._qpos[ids[:, None], pos_columns[None, :]] = positions + self._qvel[ids[:, None], vel_columns[None, :]] = velocities + self._dirty_mask[ids] = True + + def commit(self) -> dict | None: + """Commit all staged rows through one public backend call.""" + self._require_active() + dirty_ids = np.flatnonzero(self._dirty_mask).astype(np.int32, copy=False) + try: + if dirty_ids.size == 0: + return None + assert self._qpos is not None + assert self._qvel is not None + try: + return self._backend.set_state( + dirty_ids, + self._qpos[dirty_ids], + self._qvel[dirty_ids], + ) + except (AttributeError, NotImplementedError) as exc: + terms = ", ".join(sorted(self._requesting_terms)) + raise NotImplementedError( + "EventManager reset-state capability 'SimBackend.set_state' is unavailable " + f"for term(s) [{terms}] on backend '{self._backend.backend_type}': {exc}" + ) from exc + finally: + self._finish() + + def abort(self) -> None: + """Discard staged rows without touching the backend.""" + if self._active: + self._finish() + + def _materialize_default_state(self, term_name: str) -> None: + if self._default_qpos is not None: + return + try: + qpos = self._backend.get_default_qpos() + except (AttributeError, NotImplementedError) as exc: + raise self._capability_error(term_name, "default qpos", exc) from exc + try: + qvel = self._backend.get_init_qvel() + except (AttributeError, NotImplementedError) as exc: + raise self._capability_error(term_name, "initial qvel", exc) from exc + + default_qpos = self._validate_state_vector(qpos, "default qpos", term_name) + default_qvel = self._validate_state_vector(qvel, "initial qvel", term_name) + self._default_qpos = default_qpos + self._default_qvel = default_qvel + 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 _validate_state_vector( + self, + value: np.ndarray, + capability: str, + term_name: str, + ) -> np.ndarray: + if not isinstance(value, np.ndarray): + raise TypeError( + f"EventManager term '{term_name}' capability '{capability}' on backend " + f"'{self._backend.backend_type}' must return np.ndarray, got " + f"{type(value).__name__}" + ) + if value.ndim != 1: + raise ValueError( + f"EventManager term '{term_name}' capability '{capability}' on backend " + f"'{self._backend.backend_type}' returned shape {value.shape}; expected 1-D" + ) + if not np.issubdtype(value.dtype, np.floating): + raise TypeError( + f"EventManager term '{term_name}' capability '{capability}' on backend " + f"'{self._backend.backend_type}' must be floating, got {value.dtype}" + ) + if not np.isfinite(value).all(): + raise ValueError( + f"EventManager term '{term_name}' capability '{capability}' on backend " + f"'{self._backend.backend_type}' returned NaN or Inf" + ) + result = np.array(value, copy=True) + result.setflags(write=False) + return result + + def _validate_ids(self, env_ids: np.ndarray, *, capability: str) -> np.ndarray: + if not isinstance(env_ids, np.ndarray): + raise TypeError( + f"ManagerBased reset-state {capability} env_ids must be np.ndarray, " + f"got {type(env_ids).__name__}" + ) + if ( + env_ids.ndim != 1 + or not np.issubdtype(env_ids.dtype, np.integer) + or np.issubdtype(env_ids.dtype, np.bool_) + ): + raise TypeError( + f"ManagerBased reset-state {capability} env_ids must be a 1-D integer " + f"np.ndarray, got shape={env_ids.shape}, dtype={env_ids.dtype}" + ) + ids = np.asarray(env_ids, dtype=np.int32) + if np.any(ids < 0) or np.any(ids >= self._num_envs): + raise IndexError( + f"ManagerBased reset-state {capability} env_ids out of range for " + f"{self._num_envs} environments: {ids.tolist()}" + ) + if np.unique(ids).size != ids.size: + raise ValueError( + f"ManagerBased reset-state {capability} env_ids contain duplicates: {ids.tolist()}" + ) + return ids + + def _validate_columns( + self, + values: np.ndarray, + *, + width: int, + capability: str, + term_name: str, + ) -> np.ndarray: + if not isinstance(values, np.ndarray): + raise TypeError( + f"EventManager term '{term_name}' {capability} must be np.ndarray, " + f"got {type(values).__name__}" + ) + if ( + values.ndim != 1 + or not np.issubdtype(values.dtype, np.integer) + or np.issubdtype(values.dtype, np.bool_) + ): + raise TypeError( + f"EventManager term '{term_name}' {capability} must be a 1-D integer array" + ) + columns = np.asarray(values, dtype=np.intp) + if np.any(columns < 0) or np.any(columns >= width): + raise IndexError( + f"EventManager term '{term_name}' {capability} out of range for width " + f"{width}: {columns.tolist()}" + ) + if np.unique(columns).size != columns.size: + raise ValueError( + f"EventManager term '{term_name}' {capability} contain duplicates: " + f"{columns.tolist()}" + ) + return columns + + def _validate_values( + self, + values: np.ndarray, + *, + shape: tuple[int, int], + capability: str, + term_name: str, + ) -> np.ndarray: + if not isinstance(values, np.ndarray): + raise TypeError( + f"EventManager term '{term_name}' {capability} must be np.ndarray, " + f"got {type(values).__name__}" + ) + if values.shape != shape: + raise ValueError( + f"EventManager term '{term_name}' {capability} has shape {values.shape}; " + f"expected {shape}" + ) + if not np.issubdtype(values.dtype, np.floating): + raise TypeError( + f"EventManager term '{term_name}' {capability} must be floating, got {values.dtype}" + ) + if not np.isfinite(values).all(): + raise ValueError(f"EventManager term '{term_name}' {capability} contains NaN or Inf") + return values + + def _capability_error( + self, + term_name: str, + capability: str, + exc: BaseException, + ) -> NotImplementedError: + return NotImplementedError( + f"EventManager term '{term_name}' reset-state capability '{capability}' is " + f"unavailable on backend '{self._backend.backend_type}': {exc}" + ) + + def _require_active(self) -> None: + if not self._active: + raise RuntimeError("ManagerBased reset-state mutation requires an active reset event") + + def _finish(self) -> None: + self._active = False + self._active_mask.fill(False) + self._dirty_mask.fill(False) + self._requesting_terms.clear() + + +__all__ = ["ResetStateTransaction"] diff --git a/src/unilab/envs/manager_based_rl_env.py b/src/unilab/envs/manager_based_rl_env.py index e5d8c0f63..89742e3a9 100644 --- a/src/unilab/envs/manager_based_rl_env.py +++ b/src/unilab/envs/manager_based_rl_env.py @@ -18,6 +18,7 @@ from unilab.base.base import EnvCfg from unilab.base.entity import EntityScene from unilab.base.np_env import NpEnv, NpEnvState +from unilab.base.reset_state import ResetStateTransaction from unilab.base.scene import SceneCfg from unilab.dtype_config import get_global_dtype from unilab.managers import ( @@ -170,8 +171,14 @@ def __init__(self, cfg: ManagerBasedRlEnvCfg, backend: SimBackend, num_envs: int self.rng = np.random.default_rng(actual_seed) self._control = np.zeros((num_envs, backend.num_actuators), dtype=get_global_dtype()) + self._reset_state = ResetStateTransaction(backend) assert cfg.scene is not None - self.scene = EntityScene.from_scene_cfg(cfg.scene, backend, self._control) + self.scene = EntityScene.from_scene_cfg( + cfg.scene, + backend, + self._control, + reset_state=self._reset_state, + ) self.common_step_counter = 0 self._sim_step_counter = 0 @@ -431,11 +438,12 @@ def reset( log: dict[str, Any] = {} self.curriculum_manager.compute(env_ids=ids) if "reset" in self.event_manager.available_modes: - self.event_manager.apply( - mode="reset", - env_ids=ids, - global_env_step_count=self.step_counter, - ) + with self._reset_state.scoped(ids): + self.event_manager.apply( + mode="reset", + env_ids=ids, + global_env_step_count=self.step_counter, + ) for manager in ( self.observation_manager, diff --git a/src/unilab/envs/mdp/__init__.py b/src/unilab/envs/mdp/__init__.py index 925e49264..880cc9d14 100644 --- a/src/unilab/envs/mdp/__init__.py +++ b/src/unilab/envs/mdp/__init__.py @@ -4,6 +4,8 @@ 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 reset_scene_to_default as reset_scene_to_default +from unilab.envs.mdp.events import resolve_env_ids as resolve_env_ids from unilab.envs.mdp.observations import base_ang_vel as base_ang_vel from unilab.envs.mdp.observations import base_lin_vel as base_lin_vel from unilab.envs.mdp.observations import generated_commands as generated_commands @@ -48,6 +50,8 @@ "is_alive", "is_terminated", "projected_gravity", + "reset_scene_to_default", + "resolve_env_ids", "root_height_below_minimum", "time_out", "track_angular_velocity", diff --git a/src/unilab/envs/mdp/events.py b/src/unilab/envs/mdp/events.py new file mode 100644 index 000000000..568119270 --- /dev/null +++ b/src/unilab/envs/mdp/events.py @@ -0,0 +1,32 @@ +# Derived from mujocolab/mjlab v1.6.0 (0fb8a681), +# src/mjlab/envs/mdp/events.py. +# Copyright 2025, The mjlab Developers. +# Modified by UniLab for NumPy reset transactions; Apache-2.0. +"""Community-style reset event terms for UniLab's NumPy manager runtime.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np + +if TYPE_CHECKING: + from unilab.managers._types import ManagerBasedRlEnv + + +def resolve_env_ids(env: ManagerBasedRlEnv, env_ids: np.ndarray | None) -> np.ndarray: + """Return concrete NumPy environment IDs, preserving community sentinel semantics.""" + if env_ids is None: + return np.arange(env.num_envs, dtype=np.int32) + return env_ids + + +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) + if not env.scene.entities: + return + env.scene.reset_to_default(ids, term_name="reset_scene_to_default") + + +__all__ = ["reset_scene_to_default", "resolve_env_ids"] diff --git a/src/unilab/managers/_types.py b/src/unilab/managers/_types.py index eace62476..3aeb02b12 100644 --- a/src/unilab/managers/_types.py +++ b/src/unilab/managers/_types.py @@ -7,7 +7,7 @@ from __future__ import annotations -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from typing import Any, Protocol import numpy as np @@ -126,12 +126,28 @@ def find_pairs( self, keys: str | Sequence[str], preserve_order: bool = False ) -> tuple[list[int], list[str]]: ... + def write_joint_state_to_sim( + self, + position: np.ndarray, + velocity: np.ndarray, + joint_ids: np.ndarray | Sequence[int] | slice | None = None, + env_ids: np.ndarray | slice | None = None, + ) -> None: ... + class ManagerScene(Protocol): """Minimal name-addressable scene surface consumed by managers.""" + @property + def entities(self) -> Mapping[str, ManagerEntity]: ... + + @property + def env_origins(self) -> np.ndarray: ... + def __getitem__(self, name: str) -> ManagerEntity: ... + def reset_to_default(self, env_ids: np.ndarray, *, term_name: str) -> None: ... + class ManagerActionTerm(Protocol): @property diff --git a/tests/base/test_backend_conformance.py b/tests/base/test_backend_conformance.py index bee1f09e8..ad35d11dc 100644 --- a/tests/base/test_backend_conformance.py +++ b/tests/base/test_backend_conformance.py @@ -135,6 +135,10 @@ def test_actuation_metadata_defaults_fail_closed() -> None: SimBackend.get_actuator_joint_names(object()) # type: ignore[arg-type] with pytest.raises(NotImplementedError, match="default DoF positions"): SimBackend.get_default_dof_pos(object()) # type: ignore[arg-type] + with pytest.raises(NotImplementedError, match="get_joint_state_qpos_indices"): + SimBackend.get_joint_state_qpos_indices(object(), ("joint",)) # type: ignore[arg-type] + with pytest.raises(NotImplementedError, match="get_joint_state_qvel_indices"): + SimBackend.get_joint_state_qvel_indices(object(), ("joint",)) # type: ignore[arg-type] @pytest.mark.parametrize("backend_type", _BACKEND_PARAMS) @@ -153,6 +157,8 @@ def test_actuation_metadata_contract(backend_type: str) -> None: actuator_names = backend.get_actuator_names() target_joint_names = backend.get_actuator_joint_names() default_dof_pos = backend.get_default_dof_pos() + reset_qpos_ids = backend.get_joint_state_qpos_indices(target_joint_names) + reset_qvel_ids = backend.get_joint_state_qvel_indices(target_joint_names) assert len(actuator_names) == backend.num_actuators assert len(set(actuator_names)) == len(actuator_names) @@ -162,8 +168,25 @@ def test_actuation_metadata_contract(backend_type: str) -> None: assert default_dof_pos.shape == backend.get_dof_pos().shape[1:] assert np.issubdtype(default_dof_pos.dtype, np.floating) assert np.isfinite(default_dof_pos).all() + assert reset_qpos_ids.shape == (backend.num_actuators,) + assert reset_qvel_ids.shape == (backend.num_actuators,) + assert np.issubdtype(reset_qpos_ids.dtype, np.integer) + assert np.issubdtype(reset_qvel_ids.dtype, np.integer) + assert np.unique(reset_qpos_ids).size == reset_qpos_ids.size + assert np.unique(reset_qvel_ids).size == reset_qvel_ids.size + assert np.all((reset_qpos_ids >= 0) & (reset_qpos_ids < backend.get_default_qpos().size)) + assert np.all((reset_qvel_ids >= 0) & (reset_qvel_ids < backend.get_init_qvel().size)) np.testing.assert_allclose(default_dof_pos, backend.get_dof_pos()[0], atol=1e-6) + qpos = np.broadcast_to( + backend.get_default_qpos(), (NUM_ENVS, backend.get_default_qpos().size) + ).copy() + qvel = np.broadcast_to(backend.get_init_qvel(), (NUM_ENVS, backend.get_init_qvel().size)).copy() + expected_joint_pos = np.broadcast_to(default_dof_pos + 0.01, (NUM_ENVS, default_dof_pos.size)) + qpos[:, reset_qpos_ids] = expected_joint_pos + backend.set_state(np.arange(NUM_ENVS, dtype=np.int32), qpos, qvel) + np.testing.assert_allclose(backend.get_dof_pos(), expected_joint_pos, atol=1e-5) + detached = default_dof_pos.copy() default_dof_pos[:] = np.nan np.testing.assert_array_equal(backend.get_default_dof_pos(), detached) diff --git a/tests/base/test_entity_facade.py b/tests/base/test_entity_facade.py index ed048f1eb..08f4c686f 100644 --- a/tests/base/test_entity_facade.py +++ b/tests/base/test_entity_facade.py @@ -494,3 +494,16 @@ def test_scene_cfg_entity_defaults_are_not_shared() -> None: second = SceneCfg(model_file="second.xml") first.entities["robot"] = EntityCfg() assert second.entities == {} + + +def test_scene_exposes_read_only_community_entities_and_zero_origins() -> None: + _, scene = _scene() + + assert scene.entities["robot"] is scene["robot"] + with pytest.raises(TypeError): + scene.entities["other"] = scene["robot"] # type: ignore[index] + + assert scene.env_origins.shape == (3, 3) + np.testing.assert_array_equal(scene.env_origins, 0.0) + with pytest.raises(ValueError, match="read-only"): + scene.env_origins[0, 0] = 1.0 diff --git a/tests/base/test_reset_state.py b/tests/base/test_reset_state.py new file mode 100644 index 000000000..38af6d9e1 --- /dev/null +++ b/tests/base/test_reset_state.py @@ -0,0 +1,233 @@ +"""Focused tests for the base-owned Manager-Based reset transaction.""" + +from __future__ import annotations + +from typing import Any, cast + +import numpy as np +import pytest + +from unilab.base.backend.base import SimBackend +from unilab.base.reset_state import ResetStateTransaction + + +class _Backend: + backend_type = "fake" + + def __init__( + self, + *, + qpos: Any = None, + qvel: Any = None, + fail_set_state: bool = False, + ) -> None: + self.num_envs = 4 + self.qpos = np.array([1.0, 2.0, 3.0]) if qpos is None else qpos + self.qvel = np.array([0.0, 0.0]) if qvel is None else qvel + self.fail_set_state = fail_set_state + self.default_qpos_calls = 0 + self.init_qvel_calls = 0 + self.set_state_calls: list[tuple[np.ndarray, np.ndarray, np.ndarray]] = [] + + def get_default_qpos(self): + self.default_qpos_calls += 1 + return self.qpos + + def get_init_qvel(self): + self.init_qvel_calls += 1 + return self.qvel + + def set_state( + self, + env_ids: np.ndarray, + qpos: np.ndarray, + 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())) + return {"timing": {"set_state_ms": 1.0}} + + +def _transaction(backend: _Backend) -> ResetStateTransaction: + return ResetStateTransaction(cast(SimBackend, backend)) + + +def test_transaction_is_lazy_and_combines_terms_into_one_commit() -> None: + backend = _Backend() + transaction = _transaction(backend) + + with transaction.scoped(np.array([0, 2, 3], dtype=np.int32)): + assert transaction.active + assert backend.default_qpos_calls == 0 + assert backend.init_qvel_calls == 0 + assert backend.set_state_calls == [] + + with transaction.scoped(np.array([0, 2, 3], dtype=np.int32)): + transaction.reset_to_default( + np.array([2], dtype=np.int32), + term_name="first", + ) + transaction.reset_to_default( + np.array([3, 0], dtype=np.int32), + term_name="second", + ) + assert backend.set_state_calls == [] + + assert not transaction.active + assert backend.default_qpos_calls == 1 + assert backend.init_qvel_calls == 1 + assert len(backend.set_state_calls) == 1 + ids, qpos, qvel = backend.set_state_calls[0] + np.testing.assert_array_equal(ids, [0, 2, 3]) + np.testing.assert_array_equal(qpos, np.tile(backend.qpos, (3, 1))) + np.testing.assert_array_equal(qvel, np.tile(backend.qvel, (3, 1))) + + with transaction.scoped(np.array([1], dtype=np.int32)): + transaction.reset_to_default(np.array([1], dtype=np.int32), term_name="third") + assert backend.default_qpos_calls == 1 + assert backend.init_qvel_calls == 1 + assert len(backend.set_state_calls) == 2 + + +def test_exception_aborts_without_backend_mutation_and_next_reset_is_clean() -> None: + backend = _Backend() + transaction = _transaction(backend) + + with pytest.raises(RuntimeError, match="term failed"): + with transaction.scoped(np.array([0, 1], dtype=np.int32)): + transaction.reset_to_default(np.array([0], dtype=np.int32), term_name="broken") + raise RuntimeError("term failed") + + assert not transaction.active + assert backend.set_state_calls == [] + + with transaction.scoped(np.array([1], dtype=np.int32)): + transaction.reset_to_default(np.array([1], dtype=np.int32), term_name="healthy") + assert len(backend.set_state_calls) == 1 + np.testing.assert_array_equal(backend.set_state_calls[0][0], [1]) + + +def test_joint_writes_initialize_defaults_and_compose_by_column() -> None: + backend = _Backend() + transaction = _transaction(backend) + + with transaction.scoped(np.array([0, 2], dtype=np.int32)): + transaction.write_joint_state( + np.array([2, 0], dtype=np.int32), + np.array([1], dtype=np.int32), + np.array([0], dtype=np.int32), + np.array([[9.0], [8.0]], dtype=np.float32), + np.array([[-1.0], [-2.0]], dtype=np.float32), + term_name="robot.write_joint_state_to_sim", + ) + + ids, qpos, qvel = backend.set_state_calls[0] + np.testing.assert_array_equal(ids, [0, 2]) + np.testing.assert_array_equal(qpos, [[1.0, 8.0, 3.0], [1.0, 9.0, 3.0]]) + np.testing.assert_array_equal(qvel, [[-2.0, 0.0], [-1.0, 0.0]]) + + +@pytest.mark.parametrize( + ("position", "velocity", "error", "match"), + [ + (np.zeros((1, 2)), np.zeros((1, 1)), ValueError, "joint position.*expected"), + (np.zeros((1, 1)), np.zeros((2, 1)), ValueError, "joint velocity.*expected"), + (np.zeros((1, 1), dtype=np.int32), np.zeros((1, 1)), TypeError, "must be floating"), + (np.full((1, 1), np.nan), np.zeros((1, 1)), ValueError, "NaN or Inf"), + ], +) +def test_joint_write_values_fail_closed(position, velocity, error, match: str) -> None: + transaction = _transaction(_Backend()) + with pytest.raises(error, match=match): + with transaction.scoped(np.array([0], dtype=np.int32)): + transaction.write_joint_state( + np.array([0], dtype=np.int32), + np.array([1], dtype=np.int32), + np.array([0], dtype=np.int32), + position, + velocity, + term_name="joint_term", + ) + + +def test_mutation_must_stay_inside_active_reset() -> None: + transaction = _transaction(_Backend()) + with transaction.scoped(np.array([1, 2], dtype=np.int32)): + with pytest.raises(ValueError, match="outside the active reset.*3"): + transaction.reset_to_default(np.array([3], dtype=np.int32), term_name="bad") + + with pytest.raises(RuntimeError, match="requires an active reset event"): + transaction.reset_to_default(np.array([1], dtype=np.int32), term_name="late") + + +@pytest.mark.parametrize( + ("ids", "error", "match"), + [ + ([0], TypeError, "must be np.ndarray"), + (np.array([[0]], dtype=np.int32), TypeError, "1-D integer"), + (np.array([True]), TypeError, "1-D integer"), + (np.array([-1], dtype=np.int32), IndexError, "out of range"), + (np.array([4], dtype=np.int32), IndexError, "out of range"), + (np.array([1, 1], dtype=np.int32), ValueError, "duplicates"), + ], +) +def test_begin_rejects_invalid_environment_ids(ids, error, match: str) -> None: + with pytest.raises(error, match=match): + _transaction(_Backend()).begin(ids) + + +@pytest.mark.parametrize( + ("field", "value", "error", "match"), + [ + ("qpos", [1.0], TypeError, "default qpos.*np.ndarray"), + ("qpos", np.zeros((1, 1)), ValueError, "default qpos.*expected 1-D"), + ("qpos", np.array([1], dtype=np.int32), TypeError, "default qpos.*floating"), + ("qpos", np.array([np.nan]), ValueError, "default qpos.*NaN or Inf"), + ("qvel", np.array([np.inf]), ValueError, "initial qvel.*NaN or Inf"), + ], +) +def test_backend_default_state_contract_fails_at_mutation_boundary( + field: str, + value, + error, + match: str, +) -> None: + kwargs = {field: value} + transaction = _transaction(_Backend(**kwargs)) + with pytest.raises(error, match=match): + with transaction.scoped(np.array([0], dtype=np.int32)): + transaction.reset_to_default( + np.array([0], dtype=np.int32), + term_name="reset_scene_to_default", + ) + assert not transaction.active + + +def test_missing_default_and_set_state_capabilities_name_term_and_backend() -> None: + missing = type("MissingBackend", (), {"num_envs": 1, "backend_type": "missing"})() + transaction = ResetStateTransaction(cast(SimBackend, missing)) + with pytest.raises( + NotImplementedError, + match="EventManager term 'reset_scene_to_default'.*default qpos.*backend 'missing'", + ): + with transaction.scoped(np.array([0], dtype=np.int32)): + transaction.reset_to_default( + np.array([0], dtype=np.int32), + term_name="reset_scene_to_default", + ) + + backend = _Backend(fail_set_state=True) + transaction = _transaction(backend) + with pytest.raises( + NotImplementedError, + match="SimBackend.set_state.*reset_scene_to_default.*backend 'fake'", + ): + with transaction.scoped(np.array([0], dtype=np.int32)): + transaction.reset_to_default( + np.array([0], dtype=np.int32), + term_name="reset_scene_to_default", + ) + assert not transaction.active diff --git a/tests/envs/test_manager_based_rl_env.py b/tests/envs/test_manager_based_rl_env.py index 448163b1a..3a2b77250 100644 --- a/tests/envs/test_manager_based_rl_env.py +++ b/tests/envs/test_manager_based_rl_env.py @@ -17,6 +17,7 @@ ManagerBasedRlEnv, ManagerBasedRLEnvCfg, ManagerBasedRlEnvCfg, + mdp, ) from unilab.managers import ( ActionTerm, @@ -74,6 +75,63 @@ def cleanup_scene_assets(self) -> None: self.cleanup_calls += 1 +class _ResetBackend(_FakeBackend): + def __init__(self, num_envs: int) -> None: + super().__init__(num_envs) + self.default_qpos_calls = 0 + self.init_qvel_calls = 0 + self.joint_layout_calls = 0 + self.set_state_calls: list[tuple[np.ndarray, np.ndarray, np.ndarray]] = [] + + def get_default_qpos(self) -> np.ndarray: + self.default_qpos_calls += 1 + return np.array([0.0, 0.0, 0.5, 0.0], dtype=np.float64) + + def get_init_qvel(self) -> np.ndarray: + self.init_qvel_calls += 1 + return np.array([0.0, 0.0, 0.0], dtype=np.float32) + + def get_actuator_joint_names(self) -> tuple[str, ...]: + return ("joint",) + + def get_joint_dof_pos_indices(self, names) -> np.ndarray: + assert tuple(names) == ("joint",) + return np.array([0], dtype=np.int32) + + def get_joint_dof_vel_indices(self, names) -> np.ndarray: + assert tuple(names) == ("joint",) + return np.array([0], dtype=np.int32) + + def get_joint_state_qpos_indices(self, names) -> np.ndarray: + assert tuple(names) == ("joint",) + self.joint_layout_calls += 1 + return np.array([3], dtype=np.int32) + + def get_joint_state_qvel_indices(self, names) -> np.ndarray: + assert tuple(names) == ("joint",) + self.joint_layout_calls += 1 + return np.array([2], dtype=np.int32) + + def get_dof_pos(self) -> np.ndarray: + return np.zeros((self.num_envs, 1), dtype=np.float32) + + def get_default_dof_pos(self) -> np.ndarray: + return np.zeros((1,), dtype=np.float32) + + def get_dof_vel(self) -> np.ndarray: + return np.zeros((self.num_envs, 1), dtype=np.float32) + + def set_state( + self, + env_ids: np.ndarray, + qpos: np.ndarray, + qvel: np.ndarray, + randomization=None, + ) -> None: + assert randomization is None + self.set_state_calls.append((env_ids.copy(), qpos.copy(), qvel.copy())) + + @dataclass(kw_only=True) class _DriveCfg(ActionTermCfg): gain: float = 1.0 @@ -188,6 +246,22 @@ def _event(env: _TestEnv, env_ids: np.ndarray | None) -> None: env.trace.append(("event", rendered_ids)) +def _observe_uncommitted_reset(env: _TestEnv, env_ids: np.ndarray | None) -> None: + del env_ids + backend = cast(_ResetBackend, env._backend) + env.trace.append(("commit_count_during_event", len(backend.set_state_calls))) + + +def _write_reset_joint_state(env: _TestEnv, env_ids: np.ndarray | None) -> None: + assert env_ids is not None + count = len(env_ids) + env.scene["robot"].write_joint_state_to_sim( + np.full((count, 1), 0.25, dtype=np.float32), + np.full((count, 1), -0.5, dtype=np.float32), + env_ids=env_ids, + ) + + def _make_cfg( *, sim_substeps: int = 2, @@ -342,6 +416,57 @@ def test_np_env_owns_substeps_autoreset_and_final_observation() -> None: assert pre_index < post_reset_index < post_step_index +def test_reset_events_compose_then_commit_default_state_once() -> None: + cfg = _make_cfg(include_optional_managers=False) + cfg.events = { + "default_first": EventTermCfg(func=mdp.reset_scene_to_default, mode="reset"), + "joint_state": EventTermCfg(func=_write_reset_joint_state, mode="reset"), + "observe_uncommitted": EventTermCfg(func=_observe_uncommitted_reset, mode="reset"), + } + cfg.scene.entities["robot"] = EntityCfg( + joint_names=("joint",), + actuator_names=("motor",), + ) + backend = _ResetBackend(2) + env = _TestEnv(cfg, cast(SimBackend, backend), 2) + + assert backend.default_qpos_calls == 0 + assert backend.init_qvel_calls == 0 + assert backend.joint_layout_calls == 0 + env.reset() + + assert env.trace[-1] == ("commit_count_during_event", 0) + assert backend.default_qpos_calls == 1 + assert backend.init_qvel_calls == 1 + assert len(backend.set_state_calls) == 1 + ids, qpos, qvel = backend.set_state_calls[0] + np.testing.assert_array_equal(ids, [0, 1]) + np.testing.assert_array_equal( + qpos, + [[0.0, 0.0, 0.5, 0.25], [0.0, 0.0, 0.5, 0.25]], + ) + np.testing.assert_array_equal(qvel, [[0.0, 0.0, -0.5], [0.0, 0.0, -0.5]]) + assert backend.joint_layout_calls == 2 + + env.reset(env_ids=np.array([1], dtype=np.int32)) + assert ("commit_count_during_event", 1) in env.trace + assert len(backend.set_state_calls) == 2 + np.testing.assert_array_equal(backend.set_state_calls[1][0], [1]) + assert backend.default_qpos_calls == 1 + assert backend.init_qvel_calls == 1 + assert backend.joint_layout_calls == 2 + + +def test_pure_reset_event_does_not_request_backend_state_capability() -> None: + env, backend = _make_env() + + env.reset() + env.reset(env_ids=np.array([0], dtype=np.int32)) + + assert isinstance(backend, _FakeBackend) + assert not hasattr(backend, "get_default_qpos") + + def test_partial_reset_preserves_other_env_counter_and_terminal_obs() -> None: env, _ = _make_env() env.init_state() From c5e1a2f063567ab6d7bc0b5f1af0e4f5f1b86039 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 01:59:33 +0800 Subject: [PATCH 012/169] feat(config): add manager term overlays (#1069) --- src/unilab/base/config_overrides.py | 11 ++ src/unilab/base/registry.py | 103 +++++++++++++ src/unilab/envs/manager_based_rl_env.py | 29 ++-- src/unilab/managers/manager_base.py | 9 +- src/unilab/managers/observation_manager.py | 10 +- tests/base/test_manager_config_overlay.py | 165 +++++++++++++++++++++ 6 files changed, 315 insertions(+), 12 deletions(-) create mode 100644 src/unilab/base/config_overrides.py create mode 100644 tests/base/test_manager_config_overlay.py diff --git a/src/unilab/base/config_overrides.py b/src/unilab/base/config_overrides.py new file mode 100644 index 000000000..62d16d572 --- /dev/null +++ b/src/unilab/base/config_overrides.py @@ -0,0 +1,11 @@ +"""Metadata shared by config owners and the registry override engine.""" + +CONFIG_MAPPING_POLICY_KEY = "unilab_config_mapping_policy" +MANAGER_TERM_MAPPING_POLICY = "manager_terms" +MANAGER_PARAMS_MAPPING_POLICY = "manager_params" + +__all__ = [ + "CONFIG_MAPPING_POLICY_KEY", + "MANAGER_PARAMS_MAPPING_POLICY", + "MANAGER_TERM_MAPPING_POLICY", +] diff --git a/src/unilab/base/registry.py b/src/unilab/base/registry.py index 21e85bc1b..1de9ca61c 100644 --- a/src/unilab/base/registry.py +++ b/src/unilab/base/registry.py @@ -18,6 +18,11 @@ ) from .base import ABEnv, EnvCfg +from .config_overrides import ( + CONFIG_MAPPING_POLICY_KEY, + MANAGER_PARAMS_MAPPING_POLICY, + MANAGER_TERM_MAPPING_POLICY, +) TEnvCfg = TypeVar("TEnvCfg", bound=EnvCfg) _SUPPORTED_SIM_BACKENDS = ("mujoco", "mjwarp", "motrix", "drake") @@ -162,6 +167,86 @@ def _construct_dataclass_from_dict(target_type: Type[Any], values: Dict[str, Any return target_obj +def _config_mapping_policy(target_obj: Any, field_name: str) -> str | None: + if not dataclasses.is_dataclass(target_obj) or isinstance(target_obj, type): + return None + for config_field in dataclasses.fields(target_obj): + if config_field.name == field_name: + policy = config_field.metadata.get(CONFIG_MAPPING_POLICY_KEY) + return str(policy) if policy is not None else None + return None + + +def _is_manager_callable_term_cfg(target_obj: Any) -> bool: + if not dataclasses.is_dataclass(target_obj) or isinstance(target_obj, type): + return False + return any( + config_field.metadata.get(CONFIG_MAPPING_POLICY_KEY) == MANAGER_PARAMS_MAPPING_POLICY + for config_field in dataclasses.fields(target_obj) + ) + + +def _apply_manager_mapping_overrides( + target_obj: Any, + field_name: str, + existing: Any, + overrides: Any, + *, + policy: str, +) -> None: + owner = f"{type(target_obj).__name__}.{field_name}" + if not isinstance(existing, dict): + raise TypeError( + f"Config field '{owner}' declares manager mapping policy but contains " + f"{type(existing).__name__}, expected dict" + ) + if not isinstance(overrides, dict): + raise TypeError( + f"Config field '{owner}' must be overridden by a mapping, not " + f"{type(overrides).__name__}" + ) + + if policy == MANAGER_PARAMS_MAPPING_POLICY: + for param_name, value in overrides.items(): + current = existing.get(param_name) + if isinstance(value, dict) and dataclasses.is_dataclass(current): + apply_cfg_overrides(current, value) + else: + existing[param_name] = value + return + + if policy != MANAGER_TERM_MAPPING_POLICY: + raise ValueError(f"Config field '{owner}' has unknown mapping policy {policy!r}") + + for term_name, value in overrides.items(): + if term_name not in existing: + raise ValueError( + f"Config field '{owner}' has no factory-owned term '{term_name}'; " + "declare its callable/config in the task Python factory first" + ) + if value is None: + existing[term_name] = None + continue + + current = existing[term_name] + if current is None: + raise ValueError( + f"Config field '{owner}' term '{term_name}' is disabled; set a concrete " + "config in the task Python factory before overriding its fields" + ) + if not dataclasses.is_dataclass(current): + raise TypeError( + f"Config field '{owner}' term '{term_name}' contains " + f"{type(current).__name__}, expected a dataclass config" + ) + if not isinstance(value, dict): + raise TypeError( + f"Config field '{owner}' term '{term_name}' must be overridden by a " + "field mapping or None; replacing the factory-owned term is not allowed" + ) + apply_cfg_overrides(current, value) + + def apply_cfg_overrides(target_obj: Any, overrides: Dict[str, Any]) -> None: """Apply a (possibly nested) dict of overrides to ``target_obj`` in place. @@ -176,6 +261,9 @@ def apply_cfg_overrides(target_obj: Any, overrides: Dict[str, Any]) -> None: - If ``value`` is a dict and ``target_obj.key`` is currently ``None``, instantiate the field's annotated dataclass type from the dict (full-construction path). + - Fields explicitly marked as manager mappings merge only existing + factory-owned entries. ``None`` disables an entry; unknown entries and + callable/config replacement fail closed. - Otherwise ``setattr`` the value directly (scalar / list / non-dataclass). """ try: @@ -187,6 +275,21 @@ def apply_cfg_overrides(target_obj: Any, overrides: Dict[str, Any]) -> None: if not hasattr(target_obj, key): raise ValueError(f"Config class '{type(target_obj).__name__}' has no attribute '{key}'") existing = getattr(target_obj, key) + mapping_policy = _config_mapping_policy(target_obj, key) + if mapping_policy is not None: + _apply_manager_mapping_overrides( + target_obj, + key, + existing, + value, + policy=mapping_policy, + ) + continue + if key == "func" and _is_manager_callable_term_cfg(target_obj): + raise ValueError( + f"Config field '{type(target_obj).__name__}.func' is factory-owned and " + "cannot be overridden" + ) if isinstance(value, dict): if dataclasses.is_dataclass(existing) and not isinstance(existing, type): apply_cfg_overrides(existing, value) diff --git a/src/unilab/envs/manager_based_rl_env.py b/src/unilab/envs/manager_based_rl_env.py index 89742e3a9..9999e65aa 100644 --- a/src/unilab/envs/manager_based_rl_env.py +++ b/src/unilab/envs/manager_based_rl_env.py @@ -16,6 +16,10 @@ from unilab.base.backend import SimBackend from unilab.base.base import EnvCfg +from unilab.base.config_overrides import ( + CONFIG_MAPPING_POLICY_KEY, + MANAGER_TERM_MAPPING_POLICY, +) from unilab.base.entity import EntityScene from unilab.base.np_env import NpEnv, NpEnvState from unilab.base.reset_state import ResetStateTransaction @@ -47,6 +51,13 @@ ) +def _manager_terms_field() -> Any: + return field( + default_factory=dict, + metadata={CONFIG_MAPPING_POLICY_KEY: MANAGER_TERM_MAPPING_POLICY}, + ) + + @dataclass class ManagerBasedRlEnvCfg(EnvCfg): """Configuration for the manager-based NumPy environment. @@ -55,15 +66,15 @@ class ManagerBasedRlEnvCfg(EnvCfg): configs can overlay them without introducing a second configuration runtime. """ - observations: dict[str, ObservationGroupCfg | None] = field(default_factory=dict) - actions: dict[str, ActionTermCfg | None] = field(default_factory=dict) - events: dict[str, EventTermCfg | None] = field(default_factory=dict) - rewards: dict[str, RewardTermCfg | None] = field(default_factory=dict) - terminations: dict[str, TerminationTermCfg | None] = field(default_factory=dict) - commands: dict[str, CommandTermCfg | None] = field(default_factory=dict) - curriculum: dict[str, CurriculumTermCfg | None] = field(default_factory=dict) - metrics: dict[str, MetricsTermCfg | None] = field(default_factory=dict) - recorders: dict[str, RecorderTermCfg | None] = field(default_factory=dict) + observations: dict[str, ObservationGroupCfg | None] = _manager_terms_field() + actions: dict[str, ActionTermCfg | None] = _manager_terms_field() + events: dict[str, EventTermCfg | None] = _manager_terms_field() + rewards: dict[str, RewardTermCfg | None] = _manager_terms_field() + terminations: dict[str, TerminationTermCfg | None] = _manager_terms_field() + commands: dict[str, CommandTermCfg | None] = _manager_terms_field() + curriculum: dict[str, CurriculumTermCfg | None] = _manager_terms_field() + metrics: dict[str, MetricsTermCfg | None] = _manager_terms_field() + recorders: dict[str, RecorderTermCfg | None] = _manager_terms_field() seed: int | None = None is_finite_horizon: bool = False diff --git a/src/unilab/managers/manager_base.py b/src/unilab/managers/manager_base.py index dcf9fee16..0066ffcce 100644 --- a/src/unilab/managers/manager_base.py +++ b/src/unilab/managers/manager_base.py @@ -11,6 +11,10 @@ import numpy as np +from unilab.base.config_overrides import ( + CONFIG_MAPPING_POLICY_KEY, + MANAGER_PARAMS_MAPPING_POLICY, +) from unilab.managers.scene_entity_config import SceneEntityCfg if TYPE_CHECKING: @@ -60,7 +64,10 @@ def __call__(self, env, **kwargs) -> np.ndarray: """The callable that computes this term's value. Can be a function or a class. Classes are auto-instantiated with ``(cfg=term_cfg, env=env)``.""" - params: dict[str, Any] = field(default_factory=lambda: {}) + params: dict[str, Any] = field( + default_factory=dict, + metadata={CONFIG_MAPPING_POLICY_KEY: MANAGER_PARAMS_MAPPING_POLICY}, + ) """Additional keyword arguments passed to func when called.""" diff --git a/src/unilab/managers/observation_manager.py b/src/unilab/managers/observation_manager.py index b5c531ca3..7acfbc7ab 100644 --- a/src/unilab/managers/observation_manager.py +++ b/src/unilab/managers/observation_manager.py @@ -6,12 +6,16 @@ from __future__ import annotations from copy import deepcopy -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import TYPE_CHECKING, Literal, Sequence import numpy as np from prettytable import PrettyTable +from unilab.base.config_overrides import ( + CONFIG_MAPPING_POLICY_KEY, + MANAGER_TERM_MAPPING_POLICY, +) from unilab.managers._buffers import CircularBuffer, DelayBuffer from unilab.managers._noise import noise_cfg, noise_model from unilab.managers._noise.noise_cfg import NoiseCfg, NoiseModelCfg @@ -81,7 +85,9 @@ class ObservationGroupCfg: for the actor, "critic" for the value function). """ - terms: dict[str, ObservationTermCfg | None] + terms: dict[str, ObservationTermCfg | None] = field( + metadata={CONFIG_MAPPING_POLICY_KEY: MANAGER_TERM_MAPPING_POLICY} + ) """Dictionary mapping term names to their configurations.""" concatenate_terms: bool = True diff --git a/tests/base/test_manager_config_overlay.py b/tests/base/test_manager_config_overlay.py new file mode 100644 index 000000000..5cb1e9087 --- /dev/null +++ b/tests/base/test_manager_config_overlay.py @@ -0,0 +1,165 @@ +from __future__ import annotations + +from dataclasses import dataclass, field + +import numpy as np +import pytest + +from unilab.base.registry import apply_cfg_overrides +from unilab.envs.manager_based_rl_env import ManagerBasedRlEnvCfg +from unilab.envs.mdp import JointPositionActionCfg +from unilab.managers import ( + EventTermCfg, + ObservationGroupCfg, + ObservationTermCfg, + RewardTermCfg, + SceneEntityCfg, +) + + +def _first_term(_env, *, command_name: str, std: float) -> np.ndarray: + del command_name, std + return np.zeros(1, dtype=np.float32) + + +def _second_term(_env) -> np.ndarray: + return np.zeros(1, dtype=np.float32) + + +def _manager_cfg() -> ManagerBasedRlEnvCfg: + return ManagerBasedRlEnvCfg( + observations={ + "policy": ObservationGroupCfg( + terms={ + "first": ObservationTermCfg(func=_first_term), + "second": ObservationTermCfg(func=_second_term), + } + ) + }, + actions={ + "joint_pos": JointPositionActionCfg( + entity_name="robot", + actuator_names=(".*",), + scale=0.25, + ) + }, + events={ + "reset": EventTermCfg( + func=_first_term, + mode="reset", + params={"command_name": "twist", "std": 0.5}, + ) + }, + rewards={ + "tracking": RewardTermCfg( + func=_first_term, + weight=1.0, + params={ + "command_name": "twist", + "std": 0.5, + "asset_cfg": SceneEntityCfg("robot", joint_names=".*"), + }, + ), + "alive": RewardTermCfg(func=_second_term, weight=0.1), + "disabled": None, + }, + ) + + +def test_manager_mapping_overlay_preserves_factory_terms_and_order() -> None: + cfg = _manager_cfg() + reward_func = cfg.rewards["tracking"].func + + apply_cfg_overrides( + cfg, + { + "rewards": { + "tracking": { + "weight": 2.0, + "params": {"std": 0.25, "asset_cfg": {"preserve_order": True}}, + }, + "alive": None, + }, + "events": {"reset": {"min_step_count_between_reset": 3}}, + "actions": {"joint_pos": {"scale": 0.4}}, + }, + ) + + tracking = cfg.rewards["tracking"] + assert tracking is not None + assert tracking.func is reward_func + assert tracking.weight == pytest.approx(2.0) + assert tracking.params["command_name"] == "twist" + assert tracking.params["std"] == pytest.approx(0.25) + assert tracking.params["asset_cfg"].joint_names == ".*" + assert tracking.params["asset_cfg"].preserve_order is True + assert cfg.rewards["alive"] is None + assert list(cfg.rewards) == ["tracking", "alive", "disabled"] + reset = cfg.events["reset"] + assert reset is not None + assert reset.func is _first_term + assert reset.min_step_count_between_reset == 3 + action = cfg.actions["joint_pos"] + assert action is not None + assert action.entity_name == "robot" + assert action.actuator_names == (".*",) + assert action.scale == pytest.approx(0.4) + + +def test_observation_group_and_term_overlay_preserve_siblings() -> None: + cfg = _manager_cfg() + + apply_cfg_overrides( + cfg, + { + "observations": { + "policy": { + "enable_corruption": True, + "terms": {"first": {"scale": 2.0}, "second": None}, + } + } + }, + ) + + policy = cfg.observations["policy"] + assert policy is not None + assert policy.enable_corruption is True + first = policy.terms["first"] + assert first is not None + assert first.func is _first_term + assert first.scale == pytest.approx(2.0) + assert policy.terms["second"] is None + assert list(policy.terms) == ["first", "second"] + + +@pytest.mark.parametrize( + ("overrides", "match"), + [ + ({"rewards": {"missing": {"weight": 1.0}}}, "rewards.*missing"), + ({"rewards": {"disabled": {"weight": 1.0}}}, "disabled.*task Python factory"), + ({"rewards": {"tracking": _second_term}}, "tracking.*replacing"), + ({"rewards": []}, "rewards.*mapping"), + ({"rewards": {"tracking": {"func": _second_term}}}, "func.*factory-owned"), + ], +) +def test_manager_mapping_overlay_fails_closed(overrides: dict, match: str) -> None: + with pytest.raises((TypeError, ValueError), match=match): + apply_cfg_overrides(_manager_cfg(), overrides) + + +@dataclass +class _LegacyCfg: + reward_config: dict[str, object] = field( + default_factory=lambda: { + "scales": {"tracking": 1.0, "alive": 0.1}, + "tracking_sigma": 0.25, + } + ) + + +def test_legacy_plain_dict_keeps_replacement_semantics() -> None: + cfg = _LegacyCfg() + + apply_cfg_overrides(cfg, {"reward_config": {"scales": {"alive": 1.0}}}) + + assert cfg.reward_config == {"scales": {"alive": 1.0}} From 31f25c1b9705d9fc665ca8ba9ce88f7c2cbe03af Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 02:24:41 +0800 Subject: [PATCH 013/169] feat(config): route manager reward overrides (#1071) --- src/unilab/base/registry.py | 40 ++++++ src/unilab/training/backend_adapter.py | 15 ++- src/unilab/training/reward.py | 14 +- tests/config/test_manager_reward_routing.py | 139 ++++++++++++++++++++ 4 files changed, 203 insertions(+), 5 deletions(-) create mode 100644 tests/config/test_manager_reward_routing.py diff --git a/src/unilab/base/registry.py b/src/unilab/base/registry.py index 1de9ca61c..5badd68b5 100644 --- a/src/unilab/base/registry.py +++ b/src/unilab/base/registry.py @@ -8,6 +8,7 @@ Any, Callable, Dict, + Literal, Optional, Type, TypeVar, @@ -25,6 +26,7 @@ ) TEnvCfg = TypeVar("TEnvCfg", bound=EnvCfg) +RewardOverrideField = Literal["reward_config", "rewards"] _SUPPORTED_SIM_BACKENDS = ("mujoco", "mjwarp", "motrix", "drake") _DEFAULT_SIM_BACKEND_ORDER: tuple[str, ...] = ("mujoco", "motrix") _REGISTRY_MODULES_ATTR = "__unilab_registry_modules__" @@ -141,6 +143,44 @@ def find_available_sim_backend(env_name: str) -> str: return backend +def resolve_reward_override_field(env_name: str) -> RewardOverrideField: + """Resolve the Hydra root reward target declared by an env config owner. + + Legacy configs own a ``reward_config`` field. Manager-Based configs opt in + through the explicit manager-term mapping metadata on ``rewards``. The + registry resolves this on the config class without constructing an env or + backend so training adapters do not branch on task names. + """ + if env_name not in _envs: + raise ValueError(f"Environment '{env_name}' is not registered.") + + config_cls = _envs[env_name].env_cfg_cls + config_fields = { + config_field.name: config_field for config_field in dataclasses.fields(config_cls) + } + rewards_field = config_fields.get("rewards") + has_manager_rewards = ( + rewards_field is not None + and rewards_field.metadata.get(CONFIG_MAPPING_POLICY_KEY) == MANAGER_TERM_MAPPING_POLICY + ) + has_legacy_rewards = "reward_config" in config_fields + + if has_manager_rewards and has_legacy_rewards: + raise ValueError( + f"Environment '{env_name}' config owner '{config_cls.__name__}' declares both " + "Manager-Based 'rewards' and legacy 'reward_config' targets" + ) + if has_manager_rewards: + return "rewards" + if has_legacy_rewards: + return "reward_config" + raise ValueError( + f"Environment '{env_name}' config owner '{config_cls.__name__}' declares no " + "supported Hydra root reward target; expected legacy 'reward_config' or an " + "explicitly marked Manager-Based 'rewards' field" + ) + + def _resolve_dataclass_type(type_hint: Any) -> Optional[Type[Any]]: """Strip Optional/Union and return the underlying dataclass type, or None.""" if type_hint is None: diff --git a/src/unilab/training/backend_adapter.py b/src/unilab/training/backend_adapter.py index aea4687d5..fc16ce948 100644 --- a/src/unilab/training/backend_adapter.py +++ b/src/unilab/training/backend_adapter.py @@ -7,6 +7,7 @@ from omegaconf import DictConfig, OmegaConf +from unilab.base import registry from unilab.base.backend import materialize_scene_visual_override from unilab.base.scene import SceneCfg from unilab.training.reward import extract_reward_config @@ -30,8 +31,18 @@ def __init__( def build_task_env_cfg_override(self) -> dict[str, Any]: """Build env_cfg_override from the resolved reward + env sections.""" - env_cfg_override = extract_reward_config(self.cfg) - env_cfg_override.update(self._to_plain_dict(getattr(self.cfg, "env", None))) + registry.ensure_registries() + task_name = str(self.cfg.training.task_name) + reward_target = registry.resolve_reward_override_field(task_name) + env_overrides = self._to_plain_dict(getattr(self.cfg, "env", None)) + if reward_target in env_overrides: + raise ValueError( + f"Task '{task_name}' declares both Hydra root 'reward' and " + f"'env.{reward_target}'; use the root reward owner only" + ) + + env_cfg_override = extract_reward_config(self.cfg, target_field=reward_target) + env_cfg_override.update(env_overrides) return env_cfg_override diff --git a/src/unilab/training/reward.py b/src/unilab/training/reward.py index ee68d18fc..a1d2ff656 100644 --- a/src/unilab/training/reward.py +++ b/src/unilab/training/reward.py @@ -39,16 +39,24 @@ def resolve_reward_dict(cfg: DictConfig) -> RewardDict: return reward_dict -def extract_reward_config(cfg: DictConfig) -> dict[str, RewardDict]: +def extract_reward_config( + cfg: DictConfig, + *, + target_field: str = "reward_config", +) -> dict[str, RewardDict]: """Extract and validate reward config from Hydra config. Args: cfg: Hydra DictConfig containing reward section Returns: - Dictionary with reward_config key for env_cfg_override + Dictionary with ``target_field`` as the env config override key. Raises: ValueError: If reward config is missing """ - return {"reward_config": resolve_reward_dict(cfg)} + if target_field not in {"reward_config", "rewards"}: + raise ValueError( + f"Reward config target_field must be 'reward_config' or 'rewards', got {target_field!r}" + ) + return {target_field: resolve_reward_dict(cfg)} diff --git a/tests/config/test_manager_reward_routing.py b/tests/config/test_manager_reward_routing.py new file mode 100644 index 000000000..25d8ed580 --- /dev/null +++ b/tests/config/test_manager_reward_routing.py @@ -0,0 +1,139 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path + +import numpy as np +import pytest +from omegaconf import OmegaConf + +from unilab.base import registry +from unilab.base.base import EnvCfg +from unilab.base.config_overrides import ( + CONFIG_MAPPING_POLICY_KEY, + MANAGER_TERM_MAPPING_POLICY, +) +from unilab.base.registry import apply_cfg_overrides +from unilab.envs.manager_based_rl_env import ManagerBasedRlEnvCfg +from unilab.managers import RewardTermCfg +from unilab.training.backend_adapter import BackendAdapter + +_MANAGER_ENV = "_TestManagerRewardRoute" +_LEGACY_ENV = "_TestLegacyRewardRoute" +_MISSING_ENV = "_TestMissingRewardRoute" +_AMBIGUOUS_ENV = "_TestAmbiguousRewardRoute" + + +@dataclass +class _LegacyCfg(EnvCfg): + reward_config: dict[str, object] = field(default_factory=dict) + + +@dataclass +class _MissingCfg(EnvCfg): + pass + + +@dataclass +class _AmbiguousCfg(EnvCfg): + reward_config: dict[str, object] = field(default_factory=dict) + rewards: dict[str, object] = field( + default_factory=dict, + metadata={CONFIG_MAPPING_POLICY_KEY: MANAGER_TERM_MAPPING_POLICY}, + ) + + +for _name, _cfg_cls in ( + (_MANAGER_ENV, ManagerBasedRlEnvCfg), + (_LEGACY_ENV, _LegacyCfg), + (_MISSING_ENV, _MissingCfg), + (_AMBIGUOUS_ENV, _AmbiguousCfg), +): + if not registry.contains(_name): + registry.register_env_config(_name, _cfg_cls) + + +def _reward(_env, *, std: float) -> np.ndarray: + del std + return np.zeros(1, dtype=np.float32) + + +def _alive(_env) -> np.ndarray: + return np.ones(1, dtype=np.float32) + + +def _cfg(task_name: str, *, env: dict[str, object] | None = None): + return OmegaConf.create( + { + "training": {"task_name": task_name}, + "reward": { + "tracking": { + "weight": 2.0, + "params": {"std": 0.25}, + } + }, + "env": env or {"ctrl_dt": 0.02}, + } + ) + + +def test_backend_adapter_routes_manager_reward_and_preserves_factory_terms() -> None: + override = BackendAdapter( + _cfg(_MANAGER_ENV), + root_dir=Path("."), + ).build_task_env_cfg_override() + + assert "reward_config" not in override + assert override["rewards"]["tracking"]["weight"] == pytest.approx(2.0) + assert override["ctrl_dt"] == pytest.approx(0.02) + + manager_cfg = ManagerBasedRlEnvCfg( + rewards={ + "tracking": RewardTermCfg(func=_reward, weight=1.0, params={"std": 0.5}), + "alive": RewardTermCfg(func=_alive, weight=0.1), + } + ) + factory_tracking = manager_cfg.rewards["tracking"] + assert factory_tracking is not None + tracking_func = factory_tracking.func + + apply_cfg_overrides(manager_cfg, override) + + tracking = manager_cfg.rewards["tracking"] + assert tracking is not None + assert tracking.func is tracking_func + assert tracking.weight == pytest.approx(2.0) + assert tracking.params == {"std": 0.25} + assert manager_cfg.rewards["alive"] is not None + assert list(manager_cfg.rewards) == ["tracking", "alive"] + + +def test_backend_adapter_preserves_legacy_reward_target() -> None: + override = BackendAdapter( + _cfg(_LEGACY_ENV), + root_dir=Path("."), + ).build_task_env_cfg_override() + + assert "rewards" not in override + assert override["reward_config"]["tracking"]["weight"] == pytest.approx(2.0) + assert override["ctrl_dt"] == pytest.approx(0.02) + + +def test_backend_adapter_rejects_duplicate_manager_reward_sources() -> None: + cfg = _cfg(_MANAGER_ENV, env={"rewards": {"tracking": {"weight": 3.0}}}) + + with pytest.raises(ValueError, match="root 'reward'.*env.rewards"): + BackendAdapter(cfg, root_dir=Path(".")).build_task_env_cfg_override() + + +@pytest.mark.parametrize( + ("env_name", "match"), + [ + ("_UnregisteredRewardRoute", "not registered"), + (_MISSING_ENV, "declares no supported.*reward target"), + (_AMBIGUOUS_ENV, "declares both.*rewards.*reward_config"), + ], +) +def test_reward_target_resolution_fails_closed(env_name: str, match: str) -> None: + with pytest.raises(ValueError, match=match): + registry.resolve_reward_override_field(env_name) From 9630680804e2cfc98fbe39d6b244f4d9ec41f60b Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 03:06:03 +0800 Subject: [PATCH 014/169] feat: add manager root-state reset contract (#1072) (#1073) --- src/unilab/base/backend/__init__.py | 2 +- src/unilab/base/backend/base.py | 60 +++++++- src/unilab/base/backend/mjwarp/backend.py | 25 ++++ src/unilab/base/backend/motrix/backend.py | 16 +++ src/unilab/base/backend/mujoco/backend.py | 20 +++ src/unilab/base/entity.py | 158 +++++++++++++++++++++- src/unilab/base/reset_state.py | 147 +++++++++++++++++++- tests/base/test_backend_conformance.py | 135 +++++++++++++++++- tests/base/test_entity_facade.py | 113 +++++++++++++++- tests/base/test_mjwarp_backend.py | 3 + tests/base/test_motrix_backend_options.py | 23 ++++ tests/base/test_reset_state.py | 133 +++++++++++++++++- 12 files changed, 826 insertions(+), 9 deletions(-) diff --git a/src/unilab/base/backend/__init__.py b/src/unilab/base/backend/__init__.py index feb035345..0a406d077 100644 --- a/src/unilab/base/backend/__init__.py +++ b/src/unilab/base/backend/__init__.py @@ -5,7 +5,7 @@ from unilab.base.scene import SceneCfg -from .base import RenderClosedError, SimBackend +from .base import BackendRootStateLayout, RenderClosedError, SimBackend if TYPE_CHECKING: from unilab.base.base import EnvCfg diff --git a/src/unilab/base/backend/base.py b/src/unilab/base/backend/base.py index ba23ea3d7..578dd0322 100644 --- a/src/unilab/base/backend/base.py +++ b/src/unilab/base/backend/base.py @@ -51,6 +51,45 @@ def __post_init__(self) -> None: raise TypeError("sample_height must be callable") +@dataclass(frozen=True) +class BackendRootStateLayout: + """Generalized-state columns for one floating root body. + + ``qpos_indices`` address ``[x, y, z, qw, qx, qy, qz]`` in the public + :meth:`SimBackend.set_state` qpos representation. ``qvel_indices`` address + ``[linear_velocity_world, angular_velocity_body]``. Manager-facing root + states use world-frame angular velocity, so the base-owned reset + transaction performs the frame conversion before calling ``set_state``. + """ + + qpos_indices: tuple[int, ...] + qvel_indices: tuple[int, ...] + + def __post_init__(self) -> None: + for name, values, expected in ( + ("qpos_indices", self.qpos_indices, 7), + ("qvel_indices", self.qvel_indices, 6), + ): + if not isinstance(values, tuple): + raise TypeError(f"BackendRootStateLayout {name} must be a tuple") + if len(values) != expected: + raise ValueError( + f"BackendRootStateLayout {name} must contain {expected} columns; " + f"got {len(values)}" + ) + if any( + isinstance(value, (bool, np.bool_)) or not isinstance(value, (int, np.integer)) + for value in values + ): + raise TypeError(f"BackendRootStateLayout {name} must contain integer columns") + normalized = tuple(int(value) for value in values) + if any(value < 0 for value in normalized): + raise ValueError(f"BackendRootStateLayout {name} cannot contain negative columns") + if len(set(normalized)) != expected: + raise ValueError(f"BackendRootStateLayout {name} must contain unique columns") + object.__setattr__(self, name, normalized) + + @dataclass(frozen=True) class BackendPlayCapabilities: """Backend-native play/render capabilities surfaced through env contracts.""" @@ -203,6 +242,20 @@ def get_init_qvel(self) -> np.ndarray: Zero-filled qvel array. """ + def get_root_state_layout(self, root_body_name: str) -> BackendRootStateLayout: + """Resolve one body's floating-root columns on the cold path. + + Backends must verify that ``root_body_name`` owns a free/floating joint; + fixed bodies and runtimes without body-to-root metadata fail closed. + Name/model lookup is forbidden on reset and step hot paths, so callers + cache either the returned layout or the unsupported result during scene + materialization. + """ + raise NotImplementedError( + f"{self.__class__.__name__} does not expose root-state layout for " + f"body {root_body_name!r}" + ) + @abc.abstractmethod def get_body_ids(self, names: Sequence[str]) -> np.ndarray: """Resolve body/link names to backend integer IDs. @@ -358,8 +411,11 @@ def set_state( Args: env_indices: Environment indices. - qpos: Position state. - qvel: Velocity state. + qpos: Position state. Free-root columns exposed by + :meth:`get_root_state_layout` use world xyz and wxyz quaternion. + qvel: Velocity state. Free-root columns exposed by + :meth:`get_root_state_layout` use world linear velocity and + body-frame angular velocity. randomization: Optional backend randomization payload. Returns: diff --git a/src/unilab/base/backend/mjwarp/backend.py b/src/unilab/base/backend/mjwarp/backend.py index 8d2ebe77a..00e979163 100644 --- a/src/unilab/base/backend/mjwarp/backend.py +++ b/src/unilab/base/backend/mjwarp/backend.py @@ -19,6 +19,7 @@ from unilab.base.backend.base import ( BackendPlayCapabilities, BackendPlayRenderPlan, + BackendRootStateLayout, SimBackend, normalize_play_render_mode, ) @@ -346,6 +347,30 @@ def get_default_dof_pos(self) -> np.ndarray: def get_init_qvel(self) -> np.ndarray: return np.zeros((self._nv,), dtype=np.float32) + def get_root_state_layout(self, root_body_name: str) -> BackendRootStateLayout: + try: + body_id = self._body_ids[root_body_name] + except KeyError as exc: + raise ValueError(f"Body {root_body_name!r} not found in mjwarp model") from exc + joint_count = int(self._cpu_model.body_jntnum[body_id]) + joint_id = int(self._cpu_model.body_jntadr[body_id]) + free_joint = int(self._mujoco.mjtJoint.mjJNT_FREE) + if ( + joint_count != 1 + or joint_id < 0 + or int(self._cpu_model.jnt_type[joint_id]) != free_joint + ): + raise NotImplementedError( + "backend 'mjwarp' capability 'root-state layout' requires body " + f"{root_body_name!r} to own exactly one free joint" + ) + qpos_start = int(self._cpu_model.jnt_qposadr[joint_id]) + qvel_start = int(self._cpu_model.jnt_dofadr[joint_id]) + return BackendRootStateLayout( + qpos_indices=tuple(range(qpos_start, qpos_start + 7)), + qvel_indices=tuple(range(qvel_start, qvel_start + 6)), + ) + def get_body_ids(self, names: Sequence[str]) -> np.ndarray: resolved: list[int] = [] for name in names: diff --git a/src/unilab/base/backend/motrix/backend.py b/src/unilab/base/backend/motrix/backend.py index 41536d43c..a5545a5a8 100644 --- a/src/unilab/base/backend/motrix/backend.py +++ b/src/unilab/base/backend/motrix/backend.py @@ -40,6 +40,7 @@ BackendHeightScanner, BackendPlayCapabilities, BackendPlayRenderPlan, + BackendRootStateLayout, BackendTerrainSpawnData, RenderClosedError, SimBackend, @@ -448,6 +449,21 @@ def get_default_dof_pos(self) -> np.ndarray: def get_init_qvel(self) -> np.ndarray: return np.zeros((self._model.num_dof_vel,), dtype=self._np_dtype) + def get_root_state_layout(self, root_body_name: str) -> BackendRootStateLayout: + body = self._model.get_body(root_body_name) + if body is None: + raise ValueError(f"Body '{root_body_name}' not found in Motrix model") + floating_base = body.floatingbase + if floating_base is None: + raise NotImplementedError( + "backend 'motrix' capability 'root-state layout' requires body " + f"'{root_body_name}' to own a floating base" + ) + return BackendRootStateLayout( + qpos_indices=tuple(int(index) for index in floating_base.dof_pos_indices), + qvel_indices=tuple(int(index) for index in floating_base.dof_vel_indices), + ) + def get_body_ids(self, names: Sequence[str]) -> np.ndarray: ids: list[int] = [] for name in names: diff --git a/src/unilab/base/backend/mujoco/backend.py b/src/unilab/base/backend/mujoco/backend.py index d7a86afea..6cb4d4e60 100644 --- a/src/unilab/base/backend/mujoco/backend.py +++ b/src/unilab/base/backend/mujoco/backend.py @@ -37,6 +37,7 @@ BackendHeightScanner, BackendPlayCapabilities, BackendPlayRenderPlan, + BackendRootStateLayout, BackendTerrainSpawnData, SimBackend, normalize_play_render_mode, @@ -751,6 +752,25 @@ def get_default_dof_pos(self) -> np.ndarray: def get_init_qvel(self) -> np.ndarray: return np.zeros((self.nv,), dtype=self._np_dtype) + def get_root_state_layout(self, root_body_name: str) -> BackendRootStateLayout: + body_id = mujoco.mj_name2id(self._model, mujoco.mjtObj.mjOBJ_BODY, root_body_name) + if body_id < 0: + raise ValueError(f"Body '{root_body_name}' not found in MuJoCo model") + joint_count = int(self._model.body_jntnum[body_id]) + joint_id = int(self._model.body_jntadr[body_id]) + free_joint = int(mujoco.mjtJoint.mjJNT_FREE) + if joint_count != 1 or joint_id < 0 or int(self._model.jnt_type[joint_id]) != free_joint: + raise NotImplementedError( + "backend 'mujoco' capability 'root-state layout' requires body " + f"'{root_body_name}' to own exactly one free joint" + ) + qpos_start = int(self._model.jnt_qposadr[joint_id]) + qvel_start = int(self._model.jnt_dofadr[joint_id]) + return BackendRootStateLayout( + qpos_indices=tuple(range(qpos_start, qpos_start + 7)), + qvel_indices=tuple(range(qvel_start, qvel_start + 6)), + ) + def get_body_ids(self, names: "Sequence[str]") -> np.ndarray: ids: list[int] = [] for name in names: diff --git a/src/unilab/base/entity.py b/src/unilab/base/entity.py index 3dd913ca1..02a507ba1 100644 --- a/src/unilab/base/entity.py +++ b/src/unilab/base/entity.py @@ -15,8 +15,8 @@ import numpy as np -from unilab.base.backend.base import SimBackend -from unilab.utils.rotation import np_quat_apply_inverse, np_yaw_from_quat +from unilab.base.backend.base import BackendRootStateLayout, SimBackend +from unilab.utils.rotation import np_quat_apply, np_quat_apply_inverse, np_yaw_from_quat if TYPE_CHECKING: from unilab.base.reset_state import ResetStateTransaction @@ -141,6 +141,8 @@ def __init__( root_body_ids: np.ndarray | None, joint_pos_ids: np.ndarray | None, joint_vel_ids: np.ndarray | None, + default_root_state: np.ndarray | None, + default_root_state_error: str | None, default_joint_pos: np.ndarray | None, default_joint_vel: np.ndarray | None, gravity_vec_w: np.ndarray | None, @@ -157,6 +159,8 @@ def __init__( self._root_body_ids = root_body_ids self._joint_pos_index = None if joint_pos_ids is None else _as_column_index(joint_pos_ids) self._joint_vel_index = None if joint_vel_ids is None else _as_column_index(joint_vel_ids) + self._default_root_state = default_root_state + self._default_root_state_error = default_root_state_error self._default_joint_pos = default_joint_pos self._default_joint_vel = default_joint_vel self._gravity_vec_w = gravity_vec_w @@ -227,6 +231,17 @@ def root_link_pose_w(self) -> np.ndarray: def root_link_vel_w(self) -> np.ndarray: return np.concatenate((self.root_link_lin_vel_w, self.root_link_ang_vel_w), axis=-1) + @property + def default_root_state(self) -> np.ndarray: + """Read-only 13-D community root state for every environment.""" + if self._default_root_state is None: + detail = self._default_root_state_error or "root_body_name was not declared" + raise NotImplementedError( + f"Entity '{self._entity_name}' data capability 'default root state' is " + f"unavailable on backend '{self._backend_type}': {detail}" + ) + return self._default_root_state + @property def joint_pos(self) -> np.ndarray: index = self._require(self._joint_pos_index, "joint position") @@ -404,6 +419,8 @@ def __init__( self._backend_type = backend.backend_type self._backend = backend self._reset_state = reset_state + self._reset_root_layout: BackendRootStateLayout | None = None + self._reset_root_layout_error: str | None = None self._reset_joint_qpos_ids: np.ndarray | None = None self._reset_joint_qvel_ids: np.ndarray | None = None @@ -458,6 +475,11 @@ def __init__( self._validate_joint_state(backend, joint_pos_ids, joint_vel_ids) self._validate_body_state(backend, root_body_ids, body_ids) + ( + self._reset_root_layout, + default_root_state, + self._reset_root_layout_error, + ) = self._materialize_root_state(backend, cfg.root_body_name) default_joint_pos = self._materialize_default_joint_pos(backend, joint_pos_ids) default_joint_vel = self._materialize_default_joint_vel(backend, joint_vel_ids) gravity_vec_w = self._materialize_gravity_vector(backend, root_body_ids) @@ -484,6 +506,8 @@ def __init__( root_body_ids=root_body_ids, joint_pos_ids=joint_pos_ids, joint_vel_ids=joint_vel_ids, + default_root_state=default_root_state, + default_root_state_error=self._reset_root_layout_error, default_joint_pos=default_joint_pos, default_joint_vel=default_joint_vel, gravity_vec_w=gravity_vec_w, @@ -662,6 +686,78 @@ def _materialize_default_joint_pos( materialized.setflags(write=False) return materialized + def _materialize_root_state( + self, + backend: SimBackend, + root_body_name: str | None, + ) -> tuple[BackendRootStateLayout | None, np.ndarray | None, str | None]: + if root_body_name is None: + return None, None, "root_body_name was not declared in EntityCfg" + try: + layout = backend.get_root_state_layout(root_body_name) + except (AttributeError, NotImplementedError) as exc: + return None, None, str(exc) + if not isinstance(layout, BackendRootStateLayout): + raise TypeError( + f"Entity '{self.name}' capability 'root-state layout' on backend " + f"'{self._backend_type}' must return BackendRootStateLayout, got " + f"{type(layout).__name__}" + ) + try: + qpos = backend.get_default_qpos() + qvel = backend.get_init_qvel() + except (AttributeError, NotImplementedError) as exc: + return None, None, str(exc) + qpos_default = self._validate_root_default_vector(qpos, "default qpos") + qvel_default = self._validate_root_default_vector(qvel, "initial qvel") + qpos_indices = np.asarray(layout.qpos_indices, dtype=np.intp) + qvel_indices = np.asarray(layout.qvel_indices, dtype=np.intp) + if np.any(qpos_indices >= qpos_default.size): + raise ValueError( + f"Entity '{self.name}' root qpos layout exceeds backend " + f"'{self._backend_type}' width {qpos_default.size}: {qpos_indices.tolist()}" + ) + if np.any(qvel_indices >= qvel_default.size): + raise ValueError( + f"Entity '{self.name}' root qvel layout exceeds backend " + f"'{self._backend_type}' width {qvel_default.size}: {qvel_indices.tolist()}" + ) + + pose = np.asarray(qpos_default[qpos_indices]) + quaternion = pose[3:7] + norm = float(np.linalg.norm(quaternion)) + if not np.isclose(norm, 1.0, rtol=1e-5, atol=1e-6): + raise ValueError( + f"Entity '{self.name}' default root quaternion on backend " + f"'{self._backend_type}' must be unit length; norm={norm}" + ) + generalized_velocity = np.asarray(qvel_default[qvel_indices]) + velocity_w = np.array(generalized_velocity, copy=True) + velocity_w[3:6] = np_quat_apply(quaternion, generalized_velocity[3:6]) + root_state = np.concatenate((pose, velocity_w)) + materialized = np.broadcast_to(root_state, (backend.num_envs, 13)).copy() + materialized.setflags(write=False) + return layout, materialized, None + + def _validate_root_default_vector(self, value: np.ndarray, capability: str) -> np.ndarray: + if not isinstance(value, np.ndarray): + raise TypeError( + f"Entity '{self.name}' capability '{capability}' on backend " + f"'{self._backend_type}' must return np.ndarray, got {type(value).__name__}" + ) + if value.ndim != 1 or not np.issubdtype(value.dtype, np.floating): + raise TypeError( + f"Entity '{self.name}' capability '{capability}' on backend " + f"'{self._backend_type}' must be a 1-D floating array; got " + f"shape={value.shape}, dtype={value.dtype}" + ) + if not np.isfinite(value).all(): + raise ValueError( + f"Entity '{self.name}' capability '{capability}' on backend " + f"'{self._backend_type}' returned NaN or Inf" + ) + return value + def _materialize_default_joint_vel( self, backend: SimBackend, joint_vel_ids: np.ndarray | None ) -> np.ndarray | None: @@ -881,6 +977,64 @@ def set_joint_position_target( ) self.data.write_ctrl(target, env_ids, actuator_ids=actuator_ids) + def write_root_state_to_sim( + self, + root_state: np.ndarray, + env_ids: np.ndarray | slice | None = None, + ) -> None: + """Stage a 13-D world-frame root state in the active reset transaction.""" + reset_state, layout = self._require_root_state_write() + resolved_env_ids = self._normalize_reset_env_ids(env_ids) + reset_state.write_root_state( + resolved_env_ids, + layout, + root_state, + term_name=f"{self.name}.write_root_state_to_sim", + ) + + def write_root_link_pose_to_sim( + self, + root_pose: np.ndarray, + env_ids: np.ndarray | slice | None = None, + ) -> None: + """Stage world position and wxyz root orientation during reset.""" + reset_state, layout = self._require_root_state_write() + resolved_env_ids = self._normalize_reset_env_ids(env_ids) + reset_state.write_root_pose( + resolved_env_ids, + layout, + root_pose, + term_name=f"{self.name}.write_root_link_pose_to_sim", + ) + + def write_root_link_velocity_to_sim( + self, + root_velocity: np.ndarray, + env_ids: np.ndarray | slice | None = None, + ) -> None: + """Stage world linear/angular root velocity during reset.""" + reset_state, layout = self._require_root_state_write() + resolved_env_ids = self._normalize_reset_env_ids(env_ids) + reset_state.write_root_velocity( + resolved_env_ids, + layout, + root_velocity, + term_name=f"{self.name}.write_root_link_velocity_to_sim", + ) + + def _require_root_state_write( + self, + ) -> tuple[ResetStateTransaction, BackendRootStateLayout]: + if self._reset_state is None: + raise self._capability_error( + "reset root-state write", + "EntityScene was materialized without an env-owned reset transaction", + ) + if self._reset_root_layout is None: + detail = self._reset_root_layout_error or "root-state layout was not materialized" + raise self._capability_error("reset root-state layout", detail) + return self._reset_state, self._reset_root_layout + def write_joint_state_to_sim( self, position: np.ndarray, diff --git a/src/unilab/base/reset_state.py b/src/unilab/base/reset_state.py index 86a1bdffc..b28fe333f 100644 --- a/src/unilab/base/reset_state.py +++ b/src/unilab/base/reset_state.py @@ -12,7 +12,8 @@ import numpy as np -from unilab.base.backend.base import SimBackend +from unilab.base.backend.base import BackendRootStateLayout, SimBackend +from unilab.utils.rotation import np_quat_apply_inverse class ResetStateTransaction: @@ -146,6 +147,87 @@ def write_joint_state( self._qvel[ids[:, None], vel_columns[None, :]] = velocities self._dirty_mask[ids] = True + def write_root_state( + self, + env_ids: np.ndarray, + layout: BackendRootStateLayout, + root_state: np.ndarray, + *, + term_name: str, + ) -> None: + """Stage a community 13-D world-frame root state.""" + self._require_active() + ids = self._validate_ids(env_ids, capability="write_root_state") + values = self._validate_values( + root_state, + shape=(ids.size, 13), + capability="root state", + term_name=term_name, + ) + self.write_root_pose(ids, layout, values[:, :7], term_name=term_name) + self.write_root_velocity(ids, layout, values[:, 7:], term_name=term_name) + + def write_root_pose( + self, + env_ids: np.ndarray, + layout: BackendRootStateLayout, + pose: np.ndarray, + *, + term_name: str, + ) -> None: + """Stage world position and wxyz orientation for one floating root.""" + ids = self._prepare_state_write(env_ids, capability="root-pose", term_name=term_name) + positions = self._validate_values( + pose, + shape=(ids.size, 7), + capability="root pose", + term_name=term_name, + ) + self._validate_quaternions(positions[:, 3:7], term_name=term_name) + qpos_columns, _ = self._validate_root_layout(layout, term_name=term_name) + assert self._qpos is not None + if ids.size: + self._qpos[ids[:, None], qpos_columns[None, :]] = positions + self._dirty_mask[ids] = True + + def write_root_velocity( + self, + env_ids: np.ndarray, + layout: BackendRootStateLayout, + velocity_w: np.ndarray, + *, + term_name: str, + ) -> None: + """Stage world-frame root velocity in generalized qvel columns. + + The public generalized-state contract stores free-root linear velocity + in world coordinates and angular velocity in root-body coordinates. + The conversion uses the pose already staged in this transaction. + """ + ids = self._prepare_state_write(env_ids, capability="root-velocity", term_name=term_name) + velocities = self._validate_values( + velocity_w, + shape=(ids.size, 6), + capability="root velocity", + term_name=term_name, + ) + qpos_columns, qvel_columns = self._validate_root_layout(layout, term_name=term_name) + assert self._qpos is not None + assert self._qvel is not None + if ids.size: + quaternions = self._qpos[ + ids[:, None], + qpos_columns[None, 3:7], + ] + self._validate_quaternions(quaternions, term_name=term_name) + encoded_velocity = np.array(velocities, copy=True) + encoded_velocity[:, 3:6] = np_quat_apply_inverse( + quaternions, + velocities[:, 3:6], + ) + self._qvel[ids[:, None], qvel_columns[None, :]] = encoded_velocity + self._dirty_mask[ids] = True + def commit(self) -> dict | None: """Commit all staged rows through one public backend call.""" self._require_active() @@ -194,6 +276,69 @@ 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 _prepare_state_write( + self, + env_ids: np.ndarray, + *, + capability: str, + term_name: str, + ) -> np.ndarray: + self._require_active() + ids = self._validate_ids(env_ids, capability=f"write_{capability}") + outside = ids[~self._active_mask[ids]] + if outside.size: + raise ValueError( + f"EventManager term '{term_name}' attempted {capability} mutation outside " + f"the active reset: {outside.tolist()}" + ) + self._requesting_terms.add(term_name) + self._materialize_default_state(term_name) + assert self._default_qpos is not None + assert self._default_qvel is not None + assert self._qpos is not None + assert self._qvel is not None + uninitialized = ids[~self._dirty_mask[ids]] + if uninitialized.size: + self._qpos[uninitialized] = self._default_qpos + self._qvel[uninitialized] = self._default_qvel + return ids + + def _validate_root_layout( + self, + layout: BackendRootStateLayout, + *, + term_name: str, + ) -> tuple[np.ndarray, np.ndarray]: + if not isinstance(layout, BackendRootStateLayout): + raise TypeError( + f"EventManager term '{term_name}' root-state layout must be " + f"BackendRootStateLayout, got {type(layout).__name__}" + ) + assert self._default_qpos is not None + assert self._default_qvel is not None + qpos_columns = self._validate_columns( + np.asarray(layout.qpos_indices, dtype=np.intp), + width=self._default_qpos.size, + capability="root qpos indices", + term_name=term_name, + ) + qvel_columns = self._validate_columns( + np.asarray(layout.qvel_indices, dtype=np.intp), + width=self._default_qvel.size, + capability="root qvel indices", + term_name=term_name, + ) + return qpos_columns, qvel_columns + + def _validate_quaternions(self, values: np.ndarray, *, term_name: str) -> None: + norms = np.linalg.norm(values, axis=1) + invalid = ~np.isclose(norms, 1.0, rtol=1e-5, atol=1e-6) + if np.any(invalid): + raise ValueError( + f"EventManager term '{term_name}' root quaternion must be unit length; " + f"norms={norms[invalid].tolist()}" + ) + def _validate_state_vector( self, value: np.ndarray, diff --git a/tests/base/test_backend_conformance.py b/tests/base/test_backend_conformance.py index ad35d11dc..639081614 100644 --- a/tests/base/test_backend_conformance.py +++ b/tests/base/test_backend_conformance.py @@ -15,7 +15,11 @@ from unilab.assets import ASSETS_ROOT_PATH from unilab.base.backend import create_backend -from unilab.base.backend.base import BackendTerrainSpawnData, SimBackend +from unilab.base.backend.base import ( + BackendRootStateLayout, + BackendTerrainSpawnData, + SimBackend, +) from unilab.base.scene import SceneCfg REPO_ROOT = Path(__file__).resolve().parents[2] @@ -139,6 +143,29 @@ def test_actuation_metadata_defaults_fail_closed() -> None: SimBackend.get_joint_state_qpos_indices(object(), ("joint",)) # type: ignore[arg-type] with pytest.raises(NotImplementedError, match="get_joint_state_qvel_indices"): SimBackend.get_joint_state_qvel_indices(object(), ("joint",)) # type: ignore[arg-type] + with pytest.raises(NotImplementedError, match="root-state layout"): + SimBackend.get_root_state_layout(object(), "root") # type: ignore[arg-type] + + +@pytest.mark.parametrize( + ("qpos_indices", "qvel_indices", "error", "match"), + [ + (list(range(7)), tuple(range(6)), TypeError, "qpos_indices must be a tuple"), + (tuple(range(6)), tuple(range(6)), ValueError, "qpos_indices must contain 7"), + (tuple(range(7)), tuple(range(5)), ValueError, "qvel_indices must contain 6"), + ((0, 1, 2, 3, 4, 5, True), tuple(range(6)), TypeError, "integer columns"), + ((0, 1, 2, 3, 4, 5, -1), tuple(range(6)), ValueError, "negative columns"), + ((0, 1, 2, 3, 4, 5, 5), tuple(range(6)), ValueError, "unique columns"), + ], +) +def test_root_state_layout_metadata_fails_closed( + qpos_indices, + qvel_indices, + error, + match: str, +) -> None: + with pytest.raises(error, match=match): + BackendRootStateLayout(qpos_indices, qvel_indices) @pytest.mark.parametrize("backend_type", _BACKEND_PARAMS) @@ -192,6 +219,112 @@ def test_actuation_metadata_contract(backend_type: str) -> None: np.testing.assert_array_equal(backend.get_default_dof_pos(), detached) +@pytest.mark.parametrize("backend_type", _BACKEND_PARAMS) +def test_root_state_layout_contract(backend_type: str) -> None: + _require_backend(backend_type) + + backend = create_backend( + backend_type, + SceneCfg(model_file=_G1_SCENE), + NUM_ENVS, + SIM_DT, + base_name="pelvis", + ) + backend.materialize() + + if backend_type == "drake": + with pytest.raises( + NotImplementedError, + match="DrakeBackend does not expose root-state layout.*pelvis", + ): + backend.get_root_state_layout("pelvis") + return + + layout = backend.get_root_state_layout("pelvis") + assert isinstance(layout, BackendRootStateLayout) + qpos_indices = np.asarray(layout.qpos_indices) + qvel_indices = np.asarray(layout.qvel_indices) + assert qpos_indices.shape == (7,) + assert qvel_indices.shape == (6,) + assert np.all(qpos_indices < backend.get_default_qpos().size) + assert np.all(qvel_indices < backend.get_init_qvel().size) + root_pose = backend.get_default_qpos()[qpos_indices] + np.testing.assert_allclose(np.linalg.norm(root_pose[3:7]), 1.0, atol=1e-6) + + +@pytest.mark.parametrize("backend_type", ["mujoco", "motrix"]) +def test_root_qvel_body_angular_contract_reads_back_world_velocity(backend_type: str) -> None: + _require_backend(backend_type) + backend = create_backend( + backend_type, + SceneCfg(model_file=_G1_SCENE), + 1, + SIM_DT, + base_name="pelvis", + add_body_sensors=True, + ) + backend.materialize() + layout = backend.get_root_state_layout("pelvis") + qpos_indices = np.asarray(layout.qpos_indices) + qvel_indices = np.asarray(layout.qvel_indices) + qpos = backend.get_default_qpos()[None].copy() + qvel = backend.get_init_qvel()[None].copy() + half_sqrt = np.sqrt(0.5) + qpos[0, qpos_indices[3:7]] = [half_sqrt, 0.0, 0.0, half_sqrt] + qvel[0, qvel_indices[3:6]] = [0.0, -1.0, 0.0] + + backend.set_state(np.array([0], dtype=np.int32), qpos, qvel) + + root_body_id = backend.get_body_ids(["pelvis"]) + np.testing.assert_allclose( + backend.get_body_ang_vel_w(root_body_id)[0, 0], + [1.0, 0.0, 0.0], + atol=2e-6, + ) + + +def test_mujoco_root_layout_resolves_a_nonfirst_free_joint() -> None: + import mujoco + + from unilab.base.backend.mujoco.backend import MuJoCoBackend + + model = mujoco.MjModel.from_xml_string( + """ + + + + + + + + + + + + + """ + ) + backend = object.__new__(MuJoCoBackend) + backend._model = model + + layout = backend.get_root_state_layout("floating") + assert layout.qpos_indices == tuple(range(1, 8)) + assert layout.qvel_indices == tuple(range(1, 7)) + with pytest.raises(NotImplementedError, match="hinged.*exactly one free joint"): + backend.get_root_state_layout("hinged") + + +def test_drake_root_layout_is_explicitly_unsupported_without_runtime_metadata() -> None: + from unilab.base.backend.drake.backend import DrakeBackend + + backend = object.__new__(DrakeBackend) + with pytest.raises( + NotImplementedError, + match="DrakeBackend does not expose root-state layout.*trunk", + ): + backend.get_root_state_layout("trunk") + + def test_terrain_spawn_consumers_do_not_probe_private_backend_capabilities() -> None: forbidden_names = {"terrain_origins", "terrain_surface_sampler", "sample_height"} offenders: list[str] = [] diff --git a/tests/base/test_entity_facade.py b/tests/base/test_entity_facade.py index 08f4c686f..ba1f9ce34 100644 --- a/tests/base/test_entity_facade.py +++ b/tests/base/test_entity_facade.py @@ -11,8 +11,9 @@ import unilab.base.entity as entity_module from unilab.assets import ASSETS_ROOT_PATH -from unilab.base.backend.base import SimBackend +from unilab.base.backend.base import BackendRootStateLayout, SimBackend from unilab.base.entity import EntityCfg, EntityScene +from unilab.base.reset_state import ResetStateTransaction from unilab.base.scene import SceneCfg from unilab.managers import RewardManager, RewardTermCfg, SceneEntityCfg @@ -49,6 +50,15 @@ def __init__(self, backend_type: str, *, unsupported: frozenset[str] = frozenset self.body_ang_vel = self.body_pos + 300.0 self.body_lin_vel_b = self.body_pos + 400.0 self.body_ang_vel_b = self.body_pos + 500.0 + self.default_qpos = np.array( + [99.0, 1.0, 2.0, 3.0, 1.0, 0.0, 0.0, 0.0, 88.0], + dtype=np.float32, + ) + self.init_qvel = np.array( + [77.0, 66.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 55.0], + dtype=np.float32, + ) + self.set_state_calls: list[tuple[np.ndarray, np.ndarray, np.ndarray]] = [] def _check(self, capability: str) -> None: self.calls[capability] += 1 @@ -59,6 +69,30 @@ def get_body_ids(self, names) -> np.ndarray: self._check("body names") return np.asarray([self.body_ids[name] for name in names], dtype=np.int32) + def get_root_state_layout(self, root_body_name: str) -> BackendRootStateLayout: + self._check("root-state layout") + if root_body_name != "base": + raise ValueError(f"unknown root body {root_body_name}") + return BackendRootStateLayout(tuple(range(1, 8)), tuple(range(2, 8))) + + def get_default_qpos(self) -> np.ndarray: + self._check("default qpos") + return self.default_qpos.copy() + + def get_init_qvel(self) -> np.ndarray: + self._check("initial qvel") + return self.init_qvel.copy() + + def set_state( + self, + env_ids: np.ndarray, + qpos: np.ndarray, + qvel: np.ndarray, + randomization=None, + ) -> None: + assert randomization is None + self.set_state_calls.append((env_ids.copy(), qpos.copy(), qvel.copy())) + def get_joint_dof_pos_indices(self, names) -> np.ndarray: self._check("joint position names") return np.asarray([self.joint_ids[name] for name in names], dtype=np.int32) @@ -159,6 +193,11 @@ def test_backend_profiles_materialize_identical_local_entity_contract(backend_ty np.testing.assert_array_equal(robot.data.root_link_ang_vel_b, backend.body_ang_vel_b[:, 4]) np.testing.assert_array_equal(robot.data.heading_w, 0.0) np.testing.assert_array_equal(robot.data.projected_gravity_b, [[0.0, 0.0, -1.0]] * 3) + np.testing.assert_array_equal( + robot.data.default_root_state, + np.tile([1.0, 2.0, 3.0, 1.0, 0.0, 0.0, 0.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0], (3, 1)), + ) + assert not robot.data.default_root_state.flags.writeable np.testing.assert_array_equal( robot.data.actuator_ctrl_range, np.arange(10, dtype=np.float32).reshape(5, 2)[[4, 2]], @@ -237,6 +276,78 @@ def test_entity_control_write_uses_cached_actuator_columns_and_fails_closed() -> read_only_scene["robot"].data.write_ctrl(np.zeros((backend.num_envs, 2), dtype=np.float32)) +def test_entity_root_writes_use_cached_layout_and_one_reset_commit() -> None: + backend = _StrictBackendProfile("mujoco") + transaction = ResetStateTransaction(cast(SimBackend, backend)) + scene = EntityScene( + {"robot": EntityCfg(root_body_name="base")}, + cast(SimBackend, backend), + reset_state=transaction, + ) + robot = scene["robot"] + cold_layout_calls = backend.calls["root-state layout"] + half_sqrt = np.sqrt(0.5) + + with transaction.scoped(np.array([0, 2], dtype=np.int32)): + robot.write_root_link_pose_to_sim( + np.array( + [ + [10.0, 11.0, 12.0, half_sqrt, 0.0, 0.0, half_sqrt], + [20.0, 21.0, 22.0, 1.0, 0.0, 0.0, 0.0], + ], + dtype=np.float32, + ), + env_ids=np.array([2, 0], dtype=np.int32), + ) + robot.write_root_link_velocity_to_sim( + np.array( + [[1.0, 2.0, 3.0, 1.0, 0.0, 0.0], [4.0, 5.0, 6.0, 0.0, 1.0, 2.0]], + dtype=np.float32, + ), + env_ids=np.array([2, 0], dtype=np.int32), + ) + + assert backend.calls["root-state layout"] == cold_layout_calls == 1 + assert len(backend.set_state_calls) == 1 + env_ids, qpos, qvel = backend.set_state_calls[0] + np.testing.assert_array_equal(env_ids, [0, 2]) + np.testing.assert_allclose(qpos[0, 1:8], [20.0, 21.0, 22.0, 1.0, 0.0, 0.0, 0.0]) + np.testing.assert_allclose(qpos[1, 1:8], [10.0, 11.0, 12.0, half_sqrt, 0.0, 0.0, half_sqrt]) + np.testing.assert_allclose(qvel[0, 2:8], [4.0, 5.0, 6.0, 0.0, 1.0, 2.0]) + np.testing.assert_allclose(qvel[1, 2:8], [1.0, 2.0, 3.0, 0.0, -1.0, 0.0], atol=1e-6) + + +def test_entity_caches_unsupported_root_layout_without_hot_path_probe() -> None: + backend = _StrictBackendProfile( + "drake", + unsupported=frozenset({"root-state layout"}), + ) + transaction = ResetStateTransaction(cast(SimBackend, backend)) + scene = EntityScene( + {"robot": EntityCfg(root_body_name="base")}, + cast(SimBackend, backend), + reset_state=transaction, + ) + robot = scene["robot"] + + with pytest.raises( + NotImplementedError, + match="default root state.*backend 'drake'.*drake lacks root-state layout", + ): + _ = robot.data.default_root_state + with pytest.raises( + NotImplementedError, + match="reset root-state layout.*backend 'drake'.*drake lacks root-state layout", + ): + with transaction.scoped(np.array([0], dtype=np.int32)): + robot.write_root_state_to_sim( + np.array([[0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]]), + env_ids=np.array([0], dtype=np.int32), + ) + assert backend.calls["root-state layout"] == 1 + assert backend.set_state_calls == [] + + def test_entity_joint_position_target_maps_natural_joint_order_to_control_order() -> None: backend = _StrictBackendProfile("mujoco") control = np.zeros((backend.num_envs, backend.num_actuators), dtype=np.float32) diff --git a/tests/base/test_mjwarp_backend.py b/tests/base/test_mjwarp_backend.py index 2a50e0a3d..02535d219 100644 --- a/tests/base/test_mjwarp_backend.py +++ b/tests/base/test_mjwarp_backend.py @@ -53,6 +53,9 @@ def test_real_cuda_init_reset_step() -> None: assert backend.backend_type == "mjwarp" assert backend.num_actuators == 29 assert backend.num_dof_vel == 29 + root_layout = backend.get_root_state_layout("pelvis") + assert root_layout.qpos_indices == tuple(range(7)) + assert root_layout.qvel_indices == tuple(range(6)) qpos, qvel = _stand_state(backend, 2) backend.set_state(np.asarray([0, 1], dtype=np.int32), qpos, qvel) diff --git a/tests/base/test_motrix_backend_options.py b/tests/base/test_motrix_backend_options.py index dab58fc84..8a767cf9f 100644 --- a/tests/base/test_motrix_backend_options.py +++ b/tests/base/test_motrix_backend_options.py @@ -328,6 +328,29 @@ def test_motrix_backend_get_body_pose_w_slices_cached_poses_once() -> None: np.testing.assert_allclose(quat, [[[0.8, 0.5, 0.6, 0.7]]]) +def test_motrix_root_layout_uses_selected_body_floating_base_indices() -> None: + import unilab.base.backend.motrix.backend as mod + + floating_base = SimpleNamespace( + dof_pos_indices=[4, 5, 6, 7, 8, 9, 10], + dof_vel_indices=[3, 4, 5, 6, 7, 8], + ) + bodies = { + "fixed": SimpleNamespace(floatingbase=None), + "floating": SimpleNamespace(floatingbase=floating_base), + } + backend = object.__new__(mod.MotrixBackend) + backend._model = SimpleNamespace(get_body=lambda name: bodies.get(name)) + + layout = backend.get_root_state_layout("floating") + assert layout.qpos_indices == tuple(range(4, 11)) + assert layout.qvel_indices == tuple(range(3, 9)) + with pytest.raises(NotImplementedError, match="fixed.*floating base"): + backend.get_root_state_layout("fixed") + with pytest.raises(ValueError, match="missing.*not found"): + backend.get_root_state_layout("missing") + + def test_motrix_backend_applies_init_geom_size_overrides(monkeypatch, tmp_path) -> None: mod, fake_model = _install_fake_motrix(monkeypatch, tmp_path) backend = mod.MotrixBackend( diff --git a/tests/base/test_reset_state.py b/tests/base/test_reset_state.py index 38af6d9e1..5242240ac 100644 --- a/tests/base/test_reset_state.py +++ b/tests/base/test_reset_state.py @@ -7,7 +7,7 @@ import numpy as np import pytest -from unilab.base.backend.base import SimBackend +from unilab.base.backend.base import BackendRootStateLayout, SimBackend from unilab.base.reset_state import ResetStateTransaction @@ -130,6 +130,137 @@ def test_joint_writes_initialize_defaults_and_compose_by_column() -> None: np.testing.assert_array_equal(qvel, [[-2.0, 0.0], [-1.0, 0.0]]) +def test_root_pose_and_world_velocity_compose_at_nonzero_columns() -> None: + default_qpos = np.array([99.0, 1.0, 2.0, 3.0, 1.0, 0.0, 0.0, 0.0, 88.0]) + default_qvel = np.array([77.0, 66.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 55.0]) + backend = _Backend(qpos=default_qpos, qvel=default_qvel) + transaction = _transaction(backend) + layout = BackendRootStateLayout(tuple(range(1, 8)), tuple(range(2, 8))) + half_sqrt = np.sqrt(0.5) + poses = np.array( + [ + [10.0, 11.0, 12.0, half_sqrt, 0.0, 0.0, half_sqrt], + [20.0, 21.0, 22.0, 1.0, 0.0, 0.0, 0.0], + ] + ) + velocities_w = np.array( + [ + [1.0, 2.0, 3.0, 1.0, 0.0, 0.0], + [4.0, 5.0, 6.0, 0.0, 1.0, 2.0], + ] + ) + + with transaction.scoped(np.array([0, 2], dtype=np.int32)): + transaction.write_root_pose( + np.array([2, 0], dtype=np.int32), + layout, + poses, + term_name="root_pose", + ) + transaction.write_root_velocity( + np.array([2, 0], dtype=np.int32), + layout, + velocities_w, + term_name="root_velocity", + ) + + assert len(backend.set_state_calls) == 1 + ids, qpos, qvel = backend.set_state_calls[0] + np.testing.assert_array_equal(ids, [0, 2]) + np.testing.assert_array_equal(qpos[:, [0, 8]], [[99.0, 88.0], [99.0, 88.0]]) + np.testing.assert_allclose(qpos[0, 1:8], poses[1]) + np.testing.assert_allclose(qpos[1, 1:8], poses[0]) + np.testing.assert_array_equal(qvel[:, [0, 1, 8]], [[77.0, 66.0, 55.0]] * 2) + np.testing.assert_allclose(qvel[0, 2:8], velocities_w[1]) + np.testing.assert_allclose(qvel[1, 2:5], velocities_w[0, :3]) + np.testing.assert_allclose(qvel[1, 5:8], [0.0, -1.0, 0.0], atol=1e-7) + + +def test_combined_root_state_uses_staged_pose_for_angular_velocity() -> None: + backend = _Backend( + qpos=np.array([0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0]), + qvel=np.zeros(6), + ) + transaction = _transaction(backend) + layout = BackendRootStateLayout(tuple(range(7)), tuple(range(6))) + half_sqrt = np.sqrt(0.5) + root_state = np.array( + [[1.0, 2.0, 3.0, half_sqrt, 0.0, 0.0, half_sqrt, 4.0, 5.0, 6.0, 1.0, 0.0, 0.0]] + ) + + with transaction.scoped(np.array([1], dtype=np.int32)): + transaction.write_root_state( + np.array([1], dtype=np.int32), + layout, + root_state, + term_name="root_state", + ) + + _, qpos, qvel = backend.set_state_calls[0] + np.testing.assert_allclose(qpos[0], root_state[0, :7]) + np.testing.assert_allclose(qvel[0, :3], root_state[0, 7:10]) + np.testing.assert_allclose(qvel[0, 3:6], [0.0, -1.0, 0.0], atol=1e-7) + + +@pytest.mark.parametrize( + ("root_state", "error", "match"), + [ + (np.zeros((1, 12)), ValueError, "root state.*expected"), + (np.zeros((1, 13), dtype=np.int32), TypeError, "root state.*floating"), + (np.full((1, 13), np.nan), ValueError, "root state.*NaN or Inf"), + ( + np.array([[0.0, 0.0, 0.0, 2.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]]), + ValueError, + "root quaternion must be unit length", + ), + ], +) +def test_root_state_values_fail_closed(root_state, error, match: str) -> None: + backend = _Backend( + qpos=np.array([0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0]), + qvel=np.zeros(6), + ) + transaction = _transaction(backend) + layout = BackendRootStateLayout(tuple(range(7)), tuple(range(6))) + with pytest.raises(error, match=match): + with transaction.scoped(np.array([0], dtype=np.int32)): + transaction.write_root_state( + np.array([0], dtype=np.int32), + layout, + root_state, + term_name="bad_root", + ) + assert backend.set_state_calls == [] + + +def test_root_layout_bounds_and_reset_scope_fail_closed() -> None: + backend = _Backend( + qpos=np.array([0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0]), + qvel=np.zeros(6), + ) + transaction = _transaction(backend) + out_of_bounds = BackendRootStateLayout(tuple(range(1, 8)), tuple(range(6))) + + with pytest.raises(IndexError, match="root qpos indices out of range"): + with transaction.scoped(np.array([0], dtype=np.int32)): + transaction.write_root_pose( + np.array([0], dtype=np.int32), + out_of_bounds, + np.array([[0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0]]), + term_name="bad_layout", + ) + + valid = BackendRootStateLayout(tuple(range(7)), tuple(range(6))) + with pytest.raises(ValueError, match="root-pose mutation outside the active reset"): + with transaction.scoped(np.array([0], dtype=np.int32)): + transaction.write_root_pose( + np.array([1], dtype=np.int32), + valid, + np.array([[0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0]]), + term_name="outside", + ) + + @pytest.mark.parametrize( ("position", "velocity", "error", "match"), [ From 8c5b0c7c7af9a9e02177cc3279420ecdf2192d93 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 03:29:02 +0800 Subject: [PATCH 015/169] feat(events): add uniform root-state reset (#1075) --- src/unilab/envs/mdp/__init__.py | 2 + src/unilab/envs/mdp/events.py | 81 ++++++++- tests/envs/mdp/test_events.py | 300 ++++++++++++++++++++++++++++++++ 3 files changed, 381 insertions(+), 2 deletions(-) create mode 100644 tests/envs/mdp/test_events.py diff --git a/src/unilab/envs/mdp/__init__.py b/src/unilab/envs/mdp/__init__.py index 880cc9d14..1a9262dc7 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 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 from unilab.envs.mdp.observations import base_ang_vel as base_ang_vel @@ -50,6 +51,7 @@ "is_alive", "is_terminated", "projected_gravity", + "reset_root_state_uniform", "reset_scene_to_default", "resolve_env_ids", "root_height_below_minimum", diff --git a/src/unilab/envs/mdp/events.py b/src/unilab/envs/mdp/events.py index 568119270..81a9125c2 100644 --- a/src/unilab/envs/mdp/events.py +++ b/src/unilab/envs/mdp/events.py @@ -6,14 +6,55 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, cast import numpy as np +from unilab.managers.scene_entity_config import SceneEntityCfg +from unilab.utils.rotation import np_quat_from_euler_xyz, np_quat_mul + if TYPE_CHECKING: + from unilab.base.entity import Entity from unilab.managers._types import ManagerBasedRlEnv +_DEFAULT_ASSET_CFG = SceneEntityCfg("robot") +_SE3_KEYS = ("x", "y", "z", "roll", "pitch", "yaw") + + +def _sample_se3_range( + range_dict: dict[str, tuple[float, float]] | None, + shape: tuple[int, ...], + rng: np.random.Generator, +) -> np.ndarray: + """Sample uniform ``[x, y, z, roll, pitch, yaw]`` offsets with NumPy.""" + if not shape or shape[-1] != len(_SE3_KEYS): + raise ValueError( + f"reset_root_state_uniform SE(3) sample shape must end in 6; received {shape}" + ) + try: + ranges = np.asarray( + [(range_dict or {}).get(key, (0.0, 0.0)) for key in _SE3_KEYS], + dtype=np.float64, + ) + except (TypeError, ValueError) as exc: + raise ValueError( + "reset_root_state_uniform ranges must map each SE(3) key to a numeric (min, max) pair" + ) from exc + if ranges.shape != (len(_SE3_KEYS), 2): + raise ValueError( + "reset_root_state_uniform ranges must map each SE(3) key to a " + f"(min, max) pair; received shape {ranges.shape}" + ) + if not np.isfinite(ranges).all(): + raise ValueError("reset_root_state_uniform ranges must contain only finite values") + invalid = ranges[:, 0] > ranges[:, 1] + if np.any(invalid): + keys = [_SE3_KEYS[index] for index in np.flatnonzero(invalid)] + raise ValueError(f"reset_root_state_uniform range minimum exceeds maximum for keys {keys}") + return rng.uniform(ranges[:, 0], ranges[:, 1], size=shape) + + def resolve_env_ids(env: ManagerBasedRlEnv, env_ids: np.ndarray | None) -> np.ndarray: """Return concrete NumPy environment IDs, preserving community sentinel semantics.""" if env_ids is None: @@ -29,4 +70,40 @@ def reset_scene_to_default(env: ManagerBasedRlEnv, env_ids: np.ndarray | None) - env.scene.reset_to_default(ids, term_name="reset_scene_to_default") -__all__ = ["reset_scene_to_default", "resolve_env_ids"] +def reset_root_state_uniform( + env: ManagerBasedRlEnv, + env_ids: np.ndarray | None, + pose_range: dict[str, tuple[float, float]], + velocity_range: dict[str, tuple[float, float]] | None = None, + asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG, +) -> None: + """Reset a floating root from defaults plus uniformly sampled SE(3) offsets. + + This is the NumPy adaptation of the pinned mjlab event. UniLab currently has + no public mocap-pose write contract, so fixed-base/mocap requests fail through + the entity's cached floating-root capability instead of falling back. + """ + ids = resolve_env_ids(env, env_ids) + asset = cast("Entity", env.scene[asset_cfg.name]) + try: + root_states = np.array(asset.data.default_root_state[ids], copy=True) + except NotImplementedError as exc: + raise NotImplementedError( + "EventManager term 'reset_root_state_uniform' requires a floating-root " + f"state for entity '{asset_cfg.name}'; fixed-base/mocap root reset is " + f"unsupported without a formal backend contract: {exc}" + ) from exc + + pose_samples = _sample_se3_range(pose_range, (len(ids), 6), env.rng) + root_states[:, 0:3] = root_states[:, 0:3] + pose_samples[:, 0:3] + env.scene.env_origins[ids] + orientation_delta = np_quat_from_euler_xyz( + pose_samples[:, 3], pose_samples[:, 4], pose_samples[:, 5] + ) + root_states[:, 3:7] = np_quat_mul(root_states[:, 3:7], orientation_delta) + + velocity_samples = _sample_se3_range(velocity_range, (len(ids), 6), env.rng) + root_states[:, 7:13] = root_states[:, 7:13] + velocity_samples + asset.write_root_state_to_sim(root_states, env_ids=ids) + + +__all__ = ["reset_root_state_uniform", "reset_scene_to_default", "resolve_env_ids"] diff --git a/tests/envs/mdp/test_events.py b/tests/envs/mdp/test_events.py new file mode 100644 index 000000000..eb020e994 --- /dev/null +++ b/tests/envs/mdp/test_events.py @@ -0,0 +1,300 @@ +"""Upstream-derived NumPy tests for Manager-Based reset event terms.""" + +from __future__ import annotations + +import ast +from pathlib import Path +from types import SimpleNamespace +from typing import Any, cast + +import numpy as np +import pytest + +from unilab.base.backend.base import BackendRootStateLayout, SimBackend +from unilab.base.entity import EntityCfg, EntityScene +from unilab.base.reset_state import ResetStateTransaction +from unilab.envs import mdp +from unilab.managers._types import ManagerBasedRlEnv + + +class _CaptureEntity: + def __init__(self, default_root_state: np.ndarray) -> None: + self.data = SimpleNamespace(default_root_state=default_root_state) + self.writes: list[tuple[np.ndarray, np.ndarray]] = [] + + def write_root_state_to_sim( + self, + root_state: np.ndarray, + env_ids: np.ndarray | None = None, + ) -> None: + assert env_ids is not None + self.writes.append((root_state.copy(), env_ids.copy())) + + +class _CaptureScene: + def __init__(self, entity: _CaptureEntity, env_origins: np.ndarray) -> None: + self.entities = {"robot": entity} + self.env_origins = env_origins + + def __getitem__(self, name: str) -> _CaptureEntity: + return self.entities[name] + + +def _capture_env( + *, + seed: int = 17, + default_root_state: np.ndarray | None = None, + env_origins: np.ndarray | None = None, +) -> tuple[ManagerBasedRlEnv, _CaptureEntity]: + if default_root_state is None: + default_root_state = np.zeros((3, 13), dtype=np.float32) + default_root_state[:, 3] = 1.0 + default_root_state.setflags(write=False) + entity = _CaptureEntity(default_root_state) + if env_origins is None: + env_origins = np.zeros((3, 3), dtype=np.float32) + env = cast( + ManagerBasedRlEnv, + SimpleNamespace( + num_envs=3, + rng=np.random.default_rng(seed), + scene=_CaptureScene(entity, env_origins), + ), + ) + return env, entity + + +def test_uniform_root_state_applies_pinned_pose_velocity_and_origin_semantics() -> None: + half_sqrt = np.float32(np.sqrt(0.5)) + defaults = np.zeros((3, 13), dtype=np.float32) + defaults[:, :3] = [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]] + defaults[:, 3:7] = [half_sqrt, 0.0, 0.0, half_sqrt] + defaults[:, 7:13] = np.arange(18, dtype=np.float32).reshape(3, 6) + origins = np.asarray( + [[100.0, 0.0, 0.0], [0.0, 200.0, 0.0], [0.0, 0.0, 300.0]], + dtype=np.float32, + ) + env, entity = _capture_env(default_root_state=defaults, env_origins=origins) + ids = np.asarray([2, 0], dtype=np.int32) + + mdp.reset_root_state_uniform( + env, + ids, + pose_range={ + "x": (1.0, 1.0), + "y": (-2.0, -2.0), + "z": (0.5, 0.5), + "roll": (np.pi / 2.0, np.pi / 2.0), + }, + velocity_range={ + "x": (0.1, 0.1), + "y": (0.2, 0.2), + "z": (0.3, 0.3), + "roll": (0.4, 0.4), + "pitch": (0.5, 0.5), + "yaw": (0.6, 0.6), + }, + ) + + assert len(entity.writes) == 1 + root_state, written_ids = entity.writes[0] + np.testing.assert_array_equal(written_ids, ids) + np.testing.assert_allclose( + root_state[:, :3], + defaults[ids, :3] + [1.0, -2.0, 0.5] + origins[ids], + ) + np.testing.assert_allclose( + root_state[:, 3:7], + [[0.5, 0.5, 0.5, 0.5], [0.5, 0.5, 0.5, 0.5]], + atol=1e-6, + ) + np.testing.assert_allclose( + root_state[:, 7:13], + defaults[ids, 7:13] + [0.1, 0.2, 0.3, 0.4, 0.5, 0.6], + ) + np.testing.assert_array_equal(defaults[:, :3], [[1, 2, 3], [4, 5, 6], [7, 8, 9]]) + + +def test_uniform_root_state_uses_env_rng_and_preserves_single_env_shape() -> None: + pose_range = {"x": (-1.0, 1.0), "yaw": (-0.5, 0.5)} + velocity_range = {"x": (-2.0, 2.0), "roll": (-0.25, 0.25)} + left, left_entity = _capture_env(seed=9) + right, right_entity = _capture_env(seed=9) + ids = np.asarray([1], dtype=np.int32) + + mdp.reset_root_state_uniform(left, ids, pose_range, velocity_range) + mdp.reset_root_state_uniform(right, ids, pose_range, velocity_range) + + left_state, left_ids = left_entity.writes[0] + right_state, right_ids = right_entity.writes[0] + assert left_state.shape == (1, 13) + np.testing.assert_array_equal(left_ids, ids) + np.testing.assert_array_equal(right_ids, ids) + np.testing.assert_array_equal(left_state, right_state) + + +def test_uniform_root_state_none_ids_targets_all_environments() -> None: + env, entity = _capture_env() + + mdp.reset_root_state_uniform(env, None, pose_range={}) + + root_state, ids = entity.writes[0] + np.testing.assert_array_equal(ids, np.arange(3, dtype=np.int32)) + assert root_state.shape == (3, 13) + + +class _Backend: + backend_type = "fake" + num_envs = 3 + num_actuators = 0 + + def __init__(self, *, root_layout_supported: bool = True) -> None: + self.root_layout_supported = root_layout_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.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 + self.body_velocity = np.zeros((self.num_envs, 1, 3)) + + def get_body_ids(self, names) -> np.ndarray: + if tuple(names) != ("base",): + raise KeyError(names) + return np.asarray([0], dtype=np.int32) + + def get_root_state_layout(self, root_body_name: str) -> BackendRootStateLayout: + if not self.root_layout_supported: + raise NotImplementedError("fake fixed-base entity has no free-root layout") + if root_body_name != "base": + raise ValueError(root_body_name) + return BackendRootStateLayout(tuple(range(7)), tuple(range(6))) + + def get_default_qpos(self) -> np.ndarray: + return self.default_qpos.copy() + + def get_init_qvel(self) -> np.ndarray: + return self.init_qvel.copy() + + def get_dof_pos(self) -> np.ndarray: + return np.empty((self.num_envs, 0)) + + def get_dof_vel(self) -> np.ndarray: + return np.empty((self.num_envs, 0)) + + def get_body_pos_w(self, ids: np.ndarray) -> np.ndarray: + return self.body_pos[:, ids] + + def get_body_quat_w(self, ids: np.ndarray) -> np.ndarray: + return self.body_quat[:, ids] + + def get_body_lin_vel_w(self, ids: np.ndarray) -> np.ndarray: + return self.body_velocity[:, ids] + + def get_body_ang_vel_w(self, ids: np.ndarray) -> np.ndarray: + return self.body_velocity[:, ids] + + def get_body_lin_vel_b(self, ids: np.ndarray) -> np.ndarray: + return self.body_velocity[:, ids] + + def get_body_ang_vel_b(self, ids: np.ndarray) -> np.ndarray: + return self.body_velocity[:, ids] + + def set_state( + self, + env_ids: np.ndarray, + qpos: np.ndarray, + qvel: np.ndarray, + randomization=None, + ) -> None: + assert randomization is None + self.set_state_calls.append((env_ids.copy(), qpos.copy(), qvel.copy())) + + +def _transaction_env( + *, root_layout_supported: bool = True +) -> tuple[ManagerBasedRlEnv, _Backend, ResetStateTransaction]: + backend = _Backend(root_layout_supported=root_layout_supported) + transaction = ResetStateTransaction(cast(SimBackend, backend)) + scene = EntityScene( + {"robot": EntityCfg(root_body_name="base")}, + cast(SimBackend, backend), + reset_state=transaction, + ) + env = cast( + ManagerBasedRlEnv, + SimpleNamespace( + num_envs=backend.num_envs, + rng=np.random.default_rng(5), + scene=scene, + ), + ) + return env, backend, transaction + + +def test_uniform_root_state_composes_in_one_reset_transaction_commit() -> None: + env, backend, transaction = _transaction_env() + active_ids = np.asarray([0, 2], dtype=np.int32) + + with transaction.scoped(active_ids): + mdp.reset_scene_to_default(env, active_ids) + mdp.reset_root_state_uniform( + env, + np.asarray([2], dtype=np.int32), + pose_range={"x": (1.0, 1.0)}, + velocity_range={"x": (0.5, 0.5)}, + ) + assert backend.set_state_calls == [] + + assert len(backend.set_state_calls) == 1 + ids, qpos, qvel = backend.set_state_calls[0] + np.testing.assert_array_equal(ids, active_ids) + np.testing.assert_array_equal(qpos[0], backend.default_qpos) + np.testing.assert_allclose(qpos[1], backend.default_qpos + [1.0, 0, 0, 0, 0, 0, 0]) + np.testing.assert_array_equal(qvel[0], backend.init_qvel) + np.testing.assert_allclose(qvel[1], backend.init_qvel + [0.5, 0, 0, 0, 0, 0]) + + +def test_uniform_root_state_fixed_or_mocap_capability_fails_closed() -> None: + env, backend, transaction = _transaction_env(root_layout_supported=False) + + with pytest.raises( + NotImplementedError, + match="reset_root_state_uniform.*entity 'robot'.*fixed-base/mocap.*backend 'fake'", + ): + with transaction.scoped(np.asarray([1], dtype=np.int32)): + mdp.reset_root_state_uniform(env, np.asarray([1], dtype=np.int32), pose_range={}) + + assert backend.set_state_calls == [] + + +@pytest.mark.parametrize( + ("pose_range", "message"), + [ + ({"x": (np.nan, 1.0)}, "finite"), + ({"yaw": (1.0, -1.0)}, "minimum exceeds maximum"), + ({"z": (0.0, 1.0, 2.0)}, r"numeric \(min, max\) pair"), + ], +) +def test_uniform_root_state_invalid_ranges_fail_before_write( + pose_range: dict[str, Any], message: str +) -> None: + env, entity = _capture_env() + + with pytest.raises(ValueError, match=message): + mdp.reset_root_state_uniform(env, np.asarray([0], dtype=np.int32), pose_range) + + assert entity.writes == [] + + +def test_events_module_has_no_forbidden_runtime_dependencies() -> None: + path = Path(__file__).resolve().parents[3] / "src" / "unilab" / "envs" / "mdp" / "events.py" + tree = ast.parse(path.read_text(encoding="utf-8")) + forbidden = ("torch", "unilab.ipc", "unilab.algos", "unilab.training", "unilab.base.backend") + imports = [node.module or "" for node in ast.walk(tree) if isinstance(node, ast.ImportFrom)] + [ + alias.name + for node in ast.walk(tree) + if isinstance(node, ast.Import) + for alias in node.names + ] + assert not [name for name in imports if name.startswith(forbidden)] From b48b388d5fb2791bb01227f0a2fb90f6b49fd5d7 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 11:20:49 +0800 Subject: [PATCH 016/169] feat(backend): add named sensor view contract (#1077) --- ...-community-manager-api-on-numpy-runtime.md | 14 ++ .../2-contracts/2-backend_contract.md | 6 +- src/unilab/base/backend/__init__.py | 3 +- src/unilab/base/backend/base.py | 157 +++++++++++++++++- src/unilab/base/backend/drake/backend.py | 19 +++ src/unilab/base/backend/mjwarp/backend.py | 14 +- src/unilab/base/backend/motrix/backend.py | 59 +++++-- src/unilab/base/backend/motrix/scene.py | 121 ++++++++++---- src/unilab/base/backend/mujoco/backend.py | 11 +- tests/base/test_backend_conformance.py | 35 ++++ tests/base/test_backend_sensor_view.py | 126 ++++++++++++++ tests/base/test_motrix_backend_options.py | 6 +- tests/base/test_sim_backend_smoke.py | 7 + 13 files changed, 527 insertions(+), 51 deletions(-) create mode 100644 tests/base/test_backend_sensor_view.py diff --git a/docs/sphinx/source/adr/ADR-0006-community-manager-api-on-numpy-runtime.md b/docs/sphinx/source/adr/ADR-0006-community-manager-api-on-numpy-runtime.md index 51291ede7..45bc37dbd 100644 --- a/docs/sphinx/source/adr/ADR-0006-community-manager-api-on-numpy-runtime.md +++ b/docs/sphinx/source/adr/ADR-0006-community-manager-api-on-numpy-runtime.md @@ -111,6 +111,19 @@ scene composer 或通用 asset hierarchy。 - 新 backend 能力必须作为独立 child 扩展 `SimBackend` 并补 conformance tests，不能在 manager 或 env 中用 `getattr` / `hasattr` 探测私有实现。 +Named sensor 使用 `SimBackend.bind_sensor_data(names)` 在 materialization 冷路径校验名称、 +每个 sensor 的展平宽度、batch shape 与 finite 值，并返回 immutable +`BackendSensorView`。term 热路径只调用 `view.read()`；MuJoCo、Drake 和 MJWarp adapter +分别保留已解析的 host-cache slice 或数值 slot。MotrixSim 当前公开接口只提供 named +sensor accessor、没有数值 sensor ID，因此 Motrix adapter 在 scene materialization 时缓存 +可用名称，并把原生批量 accessor 与 immutable 名称 tuple 封装为 backend-owned opaque +reader；term 不接触名称解析、XML 或 model metadata，未知名称在进入原生调用前 fail-closed。 + +这是 pinned mjlab sensor-facing 语义的 intentional NumPy/backend adaptation：社区侧的 +tensor/device view 在 UniLab 表达为按请求名称顺序拼接的二维 NumPy batch +`(num_envs, sum(sensor_widths))`。名称顺序、单 sensor 宽度和当前值可见，Torch device、 +backend model/data 与原生 handle 不属于 manager contract。 + ### 5. Fail-closed capability rule 用户显式禁用与实现缺失是两种不同状态。前者允许 Null manager；后者必须失败： @@ -154,6 +167,7 @@ fallback 到旧单体 env 的永久兼容路径。 | `ManagerBasedRlEnv` return | Adapted | 保留 `NpEnvState` 与 UniLab reset/final-observation contract | | config container | Adapted | plain instances + Hydra owner YAML overlay，不引入第二套 runtime | | `SceneEntityCfg` selectors | Adapted | 语义保留；只解析 `SimBackend` 已声明能力 | +| named sensor view | Adapted | 冷路径 bind；有序展平 NumPy batch；reader 由 backend 拥有 | | event/domain randomization | Adapted | 调度语义保留；mutation 走 backend DR/capability contract | | Metrics/Recorder | Adapted | lifecycle hook 存在时启用；缺失时显式失败或显式空配置 | | Torch device、Warp mutation、viewer glue | Unsupported | 不进入 manager core，不提供静默替代 | diff --git a/docs/sphinx/source/en/4-developer_guide/2-contracts/2-backend_contract.md b/docs/sphinx/source/en/4-developer_guide/2-contracts/2-backend_contract.md index f7c021958..2012bae3b 100644 --- a/docs/sphinx/source/en/4-developer_guide/2-contracts/2-backend_contract.md +++ b/docs/sphinx/source/en/4-developer_guide/2-contracts/2-backend_contract.md @@ -17,6 +17,9 @@ Optional capabilities are explicit: physics-state playback, and native video capture support. - `BackendHeightScanner` and `create_hfield_scanner(...)` expose terrain scan support through a reusable backend-owned object. +- `BackendSensorView` and `bind_sensor_data(...)` validate ordered named sensors + on the cold path and retain a backend-owned reader for finite, shape-stable + NumPy batches. Manager hot paths do not inspect XML or model metadata. - Domain randomization support is surfaced through `get_dr_capabilities()` and the init, reset, and interval randomization methods. - Unsupported optional methods raise `NotImplementedError` from the base class. @@ -37,5 +40,6 @@ Optional capabilities are explicit: - Backend factory: `src/unilab/base/backend/__init__.py` - MuJoCo backend: `src/unilab/base/backend/mujoco/backend.py` - Motrix backend: `src/unilab/base/backend/motrix/backend.py` -- Backend contract tests: `tests/base/test_sim_backend.py`, +- Backend contract tests: `tests/base/test_backend_sensor_view.py`, + `tests/base/test_backend_conformance.py`, `tests/base/test_sim_backend.py`, `tests/base/test_backend_imports.py`, `tests/base/test_motrix_backend_options.py` diff --git a/src/unilab/base/backend/__init__.py b/src/unilab/base/backend/__init__.py index 0a406d077..299744338 100644 --- a/src/unilab/base/backend/__init__.py +++ b/src/unilab/base/backend/__init__.py @@ -5,7 +5,7 @@ from unilab.base.scene import SceneCfg -from .base import BackendRootStateLayout, RenderClosedError, SimBackend +from .base import BackendRootStateLayout, BackendSensorView, RenderClosedError, SimBackend if TYPE_CHECKING: from unilab.base.base import EnvCfg @@ -221,6 +221,7 @@ def __getattr__(name: str): __all__ = [ "SimBackend", + "BackendSensorView", "RenderClosedError", "MuJoCoBackend", "MjwarpBackend", diff --git a/src/unilab/base/backend/base.py b/src/unilab/base/backend/base.py index 578dd0322..0913342db 100644 --- a/src/unilab/base/backend/base.py +++ b/src/unilab/base/backend/base.py @@ -1,6 +1,6 @@ import abc from collections.abc import Callable, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, field from os import PathLike from typing import Any @@ -15,6 +15,7 @@ PreStepControlFn = Callable[[Any, np.ndarray], np.ndarray] TerrainHeightSampleFn = Callable[[np.ndarray], np.ndarray] +SensorReadFn = Callable[[], np.ndarray] class RenderClosedError(RuntimeError): @@ -90,6 +91,88 @@ def __post_init__(self) -> None: object.__setattr__(self, name, normalized) +@dataclass(frozen=True) +class BackendSensorView: + """Validated batch view over one or more named backend sensors. + + Sensor names and flattened per-sensor widths are resolved while the backend + is materialized. Manager terms retain this view and only call ``read`` on + the hot path; they never inspect backend model objects or resolve XML names. + The reader is intentionally backend-owned so adapters can use cached + numeric slots, stable host slices, or an opaque native batch reader without + changing the manager-facing contract. + """ + + backend_type: str + names: tuple[str, ...] + dimensions: tuple[int, ...] + num_envs: int + _reader: SensorReadFn = field(repr=False, compare=False) + + def __post_init__(self) -> None: + if not isinstance(self.backend_type, str) or not self.backend_type: + raise ValueError("BackendSensorView backend_type must be a non-empty string") + if not isinstance(self.names, tuple) or not self.names: + raise ValueError("BackendSensorView names must be a non-empty tuple") + if any(not isinstance(name, str) or not name for name in self.names): + raise ValueError("BackendSensorView names must contain non-empty strings") + if len(set(self.names)) != len(self.names): + raise ValueError(f"BackendSensorView names must be unique: {self.names}") + if not isinstance(self.dimensions, tuple) or len(self.dimensions) != len(self.names): + raise ValueError("BackendSensorView dimensions must contain one entry per sensor name") + if any( + isinstance(value, (bool, np.bool_)) + or not isinstance(value, (int, np.integer)) + or int(value) <= 0 + for value in self.dimensions + ): + raise ValueError("BackendSensorView dimensions must be positive integers") + if ( + isinstance(self.num_envs, (bool, np.bool_)) + or not isinstance(self.num_envs, (int, np.integer)) + or int(self.num_envs) <= 0 + ): + raise ValueError("BackendSensorView num_envs must be a positive integer") + if not callable(self._reader): + raise TypeError("BackendSensorView reader must be callable") + object.__setattr__(self, "dimensions", tuple(int(value) for value in self.dimensions)) + object.__setattr__(self, "num_envs", int(self.num_envs)) + + @property + def width(self) -> int: + """Total flattened sensor width in the configured name order.""" + return int(sum(self.dimensions)) + + def read(self) -> np.ndarray: + """Read the current sensor batch and enforce the stable view contract.""" + try: + value = np.asarray(self._reader()) + except (KeyError, NotImplementedError, ValueError) as exc: + raise type(exc)( + f"Backend '{self.backend_type}' sensor view {self.names} could not be read: {exc}" + ) from exc + if value.ndim != 2 or value.shape != (self.num_envs, self.width): + raise ValueError( + f"Backend '{self.backend_type}' sensor view {self.names} returned shape " + f"{value.shape}; expected ({self.num_envs}, {self.width})" + ) + if not np.issubdtype(value.dtype, np.number) and not np.issubdtype(value.dtype, np.bool_): + raise TypeError( + f"Backend '{self.backend_type}' sensor view {self.names} returned non-numeric " + f"dtype {value.dtype}" + ) + if not np.isfinite(value).all(): + raise ValueError( + f"Backend '{self.backend_type}' sensor view {self.names} returned NaN or Inf" + ) + return value + + @property + def data(self) -> np.ndarray: + """Community-style spelling for a current sensor read.""" + return self.read() + + @dataclass(frozen=True) class BackendPlayCapabilities: """Backend-native play/render capabilities surfaced through env contracts.""" @@ -898,3 +981,75 @@ def get_sensor_data_batch(self, names: Sequence[str]) -> np.ndarray: values = [np.asarray(self.get_sensor_data(name)) for name in sensor_names] flat_values = [value.reshape(value.shape[0], -1) for value in values] return np.concatenate(flat_values, axis=1) + + def bind_sensor_data(self, names: Sequence[str]) -> BackendSensorView: + """Materialize a validated view over named sensors on the cold path. + + The existing sensor getters remain the sole backend adapter surface. This + method validates each requested sensor once, records its flattened width, + and returns a stable view for manager terms. Backends override the + protected reader hook when numeric slots or stable cache slices are + available; callers do not depend on that implementation detail. + """ + if isinstance(names, (str, bytes)): + raise TypeError( + f"Backend '{self.backend_type}' sensor view names must be a sequence of strings, " + "not one string" + ) + sensor_names = tuple(names) + if not sensor_names: + raise ValueError( + f"Backend '{self.backend_type}' sensor view requires at least one name" + ) + if any(not isinstance(name, str) or not name for name in sensor_names): + raise ValueError( + f"Backend '{self.backend_type}' sensor view names must be non-empty strings" + ) + if len(set(sensor_names)) != len(sensor_names): + raise ValueError( + f"Backend '{self.backend_type}' sensor view names must be unique: {sensor_names}" + ) + + dimensions: list[int] = [] + for name in sensor_names: + try: + value = np.asarray(self.get_sensor_data(name)) + except (KeyError, NotImplementedError, ValueError) as exc: + raise type(exc)( + f"Backend '{self.backend_type}' cannot bind sensor '{name}': {exc}" + ) from exc + if value.ndim < 1 or value.shape[0] != self.num_envs: + raise ValueError( + f"Backend '{self.backend_type}' sensor '{name}' returned shape " + f"{value.shape}; expected leading dimension {self.num_envs}" + ) + width = int(np.prod(value.shape[1:], dtype=np.int64)) if value.ndim > 1 else 1 + if width <= 0: + raise ValueError( + f"Backend '{self.backend_type}' sensor '{name}' has empty data shape " + f"{value.shape}" + ) + dimensions.append(width) + + view = BackendSensorView( + backend_type=self.backend_type, + names=sensor_names, + dimensions=tuple(dimensions), + num_envs=self.num_envs, + _reader=self._bind_sensor_data_reader(sensor_names), + ) + # Validate the batch implementation at materialization as well. This + # catches adapters whose individual and batch sensor contracts disagree. + view.read() + return view + + def _bind_sensor_data_reader(self, names: tuple[str, ...]) -> SensorReadFn: + """Create the backend-owned reader retained by a materialized sensor view. + + The default keeps the existing batch getter as the compatibility path + for lightweight adapters and test doubles. Concrete backends that + expose stable numeric slots or host-cache slices override this hook so + manager hot paths never resolve model metadata. + """ + batch_reader = self.get_sensor_data_batch + return lambda: batch_reader(names) diff --git a/src/unilab/base/backend/drake/backend.py b/src/unilab/base/backend/drake/backend.py index b2591f4f5..e58251c44 100644 --- a/src/unilab/base/backend/drake/backend.py +++ b/src/unilab/base/backend/drake/backend.py @@ -733,6 +733,25 @@ def get_sensor_data(self, name: str) -> np.ndarray: return self._sensor_views[name].copy() raise KeyError(f"Unknown DrakeUni sensor: {name}") + def _bind_sensor_data_reader(self, names: tuple[str, ...]) -> Callable[[], np.ndarray]: + """Capture DrakeUni sensor addresses; read only the refreshed host cache.""" + name_to_index = {name: index for index, name in enumerate(self._sensor_names)} + slots = tuple( + ( + int(self._sensor_adr[name_to_index[name]]), + int(self._sensor_dim[name_to_index[name]]), + ) + for name in names + ) + + def read() -> np.ndarray: + values = [ + self._sensor_data[:, address : address + dimension] for address, dimension in slots + ] + return np.concatenate(values, axis=1) + + return read + # Internal helpers. def _sync_runtime_state(self, output: dict[str, Any] | None = None) -> None: # Keep UniLab's cached state/sensor views aligned after every DrakeUni update. diff --git a/src/unilab/base/backend/mjwarp/backend.py b/src/unilab/base/backend/mjwarp/backend.py index 00e979163..f749333ed 100644 --- a/src/unilab/base/backend/mjwarp/backend.py +++ b/src/unilab/base/backend/mjwarp/backend.py @@ -10,7 +10,7 @@ from __future__ import annotations import time -from collections.abc import Sequence +from collections.abc import Callable, Sequence from os import PathLike from typing import Any @@ -816,3 +816,15 @@ def get_sensor_data(self, name: str) -> np.ndarray: available = ", ".join(sorted(self._sensor_slots)) raise ValueError(f"Sensor {name!r} not found; available: {available}") from exc return self._sensor_cache[:, address : address + dimension] + + def _bind_sensor_data_reader(self, names: tuple[str, ...]) -> Callable[[], np.ndarray]: + """Capture numeric host-cache slots for a zero-metadata hot-path view.""" + slots = tuple(self._sensor_slots[name] for name in names) + + def read() -> np.ndarray: + values = [ + self._sensor_cache[:, address : address + dimension] for address, dimension in slots + ] + return np.concatenate(values, axis=1) + + return read diff --git a/src/unilab/base/backend/motrix/backend.py b/src/unilab/base/backend/motrix/backend.py index a5545a5a8..16672e3c9 100644 --- a/src/unilab/base/backend/motrix/backend.py +++ b/src/unilab/base/backend/motrix/backend.py @@ -1,7 +1,7 @@ import logging import os import time -from collections.abc import Sequence +from collections.abc import Callable, Sequence from dataclasses import dataclass from typing import Any, TypeVar, cast @@ -84,6 +84,7 @@ def _contiguous_slice(indices: np.ndarray) -> slice | None: @dataclass class _MotrixSceneContext: model: "mtx.SceneModel" + sensor_names: tuple[str, ...] terrain_origins: np.ndarray | None = None terrain_surface_sampler: object | None = None cleanup_handle: object | None = None @@ -112,8 +113,8 @@ def _build_motrix_scene_context( base_name: str, ) -> _MotrixSceneContext: from unilab.base.backend.motrix.scene import ( - materialize_motrix_hfield_attached_scene, - materialize_motrix_scene, + _materialize_motrix_hfield_attached_scene_with_sensor_names, + _materialize_motrix_scene_with_sensor_names, ) if scene is None: @@ -122,29 +123,32 @@ def _build_motrix_scene_context( raise ValueError("SceneCfg.model_file must be provided") if scene.terrain is None: - model = materialize_motrix_scene( + model, sensor_names = _materialize_motrix_scene_with_sensor_names( model_file=scene.model_file, fragment_files=scene.fragment_files, add_body_sensors=add_body_sensors, base_name=base_name, ) - return _MotrixSceneContext(model=model) + return _MotrixSceneContext(model=model, sensor_names=sensor_names) if scene.terrain.generator is None: raise ValueError("SceneCfg.terrain.generator must be configured for terrain scenes") - model, terrain_origins, terrain_surface_sampler = materialize_motrix_hfield_attached_scene( - model_file=scene.model_file, - terrain_cfg=scene.terrain.generator, - fragment_files=scene.fragment_files, - hfield_name=scene.terrain.hfield_name, - geom_name=scene.terrain.geom_name or "floor", - add_body_sensors=add_body_sensors, - base_name=base_name, - return_surface_sampler=True, + model, terrain_origins, terrain_surface_sampler, sensor_names = ( + _materialize_motrix_hfield_attached_scene_with_sensor_names( + model_file=scene.model_file, + terrain_cfg=scene.terrain.generator, + fragment_files=scene.fragment_files, + hfield_name=scene.terrain.hfield_name, + geom_name=scene.terrain.geom_name or "floor", + add_body_sensors=add_body_sensors, + base_name=base_name, + return_surface_sampler=True, + ) ) return _MotrixSceneContext( model=model, + sensor_names=sensor_names, terrain_origins=terrain_origins, terrain_surface_sampler=terrain_surface_sampler, ) @@ -192,6 +196,7 @@ def __init__( self._base_name = base_name self._model = scene_context.model + self._sensor_names = frozenset(scene_context.sensor_names) self._body_id_to_name = { # type: ignore[assignment] link.index: link.name for link in self._model.links if link.name } @@ -1166,10 +1171,20 @@ def get_body_ang_vel_b(self, body_ids: np.ndarray) -> np.ndarray: # Sensors # # ------------------------------------------------------------------ # + def _validate_sensor_names(self, names: Sequence[str]) -> tuple[str, ...]: + sensor_names = tuple(names) + missing = tuple(name for name in sensor_names if name not in self._sensor_names) + if missing: + available = ", ".join(sorted(self._sensor_names)) + raise KeyError(f"Unknown Motrix sensor(s) {missing}; available sensors: {available}") + return sensor_names + def get_sensor_data(self, name: str) -> np.ndarray: + self._validate_sensor_names((name,)) return self._model.get_sensor_value(name, self._data) # type: ignore[no-any-return] def get_sensor_data_rows(self, name: str, env_ids: np.ndarray) -> np.ndarray: + self._validate_sensor_names((name,)) rows = np.asarray(env_ids, dtype=np.intp) mask = np.zeros(self._num_envs, dtype=bool) mask[rows] = True @@ -1178,12 +1193,26 @@ def get_sensor_data_rows(self, name: str, env_ids: np.ndarray) -> np.ndarray: return selected_values[np.searchsorted(selected_rows, rows)] # type: ignore[no-any-return] def get_sensor_data_batch(self, names: Sequence[str]) -> np.ndarray: - sensor_names = tuple(names) + sensor_names = self._validate_sensor_names(names) if not sensor_names: return np.empty((self._num_envs, 0), dtype=self._np_dtype) values = self._model.get_sensor_values(sensor_names, self._data) return np.asarray(values, dtype=self._np_dtype) + def _bind_sensor_data_reader(self, names: tuple[str, ...]) -> Callable[[], np.ndarray]: + """Retain Motrix's opaque native reader after cold-path name validation. + + MotrixSim exposes named sensor access but no public numeric sensor ID; + the bound callable is therefore the narrowest backend-owned reader. + It does not inspect XML or model metadata on the manager hot path. + """ + native_reader = self._model.get_sensor_values + + def read() -> np.ndarray: + return np.asarray(native_reader(names, self._data), dtype=self._np_dtype) + + return read + # ------------------------------------------------------------------ # # MotrixSim-specific # # ------------------------------------------------------------------ # diff --git a/src/unilab/base/backend/motrix/scene.py b/src/unilab/base/backend/motrix/scene.py index 2e982bbf7..2d65d1208 100644 --- a/src/unilab/base/backend/motrix/scene.py +++ b/src/unilab/base/backend/motrix/scene.py @@ -16,6 +16,21 @@ from motrixsim.msd import Link, World +def _motrix_sensor_names(world: "World") -> tuple[str, ...]: + """Collect the names accepted by Motrix's native sensor accessor once.""" + groups = ( + world.sensors.contact, + world.sensors.frame, + world.sensors.joint, + world.sensors.subtree, + world.sensors.touch, + ) + names = tuple(str(sensor.name) for group in groups for sensor in group if sensor.name) + if len(set(names)) != len(names): + raise ValueError(f"Motrix scene contains duplicate sensor names: {names}") + return names + + def _extract_keyframes(fragment_file: Path) -> list[ET.Element]: """Return ```` child elements declared inside ``fragment_file``.""" root = ET.parse(fragment_file).getroot() @@ -146,14 +161,14 @@ def add_motrix_tracking_frame_sensors(world: World, *, base_name: str) -> None: world.sensors.frame.append(sensor) -def materialize_motrix_scene( +def _materialize_motrix_scene_with_sensor_names( *, model_file: str, fragment_files: Sequence[str] = (), add_body_sensors: bool = False, base_name: str = "base", -) -> SceneModel: - """Build a MotrixSim model through MSD scene composition.""" +) -> tuple["SceneModel", tuple[str, ...]]: + """Build a Motrix model and return its validated cold-path sensor names.""" import motrixsim.msd as msd model_path = Path(model_file).resolve() @@ -167,40 +182,29 @@ def materialize_motrix_scene( _attach_motrix_scene_fragment(world, fragment_path) if add_body_sensors: add_motrix_tracking_frame_sensors(world, base_name=base_name) - return msd.build(world) + model = msd.build(world) + return model, _motrix_sensor_names(world) finally: _cleanup_temp_xml(robot_path, model_path) -@overload -def materialize_motrix_hfield_attached_scene( - *, - model_file: str, - terrain_cfg: TerrainGeneratorCfg, - fragment_files: Sequence[str] = (), - hfield_name: str = "terrain_hfield", - geom_name: str = "floor", - add_body_sensors: bool = False, - base_name: str = "base", - return_surface_sampler: Literal[False] = False, -) -> tuple[SceneModel, np.ndarray]: ... - - -@overload -def materialize_motrix_hfield_attached_scene( +def materialize_motrix_scene( *, model_file: str, - terrain_cfg: TerrainGeneratorCfg, fragment_files: Sequence[str] = (), - hfield_name: str = "terrain_hfield", - geom_name: str = "floor", add_body_sensors: bool = False, base_name: str = "base", - return_surface_sampler: Literal[True], -) -> tuple[SceneModel, np.ndarray, object]: ... +) -> "SceneModel": + """Build a MotrixSim model through MSD scene composition.""" + return _materialize_motrix_scene_with_sensor_names( + model_file=model_file, + fragment_files=fragment_files, + add_body_sensors=add_body_sensors, + base_name=base_name, + )[0] -def materialize_motrix_hfield_attached_scene( +def _materialize_motrix_hfield_attached_scene_with_sensor_names( *, model_file: str, terrain_cfg: TerrainGeneratorCfg, @@ -210,8 +214,8 @@ def materialize_motrix_hfield_attached_scene( add_body_sensors: bool = False, base_name: str = "base", return_surface_sampler: bool = False, -) -> tuple[SceneModel, np.ndarray] | tuple[SceneModel, np.ndarray, object]: - """Build a MotrixSim model with generated hfield terrain and attached robot.""" +) -> tuple[SceneModel, np.ndarray, object | None, tuple[str, ...]]: + """Build a Motrix terrain model and return its cold-path sensor names.""" import motrixsim.msd as msd from unilab.terrains import TerrainGenerator @@ -262,6 +266,63 @@ def materialize_motrix_hfield_attached_scene( if add_body_sensors: add_motrix_tracking_frame_sensors(world, base_name=base_name) + model = msd.build(world) + sampler = generated.surface_sampler() if return_surface_sampler else None + return model, generated.terrain_origins, sampler, _motrix_sensor_names(world) + + +@overload +def materialize_motrix_hfield_attached_scene( + *, + model_file: str, + terrain_cfg: TerrainGeneratorCfg, + fragment_files: Sequence[str] = (), + hfield_name: str = "terrain_hfield", + geom_name: str = "floor", + add_body_sensors: bool = False, + base_name: str = "base", + return_surface_sampler: Literal[False] = False, +) -> tuple[SceneModel, np.ndarray]: ... + + +@overload +def materialize_motrix_hfield_attached_scene( + *, + model_file: str, + terrain_cfg: TerrainGeneratorCfg, + fragment_files: Sequence[str] = (), + hfield_name: str = "terrain_hfield", + geom_name: str = "floor", + add_body_sensors: bool = False, + base_name: str = "base", + return_surface_sampler: Literal[True], +) -> tuple[SceneModel, np.ndarray, object]: ... + + +def materialize_motrix_hfield_attached_scene( + *, + model_file: str, + terrain_cfg: TerrainGeneratorCfg, + fragment_files: Sequence[str] = (), + hfield_name: str = "terrain_hfield", + geom_name: str = "floor", + add_body_sensors: bool = False, + base_name: str = "base", + return_surface_sampler: bool = False, +) -> tuple[SceneModel, np.ndarray] | tuple[SceneModel, np.ndarray, object]: + """Build a MotrixSim model with generated hfield terrain and attached robot.""" + model, origins, sampler, _ = _materialize_motrix_hfield_attached_scene_with_sensor_names( + model_file=model_file, + terrain_cfg=terrain_cfg, + fragment_files=fragment_files, + hfield_name=hfield_name, + geom_name=geom_name, + add_body_sensors=add_body_sensors, + base_name=base_name, + return_surface_sampler=return_surface_sampler, + ) if return_surface_sampler: - return msd.build(world), generated.terrain_origins, generated.surface_sampler() - return msd.build(world), generated.terrain_origins + if sampler is None: + raise RuntimeError("Motrix terrain materialization did not produce a surface sampler") + return model, origins, sampler + return model, origins diff --git a/src/unilab/base/backend/mujoco/backend.py b/src/unilab/base/backend/mujoco/backend.py index 6cb4d4e60..682a15559 100644 --- a/src/unilab/base/backend/mujoco/backend.py +++ b/src/unilab/base/backend/mujoco/backend.py @@ -2,7 +2,7 @@ import tempfile import time import weakref -from collections.abc import Sequence +from collections.abc import Callable, Sequence from concurrent.futures import ProcessPoolExecutor from dataclasses import dataclass from multiprocessing import cpu_count, current_process, get_context @@ -1350,6 +1350,15 @@ def get_sensor_data_batch(self, names: Sequence[str]) -> np.ndarray: values = [self._sensor_views[name].reshape(self._num_envs, -1) for name in sensor_names] return np.concatenate(values, axis=1) + def _bind_sensor_data_reader(self, names: tuple[str, ...]) -> Callable[[], np.ndarray]: + """Capture MuJoCo's materialized sensor slices for hot-path reads.""" + sensor_views = tuple(self._sensor_views[name].reshape(self._num_envs, -1) for name in names) + + def read() -> np.ndarray: + return np.concatenate(sensor_views, axis=1) + + return read + def get_site_jacobian_w( self, site_id: int, diff --git a/tests/base/test_backend_conformance.py b/tests/base/test_backend_conformance.py index 639081614..f8512ef0f 100644 --- a/tests/base/test_backend_conformance.py +++ b/tests/base/test_backend_conformance.py @@ -219,6 +219,41 @@ def test_actuation_metadata_contract(backend_type: str) -> None: np.testing.assert_array_equal(backend.get_default_dof_pos(), detached) +@pytest.mark.parametrize("backend_type", _BACKEND_PARAMS) +def test_named_sensor_view_contract(backend_type: str, monkeypatch: pytest.MonkeyPatch) -> None: + """All available adapters expose the same ordered, finite sensor view.""" + _require_backend(backend_type) + + backend = create_backend( + backend_type, + SceneCfg(model_file=_G1_SCENE), + NUM_ENVS, + SIM_DT, + base_name="pelvis", + ) + backend.materialize() + names = ("pelvis_gyro", "pelvis_local_linvel") + + view = backend.bind_sensor_data(names) + assert view.backend_type == backend.backend_type + assert view.names == names + assert view.dimensions == (3, 3) + assert view.width == 6 + values = view.read() + assert values.shape == (NUM_ENVS, 6) + assert np.isfinite(values).all() + np.testing.assert_allclose(values, backend.get_sensor_data_batch(names)) + + def fail_if_public_batch_getter_is_used(*_args, **_kwargs): + raise AssertionError("materialized sensor view called the public batch getter") + + monkeypatch.setattr(backend, "get_sensor_data_batch", fail_if_public_batch_getter_is_used) + np.testing.assert_allclose(view.read(), values) + + with pytest.raises((KeyError, ValueError), match="missing_sensor|Missing|missing"): + backend.bind_sensor_data(("missing_sensor",)) + + @pytest.mark.parametrize("backend_type", _BACKEND_PARAMS) def test_root_state_layout_contract(backend_type: str) -> None: _require_backend(backend_type) diff --git a/tests/base/test_backend_sensor_view.py b/tests/base/test_backend_sensor_view.py new file mode 100644 index 000000000..b5efd3250 --- /dev/null +++ b/tests/base/test_backend_sensor_view.py @@ -0,0 +1,126 @@ +"""Focused tests for the cold-path named sensor view contract.""" + +from __future__ import annotations + +from collections import Counter + +import numpy as np +import pytest + +from unilab.base.backend import BackendSensorView +from unilab.base.backend.base import SimBackend + + +class _SensorBackend: + backend_type = "fake" + num_envs = 2 + + def __init__(self) -> None: + self.calls: Counter[str] = Counter() + self.values = { + "contact_a": np.asarray([[1.0], [0.0]], dtype=np.float32), + "contact_b": np.asarray([[0.0], [1.0]], dtype=np.float32), + } + + def get_sensor_data(self, name: str) -> np.ndarray: + self.calls["single"] += 1 + return self.values[name] + + def get_sensor_data_batch(self, names: tuple[str, ...]) -> np.ndarray: + self.calls["batch"] += 1 + return np.concatenate([self.values[name] for name in names], axis=1) + + def _bind_sensor_data_reader(self, names: tuple[str, ...]): + batch_reader = self.get_sensor_data_batch + return lambda: batch_reader(names) + + +def test_bind_sensor_data_validates_once_and_reads_through_formal_view() -> None: + backend = _SensorBackend() + + view = SimBackend.bind_sensor_data(backend, ("contact_a", "contact_b")) # type: ignore[arg-type] + + assert isinstance(view, BackendSensorView) + assert view.names == ("contact_a", "contact_b") + assert view.dimensions == (1, 1) + assert view.width == 2 + assert backend.calls == {"single": 2, "batch": 1} + + np.testing.assert_array_equal(view.read(), [[1.0, 0.0], [0.0, 1.0]]) + assert backend.calls == {"single": 2, "batch": 2} + + +@pytest.mark.parametrize( + "names, error, match", + [ + ("contact_a", TypeError, "sequence of strings"), + ((), ValueError, "at least one name"), + (("",), ValueError, "non-empty strings"), + (("contact_a", "contact_a"), ValueError, "unique"), + ], +) +def test_bind_sensor_data_rejects_invalid_name_requests( + names: str | tuple[str, ...], error: type[Exception], match: str +) -> None: + with pytest.raises(error, match=match): + SimBackend.bind_sensor_data(_SensorBackend(), names) # type: ignore[arg-type] + + +def test_bind_sensor_data_fails_closed_on_unknown_sensor() -> None: + with pytest.raises(KeyError, match="cannot bind sensor 'missing'"): + SimBackend.bind_sensor_data(_SensorBackend(), ("missing",)) # type: ignore[arg-type] + + +@pytest.mark.parametrize( + "value, error, match", + [ + (np.zeros((1, 2), dtype=np.float32), ValueError, "expected leading dimension 2"), + (np.full((2, 1), np.nan, dtype=np.float32), ValueError, "NaN or Inf"), + ], +) +def test_backend_sensor_view_rejects_invalid_runtime_data( + value: np.ndarray, error: type[Exception], match: str +) -> None: + backend = _SensorBackend() + backend.values["contact_a"] = value + + if value.shape[0] != backend.num_envs: + with pytest.raises(error, match=match): + SimBackend.bind_sensor_data(backend, ("contact_a",)) # type: ignore[arg-type] + return + + with pytest.raises(error, match=match): + SimBackend.bind_sensor_data(backend, ("contact_a",)) # type: ignore[arg-type] + + +@pytest.mark.parametrize( + "value, error, match", + [ + (np.zeros((2, 2), dtype=np.float32), ValueError, r"expected \(2, 1\)"), + (np.full((2, 1), np.inf, dtype=np.float32), ValueError, "NaN or Inf"), + (np.full((2, 1), "invalid", dtype=object), TypeError, "non-numeric dtype"), + ], +) +def test_backend_sensor_view_detects_runtime_contract_drift( + value: np.ndarray, error: type[Exception], match: str +) -> None: + backend = _SensorBackend() + view = SimBackend.bind_sensor_data(backend, ("contact_a",)) # type: ignore[arg-type] + + backend.values["contact_a"] = value + + with pytest.raises(error, match=match): + view.read() + + +def test_backend_sensor_view_metadata_is_immutable() -> None: + view = BackendSensorView( + backend_type="fake", + names=("sensor",), + dimensions=(1,), + num_envs=1, + _reader=lambda: np.zeros((1, 1), dtype=np.float32), + ) + + with pytest.raises(AttributeError): + view.names = ("other",) # type: ignore[misc] diff --git a/tests/base/test_motrix_backend_options.py b/tests/base/test_motrix_backend_options.py index 8a767cf9f..60dac0538 100644 --- a/tests/base/test_motrix_backend_options.py +++ b/tests/base/test_motrix_backend_options.py @@ -245,7 +245,11 @@ def _install_fake_motrix(monkeypatch, tmp_path): ), raising=False, ) - monkeypatch.setattr(scene_mod, "materialize_motrix_scene", lambda **kwargs: fake_model) + monkeypatch.setattr( + scene_mod, + "_materialize_motrix_scene_with_sensor_names", + lambda **kwargs: (fake_model, ()), + ) return mod, fake_model diff --git a/tests/base/test_sim_backend_smoke.py b/tests/base/test_sim_backend_smoke.py index fee15ceb3..3fa5192c6 100644 --- a/tests/base/test_sim_backend_smoke.py +++ b/tests/base/test_sim_backend_smoke.py @@ -275,6 +275,13 @@ def test_backend_batch_sensor_data_matches_individual_sensors(backend_type): np.testing.assert_allclose(bkd.get_sensor_data_batch(names), expected) _shape(bkd.get_sensor_data_batch(()), NUM_ENVS, 0) + view = bkd.bind_sensor_data(names) + assert view.names == names + assert view.dimensions == tuple( + int(np.asarray(bkd.get_sensor_data(name)).reshape(NUM_ENVS, -1).shape[1]) for name in names + ) + np.testing.assert_allclose(view.read(), expected) + def test_mujoco_model_properties_smoke(): from unilab.base.backend.mujoco.backend import MuJoCoBackend From edb5714eeedea3dffbd7ae971f11b93d4243cce0 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 12:00:59 +0800 Subject: [PATCH 017/169] feat(env): add named sensor observation terms (#1079) --- .../1-architecture/6-manager_based_api.md | 3 + .../1-architecture/6-manager_based_api.md | 2 + src/unilab/base/entity.py | 18 ++- src/unilab/envs/mdp/__init__.py | 6 + src/unilab/envs/mdp/observations.py | 73 +++++++++ src/unilab/managers/_types.py | 20 +++ tests/envs/mdp/test_observations.py | 152 +++++++++++++++++- 7 files changed, 272 insertions(+), 2 deletions(-) diff --git a/docs/sphinx/source/en/4-developer_guide/1-architecture/6-manager_based_api.md b/docs/sphinx/source/en/4-developer_guide/1-architecture/6-manager_based_api.md index 6090d9117..2c579b7c7 100644 --- a/docs/sphinx/source/en/4-developer_guide/1-architecture/6-manager_based_api.md +++ b/docs/sphinx/source/en/4-developer_guide/1-architecture/6-manager_based_api.md @@ -17,6 +17,9 @@ The normative compatibility matrix and mechanical migration example are in - `SceneEntityCfg` resolves through a base-owned scene/entity facade on the cold path. The facade uses only the public `SimBackend` contract, and hot paths reuse cached IDs and views. +- Named-sensor observation terms bind a backend-owned view through + `EntityScene.bind_sensor_data(...)` during construction; their hot path only reads + that view and never re-resolves sensor names or XML/model metadata. - Explicitly empty configuration may use a Null manager. A requested capability that is unavailable fails at the nearest boundary; it is never skipped, zero-filled, or routed back to a legacy environment. diff --git a/docs/sphinx/source/zh_CN/4-developer_guide/1-architecture/6-manager_based_api.md b/docs/sphinx/source/zh_CN/4-developer_guide/1-architecture/6-manager_based_api.md index 9cba1ead7..32d651373 100644 --- a/docs/sphinx/source/zh_CN/4-developer_guide/1-architecture/6-manager_based_api.md +++ b/docs/sphinx/source/zh_CN/4-developer_guide/1-architecture/6-manager_based_api.md @@ -14,6 +14,8 @@ NumPy，并保留现有 `NpEnvState`、Hydra owner YAML、`SimBackend`、registr Torch、Warp、runner、learner 或 IPC。 - `SceneEntityCfg` 在冷路径通过 base scene/entity facade 解析；facade 只调用正式 `SimBackend` contract，热路径复用缓存 ID/view。 +- named-sensor observation term 在构造时通过 `EntityScene.bind_sensor_data(...)` 绑定 + backend-owned view；热路径只读该 view，不重复解析 sensor 名称或 XML/model metadata。 - 用户显式空配置可以使用 Null manager；配置请求但 runtime/backend 不支持的能力必须在 最近边界报错，不能 warning、skip、返回零或回退旧 env。 - 热路径避免明显的重复解析、逐环境 Python 循环、复制和临时分配；进一步优化需要 diff --git a/src/unilab/base/entity.py b/src/unilab/base/entity.py index 02a507ba1..61949f2f2 100644 --- a/src/unilab/base/entity.py +++ b/src/unilab/base/entity.py @@ -15,7 +15,7 @@ import numpy as np -from unilab.base.backend.base import BackendRootStateLayout, SimBackend +from unilab.base.backend.base import BackendRootStateLayout, BackendSensorView, SimBackend from unilab.utils.rotation import np_quat_apply, np_quat_apply_inverse, np_yaw_from_quat if TYPE_CHECKING: @@ -1198,6 +1198,7 @@ def __init__( *, reset_state: ResetStateTransaction | None = None, ) -> None: + self._backend = backend materialized: dict[str, Entity] = {} for name, cfg in entities.items(): if not isinstance(name, str) or not name: @@ -1248,6 +1249,21 @@ def reset_to_default(self, env_ids: np.ndarray, *, term_name: str) -> None: ) self._reset_state.reset_to_default(env_ids, term_name=term_name) + def bind_sensor_data(self, names: Sequence[str]) -> BackendSensorView: + """Bind existing backend sensors for a manager term on the cold path. + + The returned view owns the backend-specific reader. Terms retain that + view and only call :meth:`BackendSensorView.read` while stepping, so the + scene facade never exposes a backend model, data object, or native handle. + """ + try: + return self._backend.bind_sensor_data(names) + except (KeyError, TypeError, ValueError, NotImplementedError) as exc: + raise type(exc)( + "Manager scene named-sensor capability on backend " + f"'{self._backend.backend_type}': {exc}" + ) from exc + def __getitem__(self, name: str) -> Entity: try: return self._entities[name] diff --git a/src/unilab/envs/mdp/__init__.py b/src/unilab/envs/mdp/__init__.py index 1a9262dc7..3b7fa400f 100644 --- a/src/unilab/envs/mdp/__init__.py +++ b/src/unilab/envs/mdp/__init__.py @@ -9,11 +9,15 @@ from unilab.envs.mdp.events import resolve_env_ids as resolve_env_ids from unilab.envs.mdp.observations import base_ang_vel as base_ang_vel from unilab.envs.mdp.observations import base_lin_vel as base_lin_vel +from unilab.envs.mdp.observations import builtin_sensor as builtin_sensor from unilab.envs.mdp.observations import generated_commands as generated_commands from unilab.envs.mdp.observations import joint_pos_rel as joint_pos_rel from unilab.envs.mdp.observations import joint_vel_rel as joint_vel_rel from unilab.envs.mdp.observations import last_action as last_action from unilab.envs.mdp.observations import projected_gravity as projected_gravity +from unilab.envs.mdp.observations import ( + projected_gravity_from_sensor as projected_gravity_from_sensor, +) from unilab.envs.mdp.rewards import action_acc_l2 as action_acc_l2 from unilab.envs.mdp.rewards import action_rate_l2 as action_rate_l2 from unilab.envs.mdp.rewards import ( @@ -40,6 +44,7 @@ "action_rate_l2", "base_ang_vel", "base_lin_vel", + "builtin_sensor", "bad_orientation", "body_angular_velocity_penalty", "flat_orientation_l2", @@ -51,6 +56,7 @@ "is_alive", "is_terminated", "projected_gravity", + "projected_gravity_from_sensor", "reset_root_state_uniform", "reset_scene_to_default", "resolve_env_ids", diff --git a/src/unilab/envs/mdp/observations.py b/src/unilab/envs/mdp/observations.py index f1de28afe..df3de0ca9 100644 --- a/src/unilab/envs/mdp/observations.py +++ b/src/unilab/envs/mdp/observations.py @@ -10,6 +10,7 @@ import numpy as np +from unilab.managers.manager_base import ManagerTermBase, ManagerTermBaseCfg from unilab.managers.scene_entity_config import SceneEntityCfg if TYPE_CHECKING: @@ -20,6 +21,76 @@ _DEFAULT_ASSET_CFG = SceneEntityCfg("robot") +class _NamedSensorObservation(ManagerTermBase): + """Cold-path binding shared by pinned named-sensor observation terms.""" + + _term_name = "named_sensor" + + def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv): + super().__init__(env) + sensor_name = cfg.params.get("sensor_name") + if not isinstance(sensor_name, str) or not sensor_name: + raise ValueError( + f"Observation term '{self._term_name}' capability 'named sensor' " + "requires a non-empty sensor_name" + ) + self._sensor_name = sensor_name + try: + self._view = env.scene.bind_sensor_data((sensor_name,)) + except (KeyError, TypeError, ValueError, NotImplementedError) as exc: + raise type(exc)( + f"Observation term '{self._term_name}' capability 'named sensor " + f"{sensor_name}' could not be materialized: {exc}" + ) from exc + + def _validate_call_name(self, sensor_name: str) -> None: + if sensor_name != self._sensor_name: + raise ValueError( + f"Observation term '{self._term_name}' was bound to sensor " + f"'{self._sensor_name}', received '{sensor_name}'" + ) + + def _read(self) -> np.ndarray: + try: + return self._view.read() + except (KeyError, TypeError, ValueError, NotImplementedError) as exc: + raise type(exc)( + f"Observation term '{self._term_name}' capability 'named sensor " + f"{self._sensor_name}' failed on backend '{self._view.backend_type}': {exc}" + ) from exc + + +class builtin_sensor(_NamedSensorObservation): + """Read one existing backend sensor through a cached NumPy view.""" + + _term_name = "builtin_sensor" + + def __call__(self, env: ManagerBasedRlEnv, sensor_name: str) -> np.ndarray: + del env + self._validate_call_name(sensor_name) + return self._read() + + +class projected_gravity_from_sensor(_NamedSensorObservation): + """Negate a cached 3-D up-vector sensor to obtain projected gravity.""" + + _term_name = "projected_gravity_from_sensor" + + def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv): + super().__init__(cfg, env) + if self._view.dimensions != (3,): + raise ValueError( + "Observation term 'projected_gravity_from_sensor' capability " + f"'3-D named sensor {self._sensor_name}' received dimensions " + f"{self._view.dimensions} on backend '{self._view.backend_type}'" + ) + + def __call__(self, env: ManagerBasedRlEnv, sensor_name: str) -> np.ndarray: + del env + self._validate_call_name(sensor_name) + return -self._read() + + def base_lin_vel( env: ManagerBasedRlEnv, asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG, @@ -88,9 +159,11 @@ def generated_commands(env: ManagerBasedRlEnv, command_name: str) -> np.ndarray: __all__ = [ "base_ang_vel", "base_lin_vel", + "builtin_sensor", "generated_commands", "joint_pos_rel", "joint_vel_rel", "last_action", "projected_gravity", + "projected_gravity_from_sensor", ] diff --git a/src/unilab/managers/_types.py b/src/unilab/managers/_types.py index 3aeb02b12..df3af5c0c 100644 --- a/src/unilab/managers/_types.py +++ b/src/unilab/managers/_types.py @@ -135,6 +135,24 @@ def write_joint_state_to_sim( ) -> None: ... +class ManagerSensorView(Protocol): + """Backend-owned named-sensor view retained by a manager term.""" + + @property + def backend_type(self) -> str: ... + + @property + def names(self) -> tuple[str, ...]: ... + + @property + def dimensions(self) -> tuple[int, ...]: ... + + @property + def data(self) -> np.ndarray: ... + + def read(self) -> np.ndarray: ... + + class ManagerScene(Protocol): """Minimal name-addressable scene surface consumed by managers.""" @@ -146,6 +164,8 @@ def env_origins(self) -> np.ndarray: ... def __getitem__(self, name: str) -> ManagerEntity: ... + def bind_sensor_data(self, names: Sequence[str]) -> ManagerSensorView: ... + def reset_to_default(self, env_ids: np.ndarray, *, term_name: str) -> None: ... diff --git a/tests/envs/mdp/test_observations.py b/tests/envs/mdp/test_observations.py index b64106587..83092c79f 100644 --- a/tests/envs/mdp/test_observations.py +++ b/tests/envs/mdp/test_observations.py @@ -3,6 +3,7 @@ from __future__ import annotations import ast +from collections import Counter from pathlib import Path from types import SimpleNamespace from typing import Any, cast @@ -24,6 +25,11 @@ class _Backend: num_actuators = 0 def __init__(self) -> None: + self.sensor_calls: Counter[str] = Counter() + self.sensor_values = { + "gyro": np.asarray([[0.1, 0.2, 0.3], [-0.1, -0.2, -0.3]], dtype=np.float32), + "upvector": np.asarray([[0.0, 0.0, 1.0], [0.0, 1.0, 0.0]], dtype=np.float32), + } self.joint_names = ("hip", "knee", "ankle") self.dof_pos = np.asarray( [[0.4, 0.1, -0.2], [0.0, 0.3, 0.7]], @@ -85,6 +91,25 @@ def get_body_lin_vel_b(self, ids: np.ndarray) -> np.ndarray: def get_body_ang_vel_b(self, ids: np.ndarray) -> np.ndarray: return self.body_ang_vel_b[:, ids] + def get_sensor_data(self, name: str) -> np.ndarray: + self.sensor_calls["single"] += 1 + try: + return self.sensor_values[name] + except KeyError as exc: + raise KeyError(f"unknown sensor {name!r}") from exc + + def get_sensor_data_batch(self, names: tuple[str, ...]) -> np.ndarray: + self.sensor_calls["batch"] += 1 + values = [self.sensor_values[name].reshape(self.num_envs, -1) for name in names] + return np.concatenate(values, axis=1) + + def _bind_sensor_data_reader(self, names: tuple[str, ...]): + return lambda: self.get_sensor_data_batch(names) + + def bind_sensor_data(self, names): + self.sensor_calls["bind"] += 1 + return SimBackend.bind_sensor_data(self, names) # type: ignore[arg-type] + class _ActionManager: def __init__(self) -> None: @@ -196,6 +221,128 @@ def test_scene_entity_selector_is_resolved_once_by_observation_manager() -> None assert resolved.joint_ids == [2, 0] +def _sensor_manager(env: ManagerBasedRlEnv) -> ObservationManager: + return ObservationManager( + { + "policy": ObservationGroupCfg( + terms={ + "gyro": ObservationTermCfg( + func=mdp.builtin_sensor, + params={"sensor_name": "gyro"}, + ), + "gravity": ObservationTermCfg( + func=mdp.projected_gravity_from_sensor, + params={"sensor_name": "upvector"}, + ), + } + ) + }, + env, + ) + + +def test_named_sensor_terms_bind_once_and_only_read_cached_views() -> None: + env, backend = _env() + + manager = _sensor_manager(env) + + assert manager.group_obs_dim == {"policy": (6,)} + assert backend.sensor_calls == {"bind": 2, "single": 2, "batch": 4} + + first = manager.compute_group("policy") + second = manager.compute_group("policy") + assert isinstance(first, np.ndarray) + assert isinstance(second, np.ndarray) + expected = np.concatenate( + [backend.sensor_values["gyro"], -backend.sensor_values["upvector"]], axis=1 + ) + np.testing.assert_array_equal(first, expected) + np.testing.assert_array_equal(second, expected) + assert backend.sensor_calls == {"bind": 2, "single": 2, "batch": 8} + + gyro_term = manager.get_term_cfg("policy", "gyro").func + gravity_term = manager.get_term_cfg("policy", "gravity").func + assert isinstance(gyro_term, mdp.builtin_sensor) + assert isinstance(gravity_term, mdp.projected_gravity_from_sensor) + + +@pytest.mark.parametrize( + ("sensor_name", "error", "message"), + [ + ("", ValueError, "builtin_sensor.*non-empty sensor_name"), + ("missing", KeyError, "builtin_sensor.*missing.*backend.*fake"), + ], +) +def test_named_sensor_term_materialization_fails_closed( + sensor_name: str, error: type[Exception], message: str +) -> None: + env, _ = _env() + with pytest.raises(error, match=message): + ObservationManager( + { + "policy": ObservationGroupCfg( + terms={ + "sensor": ObservationTermCfg( + func=mdp.builtin_sensor, + params={"sensor_name": sensor_name}, + ) + } + ) + }, + env, + ) + + +def test_named_sensor_scene_seam_rejects_duplicate_names() -> None: + env, _ = _env() + with pytest.raises(ValueError, match="named-sensor.*unique"): + env.scene.bind_sensor_data(("gyro", "gyro")) + + +def test_projected_gravity_sensor_requires_one_three_dimensional_view() -> None: + env, backend = _env() + backend.sensor_values["upvector"] = np.zeros((2, 2), dtype=np.float32) + + with pytest.raises( + ValueError, + match=r"projected_gravity_from_sensor.*3-D.*dimensions \(2,\).*backend 'fake'", + ): + ObservationManager( + { + "policy": ObservationGroupCfg( + terms={ + "gravity": ObservationTermCfg( + func=mdp.projected_gravity_from_sensor, + params={"sensor_name": "upvector"}, + ) + } + ) + }, + env, + ) + + +def test_named_sensor_runtime_drift_reports_term_and_backend() -> None: + env, backend = _env() + manager = ObservationManager( + { + "policy": ObservationGroupCfg( + terms={ + "gyro": ObservationTermCfg( + func=mdp.builtin_sensor, + params={"sensor_name": "gyro"}, + ) + } + ) + }, + env, + ) + backend.sensor_values["gyro"] = np.full((2, 3), np.nan, dtype=np.float32) + + with pytest.raises(ValueError, match="builtin_sensor.*gyro.*backend 'fake'.*NaN or Inf"): + manager.compute_group("policy") + + @pytest.mark.parametrize( ("call", "message"), [ @@ -204,7 +351,10 @@ def test_scene_entity_selector_is_resolved_once_by_observation_manager() -> None lambda env: mdp.generated_commands(env, "missing"), "Command term 'missing' not found", ), - (lambda env: mdp.joint_pos_rel(env, biased=1), "biased must be bool"), + ( + lambda env: mdp.joint_pos_rel(env, biased=1), # type: ignore[arg-type] + "biased must be bool", + ), ], ) def test_invalid_term_requests_fail_explicitly(call, message: str) -> None: From e0f4a5556702d4fabac860dc9ffaded967e185df Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 12:32:03 +0800 Subject: [PATCH 018/169] fix(env): materialize manager backend before reset (#1081) --- .../1-architecture/6-manager_based_api.md | 3 + .../1-architecture/6-manager_based_api.md | 2 + src/unilab/envs/manager_based_rl_env.py | 16 +++ tests/envs/test_manager_based_rl_env.py | 106 +++++++++++++++++- 4 files changed, 125 insertions(+), 2 deletions(-) diff --git a/docs/sphinx/source/en/4-developer_guide/1-architecture/6-manager_based_api.md b/docs/sphinx/source/en/4-developer_guide/1-architecture/6-manager_based_api.md index 2c579b7c7..d22db32e1 100644 --- a/docs/sphinx/source/en/4-developer_guide/1-architecture/6-manager_based_api.md +++ b/docs/sphinx/source/en/4-developer_guide/1-architecture/6-manager_based_api.md @@ -20,6 +20,9 @@ The normative compatibility matrix and mechanical migration example are in - Named-sensor observation terms bind a backend-owned view through `EntityScene.bind_sensor_data(...)` during construction; their hot path only reads that view and never re-resolves sensor names or XML/model metadata. +- `ManagerBasedRlEnv` owns backend materialization exactly once: manager construction + and startup events run first, then `SimBackend.materialize()` completes before any + reset or step can execute. - Explicitly empty configuration may use a Null manager. A requested capability that is unavailable fails at the nearest boundary; it is never skipped, zero-filled, or routed back to a legacy environment. diff --git a/docs/sphinx/source/zh_CN/4-developer_guide/1-architecture/6-manager_based_api.md b/docs/sphinx/source/zh_CN/4-developer_guide/1-architecture/6-manager_based_api.md index 32d651373..994c9623b 100644 --- a/docs/sphinx/source/zh_CN/4-developer_guide/1-architecture/6-manager_based_api.md +++ b/docs/sphinx/source/zh_CN/4-developer_guide/1-architecture/6-manager_based_api.md @@ -16,6 +16,8 @@ NumPy，并保留现有 `NpEnvState`、Hydra owner YAML、`SimBackend`、registr `SimBackend` contract，热路径复用缓存 ID/view。 - named-sensor observation term 在构造时通过 `EntityScene.bind_sensor_data(...)` 绑定 backend-owned view；热路径只读该 view，不重复解析 sensor 名称或 XML/model metadata。 +- `ManagerBasedRlEnv` 恰好拥有一次 backend 物化：先完成 manager 构造和 startup event， + 再调用 `SimBackend.materialize()`，任何 reset/step 都不能在物化前执行。 - 用户显式空配置可以使用 Null manager；配置请求但 runtime/backend 不支持的能力必须在 最近边界报错，不能 warning、skip、返回零或回退旧 env。 - 热路径避免明显的重复解析、逐环境 Python 循环、复制和临时分配；进一步优化需要 diff --git a/src/unilab/envs/manager_based_rl_env.py b/src/unilab/envs/manager_based_rl_env.py index 9999e65aa..26d1f8172 100644 --- a/src/unilab/envs/manager_based_rl_env.py +++ b/src/unilab/envs/manager_based_rl_env.py @@ -214,6 +214,22 @@ def __init__(self, cfg: ManagerBasedRlEnvCfg, backend: SimBackend, num_envs: int if "startup" in self.event_manager.available_modes: self.event_manager.apply(mode="startup") + self._materialize_backend() + + def _materialize_backend(self) -> None: + """Finalize backend runtime resources before the first reset or step.""" + try: + self._backend.materialize() + except NotImplementedError as exc: + raise NotImplementedError( + "ManagerBasedRlEnv lifecycle capability 'SimBackend.materialize' is " + f"unavailable on backend '{self._backend.backend_type}': {exc}" + ) from exc + except Exception as exc: + raise RuntimeError( + "ManagerBasedRlEnv failed to materialize backend " + f"'{self._backend.backend_type}' after startup events: {exc}" + ) from exc @property def physics_dt(self) -> float: diff --git a/tests/envs/test_manager_based_rl_env.py b/tests/envs/test_manager_based_rl_env.py index 3a2b77250..16e3f17fa 100644 --- a/tests/envs/test_manager_based_rl_env.py +++ b/tests/envs/test_manager_based_rl_env.py @@ -9,6 +9,8 @@ import pytest import unilab.envs.manager_based_rl_env as manager_env_module +from unilab.assets import ASSETS_ROOT_PATH +from unilab.base.backend import create_backend, env_backend_kwargs from unilab.base.backend.base import SimBackend from unilab.base.entity import EntityCfg from unilab.base.scene import SceneCfg @@ -39,13 +41,28 @@ class _FakeBackend: backend_type = "fake" - def __init__(self, num_envs: int, *, reject_pre_step: bool = False) -> None: + def __init__( + self, + num_envs: int, + *, + reject_pre_step: bool = False, + reject_materialize: bool = False, + ) -> None: self.num_envs = num_envs self.num_actuators = 1 self.reject_pre_step = reject_pre_step + self.reject_materialize = reject_materialize self.pre_step_control = None self.applied_controls: list[np.ndarray] = [] self.cleanup_calls = 0 + self.materialize_calls = 0 + self.lifecycle: list[str] = [] + + def materialize(self) -> None: + if self.reject_materialize: + raise RuntimeError("pool construction failed") + self.materialize_calls += 1 + self.lifecycle.append("materialize") def get_actuator_names(self) -> tuple[str, ...]: return ("motor",) @@ -65,6 +82,7 @@ def set_pre_step_control(self, fn) -> None: self.pre_step_control = fn def step(self, ctrl: np.ndarray, nsteps: int = 1) -> None: + assert self.materialize_calls == 1 native = ctrl for _ in range(nsteps): if self.pre_step_control is not None: @@ -129,6 +147,7 @@ def set_state( randomization=None, ) -> None: assert randomization is None + assert self.materialize_calls == 1 self.set_state_calls.append((env_ids.copy(), qpos.copy(), qvel.copy())) @@ -220,6 +239,10 @@ def _critic_obs(env: _TestEnv) -> np.ndarray: return env.episode_length_buf[:, None].astype(np.float32) +def _episode_step_observation(env: ManagerBasedRlEnv) -> np.ndarray: + return env.episode_length_buf[:, None].astype(np.float32) + + def _reward(env: _TestEnv) -> np.ndarray: return env.action_manager.action[:, 0].copy() @@ -246,6 +269,12 @@ def _event(env: _TestEnv, env_ids: np.ndarray | None) -> None: env.trace.append(("event", rendered_ids)) +def _startup_event(env: _TestEnv, env_ids: np.ndarray | None) -> None: + assert env_ids is None + backend = cast(_FakeBackend, env._backend) + backend.lifecycle.append("startup") + + def _observe_uncommitted_reset(env: _TestEnv, env_ids: np.ndarray | None) -> None: del env_ids backend = cast(_ResetBackend, env._backend) @@ -311,8 +340,13 @@ def _make_env( *, num_envs: int = 2, reject_pre_step: bool = False, + reject_materialize: bool = False, ) -> tuple[_TestEnv, _FakeBackend]: - backend = _FakeBackend(num_envs, reject_pre_step=reject_pre_step) + backend = _FakeBackend( + num_envs, + reject_pre_step=reject_pre_step, + reject_materialize=reject_materialize, + ) env = _TestEnv( cfg or _make_cfg(), cast(SimBackend, backend), @@ -377,6 +411,74 @@ def wrapped(*args, _name=name, _original=original, **kwargs): assert order == list(names) +def test_backend_materializes_once_after_startup_and_before_runtime() -> None: + cfg = _make_cfg() + cfg.events = { + "startup": EventTermCfg(func=_startup_event, mode="startup"), + "reset": EventTermCfg(func=_event, mode="reset"), + } + env, backend = _make_env(cfg) + + assert backend.lifecycle == ["startup", "materialize"] + assert backend.materialize_calls == 1 + + env.reset() + env.step(np.zeros((2, 1), dtype=np.float32)) + env.reset() + + assert backend.materialize_calls == 1 + + +def test_backend_materialization_failure_has_lifecycle_context() -> None: + with pytest.raises( + RuntimeError, + match="ManagerBasedRlEnv failed to materialize backend 'fake'.*pool construction failed", + ): + _make_env(reject_materialize=True) + + +def test_real_mujoco_backend_is_materialized_before_first_reset() -> None: + scene = SceneCfg( + model_file=str(ASSETS_ROOT_PATH / "robots" / "go2" / "scene_flat.xml"), + entities={"robot": EntityCfg()}, + ) + cfg = ManagerBasedRlEnvCfg( + scene=scene, + sim_dt=0.01, + ctrl_dt=0.02, + max_episode_seconds=1.0, + observations={ + "actor": ObservationGroupCfg( + terms={"state": ObservationTermCfg(func=_episode_step_observation)} + ) + }, + actions={}, + events={"reset_default": EventTermCfg(func=mdp.reset_scene_to_default, mode="reset")}, + rewards={"alive": RewardTermCfg(func=mdp.is_alive, weight=1.0)}, + terminations={"time_out": TerminationTermCfg(func=mdp.time_out, time_out=True)}, + policy_observation_group="actor", + ) + backend = create_backend( + "mujoco", + scene, + 2, + cfg.sim_dt, + base_name="base", + **env_backend_kwargs(cfg), + ) + env = ManagerBasedRlEnv(cfg, backend, 2) + try: + state = env.init_state() + assert state.obs["obs"].shape == (2, 1) + assert np.isfinite(state.obs["obs"]).all() + + state = env.step(np.empty((2, 0), dtype=np.float32)) + assert np.isfinite(state.obs["obs"]).all() + assert np.isfinite(state.reward).all() + finally: + env.close() + + def test_np_env_owns_substeps_autoreset_and_final_observation() -> None: env, backend = _make_env() initial_obs, initial_info = env.reset() From 247bc30bd2bf6ea822cdf1de1d81d3cc7be2586b Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 14:15:22 +0800 Subject: [PATCH 019/169] feat(env): add manager gait phase foot terms (#1083) --- .../envs/locomotion/common/manager_terms.py | 234 ++++++++++++++++++ .../locomotion/test_manager_gait_terms.py | 166 +++++++++++++ 2 files changed, 400 insertions(+) create mode 100644 src/unilab/envs/locomotion/common/manager_terms.py create mode 100644 tests/envs/locomotion/test_manager_gait_terms.py diff --git a/src/unilab/envs/locomotion/common/manager_terms.py b/src/unilab/envs/locomotion/common/manager_terms.py new file mode 100644 index 000000000..69870c96c --- /dev/null +++ b/src/unilab/envs/locomotion/common/manager_terms.py @@ -0,0 +1,234 @@ +"""Task-owned Manager-Based terms for the quadruped locomotion pilots. + +The equations come from UniLab's existing Go1/Go2 joystick tasks. The adaptation +uses community ``func + params`` terms, NumPy, and the base-owned sensor facade. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, ClassVar, Protocol + +import numpy as np + +from unilab.dtype_config import get_global_dtype +from unilab.managers.manager_base import ManagerTermBase, ManagerTermBaseCfg + +if TYPE_CHECKING: + from unilab.managers._types import ManagerBasedRlEnv, ManagerSensorView + + class _GaitEnv(ManagerBasedRlEnv, Protocol): + @property + def common_step_counter(self) -> int: ... + + +_OFFSETS = (0.0, 0.5, 0.5, 0.0) + + +def _real( + term: str, + name: str, + value: Any, + *, + minimum: float | None = None, + maximum: float | None = None, + strict_minimum: bool = False, +) -> float: + if isinstance(value, (bool, np.bool_)) or not isinstance(value, (int, float, np.number)): + raise TypeError(f"{term} {name} must be a real number") + result = float(value) + if not np.isfinite(result): + raise ValueError(f"{term} {name} must be finite") + if minimum is not None and (result <= minimum if strict_minimum else result < minimum): + relation = "greater than" if strict_minimum else "at least" + raise ValueError(f"{term} {name} must be {relation} {minimum}") + if maximum is not None and result > maximum: + raise ValueError(f"{term} {name} must be at most {maximum}") + return result + + +def _offsets(term: str, value: Any) -> np.ndarray: + if isinstance(value, (str, bytes)) or not isinstance(value, (tuple, list, np.ndarray)): + raise TypeError(f"{term} phase_offsets must be a sequence of four real numbers") + if isinstance(value, np.ndarray) and value.ndim != 1: + raise ValueError(f"{term} phase_offsets must be one-dimensional, got {value.shape}") + items = list(value) + if len(items) != 4: + raise ValueError(f"{term} phase_offsets must contain 4 values, got {len(items)}") + result = np.asarray( + [_real(term, f"phase_offsets[{index}]", item) for index, item in enumerate(items)], + dtype=get_global_dtype(), + ) + result.setflags(write=False) + return result + + +def _names(term: str, value: Any) -> tuple[str, ...]: + if isinstance(value, (str, bytes)) or not isinstance(value, (tuple, list)): + raise TypeError(f"{term} sensor_names must be a sequence of four strings") + names = tuple(value) + if len(names) != 4: + raise ValueError(f"{term} sensor_names must contain 4 names, got {len(names)}") + if any(not isinstance(name, str) or not name for name in names): + raise ValueError(f"{term} sensor_names must contain non-empty strings") + if len(set(names)) != 4: + raise ValueError(f"{term} sensor_names must be unique: {names}") + return names + + +class _GaitTerm(ManagerTermBase): + _allowed_params: ClassVar[frozenset[str]] = frozenset({"frequency", "phase_offsets"}) + + def __init__(self, cfg: ManagerTermBaseCfg, env: _GaitEnv): + super().__init__(env) + unexpected = set(cfg.params) - self._allowed_params + if unexpected: + raise TypeError(f"{self.name} received unsupported parameters: {sorted(unexpected)}") + self._frequency = _real( + self.name, "frequency", cfg.params.get("frequency", 2.0), minimum=0.0 + ) + self._offsets = _offsets(self.name, cfg.params.get("phase_offsets", _OFFSETS)) + self._step_dt = _real(self.name, "step_dt", env.step_dt, minimum=0.0, strict_minimum=True) + self._phase_value = np.asarray(0.0, dtype=get_global_dtype()) + self._last_counter = 0 + self._advance_to(self._counter(env)) + + def _counter(self, env: _GaitEnv) -> int: + counter = env.common_step_counter + if isinstance(counter, (bool, np.bool_)) or not isinstance(counter, (int, np.integer)): + raise TypeError(f"{self.name} common_step_counter must be an integer") + if counter < 0: + raise ValueError(f"{self.name} common_step_counter must be non-negative") + return int(counter) + + def _advance_to(self, counter: int) -> None: + delta = counter - self._last_counter + if delta < 0: + raise ValueError(f"{self.name} common_step_counter cannot move backwards") + increment = np.asarray(self._step_dt * self._frequency, dtype=get_global_dtype()) + if delta == 1: # Hot path: preserve the legacy float32 iterative phase exactly. + self._phase_value = np.fmod(self._phase_value + increment, 1.0) + elif delta > 1: # Cold catch-up for a term constructed or inspected between steps. + for _ in range(delta): + self._phase_value = np.fmod(self._phase_value + increment, 1.0) + self._last_counter = counter + + def _phase(self, env: _GaitEnv) -> np.ndarray: + self._advance_to(self._counter(env)) + phase = np.remainder(self._phase_value + self._offsets, 1.0).astype( + get_global_dtype(), copy=False + ) + return np.broadcast_to(phase, (env.num_envs, 4)).copy() + + +class quadruped_gait_phase(_GaitTerm): + """Four-foot phase observation with the legacy diagonal-trot ordering.""" + + def __call__(self, env: _GaitEnv, **params: Any) -> np.ndarray: + del params + return self._phase(env) + + +class _FootSensorTerm(_GaitTerm): + def __init__(self, cfg: ManagerTermBaseCfg, env: _GaitEnv): + super().__init__(cfg, env) + sensor_names = _names(self.name, cfg.params.get("sensor_names")) + try: + self._view = env.scene.bind_sensor_data(sensor_names) + except (KeyError, TypeError, ValueError, NotImplementedError) as exc: + raise type(exc)( + f"Manager term '{self.name}' named-foot-sensor capability could not be " + f"materialized for {sensor_names}: {exc}" + ) from exc + + def _read(self) -> np.ndarray: + try: + return self._view.read() + except (KeyError, TypeError, ValueError, NotImplementedError) as exc: + raise type(exc)( + f"Manager term '{self.name}' named-foot-sensor capability failed on " + f"backend '{self._view.backend_type}': {exc}" + ) from exc + + +class feet_phase_contact(_FootSensorTerm): + """Reward foot contact matching the configured stance portion of gait phase.""" + + _allowed_params = _GaitTerm._allowed_params | { + "sensor_names", + "contact_threshold", + "stance_threshold", + } + + def __init__(self, cfg: ManagerTermBaseCfg, env: _GaitEnv): + super().__init__(cfg, env) + self._contact_threshold = _real( + self.name, "contact_threshold", cfg.params.get("contact_threshold", 0.1), minimum=0.0 + ) + self._stance_threshold = _real( + self.name, + "stance_threshold", + cfg.params.get("stance_threshold", 0.6), + minimum=0.0, + maximum=1.0, + ) + if any(width not in (1, 3) for width in self._view.dimensions): + raise ValueError( + f"{self.name} contact sensors must each expose 1-D found or 3-D force; " + f"received {self._view.dimensions} on backend '{self._view.backend_type}'" + ) + starts = np.cumsum((0, *self._view.dimensions[:-1]), dtype=np.int64) + self._columns = starts + [0 if width == 1 else 2 for width in self._view.dimensions] + + def __call__(self, env: _GaitEnv, **params: Any) -> np.ndarray: + del params + contact = self._read()[:, self._columns] > self._contact_threshold + expected = self._phase(env) < self._stance_threshold + if self._frequency < 1.0e-8: + expected.fill(True) + return np.mean(contact == expected, axis=1).astype(get_global_dtype(), copy=False) + + +class feet_phase_swing_height(_FootSensorTerm): + """Reward foot height near a target during the configured swing phase.""" + + _allowed_params = _GaitTerm._allowed_params | { + "sensor_names", + "target_height", + "kernel", + "swing_start", + } + + def __init__(self, cfg: ManagerTermBaseCfg, env: _GaitEnv): + super().__init__(cfg, env) + self._target = _real( + self.name, "target_height", cfg.params.get("target_height", 0.1), minimum=0.0 + ) + self._kernel = _real( + self.name, + "kernel", + cfg.params.get("kernel", 0.01), + minimum=0.0, + strict_minimum=True, + ) + self._swing_start = _real( + self.name, + "swing_start", + cfg.params.get("swing_start", 0.6), + minimum=0.0, + maximum=1.0, + ) + if self._view.dimensions != (3, 3, 3, 3): + raise ValueError( + f"{self.name} position sensors must each expose 3-D xyz; received " + f"{self._view.dimensions} on backend '{self._view.backend_type}'" + ) + + def __call__(self, env: _GaitEnv, **params: Any) -> np.ndarray: + del params + heights = self._read()[:, (2, 5, 8, 11)] + swing = self._phase(env) >= self._swing_start + reward = np.exp(-np.square(heights - self._target) / self._kernel) * swing + return np.mean(reward, axis=1).astype(get_global_dtype(), copy=False) + + +__all__ = ["feet_phase_contact", "feet_phase_swing_height", "quadruped_gait_phase"] diff --git a/tests/envs/locomotion/test_manager_gait_terms.py b/tests/envs/locomotion/test_manager_gait_terms.py new file mode 100644 index 000000000..b45389eba --- /dev/null +++ b/tests/envs/locomotion/test_manager_gait_terms.py @@ -0,0 +1,166 @@ +"""Focused parity tests for task-owned quadruped Manager-Based terms.""" + +from __future__ import annotations + +from collections import Counter +from types import SimpleNamespace +from typing import Any, cast + +import numpy as np +import pytest + +from unilab.base.backend.base import BackendSensorView +from unilab.envs.locomotion.common import manager_terms +from unilab.managers import ( + ObservationGroupCfg, + ObservationManager, + ObservationTermCfg, + RewardManager, + RewardTermCfg, +) +from unilab.managers._types import ManagerBasedRlEnv + +CONTACTS = ("fl_contact", "fr_contact", "rl_contact", "rr_contact") +POSITIONS = ("fl_pos", "fr_pos", "rl_pos", "rr_pos") + + +class _Scene: + def __init__(self) -> None: + self.calls: Counter[str] = Counter() + self.values = { + "fl_contact": np.array([[0.2], [0.0]], dtype=np.float32), + "fr_contact": np.array([[0, 0, 0], [0, 0, 0.2]], dtype=np.float32), + "rl_contact": np.array([[0.0], [0.2]], dtype=np.float32), + "rr_contact": np.array([[0, 0, 0.2], [0, 0, 0]], dtype=np.float32), + "fl_pos": np.array([[0, 0, 0.1], [0, 0, 0]], dtype=np.float32), + "fr_pos": np.array([[0, 0, 0.2], [0, 0, 0.1]], dtype=np.float32), + "rl_pos": np.array([[0, 0, 0.1], [0, 0, 0.2]], dtype=np.float32), + "rr_pos": np.array([[0, 0, 0], [0, 0, 0.1]], dtype=np.float32), + } + + def bind_sensor_data(self, names) -> BackendSensorView: + names = tuple(names) + self.calls["bind"] += 1 + dimensions = tuple(self.values[name].shape[1] for name in names) + + def read() -> np.ndarray: + self.calls["read"] += 1 + return np.concatenate([self.values[name] for name in names], axis=1) + + view = BackendSensorView("fake", names, dimensions, 2, read) + view.read() # Mirror SimBackend.bind_sensor_data materialization validation. + return view + + +def _env(counter: int = 0, scene: _Scene | None = None) -> ManagerBasedRlEnv: + return cast( + ManagerBasedRlEnv, + SimpleNamespace( + num_envs=2, + common_step_counter=counter, + episode_length_buf=np.array([counter, 0]), + step_dt=0.02, + scene=scene or _Scene(), + ), + ) + + +def _observations(env: ManagerBasedRlEnv, **params: Any) -> ObservationManager: + term = ObservationTermCfg( + func=manager_terms.quadruped_gait_phase, + params={"frequency": 2.0, **params}, + ) + return ObservationManager({"policy": ObservationGroupCfg(terms={"gait_phase": term})}, env) + + +def _rewards(env: ManagerBasedRlEnv) -> RewardManager: + terms: dict[str, RewardTermCfg | None] = { + "contact": RewardTermCfg( + func=manager_terms.feet_phase_contact, + weight=1.0, + params={"sensor_names": CONTACTS, "frequency": 2.0}, + ), + "swing": RewardTermCfg( + func=manager_terms.feet_phase_swing_height, + weight=1.0, + params={"sensor_names": POSITIONS, "frequency": 2.0}, + ), + } + return RewardManager(terms, env, scale_by_dt=False) + + +def test_gait_phase_matches_global_legacy_clock_and_ignores_episode_reset() -> None: + env = _env() + manager = _observations(env) + assert manager.group_obs_dim == {"policy": (4,)} + initial = manager.compute_group("policy") + assert isinstance(initial, np.ndarray) + np.testing.assert_array_equal(initial, [[0.0, 0.5, 0.5, 0.0]] * 2) + + cast(Any, env).common_step_counter = 25 + env.episode_length_buf[:] = [0, 999] # Partial reset does not reset the global clock. + advanced = manager.compute_group("policy") + assert isinstance(advanced, np.ndarray) + np.testing.assert_allclose( + advanced, [[1.1920929e-7, 0.5000001, 0.5000001, 1.1920929e-7]] * 2, atol=1e-8 + ) + + +def test_foot_rewards_match_legacy_equations_and_read_only_bound_views() -> None: + scene = _Scene() + manager = _rewards(_env(5, scene)) + assert scene.calls == {"bind": 2, "read": 2} + result = manager.compute(dt=0.02) + assert scene.calls == {"bind": 2, "read": 4} + + phase = np.array([[0.2, 0.7, 0.7, 0.2]] * 2) + contact = np.array([[True, False, False, True], [False, True, True, False]]) + heights = np.array([[0.1, 0.2, 0.1, 0.0], [0.0, 0.1, 0.2, 0.1]]) + expected = np.mean(contact == (phase < 0.6), axis=1) + expected += np.mean(np.exp(-np.square(heights - 0.1) / 0.01) * (phase >= 0.6), axis=1) + np.testing.assert_allclose(result, expected, atol=1e-7) + manager.compute(dt=0.02) + assert scene.calls == {"bind": 2, "read": 6} + + +@pytest.mark.parametrize( + ("params", "error", "match"), + [ + ({"frequency": -1}, ValueError, "frequency must be at least 0.0"), + ({"phase_offsets": (0, 0.5)}, ValueError, "must contain 4 values"), + ({"phase_offsets": (0, True, 0.5, 0)}, TypeError, "must be a real number"), + ({"unknown": 1}, TypeError, "unsupported parameters"), + ], +) +def test_gait_phase_invalid_config_fails_at_construction(params, error, match: str) -> None: + with pytest.raises(error, match=match): + _observations(_env(), **params) + + +def test_foot_sensor_contracts_fail_at_construction() -> None: + env = _env() + short = RewardTermCfg( + func=manager_terms.feet_phase_contact, + weight=1, + params={"sensor_names": CONTACTS[:3]}, + ) + with pytest.raises(ValueError, match="must contain 4 names"): + RewardManager({"feet": short}, env) + + scene = _Scene() + scene.values["fr_contact"] = np.zeros((2, 2), dtype=np.float32) + with pytest.raises(ValueError, match="must each expose 1-D found or 3-D force"): + _rewards(_env(scene=scene)) + + scene = _Scene() + scene.values["rl_pos"] = np.zeros((2, 2), dtype=np.float32) + with pytest.raises(ValueError, match="must each expose 3-D xyz"): + _rewards(_env(scene=scene)) + + +def test_runtime_nonfinite_sensor_data_reports_term_and_backend() -> None: + scene = _Scene() + manager = _rewards(_env(scene=scene)) + scene.values["fl_contact"][1, 0] = np.nan + with pytest.raises(ValueError, match="feet_phase_contact.*backend 'fake'.*NaN or Inf"): + manager.compute(dt=0.02) From 58d9d225f8a1a5b5045ca911f8a40947d3243f2c Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 14:34:03 +0800 Subject: [PATCH 020/169] feat(env): add Go2 manager reward parity terms (#1085) --- .../envs/locomotion/common/manager_terms.py | 149 +++++++++++++++++- .../locomotion/test_manager_gait_terms.py | 123 +++++++++++++++ 2 files changed, 270 insertions(+), 2 deletions(-) diff --git a/src/unilab/envs/locomotion/common/manager_terms.py b/src/unilab/envs/locomotion/common/manager_terms.py index 69870c96c..33e543e2b 100644 --- a/src/unilab/envs/locomotion/common/manager_terms.py +++ b/src/unilab/envs/locomotion/common/manager_terms.py @@ -6,14 +6,16 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, ClassVar, Protocol +from typing import TYPE_CHECKING, Any, ClassVar, Protocol, cast import numpy as np from unilab.dtype_config import get_global_dtype from unilab.managers.manager_base import ManagerTermBase, ManagerTermBaseCfg +from unilab.managers.scene_entity_config import SceneEntityCfg if TYPE_CHECKING: + from unilab.base.entity import Entity from unilab.managers._types import ManagerBasedRlEnv, ManagerSensorView class _GaitEnv(ManagerBasedRlEnv, Protocol): @@ -22,6 +24,7 @@ def common_step_counter(self) -> int: ... _OFFSETS = (0.0, 0.5, 0.5, 0.0) +_DEFAULT_ASSET_CFG = SceneEntityCfg("robot") def _real( @@ -75,6 +78,138 @@ def _names(term: str, value: Any) -> tuple[str, ...]: return names +def _state(term: str, capability: str, value: Any, shape: tuple[int, ...]) -> np.ndarray: + if not isinstance(value, np.ndarray): + raise TypeError(f"{term} {capability} must be an np.ndarray") + if value.shape != shape: + raise ValueError(f"{term} {capability} must have shape {shape}, got {value.shape}") + if not np.isfinite(value).all(): + env_ids = np.flatnonzero(~np.isfinite(value).reshape(shape[0], -1).all(axis=1)).tolist() + raise ValueError(f"{term} {capability} contains NaN or Inf for environments {env_ids[:10]}") + return value + + +def _command(env: ManagerBasedRlEnv, term: str, command_name: str) -> np.ndarray: + if not isinstance(command_name, str) or not command_name: + raise ValueError(f"{term} command_name must be a non-empty string") + try: + command = env.command_manager.get_command(command_name) + except KeyError as exc: + raise KeyError(f"{term} command capability '{command_name}' is unavailable") from exc + if command is None: + raise KeyError(f"{term} command capability '{command_name}' is unavailable") + return _state(term, f"command '{command_name}'", command, (env.num_envs, 3)) + + +def _asset(env: ManagerBasedRlEnv, asset_cfg: SceneEntityCfg) -> Entity: + return cast("Entity", env.scene[asset_cfg.name]) + + +def track_lin_vel_xy_exp( + env: ManagerBasedRlEnv, + std: float, + command_name: str, + asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG, +) -> np.ndarray: + """Track commanded planar velocity with the legacy independent exponential kernel.""" + scale = _real("track_lin_vel_xy_exp", "std", std, minimum=0.0, strict_minimum=True) + actual = _state( + "track_lin_vel_xy_exp", + "root linear velocity", + _asset(env, asset_cfg).data.root_link_lin_vel_b, + (env.num_envs, 3), + ) + error = np.sum( + np.square(_command(env, "track_lin_vel_xy_exp", command_name)[:, :2] - actual[:, :2]), + axis=1, + ) + return np.asarray(np.exp(-error / scale**2), dtype=get_global_dtype()) + + +def track_ang_vel_z_exp( + env: ManagerBasedRlEnv, + std: float, + command_name: str, + asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG, +) -> np.ndarray: + """Track commanded yaw velocity without folding roll/pitch into the kernel.""" + scale = _real("track_ang_vel_z_exp", "std", std, minimum=0.0, strict_minimum=True) + actual = _state( + "track_ang_vel_z_exp", + "root angular velocity", + _asset(env, asset_cfg).data.root_link_ang_vel_b, + (env.num_envs, 3), + ) + error = np.square(_command(env, "track_ang_vel_z_exp", command_name)[:, 2] - actual[:, 2]) + return np.asarray(np.exp(-error / scale**2), dtype=get_global_dtype()) + + +def lin_vel_z_l2( + env: ManagerBasedRlEnv, + asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG, +) -> np.ndarray: + """Penalize vertical root velocity independently from planar tracking.""" + velocity = _state( + "lin_vel_z_l2", + "root linear velocity", + _asset(env, asset_cfg).data.root_link_lin_vel_b, + (env.num_envs, 3), + ) + return np.asarray(np.square(velocity[:, 2]), dtype=get_global_dtype()) + + +def ang_vel_xy_l2( + env: ManagerBasedRlEnv, + asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG, +) -> np.ndarray: + """Penalize root roll/pitch angular velocity independently from yaw tracking.""" + velocity = _state( + "ang_vel_xy_l2", + "root angular velocity", + _asset(env, asset_cfg).data.root_link_ang_vel_b, + (env.num_envs, 3), + ) + return np.asarray(np.sum(np.square(velocity[:, :2]), axis=1), dtype=get_global_dtype()) + + +def base_height_l2( + env: ManagerBasedRlEnv, + target_height: float, + asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG, +) -> np.ndarray: + """Penalize world-frame root height error for the flat-ground pilot.""" + target = _real("base_height_l2", "target_height", target_height) + position = _state( + "base_height_l2", + "root position", + _asset(env, asset_cfg).data.root_link_pos_w, + (env.num_envs, 3), + ) + return np.asarray(np.square(position[:, 2] - target), dtype=get_global_dtype()) + + +def joint_deviation_l1( + env: ManagerBasedRlEnv, + asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG, +) -> np.ndarray: + """Penalize selected joint displacement from the default pose with an L1 kernel.""" + asset = _asset(env, asset_cfg) + position = asset.data.joint_pos[:, asset_cfg.joint_ids] + default = asset.data.default_joint_pos[:, asset_cfg.joint_ids] + if ( + not isinstance(position, np.ndarray) + or position.ndim != 2 + or position.shape[0] != env.num_envs + ): + shape = getattr(position, "shape", None) + raise ValueError( + f"joint_deviation_l1 joint position must be 2-D with leading dimension {env.num_envs}, got {shape}" + ) + default = _state("joint_deviation_l1", "default joint position", default, position.shape) + position = _state("joint_deviation_l1", "joint position", position, position.shape) + return np.asarray(np.sum(np.abs(position - default), axis=1), dtype=get_global_dtype()) + + class _GaitTerm(ManagerTermBase): _allowed_params: ClassVar[frozenset[str]] = frozenset({"frequency", "phase_offsets"}) @@ -231,4 +366,14 @@ def __call__(self, env: _GaitEnv, **params: Any) -> np.ndarray: return np.mean(reward, axis=1).astype(get_global_dtype(), copy=False) -__all__ = ["feet_phase_contact", "feet_phase_swing_height", "quadruped_gait_phase"] +__all__ = [ + "ang_vel_xy_l2", + "base_height_l2", + "feet_phase_contact", + "feet_phase_swing_height", + "joint_deviation_l1", + "lin_vel_z_l2", + "quadruped_gait_phase", + "track_ang_vel_z_exp", + "track_lin_vel_xy_exp", +] diff --git a/tests/envs/locomotion/test_manager_gait_terms.py b/tests/envs/locomotion/test_manager_gait_terms.py index b45389eba..d1cb00b5a 100644 --- a/tests/envs/locomotion/test_manager_gait_terms.py +++ b/tests/envs/locomotion/test_manager_gait_terms.py @@ -10,6 +10,7 @@ import pytest from unilab.base.backend.base import BackendSensorView +from unilab.dtype_config import get_global_dtype from unilab.envs.locomotion.common import manager_terms from unilab.managers import ( ObservationGroupCfg, @@ -19,6 +20,7 @@ RewardTermCfg, ) from unilab.managers._types import ManagerBasedRlEnv +from unilab.managers.scene_entity_config import SceneEntityCfg CONTACTS = ("fl_contact", "fr_contact", "rl_contact", "rr_contact") POSITIONS = ("fl_pos", "fr_pos", "rl_pos", "rr_pos") @@ -52,6 +54,38 @@ def read() -> np.ndarray: return view +class _ParityScene: + def __init__(self) -> None: + self.robot = SimpleNamespace( + data=SimpleNamespace( + root_link_lin_vel_b=np.array( + [[0.2, 0.1, -0.3], [0.3, -0.2, 0.5]], dtype=np.float32 + ), + root_link_ang_vel_b=np.array( + [[0.1, -0.2, 0.4], [-0.3, 0.2, -0.1]], dtype=np.float32 + ), + root_link_pos_w=np.array([[1.0, 2.0, 0.4], [-1.0, 0.5, 0.2]], dtype=np.float32), + joint_pos=np.array([[0.2, -0.1, 0.5], [-0.4, 0.3, 0.1]], dtype=np.float32), + default_joint_pos=np.array([[0.1, -0.2, 0.5], [-0.1, 0.1, 0.0]], dtype=np.float32), + ) + ) + + def __getitem__(self, name: str): + if name != "robot": + raise KeyError(name) + return self.robot + + +class _Commands: + def __init__(self) -> None: + self.command = np.array([[0.5, -0.2, 0.3], [-0.1, 0.4, -0.2]], dtype=np.float32) + + def get_command(self, name: str) -> np.ndarray: + if name != "twist": + raise KeyError(name) + return self.command + + def _env(counter: int = 0, scene: _Scene | None = None) -> ManagerBasedRlEnv: return cast( ManagerBasedRlEnv, @@ -89,6 +123,23 @@ def _rewards(env: ManagerBasedRlEnv) -> RewardManager: return RewardManager(terms, env, scale_by_dt=False) +def _parity_env() -> ManagerBasedRlEnv: + return cast( + ManagerBasedRlEnv, + SimpleNamespace( + num_envs=2, + scene=_ParityScene(), + command_manager=_Commands(), + max_episode_length_s=20.0, + ), + ) + + +def _reward_value(env: ManagerBasedRlEnv, func, **params: Any) -> np.ndarray: + cfg = RewardTermCfg(func=func, weight=1.0, params=params) + return RewardManager({"parity": cfg}, env, scale_by_dt=False).compute(dt=0.02) + + def test_gait_phase_matches_global_legacy_clock_and_ignores_episode_reset() -> None: env = _env() manager = _observations(env) @@ -164,3 +215,75 @@ def test_runtime_nonfinite_sensor_data_reports_term_and_backend() -> None: scene.values["fl_contact"][1, 0] = np.nan with pytest.raises(ValueError, match="feet_phase_contact.*backend 'fake'.*NaN or Inf"): manager.compute(dt=0.02) + + +def test_base_reward_terms_match_go2_flat_equations() -> None: + env = _parity_env() + asset = SceneEntityCfg("robot") + data = cast(Any, env).scene.robot.data + command = cast(Any, env).command_manager.command + std = 0.5 + actual = { + "track_xy": _reward_value( + env, + manager_terms.track_lin_vel_xy_exp, + std=std, + command_name="twist", + asset_cfg=asset, + ), + "track_yaw": _reward_value( + env, + manager_terms.track_ang_vel_z_exp, + std=std, + command_name="twist", + asset_cfg=SceneEntityCfg("robot"), + ), + "lin_z": _reward_value(env, manager_terms.lin_vel_z_l2, asset_cfg=SceneEntityCfg("robot")), + "ang_xy": _reward_value( + env, manager_terms.ang_vel_xy_l2, asset_cfg=SceneEntityCfg("robot") + ), + "height": _reward_value( + env, + manager_terms.base_height_l2, + target_height=0.3, + asset_cfg=SceneEntityCfg("robot"), + ), + "pose": _reward_value( + env, manager_terms.joint_deviation_l1, asset_cfg=SceneEntityCfg("robot") + ), + } + expected = { + "track_xy": np.exp( + -np.sum(np.square(command[:, :2] - data.root_link_lin_vel_b[:, :2]), axis=1) / std**2 + ), + "track_yaw": np.exp(-np.square(command[:, 2] - data.root_link_ang_vel_b[:, 2]) / std**2), + "lin_z": np.square(data.root_link_lin_vel_b[:, 2]), + "ang_xy": np.sum(np.square(data.root_link_ang_vel_b[:, :2]), axis=1), + "height": np.square(data.root_link_pos_w[:, 2] - 0.3), + "pose": np.sum(np.abs(data.joint_pos - data.default_joint_pos), axis=1), + } + for name in expected: + assert actual[name].shape == (2,) + assert actual[name].dtype == np.dtype(get_global_dtype()) + np.testing.assert_allclose(actual[name], expected[name], rtol=1e-6, atol=1e-7) + + +def test_base_reward_terms_fail_closed_at_nearest_boundary() -> None: + env = _parity_env() + with pytest.raises(ValueError, match="track_lin_vel_xy_exp std must be greater than 0.0"): + _reward_value(env, manager_terms.track_lin_vel_xy_exp, std=0.0, command_name="twist") + with pytest.raises(KeyError, match="track_ang_vel_z_exp command capability 'missing'"): + _reward_value(env, manager_terms.track_ang_vel_z_exp, std=0.5, command_name="missing") + with pytest.raises(ValueError, match="base_height_l2 target_height must be finite"): + _reward_value(env, manager_terms.base_height_l2, target_height=np.nan) + + cast(Any, env).scene.robot.data.root_link_lin_vel_b[1, 2] = np.inf + with pytest.raises(ValueError, match="lin_vel_z_l2 root linear velocity contains NaN or Inf"): + _reward_value(env, manager_terms.lin_vel_z_l2) + + with pytest.raises(KeyError, match="RewardManager term 'parity'.*asset_cfg.*missing"): + _reward_value( + _parity_env(), + manager_terms.joint_deviation_l1, + asset_cfg=SceneEntityCfg("missing"), + ) From 5cdb458784e5ffb0aca907cbb5fad5c557b9e15f Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 15:40:58 +0800 Subject: [PATCH 021/169] feat(env): support named keyframe manager defaults (#1088) --- src/unilab/base/entity.py | 84 +++++++-- src/unilab/base/reset_state.py | 18 +- src/unilab/base/scene.py | 59 +++++++ src/unilab/envs/manager_based_rl_env.py | 8 +- tests/envs/test_manager_based_rl_env.py | 224 ++++++++++++++++++++++++ 5 files changed, 369 insertions(+), 24 deletions(-) diff --git a/src/unilab/base/entity.py b/src/unilab/base/entity.py index 61949f2f2..51487acf4 100644 --- a/src/unilab/base/entity.py +++ b/src/unilab/base/entity.py @@ -412,6 +412,8 @@ def __init__( backend: SimBackend, control_buffer: np.ndarray | None = None, reset_state: ResetStateTransaction | None = None, + *, + default_qpos: np.ndarray | None = None, ) -> None: if not name: raise ValueError("Entity name must be a non-empty string") @@ -479,8 +481,12 @@ def __init__( self._reset_root_layout, default_root_state, self._reset_root_layout_error, - ) = self._materialize_root_state(backend, cfg.root_body_name) - default_joint_pos = self._materialize_default_joint_pos(backend, joint_pos_ids) + ) = self._materialize_root_state(backend, cfg.root_body_name, default_qpos) + default_joint_pos = self._materialize_default_joint_pos( + backend, + joint_pos_ids, + default_qpos, + ) default_joint_vel = self._materialize_default_joint_vel(backend, joint_vel_ids) gravity_vec_w = self._materialize_gravity_vector(backend, root_body_ids) actuator_ctrl_range = self._materialize_actuator_ctrl_range(backend, actuator_ids) @@ -669,20 +675,46 @@ def _materialize_actuator_ctrl_range( return selected def _materialize_default_joint_pos( - self, backend: SimBackend, joint_pos_ids: np.ndarray | None + self, + backend: SimBackend, + joint_pos_ids: np.ndarray | None, + default_qpos: np.ndarray | None, ) -> np.ndarray | None: if joint_pos_ids is None: return None - defaults = self._read_state("default joint position", backend.get_default_dof_pos) current = self._read_state("joint position state", backend.get_dof_pos) - if defaults.shape != current.shape[1:]: - raise ValueError( - f"Entity '{self.name}' capability 'default joint position' on backend " - f"'{self._backend_type}' returned shape {defaults.shape}; expected " - f"{current.shape[1:]} to match get_dof_pos()" + if default_qpos is None: + defaults = self._read_state("default joint position", backend.get_default_dof_pos) + if defaults.shape != current.shape[1:]: + raise ValueError( + f"Entity '{self.name}' capability 'default joint position' on backend " + f"'{self._backend_type}' returned shape {defaults.shape}; expected " + f"{current.shape[1:]} to match get_dof_pos()" + ) + selected = np.asarray(defaults[_as_column_index(joint_pos_ids)]) + else: + assert self._joint_names is not None + defaults = self._validate_root_default_vector(default_qpos, "selected default qpos") + try: + state_qpos_ids = backend.get_joint_state_qpos_indices(self._joint_names) + except (AttributeError, NotImplementedError) as exc: + raise self._capability_error("default joint-state layout", str(exc)) from exc + resolved_qpos_ids = _readonly_ids( + state_qpos_ids, + expected=len(self._joint_names), + label=f"Entity '{self.name}' default qpos", ) - selected = np.asarray(defaults[_as_column_index(joint_pos_ids)]) - materialized = np.broadcast_to(selected, (backend.num_envs, len(joint_pos_ids))).copy() + if resolved_qpos_ids.size and int(np.max(resolved_qpos_ids)) >= defaults.size: + raise ValueError( + f"Entity '{self.name}' default qpos layout exceeds backend " + f"'{self._backend_type}' width {defaults.size}: {resolved_qpos_ids.tolist()}" + ) + selected = np.asarray(defaults[_as_column_index(resolved_qpos_ids)]) + self._reset_joint_qpos_ids = resolved_qpos_ids + materialized = np.broadcast_to( + selected, + (backend.num_envs, len(joint_pos_ids)), + ).astype(current.dtype, copy=True) materialized.setflags(write=False) return materialized @@ -690,6 +722,7 @@ def _materialize_root_state( self, backend: SimBackend, root_body_name: str | None, + default_qpos: np.ndarray | None, ) -> tuple[BackendRootStateLayout | None, np.ndarray | None, str | None]: if root_body_name is None: return None, None, "root_body_name was not declared in EntityCfg" @@ -704,7 +737,7 @@ def _materialize_root_state( f"{type(layout).__name__}" ) try: - qpos = backend.get_default_qpos() + qpos = backend.get_default_qpos() if default_qpos is None else default_qpos qvel = backend.get_init_qvel() except (AttributeError, NotImplementedError) as exc: return None, None, str(exc) @@ -1071,11 +1104,15 @@ def write_joint_state_to_sim( ) def _materialize_reset_joint_indices(self) -> None: - if self._reset_joint_qpos_ids is not None: + if self._reset_joint_qpos_ids is not None and self._reset_joint_qvel_ids is not None: return assert self._joint_names is not None try: - qpos_ids = self._backend.get_joint_state_qpos_indices(self._joint_names) + qpos_ids = ( + self._reset_joint_qpos_ids + if self._reset_joint_qpos_ids is not None + else self._backend.get_joint_state_qpos_indices(self._joint_names) + ) qvel_ids = self._backend.get_joint_state_qvel_indices(self._joint_names) except (AttributeError, NotImplementedError) as exc: raise self._capability_error("reset joint-state layout", str(exc)) from exc @@ -1197,6 +1234,7 @@ def __init__( control_buffer: np.ndarray | None = None, *, reset_state: ResetStateTransaction | None = None, + default_qpos: np.ndarray | None = None, ) -> None: self._backend = backend materialized: dict[str, Entity] = {} @@ -1207,7 +1245,14 @@ def __init__( raise TypeError( f"Scene entity '{name}' must be EntityCfg, got {type(cfg).__name__}" ) - materialized[name] = Entity(name, cfg, backend, control_buffer, reset_state) + materialized[name] = Entity( + name, + cfg, + backend, + control_buffer, + reset_state, + default_qpos=default_qpos, + ) self._entities = MappingProxyType(materialized) self._reset_state = reset_state env_origins = np.zeros((backend.num_envs, 3), dtype=np.float32) @@ -1222,8 +1267,15 @@ def from_scene_cfg( control_buffer: np.ndarray | None = None, *, reset_state: ResetStateTransaction | None = None, + default_qpos: np.ndarray | None = None, ) -> EntityScene: - return cls(cfg.entities, backend, control_buffer, reset_state=reset_state) + return cls( + cfg.entities, + backend, + control_buffer, + reset_state=reset_state, + default_qpos=default_qpos, + ) @property def entities(self) -> Mapping[str, Entity]: diff --git a/src/unilab/base/reset_state.py b/src/unilab/base/reset_state.py index b28fe333f..6c94f2027 100644 --- a/src/unilab/base/reset_state.py +++ b/src/unilab/base/reset_state.py @@ -19,9 +19,15 @@ class ResetStateTransaction: """Reusable, fail-closed transaction for reset-mode state mutation.""" - def __init__(self, backend: SimBackend) -> None: + def __init__( + self, + backend: SimBackend, + *, + default_qpos: np.ndarray | None = None, + ) -> None: self._backend = backend self._num_envs = backend.num_envs + self._selected_default_qpos = default_qpos self._active = False self._active_mask = np.zeros(self._num_envs, dtype=np.bool_) self._dirty_mask = np.zeros(self._num_envs, dtype=np.bool_) @@ -260,10 +266,12 @@ def abort(self) -> None: def _materialize_default_state(self, term_name: str) -> None: if self._default_qpos is not None: return - try: - qpos = self._backend.get_default_qpos() - except (AttributeError, NotImplementedError) as exc: - raise self._capability_error(term_name, "default qpos", exc) from exc + qpos = self._selected_default_qpos + if qpos is None: + try: + qpos = self._backend.get_default_qpos() + except (AttributeError, NotImplementedError) as exc: + raise self._capability_error(term_name, "default qpos", exc) from exc try: qvel = self._backend.get_init_qvel() except (AttributeError, NotImplementedError) as exc: diff --git a/src/unilab/base/scene.py b/src/unilab/base/scene.py index 4063e45f1..641e70d38 100644 --- a/src/unilab/base/scene.py +++ b/src/unilab/base/scene.py @@ -2,10 +2,16 @@ from dataclasses import dataclass, field from pathlib import Path +from typing import TYPE_CHECKING + +import numpy as np from unilab.base.entity import EntityCfg from unilab.terrains.terrain_generator import TerrainGeneratorCfg +if TYPE_CHECKING: + from unilab.base.backend.base import SimBackend + def resolve_scene_fragment_path(fragment_file: str, model_file: Path) -> Path: """Resolve a ``SceneCfg.fragment_files`` entry against the scene model file. @@ -47,3 +53,56 @@ class SceneCfg: # scene (e.g. a per-env replicable obstacle) without touching the trained # collision model. ``None`` => render with ``model_file`` (unchanged). visual_model_file: str | None = None + default_keyframe_name: str | None = None + """Optional named keyframe used as the Manager-Based default state.""" + + +def resolve_scene_default_qpos(cfg: SceneCfg, backend: SimBackend) -> np.ndarray | None: + """Resolve one named default-qpos snapshot without changing the qpos0 path.""" + keyframe_name = cfg.default_keyframe_name + if keyframe_name is not None and not isinstance(keyframe_name, str): + raise TypeError( + "SceneCfg default_keyframe_name must be a non-empty string or None, " + f"got {type(keyframe_name).__name__}" + ) + if keyframe_name == "": + raise ValueError("SceneCfg default_keyframe_name must be a non-empty string or None") + if keyframe_name is None: + return None + + capability = f"default keyframe {keyframe_name!r} qpos" + try: + value = backend.get_keyframe_qpos(keyframe_name) + except (AttributeError, NotImplementedError) as exc: + raise NotImplementedError( + f"Manager scene default keyframe {keyframe_name!r} is unavailable on " + f"backend '{backend.backend_type}': {exc}" + ) from exc + except (KeyError, ValueError) as exc: + raise ValueError( + f"Manager scene could not resolve default keyframe {keyframe_name!r} on " + f"backend '{backend.backend_type}': {exc}" + ) from exc + + if not isinstance(value, np.ndarray): + raise TypeError( + f"Manager scene {capability} on backend '{backend.backend_type}' must return " + f"np.ndarray, got {type(value).__name__}" + ) + if value.ndim != 1: + raise ValueError( + f"Manager scene {capability} on backend '{backend.backend_type}' returned shape " + f"{value.shape}; expected 1-D" + ) + if not np.issubdtype(value.dtype, np.floating): + raise TypeError( + f"Manager scene {capability} on backend '{backend.backend_type}' must be " + f"floating, got {value.dtype}" + ) + if not np.isfinite(value).all(): + raise ValueError( + f"Manager scene {capability} on backend '{backend.backend_type}' returned NaN or Inf" + ) + resolved = np.array(value, copy=True) + resolved.setflags(write=False) + return resolved diff --git a/src/unilab/envs/manager_based_rl_env.py b/src/unilab/envs/manager_based_rl_env.py index 26d1f8172..a7b30fa1b 100644 --- a/src/unilab/envs/manager_based_rl_env.py +++ b/src/unilab/envs/manager_based_rl_env.py @@ -23,7 +23,7 @@ from unilab.base.entity import EntityScene from unilab.base.np_env import NpEnv, NpEnvState from unilab.base.reset_state import ResetStateTransaction -from unilab.base.scene import SceneCfg +from unilab.base.scene import SceneCfg, resolve_scene_default_qpos from unilab.dtype_config import get_global_dtype from unilab.managers import ( ActionManager, @@ -181,14 +181,16 @@ def __init__(self, cfg: ManagerBasedRlEnvCfg, backend: SimBackend, num_envs: int cfg.seed = actual_seed self.rng = np.random.default_rng(actual_seed) - self._control = np.zeros((num_envs, backend.num_actuators), dtype=get_global_dtype()) - self._reset_state = ResetStateTransaction(backend) assert cfg.scene is not None + default_qpos = resolve_scene_default_qpos(cfg.scene, backend) + self._control = np.zeros((num_envs, backend.num_actuators), dtype=get_global_dtype()) + self._reset_state = ResetStateTransaction(backend, default_qpos=default_qpos) self.scene = EntityScene.from_scene_cfg( cfg.scene, backend, self._control, reset_state=self._reset_state, + default_qpos=default_qpos, ) self.common_step_counter = 0 diff --git a/tests/envs/test_manager_based_rl_env.py b/tests/envs/test_manager_based_rl_env.py index 16e3f17fa..107f482d1 100644 --- a/tests/envs/test_manager_based_rl_env.py +++ b/tests/envs/test_manager_based_rl_env.py @@ -151,6 +151,32 @@ def set_state( self.set_state_calls.append((env_ids.copy(), qpos.copy(), qvel.copy())) +class _KeyframeBackend(_ResetBackend): + def __init__( + self, + num_envs: int, + *, + keyframe_qpos: Any = None, + keyframe_error: Exception | None = None, + ) -> None: + super().__init__(num_envs) + self.keyframe_qpos = ( + np.array([0.0, 0.0, 0.3, 0.75], dtype=np.float64) + if keyframe_qpos is None + else keyframe_qpos + ) + self.keyframe_error = keyframe_error + self.keyframe_qpos_calls = 0 + + def get_keyframe_qpos(self, name: str) -> Any: + self.keyframe_qpos_calls += 1 + if self.keyframe_error is not None: + raise self.keyframe_error + if name != "home": + raise ValueError(f"unknown keyframe {name!r}") + return self.keyframe_qpos + + @dataclass(kw_only=True) class _DriveCfg(ActionTermCfg): gain: float = 1.0 @@ -437,6 +463,121 @@ def test_backend_materialization_failure_has_lifecycle_context() -> None: _make_env(reject_materialize=True) +@pytest.mark.parametrize( + ("selector", "error"), + [ + (1, TypeError), + ("", ValueError), + ], +) +def test_default_keyframe_selector_fails_at_env_initialization( + selector: Any, + error: type[Exception], +) -> None: + cfg = _make_cfg(include_optional_managers=False) + cfg.scene.default_keyframe_name = selector + backend = _KeyframeBackend(2) + + with pytest.raises(error, match="default_keyframe_name.*non-empty string or None"): + _TestEnv(cfg, cast(SimBackend, backend), 2) + + assert backend.keyframe_qpos_calls == 0 + + +@pytest.mark.parametrize( + ("value", "error", "match"), + [ + ([0.0], TypeError, "must return np.ndarray"), + (np.zeros((1, 1), dtype=np.float32), ValueError, "expected 1-D"), + (np.zeros(1, dtype=np.int32), TypeError, "must be floating"), + (np.array([np.nan], dtype=np.float32), ValueError, "NaN or Inf"), + ], +) +def test_default_keyframe_qpos_contract_fails_at_env_initialization( + value: Any, + error: type[Exception], + match: str, +) -> None: + cfg = _make_cfg(include_optional_managers=False) + cfg.scene.default_keyframe_name = "home" + backend = _KeyframeBackend(2, keyframe_qpos=value) + + with pytest.raises(error, match=f"default keyframe 'home'.*backend 'fake'.*{match}"): + _TestEnv(cfg, cast(SimBackend, backend), 2) + + assert backend.keyframe_qpos_calls == 1 + + +@pytest.mark.parametrize( + ("failure", "error", "match"), + [ + ( + NotImplementedError("named keyframes disabled"), + NotImplementedError, + "default keyframe 'home'.*backend 'fake'.*named keyframes disabled", + ), + ( + ValueError("keyframe missing"), + ValueError, + "resolve default keyframe 'home'.*backend 'fake'.*keyframe missing", + ), + ], +) +def test_default_keyframe_resolution_names_backend_and_keyframe( + failure: Exception, + error: type[Exception], + match: str, +) -> None: + cfg = _make_cfg(include_optional_managers=False) + cfg.scene.default_keyframe_name = "home" + backend = _KeyframeBackend(2, keyframe_error=failure) + + with pytest.raises(error, match=match): + _TestEnv(cfg, cast(SimBackend, backend), 2) + + assert backend.keyframe_qpos_calls == 1 + + +def test_named_keyframe_snapshot_is_shared_by_entity_and_reset_cold_path() -> None: + source_qpos = np.array([0.0, 0.0, 0.3, 0.75], dtype=np.float64) + backend = _KeyframeBackend(2, keyframe_qpos=source_qpos) + cfg = _make_cfg(include_optional_managers=False) + cfg.scene.default_keyframe_name = "home" + cfg.scene.entities["robot"] = EntityCfg( + joint_names=("joint",), + actuator_names=("motor",), + ) + cfg.events = { + "reset_default": EventTermCfg(func=mdp.reset_scene_to_default, mode="reset"), + } + + env = _TestEnv(cfg, cast(SimBackend, backend), 2) + selected_qpos = env._reset_state._selected_default_qpos + + assert backend.keyframe_qpos_calls == 1 + assert backend.default_qpos_calls == 0 + assert backend.joint_layout_calls == 1 + assert selected_qpos is not source_qpos + assert selected_qpos is not None + assert not selected_qpos.flags.writeable + np.testing.assert_array_equal(env.scene["robot"].data.default_joint_pos, 0.75) + assert not env.scene["robot"].data.default_joint_pos.flags.writeable + + source_qpos[:] = 9.0 + env.reset() + env.reset(env_ids=np.array([1], dtype=np.int32)) + + assert backend.keyframe_qpos_calls == 1 + assert backend.default_qpos_calls == 0 + assert backend.joint_layout_calls == 1 + assert len(backend.set_state_calls) == 2 + np.testing.assert_array_equal( + backend.set_state_calls[0][1], + [[0.0, 0.0, 0.3, 0.75], [0.0, 0.0, 0.3, 0.75]], + ) + np.testing.assert_array_equal(backend.set_state_calls[1][1], [[0.0, 0.0, 0.3, 0.75]]) + + def test_real_mujoco_backend_is_materialized_before_first_reset() -> None: scene = SceneCfg( model_file=str(ASSETS_ROOT_PATH / "robots" / "go2" / "scene_flat.xml"), @@ -479,6 +620,89 @@ def test_real_mujoco_backend_is_materialized_before_first_reset() -> None: env.close() +@pytest.mark.parametrize( + ("default_keyframe_name", "expected_root_z", "expected_joint_pos"), + [ + (None, 0.445, np.zeros(12, dtype=np.float32)), + ( + "home", + 0.3, + np.array( + [0.0, 0.8, -1.5, 0.0, 0.8, -1.5, 0.0, 1.0, -1.5, 0.0, 1.0, -1.5], + dtype=np.float32, + ), + ), + ], +) +def test_real_mujoco_default_state_matches_qpos0_or_named_home( + default_keyframe_name: str | None, + expected_root_z: float, + expected_joint_pos: np.ndarray, +) -> None: + joint_names = ( + "FL_hip_joint", + "FL_thigh_joint", + "FL_calf_joint", + "FR_hip_joint", + "FR_thigh_joint", + "FR_calf_joint", + "RL_hip_joint", + "RL_thigh_joint", + "RL_calf_joint", + "RR_hip_joint", + "RR_thigh_joint", + "RR_calf_joint", + ) + scene = SceneCfg( + model_file=str(ASSETS_ROOT_PATH / "robots" / "go2" / "scene_flat.xml"), + entities={ + "robot": EntityCfg( + root_body_name="base", + joint_names=joint_names, + ) + }, + default_keyframe_name=default_keyframe_name, + ) + cfg = ManagerBasedRlEnvCfg( + scene=scene, + sim_dt=0.01, + ctrl_dt=0.02, + max_episode_seconds=1.0, + observations={ + "actor": ObservationGroupCfg( + terms={"state": ObservationTermCfg(func=_episode_step_observation)} + ) + }, + actions={}, + events={"reset_default": EventTermCfg(func=mdp.reset_scene_to_default, mode="reset")}, + rewards={"alive": RewardTermCfg(func=mdp.is_alive, weight=1.0)}, + terminations={"time_out": TerminationTermCfg(func=mdp.time_out, time_out=True)}, + policy_observation_group="actor", + ) + backend = create_backend( + "mujoco", + scene, + 2, + cfg.sim_dt, + base_name="base", + add_body_sensors=True, + **env_backend_kwargs(cfg), + ) + env = ManagerBasedRlEnv(cfg, backend, 2) + try: + robot = env.scene["robot"] + np.testing.assert_allclose(robot.data.default_root_state[:, 2], expected_root_z) + expected_batch = np.broadcast_to(expected_joint_pos, (2, expected_joint_pos.size)) + np.testing.assert_allclose(robot.data.default_joint_pos, expected_batch) + + env.init_state() + + np.testing.assert_allclose(robot.data.root_link_pos_w[:, 2], expected_root_z) + np.testing.assert_allclose(robot.data.joint_pos, expected_batch) + finally: + env.close() + + def test_np_env_owns_substeps_autoreset_and_final_observation() -> None: env, backend = _make_env() initial_obs, initial_info = env.reset() From 5ec83af43171a6efd2bbe35472e4806184b1d9eb Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 16:01:44 +0800 Subject: [PATCH 022/169] feat: add Go2 manager config fixture (#1089) Closes #1086. Umbrella #1042. --- .../envs/locomotion/go2/manager_based_cfg.py | 223 +++++++++++++ .../locomotion/go2/test_manager_based_cfg.py | 296 ++++++++++++++++++ 2 files changed, 519 insertions(+) create mode 100644 src/unilab/envs/locomotion/go2/manager_based_cfg.py create mode 100644 tests/envs/locomotion/go2/test_manager_based_cfg.py diff --git a/src/unilab/envs/locomotion/go2/manager_based_cfg.py b/src/unilab/envs/locomotion/go2/manager_based_cfg.py new file mode 100644 index 000000000..08ec3fcf9 --- /dev/null +++ b/src/unilab/envs/locomotion/go2/manager_based_cfg.py @@ -0,0 +1,223 @@ +"""Unregistered Manager-Based configuration for the Go2 flat pilot.""" + +from __future__ import annotations + +import math + +from unilab.assets import ASSETS_ROOT_PATH +from unilab.base.entity import EntityCfg +from unilab.base.scene import SceneCfg +from unilab.envs import ManagerBasedRlEnvCfg, mdp +from unilab.envs.locomotion.common import manager_terms +from unilab.managers import ( + EventTermCfg, + ObservationGroupCfg, + ObservationTermCfg, + RewardTermCfg, + TerminationTermCfg, +) + +_JOINT_NAMES = ( + "FL_hip_joint", + "FL_thigh_joint", + "FL_calf_joint", + "FR_hip_joint", + "FR_thigh_joint", + "FR_calf_joint", + "RL_hip_joint", + "RL_thigh_joint", + "RL_calf_joint", + "RR_hip_joint", + "RR_thigh_joint", + "RR_calf_joint", +) +_ACTUATOR_NAMES = ( + "FR_hip", + "FR_thigh", + "FR_calf", + "FL_hip", + "FL_thigh", + "FL_calf", + "RR_hip", + "RR_thigh", + "RR_calf", + "RL_hip", + "RL_thigh", + "RL_calf", +) +_FOOT_CONTACT_SENSORS = ( + "FL_foot_contact", + "FR_foot_contact", + "RL_foot_contact", + "RR_foot_contact", +) +_FOOT_POSITION_SENSORS = ("FL_pos", "FR_pos", "RL_pos", "RR_pos") +_GAIT_FREQUENCY = 2.0 + + +def make_go2_joystick_flat_manager_cfg() -> ManagerBasedRlEnvCfg: + """Build the NumPy Manager-Based equivalent of the legacy Go2 flat task. + + The factory is intentionally not registered. It proves the task-owned community + config surface without changing the production Go2 registry or Hydra owners. + """ + policy_terms: dict[str, ObservationTermCfg | None] = { + "base_ang_vel": ObservationTermCfg( + func=mdp.builtin_sensor, + params={"sensor_name": "gyro"}, + ), + "projected_gravity": ObservationTermCfg( + func=mdp.projected_gravity_from_sensor, + params={"sensor_name": "upvector"}, + ), + "joint_pos": ObservationTermCfg(func=mdp.joint_pos_rel), + "joint_vel": ObservationTermCfg(func=mdp.joint_vel_rel), + "actions": ObservationTermCfg(func=mdp.last_action), + "command": ObservationTermCfg( + func=mdp.generated_commands, + params={"command_name": "twist"}, + ), + "gait_phase": ObservationTermCfg( + func=manager_terms.quadruped_gait_phase, + params={"frequency": _GAIT_FREQUENCY}, + ), + } + critic_terms: dict[str, ObservationTermCfg | None] = { + **policy_terms, + "base_lin_vel": ObservationTermCfg( + func=mdp.builtin_sensor, + params={"sensor_name": "local_linvel"}, + ), + } + + return ManagerBasedRlEnvCfg( + scene=SceneCfg( + model_file=str(ASSETS_ROOT_PATH / "robots" / "go2" / "scene_flat.xml"), + entities={ + "robot": EntityCfg( + root_body_name="base", + joint_names=_JOINT_NAMES, + actuator_names=_ACTUATOR_NAMES, + ) + }, + default_keyframe_name="home", + ), + sim_dt=0.01, + ctrl_dt=0.02, + max_episode_seconds=20.0, + observations={ + "policy": ObservationGroupCfg(terms=policy_terms), + "critic": ObservationGroupCfg(terms=critic_terms), + }, + actions={ + "joint_pos": mdp.JointPositionActionCfg( + entity_name="robot", + actuator_names=(".*",), + scale=0.25, + use_default_offset=True, + ) + }, + commands={ + "twist": mdp.UniformVelocityCommandCfg( + entity_name="robot", + resampling_time_range=(20.0, 20.0), + heading_command=False, + heading_control_stiffness=0.5, + rel_standing_envs=0.0, + rel_heading_envs=0.0, + rel_world_envs=0.0, + rel_forward_envs=0.0, + init_velocity_prob=0.0, + ranges=mdp.UniformVelocityCommandCfg.Ranges( + lin_vel_x=(-0.6, 1.0), + lin_vel_y=(-0.4, 0.4), + ang_vel_z=(-0.8, 0.8), + ), + ) + }, + events={ + "reset_scene_to_default": EventTermCfg( + func=mdp.reset_scene_to_default, + mode="reset", + ), + "reset_root_state_uniform": EventTermCfg( + func=mdp.reset_root_state_uniform, + mode="reset", + params={ + "pose_range": { + "x": (-0.5, 0.5), + "y": (-0.5, 0.5), + "z": (0.0, 0.0), + "roll": (0.0, 0.0), + "pitch": (0.0, 0.0), + "yaw": (-math.pi, math.pi), + }, + "velocity_range": { + "x": (-0.5, 0.5), + "y": (-0.5, 0.5), + "z": (-0.5, 0.5), + "roll": (-0.5, 0.5), + "pitch": (-0.5, 0.5), + "yaw": (-0.5, 0.5), + }, + }, + ), + }, + rewards={ + "tracking_lin_vel": RewardTermCfg( + func=manager_terms.track_lin_vel_xy_exp, + weight=1.0, + params={"std": math.sqrt(0.25), "command_name": "twist"}, + ), + "tracking_ang_vel": RewardTermCfg( + func=manager_terms.track_ang_vel_z_exp, + weight=0.2, + params={"std": math.sqrt(0.25), "command_name": "twist"}, + ), + "lin_vel_z": RewardTermCfg(func=manager_terms.lin_vel_z_l2, weight=-5.0), + "ang_vel_xy": RewardTermCfg(func=manager_terms.ang_vel_xy_l2, weight=-0.1), + "base_height": RewardTermCfg( + func=manager_terms.base_height_l2, + weight=-100.0, + params={"target_height": 0.3}, + ), + "action_rate": RewardTermCfg(func=mdp.action_rate_l2, weight=-0.005), + "similar_to_default": RewardTermCfg( + func=manager_terms.joint_deviation_l1, + weight=-0.1, + ), + "contact": RewardTermCfg( + func=manager_terms.feet_phase_contact, + weight=0.24, + params={ + "frequency": _GAIT_FREQUENCY, + "sensor_names": _FOOT_CONTACT_SENSORS, + "contact_threshold": 0.1, + "stance_threshold": 0.6, + }, + ), + "swing_feet_z": RewardTermCfg( + func=manager_terms.feet_phase_swing_height, + weight=4.0, + params={ + "frequency": _GAIT_FREQUENCY, + "sensor_names": _FOOT_POSITION_SENSORS, + "target_height": 0.1, + "kernel": 0.01, + "swing_start": 0.6, + }, + ), + }, + terminations={ + "time_out": TerminationTermCfg(func=mdp.time_out, time_out=True), + "bad_orientation": TerminationTermCfg( + func=mdp.bad_orientation, + params={"limit_angle": math.acos(0.5)}, + ), + }, + policy_observation_group="policy", + critic_observation_group="critic", + ) + + +__all__ = ["make_go2_joystick_flat_manager_cfg"] diff --git a/tests/envs/locomotion/go2/test_manager_based_cfg.py b/tests/envs/locomotion/go2/test_manager_based_cfg.py new file mode 100644 index 000000000..307912177 --- /dev/null +++ b/tests/envs/locomotion/go2/test_manager_based_cfg.py @@ -0,0 +1,296 @@ +"""Go2 flat Manager-Based config and real-runtime fixture tests.""" + +from __future__ import annotations + +import math +from typing import TypeVar, cast + +import numpy as np +import pytest + +from unilab.base.backend import create_backend, env_backend_kwargs +from unilab.base.np_env import NpEnvState +from unilab.envs import ManagerBasedRlEnv, mdp +from unilab.envs.locomotion.common import manager_terms +from unilab.envs.locomotion.go2.manager_based_cfg import ( + make_go2_joystick_flat_manager_cfg, +) +from unilab.envs.mdp import JointPositionAction, JointPositionActionCfg + +_JOINT_NAMES = ( + "FL_hip_joint", + "FL_thigh_joint", + "FL_calf_joint", + "FR_hip_joint", + "FR_thigh_joint", + "FR_calf_joint", + "RL_hip_joint", + "RL_thigh_joint", + "RL_calf_joint", + "RR_hip_joint", + "RR_thigh_joint", + "RR_calf_joint", +) +_ACTUATOR_NAMES = ( + "FR_hip", + "FR_thigh", + "FR_calf", + "FL_hip", + "FL_thigh", + "FL_calf", + "RR_hip", + "RR_thigh", + "RR_calf", + "RL_hip", + "RL_thigh", + "RL_calf", +) +_HOME_JOINT_POS = np.array( + [0.0, 0.8, -1.5, 0.0, 0.8, -1.5, 0.0, 1.0, -1.5, 0.0, 1.0, -1.5], + dtype=np.float32, +) +_TermCfg = TypeVar("_TermCfg") + + +def _active(terms: dict[str, _TermCfg | None]) -> dict[str, _TermCfg]: + assert all(term is not None for term in terms.values()) + return cast(dict[str, _TermCfg], terms) + + +def test_go2_manager_factory_preserves_legacy_config_surface() -> None: + cfg = make_go2_joystick_flat_manager_cfg() + + cfg.validate() + assert cfg.sim_dt == pytest.approx(0.01) + assert cfg.ctrl_dt == pytest.approx(0.02) + assert cfg.max_episode_seconds == pytest.approx(20.0) + assert cfg.policy_observation_group == "policy" + assert cfg.critic_observation_group == "critic" + assert cfg.scene is not None + assert cfg.scene.default_keyframe_name == "home" + + robot = cfg.scene.entities["robot"] + assert robot.root_body_name == "base" + assert robot.joint_names == _JOINT_NAMES + assert robot.actuator_names == _ACTUATOR_NAMES + + expected_policy = [ + "base_ang_vel", + "projected_gravity", + "joint_pos", + "joint_vel", + "actions", + "command", + "gait_phase", + ] + observations = _active(cfg.observations) + assert list(observations) == ["policy", "critic"] + policy_terms = _active(observations["policy"].terms) + critic_terms = _active(observations["critic"].terms) + expected_policy_terms = [ + ("base_ang_vel", mdp.builtin_sensor, {"sensor_name": "gyro"}), + ( + "projected_gravity", + mdp.projected_gravity_from_sensor, + {"sensor_name": "upvector"}, + ), + ("joint_pos", mdp.joint_pos_rel, {}), + ("joint_vel", mdp.joint_vel_rel, {}), + ("actions", mdp.last_action, {}), + ("command", mdp.generated_commands, {"command_name": "twist"}), + ("gait_phase", manager_terms.quadruped_gait_phase, {"frequency": 2.0}), + ] + assert list(policy_terms) == expected_policy + assert [(name, term.func, term.params) for name, term in policy_terms.items()] == ( + expected_policy_terms + ) + assert [(name, term.func, term.params) for name, term in critic_terms.items()] == [ + *expected_policy_terms, + ("base_lin_vel", mdp.builtin_sensor, {"sensor_name": "local_linvel"}), + ] + + actions = _active(cfg.actions) + assert list(actions) == ["joint_pos"] + action = actions["joint_pos"] + assert isinstance(action, JointPositionActionCfg) + assert action.entity_name == "robot" + assert action.actuator_names == (".*",) + assert action.scale == pytest.approx(0.25) + assert action.use_default_offset is True + + commands = _active(cfg.commands) + assert list(commands) == ["twist"] + command = commands["twist"] + assert isinstance(command, mdp.UniformVelocityCommandCfg) + assert command.resampling_time_range == (20.0, 20.0) + assert command.heading_command is False + assert command.heading_control_stiffness == pytest.approx(0.5) + assert command.rel_standing_envs == 0.0 + assert command.rel_heading_envs == 0.0 + assert command.rel_world_envs == 0.0 + assert command.rel_forward_envs == 0.0 + assert command.init_velocity_prob == 0.0 + assert command.ranges.lin_vel_x == (-0.6, 1.0) + assert command.ranges.lin_vel_y == (-0.4, 0.4) + assert command.ranges.ang_vel_z == (-0.8, 0.8) + assert command.ranges.heading is None + + events = _active(cfg.events) + assert list(events) == ["reset_scene_to_default", "reset_root_state_uniform"] + assert events["reset_scene_to_default"].func is mdp.reset_scene_to_default + assert events["reset_scene_to_default"].mode == "reset" + assert events["reset_scene_to_default"].params == {} + root_reset = events["reset_root_state_uniform"] + assert root_reset.func is mdp.reset_root_state_uniform + assert root_reset.mode == "reset" + assert root_reset.params["pose_range"] == { + "x": (-0.5, 0.5), + "y": (-0.5, 0.5), + "z": (0.0, 0.0), + "roll": (0.0, 0.0), + "pitch": (0.0, 0.0), + "yaw": (-math.pi, math.pi), + } + assert root_reset.params["velocity_range"] == { + key: (-0.5, 0.5) for key in ("x", "y", "z", "roll", "pitch", "yaw") + } + + expected_rewards = { + "tracking_lin_vel": 1.0, + "tracking_ang_vel": 0.2, + "lin_vel_z": -5.0, + "ang_vel_xy": -0.1, + "base_height": -100.0, + "action_rate": -0.005, + "similar_to_default": -0.1, + "contact": 0.24, + "swing_feet_z": 4.0, + } + rewards = _active(cfg.rewards) + assert list(rewards) == list(expected_rewards) + assert {name: term.weight for name, term in rewards.items()} == expected_rewards + assert [(name, term.func, term.params) for name, term in rewards.items()] == [ + ( + "tracking_lin_vel", + manager_terms.track_lin_vel_xy_exp, + {"std": 0.5, "command_name": "twist"}, + ), + ( + "tracking_ang_vel", + manager_terms.track_ang_vel_z_exp, + {"std": 0.5, "command_name": "twist"}, + ), + ("lin_vel_z", manager_terms.lin_vel_z_l2, {}), + ("ang_vel_xy", manager_terms.ang_vel_xy_l2, {}), + ("base_height", manager_terms.base_height_l2, {"target_height": 0.3}), + ("action_rate", mdp.action_rate_l2, {}), + ("similar_to_default", manager_terms.joint_deviation_l1, {}), + ( + "contact", + manager_terms.feet_phase_contact, + { + "frequency": 2.0, + "sensor_names": ( + "FL_foot_contact", + "FR_foot_contact", + "RL_foot_contact", + "RR_foot_contact", + ), + "contact_threshold": 0.1, + "stance_threshold": 0.6, + }, + ), + ( + "swing_feet_z", + manager_terms.feet_phase_swing_height, + { + "frequency": 2.0, + "sensor_names": ("FL_pos", "FR_pos", "RL_pos", "RR_pos"), + "target_height": 0.1, + "kernel": 0.01, + "swing_start": 0.6, + }, + ), + ] + + terminations = _active(cfg.terminations) + assert list(terminations) == ["time_out", "bad_orientation"] + assert terminations["time_out"].func is mdp.time_out + assert terminations["time_out"].time_out is True + assert terminations["time_out"].params == {} + assert terminations["bad_orientation"].func is mdp.bad_orientation + assert terminations["bad_orientation"].time_out is False + assert terminations["bad_orientation"].params["limit_angle"] == pytest.approx(math.pi / 3.0) + + +def test_go2_manager_factory_executes_on_real_mujoco() -> None: + cfg = make_go2_joystick_flat_manager_cfg() + assert cfg.scene is not None + backend = create_backend( + "mujoco", + cfg.scene, + 2, + cfg.sim_dt, + base_name="base", + add_body_sensors=True, + **env_backend_kwargs(cfg), + ) + env = ManagerBasedRlEnv(cfg, backend, 2) + try: + assert env.obs_groups_spec == {"obs": 49, "critic": 52} + assert env.action_space.shape == (12,) + assert env.observation_manager.active_terms == { + "policy": [ + "base_ang_vel", + "projected_gravity", + "joint_pos", + "joint_vel", + "actions", + "command", + "gait_phase", + ], + "critic": [ + "base_ang_vel", + "projected_gravity", + "joint_pos", + "joint_vel", + "actions", + "command", + "gait_phase", + "base_lin_vel", + ], + } + assert env.reward_manager.active_terms == list(cfg.rewards) + assert env.termination_manager.active_terms == list(cfg.terminations) + assert env.event_manager.active_terms["reset"] == list(cfg.events) + + action = env.action_manager.get_term("joint_pos") + assert isinstance(action, JointPositionAction) + assert action.target_names == list(_JOINT_NAMES) + np.testing.assert_allclose(action.offset, np.broadcast_to(_HOME_JOINT_POS, (2, 12))) + + obs, info = env.reset(seed=7) + assert set(obs) == {"obs", "critic"} + assert obs["obs"].shape == (2, 49) + assert obs["critic"].shape == (2, 52) + assert isinstance(info, dict) + np.testing.assert_allclose( + env.scene["robot"].data.default_joint_pos, + np.broadcast_to(_HOME_JOINT_POS, (2, 12)), + ) + np.testing.assert_allclose( + env.scene["robot"].data.joint_pos, + np.broadcast_to(_HOME_JOINT_POS, (2, 12)), + ) + + state = env.step(np.zeros((2, 12), dtype=np.float32)) + assert isinstance(state, NpEnvState) + assert state.obs["obs"].shape == (2, 49) + assert state.obs["critic"].shape == (2, 52) + for value in (*state.obs.values(), state.reward): + assert isinstance(value, np.ndarray) + assert np.isfinite(value).all() + assert state.terminated.dtype == np.bool_ + assert state.truncated.dtype == np.bool_ + finally: + env.close() From 66cc550ad6e7a63dc243558fb943f0c2f1c2003e Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 16:15:11 +0800 Subject: [PATCH 023/169] feat(config): support registry config factories (#1091) [skip ci] --- scripts/visualize_task_env.py | 9 +-- src/unilab/base/registry.py | 66 +++++++++++++++------ tests/base/test_registry.py | 58 ++++++++++++++++++ tests/config/test_manager_reward_routing.py | 17 ++++-- 4 files changed, 123 insertions(+), 27 deletions(-) diff --git a/scripts/visualize_task_env.py b/scripts/visualize_task_env.py index 60bfd2124..12a98a533 100644 --- a/scripts/visualize_task_env.py +++ b/scripts/visualize_task_env.py @@ -284,13 +284,14 @@ def _run_mujoco(env, num_envs: int) -> None: def _build_env_cfg_override(task_name: str) -> dict[str, Any]: """Build the env_cfg_override dict from CLI args alone — no Hydra.""" - if task_name not in registry._envs: + if not registry.contains(task_name): raise SystemExit( - f"Task '{task_name}' is not registered. Available: {sorted(registry._envs.keys())}" + f"Task '{task_name}' is not registered. " + f"Available: {sorted(registry.list_registered_envs())}" ) - env_cfg_cls = registry._envs[task_name].env_cfg_cls + env_cfg = registry.materialize_env_config(task_name) override: dict[str, Any] = {} - reward_stub = _build_reward_stub(env_cfg_cls) + reward_stub = _build_reward_stub(type(env_cfg)) if reward_stub is not None: override["reward_config"] = reward_stub return override diff --git a/src/unilab/base/registry.py b/src/unilab/base/registry.py index 5badd68b5..dd785692e 100644 --- a/src/unilab/base/registry.py +++ b/src/unilab/base/registry.py @@ -25,7 +25,8 @@ MANAGER_TERM_MAPPING_POLICY, ) -TEnvCfg = TypeVar("TEnvCfg", bound=EnvCfg) +EnvCfgFactory = Callable[[], EnvCfg] +TEnvCfgFactory = TypeVar("TEnvCfgFactory", bound=EnvCfgFactory) RewardOverrideField = Literal["reward_config", "rewards"] _SUPPORTED_SIM_BACKENDS = ("mujoco", "mjwarp", "motrix", "drake") _DEFAULT_SIM_BACKEND_ORDER: tuple[str, ...] = ("mujoco", "motrix") @@ -45,7 +46,7 @@ @dataclass class EnvMeta: - env_cfg_cls: Type[EnvCfg] + env_cfg_factory: EnvCfgFactory env_cls_dict: Dict[str, Type[ABEnv]] = field(default_factory=dict) def available_sim_backend(self) -> Optional[str]: @@ -68,31 +69,59 @@ def contains(name: str) -> bool: return name in _envs -def register_env_config(name: str, env_cfg_cls: Type[EnvCfg]): - """Register an environment configuration class with a name.""" +def _config_factory_name(factory: EnvCfgFactory) -> str: + return str(getattr(factory, "__qualname__", type(factory).__qualname__)) + + +def register_env_config(name: str, env_cfg_factory: EnvCfgFactory) -> None: + """Register a zero-argument environment configuration factory.""" if name in _envs.keys(): raise ValueError(f"Environment '{name}' is already registered.") - _envs[name] = EnvMeta(env_cfg_cls=env_cfg_cls) + if not callable(env_cfg_factory): + raise TypeError( + f"Environment '{name}' config factory must be callable, got " + f"{type(env_cfg_factory).__name__}" + ) + _envs[name] = EnvMeta(env_cfg_factory=env_cfg_factory) -def envcfg(name: str) -> Callable[[Type[TEnvCfg]], Type[TEnvCfg]]: +def envcfg(name: str) -> Callable[[TEnvCfgFactory], TEnvCfgFactory]: """ - Decorator to register an environment configuration class with a name. + Decorator to register an environment configuration class or factory. Usage: - @register_env_config_decorator("my-env") + @envcfg("my-env") @dataclass class MyEnvCfg(EnvCfg): ... + + @envcfg("my-manager-env") + def make_my_env_cfg() -> EnvCfg: + ... """ - def decorator(cls: Type[TEnvCfg]) -> Type[TEnvCfg]: - register_env_config(name, cls) - return cls + def decorator(factory: TEnvCfgFactory) -> TEnvCfgFactory: + register_env_config(name, factory) + return factory return decorator +def materialize_env_config(name: str) -> EnvCfg: + """Construct one config instance from the registered cold-path factory.""" + if name not in _envs: + raise ValueError(f"Environment '{name}' is not registered.") + + factory = _envs[name].env_cfg_factory + env_cfg = factory() + if not isinstance(env_cfg, EnvCfg): + raise TypeError( + f"Environment '{name}' config factory '{_config_factory_name(factory)}' returned " + f"{type(env_cfg).__name__}, expected an EnvCfg instance" + ) + return env_cfg + + def register_env(name: str, env_cls: Type[ABEnv], sim_backend: str): """Register an environment class with a name and simulation backend.""" if sim_backend not in _SUPPORTED_SIM_BACKENDS: @@ -154,20 +183,19 @@ def resolve_reward_override_field(env_name: str) -> RewardOverrideField: if env_name not in _envs: raise ValueError(f"Environment '{env_name}' is not registered.") - config_cls = _envs[env_name].env_cfg_cls - config_fields = { - config_field.name: config_field for config_field in dataclasses.fields(config_cls) - } + config = materialize_env_config(env_name) + config_fields = {config_field.name: config_field for config_field in dataclasses.fields(config)} rewards_field = config_fields.get("rewards") has_manager_rewards = ( rewards_field is not None and rewards_field.metadata.get(CONFIG_MAPPING_POLICY_KEY) == MANAGER_TERM_MAPPING_POLICY ) has_legacy_rewards = "reward_config" in config_fields + config_owner = type(config).__name__ if has_manager_rewards and has_legacy_rewards: raise ValueError( - f"Environment '{env_name}' config owner '{config_cls.__name__}' declares both " + f"Environment '{env_name}' config owner '{config_owner}' declares both " "Manager-Based 'rewards' and legacy 'reward_config' targets" ) if has_manager_rewards: @@ -175,7 +203,7 @@ def resolve_reward_override_field(env_name: str) -> RewardOverrideField: if has_legacy_rewards: return "reward_config" raise ValueError( - f"Environment '{env_name}' config owner '{config_cls.__name__}' declares no " + f"Environment '{env_name}' config owner '{config_owner}' declares no " "supported Hydra root reward target; expected legacy 'reward_config' or an " "explicitly marked Manager-Based 'rewards' field" ) @@ -366,7 +394,7 @@ def make( meta: EnvMeta = _envs[name] # Create environment config - env_cfg = meta.env_cfg_cls() + env_cfg = materialize_env_config(name) if env_cfg_override is not None: apply_cfg_overrides(env_cfg, env_cfg_override) @@ -395,7 +423,7 @@ def list_registered_envs() -> Dict[str, Dict[str, Any]]: result = {} for name, meta in _envs.items(): result[name] = { - "config_class": meta.env_cfg_cls.__name__, + "config_factory": _config_factory_name(meta.env_cfg_factory), "available_backends": list(meta.env_cls_dict.keys()), } return result diff --git a/tests/base/test_registry.py b/tests/base/test_registry.py index 05570f9c7..4a2edc72c 100644 --- a/tests/base/test_registry.py +++ b/tests/base/test_registry.py @@ -104,6 +104,38 @@ class _DecCfg(EnvCfg): assert registry_mod.contains(_name) +def test_envcfg_decorator_registers_plain_factory_and_returns_it_unchanged(): + _name = "_TestDecoratorFactoryEnv" + + def make_cfg() -> EnvCfg: + return _TestCfgA() + + registered = registry_mod.envcfg(_name)(make_cfg) + + assert registered is make_cfg + assert registry_mod.materialize_env_config(_name) == _TestCfgA() + + +def test_materialize_env_config_rejects_invalid_factory_output(): + _name = "_TestInvalidFactoryOutput" + + def make_invalid_cfg(): + return {"sim_dt": 0.01} + + registry_mod.register_env_config(_name, make_invalid_cfg) + + with pytest.raises( + TypeError, + match=r"_TestInvalidFactoryOutput.*make_invalid_cfg.*dict.*EnvCfg", + ): + registry_mod.materialize_env_config(_name) + + +def test_register_env_config_rejects_non_callable(): + with pytest.raises(TypeError, match="config factory must be callable"): + registry_mod.register_env_config("_TestNonCallableFactory", None) # type: ignore[arg-type] + + def test_env_decorator_registers(): _name = "_TestDecoratorEnv2" if not registry_mod.contains(_name): @@ -159,6 +191,7 @@ def test_make_unregistered_raises(): def test_list_registered_envs_includes_registered(): listed = registry_mod.list_registered_envs() assert _TEST_ENV_A in listed + assert listed[_TEST_ENV_A]["config_factory"] == "_TestCfgA" assert "mujoco" in listed[_TEST_ENV_A]["available_backends"] @@ -240,6 +273,31 @@ class _CfgWithField(EnvCfg): assert env.cfg.ctrl_dt == 0.05 +def test_make_materializes_fresh_function_owned_config_for_each_call(): + @dataclass + class _FactoryCfg(EnvCfg): + ctrl_dt: float = 0.02 + + _name = "_TestFunctionFactoryOverrideEnv" + + def make_cfg() -> EnvCfg: + return _FactoryCfg() + + registry_mod.register_env_config(_name, make_cfg) + registry_mod.register_env(_name, _TestEnvA, "mujoco") + + overridden = registry_mod.make( + _name, + sim_backend="mujoco", + env_cfg_override={"ctrl_dt": 0.05}, + ) + default = registry_mod.make(_name, sim_backend="mujoco") + + assert overridden.cfg.ctrl_dt == 0.05 + assert default.cfg.ctrl_dt == 0.02 + assert overridden.cfg is not default.cfg + + def test_make_with_invalid_cfg_override_raises(): """make() with a config key that doesn't exist raises ValueError.""" with pytest.raises(ValueError, match="has no attribute"): diff --git a/tests/config/test_manager_reward_routing.py b/tests/config/test_manager_reward_routing.py index 25d8ed580..3b659f308 100644 --- a/tests/config/test_manager_reward_routing.py +++ b/tests/config/test_manager_reward_routing.py @@ -19,6 +19,7 @@ from unilab.training.backend_adapter import BackendAdapter _MANAGER_ENV = "_TestManagerRewardRoute" +_MANAGER_FACTORY_ENV = "_TestManagerFactoryRewardRoute" _LEGACY_ENV = "_TestLegacyRewardRoute" _MISSING_ENV = "_TestMissingRewardRoute" _AMBIGUOUS_ENV = "_TestAmbiguousRewardRoute" @@ -43,14 +44,19 @@ class _AmbiguousCfg(EnvCfg): ) -for _name, _cfg_cls in ( +def _make_manager_cfg() -> ManagerBasedRlEnvCfg: + return ManagerBasedRlEnvCfg() + + +for _name, _cfg_factory in ( (_MANAGER_ENV, ManagerBasedRlEnvCfg), + (_MANAGER_FACTORY_ENV, _make_manager_cfg), (_LEGACY_ENV, _LegacyCfg), (_MISSING_ENV, _MissingCfg), (_AMBIGUOUS_ENV, _AmbiguousCfg), ): if not registry.contains(_name): - registry.register_env_config(_name, _cfg_cls) + registry.register_env_config(_name, _cfg_factory) def _reward(_env, *, std: float) -> np.ndarray: @@ -77,9 +83,12 @@ def _cfg(task_name: str, *, env: dict[str, object] | None = None): ) -def test_backend_adapter_routes_manager_reward_and_preserves_factory_terms() -> None: +@pytest.mark.parametrize("env_name", [_MANAGER_ENV, _MANAGER_FACTORY_ENV]) +def test_backend_adapter_routes_manager_reward_and_preserves_factory_terms( + env_name: str, +) -> None: override = BackendAdapter( - _cfg(_MANAGER_ENV), + _cfg(env_name), root_dir=Path("."), ).build_task_env_cfg_override() From f3d08fa1ef53f391cd4fed91e083102189d982e3 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 16:41:11 +0800 Subject: [PATCH 024/169] feat(env): add manager pd gain reset event (#1093) 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) From 80aae37b5e356fcf4820c93e1c61a8f3abdb6ca5 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 16:50:39 +0800 Subject: [PATCH 025/169] feat(config): enable Go2 manager pd gain event (#1095) Refs #1094 [skip ci] --- .../envs/locomotion/go2/manager_based_cfg.py | 9 +++++++++ .../locomotion/go2/test_manager_based_cfg.py | 20 +++++++++---------- 2 files changed, 18 insertions(+), 11 deletions(-) diff --git a/src/unilab/envs/locomotion/go2/manager_based_cfg.py b/src/unilab/envs/locomotion/go2/manager_based_cfg.py index 08ec3fcf9..3da473da3 100644 --- a/src/unilab/envs/locomotion/go2/manager_based_cfg.py +++ b/src/unilab/envs/locomotion/go2/manager_based_cfg.py @@ -162,6 +162,15 @@ def make_go2_joystick_flat_manager_cfg() -> ManagerBasedRlEnvCfg: }, }, ), + "pd_gains": EventTermCfg( + func=mdp.pd_gains, + mode="reset", + params={ + "kp_range": (31.5, 38.5), + "kd_range": (0.45, 0.55), + "operation": "abs", + }, + ), }, rewards={ "tracking_lin_vel": RewardTermCfg( diff --git a/tests/envs/locomotion/go2/test_manager_based_cfg.py b/tests/envs/locomotion/go2/test_manager_based_cfg.py index b70abbbea..77af704ed 100644 --- a/tests/envs/locomotion/go2/test_manager_based_cfg.py +++ b/tests/envs/locomotion/go2/test_manager_based_cfg.py @@ -16,7 +16,6 @@ make_go2_joystick_flat_manager_cfg, ) from unilab.envs.mdp import JointPositionAction, JointPositionActionCfg -from unilab.managers import EventTermCfg _JOINT_NAMES = ( "FL_hip_joint", @@ -137,7 +136,7 @@ def test_go2_manager_factory_preserves_legacy_config_surface() -> None: assert command.ranges.heading is None events = _active(cfg.events) - assert list(events) == ["reset_scene_to_default", "reset_root_state_uniform"] + assert list(events) == ["reset_scene_to_default", "reset_root_state_uniform", "pd_gains"] assert events["reset_scene_to_default"].func is mdp.reset_scene_to_default assert events["reset_scene_to_default"].mode == "reset" assert events["reset_scene_to_default"].params == {} @@ -155,6 +154,14 @@ def test_go2_manager_factory_preserves_legacy_config_surface() -> None: assert root_reset.params["velocity_range"] == { key: (-0.5, 0.5) for key in ("x", "y", "z", "roll", "pitch", "yaw") } + pd_gains = events["pd_gains"] + assert pd_gains.func is mdp.pd_gains + assert pd_gains.mode == "reset" + assert pd_gains.params == { + "kp_range": (31.5, 38.5), + "kd_range": (0.45, 0.55), + "operation": "abs", + } expected_rewards = { "tracking_lin_vel": 1.0, @@ -317,15 +324,6 @@ def _read_runtime_actuator_gains(backend_type: str, backend) -> tuple[np.ndarray @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, From a1282f5c148aaec516043768fcc06881ab4a5d97 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 17:23:19 +0800 Subject: [PATCH 026/169] feat(registry): support callable env factories [skip ci] (#1097) --- src/unilab/base/registry.py | 70 +++++++++++++---- tests/base/test_registry.py | 77 +++++++++++++++++-- tests/config/test_locomotion_params.py | 2 +- .../locomotion/g1/test_g1_owner_contract.py | 3 +- 4 files changed, 128 insertions(+), 24 deletions(-) diff --git a/src/unilab/base/registry.py b/src/unilab/base/registry.py index dd785692e..48ed19977 100644 --- a/src/unilab/base/registry.py +++ b/src/unilab/base/registry.py @@ -10,6 +10,7 @@ Dict, Literal, Optional, + Protocol, Type, TypeVar, cast, @@ -27,6 +28,21 @@ EnvCfgFactory = Callable[[], EnvCfg] TEnvCfgFactory = TypeVar("TEnvCfgFactory", bound=EnvCfgFactory) + + +class EnvFactory(Protocol): + """Construct an environment for one materialized config and backend.""" + + def __call__( + self, + cfg: Any, + *, + num_envs: int = 1, + backend_type: str = "mujoco", + ) -> ABEnv: ... + + +TEnvFactory = TypeVar("TEnvFactory", bound=EnvFactory) RewardOverrideField = Literal["reward_config", "rewards"] _SUPPORTED_SIM_BACKENDS = ("mujoco", "mjwarp", "motrix", "drake") _DEFAULT_SIM_BACKEND_ORDER: tuple[str, ...] = ("mujoco", "motrix") @@ -47,18 +63,18 @@ @dataclass class EnvMeta: env_cfg_factory: EnvCfgFactory - env_cls_dict: Dict[str, Type[ABEnv]] = field(default_factory=dict) + env_factory_dict: Dict[str, EnvFactory] = field(default_factory=dict) def available_sim_backend(self) -> Optional[str]: """Return the explicit default simulation backend for this environment.""" for backend in _DEFAULT_SIM_BACKEND_ORDER: - if backend in self.env_cls_dict: + if backend in self.env_factory_dict: return backend - return next(iter(self.env_cls_dict), None) + return next(iter(self.env_factory_dict), None) def support_sim_backend(self, sim_backend: str) -> bool: """Check if the environment supports a specific simulation backend.""" - return sim_backend in self.env_cls_dict + return sim_backend in self.env_factory_dict _envs: Dict[str, EnvMeta] = {} @@ -73,6 +89,10 @@ def _config_factory_name(factory: EnvCfgFactory) -> str: return str(getattr(factory, "__qualname__", type(factory).__qualname__)) +def _env_factory_name(factory: EnvFactory) -> str: + return str(getattr(factory, "__qualname__", type(factory).__qualname__)) + + def register_env_config(name: str, env_cfg_factory: EnvCfgFactory) -> None: """Register a zero-argument environment configuration factory.""" if name in _envs.keys(): @@ -122,8 +142,8 @@ def materialize_env_config(name: str) -> EnvCfg: return env_cfg -def register_env(name: str, env_cls: Type[ABEnv], sim_backend: str): - """Register an environment class with a name and simulation backend.""" +def register_env(name: str, env_factory: TEnvFactory, sim_backend: str) -> TEnvFactory: + """Register and return an environment class or function factory.""" if sim_backend not in _SUPPORTED_SIM_BACKENDS: raise ValueError( f"Unsupported simulation backend: {sim_backend}. " @@ -135,27 +155,37 @@ def register_env(name: str, env_cls: Type[ABEnv], sim_backend: str): f"Environment '{name}' is not registered. Please register the config first." ) - if sim_backend in _envs[name].env_cls_dict: + if not callable(env_factory): + raise TypeError( + f"Environment '{name}' backend '{sim_backend}' factory must be callable, got " + f"{type(env_factory).__name__}" + ) + + if sim_backend in _envs[name].env_factory_dict: raise ValueError( f"Environment '{name}' with sim backend '{sim_backend}' is already registered." ) - _envs[name].env_cls_dict[sim_backend] = env_cls + _envs[name].env_factory_dict[sim_backend] = env_factory + return env_factory -def env(name: str, sim_backend: str) -> Callable[[Type[ABEnv]], Type[ABEnv]]: +def env(name: str, sim_backend: str) -> Callable[[TEnvFactory], TEnvFactory]: """ - Decorator to register an environment class with a name and simulation backend. + Decorator to register an environment class or function factory. Usage: - @register_env_decorator("my-env", "np") + @env("my-env", "mujoco") class MyEnv(ABEnv): ... + + @env("my-manager-env", "mujoco") + def make_my_env(cfg, num_envs=1, backend_type="mujoco"): + ... """ - def decorator(cls: Type[ABEnv]) -> Type[ABEnv]: - register_env(name, cls, sim_backend) - return cls + def decorator(factory: TEnvFactory) -> TEnvFactory: + return register_env(name, factory, sim_backend) return decorator @@ -413,8 +443,14 @@ def make( ) # Create environment instance - env_cls_any: Any = meta.env_cls_dict[sim_backend] - env: ABEnv = env_cls_any(env_cfg, num_envs=num_envs, backend_type=sim_backend) + factory = meta.env_factory_dict[sim_backend] + env = factory(env_cfg, num_envs=num_envs, backend_type=sim_backend) + if not isinstance(env, ABEnv): + raise TypeError( + f"Environment '{name}' backend '{sim_backend}' factory " + f"'{_env_factory_name(factory)}' returned {type(env).__name__}, " + "expected an ABEnv instance" + ) return env @@ -424,7 +460,7 @@ def list_registered_envs() -> Dict[str, Dict[str, Any]]: for name, meta in _envs.items(): result[name] = { "config_factory": _config_factory_name(meta.env_cfg_factory), - "available_backends": list(meta.env_cls_dict.keys()), + "available_backends": list(meta.env_factory_dict.keys()), } return result diff --git a/tests/base/test_registry.py b/tests/base/test_registry.py index 4a2edc72c..f64961153 100644 --- a/tests/base/test_registry.py +++ b/tests/base/test_registry.py @@ -150,6 +150,19 @@ class _DecEnv(_TestEnvA): assert "mujoco" in listed[_name]["available_backends"] +def test_env_decorator_registers_plain_factory_and_returns_it_unchanged(): + _name = "_TestDecoratorEnvFactory" + registry_mod.register_env_config(_name, _TestCfgA) + + def make_env(cfg: EnvCfg, num_envs: int = 1, backend_type: str = "mujoco") -> ABEnv: + return _TestEnvA(cfg, num_envs=num_envs, backend_type=backend_type) + + registered = registry_mod.env(_name, "mujoco")(make_env) + + assert registered is make_env + assert registry_mod._envs[_name].env_factory_dict["mujoco"] is make_env + + def test_contains_before_and_after_registration(): _name = "_TestContainsDynamic" assert not registry_mod.contains(_name) @@ -222,9 +235,17 @@ def test_register_env_duplicate_backend_raises(): registry_mod.register_env(_TEST_ENV_A, _TestEnvA, "mujoco") -def test_find_available_sim_backend_no_env_cls_raises(): - """find_available_sim_backend() raises when config exists but no env_cls registered.""" - # _TEST_ENV_B has config but no env class (registered above without env) +def test_register_env_rejects_non_callable_factory(): + _name = "_TestNonCallableEnvFactory" + registry_mod.register_env_config(_name, _TestCfgA) + + with pytest.raises(TypeError, match=r"_TestNonCallableEnvFactory.*mujoco.*callable.*object"): + registry_mod.register_env(_name, object(), "mujoco") # type: ignore[arg-type] + + +def test_find_available_sim_backend_no_env_factory_raises(): + """find_available_sim_backend() raises when config exists but no env factory is registered.""" + # _TEST_ENV_B has config but no env factory (registered above without env) with pytest.raises(ValueError, match="does not support any simulation backend"): registry_mod.find_available_sim_backend(_TEST_ENV_B) @@ -245,14 +266,60 @@ def test_make_auto_selects_default_backend_independent_of_registration_order(): assert isinstance(env, _TestEnvA) +def test_make_calls_plain_factory_with_cfg_num_envs_and_backend(): + @dataclass + class _CallableFactoryCfg(EnvCfg): + ctrl_dt: float = 0.02 + + _name = "_TestCallableEnvFactory" + received: dict[str, object] = {} + + def make_env(cfg: EnvCfg, num_envs: int = 1, backend_type: str = "mujoco") -> ABEnv: + received.update(cfg=cfg, num_envs=num_envs, backend_type=backend_type) + return _TestEnvA(cfg, num_envs=num_envs, backend_type=backend_type) + + registry_mod.register_env_config(_name, _CallableFactoryCfg) + registered = registry_mod.register_env(_name, make_env, "motrix") + + made = registry_mod.make( + _name, + sim_backend=None, + env_cfg_override={"ctrl_dt": 0.05}, + num_envs=7, + ) + + assert registered is make_env + assert isinstance(made, _TestEnvA) + assert received["cfg"] is made.cfg + assert received["num_envs"] == 7 + assert received["backend_type"] == "motrix" + assert made.cfg.ctrl_dt == pytest.approx(0.05) + + +def test_make_rejects_invalid_factory_output_at_registry_boundary(): + _name = "_TestInvalidEnvFactoryOutput" + + def make_invalid_env(cfg: EnvCfg, num_envs: int = 1, backend_type: str = "mujoco") -> object: + return object() + + registry_mod.register_env_config(_name, _TestCfgA) + registry_mod.register_env(_name, make_invalid_env, "mujoco") # type: ignore[arg-type] + + with pytest.raises( + TypeError, + match=r"_TestInvalidEnvFactoryOutput.*mujoco.*make_invalid_env.*object.*ABEnv", + ): + registry_mod.make(_name, sim_backend="mujoco") + + def test_make_unsupported_backend_raises(): """make() with an unsupported backend name raises ValueError.""" with pytest.raises(ValueError, match="does not support simulation backend"): registry_mod.make(_TEST_ENV_A, sim_backend="motrix") -def test_make_no_env_cls_raises(): - """make() when no env class registered (only config) raises ValueError.""" +def test_make_no_env_factory_raises(): + """make() when no env factory is registered (only config) raises ValueError.""" with pytest.raises(ValueError, match="does not support any simulation backend"): registry_mod.make(_TEST_ENV_B, sim_backend=None) diff --git a/tests/config/test_locomotion_params.py b/tests/config/test_locomotion_params.py index f6990b9af..6cb5e7c7c 100644 --- a/tests/config/test_locomotion_params.py +++ b/tests/config/test_locomotion_params.py @@ -350,7 +350,7 @@ def test_go2_joystick_rough_motrix_registers_rough_env(): from unilab.base import registry from unilab.envs.locomotion.go2.rough import Go2JoystickRoughEnv - assert registry._envs["Go2JoystickRough"].env_cls_dict["motrix"] is Go2JoystickRoughEnv + assert registry._envs["Go2JoystickRough"].env_factory_dict["motrix"] is Go2JoystickRoughEnv def test_offpolicy_g1_rough_terrain_task_overrides(): diff --git a/tests/envs/locomotion/g1/test_g1_owner_contract.py b/tests/envs/locomotion/g1/test_g1_owner_contract.py index 9fd3a6e0a..7da05e97c 100644 --- a/tests/envs/locomotion/g1/test_g1_owner_contract.py +++ b/tests/envs/locomotion/g1/test_g1_owner_contract.py @@ -183,7 +183,8 @@ def test_g1_owner_yaml_regression_contract(case: dict[str, Any]): assert full_env_cfg.curriculum.enabled is case["curriculum_enabled"] assert env._uses_walk_observation_profile() is (case["profile"] == "walk") assert ( - registry._envs[cfg.training.task_name].env_cls_dict[cfg.training.sim_backend] is G1WalkEnv + registry._envs[cfg.training.task_name].env_factory_dict[cfg.training.sim_backend] + is G1WalkEnv ) if "model_suffix" in case: From 9ba6d81ff9930441e9b93f36e60eab8d50235892 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 18:10:36 +0800 Subject: [PATCH 027/169] feat(config): materialize Hydra-owned manager configs [skip ci] (#1099) --- ...-task-owner-and-config-compose-contract.md | 5 + ...-community-manager-api-on-numpy-runtime.md | 23 +- .../1-architecture/6-manager_based_api.md | 5 + src/unilab/base/config_materialization.py | 377 ++++++++++++++++++ src/unilab/base/registry.py | 168 +------- src/unilab/envs/manager_based_rl_env.py | 5 +- tests/base/test_manager_config_overlay.py | 251 +++++++++++- 7 files changed, 652 insertions(+), 182 deletions(-) create mode 100644 src/unilab/base/config_materialization.py diff --git a/docs/sphinx/source/adr/ADR-0003-task-owner-and-config-compose-contract.md b/docs/sphinx/source/adr/ADR-0003-task-owner-and-config-compose-contract.md index 99d5525f8..8db677e72 100644 --- a/docs/sphinx/source/adr/ADR-0003-task-owner-and-config-compose-contract.md +++ b/docs/sphinx/source/adr/ADR-0003-task-owner-and-config-compose-contract.md @@ -25,12 +25,17 @@ orphan: true 2. owner YAML 直接持有 `training.task_name`、`training.sim_backend`、`reward`、`env` 及 task-specific `algo`。 3. `training.sim_backend` 是 owner 身份字段，不是独立 backend switch。 4. CLI override 允许参数覆盖，但不能破坏 task owner 的 backend identity。 +5. Manager-Based production task 也不例外：owner YAML 完整持有 manager/term/callable 与 + observation mapping，compose 后在 Registry 冷路径物化为 typed cfg；Python 不保存 + task-specific config mirror。 ## Stable Contracts - PPO/APPO owner 路径: `conf/{ppo,appo}/task//.yaml` - Offpolicy owner 路径: `conf/offpolicy/task///.yaml` - reward 注入与 backend 差异表达必须在 owner YAML 层显式存在。 +- Manager-Based cfg 使用 Hydra `_target_` 与 dotted callable reference；解析失败或类型错误 + 必须在 env/backend 构造前报错。 ## Consequences diff --git a/docs/sphinx/source/adr/ADR-0006-community-manager-api-on-numpy-runtime.md b/docs/sphinx/source/adr/ADR-0006-community-manager-api-on-numpy-runtime.md index 45bc37dbd..c7eb30d17 100644 --- a/docs/sphinx/source/adr/ADR-0006-community-manager-api-on-numpy-runtime.md +++ b/docs/sphinx/source/adr/ADR-0006-community-manager-api-on-numpy-runtime.md @@ -89,9 +89,18 @@ Manager 内可以使用社区常见的 `policy` / `actor` / `critic` observation 映射为 `NpEnvState.obs["critic"]`。runner、learner 与 IPC 不推断、不拼接 group。 `reset() -> (obs_dict, info_dict)`、final observation 与 `obs_groups_spec` 保持现有 contract。 -Callable 和 typed config instance 由 task-owned Python factory 声明；Hydra owner YAML 只按 -manager/term 名覆盖可序列化字段。现有 CLI、registry、algorithm YAML 和 -`task=/` owner compose 保持唯一配置入口。Scripts 不解释 term 业务规则。 +Production task 的唯一配置 source of truth 是 Hydra owner YAML。它完整声明 scene/backend +tuning、manager/group/term 的顺序与启停、具体 cfg 类型、callable、params、weight 和 +observation group mapping；compose 后按以下冷路径进入现有 registry： + +`owner YAML -> DictConfig -> typed config materialization -> Registry factory -> ManagerBasedRlEnv` + +具体 cfg 类型使用 Hydra `_target_`，term callable 使用完整 dotted reference。通用 +materializer 将其解析为 plain dataclass instance，并在未知字段、target/callable 解析失败、 +抽象或错误 term cfg 类型及缺少必填字段时 fail-closed。Python 只拥有 term 实现、公共 cfg +类型和通用 factory，不保存第二份 task-specific term 清单或默认值；直接构造 typed cfg +只用于底层单测。DictConfig 与解析逻辑不能进入 reset/step 热路径，scripts 不解释 term +业务规则。 ### 4. Scene/entity owner boundary @@ -165,7 +174,7 @@ fallback 到旧单体 env 的永久兼容路径。 | observation groups、clip/scale/noise/delay/history | Adapted | 数值为 NumPy；group 在 env boundary 显式映射 | | manager buffers、env IDs、RNG | Adapted | Torch→NumPy；无 device API | | `ManagerBasedRlEnv` return | Adapted | 保留 `NpEnvState` 与 UniLab reset/final-observation contract | -| config container | Adapted | plain instances + Hydra owner YAML overlay，不引入第二套 runtime | +| config container | Adapted | Hydra owner YAML 唯一持有 task 配置，冷路径物化为 plain typed instances | | `SceneEntityCfg` selectors | Adapted | 语义保留；只解析 `SimBackend` 已声明能力 | | named sensor view | Adapted | 冷路径 bind；有序展平 NumPy batch；reader 由 backend 拥有 | | event/domain randomization | Adapted | 调度语义保留；mutation 走 backend DR/capability contract | @@ -225,6 +234,8 @@ Apache-2.0 和 UniLab 的修改类别。实现 PR 分别报告： - 先设计 compiler、fused term protocol 或专用 fast path。拒绝：在 benchmark 证明瓶颈前 增加结构复杂度，并可能牺牲社区 term 语义。 - 缺失能力 warning + skip 或回退旧 env。拒绝：配置表面与真实执行不一致，不能用于生产。 +- task-owned Python factory 声明 callable/term，Hydra 只做字段 overlay。拒绝：会让同一 task + 在 Python 和 YAML 中拥有两份配置，增加迁移 friction 和语义漂移。 ## Consequences @@ -233,6 +244,8 @@ Apache-2.0 和 UniLab 的修改类别。实现 PR 分别报告： compatibility matrix 中先记录。 - Scene/entity 采用最小 base facade，#586 不再阻塞 manager port；真实 backend 能力仍按 独立 child 和 conformance evidence 接入。 +- Config/Registry 永久维护一个通用 Hydra `_target_` / dotted callable 到 typed manager cfg + 的冷路径 materializer；production task 不维护 Python config mirror。 - 迁移初期允许 production 旧 task 与未接入的 manager package 同时存在，但 task 一旦迁移 就必须删除对应旧实现；umbrella 结束时不能保留双 lifecycle。 - 性能 gate 关注明显低效与实测瓶颈，不以复杂度换取未经证明的小收益。 @@ -243,7 +256,7 @@ Apache-2.0 和 UniLab 的修改类别。实现 PR 分别报告： - Backend contract: `src/unilab/base/backend/base.py` - Scene config owner: `src/unilab/base/scene.py` - Config schema and registry: `src/unilab/structured_configs.py`, - `src/unilab/base/registry.py`, `conf/` + `src/unilab/base/config_materialization.py`, `src/unilab/base/registry.py`, `conf/` - Observation/IPC contract: `docs/sphinx/source/adr/ADR-0005-unified-obs-critic-env-and-ipc-contract.md` - Layer boundary: `docs/sphinx/source/adr/ADR-0001-runtime-model-and-layer-boundaries.md` - Upstream checkout used for the decision: diff --git a/docs/sphinx/source/zh_CN/4-developer_guide/1-architecture/6-manager_based_api.md b/docs/sphinx/source/zh_CN/4-developer_guide/1-architecture/6-manager_based_api.md index 994c9623b..dcc9be804 100644 --- a/docs/sphinx/source/zh_CN/4-developer_guide/1-architecture/6-manager_based_api.md +++ b/docs/sphinx/source/zh_CN/4-developer_guide/1-architecture/6-manager_based_api.md @@ -10,6 +10,11 @@ NumPy，并保留现有 `NpEnvState`、Hydra owner YAML、`SimBackend`、registr ## 不变量 - 公共 manager 结构优先保持社区语义；不为局部性能制造 UniLab-only term API。 +- Production task 只从 Hydra owner YAML 配置：YAML 完整声明 manager/group/term、具体 cfg + `_target_`、dotted callable、params、weight 与 observation mapping；Registry 冷路径将其 + 物化为 plain typed cfg，Python 不保留 task config mirror。 +- 未知字段、无法解析的 target/callable、抽象或错误 cfg 类型直接报错；DictConfig 和解析 + 不进入 reset/step，scripts 不解释 task 业务规则。 - manager buffer、term return、env ID 和 entity view 使用 `np.ndarray` / `slice`，core 不依赖 Torch、Warp、runner、learner 或 IPC。 - `SceneEntityCfg` 在冷路径通过 base scene/entity facade 解析；facade 只调用正式 diff --git a/src/unilab/base/config_materialization.py b/src/unilab/base/config_materialization.py new file mode 100644 index 000000000..7e9e23517 --- /dev/null +++ b/src/unilab/base/config_materialization.py @@ -0,0 +1,377 @@ +"""Cold-path materialization of Hydra-owned typed configuration.""" + +from __future__ import annotations + +import dataclasses +import inspect +import types +from collections.abc import Mapping +from typing import Any, Union, get_args, get_origin, get_type_hints + +from hydra.utils import get_object, instantiate +from omegaconf import OmegaConf + +from .config_overrides import ( + CONFIG_MAPPING_POLICY_KEY, + MANAGER_PARAMS_MAPPING_POLICY, + MANAGER_TERM_MAPPING_POLICY, +) + +HYDRA_TARGET_KEY = "_target_" +_MISSING = object() + + +def _plain(value: Any) -> Any: + if OmegaConf.is_config(value): + return OmegaConf.to_container(value, resolve=True) + return value + + +def _hints(target_type: type[Any]) -> dict[str, Any]: + try: + return get_type_hints(target_type) + except (NameError, TypeError): + return {field.name: field.type for field in dataclasses.fields(target_type)} + + +def _dataclass_types(annotation: Any) -> tuple[type[Any], ...]: + if annotation in (Any, None): + return () + origin = get_origin(annotation) + if origin in (types.UnionType, Union): + return tuple( + target + for item in get_args(annotation) + if item is not type(None) + for target in _dataclass_types(item) + ) + if isinstance(annotation, type) and dataclasses.is_dataclass(annotation): + return (annotation,) + return () + + +def _dict_value_type(annotation: Any) -> Any: + origin = get_origin(annotation) + if origin is dict: + args = get_args(annotation) + return args[1] if len(args) == 2 else Any + if origin in (types.UnionType, Union): + for item in get_args(annotation): + value_type = _dict_value_type(item) + if value_type is not _MISSING: + return value_type + return _MISSING + + +def _resolve(reference: Any, *, path: str) -> Any: + if not isinstance(reference, str) or not reference.strip(): + raise TypeError(f"Config field '{path}' must be a non-empty dotted string") + try: + return get_object(reference) + except Exception as exc: + raise ValueError( + f"Config field '{path}' could not resolve dotted reference {reference!r}: {exc}" + ) from exc + + +def _resolve_target( + reference: Any, + *, + expected: Any, + path: str, +) -> type[Any]: + target = _resolve(reference, path=f"{path}.{HYDRA_TARGET_KEY}") + if not isinstance(target, type) or not dataclasses.is_dataclass(target): + raise TypeError( + f"Config field '{path}.{HYDRA_TARGET_KEY}' must resolve to a dataclass type" + ) + expected_types = _dataclass_types(expected) + if expected_types and not any(issubclass(target, item) for item in expected_types): + names = ", ".join(item.__qualname__ for item in expected_types) + raise TypeError( + f"Config field '{path}.{HYDRA_TARGET_KEY}' resolved to {target.__qualname__}, " + f"expected a subclass of {names}" + ) + if inspect.isabstract(target): + raise TypeError( + f"Config field '{path}.{HYDRA_TARGET_KEY}' resolved to abstract config " + f"{target.__qualname__}; select a concrete term config" + ) + return target + + +def _target_path(target: type[Any]) -> str: + return f"{target.__module__}.{target.__qualname__}" + + +def _prepare_value(value: Any, *, annotation: Any, path: str) -> Any: + value = _plain(value) + if isinstance(value, Mapping): + values = dict(value) + if HYDRA_TARGET_KEY in values: + return _prepare_dataclass(values, expected=annotation, path=path, require_target=True) + + candidates = _dataclass_types(annotation) + if len(candidates) == 1 and not inspect.isabstract(candidates[0]): + return _prepare_dataclass(values, expected=annotation, path=path, require_target=False) + + value_type = _dict_value_type(annotation) + return { + key: _prepare_value( + item, + annotation=Any if value_type is _MISSING else value_type, + path=f"{path}.{key}", + ) + for key, item in values.items() + } + if isinstance(value, list): + return [ + _prepare_value(item, annotation=Any, path=f"{path}[{index}]") + for index, item in enumerate(value) + ] + return value + + +def _prepare_manager_mapping(value: Any, *, annotation: Any, path: str) -> dict[str, Any]: + value = _plain(value) + if not isinstance(value, Mapping): + raise TypeError(f"Config field '{path}' must be a mapping") + value_type = _dict_value_type(annotation) + if value_type is _MISSING: + raise TypeError(f"Config field '{path}' manager policy requires a typed dict") + + result: dict[str, Any] = {} + for name, raw_entry in value.items(): + if not isinstance(name, str) or not name: + raise TypeError(f"Config field '{path}' term names must be non-empty strings") + entry_path = f"{path}.{name}" + entry = _plain(raw_entry) + if entry is None: + result[name] = None + elif not isinstance(entry, Mapping): + raise TypeError(f"Config field '{entry_path}' must be a mapping or None") + else: + result[name] = _prepare_dataclass( + dict(entry), + expected=value_type, + path=entry_path, + require_target=True, + ) + return result + + +def _prepare_dataclass( + values: Mapping[str, Any], + *, + expected: Any, + path: str, + require_target: bool, +) -> dict[str, Any]: + values = dict(values) + reference = values.pop(HYDRA_TARGET_KEY, _MISSING) + if reference is _MISSING: + candidates = _dataclass_types(expected) + if require_target: + raise ValueError( + f"Config field '{path}' is a new Manager-Based entry and must declare " + f"'{HYDRA_TARGET_KEY}'" + ) + if len(candidates) != 1 or inspect.isabstract(candidates[0]): + raise ValueError( + f"Config field '{path}' cannot infer one concrete dataclass type; " + f"declare '{HYDRA_TARGET_KEY}'" + ) + target = candidates[0] + reference = _target_path(target) + else: + target = _resolve_target(reference, expected=expected, path=path) + + fields = {field.name: field for field in dataclasses.fields(target) if field.init} + unknown = [name for name in values if name not in fields] + if unknown: + raise ValueError( + f"Config field '{path}' target {target.__qualname__} has no fields {unknown}" + ) + + hints = _hints(target) + prepared: dict[str, Any] = {HYDRA_TARGET_KEY: reference} + for name, raw_value in values.items(): + field = fields[name] + field_path = f"{path}.{name}" + annotation = hints.get(name, field.type) + policy = field.metadata.get(CONFIG_MAPPING_POLICY_KEY) + if policy == MANAGER_TERM_MAPPING_POLICY: + prepared[name] = _prepare_manager_mapping( + raw_value, + annotation=annotation, + path=field_path, + ) + elif name == "func": + resolved = _resolve(raw_value, path=field_path) + if not callable(resolved): + raise TypeError( + f"Config field '{field_path}' resolved to {type(resolved).__name__}, " + "expected a callable" + ) + prepared[name] = resolved + else: + prepared[name] = _prepare_value(raw_value, annotation=annotation, path=field_path) + return prepared + + +def _materialize_entry(value: Mapping[str, Any], *, expected: Any, path: str) -> Any: + prepared = _prepare_dataclass(value, expected=expected, path=path, require_target=True) + try: + result = instantiate(prepared, _convert_="all") + except Exception as exc: + raise TypeError( + f"Config field '{path}' could not construct its typed config: {exc}" + ) from exc + expected_types = _dataclass_types(expected) + if expected_types and not isinstance(result, expected_types): + raise TypeError( + f"Config field '{path}' materialized {type(result).__name__}, expected " + f"{', '.join(item.__qualname__ for item in expected_types)}" + ) + return result + + +def _policy(target_obj: Any, name: str) -> str | None: + if not dataclasses.is_dataclass(target_obj) or isinstance(target_obj, type): + return None + field = next((item for item in dataclasses.fields(target_obj) if item.name == name), None) + value = None if field is None else field.metadata.get(CONFIG_MAPPING_POLICY_KEY) + return str(value) if value is not None else None + + +def _is_term_cfg(target_obj: Any) -> bool: + return ( + dataclasses.is_dataclass(target_obj) + and not isinstance(target_obj, type) + and any( + field.metadata.get(CONFIG_MAPPING_POLICY_KEY) == MANAGER_PARAMS_MAPPING_POLICY + for field in dataclasses.fields(target_obj) + ) + ) + + +def _apply_manager_mapping( + target_obj: Any, + name: str, + overrides: Any, + *, + annotation: Any, + policy: str, +) -> None: + owner = f"{type(target_obj).__name__}.{name}" + existing = getattr(target_obj, name) + overrides = _plain(overrides) + if not isinstance(existing, dict) or not isinstance(overrides, Mapping): + raise TypeError(f"Config field '{owner}' manager value and override must be mappings") + + if policy == MANAGER_PARAMS_MAPPING_POLICY: + for param_name, raw_value in overrides.items(): + path = f"{owner}.{param_name}" + value = _plain(raw_value) + current = existing.get(param_name, _MISSING) + if isinstance(value, Mapping) and HYDRA_TARGET_KEY in value: + existing[param_name] = _materialize_entry(value, expected=Any, path=path) + elif isinstance(value, Mapping) and dataclasses.is_dataclass(current): + apply_cfg_overrides(current, value, _path=path) + else: + existing[param_name] = _prepare_value(value, annotation=Any, path=path) + return + if policy != MANAGER_TERM_MAPPING_POLICY: + raise ValueError(f"Config field '{owner}' has unknown mapping policy {policy!r}") + + value_type = _dict_value_type(annotation) + if value_type is _MISSING: + raise TypeError(f"Config field '{owner}' manager policy requires a typed dict") + for term_name, raw_value in overrides.items(): + path = f"{owner}.{term_name}" + value = _plain(raw_value) + if value is None: + existing[term_name] = None + elif not isinstance(value, Mapping): + raise TypeError(f"Config field '{path}' must be a field mapping or None") + elif HYDRA_TARGET_KEY in value: + existing[term_name] = _materialize_entry(value, expected=value_type, path=path) + else: + current = existing.get(term_name, _MISSING) + if current is _MISSING or current is None: + raise ValueError( + f"Config field '{path}' is a new Manager-Based entry and must declare " + f"'{HYDRA_TARGET_KEY}'" + ) + if not dataclasses.is_dataclass(current) or isinstance(current, type): + raise TypeError(f"Config field '{path}' does not contain a dataclass config") + apply_cfg_overrides(current, value, _path=path) + + +def apply_cfg_overrides( + target_obj: Any, + overrides: Mapping[str, Any], + *, + _path: str | None = None, +) -> None: + """Apply overrides and materialize explicitly typed Manager-Based entries.""" + overrides = _plain(overrides) + if not isinstance(overrides, Mapping): + raise TypeError(f"Config overrides for {type(target_obj).__name__} must be a mapping") + hints = _hints(type(target_obj)) if dataclasses.is_dataclass(target_obj) else {} + fields = ( + {field.name: field for field in dataclasses.fields(target_obj)} + if dataclasses.is_dataclass(target_obj) and not isinstance(target_obj, type) + else {} + ) + owner_path = _path or type(target_obj).__name__ + + for key, raw_value in overrides.items(): + if not isinstance(key, str) or not hasattr(target_obj, key): + raise ValueError(f"Config class '{type(target_obj).__name__}' has no attribute '{key}'") + value = _plain(raw_value) + existing = getattr(target_obj, key) + annotation = hints.get(key, fields[key].type if key in fields else Any) + policy = _policy(target_obj, key) + if policy is not None: + _apply_manager_mapping( + target_obj, + key, + value, + annotation=annotation, + policy=policy, + ) + continue + if key == "func" and _is_term_cfg(target_obj): + raise ValueError( + f"Config field '{type(target_obj).__name__}.func' belongs to the typed term " + "declaration and cannot be changed by a partial override" + ) + + path = f"{owner_path}.{key}" + if isinstance(value, Mapping): + if HYDRA_TARGET_KEY in value: + setattr(target_obj, key, _materialize_entry(value, expected=annotation, path=path)) + continue + if dataclasses.is_dataclass(existing) and not isinstance(existing, type): + apply_cfg_overrides(existing, value, _path=path) + continue + candidates = _dataclass_types(annotation) + if existing is None and len(candidates) == 1: + prepared = _prepare_dataclass( + value, + expected=annotation, + path=path, + require_target=False, + ) + try: + setattr(target_obj, key, instantiate(prepared, _convert_="all")) + except Exception as exc: + raise TypeError( + f"Config field '{path}' could not construct its typed config: {exc}" + ) from exc + continue + setattr(target_obj, key, _prepare_value(value, annotation=Any, path=path)) + + +__all__ = ["HYDRA_TARGET_KEY", "apply_cfg_overrides"] diff --git a/src/unilab/base/registry.py b/src/unilab/base/registry.py index 48ed19977..cf5089321 100644 --- a/src/unilab/base/registry.py +++ b/src/unilab/base/registry.py @@ -11,18 +11,13 @@ Literal, Optional, Protocol, - Type, TypeVar, - cast, - get_args, - get_origin, - get_type_hints, ) from .base import ABEnv, EnvCfg +from .config_materialization import apply_cfg_overrides from .config_overrides import ( CONFIG_MAPPING_POLICY_KEY, - MANAGER_PARAMS_MAPPING_POLICY, MANAGER_TERM_MAPPING_POLICY, ) @@ -239,167 +234,6 @@ def resolve_reward_override_field(env_name: str) -> RewardOverrideField: ) -def _resolve_dataclass_type(type_hint: Any) -> Optional[Type[Any]]: - """Strip Optional/Union and return the underlying dataclass type, or None.""" - if type_hint is None: - return None - origin = get_origin(type_hint) - if origin is not None: - args = get_args(type_hint) - type_hint = next((arg for arg in args if arg is not type(None)), None) - if ( - type_hint is not None - and dataclasses.is_dataclass(type_hint) - and isinstance(type_hint, type) - ): - return cast(Type[Any], type_hint) - return None - - -def _construct_dataclass_from_dict(target_type: Type[Any], values: Dict[str, Any]) -> Any: - try: - target_obj = target_type() - except TypeError: - return target_type(**values) - apply_cfg_overrides(target_obj, values) - return target_obj - - -def _config_mapping_policy(target_obj: Any, field_name: str) -> str | None: - if not dataclasses.is_dataclass(target_obj) or isinstance(target_obj, type): - return None - for config_field in dataclasses.fields(target_obj): - if config_field.name == field_name: - policy = config_field.metadata.get(CONFIG_MAPPING_POLICY_KEY) - return str(policy) if policy is not None else None - return None - - -def _is_manager_callable_term_cfg(target_obj: Any) -> bool: - if not dataclasses.is_dataclass(target_obj) or isinstance(target_obj, type): - return False - return any( - config_field.metadata.get(CONFIG_MAPPING_POLICY_KEY) == MANAGER_PARAMS_MAPPING_POLICY - for config_field in dataclasses.fields(target_obj) - ) - - -def _apply_manager_mapping_overrides( - target_obj: Any, - field_name: str, - existing: Any, - overrides: Any, - *, - policy: str, -) -> None: - owner = f"{type(target_obj).__name__}.{field_name}" - if not isinstance(existing, dict): - raise TypeError( - f"Config field '{owner}' declares manager mapping policy but contains " - f"{type(existing).__name__}, expected dict" - ) - if not isinstance(overrides, dict): - raise TypeError( - f"Config field '{owner}' must be overridden by a mapping, not " - f"{type(overrides).__name__}" - ) - - if policy == MANAGER_PARAMS_MAPPING_POLICY: - for param_name, value in overrides.items(): - current = existing.get(param_name) - if isinstance(value, dict) and dataclasses.is_dataclass(current): - apply_cfg_overrides(current, value) - else: - existing[param_name] = value - return - - if policy != MANAGER_TERM_MAPPING_POLICY: - raise ValueError(f"Config field '{owner}' has unknown mapping policy {policy!r}") - - for term_name, value in overrides.items(): - if term_name not in existing: - raise ValueError( - f"Config field '{owner}' has no factory-owned term '{term_name}'; " - "declare its callable/config in the task Python factory first" - ) - if value is None: - existing[term_name] = None - continue - - current = existing[term_name] - if current is None: - raise ValueError( - f"Config field '{owner}' term '{term_name}' is disabled; set a concrete " - "config in the task Python factory before overriding its fields" - ) - if not dataclasses.is_dataclass(current): - raise TypeError( - f"Config field '{owner}' term '{term_name}' contains " - f"{type(current).__name__}, expected a dataclass config" - ) - if not isinstance(value, dict): - raise TypeError( - f"Config field '{owner}' term '{term_name}' must be overridden by a " - "field mapping or None; replacing the factory-owned term is not allowed" - ) - apply_cfg_overrides(current, value) - - -def apply_cfg_overrides(target_obj: Any, overrides: Dict[str, Any]) -> None: - """Apply a (possibly nested) dict of overrides to ``target_obj`` in place. - - Behavior: - - For each ``key, value`` in ``overrides``, ``target_obj.key`` must exist - (otherwise ``ValueError``). - - If ``value`` is a dict and ``target_obj.key`` is already a dataclass - instance, recurse into it (deep merge — preserves fields not present - in ``value``). This is what lets Hydra-style partial overrides like - ``env.scene.terrain.generator.num_rows=4`` keep ``sub_terrains`` and other - defaults intact. - - If ``value`` is a dict and ``target_obj.key`` is currently ``None``, - instantiate the field's annotated dataclass type from the dict - (full-construction path). - - Fields explicitly marked as manager mappings merge only existing - factory-owned entries. ``None`` disables an entry; unknown entries and - callable/config replacement fail closed. - - Otherwise ``setattr`` the value directly (scalar / list / non-dataclass). - """ - try: - type_hints = get_type_hints(type(target_obj)) - except Exception: - type_hints = {} - - for key, value in overrides.items(): - if not hasattr(target_obj, key): - raise ValueError(f"Config class '{type(target_obj).__name__}' has no attribute '{key}'") - existing = getattr(target_obj, key) - mapping_policy = _config_mapping_policy(target_obj, key) - if mapping_policy is not None: - _apply_manager_mapping_overrides( - target_obj, - key, - existing, - value, - policy=mapping_policy, - ) - continue - if key == "func" and _is_manager_callable_term_cfg(target_obj): - raise ValueError( - f"Config field '{type(target_obj).__name__}.func' is factory-owned and " - "cannot be overridden" - ) - if isinstance(value, dict): - if dataclasses.is_dataclass(existing) and not isinstance(existing, type): - apply_cfg_overrides(existing, value) - continue - if existing is None: - target_type = _resolve_dataclass_type(type_hints.get(key)) - if target_type is not None: - setattr(target_obj, key, _construct_dataclass_from_dict(target_type, value)) - continue - setattr(target_obj, key, value) - - def make( name: str, sim_backend: Optional[str] = None, diff --git a/src/unilab/envs/manager_based_rl_env.py b/src/unilab/envs/manager_based_rl_env.py index a7b30fa1b..a99552231 100644 --- a/src/unilab/envs/manager_based_rl_env.py +++ b/src/unilab/envs/manager_based_rl_env.py @@ -62,8 +62,9 @@ def _manager_terms_field() -> Any: class ManagerBasedRlEnvCfg(EnvCfg): """Configuration for the manager-based NumPy environment. - Serializable values remain ordinary dataclass fields so task-owned Hydra owner - configs can overlay them without introducing a second configuration runtime. + Production task owners declare these fields in Hydra. The Registry materializes + them into this plain typed config on the cold path; Python factories do not mirror + task-specific manager or term declarations. """ observations: dict[str, ObservationGroupCfg | None] = _manager_terms_field() diff --git a/tests/base/test_manager_config_overlay.py b/tests/base/test_manager_config_overlay.py index 5cb1e9087..443ce5ec2 100644 --- a/tests/base/test_manager_config_overlay.py +++ b/tests/base/test_manager_config_overlay.py @@ -1,19 +1,24 @@ from __future__ import annotations -from dataclasses import dataclass, field +from dataclasses import dataclass, field, fields, is_dataclass import numpy as np import pytest +from omegaconf import OmegaConf from unilab.base.registry import apply_cfg_overrides from unilab.envs.manager_based_rl_env import ManagerBasedRlEnvCfg -from unilab.envs.mdp import JointPositionActionCfg +from unilab.envs.mdp import JointPositionActionCfg, UniformVelocityCommandCfg from unilab.managers import ( + CurriculumTermCfg, EventTermCfg, + MetricsTermCfg, ObservationGroupCfg, ObservationTermCfg, + RecorderTermCfg, RewardTermCfg, SceneEntityCfg, + TerminationTermCfg, ) @@ -68,7 +73,9 @@ def _manager_cfg() -> ManagerBasedRlEnvCfg: def test_manager_mapping_overlay_preserves_factory_terms_and_order() -> None: cfg = _manager_cfg() - reward_func = cfg.rewards["tracking"].func + initial_tracking = cfg.rewards["tracking"] + assert initial_tracking is not None + reward_func = initial_tracking.func apply_cfg_overrides( cfg, @@ -100,7 +107,7 @@ def test_manager_mapping_overlay_preserves_factory_terms_and_order() -> None: assert reset.func is _first_term assert reset.min_step_count_between_reset == 3 action = cfg.actions["joint_pos"] - assert action is not None + assert isinstance(action, JointPositionActionCfg) assert action.entity_name == "robot" assert action.actuator_names == (".*",) assert action.scale == pytest.approx(0.4) @@ -132,14 +139,152 @@ def test_observation_group_and_term_overlay_preserve_siblings() -> None: assert list(policy.terms) == ["first", "second"] +def _assert_no_omegaconf(value: object) -> None: + assert not OmegaConf.is_config(value) + if is_dataclass(value) and not isinstance(value, type): + for config_field in fields(value): + _assert_no_omegaconf(getattr(value, config_field.name)) + elif isinstance(value, dict): + for item in value.values(): + _assert_no_omegaconf(item) + elif isinstance(value, (list, tuple)): + for item in value: + _assert_no_omegaconf(item) + + +def test_hydra_mapping_fully_materializes_empty_manager_config() -> None: + cfg = ManagerBasedRlEnvCfg() + hydra_mapping = OmegaConf.create( + { + "scene": { + "model_file": "robot.xml", + "entities": { + "robot": { + "root_body_name": "base", + "joint_names": ["joint"], + "actuator_names": ["motor"], + } + }, + }, + "sim_dt": 0.01, + "ctrl_dt": 0.02, + "max_episode_seconds": 20.0, + "observations": { + "policy": { + "_target_": "unilab.managers.ObservationGroupCfg", + "terms": { + "joint_pos": { + "_target_": "unilab.managers.ObservationTermCfg", + "func": "unilab.envs.mdp.joint_pos_rel", + "params": { + "asset_cfg": { + "_target_": "unilab.managers.SceneEntityCfg", + "name": "robot", + "joint_names": ".*", + } + }, + }, + "disabled": None, + }, + } + }, + "actions": { + "joint_pos": { + "_target_": "unilab.envs.mdp.JointPositionActionCfg", + "entity_name": "robot", + "actuator_names": [".*"], + "scale": 0.25, + } + }, + "commands": { + "twist": { + "_target_": "unilab.envs.mdp.UniformVelocityCommandCfg", + "entity_name": "robot", + "resampling_time_range": [1.0, 1.0], + "ranges": { + "lin_vel_x": [-1.0, 1.0], + "lin_vel_y": [-0.5, 0.5], + "ang_vel_z": [-1.0, 1.0], + }, + } + }, + "events": { + "reset": { + "_target_": "unilab.managers.EventTermCfg", + "func": "unilab.envs.mdp.reset_scene_to_default", + "mode": "reset", + } + }, + "rewards": { + "alive": { + "_target_": "unilab.managers.RewardTermCfg", + "func": "unilab.envs.mdp.is_alive", + "weight": 1.0, + } + }, + "terminations": { + "time_out": { + "_target_": "unilab.managers.TerminationTermCfg", + "func": "unilab.envs.mdp.time_out", + "time_out": True, + } + }, + "curriculum": { + "difficulty": { + "_target_": "unilab.managers.CurriculumTermCfg", + "func": "unilab.envs.mdp.is_alive", + } + }, + "metrics": { + "alive": { + "_target_": "unilab.managers.MetricsTermCfg", + "func": "unilab.envs.mdp.is_alive", + } + }, + "recorders": { + "trace": { + "_target_": "unilab.managers.RecorderTermCfg", + "func": "unilab.managers.RecorderTerm", + } + }, + "policy_observation_group": "policy", + } + ) + + apply_cfg_overrides(cfg, hydra_mapping) + cfg.validate() + + assert cfg.scene is not None + assert cfg.scene.entities["robot"].joint_names == ["joint"] + assert list(cfg.observations) == ["policy"] + policy = cfg.observations["policy"] + assert isinstance(policy, ObservationGroupCfg) + assert list(policy.terms) == ["joint_pos", "disabled"] + joint_obs = policy.terms["joint_pos"] + assert isinstance(joint_obs, ObservationTermCfg) + assert joint_obs.func is not None and callable(joint_obs.func) + assert isinstance(joint_obs.params["asset_cfg"], SceneEntityCfg) + assert isinstance(cfg.actions["joint_pos"], JointPositionActionCfg) + twist = cfg.commands["twist"] + assert isinstance(twist, UniformVelocityCommandCfg) + assert isinstance(twist.ranges, UniformVelocityCommandCfg.Ranges) + assert isinstance(cfg.events["reset"], EventTermCfg) + assert isinstance(cfg.rewards["alive"], RewardTermCfg) + assert isinstance(cfg.terminations["time_out"], TerminationTermCfg) + assert isinstance(cfg.curriculum["difficulty"], CurriculumTermCfg) + assert isinstance(cfg.metrics["alive"], MetricsTermCfg) + assert isinstance(cfg.recorders["trace"], RecorderTermCfg) + _assert_no_omegaconf(cfg) + + @pytest.mark.parametrize( ("overrides", "match"), [ - ({"rewards": {"missing": {"weight": 1.0}}}, "rewards.*missing"), - ({"rewards": {"disabled": {"weight": 1.0}}}, "disabled.*task Python factory"), - ({"rewards": {"tracking": _second_term}}, "tracking.*replacing"), + ({"rewards": {"missing": {"weight": 1.0}}}, "missing.*_target_"), + ({"rewards": {"disabled": {"weight": 1.0}}}, "disabled.*_target_"), + ({"rewards": {"tracking": _second_term}}, "tracking.*field mapping"), ({"rewards": []}, "rewards.*mapping"), - ({"rewards": {"tracking": {"func": _second_term}}}, "func.*factory-owned"), + ({"rewards": {"tracking": {"func": _second_term}}}, "func.*typed term"), ], ) def test_manager_mapping_overlay_fails_closed(overrides: dict, match: str) -> None: @@ -147,6 +292,96 @@ def test_manager_mapping_overlay_fails_closed(overrides: dict, match: str) -> No apply_cfg_overrides(_manager_cfg(), overrides) +@pytest.mark.parametrize( + ("entry", "match"), + [ + ( + { + "_target_": "unilab.managers.EventTermCfg", + "func": "unilab.envs.mdp.is_alive", + "mode": "reset", + }, + "EventTermCfg.*RewardTermCfg", + ), + ( + {"_target_": "unilab.envs.mdp.is_alive", "func": "unilab.envs.mdp.is_alive"}, + "dataclass type", + ), + ( + {"_target_": "unilab.managers.MissingCfg", "func": "unilab.envs.mdp.is_alive"}, + "_target_.*could not resolve", + ), + ( + { + "_target_": "unilab.managers.RewardTermCfg", + "func": "unilab.envs.mdp.missing", + "weight": 1.0, + }, + "func.*could not resolve", + ), + ( + { + "_target_": "unilab.managers.RewardTermCfg", + "func": "unilab.base.config_overrides.CONFIG_MAPPING_POLICY_KEY", + "weight": 1.0, + }, + "expected a callable", + ), + ( + { + "_target_": "unilab.managers.RewardTermCfg", + "func": "unilab.envs.mdp.is_alive", + "weight": 1.0, + "unknown": 1, + }, + "has no fields.*unknown", + ), + ( + {"_target_": "unilab.managers.RewardTermCfg", "func": "unilab.envs.mdp.is_alive"}, + "could not construct", + ), + ], +) +def test_hydra_manager_materialization_fails_closed(entry: dict, match: str) -> None: + with pytest.raises((TypeError, ValueError), match=match): + apply_cfg_overrides(ManagerBasedRlEnvCfg(), {"rewards": {"term": entry}}) + + +def test_hydra_manager_materialization_rejects_abstract_action_config() -> None: + with pytest.raises(TypeError, match="abstract config.*ActionTermCfg"): + apply_cfg_overrides( + ManagerBasedRlEnvCfg(), + { + "actions": { + "joint": { + "_target_": "unilab.managers.ActionTermCfg", + "entity_name": "robot", + } + } + }, + ) + + +def test_hydra_manager_materialization_rejects_wrong_nested_config_type() -> None: + with pytest.raises(TypeError, match="SceneCfg.*Ranges"): + apply_cfg_overrides( + ManagerBasedRlEnvCfg(), + { + "commands": { + "twist": { + "_target_": "unilab.envs.mdp.UniformVelocityCommandCfg", + "entity_name": "robot", + "resampling_time_range": [1.0, 1.0], + "ranges": { + "_target_": "unilab.base.scene.SceneCfg", + "model_file": "robot.xml", + }, + } + } + }, + ) + + @dataclass class _LegacyCfg: reward_config: dict[str, object] = field( From d81e8bd7da5388ce265f3a12422476189ecdf584 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 18:40:20 +0800 Subject: [PATCH 028/169] feat(env): add Manager-Based randomization events [skip ci] (#1101) Closes #1100. --- src/unilab/base/entity.py | 230 +++++++++ src/unilab/base/reset_state.py | 299 +++++++++++- src/unilab/envs/mdp/__init__.py | 10 + src/unilab/envs/mdp/events.py | 443 +++++++++++++++++- src/unilab/managers/_types.py | 6 + .../locomotion/go2/test_manager_based_cfg.py | 70 ++- tests/envs/mdp/test_events.py | 237 +++++++++- 7 files changed, 1276 insertions(+), 19 deletions(-) diff --git a/src/unilab/base/entity.py b/src/unilab/base/entity.py index 6e51998b8..4bfd980a5 100644 --- a/src/unilab/base/entity.py +++ b/src/unilab/base/entity.py @@ -16,6 +16,7 @@ import numpy as np from unilab.base.backend.base import BackendRootStateLayout, BackendSensorView, SimBackend +from unilab.dr.types import IntervalRandomizationPlan from unilab.utils.rotation import np_quat_apply, np_quat_apply_inverse, np_yaw_from_quat if TYPE_CHECKING: @@ -441,6 +442,7 @@ def __init__( (cfg.root_body_name,), backend.get_body_ids, ) + self._root_body_ids = root_body_ids joint_pos_ids = joint_vel_ids = None if self._joint_names is not None: @@ -458,6 +460,7 @@ def __init__( body_ids = None if self._body_names is not None: body_ids = self._resolve_ids("body", self._body_names, backend.get_body_ids) + self._body_ids = body_ids self._geom_ids = None if self._geom_names is not None: @@ -1088,6 +1091,144 @@ def write_actuator_gains_to_sim( term_name=f"{term_name}:{self.name}", ) + def bind_body_mass_write( + self, + body_ids: np.ndarray | Sequence[int] | slice | None = None, + *, + term_name: str, + ) -> tuple[np.ndarray, np.ndarray]: + """Bind entity-local body columns and immutable default masses.""" + reset_state, local_ids, backend_ids = self._bind_body_randomization( + body_ids, + capability="reset body-mass write", + ) + _, defaults = reset_state.bind_body_mass_write( + backend_ids, + term_name=f"{term_name}:{self.name}", + ) + return self._readonly_local_binding(local_ids, defaults) + + def write_body_mass_to_sim( + self, + values: np.ndarray, + body_ids: np.ndarray | Sequence[int] | slice | None = None, + env_ids: np.ndarray | slice | None = None, + *, + term_name: str = "randomize_rigid_body_mass", + ) -> None: + """Stage selected entity body masses in the active reset transaction.""" + reset_state, _, backend_ids = self._bind_body_randomization( + body_ids, + capability="reset body-mass write", + ) + reset_state.write_body_mass( + self._normalize_reset_env_ids(env_ids), + backend_ids, + values, + term_name=f"{term_name}:{self.name}", + ) + + def bind_body_ipos_write( + self, + body_ids: np.ndarray | Sequence[int] | slice | None = None, + *, + term_name: str, + ) -> tuple[np.ndarray, np.ndarray]: + """Bind entity-local body columns and immutable inertial positions.""" + reset_state, local_ids, backend_ids = self._bind_body_randomization( + body_ids, + capability="reset body-ipos write", + ) + _, defaults = reset_state.bind_body_ipos_write( + backend_ids, + term_name=f"{term_name}:{self.name}", + ) + return self._readonly_local_binding(local_ids, defaults) + + def write_body_ipos_to_sim( + self, + values: np.ndarray, + body_ids: np.ndarray | Sequence[int] | slice | None = None, + env_ids: np.ndarray | slice | None = None, + *, + term_name: str = "randomize_rigid_body_com", + ) -> None: + """Stage selected entity body inertial positions in the reset transaction.""" + reset_state, _, backend_ids = self._bind_body_randomization( + body_ids, + capability="reset body-ipos write", + ) + reset_state.write_body_ipos( + self._normalize_reset_env_ids(env_ids), + backend_ids, + values, + term_name=f"{term_name}:{self.name}", + ) + + def bind_root_linear_velocity_delta(self, *, term_name: str) -> None: + """Validate the interval root-velocity capability on the cold path.""" + if self._root_body_ids is None: + raise self._capability_error( + "interval root velocity delta", + "root_body_name was not declared in EntityCfg", + ) + try: + capabilities = self._backend.get_dr_capabilities() + except (AttributeError, NotImplementedError) as exc: + raise self._capability_error("interval root velocity delta", str(exc)) from exc + if not capabilities.supports_interval_body_velocity_delta: + raise self._capability_error( + "interval root velocity delta", + f"EventManager term '{term_name}' requested an unsupported backend capability", + ) + + def apply_root_linear_velocity_delta_to_sim( + self, + values: np.ndarray, + env_ids: np.ndarray | slice | None = None, + *, + term_name: str = "push_by_setting_velocity", + ) -> None: + """Dispatch a cached root linear-velocity delta through the formal interval plan.""" + if self._root_body_ids is None: + raise self._capability_error( + "interval root velocity delta", + "root_body_name was not declared in EntityCfg", + ) + ids = self._normalize_reset_env_ids(env_ids) + if not isinstance(values, np.ndarray): + raise TypeError( + f"EventManager term '{term_name}' root velocity delta must be np.ndarray, " + f"got {type(values).__name__}" + ) + expected = (ids.size, 3) + if values.shape != expected: + raise ValueError( + f"EventManager term '{term_name}' root velocity delta has shape " + f"{values.shape}; expected {expected}" + ) + if not np.issubdtype(values.dtype, np.floating) or not np.isfinite(values).all(): + raise ValueError( + f"EventManager term '{term_name}' root velocity delta must be finite floating data" + ) + delta = np.zeros( + (self._backend.num_envs, len(self._root_body_ids), 3), + dtype=values.dtype, + ) + delta[ids, 0, :] = values + try: + self._backend.apply_interval_randomization( + IntervalRandomizationPlan( + body_ids=self._root_body_ids, + body_linear_velocity_delta=delta, + ) + ) + except NotImplementedError as exc: + raise self._capability_error( + "interval root velocity delta", + f"EventManager term '{term_name}': {exc}", + ) from exc + def write_root_link_pose_to_sim( self, root_pose: np.ndarray, @@ -1190,6 +1331,71 @@ def _materialize_reset_joint_indices(self) -> None: label=f"Entity '{self.name}' reset qvel", ) + def _bind_body_randomization( + self, + body_ids: np.ndarray | Sequence[int] | slice | None, + *, + capability: str, + ) -> tuple[ResetStateTransaction, np.ndarray, np.ndarray]: + if self._reset_state is None: + raise self._capability_error( + capability, + "EntityScene was materialized without an env-owned reset transaction", + ) + if self._body_ids is None: + raise self._capability_error( + capability, + "body_names were not declared in EntityCfg", + ) + local_ids = self._normalize_local_body_ids(body_ids, capability=capability) + if local_ids.size == 0: + raise ValueError(f"Entity '{self.name}' {capability} selected no bodies") + return self._reset_state, local_ids, self._body_ids[local_ids] + + def _readonly_local_binding( + self, + local_ids: np.ndarray, + defaults: np.ndarray, + ) -> tuple[np.ndarray, np.ndarray]: + bound_ids = np.array(local_ids, copy=True) + bound_ids.setflags(write=False) + bound_defaults = np.array(defaults, copy=True) + bound_defaults.setflags(write=False) + return bound_ids, bound_defaults + + def _normalize_local_body_ids( + self, + body_ids: np.ndarray | Sequence[int] | slice | None, + *, + capability: str, + ) -> np.ndarray: + if body_ids is None: + ids = np.arange(self.num_bodies, dtype=np.intp) + elif isinstance(body_ids, slice): + ids = np.arange(self.num_bodies, dtype=np.intp)[body_ids] + else: + raw = np.asarray(body_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} body_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_bodies): + raise IndexError( + f"Entity '{self.name}' {capability} body_ids out of range for " + f"{self.num_bodies} bodies: {ids.tolist()}" + ) + if np.unique(ids).size != ids.size: + raise ValueError( + f"Entity '{self.name}' {capability} body_ids contain duplicates: {ids.tolist()}" + ) + return ids + def _normalize_local_joint_ids( self, joint_ids: np.ndarray | Sequence[int] | slice | None, @@ -1397,6 +1603,30 @@ def reset_to_default(self, env_ids: np.ndarray, *, term_name: str) -> None: ) self._reset_state.reset_to_default(env_ids, term_name=term_name) + def bind_gravity_write(self, *, term_name: str) -> np.ndarray: + """Bind immutable backend gravity for a reset event on the cold path.""" + if self._reset_state is None: + raise NotImplementedError( + f"EventManager term '{term_name}' gravity capability is unavailable: " + "EntityScene was materialized without an env-owned reset transaction" + ) + return self._reset_state.bind_gravity_write(term_name=term_name) + + def write_gravity_to_sim( + self, + values: np.ndarray, + env_ids: np.ndarray, + *, + term_name: str, + ) -> None: + """Stage gravity values in the exactly-once reset transaction.""" + if self._reset_state is None: + raise NotImplementedError( + f"EventManager term '{term_name}' gravity capability is unavailable: " + "EntityScene was materialized without an env-owned reset transaction" + ) + self._reset_state.write_gravity(env_ids, values, term_name=term_name) + def bind_sensor_data(self, names: Sequence[str]) -> BackendSensorView: """Bind existing backend sensors for a manager term on the cold path. diff --git a/src/unilab/base/reset_state.py b/src/unilab/base/reset_state.py index 9c48d4520..bd730c710 100644 --- a/src/unilab/base/reset_state.py +++ b/src/unilab/base/reset_state.py @@ -13,7 +13,14 @@ 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.dr.types import ( + RESET_TERM_BODY_IPOS, + RESET_TERM_BODY_MASS, + RESET_TERM_GRAVITY, + RESET_TERM_KD, + RESET_TERM_KP, + ResetRandomizationPayload, +) from unilab.utils.rotation import np_quat_apply_inverse @@ -41,6 +48,9 @@ def __init__( 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._randomization_defaults: dict[str, np.ndarray] = {} + self._randomization_values: dict[str, np.ndarray] = {} + self._randomization_dirty_masks: dict[str, np.ndarray] = {} self._requesting_terms: set[str] = set() @property @@ -69,9 +79,127 @@ def begin(self, env_ids: np.ndarray) -> None: self._active_mask[ids] = True self._dirty_mask.fill(False) self._gain_dirty_mask.fill(False) + for mask in self._randomization_dirty_masks.values(): + mask.fill(False) self._requesting_terms.clear() self._active = True + def bind_body_mass_write( + self, + body_ids: np.ndarray, + *, + term_name: str, + ) -> tuple[np.ndarray, np.ndarray]: + """Bind body-mass columns and immutable backend defaults on the cold path.""" + default = self._materialize_randomization_default( + RESET_TERM_BODY_MASS, + getter=self._backend.get_body_mass, + expected_tail=None, + term_name=term_name, + ) + columns = self._validate_columns( + body_ids, + width=default.shape[0], + capability="body mass IDs", + term_name=term_name, + ) + return self._readonly_binding(columns, default[columns]) + + def bind_body_ipos_write( + self, + body_ids: np.ndarray, + *, + term_name: str, + ) -> tuple[np.ndarray, np.ndarray]: + """Bind body inertial-position columns and immutable backend defaults.""" + default = self._materialize_randomization_default( + RESET_TERM_BODY_IPOS, + getter=self._backend.get_body_ipos, + expected_tail=(3,), + term_name=term_name, + ) + columns = self._validate_columns( + body_ids, + width=default.shape[0], + capability="body ipos IDs", + term_name=term_name, + ) + return self._readonly_binding(columns, default[columns]) + + def bind_gravity_write(self, *, term_name: str) -> np.ndarray: + """Bind the immutable backend gravity vector on the cold path.""" + return self._materialize_randomization_default( + RESET_TERM_GRAVITY, + getter=self._backend.get_gravity, + expected_tail=(), + term_name=term_name, + ) + + def write_body_mass( + self, + env_ids: np.ndarray, + body_ids: np.ndarray, + values: np.ndarray, + *, + term_name: str, + ) -> None: + """Stage selected body masses in the exactly-once reset payload.""" + self._write_body_randomization( + RESET_TERM_BODY_MASS, + env_ids, + body_ids, + values, + value_tail=(), + term_name=term_name, + ) + + def write_body_ipos( + self, + env_ids: np.ndarray, + body_ids: np.ndarray, + values: np.ndarray, + *, + term_name: str, + ) -> None: + """Stage selected body inertial positions in the reset payload.""" + self._write_body_randomization( + RESET_TERM_BODY_IPOS, + env_ids, + body_ids, + values, + value_tail=(3,), + term_name=term_name, + ) + + def write_gravity( + self, + env_ids: np.ndarray, + values: np.ndarray, + *, + term_name: str, + ) -> None: + """Stage per-environment gravity vectors in the reset payload.""" + ids = self._prepare_state_write( + env_ids, + capability="gravity", + term_name=term_name, + ) + default = self._require_randomization_default(RESET_TERM_GRAVITY, term_name) + gravity = self._validate_values( + values, + shape=(ids.size, 3), + capability="gravity", + term_name=term_name, + ) + buffer = self._randomization_values[RESET_TERM_GRAVITY] + mask = self._randomization_dirty_masks[RESET_TERM_GRAVITY] + uninitialized = ids[~mask[ids]] + if uninitialized.size: + buffer[uninitialized] = default + buffer[ids] = gravity + mask[ids] = True + self._dirty_mask[ids] = True + def bind_actuator_gain_write( self, actuator_ids: np.ndarray, @@ -399,27 +527,170 @@ def _materialize_default_actuator_gains(self, term_name: str) -> None: dtype=default_kd.dtype, ) + def _materialize_randomization_default( + self, + field: str, + *, + getter, + expected_tail: tuple[int, ...] | None, + term_name: str, + ) -> np.ndarray: + cached = self._randomization_defaults.get(field) + if cached is not None: + return cached + try: + capabilities = self._backend.get_dr_capabilities() + except (AttributeError, NotImplementedError) as exc: + raise self._capability_error(term_name, f"{field} randomization", exc) from exc + unsupported = capabilities.get_unsupported_reset_terms(frozenset((field,))) + if unsupported: + raise self._capability_error( + term_name, + f"{field} randomization", + NotImplementedError(f"unsupported reset payload field: {field}"), + ) + try: + value = getter() + except (AttributeError, NotImplementedError) as exc: + raise self._capability_error(term_name, f"default {field}", exc) from exc + if not isinstance(value, np.ndarray): + raise TypeError( + f"EventManager term '{term_name}' capability 'default {field}' on backend " + f"'{self._backend.backend_type}' must return np.ndarray, got " + f"{type(value).__name__}" + ) + expected_ndim = 1 if expected_tail is None else 1 + len(expected_tail) + if value.ndim != expected_ndim: + raise ValueError( + f"EventManager term '{term_name}' capability 'default {field}' on backend " + f"'{self._backend.backend_type}' returned shape {value.shape}; expected " + f"{expected_ndim}-D" + ) + if expected_tail is not None and value.shape[1:] != expected_tail: + raise ValueError( + f"EventManager term '{term_name}' capability 'default {field}' on backend " + f"'{self._backend.backend_type}' returned shape {value.shape}; expected tail " + f"{expected_tail}" + ) + if not np.issubdtype(value.dtype, np.floating): + raise TypeError( + f"EventManager term '{term_name}' capability 'default {field}' on backend " + f"'{self._backend.backend_type}' must be floating, got {value.dtype}" + ) + if not np.isfinite(value).all(): + raise ValueError( + f"EventManager term '{term_name}' capability 'default {field}' on backend " + f"'{self._backend.backend_type}' returned NaN or Inf" + ) + default = np.array(value, copy=True) + default.setflags(write=False) + self._randomization_defaults[field] = default + self._randomization_values[field] = np.empty( + (self._num_envs, *default.shape), + dtype=default.dtype, + ) + self._randomization_dirty_masks[field] = np.zeros(self._num_envs, dtype=np.bool_) + return default + + def _require_randomization_default(self, field: str, term_name: str) -> np.ndarray: + try: + return self._randomization_defaults[field] + except KeyError as exc: + raise RuntimeError( + f"EventManager term '{term_name}' must bind reset field '{field}' " + "during manager construction before writing it" + ) from exc + + def _readonly_binding( + self, + columns: np.ndarray, + selected_default: np.ndarray, + ) -> tuple[np.ndarray, np.ndarray]: + bound_columns = np.array(columns, copy=True) + bound_columns.setflags(write=False) + selected = np.array(selected_default, copy=True) + selected.setflags(write=False) + return bound_columns, selected + + def _write_body_randomization( + self, + field: str, + env_ids: np.ndarray, + body_ids: np.ndarray, + values: np.ndarray, + *, + value_tail: tuple[int, ...], + term_name: str, + ) -> None: + ids = self._prepare_state_write( + env_ids, + capability=field, + term_name=term_name, + ) + default = self._require_randomization_default(field, term_name) + columns = self._validate_columns( + body_ids, + width=default.shape[0], + capability=f"{field} body IDs", + term_name=term_name, + ) + selected = self._validate_values( + values, + shape=(ids.size, columns.size, *value_tail), + capability=field, + term_name=term_name, + ) + buffer = self._randomization_values[field] + mask = self._randomization_dirty_masks[field] + uninitialized = ids[~mask[ids]] + if uninitialized.size: + buffer[uninitialized] = default + if ids.size and columns.size: + buffer[ids[:, None], columns[None, :]] = selected + mask[ids] = True + self._dirty_mask[ids] = True + def _build_randomization_payload( self, dirty_ids: np.ndarray, ) -> ResetRandomizationPayload | None: + payload = ResetRandomizationPayload() + for field in (RESET_TERM_BODY_MASS, RESET_TERM_BODY_IPOS, RESET_TERM_GRAVITY): + mask = self._randomization_dirty_masks.get(field) + if mask is None or not np.any(mask): + continue + self._require_dense_randomization_rows(field, mask, dirty_ids) + setattr( + payload, + field, + np.array(self._randomization_values[field][dirty_ids], copy=True), + ) + 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 gain_ids.size: + self._require_dense_randomization_rows( + "actuator gains", self._gain_dirty_mask, dirty_ids + ) + assert self._kp is not None + assert self._kd is not None + payload.kp = np.array(self._kp[dirty_ids], copy=True) + payload.kd = np.array(self._kd[dirty_ids], copy=True) + return None if payload.is_empty() else payload + + def _require_dense_randomization_rows( + self, + field: str, + mask: np.ndarray, + dirty_ids: np.ndarray, + ) -> None: + missing = dirty_ids[~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"EventManager reset {field} payload cannot represent sparse rows in one " + f"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, @@ -595,7 +866,7 @@ def _validate_values( self, values: np.ndarray, *, - shape: tuple[int, int], + shape: tuple[int, ...], capability: str, term_name: str, ) -> np.ndarray: @@ -637,6 +908,8 @@ def _finish(self) -> None: self._active_mask.fill(False) self._dirty_mask.fill(False) self._gain_dirty_mask.fill(False) + for mask in self._randomization_dirty_masks.values(): + mask.fill(False) self._requesting_terms.clear() diff --git a/src/unilab/envs/mdp/__init__.py b/src/unilab/envs/mdp/__init__.py index 7e61bd0be..f3627001b 100644 --- a/src/unilab/envs/mdp/__init__.py +++ b/src/unilab/envs/mdp/__init__.py @@ -5,6 +5,12 @@ 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 push_by_setting_velocity as push_by_setting_velocity +from unilab.envs.mdp.events import ( + randomize_physics_scene_gravity as randomize_physics_scene_gravity, +) +from unilab.envs.mdp.events import randomize_rigid_body_com as randomize_rigid_body_com +from unilab.envs.mdp.events import randomize_rigid_body_mass as randomize_rigid_body_mass 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 @@ -55,6 +61,10 @@ "joint_vel_l2", "last_action", "pd_gains", + "push_by_setting_velocity", + "randomize_physics_scene_gravity", + "randomize_rigid_body_com", + "randomize_rigid_body_mass", "is_alive", "is_terminated", "projected_gravity", diff --git a/src/unilab/envs/mdp/events.py b/src/unilab/envs/mdp/events.py index d86a382e4..e370fa159 100644 --- a/src/unilab/envs/mdp/events.py +++ b/src/unilab/envs/mdp/events.py @@ -2,6 +2,8 @@ # src/mjlab/envs/mdp/events.py. # Copyright 2025, The mjlab Developers. # Modified by UniLab for NumPy reset transactions; Apache-2.0. +# The mass/CoM/gravity public names and signatures follow Isaac Lab v2.2.0; +# their implementation here is UniLab's original payload adapter, not vendored PhysX code. """Community-style reset event terms for UniLab's NumPy manager runtime.""" from __future__ import annotations @@ -22,7 +24,10 @@ _DEFAULT_ASSET_CFG = SceneEntityCfg("robot") _SE3_KEYS = ("x", "y", "z", "roll", "pitch", "yaw") +_XYZ_KEYS = ("x", "y", "z") _PD_GAIN_PARAM_NAMES = frozenset(("kp_range", "kd_range", "asset_cfg", "distribution", "operation")) +_DISTRIBUTIONS = ("uniform", "log_uniform", "gaussian") +_OPERATIONS = ("add", "scale", "abs") def _gain_range( @@ -104,6 +109,139 @@ def _sample_se3_range( return rng.uniform(ranges[:, 0], ranges[:, 1], size=shape) +def _event_choice(value: Any, *, term_name: str, name: str, choices: tuple[str, ...]) -> str: + if not isinstance(value, str): + raise TypeError( + f"EventManager term '{term_name}' parameter '{name}' must be a string, " + f"got {type(value).__name__}" + ) + if value not in choices: + raise ValueError( + f"EventManager term '{term_name}' parameter '{name}' must be one of " + f"{choices}, got {value!r}" + ) + return value + + +def _distribution_parameters( + value: Any, + *, + term_name: str, + name: str, + width: int | None = None, + distribution: str, +) -> np.ndarray: + try: + params = np.asarray(value, dtype=np.float64) + except (TypeError, ValueError) as exc: + raise TypeError( + f"EventManager term '{term_name}' parameter '{name}' must contain numeric values" + ) from exc + expected = (2,) if width is None else (2, width) + if params.shape != expected: + raise ValueError( + f"EventManager term '{term_name}' parameter '{name}' has shape {params.shape}; " + f"expected {expected}" + ) + if not np.isfinite(params).all(): + raise ValueError( + f"EventManager term '{term_name}' parameter '{name}' must contain only finite values" + ) + if distribution == "gaussian": + if np.any(params[1] < 0.0): + raise ValueError( + f"EventManager term '{term_name}' parameter '{name}' standard deviation " + "must be non-negative" + ) + else: + if np.any(params[0] > params[1]): + raise ValueError( + f"EventManager term '{term_name}' parameter '{name}' lower bound exceeds upper bound" + ) + if distribution == "log_uniform" and np.any(params <= 0.0): + raise ValueError( + f"EventManager term '{term_name}' parameter '{name}' must be positive " + "for log_uniform sampling" + ) + result = np.array(params, copy=True) + result.setflags(write=False) + return result + + +def _sample_distribution( + rng: np.random.Generator, + params: np.ndarray, + shape: tuple[int, ...], + distribution: str, +) -> np.ndarray: + if distribution == "gaussian": + return rng.normal(params[0], params[1], size=shape) + if distribution == "log_uniform": + return np.exp(rng.uniform(np.log(params[0]), np.log(params[1]), size=shape)) + return rng.uniform(params[0], params[1], size=shape) + + +def _apply_randomization_operation( + default: np.ndarray, + samples: np.ndarray, + operation: str, +) -> np.ndarray: + if operation == "add": + return default + samples + if operation == "scale": + return default * samples + return samples + + +def _axis_ranges( + value: Any, + *, + term_name: str, + name: str, + keys: tuple[str, ...], +) -> np.ndarray: + if not isinstance(value, dict): + raise TypeError(f"EventManager term '{term_name}' parameter '{name}' must be a dict") + unknown = sorted(set(value) - set(keys)) + if unknown: + raise ValueError( + f"EventManager term '{term_name}' parameter '{name}' has unknown axes {unknown}" + ) + parameters = np.asarray([value.get(key, (0.0, 0.0)) for key in keys], dtype=object).T + return _distribution_parameters( + parameters, + term_name=term_name, + name=name, + width=len(keys), + distribution="uniform", + ).T + + +def _validate_event_term( + cfg: EventTermCfg, + *, + term_name: str, + mode: str, + allowed_params: frozenset[str], + required_params: tuple[str, ...], +) -> None: + if cfg.mode != mode: + raise NotImplementedError( + f"EventManager term '{term_name}' only supports mode='{mode}' on the UniLab runtime" + ) + if mode == "reset" and cfg.min_step_count_between_reset != 0: + raise NotImplementedError( + f"EventManager term '{term_name}' requires min_step_count_between_reset=0 " + "because sparse reset payload rows cannot be represented" + ) + unknown = sorted(set(cfg.params) - allowed_params) + if unknown: + raise ValueError(f"EventManager term '{term_name}' has unknown parameters {unknown}") + missing = [name for name in required_params if name not in cfg.params] + if missing: + raise ValueError(f"EventManager term '{term_name}' is missing parameters {missing}") + + def resolve_env_ids(env: ManagerBasedRlEnv, env_ids: np.ndarray | None) -> np.ndarray: """Return concrete NumPy environment IDs, preserving community sentinel semantics.""" if env_ids is None: @@ -205,6 +343,300 @@ def __call__( pd_gains = PdGains +class RandomizeRigidBodyMass(ManagerTermBase): + """Community-compatible body-mass randomization via the reset payload.""" + + _PARAMS = frozenset( + ( + "asset_cfg", + "mass_distribution_params", + "operation", + "distribution", + "recompute_inertia", + "min_mass", + ) + ) + + def __init__(self, cfg: EventTermCfg, env: ManagerBasedRlEnv): + super().__init__(env) + term_name = "randomize_rigid_body_mass" + _validate_event_term( + cfg, + term_name=term_name, + mode="reset", + allowed_params=self._PARAMS, + required_params=("asset_cfg", "mass_distribution_params", "operation"), + ) + asset_cfg = cfg.params["asset_cfg"] + if not isinstance(asset_cfg, SceneEntityCfg): + raise TypeError( + f"EventManager term '{term_name}' asset_cfg must be SceneEntityCfg, " + f"got {type(asset_cfg).__name__}" + ) + recompute_inertia = cfg.params.get("recompute_inertia", True) + if not isinstance(recompute_inertia, bool): + raise TypeError(f"EventManager term '{term_name}' recompute_inertia must be bool") + if recompute_inertia: + raise NotImplementedError( + f"EventManager term '{term_name}' recompute_inertia=True is unavailable: " + "ResetRandomizationPayload has no inertia-recomputation contract; set it " + "explicitly to false or do not configure this term" + ) + self._operation = _event_choice( + cfg.params["operation"], + term_name=term_name, + name="operation", + choices=_OPERATIONS, + ) + self._distribution = _event_choice( + cfg.params.get("distribution", "uniform"), + term_name=term_name, + name="distribution", + choices=_DISTRIBUTIONS, + ) + self._distribution_params = _distribution_parameters( + cfg.params["mass_distribution_params"], + term_name=term_name, + name="mass_distribution_params", + distribution=self._distribution, + ) + min_mass = cfg.params.get("min_mass", 1e-6) + if isinstance(min_mass, bool) or not isinstance(min_mass, (int, float)): + raise TypeError(f"EventManager term '{term_name}' min_mass must be numeric") + self._min_mass = float(min_mass) + if not np.isfinite(self._min_mass) or self._min_mass < 1e-6: + raise ValueError( + f"EventManager term '{term_name}' min_mass must be finite and at least 1e-6" + ) + self._entity = cast("Entity", env.scene[asset_cfg.name]) + self._body_ids, self._default_mass = self._entity.bind_body_mass_write( + asset_cfg.body_ids, + term_name=term_name, + ) + + def __call__( + self, + env: ManagerBasedRlEnv, + env_ids: np.ndarray | None, + asset_cfg: SceneEntityCfg, + mass_distribution_params: tuple[float, float], + operation: Literal["add", "scale", "abs"], + distribution: Literal["uniform", "log_uniform", "gaussian"] = "uniform", + recompute_inertia: bool = True, + min_mass: float = 1e-6, + ) -> None: + del ( + asset_cfg, + mass_distribution_params, + operation, + distribution, + recompute_inertia, + min_mass, + ) + ids = resolve_env_ids(env, env_ids) + samples = _sample_distribution( + env.rng, + self._distribution_params, + (ids.size, self._body_ids.size), + self._distribution, + ) + values = _apply_randomization_operation( + self._default_mass[None, :], + samples, + self._operation, + ) + np.maximum(values, self._min_mass, out=values) + self._entity.write_body_mass_to_sim( + values, + body_ids=self._body_ids, + env_ids=ids, + term_name="randomize_rigid_body_mass", + ) + + +randomize_rigid_body_mass = RandomizeRigidBodyMass + + +class RandomizeRigidBodyCom(ManagerTermBase): + """Community-compatible additive rigid-body CoM randomization.""" + + def __init__(self, cfg: EventTermCfg, env: ManagerBasedRlEnv): + super().__init__(env) + term_name = "randomize_rigid_body_com" + _validate_event_term( + cfg, + term_name=term_name, + mode="reset", + allowed_params=frozenset(("com_range", "asset_cfg")), + required_params=("com_range", "asset_cfg"), + ) + asset_cfg = cfg.params["asset_cfg"] + if not isinstance(asset_cfg, SceneEntityCfg): + raise TypeError( + f"EventManager term '{term_name}' asset_cfg must be SceneEntityCfg, " + f"got {type(asset_cfg).__name__}" + ) + self._ranges = _axis_ranges( + cfg.params["com_range"], + term_name=term_name, + name="com_range", + keys=_XYZ_KEYS, + ) + self._entity = cast("Entity", env.scene[asset_cfg.name]) + self._body_ids, self._default_ipos = self._entity.bind_body_ipos_write( + asset_cfg.body_ids, + term_name=term_name, + ) + + def __call__( + self, + env: ManagerBasedRlEnv, + env_ids: np.ndarray | None, + com_range: dict[str, tuple[float, float]], + asset_cfg: SceneEntityCfg, + ) -> None: + del com_range, asset_cfg + ids = resolve_env_ids(env, env_ids) + offsets = env.rng.uniform( + self._ranges[:, 0], + self._ranges[:, 1], + size=(ids.size, 3), + ) + values = self._default_ipos[None, :, :] + offsets[:, None, :] + self._entity.write_body_ipos_to_sim( + values, + body_ids=self._body_ids, + env_ids=ids, + term_name="randomize_rigid_body_com", + ) + + +randomize_rigid_body_com = RandomizeRigidBodyCom + + +class RandomizePhysicsSceneGravity(ManagerTermBase): + """Community-compatible gravity randomization through reset transactions.""" + + def __init__(self, cfg: EventTermCfg, env: ManagerBasedRlEnv): + super().__init__(env) + term_name = "randomize_physics_scene_gravity" + _validate_event_term( + cfg, + term_name=term_name, + mode="reset", + allowed_params=frozenset(("gravity_distribution_params", "operation", "distribution")), + required_params=("gravity_distribution_params", "operation"), + ) + self._operation = _event_choice( + cfg.params["operation"], + term_name=term_name, + name="operation", + choices=_OPERATIONS, + ) + self._distribution = _event_choice( + cfg.params.get("distribution", "uniform"), + term_name=term_name, + name="distribution", + choices=_DISTRIBUTIONS, + ) + self._distribution_params = _distribution_parameters( + cfg.params["gravity_distribution_params"], + term_name=term_name, + name="gravity_distribution_params", + width=3, + distribution=self._distribution, + ) + self._default_gravity = env.scene.bind_gravity_write(term_name=term_name) + + def __call__( + self, + env: ManagerBasedRlEnv, + env_ids: np.ndarray | None, + gravity_distribution_params: tuple[list[float], list[float]], + operation: Literal["add", "scale", "abs"], + distribution: Literal["uniform", "log_uniform", "gaussian"] = "uniform", + ) -> None: + del gravity_distribution_params, operation, distribution + ids = resolve_env_ids(env, env_ids) + samples = _sample_distribution( + env.rng, + self._distribution_params, + (ids.size, 3), + self._distribution, + ) + values = _apply_randomization_operation( + self._default_gravity[None, :], + samples, + self._operation, + ) + env.scene.write_gravity_to_sim( + values, + ids, + term_name="randomize_physics_scene_gravity", + ) + + +randomize_physics_scene_gravity = RandomizePhysicsSceneGravity + + +class PushBySettingVelocity(ManagerTermBase): + """Pinned community velocity kick dispatched through the interval plan.""" + + def __init__(self, cfg: EventTermCfg, env: ManagerBasedRlEnv): + super().__init__(env) + term_name = "push_by_setting_velocity" + _validate_event_term( + cfg, + term_name=term_name, + mode="interval", + allowed_params=frozenset(("velocity_range", "asset_cfg")), + required_params=("velocity_range",), + ) + asset_cfg = cfg.params.get("asset_cfg", _DEFAULT_ASSET_CFG) + if not isinstance(asset_cfg, SceneEntityCfg): + raise TypeError( + f"EventManager term '{term_name}' asset_cfg must be SceneEntityCfg, " + f"got {type(asset_cfg).__name__}" + ) + ranges = _axis_ranges( + cfg.params["velocity_range"], + term_name=term_name, + name="velocity_range", + keys=_SE3_KEYS, + ) + if np.any(ranges[3:] != 0.0): + raise NotImplementedError( + f"EventManager term '{term_name}' angular velocity ranges are unsupported: " + "IntervalRandomizationPlan only declares body_linear_velocity_delta" + ) + self._ranges = ranges[:3] + self._entity = cast("Entity", env.scene[asset_cfg.name]) + self._entity.bind_root_linear_velocity_delta(term_name=term_name) + + def __call__( + self, + env: ManagerBasedRlEnv, + env_ids: np.ndarray | None, + velocity_range: dict[str, tuple[float, float]], + asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG, + ) -> None: + del velocity_range, asset_cfg + ids = resolve_env_ids(env, env_ids) + delta = env.rng.uniform( + self._ranges[:, 0], + self._ranges[:, 1], + size=(ids.size, 3), + ) + self._entity.apply_root_linear_velocity_delta_to_sim( + delta, + env_ids=ids, + term_name="push_by_setting_velocity", + ) + + +push_by_setting_velocity = PushBySettingVelocity + + 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) @@ -249,4 +681,13 @@ def reset_root_state_uniform( asset.write_root_state_to_sim(root_states, env_ids=ids) -__all__ = ["pd_gains", "reset_root_state_uniform", "reset_scene_to_default", "resolve_env_ids"] +__all__ = [ + "pd_gains", + "push_by_setting_velocity", + "randomize_physics_scene_gravity", + "randomize_rigid_body_com", + "randomize_rigid_body_mass", + "reset_root_state_uniform", + "reset_scene_to_default", + "resolve_env_ids", +] diff --git a/src/unilab/managers/_types.py b/src/unilab/managers/_types.py index df3af5c0c..6971fcf6c 100644 --- a/src/unilab/managers/_types.py +++ b/src/unilab/managers/_types.py @@ -168,6 +168,12 @@ def bind_sensor_data(self, names: Sequence[str]) -> ManagerSensorView: ... def reset_to_default(self, env_ids: np.ndarray, *, term_name: str) -> None: ... + def bind_gravity_write(self, *, term_name: str) -> np.ndarray: ... + + def write_gravity_to_sim( + self, values: np.ndarray, env_ids: np.ndarray, *, term_name: str + ) -> None: ... + class ManagerActionTerm(Protocol): @property diff --git a/tests/envs/locomotion/go2/test_manager_based_cfg.py b/tests/envs/locomotion/go2/test_manager_based_cfg.py index 77af704ed..221ea5c6c 100644 --- a/tests/envs/locomotion/go2/test_manager_based_cfg.py +++ b/tests/envs/locomotion/go2/test_manager_based_cfg.py @@ -3,12 +3,13 @@ from __future__ import annotations import math -from typing import TypeVar, cast +from typing import Any, TypeVar, cast import numpy as np import pytest from unilab.base.backend import create_backend, env_backend_kwargs +from unilab.base.entity import EntityCfg from unilab.base.np_env import NpEnvState from unilab.envs import ManagerBasedRlEnv, mdp from unilab.envs.locomotion.common import manager_terms @@ -16,6 +17,7 @@ make_go2_joystick_flat_manager_cfg, ) from unilab.envs.mdp import JointPositionAction, JointPositionActionCfg +from unilab.managers import EventTermCfg, SceneEntityCfg _JOINT_NAMES = ( "FL_hip_joint", @@ -304,6 +306,72 @@ def test_go2_manager_factory_executes_on_real_mujoco() -> None: env.close() +def test_go2_manager_reset_randomization_mutates_real_mujoco_payload() -> None: + cfg = make_go2_joystick_flat_manager_cfg() + assert cfg.scene is not None + robot = cfg.scene.entities["robot"] + cfg.scene.entities["robot"] = EntityCfg( + root_body_name=robot.root_body_name, + joint_names=robot.joint_names, + body_names=("base",), + actuator_names=robot.actuator_names, + ) + asset_cfg = SceneEntityCfg("robot", body_names=("base",)) + cfg.events.update( + { + "mass": EventTermCfg( + func=mdp.randomize_rigid_body_mass, + mode="reset", + params={ + "asset_cfg": asset_cfg, + "mass_distribution_params": (1.25, 1.25), + "operation": "scale", + "recompute_inertia": False, + }, + ), + "com": EventTermCfg( + func=mdp.randomize_rigid_body_com, + mode="reset", + params={"asset_cfg": asset_cfg, "com_range": {"x": (0.02, 0.02)}}, + ), + "gravity": EventTermCfg( + func=mdp.randomize_physics_scene_gravity, + mode="reset", + params={ + "gravity_distribution_params": ([0.0, 0.0, -9.7],) * 2, + "operation": "abs", + }, + ), + } + ) + backend = create_backend( + "mujoco", + cfg.scene, + 2, + cfg.sim_dt, + base_name="base", + add_body_sensors=True, + **env_backend_kwargs(cfg), + ) + base_id = int(backend.get_body_ids(("base",))[0]) + default_mass = backend.get_body_mass() + default_ipos = backend.get_body_ipos() + env = ManagerBasedRlEnv(cfg, backend, 2) + try: + env.reset(seed=31) + pool = cast(Any, backend)._pool + assert pool is not None + for env_id in range(2): + mass = pool.get_field(env_id, "body_mass") + ipos = pool.get_field(env_id, "body_ipos").reshape(-1, 3) + gravity = pool.get_field(env_id, "gravity") + assert mass[base_id] == pytest.approx(default_mass[base_id] * 1.25) + np.testing.assert_allclose(ipos[base_id], default_ipos[base_id] + [0.02, 0.0, 0.0]) + np.testing.assert_allclose(gravity, [0.0, 0.0, -9.7]) + 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 diff --git a/tests/envs/mdp/test_events.py b/tests/envs/mdp/test_events.py index e73ace400..8f2000276 100644 --- a/tests/envs/mdp/test_events.py +++ b/tests/envs/mdp/test_events.py @@ -14,9 +14,13 @@ from unilab.base.entity import EntityCfg, EntityScene from unilab.base.reset_state import ResetStateTransaction from unilab.dr.types import ( + RESET_TERM_BODY_IPOS, + RESET_TERM_BODY_MASS, + RESET_TERM_GRAVITY, RESET_TERM_KD, RESET_TERM_KP, DomainRandomizationCapabilities, + IntervalRandomizationPlan, ResetRandomizationPayload, ) from unilab.envs import mdp @@ -160,9 +164,13 @@ def __init__( *, root_layout_supported: bool = True, gain_supported: bool = True, + randomization_supported: bool = True, + interval_velocity_supported: bool = True, ) -> None: self.root_layout_supported = root_layout_supported self.gain_supported = gain_supported + self.randomization_supported = randomization_supported + self.interval_velocity_supported = interval_velocity_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]] = [] @@ -171,6 +179,10 @@ def __init__( self.body_quat = np.zeros((self.num_envs, 1, 4)) self.body_quat[:, :, 0] = 1.0 self.body_velocity = np.zeros((self.num_envs, 1, 3)) + self.body_mass = np.array([10.0]) + self.body_ipos = np.array([[0.0, 0.0, 0.0]]) + self.gravity = np.array([0.0, 0.0, -9.81]) + self.interval_plans: list[IntervalRandomizationPlan] = [] def get_body_ids(self, names) -> np.ndarray: if tuple(names) != ("base",): @@ -203,12 +215,29 @@ 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) + terms: set[str] = set((RESET_TERM_KP, RESET_TERM_KD)) if self.gain_supported else set() + if self.randomization_supported: + terms.update((RESET_TERM_BODY_MASS, RESET_TERM_BODY_IPOS, RESET_TERM_GRAVITY)) + return DomainRandomizationCapabilities( + supported_reset_terms=frozenset(terms), + supports_interval_body_velocity_delta=self.interval_velocity_supported, + ) 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_mass(self) -> np.ndarray: + return self.body_mass.copy() + + def get_body_ipos(self) -> np.ndarray: + return self.body_ipos.copy() + + def get_gravity(self) -> np.ndarray: + return self.gravity.copy() + + def apply_interval_randomization(self, plan: IntervalRandomizationPlan) -> None: + self.interval_plans.append(plan) + def get_body_pos_w(self, ids: np.ndarray) -> np.ndarray: return self.body_pos[:, ids] @@ -239,15 +268,29 @@ def set_state( def _transaction_env( - *, root_layout_supported: bool = True, gain_supported: bool = True, rng_seed: int = 5 + *, + root_layout_supported: bool = True, + gain_supported: bool = True, + randomization_supported: bool = True, + interval_velocity_supported: bool = True, + body_names: tuple[str, ...] | None = ("base",), + rng_seed: int = 5, ) -> tuple[ManagerBasedRlEnv, _Backend, ResetStateTransaction]: backend = _Backend( root_layout_supported=root_layout_supported, gain_supported=gain_supported, + randomization_supported=randomization_supported, + interval_velocity_supported=interval_velocity_supported, ) transaction = ResetStateTransaction(cast(SimBackend, backend)) scene = EntityScene( - {"robot": EntityCfg(root_body_name="base", actuator_names=("a0", "a1", "a2"))}, + { + "robot": EntityCfg( + root_body_name="base", + body_names=body_names, + actuator_names=("a0", "a1", "a2"), + ) + }, cast(SimBackend, backend), reset_state=transaction, ) @@ -387,6 +430,192 @@ def test_pd_gains_missing_backend_capability_fails_during_manager_construction() assert backend.set_state_calls == [] +def test_reset_randomization_terms_compose_with_state_and_gains_exactly_once() -> None: + env, backend, transaction = _transaction_env(rng_seed=13) + asset_cfg = SceneEntityCfg("robot", body_names=("base",)) + manager = EventManager( + { + "mass": EventTermCfg( + func=mdp.randomize_rigid_body_mass, + mode="reset", + params={ + "asset_cfg": asset_cfg, + "mass_distribution_params": (1.5, 1.5), + "operation": "scale", + "recompute_inertia": False, + }, + ), + "com": EventTermCfg( + func=mdp.randomize_rigid_body_com, + mode="reset", + params={ + "asset_cfg": asset_cfg, + "com_range": {"x": (0.1, 0.1), "z": (-0.2, -0.2)}, + }, + ), + "gravity": EventTermCfg( + func=mdp.randomize_physics_scene_gravity, + mode="reset", + params={ + "gravity_distribution_params": ( + [0.0, 0.0, -10.0], + [0.0, 0.0, -10.0], + ), + "operation": "abs", + }, + ), + "gains": EventTermCfg( + func=mdp.pd_gains, + mode="reset", + params={"kp_range": (2.0, 2.0), "kd_range": (3.0, 3.0)}, + ), + }, + env, + ) + ids = np.array([0, 2], dtype=np.int32) + + with transaction.scoped(ids): + mdp.reset_scene_to_default(env, 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 + assert payload.body_mass is not None + assert payload.body_ipos is not None + assert payload.gravity is not None + assert payload.kp is not None + assert payload.kd is not None + np.testing.assert_allclose(payload.body_mass, [[15.0], [15.0]]) + np.testing.assert_allclose(payload.body_ipos, [[[0.1, 0.0, -0.2]]] * 2) + np.testing.assert_allclose(payload.gravity, [[0.0, 0.0, -10.0]] * 2) + np.testing.assert_allclose(payload.kp, [[20.0, 40.0, 60.0]] * 2) + np.testing.assert_allclose(payload.kd, [[3.0, 6.0, 9.0]] * 2) + + +@pytest.mark.parametrize( + ("func", "params", "match"), + [ + ( + mdp.randomize_rigid_body_mass, + { + "asset_cfg": SceneEntityCfg("robot", body_names=("base",)), + "mass_distribution_params": (0.9, 1.1), + "operation": "scale", + }, + "recompute_inertia=True", + ), + ( + mdp.randomize_rigid_body_mass, + { + "asset_cfg": SceneEntityCfg("robot", body_names=("base",)), + "mass_distribution_params": (0.9, 1.1), + "operation": "scale", + "recompute_inertia": False, + }, + "body_mass randomization.*unsupported", + ), + ], +) +def test_mass_randomization_capability_gaps_fail_during_construction( + func, params: dict[str, Any], match: str +) -> None: + env, backend, _ = _transaction_env(randomization_supported=False) + with pytest.raises(NotImplementedError, match=match): + EventManager( + {"mass": EventTermCfg(func=func, mode="reset", params=params)}, + env, + ) + assert backend.set_state_calls == [] + + +def test_reset_randomization_sparse_rows_abort_without_backend_mutation() -> None: + env, backend, transaction = _transaction_env() + manager = EventManager( + { + "gravity": EventTermCfg( + func=mdp.randomize_physics_scene_gravity, + mode="reset", + params={ + "gravity_distribution_params": ([0.0, 0.0, -10.0],) * 2, + "operation": "abs", + }, + ) + }, + env, + ) + cfg = manager.get_term_cfg("gravity") + ids = np.array([0, 1], dtype=np.int32) + + with pytest.raises(RuntimeError, match=r"gravity payload.*sparse rows.*missing env IDs \[1\]"): + with transaction.scoped(ids): + mdp.reset_scene_to_default(env, ids) + cfg.func(env, np.array([0], dtype=np.int32), **cfg.params) + assert backend.set_state_calls == [] + + +def test_velocity_push_uses_env_rng_and_interval_subset_plan() -> None: + env, backend, _ = _transaction_env(rng_seed=19) + manager = EventManager( + { + "push": EventTermCfg( + func=mdp.push_by_setting_velocity, + mode="interval", + interval_range_s=(1.0, 1.0), + params={ + "velocity_range": { + "x": (0.2, 0.2), + "y": (-0.3, -0.3), + "z": (0.4, 0.4), + } + }, + ) + }, + env, + ) + manager._interval_term_time_left[0][:] = [0.0, 1.0, 0.0] + + manager.apply(mode="interval", dt=0.1) + + assert len(backend.interval_plans) == 1 + plan = backend.interval_plans[0] + np.testing.assert_array_equal(plan.body_ids, [0]) + assert plan.body_linear_velocity_delta is not None + np.testing.assert_allclose( + plan.body_linear_velocity_delta[:, 0], + [[0.2, -0.3, 0.4], [0.0, 0.0, 0.0], [0.2, -0.3, 0.4]], + ) + + +@pytest.mark.parametrize( + ("supported", "velocity_range", "match"), + [ + (False, {"x": (-0.1, 0.1)}, "unsupported backend capability"), + (True, {"yaw": (-0.1, 0.1)}, "angular velocity ranges are unsupported"), + ], +) +def test_velocity_push_capability_gaps_fail_during_construction( + supported: bool, + velocity_range: dict[str, tuple[float, float]], + match: str, +) -> None: + env, backend, _ = _transaction_env(interval_velocity_supported=supported) + with pytest.raises(NotImplementedError, match=match): + EventManager( + { + "push": EventTermCfg( + func=mdp.push_by_setting_velocity, + mode="interval", + interval_range_s=(1.0, 1.0), + params={"velocity_range": velocity_range}, + ) + }, + env, + ) + assert backend.interval_plans == [] + + def test_uniform_root_state_fixed_or_mocap_capability_fails_closed() -> None: env, backend, transaction = _transaction_env(root_layout_supported=False) From 66e59a973c02f289981d8f2a28e04b42c2e1f25f Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 18:53:54 +0800 Subject: [PATCH 029/169] feat(env): add generic Manager-Based factory [skip ci] (#1103) --- src/unilab/envs/__init__.py | 2 + src/unilab/envs/manager_based_rl_env.py | 85 +++++++- tests/envs/test_manager_based_rl_env.py | 245 +++++++++++++++++++++++- 3 files changed, 320 insertions(+), 12 deletions(-) diff --git a/src/unilab/envs/__init__.py b/src/unilab/envs/__init__.py index 4542de1a5..0b2d6a246 100644 --- a/src/unilab/envs/__init__.py +++ b/src/unilab/envs/__init__.py @@ -4,10 +4,12 @@ from unilab.envs.manager_based_rl_env import ManagerBasedRlEnv as ManagerBasedRlEnv from unilab.envs.manager_based_rl_env import ManagerBasedRLEnvCfg as ManagerBasedRLEnvCfg from unilab.envs.manager_based_rl_env import ManagerBasedRlEnvCfg as ManagerBasedRlEnvCfg +from unilab.envs.manager_based_rl_env import make_manager_based_rl_env as make_manager_based_rl_env __all__ = [ "ManagerBasedRLEnv", "ManagerBasedRLEnvCfg", "ManagerBasedRlEnv", "ManagerBasedRlEnvCfg", + "make_manager_based_rl_env", ] diff --git a/src/unilab/envs/manager_based_rl_env.py b/src/unilab/envs/manager_based_rl_env.py index a99552231..eb2b5af10 100644 --- a/src/unilab/envs/manager_based_rl_env.py +++ b/src/unilab/envs/manager_based_rl_env.py @@ -14,13 +14,13 @@ import gymnasium as gym import numpy as np -from unilab.base.backend import SimBackend +from unilab.base.backend import SimBackend, create_backend, env_backend_kwargs from unilab.base.base import EnvCfg from unilab.base.config_overrides import ( CONFIG_MAPPING_POLICY_KEY, MANAGER_TERM_MAPPING_POLICY, ) -from unilab.base.entity import EntityScene +from unilab.base.entity import EntityCfg, EntityScene from unilab.base.np_env import NpEnv, NpEnvState from unilab.base.reset_state import ResetStateTransaction from unilab.base.scene import SceneCfg, resolve_scene_default_qpos @@ -146,6 +146,43 @@ def validate(self) -> None: ) +def _resolve_backend_entity_contract(cfg: ManagerBasedRlEnvCfg) -> tuple[str, bool]: + """Resolve task-independent backend inputs from declared scene entities.""" + assert cfg.scene is not None + root_entities: list[tuple[str, str]] = [] + body_state_requested = False + for entity_name, entity_cfg in cfg.scene.entities.items(): + if not isinstance(entity_name, str) or not entity_name: + raise TypeError( + f"ManagerBasedRlEnv scene entity names must be non-empty strings; " + f"got {entity_name!r}" + ) + if not isinstance(entity_cfg, EntityCfg): + raise TypeError( + f"ManagerBasedRlEnv scene entity '{entity_name}' must be EntityCfg, " + f"got {type(entity_cfg).__name__}" + ) + root_body_name = entity_cfg.root_body_name + if root_body_name is not None: + if not isinstance(root_body_name, str) or not root_body_name: + raise TypeError( + f"ManagerBasedRlEnv root entity '{entity_name}' root_body_name must be " + "a non-empty string" + ) + root_entities.append((entity_name, root_body_name)) + body_state_requested = True + if entity_cfg.body_names is not None: + body_state_requested = True + + if len(root_entities) != 1: + declared = [name for name, _ in root_entities] + raise ValueError( + "ManagerBasedRlEnv factory requires exactly one scene entity with an explicit " + f"root_body_name; found {len(root_entities)} root entities {declared}" + ) + return root_entities[0][1], body_state_requested + + class ManagerBasedRlEnv(NpEnv): """Manager-Based API adapter that reuses the single :class:`NpEnv` lifecycle.""" @@ -591,6 +628,49 @@ def close(self) -> None: super().close() +def make_manager_based_rl_env( + cfg: ManagerBasedRlEnvCfg, + num_envs: int = 1, + backend_type: str = "mujoco", +) -> ManagerBasedRlEnv: + """Construct the generic Registry-owned Manager-Based production runtime.""" + if not isinstance(cfg, ManagerBasedRlEnvCfg): + raise TypeError( + "make_manager_based_rl_env expected ManagerBasedRlEnvCfg, " + f"received {type(cfg).__name__}" + ) + if isinstance(num_envs, bool) or not isinstance(num_envs, int) or num_envs <= 0: + raise ValueError( + f"make_manager_based_rl_env num_envs must be a positive integer, got {num_envs!r}" + ) + if not isinstance(backend_type, str) or not backend_type: + raise ValueError( + "make_manager_based_rl_env backend_type must be a non-empty string, " + f"got {backend_type!r}" + ) + + cfg.validate() + assert cfg.scene is not None + base_name, body_state_requested = _resolve_backend_entity_contract(cfg) + backend_kwargs = env_backend_kwargs(cfg) + backend_kwargs["base_name"] = base_name + if backend_type in {"mujoco", "motrix"}: + backend_kwargs["add_body_sensors"] = body_state_requested + + backend = create_backend( + backend_type, + cfg.scene, + num_envs, + cfg.sim_dt, + **backend_kwargs, + ) + try: + return ManagerBasedRlEnv(cfg, backend, num_envs) + except Exception: + backend.cleanup_scene_assets() + raise + + # Isaac Lab capitalization is a spelling-only alias. There is one implementation. ManagerBasedRLEnv = ManagerBasedRlEnv ManagerBasedRLEnvCfg = ManagerBasedRlEnvCfg @@ -600,4 +680,5 @@ def close(self) -> None: "ManagerBasedRLEnvCfg", "ManagerBasedRlEnv", "ManagerBasedRlEnvCfg", + "make_manager_based_rl_env", ] diff --git a/tests/envs/test_manager_based_rl_env.py b/tests/envs/test_manager_based_rl_env.py index 107f482d1..fb6d9000c 100644 --- a/tests/envs/test_manager_based_rl_env.py +++ b/tests/envs/test_manager_based_rl_env.py @@ -19,6 +19,7 @@ ManagerBasedRlEnv, ManagerBasedRLEnvCfg, ManagerBasedRlEnvCfg, + make_manager_based_rl_env, mdp, ) from unilab.managers import ( @@ -384,6 +385,238 @@ def _make_env( def test_public_names_are_spelling_only_aliases() -> None: assert ManagerBasedRLEnv is ManagerBasedRlEnv assert ManagerBasedRLEnvCfg is ManagerBasedRlEnvCfg + assert make_manager_based_rl_env is manager_env_module.make_manager_based_rl_env + + +@pytest.mark.parametrize( + ("backend_type", "expects_body_materialization"), + [ + ("mujoco", True), + ("motrix", True), + ("mjwarp", False), + ("drake", False), + ], +) +def test_generic_factory_routes_only_public_backend_contract( + monkeypatch: pytest.MonkeyPatch, + backend_type: str, + expects_body_materialization: bool, +) -> None: + cfg = _make_cfg(include_optional_managers=False) + assert cfg.scene is not None + cfg.scene.entities["robot"] = EntityCfg( + root_body_name="base", + actuator_names=("motor",), + ) + backend = _FakeBackend(3) + constructed: dict[str, Any] = {} + + def fake_create_backend( + requested_backend: str, + scene: SceneCfg, + num_envs: int, + sim_dt: float, + **kwargs: Any, + ) -> SimBackend: + constructed.update( + backend_type=requested_backend, + scene=scene, + num_envs=num_envs, + sim_dt=sim_dt, + kwargs=kwargs, + ) + return cast(SimBackend, backend) + + sentinel = object() + + def fake_make_env( + received_cfg: ManagerBasedRlEnvCfg, + received_backend: SimBackend, + received_num_envs: int, + ) -> Any: + assert received_cfg is cfg + assert received_backend is backend + assert received_num_envs == 3 + return sentinel + + monkeypatch.setattr(manager_env_module, "create_backend", fake_create_backend) + monkeypatch.setattr(manager_env_module, "ManagerBasedRlEnv", fake_make_env) + + result = make_manager_based_rl_env(cfg, num_envs=3, backend_type=backend_type) + + assert result is sentinel + assert constructed["backend_type"] == backend_type + assert constructed["scene"] is cfg.scene + assert constructed["num_envs"] == 3 + assert constructed["sim_dt"] == cfg.sim_dt + kwargs = constructed["kwargs"] + assert kwargs["base_name"] == "base" + if expects_body_materialization: + assert kwargs["add_body_sensors"] is True + else: + assert "add_body_sensors" not in kwargs + for key, value in env_backend_kwargs(cfg).items(): + assert kwargs[key] == value + + +@pytest.mark.parametrize( + ("entities", "match"), + [ + ({"robot": EntityCfg(actuator_names=("motor",))}, "found 0 root entities"), + ( + { + "robot": EntityCfg(root_body_name="base", actuator_names=("motor",)), + "payload": EntityCfg(root_body_name="box"), + }, + "found 2 root entities.*robot.*payload", + ), + ], +) +def test_generic_factory_requires_exactly_one_explicit_root_entity( + monkeypatch: pytest.MonkeyPatch, + entities: dict[str, EntityCfg], + match: str, +) -> None: + cfg = _make_cfg(include_optional_managers=False) + assert cfg.scene is not None + cfg.scene.entities = entities + backend_constructed = False + + def reject_backend_construction(*args: Any, **kwargs: Any) -> SimBackend: + nonlocal backend_constructed + backend_constructed = True + raise AssertionError("backend construction must not run") + + monkeypatch.setattr(manager_env_module, "create_backend", reject_backend_construction) + + with pytest.raises(ValueError, match=match): + make_manager_based_rl_env(cfg, num_envs=2, backend_type="mujoco") + assert not backend_constructed + + +@pytest.mark.parametrize( + ("cfg_value", "num_envs", "backend_type", "error", "match"), + [ + (object(), 2, "mujoco", TypeError, "expected ManagerBasedRlEnvCfg"), + (None, 0, "mujoco", ValueError, "num_envs must be a positive integer"), + (None, True, "mujoco", ValueError, "num_envs must be a positive integer"), + (None, 2, "", ValueError, "backend_type must be a non-empty string"), + ], +) +def test_generic_factory_rejects_invalid_public_arguments_before_backend( + monkeypatch: pytest.MonkeyPatch, + cfg_value: object | None, + num_envs: int, + backend_type: str, + error: type[Exception], + match: str, +) -> None: + cfg = _make_cfg(include_optional_managers=False) if cfg_value is None else cfg_value + backend_constructed = False + + def reject_backend_construction(*args: Any, **kwargs: Any) -> SimBackend: + nonlocal backend_constructed + backend_constructed = True + raise AssertionError("backend construction must not run") + + monkeypatch.setattr(manager_env_module, "create_backend", reject_backend_construction) + + with pytest.raises(error, match=match): + make_manager_based_rl_env( + cast(ManagerBasedRlEnvCfg, cfg), + num_envs=num_envs, + backend_type=backend_type, + ) + assert not backend_constructed + + +@pytest.mark.parametrize( + ("entities", "error", "match"), + [ + ( + cast(dict[str, EntityCfg], {1: EntityCfg(root_body_name="base")}), + TypeError, + "entity names must be non-empty strings", + ), + ( + cast(dict[str, EntityCfg], {"robot": object()}), + TypeError, + "scene entity 'robot' must be EntityCfg", + ), + ( + {"robot": EntityCfg(root_body_name="")}, + TypeError, + "root_body_name must be a non-empty string", + ), + ], +) +def test_generic_factory_rejects_invalid_entity_contract_before_backend( + monkeypatch: pytest.MonkeyPatch, + entities: dict[str, EntityCfg], + error: type[Exception], + match: str, +) -> None: + cfg = _make_cfg(include_optional_managers=False) + assert cfg.scene is not None + cfg.scene.entities = entities + backend_constructed = False + + def reject_backend_construction(*args: Any, **kwargs: Any) -> SimBackend: + nonlocal backend_constructed + backend_constructed = True + raise AssertionError("backend construction must not run") + + monkeypatch.setattr(manager_env_module, "create_backend", reject_backend_construction) + + with pytest.raises(error, match=match): + make_manager_based_rl_env(cfg, num_envs=2, backend_type="mujoco") + assert not backend_constructed + + +def test_generic_factory_preserves_backend_construction_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + cfg = _make_cfg(include_optional_managers=False) + assert cfg.scene is not None + cfg.scene.entities["robot"] = EntityCfg( + root_body_name="base", + actuator_names=("motor",), + ) + + def fail_backend_construction(*args: Any, **kwargs: Any) -> SimBackend: + raise NotImplementedError("body-state capability is unavailable") + + monkeypatch.setattr(manager_env_module, "create_backend", fail_backend_construction) + + with pytest.raises(NotImplementedError, match="body-state capability is unavailable"): + make_manager_based_rl_env(cfg, num_envs=2, backend_type="mjwarp") + + +def test_generic_factory_cleans_backend_when_env_construction_fails( + monkeypatch: pytest.MonkeyPatch, +) -> None: + cfg = _make_cfg(include_optional_managers=False) + assert cfg.scene is not None + cfg.scene.entities["robot"] = EntityCfg( + root_body_name="base", + actuator_names=("motor",), + ) + backend = _FakeBackend(2) + + monkeypatch.setattr( + manager_env_module, + "create_backend", + lambda *args, **kwargs: cast(SimBackend, backend), + ) + + def fail_env_construction(*args: Any, **kwargs: Any) -> Any: + raise RuntimeError("manager initialization failed") + + monkeypatch.setattr(manager_env_module, "ManagerBasedRlEnv", fail_env_construction) + + with pytest.raises(RuntimeError, match="manager initialization failed"): + make_manager_based_rl_env(cfg, num_envs=2, backend_type="mujoco") + assert backend.cleanup_calls == 1 @pytest.mark.parametrize( @@ -581,7 +814,7 @@ def test_named_keyframe_snapshot_is_shared_by_entity_and_reset_cold_path() -> No def test_real_mujoco_backend_is_materialized_before_first_reset() -> None: scene = SceneCfg( model_file=str(ASSETS_ROOT_PATH / "robots" / "go2" / "scene_flat.xml"), - entities={"robot": EntityCfg()}, + entities={"robot": EntityCfg(root_body_name="base")}, ) cfg = ManagerBasedRlEnvCfg( scene=scene, @@ -599,15 +832,7 @@ def test_real_mujoco_backend_is_materialized_before_first_reset() -> None: terminations={"time_out": TerminationTermCfg(func=mdp.time_out, time_out=True)}, policy_observation_group="actor", ) - backend = create_backend( - "mujoco", - scene, - 2, - cfg.sim_dt, - base_name="base", - **env_backend_kwargs(cfg), - ) - env = ManagerBasedRlEnv(cfg, backend, 2) + env = make_manager_based_rl_env(cfg, num_envs=2, backend_type="mujoco") try: state = env.init_state() assert state.obs["obs"].shape == (2, 1) From eed7a901568f815c238ea49933a08770d932f1d0 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 19:07:24 +0800 Subject: [PATCH 030/169] feat(backend): support MuJoCo root velocity kicks [skip ci] (#1105) --- src/unilab/base/backend/mujoco/backend.py | 92 +++++++++++++++-- tests/base/test_sim_backend_smoke.py | 114 +++++++++++++++++++++- 2 files changed, 199 insertions(+), 7 deletions(-) diff --git a/src/unilab/base/backend/mujoco/backend.py b/src/unilab/base/backend/mujoco/backend.py index 682a15559..b14a3f03d 100644 --- a/src/unilab/base/backend/mujoco/backend.py +++ b/src/unilab/base/backend/mujoco/backend.py @@ -353,6 +353,7 @@ def __init__( self._root_qpos_dim, self._root_qvel_dim = _root_state_dims(self._model) self._num_dof_pos = self.nq - self._root_qpos_dim self._num_dof_vel = self.nv - self._root_qvel_dim + self._interval_root_velocity_qvel_ids = self._resolve_interval_root_velocity_qvel_ids() # State storage. nstate = mujoco.mj_stateSize(self._model, mujoco.mjtState.mjSTATE_FULLPHYSICS) @@ -652,6 +653,7 @@ def _apply_model_assignments( ) self._push_body_id = self._resolve_push_body_id(self._model) self._push_body_force_slice = self._resolve_push_body_force_slice(self._push_body_id) + self._interval_root_velocity_qvel_ids = self._resolve_interval_root_velocity_qvel_ids() self._base_body_mass = np.asarray(self._model.body_mass).copy() self._base_body_ipos = np.asarray(self._model.body_ipos).copy() self._pending_xfrc_applied = np.zeros( @@ -771,6 +773,19 @@ def get_root_state_layout(self, root_body_name: str) -> BackendRootStateLayout: qvel_indices=tuple(range(qvel_start, qvel_start + 6)), ) + def _resolve_interval_root_velocity_qvel_ids(self) -> tuple[int, int, int] | None: + """Bind the configured free root's world-linear qvel columns on the cold path.""" + if self._base_name is None or self._base_body_id < 0: + return None + try: + layout = self.get_root_state_layout(self._base_name) + except (NotImplementedError, ValueError): + return None + linear_ids = layout.qvel_indices[:3] + if linear_ids != tuple(range(linear_ids[0], linear_ids[0] + 3)): + return None + return int(linear_ids[0]), int(linear_ids[1]), int(linear_ids[2]) + def get_body_ids(self, names: "Sequence[str]") -> np.ndarray: ids: list[int] = [] for name in names: @@ -1074,6 +1089,9 @@ def get_dr_capabilities(self) -> DomainRandomizationCapabilities: } ), supports_interval_push=self._push_body_id >= 0, + supports_interval_body_velocity_delta=( + self._interval_root_velocity_qvel_ids is not None + ), supports_interval_body_force=True, ) @@ -1128,14 +1146,76 @@ def _apply_body_linear_velocity_delta( body_ids: np.ndarray, velocity_delta: np.ndarray, ) -> None: - """Apply a world-frame linear-velocity delta to specific bodies. + """Apply a row-selective world-frame velocity kick to the configured free root.""" + qvel_ids = self._interval_root_velocity_qvel_ids + if qvel_ids is None: + raise NotImplementedError( + "MuJoCo interval body velocity perturbation requires base_name to identify " + "a body with exactly one free joint" + ) - Backend-internal hook for ``apply_interval_randomization``; it is not - part of the public ``SimBackend`` surface. - """ - raise NotImplementedError( - f"{self.__class__.__name__} does not support interval body velocity perturbation" + raw_body_ids = np.asarray(body_ids) + if ( + raw_body_ids.ndim != 1 + or not np.issubdtype(raw_body_ids.dtype, np.integer) + or np.issubdtype(raw_body_ids.dtype, np.bool_) + ): + raise TypeError( + "MuJoCo interval body velocity perturbation body_ids must be a 1-D " + f"integer array, got shape={raw_body_ids.shape}, dtype={raw_body_ids.dtype}" + ) + resolved_body_ids = np.asarray(raw_body_ids, dtype=np.int32) + expected_body_ids = np.asarray([self._base_body_id], dtype=np.int32) + if not np.array_equal(resolved_body_ids, expected_body_ids): + raise NotImplementedError( + "MuJoCo interval body velocity perturbation only supports the configured " + f"free root body '{self._base_name}' (id={self._base_body_id}); " + f"received body_ids={resolved_body_ids.tolist()}" + ) + + if not isinstance(velocity_delta, np.ndarray): + raise TypeError( + "MuJoCo interval body velocity perturbation must be an np.ndarray, " + f"got {type(velocity_delta).__name__}" + ) + expected_shape = (self._num_envs, 1, 3) + if velocity_delta.shape != expected_shape: + raise ValueError( + "MuJoCo interval body velocity perturbation has shape " + f"{velocity_delta.shape}; expected {expected_shape}" + ) + if not np.issubdtype(velocity_delta.dtype, np.floating): + raise TypeError( + "MuJoCo interval body velocity perturbation must have floating dtype, " + f"got {velocity_delta.dtype}" + ) + if not np.isfinite(velocity_delta).all(): + raise ValueError("MuJoCo interval body velocity perturbation contains NaN or Inf") + if self._pool is None: + raise RuntimeError( + "MuJoCo interval body velocity perturbation requires a materialized backend" + ) + + active_rows = np.flatnonzero(np.any(velocity_delta[:, 0, :] != 0.0, axis=1)).astype( + np.int32, + copy=False, + ) + if active_rows.size == 0: + return + + state_rows = np.asarray(self._physics_state[active_rows], dtype=np.float64).copy() + state_qvel_ids = np.asarray( + [self._idx_qvel + qvel_id for qvel_id in qvel_ids], + dtype=np.intp, + ) + state_rows[:, state_qvel_ids] += velocity_delta[active_rows, 0, :] + state_out, sensor_out = self._pool.reset( + env_ids=active_rows, + initial_state=state_rows, + chunk_size=self._chunk_size, ) + self._physics_state[active_rows] = state_out.astype(self._np_dtype) + self._sensor_data[active_rows] = sensor_out.astype(self._np_dtype) def push_robots(self, force_range: Sequence[float] | np.ndarray) -> None: self._pending_xfrc_applied.fill(0.0) diff --git a/tests/base/test_sim_backend_smoke.py b/tests/base/test_sim_backend_smoke.py index 3fa5192c6..bc4ab6322 100644 --- a/tests/base/test_sim_backend_smoke.py +++ b/tests/base/test_sim_backend_smoke.py @@ -14,7 +14,7 @@ from unilab.assets import ASSETS_ROOT_PATH from unilab.base.backend.mujoco.xml import get_named_body_ids from unilab.base.scene import SceneCfg -from unilab.dr.types import ResetRandomizationPayload +from unilab.dr.types import IntervalRandomizationPlan, ResetRandomizationPayload pytest.importorskip("mujoco", reason="mujoco not installed") @@ -108,6 +108,7 @@ def test_mujoco_backend_smoke_contract(robot): "kd", }.issubset(caps.supported_reset_terms) assert caps.supports_interval_push + assert caps.supports_interval_body_velocity_delta def test_mujoco_backend_fixed_base_dof_views_do_not_skip_first_joint(): @@ -129,6 +130,117 @@ def test_mujoco_backend_fixed_base_dof_views_do_not_skip_first_joint(): np.testing.assert_allclose(bkd.get_base_lin_vel(), 0.0, atol=1e-8) np.testing.assert_allclose(bkd.get_base_ang_vel(), 0.0, atol=1e-8) _unit_quat(bkd.get_base_quat(), "MuJoCo fixed-base smoke") + assert not bkd.get_dr_capabilities().supports_interval_body_velocity_delta + + +def test_mujoco_interval_root_velocity_kick_is_row_selective_and_refreshes_sensors(): + from unilab.base.backend.mujoco.backend import MuJoCoBackend + + bkd = MuJoCoBackend( + SceneCfg(model_file=_xml("go2")), + NUM_ENVS, + SIM_DT, + base_name="base", + add_body_sensors=True, + ) + bkd.materialize() + qpos = np.broadcast_to(bkd.get_default_qpos(), (NUM_ENVS, bkd.model.nq)).copy() + qvel = np.zeros((NUM_ENVS, bkd.model.nv), dtype=np.float64) + bkd.set_state(np.arange(NUM_ENVS, dtype=np.int32), qpos, qvel) + bkd.step(np.zeros((NUM_ENVS, bkd.model.nu)), nsteps=1) + assert bkd._pool is not None + bkd._sensor_data[:] = bkd._pool.forward(bkd.get_physics_state()) + + body_ids = bkd.get_body_ids(("base",)) + state_before = bkd.get_physics_state().copy() + body_velocity_before = bkd.get_body_lin_vel_w(body_ids).copy() + delta = np.zeros((NUM_ENVS, 1, 3), dtype=np.float64) + delta[1, 0] = (0.25, -0.4, 0.15) + + bkd.apply_interval_randomization( + IntervalRandomizationPlan( + body_ids=body_ids, + body_linear_velocity_delta=delta, + ) + ) + + state_after = bkd.get_physics_state() + body_velocity_after = bkd.get_body_lin_vel_w(body_ids) + np.testing.assert_array_equal(state_after[0], state_before[0]) + np.testing.assert_allclose(state_after[1, 0], state_before[1, 0], atol=1e-12) + np.testing.assert_allclose( + state_after[1, bkd._idx_qpos : bkd._idx_qvel], + state_before[1, bkd._idx_qpos : bkd._idx_qvel], + atol=1e-12, + ) + np.testing.assert_allclose( + state_after[1, bkd._idx_qvel + 3 :], + state_before[1, bkd._idx_qvel + 3 :], + atol=1e-12, + ) + np.testing.assert_allclose( + state_after[1, bkd._idx_qvel : bkd._idx_qvel + 3], + state_before[1, bkd._idx_qvel : bkd._idx_qvel + 3] + delta[1, 0], + atol=1e-7, + ) + np.testing.assert_allclose( + body_velocity_after[1, 0], + body_velocity_before[1, 0] + delta[1, 0], + atol=1e-6, + ) + + +def test_mujoco_interval_root_velocity_kick_rejects_invalid_contracts(): + from unilab.base.backend.mujoco.backend import MuJoCoBackend + + bkd = MuJoCoBackend( + SceneCfg(model_file=_xml("go2")), + NUM_ENVS, + SIM_DT, + base_name="base", + ) + body_ids = bkd.get_body_ids(("base",)) + delta = np.zeros((NUM_ENVS, 1, 3), dtype=np.float64) + + with pytest.raises(RuntimeError, match="requires a materialized backend"): + bkd.apply_interval_randomization( + IntervalRandomizationPlan( + body_ids=body_ids, + body_linear_velocity_delta=delta, + ) + ) + + bkd.materialize() + with pytest.raises(TypeError, match="body_ids must be a 1-D integer array"): + bkd.apply_interval_randomization( + IntervalRandomizationPlan( + body_ids=body_ids.astype(np.float64), + body_linear_velocity_delta=delta, + ) + ) + with pytest.raises(NotImplementedError, match="only supports the configured free root"): + bkd.apply_interval_randomization( + IntervalRandomizationPlan( + body_ids=np.asarray([body_ids[0] + 1], dtype=np.int32), + body_linear_velocity_delta=delta, + ) + ) + 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, 3), dtype=np.float64), + ) + ) + invalid = delta.copy() + invalid[0, 0, 0] = np.nan + with pytest.raises(ValueError, match="contains NaN or Inf"): + bkd.apply_interval_randomization( + IntervalRandomizationPlan( + body_ids=body_ids, + body_linear_velocity_delta=invalid, + ) + ) @pytest.mark.parametrize("robot", BASIC_ROBOTS) From 9cf26b76469da5c641e2a7476a7c8fb250deea10 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 19:16:09 +0800 Subject: [PATCH 031/169] refactor(backend): isolate body-state materialization [skip ci] (#1107) --- src/unilab/base/backend/__init__.py | 14 +++ src/unilab/envs/manager_based_rl_env.py | 3 +- tests/base/test_motrix_backend_options.py | 100 ++++++++++++++++++++++ tests/envs/test_manager_based_rl_env.py | 17 +--- 4 files changed, 118 insertions(+), 16 deletions(-) diff --git a/src/unilab/base/backend/__init__.py b/src/unilab/base/backend/__init__.py index 299744338..ba23a3ffe 100644 --- a/src/unilab/base/backend/__init__.py +++ b/src/unilab/base/backend/__init__.py @@ -98,6 +98,8 @@ def create_backend( scene: SceneCfg, num_envs: int, sim_dt: float, + *, + body_state_required: bool = False, **kwargs, ) -> SimBackend: """Create a simulation backend. @@ -108,6 +110,9 @@ def create_backend( scene: SceneCfg for either static or composed scenes. num_envs: Number of environments. sim_dt: Simulation timestep. + body_state_required: Whether the caller requires public body-state views. + Backend adapters decide whether satisfying this request requires extra + cold-path scene materialization. **kwargs: Additional backend options such as ``position_actuator_gains``, ``iterations``, or ``motrix_max_iterations``. @@ -116,6 +121,11 @@ def create_backend( """ if scene is None: raise ValueError("SceneCfg must be provided") + if not isinstance(body_state_required, bool): + raise TypeError( + "create_backend body_state_required must be bool, " + f"got {type(body_state_required).__name__}" + ) position_actuator_gains = kwargs.pop("position_actuator_gains", None) motrix_max_iterations = kwargs.pop("motrix_max_iterations", None) @@ -131,6 +141,8 @@ def create_backend( drake_nthread = kwargs.pop("drake_nthread", None) if backend_type == "mujoco": MuJoCoBackend = _load_mujoco_backend() + if body_state_required: + kwargs["add_body_sensors"] = True if position_actuator_gains is not None: kwargs["position_actuator_gains"] = position_actuator_gains if post_step_forward_sensor is not None: @@ -177,6 +189,8 @@ def create_backend( MotrixBackend, motrix_available = _load_motrix_backend() if not motrix_available: raise ImportError("MotrixSim not available, install motrixsim package") + if body_state_required: + kwargs["add_body_sensors"] = True if motrix_max_iterations is not None: kwargs["max_iterations"] = motrix_max_iterations return cast(SimBackend, MotrixBackend(scene, num_envs, sim_dt, **kwargs)) diff --git a/src/unilab/envs/manager_based_rl_env.py b/src/unilab/envs/manager_based_rl_env.py index eb2b5af10..21d8edd0a 100644 --- a/src/unilab/envs/manager_based_rl_env.py +++ b/src/unilab/envs/manager_based_rl_env.py @@ -654,14 +654,13 @@ def make_manager_based_rl_env( base_name, body_state_requested = _resolve_backend_entity_contract(cfg) backend_kwargs = env_backend_kwargs(cfg) backend_kwargs["base_name"] = base_name - if backend_type in {"mujoco", "motrix"}: - backend_kwargs["add_body_sensors"] = body_state_requested backend = create_backend( backend_type, cfg.scene, num_envs, cfg.sim_dt, + body_state_required=body_state_requested, **backend_kwargs, ) try: diff --git a/tests/base/test_motrix_backend_options.py b/tests/base/test_motrix_backend_options.py index 60dac0538..aeefa158c 100644 --- a/tests/base/test_motrix_backend_options.py +++ b/tests/base/test_motrix_backend_options.py @@ -772,6 +772,106 @@ def __init__(self, scene: SceneCfg, num_envs: int, sim_dt: float, **kwargs: Any) assert captured["kwargs"]["post_step_forward_sensor"] is False +@pytest.mark.parametrize("body_state_required", [False, True]) +def test_create_backend_maps_body_state_request_inside_mujoco_adapter( + monkeypatch, body_state_required: bool +) -> None: + import unilab.base.backend as backend_factory + from unilab.base.scene import SceneCfg + + captured: dict[str, Any] = {} + + class FakeMuJoCoBackend: + def __init__(self, scene: SceneCfg, num_envs: int, sim_dt: float, **kwargs: Any) -> None: + captured["kwargs"] = kwargs + + monkeypatch.setattr(backend_factory, "_load_mujoco_backend", lambda: FakeMuJoCoBackend) + + backend_factory.create_backend( + "mujoco", + SceneCfg(model_file="model.xml"), + num_envs=1, + sim_dt=0.01, + body_state_required=body_state_required, + ) + + if body_state_required: + assert captured["kwargs"]["add_body_sensors"] is True + else: + assert "add_body_sensors" not in captured["kwargs"] + + +def test_create_backend_maps_body_state_request_inside_motrix_adapter(monkeypatch) -> None: + import unilab.base.backend as backend_factory + from unilab.base.scene import SceneCfg + + captured: dict[str, Any] = {} + + class FakeMotrixBackend: + def __init__(self, scene: SceneCfg, num_envs: int, sim_dt: float, **kwargs: Any) -> None: + captured["kwargs"] = kwargs + + monkeypatch.setattr( + backend_factory, + "_load_motrix_backend", + lambda: (FakeMotrixBackend, True), + ) + + backend_factory.create_backend( + "motrix", + SceneCfg(model_file="model.xml"), + num_envs=1, + sim_dt=0.01, + body_state_required=True, + ) + + assert captured["kwargs"]["add_body_sensors"] is True + + +@pytest.mark.parametrize("backend_type", ["drake", "mjwarp"]) +def test_create_backend_keeps_body_state_request_out_of_native_state_adapters( + monkeypatch, backend_type: str +) -> None: + import unilab.base.backend as backend_factory + from unilab.base.scene import SceneCfg + + captured: dict[str, Any] = {} + + class FakeBackend: + def __init__(self, scene: SceneCfg, num_envs: int, sim_dt: float, **kwargs: Any) -> None: + captured["kwargs"] = kwargs + + if backend_type == "drake": + monkeypatch.setattr(backend_factory, "_load_drake_backend", lambda: FakeBackend) + else: + monkeypatch.setattr(backend_factory, "_load_mjwarp_backend", lambda: FakeBackend) + + backend_factory.create_backend( + backend_type, + SceneCfg(model_file="model.xml"), + num_envs=1, + sim_dt=0.01, + body_state_required=True, + ) + + assert "body_state_required" not in captured["kwargs"] + assert "add_body_sensors" not in captured["kwargs"] + + +def test_create_backend_rejects_non_bool_body_state_request() -> None: + import unilab.base.backend as backend_factory + from unilab.base.scene import SceneCfg + + with pytest.raises(TypeError, match="body_state_required must be bool"): + backend_factory.create_backend( + "mujoco", + SceneCfg(model_file="model.xml"), + num_envs=1, + sim_dt=0.01, + body_state_required=1, # type: ignore[arg-type] + ) + + def test_create_backend_does_not_route_post_step_forward_sensor_to_motrix(monkeypatch) -> None: import unilab.base.backend as backend_factory from unilab.base.scene import SceneCfg diff --git a/tests/envs/test_manager_based_rl_env.py b/tests/envs/test_manager_based_rl_env.py index fb6d9000c..0cbf39fd7 100644 --- a/tests/envs/test_manager_based_rl_env.py +++ b/tests/envs/test_manager_based_rl_env.py @@ -388,19 +388,10 @@ def test_public_names_are_spelling_only_aliases() -> None: assert make_manager_based_rl_env is manager_env_module.make_manager_based_rl_env -@pytest.mark.parametrize( - ("backend_type", "expects_body_materialization"), - [ - ("mujoco", True), - ("motrix", True), - ("mjwarp", False), - ("drake", False), - ], -) +@pytest.mark.parametrize("backend_type", ["mujoco", "motrix", "mjwarp", "drake"]) def test_generic_factory_routes_only_public_backend_contract( monkeypatch: pytest.MonkeyPatch, backend_type: str, - expects_body_materialization: bool, ) -> None: cfg = _make_cfg(include_optional_managers=False) assert cfg.scene is not None @@ -451,10 +442,8 @@ def fake_make_env( assert constructed["sim_dt"] == cfg.sim_dt kwargs = constructed["kwargs"] assert kwargs["base_name"] == "base" - if expects_body_materialization: - assert kwargs["add_body_sensors"] is True - else: - assert "add_body_sensors" not in kwargs + assert kwargs["body_state_required"] is True + assert "add_body_sensors" not in kwargs for key, value in env_backend_kwargs(cfg).items(): assert kwargs[key] == value From 89400bd24a44f46145ed8b1f625e567a86909cd4 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 19:26:24 +0800 Subject: [PATCH 032/169] feat(sim2sim): guard manager policy contracts [skip ci] (#1109) --- .../2-sim_to_sim/7-config_guard.md | 8 +- .../2-sim_to_sim/7-config_guard.md | 6 +- src/unilab/training/sim2sim.py | 4 + tests/training/test_sim2sim_resolver.py | 90 ++++++++++++++++++- 4 files changed, 105 insertions(+), 3 deletions(-) diff --git a/docs/sphinx/source/en/3-deployment/2-sim_to_sim/7-config_guard.md b/docs/sphinx/source/en/3-deployment/2-sim_to_sim/7-config_guard.md index ae58f3f13..eaa655e34 100644 --- a/docs/sphinx/source/en/3-deployment/2-sim_to_sim/7-config_guard.md +++ b/docs/sphinx/source/en/3-deployment/2-sim_to_sim/7-config_guard.md @@ -27,7 +27,7 @@ Fields are classified by dotted path into three tiers (see `src/unilab/training/ | Tier | Behavior | Fields | |---|---|---| -| **DENYLIST** | Mismatch → `CrossBackendIncompatibleError`, aborts | `algo.obs_groups`, `env.control_config.action_scale`, `algo.policy.actor_hidden_dims` / `critic_hidden_dims`, `algo.empirical_normalization` / `algo.obs_normalization`, `env.sampling_mode` | +| **DENYLIST** | Mismatch → `CrossBackendIncompatibleError`, aborts | `algo.obs_groups`, legacy `env.control_config.action_scale`, Manager-Based `env.observations` / `env.actions` / policy and critic group mapping, `algo.policy.actor_hidden_dims` / `critic_hidden_dims`, `algo.empirical_normalization` / `algo.obs_normalization`, `env.sampling_mode` | | **WARNING_LIST** | Prints a warning, continues | `reward.*`, `env.control_config.simulate_action_latency`, `env.ctrl_dt` | | **ALLOWLIST** | Free to override, not checked | `training.sim_backend`, `env.scene`, `training.play_steps`, `env.domain_rand`, `env.noise_config`, `env.commands.vel_limit` | @@ -40,6 +40,12 @@ If the target backend's DENYLIST fields differ from training (e.g. a task whose > Legacy runs: if `run_config.json` has no `contract_snapshot` (older training), the guard skips with a warning instead of breaking your workflow. +Manager-Based snapshots store the complete typed observation and action declarations from +Hydra. A snapshot from before those fields existed cannot prove that its policy I/O is +equivalent to a Manager-Based target, so asymmetric presence fails closed. Set +`training.sim2sim_strict=false` only as an explicit user override; the load-time dimension +guard still remains active. + ## See also - {doc}`1-backend_swap` diff --git a/docs/sphinx/source/zh_CN/3-deployment/2-sim_to_sim/7-config_guard.md b/docs/sphinx/source/zh_CN/3-deployment/2-sim_to_sim/7-config_guard.md index 7c94210e1..6069ca1fe 100644 --- a/docs/sphinx/source/zh_CN/3-deployment/2-sim_to_sim/7-config_guard.md +++ b/docs/sphinx/source/zh_CN/3-deployment/2-sim_to_sim/7-config_guard.md @@ -27,7 +27,7 @@ uv run eval --algo ppo --task go2_joystick_flat --sim motrix --load-run -1 | 档位 | 行为 | 字段 | |---|---|---| -| **DENYLIST** | 差异即 `CrossBackendIncompatibleError`，中断 | `algo.obs_groups`、`env.control_config.action_scale`、`algo.policy.actor_hidden_dims` / `critic_hidden_dims`、`algo.empirical_normalization` / `algo.obs_normalization`、`env.sampling_mode` | +| **DENYLIST** | 差异即 `CrossBackendIncompatibleError`，中断 | `algo.obs_groups`、legacy `env.control_config.action_scale`、Manager-Based `env.observations` / `env.actions` / policy 与 critic group mapping、`algo.policy.actor_hidden_dims` / `critic_hidden_dims`、`algo.empirical_normalization` / `algo.obs_normalization`、`env.sampling_mode` | | **WARNING_LIST** | 仅打印 warning，继续 | `reward.*`、`env.control_config.simulate_action_latency`、`env.ctrl_dt` | | **ALLOWLIST** | 自由覆盖，不检查 | `training.sim_backend`、`env.scene`、`training.play_steps`、`env.domain_rand`、`env.noise_config`、`env.commands.vel_limit` | @@ -40,6 +40,10 @@ uv run eval --algo ppo --task go2_joystick_flat --sim motrix --load-run -1 > 兼容旧 run：若 `run_config.json` 没有 `contract_snapshot`（早期训练），守卫自动跳过并打印 warning，不会中断现有工作流。 +Manager-Based snapshot 会保存 Hydra 中完整的 typed observation/action 声明。缺少这些字段的旧 +snapshot 无法证明其 policy I/O 与 Manager-Based 目标等价，因此不对称出现时会 fail-closed。 +只有用户显式设置 `training.sim2sim_strict=false` 才会继续；加载权重时的维度守卫仍然生效。 + ## 另请参阅 - {doc}`1-backend_swap` diff --git a/src/unilab/training/sim2sim.py b/src/unilab/training/sim2sim.py index 9b81224e7..d7e35426d 100644 --- a/src/unilab/training/sim2sim.py +++ b/src/unilab/training/sim2sim.py @@ -36,6 +36,10 @@ class CrossBackendIncompatibleError(RuntimeError): DENYLIST: list[str] = [ "algo.obs_groups", "env.control_config.action_scale", + "env.observations", + "env.actions", + "env.policy_observation_group", + "env.critic_observation_group", "algo.policy.actor_hidden_dims", "algo.policy.critic_hidden_dims", "algo.empirical_normalization", diff --git a/tests/training/test_sim2sim_resolver.py b/tests/training/test_sim2sim_resolver.py index 75edb4d0c..fc971bbba 100644 --- a/tests/training/test_sim2sim_resolver.py +++ b/tests/training/test_sim2sim_resolver.py @@ -56,6 +56,47 @@ def _mujoco_cfg() -> Any: ) +def _manager_cfg() -> Any: + return OmegaConf.create( + { + "training": {"sim_backend": "mujoco"}, + "algo": {"obs_groups": {"actor": ["actor"]}}, + "env": { + "observations": { + "policy": { + "_target_": "unilab.managers.ObservationGroupCfg", + "terms": { + "joint_pos": { + "_target_": "unilab.managers.ObservationTermCfg", + "func": "unilab.envs.mdp.joint_pos_rel", + } + }, + }, + "critic": { + "_target_": "unilab.managers.ObservationGroupCfg", + "terms": { + "joint_pos": { + "_target_": "unilab.managers.ObservationTermCfg", + "func": "unilab.envs.mdp.joint_pos_rel", + } + }, + }, + }, + "actions": { + "joint_pos": { + "_target_": "unilab.envs.mdp.JointPositionActionCfg", + "entity_name": "robot", + "actuator_names": [".*"], + "scale": 0.25, + } + }, + "policy_observation_group": "policy", + "critic_observation_group": "critic", + }, + } + ) + + def test_field_lists_are_disjoint(): deny, warn, allow = set(DENYLIST), set(WARNING_LIST), set(ALLOWLIST) assert deny.isdisjoint(warn) @@ -83,6 +124,46 @@ def test_snapshot_json_round_trips(): assert json.loads(json.dumps(snapshot)) == snapshot +def test_extract_snapshot_captures_manager_policy_io_contract() -> None: + cfg = _manager_cfg() + + snapshot = extract_contract_snapshot(cfg) + + assert snapshot["env.actions"]["joint_pos"]["scale"] == pytest.approx(0.25) + assert snapshot["env.observations"]["policy"]["terms"]["joint_pos"]["func"] == ( + "unilab.envs.mdp.joint_pos_rel" + ) + assert snapshot["env.policy_observation_group"] == "policy" + assert snapshot["env.critic_observation_group"] == "critic" + assert "env.control_config.action_scale" not in snapshot + + +def test_matching_manager_policy_io_contract_passes(tmp_path: Path) -> None: + _write_sidecar(tmp_path, extract_contract_snapshot(_manager_cfg())) + target = _manager_cfg() + + assert resolve_sim2sim_config(tmp_path, target) is target + + +@pytest.mark.parametrize( + ("path", "value"), + [ + ("env.actions.joint_pos.scale", 0.5), + ("env.observations.policy.terms.joint_pos.func", "unilab.envs.mdp.joint_vel_rel"), + ("env.policy_observation_group", "critic"), + ("env.critic_observation_group", "policy"), + ], +) +def test_manager_policy_io_mismatch_fails_closed(tmp_path: Path, path: str, value: object) -> None: + source = _manager_cfg() + _write_sidecar(tmp_path, extract_contract_snapshot(source)) + target = _manager_cfg() + OmegaConf.update(target, path, value, merge=False) + + with pytest.raises(CrossBackendIncompatibleError, match="env\\."): + resolve_sim2sim_config(tmp_path, target) + + def test_matching_contract_returns_same_cfg(tmp_path): _write_sidecar(tmp_path, extract_contract_snapshot(_mujoco_cfg())) target = _mujoco_cfg() @@ -185,7 +266,14 @@ def test_action_scale_list_form(tmp_path): def test_env_structural_denylist_is_the_env_subset(): - assert ENV_STRUCTURAL_DENYLIST == ["env.control_config.action_scale", "env.sampling_mode"] + assert ENV_STRUCTURAL_DENYLIST == [ + "env.control_config.action_scale", + "env.observations", + "env.actions", + "env.policy_observation_group", + "env.critic_observation_group", + "env.sampling_mode", + ] assert set(ENV_STRUCTURAL_DENYLIST) <= set(DENYLIST) From a812afe5e76b27a9aa723b01009d1578845d0185 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 19:56:56 +0800 Subject: [PATCH 033/169] feat(env): migrate Go2 flat production to Hydra MBA [skip ci] (#1111) Closes #1110 --- conf/appo/task/go2_joystick_flat/base.yaml | 241 +++++++ conf/appo/task/go2_joystick_flat/motrix.yaml | 23 +- conf/appo/task/go2_joystick_flat/mujoco.yaml | 18 +- .../flashsac/go2_joystick_flat/mujoco.yaml | 135 +++- .../task/go2_joystick_flat/base.yaml | 237 +++++++ .../task/sac/go2_joystick_flat/drake.yaml | 26 +- .../task/td3/go2_joystick_flat/motrix.yaml | 30 +- conf/ppo/task/go2_joystick_flat/base.yaml | 237 +++++++ conf/ppo/task/go2_joystick_flat/drake.yaml | 28 +- conf/ppo/task/go2_joystick_flat/motrix.yaml | 32 +- conf/ppo/task/go2_joystick_flat/mujoco.yaml | 19 +- src/unilab/envs/locomotion/go2/joystick.py | 17 +- .../envs/locomotion/go2/manager_based_cfg.py | 232 ------- tests/algos/test_appo_runner.py | 14 +- tests/algos/test_rsl_rl_runner.py | 19 +- tests/config/test_config_system.py | 51 +- tests/config/test_locomotion_params.py | 4 +- tests/config/test_reward_injection.py | 5 +- .../locomotion/go2/test_manager_based_cfg.py | 617 +++++++++--------- tests/envs/test_env_configs.py | 72 -- tests/scripts/test_train_scripts.py | 24 +- 21 files changed, 1267 insertions(+), 814 deletions(-) create mode 100644 conf/appo/task/go2_joystick_flat/base.yaml create mode 100644 conf/offpolicy/task/go2_joystick_flat/base.yaml create mode 100644 conf/ppo/task/go2_joystick_flat/base.yaml delete mode 100644 src/unilab/envs/locomotion/go2/manager_based_cfg.py diff --git a/conf/appo/task/go2_joystick_flat/base.yaml b/conf/appo/task/go2_joystick_flat/base.yaml new file mode 100644 index 000000000..ad7eb2fba --- /dev/null +++ b/conf/appo/task/go2_joystick_flat/base.yaml @@ -0,0 +1,241 @@ +# @package _global_ +# Canonical Go2 flat Manager-Based task declaration. Backend owner leaves inherit +# this file and only override backend/algo tuning or explicitly disabled terms. +env: + scene: + model_file: src/unilab/assets/robots/go2/scene_flat.xml + default_keyframe_name: home + entities: + robot: + root_body_name: base + joint_names: + - FL_hip_joint + - FL_thigh_joint + - FL_calf_joint + - FR_hip_joint + - FR_thigh_joint + - FR_calf_joint + - RL_hip_joint + - RL_thigh_joint + - RL_calf_joint + - RR_hip_joint + - RR_thigh_joint + - RR_calf_joint + actuator_names: + - FR_hip + - FR_thigh + - FR_calf + - FL_hip + - FL_thigh + - FL_calf + - RR_hip + - RR_thigh + - RR_calf + - RL_hip + - RL_thigh + - RL_calf + sim_dt: 0.01 + ctrl_dt: 0.02 + max_episode_seconds: 20.0 + observations: + policy: + _target_: unilab.managers.ObservationGroupCfg + terms: + base_ang_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.builtin_sensor + params: + sensor_name: gyro + projected_gravity: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.projected_gravity_from_sensor + params: + sensor_name: upvector + joint_pos: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_pos_rel + joint_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_vel_rel + actions: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.last_action + command: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.generated_commands + params: + command_name: twist + gait_phase: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.locomotion.common.manager_terms.quadruped_gait_phase + params: + frequency: 2.0 + critic: + _target_: unilab.managers.ObservationGroupCfg + terms: + base_ang_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.builtin_sensor + params: + sensor_name: gyro + projected_gravity: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.projected_gravity_from_sensor + params: + sensor_name: upvector + joint_pos: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_pos_rel + joint_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_vel_rel + actions: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.last_action + command: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.generated_commands + params: + command_name: twist + gait_phase: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.locomotion.common.manager_terms.quadruped_gait_phase + params: + frequency: 2.0 + base_lin_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.builtin_sensor + params: + sensor_name: local_linvel + actions: + joint_pos: + _target_: unilab.envs.mdp.JointPositionActionCfg + entity_name: robot + actuator_names: [".*"] + scale: 0.25 + use_default_offset: true + commands: + twist: + _target_: unilab.envs.mdp.UniformVelocityCommandCfg + entity_name: robot + resampling_time_range: [20.0, 20.0] + heading_command: false + heading_control_stiffness: 0.5 + rel_standing_envs: 0.0 + rel_heading_envs: 0.0 + rel_world_envs: 0.0 + rel_forward_envs: 0.0 + init_velocity_prob: 0.0 + ranges: + lin_vel_x: [-0.6, 1.0] + lin_vel_y: [-0.4, 0.4] + ang_vel_z: [-0.8, 0.8] + events: + reset_scene_to_default: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.reset_scene_to_default + mode: reset + reset_root_state_uniform: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.reset_root_state_uniform + mode: reset + params: + pose_range: + x: [-0.5, 0.5] + y: [-0.5, 0.5] + z: [0.0, 0.0] + roll: [0.0, 0.0] + pitch: [0.0, 0.0] + yaw: [-3.141592653589793, 3.141592653589793] + velocity_range: + x: [-0.5, 0.5] + y: [-0.5, 0.5] + z: [-0.5, 0.5] + roll: [-0.5, 0.5] + pitch: [-0.5, 0.5] + yaw: [-0.5, 0.5] + pd_gains: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.pd_gains + mode: reset + params: + kp_range: [31.5, 38.5] + kd_range: [0.45, 0.55] + operation: abs + terminations: + time_out: + _target_: unilab.managers.TerminationTermCfg + func: unilab.envs.mdp.time_out + time_out: true + bad_orientation: + _target_: unilab.managers.TerminationTermCfg + func: unilab.envs.mdp.bad_orientation + params: + limit_angle: 1.0471975511965976 + policy_observation_group: policy + critic_observation_group: critic + +reward: + tracking_lin_vel: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.locomotion.common.manager_terms.track_lin_vel_xy_exp + weight: 1.0 + params: + std: 0.5 + command_name: twist + tracking_ang_vel: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.locomotion.common.manager_terms.track_ang_vel_z_exp + weight: 0.2 + params: + std: 0.5 + command_name: twist + lin_vel_z: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.locomotion.common.manager_terms.lin_vel_z_l2 + weight: -5.0 + ang_vel_xy: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.locomotion.common.manager_terms.ang_vel_xy_l2 + weight: -0.1 + base_height: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.locomotion.common.manager_terms.base_height_l2 + weight: -100.0 + params: + target_height: 0.3 + action_rate: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.mdp.action_rate_l2 + weight: -0.005 + similar_to_default: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.locomotion.common.manager_terms.joint_deviation_l1 + weight: -0.1 + alive: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.mdp.is_alive + weight: 0.0 + contact: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.locomotion.common.manager_terms.feet_phase_contact + weight: 0.24 + params: + frequency: 2.0 + sensor_names: + - FL_foot_contact + - FR_foot_contact + - RL_foot_contact + - RR_foot_contact + contact_threshold: 0.1 + stance_threshold: 0.6 + swing_feet_z: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.locomotion.common.manager_terms.feet_phase_swing_height + weight: 4.0 + params: + frequency: 2.0 + sensor_names: [FL_pos, FR_pos, RL_pos, RR_pos] + target_height: 0.1 + kernel: 0.01 + swing_start: 0.6 diff --git a/conf/appo/task/go2_joystick_flat/motrix.yaml b/conf/appo/task/go2_joystick_flat/motrix.yaml index fc4597bbd..b46281fe6 100644 --- a/conf/appo/task/go2_joystick_flat/motrix.yaml +++ b/conf/appo/task/go2_joystick_flat/motrix.yaml @@ -1,4 +1,8 @@ # @package _global_ +defaults: + - /task/go2_joystick_flat/base + - _self_ + training: task_name: Go2JoystickFlat sim_backend: motrix @@ -7,18 +11,7 @@ algo: steps_per_env: 24 max_iterations: 180 env: - sim_dt: 0.015 -reward: - scales: - tracking_lin_vel: 1.0 - tracking_ang_vel: 0.2 - lin_vel_z: -5.0 - ang_vel_xy: -0.1 - base_height: -100.0 - action_rate: -0.005 - similar_to_default: -0.1 - alive: 0.0 - contact: 0.24 - swing_feet_z: 4.0 - tracking_sigma: 0.25 - base_height_target: 0.3 + # The former 0.015/0.02 ratio was not an integer number of physics substeps. + # MBA keeps the shared 0.01/0.02 contract and fails on fractional ratios. + events: + pd_gains: null diff --git a/conf/appo/task/go2_joystick_flat/mujoco.yaml b/conf/appo/task/go2_joystick_flat/mujoco.yaml index d25c356ad..aec4a0437 100644 --- a/conf/appo/task/go2_joystick_flat/mujoco.yaml +++ b/conf/appo/task/go2_joystick_flat/mujoco.yaml @@ -1,20 +1,10 @@ # @package _global_ +defaults: + - /task/go2_joystick_flat/base + - _self_ + training: task_name: Go2JoystickFlat sim_backend: mujoco algo: max_iterations: 150 -reward: - scales: - tracking_lin_vel: 1.0 - tracking_ang_vel: 0.2 - lin_vel_z: -5.0 - ang_vel_xy: -0.1 - base_height: -100.0 - action_rate: -0.005 - similar_to_default: -0.1 - alive: 0.0 - contact: 0.24 - swing_feet_z: 4.0 - tracking_sigma: 0.25 - base_height_target: 0.3 diff --git a/conf/offpolicy/task/flashsac/go2_joystick_flat/mujoco.yaml b/conf/offpolicy/task/flashsac/go2_joystick_flat/mujoco.yaml index 031d8e2b7..ba0ddaa30 100644 --- a/conf/offpolicy/task/flashsac/go2_joystick_flat/mujoco.yaml +++ b/conf/offpolicy/task/flashsac/go2_joystick_flat/mujoco.yaml @@ -1,4 +1,8 @@ # @package _global_ +defaults: + - /task/go2_joystick_flat/base + - _self_ + training: task_name: Go2JoystickFlat sim_backend: mujoco @@ -12,32 +16,107 @@ algo: replay_buffer_n: 4096 tau: 0.05 env: - control_config: - action_scale: 0.4 - domain_rand: - randomize_kp: true - randomize_kd: true - randomize_base_mass: true - random_com: true - randomize_gravity: true - push_robots: true - noise_config: - level: 1.0 - scale_joint_angle: 0.01 - scale_joint_vel: 0.1 - scale_gyro: 0.0 - scale_gravity: 0.0 - scale_linvel: 0.0 + scene: + entities: + robot: + body_names: [base] + actions: + joint_pos: + scale: 0.4 + observations: + policy: + enable_corruption: true + terms: + joint_pos: + noise: + _target_: unilab.managers._noise.UniformNoiseCfg + n_min: -0.01 + n_max: 0.01 + joint_vel: + noise: + _target_: unilab.managers._noise.UniformNoiseCfg + n_min: -0.1 + n_max: 0.1 + critic: + enable_corruption: true + terms: + joint_pos: + noise: + _target_: unilab.managers._noise.UniformNoiseCfg + n_min: -0.01 + n_max: 0.01 + joint_vel: + noise: + _target_: unilab.managers._noise.UniformNoiseCfg + n_min: -0.1 + n_max: 0.1 + events: + randomize_rigid_body_mass: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.randomize_rigid_body_mass + mode: reset + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + body_names: base + mass_distribution_params: [-1.5, 1.5] + operation: add + recompute_inertia: false + randomize_rigid_body_com: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.randomize_rigid_body_com + mode: reset + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + body_names: base + com_range: + x: [-0.05, 0.05] + randomize_physics_scene_gravity: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.randomize_physics_scene_gravity + mode: reset + params: + gravity_distribution_params: + - [0.0, 0.0, -9.81] + - [0.0, 0.0, -9.81] + operation: abs + push_by_setting_velocity: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.push_by_setting_velocity + mode: interval + interval_range_s: [15.0, 15.0] + is_global_time: true + params: + velocity_range: + x: [-1.0, 1.0] + y: [-1.0, 1.0] + z: [-0.5, 0.5] + roll: [0.0, 0.0] + pitch: [0.0, 0.0] + yaw: [0.0, 0.0] reward: - scales: - tracking_lin_vel: 1.0 - tracking_ang_vel: 0.2 - lin_vel_z: -5.0 - ang_vel_xy: -0.1 - base_height: -20.0 - action_rate: -0.02 - similar_to_default: -0.4 - contact: 1.5 - swing_feet_z: 4.0 - tracking_sigma: 0.4 - base_height_target: 0.3 \ No newline at end of file + tracking_lin_vel: + weight: 1.0 + params: + std: 0.6324555320336759 + tracking_ang_vel: + weight: 0.2 + params: + std: 0.6324555320336759 + lin_vel_z: + weight: -5.0 + ang_vel_xy: + weight: -0.1 + base_height: + weight: -20.0 + action_rate: + weight: -0.02 + similar_to_default: + weight: -0.4 + contact: + weight: 1.5 + swing_feet_z: + weight: 4.0 diff --git a/conf/offpolicy/task/go2_joystick_flat/base.yaml b/conf/offpolicy/task/go2_joystick_flat/base.yaml new file mode 100644 index 000000000..3c51d7605 --- /dev/null +++ b/conf/offpolicy/task/go2_joystick_flat/base.yaml @@ -0,0 +1,237 @@ +# @package _global_ +# Canonical Go2 flat Manager-Based task declaration. Backend owner leaves inherit +# this file and only override backend/algo tuning or explicitly disabled terms. +env: + scene: + model_file: src/unilab/assets/robots/go2/scene_flat.xml + default_keyframe_name: home + entities: + robot: + root_body_name: base + joint_names: + - FL_hip_joint + - FL_thigh_joint + - FL_calf_joint + - FR_hip_joint + - FR_thigh_joint + - FR_calf_joint + - RL_hip_joint + - RL_thigh_joint + - RL_calf_joint + - RR_hip_joint + - RR_thigh_joint + - RR_calf_joint + actuator_names: + - FR_hip + - FR_thigh + - FR_calf + - FL_hip + - FL_thigh + - FL_calf + - RR_hip + - RR_thigh + - RR_calf + - RL_hip + - RL_thigh + - RL_calf + sim_dt: 0.01 + ctrl_dt: 0.02 + max_episode_seconds: 20.0 + observations: + policy: + _target_: unilab.managers.ObservationGroupCfg + terms: + base_ang_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.builtin_sensor + params: + sensor_name: gyro + projected_gravity: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.projected_gravity_from_sensor + params: + sensor_name: upvector + joint_pos: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_pos_rel + joint_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_vel_rel + actions: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.last_action + command: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.generated_commands + params: + command_name: twist + gait_phase: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.locomotion.common.manager_terms.quadruped_gait_phase + params: + frequency: 2.0 + critic: + _target_: unilab.managers.ObservationGroupCfg + terms: + base_ang_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.builtin_sensor + params: + sensor_name: gyro + projected_gravity: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.projected_gravity_from_sensor + params: + sensor_name: upvector + joint_pos: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_pos_rel + joint_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_vel_rel + actions: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.last_action + command: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.generated_commands + params: + command_name: twist + gait_phase: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.locomotion.common.manager_terms.quadruped_gait_phase + params: + frequency: 2.0 + base_lin_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.builtin_sensor + params: + sensor_name: local_linvel + actions: + joint_pos: + _target_: unilab.envs.mdp.JointPositionActionCfg + entity_name: robot + actuator_names: [".*"] + scale: 0.25 + use_default_offset: true + commands: + twist: + _target_: unilab.envs.mdp.UniformVelocityCommandCfg + entity_name: robot + resampling_time_range: [20.0, 20.0] + heading_command: false + heading_control_stiffness: 0.5 + rel_standing_envs: 0.0 + rel_heading_envs: 0.0 + rel_world_envs: 0.0 + rel_forward_envs: 0.0 + init_velocity_prob: 0.0 + ranges: + lin_vel_x: [-0.6, 1.0] + lin_vel_y: [-0.4, 0.4] + ang_vel_z: [-0.8, 0.8] + events: + reset_scene_to_default: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.reset_scene_to_default + mode: reset + reset_root_state_uniform: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.reset_root_state_uniform + mode: reset + params: + pose_range: + x: [-0.5, 0.5] + y: [-0.5, 0.5] + z: [0.0, 0.0] + roll: [0.0, 0.0] + pitch: [0.0, 0.0] + yaw: [-3.141592653589793, 3.141592653589793] + velocity_range: + x: [-0.5, 0.5] + y: [-0.5, 0.5] + z: [-0.5, 0.5] + roll: [-0.5, 0.5] + pitch: [-0.5, 0.5] + yaw: [-0.5, 0.5] + pd_gains: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.pd_gains + mode: reset + params: + kp_range: [31.5, 38.5] + kd_range: [0.45, 0.55] + operation: abs + terminations: + time_out: + _target_: unilab.managers.TerminationTermCfg + func: unilab.envs.mdp.time_out + time_out: true + bad_orientation: + _target_: unilab.managers.TerminationTermCfg + func: unilab.envs.mdp.bad_orientation + params: + limit_angle: 1.0471975511965976 + policy_observation_group: policy + critic_observation_group: critic + +reward: + tracking_lin_vel: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.locomotion.common.manager_terms.track_lin_vel_xy_exp + weight: 1.0 + params: + std: 0.5 + command_name: twist + tracking_ang_vel: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.locomotion.common.manager_terms.track_ang_vel_z_exp + weight: 0.2 + params: + std: 0.5 + command_name: twist + lin_vel_z: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.locomotion.common.manager_terms.lin_vel_z_l2 + weight: -5.0 + ang_vel_xy: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.locomotion.common.manager_terms.ang_vel_xy_l2 + weight: -0.1 + base_height: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.locomotion.common.manager_terms.base_height_l2 + weight: -100.0 + params: + target_height: 0.3 + action_rate: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.mdp.action_rate_l2 + weight: -0.005 + similar_to_default: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.locomotion.common.manager_terms.joint_deviation_l1 + weight: -0.1 + contact: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.locomotion.common.manager_terms.feet_phase_contact + weight: 0.24 + params: + frequency: 2.0 + sensor_names: + - FL_foot_contact + - FR_foot_contact + - RL_foot_contact + - RR_foot_contact + contact_threshold: 0.1 + stance_threshold: 0.6 + swing_feet_z: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.locomotion.common.manager_terms.feet_phase_swing_height + weight: 4.0 + params: + frequency: 2.0 + sensor_names: [FL_pos, FR_pos, RL_pos, RR_pos] + target_height: 0.1 + kernel: 0.01 + swing_start: 0.6 diff --git a/conf/offpolicy/task/sac/go2_joystick_flat/drake.yaml b/conf/offpolicy/task/sac/go2_joystick_flat/drake.yaml index 6829cfe93..b709df1b7 100644 --- a/conf/offpolicy/task/sac/go2_joystick_flat/drake.yaml +++ b/conf/offpolicy/task/sac/go2_joystick_flat/drake.yaml @@ -1,4 +1,8 @@ # @package _global_ +defaults: + - /task/go2_joystick_flat/base + - _self_ + training: task_name: Go2JoystickFlat sim_backend: drake @@ -28,23 +32,5 @@ algo: env: drake_backend_mode: batch drake_nthread: 20 - scene: - model_file: src/unilab/assets/robots/go2/scene_flat.xml - domain_rand: - randomize_kp: false - randomize_kd: false - push_robots: false - -reward: - scales: - tracking_lin_vel: 1.0 - tracking_ang_vel: 0.2 - lin_vel_z: -5.0 - ang_vel_xy: -0.1 - base_height: -100.0 - action_rate: -0.005 - similar_to_default: -0.1 - contact: 0.24 - swing_feet_z: 4.0 - tracking_sigma: 0.25 - base_height_target: 0.3 + events: + pd_gains: null diff --git a/conf/offpolicy/task/td3/go2_joystick_flat/motrix.yaml b/conf/offpolicy/task/td3/go2_joystick_flat/motrix.yaml index 8b32bd78a..a12a474be 100644 --- a/conf/offpolicy/task/td3/go2_joystick_flat/motrix.yaml +++ b/conf/offpolicy/task/td3/go2_joystick_flat/motrix.yaml @@ -1,4 +1,8 @@ # @package _global_ +defaults: + - /task/go2_joystick_flat/base + - _self_ + training: task_name: Go2JoystickFlat sim_backend: motrix @@ -12,22 +16,10 @@ algo: replay_buffer_n: 1024 env: commands: - vel_limit: - - [0.5, 0.0, 0.0] - - [0.5, 0.0, 0.0] - domain_rand: - randomize_kp: false - randomize_kd: false -reward: - scales: - tracking_lin_vel: 1.0 - tracking_ang_vel: 0.2 - lin_vel_z: -5.0 - ang_vel_xy: -0.1 - base_height: -100.0 - action_rate: -0.005 - similar_to_default: -0.1 - contact: 0.24 - swing_feet_z: 4.0 - tracking_sigma: 0.25 - base_height_target: 0.3 + twist: + ranges: + lin_vel_x: [0.5, 0.5] + lin_vel_y: [0.0, 0.0] + ang_vel_z: [0.0, 0.0] + events: + pd_gains: null diff --git a/conf/ppo/task/go2_joystick_flat/base.yaml b/conf/ppo/task/go2_joystick_flat/base.yaml new file mode 100644 index 000000000..3c51d7605 --- /dev/null +++ b/conf/ppo/task/go2_joystick_flat/base.yaml @@ -0,0 +1,237 @@ +# @package _global_ +# Canonical Go2 flat Manager-Based task declaration. Backend owner leaves inherit +# this file and only override backend/algo tuning or explicitly disabled terms. +env: + scene: + model_file: src/unilab/assets/robots/go2/scene_flat.xml + default_keyframe_name: home + entities: + robot: + root_body_name: base + joint_names: + - FL_hip_joint + - FL_thigh_joint + - FL_calf_joint + - FR_hip_joint + - FR_thigh_joint + - FR_calf_joint + - RL_hip_joint + - RL_thigh_joint + - RL_calf_joint + - RR_hip_joint + - RR_thigh_joint + - RR_calf_joint + actuator_names: + - FR_hip + - FR_thigh + - FR_calf + - FL_hip + - FL_thigh + - FL_calf + - RR_hip + - RR_thigh + - RR_calf + - RL_hip + - RL_thigh + - RL_calf + sim_dt: 0.01 + ctrl_dt: 0.02 + max_episode_seconds: 20.0 + observations: + policy: + _target_: unilab.managers.ObservationGroupCfg + terms: + base_ang_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.builtin_sensor + params: + sensor_name: gyro + projected_gravity: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.projected_gravity_from_sensor + params: + sensor_name: upvector + joint_pos: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_pos_rel + joint_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_vel_rel + actions: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.last_action + command: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.generated_commands + params: + command_name: twist + gait_phase: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.locomotion.common.manager_terms.quadruped_gait_phase + params: + frequency: 2.0 + critic: + _target_: unilab.managers.ObservationGroupCfg + terms: + base_ang_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.builtin_sensor + params: + sensor_name: gyro + projected_gravity: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.projected_gravity_from_sensor + params: + sensor_name: upvector + joint_pos: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_pos_rel + joint_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_vel_rel + actions: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.last_action + command: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.generated_commands + params: + command_name: twist + gait_phase: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.locomotion.common.manager_terms.quadruped_gait_phase + params: + frequency: 2.0 + base_lin_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.builtin_sensor + params: + sensor_name: local_linvel + actions: + joint_pos: + _target_: unilab.envs.mdp.JointPositionActionCfg + entity_name: robot + actuator_names: [".*"] + scale: 0.25 + use_default_offset: true + commands: + twist: + _target_: unilab.envs.mdp.UniformVelocityCommandCfg + entity_name: robot + resampling_time_range: [20.0, 20.0] + heading_command: false + heading_control_stiffness: 0.5 + rel_standing_envs: 0.0 + rel_heading_envs: 0.0 + rel_world_envs: 0.0 + rel_forward_envs: 0.0 + init_velocity_prob: 0.0 + ranges: + lin_vel_x: [-0.6, 1.0] + lin_vel_y: [-0.4, 0.4] + ang_vel_z: [-0.8, 0.8] + events: + reset_scene_to_default: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.reset_scene_to_default + mode: reset + reset_root_state_uniform: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.reset_root_state_uniform + mode: reset + params: + pose_range: + x: [-0.5, 0.5] + y: [-0.5, 0.5] + z: [0.0, 0.0] + roll: [0.0, 0.0] + pitch: [0.0, 0.0] + yaw: [-3.141592653589793, 3.141592653589793] + velocity_range: + x: [-0.5, 0.5] + y: [-0.5, 0.5] + z: [-0.5, 0.5] + roll: [-0.5, 0.5] + pitch: [-0.5, 0.5] + yaw: [-0.5, 0.5] + pd_gains: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.pd_gains + mode: reset + params: + kp_range: [31.5, 38.5] + kd_range: [0.45, 0.55] + operation: abs + terminations: + time_out: + _target_: unilab.managers.TerminationTermCfg + func: unilab.envs.mdp.time_out + time_out: true + bad_orientation: + _target_: unilab.managers.TerminationTermCfg + func: unilab.envs.mdp.bad_orientation + params: + limit_angle: 1.0471975511965976 + policy_observation_group: policy + critic_observation_group: critic + +reward: + tracking_lin_vel: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.locomotion.common.manager_terms.track_lin_vel_xy_exp + weight: 1.0 + params: + std: 0.5 + command_name: twist + tracking_ang_vel: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.locomotion.common.manager_terms.track_ang_vel_z_exp + weight: 0.2 + params: + std: 0.5 + command_name: twist + lin_vel_z: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.locomotion.common.manager_terms.lin_vel_z_l2 + weight: -5.0 + ang_vel_xy: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.locomotion.common.manager_terms.ang_vel_xy_l2 + weight: -0.1 + base_height: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.locomotion.common.manager_terms.base_height_l2 + weight: -100.0 + params: + target_height: 0.3 + action_rate: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.mdp.action_rate_l2 + weight: -0.005 + similar_to_default: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.locomotion.common.manager_terms.joint_deviation_l1 + weight: -0.1 + contact: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.locomotion.common.manager_terms.feet_phase_contact + weight: 0.24 + params: + frequency: 2.0 + sensor_names: + - FL_foot_contact + - FR_foot_contact + - RL_foot_contact + - RR_foot_contact + contact_threshold: 0.1 + stance_threshold: 0.6 + swing_feet_z: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.locomotion.common.manager_terms.feet_phase_swing_height + weight: 4.0 + params: + frequency: 2.0 + sensor_names: [FL_pos, FR_pos, RL_pos, RR_pos] + target_height: 0.1 + kernel: 0.01 + swing_start: 0.6 diff --git a/conf/ppo/task/go2_joystick_flat/drake.yaml b/conf/ppo/task/go2_joystick_flat/drake.yaml index e246865f7..9f200f570 100644 --- a/conf/ppo/task/go2_joystick_flat/drake.yaml +++ b/conf/ppo/task/go2_joystick_flat/drake.yaml @@ -1,4 +1,8 @@ # @package _global_ +defaults: + - /task/go2_joystick_flat/base + - _self_ + training: task_name: Go2JoystickFlat sim_backend: drake @@ -10,6 +14,8 @@ algo: obs_groups: actor: - actor + critic: + - critic policy: init_noise_std: 0.5 algorithm: @@ -19,23 +25,5 @@ algo: env: drake_backend_mode: batch drake_nthread: 0 - scene: - model_file: src/unilab/assets/robots/go2/scene_flat.xml - domain_rand: - randomize_kp: false - randomize_kd: false - push_robots: false - -reward: - scales: - tracking_lin_vel: 1.0 - tracking_ang_vel: 0.2 - lin_vel_z: -5.0 - ang_vel_xy: -0.1 - base_height: -100.0 - action_rate: -0.005 - similar_to_default: -0.1 - contact: 0.24 - swing_feet_z: 4.0 - tracking_sigma: 0.25 - base_height_target: 0.3 + events: + pd_gains: null diff --git a/conf/ppo/task/go2_joystick_flat/motrix.yaml b/conf/ppo/task/go2_joystick_flat/motrix.yaml index aca13216d..84282c9cb 100644 --- a/conf/ppo/task/go2_joystick_flat/motrix.yaml +++ b/conf/ppo/task/go2_joystick_flat/motrix.yaml @@ -1,4 +1,8 @@ # @package _global_ +defaults: + - /task/go2_joystick_flat/base + - _self_ + training: task_name: Go2JoystickFlat sim_backend: motrix @@ -9,6 +13,8 @@ algo: obs_groups: actor: - actor + critic: + - critic policy: init_noise_std: 0.5 algorithm: @@ -16,25 +22,13 @@ algo: entropy_coef: 1.0e-3 env: commands: - vel_limit: - - [0.5, 0.0, 0.0] - - [0.5, 0.0, 0.0] - domain_rand: - randomize_kp: false - randomize_kd: false -reward: - scales: - tracking_lin_vel: 1.0 - tracking_ang_vel: 0.2 - lin_vel_z: -5.0 - ang_vel_xy: -0.1 - base_height: -100.0 - action_rate: -0.005 - similar_to_default: -0.1 - contact: 0.24 - swing_feet_z: 4.0 - tracking_sigma: 0.25 - base_height_target: 0.3 + twist: + ranges: + lin_vel_x: [0.5, 0.5] + lin_vel_y: [0.0, 0.0] + ang_vel_z: [0.0, 0.0] + events: + pd_gains: null play_profile: enabled: true env: diff --git a/conf/ppo/task/go2_joystick_flat/mujoco.yaml b/conf/ppo/task/go2_joystick_flat/mujoco.yaml index 35cf09bed..4201cf978 100644 --- a/conf/ppo/task/go2_joystick_flat/mujoco.yaml +++ b/conf/ppo/task/go2_joystick_flat/mujoco.yaml @@ -1,4 +1,8 @@ # @package _global_ +defaults: + - /task/go2_joystick_flat/base + - _self_ + training: task_name: Go2JoystickFlat sim_backend: mujoco @@ -9,24 +13,13 @@ algo: obs_groups: actor: - actor + critic: + - critic policy: init_noise_std: 0.5 algorithm: learning_rate: 3.0e-4 entropy_coef: 1.0e-3 -reward: - scales: - tracking_lin_vel: 1.0 - tracking_ang_vel: 0.2 - lin_vel_z: -5.0 - ang_vel_xy: -0.1 - base_height: -100.0 - action_rate: -0.005 - similar_to_default: -0.1 - contact: 0.24 - swing_feet_z: 4.0 - tracking_sigma: 0.25 - base_height_target: 0.3 play_profile: enabled: true env: diff --git a/src/unilab/envs/locomotion/go2/joystick.py b/src/unilab/envs/locomotion/go2/joystick.py index 5fdd87da8..98f0bbefc 100644 --- a/src/unilab/envs/locomotion/go2/joystick.py +++ b/src/unilab/envs/locomotion/go2/joystick.py @@ -22,6 +22,10 @@ TerrainSpawnManager, ) from unilab.envs.locomotion.go2.base import Go2BaseCfg, Go2BaseEnv +from unilab.envs.manager_based_rl_env import ( + ManagerBasedRlEnvCfg, + make_manager_based_rl_env, +) @dataclass @@ -54,7 +58,6 @@ class JoystickSensor(Sensor): feet_pos = ["FL_pos", "FR_pos", "RL_pos", "RR_pos"] -@registry.envcfg("Go2JoystickFlat") @dataclass class Go2JoystickCfg(Go2BaseCfg): scene: SceneCfg = field( @@ -91,9 +94,6 @@ def _compute_reset_obs( ) -@registry.env("Go2JoystickFlat", sim_backend="mujoco") -@registry.env("Go2JoystickFlat", sim_backend="motrix") -@registry.env("Go2JoystickFlat", sim_backend="drake") class Go2WalkTask(Go2BaseEnv): _cfg: Go2JoystickCfg @@ -316,3 +316,12 @@ def _reward_contact(self, ctx: RewardContext) -> np.ndarray: is_contact = (self.feet_phase[:, i] < 0.6) | (self.gait_frequency < 1.0e-8) res += (contact[:, i] == is_contact).astype(np.float32) return res / len(self._cfg.sensor.feet_force) + + +# Go2JoystickCfg and Go2WalkTask remain as implementation bases for legacy rough +# and A2 tasks. The flat production identity is Hydra-owned and uses only the +# generic Manager-Based config/runtime factories. +registry.register_env_config("Go2JoystickFlat", ManagerBasedRlEnvCfg) +registry.register_env("Go2JoystickFlat", make_manager_based_rl_env, sim_backend="mujoco") +registry.register_env("Go2JoystickFlat", make_manager_based_rl_env, sim_backend="motrix") +registry.register_env("Go2JoystickFlat", make_manager_based_rl_env, sim_backend="drake") diff --git a/src/unilab/envs/locomotion/go2/manager_based_cfg.py b/src/unilab/envs/locomotion/go2/manager_based_cfg.py deleted file mode 100644 index 3da473da3..000000000 --- a/src/unilab/envs/locomotion/go2/manager_based_cfg.py +++ /dev/null @@ -1,232 +0,0 @@ -"""Unregistered Manager-Based configuration for the Go2 flat pilot.""" - -from __future__ import annotations - -import math - -from unilab.assets import ASSETS_ROOT_PATH -from unilab.base.entity import EntityCfg -from unilab.base.scene import SceneCfg -from unilab.envs import ManagerBasedRlEnvCfg, mdp -from unilab.envs.locomotion.common import manager_terms -from unilab.managers import ( - EventTermCfg, - ObservationGroupCfg, - ObservationTermCfg, - RewardTermCfg, - TerminationTermCfg, -) - -_JOINT_NAMES = ( - "FL_hip_joint", - "FL_thigh_joint", - "FL_calf_joint", - "FR_hip_joint", - "FR_thigh_joint", - "FR_calf_joint", - "RL_hip_joint", - "RL_thigh_joint", - "RL_calf_joint", - "RR_hip_joint", - "RR_thigh_joint", - "RR_calf_joint", -) -_ACTUATOR_NAMES = ( - "FR_hip", - "FR_thigh", - "FR_calf", - "FL_hip", - "FL_thigh", - "FL_calf", - "RR_hip", - "RR_thigh", - "RR_calf", - "RL_hip", - "RL_thigh", - "RL_calf", -) -_FOOT_CONTACT_SENSORS = ( - "FL_foot_contact", - "FR_foot_contact", - "RL_foot_contact", - "RR_foot_contact", -) -_FOOT_POSITION_SENSORS = ("FL_pos", "FR_pos", "RL_pos", "RR_pos") -_GAIT_FREQUENCY = 2.0 - - -def make_go2_joystick_flat_manager_cfg() -> ManagerBasedRlEnvCfg: - """Build the NumPy Manager-Based equivalent of the legacy Go2 flat task. - - The factory is intentionally not registered. It proves the task-owned community - config surface without changing the production Go2 registry or Hydra owners. - """ - policy_terms: dict[str, ObservationTermCfg | None] = { - "base_ang_vel": ObservationTermCfg( - func=mdp.builtin_sensor, - params={"sensor_name": "gyro"}, - ), - "projected_gravity": ObservationTermCfg( - func=mdp.projected_gravity_from_sensor, - params={"sensor_name": "upvector"}, - ), - "joint_pos": ObservationTermCfg(func=mdp.joint_pos_rel), - "joint_vel": ObservationTermCfg(func=mdp.joint_vel_rel), - "actions": ObservationTermCfg(func=mdp.last_action), - "command": ObservationTermCfg( - func=mdp.generated_commands, - params={"command_name": "twist"}, - ), - "gait_phase": ObservationTermCfg( - func=manager_terms.quadruped_gait_phase, - params={"frequency": _GAIT_FREQUENCY}, - ), - } - critic_terms: dict[str, ObservationTermCfg | None] = { - **policy_terms, - "base_lin_vel": ObservationTermCfg( - func=mdp.builtin_sensor, - params={"sensor_name": "local_linvel"}, - ), - } - - return ManagerBasedRlEnvCfg( - scene=SceneCfg( - model_file=str(ASSETS_ROOT_PATH / "robots" / "go2" / "scene_flat.xml"), - entities={ - "robot": EntityCfg( - root_body_name="base", - joint_names=_JOINT_NAMES, - actuator_names=_ACTUATOR_NAMES, - ) - }, - default_keyframe_name="home", - ), - sim_dt=0.01, - ctrl_dt=0.02, - max_episode_seconds=20.0, - observations={ - "policy": ObservationGroupCfg(terms=policy_terms), - "critic": ObservationGroupCfg(terms=critic_terms), - }, - actions={ - "joint_pos": mdp.JointPositionActionCfg( - entity_name="robot", - actuator_names=(".*",), - scale=0.25, - use_default_offset=True, - ) - }, - commands={ - "twist": mdp.UniformVelocityCommandCfg( - entity_name="robot", - resampling_time_range=(20.0, 20.0), - heading_command=False, - heading_control_stiffness=0.5, - rel_standing_envs=0.0, - rel_heading_envs=0.0, - rel_world_envs=0.0, - rel_forward_envs=0.0, - init_velocity_prob=0.0, - ranges=mdp.UniformVelocityCommandCfg.Ranges( - lin_vel_x=(-0.6, 1.0), - lin_vel_y=(-0.4, 0.4), - ang_vel_z=(-0.8, 0.8), - ), - ) - }, - events={ - "reset_scene_to_default": EventTermCfg( - func=mdp.reset_scene_to_default, - mode="reset", - ), - "reset_root_state_uniform": EventTermCfg( - func=mdp.reset_root_state_uniform, - mode="reset", - params={ - "pose_range": { - "x": (-0.5, 0.5), - "y": (-0.5, 0.5), - "z": (0.0, 0.0), - "roll": (0.0, 0.0), - "pitch": (0.0, 0.0), - "yaw": (-math.pi, math.pi), - }, - "velocity_range": { - "x": (-0.5, 0.5), - "y": (-0.5, 0.5), - "z": (-0.5, 0.5), - "roll": (-0.5, 0.5), - "pitch": (-0.5, 0.5), - "yaw": (-0.5, 0.5), - }, - }, - ), - "pd_gains": EventTermCfg( - func=mdp.pd_gains, - mode="reset", - params={ - "kp_range": (31.5, 38.5), - "kd_range": (0.45, 0.55), - "operation": "abs", - }, - ), - }, - rewards={ - "tracking_lin_vel": RewardTermCfg( - func=manager_terms.track_lin_vel_xy_exp, - weight=1.0, - params={"std": math.sqrt(0.25), "command_name": "twist"}, - ), - "tracking_ang_vel": RewardTermCfg( - func=manager_terms.track_ang_vel_z_exp, - weight=0.2, - params={"std": math.sqrt(0.25), "command_name": "twist"}, - ), - "lin_vel_z": RewardTermCfg(func=manager_terms.lin_vel_z_l2, weight=-5.0), - "ang_vel_xy": RewardTermCfg(func=manager_terms.ang_vel_xy_l2, weight=-0.1), - "base_height": RewardTermCfg( - func=manager_terms.base_height_l2, - weight=-100.0, - params={"target_height": 0.3}, - ), - "action_rate": RewardTermCfg(func=mdp.action_rate_l2, weight=-0.005), - "similar_to_default": RewardTermCfg( - func=manager_terms.joint_deviation_l1, - weight=-0.1, - ), - "contact": RewardTermCfg( - func=manager_terms.feet_phase_contact, - weight=0.24, - params={ - "frequency": _GAIT_FREQUENCY, - "sensor_names": _FOOT_CONTACT_SENSORS, - "contact_threshold": 0.1, - "stance_threshold": 0.6, - }, - ), - "swing_feet_z": RewardTermCfg( - func=manager_terms.feet_phase_swing_height, - weight=4.0, - params={ - "frequency": _GAIT_FREQUENCY, - "sensor_names": _FOOT_POSITION_SENSORS, - "target_height": 0.1, - "kernel": 0.01, - "swing_start": 0.6, - }, - ), - }, - terminations={ - "time_out": TerminationTermCfg(func=mdp.time_out, time_out=True), - "bad_orientation": TerminationTermCfg( - func=mdp.bad_orientation, - params={"limit_angle": math.acos(0.5)}, - ), - }, - policy_observation_group="policy", - critic_observation_group="critic", - ) - - -__all__ = ["make_go2_joystick_flat_manager_cfg"] diff --git a/tests/algos/test_appo_runner.py b/tests/algos/test_appo_runner.py index ca35227d0..4c3ceb4e1 100644 --- a/tests/algos/test_appo_runner.py +++ b/tests/algos/test_appo_runner.py @@ -7,13 +7,17 @@ from __future__ import annotations import tempfile +from pathlib import Path import pytest +from hydra import compose, initialize_config_dir +from hydra.core.global_hydra import GlobalHydra pytest.importorskip("mujoco") from unilab.algos.torch.appo.runner import APPORunner from unilab.structured_configs import APPOConfig +from unilab.training.backend_adapter import BackendAdapter @pytest.mark.slow @@ -34,7 +38,7 @@ def test_appo_runner_init_no_crash(mock_env_name): @pytest.mark.slow @pytest.mark.parametrize("env_name", ["Go2JoystickFlat"]) -def test_appo_runner_learn_two_iterations(env_name, default_go2_reward_config): +def test_appo_runner_learn_two_iterations(env_name): """APPO learn test must use a real env — DummyFlatTest is not registered in the collector subprocess (mp.spawn) so registry.make() would fail there.""" cfg = APPOConfig().to_dict() @@ -46,9 +50,15 @@ def test_appo_runner_learn_two_iterations(env_name, default_go2_reward_config): cfg["algorithm"]["num_learning_epochs"] = 1 cfg["algorithm"]["num_mini_batches"] = 2 + root_dir = Path(__file__).parents[2] + GlobalHydra.instance().clear() + with initialize_config_dir(config_dir=str(root_dir / "conf" / "appo"), version_base="1.3"): + hydra_cfg = compose("config", overrides=["task=go2_joystick_flat/mujoco"]) + env_cfg_overrides = BackendAdapter(hydra_cfg, root_dir=root_dir).build_task_env_cfg_override() + runner = APPORunner( env_name=env_name, - env_cfg_overrides={"reward_config": default_go2_reward_config}, + env_cfg_overrides=env_cfg_overrides, rl_cfg=cfg, num_envs=128, steps_per_env=8, diff --git a/tests/algos/test_rsl_rl_runner.py b/tests/algos/test_rsl_rl_runner.py index a7e2d7f84..582780ed3 100644 --- a/tests/algos/test_rsl_rl_runner.py +++ b/tests/algos/test_rsl_rl_runner.py @@ -8,9 +8,12 @@ import sys import tempfile +from pathlib import Path from typing import Any, cast import pytest +from hydra import compose, initialize_config_dir +from hydra.core.global_hydra import GlobalHydra pytest.importorskip("mujoco") rsl_rl = pytest.importorskip("rsl_rl") @@ -22,6 +25,7 @@ from unilab.base import registry from unilab.base.registry import ensure_registries from unilab.structured_configs import PPOConfig +from unilab.training.backend_adapter import BackendAdapter from unilab.training.rsl_rl import normalize_ppo_train_cfg from unilab.utils.tensor import to_torch @@ -114,7 +118,6 @@ def get_privileged_observations(self): ) def test_rsl_rl_ppo_one_iteration( env_name: str, - default_go2_reward_config, default_g1_reward_config, default_allegro_reward_config, ): @@ -122,26 +125,36 @@ def test_rsl_rl_ppo_one_iteration( from rsl_rl.runners import OnPolicyRunner if "Go2" in env_name: - reward_cfg = default_go2_reward_config num_envs = 256 + root_dir = Path(__file__).parents[2] + GlobalHydra.instance().clear() + with initialize_config_dir(config_dir=str(root_dir / "conf" / "ppo"), version_base="1.3"): + hydra_cfg = compose("config", overrides=["task=go2_joystick_flat/mujoco"]) + env_cfg_override = BackendAdapter( + hydra_cfg, root_dir=root_dir + ).build_task_env_cfg_override() elif "G1" in env_name: reward_cfg = default_g1_reward_config num_envs = 256 + env_cfg_override = {"reward_config": reward_cfg} else: reward_cfg = default_allegro_reward_config num_envs = 128 + env_cfg_override = {"reward_config": reward_cfg} env = registry.make( env_name, num_envs=num_envs, sim_backend="mujoco", - env_cfg_override={"reward_config": reward_cfg}, + env_cfg_override=env_cfg_override, ) wrapped = _RslRlVecEnvWrapper(env, device="cpu") cfg = PPOConfig() train_cfg = cfg.to_dict() train_cfg["runner"] = {"logger": "none"} + if env_name == "Go2JoystickFlat": + train_cfg["obs_groups"] = {"actor": ["actor"], "critic": ["critic"]} # Small network + short loop; large num_envs to saturate CPU train_cfg["num_steps_per_env"] = 8 train_cfg["policy"] = { diff --git a/tests/config/test_config_system.py b/tests/config/test_config_system.py index 359f43b0b..d347efe37 100644 --- a/tests/config/test_config_system.py +++ b/tests/config/test_config_system.py @@ -64,8 +64,16 @@ def _assert_reward_populated(cfg, label: str): assert hasattr(cfg, "reward"), f"{label} missing cfg.reward" reward_dict = OmegaConf.to_container(cfg.reward, resolve=True) assert isinstance(reward_dict, dict), f"{label} reward must resolve to mapping" - assert "scales" in reward_dict, f"{label} reward must contain scales" - assert len(reward_dict["scales"]) > 0, f"{label} reward.scales must be non-empty" + if "scales" in reward_dict: + assert len(reward_dict["scales"]) > 0, f"{label} reward.scales must be non-empty" + return + + assert reward_dict, f"{label} Manager-Based reward terms must be non-empty" + for term_name, term in reward_dict.items(): + assert isinstance(term, dict), f"{label} reward.{term_name} must be a mapping" + assert "_target_" in term, f"{label} reward.{term_name} must declare _target_" + assert "func" in term, f"{label} reward.{term_name} must declare func" + assert "weight" in term, f"{label} reward.{term_name} must declare weight" def _supported_task_cases() -> list[tuple[str, str, str, str, str, list[str]]]: @@ -237,8 +245,8 @@ def test_offpolicy_td3_go2_joystick_flat_motrix_composes(): assert cfg.algo.tau == pytest.approx(0.1) assert cfg.algo.algo_params.weight_decay == pytest.approx(0.1) assert cfg.algo.algo_params.policy_noise == pytest.approx(0.2) - assert cfg.reward.scales.tracking_lin_vel == pytest.approx(1.0) - assert cfg.reward.base_height_target == pytest.approx(0.3) + assert cfg.reward.tracking_lin_vel.weight == pytest.approx(1.0) + assert cfg.reward.base_height.params.target_height == pytest.approx(0.3) def test_offpolicy_td3_go1_joystick_flat_motrix_composes(): @@ -402,8 +410,10 @@ def test_ppo_go2_motrix_preserves_backend_env_overrides(): assert cfg.algo.num_envs == 1024 assert cfg.algo.empirical_normalization is True - assert cfg.env.domain_rand.randomize_kp is False - assert cfg.env.domain_rand.randomize_kd is False + assert cfg.env.events.pd_gains is None + assert cfg.env.commands.twist.ranges.lin_vel_x == [0.5, 0.5] + assert cfg.env.commands.twist.ranges.lin_vel_y == [0.0, 0.0] + assert cfg.env.commands.twist.ranges.ang_vel_z == [0.0, 0.0] def test_ppo_go2w_mujoco_uses_motor_owner_dr_path(): @@ -523,13 +533,28 @@ def test_offpolicy_flashsac_go2_joystick_mujoco_enables_full_dr_stack(): assert mujoco_cfg.training.task_name == "Go2JoystickFlat" assert mujoco_cfg.training.sim_backend == "mujoco" - assert mujoco_cfg.env.domain_rand.randomize_kp is True - assert mujoco_cfg.env.domain_rand.randomize_kd is True - assert mujoco_cfg.env.domain_rand.randomize_base_mass is True - assert mujoco_cfg.env.domain_rand.random_com is True - assert mujoco_cfg.env.domain_rand.randomize_gravity is True - assert mujoco_cfg.env.domain_rand.push_robots is True - assert mujoco_cfg.env.noise_config.level == pytest.approx(1.0) + assert mujoco_cfg.env.events.pd_gains.func == "unilab.envs.mdp.pd_gains" + assert ( + mujoco_cfg.env.events.randomize_rigid_body_mass.func + == "unilab.envs.mdp.randomize_rigid_body_mass" + ) + assert ( + mujoco_cfg.env.events.randomize_rigid_body_com.func + == "unilab.envs.mdp.randomize_rigid_body_com" + ) + assert ( + mujoco_cfg.env.events.randomize_physics_scene_gravity.func + == "unilab.envs.mdp.randomize_physics_scene_gravity" + ) + assert ( + mujoco_cfg.env.events.push_by_setting_velocity.func + == "unilab.envs.mdp.push_by_setting_velocity" + ) + assert mujoco_cfg.env.events.push_by_setting_velocity.mode == "interval" + assert mujoco_cfg.env.events.push_by_setting_velocity.is_global_time is True + assert mujoco_cfg.env.observations.policy.enable_corruption is True + assert mujoco_cfg.env.observations.policy.terms.joint_pos.noise.n_min == pytest.approx(-0.01) + assert mujoco_cfg.env.observations.policy.terms.joint_vel.noise.n_min == pytest.approx(-0.1) def test_cli_override_beats_task_defaults(): diff --git a/tests/config/test_locomotion_params.py b/tests/config/test_locomotion_params.py index 6cb5e7c7c..376ce31a2 100644 --- a/tests/config/test_locomotion_params.py +++ b/tests/config/test_locomotion_params.py @@ -220,8 +220,8 @@ def test_offpolicy_flashsac_go2_task_overrides(): assert cfg.algo.tau == pytest.approx(0.05) assert cfg.algo.replay_buffer_n == 4096 assert cfg.algo.updates_per_step == 2 - assert cfg.reward.scales.swing_feet_z == pytest.approx(4.0) - assert cfg.env.control_config.action_scale == pytest.approx(0.4) + assert cfg.reward.swing_feet_z.weight == pytest.approx(4.0) + assert cfg.env.actions.joint_pos.scale == pytest.approx(0.4) def test_go2_joystick_rough_uses_terrain_generator(): diff --git a/tests/config/test_reward_injection.py b/tests/config/test_reward_injection.py index 397093136..78765e3cb 100644 --- a/tests/config/test_reward_injection.py +++ b/tests/config/test_reward_injection.py @@ -38,8 +38,9 @@ def test_resolve_reward_dict_reads_task_reward(): reward_dict = resolve_reward_dict(cfg) - assert reward_dict["scales"]["tracking_lin_vel"] == 1.0 - assert reward_dict["scales"]["tracking_ang_vel"] == 0.2 + assert reward_dict["tracking_lin_vel"]["weight"] == 1.0 + assert reward_dict["tracking_ang_vel"]["weight"] == 0.2 + assert reward_dict["tracking_lin_vel"]["func"].endswith("track_lin_vel_xy_exp") def test_reward_config_conversion(): diff --git a/tests/envs/locomotion/go2/test_manager_based_cfg.py b/tests/envs/locomotion/go2/test_manager_based_cfg.py index 221ea5c6c..8948fb1d0 100644 --- a/tests/envs/locomotion/go2/test_manager_based_cfg.py +++ b/tests/envs/locomotion/go2/test_manager_based_cfg.py @@ -1,23 +1,25 @@ -"""Go2 flat Manager-Based config and real-runtime fixture tests.""" +"""Hydra-owned production contract for the Go2 flat Manager-Based task.""" from __future__ import annotations -import math -from typing import Any, TypeVar, cast +from collections.abc import Mapping, Sequence +from dataclasses import fields, is_dataclass +from pathlib import Path +from typing import Any import numpy as np import pytest +from hydra import compose, initialize_config_dir +from hydra.core.global_hydra import GlobalHydra +from omegaconf import DictConfig, OmegaConf -from unilab.base.backend import create_backend, env_backend_kwargs -from unilab.base.entity import EntityCfg -from unilab.base.np_env import NpEnvState -from unilab.envs import ManagerBasedRlEnv, mdp -from unilab.envs.locomotion.common import manager_terms -from unilab.envs.locomotion.go2.manager_based_cfg import ( - make_go2_joystick_flat_manager_cfg, -) -from unilab.envs.mdp import JointPositionAction, JointPositionActionCfg -from unilab.managers import EventTermCfg, SceneEntityCfg +from unilab.base import registry +from unilab.base.config_materialization import apply_cfg_overrides +from unilab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg, mdp +from unilab.training.backend_adapter import BackendAdapter + +ROOT_DIR = Path(__file__).parents[4] +CONF_DIR = ROOT_DIR / "conf" _JOINT_NAMES = ( "FL_hip_joint", @@ -51,32 +53,154 @@ [0.0, 0.8, -1.5, 0.0, 0.8, -1.5, 0.0, 1.0, -1.5, 0.0, 1.0, -1.5], dtype=np.float32, ) -_TermCfg = TypeVar("_TermCfg") + +_OWNER_CASES = ( + pytest.param( + "ppo", + ("task=go2_joystick_flat/mujoco",), + "mujoco", + 0.25, + True, + False, + False, + id="ppo-mujoco", + ), + pytest.param( + "ppo", + ("task=go2_joystick_flat/motrix",), + "motrix", + 0.25, + False, + False, + True, + id="ppo-motrix", + ), + pytest.param( + "ppo", + ("task=go2_joystick_flat/drake",), + "drake", + 0.25, + False, + False, + False, + id="ppo-drake", + ), + pytest.param( + "appo", + ("task=go2_joystick_flat/mujoco",), + "mujoco", + 0.25, + True, + True, + False, + id="appo-mujoco", + ), + pytest.param( + "appo", + ("task=go2_joystick_flat/motrix",), + "motrix", + 0.25, + False, + True, + False, + id="appo-motrix", + ), + pytest.param( + "offpolicy", + ("algo=flashsac", "task=flashsac/go2_joystick_flat/mujoco"), + "mujoco", + 0.4, + True, + False, + False, + id="flashsac-mujoco", + ), + pytest.param( + "offpolicy", + ("algo=td3", "task=td3/go2_joystick_flat/motrix"), + "motrix", + 0.25, + False, + False, + True, + id="td3-motrix", + ), + pytest.param( + "offpolicy", + ("algo=sac", "task=sac/go2_joystick_flat/drake"), + "drake", + 0.25, + False, + False, + False, + id="sac-drake", + ), +) -def _active(terms: dict[str, _TermCfg | None]) -> dict[str, _TermCfg]: - assert all(term is not None for term in terms.values()) - return cast(dict[str, _TermCfg], terms) +def _compose(config_group: str, overrides: Sequence[str]) -> DictConfig: + GlobalHydra.instance().clear() + with initialize_config_dir(config_dir=str(CONF_DIR / config_group), version_base="1.3"): + return compose("config", overrides=list(overrides)) -def test_go2_manager_factory_preserves_legacy_config_surface() -> None: - cfg = make_go2_joystick_flat_manager_cfg() +def _materialize( + config_group: str, overrides: Sequence[str] +) -> tuple[DictConfig, ManagerBasedRlEnvCfg, dict[str, Any]]: + hydra_cfg = _compose(config_group, overrides) + env_override = BackendAdapter(hydra_cfg, root_dir=ROOT_DIR).build_task_env_cfg_override() + env_cfg = registry.materialize_env_config("Go2JoystickFlat") + assert isinstance(env_cfg, ManagerBasedRlEnvCfg) + apply_cfg_overrides(env_cfg, env_override) + env_cfg.validate() + return hydra_cfg, env_cfg, env_override - cfg.validate() - assert cfg.sim_dt == pytest.approx(0.01) - assert cfg.ctrl_dt == pytest.approx(0.02) - assert cfg.max_episode_seconds == pytest.approx(20.0) - assert cfg.policy_observation_group == "policy" - assert cfg.critic_observation_group == "critic" - assert cfg.scene is not None - assert cfg.scene.default_keyframe_name == "home" - robot = cfg.scene.entities["robot"] - assert robot.root_body_name == "base" - assert robot.joint_names == _JOINT_NAMES - assert robot.actuator_names == _ACTUATOR_NAMES +def _assert_no_omegaconf(value: Any) -> None: + assert not OmegaConf.is_config(value) + if is_dataclass(value) and not isinstance(value, type): + for item in fields(value): + _assert_no_omegaconf(getattr(value, item.name)) + elif isinstance(value, Mapping): + for key, item in value.items(): + _assert_no_omegaconf(key) + _assert_no_omegaconf(item) + elif isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + for item in value: + _assert_no_omegaconf(item) - expected_policy = [ + +@pytest.mark.parametrize( + "config_group,overrides,backend,action_scale,pd_enabled,alive_declared,fixed_command", + _OWNER_CASES, +) +def test_go2_flat_owner_materializes_complete_plain_manager_cfg( + config_group: str, + overrides: tuple[str, ...], + backend: str, + action_scale: float, + pd_enabled: bool, + alive_declared: bool, + fixed_command: bool, +) -> None: + registry.ensure_registries() + hydra_cfg, env_cfg, _ = _materialize(config_group, overrides) + + assert hydra_cfg.training.task_name == "Go2JoystickFlat" + assert hydra_cfg.training.sim_backend == backend + assert env_cfg.sim_dt == pytest.approx(0.01) + assert env_cfg.ctrl_dt == pytest.approx(0.02) + assert env_cfg.max_episode_seconds == pytest.approx(20.0) + assert env_cfg.policy_observation_group == "policy" + assert env_cfg.critic_observation_group == "critic" + assert env_cfg.scene is not None + assert env_cfg.scene.default_keyframe_name == "home" + robot_cfg = env_cfg.scene.entities["robot"] + assert robot_cfg.root_body_name == "base" + assert tuple(robot_cfg.joint_names) == _JOINT_NAMES + assert tuple(robot_cfg.actuator_names) == _ACTUATOR_NAMES + + expected_policy_terms = [ "base_ang_vel", "projected_gravity", "joint_pos", @@ -85,204 +209,103 @@ def test_go2_manager_factory_preserves_legacy_config_surface() -> None: "command", "gait_phase", ] - observations = _active(cfg.observations) - assert list(observations) == ["policy", "critic"] - policy_terms = _active(observations["policy"].terms) - critic_terms = _active(observations["critic"].terms) - expected_policy_terms = [ - ("base_ang_vel", mdp.builtin_sensor, {"sensor_name": "gyro"}), - ( - "projected_gravity", - mdp.projected_gravity_from_sensor, - {"sensor_name": "upvector"}, - ), - ("joint_pos", mdp.joint_pos_rel, {}), - ("joint_vel", mdp.joint_vel_rel, {}), - ("actions", mdp.last_action, {}), - ("command", mdp.generated_commands, {"command_name": "twist"}), - ("gait_phase", manager_terms.quadruped_gait_phase, {"frequency": 2.0}), - ] - assert list(policy_terms) == expected_policy - assert [(name, term.func, term.params) for name, term in policy_terms.items()] == ( - expected_policy_terms - ) - assert [(name, term.func, term.params) for name, term in critic_terms.items()] == [ + assert list(env_cfg.observations) == ["policy", "critic"] + assert list(env_cfg.observations["policy"].terms) == expected_policy_terms + assert list(env_cfg.observations["critic"].terms) == [ *expected_policy_terms, - ("base_lin_vel", mdp.builtin_sensor, {"sensor_name": "local_linvel"}), + "base_lin_vel", ] + assert list(env_cfg.actions) == ["joint_pos"] + assert env_cfg.actions["joint_pos"].scale == pytest.approx(action_scale) + assert list(env_cfg.commands) == ["twist"] + assert list(env_cfg.terminations) == ["time_out", "bad_orientation"] + assert (env_cfg.events["pd_gains"] is not None) is pd_enabled + assert ("alive" in env_cfg.rewards) is alive_declared - actions = _active(cfg.actions) - assert list(actions) == ["joint_pos"] - action = actions["joint_pos"] - assert isinstance(action, JointPositionActionCfg) - assert action.entity_name == "robot" - assert action.actuator_names == (".*",) - assert action.scale == pytest.approx(0.25) - assert action.use_default_offset is True - - commands = _active(cfg.commands) - assert list(commands) == ["twist"] - command = commands["twist"] - assert isinstance(command, mdp.UniformVelocityCommandCfg) - assert command.resampling_time_range == (20.0, 20.0) - assert command.heading_command is False - assert command.heading_control_stiffness == pytest.approx(0.5) - assert command.rel_standing_envs == 0.0 - assert command.rel_heading_envs == 0.0 - assert command.rel_world_envs == 0.0 - assert command.rel_forward_envs == 0.0 - assert command.init_velocity_prob == 0.0 - assert command.ranges.lin_vel_x == (-0.6, 1.0) - assert command.ranges.lin_vel_y == (-0.4, 0.4) - assert command.ranges.ang_vel_z == (-0.8, 0.8) - assert command.ranges.heading is None - - events = _active(cfg.events) - assert list(events) == ["reset_scene_to_default", "reset_root_state_uniform", "pd_gains"] - assert events["reset_scene_to_default"].func is mdp.reset_scene_to_default - assert events["reset_scene_to_default"].mode == "reset" - assert events["reset_scene_to_default"].params == {} - root_reset = events["reset_root_state_uniform"] - assert root_reset.func is mdp.reset_root_state_uniform - assert root_reset.mode == "reset" - assert root_reset.params["pose_range"] == { - "x": (-0.5, 0.5), - "y": (-0.5, 0.5), - "z": (0.0, 0.0), - "roll": (0.0, 0.0), - "pitch": (0.0, 0.0), - "yaw": (-math.pi, math.pi), - } - assert root_reset.params["velocity_range"] == { - key: (-0.5, 0.5) for key in ("x", "y", "z", "roll", "pitch", "yaw") - } - pd_gains = events["pd_gains"] - assert pd_gains.func is mdp.pd_gains - assert pd_gains.mode == "reset" - assert pd_gains.params == { - "kp_range": (31.5, 38.5), - "kd_range": (0.45, 0.55), - "operation": "abs", - } + if action_scale == 0.25: + expected_weights = { + "tracking_lin_vel": 1.0, + "tracking_ang_vel": 0.2, + "lin_vel_z": -5.0, + "ang_vel_xy": -0.1, + "base_height": -100.0, + "action_rate": -0.005, + "similar_to_default": -0.1, + "contact": 0.24, + "swing_feet_z": 4.0, + } + if alive_declared: + expected_weights["alive"] = 0.0 + assert {name: term.weight for name, term in env_cfg.rewards.items()} == expected_weights - expected_rewards = { - "tracking_lin_vel": 1.0, - "tracking_ang_vel": 0.2, - "lin_vel_z": -5.0, - "ang_vel_xy": -0.1, - "base_height": -100.0, - "action_rate": -0.005, - "similar_to_default": -0.1, - "contact": 0.24, - "swing_feet_z": 4.0, - } - rewards = _active(cfg.rewards) - assert list(rewards) == list(expected_rewards) - assert {name: term.weight for name, term in rewards.items()} == expected_rewards - assert [(name, term.func, term.params) for name, term in rewards.items()] == [ - ( - "tracking_lin_vel", - manager_terms.track_lin_vel_xy_exp, - {"std": 0.5, "command_name": "twist"}, - ), - ( - "tracking_ang_vel", - manager_terms.track_ang_vel_z_exp, - {"std": 0.5, "command_name": "twist"}, - ), - ("lin_vel_z", manager_terms.lin_vel_z_l2, {}), - ("ang_vel_xy", manager_terms.ang_vel_xy_l2, {}), - ("base_height", manager_terms.base_height_l2, {"target_height": 0.3}), - ("action_rate", mdp.action_rate_l2, {}), - ("similar_to_default", manager_terms.joint_deviation_l1, {}), - ( - "contact", - manager_terms.feet_phase_contact, - { - "frequency": 2.0, - "sensor_names": ( - "FL_foot_contact", - "FR_foot_contact", - "RL_foot_contact", - "RR_foot_contact", - ), - "contact_threshold": 0.1, - "stance_threshold": 0.6, - }, - ), - ( - "swing_feet_z", - manager_terms.feet_phase_swing_height, - { - "frequency": 2.0, - "sensor_names": ("FL_pos", "FR_pos", "RL_pos", "RR_pos"), - "target_height": 0.1, - "kernel": 0.01, - "swing_start": 0.6, - }, - ), - ] + if fixed_command: + ranges = env_cfg.commands["twist"].ranges + assert tuple(ranges.lin_vel_x) == (0.5, 0.5) + assert tuple(ranges.lin_vel_y) == (0.0, 0.0) + assert tuple(ranges.ang_vel_z) == (0.0, 0.0) - terminations = _active(cfg.terminations) - assert list(terminations) == ["time_out", "bad_orientation"] - assert terminations["time_out"].func is mdp.time_out - assert terminations["time_out"].time_out is True - assert terminations["time_out"].params == {} - assert terminations["bad_orientation"].func is mdp.bad_orientation - assert terminations["bad_orientation"].time_out is False - assert terminations["bad_orientation"].params["limit_angle"] == pytest.approx(math.pi / 3.0) + for manager_name in ("observations", "events", "rewards", "terminations"): + for term in getattr(env_cfg, manager_name).values(): + if term is None: + continue + terms = term.terms.values() if manager_name == "observations" else (term,) + for nested_term in terms: + if nested_term is None: + continue + module = nested_term.func.__module__ + assert ".backend." not in module + assert not any(name in module for name in (".mujoco", ".motrix", ".drake")) + _assert_no_omegaconf(env_cfg) -def test_go2_manager_factory_executes_on_real_mujoco() -> None: - cfg = make_go2_joystick_flat_manager_cfg() - assert cfg.scene is not None - backend = create_backend( - "mujoco", - cfg.scene, - 2, - cfg.sim_dt, - base_name="base", - add_body_sensors=True, - **env_backend_kwargs(cfg), + +def test_go2_flat_registry_has_no_legacy_config_fallback() -> None: + registry.ensure_registries() + bare_cfg = registry.materialize_env_config("Go2JoystickFlat") + + assert isinstance(bare_cfg, ManagerBasedRlEnvCfg) + assert bare_cfg.observations == {} + assert bare_cfg.actions == {} + assert bare_cfg.rewards == {} + assert registry.list_registered_envs()["Go2JoystickFlat"] == { + "config_factory": "ManagerBasedRlEnvCfg", + "available_backends": ["mujoco", "motrix", "drake"], + } + for legacy_override in ( + {"reward_config": {}}, + {"domain_rand": {"randomize_kp": True}}, + {"control_config": {"action_scale": 0.4}}, + ): + with pytest.raises(ValueError, match="has no attribute"): + apply_cfg_overrides(ManagerBasedRlEnvCfg(), legacy_override) + + +@pytest.mark.parametrize( + ("backend", "owner"), + (("mujoco", "task=go2_joystick_flat/mujoco"), ("motrix", "task=go2_joystick_flat/motrix")), +) +def test_go2_flat_registry_executes_real_manager_runtime(backend: str, owner: str) -> None: + registry.ensure_registries() + hydra_cfg, _, env_override = _materialize("ppo", (owner,)) + env = registry.make( + str(hydra_cfg.training.task_name), + sim_backend=backend, + env_cfg_override=env_override, + num_envs=2, ) - env = ManagerBasedRlEnv(cfg, backend, 2) try: + assert isinstance(env, ManagerBasedRlEnv) assert env.obs_groups_spec == {"obs": 49, "critic": 52} assert env.action_space.shape == (12,) - assert env.observation_manager.active_terms == { - "policy": [ - "base_ang_vel", - "projected_gravity", - "joint_pos", - "joint_vel", - "actions", - "command", - "gait_phase", - ], - "critic": [ - "base_ang_vel", - "projected_gravity", - "joint_pos", - "joint_vel", - "actions", - "command", - "gait_phase", - "base_lin_vel", - ], - } - assert env.reward_manager.active_terms == list(cfg.rewards) - assert env.termination_manager.active_terms == list(cfg.terminations) - assert env.event_manager.active_terms["reset"] == list(cfg.events) - action = env.action_manager.get_term("joint_pos") - assert isinstance(action, JointPositionAction) assert action.target_names == list(_JOINT_NAMES) np.testing.assert_allclose(action.offset, np.broadcast_to(_HOME_JOINT_POS, (2, 12))) obs, info = env.reset(seed=7) - assert set(obs) == {"obs", "critic"} - assert obs["obs"].shape == (2, 49) - assert obs["critic"].shape == (2, 52) + assert {name: value.shape for name, value in obs.items()} == { + "obs": (2, 49), + "critic": (2, 52), + } assert isinstance(info, dict) np.testing.assert_allclose( env.scene["robot"].data.default_joint_pos, @@ -294,123 +317,69 @@ def test_go2_manager_factory_executes_on_real_mujoco() -> None: ) state = env.step(np.zeros((2, 12), dtype=np.float32)) - assert isinstance(state, NpEnvState) - assert state.obs["obs"].shape == (2, 49) - assert state.obs["critic"].shape == (2, 52) + assert {name: value.shape for name, value in state.obs.items()} == { + "obs": (2, 49), + "critic": (2, 52), + } for value in (*state.obs.values(), state.reward): assert isinstance(value, np.ndarray) assert np.isfinite(value).all() - assert state.terminated.dtype == np.bool_ - assert state.truncated.dtype == np.bool_ finally: env.close() -def test_go2_manager_reset_randomization_mutates_real_mujoco_payload() -> None: - cfg = make_go2_joystick_flat_manager_cfg() - assert cfg.scene is not None - robot = cfg.scene.entities["robot"] - cfg.scene.entities["robot"] = EntityCfg( - root_body_name=robot.root_body_name, - joint_names=robot.joint_names, - body_names=("base",), - actuator_names=robot.actuator_names, +def test_go2_flat_flashsac_uses_canonical_manager_events_and_numpy_noise() -> None: + registry.ensure_registries() + _, env_cfg, _ = _materialize( + "offpolicy", + ("algo=flashsac", "task=flashsac/go2_joystick_flat/mujoco"), ) - asset_cfg = SceneEntityCfg("robot", body_names=("base",)) - cfg.events.update( - { - "mass": EventTermCfg( - func=mdp.randomize_rigid_body_mass, - mode="reset", - params={ - "asset_cfg": asset_cfg, - "mass_distribution_params": (1.25, 1.25), - "operation": "scale", - "recompute_inertia": False, - }, - ), - "com": EventTermCfg( - func=mdp.randomize_rigid_body_com, - mode="reset", - params={"asset_cfg": asset_cfg, "com_range": {"x": (0.02, 0.02)}}, - ), - "gravity": EventTermCfg( - func=mdp.randomize_physics_scene_gravity, - mode="reset", - params={ - "gravity_distribution_params": ([0.0, 0.0, -9.7],) * 2, - "operation": "abs", - }, - ), - } - ) - backend = create_backend( - "mujoco", - cfg.scene, - 2, - cfg.sim_dt, - base_name="base", - add_body_sensors=True, - **env_backend_kwargs(cfg), - ) - base_id = int(backend.get_body_ids(("base",))[0]) - default_mass = backend.get_body_mass() - default_ipos = backend.get_body_ipos() - env = ManagerBasedRlEnv(cfg, backend, 2) - try: - env.reset(seed=31) - pool = cast(Any, backend)._pool - assert pool is not None - for env_id in range(2): - mass = pool.get_field(env_id, "body_mass") - ipos = pool.get_field(env_id, "body_ipos").reshape(-1, 3) - gravity = pool.get_field(env_id, "gravity") - assert mass[base_id] == pytest.approx(default_mass[base_id] * 1.25) - np.testing.assert_allclose(ipos[base_id], default_ipos[base_id] + [0.02, 0.0, 0.0]) - np.testing.assert_allclose(gravity, [0.0, 0.0, -9.7]) - finally: - env.close() + assert env_cfg.scene is not None + assert env_cfg.scene.entities["robot"].body_names == ["base"] + assert list(env_cfg.events) == [ + "reset_scene_to_default", + "reset_root_state_uniform", + "pd_gains", + "randomize_rigid_body_mass", + "randomize_rigid_body_com", + "randomize_physics_scene_gravity", + "push_by_setting_velocity", + ] + push = env_cfg.events["push_by_setting_velocity"] + assert push.func is mdp.push_by_setting_velocity + assert push.mode == "interval" + assert push.interval_range_s == [15.0, 15.0] + assert push.is_global_time is True -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() - 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() + for group_name in ("policy", "critic"): + group = env_cfg.observations[group_name] + assert group.enable_corruption is True + assert type(group.terms["joint_pos"].noise).__name__ == "UniformNoiseCfg" + assert group.terms["joint_pos"].noise.n_min == pytest.approx(-0.01) + assert group.terms["joint_pos"].noise.n_max == pytest.approx(0.01) + assert type(group.terms["joint_vel"].noise).__name__ == "UniformNoiseCfg" + assert group.terms["joint_vel"].noise.n_min == pytest.approx(-0.1) + assert group.terms["joint_vel"].noise.n_max == pytest.approx(0.1) + + assert env_cfg.rewards["tracking_lin_vel"].params["std"] == pytest.approx(0.4**0.5) + assert env_cfg.rewards["base_height"].weight == pytest.approx(-20.0) + assert env_cfg.rewards["contact"].weight == pytest.approx(1.5) + + +def test_go2_flat_drake_missing_dependency_is_explicit() -> None: + from unilab.base.backend.drake.backend import ensure_drake_batch_available + + available, _ = ensure_drake_batch_available() + if available: + pytest.skip("DrakeUni is installed; missing-dependency behavior is not applicable") + + registry.ensure_registries() + hydra_cfg, _, env_override = _materialize("ppo", ("task=go2_joystick_flat/drake",)) + with pytest.raises(ImportError, match="[Dd]rake"): + registry.make( + str(hydra_cfg.training.task_name), + sim_backend="drake", + env_cfg_override=env_override, + num_envs=1, + ) diff --git a/tests/envs/test_env_configs.py b/tests/envs/test_env_configs.py index 55083925d..746fa7a4f 100644 --- a/tests/envs/test_env_configs.py +++ b/tests/envs/test_env_configs.py @@ -2049,7 +2049,6 @@ def test_g1_motion_tracking_clip_end_does_not_override_true_termination(): _STANDARD_ENVS = [ "Go1JoystickFlat", "Go1JoystickRough", - "Go2JoystickFlat", "Go2WJoystickFlat", "Go2WJoystickRough", "G1WalkFlat", @@ -2167,42 +2166,6 @@ def test_go1_env_initializes_kp_kd_into_pool(default_go1_reward_config): env.close() -def test_go2_env_initializes_kp_kd_into_pool(): - _require_mujoco_runtime() - ensure_registries() - from unilab.base import registry - from unilab.envs.locomotion.go2.joystick import RewardConfig - - env = cast( - Any, - registry.make( - "Go2JoystickFlat", - num_envs=2, - sim_backend="mujoco", - env_cfg_override={ - "reward_config": RewardConfig( - scales={ - "tracking_lin_vel": 1.0, - "tracking_ang_vel": 0.2, - "lin_vel_z": -5.0, - "ang_vel_xy": -0.02, - "base_height": -100.0, - "action_rate": -0.005, - "similar_to_default": -0.1, - }, - tracking_sigma=0.25, - base_height_target=0.3, - ), - "control_config": {"Kp": 18.0, "Kd": 0.9}, - }, - ), - ) - try: - _assert_mujoco_position_gains(env, kp=18.0, kd=0.9) - finally: - env.close() - - def test_allegro_env_initializes_kp_kd_into_pool(default_allegro_reward_config): _require_mujoco_runtime() ensure_registries() @@ -2328,41 +2291,6 @@ def test_g1_motion_tracking_deploy_reset_and_step_mujoco(): env.close() -def test_go2_mujoco_reset_applies_kp_kd_domain_randomization(default_go2_reward_config): - _require_mujoco_runtime() - ensure_registries() - - from unilab.base import registry - - env = cast( - Any, - registry.make( - "Go2JoystickFlat", - num_envs=4, - sim_backend="mujoco", - env_cfg_override={"reward_config": default_go2_reward_config}, - ), - ) - try: - env.init_state() - backend = env._backend - kp = np.stack([backend._pool.get_field(i, "kp") for i in range(env.num_envs)]) - kd = np.stack([backend._pool.get_field(i, "kd") for i in range(env.num_envs)]) - base_kp = float(env.cfg.control_config.Kp) - base_kd = float(env.cfg.control_config.Kd) - - assert np.unique(np.round(kp[:, 0], 6)).size > 1 - assert np.unique(np.round(kd[:, 0], 6)).size > 1 - np.testing.assert_allclose(kp / base_kp, np.broadcast_to(kp[:, :1] / base_kp, kp.shape)) - np.testing.assert_allclose(kd / base_kd, np.broadcast_to(kd[:, :1] / base_kd, kd.shape)) - assert np.all(kp >= base_kp * 0.9) - assert np.all(kp <= base_kp * 1.1) - assert np.all(kd >= base_kd * 0.9) - assert np.all(kd <= base_kd * 1.1) - finally: - env.close() - - def test_go2w_mujoco_keeps_kp_kd_out_of_backend_position_actuator_path(): _require_mujoco_runtime() ensure_registries() diff --git a/tests/scripts/test_train_scripts.py b/tests/scripts/test_train_scripts.py index cb495cf6c..99f97bc3d 100644 --- a/tests/scripts/test_train_scripts.py +++ b/tests/scripts/test_train_scripts.py @@ -675,11 +675,13 @@ def test_ppo_task_go2_aligns_mujoco_with_motrix_defaults(): cfg = _ppo_cfg(["task=go2_joystick_flat/mujoco"]) assert cfg.algo.num_envs == 1024 - assert cfg.reward.scales.tracking_lin_vel == pytest.approx(1.0) - assert cfg.reward.scales.tracking_ang_vel == pytest.approx(0.2) - assert cfg.reward.scales.lin_vel_z == pytest.approx(-5.0) - assert cfg.reward.scales.ang_vel_xy == pytest.approx(-0.1) + assert cfg.reward.tracking_lin_vel.weight == pytest.approx(1.0) + assert cfg.reward.tracking_ang_vel.weight == pytest.approx(0.2) + assert cfg.reward.lin_vel_z.weight == pytest.approx(-5.0) + assert cfg.reward.ang_vel_xy.weight == pytest.approx(-0.1) assert cfg.algo.empirical_normalization is True + assert cfg.algo.obs_groups.actor == ["actor"] + assert cfg.algo.obs_groups.critic == ["critic"] assert cfg.algo.policy.init_noise_std == pytest.approx(0.5) assert cfg.algo.algorithm.learning_rate == pytest.approx(3.0e-4) assert cfg.algo.algorithm.entropy_coef == pytest.approx(1.0e-3) @@ -707,9 +709,8 @@ def test_ppo_go2_drake_batch_config_matches_go2_training_defaults(): assert cfg.env.drake_backend_mode == "batch" assert cfg.env.drake_nthread == 0 assert cfg.env.scene.model_file == "src/unilab/assets/robots/go2/scene_flat.xml" - assert cfg.env.domain_rand.randomize_kp is False - assert cfg.env.domain_rand.randomize_kd is False - assert cfg.reward.scales.contact == pytest.approx(0.24) + assert cfg.env.events.pd_gains is None + assert cfg.reward.contact.weight == pytest.approx(0.24) def test_build_ppo_env_cfg_override_go1_motrix( @@ -790,12 +791,11 @@ def test_build_ppo_env_cfg_override_applies_go2_motrix_reward( env_cfg_override = mod.build_ppo_env_cfg_override(cfg) - assert cfg.reward.scales.tracking_lin_vel == pytest.approx(1.0) + assert cfg.reward.tracking_lin_vel.weight == pytest.approx(1.0) assert cfg.algo.num_envs == 1024 - assert env_cfg_override["domain_rand"]["randomize_kp"] is False - assert env_cfg_override["domain_rand"]["randomize_kd"] is False - assert env_cfg_override["reward_config"]["scales"]["tracking_lin_vel"] == pytest.approx(1.0) - assert env_cfg_override["reward_config"]["scales"]["tracking_ang_vel"] == pytest.approx(0.2) + assert env_cfg_override["events"]["pd_gains"] is None + assert env_cfg_override["rewards"]["tracking_lin_vel"]["weight"] == pytest.approx(1.0) + assert env_cfg_override["rewards"]["tracking_ang_vel"]["weight"] == pytest.approx(0.2) def test_build_ppo_env_cfg_override_allegro_mujoco( From 470be21d31ab01bf9e881628a136506ee702b0ea Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 20:14:39 +0800 Subject: [PATCH 034/169] refactor(tasks): establish registry bootstrap boundary [skip ci] Part of #1112 and #1042. Local make test-all passed; remote child CI intentionally skipped per maintainer workflow. --- src/unilab/base/registry.py | 6 +--- src/unilab/tasks/__init__.py | 24 +++++++++++++ tests/tasks/test_package_boundary.py | 53 ++++++++++++++++++++++++++++ 3 files changed, 78 insertions(+), 5 deletions(-) create mode 100644 src/unilab/tasks/__init__.py create mode 100644 tests/tasks/test_package_boundary.py diff --git a/src/unilab/base/registry.py b/src/unilab/base/registry.py index cf5089321..2ad66f781 100644 --- a/src/unilab/base/registry.py +++ b/src/unilab/base/registry.py @@ -42,11 +42,7 @@ def __call__( _SUPPORTED_SIM_BACKENDS = ("mujoco", "mjwarp", "motrix", "drake") _DEFAULT_SIM_BACKEND_ORDER: tuple[str, ...] = ("mujoco", "motrix") _REGISTRY_MODULES_ATTR = "__unilab_registry_modules__" -_DEFAULT_REGISTRY_PACKAGES = ( - "unilab.envs.locomotion", - "unilab.envs.manipulation", - "unilab.envs.motion_tracking", -) +_DEFAULT_REGISTRY_PACKAGES = ("unilab.tasks",) # Environment variable used to extend ensure_registries() with extra packages. # Mainly intended for test setups that need to ship a fixture-only registry into # spawn subprocesses (which do not inherit pytest conftest state). diff --git a/src/unilab/tasks/__init__.py b/src/unilab/tasks/__init__.py new file mode 100644 index 000000000..17a2ac2fa --- /dev/null +++ b/src/unilab/tasks/__init__.py @@ -0,0 +1,24 @@ +"""Production task registry bootstrap. + +Concrete task implementations are moving from :mod:`unilab.envs` into this +package under issue #1112. Until each task family moves, this explicit list +records its legacy module as the last remaining consumer of that path. The +registry imports these leaf modules directly, so registration stays explicit +and deterministic throughout the migration. +""" + +__unilab_registry_modules__ = ( + "unilab.envs.locomotion.go1", + "unilab.envs.locomotion.go2", + "unilab.envs.locomotion.go2w", + "unilab.envs.locomotion.g1", + "unilab.envs.locomotion.go2_arm", + "unilab.envs.locomotion.a2", + "unilab.envs.manipulation.allegro_inhand", + "unilab.envs.manipulation.sharpa_inhand", + "unilab.envs.manipulation.stewart", + "unilab.envs.motion_tracking.g1", + "unilab.envs.motion_tracking.x2", +) + +__all__ = ["__unilab_registry_modules__"] diff --git a/tests/tasks/test_package_boundary.py b/tests/tasks/test_package_boundary.py new file mode 100644 index 000000000..fc3156e20 --- /dev/null +++ b/tests/tasks/test_package_boundary.py @@ -0,0 +1,53 @@ +"""Task bootstrap and package dependency boundary tests.""" + +from __future__ import annotations + +import ast +from pathlib import Path + +from unilab.base import registry +from unilab.tasks import __unilab_registry_modules__ + +_REPO_ROOT = Path(__file__).resolve().parents[2] +_ENV_PACKAGE = _REPO_ROOT / "src" / "unilab" / "envs" + +_LEGACY_TASK_MODULES = ( + "unilab.envs.locomotion.go1", + "unilab.envs.locomotion.go2", + "unilab.envs.locomotion.go2w", + "unilab.envs.locomotion.g1", + "unilab.envs.locomotion.go2_arm", + "unilab.envs.locomotion.a2", + "unilab.envs.manipulation.allegro_inhand", + "unilab.envs.manipulation.sharpa_inhand", + "unilab.envs.manipulation.stewart", + "unilab.envs.motion_tracking.g1", + "unilab.envs.motion_tracking.x2", +) + + +def _imports(path: Path) -> set[str]: + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + modules: set[str] = set() + for node in ast.walk(tree): + if isinstance(node, ast.ImportFrom) and node.module: + modules.add(node.module) + elif isinstance(node, ast.Import): + modules.update(alias.name for alias in node.names) + return modules + + +def test_tasks_is_the_only_default_registry_bootstrap() -> None: + assert registry._DEFAULT_REGISTRY_PACKAGES == ("unilab.tasks",) + assert __unilab_registry_modules__ == _LEGACY_TASK_MODULES + + +def test_env_runtime_does_not_depend_on_tasks() -> None: + violations = [ + (path.relative_to(_REPO_ROOT).as_posix(), module) + for path in sorted(_ENV_PACKAGE.rglob("*.py")) + for module in sorted(_imports(path)) + if module == "unilab.tasks" or module.startswith("unilab.tasks.") + ] + + assert violations == [], "unilab.envs must not import concrete unilab.tasks modules" From b94d9a829d4a6ff3c179c2d8da264b2237c557e3 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 20:23:09 +0800 Subject: [PATCH 035/169] refactor(tasks): move A2 locomotion owner [skip ci] Part of #1112 and #1042. Final local make test-all passed; remote child CI intentionally skipped. --- src/unilab/envs/locomotion/__init__.py | 1 - src/unilab/tasks/__init__.py | 2 +- src/unilab/tasks/locomotion/__init__.py | 1 + .../{envs => tasks}/locomotion/a2/__init__.py | 0 .../{envs => tasks}/locomotion/a2/joystick.py | 1 - .../a2/test_a2_joystick_contract.py | 24 +++++++++---------- tests/tasks/test_package_boundary.py | 6 ++--- 7 files changed, 17 insertions(+), 18 deletions(-) create mode 100644 src/unilab/tasks/locomotion/__init__.py rename src/unilab/{envs => tasks}/locomotion/a2/__init__.py (100%) rename src/unilab/{envs => tasks}/locomotion/a2/joystick.py (99%) diff --git a/src/unilab/envs/locomotion/__init__.py b/src/unilab/envs/locomotion/__init__.py index 59a8b0270..a101b48c1 100644 --- a/src/unilab/envs/locomotion/__init__.py +++ b/src/unilab/envs/locomotion/__init__.py @@ -6,5 +6,4 @@ "unilab.envs.locomotion.go2w", "unilab.envs.locomotion.g1", "unilab.envs.locomotion.go2_arm", - "unilab.envs.locomotion.a2", ) diff --git a/src/unilab/tasks/__init__.py b/src/unilab/tasks/__init__.py index 17a2ac2fa..13c4f91df 100644 --- a/src/unilab/tasks/__init__.py +++ b/src/unilab/tasks/__init__.py @@ -13,7 +13,7 @@ "unilab.envs.locomotion.go2w", "unilab.envs.locomotion.g1", "unilab.envs.locomotion.go2_arm", - "unilab.envs.locomotion.a2", + "unilab.tasks.locomotion.a2", "unilab.envs.manipulation.allegro_inhand", "unilab.envs.manipulation.sharpa_inhand", "unilab.envs.manipulation.stewart", diff --git a/src/unilab/tasks/locomotion/__init__.py b/src/unilab/tasks/locomotion/__init__.py new file mode 100644 index 000000000..3439ae657 --- /dev/null +++ b/src/unilab/tasks/locomotion/__init__.py @@ -0,0 +1 @@ +"""Locomotion task packages.""" diff --git a/src/unilab/envs/locomotion/a2/__init__.py b/src/unilab/tasks/locomotion/a2/__init__.py similarity index 100% rename from src/unilab/envs/locomotion/a2/__init__.py rename to src/unilab/tasks/locomotion/a2/__init__.py diff --git a/src/unilab/envs/locomotion/a2/joystick.py b/src/unilab/tasks/locomotion/a2/joystick.py similarity index 99% rename from src/unilab/envs/locomotion/a2/joystick.py rename to src/unilab/tasks/locomotion/a2/joystick.py index e9f675fd8..59df62acc 100644 --- a/src/unilab/envs/locomotion/a2/joystick.py +++ b/src/unilab/tasks/locomotion/a2/joystick.py @@ -169,7 +169,6 @@ class A2JoystickFlatEnv(Go2WalkTask): zero-command standstill (phase freeze + gated gait rewards + standing resample) gated by A2RewardConfig.command_threshold.""" - _cfg: A2JoystickCfg _reward_cfg: A2RewardConfig def _make_dr_provider(self) -> LocomotionDRProvider: diff --git a/tests/envs/locomotion/a2/test_a2_joystick_contract.py b/tests/envs/locomotion/a2/test_a2_joystick_contract.py index 14fcc656f..0e3fbf520 100644 --- a/tests/envs/locomotion/a2/test_a2_joystick_contract.py +++ b/tests/envs/locomotion/a2/test_a2_joystick_contract.py @@ -126,7 +126,7 @@ def test_a2_home_keyframe_matches_mjlab_pose(): def test_a2_control_config_per_joint_gains(): """A2JoystickControlConfig.position_gains() yields per-joint arrays matching mjlab (calf 150/6, hip/thigh 100/4) in actuator order.""" - from unilab.envs.locomotion.a2.joystick import A2JoystickControlConfig + from unilab.tasks.locomotion.a2.joystick import A2JoystickControlConfig gains = A2JoystickControlConfig().position_gains() np.testing.assert_allclose(np.asarray(gains["kp"]), _MJLAB_KP) @@ -144,7 +144,7 @@ def test_pd_control_config_position_gains_default_is_scalar(): def test_a2_dr_provider_returns_per_joint_base_gains(): """The A2 DR provider exposes per-joint base kp/kd so randomize_kp/kd scales each joint off the correct baseline (calf off 150, not 100).""" - from unilab.envs.locomotion.a2.joystick import ( + from unilab.tasks.locomotion.a2.joystick import ( A2JoystickControlConfig, A2JoystickDomainRandomizationProvider, ) @@ -179,7 +179,7 @@ def test_a2_dr_provider_caches_friction_and_armature_baselines(): (and the floor geom id) from the backend so randomize_ground_friction / randomize_dof_armature can multiply against them. body_mass stays uncached.""" mujoco = pytest.importorskip("mujoco") - from unilab.envs.locomotion.a2.joystick import A2JoystickDomainRandomizationProvider + from unilab.tasks.locomotion.a2.joystick import A2JoystickDomainRandomizationProvider xml = ASSETS_ROOT_PATH / "robots" / "a2" / "scene_flat.xml" model = mujoco.MjModel.from_xml_path(str(xml)) @@ -215,7 +215,7 @@ def _ensure_registered() -> None: registry.ensure_registries() if not registry.contains("A2JoystickFlat"): - importlib.import_module("unilab.envs.locomotion.a2.joystick") + importlib.import_module("unilab.tasks.locomotion.a2.joystick") def test_a2_joystick_registered(): @@ -239,7 +239,7 @@ def test_a2_joystick_yaml_composes_and_targets_a2(): def _default_reward_cfg(): - from unilab.envs.locomotion.a2.joystick import A2RewardConfig + from unilab.tasks.locomotion.a2.joystick import A2RewardConfig return A2RewardConfig( scales={ @@ -340,7 +340,7 @@ def test_a2_joystick_dr_on_constructs_and_steps_finite(): randomize_body_mass stays off (base_body_mass baseline not cached). The YAML-surface is covered by test_a2_joystick_domain_rand_fully_configured.""" _skip_if_no_mujoco() - from unilab.envs.locomotion.a2.joystick import A2JoystickDomainRandConfig + from unilab.tasks.locomotion.a2.joystick import A2JoystickDomainRandConfig dr_on = A2JoystickDomainRandConfig( randomize_base_mass=True, @@ -423,7 +423,7 @@ def test_a2_joystick_domain_rand_fully_configured(): def test_a2_reward_config_declares_command_threshold(): import dataclasses - from unilab.envs.locomotion.a2.joystick import A2RewardConfig + from unilab.tasks.locomotion.a2.joystick import A2RewardConfig names = {f.name for f in dataclasses.fields(A2RewardConfig)} assert "command_threshold" in names @@ -433,7 +433,7 @@ def test_a2_reward_config_declares_command_threshold(): def test_a2_cfg_reward_config_annotation_is_a2_type(): from typing import get_type_hints - from unilab.envs.locomotion.a2.joystick import A2JoystickCfg, A2RewardConfig + from unilab.tasks.locomotion.a2.joystick import A2JoystickCfg, A2RewardConfig hints = get_type_hints(A2JoystickCfg) assert A2RewardConfig in getattr(hints["reward_config"], "__args__", (hints["reward_config"],)) @@ -457,7 +457,7 @@ def _a2_ctx(commands, dof_pos=None): def test_a2_advance_phase_freezes_standing_envs(): - from unilab.envs.locomotion.a2.joystick import A2JoystickFlatEnv + from unilab.tasks.locomotion.a2.joystick import A2JoystickFlatEnv stub = SimpleNamespace( _cfg=SimpleNamespace(ctrl_dt=0.02), @@ -472,7 +472,7 @@ def test_a2_advance_phase_freezes_standing_envs(): def test_a2_hip_deviation_l1_over_hip_indices(): - from unilab.envs.locomotion.a2.joystick import A2JoystickFlatEnv + from unilab.tasks.locomotion.a2.joystick import A2JoystickFlatEnv dof_pos = np.zeros((1, 12)) dof_pos[0, [0, 3, 6, 9]] = [0.1, -0.2, 0.3, -0.4] @@ -482,7 +482,7 @@ def test_a2_hip_deviation_l1_over_hip_indices(): def test_a2_stand_feet_air_counts_lifted_feet_when_standing(): - from unilab.envs.locomotion.a2.joystick import A2JoystickFlatEnv + from unilab.tasks.locomotion.a2.joystick import A2JoystickFlatEnv stub = SimpleNamespace( _reward_cfg=SimpleNamespace(command_threshold=0.1), @@ -493,7 +493,7 @@ def test_a2_stand_feet_air_counts_lifted_feet_when_standing(): def test_a2_stand_feet_air_inactive_during_locomotion(): - from unilab.envs.locomotion.a2.joystick import A2JoystickFlatEnv + from unilab.tasks.locomotion.a2.joystick import A2JoystickFlatEnv stub = SimpleNamespace( _reward_cfg=SimpleNamespace(command_threshold=0.1), diff --git a/tests/tasks/test_package_boundary.py b/tests/tasks/test_package_boundary.py index fc3156e20..faa43c75a 100644 --- a/tests/tasks/test_package_boundary.py +++ b/tests/tasks/test_package_boundary.py @@ -11,13 +11,13 @@ _REPO_ROOT = Path(__file__).resolve().parents[2] _ENV_PACKAGE = _REPO_ROOT / "src" / "unilab" / "envs" -_LEGACY_TASK_MODULES = ( +_TASK_REGISTRY_MODULES = ( "unilab.envs.locomotion.go1", "unilab.envs.locomotion.go2", "unilab.envs.locomotion.go2w", "unilab.envs.locomotion.g1", "unilab.envs.locomotion.go2_arm", - "unilab.envs.locomotion.a2", + "unilab.tasks.locomotion.a2", "unilab.envs.manipulation.allegro_inhand", "unilab.envs.manipulation.sharpa_inhand", "unilab.envs.manipulation.stewart", @@ -39,7 +39,7 @@ def _imports(path: Path) -> set[str]: def test_tasks_is_the_only_default_registry_bootstrap() -> None: assert registry._DEFAULT_REGISTRY_PACKAGES == ("unilab.tasks",) - assert __unilab_registry_modules__ == _LEGACY_TASK_MODULES + assert __unilab_registry_modules__ == _TASK_REGISTRY_MODULES def test_env_runtime_does_not_depend_on_tasks() -> None: From c63ac0e248e84babeae6103c8c4600c1d3b29bde Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 20:29:10 +0800 Subject: [PATCH 036/169] refactor(tasks): move Go2 FootStand owner [skip ci] Part of #1112 and #1042. Local make test-all passed; remote child CI intentionally skipped. --- src/unilab/envs/locomotion/go2/__init__.py | 1 - src/unilab/tasks/__init__.py | 1 + src/unilab/tasks/locomotion/go2/__init__.py | 3 +++ src/unilab/{envs => tasks}/locomotion/go2/footstand.py | 8 ++++---- tests/envs/locomotion/test_go2_footstand.py | 2 +- tests/tasks/test_package_boundary.py | 1 + 6 files changed, 10 insertions(+), 6 deletions(-) create mode 100644 src/unilab/tasks/locomotion/go2/__init__.py rename src/unilab/{envs => tasks}/locomotion/go2/footstand.py (99%) diff --git a/src/unilab/envs/locomotion/go2/__init__.py b/src/unilab/envs/locomotion/go2/__init__.py index 4a4eded57..4844e53be 100644 --- a/src/unilab/envs/locomotion/go2/__init__.py +++ b/src/unilab/envs/locomotion/go2/__init__.py @@ -1,3 +1,2 @@ -from .footstand import Go2FootStandCfg, Go2FootStandTask from .joystick import Go2JoystickCfg, Go2WalkTask from .rough import Go2JoystickRoughCfg, Go2JoystickRoughEnv diff --git a/src/unilab/tasks/__init__.py b/src/unilab/tasks/__init__.py index 13c4f91df..5ee1095a5 100644 --- a/src/unilab/tasks/__init__.py +++ b/src/unilab/tasks/__init__.py @@ -9,6 +9,7 @@ __unilab_registry_modules__ = ( "unilab.envs.locomotion.go1", + "unilab.tasks.locomotion.go2.footstand", "unilab.envs.locomotion.go2", "unilab.envs.locomotion.go2w", "unilab.envs.locomotion.g1", diff --git a/src/unilab/tasks/locomotion/go2/__init__.py b/src/unilab/tasks/locomotion/go2/__init__.py new file mode 100644 index 000000000..cca0f6c82 --- /dev/null +++ b/src/unilab/tasks/locomotion/go2/__init__.py @@ -0,0 +1,3 @@ +from .footstand import Go2FootStandCfg, Go2FootStandTask + +__all__ = ["Go2FootStandCfg", "Go2FootStandTask"] diff --git a/src/unilab/envs/locomotion/go2/footstand.py b/src/unilab/tasks/locomotion/go2/footstand.py similarity index 99% rename from src/unilab/envs/locomotion/go2/footstand.py rename to src/unilab/tasks/locomotion/go2/footstand.py index 6e3242f7e..8aae60c37 100644 --- a/src/unilab/envs/locomotion/go2/footstand.py +++ b/src/unilab/tasks/locomotion/go2/footstand.py @@ -83,12 +83,12 @@ class JoystickSensor: @dataclass class Go2HandStandCfg(Go2BaseCfg): - scene: SceneCfg = field( + scene: SceneCfg = field( # pyright: ignore[reportIncompatibleVariableOverride] default_factory=lambda: SceneCfg( model_file=str(ASSETS_ROOT_PATH / "robots" / "go2" / "scene_flat.xml") ) ) - max_episode_seconds: float = 20.0 + max_episode_seconds: float = 20.0 # pyright: ignore[reportIncompatibleVariableOverride] init_state: InitState = field(default_factory=InitState) commands: Commands = field(default_factory=Commands) reward_config: RewardConfig | None = None @@ -127,7 +127,7 @@ def _compute_reset_obs( class Go2HandStandTask(Go2BaseEnv): - _cfg: Go2HandStandCfg + _cfg: Go2HandStandCfg # pyright: ignore[reportIncompatibleVariableOverride] def __init__(self, cfg: Go2HandStandCfg, num_envs=1, backend_type="mujoco"): if cfg.reward_config is None: @@ -449,7 +449,7 @@ def _compute_reset_obs( @registry.env("Go2FootStand", sim_backend="motrix") @registry.env("Go2FootStand", sim_backend="drake") class Go2FootStandTask(Go2HandStandTask): - _cfg: Go2FootStandCfg + _cfg: Go2FootStandCfg # pyright: ignore[reportIncompatibleVariableOverride] def __init__(self, cfg: Go2FootStandCfg, num_envs=1, backend_type="mujoco"): super().__init__(cfg, num_envs=num_envs, backend_type=backend_type) diff --git a/tests/envs/locomotion/test_go2_footstand.py b/tests/envs/locomotion/test_go2_footstand.py index b99b16bef..ab4356a8b 100644 --- a/tests/envs/locomotion/test_go2_footstand.py +++ b/tests/envs/locomotion/test_go2_footstand.py @@ -8,7 +8,7 @@ from unilab.base.registry import ensure_registries from unilab.dr import ResetRandomizationPayload from unilab.envs.locomotion.common.rewards import RewardContext -from unilab.envs.locomotion.go2.footstand import ( +from unilab.tasks.locomotion.go2.footstand import ( FootstandControlConfig, FootstandSensor, Go2FootStandCfg, diff --git a/tests/tasks/test_package_boundary.py b/tests/tasks/test_package_boundary.py index faa43c75a..63eea6c92 100644 --- a/tests/tasks/test_package_boundary.py +++ b/tests/tasks/test_package_boundary.py @@ -13,6 +13,7 @@ _TASK_REGISTRY_MODULES = ( "unilab.envs.locomotion.go1", + "unilab.tasks.locomotion.go2.footstand", "unilab.envs.locomotion.go2", "unilab.envs.locomotion.go2w", "unilab.envs.locomotion.g1", From d9c099934f968c2e2cbcb769fb165a42bedb39eb Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 20:33:52 +0800 Subject: [PATCH 037/169] refactor(tasks): move Go2 rough owner [skip ci] Part of #1112 and #1042. Local make test-all passed; remote child CI intentionally skipped. --- scripts/benchmark/env/benchmark_env_step.py | 4 ++-- src/unilab/envs/locomotion/go2/__init__.py | 1 - src/unilab/tasks/__init__.py | 1 + src/unilab/tasks/locomotion/go2/__init__.py | 8 +++++++- .../{envs => tasks}/locomotion/go2/rough.py | 0 tests/config/test_locomotion_params.py | 12 ++++++------ .../locomotion/test_go2_rough_height_scan.py | 2 +- .../envs/locomotion/test_go2_terrain_spawn.py | 18 +++++++++--------- tests/tasks/test_package_boundary.py | 1 + 9 files changed, 27 insertions(+), 20 deletions(-) rename src/unilab/{envs => tasks}/locomotion/go2/rough.py (100%) diff --git a/scripts/benchmark/env/benchmark_env_step.py b/scripts/benchmark/env/benchmark_env_step.py index 286514b81..c1a9ffe59 100644 --- a/scripts/benchmark/env/benchmark_env_step.py +++ b/scripts/benchmark/env/benchmark_env_step.py @@ -300,13 +300,13 @@ def _go2_env_cls() -> type: def _go2_rough_cfg(backend: str, config_overrides: list[str]) -> Any: - from unilab.envs.locomotion.go2.rough import Go2JoystickRoughCfg + from unilab.tasks.locomotion.go2.rough import Go2JoystickRoughCfg return _ppo_owner_yaml_cfg("go2_joystick_rough", backend, Go2JoystickRoughCfg, config_overrides) def _go2_rough_env_cls() -> type: - from unilab.envs.locomotion.go2.rough import Go2JoystickRoughEnv + from unilab.tasks.locomotion.go2.rough import Go2JoystickRoughEnv return Go2JoystickRoughEnv diff --git a/src/unilab/envs/locomotion/go2/__init__.py b/src/unilab/envs/locomotion/go2/__init__.py index 4844e53be..d6f14dbb3 100644 --- a/src/unilab/envs/locomotion/go2/__init__.py +++ b/src/unilab/envs/locomotion/go2/__init__.py @@ -1,2 +1 @@ from .joystick import Go2JoystickCfg, Go2WalkTask -from .rough import Go2JoystickRoughCfg, Go2JoystickRoughEnv diff --git a/src/unilab/tasks/__init__.py b/src/unilab/tasks/__init__.py index 5ee1095a5..9e1b0bad1 100644 --- a/src/unilab/tasks/__init__.py +++ b/src/unilab/tasks/__init__.py @@ -11,6 +11,7 @@ "unilab.envs.locomotion.go1", "unilab.tasks.locomotion.go2.footstand", "unilab.envs.locomotion.go2", + "unilab.tasks.locomotion.go2.rough", "unilab.envs.locomotion.go2w", "unilab.envs.locomotion.g1", "unilab.envs.locomotion.go2_arm", diff --git a/src/unilab/tasks/locomotion/go2/__init__.py b/src/unilab/tasks/locomotion/go2/__init__.py index cca0f6c82..9b06bc55a 100644 --- a/src/unilab/tasks/locomotion/go2/__init__.py +++ b/src/unilab/tasks/locomotion/go2/__init__.py @@ -1,3 +1,9 @@ from .footstand import Go2FootStandCfg, Go2FootStandTask +from .rough import Go2JoystickRoughCfg, Go2JoystickRoughEnv -__all__ = ["Go2FootStandCfg", "Go2FootStandTask"] +__all__ = [ + "Go2FootStandCfg", + "Go2FootStandTask", + "Go2JoystickRoughCfg", + "Go2JoystickRoughEnv", +] diff --git a/src/unilab/envs/locomotion/go2/rough.py b/src/unilab/tasks/locomotion/go2/rough.py similarity index 100% rename from src/unilab/envs/locomotion/go2/rough.py rename to src/unilab/tasks/locomotion/go2/rough.py diff --git a/tests/config/test_locomotion_params.py b/tests/config/test_locomotion_params.py index 376ce31a2..92b55df59 100644 --- a/tests/config/test_locomotion_params.py +++ b/tests/config/test_locomotion_params.py @@ -227,7 +227,7 @@ def test_offpolicy_flashsac_go2_task_overrides(): def test_go2_joystick_rough_uses_terrain_generator(): from unilab.assets import ASSETS_ROOT_PATH from unilab.base.scene import SceneCfg, TerrainSceneCfg - from unilab.envs.locomotion.go2.rough import Go2JoystickRoughCfg + from unilab.tasks.locomotion.go2.rough import Go2JoystickRoughCfg from unilab.terrains import TerrainGeneratorCfg cfg = Go2JoystickRoughCfg() @@ -245,7 +245,7 @@ def test_go2_joystick_rough_uses_terrain_generator(): def test_go2_joystick_rough_terrain_cfg_is_independent_per_instance(): """Confirm rough terrain cfg defaults are not shared across instances.""" - from unilab.envs.locomotion.go2.rough import Go2JoystickRoughCfg + from unilab.tasks.locomotion.go2.rough import Go2JoystickRoughCfg a = Go2JoystickRoughCfg() b = Go2JoystickRoughCfg() @@ -260,7 +260,7 @@ def test_go2_joystick_rough_playback_model_uses_backend_scene(tmp_path): import mujoco from unilab.envs.locomotion.go2.joystick import RewardConfig - from unilab.envs.locomotion.go2.rough import Go2JoystickRoughCfg, Go2JoystickRoughEnv + from unilab.tasks.locomotion.go2.rough import Go2JoystickRoughCfg, Go2JoystickRoughEnv from unilab.visualization.playback import _resolve_render_play_model_files cfg = Go2JoystickRoughCfg( @@ -348,7 +348,7 @@ def test_ppo_go2_joystick_rough_motrix_task_compose(): def test_go2_joystick_rough_motrix_registers_rough_env(): from unilab.base import registry - from unilab.envs.locomotion.go2.rough import Go2JoystickRoughEnv + from unilab.tasks.locomotion.go2.rough import Go2JoystickRoughEnv assert registry._envs["Go2JoystickRough"].env_factory_dict["motrix"] is Go2JoystickRoughEnv @@ -659,7 +659,7 @@ def test_apply_cfg_overrides_deep_merges_dataclass_field(): instances rather than re-instantiating them, so partial overrides like `scene.terrain.generator.num_rows=4` keep `sub_terrains` and other defaults.""" from unilab.base.registry import apply_cfg_overrides - from unilab.envs.locomotion.go2.rough import Go2JoystickRoughCfg + from unilab.tasks.locomotion.go2.rough import Go2JoystickRoughCfg cfg = Go2JoystickRoughCfg() cfg.scene.terrain.generator.num_cols = 3 @@ -693,7 +693,7 @@ def test_ppo_go2_joystick_rough_hydra_terrain_override(): from hydra.core.global_hydra import GlobalHydra from unilab.base.registry import apply_cfg_overrides - from unilab.envs.locomotion.go2.rough import Go2JoystickRoughCfg + from unilab.tasks.locomotion.go2.rough import Go2JoystickRoughCfg from unilab.training.backend_adapter import BackendAdapter GlobalHydra.instance().clear() diff --git a/tests/envs/locomotion/test_go2_rough_height_scan.py b/tests/envs/locomotion/test_go2_rough_height_scan.py index b1ec5ec27..ba62e62bc 100644 --- a/tests/envs/locomotion/test_go2_rough_height_scan.py +++ b/tests/envs/locomotion/test_go2_rough_height_scan.py @@ -5,7 +5,7 @@ import numpy as np from unilab.envs.locomotion.common.height_scan import height_scan_offsets as _height_scan_offsets -from unilab.envs.locomotion.go2.rough import Go2JoystickRoughEnv +from unilab.tasks.locomotion.go2.rough import Go2JoystickRoughEnv def test_go2_rough_height_scan_uses_backend_native_sampling() -> None: diff --git a/tests/envs/locomotion/test_go2_terrain_spawn.py b/tests/envs/locomotion/test_go2_terrain_spawn.py index bf3c10175..4444d4d78 100644 --- a/tests/envs/locomotion/test_go2_terrain_spawn.py +++ b/tests/envs/locomotion/test_go2_terrain_spawn.py @@ -30,7 +30,7 @@ def _configure_small_terrain(cfg, *, seed: int = 0) -> None: def _rough_cfg(*, curriculum_enabled: bool = False, seed: int = 0): - from unilab.envs.locomotion.go2.rough import Go2JoystickRoughCfg, RoughRewardConfig + from unilab.tasks.locomotion.go2.rough import Go2JoystickRoughCfg, RoughRewardConfig cfg = Go2JoystickRoughCfg( reward_config=RoughRewardConfig(scales={}, tracking_sigma=0.25, base_height_target=0.3) @@ -41,7 +41,7 @@ def _rough_cfg(*, curriculum_enabled: bool = False, seed: int = 0): def test_terrain_spawn_attached_when_rough(): - from unilab.envs.locomotion.go2.rough import Go2JoystickRoughEnv + from unilab.tasks.locomotion.go2.rough import Go2JoystickRoughEnv cfg = _rough_cfg() env = Go2JoystickRoughEnv(cfg, num_envs=4, backend_type="mujoco") @@ -204,7 +204,7 @@ def sample_height(self, xy): def test_go2_rough_pd_torque_estimate_returns_dof_order(): - from unilab.envs.locomotion.go2.rough import ( + from unilab.tasks.locomotion.go2.rough import ( GO2_ACTUATOR_TO_DOF_INDICES, Go2JoystickRoughEnv, ) @@ -250,7 +250,7 @@ def test_default_spawn_used_when_flat(): def test_curriculum_disabled_distributes_levels_uniformly(): - from unilab.envs.locomotion.go2.rough import Go2JoystickRoughEnv + from unilab.tasks.locomotion.go2.rough import Go2JoystickRoughEnv cfg = _rough_cfg(curriculum_enabled=False, seed=0) env = Go2JoystickRoughEnv(cfg, num_envs=64, backend_type="mujoco") @@ -266,7 +266,7 @@ def test_curriculum_disabled_distributes_levels_uniformly(): def test_curriculum_enabled_levels_start_at_zero(): - from unilab.envs.locomotion.go2.rough import Go2JoystickRoughEnv + from unilab.tasks.locomotion.go2.rough import Go2JoystickRoughEnv cfg = _rough_cfg(curriculum_enabled=True, seed=0) env = Go2JoystickRoughEnv(cfg, num_envs=8, backend_type="mujoco") @@ -279,7 +279,7 @@ def test_curriculum_enabled_levels_start_at_zero(): def test_reset_qpos_xy_matches_terrain_origins(): - from unilab.envs.locomotion.go2.rough import Go2JoystickRoughEnv + from unilab.tasks.locomotion.go2.rough import Go2JoystickRoughEnv cfg = _rough_cfg(curriculum_enabled=False, seed=0) env = Go2JoystickRoughEnv(cfg, num_envs=4, backend_type="mujoco") @@ -299,7 +299,7 @@ def test_reset_qpos_xy_matches_terrain_origins(): def test_rough_reset_spawns_above_sampled_terrain(): - from unilab.envs.locomotion.go2.rough import Go2JoystickRoughEnv + from unilab.tasks.locomotion.go2.rough import Go2JoystickRoughEnv cfg = _rough_cfg(curriculum_enabled=False, seed=0) env = Go2JoystickRoughEnv(cfg, num_envs=64, backend_type="mujoco") @@ -315,7 +315,7 @@ def test_rough_reset_spawns_above_sampled_terrain(): def test_curriculum_logs_appear_after_done(): - from unilab.envs.locomotion.go2.rough import Go2JoystickRoughEnv + from unilab.tasks.locomotion.go2.rough import Go2JoystickRoughEnv cfg = _rough_cfg(curriculum_enabled=True, seed=0) env = Go2JoystickRoughEnv(cfg, num_envs=4, backend_type="mujoco") @@ -345,7 +345,7 @@ def test_episode_start_recorded_after_reset(preset): Go2WalkTask, RewardConfig, ) - from unilab.envs.locomotion.go2.rough import Go2JoystickRoughEnv + from unilab.tasks.locomotion.go2.rough import Go2JoystickRoughEnv if preset == "flat": cfg = Go2JoystickCfg( diff --git a/tests/tasks/test_package_boundary.py b/tests/tasks/test_package_boundary.py index 63eea6c92..d443d73d5 100644 --- a/tests/tasks/test_package_boundary.py +++ b/tests/tasks/test_package_boundary.py @@ -15,6 +15,7 @@ "unilab.envs.locomotion.go1", "unilab.tasks.locomotion.go2.footstand", "unilab.envs.locomotion.go2", + "unilab.tasks.locomotion.go2.rough", "unilab.envs.locomotion.go2w", "unilab.envs.locomotion.g1", "unilab.envs.locomotion.go2_arm", From 36cfd744db5d69f716078bd7a999a292a1cd102b Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 20:45:19 +0800 Subject: [PATCH 038/169] test(backend): follow task owner source roots [skip ci] Closes #1122. Roadmap #1042. --- tests/base/test_backend_conformance.py | 41 ++++++++++++++++---------- 1 file changed, 26 insertions(+), 15 deletions(-) diff --git a/tests/base/test_backend_conformance.py b/tests/base/test_backend_conformance.py index f8512ef0f..c9c8f06c3 100644 --- a/tests/base/test_backend_conformance.py +++ b/tests/base/test_backend_conformance.py @@ -28,11 +28,15 @@ _BACKEND_CLASS_NAMES = frozenset( {"MuJoCoBackend", "MotrixBackend", "DrakeBackend", "MjwarpBackend"} ) -_TERRAIN_CONSUMER_FILES = ( - SRC_ROOT / "unilab" / "envs" / "locomotion" / "common" / "terrain_spawn.py", - SRC_ROOT / "unilab" / "envs" / "locomotion" / "go1" / "joystick.py", - SRC_ROOT / "unilab" / "envs" / "locomotion" / "go2" / "joystick.py", - SRC_ROOT / "unilab" / "envs" / "locomotion" / "go2w" / "rough.py", +_TASK_SOURCE_ROOTS = ( + SRC_ROOT / "unilab" / "envs", + SRC_ROOT / "unilab" / "tasks", +) +_TERRAIN_CONSUMER_PATHS = ( + Path("locomotion/common/terrain_spawn.py"), + Path("locomotion/go1/joystick.py"), + Path("locomotion/go2/joystick.py"), + Path("locomotion/go2w/rough.py"), ) NUM_ENVS = 2 @@ -363,16 +367,23 @@ def test_drake_root_layout_is_explicitly_unsupported_without_runtime_metadata() def test_terrain_spawn_consumers_do_not_probe_private_backend_capabilities() -> None: forbidden_names = {"terrain_origins", "terrain_surface_sampler", "sample_height"} offenders: list[str] = [] - for path in _TERRAIN_CONSUMER_FILES: - tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) - for node in ast.walk(tree): - if not isinstance(node, ast.Call) or not isinstance(node.func, ast.Name): - continue - if node.func.id not in {"getattr", "hasattr"}: - continue - for arg in node.args[1:]: - if isinstance(arg, ast.Constant) and arg.value in forbidden_names: - offenders.append(f"{path.relative_to(REPO_ROOT)}:{node.lineno}: {arg.value}") + for relative_path in _TERRAIN_CONSUMER_PATHS: + owner_paths = tuple( + root / relative_path for root in _TASK_SOURCE_ROOTS if (root / relative_path).is_file() + ) + assert owner_paths, f"terrain consumer source not found: {relative_path}" + for path in owner_paths: + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + for node in ast.walk(tree): + if not isinstance(node, ast.Call) or not isinstance(node.func, ast.Name): + continue + if node.func.id not in {"getattr", "hasattr"}: + continue + for arg in node.args[1:]: + if isinstance(arg, ast.Constant) and arg.value in forbidden_names: + offenders.append( + f"{path.relative_to(REPO_ROOT)}:{node.lineno}: {arg.value}" + ) assert not offenders, "private terrain capability probes:\n" + "\n".join(offenders) From c8c072eab29bc5e8eeaea08f2f37586c594eda39 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 20:47:52 +0800 Subject: [PATCH 039/169] refactor(tasks): move Go2 joystick owner [skip ci] Closes #1121. Roadmap #1042. --- scripts/benchmark/benchmark_drake_performance.py | 4 ++-- scripts/benchmark/core/task_names.py | 2 +- scripts/benchmark/env/benchmark_env_step.py | 4 ++-- scripts/play_interactive.py | 8 +++++--- src/unilab/envs/locomotion/go2/__init__.py | 2 +- src/unilab/tasks/__init__.py | 4 +--- src/unilab/tasks/locomotion/a2/joystick.py | 2 +- src/unilab/tasks/locomotion/go2/__init__.py | 3 +++ src/unilab/{envs => tasks}/locomotion/go2/joystick.py | 10 ++++++---- src/unilab/tasks/locomotion/go2/rough.py | 2 +- tests/config/test_locomotion_params.py | 4 ++-- tests/envs/locomotion/test_go2_joystick_stand_still.py | 4 ++-- tests/envs/locomotion/test_go2_terrain_spawn.py | 8 ++++---- tests/scripts/test_visualization_entrypoints.py | 4 ++-- tests/tasks/test_package_boundary.py | 4 +--- 15 files changed, 34 insertions(+), 31 deletions(-) rename src/unilab/{envs => tasks}/locomotion/go2/joystick.py (97%) diff --git a/scripts/benchmark/benchmark_drake_performance.py b/scripts/benchmark/benchmark_drake_performance.py index 2886bebf6..a5e405eee 100644 --- a/scripts/benchmark/benchmark_drake_performance.py +++ b/scripts/benchmark/benchmark_drake_performance.py @@ -111,12 +111,12 @@ def go1_env() -> type: return Go1WalkTask def go2_cfg() -> Any: - from unilab.envs.locomotion.go2.joystick import Go2JoystickCfg + from unilab.tasks.locomotion.go2.joystick import Go2JoystickCfg return Go2JoystickCfg() def go2_env() -> type: - from unilab.envs.locomotion.go2.joystick import Go2WalkTask + from unilab.tasks.locomotion.go2.joystick import Go2WalkTask return Go2WalkTask diff --git a/scripts/benchmark/core/task_names.py b/scripts/benchmark/core/task_names.py index b922c06f5..1a914fc53 100644 --- a/scripts/benchmark/core/task_names.py +++ b/scripts/benchmark/core/task_names.py @@ -4,8 +4,8 @@ from unilab.envs.locomotion.g1.joystick import G1WalkFlatCfg from unilab.envs.locomotion.go1.joystick import Go1JoystickCfg -from unilab.envs.locomotion.go2.joystick import Go2JoystickCfg from unilab.envs.manipulation.sharpa_inhand.rotation import SharpaInhandRotationCfg +from unilab.tasks.locomotion.go2.joystick import Go2JoystickCfg @dataclass(frozen=True) diff --git a/scripts/benchmark/env/benchmark_env_step.py b/scripts/benchmark/env/benchmark_env_step.py index c1a9ffe59..bfc45a5fc 100644 --- a/scripts/benchmark/env/benchmark_env_step.py +++ b/scripts/benchmark/env/benchmark_env_step.py @@ -288,13 +288,13 @@ def _go1_env_cls() -> type: def _go2_cfg(backend: str, config_overrides: list[str]) -> Any: - from unilab.envs.locomotion.go2.joystick import Go2JoystickCfg + from unilab.tasks.locomotion.go2.joystick import Go2JoystickCfg return _ppo_owner_yaml_cfg("go2_joystick_flat", backend, Go2JoystickCfg, config_overrides) def _go2_env_cls() -> type: - from unilab.envs.locomotion.go2.joystick import Go2WalkTask + from unilab.tasks.locomotion.go2.joystick import Go2WalkTask return Go2WalkTask diff --git a/scripts/play_interactive.py b/scripts/play_interactive.py index e812c11e1..c0bff8328 100644 --- a/scripts/play_interactive.py +++ b/scripts/play_interactive.py @@ -894,12 +894,14 @@ def _state_has_velocity_commands(env: Any) -> bool: ) -def _is_locomotion_env(env: Any) -> bool: - return type(env).__module__.startswith("unilab.envs.locomotion") +def _has_velocity_command_config(env: Any) -> bool: + cfg = getattr(env, "cfg", None) + commands_cfg = getattr(cfg, "commands", None) if cfg is not None else None + return getattr(commands_cfg, "vel_limit", None) is not None def _is_velocity_command_locomotion_task(env: Any) -> bool: - if not _is_locomotion_env(env): + if not _has_velocity_command_config(env): return False cfg = getattr(env, "cfg", None) candidate_names = [ diff --git a/src/unilab/envs/locomotion/go2/__init__.py b/src/unilab/envs/locomotion/go2/__init__.py index d6f14dbb3..49122f621 100644 --- a/src/unilab/envs/locomotion/go2/__init__.py +++ b/src/unilab/envs/locomotion/go2/__init__.py @@ -1 +1 @@ -from .joystick import Go2JoystickCfg, Go2WalkTask +"""Legacy Go2 shared base pending migration to :mod:`unilab.tasks`.""" diff --git a/src/unilab/tasks/__init__.py b/src/unilab/tasks/__init__.py index 9e1b0bad1..448745635 100644 --- a/src/unilab/tasks/__init__.py +++ b/src/unilab/tasks/__init__.py @@ -9,9 +9,7 @@ __unilab_registry_modules__ = ( "unilab.envs.locomotion.go1", - "unilab.tasks.locomotion.go2.footstand", - "unilab.envs.locomotion.go2", - "unilab.tasks.locomotion.go2.rough", + "unilab.tasks.locomotion.go2", "unilab.envs.locomotion.go2w", "unilab.envs.locomotion.g1", "unilab.envs.locomotion.go2_arm", diff --git a/src/unilab/tasks/locomotion/a2/joystick.py b/src/unilab/tasks/locomotion/a2/joystick.py index 59df62acc..cf7f03448 100644 --- a/src/unilab/tasks/locomotion/a2/joystick.py +++ b/src/unilab/tasks/locomotion/a2/joystick.py @@ -28,7 +28,7 @@ from unilab.envs.locomotion.common.dr_provider import LocomotionDRProvider from unilab.envs.locomotion.common.rewards import RewardContext from unilab.envs.locomotion.go2.base import Asset, ControlConfig -from unilab.envs.locomotion.go2.joystick import ( +from unilab.tasks.locomotion.go2.joystick import ( Go2DomainRandConfig, Go2JoystickCfg, Go2JoystickDomainRandomizationProvider, diff --git a/src/unilab/tasks/locomotion/go2/__init__.py b/src/unilab/tasks/locomotion/go2/__init__.py index 9b06bc55a..d15ae6ac9 100644 --- a/src/unilab/tasks/locomotion/go2/__init__.py +++ b/src/unilab/tasks/locomotion/go2/__init__.py @@ -1,9 +1,12 @@ from .footstand import Go2FootStandCfg, Go2FootStandTask +from .joystick import Go2JoystickCfg, Go2WalkTask from .rough import Go2JoystickRoughCfg, Go2JoystickRoughEnv __all__ = [ "Go2FootStandCfg", "Go2FootStandTask", + "Go2JoystickCfg", "Go2JoystickRoughCfg", "Go2JoystickRoughEnv", + "Go2WalkTask", ] diff --git a/src/unilab/envs/locomotion/go2/joystick.py b/src/unilab/tasks/locomotion/go2/joystick.py similarity index 97% rename from src/unilab/envs/locomotion/go2/joystick.py rename to src/unilab/tasks/locomotion/go2/joystick.py index 98f0bbefc..855c999f0 100644 --- a/src/unilab/envs/locomotion/go2/joystick.py +++ b/src/unilab/tasks/locomotion/go2/joystick.py @@ -60,16 +60,18 @@ class JoystickSensor(Sensor): @dataclass class Go2JoystickCfg(Go2BaseCfg): - scene: SceneCfg = field( + scene: SceneCfg = field( # pyright: ignore[reportIncompatibleVariableOverride] default_factory=lambda: SceneCfg( model_file=str(ASSETS_ROOT_PATH / "robots" / "go2" / "scene_flat.xml") ) ) - max_episode_seconds: float = 20.0 + max_episode_seconds: float = 20.0 # pyright: ignore[reportIncompatibleVariableOverride] init_state: InitState = field(default_factory=InitState) commands: Commands = field(default_factory=Commands) reward_config: RewardConfig | None = None - sensor: JoystickSensor = field(default_factory=JoystickSensor) + sensor: JoystickSensor = field( # pyright: ignore[reportIncompatibleVariableOverride] + default_factory=JoystickSensor + ) domain_rand: Go2DomainRandConfig = field(default_factory=Go2DomainRandConfig) terrain_curriculum: TerrainCurriculumCfg = field(default_factory=TerrainCurriculumCfg) @@ -95,7 +97,7 @@ def _compute_reset_obs( class Go2WalkTask(Go2BaseEnv): - _cfg: Go2JoystickCfg + _cfg: Go2JoystickCfg # pyright: ignore[reportIncompatibleVariableOverride] def __init__(self, cfg: Go2JoystickCfg, num_envs=1, backend_type="mujoco"): if cfg.reward_config is None: diff --git a/src/unilab/tasks/locomotion/go2/rough.py b/src/unilab/tasks/locomotion/go2/rough.py index 0d7d5975e..9f516b290 100644 --- a/src/unilab/tasks/locomotion/go2/rough.py +++ b/src/unilab/tasks/locomotion/go2/rough.py @@ -30,7 +30,7 @@ ) from unilab.envs.locomotion.common.rewards import RewardContext from unilab.envs.locomotion.go2.base import ControlConfig -from unilab.envs.locomotion.go2.joystick import ( +from unilab.tasks.locomotion.go2.joystick import ( Commands, Go2JoystickCfg, Go2JoystickDomainRandomizationProvider, diff --git a/tests/config/test_locomotion_params.py b/tests/config/test_locomotion_params.py index 92b55df59..1cd789188 100644 --- a/tests/config/test_locomotion_params.py +++ b/tests/config/test_locomotion_params.py @@ -259,7 +259,7 @@ def test_go2_joystick_rough_playback_model_uses_backend_scene(tmp_path): """Offline playback / video rendering must reuse the backend-compiled scene model.""" import mujoco - from unilab.envs.locomotion.go2.joystick import RewardConfig + from unilab.tasks.locomotion.go2.joystick import RewardConfig from unilab.tasks.locomotion.go2.rough import Go2JoystickRoughCfg, Go2JoystickRoughEnv from unilab.visualization.playback import _resolve_render_play_model_files @@ -296,7 +296,7 @@ def test_go2_joystick_rough_playback_model_uses_backend_scene(tmp_path): def test_go2_joystick_flat_no_terrain_materialized(): """Flat task keeps the static scene source and has no terrain origins.""" - from unilab.envs.locomotion.go2.joystick import ( + from unilab.tasks.locomotion.go2.joystick import ( Go2JoystickCfg, Go2WalkTask, RewardConfig, diff --git a/tests/envs/locomotion/test_go2_joystick_stand_still.py b/tests/envs/locomotion/test_go2_joystick_stand_still.py index 7e8f01f2d..8f27eb2d3 100644 --- a/tests/envs/locomotion/test_go2_joystick_stand_still.py +++ b/tests/envs/locomotion/test_go2_joystick_stand_still.py @@ -67,7 +67,7 @@ def test_sample_commands_with_standing_matches_rough_block(): def test_go2_advance_phase_is_unconditional(): """Go2WalkTask advances the gait clock every step regardless of command — the A2 standing freeze must not have leaked into the Go2 owner.""" - from unilab.envs.locomotion.go2.joystick import Go2WalkTask + from unilab.tasks.locomotion.go2.joystick import Go2WalkTask stub = SimpleNamespace(_cfg=SimpleNamespace(ctrl_dt=0.02), gait_frequency=2.0) phase = np.array([0.3, 0.3]) @@ -81,7 +81,7 @@ def test_go2_reward_config_has_no_command_threshold(): must not declare it.""" import dataclasses - from unilab.envs.locomotion.go2.joystick import RewardConfig + from unilab.tasks.locomotion.go2.joystick import RewardConfig names = {f.name for f in dataclasses.fields(RewardConfig)} assert "command_threshold" not in names diff --git a/tests/envs/locomotion/test_go2_terrain_spawn.py b/tests/envs/locomotion/test_go2_terrain_spawn.py index 4444d4d78..4587c0a3c 100644 --- a/tests/envs/locomotion/test_go2_terrain_spawn.py +++ b/tests/envs/locomotion/test_go2_terrain_spawn.py @@ -63,7 +63,7 @@ def test_terrain_spawn_attached_when_rough(): def test_terrain_spawn_attached_when_rough_motrix(): pytest.importorskip("motrixsim") - from unilab.envs.locomotion.go2.joystick import Go2WalkTask + from unilab.tasks.locomotion.go2.joystick import Go2WalkTask cfg = _rough_cfg() cfg.domain_rand.randomize_kp = False @@ -178,7 +178,7 @@ def test_terrain_spawn_rejects_non_callable_height_sampler_on_init(): def test_go2_rough_base_height_reward_uses_terrain_relative_height(): - from unilab.envs.locomotion.go2.joystick import Go2WalkTask + from unilab.tasks.locomotion.go2.joystick import Go2WalkTask class FakeBackend: def get_base_pos(self): @@ -227,7 +227,7 @@ def test_go2_rough_pd_torque_estimate_returns_dof_order(): def test_default_spawn_used_when_flat(): - from unilab.envs.locomotion.go2.joystick import ( + from unilab.tasks.locomotion.go2.joystick import ( Go2JoystickCfg, Go2WalkTask, RewardConfig, @@ -340,7 +340,7 @@ def test_curriculum_logs_appear_after_done(): @pytest.mark.parametrize("preset", ["flat", "rough"]) def test_episode_start_recorded_after_reset(preset): - from unilab.envs.locomotion.go2.joystick import ( + from unilab.tasks.locomotion.go2.joystick import ( Go2JoystickCfg, Go2WalkTask, RewardConfig, diff --git a/tests/scripts/test_visualization_entrypoints.py b/tests/scripts/test_visualization_entrypoints.py index a5006b0ab..15175ee69 100644 --- a/tests/scripts/test_visualization_entrypoints.py +++ b/tests/scripts/test_visualization_entrypoints.py @@ -201,7 +201,7 @@ def test_velocity_arrows_require_velocity_command_task_and_policy_obs(): joystick_env = _keyboard_env( env_cls_name="Go2WalkTask", cfg_cls_name="Go2JoystickCfg", - module="unilab.envs.locomotion.go2.joystick", + module="unilab.tasks.locomotion.go2.joystick", obs_contains_command=True, ) manip_loco_env = _keyboard_env( @@ -213,7 +213,7 @@ def test_velocity_arrows_require_velocity_command_task_and_policy_obs(): missing_obs_command_env = _keyboard_env( env_cls_name="Go2WalkTask", cfg_cls_name="Go2JoystickCfg", - module="unilab.envs.locomotion.go2.joystick", + module="unilab.tasks.locomotion.go2.joystick", obs_contains_command=False, ) diff --git a/tests/tasks/test_package_boundary.py b/tests/tasks/test_package_boundary.py index d443d73d5..d1d60b051 100644 --- a/tests/tasks/test_package_boundary.py +++ b/tests/tasks/test_package_boundary.py @@ -13,9 +13,7 @@ _TASK_REGISTRY_MODULES = ( "unilab.envs.locomotion.go1", - "unilab.tasks.locomotion.go2.footstand", - "unilab.envs.locomotion.go2", - "unilab.tasks.locomotion.go2.rough", + "unilab.tasks.locomotion.go2", "unilab.envs.locomotion.go2w", "unilab.envs.locomotion.g1", "unilab.envs.locomotion.go2_arm", From 955cae0271ed2449e6251a7807864172fe7c65d6 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 20:53:54 +0800 Subject: [PATCH 040/169] refactor(tasks): move Go2 shared base [skip ci] Closes #1125. Roadmap #1042. --- docs/sphinx/source/api_reference/envs/locomotion.md | 2 +- src/unilab/envs/locomotion/__init__.py | 1 - src/unilab/envs/locomotion/go2/__init__.py | 1 - src/unilab/tasks/locomotion/a2/joystick.py | 2 +- src/unilab/{envs => tasks}/locomotion/go2/base.py | 2 +- src/unilab/tasks/locomotion/go2/footstand.py | 2 +- src/unilab/tasks/locomotion/go2/joystick.py | 2 +- src/unilab/tasks/locomotion/go2/rough.py | 2 +- tests/envs/test_go2_obs_noise.py | 2 +- 9 files changed, 7 insertions(+), 9 deletions(-) delete mode 100644 src/unilab/envs/locomotion/go2/__init__.py rename src/unilab/{envs => tasks}/locomotion/go2/base.py (95%) diff --git a/docs/sphinx/source/api_reference/envs/locomotion.md b/docs/sphinx/source/api_reference/envs/locomotion.md index 1a9d5107a..66ddebc84 100644 --- a/docs/sphinx/source/api_reference/envs/locomotion.md +++ b/docs/sphinx/source/api_reference/envs/locomotion.md @@ -9,7 +9,7 @@ unilab.envs.locomotion.common unilab.envs.locomotion.g1 unilab.envs.locomotion.go1 - unilab.envs.locomotion.go2 + unilab.tasks.locomotion.go2 unilab.envs.locomotion.go2_arm unilab.envs.locomotion.go2w ``` diff --git a/src/unilab/envs/locomotion/__init__.py b/src/unilab/envs/locomotion/__init__.py index a101b48c1..3daf9340e 100644 --- a/src/unilab/envs/locomotion/__init__.py +++ b/src/unilab/envs/locomotion/__init__.py @@ -2,7 +2,6 @@ __unilab_registry_modules__ = ( "unilab.envs.locomotion.go1", - "unilab.envs.locomotion.go2", "unilab.envs.locomotion.go2w", "unilab.envs.locomotion.g1", "unilab.envs.locomotion.go2_arm", diff --git a/src/unilab/envs/locomotion/go2/__init__.py b/src/unilab/envs/locomotion/go2/__init__.py deleted file mode 100644 index 49122f621..000000000 --- a/src/unilab/envs/locomotion/go2/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Legacy Go2 shared base pending migration to :mod:`unilab.tasks`.""" diff --git a/src/unilab/tasks/locomotion/a2/joystick.py b/src/unilab/tasks/locomotion/a2/joystick.py index cf7f03448..63aabec03 100644 --- a/src/unilab/tasks/locomotion/a2/joystick.py +++ b/src/unilab/tasks/locomotion/a2/joystick.py @@ -27,7 +27,7 @@ from unilab.envs.locomotion.common.commands import sample_commands_with_standing from unilab.envs.locomotion.common.dr_provider import LocomotionDRProvider from unilab.envs.locomotion.common.rewards import RewardContext -from unilab.envs.locomotion.go2.base import Asset, ControlConfig +from unilab.tasks.locomotion.go2.base import Asset, ControlConfig from unilab.tasks.locomotion.go2.joystick import ( Go2DomainRandConfig, Go2JoystickCfg, diff --git a/src/unilab/envs/locomotion/go2/base.py b/src/unilab/tasks/locomotion/go2/base.py similarity index 95% rename from src/unilab/envs/locomotion/go2/base.py rename to src/unilab/tasks/locomotion/go2/base.py index 08c9bbe7a..e00dc0eb5 100644 --- a/src/unilab/envs/locomotion/go2/base.py +++ b/src/unilab/tasks/locomotion/go2/base.py @@ -39,7 +39,7 @@ class Go2BaseCfg(LocomotionBaseCfg): class Go2BaseEnv(LocomotionBaseEnv): - _cfg: Go2BaseCfg + _cfg: Go2BaseCfg # pyright: ignore[reportIncompatibleVariableOverride] def get_foot_pos(self) -> np.ndarray: """Get foot positions. Returns shape (num_envs, 4, 3)""" diff --git a/src/unilab/tasks/locomotion/go2/footstand.py b/src/unilab/tasks/locomotion/go2/footstand.py index 8aae60c37..f306797d2 100644 --- a/src/unilab/tasks/locomotion/go2/footstand.py +++ b/src/unilab/tasks/locomotion/go2/footstand.py @@ -17,7 +17,7 @@ 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 ControlConfig, Go2BaseCfg, Go2BaseEnv, NoiseConfig +from unilab.tasks.locomotion.go2.base import ControlConfig, Go2BaseCfg, Go2BaseEnv, NoiseConfig from unilab.utils.rotation import np_quat_apply, np_quat_apply_inverse diff --git a/src/unilab/tasks/locomotion/go2/joystick.py b/src/unilab/tasks/locomotion/go2/joystick.py index 855c999f0..303f82542 100644 --- a/src/unilab/tasks/locomotion/go2/joystick.py +++ b/src/unilab/tasks/locomotion/go2/joystick.py @@ -21,11 +21,11 @@ TerrainCurriculumCfg, TerrainSpawnManager, ) -from unilab.envs.locomotion.go2.base import Go2BaseCfg, Go2BaseEnv from unilab.envs.manager_based_rl_env import ( ManagerBasedRlEnvCfg, make_manager_based_rl_env, ) +from unilab.tasks.locomotion.go2.base import Go2BaseCfg, Go2BaseEnv @dataclass diff --git a/src/unilab/tasks/locomotion/go2/rough.py b/src/unilab/tasks/locomotion/go2/rough.py index 9f516b290..badde1fda 100644 --- a/src/unilab/tasks/locomotion/go2/rough.py +++ b/src/unilab/tasks/locomotion/go2/rough.py @@ -29,7 +29,7 @@ terrain_out_of_bounds, ) from unilab.envs.locomotion.common.rewards import RewardContext -from unilab.envs.locomotion.go2.base import ControlConfig +from unilab.tasks.locomotion.go2.base import ControlConfig from unilab.tasks.locomotion.go2.joystick import ( Commands, Go2JoystickCfg, diff --git a/tests/envs/test_go2_obs_noise.py b/tests/envs/test_go2_obs_noise.py index 3d179f5c0..0eb1840a2 100644 --- a/tests/envs/test_go2_obs_noise.py +++ b/tests/envs/test_go2_obs_noise.py @@ -4,7 +4,7 @@ import numpy as np -from unilab.envs.locomotion.go2.base import Go2BaseCfg, Go2BaseEnv, NoiseConfig +from unilab.tasks.locomotion.go2.base import Go2BaseCfg, Go2BaseEnv, NoiseConfig class _ConcreteGo2Env(Go2BaseEnv): From 7e1e79101ba207b8787510a36725a3c71404088d Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 20:58:41 +0800 Subject: [PATCH 041/169] refactor(tasks): move Go1 rough owner [skip ci] Closes #1127. Roadmap #1042. --- src/unilab/envs/locomotion/go1/__init__.py | 1 - src/unilab/tasks/__init__.py | 1 + src/unilab/tasks/locomotion/go1/__init__.py | 3 +++ src/unilab/{envs => tasks}/locomotion/go1/rough.py | 0 tests/envs/locomotion/test_go2_terrain_spawn.py | 2 +- tests/tasks/test_package_boundary.py | 1 + 6 files changed, 6 insertions(+), 2 deletions(-) create mode 100644 src/unilab/tasks/locomotion/go1/__init__.py rename src/unilab/{envs => tasks}/locomotion/go1/rough.py (100%) diff --git a/src/unilab/envs/locomotion/go1/__init__.py b/src/unilab/envs/locomotion/go1/__init__.py index b0608c22b..40667c153 100644 --- a/src/unilab/envs/locomotion/go1/__init__.py +++ b/src/unilab/envs/locomotion/go1/__init__.py @@ -1,2 +1 @@ from .joystick import Go1JoystickCfg, Go1WalkTask -from .rough import Go1JoystickRoughCfg, Go1JoystickRoughEnv diff --git a/src/unilab/tasks/__init__.py b/src/unilab/tasks/__init__.py index 448745635..096222dd0 100644 --- a/src/unilab/tasks/__init__.py +++ b/src/unilab/tasks/__init__.py @@ -9,6 +9,7 @@ __unilab_registry_modules__ = ( "unilab.envs.locomotion.go1", + "unilab.tasks.locomotion.go1.rough", "unilab.tasks.locomotion.go2", "unilab.envs.locomotion.go2w", "unilab.envs.locomotion.g1", diff --git a/src/unilab/tasks/locomotion/go1/__init__.py b/src/unilab/tasks/locomotion/go1/__init__.py new file mode 100644 index 000000000..3081345cf --- /dev/null +++ b/src/unilab/tasks/locomotion/go1/__init__.py @@ -0,0 +1,3 @@ +from .rough import Go1JoystickRoughCfg, Go1JoystickRoughEnv + +__all__ = ["Go1JoystickRoughCfg", "Go1JoystickRoughEnv"] diff --git a/src/unilab/envs/locomotion/go1/rough.py b/src/unilab/tasks/locomotion/go1/rough.py similarity index 100% rename from src/unilab/envs/locomotion/go1/rough.py rename to src/unilab/tasks/locomotion/go1/rough.py diff --git a/tests/envs/locomotion/test_go2_terrain_spawn.py b/tests/envs/locomotion/test_go2_terrain_spawn.py index 4587c0a3c..0c3c7eafe 100644 --- a/tests/envs/locomotion/test_go2_terrain_spawn.py +++ b/tests/envs/locomotion/test_go2_terrain_spawn.py @@ -87,7 +87,7 @@ def test_terrain_spawn_attached_when_rough_motrix(): def test_go1_rough_initialization_and_reset_use_backend_terrain_contract(): - from unilab.envs.locomotion.go1.rough import ( + from unilab.tasks.locomotion.go1.rough import ( Go1JoystickRoughCfg, Go1JoystickRoughEnv, RoughRewardConfig, diff --git a/tests/tasks/test_package_boundary.py b/tests/tasks/test_package_boundary.py index d1d60b051..ddfcdd05d 100644 --- a/tests/tasks/test_package_boundary.py +++ b/tests/tasks/test_package_boundary.py @@ -13,6 +13,7 @@ _TASK_REGISTRY_MODULES = ( "unilab.envs.locomotion.go1", + "unilab.tasks.locomotion.go1.rough", "unilab.tasks.locomotion.go2", "unilab.envs.locomotion.go2w", "unilab.envs.locomotion.g1", From 605f8bf1832af00a6969f6e0a3311cfa50c11af5 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 21:03:59 +0800 Subject: [PATCH 042/169] refactor(tasks): move Go1 joystick owner [skip ci] Closes #1129. Roadmap #1042. --- scripts/benchmark/benchmark_drake_performance.py | 4 ++-- scripts/benchmark/core/task_names.py | 2 +- scripts/benchmark/env/benchmark_env_step.py | 4 ++-- src/unilab/envs/locomotion/go1/__init__.py | 2 +- src/unilab/tasks/__init__.py | 3 +-- src/unilab/tasks/locomotion/go1/__init__.py | 8 +++++++- src/unilab/{envs => tasks}/locomotion/go1/joystick.py | 6 +++--- src/unilab/tasks/locomotion/go1/rough.py | 2 +- tests/base/test_reward_override.py | 2 +- tests/integration/test_reward_injection_integration.py | 4 ++-- tests/tasks/test_package_boundary.py | 3 +-- 11 files changed, 22 insertions(+), 18 deletions(-) rename src/unilab/{envs => tasks}/locomotion/go1/joystick.py (97%) diff --git a/scripts/benchmark/benchmark_drake_performance.py b/scripts/benchmark/benchmark_drake_performance.py index a5e405eee..3b391546b 100644 --- a/scripts/benchmark/benchmark_drake_performance.py +++ b/scripts/benchmark/benchmark_drake_performance.py @@ -101,12 +101,12 @@ def wrapped(*args: Any, **kwargs: Any) -> Any: def _task_specs() -> dict[str, TaskSpec]: def go1_cfg() -> Any: - from unilab.envs.locomotion.go1.joystick import Go1JoystickCfg + from unilab.tasks.locomotion.go1.joystick import Go1JoystickCfg return Go1JoystickCfg() def go1_env() -> type: - from unilab.envs.locomotion.go1.joystick import Go1WalkTask + from unilab.tasks.locomotion.go1.joystick import Go1WalkTask return Go1WalkTask diff --git a/scripts/benchmark/core/task_names.py b/scripts/benchmark/core/task_names.py index 1a914fc53..52f348e31 100644 --- a/scripts/benchmark/core/task_names.py +++ b/scripts/benchmark/core/task_names.py @@ -3,8 +3,8 @@ from dataclasses import dataclass from unilab.envs.locomotion.g1.joystick import G1WalkFlatCfg -from unilab.envs.locomotion.go1.joystick import Go1JoystickCfg from unilab.envs.manipulation.sharpa_inhand.rotation import SharpaInhandRotationCfg +from unilab.tasks.locomotion.go1.joystick import Go1JoystickCfg from unilab.tasks.locomotion.go2.joystick import Go2JoystickCfg diff --git a/scripts/benchmark/env/benchmark_env_step.py b/scripts/benchmark/env/benchmark_env_step.py index bfc45a5fc..e96e29c6a 100644 --- a/scripts/benchmark/env/benchmark_env_step.py +++ b/scripts/benchmark/env/benchmark_env_step.py @@ -276,13 +276,13 @@ def _materialize_sharpa_motrix_scene() -> str: def _go1_cfg(backend: str, config_overrides: list[str]) -> Any: - from unilab.envs.locomotion.go1.joystick import Go1JoystickCfg + from unilab.tasks.locomotion.go1.joystick import Go1JoystickCfg return _ppo_owner_yaml_cfg("go1_joystick_flat", backend, Go1JoystickCfg, config_overrides) def _go1_env_cls() -> type: - from unilab.envs.locomotion.go1.joystick import Go1WalkTask + from unilab.tasks.locomotion.go1.joystick import Go1WalkTask return Go1WalkTask diff --git a/src/unilab/envs/locomotion/go1/__init__.py b/src/unilab/envs/locomotion/go1/__init__.py index 40667c153..7d4b2484f 100644 --- a/src/unilab/envs/locomotion/go1/__init__.py +++ b/src/unilab/envs/locomotion/go1/__init__.py @@ -1 +1 @@ -from .joystick import Go1JoystickCfg, Go1WalkTask +"""Legacy Go1 shared base pending migration to :mod:`unilab.tasks`.""" diff --git a/src/unilab/tasks/__init__.py b/src/unilab/tasks/__init__.py index 096222dd0..64fd320bf 100644 --- a/src/unilab/tasks/__init__.py +++ b/src/unilab/tasks/__init__.py @@ -8,8 +8,7 @@ """ __unilab_registry_modules__ = ( - "unilab.envs.locomotion.go1", - "unilab.tasks.locomotion.go1.rough", + "unilab.tasks.locomotion.go1", "unilab.tasks.locomotion.go2", "unilab.envs.locomotion.go2w", "unilab.envs.locomotion.g1", diff --git a/src/unilab/tasks/locomotion/go1/__init__.py b/src/unilab/tasks/locomotion/go1/__init__.py index 3081345cf..4f26efac3 100644 --- a/src/unilab/tasks/locomotion/go1/__init__.py +++ b/src/unilab/tasks/locomotion/go1/__init__.py @@ -1,3 +1,9 @@ +from .joystick import Go1JoystickCfg, Go1WalkTask from .rough import Go1JoystickRoughCfg, Go1JoystickRoughEnv -__all__ = ["Go1JoystickRoughCfg", "Go1JoystickRoughEnv"] +__all__ = [ + "Go1JoystickCfg", + "Go1JoystickRoughCfg", + "Go1JoystickRoughEnv", + "Go1WalkTask", +] diff --git a/src/unilab/envs/locomotion/go1/joystick.py b/src/unilab/tasks/locomotion/go1/joystick.py similarity index 97% rename from src/unilab/envs/locomotion/go1/joystick.py rename to src/unilab/tasks/locomotion/go1/joystick.py index 9b798ecc2..b635cfd34 100644 --- a/src/unilab/envs/locomotion/go1/joystick.py +++ b/src/unilab/tasks/locomotion/go1/joystick.py @@ -46,12 +46,12 @@ class JoystickSensor: @registry.envcfg("Go1JoystickFlat") @dataclass class Go1JoystickCfg(Go1BaseCfg): - scene: SceneCfg = field( + scene: SceneCfg = field( # pyright: ignore[reportIncompatibleVariableOverride] default_factory=lambda: SceneCfg( model_file=str(ASSETS_ROOT_PATH / "robots" / "go1" / "scene_flat.xml") ) ) - max_episode_seconds: float = 20.0 + max_episode_seconds: float = 20.0 # pyright: ignore[reportIncompatibleVariableOverride] init_state: InitState = field(default_factory=InitState) commands: Commands = field(default_factory=Commands) reward_config: RewardConfig | None = None @@ -86,7 +86,7 @@ def _compute_reset_obs( @registry.env("Go1JoystickFlat", sim_backend="motrix") @registry.env("Go1JoystickFlat", sim_backend="drake") class Go1WalkTask(Go1BaseEnv): - _cfg: Go1JoystickCfg + _cfg: Go1JoystickCfg # pyright: ignore[reportIncompatibleVariableOverride] def __init__(self, cfg: Go1JoystickCfg, num_envs=1, backend_type="mujoco"): if cfg.reward_config is None: diff --git a/src/unilab/tasks/locomotion/go1/rough.py b/src/unilab/tasks/locomotion/go1/rough.py index 9f67c2305..92b5520b5 100644 --- a/src/unilab/tasks/locomotion/go1/rough.py +++ b/src/unilab/tasks/locomotion/go1/rough.py @@ -35,7 +35,7 @@ TerrainCurriculumCfg, ) from unilab.envs.locomotion.go1.base import ControlConfig -from unilab.envs.locomotion.go1.joystick import ( +from unilab.tasks.locomotion.go1.joystick import ( Go1JoystickCfg, Go1JoystickDomainRandomizationProvider, Go1WalkTask, diff --git a/tests/base/test_reward_override.py b/tests/base/test_reward_override.py index cd9b17efa..f8d9be2fa 100644 --- a/tests/base/test_reward_override.py +++ b/tests/base/test_reward_override.py @@ -12,7 +12,7 @@ def test_reward_override_go1(): """Test Go1 reward config override.""" ensure_registries() - from unilab.envs.locomotion.go1.joystick import RewardConfig + from unilab.tasks.locomotion.go1.joystick import RewardConfig override_config = RewardConfig( scales={"tracking_lin_vel": 999.0}, diff --git a/tests/integration/test_reward_injection_integration.py b/tests/integration/test_reward_injection_integration.py index 7b5c7ae96..e549e0a90 100644 --- a/tests/integration/test_reward_injection_integration.py +++ b/tests/integration/test_reward_injection_integration.py @@ -38,7 +38,7 @@ def test_reward_override_propagation(): """Test reward override propagates through multiprocess collector.""" from unilab.base import registry from unilab.base.registry import ensure_registries - from unilab.envs.locomotion.go1.joystick import RewardConfig + from unilab.tasks.locomotion.go1.joystick import RewardConfig ensure_registries() @@ -102,7 +102,7 @@ def test_zero_scale_skips_computation(): """Test that reward functions with scale=0 are skipped.""" from unilab.base import registry from unilab.base.registry import ensure_registries - from unilab.envs.locomotion.go1.joystick import RewardConfig + from unilab.tasks.locomotion.go1.joystick import RewardConfig ensure_registries() diff --git a/tests/tasks/test_package_boundary.py b/tests/tasks/test_package_boundary.py index ddfcdd05d..d7d0bc14f 100644 --- a/tests/tasks/test_package_boundary.py +++ b/tests/tasks/test_package_boundary.py @@ -12,8 +12,7 @@ _ENV_PACKAGE = _REPO_ROOT / "src" / "unilab" / "envs" _TASK_REGISTRY_MODULES = ( - "unilab.envs.locomotion.go1", - "unilab.tasks.locomotion.go1.rough", + "unilab.tasks.locomotion.go1", "unilab.tasks.locomotion.go2", "unilab.envs.locomotion.go2w", "unilab.envs.locomotion.g1", From 36fab0c4777b421d8a4a8e1c20467d858c938911 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 21:11:13 +0800 Subject: [PATCH 043/169] refactor(tasks): move Go1 shared base [skip ci] Closes #1131. Roadmap #1042. --- docs/sphinx/source/api_reference/envs/locomotion.md | 2 +- src/unilab/envs/locomotion/__init__.py | 1 - src/unilab/envs/locomotion/go1/__init__.py | 1 - src/unilab/{envs => tasks}/locomotion/go1/base.py | 2 +- src/unilab/tasks/locomotion/go1/joystick.py | 2 +- src/unilab/tasks/locomotion/go1/rough.py | 2 +- tests/envs/test_go1_obs_noise.py | 2 +- 7 files changed, 5 insertions(+), 7 deletions(-) delete mode 100644 src/unilab/envs/locomotion/go1/__init__.py rename src/unilab/{envs => tasks}/locomotion/go1/base.py (91%) diff --git a/docs/sphinx/source/api_reference/envs/locomotion.md b/docs/sphinx/source/api_reference/envs/locomotion.md index 66ddebc84..570a57ae7 100644 --- a/docs/sphinx/source/api_reference/envs/locomotion.md +++ b/docs/sphinx/source/api_reference/envs/locomotion.md @@ -8,7 +8,7 @@ unilab.envs.locomotion.common unilab.envs.locomotion.g1 - unilab.envs.locomotion.go1 + unilab.tasks.locomotion.go1 unilab.tasks.locomotion.go2 unilab.envs.locomotion.go2_arm unilab.envs.locomotion.go2w diff --git a/src/unilab/envs/locomotion/__init__.py b/src/unilab/envs/locomotion/__init__.py index 3daf9340e..7e06f1036 100644 --- a/src/unilab/envs/locomotion/__init__.py +++ b/src/unilab/envs/locomotion/__init__.py @@ -1,7 +1,6 @@ """Locomotion env registry bootstrap contract.""" __unilab_registry_modules__ = ( - "unilab.envs.locomotion.go1", "unilab.envs.locomotion.go2w", "unilab.envs.locomotion.g1", "unilab.envs.locomotion.go2_arm", diff --git a/src/unilab/envs/locomotion/go1/__init__.py b/src/unilab/envs/locomotion/go1/__init__.py deleted file mode 100644 index 7d4b2484f..000000000 --- a/src/unilab/envs/locomotion/go1/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Legacy Go1 shared base pending migration to :mod:`unilab.tasks`.""" diff --git a/src/unilab/envs/locomotion/go1/base.py b/src/unilab/tasks/locomotion/go1/base.py similarity index 91% rename from src/unilab/envs/locomotion/go1/base.py rename to src/unilab/tasks/locomotion/go1/base.py index 06b841baa..3869a9fd6 100644 --- a/src/unilab/envs/locomotion/go1/base.py +++ b/src/unilab/tasks/locomotion/go1/base.py @@ -37,4 +37,4 @@ class Go1BaseCfg(LocomotionBaseCfg): class Go1BaseEnv(LocomotionBaseEnv): - _cfg: Go1BaseCfg + _cfg: Go1BaseCfg # pyright: ignore[reportIncompatibleVariableOverride] diff --git a/src/unilab/tasks/locomotion/go1/joystick.py b/src/unilab/tasks/locomotion/go1/joystick.py index b635cfd34..75289333a 100644 --- a/src/unilab/tasks/locomotion/go1/joystick.py +++ b/src/unilab/tasks/locomotion/go1/joystick.py @@ -20,7 +20,7 @@ TerrainCurriculumCfg, TerrainSpawnManager, ) -from unilab.envs.locomotion.go1.base import Go1BaseCfg, Go1BaseEnv +from unilab.tasks.locomotion.go1.base import Go1BaseCfg, Go1BaseEnv @dataclass diff --git a/src/unilab/tasks/locomotion/go1/rough.py b/src/unilab/tasks/locomotion/go1/rough.py index 92b5520b5..82b12f6da 100644 --- a/src/unilab/tasks/locomotion/go1/rough.py +++ b/src/unilab/tasks/locomotion/go1/rough.py @@ -34,7 +34,7 @@ from unilab.envs.locomotion.common.terrain_spawn import ( TerrainCurriculumCfg, ) -from unilab.envs.locomotion.go1.base import ControlConfig +from unilab.tasks.locomotion.go1.base import ControlConfig from unilab.tasks.locomotion.go1.joystick import ( Go1JoystickCfg, Go1JoystickDomainRandomizationProvider, diff --git a/tests/envs/test_go1_obs_noise.py b/tests/envs/test_go1_obs_noise.py index cc8815a23..a70eed7ae 100644 --- a/tests/envs/test_go1_obs_noise.py +++ b/tests/envs/test_go1_obs_noise.py @@ -4,7 +4,7 @@ import numpy as np -from unilab.envs.locomotion.go1.base import Go1BaseCfg, Go1BaseEnv, NoiseConfig +from unilab.tasks.locomotion.go1.base import Go1BaseCfg, Go1BaseEnv, NoiseConfig class _ConcreteGo1Env(Go1BaseEnv): From 27d3a8ffb3526a29d5c05124d586b4c243be4f9c Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 21:16:01 +0800 Subject: [PATCH 044/169] refactor(tasks): move Go2W rough owner [skip ci] Closes #1133. Roadmap #1042. --- scripts/benchmark/env/benchmark_env_step.py | 4 ++-- src/unilab/envs/locomotion/go2w/__init__.py | 1 - src/unilab/tasks/__init__.py | 1 + src/unilab/tasks/locomotion/go2w/__init__.py | 3 +++ src/unilab/{envs => tasks}/locomotion/go2w/rough.py | 0 tests/envs/locomotion/go2w/test_go2w_height_scan.py | 2 +- tests/envs/locomotion/test_go2_terrain_spawn.py | 2 +- tests/tasks/test_package_boundary.py | 1 + 8 files changed, 9 insertions(+), 5 deletions(-) create mode 100644 src/unilab/tasks/locomotion/go2w/__init__.py rename src/unilab/{envs => tasks}/locomotion/go2w/rough.py (100%) diff --git a/scripts/benchmark/env/benchmark_env_step.py b/scripts/benchmark/env/benchmark_env_step.py index e96e29c6a..02f1a20ea 100644 --- a/scripts/benchmark/env/benchmark_env_step.py +++ b/scripts/benchmark/env/benchmark_env_step.py @@ -318,7 +318,7 @@ def _go2w_cfg(backend: str, config_overrides: list[str]) -> Any: def _go2w_rough_cfg(backend: str, config_overrides: list[str]) -> Any: - from unilab.envs.locomotion.go2w.rough import Go2WJoystickRoughCfg + from unilab.tasks.locomotion.go2w.rough import Go2WJoystickRoughCfg return _ppo_owner_yaml_cfg( "go2w_joystick_rough", backend, Go2WJoystickRoughCfg, config_overrides @@ -332,7 +332,7 @@ def _go2w_env_cls() -> type: def _go2w_rough_env_cls() -> type: - from unilab.envs.locomotion.go2w.rough import Go2WJoystickRoughEnv + from unilab.tasks.locomotion.go2w.rough import Go2WJoystickRoughEnv return Go2WJoystickRoughEnv diff --git a/src/unilab/envs/locomotion/go2w/__init__.py b/src/unilab/envs/locomotion/go2w/__init__.py index 35727958a..4eeda31b1 100644 --- a/src/unilab/envs/locomotion/go2w/__init__.py +++ b/src/unilab/envs/locomotion/go2w/__init__.py @@ -1,2 +1 @@ from .joystick import Go2WJoystickCfg, Go2WJoystickEnv -from .rough import Go2WJoystickRoughCfg, Go2WJoystickRoughEnv diff --git a/src/unilab/tasks/__init__.py b/src/unilab/tasks/__init__.py index 64fd320bf..96c3f6feb 100644 --- a/src/unilab/tasks/__init__.py +++ b/src/unilab/tasks/__init__.py @@ -11,6 +11,7 @@ "unilab.tasks.locomotion.go1", "unilab.tasks.locomotion.go2", "unilab.envs.locomotion.go2w", + "unilab.tasks.locomotion.go2w.rough", "unilab.envs.locomotion.g1", "unilab.envs.locomotion.go2_arm", "unilab.tasks.locomotion.a2", diff --git a/src/unilab/tasks/locomotion/go2w/__init__.py b/src/unilab/tasks/locomotion/go2w/__init__.py new file mode 100644 index 000000000..cbd7059ab --- /dev/null +++ b/src/unilab/tasks/locomotion/go2w/__init__.py @@ -0,0 +1,3 @@ +from .rough import Go2WJoystickRoughCfg, Go2WJoystickRoughEnv + +__all__ = ["Go2WJoystickRoughCfg", "Go2WJoystickRoughEnv"] diff --git a/src/unilab/envs/locomotion/go2w/rough.py b/src/unilab/tasks/locomotion/go2w/rough.py similarity index 100% rename from src/unilab/envs/locomotion/go2w/rough.py rename to src/unilab/tasks/locomotion/go2w/rough.py diff --git a/tests/envs/locomotion/go2w/test_go2w_height_scan.py b/tests/envs/locomotion/go2w/test_go2w_height_scan.py index 538c0410e..c9e1a2c87 100644 --- a/tests/envs/locomotion/go2w/test_go2w_height_scan.py +++ b/tests/envs/locomotion/go2w/test_go2w_height_scan.py @@ -21,7 +21,7 @@ allow_module_level=True, ) -from unilab.envs.locomotion.go2w.rough import Go2WJoystickRoughEnv +from unilab.tasks.locomotion.go2w.rough import Go2WJoystickRoughEnv def test_go2w_rough_height_scan_uses_backend_native_sampling() -> None: diff --git a/tests/envs/locomotion/test_go2_terrain_spawn.py b/tests/envs/locomotion/test_go2_terrain_spawn.py index 0c3c7eafe..eb58ac392 100644 --- a/tests/envs/locomotion/test_go2_terrain_spawn.py +++ b/tests/envs/locomotion/test_go2_terrain_spawn.py @@ -116,7 +116,7 @@ def test_go1_rough_initialization_and_reset_use_backend_terrain_contract(): def test_go2w_rough_initialization_and_reset_use_backend_terrain_contract(): from unilab.envs.locomotion.go2w.joystick import RewardConfig - from unilab.envs.locomotion.go2w.rough import Go2WJoystickRoughCfg, Go2WJoystickRoughEnv + from unilab.tasks.locomotion.go2w.rough import Go2WJoystickRoughCfg, Go2WJoystickRoughEnv cfg = Go2WJoystickRoughCfg( reward_config=RewardConfig(scales={}, tracking_sigma=0.25, base_height_target=0.3) diff --git a/tests/tasks/test_package_boundary.py b/tests/tasks/test_package_boundary.py index d7d0bc14f..4cfa7ab75 100644 --- a/tests/tasks/test_package_boundary.py +++ b/tests/tasks/test_package_boundary.py @@ -15,6 +15,7 @@ "unilab.tasks.locomotion.go1", "unilab.tasks.locomotion.go2", "unilab.envs.locomotion.go2w", + "unilab.tasks.locomotion.go2w.rough", "unilab.envs.locomotion.g1", "unilab.envs.locomotion.go2_arm", "unilab.tasks.locomotion.a2", From 8a6e4dad8445d5319fc5fdae14a9904620ce4642 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 21:21:13 +0800 Subject: [PATCH 045/169] refactor(tasks): move Go2W joystick owner [skip ci] Closes #1135. Roadmap #1042. --- scripts/benchmark/env/benchmark_env_step.py | 4 ++-- src/unilab/envs/locomotion/go2w/__init__.py | 2 +- src/unilab/tasks/__init__.py | 3 +-- src/unilab/tasks/locomotion/go2w/__init__.py | 8 +++++++- src/unilab/{envs => tasks}/locomotion/go2w/joystick.py | 6 +++--- src/unilab/tasks/locomotion/go2w/rough.py | 2 +- tests/envs/locomotion/go2w/test_go2w_motor_control.py | 4 ++-- tests/envs/locomotion/test_go2_terrain_spawn.py | 2 +- tests/tasks/test_package_boundary.py | 3 +-- 9 files changed, 19 insertions(+), 15 deletions(-) rename src/unilab/{envs => tasks}/locomotion/go2w/joystick.py (99%) diff --git a/scripts/benchmark/env/benchmark_env_step.py b/scripts/benchmark/env/benchmark_env_step.py index 02f1a20ea..77b137026 100644 --- a/scripts/benchmark/env/benchmark_env_step.py +++ b/scripts/benchmark/env/benchmark_env_step.py @@ -312,7 +312,7 @@ def _go2_rough_env_cls() -> type: def _go2w_cfg(backend: str, config_overrides: list[str]) -> Any: - from unilab.envs.locomotion.go2w.joystick import Go2WJoystickCfg + from unilab.tasks.locomotion.go2w.joystick import Go2WJoystickCfg return _ppo_owner_yaml_cfg("go2w_joystick_flat", backend, Go2WJoystickCfg, config_overrides) @@ -326,7 +326,7 @@ def _go2w_rough_cfg(backend: str, config_overrides: list[str]) -> Any: def _go2w_env_cls() -> type: - from unilab.envs.locomotion.go2w.joystick import Go2WJoystickEnv + from unilab.tasks.locomotion.go2w.joystick import Go2WJoystickEnv return Go2WJoystickEnv diff --git a/src/unilab/envs/locomotion/go2w/__init__.py b/src/unilab/envs/locomotion/go2w/__init__.py index 4eeda31b1..f8e0e66b7 100644 --- a/src/unilab/envs/locomotion/go2w/__init__.py +++ b/src/unilab/envs/locomotion/go2w/__init__.py @@ -1 +1 @@ -from .joystick import Go2WJoystickCfg, Go2WJoystickEnv +"""Legacy Go2W shared base pending migration to :mod:`unilab.tasks`.""" diff --git a/src/unilab/tasks/__init__.py b/src/unilab/tasks/__init__.py index 96c3f6feb..d07c408de 100644 --- a/src/unilab/tasks/__init__.py +++ b/src/unilab/tasks/__init__.py @@ -10,8 +10,7 @@ __unilab_registry_modules__ = ( "unilab.tasks.locomotion.go1", "unilab.tasks.locomotion.go2", - "unilab.envs.locomotion.go2w", - "unilab.tasks.locomotion.go2w.rough", + "unilab.tasks.locomotion.go2w", "unilab.envs.locomotion.g1", "unilab.envs.locomotion.go2_arm", "unilab.tasks.locomotion.a2", diff --git a/src/unilab/tasks/locomotion/go2w/__init__.py b/src/unilab/tasks/locomotion/go2w/__init__.py index cbd7059ab..d82303136 100644 --- a/src/unilab/tasks/locomotion/go2w/__init__.py +++ b/src/unilab/tasks/locomotion/go2w/__init__.py @@ -1,3 +1,9 @@ +from .joystick import Go2WJoystickCfg, Go2WJoystickEnv from .rough import Go2WJoystickRoughCfg, Go2WJoystickRoughEnv -__all__ = ["Go2WJoystickRoughCfg", "Go2WJoystickRoughEnv"] +__all__ = [ + "Go2WJoystickCfg", + "Go2WJoystickEnv", + "Go2WJoystickRoughCfg", + "Go2WJoystickRoughEnv", +] diff --git a/src/unilab/envs/locomotion/go2w/joystick.py b/src/unilab/tasks/locomotion/go2w/joystick.py similarity index 99% rename from src/unilab/envs/locomotion/go2w/joystick.py rename to src/unilab/tasks/locomotion/go2w/joystick.py index fa1f9978c..0cf646207 100644 --- a/src/unilab/envs/locomotion/go2w/joystick.py +++ b/src/unilab/tasks/locomotion/go2w/joystick.py @@ -88,12 +88,12 @@ class JoystickSensor: @registry.envcfg("Go2WJoystickFlat") @dataclass class Go2WJoystickCfg(Go2WBaseCfg): - scene: SceneCfg = field( + scene: SceneCfg = field( # pyright: ignore[reportIncompatibleVariableOverride] default_factory=lambda: SceneCfg( model_file=str(ASSETS_ROOT_PATH / "robots" / "go2w" / "scene_flat.xml") ) ) - max_episode_seconds: float = 20.0 + max_episode_seconds: float = 20.0 # pyright: ignore[reportIncompatibleVariableOverride] init_state: InitState = field(default_factory=InitState) commands: Commands = field(default_factory=Commands) reward_config: RewardConfig | None = None @@ -229,7 +229,7 @@ def _sample_commands(self, env: Any, num_reset: int) -> np.ndarray: @registry.env("Go2WJoystickFlat", sim_backend="mujoco") @registry.env("Go2WJoystickFlat", sim_backend="drake") class Go2WJoystickEnv(Go2WBaseEnv): - _cfg: Go2WJoystickCfg + _cfg: Go2WJoystickCfg # pyright: ignore[reportIncompatibleVariableOverride] def __init__(self, cfg: Go2WJoystickCfg, num_envs=1, backend_type="mujoco"): if cfg.reward_config is None: diff --git a/src/unilab/tasks/locomotion/go2w/rough.py b/src/unilab/tasks/locomotion/go2w/rough.py index da11d4321..c719d2462 100644 --- a/src/unilab/tasks/locomotion/go2w/rough.py +++ b/src/unilab/tasks/locomotion/go2w/rough.py @@ -32,7 +32,7 @@ TerrainSpawnManager, ) from unilab.envs.locomotion.go2w.base import NUM_GO2W_ACTIONS, NUM_LEG_ACTIONS -from unilab.envs.locomotion.go2w.joystick import ( +from unilab.tasks.locomotion.go2w.joystick import ( Go2WJoystickCfg, Go2WJoystickDomainRandomizationProvider, Go2WJoystickEnv, diff --git a/tests/envs/locomotion/go2w/test_go2w_motor_control.py b/tests/envs/locomotion/go2w/test_go2w_motor_control.py index 2eec857f5..453be7bf0 100644 --- a/tests/envs/locomotion/go2w/test_go2w_motor_control.py +++ b/tests/envs/locomotion/go2w/test_go2w_motor_control.py @@ -16,7 +16,7 @@ Go2WBaseEnv, compute_go2w_motor_ctrl, ) -from unilab.envs.locomotion.go2w.joystick import ( +from unilab.tasks.locomotion.go2w.joystick import ( Go2WJoystickCfg, Go2WJoystickDomainRandomizationProvider, Go2WJoystickEnv, @@ -135,7 +135,7 @@ def test_go2w_reset_plan_can_disable_initial_yaw_randomization() -> None: def test_go2w_init_does_not_pass_position_actuator_gains(monkeypatch: pytest.MonkeyPatch) -> None: - from unilab.envs.locomotion.go2w import joystick as go2w_module + from unilab.tasks.locomotion.go2w import joystick as go2w_module captured: dict[str, Any] = {} diff --git a/tests/envs/locomotion/test_go2_terrain_spawn.py b/tests/envs/locomotion/test_go2_terrain_spawn.py index eb58ac392..5d357ed8e 100644 --- a/tests/envs/locomotion/test_go2_terrain_spawn.py +++ b/tests/envs/locomotion/test_go2_terrain_spawn.py @@ -115,7 +115,7 @@ def test_go1_rough_initialization_and_reset_use_backend_terrain_contract(): def test_go2w_rough_initialization_and_reset_use_backend_terrain_contract(): - from unilab.envs.locomotion.go2w.joystick import RewardConfig + from unilab.tasks.locomotion.go2w.joystick import RewardConfig from unilab.tasks.locomotion.go2w.rough import Go2WJoystickRoughCfg, Go2WJoystickRoughEnv cfg = Go2WJoystickRoughCfg( diff --git a/tests/tasks/test_package_boundary.py b/tests/tasks/test_package_boundary.py index 4cfa7ab75..6bf5dae8e 100644 --- a/tests/tasks/test_package_boundary.py +++ b/tests/tasks/test_package_boundary.py @@ -14,8 +14,7 @@ _TASK_REGISTRY_MODULES = ( "unilab.tasks.locomotion.go1", "unilab.tasks.locomotion.go2", - "unilab.envs.locomotion.go2w", - "unilab.tasks.locomotion.go2w.rough", + "unilab.tasks.locomotion.go2w", "unilab.envs.locomotion.g1", "unilab.envs.locomotion.go2_arm", "unilab.tasks.locomotion.a2", From 1a7c452cd50fe4291a2dbe0ea1625322b3ef7cc4 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 21:25:58 +0800 Subject: [PATCH 046/169] refactor(tasks): move Go2W shared base [skip ci] Closes #1137. Roadmap #1042. --- docs/sphinx/source/api_reference/envs/locomotion.md | 2 +- src/unilab/envs/locomotion/__init__.py | 1 - src/unilab/envs/locomotion/go2w/__init__.py | 1 - src/unilab/{envs => tasks}/locomotion/go2w/base.py | 2 +- src/unilab/tasks/locomotion/go2w/joystick.py | 2 +- src/unilab/tasks/locomotion/go2w/rough.py | 2 +- tests/base/test_sim_backend_smoke.py | 2 +- tests/envs/locomotion/go2w/test_go2w_motor_control.py | 2 +- 8 files changed, 6 insertions(+), 8 deletions(-) delete mode 100644 src/unilab/envs/locomotion/go2w/__init__.py rename src/unilab/{envs => tasks}/locomotion/go2w/base.py (98%) diff --git a/docs/sphinx/source/api_reference/envs/locomotion.md b/docs/sphinx/source/api_reference/envs/locomotion.md index 570a57ae7..490cdd8ee 100644 --- a/docs/sphinx/source/api_reference/envs/locomotion.md +++ b/docs/sphinx/source/api_reference/envs/locomotion.md @@ -11,5 +11,5 @@ unilab.tasks.locomotion.go1 unilab.tasks.locomotion.go2 unilab.envs.locomotion.go2_arm - unilab.envs.locomotion.go2w + unilab.tasks.locomotion.go2w ``` diff --git a/src/unilab/envs/locomotion/__init__.py b/src/unilab/envs/locomotion/__init__.py index 7e06f1036..4e1c8a3eb 100644 --- a/src/unilab/envs/locomotion/__init__.py +++ b/src/unilab/envs/locomotion/__init__.py @@ -1,7 +1,6 @@ """Locomotion env registry bootstrap contract.""" __unilab_registry_modules__ = ( - "unilab.envs.locomotion.go2w", "unilab.envs.locomotion.g1", "unilab.envs.locomotion.go2_arm", ) diff --git a/src/unilab/envs/locomotion/go2w/__init__.py b/src/unilab/envs/locomotion/go2w/__init__.py deleted file mode 100644 index f8e0e66b7..000000000 --- a/src/unilab/envs/locomotion/go2w/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Legacy Go2W shared base pending migration to :mod:`unilab.tasks`.""" diff --git a/src/unilab/envs/locomotion/go2w/base.py b/src/unilab/tasks/locomotion/go2w/base.py similarity index 98% rename from src/unilab/envs/locomotion/go2w/base.py rename to src/unilab/tasks/locomotion/go2w/base.py index b48e5b07c..3875e1919 100644 --- a/src/unilab/envs/locomotion/go2w/base.py +++ b/src/unilab/tasks/locomotion/go2w/base.py @@ -117,7 +117,7 @@ def compute_go2w_motor_ctrl( class Go2WBaseEnv(LocomotionBaseEnv): - _cfg: Go2WBaseCfg + _cfg: Go2WBaseCfg # pyright: ignore[reportIncompatibleVariableOverride] def _init_action_space(self) -> None: self._action_space = gym.spaces.Box( diff --git a/src/unilab/tasks/locomotion/go2w/joystick.py b/src/unilab/tasks/locomotion/go2w/joystick.py index 0cf646207..dc13a3e78 100644 --- a/src/unilab/tasks/locomotion/go2w/joystick.py +++ b/src/unilab/tasks/locomotion/go2w/joystick.py @@ -29,7 +29,7 @@ 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.go2w.base import ( +from unilab.tasks.locomotion.go2w.base import ( DEFAULT_GO2W_ANGLES, NUM_GO2W_ACTIONS, NUM_LEG_ACTIONS, diff --git a/src/unilab/tasks/locomotion/go2w/rough.py b/src/unilab/tasks/locomotion/go2w/rough.py index c719d2462..de1db4ee5 100644 --- a/src/unilab/tasks/locomotion/go2w/rough.py +++ b/src/unilab/tasks/locomotion/go2w/rough.py @@ -31,7 +31,7 @@ TerrainCurriculumCfg, TerrainSpawnManager, ) -from unilab.envs.locomotion.go2w.base import NUM_GO2W_ACTIONS, NUM_LEG_ACTIONS +from unilab.tasks.locomotion.go2w.base import NUM_GO2W_ACTIONS, NUM_LEG_ACTIONS from unilab.tasks.locomotion.go2w.joystick import ( Go2WJoystickCfg, Go2WJoystickDomainRandomizationProvider, diff --git a/tests/base/test_sim_backend_smoke.py b/tests/base/test_sim_backend_smoke.py index bc4ab6322..ee3b9fb6f 100644 --- a/tests/base/test_sim_backend_smoke.py +++ b/tests/base/test_sim_backend_smoke.py @@ -367,7 +367,7 @@ def test_backend_batch_sensor_data_matches_individual_sensors(backend_type): pytest.importorskip("motrixsim") from unilab.base.backend import create_backend - from unilab.envs.locomotion.go2w.base import JOINT_SENSOR_PREFIXES + from unilab.tasks.locomotion.go2w.base import JOINT_SENSOR_PREFIXES bkd = create_backend( backend_type, diff --git a/tests/envs/locomotion/go2w/test_go2w_motor_control.py b/tests/envs/locomotion/go2w/test_go2w_motor_control.py index 453be7bf0..96be2d5a5 100644 --- a/tests/envs/locomotion/go2w/test_go2w_motor_control.py +++ b/tests/envs/locomotion/go2w/test_go2w_motor_control.py @@ -7,7 +7,7 @@ import pytest from unilab.base.np_env import NpEnvState -from unilab.envs.locomotion.go2w.base import ( +from unilab.tasks.locomotion.go2w.base import ( DEFAULT_GO2W_ANGLES, JOINT_SENSOR_PREFIXES, NUM_GO2W_ACTIONS, From e9dffb29ca4206464041371fce2c166d19f54df2 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 21:32:04 +0800 Subject: [PATCH 047/169] refactor(tasks): move Go2 Arm manip-loco owner [skip ci] Closes #1139. Roadmap #1042. --- scripts/manip_loco/diagnose_go2_arm_ik.py | 2 +- scripts/manip_loco/play_go2_arm_ik_only.py | 2 +- src/unilab/envs/locomotion/go2_arm/__init__.py | 2 +- src/unilab/tasks/__init__.py | 2 +- .../tasks/locomotion/go2_arm/__init__.py | 3 +++ .../locomotion/go2_arm/manip_loco.py | 8 +++++--- .../go2_arm/test_manip_loco_contract.py | 18 +++++++++--------- .../scripts/test_visualization_entrypoints.py | 2 +- tests/tasks/test_package_boundary.py | 2 +- 9 files changed, 23 insertions(+), 18 deletions(-) create mode 100644 src/unilab/tasks/locomotion/go2_arm/__init__.py rename src/unilab/{envs => tasks}/locomotion/go2_arm/manip_loco.py (99%) diff --git a/scripts/manip_loco/diagnose_go2_arm_ik.py b/scripts/manip_loco/diagnose_go2_arm_ik.py index 601eb70b7..d859fba07 100644 --- a/scripts/manip_loco/diagnose_go2_arm_ik.py +++ b/scripts/manip_loco/diagnose_go2_arm_ik.py @@ -16,7 +16,7 @@ from unilab.base import registry from unilab.base.registry import ensure_registries -from unilab.envs.locomotion.go2_arm.manip_loco import RewardConfig +from unilab.tasks.locomotion.go2_arm.manip_loco import RewardConfig from unilab.utils.rotation import np_matrix_from_quat diff --git a/scripts/manip_loco/play_go2_arm_ik_only.py b/scripts/manip_loco/play_go2_arm_ik_only.py index dce252814..2fbcb4572 100644 --- a/scripts/manip_loco/play_go2_arm_ik_only.py +++ b/scripts/manip_loco/play_go2_arm_ik_only.py @@ -19,7 +19,7 @@ sys.path.insert(0, str(ROOT_DIR)) from unilab.envs.locomotion.go2_arm.base import build_go2_arm_position_gains -from unilab.envs.locomotion.go2_arm.manip_loco import Go2ArmManipLocoCfg +from unilab.tasks.locomotion.go2_arm.manip_loco import Go2ArmManipLocoCfg TARGET_BODY = "ik_mocap_target" diff --git a/src/unilab/envs/locomotion/go2_arm/__init__.py b/src/unilab/envs/locomotion/go2_arm/__init__.py index cf26aa650..4626e666f 100644 --- a/src/unilab/envs/locomotion/go2_arm/__init__.py +++ b/src/unilab/envs/locomotion/go2_arm/__init__.py @@ -1 +1 @@ -from .manip_loco import Go2ArmManipLocoCfg, Go2ArmManipLocoEnv +"""Legacy Go2 Arm shared base pending migration to :mod:`unilab.tasks`.""" diff --git a/src/unilab/tasks/__init__.py b/src/unilab/tasks/__init__.py index d07c408de..14c8be640 100644 --- a/src/unilab/tasks/__init__.py +++ b/src/unilab/tasks/__init__.py @@ -12,7 +12,7 @@ "unilab.tasks.locomotion.go2", "unilab.tasks.locomotion.go2w", "unilab.envs.locomotion.g1", - "unilab.envs.locomotion.go2_arm", + "unilab.tasks.locomotion.go2_arm", "unilab.tasks.locomotion.a2", "unilab.envs.manipulation.allegro_inhand", "unilab.envs.manipulation.sharpa_inhand", diff --git a/src/unilab/tasks/locomotion/go2_arm/__init__.py b/src/unilab/tasks/locomotion/go2_arm/__init__.py new file mode 100644 index 000000000..090283ecc --- /dev/null +++ b/src/unilab/tasks/locomotion/go2_arm/__init__.py @@ -0,0 +1,3 @@ +from .manip_loco import Go2ArmManipLocoCfg, Go2ArmManipLocoEnv + +__all__ = ["Go2ArmManipLocoCfg", "Go2ArmManipLocoEnv"] diff --git a/src/unilab/envs/locomotion/go2_arm/manip_loco.py b/src/unilab/tasks/locomotion/go2_arm/manip_loco.py similarity index 99% rename from src/unilab/envs/locomotion/go2_arm/manip_loco.py rename to src/unilab/tasks/locomotion/go2_arm/manip_loco.py index 94c9b26b6..ec82e64fc 100644 --- a/src/unilab/envs/locomotion/go2_arm/manip_loco.py +++ b/src/unilab/tasks/locomotion/go2_arm/manip_loco.py @@ -150,9 +150,11 @@ class ArmStageConfig: @registry.envcfg("Go2ArmManipLoco") @dataclass class Go2ArmManipLocoCfg(Go2ArmBaseCfg): - scene: SceneCfg = field(default_factory=_default_go2_arm_scene) + scene: SceneCfg = field( # pyright: ignore[reportIncompatibleVariableOverride] + default_factory=_default_go2_arm_scene + ) model_file: str = field(default_factory=_default_go2_arm_model_file) - max_episode_seconds: float = 20.0 + max_episode_seconds: float = 20.0 # pyright: ignore[reportIncompatibleVariableOverride] init_state: InitState = field(default_factory=InitState) commands: CommandsConfig = field(default_factory=CommandsConfig) # type: ignore[assignment] reward_config: RewardConfig | None = None @@ -266,7 +268,7 @@ def _compute_reset_obs( @registry.env("Go2ArmManipLoco", sim_backend="drake") @registry.env("Go2ArmManipLoco", sim_backend="mujoco") class Go2ArmManipLocoEnv(Go2ArmBaseEnv): - _cfg: Go2ArmManipLocoCfg + _cfg: Go2ArmManipLocoCfg # pyright: ignore[reportIncompatibleVariableOverride] def __init__(self, cfg: Go2ArmManipLocoCfg, num_envs=1, backend_type="mujoco"): if cfg.reward_config is None: diff --git a/tests/envs/locomotion/go2_arm/test_manip_loco_contract.py b/tests/envs/locomotion/go2_arm/test_manip_loco_contract.py index e4a225c3c..b083b9345 100644 --- a/tests/envs/locomotion/go2_arm/test_manip_loco_contract.py +++ b/tests/envs/locomotion/go2_arm/test_manip_loco_contract.py @@ -11,7 +11,7 @@ from unilab.base.np_env import NpEnvState -_GO2_ARM_MANIP_LOCO_MODULE = "unilab.envs.locomotion.go2_arm.manip_loco" +_GO2_ARM_MANIP_LOCO_MODULE = "unilab.tasks.locomotion.go2_arm.manip_loco" _REGISTRY_MODULE = "unilab.base.registry" @@ -24,7 +24,7 @@ def _skip_if_no_mujoco(): def _default_reward_cfg(): - from unilab.envs.locomotion.go2_arm.manip_loco import RewardConfig + from unilab.tasks.locomotion.go2_arm.manip_loco import RewardConfig return RewardConfig( scales={ @@ -88,7 +88,7 @@ def test_go2_arm_manip_loco_registers_motrix_backend(): def test_go2_arm_manip_loco_cfg_declares_scene_for_playback(): """MuJoCo video playback needs the original visual scene, not only legacy model_file.""" from unilab.base.scene import SceneCfg - from unilab.envs.locomotion.go2_arm.manip_loco import ( + from unilab.tasks.locomotion.go2_arm.manip_loco import ( Go2ArmManipLocoCfg, _resolve_go2_arm_scene, ) @@ -107,7 +107,7 @@ def test_go2_arm_manip_loco_cfg_declares_scene_for_playback(): def test_go2_arm_ee_goal_collision_check_matches_reference_semantics(): """Any EE goal path sample inside the collision box or below ground is unsafe.""" - from unilab.envs.locomotion.go2_arm.manip_loco import ( + from unilab.tasks.locomotion.go2_arm.manip_loco import ( EEGoalConfig, Go2ArmManipLocoCfg, Go2ArmManipLocoEnv, @@ -131,7 +131,7 @@ def test_go2_arm_ee_goal_collision_check_matches_reference_semantics(): def test_go2_arm_command_moving_mask_includes_all_velocity_axes(): """A command is moving when vx, vy, or vyaw exceeds the motion threshold.""" - from unilab.envs.locomotion.go2_arm.manip_loco import Go2ArmManipLocoEnv + from unilab.tasks.locomotion.go2_arm.manip_loco import Go2ArmManipLocoEnv env = object.__new__(Go2ArmManipLocoEnv) clip = env._CMD_CLIP @@ -156,7 +156,7 @@ def test_go2_arm_command_moving_mask_includes_all_velocity_axes(): def test_go2_arm_command_postprocess_can_force_zero_commands(): """zero_command_prob should inject exact zero commands after small-command zeroing.""" - from unilab.envs.locomotion.go2_arm.manip_loco import ( + from unilab.tasks.locomotion.go2_arm.manip_loco import ( Go2ArmManipLocoCfg, Go2ArmManipLocoEnv, ) @@ -184,7 +184,7 @@ def test_go2_arm_command_postprocess_can_force_zero_commands(): def test_go2_arm_stand_still_reward_uses_same_command_mask(): """stand_still should not penalize leg pose under lateral, yaw, or forward commands.""" from unilab.envs.locomotion.common.rewards import RewardContext - from unilab.envs.locomotion.go2_arm.manip_loco import Go2ArmManipLocoEnv + from unilab.tasks.locomotion.go2_arm.manip_loco import Go2ArmManipLocoEnv env = object.__new__(Go2ArmManipLocoEnv) clip = env._CMD_CLIP @@ -212,7 +212,7 @@ def test_go2_arm_stand_still_reward_uses_same_command_mask(): def test_go2_arm_write_feet_phase_updates_indexed_envs(): """Resetting env subsets must write back feet_phase instead of losing fancy-index copies.""" - from unilab.envs.locomotion.go2_arm.manip_loco import Go2ArmManipLocoEnv + from unilab.tasks.locomotion.go2_arm.manip_loco import Go2ArmManipLocoEnv env = object.__new__(Go2ArmManipLocoEnv) env.phase = np.asarray([0.2, 0.4, 0.6], dtype=np.float32) @@ -229,7 +229,7 @@ def test_go2_arm_write_feet_phase_updates_indexed_envs(): def test_go2_arm_apply_action_uses_arm_action_scale_for_arm_residual(): """Leg residuals use action_scale while arm residuals use arm_action_scale.""" - from unilab.envs.locomotion.go2_arm.manip_loco import ( + from unilab.tasks.locomotion.go2_arm.manip_loco import ( Go2ArmManipLocoCfg, Go2ArmManipLocoEnv, ) diff --git a/tests/scripts/test_visualization_entrypoints.py b/tests/scripts/test_visualization_entrypoints.py index 15175ee69..d34885b0f 100644 --- a/tests/scripts/test_visualization_entrypoints.py +++ b/tests/scripts/test_visualization_entrypoints.py @@ -207,7 +207,7 @@ def test_velocity_arrows_require_velocity_command_task_and_policy_obs(): manip_loco_env = _keyboard_env( env_cls_name="Go2ArmManipLocoEnv", cfg_cls_name="Go2ArmManipLocoCfg", - module="unilab.envs.locomotion.go2_arm.manip_loco", + module="unilab.tasks.locomotion.go2_arm.manip_loco", obs_contains_command=True, ) missing_obs_command_env = _keyboard_env( diff --git a/tests/tasks/test_package_boundary.py b/tests/tasks/test_package_boundary.py index 6bf5dae8e..bba50b0b4 100644 --- a/tests/tasks/test_package_boundary.py +++ b/tests/tasks/test_package_boundary.py @@ -16,7 +16,7 @@ "unilab.tasks.locomotion.go2", "unilab.tasks.locomotion.go2w", "unilab.envs.locomotion.g1", - "unilab.envs.locomotion.go2_arm", + "unilab.tasks.locomotion.go2_arm", "unilab.tasks.locomotion.a2", "unilab.envs.manipulation.allegro_inhand", "unilab.envs.manipulation.sharpa_inhand", From 001c455d423e9cd17c263476c5c52cf8142a6bbe Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 21:40:43 +0800 Subject: [PATCH 048/169] refactor(tasks): move Go2 Arm shared base [skip ci] (#1142) --- docs/sphinx/source/api_reference/envs/locomotion.md | 2 +- scripts/manip_loco/play_go2_arm_ik_only.py | 2 +- src/unilab/envs/locomotion/__init__.py | 5 +---- src/unilab/envs/locomotion/go2_arm/__init__.py | 1 - src/unilab/{envs => tasks}/locomotion/go2_arm/base.py | 2 +- src/unilab/tasks/locomotion/go2_arm/manip_loco.py | 2 +- tests/envs/locomotion/go2_arm/test_base_ik.py | 2 +- 7 files changed, 6 insertions(+), 10 deletions(-) delete mode 100644 src/unilab/envs/locomotion/go2_arm/__init__.py rename src/unilab/{envs => tasks}/locomotion/go2_arm/base.py (99%) diff --git a/docs/sphinx/source/api_reference/envs/locomotion.md b/docs/sphinx/source/api_reference/envs/locomotion.md index 490cdd8ee..9c2876c85 100644 --- a/docs/sphinx/source/api_reference/envs/locomotion.md +++ b/docs/sphinx/source/api_reference/envs/locomotion.md @@ -10,6 +10,6 @@ unilab.envs.locomotion.g1 unilab.tasks.locomotion.go1 unilab.tasks.locomotion.go2 - unilab.envs.locomotion.go2_arm + unilab.tasks.locomotion.go2_arm unilab.tasks.locomotion.go2w ``` diff --git a/scripts/manip_loco/play_go2_arm_ik_only.py b/scripts/manip_loco/play_go2_arm_ik_only.py index 2fbcb4572..c1090ec84 100644 --- a/scripts/manip_loco/play_go2_arm_ik_only.py +++ b/scripts/manip_loco/play_go2_arm_ik_only.py @@ -18,7 +18,7 @@ if str(ROOT_DIR) not in sys.path: sys.path.insert(0, str(ROOT_DIR)) -from unilab.envs.locomotion.go2_arm.base import build_go2_arm_position_gains +from unilab.tasks.locomotion.go2_arm.base import build_go2_arm_position_gains from unilab.tasks.locomotion.go2_arm.manip_loco import Go2ArmManipLocoCfg TARGET_BODY = "ik_mocap_target" diff --git a/src/unilab/envs/locomotion/__init__.py b/src/unilab/envs/locomotion/__init__.py index 4e1c8a3eb..2792c0897 100644 --- a/src/unilab/envs/locomotion/__init__.py +++ b/src/unilab/envs/locomotion/__init__.py @@ -1,6 +1,3 @@ """Locomotion env registry bootstrap contract.""" -__unilab_registry_modules__ = ( - "unilab.envs.locomotion.g1", - "unilab.envs.locomotion.go2_arm", -) +__unilab_registry_modules__ = ("unilab.envs.locomotion.g1",) diff --git a/src/unilab/envs/locomotion/go2_arm/__init__.py b/src/unilab/envs/locomotion/go2_arm/__init__.py deleted file mode 100644 index 4626e666f..000000000 --- a/src/unilab/envs/locomotion/go2_arm/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Legacy Go2 Arm shared base pending migration to :mod:`unilab.tasks`.""" diff --git a/src/unilab/envs/locomotion/go2_arm/base.py b/src/unilab/tasks/locomotion/go2_arm/base.py similarity index 99% rename from src/unilab/envs/locomotion/go2_arm/base.py rename to src/unilab/tasks/locomotion/go2_arm/base.py index 454958322..dee8a0b68 100644 --- a/src/unilab/envs/locomotion/go2_arm/base.py +++ b/src/unilab/tasks/locomotion/go2_arm/base.py @@ -128,7 +128,7 @@ def build_go2_arm_position_gains(cfg: ControlConfig) -> dict[str, np.ndarray]: class Go2ArmBaseEnv(LocomotionBaseEnv): - _cfg: Go2ArmBaseCfg + _cfg: Go2ArmBaseCfg # pyright: ignore[reportIncompatibleVariableOverride] def __init__(self, cfg: Go2ArmBaseCfg, backend: SimBackend, num_envs: int = 1): super().__init__(cfg, backend, num_envs) diff --git a/src/unilab/tasks/locomotion/go2_arm/manip_loco.py b/src/unilab/tasks/locomotion/go2_arm/manip_loco.py index ec82e64fc..6e5a3eda0 100644 --- a/src/unilab/tasks/locomotion/go2_arm/manip_loco.py +++ b/src/unilab/tasks/locomotion/go2_arm/manip_loco.py @@ -17,7 +17,7 @@ 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_arm.base import ( +from unilab.tasks.locomotion.go2_arm.base import ( Go2ArmBaseCfg, Go2ArmBaseEnv, Go2ArmSensor, diff --git a/tests/envs/locomotion/go2_arm/test_base_ik.py b/tests/envs/locomotion/go2_arm/test_base_ik.py index 0fc1abddd..044f9f129 100644 --- a/tests/envs/locomotion/go2_arm/test_base_ik.py +++ b/tests/envs/locomotion/go2_arm/test_base_ik.py @@ -4,7 +4,7 @@ import numpy as np -from unilab.envs.locomotion.go2_arm.base import Go2ArmBaseCfg, Go2ArmBaseEnv +from unilab.tasks.locomotion.go2_arm.base import Go2ArmBaseCfg, Go2ArmBaseEnv class _IkHarness(Go2ArmBaseEnv): From 31bfdef29978608fcfd8089cdc35212034690f71 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 21:50:45 +0800 Subject: [PATCH 049/169] refactor(tasks): move G1 joystick owner [skip ci] (#1144) --- scripts/benchmark/core/task_names.py | 2 +- scripts/benchmark/env/benchmark_env_step.py | 6 +++--- src/unilab/envs/locomotion/__init__.py | 4 ++-- src/unilab/envs/locomotion/g1/__init__.py | 18 +----------------- src/unilab/tasks/__init__.py | 2 +- src/unilab/tasks/locomotion/g1/__init__.py | 19 +++++++++++++++++++ .../{envs => tasks}/locomotion/g1/joystick.py | 8 ++++---- .../{envs => tasks}/locomotion/g1/symmetry.py | 0 tests/base/test_reward_override.py | 2 +- tests/config/test_locomotion_params.py | 4 ++-- .../locomotion/g1/test_g1_owner_contract.py | 8 ++++---- .../locomotion/g1/test_symmetry_contract.py | 2 +- tests/envs/test_env_configs.py | 18 +++++++++--------- tests/tasks/test_package_boundary.py | 2 +- 14 files changed, 49 insertions(+), 46 deletions(-) create mode 100644 src/unilab/tasks/locomotion/g1/__init__.py rename src/unilab/{envs => tasks}/locomotion/g1/joystick.py (98%) rename src/unilab/{envs => tasks}/locomotion/g1/symmetry.py (100%) diff --git a/scripts/benchmark/core/task_names.py b/scripts/benchmark/core/task_names.py index 52f348e31..588dd15f7 100644 --- a/scripts/benchmark/core/task_names.py +++ b/scripts/benchmark/core/task_names.py @@ -2,8 +2,8 @@ from dataclasses import dataclass -from unilab.envs.locomotion.g1.joystick import G1WalkFlatCfg from unilab.envs.manipulation.sharpa_inhand.rotation import SharpaInhandRotationCfg +from unilab.tasks.locomotion.g1.joystick import G1WalkFlatCfg from unilab.tasks.locomotion.go1.joystick import Go1JoystickCfg from unilab.tasks.locomotion.go2.joystick import Go2JoystickCfg diff --git a/scripts/benchmark/env/benchmark_env_step.py b/scripts/benchmark/env/benchmark_env_step.py index 77b137026..c77d370fb 100644 --- a/scripts/benchmark/env/benchmark_env_step.py +++ b/scripts/benchmark/env/benchmark_env_step.py @@ -338,13 +338,13 @@ def _go2w_rough_env_cls() -> type: def _g1_flat_cfg(backend: str, config_overrides: list[str]) -> Any: - from unilab.envs.locomotion.g1.joystick import G1WalkFlatCfg + from unilab.tasks.locomotion.g1.joystick import G1WalkFlatCfg return _ppo_owner_yaml_cfg("g1_walk_flat", backend, G1WalkFlatCfg, config_overrides) def _g1_rough_cfg(backend: str, config_overrides: list[str]) -> Any: - from unilab.envs.locomotion.g1.joystick import G1WalkRoughCfg + from unilab.tasks.locomotion.g1.joystick import G1WalkRoughCfg return _sac_owner_yaml_cfg("g1_walk_rough", backend, G1WalkRoughCfg, config_overrides) @@ -416,7 +416,7 @@ def _ensure_sharpa_benchmark_grasp_cache(cfg: Any, _: str) -> None: def _g1_walk_env_cls() -> type: - from unilab.envs.locomotion.g1.joystick import G1WalkEnv + from unilab.tasks.locomotion.g1.joystick import G1WalkEnv return G1WalkEnv diff --git a/src/unilab/envs/locomotion/__init__.py b/src/unilab/envs/locomotion/__init__.py index 2792c0897..11c9f906a 100644 --- a/src/unilab/envs/locomotion/__init__.py +++ b/src/unilab/envs/locomotion/__init__.py @@ -1,3 +1,3 @@ -"""Locomotion env registry bootstrap contract.""" +"""Generic locomotion environment components.""" -__unilab_registry_modules__ = ("unilab.envs.locomotion.g1",) +__unilab_registry_modules__: tuple[str, ...] = () diff --git a/src/unilab/envs/locomotion/g1/__init__.py b/src/unilab/envs/locomotion/g1/__init__.py index 55806de55..b4901c6b7 100644 --- a/src/unilab/envs/locomotion/g1/__init__.py +++ b/src/unilab/envs/locomotion/g1/__init__.py @@ -1,17 +1 @@ -from .joystick import ( - G1WalkControlConfig, - G1WalkEnv, - G1WalkEnvCfg, - G1WalkFlatCfg, - G1WalkRewardConfig, - G1WalkRoughCfg, -) - -__all__ = [ - "G1WalkControlConfig", - "G1WalkEnv", - "G1WalkEnvCfg", - "G1WalkFlatCfg", - "G1WalkRewardConfig", - "G1WalkRoughCfg", -] +"""Legacy G1 shared base pending motion-tracking migration.""" diff --git a/src/unilab/tasks/__init__.py b/src/unilab/tasks/__init__.py index 14c8be640..4d5959bee 100644 --- a/src/unilab/tasks/__init__.py +++ b/src/unilab/tasks/__init__.py @@ -11,7 +11,7 @@ "unilab.tasks.locomotion.go1", "unilab.tasks.locomotion.go2", "unilab.tasks.locomotion.go2w", - "unilab.envs.locomotion.g1", + "unilab.tasks.locomotion.g1", "unilab.tasks.locomotion.go2_arm", "unilab.tasks.locomotion.a2", "unilab.envs.manipulation.allegro_inhand", diff --git a/src/unilab/tasks/locomotion/g1/__init__.py b/src/unilab/tasks/locomotion/g1/__init__.py new file mode 100644 index 000000000..fbd4f6660 --- /dev/null +++ b/src/unilab/tasks/locomotion/g1/__init__.py @@ -0,0 +1,19 @@ +from .joystick import ( + G1WalkControlConfig, + G1WalkEnv, + G1WalkEnvCfg, + G1WalkFlatCfg, + G1WalkRewardConfig, + G1WalkRoughCfg, +) +from .symmetry import G1SymmetryAugmentation + +__all__ = [ + "G1SymmetryAugmentation", + "G1WalkControlConfig", + "G1WalkEnv", + "G1WalkEnvCfg", + "G1WalkFlatCfg", + "G1WalkRewardConfig", + "G1WalkRoughCfg", +] diff --git a/src/unilab/envs/locomotion/g1/joystick.py b/src/unilab/tasks/locomotion/g1/joystick.py similarity index 98% rename from src/unilab/envs/locomotion/g1/joystick.py rename to src/unilab/tasks/locomotion/g1/joystick.py index 2a74dd86d..902eaafde 100644 --- a/src/unilab/envs/locomotion/g1/joystick.py +++ b/src/unilab/tasks/locomotion/g1/joystick.py @@ -174,7 +174,7 @@ class CurriculumConfig: @dataclass class G1WalkEnvCfg(G1BaseCfg): - scene: SceneCfg = field( + scene: SceneCfg = field( # pyright: ignore[reportIncompatibleVariableOverride] default_factory=lambda: SceneCfg( model_file=str(ASSETS_ROOT_PATH / "robots" / "g1" / "scene_flat.xml") ) @@ -254,7 +254,7 @@ def _compute_reset_obs( class G1WalkEnv(G1BaseEnv): - _cfg: G1WalkEnvCfg + _cfg: G1WalkEnvCfg # pyright: ignore[reportIncompatibleVariableOverride] _reward_cfg: Any def __init__(self, cfg: G1WalkEnvCfg, num_envs=1, backend_type="mujoco"): @@ -485,7 +485,7 @@ def build_symmetry_augmentation(self, *, device: str): actuator_names = self._backend.get_actuator_names() except NotImplementedError: return None - from unilab.envs.locomotion.g1.symmetry import G1SymmetryAugmentation + from unilab.tasks.locomotion.g1.symmetry import G1SymmetryAugmentation return G1SymmetryAugmentation( actuator_names, @@ -656,7 +656,7 @@ class G1WalkRewardConfig(G1RewardConfig): @registry.envcfg("G1WalkFlat") @dataclass class G1WalkFlatCfg(G1WalkEnvCfg): - reward_config: G1WalkRewardConfig | None = None + reward_config: G1WalkRewardConfig | None = None # pyright: ignore[reportIncompatibleVariableOverride] scene: SceneCfg = field( default_factory=lambda: SceneCfg( model_file=str(ASSETS_ROOT_PATH / "robots" / "g1" / "scene_flat.xml") diff --git a/src/unilab/envs/locomotion/g1/symmetry.py b/src/unilab/tasks/locomotion/g1/symmetry.py similarity index 100% rename from src/unilab/envs/locomotion/g1/symmetry.py rename to src/unilab/tasks/locomotion/g1/symmetry.py diff --git a/tests/base/test_reward_override.py b/tests/base/test_reward_override.py index f8d9be2fa..57e5e86f0 100644 --- a/tests/base/test_reward_override.py +++ b/tests/base/test_reward_override.py @@ -38,7 +38,7 @@ def test_reward_override_g1(): """Test G1 reward config override.""" ensure_registries() - from unilab.envs.locomotion.g1.joystick import G1WalkRewardConfig + from unilab.tasks.locomotion.g1.joystick import G1WalkRewardConfig override_config = G1WalkRewardConfig( scales={"tracking_lin_vel": 888.0, "alive": 20.0}, diff --git a/tests/config/test_locomotion_params.py b/tests/config/test_locomotion_params.py index 1cd789188..b7cae7a2e 100644 --- a/tests/config/test_locomotion_params.py +++ b/tests/config/test_locomotion_params.py @@ -357,7 +357,7 @@ def test_offpolicy_g1_rough_terrain_task_overrides(): from hydra import compose, initialize_config_dir from hydra.core.global_hydra import GlobalHydra - from unilab.envs.locomotion.g1.joystick import G1WalkRoughCfg + from unilab.tasks.locomotion.g1.joystick import G1WalkRoughCfg GlobalHydra.instance().clear() with initialize_config_dir(config_dir=str(CONF_DIR / "offpolicy"), version_base="1.3"): @@ -375,7 +375,7 @@ def test_g1_task_owner_yamls_preserve_legacy_and_walk_observation_profiles(): from hydra import compose, initialize_config_dir from hydra.core.global_hydra import GlobalHydra - from unilab.envs.locomotion.g1.joystick import G1WalkEnv + from unilab.tasks.locomotion.g1.joystick import G1WalkEnv def uses_walk_profile(config_group: str, overrides: list[str]) -> bool: GlobalHydra.instance().clear() diff --git a/tests/envs/locomotion/g1/test_g1_owner_contract.py b/tests/envs/locomotion/g1/test_g1_owner_contract.py index 7da05e97c..2b22c9d97 100644 --- a/tests/envs/locomotion/g1/test_g1_owner_contract.py +++ b/tests/envs/locomotion/g1/test_g1_owner_contract.py @@ -129,14 +129,14 @@ def _compose_cfg(config_group: str, overrides: list[str]): def _materialize_env_cfg(cfg: Any): - from unilab.envs.locomotion.g1.joystick import G1WalkFlatCfg, G1WalkRoughCfg + from unilab.tasks.locomotion.g1.joystick import G1WalkFlatCfg, G1WalkRoughCfg env_cfg_cls = G1WalkRoughCfg if cfg.training.task_name == "G1WalkRough" else G1WalkFlatCfg return OmegaConf.merge(OmegaConf.structured(env_cfg_cls()), cfg.env) def _build_probe_env(cfg: Any): - from unilab.envs.locomotion.g1.joystick import G1WalkEnv + from unilab.tasks.locomotion.g1.joystick import G1WalkEnv env = cast(Any, object.__new__(G1WalkEnv)) env._num_envs = 1 @@ -168,7 +168,7 @@ def _compute_probe_obs(cfg: Any) -> dict[str, np.ndarray]: @pytest.mark.parametrize("case", _G1_OWNER_CASES, ids=[case["id"] for case in _G1_OWNER_CASES]) def test_g1_owner_yaml_regression_contract(case: dict[str, Any]): - from unilab.envs.locomotion.g1.joystick import G1WalkEnv + from unilab.tasks.locomotion.g1.joystick import G1WalkEnv cfg = _compose_cfg(case["config_group"], case["overrides"]) full_env_cfg = _materialize_env_cfg(cfg) @@ -222,7 +222,7 @@ def test_g1_owner_yaml_observation_profiles_match_expected_family(case: dict[str def test_g1_observation_profile_selection_prefers_reward_family_over_curriculum_flag(): - from unilab.envs.locomotion.g1.joystick import G1WalkEnv + from unilab.tasks.locomotion.g1.joystick import G1WalkEnv env = cast(Any, object.__new__(G1WalkEnv)) diff --git a/tests/envs/locomotion/g1/test_symmetry_contract.py b/tests/envs/locomotion/g1/test_symmetry_contract.py index fd16ad1ed..3594e94ca 100644 --- a/tests/envs/locomotion/g1/test_symmetry_contract.py +++ b/tests/envs/locomotion/g1/test_symmetry_contract.py @@ -7,7 +7,7 @@ from unilab.base import registry from unilab.base.registry import ensure_registries -from unilab.envs.locomotion.g1.joystick import G1WalkRewardConfig +from unilab.tasks.locomotion.g1.joystick import G1WalkRewardConfig pytest.importorskip("mujoco", reason="mujoco is required for G1 symmetry contract tests") diff --git a/tests/envs/test_env_configs.py b/tests/envs/test_env_configs.py index 746fa7a4f..ae409bae8 100644 --- a/tests/envs/test_env_configs.py +++ b/tests/envs/test_env_configs.py @@ -86,7 +86,7 @@ def blocked_import(name, globals=None, locals=None, fromlist=(), level=0): def test_g1_walk_env_cfg_obs_groups_spec(): """G1WalkEnv must declare obs_groups_spec with actor and critic groups.""" - from unilab.envs.locomotion.g1.joystick import G1WalkEnvCfg + from unilab.tasks.locomotion.g1.joystick import G1WalkEnvCfg cfg = G1WalkEnvCfg() assert not hasattr(cfg, "obs_config"), "obs_config should have been removed" @@ -94,7 +94,7 @@ def test_g1_walk_env_cfg_obs_groups_spec(): def test_g1_walk_flat_cfg_no_obs_config(): """G1WalkFlatCfg should no longer have obs_config after dict obs refactor.""" - from unilab.envs.locomotion.g1.joystick import G1WalkFlatCfg + from unilab.tasks.locomotion.g1.joystick import G1WalkFlatCfg cfg = G1WalkFlatCfg() assert not hasattr(cfg, "obs_config"), ( @@ -103,7 +103,7 @@ def test_g1_walk_flat_cfg_no_obs_config(): def test_g1_walk_flat_cfg_has_domain_rand_for_motrix(): - from unilab.envs.locomotion.g1.joystick import G1WalkFlatCfg + from unilab.tasks.locomotion.g1.joystick import G1WalkFlatCfg cfg = G1WalkFlatCfg() assert hasattr(cfg, "domain_rand") @@ -116,7 +116,7 @@ def test_g1_walk_flat_cfg_has_domain_rand_for_motrix(): def test_g1_walk_flat_cfg_defaults_match_walk_profile(): - from unilab.envs.locomotion.g1.joystick import G1WalkFlatCfg + from unilab.tasks.locomotion.g1.joystick import G1WalkFlatCfg cfg = G1WalkFlatCfg() assert not hasattr(cfg, "obs_profile") @@ -125,7 +125,7 @@ def test_g1_walk_flat_cfg_defaults_match_walk_profile(): def test_g1_walk_tasks_register_to_algorithm_agnostic_env_base(): from unilab.base import registry - from unilab.envs.locomotion.g1.joystick import G1WalkEnv, G1WalkRewardConfig + from unilab.tasks.locomotion.g1.joystick import G1WalkEnv, G1WalkRewardConfig env = cast( Any, @@ -156,7 +156,7 @@ def test_g1_walk_tasks_register_to_algorithm_agnostic_env_base(): def test_g1_walk_flat_observation_construction_is_hardcoded_for_legacy_and_walk_modes(): - from unilab.envs.locomotion.g1.joystick import G1WalkEnv + from unilab.tasks.locomotion.g1.joystick import G1WalkEnv class NoiseCfg: level = 0.0 @@ -219,7 +219,7 @@ def test_g1_walk_env_obs_groups_spec_dims(): + last_actions(29) + command(3) + gait_phase(2) = 98 critic: actor(98) + linvel(3) = 101 """ - from unilab.envs.locomotion.g1.joystick import G1WalkEnv + from unilab.tasks.locomotion.g1.joystick import G1WalkEnv # obs_groups_spec is a @property; access via descriptor protocol spec = G1WalkEnv.obs_groups_spec.fget(None) # type: ignore[union-attr] @@ -229,7 +229,7 @@ def test_g1_walk_env_obs_groups_spec_dims(): def test_g1_walk_env_reward_dispatch_restores_motrix_terms(): - from unilab.envs.locomotion.g1.joystick import G1WalkEnv + from unilab.tasks.locomotion.g1.joystick import G1WalkEnv env = cast(Any, object.__new__(G1WalkEnv)) env._reward_fns = {} @@ -244,7 +244,7 @@ def test_g1_walk_env_reward_dispatch_restores_motrix_terms(): def test_g1_walk_env_feet_phase_reward_is_gated_by_forward_speed(): from unilab.envs.locomotion.common.rewards import RewardContext - from unilab.envs.locomotion.g1.joystick import G1WalkEnv + from unilab.tasks.locomotion.g1.joystick import G1WalkEnv class FakeBackend: def get_sensor_data(self, name: str) -> np.ndarray: diff --git a/tests/tasks/test_package_boundary.py b/tests/tasks/test_package_boundary.py index bba50b0b4..3b208ce19 100644 --- a/tests/tasks/test_package_boundary.py +++ b/tests/tasks/test_package_boundary.py @@ -15,7 +15,7 @@ "unilab.tasks.locomotion.go1", "unilab.tasks.locomotion.go2", "unilab.tasks.locomotion.go2w", - "unilab.envs.locomotion.g1", + "unilab.tasks.locomotion.g1", "unilab.tasks.locomotion.go2_arm", "unilab.tasks.locomotion.a2", "unilab.envs.manipulation.allegro_inhand", From 8b74d4697fa0921138d02b4423d05674995249e5 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 21:55:16 +0800 Subject: [PATCH 050/169] docs(tasks): align G1 task owner paths [skip ci] (#1146) --- docs/sphinx/source/api_reference/envs/locomotion.md | 2 +- .../2-user_guide/5-domain_randomization/2-writing_providers.md | 2 +- .../source/en/4-developer_guide/2-contracts/4-dr_contract.md | 2 +- .../2-user_guide/5-domain_randomization/2-writing_providers.md | 2 +- .../source/zh_CN/4-developer_guide/2-contracts/4-dr_contract.md | 2 +- scripts/benchmark/env/benchmark_env_step.py | 2 +- scripts/benchmark/torch_env/walk_flat.py | 2 +- 7 files changed, 7 insertions(+), 7 deletions(-) diff --git a/docs/sphinx/source/api_reference/envs/locomotion.md b/docs/sphinx/source/api_reference/envs/locomotion.md index 9c2876c85..cf185080b 100644 --- a/docs/sphinx/source/api_reference/envs/locomotion.md +++ b/docs/sphinx/source/api_reference/envs/locomotion.md @@ -7,7 +7,7 @@ :recursive: unilab.envs.locomotion.common - unilab.envs.locomotion.g1 + unilab.tasks.locomotion.g1 unilab.tasks.locomotion.go1 unilab.tasks.locomotion.go2 unilab.tasks.locomotion.go2_arm 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 36f74c627..9f5328e47 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 @@ -28,7 +28,7 @@ The shared types live in `src/unilab/dr/types.py`, and the manager lives in Representative provider implementations are in: - `src/unilab/envs/locomotion/go1/joystick.py` -- `src/unilab/envs/locomotion/g1/joystick.py` +- `src/unilab/tasks/locomotion/g1/joystick.py` - `src/unilab/envs/motion_tracking/g1/tracking.py` - `src/unilab/envs/manipulation/allegro_inhand/rotation.py` - `src/unilab/envs/manipulation/sharpa_inhand/rotation.py` 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 39fd08277..1682f1886 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 @@ -100,6 +100,6 @@ payloads. - DR types: `src/unilab/dr/types.py` - DR manager: `src/unilab/dr/manager.py` - Backend interface: `src/unilab/base/backend/base.py` -- Example providers: `src/unilab/envs/locomotion/g1/joystick.py`, +- Example providers: `src/unilab/tasks/locomotion/g1/joystick.py`, `src/unilab/envs/motion_tracking/g1/tracking.py`, `src/unilab/envs/manipulation/sharpa_inhand/rotation.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 16c242981..bd6a72ef8 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 @@ -27,7 +27,7 @@ 具有代表性的 provider 实现位于： - `src/unilab/envs/locomotion/go1/joystick.py` -- `src/unilab/envs/locomotion/g1/joystick.py` +- `src/unilab/tasks/locomotion/g1/joystick.py` - `src/unilab/envs/motion_tracking/g1/tracking.py` - `src/unilab/envs/manipulation/allegro_inhand/rotation.py` - `src/unilab/envs/manipulation/sharpa_inhand/rotation.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 cdb493364..1f7f30525 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 @@ -90,6 +90,6 @@ actuator 的机制泄漏到共享 payload 里。 - DR 类型：`src/unilab/dr/types.py` - DR manager：`src/unilab/dr/manager.py` - Backend 接口：`src/unilab/base/backend/base.py` -- 示例 provider：`src/unilab/envs/locomotion/g1/joystick.py`、 +- 示例 provider：`src/unilab/tasks/locomotion/g1/joystick.py`、 `src/unilab/envs/motion_tracking/g1/tracking.py`、 `src/unilab/envs/manipulation/sharpa_inhand/rotation.py` diff --git a/scripts/benchmark/env/benchmark_env_step.py b/scripts/benchmark/env/benchmark_env_step.py index c77d370fb..fb9dedd8a 100644 --- a/scripts/benchmark/env/benchmark_env_step.py +++ b/scripts/benchmark/env/benchmark_env_step.py @@ -78,7 +78,7 @@ def _load_helper_module(module_name: str, relative_path: str): def _install_mjwarp_patch() -> bool: """Route ``backend_type == "mjwarp"`` to ``scripts/benchmark/mjwarp`` via factory patch. - Must run before any task env module (e.g. ``unilab.envs.locomotion.g1.joystick``) + Must run before any task env module (e.g. ``unilab.tasks.locomotion.g1.joystick``) is imported, because those modules bind ``create_backend`` at module load time via ``from unilab.base.backend import create_backend``. diff --git a/scripts/benchmark/torch_env/walk_flat.py b/scripts/benchmark/torch_env/walk_flat.py index d93ce437d..12d02b6e6 100644 --- a/scripts/benchmark/torch_env/walk_flat.py +++ b/scripts/benchmark/torch_env/walk_flat.py @@ -3,7 +3,7 @@ Faithful xp-port of the NumPy computation in the collector-timed sections of `uv run train --algo sac --task g1_walk_flat --sim mujoco` (num_envs=2048): -- `G1WalkEnv.update_state` (src/unilab/envs/locomotion/g1/joystick.py): +- `G1WalkEnv.update_state` (src/unilab/tasks/locomotion/g1/joystick.py): termination, `_compute_reward` (9 active terms under the SAC scales incl. per-term logging every 4 steps), `_compute_obs` (noise + concat, walk profile), and the done-triggered curriculum bookkeeping. From 3cac8ad27f6b33e5c75bd62b8eb8fcd5e43b2d32 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 22:01:39 +0800 Subject: [PATCH 051/169] refactor(tasks): move Stewart task owner [skip ci] (#1148) --- docs/sphinx/source/api_reference/envs/manipulation.md | 2 +- src/unilab/envs/manipulation/__init__.py | 1 - src/unilab/tasks/__init__.py | 2 +- src/unilab/tasks/manipulation/__init__.py | 1 + src/unilab/{envs => tasks}/manipulation/stewart/__init__.py | 0 src/unilab/{envs => tasks}/manipulation/stewart/balance.py | 4 ++-- tests/envs/test_stewart.py | 2 +- tests/tasks/test_package_boundary.py | 2 +- 8 files changed, 7 insertions(+), 7 deletions(-) create mode 100644 src/unilab/tasks/manipulation/__init__.py rename src/unilab/{envs => tasks}/manipulation/stewart/__init__.py (100%) rename src/unilab/{envs => tasks}/manipulation/stewart/balance.py (99%) diff --git a/docs/sphinx/source/api_reference/envs/manipulation.md b/docs/sphinx/source/api_reference/envs/manipulation.md index 5e2b2b73b..2667506f4 100644 --- a/docs/sphinx/source/api_reference/envs/manipulation.md +++ b/docs/sphinx/source/api_reference/envs/manipulation.md @@ -8,5 +8,5 @@ unilab.envs.manipulation.allegro_inhand unilab.envs.manipulation.sharpa_inhand - unilab.envs.manipulation.stewart + unilab.tasks.manipulation.stewart ``` diff --git a/src/unilab/envs/manipulation/__init__.py b/src/unilab/envs/manipulation/__init__.py index 16b529f14..02b0f854f 100644 --- a/src/unilab/envs/manipulation/__init__.py +++ b/src/unilab/envs/manipulation/__init__.py @@ -3,5 +3,4 @@ __unilab_registry_modules__ = ( "unilab.envs.manipulation.allegro_inhand", "unilab.envs.manipulation.sharpa_inhand", - "unilab.envs.manipulation.stewart", ) diff --git a/src/unilab/tasks/__init__.py b/src/unilab/tasks/__init__.py index 4d5959bee..2c64cb3ba 100644 --- a/src/unilab/tasks/__init__.py +++ b/src/unilab/tasks/__init__.py @@ -16,7 +16,7 @@ "unilab.tasks.locomotion.a2", "unilab.envs.manipulation.allegro_inhand", "unilab.envs.manipulation.sharpa_inhand", - "unilab.envs.manipulation.stewart", + "unilab.tasks.manipulation.stewart", "unilab.envs.motion_tracking.g1", "unilab.envs.motion_tracking.x2", ) diff --git a/src/unilab/tasks/manipulation/__init__.py b/src/unilab/tasks/manipulation/__init__.py new file mode 100644 index 000000000..530e1ef29 --- /dev/null +++ b/src/unilab/tasks/manipulation/__init__.py @@ -0,0 +1 @@ +"""Manipulation task packages.""" diff --git a/src/unilab/envs/manipulation/stewart/__init__.py b/src/unilab/tasks/manipulation/stewart/__init__.py similarity index 100% rename from src/unilab/envs/manipulation/stewart/__init__.py rename to src/unilab/tasks/manipulation/stewart/__init__.py diff --git a/src/unilab/envs/manipulation/stewart/balance.py b/src/unilab/tasks/manipulation/stewart/balance.py similarity index 99% rename from src/unilab/envs/manipulation/stewart/balance.py rename to src/unilab/tasks/manipulation/stewart/balance.py index ff3f0f38d..0a35290c2 100644 --- a/src/unilab/envs/manipulation/stewart/balance.py +++ b/src/unilab/tasks/manipulation/stewart/balance.py @@ -62,7 +62,7 @@ class StewartRewardConfig: @registry.envcfg("StewartBalance") @dataclass class StewartBalanceCfg(EnvCfg): - scene: SceneCfg = field( + scene: SceneCfg = field( # pyright: ignore[reportIncompatibleVariableOverride] default_factory=lambda: SceneCfg( model_file=str(ASSETS_ROOT_PATH / "robots" / "stewart" / "scene.xml") ) @@ -70,7 +70,7 @@ class StewartBalanceCfg(EnvCfg): # The XML model is stiff; do not raise sim_dt above ~0.005. sim_dt: float = 0.004 ctrl_dt: float = 0.02 - max_episode_seconds: float = 24.0 + max_episode_seconds: float = 24.0 # pyright: ignore[reportIncompatibleVariableOverride] render_spacing: float = 4.5 # Body the backend treats as the kinematic base for its base-pose accessors. diff --git a/tests/envs/test_stewart.py b/tests/envs/test_stewart.py index a32c177c6..a7ed673a9 100644 --- a/tests/envs/test_stewart.py +++ b/tests/envs/test_stewart.py @@ -26,7 +26,7 @@ def _make_env(num_envs: int = 2): def test_stewart_env_uses_backend_contract() -> None: """The task must go through the backend contract, not raw sim internals.""" - source = (_SRC_DIR / "unilab" / "envs" / "manipulation" / "stewart" / "balance.py").read_text( + source = (_SRC_DIR / "unilab" / "tasks" / "manipulation" / "stewart" / "balance.py").read_text( encoding="utf-8" ) assert "import motrixsim" not in source diff --git a/tests/tasks/test_package_boundary.py b/tests/tasks/test_package_boundary.py index 3b208ce19..fcdf8240b 100644 --- a/tests/tasks/test_package_boundary.py +++ b/tests/tasks/test_package_boundary.py @@ -20,7 +20,7 @@ "unilab.tasks.locomotion.a2", "unilab.envs.manipulation.allegro_inhand", "unilab.envs.manipulation.sharpa_inhand", - "unilab.envs.manipulation.stewart", + "unilab.tasks.manipulation.stewart", "unilab.envs.motion_tracking.g1", "unilab.envs.motion_tracking.x2", ) From 069aa41ee1b7578311130002ea2c1d2429787180 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 22:08:33 +0800 Subject: [PATCH 052/169] refactor(tasks): move Allegro in-hand owner [skip ci] (#1150) --- .../source/api_reference/envs/manipulation.md | 2 +- .../2-writing_providers.md | 2 +- .../1-sim_to_real/4-allegro_inhand.md | 4 ++-- .../3-extending/1-new_task.md | 2 +- .../2-writing_providers.md | 2 +- .../1-sim_to_real/4-allegro_inhand.md | 4 ++-- .../3-extending/1-new_task.md | 2 +- src/unilab/envs/manipulation/__init__.py | 5 +---- src/unilab/tasks/__init__.py | 2 +- .../manipulation/allegro_inhand/__init__.py | 0 .../manipulation/allegro_inhand/base.py | 0 .../manipulation/allegro_inhand/grasp_gen.py | 4 ++-- .../manipulation/allegro_inhand/rotation.py | 6 +++--- tests/envs/test_env_configs.py | 18 +++++++++--------- tests/tasks/test_package_boundary.py | 2 +- 15 files changed, 26 insertions(+), 29 deletions(-) rename src/unilab/{envs => tasks}/manipulation/allegro_inhand/__init__.py (100%) rename src/unilab/{envs => tasks}/manipulation/allegro_inhand/base.py (100%) rename src/unilab/{envs => tasks}/manipulation/allegro_inhand/grasp_gen.py (98%) rename src/unilab/{envs => tasks}/manipulation/allegro_inhand/rotation.py (98%) diff --git a/docs/sphinx/source/api_reference/envs/manipulation.md b/docs/sphinx/source/api_reference/envs/manipulation.md index 2667506f4..7734224c0 100644 --- a/docs/sphinx/source/api_reference/envs/manipulation.md +++ b/docs/sphinx/source/api_reference/envs/manipulation.md @@ -6,7 +6,7 @@ :template: autosummary/module.rst :recursive: - unilab.envs.manipulation.allegro_inhand + unilab.tasks.manipulation.allegro_inhand unilab.envs.manipulation.sharpa_inhand unilab.tasks.manipulation.stewart ``` 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 9f5328e47..1678df319 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 @@ -30,7 +30,7 @@ Representative provider implementations are in: - `src/unilab/envs/locomotion/go1/joystick.py` - `src/unilab/tasks/locomotion/g1/joystick.py` - `src/unilab/envs/motion_tracking/g1/tracking.py` -- `src/unilab/envs/manipulation/allegro_inhand/rotation.py` +- `src/unilab/tasks/manipulation/allegro_inhand/rotation.py` - `src/unilab/envs/manipulation/sharpa_inhand/rotation.py` Developer contract details are in diff --git a/docs/sphinx/source/en/3-deployment/1-sim_to_real/4-allegro_inhand.md b/docs/sphinx/source/en/3-deployment/1-sim_to_real/4-allegro_inhand.md index 604b3b72e..a42be5ff1 100644 --- a/docs/sphinx/source/en/3-deployment/1-sim_to_real/4-allegro_inhand.md +++ b/docs/sphinx/source/en/3-deployment/1-sim_to_real/4-allegro_inhand.md @@ -60,7 +60,7 @@ that samples plausible initial hand configurations. The hardware-side equivalent is the operator placing the cube in the hand — verify your distribution of starting configurations matches the trained env's grasp generator output (see -`unilab.envs.manipulation.allegro_inhand.grasp_gen`). +`unilab.tasks.manipulation.allegro_inhand.grasp_gen`). If your real-world starting grip differs systematically, **add those poses to the grasp generator**, retrain, and try again. @@ -68,7 +68,7 @@ to the grasp generator**, retrain, and try again. ## Action interface The manipulation envs map policy actions to joint position targets through the -task control config (`src/unilab/envs/manipulation/allegro_inhand/base.py` and +task control config (`src/unilab/tasks/manipulation/allegro_inhand/base.py` and `src/unilab/envs/manipulation/sharpa_inhand/base.py`). The deploy controller must use the same joint order, action scale, and limit policy. diff --git a/docs/sphinx/source/en/4-developer_guide/3-extending/1-new_task.md b/docs/sphinx/source/en/4-developer_guide/3-extending/1-new_task.md index 37924785a..6ee3871a7 100644 --- a/docs/sphinx/source/en/4-developer_guide/3-extending/1-new_task.md +++ b/docs/sphinx/source/en/4-developer_guide/3-extending/1-new_task.md @@ -39,4 +39,4 @@ Start from the contracts: {doc}`../2-contracts/1-env_contract`, - Env state contract: `src/unilab/base/np_env.py` - Scene config: `src/unilab/base/scene.py` - Existing task examples: `src/unilab/envs/locomotion/go2/joystick.py`, - `src/unilab/envs/manipulation/allegro_inhand/rotation.py` + `src/unilab/tasks/manipulation/allegro_inhand/rotation.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 bd6a72ef8..e045f4fd6 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 @@ -29,7 +29,7 @@ - `src/unilab/envs/locomotion/go1/joystick.py` - `src/unilab/tasks/locomotion/g1/joystick.py` - `src/unilab/envs/motion_tracking/g1/tracking.py` -- `src/unilab/envs/manipulation/allegro_inhand/rotation.py` +- `src/unilab/tasks/manipulation/allegro_inhand/rotation.py` - `src/unilab/envs/manipulation/sharpa_inhand/rotation.py` 开发者 contract 详情见 diff --git a/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/4-allegro_inhand.md b/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/4-allegro_inhand.md index d5df85edc..4edc57715 100644 --- a/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/4-allegro_inhand.md +++ b/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/4-allegro_inhand.md @@ -55,7 +55,7 @@ owner 与部署运行时在观测时序上达成一致。见 `4-allegro_inhand` 与 `sharpa_inhand` 两个环境都自带一个**抓取生成器**，用于采样 合理的初始手部构型。硬件侧的等价物是操作员把方块放到手里 —— 请核实你的起始构型分布 与训练环境的抓取生成器输出相匹配（参见 -`unilab.envs.manipulation.allegro_inhand.grasp_gen`）。 +`unilab.tasks.manipulation.allegro_inhand.grasp_gen`）。 如果你真实世界的起始握持存在系统性差异，**把这些位姿加入抓取生成器**，重新训练， 然后再试。 @@ -63,7 +63,7 @@ owner 与部署运行时在观测时序上达成一致。见 ## 动作接口 操作类环境通过任务控制配置把策略动作映射为关节位置目标 -（`src/unilab/envs/manipulation/allegro_inhand/base.py` 和 +（`src/unilab/tasks/manipulation/allegro_inhand/base.py` 和 `src/unilab/envs/manipulation/sharpa_inhand/base.py`）。部署控制器必须使用相同的 关节顺序、动作缩放与限位策略。 diff --git a/docs/sphinx/source/zh_CN/4-developer_guide/3-extending/1-new_task.md b/docs/sphinx/source/zh_CN/4-developer_guide/3-extending/1-new_task.md index a60aeea32..49e7c39e7 100644 --- a/docs/sphinx/source/zh_CN/4-developer_guide/3-extending/1-new_task.md +++ b/docs/sphinx/source/zh_CN/4-developer_guide/3-extending/1-new_task.md @@ -39,4 +39,4 @@ - Env 状态契约：`src/unilab/base/np_env.py` - 场景配置：`src/unilab/base/scene.py` - 现有任务示例：`src/unilab/envs/locomotion/go2/joystick.py`、 - `src/unilab/envs/manipulation/allegro_inhand/rotation.py` + `src/unilab/tasks/manipulation/allegro_inhand/rotation.py` diff --git a/src/unilab/envs/manipulation/__init__.py b/src/unilab/envs/manipulation/__init__.py index 02b0f854f..a0fc7c4f3 100644 --- a/src/unilab/envs/manipulation/__init__.py +++ b/src/unilab/envs/manipulation/__init__.py @@ -1,6 +1,3 @@ """Manipulation env registry bootstrap contract.""" -__unilab_registry_modules__ = ( - "unilab.envs.manipulation.allegro_inhand", - "unilab.envs.manipulation.sharpa_inhand", -) +__unilab_registry_modules__ = ("unilab.envs.manipulation.sharpa_inhand",) diff --git a/src/unilab/tasks/__init__.py b/src/unilab/tasks/__init__.py index 2c64cb3ba..d060155b3 100644 --- a/src/unilab/tasks/__init__.py +++ b/src/unilab/tasks/__init__.py @@ -14,7 +14,7 @@ "unilab.tasks.locomotion.g1", "unilab.tasks.locomotion.go2_arm", "unilab.tasks.locomotion.a2", - "unilab.envs.manipulation.allegro_inhand", + "unilab.tasks.manipulation.allegro_inhand", "unilab.envs.manipulation.sharpa_inhand", "unilab.tasks.manipulation.stewart", "unilab.envs.motion_tracking.g1", diff --git a/src/unilab/envs/manipulation/allegro_inhand/__init__.py b/src/unilab/tasks/manipulation/allegro_inhand/__init__.py similarity index 100% rename from src/unilab/envs/manipulation/allegro_inhand/__init__.py rename to src/unilab/tasks/manipulation/allegro_inhand/__init__.py diff --git a/src/unilab/envs/manipulation/allegro_inhand/base.py b/src/unilab/tasks/manipulation/allegro_inhand/base.py similarity index 100% rename from src/unilab/envs/manipulation/allegro_inhand/base.py rename to src/unilab/tasks/manipulation/allegro_inhand/base.py diff --git a/src/unilab/envs/manipulation/allegro_inhand/grasp_gen.py b/src/unilab/tasks/manipulation/allegro_inhand/grasp_gen.py similarity index 98% rename from src/unilab/envs/manipulation/allegro_inhand/grasp_gen.py rename to src/unilab/tasks/manipulation/allegro_inhand/grasp_gen.py index 238d7df58..2e9c30ec6 100644 --- a/src/unilab/envs/manipulation/allegro_inhand/grasp_gen.py +++ b/src/unilab/tasks/manipulation/allegro_inhand/grasp_gen.py @@ -22,7 +22,7 @@ class AllegroRotationGraspCfg(AllegroRotationPPOCfg): # conf/ppo/task/allegro_inhand_grasp/mujoco.yaml and CLI env.*) # are applied at env construction and take precedence. max_episode_seconds: float = 2.0 - reward_config: RewardConfigPPO = field( + reward_config: RewardConfigPPO = field( # pyright: ignore[reportIncompatibleVariableOverride] default_factory=lambda: RewardConfigPPO( scales={ "rotate": 0.0, @@ -47,7 +47,7 @@ class AllegroRotationGraspCfg(AllegroRotationPPOCfg): @registry.env("AllegroInhandRotationGrasp", sim_backend="mujoco") @registry.env("AllegroInhandRotationGrasp", sim_backend="motrix") class AllegroRotationGrasp(AllegroRotationPPO): - _cfg: AllegroRotationGraspCfg + _cfg: AllegroRotationGraspCfg # pyright: ignore[reportIncompatibleVariableOverride] _CONTACT_SENSORS = ("ff_contact", "mf_contact", "rf_contact", "th_contact") def __init__( diff --git a/src/unilab/envs/manipulation/allegro_inhand/rotation.py b/src/unilab/tasks/manipulation/allegro_inhand/rotation.py similarity index 98% rename from src/unilab/envs/manipulation/allegro_inhand/rotation.py rename to src/unilab/tasks/manipulation/allegro_inhand/rotation.py index 87faf5952..df4ade827 100644 --- a/src/unilab/envs/manipulation/allegro_inhand/rotation.py +++ b/src/unilab/tasks/manipulation/allegro_inhand/rotation.py @@ -145,12 +145,12 @@ class DomainRandConfig: @registry.envcfg("AllegroInhandRotation") @dataclass class AllegroRotationPPOCfg(AllegroBaseCfg): - scene: SceneCfg = field( + scene: SceneCfg = field( # pyright: ignore[reportIncompatibleVariableOverride] default_factory=lambda: SceneCfg( model_file=str(ASSETS_ROOT_PATH / "robots" / "allegro_hand" / "scene.xml") ) ) - max_episode_seconds: float = 20.0 + max_episode_seconds: float = 20.0 # pyright: ignore[reportIncompatibleVariableOverride] reward_config: RewardConfigPPO | None = None domain_rand: DomainRandConfig = field(default_factory=DomainRandConfig) rotation_axis: tuple[float, float, float] = (0.0, 0.0, 1.0) @@ -260,7 +260,7 @@ def build_reset_observation( @registry.env("AllegroInhandRotation", sim_backend="mujoco") @registry.env("AllegroInhandRotation", sim_backend="motrix") class AllegroRotationPPO(AllegroBaseEnv): - _cfg: AllegroRotationPPOCfg + _cfg: AllegroRotationPPOCfg # pyright: ignore[reportIncompatibleVariableOverride] _reward_cfg: RewardConfigPPO _NUM_OBS_PER_STEP = 35 diff --git a/tests/envs/test_env_configs.py b/tests/envs/test_env_configs.py index ae409bae8..147ea1754 100644 --- a/tests/envs/test_env_configs.py +++ b/tests/envs/test_env_configs.py @@ -52,7 +52,7 @@ def blocked_import(name, globals=None, locals=None, fromlist=(), level=0): from unilab.base import registry from unilab.base.backend import create_backend - from unilab.envs.manipulation.allegro_inhand.rotation import AllegroRotationCfg + from unilab.tasks.manipulation.allegro_inhand.rotation import AllegroRotationCfg from unilab.envs.motion_tracking.g1.tracking import ( G1MotionTrackingCfg, G1MotionTrackingDeployEnvCfg, @@ -396,7 +396,7 @@ def test_g1_box_tracking_scene_uses_sphere_hand_and_box_tracking_mesh(): def test_allegro_rotation_obs_groups_spec_dims(): """Allegro rotation obs_groups_spec should expose single actor obs group.""" - from unilab.envs.manipulation.allegro_inhand.rotation import AllegroRotationPPO + from unilab.tasks.manipulation.allegro_inhand.rotation import AllegroRotationPPO env = cast(Any, object.__new__(AllegroRotationPPO)) spec = env.obs_groups_spec @@ -406,7 +406,7 @@ def test_allegro_rotation_obs_groups_spec_dims(): def test_allegro_grasp_obs_groups_spec_dims(): """Allegro grasp task inherits the same obs group layout as rotation.""" - from unilab.envs.manipulation.allegro_inhand.grasp_gen import AllegroRotationGrasp + from unilab.tasks.manipulation.allegro_inhand.grasp_gen import AllegroRotationGrasp env = cast(Any, object.__new__(AllegroRotationGrasp)) spec = env.obs_groups_spec @@ -417,7 +417,7 @@ def test_allegro_grasp_obs_groups_spec_dims(): def test_allegro_missing_grasp_cache_logs_local_generation_notice( tmp_path: Path, caplog: pytest.LogCaptureFixture ) -> None: - from unilab.envs.manipulation.allegro_inhand.rotation import ( + from unilab.tasks.manipulation.allegro_inhand.rotation import ( AllegroRotationPPOCfg, _materialize_grasp_cache, ) @@ -429,7 +429,7 @@ def test_allegro_missing_grasp_cache_logs_local_generation_notice( ) with caplog.at_level( - logging.WARNING, logger="unilab.envs.manipulation.allegro_inhand.rotation" + logging.WARNING, logger="unilab.tasks.manipulation.allegro_inhand.rotation" ): assert _materialize_grasp_cache(cfg) is None notice = caplog.text @@ -443,7 +443,7 @@ def test_allegro_missing_grasp_cache_logs_local_generation_notice( def test_allegro_grasp_generation_skips_cache_materialization( monkeypatch: pytest.MonkeyPatch, ) -> None: - from unilab.envs.manipulation.allegro_inhand import rotation + from unilab.tasks.manipulation.allegro_inhand import rotation def fail_io(*_args: Any, **_kwargs: Any) -> None: raise AssertionError("grasp generation must not resolve or load a rotation cache") @@ -459,7 +459,7 @@ def test_allegro_grasp_target_raises_run_complete_without_resaving_on_close( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: from unilab.base.run_control import RunComplete - from unilab.envs.manipulation.allegro_inhand import grasp_gen + from unilab.tasks.manipulation.allegro_inhand import grasp_gen cache_path = tmp_path / "allegro.npy" env = cast(Any, object.__new__(grasp_gen.AllegroRotationGrasp)) @@ -528,7 +528,7 @@ def save_once(path: str | Path, values: np.ndarray) -> None: def test_allegro_grasp_save_failure_does_not_signal_completion( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: - from unilab.envs.manipulation.allegro_inhand import grasp_gen + from unilab.tasks.manipulation.allegro_inhand import grasp_gen env = cast(Any, object.__new__(grasp_gen.AllegroRotationGrasp)) env._cfg = grasp_gen.AllegroRotationGraspCfg( @@ -555,7 +555,7 @@ def test_allegro_grasp_save_failure_does_not_signal_completion( def test_allegro_reset_samples_materialized_cache_without_file_io( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: - from unilab.envs.manipulation.allegro_inhand import rotation + from unilab.tasks.manipulation.allegro_inhand import rotation cache_path = tmp_path / "allegro.npy" cache = np.zeros((4, 23), dtype=np.float64) diff --git a/tests/tasks/test_package_boundary.py b/tests/tasks/test_package_boundary.py index fcdf8240b..c064c3dfd 100644 --- a/tests/tasks/test_package_boundary.py +++ b/tests/tasks/test_package_boundary.py @@ -18,7 +18,7 @@ "unilab.tasks.locomotion.g1", "unilab.tasks.locomotion.go2_arm", "unilab.tasks.locomotion.a2", - "unilab.envs.manipulation.allegro_inhand", + "unilab.tasks.manipulation.allegro_inhand", "unilab.envs.manipulation.sharpa_inhand", "unilab.tasks.manipulation.stewart", "unilab.envs.motion_tracking.g1", From ed2300eafcd056e37cce8f84d34672cdd69cceb8 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 22:15:41 +0800 Subject: [PATCH 053/169] refactor(tasks): move Sharpa in-hand runtime [skip ci] (#1152) --- scripts/benchmark/core/task_names.py | 2 +- scripts/benchmark/env/benchmark_env_step.py | 6 +++--- .../env/benchmark_sharpa_init_dr_construct.py | 4 ++-- src/unilab/envs/manipulation/__init__.py | 4 ++-- src/unilab/tasks/__init__.py | 2 +- .../manipulation/sharpa_inhand/__init__.py | 0 .../manipulation/sharpa_inhand/base.py | 4 ++-- .../manipulation/sharpa_inhand/grasp_gen.py | 8 ++++---- .../manipulation/sharpa_inhand/rotation.py | 4 ++-- tests/envs/test_env_configs.py | 16 ++++++++-------- tests/envs/test_sharpa.py | 8 ++++---- tests/scripts/test_train_script_configs.py | 2 +- tests/tasks/test_package_boundary.py | 2 +- tests/test_sharpa.py | 6 +++--- 14 files changed, 34 insertions(+), 34 deletions(-) rename src/unilab/{envs => tasks}/manipulation/sharpa_inhand/__init__.py (100%) rename src/unilab/{envs => tasks}/manipulation/sharpa_inhand/base.py (99%) rename src/unilab/{envs => tasks}/manipulation/sharpa_inhand/grasp_gen.py (97%) rename src/unilab/{envs => tasks}/manipulation/sharpa_inhand/rotation.py (99%) diff --git a/scripts/benchmark/core/task_names.py b/scripts/benchmark/core/task_names.py index 588dd15f7..409ca6137 100644 --- a/scripts/benchmark/core/task_names.py +++ b/scripts/benchmark/core/task_names.py @@ -2,10 +2,10 @@ from dataclasses import dataclass -from unilab.envs.manipulation.sharpa_inhand.rotation import SharpaInhandRotationCfg from unilab.tasks.locomotion.g1.joystick import G1WalkFlatCfg from unilab.tasks.locomotion.go1.joystick import Go1JoystickCfg from unilab.tasks.locomotion.go2.joystick import Go2JoystickCfg +from unilab.tasks.manipulation.sharpa_inhand.rotation import SharpaInhandRotationCfg @dataclass(frozen=True) diff --git a/scripts/benchmark/env/benchmark_env_step.py b/scripts/benchmark/env/benchmark_env_step.py index fb9dedd8a..e325d8026 100644 --- a/scripts/benchmark/env/benchmark_env_step.py +++ b/scripts/benchmark/env/benchmark_env_step.py @@ -365,7 +365,7 @@ def _sharpa_inhand_cfg(backend: str, config_overrides: list[str]) -> Any: from hydra.core.global_hydra import GlobalHydra from unilab.base.registry import apply_cfg_overrides - from unilab.envs.manipulation.sharpa_inhand.rotation import SharpaInhandRotationCfg + from unilab.tasks.manipulation.sharpa_inhand.rotation import SharpaInhandRotationCfg from unilab.training import BackendAdapter yaml_backend = _hydra_yaml_backend(backend) @@ -396,7 +396,7 @@ def _sharpa_inhand_cfg(backend: str, config_overrides: list[str]) -> Any: def _ensure_sharpa_benchmark_grasp_cache(cfg: Any, _: str) -> None: - from unilab.envs.manipulation.sharpa_inhand.base import ( + from unilab.tasks.manipulation.sharpa_inhand.base import ( SOURCE_DEFAULT_HAND_JOINT_POS_DEG, resolve_grasp_cache_file, ) @@ -428,7 +428,7 @@ def _g1_motion_tracking_env_cls() -> type: def _sharpa_inhand_env_cls() -> type: - from unilab.envs.manipulation.sharpa_inhand.rotation import SharpaInhandRotationEnv + from unilab.tasks.manipulation.sharpa_inhand.rotation import SharpaInhandRotationEnv return SharpaInhandRotationEnv diff --git a/scripts/benchmark/env/benchmark_sharpa_init_dr_construct.py b/scripts/benchmark/env/benchmark_sharpa_init_dr_construct.py index 31017047f..4f98687b3 100644 --- a/scripts/benchmark/env/benchmark_sharpa_init_dr_construct.py +++ b/scripts/benchmark/env/benchmark_sharpa_init_dr_construct.py @@ -140,7 +140,7 @@ def _compose_cfg(task: str, *, lower: float, upper: float, variant_count: int): @contextmanager def _init_dr_mode(enabled: bool) -> Iterator[None]: - from unilab.envs.manipulation.sharpa_inhand.rotation import SharpaInhandRotationDRProvider + from unilab.tasks.manipulation.sharpa_inhand.rotation import SharpaInhandRotationDRProvider original = SharpaInhandRotationDRProvider.build_init_randomization_plan if enabled: @@ -162,7 +162,7 @@ def disabled_build_init_randomization_plan(self: Any, env: Any) -> None: @contextmanager def _synthetic_grasp_cache_mode(enabled: bool) -> Iterator[None]: - from unilab.envs.manipulation.sharpa_inhand.rotation import SharpaInhandRotationDRProvider + from unilab.tasks.manipulation.sharpa_inhand.rotation import SharpaInhandRotationDRProvider original = SharpaInhandRotationDRProvider._load_grasp_cache if not enabled: diff --git a/src/unilab/envs/manipulation/__init__.py b/src/unilab/envs/manipulation/__init__.py index a0fc7c4f3..15c01b96c 100644 --- a/src/unilab/envs/manipulation/__init__.py +++ b/src/unilab/envs/manipulation/__init__.py @@ -1,3 +1,3 @@ -"""Manipulation env registry bootstrap contract.""" +"""Generic manipulation environment components.""" -__unilab_registry_modules__ = ("unilab.envs.manipulation.sharpa_inhand",) +__unilab_registry_modules__: tuple[str, ...] = () diff --git a/src/unilab/tasks/__init__.py b/src/unilab/tasks/__init__.py index d060155b3..daf7d36ab 100644 --- a/src/unilab/tasks/__init__.py +++ b/src/unilab/tasks/__init__.py @@ -15,7 +15,7 @@ "unilab.tasks.locomotion.go2_arm", "unilab.tasks.locomotion.a2", "unilab.tasks.manipulation.allegro_inhand", - "unilab.envs.manipulation.sharpa_inhand", + "unilab.tasks.manipulation.sharpa_inhand", "unilab.tasks.manipulation.stewart", "unilab.envs.motion_tracking.g1", "unilab.envs.motion_tracking.x2", diff --git a/src/unilab/envs/manipulation/sharpa_inhand/__init__.py b/src/unilab/tasks/manipulation/sharpa_inhand/__init__.py similarity index 100% rename from src/unilab/envs/manipulation/sharpa_inhand/__init__.py rename to src/unilab/tasks/manipulation/sharpa_inhand/__init__.py diff --git a/src/unilab/envs/manipulation/sharpa_inhand/base.py b/src/unilab/tasks/manipulation/sharpa_inhand/base.py similarity index 99% rename from src/unilab/envs/manipulation/sharpa_inhand/base.py rename to src/unilab/tasks/manipulation/sharpa_inhand/base.py index e5e323b31..677ae7bf6 100644 --- a/src/unilab/envs/manipulation/sharpa_inhand/base.py +++ b/src/unilab/tasks/manipulation/sharpa_inhand/base.py @@ -152,12 +152,12 @@ class SharpaDomainRandConfig: @dataclass class SharpaInhandBaseCfg(EnvCfg): - scene: SceneCfg = field( + scene: SceneCfg = field( # pyright: ignore[reportIncompatibleVariableOverride] default_factory=lambda: SceneCfg( model_file=str(ASSETS_ROOT_PATH / "robots" / "sharpa_wave" / "scene.xml") ) ) - max_episode_seconds: float = 20.0 + max_episode_seconds: float = 20.0 # pyright: ignore[reportIncompatibleVariableOverride] sim_dt: float = 1.0 / 240.0 ctrl_dt: float = 12.0 / 240.0 diff --git a/src/unilab/envs/manipulation/sharpa_inhand/grasp_gen.py b/src/unilab/tasks/manipulation/sharpa_inhand/grasp_gen.py similarity index 97% rename from src/unilab/envs/manipulation/sharpa_inhand/grasp_gen.py rename to src/unilab/tasks/manipulation/sharpa_inhand/grasp_gen.py index 8e85e1253..4a6583451 100644 --- a/src/unilab/envs/manipulation/sharpa_inhand/grasp_gen.py +++ b/src/unilab/tasks/manipulation/sharpa_inhand/grasp_gen.py @@ -12,12 +12,12 @@ from unilab.base.run_control import RunComplete from unilab.dr import ResetPlan from unilab.dr.dr_utils import build_common_reset_randomization -from unilab.envs.manipulation.sharpa_inhand.base import ( +from unilab.tasks.manipulation.sharpa_inhand.base import ( SOURCE_DEFAULT_HAND_JOINT_POS_DEG, SharpaDomainRandConfig, resolve_grasp_cache_file, ) -from unilab.envs.manipulation.sharpa_inhand.rotation import ( +from unilab.tasks.manipulation.sharpa_inhand.rotation import ( RewardConfig, SharpaInhandRotationCfg, SharpaInhandRotationDRProvider, @@ -57,7 +57,7 @@ class SharpaInhandRotationGraspCfg(SharpaInhandRotationCfg): grasp_cache_path: str = "" domain_rand: SharpaDomainRandConfig = field(default_factory=_default_sharpa_grasp_domain_rand) - reward_config: RewardConfig = field( + reward_config: RewardConfig = field( # pyright: ignore[reportIncompatibleVariableOverride] default_factory=lambda: RewardConfig( scales={ "rotate": 0.0, @@ -141,7 +141,7 @@ def build_reset_plan(self, env: Any, env_ids: np.ndarray) -> ResetPlan: @registry.env("SharpaInhandRotationGrasp", sim_backend="mujoco") @registry.env("SharpaInhandRotationGrasp", sim_backend="motrix") class SharpaInhandRotationGraspEnv(SharpaInhandRotationEnv): - _cfg: SharpaInhandRotationGraspCfg + _cfg: SharpaInhandRotationGraspCfg # pyright: ignore[reportIncompatibleVariableOverride] _MATERIALIZE_ROTATION_GRASP_CACHE = False def __init__( diff --git a/src/unilab/envs/manipulation/sharpa_inhand/rotation.py b/src/unilab/tasks/manipulation/sharpa_inhand/rotation.py similarity index 99% rename from src/unilab/envs/manipulation/sharpa_inhand/rotation.py rename to src/unilab/tasks/manipulation/sharpa_inhand/rotation.py index 252330e8d..8a6c5ca37 100644 --- a/src/unilab/envs/manipulation/sharpa_inhand/rotation.py +++ b/src/unilab/tasks/manipulation/sharpa_inhand/rotation.py @@ -30,7 +30,7 @@ ResetRandomizationPayload, ) from unilab.dtype_config import get_global_dtype -from unilab.envs.manipulation.sharpa_inhand.base import ( +from unilab.tasks.manipulation.sharpa_inhand.base import ( SharpaInhandBaseCfg, SharpaInhandBaseEnv, repeat_obs_history, @@ -444,7 +444,7 @@ def build_interval_randomization_plan( @registry.env("SharpaInhandRotation", sim_backend="mujoco") @registry.env("SharpaInhandRotation", sim_backend="motrix") class SharpaInhandRotationEnv(SharpaInhandBaseEnv): - _cfg: SharpaInhandRotationCfg + _cfg: SharpaInhandRotationCfg # pyright: ignore[reportIncompatibleVariableOverride] _reward_cfg: RewardConfig _MATERIALIZE_ROTATION_GRASP_CACHE = True _OBS_MODE_ALIASES: dict[str, str] = { diff --git a/tests/envs/test_env_configs.py b/tests/envs/test_env_configs.py index 147ea1754..833292ee3 100644 --- a/tests/envs/test_env_configs.py +++ b/tests/envs/test_env_configs.py @@ -1480,8 +1480,8 @@ def _patch_sharpa_rotation_constructor( monkeypatch: pytest.MonkeyPatch, initialized_providers: list[Any], ) -> Any: - from unilab.envs.manipulation.sharpa_inhand import rotation as sharpa_rotation_module - from unilab.envs.manipulation.sharpa_inhand.base import SharpaInhandBaseEnv + from unilab.tasks.manipulation.sharpa_inhand import rotation as sharpa_rotation_module + from unilab.tasks.manipulation.sharpa_inhand.base import SharpaInhandBaseEnv def fake_base_init(self, cfg, backend, num_envs): self._cfg = cfg @@ -1529,7 +1529,7 @@ def unsupported_backend_metadata(*args, **kwargs): def test_sharpa_rotation_explicit_default_provider_materializes_cache(monkeypatch): - from unilab.envs.manipulation.sharpa_inhand.rotation import ( + from unilab.tasks.manipulation.sharpa_inhand.rotation import ( RewardConfig, SharpaInhandRotationCfg, SharpaInhandRotationDRProvider, @@ -1564,7 +1564,7 @@ def materialize(path: str, scale_values: np.ndarray) -> tuple[np.ndarray, ...]: def test_sharpa_grasp_env_initializes_dr_once_with_grasp_provider(monkeypatch): - from unilab.envs.manipulation.sharpa_inhand.grasp_gen import ( + from unilab.tasks.manipulation.sharpa_inhand.grasp_gen import ( SharpaInhandRotationGraspCfg, SharpaInhandRotationGraspEnv, ) @@ -1613,7 +1613,7 @@ def test_sharpa_grasp_target_saves_cache_then_raises_run_complete( caplog: pytest.LogCaptureFixture, ) -> None: from unilab.base.run_control import RunComplete - from unilab.envs.manipulation.sharpa_inhand.grasp_gen import ( + from unilab.tasks.manipulation.sharpa_inhand.grasp_gen import ( SharpaInhandRotationGraspCfg, SharpaInhandRotationGraspEnv, ) @@ -1651,13 +1651,13 @@ def save_once(path: str | Path, values: np.ndarray) -> None: real_save(path, values) monkeypatch.setattr( - "unilab.envs.manipulation.sharpa_inhand.grasp_gen.np.save", + "unilab.tasks.manipulation.sharpa_inhand.grasp_gen.np.save", save_once, ) with caplog.at_level( logging.INFO, - logger="unilab.envs.manipulation.sharpa_inhand.grasp_gen", + logger="unilab.tasks.manipulation.sharpa_inhand.grasp_gen", ): with pytest.raises(RunComplete) as caught: env._collect_successful_grasps(np.asarray([0], dtype=np.int32)) @@ -1692,7 +1692,7 @@ def save_once(path: str | Path, values: np.ndarray) -> None: def test_sharpa_run_complete_reports_effective_collection_target() -> None: from unilab.base.run_control import RunComplete - from unilab.envs.manipulation.sharpa_inhand.grasp_gen import ( + from unilab.tasks.manipulation.sharpa_inhand.grasp_gen import ( SharpaInhandRotationGraspCfg, SharpaInhandRotationGraspEnv, ) diff --git a/tests/envs/test_sharpa.py b/tests/envs/test_sharpa.py index 2b46eae0c..9639a3f25 100644 --- a/tests/envs/test_sharpa.py +++ b/tests/envs/test_sharpa.py @@ -12,12 +12,12 @@ from omegaconf import OmegaConf from unilab.base.registry import ensure_registries -from unilab.envs.manipulation.sharpa_inhand.base import ( +from unilab.tasks.manipulation.sharpa_inhand.base import ( SOURCE_DEFAULT_HAND_JOINT_POS_DEG, SharpaInhandBaseEnv, resolve_grasp_cache_file, ) -from unilab.envs.manipulation.sharpa_inhand.rotation import ( +from unilab.tasks.manipulation.sharpa_inhand.rotation import ( SharpaInhandRotationDRProvider, _materialize_grasp_caches, ) @@ -29,7 +29,7 @@ def test_sharpa_env_uses_backend_contract_for_mujoco_metadata() -> None: """Sharpa env code should not read MuJoCo model internals directly.""" source = "\n".join( - (_SRC_DIR / "unilab" / "envs" / "manipulation" / "sharpa_inhand" / path).read_text( + (_SRC_DIR / "unilab" / "tasks" / "manipulation" / "sharpa_inhand" / path).read_text( encoding="utf-8" ) for path in ("base.py", "rotation.py") @@ -193,7 +193,7 @@ def test_sharpa_grasp_cache_is_materialized_once_and_reset_samples_memory( scale_values = np.asarray([0.8, 1.0], dtype=np.float64) _write_sharpa_grasp_cache(cache_prefix, [0.8, 1.0]) - from unilab.envs.manipulation.sharpa_inhand import rotation + from unilab.tasks.manipulation.sharpa_inhand import rotation path_resolve_calls: list[tuple[str, float]] = [] hf_resolve_calls: list[str] = [] diff --git a/tests/scripts/test_train_script_configs.py b/tests/scripts/test_train_script_configs.py index f1f1575d9..ba0ea43d8 100644 --- a/tests/scripts/test_train_script_configs.py +++ b/tests/scripts/test_train_script_configs.py @@ -20,7 +20,7 @@ def _write_sharpa_smoke_cache(cache_prefix, scale_values: list[float]) -> None: - from unilab.envs.manipulation.sharpa_inhand.base import ( + from unilab.tasks.manipulation.sharpa_inhand.base import ( SOURCE_DEFAULT_HAND_JOINT_POS_DEG, resolve_grasp_cache_file, ) diff --git a/tests/tasks/test_package_boundary.py b/tests/tasks/test_package_boundary.py index c064c3dfd..70a6b6ba6 100644 --- a/tests/tasks/test_package_boundary.py +++ b/tests/tasks/test_package_boundary.py @@ -19,7 +19,7 @@ "unilab.tasks.locomotion.go2_arm", "unilab.tasks.locomotion.a2", "unilab.tasks.manipulation.allegro_inhand", - "unilab.envs.manipulation.sharpa_inhand", + "unilab.tasks.manipulation.sharpa_inhand", "unilab.tasks.manipulation.stewart", "unilab.envs.motion_tracking.g1", "unilab.envs.motion_tracking.x2", diff --git a/tests/test_sharpa.py b/tests/test_sharpa.py index 5b4618f02..c48c376dc 100644 --- a/tests/test_sharpa.py +++ b/tests/test_sharpa.py @@ -5,12 +5,12 @@ import numpy as np import pytest -from unilab.envs.manipulation.sharpa_inhand.base import SharpaDomainRandConfig -from unilab.envs.manipulation.sharpa_inhand.grasp_gen import ( +from unilab.tasks.manipulation.sharpa_inhand.base import SharpaDomainRandConfig +from unilab.tasks.manipulation.sharpa_inhand.grasp_gen import ( SharpaInhandRotationGraspCfg, SharpaInhandRotationGraspEnv, ) -from unilab.envs.manipulation.sharpa_inhand.rotation import ( +from unilab.tasks.manipulation.sharpa_inhand.rotation import ( SharpaInhandRotationEnv, sample_random_quaternion, ) From e0ce62809e01ef225ef63ed8d2dcfa2b4cdb6d40 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 22:22:01 +0800 Subject: [PATCH 054/169] docs(tasks): align manipulation owner paths [skip ci] (#1154) --- .../adr/ADR-0004-registry-bootstrap-contract.md | 4 ++-- .../source/api_reference/envs/manipulation.md | 4 ++-- .../en/2-user_guide/4-tasks/3-manipulation.md | 4 ++-- .../2-writing_providers.md | 2 +- .../1-sim_to_real/4-allegro_inhand.md | 2 +- .../1-architecture/5-registry.md | 17 ++++++----------- .../2-contracts/4-dr_contract.md | 2 +- .../8-motrix_contact_sensor.md | 6 +++--- .../2-user_guide/4-tasks/3-manipulation.md | 4 ++-- .../2-writing_providers.md | 2 +- .../1-sim_to_real/4-allegro_inhand.md | 2 +- .../1-architecture/5-registry.md | 15 ++++++--------- .../2-contracts/4-dr_contract.md | 2 +- .../8-motrix_contact_sensor.md | 6 +++--- 14 files changed, 32 insertions(+), 40 deletions(-) diff --git a/docs/sphinx/source/adr/ADR-0004-registry-bootstrap-contract.md b/docs/sphinx/source/adr/ADR-0004-registry-bootstrap-contract.md index 7c96f790c..5e0691f82 100644 --- a/docs/sphinx/source/adr/ADR-0004-registry-bootstrap-contract.md +++ b/docs/sphinx/source/adr/ADR-0004-registry-bootstrap-contract.md @@ -35,7 +35,7 @@ UniLab 的 env 注册依赖 `@registry.envcfg(...)` 与 `@registry.env(...)` dec ## Consequences -- 新增 env package 时，需要同步声明 bootstrap modules。 +- 新增 task leaf 时，需要在 `unilab.tasks` 中同步声明 bootstrap module。 - registry 相关回归可以在 `ensure_registries()` 边界直接测试，不必依赖顶层训练脚本间接发现。 - 文档可以把 registry bootstrap 作为正式架构引用，而不是“当前实现细节”。 @@ -48,7 +48,7 @@ UniLab 的 env 注册依赖 `@registry.envcfg(...)` 与 `@registry.env(...)` dec - Registry 入口: `src/unilab/base/registry.py` - Bootstrap helper: `src/unilab/base/registry.py` -- Env package 入口: `src/unilab/envs/locomotion/__init__.py`, `src/unilab/envs/motion_tracking/__init__.py`, `src/unilab/envs/manipulation/__init__.py` +- Task package 入口: `src/unilab/tasks/__init__.py` - Bootstrap tests: `tests/utils/test_algo_utils.py`, `tests/base/test_registry.py` ## Related Documents diff --git a/docs/sphinx/source/api_reference/envs/manipulation.md b/docs/sphinx/source/api_reference/envs/manipulation.md index 7734224c0..328e820de 100644 --- a/docs/sphinx/source/api_reference/envs/manipulation.md +++ b/docs/sphinx/source/api_reference/envs/manipulation.md @@ -1,4 +1,4 @@ -# `unilab.envs.manipulation` +# `unilab.tasks.manipulation` ```{eval-rst} .. autosummary:: @@ -7,6 +7,6 @@ :recursive: unilab.tasks.manipulation.allegro_inhand - unilab.envs.manipulation.sharpa_inhand + unilab.tasks.manipulation.sharpa_inhand unilab.tasks.manipulation.stewart ``` diff --git a/docs/sphinx/source/en/2-user_guide/4-tasks/3-manipulation.md b/docs/sphinx/source/en/2-user_guide/4-tasks/3-manipulation.md index 33d577191..c1a60a289 100644 --- a/docs/sphinx/source/en/2-user_guide/4-tasks/3-manipulation.md +++ b/docs/sphinx/source/en/2-user_guide/4-tasks/3-manipulation.md @@ -1,7 +1,7 @@ # Manipulation -Manipulation tasks live in `src/unilab/envs/manipulation/` and the Go2 arm -manip-loco env lives in `src/unilab/envs/locomotion/go2_arm/`. +Manipulation tasks live in `src/unilab/tasks/manipulation/` and the Go2 arm +manip-loco env lives in `src/unilab/tasks/locomotion/go2_arm/`. ## In-Hand 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 1678df319..1491c0ea6 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 @@ -31,7 +31,7 @@ Representative provider implementations are in: - `src/unilab/tasks/locomotion/g1/joystick.py` - `src/unilab/envs/motion_tracking/g1/tracking.py` - `src/unilab/tasks/manipulation/allegro_inhand/rotation.py` -- `src/unilab/envs/manipulation/sharpa_inhand/rotation.py` +- `src/unilab/tasks/manipulation/sharpa_inhand/rotation.py` Developer contract details are in {doc}`../../4-developer_guide/2-contracts/4-dr_contract`. diff --git a/docs/sphinx/source/en/3-deployment/1-sim_to_real/4-allegro_inhand.md b/docs/sphinx/source/en/3-deployment/1-sim_to_real/4-allegro_inhand.md index a42be5ff1..e5105b9b0 100644 --- a/docs/sphinx/source/en/3-deployment/1-sim_to_real/4-allegro_inhand.md +++ b/docs/sphinx/source/en/3-deployment/1-sim_to_real/4-allegro_inhand.md @@ -69,7 +69,7 @@ to the grasp generator**, retrain, and try again. The manipulation envs map policy actions to joint position targets through the task control config (`src/unilab/tasks/manipulation/allegro_inhand/base.py` and -`src/unilab/envs/manipulation/sharpa_inhand/base.py`). The deploy controller +`src/unilab/tasks/manipulation/sharpa_inhand/base.py`). The deploy controller must use the same joint order, action scale, and limit policy. ## Failure recovery diff --git a/docs/sphinx/source/en/4-developer_guide/1-architecture/5-registry.md b/docs/sphinx/source/en/4-developer_guide/1-architecture/5-registry.md index 4b6538cd1..3bf01b4eb 100644 --- a/docs/sphinx/source/en/4-developer_guide/1-architecture/5-registry.md +++ b/docs/sphinx/source/en/4-developer_guide/1-architecture/5-registry.md @@ -8,11 +8,9 @@ defined by {doc}`/adr/ADR-0004-registry-bootstrap-contract` and implemented in 1. Training entrypoints call `unilab.training.common.ensure_registries()`. 2. That helper delegates to `unilab.base.registry.ensure_registries()`. -3. The registry imports declared bootstrap packages: - `unilab.envs.locomotion`, `unilab.envs.manipulation`, and - `unilab.envs.motion_tracking`. -4. Each package exposes `__unilab_registry_modules__`, a tuple of modules that - contain registration side effects. +3. The registry imports its sole declared bootstrap package, `unilab.tasks`. +4. `unilab.tasks` exposes `__unilab_registry_modules__`, an explicit tuple of + task leaf modules that contain registration side effects. 5. Imported modules register configs with `@registry.envcfg(...)` and env implementations with `@registry.env(..., sim_backend=...)` or `registry.register_env(...)`. @@ -22,9 +20,8 @@ defined by {doc}`/adr/ADR-0004-registry-bootstrap-contract` and implemented in ## Extension Rules -- Add new env modules to the package-level `__unilab_registry_modules__` tuple - if they live in a new module that is not imported by an existing bootstrap - entry. +- Add new task leaves to `unilab.tasks.__unilab_registry_modules__` when they + are not imported by an existing bootstrap entry. - Keep registration cheap. Scene materialization, XML processing, asset access, and backend construction belong after `registry.make(...)`, not in decorator registration. @@ -35,7 +32,5 @@ defined by {doc}`/adr/ADR-0004-registry-bootstrap-contract` and implemented in - Bootstrap helper: `src/unilab/base/registry.py` - Training helper: `src/unilab/training/common.py` -- Package declarations: `src/unilab/envs/locomotion/__init__.py`, - `src/unilab/envs/manipulation/__init__.py`, - `src/unilab/envs/motion_tracking/__init__.py` +- Task bootstrap declaration: `src/unilab/tasks/__init__.py` - Tests: `tests/base/test_registry.py`, `tests/utils/test_algo_utils.py` 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 1682f1886..723762745 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 @@ -102,4 +102,4 @@ payloads. - Backend interface: `src/unilab/base/backend/base.py` - Example providers: `src/unilab/tasks/locomotion/g1/joystick.py`, `src/unilab/envs/motion_tracking/g1/tracking.py`, - `src/unilab/envs/manipulation/sharpa_inhand/rotation.py` + `src/unilab/tasks/manipulation/sharpa_inhand/rotation.py` diff --git a/docs/sphinx/source/en/4-developer_guide/8-motrix_contact_sensor.md b/docs/sphinx/source/en/4-developer_guide/8-motrix_contact_sensor.md index 2dd6d323c..805493dce 100644 --- a/docs/sphinx/source/en/4-developer_guide/8-motrix_contact_sensor.md +++ b/docs/sphinx/source/en/4-developer_guide/8-motrix_contact_sensor.md @@ -91,7 +91,7 @@ contact-frame data (one normal scalar plus two tangent scalars). The env reads tactile force through `_read_tactile_force()` → `_extract_sensor_scalar()` in -`src/unilab/envs/manipulation/sharpa_inhand/base.py`. That helper currently +`src/unilab/tasks/manipulation/sharpa_inhand/base.py`. That helper currently collapses any `(N, >=3)` array with `np.linalg.norm(data[:, :3], axis=1)`. If the env still routes both backend shapes through that one branch, the @@ -131,8 +131,8 @@ backend subclass. | File | Role | | --- | --- | -| `src/unilab/envs/manipulation/sharpa_inhand/base.py` | `_extract_sensor_scalar()`, `_read_tactile_force()` | -| `src/unilab/envs/manipulation/sharpa_inhand/rotation.py` | reward computation, virtual torque | +| `src/unilab/tasks/manipulation/sharpa_inhand/base.py` | `_extract_sensor_scalar()`, `_read_tactile_force()` | +| `src/unilab/tasks/manipulation/sharpa_inhand/rotation.py` | reward computation, virtual torque | | `src/unilab/assets/robots/sharpa_wave/right_sharpa_wave.xml` | contact-sensor XML definitions | | `src/unilab/base/backend/motrix/backend.py` | Motrix `get_sensor_data()` | | `src/unilab/base/backend/mujoco/backend.py` | MuJoCo `get_sensor_data()` | diff --git a/docs/sphinx/source/zh_CN/2-user_guide/4-tasks/3-manipulation.md b/docs/sphinx/source/zh_CN/2-user_guide/4-tasks/3-manipulation.md index 6352b4f2a..c6d98ff36 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/4-tasks/3-manipulation.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/4-tasks/3-manipulation.md @@ -1,7 +1,7 @@ # 操作 -操作任务位于 `src/unilab/envs/manipulation/` 中，Go2 机械臂 manip-loco -env 位于 `src/unilab/envs/locomotion/go2_arm/` 中。 +操作任务位于 `src/unilab/tasks/manipulation/` 中，Go2 机械臂 manip-loco +env 位于 `src/unilab/tasks/locomotion/go2_arm/` 中。 ## 手内操作 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 e045f4fd6..19227a448 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 @@ -30,7 +30,7 @@ - `src/unilab/tasks/locomotion/g1/joystick.py` - `src/unilab/envs/motion_tracking/g1/tracking.py` - `src/unilab/tasks/manipulation/allegro_inhand/rotation.py` -- `src/unilab/envs/manipulation/sharpa_inhand/rotation.py` +- `src/unilab/tasks/manipulation/sharpa_inhand/rotation.py` 开发者 contract 详情见 {doc}`../../4-developer_guide/2-contracts/4-dr_contract`。 diff --git a/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/4-allegro_inhand.md b/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/4-allegro_inhand.md index 4edc57715..a202759f8 100644 --- a/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/4-allegro_inhand.md +++ b/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/4-allegro_inhand.md @@ -64,7 +64,7 @@ owner 与部署运行时在观测时序上达成一致。见 操作类环境通过任务控制配置把策略动作映射为关节位置目标 （`src/unilab/tasks/manipulation/allegro_inhand/base.py` 和 -`src/unilab/envs/manipulation/sharpa_inhand/base.py`）。部署控制器必须使用相同的 +`src/unilab/tasks/manipulation/sharpa_inhand/base.py`）。部署控制器必须使用相同的 关节顺序、动作缩放与限位策略。 ## 失败恢复 diff --git a/docs/sphinx/source/zh_CN/4-developer_guide/1-architecture/5-registry.md b/docs/sphinx/source/zh_CN/4-developer_guide/1-architecture/5-registry.md index 0df987d96..c96662e06 100644 --- a/docs/sphinx/source/zh_CN/4-developer_guide/1-architecture/5-registry.md +++ b/docs/sphinx/source/zh_CN/4-developer_guide/1-architecture/5-registry.md @@ -8,10 +8,9 @@ Registry bootstrap 是一个针对环境的显式导入契约。它由 1. 训练入口调用 `unilab.training.common.ensure_registries()`。 2. 该 helper 委托给 `unilab.base.registry.ensure_registries()`。 -3. registry 导入已声明的 bootstrap 包： - `unilab.envs.locomotion`、`unilab.envs.manipulation` 与 - `unilab.envs.motion_tracking`。 -4. 每个包都暴露 `__unilab_registry_modules__`，即一个包含注册副作用的模块元组。 +3. registry 导入唯一声明的 bootstrap 包 `unilab.tasks`。 +4. `unilab.tasks` 暴露 `__unilab_registry_modules__`，即一个包含注册副作用的 + task leaf module 显式元组。 5. 被导入的模块通过 `@registry.envcfg(...)` 注册 config，并通过 `@registry.env(..., sim_backend=...)` 或 `registry.register_env(...)` 注册 env 实现。 @@ -20,8 +19,8 @@ Registry bootstrap 是一个针对环境的显式导入契约。它由 ## 扩展规则 -- 如果新的 env 模块位于某个尚未被现有 bootstrap 条目导入的新模块中，需将其加入 - 包级别的 `__unilab_registry_modules__` 元组。 +- 如果新的 task leaf 尚未被现有 bootstrap 条目导入，需将其加入 + `unilab.tasks.__unilab_registry_modules__`。 - 保持注册过程轻量。场景 materialization、XML 处理、资源访问以及 backend 构造 应放在 `registry.make(...)` 之后，而不是放在装饰器注册中。 - 重复的 env config 以及重复的 `(env, sim_backend)` 注册会在 @@ -31,7 +30,5 @@ Registry bootstrap 是一个针对环境的显式导入契约。它由 - Bootstrap helper：`src/unilab/base/registry.py` - 训练 helper：`src/unilab/training/common.py` -- 包声明：`src/unilab/envs/locomotion/__init__.py`、 - `src/unilab/envs/manipulation/__init__.py`、 - `src/unilab/envs/motion_tracking/__init__.py` +- Task bootstrap 声明：`src/unilab/tasks/__init__.py` - 测试：`tests/base/test_registry.py`、`tests/utils/test_algo_utils.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 1f7f30525..d6120a983 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 @@ -92,4 +92,4 @@ actuator 的机制泄漏到共享 payload 里。 - Backend 接口：`src/unilab/base/backend/base.py` - 示例 provider：`src/unilab/tasks/locomotion/g1/joystick.py`、 `src/unilab/envs/motion_tracking/g1/tracking.py`、 - `src/unilab/envs/manipulation/sharpa_inhand/rotation.py` + `src/unilab/tasks/manipulation/sharpa_inhand/rotation.py` diff --git a/docs/sphinx/source/zh_CN/4-developer_guide/8-motrix_contact_sensor.md b/docs/sphinx/source/zh_CN/4-developer_guide/8-motrix_contact_sensor.md index 2746b8918..656914a93 100644 --- a/docs/sphinx/source/zh_CN/4-developer_guide/8-motrix_contact_sensor.md +++ b/docs/sphinx/source/zh_CN/4-developer_guide/8-motrix_contact_sensor.md @@ -80,7 +80,7 @@ shape = (num_envs, 1 + 4 * 12) = (num_envs, 49) ## 单一 norm 分支为何不够 -env 通过 `src/unilab/envs/manipulation/sharpa_inhand/base.py` 中的 +env 通过 `src/unilab/tasks/manipulation/sharpa_inhand/base.py` 中的 `_read_tactile_force()` → `_extract_sensor_scalar()` 读取触觉力。该 helper 目前对任意 `(N, >=3)` 数组都用 `np.linalg.norm(data[:, :3], axis=1)` 折叠。 如果 env 仍把两种后端形状都走这一个分支，MuJoCo 的 `(N, 3)` 是正确的（对真实力向量取 norm），但 Motrix 的 `(N, 4)` 会出错：`data[:, :3]` 取到的是 `[count, fx, fy]`——把接触点数当成了力分量，并且漏掉了 `fz`。正确做法不是在 env 里按形状特判，而是把每个后端的布局知识下沉到 backend 方法。 @@ -104,8 +104,8 @@ env 层的 `_read_tactile_force()` 对 contact sensor 走 `get_contact_force_mag | 文件 | 说明 | | --- | --- | -| `src/unilab/envs/manipulation/sharpa_inhand/base.py` | `_extract_sensor_scalar()`, `_read_tactile_force()` | -| `src/unilab/envs/manipulation/sharpa_inhand/rotation.py` | reward 计算，virtual torque | +| `src/unilab/tasks/manipulation/sharpa_inhand/base.py` | `_extract_sensor_scalar()`, `_read_tactile_force()` | +| `src/unilab/tasks/manipulation/sharpa_inhand/rotation.py` | reward 计算，virtual torque | | `src/unilab/assets/robots/sharpa_wave/right_sharpa_wave.xml` | contact sensor XML 定义 | | `src/unilab/base/backend/motrix/backend.py` | Motrix `get_sensor_data()` | | `src/unilab/base/backend/mujoco/backend.py` | MuJoCo `get_sensor_data()` | From 6d44de06a55a8cf2ff21586078d55e407e22ae82 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 22:29:03 +0800 Subject: [PATCH 055/169] refactor(tasks): move X2 tracking owner [skip ci] (#1156) --- src/unilab/envs/motion_tracking/__init__.py | 15 +-------------- src/unilab/tasks/__init__.py | 2 +- src/unilab/tasks/motion_tracking/__init__.py | 15 +++++++++++++++ .../motion_tracking/x2/__init__.py | 0 .../motion_tracking/x2/flip_tracking.py | 2 +- .../backend/test_mujoco_scene_context_visual.py | 2 +- tests/envs/test_env_configs.py | 4 ++-- tests/tasks/test_package_boundary.py | 2 +- 8 files changed, 22 insertions(+), 20 deletions(-) create mode 100644 src/unilab/tasks/motion_tracking/__init__.py rename src/unilab/{envs => tasks}/motion_tracking/x2/__init__.py (100%) rename src/unilab/{envs => tasks}/motion_tracking/x2/flip_tracking.py (98%) diff --git a/src/unilab/envs/motion_tracking/__init__.py b/src/unilab/envs/motion_tracking/__init__.py index 1d392bc92..cfa1affcf 100644 --- a/src/unilab/envs/motion_tracking/__init__.py +++ b/src/unilab/envs/motion_tracking/__init__.py @@ -1,9 +1,6 @@ """Motion tracking environments.""" -__unilab_registry_modules__ = ( - "unilab.envs.motion_tracking.g1", - "unilab.envs.motion_tracking.x2", -) +__unilab_registry_modules__ = ("unilab.envs.motion_tracking.g1",) from .g1 import ( BoxMotionData, @@ -44,12 +41,6 @@ G1WBTObsCfg, G1WBTObsEnv, ) -from .x2 import ( - X2MotionTrackingCfg, - X2WallFlipTrackingCfg, - X2WallFlipTrackingEnv, - X2WallFlipTrackingEnvCfg, -) __all__ = [ "G1MotionTrackingCfg", @@ -73,10 +64,6 @@ "G1BoxTrackingEnvCfg", "BoxMotionData", "BoxMotionLoader", - "X2MotionTrackingCfg", - "X2WallFlipTrackingCfg", - "X2WallFlipTrackingEnv", - "X2WallFlipTrackingEnvCfg", # 23-DoF variants (from parent UniLab/, under testing) "G1MotionTracking23DofCfg", "G1MotionTracking23DofDeployEnvCfg", diff --git a/src/unilab/tasks/__init__.py b/src/unilab/tasks/__init__.py index daf7d36ab..9cb175d21 100644 --- a/src/unilab/tasks/__init__.py +++ b/src/unilab/tasks/__init__.py @@ -18,7 +18,7 @@ "unilab.tasks.manipulation.sharpa_inhand", "unilab.tasks.manipulation.stewart", "unilab.envs.motion_tracking.g1", - "unilab.envs.motion_tracking.x2", + "unilab.tasks.motion_tracking.x2", ) __all__ = ["__unilab_registry_modules__"] diff --git a/src/unilab/tasks/motion_tracking/__init__.py b/src/unilab/tasks/motion_tracking/__init__.py new file mode 100644 index 000000000..8144b1da3 --- /dev/null +++ b/src/unilab/tasks/motion_tracking/__init__.py @@ -0,0 +1,15 @@ +"""Motion-tracking task packages.""" + +from .x2 import ( + X2MotionTrackingCfg, + X2WallFlipTrackingCfg, + X2WallFlipTrackingEnv, + X2WallFlipTrackingEnvCfg, +) + +__all__ = [ + "X2MotionTrackingCfg", + "X2WallFlipTrackingCfg", + "X2WallFlipTrackingEnv", + "X2WallFlipTrackingEnvCfg", +] diff --git a/src/unilab/envs/motion_tracking/x2/__init__.py b/src/unilab/tasks/motion_tracking/x2/__init__.py similarity index 100% rename from src/unilab/envs/motion_tracking/x2/__init__.py rename to src/unilab/tasks/motion_tracking/x2/__init__.py diff --git a/src/unilab/envs/motion_tracking/x2/flip_tracking.py b/src/unilab/tasks/motion_tracking/x2/flip_tracking.py similarity index 98% rename from src/unilab/envs/motion_tracking/x2/flip_tracking.py rename to src/unilab/tasks/motion_tracking/x2/flip_tracking.py index da7a0ff1c..367279177 100644 --- a/src/unilab/envs/motion_tracking/x2/flip_tracking.py +++ b/src/unilab/tasks/motion_tracking/x2/flip_tracking.py @@ -122,7 +122,7 @@ class X2WallFlipTrackingEnvCfg(X2WallFlipTrackingCfg): class X2WallFlipTrackingEnv(MotionTrackingDeployEnv): """X2 wall flip-tracking environment implementation.""" - _cfg: X2WallFlipTrackingCfg + _cfg: X2WallFlipTrackingCfg # pyright: ignore[reportIncompatibleVariableOverride] _keyframe_name = "home" def __init__(self, cfg: X2WallFlipTrackingCfg, num_envs: int = 1, backend_type: str = "mujoco"): diff --git a/tests/base/backend/test_mujoco_scene_context_visual.py b/tests/base/backend/test_mujoco_scene_context_visual.py index e1c3d5e14..2aa276f88 100644 --- a/tests/base/backend/test_mujoco_scene_context_visual.py +++ b/tests/base/backend/test_mujoco_scene_context_visual.py @@ -35,7 +35,7 @@ def test_visual_model_file_override_is_render_only() -> None: def test_x2_wall_flip_wires_render_only_visual_twin() -> None: - from unilab.envs.motion_tracking.x2.flip_tracking import X2WallFlipTrackingEnvCfg + from unilab.tasks.motion_tracking.x2.flip_tracking import X2WallFlipTrackingEnvCfg cfg = X2WallFlipTrackingEnvCfg() diff --git a/tests/envs/test_env_configs.py b/tests/envs/test_env_configs.py index 833292ee3..4449f727f 100644 --- a/tests/envs/test_env_configs.py +++ b/tests/envs/test_env_configs.py @@ -57,7 +57,7 @@ def blocked_import(name, globals=None, locals=None, fromlist=(), level=0): G1MotionTrackingCfg, G1MotionTrackingDeployEnvCfg, ) - from unilab.envs.motion_tracking.x2 import X2WallFlipTrackingCfg + from unilab.tasks.motion_tracking.x2 import X2WallFlipTrackingCfg from unilab.base.registry import ensure_registries ensure_registries() @@ -1749,7 +1749,7 @@ def test_g1_wall_flip_tracking_cfg_uses_wall_flip_profile(): def test_x2_wall_flip_tracking_cfg_uses_x2_wall_flip_profile(): - from unilab.envs.motion_tracking.x2 import X2WallFlipTrackingCfg + from unilab.tasks.motion_tracking.x2 import X2WallFlipTrackingCfg cfg = X2WallFlipTrackingCfg() diff --git a/tests/tasks/test_package_boundary.py b/tests/tasks/test_package_boundary.py index 70a6b6ba6..bcd660581 100644 --- a/tests/tasks/test_package_boundary.py +++ b/tests/tasks/test_package_boundary.py @@ -22,7 +22,7 @@ "unilab.tasks.manipulation.sharpa_inhand", "unilab.tasks.manipulation.stewart", "unilab.envs.motion_tracking.g1", - "unilab.envs.motion_tracking.x2", + "unilab.tasks.motion_tracking.x2", ) From 36e610aeed1da6ca1563be40b65b145bda543a03 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 22:37:43 +0800 Subject: [PATCH 056/169] refactor(tasks): move G1 box-tracking owner [skip ci] (#1158) --- src/unilab/envs/motion_tracking/__init__.py | 14 ------------ .../envs/motion_tracking/g1/__init__.py | 15 ------------- src/unilab/tasks/__init__.py | 1 + src/unilab/tasks/motion_tracking/__init__.py | 16 ++++++++++++++ .../tasks/motion_tracking/g1/__init__.py | 20 +++++++++++++++++ .../motion_tracking/g1/box_tracking.py | 20 +++++++++-------- .../motion_tracking/g1/motion_box_loader.py | 2 +- tests/envs/test_env_configs.py | 22 +++++++++---------- tests/envs/test_motion_loader.py | 6 ++--- tests/tasks/test_package_boundary.py | 1 + 10 files changed, 64 insertions(+), 53 deletions(-) create mode 100644 src/unilab/tasks/motion_tracking/g1/__init__.py rename src/unilab/{envs => tasks}/motion_tracking/g1/box_tracking.py (97%) rename src/unilab/{envs => tasks}/motion_tracking/g1/motion_box_loader.py (98%) diff --git a/src/unilab/envs/motion_tracking/__init__.py b/src/unilab/envs/motion_tracking/__init__.py index cfa1affcf..5143050d1 100644 --- a/src/unilab/envs/motion_tracking/__init__.py +++ b/src/unilab/envs/motion_tracking/__init__.py @@ -3,13 +3,6 @@ __unilab_registry_modules__ = ("unilab.envs.motion_tracking.g1",) from .g1 import ( - BoxMotionData, - BoxMotionLoader, - G1BoxTracking23DofCfg, - G1BoxTracking23DofEnvCfg, - G1BoxTrackingCfg, - G1BoxTrackingEnv, - G1BoxTrackingEnvCfg, G1ClimbTrackingCfg, G1ClimbTrackingEnv, G1ClimbTrackingEnvCfg, @@ -59,11 +52,6 @@ "G1ClimbTrackingCfg", "G1ClimbTrackingEnv", "G1ClimbTrackingEnvCfg", - "G1BoxTrackingCfg", - "G1BoxTrackingEnv", - "G1BoxTrackingEnvCfg", - "BoxMotionData", - "BoxMotionLoader", # 23-DoF variants (from parent UniLab/, under testing) "G1MotionTracking23DofCfg", "G1MotionTracking23DofDeployEnvCfg", @@ -79,6 +67,4 @@ "G1WallFlipTracking23DofEnvCfg", "G1WallFlipTrackingSAC23DofCfg", "G1WallFlipTrackingSAC23DofEnv", - "G1BoxTracking23DofCfg", - "G1BoxTracking23DofEnvCfg", ] diff --git a/src/unilab/envs/motion_tracking/g1/__init__.py b/src/unilab/envs/motion_tracking/g1/__init__.py index c97e3866a..694e68627 100644 --- a/src/unilab/envs/motion_tracking/g1/__init__.py +++ b/src/unilab/envs/motion_tracking/g1/__init__.py @@ -1,12 +1,5 @@ """Motion tracking environments for Unitree G1.""" -from .box_tracking import ( - G1BoxTracking23DofCfg, - G1BoxTracking23DofEnvCfg, - G1BoxTrackingCfg, - G1BoxTrackingEnv, - G1BoxTrackingEnvCfg, -) from .flip_tracking import ( G1ClimbTrackingCfg, G1ClimbTrackingEnv, @@ -32,7 +25,6 @@ G1WallFlipTrackingSACCfg, G1WallFlipTrackingSACEnv, ) -from .motion_box_loader import BoxMotionData, BoxMotionLoader from .tracking import ( G1MotionTracking23DofCfg, G1MotionTracking23DofDeployEnvCfg, @@ -88,11 +80,4 @@ "G1ClimbTrackingCfg", "G1ClimbTrackingEnv", "G1ClimbTrackingEnvCfg", - "G1BoxTracking23DofCfg", - "G1BoxTracking23DofEnvCfg", - "G1BoxTrackingCfg", - "G1BoxTrackingEnv", - "G1BoxTrackingEnvCfg", - "BoxMotionData", - "BoxMotionLoader", ] diff --git a/src/unilab/tasks/__init__.py b/src/unilab/tasks/__init__.py index 9cb175d21..53c3ac94a 100644 --- a/src/unilab/tasks/__init__.py +++ b/src/unilab/tasks/__init__.py @@ -18,6 +18,7 @@ "unilab.tasks.manipulation.sharpa_inhand", "unilab.tasks.manipulation.stewart", "unilab.envs.motion_tracking.g1", + "unilab.tasks.motion_tracking.g1.box_tracking", "unilab.tasks.motion_tracking.x2", ) diff --git a/src/unilab/tasks/motion_tracking/__init__.py b/src/unilab/tasks/motion_tracking/__init__.py index 8144b1da3..73c9be54a 100644 --- a/src/unilab/tasks/motion_tracking/__init__.py +++ b/src/unilab/tasks/motion_tracking/__init__.py @@ -1,5 +1,14 @@ """Motion-tracking task packages.""" +from .g1 import ( + BoxMotionData, + BoxMotionLoader, + G1BoxTracking23DofCfg, + G1BoxTracking23DofEnvCfg, + G1BoxTrackingCfg, + G1BoxTrackingEnv, + G1BoxTrackingEnvCfg, +) from .x2 import ( X2MotionTrackingCfg, X2WallFlipTrackingCfg, @@ -8,6 +17,13 @@ ) __all__ = [ + "BoxMotionData", + "BoxMotionLoader", + "G1BoxTracking23DofCfg", + "G1BoxTracking23DofEnvCfg", + "G1BoxTrackingCfg", + "G1BoxTrackingEnv", + "G1BoxTrackingEnvCfg", "X2MotionTrackingCfg", "X2WallFlipTrackingCfg", "X2WallFlipTrackingEnv", diff --git a/src/unilab/tasks/motion_tracking/g1/__init__.py b/src/unilab/tasks/motion_tracking/g1/__init__.py new file mode 100644 index 000000000..376b85adf --- /dev/null +++ b/src/unilab/tasks/motion_tracking/g1/__init__.py @@ -0,0 +1,20 @@ +"""G1 box-tracking tasks.""" + +from .box_tracking import ( + G1BoxTracking23DofCfg, + G1BoxTracking23DofEnvCfg, + G1BoxTrackingCfg, + G1BoxTrackingEnv, + G1BoxTrackingEnvCfg, +) +from .motion_box_loader import BoxMotionData, BoxMotionLoader + +__all__ = [ + "BoxMotionData", + "BoxMotionLoader", + "G1BoxTracking23DofCfg", + "G1BoxTracking23DofEnvCfg", + "G1BoxTrackingCfg", + "G1BoxTrackingEnv", + "G1BoxTrackingEnvCfg", +] diff --git a/src/unilab/envs/motion_tracking/g1/box_tracking.py b/src/unilab/tasks/motion_tracking/g1/box_tracking.py similarity index 97% rename from src/unilab/envs/motion_tracking/g1/box_tracking.py rename to src/unilab/tasks/motion_tracking/g1/box_tracking.py index b49a98afc..7b3fb0bd6 100644 --- a/src/unilab/envs/motion_tracking/g1/box_tracking.py +++ b/src/unilab/tasks/motion_tracking/g1/box_tracking.py @@ -13,6 +13,13 @@ from unilab.dr import DomainRandomizationManager, ResetPlan from unilab.dr.dr_utils import build_common_reset_randomization, zero_actions from unilab.dtype_config import get_global_dtype +from unilab.envs.motion_tracking.common.rewards import RewardContext +from unilab.envs.motion_tracking.g1.tracking import ( + G1MotionTrackingCfg, + G1MotionTrackingDomainRandomizationProvider, + G1MotionTrackingEnv, + RewardConfig, +) from unilab.utils.geometry import np_sample_uniform from unilab.utils.rotation import ( np_matrix_from_quat, @@ -24,14 +31,7 @@ np_subtract_frame_transforms, ) -from ..common.rewards import RewardContext from .motion_box_loader import BoxMotionData, BoxMotionLoader -from .tracking import ( - G1MotionTrackingCfg, - G1MotionTrackingDomainRandomizationProvider, - G1MotionTrackingEnv, - RewardConfig, -) @dataclass @@ -65,7 +65,9 @@ class G1BoxTrackingCfg(G1MotionTrackingCfg): object_body_name: str = "largebox" object_pos_threshold: float = 0.25 object_ori_threshold: float = 0.8 - reward_config: BoxRewardConfig = field(default_factory=BoxRewardConfig) + reward_config: BoxRewardConfig = field( # pyright: ignore[reportIncompatibleVariableOverride] + default_factory=BoxRewardConfig + ) @registry.envcfg("G1BoxTracking") @@ -227,7 +229,7 @@ def build_reset_plan(self, env: Any, env_ids: np.ndarray) -> ResetPlan: class G1BoxTrackingEnv(G1MotionTrackingEnv): """Motion tracking env extended with large-box state and rewards.""" - _cfg: G1BoxTrackingCfg + _cfg: G1BoxTrackingCfg # pyright: ignore[reportIncompatibleVariableOverride] def __init__(self, cfg: G1BoxTrackingCfg, num_envs=1, backend_type="mujoco"): super().__init__(cfg, num_envs, backend_type) diff --git a/src/unilab/envs/motion_tracking/g1/motion_box_loader.py b/src/unilab/tasks/motion_tracking/g1/motion_box_loader.py similarity index 98% rename from src/unilab/envs/motion_tracking/g1/motion_box_loader.py rename to src/unilab/tasks/motion_tracking/g1/motion_box_loader.py index 795f13274..0ef5e421d 100644 --- a/src/unilab/envs/motion_tracking/g1/motion_box_loader.py +++ b/src/unilab/tasks/motion_tracking/g1/motion_box_loader.py @@ -6,7 +6,7 @@ import numpy as np -from ..common.motion_loader import MotionData, MotionLoader +from unilab.envs.motion_tracking.common.motion_loader import MotionData, MotionLoader @dataclass diff --git a/tests/envs/test_env_configs.py b/tests/envs/test_env_configs.py index 4449f727f..3e34cdba0 100644 --- a/tests/envs/test_env_configs.py +++ b/tests/envs/test_env_configs.py @@ -1086,7 +1086,7 @@ def test_g1_motion_tracking_deploy_actor_matches_unitree_mimic_terms(): def test_g1_box_tracking_cfg_uses_largebox_scene_and_motion_defaults(): - from unilab.envs.motion_tracking.g1.box_tracking import BoxRewardConfig, G1BoxTrackingCfg + from unilab.tasks.motion_tracking.g1.box_tracking import BoxRewardConfig, G1BoxTrackingCfg cfg = G1BoxTrackingCfg() @@ -1101,22 +1101,22 @@ def test_g1_box_tracking_cfg_uses_largebox_scene_and_motion_defaults(): def test_g1_box_tracking_is_exported_from_g1_and_motion_tracking_packages(): - from unilab.envs.motion_tracking import ( + from unilab.tasks.motion_tracking import ( G1BoxTrackingCfg as TopLevelCfg, ) - from unilab.envs.motion_tracking import ( + from unilab.tasks.motion_tracking import ( G1BoxTrackingEnv as TopLevelEnv, ) - from unilab.envs.motion_tracking import ( + from unilab.tasks.motion_tracking import ( G1BoxTrackingEnvCfg as TopLevelEnvCfg, ) - from unilab.envs.motion_tracking.g1 import ( + from unilab.tasks.motion_tracking.g1 import ( G1BoxTrackingCfg as G1PkgCfg, ) - from unilab.envs.motion_tracking.g1 import ( + from unilab.tasks.motion_tracking.g1 import ( G1BoxTrackingEnv as G1PkgEnv, ) - from unilab.envs.motion_tracking.g1 import ( + from unilab.tasks.motion_tracking.g1 import ( G1BoxTrackingEnvCfg as G1PkgEnvCfg, ) @@ -1126,8 +1126,8 @@ def test_g1_box_tracking_is_exported_from_g1_and_motion_tracking_packages(): def _compute_g1_box_tracking_obs_stub(): - from unilab.envs.motion_tracking.g1.box_tracking import G1BoxTrackingEnv - from unilab.envs.motion_tracking.g1.motion_box_loader import BoxMotionData + from unilab.tasks.motion_tracking.g1.box_tracking import G1BoxTrackingEnv + from unilab.tasks.motion_tracking.g1.motion_box_loader import BoxMotionData env = cast(Any, object.__new__(G1BoxTrackingEnv)) env._num_envs = 1 @@ -1219,8 +1219,8 @@ def test_g1_box_tracking_actor_matches_deploy_and_critic_adds_object_state(): def test_g1_box_tracking_critic_object_state_respects_subset_env_order(): - from unilab.envs.motion_tracking.g1.box_tracking import G1BoxTrackingEnv - from unilab.envs.motion_tracking.g1.motion_box_loader import BoxMotionData + from unilab.tasks.motion_tracking.g1.box_tracking import G1BoxTrackingEnv + from unilab.tasks.motion_tracking.g1.motion_box_loader import BoxMotionData env = cast(Any, object.__new__(G1BoxTrackingEnv)) env._num_envs = 4 diff --git a/tests/envs/test_motion_loader.py b/tests/envs/test_motion_loader.py index 1b5069faf..e717fe99f 100644 --- a/tests/envs/test_motion_loader.py +++ b/tests/envs/test_motion_loader.py @@ -185,7 +185,7 @@ def test_motion_sampler_step_respects_current_clip_end(tmp_path): def test_box_motion_loader_reads_object_state_and_trims_robot_joints(tmp_path): - from unilab.envs.motion_tracking.g1.motion_box_loader import BoxMotionLoader + from unilab.tasks.motion_tracking.g1.motion_box_loader import BoxMotionLoader motion = tmp_path / "motion_box.npz" _write_box_motion_npz(motion, base_value=1.0, num_frames=2, num_joints=2) @@ -207,7 +207,7 @@ def test_box_motion_loader_reads_object_state_and_trims_robot_joints(tmp_path): def test_box_motion_loader_rejects_partial_object_key_sets(tmp_path): - from unilab.envs.motion_tracking.g1.motion_box_loader import BoxMotionLoader + from unilab.tasks.motion_tracking.g1.motion_box_loader import BoxMotionLoader motion = tmp_path / "motion_box_missing_keys.npz" _write_box_motion_npz(motion, base_value=1.0, num_frames=2, num_joints=2) @@ -221,7 +221,7 @@ def test_box_motion_loader_rejects_partial_object_key_sets(tmp_path): def test_box_motion_loader_rejects_multi_clip_object_presence_mismatch(tmp_path): - from unilab.envs.motion_tracking.g1.motion_box_loader import BoxMotionLoader + from unilab.tasks.motion_tracking.g1.motion_box_loader import BoxMotionLoader motion_without_object = tmp_path / "motion_without_object.npz" motion_with_object = tmp_path / "motion_with_object.npz" diff --git a/tests/tasks/test_package_boundary.py b/tests/tasks/test_package_boundary.py index bcd660581..7d0667f5d 100644 --- a/tests/tasks/test_package_boundary.py +++ b/tests/tasks/test_package_boundary.py @@ -22,6 +22,7 @@ "unilab.tasks.manipulation.sharpa_inhand", "unilab.tasks.manipulation.stewart", "unilab.envs.motion_tracking.g1", + "unilab.tasks.motion_tracking.g1.box_tracking", "unilab.tasks.motion_tracking.x2", ) From 8afbf6f69cd047fffdb4b518b333a1bf3f1f8a72 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 22:45:49 +0800 Subject: [PATCH 057/169] refactor(tasks): move G1 SAC tracking leaves [skip ci] (#1160) --- src/unilab/envs/motion_tracking/__init__.py | 22 ------------ .../envs/motion_tracking/g1/__init__.py | 32 ----------------- src/unilab/tasks/__init__.py | 2 +- src/unilab/tasks/motion_tracking/__init__.py | 30 ++++++++++++++++ .../tasks/motion_tracking/g1/__init__.py | 34 ++++++++++++++++++- .../motion_tracking/g1/flip_tracking_sac.py | 12 +++---- .../motion_tracking/g1/tracking_obs.py | 8 ++--- .../motion_tracking/g1/tracking_sac.py | 9 +++-- tests/tasks/test_package_boundary.py | 2 +- 9 files changed, 81 insertions(+), 70 deletions(-) rename src/unilab/{envs => tasks}/motion_tracking/g1/flip_tracking_sac.py (86%) rename src/unilab/{envs => tasks}/motion_tracking/g1/tracking_obs.py (99%) rename src/unilab/{envs => tasks}/motion_tracking/g1/tracking_sac.py (92%) diff --git a/src/unilab/envs/motion_tracking/__init__.py b/src/unilab/envs/motion_tracking/__init__.py index 5143050d1..9e1410ebd 100644 --- a/src/unilab/envs/motion_tracking/__init__.py +++ b/src/unilab/envs/motion_tracking/__init__.py @@ -11,38 +11,23 @@ G1FlipTrackingCfg, G1FlipTrackingEnv, G1FlipTrackingEnvCfg, - G1FlipTrackingSAC23DofCfg, - G1FlipTrackingSAC23DofEnv, G1MotionTracking23DofCfg, G1MotionTracking23DofDeployEnvCfg, G1MotionTracking23DofEnvCfg, G1MotionTrackingCfg, G1MotionTrackingEnv, G1MotionTrackingEnvCfg, - G1MotionTrackingSAC23DofCfg, - G1MotionTrackingSAC23DofEnv, - G1MotionTrackingSACCfg, - G1MotionTrackingSACEnv, G1WallFlipTracking23DofCfg, G1WallFlipTracking23DofEnvCfg, G1WallFlipTrackingCfg, G1WallFlipTrackingEnv, G1WallFlipTrackingEnvCfg, - G1WallFlipTrackingSAC23DofCfg, - G1WallFlipTrackingSAC23DofEnv, - G1WBTObs23DofCfg, - G1WBTObsCfg, - G1WBTObsEnv, ) __all__ = [ "G1MotionTrackingCfg", "G1MotionTrackingEnv", "G1MotionTrackingEnvCfg", - "G1MotionTrackingSACCfg", - "G1MotionTrackingSACEnv", - "G1WBTObsCfg", - "G1WBTObsEnv", "G1FlipTrackingCfg", "G1FlipTrackingEnv", "G1FlipTrackingEnvCfg", @@ -56,15 +41,8 @@ "G1MotionTracking23DofCfg", "G1MotionTracking23DofDeployEnvCfg", "G1MotionTracking23DofEnvCfg", - "G1MotionTrackingSAC23DofCfg", - "G1MotionTrackingSAC23DofEnv", - "G1WBTObs23DofCfg", "G1FlipTracking23DofCfg", "G1FlipTracking23DofEnvCfg", - "G1FlipTrackingSAC23DofCfg", - "G1FlipTrackingSAC23DofEnv", "G1WallFlipTracking23DofCfg", "G1WallFlipTracking23DofEnvCfg", - "G1WallFlipTrackingSAC23DofCfg", - "G1WallFlipTrackingSAC23DofEnv", ] diff --git a/src/unilab/envs/motion_tracking/g1/__init__.py b/src/unilab/envs/motion_tracking/g1/__init__.py index 694e68627..eb9e235a7 100644 --- a/src/unilab/envs/motion_tracking/g1/__init__.py +++ b/src/unilab/envs/motion_tracking/g1/__init__.py @@ -15,16 +15,6 @@ G1WallFlipTrackingEnv, G1WallFlipTrackingEnvCfg, ) -from .flip_tracking_sac import ( - G1FlipTrackingSAC23DofCfg, - G1FlipTrackingSAC23DofEnv, - G1FlipTrackingSACCfg, - G1FlipTrackingSACEnv, - G1WallFlipTrackingSAC23DofCfg, - G1WallFlipTrackingSAC23DofEnv, - G1WallFlipTrackingSACCfg, - G1WallFlipTrackingSACEnv, -) from .tracking import ( G1MotionTracking23DofCfg, G1MotionTracking23DofDeployEnvCfg, @@ -35,13 +25,6 @@ G1MotionTrackingEnv, G1MotionTrackingEnvCfg, ) -from .tracking_obs import G1WBTObs23DofCfg, G1WBTObsCfg, G1WBTObsEnv -from .tracking_sac import ( - G1MotionTrackingSAC23DofCfg, - G1MotionTrackingSAC23DofEnv, - G1MotionTrackingSACCfg, - G1MotionTrackingSACEnv, -) __all__ = [ "G1MotionTracking23DofCfg", @@ -52,31 +35,16 @@ "G1MotionTrackingDeployEnvCfg", "G1MotionTrackingEnv", "G1MotionTrackingEnvCfg", - "G1MotionTrackingSAC23DofCfg", - "G1MotionTrackingSAC23DofEnv", - "G1MotionTrackingSACCfg", - "G1MotionTrackingSACEnv", - "G1WBTObs23DofCfg", - "G1WBTObsCfg", - "G1WBTObsEnv", "G1FlipTracking23DofCfg", "G1FlipTracking23DofEnvCfg", "G1FlipTrackingCfg", "G1FlipTrackingEnv", "G1FlipTrackingEnvCfg", - "G1FlipTrackingSAC23DofCfg", - "G1FlipTrackingSAC23DofEnv", - "G1FlipTrackingSACCfg", - "G1FlipTrackingSACEnv", "G1WallFlipTracking23DofCfg", "G1WallFlipTracking23DofEnvCfg", "G1WallFlipTrackingCfg", "G1WallFlipTrackingEnv", "G1WallFlipTrackingEnvCfg", - "G1WallFlipTrackingSAC23DofCfg", - "G1WallFlipTrackingSAC23DofEnv", - "G1WallFlipTrackingSACCfg", - "G1WallFlipTrackingSACEnv", "G1ClimbTrackingCfg", "G1ClimbTrackingEnv", "G1ClimbTrackingEnvCfg", diff --git a/src/unilab/tasks/__init__.py b/src/unilab/tasks/__init__.py index 53c3ac94a..f103c90bd 100644 --- a/src/unilab/tasks/__init__.py +++ b/src/unilab/tasks/__init__.py @@ -18,7 +18,7 @@ "unilab.tasks.manipulation.sharpa_inhand", "unilab.tasks.manipulation.stewart", "unilab.envs.motion_tracking.g1", - "unilab.tasks.motion_tracking.g1.box_tracking", + "unilab.tasks.motion_tracking.g1", "unilab.tasks.motion_tracking.x2", ) diff --git a/src/unilab/tasks/motion_tracking/__init__.py b/src/unilab/tasks/motion_tracking/__init__.py index 73c9be54a..a5446750e 100644 --- a/src/unilab/tasks/motion_tracking/__init__.py +++ b/src/unilab/tasks/motion_tracking/__init__.py @@ -8,6 +8,21 @@ G1BoxTrackingCfg, G1BoxTrackingEnv, G1BoxTrackingEnvCfg, + G1FlipTrackingSAC23DofCfg, + G1FlipTrackingSAC23DofEnv, + G1FlipTrackingSACCfg, + G1FlipTrackingSACEnv, + G1MotionTrackingSAC23DofCfg, + G1MotionTrackingSAC23DofEnv, + G1MotionTrackingSACCfg, + G1MotionTrackingSACEnv, + G1WallFlipTrackingSAC23DofCfg, + G1WallFlipTrackingSAC23DofEnv, + G1WallFlipTrackingSACCfg, + G1WallFlipTrackingSACEnv, + G1WBTObs23DofCfg, + G1WBTObsCfg, + G1WBTObsEnv, ) from .x2 import ( X2MotionTrackingCfg, @@ -24,6 +39,21 @@ "G1BoxTrackingCfg", "G1BoxTrackingEnv", "G1BoxTrackingEnvCfg", + "G1FlipTrackingSAC23DofCfg", + "G1FlipTrackingSAC23DofEnv", + "G1FlipTrackingSACCfg", + "G1FlipTrackingSACEnv", + "G1MotionTrackingSAC23DofCfg", + "G1MotionTrackingSAC23DofEnv", + "G1MotionTrackingSACCfg", + "G1MotionTrackingSACEnv", + "G1WallFlipTrackingSAC23DofCfg", + "G1WallFlipTrackingSAC23DofEnv", + "G1WallFlipTrackingSACCfg", + "G1WallFlipTrackingSACEnv", + "G1WBTObs23DofCfg", + "G1WBTObsCfg", + "G1WBTObsEnv", "X2MotionTrackingCfg", "X2WallFlipTrackingCfg", "X2WallFlipTrackingEnv", diff --git a/src/unilab/tasks/motion_tracking/g1/__init__.py b/src/unilab/tasks/motion_tracking/g1/__init__.py index 376b85adf..c469de565 100644 --- a/src/unilab/tasks/motion_tracking/g1/__init__.py +++ b/src/unilab/tasks/motion_tracking/g1/__init__.py @@ -1,4 +1,4 @@ -"""G1 box-tracking tasks.""" +"""G1 motion-tracking tasks.""" from .box_tracking import ( G1BoxTracking23DofCfg, @@ -7,7 +7,24 @@ G1BoxTrackingEnv, G1BoxTrackingEnvCfg, ) +from .flip_tracking_sac import ( + G1FlipTrackingSAC23DofCfg, + G1FlipTrackingSAC23DofEnv, + G1FlipTrackingSACCfg, + G1FlipTrackingSACEnv, + G1WallFlipTrackingSAC23DofCfg, + G1WallFlipTrackingSAC23DofEnv, + G1WallFlipTrackingSACCfg, + G1WallFlipTrackingSACEnv, +) from .motion_box_loader import BoxMotionData, BoxMotionLoader +from .tracking_obs import G1WBTObs23DofCfg, G1WBTObsCfg, G1WBTObsEnv +from .tracking_sac import ( + G1MotionTrackingSAC23DofCfg, + G1MotionTrackingSAC23DofEnv, + G1MotionTrackingSACCfg, + G1MotionTrackingSACEnv, +) __all__ = [ "BoxMotionData", @@ -17,4 +34,19 @@ "G1BoxTrackingCfg", "G1BoxTrackingEnv", "G1BoxTrackingEnvCfg", + "G1FlipTrackingSAC23DofCfg", + "G1FlipTrackingSAC23DofEnv", + "G1FlipTrackingSACCfg", + "G1FlipTrackingSACEnv", + "G1MotionTrackingSAC23DofCfg", + "G1MotionTrackingSAC23DofEnv", + "G1MotionTrackingSACCfg", + "G1MotionTrackingSACEnv", + "G1WBTObs23DofCfg", + "G1WBTObsCfg", + "G1WBTObsEnv", + "G1WallFlipTrackingSAC23DofCfg", + "G1WallFlipTrackingSAC23DofEnv", + "G1WallFlipTrackingSACCfg", + "G1WallFlipTrackingSACEnv", ] diff --git a/src/unilab/envs/motion_tracking/g1/flip_tracking_sac.py b/src/unilab/tasks/motion_tracking/g1/flip_tracking_sac.py similarity index 86% rename from src/unilab/envs/motion_tracking/g1/flip_tracking_sac.py rename to src/unilab/tasks/motion_tracking/g1/flip_tracking_sac.py index 894939d2f..f43dc4c15 100644 --- a/src/unilab/envs/motion_tracking/g1/flip_tracking_sac.py +++ b/src/unilab/tasks/motion_tracking/g1/flip_tracking_sac.py @@ -17,13 +17,13 @@ from dataclasses import dataclass from unilab.base import registry - -from .flip_tracking import ( +from unilab.envs.motion_tracking.g1.flip_tracking import ( G1FlipTracking23DofEnvCfg, G1FlipTrackingCfg, G1WallFlipTracking23DofEnvCfg, G1WallFlipTrackingCfg, ) + from .tracking_sac import G1MotionTrackingSACEnv @@ -41,7 +41,7 @@ class G1FlipTrackingSACEnv(G1MotionTrackingSACEnv): The motrix backend is registered for sim2sim eval/playback only. """ - _cfg: G1FlipTrackingSACCfg + _cfg: G1FlipTrackingSACCfg # pyright: ignore[reportIncompatibleVariableOverride] @registry.envcfg("G1WallFlipTrackingSAC") @@ -58,7 +58,7 @@ class G1WallFlipTrackingSACEnv(G1MotionTrackingSACEnv): The motrix backend is registered for sim2sim eval/playback only. """ - _cfg: G1WallFlipTrackingSACCfg + _cfg: G1WallFlipTrackingSACCfg # pyright: ignore[reportIncompatibleVariableOverride] @registry.envcfg("G1FlipTrackingSAC23Dof") @@ -70,7 +70,7 @@ class G1FlipTrackingSAC23DofCfg(G1FlipTracking23DofEnvCfg): @registry.env("G1FlipTrackingSAC23Dof", sim_backend="mujoco") @registry.env("G1FlipTrackingSAC23Dof", sim_backend="motrix") class G1FlipTrackingSAC23DofEnv(G1MotionTrackingSACEnv): - _cfg: G1FlipTrackingSAC23DofCfg + _cfg: G1FlipTrackingSAC23DofCfg # pyright: ignore[reportIncompatibleVariableOverride] @registry.envcfg("G1WallFlipTrackingSAC23Dof") @@ -82,4 +82,4 @@ class G1WallFlipTrackingSAC23DofCfg(G1WallFlipTracking23DofEnvCfg): @registry.env("G1WallFlipTrackingSAC23Dof", sim_backend="mujoco") @registry.env("G1WallFlipTrackingSAC23Dof", sim_backend="motrix") class G1WallFlipTrackingSAC23DofEnv(G1MotionTrackingSACEnv): - _cfg: G1WallFlipTrackingSAC23DofCfg + _cfg: G1WallFlipTrackingSAC23DofCfg # pyright: ignore[reportIncompatibleVariableOverride] diff --git a/src/unilab/envs/motion_tracking/g1/tracking_obs.py b/src/unilab/tasks/motion_tracking/g1/tracking_obs.py similarity index 99% rename from src/unilab/envs/motion_tracking/g1/tracking_obs.py rename to src/unilab/tasks/motion_tracking/g1/tracking_obs.py index 270c4bdc6..2efe1cc77 100644 --- a/src/unilab/envs/motion_tracking/g1/tracking_obs.py +++ b/src/unilab/tasks/motion_tracking/g1/tracking_obs.py @@ -44,13 +44,13 @@ from unilab.dr.types import RESET_TERM_GEOM_FRICTION from unilab.dtype_config import get_global_dtype from unilab.envs.locomotion.g1.base import NoiseConfig - -from ..common.rewards import RewardContext -from .tracking import ( +from unilab.envs.motion_tracking.common.rewards import RewardContext +from unilab.envs.motion_tracking.g1.tracking import ( Domain_Rand, G1MotionTrackingDomainRandomizationProvider, _build_motion_reference_state, ) + from .tracking_sac import G1MotionTrackingSACCfg, G1MotionTrackingSACEnv # --------------------------------------------------------------------------- # @@ -230,7 +230,7 @@ class G1WBTObsEnv(G1MotionTrackingSACEnv): on ``G1WBTObsCfg`` are toggled from the task yaml. """ - _cfg: G1WBTObsCfg + _cfg: G1WBTObsCfg # pyright: ignore[reportIncompatibleVariableOverride] def __init__(self, cfg: G1WBTObsCfg, num_envs: int = 1, backend_type: str = "mujoco"): super().__init__(cfg, num_envs=num_envs, backend_type=backend_type) diff --git a/src/unilab/envs/motion_tracking/g1/tracking_sac.py b/src/unilab/tasks/motion_tracking/g1/tracking_sac.py similarity index 92% rename from src/unilab/envs/motion_tracking/g1/tracking_sac.py rename to src/unilab/tasks/motion_tracking/g1/tracking_sac.py index db968b287..78dce9ab4 100644 --- a/src/unilab/envs/motion_tracking/g1/tracking_sac.py +++ b/src/unilab/tasks/motion_tracking/g1/tracking_sac.py @@ -15,8 +15,11 @@ from unilab.base import registry from unilab.dtype_config import get_global_dtype - -from .tracking import G1MotionTracking23DofCfg, G1MotionTrackingCfg, G1MotionTrackingEnv +from unilab.envs.motion_tracking.g1.tracking import ( + G1MotionTracking23DofCfg, + G1MotionTrackingCfg, + G1MotionTrackingEnv, +) @registry.envcfg("G1MotionTrackingSAC") @@ -80,4 +83,4 @@ class G1MotionTrackingSAC23DofCfg(G1MotionTracking23DofCfg): @registry.env("G1MotionTrackingSAC23Dof", sim_backend="mujoco") @registry.env("G1MotionTrackingSAC23Dof", sim_backend="motrix") class G1MotionTrackingSAC23DofEnv(G1MotionTrackingSACEnv): - _cfg: G1MotionTrackingSAC23DofCfg + _cfg: G1MotionTrackingSAC23DofCfg # pyright: ignore[reportIncompatibleVariableOverride] diff --git a/tests/tasks/test_package_boundary.py b/tests/tasks/test_package_boundary.py index 7d0667f5d..296d4bfc5 100644 --- a/tests/tasks/test_package_boundary.py +++ b/tests/tasks/test_package_boundary.py @@ -22,7 +22,7 @@ "unilab.tasks.manipulation.sharpa_inhand", "unilab.tasks.manipulation.stewart", "unilab.envs.motion_tracking.g1", - "unilab.tasks.motion_tracking.g1.box_tracking", + "unilab.tasks.motion_tracking.g1", "unilab.tasks.motion_tracking.x2", ) From 432d70a3ec08cb381ff0e770f81bf2c132cd1ab2 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 22:51:33 +0800 Subject: [PATCH 058/169] refactor(tasks): move G1 flip tracking owner [skip ci] (#1162) --- src/unilab/envs/motion_tracking/__init__.py | 26 ----------------- .../envs/motion_tracking/g1/__init__.py | 28 ------------------- src/unilab/tasks/motion_tracking/__init__.py | 26 +++++++++++++++++ .../tasks/motion_tracking/g1/__init__.py | 28 +++++++++++++++++++ .../motion_tracking/g1/flip_tracking.py | 14 ++++++---- .../motion_tracking/g1/flip_tracking_sac.py | 4 +-- tests/envs/test_env_configs.py | 4 +-- 7 files changed, 66 insertions(+), 64 deletions(-) rename src/unilab/{envs => tasks}/motion_tracking/g1/flip_tracking.py (94%) diff --git a/src/unilab/envs/motion_tracking/__init__.py b/src/unilab/envs/motion_tracking/__init__.py index 9e1410ebd..14d310250 100644 --- a/src/unilab/envs/motion_tracking/__init__.py +++ b/src/unilab/envs/motion_tracking/__init__.py @@ -3,46 +3,20 @@ __unilab_registry_modules__ = ("unilab.envs.motion_tracking.g1",) from .g1 import ( - G1ClimbTrackingCfg, - G1ClimbTrackingEnv, - G1ClimbTrackingEnvCfg, - G1FlipTracking23DofCfg, - G1FlipTracking23DofEnvCfg, - G1FlipTrackingCfg, - G1FlipTrackingEnv, - G1FlipTrackingEnvCfg, G1MotionTracking23DofCfg, G1MotionTracking23DofDeployEnvCfg, G1MotionTracking23DofEnvCfg, G1MotionTrackingCfg, G1MotionTrackingEnv, G1MotionTrackingEnvCfg, - G1WallFlipTracking23DofCfg, - G1WallFlipTracking23DofEnvCfg, - G1WallFlipTrackingCfg, - G1WallFlipTrackingEnv, - G1WallFlipTrackingEnvCfg, ) __all__ = [ "G1MotionTrackingCfg", "G1MotionTrackingEnv", "G1MotionTrackingEnvCfg", - "G1FlipTrackingCfg", - "G1FlipTrackingEnv", - "G1FlipTrackingEnvCfg", - "G1WallFlipTrackingCfg", - "G1WallFlipTrackingEnv", - "G1WallFlipTrackingEnvCfg", - "G1ClimbTrackingCfg", - "G1ClimbTrackingEnv", - "G1ClimbTrackingEnvCfg", # 23-DoF variants (from parent UniLab/, under testing) "G1MotionTracking23DofCfg", "G1MotionTracking23DofDeployEnvCfg", "G1MotionTracking23DofEnvCfg", - "G1FlipTracking23DofCfg", - "G1FlipTracking23DofEnvCfg", - "G1WallFlipTracking23DofCfg", - "G1WallFlipTracking23DofEnvCfg", ] diff --git a/src/unilab/envs/motion_tracking/g1/__init__.py b/src/unilab/envs/motion_tracking/g1/__init__.py index eb9e235a7..0df469931 100644 --- a/src/unilab/envs/motion_tracking/g1/__init__.py +++ b/src/unilab/envs/motion_tracking/g1/__init__.py @@ -1,20 +1,5 @@ """Motion tracking environments for Unitree G1.""" -from .flip_tracking import ( - G1ClimbTrackingCfg, - G1ClimbTrackingEnv, - G1ClimbTrackingEnvCfg, - G1FlipTracking23DofCfg, - G1FlipTracking23DofEnvCfg, - G1FlipTrackingCfg, - G1FlipTrackingEnv, - G1FlipTrackingEnvCfg, - G1WallFlipTracking23DofCfg, - G1WallFlipTracking23DofEnvCfg, - G1WallFlipTrackingCfg, - G1WallFlipTrackingEnv, - G1WallFlipTrackingEnvCfg, -) from .tracking import ( G1MotionTracking23DofCfg, G1MotionTracking23DofDeployEnvCfg, @@ -35,17 +20,4 @@ "G1MotionTrackingDeployEnvCfg", "G1MotionTrackingEnv", "G1MotionTrackingEnvCfg", - "G1FlipTracking23DofCfg", - "G1FlipTracking23DofEnvCfg", - "G1FlipTrackingCfg", - "G1FlipTrackingEnv", - "G1FlipTrackingEnvCfg", - "G1WallFlipTracking23DofCfg", - "G1WallFlipTracking23DofEnvCfg", - "G1WallFlipTrackingCfg", - "G1WallFlipTrackingEnv", - "G1WallFlipTrackingEnvCfg", - "G1ClimbTrackingCfg", - "G1ClimbTrackingEnv", - "G1ClimbTrackingEnvCfg", ] diff --git a/src/unilab/tasks/motion_tracking/__init__.py b/src/unilab/tasks/motion_tracking/__init__.py index a5446750e..47e830b17 100644 --- a/src/unilab/tasks/motion_tracking/__init__.py +++ b/src/unilab/tasks/motion_tracking/__init__.py @@ -8,6 +8,14 @@ G1BoxTrackingCfg, G1BoxTrackingEnv, G1BoxTrackingEnvCfg, + G1ClimbTrackingCfg, + G1ClimbTrackingEnv, + G1ClimbTrackingEnvCfg, + G1FlipTracking23DofCfg, + G1FlipTracking23DofEnvCfg, + G1FlipTrackingCfg, + G1FlipTrackingEnv, + G1FlipTrackingEnvCfg, G1FlipTrackingSAC23DofCfg, G1FlipTrackingSAC23DofEnv, G1FlipTrackingSACCfg, @@ -16,6 +24,11 @@ G1MotionTrackingSAC23DofEnv, G1MotionTrackingSACCfg, G1MotionTrackingSACEnv, + G1WallFlipTracking23DofCfg, + G1WallFlipTracking23DofEnvCfg, + G1WallFlipTrackingCfg, + G1WallFlipTrackingEnv, + G1WallFlipTrackingEnvCfg, G1WallFlipTrackingSAC23DofCfg, G1WallFlipTrackingSAC23DofEnv, G1WallFlipTrackingSACCfg, @@ -39,6 +52,14 @@ "G1BoxTrackingCfg", "G1BoxTrackingEnv", "G1BoxTrackingEnvCfg", + "G1ClimbTrackingCfg", + "G1ClimbTrackingEnv", + "G1ClimbTrackingEnvCfg", + "G1FlipTracking23DofCfg", + "G1FlipTracking23DofEnvCfg", + "G1FlipTrackingCfg", + "G1FlipTrackingEnv", + "G1FlipTrackingEnvCfg", "G1FlipTrackingSAC23DofCfg", "G1FlipTrackingSAC23DofEnv", "G1FlipTrackingSACCfg", @@ -47,6 +68,11 @@ "G1MotionTrackingSAC23DofEnv", "G1MotionTrackingSACCfg", "G1MotionTrackingSACEnv", + "G1WallFlipTracking23DofCfg", + "G1WallFlipTracking23DofEnvCfg", + "G1WallFlipTrackingCfg", + "G1WallFlipTrackingEnv", + "G1WallFlipTrackingEnvCfg", "G1WallFlipTrackingSAC23DofCfg", "G1WallFlipTrackingSAC23DofEnv", "G1WallFlipTrackingSACCfg", diff --git a/src/unilab/tasks/motion_tracking/g1/__init__.py b/src/unilab/tasks/motion_tracking/g1/__init__.py index c469de565..f85ce164a 100644 --- a/src/unilab/tasks/motion_tracking/g1/__init__.py +++ b/src/unilab/tasks/motion_tracking/g1/__init__.py @@ -7,6 +7,21 @@ G1BoxTrackingEnv, G1BoxTrackingEnvCfg, ) +from .flip_tracking import ( + G1ClimbTrackingCfg, + G1ClimbTrackingEnv, + G1ClimbTrackingEnvCfg, + G1FlipTracking23DofCfg, + G1FlipTracking23DofEnvCfg, + G1FlipTrackingCfg, + G1FlipTrackingEnv, + G1FlipTrackingEnvCfg, + G1WallFlipTracking23DofCfg, + G1WallFlipTracking23DofEnvCfg, + G1WallFlipTrackingCfg, + G1WallFlipTrackingEnv, + G1WallFlipTrackingEnvCfg, +) from .flip_tracking_sac import ( G1FlipTrackingSAC23DofCfg, G1FlipTrackingSAC23DofEnv, @@ -34,6 +49,14 @@ "G1BoxTrackingCfg", "G1BoxTrackingEnv", "G1BoxTrackingEnvCfg", + "G1ClimbTrackingCfg", + "G1ClimbTrackingEnv", + "G1ClimbTrackingEnvCfg", + "G1FlipTracking23DofCfg", + "G1FlipTracking23DofEnvCfg", + "G1FlipTrackingCfg", + "G1FlipTrackingEnv", + "G1FlipTrackingEnvCfg", "G1FlipTrackingSAC23DofCfg", "G1FlipTrackingSAC23DofEnv", "G1FlipTrackingSACCfg", @@ -45,6 +68,11 @@ "G1WBTObs23DofCfg", "G1WBTObsCfg", "G1WBTObsEnv", + "G1WallFlipTracking23DofCfg", + "G1WallFlipTracking23DofEnvCfg", + "G1WallFlipTrackingCfg", + "G1WallFlipTrackingEnv", + "G1WallFlipTrackingEnvCfg", "G1WallFlipTrackingSAC23DofCfg", "G1WallFlipTrackingSAC23DofEnv", "G1WallFlipTrackingSACCfg", diff --git a/src/unilab/envs/motion_tracking/g1/flip_tracking.py b/src/unilab/tasks/motion_tracking/g1/flip_tracking.py similarity index 94% rename from src/unilab/envs/motion_tracking/g1/flip_tracking.py rename to src/unilab/tasks/motion_tracking/g1/flip_tracking.py index a8d2aef00..2ce3fe4b0 100644 --- a/src/unilab/envs/motion_tracking/g1/flip_tracking.py +++ b/src/unilab/tasks/motion_tracking/g1/flip_tracking.py @@ -12,9 +12,11 @@ from unilab.assets import ASSETS_ROOT_PATH from unilab.base import registry from unilab.base.scene import SceneCfg - -from ..common.config import _zero_pose_randomization, _zero_velocity_randomization -from .tracking import ( +from unilab.envs.motion_tracking.common.config import ( + _zero_pose_randomization, + _zero_velocity_randomization, +) +from unilab.envs.motion_tracking.g1.tracking import ( G1MotionTrackingCfg, G1MotionTrackingEnv, PoseRandomization, @@ -58,7 +60,7 @@ class G1FlipTrackingEnvCfg(G1FlipTrackingCfg): class G1FlipTrackingEnv(G1MotionTrackingEnv): """G1 flip-tracking environment implementation.""" - _cfg: G1FlipTrackingCfg + _cfg: G1FlipTrackingCfg # pyright: ignore[reportIncompatibleVariableOverride] @dataclass @@ -91,7 +93,7 @@ class G1WallFlipTrackingEnvCfg(G1WallFlipTrackingCfg): class G1WallFlipTrackingEnv(G1MotionTrackingEnv): """G1 wall flip-tracking environment implementation.""" - _cfg: G1WallFlipTrackingCfg + _cfg: G1WallFlipTrackingCfg # pyright: ignore[reportIncompatibleVariableOverride] @dataclass @@ -124,7 +126,7 @@ class G1ClimbTrackingEnvCfg(G1ClimbTrackingCfg): class G1ClimbTrackingEnv(G1MotionTrackingEnv): """G1 climb-tracking environment implementation.""" - _cfg: G1ClimbTrackingCfg + _cfg: G1ClimbTrackingCfg # pyright: ignore[reportIncompatibleVariableOverride] @dataclass diff --git a/src/unilab/tasks/motion_tracking/g1/flip_tracking_sac.py b/src/unilab/tasks/motion_tracking/g1/flip_tracking_sac.py index f43dc4c15..f3ec25e5d 100644 --- a/src/unilab/tasks/motion_tracking/g1/flip_tracking_sac.py +++ b/src/unilab/tasks/motion_tracking/g1/flip_tracking_sac.py @@ -17,13 +17,13 @@ from dataclasses import dataclass from unilab.base import registry -from unilab.envs.motion_tracking.g1.flip_tracking import ( + +from .flip_tracking import ( G1FlipTracking23DofEnvCfg, G1FlipTrackingCfg, G1WallFlipTracking23DofEnvCfg, G1WallFlipTrackingCfg, ) - from .tracking_sac import G1MotionTrackingSACEnv diff --git a/tests/envs/test_env_configs.py b/tests/envs/test_env_configs.py index 3e34cdba0..57b9e6ff2 100644 --- a/tests/envs/test_env_configs.py +++ b/tests/envs/test_env_configs.py @@ -1715,7 +1715,7 @@ def test_sharpa_run_complete_reports_effective_collection_target() -> None: def test_g1_flip_tracking_cfg_uses_flip_profile(): - from unilab.envs.motion_tracking.g1.flip_tracking import G1FlipTrackingCfg + from unilab.tasks.motion_tracking.g1.flip_tracking import G1FlipTrackingCfg cfg = G1FlipTrackingCfg() @@ -1731,7 +1731,7 @@ def test_g1_flip_tracking_cfg_uses_flip_profile(): def test_g1_wall_flip_tracking_cfg_uses_wall_flip_profile(): - from unilab.envs.motion_tracking.g1.flip_tracking import G1WallFlipTrackingCfg + from unilab.tasks.motion_tracking.g1.flip_tracking import G1WallFlipTrackingCfg cfg = G1WallFlipTrackingCfg() From 4ffc37394c931235f14e762a0efd6ed524248878 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 23:05:28 +0800 Subject: [PATCH 059/169] refactor(tasks): move G1 core tracking owner [skip ci] (#1164) --- .../benchmark/benchmark_drake_performance.py | 4 +-- scripts/benchmark/env/benchmark_env_step.py | 4 +-- src/unilab/envs/motion_tracking/__init__.py | 23 +----------- .../envs/motion_tracking/g1/__init__.py | 23 ------------ src/unilab/tasks/__init__.py | 1 - src/unilab/tasks/motion_tracking/__init__.py | 12 +++++++ .../tasks/motion_tracking/g1/__init__.py | 18 ++++++++++ .../tasks/motion_tracking/g1/box_tracking.py | 12 +++---- .../tasks/motion_tracking/g1/flip_tracking.py | 3 +- .../motion_tracking/g1/tracking.py | 16 +++++---- .../tasks/motion_tracking/g1/tracking_obs.py | 4 +-- .../tasks/motion_tracking/g1/tracking_sac.py | 3 +- tests/envs/test_env_configs.py | 35 +++++++++++-------- tests/tasks/test_package_boundary.py | 1 - 14 files changed, 77 insertions(+), 82 deletions(-) delete mode 100644 src/unilab/envs/motion_tracking/g1/__init__.py rename src/unilab/{envs => tasks}/motion_tracking/g1/tracking.py (91%) diff --git a/scripts/benchmark/benchmark_drake_performance.py b/scripts/benchmark/benchmark_drake_performance.py index 3b391546b..78bddfe80 100644 --- a/scripts/benchmark/benchmark_drake_performance.py +++ b/scripts/benchmark/benchmark_drake_performance.py @@ -121,12 +121,12 @@ def go2_env() -> type: return Go2WalkTask def g1_tracking_cfg() -> Any: - from unilab.envs.motion_tracking.g1.tracking import G1MotionTrackingEnvCfg + from unilab.tasks.motion_tracking.g1.tracking import G1MotionTrackingEnvCfg return G1MotionTrackingEnvCfg() def g1_tracking_env() -> type: - from unilab.envs.motion_tracking.g1.tracking import G1MotionTrackingEnv + from unilab.tasks.motion_tracking.g1.tracking import G1MotionTrackingEnv return G1MotionTrackingEnv diff --git a/scripts/benchmark/env/benchmark_env_step.py b/scripts/benchmark/env/benchmark_env_step.py index e325d8026..632ef33fc 100644 --- a/scripts/benchmark/env/benchmark_env_step.py +++ b/scripts/benchmark/env/benchmark_env_step.py @@ -350,7 +350,7 @@ def _g1_rough_cfg(backend: str, config_overrides: list[str]) -> Any: def _g1_motion_tracking_cfg(backend: str, config_overrides: list[str]) -> Any: - from unilab.envs.motion_tracking.g1.tracking import G1MotionTrackingEnvCfg + from unilab.tasks.motion_tracking.g1.tracking import G1MotionTrackingEnvCfg return _ppo_owner_yaml_cfg( "g1_motion_tracking", @@ -422,7 +422,7 @@ def _g1_walk_env_cls() -> type: def _g1_motion_tracking_env_cls() -> type: - from unilab.envs.motion_tracking.g1.tracking import G1MotionTrackingEnv + from unilab.tasks.motion_tracking.g1.tracking import G1MotionTrackingEnv return G1MotionTrackingEnv diff --git a/src/unilab/envs/motion_tracking/__init__.py b/src/unilab/envs/motion_tracking/__init__.py index 14d310250..492849929 100644 --- a/src/unilab/envs/motion_tracking/__init__.py +++ b/src/unilab/envs/motion_tracking/__init__.py @@ -1,22 +1 @@ -"""Motion tracking environments.""" - -__unilab_registry_modules__ = ("unilab.envs.motion_tracking.g1",) - -from .g1 import ( - G1MotionTracking23DofCfg, - G1MotionTracking23DofDeployEnvCfg, - G1MotionTracking23DofEnvCfg, - G1MotionTrackingCfg, - G1MotionTrackingEnv, - G1MotionTrackingEnvCfg, -) - -__all__ = [ - "G1MotionTrackingCfg", - "G1MotionTrackingEnv", - "G1MotionTrackingEnvCfg", - # 23-DoF variants (from parent UniLab/, under testing) - "G1MotionTracking23DofCfg", - "G1MotionTracking23DofDeployEnvCfg", - "G1MotionTracking23DofEnvCfg", -] +"""Shared motion-tracking runtime package.""" diff --git a/src/unilab/envs/motion_tracking/g1/__init__.py b/src/unilab/envs/motion_tracking/g1/__init__.py deleted file mode 100644 index 0df469931..000000000 --- a/src/unilab/envs/motion_tracking/g1/__init__.py +++ /dev/null @@ -1,23 +0,0 @@ -"""Motion tracking environments for Unitree G1.""" - -from .tracking import ( - G1MotionTracking23DofCfg, - G1MotionTracking23DofDeployEnvCfg, - G1MotionTracking23DofEnvCfg, - G1MotionTrackingCfg, - G1MotionTrackingDeployEnv, - G1MotionTrackingDeployEnvCfg, - G1MotionTrackingEnv, - G1MotionTrackingEnvCfg, -) - -__all__ = [ - "G1MotionTracking23DofCfg", - "G1MotionTracking23DofDeployEnvCfg", - "G1MotionTracking23DofEnvCfg", - "G1MotionTrackingCfg", - "G1MotionTrackingDeployEnv", - "G1MotionTrackingDeployEnvCfg", - "G1MotionTrackingEnv", - "G1MotionTrackingEnvCfg", -] diff --git a/src/unilab/tasks/__init__.py b/src/unilab/tasks/__init__.py index f103c90bd..6e0d95230 100644 --- a/src/unilab/tasks/__init__.py +++ b/src/unilab/tasks/__init__.py @@ -17,7 +17,6 @@ "unilab.tasks.manipulation.allegro_inhand", "unilab.tasks.manipulation.sharpa_inhand", "unilab.tasks.manipulation.stewart", - "unilab.envs.motion_tracking.g1", "unilab.tasks.motion_tracking.g1", "unilab.tasks.motion_tracking.x2", ) diff --git a/src/unilab/tasks/motion_tracking/__init__.py b/src/unilab/tasks/motion_tracking/__init__.py index 47e830b17..c121df023 100644 --- a/src/unilab/tasks/motion_tracking/__init__.py +++ b/src/unilab/tasks/motion_tracking/__init__.py @@ -20,6 +20,12 @@ G1FlipTrackingSAC23DofEnv, G1FlipTrackingSACCfg, G1FlipTrackingSACEnv, + G1MotionTracking23DofCfg, + G1MotionTracking23DofDeployEnvCfg, + G1MotionTracking23DofEnvCfg, + G1MotionTrackingCfg, + G1MotionTrackingEnv, + G1MotionTrackingEnvCfg, G1MotionTrackingSAC23DofCfg, G1MotionTrackingSAC23DofEnv, G1MotionTrackingSACCfg, @@ -64,6 +70,12 @@ "G1FlipTrackingSAC23DofEnv", "G1FlipTrackingSACCfg", "G1FlipTrackingSACEnv", + "G1MotionTracking23DofCfg", + "G1MotionTracking23DofDeployEnvCfg", + "G1MotionTracking23DofEnvCfg", + "G1MotionTrackingCfg", + "G1MotionTrackingEnv", + "G1MotionTrackingEnvCfg", "G1MotionTrackingSAC23DofCfg", "G1MotionTrackingSAC23DofEnv", "G1MotionTrackingSACCfg", diff --git a/src/unilab/tasks/motion_tracking/g1/__init__.py b/src/unilab/tasks/motion_tracking/g1/__init__.py index f85ce164a..117e74579 100644 --- a/src/unilab/tasks/motion_tracking/g1/__init__.py +++ b/src/unilab/tasks/motion_tracking/g1/__init__.py @@ -33,6 +33,16 @@ G1WallFlipTrackingSACEnv, ) from .motion_box_loader import BoxMotionData, BoxMotionLoader +from .tracking import ( + G1MotionTracking23DofCfg, + G1MotionTracking23DofDeployEnvCfg, + G1MotionTracking23DofEnvCfg, + G1MotionTrackingCfg, + G1MotionTrackingDeployEnv, + G1MotionTrackingDeployEnvCfg, + G1MotionTrackingEnv, + G1MotionTrackingEnvCfg, +) from .tracking_obs import G1WBTObs23DofCfg, G1WBTObsCfg, G1WBTObsEnv from .tracking_sac import ( G1MotionTrackingSAC23DofCfg, @@ -61,6 +71,14 @@ "G1FlipTrackingSAC23DofEnv", "G1FlipTrackingSACCfg", "G1FlipTrackingSACEnv", + "G1MotionTracking23DofCfg", + "G1MotionTracking23DofDeployEnvCfg", + "G1MotionTracking23DofEnvCfg", + "G1MotionTrackingCfg", + "G1MotionTrackingDeployEnv", + "G1MotionTrackingDeployEnvCfg", + "G1MotionTrackingEnv", + "G1MotionTrackingEnvCfg", "G1MotionTrackingSAC23DofCfg", "G1MotionTrackingSAC23DofEnv", "G1MotionTrackingSACCfg", diff --git a/src/unilab/tasks/motion_tracking/g1/box_tracking.py b/src/unilab/tasks/motion_tracking/g1/box_tracking.py index 7b3fb0bd6..3612c31f8 100644 --- a/src/unilab/tasks/motion_tracking/g1/box_tracking.py +++ b/src/unilab/tasks/motion_tracking/g1/box_tracking.py @@ -14,12 +14,6 @@ from unilab.dr.dr_utils import build_common_reset_randomization, zero_actions from unilab.dtype_config import get_global_dtype from unilab.envs.motion_tracking.common.rewards import RewardContext -from unilab.envs.motion_tracking.g1.tracking import ( - G1MotionTrackingCfg, - G1MotionTrackingDomainRandomizationProvider, - G1MotionTrackingEnv, - RewardConfig, -) from unilab.utils.geometry import np_sample_uniform from unilab.utils.rotation import ( np_matrix_from_quat, @@ -32,6 +26,12 @@ ) from .motion_box_loader import BoxMotionData, BoxMotionLoader +from .tracking import ( + G1MotionTrackingCfg, + G1MotionTrackingDomainRandomizationProvider, + G1MotionTrackingEnv, + RewardConfig, +) @dataclass diff --git a/src/unilab/tasks/motion_tracking/g1/flip_tracking.py b/src/unilab/tasks/motion_tracking/g1/flip_tracking.py index 2ce3fe4b0..c640a912c 100644 --- a/src/unilab/tasks/motion_tracking/g1/flip_tracking.py +++ b/src/unilab/tasks/motion_tracking/g1/flip_tracking.py @@ -16,7 +16,8 @@ _zero_pose_randomization, _zero_velocity_randomization, ) -from unilab.envs.motion_tracking.g1.tracking import ( + +from .tracking import ( G1MotionTrackingCfg, G1MotionTrackingEnv, PoseRandomization, diff --git a/src/unilab/envs/motion_tracking/g1/tracking.py b/src/unilab/tasks/motion_tracking/g1/tracking.py similarity index 91% rename from src/unilab/envs/motion_tracking/g1/tracking.py rename to src/unilab/tasks/motion_tracking/g1/tracking.py index c99648f34..87c90500e 100644 --- a/src/unilab/envs/motion_tracking/g1/tracking.py +++ b/src/unilab/tasks/motion_tracking/g1/tracking.py @@ -14,8 +14,7 @@ from unilab.assets import ASSETS_ROOT_PATH from unilab.base import registry from unilab.base.scene import SceneCfg - -from ..common.config import ( +from unilab.envs.motion_tracking.common.config import ( Domain_Rand, DomainRand, MotionTrackingCfg, @@ -25,10 +24,15 @@ _zero_pose_randomization, _zero_velocity_randomization, ) -from ..common.domain_randomization import MotionTrackingDomainRandomizationProvider -from ..common.reset import build_motion_reference_state -from ..common.rewards import RewardConfig -from ..common.tracking import MotionTrackingDeployEnv, MotionTrackingEnv +from unilab.envs.motion_tracking.common.domain_randomization import ( + MotionTrackingDomainRandomizationProvider, +) +from unilab.envs.motion_tracking.common.reset import build_motion_reference_state +from unilab.envs.motion_tracking.common.rewards import RewardConfig +from unilab.envs.motion_tracking.common.tracking import ( + MotionTrackingDeployEnv, + MotionTrackingEnv, +) # ── backward-compatible aliases (historical G1* symbol names) ──────── G1MotionTrackingCfg = MotionTrackingCfg diff --git a/src/unilab/tasks/motion_tracking/g1/tracking_obs.py b/src/unilab/tasks/motion_tracking/g1/tracking_obs.py index 2efe1cc77..84bf22f99 100644 --- a/src/unilab/tasks/motion_tracking/g1/tracking_obs.py +++ b/src/unilab/tasks/motion_tracking/g1/tracking_obs.py @@ -45,12 +45,12 @@ from unilab.dtype_config import get_global_dtype from unilab.envs.locomotion.g1.base import NoiseConfig from unilab.envs.motion_tracking.common.rewards import RewardContext -from unilab.envs.motion_tracking.g1.tracking import ( + +from .tracking import ( Domain_Rand, G1MotionTrackingDomainRandomizationProvider, _build_motion_reference_state, ) - from .tracking_sac import G1MotionTrackingSACCfg, G1MotionTrackingSACEnv # --------------------------------------------------------------------------- # diff --git a/src/unilab/tasks/motion_tracking/g1/tracking_sac.py b/src/unilab/tasks/motion_tracking/g1/tracking_sac.py index 78dce9ab4..1654838f4 100644 --- a/src/unilab/tasks/motion_tracking/g1/tracking_sac.py +++ b/src/unilab/tasks/motion_tracking/g1/tracking_sac.py @@ -15,7 +15,8 @@ from unilab.base import registry from unilab.dtype_config import get_global_dtype -from unilab.envs.motion_tracking.g1.tracking import ( + +from .tracking import ( G1MotionTracking23DofCfg, G1MotionTrackingCfg, G1MotionTrackingEnv, diff --git a/tests/envs/test_env_configs.py b/tests/envs/test_env_configs.py index 57b9e6ff2..b909ea7e0 100644 --- a/tests/envs/test_env_configs.py +++ b/tests/envs/test_env_configs.py @@ -53,7 +53,7 @@ def blocked_import(name, globals=None, locals=None, fromlist=(), level=0): from unilab.base import registry from unilab.base.backend import create_backend from unilab.tasks.manipulation.allegro_inhand.rotation import AllegroRotationCfg - from unilab.envs.motion_tracking.g1.tracking import ( + from unilab.tasks.motion_tracking.g1.tracking import ( G1MotionTrackingCfg, G1MotionTrackingDeployEnvCfg, ) @@ -596,7 +596,7 @@ def test_allegro_reset_samples_materialized_cache_without_file_io( def test_g1_motion_tracking_uses_combined_body_pose_query(): """G1MotionTracking should query pos/quat via the stable combined backend API.""" - from unilab.envs.motion_tracking.g1.tracking import G1MotionTrackingEnv + from unilab.tasks.motion_tracking.g1.tracking import G1MotionTrackingEnv class FakeBackend: def __init__(self) -> None: @@ -620,7 +620,9 @@ def get_body_pose_w(self, body_ids: np.ndarray) -> tuple[np.ndarray, np.ndarray] def test_g1_motion_tracking_reset_observation_uses_sparse_body_pose_rows(): from unilab.envs.motion_tracking.common.motion_loader import MotionData - from unilab.envs.motion_tracking.g1.tracking import G1MotionTrackingDomainRandomizationProvider + from unilab.tasks.motion_tracking.g1.tracking import ( + G1MotionTrackingDomainRandomizationProvider, + ) class FakeBackend: def __init__(self) -> None: @@ -758,7 +760,7 @@ def _compute_g1_motion_tracking_obs_stub(env_cls: type): def test_g1_motion_tracking_critic_uses_clean_beyondmimic_aligned_terms(): - from unilab.envs.motion_tracking.g1.tracking import G1MotionTrackingEnv + from unilab.tasks.motion_tracking.g1.tracking import G1MotionTrackingEnv env, obs, motion_data, linvel, gyro, dof_pos, dof_vel, info = ( _compute_g1_motion_tracking_obs_stub(G1MotionTrackingEnv) @@ -795,7 +797,7 @@ def test_g1_motion_tracking_critic_uses_clean_beyondmimic_aligned_terms(): def test_g1_motion_tracking_anchor_frame_writers_match_reference(): - from unilab.envs.motion_tracking.g1.tracking import G1MotionTrackingEnv + from unilab.tasks.motion_tracking.g1.tracking import G1MotionTrackingEnv from unilab.utils.rotation import ( np_matrix_from_quat, np_quat_apply, @@ -845,7 +847,7 @@ def random_quat(shape: tuple[int, ...]) -> np.ndarray: def test_g1_motion_tracking_relative_transform_fast_path_matches_reference(): - from unilab.envs.motion_tracking.g1.tracking import G1MotionTrackingEnv + from unilab.tasks.motion_tracking.g1.tracking import G1MotionTrackingEnv from unilab.utils.rotation import np_quat_apply, np_quat_inv, np_quat_mul, np_yaw_quat rng = np.random.default_rng(321) @@ -903,7 +905,7 @@ def random_quat(shape: tuple[int, ...]) -> np.ndarray: def test_g1_motion_tracking_reward_fast_path_matches_reference(): from unilab.envs.motion_tracking.common.motion_loader import MotionData - from unilab.envs.motion_tracking.g1.tracking import G1MotionTrackingEnv, RewardConfig + from unilab.tasks.motion_tracking.g1.tracking import G1MotionTrackingEnv, RewardConfig from unilab.utils.rotation import np_quat_error_magnitude rng = np.random.default_rng(456) @@ -1071,7 +1073,7 @@ def random_quat(shape: tuple[int, ...]) -> np.ndarray: def test_g1_motion_tracking_deploy_actor_matches_unitree_mimic_terms(): - from unilab.envs.motion_tracking.g1.tracking import G1MotionTrackingDeployEnv + from unilab.tasks.motion_tracking.g1.tracking import G1MotionTrackingDeployEnv env, obs, _motion_data, _linvel, gyro, dof_pos, dof_vel, info = ( _compute_g1_motion_tracking_obs_stub(G1MotionTrackingDeployEnv) @@ -1321,7 +1323,7 @@ def get_body_lin_vel_w(self, body_ids: np.ndarray) -> np.ndarray: def test_g1_motion_tracking_can_terminate_on_undesired_contacts(): - from unilab.envs.motion_tracking.g1.tracking import G1MotionTrackingEnv + from unilab.tasks.motion_tracking.g1.tracking import G1MotionTrackingEnv env = cast(Any, object.__new__(G1MotionTrackingEnv)) env._num_envs = 2 @@ -1365,7 +1367,7 @@ def test_g1_motion_tracking_can_terminate_on_undesired_contacts(): def test_g1_motion_tracking_cfg_has_domain_rand_for_motrix(): - from unilab.envs.motion_tracking.g1.tracking import G1MotionTrackingCfg + from unilab.tasks.motion_tracking.g1.tracking import G1MotionTrackingCfg cfg = G1MotionTrackingCfg() assert hasattr(cfg, "domain_rand") @@ -1376,7 +1378,7 @@ def test_g1_motion_tracking_cfg_has_domain_rand_for_motrix(): def test_g1_motion_tracking_cfg_preserves_legacy_defaults(): - from unilab.envs.motion_tracking.g1.tracking import G1MotionTrackingCfg + from unilab.tasks.motion_tracking.g1.tracking import G1MotionTrackingCfg cfg = G1MotionTrackingCfg() @@ -1392,7 +1394,10 @@ def test_g1_motion_tracking_cfg_preserves_legacy_defaults(): def test_g1_motion_tracking_init_delegates_motion_body_ids_to_backend(monkeypatch): from unilab.envs.locomotion.g1.base import G1BaseEnv from unilab.envs.motion_tracking.common import tracking as tracking_module - from unilab.envs.motion_tracking.g1.tracking import G1MotionTrackingCfg, G1MotionTrackingEnv + from unilab.tasks.motion_tracking.g1.tracking import ( + G1MotionTrackingCfg, + G1MotionTrackingEnv, + ) calls: dict[str, Any] = {} @@ -1773,7 +1778,7 @@ def test_x2_wall_flip_tracking_cfg_uses_x2_wall_flip_profile(): def test_g1_motion_tracking_apply_action_accepts_per_joint_action_scale(): - from unilab.envs.motion_tracking.g1.tracking import G1MotionTrackingEnv + from unilab.tasks.motion_tracking.g1.tracking import G1MotionTrackingEnv env = cast(Any, object.__new__(G1MotionTrackingEnv)) env.default_angles = np.array([0.5, -0.5, 1.0], dtype=np.float32) @@ -1800,7 +1805,7 @@ def _make_g1_motion_tracking_clip_end_stub( ): from unilab.base.np_env import NpEnvState from unilab.envs.motion_tracking.common.motion_loader import MotionData - from unilab.envs.motion_tracking.g1.tracking import G1MotionTrackingEnv + from unilab.tasks.motion_tracking.g1.tracking import G1MotionTrackingEnv class FakeBackend: def __init__(self) -> None: @@ -2022,7 +2027,7 @@ def test_g1_motion_tracking_clip_end_resample_keeps_terminated_final_obs_valid() def test_g1_motion_tracking_clip_end_does_not_override_true_termination(): from unilab.base.np_env import NpEnvState - from unilab.envs.motion_tracking.g1.tracking import G1MotionTrackingEnv + from unilab.tasks.motion_tracking.g1.tracking import G1MotionTrackingEnv env = cast(Any, object.__new__(G1MotionTrackingEnv)) env._num_envs = 2 diff --git a/tests/tasks/test_package_boundary.py b/tests/tasks/test_package_boundary.py index 296d4bfc5..c75ba15c8 100644 --- a/tests/tasks/test_package_boundary.py +++ b/tests/tasks/test_package_boundary.py @@ -21,7 +21,6 @@ "unilab.tasks.manipulation.allegro_inhand", "unilab.tasks.manipulation.sharpa_inhand", "unilab.tasks.manipulation.stewart", - "unilab.envs.motion_tracking.g1", "unilab.tasks.motion_tracking.g1", "unilab.tasks.motion_tracking.x2", ) From f5222e86ee334c3b3ae0e9da4110005c9dedcf80 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 23:13:31 +0800 Subject: [PATCH 060/169] refactor(tasks): move motion tracking orchestration [skip ci] (#1166) --- .../tasks/motion_tracking/common/__init__.py | 1 + .../motion_tracking/common/config.py | 7 +++---- .../common/domain_randomization.py | 3 +-- .../motion_tracking/common/tracking.py | 20 +++++++++++-------- .../tasks/motion_tracking/g1/flip_tracking.py | 4 ++-- .../tasks/motion_tracking/g1/tracking.py | 15 +++++++------- .../tasks/motion_tracking/x2/flip_tracking.py | 5 +++-- tests/envs/test_env_configs.py | 2 +- 8 files changed, 31 insertions(+), 26 deletions(-) create mode 100644 src/unilab/tasks/motion_tracking/common/__init__.py rename src/unilab/{envs => tasks}/motion_tracking/common/config.py (95%) rename src/unilab/{envs => tasks}/motion_tracking/common/domain_randomization.py (98%) rename src/unilab/{envs => tasks}/motion_tracking/common/tracking.py (96%) diff --git a/src/unilab/tasks/motion_tracking/common/__init__.py b/src/unilab/tasks/motion_tracking/common/__init__.py new file mode 100644 index 000000000..9e4fc445c --- /dev/null +++ b/src/unilab/tasks/motion_tracking/common/__init__.py @@ -0,0 +1 @@ +"""Shared motion-tracking task engine.""" diff --git a/src/unilab/envs/motion_tracking/common/config.py b/src/unilab/tasks/motion_tracking/common/config.py similarity index 95% rename from src/unilab/envs/motion_tracking/common/config.py rename to src/unilab/tasks/motion_tracking/common/config.py index 3845caeba..d03141cac 100644 --- a/src/unilab/envs/motion_tracking/common/config.py +++ b/src/unilab/tasks/motion_tracking/common/config.py @@ -14,8 +14,7 @@ from unilab.assets import ASSETS_ROOT_PATH from unilab.base.scene import SceneCfg from unilab.envs.locomotion.g1.base import G1BaseCfg - -from .rewards import RewardConfig +from unilab.envs.motion_tracking.common.rewards import RewardConfig @dataclass @@ -108,7 +107,7 @@ def _zero_velocity_randomization() -> VelocityRandomization: class MotionTrackingCfg(G1BaseCfg): """Configuration for the motion tracking environment.""" - scene: SceneCfg = field( + scene: SceneCfg = field( # pyright: ignore[reportIncompatibleVariableOverride] default_factory=lambda: SceneCfg( model_file=str(ASSETS_ROOT_PATH / "robots" / "g1" / "scene_flat.xml") ) @@ -141,7 +140,7 @@ class MotionTrackingCfg(G1BaseCfg): sampling_mode: Literal["start", "clip_start", "uniform", "adaptive", "mixed"] = "adaptive" sampling_start_ratio: float = 0.0 truncate_on_clip_end: bool = False - max_episode_seconds: float = 10.0 + max_episode_seconds: float = 10.0 # pyright: ignore[reportIncompatibleVariableOverride] reward_config: RewardConfig = field(default_factory=RewardConfig) pose_randomization: PoseRandomization = field(default_factory=PoseRandomization) velocity_randomization: VelocityRandomization = field(default_factory=VelocityRandomization) diff --git a/src/unilab/envs/motion_tracking/common/domain_randomization.py b/src/unilab/tasks/motion_tracking/common/domain_randomization.py similarity index 98% rename from src/unilab/envs/motion_tracking/common/domain_randomization.py rename to src/unilab/tasks/motion_tracking/common/domain_randomization.py index aa74e50e8..435ad56c3 100644 --- a/src/unilab/envs/motion_tracking/common/domain_randomization.py +++ b/src/unilab/tasks/motion_tracking/common/domain_randomization.py @@ -22,8 +22,7 @@ ) from unilab.dr.types import RESET_TERM_GEOM_FRICTION, ResetRandomizationPayload from unilab.dtype_config import get_global_dtype - -from .reset import build_motion_reference_state +from unilab.envs.motion_tracking.common.reset import build_motion_reference_state class MotionTrackingDomainRandomizationProvider(DomainRandomizationProvider): diff --git a/src/unilab/envs/motion_tracking/common/tracking.py b/src/unilab/tasks/motion_tracking/common/tracking.py similarity index 96% rename from src/unilab/envs/motion_tracking/common/tracking.py rename to src/unilab/tasks/motion_tracking/common/tracking.py index 9269d224a..02eda301a 100644 --- a/src/unilab/envs/motion_tracking/common/tracking.py +++ b/src/unilab/tasks/motion_tracking/common/tracking.py @@ -18,21 +18,25 @@ from unilab.base.np_env import NpEnvState from unilab.dtype_config import get_global_dtype from unilab.envs.locomotion.g1.base import G1BaseEnv +from unilab.envs.motion_tracking.common import observations +from unilab.envs.motion_tracking.common.motion_loader import MotionData, MotionLoader, MotionSampler +from unilab.envs.motion_tracking.common.reset import build_motion_reference_state +from unilab.envs.motion_tracking.common.rewards import ( + RewardContext, + build_reward_functions, + compute_reward, +) +from unilab.envs.motion_tracking.common.terminations import compute_terminations +from unilab.envs.motion_tracking.common.transforms import update_relative_transforms -from . import observations from .config import MotionTrackingCfg, MotionTrackingDeployEnvCfg from .domain_randomization import MotionTrackingDomainRandomizationProvider -from .motion_loader import MotionData, MotionLoader, MotionSampler -from .reset import build_motion_reference_state -from .rewards import RewardContext, build_reward_functions, compute_reward -from .terminations import compute_terminations -from .transforms import update_relative_transforms class MotionTrackingEnv(G1BaseEnv): """Motion Tracking Environment (robot-agnostic imitation engine).""" - _cfg: MotionTrackingCfg + _cfg: MotionTrackingCfg # pyright: ignore[reportIncompatibleVariableOverride] def __init__(self, cfg: MotionTrackingCfg, num_envs=1, backend_type="mujoco"): if not cfg.motion_file: @@ -508,7 +512,7 @@ def _compute_reward( class MotionTrackingDeployEnv(MotionTrackingEnv): """Deploy-oriented motion tracking env with unitree_rl_lab mimic actor inputs.""" - _cfg: MotionTrackingDeployEnvCfg + _cfg: MotionTrackingDeployEnvCfg # pyright: ignore[reportIncompatibleVariableOverride] def _actor_obs_dim(self, n: int) -> int: return observations.mimic_actor_obs_dim(n) diff --git a/src/unilab/tasks/motion_tracking/g1/flip_tracking.py b/src/unilab/tasks/motion_tracking/g1/flip_tracking.py index c640a912c..2bda6a55b 100644 --- a/src/unilab/tasks/motion_tracking/g1/flip_tracking.py +++ b/src/unilab/tasks/motion_tracking/g1/flip_tracking.py @@ -12,11 +12,11 @@ from unilab.assets import ASSETS_ROOT_PATH from unilab.base import registry from unilab.base.scene import SceneCfg -from unilab.envs.motion_tracking.common.config import ( + +from ..common.config import ( _zero_pose_randomization, _zero_velocity_randomization, ) - from .tracking import ( G1MotionTrackingCfg, G1MotionTrackingEnv, diff --git a/src/unilab/tasks/motion_tracking/g1/tracking.py b/src/unilab/tasks/motion_tracking/g1/tracking.py index 87c90500e..f68b2e582 100644 --- a/src/unilab/tasks/motion_tracking/g1/tracking.py +++ b/src/unilab/tasks/motion_tracking/g1/tracking.py @@ -1,7 +1,7 @@ """G1 Motion Tracking profiles — thin registry subclasses over the shared engine. -The robot-agnostic engine and owner modules now live in -:mod:`unilab.envs.motion_tracking.common`. This module keeps the G1 registry +The robot-agnostic engine and owner modules live in +:mod:`unilab.tasks.motion_tracking.common`. This module keeps the G1 registry entries (``G1MotionTracking`` / ``G1MotionTrackingDeploy``) and re-exports the historical ``G1*`` / ``Domain_Rand`` / ``_build_motion_reference_state`` symbol names so existing subclasses and tests keep importing them from ``.tracking``. @@ -14,7 +14,10 @@ from unilab.assets import ASSETS_ROOT_PATH from unilab.base import registry from unilab.base.scene import SceneCfg -from unilab.envs.motion_tracking.common.config import ( +from unilab.envs.motion_tracking.common.reset import build_motion_reference_state +from unilab.envs.motion_tracking.common.rewards import RewardConfig + +from ..common.config import ( Domain_Rand, DomainRand, MotionTrackingCfg, @@ -24,12 +27,10 @@ _zero_pose_randomization, _zero_velocity_randomization, ) -from unilab.envs.motion_tracking.common.domain_randomization import ( +from ..common.domain_randomization import ( MotionTrackingDomainRandomizationProvider, ) -from unilab.envs.motion_tracking.common.reset import build_motion_reference_state -from unilab.envs.motion_tracking.common.rewards import RewardConfig -from unilab.envs.motion_tracking.common.tracking import ( +from ..common.tracking import ( MotionTrackingDeployEnv, MotionTrackingEnv, ) diff --git a/src/unilab/tasks/motion_tracking/x2/flip_tracking.py b/src/unilab/tasks/motion_tracking/x2/flip_tracking.py index 367279177..660d9b16a 100644 --- a/src/unilab/tasks/motion_tracking/x2/flip_tracking.py +++ b/src/unilab/tasks/motion_tracking/x2/flip_tracking.py @@ -9,13 +9,14 @@ from unilab.base import registry from unilab.base.scene import SceneCfg from unilab.envs.locomotion.g1.base import Sensor -from unilab.envs.motion_tracking.common.config import ( + +from ..common.config import ( PoseRandomization, VelocityRandomization, _zero_pose_randomization, _zero_velocity_randomization, ) -from unilab.envs.motion_tracking.common.tracking import ( +from ..common.tracking import ( MotionTrackingDeployEnv, MotionTrackingDeployEnvCfg, ) diff --git a/tests/envs/test_env_configs.py b/tests/envs/test_env_configs.py index b909ea7e0..61f0eb801 100644 --- a/tests/envs/test_env_configs.py +++ b/tests/envs/test_env_configs.py @@ -1393,7 +1393,7 @@ def test_g1_motion_tracking_cfg_preserves_legacy_defaults(): def test_g1_motion_tracking_init_delegates_motion_body_ids_to_backend(monkeypatch): from unilab.envs.locomotion.g1.base import G1BaseEnv - from unilab.envs.motion_tracking.common import tracking as tracking_module + from unilab.tasks.motion_tracking.common import tracking as tracking_module from unilab.tasks.motion_tracking.g1.tracking import ( G1MotionTrackingCfg, G1MotionTrackingEnv, From 7c6fea0aea0830dc28470c591a71c07d33571f68 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 23:21:10 +0800 Subject: [PATCH 061/169] refactor(tasks): move motion loader and reset owners [skip ci] (#1168) --- .../tasks/motion_tracking/common/domain_randomization.py | 3 ++- .../motion_tracking/common/motion_loader.py | 2 +- .../{envs => tasks}/motion_tracking/common/reset.py | 2 +- src/unilab/tasks/motion_tracking/common/tracking.py | 4 ++-- src/unilab/tasks/motion_tracking/g1/motion_box_loader.py | 2 +- src/unilab/tasks/motion_tracking/g1/tracking.py | 2 +- tests/envs/test_env_configs.py | 8 ++++---- tests/envs/test_motion_loader.py | 2 +- 8 files changed, 13 insertions(+), 12 deletions(-) rename src/unilab/{envs => tasks}/motion_tracking/common/motion_loader.py (99%) rename src/unilab/{envs => tasks}/motion_tracking/common/reset.py (97%) diff --git a/src/unilab/tasks/motion_tracking/common/domain_randomization.py b/src/unilab/tasks/motion_tracking/common/domain_randomization.py index 435ad56c3..aa74e50e8 100644 --- a/src/unilab/tasks/motion_tracking/common/domain_randomization.py +++ b/src/unilab/tasks/motion_tracking/common/domain_randomization.py @@ -22,7 +22,8 @@ ) from unilab.dr.types import RESET_TERM_GEOM_FRICTION, ResetRandomizationPayload from unilab.dtype_config import get_global_dtype -from unilab.envs.motion_tracking.common.reset import build_motion_reference_state + +from .reset import build_motion_reference_state class MotionTrackingDomainRandomizationProvider(DomainRandomizationProvider): diff --git a/src/unilab/envs/motion_tracking/common/motion_loader.py b/src/unilab/tasks/motion_tracking/common/motion_loader.py similarity index 99% rename from src/unilab/envs/motion_tracking/common/motion_loader.py rename to src/unilab/tasks/motion_tracking/common/motion_loader.py index 4bc7822be..b11b497aa 100644 --- a/src/unilab/envs/motion_tracking/common/motion_loader.py +++ b/src/unilab/tasks/motion_tracking/common/motion_loader.py @@ -1,4 +1,4 @@ -"""Motion loading and sampling for motion tracking tasks.""" +"""Shared motion loading and sampling for motion-tracking tasks.""" from __future__ import annotations diff --git a/src/unilab/envs/motion_tracking/common/reset.py b/src/unilab/tasks/motion_tracking/common/reset.py similarity index 97% rename from src/unilab/envs/motion_tracking/common/reset.py rename to src/unilab/tasks/motion_tracking/common/reset.py index bea6b99e2..4959bfc47 100644 --- a/src/unilab/envs/motion_tracking/common/reset.py +++ b/src/unilab/tasks/motion_tracking/common/reset.py @@ -1,4 +1,4 @@ -"""Reset-state construction for motion tracking.""" +"""Shared reset-state construction for motion tracking.""" from __future__ import annotations diff --git a/src/unilab/tasks/motion_tracking/common/tracking.py b/src/unilab/tasks/motion_tracking/common/tracking.py index 02eda301a..4603d37f7 100644 --- a/src/unilab/tasks/motion_tracking/common/tracking.py +++ b/src/unilab/tasks/motion_tracking/common/tracking.py @@ -19,8 +19,6 @@ from unilab.dtype_config import get_global_dtype from unilab.envs.locomotion.g1.base import G1BaseEnv from unilab.envs.motion_tracking.common import observations -from unilab.envs.motion_tracking.common.motion_loader import MotionData, MotionLoader, MotionSampler -from unilab.envs.motion_tracking.common.reset import build_motion_reference_state from unilab.envs.motion_tracking.common.rewards import ( RewardContext, build_reward_functions, @@ -31,6 +29,8 @@ from .config import MotionTrackingCfg, MotionTrackingDeployEnvCfg from .domain_randomization import MotionTrackingDomainRandomizationProvider +from .motion_loader import MotionData, MotionLoader, MotionSampler +from .reset import build_motion_reference_state class MotionTrackingEnv(G1BaseEnv): diff --git a/src/unilab/tasks/motion_tracking/g1/motion_box_loader.py b/src/unilab/tasks/motion_tracking/g1/motion_box_loader.py index 0ef5e421d..795f13274 100644 --- a/src/unilab/tasks/motion_tracking/g1/motion_box_loader.py +++ b/src/unilab/tasks/motion_tracking/g1/motion_box_loader.py @@ -6,7 +6,7 @@ import numpy as np -from unilab.envs.motion_tracking.common.motion_loader import MotionData, MotionLoader +from ..common.motion_loader import MotionData, MotionLoader @dataclass diff --git a/src/unilab/tasks/motion_tracking/g1/tracking.py b/src/unilab/tasks/motion_tracking/g1/tracking.py index f68b2e582..287a91014 100644 --- a/src/unilab/tasks/motion_tracking/g1/tracking.py +++ b/src/unilab/tasks/motion_tracking/g1/tracking.py @@ -14,7 +14,6 @@ from unilab.assets import ASSETS_ROOT_PATH from unilab.base import registry from unilab.base.scene import SceneCfg -from unilab.envs.motion_tracking.common.reset import build_motion_reference_state from unilab.envs.motion_tracking.common.rewards import RewardConfig from ..common.config import ( @@ -30,6 +29,7 @@ from ..common.domain_randomization import ( MotionTrackingDomainRandomizationProvider, ) +from ..common.reset import build_motion_reference_state from ..common.tracking import ( MotionTrackingDeployEnv, MotionTrackingEnv, diff --git a/tests/envs/test_env_configs.py b/tests/envs/test_env_configs.py index 61f0eb801..bbc743234 100644 --- a/tests/envs/test_env_configs.py +++ b/tests/envs/test_env_configs.py @@ -619,7 +619,7 @@ def get_body_pose_w(self, body_ids: np.ndarray) -> tuple[np.ndarray, np.ndarray] def test_g1_motion_tracking_reset_observation_uses_sparse_body_pose_rows(): - from unilab.envs.motion_tracking.common.motion_loader import MotionData + from unilab.tasks.motion_tracking.common.motion_loader import MotionData from unilab.tasks.motion_tracking.g1.tracking import ( G1MotionTrackingDomainRandomizationProvider, ) @@ -700,7 +700,7 @@ def compute_obs( def _compute_g1_motion_tracking_obs_stub(env_cls: type): - from unilab.envs.motion_tracking.common.motion_loader import MotionData + from unilab.tasks.motion_tracking.common.motion_loader import MotionData env = cast(Any, object.__new__(env_cls)) env._num_envs = 1 @@ -904,7 +904,7 @@ def random_quat(shape: tuple[int, ...]) -> np.ndarray: def test_g1_motion_tracking_reward_fast_path_matches_reference(): - from unilab.envs.motion_tracking.common.motion_loader import MotionData + from unilab.tasks.motion_tracking.common.motion_loader import MotionData from unilab.tasks.motion_tracking.g1.tracking import G1MotionTrackingEnv, RewardConfig from unilab.utils.rotation import np_quat_error_magnitude @@ -1804,7 +1804,7 @@ def _make_g1_motion_tracking_clip_end_stub( step_env_ids: np.ndarray | None = None, ): from unilab.base.np_env import NpEnvState - from unilab.envs.motion_tracking.common.motion_loader import MotionData + from unilab.tasks.motion_tracking.common.motion_loader import MotionData from unilab.tasks.motion_tracking.g1.tracking import G1MotionTrackingEnv class FakeBackend: diff --git a/tests/envs/test_motion_loader.py b/tests/envs/test_motion_loader.py index e717fe99f..b7961cc1b 100644 --- a/tests/envs/test_motion_loader.py +++ b/tests/envs/test_motion_loader.py @@ -2,7 +2,7 @@ import numpy as np -from unilab.envs.motion_tracking.common.motion_loader import MotionLoader, MotionSampler +from unilab.tasks.motion_tracking.common.motion_loader import MotionLoader, MotionSampler def _write_motion_npz( From e816ae936bc7194080da7ddeab3f5ba28d96b389 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 23:29:01 +0800 Subject: [PATCH 062/169] refactor(tasks): finish motion tracking common owner [skip ci] (#1170) --- src/unilab/envs/motion_tracking/__init__.py | 1 - .../envs/motion_tracking/common/__init__.py | 8 -------- .../tasks/motion_tracking/common/__init__.py | 8 +++++++- .../tasks/motion_tracking/common/config.py | 3 ++- .../motion_tracking/common/observations.py | 2 +- .../motion_tracking/common/rewards.py | 2 +- .../motion_tracking/common/terminations.py | 2 +- .../tasks/motion_tracking/common/tracking.py | 16 ++++++++-------- .../motion_tracking/common/transforms.py | 2 +- .../tasks/motion_tracking/g1/box_tracking.py | 2 +- src/unilab/tasks/motion_tracking/g1/tracking.py | 2 +- .../tasks/motion_tracking/g1/tracking_obs.py | 2 +- tests/envs/test_motion_tracking_rewards.py | 4 ++-- 13 files changed, 26 insertions(+), 28 deletions(-) delete mode 100644 src/unilab/envs/motion_tracking/__init__.py delete mode 100644 src/unilab/envs/motion_tracking/common/__init__.py rename src/unilab/{envs => tasks}/motion_tracking/common/observations.py (99%) rename src/unilab/{envs => tasks}/motion_tracking/common/rewards.py (99%) rename src/unilab/{envs => tasks}/motion_tracking/common/terminations.py (97%) rename src/unilab/{envs => tasks}/motion_tracking/common/transforms.py (98%) diff --git a/src/unilab/envs/motion_tracking/__init__.py b/src/unilab/envs/motion_tracking/__init__.py deleted file mode 100644 index 492849929..000000000 --- a/src/unilab/envs/motion_tracking/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Shared motion-tracking runtime package.""" diff --git a/src/unilab/envs/motion_tracking/common/__init__.py b/src/unilab/envs/motion_tracking/common/__init__.py deleted file mode 100644 index 670d69ccb..000000000 --- a/src/unilab/envs/motion_tracking/common/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -"""Robot-agnostic motion-tracking engine and owner modules. - -This package holds the motion-tracking task engine (:class:`MotionTrackingEnv` -/ :class:`MotionTrackingDeployEnv`) and its per-concern owner modules (config, -rewards, observations, terminations, transforms, reset, domain randomization, -motion loading). Per-robot profiles live under ``g1/`` and -``x2/`` and only carry robot-specific defaults and thin registry subclasses. -""" diff --git a/src/unilab/tasks/motion_tracking/common/__init__.py b/src/unilab/tasks/motion_tracking/common/__init__.py index 9e4fc445c..941751b15 100644 --- a/src/unilab/tasks/motion_tracking/common/__init__.py +++ b/src/unilab/tasks/motion_tracking/common/__init__.py @@ -1 +1,7 @@ -"""Shared motion-tracking task engine.""" +"""Shared motion-tracking task engine and owner modules. + +This package holds the motion-tracking task engine and its per-concern owner +modules: config, rewards, observations, terminations, transforms, reset, +domain randomization, and motion loading. Per-robot profiles live under +``g1/`` and ``x2/`` and contain robot-specific defaults and registry leaves. +""" diff --git a/src/unilab/tasks/motion_tracking/common/config.py b/src/unilab/tasks/motion_tracking/common/config.py index d03141cac..c76c1569b 100644 --- a/src/unilab/tasks/motion_tracking/common/config.py +++ b/src/unilab/tasks/motion_tracking/common/config.py @@ -14,7 +14,8 @@ from unilab.assets import ASSETS_ROOT_PATH from unilab.base.scene import SceneCfg from unilab.envs.locomotion.g1.base import G1BaseCfg -from unilab.envs.motion_tracking.common.rewards import RewardConfig + +from .rewards import RewardConfig @dataclass diff --git a/src/unilab/envs/motion_tracking/common/observations.py b/src/unilab/tasks/motion_tracking/common/observations.py similarity index 99% rename from src/unilab/envs/motion_tracking/common/observations.py rename to src/unilab/tasks/motion_tracking/common/observations.py index fd5b0b6ee..ba6f6203f 100644 --- a/src/unilab/envs/motion_tracking/common/observations.py +++ b/src/unilab/tasks/motion_tracking/common/observations.py @@ -1,4 +1,4 @@ -"""Observation construction for motion tracking. +"""Shared observation construction for motion tracking. Holds the robot-agnostic observation builders. The environment classes keep a thin polymorphic method surface (``_compute_obs`` / ``_build_actor_obs`` / diff --git a/src/unilab/envs/motion_tracking/common/rewards.py b/src/unilab/tasks/motion_tracking/common/rewards.py similarity index 99% rename from src/unilab/envs/motion_tracking/common/rewards.py rename to src/unilab/tasks/motion_tracking/common/rewards.py index 20ff092c6..ffe59a191 100644 --- a/src/unilab/envs/motion_tracking/common/rewards.py +++ b/src/unilab/tasks/motion_tracking/common/rewards.py @@ -1,4 +1,4 @@ -"""Reward configuration and reward functions for motion tracking. +"""Shared reward configuration and functions for motion tracking. Reward terms are plain module-level callables ``fn(ctx: RewardContext) -> np.ndarray`` mirroring :mod:`unilab.envs.locomotion.common.rewards`. Robot-specific terms that diff --git a/src/unilab/envs/motion_tracking/common/terminations.py b/src/unilab/tasks/motion_tracking/common/terminations.py similarity index 97% rename from src/unilab/envs/motion_tracking/common/terminations.py rename to src/unilab/tasks/motion_tracking/common/terminations.py index a09bfd90c..6a9ece326 100644 --- a/src/unilab/envs/motion_tracking/common/terminations.py +++ b/src/unilab/tasks/motion_tracking/common/terminations.py @@ -1,4 +1,4 @@ -"""Termination computation for motion tracking.""" +"""Shared termination computation for motion tracking.""" from __future__ import annotations diff --git a/src/unilab/tasks/motion_tracking/common/tracking.py b/src/unilab/tasks/motion_tracking/common/tracking.py index 4603d37f7..0dc66ee1b 100644 --- a/src/unilab/tasks/motion_tracking/common/tracking.py +++ b/src/unilab/tasks/motion_tracking/common/tracking.py @@ -18,19 +18,19 @@ from unilab.base.np_env import NpEnvState from unilab.dtype_config import get_global_dtype from unilab.envs.locomotion.g1.base import G1BaseEnv -from unilab.envs.motion_tracking.common import observations -from unilab.envs.motion_tracking.common.rewards import ( - RewardContext, - build_reward_functions, - compute_reward, -) -from unilab.envs.motion_tracking.common.terminations import compute_terminations -from unilab.envs.motion_tracking.common.transforms import update_relative_transforms +from . import observations from .config import MotionTrackingCfg, MotionTrackingDeployEnvCfg from .domain_randomization import MotionTrackingDomainRandomizationProvider from .motion_loader import MotionData, MotionLoader, MotionSampler from .reset import build_motion_reference_state +from .rewards import ( + RewardContext, + build_reward_functions, + compute_reward, +) +from .terminations import compute_terminations +from .transforms import update_relative_transforms class MotionTrackingEnv(G1BaseEnv): diff --git a/src/unilab/envs/motion_tracking/common/transforms.py b/src/unilab/tasks/motion_tracking/common/transforms.py similarity index 98% rename from src/unilab/envs/motion_tracking/common/transforms.py rename to src/unilab/tasks/motion_tracking/common/transforms.py index a4b8ae5ef..31b4f31a6 100644 --- a/src/unilab/envs/motion_tracking/common/transforms.py +++ b/src/unilab/tasks/motion_tracking/common/transforms.py @@ -1,4 +1,4 @@ -"""Relative body-transform computation for motion tracking. +"""Shared relative body-transform computation for motion tracking. Fills the environment's ``body_pos_relative_w`` / ``body_quat_relative_w`` reference buffers each step. The op order and in-place ``out=`` usage are diff --git a/src/unilab/tasks/motion_tracking/g1/box_tracking.py b/src/unilab/tasks/motion_tracking/g1/box_tracking.py index 3612c31f8..64ef22763 100644 --- a/src/unilab/tasks/motion_tracking/g1/box_tracking.py +++ b/src/unilab/tasks/motion_tracking/g1/box_tracking.py @@ -13,7 +13,6 @@ from unilab.dr import DomainRandomizationManager, ResetPlan from unilab.dr.dr_utils import build_common_reset_randomization, zero_actions from unilab.dtype_config import get_global_dtype -from unilab.envs.motion_tracking.common.rewards import RewardContext from unilab.utils.geometry import np_sample_uniform from unilab.utils.rotation import ( np_matrix_from_quat, @@ -25,6 +24,7 @@ np_subtract_frame_transforms, ) +from ..common.rewards import RewardContext from .motion_box_loader import BoxMotionData, BoxMotionLoader from .tracking import ( G1MotionTrackingCfg, diff --git a/src/unilab/tasks/motion_tracking/g1/tracking.py b/src/unilab/tasks/motion_tracking/g1/tracking.py index 287a91014..58cb5d7a4 100644 --- a/src/unilab/tasks/motion_tracking/g1/tracking.py +++ b/src/unilab/tasks/motion_tracking/g1/tracking.py @@ -14,7 +14,6 @@ from unilab.assets import ASSETS_ROOT_PATH from unilab.base import registry from unilab.base.scene import SceneCfg -from unilab.envs.motion_tracking.common.rewards import RewardConfig from ..common.config import ( Domain_Rand, @@ -30,6 +29,7 @@ MotionTrackingDomainRandomizationProvider, ) from ..common.reset import build_motion_reference_state +from ..common.rewards import RewardConfig from ..common.tracking import ( MotionTrackingDeployEnv, MotionTrackingEnv, diff --git a/src/unilab/tasks/motion_tracking/g1/tracking_obs.py b/src/unilab/tasks/motion_tracking/g1/tracking_obs.py index 84bf22f99..3fd78b712 100644 --- a/src/unilab/tasks/motion_tracking/g1/tracking_obs.py +++ b/src/unilab/tasks/motion_tracking/g1/tracking_obs.py @@ -44,8 +44,8 @@ from unilab.dr.types import RESET_TERM_GEOM_FRICTION from unilab.dtype_config import get_global_dtype from unilab.envs.locomotion.g1.base import NoiseConfig -from unilab.envs.motion_tracking.common.rewards import RewardContext +from ..common.rewards import RewardContext from .tracking import ( Domain_Rand, G1MotionTrackingDomainRandomizationProvider, diff --git a/tests/envs/test_motion_tracking_rewards.py b/tests/envs/test_motion_tracking_rewards.py index c7a8fe20b..5c6783cef 100644 --- a/tests/envs/test_motion_tracking_rewards.py +++ b/tests/envs/test_motion_tracking_rewards.py @@ -10,8 +10,8 @@ import numpy as np -from unilab.envs.motion_tracking.common import rewards -from unilab.envs.motion_tracking.common.rewards import RewardConfig, RewardContext +from unilab.tasks.motion_tracking.common import rewards +from unilab.tasks.motion_tracking.common.rewards import RewardConfig, RewardContext def _make_ctx(*, scales: dict[str, float] | None = None) -> RewardContext: From 7b8f110e3f26bb395663b1b429a68b03f71e7d52 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 23:37:00 +0800 Subject: [PATCH 063/169] refactor(tasks): move G1 shared base owner [skip ci] (#1172) --- src/unilab/envs/locomotion/g1/__init__.py | 1 - src/unilab/{envs => tasks}/locomotion/g1/base.py | 8 ++++++-- src/unilab/tasks/locomotion/g1/joystick.py | 3 ++- src/unilab/tasks/motion_tracking/common/config.py | 2 +- src/unilab/tasks/motion_tracking/common/tracking.py | 2 +- src/unilab/tasks/motion_tracking/g1/tracking_obs.py | 2 +- src/unilab/tasks/motion_tracking/x2/flip_tracking.py | 2 +- tests/envs/test_env_configs.py | 2 +- tests/envs/test_g1_obs_noise.py | 2 +- 9 files changed, 14 insertions(+), 10 deletions(-) delete mode 100644 src/unilab/envs/locomotion/g1/__init__.py rename src/unilab/{envs => tasks}/locomotion/g1/base.py (85%) diff --git a/src/unilab/envs/locomotion/g1/__init__.py b/src/unilab/envs/locomotion/g1/__init__.py deleted file mode 100644 index b4901c6b7..000000000 --- a/src/unilab/envs/locomotion/g1/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Legacy G1 shared base pending motion-tracking migration.""" diff --git a/src/unilab/envs/locomotion/g1/base.py b/src/unilab/tasks/locomotion/g1/base.py similarity index 85% rename from src/unilab/envs/locomotion/g1/base.py rename to src/unilab/tasks/locomotion/g1/base.py index 461013c8f..18c23fb16 100644 --- a/src/unilab/envs/locomotion/g1/base.py +++ b/src/unilab/tasks/locomotion/g1/base.py @@ -1,3 +1,5 @@ +"""Shared runtime and configuration for G1 task owners.""" + from __future__ import annotations from dataclasses import dataclass, field @@ -45,14 +47,16 @@ class Asset: class G1BaseCfg(LocomotionBaseCfg): noise_config: NoiseConfig = field(default_factory=NoiseConfig) # type: ignore[assignment] control_config: ControlConfig = field(default_factory=ControlConfig) # type: ignore[assignment] - sensor: Sensor = field(default_factory=Sensor) + sensor: Sensor = field( # pyright: ignore[reportIncompatibleVariableOverride] + default_factory=Sensor + ) asset: Asset = field(default_factory=Asset) sim_dt: float = 0.02 / 3.0 ctrl_dt: float = 0.02 class G1BaseEnv(LocomotionBaseEnv): - _cfg: G1BaseCfg + _cfg: G1BaseCfg # pyright: ignore[reportIncompatibleVariableOverride] _keyframe_name = "stand" _use_global_dtype = False diff --git a/src/unilab/tasks/locomotion/g1/joystick.py b/src/unilab/tasks/locomotion/g1/joystick.py index 902eaafde..351a50a69 100644 --- a/src/unilab/tasks/locomotion/g1/joystick.py +++ b/src/unilab/tasks/locomotion/g1/joystick.py @@ -25,7 +25,8 @@ 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 .base import G1BaseCfg, G1BaseEnv @dataclass diff --git a/src/unilab/tasks/motion_tracking/common/config.py b/src/unilab/tasks/motion_tracking/common/config.py index c76c1569b..fb9d5419e 100644 --- a/src/unilab/tasks/motion_tracking/common/config.py +++ b/src/unilab/tasks/motion_tracking/common/config.py @@ -13,7 +13,7 @@ from unilab.assets import ASSETS_ROOT_PATH from unilab.base.scene import SceneCfg -from unilab.envs.locomotion.g1.base import G1BaseCfg +from unilab.tasks.locomotion.g1.base import G1BaseCfg from .rewards import RewardConfig diff --git a/src/unilab/tasks/motion_tracking/common/tracking.py b/src/unilab/tasks/motion_tracking/common/tracking.py index 0dc66ee1b..7da3516a6 100644 --- a/src/unilab/tasks/motion_tracking/common/tracking.py +++ b/src/unilab/tasks/motion_tracking/common/tracking.py @@ -17,7 +17,7 @@ from unilab.base.backend import create_backend, env_backend_kwargs from unilab.base.np_env import NpEnvState from unilab.dtype_config import get_global_dtype -from unilab.envs.locomotion.g1.base import G1BaseEnv +from unilab.tasks.locomotion.g1.base import G1BaseEnv from . import observations from .config import MotionTrackingCfg, MotionTrackingDeployEnvCfg diff --git a/src/unilab/tasks/motion_tracking/g1/tracking_obs.py b/src/unilab/tasks/motion_tracking/g1/tracking_obs.py index 3fd78b712..ea194f3c0 100644 --- a/src/unilab/tasks/motion_tracking/g1/tracking_obs.py +++ b/src/unilab/tasks/motion_tracking/g1/tracking_obs.py @@ -43,7 +43,7 @@ ) from unilab.dr.types import RESET_TERM_GEOM_FRICTION from unilab.dtype_config import get_global_dtype -from unilab.envs.locomotion.g1.base import NoiseConfig +from unilab.tasks.locomotion.g1.base import NoiseConfig from ..common.rewards import RewardContext from .tracking import ( diff --git a/src/unilab/tasks/motion_tracking/x2/flip_tracking.py b/src/unilab/tasks/motion_tracking/x2/flip_tracking.py index 660d9b16a..8c0a5340f 100644 --- a/src/unilab/tasks/motion_tracking/x2/flip_tracking.py +++ b/src/unilab/tasks/motion_tracking/x2/flip_tracking.py @@ -8,7 +8,7 @@ from unilab.assets import ASSETS_ROOT_PATH from unilab.base import registry from unilab.base.scene import SceneCfg -from unilab.envs.locomotion.g1.base import Sensor +from unilab.tasks.locomotion.g1.base import Sensor from ..common.config import ( PoseRandomization, diff --git a/tests/envs/test_env_configs.py b/tests/envs/test_env_configs.py index bbc743234..37f7a6ed3 100644 --- a/tests/envs/test_env_configs.py +++ b/tests/envs/test_env_configs.py @@ -1392,7 +1392,7 @@ def test_g1_motion_tracking_cfg_preserves_legacy_defaults(): def test_g1_motion_tracking_init_delegates_motion_body_ids_to_backend(monkeypatch): - from unilab.envs.locomotion.g1.base import G1BaseEnv + from unilab.tasks.locomotion.g1.base import G1BaseEnv from unilab.tasks.motion_tracking.common import tracking as tracking_module from unilab.tasks.motion_tracking.g1.tracking import ( G1MotionTrackingCfg, diff --git a/tests/envs/test_g1_obs_noise.py b/tests/envs/test_g1_obs_noise.py index 48b061204..2dbb1c4c4 100644 --- a/tests/envs/test_g1_obs_noise.py +++ b/tests/envs/test_g1_obs_noise.py @@ -4,7 +4,7 @@ import numpy as np -from unilab.envs.locomotion.g1.base import G1BaseCfg, G1BaseEnv, NoiseConfig +from unilab.tasks.locomotion.g1.base import G1BaseCfg, G1BaseEnv, NoiseConfig class _ConcreteG1Env(G1BaseEnv): From 1478465be5f98c7070438729e30c0a4064bbc2ad Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 23:44:53 +0800 Subject: [PATCH 064/169] refactor(tasks): move locomotion manager terms [skip ci] (#1174) --- conf/appo/task/go2_joystick_flat/base.yaml | 20 +++++++++---------- .../task/go2_joystick_flat/base.yaml | 20 +++++++++---------- conf/ppo/task/go2_joystick_flat/base.yaml | 20 +++++++++---------- .../tasks/locomotion/common/__init__.py | 1 + .../locomotion/common/manager_terms.py | 2 +- .../locomotion/test_manager_gait_terms.py | 2 +- 6 files changed, 33 insertions(+), 32 deletions(-) create mode 100644 src/unilab/tasks/locomotion/common/__init__.py rename src/unilab/{envs => tasks}/locomotion/common/manager_terms.py (99%) diff --git a/conf/appo/task/go2_joystick_flat/base.yaml b/conf/appo/task/go2_joystick_flat/base.yaml index ad7eb2fba..288c022c7 100644 --- a/conf/appo/task/go2_joystick_flat/base.yaml +++ b/conf/appo/task/go2_joystick_flat/base.yaml @@ -67,7 +67,7 @@ env: command_name: twist gait_phase: _target_: unilab.managers.ObservationTermCfg - func: unilab.envs.locomotion.common.manager_terms.quadruped_gait_phase + func: unilab.tasks.locomotion.common.manager_terms.quadruped_gait_phase params: frequency: 2.0 critic: @@ -99,7 +99,7 @@ env: command_name: twist gait_phase: _target_: unilab.managers.ObservationTermCfg - func: unilab.envs.locomotion.common.manager_terms.quadruped_gait_phase + func: unilab.tasks.locomotion.common.manager_terms.quadruped_gait_phase params: frequency: 2.0 base_lin_vel: @@ -178,29 +178,29 @@ env: reward: tracking_lin_vel: _target_: unilab.managers.RewardTermCfg - func: unilab.envs.locomotion.common.manager_terms.track_lin_vel_xy_exp + func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel_xy_exp weight: 1.0 params: std: 0.5 command_name: twist tracking_ang_vel: _target_: unilab.managers.RewardTermCfg - func: unilab.envs.locomotion.common.manager_terms.track_ang_vel_z_exp + func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel_z_exp weight: 0.2 params: std: 0.5 command_name: twist lin_vel_z: _target_: unilab.managers.RewardTermCfg - func: unilab.envs.locomotion.common.manager_terms.lin_vel_z_l2 + func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z_l2 weight: -5.0 ang_vel_xy: _target_: unilab.managers.RewardTermCfg - func: unilab.envs.locomotion.common.manager_terms.ang_vel_xy_l2 + func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy_l2 weight: -0.1 base_height: _target_: unilab.managers.RewardTermCfg - func: unilab.envs.locomotion.common.manager_terms.base_height_l2 + func: unilab.tasks.locomotion.common.manager_terms.base_height_l2 weight: -100.0 params: target_height: 0.3 @@ -210,7 +210,7 @@ reward: weight: -0.005 similar_to_default: _target_: unilab.managers.RewardTermCfg - func: unilab.envs.locomotion.common.manager_terms.joint_deviation_l1 + func: unilab.tasks.locomotion.common.manager_terms.joint_deviation_l1 weight: -0.1 alive: _target_: unilab.managers.RewardTermCfg @@ -218,7 +218,7 @@ reward: weight: 0.0 contact: _target_: unilab.managers.RewardTermCfg - func: unilab.envs.locomotion.common.manager_terms.feet_phase_contact + func: unilab.tasks.locomotion.common.manager_terms.feet_phase_contact weight: 0.24 params: frequency: 2.0 @@ -231,7 +231,7 @@ reward: stance_threshold: 0.6 swing_feet_z: _target_: unilab.managers.RewardTermCfg - func: unilab.envs.locomotion.common.manager_terms.feet_phase_swing_height + func: unilab.tasks.locomotion.common.manager_terms.feet_phase_swing_height weight: 4.0 params: frequency: 2.0 diff --git a/conf/offpolicy/task/go2_joystick_flat/base.yaml b/conf/offpolicy/task/go2_joystick_flat/base.yaml index 3c51d7605..639d29849 100644 --- a/conf/offpolicy/task/go2_joystick_flat/base.yaml +++ b/conf/offpolicy/task/go2_joystick_flat/base.yaml @@ -67,7 +67,7 @@ env: command_name: twist gait_phase: _target_: unilab.managers.ObservationTermCfg - func: unilab.envs.locomotion.common.manager_terms.quadruped_gait_phase + func: unilab.tasks.locomotion.common.manager_terms.quadruped_gait_phase params: frequency: 2.0 critic: @@ -99,7 +99,7 @@ env: command_name: twist gait_phase: _target_: unilab.managers.ObservationTermCfg - func: unilab.envs.locomotion.common.manager_terms.quadruped_gait_phase + func: unilab.tasks.locomotion.common.manager_terms.quadruped_gait_phase params: frequency: 2.0 base_lin_vel: @@ -178,29 +178,29 @@ env: reward: tracking_lin_vel: _target_: unilab.managers.RewardTermCfg - func: unilab.envs.locomotion.common.manager_terms.track_lin_vel_xy_exp + func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel_xy_exp weight: 1.0 params: std: 0.5 command_name: twist tracking_ang_vel: _target_: unilab.managers.RewardTermCfg - func: unilab.envs.locomotion.common.manager_terms.track_ang_vel_z_exp + func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel_z_exp weight: 0.2 params: std: 0.5 command_name: twist lin_vel_z: _target_: unilab.managers.RewardTermCfg - func: unilab.envs.locomotion.common.manager_terms.lin_vel_z_l2 + func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z_l2 weight: -5.0 ang_vel_xy: _target_: unilab.managers.RewardTermCfg - func: unilab.envs.locomotion.common.manager_terms.ang_vel_xy_l2 + func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy_l2 weight: -0.1 base_height: _target_: unilab.managers.RewardTermCfg - func: unilab.envs.locomotion.common.manager_terms.base_height_l2 + func: unilab.tasks.locomotion.common.manager_terms.base_height_l2 weight: -100.0 params: target_height: 0.3 @@ -210,11 +210,11 @@ reward: weight: -0.005 similar_to_default: _target_: unilab.managers.RewardTermCfg - func: unilab.envs.locomotion.common.manager_terms.joint_deviation_l1 + func: unilab.tasks.locomotion.common.manager_terms.joint_deviation_l1 weight: -0.1 contact: _target_: unilab.managers.RewardTermCfg - func: unilab.envs.locomotion.common.manager_terms.feet_phase_contact + func: unilab.tasks.locomotion.common.manager_terms.feet_phase_contact weight: 0.24 params: frequency: 2.0 @@ -227,7 +227,7 @@ reward: stance_threshold: 0.6 swing_feet_z: _target_: unilab.managers.RewardTermCfg - func: unilab.envs.locomotion.common.manager_terms.feet_phase_swing_height + func: unilab.tasks.locomotion.common.manager_terms.feet_phase_swing_height weight: 4.0 params: frequency: 2.0 diff --git a/conf/ppo/task/go2_joystick_flat/base.yaml b/conf/ppo/task/go2_joystick_flat/base.yaml index 3c51d7605..639d29849 100644 --- a/conf/ppo/task/go2_joystick_flat/base.yaml +++ b/conf/ppo/task/go2_joystick_flat/base.yaml @@ -67,7 +67,7 @@ env: command_name: twist gait_phase: _target_: unilab.managers.ObservationTermCfg - func: unilab.envs.locomotion.common.manager_terms.quadruped_gait_phase + func: unilab.tasks.locomotion.common.manager_terms.quadruped_gait_phase params: frequency: 2.0 critic: @@ -99,7 +99,7 @@ env: command_name: twist gait_phase: _target_: unilab.managers.ObservationTermCfg - func: unilab.envs.locomotion.common.manager_terms.quadruped_gait_phase + func: unilab.tasks.locomotion.common.manager_terms.quadruped_gait_phase params: frequency: 2.0 base_lin_vel: @@ -178,29 +178,29 @@ env: reward: tracking_lin_vel: _target_: unilab.managers.RewardTermCfg - func: unilab.envs.locomotion.common.manager_terms.track_lin_vel_xy_exp + func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel_xy_exp weight: 1.0 params: std: 0.5 command_name: twist tracking_ang_vel: _target_: unilab.managers.RewardTermCfg - func: unilab.envs.locomotion.common.manager_terms.track_ang_vel_z_exp + func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel_z_exp weight: 0.2 params: std: 0.5 command_name: twist lin_vel_z: _target_: unilab.managers.RewardTermCfg - func: unilab.envs.locomotion.common.manager_terms.lin_vel_z_l2 + func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z_l2 weight: -5.0 ang_vel_xy: _target_: unilab.managers.RewardTermCfg - func: unilab.envs.locomotion.common.manager_terms.ang_vel_xy_l2 + func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy_l2 weight: -0.1 base_height: _target_: unilab.managers.RewardTermCfg - func: unilab.envs.locomotion.common.manager_terms.base_height_l2 + func: unilab.tasks.locomotion.common.manager_terms.base_height_l2 weight: -100.0 params: target_height: 0.3 @@ -210,11 +210,11 @@ reward: weight: -0.005 similar_to_default: _target_: unilab.managers.RewardTermCfg - func: unilab.envs.locomotion.common.manager_terms.joint_deviation_l1 + func: unilab.tasks.locomotion.common.manager_terms.joint_deviation_l1 weight: -0.1 contact: _target_: unilab.managers.RewardTermCfg - func: unilab.envs.locomotion.common.manager_terms.feet_phase_contact + func: unilab.tasks.locomotion.common.manager_terms.feet_phase_contact weight: 0.24 params: frequency: 2.0 @@ -227,7 +227,7 @@ reward: stance_threshold: 0.6 swing_feet_z: _target_: unilab.managers.RewardTermCfg - func: unilab.envs.locomotion.common.manager_terms.feet_phase_swing_height + func: unilab.tasks.locomotion.common.manager_terms.feet_phase_swing_height weight: 4.0 params: frequency: 2.0 diff --git a/src/unilab/tasks/locomotion/common/__init__.py b/src/unilab/tasks/locomotion/common/__init__.py new file mode 100644 index 000000000..1ac324e95 --- /dev/null +++ b/src/unilab/tasks/locomotion/common/__init__.py @@ -0,0 +1 @@ +"""Shared task-specific locomotion components.""" diff --git a/src/unilab/envs/locomotion/common/manager_terms.py b/src/unilab/tasks/locomotion/common/manager_terms.py similarity index 99% rename from src/unilab/envs/locomotion/common/manager_terms.py rename to src/unilab/tasks/locomotion/common/manager_terms.py index 33e543e2b..a65f6c16b 100644 --- a/src/unilab/envs/locomotion/common/manager_terms.py +++ b/src/unilab/tasks/locomotion/common/manager_terms.py @@ -1,4 +1,4 @@ -"""Task-owned Manager-Based terms for the quadruped locomotion pilots. +"""Task-owned Manager-Based terms for quadruped locomotion. The equations come from UniLab's existing Go1/Go2 joystick tasks. The adaptation uses community ``func + params`` terms, NumPy, and the base-owned sensor facade. diff --git a/tests/envs/locomotion/test_manager_gait_terms.py b/tests/envs/locomotion/test_manager_gait_terms.py index d1cb00b5a..7d8fcf77d 100644 --- a/tests/envs/locomotion/test_manager_gait_terms.py +++ b/tests/envs/locomotion/test_manager_gait_terms.py @@ -11,7 +11,6 @@ from unilab.base.backend.base import BackendSensorView from unilab.dtype_config import get_global_dtype -from unilab.envs.locomotion.common import manager_terms from unilab.managers import ( ObservationGroupCfg, ObservationManager, @@ -21,6 +20,7 @@ ) from unilab.managers._types import ManagerBasedRlEnv from unilab.managers.scene_entity_config import SceneEntityCfg +from unilab.tasks.locomotion.common import manager_terms CONTACTS = ("fl_contact", "fr_contact", "rl_contact", "rr_contact") POSITIONS = ("fl_pos", "fr_pos", "rl_pos", "rr_pos") From 3d98ba2638c210dcf0e492a8e417969e2c3ab4a6 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 23:51:01 +0800 Subject: [PATCH 065/169] refactor(tasks): move locomotion command helpers [skip ci] (#1176) --- src/unilab/envs/locomotion/common/__init__.py | 12 ------------ src/unilab/tasks/locomotion/a2/joystick.py | 2 +- src/unilab/tasks/locomotion/common/__init__.py | 16 ++++++++++++++++ .../locomotion/common/commands.py | 2 ++ src/unilab/tasks/locomotion/g1/joystick.py | 8 ++++---- src/unilab/tasks/locomotion/go1/joystick.py | 2 +- src/unilab/tasks/locomotion/go1/rough.py | 12 ++++++------ src/unilab/tasks/locomotion/go2/footstand.py | 2 +- src/unilab/tasks/locomotion/go2/joystick.py | 2 +- src/unilab/tasks/locomotion/go2/rough.py | 10 +++++----- .../tasks/locomotion/go2_arm/manip_loco.py | 2 +- src/unilab/tasks/locomotion/go2w/joystick.py | 10 +++++----- src/unilab/tasks/locomotion/go2w/rough.py | 10 +++++----- .../locomotion/test_go2_joystick_stand_still.py | 2 +- 14 files changed, 49 insertions(+), 43 deletions(-) rename src/unilab/{envs => tasks}/locomotion/common/commands.py (98%) diff --git a/src/unilab/envs/locomotion/common/__init__.py b/src/unilab/envs/locomotion/common/__init__.py index 7455130cb..0c50ba3a0 100644 --- a/src/unilab/envs/locomotion/common/__init__.py +++ b/src/unilab/envs/locomotion/common/__init__.py @@ -6,13 +6,6 @@ PdControlConfig, Sensor, ) -from .commands import ( - Commands, - apply_heading_yaw_feedback, - sample_heading_commands, - sample_velocity_commands, - zero_small_xy_commands, -) from .domain_rand import DomainRandConfig from .dr_provider import LocomotionDRProvider from .height_scan import ( @@ -24,7 +17,6 @@ __all__ = [ "BaseNoiseConfig", - "Commands", "ControlConfigBase", "DEFAULT_SCAN_POINTS_X", "DEFAULT_SCAN_POINTS_Y", @@ -36,8 +28,4 @@ "PdControlConfig", "RewardContext", "Sensor", - "apply_heading_yaw_feedback", - "sample_heading_commands", - "sample_velocity_commands", - "zero_small_xy_commands", ] diff --git a/src/unilab/tasks/locomotion/a2/joystick.py b/src/unilab/tasks/locomotion/a2/joystick.py index 63aabec03..524949ec4 100644 --- a/src/unilab/tasks/locomotion/a2/joystick.py +++ b/src/unilab/tasks/locomotion/a2/joystick.py @@ -24,9 +24,9 @@ from unilab.base.scene import SceneCfg from unilab.dtype_config import get_global_dtype from unilab.envs.locomotion.common import rewards -from unilab.envs.locomotion.common.commands import sample_commands_with_standing from unilab.envs.locomotion.common.dr_provider import LocomotionDRProvider from unilab.envs.locomotion.common.rewards import RewardContext +from unilab.tasks.locomotion.common.commands import sample_commands_with_standing from unilab.tasks.locomotion.go2.base import Asset, ControlConfig from unilab.tasks.locomotion.go2.joystick import ( Go2DomainRandConfig, diff --git a/src/unilab/tasks/locomotion/common/__init__.py b/src/unilab/tasks/locomotion/common/__init__.py index 1ac324e95..d09d04c3f 100644 --- a/src/unilab/tasks/locomotion/common/__init__.py +++ b/src/unilab/tasks/locomotion/common/__init__.py @@ -1 +1,17 @@ """Shared task-specific locomotion components.""" + +from .commands import ( + Commands, + apply_heading_yaw_feedback, + sample_heading_commands, + sample_velocity_commands, + zero_small_xy_commands, +) + +__all__ = [ + "Commands", + "apply_heading_yaw_feedback", + "sample_heading_commands", + "sample_velocity_commands", + "zero_small_xy_commands", +] diff --git a/src/unilab/envs/locomotion/common/commands.py b/src/unilab/tasks/locomotion/common/commands.py similarity index 98% rename from src/unilab/envs/locomotion/common/commands.py rename to src/unilab/tasks/locomotion/common/commands.py index ea58ffbf4..3fdb2dc51 100644 --- a/src/unilab/envs/locomotion/common/commands.py +++ b/src/unilab/tasks/locomotion/common/commands.py @@ -1,3 +1,5 @@ +"""Shared command helpers for locomotion tasks.""" + from __future__ import annotations from dataclasses import dataclass, field diff --git a/src/unilab/tasks/locomotion/g1/joystick.py b/src/unilab/tasks/locomotion/g1/joystick.py index 351a50a69..d232ed706 100644 --- a/src/unilab/tasks/locomotion/g1/joystick.py +++ b/src/unilab/tasks/locomotion/g1/joystick.py @@ -17,14 +17,14 @@ from unilab.base.scene import SceneCfg from unilab.dtype_config import get_global_dtype from unilab.envs.locomotion.common import rewards -from unilab.envs.locomotion.common.commands import ( +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.tasks.locomotion.common.commands import ( Commands, sample_heading_commands, zero_small_xy_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 .base import G1BaseCfg, G1BaseEnv diff --git a/src/unilab/tasks/locomotion/go1/joystick.py b/src/unilab/tasks/locomotion/go1/joystick.py index 75289333a..dbdcfd6e7 100644 --- a/src/unilab/tasks/locomotion/go1/joystick.py +++ b/src/unilab/tasks/locomotion/go1/joystick.py @@ -12,7 +12,6 @@ from unilab.base.scene import SceneCfg from unilab.dtype_config import get_global_dtype 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 @@ -20,6 +19,7 @@ TerrainCurriculumCfg, TerrainSpawnManager, ) +from unilab.tasks.locomotion.common.commands import Commands from unilab.tasks.locomotion.go1.base import Go1BaseCfg, Go1BaseEnv diff --git a/src/unilab/tasks/locomotion/go1/rough.py b/src/unilab/tasks/locomotion/go1/rough.py index 82b12f6da..2726aafcf 100644 --- a/src/unilab/tasks/locomotion/go1/rough.py +++ b/src/unilab/tasks/locomotion/go1/rough.py @@ -15,12 +15,6 @@ from unilab.dr.dr_utils import build_common_reset_randomization, zero_actions from unilab.dtype_config import get_global_dtype from unilab.envs.locomotion.common import rewards -from unilab.envs.locomotion.common.commands import ( - Commands, - apply_heading_yaw_feedback, - sample_heading_commands, - zero_small_xy_commands, -) from unilab.envs.locomotion.common.domain_rand import DomainRandConfig from unilab.envs.locomotion.common.height_scan import ( HeightScanConfig, @@ -34,6 +28,12 @@ from unilab.envs.locomotion.common.terrain_spawn import ( TerrainCurriculumCfg, ) +from unilab.tasks.locomotion.common.commands import ( + Commands, + apply_heading_yaw_feedback, + sample_heading_commands, + zero_small_xy_commands, +) from unilab.tasks.locomotion.go1.base import ControlConfig from unilab.tasks.locomotion.go1.joystick import ( Go1JoystickCfg, diff --git a/src/unilab/tasks/locomotion/go2/footstand.py b/src/unilab/tasks/locomotion/go2/footstand.py index f306797d2..80b94473a 100644 --- a/src/unilab/tasks/locomotion/go2/footstand.py +++ b/src/unilab/tasks/locomotion/go2/footstand.py @@ -13,10 +13,10 @@ from unilab.dr import ResetPlan, ResetRandomizationPayload from unilab.dtype_config import get_global_dtype 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.tasks.locomotion.common.commands import Commands from unilab.tasks.locomotion.go2.base import ControlConfig, Go2BaseCfg, Go2BaseEnv, NoiseConfig from unilab.utils.rotation import np_quat_apply, np_quat_apply_inverse diff --git a/src/unilab/tasks/locomotion/go2/joystick.py b/src/unilab/tasks/locomotion/go2/joystick.py index 303f82542..79ef97f32 100644 --- a/src/unilab/tasks/locomotion/go2/joystick.py +++ b/src/unilab/tasks/locomotion/go2/joystick.py @@ -13,7 +13,6 @@ from unilab.dtype_config import get_global_dtype from unilab.envs.locomotion.common import rewards from unilab.envs.locomotion.common.base import Sensor -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 @@ -25,6 +24,7 @@ ManagerBasedRlEnvCfg, make_manager_based_rl_env, ) +from unilab.tasks.locomotion.common.commands import Commands from unilab.tasks.locomotion.go2.base import Go2BaseCfg, Go2BaseEnv diff --git a/src/unilab/tasks/locomotion/go2/rough.py b/src/unilab/tasks/locomotion/go2/rough.py index badde1fda..88e3c51f1 100644 --- a/src/unilab/tasks/locomotion/go2/rough.py +++ b/src/unilab/tasks/locomotion/go2/rough.py @@ -13,11 +13,6 @@ from unilab.dr.dr_utils import build_common_reset_randomization, zero_actions from unilab.dtype_config import get_global_dtype from unilab.envs.locomotion.common import rewards -from unilab.envs.locomotion.common.commands import ( - apply_heading_yaw_feedback, - sample_heading_commands, - zero_small_xy_commands, -) from unilab.envs.locomotion.common.height_scan import ( DEFAULT_SCAN_POINTS_X, DEFAULT_SCAN_POINTS_Y, @@ -29,6 +24,11 @@ terrain_out_of_bounds, ) from unilab.envs.locomotion.common.rewards import RewardContext +from unilab.tasks.locomotion.common.commands import ( + apply_heading_yaw_feedback, + sample_heading_commands, + zero_small_xy_commands, +) from unilab.tasks.locomotion.go2.base import ControlConfig from unilab.tasks.locomotion.go2.joystick import ( Commands, diff --git a/src/unilab/tasks/locomotion/go2_arm/manip_loco.py b/src/unilab/tasks/locomotion/go2_arm/manip_loco.py index 6e5a3eda0..e6818acca 100644 --- a/src/unilab/tasks/locomotion/go2_arm/manip_loco.py +++ b/src/unilab/tasks/locomotion/go2_arm/manip_loco.py @@ -13,10 +13,10 @@ from unilab.dr.types import ResetPlan from unilab.dtype_config import get_global_dtype 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.tasks.locomotion.common.commands import Commands from unilab.tasks.locomotion.go2_arm.base import ( Go2ArmBaseCfg, Go2ArmBaseEnv, diff --git a/src/unilab/tasks/locomotion/go2w/joystick.py b/src/unilab/tasks/locomotion/go2w/joystick.py index dc13a3e78..c688dc0b0 100644 --- a/src/unilab/tasks/locomotion/go2w/joystick.py +++ b/src/unilab/tasks/locomotion/go2w/joystick.py @@ -18,17 +18,17 @@ ) from unilab.dtype_config import get_global_dtype from unilab.envs.locomotion.common import rewards -from unilab.envs.locomotion.common.commands import ( +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.tasks.locomotion.common.commands import ( Commands, apply_heading_yaw_feedback, zero_small_xy_commands, ) -from unilab.envs.locomotion.common.commands import ( +from unilab.tasks.locomotion.common.commands import ( sample_heading_commands as sample_go2w_heading_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.tasks.locomotion.go2w.base import ( DEFAULT_GO2W_ANGLES, NUM_GO2W_ACTIONS, diff --git a/src/unilab/tasks/locomotion/go2w/rough.py b/src/unilab/tasks/locomotion/go2w/rough.py index de1db4ee5..fb57e5e8c 100644 --- a/src/unilab/tasks/locomotion/go2w/rough.py +++ b/src/unilab/tasks/locomotion/go2w/rough.py @@ -13,11 +13,6 @@ from unilab.dr.dr_utils import zero_actions from unilab.dtype_config import get_global_dtype from unilab.envs.locomotion.common import rewards -from unilab.envs.locomotion.common.commands import ( - Commands, - apply_heading_yaw_feedback, - zero_small_xy_commands, -) from unilab.envs.locomotion.common.height_scan import ( HeightScanConfig, base_height_from_scan, @@ -31,6 +26,11 @@ TerrainCurriculumCfg, TerrainSpawnManager, ) +from unilab.tasks.locomotion.common.commands import ( + Commands, + apply_heading_yaw_feedback, + zero_small_xy_commands, +) from unilab.tasks.locomotion.go2w.base import NUM_GO2W_ACTIONS, NUM_LEG_ACTIONS from unilab.tasks.locomotion.go2w.joystick import ( Go2WJoystickCfg, diff --git a/tests/envs/locomotion/test_go2_joystick_stand_still.py b/tests/envs/locomotion/test_go2_joystick_stand_still.py index 8f27eb2d3..021584ba5 100644 --- a/tests/envs/locomotion/test_go2_joystick_stand_still.py +++ b/tests/envs/locomotion/test_go2_joystick_stand_still.py @@ -17,7 +17,7 @@ import numpy as np -from unilab.envs.locomotion.common.commands import ( +from unilab.tasks.locomotion.common.commands import ( sample_commands_with_standing, zero_small_xy_commands, ) From fc30d98b2b5356946b30fe5f6a00fdad3b5c7cc7 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 23:55:25 +0800 Subject: [PATCH 066/169] refactor(tasks): move locomotion DR config [skip ci] (#1178) --- src/unilab/envs/locomotion/common/__init__.py | 2 -- src/unilab/tasks/locomotion/common/__init__.py | 2 ++ src/unilab/{envs => tasks}/locomotion/common/domain_rand.py | 2 ++ src/unilab/tasks/locomotion/g1/joystick.py | 2 +- src/unilab/tasks/locomotion/go1/joystick.py | 2 +- src/unilab/tasks/locomotion/go1/rough.py | 2 +- src/unilab/tasks/locomotion/go2/footstand.py | 2 +- src/unilab/tasks/locomotion/go2/joystick.py | 2 +- src/unilab/tasks/locomotion/go2_arm/manip_loco.py | 2 +- src/unilab/tasks/locomotion/go2w/joystick.py | 2 +- 10 files changed, 11 insertions(+), 9 deletions(-) rename src/unilab/{envs => tasks}/locomotion/common/domain_rand.py (94%) diff --git a/src/unilab/envs/locomotion/common/__init__.py b/src/unilab/envs/locomotion/common/__init__.py index 0c50ba3a0..f42b526a4 100644 --- a/src/unilab/envs/locomotion/common/__init__.py +++ b/src/unilab/envs/locomotion/common/__init__.py @@ -6,7 +6,6 @@ PdControlConfig, Sensor, ) -from .domain_rand import DomainRandConfig from .dr_provider import LocomotionDRProvider from .height_scan import ( DEFAULT_SCAN_POINTS_X, @@ -20,7 +19,6 @@ "ControlConfigBase", "DEFAULT_SCAN_POINTS_X", "DEFAULT_SCAN_POINTS_Y", - "DomainRandConfig", "HeightScanConfig", "LocomotionBaseCfg", "LocomotionBaseEnv", diff --git a/src/unilab/tasks/locomotion/common/__init__.py b/src/unilab/tasks/locomotion/common/__init__.py index d09d04c3f..f3b9e1163 100644 --- a/src/unilab/tasks/locomotion/common/__init__.py +++ b/src/unilab/tasks/locomotion/common/__init__.py @@ -7,9 +7,11 @@ sample_velocity_commands, zero_small_xy_commands, ) +from .domain_rand import DomainRandConfig __all__ = [ "Commands", + "DomainRandConfig", "apply_heading_yaw_feedback", "sample_heading_commands", "sample_velocity_commands", diff --git a/src/unilab/envs/locomotion/common/domain_rand.py b/src/unilab/tasks/locomotion/common/domain_rand.py similarity index 94% rename from src/unilab/envs/locomotion/common/domain_rand.py rename to src/unilab/tasks/locomotion/common/domain_rand.py index 9e10c893a..968d9a894 100644 --- a/src/unilab/envs/locomotion/common/domain_rand.py +++ b/src/unilab/tasks/locomotion/common/domain_rand.py @@ -1,3 +1,5 @@ +"""Shared domain-randomization config for locomotion tasks.""" + from __future__ import annotations from dataclasses import dataclass, field diff --git a/src/unilab/tasks/locomotion/g1/joystick.py b/src/unilab/tasks/locomotion/g1/joystick.py index d232ed706..98b46ad12 100644 --- a/src/unilab/tasks/locomotion/g1/joystick.py +++ b/src/unilab/tasks/locomotion/g1/joystick.py @@ -17,7 +17,6 @@ from unilab.base.scene import SceneCfg from unilab.dtype_config import get_global_dtype from unilab.envs.locomotion.common import rewards -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.tasks.locomotion.common.commands import ( @@ -25,6 +24,7 @@ sample_heading_commands, zero_small_xy_commands, ) +from unilab.tasks.locomotion.common.domain_rand import DomainRandConfig from .base import G1BaseCfg, G1BaseEnv diff --git a/src/unilab/tasks/locomotion/go1/joystick.py b/src/unilab/tasks/locomotion/go1/joystick.py index dbdcfd6e7..1802ccfd3 100644 --- a/src/unilab/tasks/locomotion/go1/joystick.py +++ b/src/unilab/tasks/locomotion/go1/joystick.py @@ -12,7 +12,6 @@ from unilab.base.scene import SceneCfg from unilab.dtype_config import get_global_dtype from unilab.envs.locomotion.common import rewards -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.common.terrain_spawn import ( @@ -20,6 +19,7 @@ TerrainSpawnManager, ) from unilab.tasks.locomotion.common.commands import Commands +from unilab.tasks.locomotion.common.domain_rand import DomainRandConfig from unilab.tasks.locomotion.go1.base import Go1BaseCfg, Go1BaseEnv diff --git a/src/unilab/tasks/locomotion/go1/rough.py b/src/unilab/tasks/locomotion/go1/rough.py index 2726aafcf..cd2ba886b 100644 --- a/src/unilab/tasks/locomotion/go1/rough.py +++ b/src/unilab/tasks/locomotion/go1/rough.py @@ -15,7 +15,6 @@ from unilab.dr.dr_utils import build_common_reset_randomization, zero_actions from unilab.dtype_config import get_global_dtype from unilab.envs.locomotion.common import rewards -from unilab.envs.locomotion.common.domain_rand import DomainRandConfig from unilab.envs.locomotion.common.height_scan import ( HeightScanConfig, base_height_from_scan, @@ -34,6 +33,7 @@ sample_heading_commands, zero_small_xy_commands, ) +from unilab.tasks.locomotion.common.domain_rand import DomainRandConfig from unilab.tasks.locomotion.go1.base import ControlConfig from unilab.tasks.locomotion.go1.joystick import ( Go1JoystickCfg, diff --git a/src/unilab/tasks/locomotion/go2/footstand.py b/src/unilab/tasks/locomotion/go2/footstand.py index 80b94473a..ffcda2bc6 100644 --- a/src/unilab/tasks/locomotion/go2/footstand.py +++ b/src/unilab/tasks/locomotion/go2/footstand.py @@ -13,10 +13,10 @@ from unilab.dr import ResetPlan, ResetRandomizationPayload from unilab.dtype_config import get_global_dtype from unilab.envs.locomotion.common import rewards -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.tasks.locomotion.common.commands import Commands +from unilab.tasks.locomotion.common.domain_rand import DomainRandConfig from unilab.tasks.locomotion.go2.base import ControlConfig, Go2BaseCfg, Go2BaseEnv, NoiseConfig from unilab.utils.rotation import np_quat_apply, np_quat_apply_inverse diff --git a/src/unilab/tasks/locomotion/go2/joystick.py b/src/unilab/tasks/locomotion/go2/joystick.py index 79ef97f32..cd5499954 100644 --- a/src/unilab/tasks/locomotion/go2/joystick.py +++ b/src/unilab/tasks/locomotion/go2/joystick.py @@ -13,7 +13,6 @@ from unilab.dtype_config import get_global_dtype from unilab.envs.locomotion.common import rewards from unilab.envs.locomotion.common.base import Sensor -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.common.terrain_spawn import ( @@ -25,6 +24,7 @@ make_manager_based_rl_env, ) from unilab.tasks.locomotion.common.commands import Commands +from unilab.tasks.locomotion.common.domain_rand import DomainRandConfig from unilab.tasks.locomotion.go2.base import Go2BaseCfg, Go2BaseEnv diff --git a/src/unilab/tasks/locomotion/go2_arm/manip_loco.py b/src/unilab/tasks/locomotion/go2_arm/manip_loco.py index e6818acca..a8901d911 100644 --- a/src/unilab/tasks/locomotion/go2_arm/manip_loco.py +++ b/src/unilab/tasks/locomotion/go2_arm/manip_loco.py @@ -13,10 +13,10 @@ from unilab.dr.types import ResetPlan from unilab.dtype_config import get_global_dtype from unilab.envs.locomotion.common import rewards -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.tasks.locomotion.common.commands import Commands +from unilab.tasks.locomotion.common.domain_rand import DomainRandConfig from unilab.tasks.locomotion.go2_arm.base import ( Go2ArmBaseCfg, Go2ArmBaseEnv, diff --git a/src/unilab/tasks/locomotion/go2w/joystick.py b/src/unilab/tasks/locomotion/go2w/joystick.py index c688dc0b0..51e18f158 100644 --- a/src/unilab/tasks/locomotion/go2w/joystick.py +++ b/src/unilab/tasks/locomotion/go2w/joystick.py @@ -18,7 +18,6 @@ ) from unilab.dtype_config import get_global_dtype from unilab.envs.locomotion.common import rewards -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.tasks.locomotion.common.commands import ( @@ -29,6 +28,7 @@ from unilab.tasks.locomotion.common.commands import ( sample_heading_commands as sample_go2w_heading_commands, ) +from unilab.tasks.locomotion.common.domain_rand import DomainRandConfig from unilab.tasks.locomotion.go2w.base import ( DEFAULT_GO2W_ANGLES, NUM_GO2W_ACTIONS, From 3fa30dcdb378931be5efc1dec6a155d39bfb3603 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Tue, 18 Aug 2026 23:59:49 +0800 Subject: [PATCH 067/169] refactor(tasks): move locomotion DR provider [skip ci] (#1180) --- .../en/3-deployment/1-sim_to_real/6-domain_randomization.md | 2 +- .../zh_CN/3-deployment/1-sim_to_real/6-domain_randomization.md | 2 +- src/unilab/envs/locomotion/common/__init__.py | 2 -- src/unilab/tasks/locomotion/a2/joystick.py | 2 +- src/unilab/tasks/locomotion/common/__init__.py | 2 ++ src/unilab/{envs => tasks}/locomotion/common/dr_provider.py | 2 +- src/unilab/tasks/locomotion/g1/joystick.py | 2 +- src/unilab/tasks/locomotion/go1/joystick.py | 2 +- src/unilab/tasks/locomotion/go2/footstand.py | 2 +- src/unilab/tasks/locomotion/go2/joystick.py | 2 +- src/unilab/tasks/locomotion/go2_arm/manip_loco.py | 2 +- src/unilab/tasks/locomotion/go2w/joystick.py | 2 +- 12 files changed, 12 insertions(+), 12 deletions(-) rename src/unilab/{envs => tasks}/locomotion/common/dr_provider.py (99%) diff --git a/docs/sphinx/source/en/3-deployment/1-sim_to_real/6-domain_randomization.md b/docs/sphinx/source/en/3-deployment/1-sim_to_real/6-domain_randomization.md index 1a4512a13..64de13a12 100644 --- a/docs/sphinx/source/en/3-deployment/1-sim_to_real/6-domain_randomization.md +++ b/docs/sphinx/source/en/3-deployment/1-sim_to_real/6-domain_randomization.md @@ -45,7 +45,7 @@ range in the task owner only after recording why that range is plausible. Tasks that use DR attach a provider through the env initialization path: ```python -from unilab.envs.locomotion.common.dr_provider import LocomotionDRProvider +from unilab.tasks.locomotion.common.dr_provider import LocomotionDRProvider class MyTaskEnv(NpEnv): def __init__(self, cfg): diff --git a/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/6-domain_randomization.md b/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/6-domain_randomization.md index 53d864e55..6f642cdf3 100644 --- a/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/6-domain_randomization.md +++ b/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/6-domain_randomization.md @@ -43,7 +43,7 @@ 使用 DR 的任务通过环境初始化路径挂接一个 provider： ```python -from unilab.envs.locomotion.common.dr_provider import LocomotionDRProvider +from unilab.tasks.locomotion.common.dr_provider import LocomotionDRProvider class MyTaskEnv(NpEnv): def __init__(self, cfg): diff --git a/src/unilab/envs/locomotion/common/__init__.py b/src/unilab/envs/locomotion/common/__init__.py index f42b526a4..76f2b61e6 100644 --- a/src/unilab/envs/locomotion/common/__init__.py +++ b/src/unilab/envs/locomotion/common/__init__.py @@ -6,7 +6,6 @@ PdControlConfig, Sensor, ) -from .dr_provider import LocomotionDRProvider from .height_scan import ( DEFAULT_SCAN_POINTS_X, DEFAULT_SCAN_POINTS_Y, @@ -22,7 +21,6 @@ "HeightScanConfig", "LocomotionBaseCfg", "LocomotionBaseEnv", - "LocomotionDRProvider", "PdControlConfig", "RewardContext", "Sensor", diff --git a/src/unilab/tasks/locomotion/a2/joystick.py b/src/unilab/tasks/locomotion/a2/joystick.py index 524949ec4..dd6f45c68 100644 --- a/src/unilab/tasks/locomotion/a2/joystick.py +++ b/src/unilab/tasks/locomotion/a2/joystick.py @@ -24,9 +24,9 @@ from unilab.base.scene import SceneCfg from unilab.dtype_config import get_global_dtype from unilab.envs.locomotion.common import rewards -from unilab.envs.locomotion.common.dr_provider import LocomotionDRProvider from unilab.envs.locomotion.common.rewards import RewardContext from unilab.tasks.locomotion.common.commands import sample_commands_with_standing +from unilab.tasks.locomotion.common.dr_provider import LocomotionDRProvider from unilab.tasks.locomotion.go2.base import Asset, ControlConfig from unilab.tasks.locomotion.go2.joystick import ( Go2DomainRandConfig, diff --git a/src/unilab/tasks/locomotion/common/__init__.py b/src/unilab/tasks/locomotion/common/__init__.py index f3b9e1163..1f36c0758 100644 --- a/src/unilab/tasks/locomotion/common/__init__.py +++ b/src/unilab/tasks/locomotion/common/__init__.py @@ -8,10 +8,12 @@ zero_small_xy_commands, ) from .domain_rand import DomainRandConfig +from .dr_provider import LocomotionDRProvider __all__ = [ "Commands", "DomainRandConfig", + "LocomotionDRProvider", "apply_heading_yaw_feedback", "sample_heading_commands", "sample_velocity_commands", diff --git a/src/unilab/envs/locomotion/common/dr_provider.py b/src/unilab/tasks/locomotion/common/dr_provider.py similarity index 99% rename from src/unilab/envs/locomotion/common/dr_provider.py rename to src/unilab/tasks/locomotion/common/dr_provider.py index 24f75ab54..182274e4d 100644 --- a/src/unilab/envs/locomotion/common/dr_provider.py +++ b/src/unilab/tasks/locomotion/common/dr_provider.py @@ -1,4 +1,4 @@ -"""Shared DomainRandomizationProvider for locomotion environments. +"""Shared DomainRandomizationProvider for locomotion tasks. Implements the common reset/interval randomization logic shared by G1, Go1, and Go2 joystick environments. Subclasses override hooks diff --git a/src/unilab/tasks/locomotion/g1/joystick.py b/src/unilab/tasks/locomotion/g1/joystick.py index 98b46ad12..bd715ffd6 100644 --- a/src/unilab/tasks/locomotion/g1/joystick.py +++ b/src/unilab/tasks/locomotion/g1/joystick.py @@ -17,7 +17,6 @@ from unilab.base.scene import SceneCfg from unilab.dtype_config import get_global_dtype from unilab.envs.locomotion.common import rewards -from unilab.envs.locomotion.common.dr_provider import LocomotionDRProvider from unilab.envs.locomotion.common.rewards import RewardContext from unilab.tasks.locomotion.common.commands import ( Commands, @@ -25,6 +24,7 @@ zero_small_xy_commands, ) from unilab.tasks.locomotion.common.domain_rand import DomainRandConfig +from unilab.tasks.locomotion.common.dr_provider import LocomotionDRProvider from .base import G1BaseCfg, G1BaseEnv diff --git a/src/unilab/tasks/locomotion/go1/joystick.py b/src/unilab/tasks/locomotion/go1/joystick.py index 1802ccfd3..58c022661 100644 --- a/src/unilab/tasks/locomotion/go1/joystick.py +++ b/src/unilab/tasks/locomotion/go1/joystick.py @@ -12,7 +12,6 @@ from unilab.base.scene import SceneCfg from unilab.dtype_config import get_global_dtype from unilab.envs.locomotion.common import rewards -from unilab.envs.locomotion.common.dr_provider import LocomotionDRProvider from unilab.envs.locomotion.common.rewards import RewardContext from unilab.envs.locomotion.common.terrain_spawn import ( TerrainCurriculumCfg, @@ -20,6 +19,7 @@ ) from unilab.tasks.locomotion.common.commands import Commands from unilab.tasks.locomotion.common.domain_rand import DomainRandConfig +from unilab.tasks.locomotion.common.dr_provider import LocomotionDRProvider from unilab.tasks.locomotion.go1.base import Go1BaseCfg, Go1BaseEnv diff --git a/src/unilab/tasks/locomotion/go2/footstand.py b/src/unilab/tasks/locomotion/go2/footstand.py index ffcda2bc6..e4d0b6adc 100644 --- a/src/unilab/tasks/locomotion/go2/footstand.py +++ b/src/unilab/tasks/locomotion/go2/footstand.py @@ -13,10 +13,10 @@ from unilab.dr import ResetPlan, ResetRandomizationPayload from unilab.dtype_config import get_global_dtype from unilab.envs.locomotion.common import rewards -from unilab.envs.locomotion.common.dr_provider import LocomotionDRProvider from unilab.envs.locomotion.common.rewards import RewardContext from unilab.tasks.locomotion.common.commands import Commands from unilab.tasks.locomotion.common.domain_rand import DomainRandConfig +from unilab.tasks.locomotion.common.dr_provider import LocomotionDRProvider from unilab.tasks.locomotion.go2.base import ControlConfig, Go2BaseCfg, Go2BaseEnv, NoiseConfig from unilab.utils.rotation import np_quat_apply, np_quat_apply_inverse diff --git a/src/unilab/tasks/locomotion/go2/joystick.py b/src/unilab/tasks/locomotion/go2/joystick.py index cd5499954..614006968 100644 --- a/src/unilab/tasks/locomotion/go2/joystick.py +++ b/src/unilab/tasks/locomotion/go2/joystick.py @@ -13,7 +13,6 @@ from unilab.dtype_config import get_global_dtype from unilab.envs.locomotion.common import rewards from unilab.envs.locomotion.common.base import Sensor -from unilab.envs.locomotion.common.dr_provider import LocomotionDRProvider from unilab.envs.locomotion.common.rewards import RewardContext from unilab.envs.locomotion.common.terrain_spawn import ( TerrainCurriculumCfg, @@ -25,6 +24,7 @@ ) from unilab.tasks.locomotion.common.commands import Commands from unilab.tasks.locomotion.common.domain_rand import DomainRandConfig +from unilab.tasks.locomotion.common.dr_provider import LocomotionDRProvider from unilab.tasks.locomotion.go2.base import Go2BaseCfg, Go2BaseEnv diff --git a/src/unilab/tasks/locomotion/go2_arm/manip_loco.py b/src/unilab/tasks/locomotion/go2_arm/manip_loco.py index a8901d911..0a1968863 100644 --- a/src/unilab/tasks/locomotion/go2_arm/manip_loco.py +++ b/src/unilab/tasks/locomotion/go2_arm/manip_loco.py @@ -13,10 +13,10 @@ from unilab.dr.types import ResetPlan from unilab.dtype_config import get_global_dtype from unilab.envs.locomotion.common import rewards -from unilab.envs.locomotion.common.dr_provider import LocomotionDRProvider from unilab.envs.locomotion.common.rewards import RewardContext from unilab.tasks.locomotion.common.commands import Commands from unilab.tasks.locomotion.common.domain_rand import DomainRandConfig +from unilab.tasks.locomotion.common.dr_provider import LocomotionDRProvider from unilab.tasks.locomotion.go2_arm.base import ( Go2ArmBaseCfg, Go2ArmBaseEnv, diff --git a/src/unilab/tasks/locomotion/go2w/joystick.py b/src/unilab/tasks/locomotion/go2w/joystick.py index 51e18f158..82a8715b0 100644 --- a/src/unilab/tasks/locomotion/go2w/joystick.py +++ b/src/unilab/tasks/locomotion/go2w/joystick.py @@ -18,7 +18,6 @@ ) from unilab.dtype_config import get_global_dtype from unilab.envs.locomotion.common import rewards -from unilab.envs.locomotion.common.dr_provider import LocomotionDRProvider from unilab.envs.locomotion.common.rewards import RewardContext from unilab.tasks.locomotion.common.commands import ( Commands, @@ -29,6 +28,7 @@ sample_heading_commands as sample_go2w_heading_commands, ) from unilab.tasks.locomotion.common.domain_rand import DomainRandConfig +from unilab.tasks.locomotion.common.dr_provider import LocomotionDRProvider from unilab.tasks.locomotion.go2w.base import ( DEFAULT_GO2W_ANGLES, NUM_GO2W_ACTIONS, From c854fb04baec3a368e2874821ad9abb4a5adcca6 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Wed, 19 Aug 2026 00:04:12 +0800 Subject: [PATCH 068/169] refactor(tasks): move locomotion height scan [skip ci] (#1182) --- .../3-extending/4-new_terrain.md | 2 +- .../3-extending/4-new_terrain.md | 2 +- src/unilab/envs/locomotion/common/__init__.py | 8 -------- src/unilab/tasks/locomotion/common/__init__.py | 8 ++++++++ .../locomotion/common/height_scan.py | 2 +- src/unilab/tasks/locomotion/go1/rough.py | 16 ++++++++-------- src/unilab/tasks/locomotion/go2/rough.py | 14 +++++++------- src/unilab/tasks/locomotion/go2w/rough.py | 16 ++++++++-------- .../locomotion/test_go2_rough_height_scan.py | 2 +- 9 files changed, 35 insertions(+), 35 deletions(-) rename src/unilab/{envs => tasks}/locomotion/common/height_scan.py (99%) diff --git a/docs/sphinx/source/en/4-developer_guide/3-extending/4-new_terrain.md b/docs/sphinx/source/en/4-developer_guide/3-extending/4-new_terrain.md index 23ae12c30..de2dd4cc4 100644 --- a/docs/sphinx/source/en/4-developer_guide/3-extending/4-new_terrain.md +++ b/docs/sphinx/source/en/4-developer_guide/3-extending/4-new_terrain.md @@ -35,4 +35,4 @@ materialization out of `step()`, `reset()`, and hot domain-randomization loops. - Terrain configs and presets: `src/unilab/terrains/config.py` - Terrain generator: `src/unilab/terrains/terrain_generator.py` - Heightfield terrain types: `src/unilab/terrains/heightfield_terrains.py` -- Height-scan helper: `src/unilab/envs/locomotion/common/height_scan.py` +- Height-scan helper: `src/unilab/tasks/locomotion/common/height_scan.py` diff --git a/docs/sphinx/source/zh_CN/4-developer_guide/3-extending/4-new_terrain.md b/docs/sphinx/source/zh_CN/4-developer_guide/3-extending/4-new_terrain.md index e330cd438..619a588b8 100644 --- a/docs/sphinx/source/zh_CN/4-developer_guide/3-extending/4-new_terrain.md +++ b/docs/sphinx/source/zh_CN/4-developer_guide/3-extending/4-new_terrain.md @@ -35,4 +35,4 @@ - 地形配置与 preset：`src/unilab/terrains/config.py` - 地形生成器：`src/unilab/terrains/terrain_generator.py` - Heightfield 地形类型：`src/unilab/terrains/heightfield_terrains.py` -- 高度扫描辅助工具：`src/unilab/envs/locomotion/common/height_scan.py` +- 高度扫描辅助工具：`src/unilab/tasks/locomotion/common/height_scan.py` diff --git a/src/unilab/envs/locomotion/common/__init__.py b/src/unilab/envs/locomotion/common/__init__.py index 76f2b61e6..96586a2fb 100644 --- a/src/unilab/envs/locomotion/common/__init__.py +++ b/src/unilab/envs/locomotion/common/__init__.py @@ -6,19 +6,11 @@ PdControlConfig, Sensor, ) -from .height_scan import ( - DEFAULT_SCAN_POINTS_X, - DEFAULT_SCAN_POINTS_Y, - HeightScanConfig, -) from .rewards import RewardContext __all__ = [ "BaseNoiseConfig", "ControlConfigBase", - "DEFAULT_SCAN_POINTS_X", - "DEFAULT_SCAN_POINTS_Y", - "HeightScanConfig", "LocomotionBaseCfg", "LocomotionBaseEnv", "PdControlConfig", diff --git a/src/unilab/tasks/locomotion/common/__init__.py b/src/unilab/tasks/locomotion/common/__init__.py index 1f36c0758..7e40f0c1f 100644 --- a/src/unilab/tasks/locomotion/common/__init__.py +++ b/src/unilab/tasks/locomotion/common/__init__.py @@ -9,10 +9,18 @@ ) from .domain_rand import DomainRandConfig from .dr_provider import LocomotionDRProvider +from .height_scan import ( + DEFAULT_SCAN_POINTS_X, + DEFAULT_SCAN_POINTS_Y, + HeightScanConfig, +) __all__ = [ "Commands", + "DEFAULT_SCAN_POINTS_X", + "DEFAULT_SCAN_POINTS_Y", "DomainRandConfig", + "HeightScanConfig", "LocomotionDRProvider", "apply_heading_yaw_feedback", "sample_heading_commands", diff --git a/src/unilab/envs/locomotion/common/height_scan.py b/src/unilab/tasks/locomotion/common/height_scan.py similarity index 99% rename from src/unilab/envs/locomotion/common/height_scan.py rename to src/unilab/tasks/locomotion/common/height_scan.py index 0f8bb0580..a557e8f8f 100644 --- a/src/unilab/envs/locomotion/common/height_scan.py +++ b/src/unilab/tasks/locomotion/common/height_scan.py @@ -1,4 +1,4 @@ -"""Shared height-scan and terrain-bound helpers for rough locomotion envs. +"""Shared height-scan and terrain-bound helpers for rough locomotion tasks. These functions and the ``HeightScanConfig`` dataclass are consumed by Go2, Go2W, Go1, and G1 rough environments — anywhere the policy / critic ingests diff --git a/src/unilab/tasks/locomotion/go1/rough.py b/src/unilab/tasks/locomotion/go1/rough.py index cd2ba886b..1906533b9 100644 --- a/src/unilab/tasks/locomotion/go1/rough.py +++ b/src/unilab/tasks/locomotion/go1/rough.py @@ -15,14 +15,6 @@ from unilab.dr.dr_utils import build_common_reset_randomization, zero_actions from unilab.dtype_config import get_global_dtype from unilab.envs.locomotion.common import rewards -from unilab.envs.locomotion.common.height_scan import ( - HeightScanConfig, - base_height_from_scan, - height_scan_obs, - init_height_scan_sensor, - raw_height_scan_obs, - terrain_out_of_bounds, -) from unilab.envs.locomotion.common.rewards import RewardContext from unilab.envs.locomotion.common.terrain_spawn import ( TerrainCurriculumCfg, @@ -34,6 +26,14 @@ zero_small_xy_commands, ) from unilab.tasks.locomotion.common.domain_rand import DomainRandConfig +from unilab.tasks.locomotion.common.height_scan import ( + HeightScanConfig, + base_height_from_scan, + height_scan_obs, + init_height_scan_sensor, + raw_height_scan_obs, + terrain_out_of_bounds, +) from unilab.tasks.locomotion.go1.base import ControlConfig from unilab.tasks.locomotion.go1.joystick import ( Go1JoystickCfg, diff --git a/src/unilab/tasks/locomotion/go2/rough.py b/src/unilab/tasks/locomotion/go2/rough.py index 88e3c51f1..ad675a8eb 100644 --- a/src/unilab/tasks/locomotion/go2/rough.py +++ b/src/unilab/tasks/locomotion/go2/rough.py @@ -13,7 +13,13 @@ from unilab.dr.dr_utils import build_common_reset_randomization, zero_actions from unilab.dtype_config import get_global_dtype from unilab.envs.locomotion.common import rewards -from unilab.envs.locomotion.common.height_scan import ( +from unilab.envs.locomotion.common.rewards import RewardContext +from unilab.tasks.locomotion.common.commands import ( + apply_heading_yaw_feedback, + sample_heading_commands, + zero_small_xy_commands, +) +from unilab.tasks.locomotion.common.height_scan import ( DEFAULT_SCAN_POINTS_X, DEFAULT_SCAN_POINTS_Y, HeightScanConfig, @@ -23,12 +29,6 @@ raw_height_scan_obs, terrain_out_of_bounds, ) -from unilab.envs.locomotion.common.rewards import RewardContext -from unilab.tasks.locomotion.common.commands import ( - apply_heading_yaw_feedback, - sample_heading_commands, - zero_small_xy_commands, -) from unilab.tasks.locomotion.go2.base import ControlConfig from unilab.tasks.locomotion.go2.joystick import ( Commands, diff --git a/src/unilab/tasks/locomotion/go2w/rough.py b/src/unilab/tasks/locomotion/go2w/rough.py index fb57e5e8c..8caf28ad0 100644 --- a/src/unilab/tasks/locomotion/go2w/rough.py +++ b/src/unilab/tasks/locomotion/go2w/rough.py @@ -13,14 +13,6 @@ from unilab.dr.dr_utils import zero_actions from unilab.dtype_config import get_global_dtype from unilab.envs.locomotion.common import rewards -from unilab.envs.locomotion.common.height_scan import ( - HeightScanConfig, - base_height_from_scan, - height_scan_obs, - init_height_scan_sensor, - raw_height_scan_obs, - terrain_out_of_bounds, -) from unilab.envs.locomotion.common.rewards import RewardContext from unilab.envs.locomotion.common.terrain_spawn import ( TerrainCurriculumCfg, @@ -31,6 +23,14 @@ apply_heading_yaw_feedback, zero_small_xy_commands, ) +from unilab.tasks.locomotion.common.height_scan import ( + HeightScanConfig, + base_height_from_scan, + height_scan_obs, + init_height_scan_sensor, + raw_height_scan_obs, + terrain_out_of_bounds, +) from unilab.tasks.locomotion.go2w.base import NUM_GO2W_ACTIONS, NUM_LEG_ACTIONS from unilab.tasks.locomotion.go2w.joystick import ( Go2WJoystickCfg, diff --git a/tests/envs/locomotion/test_go2_rough_height_scan.py b/tests/envs/locomotion/test_go2_rough_height_scan.py index ba62e62bc..8776f78f3 100644 --- a/tests/envs/locomotion/test_go2_rough_height_scan.py +++ b/tests/envs/locomotion/test_go2_rough_height_scan.py @@ -4,7 +4,7 @@ import numpy as np -from unilab.envs.locomotion.common.height_scan import height_scan_offsets as _height_scan_offsets +from unilab.tasks.locomotion.common.height_scan import height_scan_offsets as _height_scan_offsets from unilab.tasks.locomotion.go2.rough import Go2JoystickRoughEnv From 0e5261903d22741f0198facbd6d04c51c64ee311 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Wed, 19 Aug 2026 00:09:12 +0800 Subject: [PATCH 069/169] refactor(tasks): move locomotion shared base [skip ci] (#1184) --- .../3-framework_migration/2-from_legged_gym.md | 2 +- .../3-framework_migration/2-from_legged_gym.md | 2 +- src/unilab/envs/locomotion/common/__init__.py | 14 -------------- src/unilab/tasks/locomotion/common/__init__.py | 14 ++++++++++++++ .../{envs => tasks}/locomotion/common/base.py | 2 ++ src/unilab/tasks/locomotion/g1/base.py | 4 ++-- src/unilab/tasks/locomotion/go1/base.py | 2 +- src/unilab/tasks/locomotion/go2/base.py | 2 +- src/unilab/tasks/locomotion/go2/joystick.py | 2 +- src/unilab/tasks/locomotion/go2_arm/base.py | 2 +- src/unilab/tasks/locomotion/go2w/base.py | 2 +- .../locomotion/a2/test_a2_joystick_contract.py | 2 +- 12 files changed, 26 insertions(+), 24 deletions(-) rename src/unilab/{envs => tasks}/locomotion/common/base.py (98%) diff --git a/docs/sphinx/source/en/3-deployment/3-framework_migration/2-from_legged_gym.md b/docs/sphinx/source/en/3-deployment/3-framework_migration/2-from_legged_gym.md index 650c908ba..003a3a8d8 100644 --- a/docs/sphinx/source/en/3-deployment/3-framework_migration/2-from_legged_gym.md +++ b/docs/sphinx/source/en/3-deployment/3-framework_migration/2-from_legged_gym.md @@ -9,7 +9,7 @@ mostly mechanical. | Legged Gym | UniLab | |---|---| -| `LeggedRobot` env class | `unilab.envs.locomotion.common.base` | +| `LeggedRobot` env class | `unilab.tasks.locomotion.common.base` | | `compute_observations()` | env-side obs builder + `unilab.base.observations` | | `_reward_*` methods | env's `compute_reward()` + reward term registry | | `command_ranges` | task owner YAML's `commands` block | diff --git a/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/2-from_legged_gym.md b/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/2-from_legged_gym.md index 93c01c49f..05411ec4e 100644 --- a/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/2-from_legged_gym.md +++ b/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/2-from_legged_gym.md @@ -8,7 +8,7 @@ Legged Gym 曾是那套 GPU 常驻的 PPO 模板，教会了整个领域如何 | Legged Gym | UniLab | |---|---| -| `LeggedRobot` env 类 | `unilab.envs.locomotion.common.base` | +| `LeggedRobot` env 类 | `unilab.tasks.locomotion.common.base` | | `compute_observations()` | env 侧 obs 构建器 + `unilab.base.observations` | | `_reward_*` 方法 | env 的 `compute_reward()` + reward 项 registry | | `command_ranges` | 任务 owner YAML 的 `commands` 块 | diff --git a/src/unilab/envs/locomotion/common/__init__.py b/src/unilab/envs/locomotion/common/__init__.py index 96586a2fb..05d89eb00 100644 --- a/src/unilab/envs/locomotion/common/__init__.py +++ b/src/unilab/envs/locomotion/common/__init__.py @@ -1,19 +1,5 @@ -from .base import ( - BaseNoiseConfig, - ControlConfigBase, - LocomotionBaseCfg, - LocomotionBaseEnv, - PdControlConfig, - Sensor, -) from .rewards import RewardContext __all__ = [ - "BaseNoiseConfig", - "ControlConfigBase", - "LocomotionBaseCfg", - "LocomotionBaseEnv", - "PdControlConfig", "RewardContext", - "Sensor", ] diff --git a/src/unilab/tasks/locomotion/common/__init__.py b/src/unilab/tasks/locomotion/common/__init__.py index 7e40f0c1f..e21351fcf 100644 --- a/src/unilab/tasks/locomotion/common/__init__.py +++ b/src/unilab/tasks/locomotion/common/__init__.py @@ -1,5 +1,13 @@ """Shared task-specific locomotion components.""" +from .base import ( + BaseNoiseConfig, + ControlConfigBase, + LocomotionBaseCfg, + LocomotionBaseEnv, + PdControlConfig, + Sensor, +) from .commands import ( Commands, apply_heading_yaw_feedback, @@ -16,12 +24,18 @@ ) __all__ = [ + "BaseNoiseConfig", "Commands", + "ControlConfigBase", "DEFAULT_SCAN_POINTS_X", "DEFAULT_SCAN_POINTS_Y", "DomainRandConfig", "HeightScanConfig", "LocomotionDRProvider", + "LocomotionBaseCfg", + "LocomotionBaseEnv", + "PdControlConfig", + "Sensor", "apply_heading_yaw_feedback", "sample_heading_commands", "sample_velocity_commands", diff --git a/src/unilab/envs/locomotion/common/base.py b/src/unilab/tasks/locomotion/common/base.py similarity index 98% rename from src/unilab/envs/locomotion/common/base.py rename to src/unilab/tasks/locomotion/common/base.py index 15ef1677b..41f4e81f0 100644 --- a/src/unilab/envs/locomotion/common/base.py +++ b/src/unilab/tasks/locomotion/common/base.py @@ -1,3 +1,5 @@ +"""Shared base environment and configuration for legacy locomotion tasks.""" + from __future__ import annotations from dataclasses import dataclass, field diff --git a/src/unilab/tasks/locomotion/g1/base.py b/src/unilab/tasks/locomotion/g1/base.py index 18c23fb16..9c25bacf3 100644 --- a/src/unilab/tasks/locomotion/g1/base.py +++ b/src/unilab/tasks/locomotion/g1/base.py @@ -6,13 +6,13 @@ import numpy as np -from unilab.envs.locomotion.common.base import ( +from unilab.tasks.locomotion.common.base import ( BaseNoiseConfig, ControlConfigBase, LocomotionBaseCfg, LocomotionBaseEnv, ) -from unilab.envs.locomotion.common.base import ( +from unilab.tasks.locomotion.common.base import ( Sensor as LocomotionSensor, ) diff --git a/src/unilab/tasks/locomotion/go1/base.py b/src/unilab/tasks/locomotion/go1/base.py index 3869a9fd6..8b82c6681 100644 --- a/src/unilab/tasks/locomotion/go1/base.py +++ b/src/unilab/tasks/locomotion/go1/base.py @@ -2,7 +2,7 @@ from dataclasses import dataclass, field -from unilab.envs.locomotion.common.base import ( +from unilab.tasks.locomotion.common.base import ( BaseNoiseConfig, LocomotionBaseCfg, LocomotionBaseEnv, diff --git a/src/unilab/tasks/locomotion/go2/base.py b/src/unilab/tasks/locomotion/go2/base.py index e00dc0eb5..d3704c0a0 100644 --- a/src/unilab/tasks/locomotion/go2/base.py +++ b/src/unilab/tasks/locomotion/go2/base.py @@ -4,7 +4,7 @@ import numpy as np -from unilab.envs.locomotion.common.base import ( +from unilab.tasks.locomotion.common.base import ( BaseNoiseConfig, LocomotionBaseCfg, LocomotionBaseEnv, diff --git a/src/unilab/tasks/locomotion/go2/joystick.py b/src/unilab/tasks/locomotion/go2/joystick.py index 614006968..5e489ed1d 100644 --- a/src/unilab/tasks/locomotion/go2/joystick.py +++ b/src/unilab/tasks/locomotion/go2/joystick.py @@ -12,7 +12,6 @@ from unilab.base.scene import SceneCfg from unilab.dtype_config import get_global_dtype from unilab.envs.locomotion.common import rewards -from unilab.envs.locomotion.common.base import Sensor from unilab.envs.locomotion.common.rewards import RewardContext from unilab.envs.locomotion.common.terrain_spawn import ( TerrainCurriculumCfg, @@ -22,6 +21,7 @@ ManagerBasedRlEnvCfg, make_manager_based_rl_env, ) +from unilab.tasks.locomotion.common.base import Sensor from unilab.tasks.locomotion.common.commands import Commands from unilab.tasks.locomotion.common.domain_rand import DomainRandConfig from unilab.tasks.locomotion.common.dr_provider import LocomotionDRProvider diff --git a/src/unilab/tasks/locomotion/go2_arm/base.py b/src/unilab/tasks/locomotion/go2_arm/base.py index dee8a0b68..72cc680d6 100644 --- a/src/unilab/tasks/locomotion/go2_arm/base.py +++ b/src/unilab/tasks/locomotion/go2_arm/base.py @@ -5,7 +5,7 @@ import numpy as np from unilab.base.backend import SimBackend -from unilab.envs.locomotion.common.base import ( +from unilab.tasks.locomotion.common.base import ( ControlConfigBase, LocomotionBaseCfg, LocomotionBaseEnv, diff --git a/src/unilab/tasks/locomotion/go2w/base.py b/src/unilab/tasks/locomotion/go2w/base.py index 3875e1919..7ef1f4def 100644 --- a/src/unilab/tasks/locomotion/go2w/base.py +++ b/src/unilab/tasks/locomotion/go2w/base.py @@ -5,7 +5,7 @@ import gymnasium as gym import numpy as np -from unilab.envs.locomotion.common.base import ( +from unilab.tasks.locomotion.common.base import ( BaseNoiseConfig, LocomotionBaseCfg, LocomotionBaseEnv, diff --git a/tests/envs/locomotion/a2/test_a2_joystick_contract.py b/tests/envs/locomotion/a2/test_a2_joystick_contract.py index 0e3fbf520..c81b5dbb2 100644 --- a/tests/envs/locomotion/a2/test_a2_joystick_contract.py +++ b/tests/envs/locomotion/a2/test_a2_joystick_contract.py @@ -135,7 +135,7 @@ def test_a2_control_config_per_joint_gains(): def test_pd_control_config_position_gains_default_is_scalar(): """Base PdControlConfig keeps the scalar gain contract (Go2 path unchanged).""" - from unilab.envs.locomotion.common.base import PdControlConfig + from unilab.tasks.locomotion.common.base import PdControlConfig gains = PdControlConfig(Kp=35.0, Kd=0.5).position_gains() assert gains == {"kp": 35.0, "kd": 0.5} From 8614a86fd80697e4a07fbcb2f80c3e66454cb2fa Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Wed, 19 Aug 2026 00:13:48 +0800 Subject: [PATCH 070/169] refactor(tasks): move locomotion terrain spawn [skip ci] (#1186) --- src/unilab/tasks/locomotion/common/base.py | 2 +- .../{envs => tasks}/locomotion/common/terrain_spawn.py | 2 +- src/unilab/tasks/locomotion/go1/joystick.py | 8 ++++---- src/unilab/tasks/locomotion/go1/rough.py | 6 +++--- src/unilab/tasks/locomotion/go2/joystick.py | 8 ++++---- src/unilab/tasks/locomotion/go2w/rough.py | 8 ++++---- tests/envs/locomotion/go2w/test_go2w_motor_control.py | 2 +- tests/envs/locomotion/test_go2_terrain_spawn.py | 10 +++++----- tests/envs/locomotion/test_terrain_spawn.py | 2 +- 9 files changed, 24 insertions(+), 24 deletions(-) rename src/unilab/{envs => tasks}/locomotion/common/terrain_spawn.py (99%) diff --git a/src/unilab/tasks/locomotion/common/base.py b/src/unilab/tasks/locomotion/common/base.py index 41f4e81f0..68d82754b 100644 --- a/src/unilab/tasks/locomotion/common/base.py +++ b/src/unilab/tasks/locomotion/common/base.py @@ -12,7 +12,7 @@ from unilab.base.base import EnvCfg from unilab.base.np_env import NpEnv, NpEnvState from unilab.dtype_config import get_global_dtype -from unilab.envs.locomotion.common.terrain_spawn import BaseSpawnManager +from unilab.tasks.locomotion.common.terrain_spawn import BaseSpawnManager @dataclass diff --git a/src/unilab/envs/locomotion/common/terrain_spawn.py b/src/unilab/tasks/locomotion/common/terrain_spawn.py similarity index 99% rename from src/unilab/envs/locomotion/common/terrain_spawn.py rename to src/unilab/tasks/locomotion/common/terrain_spawn.py index cb623add2..2545fa772 100644 --- a/src/unilab/envs/locomotion/common/terrain_spawn.py +++ b/src/unilab/tasks/locomotion/common/terrain_spawn.py @@ -1,4 +1,4 @@ -"""Spawn-origin managers for locomotion envs. +"""Spawn-origin managers for locomotion tasks. ``BaseSpawnManager`` is a no-op default: every env spawns at the world origin (plus the existing per-env xy jitter from the dr_provider). Used whenever the diff --git a/src/unilab/tasks/locomotion/go1/joystick.py b/src/unilab/tasks/locomotion/go1/joystick.py index 58c022661..8445212fc 100644 --- a/src/unilab/tasks/locomotion/go1/joystick.py +++ b/src/unilab/tasks/locomotion/go1/joystick.py @@ -13,13 +13,13 @@ from unilab.dtype_config import get_global_dtype from unilab.envs.locomotion.common import rewards from unilab.envs.locomotion.common.rewards import RewardContext -from unilab.envs.locomotion.common.terrain_spawn import ( - TerrainCurriculumCfg, - TerrainSpawnManager, -) from unilab.tasks.locomotion.common.commands import Commands from unilab.tasks.locomotion.common.domain_rand import DomainRandConfig from unilab.tasks.locomotion.common.dr_provider import LocomotionDRProvider +from unilab.tasks.locomotion.common.terrain_spawn import ( + TerrainCurriculumCfg, + TerrainSpawnManager, +) from unilab.tasks.locomotion.go1.base import Go1BaseCfg, Go1BaseEnv diff --git a/src/unilab/tasks/locomotion/go1/rough.py b/src/unilab/tasks/locomotion/go1/rough.py index 1906533b9..d4c9edafc 100644 --- a/src/unilab/tasks/locomotion/go1/rough.py +++ b/src/unilab/tasks/locomotion/go1/rough.py @@ -16,9 +16,6 @@ from unilab.dtype_config import get_global_dtype from unilab.envs.locomotion.common import rewards from unilab.envs.locomotion.common.rewards import RewardContext -from unilab.envs.locomotion.common.terrain_spawn import ( - TerrainCurriculumCfg, -) from unilab.tasks.locomotion.common.commands import ( Commands, apply_heading_yaw_feedback, @@ -34,6 +31,9 @@ raw_height_scan_obs, terrain_out_of_bounds, ) +from unilab.tasks.locomotion.common.terrain_spawn import ( + TerrainCurriculumCfg, +) from unilab.tasks.locomotion.go1.base import ControlConfig from unilab.tasks.locomotion.go1.joystick import ( Go1JoystickCfg, diff --git a/src/unilab/tasks/locomotion/go2/joystick.py b/src/unilab/tasks/locomotion/go2/joystick.py index 5e489ed1d..7c7f7118d 100644 --- a/src/unilab/tasks/locomotion/go2/joystick.py +++ b/src/unilab/tasks/locomotion/go2/joystick.py @@ -13,10 +13,6 @@ from unilab.dtype_config import get_global_dtype from unilab.envs.locomotion.common import rewards from unilab.envs.locomotion.common.rewards import RewardContext -from unilab.envs.locomotion.common.terrain_spawn import ( - TerrainCurriculumCfg, - TerrainSpawnManager, -) from unilab.envs.manager_based_rl_env import ( ManagerBasedRlEnvCfg, make_manager_based_rl_env, @@ -25,6 +21,10 @@ from unilab.tasks.locomotion.common.commands import Commands from unilab.tasks.locomotion.common.domain_rand import DomainRandConfig from unilab.tasks.locomotion.common.dr_provider import LocomotionDRProvider +from unilab.tasks.locomotion.common.terrain_spawn import ( + TerrainCurriculumCfg, + TerrainSpawnManager, +) from unilab.tasks.locomotion.go2.base import Go2BaseCfg, Go2BaseEnv diff --git a/src/unilab/tasks/locomotion/go2w/rough.py b/src/unilab/tasks/locomotion/go2w/rough.py index 8caf28ad0..481c7220c 100644 --- a/src/unilab/tasks/locomotion/go2w/rough.py +++ b/src/unilab/tasks/locomotion/go2w/rough.py @@ -14,10 +14,6 @@ from unilab.dtype_config import get_global_dtype from unilab.envs.locomotion.common import rewards from unilab.envs.locomotion.common.rewards import RewardContext -from unilab.envs.locomotion.common.terrain_spawn import ( - TerrainCurriculumCfg, - TerrainSpawnManager, -) from unilab.tasks.locomotion.common.commands import ( Commands, apply_heading_yaw_feedback, @@ -31,6 +27,10 @@ raw_height_scan_obs, terrain_out_of_bounds, ) +from unilab.tasks.locomotion.common.terrain_spawn import ( + TerrainCurriculumCfg, + TerrainSpawnManager, +) from unilab.tasks.locomotion.go2w.base import NUM_GO2W_ACTIONS, NUM_LEG_ACTIONS from unilab.tasks.locomotion.go2w.joystick import ( Go2WJoystickCfg, diff --git a/tests/envs/locomotion/go2w/test_go2w_motor_control.py b/tests/envs/locomotion/go2w/test_go2w_motor_control.py index 96be2d5a5..fbd49b175 100644 --- a/tests/envs/locomotion/go2w/test_go2w_motor_control.py +++ b/tests/envs/locomotion/go2w/test_go2w_motor_control.py @@ -107,7 +107,7 @@ def test_go2w_backend_reset_randomization_excludes_kp_kd_payload() -> None: def test_go2w_reset_plan_can_disable_initial_yaw_randomization() -> None: - from unilab.envs.locomotion.common.terrain_spawn import BaseSpawnManager + from unilab.tasks.locomotion.common.terrain_spawn import BaseSpawnManager cfg = Go2WJoystickCfg(reward_config=_reward_config()) cfg.domain_rand.randomize_init_yaw = False diff --git a/tests/envs/locomotion/test_go2_terrain_spawn.py b/tests/envs/locomotion/test_go2_terrain_spawn.py index 5d357ed8e..b977e8243 100644 --- a/tests/envs/locomotion/test_go2_terrain_spawn.py +++ b/tests/envs/locomotion/test_go2_terrain_spawn.py @@ -7,7 +7,7 @@ import numpy as np import pytest -from unilab.envs.locomotion.common.terrain_spawn import ( +from unilab.tasks.locomotion.common.terrain_spawn import ( TerrainCurriculumCfg, TerrainSpawnManager, ) @@ -52,7 +52,7 @@ def test_terrain_spawn_attached_when_rough(): assert terrain_data.terrain_origins.shape == (3, 3, 3) assert not terrain_data.terrain_origins.flags.writeable assert _class_path(env._spawn) == ( - "unilab.envs.locomotion.common.terrain_spawn.TerrainSpawnManager" + "unilab.tasks.locomotion.common.terrain_spawn.TerrainSpawnManager" ) assert env._scene_terrain_origins is not None assert env._scene_terrain_origins.shape == (3, 3, 3) @@ -76,7 +76,7 @@ def test_terrain_spawn_attached_when_rough_motrix(): assert terrain_data.terrain_origins.shape == (3, 3, 3) assert not terrain_data.terrain_origins.flags.writeable assert _class_path(env._spawn) == ( - "unilab.envs.locomotion.common.terrain_spawn.TerrainSpawnManager" + "unilab.tasks.locomotion.common.terrain_spawn.TerrainSpawnManager" ) assert env._scene_terrain_origins is not None assert env._scene_terrain_origins.shape == (3, 3, 3) @@ -240,7 +240,7 @@ def test_default_spawn_used_when_flat(): try: assert env._backend.get_terrain_spawn_data() is None assert _class_path(env._spawn) == ( - "unilab.envs.locomotion.common.terrain_spawn.BaseSpawnManager" + "unilab.tasks.locomotion.common.terrain_spawn.BaseSpawnManager" ) assert env._scene_terrain_origins is None # Origins are zeros (flat scene needs no spread; per-env xy jitter still applies). @@ -359,7 +359,7 @@ def test_episode_start_recorded_after_reset(preset): try: env.init_state() sm = env._spawn - if _class_path(sm) == "unilab.envs.locomotion.common.terrain_spawn.TerrainSpawnManager": + if _class_path(sm) == "unilab.tasks.locomotion.common.terrain_spawn.TerrainSpawnManager": assert np.all(sm._has_started) finally: env.close() diff --git a/tests/envs/locomotion/test_terrain_spawn.py b/tests/envs/locomotion/test_terrain_spawn.py index 249b2d174..3e1bc60f8 100644 --- a/tests/envs/locomotion/test_terrain_spawn.py +++ b/tests/envs/locomotion/test_terrain_spawn.py @@ -5,7 +5,7 @@ import numpy as np import pytest -from unilab.envs.locomotion.common.terrain_spawn import ( +from unilab.tasks.locomotion.common.terrain_spawn import ( TerrainCurriculumCfg, TerrainSpawnManager, ) From fd366c37505be812f7507e40cf1f6909a357cbb2 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Wed, 19 Aug 2026 00:19:15 +0800 Subject: [PATCH 071/169] refactor(tasks): finish locomotion common rewards [skip ci] (#1188) --- .../3-deployment/3-framework_migration/6-reward_porting.md | 6 +++--- .../3-deployment/3-framework_migration/6-reward_porting.md | 6 +++--- src/unilab/envs/locomotion/common/__init__.py | 5 ----- src/unilab/tasks/locomotion/a2/joystick.py | 4 ++-- src/unilab/tasks/locomotion/common/__init__.py | 2 ++ src/unilab/{envs => tasks}/locomotion/common/rewards.py | 2 +- src/unilab/tasks/locomotion/g1/joystick.py | 4 ++-- src/unilab/tasks/locomotion/go1/joystick.py | 4 ++-- src/unilab/tasks/locomotion/go1/rough.py | 4 ++-- src/unilab/tasks/locomotion/go2/footstand.py | 4 ++-- src/unilab/tasks/locomotion/go2/joystick.py | 4 ++-- src/unilab/tasks/locomotion/go2/rough.py | 4 ++-- src/unilab/tasks/locomotion/go2_arm/manip_loco.py | 4 ++-- src/unilab/tasks/locomotion/go2w/joystick.py | 4 ++-- src/unilab/tasks/locomotion/go2w/rough.py | 4 ++-- src/unilab/tasks/motion_tracking/common/rewards.py | 2 +- tests/envs/locomotion/a2/test_a2_joystick_contract.py | 2 +- tests/envs/locomotion/go2_arm/test_manip_loco_contract.py | 2 +- tests/envs/locomotion/test_go2_footstand.py | 2 +- tests/envs/test_env_configs.py | 2 +- 20 files changed, 34 insertions(+), 37 deletions(-) delete mode 100644 src/unilab/envs/locomotion/common/__init__.py rename src/unilab/{envs => tasks}/locomotion/common/rewards.py (99%) diff --git a/docs/sphinx/source/en/3-deployment/3-framework_migration/6-reward_porting.md b/docs/sphinx/source/en/3-deployment/3-framework_migration/6-reward_porting.md index 72e11d3a3..78cfff3d9 100644 --- a/docs/sphinx/source/en/3-deployment/3-framework_migration/6-reward_porting.md +++ b/docs/sphinx/source/en/3-deployment/3-framework_migration/6-reward_porting.md @@ -38,7 +38,7 @@ Notes: - UniLab's `state` carries `prev_contact` so you don't need to manage edge detection yourself. See - `unilab.envs.locomotion.common.rewards`. + `unilab.tasks.locomotion.common.rewards`. ## Pattern: action smoothness penalty @@ -47,7 +47,7 @@ def reward_action_rate(self, state): return -np.sum((state.action - state.prev_action) ** 2, axis=1) ``` -Already a stock helper in `unilab.envs.locomotion.common.rewards`. +Already a stock helper in `unilab.tasks.locomotion.common.rewards`. ## Pattern: posture penalty @@ -77,4 +77,4 @@ def reward_termination(self, state): - {doc}`5-task_config_translation` - `unilab.training.reward` -- `unilab.envs.locomotion.common.rewards` +- `unilab.tasks.locomotion.common.rewards` diff --git a/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/6-reward_porting.md b/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/6-reward_porting.md index a86136ec9..3b4104df6 100644 --- a/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/6-reward_porting.md +++ b/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/6-reward_porting.md @@ -37,7 +37,7 @@ def reward_feet_air_time(self, state): 注意： - UniLab 的 `state` 携带了 `prev_contact`，因此你无需自己管理边沿检测。参见 - `unilab.envs.locomotion.common.rewards`。 + `unilab.tasks.locomotion.common.rewards`。 ## 模式：动作平滑惩罚 @@ -46,7 +46,7 @@ def reward_action_rate(self, state): return -np.sum((state.action - state.prev_action) ** 2, axis=1) ``` -它已经是 `unilab.envs.locomotion.common.rewards` 中的现成辅助函数。 +它已经是 `unilab.tasks.locomotion.common.rewards` 中的现成辅助函数。 ## 模式：姿态惩罚 @@ -75,4 +75,4 @@ def reward_termination(self, state): - {doc}`5-task_config_translation` - `unilab.training.reward` -- `unilab.envs.locomotion.common.rewards` +- `unilab.tasks.locomotion.common.rewards` diff --git a/src/unilab/envs/locomotion/common/__init__.py b/src/unilab/envs/locomotion/common/__init__.py deleted file mode 100644 index 05d89eb00..000000000 --- a/src/unilab/envs/locomotion/common/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -from .rewards import RewardContext - -__all__ = [ - "RewardContext", -] diff --git a/src/unilab/tasks/locomotion/a2/joystick.py b/src/unilab/tasks/locomotion/a2/joystick.py index dd6f45c68..88a2672e4 100644 --- a/src/unilab/tasks/locomotion/a2/joystick.py +++ b/src/unilab/tasks/locomotion/a2/joystick.py @@ -23,10 +23,10 @@ from unilab.base import registry from unilab.base.scene import SceneCfg from unilab.dtype_config import get_global_dtype -from unilab.envs.locomotion.common import rewards -from unilab.envs.locomotion.common.rewards import RewardContext +from unilab.tasks.locomotion.common import rewards from unilab.tasks.locomotion.common.commands import sample_commands_with_standing from unilab.tasks.locomotion.common.dr_provider import LocomotionDRProvider +from unilab.tasks.locomotion.common.rewards import RewardContext from unilab.tasks.locomotion.go2.base import Asset, ControlConfig from unilab.tasks.locomotion.go2.joystick import ( Go2DomainRandConfig, diff --git a/src/unilab/tasks/locomotion/common/__init__.py b/src/unilab/tasks/locomotion/common/__init__.py index e21351fcf..cc78aac1f 100644 --- a/src/unilab/tasks/locomotion/common/__init__.py +++ b/src/unilab/tasks/locomotion/common/__init__.py @@ -22,6 +22,7 @@ DEFAULT_SCAN_POINTS_Y, HeightScanConfig, ) +from .rewards import RewardContext __all__ = [ "BaseNoiseConfig", @@ -35,6 +36,7 @@ "LocomotionBaseCfg", "LocomotionBaseEnv", "PdControlConfig", + "RewardContext", "Sensor", "apply_heading_yaw_feedback", "sample_heading_commands", diff --git a/src/unilab/envs/locomotion/common/rewards.py b/src/unilab/tasks/locomotion/common/rewards.py similarity index 99% rename from src/unilab/envs/locomotion/common/rewards.py rename to src/unilab/tasks/locomotion/common/rewards.py index 77e56df50..a147a52be 100644 --- a/src/unilab/envs/locomotion/common/rewards.py +++ b/src/unilab/tasks/locomotion/common/rewards.py @@ -1,4 +1,4 @@ -"""Shared reward functions for locomotion environments. +"""Shared reward functions for locomotion tasks. Introduces ``RewardContext`` — a dataclass that bundles all state any reward function might need. Shared reward functions are plain diff --git a/src/unilab/tasks/locomotion/g1/joystick.py b/src/unilab/tasks/locomotion/g1/joystick.py index bd715ffd6..0191ca44b 100644 --- a/src/unilab/tasks/locomotion/g1/joystick.py +++ b/src/unilab/tasks/locomotion/g1/joystick.py @@ -16,8 +16,7 @@ from unilab.base.np_env import NpEnvState from unilab.base.scene import SceneCfg from unilab.dtype_config import get_global_dtype -from unilab.envs.locomotion.common import rewards -from unilab.envs.locomotion.common.rewards import RewardContext +from unilab.tasks.locomotion.common import rewards from unilab.tasks.locomotion.common.commands import ( Commands, sample_heading_commands, @@ -25,6 +24,7 @@ ) from unilab.tasks.locomotion.common.domain_rand import DomainRandConfig from unilab.tasks.locomotion.common.dr_provider import LocomotionDRProvider +from unilab.tasks.locomotion.common.rewards import RewardContext from .base import G1BaseCfg, G1BaseEnv diff --git a/src/unilab/tasks/locomotion/go1/joystick.py b/src/unilab/tasks/locomotion/go1/joystick.py index 8445212fc..e36089181 100644 --- a/src/unilab/tasks/locomotion/go1/joystick.py +++ b/src/unilab/tasks/locomotion/go1/joystick.py @@ -11,11 +11,11 @@ from unilab.base.np_env import NpEnvState from unilab.base.scene import SceneCfg from unilab.dtype_config import get_global_dtype -from unilab.envs.locomotion.common import rewards -from unilab.envs.locomotion.common.rewards import RewardContext +from unilab.tasks.locomotion.common import rewards from unilab.tasks.locomotion.common.commands import Commands from unilab.tasks.locomotion.common.domain_rand import DomainRandConfig from unilab.tasks.locomotion.common.dr_provider import LocomotionDRProvider +from unilab.tasks.locomotion.common.rewards import RewardContext from unilab.tasks.locomotion.common.terrain_spawn import ( TerrainCurriculumCfg, TerrainSpawnManager, diff --git a/src/unilab/tasks/locomotion/go1/rough.py b/src/unilab/tasks/locomotion/go1/rough.py index d4c9edafc..3a11d17c9 100644 --- a/src/unilab/tasks/locomotion/go1/rough.py +++ b/src/unilab/tasks/locomotion/go1/rough.py @@ -14,8 +14,7 @@ from unilab.dr import DomainRandomizationManager, ResetPlan from unilab.dr.dr_utils import build_common_reset_randomization, zero_actions from unilab.dtype_config import get_global_dtype -from unilab.envs.locomotion.common import rewards -from unilab.envs.locomotion.common.rewards import RewardContext +from unilab.tasks.locomotion.common import rewards from unilab.tasks.locomotion.common.commands import ( Commands, apply_heading_yaw_feedback, @@ -31,6 +30,7 @@ raw_height_scan_obs, terrain_out_of_bounds, ) +from unilab.tasks.locomotion.common.rewards import RewardContext from unilab.tasks.locomotion.common.terrain_spawn import ( TerrainCurriculumCfg, ) diff --git a/src/unilab/tasks/locomotion/go2/footstand.py b/src/unilab/tasks/locomotion/go2/footstand.py index e4d0b6adc..9bd06e914 100644 --- a/src/unilab/tasks/locomotion/go2/footstand.py +++ b/src/unilab/tasks/locomotion/go2/footstand.py @@ -12,11 +12,11 @@ from unilab.base.scene import SceneCfg from unilab.dr import ResetPlan, ResetRandomizationPayload from unilab.dtype_config import get_global_dtype -from unilab.envs.locomotion.common import rewards -from unilab.envs.locomotion.common.rewards import RewardContext +from unilab.tasks.locomotion.common import rewards from unilab.tasks.locomotion.common.commands import Commands from unilab.tasks.locomotion.common.domain_rand import DomainRandConfig from unilab.tasks.locomotion.common.dr_provider import LocomotionDRProvider +from unilab.tasks.locomotion.common.rewards import RewardContext from unilab.tasks.locomotion.go2.base import ControlConfig, Go2BaseCfg, Go2BaseEnv, NoiseConfig from unilab.utils.rotation import np_quat_apply, np_quat_apply_inverse diff --git a/src/unilab/tasks/locomotion/go2/joystick.py b/src/unilab/tasks/locomotion/go2/joystick.py index 7c7f7118d..8009226ca 100644 --- a/src/unilab/tasks/locomotion/go2/joystick.py +++ b/src/unilab/tasks/locomotion/go2/joystick.py @@ -11,16 +11,16 @@ from unilab.base.np_env import NpEnvState from unilab.base.scene import SceneCfg from unilab.dtype_config import get_global_dtype -from unilab.envs.locomotion.common import rewards -from unilab.envs.locomotion.common.rewards import RewardContext from unilab.envs.manager_based_rl_env import ( ManagerBasedRlEnvCfg, make_manager_based_rl_env, ) +from unilab.tasks.locomotion.common import rewards from unilab.tasks.locomotion.common.base import Sensor from unilab.tasks.locomotion.common.commands import Commands from unilab.tasks.locomotion.common.domain_rand import DomainRandConfig from unilab.tasks.locomotion.common.dr_provider import LocomotionDRProvider +from unilab.tasks.locomotion.common.rewards import RewardContext from unilab.tasks.locomotion.common.terrain_spawn import ( TerrainCurriculumCfg, TerrainSpawnManager, diff --git a/src/unilab/tasks/locomotion/go2/rough.py b/src/unilab/tasks/locomotion/go2/rough.py index ad675a8eb..d8a2dedf7 100644 --- a/src/unilab/tasks/locomotion/go2/rough.py +++ b/src/unilab/tasks/locomotion/go2/rough.py @@ -12,8 +12,7 @@ from unilab.dr import DomainRandomizationManager, ResetPlan from unilab.dr.dr_utils import build_common_reset_randomization, zero_actions from unilab.dtype_config import get_global_dtype -from unilab.envs.locomotion.common import rewards -from unilab.envs.locomotion.common.rewards import RewardContext +from unilab.tasks.locomotion.common import rewards from unilab.tasks.locomotion.common.commands import ( apply_heading_yaw_feedback, sample_heading_commands, @@ -29,6 +28,7 @@ raw_height_scan_obs, terrain_out_of_bounds, ) +from unilab.tasks.locomotion.common.rewards import RewardContext from unilab.tasks.locomotion.go2.base import ControlConfig from unilab.tasks.locomotion.go2.joystick import ( Commands, diff --git a/src/unilab/tasks/locomotion/go2_arm/manip_loco.py b/src/unilab/tasks/locomotion/go2_arm/manip_loco.py index 0a1968863..ec961600d 100644 --- a/src/unilab/tasks/locomotion/go2_arm/manip_loco.py +++ b/src/unilab/tasks/locomotion/go2_arm/manip_loco.py @@ -12,11 +12,11 @@ from unilab.base.scene import SceneCfg from unilab.dr.types import ResetPlan from unilab.dtype_config import get_global_dtype -from unilab.envs.locomotion.common import rewards -from unilab.envs.locomotion.common.rewards import RewardContext +from unilab.tasks.locomotion.common import rewards from unilab.tasks.locomotion.common.commands import Commands from unilab.tasks.locomotion.common.domain_rand import DomainRandConfig from unilab.tasks.locomotion.common.dr_provider import LocomotionDRProvider +from unilab.tasks.locomotion.common.rewards import RewardContext from unilab.tasks.locomotion.go2_arm.base import ( Go2ArmBaseCfg, Go2ArmBaseEnv, diff --git a/src/unilab/tasks/locomotion/go2w/joystick.py b/src/unilab/tasks/locomotion/go2w/joystick.py index 82a8715b0..c27266975 100644 --- a/src/unilab/tasks/locomotion/go2w/joystick.py +++ b/src/unilab/tasks/locomotion/go2w/joystick.py @@ -17,8 +17,7 @@ zero_actions, ) from unilab.dtype_config import get_global_dtype -from unilab.envs.locomotion.common import rewards -from unilab.envs.locomotion.common.rewards import RewardContext +from unilab.tasks.locomotion.common import rewards from unilab.tasks.locomotion.common.commands import ( Commands, apply_heading_yaw_feedback, @@ -29,6 +28,7 @@ ) from unilab.tasks.locomotion.common.domain_rand import DomainRandConfig from unilab.tasks.locomotion.common.dr_provider import LocomotionDRProvider +from unilab.tasks.locomotion.common.rewards import RewardContext from unilab.tasks.locomotion.go2w.base import ( DEFAULT_GO2W_ANGLES, NUM_GO2W_ACTIONS, diff --git a/src/unilab/tasks/locomotion/go2w/rough.py b/src/unilab/tasks/locomotion/go2w/rough.py index 481c7220c..5c9151ba4 100644 --- a/src/unilab/tasks/locomotion/go2w/rough.py +++ b/src/unilab/tasks/locomotion/go2w/rough.py @@ -12,8 +12,7 @@ from unilab.dr import DomainRandomizationManager, ResetPlan from unilab.dr.dr_utils import zero_actions from unilab.dtype_config import get_global_dtype -from unilab.envs.locomotion.common import rewards -from unilab.envs.locomotion.common.rewards import RewardContext +from unilab.tasks.locomotion.common import rewards from unilab.tasks.locomotion.common.commands import ( Commands, apply_heading_yaw_feedback, @@ -27,6 +26,7 @@ raw_height_scan_obs, terrain_out_of_bounds, ) +from unilab.tasks.locomotion.common.rewards import RewardContext from unilab.tasks.locomotion.common.terrain_spawn import ( TerrainCurriculumCfg, TerrainSpawnManager, diff --git a/src/unilab/tasks/motion_tracking/common/rewards.py b/src/unilab/tasks/motion_tracking/common/rewards.py index ffe59a191..36a169c02 100644 --- a/src/unilab/tasks/motion_tracking/common/rewards.py +++ b/src/unilab/tasks/motion_tracking/common/rewards.py @@ -1,7 +1,7 @@ """Shared reward configuration and functions for motion tracking. Reward terms are plain module-level callables ``fn(ctx: RewardContext) -> np.ndarray`` -mirroring :mod:`unilab.envs.locomotion.common.rewards`. Robot-specific terms that +mirroring :mod:`unilab.tasks.locomotion.common.rewards`. Robot-specific terms that live on env subclasses (box-object / joint-effort terms) are stored in the same ``_reward_fns`` dispatch table as bound methods and are called with the same ``ctx`` argument. diff --git a/tests/envs/locomotion/a2/test_a2_joystick_contract.py b/tests/envs/locomotion/a2/test_a2_joystick_contract.py index c81b5dbb2..eaa224dd5 100644 --- a/tests/envs/locomotion/a2/test_a2_joystick_contract.py +++ b/tests/envs/locomotion/a2/test_a2_joystick_contract.py @@ -440,7 +440,7 @@ def test_a2_cfg_reward_config_annotation_is_a2_type(): def _a2_ctx(commands, dof_pos=None): - from unilab.envs.locomotion.common.rewards import RewardContext + from unilab.tasks.locomotion.common.rewards import RewardContext n = commands.shape[0] return RewardContext( diff --git a/tests/envs/locomotion/go2_arm/test_manip_loco_contract.py b/tests/envs/locomotion/go2_arm/test_manip_loco_contract.py index b083b9345..b36b35bb5 100644 --- a/tests/envs/locomotion/go2_arm/test_manip_loco_contract.py +++ b/tests/envs/locomotion/go2_arm/test_manip_loco_contract.py @@ -183,7 +183,7 @@ def test_go2_arm_command_postprocess_can_force_zero_commands(): def test_go2_arm_stand_still_reward_uses_same_command_mask(): """stand_still should not penalize leg pose under lateral, yaw, or forward commands.""" - from unilab.envs.locomotion.common.rewards import RewardContext + from unilab.tasks.locomotion.common.rewards import RewardContext from unilab.tasks.locomotion.go2_arm.manip_loco import Go2ArmManipLocoEnv env = object.__new__(Go2ArmManipLocoEnv) diff --git a/tests/envs/locomotion/test_go2_footstand.py b/tests/envs/locomotion/test_go2_footstand.py index ab4356a8b..f58b9e2ee 100644 --- a/tests/envs/locomotion/test_go2_footstand.py +++ b/tests/envs/locomotion/test_go2_footstand.py @@ -7,7 +7,7 @@ from unilab.base.np_env import NpEnvState from unilab.base.registry import ensure_registries from unilab.dr import ResetRandomizationPayload -from unilab.envs.locomotion.common.rewards import RewardContext +from unilab.tasks.locomotion.common.rewards import RewardContext from unilab.tasks.locomotion.go2.footstand import ( FootstandControlConfig, FootstandSensor, diff --git a/tests/envs/test_env_configs.py b/tests/envs/test_env_configs.py index 37f7a6ed3..2f3776974 100644 --- a/tests/envs/test_env_configs.py +++ b/tests/envs/test_env_configs.py @@ -243,7 +243,7 @@ def test_g1_walk_env_reward_dispatch_restores_motrix_terms(): def test_g1_walk_env_feet_phase_reward_is_gated_by_forward_speed(): - from unilab.envs.locomotion.common.rewards import RewardContext + from unilab.tasks.locomotion.common.rewards import RewardContext from unilab.tasks.locomotion.g1.joystick import G1WalkEnv class FakeBackend: From e8f6782a805c737c6a04125003b9dc880cd8658e Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Wed, 19 Aug 2026 00:24:33 +0800 Subject: [PATCH 072/169] refactor(tasks): remove legacy env task packages [skip ci] (#1190) --- src/unilab/envs/locomotion/__init__.py | 3 --- src/unilab/envs/manipulation/__init__.py | 3 --- tests/tasks/test_package_boundary.py | 11 +++++++++++ tests/utils/test_algo_utils.py | 2 +- 4 files changed, 12 insertions(+), 7 deletions(-) delete mode 100644 src/unilab/envs/locomotion/__init__.py delete mode 100644 src/unilab/envs/manipulation/__init__.py diff --git a/src/unilab/envs/locomotion/__init__.py b/src/unilab/envs/locomotion/__init__.py deleted file mode 100644 index 11c9f906a..000000000 --- a/src/unilab/envs/locomotion/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -"""Generic locomotion environment components.""" - -__unilab_registry_modules__: tuple[str, ...] = () diff --git a/src/unilab/envs/manipulation/__init__.py b/src/unilab/envs/manipulation/__init__.py deleted file mode 100644 index 15c01b96c..000000000 --- a/src/unilab/envs/manipulation/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -"""Generic manipulation environment components.""" - -__unilab_registry_modules__: tuple[str, ...] = () diff --git a/tests/tasks/test_package_boundary.py b/tests/tasks/test_package_boundary.py index c75ba15c8..38cfb78b1 100644 --- a/tests/tasks/test_package_boundary.py +++ b/tests/tasks/test_package_boundary.py @@ -10,6 +10,7 @@ _REPO_ROOT = Path(__file__).resolve().parents[2] _ENV_PACKAGE = _REPO_ROOT / "src" / "unilab" / "envs" +_CONCRETE_TASK_PACKAGES = ("locomotion", "manipulation", "motion_tracking") _TASK_REGISTRY_MODULES = ( "unilab.tasks.locomotion.go1", @@ -51,3 +52,13 @@ def test_env_runtime_does_not_depend_on_tasks() -> None: ] assert violations == [], "unilab.envs must not import concrete unilab.tasks modules" + + +def test_env_runtime_does_not_own_concrete_task_packages() -> None: + violations = [ + path.relative_to(_REPO_ROOT).as_posix() + for package in _CONCRETE_TASK_PACKAGES + for path in sorted((_ENV_PACKAGE / package).rglob("*.py")) + ] + + assert violations == [], "concrete task source must be owned by unilab.tasks" diff --git a/tests/utils/test_algo_utils.py b/tests/utils/test_algo_utils.py index 9d80d5b52..bb4ba181b 100644 --- a/tests/utils/test_algo_utils.py +++ b/tests/utils/test_algo_utils.py @@ -53,7 +53,7 @@ def test_mixed_optional_and_required(self, caplog) -> None: # Use real package + fake optional package with caplog.at_level(logging.WARNING): ensure_registries( - ["unilab.envs.locomotion", "nonexistent_optional_12345"], + ["unilab.tasks.locomotion", "nonexistent_optional_12345"], optional_packages=["nonexistent_optional_12345"], ) From 763859a7efcc17550239b244ee35ceeac5d99bea Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Wed, 19 Aug 2026 00:33:57 +0800 Subject: [PATCH 073/169] docs(api): split env runtime and task references [skip ci] (#1192) --- .../sphinx/source/api_reference/envs/index.md | 22 +++++---------- .../api_reference/envs/motion_tracking.md | 13 --------- docs/sphinx/source/api_reference/index.md | 9 ++++++- .../source/api_reference/tasks/index.md | 27 +++++++++++++++++++ .../{envs => tasks}/locomotion.md | 5 ++-- .../{envs => tasks}/manipulation.md | 0 .../api_reference/tasks/motion_tracking.md | 14 ++++++++++ 7 files changed, 58 insertions(+), 32 deletions(-) delete mode 100644 docs/sphinx/source/api_reference/envs/motion_tracking.md create mode 100644 docs/sphinx/source/api_reference/tasks/index.md rename docs/sphinx/source/api_reference/{envs => tasks}/locomotion.md (74%) rename docs/sphinx/source/api_reference/{envs => tasks}/manipulation.md (100%) create mode 100644 docs/sphinx/source/api_reference/tasks/motion_tracking.md diff --git a/docs/sphinx/source/api_reference/envs/index.md b/docs/sphinx/source/api_reference/envs/index.md index 6665cd124..f7eb11999 100644 --- a/docs/sphinx/source/api_reference/envs/index.md +++ b/docs/sphinx/source/api_reference/envs/index.md @@ -1,21 +1,11 @@ -# `unilab.envs` — Tasks +# `unilab.envs` — Environment runtime -Concrete RL tasks split by family: +Task-agnostic Manager-Based environment runtime and reusable MDP terms. +Concrete task implementations are owned by {doc}`../tasks/index`. -- **locomotion** — Go1, Go2, Go2w, Go2 + Airbot, Unitree G1 -- **manipulation** — Allegro / Sharpa in-hand cube -- **motion_tracking** — G1 whole-body motion tracking + flips - -Every env inherits `NpEnv` and is registered into the task `Registry` so it -can be selected via `uv run train --algo --task --sim `. - -```{toctree} -:maxdepth: 2 - -locomotion -manipulation -motion_tracking -``` +`ManagerBasedRLEnv` preserves UniLab's NumPy `NpEnv` contract while executing +community-style action, observation, reward, termination, event, command, and +curriculum managers. ```{eval-rst} .. autosummary:: diff --git a/docs/sphinx/source/api_reference/envs/motion_tracking.md b/docs/sphinx/source/api_reference/envs/motion_tracking.md deleted file mode 100644 index c4371bd6b..000000000 --- a/docs/sphinx/source/api_reference/envs/motion_tracking.md +++ /dev/null @@ -1,13 +0,0 @@ -# `unilab.envs.motion_tracking` - -Whole-body motion tracking tasks. G1 humanoid currently ships flip tracking -plus general motion tracking (PPO + SAC variants). - -```{eval-rst} -.. autosummary:: - :toctree: _autosummary - :template: autosummary/module.rst - :recursive: - - unilab.envs.motion_tracking.g1 -``` diff --git a/docs/sphinx/source/api_reference/index.md b/docs/sphinx/source/api_reference/index.md index df6920ca1..fdbed0fb2 100644 --- a/docs/sphinx/source/api_reference/index.md +++ b/docs/sphinx/source/api_reference/index.md @@ -30,10 +30,16 @@ The contracts everything else depends on: `NpEnv`, `SimBackend`, `Registry`, :::{grid-item-card} 🧪 `unilab.envs` :link: envs/index :link-type: doc -Concrete tasks — locomotion, manipulation, motion tracking — layered on +Manager-Based environment runtime and task-agnostic MDP terms layered on top of `base`. ::: +:::{grid-item-card} 🤖 `unilab.tasks` +:link: tasks/index +:link-type: doc +Concrete locomotion, manipulation, and motion-tracking task packages. +::: + :::: ## Learning stack @@ -125,6 +131,7 @@ top_level base/index envs/index +tasks/index ``` ```{toctree} diff --git a/docs/sphinx/source/api_reference/tasks/index.md b/docs/sphinx/source/api_reference/tasks/index.md new file mode 100644 index 000000000..eb4ed03d1 --- /dev/null +++ b/docs/sphinx/source/api_reference/tasks/index.md @@ -0,0 +1,27 @@ +# `unilab.tasks` — Concrete tasks + +Concrete RL tasks split by family: + +- **locomotion** — A2, Go1, Go2, Go2w, Go2 + Airbot, and Unitree G1 +- **manipulation** — Allegro / Sharpa in-hand cube and Stewart balance +- **motion_tracking** — G1 and X2 whole-body motion tracking + +Every task is registered into the task `Registry` so it can be selected via +`uv run train --algo --task --sim `. + +```{toctree} +:maxdepth: 2 + +locomotion +manipulation +motion_tracking +``` + +```{eval-rst} +.. autosummary:: + :toctree: _autosummary + :template: autosummary/module.rst + :recursive: + + unilab.tasks +``` diff --git a/docs/sphinx/source/api_reference/envs/locomotion.md b/docs/sphinx/source/api_reference/tasks/locomotion.md similarity index 74% rename from docs/sphinx/source/api_reference/envs/locomotion.md rename to docs/sphinx/source/api_reference/tasks/locomotion.md index cf185080b..c4e504491 100644 --- a/docs/sphinx/source/api_reference/envs/locomotion.md +++ b/docs/sphinx/source/api_reference/tasks/locomotion.md @@ -1,4 +1,4 @@ -# `unilab.envs.locomotion` +# `unilab.tasks.locomotion` ```{eval-rst} .. autosummary:: @@ -6,7 +6,8 @@ :template: autosummary/module.rst :recursive: - unilab.envs.locomotion.common + unilab.tasks.locomotion.a2 + unilab.tasks.locomotion.common unilab.tasks.locomotion.g1 unilab.tasks.locomotion.go1 unilab.tasks.locomotion.go2 diff --git a/docs/sphinx/source/api_reference/envs/manipulation.md b/docs/sphinx/source/api_reference/tasks/manipulation.md similarity index 100% rename from docs/sphinx/source/api_reference/envs/manipulation.md rename to docs/sphinx/source/api_reference/tasks/manipulation.md diff --git a/docs/sphinx/source/api_reference/tasks/motion_tracking.md b/docs/sphinx/source/api_reference/tasks/motion_tracking.md new file mode 100644 index 000000000..8ca1f557d --- /dev/null +++ b/docs/sphinx/source/api_reference/tasks/motion_tracking.md @@ -0,0 +1,14 @@ +# `unilab.tasks.motion_tracking` + +Whole-body motion-tracking tasks for G1 and X2 robots. + +```{eval-rst} +.. autosummary:: + :toctree: _autosummary + :template: autosummary/module.rst + :recursive: + + unilab.tasks.motion_tracking.common + unilab.tasks.motion_tracking.g1 + unilab.tasks.motion_tracking.x2 +``` From d8e6657dbe82d1e5244e0da32322e8cbd8fd1700 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Wed, 19 Aug 2026 00:38:01 +0800 Subject: [PATCH 074/169] docs(tasks): update locomotion owner paths [skip ci] (#1194) --- docs/sphinx/source/en/2-user_guide/4-tasks/1-locomotion.md | 6 +++--- .../5-domain_randomization/2-writing_providers.md | 4 ++-- .../en/2-user_guide/6-terrain/2-heightfield_import.md | 2 +- .../source/en/2-user_guide/7-tooling/5-robot_import.md | 2 +- .../3-deployment/3-framework_migration/2-from_legged_gym.md | 2 +- .../4-developer_guide/1-architecture/4-scene_composition.md | 2 +- .../source/en/4-developer_guide/3-extending/1-new_task.md | 2 +- .../source/zh_CN/2-user_guide/4-tasks/1-locomotion.md | 6 +++--- .../5-domain_randomization/2-writing_providers.md | 4 ++-- .../zh_CN/2-user_guide/6-terrain/2-heightfield_import.md | 2 +- .../source/zh_CN/2-user_guide/7-tooling/5-robot_import.md | 2 +- .../3-deployment/3-framework_migration/2-from_legged_gym.md | 2 +- .../4-developer_guide/1-architecture/4-scene_composition.md | 2 +- .../zh_CN/4-developer_guide/3-extending/1-new_task.md | 2 +- 14 files changed, 20 insertions(+), 20 deletions(-) diff --git a/docs/sphinx/source/en/2-user_guide/4-tasks/1-locomotion.md b/docs/sphinx/source/en/2-user_guide/4-tasks/1-locomotion.md index 94818f8e7..2a786ce61 100644 --- a/docs/sphinx/source/en/2-user_guide/4-tasks/1-locomotion.md +++ b/docs/sphinx/source/en/2-user_guide/4-tasks/1-locomotion.md @@ -1,7 +1,7 @@ # Locomotion -Locomotion tasks are registered in `src/unilab/envs/locomotion/` and -`src/unilab/envs/motion_tracking/`. The available owner YAMLs under `conf/` +Locomotion tasks are registered in `src/unilab/tasks/locomotion/` and +`src/unilab/tasks/motion_tracking/`. The available owner YAMLs under `conf/` define which algorithm and backend combinations are runnable. ## Families @@ -33,7 +33,7 @@ backend: {doc}`../../5-reference/5-support_matrix`. - PPO config: `conf/ppo/task/go2_footstand/mujoco.yaml` - Registered env: `Go2FootStand` (registered for `sim_backend="mujoco"`) -- Implementation: `src/unilab/envs/locomotion/go2/footstand.py` +- Implementation: `src/unilab/tasks/locomotion/go2/footstand.py` (extends the Go2 base task) - Go2 model XML: `src/unilab/assets/robots/go2/go2.xml` 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 1491c0ea6..a0fc87009 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 @@ -27,9 +27,9 @@ The shared types live in `src/unilab/dr/types.py`, and the manager lives in Representative provider implementations are in: -- `src/unilab/envs/locomotion/go1/joystick.py` +- `src/unilab/tasks/locomotion/go1/joystick.py` - `src/unilab/tasks/locomotion/g1/joystick.py` -- `src/unilab/envs/motion_tracking/g1/tracking.py` +- `src/unilab/tasks/motion_tracking/g1/tracking.py` - `src/unilab/tasks/manipulation/allegro_inhand/rotation.py` - `src/unilab/tasks/manipulation/sharpa_inhand/rotation.py` diff --git a/docs/sphinx/source/en/2-user_guide/6-terrain/2-heightfield_import.md b/docs/sphinx/source/en/2-user_guide/6-terrain/2-heightfield_import.md index ef5d80fa6..32a17a7d0 100644 --- a/docs/sphinx/source/en/2-user_guide/6-terrain/2-heightfield_import.md +++ b/docs/sphinx/source/en/2-user_guide/6-terrain/2-heightfield_import.md @@ -10,7 +10,7 @@ example is `Go2JoystickRough`, with owners in - `src/unilab/terrains/heightfield_terrains.py` - `src/unilab/terrains/terrain_generator.py` -- `src/unilab/envs/locomotion/go2/rough.py` +- `src/unilab/tasks/locomotion/go2/rough.py` - `src/unilab/base/backend/mujoco/xml.py` - `src/unilab/base/backend/motrix/scene.py` diff --git a/docs/sphinx/source/en/2-user_guide/7-tooling/5-robot_import.md b/docs/sphinx/source/en/2-user_guide/7-tooling/5-robot_import.md index 5c657bb03..bb6569642 100644 --- a/docs/sphinx/source/en/2-user_guide/7-tooling/5-robot_import.md +++ b/docs/sphinx/source/en/2-user_guide/7-tooling/5-robot_import.md @@ -44,7 +44,7 @@ where possible. only for position-control owners. - If the robot must preserve torque/motor actuator semantics, later task extension should follow the control pattern in - `src/unilab/envs/locomotion/go2w/`: keep action interpretation, PD/torque + `src/unilab/tasks/locomotion/go2w/`: keep action interpretation, PD/torque control, and the actuator contract inside the robot owner boundary. - After conversion, `mujoco.viewer` opens automatically to show the converted result and proceed to keyframe adjustment. diff --git a/docs/sphinx/source/en/3-deployment/3-framework_migration/2-from_legged_gym.md b/docs/sphinx/source/en/3-deployment/3-framework_migration/2-from_legged_gym.md index 003a3a8d8..02117cfe0 100644 --- a/docs/sphinx/source/en/3-deployment/3-framework_migration/2-from_legged_gym.md +++ b/docs/sphinx/source/en/3-deployment/3-framework_migration/2-from_legged_gym.md @@ -32,7 +32,7 @@ mostly mechanical. ## Migration checklist 1. Copy your URDF / MJCF assets under `src/unilab/assets/robots//`. -2. Create a task module under `src/unilab/envs/locomotion//`. +2. Create a task module under `src/unilab/tasks/locomotion//`. 3. Mirror your reward terms; keep the same names so reward parity is diff-able. 4. Translate command sampling — Legged Gym's `_resample_commands` becomes diff --git a/docs/sphinx/source/en/4-developer_guide/1-architecture/4-scene_composition.md b/docs/sphinx/source/en/4-developer_guide/1-architecture/4-scene_composition.md index 872ac1a37..c921febca 100644 --- a/docs/sphinx/source/en/4-developer_guide/1-architecture/4-scene_composition.md +++ b/docs/sphinx/source/en/4-developer_guide/1-architecture/4-scene_composition.md @@ -105,7 +105,7 @@ Disallowed on hot paths: The current procedural terrain user-facing path is Go2 rough terrain: -- Env owner: `src/unilab/envs/locomotion/go2/rough.py` +- Task owner: `src/unilab/tasks/locomotion/go2/rough.py` - Terrain generator: `src/unilab/terrains/terrain_generator.py` - MuJoCo materializer: `src/unilab/base/backend/mujoco/xml.py` - Motrix materializer: `src/unilab/base/backend/motrix/scene.py` diff --git a/docs/sphinx/source/en/4-developer_guide/3-extending/1-new_task.md b/docs/sphinx/source/en/4-developer_guide/3-extending/1-new_task.md index 6ee3871a7..25b37fd60 100644 --- a/docs/sphinx/source/en/4-developer_guide/3-extending/1-new_task.md +++ b/docs/sphinx/source/en/4-developer_guide/3-extending/1-new_task.md @@ -38,5 +38,5 @@ Start from the contracts: {doc}`../2-contracts/1-env_contract`, - Registry API: `src/unilab/base/registry.py` - Env state contract: `src/unilab/base/np_env.py` - Scene config: `src/unilab/base/scene.py` -- Existing task examples: `src/unilab/envs/locomotion/go2/joystick.py`, +- Existing task examples: `src/unilab/tasks/locomotion/go2/joystick.py`, `src/unilab/tasks/manipulation/allegro_inhand/rotation.py` diff --git a/docs/sphinx/source/zh_CN/2-user_guide/4-tasks/1-locomotion.md b/docs/sphinx/source/zh_CN/2-user_guide/4-tasks/1-locomotion.md index ba6885967..006275c6c 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/4-tasks/1-locomotion.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/4-tasks/1-locomotion.md @@ -1,7 +1,7 @@ # 运动控制 -运动控制任务注册在 `src/unilab/envs/locomotion/` 和 -`src/unilab/envs/motion_tracking/` 中。`conf/` 下可用的 owner YAML +运动控制任务注册在 `src/unilab/tasks/locomotion/` 和 +`src/unilab/tasks/motion_tracking/` 中。`conf/` 下可用的 owner YAML 定义了哪些算法与后端组合是可运行的。 ## 系列 @@ -33,7 +33,7 @@ uv run train --algo sac --task g1_walk_flat --sim mujoco - PPO 配置：`conf/ppo/task/go2_footstand/mujoco.yaml` - 环境注册名：`Go2FootStand`（注册于 `sim_backend="mujoco"`） -- 环境实现：`src/unilab/envs/locomotion/go2/footstand.py`（继承 Go2 基础任务） +- 环境实现：`src/unilab/tasks/locomotion/go2/footstand.py`（继承 Go2 基础任务） - Go2 模型 XML：`src/unilab/assets/robots/go2/go2.xml` ```bash 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 19227a448..08b4fee9d 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 @@ -26,9 +26,9 @@ 具有代表性的 provider 实现位于： -- `src/unilab/envs/locomotion/go1/joystick.py` +- `src/unilab/tasks/locomotion/go1/joystick.py` - `src/unilab/tasks/locomotion/g1/joystick.py` -- `src/unilab/envs/motion_tracking/g1/tracking.py` +- `src/unilab/tasks/motion_tracking/g1/tracking.py` - `src/unilab/tasks/manipulation/allegro_inhand/rotation.py` - `src/unilab/tasks/manipulation/sharpa_inhand/rotation.py` diff --git a/docs/sphinx/source/zh_CN/2-user_guide/6-terrain/2-heightfield_import.md b/docs/sphinx/source/zh_CN/2-user_guide/6-terrain/2-heightfield_import.md index 4020a517d..64b439691 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/6-terrain/2-heightfield_import.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/6-terrain/2-heightfield_import.md @@ -6,7 +6,7 @@ - `src/unilab/terrains/heightfield_terrains.py` - `src/unilab/terrains/terrain_generator.py` -- `src/unilab/envs/locomotion/go2/rough.py` +- `src/unilab/tasks/locomotion/go2/rough.py` - `src/unilab/base/backend/mujoco/xml.py` - `src/unilab/base/backend/motrix/scene.py` diff --git a/docs/sphinx/source/zh_CN/2-user_guide/7-tooling/5-robot_import.md b/docs/sphinx/source/zh_CN/2-user_guide/7-tooling/5-robot_import.md index d74b5c620..08a48e67b 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/7-tooling/5-robot_import.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/7-tooling/5-robot_import.md @@ -36,7 +36,7 @@ visual mesh 作为 collision mesh，尽量把碰撞体简化为 box / capsule / - 默认自动导入会把 actuator 写成 `position`，这只适合位置控制 owner。 - 如果机器人必须保留 torque/motor actuator 语义，后续扩展任务时，需要参考 - `src/unilab/envs/locomotion/go2w/` 的控制方式，把 action 解释、PD/力矩控制和 + `src/unilab/tasks/locomotion/go2w/` 的控制方式，把 action 解释、PD/力矩控制和 actuator contract 放在机器人 owner 的控制边界内。 - 转换完成后，会自动弹出 `mujoco.viewer` 可视化界面展示转换结果，并进行下一步调整 Keyframe。 diff --git a/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/2-from_legged_gym.md b/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/2-from_legged_gym.md index 05411ec4e..026a674b9 100644 --- a/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/2-from_legged_gym.md +++ b/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/2-from_legged_gym.md @@ -29,7 +29,7 @@ Legged Gym 曾是那套 GPU 常驻的 PPO 模板，教会了整个领域如何 ## 迁移清单 1. 把你的 URDF / MJCF asset 复制到 `src/unilab/assets/robots//` 下。 -2. 在 `src/unilab/envs/locomotion//` 下创建一个任务模块。 +2. 在 `src/unilab/tasks/locomotion//` 下创建一个任务模块。 3. 镜像你的 reward 项；保持名称相同，以便 reward 一致性可被 diff。 4. 翻译命令采样 —— Legged Gym 的 `_resample_commands` 在 UniLab 中变成一个 curriculum provider。 diff --git a/docs/sphinx/source/zh_CN/4-developer_guide/1-architecture/4-scene_composition.md b/docs/sphinx/source/zh_CN/4-developer_guide/1-architecture/4-scene_composition.md index 0b4abbb35..f3e0e662a 100644 --- a/docs/sphinx/source/zh_CN/4-developer_guide/1-architecture/4-scene_composition.md +++ b/docs/sphinx/source/zh_CN/4-developer_guide/1-architecture/4-scene_composition.md @@ -98,7 +98,7 @@ materializer。 当前面向用户的程序化地形路径是 Go2 崎岖地形： -- Env owner：`src/unilab/envs/locomotion/go2/rough.py` +- Task owner：`src/unilab/tasks/locomotion/go2/rough.py` - 地形生成器：`src/unilab/terrains/terrain_generator.py` - MuJoCo materializer：`src/unilab/base/backend/mujoco/xml.py` - Motrix materializer：`src/unilab/base/backend/motrix/scene.py` diff --git a/docs/sphinx/source/zh_CN/4-developer_guide/3-extending/1-new_task.md b/docs/sphinx/source/zh_CN/4-developer_guide/3-extending/1-new_task.md index 49e7c39e7..441574998 100644 --- a/docs/sphinx/source/zh_CN/4-developer_guide/3-extending/1-new_task.md +++ b/docs/sphinx/source/zh_CN/4-developer_guide/3-extending/1-new_task.md @@ -38,5 +38,5 @@ - Registry API：`src/unilab/base/registry.py` - Env 状态契约：`src/unilab/base/np_env.py` - 场景配置：`src/unilab/base/scene.py` -- 现有任务示例：`src/unilab/envs/locomotion/go2/joystick.py`、 +- 现有任务示例：`src/unilab/tasks/locomotion/go2/joystick.py`、 `src/unilab/tasks/manipulation/allegro_inhand/rotation.py` From fa616e54b9288e1098986e6ddb54bc468130d8a7 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Wed, 19 Aug 2026 00:42:20 +0800 Subject: [PATCH 075/169] docs(tasks): update motion tracking owner paths [skip ci] (#1196) --- .../source/en/2-user_guide/4-tasks/2-motion_tracking.md | 6 +++--- .../source/en/3-deployment/1-sim_to_real/2-g1_whole_body.md | 2 +- .../en/4-developer_guide/2-contracts/4-dr_contract.md | 2 +- docs/sphinx/source/en/4-developer_guide/7-motion_assets.md | 4 ++-- .../source/zh_CN/2-user_guide/4-tasks/2-motion_tracking.md | 6 +++--- .../zh_CN/3-deployment/1-sim_to_real/2-g1_whole_body.md | 2 +- .../zh_CN/4-developer_guide/2-contracts/4-dr_contract.md | 2 +- .../source/zh_CN/4-developer_guide/7-motion_assets.md | 4 ++-- 8 files changed, 14 insertions(+), 14 deletions(-) diff --git a/docs/sphinx/source/en/2-user_guide/4-tasks/2-motion_tracking.md b/docs/sphinx/source/en/2-user_guide/4-tasks/2-motion_tracking.md index 7b78c58a5..3245e8c45 100644 --- a/docs/sphinx/source/en/2-user_guide/4-tasks/2-motion_tracking.md +++ b/docs/sphinx/source/en/2-user_guide/4-tasks/2-motion_tracking.md @@ -1,12 +1,12 @@ # Motion Tracking -G1 motion tracking tasks live under `src/unilab/envs/motion_tracking/` and are +G1 motion tracking tasks live under `src/unilab/tasks/motion_tracking/` and are selected through task owner YAMLs in `conf/ppo/`, `conf/appo/`, and selected off-policy paths. > **Motion assets moved to Hugging Face.** The `.npz` clips are no longer shipped > in the repository. On first use `MotionLoader` -> (`src/unilab/envs/motion_tracking/g1/motion_loader.py`) downloads them on demand +> (`src/unilab/tasks/motion_tracking/common/motion_loader.py`) downloads them on demand > from [unilabsim/unilab-motions](https://huggingface.co/datasets/unilabsim/unilab-motions) > via `src/unilab/assets/hub.py` (`_HF_MOTIONS_REPO_ID`). `uv sync` already installs > the required `huggingface_hub` dependency. @@ -75,7 +75,7 @@ uv run eval --algo sac --task g1_motion_tracking --sim motrix \ Motion NPZ files are read through `env.motion_file`, which also accepts a list of paths. A standard clip must contain the seven keys `fps`, `joint_pos`, `joint_vel`, `body_pos_w`, `body_quat_w`, `body_lin_vel_w`, and `body_ang_vel_w` -(validated in `g1/motion_loader.py`): +(validated in `common/motion_loader.py`): ```yaml env: diff --git a/docs/sphinx/source/en/3-deployment/1-sim_to_real/2-g1_whole_body.md b/docs/sphinx/source/en/3-deployment/1-sim_to_real/2-g1_whole_body.md index e4ba8b088..905ac9736 100644 --- a/docs/sphinx/source/en/3-deployment/1-sim_to_real/2-g1_whole_body.md +++ b/docs/sphinx/source/en/3-deployment/1-sim_to_real/2-g1_whole_body.md @@ -115,7 +115,7 @@ clip. On hardware you need a wall-clock → phase mapping that is: - **Bounded rate** — clip dφ/dt to the value the policy was trained with (the motion loader records this; load `reference_motion.npz`). -See `unilab.envs.motion_tracking.g1.motion_loader` for the sim-side +See `unilab.tasks.motion_tracking.common.motion_loader` for the sim-side loader you should mirror on hardware. ## 5. Safety layer 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 723762745..94ddb57d4 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 @@ -101,5 +101,5 @@ payloads. - DR manager: `src/unilab/dr/manager.py` - Backend interface: `src/unilab/base/backend/base.py` - Example providers: `src/unilab/tasks/locomotion/g1/joystick.py`, - `src/unilab/envs/motion_tracking/g1/tracking.py`, + `src/unilab/tasks/motion_tracking/g1/tracking.py`, `src/unilab/tasks/manipulation/sharpa_inhand/rotation.py` diff --git a/docs/sphinx/source/en/4-developer_guide/7-motion_assets.md b/docs/sphinx/source/en/4-developer_guide/7-motion_assets.md index ded0b552f..c32b02dfb 100644 --- a/docs/sphinx/source/en/4-developer_guide/7-motion_assets.md +++ b/docs/sphinx/source/en/4-developer_guide/7-motion_assets.md @@ -115,12 +115,12 @@ To add a new robot's meshes: - Asset resolver module: `src/unilab/assets/hub.py` (`resolve_motion_files`). - Single integration point: `MotionLoader.__init__` in - `src/unilab/envs/motion_tracking/g1/motion_loader.py`, which calls the + `src/unilab/tasks/motion_tracking/common/motion_loader.py`, which calls the resolver once on a cold path. - Hot paths (`step` / `reset`) never trigger any file download or parsing. - `ASSETS_ROOT_PATH` is unchanged, so the download target matches the original local path exactly. - Robot meshes use the same directory resolver (`resolve_robot_asset_dir`), integrated at `X2WallFlipTrackingEnv.__init__` in - `src/unilab/envs/motion_tracking/x2/flip_tracking.py`, and exposed as the + `src/unilab/tasks/motion_tracking/x2/flip_tracking.py`, and exposed as the `unilab-pull-assets` CLI. diff --git a/docs/sphinx/source/zh_CN/2-user_guide/4-tasks/2-motion_tracking.md b/docs/sphinx/source/zh_CN/2-user_guide/4-tasks/2-motion_tracking.md index cf7a16812..8812fd3c4 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/4-tasks/2-motion_tracking.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/4-tasks/2-motion_tracking.md @@ -1,10 +1,10 @@ # 动作追踪 -G1 动作追踪任务位于 `src/unilab/envs/motion_tracking/` 下，并通过 +G1 动作追踪任务位于 `src/unilab/tasks/motion_tracking/` 下，并通过 `conf/ppo/`、`conf/appo/` 以及选定的 off-policy 路径中的 task owner YAML 选择。 > **Motion 资产已迁移到 Hugging Face。** `.npz` 片段不再随仓库分发，首次使用时由 -> `MotionLoader`（`src/unilab/envs/motion_tracking/g1/motion_loader.py`）按需从 +> `MotionLoader`（`src/unilab/tasks/motion_tracking/common/motion_loader.py`）按需从 > [unilabsim/unilab-motions](https://huggingface.co/datasets/unilabsim/unilab-motions) > 下载，下载逻辑在 `src/unilab/assets/hub.py`（`_HF_MOTIONS_REPO_ID`）。`uv sync` > 已自动安装所需的 `huggingface_hub` 依赖。 @@ -71,7 +71,7 @@ uv run eval --algo sac --task g1_motion_tracking --sim motrix \ 动作 NPZ 文件通过 `env.motion_file` 读取，也支持路径列表。标准片段必须包含七个 key： `fps`、`joint_pos`、`joint_vel`、`body_pos_w`、`body_quat_w`、`body_lin_vel_w`、 -`body_ang_vel_w`（在 `g1/motion_loader.py` 中校验）： +`body_ang_vel_w`（在 `common/motion_loader.py` 中校验）： ```yaml env: diff --git a/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/2-g1_whole_body.md b/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/2-g1_whole_body.md index df79c491e..0b7f9d12b 100644 --- a/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/2-g1_whole_body.md +++ b/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/2-g1_whole_body.md @@ -110,7 +110,7 @@ G1 部署原型将 actor 输出严格映射为： - **速率有界** —— 将 dφ/dt 钳制到策略训练时所用的值（运动加载器会记录这个值；加载 `reference_motion.npz`）。 -参见 `unilab.envs.motion_tracking.g1.motion_loader`，这是你应当在硬件上镜像的仿真侧 +参见 `unilab.tasks.motion_tracking.common.motion_loader`，这是你应当在硬件上镜像的仿真侧 加载器。 ## 5. 安全层 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 d6120a983..ef5f820f4 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 @@ -91,5 +91,5 @@ actuator 的机制泄漏到共享 payload 里。 - DR manager：`src/unilab/dr/manager.py` - Backend 接口：`src/unilab/base/backend/base.py` - 示例 provider：`src/unilab/tasks/locomotion/g1/joystick.py`、 - `src/unilab/envs/motion_tracking/g1/tracking.py`、 + `src/unilab/tasks/motion_tracking/g1/tracking.py`、 `src/unilab/tasks/manipulation/sharpa_inhand/rotation.py` diff --git a/docs/sphinx/source/zh_CN/4-developer_guide/7-motion_assets.md b/docs/sphinx/source/zh_CN/4-developer_guide/7-motion_assets.md index b8df1f175..5db21042e 100644 --- a/docs/sphinx/source/zh_CN/4-developer_guide/7-motion_assets.md +++ b/docs/sphinx/source/zh_CN/4-developer_guide/7-motion_assets.md @@ -103,10 +103,10 @@ uv run unilab-pull-assets --robot x2 ## 架构说明 - 资产解析模块：`src/unilab/assets/hub.py`（`resolve_motion_files`）。 -- 唯一集成点：`src/unilab/envs/motion_tracking/g1/motion_loader.py` 中的 +- 唯一集成点：`src/unilab/tasks/motion_tracking/common/motion_loader.py` 中的 `MotionLoader.__init__`，在冷路径上调用一次 resolver。 - 热路径（`step` / `reset`）**不会**触发任何文件下载或解析。 - `ASSETS_ROOT_PATH` 定义不变，下载落盘位置与原始本地路径完全一致。 - 机器人网格使用同一目录 resolver（`resolve_robot_asset_dir`），集成点为 - `src/unilab/envs/motion_tracking/x2/flip_tracking.py` 中的 + `src/unilab/tasks/motion_tracking/x2/flip_tracking.py` 中的 `X2WallFlipTrackingEnv.__init__`，并通过 `unilab-pull-assets` CLI 暴露。 From b1c04695e26c50c6ed4629f4776d1b539c4d062f Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Wed, 19 Aug 2026 00:47:18 +0800 Subject: [PATCH 076/169] chore(tasks): finish ownership metadata cleanup [skip ci] (#1198) --- .github/CODEOWNERS | 4 ++-- scripts/benchmark/torch_env/motion_tracking.py | 2 +- scripts/deploy/export_motion_bin.py | 2 +- src/unilab/tasks/__init__.py | 8 +++----- tests/nan_injection/stage3_nan_inject.py | 2 +- 5 files changed, 8 insertions(+), 10 deletions(-) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index d30646ad1..a108a93ad 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -7,8 +7,8 @@ /scripts/train_*.py @TATP-233 @caozx1110 # Task and environment ownership -/src/unilab/envs/motion_tracking/ @caozx1110 -/src/unilab/envs/manipulation/ @Mingrui-Yu +/src/unilab/tasks/motion_tracking/ @caozx1110 +/src/unilab/tasks/manipulation/ @Mingrui-Yu /scripts/motion/ @caozx1110 # Project process and docs diff --git a/scripts/benchmark/torch_env/motion_tracking.py b/scripts/benchmark/torch_env/motion_tracking.py index 029aed9a9..7d7d07bf9 100644 --- a/scripts/benchmark/torch_env/motion_tracking.py +++ b/scripts/benchmark/torch_env/motion_tracking.py @@ -5,7 +5,7 @@ (num_envs=2048, 29-dof, 14 tracked bodies): - `MotionTrackingEnv.update_state` - (src/unilab/envs/motion_tracking/common/tracking.py): + (src/unilab/tasks/motion_tracking/common/tracking.py): motion gather, relative transforms (transforms.py), terminations (terminations.py), 9 active reward terms (rewards.py, incl. per-term logging every 4 steps), observation build (observations.py, actor 160 / critic 289 diff --git a/scripts/deploy/export_motion_bin.py b/scripts/deploy/export_motion_bin.py index 27905845c..058c82610 100644 --- a/scripts/deploy/export_motion_bin.py +++ b/scripts/deploy/export_motion_bin.py @@ -15,7 +15,7 @@ body_lin_vel_w [num_frames][num_bodies][3] body_ang_vel_w [num_frames][num_bodies][3] -NPZ source layout (per src/unilab/envs/motion_tracking/g1/motion_loader.py): +NPZ source layout (per src/unilab/tasks/motion_tracking/common/motion_loader.py): - 'fps' (int) - 'joint_pos' (N, 29) - 'joint_vel' (N, 29) diff --git a/src/unilab/tasks/__init__.py b/src/unilab/tasks/__init__.py index 6e0d95230..d9c58aa9f 100644 --- a/src/unilab/tasks/__init__.py +++ b/src/unilab/tasks/__init__.py @@ -1,10 +1,8 @@ """Production task registry bootstrap. -Concrete task implementations are moving from :mod:`unilab.envs` into this -package under issue #1112. Until each task family moves, this explicit list -records its legacy module as the last remaining consumer of that path. The -registry imports these leaf modules directly, so registration stays explicit -and deterministic throughout the migration. +Concrete task implementations live in this package. The registry imports the +explicit leaf-module list directly, so registration stays deterministic and +does not depend on package discovery or import order. """ __unilab_registry_modules__ = ( diff --git a/tests/nan_injection/stage3_nan_inject.py b/tests/nan_injection/stage3_nan_inject.py index 3440ff786..b23860558 100644 --- a/tests/nan_injection/stage3_nan_inject.py +++ b/tests/nan_injection/stage3_nan_inject.py @@ -306,7 +306,7 @@ def _print_manual_recipe(): To exercise NaN detection end-to-end inside a collector subprocess: 1. Pick a task env file, e.g. - src/unilab/envs/locomotion/go1/go1_joystick.py + src/unilab/tasks/locomotion/go1/joystick.py 2. Inside the env's update_state or apply_action, add a temporary one-shot NaN raise guarded by an env-counter, for example: From c416e252c8251be782c7d60cbcfceac74d03f9fe Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Wed, 19 Aug 2026 01:31:41 +0800 Subject: [PATCH 077/169] feat(env): add model-field reset terms [skip ci] (#1201) --- src/unilab/base/entity.py | 161 +++++++++++++++ src/unilab/base/reset_state.py | 100 +++++++++- src/unilab/envs/mdp/__init__.py | 6 + src/unilab/envs/mdp/events.py | 338 ++++++++++++++++++++++++++++++++ tests/envs/mdp/test_events.py | 144 +++++++++++++- 5 files changed, 739 insertions(+), 10 deletions(-) diff --git a/src/unilab/base/entity.py b/src/unilab/base/entity.py index 4bfd980a5..8ccbc5eea 100644 --- a/src/unilab/base/entity.py +++ b/src/unilab/base/entity.py @@ -426,6 +426,7 @@ def __init__( self._reset_root_layout_error: str | None = None self._reset_joint_qpos_ids: np.ndarray | None = None self._reset_joint_qvel_ids: np.ndarray | None = None + self._joint_model_dof_ids: np.ndarray | None = None self._joint_names = _normalize_names(name, "joint", cfg.joint_names) self._body_names = _normalize_names(name, "body", cfg.body_names) @@ -1091,6 +1092,106 @@ def write_actuator_gains_to_sim( term_name=f"{term_name}:{self.name}", ) + def bind_joint_armature_write( + self, + joint_ids: np.ndarray | Sequence[int] | slice | None = None, + *, + term_name: str, + ) -> tuple[np.ndarray, np.ndarray]: + """Bind entity-local joints and immutable default DOF armatures.""" + if self._reset_state is None: + raise self._capability_error( + "reset joint-armature write", + "EntityScene was materialized without an env-owned reset transaction", + ) + model_dof_ids = self._materialize_joint_model_dof_ids() + local_ids = self._normalize_local_joint_ids( + joint_ids, + capability="reset joint-armature write", + ) + if local_ids.size == 0: + raise ValueError(f"Entity '{self.name}' reset joint-armature write selected no joints") + _, defaults = self._reset_state.bind_dof_armature_write( + model_dof_ids[local_ids], + term_name=f"{term_name}:{self.name}", + ) + return self._readonly_local_binding(local_ids, defaults) + + def write_joint_armature_to_sim( + self, + values: np.ndarray, + joint_ids: np.ndarray | Sequence[int] | slice | None = None, + env_ids: np.ndarray | slice | None = None, + *, + term_name: str = "joint_armature", + ) -> None: + """Stage selected entity joint armatures in the active reset transaction.""" + if self._reset_state is None: + raise self._capability_error( + "reset joint-armature write", + "EntityScene was materialized without an env-owned reset transaction", + ) + model_dof_ids = self._materialize_joint_model_dof_ids() + local_ids = self._normalize_local_joint_ids( + joint_ids, + capability="reset joint-armature write", + ) + self._reset_state.write_dof_armature( + self._normalize_reset_env_ids(env_ids), + model_dof_ids[local_ids], + values, + term_name=f"{term_name}:{self.name}", + ) + + def bind_geom_friction_write( + self, + geom_ids: np.ndarray | Sequence[int] | slice | None = None, + *, + term_name: str, + ) -> tuple[np.ndarray, np.ndarray]: + """Bind entity-local geoms and immutable default friction vectors.""" + if self._reset_state is None or self._geom_ids is None: + raise self._capability_error( + "reset geom-friction write", + "geom metadata or the env-owned reset transaction was not materialized", + ) + local_ids = self._normalize_local_geom_ids( + geom_ids, + capability="reset geom-friction write", + ) + if local_ids.size == 0: + raise ValueError(f"Entity '{self.name}' reset geom-friction write selected no geoms") + _, defaults = self._reset_state.bind_geom_friction_write( + self._geom_ids[local_ids], + term_name=f"{term_name}:{self.name}", + ) + return self._readonly_local_binding(local_ids, defaults) + + def write_geom_friction_to_sim( + self, + values: np.ndarray, + geom_ids: np.ndarray | Sequence[int] | slice | None = None, + env_ids: np.ndarray | slice | None = None, + *, + term_name: str = "geom_friction", + ) -> None: + """Stage selected entity geom friction in the active reset transaction.""" + if self._reset_state is None or self._geom_ids is None: + raise self._capability_error( + "reset geom-friction write", + "geom metadata or the env-owned reset transaction was not materialized", + ) + local_ids = self._normalize_local_geom_ids( + geom_ids, + capability="reset geom-friction write", + ) + self._reset_state.write_geom_friction( + self._normalize_reset_env_ids(env_ids), + self._geom_ids[local_ids], + values, + term_name=f"{term_name}:{self.name}", + ) + def bind_body_mass_write( self, body_ids: np.ndarray | Sequence[int] | slice | None = None, @@ -1363,6 +1464,33 @@ def _readonly_local_binding( bound_defaults.setflags(write=False) return bound_ids, bound_defaults + def _materialize_joint_model_dof_ids(self) -> np.ndarray: + """Resolve full model DOF addresses once for reset-time model fields.""" + cached = self._joint_model_dof_ids + if cached is not None: + return cached + if self._joint_names is None: + raise self._capability_error( + "reset joint-armature write", + "joint_names were not declared in EntityCfg", + ) + try: + values = self._backend.get_joint_dof_indices(self._joint_names) + except (AttributeError, NotImplementedError) as exc: + raise self._capability_error("reset joint-armature write", str(exc)) from exc + except (KeyError, ValueError) as exc: + raise ValueError( + f"Entity '{self.name}' could not resolve joint model DOF names " + f"{list(self._joint_names)} on backend '{self._backend_type}': {exc}" + ) from exc + resolved = _readonly_ids( + values, + expected=len(self._joint_names), + label=f"Entity '{self.name}' joint model DOF", + ) + self._joint_model_dof_ids = resolved + return resolved + def _normalize_local_body_ids( self, body_ids: np.ndarray | Sequence[int] | slice | None, @@ -1429,6 +1557,39 @@ def _normalize_local_joint_ids( ) return ids + def _normalize_local_geom_ids( + self, + geom_ids: np.ndarray | Sequence[int] | slice | None, + *, + capability: str, + ) -> np.ndarray: + if geom_ids is None: + ids = np.arange(self.num_geoms, dtype=np.intp) + elif isinstance(geom_ids, slice): + ids = np.arange(self.num_geoms, dtype=np.intp)[geom_ids] + else: + raw = np.asarray(geom_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} geom_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_geoms): + raise IndexError( + f"Entity '{self.name}' {capability} geom_ids out of range for " + f"{self.num_geoms} geoms: {ids.tolist()}" + ) + if np.unique(ids).size != ids.size: + raise ValueError( + f"Entity '{self.name}' {capability} geom_ids contain duplicates: {ids.tolist()}" + ) + return ids + def _normalize_local_actuator_ids( self, actuator_ids: np.ndarray | Sequence[int] | slice | None, diff --git a/src/unilab/base/reset_state.py b/src/unilab/base/reset_state.py index bd730c710..96fce42f1 100644 --- a/src/unilab/base/reset_state.py +++ b/src/unilab/base/reset_state.py @@ -16,6 +16,8 @@ from unilab.dr.types import ( RESET_TERM_BODY_IPOS, RESET_TERM_BODY_MASS, + RESET_TERM_DOF_ARMATURE, + RESET_TERM_GEOM_FRICTION, RESET_TERM_GRAVITY, RESET_TERM_KD, RESET_TERM_KP, @@ -135,6 +137,48 @@ def bind_gravity_write(self, *, term_name: str) -> np.ndarray: term_name=term_name, ) + def bind_dof_armature_write( + self, + dof_ids: np.ndarray, + *, + term_name: str, + ) -> tuple[np.ndarray, np.ndarray]: + """Bind DOF-armature columns and immutable backend defaults.""" + default = self._materialize_randomization_default( + RESET_TERM_DOF_ARMATURE, + getter=self._backend.get_dof_armature, + expected_tail=None, + term_name=term_name, + ) + columns = self._validate_columns( + dof_ids, + width=default.shape[0], + capability="DOF armature IDs", + term_name=term_name, + ) + return self._readonly_binding(columns, default[columns]) + + def bind_geom_friction_write( + self, + geom_ids: np.ndarray, + *, + term_name: str, + ) -> tuple[np.ndarray, np.ndarray]: + """Bind geom-friction rows and immutable backend defaults.""" + default = self._materialize_randomization_default( + RESET_TERM_GEOM_FRICTION, + getter=self._backend.get_geom_friction, + expected_tail=(3,), + term_name=term_name, + ) + columns = self._validate_columns( + geom_ids, + width=default.shape[0], + capability="geom friction IDs", + term_name=term_name, + ) + return self._readonly_binding(columns, default[columns]) + def write_body_mass( self, env_ids: np.ndarray, @@ -144,7 +188,7 @@ def write_body_mass( term_name: str, ) -> None: """Stage selected body masses in the exactly-once reset payload.""" - self._write_body_randomization( + self._write_selected_randomization( RESET_TERM_BODY_MASS, env_ids, body_ids, @@ -162,7 +206,7 @@ def write_body_ipos( term_name: str, ) -> None: """Stage selected body inertial positions in the reset payload.""" - self._write_body_randomization( + self._write_selected_randomization( RESET_TERM_BODY_IPOS, env_ids, body_ids, @@ -200,6 +244,42 @@ def write_gravity( mask[ids] = True self._dirty_mask[ids] = True + def write_dof_armature( + self, + env_ids: np.ndarray, + dof_ids: np.ndarray, + values: np.ndarray, + *, + term_name: str, + ) -> None: + """Stage selected DOF armatures in the reset payload.""" + self._write_selected_randomization( + RESET_TERM_DOF_ARMATURE, + env_ids, + dof_ids, + values, + value_tail=(), + term_name=term_name, + ) + + def write_geom_friction( + self, + env_ids: np.ndarray, + geom_ids: np.ndarray, + values: np.ndarray, + *, + term_name: str, + ) -> None: + """Stage selected three-axis geom friction in the reset payload.""" + self._write_selected_randomization( + RESET_TERM_GEOM_FRICTION, + env_ids, + geom_ids, + values, + value_tail=(3,), + term_name=term_name, + ) + def bind_actuator_gain_write( self, actuator_ids: np.ndarray, @@ -612,11 +692,11 @@ def _readonly_binding( selected.setflags(write=False) return bound_columns, selected - def _write_body_randomization( + def _write_selected_randomization( self, field: str, env_ids: np.ndarray, - body_ids: np.ndarray, + column_ids: np.ndarray, values: np.ndarray, *, value_tail: tuple[int, ...], @@ -629,9 +709,9 @@ def _write_body_randomization( ) default = self._require_randomization_default(field, term_name) columns = self._validate_columns( - body_ids, + column_ids, width=default.shape[0], - capability=f"{field} body IDs", + capability=f"{field} column IDs", term_name=term_name, ) selected = self._validate_values( @@ -655,7 +735,13 @@ def _build_randomization_payload( dirty_ids: np.ndarray, ) -> ResetRandomizationPayload | None: payload = ResetRandomizationPayload() - for field in (RESET_TERM_BODY_MASS, RESET_TERM_BODY_IPOS, RESET_TERM_GRAVITY): + for field in ( + RESET_TERM_BODY_MASS, + RESET_TERM_BODY_IPOS, + RESET_TERM_DOF_ARMATURE, + RESET_TERM_GEOM_FRICTION, + RESET_TERM_GRAVITY, + ): mask = self._randomization_dirty_masks.get(field) if mask is None or not np.any(mask): continue diff --git a/src/unilab/envs/mdp/__init__.py b/src/unilab/envs/mdp/__init__.py index f3627001b..6d9ba3726 100644 --- a/src/unilab/envs/mdp/__init__.py +++ b/src/unilab/envs/mdp/__init__.py @@ -4,6 +4,9 @@ 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 dof_armature as dof_armature +from unilab.envs.mdp.events import geom_friction as geom_friction +from unilab.envs.mdp.events import joint_armature as joint_armature from unilab.envs.mdp.events import pd_gains as pd_gains from unilab.envs.mdp.events import push_by_setting_velocity as push_by_setting_velocity from unilab.envs.mdp.events import ( @@ -54,9 +57,12 @@ "builtin_sensor", "bad_orientation", "body_angular_velocity_penalty", + "dof_armature", "flat_orientation_l2", + "geom_friction", "generated_commands", "joint_pos_rel", + "joint_armature", "joint_vel_rel", "joint_vel_l2", "last_action", diff --git a/src/unilab/envs/mdp/events.py b/src/unilab/envs/mdp/events.py index e370fa159..12b0a7928 100644 --- a/src/unilab/envs/mdp/events.py +++ b/src/unilab/envs/mdp/events.py @@ -8,6 +8,7 @@ from __future__ import annotations +import re from typing import TYPE_CHECKING, Any, Literal, cast import numpy as np @@ -249,6 +250,340 @@ def resolve_env_ids(env: ManagerBasedRlEnv, env_ids: np.ndarray | None) -> np.nd return env_ids +class _ModelFieldRandomizer(ManagerTermBase): + """Cold-path-bound NumPy adapter for pinned mjlab model-field DR terms.""" + + _term_name = "model_field" + _field_width = 1 + _default_axes: tuple[int, ...] = (0,) + _valid_axes: tuple[int, ...] = (0,) + _PARAMS = frozenset( + ("ranges", "asset_cfg", "distribution", "operation", "axes", "shared_random") + ) + + def __init__(self, cfg: EventTermCfg, env: ManagerBasedRlEnv): + super().__init__(env) + _validate_event_term( + cfg, + term_name=self._term_name, + mode="reset", + allowed_params=self._PARAMS, + required_params=("ranges",), + ) + self._distribution = _event_choice( + cfg.params.get("distribution", "uniform"), + term_name=self._term_name, + name="distribution", + choices=_DISTRIBUTIONS, + ) + self._operation = _event_choice( + cfg.params.get("operation", "abs"), + term_name=self._term_name, + name="operation", + choices=_OPERATIONS, + ) + shared_random = cfg.params.get("shared_random", False) + if not isinstance(shared_random, bool): + raise TypeError( + f"EventManager term '{self._term_name}' parameter 'shared_random' must be bool" + ) + self._shared_random = shared_random + asset_cfg = cfg.params.get("asset_cfg", _DEFAULT_ASSET_CFG) + if not isinstance(asset_cfg, SceneEntityCfg): + raise TypeError( + f"EventManager term '{self._term_name}' asset_cfg must be SceneEntityCfg, " + f"got {type(asset_cfg).__name__}" + ) + entity = cast("Entity", env.scene[asset_cfg.name]) + local_ids, defaults, names = self._bind(entity, asset_cfg) + ranges = cfg.params["ranges"] + local_ids, defaults, names, ranges = self._select_string_ranges( + local_ids, + defaults, + names, + ranges, + ) + self._entity = entity + self._local_ids = local_ids + self._defaults = defaults + self._axes = self._resolve_axes(cfg.params.get("axes"), ranges) + self._ranges = self._resolve_ranges(ranges, names) + + def _bind( + self, + entity: Entity, + asset_cfg: SceneEntityCfg, + ) -> tuple[np.ndarray, np.ndarray, tuple[str, ...]]: + raise NotImplementedError + + def _write( + self, + values: np.ndarray, + env_ids: np.ndarray, + ) -> None: + raise NotImplementedError + + def _select_string_ranges( + self, + local_ids: np.ndarray, + defaults: np.ndarray, + names: tuple[str, ...], + ranges: Any, + ) -> tuple[np.ndarray, np.ndarray, tuple[str, ...], Any]: + if not isinstance(ranges, dict) or not ranges: + return local_ids, defaults, names, ranges + keys = tuple(ranges) + if not all(isinstance(key, str) for key in keys): + if any(isinstance(key, str) for key in keys): + raise TypeError( + f"EventManager term '{self._term_name}' ranges cannot mix string and integer keys" + ) + return local_ids, defaults, names, ranges + + assigned: list[Any | None] = [None] * len(names) + for pattern, bounds in ranges.items(): + try: + matched = [index for index, name in enumerate(names) if re.fullmatch(pattern, name)] + except re.error as exc: + raise ValueError( + f"EventManager term '{self._term_name}' ranges contains invalid regex " + f"{pattern!r}: {exc}" + ) from exc + if not matched: + raise ValueError( + f"EventManager term '{self._term_name}' ranges pattern {pattern!r} " + f"matched no selected names; available={list(names)}" + ) + for index in matched: + if assigned[index] is not None: + raise ValueError( + f"EventManager term '{self._term_name}' ranges patterns overlap for " + f"selected name '{names[index]}'" + ) + assigned[index] = bounds + selected = np.asarray( + [index for index, value in enumerate(assigned) if value is not None], + dtype=np.intp, + ) + return ( + local_ids[selected], + defaults[selected], + tuple(names[index] for index in selected), + [assigned[index] for index in selected], + ) + + def _resolve_axes(self, value: Any, ranges: Any) -> tuple[int, ...]: + if value is None and isinstance(ranges, dict) and ranges: + value = list(ranges) + if value is None: + axes = self._default_axes + else: + if not isinstance(value, (list, tuple)): + raise TypeError( + f"EventManager term '{self._term_name}' parameter 'axes' must be a sequence" + ) + if any( + isinstance(axis, bool) or not isinstance(axis, (int, np.integer)) for axis in value + ): + raise TypeError( + f"EventManager term '{self._term_name}' parameter 'axes' must contain integers" + ) + axes = tuple(int(axis) for axis in value) + if not axes: + raise ValueError( + f"EventManager term '{self._term_name}' parameter 'axes' cannot be empty" + ) + if len(set(axes)) != len(axes): + raise ValueError( + f"EventManager term '{self._term_name}' parameter 'axes' contains duplicates" + ) + invalid = sorted(set(axes) - set(self._valid_axes)) + if invalid: + raise ValueError( + f"EventManager term '{self._term_name}' has invalid axes {invalid}; " + f"valid axes are {list(self._valid_axes)}" + ) + return axes + + def _resolve_ranges(self, value: Any, names: tuple[str, ...]) -> np.ndarray: + per_entity: list[Any] + if ( + isinstance(value, list) + and len(value) == len(names) + and any(isinstance(item, (tuple, list, np.ndarray)) for item in value) + ): + per_entity = list(value) + value = None + else: + per_entity = [] + + result = np.empty((len(names), self._field_width, 2), dtype=np.float64) + result[:] = np.nan + if per_entity: + if len(self._axes) != 1: + raise ValueError( + f"EventManager term '{self._term_name}' string-keyed ranges require one axis" + ) + for index, bounds in enumerate(per_entity): + result[index, self._axes[0]] = _distribution_parameters( + bounds, + term_name=self._term_name, + name=f"ranges[{names[index]}]", + distribution=self._distribution, + ) + elif isinstance(value, dict): + unknown = sorted(set(value) - set(self._axes)) + missing = sorted(set(self._axes) - set(value)) + if unknown or missing: + raise ValueError( + f"EventManager term '{self._term_name}' ranges axes mismatch; " + f"missing={missing}, unknown={unknown}" + ) + for axis in self._axes: + result[:, axis] = _distribution_parameters( + value[axis], + term_name=self._term_name, + name=f"ranges[{axis}]", + distribution=self._distribution, + ) + else: + parameters = _distribution_parameters( + value, + term_name=self._term_name, + name="ranges", + distribution=self._distribution, + ) + for axis in self._axes: + result[:, axis] = parameters + result.setflags(write=False) + return result + + def _sample_axis(self, env: ManagerBasedRlEnv, axis: int, count: int) -> np.ndarray: + parameters = self._ranges[:, axis] + if self._shared_random: + samples = np.empty((count, len(parameters)), dtype=np.float64) + groups: dict[tuple[float, float], list[int]] = {} + for index, pair in enumerate(parameters): + groups.setdefault((float(pair[0]), float(pair[1])), []).append(index) + for (first, second), indices in groups.items(): + if self._distribution == "gaussian": + shared = env.rng.normal(first, second, size=(count, 1)) + elif self._distribution == "log_uniform": + shared = np.exp(env.rng.uniform(np.log(first), np.log(second), size=(count, 1))) + else: + shared = env.rng.uniform(first, second, size=(count, 1)) + samples[:, indices] = shared + return samples + if self._distribution == "gaussian": + return env.rng.normal(parameters[:, 0], parameters[:, 1], size=(count, len(parameters))) + if self._distribution == "log_uniform": + return np.exp( + env.rng.uniform( + np.log(parameters[:, 0]), + np.log(parameters[:, 1]), + size=(count, len(parameters)), + ) + ) + return env.rng.uniform( + parameters[:, 0], + parameters[:, 1], + size=(count, len(parameters)), + ) + + def __call__( + self, + env: ManagerBasedRlEnv, + env_ids: np.ndarray | None, + ranges: Any, + asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG, + distribution: str = "uniform", + operation: str = "abs", + axes: list[int] | None = None, + shared_random: bool = False, + ) -> None: + del ranges, asset_cfg, distribution, operation, axes, shared_random + ids = resolve_env_ids(env, env_ids) + defaults = self._defaults + scalar = defaults.ndim == 1 + default_values = defaults[:, None] if scalar else defaults + values = np.broadcast_to( + default_values, + (len(ids), *default_values.shape), + ).copy() + for axis in self._axes: + samples = self._sample_axis(env, axis, len(ids)) + values[:, :, axis] = _apply_randomization_operation( + default_values[None, :, axis], + samples, + self._operation, + ) + if np.any(values < 0.0) or not np.isfinite(values).all(): + raise ValueError( + f"EventManager term '{self._term_name}' produced negative, NaN, or Inf values" + ) + self._write(values[:, :, 0] if scalar else values, ids) + + +class GeomFriction(_ModelFieldRandomizer): + """Pinned mjlab-style geom friction randomization through the reset payload.""" + + _term_name = "geom_friction" + _field_width = 3 + _default_axes = (0,) + _valid_axes = (0, 1, 2) + + def _bind( + self, + entity: Entity, + asset_cfg: SceneEntityCfg, + ) -> tuple[np.ndarray, np.ndarray, tuple[str, ...]]: + local_ids, defaults = entity.bind_geom_friction_write( + asset_cfg.geom_ids, + term_name=self._term_name, + ) + names = tuple(entity.geom_names[int(index)] for index in local_ids) + return local_ids, defaults, names + + def _write(self, values: np.ndarray, env_ids: np.ndarray) -> None: + self._entity.write_geom_friction_to_sim( + values, + self._local_ids, + env_ids, + term_name=self._term_name, + ) + + +class JointArmature(_ModelFieldRandomizer): + """Pinned mjlab-style joint armature randomization through the reset payload.""" + + _term_name = "joint_armature" + + def _bind( + self, + entity: Entity, + asset_cfg: SceneEntityCfg, + ) -> tuple[np.ndarray, np.ndarray, tuple[str, ...]]: + local_ids, defaults = entity.bind_joint_armature_write( + asset_cfg.joint_ids, + term_name=self._term_name, + ) + names = tuple(entity.joint_names[int(index)] for index in local_ids) + return local_ids, defaults, names + + def _write(self, values: np.ndarray, env_ids: np.ndarray) -> None: + self._entity.write_joint_armature_to_sim( + values, + self._local_ids, + env_ids, + term_name=self._term_name, + ) + + +geom_friction = GeomFriction +joint_armature = JointArmature +dof_armature = joint_armature + + class PdGains(ManagerTermBase): """Pinned-mjlab-compatible PD gain randomization on UniLab reset payloads.""" @@ -682,6 +1017,9 @@ def reset_root_state_uniform( __all__ = [ + "dof_armature", + "geom_friction", + "joint_armature", "pd_gains", "push_by_setting_velocity", "randomize_physics_scene_gravity", diff --git a/tests/envs/mdp/test_events.py b/tests/envs/mdp/test_events.py index 8f2000276..3b2a3ccae 100644 --- a/tests/envs/mdp/test_events.py +++ b/tests/envs/mdp/test_events.py @@ -16,6 +16,8 @@ from unilab.dr.types import ( RESET_TERM_BODY_IPOS, RESET_TERM_BODY_MASS, + RESET_TERM_DOF_ARMATURE, + RESET_TERM_GEOM_FRICTION, RESET_TERM_GRAVITY, RESET_TERM_KD, RESET_TERM_KP, @@ -182,6 +184,8 @@ def __init__( self.body_mass = np.array([10.0]) self.body_ipos = np.array([[0.0, 0.0, 0.0]]) self.gravity = np.array([0.0, 0.0, -9.81]) + self.dof_armature = np.array([0.0] * 6 + [1.0, 2.0, 3.0]) + self.geom_friction = np.array([[0.5, 0.01, 0.001], [0.7, 0.02, 0.002], [0.9, 0.03, 0.003]]) self.interval_plans: list[IntervalRandomizationPlan] = [] def get_body_ids(self, names) -> np.ndarray: @@ -203,21 +207,49 @@ def get_init_qvel(self) -> np.ndarray: return self.init_qvel.copy() def get_dof_pos(self) -> np.ndarray: - return np.empty((self.num_envs, 0)) + return np.zeros((self.num_envs, 3)) def get_dof_vel(self) -> np.ndarray: - return np.empty((self.num_envs, 0)) + return np.zeros((self.num_envs, 3)) + + def get_default_dof_pos(self) -> np.ndarray: + return np.zeros(3) + + def get_joint_dof_indices(self, names) -> np.ndarray: + table = {"j0": 6, "j1": 7, "j2": 8} + return np.asarray([table[name] for name in names], dtype=np.int32) + + def get_joint_dof_pos_indices(self, names) -> np.ndarray: + table = {"j0": 0, "j1": 1, "j2": 2} + return np.asarray([table[name] for name in names], dtype=np.int32) + + def get_joint_dof_vel_indices(self, names) -> np.ndarray: + return self.get_joint_dof_pos_indices(names) + + def get_geom_names(self) -> tuple[str, ...]: + return ("floor", "foot", "base_geom") def get_actuator_names(self) -> tuple[str, ...]: return ("a0", "a1", "a2") + def get_actuator_joint_names(self) -> tuple[str, ...]: + return ("j0", "j1", "j2") + 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: set[str] = set((RESET_TERM_KP, RESET_TERM_KD)) if self.gain_supported else set() if self.randomization_supported: - terms.update((RESET_TERM_BODY_MASS, RESET_TERM_BODY_IPOS, RESET_TERM_GRAVITY)) + terms.update( + ( + RESET_TERM_BODY_MASS, + RESET_TERM_BODY_IPOS, + RESET_TERM_DOF_ARMATURE, + RESET_TERM_GEOM_FRICTION, + RESET_TERM_GRAVITY, + ) + ) return DomainRandomizationCapabilities( supported_reset_terms=frozenset(terms), supports_interval_body_velocity_delta=self.interval_velocity_supported, @@ -235,6 +267,12 @@ def get_body_ipos(self) -> np.ndarray: def get_gravity(self) -> np.ndarray: return self.gravity.copy() + def get_dof_armature(self) -> np.ndarray: + return self.dof_armature.copy() + + def get_geom_friction(self) -> np.ndarray: + return self.geom_friction.copy() + def apply_interval_randomization(self, plan: IntervalRandomizationPlan) -> None: self.interval_plans.append(plan) @@ -287,7 +325,9 @@ def _transaction_env( { "robot": EntityCfg( root_body_name="base", + joint_names=("j0", "j1", "j2"), body_names=body_names, + geom_names=("floor", "foot", "base_geom"), actuator_names=("a0", "a1", "a2"), ) }, @@ -494,6 +534,104 @@ def test_reset_randomization_terms_compose_with_state_and_gains_exactly_once() - np.testing.assert_allclose(payload.kd, [[3.0, 6.0, 9.0]] * 2) +def test_model_field_terms_use_cached_selectors_and_one_dense_reset_payload() -> None: + env, backend, transaction = _transaction_env(rng_seed=23) + manager = EventManager( + { + "armature": EventTermCfg( + func=mdp.joint_armature, + mode="reset", + params={ + "asset_cfg": SceneEntityCfg( + "robot", + joint_names=("j2", "j0"), + preserve_order=True, + ), + "ranges": (2.0, 2.0), + "operation": "scale", + }, + ), + "friction": EventTermCfg( + func=mdp.geom_friction, + mode="reset", + params={ + "asset_cfg": SceneEntityCfg("robot", geom_names=".*"), + "ranges": {"floor": (2.0, 2.0), "base_.*": (3.0, 3.0)}, + "operation": "scale", + "shared_random": True, + }, + ), + }, + env, + ) + ids = np.array([0, 2], dtype=np.int32) + + with transaction.scoped(ids): + mdp.reset_scene_to_default(env, 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 + assert payload.dof_armature is not None + assert payload.geom_friction is not None + np.testing.assert_allclose( + payload.dof_armature, + [[0.0] * 6 + [2.0, 2.0, 6.0]] * 2, + ) + expected_friction = backend.geom_friction.copy() + expected_friction[0, 0] *= 2.0 + expected_friction[2, 0] *= 3.0 + np.testing.assert_allclose(payload.geom_friction, [expected_friction] * 2) + + +def test_model_field_aliases_are_identical_and_capability_gaps_fail_cold() -> None: + assert mdp.dof_armature is mdp.joint_armature + env, backend, _ = _transaction_env(randomization_supported=False) + + with pytest.raises( + NotImplementedError, + match="joint_armature:robot.*dof_armature randomization.*backend 'fake'", + ): + EventManager( + { + "armature": EventTermCfg( + func=mdp.dof_armature, + mode="reset", + params={"ranges": (0.9, 1.1)}, + ) + }, + env, + ) + assert backend.set_state_calls == [] + + +@pytest.mark.parametrize( + ("func", "params", "match"), + [ + (mdp.geom_friction, {"ranges": (0.5, 1.0), "axes": [3]}, "invalid axes"), + (mdp.joint_armature, {"ranges": (-1.0, -0.5)}, "produced negative"), + ], +) +def test_model_field_invalid_requests_fail_explicitly(func, params, match: str) -> None: + env, backend, transaction = _transaction_env() + if func is mdp.geom_friction: + with pytest.raises(ValueError, match=match): + EventManager({"field": EventTermCfg(func=func, mode="reset", params=params)}, env) + else: + manager = EventManager( + {"field": EventTermCfg(func=func, mode="reset", params=params)}, + env, + ) + ids = np.array([0, 1], dtype=np.int32) + with pytest.raises(ValueError, match=match): + with transaction.scoped(ids): + mdp.reset_scene_to_default(env, ids) + manager.apply(mode="reset", env_ids=ids, global_env_step_count=0) + assert backend.set_state_calls == [] + + @pytest.mark.parametrize( ("func", "params", "match"), [ From 5713d7056a09fbe5169b1c5cf2410f3c8720aefd Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Wed, 19 Aug 2026 02:25:09 +0800 Subject: [PATCH 078/169] feat(tasks): migrate A2 to Manager-Based runtime [skip ci] (#1202) --- conf/ppo/task/a2_joystick_flat/base.yaml | 335 ++++++++ conf/ppo/task/a2_joystick_flat/mujoco.yaml | 76 +- src/unilab/assets/robots/a2/a2.xml | 16 +- src/unilab/tasks/locomotion/a2/__init__.py | 4 +- src/unilab/tasks/locomotion/a2/joystick.py | 251 +----- .../tasks/locomotion/common/commands.py | 30 - .../tasks/locomotion/common/manager_terms.py | 126 ++- .../a2/test_a2_joystick_contract.py | 765 +++++++----------- .../test_go2_joystick_stand_still.py | 62 +- .../locomotion/test_manager_gait_terms.py | 81 +- 10 files changed, 855 insertions(+), 891 deletions(-) create mode 100644 conf/ppo/task/a2_joystick_flat/base.yaml diff --git a/conf/ppo/task/a2_joystick_flat/base.yaml b/conf/ppo/task/a2_joystick_flat/base.yaml new file mode 100644 index 000000000..8ac59f4dc --- /dev/null +++ b/conf/ppo/task/a2_joystick_flat/base.yaml @@ -0,0 +1,335 @@ +# @package _global_ +# Canonical A2 flat Manager-Based task declaration. The MuJoCo owner leaf adds +# only training identity/tuning; all task behavior is declared here for Hydra. +env: + scene: + model_file: src/unilab/assets/robots/a2/scene_flat.xml + default_keyframe_name: home + entities: + robot: + root_body_name: base_link + joint_names: + - FL_hip_joint + - FL_thigh_joint + - FL_calf_joint + - FR_hip_joint + - FR_thigh_joint + - FR_calf_joint + - RL_hip_joint + - RL_thigh_joint + - RL_calf_joint + - RR_hip_joint + - RR_thigh_joint + - RR_calf_joint + body_names: [base_link] + geom_names: [floor] + actuator_names: + - FL_hip + - FL_thigh + - FL_calf + - FR_hip + - FR_thigh + - FR_calf + - RL_hip + - RL_thigh + - RL_calf + - RR_hip + - RR_thigh + - RR_calf + sim_dt: 0.01 + ctrl_dt: 0.02 + max_episode_seconds: 20.0 + observations: + policy: + _target_: unilab.managers.ObservationGroupCfg + terms: + base_ang_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.builtin_sensor + params: + sensor_name: gyro + projected_gravity: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.projected_gravity_from_sensor + params: + sensor_name: upvector + joint_pos: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_pos_rel + joint_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_vel_rel + actions: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.last_action + command: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.generated_commands + params: + command_name: twist + gait_phase: + _target_: unilab.managers.ObservationTermCfg + func: unilab.tasks.locomotion.common.manager_terms.quadruped_gait_phase + params: + frequency: 2.0 + command_name: twist + command_threshold: 0.1 + critic: + _target_: unilab.managers.ObservationGroupCfg + terms: + base_ang_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.builtin_sensor + params: + sensor_name: gyro + projected_gravity: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.projected_gravity_from_sensor + params: + sensor_name: upvector + joint_pos: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_pos_rel + joint_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_vel_rel + actions: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.last_action + command: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.generated_commands + params: + command_name: twist + gait_phase: + _target_: unilab.managers.ObservationTermCfg + func: unilab.tasks.locomotion.common.manager_terms.quadruped_gait_phase + params: + frequency: 2.0 + command_name: twist + command_threshold: 0.1 + base_lin_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.builtin_sensor + params: + sensor_name: local_linvel + actions: + joint_pos: + _target_: unilab.envs.mdp.JointPositionActionCfg + entity_name: robot + actuator_names: [".*"] + scale: 0.25 + use_default_offset: true + commands: + twist: + _target_: unilab.envs.mdp.UniformVelocityCommandCfg + entity_name: robot + resampling_time_range: [5.0, 5.0] + heading_command: false + heading_control_stiffness: 0.5 + rel_standing_envs: 0.1 + rel_heading_envs: 0.0 + rel_world_envs: 0.0 + rel_forward_envs: 0.0 + init_velocity_prob: 0.0 + ranges: + lin_vel_x: [-0.6, 1.0] + lin_vel_y: [-0.4, 0.4] + ang_vel_z: [-0.8, 0.8] + events: + reset_scene_to_default: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.reset_scene_to_default + mode: reset + reset_root_state_uniform: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.reset_root_state_uniform + mode: reset + params: + pose_range: + x: [-0.5, 0.5] + y: [-0.5, 0.5] + z: [0.0, 0.0] + roll: [0.0, 0.0] + pitch: [0.0, 0.0] + yaw: [-3.141592653589793, 3.141592653589793] + velocity_range: + x: [-0.5, 0.5] + y: [-0.5, 0.5] + z: [-0.5, 0.5] + roll: [-0.5, 0.5] + pitch: [-0.5, 0.5] + yaw: [-0.5, 0.5] + base_mass: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.randomize_rigid_body_mass + mode: reset + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + body_names: base_link + mass_distribution_params: [0.0, 8.0] + operation: add + recompute_inertia: false + base_com: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.randomize_rigid_body_com + mode: reset + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + body_names: base_link + com_range: + x: [-0.08, 0.08] + y: [-0.08, 0.08] + z: [-0.08, 0.08] + foot_friction: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.geom_friction + mode: reset + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + geom_names: floor + ranges: [0.3, 1.6] + operation: scale + axes: [0] + shared_random: true + joint_armature: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.joint_armature + mode: reset + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + joint_names: ".*" + ranges: [0.9, 1.1] + operation: scale + pd_gains: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.pd_gains + mode: reset + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + actuator_names: ".*" + kp_range: [0.9, 1.1] + kd_range: [0.9, 1.1] + operation: scale + push_robot: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.push_by_setting_velocity + mode: interval + interval_range_s: [8.0, 8.0] + is_global_time: true + params: + velocity_range: + x: [-1.0, 1.0] + y: [-1.0, 1.0] + z: [-0.5, 0.5] + roll: [0.0, 0.0] + pitch: [0.0, 0.0] + yaw: [0.0, 0.0] + terminations: + time_out: + _target_: unilab.managers.TerminationTermCfg + func: unilab.envs.mdp.time_out + time_out: true + bad_orientation: + _target_: unilab.managers.TerminationTermCfg + func: unilab.envs.mdp.bad_orientation + params: + limit_angle: 1.0471975511965976 + policy_observation_group: policy + critic_observation_group: critic + +reward: + tracking_lin_vel: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel_xy_exp + weight: 1.0 + params: + std: 0.5 + command_name: twist + tracking_ang_vel: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel_z_exp + weight: 0.4 + params: + std: 0.5 + command_name: twist + lin_vel_z: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z_l2 + weight: -5.0 + ang_vel_xy: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy_l2 + weight: -0.1 + base_height: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.base_height_l2 + weight: -100.0 + params: + target_height: 0.4 + action_rate: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.mdp.action_rate_l2 + weight: -0.02 + similar_to_default: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.joint_deviation_l1 + weight: -0.25 + contact: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.feet_phase_contact + weight: 0.5 + params: + frequency: 2.0 + command_name: twist + command_threshold: 0.1 + sensor_names: [FL_foot_contact, FR_foot_contact, RL_foot_contact, RR_foot_contact] + contact_threshold: 0.1 + stance_threshold: 0.6 + swing_feet_z: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.feet_phase_swing_height + weight: 4.0 + params: + frequency: 2.0 + command_name: twist + command_threshold: 0.1 + sensor_names: [FL_pos, FR_pos, RL_pos, RR_pos] + target_height: 0.1 + kernel: 0.01 + swing_start: 0.6 + stand_still: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.stand_still_l1 + weight: -4.0 + params: + command_name: twist + command_threshold: 0.1 + hip_deviation: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.joint_deviation_l1 + weight: -1.0 + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + joint_names: ".*_hip_joint" + stand_feet_air: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.feet_air_while_standing + weight: -1.0 + params: + command_name: twist + command_threshold: 0.1 + sensor_names: [FL_foot_contact, FR_foot_contact, RL_foot_contact, RR_foot_contact] + contact_threshold: 0.1 diff --git a/conf/ppo/task/a2_joystick_flat/mujoco.yaml b/conf/ppo/task/a2_joystick_flat/mujoco.yaml index b0863e9db..8a5f336ca 100644 --- a/conf/ppo/task/a2_joystick_flat/mujoco.yaml +++ b/conf/ppo/task/a2_joystick_flat/mujoco.yaml @@ -1,84 +1,22 @@ # @package _global_ -# A2 (leg-only Unitree A2) joystick flat task. Same isomorphic task as -# Go2JoystickFlat: 12-DOF velocity tracking with a gait phase. Robot identity -# (asset path, standing height 0.465, A2 leg PD gains) lives in A2JoystickCfg; -# this YAML carries training + reward only, mirroring the Go2 task. +defaults: + - /task/a2_joystick_flat/base + - _self_ + training: task_name: A2JoystickFlat sim_backend: mujoco algo: num_envs: 1024 - max_iterations: 500 # A2 (19.6 kg, ~2.8x Go2) + full DR needs more budget than Go2's flat task + max_iterations: 500 empirical_normalization: true obs_groups: actor: - actor + critic: + - critic policy: init_noise_std: 0.5 algorithm: learning_rate: 3.0e-4 entropy_coef: 1.0e-3 -env: - # Domain randomization for sim2real deployment. A2JoystickDomainRandomizationProvider - # caches the dof-armature + geom-friction baselines, so randomize_dof_armature and - # randomize_ground_friction are ON. Ground friction is effective because the floor - # geom is the priority geom (scene_flat.xml). Ranges reference unitree_rl_mjlab A2 - # events: joint_armature scale [0.9,1.1], foot friction [0.3,1.6]. randomize_body_mass - # stays off (base_body_mass not cached); gravity OFF (constant on flat ground). - # env.domain_rand is sim2sim ALLOWLIST (free override). - domain_rand: - randomize_base_mass: true - added_mass_range: [0.0, 8.0] - - randomize_body_mass: false # provider does not cache base_body_mass - body_mass_multiplier_range: [0.9, 1.1] - - random_com: true - com_offset_x: [-0.08, 0.08] - com_offset_y: [-0.08, 0.08] - com_offset_z: [-0.08, 0.08] - - randomize_gravity: false # flat ground: gravity constant on the real robot; randomizing only slows training. - - randomize_ground_friction: true # floor is the priority geom, so this moves the foot-ground friction - ground_friction_multiplier_range: [0.3, 1.6] # mjlab foot_friction range - - randomize_dof_armature: true - dof_armature_multiplier_range: [0.9, 1.1] # mjlab joint_armature scale - - randomize_kp: true - kp_multiplier_range: [0.9, 1.1] - - randomize_kd: true - kd_multiplier_range: [0.9, 1.1] - - push_robots: true - push_interval: 400 # control steps between base velocity pushes - max_force: [1.0, 1.0, 0.5] - push_body_name: base_link # A2 base body (Go2 uses "base"); required or push has no target. - # Standing-aware commands so the policy trains on genuine zero-command samples - # (rel_standing_envs fraction forced to stand) and resamples mid-episode every 5s. env.commands is - # a sim2sim ALLOWLIST subset (vel_limit) / free fields; rel_standing_envs and - # resampling_time are declared fields on Commands so Hydra struct mode accepts them. - commands: - rel_standing_envs: 0.1 - resampling_time: 5.0 -reward: - # command_threshold gates the phase-driven gait rewards (swing_feet_z / contact) - # so the A2 stands still at zero command instead of marching in place. - command_threshold: 0.1 - scales: - tracking_lin_vel: 1.0 - tracking_ang_vel: 0.4 - lin_vel_z: -5.0 - ang_vel_xy: -0.1 - base_height: -100.0 - action_rate: -0.02 - similar_to_default: -0.25 - contact: 0.5 - swing_feet_z: 4.0 - stand_still: -4.0 - hip_deviation: -1.0 - stand_feet_air: -1.0 # penalize feet leaving the ground at zero command (gated off during locomotion) - tracking_sigma: 0.25 - base_height_target: 0.40 diff --git a/src/unilab/assets/robots/a2/a2.xml b/src/unilab/assets/robots/a2/a2.xml index debefe71e..92a6ccb15 100644 --- a/src/unilab/assets/robots/a2/a2.xml +++ b/src/unilab/assets/robots/a2/a2.xml @@ -18,20 +18,18 @@ - + - + - + - + diff --git a/src/unilab/tasks/locomotion/a2/__init__.py b/src/unilab/tasks/locomotion/a2/__init__.py index 48a5e3355..e078a7e20 100644 --- a/src/unilab/tasks/locomotion/a2/__init__.py +++ b/src/unilab/tasks/locomotion/a2/__init__.py @@ -1,3 +1,3 @@ -from .joystick import A2JoystickCfg, A2JoystickFlatEnv +from . import joystick as joystick -__all__ = ["A2JoystickCfg", "A2JoystickFlatEnv"] +__all__ = ["joystick"] diff --git a/src/unilab/tasks/locomotion/a2/joystick.py b/src/unilab/tasks/locomotion/a2/joystick.py index 88a2672e4..f5b22df20 100644 --- a/src/unilab/tasks/locomotion/a2/joystick.py +++ b/src/unilab/tasks/locomotion/a2/joystick.py @@ -1,250 +1,7 @@ -"""A2 joystick task (leg-only Unitree A2). +"""Hydra-owned A2 flat Manager-Based production registration.""" -The A2 leg-only MJCF (robots/a2/scene_flat.xml) mirrors the Go2 joystick -sensor/geom/leg-ordering contract and uses actuators, so this -task reuses Go2WalkTask unchanged. Only the A2 identity differs: scene path, -standing pose, and per-joint PD gains. - -Asset values are aligned to the official unitree_rl_mjlab A2 (a2_constants.py): -the home keyframe matches its INIT_STATE (height 0.4, thigh 0.9, calf -1.8, -hips +-0.1), and the PD gains match its BuiltinPositionActuatorCfg — hip/thigh -kp=100/kd=4, calf kp=150/kd=6 — applied per joint at init via -position_actuator_gains and used as the per-joint baseline for kp/kd domain -randomization.""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from typing import Any - -import numpy as np - -from unilab.assets import ASSETS_ROOT_PATH from unilab.base import registry -from unilab.base.scene import SceneCfg -from unilab.dtype_config import get_global_dtype -from unilab.tasks.locomotion.common import rewards -from unilab.tasks.locomotion.common.commands import sample_commands_with_standing -from unilab.tasks.locomotion.common.dr_provider import LocomotionDRProvider -from unilab.tasks.locomotion.common.rewards import RewardContext -from unilab.tasks.locomotion.go2.base import Asset, ControlConfig -from unilab.tasks.locomotion.go2.joystick import ( - Go2DomainRandConfig, - Go2JoystickCfg, - Go2JoystickDomainRandomizationProvider, - Go2WalkTask, - RewardConfig, -) - -# Actuator/keyframe leg order: FL, FR, RL, RR x (hip, thigh, calf). -_NUM_LEGS = 4 - - -def _per_leg_gains(hip: float, thigh: float, calf: float) -> np.ndarray: - """Tile (hip, thigh, calf) gains across the four legs in actuator order.""" - return np.asarray([hip, thigh, calf] * _NUM_LEGS, dtype=np.float64) - - -@dataclass -class A2InitState: - pos = [0.0, 0.0, 0.4] - - -@dataclass -class A2Asset(Asset): - # The A2 base body is named "base_link" in its MJCF, whereas Go2 uses "base". - base_name: str = "base_link" # type: ignore[assignment] - - -@dataclass -class A2JoystickControlConfig(ControlConfig): - # Per-joint PD gains aligned to unitree_rl_mjlab: hip/thigh share Kp/Kd, the - # calf is stiffer. position_gains() expands these into 12-actuator arrays - # forwarded to the backend (overriding the static per-class kp in a2.xml). - Kp: float = 100.0 - Kd: float = 4.0 - calf_Kp: float = 150.0 # noqa: N815 - matches the Kp/Kd Hydra config convention. - calf_Kd: float = 6.0 # noqa: N815 - matches the Kp/Kd Hydra config convention. - - def position_gains(self) -> dict[str, float | np.ndarray]: - return { - "kp": _per_leg_gains(self.Kp, self.Kp, self.calf_Kp), - "kd": _per_leg_gains(self.Kd, self.Kd, self.calf_Kd), - } - - -def _a2_scene() -> SceneCfg: - return SceneCfg(model_file=str(ASSETS_ROOT_PATH / "robots" / "a2" / "scene_flat.xml")) - - -@dataclass -class A2JoystickDomainRandConfig(Go2DomainRandConfig): - # A2's base COM is uncertain in all 3 axes on the real robot. dr_utils reads - # com_offset_y/z via getattr, so they must be declared here to be settable - # from the owner YAML (Hydra struct mode rejects undeclared keys). Inherits - # com_offset_x + every other DR switch/range from Go2DomainRandConfig and - # the base DomainRandConfig; on/off + ranges are set in the owner YAML. - com_offset_y: list[float] = field(default_factory=lambda: [-0.08, 0.08]) - com_offset_z: list[float] = field(default_factory=lambda: [-0.08, 0.08]) - - -class A2JoystickDomainRandomizationProvider(Go2JoystickDomainRandomizationProvider): - """A2 reuses the Go2 joystick DR logic but supplies per-joint base gains so - randomize_kp/kd scales each actuator off its true baseline (calf off 150, - not the shared scalar 100). Without this, kp/kd DR would fall back to a - uniform ``control_config.Kp``/``Kd`` and silently weaken the calf.""" - - def _sample_commands(self, env: Any, num_reset: int) -> np.ndarray: - """Standing-aware reset commands for A2 (Go2 base stays pure-uniform). - - Mirrors rough.py's provider override: draw from ``commands.vel_limit`` and - zero a ``rel_standing_envs`` fraction so the policy trains on genuine - zero-command samples. Uses the shared ``sample_commands_with_standing`` so - the reset and mid-episode resampling paths stay a single source of truth.""" - low = np.asarray(env.cfg.commands.vel_limit[0], dtype=np.float64) - high = np.asarray(env.cfg.commands.vel_limit[1], dtype=np.float64) - return sample_commands_with_standing( - low, high, num_reset, rel_standing_envs=env.cfg.commands.rel_standing_envs - ) - - def _get_base_actuator_gains(self, env: Any) -> tuple[np.ndarray | None, np.ndarray | None]: - gains = env.cfg.control_config.position_gains() - num_actuators = env._num_action - base_kp = np.broadcast_to( - np.asarray(gains["kp"], dtype=np.float64), (num_actuators,) - ).copy() - base_kd = np.broadcast_to( - np.asarray(gains["kd"], dtype=np.float64), (num_actuators,) - ).copy() - return base_kp, base_kd - - def _get_reset_randomization_baselines( - self, env: Any - ) -> tuple[np.ndarray | None, np.ndarray | None, int | None, np.ndarray | None]: - """Snapshot the pristine model tables that reset-time DR multiplies against. - - Caches once per env (the base model is not mutated by per-env reset - randomization) via the public backend getters — no infra change, no - feature leak. Enables randomize_ground_friction (floor geom is the - priority geom, see scene_flat.xml) and randomize_dof_armature. - body_mass stays uncached (that DR switch is intentionally off).""" - cached = getattr(self, "_reset_baselines", None) - if cached is None: - backend = env._backend - base_geom_friction = backend.get_geom_friction() - ground_geom_id = backend.get_geom_id(env.cfg.asset.ground) - base_dof_armature = backend.get_dof_armature() - cached = (None, base_geom_friction, ground_geom_id, base_dof_armature) - self._reset_baselines = cached - return cached - - -@dataclass -class A2RewardConfig(RewardConfig): - # Command norm below which the phase-driven gait rewards (swing_feet_z / - # contact) switch to standing behaviour and the gait clock freezes, so the - # A2 stands still at zero command instead of marching in place. Set via the - # A2 owner YAML; default 0.0 leaves gating off. - command_threshold: float = 0.0 - - -@registry.envcfg("A2JoystickFlat") -@dataclass -class A2JoystickCfg(Go2JoystickCfg): - scene: SceneCfg = field(default_factory=_a2_scene) - init_state: A2InitState = field(default_factory=A2InitState) # type: ignore[assignment] - asset: A2Asset = field(default_factory=A2Asset) # type: ignore[assignment] - control_config: A2JoystickControlConfig = field( # type: ignore[assignment] - default_factory=A2JoystickControlConfig - ) - domain_rand: A2JoystickDomainRandConfig = field( # type: ignore[assignment] - default_factory=A2JoystickDomainRandConfig - ) - reward_config: A2RewardConfig | None = None # type: ignore[assignment] - - -@registry.env("A2JoystickFlat", sim_backend="mujoco") -class A2JoystickFlatEnv(Go2WalkTask): - """Leg-only A2 joystick task. Reuses Go2WalkTask locomotion; adds - zero-command standstill (phase freeze + gated gait rewards + standing - resample) gated by A2RewardConfig.command_threshold.""" - - _reward_cfg: A2RewardConfig - - def _make_dr_provider(self) -> LocomotionDRProvider: - return A2JoystickDomainRandomizationProvider() - - def _advance_phase(self, phase: np.ndarray) -> np.ndarray: - """Advance the gait phase, freezing envs whose command is at/below - ``command_threshold`` so a standing A2 holds phase instead of swaying.""" - cmd_norm = np.linalg.norm(self._latest_commands, axis=1) - moving = cmd_norm > self._reward_cfg.command_threshold - increment = self._cfg.ctrl_dt * self.gait_frequency * moving - return np.fmod(phase + increment, 1.0) - - def _init_reward_functions(self) -> None: - super()._init_reward_functions() - self._reward_fns.update( - { - "stand_still": rewards.stand_still, - "hip_deviation": self._reward_hip_deviation, - "stand_feet_air": self._reward_stand_feet_air, - "swing_feet_z": self._gated_swing_feet_z, - "contact": self._gated_contact, - } - ) - - def _gated_swing_feet_z(self, ctx: RewardContext) -> np.ndarray: - """Base swing reward, zeroed while standing (command at/below threshold).""" - reward = super()._reward_swing_feet_z(ctx) - cmd_norm = np.linalg.norm(ctx.info["commands"], axis=1) - active = cmd_norm > self._reward_cfg.command_threshold - return reward * active - - def _gated_contact(self, ctx: RewardContext) -> np.ndarray: - """Contact reward; while standing every foot is expected planted so a - planted robot earns full contact reward (standing branch is interleaved - per-foot, so this re-implements rather than wraps the base loop).""" - contact = self.feet_force[:, :, 2] > 0.1 - cmd_norm = np.linalg.norm(ctx.info["commands"], axis=1) - standing = cmd_norm <= self._reward_cfg.command_threshold - res = np.zeros(self._num_envs, dtype=np.float32) - for i in range(len(self._cfg.sensor.feet_force)): - is_contact = (self.feet_phase[:, i] < 0.6) | (self.gait_frequency < 1.0e-8) | standing - res += (contact[:, i] == is_contact).astype(np.float32) - return res / len(self._cfg.sensor.feet_force) - - def _reward_hip_deviation(self, ctx: RewardContext) -> np.ndarray: - """L1 deviation of the hip DOFs ([0, 3, 6, 9]) from the default pose.""" - hip_indices = [0, 3, 6, 9] - diff = ctx.dof_pos[:, hip_indices] - self.default_angles[hip_indices] - return np.asarray(np.sum(np.abs(diff), axis=1), dtype=get_global_dtype()) - - def _reward_stand_feet_air(self, ctx: RewardContext) -> np.ndarray: - """Penalize feet leaving the ground while standing (||command|| <= threshold).""" - cmd_norm = np.linalg.norm(ctx.info["commands"], axis=1) - standing = cmd_norm <= self._reward_cfg.command_threshold - in_air = np.sum(self.feet_force[:, :, 2] <= 0.1, axis=1) - return np.asarray(in_air * standing, dtype=get_global_dtype()) +from unilab.envs import ManagerBasedRlEnvCfg, make_manager_based_rl_env - def _update_commands(self, info: dict) -> None: - """Standing-aware mid-episode resample (gated by ``resampling_time``), - then stamp ``self._latest_commands`` for the phase-freeze read.""" - resampling_time = float(self._cfg.commands.resampling_time) - if resampling_time > 0.0: - commands_arr = np.asarray(info["commands"], dtype=get_global_dtype()) - interval_steps = max(int(round(resampling_time / self._cfg.ctrl_dt)), 1) - steps = np.asarray(info.get("steps", np.zeros((self._num_envs,), dtype=np.uint32))) - resample_mask = (steps > 0) & ((steps % interval_steps) == 0) - if np.any(resample_mask): - num_resample = int(np.count_nonzero(resample_mask)) - low = np.asarray(self._cfg.commands.vel_limit[0], dtype=get_global_dtype()) - high = np.asarray(self._cfg.commands.vel_limit[1], dtype=get_global_dtype()) - sampled = sample_commands_with_standing( - low, high, num_resample, rel_standing_envs=self._cfg.commands.rel_standing_envs - ) - commands_arr[resample_mask] = sampled - if self._cfg.commands.heading_command: - commands_arr[resample_mask, 2] = 0.0 - info["commands"] = commands_arr - self._latest_commands = np.asarray(info["commands"], dtype=get_global_dtype()) +registry.register_env_config("A2JoystickFlat", ManagerBasedRlEnvCfg) +registry.register_env("A2JoystickFlat", make_manager_based_rl_env, sim_backend="mujoco") diff --git a/src/unilab/tasks/locomotion/common/commands.py b/src/unilab/tasks/locomotion/common/commands.py index 3fdb2dc51..08b85fcaa 100644 --- a/src/unilab/tasks/locomotion/common/commands.py +++ b/src/unilab/tasks/locomotion/common/commands.py @@ -40,36 +40,6 @@ def zero_small_xy_commands(commands: np.ndarray, *, threshold: float = 0.2) -> N commands[:, :2] *= moving[:, None] -def sample_commands_with_standing( - low: np.ndarray, - high: np.ndarray, - num_samples: int, - *, - rel_standing_envs: float, - zero_xy_threshold: float = 0.08, -) -> np.ndarray: - """Sample velocity commands, zero small xy, then force a standing fraction. - - Single source of truth for the standing-aware command distribution shared by - the reset path (DR provider ``_sample_commands`` override) and the - mid-episode resampling path (``Go2WalkTask._update_commands``). Mirrors - rough.py's provider standing block: uniform sample in ``[low, high]``, - ``zero_small_xy_commands`` to suppress tiny lateral drift, then zero the full - command for a random ``rel_standing_envs`` fraction of rows so the policy sees - genuine zero-command (stand still) samples during training. - """ - low = np.asarray(low, dtype=get_global_dtype()) - high = np.asarray(high, dtype=get_global_dtype()) - commands = np.asarray( - np.random.uniform(low=low, high=high, size=(num_samples, 3)), dtype=get_global_dtype() - ) - zero_small_xy_commands(commands, threshold=zero_xy_threshold) - if rel_standing_envs > 0.0: - standing = np.random.uniform(size=(num_samples,)) < min(rel_standing_envs, 1.0) - commands[standing] = 0.0 - return commands - - def sample_heading_commands(env: Any, num_samples: int) -> np.ndarray: """Uniformly sample heading targets from ``env.cfg.commands.heading_range``.""" heading_range = np.asarray(env.cfg.commands.heading_range, dtype=get_global_dtype()) diff --git a/src/unilab/tasks/locomotion/common/manager_terms.py b/src/unilab/tasks/locomotion/common/manager_terms.py index a65f6c16b..25e743c1a 100644 --- a/src/unilab/tasks/locomotion/common/manager_terms.py +++ b/src/unilab/tasks/locomotion/common/manager_terms.py @@ -210,8 +210,30 @@ def joint_deviation_l1( return np.asarray(np.sum(np.abs(position - default), axis=1), dtype=get_global_dtype()) +def stand_still_l1( + env: ManagerBasedRlEnv, + command_name: str, + command_threshold: float = 0.1, + asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG, +) -> np.ndarray: + """Penalize selected joint deviation only below the commanded-motion threshold.""" + threshold = _real( + "stand_still_l1", + "command_threshold", + command_threshold, + minimum=0.0, + ) + stopped = np.linalg.norm(_command(env, "stand_still_l1", command_name), axis=1) < threshold + return np.asarray( + joint_deviation_l1(env, asset_cfg=asset_cfg) * stopped, + dtype=get_global_dtype(), + ) + + class _GaitTerm(ManagerTermBase): - _allowed_params: ClassVar[frozenset[str]] = frozenset({"frequency", "phase_offsets"}) + _allowed_params: ClassVar[frozenset[str]] = frozenset( + {"frequency", "phase_offsets", "command_name", "command_threshold"} + ) def __init__(self, cfg: ManagerTermBaseCfg, env: _GaitEnv): super().__init__(env) @@ -223,9 +245,19 @@ def __init__(self, cfg: ManagerTermBaseCfg, env: _GaitEnv): ) self._offsets = _offsets(self.name, cfg.params.get("phase_offsets", _OFFSETS)) self._step_dt = _real(self.name, "step_dt", env.step_dt, minimum=0.0, strict_minimum=True) - self._phase_value = np.asarray(0.0, dtype=get_global_dtype()) + command_name = cfg.params.get("command_name") + if command_name is not None and (not isinstance(command_name, str) or not command_name): + raise ValueError(f"{self.name} command_name must be a non-empty string or None") + self._command_name = command_name + self._command_threshold = _real( + self.name, + "command_threshold", + cfg.params.get("command_threshold", 0.0), + minimum=0.0, + ) + self._phase_value = np.zeros(env.num_envs, dtype=get_global_dtype()) self._last_counter = 0 - self._advance_to(self._counter(env)) + self._advance_to(env, self._counter(env)) def _counter(self, env: _GaitEnv) -> int: counter = env.common_step_counter @@ -235,24 +267,33 @@ def _counter(self, env: _GaitEnv) -> int: raise ValueError(f"{self.name} common_step_counter must be non-negative") return int(counter) - def _advance_to(self, counter: int) -> None: + def _moving(self, env: _GaitEnv) -> np.ndarray: + if self._command_name is None: + return np.ones(env.num_envs, dtype=np.bool_) + command = _command(env, self.name, self._command_name) + return np.linalg.norm(command, axis=1) > self._command_threshold + + def _advance_to(self, env: _GaitEnv, counter: int) -> None: delta = counter - self._last_counter if delta < 0: raise ValueError(f"{self.name} common_step_counter cannot move backwards") increment = np.asarray(self._step_dt * self._frequency, dtype=get_global_dtype()) if delta == 1: # Hot path: preserve the legacy float32 iterative phase exactly. - self._phase_value = np.fmod(self._phase_value + increment, 1.0) + self._phase_value = np.fmod(self._phase_value + increment * self._moving(env), 1.0) elif delta > 1: # Cold catch-up for a term constructed or inspected between steps. for _ in range(delta): - self._phase_value = np.fmod(self._phase_value + increment, 1.0) + self._phase_value = np.fmod( + self._phase_value + increment * self._moving(env), + 1.0, + ) self._last_counter = counter def _phase(self, env: _GaitEnv) -> np.ndarray: - self._advance_to(self._counter(env)) - phase = np.remainder(self._phase_value + self._offsets, 1.0).astype( + self._advance_to(env, self._counter(env)) + phase = np.remainder(self._phase_value[:, None] + self._offsets[None, :], 1.0).astype( get_global_dtype(), copy=False ) - return np.broadcast_to(phase, (env.num_envs, 4)).copy() + return phase class quadruped_gait_phase(_GaitTerm): @@ -320,6 +361,8 @@ def __call__(self, env: _GaitEnv, **params: Any) -> np.ndarray: expected = self._phase(env) < self._stance_threshold if self._frequency < 1.0e-8: expected.fill(True) + elif self._command_name is not None: + expected |= ~self._moving(env)[:, None] return np.mean(contact == expected, axis=1).astype(get_global_dtype(), copy=False) @@ -362,18 +405,83 @@ def __call__(self, env: _GaitEnv, **params: Any) -> np.ndarray: del params heights = self._read()[:, (2, 5, 8, 11)] swing = self._phase(env) >= self._swing_start + if self._command_name is not None: + swing &= self._moving(env)[:, None] reward = np.exp(-np.square(heights - self._target) / self._kernel) * swing return np.mean(reward, axis=1).astype(get_global_dtype(), copy=False) +class feet_air_while_standing(ManagerTermBase): + """Count feet without contact while the configured velocity command is standing.""" + + _allowed_params = frozenset( + {"sensor_names", "command_name", "command_threshold", "contact_threshold"} + ) + + def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv): + super().__init__(env) + unexpected = set(cfg.params) - self._allowed_params + if unexpected: + raise TypeError(f"{self.name} received unsupported parameters: {sorted(unexpected)}") + sensor_names = _names(self.name, cfg.params.get("sensor_names")) + command_name = cfg.params.get("command_name") + if not isinstance(command_name, str) or not command_name: + raise ValueError(f"{self.name} command_name must be a non-empty string") + self._command_name = command_name + self._command_threshold = _real( + self.name, + "command_threshold", + cfg.params.get("command_threshold", 0.1), + minimum=0.0, + ) + self._contact_threshold = _real( + self.name, + "contact_threshold", + cfg.params.get("contact_threshold", 0.1), + minimum=0.0, + ) + try: + self._view = env.scene.bind_sensor_data(sensor_names) + except (KeyError, TypeError, ValueError, NotImplementedError) as exc: + raise type(exc)( + f"Manager term '{self.name}' named-foot-sensor capability could not be " + f"materialized for {sensor_names}: {exc}" + ) from exc + if any(width not in (1, 3) for width in self._view.dimensions): + raise ValueError( + f"{self.name} contact sensors must each expose 1-D found or 3-D force; " + f"received {self._view.dimensions} on backend '{self._view.backend_type}'" + ) + starts = np.cumsum((0, *self._view.dimensions[:-1]), dtype=np.int64) + self._columns = starts + [0 if width == 1 else 2 for width in self._view.dimensions] + + def __call__(self, env: ManagerBasedRlEnv, **params: Any) -> np.ndarray: + del params + try: + values = self._view.read() + except (KeyError, TypeError, ValueError, NotImplementedError) as exc: + raise type(exc)( + f"Manager term '{self.name}' named-foot-sensor capability failed on " + f"backend '{self._view.backend_type}': {exc}" + ) from exc + contact = values[:, self._columns] > self._contact_threshold + standing = ( + np.linalg.norm(_command(env, self.name, self._command_name), axis=1) + <= self._command_threshold + ) + return np.asarray(np.sum(~contact, axis=1) * standing, dtype=get_global_dtype()) + + __all__ = [ "ang_vel_xy_l2", "base_height_l2", + "feet_air_while_standing", "feet_phase_contact", "feet_phase_swing_height", "joint_deviation_l1", "lin_vel_z_l2", "quadruped_gait_phase", + "stand_still_l1", "track_ang_vel_z_exp", "track_lin_vel_xy_exp", ] diff --git a/tests/envs/locomotion/a2/test_a2_joystick_contract.py b/tests/envs/locomotion/a2/test_a2_joystick_contract.py index eaa224dd5..0ce68ddef 100644 --- a/tests/envs/locomotion/a2/test_a2_joystick_contract.py +++ b/tests/envs/locomotion/a2/test_a2_joystick_contract.py @@ -1,87 +1,125 @@ -"""Contract tests for the A2JoystickFlat environment (leg-only Unitree A2). - -The A2 leg-only MJCF mirrors the Go2 joystick sensor/geom/leg-ordering -contract (legs FL,FR,RL,RR; foot geoms+sites FL/FR/RL/RR; Go2-named IMU/foot -sensors) and uses actuators, so the env reuses Go2WalkTask -unchanged. These tests prove the A2 model + scene + config + env chain -constructs and steps in MuJoCo as a 12-DOF joystick task.""" +"""Hydra-owned Manager-Based production contract for A2JoystickFlat.""" from __future__ import annotations import importlib -from types import SimpleNamespace +from collections.abc import Mapping, Sequence +from dataclasses import fields, is_dataclass +from pathlib import Path +from typing import Any import numpy as np import pytest +from hydra import compose, initialize_config_dir +from hydra.core.global_hydra import GlobalHydra +from omegaconf import DictConfig, OmegaConf from unilab.assets import ASSETS_ROOT_PATH - -# mjlab INIT_STATE: pos z=0.4, thigh=0.9 (all), calf=-1.8 (all), R_hip=+0.1, L_hip=-0.1. -# Asset/actuator order is FL,FR,RL,RR x (hip,thigh,calf). -_MJLAB_HOME_HEIGHT = 0.4 -_MJLAB_LEG_ANGLES = [ - -0.1, 0.9, -1.8, # FL - 0.1, 0.9, -1.8, # FR - -0.1, 0.9, -1.8, # RL - 0.1, 0.9, -1.8, # RR -] # fmt: skip -# mjlab per-joint PD gains: hip/thigh kp=100/kd=4, calf kp=150/kd=6. -_MJLAB_KP = [100.0, 100.0, 150.0] * 4 -_MJLAB_KD = [4.0, 4.0, 6.0] * 4 -# DR ranges referencing mjlab events: joint_armature scale [0.9,1.1], foot -# friction [0.3,1.6] (UniLab realises it as a multiplier on the floor geom, -# which is made the priority geom so it dictates the foot-ground friction). -_MJLAB_ARMATURE_RANGE = [0.9, 1.1] -_MJLAB_FRICTION_RANGE = [0.3, 1.6] - - -def _skip_if_no_mujoco(): - pytest.importorskip("mujoco", reason="mujoco not installed") - try: - from mujoco_uni.batch_env import BatchEnvPool # noqa: F401 - except Exception: - pytest.skip("mujoco_uni.batch_env not available") - - -def test_a2_robot_xml_compiles_with_12_position_actuators(): - """a2.xml loads standalone and exposes exactly 12 position-style leg - actuators in the FL,FR,RL,RR x hip,thigh,calf order.""" +from unilab.base import registry +from unilab.base.config_materialization import apply_cfg_overrides +from unilab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg, mdp +from unilab.tasks.locomotion.common import manager_terms +from unilab.training.backend_adapter import BackendAdapter + +ROOT_DIR = Path(__file__).parents[4] +CONF_DIR = ROOT_DIR / "conf" + +_JOINT_NAMES = ( + "FL_hip_joint", + "FL_thigh_joint", + "FL_calf_joint", + "FR_hip_joint", + "FR_thigh_joint", + "FR_calf_joint", + "RL_hip_joint", + "RL_thigh_joint", + "RL_calf_joint", + "RR_hip_joint", + "RR_thigh_joint", + "RR_calf_joint", +) +_ACTUATOR_NAMES = ( + "FL_hip", + "FL_thigh", + "FL_calf", + "FR_hip", + "FR_thigh", + "FR_calf", + "RL_hip", + "RL_thigh", + "RL_calf", + "RR_hip", + "RR_thigh", + "RR_calf", +) +_HOME_JOINT_POS = np.array( + [-0.1, 0.9, -1.8, 0.1, 0.9, -1.8, -0.1, 0.9, -1.8, 0.1, 0.9, -1.8], + dtype=np.float32, +) +_KP = np.array([100.0, 100.0, 150.0] * 4) +_KD = np.array([4.0, 4.0, 6.0] * 4) + + +def _compose() -> DictConfig: + GlobalHydra.instance().clear() + with initialize_config_dir(config_dir=str(CONF_DIR / "ppo"), version_base="1.3"): + return compose("config", overrides=["task=a2_joystick_flat/mujoco"]) + + +def _materialize() -> tuple[DictConfig, ManagerBasedRlEnvCfg, dict[str, Any]]: + hydra_cfg = _compose() + env_override = BackendAdapter(hydra_cfg, root_dir=ROOT_DIR).build_task_env_cfg_override() + env_cfg = registry.materialize_env_config("A2JoystickFlat") + assert isinstance(env_cfg, ManagerBasedRlEnvCfg) + apply_cfg_overrides(env_cfg, env_override) + env_cfg.validate() + return hydra_cfg, env_cfg, env_override + + +def _assert_no_omegaconf(value: Any) -> None: + assert not OmegaConf.is_config(value) + if is_dataclass(value) and not isinstance(value, type): + for item in fields(value): + _assert_no_omegaconf(getattr(value, item.name)) + elif isinstance(value, Mapping): + for key, item in value.items(): + _assert_no_omegaconf(key) + _assert_no_omegaconf(item) + elif isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + for item in value: + _assert_no_omegaconf(item) + + +def test_a2_asset_declares_home_pose_and_per_joint_pd_defaults() -> None: mujoco = pytest.importorskip("mujoco") - xml = ASSETS_ROOT_PATH / "robots" / "a2" / "a2.xml" - model = mujoco.MjModel.from_xml_path(str(xml)) - assert model.nu == 12 - names = [mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_ACTUATOR, i) for i in range(model.nu)] - assert names == [ - "FL_hip", - "FL_thigh", - "FL_calf", - "FR_hip", - "FR_thigh", - "FR_calf", - "RL_hip", - "RL_thigh", - "RL_calf", - "RR_hip", - "RR_thigh", - "RR_calf", - ] - # Position actuators carry an affine bias (kp in gainprm[0]); motor actuators do not. + model = mujoco.MjModel.from_xml_path(str(ASSETS_ROOT_PATH / "robots" / "a2" / "scene_flat.xml")) + + actuator_names = tuple( + mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_ACTUATOR, index) for index in range(model.nu) + ) + assert actuator_names == _ACTUATOR_NAMES affine = int(mujoco.mjtBias.mjBIAS_AFFINE) - assert all(int(model.actuator_biastype[i]) == affine for i in range(model.nu)) + assert all(int(value) == affine for value in model.actuator_biastype) + np.testing.assert_allclose(model.actuator_gainprm[:, 0], _KP) + np.testing.assert_allclose(model.actuator_biasprm[:, 1], -_KP) + np.testing.assert_allclose(model.actuator_biasprm[:, 2], -_KD) + key_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_KEY, "home") + assert key_id >= 0 + assert model.nq == 19 + assert model.key_qpos[key_id, 2] == pytest.approx(0.4) + np.testing.assert_allclose(model.key_qpos[key_id, 7:19], _HOME_JOINT_POS) + np.testing.assert_allclose(model.key_ctrl[key_id], _HOME_JOINT_POS) -def test_a2_scene_loads_with_foot_contacts_and_home_keyframe(): - """scene_flat.xml includes a2.xml + floor, exposes the four foot-contact - sensors and the joystick foot-pos/IMU sensors, and a home keyframe whose - qpos is base(7)+12 leg = 19.""" - mujoco = pytest.importorskip("mujoco") - xml = ASSETS_ROOT_PATH / "robots" / "a2" / "scene_flat.xml" - model = mujoco.MjModel.from_xml_path(str(xml)) - sensor_names = { - mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_SENSOR, i) for i in range(model.nsensor) +def test_a2_asset_exposes_manager_sensor_and_floor_friction_contract() -> None: + mujoco = pytest.importorskip("mujoco") + model = mujoco.MjModel.from_xml_path(str(ASSETS_ROOT_PATH / "robots" / "a2" / "scene_flat.xml")) + sensors = { + mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_SENSOR, index) + for index in range(model.nsensor) } - for required in [ + assert { "gyro", "local_linvel", "upvector", @@ -93,411 +131,212 @@ def test_a2_scene_loads_with_foot_contacts_and_home_keyframe(): "FR_foot_contact", "RL_foot_contact", "RR_foot_contact", - ]: - assert required in sensor_names, f"missing sensor {required}" - - # home keyframe present, qpos length = 7 (free base) + 12 (legs). - assert model.nkey >= 1 - key_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_KEY, "home") - assert key_id >= 0 - assert model.nq == 19 - # foot geoms used by the contact sensors exist. - for g in ["FL", "FR", "RL", "RR", "floor"]: - assert mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, g) >= 0 + } <= sensors - -def test_a2_home_keyframe_matches_mjlab_pose(): - """The home keyframe is aligned to mjlab's INIT_STATE: base height 0.4, - thigh 0.9, calf -1.8 on all legs, hips +-0.1 (R/L).""" - mujoco = pytest.importorskip("mujoco") - xml = ASSETS_ROOT_PATH / "robots" / "a2" / "scene_flat.xml" - model = mujoco.MjModel.from_xml_path(str(xml)) - key_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_KEY, "home") - assert key_id >= 0 - qpos = np.asarray(model.key_qpos[key_id]) - assert qpos.shape == (19,) - assert qpos[2] == pytest.approx(_MJLAB_HOME_HEIGHT) - np.testing.assert_allclose(qpos[7:19], _MJLAB_LEG_ANGLES) - # ctrl targets the same standing pose so position actuators hold it at reset. - ctrl = np.asarray(model.key_ctrl[key_id]) - np.testing.assert_allclose(ctrl, _MJLAB_LEG_ANGLES) - - -def test_a2_control_config_per_joint_gains(): - """A2JoystickControlConfig.position_gains() yields per-joint arrays matching - mjlab (calf 150/6, hip/thigh 100/4) in actuator order.""" - from unilab.tasks.locomotion.a2.joystick import A2JoystickControlConfig - - gains = A2JoystickControlConfig().position_gains() - np.testing.assert_allclose(np.asarray(gains["kp"]), _MJLAB_KP) - np.testing.assert_allclose(np.asarray(gains["kd"]), _MJLAB_KD) - - -def test_pd_control_config_position_gains_default_is_scalar(): - """Base PdControlConfig keeps the scalar gain contract (Go2 path unchanged).""" - from unilab.tasks.locomotion.common.base import PdControlConfig - - gains = PdControlConfig(Kp=35.0, Kd=0.5).position_gains() - assert gains == {"kp": 35.0, "kd": 0.5} - - -def test_a2_dr_provider_returns_per_joint_base_gains(): - """The A2 DR provider exposes per-joint base kp/kd so randomize_kp/kd scales - each joint off the correct baseline (calf off 150, not 100).""" - from unilab.tasks.locomotion.a2.joystick import ( - A2JoystickControlConfig, - A2JoystickDomainRandomizationProvider, - ) - - env = SimpleNamespace( - cfg=SimpleNamespace(control_config=A2JoystickControlConfig()), - _num_action=12, - ) - base_kp, base_kd = A2JoystickDomainRandomizationProvider()._get_base_actuator_gains(env) - np.testing.assert_allclose(np.asarray(base_kp), _MJLAB_KP) - np.testing.assert_allclose(np.asarray(base_kd), _MJLAB_KD) - - -def test_a2_floor_geom_dominates_for_friction_dr(): - """The floor geom is the priority geom (priority=2 > feet's 1) and carries - condim=6, so it dictates the foot-ground friction. This makes randomizing the - floor geom's friction actually move the contact friction (otherwise the - priority-1 feet would override it).""" - mujoco = pytest.importorskip("mujoco") - xml = ASSETS_ROOT_PATH / "robots" / "a2" / "scene_flat.xml" - model = mujoco.MjModel.from_xml_path(str(xml)) floor = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, "floor") foot = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, "FL") - assert int(model.geom_priority[floor]) == 2 + assert floor >= 0 and foot >= 0 assert int(model.geom_priority[floor]) > int(model.geom_priority[foot]) - # condim must stay 6 (feet use 6); floor would default to 3 without this. assert int(model.geom_condim[floor]) == 6 -def test_a2_dr_provider_caches_friction_and_armature_baselines(): - """The A2 DR provider caches the pristine geom-friction + dof-armature tables - (and the floor geom id) from the backend so randomize_ground_friction / - randomize_dof_armature can multiply against them. body_mass stays uncached.""" - mujoco = pytest.importorskip("mujoco") - from unilab.tasks.locomotion.a2.joystick import A2JoystickDomainRandomizationProvider - - xml = ASSETS_ROOT_PATH / "robots" / "a2" / "scene_flat.xml" - model = mujoco.MjModel.from_xml_path(str(xml)) - floor_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, "floor") - - class _StubBackend: - def get_geom_friction(self): - return np.asarray(model.geom_friction, dtype=np.float64).copy() - - def get_dof_armature(self): - return np.asarray(model.dof_armature, dtype=np.float64).copy() - - def get_geom_id(self, name): - return int(mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, name)) - - env = SimpleNamespace( - _backend=_StubBackend(), - cfg=SimpleNamespace(asset=SimpleNamespace(ground="floor")), - ) - base_body_mass, base_geom_friction, ground_geom_id, base_dof_armature = ( - A2JoystickDomainRandomizationProvider()._get_reset_randomization_baselines(env) - ) - assert base_body_mass is None # body_mass DR not enabled - assert ground_geom_id == floor_id - assert base_geom_friction.shape == (model.ngeom, 3) - np.testing.assert_allclose(base_geom_friction, model.geom_friction) - assert base_dof_armature.shape == (model.nv,) - np.testing.assert_allclose(base_dof_armature, model.dof_armature) - - -def _ensure_registered() -> None: - from unilab.base import registry +def test_a2_owner_materializes_complete_plain_manager_config() -> None: + registry.ensure_registries() + hydra_cfg, env_cfg, _ = _materialize() + + assert hydra_cfg.training.task_name == "A2JoystickFlat" + assert hydra_cfg.training.sim_backend == "mujoco" + assert hydra_cfg.algo.max_iterations == 500 + assert list(hydra_cfg.algo.obs_groups.actor) == ["actor"] + assert list(hydra_cfg.algo.obs_groups.critic) == ["critic"] + assert env_cfg.sim_dt == pytest.approx(0.01) + assert env_cfg.ctrl_dt == pytest.approx(0.02) + assert env_cfg.max_episode_seconds == pytest.approx(20.0) + assert env_cfg.policy_observation_group == "policy" + assert env_cfg.critic_observation_group == "critic" + + assert env_cfg.scene is not None + assert env_cfg.scene.model_file.endswith("robots/a2/scene_flat.xml") + assert env_cfg.scene.default_keyframe_name == "home" + robot = env_cfg.scene.entities["robot"] + assert robot.root_body_name == "base_link" + assert tuple(robot.joint_names or ()) == _JOINT_NAMES + assert tuple(robot.actuator_names or ()) == _ACTUATOR_NAMES + assert robot.body_names == ["base_link"] + assert robot.geom_names == ["floor"] + + policy_terms = [ + "base_ang_vel", + "projected_gravity", + "joint_pos", + "joint_vel", + "actions", + "command", + "gait_phase", + ] + assert list(env_cfg.observations) == ["policy", "critic"] + assert list(env_cfg.observations["policy"].terms) == policy_terms + assert list(env_cfg.observations["critic"].terms) == [*policy_terms, "base_lin_vel"] + assert env_cfg.observations["policy"].terms["gait_phase"].params == { + "frequency": 2.0, + "command_name": "twist", + "command_threshold": 0.1, + } + assert list(env_cfg.actions) == ["joint_pos"] + assert env_cfg.actions["joint_pos"].scale == pytest.approx(0.25) + command = env_cfg.commands["twist"] + assert command.resampling_time_range == [5.0, 5.0] + assert command.rel_standing_envs == pytest.approx(0.1) + assert tuple(command.ranges.lin_vel_x) == (-0.6, 1.0) + + expected_weights = { + "tracking_lin_vel": 1.0, + "tracking_ang_vel": 0.4, + "lin_vel_z": -5.0, + "ang_vel_xy": -0.1, + "base_height": -100.0, + "action_rate": -0.02, + "similar_to_default": -0.25, + "contact": 0.5, + "swing_feet_z": 4.0, + "stand_still": -4.0, + "hip_deviation": -1.0, + "stand_feet_air": -1.0, + } + assert {name: term.weight for name, term in env_cfg.rewards.items()} == expected_weights + assert env_cfg.rewards["stand_still"].func is manager_terms.stand_still_l1 + assert env_cfg.rewards["stand_feet_air"].func is manager_terms.feet_air_while_standing + assert env_cfg.rewards["hip_deviation"].params["asset_cfg"].joint_names == ".*_hip_joint" + for name in ("contact", "swing_feet_z"): + assert env_cfg.rewards[name].params["command_name"] == "twist" + assert env_cfg.rewards[name].params["command_threshold"] == pytest.approx(0.1) + + for manager_name in ("observations", "events", "rewards", "terminations"): + for term in getattr(env_cfg, manager_name).values(): + if term is None: + continue + terms = term.terms.values() if manager_name == "observations" else (term,) + for nested in terms: + if nested is None: + continue + module = nested.func.__module__ + assert ".backend." not in module + assert not any(name in module for name in (".mujoco", ".motrix", ".drake")) + + _assert_no_omegaconf(env_cfg) + + +def test_a2_owner_declares_all_randomization_as_manager_events() -> None: + _, env_cfg, _ = _materialize() + assert list(env_cfg.events) == [ + "reset_scene_to_default", + "reset_root_state_uniform", + "base_mass", + "base_com", + "foot_friction", + "joint_armature", + "pd_gains", + "push_robot", + ] + mass = env_cfg.events["base_mass"] + assert mass.func is mdp.randomize_rigid_body_mass + assert mass.params["mass_distribution_params"] == [0.0, 8.0] + assert mass.params["recompute_inertia"] is False + com = env_cfg.events["base_com"] + assert com.func is mdp.randomize_rigid_body_com + assert com.params["com_range"] == { + "x": [-0.08, 0.08], + "y": [-0.08, 0.08], + "z": [-0.08, 0.08], + } + friction = env_cfg.events["foot_friction"] + assert friction.func is mdp.geom_friction + assert friction.params["ranges"] == [0.3, 1.6] + assert friction.params["operation"] == "scale" + assert friction.params["shared_random"] is True + armature = env_cfg.events["joint_armature"] + assert armature.func is mdp.joint_armature + assert armature.params["ranges"] == [0.9, 1.1] + gains = env_cfg.events["pd_gains"] + assert gains.func is mdp.pd_gains + assert gains.params["kp_range"] == [0.9, 1.1] + assert gains.params["kd_range"] == [0.9, 1.1] + push = env_cfg.events["push_robot"] + assert push.func is mdp.push_by_setting_velocity + assert push.mode == "interval" + assert push.interval_range_s == [8.0, 8.0] + assert push.is_global_time is True + + +def test_a2_registry_has_no_legacy_config_or_runtime_fallback() -> None: registry.ensure_registries() - if not registry.contains("A2JoystickFlat"): - importlib.import_module("unilab.tasks.locomotion.a2.joystick") - - -def test_a2_joystick_registered(): - """Registers without MuJoCo (decorators run on module import).""" - from unilab.base import registry - - _ensure_registered() - assert registry.contains("A2JoystickFlat") - - -def test_a2_joystick_yaml_composes_and_targets_a2(): - """The owner YAML composes under Hydra and selects the A2JoystickFlat task - with a reward block that injects into the env's reward_config.""" - from hydra import compose, initialize - - with initialize(config_path="../../../../conf/ppo", version_base="1.3"): - cfg = compose(config_name="config", overrides=["task=a2_joystick_flat/mujoco"]) - assert cfg.training.task_name == "A2JoystickFlat" - assert cfg.training.sim_backend == "mujoco" - assert "tracking_lin_vel" in cfg.reward.scales - - -def _default_reward_cfg(): - from unilab.tasks.locomotion.a2.joystick import A2RewardConfig - - return A2RewardConfig( - scales={ - "tracking_lin_vel": 1.0, - "tracking_ang_vel": 0.2, - "lin_vel_z": -5.0, - "ang_vel_xy": -0.1, - "base_height": -100.0, - "action_rate": -0.005, - "similar_to_default": -0.1, - "contact": 0.24, - "swing_feet_z": 4.0, - "stand_still": -4.0, - "hip_deviation": -1.0, - "stand_feet_air": -1.0, - }, - tracking_sigma=0.25, - base_height_target=0.45, - command_threshold=0.1, - ) + module = importlib.import_module("unilab.tasks.locomotion.a2.joystick") + assert not hasattr(module, "A2JoystickCfg") + assert not hasattr(module, "A2JoystickFlatEnv") + assert not hasattr(module, "A2JoystickDomainRandomizationProvider") + assert registry.list_registered_envs()["A2JoystickFlat"] == { + "config_factory": "ManagerBasedRlEnvCfg", + "available_backends": ["mujoco"], + } + for legacy_override in ( + {"reward_config": {}}, + {"domain_rand": {"randomize_kp": True}}, + {"control_config": {"action_scale": 0.4}}, + ): + with pytest.raises(ValueError, match="has no attribute"): + apply_cfg_overrides(ManagerBasedRlEnvCfg(), legacy_override) -def _make_a2_env(num_envs: int = 2, domain_rand=None): - from unilab.base import registry +def test_a2_registry_executes_real_manager_runtime() -> None: + pytest.importorskip("mujoco") + try: + from mujoco_uni.batch_env import BatchEnvPool # noqa: F401 + except Exception: + pytest.skip("mujoco_uni.batch_env not available") - _ensure_registered() - override = {"reward_config": _default_reward_cfg()} - if domain_rand is not None: - override["domain_rand"] = domain_rand - return registry.make( - "A2JoystickFlat", + hydra_cfg, _, env_override = _materialize() + env = registry.make( + str(hydra_cfg.training.task_name), sim_backend="mujoco", - num_envs=num_envs, - env_cfg_override=override, + env_cfg_override=env_override, + num_envs=2, ) - - -@pytest.mark.slow -def test_a2_joystick_obs_layout_and_12_dof(): - _skip_if_no_mujoco() - env = _make_a2_env(num_envs=2) - assert env._num_action == 12 - assert env.default_angles.shape == (12,) - assert env.obs_groups_spec == {"obs": 49, "critic": 52} - - -@pytest.mark.slow -def test_a2_joystick_model_gains_are_per_joint(): - """End-to-end: the per-joint gains reach the compiled MuJoCo model — calf - actuators carry kp=150/kd=6, hip/thigh kp=100/kd=4 — and default_angles - (derived from the home keyframe) match the mjlab standing pose.""" - _skip_if_no_mujoco() - env = _make_a2_env(num_envs=2) - - # default_angles come from the home keyframe -> mjlab pose. - np.testing.assert_allclose(env.default_angles, _MJLAB_LEG_ANGLES) - - # The env forwarded per-joint arrays to the backend. - stored = env._backend._position_actuator_gains - np.testing.assert_allclose(np.asarray(stored["kp"]), _MJLAB_KP) - np.testing.assert_allclose(np.asarray(stored["kd"]), _MJLAB_KD) - - # ...and they are written into the compiled model: position actuators store - # kp in gainprm[0], and -kp / -kd in biasprm[1] / biasprm[2]. - model = env._backend._model - np.testing.assert_allclose(model.actuator_gainprm[:, 0], _MJLAB_KP) - np.testing.assert_allclose(model.actuator_biasprm[:, 1], [-v for v in _MJLAB_KP]) - np.testing.assert_allclose(model.actuator_biasprm[:, 2], [-v for v in _MJLAB_KD]) - - -@pytest.mark.slow -def test_a2_joystick_init_step_runs_finite(): - """End-to-end: init + steps must run (all A2 sensors/geoms resolve) with - finite obs/reward, proving the leg-only A2 asset satisfies the joystick - sensor contract on the hot path.""" - _skip_if_no_mujoco() - - env = _make_a2_env(num_envs=2) - state = env.init_state() - assert state.obs["obs"].shape == (2, 49) - assert state.obs["critic"].shape == (2, 52) - for _ in range(10): - state = env.step(np.zeros((2, 12), dtype=np.float64)) - assert np.isfinite(state.reward).all() - assert np.isfinite(state.obs["obs"]).all() - assert np.isfinite(state.obs["critic"]).all() - - -@pytest.mark.slow -def test_a2_joystick_dr_on_constructs_and_steps_finite(): - """With DR on (incl. base_link interval push, dof-armature and ground-friction - randomization), the env constructs and steps with finite obs/reward — proving - push_body_name resolves to a real body and the mass/COM/kp-kd/armature/friction - randomization path is sound. - - A2JoystickDomainRandomizationProvider caches the dof-armature + geom-friction - baselines, so randomize_dof_armature / randomize_ground_friction are now ON. - randomize_body_mass stays off (base_body_mass baseline not cached). The - YAML-surface is covered by test_a2_joystick_domain_rand_fully_configured.""" - _skip_if_no_mujoco() - from unilab.tasks.locomotion.a2.joystick import A2JoystickDomainRandConfig - - dr_on = A2JoystickDomainRandConfig( - randomize_base_mass=True, - added_mass_range=[0.0, 8.0], - randomize_body_mass=False, # provider does not cache base_body_mass - random_com=True, - com_offset_x=[-0.08, 0.08], - com_offset_y=[-0.08, 0.08], - com_offset_z=[-0.08, 0.08], - randomize_ground_friction=True, - ground_friction_multiplier_range=_MJLAB_FRICTION_RANGE, - randomize_dof_armature=True, - dof_armature_multiplier_range=_MJLAB_ARMATURE_RANGE, - randomize_kp=True, - randomize_kd=True, - push_robots=True, - push_interval=400, - push_body_name="base_link", - ) - env = _make_a2_env(num_envs=4, domain_rand=dr_on) - - # DR fields are active on the constructed config. - assert env._cfg.domain_rand.push_robots is True - assert env._cfg.domain_rand.push_body_name == "base_link" - assert env._cfg.domain_rand.randomize_base_mass is True - assert env._cfg.domain_rand.randomize_dof_armature is True - assert env._cfg.domain_rand.randomize_ground_friction is True - assert list(env._cfg.domain_rand.com_offset_z) == [-0.08, 0.08] - - state = env.init_state() - assert state.obs["obs"].shape == (4, 49) - # 10 steps exercises reset-time DR (mass/friction/COM/armature/kp-kd) + stepping. - for _ in range(10): - state = env.step(np.zeros((4, 12), dtype=np.float64)) - assert np.isfinite(state.reward).all() - assert np.isfinite(state.obs["obs"]).all() - assert np.isfinite(state.obs["critic"]).all() - - -def test_a2_joystick_domain_rand_fully_configured(): - """Owner YAML enables DR switches supported by A2JoystickDomainRandomizationProvider, - with mjlab-referenced ranges, 3-axis COM, base_link push target, and the - 500-iteration budget. - - randomize_dof_armature + randomize_ground_friction are ON: the A2 provider caches - their baselines (dof_armature, geom_friction + floor geom id). randomize_body_mass - stays OFF (base_body_mass not cached); gravity OFF (constant on flat ground).""" - from hydra import compose, initialize - - with initialize(config_path="../../../../conf/ppo", version_base="1.3"): - cfg = compose(config_name="config", overrides=["task=a2_joystick_flat/mujoco"]) - - dr = cfg.env.domain_rand - assert dr.randomize_base_mass is True - assert dr.randomize_body_mass is False # provider does not cache base_body_mass - assert dr.random_com is True - assert dr.randomize_gravity is False - assert dr.randomize_ground_friction is True - assert dr.randomize_dof_armature is True - assert dr.randomize_kp is True - assert dr.randomize_kd is True - assert dr.push_robots is True - # 3-axis COM present + A2-scale value - assert list(dr.com_offset_x) == [-0.08, 0.08] - assert list(dr.com_offset_y) == [-0.08, 0.08] - assert list(dr.com_offset_z) == [-0.08, 0.08] - # mjlab-referenced ranges + push target - assert list(dr.added_mass_range) == [0.0, 8.0] - assert list(dr.ground_friction_multiplier_range) == _MJLAB_FRICTION_RANGE - assert list(dr.dof_armature_multiplier_range) == _MJLAB_ARMATURE_RANGE - assert dr.push_interval == 400 - assert dr.push_body_name == "base_link" - # bumped budget - assert cfg.algo.max_iterations == 500 - - -# ── zero-command standstill (A2-owned) ─────────────────────────────── - - -def test_a2_reward_config_declares_command_threshold(): - import dataclasses - - from unilab.tasks.locomotion.a2.joystick import A2RewardConfig - - names = {f.name for f in dataclasses.fields(A2RewardConfig)} - assert "command_threshold" in names - assert {"scales", "tracking_sigma", "base_height_target"} <= names - - -def test_a2_cfg_reward_config_annotation_is_a2_type(): - from typing import get_type_hints - - from unilab.tasks.locomotion.a2.joystick import A2JoystickCfg, A2RewardConfig - - hints = get_type_hints(A2JoystickCfg) - assert A2RewardConfig in getattr(hints["reward_config"], "__args__", (hints["reward_config"],)) - - -def _a2_ctx(commands, dof_pos=None): - from unilab.tasks.locomotion.common.rewards import RewardContext - - n = commands.shape[0] - return RewardContext( - info={"commands": commands}, - linvel=np.zeros((n, 3)), - gyro=np.zeros((n, 3)), - dof_pos=np.zeros((n, 12)) if dof_pos is None else dof_pos, - num_envs=n, - default_angles=np.zeros(12), - tracking_sigma=0.25, - base_height_target=0.4, - base_height=np.zeros(n), - ) - - -def test_a2_advance_phase_freezes_standing_envs(): - from unilab.tasks.locomotion.a2.joystick import A2JoystickFlatEnv - - stub = SimpleNamespace( - _cfg=SimpleNamespace(ctrl_dt=0.02), - gait_frequency=2.0, - _reward_cfg=SimpleNamespace(command_threshold=0.1), - _latest_commands=np.array([[0.0, 0.0, 0.0], [0.5, 0.0, 0.0]]), - ) - phase = np.array([0.3, 0.3]) - out = A2JoystickFlatEnv._advance_phase(stub, phase) - assert out[0] == 0.3 - assert out[1] > 0.3 - - -def test_a2_hip_deviation_l1_over_hip_indices(): - from unilab.tasks.locomotion.a2.joystick import A2JoystickFlatEnv - - dof_pos = np.zeros((1, 12)) - dof_pos[0, [0, 3, 6, 9]] = [0.1, -0.2, 0.3, -0.4] - stub = SimpleNamespace(default_angles=np.zeros(12)) - out = A2JoystickFlatEnv._reward_hip_deviation(stub, _a2_ctx(np.zeros((1, 3)), dof_pos=dof_pos)) - assert np.isclose(out[0], 1.0) - - -def test_a2_stand_feet_air_counts_lifted_feet_when_standing(): - from unilab.tasks.locomotion.a2.joystick import A2JoystickFlatEnv - - stub = SimpleNamespace( - _reward_cfg=SimpleNamespace(command_threshold=0.1), - feet_force=np.zeros((1, 4, 3)), - ) - out = A2JoystickFlatEnv._reward_stand_feet_air(stub, _a2_ctx(np.zeros((1, 3)))) - assert out[0] == 4.0 - - -def test_a2_stand_feet_air_inactive_during_locomotion(): - from unilab.tasks.locomotion.a2.joystick import A2JoystickFlatEnv - - stub = SimpleNamespace( - _reward_cfg=SimpleNamespace(command_threshold=0.1), - feet_force=np.zeros((1, 4, 3)), - ) - out = A2JoystickFlatEnv._reward_stand_feet_air(stub, _a2_ctx(np.array([[0.5, 0.0, 0.0]]))) - assert out[0] == 0.0 + try: + assert isinstance(env, ManagerBasedRlEnv) + assert env.obs_groups_spec == {"obs": 49, "critic": 52} + assert env.action_space.shape == (12,) + action = env.action_manager.get_term("joint_pos") + assert action.target_names == list(_JOINT_NAMES) + np.testing.assert_allclose(action.offset, np.broadcast_to(_HOME_JOINT_POS, (2, 12))) + assert env.event_manager.active_terms == { + "reset": [ + "reset_scene_to_default", + "reset_root_state_uniform", + "base_mass", + "base_com", + "foot_friction", + "joint_armature", + "pd_gains", + ], + "interval": ["push_robot"], + } + + obs, info = env.reset(seed=7) + assert {name: value.shape for name, value in obs.items()} == { + "obs": (2, 49), + "critic": (2, 52), + } + assert isinstance(info, dict) + np.testing.assert_allclose( + env.scene["robot"].data.default_joint_pos, + np.broadcast_to(_HOME_JOINT_POS, (2, 12)), + ) + for _ in range(10): + state = env.step(np.zeros((2, 12), dtype=np.float32)) + assert state.reward.shape == (2,) + for value in (*state.obs.values(), state.reward): + assert isinstance(value, np.ndarray) + assert np.isfinite(value).all() + finally: + env.close() diff --git a/tests/envs/locomotion/test_go2_joystick_stand_still.py b/tests/envs/locomotion/test_go2_joystick_stand_still.py index 021584ba5..26850eddf 100644 --- a/tests/envs/locomotion/test_go2_joystick_stand_still.py +++ b/tests/envs/locomotion/test_go2_joystick_stand_still.py @@ -1,15 +1,4 @@ -"""Tests for the shared standing-command helper and Go2 flat regressions. - -The zero-command standstill behaviour itself is A2-specific and lives on -``A2JoystickFlatEnv`` (see ``tests/envs/locomotion/a2/test_a2_joystick_contract.py``). -This file keeps only: - -- ``sample_commands_with_standing`` (zero-xy + standing fraction helper) — a - shared building block used by A2 reset + A2 resample. -- Go2-flat-unchanged regressions asserting ``Go2WalkTask`` reverted to the main - baseline: the phase clock advances unconditionally and ``RewardConfig`` carries - no ``command_threshold``. -""" +"""Go2 legacy-base regressions around the A2 Manager-Based migration.""" from __future__ import annotations @@ -17,52 +6,6 @@ import numpy as np -from unilab.tasks.locomotion.common.commands import ( - sample_commands_with_standing, - zero_small_xy_commands, -) - -# ── sample_commands_with_standing (shared helper) ───────────────────── - - -def test_sample_commands_with_standing_all_standing(): - low = np.array([-1.0, -1.0, -1.0]) - high = np.array([1.0, 1.0, 1.0]) - out = sample_commands_with_standing(low, high, 64, rel_standing_envs=1.0) - assert out.shape == (64, 3) - np.testing.assert_array_equal(out, np.zeros((64, 3))) - - -def test_sample_commands_with_standing_none_standing_zeroes_small_xy(): - np.random.seed(7) - low = np.array([-1.0, -1.0, -1.0]) - high = np.array([1.0, 1.0, 1.0]) - out = sample_commands_with_standing(low, high, 2000, rel_standing_envs=0.0) - xy_norm = np.linalg.norm(out[:, :2], axis=1) - # no row below the zero-xy threshold may retain a nonzero xy. - assert np.all((xy_norm == 0.0) | (xy_norm >= 0.08)) - # with no standing fraction, not every row is fully zero (yaw survives). - assert np.any(np.abs(out[:, 2]) > 0.0) - - -def test_sample_commands_with_standing_matches_rough_block(): - """Same construction as rough.py's provider standing block, so reset (provider) - and resampling stay a single source of truth.""" - np.random.seed(3) - low = np.array([-1.0, -1.0, -1.0]) - high = np.array([1.0, 1.0, 1.0]) - out = sample_commands_with_standing(low, high, 5, rel_standing_envs=0.5) - - np.random.seed(3) - ref = np.asarray(np.random.uniform(low=low, high=high, size=(5, 3))) - zero_small_xy_commands(ref, threshold=0.08) - standing = np.random.uniform(size=(5,)) < 0.5 - ref[standing] = 0.0 - np.testing.assert_allclose(out, ref) - - -# ── Go2 flat regressions (reverted to main baseline) ────────────────── - def test_go2_advance_phase_is_unconditional(): """Go2WalkTask advances the gait clock every step regardless of command — @@ -77,8 +20,7 @@ def test_go2_advance_phase_is_unconditional(): def test_go2_reward_config_has_no_command_threshold(): - """command_threshold is A2-owned (A2RewardConfig); the Go2 base RewardConfig - must not declare it.""" + """A2's threshold stays in Hydra Manager terms, not the shared legacy config.""" import dataclasses from unilab.tasks.locomotion.go2.joystick import RewardConfig diff --git a/tests/envs/locomotion/test_manager_gait_terms.py b/tests/envs/locomotion/test_manager_gait_terms.py index 7d8fcf77d..d91bf2eb6 100644 --- a/tests/envs/locomotion/test_manager_gait_terms.py +++ b/tests/envs/locomotion/test_manager_gait_terms.py @@ -77,8 +77,12 @@ def __getitem__(self, name: str): class _Commands: - def __init__(self) -> None: - self.command = np.array([[0.5, -0.2, 0.3], [-0.1, 0.4, -0.2]], dtype=np.float32) + def __init__(self, command: np.ndarray | None = None) -> None: + self.command = ( + np.array([[0.5, -0.2, 0.3], [-0.1, 0.4, -0.2]], dtype=np.float32) + if command is None + else command + ) def get_command(self, name: str) -> np.ndarray: if name != "twist": @@ -95,6 +99,7 @@ def _env(counter: int = 0, scene: _Scene | None = None) -> ManagerBasedRlEnv: episode_length_buf=np.array([counter, 0]), step_dt=0.02, scene=scene or _Scene(), + command_manager=_Commands(), ), ) @@ -157,6 +162,76 @@ def test_gait_phase_matches_global_legacy_clock_and_ignores_episode_reset() -> N ) +def test_standing_aware_gait_freezes_only_standing_environments() -> None: + env = _env() + cast(Any, env).command_manager.command[:] = [[0.0, 0.0, 0.0], [0.5, 0.0, 0.0]] + manager = _observations(env, command_name="twist", command_threshold=0.1) + + cast(Any, env).common_step_counter = 1 + phase = manager.compute_group("policy") + assert isinstance(phase, np.ndarray) + np.testing.assert_allclose( + phase, + [[0.0, 0.5, 0.5, 0.0], [0.04, 0.54, 0.54, 0.04]], + atol=1e-7, + ) + + +def test_standing_aware_foot_rewards_gate_swing_and_expect_planted_feet() -> None: + env = _env(counter=5) + cast(Any, env).command_manager.command[:] = [[0.0, 0.0, 0.0], [0.5, 0.0, 0.0]] + params = {"frequency": 2.0, "command_name": "twist", "command_threshold": 0.1} + manager = RewardManager( + { + "contact": RewardTermCfg( + func=manager_terms.feet_phase_contact, + weight=1.0, + params={"sensor_names": CONTACTS, **params}, + ), + "swing": RewardTermCfg( + func=manager_terms.feet_phase_swing_height, + weight=1.0, + params={"sensor_names": POSITIONS, **params}, + ), + }, + env, + scale_by_dt=False, + ) + + value = manager.compute(dt=0.02) + expected_moving_swing = (1.0 + np.exp(-1.0)) / 4.0 + np.testing.assert_allclose(value, [0.5, expected_moving_swing], atol=1e-7) + + +def test_standing_penalties_match_a2_legacy_gates() -> None: + parity_env = _parity_env() + cast(Any, parity_env).command_manager.command[:] = [ + [0.0, 0.0, 0.0], + [0.5, 0.0, 0.0], + ] + stand_still = _reward_value( + parity_env, + manager_terms.stand_still_l1, + command_name="twist", + command_threshold=0.1, + ) + np.testing.assert_allclose(stand_still, [0.2, 0.0], atol=1e-7) + + foot_env = _env() + cast(Any, foot_env).command_manager.command[:] = [ + [0.0, 0.0, 0.0], + [0.5, 0.0, 0.0], + ] + feet_air = _reward_value( + foot_env, + manager_terms.feet_air_while_standing, + sensor_names=CONTACTS, + command_name="twist", + command_threshold=0.1, + ) + np.testing.assert_array_equal(feet_air, [2.0, 0.0]) + + def test_foot_rewards_match_legacy_equations_and_read_only_bound_views() -> None: scene = _Scene() manager = _rewards(_env(5, scene)) @@ -180,6 +255,8 @@ def test_foot_rewards_match_legacy_equations_and_read_only_bound_views() -> None ({"frequency": -1}, ValueError, "frequency must be at least 0.0"), ({"phase_offsets": (0, 0.5)}, ValueError, "must contain 4 values"), ({"phase_offsets": (0, True, 0.5, 0)}, TypeError, "must be a real number"), + ({"command_name": ""}, ValueError, "command_name must be a non-empty string"), + ({"command_threshold": -0.1}, ValueError, "must be at least 0.0"), ({"unknown": 1}, TypeError, "unsupported parameters"), ], ) From e37c5c42fcd9db911343ac91bbdf6e263c63a874 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Wed, 19 Aug 2026 03:39:38 +0800 Subject: [PATCH 079/169] feat(tasks): migrate Go1 flat to Manager-Based runtime [skip ci] (#1204) --- conf/appo/task/go1_joystick_flat/base.yaml | 278 ++++++++++++++ conf/appo/task/go1_joystick_flat/motrix.yaml | 38 +- conf/appo/task/go1_joystick_flat/mujoco.yaml | 16 +- .../task/go1_joystick_flat/base.yaml | 278 ++++++++++++++ .../task/td3/go1_joystick_flat/motrix.yaml | 26 +- conf/ppo/task/go1_joystick_flat/base.yaml | 278 ++++++++++++++ conf/ppo/task/go1_joystick_flat/drake.yaml | 29 +- conf/ppo/task/go1_joystick_flat/motrix.yaml | 29 +- conf/ppo/task/go1_joystick_flat/mujoco.yaml | 19 +- src/unilab/tasks/locomotion/go1/joystick.py | 14 +- tests/base/test_reward_override.py | 28 -- tests/config/test_config_system.py | 17 +- tests/config/test_reward_injection.py | 19 - .../go1/test_manager_based_flat_cfg.py | 356 ++++++++++++++++++ tests/envs/test_env_configs.py | 24 -- tests/envs/test_go1_domain_randomization.py | 61 --- tests/integration/test_appo_rsl_reward.py | 51 --- .../test_reward_injection_integration.py | 105 ------ tests/scripts/test_train_scripts.py | 17 +- 19 files changed, 1297 insertions(+), 386 deletions(-) create mode 100644 conf/appo/task/go1_joystick_flat/base.yaml create mode 100644 conf/offpolicy/task/go1_joystick_flat/base.yaml create mode 100644 conf/ppo/task/go1_joystick_flat/base.yaml create mode 100644 tests/envs/locomotion/go1/test_manager_based_flat_cfg.py delete mode 100644 tests/envs/test_go1_domain_randomization.py delete mode 100644 tests/integration/test_appo_rsl_reward.py diff --git a/conf/appo/task/go1_joystick_flat/base.yaml b/conf/appo/task/go1_joystick_flat/base.yaml new file mode 100644 index 000000000..07c3f2990 --- /dev/null +++ b/conf/appo/task/go1_joystick_flat/base.yaml @@ -0,0 +1,278 @@ +# @package _global_ +# Canonical Go1 flat Manager-Based task declaration. Backend owner leaves inherit +# this file and only override backend/algo tuning or explicitly disabled terms. +env: + scene: + model_file: src/unilab/assets/robots/go1/scene_flat.xml + default_keyframe_name: home + entities: + robot: + root_body_name: trunk + joint_names: + - FR_hip_joint + - FR_thigh_joint + - FR_calf_joint + - FL_hip_joint + - FL_thigh_joint + - FL_calf_joint + - RR_hip_joint + - RR_thigh_joint + - RR_calf_joint + - RL_hip_joint + - RL_thigh_joint + - RL_calf_joint + actuator_names: + - FR_hip + - FR_thigh + - FR_calf + - FL_hip + - FL_thigh + - FL_calf + - RR_hip + - RR_thigh + - RR_calf + - RL_hip + - RL_thigh + - RL_calf + body_names: [trunk] + sim_dt: 0.01 + ctrl_dt: 0.02 + max_episode_seconds: 20.0 + observations: + policy: + _target_: unilab.managers.ObservationGroupCfg + terms: + base_ang_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.builtin_sensor + params: + sensor_name: gyro + projected_gravity: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.projected_gravity_from_sensor + params: + sensor_name: upvector + joint_pos: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_pos_rel + joint_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_vel_rel + actions: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.last_action + command: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.generated_commands + params: + command_name: twist + gait_phase: + _target_: unilab.managers.ObservationTermCfg + func: unilab.tasks.locomotion.common.manager_terms.quadruped_gait_phase + params: + frequency: 2.0 + critic: + _target_: unilab.managers.ObservationGroupCfg + terms: + base_ang_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.builtin_sensor + params: + sensor_name: gyro + projected_gravity: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.projected_gravity_from_sensor + params: + sensor_name: upvector + joint_pos: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_pos_rel + joint_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_vel_rel + actions: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.last_action + command: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.generated_commands + params: + command_name: twist + gait_phase: + _target_: unilab.managers.ObservationTermCfg + func: unilab.tasks.locomotion.common.manager_terms.quadruped_gait_phase + params: + frequency: 2.0 + base_lin_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.builtin_sensor + params: + sensor_name: local_linvel + actions: + joint_pos: + _target_: unilab.envs.mdp.JointPositionActionCfg + entity_name: robot + actuator_names: [".*"] + scale: 0.25 + use_default_offset: true + commands: + twist: + _target_: unilab.envs.mdp.UniformVelocityCommandCfg + entity_name: robot + resampling_time_range: [20.0, 20.0] + heading_command: false + heading_control_stiffness: 0.5 + rel_standing_envs: 0.0 + rel_heading_envs: 0.0 + rel_world_envs: 0.0 + rel_forward_envs: 0.0 + init_velocity_prob: 0.0 + ranges: + lin_vel_x: [-0.6, 1.0] + lin_vel_y: [-0.4, 0.4] + ang_vel_z: [-0.8, 0.8] + events: + reset_scene_to_default: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.reset_scene_to_default + mode: reset + reset_root_state_uniform: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.reset_root_state_uniform + mode: reset + params: + pose_range: + x: [-0.5, 0.5] + y: [-0.5, 0.5] + z: [0.0, 0.0] + roll: [0.0, 0.0] + pitch: [0.0, 0.0] + yaw: [-3.141592653589793, 3.141592653589793] + velocity_range: + x: [-0.5, 0.5] + y: [-0.5, 0.5] + z: [-0.5, 0.5] + roll: [-0.5, 0.5] + pitch: [-0.5, 0.5] + yaw: [-0.5, 0.5] + base_mass: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.randomize_rigid_body_mass + mode: reset + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + body_names: trunk + mass_distribution_params: [-1.5, 1.5] + operation: add + recompute_inertia: false + base_com: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.randomize_rigid_body_com + mode: reset + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + body_names: trunk + com_range: + x: [-0.05, 0.05] + y: [0.0, 0.0] + z: [0.0, 0.0] + pd_gains: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.pd_gains + mode: reset + params: + kp_range: [35.0, 35.0] + kd_range: [0.5, 0.5] + operation: abs + push_robot: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.push_by_setting_velocity + mode: interval + interval_range_s: [15.0, 15.0] + is_global_time: true + params: + velocity_range: + x: [-1.0, 1.0] + y: [-1.0, 1.0] + z: [-0.5, 0.5] + roll: [0.0, 0.0] + pitch: [0.0, 0.0] + yaw: [0.0, 0.0] + terminations: + time_out: + _target_: unilab.managers.TerminationTermCfg + func: unilab.envs.mdp.time_out + time_out: true + bad_orientation: + _target_: unilab.managers.TerminationTermCfg + func: unilab.envs.mdp.bad_orientation + params: + limit_angle: 1.0471975511965976 + policy_observation_group: policy + critic_observation_group: critic + +reward: + tracking_lin_vel: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel_xy_exp + weight: 1.0 + params: + std: 0.5 + command_name: twist + tracking_ang_vel: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel_z_exp + weight: 0.2 + params: + std: 0.5 + command_name: twist + lin_vel_z: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z_l2 + weight: -5.0 + ang_vel_xy: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy_l2 + weight: -0.1 + base_height: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.base_height_l2 + weight: -100.0 + params: + target_height: 0.3 + action_rate: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.mdp.action_rate_l2 + weight: -0.005 + similar_to_default: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.joint_deviation_l1 + weight: -0.1 + contact: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.feet_phase_contact + # Legacy Go1 sums four matching feet while this community term returns their mean. + weight: 0.96 + params: + frequency: 2.0 + sensor_names: + - FL_foot_contact + - FR_foot_contact + - RL_foot_contact + - RR_foot_contact + contact_threshold: 0.1 + stance_threshold: 0.6 + swing_feet_z: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.feet_phase_swing_height + weight: 4.0 + params: + frequency: 2.0 + sensor_names: [FL_pos, FR_pos, RL_pos, RR_pos] + target_height: 0.1 + kernel: 0.01 + swing_start: 0.6 diff --git a/conf/appo/task/go1_joystick_flat/motrix.yaml b/conf/appo/task/go1_joystick_flat/motrix.yaml index f7343fd5d..09b622ce6 100644 --- a/conf/appo/task/go1_joystick_flat/motrix.yaml +++ b/conf/appo/task/go1_joystick_flat/motrix.yaml @@ -1,4 +1,8 @@ # @package _global_ +defaults: + - /task/go1_joystick_flat/base + - _self_ + training: task_name: Go1JoystickFlat sim_backend: motrix @@ -14,21 +18,23 @@ algo: entropy_coef: 1.0e-3 desired_kl: 0.008 env: - sim_dt: 0.01 commands: - vel_limit: - - [0.5, 0.0, 0.0] - - [0.5, 0.0, 0.0] + twist: + ranges: + lin_vel_x: [0.5, 0.5] + lin_vel_y: [0.0, 0.0] + ang_vel_z: [0.0, 0.0] + events: + push_robot: null reward: - scales: - tracking_lin_vel: 1.0 - tracking_ang_vel: 0.2 - lin_vel_z: -5.0 - ang_vel_xy: -0.1 - base_height: -100.0 - action_rate: -0.015 - action_smooth: -0.01 - similar_to_default: -0.15 - swing_feet_z: 2.0 - tracking_sigma: 0.25 - base_height_target: 0.3 + action_rate: + weight: -0.015 + action_smooth: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.mdp.action_acc_l2 + weight: -0.01 + similar_to_default: + weight: -0.15 + contact: null + swing_feet_z: + weight: 2.0 diff --git a/conf/appo/task/go1_joystick_flat/mujoco.yaml b/conf/appo/task/go1_joystick_flat/mujoco.yaml index 8646817a5..5443809b6 100644 --- a/conf/appo/task/go1_joystick_flat/mujoco.yaml +++ b/conf/appo/task/go1_joystick_flat/mujoco.yaml @@ -1,18 +1,12 @@ # @package _global_ +defaults: + - /task/go1_joystick_flat/base + - _self_ + training: task_name: Go1JoystickFlat sim_backend: mujoco algo: max_iterations: 150 reward: - scales: - tracking_lin_vel: 1.0 - tracking_ang_vel: 0.2 - lin_vel_z: -5.0 - ang_vel_xy: -0.1 - base_height: -100.0 - action_rate: -0.005 - similar_to_default: -0.1 - contact: 0.24 - tracking_sigma: 0.25 - base_height_target: 0.3 + swing_feet_z: null diff --git a/conf/offpolicy/task/go1_joystick_flat/base.yaml b/conf/offpolicy/task/go1_joystick_flat/base.yaml new file mode 100644 index 000000000..07c3f2990 --- /dev/null +++ b/conf/offpolicy/task/go1_joystick_flat/base.yaml @@ -0,0 +1,278 @@ +# @package _global_ +# Canonical Go1 flat Manager-Based task declaration. Backend owner leaves inherit +# this file and only override backend/algo tuning or explicitly disabled terms. +env: + scene: + model_file: src/unilab/assets/robots/go1/scene_flat.xml + default_keyframe_name: home + entities: + robot: + root_body_name: trunk + joint_names: + - FR_hip_joint + - FR_thigh_joint + - FR_calf_joint + - FL_hip_joint + - FL_thigh_joint + - FL_calf_joint + - RR_hip_joint + - RR_thigh_joint + - RR_calf_joint + - RL_hip_joint + - RL_thigh_joint + - RL_calf_joint + actuator_names: + - FR_hip + - FR_thigh + - FR_calf + - FL_hip + - FL_thigh + - FL_calf + - RR_hip + - RR_thigh + - RR_calf + - RL_hip + - RL_thigh + - RL_calf + body_names: [trunk] + sim_dt: 0.01 + ctrl_dt: 0.02 + max_episode_seconds: 20.0 + observations: + policy: + _target_: unilab.managers.ObservationGroupCfg + terms: + base_ang_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.builtin_sensor + params: + sensor_name: gyro + projected_gravity: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.projected_gravity_from_sensor + params: + sensor_name: upvector + joint_pos: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_pos_rel + joint_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_vel_rel + actions: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.last_action + command: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.generated_commands + params: + command_name: twist + gait_phase: + _target_: unilab.managers.ObservationTermCfg + func: unilab.tasks.locomotion.common.manager_terms.quadruped_gait_phase + params: + frequency: 2.0 + critic: + _target_: unilab.managers.ObservationGroupCfg + terms: + base_ang_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.builtin_sensor + params: + sensor_name: gyro + projected_gravity: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.projected_gravity_from_sensor + params: + sensor_name: upvector + joint_pos: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_pos_rel + joint_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_vel_rel + actions: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.last_action + command: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.generated_commands + params: + command_name: twist + gait_phase: + _target_: unilab.managers.ObservationTermCfg + func: unilab.tasks.locomotion.common.manager_terms.quadruped_gait_phase + params: + frequency: 2.0 + base_lin_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.builtin_sensor + params: + sensor_name: local_linvel + actions: + joint_pos: + _target_: unilab.envs.mdp.JointPositionActionCfg + entity_name: robot + actuator_names: [".*"] + scale: 0.25 + use_default_offset: true + commands: + twist: + _target_: unilab.envs.mdp.UniformVelocityCommandCfg + entity_name: robot + resampling_time_range: [20.0, 20.0] + heading_command: false + heading_control_stiffness: 0.5 + rel_standing_envs: 0.0 + rel_heading_envs: 0.0 + rel_world_envs: 0.0 + rel_forward_envs: 0.0 + init_velocity_prob: 0.0 + ranges: + lin_vel_x: [-0.6, 1.0] + lin_vel_y: [-0.4, 0.4] + ang_vel_z: [-0.8, 0.8] + events: + reset_scene_to_default: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.reset_scene_to_default + mode: reset + reset_root_state_uniform: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.reset_root_state_uniform + mode: reset + params: + pose_range: + x: [-0.5, 0.5] + y: [-0.5, 0.5] + z: [0.0, 0.0] + roll: [0.0, 0.0] + pitch: [0.0, 0.0] + yaw: [-3.141592653589793, 3.141592653589793] + velocity_range: + x: [-0.5, 0.5] + y: [-0.5, 0.5] + z: [-0.5, 0.5] + roll: [-0.5, 0.5] + pitch: [-0.5, 0.5] + yaw: [-0.5, 0.5] + base_mass: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.randomize_rigid_body_mass + mode: reset + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + body_names: trunk + mass_distribution_params: [-1.5, 1.5] + operation: add + recompute_inertia: false + base_com: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.randomize_rigid_body_com + mode: reset + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + body_names: trunk + com_range: + x: [-0.05, 0.05] + y: [0.0, 0.0] + z: [0.0, 0.0] + pd_gains: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.pd_gains + mode: reset + params: + kp_range: [35.0, 35.0] + kd_range: [0.5, 0.5] + operation: abs + push_robot: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.push_by_setting_velocity + mode: interval + interval_range_s: [15.0, 15.0] + is_global_time: true + params: + velocity_range: + x: [-1.0, 1.0] + y: [-1.0, 1.0] + z: [-0.5, 0.5] + roll: [0.0, 0.0] + pitch: [0.0, 0.0] + yaw: [0.0, 0.0] + terminations: + time_out: + _target_: unilab.managers.TerminationTermCfg + func: unilab.envs.mdp.time_out + time_out: true + bad_orientation: + _target_: unilab.managers.TerminationTermCfg + func: unilab.envs.mdp.bad_orientation + params: + limit_angle: 1.0471975511965976 + policy_observation_group: policy + critic_observation_group: critic + +reward: + tracking_lin_vel: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel_xy_exp + weight: 1.0 + params: + std: 0.5 + command_name: twist + tracking_ang_vel: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel_z_exp + weight: 0.2 + params: + std: 0.5 + command_name: twist + lin_vel_z: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z_l2 + weight: -5.0 + ang_vel_xy: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy_l2 + weight: -0.1 + base_height: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.base_height_l2 + weight: -100.0 + params: + target_height: 0.3 + action_rate: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.mdp.action_rate_l2 + weight: -0.005 + similar_to_default: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.joint_deviation_l1 + weight: -0.1 + contact: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.feet_phase_contact + # Legacy Go1 sums four matching feet while this community term returns their mean. + weight: 0.96 + params: + frequency: 2.0 + sensor_names: + - FL_foot_contact + - FR_foot_contact + - RL_foot_contact + - RR_foot_contact + contact_threshold: 0.1 + stance_threshold: 0.6 + swing_feet_z: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.feet_phase_swing_height + weight: 4.0 + params: + frequency: 2.0 + sensor_names: [FL_pos, FR_pos, RL_pos, RR_pos] + target_height: 0.1 + kernel: 0.01 + swing_start: 0.6 diff --git a/conf/offpolicy/task/td3/go1_joystick_flat/motrix.yaml b/conf/offpolicy/task/td3/go1_joystick_flat/motrix.yaml index 40c9f28b7..fb408b50b 100644 --- a/conf/offpolicy/task/td3/go1_joystick_flat/motrix.yaml +++ b/conf/offpolicy/task/td3/go1_joystick_flat/motrix.yaml @@ -1,4 +1,8 @@ # @package _global_ +defaults: + - /task/go1_joystick_flat/base + - _self_ + training: task_name: Go1JoystickFlat sim_backend: motrix @@ -12,18 +16,12 @@ algo: replay_buffer_n: 1024 env: commands: - vel_limit: - - [0.5, 0.0, 0.0] - - [0.5, 0.0, 0.0] + twist: + ranges: + lin_vel_x: [0.5, 0.5] + lin_vel_y: [0.0, 0.0] + ang_vel_z: [0.0, 0.0] + events: + push_robot: null reward: - scales: - tracking_lin_vel: 1.0 - tracking_ang_vel: 0.2 - lin_vel_z: -5.0 - ang_vel_xy: -0.1 - base_height: -100.0 - action_rate: -0.005 - similar_to_default: -0.1 - swing_feet_z: 4.0 - tracking_sigma: 0.25 - base_height_target: 0.3 + contact: null diff --git a/conf/ppo/task/go1_joystick_flat/base.yaml b/conf/ppo/task/go1_joystick_flat/base.yaml new file mode 100644 index 000000000..07c3f2990 --- /dev/null +++ b/conf/ppo/task/go1_joystick_flat/base.yaml @@ -0,0 +1,278 @@ +# @package _global_ +# Canonical Go1 flat Manager-Based task declaration. Backend owner leaves inherit +# this file and only override backend/algo tuning or explicitly disabled terms. +env: + scene: + model_file: src/unilab/assets/robots/go1/scene_flat.xml + default_keyframe_name: home + entities: + robot: + root_body_name: trunk + joint_names: + - FR_hip_joint + - FR_thigh_joint + - FR_calf_joint + - FL_hip_joint + - FL_thigh_joint + - FL_calf_joint + - RR_hip_joint + - RR_thigh_joint + - RR_calf_joint + - RL_hip_joint + - RL_thigh_joint + - RL_calf_joint + actuator_names: + - FR_hip + - FR_thigh + - FR_calf + - FL_hip + - FL_thigh + - FL_calf + - RR_hip + - RR_thigh + - RR_calf + - RL_hip + - RL_thigh + - RL_calf + body_names: [trunk] + sim_dt: 0.01 + ctrl_dt: 0.02 + max_episode_seconds: 20.0 + observations: + policy: + _target_: unilab.managers.ObservationGroupCfg + terms: + base_ang_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.builtin_sensor + params: + sensor_name: gyro + projected_gravity: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.projected_gravity_from_sensor + params: + sensor_name: upvector + joint_pos: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_pos_rel + joint_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_vel_rel + actions: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.last_action + command: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.generated_commands + params: + command_name: twist + gait_phase: + _target_: unilab.managers.ObservationTermCfg + func: unilab.tasks.locomotion.common.manager_terms.quadruped_gait_phase + params: + frequency: 2.0 + critic: + _target_: unilab.managers.ObservationGroupCfg + terms: + base_ang_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.builtin_sensor + params: + sensor_name: gyro + projected_gravity: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.projected_gravity_from_sensor + params: + sensor_name: upvector + joint_pos: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_pos_rel + joint_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_vel_rel + actions: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.last_action + command: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.generated_commands + params: + command_name: twist + gait_phase: + _target_: unilab.managers.ObservationTermCfg + func: unilab.tasks.locomotion.common.manager_terms.quadruped_gait_phase + params: + frequency: 2.0 + base_lin_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.builtin_sensor + params: + sensor_name: local_linvel + actions: + joint_pos: + _target_: unilab.envs.mdp.JointPositionActionCfg + entity_name: robot + actuator_names: [".*"] + scale: 0.25 + use_default_offset: true + commands: + twist: + _target_: unilab.envs.mdp.UniformVelocityCommandCfg + entity_name: robot + resampling_time_range: [20.0, 20.0] + heading_command: false + heading_control_stiffness: 0.5 + rel_standing_envs: 0.0 + rel_heading_envs: 0.0 + rel_world_envs: 0.0 + rel_forward_envs: 0.0 + init_velocity_prob: 0.0 + ranges: + lin_vel_x: [-0.6, 1.0] + lin_vel_y: [-0.4, 0.4] + ang_vel_z: [-0.8, 0.8] + events: + reset_scene_to_default: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.reset_scene_to_default + mode: reset + reset_root_state_uniform: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.reset_root_state_uniform + mode: reset + params: + pose_range: + x: [-0.5, 0.5] + y: [-0.5, 0.5] + z: [0.0, 0.0] + roll: [0.0, 0.0] + pitch: [0.0, 0.0] + yaw: [-3.141592653589793, 3.141592653589793] + velocity_range: + x: [-0.5, 0.5] + y: [-0.5, 0.5] + z: [-0.5, 0.5] + roll: [-0.5, 0.5] + pitch: [-0.5, 0.5] + yaw: [-0.5, 0.5] + base_mass: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.randomize_rigid_body_mass + mode: reset + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + body_names: trunk + mass_distribution_params: [-1.5, 1.5] + operation: add + recompute_inertia: false + base_com: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.randomize_rigid_body_com + mode: reset + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + body_names: trunk + com_range: + x: [-0.05, 0.05] + y: [0.0, 0.0] + z: [0.0, 0.0] + pd_gains: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.pd_gains + mode: reset + params: + kp_range: [35.0, 35.0] + kd_range: [0.5, 0.5] + operation: abs + push_robot: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.push_by_setting_velocity + mode: interval + interval_range_s: [15.0, 15.0] + is_global_time: true + params: + velocity_range: + x: [-1.0, 1.0] + y: [-1.0, 1.0] + z: [-0.5, 0.5] + roll: [0.0, 0.0] + pitch: [0.0, 0.0] + yaw: [0.0, 0.0] + terminations: + time_out: + _target_: unilab.managers.TerminationTermCfg + func: unilab.envs.mdp.time_out + time_out: true + bad_orientation: + _target_: unilab.managers.TerminationTermCfg + func: unilab.envs.mdp.bad_orientation + params: + limit_angle: 1.0471975511965976 + policy_observation_group: policy + critic_observation_group: critic + +reward: + tracking_lin_vel: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel_xy_exp + weight: 1.0 + params: + std: 0.5 + command_name: twist + tracking_ang_vel: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel_z_exp + weight: 0.2 + params: + std: 0.5 + command_name: twist + lin_vel_z: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z_l2 + weight: -5.0 + ang_vel_xy: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy_l2 + weight: -0.1 + base_height: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.base_height_l2 + weight: -100.0 + params: + target_height: 0.3 + action_rate: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.mdp.action_rate_l2 + weight: -0.005 + similar_to_default: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.joint_deviation_l1 + weight: -0.1 + contact: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.feet_phase_contact + # Legacy Go1 sums four matching feet while this community term returns their mean. + weight: 0.96 + params: + frequency: 2.0 + sensor_names: + - FL_foot_contact + - FR_foot_contact + - RL_foot_contact + - RR_foot_contact + contact_threshold: 0.1 + stance_threshold: 0.6 + swing_feet_z: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.feet_phase_swing_height + weight: 4.0 + params: + frequency: 2.0 + sensor_names: [FL_pos, FR_pos, RL_pos, RR_pos] + target_height: 0.1 + kernel: 0.01 + swing_start: 0.6 diff --git a/conf/ppo/task/go1_joystick_flat/drake.yaml b/conf/ppo/task/go1_joystick_flat/drake.yaml index 3c9f63c80..f1251763b 100644 --- a/conf/ppo/task/go1_joystick_flat/drake.yaml +++ b/conf/ppo/task/go1_joystick_flat/drake.yaml @@ -1,4 +1,8 @@ # @package _global_ +defaults: + - /task/go1_joystick_flat/base + - _self_ + training: task_name: Go1JoystickFlat sim_backend: drake @@ -23,26 +27,17 @@ algo: obs_groups: actor: - actor + critic: + - critic env: drake_backend_mode: batch drake_nthread: 0 - scene: - model_file: src/unilab/assets/robots/go1/scene_flat.xml - domain_rand: - randomize_base_mass: false - random_com: false - push_robots: false + events: + base_mass: null + base_com: null + pd_gains: null + push_robot: null reward: - scales: - tracking_lin_vel: 1.0 - tracking_ang_vel: 0.2 - lin_vel_z: -5.0 - ang_vel_xy: -0.1 - base_height: -100.0 - action_rate: -0.005 - similar_to_default: -0.1 - swing_feet_z: 4.0 - tracking_sigma: 0.25 - base_height_target: 0.3 + contact: null diff --git a/conf/ppo/task/go1_joystick_flat/motrix.yaml b/conf/ppo/task/go1_joystick_flat/motrix.yaml index d84ba1b03..e05a46d7b 100644 --- a/conf/ppo/task/go1_joystick_flat/motrix.yaml +++ b/conf/ppo/task/go1_joystick_flat/motrix.yaml @@ -1,4 +1,8 @@ # @package _global_ +defaults: + - /task/go1_joystick_flat/base + - _self_ + training: task_name: Go1JoystickFlat sim_backend: motrix @@ -20,6 +24,8 @@ algo: obs_groups: actor: - actor + critic: + - critic empirical_normalization: true policy: init_noise_std: 0.5 @@ -28,21 +34,16 @@ algo: entropy_coef: 1.0e-3 env: commands: - vel_limit: - - [0.5, 0.0, 0.0] - - [0.5, 0.0, 0.0] + twist: + ranges: + lin_vel_x: [0.5, 0.5] + lin_vel_y: [0.0, 0.0] + ang_vel_z: [0.0, 0.0] + events: + # Motrix has no formal root velocity-delta capability; do not fall back to force push. + push_robot: null reward: - scales: - tracking_lin_vel: 1.0 - tracking_ang_vel: 0.2 - lin_vel_z: -5.0 - ang_vel_xy: -0.1 - base_height: -100.0 - action_rate: -0.005 - similar_to_default: -0.1 - swing_feet_z: 4.0 - tracking_sigma: 0.25 - base_height_target: 0.3 + contact: null play_profile: enabled: true env: diff --git a/conf/ppo/task/go1_joystick_flat/mujoco.yaml b/conf/ppo/task/go1_joystick_flat/mujoco.yaml index 794266ad3..867b319b8 100644 --- a/conf/ppo/task/go1_joystick_flat/mujoco.yaml +++ b/conf/ppo/task/go1_joystick_flat/mujoco.yaml @@ -1,4 +1,8 @@ # @package _global_ +defaults: + - /task/go1_joystick_flat/base + - _self_ + training: task_name: Go1JoystickFlat sim_backend: mujoco @@ -21,19 +25,8 @@ algo: obs_groups: actor: - actor -reward: - scales: - tracking_lin_vel: 1.0 - tracking_ang_vel: 0.2 - lin_vel_z: -5.0 - ang_vel_xy: -0.1 - base_height: -100.0 - action_rate: -0.005 - similar_to_default: -0.1 - contact: 0.24 - swing_feet_z: 4.0 - tracking_sigma: 0.25 - base_height_target: 0.3 + critic: + - critic play_profile: enabled: true env: diff --git a/src/unilab/tasks/locomotion/go1/joystick.py b/src/unilab/tasks/locomotion/go1/joystick.py index e36089181..8df9d23d3 100644 --- a/src/unilab/tasks/locomotion/go1/joystick.py +++ b/src/unilab/tasks/locomotion/go1/joystick.py @@ -11,6 +11,7 @@ from unilab.base.np_env import NpEnvState from unilab.base.scene import SceneCfg from unilab.dtype_config import get_global_dtype +from unilab.envs import ManagerBasedRlEnvCfg, make_manager_based_rl_env from unilab.tasks.locomotion.common import rewards from unilab.tasks.locomotion.common.commands import Commands from unilab.tasks.locomotion.common.domain_rand import DomainRandConfig @@ -43,7 +44,6 @@ class JoystickSensor: feet_pos = ["FL_pos", "FR_pos", "RL_pos", "RR_pos"] -@registry.envcfg("Go1JoystickFlat") @dataclass class Go1JoystickCfg(Go1BaseCfg): scene: SceneCfg = field( # pyright: ignore[reportIncompatibleVariableOverride] @@ -82,9 +82,6 @@ def _compute_reset_obs( ) -@registry.env("Go1JoystickFlat", sim_backend="mujoco") -@registry.env("Go1JoystickFlat", sim_backend="motrix") -@registry.env("Go1JoystickFlat", sim_backend="drake") class Go1WalkTask(Go1BaseEnv): _cfg: Go1JoystickCfg # pyright: ignore[reportIncompatibleVariableOverride] @@ -254,3 +251,12 @@ def apply_action(self, actions: np.ndarray, state: NpEnvState) -> np.ndarray: exec_actions * self._cfg.control_config.action_scale + self.default_angles ) return ctrl + + +# Go1JoystickCfg and Go1WalkTask remain as the rough-task bridge and as inputs +# for pre-migration A/B benchmark adapters. Neither owns the flat production +# identity, which is Hydra-owned and uses the generic Manager-Based factories. +registry.register_env_config("Go1JoystickFlat", ManagerBasedRlEnvCfg) +registry.register_env("Go1JoystickFlat", make_manager_based_rl_env, sim_backend="mujoco") +registry.register_env("Go1JoystickFlat", make_manager_based_rl_env, sim_backend="motrix") +registry.register_env("Go1JoystickFlat", make_manager_based_rl_env, sim_backend="drake") diff --git a/tests/base/test_reward_override.py b/tests/base/test_reward_override.py index 57e5e86f0..3eca81e67 100644 --- a/tests/base/test_reward_override.py +++ b/tests/base/test_reward_override.py @@ -2,38 +2,10 @@ from typing import Any, cast -import pytest - from unilab.base import registry from unilab.base.registry import ensure_registries -def test_reward_override_go1(): - """Test Go1 reward config override.""" - ensure_registries() - - from unilab.tasks.locomotion.go1.joystick import RewardConfig - - override_config = RewardConfig( - scales={"tracking_lin_vel": 999.0}, - tracking_sigma=0.5, - base_height_target=0.5, - ) - - env = cast( - Any, - registry.make( - "Go1JoystickFlat", - num_envs=1, - sim_backend="mujoco", - env_cfg_override={"reward_config": override_config}, - ), - ) - - assert env._cfg.reward_config.scales["tracking_lin_vel"] == 999.0 - env.close() - - def test_reward_override_g1(): """Test G1 reward config override.""" ensure_registries() diff --git a/tests/config/test_config_system.py b/tests/config/test_config_system.py index d347efe37..b94883e78 100644 --- a/tests/config/test_config_system.py +++ b/tests/config/test_config_system.py @@ -68,8 +68,9 @@ def _assert_reward_populated(cfg, label: str): assert len(reward_dict["scales"]) > 0, f"{label} reward.scales must be non-empty" return - assert reward_dict, f"{label} Manager-Based reward terms must be non-empty" - for term_name, term in reward_dict.items(): + active_terms = {name: term for name, term in reward_dict.items() if term is not None} + assert active_terms, f"{label} Manager-Based reward terms must be non-empty" + for term_name, term in active_terms.items(): assert isinstance(term, dict), f"{label} reward.{term_name} must be a mapping" assert "_target_" in term, f"{label} reward.{term_name} must declare _target_" assert "func" in term, f"{label} reward.{term_name} must declare func" @@ -258,7 +259,9 @@ def test_offpolicy_td3_go1_joystick_flat_motrix_composes(): assert cfg.training.task_name == "Go1JoystickFlat" assert cfg.training.sim_backend == "motrix" assert cfg.algo.algo == "td3" - assert cfg.reward.scales.tracking_lin_vel == pytest.approx(1.0) + assert cfg.reward.tracking_lin_vel.weight == pytest.approx(1.0) + assert cfg.reward.contact is None + assert cfg.env.events.push_robot is None def test_offpolicy_g1_walk_flat_motrix_preserves_backend_specific_algo_value(): @@ -401,8 +404,12 @@ def test_ppo_go1_motrix_preserves_reward_and_algo_values(): assert cfg.algo.empirical_normalization is True assert cfg.algo.policy.init_noise_std == pytest.approx(0.5) assert cfg.algo.algorithm.learning_rate == pytest.approx(3.0e-4) - assert cfg.reward.scales.tracking_lin_vel == pytest.approx(1.0) - assert cfg.env.commands.vel_limit == [[0.5, 0.0, 0.0], [0.5, 0.0, 0.0]] + assert cfg.reward.tracking_lin_vel.weight == pytest.approx(1.0) + assert cfg.reward.contact is None + assert cfg.env.commands.twist.ranges.lin_vel_x == [0.5, 0.5] + assert cfg.env.commands.twist.ranges.lin_vel_y == [0.0, 0.0] + assert cfg.env.commands.twist.ranges.ang_vel_z == [0.0, 0.0] + assert cfg.env.events.push_robot is None def test_ppo_go2_motrix_preserves_backend_env_overrides(): diff --git a/tests/config/test_reward_injection.py b/tests/config/test_reward_injection.py index 78765e3cb..0c9e5e8fc 100644 --- a/tests/config/test_reward_injection.py +++ b/tests/config/test_reward_injection.py @@ -75,22 +75,3 @@ def test_reward_config_conversion(): assert hasattr(env._cfg.reward_config, "scales") assert env._cfg.reward_config.scales["tracking_lin_vel"] == 2.0 env.close() - - # Test Go1 config - registry auto-converts dict to RewardConfig - go1_dict = { - "scales": {"tracking_lin_vel": 1.0, "base_height": -100.0}, - "tracking_sigma": 0.25, - "base_height_target": 0.3, - } - env = cast( - Any, - registry.make( - "Go1JoystickFlat", - num_envs=1, - sim_backend="mujoco", - env_cfg_override={"reward_config": go1_dict}, - ), - ) - assert hasattr(env._cfg.reward_config, "scales") - assert env._cfg.reward_config.scales["tracking_lin_vel"] == 1.0 - env.close() diff --git a/tests/envs/locomotion/go1/test_manager_based_flat_cfg.py b/tests/envs/locomotion/go1/test_manager_based_flat_cfg.py new file mode 100644 index 000000000..3b071ddd9 --- /dev/null +++ b/tests/envs/locomotion/go1/test_manager_based_flat_cfg.py @@ -0,0 +1,356 @@ +"""Hydra-owned production contract for the Go1 flat Manager-Based task.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from copy import deepcopy +from dataclasses import fields, is_dataclass +from pathlib import Path +from typing import Any + +import numpy as np +import pytest +from hydra import compose, initialize_config_dir +from hydra.core.global_hydra import GlobalHydra +from omegaconf import DictConfig, OmegaConf + +from unilab.base import registry +from unilab.base.config_materialization import apply_cfg_overrides +from unilab.envs import ( + ManagerBasedRlEnv, + ManagerBasedRlEnvCfg, + make_manager_based_rl_env, + mdp, +) +from unilab.tasks.locomotion.common import manager_terms +from unilab.training.backend_adapter import BackendAdapter + +ROOT_DIR = Path(__file__).parents[4] +CONF_DIR = ROOT_DIR / "conf" + +_JOINT_NAMES = ( + "FR_hip_joint", + "FR_thigh_joint", + "FR_calf_joint", + "FL_hip_joint", + "FL_thigh_joint", + "FL_calf_joint", + "RR_hip_joint", + "RR_thigh_joint", + "RR_calf_joint", + "RL_hip_joint", + "RL_thigh_joint", + "RL_calf_joint", +) +_ACTUATOR_NAMES = tuple(name.removesuffix("_joint") for name in _JOINT_NAMES) +_HOME_JOINT_POS = np.array( + [0.0, 0.9, -1.8, 0.0, 0.9, -1.8, 0.0, 1.0, -1.8, 0.0, 1.0, -1.8], + dtype=np.float32, +) +_RESET_EVENTS = ("reset_scene_to_default", "reset_root_state_uniform") +_DR_EVENTS = ("base_mass", "base_com", "pd_gains") +_BASE_REWARDS = ( + "tracking_lin_vel", + "tracking_ang_vel", + "lin_vel_z", + "ang_vel_xy", + "base_height", + "action_rate", + "similar_to_default", +) + +_OWNER_CASES = ( + pytest.param( + "ppo", + ("task=go1_joystick_flat/mujoco",), + "mujoco", + (*_RESET_EVENTS, *_DR_EVENTS, "push_robot"), + (*_BASE_REWARDS, "contact", "swing_feet_z"), + False, + id="ppo-mujoco", + ), + pytest.param( + "ppo", + ("task=go1_joystick_flat/motrix",), + "motrix", + (*_RESET_EVENTS, *_DR_EVENTS), + (*_BASE_REWARDS, "swing_feet_z"), + True, + id="ppo-motrix", + ), + pytest.param( + "ppo", + ("task=go1_joystick_flat/drake",), + "drake", + _RESET_EVENTS, + (*_BASE_REWARDS, "swing_feet_z"), + False, + id="ppo-drake", + ), + pytest.param( + "appo", + ("task=go1_joystick_flat/mujoco",), + "mujoco", + (*_RESET_EVENTS, *_DR_EVENTS, "push_robot"), + (*_BASE_REWARDS, "contact"), + False, + id="appo-mujoco", + ), + pytest.param( + "appo", + ("task=go1_joystick_flat/motrix",), + "motrix", + (*_RESET_EVENTS, *_DR_EVENTS), + (*_BASE_REWARDS, "swing_feet_z", "action_smooth"), + True, + id="appo-motrix", + ), + pytest.param( + "offpolicy", + ("algo=td3", "task=td3/go1_joystick_flat/motrix"), + "motrix", + (*_RESET_EVENTS, *_DR_EVENTS), + (*_BASE_REWARDS, "swing_feet_z"), + True, + id="td3-motrix", + ), +) + + +def _compose(config_group: str, overrides: Sequence[str]) -> DictConfig: + GlobalHydra.instance().clear() + with initialize_config_dir(config_dir=str(CONF_DIR / config_group), version_base="1.3"): + return compose("config", overrides=list(overrides)) + + +def _materialize( + config_group: str, overrides: Sequence[str] +) -> tuple[DictConfig, ManagerBasedRlEnvCfg, dict[str, Any]]: + hydra_cfg = _compose(config_group, overrides) + env_override = BackendAdapter(hydra_cfg, root_dir=ROOT_DIR).build_task_env_cfg_override() + env_cfg = registry.materialize_env_config("Go1JoystickFlat") + assert isinstance(env_cfg, ManagerBasedRlEnvCfg) + apply_cfg_overrides(env_cfg, env_override) + env_cfg.validate() + return hydra_cfg, env_cfg, env_override + + +def _assert_no_omegaconf(value: Any) -> None: + assert not OmegaConf.is_config(value) + if is_dataclass(value) and not isinstance(value, type): + for item in fields(value): + _assert_no_omegaconf(getattr(value, item.name)) + elif isinstance(value, Mapping): + for key, item in value.items(): + _assert_no_omegaconf(key) + _assert_no_omegaconf(item) + elif isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + for item in value: + _assert_no_omegaconf(item) + + +@pytest.mark.parametrize( + "config_group,overrides,backend,expected_events,expected_rewards,fixed_command", + _OWNER_CASES, +) +def test_go1_flat_owner_materializes_complete_plain_manager_cfg( + config_group: str, + overrides: tuple[str, ...], + backend: str, + expected_events: tuple[str, ...], + expected_rewards: tuple[str, ...], + fixed_command: bool, +) -> None: + registry.ensure_registries() + hydra_cfg, env_cfg, _ = _materialize(config_group, overrides) + + assert hydra_cfg.training.task_name == "Go1JoystickFlat" + assert hydra_cfg.training.sim_backend == backend + assert env_cfg.sim_dt == pytest.approx(0.01) + assert env_cfg.ctrl_dt == pytest.approx(0.02) + assert env_cfg.max_episode_seconds == pytest.approx(20.0) + assert env_cfg.policy_observation_group == "policy" + assert env_cfg.critic_observation_group == "critic" + assert env_cfg.scale_rewards_by_dt is True + + assert env_cfg.scene is not None + assert env_cfg.scene.model_file.endswith("robots/go1/scene_flat.xml") + assert env_cfg.scene.default_keyframe_name == "home" + robot = env_cfg.scene.entities["robot"] + assert robot.root_body_name == "trunk" + assert tuple(robot.joint_names or ()) == _JOINT_NAMES + assert tuple(robot.actuator_names or ()) == _ACTUATOR_NAMES + assert robot.body_names == ["trunk"] + + observation_terms = [ + "base_ang_vel", + "projected_gravity", + "joint_pos", + "joint_vel", + "actions", + "command", + "gait_phase", + ] + assert list(env_cfg.observations) == ["policy", "critic"] + assert list(env_cfg.observations["policy"].terms) == observation_terms + assert list(env_cfg.observations["critic"].terms) == [*observation_terms, "base_lin_vel"] + assert list(env_cfg.actions) == ["joint_pos"] + assert env_cfg.actions["joint_pos"].scale == pytest.approx(0.25) + assert list(env_cfg.terminations) == ["time_out", "bad_orientation"] + assert ( + tuple(name for name, term in env_cfg.events.items() if term is not None) == expected_events + ) + assert tuple(name for name, term in env_cfg.rewards.items() if term is not None) == ( + expected_rewards + ) + + command = env_cfg.commands["twist"] + assert command.resampling_time_range == [20.0, 20.0] + ranges = command.ranges + if fixed_command: + assert tuple(ranges.lin_vel_x) == (0.5, 0.5) + assert tuple(ranges.lin_vel_y) == (0.0, 0.0) + assert tuple(ranges.ang_vel_z) == (0.0, 0.0) + else: + assert tuple(ranges.lin_vel_x) == (-0.6, 1.0) + + if env_cfg.events["base_mass"] is not None: + mass = env_cfg.events["base_mass"] + assert mass.func is mdp.randomize_rigid_body_mass + assert mass.params["mass_distribution_params"] == [-1.5, 1.5] + assert mass.params["recompute_inertia"] is False + com = env_cfg.events["base_com"] + assert com.func is mdp.randomize_rigid_body_com + assert com.params["com_range"] == { + "x": [-0.05, 0.05], + "y": [0.0, 0.0], + "z": [0.0, 0.0], + } + gains = env_cfg.events["pd_gains"] + assert gains.func is mdp.pd_gains + assert gains.params["kp_range"] == [35.0, 35.0] + assert gains.params["kd_range"] == [0.5, 0.5] + assert gains.params["operation"] == "abs" + + push = env_cfg.events["push_robot"] + if push is not None: + assert push.func is mdp.push_by_setting_velocity + assert push.interval_range_s == [15.0, 15.0] + assert push.is_global_time is True + + contact = env_cfg.rewards["contact"] + if contact is not None: + assert contact.func is manager_terms.feet_phase_contact + # Legacy Go1 returned a four-foot sum; the community term returns a mean. + assert contact.weight == pytest.approx(4.0 * 0.24) + action_smooth = env_cfg.rewards.get("action_smooth") + if action_smooth is not None: + assert action_smooth.func is mdp.action_acc_l2 + assert action_smooth.weight == pytest.approx(-0.01) + + for manager_name in ("observations", "events", "rewards", "terminations"): + for term in getattr(env_cfg, manager_name).values(): + if term is None: + continue + nested_terms = term.terms.values() if manager_name == "observations" else (term,) + for nested in nested_terms: + if nested is None: + continue + module = nested.func.__module__ + assert ".backend." not in module + assert not any(name in module for name in (".mujoco", ".motrix", ".drake")) + + _assert_no_omegaconf(env_cfg) + + +def test_go1_flat_registry_is_manager_only_and_rough_owns_the_legacy_bridge() -> None: + registry.ensure_registries() + from unilab.tasks.locomotion.go1.joystick import Go1JoystickCfg, Go1WalkTask + from unilab.tasks.locomotion.go1.rough import Go1JoystickRoughCfg, Go1JoystickRoughEnv + + assert registry.list_registered_envs()["Go1JoystickFlat"] == { + "config_factory": "ManagerBasedRlEnvCfg", + "available_backends": ["mujoco", "motrix", "drake"], + } + assert Go1JoystickRoughCfg.__bases__ == (Go1JoystickCfg,) + assert Go1JoystickRoughEnv.__bases__ == (Go1WalkTask,) + + for legacy_override in ( + {"reward_config": {}}, + {"domain_rand": {"randomize_base_mass": True}}, + {"control_config": {"Kp": 35.0}}, + ): + with pytest.raises(ValueError, match="has no attribute"): + apply_cfg_overrides(ManagerBasedRlEnvCfg(), legacy_override) + + +@pytest.mark.parametrize( + ("backend", "owner", "expected_events"), + ( + ( + "mujoco", + "task=go1_joystick_flat/mujoco", + {"reset": [*_RESET_EVENTS, *_DR_EVENTS], "interval": ["push_robot"]}, + ), + ( + "motrix", + "task=go1_joystick_flat/motrix", + {"reset": [*_RESET_EVENTS, *_DR_EVENTS]}, + ), + ), +) +def test_go1_flat_registry_executes_real_manager_runtime( + backend: str, + owner: str, + expected_events: dict[str, list[str]], +) -> None: + registry.ensure_registries() + hydra_cfg, _, env_override = _materialize("ppo", (owner,)) + try: + env = registry.make( + str(hydra_cfg.training.task_name), + sim_backend=backend, + env_cfg_override=env_override, + num_envs=2, + ) + except ImportError as exc: + pytest.skip(f"{backend} runtime unavailable: {exc}") + + try: + assert isinstance(env, ManagerBasedRlEnv) + assert env.obs_groups_spec == {"obs": 49, "critic": 52} + assert env.action_space.shape == (12,) + action = env.action_manager.get_term("joint_pos") + assert action.target_names == list(_JOINT_NAMES) + np.testing.assert_allclose(action.offset, np.broadcast_to(_HOME_JOINT_POS, (2, 12))) + assert env.event_manager.active_terms == expected_events + + obs, info = env.reset(seed=7) + assert {name: value.shape for name, value in obs.items()} == { + "obs": (2, 49), + "critic": (2, 52), + } + assert isinstance(info, dict) + np.testing.assert_allclose( + env.scene["robot"].data.default_joint_pos, + np.broadcast_to(_HOME_JOINT_POS, (2, 12)), + ) + for _ in range(5): + state = env.step(np.zeros((2, 12), dtype=np.float32)) + for value in (*state.obs.values(), state.reward): + assert isinstance(value, np.ndarray) + assert np.isfinite(value).all() + finally: + env.close() + + +def test_go1_motrix_velocity_push_request_fails_closed() -> None: + _, motrix_cfg, _ = _materialize("ppo", ("task=go1_joystick_flat/motrix",)) + _, mujoco_cfg, _ = _materialize("ppo", ("task=go1_joystick_flat/mujoco",)) + motrix_cfg.events["push_robot"] = deepcopy(mujoco_cfg.events["push_robot"]) + + try: + with pytest.raises(NotImplementedError, match="interval root velocity delta.*motrix"): + make_manager_based_rl_env(motrix_cfg, num_envs=1, backend_type="motrix") + except ImportError as exc: + pytest.skip(f"motrix runtime unavailable: {exc}") diff --git a/tests/envs/test_env_configs.py b/tests/envs/test_env_configs.py index 2f3776974..bfd00f775 100644 --- a/tests/envs/test_env_configs.py +++ b/tests/envs/test_env_configs.py @@ -2052,7 +2052,6 @@ def test_g1_motion_tracking_clip_end_does_not_override_true_termination(): # Environments that don't need special config overrides _STANDARD_ENVS = [ - "Go1JoystickFlat", "Go1JoystickRough", "Go2WJoystickFlat", "Go2WJoystickRough", @@ -2148,29 +2147,6 @@ def _assert_mujoco_position_gains( np.testing.assert_allclose(pool.get_field(0, "kd")[actuator_ids], kd) -def test_go1_env_initializes_kp_kd_into_pool(default_go1_reward_config): - _require_mujoco_runtime() - ensure_registries() - from unilab.base import registry - - env = cast( - Any, - registry.make( - "Go1JoystickFlat", - num_envs=2, - sim_backend="mujoco", - env_cfg_override={ - "reward_config": default_go1_reward_config, - "control_config": {"Kp": 12.0, "Kd": 0.7}, - }, - ), - ) - try: - _assert_mujoco_position_gains(env, kp=12.0, kd=0.7) - finally: - env.close() - - def test_allegro_env_initializes_kp_kd_into_pool(default_allegro_reward_config): _require_mujoco_runtime() ensure_registries() diff --git a/tests/envs/test_go1_domain_randomization.py b/tests/envs/test_go1_domain_randomization.py deleted file mode 100644 index 3004ec5f4..000000000 --- a/tests/envs/test_go1_domain_randomization.py +++ /dev/null @@ -1,61 +0,0 @@ -from __future__ import annotations - -from typing import Any - -import numpy as np -import pytest - -pytest.importorskip("mujoco", reason="mujoco not installed") - -try: - from mujoco_uni.batch_env import BatchEnvPool as _ # noqa: F401 -except Exception: - pytest.skip( - "mujoco_uni.batch_env not available (platform/libstdc++ issue)", allow_module_level=True - ) - -from unilab.base.registry import ensure_registries - - -def test_go1_mujoco_reset_applies_base_mass_and_com_domain_randomization( - default_go1_reward_config: dict[str, Any], -) -> None: - ensure_registries() - - from unilab.base import registry - - env = registry.make( - "Go1JoystickFlat", - num_envs=4, - sim_backend="mujoco", - env_cfg_override={"reward_config": default_go1_reward_config}, - ) - env_obj: Any = env - try: - env_obj.init_state() - backend: Any = env_obj._backend - base_body_id = int(backend._base_body_id) - body_mass = np.stack( - [backend._pool.get_field(i, "body_mass") for i in range(env_obj.num_envs)] - ) - body_ipos = np.stack( - [backend._pool.get_field(i, "body_ipos") for i in range(env_obj.num_envs)] - ) - body_ipos = body_ipos.reshape(env_obj.num_envs, -1, 3) - - base_mass = float(backend._base_body_mass[base_body_id]) - base_ipos = np.asarray(backend._base_body_ipos[base_body_id]) - randomized_base_mass = body_mass[:, base_body_id] - randomized_base_ipos = body_ipos[:, base_body_id] - delta_x = randomized_base_ipos[:, 0] - base_ipos[0] - delta_yz = randomized_base_ipos[:, 1:] - base_ipos[1:] - - assert np.unique(np.round(randomized_base_mass, 6)).size > 1 - assert np.unique(np.round(delta_x, 6)).size > 1 - assert np.all(randomized_base_mass >= base_mass - 1.5) - assert np.all(randomized_base_mass <= base_mass + 1.5) - assert np.all(delta_x >= -0.05) - assert np.all(delta_x <= 0.05) - np.testing.assert_allclose(delta_yz, 0.0) - finally: - env_obj.close() diff --git a/tests/integration/test_appo_rsl_reward.py b/tests/integration/test_appo_rsl_reward.py deleted file mode 100644 index cecefdbc3..000000000 --- a/tests/integration/test_appo_rsl_reward.py +++ /dev/null @@ -1,51 +0,0 @@ -"""Test APPO reward injection.""" - -import pytest - - -def test_appo_reward_override(): - """Test APPO with reward override.""" - from unilab.base import registry - from unilab.base.registry import ensure_registries - - ensure_registries() - - reward_dict = { - "scales": {"tracking_lin_vel": 888.0}, - "tracking_sigma": 0.3, - "base_height_target": 0.35, - } - - env = registry.make( - "Go1JoystickFlat", - num_envs=2, - sim_backend="mujoco", - env_cfg_override={"reward_config": reward_dict}, - ) - - assert env._cfg.reward_config.scales["tracking_lin_vel"] == 888.0 - env.close() - - -def test_rsl_rl_reward_override(): - """Test RSL-RL with reward override.""" - from unilab.base import registry - from unilab.base.registry import ensure_registries - - ensure_registries() - - reward_dict = { - "scales": {"tracking_lin_vel": 777.0}, - "tracking_sigma": 0.2, - "base_height_target": 0.32, - } - - env = registry.make( - "Go1JoystickFlat", - num_envs=2, - sim_backend="mujoco", - env_cfg_override={"reward_config": reward_dict}, - ) - - assert env._cfg.reward_config.scales["tracking_lin_vel"] == 777.0 - env.close() diff --git a/tests/integration/test_reward_injection_integration.py b/tests/integration/test_reward_injection_integration.py index e549e0a90..d884a2eb3 100644 --- a/tests/integration/test_reward_injection_integration.py +++ b/tests/integration/test_reward_injection_integration.py @@ -1,10 +1,5 @@ """Integration test for reward config injection in training.""" -from typing import Any, cast - -import numpy as np -import pytest - def test_reward_injection_in_offpolicy_env_override(): """Test reward config is injected without requiring accelerator hardware.""" @@ -32,103 +27,3 @@ def test_reward_injection_in_offpolicy_env_override(): reward_dict = env_cfg_override["reward_config"] assert reward_dict["scales"]["tracking_lin_vel"] == 2.0 assert reward_dict["scales"]["alive"] == 10.0 - - -def test_reward_override_propagation(): - """Test reward override propagates through multiprocess collector.""" - from unilab.base import registry - from unilab.base.registry import ensure_registries - from unilab.tasks.locomotion.go1.joystick import RewardConfig - - ensure_registries() - - # Create custom reward config - custom_config = RewardConfig( - scales={ - "tracking_lin_vel": 5.0, - "tracking_ang_vel": 0.5, - "lin_vel_z": -10.0, - }, - tracking_sigma=0.5, - base_height_target=0.4, - ) - - # Create env with override - env = cast( - Any, - registry.make( - "Go1JoystickFlat", - num_envs=4, - sim_backend="mujoco", - env_cfg_override={"reward_config": custom_config}, - ), - ) - - # Verify override was applied - assert env._cfg.reward_config.scales["tracking_lin_vel"] == 5.0 - assert env._cfg.reward_config.tracking_sigma == 0.5 - - # Test reward computation uses overridden scales - env.init_state() - env.reset(np.array([0, 1, 2, 3], dtype=np.int32)) - - # Take a step and verify reward is computed - actions = np.zeros((4, env.action_space.shape[0]), dtype=np.float32) - state = env.step(actions) - - assert state.reward is not None - assert len(state.reward) == 4 - - env.close() - - -def test_backward_compatibility_no_reward_config(): - """Test env requires reward config - should fail without it.""" - from unilab.base import registry - from unilab.base.registry import ensure_registries - - ensure_registries() - - # Should fail without reward_config - with pytest.raises(ValueError, match="reward_config must be provided"): - registry.make( - "Go1JoystickFlat", - num_envs=2, - sim_backend="mujoco", - ) - - -def test_zero_scale_skips_computation(): - """Test that reward functions with scale=0 are skipped.""" - from unilab.base import registry - from unilab.base.registry import ensure_registries - from unilab.tasks.locomotion.go1.joystick import RewardConfig - - ensure_registries() - - # Set all scales to 0 except one - custom_config = RewardConfig( - scales={ - "tracking_lin_vel": 1.0, - "tracking_ang_vel": 0.0, # Should be skipped - "lin_vel_z": 0.0, # Should be skipped - }, - tracking_sigma=0.25, - base_height_target=0.3, - ) - - env = cast( - Any, - registry.make( - "Go1JoystickFlat", - num_envs=2, - sim_backend="mujoco", - env_cfg_override={"reward_config": custom_config}, - ), - ) - - # Verify only non-zero scales are in config - assert env._cfg.reward_config.scales["tracking_lin_vel"] == 1.0 - assert env._cfg.reward_config.scales["tracking_ang_vel"] == 0.0 - - env.close() diff --git a/tests/scripts/test_train_scripts.py b/tests/scripts/test_train_scripts.py index 99f97bc3d..017f41732 100644 --- a/tests/scripts/test_train_scripts.py +++ b/tests/scripts/test_train_scripts.py @@ -692,7 +692,11 @@ def test_ppo_go1_drake_batch_config_matches_current_contact_support(): assert cfg.env.drake_backend_mode == "batch" assert cfg.env.drake_nthread == 0 - assert "contact" not in cfg.reward.scales + assert cfg.reward.contact is None + assert cfg.env.events.base_mass is None + assert cfg.env.events.base_com is None + assert cfg.env.events.pd_gains is None + assert cfg.env.events.push_robot is None def test_ppo_go2_drake_batch_config_matches_go2_training_defaults(): @@ -721,9 +725,14 @@ def test_build_ppo_env_cfg_override_go1_motrix( env_cfg_override = mod.build_ppo_env_cfg_override(cfg) - # env_cfg_override has reward + env preset commands - assert env_cfg_override["reward_config"]["scales"]["tracking_lin_vel"] == pytest.approx(1.0) - assert env_cfg_override["commands"]["vel_limit"] == [[0.5, 0.0, 0.0], [0.5, 0.0, 0.0]] + assert env_cfg_override["rewards"]["tracking_lin_vel"]["weight"] == pytest.approx(1.0) + assert env_cfg_override["rewards"]["contact"] is None + assert env_cfg_override["commands"]["twist"]["ranges"] == { + "lin_vel_x": [0.5, 0.5], + "lin_vel_y": [0.0, 0.0], + "ang_vel_z": [0.0, 0.0], + } + assert env_cfg_override["events"]["push_robot"] is None def test_build_ppo_env_cfg_override_g1_motrix( From c87d3f4faea6c856b255db6b49dc3b24b457c0b0 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Wed, 19 Aug 2026 04:55:18 +0800 Subject: [PATCH 080/169] feat(tasks): migrate Stewart balance to Manager-Based runtime [skip ci] (#1206) --- .../task/sac/stewart_balance/drake.yaml | 11 +- conf/offpolicy/task/stewart_balance/base.yaml | 126 ++ conf/ppo/task/stewart_balance/base.yaml | 128 ++ conf/ppo/task/stewart_balance/drake.yaml | 11 +- conf/ppo/task/stewart_balance/motrix.yaml | 10 +- src/unilab/base/backend/mujoco/backend.py | 20 +- .../tasks/manipulation/stewart/__init__.py | 17 +- .../tasks/manipulation/stewart/balance.py | 1160 +++++++++++------ tests/envs/test_stewart.py | 450 +++++-- 9 files changed, 1412 insertions(+), 521 deletions(-) create mode 100644 conf/offpolicy/task/stewart_balance/base.yaml create mode 100644 conf/ppo/task/stewart_balance/base.yaml diff --git a/conf/offpolicy/task/sac/stewart_balance/drake.yaml b/conf/offpolicy/task/sac/stewart_balance/drake.yaml index 5a672e516..4992f2146 100644 --- a/conf/offpolicy/task/sac/stewart_balance/drake.yaml +++ b/conf/offpolicy/task/sac/stewart_balance/drake.yaml @@ -1,4 +1,8 @@ # @package _global_ +defaults: + - /task/stewart_balance/base + - _self_ + training: task_name: StewartBalance sim_backend: drake @@ -29,10 +33,3 @@ algo: env: drake_backend_mode: batch drake_nthread: 20 - -reward: - scales: - center: 0.7 - progress: 0.6 - still: 3.0 - fall_penalty: -6.0 diff --git a/conf/offpolicy/task/stewart_balance/base.yaml b/conf/offpolicy/task/stewart_balance/base.yaml new file mode 100644 index 000000000..5dfdf4279 --- /dev/null +++ b/conf/offpolicy/task/stewart_balance/base.yaml @@ -0,0 +1,126 @@ +# @package _global_ +# Off-policy copy of the canonical Stewart Manager-Based declaration. Hydra +# config groups have separate search roots, so this intentionally mirrors PPO. +env: + scene: + model_file: src/unilab/assets/robots/stewart/scene.xml + entities: + stewart: + root_body_name: ball + actuator_names: [a0, a1, a2, a3, a4, a5] + body_names: + - ball + - top + - leg00 + - leg10 + - leg01 + - leg11 + - leg02 + - leg12 + - top_connect00 + - top_connect10 + - top_connect01 + - top_connect11 + - top_connect02 + - top_connect12 + sim_dt: 0.004 + ctrl_dt: 0.02 + max_episode_seconds: 24.0 + render_spacing: 4.5 + observations: + policy: + _target_: unilab.managers.ObservationGroupCfg + terms: + balance: + _target_: unilab.managers.ObservationTermCfg + func: unilab.tasks.manipulation.stewart.balance.StewartObservation + params: + entity_name: stewart + action_name: tilt + ball_body_name: ball + top_body_name: top + target_rotation_limit_deg: 6.0 + vel_smooth: 0.25 + actions: + tilt: + _target_: unilab.tasks.manipulation.stewart.balance.StewartTiltActionCfg + entity_name: stewart + actuator_names: [a0, a1, a2, a3, a4, a5] + top_body_name: top + ball_body_name: ball + leg_body_names: [leg00, leg10, leg01, leg11, leg02, leg12] + top_connect_body_names: + - top_connect00 + - top_connect10 + - top_connect01 + - top_connect11 + - top_connect02 + - top_connect12 + raw_action_clip: [-1.0, 1.0] + target_rotation_limit_deg: 6.0 + action_smooth: 0.60 + center_control_radius: 0.25 + center_control_min_gain: 0.15 + events: + reset_scene_to_default: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.reset_scene_to_default + mode: reset + reset_ball: + _target_: unilab.managers.EventTermCfg + func: unilab.tasks.manipulation.stewart.balance.StewartBallReset + mode: reset + params: + entity_name: stewart + platform_radius: 0.8 + init_ball_radius_ratio: 0.18 + ball_home_z: 1.2 + terminations: + balance_state: + _target_: unilab.managers.TerminationTermCfg + func: unilab.tasks.manipulation.stewart.balance.StewartBalanceState + params: + observation_group: policy + observation_term: balance + platform_radius: 0.8 + fall_radius: 0.5 + top_center_z: 1.0 + still_xy: 0.12 + still_vel: 0.07 + still_xy_hysteresis: 1.15 + still_vel_hysteresis: 1.20 + zero_vel_thresh: 0.07 + still_steps_needed: 5 + time_out: + _target_: unilab.managers.TerminationTermCfg + func: unilab.envs.mdp.time_out + time_out: true + scale_rewards_by_dt: false + policy_observation_group: policy + critic_observation_group: null + +reward: + center: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.manipulation.stewart.balance.center_reward + weight: 0.7 + params: + state_term_name: balance_state + progress: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.manipulation.stewart.balance.progress_reward + weight: 0.6 + params: + state_term_name: balance_state + still: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.manipulation.stewart.balance.still_reward + weight: 3.0 + params: + state_term_name: balance_state + fall: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.manipulation.stewart.balance.fall_reward + weight: -6.0 + params: + state_term_name: balance_state diff --git a/conf/ppo/task/stewart_balance/base.yaml b/conf/ppo/task/stewart_balance/base.yaml new file mode 100644 index 000000000..87db4efc5 --- /dev/null +++ b/conf/ppo/task/stewart_balance/base.yaml @@ -0,0 +1,128 @@ +# @package _global_ +# Canonical Stewart Manager-Based task declaration. Backend leaves own only +# backend identity and algorithm/runtime tuning. +env: + scene: + model_file: src/unilab/assets/robots/stewart/scene.xml + entities: + stewart: + root_body_name: ball + actuator_names: [a0, a1, a2, a3, a4, a5] + body_names: + - ball + - top + - leg00 + - leg10 + - leg01 + - leg11 + - leg02 + - leg12 + - top_connect00 + - top_connect10 + - top_connect01 + - top_connect11 + - top_connect02 + - top_connect12 + # The stiff closed-loop model requires a physics step no larger than ~0.005 s. + sim_dt: 0.004 + ctrl_dt: 0.02 + max_episode_seconds: 24.0 + render_spacing: 4.5 + observations: + policy: + _target_: unilab.managers.ObservationGroupCfg + terms: + balance: + _target_: unilab.managers.ObservationTermCfg + func: unilab.tasks.manipulation.stewart.balance.StewartObservation + params: + entity_name: stewart + action_name: tilt + ball_body_name: ball + top_body_name: top + target_rotation_limit_deg: 6.0 + vel_smooth: 0.25 + actions: + tilt: + _target_: unilab.tasks.manipulation.stewart.balance.StewartTiltActionCfg + entity_name: stewart + actuator_names: [a0, a1, a2, a3, a4, a5] + top_body_name: top + ball_body_name: ball + leg_body_names: [leg00, leg10, leg01, leg11, leg02, leg12] + top_connect_body_names: + - top_connect00 + - top_connect10 + - top_connect01 + - top_connect11 + - top_connect02 + - top_connect12 + raw_action_clip: [-1.0, 1.0] + target_rotation_limit_deg: 6.0 + action_smooth: 0.60 + center_control_radius: 0.25 + center_control_min_gain: 0.15 + events: + reset_scene_to_default: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.reset_scene_to_default + mode: reset + reset_ball: + _target_: unilab.managers.EventTermCfg + func: unilab.tasks.manipulation.stewart.balance.StewartBallReset + mode: reset + params: + entity_name: stewart + platform_radius: 0.8 + init_ball_radius_ratio: 0.18 + ball_home_z: 1.2 + terminations: + balance_state: + _target_: unilab.managers.TerminationTermCfg + func: unilab.tasks.manipulation.stewart.balance.StewartBalanceState + params: + observation_group: policy + observation_term: balance + platform_radius: 0.8 + fall_radius: 0.5 + top_center_z: 1.0 + still_xy: 0.12 + still_vel: 0.07 + still_xy_hysteresis: 1.15 + still_vel_hysteresis: 1.20 + zero_vel_thresh: 0.07 + still_steps_needed: 5 + time_out: + _target_: unilab.managers.TerminationTermCfg + func: unilab.envs.mdp.time_out + time_out: true + # Legacy Stewart rewards were discrete per-control-step values, not rates. + scale_rewards_by_dt: false + policy_observation_group: policy + critic_observation_group: null + +reward: + center: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.manipulation.stewart.balance.center_reward + weight: 0.7 + params: + state_term_name: balance_state + progress: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.manipulation.stewart.balance.progress_reward + weight: 0.6 + params: + state_term_name: balance_state + still: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.manipulation.stewart.balance.still_reward + weight: 3.0 + params: + state_term_name: balance_state + fall: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.manipulation.stewart.balance.fall_reward + weight: -6.0 + params: + state_term_name: balance_state diff --git a/conf/ppo/task/stewart_balance/drake.yaml b/conf/ppo/task/stewart_balance/drake.yaml index cdd440832..6be5f7d0d 100644 --- a/conf/ppo/task/stewart_balance/drake.yaml +++ b/conf/ppo/task/stewart_balance/drake.yaml @@ -1,4 +1,8 @@ # @package _global_ +defaults: + - /task/stewart_balance/base + - _self_ + training: task_name: StewartBalance sim_backend: drake @@ -42,10 +46,3 @@ algo: gamma: 0.99 lam: 0.95 save_interval: 50 - -reward: - scales: - center: 0.7 - progress: 0.6 - still: 3.0 - fall_penalty: -6.0 diff --git a/conf/ppo/task/stewart_balance/motrix.yaml b/conf/ppo/task/stewart_balance/motrix.yaml index f018225f8..e15d86cf5 100644 --- a/conf/ppo/task/stewart_balance/motrix.yaml +++ b/conf/ppo/task/stewart_balance/motrix.yaml @@ -2,6 +2,10 @@ # Stewart-platform ball-balancing (motrix). A short, runnable PPO baseline, not # tuned for best final performance (raise max_iterations / num_envs for higher # success rates). +defaults: + - /task/stewart_balance/base + - _self_ + training: task_name: StewartBalance sim_backend: motrix @@ -40,12 +44,6 @@ algo: gamma: 0.99 lam: 0.95 save_interval: 50 -reward: - scales: - center: 0.7 - progress: 0.6 - still: 3.0 - fall_penalty: -6.0 play_profile: enabled: true env: diff --git a/src/unilab/base/backend/mujoco/backend.py b/src/unilab/base/backend/mujoco/backend.py index b14a3f03d..1decc5389 100644 --- a/src/unilab/base/backend/mujoco/backend.py +++ b/src/unilab/base/backend/mujoco/backend.py @@ -445,11 +445,29 @@ def _load_base_model(self) -> mujoco.MjModel: for tmp_path in reversed(tmp_paths): os.remove(tmp_path) - self._tracked_body_ids = tracked_body_ids if self.add_body_sensors: + # MjSpec compilation can reorder bodies expanded from . + # Sensor columns follow ``valid_bnames`` insertion order, so rebuild + # the name-to-column map from the final compiled model instead of + # retaining IDs from the pre-injection source model. + self._tracked_body_ids = [ + mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, name) for name in valid_bnames + ] + missing = [ + name + for name, body_id in zip(valid_bnames, self._tracked_body_ids, strict=True) + if body_id < 0 + ] + if missing: + raise ValueError( + "Injected MuJoCo body tracking sensors reference bodies missing from " + f"the compiled model: {missing}" + ) self._body_id_to_tracked_idx = np.full(model.nbody, -1, dtype=int) for idx, bid in enumerate(self._tracked_body_ids): self._body_id_to_tracked_idx[bid] = idx + else: + self._tracked_body_ids = tracked_body_ids self._valid_bnames = valid_bnames self._configure_model(model) return model diff --git a/src/unilab/tasks/manipulation/stewart/__init__.py b/src/unilab/tasks/manipulation/stewart/__init__.py index 0d377c02e..5b5f99fa4 100644 --- a/src/unilab/tasks/manipulation/stewart/__init__.py +++ b/src/unilab/tasks/manipulation/stewart/__init__.py @@ -1,4 +1,15 @@ -from . import balance # registers StewartBalance via @registry decorators -from .balance import StewartBalanceCfg, StewartBalanceEnv +from . import balance as balance +from .balance import StewartBalanceState as StewartBalanceState +from .balance import StewartBallReset as StewartBallReset +from .balance import StewartObservation as StewartObservation +from .balance import StewartTiltAction as StewartTiltAction +from .balance import StewartTiltActionCfg as StewartTiltActionCfg -__all__ = ["StewartBalanceCfg", "StewartBalanceEnv"] +__all__ = [ + "StewartBalanceState", + "StewartBallReset", + "StewartObservation", + "StewartTiltAction", + "StewartTiltActionCfg", + "balance", +] diff --git a/src/unilab/tasks/manipulation/stewart/balance.py b/src/unilab/tasks/manipulation/stewart/balance.py index 0a35290c2..03bea58ed 100644 --- a/src/unilab/tasks/manipulation/stewart/balance.py +++ b/src/unilab/tasks/manipulation/stewart/balance.py @@ -1,28 +1,21 @@ -"""Stewart-platform ball-balancing task. +"""Manager-Based terms for Stewart-platform ball balancing. -A 6-DOF parallel (Stewart) platform balances a free ball on its top plate. The -policy commands a 2-D platform tilt (roll, pitch); an inverse-kinematics step -converts the commanded plate pose into the six prismatic leg lengths that the -position actuators track. The objective is to bring the ball to the plate center -and hold it still. The platform base is welded to the world. +Hydra owns the production task declaration. This module contains only the +task-specific NumPy terms and the generic Manager-Based registry binding. """ from __future__ import annotations -from dataclasses import dataclass, field +from dataclasses import dataclass +from numbers import Real +from typing import TYPE_CHECKING, Any, Protocol, cast -import gymnasium as gym import numpy as np -from unilab.assets import ASSETS_ROOT_PATH from unilab.base import registry -from unilab.base.backend import create_backend, env_backend_kwargs -from unilab.base.base import EnvCfg -from unilab.base.np_env import NpEnv, NpEnvState -from unilab.base.scene import SceneCfg -from unilab.dr.provider import DomainRandomizationProvider -from unilab.dr.types import DomainRandomizationCapabilities, ResetPlan from unilab.dtype_config import get_global_dtype +from unilab.envs import ManagerBasedRlEnvCfg, make_manager_based_rl_env +from unilab.managers import ActionTerm, ActionTermCfg, ManagerTermBase, ManagerTermBaseCfg from unilab.utils.geometry import np_roll_pitch_from_quat from unilab.utils.rotation import ( np_quat_apply_batched, @@ -33,414 +26,771 @@ np_quat_to_axis_angle, ) -# Leg base / top-connect bodies in actuator order (a0..a5 -> slide00,slide10, -# slide01,slide11,slide02,slide12), produced by the `replicate count=3` in the XML. -_LEG_BODY_NAMES = ["leg00", "leg10", "leg01", "leg11", "leg02", "leg12"] -_TOP_CONNECT_NAMES = [ - "top_connect00", - "top_connect10", - "top_connect01", - "top_connect11", - "top_connect02", - "top_connect12", -] - -_OBS_DIM = 15 -_ACTION_DIM = 2 +if TYPE_CHECKING: + from unilab.base.entity import Entity + from unilab.managers._types import ManagerBasedRlEnv + from unilab.managers.observation_manager import ObservationManager + from unilab.managers.termination_manager import TerminationManager + class _StewartEnv(ManagerBasedRlEnv, Protocol): + common_step_counter: int + observation_manager: ObservationManager -@dataclass -class StewartRewardConfig: - """Reward shaping for the ball-balancing task (see `_compute_reward`).""" - scales: dict[str, float] = field( - default_factory=lambda: {"center": 0.7, "progress": 0.6, "still": 3.0} - ) - fall_penalty: float = -6.0 +_ACTION_DIM = 2 +_LEG_COUNT = 6 + + +def _real( + term: str, + name: str, + value: Any, + *, + minimum: float | None = None, + maximum: float | None = None, + strict_minimum: bool = False, +) -> float: + if isinstance(value, (bool, np.bool_)) or not isinstance(value, Real): + raise TypeError(f"{term} {name} must be a real number") + result = float(value) + if not np.isfinite(result): + raise ValueError(f"{term} {name} must be finite") + if minimum is not None and (result <= minimum if strict_minimum else result < minimum): + relation = "greater than" if strict_minimum else "at least" + raise ValueError(f"{term} {name} must be {relation} {minimum}") + if maximum is not None and result > maximum: + raise ValueError(f"{term} {name} must be at most {maximum}") + return result + + +def _name(term: str, name: str, value: Any) -> str: + if not isinstance(value, str) or not value: + raise ValueError(f"{term} {name} must be a non-empty string") + return value + + +def _names(term: str, name: str, value: Any, *, count: int) -> tuple[str, ...]: + if isinstance(value, (str, bytes)) or not isinstance(value, (tuple, list)): + raise TypeError(f"{term} {name} must be a sequence of {count} strings") + result = tuple(value) + if len(result) != count: + raise ValueError(f"{term} {name} must contain {count} names, got {len(result)}") + if any(not isinstance(item, str) or not item for item in result): + raise ValueError(f"{term} {name} must contain non-empty strings") + if len(set(result)) != count: + raise ValueError(f"{term} {name} must contain unique names: {result}") + return result + + +def _pair(term: str, name: str, value: Any) -> tuple[float, float]: + if isinstance(value, (str, bytes)) or not isinstance(value, (tuple, list)): + raise TypeError(f"{term} {name} must be a numeric (min, max) pair") + if len(value) != 2: + raise ValueError(f"{term} {name} must contain two values") + lower = _real(term, f"{name}[0]", value[0]) + upper = _real(term, f"{name}[1]", value[1]) + if lower > upper: + raise ValueError(f"{term} {name} lower bound {lower} exceeds upper bound {upper}") + return lower, upper + + +def _env_ids(env: ManagerBasedRlEnv, env_ids: np.ndarray | slice | None) -> np.ndarray: + if env_ids is None: + return np.arange(env.num_envs, dtype=np.int32) + if isinstance(env_ids, slice): + return np.arange(env.num_envs, dtype=np.int32)[env_ids] + return env_ids + + +def _body_id(entity: Entity, name: str, *, term: str) -> int: + ids, resolved = entity.find_bodies(name) + if len(ids) != 1 or resolved != [name]: + raise ValueError(f"{term} body selector {name!r} did not resolve exactly once") + return ids[0] + + +def _body_ids(entity: Entity, names: tuple[str, ...], *, term: str) -> np.ndarray: + ids, resolved = entity.find_bodies(names, preserve_order=True) + if tuple(resolved) != names: + raise ValueError(f"{term} body selectors resolved in an unexpected order: {resolved}") + result = np.asarray(ids, dtype=np.intp) + result.setflags(write=False) + return result + + +def _relative_ball_state( + entity: Entity, + *, + ball_body_id: int, + top_body_id: int, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + body_pos = entity.data.body_link_pos_w + body_quat = entity.data.body_link_quat_w + top_pos = body_pos[:, top_body_id] + top_quat = body_quat[:, top_body_id] + ball_pos = body_pos[:, ball_body_id] + relative = np_quat_apply_inverse(top_quat, ball_pos - top_pos) + return relative, top_quat, ball_pos + + +@dataclass(kw_only=True) +class StewartTiltActionCfg(ActionTermCfg): + """Two-axis tilt action converted to six Stewart actuator targets.""" + + actuator_names: tuple[str, ...] | list[str] + top_body_name: str + ball_body_name: str + leg_body_names: tuple[str, ...] | list[str] + top_connect_body_names: tuple[str, ...] | list[str] + raw_action_clip: tuple[float, float] | list[float] + target_rotation_limit_deg: float + action_smooth: float + center_control_radius: float + center_control_min_gain: float + + def build(self, env: ManagerBasedRlEnv) -> StewartTiltAction: + return StewartTiltAction(self, env) + + +class StewartTiltAction(ActionTerm): + """Vectorized tilt IK using only the public entity state/control facade.""" + + cfg: StewartTiltActionCfg + _entity: Entity + + def __init__(self, cfg: StewartTiltActionCfg, env: ManagerBasedRlEnv): + super().__init__(cfg, env) + term = type(self).__name__ + if cfg.clip is not None: + raise NotImplementedError( + f"{term} does not support the actuator-name clip field; use raw_action_clip" + ) + actuator_names = _names(term, "actuator_names", cfg.actuator_names, count=_LEG_COUNT) + actuator_ids, resolved = self._entity.find_actuators(actuator_names, preserve_order=True) + if tuple(resolved) != actuator_names: + raise ValueError(f"{term} actuator selectors resolved out of order: {resolved}") + self._actuator_ids = np.asarray(actuator_ids, dtype=np.intp) + self._actuator_ids.setflags(write=False) + + top_name = _name(term, "top_body_name", cfg.top_body_name) + ball_name = _name(term, "ball_body_name", cfg.ball_body_name) + leg_names = _names(term, "leg_body_names", cfg.leg_body_names, count=_LEG_COUNT) + connect_names = _names( + term, + "top_connect_body_names", + cfg.top_connect_body_names, + count=_LEG_COUNT, + ) + self._top_body_id = _body_id(self._entity, top_name, term=term) + self._ball_body_id = _body_id(self._entity, ball_name, term=term) + self._leg_body_ids = _body_ids(self._entity, leg_names, term=term) + self._top_connect_body_ids = _body_ids(self._entity, connect_names, term=term) + + self._raw_clip = _pair(term, "raw_action_clip", cfg.raw_action_clip) + self._tilt_limit_deg = _real( + term, + "target_rotation_limit_deg", + cfg.target_rotation_limit_deg, + minimum=0.0, + strict_minimum=True, + ) + self._action_smooth = _real( + term, "action_smooth", cfg.action_smooth, minimum=0.0, maximum=1.0 + ) + self._center_radius = _real( + term, "center_control_radius", cfg.center_control_radius, minimum=0.0 + ) + self._center_min_gain = _real( + term, + "center_control_min_gain", + cfg.center_control_min_gain, + minimum=0.0, + maximum=1.0, + ) + ranges = np.asarray(self._entity.data.actuator_ctrl_range, dtype=get_global_dtype()) + self._ctrl_lower = ranges[self._actuator_ids, 0] + self._ctrl_upper = ranges[self._actuator_ids, 1] + dtype = get_global_dtype() + self._raw_action = np.zeros((env.num_envs, _ACTION_DIM), dtype=dtype) + self._clipped_action = np.zeros_like(self._raw_action) + self._executed_action = np.zeros_like(self._raw_action) + self._previous_executed_action = np.zeros_like(self._raw_action) + self._effective_action = np.zeros_like(self._raw_action) + self._target_tilt_deg = np.zeros_like(self._raw_action) + self._target_tilt_rad = np.zeros_like(self._raw_action) + self._control = np.zeros((env.num_envs, _LEG_COUNT), dtype=dtype) -@registry.envcfg("StewartBalance") -@dataclass -class StewartBalanceCfg(EnvCfg): - scene: SceneCfg = field( # pyright: ignore[reportIncompatibleVariableOverride] - default_factory=lambda: SceneCfg( - model_file=str(ASSETS_ROOT_PATH / "robots" / "stewart" / "scene.xml") - ) - ) - # The XML model is stiff; do not raise sim_dt above ~0.005. - sim_dt: float = 0.004 - ctrl_dt: float = 0.02 - max_episode_seconds: float = 24.0 # pyright: ignore[reportIncompatibleVariableOverride] - render_spacing: float = 4.5 - - # Body the backend treats as the kinematic base for its base-pose accessors. - # The task reads explicit body ids instead, so any real body works; the moving - # plate is the natural choice. - base_name: str = "top" - - # Geometry (platform centered at the world origin, plate center at z=1). - platform_radius: float = 0.8 - # The episode ends (ball "fallen") once it strays this far from the plate - # center. Kept inside the physical rim (platform_radius) so the ball never - # reaches the edge-contact regime that destabilizes the stiff closed-loop solver. - fall_radius: float = 0.5 - top_center_z: float = 1.0 - top_surface_offset: float = 0.1 - ball_radius: float = 0.10 - init_ball_radius_ratio: float = 0.18 - - # Control. - target_rotation_limit_deg: float = 6.0 - action_smooth: float = 0.60 - center_control_radius: float = 0.25 - center_control_min_gain: float = 0.15 - vel_smooth: float = 0.25 - - # Success / stillness window. - still_xy: float = 0.12 - still_vel: float = 0.07 - still_xy_hysteresis: float = 1.15 - still_vel_hysteresis: float = 1.20 - zero_vel_thresh: float = 0.07 - still_steps_needed: int = 5 - - reward_config: StewartRewardConfig = field(default_factory=StewartRewardConfig) - - def validate(self) -> None: - super().validate() - if not 0.0 <= self.init_ball_radius_ratio <= 1.0: - raise ValueError("init_ball_radius_ratio must be in [0, 1]") - if not 0.0 <= self.action_smooth <= 1.0: - raise ValueError("action_smooth must be in [0, 1]") - if not 0.0 <= self.center_control_min_gain <= 1.0: - raise ValueError("center_control_min_gain must be in [0, 1]") - - -def _ball_home_z(cfg: StewartBalanceCfg) -> float: - return cfg.top_center_z + cfg.top_surface_offset + cfg.ball_radius - - -def _roll_pitch_from_quat(quat: np.ndarray) -> tuple[np.ndarray, np.ndarray]: - """Extract roll/pitch (rad) from a wxyz quaternion, cast to float32.""" - roll, pitch = np_roll_pitch_from_quat(quat) - return roll.astype(np.float32), pitch.astype(np.float32) - - -# The mujoco backend constructs, resets, and steps correctly. Its closed-loop -# constraint solver is, however, less forgiving than motrix's under load (stiff -# parallel mechanism + ball contact), so a trained-grade policy is not yet stable -# there — closed-loop stability tuning for mujoco is a follow-up. The motrix -# backend is the validated training path. -@registry.env("StewartBalance", sim_backend="mujoco") -@registry.env("StewartBalance", sim_backend="motrix") -@registry.env("StewartBalance", sim_backend="drake") -class StewartBalanceEnv(NpEnv): - _cfg: StewartBalanceCfg - - def __init__( - self, - cfg: StewartBalanceCfg, - num_envs: int = 1, - backend_type: str = "motrix", - dr_provider: DomainRandomizationProvider | None = None, - ) -> None: - # add_body_sensors=True injects body-pose tracking sensors the MuJoCo - # backend needs for get_body_pos_w/quat_w on arbitrary bodies (the IK + obs - # read top/ball/leg poses). The motrix backend reads poses natively and - # ignores the flag. - backend = create_backend( - backend_type, - cfg.scene, - num_envs, - cfg.sim_dt, - base_name=cfg.base_name, - add_body_sensors=True, - **env_backend_kwargs(cfg), - ) - super().__init__(cfg, backend, num_envs) - - self._np_dtype = get_global_dtype() - self._action_space = gym.spaces.Box(-1.0, 1.0, (_ACTION_DIM,), dtype=np.float32) - - if self._backend.num_actuators != 6: - raise ValueError(f"Stewart model needs 6 actuators, got {self._backend.num_actuators}") - ctrl_range = np.asarray(self._backend.get_actuator_ctrl_range(), dtype=np.float32) - self._ctrl_lo = ctrl_range[:, 0] - self._ctrl_hi = ctrl_range[:, 1] - - self._top_body_ids = self._backend.get_body_ids(["top"]) - self._ball_body_ids = self._backend.get_body_ids(["ball"]) - self._leg_body_ids = self._backend.get_body_ids(_LEG_BODY_NAMES) - self._top_connect_ids = self._backend.get_body_ids(_TOP_CONNECT_NAMES) - # The ball free joint is the first jointed body in the scene, so its - # position occupies qpos[0:3] (validated against the default qpos in reset). - self._ball_pos_qpos_idx = np.array([0, 1, 2], dtype=np.int64) - - # IK calibration (top home center, connect offsets, neutral leg lengths) is - # resolved lazily on first use, once reset has placed the home state. self._ik_ready = False - self._top_pos0 = np.zeros(3, dtype=np.float32) - self._connect_offsets = np.zeros((6, 3), dtype=np.float32) - self._leg0 = np.zeros(6, dtype=np.float32) + self._top_home_pos = np.zeros(3, dtype=dtype) + self._connect_offsets = np.zeros((_LEG_COUNT, 3), dtype=dtype) + self._neutral_leg_lengths = np.zeros(_LEG_COUNT, dtype=dtype) - self._init_domain_randomization( - dr_provider if dr_provider is not None else StewartBalanceDRProvider() - ) + @property + def action_dim(self) -> int: + return _ACTION_DIM + + @property + def raw_action(self) -> np.ndarray: + return self._raw_action @property - def action_space(self) -> gym.spaces.Box: - return self._action_space + def executed_action(self) -> np.ndarray: + return self._executed_action @property - def obs_groups_spec(self) -> dict[str, int]: - return {"obs": _OBS_DIM} + def target_tilt_deg(self) -> np.ndarray: + return self._target_tilt_deg + + @property + def neutral_leg_lengths(self) -> np.ndarray: + self._ensure_ik_calibration() + return self._neutral_leg_lengths def _ensure_ik_calibration(self) -> None: if self._ik_ready: return - top_pos = np.asarray( - self._backend.get_body_pos_w(self._top_body_ids)[:, 0, :], dtype=np.float32 - ) - connects = np.asarray(self._backend.get_body_pos_w(self._top_connect_ids), dtype=np.float32) - legs = np.asarray(self._backend.get_body_pos_w(self._leg_body_ids), dtype=np.float32) - # Use env 0 as the reference home configuration (all envs share the model). - self._top_pos0 = top_pos[0] - self._connect_offsets = (connects[0] - self._top_pos0).astype(np.float32) # (6,3) - self._leg0 = np.linalg.norm(connects[0] - legs[0], axis=-1).astype(np.float32) # (6,) + positions = np.asarray(self._entity.data.body_link_pos_w, dtype=get_global_dtype()) + top = positions[:, self._top_body_id] + connects = positions[:, self._top_connect_body_ids] + legs = positions[:, self._leg_body_ids] + self._top_home_pos[:] = top[0] + self._connect_offsets[:] = connects[0] - self._top_home_pos + self._neutral_leg_lengths[:] = np.linalg.norm(connects[0] - legs[0], axis=-1) self._ik_ready = True - def _leg_ctrl_for_tilt(self, target_tilt_rad: np.ndarray) -> np.ndarray: - """Inverse kinematics: tilt command -> six prismatic leg targets. - - ``target_tilt_rad`` has shape (N, 2) = (roll, pitch). The plate stays at its - home center; only its orientation tracks the commanded tilt. - """ + def leg_control_for_tilt(self, target_tilt_rad: np.ndarray) -> np.ndarray: + """Return six actuator controls for ``(roll, pitch)`` radians.""" + expected = (self.num_envs, _ACTION_DIM) + if not isinstance(target_tilt_rad, np.ndarray) or target_tilt_rad.shape != expected: + shape = getattr(target_tilt_rad, "shape", None) + raise ValueError(f"{type(self).__name__} tilt must have shape {expected}, got {shape}") + if not np.isfinite(target_tilt_rad).all(): + raise ValueError(f"{type(self).__name__} tilt contains NaN or Inf") self._ensure_ik_calibration() - num = target_tilt_rad.shape[0] - zeros = np.zeros((num,), dtype=np.float32) - target_quat = np_quat_from_euler_xyz( - target_tilt_rad[:, 0], target_tilt_rad[:, 1], zeros - ).reshape(num, 4) - # Rotate each connect offset by the per-env target quat: (N,1,4) x (1,6,3) -> (N,6,3). + zeros = np.zeros(self.num_envs, dtype=target_tilt_rad.dtype) + target_quat = np_quat_from_euler_xyz(target_tilt_rad[:, 0], target_tilt_rad[:, 1], zeros) rotated = np_quat_apply_batched(target_quat[:, None, :], self._connect_offsets[None, :, :]) - expected = self._top_pos0[None, None, :] + rotated # (N,6,3) connect targets - bottoms = np.asarray( - self._backend.get_body_pos_w(self._leg_body_ids), dtype=np.float32 - ) # (N,6,3) - leg_len = np.linalg.norm(expected - bottoms, axis=-1) - self._leg0[None, :] - return np.clip(leg_len, self._ctrl_lo, self._ctrl_hi).astype(np.float32) - - def apply_action(self, actions: np.ndarray, state: NpEnvState) -> np.ndarray: - cfg = self._cfg - info = state.info - raw = np.clip(np.asarray(actions, dtype=np.float32), -1.0, 1.0).reshape( - self._num_envs, _ACTION_DIM - ) - - # Exponential action smoothing. - prev = info["prev_action_exec"] - alpha = float(cfg.action_smooth) - action_exec = (alpha * raw + (1.0 - alpha) * prev).astype(np.float32) - info["prev_action_exec"] = action_exec - info["action_exec"] = action_exec - - # Soften authority while the ball is already near the center. - rel_xy = info["last_rel_xy"] - if cfg.center_control_radius > 0.0 and cfg.center_control_min_gain < 1.0: - ratio = np.clip(rel_xy / max(cfg.center_control_radius, 1e-6), 0.0, 1.0) - gain = cfg.center_control_min_gain + (1.0 - cfg.center_control_min_gain) * ratio + expected_connects = self._top_home_pos[None, None, :] + rotated + leg_positions = np.asarray( + self._entity.data.body_link_pos_w[:, self._leg_body_ids], + dtype=get_global_dtype(), + ) + controls = ( + np.linalg.norm(expected_connects - leg_positions, axis=-1) + - self._neutral_leg_lengths[None, :] + ) + return np.asarray( + np.clip(controls, self._ctrl_lower, self._ctrl_upper), + dtype=get_global_dtype(), + ) + + def process_actions(self, actions: np.ndarray) -> None: + expected = self._raw_action.shape + if not isinstance(actions, np.ndarray): + raise TypeError( + f"{type(self).__name__} expected np.ndarray, got {type(actions).__name__}" + ) + if actions.shape != expected: + raise ValueError( + f"{type(self).__name__} expected action shape {expected}, got {actions.shape}" + ) + if not np.isfinite(actions).all(): + raise ValueError(f"{type(self).__name__} received NaN or Inf actions") + self._raw_action[:] = actions + np.clip(actions, self._raw_clip[0], self._raw_clip[1], out=self._clipped_action) + np.multiply(self._clipped_action, self._action_smooth, out=self._executed_action) + self._executed_action += (1.0 - self._action_smooth) * self._previous_executed_action + self._previous_executed_action[:] = self._executed_action + + relative, _, _ = _relative_ball_state( + self._entity, + ball_body_id=self._ball_body_id, + top_body_id=self._top_body_id, + ) + relative_xy = np.linalg.norm(relative[:, :2], axis=-1) + if self._center_radius > 0.0 and self._center_min_gain < 1.0: + ratio = np.clip(relative_xy / self._center_radius, 0.0, 1.0) + gain = self._center_min_gain + (1.0 - self._center_min_gain) * ratio + np.multiply(self._executed_action, gain[:, None], out=self._effective_action) else: - gain = np.ones((self._num_envs,), dtype=np.float32) - effective = action_exec * gain[:, None] - - target_tilt_deg = effective * cfg.target_rotation_limit_deg - info["target_tilt_cmd"] = target_tilt_deg.astype(np.float32) - return self._leg_ctrl_for_tilt(np.deg2rad(target_tilt_deg).astype(np.float32)) - - def _read_ball_rel(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]: - top_pos = np.asarray( - self._backend.get_body_pos_w(self._top_body_ids)[:, 0, :], dtype=np.float32 - ) - top_quat = np.asarray( - self._backend.get_body_quat_w(self._top_body_ids)[:, 0, :], dtype=np.float32 - ) - ball_pos = np.asarray( - self._backend.get_body_pos_w(self._ball_body_ids)[:, 0, :], dtype=np.float32 - ) - rel = np_quat_apply_inverse(top_quat, (ball_pos - top_pos)).astype(np.float32) - return rel, top_quat, ball_pos - - def update_state(self, state: NpEnvState) -> NpEnvState: - cfg = self._cfg - info = state.info - dt = float(cfg.ctrl_dt) - - rel, top_quat, ball_pos = self._read_ball_rel() - - rel_vel = (rel - info["prev_rel"]) / dt - filt_rel_vel = ( - cfg.vel_smooth * rel_vel + (1.0 - cfg.vel_smooth) * info["filtered_rel_vel"] - ).astype(np.float32) - info["prev_rel"] = rel - info["filtered_rel_vel"] = filt_rel_vel - - quat_delta = np_quat_mul_batched(top_quat, np_quat_conjugate_batched(info["prev_top_quat"])) - top_ang_vel = (np_quat_to_axis_angle(quat_delta) / dt).astype(np.float32) - filt_ang_vel = ( - cfg.vel_smooth * top_ang_vel + (1.0 - cfg.vel_smooth) * info["filtered_top_ang_vel"] - ).astype(np.float32) - info["prev_top_quat"] = top_quat - info["filtered_top_ang_vel"] = filt_ang_vel - ang_vel_local = np_quat_apply_inverse(top_quat, filt_ang_vel).astype(np.float32) - - roll, pitch = _roll_pitch_from_quat(top_quat) - rel_xy = np.linalg.norm(rel[:, :2], axis=-1).astype(np.float32) - vel_xy = np.linalg.norm(filt_rel_vel[:, :2], axis=-1).astype(np.float32) - info["last_rel_xy"] = rel_xy - - limit = max(cfg.target_rotation_limit_deg, 1e-6) - obs = np.concatenate( - [ - rel, - filt_rel_vel, - np.stack([np.rad2deg(roll) / limit, np.rad2deg(pitch) / limit], axis=-1).astype( - np.float32 - ), - ang_vel_local, - info["target_tilt_cmd"] / limit, - info["action_exec"], - ], - axis=-1, - ).astype(self._np_dtype) - - reward, terminated = self._compute_reward(cfg, info, rel_xy, vel_xy, ball_pos) - return state.replace(obs={"obs": obs}, reward=reward, terminated=terminated) - - def _compute_reward(self, cfg, info, rel_xy, vel_xy, ball_pos): - rc = cfg.reward_config - scales = rc.scales - - fall_z = cfg.top_center_z - np.sin(np.deg2rad(30.0)) * cfg.platform_radius - fallen = (rel_xy > cfg.fall_radius) | (ball_pos[:, 2] < fall_z) - - center_score = np.clip(1.0 - rel_xy / max(cfg.fall_radius, 1e-6), 0.0, 1.0) - term_center = scales["center"] * center_score - - # Reward shrinking the stop-radius: progress toward the center between - # near-zero-velocity moments (a stable "settled closer than before" event). - prev_zero = info["prev_zero_vel_rel_xy"] - zero_event = vel_xy <= cfg.zero_vel_thresh - improve = np.maximum(prev_zero - rel_xy, 0.0) - improve_norm = np.clip(improve / max(cfg.platform_radius, 1e-6), 0.0, 1.0) - term_progress = np.where( - zero_event & (rel_xy < prev_zero), scales["progress"] * improve_norm, 0.0 - ) - next_zero = prev_zero.copy() - next_zero[zero_event] = rel_xy[zero_event] - info["prev_zero_vel_rel_xy"] = next_zero.astype(np.float32) - - still_steps = self._update_stillness(cfg, info, rel_xy, vel_xy) - success = still_steps >= cfg.still_steps_needed - term_still = np.where(success, scales["still"], 0.0) - - reward = (term_center + term_progress + term_still).astype(self._np_dtype) - reward = np.where(fallen, rc.fall_penalty, reward).astype(self._np_dtype) - terminated = (fallen | success).astype(bool) - return reward, terminated - - def _update_stillness(self, cfg, info, rel_xy, vel_xy) -> np.ndarray: - xy_enter, vel_enter = cfg.still_xy, cfg.still_vel - xy_exit = cfg.still_xy * cfg.still_xy_hysteresis - vel_exit = cfg.still_vel * cfg.still_vel_hysteresis - - active = info["still_window_active"] - steps = info["still_steps"] - keep = active & (rel_xy <= xy_exit) & (vel_xy <= vel_exit) - enter = (~active) & (rel_xy <= xy_enter) & (vel_xy <= vel_enter) - steps = np.where(keep, steps + 1, np.where(enter, 1, 0)).astype(np.int32) - active = keep | enter - info["still_window_active"] = active - info["still_steps"] = steps - return steps - - -class StewartBalanceDRProvider(DomainRandomizationProvider): - """Resets the plate to its level home and drops the ball near the center. - - No physical randomization terms are used; the reset variety comes from the - randomized ball position, which is the balancing challenge. - """ - - def validate(self, env, capabilities: DomainRandomizationCapabilities) -> None: # noqa: D102 - return None - - def build_reset_plan(self, env: StewartBalanceEnv, env_ids: np.ndarray) -> ResetPlan: - cfg: StewartBalanceCfg = env._cfg - n = int(env_ids.shape[0]) - default_qpos = np.asarray(env._backend.get_default_qpos(), dtype=np.float64) - # Ball free joint is first in the scene -> position is qpos[0:3]. - if not np.allclose(default_qpos[0:2], 0.0, atol=1e-3): - raise ValueError("Unexpected qpos layout: ball position is not at qpos[0:3]") - qpos = np.broadcast_to(default_qpos, (n, default_qpos.shape[0])).copy() - - # Ball: uniform within a disk near the plate center, resting on the surface. - radius = ( - cfg.platform_radius - * cfg.init_ball_radius_ratio - * np.sqrt(np.random.uniform(0.0, 1.0, size=n)) - ) - theta = np.random.uniform(0.0, 2.0 * np.pi, size=n) - ball_xyz = np.stack( - [radius * np.cos(theta), radius * np.sin(theta), np.full((n,), _ball_home_z(cfg))], - axis=-1, - ) - qpos[:, env._ball_pos_qpos_idx] = ball_xyz - - init_qvel = np.asarray(env._backend.get_init_qvel(), dtype=np.float64) - qvel = np.broadcast_to(init_qvel, (n, init_qvel.shape[0])).copy() - - zeros3 = np.zeros((n, 3), dtype=np.float32) - zeros2 = np.zeros((n, 2), dtype=np.float32) - identity_quat = np.tile(np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32), (n, 1)) - info_updates: dict = { - "prev_rel": zeros3.copy(), - "filtered_rel_vel": zeros3.copy(), - "prev_top_quat": identity_quat, - "filtered_top_ang_vel": zeros3.copy(), - "target_tilt_cmd": zeros2.copy(), - "action_exec": zeros2.copy(), - "prev_action_exec": zeros2.copy(), - "last_rel_xy": np.zeros((n,), dtype=np.float32), - "prev_zero_vel_rel_xy": np.full((n,), cfg.platform_radius, dtype=np.float32), - "still_steps": np.zeros((n,), dtype=np.int32), - "still_window_active": np.zeros((n,), dtype=bool), + self._effective_action[:] = self._executed_action + np.multiply(self._effective_action, self._tilt_limit_deg, out=self._target_tilt_deg) + np.deg2rad(self._target_tilt_deg, out=self._target_tilt_rad) + self._control[:] = self.leg_control_for_tilt(self._target_tilt_rad) + + def apply_actions(self) -> None: + self._entity.data.write_ctrl(self._control, actuator_ids=self._actuator_ids) + + def reset(self, env_ids: np.ndarray | slice | None = None) -> None: + ids = slice(None) if env_ids is None else env_ids + for value in ( + self._raw_action, + self._clipped_action, + self._executed_action, + self._previous_executed_action, + self._effective_action, + self._target_tilt_deg, + self._target_tilt_rad, + self._control, + ): + value[ids] = 0.0 + + +class StewartObservation(ManagerTermBase): + """Legacy 15-D observation with per-environment filtered finite differences.""" + + _ALLOWED_PARAMS = frozenset( + { + "entity_name", + "action_name", + "ball_body_name", + "top_body_name", + "target_rotation_limit_deg", + "vel_smooth", + } + ) + + def __init__(self, cfg: ManagerTermBaseCfg, env: _StewartEnv): + super().__init__(env) + term = type(self).__name__ + unexpected = set(cfg.params) - self._ALLOWED_PARAMS + if unexpected: + raise TypeError(f"{term} received unsupported parameters: {sorted(unexpected)}") + entity_name = _name(term, "entity_name", cfg.params.get("entity_name")) + action_name = _name(term, "action_name", cfg.params.get("action_name")) + self._entity = cast("Entity", env.scene[entity_name]) + self._ball_body_id = _body_id( + self._entity, + _name(term, "ball_body_name", cfg.params.get("ball_body_name")), + term=term, + ) + self._top_body_id = _body_id( + self._entity, + _name(term, "top_body_name", cfg.params.get("top_body_name")), + term=term, + ) + action = env.action_manager.get_term(action_name) + if not isinstance(action, StewartTiltAction): + raise TypeError( + f"{term} action term {action_name!r} must be StewartTiltAction, " + f"got {type(action).__name__}" + ) + self._action = action + self._tilt_limit_deg = _real( + term, + "target_rotation_limit_deg", + cfg.params.get("target_rotation_limit_deg"), + minimum=0.0, + strict_minimum=True, + ) + self._vel_smooth = _real( + term, + "vel_smooth", + cfg.params.get("vel_smooth"), + minimum=0.0, + maximum=1.0, + ) + self._step_dt = _real(term, "step_dt", env.step_dt, minimum=0.0, strict_minimum=True) + + dtype = get_global_dtype() + self._relative = np.zeros((env.num_envs, 3), dtype=dtype) + self._previous_relative = np.zeros_like(self._relative) + self._filtered_relative_velocity = np.zeros_like(self._relative) + self._top_quat = np.zeros((env.num_envs, 4), dtype=dtype) + self._top_quat[:, 0] = 1.0 + self._previous_top_quat = self._top_quat.copy() + self._filtered_top_angular_velocity = np.zeros_like(self._relative) + self._local_top_angular_velocity = np.zeros_like(self._relative) + self._ball_pos = np.zeros_like(self._relative) + self._relative_xy = np.zeros(env.num_envs, dtype=dtype) + self._velocity_xy = np.zeros(env.num_envs, dtype=dtype) + self._obs = np.zeros((env.num_envs, 15), dtype=dtype) + self._last_counter = self._counter(env) + self.reset(None) + + @staticmethod + def _counter(env: _StewartEnv) -> int: + counter = env.common_step_counter + if isinstance(counter, (bool, np.bool_)) or not isinstance(counter, (int, np.integer)): + raise TypeError("StewartObservation common_step_counter must be an integer") + if counter < 0: + raise ValueError("StewartObservation common_step_counter must be non-negative") + return int(counter) + + @property + def relative_xy(self) -> np.ndarray: + return self._relative_xy + + @property + def velocity_xy(self) -> np.ndarray: + return self._velocity_xy + + @property + def ball_pos(self) -> np.ndarray: + return self._ball_pos + + def _write_observation_rows(self, ids: np.ndarray, *, reset_actions: bool) -> None: + roll, pitch = np_roll_pitch_from_quat(self._top_quat[ids]) + self._obs[ids, 0:3] = self._relative[ids] + self._obs[ids, 3:6] = self._filtered_relative_velocity[ids] + self._obs[ids, 6] = np.rad2deg(roll) / self._tilt_limit_deg + self._obs[ids, 7] = np.rad2deg(pitch) / self._tilt_limit_deg + self._obs[ids, 8:11] = self._local_top_angular_velocity[ids] + if reset_actions: + self._obs[ids, 11:15] = 0.0 + else: + self._obs[ids, 11:13] = self._action.target_tilt_deg[ids] / self._tilt_limit_deg + self._obs[ids, 13:15] = self._action.executed_action[ids] + + def reset(self, env_ids: np.ndarray | slice | None = None) -> None: + ids = _env_ids(self._env, env_ids) + relative, top_quat, ball_pos = _relative_ball_state( + self._entity, + ball_body_id=self._ball_body_id, + top_body_id=self._top_body_id, + ) + self._relative[ids] = relative[ids] + self._previous_relative[ids] = relative[ids] + self._filtered_relative_velocity[ids] = 0.0 + self._top_quat[ids] = top_quat[ids] + self._previous_top_quat[ids] = top_quat[ids] + self._filtered_top_angular_velocity[ids] = 0.0 + self._local_top_angular_velocity[ids] = 0.0 + self._ball_pos[ids] = ball_pos[ids] + self._relative_xy[ids] = np.linalg.norm(relative[ids, :2], axis=-1) + self._velocity_xy[ids] = 0.0 + self._write_observation_rows(ids, reset_actions=True) + + def _advance(self, env: _StewartEnv) -> None: + counter = self._counter(env) + if counter == self._last_counter: + return + if counter != self._last_counter + 1: + raise RuntimeError( + "StewartObservation missed a control-step update: " + f"last={self._last_counter}, current={counter}" + ) + relative, top_quat, ball_pos = _relative_ball_state( + self._entity, + ball_body_id=self._ball_body_id, + top_body_id=self._top_body_id, + ) + relative_velocity = (relative - self._previous_relative) / self._step_dt + self._filtered_relative_velocity[:] = ( + self._vel_smooth * relative_velocity + + (1.0 - self._vel_smooth) * self._filtered_relative_velocity + ) + quaternion_delta = np_quat_mul_batched( + top_quat, np_quat_conjugate_batched(self._previous_top_quat) + ) + top_angular_velocity = np_quat_to_axis_angle(quaternion_delta) / self._step_dt + self._filtered_top_angular_velocity[:] = ( + self._vel_smooth * top_angular_velocity + + (1.0 - self._vel_smooth) * self._filtered_top_angular_velocity + ) + self._local_top_angular_velocity[:] = np_quat_apply_inverse( + top_quat, self._filtered_top_angular_velocity + ) + self._relative[:] = relative + self._previous_relative[:] = relative + self._top_quat[:] = top_quat + self._previous_top_quat[:] = top_quat + self._ball_pos[:] = ball_pos + self._relative_xy[:] = np.linalg.norm(relative[:, :2], axis=-1) + self._velocity_xy[:] = np.linalg.norm(self._filtered_relative_velocity[:, :2], axis=-1) + all_ids = np.arange(env.num_envs, dtype=np.int32) + self._write_observation_rows(all_ids, reset_actions=False) + self._last_counter = counter + + def snapshot(self, env: _StewartEnv) -> np.ndarray: + self._advance(env) + return self._obs + + def __call__(self, env: _StewartEnv, **params: Any) -> np.ndarray: + del params + return self.snapshot(env) + + +class StewartBalanceState(ManagerTermBase): + """Termination-owned progress, stillness, success, and fall state.""" + + _ALLOWED_PARAMS = frozenset( + { + "observation_group", + "observation_term", + "platform_radius", + "fall_radius", + "top_center_z", + "still_xy", + "still_vel", + "still_xy_hysteresis", + "still_vel_hysteresis", + "zero_vel_thresh", + "still_steps_needed", } - return ResetPlan(env_ids=env_ids, qpos=qpos, qvel=qvel, info_updates=info_updates) - - def build_reset_observation( - self, env: StewartBalanceEnv, env_ids: np.ndarray, info_updates: dict - ) -> dict: - rel, top_quat, _ = env._read_ball_rel() - rel = rel[env_ids] - top_quat = top_quat[env_ids] - roll, pitch = _roll_pitch_from_quat(top_quat) - limit = max(env._cfg.target_rotation_limit_deg, 1e-6) - n = int(env_ids.shape[0]) - info_updates["prev_rel"] = rel.astype(np.float32) - info_updates["prev_top_quat"] = top_quat.astype(np.float32) - info_updates["last_rel_xy"] = np.linalg.norm(rel[:, :2], axis=-1).astype(np.float32) - info_updates["prev_zero_vel_rel_xy"] = info_updates["last_rel_xy"].copy() - obs = np.concatenate( - [ - rel, - np.zeros((n, 3), dtype=np.float32), - np.stack([np.rad2deg(roll) / limit, np.rad2deg(pitch) / limit], axis=-1).astype( - np.float32 - ), - np.zeros((n, 3), dtype=np.float32), - np.zeros((n, 2), dtype=np.float32), - np.zeros((n, 2), dtype=np.float32), - ], - axis=-1, - ).astype(get_global_dtype()) - return {"obs": obs} + ) + + def __init__(self, cfg: ManagerTermBaseCfg, env: _StewartEnv): + super().__init__(env) + term = type(self).__name__ + unexpected = set(cfg.params) - self._ALLOWED_PARAMS + if unexpected: + raise TypeError(f"{term} received unsupported parameters: {sorted(unexpected)}") + group_name = _name(term, "observation_group", cfg.params.get("observation_group")) + observation_name = _name(term, "observation_term", cfg.params.get("observation_term")) + observation = env.observation_manager.get_term_cfg(group_name, observation_name).func + if not isinstance(observation, StewartObservation): + raise TypeError( + f"{term} observation {group_name}/{observation_name} must be " + f"StewartObservation, got {type(observation).__name__}" + ) + self._observation = observation + self._platform_radius = _real( + term, + "platform_radius", + cfg.params.get("platform_radius"), + minimum=0.0, + strict_minimum=True, + ) + self._fall_radius = _real( + term, + "fall_radius", + cfg.params.get("fall_radius"), + minimum=0.0, + strict_minimum=True, + ) + self._top_center_z = _real(term, "top_center_z", cfg.params.get("top_center_z")) + self._still_xy = _real(term, "still_xy", cfg.params.get("still_xy"), minimum=0.0) + self._still_vel = _real(term, "still_vel", cfg.params.get("still_vel"), minimum=0.0) + self._still_xy_hysteresis = _real( + term, + "still_xy_hysteresis", + cfg.params.get("still_xy_hysteresis"), + minimum=1.0, + ) + self._still_vel_hysteresis = _real( + term, + "still_vel_hysteresis", + cfg.params.get("still_vel_hysteresis"), + minimum=1.0, + ) + self._zero_vel_thresh = _real( + term, + "zero_vel_thresh", + cfg.params.get("zero_vel_thresh"), + minimum=0.0, + ) + steps = cfg.params.get("still_steps_needed") + if isinstance(steps, (bool, np.bool_)) or not isinstance(steps, (int, np.integer)): + raise TypeError(f"{term} still_steps_needed must be an integer") + if steps <= 0: + raise ValueError(f"{term} still_steps_needed must be positive") + self._still_steps_needed = int(steps) + + dtype = get_global_dtype() + self.fallen = np.zeros(env.num_envs, dtype=np.bool_) + self.success = np.zeros(env.num_envs, dtype=np.bool_) + self.center_score = np.zeros(env.num_envs, dtype=dtype) + self.progress = np.zeros(env.num_envs, dtype=dtype) + self.still_steps = np.zeros(env.num_envs, dtype=np.int32) + self.still_window_active = np.zeros(env.num_envs, dtype=np.bool_) + self._previous_zero_velocity_xy = np.zeros(env.num_envs, dtype=dtype) + self._done = np.zeros(env.num_envs, dtype=np.bool_) + self._last_counter = int(env.common_step_counter) + + @property + def last_counter(self) -> int: + return self._last_counter + + def reset(self, env_ids: np.ndarray | slice | None = None) -> None: + ids = _env_ids(self._env, env_ids) + self.fallen[ids] = False + self.success[ids] = False + self.center_score[ids] = np.clip( + 1.0 - self._observation.relative_xy[ids] / self._fall_radius, + 0.0, + 1.0, + ) + self.progress[ids] = 0.0 + self.still_steps[ids] = 0 + self.still_window_active[ids] = False + self._previous_zero_velocity_xy[ids] = self._observation.relative_xy[ids] + self._done[ids] = False + + def _update( + self, + relative_xy: np.ndarray, + velocity_xy: np.ndarray, + ball_pos: np.ndarray, + ) -> None: + fall_z = self._top_center_z - np.sin(np.deg2rad(30.0)) * self._platform_radius + self.fallen[:] = (relative_xy > self._fall_radius) | (ball_pos[:, 2] < fall_z) + self.center_score[:] = np.clip( + 1.0 - relative_xy / self._fall_radius, + 0.0, + 1.0, + ) + + zero_event = velocity_xy <= self._zero_vel_thresh + improvement = np.maximum(self._previous_zero_velocity_xy - relative_xy, 0.0) + self.progress[:] = np.where( + zero_event & (relative_xy < self._previous_zero_velocity_xy), + np.clip(improvement / self._platform_radius, 0.0, 1.0), + 0.0, + ) + self._previous_zero_velocity_xy[zero_event] = relative_xy[zero_event] + + keep = ( + self.still_window_active + & (relative_xy <= self._still_xy * self._still_xy_hysteresis) + & (velocity_xy <= self._still_vel * self._still_vel_hysteresis) + ) + enter = ( + ~self.still_window_active + & (relative_xy <= self._still_xy) + & (velocity_xy <= self._still_vel) + ) + self.still_steps[:] = np.where( + keep, + self.still_steps + 1, + np.where(enter, 1, 0), + ) + self.still_window_active[:] = keep | enter + self.success[:] = self.still_steps >= self._still_steps_needed + self._done[:] = self.fallen | self.success + + def __call__(self, env: _StewartEnv, **params: Any) -> np.ndarray: + del params + counter = int(env.common_step_counter) + if counter == self._last_counter: + return self._done + if counter != self._last_counter + 1: + raise RuntimeError( + "StewartBalanceState missed a control-step update: " + f"last={self._last_counter}, current={counter}" + ) + self._observation.snapshot(env) + self._update( + self._observation.relative_xy, + self._observation.velocity_xy, + self._observation.ball_pos, + ) + self._last_counter = counter + return self._done + + +def _balance_state(env: _StewartEnv, state_term_name: str) -> StewartBalanceState: + name = _name("Stewart reward", "state_term_name", state_term_name) + termination_manager = cast("TerminationManager", env.termination_manager) + state = termination_manager.get_term_cfg(name).func + if not isinstance(state, StewartBalanceState): + raise TypeError( + f"Stewart reward termination term {name!r} must be StewartBalanceState, " + f"got {type(state).__name__}" + ) + if state.last_counter != int(env.common_step_counter): + raise RuntimeError( + f"Stewart reward state {name!r} was not computed for control step " + f"{env.common_step_counter}" + ) + return state + + +def center_reward(env: _StewartEnv, state_term_name: str) -> np.ndarray: + state = _balance_state(env, state_term_name) + return np.asarray(np.where(state.fallen, 0.0, state.center_score), dtype=get_global_dtype()) + + +def progress_reward(env: _StewartEnv, state_term_name: str) -> np.ndarray: + state = _balance_state(env, state_term_name) + return np.asarray(np.where(state.fallen, 0.0, state.progress), dtype=get_global_dtype()) + + +def still_reward(env: _StewartEnv, state_term_name: str) -> np.ndarray: + state = _balance_state(env, state_term_name) + return np.asarray(state.success & ~state.fallen, dtype=get_global_dtype()) + + +def fall_reward(env: _StewartEnv, state_term_name: str) -> np.ndarray: + state = _balance_state(env, state_term_name) + return np.asarray(state.fallen, dtype=get_global_dtype()) + + +class StewartBallReset(ManagerTermBase): + """Sample the ball uniformly within a disk via root-state entity writes.""" + + _ALLOWED_PARAMS = frozenset( + {"entity_name", "platform_radius", "init_ball_radius_ratio", "ball_home_z"} + ) + + def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv): + super().__init__(env) + term = type(self).__name__ + unexpected = set(cfg.params) - self._ALLOWED_PARAMS + if unexpected: + raise TypeError(f"{term} received unsupported parameters: {sorted(unexpected)}") + entity_name = _name(term, "entity_name", cfg.params.get("entity_name")) + self._entity = cast("Entity", env.scene[entity_name]) + self._platform_radius = _real( + term, + "platform_radius", + cfg.params.get("platform_radius"), + minimum=0.0, + strict_minimum=True, + ) + self._radius_ratio = _real( + term, + "init_ball_radius_ratio", + cfg.params.get("init_ball_radius_ratio"), + minimum=0.0, + maximum=1.0, + ) + self._ball_home_z = _real(term, "ball_home_z", cfg.params.get("ball_home_z")) + # Resolve the complete floating-root capability on the cold path. + default_state = self._entity.data.default_root_state + if default_state.shape != (env.num_envs, 13): + raise ValueError( + f"{term} default root state must have shape ({env.num_envs}, 13), " + f"got {default_state.shape}" + ) + + def __call__( + self, + env: ManagerBasedRlEnv, + env_ids: np.ndarray | None, + **params: Any, + ) -> None: + del params + ids = _env_ids(env, env_ids) + root_state = np.array(self._entity.data.default_root_state[ids], copy=True) + radius = ( + self._platform_radius + * self._radius_ratio + * np.sqrt(env.rng.uniform(0.0, 1.0, size=ids.size)) + ) + theta = env.rng.uniform(0.0, 2.0 * np.pi, size=ids.size) + root_state[:, 0] = radius * np.cos(theta) + root_state[:, 1] = radius * np.sin(theta) + root_state[:, 2] = self._ball_home_z + self._entity.write_root_link_pose_to_sim(root_state[:, :7], env_ids=ids) + self._entity.write_root_link_velocity_to_sim(root_state[:, 7:], env_ids=ids) + + +registry.register_env_config("StewartBalance", ManagerBasedRlEnvCfg) +registry.register_env("StewartBalance", make_manager_based_rl_env, sim_backend="mujoco") +registry.register_env("StewartBalance", make_manager_based_rl_env, sim_backend="motrix") +registry.register_env("StewartBalance", make_manager_based_rl_env, sim_backend="drake") + + +__all__ = [ + "StewartBalanceState", + "StewartBallReset", + "StewartObservation", + "StewartTiltAction", + "StewartTiltActionCfg", + "center_reward", + "fall_reward", + "progress_reward", + "still_reward", +] diff --git a/tests/envs/test_stewart.py b/tests/envs/test_stewart.py index a7ed673a9..c5d2eb4d0 100644 --- a/tests/envs/test_stewart.py +++ b/tests/envs/test_stewart.py @@ -1,116 +1,382 @@ +"""Hydra-owned production contract for the Stewart Manager-Based task.""" + from __future__ import annotations +from collections.abc import Mapping, Sequence +from dataclasses import fields, is_dataclass from pathlib import Path +from typing import Any import numpy as np import pytest from hydra import compose, initialize_config_dir from hydra.core.global_hydra import GlobalHydra -from omegaconf import OmegaConf +from omegaconf import DictConfig, OmegaConf from unilab.base import registry -from unilab.base.registry import ensure_registries +from unilab.base.config_materialization import apply_cfg_overrides +from unilab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg, make_manager_based_rl_env +from unilab.managers import ObservationTermCfg, RewardTermCfg, TerminationTermCfg +from unilab.tasks.manipulation.stewart.balance import ( + StewartBalanceState, + StewartBallReset, + StewartObservation, + StewartTiltAction, + StewartTiltActionCfg, +) +from unilab.training.backend_adapter import BackendAdapter -_CONF_DIR = Path(__file__).resolve().parents[2] / "conf" -_SRC_DIR = Path(__file__).resolve().parents[2] / "src" +ROOT_DIR = Path(__file__).parents[2] +CONF_DIR = ROOT_DIR / "conf" -_OBS_DIM = 15 -_ACTION_DIM = 2 +_BODY_NAMES = ( + "ball", + "top", + "leg00", + "leg10", + "leg01", + "leg11", + "leg02", + "leg12", + "top_connect00", + "top_connect10", + "top_connect01", + "top_connect11", + "top_connect02", + "top_connect12", +) +_ACTUATOR_NAMES = ("a0", "a1", "a2", "a3", "a4", "a5") +_OWNER_CASES = ( + pytest.param("ppo", ("task=stewart_balance/motrix",), "motrix", id="ppo-motrix"), + pytest.param("ppo", ("task=stewart_balance/mujoco",), "mujoco", id="ppo-mujoco"), + pytest.param("ppo", ("task=stewart_balance/drake",), "drake", id="ppo-drake"), + pytest.param( + "offpolicy", + ("algo=sac", "task=sac/stewart_balance/drake"), + "drake", + id="sac-drake", + ), +) -def _make_env(num_envs: int = 2): - pytest.importorskip("motrixsim", reason="motrixsim not installed") - ensure_registries() - return registry.make("StewartBalance", sim_backend="motrix", num_envs=num_envs) +def _compose(config_group: str, overrides: Sequence[str]) -> DictConfig: + GlobalHydra.instance().clear() + with initialize_config_dir(config_dir=str(CONF_DIR / config_group), version_base="1.3"): + return compose("config", overrides=list(overrides)) -def test_stewart_env_uses_backend_contract() -> None: - """The task must go through the backend contract, not raw sim internals.""" - source = (_SRC_DIR / "unilab" / "tasks" / "manipulation" / "stewart" / "balance.py").read_text( - encoding="utf-8" + +def _materialize( + config_group: str, + overrides: Sequence[str], +) -> tuple[DictConfig, ManagerBasedRlEnvCfg, dict[str, Any]]: + hydra_cfg = _compose(config_group, overrides) + env_override = BackendAdapter(hydra_cfg, root_dir=ROOT_DIR).build_task_env_cfg_override() + env_cfg = registry.materialize_env_config("StewartBalance") + assert isinstance(env_cfg, ManagerBasedRlEnvCfg) + apply_cfg_overrides(env_cfg, env_override) + env_cfg.validate() + return hydra_cfg, env_cfg, env_override + + +def _make_env(backend: str, *, num_envs: int = 2) -> ManagerBasedRlEnv: + hydra_cfg, _, env_override = _materialize( + "ppo", + (f"task=stewart_balance/{backend}", f"algo.num_envs={num_envs}"), ) - assert "import motrixsim" not in source - assert "import mujoco" not in source - assert "_backend.model" not in source - - -def test_stewart_registered_backends() -> None: - ensure_registries() - registered = registry.list_registered_envs() - assert "StewartBalance" in registered - # motrix is the original validated backend; mujoco and drake are - # construct/step-capable comparison backends. - assert set(registered["StewartBalance"]["available_backends"]) == { - "drake", - "motrix", - "mujoco", + env = registry.make( + str(hydra_cfg.training.task_name), + sim_backend=backend, + env_cfg_override=env_override, + num_envs=num_envs, + ) + assert isinstance(env, ManagerBasedRlEnv) + return env + + +def _assert_no_omegaconf(value: Any) -> None: + assert not OmegaConf.is_config(value) + if is_dataclass(value) and not isinstance(value, type): + for item in fields(value): + _assert_no_omegaconf(getattr(value, item.name)) + elif isinstance(value, Mapping): + for key, item in value.items(): + _assert_no_omegaconf(key) + _assert_no_omegaconf(item) + elif isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + for item in value: + _assert_no_omegaconf(item) + + +@pytest.mark.parametrize("config_group,overrides,backend", _OWNER_CASES) +def test_stewart_owner_materializes_complete_plain_manager_cfg( + config_group: str, + overrides: tuple[str, ...], + backend: str, +) -> None: + registry.ensure_registries() + hydra_cfg, env_cfg, _ = _materialize(config_group, overrides) + + assert hydra_cfg.training.task_name == "StewartBalance" + assert hydra_cfg.training.sim_backend == backend + assert env_cfg.sim_dt == pytest.approx(0.004) + assert env_cfg.ctrl_dt == pytest.approx(0.02) + assert env_cfg.max_episode_seconds == pytest.approx(24.0) + assert env_cfg.max_episode_steps == 1200 + assert env_cfg.scale_rewards_by_dt is False + assert env_cfg.policy_observation_group == "policy" + assert env_cfg.critic_observation_group is None + + assert env_cfg.scene is not None + assert env_cfg.scene.model_file.endswith("robots/stewart/scene.xml") + assert list(env_cfg.scene.entities) == ["stewart"] + entity = env_cfg.scene.entities["stewart"] + assert entity.root_body_name == "ball" + assert entity.joint_names is None + assert tuple(entity.actuator_names or ()) == _ACTUATOR_NAMES + assert tuple(entity.body_names or ()) == _BODY_NAMES + + assert list(env_cfg.actions) == ["tilt"] + action_cfg = env_cfg.actions["tilt"] + assert isinstance(action_cfg, StewartTiltActionCfg) + assert action_cfg.entity_name == "stewart" + assert tuple(action_cfg.actuator_names) == _ACTUATOR_NAMES + assert action_cfg.target_rotation_limit_deg == pytest.approx(6.0) + assert action_cfg.action_smooth == pytest.approx(0.60) + assert action_cfg.center_control_radius == pytest.approx(0.25) + assert action_cfg.center_control_min_gain == pytest.approx(0.15) + + assert list(env_cfg.observations) == ["policy"] + policy_group = env_cfg.observations["policy"] + assert policy_group is not None + assert list(policy_group.terms) == ["balance"] + observation_cfg = policy_group.terms["balance"] + assert isinstance(observation_cfg, ObservationTermCfg) + assert observation_cfg.func is StewartObservation + assert observation_cfg.params["vel_smooth"] == pytest.approx(0.25) + + assert list(env_cfg.events) == ["reset_scene_to_default", "reset_ball"] + reset_ball = env_cfg.events["reset_ball"] + assert reset_ball is not None + assert reset_ball.func is StewartBallReset + assert reset_ball.params == { + "entity_name": "stewart", + "platform_radius": 0.8, + "init_ball_radius_ratio": 0.18, + "ball_home_z": 1.2, + } + assert list(env_cfg.terminations) == ["balance_state", "time_out"] + state_cfg = env_cfg.terminations["balance_state"] + assert isinstance(state_cfg, TerminationTermCfg) + assert state_cfg.func is StewartBalanceState + assert state_cfg.params["still_steps_needed"] == 5 + time_out = env_cfg.terminations["time_out"] + assert time_out is not None + assert time_out.time_out is True + + assert list(env_cfg.rewards) == ["center", "progress", "still", "fall"] + weights = { + name: term.weight + for name, term in env_cfg.rewards.items() + if isinstance(term, RewardTermCfg) } + assert weights == {"center": 0.7, "progress": 0.6, "still": 3.0, "fall": -6.0} + + for manager_name in ("observations", "events", "rewards", "terminations"): + for manager_entry in getattr(env_cfg, manager_name).values(): + if manager_entry is None: + continue + terms = ( + manager_entry.terms.values() if manager_name == "observations" else (manager_entry,) + ) + for term in terms: + if term is None: + continue + module = term.func.__module__ + assert ".backend." not in module + assert not any(name in module for name in (".mujoco", ".motrix", ".drake")) + + _assert_no_omegaconf(env_cfg) -def test_stewart_motrix_owner_cfg_composes() -> None: - if GlobalHydra().is_initialized(): - GlobalHydra.instance().clear() - with initialize_config_dir(config_dir=str(_CONF_DIR / "ppo"), version_base="1.3"): - cfg = compose("config", overrides=["task=stewart_balance/motrix", "algo.num_envs=2"]) - assert cfg.training.task_name == "StewartBalance" - assert cfg.training.sim_backend == "motrix" - # Reward block maps onto the env's StewartRewardConfig. - reward = OmegaConf.to_container(cfg.reward, resolve=True) - assert set(reward) == {"scales", "fall_penalty"} - assert set(reward["scales"]) == {"center", "progress", "still"} - assert cfg.algo.obs_groups.actor == ["policy"] - - -def test_stewart_mujoco_owner_cfg_composes() -> None: - if GlobalHydra().is_initialized(): - GlobalHydra.instance().clear() - with initialize_config_dir(config_dir=str(_CONF_DIR / "ppo"), version_base="1.3"): - cfg = compose("config", overrides=["task=stewart_balance/mujoco", "algo.num_envs=2"]) - # Inherits the motrix owner config, only switching the backend. - assert cfg.training.task_name == "StewartBalance" - assert cfg.training.sim_backend == "mujoco" - assert cfg.algo.obs_groups.actor == ["policy"] - - -def test_stewart_env_constructs_and_steps() -> None: - pytest.importorskip("motrixsim", reason="motrixsim not installed") - env = _make_env(num_envs=2) - assert env.obs_groups_spec == {"obs": _OBS_DIM} - assert env.action_space.shape == (_ACTION_DIM,) - - state = None - for _ in range(20): - state = env.step(np.zeros((2, _ACTION_DIM), dtype=np.float32)) - assert state is not None - obs = state.obs["obs"] - assert obs.shape == (2, _OBS_DIM) - assert np.isfinite(obs).all() - assert np.isfinite(state.reward).all() - assert state.terminated.dtype == bool - - -def test_stewart_ik_holds_level_platform() -> None: - """At zero action the IK should hold the plate near its home height (z=1).""" - pytest.importorskip("motrixsim", reason="motrixsim not installed") - env = _make_env(num_envs=2) - env.step(np.zeros((2, _ACTION_DIM), dtype=np.float32)) # triggers reset + calibration - # Level-hold control should be ~zero leg displacement and ~1.1 m neutral legs. - np.testing.assert_allclose(env._leg0, 1.1, atol=1e-2) - ctrl = env._leg_ctrl_for_tilt(np.zeros((2, _ACTION_DIM), dtype=np.float32)) - assert np.allclose(ctrl, 0.0, atol=1e-2) - for _ in range(20): - env.step(np.zeros((2, _ACTION_DIM), dtype=np.float32)) - top_z = env._backend.get_body_pos_w(env._top_body_ids)[:, 0, 2] - assert np.all(np.abs(top_z - 1.0) < 0.1) +def test_stewart_registry_is_manager_only_and_legacy_overrides_fail_closed() -> None: + registry.ensure_registries() + assert registry.list_registered_envs()["StewartBalance"] == { + "config_factory": "ManagerBasedRlEnvCfg", + "available_backends": ["mujoco", "motrix", "drake"], + } + + for legacy_override in ( + {"reward_config": {}}, + {"platform_radius": 0.8}, + {"action_smooth": 0.6}, + ): + with pytest.raises(ValueError, match="has no attribute"): + apply_cfg_overrides(ManagerBasedRlEnvCfg(), legacy_override) + + +def test_stewart_terms_do_not_access_physics_implementations() -> None: + source = (ROOT_DIR / "src/unilab/tasks/manipulation/stewart/balance.py").read_text( + encoding="utf-8" + ) + for forbidden in ( + "import mujoco", + "import motrixsim", + "create_backend", + "env_backend_kwargs", + "._backend", + "get_body_pos_w", + "get_body_quat_w", + "get_default_qpos", + "get_init_qvel", + "set_state(", + ): + assert forbidden not in source + + +@pytest.mark.parametrize("backend", ("motrix", "mujoco")) +def test_stewart_real_manager_runtime_preserves_io_reset_and_level_ik(backend: str) -> None: + try: + env = _make_env(backend, num_envs=2) + except ImportError as exc: + pytest.skip(f"{backend} runtime unavailable: {exc}") + + try: + assert env.obs_groups_spec == {"obs": 15} + assert env.action_space.shape == (2,) + assert env.event_manager.active_terms == {"reset": ["reset_scene_to_default", "reset_ball"]} + obs, info = env.reset(seed=7) + assert {name: value.shape for name, value in obs.items()} == {"obs": (2, 15)} + assert isinstance(info, dict) + assert np.isfinite(obs["obs"]).all() + + entity = env.scene["stewart"] + ball_id = entity.find_bodies("ball")[0][0] + top_id = entity.find_bodies("top")[0][0] + ball_pos = entity.data.body_link_pos_w[:, ball_id] + assert np.all(np.linalg.norm(ball_pos[:, :2], axis=-1) <= 0.8 * 0.18 + 1e-6) + np.testing.assert_allclose(ball_pos[:, 2], 1.2, atol=1e-6) + + action = env.action_manager.get_term("tilt") + assert isinstance(action, StewartTiltAction) + np.testing.assert_allclose(action.neutral_leg_lengths, 1.1, atol=1e-4) + level_control = action.leg_control_for_tilt(np.zeros((2, 2), dtype=np.float32)) + np.testing.assert_allclose(level_control, 0.0, atol=1e-4) + + state = env.step(np.zeros((2, 2), dtype=np.float32)) + for _ in range(19): + state = env.step(np.zeros((2, 2), dtype=np.float32)) + assert state.obs["obs"].shape == (2, 15) + assert np.isfinite(state.obs["obs"]).all() + assert np.isfinite(state.reward).all() + assert state.terminated.dtype == np.bool_ + top_z = entity.data.body_link_pos_w[:, top_id, 2] + assert np.all(np.abs(top_z - 1.0) < 0.1) + finally: + env.close() + + +def test_stewart_action_smoothing_and_center_authority_match_legacy_equations() -> None: + try: + env = _make_env("motrix", num_envs=2) + except ImportError as exc: + pytest.skip(f"motrix runtime unavailable: {exc}") + + try: + env.reset(seed=11) + action = env.action_manager.get_term("tilt") + observation = env.observation_manager.get_term_cfg("policy", "balance").func + assert isinstance(action, StewartTiltAction) + assert isinstance(observation, StewartObservation) + + action.process_actions(np.full((2, 2), 2.0, dtype=np.float32)) + np.testing.assert_allclose(action.executed_action, 0.6, atol=1e-6) + ratio = np.clip(observation.relative_xy / 0.25, 0.0, 1.0) + expected_gain = 0.15 + 0.85 * ratio + np.testing.assert_allclose( + action.target_tilt_deg, + np.broadcast_to(0.6 * expected_gain[:, None] * 6.0, (2, 2)), + atol=1e-6, + ) + + action.process_actions(np.ones((2, 2), dtype=np.float32)) + np.testing.assert_allclose(action.executed_action, 0.84, atol=1e-6) + action.reset(np.array([1], dtype=np.int32)) + np.testing.assert_allclose(action.executed_action[1], 0.0) + np.testing.assert_allclose(action.executed_action[0], 0.84) + finally: + env.close() + + +def test_stewart_state_machine_and_fall_reward_are_exact() -> None: + try: + env = _make_env("motrix", num_envs=2) + except ImportError as exc: + pytest.skip(f"motrix runtime unavailable: {exc}") + + try: + env.reset(seed=3) + state_term = env.termination_manager.get_term_cfg("balance_state").func + assert isinstance(state_term, StewartBalanceState) + state_term._previous_zero_velocity_xy[:] = 0.4 + state_term._update( + np.array([0.2, 0.6], dtype=np.float32), + np.array([0.0, 0.2], dtype=np.float32), + np.array([[0.0, 0.0, 1.2], [0.0, 0.0, 1.2]], dtype=np.float32), + ) + np.testing.assert_array_equal(state_term.fallen, [False, True]) + np.testing.assert_allclose(state_term.center_score, [0.6, 0.0], atol=1e-6) + np.testing.assert_allclose(state_term.progress, [0.25, 0.0], atol=1e-6) + + # Prove fall masking independently of the geometric scores themselves. + state_term.center_score[1] = 0.9 + state_term.progress[1] = 0.5 + state_term.success[1] = True + reward = env.reward_manager.compute(dt=env.step_dt) + assert reward[0] == pytest.approx(0.7 * 0.6 + 0.6 * 0.25) + # Positive terms are explicitly masked for fallen environments. + assert reward[1] == pytest.approx(-6.0) + + state_term.fallen[:] = False + state_term.success[:] = False + state_term.still_steps[:] = 0 + state_term.still_window_active[:] = False + state_term._previous_zero_velocity_xy[:] = 0.1 + centered = np.full(2, 0.1, dtype=np.float32) + slow = np.full(2, 0.05, dtype=np.float32) + ball_pos = np.full((2, 3), (0.0, 0.0, 1.2), dtype=np.float32) + for _ in range(5): + state_term._update(centered, slow, ball_pos) + np.testing.assert_array_equal(state_term.still_steps, [5, 5]) + np.testing.assert_array_equal(state_term.success, [True, True]) + finally: + env.close() + + +def test_stewart_drake_materializes_or_fails_at_optional_runtime_boundary() -> None: + _, env_cfg, _ = _materialize("ppo", ("task=stewart_balance/drake",)) + try: + env = make_manager_based_rl_env(env_cfg, num_envs=1, backend_type="drake") + except ImportError as exc: + assert "DrakeUni batch runtime is not installed" in str(exc) + return + try: + obs, _ = env.reset(seed=5) + assert obs["obs"].shape == (1, 15) + finally: + env.close() @pytest.mark.slow def test_stewart_solver_stable_under_random_actions() -> None: - """Regression for the closed-loop solver blow-up (NotPositiveDefinite): - the fall-radius margin + softened actuator kp must keep it stable.""" - pytest.importorskip("motrixsim", reason="motrixsim not installed") - env = _make_env(num_envs=8) - np.random.seed(0) - for _ in range(400): - state = env.step(np.random.uniform(-1.0, 1.0, (8, _ACTION_DIM)).astype(np.float32)) - assert np.isfinite(state.obs["obs"]).all() + try: + env = _make_env("motrix", num_envs=8) + except ImportError as exc: + pytest.skip(f"motrix runtime unavailable: {exc}") + try: + rng = np.random.default_rng(0) + for _ in range(400): + state = env.step(rng.uniform(-1.0, 1.0, (8, 2)).astype(np.float32)) + assert np.isfinite(state.obs["obs"]).all() + finally: + env.close() From 50b24992d879e653b37f8a76d9d5094e1e832202 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Wed, 19 Aug 2026 06:11:39 +0800 Subject: [PATCH 081/169] feat(tasks): migrate Allegro rotation to Manager-Based runtime [skip ci] (#1208) --- conf/appo/task/allegro_inhand/base.yaml | 151 +++++ conf/appo/task/allegro_inhand/drake.yaml | 50 ++ conf/appo/task/allegro_inhand/motrix.yaml | 31 +- conf/appo/task/allegro_inhand/mujoco.yaml | 28 +- conf/ppo/task/allegro_inhand/base.yaml | 152 +++++ conf/ppo/task/allegro_inhand/drake.yaml | 46 ++ conf/ppo/task/allegro_inhand/motrix.yaml | 33 +- conf/ppo/task/allegro_inhand/mujoco.yaml | 28 +- .../ppo/task/allegro_inhand_grasp/motrix.yaml | 32 +- .../ppo/task/allegro_inhand_grasp/mujoco.yaml | 36 +- .../allegro_inhand/manager_terms.py | 577 ++++++++++++++++++ .../manipulation/allegro_inhand/rotation.py | 13 +- tests/algos/test_rsl_rl_runner.py | 10 +- .../envs/test_allegro_domain_randomization.py | 70 --- tests/envs/test_env_configs.py | 154 ++++- tests/scripts/test_train_scripts.py | 43 +- 16 files changed, 1224 insertions(+), 230 deletions(-) create mode 100644 conf/appo/task/allegro_inhand/base.yaml create mode 100644 conf/appo/task/allegro_inhand/drake.yaml create mode 100644 conf/ppo/task/allegro_inhand/base.yaml create mode 100644 conf/ppo/task/allegro_inhand/drake.yaml create mode 100644 src/unilab/tasks/manipulation/allegro_inhand/manager_terms.py delete mode 100644 tests/envs/test_allegro_domain_randomization.py diff --git a/conf/appo/task/allegro_inhand/base.yaml b/conf/appo/task/allegro_inhand/base.yaml new file mode 100644 index 000000000..7ee1f9deb --- /dev/null +++ b/conf/appo/task/allegro_inhand/base.yaml @@ -0,0 +1,151 @@ +# @package _global_ +# Canonical Allegro rotation Manager-Based task declaration. Backend leaves own +# only backend identity and algorithm/runtime tuning. +env: + scene: + model_file: src/unilab/assets/robots/allegro_hand/scene.xml + default_keyframe_name: home + entities: + robot: + root_body_name: ball + joint_names: + - ffj0 + - ffj1 + - ffj2 + - ffj3 + - mfj0 + - mfj1 + - mfj2 + - mfj3 + - rfj0 + - rfj1 + - rfj2 + - rfj3 + - thj0 + - thj1 + - thj2 + - thj3 + body_names: [ball, ff_tip, mf_tip, rf_tip, th_tip] + actuator_names: + - ffa0 + - ffa1 + - ffa2 + - ffa3 + - mfa0 + - mfa1 + - mfa2 + - mfa3 + - rfa0 + - rfa1 + - rfa2 + - rfa3 + - tha0 + - tha1 + - tha2 + - tha3 + sim_dt: 0.005 + ctrl_dt: 0.05 + max_episode_seconds: 20.0 + observations: + policy: + _target_: unilab.managers.ObservationGroupCfg + history_length: 3 + flatten_history_dim: true + terms: + rotation: + _target_: unilab.managers.ObservationTermCfg + func: unilab.tasks.manipulation.allegro_inhand.manager_terms.AllegroRotationObservation + params: + entity_name: robot + action_name: hand + joint_noise: 0.02 + torque_estimate_kp: 1.0 + torque_estimate_kd: 0.1 + actions: + hand: + _target_: unilab.tasks.manipulation.allegro_inhand.manager_terms.AllegroIncrementalPositionActionCfg + entity_name: robot + actuator_names: [".*"] + action_scale: 0.041666666666666664 + raw_action_clip: [-1.0, 1.0] + events: + reset_scene_to_default: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.reset_scene_to_default + mode: reset + reset_hand_ball: + _target_: unilab.managers.EventTermCfg + func: unilab.tasks.manipulation.allegro_inhand.manager_terms.AllegroHandBallReset + mode: reset + params: + entity_name: robot + # null explicitly selects the model home pose. A configured path is + # fail-closed when missing or malformed. + grasp_cache_path: null + joint_noise: 0.0 + ball_velocity_noise: 0.0 + ball_z_offset: 0.0 + pd_gains: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.pd_gains + mode: reset + params: + kp_range: [1.0, 1.0] + kd_range: [0.1, 0.1] + operation: abs + terminations: + dropped: + _target_: unilab.managers.TerminationTermCfg + func: unilab.tasks.manipulation.allegro_inhand.manager_terms.AllegroDropTermination + params: + observation_group: policy + observation_term: rotation + minimum_ball_height: 0.125 + time_out: + _target_: unilab.managers.TerminationTermCfg + func: unilab.envs.mdp.time_out + time_out: true + scale_rewards_by_dt: true + policy_observation_group: policy + critic_observation_group: null + +reward: + rotate: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.manipulation.allegro_inhand.manager_terms.AllegroRotateReward + weight: 1.25 + params: + state_term_name: dropped + rotation_axis: [0.0, 0.0, 1.0] + clip_min: -0.5 + clip_max: 0.5 + obj_linvel: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.manipulation.allegro_inhand.manager_terms.object_linear_velocity_l1 + weight: -0.3 + params: + state_term_name: dropped + pose_diff: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.manipulation.allegro_inhand.manager_terms.hand_pose_deviation_l2 + weight: -0.3 + params: + state_term_name: dropped + torque: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.manipulation.allegro_inhand.manager_terms.estimated_torque_l2 + weight: -0.1 + params: + state_term_name: dropped + work: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.manipulation.allegro_inhand.manager_terms.estimated_work_l2 + weight: -2.0 + params: + state_term_name: dropped + drop: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.manipulation.allegro_inhand.manager_terms.dropped + weight: 0.0 + params: + state_term_name: dropped diff --git a/conf/appo/task/allegro_inhand/drake.yaml b/conf/appo/task/allegro_inhand/drake.yaml new file mode 100644 index 000000000..edbb75562 --- /dev/null +++ b/conf/appo/task/allegro_inhand/drake.yaml @@ -0,0 +1,50 @@ +# @package _global_ +defaults: + - /task/allegro_inhand/base + - _self_ + +training: + task_name: AllegroInhandRotation + sim_backend: drake + play_steps: 200 + render_spacing: 0.5 + cam_distance: 1.5 + cam_lookat: [0.75, 0.75, 0] + cam_elevation: -20.0 + replay_queue_size: 4 +algo: + num_envs: 1024 + steps_per_env: 8 + max_iterations: 3000 + save_interval: 500 + algorithm: + value_loss_coef: 4.0 + entropy_coef: 0.01 + learning_rate: 0.001 + desired_kl: 0.025 + adaptive_kl_factor: 2.0 + adaptive_lr_factor: 1.5 + num_learning_epochs: 5 + num_mini_batches: 4 + clip_param: 0.2 + gamma: 0.99 + lam: 0.95 + max_grad_norm: 1.0 + use_clipped_value_loss: true + schedule: adaptive + actor: + hidden_dims: [512, 256, 128] + activation: elu + obs_normalization: true + distribution_cfg: + class_name: rsl_rl.modules.distribution.GaussianDistribution + init_std: 1.0 + std_type: scalar + critic: + hidden_dims: [512, 256, 128] + activation: elu + obs_normalization: true +env: + drake_backend_mode: batch + events: + pd_gains: null diff --git a/conf/appo/task/allegro_inhand/motrix.yaml b/conf/appo/task/allegro_inhand/motrix.yaml index c027729a8..1a89e3b37 100644 --- a/conf/appo/task/allegro_inhand/motrix.yaml +++ b/conf/appo/task/allegro_inhand/motrix.yaml @@ -1,4 +1,8 @@ # @package _global_ +defaults: + - /task/allegro_inhand/base + - _self_ + training: task_name: AllegroInhandRotation sim_backend: motrix @@ -7,6 +11,9 @@ training: cam_distance: 1.5 cam_lookat: [0.75, 0.75, 0] cam_elevation: -20.0 +env: + events: + pd_gains: null algo: num_envs: 16384 steps_per_env: 8 @@ -39,27 +46,3 @@ algo: hidden_dims: [512, 256, 128] activation: elu obs_normalization: true -reward: - scales: - rotate: 1.25 - obj_linvel: -0.3 - pose_diff: -0.3 - torque: -0.1 - work: -2.0 - drop: 0.0 - angvel_clip_min: -0.5 - angvel_clip_max: 0.5 - reset_z_threshold: 0.125 -env: - gen_grasp: false - max_episode_seconds: 20.0 - grasp_cache_path: caches/allegro_grasp_50k.npy - # Keep only grasp/pose reset variation. All online DR terms stay disabled. - domain_rand: - randomize_base_mass: false - random_com: false - randomize_gravity: false - push_robots: false - joint_noise: 0.0 - ball_vel_noise: 0.0 - ball_z_offset: 0.0 diff --git a/conf/appo/task/allegro_inhand/mujoco.yaml b/conf/appo/task/allegro_inhand/mujoco.yaml index c24a442bc..7d4bcfbd0 100644 --- a/conf/appo/task/allegro_inhand/mujoco.yaml +++ b/conf/appo/task/allegro_inhand/mujoco.yaml @@ -1,4 +1,8 @@ # @package _global_ +defaults: + - /task/allegro_inhand/base + - _self_ + training: task_name: AllegroInhandRotation sim_backend: mujoco @@ -41,27 +45,3 @@ algo: hidden_dims: [512, 256, 128] activation: elu obs_normalization: true -reward: - scales: - rotate: 1.25 - obj_linvel: -0.3 - pose_diff: -0.3 - torque: -0.1 - work: -2.0 - drop: 0.0 - angvel_clip_min: -0.5 - angvel_clip_max: 0.5 - reset_z_threshold: 0.125 -env: - gen_grasp: false - max_episode_seconds: 20.0 - grasp_cache_path: caches/allegro_grasp_50k.npy - # Keep only grasp/pose reset variation. All online DR terms stay disabled. - domain_rand: - randomize_base_mass: false - random_com: false - randomize_gravity: false - push_robots: false - joint_noise: 0.0 - ball_vel_noise: 0.0 - ball_z_offset: 0.0 diff --git a/conf/ppo/task/allegro_inhand/base.yaml b/conf/ppo/task/allegro_inhand/base.yaml new file mode 100644 index 000000000..6c7851171 --- /dev/null +++ b/conf/ppo/task/allegro_inhand/base.yaml @@ -0,0 +1,152 @@ +# @package _global_ +# Canonical Allegro rotation Manager-Based task declaration. Backend leaves own +# only backend identity and algorithm/runtime tuning. +env: + scene: + model_file: src/unilab/assets/robots/allegro_hand/scene.xml + default_keyframe_name: home + entities: + robot: + # The hand is fixed; the free ball is the task's single root-state entity. + root_body_name: ball + joint_names: + - ffj0 + - ffj1 + - ffj2 + - ffj3 + - mfj0 + - mfj1 + - mfj2 + - mfj3 + - rfj0 + - rfj1 + - rfj2 + - rfj3 + - thj0 + - thj1 + - thj2 + - thj3 + body_names: [ball, ff_tip, mf_tip, rf_tip, th_tip] + actuator_names: + - ffa0 + - ffa1 + - ffa2 + - ffa3 + - mfa0 + - mfa1 + - mfa2 + - mfa3 + - rfa0 + - rfa1 + - rfa2 + - rfa3 + - tha0 + - tha1 + - tha2 + - tha3 + sim_dt: 0.005 + ctrl_dt: 0.05 + max_episode_seconds: 20.0 + observations: + policy: + _target_: unilab.managers.ObservationGroupCfg + history_length: 3 + flatten_history_dim: true + terms: + rotation: + _target_: unilab.managers.ObservationTermCfg + func: unilab.tasks.manipulation.allegro_inhand.manager_terms.AllegroRotationObservation + params: + entity_name: robot + action_name: hand + joint_noise: 0.02 + torque_estimate_kp: 1.0 + torque_estimate_kd: 0.1 + actions: + hand: + _target_: unilab.tasks.manipulation.allegro_inhand.manager_terms.AllegroIncrementalPositionActionCfg + entity_name: robot + actuator_names: [".*"] + action_scale: 0.041666666666666664 + raw_action_clip: [-1.0, 1.0] + events: + reset_scene_to_default: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.reset_scene_to_default + mode: reset + reset_hand_ball: + _target_: unilab.managers.EventTermCfg + func: unilab.tasks.manipulation.allegro_inhand.manager_terms.AllegroHandBallReset + mode: reset + params: + entity_name: robot + # null explicitly selects the model home pose. A configured path is + # fail-closed when missing or malformed. + grasp_cache_path: null + joint_noise: 0.0 + ball_velocity_noise: 0.0 + ball_z_offset: 0.0 + pd_gains: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.pd_gains + mode: reset + params: + kp_range: [1.0, 1.0] + kd_range: [0.1, 0.1] + operation: abs + terminations: + dropped: + _target_: unilab.managers.TerminationTermCfg + func: unilab.tasks.manipulation.allegro_inhand.manager_terms.AllegroDropTermination + params: + observation_group: policy + observation_term: rotation + minimum_ball_height: 0.125 + time_out: + _target_: unilab.managers.TerminationTermCfg + func: unilab.envs.mdp.time_out + time_out: true + scale_rewards_by_dt: true + policy_observation_group: policy + critic_observation_group: null + +reward: + rotate: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.manipulation.allegro_inhand.manager_terms.AllegroRotateReward + weight: 1.25 + params: + state_term_name: dropped + rotation_axis: [0.0, 0.0, 1.0] + clip_min: -0.5 + clip_max: 0.5 + obj_linvel: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.manipulation.allegro_inhand.manager_terms.object_linear_velocity_l1 + weight: -0.3 + params: + state_term_name: dropped + pose_diff: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.manipulation.allegro_inhand.manager_terms.hand_pose_deviation_l2 + weight: -0.3 + params: + state_term_name: dropped + torque: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.manipulation.allegro_inhand.manager_terms.estimated_torque_l2 + weight: -0.1 + params: + state_term_name: dropped + work: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.manipulation.allegro_inhand.manager_terms.estimated_work_l2 + weight: -2.0 + params: + state_term_name: dropped + drop: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.manipulation.allegro_inhand.manager_terms.dropped + weight: 0.0 + params: + state_term_name: dropped diff --git a/conf/ppo/task/allegro_inhand/drake.yaml b/conf/ppo/task/allegro_inhand/drake.yaml new file mode 100644 index 000000000..dc7f0e09c --- /dev/null +++ b/conf/ppo/task/allegro_inhand/drake.yaml @@ -0,0 +1,46 @@ +# @package _global_ +defaults: + - /task/allegro_inhand/base + - _self_ + +training: + task_name: AllegroInhandRotation + sim_backend: drake + render_spacing: 0.5 + cam_distance: 1.5 + cam_lookat: [0.75, 0.75, 0] + cam_elevation: -20.0 +algo: + num_envs: 16384 + num_steps_per_env: 8 + max_iterations: 201 + obs_groups: + actor: [policy] + critic: [policy] + actor: + class_name: rsl_rl.models.MLPModel + hidden_dims: [512, 256, 128] + activation: elu + obs_normalization: true + distribution_cfg: + class_name: rsl_rl.modules.distribution.GaussianDistribution + init_std: 1.0 + std_type: scalar + critic: + class_name: rsl_rl.models.MLPModel + hidden_dims: [512, 256, 128] + activation: elu + obs_normalization: true + algorithm: + value_loss_coef: 4.0 + desired_kl: 0.02 +env: + drake_backend_mode: batch + events: + # Drake consumes the fixed MJCF actuator gains and does not expose reset + # gain mutation in the installed production adapter. + pd_gains: null +play_profile: + enabled: true + env: + render_spacing: 2.0 diff --git a/conf/ppo/task/allegro_inhand/motrix.yaml b/conf/ppo/task/allegro_inhand/motrix.yaml index 3608d9743..8df1fe13b 100644 --- a/conf/ppo/task/allegro_inhand/motrix.yaml +++ b/conf/ppo/task/allegro_inhand/motrix.yaml @@ -1,7 +1,16 @@ # @package _global_ +defaults: + - /task/allegro_inhand/base + - _self_ + training: task_name: AllegroInhandRotation sim_backend: motrix +env: + # Legacy Motrix consumed the MJCF-native gains; only MuJoCo used the explicit + # host-side position-gain override. + events: + pd_gains: null algo: num_envs: 16384 num_steps_per_env: 8 @@ -26,30 +35,6 @@ algo: algorithm: value_loss_coef: 4.0 desired_kl: 0.02 -reward: - scales: - rotate: 1.25 - obj_linvel: -0.3 - pose_diff: -0.3 - torque: -0.1 - work: -2.0 - drop: 0.0 - angvel_clip_min: -0.5 - angvel_clip_max: 0.5 - reset_z_threshold: 0.125 -env: - gen_grasp: false - max_episode_seconds: 20.0 - grasp_cache_path: caches/allegro_grasp_50k.npy - # Keep only grasp/pose reset variation. All online DR terms stay disabled. - domain_rand: - randomize_base_mass: false - random_com: false - randomize_gravity: false - push_robots: false - joint_noise: 0.0 - ball_vel_noise: 0.0 - ball_z_offset: 0.0 play_profile: enabled: true env: diff --git a/conf/ppo/task/allegro_inhand/mujoco.yaml b/conf/ppo/task/allegro_inhand/mujoco.yaml index 17d639913..3538defdc 100644 --- a/conf/ppo/task/allegro_inhand/mujoco.yaml +++ b/conf/ppo/task/allegro_inhand/mujoco.yaml @@ -1,4 +1,8 @@ # @package _global_ +defaults: + - /task/allegro_inhand/base + - _self_ + training: task_name: AllegroInhandRotation sim_backend: mujoco @@ -30,30 +34,6 @@ algo: algorithm: value_loss_coef: 4.0 desired_kl: 0.02 -reward: - scales: - rotate: 1.25 - obj_linvel: -0.3 - pose_diff: -0.3 - torque: -0.1 - work: -2.0 - drop: 0.0 - angvel_clip_min: -0.5 - angvel_clip_max: 0.5 - reset_z_threshold: 0.125 -env: - gen_grasp: false - max_episode_seconds: 20.0 - grasp_cache_path: caches/allegro_grasp_50k.npy - # Keep only grasp/pose reset variation. All online DR terms stay disabled. - domain_rand: - randomize_base_mass: false - random_com: false - randomize_gravity: false - push_robots: false - joint_noise: 0.0 - ball_vel_noise: 0.0 - ball_z_offset: 0.0 play_profile: enabled: true env: diff --git a/conf/ppo/task/allegro_inhand_grasp/motrix.yaml b/conf/ppo/task/allegro_inhand_grasp/motrix.yaml index 6b5cdca31..803c49e7e 100644 --- a/conf/ppo/task/allegro_inhand_grasp/motrix.yaml +++ b/conf/ppo/task/allegro_inhand_grasp/motrix.yaml @@ -1,14 +1,35 @@ # @package _global_ -defaults: - - /task/allegro_inhand/motrix - - _self_ +# Temporary last-consumer owner for the legacy grasp-generation env. Remove +# with the dedicated AllegroInhandRotationGrasp Manager-Based migration child. training: task_name: AllegroInhandRotationGrasp sim_backend: motrix no_play: true algo: + num_envs: 16384 + num_steps_per_env: 8 max_iterations: 1000 # infinite rollout + obs_groups: + actor: [policy] + critic: [policy] + actor: + class_name: rsl_rl.models.MLPModel + hidden_dims: [512, 256, 128] + activation: elu + obs_normalization: true + distribution_cfg: + class_name: rsl_rl.modules.distribution.GaussianDistribution + init_std: 1.0 + std_type: scalar + critic: + class_name: rsl_rl.models.MLPModel + hidden_dims: [512, 256, 128] + activation: elu + obs_normalization: true + algorithm: + value_loss_coef: 4.0 + desired_kl: 0.02 reward: scales: rotate: 0.0 @@ -17,6 +38,9 @@ reward: torque: 0.0 work: 0.0 drop: 0.0 + angvel_clip_min: -0.5 + angvel_clip_max: 0.5 + reset_z_threshold: 0.125 env: gen_grasp: true max_episode_seconds: 3.0 @@ -28,7 +52,9 @@ env: domain_rand: randomize_base_mass: false random_com: false + randomize_gravity: false push_robots: false + ball_z_offset: 0.0 ball_vel_noise: 0.0 joint_noise: 0.25 # random sampling of the grasp poses play_profile: diff --git a/conf/ppo/task/allegro_inhand_grasp/mujoco.yaml b/conf/ppo/task/allegro_inhand_grasp/mujoco.yaml index e6e62f2f6..e930d683d 100644 --- a/conf/ppo/task/allegro_inhand_grasp/mujoco.yaml +++ b/conf/ppo/task/allegro_inhand_grasp/mujoco.yaml @@ -1,14 +1,39 @@ # @package _global_ -defaults: - - /task/allegro_inhand/mujoco - - _self_ +# Temporary last-consumer owner for the legacy grasp-generation env. Remove +# with the dedicated AllegroInhandRotationGrasp Manager-Based migration child. training: task_name: AllegroInhandRotationGrasp sim_backend: mujoco no_play: true + render_spacing: 0.5 + cam_distance: 1.5 + cam_lookat: [0.75, 0.75, 0] + cam_elevation: -20.0 algo: + num_envs: 16384 + num_steps_per_env: 8 max_iterations: 1000 # infinite rollout + obs_groups: + actor: [policy] + critic: [policy] + actor: + class_name: rsl_rl.models.MLPModel + hidden_dims: [512, 256, 128] + activation: elu + obs_normalization: true + distribution_cfg: + class_name: rsl_rl.modules.distribution.GaussianDistribution + init_std: 1.0 + std_type: scalar + critic: + class_name: rsl_rl.models.MLPModel + hidden_dims: [512, 256, 128] + activation: elu + obs_normalization: true + algorithm: + value_loss_coef: 4.0 + desired_kl: 0.02 reward: scales: rotate: 0.0 @@ -17,6 +42,9 @@ reward: torque: 0.0 work: 0.0 drop: 0.0 + angvel_clip_min: -0.5 + angvel_clip_max: 0.5 + reset_z_threshold: 0.125 env: gen_grasp: true max_episode_seconds: 3.0 @@ -28,7 +56,9 @@ env: domain_rand: randomize_base_mass: false random_com: false + randomize_gravity: false push_robots: false + ball_z_offset: 0.0 ball_vel_noise: 0.0 joint_noise: 0.25 # random sampling of the grasp poses play_profile: diff --git a/src/unilab/tasks/manipulation/allegro_inhand/manager_terms.py b/src/unilab/tasks/manipulation/allegro_inhand/manager_terms.py new file mode 100644 index 000000000..17e7eb08b --- /dev/null +++ b/src/unilab/tasks/manipulation/allegro_inhand/manager_terms.py @@ -0,0 +1,577 @@ +"""Manager-Based terms for Allegro in-hand ball rotation. + +Hydra owns the production task declaration. These terms use only the public +Entity facade and the community manager lifecycle; they do not inspect backend +objects or physical state layouts. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from numbers import Real +from typing import TYPE_CHECKING, Any, Protocol, cast + +import numpy as np +from etils import epath + +from unilab.assets import ASSETS_ROOT_PATH +from unilab.dtype_config import get_global_dtype +from unilab.managers import ActionTerm, ActionTermCfg, ManagerTermBase, ManagerTermBaseCfg +from unilab.utils.geometry import np_normalize_axis, np_quat_angular_velocity_from_pair + +if TYPE_CHECKING: + from unilab.base.entity import Entity + from unilab.managers._types import ManagerBasedRlEnv + from unilab.managers.action_manager import ActionManager + from unilab.managers.observation_manager import ObservationManager + from unilab.managers.termination_manager import TerminationManager + + class _AllegroEnv(ManagerBasedRlEnv, Protocol): + @property + def common_step_counter(self) -> int: ... + + @property + def action_manager(self) -> ActionManager: ... + + @property + def observation_manager(self) -> ObservationManager: ... + + @property + def termination_manager(self) -> TerminationManager: ... + + +def _real( + term: str, + name: str, + value: Any, + *, + minimum: float | None = None, + strict_minimum: bool = False, +) -> float: + if isinstance(value, (bool, np.bool_)) or not isinstance(value, Real): + raise TypeError(f"{term} {name} must be a real number") + result = float(value) + if not np.isfinite(result): + raise ValueError(f"{term} {name} must be finite") + if minimum is not None and (result <= minimum if strict_minimum else result < minimum): + relation = "greater than" if strict_minimum else "at least" + raise ValueError(f"{term} {name} must be {relation} {minimum}") + return result + + +def _name(term: str, name: str, value: Any) -> str: + if not isinstance(value, str) or not value: + raise ValueError(f"{term} {name} must be a non-empty string") + return value + + +def _pair(term: str, name: str, value: Any) -> tuple[float, float]: + if isinstance(value, (str, bytes)) or not isinstance(value, (tuple, list)): + raise TypeError(f"{term} {name} must be a numeric (min, max) pair") + if len(value) != 2: + raise ValueError(f"{term} {name} must contain two values") + lower = _real(term, f"{name}[0]", value[0]) + upper = _real(term, f"{name}[1]", value[1]) + if lower > upper: + raise ValueError(f"{term} {name} lower bound {lower} exceeds upper bound {upper}") + return lower, upper + + +def _env_ids(env: ManagerBasedRlEnv, env_ids: np.ndarray | slice | None) -> np.ndarray: + if env_ids is None: + return np.arange(env.num_envs, dtype=np.int32) + if isinstance(env_ids, slice): + return np.arange(env.num_envs, dtype=np.int32)[env_ids] + return env_ids + + +def _resolve_grasp_cache(cache_path: str) -> epath.Path: + path = epath.Path(cache_path) + if path.is_absolute() or path.exists(): + return path + return epath.Path(ASSETS_ROOT_PATH / cache_path) + + +@dataclass(kw_only=True) +class AllegroIncrementalPositionActionCfg(ActionTermCfg): + """Incremental position targets used by the original Allegro policy.""" + + actuator_names: tuple[str, ...] | list[str] + action_scale: float + raw_action_clip: tuple[float, float] | list[float] + + def build(self, env: ManagerBasedRlEnv) -> AllegroIncrementalPositionAction: + return AllegroIncrementalPositionAction(self, env) + + +class AllegroIncrementalPositionAction(ActionTerm): + """Integrate clipped policy deltas into bounded hand-joint targets.""" + + cfg: AllegroIncrementalPositionActionCfg + _entity: Entity + _raw_action: np.ndarray + _clipped_action: np.ndarray + _target: np.ndarray + + def __init__(self, cfg: AllegroIncrementalPositionActionCfg, env: ManagerBasedRlEnv): + super().__init__(cfg, env) + term = type(self).__name__ + if cfg.clip is not None: + raise NotImplementedError( + f"{term} does not support actuator-name clip; use raw_action_clip" + ) + if isinstance(cfg.actuator_names, (str, bytes)) or not isinstance( + cfg.actuator_names, (tuple, list) + ): + raise TypeError(f"{term} actuator_names must be a sequence of patterns") + self._joint_ids, target_names = self._entity.find_joints_by_actuator_names( + cfg.actuator_names + ) + actuator_ids, actuator_names = self._entity.find_actuators( + cfg.actuator_names, preserve_order=True + ) + if len(self._joint_ids) != len(actuator_ids) or target_names != list( + self._entity.joint_names[index] for index in self._joint_ids + ): + raise ValueError(f"{term} actuator-to-joint mapping is incomplete") + if len(self._joint_ids) != 16: + raise ValueError(f"{term} requires 16 hand actuators, got {len(self._joint_ids)}") + if len(set(actuator_names)) != len(actuator_names): + raise ValueError(f"{term} actuator selector resolved duplicate names") + + self._joint_ids_array = np.asarray(self._joint_ids, dtype=np.intp) + self._joint_ids_array.setflags(write=False) + local_actuator_ids = np.asarray(actuator_ids, dtype=np.intp) + ranges = np.asarray(self._entity.data.actuator_ctrl_range, dtype=get_global_dtype()) + self._ctrl_lower = np.array(ranges[local_actuator_ids, 0], copy=True) + self._ctrl_upper = np.array(ranges[local_actuator_ids, 1], copy=True) + if np.any(self._ctrl_lower >= self._ctrl_upper): + raise ValueError(f"{term} actuator control ranges must have lower < upper") + + self._scale = _real(term, "action_scale", cfg.action_scale, minimum=0.0) + self._raw_clip = _pair(term, "raw_action_clip", cfg.raw_action_clip) + dtype = get_global_dtype() + self._raw_action = np.zeros((env.num_envs, len(self._joint_ids)), dtype=dtype) + self._clipped_action = np.zeros_like(self._raw_action) + self._target = np.asarray( + self._entity.data.default_joint_pos[:, self._joint_ids_array], dtype=dtype + ).copy() + + @property + def action_dim(self) -> int: + return int(self._raw_action.shape[1]) + + @property + def raw_action(self) -> np.ndarray: + return self._raw_action + + @property + def target(self) -> np.ndarray: + return self._target + + @property + def ctrl_lower(self) -> np.ndarray: + return self._ctrl_lower + + @property + def ctrl_upper(self) -> np.ndarray: + return self._ctrl_upper + + @property + def joint_ids(self) -> np.ndarray: + return self._joint_ids_array + + def process_actions(self, actions: np.ndarray) -> None: + if not isinstance(actions, np.ndarray): + raise TypeError(f"expected np.ndarray actions, got {type(actions).__name__}") + if actions.shape != self._raw_action.shape: + raise ValueError(f"expected action shape {self._raw_action.shape}, got {actions.shape}") + if not np.isfinite(actions).all(): + raise ValueError("received NaN or Inf actions") + self._raw_action[:] = actions + np.clip(actions, self._raw_clip[0], self._raw_clip[1], out=self._clipped_action) + self._target += self._scale * self._clipped_action + np.clip(self._target, self._ctrl_lower, self._ctrl_upper, out=self._target) + + def apply_actions(self) -> None: + self._entity.set_joint_position_target(self._target, joint_ids=self._joint_ids_array) + + def reset(self, env_ids: np.ndarray | slice | None = None) -> None: + ids = _env_ids(self._env, env_ids) + self._raw_action[ids] = 0.0 + self._clipped_action[ids] = 0.0 + self._target[ids] = self._entity.data.joint_pos[ids][:, self._joint_ids_array] + + +class AllegroRotationObservation(ManagerTermBase): + """One 35-D frame plus state shared by termination and reward terms.""" + + _ALLOWED_PARAMS = frozenset( + { + "entity_name", + "action_name", + "joint_noise", + "torque_estimate_kp", + "torque_estimate_kd", + } + ) + + def __init__(self, cfg: ManagerTermBaseCfg, env: _AllegroEnv): + super().__init__(env) + term = type(self).__name__ + unexpected = set(cfg.params) - self._ALLOWED_PARAMS + if unexpected: + raise TypeError(f"{term} received unsupported parameters: {sorted(unexpected)}") + entity_name = _name(term, "entity_name", cfg.params.get("entity_name")) + action_name = _name(term, "action_name", cfg.params.get("action_name")) + self._entity = cast("Entity", env.scene[entity_name]) + action = env.action_manager.get_term(action_name) + if not isinstance(action, AllegroIncrementalPositionAction): + raise TypeError( + f"{term} action {action_name!r} must be AllegroIncrementalPositionAction, " + f"got {type(action).__name__}" + ) + self._action = action + self._joint_noise = _real(term, "joint_noise", cfg.params.get("joint_noise"), minimum=0.0) + self._torque_kp = _real( + term, "torque_estimate_kp", cfg.params.get("torque_estimate_kp"), minimum=0.0 + ) + self._torque_kd = _real( + term, "torque_estimate_kd", cfg.params.get("torque_estimate_kd"), minimum=0.0 + ) + + dtype = get_global_dtype() + self.dof_pos = np.asarray( + self._entity.data.joint_pos[:, self._action.joint_ids], dtype=dtype + ).copy() + self.dof_vel = np.zeros_like(self.dof_pos) + self.ball_pos = np.asarray(self._entity.data.root_link_pos_w, dtype=dtype).copy() + self.ball_quat = np.asarray(self._entity.data.root_link_quat_w, dtype=dtype).copy() + self.ball_linvel = np.zeros_like(self.ball_pos) + self.ball_angvel = np.zeros_like(self.ball_pos) + self.torques = np.zeros_like(self.dof_pos) + self.init_pose = self.dof_pos.copy() + self._previous_dof_pos = self.dof_pos.copy() + self._previous_ball_pos = self.ball_pos.copy() + self._previous_ball_quat = self.ball_quat.copy() + self._just_reset = np.ones(env.num_envs, dtype=np.bool_) + self._last_counter = int(env.common_step_counter) + + self._dof_mid = (self._action.ctrl_upper + self._action.ctrl_lower) / 2.0 + self._dof_range = self._action.ctrl_upper - self._action.ctrl_lower + + @property + def last_counter(self) -> int: + return self._last_counter + + def reset(self, env_ids: np.ndarray | slice | None = None) -> None: + ids = _env_ids(self._env, env_ids) + dof_pos = np.asarray( + self._entity.data.joint_pos[:, self._action.joint_ids], dtype=get_global_dtype() + ) + ball_pos = np.asarray(self._entity.data.root_link_pos_w, dtype=get_global_dtype()) + ball_quat = np.asarray(self._entity.data.root_link_quat_w, dtype=get_global_dtype()) + self.dof_pos[ids] = dof_pos[ids] + self.dof_vel[ids] = 0.0 + self.ball_pos[ids] = ball_pos[ids] + self.ball_quat[ids] = ball_quat[ids] + self.ball_linvel[ids] = 0.0 + self.ball_angvel[ids] = 0.0 + self.torques[ids] = 0.0 + self.init_pose[ids] = dof_pos[ids] + self._previous_dof_pos[ids] = dof_pos[ids] + self._previous_ball_pos[ids] = ball_pos[ids] + self._previous_ball_quat[ids] = ball_quat[ids] + self._just_reset[ids] = True + self._last_counter = int(cast("_AllegroEnv", self._env).common_step_counter) + + def snapshot(self, env: _AllegroEnv) -> AllegroRotationObservation: + counter = int(env.common_step_counter) + if counter == self._last_counter: + return self + if counter != self._last_counter + 1: + raise RuntimeError( + f"AllegroRotationObservation missed a control-step update: " + f"last={self._last_counter}, current={counter}" + ) + + dtype = get_global_dtype() + dof_pos = np.asarray(self._entity.data.joint_pos[:, self._action.joint_ids], dtype=dtype) + ball_pos = np.asarray(self._entity.data.root_link_pos_w, dtype=dtype) + ball_quat = np.asarray(self._entity.data.root_link_quat_w, dtype=dtype) + np.subtract(dof_pos, self._previous_dof_pos, out=self.dof_vel) + self.dof_vel /= env.step_dt + np.subtract(ball_pos, self._previous_ball_pos, out=self.ball_linvel) + self.ball_linvel /= env.step_dt + self.ball_angvel[:] = np_quat_angular_velocity_from_pair( + ball_quat, self._previous_ball_quat, env.step_dt + ) + self.dof_pos[:] = dof_pos + self.ball_pos[:] = ball_pos + self.ball_quat[:] = ball_quat + self.torques[:] = self._torque_kp * (self._action.target - self.dof_pos) + self.torques -= self._torque_kd * self.dof_vel + np.clip(self.torques, -0.5, 0.5, out=self.torques) + self._previous_dof_pos[:] = dof_pos + self._previous_ball_pos[:] = ball_pos + self._previous_ball_quat[:] = ball_quat + self._last_counter = counter + return self + + def __call__(self, env: _AllegroEnv, **params: Any) -> np.ndarray: + del params + self.snapshot(env) + dof_pos_norm = 2.0 * (self.dof_pos - self._dof_mid) / (self._dof_range + 1.0e-8) + if self._joint_noise > 0.0: + active = ~self._just_reset + if np.any(active): + dof_pos_norm = dof_pos_norm.copy() + dof_pos_norm[active] += env.rng.uniform( + -self._joint_noise, + self._joint_noise, + size=(int(np.count_nonzero(active)), self._action.action_dim), + ) + self._just_reset[:] = False + return np.concatenate( + (dof_pos_norm, self._action.target, self.ball_pos), + axis=1, + dtype=get_global_dtype(), + ) + + +class AllegroDropTermination(ManagerTermBase): + """Termination-owned drop state, computed before reward terms.""" + + def __init__(self, cfg: ManagerTermBaseCfg, env: _AllegroEnv): + super().__init__(env) + term = type(self).__name__ + allowed = {"observation_group", "observation_term", "minimum_ball_height"} + unexpected = set(cfg.params) - allowed + if unexpected: + raise TypeError(f"{term} received unsupported parameters: {sorted(unexpected)}") + group = _name(term, "observation_group", cfg.params.get("observation_group")) + name = _name(term, "observation_term", cfg.params.get("observation_term")) + observation = env.observation_manager.get_term_cfg(group, name).func + if not isinstance(observation, AllegroRotationObservation): + raise TypeError( + f"{term} observation {group}/{name} must be AllegroRotationObservation, " + f"got {type(observation).__name__}" + ) + self.observation = observation + self._minimum_height = _real( + term, "minimum_ball_height", cfg.params.get("minimum_ball_height") + ) + self.dropped = np.zeros(env.num_envs, dtype=np.bool_) + self._last_counter = int(env.common_step_counter) + + @property + def last_counter(self) -> int: + return self._last_counter + + def reset(self, env_ids: np.ndarray | slice | None = None) -> None: + self.dropped[_env_ids(self._env, env_ids)] = False + self._last_counter = int(cast("_AllegroEnv", self._env).common_step_counter) + + def __call__(self, env: _AllegroEnv, **params: Any) -> np.ndarray: + del params + self.observation.snapshot(env) + self.dropped[:] = self.observation.ball_pos[:, 2] < self._minimum_height + self._last_counter = int(env.common_step_counter) + return self.dropped + + +def _rotation_state( + env: _AllegroEnv, + state_term_name: str, +) -> tuple[AllegroDropTermination, AllegroRotationObservation]: + name = _name("Allegro reward", "state_term_name", state_term_name) + state = env.termination_manager.get_term_cfg(name).func + if not isinstance(state, AllegroDropTermination): + raise TypeError( + f"Allegro reward termination term {name!r} must be AllegroDropTermination, " + f"got {type(state).__name__}" + ) + if state.last_counter != int(env.common_step_counter): + raise RuntimeError( + f"Allegro reward state {name!r} was not computed for control step " + f"{env.common_step_counter}" + ) + return state, state.observation + + +class AllegroRotateReward(ManagerTermBase): + """Reward angular velocity projected onto a cold-path-normalized axis.""" + + def __init__(self, cfg: ManagerTermBaseCfg, env: _AllegroEnv): + super().__init__(env) + term = type(self).__name__ + allowed = {"state_term_name", "rotation_axis", "clip_min", "clip_max"} + unexpected = set(cfg.params) - allowed + if unexpected: + raise TypeError(f"{term} received unsupported parameters: {sorted(unexpected)}") + self._state_term_name = _name(term, "state_term_name", cfg.params.get("state_term_name")) + try: + axis = np.asarray(cfg.params.get("rotation_axis"), dtype=get_global_dtype()) + except (TypeError, ValueError) as exc: + raise TypeError(f"{term} rotation_axis must contain three numeric values") from exc + if axis.shape != (3,) or not np.isfinite(axis).all(): + raise ValueError(f"{term} rotation_axis must be a finite 3-D vector") + self._axis = np.asarray(np_normalize_axis(axis), dtype=get_global_dtype()) + self._clip = _pair( + term, + "clip", + (cfg.params.get("clip_min"), cfg.params.get("clip_max")), + ) + + def __call__(self, env: _AllegroEnv, **params: Any) -> np.ndarray: + del params + _, state = _rotation_state(env, self._state_term_name) + return np.asarray( + np.clip(state.ball_angvel @ self._axis, self._clip[0], self._clip[1]), + dtype=get_global_dtype(), + ) + + +def object_linear_velocity_l1(env: _AllegroEnv, state_term_name: str) -> np.ndarray: + _, state = _rotation_state(env, state_term_name) + return np.asarray(np.sum(np.abs(state.ball_linvel), axis=1), dtype=get_global_dtype()) + + +def hand_pose_deviation_l2(env: _AllegroEnv, state_term_name: str) -> np.ndarray: + _, state = _rotation_state(env, state_term_name) + return np.asarray( + np.sum(np.square(state.dof_pos - state.init_pose), axis=1), dtype=get_global_dtype() + ) + + +def estimated_torque_l2(env: _AllegroEnv, state_term_name: str) -> np.ndarray: + _, state = _rotation_state(env, state_term_name) + return np.asarray(np.sum(np.square(state.torques), axis=1), dtype=get_global_dtype()) + + +def estimated_work_l2(env: _AllegroEnv, state_term_name: str) -> np.ndarray: + _, state = _rotation_state(env, state_term_name) + work = np.sum(state.torques * state.dof_vel, axis=1) + return np.asarray(np.square(work), dtype=get_global_dtype()) + + +def dropped(env: _AllegroEnv, state_term_name: str) -> np.ndarray: + state, _ = _rotation_state(env, state_term_name) + return np.asarray(state.dropped, dtype=get_global_dtype()) + + +class AllegroHandBallReset(ManagerTermBase): + """Reset hand joints and the ball root without exposing qpos layout.""" + + def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv): + super().__init__(env) + term = type(self).__name__ + allowed = { + "entity_name", + "grasp_cache_path", + "joint_noise", + "ball_velocity_noise", + "ball_z_offset", + } + unexpected = set(cfg.params) - allowed + if unexpected: + raise TypeError(f"{term} received unsupported parameters: {sorted(unexpected)}") + entity_name = _name(term, "entity_name", cfg.params.get("entity_name")) + self._entity = cast("Entity", env.scene[entity_name]) + if self._entity.num_joints != 16: + raise ValueError(f"{term} requires 16 hand joints, got {self._entity.num_joints}") + if self._entity.data.default_root_state.shape != (env.num_envs, 13): + raise ValueError(f"{term} requires a 13-D floating ball root for every environment") + self._joint_noise = _real(term, "joint_noise", cfg.params.get("joint_noise"), minimum=0.0) + self._ball_velocity_noise = _real( + term, + "ball_velocity_noise", + cfg.params.get("ball_velocity_noise"), + minimum=0.0, + ) + self._ball_z_offset = _real(term, "ball_z_offset", cfg.params.get("ball_z_offset")) + + cache_value = cfg.params.get("grasp_cache_path") + self._grasp_cache: np.ndarray | None = None + if cache_value is not None: + cache_path = _resolve_grasp_cache(_name(term, "grasp_cache_path", cache_value)) + if not cache_path.exists(): + raise FileNotFoundError( + f"{term} configured grasp cache does not exist: {cache_path}. " + "Set grasp_cache_path to null to explicitly use the model home pose, " + "or generate a cache with `uv run train --algo ppo " + "--task allegro_inhand_grasp --sim mujoco training.no_play=true`." + ) + cache = np.asarray(np.load(cache_path), dtype=np.float64) + if cache.ndim != 2 or cache.shape[1] != 23 or cache.shape[0] == 0: + raise ValueError( + f"{term} grasp cache {cache_path} must have shape (N, 23), got {cache.shape}" + ) + if not np.isfinite(cache).all(): + raise ValueError(f"{term} grasp cache {cache_path} contains NaN or Inf") + self._grasp_cache = cache + + ranges = np.asarray(self._entity.data.actuator_ctrl_range, dtype=np.float64) + if ranges.shape != (16, 2): + raise ValueError(f"{term} actuator control range must have shape (16, 2)") + self._ctrl_lower = ranges[:, 0] + self._ctrl_upper = ranges[:, 1] + + def __call__( + self, + env: ManagerBasedRlEnv, + env_ids: np.ndarray | None, + **params: Any, + ) -> None: + del params + ids = _env_ids(env, env_ids) + count = ids.size + dtype = get_global_dtype() + root_state = np.array(self._entity.data.default_root_state[ids], copy=True) + if self._grasp_cache is not None: + rows = self._grasp_cache[env.rng.integers(0, self._grasp_cache.shape[0], size=count)] + joint_pos = np.array(rows[:, :16], copy=True) + root_state[:, :3] = rows[:, 16:19] + root_state[:, 3:7] = rows[:, 19:23] + else: + joint_pos = np.array(self._entity.data.default_joint_pos[ids], copy=True) + if self._joint_noise > 0.0: + joint_pos += env.rng.uniform( + -self._joint_noise, self._joint_noise, size=joint_pos.shape + ) + root_state[:, 2] += self._ball_z_offset + np.clip(joint_pos, self._ctrl_lower, self._ctrl_upper, out=joint_pos) + joint_vel = np.zeros_like(joint_pos) + root_state[:, 7:] = 0.0 + if self._ball_velocity_noise > 0.0: + root_state[:, 7:10] = env.rng.uniform( + -self._ball_velocity_noise, + self._ball_velocity_noise, + size=(count, 3), + ) + self._entity.write_joint_state_to_sim( + np.asarray(joint_pos, dtype=dtype), + np.asarray(joint_vel, dtype=dtype), + env_ids=ids, + ) + self._entity.write_root_link_pose_to_sim( + np.asarray(root_state[:, :7], dtype=dtype), env_ids=ids + ) + self._entity.write_root_link_velocity_to_sim( + np.asarray(root_state[:, 7:], dtype=dtype), env_ids=ids + ) + + +__all__ = [ + "AllegroDropTermination", + "AllegroHandBallReset", + "AllegroIncrementalPositionAction", + "AllegroIncrementalPositionActionCfg", + "AllegroRotateReward", + "AllegroRotationObservation", + "dropped", + "estimated_torque_l2", + "estimated_work_l2", + "hand_pose_deviation_l2", + "object_linear_velocity_l1", +] diff --git a/src/unilab/tasks/manipulation/allegro_inhand/rotation.py b/src/unilab/tasks/manipulation/allegro_inhand/rotation.py index df4ade827..928cdd500 100644 --- a/src/unilab/tasks/manipulation/allegro_inhand/rotation.py +++ b/src/unilab/tasks/manipulation/allegro_inhand/rotation.py @@ -28,6 +28,7 @@ zero_actions, ) from unilab.dtype_config import get_global_dtype +from unilab.envs import ManagerBasedRlEnvCfg, make_manager_based_rl_env from unilab.utils.geometry import ( np_normalize_axis, np_quat_angular_velocity_from_pair, @@ -142,7 +143,6 @@ class DomainRandConfig: ball_z_offset: float = 0.0 -@registry.envcfg("AllegroInhandRotation") @dataclass class AllegroRotationPPOCfg(AllegroBaseCfg): scene: SceneCfg = field( # pyright: ignore[reportIncompatibleVariableOverride] @@ -256,9 +256,6 @@ def build_reset_observation( # ─────────────────────────── Environment ────────────────────────────── -@registry.env("AllegroInhandRotation", sim_backend="drake") -@registry.env("AllegroInhandRotation", sim_backend="mujoco") -@registry.env("AllegroInhandRotation", sim_backend="motrix") class AllegroRotationPPO(AllegroBaseEnv): _cfg: AllegroRotationPPOCfg # pyright: ignore[reportIncompatibleVariableOverride] _reward_cfg: RewardConfigPPO @@ -515,3 +512,11 @@ def _compute_obs( Domain_Rand = DomainRandConfig AllegroRotationCfg = AllegroRotationPPOCfg AllegroRotationMj = AllegroRotationPPO + +# The legacy class/config above remain only as the implementation base for +# AllegroInhandRotationGrasp. The production rotation identity is Hydra-owned +# and has one generic Manager-Based runtime across all registered backends. +registry.register_env_config("AllegroInhandRotation", ManagerBasedRlEnvCfg) +registry.register_env("AllegroInhandRotation", make_manager_based_rl_env, sim_backend="mujoco") +registry.register_env("AllegroInhandRotation", make_manager_based_rl_env, sim_backend="motrix") +registry.register_env("AllegroInhandRotation", make_manager_based_rl_env, sim_backend="drake") diff --git a/tests/algos/test_rsl_rl_runner.py b/tests/algos/test_rsl_rl_runner.py index 582780ed3..3585d577f 100644 --- a/tests/algos/test_rsl_rl_runner.py +++ b/tests/algos/test_rsl_rl_runner.py @@ -119,7 +119,6 @@ def get_privileged_observations(self): def test_rsl_rl_ppo_one_iteration( env_name: str, default_g1_reward_config, - default_allegro_reward_config, ): """RSL-RL PPO can complete 1 training iteration on a real env.""" from rsl_rl.runners import OnPolicyRunner @@ -138,9 +137,14 @@ def test_rsl_rl_ppo_one_iteration( num_envs = 256 env_cfg_override = {"reward_config": reward_cfg} else: - reward_cfg = default_allegro_reward_config num_envs = 128 - env_cfg_override = {"reward_config": reward_cfg} + root_dir = Path(__file__).parents[2] + GlobalHydra.instance().clear() + with initialize_config_dir(config_dir=str(root_dir / "conf" / "ppo"), version_base="1.3"): + hydra_cfg = compose("config", overrides=["task=allegro_inhand/mujoco"]) + env_cfg_override = BackendAdapter( + hydra_cfg, root_dir=root_dir + ).build_task_env_cfg_override() env = registry.make( env_name, diff --git a/tests/envs/test_allegro_domain_randomization.py b/tests/envs/test_allegro_domain_randomization.py deleted file mode 100644 index 0df12763b..000000000 --- a/tests/envs/test_allegro_domain_randomization.py +++ /dev/null @@ -1,70 +0,0 @@ -from __future__ import annotations - -from typing import Any - -import numpy as np -import pytest - -pytest.importorskip("mujoco", reason="mujoco not installed") - -try: - from mujoco_uni.batch_env import BatchEnvPool as _ # noqa: F401 -except Exception: - pytest.skip( - "mujoco_uni.batch_env not available (platform/libstdc++ issue)", allow_module_level=True - ) - -from unilab.base.registry import ensure_registries - - -def test_allegro_mujoco_reset_applies_base_mass_and_com_domain_randomization( - default_allegro_reward_config: dict[str, Any], -) -> None: - ensure_registries() - - from unilab.base import registry - - env = registry.make( - "AllegroInhandRotation", - num_envs=4, - sim_backend="mujoco", - env_cfg_override={ - "reward_config": default_allegro_reward_config, - "domain_rand": { - "randomize_base_mass": True, - "added_mass_range": [-0.02, 0.02], - "random_com": True, - "com_offset_x": [-0.005, 0.005], - }, - }, - ) - env_obj: Any = env - try: - env_obj.init_state() - backend: Any = env_obj._backend - pool = backend._pool - base_body_id = int(backend._base_body_id) - body_mass = np.stack([pool.get_field(i, "body_mass") for i in range(env_obj.num_envs)]) - body_ipos = np.stack([pool.get_field(i, "body_ipos") for i in range(env_obj.num_envs)]) - body_ipos = body_ipos.reshape(env_obj.num_envs, -1, 3) - - base_mass = float(backend._base_body_mass[base_body_id]) - base_ipos = np.asarray(backend._base_body_ipos[base_body_id]) - - randomized_base_mass = body_mass[:, base_body_id] - randomized_base_ipos = body_ipos[:, base_body_id] - - delta_mass = randomized_base_mass - base_mass - delta_x = randomized_base_ipos[:, 0] - base_ipos[0] - delta_yz = randomized_base_ipos[:, 1:] - base_ipos[1:] - - assert np.unique(np.round(delta_mass, 6)).size > 1 - assert np.unique(np.round(delta_x, 6)).size > 1 - - assert np.all(delta_mass >= -0.02) - assert np.all(delta_mass <= 0.02) - assert np.all(delta_x >= -0.005) - assert np.all(delta_x <= 0.005) - np.testing.assert_allclose(delta_yz, 0.0) - finally: - env_obj.close() diff --git a/tests/envs/test_env_configs.py b/tests/envs/test_env_configs.py index bfd00f775..6bea250aa 100644 --- a/tests/envs/test_env_configs.py +++ b/tests/envs/test_env_configs.py @@ -30,6 +30,23 @@ def _require_mujoco_runtime() -> None: pytest.skip("mujoco_uni.batch_env not available (platform/libstdc++ issue)") +def _allegro_manager_override( + backend: str = "mujoco", *, config_root: str = "ppo" +) -> dict[str, Any]: + from hydra import compose, initialize_config_dir + + from unilab.training.backend_adapter import BackendAdapter + + repo_root = Path(__file__).parents[2] + with initialize_config_dir( + config_dir=str(repo_root / "conf" / config_root), version_base="1.3" + ): + cfg = compose("config", overrides=[f"task=allegro_inhand/{backend}"]) + return BackendAdapter( + cfg, root_dir=repo_root, algo_name=config_root + ).build_task_env_cfg_override() + + # --------------------------------------------------------------------------- # Non-slow: config attribute completeness (no env.step(), no MuJoCo sim) # --------------------------------------------------------------------------- @@ -394,14 +411,64 @@ def test_g1_box_tracking_scene_uses_sphere_hand_and_box_tracking_mesh(): assert name in scene_text -def test_allegro_rotation_obs_groups_spec_dims(): - """Allegro rotation obs_groups_spec should expose single actor obs group.""" - from unilab.tasks.manipulation.allegro_inhand.rotation import AllegroRotationPPO +def test_allegro_rotation_registry_is_manager_only_and_grasp_owns_legacy_bridge(): + from unilab.base import registry + from unilab.base.config_materialization import apply_cfg_overrides + from unilab.envs import ManagerBasedRlEnvCfg - env = cast(Any, object.__new__(AllegroRotationPPO)) - spec = env.obs_groups_spec + ensure_registries() + metadata = registry.list_registered_envs() + assert metadata["AllegroInhandRotation"] == { + "config_factory": "ManagerBasedRlEnvCfg", + "available_backends": ["mujoco", "motrix", "drake"], + } + assert metadata["AllegroInhandRotationGrasp"]["config_factory"] == ("AllegroRotationGraspCfg") + + cfg = registry.materialize_env_config("AllegroInhandRotation") + assert isinstance(cfg, ManagerBasedRlEnvCfg) + apply_cfg_overrides(cfg, _allegro_manager_override()) + assert cfg.policy_observation_group == "policy" + assert cfg.critic_observation_group is None + assert cfg.observations["policy"].history_length == 3 + assert list(cfg.actions) == ["hand"] + assert list(cfg.terminations) == ["dropped", "time_out"] + assert list(cfg.rewards) == [ + "rotate", + "obj_linvel", + "pose_diff", + "torque", + "work", + "drop", + ] + assert not hasattr(cfg, "reward_config") - assert spec == {"obs": 105} + +def test_allegro_manager_configured_missing_grasp_cache_fails_closed(tmp_path: Path): + from unilab.managers import EventTermCfg + from unilab.tasks.manipulation.allegro_inhand.manager_terms import AllegroHandBallReset + + entity = SimpleNamespace( + num_joints=16, + data=SimpleNamespace( + default_root_state=np.zeros((2, 13), dtype=np.float32), + actuator_ctrl_range=np.tile([-1.0, 1.0], (16, 1)), + ), + ) + env = SimpleNamespace(num_envs=2, scene={"robot": entity}) + cfg = EventTermCfg( + func=AllegroHandBallReset, + mode="reset", + params={ + "entity_name": "robot", + "grasp_cache_path": str(tmp_path / "missing.npy"), + "joint_noise": 0.0, + "ball_velocity_noise": 0.0, + "ball_z_offset": 0.0, + }, + ) + + with pytest.raises(FileNotFoundError, match="configured grasp cache does not exist"): + AllegroHandBallReset(cfg, cast(Any, env)) def test_allegro_grasp_obs_groups_spec_dims(): @@ -2092,7 +2159,9 @@ def test_env_reset_and_step( env_cfg_override = {"reward_config": default_g1_walk_flat_reward_config} elif "G1" in env_name: env_cfg_override = {"reward_config": default_g1_reward_config} - elif "Allegro" in env_name: + elif env_name == "AllegroInhandRotation": + env_cfg_override = _allegro_manager_override() + elif env_name == "AllegroInhandRotationGrasp": env_cfg_override = {"reward_config": default_allegro_reward_config} env = cast( @@ -2135,39 +2204,60 @@ def test_env_reset_and_step( env.close() -def _assert_mujoco_position_gains( - env: Any, *, kp: float, kd: float, actuator_ids=slice(None) -) -> None: - model = env._backend.model - pool = env._backend._pool - np.testing.assert_allclose(model.actuator_gainprm[actuator_ids, 0], kp) - np.testing.assert_allclose(model.actuator_biasprm[actuator_ids, 1], -kp) - np.testing.assert_allclose(model.actuator_biasprm[actuator_ids, 2], -kd) - np.testing.assert_allclose(pool.get_field(0, "kp")[actuator_ids], kp) - np.testing.assert_allclose(pool.get_field(0, "kd")[actuator_ids], kd) - - -def test_allegro_env_initializes_kp_kd_into_pool(default_allegro_reward_config): +def test_allegro_manager_runtime_matches_legacy_rotation_transition( + default_allegro_reward_config, +): _require_mujoco_runtime() ensure_registries() from unilab.base import registry + from unilab.envs import ManagerBasedRlEnv + from unilab.tasks.manipulation.allegro_inhand.manager_terms import ( + AllegroIncrementalPositionAction, + ) + from unilab.tasks.manipulation.allegro_inhand.rotation import ( + AllegroRotationPPO, + AllegroRotationPPOCfg, + RewardConfigPPO, + ) - env = cast( - Any, - registry.make( - "AllegroInhandRotation", - num_envs=2, - sim_backend="mujoco", - env_cfg_override={ - "reward_config": default_allegro_reward_config, - "control_config": {"kp": 2.5, "kd": 0.4}, - }, - ), + legacy_cfg = AllegroRotationPPOCfg( + reward_config=RewardConfigPPO(**default_allegro_reward_config) + ) + legacy_cfg.noise_config.level = 0.0 + legacy = AllegroRotationPPO(legacy_cfg, num_envs=2, backend_type="mujoco") + + manager_override = _allegro_manager_override() + manager_override["observations"]["policy"]["terms"]["rotation"]["params"]["joint_noise"] = 0.0 + env = registry.make( + "AllegroInhandRotation", + num_envs=2, + sim_backend="mujoco", + env_cfg_override=manager_override, ) + assert isinstance(env, ManagerBasedRlEnv) try: - _assert_mujoco_position_gains(env, kp=2.5, kd=0.4, actuator_ids=slice(0, 16)) + legacy_initial = legacy.init_state() + manager_initial = env.init_state() + np.testing.assert_allclose(manager_initial.obs["obs"], legacy_initial.obs["obs"]) + + actions = np.full((2, 16), 0.25, dtype=np.float32) + legacy_state = legacy.step(actions) + manager_state = env.step(actions) + np.testing.assert_allclose( + manager_state.obs["obs"], legacy_state.obs["obs"], rtol=1.0e-5, atol=2.0e-6 + ) + np.testing.assert_allclose( + manager_state.reward, legacy_state.reward, rtol=1.0e-5, atol=2.0e-8 + ) + np.testing.assert_array_equal(manager_state.terminated, legacy_state.terminated) + + action = env.action_manager.get_term("hand") + assert isinstance(action, AllegroIncrementalPositionAction) + assert action.action_dim == 16 + assert env.obs_groups_spec == {"obs": 105} finally: env.close() + legacy.close() @pytest.mark.parametrize("sim_backend", ["mujoco", "motrix"]) diff --git a/tests/scripts/test_train_scripts.py b/tests/scripts/test_train_scripts.py index 017f41732..1e2ab80c9 100644 --- a/tests/scripts/test_train_scripts.py +++ b/tests/scripts/test_train_scripts.py @@ -813,8 +813,10 @@ def test_build_ppo_env_cfg_override_allegro_mujoco( mod = _train_rsl_rl(monkeypatch) cfg = _ppo_cfg(["task=allegro_inhand/mujoco"]) ppo_motrix_cfg = _ppo_cfg(["task=allegro_inhand/motrix"]) + ppo_drake_cfg = _ppo_cfg(["task=allegro_inhand/drake"]) appo_cfg = _appo_cfg(["task=allegro_inhand/mujoco"]) appo_motrix_cfg = _appo_cfg(["task=allegro_inhand/motrix"]) + appo_drake_cfg = _appo_cfg(["task=allegro_inhand/drake"]) env_cfg_override = mod.build_ppo_env_cfg_override(cfg) @@ -822,18 +824,20 @@ def test_build_ppo_env_cfg_override_allegro_mujoco( assert cfg.algo.empirical_normalization is False assert cfg.algo.actor.obs_normalization is True assert cfg.algo.critic.obs_normalization is True - assert env_cfg_override["reward_config"]["scales"]["rotate"] == pytest.approx(1.25) - assert env_cfg_override["reward_config"]["reset_z_threshold"] == pytest.approx(0.125) - assert env_cfg_override["gen_grasp"] is False + assert env_cfg_override["rewards"]["rotate"]["weight"] == pytest.approx(1.25) + assert env_cfg_override["terminations"]["dropped"]["params"][ + "minimum_ball_height" + ] == pytest.approx(0.125) assert env_cfg_override["max_episode_seconds"] == pytest.approx(20.0) - assert env_cfg_override["grasp_cache_path"] == "caches/allegro_grasp_50k.npy" - assert env_cfg_override["domain_rand"]["randomize_base_mass"] is False - assert env_cfg_override["domain_rand"]["random_com"] is False - assert env_cfg_override["domain_rand"]["randomize_gravity"] is False - assert env_cfg_override["domain_rand"]["push_robots"] is False - assert env_cfg_override["domain_rand"]["joint_noise"] == pytest.approx(0.0) - assert env_cfg_override["domain_rand"]["ball_vel_noise"] == pytest.approx(0.0) - assert env_cfg_override["domain_rand"]["ball_z_offset"] == pytest.approx(0.0) + reset_params = env_cfg_override["events"]["reset_hand_ball"]["params"] + assert reset_params["grasp_cache_path"] is None + assert reset_params["joint_noise"] == pytest.approx(0.0) + assert reset_params["ball_velocity_noise"] == pytest.approx(0.0) + assert reset_params["ball_z_offset"] == pytest.approx(0.0) + assert env_cfg_override["observations"]["policy"]["history_length"] == 3 + assert env_cfg_override["actions"]["hand"]["action_scale"] == pytest.approx(1.0 / 24.0) + assert "reward_config" not in env_cfg_override + assert "domain_rand" not in env_cfg_override assert appo_cfg.algo.steps_per_env == cfg.algo.num_steps_per_env assert list(appo_cfg.algo.actor.hidden_dims) == list(cfg.algo.actor.hidden_dims) assert appo_cfg.algo.actor.activation == cfg.algo.actor.activation @@ -859,14 +863,15 @@ def test_build_ppo_env_cfg_override_allegro_mujoco( assert appo_motrix_cfg.training.sim_backend == ppo_motrix_cfg.training.sim_backend assert appo_motrix_cfg.algo.actor.obs_normalization is True assert appo_motrix_cfg.algo.critic.obs_normalization is True - assert appo_motrix_cfg.reward.scales.rotate == pytest.approx( - ppo_motrix_cfg.reward.scales.rotate - ) - assert appo_motrix_cfg.env.gen_grasp is ppo_motrix_cfg.env.gen_grasp - assert appo_motrix_cfg.env.domain_rand.randomize_base_mass is False - assert appo_motrix_cfg.env.domain_rand.random_com is False - assert appo_motrix_cfg.env.domain_rand.randomize_gravity is False - assert appo_motrix_cfg.env.domain_rand.push_robots is False + assert appo_motrix_cfg.reward.rotate.weight == pytest.approx( + ppo_motrix_cfg.reward.rotate.weight + ) + assert appo_motrix_cfg.env.events.pd_gains is None + assert ppo_motrix_cfg.env.events.pd_gains is None + assert ppo_drake_cfg.training.sim_backend == "drake" + assert appo_drake_cfg.training.sim_backend == "drake" + assert ppo_drake_cfg.env.events.pd_gains is None + assert appo_drake_cfg.env.events.pd_gains is None def test_build_ppo_env_cfg_override_allegro_grasp_mujoco( From 9ba3d8e63719a87b1640bc2c9db66d062cd07748 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Wed, 19 Aug 2026 07:04:40 +0800 Subject: [PATCH 082/169] feat(tasks): migrate Allegro grasp generation to Manager-Based runtime [skip ci] (#1210) --- .../ppo/task/allegro_inhand_grasp/motrix.yaml | 94 ++- .../ppo/task/allegro_inhand_grasp/mujoco.yaml | 94 ++- .../5-domain_randomization/0-index.md | 6 +- .../2-writing_providers.md | 2 +- .../1-sim_to_real/4-allegro_inhand.md | 10 +- .../5-domain_randomization/0-index.md | 6 +- .../2-writing_providers.md | 2 +- .../1-sim_to_real/4-allegro_inhand.md | 8 +- .../manipulation/allegro_inhand/__init__.py | 4 +- .../tasks/manipulation/allegro_inhand/base.py | 147 ----- .../manipulation/allegro_inhand/grasp_gen.py | 552 +++++++++++------- .../manipulation/allegro_inhand/rotation.py | 515 +--------------- tests/conftest.py | 18 - tests/envs/test_env_configs.py | 398 +++++++------ tests/scripts/test_train_scripts.py | 33 +- 15 files changed, 753 insertions(+), 1136 deletions(-) delete mode 100644 src/unilab/tasks/manipulation/allegro_inhand/base.py diff --git a/conf/ppo/task/allegro_inhand_grasp/motrix.yaml b/conf/ppo/task/allegro_inhand_grasp/motrix.yaml index 803c49e7e..0006f9b43 100644 --- a/conf/ppo/task/allegro_inhand_grasp/motrix.yaml +++ b/conf/ppo/task/allegro_inhand_grasp/motrix.yaml @@ -1,15 +1,17 @@ # @package _global_ -# Temporary last-consumer owner for the legacy grasp-generation env. Remove -# with the dedicated AllegroInhandRotationGrasp Manager-Based migration child. +defaults: + - /task/allegro_inhand/base + - _self_ training: task_name: AllegroInhandRotationGrasp sim_backend: motrix no_play: true + algo: num_envs: 16384 num_steps_per_env: 8 - max_iterations: 1000 # infinite rollout + max_iterations: 1000 # infinite rollout until the recorder raises RunComplete obs_groups: actor: [policy] critic: [policy] @@ -30,33 +32,69 @@ algo: algorithm: value_loss_coef: 4.0 desired_kl: 0.02 -reward: - scales: - rotate: 0.0 - obj_linvel: 0.0 - pose_diff: 0.0 - torque: 0.0 - work: 0.0 - drop: 0.0 - angvel_clip_min: -0.5 - angvel_clip_max: 0.5 - reset_z_threshold: 0.125 + env: - gen_grasp: true max_episode_seconds: 3.0 - grasp_cache_path: caches/allegro_grasp_50k.npy - grasp_collection_target: 50000 - grasp_auto_save: true - grasp_quality_check: true - grasp_min_contacts: 2 - domain_rand: - randomize_base_mass: false - random_com: false - randomize_gravity: false - push_robots: false - ball_z_offset: 0.0 - ball_vel_noise: 0.0 - joint_noise: 0.25 # random sampling of the grasp poses + actions: + hand: + action_scale: 0.0 + events: + pd_gains: null + reset_hand_ball: + params: + grasp_cache_path: null + joint_noise: 0.25 + ball_velocity_noise: 0.0 + ball_z_offset: 0.0 + terminations: + invalid_grasp: + _target_: unilab.managers.TerminationTermCfg + func: unilab.tasks.manipulation.allegro_inhand.grasp_gen.AllegroGraspQualityTermination + params: + entity_name: robot + observation_group: policy + observation_term: rotation + fingertip_body_names: [ff_tip, mf_tip, rf_tip, th_tip] + contact_sensor_names: [ff_contact, mf_contact, rf_contact, th_contact] + max_fingertip_distance: 0.1 + minimum_contacts: 2 + minimum_ball_height: 0.125 + enabled: true + metrics: + fingertips_close: + _target_: unilab.managers.MetricsTermCfg + func: unilab.tasks.manipulation.allegro_inhand.grasp_gen.AllegroGraspQualityMetric + params: {quality_term_name: invalid_grasp, condition: fingertips_close} + enough_contacts: + _target_: unilab.managers.MetricsTermCfg + func: unilab.tasks.manipulation.allegro_inhand.grasp_gen.AllegroGraspQualityMetric + params: {quality_term_name: invalid_grasp, condition: enough_contacts} + ball_held: + _target_: unilab.managers.MetricsTermCfg + func: unilab.tasks.manipulation.allegro_inhand.grasp_gen.AllegroGraspQualityMetric + params: {quality_term_name: invalid_grasp, condition: ball_held} + valid: + _target_: unilab.managers.MetricsTermCfg + func: unilab.tasks.manipulation.allegro_inhand.grasp_gen.AllegroGraspQualityMetric + params: {quality_term_name: invalid_grasp, condition: valid} + recorders: + grasp_cache: + _target_: unilab.managers.RecorderTermCfg + func: unilab.tasks.manipulation.allegro_inhand.grasp_gen.AllegroGraspRecorder + params: + quality_term_name: invalid_grasp + output_path: caches/allegro_grasp_50k.npy + collection_target: 50000 + auto_save: true + +reward: + rotate: {weight: 0.0} + obj_linvel: {weight: 0.0} + pose_diff: {weight: 0.0} + torque: {weight: 0.0} + work: {weight: 0.0} + drop: {weight: 0.0} + play_profile: enabled: true env: diff --git a/conf/ppo/task/allegro_inhand_grasp/mujoco.yaml b/conf/ppo/task/allegro_inhand_grasp/mujoco.yaml index e930d683d..9d16b9809 100644 --- a/conf/ppo/task/allegro_inhand_grasp/mujoco.yaml +++ b/conf/ppo/task/allegro_inhand_grasp/mujoco.yaml @@ -1,6 +1,7 @@ # @package _global_ -# Temporary last-consumer owner for the legacy grasp-generation env. Remove -# with the dedicated AllegroInhandRotationGrasp Manager-Based migration child. +defaults: + - /task/allegro_inhand/base + - _self_ training: task_name: AllegroInhandRotationGrasp @@ -10,10 +11,11 @@ training: cam_distance: 1.5 cam_lookat: [0.75, 0.75, 0] cam_elevation: -20.0 + algo: num_envs: 16384 num_steps_per_env: 8 - max_iterations: 1000 # infinite rollout + max_iterations: 1000 # infinite rollout until the recorder raises RunComplete obs_groups: actor: [policy] critic: [policy] @@ -34,33 +36,69 @@ algo: algorithm: value_loss_coef: 4.0 desired_kl: 0.02 -reward: - scales: - rotate: 0.0 - obj_linvel: 0.0 - pose_diff: 0.0 - torque: 0.0 - work: 0.0 - drop: 0.0 - angvel_clip_min: -0.5 - angvel_clip_max: 0.5 - reset_z_threshold: 0.125 + env: - gen_grasp: true max_episode_seconds: 3.0 - grasp_cache_path: caches/allegro_grasp_50k.npy - grasp_collection_target: 50000 - grasp_auto_save: true - grasp_quality_check: true - grasp_min_contacts: 2 - domain_rand: - randomize_base_mass: false - random_com: false - randomize_gravity: false - push_robots: false - ball_z_offset: 0.0 - ball_vel_noise: 0.0 - joint_noise: 0.25 # random sampling of the grasp poses + actions: + # The collector holds each sampled reset pose; policy output is intentionally ignored. + hand: + action_scale: 0.0 + events: + reset_hand_ball: + params: + grasp_cache_path: null + joint_noise: 0.25 + ball_velocity_noise: 0.0 + ball_z_offset: 0.0 + terminations: + invalid_grasp: + _target_: unilab.managers.TerminationTermCfg + func: unilab.tasks.manipulation.allegro_inhand.grasp_gen.AllegroGraspQualityTermination + params: + entity_name: robot + observation_group: policy + observation_term: rotation + fingertip_body_names: [ff_tip, mf_tip, rf_tip, th_tip] + contact_sensor_names: [ff_contact, mf_contact, rf_contact, th_contact] + max_fingertip_distance: 0.1 + minimum_contacts: 2 + minimum_ball_height: 0.125 + enabled: true + metrics: + fingertips_close: + _target_: unilab.managers.MetricsTermCfg + func: unilab.tasks.manipulation.allegro_inhand.grasp_gen.AllegroGraspQualityMetric + params: {quality_term_name: invalid_grasp, condition: fingertips_close} + enough_contacts: + _target_: unilab.managers.MetricsTermCfg + func: unilab.tasks.manipulation.allegro_inhand.grasp_gen.AllegroGraspQualityMetric + params: {quality_term_name: invalid_grasp, condition: enough_contacts} + ball_held: + _target_: unilab.managers.MetricsTermCfg + func: unilab.tasks.manipulation.allegro_inhand.grasp_gen.AllegroGraspQualityMetric + params: {quality_term_name: invalid_grasp, condition: ball_held} + valid: + _target_: unilab.managers.MetricsTermCfg + func: unilab.tasks.manipulation.allegro_inhand.grasp_gen.AllegroGraspQualityMetric + params: {quality_term_name: invalid_grasp, condition: valid} + recorders: + grasp_cache: + _target_: unilab.managers.RecorderTermCfg + func: unilab.tasks.manipulation.allegro_inhand.grasp_gen.AllegroGraspRecorder + params: + quality_term_name: invalid_grasp + output_path: caches/allegro_grasp_50k.npy + collection_target: 50000 + auto_save: true + +reward: + rotate: {weight: 0.0} + obj_linvel: {weight: 0.0} + pose_diff: {weight: 0.0} + torque: {weight: 0.0} + work: {weight: 0.0} + drop: {weight: 0.0} + play_profile: enabled: true env: diff --git a/docs/sphinx/source/en/2-user_guide/5-domain_randomization/0-index.md b/docs/sphinx/source/en/2-user_guide/5-domain_randomization/0-index.md index 3ee8a2c20..8d50e3800 100644 --- a/docs/sphinx/source/en/2-user_guide/5-domain_randomization/0-index.md +++ b/docs/sphinx/source/en/2-user_guide/5-domain_randomization/0-index.md @@ -33,7 +33,8 @@ These three paths correspond to three lifecycle classes: | `G1WalkFlat` | Yes | Yes: `Domain_Rand + Provider + ResetPlan` | task state sampling + common payload | push | `g1/joystick.py` | | `G1WalkRough` | Yes | Yes: reuses `G1WalkDomainRandomizationProvider` | task state sampling + common payload | push | `g1/joystick.py` | | `G1MotionTracking` | Yes | Yes: `Domain_Rand + Provider + ResetPlan` | extensive task-specific reset sampling + common payload | push | `motion_tracking/g1/tracking.py` | -| `AllegroInhandRotation` | Yes | Yes: `DomainRandConfig + Provider + ResetPlan` | task-specific reset sampling + common payload | none | `allegro_inhand/rotation.py` | +| `AllegroInhandRotation` | Yes | Yes: Hydra `EventTermCfg` + Manager-Based reset term | entity-scoped hand/ball reset | none | `allegro_inhand/manager_terms.py` | +| `AllegroInhandRotationGrasp` | Yes | Yes: reuses the rotation reset event + `RecorderTermCfg` | noisy hand reset + grasp collection | none | `allegro_inhand/grasp_gen.py` | | `SharpaInhandRotation` | Yes | Yes: `InitRandomizationPlan + ResetPlan + IntervalRandomizationPlan` | grasp cache sampling + common payload | object `body_force` | `sharpa_inhand/rotation.py` | | `SharpaInhandRotationGrasp` | Yes | Yes: reuses the Sharpa rotation provider and overrides reset sampling | grasp collection reset + common payload | none | `sharpa_inhand/grasp_gen.py` | @@ -46,7 +47,8 @@ These three paths correspond to three lifecycle classes: | `G1WalkFlat` | base xy; base yaw; base qvel sampled by `reset_base_qvel_limit`; command sampling; `gait_phase` sampling; `current_actions/last_actions` zeroed; kp/kd randomization (enabled by default); optional `base_mass_delta`; optional `base_com_offset`; optional `gravity` | `push_robots` | kp/kd enabled by default; common payload and push disabled by default | | `G1WalkRough` | Same as `G1WalkFlat`, directly reuses the same provider | `push_robots` | kp/kd enabled by default; common payload and push disabled by default | | `G1MotionTracking` | motion frame sampling; root pose perturbation `x/y/z/roll/pitch/yaw`; root velocity perturbation `x/y/z/roll/pitch/yaw`; joint position noise; under MuJoCo clipped by joint range; `current_actions/last_actions` zeroed; optional `base_mass_delta`; optional `base_com_offset`; optional `gravity` | `push_robots` | `pose_randomization`, `velocity_randomization`, `joint_position_range` have non-zero perturbations by default; common payload and push disabled by default | -| `AllegroInhandRotation` | If a grasp cache exists, sample a grasp randomly; otherwise apply `joint_noise` to hand joints and `ball_z_offset` to the ball; always apply `ball_vel_noise` to ball linear velocity; optional common reset randomization payload (incl. `gravity`) | none | If the grasp cache path is available it is sampled by default; `joint_noise`, `ball_vel_noise`, `ball_z_offset` default to 0; common payload disabled by default | +| `AllegroInhandRotation` | Entity-scoped hand/ball reset; an explicitly configured grasp cache is sampled, otherwise `null` explicitly selects the model home pose; optional `joint_noise`, `ball_velocity_noise`, and `ball_z_offset` | none | owner YAML explicitly selects the home pose and zero reset noise; a configured missing or malformed cache fails closed | +| `AllegroInhandRotationGrasp` | Reuses the rotation reset with `joint_noise=0.25`; Manager-Based termination checks fingertip distance, contact count, and ball height; recorder stores successful timeout rows | none | generates the 50k-row Allegro grasp cache and raises `RunComplete` after a successful save | | `SharpaInhandRotation` | grasp cache bucketed sampling by `scale_ids`; object pose / quat reset; optional common reset randomization payload (incl. `gravity`) | object `body_force` direct force disturbance | `domain_rand.scale_list` defaults come from the owner YAML; under MuJoCo, object geom scale is materialized during init; common payload disabled by default; object force enabled by default via the Sharpa owner YAML | | `SharpaInhandRotationGrasp` | hand pose reset; object pose / quat reset; collects successful grasps and stores them bucketed by `scale_ids`; optional `base_mass_delta`; optional `base_com_offset`; optional `gravity` | none | Used by default to generate the Sharpa grasp cache; cache filename includes the single scale value; common payload disabled by default | 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 a0fc87009..4f277d09d 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 @@ -30,7 +30,7 @@ Representative provider implementations are in: - `src/unilab/tasks/locomotion/go1/joystick.py` - `src/unilab/tasks/locomotion/g1/joystick.py` - `src/unilab/tasks/motion_tracking/g1/tracking.py` -- `src/unilab/tasks/manipulation/allegro_inhand/rotation.py` +- `src/unilab/tasks/locomotion/go2/footstand.py` - `src/unilab/tasks/manipulation/sharpa_inhand/rotation.py` Developer contract details are in diff --git a/docs/sphinx/source/en/3-deployment/1-sim_to_real/4-allegro_inhand.md b/docs/sphinx/source/en/3-deployment/1-sim_to_real/4-allegro_inhand.md index e5105b9b0..e6351dde7 100644 --- a/docs/sphinx/source/en/3-deployment/1-sim_to_real/4-allegro_inhand.md +++ b/docs/sphinx/source/en/3-deployment/1-sim_to_real/4-allegro_inhand.md @@ -67,10 +67,12 @@ to the grasp generator**, retrain, and try again. ## Action interface -The manipulation envs map policy actions to joint position targets through the -task control config (`src/unilab/tasks/manipulation/allegro_inhand/base.py` and -`src/unilab/tasks/manipulation/sharpa_inhand/base.py`). The deploy controller -must use the same joint order, action scale, and limit policy. +The manipulation envs map policy actions to joint position targets through their +task control config. Allegro owns this declaration in +`conf/ppo/task/allegro_inhand/base.yaml` and its Manager-Based action term; +Sharpa currently owns it in `src/unilab/tasks/manipulation/sharpa_inhand/base.py`. +The deploy controller must use the same joint order, action scale, and limit +policy. ## Failure recovery diff --git a/docs/sphinx/source/zh_CN/2-user_guide/5-domain_randomization/0-index.md b/docs/sphinx/source/zh_CN/2-user_guide/5-domain_randomization/0-index.md index 2324ced4e..e8dd9c389 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/5-domain_randomization/0-index.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/5-domain_randomization/0-index.md @@ -33,7 +33,8 @@ | `G1WalkFlat` | 是 | 是：`Domain_Rand + Provider + ResetPlan` | task 状态采样 + common payload | push | `g1/joystick.py` | | `G1WalkRough` | 是 | 是：复用 `G1WalkDomainRandomizationProvider` | task 状态采样 + common payload | push | `g1/joystick.py` | | `G1MotionTracking` | 是 | 是：`Domain_Rand + Provider + ResetPlan` | 大量 task 专属的 reset 采样 + common payload | push | `motion_tracking/g1/tracking.py` | -| `AllegroInhandRotation` | 是 | 是：`DomainRandConfig + Provider + ResetPlan` | task 专属的 reset 采样 + common payload | 无 | `allegro_inhand/rotation.py` | +| `AllegroInhandRotation` | 是 | 是：Hydra `EventTermCfg` + Manager-Based reset term | entity 范围的手/球 reset | 无 | `allegro_inhand/manager_terms.py` | +| `AllegroInhandRotationGrasp` | 是 | 是：复用 rotation reset event + `RecorderTermCfg` | 带噪声的手部 reset + grasp 收集 | 无 | `allegro_inhand/grasp_gen.py` | | `SharpaInhandRotation` | 是 | 是：`InitRandomizationPlan + ResetPlan + IntervalRandomizationPlan` | grasp cache 采样 + common payload | 物体 `body_force` | `sharpa_inhand/rotation.py` | | `SharpaInhandRotationGrasp` | 是 | 是：复用 Sharpa rotation provider 并 override reset 采样 | grasp 收集 reset + common payload | 无 | `sharpa_inhand/grasp_gen.py` | @@ -46,7 +47,8 @@ | `G1WalkFlat` | base xy；base yaw；由 `reset_base_qvel_limit` 采样的 base qvel；command 采样；`gait_phase` 采样；`current_actions/last_actions` 清零；kp/kd 随机化（默认启用）；可选 `base_mass_delta`；可选 `base_com_offset`；可选 `gravity` | `push_robots` | kp/kd 默认启用；common payload 和 push 默认禁用 | | `G1WalkRough` | 与 `G1WalkFlat` 相同，直接复用同一 provider | `push_robots` | kp/kd 默认启用；common payload 和 push 默认禁用 | | `G1MotionTracking` | 动作帧采样；root 位姿扰动 `x/y/z/roll/pitch/yaw`；root 速度扰动 `x/y/z/roll/pitch/yaw`；关节位置噪声；在 MuJoCo 下被关节范围 clip；`current_actions/last_actions` 清零；可选 `base_mass_delta`；可选 `base_com_offset`；可选 `gravity` | `push_robots` | `pose_randomization`、`velocity_randomization`、`joint_position_range` 默认有非零扰动；common payload 和 push 默认禁用 | -| `AllegroInhandRotation` | 若存在 grasp cache，则随机采样一个 grasp；否则对手部关节施加 `joint_noise` 并对球施加 `ball_z_offset`；始终对球的线速度施加 `ball_vel_noise`；可选 common reset 随机化 payload（含 `gravity`） | 无 | 若 grasp cache 路径可用则默认采样；`joint_noise`、`ball_vel_noise`、`ball_z_offset` 默认为 0；common payload 默认禁用 | +| `AllegroInhandRotation` | entity 范围的手/球 reset；显式配置 grasp cache 时进行采样，否则以 `null` 显式选择模型 home pose；可选 `joint_noise`、`ball_velocity_noise` 与 `ball_z_offset` | 无 | owner YAML 显式选择 home pose 与零 reset 噪声；配置的 cache 缺失或格式错误时 fail-closed | +| `AllegroInhandRotationGrasp` | 复用 rotation reset 并设置 `joint_noise=0.25`；Manager-Based termination 检查指尖距离、接触数和球高度；recorder 保存成功 timeout rows | 无 | 生成 5 万行 Allegro grasp cache，成功保存后抛出 `RunComplete` | | `SharpaInhandRotation` | grasp cache 按 `scale_ids` 分桶采样；物体位姿 / quat reset；可选 common reset 随机化 payload（含 `gravity`） | 物体 `body_force` 直接力扰动 | `domain_rand.scale_list` 默认值来自 owner YAML；在 MuJoCo 下，物体 geom 缩放在 init 期间 materialize；common payload 默认禁用；物体 force 通过 Sharpa owner YAML 默认启用 | | `SharpaInhandRotationGrasp` | 手部位姿 reset；物体位姿 / quat reset；收集成功的 grasp 并按 `scale_ids` 分桶存储；可选 `base_mass_delta`；可选 `base_com_offset`；可选 `gravity` | 无 | 默认用于生成 Sharpa grasp cache；cache 文件名包含单个 scale 值；common payload 默认禁用 | 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 08b4fee9d..6f36f1c30 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 @@ -29,7 +29,7 @@ - `src/unilab/tasks/locomotion/go1/joystick.py` - `src/unilab/tasks/locomotion/g1/joystick.py` - `src/unilab/tasks/motion_tracking/g1/tracking.py` -- `src/unilab/tasks/manipulation/allegro_inhand/rotation.py` +- `src/unilab/tasks/locomotion/go2/footstand.py` - `src/unilab/tasks/manipulation/sharpa_inhand/rotation.py` 开发者 contract 详情见 diff --git a/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/4-allegro_inhand.md b/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/4-allegro_inhand.md index a202759f8..c325fdac2 100644 --- a/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/4-allegro_inhand.md +++ b/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/4-allegro_inhand.md @@ -62,10 +62,10 @@ owner 与部署运行时在观测时序上达成一致。见 ## 动作接口 -操作类环境通过任务控制配置把策略动作映射为关节位置目标 -（`src/unilab/tasks/manipulation/allegro_inhand/base.py` 和 -`src/unilab/tasks/manipulation/sharpa_inhand/base.py`）。部署控制器必须使用相同的 -关节顺序、动作缩放与限位策略。 +操作类环境通过任务控制配置把策略动作映射为关节位置目标。Allegro 的声明由 +`conf/ppo/task/allegro_inhand/base.yaml` 与其 Manager-Based action term 持有； +Sharpa 当前仍由 `src/unilab/tasks/manipulation/sharpa_inhand/base.py` 持有。 +部署控制器必须使用相同的关节顺序、动作缩放与限位策略。 ## 失败恢复 diff --git a/src/unilab/tasks/manipulation/allegro_inhand/__init__.py b/src/unilab/tasks/manipulation/allegro_inhand/__init__.py index bb042cf8a..289cb394f 100644 --- a/src/unilab/tasks/manipulation/allegro_inhand/__init__.py +++ b/src/unilab/tasks/manipulation/allegro_inhand/__init__.py @@ -1,4 +1,4 @@ from . import ( - grasp_gen, # registers AllegroInhandRotationGrasp via @registry decorators - rotation, # registers AllegroInhandRotation via @registry decorators + grasp_gen, # registers the Manager-Based grasp-generation identity + rotation, # registers the Manager-Based rotation identity ) diff --git a/src/unilab/tasks/manipulation/allegro_inhand/base.py b/src/unilab/tasks/manipulation/allegro_inhand/base.py deleted file mode 100644 index 2a93735a8..000000000 --- a/src/unilab/tasks/manipulation/allegro_inhand/base.py +++ /dev/null @@ -1,147 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass, field - -import gymnasium as gym -import numpy as np - -from unilab.base.backend import SimBackend -from unilab.base.base import EnvCfg -from unilab.base.np_env import NpEnv, NpEnvState -from unilab.dtype_config import get_global_dtype - - -@dataclass -class NoiseConfig: - level: float = 1.0 - scale_joint_angle: float = 0.02 - - -@dataclass -class ControlConfig: - action_scale: float = 1.0 / 24.0 - kp: float = 1.0 - kd: float = 0.1 - - -@dataclass -class AllegroBaseCfg(EnvCfg): - sim_dt: float = 0.005 - ctrl_dt: float = 0.05 - noise_config: NoiseConfig = field(default_factory=NoiseConfig) - control_config: ControlConfig = field(default_factory=ControlConfig) - - -class AllegroBaseEnv(NpEnv): - _NUM_HAND_DOF: int = 16 - _FINGERTIP_BODY_NAMES: tuple[str, ...] = ("ff_tip", "mf_tip", "rf_tip", "th_tip") - _cfg: AllegroBaseCfg - _init_qpos: np.ndarray - _init_qvel: np.ndarray - - def __init__(self, cfg: AllegroBaseCfg, backend: SimBackend, num_envs: int = 1): - super().__init__(cfg, backend, num_envs) - - self._np_dtype = get_global_dtype() - actuator_range = np.asarray(self._backend.get_actuator_ctrl_range(), dtype=self._np_dtype) - if actuator_range.shape[0] < self._NUM_HAND_DOF: - raise ValueError( - f"Model has {actuator_range.shape[0]} actuators, expected at least {self._NUM_HAND_DOF}" - ) - self._ctrl_lower = np.asarray(actuator_range[: self._NUM_HAND_DOF, 0], dtype=self._np_dtype) - self._ctrl_upper = np.asarray(actuator_range[: self._NUM_HAND_DOF, 1], dtype=self._np_dtype) - - self._init_action_space() - self._num_action = self._action_space.shape[0] - if self._num_action != self._NUM_HAND_DOF: - raise ValueError(f"Expected {self._NUM_HAND_DOF} actuators, got {self._num_action}") - - self._init_buffers() - self.nq = int(self._init_qpos.shape[0]) - self.nv = int(self._init_qvel.shape[0]) - - self._ball_body_ids = self._backend.get_body_ids(["ball"]) - self._fingertip_body_ids = self._backend.get_body_ids(self._FINGERTIP_BODY_NAMES) - - def _init_action_space(self) -> None: - self._action_space = gym.spaces.Box( - low=-1.0, - high=1.0, - shape=(self._NUM_HAND_DOF,), - dtype=np.float32, - ) - - @property - def action_space(self) -> gym.spaces.Box: - return self._action_space # type: ignore[no-any-return] - - def _init_buffers(self) -> None: - self.default_angles = np.zeros((self._num_action,), dtype=self._np_dtype) - self._init_qpos = self._resolve_init_qpos() - self.default_angles = np.asarray( - self._init_qpos[: self._NUM_HAND_DOF], dtype=self._np_dtype - ) - self._init_qvel = np.asarray(self._backend.get_init_qvel(), dtype=self._np_dtype) - - def _resolve_init_qpos(self) -> np.ndarray: - for key_name in ("home", "stand", "default"): - try: - return np.asarray(self._backend.get_keyframe_qpos(key_name), dtype=self._np_dtype) - except Exception: - continue - raise ValueError("Could not resolve initial qpos from backend keyframes") - - def apply_action(self, actions: np.ndarray, state: NpEnvState) -> np.ndarray: - clipped_actions = np.asarray(np.clip(actions, -1.0, 1.0), dtype=self._np_dtype) - state.info["last_actions"] = state.info.get( - "current_actions", np.zeros_like(clipped_actions) - ) - state.info["current_actions"] = clipped_actions - - prev_ctrl = state.info.get( - "prev_ctrl", - np.broadcast_to( - self.default_angles, (clipped_actions.shape[0], self._num_action) - ).copy(), - ) - new_ctrl = prev_ctrl + self._cfg.control_config.action_scale * clipped_actions - new_ctrl = np.clip(new_ctrl, self._ctrl_lower, self._ctrl_upper) - prev_ctrl = np.asarray(new_ctrl, dtype=self._np_dtype) - state.info["prev_ctrl"] = prev_ctrl - return prev_ctrl - - def get_hand_dof_pos(self) -> np.ndarray: - return np.asarray( - self._backend.get_dof_pos()[:, : self._NUM_HAND_DOF], - dtype=self._np_dtype, - ) - - def get_hand_dof_vel(self) -> np.ndarray: - return np.asarray( - self._backend.get_dof_vel()[:, : self._NUM_HAND_DOF], - dtype=self._np_dtype, - ) - - def get_ball_pos(self) -> np.ndarray: - return np.asarray( - self._backend.get_body_pos_w(self._ball_body_ids)[:, 0, :], - dtype=self._np_dtype, - ) - - def get_ball_quat(self) -> np.ndarray: - return np.asarray( - self._backend.get_body_quat_w(self._ball_body_ids)[:, 0, :], - dtype=self._np_dtype, - ) - - def get_fingertip_pos(self) -> np.ndarray: - return np.asarray( - self._backend.get_body_pos_w(self._fingertip_body_ids), - dtype=self._np_dtype, - ) - - def get_sensor_data(self, name: str) -> np.ndarray: - return np.asarray(self._backend.get_sensor_data(name), dtype=self._np_dtype) - - -AllegroBaseMjEnv = AllegroBaseEnv diff --git a/src/unilab/tasks/manipulation/allegro_inhand/grasp_gen.py b/src/unilab/tasks/manipulation/allegro_inhand/grasp_gen.py index 2e9c30ec6..d3c610b45 100644 --- a/src/unilab/tasks/manipulation/allegro_inhand/grasp_gen.py +++ b/src/unilab/tasks/manipulation/allegro_inhand/grasp_gen.py @@ -1,238 +1,382 @@ -"""Allegro grasp-generation task built on top of rotation env.""" +"""Manager-Based terms and registry owner for Allegro grasp generation.""" from __future__ import annotations -from dataclasses import dataclass, field +from numbers import Integral, Real from pathlib import Path +from typing import TYPE_CHECKING, Any, Protocol, cast import numpy as np from unilab.assets import ASSETS_ROOT_PATH from unilab.base import registry -from unilab.base.np_env import NpEnvState from unilab.base.run_control import RunComplete +from unilab.dtype_config import get_global_dtype +from unilab.envs import ManagerBasedRlEnvCfg, make_manager_based_rl_env +from unilab.managers import ManagerTermBase, ManagerTermBaseCfg, RecorderTerm, RecorderTermCfg + +from .manager_terms import AllegroRotationObservation + +if TYPE_CHECKING: + from unilab.base.entity import Entity + from unilab.managers._types import ManagerBasedRlEnv + from unilab.managers.action_manager import ActionManager + from unilab.managers.observation_manager import ObservationManager + from unilab.managers.termination_manager import TerminationManager + + class _GraspEnv(ManagerBasedRlEnv, Protocol): + @property + def common_step_counter(self) -> int: ... + + @property + def action_manager(self) -> ActionManager: ... + + @property + def observation_manager(self) -> ObservationManager: ... + + @property + def termination_manager(self) -> TerminationManager: ... + + @property + def reset_terminated(self) -> np.ndarray: ... + + @property + def reset_time_outs(self) -> np.ndarray: ... + + @property + def extras(self) -> dict[str, Any]: ... + + +def _name(term: str, field: str, value: Any) -> str: + if not isinstance(value, str) or not value: + raise ValueError(f"{term} {field} must be a non-empty string") + return value + + +def _names(term: str, field: str, value: Any) -> tuple[str, ...]: + if isinstance(value, (str, bytes)) or not isinstance(value, (tuple, list)): + raise TypeError(f"{term} {field} must be a sequence of names") + result = tuple(_name(term, field, item) for item in value) + if not result: + raise ValueError(f"{term} {field} must not be empty") + if len(set(result)) != len(result): + raise ValueError(f"{term} {field} must contain unique names") + return result + + +def _real(term: str, field: str, value: Any, *, positive: bool = False) -> float: + if isinstance(value, (bool, np.bool_)) or not isinstance(value, Real): + raise TypeError(f"{term} {field} must be a real number") + result = float(value) + if not np.isfinite(result): + raise ValueError(f"{term} {field} must be finite") + if positive and result <= 0.0: + raise ValueError(f"{term} {field} must be positive") + return result + + +def _positive_int(term: str, field: str, value: Any) -> int: + if isinstance(value, (bool, np.bool_)) or not isinstance(value, Integral): + raise TypeError(f"{term} {field} must be an integer") + result = int(value) + if result <= 0: + raise ValueError(f"{term} {field} must be positive") + return result + + +def _bool(term: str, field: str, value: Any) -> bool: + if not isinstance(value, (bool, np.bool_)): + raise TypeError(f"{term} {field} must be boolean") + return bool(value) + + +class AllegroGraspQualityTermination(ManagerTermBase): + """Reject timeouts that are not stable multi-finger grasps.""" + + _ALLOWED_PARAMS = frozenset( + { + "entity_name", + "observation_group", + "observation_term", + "fingertip_body_names", + "contact_sensor_names", + "max_fingertip_distance", + "minimum_contacts", + "minimum_ball_height", + "enabled", + } + ) -from .rotation import AllegroRotationPPO, AllegroRotationPPOCfg, RewardConfigPPO - - -@registry.envcfg("AllegroInhandRotationGrasp") -@dataclass -class AllegroRotationGraspCfg(AllegroRotationPPOCfg): - # These are fallback defaults. Hydra task env overrides (e.g. - # conf/ppo/task/allegro_inhand_grasp/mujoco.yaml and CLI env.*) - # are applied at env construction and take precedence. - max_episode_seconds: float = 2.0 - reward_config: RewardConfigPPO = field( # pyright: ignore[reportIncompatibleVariableOverride] - default_factory=lambda: RewardConfigPPO( - scales={ - "rotate": 0.0, - "obj_linvel": 0.0, - "pose_diff": 0.0, - "torque": 0.0, - "work": 0.0, - "drop": 0.0, - }, - angvel_clip_min=-0.5, - angvel_clip_max=0.5, - reset_z_threshold=0.125, + def __init__(self, cfg: ManagerTermBaseCfg, env: _GraspEnv): + super().__init__(env) + term = type(self).__name__ + unexpected = set(cfg.params) - self._ALLOWED_PARAMS + if unexpected: + raise TypeError(f"{term} received unsupported parameters: {sorted(unexpected)}") + + entity_name = _name(term, "entity_name", cfg.params.get("entity_name")) + self._entity = cast("Entity", env.scene[entity_name]) + fingertip_names = _names( + term, "fingertip_body_names", cfg.params.get("fingertip_body_names") ) - ) - gen_grasp: bool = True - grasp_collection_target: int = 50_000 - grasp_auto_save: bool = True - grasp_quality_check: bool = True - grasp_min_contacts: int = 2 - - -@registry.env("AllegroInhandRotationGrasp", sim_backend="mujoco") -@registry.env("AllegroInhandRotationGrasp", sim_backend="motrix") -class AllegroRotationGrasp(AllegroRotationPPO): - _cfg: AllegroRotationGraspCfg # pyright: ignore[reportIncompatibleVariableOverride] - _CONTACT_SENSORS = ("ff_contact", "mf_contact", "rf_contact", "th_contact") - - def __init__( - self, cfg: AllegroRotationGraspCfg, num_envs: int = 1, backend_type: str = "mujoco" - ) -> None: - super().__init__(cfg, num_envs=num_envs, backend_type=backend_type) - self._saved_grasping_states: list[np.ndarray] = [] - self._grasp_cache_saved = False - self._grasp_target_reached_notified = False - - def apply_action(self, actions: np.ndarray, state: NpEnvState) -> np.ndarray: - del actions - zero_actions = np.zeros((self._num_envs, self._NUM_HAND_DOF), dtype=self._np_dtype) - return super().apply_action(zero_actions, state) - - @staticmethod - def _sensor_scalar(sensor_data: np.ndarray) -> np.ndarray: - sensor_data = np.asarray(sensor_data) - if sensor_data.ndim == 1: - return sensor_data - return sensor_data.reshape(sensor_data.shape[0], -1)[:, 0] - - def _contact_count(self) -> np.ndarray: - contacts = np.stack( - [self._sensor_scalar(self.get_sensor_data(name)) for name in self._CONTACT_SENSORS], - axis=1, + fingertip_ids, matched_names = self._entity.find_bodies( + fingertip_names, preserve_order=True ) - return np.asarray(np.sum(contacts > 0.5, axis=1), dtype=np.int32) - - def _compute_grasp_conditions(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]: - # Hora-style good grasp criteria from compute_reward: - # 1) all fingertips near object, 2) enough fingertip contacts, 3) object not dropped. - ball_pos = self.get_ball_pos() - fingertip_pos = self.get_fingertip_pos() - - cond1 = np.all(np.linalg.norm(fingertip_pos - ball_pos[:, None, :], axis=-1) < 0.1, axis=1) - cond2 = self._contact_count() >= int(self._cfg.grasp_min_contacts) - cond3 = ball_pos[:, 2] > float(self._reward_cfg.reset_z_threshold) - return ( - np.asarray(cond1, dtype=bool), - np.asarray(cond2, dtype=bool), - np.asarray(cond3, dtype=bool), + if tuple(matched_names) != fingertip_names: + raise ValueError( + f"{term} fingertip body order mismatch: expected {fingertip_names}, " + f"got {tuple(matched_names)}" + ) + self._fingertip_ids = np.asarray(fingertip_ids, dtype=np.intp) + self._fingertip_ids.setflags(write=False) + + contact_names = _names(term, "contact_sensor_names", cfg.params.get("contact_sensor_names")) + self._contact_view = env.scene.bind_sensor_data(contact_names) + if self._contact_view.dimensions != (1,) * len(contact_names): + raise ValueError( + f"{term} contact sensors must each be scalar, got " + f"{dict(zip(contact_names, self._contact_view.dimensions, strict=True))}" + ) + + group = _name(term, "observation_group", cfg.params.get("observation_group")) + name = _name(term, "observation_term", cfg.params.get("observation_term")) + observation = env.observation_manager.get_term_cfg(group, name).func + if not isinstance(observation, AllegroRotationObservation): + raise TypeError( + f"{term} observation {group}/{name} must be AllegroRotationObservation, " + f"got {type(observation).__name__}" + ) + self.observation = observation + self._maximum_distance = _real( + term, + "max_fingertip_distance", + cfg.params.get("max_fingertip_distance"), + positive=True, ) + self._minimum_contacts = _positive_int( + term, "minimum_contacts", cfg.params.get("minimum_contacts") + ) + if self._minimum_contacts > len(contact_names): + raise ValueError( + f"{term} minimum_contacts={self._minimum_contacts} exceeds " + f"{len(contact_names)} configured contact sensors" + ) + self._minimum_height = _real( + term, "minimum_ball_height", cfg.params.get("minimum_ball_height") + ) + self._enabled = _bool(term, "enabled", cfg.params.get("enabled")) + + self.fingertips_close = np.zeros(env.num_envs, dtype=np.bool_) + self.enough_contacts = np.zeros(env.num_envs, dtype=np.bool_) + self.ball_held = np.zeros(env.num_envs, dtype=np.bool_) + self.valid = np.zeros(env.num_envs, dtype=np.bool_) + self._disabled = np.zeros(env.num_envs, dtype=np.bool_) + self._last_counter = int(env.common_step_counter) + + @property + def last_counter(self) -> int: + return self._last_counter + + def reset(self, env_ids: np.ndarray | slice | None = None) -> None: + ids = slice(None) if env_ids is None else env_ids + self.fingertips_close[ids] = False + self.enough_contacts[ids] = False + self.ball_held[ids] = False + self.valid[ids] = False + self._last_counter = int(cast("_GraspEnv", self._env).common_step_counter) + + def __call__(self, env: _GraspEnv, **params: Any) -> np.ndarray: + del params + self.observation.snapshot(env) + fingertip_pos = self._entity.data.body_link_pos_w[:, self._fingertip_ids] + distance = np.linalg.norm(fingertip_pos - self.observation.ball_pos[:, None, :], axis=-1) + self.fingertips_close[:] = np.all(distance < self._maximum_distance, axis=1) + contacts = self._contact_view.read() + self.enough_contacts[:] = np.count_nonzero(contacts > 0.5, axis=1) >= self._minimum_contacts + self.ball_held[:] = self.observation.ball_pos[:, 2] > self._minimum_height + np.logical_and(self.fingertips_close, self.enough_contacts, out=self.valid) + np.logical_and(self.valid, self.ball_held, out=self.valid) + self._last_counter = int(env.common_step_counter) + if not self._enabled: + return self._disabled + return np.logical_not(self.valid) + + +class AllegroGraspQualityMetric(ManagerTermBase): + """Expose one cached quality condition through the community MetricsManager.""" + + def __init__(self, cfg: ManagerTermBaseCfg, env: _GraspEnv): + super().__init__(env) + term_name = type(self).__name__ + unexpected = set(cfg.params) - {"quality_term_name", "condition"} + if unexpected: + raise TypeError(f"{term_name} received unsupported parameters: {sorted(unexpected)}") + quality_name = _name(term_name, "quality_term_name", cfg.params.get("quality_term_name")) + quality = env.termination_manager.get_term_cfg(quality_name).func + if not isinstance(quality, AllegroGraspQualityTermination): + raise TypeError( + f"{term_name} termination term {quality_name!r} must be " + f"AllegroGraspQualityTermination, got {type(quality).__name__}" + ) + condition = _name(term_name, "condition", cfg.params.get("condition")) + values = { + "fingertips_close": quality.fingertips_close, + "enough_contacts": quality.enough_contacts, + "ball_held": quality.ball_held, + "valid": quality.valid, + } + try: + self._value = values[condition] + except KeyError: + raise ValueError( + f"{term_name} condition must be one of {sorted(values)}, got {condition!r}" + ) from None + self._quality = quality + self._quality_name = quality_name + + def __call__(self, env: _GraspEnv, **params: Any) -> np.ndarray: + del params + if self._quality.last_counter != int(env.common_step_counter): + raise RuntimeError( + f"{type(self).__name__} term {self._quality_name!r} was not computed for " + f"control step {env.common_step_counter}" + ) + return np.asarray(self._value, dtype=get_global_dtype()) + + +class AllegroGraspRecorder(RecorderTerm): + """Collect successful timeout states and persist the canonical 23-D cache.""" + + _ALLOWED_PARAMS = frozenset( + {"quality_term_name", "output_path", "collection_target", "auto_save"} + ) - def _check_grasp_quality(self, env_ids: np.ndarray) -> np.ndarray: - cond1, cond2, cond3 = self._compute_grasp_conditions() - return np.asarray(cond1[env_ids] & cond2[env_ids] & cond3[env_ids], dtype=bool) - - def _total_saved_grasps(self) -> int: - if not self._saved_grasping_states: - return 0 - return int(sum(states.shape[0] for states in self._saved_grasping_states)) - - def _stop_collection(self) -> None: - if self._grasp_target_reached_notified: + def __init__(self, cfg: RecorderTermCfg, env: _GraspEnv): + super().__init__(cfg, env) + term = type(self).__name__ + unexpected = set(cfg.params) - self._ALLOWED_PARAMS + if unexpected: + raise TypeError(f"{term} received unsupported parameters: {sorted(unexpected)}") + + quality_name = _name(term, "quality_term_name", cfg.params.get("quality_term_name")) + quality = env.termination_manager.get_term_cfg(quality_name).func + if not isinstance(quality, AllegroGraspQualityTermination): + raise TypeError( + f"{term} termination term {quality_name!r} must be " + f"AllegroGraspQualityTermination, got {type(quality).__name__}" + ) + self._quality = quality + output = Path(_name(term, "output_path", cfg.params.get("output_path"))) + self._output_path = output if output.is_absolute() else Path(ASSETS_ROOT_PATH) / output + self._target = _positive_int(term, "collection_target", cfg.params.get("collection_target")) + self._auto_save = _bool(term, "auto_save", cfg.params.get("auto_save")) + self._saved_states: list[np.ndarray] = [] + self._cache_saved = False + self._target_notified = False + + @property + def total_saved_grasps(self) -> int: + return int(sum(states.shape[0] for states in self._saved_states)) + + @property + def cache_saved(self) -> bool: + return self._cache_saved + + @property + def output_path(self) -> Path: + return self._output_path + + def _log(self, name: str, value: float) -> None: + env = cast("_GraspEnv", self._env) + log = env.extras.setdefault("log", {}) + log[name] = value + + def _save_cache(self, *, force: bool = False) -> None: + if self._cache_saved: return - - target = int(self._cfg.grasp_collection_target) - if target <= 0: + total = self.total_saved_grasps + if not force and total < self._target: return - - total = self._total_saved_grasps() - if total < target: + if total == 0: return - self._grasp_target_reached_notified = True + all_states = np.concatenate(self._saved_states, axis=0).astype(np.float32) + all_states = all_states[: self._target] + if all_states.ndim != 2 or all_states.shape[1] != 23: + raise ValueError( + f"{type(self).__name__} collected cache must have shape (N, 23), " + f"got {all_states.shape}" + ) + if not np.isfinite(all_states).all(): + raise ValueError(f"{type(self).__name__} collected cache contains NaN or Inf") + self._output_path.parent.mkdir(parents=True, exist_ok=True) + np.save(self._output_path, all_states) + self._cache_saved = True + self._log("grasp_cache/saved", 1.0) + self._log("grasp_cache/num_states", float(all_states.shape[0])) + + def _stop_collection(self) -> None: + if self._target_notified or self.total_saved_grasps < self._target: + return + total = self.total_saved_grasps + self._target_notified = True + self._log("grasp/target_reached", 1.0) print( "[AllegroInhandRotationGrasp] Grasp collection target reached " - f"({total}/{target}). Collection completed." + f"({total}/{self._target}). Collection completed." ) - - if self.state is not None: - log = self.state.info.get("log", {}) - log["grasp/target_reached"] = 1.0 - self.state.info["log"] = log - raise RunComplete( reason="grasp_collection_target_reached", summary={ - "collected_grasps": int(total), - "saved_grasps": int(min(total, target)), - "grasp_collection_target": target, + "collected_grasps": total, + "saved_grasps": min(total, self._target), + "grasp_collection_target": self._target, }, ) - def _save_grasp_cache(self, force: bool = False) -> None: - if self._grasp_cache_saved: - return - - total = self._total_saved_grasps() - target = int(self._cfg.grasp_collection_target) - if not force and total < target: - return - - if total == 0: - return - - all_states = np.concatenate(self._saved_grasping_states, axis=0).astype(np.float32) - if target > 0: - all_states = all_states[:target] - - output_file = Path(self._cfg.grasp_cache_path or "caches/allegro_grasp_50k.npy") - if not output_file.is_absolute(): - output_file = ASSETS_ROOT_PATH / output_file - output_file.parent.mkdir(parents=True, exist_ok=True) - np.save(output_file, all_states) - - self._grasp_cache_saved = True - if self.state is not None: - log = self.state.info.get("log", {}) - log["grasp_cache/saved"] = 1.0 - log["grasp_cache/num_states"] = float(all_states.shape[0]) - self.state.info["log"] = log - - def _collect_successful_grasps(self, env_ids: np.ndarray) -> None: - if self.state is None or env_ids.size == 0: - return - - success_mask = self.state.truncated[env_ids] & ~self.state.terminated[env_ids] - if not np.any(success_mask): + def record_pre_reset(self, env_ids: np.ndarray) -> None: + env = cast("_GraspEnv", self._env) + ids = np.asarray(env_ids, dtype=np.intp) + success = env.reset_time_outs[ids] & ~env.reset_terminated[ids] + success_ids = ids[np.flatnonzero(success)] + if success_ids.size == 0: return - - success_env_ids = env_ids[np.flatnonzero(success_mask)] - if self._cfg.grasp_quality_check: - quality_mask = self._check_grasp_quality(success_env_ids) - success_env_ids = success_env_ids[np.flatnonzero(quality_mask)] - - if success_env_ids.size == 0: - return - - curr_dof_pos = np.asarray(self.state.info.get("curr_dof_pos", self.get_hand_dof_pos())) - curr_ball_pos = np.asarray(self.state.info.get("curr_ball_pos", self.get_ball_pos())) - curr_ball_quat = np.asarray(self.state.info.get("curr_ball_quat", self.get_ball_quat())) - - hand_qpos = curr_dof_pos[success_env_ids, : self._NUM_HAND_DOF] - ball_pos = curr_ball_pos[success_env_ids] - ball_quat = curr_ball_quat[success_env_ids] - states = np.concatenate([hand_qpos, ball_pos, ball_quat], axis=1).astype(np.float32) - - self._saved_grasping_states.append(states) - self._save_grasp_cache() + if self._quality.last_counter != int(env.common_step_counter): + raise RuntimeError( + f"{type(self).__name__} quality state was not computed for control step " + f"{env.common_step_counter}" + ) + state = self._quality.observation + rows = np.concatenate( + ( + state.dof_pos[success_ids], + state.ball_pos[success_ids], + state.ball_quat[success_ids], + ), + axis=1, + dtype=np.float32, + ) + if rows.shape != (success_ids.size, 23): + raise ValueError( + f"{type(self).__name__} expected collected rows shape " + f"({success_ids.size}, 23), got {rows.shape}" + ) + self._saved_states.append(rows) + self._save_cache() self._stop_collection() - - if self.state is not None: - log = self.state.info.get("log", {}) - log["grasp/cache_size"] = float(self._total_saved_grasps()) - self.state.info["log"] = log - - def update_state(self, state: NpEnvState) -> NpEnvState: - next_state = super().update_state(state) - reward = np.zeros((self._num_envs,), dtype=self._np_dtype) - - cond1, cond2, cond3 = self._compute_grasp_conditions() - if self._cfg.grasp_quality_check: - grasp_valid = cond1 & cond2 & cond3 - terminated = np.asarray(next_state.terminated | (~grasp_valid), dtype=bool) - else: - grasp_valid = np.ones((self._num_envs,), dtype=bool) - terminated = np.asarray(next_state.terminated, dtype=bool) - - step_count = next_state.info.get("steps", np.zeros((self._num_envs,), dtype=np.uint32)) - should_log = self._enable_reward_log and (int(step_count[0]) % 4 == 0) - if should_log: - log = next_state.info.get("log", {}) - log["grasp/cond1"] = float(np.mean(cond1.astype(np.float32))) - log["grasp/cond2"] = float(np.mean(cond2.astype(np.float32))) - log["grasp/cond3"] = float(np.mean(cond3.astype(np.float32))) - log["grasp/valid"] = float(np.mean(grasp_valid.astype(np.float32))) - log["grasp/cache_size"] = float(self._total_saved_grasps()) - next_state.info["log"] = log - - return next_state.replace(reward=reward, terminated=terminated) - - def _reset_done_envs(self) -> None: - if self.state is not None: - done = self.state.terminated | self.state.truncated - if np.any(done): - env_ids = np.flatnonzero(done).astype(np.int32) - self._collect_successful_grasps(env_ids) - super()._reset_done_envs() + self._log("grasp/cache_size", float(self.total_saved_grasps)) def close(self) -> None: - self._save_grasp_cache(force=bool(self._cfg.grasp_auto_save)) - super().close() + self._save_cache(force=self._auto_save) + + +registry.register_env_config("AllegroInhandRotationGrasp", ManagerBasedRlEnvCfg) +registry.register_env("AllegroInhandRotationGrasp", make_manager_based_rl_env, sim_backend="mujoco") +registry.register_env("AllegroInhandRotationGrasp", make_manager_based_rl_env, sim_backend="motrix") -AllegroRotationGraspEnv = AllegroRotationGrasp -AllegroRotationGraspCfgAlias = AllegroRotationGraspCfg +__all__ = [ + "AllegroGraspQualityMetric", + "AllegroGraspQualityTermination", + "AllegroGraspRecorder", +] diff --git a/src/unilab/tasks/manipulation/allegro_inhand/rotation.py b/src/unilab/tasks/manipulation/allegro_inhand/rotation.py index 928cdd500..0dba2a728 100644 --- a/src/unilab/tasks/manipulation/allegro_inhand/rotation.py +++ b/src/unilab/tasks/manipulation/allegro_inhand/rotation.py @@ -1,521 +1,8 @@ -"""Allegro in-hand rotation environment.""" +"""Registry owner for the Allegro in-hand rotation Manager-Based task.""" -from __future__ import annotations - -import logging -from dataclasses import dataclass, field -from typing import Any, cast - -import numpy as np -from etils import epath - -from unilab.assets import ASSETS_ROOT_PATH from unilab.base import registry -from unilab.base.backend import create_backend, env_backend_kwargs -from unilab.base.np_env import NpEnvState -from unilab.base.scene import SceneCfg -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.dtype_config import get_global_dtype from unilab.envs import ManagerBasedRlEnvCfg, make_manager_based_rl_env -from unilab.utils.geometry import ( - np_normalize_axis, - np_quat_angular_velocity_from_pair, -) - -from .base import AllegroBaseCfg, AllegroBaseEnv - -logger = logging.getLogger(__name__) - - -def resolve_grasp_cache_path(cache_path: str) -> epath.Path: - """Resolve Allegro grasp cache paths using the asset-root convention.""" - path = epath.Path(cache_path) - if path.is_absolute() or path.exists(): - return path - return epath.Path(ASSETS_ROOT_PATH / cache_path) - - -def normalize_rotation_axis(rotation_axis: tuple[float, float, float]) -> np.ndarray: - # Cast to the training dtype first so the norm and division happen at that - # precision, matching the pre-refactor bit-exact behavior for float32 runs. - axis = np.asarray(rotation_axis, dtype=get_global_dtype()) - return np.asarray(np_normalize_axis(axis), dtype=get_global_dtype()) - - -def compute_ball_angvel( - ball_quat: np.ndarray, prev_ball_quat: np.ndarray, ctrl_dt: float -) -> np.ndarray: - return np.asarray( - np_quat_angular_velocity_from_pair(ball_quat, prev_ball_quat, ctrl_dt), - dtype=get_global_dtype(), - ) - - -def compute_pd_torques( - targets: np.ndarray, dof_pos: np.ndarray, dof_vel: np.ndarray, kp: float, kd: float -) -> np.ndarray: - torques = kp * (targets - dof_pos) - kd * dof_vel - return np.asarray(np.clip(torques, -0.5, 0.5), dtype=get_global_dtype()) - - -def build_obs_lag_history( - init_obs: np.ndarray, num_lag_steps: int, num_obs_per_step: int -) -> np.ndarray: - num_envs = init_obs.shape[0] - history = np.broadcast_to( - init_obs[:, None, :], - (num_envs, num_lag_steps, num_obs_per_step), - ).copy() - return np.asarray(history, dtype=init_obs.dtype) - - -def sample_cached_grasps( - grasp_cache: np.ndarray, num_reset: int -) -> tuple[np.ndarray, np.ndarray, np.ndarray]: - idx = np.random.randint(0, len(grasp_cache), size=num_reset) - sampled = grasp_cache[idx] - return sampled[:, :16], sampled[:, 16:19], sampled[:, 19:23] - - -def _materialize_grasp_cache(cfg: AllegroRotationPPOCfg) -> np.ndarray | None: - """Load the optional Allegro grasp cache on the env construction path.""" - if cfg.gen_grasp: - return None - - cache_path = resolve_grasp_cache_path(cfg.grasp_cache_path) - if not cache_path.exists(): - logger.warning( - "[allegro_inhand] Grasp cache is missing; no Hugging Face download will be " - "attempted. Expected local cache: %s. Generate one with " - "`uv run train --algo ppo --task allegro_inhand_grasp --sim mujoco " - "training.no_play=true`, or point `env.grasp_cache_path` at an existing " - "local cache.", - cache_path, - ) - return None - - grasp_cache = np.load(cache_path).astype(np.float64) - logger.info( - "[allegro_inhand] Loaded grasp cache: %s, shape=%s, dtype=%s", - cache_path, - grasp_cache.shape, - grasp_cache.dtype, - ) - return grasp_cache - - -@dataclass -class RewardConfigPPO: - scales: dict[str, float] - angvel_clip_min: float - angvel_clip_max: float - reset_z_threshold: float - - -@dataclass -class DomainRandConfig: - randomize_base_mass: bool = False - added_mass_range: list[float] = field(default_factory=lambda: [0.0, 0.0]) - random_com: bool = False - com_offset_x: list[float] = field(default_factory=lambda: [0.0, 0.0]) - randomize_gravity: bool = False - gravity_range: list[list[float]] = field( - default_factory=lambda: [[0.0, 0.0, -9.81], [0.0, 0.0, -9.81]] - ) - push_robots: bool = False - push_interval: int = 750 - max_force: list[float] = field(default_factory=lambda: [1.0, 1.0, 0.5]) - push_body_name: str | None = None - joint_noise: float = 0.0 - ball_vel_noise: float = 0.0 - ball_z_offset: float = 0.0 - - -@dataclass -class AllegroRotationPPOCfg(AllegroBaseCfg): - scene: SceneCfg = field( # pyright: ignore[reportIncompatibleVariableOverride] - default_factory=lambda: SceneCfg( - model_file=str(ASSETS_ROOT_PATH / "robots" / "allegro_hand" / "scene.xml") - ) - ) - max_episode_seconds: float = 20.0 # pyright: ignore[reportIncompatibleVariableOverride] - reward_config: RewardConfigPPO | None = None - domain_rand: DomainRandConfig = field(default_factory=DomainRandConfig) - rotation_axis: tuple[float, float, float] = (0.0, 0.0, 1.0) - grasp_cache_path: str = "caches/allegro_grasp_50k.npy" - gen_grasp: bool = False - - -class AllegroRotationDomainRandomizationProvider(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_reset_state( - self, env: Any, num_reset: int - ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: - dr = env.cfg.domain_rand - grasp_cache = cast(np.ndarray | None, env._grasp_cache) - if grasp_cache is not None: - hand_qpos, ball_pos, ball_quat = sample_cached_grasps(grasp_cache, num_reset) - else: - hand_qpos = np.broadcast_to(env.default_angles, (num_reset, env._NUM_HAND_DOF)).copy() - hand_qpos += np.random.uniform(-dr.joint_noise, dr.joint_noise, hand_qpos.shape).astype( - np.float64 - ) - hand_qpos = np.clip( - hand_qpos, - env._ctrl_lower.astype(np.float64), - env._ctrl_upper.astype(np.float64), - ) - ball_init_pos = env._init_qpos[env._NUM_HAND_DOF : env._NUM_HAND_DOF + 3] - ball_pos = np.broadcast_to(ball_init_pos, (num_reset, 3)).copy() - ball_pos[:, 2] += dr.ball_z_offset - ball_quat = np.tile([1.0, 0.0, 0.0, 0.0], (num_reset, 1)) - - qvel = np.zeros((num_reset, env.nv), dtype=np.float64) - qvel[:, env._NUM_HAND_DOF : env._NUM_HAND_DOF + 3] = np.random.uniform( - -dr.ball_vel_noise, - dr.ball_vel_noise, - (num_reset, 3), - ) - return hand_qpos, ball_pos, ball_quat, qvel - - def _build_info_updates( - self, - env: Any, - hand_qpos: np.ndarray, - ball_pos: np.ndarray, - ball_quat: np.ndarray, - ) -> dict[str, np.ndarray]: - num_reset = hand_qpos.shape[0] - dtype = get_global_dtype() - - init_ctrl = np.asarray(hand_qpos, dtype=dtype) - init_ball_pos = np.asarray(ball_pos, dtype=dtype) - dof_pos_norm = 2.0 * (init_ctrl - env._dof_mid) / (env._dof_range + 1e-8) - init_obs = np.concatenate([dof_pos_norm, init_ctrl, init_ball_pos], axis=1, dtype=dtype) - obs_lag_history = build_obs_lag_history(init_obs, env._NUM_LAG_STEPS, env._NUM_OBS_PER_STEP) - - return { - "current_actions": zero_actions(num_reset, env._num_action), - "last_actions": zero_actions(num_reset, env._num_action), - "prev_ctrl": init_ctrl, - "init_pose": init_ctrl.copy(), - "prev_dof_pos": init_ctrl.copy(), - "prev_ball_pos": init_ball_pos.copy(), - "prev_ball_quat": np.asarray(ball_quat, dtype=dtype).copy(), - "obs_lag_history": obs_lag_history, - } - - def build_reset_plan(self, env: Any, env_ids: np.ndarray) -> ResetPlan: - num_reset = len(env_ids) - hand_qpos, ball_pos, ball_quat, qvel = self._sample_reset_state(env, num_reset) - qpos = np.concatenate([hand_qpos, ball_pos, ball_quat], axis=1, dtype=np.float64) - info_updates = self._build_info_updates(env, hand_qpos, ball_pos, ball_quat) - - 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] - ) -> dict[str, np.ndarray]: - del env_ids - return cast( - dict[str, np.ndarray], - env._compute_obs( - info_updates, - info_updates["prev_ctrl"], - info_updates["prev_ball_pos"], - ), - ) - - -# ─────────────────────────── Environment ────────────────────────────── - - -class AllegroRotationPPO(AllegroBaseEnv): - _cfg: AllegroRotationPPOCfg # pyright: ignore[reportIncompatibleVariableOverride] - _reward_cfg: RewardConfigPPO - - _NUM_OBS_PER_STEP = 35 - _NUM_LAG_STEPS = 3 - - def __init__( - self, cfg: AllegroRotationPPOCfg, num_envs: int = 1, backend_type: str = "mujoco" - ) -> None: - if cfg.reward_config is None: - raise ValueError("reward_config must be provided via Hydra configuration") - backend = create_backend( - backend_type, - cfg.scene, - num_envs, - cfg.sim_dt, - base_name="palm", - push_body_name=cfg.domain_rand.push_body_name, - add_body_sensors=True, - position_actuator_gains={ - "kp": cfg.control_config.kp, - "kd": cfg.control_config.kd, - "actuator_ids": slice(0, 16), - }, - **env_backend_kwargs(cfg), - ) - super().__init__(cfg, backend, num_envs) - self._enable_reward_log = True - self._reward_cfg = cfg.reward_config - - self._dof_range = self._ctrl_upper - self._ctrl_lower - self._dof_mid = (self._ctrl_upper + self._ctrl_lower) / 2.0 - self._rot_axis = normalize_rotation_axis(cfg.rotation_axis) - self._grasp_cache = _materialize_grasp_cache(cfg) - - self._init_reward_functions() - self._init_domain_randomization(AllegroRotationDomainRandomizationProvider()) - - @property - def obs_groups_spec(self) -> dict[str, int]: - return {"obs": self._NUM_OBS_PER_STEP * self._NUM_LAG_STEPS} - - def _init_reward_functions(self) -> None: - self._reward_fns = { - "rotate": self._reward_rotate, - "obj_linvel": self._reward_obj_linvel, - "pose_diff": self._reward_pose_diff, - "torque": self._reward_torque, - "work": self._reward_work, - "drop": self._reward_drop, - } - - def _reward_rotate( - self, - info: dict[str, Any], - dof_pos: np.ndarray, - dof_vel: np.ndarray, - ball_pos: np.ndarray, - ball_linvel: np.ndarray, - ball_angvel: np.ndarray, - torques: np.ndarray, - terminated: np.ndarray, - ) -> np.ndarray: - del info, dof_pos, dof_vel, ball_pos, ball_linvel, torques, terminated - vec_dot = ball_angvel @ self._rot_axis - reward: np.ndarray = np.clip( - vec_dot, self._reward_cfg.angvel_clip_min, self._reward_cfg.angvel_clip_max - ) - return reward - - def _reward_obj_linvel( - self, - info: dict[str, Any], - dof_pos: np.ndarray, - dof_vel: np.ndarray, - ball_pos: np.ndarray, - ball_linvel: np.ndarray, - ball_angvel: np.ndarray, - torques: np.ndarray, - terminated: np.ndarray, - ) -> np.ndarray: - del info, dof_pos, dof_vel, ball_pos, ball_angvel, torques, terminated - penalty: np.ndarray = np.sum(np.abs(ball_linvel), axis=1) - return penalty - - def _reward_pose_diff( - self, - info: dict[str, Any], - dof_pos: np.ndarray, - dof_vel: np.ndarray, - ball_pos: np.ndarray, - ball_linvel: np.ndarray, - ball_angvel: np.ndarray, - torques: np.ndarray, - terminated: np.ndarray, - ) -> np.ndarray: - del dof_vel, ball_pos, ball_linvel, ball_angvel, torques, terminated - diff = dof_pos - info["init_pose"] - penalty: np.ndarray = np.sum(np.square(diff), axis=1) - return penalty - - def _reward_torque( - self, - info: dict[str, Any], - dof_pos: np.ndarray, - dof_vel: np.ndarray, - ball_pos: np.ndarray, - ball_linvel: np.ndarray, - ball_angvel: np.ndarray, - torques: np.ndarray, - terminated: np.ndarray, - ) -> np.ndarray: - del info, dof_pos, dof_vel, ball_pos, ball_linvel, ball_angvel, terminated - penalty: np.ndarray = np.sum(np.square(torques), axis=1) - return penalty - - def _reward_work( - self, - info: dict[str, Any], - dof_pos: np.ndarray, - dof_vel: np.ndarray, - ball_pos: np.ndarray, - ball_linvel: np.ndarray, - ball_angvel: np.ndarray, - torques: np.ndarray, - terminated: np.ndarray, - ) -> np.ndarray: - del info, dof_pos, ball_pos, ball_linvel, ball_angvel, terminated - work = np.sum(torques * dof_vel, axis=1) - penalty: np.ndarray = np.square(work) - return penalty - - def _reward_drop( - self, - info: dict[str, Any], - dof_pos: np.ndarray, - dof_vel: np.ndarray, - ball_pos: np.ndarray, - ball_linvel: np.ndarray, - ball_angvel: np.ndarray, - torques: np.ndarray, - terminated: np.ndarray, - ) -> np.ndarray: - del info, dof_pos, dof_vel, ball_pos, ball_linvel, ball_angvel, torques - return np.asarray(terminated, dtype=get_global_dtype()) - - def update_state(self, state: NpEnvState) -> NpEnvState: - dof_pos = self.get_hand_dof_pos() - ball_pos = self.get_ball_pos() - ball_quat = self.get_ball_quat() - - dof_vel = (dof_pos - state.info.get("prev_dof_pos", dof_pos)) / self._cfg.ctrl_dt - ball_linvel = (ball_pos - state.info.get("prev_ball_pos", ball_pos)) / self._cfg.ctrl_dt - - prev_ball_quat = state.info.get("prev_ball_quat", ball_quat) - ball_angvel = compute_ball_angvel(ball_quat, prev_ball_quat, self._cfg.ctrl_dt) - - state.info["curr_dof_pos"] = dof_pos.copy() - state.info["curr_ball_pos"] = ball_pos.copy() - state.info["curr_ball_quat"] = ball_quat.copy() - - state.info["prev_dof_pos"] = dof_pos.copy() - state.info["prev_ball_pos"] = ball_pos.copy() - state.info["prev_ball_quat"] = ball_quat.copy() - - targets = state.info["prev_ctrl"] - torques = compute_pd_torques( - targets=targets, - dof_pos=dof_pos, - dof_vel=dof_vel, - kp=self._cfg.control_config.kp, - kd=self._cfg.control_config.kd, - ) - terminated = ball_pos[:, 2] < self._reward_cfg.reset_z_threshold - - reward = self._compute_reward( - state.info, dof_pos, dof_vel, ball_pos, ball_linvel, ball_angvel, torques, terminated - ) - obs = self._compute_obs(state.info, dof_pos, ball_pos) - return state.replace(obs=obs, reward=reward, terminated=terminated) - - def _compute_reward( - self, - info: dict[str, Any], - dof_pos: np.ndarray, - dof_vel: np.ndarray, - ball_pos: np.ndarray, - ball_linvel: np.ndarray, - ball_angvel: np.ndarray, - torques: np.ndarray, - terminated: np.ndarray, - ) -> np.ndarray: - dtype = get_global_dtype() - reward = np.zeros(self._num_envs, dtype=dtype) - step_count = info.get("steps", np.zeros(self._num_envs, dtype=np.uint32)) - should_log = self._enable_reward_log and (int(step_count[0]) % 4 == 0) - log = {} if should_log else info.get("log", {}) - - for name, scale in self._reward_cfg.scales.items(): - if scale == 0 or name not in self._reward_fns: - continue - rew = self._reward_fns[name]( - info, dof_pos, dof_vel, ball_pos, ball_linvel, ball_angvel, torques, terminated - ) - weighted_rew = rew * scale - reward += weighted_rew - if should_log: - log[f"reward/{name}"] = float(np.mean(weighted_rew)) - - if should_log: - log["reward/total"] = float(np.mean(reward)) - - info["log"] = log - return reward * self._cfg.ctrl_dt - - def _compute_obs( - self, info: dict[str, Any], dof_pos: np.ndarray, ball_pos: np.ndarray - ) -> dict[str, np.ndarray]: - dtype = get_global_dtype() - targets = info["prev_ctrl"] - dof_pos_norm = 2.0 * (dof_pos - self._dof_mid) / (self._dof_range + 1e-8) - - noise_cfg = self._cfg.noise_config - if noise_cfg.level > 0.0: - dof_pos_norm += ( - np.random.uniform(-1.0, 1.0, dof_pos_norm.shape).astype(dtype) - * noise_cfg.level - * noise_cfg.scale_joint_angle - ) - - current_obs = np.concatenate( - [dof_pos_norm, targets, ball_pos.astype(dtype)], axis=1, dtype=dtype - ) - - num_envs = dof_pos.shape[0] - obs_lag_history = info.get( - "obs_lag_history", - np.zeros( - (num_envs, self._NUM_LAG_STEPS, self._NUM_OBS_PER_STEP), - dtype=dtype, - ), - ) - obs_lag_history[:, :-1] = obs_lag_history[:, 1:] - obs_lag_history[:, -1] = current_obs - info["obs_lag_history"] = obs_lag_history - - return { - "obs": np.asarray(obs_lag_history.reshape(num_envs, -1), dtype=dtype), - } - - -RewardConfig = RewardConfigPPO -Domain_Rand = DomainRandConfig -AllegroRotationCfg = AllegroRotationPPOCfg -AllegroRotationMj = AllegroRotationPPO -# The legacy class/config above remain only as the implementation base for -# AllegroInhandRotationGrasp. The production rotation identity is Hydra-owned -# and has one generic Manager-Based runtime across all registered backends. registry.register_env_config("AllegroInhandRotation", ManagerBasedRlEnvCfg) registry.register_env("AllegroInhandRotation", make_manager_based_rl_env, sim_backend="mujoco") registry.register_env("AllegroInhandRotation", make_manager_based_rl_env, sim_backend="motrix") diff --git a/tests/conftest.py b/tests/conftest.py index 5537ab397..7e5dc819e 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -229,24 +229,6 @@ def default_g1_reward_config(): } -@pytest.fixture -def default_allegro_reward_config(): - """Default reward config for AllegroInhandRotation testing.""" - return { - "scales": { - "rotate": 1.25, - "obj_linvel": -0.3, - "pose_diff": -0.3, - "torque": -0.1, - "work": -2.0, - "drop": 0.0, - }, - "angvel_clip_min": -0.5, - "angvel_clip_max": 0.5, - "reset_z_threshold": 0.125, - } - - @pytest.fixture def default_g1_walk_flat_reward_config(): """Default reward config for G1 SAC testing.""" diff --git a/tests/envs/test_env_configs.py b/tests/envs/test_env_configs.py index 6bea250aa..48bc874a8 100644 --- a/tests/envs/test_env_configs.py +++ b/tests/envs/test_env_configs.py @@ -31,7 +31,10 @@ def _require_mujoco_runtime() -> None: def _allegro_manager_override( - backend: str = "mujoco", *, config_root: str = "ppo" + backend: str = "mujoco", + *, + config_root: str = "ppo", + task: str = "allegro_inhand", ) -> dict[str, Any]: from hydra import compose, initialize_config_dir @@ -41,7 +44,7 @@ def _allegro_manager_override( with initialize_config_dir( config_dir=str(repo_root / "conf" / config_root), version_base="1.3" ): - cfg = compose("config", overrides=[f"task=allegro_inhand/{backend}"]) + cfg = compose("config", overrides=[f"task={task}/{backend}"]) return BackendAdapter( cfg, root_dir=repo_root, algo_name=config_root ).build_task_env_cfg_override() @@ -69,7 +72,6 @@ def blocked_import(name, globals=None, locals=None, fromlist=(), level=0): from unilab.base import registry from unilab.base.backend import create_backend - from unilab.tasks.manipulation.allegro_inhand.rotation import AllegroRotationCfg from unilab.tasks.motion_tracking.g1.tracking import ( G1MotionTrackingCfg, G1MotionTrackingDeployEnvCfg, @@ -86,7 +88,6 @@ def blocked_import(name, globals=None, locals=None, fromlist=(), level=0): G1MotionTrackingCfg() G1MotionTrackingDeployEnvCfg() X2WallFlipTrackingCfg() - AllegroRotationCfg() """ ) @@ -411,7 +412,7 @@ def test_g1_box_tracking_scene_uses_sphere_hand_and_box_tracking_mesh(): assert name in scene_text -def test_allegro_rotation_registry_is_manager_only_and_grasp_owns_legacy_bridge(): +def test_allegro_rotation_and_grasp_registries_are_manager_only(): from unilab.base import registry from unilab.base.config_materialization import apply_cfg_overrides from unilab.envs import ManagerBasedRlEnvCfg @@ -422,7 +423,10 @@ def test_allegro_rotation_registry_is_manager_only_and_grasp_owns_legacy_bridge( "config_factory": "ManagerBasedRlEnvCfg", "available_backends": ["mujoco", "motrix", "drake"], } - assert metadata["AllegroInhandRotationGrasp"]["config_factory"] == ("AllegroRotationGraspCfg") + assert metadata["AllegroInhandRotationGrasp"] == { + "config_factory": "ManagerBasedRlEnvCfg", + "available_backends": ["mujoco", "motrix"], + } cfg = registry.materialize_env_config("AllegroInhandRotation") assert isinstance(cfg, ManagerBasedRlEnvCfg) @@ -442,6 +446,23 @@ def test_allegro_rotation_registry_is_manager_only_and_grasp_owns_legacy_bridge( ] assert not hasattr(cfg, "reward_config") + grasp_cfg = registry.materialize_env_config("AllegroInhandRotationGrasp") + assert isinstance(grasp_cfg, ManagerBasedRlEnvCfg) + apply_cfg_overrides( + grasp_cfg, + _allegro_manager_override(task="allegro_inhand_grasp"), + ) + assert grasp_cfg.observations["policy"].history_length == 3 + assert grasp_cfg.actions["hand"].action_scale == 0.0 + assert list(grasp_cfg.terminations) == ["dropped", "time_out", "invalid_grasp"] + assert list(grasp_cfg.metrics) == [ + "fingertips_close", + "enough_contacts", + "ball_held", + "valid", + ] + assert list(grasp_cfg.recorders) == ["grasp_cache"] + def test_allegro_manager_configured_missing_grasp_cache_fails_closed(tmp_path: Path): from unilab.managers import EventTermCfg @@ -471,88 +492,118 @@ def test_allegro_manager_configured_missing_grasp_cache_fails_closed(tmp_path: P AllegroHandBallReset(cfg, cast(Any, env)) -def test_allegro_grasp_obs_groups_spec_dims(): - """Allegro grasp task inherits the same obs group layout as rotation.""" - from unilab.tasks.manipulation.allegro_inhand.grasp_gen import AllegroRotationGrasp - - env = cast(Any, object.__new__(AllegroRotationGrasp)) - spec = env.obs_groups_spec +def _allegro_grasp_term_fixture() -> tuple[Any, Any, np.ndarray]: + from unilab.managers import TerminationTermCfg + from unilab.tasks.manipulation.allegro_inhand.grasp_gen import ( + AllegroGraspQualityTermination, + ) + from unilab.tasks.manipulation.allegro_inhand.manager_terms import ( + AllegroRotationObservation, + ) - assert spec == {"obs": 105} + num_envs = 3 + states = np.arange(num_envs * 23, dtype=np.float32).reshape(num_envs, 23) + observation = object.__new__(AllegroRotationObservation) + observation.dof_pos = states[:, :16] + observation.ball_pos = np.array( + [[0.0, 0.0, 0.2], [0.0, 0.0, 0.2], [0.0, 0.0, 0.1]], dtype=np.float32 + ) + observation.ball_quat = states[:, 19:23] + observation._last_counter = 1 + + body_pos = np.repeat(observation.ball_pos[:, None, :], 4, axis=1) + body_pos[:, :, 0] += 0.05 + body_pos[1, 0, 0] += 0.2 + contacts = np.array( + [[1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0], [1.0, 0.0, 0.0, 0.0]], + dtype=np.float32, + ) + entity = SimpleNamespace( + data=SimpleNamespace(body_link_pos_w=body_pos), + find_bodies=lambda names, preserve_order: (list(range(4)), list(names)), + ) + class _Scene(dict): + def bind_sensor_data(self, names): + assert tuple(names) == ("ff_contact", "mf_contact", "rf_contact", "th_contact") + return SimpleNamespace(dimensions=(1, 1, 1, 1), read=lambda: contacts) -def test_allegro_missing_grasp_cache_logs_local_generation_notice( - tmp_path: Path, caplog: pytest.LogCaptureFixture -) -> None: - from unilab.tasks.manipulation.allegro_inhand.rotation import ( - AllegroRotationPPOCfg, - _materialize_grasp_cache, + env = SimpleNamespace( + num_envs=num_envs, + common_step_counter=1, + scene=_Scene(robot=entity), + observation_manager=SimpleNamespace( + get_term_cfg=lambda group, name: SimpleNamespace(func=observation) + ), + action_manager=SimpleNamespace(), + reset_time_outs=np.ones(num_envs, dtype=np.bool_), + reset_terminated=np.zeros(num_envs, dtype=np.bool_), + extras={"log": {}}, ) - - missing_cache = tmp_path / "missing_allegro_cache.npy" - cfg = AllegroRotationPPOCfg( - grasp_cache_path=str(missing_cache), - gen_grasp=False, + cfg = TerminationTermCfg( + func=AllegroGraspQualityTermination, + params={ + "entity_name": "robot", + "observation_group": "policy", + "observation_term": "rotation", + "fingertip_body_names": ["ff_tip", "mf_tip", "rf_tip", "th_tip"], + "contact_sensor_names": [ + "ff_contact", + "mf_contact", + "rf_contact", + "th_contact", + ], + "max_fingertip_distance": 0.1, + "minimum_contacts": 2, + "minimum_ball_height": 0.125, + "enabled": True, + }, ) + term = AllegroGraspQualityTermination(cfg, cast(Any, env)) + env.termination_manager = SimpleNamespace(get_term_cfg=lambda name: SimpleNamespace(func=term)) + return env, term, states - with caplog.at_level( - logging.WARNING, logger="unilab.tasks.manipulation.allegro_inhand.rotation" - ): - assert _materialize_grasp_cache(cfg) is None - notice = caplog.text - - assert str(missing_cache) in notice - assert "no Hugging Face download will be attempted" in notice - assert "uv run train --algo ppo --task allegro_inhand_grasp --sim mujoco" in notice - assert "env.grasp_cache_path" in notice +def test_allegro_grasp_quality_term_matches_legacy_conditions(): + from unilab.managers import MetricsTermCfg + from unilab.tasks.manipulation.allegro_inhand.grasp_gen import AllegroGraspQualityMetric -def test_allegro_grasp_generation_skips_cache_materialization( - monkeypatch: pytest.MonkeyPatch, -) -> None: - from unilab.tasks.manipulation.allegro_inhand import rotation + env, term, _ = _allegro_grasp_term_fixture() - def fail_io(*_args: Any, **_kwargs: Any) -> None: - raise AssertionError("grasp generation must not resolve or load a rotation cache") - - monkeypatch.setattr(rotation, "resolve_grasp_cache_path", fail_io) - monkeypatch.setattr(rotation.np, "load", fail_io) - - cfg = rotation.AllegroRotationPPOCfg(gen_grasp=True) - assert rotation._materialize_grasp_cache(cfg) is None + np.testing.assert_array_equal(term(cast(Any, env)), [False, True, True]) + np.testing.assert_array_equal(term.fingertips_close, [True, False, True]) + np.testing.assert_array_equal(term.enough_contacts, [True, True, False]) + np.testing.assert_array_equal(term.ball_held, [True, True, False]) + metric = AllegroGraspQualityMetric( + MetricsTermCfg( + func=AllegroGraspQualityMetric, + params={"quality_term_name": "invalid_grasp", "condition": "valid"}, + ), + cast(Any, env), + ) + np.testing.assert_array_equal(metric(cast(Any, env)), [1.0, 0.0, 0.0]) -def test_allegro_grasp_target_raises_run_complete_without_resaving_on_close( +def test_allegro_grasp_recorder_saves_target_and_raises_run_complete( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: from unilab.base.run_control import RunComplete + from unilab.managers import RecorderTermCfg from unilab.tasks.manipulation.allegro_inhand import grasp_gen + env, term, states = _allegro_grasp_term_fixture() + term(cast(Any, env)) cache_path = tmp_path / "allegro.npy" - env = cast(Any, object.__new__(grasp_gen.AllegroRotationGrasp)) - env._cfg = grasp_gen.AllegroRotationGraspCfg( - grasp_collection_target=2, - grasp_cache_path=str(cache_path), - grasp_auto_save=True, - grasp_quality_check=False, - ) - states = np.arange(3 * 23, dtype=np.float64).reshape(3, 23) - env._saved_grasping_states = [] - env._grasp_cache_saved = False - env._grasp_target_reached_notified = False - env._state = SimpleNamespace( - truncated=np.ones((3,), dtype=bool), - terminated=np.zeros((3,), dtype=bool), - info={ - "curr_dof_pos": states[:, :16], - "curr_ball_pos": states[:, 16:19], - "curr_ball_quat": states[:, 19:23], + cfg = RecorderTermCfg( + func=grasp_gen.AllegroGraspRecorder, + params={ + "quality_term_name": "invalid_grasp", + "output_path": str(cache_path), + "collection_target": 2, + "auto_save": True, }, ) - env.get_hand_dof_pos = lambda: states[:, :16] - env.get_ball_pos = lambda: states[:, 16:19] - env.get_ball_quat = lambda: states[:, 19:23] - + recorder = grasp_gen.AllegroGraspRecorder(cfg, cast(Any, env)) save_calls: list[Path] = [] real_save = grasp_gen.np.save @@ -560,22 +611,16 @@ def save_once(path: str | Path, values: np.ndarray) -> None: save_calls.append(Path(path)) real_save(path, values) - parent_closes: list[Any] = [] monkeypatch.setattr(grasp_gen.np, "save", save_once) - monkeypatch.setattr( - grasp_gen.AllegroRotationPPO, - "close", - lambda instance: parent_closes.append(instance), - ) - with pytest.raises(RunComplete) as caught: - env._collect_successful_grasps(np.arange(3, dtype=np.int32)) + recorder.record_pre_reset(np.arange(3, dtype=np.int32)) - saved = np.load(cache_path) - np.testing.assert_array_equal(saved, states[:2].astype(np.float32)) - assert saved.dtype == np.float32 + expected = np.concatenate( + (states[:, :16], term.observation.ball_pos, states[:, 19:23]), axis=1, dtype=np.float32 + ) + np.testing.assert_array_equal(np.load(cache_path), expected[:2]) assert save_calls == [cache_path] - assert env.state.info["log"] == { + assert env.extras["log"] == { "grasp_cache/saved": 1.0, "grasp_cache/num_states": 2.0, "grasp/target_reached": 1.0, @@ -586,79 +631,55 @@ def save_once(path: str | Path, values: np.ndarray) -> None: "grasp_collection_target": 2, } - env.close() - env._stop_collection() + recorder.close() assert save_calls == [cache_path] - assert parent_closes == [env] -def test_allegro_grasp_save_failure_does_not_signal_completion( +def test_allegro_grasp_recorder_close_autosaves_and_io_failure_is_fail_closed( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: + from unilab.managers import RecorderTermCfg from unilab.tasks.manipulation.allegro_inhand import grasp_gen - env = cast(Any, object.__new__(grasp_gen.AllegroRotationGrasp)) - env._cfg = grasp_gen.AllegroRotationGraspCfg( - grasp_collection_target=1, - grasp_cache_path=str(tmp_path / "allegro.npy"), + env, term, _ = _allegro_grasp_term_fixture() + term(cast(Any, env)) + cache_path = tmp_path / "allegro.npy" + cfg = RecorderTermCfg( + func=grasp_gen.AllegroGraspRecorder, + params={ + "quality_term_name": "invalid_grasp", + "output_path": str(cache_path), + "collection_target": 3, + "auto_save": True, + }, ) - env._saved_grasping_states = [np.zeros((1, 23), dtype=np.float32)] - env._grasp_cache_saved = False - env._grasp_target_reached_notified = False - env._state = SimpleNamespace(info={}) + recorder = grasp_gen.AllegroGraspRecorder(cfg, cast(Any, env)) + env.reset_terminated[1] = True + recorder.record_pre_reset(np.array([0, 1], dtype=np.int32)) + assert recorder.total_saved_grasps == 1 + assert not cache_path.exists() + recorder.close() + assert np.load(cache_path).shape == (1, 23) + + failed_path = tmp_path / "failed.npy" + failed_cfg = RecorderTermCfg( + func=grasp_gen.AllegroGraspRecorder, + params={ + "quality_term_name": "invalid_grasp", + "output_path": str(failed_path), + "collection_target": 1, + "auto_save": True, + }, + ) + failed = grasp_gen.AllegroGraspRecorder(failed_cfg, cast(Any, env)) sentinel = OSError("disk full") monkeypatch.setattr( grasp_gen.np, "save", lambda *_args, **_kwargs: (_ for _ in ()).throw(sentinel) ) - with pytest.raises(OSError) as caught: - env._save_grasp_cache() - + failed.record_pre_reset(np.array([0], dtype=np.int32)) assert caught.value is sentinel - assert env._grasp_cache_saved is False - assert env._grasp_target_reached_notified is False - - -def test_allegro_reset_samples_materialized_cache_without_file_io( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -) -> None: - from unilab.tasks.manipulation.allegro_inhand import rotation - - cache_path = tmp_path / "allegro.npy" - cache = np.zeros((4, 23), dtype=np.float64) - cache[:, 16] = 0.1 - cache[:, 19] = 1.0 - np.save(cache_path, cache) - cfg = rotation.AllegroRotationPPOCfg(grasp_cache_path=str(cache_path)) - materialized = rotation._materialize_grasp_cache(cfg) - assert materialized is not None - - monkeypatch.setattr( - rotation, "resolve_grasp_cache_path", lambda _: (_ for _ in ()).throw(AssertionError) - ) - monkeypatch.setattr( - rotation.np, "load", lambda *_args, **_kwargs: (_ for _ in ()).throw(AssertionError) - ) - env = SimpleNamespace( - _grasp_cache=materialized, - cfg=SimpleNamespace( - domain_rand=SimpleNamespace(joint_noise=0.0, ball_vel_noise=0.0, ball_z_offset=0.0) - ), - default_angles=np.zeros(16), - _NUM_HAND_DOF=16, - _ctrl_lower=-np.ones(16), - _ctrl_upper=np.ones(16), - _init_qpos=np.zeros(23), - nv=23, - ) - - provider = rotation.AllegroRotationDomainRandomizationProvider() - for _ in range(2): - hand_qpos, ball_pos, ball_quat, qvel = provider._sample_reset_state(env, 2) - assert hand_qpos.shape == (2, 16) - assert ball_pos.shape == (2, 3) - assert ball_quat.shape == (2, 4) - assert qvel.shape == (2, 23) + assert failed.cache_saved is False def test_g1_motion_tracking_uses_combined_body_pose_query(): @@ -2136,7 +2157,6 @@ def test_env_reset_and_step( default_go2_reward_config, default_g1_reward_config, default_g1_walk_flat_reward_config, - default_allegro_reward_config, ): """Every registered env must be constructible, resetable, and steppable. @@ -2162,7 +2182,7 @@ def test_env_reset_and_step( elif env_name == "AllegroInhandRotation": env_cfg_override = _allegro_manager_override() elif env_name == "AllegroInhandRotationGrasp": - env_cfg_override = {"reward_config": default_allegro_reward_config} + env_cfg_override = _allegro_manager_override(task="allegro_inhand_grasp") env = cast( Any, @@ -2204,28 +2224,16 @@ def test_env_reset_and_step( env.close() -def test_allegro_manager_runtime_matches_legacy_rotation_transition( - default_allegro_reward_config, -): +def test_allegro_manager_runtime_transition_contract(): _require_mujoco_runtime() ensure_registries() from unilab.base import registry from unilab.envs import ManagerBasedRlEnv from unilab.tasks.manipulation.allegro_inhand.manager_terms import ( AllegroIncrementalPositionAction, - ) - from unilab.tasks.manipulation.allegro_inhand.rotation import ( - AllegroRotationPPO, - AllegroRotationPPOCfg, - RewardConfigPPO, + AllegroRotationObservation, ) - legacy_cfg = AllegroRotationPPOCfg( - reward_config=RewardConfigPPO(**default_allegro_reward_config) - ) - legacy_cfg.noise_config.level = 0.0 - legacy = AllegroRotationPPO(legacy_cfg, num_envs=2, backend_type="mujoco") - manager_override = _allegro_manager_override() manager_override["observations"]["policy"]["terms"]["rotation"]["params"]["joint_noise"] = 0.0 env = registry.make( @@ -2236,28 +2244,86 @@ def test_allegro_manager_runtime_matches_legacy_rotation_transition( ) assert isinstance(env, ManagerBasedRlEnv) try: - legacy_initial = legacy.init_state() manager_initial = env.init_state() - np.testing.assert_allclose(manager_initial.obs["obs"], legacy_initial.obs["obs"]) + history = manager_initial.obs["obs"].reshape(2, 3, 35) + np.testing.assert_array_equal(history[:, 0], history[:, 1]) + np.testing.assert_array_equal(history[:, 1], history[:, 2]) + action = env.action_manager.get_term("hand") + assert isinstance(action, AllegroIncrementalPositionAction) + target_before = action.target.copy() actions = np.full((2, 16), 0.25, dtype=np.float32) - legacy_state = legacy.step(actions) manager_state = env.step(actions) - np.testing.assert_allclose( - manager_state.obs["obs"], legacy_state.obs["obs"], rtol=1.0e-5, atol=2.0e-6 + expected_target = np.clip( + target_before + 0.25 / 24.0, + action.ctrl_lower, + action.ctrl_upper, ) + np.testing.assert_allclose(action.target, expected_target, rtol=0.0, atol=1.0e-7) + + observation = env.observation_manager.get_term_cfg("policy", "rotation").func + assert isinstance(observation, AllegroRotationObservation) + current_frame = observation(env) np.testing.assert_allclose( - manager_state.reward, legacy_state.reward, rtol=1.0e-5, atol=2.0e-8 + manager_state.obs["obs"][:, -35:], current_frame, rtol=0.0, atol=1.0e-7 ) - np.testing.assert_array_equal(manager_state.terminated, legacy_state.terminated) + assert np.isfinite(manager_state.obs["obs"]).all() + assert np.isfinite(manager_state.reward).all() + assert manager_state.terminated.dtype == np.bool_ + assert action.action_dim == 16 + assert env.obs_groups_spec == {"obs": 105} + finally: + env.close() + + +@pytest.mark.parametrize("sim_backend", ["mujoco", "motrix"]) +def test_allegro_grasp_manager_runtime_uses_zero_increment_action(sim_backend: str, tmp_path: Path): + if sim_backend == "mujoco": + _require_mujoco_runtime() + else: + pytest.importorskip("motrixsim") + ensure_registries() + from unilab.base import registry + from unilab.envs import ManagerBasedRlEnv + from unilab.tasks.manipulation.allegro_inhand.grasp_gen import ( + AllegroGraspQualityTermination, + AllegroGraspRecorder, + ) + from unilab.tasks.manipulation.allegro_inhand.manager_terms import ( + AllegroIncrementalPositionAction, + ) + override = _allegro_manager_override(sim_backend, task="allegro_inhand_grasp") + override["auto_reset"] = False + override["terminations"]["invalid_grasp"]["params"]["enabled"] = False + override["recorders"]["grasp_cache"]["params"].update( + {"output_path": str(tmp_path / f"{sim_backend}.npy"), "auto_save": False} + ) + env = registry.make( + "AllegroInhandRotationGrasp", + num_envs=2, + sim_backend=sim_backend, + env_cfg_override=override, + ) + assert isinstance(env, ManagerBasedRlEnv) + try: + initial = env.init_state() action = env.action_manager.get_term("hand") + quality = env.termination_manager.get_term_cfg("invalid_grasp").func + recorder = env.recorder_manager.get_term("grasp_cache") assert isinstance(action, AllegroIncrementalPositionAction) - assert action.action_dim == 16 - assert env.obs_groups_spec == {"obs": 105} + assert isinstance(quality, AllegroGraspQualityTermination) + assert isinstance(recorder, AllegroGraspRecorder) + target = action.target.copy() + + state = env.step(np.ones((2, 16), dtype=np.float32)) + np.testing.assert_array_equal(action.target, target) + np.testing.assert_array_equal(state.reward, np.zeros(2, dtype=state.reward.dtype)) + assert initial.obs["obs"].shape == (2, 105) + assert state.obs["obs"].shape == (2, 105) + assert quality.last_counter == env.common_step_counter finally: env.close() - legacy.close() @pytest.mark.parametrize("sim_backend", ["mujoco", "motrix"]) diff --git a/tests/scripts/test_train_scripts.py b/tests/scripts/test_train_scripts.py index 1e2ab80c9..8d2bd510a 100644 --- a/tests/scripts/test_train_scripts.py +++ b/tests/scripts/test_train_scripts.py @@ -886,16 +886,20 @@ def test_build_ppo_env_cfg_override_allegro_grasp_mujoco( assert cfg.algo.empirical_normalization is False assert cfg.algo.actor.obs_normalization is True assert cfg.algo.critic.obs_normalization is True - assert env_cfg_override["reward_config"]["scales"]["rotate"] == pytest.approx(0.0) - assert env_cfg_override["gen_grasp"] is True - assert env_cfg_override["grasp_collection_target"] == 50000 - assert env_cfg_override["grasp_quality_check"] is True - assert env_cfg_override["domain_rand"]["randomize_base_mass"] is False - assert env_cfg_override["domain_rand"]["random_com"] is False - assert env_cfg_override["domain_rand"]["randomize_gravity"] is False - assert env_cfg_override["domain_rand"]["push_robots"] is False - assert env_cfg_override["domain_rand"]["ball_vel_noise"] == pytest.approx(0.0) - assert env_cfg_override["domain_rand"]["joint_noise"] == pytest.approx(0.25) + assert env_cfg_override["rewards"]["rotate"]["weight"] == pytest.approx(0.0) + assert env_cfg_override["actions"]["hand"]["action_scale"] == pytest.approx(0.0) + reset = env_cfg_override["events"]["reset_hand_ball"]["params"] + assert reset["grasp_cache_path"] is None + assert reset["ball_velocity_noise"] == pytest.approx(0.0) + assert reset["joint_noise"] == pytest.approx(0.25) + quality = env_cfg_override["terminations"]["invalid_grasp"]["params"] + assert quality["enabled"] is True + assert quality["minimum_contacts"] == 2 + recorder = env_cfg_override["recorders"]["grasp_cache"]["params"] + assert recorder["collection_target"] == 50000 + assert recorder["auto_save"] is True + assert "reward_config" not in env_cfg_override + assert "domain_rand" not in env_cfg_override def test_build_ppo_env_cfg_override_allegro_grasp_cli_override_wins( @@ -906,17 +910,16 @@ def test_build_ppo_env_cfg_override_allegro_grasp_cli_override_wins( [ "task=allegro_inhand_grasp/mujoco", "algo.max_iterations=1", - "env.grasp_collection_target=128", - "reward.scales.rotate=0.3", + "env.recorders.grasp_cache.params.collection_target=128", + "reward.rotate.weight=0.3", ] ) env_cfg_override = mod.build_ppo_env_cfg_override(cfg) assert cfg.algo.max_iterations == 1 - assert env_cfg_override["grasp_collection_target"] == 128 - assert env_cfg_override["reward_config"]["scales"]["rotate"] == pytest.approx(0.3) - assert env_cfg_override["gen_grasp"] is True + assert env_cfg_override["recorders"]["grasp_cache"]["params"]["collection_target"] == 128 + assert env_cfg_override["rewards"]["rotate"]["weight"] == pytest.approx(0.3) def test_build_ppo_env_cfg_override_sharpa_grasp_cli_override_wins( From 48fe41164b741c4773c5a9b49dea3f08e0b314c1 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Wed, 19 Aug 2026 08:14:43 +0800 Subject: [PATCH 083/169] feat(tasks): migrate Go2W flat to Manager-Based runtime [skip ci] (#1212) --- .../task/go2w_joystick_flat/base.yaml | 298 +++++++++++++++ .../task/sac/go2w_joystick_flat/drake.yaml | 37 +- conf/ppo/task/go2w_joystick_flat/base.yaml | 298 +++++++++++++++ conf/ppo/task/go2w_joystick_flat/drake.yaml | 39 +- conf/ppo/task/go2w_joystick_flat/motrix.yaml | 35 +- conf/ppo/task/go2w_joystick_flat/mujoco.yaml | 36 +- scripts/benchmark/env/benchmark_env_step.py | 28 +- src/unilab/tasks/locomotion/go2w/joystick.py | 12 +- .../tasks/locomotion/go2w/manager_terms.py | 361 ++++++++++++++++++ .../test_env_step_config_contract.py | 12 + tests/config/test_config_system.py | 66 ++-- .../go2w/test_go2w_manager_based_flat_cfg.py | 299 +++++++++++++++ tests/envs/test_env_configs.py | 32 -- 13 files changed, 1352 insertions(+), 201 deletions(-) create mode 100644 conf/offpolicy/task/go2w_joystick_flat/base.yaml create mode 100644 conf/ppo/task/go2w_joystick_flat/base.yaml create mode 100644 src/unilab/tasks/locomotion/go2w/manager_terms.py create mode 100644 tests/envs/locomotion/go2w/test_go2w_manager_based_flat_cfg.py diff --git a/conf/offpolicy/task/go2w_joystick_flat/base.yaml b/conf/offpolicy/task/go2w_joystick_flat/base.yaml new file mode 100644 index 000000000..0f13a638c --- /dev/null +++ b/conf/offpolicy/task/go2w_joystick_flat/base.yaml @@ -0,0 +1,298 @@ +# @package _global_ +# Canonical Go2W flat Manager-Based task declaration. Backend leaves only own +# backend identity and backend-specific rendering/runtime settings. +env: + scene: + model_file: src/unilab/assets/robots/go2w/scene_flat.xml + default_keyframe_name: home + entities: + robot: + root_body_name: base_link + joint_names: + - FR_hip_joint + - FR_thigh_joint + - FR_calf_joint + - FL_hip_joint + - FL_thigh_joint + - FL_calf_joint + - RR_hip_joint + - RR_thigh_joint + - RR_calf_joint + - RL_hip_joint + - RL_thigh_joint + - RL_calf_joint + - FR_wheel_joint + - FL_wheel_joint + - RR_wheel_joint + - RL_wheel_joint + actuator_names: + - FR_hip + - FR_thigh + - FR_calf + - FL_hip + - FL_thigh + - FL_calf + - RR_hip + - RR_thigh + - RR_calf + - RL_hip + - RL_thigh + - RL_calf + - FR_wheel + - FL_wheel + - RR_wheel + - RL_wheel + body_names: [base_link] + sim_dt: 0.005 + ctrl_dt: 0.02 + max_episode_seconds: 20.0 + observations: + policy: + _target_: unilab.managers.ObservationGroupCfg + terms: + base_ang_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.builtin_sensor + params: + sensor_name: gyro + projected_gravity: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.projected_gravity_from_sensor + params: + sensor_name: upvector + leg_joint_pos: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_pos_rel + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + joint_names: ".*_(hip|thigh|calf)_joint" + leg_joint_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_vel_rel + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + joint_names: ".*_(hip|thigh|calf)_joint" + wheel_joint_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_vel_rel + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + joint_names: ".*_wheel_joint" + actions: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.last_action + params: + action_name: motor + command: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.generated_commands + params: + command_name: twist + critic: + _target_: unilab.managers.ObservationGroupCfg + terms: + base_ang_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.builtin_sensor + params: + sensor_name: gyro + projected_gravity: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.projected_gravity_from_sensor + params: + sensor_name: upvector + leg_joint_pos: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_pos_rel + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + joint_names: ".*_(hip|thigh|calf)_joint" + leg_joint_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_vel_rel + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + joint_names: ".*_(hip|thigh|calf)_joint" + wheel_joint_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_vel_rel + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + joint_names: ".*_wheel_joint" + actions: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.last_action + params: + action_name: motor + command: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.generated_commands + params: + command_name: twist + base_lin_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.builtin_sensor + params: + sensor_name: local_linvel + motor_torque: + _target_: unilab.managers.ObservationTermCfg + func: unilab.tasks.locomotion.go2w.manager_terms.motor_torque + params: + action_name: motor + actions: + motor: + _target_: unilab.tasks.locomotion.go2w.manager_terms.Go2WMixedActionCfg + entity_name: robot + actuator_names: [".*"] + leg_action_scale: 0.5 + wheel_action_scale: 10.0 + leg_kp: 50.0 + leg_kd: 1.5 + wheel_kd: 0.5 + clip_actions: 1.0 + simulate_action_latency: false + commands: + twist: + _target_: unilab.tasks.locomotion.go2w.manager_terms.Go2WVelocityCommandCfg + entity_name: robot + resampling_time_range: [20.0, 20.0] + heading_command: false + heading_control_stiffness: 0.5 + rel_standing_envs: 0.0 + rel_heading_envs: 0.0 + rel_world_envs: 0.0 + rel_forward_envs: 0.0 + init_velocity_prob: 0.0 + planar_dead_zone: 0.2 + ranges: + lin_vel_x: [0.0, 1.0] + lin_vel_y: [0.0, 0.0] + ang_vel_z: [-1.0, 1.0] + events: + reset_scene_to_default: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.reset_scene_to_default + mode: reset + reset_root_state_uniform: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.reset_root_state_uniform + mode: reset + params: + pose_range: + x: [-0.5, 0.5] + y: [-0.5, 0.5] + z: [0.0, 0.0] + roll: [0.0, 0.0] + pitch: [0.0, 0.0] + yaw: [-3.141592653589793, 3.141592653589793] + velocity_range: + x: [-0.5, 0.5] + y: [-0.5, 0.5] + z: [-0.5, 0.5] + roll: [-0.5, 0.5] + pitch: [-0.5, 0.5] + yaw: [-0.5, 0.5] + motor_gains: + _target_: unilab.managers.EventTermCfg + func: unilab.tasks.locomotion.go2w.manager_terms.randomize_motor_gains + mode: reset + params: + action_name: motor + kp_multiplier_range: [1.0, 1.0] + kd_multiplier_range: [1.0, 1.0] + terminations: + time_out: + _target_: unilab.managers.TerminationTermCfg + func: unilab.envs.mdp.time_out + time_out: true + bad_orientation: + _target_: unilab.managers.TerminationTermCfg + func: unilab.envs.mdp.bad_orientation + params: + limit_angle: 1.0471975511965976 + policy_observation_group: policy + critic_observation_group: critic + +reward: + tracking_lin_vel: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel_xy_exp + weight: 1.0 + params: + std: 0.5 + command_name: twist + tracking_ang_vel: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel_z_exp + weight: 0.75 + params: + std: 0.5 + command_name: twist + lin_vel_z: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z_l2 + weight: -5.0 + ang_vel_xy: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy_l2 + weight: -0.1 + base_height: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.base_height_l2 + weight: -100.0 + params: + target_height: 0.4 + orientation: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.mdp.flat_orientation_l2 + weight: -2.0 + action_rate: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.go2w.manager_terms.clipped_action_rate_l2 + weight: -0.005 + params: + action_name: motor + similar_to_default: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.joint_deviation_l1 + weight: -0.5 + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + joint_names: ".*_(hip|thigh|calf)_joint" + torques: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.go2w.manager_terms.motor_torque_l2 + weight: -0.0002 + params: + action_name: motor + wheel_vel: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.mdp.joint_vel_l2 + weight: 0.0 + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + joint_names: ".*_wheel_joint" + alive: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.go2w.manager_terms.constant_alive + weight: 0.5 + upward: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.go2w.manager_terms.upward_l2 + weight: 1.0 diff --git a/conf/offpolicy/task/sac/go2w_joystick_flat/drake.yaml b/conf/offpolicy/task/sac/go2w_joystick_flat/drake.yaml index fa05273e4..644465051 100644 --- a/conf/offpolicy/task/sac/go2w_joystick_flat/drake.yaml +++ b/conf/offpolicy/task/sac/go2w_joystick_flat/drake.yaml @@ -1,4 +1,8 @@ # @package _global_ +defaults: + - /task/go2w_joystick_flat/base + - _self_ + training: task_name: Go2WJoystickFlat sim_backend: drake @@ -28,36 +32,3 @@ algo: env: drake_backend_mode: batch drake_nthread: 20 - scene: - model_file: src/unilab/assets/robots/go2w/scene_flat.xml - commands: - vel_limit: - - [0.0, 0.0, -1.0] - - [1.0, 0.0, 1.0] - control_config: - action_scale: 0.5 - wheel_action_scale: 10.0 - Kp: 50.0 - Kd: 1.5 - wheel_Kd: 0.5 - domain_rand: - randomize_kp: false - randomize_kd: false - push_robots: false - -reward: - scales: - tracking_lin_vel: 1.0 - tracking_ang_vel: 0.75 - lin_vel_z: -5.0 - ang_vel_xy: -0.1 - base_height: -100.0 - orientation: -2.0 - action_rate: -0.005 - similar_to_default: -0.5 - torques: -0.0002 - wheel_vel: 0.0 - alive: 0.5 - upward: 1.0 - tracking_sigma: 0.25 - base_height_target: 0.4 diff --git a/conf/ppo/task/go2w_joystick_flat/base.yaml b/conf/ppo/task/go2w_joystick_flat/base.yaml new file mode 100644 index 000000000..0f13a638c --- /dev/null +++ b/conf/ppo/task/go2w_joystick_flat/base.yaml @@ -0,0 +1,298 @@ +# @package _global_ +# Canonical Go2W flat Manager-Based task declaration. Backend leaves only own +# backend identity and backend-specific rendering/runtime settings. +env: + scene: + model_file: src/unilab/assets/robots/go2w/scene_flat.xml + default_keyframe_name: home + entities: + robot: + root_body_name: base_link + joint_names: + - FR_hip_joint + - FR_thigh_joint + - FR_calf_joint + - FL_hip_joint + - FL_thigh_joint + - FL_calf_joint + - RR_hip_joint + - RR_thigh_joint + - RR_calf_joint + - RL_hip_joint + - RL_thigh_joint + - RL_calf_joint + - FR_wheel_joint + - FL_wheel_joint + - RR_wheel_joint + - RL_wheel_joint + actuator_names: + - FR_hip + - FR_thigh + - FR_calf + - FL_hip + - FL_thigh + - FL_calf + - RR_hip + - RR_thigh + - RR_calf + - RL_hip + - RL_thigh + - RL_calf + - FR_wheel + - FL_wheel + - RR_wheel + - RL_wheel + body_names: [base_link] + sim_dt: 0.005 + ctrl_dt: 0.02 + max_episode_seconds: 20.0 + observations: + policy: + _target_: unilab.managers.ObservationGroupCfg + terms: + base_ang_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.builtin_sensor + params: + sensor_name: gyro + projected_gravity: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.projected_gravity_from_sensor + params: + sensor_name: upvector + leg_joint_pos: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_pos_rel + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + joint_names: ".*_(hip|thigh|calf)_joint" + leg_joint_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_vel_rel + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + joint_names: ".*_(hip|thigh|calf)_joint" + wheel_joint_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_vel_rel + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + joint_names: ".*_wheel_joint" + actions: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.last_action + params: + action_name: motor + command: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.generated_commands + params: + command_name: twist + critic: + _target_: unilab.managers.ObservationGroupCfg + terms: + base_ang_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.builtin_sensor + params: + sensor_name: gyro + projected_gravity: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.projected_gravity_from_sensor + params: + sensor_name: upvector + leg_joint_pos: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_pos_rel + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + joint_names: ".*_(hip|thigh|calf)_joint" + leg_joint_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_vel_rel + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + joint_names: ".*_(hip|thigh|calf)_joint" + wheel_joint_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_vel_rel + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + joint_names: ".*_wheel_joint" + actions: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.last_action + params: + action_name: motor + command: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.generated_commands + params: + command_name: twist + base_lin_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.builtin_sensor + params: + sensor_name: local_linvel + motor_torque: + _target_: unilab.managers.ObservationTermCfg + func: unilab.tasks.locomotion.go2w.manager_terms.motor_torque + params: + action_name: motor + actions: + motor: + _target_: unilab.tasks.locomotion.go2w.manager_terms.Go2WMixedActionCfg + entity_name: robot + actuator_names: [".*"] + leg_action_scale: 0.5 + wheel_action_scale: 10.0 + leg_kp: 50.0 + leg_kd: 1.5 + wheel_kd: 0.5 + clip_actions: 1.0 + simulate_action_latency: false + commands: + twist: + _target_: unilab.tasks.locomotion.go2w.manager_terms.Go2WVelocityCommandCfg + entity_name: robot + resampling_time_range: [20.0, 20.0] + heading_command: false + heading_control_stiffness: 0.5 + rel_standing_envs: 0.0 + rel_heading_envs: 0.0 + rel_world_envs: 0.0 + rel_forward_envs: 0.0 + init_velocity_prob: 0.0 + planar_dead_zone: 0.2 + ranges: + lin_vel_x: [0.0, 1.0] + lin_vel_y: [0.0, 0.0] + ang_vel_z: [-1.0, 1.0] + events: + reset_scene_to_default: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.reset_scene_to_default + mode: reset + reset_root_state_uniform: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.reset_root_state_uniform + mode: reset + params: + pose_range: + x: [-0.5, 0.5] + y: [-0.5, 0.5] + z: [0.0, 0.0] + roll: [0.0, 0.0] + pitch: [0.0, 0.0] + yaw: [-3.141592653589793, 3.141592653589793] + velocity_range: + x: [-0.5, 0.5] + y: [-0.5, 0.5] + z: [-0.5, 0.5] + roll: [-0.5, 0.5] + pitch: [-0.5, 0.5] + yaw: [-0.5, 0.5] + motor_gains: + _target_: unilab.managers.EventTermCfg + func: unilab.tasks.locomotion.go2w.manager_terms.randomize_motor_gains + mode: reset + params: + action_name: motor + kp_multiplier_range: [1.0, 1.0] + kd_multiplier_range: [1.0, 1.0] + terminations: + time_out: + _target_: unilab.managers.TerminationTermCfg + func: unilab.envs.mdp.time_out + time_out: true + bad_orientation: + _target_: unilab.managers.TerminationTermCfg + func: unilab.envs.mdp.bad_orientation + params: + limit_angle: 1.0471975511965976 + policy_observation_group: policy + critic_observation_group: critic + +reward: + tracking_lin_vel: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel_xy_exp + weight: 1.0 + params: + std: 0.5 + command_name: twist + tracking_ang_vel: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel_z_exp + weight: 0.75 + params: + std: 0.5 + command_name: twist + lin_vel_z: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z_l2 + weight: -5.0 + ang_vel_xy: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy_l2 + weight: -0.1 + base_height: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.base_height_l2 + weight: -100.0 + params: + target_height: 0.4 + orientation: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.mdp.flat_orientation_l2 + weight: -2.0 + action_rate: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.go2w.manager_terms.clipped_action_rate_l2 + weight: -0.005 + params: + action_name: motor + similar_to_default: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.joint_deviation_l1 + weight: -0.5 + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + joint_names: ".*_(hip|thigh|calf)_joint" + torques: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.go2w.manager_terms.motor_torque_l2 + weight: -0.0002 + params: + action_name: motor + wheel_vel: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.mdp.joint_vel_l2 + weight: 0.0 + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + joint_names: ".*_wheel_joint" + alive: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.go2w.manager_terms.constant_alive + weight: 0.5 + upward: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.go2w.manager_terms.upward_l2 + weight: 1.0 diff --git a/conf/ppo/task/go2w_joystick_flat/drake.yaml b/conf/ppo/task/go2w_joystick_flat/drake.yaml index 6a5b59ef5..ddbbc1bfa 100644 --- a/conf/ppo/task/go2w_joystick_flat/drake.yaml +++ b/conf/ppo/task/go2w_joystick_flat/drake.yaml @@ -1,4 +1,8 @@ # @package _global_ +defaults: + - /task/go2w_joystick_flat/base + - _self_ + training: task_name: Go2WJoystickFlat sim_backend: drake @@ -10,6 +14,8 @@ algo: obs_groups: actor: - actor + critic: + - critic policy: init_noise_std: 0.5 algorithm: @@ -19,36 +25,3 @@ algo: env: drake_backend_mode: batch drake_nthread: 0 - scene: - model_file: src/unilab/assets/robots/go2w/scene_flat.xml - commands: - vel_limit: - - [0.0, 0.0, -1.0] - - [1.0, 0.0, 1.0] - control_config: - action_scale: 0.5 - wheel_action_scale: 10.0 - Kp: 50.0 - Kd: 1.5 - wheel_Kd: 0.5 - domain_rand: - randomize_kp: false - randomize_kd: false - push_robots: false - -reward: - scales: - tracking_lin_vel: 1.0 - tracking_ang_vel: 0.75 - lin_vel_z: -5.0 - ang_vel_xy: -0.1 - base_height: -100.0 - orientation: -2.0 - action_rate: -0.005 - similar_to_default: -0.5 - torques: -0.0002 - wheel_vel: 0.0 - alive: 0.5 - upward: 1.0 - tracking_sigma: 0.25 - base_height_target: 0.4 diff --git a/conf/ppo/task/go2w_joystick_flat/motrix.yaml b/conf/ppo/task/go2w_joystick_flat/motrix.yaml index 0e86b63b0..01af0050d 100644 --- a/conf/ppo/task/go2w_joystick_flat/motrix.yaml +++ b/conf/ppo/task/go2w_joystick_flat/motrix.yaml @@ -1,4 +1,8 @@ # @package _global_ +defaults: + - /task/go2w_joystick_flat/base + - _self_ + training: task_name: Go2WJoystickFlat sim_backend: motrix @@ -9,6 +13,8 @@ algo: obs_groups: actor: - actor + critic: + - critic policy: init_noise_std: 0.5 algorithm: @@ -16,35 +22,6 @@ algo: entropy_coef: 1.0e-3 env: render_offset_mode: zero - commands: - vel_limit: - - [0.0, 0.0, -1.0] - - [1.0, 0.0, 1.0] - control_config: - action_scale: 0.5 - wheel_action_scale: 10.0 - Kp: 50.0 - Kd: 1.5 - wheel_Kd: 0.5 - domain_rand: - randomize_kp: false - randomize_kd: false -reward: - scales: - tracking_lin_vel: 1.0 - tracking_ang_vel: 0.75 - lin_vel_z: -5.0 - ang_vel_xy: -0.1 - base_height: -100.0 - orientation: -2.0 - action_rate: -0.005 - similar_to_default: -0.5 - torques: -0.0002 - wheel_vel: 0.0 - alive: 0.5 - upward: 1.0 - tracking_sigma: 0.25 - base_height_target: 0.4 play_profile: enabled: true env: diff --git a/conf/ppo/task/go2w_joystick_flat/mujoco.yaml b/conf/ppo/task/go2w_joystick_flat/mujoco.yaml index 602b95a27..3e092334c 100644 --- a/conf/ppo/task/go2w_joystick_flat/mujoco.yaml +++ b/conf/ppo/task/go2w_joystick_flat/mujoco.yaml @@ -1,4 +1,8 @@ # @package _global_ +defaults: + - /task/go2w_joystick_flat/base + - _self_ + training: task_name: Go2WJoystickFlat sim_backend: mujoco @@ -9,41 +13,13 @@ algo: obs_groups: actor: - actor + critic: + - critic policy: init_noise_std: 0.5 algorithm: learning_rate: 3.0e-4 entropy_coef: 1.0e-3 -env: - commands: - vel_limit: - - [0.0, 0.0, -1.0] - - [1.0, 0.0, 1.0] - control_config: - action_scale: 0.5 - wheel_action_scale: 10.0 - Kp: 50.0 - Kd: 1.5 - wheel_Kd: 0.5 - domain_rand: - randomize_kp: false - randomize_kd: false -reward: - scales: - tracking_lin_vel: 1.0 - tracking_ang_vel: 0.75 - lin_vel_z: -5.0 - ang_vel_xy: -0.1 - base_height: -100.0 - orientation: -2.0 - action_rate: -0.005 - similar_to_default: -0.5 - torques: -0.0002 - wheel_vel: 0.0 - alive: 0.5 - upward: 1.0 - tracking_sigma: 0.25 - base_height_target: 0.4 play_profile: enabled: true env: diff --git a/scripts/benchmark/env/benchmark_env_step.py b/scripts/benchmark/env/benchmark_env_step.py index 632ef33fc..e1ec4e312 100644 --- a/scripts/benchmark/env/benchmark_env_step.py +++ b/scripts/benchmark/env/benchmark_env_step.py @@ -129,7 +129,7 @@ class TaskConfig: task_id: str env_name: str cfg_factory: Callable[[str, list[str]], Any] - env_cls_factory: Callable[[], type] + env_cls_factory: Callable[[], Callable[..., Any]] backends: tuple[str, ...] = ("mujoco", "motrix") aliases: tuple[str, ...] = () cfg_finalizer: Callable[[Any, str], None] | None = None @@ -312,9 +312,11 @@ def _go2_rough_env_cls() -> type: def _go2w_cfg(backend: str, config_overrides: list[str]) -> Any: - from unilab.tasks.locomotion.go2w.joystick import Go2WJoystickCfg + from unilab.envs import ManagerBasedRlEnvCfg - return _ppo_owner_yaml_cfg("go2w_joystick_flat", backend, Go2WJoystickCfg, config_overrides) + return _ppo_owner_yaml_cfg( + "go2w_joystick_flat", backend, ManagerBasedRlEnvCfg, config_overrides + ) def _go2w_rough_cfg(backend: str, config_overrides: list[str]) -> Any: @@ -325,10 +327,10 @@ def _go2w_rough_cfg(backend: str, config_overrides: list[str]) -> Any: ) -def _go2w_env_cls() -> type: - from unilab.tasks.locomotion.go2w.joystick import Go2WJoystickEnv +def _go2w_env_cls() -> Callable[..., Any]: + from unilab.envs import make_manager_based_rl_env - return Go2WJoystickEnv + return make_manager_based_rl_env def _go2w_rough_env_cls() -> type: @@ -502,7 +504,7 @@ def _sharpa_inhand_env_cls() -> type: } # Default benchmark parameters -DEFAULT_NUM_ENVS = 2048 +DEFAULT_NUM_ENVS = 4096 DEFAULT_NUM_STEPS = 20 DEFAULT_WARMUP_STEPS = 5 @@ -708,15 +710,21 @@ def _run_single(extra_args: list[str]) -> dict[str, Any]: env_cls = task_config.env_cls_factory() env = env_cls(cfg, num_envs=num_envs, backend_type=sim_backend) - nu = env._backend.num_actuators # type: ignore[reportAttributeAccessIssue] + action_shape = env.action_space.shape + if action_shape is None or len(action_shape) != 1: + raise ValueError( + f"Benchmark task {task_config.env_name!r} requires a flat action space, " + f"got {action_shape}" + ) + action_dim = int(action_shape[0]) env.init_state() for _ in range(warmup_steps): - actions = np.random.uniform(-1, 1, size=(num_envs, nu)).astype(np.float32) + actions = np.random.uniform(-1, 1, size=(num_envs, action_dim)).astype(np.float32) env.step(actions) for _ in range(num_steps): - actions = np.random.uniform(-1, 1, size=(num_envs, nu)).astype(np.float32) + actions = np.random.uniform(-1, 1, size=(num_envs, action_dim)).astype(np.float32) state = env.step(actions) timing = state.info.get("timing", {}) for k, v in timing.items(): diff --git a/src/unilab/tasks/locomotion/go2w/joystick.py b/src/unilab/tasks/locomotion/go2w/joystick.py index c27266975..9a61a063f 100644 --- a/src/unilab/tasks/locomotion/go2w/joystick.py +++ b/src/unilab/tasks/locomotion/go2w/joystick.py @@ -17,6 +17,7 @@ zero_actions, ) from unilab.dtype_config import get_global_dtype +from unilab.envs import ManagerBasedRlEnvCfg, make_manager_based_rl_env from unilab.tasks.locomotion.common import rewards from unilab.tasks.locomotion.common.commands import ( Commands, @@ -85,7 +86,6 @@ class JoystickSensor: gravity = "upvector" -@registry.envcfg("Go2WJoystickFlat") @dataclass class Go2WJoystickCfg(Go2WBaseCfg): scene: SceneCfg = field( # pyright: ignore[reportIncompatibleVariableOverride] @@ -226,8 +226,6 @@ def _sample_commands(self, env: Any, num_reset: int) -> np.ndarray: return commands -@registry.env("Go2WJoystickFlat", sim_backend="mujoco") -@registry.env("Go2WJoystickFlat", sim_backend="drake") class Go2WJoystickEnv(Go2WBaseEnv): _cfg: Go2WJoystickCfg # pyright: ignore[reportIncompatibleVariableOverride] @@ -644,4 +642,10 @@ def _reward_joint_mirror(self, ctx: RewardContext) -> np.ndarray: return np.asarray(mirror, dtype=get_global_dtype()) -registry.register_env("Go2WJoystickFlat", Go2WJoystickEnv, sim_backend="motrix") +# Go2WJoystickCfg and Go2WJoystickEnv remain solely as the rough-task bridge. +# The flat production identity is Hydra-owned and uses the generic Manager-Based +# factory; the bridge is deleted with the rough-task migration. +registry.register_env_config("Go2WJoystickFlat", ManagerBasedRlEnvCfg) +registry.register_env("Go2WJoystickFlat", make_manager_based_rl_env, sim_backend="mujoco") +registry.register_env("Go2WJoystickFlat", make_manager_based_rl_env, sim_backend="motrix") +registry.register_env("Go2WJoystickFlat", make_manager_based_rl_env, sim_backend="drake") diff --git a/src/unilab/tasks/locomotion/go2w/manager_terms.py b/src/unilab/tasks/locomotion/go2w/manager_terms.py new file mode 100644 index 000000000..a58d64507 --- /dev/null +++ b/src/unilab/tasks/locomotion/go2w/manager_terms.py @@ -0,0 +1,361 @@ +"""Manager-Based terms owned by the Go2W flat task.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from numbers import Real +from typing import TYPE_CHECKING, Any, cast + +import numpy as np + +from unilab.dtype_config import get_global_dtype +from unilab.envs.mdp.commands.velocity_command import ( + UniformVelocityCommand, + UniformVelocityCommandCfg, +) +from unilab.managers import ActionTerm, ActionTermCfg +from unilab.managers.scene_entity_config import SceneEntityCfg +from unilab.tasks.locomotion.go2w.base import ( + NUM_GO2W_ACTIONS, + NUM_LEG_ACTIONS, + NUM_WHEEL_ACTIONS, + compute_go2w_motor_ctrl, +) + +if TYPE_CHECKING: + from unilab.base.entity import Entity + from unilab.managers._types import ManagerBasedRlEnv + + +_HIP_INDICES = np.asarray([0, 3, 6, 9], dtype=np.intp) +_DEFAULT_ASSET_CFG = SceneEntityCfg("robot") + + +def _real( + value: Any, + *, + label: str, + minimum: float | None = None, + strict_minimum: bool = False, +) -> float: + if isinstance(value, (bool, np.bool_)) or not isinstance(value, Real): + raise TypeError(f"{label} must be a real number, got {type(value).__name__}") + result = float(value) + if not math.isfinite(result): + raise ValueError(f"{label} must be finite") + if minimum is not None and (result <= minimum if strict_minimum else result < minimum): + relation = "greater than" if strict_minimum else "at least" + raise ValueError(f"{label} must be {relation} {minimum}") + return result + + +def _range(value: Any, *, label: str) -> tuple[float, float]: + if not isinstance(value, (tuple, list)) or len(value) != 2: + raise TypeError(f"{label} must be a two-value range") + lower = _real(value[0], label=f"{label} lower", minimum=0.0) + upper = _real(value[1], label=f"{label} upper", minimum=0.0) + if lower > upper: + raise ValueError(f"{label} lower {lower} exceeds upper {upper}") + return lower, upper + + +@dataclass(kw_only=True) +class Go2WMixedActionCfg(ActionTermCfg): + """Configure the Go2W leg-position and wheel-velocity motor action.""" + + actuator_names: tuple[str, ...] | list[str] + leg_action_scale: float = 0.25 + hip_action_scale: float | None = None + wheel_action_scale: float = 10.0 + leg_kp: float = 35.0 + leg_kd: float = 0.5 + wheel_kd: float = 0.5 + clip_actions: float = 1.0 + simulate_action_latency: bool = False + + def build(self, env: ManagerBasedRlEnv) -> Go2WMixedAction: + return Go2WMixedAction(self, env) + + +class Go2WMixedAction(ActionTerm): + """Convert one community action term into Go2W motor torques per substep.""" + + cfg: Go2WMixedActionCfg + _entity: Entity + + def __init__(self, cfg: Go2WMixedActionCfg, env: ManagerBasedRlEnv): + self._validate_cfg(cfg) + super().__init__(cfg=cfg, env=env) + actuator_ids, actuator_names = self._entity.find_actuators(cfg.actuator_names) + joint_ids, joint_names = self._entity.find_joints_by_actuator_names(cfg.actuator_names) + if len(actuator_ids) != NUM_GO2W_ACTIONS or len(joint_ids) != NUM_GO2W_ACTIONS: + raise ValueError( + "Go2WMixedAction requires exactly " + f"{NUM_GO2W_ACTIONS} actuators and target joints; received " + f"actuators={actuator_names}, joints={joint_names}" + ) + self._actuator_ids = np.asarray(actuator_ids, dtype=np.intp) + self._joint_ids = np.asarray(joint_ids, dtype=np.intp) + self._actuator_ids.setflags(write=False) + self._joint_ids.setflags(write=False) + + dtype = get_global_dtype() + shape = (self.num_envs, NUM_GO2W_ACTIONS) + self._raw_action = np.zeros(shape, dtype=dtype) + self._previous_raw_action = np.zeros_like(self._raw_action) + self._processed_action = np.zeros_like(self._raw_action) + self._motor_torque = np.zeros_like(self._raw_action) + + self._leg_action_scale = np.full( + (NUM_LEG_ACTIONS,), float(cfg.leg_action_scale), dtype=dtype + ) + if cfg.hip_action_scale is not None: + self._leg_action_scale[_HIP_INDICES] = float(cfg.hip_action_scale) + self._base_leg_kp = np.full((NUM_LEG_ACTIONS,), float(cfg.leg_kp), dtype=dtype) + self._base_leg_kd = np.full((NUM_LEG_ACTIONS,), float(cfg.leg_kd), dtype=dtype) + self._leg_kp = np.broadcast_to(self._base_leg_kp, (self.num_envs, NUM_LEG_ACTIONS)).copy() + self._leg_kd = np.broadcast_to(self._base_leg_kd, (self.num_envs, NUM_LEG_ACTIONS)).copy() + self._wheel_kd = np.full( + (self.num_envs, NUM_WHEEL_ACTIONS), float(cfg.wheel_kd), dtype=dtype + ) + + ctrl_range = self._entity.data.actuator_ctrl_range[self._actuator_ids] + expected_range_shape = (NUM_GO2W_ACTIONS, 2) + if ctrl_range.shape != expected_range_shape: + raise ValueError( + "Go2WMixedAction actuator control range must have shape " + f"{expected_range_shape}, got {ctrl_range.shape}" + ) + self._ctrl_lower = np.asarray(ctrl_range[:, 0], dtype=dtype) + self._ctrl_upper = np.asarray(ctrl_range[:, 1], dtype=dtype) + + @staticmethod + def _validate_cfg(cfg: Go2WMixedActionCfg) -> None: + if not isinstance(cfg.simulate_action_latency, bool): + raise TypeError("Go2WMixedActionCfg simulate_action_latency must be bool") + for name, value in ( + ("leg_action_scale", cfg.leg_action_scale), + ("wheel_action_scale", cfg.wheel_action_scale), + ("leg_kp", cfg.leg_kp), + ("leg_kd", cfg.leg_kd), + ("wheel_kd", cfg.wheel_kd), + ("clip_actions", cfg.clip_actions), + ): + _real( + value, + label=f"Go2WMixedActionCfg {name}", + minimum=0.0, + strict_minimum=name == "clip_actions", + ) + if cfg.hip_action_scale is not None: + _real( + cfg.hip_action_scale, + label="Go2WMixedActionCfg hip_action_scale", + minimum=0.0, + ) + + @property + def action_dim(self) -> int: + return NUM_GO2W_ACTIONS + + @property + def raw_action(self) -> np.ndarray: + """Clipped policy action, matching the legacy observable action buffer.""" + return self._raw_action + + @property + def previous_raw_action(self) -> np.ndarray: + return self._previous_raw_action + + @property + def processed_action(self) -> np.ndarray: + return self._processed_action + + @property + def motor_torque(self) -> np.ndarray: + return self._motor_torque + + @property + def leg_kp(self) -> np.ndarray: + return self._leg_kp + + @property + def leg_kd(self) -> np.ndarray: + return self._leg_kd + + def process_actions(self, actions: np.ndarray) -> None: + if not isinstance(actions, np.ndarray): + raise TypeError(f"Go2WMixedAction expected np.ndarray, got {type(actions).__name__}") + if actions.shape != self._raw_action.shape: + raise ValueError( + f"Go2WMixedAction expected action shape {self._raw_action.shape}, " + f"got {actions.shape}" + ) + if not np.isfinite(actions).all(): + raise ValueError("Go2WMixedAction received NaN or Inf actions") + + self._previous_raw_action[:] = self._raw_action + np.clip( + actions, + -float(self.cfg.clip_actions), + float(self.cfg.clip_actions), + out=self._raw_action, + ) + executed = ( + self._previous_raw_action if self.cfg.simulate_action_latency else self._raw_action + ) + np.multiply( + executed[:, :NUM_LEG_ACTIONS], + self._leg_action_scale, + out=self._processed_action[:, :NUM_LEG_ACTIONS], + ) + self._processed_action[:, :NUM_LEG_ACTIONS] += self._entity.data.default_joint_pos[ + :, self._joint_ids[:NUM_LEG_ACTIONS] + ] + np.multiply( + executed[:, NUM_LEG_ACTIONS:], + float(self.cfg.wheel_action_scale), + out=self._processed_action[:, NUM_LEG_ACTIONS:], + ) + + def apply_actions(self) -> None: + joint_pos = self._entity.data.joint_pos[:, self._joint_ids] + joint_vel = self._entity.data.joint_vel[:, self._joint_ids] + compute_go2w_motor_ctrl( + self._processed_action, + joint_pos, + joint_vel, + self._leg_kp, + self._leg_kd, + self._wheel_kd, + self._ctrl_lower, + self._ctrl_upper, + self._motor_torque, + ) + self._entity.data.write_ctrl(self._motor_torque, actuator_ids=self._actuator_ids) + + def set_motor_gain_multipliers( + self, + env_ids: np.ndarray, + kp_multiplier: np.ndarray, + kd_multiplier: np.ndarray, + ) -> None: + expected = (len(env_ids), 1) + if kp_multiplier.shape != expected or kd_multiplier.shape != expected: + raise ValueError( + "Go2WMixedAction motor gain multipliers must have shape " + f"{expected}, got kp={kp_multiplier.shape}, kd={kd_multiplier.shape}" + ) + self._leg_kp[env_ids] = self._base_leg_kp * kp_multiplier + self._leg_kd[env_ids] = self._base_leg_kd * kd_multiplier + + def reset(self, env_ids: np.ndarray | slice | None = None) -> None: + if env_ids is None: + env_ids = slice(None) + self._raw_action[env_ids] = 0.0 + self._previous_raw_action[env_ids] = 0.0 + self._processed_action[env_ids] = 0.0 + self._motor_torque[env_ids] = 0.0 + + +@dataclass(kw_only=True) +class Go2WVelocityCommandCfg(UniformVelocityCommandCfg): + """Velocity command with the legacy Go2W planar dead zone.""" + + planar_dead_zone: float = 0.2 + + def build(self, env: ManagerBasedRlEnv) -> Go2WVelocityCommand: + return Go2WVelocityCommand(self, env) + + +class Go2WVelocityCommand(UniformVelocityCommand): + cfg: Go2WVelocityCommandCfg # pyright: ignore[reportIncompatibleVariableOverride] + + def __init__(self, cfg: Go2WVelocityCommandCfg, env: ManagerBasedRlEnv): + self._planar_dead_zone = _real( + cfg.planar_dead_zone, + label="Go2WVelocityCommandCfg planar_dead_zone", + minimum=0.0, + ) + super().__init__(cfg, env) + + def _resample_command(self, env_ids: np.ndarray) -> None: + super()._resample_command(env_ids) + planar = self.vel_command_b[env_ids, :2] + moving = np.linalg.norm(planar, axis=1) > self._planar_dead_zone + self.vel_command_b[env_ids, :2] = planar * moving[:, None] + + +def _action(env: ManagerBasedRlEnv, action_name: str) -> Go2WMixedAction: + if not isinstance(action_name, str) or not action_name: + raise ValueError("Go2W manager term action_name must be a non-empty string") + try: + term = env.action_manager.get_term(action_name) + except KeyError as exc: + raise KeyError(f"Go2W action term '{action_name}' is unavailable") from exc + if not isinstance(term, Go2WMixedAction): + raise TypeError( + f"Go2W action term '{action_name}' must be Go2WMixedAction, got {type(term).__name__}" + ) + return term + + +def randomize_motor_gains( + env: ManagerBasedRlEnv, + env_ids: np.ndarray | None, + action_name: str, + kp_multiplier_range: tuple[float, float] | list[float], + kd_multiplier_range: tuple[float, float] | list[float], +) -> None: + """Sample owner-level motor gains without mutating backend actuator models.""" + ids = ( + np.arange(env.num_envs, dtype=np.int32) + if env_ids is None + else np.asarray(env_ids, dtype=np.int32) + ) + kp_range = _range(kp_multiplier_range, label="randomize_motor_gains kp_multiplier_range") + kd_range = _range(kd_multiplier_range, label="randomize_motor_gains kd_multiplier_range") + shape = (len(ids), 1) + kp = env.rng.uniform(*kp_range, size=shape).astype(get_global_dtype(), copy=False) + kd = env.rng.uniform(*kd_range, size=shape).astype(get_global_dtype(), copy=False) + _action(env, action_name).set_motor_gain_multipliers(ids, kp, kd) + + +def motor_torque(env: ManagerBasedRlEnv, action_name: str) -> np.ndarray: + return _action(env, action_name).motor_torque + + +def motor_torque_l2(env: ManagerBasedRlEnv, action_name: str) -> np.ndarray: + return np.sum(np.square(motor_torque(env, action_name)), axis=1) + + +def clipped_action_rate_l2(env: ManagerBasedRlEnv, action_name: str) -> np.ndarray: + action = _action(env, action_name) + return np.sum(np.square(action.raw_action - action.previous_raw_action), axis=1) + + +def upward_l2( + env: ManagerBasedRlEnv, + asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG, +) -> np.ndarray: + asset = cast("Entity", env.scene[asset_cfg.name]) + return np.square(1.0 - asset.data.projected_gravity_b[:, 2]) + + +def constant_alive(env: ManagerBasedRlEnv) -> np.ndarray: + return np.ones((env.num_envs,), dtype=get_global_dtype()) + + +__all__ = [ + "Go2WMixedAction", + "Go2WMixedActionCfg", + "Go2WVelocityCommand", + "Go2WVelocityCommandCfg", + "clipped_action_rate_l2", + "constant_alive", + "motor_torque", + "motor_torque_l2", + "randomize_motor_gains", + "upward_l2", +] diff --git a/tests/benchmark/test_env_step_config_contract.py b/tests/benchmark/test_env_step_config_contract.py index eef8c8f34..b0d33e862 100644 --- a/tests/benchmark/test_env_step_config_contract.py +++ b/tests/benchmark/test_env_step_config_contract.py @@ -7,6 +7,18 @@ from omegaconf import OmegaConf from scripts.benchmark.env import benchmark_env_step as bench +from unilab.envs import ManagerBasedRlEnvCfg, make_manager_based_rl_env + + +def test_go2w_flat_benchmark_uses_production_manager_owner() -> None: + cfg = bench.TASK_CONFIGS["go2w"].build_cfg("mujoco") + + assert isinstance(cfg, ManagerBasedRlEnvCfg) + assert list(cfg.actions) == ["motor"] + assert cfg.critic_observation_group == "critic" + assert bench.TASK_CONFIGS["go2w"].env_cls_factory() is make_manager_based_rl_env + assert bench.DEFAULT_NUM_ENVS == 4096 + def test_go2w_rough_cfg_matches_ppo_owner_yaml() -> None: cfg = bench.TASK_CONFIGS["go2w_rough"].build_cfg("mujoco") diff --git a/tests/config/test_config_system.py b/tests/config/test_config_system.py index b94883e78..1fa7f78d8 100644 --- a/tests/config/test_config_system.py +++ b/tests/config/test_config_system.py @@ -428,19 +428,24 @@ def test_ppo_go2w_mujoco_uses_motor_owner_dr_path(): assert cfg.training.task_name == "Go2WJoystickFlat" assert cfg.training.sim_backend == "mujoco" - assert cfg.env.commands.vel_limit == [[0.0, 0.0, -1.0], [1.0, 0.0, 1.0]] - assert cfg.env.domain_rand.randomize_kp is False - assert cfg.env.domain_rand.randomize_kd is False - assert cfg.env.control_config.action_scale == pytest.approx(0.5) - assert cfg.env.control_config.Kp == pytest.approx(50.0) - assert cfg.env.control_config.Kd == pytest.approx(1.5) - assert cfg.env.control_config.wheel_action_scale == pytest.approx(10.0) - assert cfg.env.control_config.wheel_Kd == pytest.approx(0.5) - assert cfg.reward.scales.tracking_ang_vel == pytest.approx(0.75) - assert cfg.reward.scales.orientation == pytest.approx(-2.0) - assert cfg.reward.scales.upward == pytest.approx(1.0) - assert cfg.reward.base_height_target == pytest.approx(0.4) - assert cfg.reward.scales.torques < 0.0 + command = cfg.env.commands.twist + assert command.ranges.lin_vel_x == [0.0, 1.0] + assert command.ranges.lin_vel_y == [0.0, 0.0] + assert command.ranges.ang_vel_z == [-1.0, 1.0] + action = cfg.env.actions.motor + assert action.leg_action_scale == pytest.approx(0.5) + assert action.leg_kp == pytest.approx(50.0) + assert action.leg_kd == pytest.approx(1.5) + assert action.wheel_action_scale == pytest.approx(10.0) + assert action.wheel_kd == pytest.approx(0.5) + gains = cfg.env.events.motor_gains.params + assert gains.kp_multiplier_range == [1.0, 1.0] + assert gains.kd_multiplier_range == [1.0, 1.0] + assert cfg.reward.tracking_ang_vel.weight == pytest.approx(0.75) + assert cfg.reward.orientation.weight == pytest.approx(-2.0) + assert cfg.reward.upward.weight == pytest.approx(1.0) + assert cfg.reward.base_height.params.target_height == pytest.approx(0.4) + assert cfg.reward.torques.weight < 0.0 def test_ppo_go2w_motrix_uses_motor_owner_dr_path(): @@ -449,18 +454,20 @@ def test_ppo_go2w_motrix_uses_motor_owner_dr_path(): assert cfg.training.task_name == "Go2WJoystickFlat" assert cfg.training.sim_backend == "motrix" assert cfg.env.render_offset_mode == "zero" - assert cfg.env.commands.vel_limit == [[0.0, 0.0, -1.0], [1.0, 0.0, 1.0]] - assert cfg.env.domain_rand.randomize_kp is False - assert cfg.env.domain_rand.randomize_kd is False - assert cfg.env.control_config.action_scale == pytest.approx(0.5) - assert cfg.env.control_config.Kp == pytest.approx(50.0) - assert cfg.env.control_config.Kd == pytest.approx(1.5) - assert cfg.env.control_config.wheel_action_scale == pytest.approx(10.0) - assert cfg.env.control_config.wheel_Kd == pytest.approx(0.5) - assert cfg.reward.scales.tracking_ang_vel == pytest.approx(0.75) - assert cfg.reward.scales.orientation == pytest.approx(-2.0) - assert cfg.reward.scales.upward == pytest.approx(1.0) - assert cfg.reward.scales.torques < 0.0 + command = cfg.env.commands.twist + assert command.ranges.lin_vel_x == [0.0, 1.0] + assert command.ranges.lin_vel_y == [0.0, 0.0] + assert command.ranges.ang_vel_z == [-1.0, 1.0] + action = cfg.env.actions.motor + assert action.leg_action_scale == pytest.approx(0.5) + assert action.leg_kp == pytest.approx(50.0) + assert action.leg_kd == pytest.approx(1.5) + assert action.wheel_action_scale == pytest.approx(10.0) + assert action.wheel_kd == pytest.approx(0.5) + assert cfg.reward.tracking_ang_vel.weight == pytest.approx(0.75) + assert cfg.reward.orientation.weight == pytest.approx(-2.0) + assert cfg.reward.upward.weight == pytest.approx(1.0) + assert cfg.reward.torques.weight < 0.0 def test_ppo_go2w_motrix_uses_motor_owner_scene_path(): @@ -468,11 +475,10 @@ def test_ppo_go2w_motrix_uses_motor_owner_scene_path(): assert cfg.training.task_name == "Go2WJoystickFlat" assert cfg.training.sim_backend == "motrix" - assert "model_file" not in cfg.env - assert cfg.env.domain_rand.randomize_kp is False - assert cfg.env.domain_rand.randomize_kd is False - assert cfg.env.control_config.wheel_action_scale == pytest.approx(10.0) - assert cfg.reward.scales.torques < 0.0 + assert str(cfg.env.scene.model_file).endswith("src/unilab/assets/robots/go2w/scene_flat.xml") + assert cfg.env.scene.default_keyframe_name == "home" + assert cfg.env.actions.motor.wheel_action_scale == pytest.approx(10.0) + assert cfg.reward.torques.weight < 0.0 def test_ppo_go2w_rough_mujoco_uses_terrain_generator(): diff --git a/tests/envs/locomotion/go2w/test_go2w_manager_based_flat_cfg.py b/tests/envs/locomotion/go2w/test_go2w_manager_based_flat_cfg.py new file mode 100644 index 000000000..9e33af9ff --- /dev/null +++ b/tests/envs/locomotion/go2w/test_go2w_manager_based_flat_cfg.py @@ -0,0 +1,299 @@ +"""Hydra-owned production contract for the Go2W flat Manager-Based task.""" + +from __future__ import annotations + +import inspect +from collections.abc import Mapping, Sequence +from dataclasses import fields, is_dataclass +from pathlib import Path +from typing import Any + +import numpy as np +import pytest +from hydra import compose, initialize_config_dir +from hydra.core.global_hydra import GlobalHydra +from omegaconf import DictConfig, OmegaConf + +from unilab.base import registry +from unilab.base.config_materialization import apply_cfg_overrides +from unilab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg, make_manager_based_rl_env +from unilab.tasks.locomotion.go2w import manager_terms +from unilab.tasks.locomotion.go2w.manager_terms import ( + Go2WMixedAction, + Go2WMixedActionCfg, + Go2WVelocityCommandCfg, +) +from unilab.training.backend_adapter import BackendAdapter + +ROOT_DIR = Path(__file__).parents[4] +CONF_DIR = ROOT_DIR / "conf" + +_LEG_JOINT_NAMES = tuple( + f"{leg}_{joint}_joint" for leg in ("FR", "FL", "RR", "RL") for joint in ("hip", "thigh", "calf") +) +_WHEEL_JOINT_NAMES = tuple(f"{leg}_wheel_joint" for leg in ("FR", "FL", "RR", "RL")) +_JOINT_NAMES = (*_LEG_JOINT_NAMES, *_WHEEL_JOINT_NAMES) +_ACTUATOR_NAMES = tuple(name.removesuffix("_joint") for name in _JOINT_NAMES) +_HOME_JOINT_POS = np.asarray([0.0, 0.8, -1.5] * 4 + [0.0] * 4, dtype=np.float32) + + +def _compose(backend: str) -> DictConfig: + GlobalHydra.instance().clear() + with initialize_config_dir(config_dir=str(CONF_DIR / "ppo"), version_base="1.3"): + return compose("config", overrides=[f"task=go2w_joystick_flat/{backend}"]) + + +def _materialize( + backend: str, +) -> tuple[DictConfig, ManagerBasedRlEnvCfg, dict[str, Any]]: + hydra_cfg = _compose(backend) + env_override = BackendAdapter(hydra_cfg, root_dir=ROOT_DIR).build_task_env_cfg_override() + env_cfg = registry.materialize_env_config("Go2WJoystickFlat") + assert isinstance(env_cfg, ManagerBasedRlEnvCfg) + apply_cfg_overrides(env_cfg, env_override) + env_cfg.validate() + return hydra_cfg, env_cfg, env_override + + +def _assert_no_omegaconf(value: Any) -> None: + assert not OmegaConf.is_config(value) + if is_dataclass(value) and not isinstance(value, type): + for item in fields(value): + _assert_no_omegaconf(getattr(value, item.name)) + elif isinstance(value, Mapping): + for key, item in value.items(): + _assert_no_omegaconf(key) + _assert_no_omegaconf(item) + elif isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + for item in value: + _assert_no_omegaconf(item) + + +@pytest.mark.parametrize("backend", ["mujoco", "motrix", "drake"]) +def test_go2w_flat_owner_materializes_complete_plain_manager_cfg(backend: str) -> None: + registry.ensure_registries() + hydra_cfg, env_cfg, _ = _materialize(backend) + + assert hydra_cfg.training.task_name == "Go2WJoystickFlat" + assert hydra_cfg.training.sim_backend == backend + assert list(hydra_cfg.algo.obs_groups.actor) == ["actor"] + assert list(hydra_cfg.algo.obs_groups.critic) == ["critic"] + assert env_cfg.sim_dt == pytest.approx(0.005) + assert env_cfg.ctrl_dt == pytest.approx(0.02) + assert env_cfg.max_episode_seconds == pytest.approx(20.0) + assert env_cfg.policy_observation_group == "policy" + assert env_cfg.critic_observation_group == "critic" + + assert env_cfg.scene is not None + assert env_cfg.scene.model_file.endswith("robots/go2w/scene_flat.xml") + assert env_cfg.scene.default_keyframe_name == "home" + robot = env_cfg.scene.entities["robot"] + assert robot.root_body_name == "base_link" + assert tuple(robot.joint_names or ()) == _JOINT_NAMES + assert tuple(robot.actuator_names or ()) == _ACTUATOR_NAMES + assert robot.body_names == ["base_link"] + + policy_terms = [ + "base_ang_vel", + "projected_gravity", + "leg_joint_pos", + "leg_joint_vel", + "wheel_joint_vel", + "actions", + "command", + ] + assert list(env_cfg.observations) == ["policy", "critic"] + assert list(env_cfg.observations["policy"].terms) == policy_terms + assert list(env_cfg.observations["critic"].terms) == [ + *policy_terms, + "base_lin_vel", + "motor_torque", + ] + + assert list(env_cfg.actions) == ["motor"] + action = env_cfg.actions["motor"] + assert isinstance(action, Go2WMixedActionCfg) + assert action.leg_action_scale == pytest.approx(0.5) + assert action.wheel_action_scale == pytest.approx(10.0) + assert action.leg_kp == pytest.approx(50.0) + assert action.leg_kd == pytest.approx(1.5) + assert action.wheel_kd == pytest.approx(0.5) + + command = env_cfg.commands["twist"] + assert isinstance(command, Go2WVelocityCommandCfg) + assert command.resampling_time_range == [20.0, 20.0] + assert command.planar_dead_zone == pytest.approx(0.2) + assert tuple(command.ranges.lin_vel_x) == (0.0, 1.0) + assert tuple(command.ranges.lin_vel_y) == (0.0, 0.0) + assert tuple(command.ranges.ang_vel_z) == (-1.0, 1.0) + + assert list(env_cfg.events) == [ + "reset_scene_to_default", + "reset_root_state_uniform", + "motor_gains", + ] + gains = env_cfg.events["motor_gains"] + assert gains.func is manager_terms.randomize_motor_gains + assert gains.params["kp_multiplier_range"] == [1.0, 1.0] + assert gains.params["kd_multiplier_range"] == [1.0, 1.0] + assert list(env_cfg.terminations) == ["time_out", "bad_orientation"] + assert {name: term.weight for name, term in env_cfg.rewards.items()} == { + "tracking_lin_vel": 1.0, + "tracking_ang_vel": 0.75, + "lin_vel_z": -5.0, + "ang_vel_xy": -0.1, + "base_height": -100.0, + "orientation": -2.0, + "action_rate": -0.005, + "similar_to_default": -0.5, + "torques": -0.0002, + "wheel_vel": 0.0, + "alive": 0.5, + "upward": 1.0, + } + + for manager_name in ("observations", "events", "rewards", "terminations"): + for term in getattr(env_cfg, manager_name).values(): + if term is None: + continue + terms = term.terms.values() if manager_name == "observations" else (term,) + for nested in terms: + if nested is None: + continue + module = nested.func.__module__ + assert ".backend." not in module + assert not any(name in module for name in (".mujoco", ".motrix", ".drake")) + + _assert_no_omegaconf(env_cfg) + + +def test_go2w_sac_drake_owner_uses_the_same_manager_contract() -> None: + GlobalHydra.instance().clear() + with initialize_config_dir(config_dir=str(CONF_DIR / "offpolicy"), version_base="1.3"): + hydra_cfg = compose( + "config", + overrides=["algo=sac", "task=sac/go2w_joystick_flat/drake"], + ) + env_override = BackendAdapter(hydra_cfg, root_dir=ROOT_DIR).build_task_env_cfg_override() + env_cfg = registry.materialize_env_config("Go2WJoystickFlat") + apply_cfg_overrides(env_cfg, env_override) + + assert isinstance(env_cfg, ManagerBasedRlEnvCfg) + assert hydra_cfg.training.task_name == "Go2WJoystickFlat" + assert hydra_cfg.training.sim_backend == "drake" + assert list(env_cfg.actions) == ["motor"] + assert env_cfg.scene is not None + assert env_cfg.scene.default_keyframe_name == "home" + env_cfg.validate() + + +def test_go2w_flat_registry_is_manager_only_and_rough_owns_legacy_bridge() -> None: + registry.ensure_registries() + from unilab.tasks.locomotion.go2w.joystick import Go2WJoystickCfg, Go2WJoystickEnv + from unilab.tasks.locomotion.go2w.rough import Go2WJoystickRoughCfg, Go2WJoystickRoughEnv + + assert registry.list_registered_envs()["Go2WJoystickFlat"] == { + "config_factory": "ManagerBasedRlEnvCfg", + "available_backends": ["mujoco", "motrix", "drake"], + } + assert Go2WJoystickRoughCfg.__bases__ == (Go2WJoystickCfg,) + assert Go2WJoystickRoughEnv.__bases__ == (Go2WJoystickEnv,) + for legacy_override in ( + {"reward_config": {}}, + {"domain_rand": {"randomize_kp": False}}, + {"control_config": {"action_scale": 0.5}}, + ): + with pytest.raises(ValueError, match="has no attribute"): + apply_cfg_overrides(ManagerBasedRlEnvCfg(), legacy_override) + + +@pytest.mark.parametrize("backend", ["mujoco", "motrix"]) +def test_go2w_flat_registry_executes_real_manager_runtime(backend: str) -> None: + registry.ensure_registries() + _, _, env_override = _materialize(backend) + try: + env = registry.make( + "Go2WJoystickFlat", + sim_backend=backend, + env_cfg_override=env_override, + num_envs=2, + ) + except ImportError as exc: + pytest.skip(f"{backend} runtime unavailable: {exc}") + + try: + assert isinstance(env, ManagerBasedRlEnv) + assert env.obs_groups_spec == {"obs": 53, "critic": 72} + assert env.action_space.shape == (16,) + action = env.action_manager.get_term("motor") + assert isinstance(action, Go2WMixedAction) + np.testing.assert_allclose(action.leg_kp, 50.0) + np.testing.assert_allclose(action.leg_kd, 1.5) + + obs, info = env.reset(seed=7) + assert {name: value.shape for name, value in obs.items()} == { + "obs": (2, 53), + "critic": (2, 72), + } + assert isinstance(info, dict) + np.testing.assert_allclose( + env.scene["robot"].data.default_joint_pos, + np.broadcast_to(_HOME_JOINT_POS, (2, 16)), + ) + + state = env.step(np.full((2, 16), 2.0, dtype=np.float32)) + np.testing.assert_allclose(action.raw_action, 1.0) + np.testing.assert_allclose(action.previous_raw_action, 0.0) + np.testing.assert_allclose( + action.processed_action[:, :12], + np.broadcast_to(_HOME_JOINT_POS[:12] + 0.5, (2, 12)), + ) + np.testing.assert_allclose(action.processed_action[:, 12:], 10.0) + np.testing.assert_allclose(state.obs["critic"][:, -16:], action.motor_torque) + for value in (*state.obs.values(), state.reward, action.motor_torque): + assert np.isfinite(value).all() + + env.reset(np.asarray([0], dtype=np.int32)) + np.testing.assert_allclose(action.raw_action[0], 0.0) + np.testing.assert_allclose(action.raw_action[1], 1.0) + np.testing.assert_allclose(action.motor_torque[0], 0.0) + finally: + env.close() + + +def test_go2w_flat_dead_zone_and_motor_gain_overrides_are_manager_owned() -> None: + _, env_cfg, _ = _materialize("mujoco") + command = env_cfg.commands["twist"] + command.ranges.lin_vel_x = (0.1, 0.1) + command.ranges.lin_vel_y = (0.0, 0.0) + command.ranges.ang_vel_z = (0.0, 0.0) + gains = env_cfg.events["motor_gains"].params + gains["kp_multiplier_range"] = (0.5, 0.5) + gains["kd_multiplier_range"] = (2.0, 2.0) + + env = make_manager_based_rl_env(env_cfg, num_envs=2, backend_type="mujoco") + try: + env.reset(seed=11) + np.testing.assert_allclose(env.command_manager.get_command("twist"), 0.0) + action = env.action_manager.get_term("motor") + assert isinstance(action, Go2WMixedAction) + np.testing.assert_allclose(action.leg_kp, 25.0) + np.testing.assert_allclose(action.leg_kd, 3.0) + finally: + env.close() + + +def test_go2w_flat_incomplete_motor_selection_fails_closed() -> None: + _, env_cfg, _ = _materialize("mujoco") + action = env_cfg.actions["motor"] + assert isinstance(action, Go2WMixedActionCfg) + action.actuator_names = ["FR_.*"] + + with pytest.raises(ValueError, match="requires exactly 16 actuators and target joints"): + make_manager_based_rl_env(env_cfg, num_envs=1, backend_type="mujoco") + + +def test_go2w_manager_terms_do_not_leak_backend_or_physical_layout() -> None: + source = inspect.getsource(manager_terms) + for forbidden in ("._backend", "getattr(", "hasattr(", "qpos", "qvel", "ASSETS_ROOT_PATH"): + assert forbidden not in source diff --git a/tests/envs/test_env_configs.py b/tests/envs/test_env_configs.py index 48bc874a8..7c9279249 100644 --- a/tests/envs/test_env_configs.py +++ b/tests/envs/test_env_configs.py @@ -2141,7 +2141,6 @@ def test_g1_motion_tracking_clip_end_does_not_override_true_termination(): # Environments that don't need special config overrides _STANDARD_ENVS = [ "Go1JoystickRough", - "Go2WJoystickFlat", "Go2WJoystickRough", "G1WalkFlat", "G1WalkRough", @@ -2426,34 +2425,3 @@ def test_g1_motion_tracking_deploy_reset_and_step_mujoco(): assert state.obs["critic"].shape == (2, 286) finally: env.close() - - -def test_go2w_mujoco_keeps_kp_kd_out_of_backend_position_actuator_path(): - _require_mujoco_runtime() - ensure_registries() - - from unilab.base import registry - - env = cast( - Any, - registry.make( - "Go2WJoystickFlat", - num_envs=2, - sim_backend="mujoco", - env_cfg_override={ - "reward_config": { - "scales": {"alive": 1.0, "torques": -0.0002}, - "tracking_sigma": 0.25, - "base_height_target": 0.3, - } - }, - ), - ) - try: - env.init_state() - assert env._backend._position_actuator_gains is None - assert env._backend._pre_step_control_fn.__self__ is env - assert env._backend._pre_step_control_fn.__func__ is env._pre_step_motor_control.__func__ - assert env._last_motor_ctrl.shape == (2, env.action_space.shape[0]) - finally: - env.close() From 7e86dc736c78581437b66c746cf543d7fe930700 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Wed, 19 Aug 2026 10:13:09 +0800 Subject: [PATCH 084/169] feat: migrate go2 footstand to manager api [skip ci] (#1214) --- .../task/sac/go2_footstand/base.yaml | 229 ++ .../task/sac/go2_footstand/drake.yaml | 58 +- conf/ppo/task/go2_footstand/base.yaml | 229 ++ conf/ppo/task/go2_footstand/drake.yaml | 58 +- conf/ppo/task/go2_footstand/motrix.yaml | 113 +- conf/ppo/task/go2_footstand/mujoco.yaml | 105 +- .../en/2-user_guide/4-tasks/1-locomotion.md | 33 +- .../source/en/5-reference/5-support_matrix.md | 2 +- .../2-user_guide/4-tasks/1-locomotion.md | 33 +- .../9-sim2sim_contract_status.md | 4 +- src/unilab/tasks/locomotion/go2/__init__.py | 5 +- src/unilab/tasks/locomotion/go2/footstand.py | 2007 ++++++++--------- tests/config/test_locomotion_params.py | 35 +- tests/envs/locomotion/test_go2_footstand.py | 1321 ++++------- 14 files changed, 2046 insertions(+), 2186 deletions(-) create mode 100644 conf/offpolicy/task/sac/go2_footstand/base.yaml create mode 100644 conf/ppo/task/go2_footstand/base.yaml diff --git a/conf/offpolicy/task/sac/go2_footstand/base.yaml b/conf/offpolicy/task/sac/go2_footstand/base.yaml new file mode 100644 index 000000000..814e6d7a2 --- /dev/null +++ b/conf/offpolicy/task/sac/go2_footstand/base.yaml @@ -0,0 +1,229 @@ +# @package _global_ +# Canonical Go2 footstand Manager-Based declaration for off-policy training. +# Keep this task surface aligned with the PPO owner; Hydra remains the sole entry. +env: + scene: + model_file: src/unilab/assets/robots/go2/scene_flat.xml + default_keyframe_name: home + entities: + robot: + root_body_name: base + joint_names: + - FL_hip_joint + - FL_thigh_joint + - FL_calf_joint + - FR_hip_joint + - FR_thigh_joint + - FR_calf_joint + - RL_hip_joint + - RL_thigh_joint + - RL_calf_joint + - RR_hip_joint + - RR_thigh_joint + - RR_calf_joint + body_names: + - base + - FL_hip + - FL_thigh + - FL_calf + - FR_hip + - FR_thigh + - FR_calf + - RL_hip + - RL_thigh + - RL_calf + - RR_hip + - RR_thigh + - RR_calf + geom_names: [floor] + actuator_names: + - FR_hip + - FR_thigh + - FR_calf + - FL_hip + - FL_thigh + - FL_calf + - RR_hip + - RR_thigh + - RR_calf + - RL_hip + - RL_thigh + - RL_calf + sim_dt: 0.004 + ctrl_dt: 0.02 + max_episode_seconds: 10.0 + adaptive_chunk_size: false + observations: + policy: + _target_: unilab.managers.ObservationGroupCfg + enable_corruption: true + terms: + frame: + _target_: unilab.managers.ObservationTermCfg + func: unilab.tasks.locomotion.go2.footstand.frame_observation + params: + action_name: joint_pos + noise: + _target_: unilab.managers._noise.UniformNoiseCfg + n_min: [-0.1, -0.1, -0.1, -0.2, -0.2, -0.2, -0.05, -0.05, -0.05, -0.01, -0.01, -0.01, -0.01, -0.01, -0.01, -0.01, -0.01, -0.01, -0.01, -0.01, -0.01, -1.5, -1.5, -1.5, -1.5, -1.5, -1.5, -1.5, -1.5, -1.5, -1.5, -1.5, -1.5, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] + n_max: [0.1, 0.1, 0.1, 0.2, 0.2, 0.2, 0.05, 0.05, 0.05, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] + history_length: 15 + critic: + _target_: unilab.managers.ObservationGroupCfg + terms: + frame: + _target_: unilab.managers.ObservationTermCfg + func: unilab.tasks.locomotion.go2.footstand.frame_observation + params: + action_name: joint_pos + history_length: 15 + privileged: + _target_: unilab.managers.ObservationTermCfg + func: unilab.tasks.locomotion.go2.footstand.privileged_observation + params: + action_name: joint_pos + actions: + joint_pos: + _target_: unilab.tasks.locomotion.go2.footstand.FootstandIncrementalActionCfg + entity_name: robot + actuator_names: + - FR_hip + - FR_thigh + - FR_calf + - FL_hip + - FL_thigh + - FL_calf + - RR_hip + - RR_thigh + - RR_calf + - RL_hip + - RL_thigh + - RL_calf + joint_names: + - FR_hip_joint + - FR_thigh_joint + - FR_calf_joint + - FL_hip_joint + - FL_thigh_joint + - FL_calf_joint + - RR_hip_joint + - RR_thigh_joint + - RR_calf_joint + - RL_hip_joint + - RL_thigh_joint + - RL_calf_joint + joint_position_limits: + - [-1.0472, 1.0472] + - [-1.5708, 3.4907] + - [-2.7227, -0.83776] + - [-1.0472, 1.0472] + - [-1.5708, 3.4907] + - [-2.7227, -0.83776] + - [-1.0472, 1.0472] + - [-0.5236, 4.5379] + - [-2.7227, -0.83776] + - [-1.0472, 1.0472] + - [-0.5236, 4.5379] + - [-2.7227, -0.83776] + action_scale: 0.3 + clip_actions: 1.0 + kp: 35.0 + kd: 0.5 + simulate_action_latency: false + events: + reset_scene_to_default: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.reset_scene_to_default + mode: reset + reset_root_state_uniform: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.reset_root_state_uniform + mode: reset + params: + pose_range: + x: [-0.5, 0.5] + y: [-0.5, 0.5] + z: [0.0, 0.0] + roll: [0.0, 0.0] + pitch: [0.0, 0.0] + yaw: [-3.141592653589793, 3.141592653589793] + velocity_range: + x: [-0.5, 0.5] + y: [-0.5, 0.5] + z: [-0.5, 0.5] + roll: [-0.5, 0.5] + pitch: [-0.5, 0.5] + yaw: [-0.5, 0.5] + reset_joints: + _target_: unilab.managers.EventTermCfg + func: unilab.tasks.locomotion.go2.footstand.FootstandJointReset + mode: reset + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + joint_names: ".*" + position_offset_range: [-0.05, 0.05] + pd_gains: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.pd_gains + mode: reset + params: + kp_range: [35.0, 35.0] + kd_range: [0.5, 0.5] + operation: abs + floor_friction: null + link_mass: null + torso_com: null + joint_armature: null + terminations: + footstand: + _target_: unilab.managers.TerminationTermCfg + func: unilab.tasks.locomotion.go2.footstand.FootstandTermination + params: + action_name: joint_pos + grace_steps: 100 + height_fraction: 0.8 + orientation_threshold: 0.2 + energy_threshold: 200.0 + time_out: + _target_: unilab.managers.TerminationTermCfg + func: unilab.envs.mdp.time_out + time_out: true + policy_observation_group: policy + critic_observation_group: critic + scale_rewards_by_dt: true + +reward: + footstand: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.go2.footstand.FootstandReward + weight: 1.0 + params: + state_term_name: footstand + scales: + height: 2.0 + orientation: 2.0 + contact: -1.0 + action_rate: -0.01 + termination: -2.0 + dof_pos_limits: -0.5 + torques: 0.0 + pose: -0.1 + penalty_contact: -0.2 + tar: 0.8 + rear_feet_contact: 0.5 + rear_leg_symmetry: -0.2 + front_leg_motion: -0.05 + upright_stability: -0.2 + knee_clearance: -0.5 + stay_still: -0.1 + energy: -0.003 + dof_acc: -2.5e-7 + soft_joint_pos_limit_factor: 0.9 + knee_height_target: 0.08 + front_feet_min_separation: 0.16 + front_feet_side_margin: 0.04 + rear_hip_abduction_margin: 0.25 + rear_foot_slip_deadband: 0.02 + rear_foot_anchor_radius: 0.03 diff --git a/conf/offpolicy/task/sac/go2_footstand/drake.yaml b/conf/offpolicy/task/sac/go2_footstand/drake.yaml index 31035486b..3c564a027 100644 --- a/conf/offpolicy/task/sac/go2_footstand/drake.yaml +++ b/conf/offpolicy/task/sac/go2_footstand/drake.yaml @@ -1,4 +1,8 @@ # @package _global_ +defaults: + - /task/sac/go2_footstand/base + - _self_ + training: task_name: Go2FootStand sim_backend: drake @@ -9,7 +13,7 @@ training: algo: algo_log_name: fast_sac_drake - num_envs: 512 + num_envs: 4096 batch_size: 1024 replay_buffer_n: 512 updates_per_step: 2 @@ -26,50 +30,12 @@ algo: use_compile: false env: - sim_dt: 0.004 drake_backend_mode: batch drake_nthread: 20 - add_body_sensors: true - obs_history_len: 15 - energy_termination_threshold: 200.0 - noise_config: - level: 1.0 - scale_joint_angle: 0.01 - scale_joint_vel: 1.5 - scale_gyro: 0.2 - scale_gravity: 0.05 - scale_linvel: 0.1 - control_config: - action_scale: 0.3 - domain_rand: - randomize_floor_friction: false - randomize_link_mass: false - torso_added_mass_range: null - randomize_torso_com: false - randomize_dof_armature: false - randomize_reset_joint_qpos: true - reset_joint_qpos_range: [-0.05, 0.05] - -reward: - scales: - height: 2.0 - orientation: 2.0 - contact: -1.0 - action_rate: -0.01 - termination: -2.0 - dof_pos_limits: -0.5 - torques: 0.0 - pose: -0.1 - penalty_contact: -0.2 - tar: 0.8 - rear_feet_contact: 0.5 - rear_leg_symmetry: -0.2 - front_leg_motion: -0.05 - upright_stability: -0.2 - knee_clearance: -0.5 - stay_still: -0.1 - energy: -0.003 - dof_acc: -2.5e-7 - tracking_sigma: 0.25 - base_height_target: 0.3 - knee_height_target: 0.08 + events: + # Drake keeps the task joint reset and rejects unsupported payload fields. + pd_gains: null + floor_friction: null + link_mass: null + torso_com: null + joint_armature: null diff --git a/conf/ppo/task/go2_footstand/base.yaml b/conf/ppo/task/go2_footstand/base.yaml new file mode 100644 index 000000000..8b7abd5d2 --- /dev/null +++ b/conf/ppo/task/go2_footstand/base.yaml @@ -0,0 +1,229 @@ +# @package _global_ +# Canonical Go2 footstand Manager-Based declaration. Hydra owns every task term; +# backend leaves only select identity, tuning, and supported reset capabilities. +env: + scene: + model_file: src/unilab/assets/robots/go2/scene_flat.xml + default_keyframe_name: home + entities: + robot: + root_body_name: base + joint_names: + - FL_hip_joint + - FL_thigh_joint + - FL_calf_joint + - FR_hip_joint + - FR_thigh_joint + - FR_calf_joint + - RL_hip_joint + - RL_thigh_joint + - RL_calf_joint + - RR_hip_joint + - RR_thigh_joint + - RR_calf_joint + body_names: + - base + - FL_hip + - FL_thigh + - FL_calf + - FR_hip + - FR_thigh + - FR_calf + - RL_hip + - RL_thigh + - RL_calf + - RR_hip + - RR_thigh + - RR_calf + geom_names: [floor] + actuator_names: + - FR_hip + - FR_thigh + - FR_calf + - FL_hip + - FL_thigh + - FL_calf + - RR_hip + - RR_thigh + - RR_calf + - RL_hip + - RL_thigh + - RL_calf + sim_dt: 0.004 + ctrl_dt: 0.02 + max_episode_seconds: 10.0 + adaptive_chunk_size: false + observations: + policy: + _target_: unilab.managers.ObservationGroupCfg + enable_corruption: true + terms: + frame: + _target_: unilab.managers.ObservationTermCfg + func: unilab.tasks.locomotion.go2.footstand.frame_observation + params: + action_name: joint_pos + noise: + _target_: unilab.managers._noise.UniformNoiseCfg + n_min: [-0.1, -0.1, -0.1, -0.2, -0.2, -0.2, -0.05, -0.05, -0.05, -0.01, -0.01, -0.01, -0.01, -0.01, -0.01, -0.01, -0.01, -0.01, -0.01, -0.01, -0.01, -1.5, -1.5, -1.5, -1.5, -1.5, -1.5, -1.5, -1.5, -1.5, -1.5, -1.5, -1.5, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] + n_max: [0.1, 0.1, 0.1, 0.2, 0.2, 0.2, 0.05, 0.05, 0.05, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] + history_length: 15 + critic: + _target_: unilab.managers.ObservationGroupCfg + terms: + frame: + _target_: unilab.managers.ObservationTermCfg + func: unilab.tasks.locomotion.go2.footstand.frame_observation + params: + action_name: joint_pos + history_length: 15 + privileged: + _target_: unilab.managers.ObservationTermCfg + func: unilab.tasks.locomotion.go2.footstand.privileged_observation + params: + action_name: joint_pos + actions: + joint_pos: + _target_: unilab.tasks.locomotion.go2.footstand.FootstandIncrementalActionCfg + entity_name: robot + actuator_names: + - FR_hip + - FR_thigh + - FR_calf + - FL_hip + - FL_thigh + - FL_calf + - RR_hip + - RR_thigh + - RR_calf + - RL_hip + - RL_thigh + - RL_calf + joint_names: + - FR_hip_joint + - FR_thigh_joint + - FR_calf_joint + - FL_hip_joint + - FL_thigh_joint + - FL_calf_joint + - RR_hip_joint + - RR_thigh_joint + - RR_calf_joint + - RL_hip_joint + - RL_thigh_joint + - RL_calf_joint + joint_position_limits: + - [-1.0472, 1.0472] + - [-1.5708, 3.4907] + - [-2.7227, -0.83776] + - [-1.0472, 1.0472] + - [-1.5708, 3.4907] + - [-2.7227, -0.83776] + - [-1.0472, 1.0472] + - [-0.5236, 4.5379] + - [-2.7227, -0.83776] + - [-1.0472, 1.0472] + - [-0.5236, 4.5379] + - [-2.7227, -0.83776] + action_scale: 0.3 + clip_actions: 1.0 + kp: 35.0 + kd: 0.5 + simulate_action_latency: false + events: + reset_scene_to_default: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.reset_scene_to_default + mode: reset + reset_root_state_uniform: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.reset_root_state_uniform + mode: reset + params: + pose_range: + x: [-0.5, 0.5] + y: [-0.5, 0.5] + z: [0.0, 0.0] + roll: [0.0, 0.0] + pitch: [0.0, 0.0] + yaw: [-3.141592653589793, 3.141592653589793] + velocity_range: + x: [-0.5, 0.5] + y: [-0.5, 0.5] + z: [-0.5, 0.5] + roll: [-0.5, 0.5] + pitch: [-0.5, 0.5] + yaw: [-0.5, 0.5] + reset_joints: + _target_: unilab.managers.EventTermCfg + func: unilab.tasks.locomotion.go2.footstand.FootstandJointReset + mode: reset + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + joint_names: ".*" + position_offset_range: [-0.05, 0.05] + pd_gains: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.pd_gains + mode: reset + params: + kp_range: [35.0, 35.0] + kd_range: [0.5, 0.5] + operation: abs + floor_friction: null + link_mass: null + torso_com: null + joint_armature: null + terminations: + footstand: + _target_: unilab.managers.TerminationTermCfg + func: unilab.tasks.locomotion.go2.footstand.FootstandTermination + params: + action_name: joint_pos + grace_steps: 100 + height_fraction: 0.8 + orientation_threshold: 0.2 + energy_threshold: 200.0 + time_out: + _target_: unilab.managers.TerminationTermCfg + func: unilab.envs.mdp.time_out + time_out: true + policy_observation_group: policy + critic_observation_group: critic + scale_rewards_by_dt: true + +reward: + footstand: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.go2.footstand.FootstandReward + weight: 1.0 + params: + state_term_name: footstand + scales: + height: 2.0 + orientation: 2.0 + contact: -1.0 + action_rate: -0.01 + termination: -2.0 + dof_pos_limits: -0.5 + torques: 0.0 + pose: -0.1 + penalty_contact: -0.2 + tar: 0.8 + rear_feet_contact: 0.5 + rear_leg_symmetry: -0.2 + front_leg_motion: -0.05 + upright_stability: -0.2 + knee_clearance: -0.5 + stay_still: -0.1 + energy: -0.003 + dof_acc: -2.5e-7 + soft_joint_pos_limit_factor: 0.9 + knee_height_target: 0.08 + front_feet_min_separation: 0.16 + front_feet_side_margin: 0.04 + rear_hip_abduction_margin: 0.25 + rear_foot_slip_deadband: 0.02 + rear_foot_anchor_radius: 0.03 diff --git a/conf/ppo/task/go2_footstand/drake.yaml b/conf/ppo/task/go2_footstand/drake.yaml index b8c00f148..3d7c02a40 100644 --- a/conf/ppo/task/go2_footstand/drake.yaml +++ b/conf/ppo/task/go2_footstand/drake.yaml @@ -1,36 +1,26 @@ # @package _global_ +defaults: + - /task/go2_footstand/base + - _self_ + training: task_name: Go2FootStand sim_backend: drake env: - sim_dt: 0.004 drake_backend_mode: batch drake_nthread: 0 - add_body_sensors: true - obs_history_len: 15 - energy_termination_threshold: 200.0 - noise_config: - level: 1.0 - scale_joint_angle: 0.01 - scale_joint_vel: 1.5 - scale_gyro: 0.2 - scale_gravity: 0.05 - scale_linvel: 0.1 - control_config: - action_scale: 0.3 - domain_rand: - randomize_floor_friction: false - randomize_link_mass: false - torso_added_mass_range: null - randomize_torso_com: false - randomize_dof_armature: false - randomize_reset_joint_qpos: true - reset_joint_qpos_range: [-0.05, 0.05] + events: + # Drake supports the joint-state reset, but not these reset payload fields. + pd_gains: null + floor_friction: null + link_mass: null + torso_com: null + joint_armature: null algo: empirical_normalization: true - num_envs: 1024 + num_envs: 4096 max_iterations: 10000 obs_groups: actor: @@ -41,27 +31,3 @@ algo: init_noise_std: 0.5 algorithm: entropy_coef: 0.005 - -reward: - scales: - height: 2.0 - orientation: 2.0 - contact: -1.0 - action_rate: -0.01 - termination: -2.0 - dof_pos_limits: -0.5 - torques: 0.0 - pose: -0.1 - penalty_contact: -0.2 - tar: 0.8 - rear_feet_contact: 0.5 - rear_leg_symmetry: -0.2 - front_leg_motion: -0.05 - upright_stability: -0.2 - knee_clearance: -0.5 - stay_still: -0.1 - energy: -0.003 - dof_acc: -2.5e-7 - tracking_sigma: 0.25 - base_height_target: 0.3 - knee_height_target: 0.08 diff --git a/conf/ppo/task/go2_footstand/motrix.yaml b/conf/ppo/task/go2_footstand/motrix.yaml index 98d5b3a4e..c53005fbd 100644 --- a/conf/ppo/task/go2_footstand/motrix.yaml +++ b/conf/ppo/task/go2_footstand/motrix.yaml @@ -1,40 +1,27 @@ # @package _global_ +defaults: + - /task/go2_footstand/base + - _self_ + training: task_name: Go2FootStand sim_backend: motrix no_play: true + env: - sim_dt: 0.004 - add_body_sensors: true - obs_history_len: 15 - soft_joint_pos_limit_factor: 0.9 - energy_termination_threshold: 200.0 - noise_config: - level: 1.0 - scale_joint_angle: 0.01 - scale_joint_vel: 1.5 - scale_gyro: 0.2 - scale_gravity: 0.05 - scale_linvel: 0.1 - control_config: - action_scale: 0.3 - clip_actions: 1.0 - Kd: 0.5 - domain_rand: - randomize_floor_friction: false - floor_friction_range: [0.6, 1.0] - # Motrix does not implement dof_armature randomization; disabled. - randomize_dof_armature: false - randomize_link_mass: false - link_mass_scale_range: [0.95, 1.05] - torso_added_mass_range: [0.0, 0.0] - randomize_torso_com: false - torso_com_offset_range: [-0.02, 0.02] - randomize_reset_joint_qpos: true - reset_joint_qpos_range: [-0.02, 0.02] + events: + reset_joints: + params: + position_offset_range: [-0.02, 0.02] + # Unsupported model-field payloads remain explicitly disabled for Motrix. + floor_friction: null + link_mass: null + torso_com: null + joint_armature: null + algo: empirical_normalization: true - num_envs: 1024 + num_envs: 4096 max_iterations: 10000 obs_groups: actor: @@ -45,38 +32,38 @@ algo: init_noise_std: 0.5 algorithm: entropy_coef: 0.005 + reward: - scales: - height: 2.0 - orientation: 3.0 - contact: -1.0 - action_rate: -0.01 - termination: -2.0 - dof_pos_limits: -0.5 - torques: 0.0 - pose: -0.1 - penalty_contact: -0.2 - tar: 1.3 - rear_feet_contact: 0.5 - both_rear_feet_contact: 0.25 - rear_foot_slip: -1.0 - rear_foot_anchor: -0.15 - front_feet_air: 0.0 - balanced_footstand: 0.0 - rear_leg_symmetry: -0.2 - rear_leg_splay: -0.25 - front_leg_motion: -0.06 - front_leg_crossing: -2.0 - upright_stability: -0.2 - knee_clearance: -0.5 - stay_still: -0.12 - energy: -0.003 - dof_acc: -2.5e-7 - tracking_sigma: 0.25 - base_height_target: 0.3 - knee_height_target: 0.08 - front_feet_min_separation: 0.18 - front_feet_side_margin: 0.06 - rear_hip_abduction_margin: 0.25 - rear_foot_slip_deadband: 0.012 - rear_foot_anchor_radius: 0.04 + footstand: + params: + scales: + height: 2.0 + orientation: 3.0 + contact: -1.0 + action_rate: -0.01 + termination: -2.0 + dof_pos_limits: -0.5 + torques: 0.0 + pose: -0.1 + penalty_contact: -0.2 + tar: 1.3 + rear_feet_contact: 0.5 + both_rear_feet_contact: 0.25 + rear_foot_slip: -1.0 + rear_foot_anchor: -0.15 + front_feet_air: 0.0 + balanced_footstand: 0.0 + rear_leg_symmetry: -0.2 + rear_leg_splay: -0.25 + front_leg_motion: -0.06 + front_leg_crossing: -2.0 + upright_stability: -0.2 + knee_clearance: -0.5 + stay_still: -0.12 + energy: -0.003 + dof_acc: -2.5e-7 + front_feet_min_separation: 0.18 + front_feet_side_margin: 0.06 + rear_hip_abduction_margin: 0.25 + rear_foot_slip_deadband: 0.012 + rear_foot_anchor_radius: 0.04 diff --git a/conf/ppo/task/go2_footstand/mujoco.yaml b/conf/ppo/task/go2_footstand/mujoco.yaml index cca9fb89d..6e8e80c94 100644 --- a/conf/ppo/task/go2_footstand/mujoco.yaml +++ b/conf/ppo/task/go2_footstand/mujoco.yaml @@ -1,37 +1,65 @@ # @package _global_ +defaults: + - /task/go2_footstand/base + - _self_ + training: task_name: Go2FootStand sim_backend: mujoco + env: - sim_dt: 0.004 - add_body_sensors: true - obs_history_len: 15 - energy_termination_threshold: 200.0 - noise_config: - level: 1.0 - scale_joint_angle: 0.01 - scale_joint_vel: 1.5 - scale_gyro: 0.2 - scale_gravity: 0.05 - scale_linvel: 0.1 - control_config: - action_scale: 0.3 - domain_rand: - randomize_floor_friction: true - floor_friction_range: [0.4, 1.0] - randomize_link_mass: true - link_mass_scale_range: [0.9, 1.1] - torso_added_mass_range: [-1.0, 1.0] - randomize_torso_com: true - torso_com_offset_range: [-0.05, 0.05] - randomize_dof_armature: true - dof_armature_scale_range: [1.0, 1.05] - randomize_reset_joint_qpos: true - reset_joint_qpos_range: [-0.05, 0.05] + events: + floor_friction: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.geom_friction + mode: reset + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + geom_names: floor + ranges: [0.4, 1.0] + operation: abs + link_mass: + _target_: unilab.managers.EventTermCfg + func: unilab.tasks.locomotion.go2.footstand.FootstandMassRandomization + mode: reset + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + body_names: ".*" + torso_body_name: base + link_mass_scale_range: [0.9, 1.1] + torso_added_mass_range: [-1.0, 1.0] + torso_com: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.randomize_rigid_body_com + mode: reset + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + body_names: base + com_range: + x: [-0.05, 0.05] + y: [-0.05, 0.05] + z: [-0.05, 0.05] + joint_armature: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.joint_armature + mode: reset + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + joint_names: ".*" + ranges: [1.0, 1.05] + operation: scale + algo: empirical_normalization: true - num_envs: 1024 # 4096 # 1024 - # max_iterations: 3000 + num_envs: 4096 max_iterations: 10000 obs_groups: actor: @@ -42,26 +70,3 @@ algo: init_noise_std: 0.5 algorithm: entropy_coef: 0.005 -reward: - scales: - height: 2.0 - orientation: 2.0 - contact: -1.0 - action_rate: -0.01 - termination: -2.0 - dof_pos_limits: -0.5 - torques: 0.0 - pose: -0.1 - penalty_contact: -0.2 - tar: 0.8 - rear_feet_contact: 0.5 - rear_leg_symmetry: -0.2 - front_leg_motion: -0.05 - upright_stability: -0.2 - knee_clearance: -0.5 - stay_still: -0.1 - energy: -0.003 - dof_acc: -2.5e-7 - tracking_sigma: 0.25 - base_height_target: 0.3 - knee_height_target: 0.08 diff --git a/docs/sphinx/source/en/2-user_guide/4-tasks/1-locomotion.md b/docs/sphinx/source/en/2-user_guide/4-tasks/1-locomotion.md index 2a786ce61..ed04ce19e 100644 --- a/docs/sphinx/source/en/2-user_guide/4-tasks/1-locomotion.md +++ b/docs/sphinx/source/en/2-user_guide/4-tasks/1-locomotion.md @@ -29,12 +29,14 @@ backend: {doc}`../../5-reference/5-support_matrix`. ## Go2 FootStand -`go2_footstand` is the Go2 front-feet-stand task. It is **MuJoCo-only**. +`go2_footstand` is the Go2 front-feet-stand task. Its PPO owner YAMLs +register MuJoCo, Motrix, and Drake; the SAC owner currently targets Drake. -- PPO config: `conf/ppo/task/go2_footstand/mujoco.yaml` -- Registered env: `Go2FootStand` (registered for `sim_backend="mujoco"`) +- Canonical PPO task config: `conf/ppo/task/go2_footstand/base.yaml` +- Backend owners: `conf/ppo/task/go2_footstand/{mujoco,motrix,drake}.yaml` +- Registered env: `Go2FootStand` (MuJoCo, Motrix, and Drake) - Implementation: `src/unilab/tasks/locomotion/go2/footstand.py` - (extends the Go2 base task) + (task-owned NumPy manager terms on the generic Manager-Based runtime) - Go2 model XML: `src/unilab/assets/robots/go2/go2.xml` ```bash @@ -66,14 +68,14 @@ The full FootStand recipe is a three-stage teacher-student pipeline; the shipped ### Observation Layout The `Go2FootStand` policy (actor) observation uses 15 history frames of 45 dims -each (`_FOOTSTAND_FRAME_OBS_DIM = 45`): +each (`FRAME_OBS_DIM = 45`): ```text linvel(3) + gyro(3) + gravity(3) + joint_position_delta(12) + joint_velocity(12) + last_action(12) ``` So the policy observation is `45 * 15 = 675`. The value (critic) observation -appends the current-step privileged tail (`_FOOTSTAND_PRIVILEGED_TAIL_DIM = 49`) +appends the current-step privileged tail (`PRIVILEGED_OBS_DIM = 49`) after that history: ```text @@ -84,23 +86,24 @@ The value observation is therefore `675 + 49 = 724`. ### Rewards And Terminations -Defaults come from `conf/ppo/task/go2_footstand/mujoco.yaml`. The reward scales +Defaults come from `conf/ppo/task/go2_footstand/base.yaml`; backend leaves only +override backend-specific terms and tuning. The reward scales include stand `height`, `orientation`, `rear_feet_contact`, target front-leg angle (`tar`), `action_rate`, `dof_pos_limits`, `front_leg_motion`, `rear_leg_symmetry`, `knee_clearance`, `upright_stability`, `stay_still`, `pose`, plus `energy` and `dof_acc` penalties; `termination` and `penalty_contact` drive the termination / penalty paths (front-leg / front-body contact, low height, bad orientation, and a -high-energy cutoff via `energy_termination_threshold`). +high-energy cutoff in the `footstand` termination term). ### Tuning Keys -- `env.obs_history_len`: policy observation history length; config default is `15`. -- `env.energy_termination_threshold`: high-energy termination cutoff; config - default is `200.0`. -- `env.domain_rand`: floor friction, link mass, torso CoM, dof armature, and reset - joint qpos randomization. -- `reward.scales.height` / `orientation` / `rear_feet_contact`: stand pose and - rear-foot contact weights. +- `env.observations.policy.terms.frame.history_length`: policy history length + (default `15`). +- `env.terminations.footstand.params.energy_threshold`: high-energy cutoff + (default `200.0`). +- `env.events`: reset and domain-randomization terms. Backend owners explicitly + set unsupported model-field terms to `null`. +- `reward.footstand.params.scales`: stand, contact, motion, and energy weights. ### Near-Risk Validation diff --git a/docs/sphinx/source/en/5-reference/5-support_matrix.md b/docs/sphinx/source/en/5-reference/5-support_matrix.md index 5c9c3119a..bae44e4f4 100644 --- a/docs/sphinx/source/en/5-reference/5-support_matrix.md +++ b/docs/sphinx/source/en/5-reference/5-support_matrix.md @@ -54,7 +54,7 @@ recommendation metadata, so rows do not auto-promote to `Benchmarked` or | PPO (torch) | `g1_motion_tracking_deploy` | Tested | Registered | | PPO (torch) | `go1_joystick_rough` | Tested | Tested | | PPO (torch) | `go2_arm_manip_loco` | Tested | - | -| PPO (torch) | `go2_footstand` | Tested | - | +| PPO (torch) | `go2_footstand` | Tested | Tested | | PPO (torch) | `go2w_joystick_flat` | Tested | Tested | | PPO (torch) | `go2w_joystick_rough` | Tested | Tested | | APPO (torch) | `go1_joystick_flat` | Tested | Registered | diff --git a/docs/sphinx/source/zh_CN/2-user_guide/4-tasks/1-locomotion.md b/docs/sphinx/source/zh_CN/2-user_guide/4-tasks/1-locomotion.md index 006275c6c..f8987b5da 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/4-tasks/1-locomotion.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/4-tasks/1-locomotion.md @@ -29,11 +29,14 @@ uv run train --algo sac --task g1_walk_flat --sim mujoco ## Go2 FootStand -`go2_footstand` 是 Go2 前足站立任务，**仅支持 MuJoCo**。 - -- PPO 配置：`conf/ppo/task/go2_footstand/mujoco.yaml` -- 环境注册名：`Go2FootStand`（注册于 `sim_backend="mujoco"`） -- 环境实现：`src/unilab/tasks/locomotion/go2/footstand.py`（继承 Go2 基础任务） +`go2_footstand` 是 Go2 前足站立任务。PPO owner YAML 已注册 MuJoCo、 +Motrix 和 Drake；当前 SAC owner 使用 Drake。 + +- PPO canonical 配置：`conf/ppo/task/go2_footstand/base.yaml` +- 后端 owner：`conf/ppo/task/go2_footstand/{mujoco,motrix,drake}.yaml` +- 环境注册名：`Go2FootStand`（MuJoCo、Motrix、Drake） +- 环境实现：`src/unilab/tasks/locomotion/go2/footstand.py` + （通用 Manager-Based runtime 上的 task-owned NumPy manager terms） - Go2 模型 XML：`src/unilab/assets/robots/go2/go2.xml` ```bash @@ -59,14 +62,14 @@ FootStand 的完整流程是三阶段教师-学生 pipeline；当前仓库里的 ### 观测口径 `Go2FootStand` 的策略（actor）网络观测使用 15 帧历史，每帧 45 维 -（`_FOOTSTAND_FRAME_OBS_DIM = 45`）： +（`FRAME_OBS_DIM = 45`）： ```text linvel(3) + gyro(3) + gravity(3) + joint_position_delta(12) + joint_velocity(12) + last_action(12) ``` 因此策略网络观测维度是 `45 * 15 = 675`。价值（critic）网络在这段历史观测后追加当前时刻的 -特权观测尾部（`_FOOTSTAND_PRIVILEGED_TAIL_DIM = 49`）： +特权观测尾部（`PRIVILEGED_OBS_DIM = 49`）： ```text gyro(3) + accelerometer(3) + linvel(3) + global_angvel(3) + dof_pos(12) + dof_vel(12) + torques(12) + height(1) @@ -76,19 +79,23 @@ gyro(3) + accelerometer(3) + linvel(3) + global_angvel(3) + dof_pos(12) + dof_ve ### 奖励与终止项 -默认奖励来自 `conf/ppo/task/go2_footstand/mujoco.yaml`。奖励权重包括站立 `height`、 +默认奖励来自 `conf/ppo/task/go2_footstand/base.yaml`，后端 leaf 只覆盖 +后端专属 term 和调优项。奖励权重包括站立 `height`、 `orientation`、`rear_feet_contact`、前腿目标角度（`tar`）、`action_rate`、 `dof_pos_limits`、`front_leg_motion`、`rear_leg_symmetry`、`knee_clearance`、 `upright_stability`、`stay_still`、`pose`，以及 `energy` 和 `dof_acc` 惩罚； `termination` 与 `penalty_contact` 驱动终止/惩罚路径（前腿/前身体接触、低高度、坏朝向， -以及由 `energy_termination_threshold` 控制的高能耗截断）。 +以及 `footstand` termination term 中的高能耗截断）。 ### 调参提示 -- `env.obs_history_len`：策略观测历史长度，配置默认为 `15`。 -- `env.energy_termination_threshold`：高能耗终止阈值，配置默认为 `200.0`。 -- `env.domain_rand`：地面摩擦、连杆质量、机身质心、关节惯量和重置关节位置随机化。 -- `reward.scales.height` / `orientation` / `rear_feet_contact`：站立姿态和后脚接触权重。 +- `env.observations.policy.terms.frame.history_length`：策略观测历史长度， + 默认为 `15`。 +- `env.terminations.footstand.params.energy_threshold`：高能耗终止阈值， + 默认为 `200.0`。 +- `env.events`：重置和 domain randomization terms；后端 owner 对不支持的 + model-field term 显式设为 `null`。 +- `reward.footstand.params.scales`：站立、接触、运动和能耗权重。 ### 近风险检查 diff --git a/docs/sphinx/source/zh_CN/4-developer_guide/9-sim2sim_contract_status.md b/docs/sphinx/source/zh_CN/4-developer_guide/9-sim2sim_contract_status.md index 18459b02d..f981f54c5 100644 --- a/docs/sphinx/source/zh_CN/4-developer_guide/9-sim2sim_contract_status.md +++ b/docs/sphinx/source/zh_CN/4-developer_guide/9-sim2sim_contract_status.md @@ -28,12 +28,12 @@ uv run scripts/audit_sim2sim_contracts.py | Task | 判定 | 分歧 | |---|---|---| -| allegro_inhand · allegro_inhand_grasp · g1_climb_tracking · g1_motion_tracking · g1_wall_flip_tracking · go1_joystick_rough · go2_arm_manip_loco · go2_handstand · go2_joystick_flat · go2_joystick_rough · go2w_joystick_flat · go2w_joystick_rough · sharpa_inhand · sharpa_inhand_grasp | ✅ | 无 | +| allegro_inhand · allegro_inhand_grasp · g1_climb_tracking · g1_motion_tracking · g1_wall_flip_tracking · go1_joystick_rough · go2_arm_manip_loco · go2_footstand · go2_handstand · go2_joystick_flat · go2_joystick_rough · go2w_joystick_flat · go2w_joystick_rough · sharpa_inhand · sharpa_inhand_grasp | ✅ | 无 | | g1_box_tracking | ❌ | `empirical_normalization` false↔true；`obs_groups` critic 组差异 | | g1_flip_tracking | ❌ | `empirical_normalization` true↔false；`obs_groups`；`action_scale` 29 维↔默认 0.25；`sampling_mode` 两后端运行时同为 `start`（无害） | | g1_walk_flat | ❌ | `action_scale` 0.25↔0.5；`empirical_normalization` false↔true；`obs_groups` | | go1_joystick_flat | ❌ | `empirical_normalization` false↔true | -| g1_motion_tracking_deploy · go2_footstand | ⚪ | 仅 mujoco | +| g1_motion_tracking_deploy | ⚪ | 仅 mujoco | ## `conf/appo/task/` diff --git a/src/unilab/tasks/locomotion/go2/__init__.py b/src/unilab/tasks/locomotion/go2/__init__.py index d15ae6ac9..41578562f 100644 --- a/src/unilab/tasks/locomotion/go2/__init__.py +++ b/src/unilab/tasks/locomotion/go2/__init__.py @@ -1,12 +1,11 @@ -from .footstand import Go2FootStandCfg, Go2FootStandTask +from . import footstand as footstand from .joystick import Go2JoystickCfg, Go2WalkTask from .rough import Go2JoystickRoughCfg, Go2JoystickRoughEnv __all__ = [ - "Go2FootStandCfg", - "Go2FootStandTask", "Go2JoystickCfg", "Go2JoystickRoughCfg", "Go2JoystickRoughEnv", "Go2WalkTask", + "footstand", ] diff --git a/src/unilab/tasks/locomotion/go2/footstand.py b/src/unilab/tasks/locomotion/go2/footstand.py index 9bd06e914..011485c41 100644 --- a/src/unilab/tasks/locomotion/go2/footstand.py +++ b/src/unilab/tasks/locomotion/go2/footstand.py @@ -1,1087 +1,1068 @@ +"""Hydra-owned Manager-Based terms for the Go2 footstand task. + +The task keeps its historical NumPy observation, action, reward, termination, +and reset semantics while using only the public manager/entity facade. +""" + from __future__ import annotations -from dataclasses import dataclass, field -from typing import Any, cast +from dataclasses import dataclass +from numbers import Real +from typing import TYPE_CHECKING, Any, Protocol, cast import numpy as np -from unilab.assets import ASSETS_ROOT_PATH from unilab.base import registry -from unilab.base.backend import create_backend, env_backend_kwargs -from unilab.base.np_env import NpEnvState -from unilab.base.scene import SceneCfg -from unilab.dr import ResetPlan, ResetRandomizationPayload from unilab.dtype_config import get_global_dtype -from unilab.tasks.locomotion.common import rewards -from unilab.tasks.locomotion.common.commands import Commands -from unilab.tasks.locomotion.common.domain_rand import DomainRandConfig -from unilab.tasks.locomotion.common.dr_provider import LocomotionDRProvider -from unilab.tasks.locomotion.common.rewards import RewardContext -from unilab.tasks.locomotion.go2.base import ControlConfig, Go2BaseCfg, Go2BaseEnv, NoiseConfig +from unilab.envs import ManagerBasedRlEnvCfg, make_manager_based_rl_env +from unilab.managers import ActionTerm, ActionTermCfg, ManagerTermBase, ManagerTermBaseCfg +from unilab.managers.scene_entity_config import SceneEntityCfg from unilab.utils.rotation import np_quat_apply, np_quat_apply_inverse - -@dataclass -class InitState: - pos = [0.0, 0.0, 0.42] - - -@dataclass -class Go2DomainRandConfig(DomainRandConfig): - randomize_kp: bool = True - kp_multiplier_range: list[float] = field(default_factory=lambda: [0.9, 1.1]) - - randomize_kd: bool = True - kd_multiplier_range: list[float] = field(default_factory=lambda: [0.9, 1.1]) - - -@dataclass -class RewardConfig: - scales: dict[str, float] - tracking_sigma: float - base_height_target: float - target_foot_height: float = 0.1 - knee_height_target: float = 0.08 - front_feet_min_separation: float = 0.16 - front_feet_side_margin: float = 0.04 - rear_hip_abduction_margin: float = 0.25 - rear_foot_slip_deadband: float = 0.02 - rear_foot_anchor_radius: float = 0.03 - - -@dataclass -class JoystickSensor: - local_linvel = "local_linvel" - gyro = "gyro" - feet_force = ["FL_foot_contact", "FR_foot_contact", "RL_foot_contact", "RR_foot_contact"] - feet_pos = ["FL_pos", "FR_pos", "RL_pos", "RR_pos"] - global_pos = "global_position" - ternamate_contact = [ - "base1_contact", - "base2_contact", - "base3_contact", - "FL_hip_contact", - "FR_hip_contact", - "FL_thigh_contact", - "FR_thigh_contact", - "FL_calf_contact1", - "FL_calf_contact2", - "FR_calf_contact1", - "FR_calf_contact2", - ] - penalty_contact = [ - "RL_hip_contact", - "RR_hip_contact", - "RL_thigh_contact", - "RR_thigh_contact", - "RL_calf_contact1", - "RL_calf_contact2", - "RR_calf_contact1", - "RR_calf_contact2", - ] - - -@dataclass -class Go2HandStandCfg(Go2BaseCfg): - scene: SceneCfg = field( # pyright: ignore[reportIncompatibleVariableOverride] - default_factory=lambda: SceneCfg( - model_file=str(ASSETS_ROOT_PATH / "robots" / "go2" / "scene_flat.xml") +if TYPE_CHECKING: + from unilab.base.entity import Entity + from unilab.managers._types import ManagerBasedRlEnv, ManagerSensorView + from unilab.managers.action_manager import ActionManager + from unilab.managers.termination_manager import TerminationManager + + class _FootstandEnv(ManagerBasedRlEnv, Protocol): + @property + def common_step_counter(self) -> int: ... + + @property + def action_manager(self) -> ActionManager: ... + + @property + def termination_manager(self) -> TerminationManager: ... + + +NUM_ACTIONS = 12 +FRAME_OBS_DIM = 45 +PRIVILEGED_OBS_DIM = 49 + +_WORLD_GRAVITY = np.asarray([0.0, 0.0, -1.0], dtype=np.float32) +_BODY_FORWARD = np.asarray([1.0, 0.0, 0.0], dtype=np.float32) +_TARGET_HEIGHT = 0.53 +_CONTACT_THRESHOLD = 0.1 +_STAND_HEIGHT_FRACTION = 0.8 +_STAND_ORIENTATION_THRESHOLD = 0.5 + +_FRONT_FEET = np.asarray([0, 1], dtype=np.intp) +_REAR_FEET = np.asarray([2, 3], dtype=np.intp) +_FRONT_LEGS = np.arange(0, 6, dtype=np.intp) +_REAR_LEGS = np.arange(6, 12, dtype=np.intp) +_REAR_HIPS = np.asarray([6, 9], dtype=np.intp) +_REAR_LEFT = np.asarray([6, 7, 8], dtype=np.intp) +_REAR_RIGHT = np.asarray([9, 10, 11], dtype=np.intp) +_REAR_MIRROR_SIGNS = np.asarray([-1.0, 1.0, 1.0], dtype=np.float32) +_FRONT_LEG_TARGET = np.asarray([0.0, 1.82, -1.16, 0.0, 1.82, -1.16], dtype=np.float32) + +_TRACKED_BODY_NAMES = ( + "FL_thigh", + "FR_thigh", + "FL_calf", + "FR_calf", + "RL_calf", + "RR_calf", +) +_FRONT_LEFT_BODY_INDICES = np.asarray([0, 2], dtype=np.intp) +_FRONT_RIGHT_BODY_INDICES = np.asarray([1, 3], dtype=np.intp) +_KNEE_BODY_INDICES = np.asarray([2, 3, 4, 5], dtype=np.intp) + +_SENSOR_SPECS = ( + ("local_linvel", 3), + ("gyro", 3), + ("upvector", 3), + ("global_position", 3), + ("accelerometer", 3), + ("global_angvel", 3), + ("FL_foot_contact", 1), + ("FR_foot_contact", 1), + ("RL_foot_contact", 1), + ("RR_foot_contact", 1), + ("FL_pos", 3), + ("FR_pos", 3), + ("RL_pos", 3), + ("RR_pos", 3), + ("base1_contact", 1), + ("base2_contact", 1), + ("base3_contact", 1), + ("RL_hip_contact", 1), + ("RR_hip_contact", 1), + ("RL_thigh_contact", 1), + ("RR_thigh_contact", 1), + ("RL_calf_contact1", 1), + ("RL_calf_contact2", 1), + ("RR_calf_contact1", 1), + ("RR_calf_contact2", 1), + ("FL_hip_contact", 1), + ("FR_hip_contact", 1), + ("FL_thigh_contact", 1), + ("FR_thigh_contact", 1), + ("FL_calf_contact1", 1), + ("FL_calf_contact2", 1), + ("FR_calf_contact1", 1), + ("FR_calf_contact2", 1), +) +_FOOT_CONTACT_NAMES = tuple(name for name, _ in _SENSOR_SPECS[6:10]) +_FOOT_POSITION_NAMES = tuple(name for name, _ in _SENSOR_SPECS[10:14]) +_TERMINATION_CONTACT_NAMES = tuple(name for name, _ in _SENSOR_SPECS[14:25]) +_PENALTY_CONTACT_NAMES = tuple(name for name, _ in _SENSOR_SPECS[25:33]) + + +def _real( + term: str, + name: str, + value: Any, + *, + minimum: float | None = None, + strict_minimum: bool = False, +) -> float: + if isinstance(value, (bool, np.bool_)) or not isinstance(value, Real): + raise TypeError(f"{term} {name} must be a real number") + result = float(value) + if not np.isfinite(result): + raise ValueError(f"{term} {name} must be finite") + if minimum is not None and (result <= minimum if strict_minimum else result < minimum): + relation = "greater than" if strict_minimum else "at least" + raise ValueError(f"{term} {name} must be {relation} {minimum}") + return result + + +def _pair( + term: str, + name: str, + value: Any, + *, + minimum: float | None = None, +) -> tuple[float, float]: + if isinstance(value, (str, bytes)) or not isinstance(value, (tuple, list)): + raise TypeError(f"{term} {name} must be a two-value range") + if len(value) != 2: + raise ValueError(f"{term} {name} must contain two values") + lower = _real(term, f"{name}[0]", value[0], minimum=minimum) + upper = _real(term, f"{name}[1]", value[1], minimum=minimum) + if lower > upper: + raise ValueError(f"{term} {name} lower bound {lower} exceeds upper bound {upper}") + return lower, upper + + +def _name(term: str, field: str, value: Any) -> str: + if not isinstance(value, str) or not value: + raise ValueError(f"{term} {field} must be a non-empty string") + return value + + +def _env_ids(env: ManagerBasedRlEnv, env_ids: np.ndarray | slice | None) -> np.ndarray: + if env_ids is None: + return np.arange(env.num_envs, dtype=np.int32) + if isinstance(env_ids, slice): + return np.arange(env.num_envs, dtype=np.int32)[env_ids] + return np.asarray(env_ids, dtype=np.int32) + + +@dataclass(kw_only=True) +class FootstandIncrementalActionCfg(ActionTermCfg): + """Incremental position action in the historical actuator/policy order.""" + + actuator_names: tuple[str, ...] | list[str] + joint_names: tuple[str, ...] | list[str] + joint_position_limits: tuple[tuple[float, float], ...] | list[list[float]] + action_scale: float = 0.3 + clip_actions: float = 1.0 + kp: float = 35.0 + kd: float = 0.5 + simulate_action_latency: bool = False + + def build(self, env: ManagerBasedRlEnv) -> FootstandIncrementalAction: + return FootstandIncrementalAction(self, env) + + +class FootstandIncrementalAction(ActionTerm): + """Integrate clipped policy deltas and write position targets each substep.""" + + cfg: FootstandIncrementalActionCfg + _entity: Entity + + def __init__(self, cfg: FootstandIncrementalActionCfg, env: ManagerBasedRlEnv): + super().__init__(cfg, env) + term = type(self).__name__ + if cfg.clip is not None: + raise NotImplementedError(f"{term} does not support actuator-name clip") + for field_name, patterns in ( + ("actuator_names", cfg.actuator_names), + ("joint_names", cfg.joint_names), + ): + if isinstance(patterns, (str, bytes)) or not isinstance(patterns, (tuple, list)): + raise TypeError(f"{term} {field_name} must be an ordered sequence of patterns") + if len(patterns) != NUM_ACTIONS: + raise ValueError(f"{term} requires {NUM_ACTIONS} ordered {field_name} patterns") + self._scale = _real(term, "action_scale", cfg.action_scale, minimum=0.0) + self._clip_actions = _real( + term, "clip_actions", cfg.clip_actions, minimum=0.0, strict_minimum=True ) - ) - max_episode_seconds: float = 20.0 # pyright: ignore[reportIncompatibleVariableOverride] - init_state: InitState = field(default_factory=InitState) - commands: Commands = field(default_factory=Commands) - reward_config: RewardConfig | None = None - sensor: JoystickSensor = field(default_factory=JoystickSensor) # type: ignore[assignment] - domain_rand: Go2DomainRandConfig = field(default_factory=Go2DomainRandConfig) + self._kp = _real(term, "kp", cfg.kp, minimum=0.0) + self._kd = _real(term, "kd", cfg.kd, minimum=0.0) + if not isinstance(cfg.simulate_action_latency, bool): + raise TypeError(f"{term} simulate_action_latency must be bool") + + actuator_ids: list[int] = [] + joint_ids: list[int] = [] + actuator_names: list[str] = [] + joint_names: list[str] = [] + for actuator_pattern, joint_pattern in zip( + cfg.actuator_names, cfg.joint_names, strict=True + ): + if not isinstance(actuator_pattern, str) or not actuator_pattern: + raise ValueError(f"{term} actuator patterns must be non-empty strings") + if not isinstance(joint_pattern, str) or not joint_pattern: + raise ValueError(f"{term} joint patterns must be non-empty strings") + found_actuator_ids, found_actuator_names = self._entity.find_actuators( + (actuator_pattern,), preserve_order=True + ) + found_joint_ids, found_joint_names = self._entity.find_joints( + (joint_pattern,), preserve_order=True + ) + if len(found_actuator_ids) != 1 or len(found_joint_ids) != 1: + raise ValueError( + f"{term} patterns actuator={actuator_pattern!r}, joint={joint_pattern!r} " + "must each resolve exactly once; " + f"got actuators={found_actuator_names}, joints={found_joint_names}" + ) + actuator_ids.append(found_actuator_ids[0]) + joint_ids.append(found_joint_ids[0]) + actuator_names.extend(found_actuator_names) + joint_names.extend(found_joint_names) + if len(set(actuator_ids)) != NUM_ACTIONS or len(set(joint_ids)) != NUM_ACTIONS: + raise ValueError(f"{term} actuator-to-joint mapping must be one-to-one") + if set(joint_ids) != set(range(self._entity.num_joints)): + raise ValueError(f"{term} must control every declared Go2 joint exactly once") + + self._actuator_ids = np.asarray(actuator_ids, dtype=np.intp) + self._joint_ids = np.asarray(joint_ids, dtype=np.intp) + self._actuator_ids.setflags(write=False) + self._joint_ids.setflags(write=False) + self._actuator_names = tuple(actuator_names) + self._joint_names = tuple(joint_names) + + try: + selected_ranges = np.asarray(cfg.joint_position_limits, dtype=get_global_dtype()) + except (TypeError, ValueError) as exc: + raise TypeError(f"{term} joint_position_limits must be numeric") from exc + if selected_ranges.shape != (NUM_ACTIONS, 2): + raise ValueError(f"{term} joint_position_limits must have shape ({NUM_ACTIONS}, 2)") + if not np.isfinite(selected_ranges).all(): + raise ValueError(f"{term} joint_position_limits must be finite") + if np.any(selected_ranges[:, 0] >= selected_ranges[:, 1]): + raise ValueError(f"{term} joint_position_limits must have lower < upper") + self._target_lower = np.asarray(selected_ranges[:, 0], dtype=get_global_dtype()) + self._target_upper = np.asarray(selected_ranges[:, 1], dtype=get_global_dtype()) + self._joint_lower = np.empty((NUM_ACTIONS,), dtype=get_global_dtype()) + self._joint_upper = np.empty_like(self._joint_lower) + self._joint_lower[self._joint_ids] = self._target_lower + self._joint_upper[self._joint_ids] = self._target_upper + dtype = get_global_dtype() + shape = (env.num_envs, NUM_ACTIONS) + self._raw_action = np.zeros(shape, dtype=dtype) + self._previous_raw_action = np.zeros_like(self._raw_action) + self._target = np.asarray( + self._entity.data.joint_pos[:, self._joint_ids], dtype=dtype + ).copy() + self._state = FootstandState(cast("_FootstandEnv", env), self) -class Go2HandStandDomainRandomizationProvider(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]: - height = env.torso_height[env_ids].reshape(-1, 1) - env.feet_phase[env_ids, :] = 0 - env.feet_phase[:, 2] = 0.0 - env.feet_phase[:, 3] = 0.5 - env._feet_air_time[env_ids, :] = 0.0 - env._last_contacts[env_ids, :] = False - - return env._compute_obs( # type: ignore[no-any-return] - info_updates, - linvel, - gyro, - gravity, - dof_pos, - dof_vel, - height, - ) + @property + def action_dim(self) -> int: + return NUM_ACTIONS + @property + def raw_action(self) -> np.ndarray: + return self._raw_action -class Go2HandStandTask(Go2BaseEnv): - _cfg: Go2HandStandCfg # pyright: ignore[reportIncompatibleVariableOverride] - - def __init__(self, cfg: Go2HandStandCfg, num_envs=1, backend_type="mujoco"): - if cfg.reward_config is None: - raise ValueError("reward_config must be provided via Hydra configuration") - # Footstand historically left the MuJoCo pool tuning knobs at the - # create_backend defaults; keep them pinned so this convergence does - # not silently enable adaptive chunk tuning for the task. - backend_kwargs: dict[str, Any] = { - "base_name": cfg.asset.base_name, - "push_body_name": cfg.domain_rand.push_body_name, - "add_body_sensors": bool(getattr(cfg, "add_body_sensors", False)), - "position_actuator_gains": {"kp": cfg.control_config.Kp, "kd": cfg.control_config.Kd}, - **env_backend_kwargs(cfg), - "adaptive_chunk_size": False, - "bench_nsteps": 1, - } - backend = create_backend( - backend_type, - cfg.scene, - num_envs, - cfg.sim_dt, - **backend_kwargs, - ) - super().__init__(cfg, backend, num_envs) - self._enable_reward_log = True - self._reward_cfg = cfg.reward_config - self._init_reward_functions() - self._init_task_domain_randomization() - self.phase = np.zeros((num_envs,), dtype=np.float32) - self.feet_phase = np.zeros((num_envs, len(cfg.sensor.feet_force)), dtype=np.float32) - self.feet_phase[:, 2] = 0.0 - self.feet_phase[:, 3] = 0.5 - self.gait_frequency = 2 - self.feet_force = np.zeros((num_envs, len(cfg.sensor.feet_force), 1), dtype=np.float32) - self._feet_air_time = np.zeros((num_envs, len(cfg.sensor.feet_force)), dtype=np.float32) - self._last_contacts = np.zeros((num_envs, 2), dtype=bool) - self.feet_pos = np.zeros((num_envs, len(cfg.sensor.feet_pos), 3), dtype=np.float32) - self.torso_height = np.zeros((num_envs,), dtype=np.float32) - self._z_des = 0.55 - self._desired_gravity = np.array([-1, 0, 0]) - self.feet_geom_names = [0, 1] - self._joint_ids = [0, 1, 2, 3, 4, 5, 6, 9] - self._tar_ids = [6, 7, 8, 9, 10, 11] - self.target_angle = np.array([0, 1.82, -1.16, 0.0, 1.82, -1.16]) - - def _init_task_domain_randomization(self) -> None: - self._init_domain_randomization(Go2HandStandDomainRandomizationProvider()) + @property + def previous_raw_action(self) -> np.ndarray: + return self._previous_raw_action @property - def obs_groups_spec(self) -> dict[str, int]: - return {"obs": 42, "critic": 46} + def target(self) -> np.ndarray: + return self._target - def _init_reward_functions(self): - self._reward_fns: dict[str, Any] = {} + @property + def joint_ids(self) -> np.ndarray: + return self._joint_ids - def update_state(self, state: NpEnvState) -> NpEnvState: - return state + @property + def joint_names(self) -> tuple[str, ...]: + return self._joint_names - def _compute_obs( - self, - info: dict, - linvel, - gyro, - gravity, - dof_pos, - dof_vel, - height, - ) -> dict[str, np.ndarray]: - return {"obs": np.zeros((self._num_envs, 1)), "critic": np.zeros((self._num_envs, 1))} - - def _compute_reward(self, info: dict, linvel, gyro, dof_pos) -> np.ndarray: - return np.zeros((self._num_envs,), dtype=get_global_dtype()) - - def _cost_pose(self, ctx: RewardContext) -> np.ndarray: - dof_pos = self.get_dof_pos() - error = dof_pos[:, self._joint_ids] - self.default_angles[self._joint_ids] - return cast(np.ndarray, np.sum(np.square(error), axis=1)) - - def _reward_penalty_contact(self, ctx: RewardContext) -> np.ndarray: - contact_arrays = [] - for name in self._cfg.sensor.penalty_contact: - arr = self._backend.get_sensor_data(name) - contact_arrays.append(arr) - result = np.concatenate(contact_arrays, axis=1) - return np.asarray(np.any(result, axis=1)) - - def _reward_tar(self, ctx: RewardContext) -> np.ndarray: - dof_pos = self.get_dof_pos() - error = dof_pos[:, self._tar_ids] - self.target_angle - error = np.sum(np.square(error), axis=1) - mask = (self.torso_height >= self._z_des * 0.8).astype(np.float32) - return cast(np.ndarray, np.exp(-error / 1) * mask) - - -_GO2_DOF_TO_CTRL = np.array([3, 4, 5, 0, 1, 2, 9, 10, 11, 6, 7, 8], dtype=np.int32) -_WORLD_GRAVITY = np.array([0.0, 0.0, -1.0], dtype=np.float32) -_BODY_FORWARD = np.array([1.0, 0.0, 0.0], dtype=np.float32) -_FOOTSTAND_FRAME_OBS_DIM = 45 -_FOOTSTAND_PRIVILEGED_TAIL_DIM = 49 -_FOOTSTAND_MIN_OBS_HISTORY_LEN = 15 -_FOOTSTAND_FRONT_FEET = [0, 1] -_FOOTSTAND_REAR_FEET = [2, 3] -_FOOTSTAND_FRONT_LEG_IDS = [0, 1, 2, 3, 4, 5] -_FOOTSTAND_REAR_LEG_IDS = [6, 7, 8, 9, 10, 11] -_FOOTSTAND_REAR_HIP_IDS = [6, 9] -_FOOTSTAND_REAR_LEFT_LEG_IDS = [6, 7, 8] -_FOOTSTAND_REAR_RIGHT_LEG_IDS = [9, 10, 11] -_FOOTSTAND_REAR_MIRROR_SIGNS = np.array([-1.0, 1.0, 1.0], dtype=np.float32) -_FOOTSTAND_FRONT_LEG_TARGET = np.array([0.0, 1.82, -1.16, 0.0, 1.82, -1.16]) -_FOOTSTAND_TRACKED_BODY_NAMES = ("FL_thigh", "FR_thigh", "FL_calf", "FR_calf", "RL_calf", "RR_calf") -_FOOTSTAND_FRONT_LEFT_BODY_INDICES = [0, 2] -_FOOTSTAND_FRONT_RIGHT_BODY_INDICES = [1, 3] -_FOOTSTAND_KNEE_BODY_INDICES = [2, 3, 4, 5] -_FOOTSTAND_CONTACT_THRESHOLD = 0.1 -_FOOTSTAND_STAND_HEIGHT_FRACTION = 0.8 -_FOOTSTAND_STAND_ORIENTATION_THRESHOLD = 0.5 - - -@dataclass -class FootstandNoiseConfig(NoiseConfig): - level: float = 1.0 - scale_joint_angle: float = 0.01 - scale_joint_vel: float = 1.5 - scale_gyro: float = 0.2 - scale_gravity: float = 0.05 - scale_linvel: float = 0.1 - - -@dataclass -class FootstandControlConfig(ControlConfig): - clip_actions: float = 1.0 + @property + def actuator_names(self) -> tuple[str, ...]: + return self._actuator_names + @property + def joint_lower(self) -> np.ndarray: + return self._joint_lower -@dataclass -class Go2FootStandDomainRandConfig(Go2DomainRandConfig): - randomize_kp: bool = False - randomize_kd: bool = False - randomize_base_mass: bool = False - random_com: bool = False - push_robots: bool = False - - randomize_floor_friction: bool = True - floor_friction_range: list[float] = field(default_factory=lambda: [0.4, 1.0]) - - randomize_link_mass: bool = True - link_mass_scale_range: list[float] = field(default_factory=lambda: [0.9, 1.1]) - torso_added_mass_range: list[float] = field(default_factory=lambda: [-1.0, 1.0]) - - randomize_torso_com: bool = True - torso_com_offset_range: list[float] = field(default_factory=lambda: [-0.05, 0.05]) - - randomize_dof_armature: bool = True - dof_armature_scale_range: list[float] = field(default_factory=lambda: [1.0, 1.05]) - - randomize_reset_joint_qpos: bool = True - reset_joint_qpos_range: list[float] = field(default_factory=lambda: [-0.05, 0.05]) - - -@dataclass -class FootstandSensor(JoystickSensor): - accelerometer = "accelerometer" - global_angvel = "global_angvel" - ternamate_contact = [ - "base1_contact", - "base2_contact", - "base3_contact", - "RL_hip_contact", - "RR_hip_contact", - "RL_thigh_contact", - "RR_thigh_contact", - "RL_calf_contact1", - "RL_calf_contact2", - "RR_calf_contact1", - "RR_calf_contact2", - ] - penalty_contact = [ - "FL_hip_contact", - "FR_hip_contact", - "FL_thigh_contact", - "FR_thigh_contact", - "FL_calf_contact1", - "FL_calf_contact2", - "FR_calf_contact1", - "FR_calf_contact2", - ] - - -@registry.envcfg("Go2FootStand") -@dataclass -class Go2FootStandCfg(Go2HandStandCfg): - max_episode_seconds: float = 10.0 - add_body_sensors: bool = True - obs_history_len: int = _FOOTSTAND_MIN_OBS_HISTORY_LEN - soft_joint_pos_limit_factor: float = 0.9 - energy_termination_threshold: float = np.inf - termination_grace_steps: int = 100 - termination_height_fraction: float = 0.8 - termination_orientation_threshold: float = 0.2 - noise_config: FootstandNoiseConfig = field(default_factory=FootstandNoiseConfig) # type: ignore[assignment] - control_config: FootstandControlConfig = field( # type: ignore[assignment] - default_factory=lambda: FootstandControlConfig(action_scale=0.3) - ) - sensor: FootstandSensor = field(default_factory=FootstandSensor) # type: ignore[assignment] - domain_rand: Go2FootStandDomainRandConfig = field(default_factory=Go2FootStandDomainRandConfig) # type: ignore[assignment] + @property + def joint_upper(self) -> np.ndarray: + return self._joint_upper + @property + def state(self) -> FootstandState: + return self._state -class Go2FootStandDomainRandomizationProvider(Go2HandStandDomainRandomizationProvider): - def _get_reset_randomization_baselines( - self, env: Any - ) -> tuple[np.ndarray | None, np.ndarray | None, int | None, np.ndarray | None]: - return ( - env._base_body_mass, - env._base_geom_friction, - env._floor_geom_id, - env._base_dof_armature, - ) + @property + def entity(self) -> Entity: + return self._entity - def build_reset_plan(self, env: Any, env_ids: np.ndarray) -> ResetPlan: - plan = super().build_reset_plan(env, env_ids) - qpos = np.asarray(plan.qpos, dtype=get_global_dtype()).copy() - domain_rand = env.cfg.domain_rand - if domain_rand.randomize_reset_joint_qpos: - low, high = domain_rand.reset_joint_qpos_range - qpos[:, -env._num_action :] += np.random.uniform( - low, high, size=(len(env_ids), env._num_action) - ).astype(qpos.dtype) - - return ResetPlan( - env_ids=plan.env_ids, - qpos=qpos, - qvel=plan.qvel, - info_updates=plan.info_updates, - randomization=self._merge_reset_randomization( - plan.randomization, - env._build_playground_reset_randomization(len(env_ids)), - ), + @property + def estimated_torque(self) -> np.ndarray: + return self._state.torques + + def process_actions(self, actions: np.ndarray) -> None: + if not isinstance(actions, np.ndarray): + raise TypeError(f"expected np.ndarray actions, got {type(actions).__name__}") + if actions.shape != self._raw_action.shape: + raise ValueError(f"expected action shape {self._raw_action.shape}, got {actions.shape}") + if not np.isfinite(actions).all(): + raise ValueError("received NaN or Inf actions") + self._previous_raw_action[:] = self._raw_action + np.clip(actions, -self._clip_actions, self._clip_actions, out=self._raw_action) + executed = ( + self._previous_raw_action if self.cfg.simulate_action_latency else self._raw_action ) - - @staticmethod - def _merge_reset_randomization( - base: ResetRandomizationPayload | None, - override: ResetRandomizationPayload | None, - ) -> ResetRandomizationPayload | None: - if base is None or base.is_empty(): - return override - if override is None or override.is_empty(): - return base - return ResetRandomizationPayload( - base_mass_delta=base.base_mass_delta, - base_com_offset=base.base_com_offset, - gravity=base.gravity, - body_iquat=override.body_iquat if override.body_iquat is not None else base.body_iquat, - body_inertia=override.body_inertia - if override.body_inertia is not None - else base.body_inertia, - body_ipos=override.body_ipos if override.body_ipos is not None else base.body_ipos, - body_mass=override.body_mass if override.body_mass is not None else base.body_mass, - dof_armature=override.dof_armature - if override.dof_armature is not None - else base.dof_armature, - geom_friction=override.geom_friction - if override.geom_friction is not None - else base.geom_friction, - kp=base.kp, - kd=base.kd, + self._target += self._scale * executed + np.clip(self._target, self._target_lower, self._target_upper, out=self._target) + + def apply_actions(self) -> None: + self._entity.set_joint_position_target(self._target, joint_ids=self._joint_ids) + + def estimate_torque( + self, joint_pos: np.ndarray, joint_vel: np.ndarray, out: np.ndarray + ) -> None: + out.fill(0.0) + selected_pos = joint_pos[:, self._joint_ids] + selected_vel = joint_vel[:, self._joint_ids] + out[:, self._joint_ids] = self._kp * (self._target - selected_pos) - self._kd * selected_vel + + def reset(self, env_ids: np.ndarray | slice | None = None) -> None: + ids = _env_ids(self._env, env_ids) + self._raw_action[ids] = 0.0 + self._previous_raw_action[ids] = 0.0 + joint_pos = self._entity.data.joint_pos + self._target[ids] = joint_pos[ids][:, self._joint_ids] + self._state.reset(ids) + + +class FootstandState: + """One per-control-step snapshot shared by termination, reward, and observations.""" + + def __init__(self, env: _FootstandEnv, action: FootstandIncrementalAction): + self._env = env + self._action = action + self._entity = action.entity + names = tuple(name for name, _ in _SENSOR_SPECS) + self._sensor_view: ManagerSensorView = env.scene.bind_sensor_data(names) + expected_dims = tuple(width for _, width in _SENSOR_SPECS) + if self._sensor_view.dimensions != expected_dims: + raise ValueError( + "Footstand named-sensor dimensions differ from the task contract: " + f"expected={expected_dims}, got={self._sensor_view.dimensions}" + ) + offsets = np.cumsum((0, *expected_dims), dtype=np.intp) + self._sensor_slices = { + name: slice(int(offsets[index]), int(offsets[index + 1])) + for index, (name, _) in enumerate(_SENSOR_SPECS) + } + tracked_ids, tracked_names = self._entity.find_bodies( + _TRACKED_BODY_NAMES, preserve_order=True ) + if tuple(tracked_names) != _TRACKED_BODY_NAMES: + raise ValueError( + f"Footstand tracked body order differs from the task contract: {tracked_names}" + ) + self._tracked_body_ids = np.asarray(tracked_ids, dtype=np.intp) + self._tracked_body_ids.setflags(write=False) - 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]: - height = env._backend.get_sensor_data(env._cfg.sensor.global_pos)[env_ids, -1].reshape( - -1, 1 + dtype = get_global_dtype() + num_envs = env.num_envs + self.linvel = np.zeros((num_envs, 3), dtype=dtype) + self.gyro = np.zeros_like(self.linvel) + self.gravity = np.broadcast_to(_WORLD_GRAVITY, (num_envs, 3)).astype(dtype, copy=True) + self.upvector = -self.gravity.copy() + self.accelerometer = np.zeros_like(self.linvel) + self.global_angvel = np.zeros_like(self.linvel) + self.root_pos = np.zeros_like(self.linvel) + self.root_quat = np.zeros((num_envs, 4), dtype=dtype) + self.root_quat[:, 0] = 1.0 + self.root_linvel_w = np.zeros_like(self.linvel) + self.root_angvel_w = np.zeros_like(self.linvel) + self.joint_pos = np.asarray(self._entity.data.default_joint_pos, dtype=dtype).copy() + self.joint_vel = np.zeros_like(self.joint_pos) + self.qacc = np.zeros_like(self.joint_pos) + self.torques = np.zeros_like(self.joint_pos) + self.height = np.zeros((num_envs,), dtype=dtype) + self.orientation = np.zeros((num_envs,), dtype=dtype) + self.foot_contact = np.zeros((num_envs, 4), dtype=np.bool_) + self.foot_pos = np.zeros((num_envs, 4, 3), dtype=dtype) + self.termination_contact = np.zeros((num_envs,), dtype=np.bool_) + self.penalty_contact = np.zeros((num_envs,), dtype=np.bool_) + self.tracked_body_pos = np.zeros((num_envs, len(_TRACKED_BODY_NAMES), 3), dtype=dtype) + self.rear_speed = np.zeros((num_envs, 2), dtype=dtype) + self.rear_anchor_drift = np.zeros((num_envs, 2), dtype=dtype) + self.rear_anchor_contact = np.zeros((num_envs, 2), dtype=np.bool_) + self._last_foot_pos = np.zeros_like(self.foot_pos) + self._rear_anchor_pos = np.zeros((num_envs, 2, 2), dtype=dtype) + self._last_counter = int(env.common_step_counter) + + @property + def last_counter(self) -> int: + return self._last_counter + + @property + def default_joint_pos(self) -> np.ndarray: + return self._entity.data.default_joint_pos + + @property + def action(self) -> FootstandIncrementalAction: + return self._action + + def _sensor(self, values: np.ndarray, name: str) -> np.ndarray: + return values[:, self._sensor_slices[name]] + + def _capture(self) -> dict[str, np.ndarray]: + dtype = get_global_dtype() + sensors = np.asarray(self._sensor_view.read(), dtype=dtype) + root_quat = np.asarray(self._entity.data.root_link_quat_w, dtype=dtype) + gravity_w = np.broadcast_to(_WORLD_GRAVITY, (self._env.num_envs, 3)) + gravity = np.asarray(np_quat_apply_inverse(root_quat, gravity_w), dtype=dtype) + forward_w = np_quat_apply( + root_quat, np.broadcast_to(_BODY_FORWARD, (self._env.num_envs, 3)) ) - local_gravity = env._get_local_gravity()[env_ids] - accelerometer = env._backend.get_sensor_data(env._cfg.sensor.accelerometer)[env_ids] - global_angvel = env._backend.get_sensor_data(env._cfg.sensor.global_angvel)[env_ids] - env.torso_height[env_ids] = height[:, 0] - env._last_dof_vel_for_acc[env_ids, :] = dof_vel - env._last_terminated[env_ids] = False - env._motor_targets[env_ids] = env._dof_to_ctrl_order(dof_pos) - feet_pos = env._backend.get_sensor_data_batch(env._cfg.sensor.feet_pos).reshape( - env._num_envs, len(env._cfg.sensor.feet_pos), 3 + orientation = np.asarray(np.square(0.5 * forward_w[:, 2] + 0.5), dtype=dtype) + foot_contact = ( + np.concatenate([self._sensor(sensors, name) for name in _FOOT_CONTACT_NAMES], axis=1) + > _CONTACT_THRESHOLD ) - env._last_feet_pos[env_ids] = feet_pos[env_ids] - env._rear_foot_slip[env_ids] = 0.0 - env._rear_foot_anchor_pos[env_ids] = feet_pos[env_ids][:, _FOOTSTAND_REAR_FEET, :2] - env._rear_foot_anchor_contact[env_ids] = False - env._rear_foot_anchor[env_ids] = 0.0 - target_dof = env._ctrl_to_dof_order(env._motor_targets[env_ids]) - info_updates["torques"] = np.asarray( - env._cfg.control_config.Kp * (target_dof - dof_pos) - - env._cfg.control_config.Kd * dof_vel, - dtype=get_global_dtype(), + foot_pos = np.stack([self._sensor(sensors, name) for name in _FOOT_POSITION_NAMES], axis=1) + termination_contact = np.any( + np.concatenate( + [self._sensor(sensors, name) for name in _TERMINATION_CONTACT_NAMES], axis=1 + ), + axis=1, ) - - return env._compute_obs( # type: ignore[no-any-return] - info_updates, - linvel, - gyro, - local_gravity, - dof_pos, - dof_vel, - height, - accelerometer, - global_angvel, - env_ids=env_ids, + penalty_contact = np.any( + np.concatenate( + [self._sensor(sensors, name) for name in _PENALTY_CONTACT_NAMES], axis=1 + ), + axis=1, ) + return { + "linvel": self._sensor(sensors, "local_linvel"), + "gyro": self._sensor(sensors, "gyro"), + "gravity": gravity, + "upvector": self._sensor(sensors, "upvector"), + "accelerometer": self._sensor(sensors, "accelerometer"), + "global_angvel": self._sensor(sensors, "global_angvel"), + "root_pos": np.asarray(self._entity.data.root_link_pos_w, dtype=dtype), + "root_quat": root_quat, + "root_linvel_w": np.asarray(self._entity.data.root_link_lin_vel_w, dtype=dtype), + "root_angvel_w": np.asarray(self._entity.data.root_link_ang_vel_w, dtype=dtype), + "joint_pos": np.asarray(self._entity.data.joint_pos, dtype=dtype), + "joint_vel": np.asarray(self._entity.data.joint_vel, dtype=dtype), + "height": self._sensor(sensors, "global_position")[:, 2], + "orientation": orientation, + "foot_contact": foot_contact, + "foot_pos": foot_pos, + "termination_contact": termination_contact, + "penalty_contact": penalty_contact, + "tracked_body_pos": np.asarray( + self._entity.data.body_link_pos_w[:, self._tracked_body_ids], dtype=dtype + ), + } + def reset(self, env_ids: np.ndarray) -> None: + values = self._capture() + self.linvel[env_ids] = values["linvel"][env_ids] + self.gyro[env_ids] = values["gyro"][env_ids] + self.gravity[env_ids] = values["gravity"][env_ids] + self.upvector[env_ids] = values["upvector"][env_ids] + self.accelerometer[env_ids] = values["accelerometer"][env_ids] + self.global_angvel[env_ids] = values["global_angvel"][env_ids] + self.root_pos[env_ids] = values["root_pos"][env_ids] + self.root_quat[env_ids] = values["root_quat"][env_ids] + self.root_linvel_w[env_ids] = values["root_linvel_w"][env_ids] + self.root_angvel_w[env_ids] = values["root_angvel_w"][env_ids] + self.joint_pos[env_ids] = values["joint_pos"][env_ids] + self.joint_vel[env_ids] = values["joint_vel"][env_ids] + self.height[env_ids] = values["height"][env_ids] + self.orientation[env_ids] = values["orientation"][env_ids] + self.foot_contact[env_ids] = values["foot_contact"][env_ids] + self.foot_pos[env_ids] = values["foot_pos"][env_ids] + self.termination_contact[env_ids] = values["termination_contact"][env_ids] + self.penalty_contact[env_ids] = values["penalty_contact"][env_ids] + self.tracked_body_pos[env_ids] = values["tracked_body_pos"][env_ids] + self.qacc[env_ids] = 0.0 + torque = np.empty_like(self.torques) + self._action.estimate_torque(self.joint_pos, self.joint_vel, torque) + self.torques[env_ids] = torque[env_ids] + self._last_foot_pos[env_ids] = self.foot_pos[env_ids] + self.rear_speed[env_ids] = 0.0 + self._rear_anchor_pos[env_ids] = self.foot_pos[env_ids][:, _REAR_FEET, :2] + self.rear_anchor_contact[env_ids] = False + self.rear_anchor_drift[env_ids] = 0.0 + self._last_counter = int(self._env.common_step_counter) + + def snapshot(self, env: _FootstandEnv) -> FootstandState: + counter = int(env.common_step_counter) + if counter == self._last_counter: + return self + if counter != self._last_counter + 1: + raise RuntimeError( + "FootstandState missed a control-step update: " + f"last={self._last_counter}, current={counter}" + ) + values = self._capture() + new_joint_vel = values["joint_vel"] + np.subtract(new_joint_vel, self.joint_vel, out=self.qacc) + self.qacc /= env.step_dt -@registry.env("Go2FootStand", sim_backend="mujoco") -@registry.env("Go2FootStand", sim_backend="motrix") -@registry.env("Go2FootStand", sim_backend="drake") -class Go2FootStandTask(Go2HandStandTask): - _cfg: Go2FootStandCfg # pyright: ignore[reportIncompatibleVariableOverride] - - def __init__(self, cfg: Go2FootStandCfg, num_envs=1, backend_type="mujoco"): - super().__init__(cfg, num_envs=num_envs, backend_type=backend_type) - self._z_des = 0.53 - self._desired_forward_vec = np.array([0.0, 0.0, 1.0], dtype=np.float32) - self._init_footstand_pose_targets() - self._init_soft_joint_limits() - self._init_motor_target_limits() - self._last_dof_vel_for_acc = np.zeros((num_envs, self._num_action), dtype=np.float32) - self._last_terminated = np.zeros((num_envs,), dtype=bool) - self._motor_targets = np.zeros((num_envs, self._num_action), dtype=get_global_dtype()) - self._last_feet_pos = np.zeros((num_envs, len(cfg.sensor.feet_pos), 3), dtype=np.float32) - self._rear_foot_slip = np.zeros((num_envs,), dtype=get_global_dtype()) - self._rear_foot_anchor_pos = np.zeros( - (num_envs, len(_FOOTSTAND_REAR_FEET), 2), dtype=get_global_dtype() + new_foot_pos = values["foot_pos"] + rear_delta = new_foot_pos[:, _REAR_FEET, :2] - self._last_foot_pos[:, _REAR_FEET, :2] + self.rear_speed[:] = np.linalg.norm(rear_delta / env.step_dt, axis=2) + + standing = (values["height"] >= _TARGET_HEIGHT * _STAND_HEIGHT_FRACTION) & ( + values["orientation"] >= _STAND_ORIENTATION_THRESHOLD ) - self._rear_foot_anchor_contact = np.zeros((num_envs, len(_FOOTSTAND_REAR_FEET)), dtype=bool) - self._rear_foot_anchor = np.zeros((num_envs,), dtype=get_global_dtype()) - self._obs_history = np.zeros( - (num_envs, self._obs_history_len, _FOOTSTAND_FRAME_OBS_DIM), + anchor_contact = values["foot_contact"][:, _REAR_FEET] & standing[:, None] + rear_xy = new_foot_pos[:, _REAR_FEET, :2] + new_contact = anchor_contact & ~self.rear_anchor_contact + self._rear_anchor_pos[new_contact] = rear_xy[new_contact] + self.rear_anchor_contact[:] = anchor_contact + self.rear_anchor_drift[:] = np.linalg.norm(rear_xy - self._rear_anchor_pos, axis=2) + + self.linvel[:] = values["linvel"] + self.gyro[:] = values["gyro"] + self.gravity[:] = values["gravity"] + self.upvector[:] = values["upvector"] + self.accelerometer[:] = values["accelerometer"] + self.global_angvel[:] = values["global_angvel"] + self.root_pos[:] = values["root_pos"] + self.root_quat[:] = values["root_quat"] + self.root_linvel_w[:] = values["root_linvel_w"] + self.root_angvel_w[:] = values["root_angvel_w"] + self.joint_pos[:] = values["joint_pos"] + self.joint_vel[:] = values["joint_vel"] + self.height[:] = values["height"] + self.orientation[:] = values["orientation"] + self.foot_contact[:] = values["foot_contact"] + self.foot_pos[:] = values["foot_pos"] + self.termination_contact[:] = values["termination_contact"] + self.penalty_contact[:] = values["penalty_contact"] + self.tracked_body_pos[:] = values["tracked_body_pos"] + self._action.estimate_torque(self.joint_pos, self.joint_vel, self.torques) + self._last_foot_pos[:] = self.foot_pos + self._last_counter = counter + return self + + def frame(self, env: _FootstandEnv) -> np.ndarray: + self.snapshot(env) + return np.concatenate( + ( + self.linvel, + self.gyro, + self.gravity, + self.joint_pos - self.default_joint_pos, + self.joint_vel, + self._action.previous_raw_action, + ), + axis=1, dtype=get_global_dtype(), ) - self._critic_obs_history = np.zeros_like(self._obs_history) - domain_rand = self._cfg.domain_rand - needs_floor_friction = bool(domain_rand.randomize_floor_friction) - needs_body_mass = bool( - domain_rand.randomize_link_mass or domain_rand.torso_added_mass_range is not None - ) - needs_body_ipos = bool(domain_rand.randomize_torso_com) - needs_dof_armature = bool(domain_rand.randomize_dof_armature) - self._base_geom_friction = ( - self._backend.get_geom_friction() - if needs_floor_friction - else np.zeros((0, 3), dtype=np.float64) - ) - self._floor_geom_id = ( - self._backend.get_geom_id(self._cfg.asset.ground) if needs_floor_friction else -1 - ) - self._base_body_id = ( - self._backend.get_body_id(self._cfg.asset.base_name) - if needs_body_mass or needs_body_ipos - else -1 + + def privileged(self, env: _FootstandEnv) -> np.ndarray: + self.snapshot(env) + return np.concatenate( + ( + self.gyro, + self.accelerometer, + self.linvel, + self.global_angvel, + self.joint_pos, + self.joint_vel, + self.torques, + self.height[:, None], + ), + axis=1, + dtype=get_global_dtype(), ) - self._base_body_mass = ( - self._backend.get_body_mass() if needs_body_mass else np.zeros((0,), dtype=np.float64) + + +def _action(env: _FootstandEnv, action_name: str) -> FootstandIncrementalAction: + name = _name("Footstand manager term", "action_name", action_name) + try: + action = env.action_manager.get_term(name) + except KeyError as exc: + raise KeyError(f"Footstand action term {name!r} is unavailable") from exc + if not isinstance(action, FootstandIncrementalAction): + raise TypeError( + f"Footstand action term {name!r} must be FootstandIncrementalAction, " + f"got {type(action).__name__}" ) - self._base_body_ipos = ( - self._backend.get_body_ipos() if needs_body_ipos else np.zeros((0, 3), dtype=np.float64) + return action + + +def frame_observation(env: _FootstandEnv, action_name: str) -> np.ndarray: + return _action(env, action_name).state.frame(env) + + +def privileged_observation(env: _FootstandEnv, action_name: str) -> np.ndarray: + return _action(env, action_name).state.privileged(env) + + +class FootstandTermination(ManagerTermBase): + """Aggregate the historical non-timeout termination state before rewards.""" + + def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv): + super().__init__(env) + task_env = cast("_FootstandEnv", env) + term = type(self).__name__ + allowed = { + "action_name", + "grace_steps", + "height_fraction", + "orientation_threshold", + "energy_threshold", + } + unknown = sorted(set(cfg.params) - allowed) + if unknown: + raise TypeError(f"{term} received unsupported parameters: {unknown}") + self._state = _action( + task_env, _name(term, "action_name", cfg.params.get("action_name")) + ).state + grace = cfg.params.get("grace_steps") + if isinstance(grace, (bool, np.bool_)) or not isinstance(grace, (int, np.integer)): + raise TypeError(f"{term} grace_steps must be an integer") + if int(grace) < 0: + raise ValueError(f"{term} grace_steps must be non-negative") + self._grace_steps = int(grace) + self._height_fraction = _real( + term, "height_fraction", cfg.params.get("height_fraction"), minimum=0.0 ) - self._base_dof_armature = ( - self._backend.get_dof_armature() - if needs_dof_armature - else np.zeros((0,), dtype=np.float64) + self._orientation_threshold = _real( + term, "orientation_threshold", cfg.params.get("orientation_threshold"), minimum=0.0 ) - self._tracked_body_ids = self._backend.get_body_ids(_FOOTSTAND_TRACKED_BODY_NAMES) - self._tracked_body_pos = np.zeros( - (num_envs, len(_FOOTSTAND_TRACKED_BODY_NAMES), 3), dtype=get_global_dtype() + self._energy_threshold = _real( + term, "energy_threshold", cfg.params.get("energy_threshold"), minimum=0.0 ) - self._init_domain_randomization(Go2FootStandDomainRandomizationProvider()) - - def _init_task_domain_randomization(self) -> None: - pass - - def _init_footstand_pose_targets(self) -> None: - self.feet_geom_names = list(_FOOTSTAND_FRONT_FEET) - self._joint_ids = list(_FOOTSTAND_REAR_LEG_IDS) - self._tar_ids = list(_FOOTSTAND_FRONT_LEG_IDS) - self.target_angle = np.asarray(_FOOTSTAND_FRONT_LEG_TARGET, dtype=get_global_dtype()) + self.terminated = np.zeros(env.num_envs, dtype=np.bool_) + self._last_counter = int(task_env.common_step_counter) @property - def obs_groups_spec(self) -> dict[str, int]: - # Playground state: - # linvel(3) + gyro(3) + gravity(3) + diff(12) + dof_vel(12) + last_action(12) = 45. - # UniLab stacks the actor state for short-horizon dynamics; critic appends current privileged tail. - obs_dim = _FOOTSTAND_FRAME_OBS_DIM * self._obs_history_len - return {"obs": obs_dim, "critic": obs_dim + _FOOTSTAND_PRIVILEGED_TAIL_DIM} + def state(self) -> FootstandState: + return self._state @property - def _obs_history_len(self) -> int: - return max(_FOOTSTAND_MIN_OBS_HISTORY_LEN, int(self._cfg.obs_history_len)) - - def _build_playground_reset_randomization( - self, num_reset: int - ) -> ResetRandomizationPayload | None: - domain_rand = self._cfg.domain_rand - payload = ResetRandomizationPayload() - - if domain_rand.randomize_floor_friction: - base_geom_friction = self._base_geom_friction - assert base_geom_friction is not None - low, high = domain_rand.floor_friction_range - geom_friction = np.broadcast_to( - base_geom_friction, (num_reset, *base_geom_friction.shape) - ).copy() - geom_friction[:, self._floor_geom_id, 0] = np.random.uniform( - low, high, size=(num_reset,) - ) - payload.geom_friction = geom_friction - - body_mass = None - if domain_rand.randomize_link_mass: - low, high = domain_rand.link_mass_scale_range - scale = np.random.uniform(low, high, size=(num_reset, self._base_body_mass.size)) - body_mass = self._base_body_mass.reshape(1, -1) * scale - if domain_rand.torso_added_mass_range is not None: - low, high = domain_rand.torso_added_mass_range - if body_mass is None: - body_mass = np.broadcast_to( - self._base_body_mass, (num_reset, self._base_body_mass.size) - ).copy() - body_mass[:, self._base_body_id] += np.random.uniform(low, high, size=(num_reset,)) - if body_mass is not None: - payload.body_mass = body_mass.astype(np.float64, copy=False) - - if domain_rand.randomize_torso_com: - low, high = domain_rand.torso_com_offset_range - body_ipos = np.broadcast_to( - self._base_body_ipos, (num_reset, *self._base_body_ipos.shape) - ).copy() - body_ipos[:, self._base_body_id, :] += np.random.uniform(low, high, size=(num_reset, 3)) - payload.body_ipos = body_ipos - - if domain_rand.randomize_dof_armature: - base_dof_armature = self._base_dof_armature - assert base_dof_armature is not None - low, high = domain_rand.dof_armature_scale_range - dof_armature = np.broadcast_to( - base_dof_armature, (num_reset, base_dof_armature.size) - ).copy() - dof_armature[:, -self._num_action :] *= np.random.uniform( - low, high, size=(num_reset, self._num_action) - ) - payload.dof_armature = dof_armature - - return None if payload.is_empty() else payload - - def _init_reward_functions(self): - self._reward_fns: dict[str, Any] = { - "height": self._reward_height, - "contact": self._cost_contact, - "orientation": self._reward_orientation, - "oritentation": self._reward_orientation, - "action_rate": rewards.action_rate, - "termination": self._reward_termination, - "dof_pos_limits": self._cost_joint_pos_limits, - "torques": self._cost_torques, - "pose": self._cost_pose, - "penalty_contact": self._reward_penalty_contact, - "tar": self._reward_tar, - "rear_feet_contact": self._reward_rear_feet_contact, - "both_rear_feet_contact": self._reward_both_rear_feet_contact, - "rear_foot_slip": self._cost_rear_foot_slip, - "rear_foot_anchor": self._cost_rear_foot_anchor, - "front_feet_air": self._reward_front_feet_air, - "balanced_footstand": self._reward_balanced_footstand, - "rear_leg_symmetry": self._cost_rear_leg_symmetry, - "rear_leg_splay": self._cost_rear_leg_splay, - "front_leg_motion": self._cost_front_leg_motion, - "front_feet_crossing": self._cost_front_leg_crossing, - "front_leg_crossing": self._cost_front_leg_crossing, - "upright_stability": self._cost_upright_stability, - "knee_clearance": self._cost_knee_clearance, - "stay_still": self._cost_stay_still, - "energy": rewards.energy, - "dof_acc": rewards.dof_acc, - } - - def _dof_to_ctrl_order(self, values: np.ndarray) -> np.ndarray: - return np.asarray(values[:, _GO2_DOF_TO_CTRL], dtype=get_global_dtype()) + def last_counter(self) -> int: + return self._last_counter + + def reset(self, env_ids: np.ndarray | slice | None = None) -> None: + self.terminated[_env_ids(self._env, env_ids)] = False + self._last_counter = int(cast("_FootstandEnv", self._env).common_step_counter) + + def __call__(self, env: _FootstandEnv, **params: Any) -> np.ndarray: + del params + state = self._state.snapshot(env) + previous_steps = np.maximum(env.episode_length_buf - 1, 0) + grace_elapsed = previous_steps >= self._grace_steps + low_height = state.height < _TARGET_HEIGHT * self._height_fraction + bad_orientation = state.orientation < self._orientation_threshold + pose_failure = grace_elapsed & (low_height | bad_orientation) + energy = np.sum(np.abs(state.torques) * np.abs(state.joint_vel), axis=1) + energy_failure = energy > self._energy_threshold + upside_down = state.upvector[:, 2] < -0.25 + self.terminated[:] = np.logical_or.reduce( + (state.termination_contact, upside_down, energy_failure, pose_failure) + ) + self._last_counter = int(env.common_step_counter) + return self.terminated - def _ctrl_to_dof_order(self, values: np.ndarray) -> np.ndarray: - return np.asarray(values[:, _GO2_DOF_TO_CTRL], dtype=get_global_dtype()) - def _get_local_gravity(self) -> np.ndarray: - gravity = np.broadcast_to(_WORLD_GRAVITY, (self._num_envs, 3)) - return np.asarray( - np_quat_apply_inverse(self._backend.get_base_quat(), gravity), dtype=get_global_dtype() +def _termination(env: _FootstandEnv, state_term_name: str) -> FootstandTermination: + name = _name("Footstand reward", "state_term_name", state_term_name) + state_term = env.termination_manager.get_term_cfg(name).func + if not isinstance(state_term, FootstandTermination): + raise TypeError( + f"Footstand termination term {name!r} must be FootstandTermination, " + f"got {type(state_term).__name__}" ) - - def _get_body_forward(self) -> np.ndarray: - forward = np.broadcast_to(_BODY_FORWARD, (self._num_envs, 3)) - return np.asarray( - np_quat_apply(self._backend.get_base_quat(), forward), dtype=get_global_dtype() + if state_term.last_counter != int(env.common_step_counter): + raise RuntimeError( + f"Footstand termination state {name!r} was not computed for control step " + f"{env.common_step_counter}" ) + return state_term + + +class FootstandReward(ManagerTermBase): + """Historical positive-clipped reward aggregate backed by one state snapshot.""" + + _REWARD_NAMES = frozenset( + { + "height", + "contact", + "orientation", + "oritentation", + "action_rate", + "termination", + "dof_pos_limits", + "torques", + "pose", + "penalty_contact", + "tar", + "rear_feet_contact", + "both_rear_feet_contact", + "rear_foot_slip", + "rear_foot_anchor", + "front_feet_air", + "balanced_footstand", + "rear_leg_symmetry", + "rear_leg_splay", + "front_leg_motion", + "front_feet_crossing", + "front_leg_crossing", + "upright_stability", + "knee_clearance", + "stay_still", + "energy", + "dof_acc", + } + ) - def _reward_height(self, ctx: RewardContext) -> np.ndarray: - del ctx - error = np.abs(self._z_des - self.torso_height) - return np.asarray(np.exp(-error / 0.1), dtype=get_global_dtype()) - - def _standing_mask(self) -> np.ndarray: - height_ready = self.torso_height >= self._z_des * _FOOTSTAND_STAND_HEIGHT_FRACTION - orientation_ready = self._orientation_score() >= _FOOTSTAND_STAND_ORIENTATION_THRESHOLD - return np.asarray(height_ready & orientation_ready, dtype=get_global_dtype()) - - def apply_action(self, actions: np.ndarray, state: NpEnvState) -> np.ndarray: - clip_actions = float(getattr(self._cfg.control_config, "clip_actions", np.inf)) - actions_np = np.asarray(actions, dtype=get_global_dtype()) - if np.isfinite(clip_actions): - actions_np = np.clip(actions_np, -clip_actions, clip_actions) - - state.info["last_actions"] = state.info.get("current_actions", np.zeros_like(actions_np)) - state.info["current_actions"] = actions_np - exec_actions = ( - state.info["last_actions"] - if self._cfg.control_config.simulate_action_latency - else actions_np + def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv): + super().__init__(env) + task_env = cast("_FootstandEnv", env) + term = type(self).__name__ + allowed = { + "state_term_name", + "scales", + "soft_joint_pos_limit_factor", + "knee_height_target", + "front_feet_min_separation", + "front_feet_side_margin", + "rear_hip_abduction_margin", + "rear_foot_slip_deadband", + "rear_foot_anchor_radius", + } + unknown = sorted(set(cfg.params) - allowed) + if unknown: + raise TypeError(f"{term} received unsupported parameters: {unknown}") + self._state_term_name = _name(term, "state_term_name", cfg.params.get("state_term_name")) + scales = cfg.params.get("scales") + if not isinstance(scales, dict) or not scales: + raise TypeError(f"{term} scales must be a non-empty mapping") + unknown_rewards = sorted(set(scales) - self._REWARD_NAMES) + if unknown_rewards: + raise ValueError(f"{term} scales contains unknown rewards: {unknown_rewards}") + self._scales = { + name: _real(term, f"scales.{name}", value) for name, value in scales.items() + } + self._soft_limit_factor = _real( + term, + "soft_joint_pos_limit_factor", + cfg.params.get("soft_joint_pos_limit_factor"), + minimum=0.0, ) - self._motor_targets += exec_actions * self._cfg.control_config.action_scale - self._clip_motor_targets() - return np.asarray(self._motor_targets, dtype=get_global_dtype()) - - def update_state(self, state: NpEnvState) -> NpEnvState: - linvel = self.get_local_linvel() - gyro = self.get_gyro() - upvector = self._backend.get_sensor_data("upvector") - gravity = self._get_local_gravity() - accelerometer = self._backend.get_sensor_data(self._cfg.sensor.accelerometer) - global_angvel = self._backend.get_sensor_data(self._cfg.sensor.global_angvel) - dof_pos = self.get_dof_pos() - dof_vel = self.get_dof_vel() - self.feet_force = self._backend.get_sensor_data_batch(self._cfg.sensor.feet_force).reshape( - self._num_envs, len(self._cfg.sensor.feet_force), 1 + self._knee_height_target = _real( + term, "knee_height_target", cfg.params.get("knee_height_target"), minimum=0.0 ) - self.feet_pos = self._backend.get_sensor_data_batch(self._cfg.sensor.feet_pos).reshape( - self._num_envs, len(self._cfg.sensor.feet_pos), 3 + self._front_min_separation = _real( + term, + "front_feet_min_separation", + cfg.params.get("front_feet_min_separation"), + minimum=0.0, ) - self.torso_height = self._backend.get_sensor_data(self._cfg.sensor.global_pos)[:, -1] - self._update_rear_foot_slip() - self._update_rear_foot_anchor() - self._tracked_body_pos = self._backend.get_body_pos_w(self._tracked_body_ids) - result = self._backend.get_sensor_data_batch(self._cfg.sensor.ternamate_contact) - - state.info["qacc"] = self._estimate_dof_acc(dof_vel) - state.info["torques"] = self._estimate_pd_torques(state.info, dof_pos, dof_vel) - orientation = self._orientation_score() - step_count = state.info.get("steps", np.zeros((self._num_envs,), dtype=np.uint32)) - grace_elapsed = step_count >= self._cfg.termination_grace_steps - terminated_z = upvector[:, 2] < -0.25 - terminated_contact = np.any(result, axis=1) - terminated_low_height = ( - self.torso_height < self._z_des * self._cfg.termination_height_fraction + self._front_side_margin = _real( + term, + "front_feet_side_margin", + cfg.params.get("front_feet_side_margin"), + minimum=0.0, ) - terminated_bad_orientation = orientation < self._cfg.termination_orientation_threshold - terminated_pose = grace_elapsed & (terminated_low_height | terminated_bad_orientation) - energy = np.sum(np.abs(state.info["torques"]) * np.abs(dof_vel), axis=1) - terminated_energy = energy > self._cfg.energy_termination_threshold - terminated = np.logical_or.reduce( - (terminated_contact, terminated_z, terminated_energy, terminated_pose) + self._rear_hip_margin = _real( + term, + "rear_hip_abduction_margin", + cfg.params.get("rear_hip_abduction_margin"), + minimum=0.0, ) - self._last_terminated = terminated.copy() - reward = self._compute_reward(state.info, linvel, gyro, dof_pos, dof_vel) - self._last_feet_pos = self.feet_pos.copy() - obs = self._compute_obs( - state.info, - linvel, - gyro, - gravity, - dof_pos, - dof_vel, - self.torso_height.reshape(-1, 1), - accelerometer, - global_angvel, + self._rear_slip_deadband = _real( + term, + "rear_foot_slip_deadband", + cfg.params.get("rear_foot_slip_deadband"), + minimum=0.0, ) - return state.replace(obs=obs, reward=reward, terminated=terminated) - - def _compute_reward( - self, - info: dict, - linvel: np.ndarray, - gyro: np.ndarray, - dof_pos: np.ndarray, - dof_vel: np.ndarray | None = None, - ) -> np.ndarray: - dtype = get_global_dtype() - reward = np.zeros((self._num_envs,), dtype=dtype) - cfg = self._reward_cfg - if dof_vel is None: - dof_vel = self.get_dof_vel() - - ctx = RewardContext( - info=info, - linvel=linvel, - gyro=gyro, - dof_pos=dof_pos, - dof_vel=dof_vel, - num_envs=self._num_envs, - default_angles=self.default_angles, - tracking_sigma=cfg.tracking_sigma, - base_height_target=cfg.base_height_target, - base_height=self._backend.get_base_pos()[:, 2], + self._rear_anchor_radius = _real( + term, + "rear_foot_anchor_radius", + cfg.params.get("rear_foot_anchor_radius"), + minimum=0.0, + strict_minimum=True, ) + state_term = _termination(task_env, self._state_term_name) + action = state_term.state.action + centers = (action.joint_lower + action.joint_upper) / 2.0 + widths = action.joint_upper - action.joint_lower + self._soft_lower = centers - 0.5 * widths * self._soft_limit_factor + self._soft_upper = centers + 0.5 * widths * self._soft_limit_factor - step_count = info.get("steps", np.zeros((self._num_envs,), dtype=np.uint32)) - should_log = self._enable_reward_log and (int(step_count[0]) % 4 == 0) - log = {} if should_log else info.get("log", {}) - - for name, scale in cfg.scales.items(): - if scale == 0 or name not in self._reward_fns: - continue - rew = self._reward_fns[name](ctx) - weighted_rew = rew * scale - reward += weighted_rew - if should_log: - log[f"reward/{name}"] = float(np.mean(weighted_rew)) - - info["log"] = log - return np.clip(reward * self._cfg.ctrl_dt, 0.0, 10000.0) + @staticmethod + def _standing(state: FootstandState) -> np.ndarray: + return ( + (state.height >= _TARGET_HEIGHT * _STAND_HEIGHT_FRACTION) + & (state.orientation >= _STAND_ORIENTATION_THRESHOLD) + ).astype(get_global_dtype(), copy=False) - def _compute_obs( + def _value( self, - info: dict, - linvel: np.ndarray, - gyro: np.ndarray, - gravity: np.ndarray, - dof_pos: np.ndarray, - dof_vel: np.ndarray, - height: np.ndarray, - accelerometer: np.ndarray | None = None, - global_angvel: np.ndarray | None = None, - env_ids: np.ndarray | None = None, - ) -> dict[str, np.ndarray]: - noise_cfg = self._cfg.noise_config - diff = dof_pos - self.default_angles - noisy_linvel = self._obs_noise(linvel, noise_cfg.scale_linvel) - noisy_gyro = self._obs_noise(gyro, noise_cfg.scale_gyro) - noisy_gravity = self._obs_noise(gravity, noise_cfg.scale_gravity) - noisy_diff = self._obs_noise(diff, noise_cfg.scale_joint_angle) - noisy_dof_vel = self._obs_noise(dof_vel, noise_cfg.scale_joint_vel) - last_actions = info.get("last_actions", np.zeros_like(diff)) - - obs_frame = np.concatenate( - [noisy_linvel, noisy_gyro, noisy_gravity, noisy_diff, noisy_dof_vel, last_actions], + name: str, + state_term: FootstandTermination, + state: FootstandState, + ) -> np.ndarray: + dtype = get_global_dtype() + standing = self._standing(state) + default = state.default_joint_pos + if name == "height": + return np.asarray(np.exp(-np.abs(_TARGET_HEIGHT - state.height) / 0.1), dtype=dtype) + if name == "contact": + return np.any(state.foot_contact[:, _FRONT_FEET], axis=1).astype(dtype) + if name in ("orientation", "oritentation"): + return state.orientation + if name == "action_rate": + action = state.action + return np.sum(np.square(action.raw_action - action.previous_raw_action), axis=1) + if name == "termination": + return state_term.terminated.astype(dtype) + if name == "dof_pos_limits": + below = np.clip(self._soft_lower - state.joint_pos, 0.0, None) + above = np.clip(state.joint_pos - self._soft_upper, 0.0, None) + return np.sum(below + above, axis=1) + if name == "torques": + return np.sum(np.square(state.torques), axis=1) + if name == "pose": + return np.sum( + np.square(state.joint_pos[:, _REAR_LEGS] - default[:, _REAR_LEGS]), axis=1 + ) + if name == "penalty_contact": + return state.penalty_contact.astype(dtype) + if name == "tar": + error = np.sum(np.square(state.joint_pos[:, _FRONT_LEGS] - _FRONT_LEG_TARGET), axis=1) + height_mask = (state.height >= _TARGET_HEIGHT * _STAND_HEIGHT_FRACTION).astype(dtype) + return np.asarray(np.exp(-error) * height_mask, dtype=dtype) + if name == "rear_feet_contact": + return np.mean(state.foot_contact[:, _REAR_FEET], axis=1, dtype=dtype) + if name == "both_rear_feet_contact": + return np.all(state.foot_contact[:, _REAR_FEET], axis=1).astype(dtype) + if name == "rear_foot_slip": + slip = np.square(np.clip(state.rear_speed - self._rear_slip_deadband, 0.0, None)) + slip *= state.foot_contact[:, _REAR_FEET] + return np.mean(slip, axis=1, dtype=dtype) + if name == "rear_foot_anchor": + drift = np.square( + np.clip(state.rear_anchor_drift - self._rear_anchor_radius, 0.0, None) + / self._rear_anchor_radius + ) + drift *= state.rear_anchor_contact + return np.mean(drift, axis=1, dtype=dtype) + if name == "front_feet_air": + return (~np.any(state.foot_contact[:, _FRONT_FEET], axis=1)).astype(dtype) + if name == "balanced_footstand": + support = np.all(state.foot_contact[:, _REAR_FEET], axis=1) + support &= ~np.any(state.foot_contact[:, _FRONT_FEET], axis=1) + return support.astype(dtype) * standing + if name == "rear_leg_symmetry": + mirrored = state.joint_pos[:, _REAR_RIGHT] * _REAR_MIRROR_SIGNS + cost = np.mean(np.square(state.joint_pos[:, _REAR_LEFT] - mirrored), axis=1) + return cost * (1.0 - standing) + if name == "rear_leg_splay": + error = state.joint_pos[:, _REAR_HIPS] - default[:, _REAR_HIPS] + splay = np.clip(np.abs(error) - self._rear_hip_margin, 0.0, None) + return np.mean(np.square(splay), axis=1) * standing + if name == "front_leg_motion": + return np.mean(np.square(state.joint_vel[:, _FRONT_LEGS]), axis=1) * standing + if name in ("front_feet_crossing", "front_leg_crossing"): + return self._front_crossing(state) + if name == "upright_stability": + cost = np.sum(np.square(state.root_linvel_w), axis=1) + cost += 0.25 * np.sum(np.square(state.root_angvel_w), axis=1) + return cost * standing + if name == "knee_clearance": + target = max(self._knee_height_target, 1.0e-6) + height = state.tracked_body_pos[:, _KNEE_BODY_INDICES, 2] + return np.mean(np.square(np.clip(target - height, 0.0, None) / target), axis=1) + if name == "stay_still": + return np.sum(np.square(state.root_linvel_w[:, :2]), axis=1) + np.square( + state.root_angvel_w[:, 2] + ) + if name == "energy": + return np.sum(np.abs(state.joint_vel) * np.abs(state.torques), axis=1) + if name == "dof_acc": + return np.sum(np.square(state.qacc), axis=1) + raise RuntimeError(f"Footstand reward dispatch is incomplete for {name!r}") + + def _front_crossing(self, state: FootstandState) -> np.ndarray: + left = np.concatenate( + ( + state.foot_pos[:, [0], :], + state.tracked_body_pos[:, _FRONT_LEFT_BODY_INDICES, :], + ), axis=1, - dtype=get_global_dtype(), ) - critic_frame = np.concatenate( - [linvel, gyro, gravity, diff, dof_vel, last_actions], + right = np.concatenate( + ( + state.foot_pos[:, [1], :], + state.tracked_body_pos[:, _FRONT_RIGHT_BODY_INDICES, :], + ), axis=1, - dtype=get_global_dtype(), ) - obs = self._update_obs_history(obs_frame, env_ids=env_ids, history_attr="_obs_history") - critic_obs = self._update_obs_history( - critic_frame, - env_ids=env_ids, - history_attr="_critic_obs_history", + points = np.concatenate((left, right), axis=1) + relative = (points - state.root_pos[:, None, :]).reshape(-1, 3) + quaternions = np.repeat(state.root_quat, points.shape[1], axis=0) + body_points = np_quat_apply_inverse(quaternions, relative).reshape( + state.root_pos.shape[0], points.shape[1], 3 ) - torques = np.asarray(info.get("torques", np.zeros_like(dof_pos)), dtype=get_global_dtype()) - if accelerometer is None: - accelerometer = np.zeros_like(gyro) - if global_angvel is None: - global_angvel = np.zeros_like(gyro) - critic = np.concatenate( - [ - critic_obs, - gyro, - accelerometer, - linvel, - global_angvel, - dof_pos, - dof_vel, - torques, - height, - ], - axis=1, - dtype=get_global_dtype(), + left_y = body_points[:, : left.shape[1], 1] + right_y = body_points[:, left.shape[1] :, 1] + left_error = np.clip(self._front_side_margin - left_y, 0.0, None) + right_error = np.clip(right_y + self._front_side_margin, 0.0, None) + separation_error = np.clip(self._front_min_separation - (left_y - right_y), 0.0, None) + return np.mean( + np.square(left_error) + np.square(right_error) + np.square(separation_error), axis=1 ) - return {"obs": obs, "critic": critic} - def _update_obs_history( - self, - frame_obs: np.ndarray, - *, - env_ids: np.ndarray | None = None, - history_attr: str = "_obs_history", - ) -> np.ndarray: - frame_obs = np.asarray(frame_obs, dtype=get_global_dtype()) - batch_size = int(frame_obs.shape[0]) - history_len = self._obs_history_len - expected_shape = (self._num_envs, history_len, _FOOTSTAND_FRAME_OBS_DIM) - history = getattr(self, history_attr, None) - if history is None or history.shape != expected_shape: - if env_ids is None and batch_size == self._num_envs: - history = np.zeros(expected_shape, dtype=get_global_dtype()) - setattr(self, history_attr, history) - elif env_ids is not None: - history = np.zeros(expected_shape, dtype=get_global_dtype()) - setattr(self, history_attr, history) - else: - repeated = np.broadcast_to( - frame_obs[:, None, :], (batch_size, history_len, frame_obs.shape[1]) - ).copy() - return np.asarray(repeated.reshape(batch_size, -1), dtype=get_global_dtype()) - - assert history is not None - if env_ids is None: - if batch_size != self._num_envs: - repeated = np.broadcast_to( - frame_obs[:, None, :], (batch_size, history_len, frame_obs.shape[1]) - ).copy() - return np.asarray(repeated.reshape(batch_size, -1), dtype=get_global_dtype()) - history[:, :-1] = history[:, 1:] - history[:, -1] = frame_obs - selected_history = history - else: - env_ids = np.asarray(env_ids, dtype=np.int32) - selected_history = np.broadcast_to( - frame_obs[:, None, :], (batch_size, history_len, frame_obs.shape[1]) - ).copy() - history[env_ids] = selected_history - - return np.asarray(selected_history.reshape(batch_size, -1), dtype=get_global_dtype()) - - def _init_soft_joint_limits(self) -> None: - joint_range = self._backend.get_joint_range() - if joint_range is None: - self._soft_lowers = np.full((self._num_action,), -np.inf, dtype=np.float32) - self._soft_uppers = np.full((self._num_action,), np.inf, dtype=np.float32) - return - - joint_range = np.asarray(joint_range, dtype=np.float32) - centers = (joint_range[:, 0] + joint_range[:, 1]) / 2.0 - widths = joint_range[:, 1] - joint_range[:, 0] - factor = self._cfg.soft_joint_pos_limit_factor - self._soft_lowers = centers - 0.5 * widths * factor - self._soft_uppers = centers + 0.5 * widths * factor - - def _init_motor_target_limits(self) -> None: - joint_range = self._backend.get_joint_range() - if joint_range is None: - self._target_lowers = np.full((self._num_action,), -np.inf, dtype=get_global_dtype()) - self._target_uppers = np.full((self._num_action,), np.inf, dtype=get_global_dtype()) - return - - joint_range = np.asarray(joint_range, dtype=get_global_dtype()) - lowers = joint_range[:, 0] - uppers = joint_range[:, 1] - if lowers.size == _GO2_DOF_TO_CTRL.size: - lowers = self._dof_to_ctrl_order(lowers.reshape(1, -1))[0] - uppers = self._dof_to_ctrl_order(uppers.reshape(1, -1))[0] - self._target_lowers = np.asarray(lowers, dtype=get_global_dtype()) - self._target_uppers = np.asarray(uppers, dtype=get_global_dtype()) - - def _clip_motor_targets(self) -> None: - lowers = getattr(self, "_target_lowers", None) - uppers = getattr(self, "_target_uppers", None) - if lowers is None or uppers is None: - return - np.clip(self._motor_targets, lowers, uppers, out=self._motor_targets) - - def _reward_termination(self, ctx: RewardContext) -> np.ndarray: - return self._last_terminated.astype(get_global_dtype()) - - def _cost_joint_pos_limits(self, ctx: RewardContext) -> np.ndarray: - out_of_limits = -np.clip(ctx.dof_pos - self._soft_lowers, None, 0.0) - out_of_limits += np.clip(ctx.dof_pos - self._soft_uppers, 0.0, None) - return cast(np.ndarray, np.sum(out_of_limits, axis=1)) - - def _cost_stay_still(self, ctx: RewardContext) -> np.ndarray: - linvel = self._backend.get_base_lin_vel() - angvel = self._backend.get_base_ang_vel() - return cast(np.ndarray, np.sum(np.square(linvel[:, :2]), axis=1) + np.square(angvel[:, 2])) - - def _cost_torques(self, ctx: RewardContext) -> np.ndarray: - torques = np.asarray(ctx.info.get("torques", np.zeros_like(ctx.dof_pos))) - return cast(np.ndarray, np.sum(np.square(torques), axis=1)) - - def _estimate_dof_acc(self, dof_vel: np.ndarray) -> np.ndarray: - qacc = np.asarray((dof_vel - self._last_dof_vel_for_acc) / self._cfg.ctrl_dt) - self._last_dof_vel_for_acc = np.asarray(dof_vel, dtype=np.float32).copy() - return np.asarray(qacc, dtype=get_global_dtype()) - - def _estimate_pd_torques( - self, info: dict, dof_pos: np.ndarray, dof_vel: np.ndarray - ) -> np.ndarray: - del info - targets = self._ctrl_to_dof_order(self._motor_targets) - torques = self._cfg.control_config.Kp * (targets - dof_pos) - torques -= self._cfg.control_config.Kd * dof_vel - return np.asarray(torques, dtype=get_global_dtype()) - - def _reward_orientation(self, ctx: RewardContext) -> np.ndarray: - del ctx - return self._orientation_score() - - def _orientation_score(self) -> np.ndarray: - forward = self._get_body_forward() - cos_dist = forward @ self._desired_forward_vec - normalized = 0.5 * cos_dist + 0.5 - return np.asarray(np.square(normalized), dtype=get_global_dtype()) - - def _cost_contact(self, ctx: RewardContext) -> np.ndarray: - del ctx - feet_contact = self.feet_force[:, self.feet_geom_names, 0] > _FOOTSTAND_CONTACT_THRESHOLD - return np.asarray(np.any(feet_contact, axis=1), dtype=get_global_dtype()) - - def _reward_rear_feet_contact(self, ctx: RewardContext) -> np.ndarray: - del ctx - rear_contact = self.feet_force[:, _FOOTSTAND_REAR_FEET, 0] > _FOOTSTAND_CONTACT_THRESHOLD - return np.asarray(np.mean(rear_contact, axis=1), dtype=get_global_dtype()) - - def _reward_both_rear_feet_contact(self, ctx: RewardContext) -> np.ndarray: - del ctx - rear_contact = self.feet_force[:, _FOOTSTAND_REAR_FEET, 0] > _FOOTSTAND_CONTACT_THRESHOLD - return np.asarray(np.all(rear_contact, axis=1), dtype=get_global_dtype()) - - def _update_rear_foot_slip(self) -> None: - if not hasattr(self, "_last_feet_pos"): - self._last_feet_pos = self.feet_pos.copy() - rear_contact = self.feet_force[:, _FOOTSTAND_REAR_FEET, 0] > _FOOTSTAND_CONTACT_THRESHOLD - rear_delta_xy = ( - self.feet_pos[:, _FOOTSTAND_REAR_FEET, :2] - - self._last_feet_pos[:, _FOOTSTAND_REAR_FEET, :2] - ) - rear_vel_xy = rear_delta_xy / max(float(self._cfg.ctrl_dt), 1e-6) - speed = np.linalg.norm(rear_vel_xy, axis=2) - deadband = float(self._reward_cfg.rear_foot_slip_deadband) - slip = np.square(np.clip(speed - deadband, 0.0, None)) * rear_contact - self._rear_foot_slip = np.asarray(np.mean(slip, axis=1), dtype=get_global_dtype()) - - def _cost_rear_foot_slip(self, ctx: RewardContext) -> np.ndarray: - del ctx - return np.asarray(self._rear_foot_slip, dtype=get_global_dtype()) - - def _update_rear_foot_anchor(self) -> None: - if not hasattr(self, "_rear_foot_anchor_pos"): - self._rear_foot_anchor_pos = self.feet_pos[:, _FOOTSTAND_REAR_FEET, :2].copy() - self._rear_foot_anchor_contact = np.zeros( - (self._num_envs, len(_FOOTSTAND_REAR_FEET)), dtype=bool + def __call__(self, env: _FootstandEnv, **params: Any) -> np.ndarray: + del params + state_term = _termination(env, self._state_term_name) + state = state_term.state + reward = np.zeros((env.num_envs,), dtype=get_global_dtype()) + for name, scale in self._scales.items(): + if scale != 0.0: + reward += scale * self._value(name, state_term, state) + max_rate = 10000.0 / env.step_dt + return np.clip(reward, 0.0, max_rate) + + +class FootstandJointReset(ManagerTermBase): + """Reset all Go2 joints to the home pose plus a uniform offset.""" + + def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv): + super().__init__(env) + term = type(self).__name__ + if set(cfg.params) != {"asset_cfg", "position_offset_range"}: + raise ValueError( + f"{term} requires exactly asset_cfg and position_offset_range parameters" ) - rear_contact = self.feet_force[:, _FOOTSTAND_REAR_FEET, 0] > _FOOTSTAND_CONTACT_THRESHOLD - standing = self._standing_mask().astype(bool)[:, None] - anchor_contact = rear_contact & standing - rear_xy = self.feet_pos[:, _FOOTSTAND_REAR_FEET, :2] - new_contact = anchor_contact & ~self._rear_foot_anchor_contact - self._rear_foot_anchor_pos[new_contact] = rear_xy[new_contact] - self._rear_foot_anchor_contact = anchor_contact.copy() - - drift = np.linalg.norm(rear_xy - self._rear_foot_anchor_pos, axis=2) - radius = max(float(self._reward_cfg.rear_foot_anchor_radius), 1e-6) - anchored_drift = np.square(np.clip(drift - radius, 0.0, None) / radius) * anchor_contact - self._rear_foot_anchor = np.asarray( - np.mean(anchored_drift, axis=1), dtype=get_global_dtype() + asset_cfg = cfg.params["asset_cfg"] + if not isinstance(asset_cfg, SceneEntityCfg): + raise TypeError(f"{term} asset_cfg must be SceneEntityCfg") + self._entity = cast("Entity", env.scene[asset_cfg.name]) + self._joint_ids = asset_cfg.joint_ids + selected = self._entity.data.default_joint_pos[:, self._joint_ids] + if selected.shape != (env.num_envs, NUM_ACTIONS): + raise ValueError(f"{term} requires exactly {NUM_ACTIONS} selected joints") + self._range = _pair(term, "position_offset_range", cfg.params["position_offset_range"]) + + def __call__( + self, + env: ManagerBasedRlEnv, + env_ids: np.ndarray | None, + **params: Any, + ) -> None: + del params + ids = _env_ids(env, env_ids) + position = np.array(self._entity.data.default_joint_pos[ids][:, self._joint_ids], copy=True) + position += env.rng.uniform(*self._range, size=position.shape) + velocity = np.array(self._entity.data.default_joint_vel[ids][:, self._joint_ids], copy=True) + self._entity.write_joint_state_to_sim( + np.asarray(position, dtype=get_global_dtype()), + np.asarray(velocity, dtype=get_global_dtype()), + joint_ids=self._joint_ids, + env_ids=ids, ) - def _cost_rear_foot_anchor(self, ctx: RewardContext) -> np.ndarray: - del ctx - return np.asarray(self._rear_foot_anchor, dtype=get_global_dtype()) - - def _reward_front_feet_air(self, ctx: RewardContext) -> np.ndarray: - del ctx - front_contact = self.feet_force[:, _FOOTSTAND_FRONT_FEET, 0] > _FOOTSTAND_CONTACT_THRESHOLD - return np.asarray(~np.any(front_contact, axis=1), dtype=get_global_dtype()) - - def _reward_balanced_footstand(self, ctx: RewardContext) -> np.ndarray: - del ctx - rear_contact = self.feet_force[:, _FOOTSTAND_REAR_FEET, 0] > _FOOTSTAND_CONTACT_THRESHOLD - front_contact = self.feet_force[:, _FOOTSTAND_FRONT_FEET, 0] > _FOOTSTAND_CONTACT_THRESHOLD - valid_support = np.all(rear_contact, axis=1) & ~np.any(front_contact, axis=1) - return np.asarray(valid_support * self._standing_mask(), dtype=get_global_dtype()) - - def _cost_rear_leg_symmetry(self, ctx: RewardContext) -> np.ndarray: - rear_left = ctx.dof_pos[:, _FOOTSTAND_REAR_LEFT_LEG_IDS] - rear_right = ctx.dof_pos[:, _FOOTSTAND_REAR_RIGHT_LEG_IDS] - mirrored_right = rear_right * _FOOTSTAND_REAR_MIRROR_SIGNS - cost = np.mean(np.square(rear_left - mirrored_right), axis=1) - rising_mask = 1.0 - self._standing_mask() - return np.asarray(cost * rising_mask, dtype=get_global_dtype()) - - def _cost_rear_leg_splay(self, ctx: RewardContext) -> np.ndarray: - default_angles = np.asarray(ctx.default_angles, dtype=get_global_dtype()).reshape(-1) - rear_hip_error = ( - ctx.dof_pos[:, _FOOTSTAND_REAR_HIP_IDS] - default_angles[_FOOTSTAND_REAR_HIP_IDS] - ) - margin = float(self._reward_cfg.rear_hip_abduction_margin) - splay = np.clip(np.abs(rear_hip_error) - margin, 0.0, None) - cost = np.mean(np.square(splay), axis=1) - return np.asarray(cost * self._standing_mask(), dtype=get_global_dtype()) - - def _cost_front_leg_motion(self, ctx: RewardContext) -> np.ndarray: - assert ctx.dof_vel is not None - front_leg_vel = ctx.dof_vel[:, _FOOTSTAND_FRONT_LEG_IDS] - cost = np.mean(np.square(front_leg_vel), axis=1) - return np.asarray(cost * self._standing_mask(), dtype=get_global_dtype()) - - def _cost_front_leg_crossing(self, ctx: RewardContext) -> np.ndarray: - del ctx - base_pos = self._backend.get_base_pos() - base_quat = self._backend.get_base_quat() - left_points = np.concatenate( - [ - self.feet_pos[:, [0], :], - self._tracked_body_pos[:, _FOOTSTAND_FRONT_LEFT_BODY_INDICES, :], - ], - axis=1, + +class FootstandMassRandomization(ManagerTermBase): + """Compose all-link mass scaling and torso additive mass in one reset write.""" + + def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv): + super().__init__(env) + term = type(self).__name__ + allowed = { + "asset_cfg", + "torso_body_name", + "link_mass_scale_range", + "torso_added_mass_range", + } + if set(cfg.params) != allowed: + raise ValueError(f"{term} requires parameters {sorted(allowed)}") + asset_cfg = cfg.params["asset_cfg"] + if not isinstance(asset_cfg, SceneEntityCfg): + raise TypeError(f"{term} asset_cfg must be SceneEntityCfg") + self._entity = cast("Entity", env.scene[asset_cfg.name]) + self._body_ids, self._default_mass = self._entity.bind_body_mass_write( + asset_cfg.body_ids, term_name="footstand_mass" ) - right_points = np.concatenate( - [ - self.feet_pos[:, [1], :], - self._tracked_body_pos[:, _FOOTSTAND_FRONT_RIGHT_BODY_INDICES, :], - ], - axis=1, + torso_name = _name(term, "torso_body_name", cfg.params["torso_body_name"]) + torso_ids, _ = self._entity.find_bodies((torso_name,)) + if len(torso_ids) != 1: + raise ValueError(f"{term} torso_body_name must resolve exactly one body") + selected = np.flatnonzero(self._body_ids == torso_ids[0]) + if selected.size != 1: + raise ValueError(f"{term} torso body must be included in asset_cfg") + self._torso_index = int(selected[0]) + self._scale_range = _pair( + term, "link_mass_scale_range", cfg.params["link_mass_scale_range"], minimum=0.0 ) - points = np.concatenate([left_points, right_points], axis=1) - rel_points = (points - base_pos[:, None, :]).reshape(-1, 3) - rel_quat = np.repeat(base_quat, points.shape[1], axis=0) - body_points = np_quat_apply_inverse(rel_quat, rel_points).reshape( - self._num_envs, points.shape[1], 3 + self._added_range = _pair( + term, "torso_added_mass_range", cfg.params["torso_added_mass_range"] ) - left_y = body_points[:, : left_points.shape[1], 1] - right_y = body_points[:, left_points.shape[1] :, 1] - side_margin = float(self._reward_cfg.front_feet_side_margin) - min_separation = float(self._reward_cfg.front_feet_min_separation) - left_error = np.clip(side_margin - left_y, 0.0, None) - right_error = np.clip(right_y + side_margin, 0.0, None) - separation_error = np.clip(min_separation - (left_y - right_y), 0.0, None) - cost = np.mean( - np.square(left_error) + np.square(right_error) + np.square(separation_error), - axis=1, + + def __call__( + self, + env: ManagerBasedRlEnv, + env_ids: np.ndarray | None, + **params: Any, + ) -> None: + del params + ids = _env_ids(env, env_ids) + scale = env.rng.uniform(*self._scale_range, size=(ids.size, self._default_mass.size)) + mass = self._default_mass[None, :] * scale + mass[:, self._torso_index] += env.rng.uniform(*self._added_range, size=ids.size) + if np.any(mass <= 0.0): + raise ValueError("FootstandMassRandomization produced a non-positive body mass") + self._entity.write_body_mass_to_sim( + mass, + body_ids=self._body_ids, + env_ids=ids, + term_name="footstand_mass", ) - return np.asarray(cost, dtype=get_global_dtype()) - - def _cost_upright_stability(self, ctx: RewardContext) -> np.ndarray: - del ctx - linvel = self._backend.get_base_lin_vel() - angvel = self._backend.get_base_ang_vel() - cost = np.sum(np.square(linvel), axis=1) + 0.25 * np.sum(np.square(angvel), axis=1) - return np.asarray(cost * self._standing_mask(), dtype=get_global_dtype()) - - def _cost_knee_clearance(self, ctx: RewardContext) -> np.ndarray: - del ctx - target = max(float(self._reward_cfg.knee_height_target), 1e-6) - knee_height = self._tracked_body_pos[:, _FOOTSTAND_KNEE_BODY_INDICES, 2] - clearance_error = np.clip(target - knee_height, 0.0, None) / target - return np.asarray(np.mean(np.square(clearance_error), axis=1), dtype=get_global_dtype()) + + +registry.register_env_config("Go2FootStand", ManagerBasedRlEnvCfg) +registry.register_env("Go2FootStand", make_manager_based_rl_env, sim_backend="mujoco") +registry.register_env("Go2FootStand", make_manager_based_rl_env, sim_backend="motrix") +registry.register_env("Go2FootStand", make_manager_based_rl_env, sim_backend="drake") + + +__all__ = [ + "FRAME_OBS_DIM", + "NUM_ACTIONS", + "PRIVILEGED_OBS_DIM", + "FootstandIncrementalAction", + "FootstandIncrementalActionCfg", + "FootstandJointReset", + "FootstandMassRandomization", + "FootstandReward", + "FootstandState", + "FootstandTermination", + "frame_observation", + "privileged_observation", +] diff --git a/tests/config/test_locomotion_params.py b/tests/config/test_locomotion_params.py index b7cae7a2e..62a03f468 100644 --- a/tests/config/test_locomotion_params.py +++ b/tests/config/test_locomotion_params.py @@ -487,7 +487,7 @@ def test_ppo_go2_num_envs(): assert cfg.algo.max_iterations == 151 -def test_ppo_go2_footstand_uses_teacher_linvel_task(): +def test_ppo_go2_footstand_uses_hydra_owned_manager_task(): from hydra import compose, initialize_config_dir from hydra.core.global_hydra import GlobalHydra @@ -497,19 +497,26 @@ def test_ppo_go2_footstand_uses_teacher_linvel_task(): assert cfg.training.task_name == "Go2FootStand" assert cfg.training.sim_backend == "mujoco" - assert cfg.env.add_body_sensors is True - assert cfg.env.obs_history_len == 15 - assert cfg.env.energy_termination_threshold == pytest.approx(200.0) - assert cfg.reward.scales.energy == pytest.approx(-0.003) - assert cfg.reward.scales.dof_acc == pytest.approx(-2.5e-7) - assert cfg.reward.scales.rear_leg_symmetry == pytest.approx(-0.2) - assert cfg.reward.scales.knee_clearance == pytest.approx(-0.5) - assert cfg.reward.knee_height_target == pytest.approx(0.08) - assert cfg.env.domain_rand.randomize_floor_friction is True - assert cfg.env.domain_rand.randomize_link_mass is True - assert cfg.env.domain_rand.randomize_torso_com is True - assert cfg.env.domain_rand.randomize_dof_armature is True - assert cfg.env.domain_rand.randomize_reset_joint_qpos is True + assert cfg.algo.num_envs == 4096 + assert cfg.env.sim_dt == pytest.approx(0.004) + assert cfg.env.ctrl_dt == pytest.approx(0.02) + assert cfg.env.max_episode_seconds == pytest.approx(10.0) + assert cfg.env.adaptive_chunk_size is False + assert cfg.env.observations.policy.terms.frame.history_length == 15 + assert cfg.env.observations.critic.terms.frame.history_length == 15 + assert cfg.env.actions.joint_pos.action_scale == pytest.approx(0.3) + assert cfg.env.actions.joint_pos.clip_actions == pytest.approx(1.0) + assert cfg.env.terminations.footstand.params.energy_threshold == pytest.approx(200.0) + assert cfg.reward.footstand.params.scales.energy == pytest.approx(-0.003) + assert cfg.reward.footstand.params.scales.dof_acc == pytest.approx(-2.5e-7) + assert cfg.reward.footstand.params.scales.rear_leg_symmetry == pytest.approx(-0.2) + assert cfg.reward.footstand.params.scales.knee_clearance == pytest.approx(-0.5) + assert cfg.reward.footstand.params.knee_height_target == pytest.approx(0.08) + assert cfg.env.events.floor_friction is not None + assert cfg.env.events.link_mass is not None + assert cfg.env.events.torso_com is not None + assert cfg.env.events.joint_armature is not None + assert cfg.env.events.reset_joints is not None def test_ppo_g1_motion_tracking(): diff --git a/tests/envs/locomotion/test_go2_footstand.py b/tests/envs/locomotion/test_go2_footstand.py index f58b9e2ee..dd404db1c 100644 --- a/tests/envs/locomotion/test_go2_footstand.py +++ b/tests/envs/locomotion/test_go2_footstand.py @@ -1,957 +1,438 @@ +"""Hydra-owned Manager-Based contract for Go2 footstand.""" + from __future__ import annotations +import importlib.util +from collections.abc import Mapping, Sequence +from dataclasses import fields, is_dataclass +from pathlib import Path +from typing import Any + import numpy as np import pytest +from hydra import compose, initialize_config_dir +from hydra.core.global_hydra import GlobalHydra +from omegaconf import DictConfig, OmegaConf from unilab.base import registry -from unilab.base.np_env import NpEnvState -from unilab.base.registry import ensure_registries -from unilab.dr import ResetRandomizationPayload -from unilab.tasks.locomotion.common.rewards import RewardContext +from unilab.base.config_materialization import apply_cfg_overrides +from unilab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg, make_manager_based_rl_env +from unilab.managers._noise import UniformNoiseCfg from unilab.tasks.locomotion.go2.footstand import ( - FootstandControlConfig, - FootstandSensor, - Go2FootStandCfg, - Go2FootStandDomainRandConfig, - Go2FootStandDomainRandomizationProvider, - Go2FootStandTask, - RewardConfig, + FRAME_OBS_DIM, + NUM_ACTIONS, + PRIVILEGED_OBS_DIM, + FootstandIncrementalAction, + FootstandIncrementalActionCfg, + FootstandMassRandomization, + FootstandReward, + FootstandTermination, +) +from unilab.training.backend_adapter import BackendAdapter + +ROOT_DIR = Path(__file__).parents[3] +CONF_DIR = ROOT_DIR / "conf" + +_JOINT_NAMES = ( + "FL_hip_joint", + "FL_thigh_joint", + "FL_calf_joint", + "FR_hip_joint", + "FR_thigh_joint", + "FR_calf_joint", + "RL_hip_joint", + "RL_thigh_joint", + "RL_calf_joint", + "RR_hip_joint", + "RR_thigh_joint", + "RR_calf_joint", +) +_ACTION_JOINT_NAMES = ( + "FR_hip_joint", + "FR_thigh_joint", + "FR_calf_joint", + "FL_hip_joint", + "FL_thigh_joint", + "FL_calf_joint", + "RR_hip_joint", + "RR_thigh_joint", + "RR_calf_joint", + "RL_hip_joint", + "RL_thigh_joint", + "RL_calf_joint", +) +_ACTUATOR_NAMES = tuple(name.removesuffix("_joint") for name in _ACTION_JOINT_NAMES) +_BODY_NAMES = ( + "base", + "FL_hip", + "FL_thigh", + "FL_calf", + "FR_hip", + "FR_thigh", + "FR_calf", + "RL_hip", + "RL_thigh", + "RL_calf", + "RR_hip", + "RR_thigh", + "RR_calf", +) +_OWNER_CASES = ( + pytest.param( + "ppo", + ("task=go2_footstand/mujoco",), + "mujoco", + frozenset(("pd_gains", "floor_friction", "link_mass", "torso_com", "joint_armature")), + 0.05, + 2.0, + id="ppo-mujoco", + ), + pytest.param( + "ppo", + ("task=go2_footstand/motrix",), + "motrix", + frozenset(("pd_gains",)), + 0.02, + 3.0, + id="ppo-motrix", + ), + pytest.param( + "ppo", + ("task=go2_footstand/drake",), + "drake", + frozenset(), + 0.05, + 2.0, + id="ppo-drake", + ), + pytest.param( + "offpolicy", + ("algo=sac", "task=sac/go2_footstand/drake"), + "drake", + frozenset(), + 0.05, + 2.0, + id="sac-drake", + ), ) -def test_go2_footstand_registers_supported_backends() -> None: - ensure_registries() - - meta = registry.list_registered_envs()["Go2FootStand"] - - assert set(meta["available_backends"]) == {"drake", "motrix", "mujoco"} - - -class _OrientationBackend: - pass - - -class _JointRangeBackend: - def get_joint_range(self) -> np.ndarray: - return np.array([[-2.0, 2.0], [0.0, 2.0]], dtype=np.float32) - - -class _BaseMotionBackend: - def get_base_lin_vel(self) -> np.ndarray: - return np.array([[1.0, 2.0, 0.0], [1.0, 0.0, 0.0]], dtype=np.float32) - - def get_base_ang_vel(self) -> np.ndarray: - return np.array([[0.0, 0.0, 2.0], [0.0, 2.0, 0.0]], dtype=np.float32) - - -class _KneeHeightBackend: - def __init__(self, knee_height: np.ndarray) -> None: - self._knee_height = knee_height - - def get_body_pos_w(self, body_ids: np.ndarray) -> np.ndarray: - assert body_ids.shape == (4,) - out = np.zeros((*self._knee_height.shape, 3), dtype=np.float32) - out[:, :, 2] = self._knee_height - return out - - -class _BasePoseBackend: - def get_base_pos(self) -> np.ndarray: - return np.zeros((3, 3), dtype=np.float32) - - def get_base_quat(self) -> np.ndarray: - return np.tile(np.array([[1.0, 0.0, 0.0, 0.0]], dtype=np.float32), (3, 1)) - - -def test_go2_footstand_cfg_uses_rear_body_contact_termination() -> None: - cfg = Go2FootStandCfg() - - assert isinstance(cfg.sensor, FootstandSensor) - assert "base1_contact" in cfg.sensor.ternamate_contact - assert "base2_contact" in cfg.sensor.ternamate_contact - assert "base3_contact" in cfg.sensor.ternamate_contact - assert "RL_calf_contact1" in cfg.sensor.ternamate_contact - assert "RR_calf_contact2" in cfg.sensor.ternamate_contact - assert "FL_calf_contact1" in cfg.sensor.penalty_contact - assert "FR_calf_contact2" in cfg.sensor.penalty_contact - assert cfg.noise_config.level == pytest.approx(1.0) - assert cfg.noise_config.scale_joint_angle == pytest.approx(0.01) - assert cfg.noise_config.scale_joint_vel == pytest.approx(1.5) - assert cfg.add_body_sensors is True - assert isinstance(cfg.control_config, FootstandControlConfig) - assert cfg.control_config.action_scale == pytest.approx(0.3) - assert cfg.control_config.clip_actions == pytest.approx(1.0) - assert isinstance(cfg.domain_rand, Go2FootStandDomainRandConfig) - assert cfg.domain_rand.randomize_kp is False - assert cfg.domain_rand.randomize_floor_friction is True - assert cfg.obs_history_len == 15 - assert cfg.soft_joint_pos_limit_factor == pytest.approx(0.9) - assert cfg.energy_termination_threshold == np.inf - assert cfg.termination_grace_steps == 100 - assert cfg.termination_height_fraction == pytest.approx(0.8) - assert cfg.termination_orientation_threshold == pytest.approx(0.2) - assert cfg.max_episode_seconds == pytest.approx(10.0) - - -def test_go2_footstand_orientation_flips_handstand_target() -> None: - env = object.__new__(Go2FootStandTask) - env._backend = _OrientationBackend() - env._desired_forward_vec = np.array([0.0, 0.0, 1.0], dtype=np.float32) - env._get_body_forward = lambda: np.array( # type: ignore[method-assign] - [[0.0, 0.0, 1.0], [0.0, 0.0, -1.0]], dtype=np.float32 - ) - - reward = env._reward_orientation( - RewardContext( - info={}, - linvel=np.zeros((2, 3), dtype=np.float32), - gyro=np.zeros((2, 3), dtype=np.float32), - dof_pos=np.zeros((2, 12), dtype=np.float32), - ) - ) - - np.testing.assert_allclose(reward, np.array([1.0, 0.0], dtype=np.float32)) - - -def test_go2_footstand_soft_joint_limits_use_playground_factor() -> None: - env = object.__new__(Go2FootStandTask) - cfg = Go2FootStandCfg() - cfg.soft_joint_pos_limit_factor = 0.5 - env._cfg = cfg - env._num_action = 2 - env._backend = _JointRangeBackend() - - env._init_soft_joint_limits() - - np.testing.assert_allclose(env._soft_lowers, np.array([-1.0, 0.5], dtype=np.float32)) - np.testing.assert_allclose(env._soft_uppers, np.array([1.0, 1.5], dtype=np.float32)) - - -def test_go2_footstand_reward_functions_include_stability_terms() -> None: - env = object.__new__(Go2FootStandTask) - - env._init_reward_functions() - - assert "tar" in env._reward_fns - assert "penalty_contact" in env._reward_fns - assert "termination" in env._reward_fns - assert "rear_feet_contact" in env._reward_fns - assert "rear_leg_symmetry" in env._reward_fns - assert "rear_leg_splay" in env._reward_fns - assert "rear_foot_slip" in env._reward_fns - assert "rear_foot_anchor" in env._reward_fns - assert "front_leg_motion" in env._reward_fns - assert "front_feet_crossing" in env._reward_fns - assert "front_leg_crossing" in env._reward_fns - assert "upright_stability" in env._reward_fns - assert "knee_clearance" in env._reward_fns - - -def test_go2_footstand_pose_targets_front_legs_and_supports_rear() -> None: - env = object.__new__(Go2FootStandTask) - - env._init_footstand_pose_targets() - - assert env.feet_geom_names == [0, 1] - assert env._joint_ids == [6, 7, 8, 9, 10, 11] - assert env._tar_ids == [0, 1, 2, 3, 4, 5] - np.testing.assert_allclose( - env.target_angle, np.array([0.0, 1.82, -1.16, 0.0, 1.82, -1.16], dtype=np.float32) - ) - - -def test_go2_footstand_obs_matches_playground_state_layout() -> None: - env = object.__new__(Go2FootStandTask) - cfg = Go2FootStandCfg() - cfg.noise_config.level = 0.0 - env._cfg = cfg - env._num_envs = 1 - env.default_angles = np.arange(12, dtype=np.float32).reshape(1, 12) - env._obs_history = np.zeros((1, cfg.obs_history_len, 45), dtype=np.float32) - - linvel = np.array([[1.0, 2.0, 3.0]], dtype=np.float32) - gyro = np.array([[4.0, 5.0, 6.0]], dtype=np.float32) - gravity = np.array([[0.0, 0.0, -1.0]], dtype=np.float32) - dof_pos = env.default_angles + 0.5 - dof_vel = np.arange(12, dtype=np.float32).reshape(1, 12) + 10.0 - last_actions = np.full((1, 12), 0.25, dtype=np.float32) - current_actions = np.full((1, 12), 0.75, dtype=np.float32) - accelerometer = np.array([[7.0, 8.0, 9.0]], dtype=np.float32) - global_angvel = np.array([[10.0, 11.0, 12.0]], dtype=np.float32) - torques = np.arange(12, dtype=np.float32).reshape(1, 12) + 20.0 - - obs = env._compute_obs( - {"last_actions": last_actions, "current_actions": current_actions, "torques": torques}, - linvel, - gyro, - gravity, - dof_pos, - dof_vel, - np.array([[0.53]], dtype=np.float32), - accelerometer, - global_angvel, - ) - - current_frame = obs["obs"][:, -45:] - assert obs["obs"].shape == (1, 675) - assert obs["critic"].shape == (1, 724) - np.testing.assert_allclose(obs["obs"][:, : 14 * 45], 0.0) - np.testing.assert_allclose(current_frame[:, 0:3], linvel) - np.testing.assert_allclose(current_frame[:, 3:6], gyro) - np.testing.assert_allclose(current_frame[:, 6:9], gravity) - np.testing.assert_allclose(current_frame[:, -12:], last_actions) - - -def test_go2_footstand_reset_obs_fills_history_with_current_frame() -> None: - env = object.__new__(Go2FootStandTask) - cfg = Go2FootStandCfg() - cfg.noise_config.level = 0.0 - env._cfg = cfg - env._num_envs = 2 - env.default_angles = np.zeros((1, 12), dtype=np.float32) - env._obs_history = np.zeros((2, cfg.obs_history_len, 45), dtype=np.float32) - - linvel = np.array([[1.0, 0.0, 0.0]], dtype=np.float32) - gyro = np.array([[0.0, 2.0, 0.0]], dtype=np.float32) - gravity = np.array([[0.0, 0.0, -1.0]], dtype=np.float32) - dof_pos = np.ones((1, 12), dtype=np.float32) - dof_vel = np.full((1, 12), 3.0, dtype=np.float32) - - obs = env._compute_obs( - {"last_actions": np.full((1, 12), 0.5, dtype=np.float32)}, - linvel, - gyro, - gravity, - dof_pos, - dof_vel, - np.array([[0.53]], dtype=np.float32), - np.zeros((1, 3), dtype=np.float32), - np.zeros((1, 3), dtype=np.float32), - env_ids=np.array([1], dtype=np.int32), - ) - - frames = obs["obs"].reshape(1, cfg.obs_history_len, 45) - np.testing.assert_allclose(frames[:, 0, :], frames[:, -1, :]) - np.testing.assert_allclose(env._obs_history[0], 0.0) - np.testing.assert_allclose(env._obs_history[1], frames[0]) - - -class _EnergyTerminationBackend: - def __init__(self) -> None: - self._sensors = { - "local_linvel": np.zeros((1, 3), dtype=np.float32), - "gyro": np.zeros((1, 3), dtype=np.float32), - "upvector": np.array([[0.0, 0.0, 1.0]], dtype=np.float32), - "accelerometer": np.zeros((1, 3), dtype=np.float32), - "global_angvel": np.zeros((1, 3), dtype=np.float32), - "global_position": np.array([[0.0, 0.0, 0.53]], dtype=np.float32), - } - for name in FootstandSensor.feet_force: - self._sensors[name] = np.zeros((1, 1), dtype=np.float32) - for name in FootstandSensor.feet_pos: - self._sensors[name] = np.zeros((1, 3), dtype=np.float32) - for name in FootstandSensor.ternamate_contact: - self._sensors[name] = np.zeros((1, 1), dtype=np.float32) - - def get_sensor_data(self, name: str) -> np.ndarray: - return self._sensors[name] - - def get_sensor_data_batch(self, names: list[str]) -> np.ndarray: - values = [self.get_sensor_data(name).reshape(1, -1) for name in names] - return np.concatenate(values, axis=1) - - def get_dof_pos(self) -> np.ndarray: - return np.zeros((1, 12), dtype=np.float32) - - def get_dof_vel(self) -> np.ndarray: - return np.full((1, 12), 10.0, dtype=np.float32) - - def get_base_quat(self) -> np.ndarray: - return np.array([[1.0, 0.0, 0.0, 0.0]], dtype=np.float32) - - def get_base_pos(self) -> np.ndarray: - return np.array([[0.0, 0.0, 0.53]], dtype=np.float32) - - def get_body_pos_w(self, body_ids: np.ndarray) -> np.ndarray: - return np.zeros((1, len(body_ids), 3), dtype=np.float32) - - -def test_go2_footstand_energy_threshold_terminates() -> None: - env = object.__new__(Go2FootStandTask) - cfg = Go2FootStandCfg( - reward_config=RewardConfig(scales={}, tracking_sigma=0.25, base_height_target=0.3) - ) - cfg.noise_config.level = 0.0 - cfg.energy_termination_threshold = 1.0 - env._cfg = cfg - env._reward_cfg = cfg.reward_config - env._backend = _EnergyTerminationBackend() - env._num_envs = 1 - env._num_action = 12 - env._z_des = 0.53 - env._desired_forward_vec = np.array([0.0, 0.0, 1.0], dtype=np.float32) - env.default_angles = np.zeros((1, 12), dtype=np.float32) - env.feet_force = np.zeros((1, 4, 1), dtype=np.float32) - env.feet_pos = np.zeros((1, 4, 3), dtype=np.float32) - env.torso_height = np.zeros((1,), dtype=np.float32) - env._last_dof_vel_for_acc = np.zeros((1, 12), dtype=np.float32) - env._motor_targets = np.zeros((1, 12), dtype=np.float32) - env._last_terminated = np.zeros((1,), dtype=bool) - env._tracked_body_ids = np.arange(6, dtype=np.int32) - env._tracked_body_pos = np.zeros((1, 6, 3), dtype=np.float32) - env._enable_reward_log = False - - state = NpEnvState( - obs={}, - reward=np.zeros((1,), dtype=np.float32), - terminated=np.zeros((1,), dtype=bool), - truncated=np.zeros((1,), dtype=bool), - info={ - "current_actions": np.zeros((1, 12), dtype=np.float32), - "last_actions": np.zeros((1, 12), dtype=np.float32), - }, - ) - - updated = env.update_state(state) - - assert updated.terminated[0] - - -def test_go2_footstand_action_updates_incremental_motor_targets() -> None: - env = object.__new__(Go2FootStandTask) - env._cfg = Go2FootStandCfg() - env._motor_targets = np.zeros((1, 12), dtype=np.float32) - state = NpEnvState( - obs={}, - reward=np.zeros((1,), dtype=np.float32), - terminated=np.zeros((1,), dtype=bool), - truncated=np.zeros((1,), dtype=bool), - info={"current_actions": np.full((1, 12), 0.1, dtype=np.float32)}, - ) - - ctrl = env.apply_action(np.full((1, 12), 0.5, dtype=np.float32), state) - - np.testing.assert_allclose(ctrl, np.full((1, 12), 0.15, dtype=np.float32)) - np.testing.assert_allclose(state.info["last_actions"], np.full((1, 12), 0.1, dtype=np.float32)) - np.testing.assert_allclose( - state.info["current_actions"], np.full((1, 12), 0.5, dtype=np.float32) - ) - - -def test_go2_footstand_action_clips_policy_actions_and_motor_targets() -> None: - env = object.__new__(Go2FootStandTask) - env._cfg = Go2FootStandCfg() - env._motor_targets = np.zeros((1, 2), dtype=np.float32) - env._target_lowers = np.array([-0.2, -0.4], dtype=np.float32) - env._target_uppers = np.array([0.2, 0.4], dtype=np.float32) - state = NpEnvState( - obs={}, - reward=np.zeros((1,), dtype=np.float32), - terminated=np.zeros((1,), dtype=bool), - truncated=np.zeros((1,), dtype=bool), - info={}, - ) - - ctrl = env.apply_action(np.array([[10.0, -10.0]], dtype=np.float32), state) - - np.testing.assert_allclose( - state.info["current_actions"], np.array([[1.0, -1.0]], dtype=np.float32) - ) - np.testing.assert_allclose(ctrl, np.array([[0.2, -0.3]], dtype=np.float32)) - - ctrl = env.apply_action(np.array([[10.0, -10.0]], dtype=np.float32), state) - - np.testing.assert_allclose(ctrl, np.array([[0.2, -0.4]], dtype=np.float32)) - - -def test_go2_footstand_playground_reset_randomization_payload_shapes() -> None: - np.random.seed(0) - env = object.__new__(Go2FootStandTask) - env._cfg = Go2FootStandCfg() - env._num_action = 12 - env._floor_geom_id = 0 - env._base_body_id = 1 - env._base_geom_friction = np.ones((3, 3), dtype=np.float64) - env._base_body_mass = np.ones((4,), dtype=np.float64) - env._base_body_ipos = np.zeros((4, 3), dtype=np.float64) - env._base_dof_armature = np.ones((18,), dtype=np.float64) - - payload = env._build_playground_reset_randomization(num_reset=2) - - assert payload is not None - assert payload.geom_friction is not None and payload.geom_friction.shape == (2, 3, 3) - assert payload.body_mass is not None and payload.body_mass.shape == (2, 4) - assert payload.body_ipos is not None and payload.body_ipos.shape == (2, 4, 3) - assert payload.dof_armature is not None and payload.dof_armature.shape == (2, 18) - assert np.all((payload.geom_friction[:, 0, 0] >= 0.4) & (payload.geom_friction[:, 0, 0] <= 1.0)) - np.testing.assert_allclose(payload.dof_armature[:, :6], 1.0) - - -def test_go2_footstand_reset_randomization_merges_common_and_playground_terms() -> None: - base = ResetRandomizationPayload( - base_mass_delta=np.array([0.1], dtype=np.float64), - kp=np.ones((1, 12), dtype=np.float64), - ) - playground = ResetRandomizationPayload( - body_mass=np.full((1, 4), 2.0, dtype=np.float64), - geom_friction=np.full((1, 3, 3), 0.7, dtype=np.float64), - ) - - merged = Go2FootStandDomainRandomizationProvider._merge_reset_randomization(base, playground) - - assert merged is not None - np.testing.assert_allclose(merged.base_mass_delta, base.base_mass_delta) - np.testing.assert_allclose(merged.kp, base.kp) - np.testing.assert_allclose(merged.body_mass, playground.body_mass) - np.testing.assert_allclose(merged.geom_friction, playground.geom_friction) - - -def test_go2_footstand_height_reward_matches_playground_shape() -> None: - env = object.__new__(Go2FootStandTask) - env._z_des = 0.53 - env.torso_height = np.array([0.53, 0.33, 0.63], dtype=np.float32) - - reward = env._reward_height( - RewardContext( - info={}, - linvel=np.zeros((3, 3), dtype=np.float32), - gyro=np.zeros((3, 3), dtype=np.float32), - dof_pos=np.zeros((3, 12), dtype=np.float32), - ) - ) - - np.testing.assert_allclose( - reward, np.array([1.0, np.exp(-2.0), np.exp(-1.0)], dtype=np.float32), rtol=1e-6 - ) - - -def test_go2_footstand_contact_cost_only_penalizes_front_feet() -> None: - env = object.__new__(Go2FootStandTask) - env.feet_geom_names = [0, 1] - env.feet_force = np.zeros((3, 4, 1), dtype=np.float32) - env.feet_force[0, 0, 0] = 1.0 - env.feet_force[1, 2, 0] = 1.0 - env.feet_force[2, 3, 0] = 1.0 - - cost = env._cost_contact( - RewardContext( - info={}, - linvel=np.zeros((3, 3), dtype=np.float32), - gyro=np.zeros((3, 3), dtype=np.float32), - dof_pos=np.zeros((3, 12), dtype=np.float32), - ) - ) - - np.testing.assert_allclose(cost, np.array([1.0, 0.0, 0.0], dtype=np.float32)) - - -def test_go2_footstand_rear_feet_contact_rewards_support_feet() -> None: - env = object.__new__(Go2FootStandTask) - env.feet_force = np.zeros((3, 4, 1), dtype=np.float32) - env.feet_force[0, 2:4, 0] = 1.0 - env.feet_force[1, 2, 0] = 1.0 - env.feet_force[2, 0:2, 0] = 1.0 - - reward = env._reward_rear_feet_contact( - RewardContext( - info={}, - linvel=np.zeros((3, 3), dtype=np.float32), - gyro=np.zeros((3, 3), dtype=np.float32), - dof_pos=np.zeros((3, 12), dtype=np.float32), - ) - ) - - np.testing.assert_allclose(reward, np.array([1.0, 0.5, 0.0], dtype=np.float32)) - - -def test_go2_footstand_both_rear_feet_contact_requires_two_support_feet() -> None: - env = object.__new__(Go2FootStandTask) - env.feet_force = np.zeros((3, 4, 1), dtype=np.float32) - env.feet_force[0, 2:4, 0] = 1.0 - env.feet_force[1, 2, 0] = 1.0 - env.feet_force[2, 0:2, 0] = 1.0 - - reward = env._reward_both_rear_feet_contact( - RewardContext( - info={}, - linvel=np.zeros((3, 3), dtype=np.float32), - gyro=np.zeros((3, 3), dtype=np.float32), - dof_pos=np.zeros((3, 12), dtype=np.float32), - ) - ) - - np.testing.assert_allclose(reward, np.array([1.0, 0.0, 0.0], dtype=np.float32)) - - -def test_go2_footstand_rear_foot_slip_penalizes_contact_foot_xy_motion() -> None: - env = object.__new__(Go2FootStandTask) - env._cfg = Go2FootStandCfg( - reward_config=RewardConfig( - scales={}, - tracking_sigma=0.25, - base_height_target=0.3, - rear_foot_slip_deadband=0.02, - ) - ) - env._reward_cfg = env._cfg.reward_config - env._num_envs = 3 - env.feet_force = np.zeros((3, 4, 1), dtype=np.float32) - env.feet_force[:, 2:4, 0] = 1.0 - env.feet_force[2, 3, 0] = 0.0 - env._last_feet_pos = np.zeros((3, 4, 3), dtype=np.float32) - env.feet_pos = np.zeros((3, 4, 3), dtype=np.float32) - env.feet_pos[0, 2, 0] = 0.02 - env.feet_pos[0, 3, 0] = 0.02 - env.feet_pos[1, 2, 0] = 0.04 - env.feet_pos[1, 3, 1] = 0.03 - env.feet_pos[2, 2, 0] = 0.04 - env.feet_pos[2, 3, 0] = 0.10 - - env._update_rear_foot_slip() - cost = env._cost_rear_foot_slip( - RewardContext( - info={}, - linvel=np.zeros((3, 3), dtype=np.float32), - gyro=np.zeros((3, 3), dtype=np.float32), - dof_pos=np.zeros((3, 12), dtype=np.float32), - ) - ) - - np.testing.assert_allclose( - cost, - np.array([0.9604, 3.0554, 1.9602], dtype=np.float32), - rtol=1e-6, - ) - - -def test_go2_footstand_rear_foot_anchor_penalizes_cumulative_contact_drift() -> None: - env = object.__new__(Go2FootStandTask) - env._reward_cfg = RewardConfig( - scales={}, - tracking_sigma=0.25, - base_height_target=0.3, - rear_foot_anchor_radius=0.03, - ) - env._num_envs = 3 - env._standing_mask = lambda: np.ones((3,), dtype=np.float32) # type: ignore[method-assign] - env.feet_force = np.zeros((3, 4, 1), dtype=np.float32) - env.feet_force[:, 2:4, 0] = 1.0 - env.feet_force[2, 3, 0] = 0.0 - env.feet_pos = np.zeros((3, 4, 3), dtype=np.float32) - env.feet_pos[:, 2:4, :2] = np.array( - [ - [[0.0, 0.0], [0.1, 0.0]], - [[0.0, 0.0], [0.1, 0.0]], - [[0.0, 0.0], [0.1, 0.0]], - ], - dtype=np.float32, - ) - - env._update_rear_foot_anchor() - np.testing.assert_allclose(env._rear_foot_anchor, 0.0) - - env.feet_pos[0, 2, 0] += 0.02 - env.feet_pos[0, 3, 0] += 0.02 - env.feet_pos[1, 2, 0] += 0.06 - env.feet_pos[1, 3, 1] += 0.09 - env.feet_pos[2, 2, 0] += 0.06 - env.feet_pos[2, 3, 0] += 0.20 - - env._update_rear_foot_anchor() - cost = env._cost_rear_foot_anchor( - RewardContext( - info={}, - linvel=np.zeros((3, 3), dtype=np.float32), - gyro=np.zeros((3, 3), dtype=np.float32), - dof_pos=np.zeros((3, 12), dtype=np.float32), - ) - ) - - np.testing.assert_allclose( - cost, - np.array([0.0, 2.5, 0.5], dtype=np.float32), - rtol=1e-6, - ) - - env.feet_force[1, 2:4, 0] = 0.0 - env._update_rear_foot_anchor() - assert env._rear_foot_anchor_contact[1].tolist() == [False, False] - np.testing.assert_allclose(env._rear_foot_anchor[1], 0.0) - - env.feet_pos[1, 2:4, :2] = np.array([[0.5, 0.0], [0.6, 0.0]], dtype=np.float32) - env.feet_force[1, 2:4, 0] = 1.0 - env._update_rear_foot_anchor() - np.testing.assert_allclose(env._rear_foot_anchor[1], 0.0) - np.testing.assert_allclose( - env._rear_foot_anchor_pos[1], - np.array([[0.5, 0.0], [0.6, 0.0]], dtype=np.float32), - ) - - -def test_go2_footstand_rear_foot_anchor_is_standing_gated() -> None: - env = object.__new__(Go2FootStandTask) - env._reward_cfg = RewardConfig( - scales={}, - tracking_sigma=0.25, - base_height_target=0.3, - rear_foot_anchor_radius=0.03, - ) - env._num_envs = 2 - env._standing_mask = lambda: np.array([1.0, 0.0], dtype=np.float32) # type: ignore[method-assign] - env.feet_force = np.zeros((2, 4, 1), dtype=np.float32) - env.feet_force[:, 2:4, 0] = 1.0 - env.feet_pos = np.zeros((2, 4, 3), dtype=np.float32) - env.feet_pos[:, 2:4, :2] = np.array( - [ - [[0.0, 0.0], [0.1, 0.0]], - [[0.0, 0.0], [0.1, 0.0]], - ], - dtype=np.float32, - ) - - env._update_rear_foot_anchor() - env.feet_pos[:, 2:4, 0] += 0.09 - env._update_rear_foot_anchor() - - np.testing.assert_allclose(env._rear_foot_anchor[0], 4.0) - np.testing.assert_allclose(env._rear_foot_anchor[1], 0.0) - assert env._rear_foot_anchor_contact[0].tolist() == [True, True] - assert env._rear_foot_anchor_contact[1].tolist() == [False, False] - - -def test_go2_footstand_front_feet_air_rewards_front_feet_off_ground() -> None: - env = object.__new__(Go2FootStandTask) - env.feet_force = np.zeros((3, 4, 1), dtype=np.float32) - env.feet_force[0, 2:4, 0] = 1.0 - env.feet_force[1, 0, 0] = 1.0 - env.feet_force[2, 0:2, 0] = 1.0 - - reward = env._reward_front_feet_air( - RewardContext( - info={}, - linvel=np.zeros((3, 3), dtype=np.float32), - gyro=np.zeros((3, 3), dtype=np.float32), - dof_pos=np.zeros((3, 12), dtype=np.float32), - ) - ) - - np.testing.assert_allclose(reward, np.array([1.0, 0.0, 0.0], dtype=np.float32)) - - -def test_go2_footstand_balanced_footstand_requires_support_air_and_standing() -> None: - env = object.__new__(Go2FootStandTask) - env.feet_force = np.zeros((4, 4, 1), dtype=np.float32) - env.feet_force[0, 2:4, 0] = 1.0 - env.feet_force[1, 2:4, 0] = 1.0 - env.feet_force[1, 0, 0] = 1.0 - env.feet_force[2, 2, 0] = 1.0 - env.feet_force[3, 2:4, 0] = 1.0 - env._standing_mask = lambda: np.array([1.0, 1.0, 1.0, 0.0], dtype=np.float32) # type: ignore[method-assign] - - reward = env._reward_balanced_footstand( - RewardContext( - info={}, - linvel=np.zeros((4, 3), dtype=np.float32), - gyro=np.zeros((4, 3), dtype=np.float32), - dof_pos=np.zeros((4, 12), dtype=np.float32), - ) - ) +def _compose(config_group: str, overrides: Sequence[str]) -> DictConfig: + GlobalHydra.instance().clear() + with initialize_config_dir(config_dir=str(CONF_DIR / config_group), version_base="1.3"): + return compose("config", overrides=list(overrides)) - np.testing.assert_allclose(reward, np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32)) - - -def test_go2_footstand_rear_leg_symmetry_mirrors_hip_sign_only() -> None: - env = object.__new__(Go2FootStandTask) - env._standing_mask = lambda: np.array([0.0, 0.0, 1.0], dtype=np.float32) # type: ignore[method-assign] - dof_pos = np.zeros((3, 12), dtype=np.float32) - dof_pos[0, 6:9] = np.array([0.2, 1.0, -1.5], dtype=np.float32) - dof_pos[0, 9:12] = np.array([-0.2, 1.0, -1.5], dtype=np.float32) - dof_pos[1, 6:9] = np.array([0.2, 1.0, -1.5], dtype=np.float32) - dof_pos[1, 9:12] = np.array([0.3, 1.0, -1.5], dtype=np.float32) - dof_pos[2, 6:9] = np.array([0.0, 1.2, -1.1], dtype=np.float32) - dof_pos[2, 9:12] = np.array([0.0, 0.9, -1.7], dtype=np.float32) - - cost = env._cost_rear_leg_symmetry( - RewardContext( - info={}, - linvel=np.zeros((3, 3), dtype=np.float32), - gyro=np.zeros((3, 3), dtype=np.float32), - dof_pos=dof_pos, - ) - ) - np.testing.assert_allclose( - cost, - np.array([0.0, 0.25 / 3.0, 0.0], dtype=np.float32), - rtol=1e-6, - ) +def _materialize( + config_group: str, + overrides: Sequence[str], +) -> tuple[DictConfig, ManagerBasedRlEnvCfg, dict[str, Any]]: + hydra_cfg = _compose(config_group, overrides) + env_override = BackendAdapter(hydra_cfg, root_dir=ROOT_DIR).build_task_env_cfg_override() + env_cfg = registry.materialize_env_config("Go2FootStand") + assert isinstance(env_cfg, ManagerBasedRlEnvCfg) + apply_cfg_overrides(env_cfg, env_override) + env_cfg.validate() + return hydra_cfg, env_cfg, env_override -def test_go2_footstand_rear_leg_splay_penalizes_standing_hip_abduction() -> None: - env = object.__new__(Go2FootStandTask) - env._reward_cfg = RewardConfig( - scales={}, - tracking_sigma=0.25, - base_height_target=0.3, - rear_hip_abduction_margin=0.25, - ) - env._standing_mask = lambda: np.array([1.0, 1.0, 0.0], dtype=np.float32) # type: ignore[method-assign] - dof_pos = np.zeros((3, 12), dtype=np.float32) - dof_pos[0, [6, 9]] = np.array([0.2, -0.2], dtype=np.float32) - dof_pos[1, [6, 9]] = np.array([0.5, -0.4], dtype=np.float32) - dof_pos[2, [6, 9]] = np.array([0.6, -0.6], dtype=np.float32) - - cost = env._cost_rear_leg_splay( - RewardContext( - info={}, - linvel=np.zeros((3, 3), dtype=np.float32), - gyro=np.zeros((3, 3), dtype=np.float32), - dof_pos=dof_pos, - default_angles=np.zeros((1, 12), dtype=np.float32), - ) +def _make_env(backend: str, *, num_envs: int = 2) -> ManagerBasedRlEnv: + hydra_cfg, _, env_override = _materialize( + "ppo", + (f"task=go2_footstand/{backend}", f"algo.num_envs={num_envs}"), ) - - np.testing.assert_allclose( - cost, - np.array([0.0, (0.25**2 + 0.15**2) / 2.0, 0.0], dtype=np.float32), - rtol=1e-6, - ) - - -def test_go2_footstand_front_leg_motion_only_penalizes_standing_pose() -> None: - env = object.__new__(Go2FootStandTask) - env._z_des = 0.53 - env.torso_height = np.array([0.53, 0.53, 0.2], dtype=np.float32) - env._orientation_score = lambda: np.array([1.0, 0.4, 1.0], dtype=np.float32) # type: ignore[method-assign] - dof_vel = np.zeros((3, 12), dtype=np.float32) - dof_vel[:, 0:6] = 2.0 - - cost = env._cost_front_leg_motion( - RewardContext( - info={}, - linvel=np.zeros((3, 3), dtype=np.float32), - gyro=np.zeros((3, 3), dtype=np.float32), - dof_pos=np.zeros((3, 12), dtype=np.float32), - dof_vel=dof_vel, - ) + env = registry.make( + str(hydra_cfg.training.task_name), + sim_backend=backend, + env_cfg_override=env_override, + num_envs=num_envs, ) + assert isinstance(env, ManagerBasedRlEnv) + return env - np.testing.assert_allclose(cost, np.array([4.0, 0.0, 0.0], dtype=np.float32)) +def _assert_plain(value: Any) -> None: + assert not OmegaConf.is_config(value) + if is_dataclass(value) and not isinstance(value, type): + for item in fields(value): + _assert_plain(getattr(value, item.name)) + elif isinstance(value, Mapping): + for key, item in value.items(): + _assert_plain(key) + _assert_plain(item) + elif isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + for item in value: + _assert_plain(item) -def test_go2_footstand_front_leg_crossing_penalizes_body_frame_overlap() -> None: - env = object.__new__(Go2FootStandTask) - env._num_envs = 3 - env._backend = _BasePoseBackend() - env._reward_cfg = RewardConfig( - scales={}, - tracking_sigma=0.25, - base_height_target=0.3, - front_feet_min_separation=0.16, - front_feet_side_margin=0.04, - ) - env.feet_pos = np.zeros((3, 4, 3), dtype=np.float32) - env.feet_pos[0, 0, 1] = 0.10 - env.feet_pos[0, 1, 1] = -0.10 - env.feet_pos[1, 0, 1] = -0.02 - env.feet_pos[1, 1, 1] = 0.02 - env.feet_pos[2, 0, 1] = 0.05 - env.feet_pos[2, 1, 1] = -0.05 - env._tracked_body_pos = np.zeros((3, 6, 3), dtype=np.float32) - env._tracked_body_pos[:, 0, 1] = env.feet_pos[:, 0, 1] - env._tracked_body_pos[:, 1, 1] = env.feet_pos[:, 1, 1] - env._tracked_body_pos[:, 2, 1] = env.feet_pos[:, 0, 1] - env._tracked_body_pos[:, 3, 1] = env.feet_pos[:, 1, 1] - - cost = env._cost_front_leg_crossing( - RewardContext( - info={}, - linvel=np.zeros((3, 3), dtype=np.float32), - gyro=np.zeros((3, 3), dtype=np.float32), - dof_pos=np.zeros((3, 12), dtype=np.float32), - ) - ) - - np.testing.assert_allclose( - cost, - np.array([0.0, 0.0472, 0.0036], dtype=np.float32), - rtol=1e-6, - ) +def _action(env: ManagerBasedRlEnv) -> FootstandIncrementalAction: + term = env.action_manager.get_term("joint_pos") + assert isinstance(term, FootstandIncrementalAction) + return term -def test_go2_footstand_upright_stability_is_standing_gated() -> None: - env = object.__new__(Go2FootStandTask) - env._backend = _BaseMotionBackend() - env._z_des = 0.53 - env.torso_height = np.array([0.53, 0.53], dtype=np.float32) - env._orientation_score = lambda: np.array([1.0, 0.4], dtype=np.float32) # type: ignore[method-assign] - cost = env._cost_upright_stability( - RewardContext( - info={}, - linvel=np.zeros((2, 3), dtype=np.float32), - gyro=np.zeros((2, 3), dtype=np.float32), - dof_pos=np.zeros((2, 12), dtype=np.float32), - ) - ) +def _termination(env: ManagerBasedRlEnv) -> FootstandTermination: + term = env.termination_manager.get_term_cfg("footstand").func + assert isinstance(term, FootstandTermination) + return term - np.testing.assert_allclose(cost, np.array([6.0, 0.0], dtype=np.float32)) +def _reward(env: ManagerBasedRlEnv) -> FootstandReward: + term = env.reward_manager.get_term_cfg("footstand").func + assert isinstance(term, FootstandReward) + return term -def test_go2_footstand_knee_clearance_penalizes_low_knees() -> None: - env = object.__new__(Go2FootStandTask) - env._reward_cfg = RewardConfig( - scales={}, - tracking_sigma=0.25, - base_height_target=0.3, - knee_height_target=0.1, - ) - env._tracked_body_pos = np.zeros((3, 6, 3), dtype=np.float32) - env._tracked_body_pos[:, [2, 3, 4, 5], 2] = np.array( - [ - [0.1, 0.12, 0.2, 0.3], - [0.05, 0.05, 0.05, 0.05], - [0.0, 0.05, 0.1, 0.15], - ], - dtype=np.float32, - ) - cost = env._cost_knee_clearance( - RewardContext( - info={}, - linvel=np.zeros((3, 3), dtype=np.float32), - gyro=np.zeros((3, 3), dtype=np.float32), - dof_pos=np.zeros((3, 12), dtype=np.float32), +@pytest.mark.parametrize( + "config_group,overrides,backend,backend_events,joint_reset_radius,orientation_scale", + _OWNER_CASES, +) +def test_footstand_owner_materializes_complete_plain_manager_cfg( + config_group: str, + overrides: tuple[str, ...], + backend: str, + backend_events: frozenset[str], + joint_reset_radius: float, + orientation_scale: float, +) -> None: + registry.ensure_registries() + hydra_cfg, env_cfg, _ = _materialize(config_group, overrides) + + assert hydra_cfg.training.task_name == "Go2FootStand" + assert hydra_cfg.training.sim_backend == backend + assert hydra_cfg.algo.num_envs == 4096 + assert env_cfg.sim_dt == pytest.approx(0.004) + assert env_cfg.ctrl_dt == pytest.approx(0.02) + assert env_cfg.max_episode_seconds == pytest.approx(10.0) + assert env_cfg.adaptive_chunk_size is False + assert env_cfg.policy_observation_group == "policy" + assert env_cfg.critic_observation_group == "critic" + + assert env_cfg.scene is not None + assert env_cfg.scene.default_keyframe_name == "home" + robot = env_cfg.scene.entities["robot"] + assert robot.root_body_name == "base" + assert tuple(robot.joint_names or ()) == _JOINT_NAMES + assert tuple(robot.body_names or ()) == _BODY_NAMES + assert tuple(robot.geom_names or ()) == ("floor",) + assert tuple(robot.actuator_names or ()) == _ACTUATOR_NAMES + + policy = env_cfg.observations["policy"] + critic = env_cfg.observations["critic"] + assert policy is not None and critic is not None + assert list(policy.terms) == ["frame"] + assert list(critic.terms) == ["frame", "privileged"] + assert policy.enable_corruption is True + assert policy.terms["frame"] is not None + assert critic.terms["frame"] is not None + assert critic.terms["privileged"] is not None + assert policy.terms["frame"].history_length == 15 + assert critic.terms["frame"].history_length == 15 + assert critic.terms["privileged"].history_length == 0 + noise = policy.terms["frame"].noise + assert isinstance(noise, UniformNoiseCfg) + assert np.asarray(noise.n_min).shape == (FRAME_OBS_DIM,) + assert np.asarray(noise.n_max).shape == (FRAME_OBS_DIM,) + + action = env_cfg.actions["joint_pos"] + assert isinstance(action, FootstandIncrementalActionCfg) + assert tuple(action.actuator_names) == _ACTUATOR_NAMES + assert tuple(action.joint_names) == _ACTION_JOINT_NAMES + assert action.action_scale == pytest.approx(0.3) + assert action.clip_actions == pytest.approx(1.0) + assert action.kp == pytest.approx(35.0) + assert action.kd == pytest.approx(0.5) + + always_enabled = { + "reset_scene_to_default", + "reset_root_state_uniform", + "reset_joints", + } + active_events = {name for name, term in env_cfg.events.items() if term is not None} + assert active_events == always_enabled | backend_events + reset_joints = env_cfg.events["reset_joints"] + assert reset_joints is not None + assert tuple(reset_joints.params["position_offset_range"]) == pytest.approx( + (-joint_reset_radius, joint_reset_radius) + ) + if "link_mass" in backend_events: + link_mass = env_cfg.events["link_mass"] + assert link_mass is not None + assert link_mass.func is FootstandMassRandomization + + termination = env_cfg.terminations["footstand"] + assert termination is not None + assert termination.params["grace_steps"] == 100 + assert termination.params["energy_threshold"] == pytest.approx(200.0) + reward = env_cfg.rewards["footstand"] + assert reward is not None + assert reward.func is FootstandReward + assert reward.params["scales"]["orientation"] == pytest.approx(orientation_scale) + + assert FRAME_OBS_DIM * 15 == 675 + assert FRAME_OBS_DIM * 15 + PRIVILEGED_OBS_DIM == 724 + _assert_plain(env_cfg) + + +def test_footstand_registry_has_no_legacy_config_or_factory() -> None: + registry.ensure_registries() + bare_cfg = registry.materialize_env_config("Go2FootStand") + + assert isinstance(bare_cfg, ManagerBasedRlEnvCfg) + assert bare_cfg.actions == {} + assert bare_cfg.observations == {} + assert registry.list_registered_envs()["Go2FootStand"] == { + "config_factory": "ManagerBasedRlEnvCfg", + "available_backends": ["mujoco", "motrix", "drake"], + } + meta = registry._envs["Go2FootStand"] + assert all(factory is make_manager_based_rl_env for factory in meta.env_factory_dict.values()) + for legacy_override in ( + {"reward_config": {}}, + {"domain_rand": {"randomize_floor_friction": True}}, + {"control_config": {"action_scale": 0.4}}, + {"obs_history_len": 15}, + ): + with pytest.raises(ValueError, match="has no attribute"): + apply_cfg_overrides(ManagerBasedRlEnvCfg(), legacy_override) + + +@pytest.mark.parametrize("backend", ("mujoco", "motrix")) +def test_footstand_real_runtime_preserves_history_action_and_partial_reset( + backend: str, +) -> None: + registry.ensure_registries() + env = _make_env(backend) + try: + obs, info = env.reset(seed=7) + assert set(info) == {"log"} + assert env.obs_groups_spec == {"obs": 675, "critic": 724} + assert env.action_space.shape == (NUM_ACTIONS,) + assert obs["obs"].shape == (2, 675) + assert obs["critic"].shape == (2, 724) + assert np.isfinite(obs["obs"]).all() + assert np.isfinite(obs["critic"]).all() + + reset_frames = obs["obs"].reshape(2, 15, FRAME_OBS_DIM) + np.testing.assert_allclose(reset_frames, np.repeat(reset_frames[:, :1], 15, axis=1)) + clean_history = obs["critic"][:, :675].reshape(2, 15, FRAME_OBS_DIM) + np.testing.assert_allclose(clean_history, np.repeat(clean_history[:, :1], 15, axis=1)) + + action = _action(env) + initial_target = action.target.copy() + first_policy_action = np.full((2, NUM_ACTIONS), 2.0, dtype=np.float32) + action.process_actions(first_policy_action) + expected_target = np.clip( + initial_target + 0.3, + action.joint_lower[action.joint_ids], + action.joint_upper[action.joint_ids], ) - ) - - np.testing.assert_allclose(cost, np.array([0.0, 0.25, 0.3125], dtype=np.float32)) - - -def test_go2_footstand_post_grace_low_height_terminates() -> None: - env = object.__new__(Go2FootStandTask) - cfg = Go2FootStandCfg( - reward_config=RewardConfig(scales={}, tracking_sigma=0.25, base_height_target=0.3) - ) - cfg.noise_config.level = 0.0 - cfg.termination_grace_steps = 10 - env._cfg = cfg - env._reward_cfg = cfg.reward_config - env._backend = _EnergyTerminationBackend() - env._backend._sensors["global_position"] = np.array([[0.0, 0.0, 0.2]], dtype=np.float32) - env._num_envs = 1 - env._num_action = 12 - env._z_des = 0.53 - env.default_angles = np.zeros((1, 12), dtype=np.float32) - env.feet_force = np.zeros((1, 4, 1), dtype=np.float32) - env.feet_pos = np.zeros((1, 4, 3), dtype=np.float32) - env.torso_height = np.zeros((1,), dtype=np.float32) - env._last_dof_vel_for_acc = np.zeros((1, 12), dtype=np.float32) - env._motor_targets = np.zeros((1, 12), dtype=np.float32) - env._last_terminated = np.zeros((1,), dtype=bool) - env._tracked_body_ids = np.arange(6, dtype=np.int32) - env._tracked_body_pos = np.zeros((1, 6, 3), dtype=np.float32) - env._enable_reward_log = False - env._orientation_score = lambda: np.array([1.0], dtype=np.float32) # type: ignore[method-assign] - state = NpEnvState( - obs={}, - reward=np.zeros((1,), dtype=np.float32), - terminated=np.zeros((1,), dtype=bool), - truncated=np.zeros((1,), dtype=bool), - info={ - "steps": np.array([10], dtype=np.uint32), - "current_actions": np.zeros((1, 12), dtype=np.float32), - "last_actions": np.zeros((1, 12), dtype=np.float32), - }, - ) - - updated = env.update_state(state) - - assert updated.terminated[0] - - -def test_go2_footstand_returned_termination_does_not_alias_reset_bookkeeping() -> None: - env = object.__new__(Go2FootStandTask) - cfg = Go2FootStandCfg( - reward_config=RewardConfig(scales={}, tracking_sigma=0.25, base_height_target=0.3) - ) - cfg.noise_config.level = 0.0 - cfg.termination_grace_steps = 10 - env._cfg = cfg - env._reward_cfg = cfg.reward_config - env._backend = _EnergyTerminationBackend() - env._backend._sensors["global_position"] = np.array([[0.0, 0.0, 0.2]], dtype=np.float32) - env._num_envs = 1 - env._num_action = 12 - env._z_des = 0.53 - env.default_angles = np.zeros((1, 12), dtype=np.float32) - env.feet_force = np.zeros((1, 4, 1), dtype=np.float32) - env.feet_pos = np.zeros((1, 4, 3), dtype=np.float32) - env.torso_height = np.zeros((1,), dtype=np.float32) - env._last_dof_vel_for_acc = np.zeros((1, 12), dtype=np.float32) - env._motor_targets = np.zeros((1, 12), dtype=np.float32) - env._last_terminated = np.zeros((1,), dtype=bool) - env._tracked_body_ids = np.arange(6, dtype=np.int32) - env._tracked_body_pos = np.zeros((1, 6, 3), dtype=np.float32) - env._enable_reward_log = False - env._orientation_score = lambda: np.array([1.0], dtype=np.float32) # type: ignore[method-assign] - state = NpEnvState( - obs={}, - reward=np.zeros((1,), dtype=np.float32), - terminated=np.zeros((1,), dtype=bool), - truncated=np.zeros((1,), dtype=bool), - info={ - "steps": np.array([10], dtype=np.uint32), - "current_actions": np.zeros((1, 12), dtype=np.float32), - "last_actions": np.zeros((1, 12), dtype=np.float32), - }, - ) - - updated = env.update_state(state) - env._last_terminated[0] = False + np.testing.assert_allclose(action.target, expected_target, atol=1e-6) + + obs, _ = env.reset(seed=7) + first_policy_action = np.full((2, NUM_ACTIONS), 0.05, dtype=np.float32) + first_state = env.step(first_policy_action) + assert not first_state.terminated.any() + first_frames = first_state.obs["obs"].reshape(2, 15, FRAME_OBS_DIM) + np.testing.assert_allclose(first_frames[:, -1, -NUM_ACTIONS:], 0.0) + + second_policy_action = np.full((2, NUM_ACTIONS), 0.025, dtype=np.float32) + second_state = env.step(second_policy_action) + assert not second_state.terminated.any() + second_frames = second_state.obs["obs"].reshape(2, 15, FRAME_OBS_DIM) + np.testing.assert_allclose(second_frames[:, -1, -NUM_ACTIONS:], 0.05) + assert np.isfinite(second_state.reward).all() + assert np.all(second_state.reward >= 0.0) + + untouched_target = action.target[1].copy() + untouched_torque = action.state.torques[1].copy() + assert env.state is not None + untouched_obs = {name: value[1].copy() for name, value in env.state.obs.items()} + reset_obs, _ = env.reset(env_ids=np.asarray([0], dtype=np.int32)) + assert reset_obs["obs"].shape == (1, 675) + partial_frames = reset_obs["obs"].reshape(1, 15, FRAME_OBS_DIM) + np.testing.assert_allclose(partial_frames, np.repeat(partial_frames[:, :1], 15, axis=1)) + np.testing.assert_allclose(action.target[1], untouched_target) + np.testing.assert_allclose(action.state.torques[1], untouched_torque) + assert env.state is not None + for name, expected in untouched_obs.items(): + np.testing.assert_allclose(env.state.obs[name][1], expected) + finally: + env.close() - assert updated.terminated[0] +def test_footstand_termination_uses_grace_boundary() -> None: + registry.ensure_registries() + env = _make_env("mujoco") + try: + env.reset(seed=11) + term = _termination(env) + state = term.state + state.height[:] = 0.0 + state.orientation[:] = 0.0 + state.upvector[:] = (0.0, 0.0, 1.0) + state.termination_contact[:] = False + state.torques[:] = 0.0 + state.joint_vel[:] = 0.0 + + env.episode_length_buf[:] = 100 + assert not term(env).any() + env.episode_length_buf[:] = 101 + assert term(env).all() + finally: + env.close() -def test_go2_footstand_joint_limit_cost_uses_soft_limits() -> None: - env = object.__new__(Go2FootStandTask) - env._soft_lowers = np.array([-1.0, -1.0], dtype=np.float32) - env._soft_uppers = np.array([1.0, 1.0], dtype=np.float32) - cost = env._cost_joint_pos_limits( - RewardContext( - info={}, - linvel=np.zeros((2, 3), dtype=np.float32), - gyro=np.zeros((2, 3), dtype=np.float32), - dof_pos=np.array([[0.0, 1.5], [-1.25, 0.0]], dtype=np.float32), +def test_footstand_reward_clips_aggregate_before_dt_scaling() -> None: + registry.ensure_registries() + env = _make_env("mujoco") + try: + env.reset(seed=13) + term = _termination(env) + reward = _reward(env) + state = term.state + term.terminated[:] = False + state.height[:] = 0.53 + state.foot_contact[:] = False + state.foot_contact[0, 0] = True + reward._scales = {"height": 2.0, "contact": -100.0} + + np.testing.assert_allclose(reward(env), np.asarray([0.0, 2.0], dtype=np.float32)) + np.testing.assert_allclose( + env.reward_manager.compute(dt=env.step_dt), + np.asarray([0.0, 0.04], dtype=np.float32), ) - ) - - np.testing.assert_allclose(cost, np.array([0.5, 0.25], dtype=np.float32)) - + finally: + env.close() -def test_go2_footstand_reset_critic_height_uses_backend_sensor() -> None: - pytest.importorskip("mujoco", reason="mujoco not installed") - try: - from mujoco_uni.batch_env import BatchEnvPool as _ # noqa: F401 - except Exception: - pytest.skip("mujoco_uni.batch_env not available") - ensure_registries() - env = registry.make( - "Go2FootStand", - sim_backend="mujoco", - num_envs=1, - env_cfg_override={ - "reward_config": RewardConfig( - scales={"height": 1.0}, - tracking_sigma=0.25, - base_height_target=0.3, - ) - }, - ) +@pytest.mark.skipif( + importlib.util.find_spec("drakeuni") is None, + reason="optional DrakeUni batch runtime is not installed", +) +def test_footstand_drake_real_runtime_when_available() -> None: + registry.ensure_registries() + env = _make_env("drake") try: - state = env.init_state() - assert state.obs["critic"][0, -1] > 0.1 + obs, _ = env.reset(seed=17) + assert {name: value.shape for name, value in obs.items()} == { + "obs": (2, 675), + "critic": (2, 724), + } + state = env.step(np.zeros((2, NUM_ACTIONS), dtype=np.float32)) + assert np.isfinite(state.reward).all() finally: env.close() + + +def test_footstand_production_terms_do_not_leak_backend_or_layout() -> None: + source = (ROOT_DIR / "src/unilab/tasks/locomotion/go2/footstand.py").read_text(encoding="utf-8") + for forbidden in ( + "._backend", + "getattr(", + "hasattr(", + "ASSETS_ROOT_PATH", + " qpos", + " qvel", + ): + assert forbidden not in source From d2ff3df0481fe62b969b394f99a9fc8662238482 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Wed, 19 Aug 2026 11:02:26 +0800 Subject: [PATCH 085/169] test: add Isaac Lab manager migration fixture [skip ci] (#1216) --- tests/fixtures/__init__.py | 1 + tests/fixtures/isaac_lab_cartpole/README.md | 22 ++ tests/fixtures/isaac_lab_cartpole/__init__.py | 3 + .../fixtures/isaac_lab_cartpole/cartpole.xml | 16 ++ .../isaac_lab_cartpole/conf/config.yaml | 117 ++++++++++ tests/fixtures/isaac_lab_cartpole/task.py | 220 ++++++++++++++++++ .../test_isaac_lab_migration_fixture.py | 188 +++++++++++++++ 7 files changed, 567 insertions(+) create mode 100644 tests/fixtures/__init__.py create mode 100644 tests/fixtures/isaac_lab_cartpole/README.md create mode 100644 tests/fixtures/isaac_lab_cartpole/__init__.py create mode 100644 tests/fixtures/isaac_lab_cartpole/cartpole.xml create mode 100644 tests/fixtures/isaac_lab_cartpole/conf/config.yaml create mode 100644 tests/fixtures/isaac_lab_cartpole/task.py create mode 100644 tests/managers/test_isaac_lab_migration_fixture.py diff --git a/tests/fixtures/__init__.py b/tests/fixtures/__init__.py new file mode 100644 index 000000000..fae6e65c8 --- /dev/null +++ b/tests/fixtures/__init__.py @@ -0,0 +1 @@ +"""Runnable external-task migration fixtures.""" diff --git a/tests/fixtures/isaac_lab_cartpole/README.md b/tests/fixtures/isaac_lab_cartpole/README.md new file mode 100644 index 000000000..4e9d7a9b6 --- /dev/null +++ b/tests/fixtures/isaac_lab_cartpole/README.md @@ -0,0 +1,22 @@ +# Isaac Lab Cartpole migration fixture + +This test-only fixture is derived from Isaac Lab commit +`b0542fe2d45bf91c4e1d9ef6952b9c709c80b4e8`, specifically its manager-based +Cartpole config and `joint_pos_target_l2` reward. The source is BSD-3-Clause; +the retained notice is in `task.py`. + +The fixture is evidence for #1042, not a production task or a claim that all +Isaac Lab tasks run unchanged. + +| Surface | Status | Migration delta | +| --- | --- | --- | +| Manager/term names, `func + params`, dict order | Compatible | Imports change from `isaaclab` to `unilab`; all 12 terms keep their source names and order. | +| Term math and buffers | Adapted | `torch.Tensor` and Torch ops become `np.ndarray` and NumPy ops. | +| Config container | Adapted | Nested `@configclass` objects become one Hydra owner YAML, materialized as a plain `ManagerBasedRlEnvCfg`. | +| Scene | Adapted | USD/`InteractiveSceneCfg` becomes a minimal task-owned MJCF plus `SceneCfg`/`EntityCfg`. | +| Joint effort action | Adapted, fixture-local | A thin action adapter writes the entity control buffer; it is deliberately not exported as a public built-in. | +| Reset mutation | Adapted, fixture-local | The same two reset terms write through the scoped entity reset transaction. | +| PhysX, USD, Isaac renderer | Unsupported | No dependency or fallback is provided. | + +The fixture adds no production registry entry, public manager/backend/lifecycle +contract, training-script branch, runner/IPC path, or external runtime dependency. diff --git a/tests/fixtures/isaac_lab_cartpole/__init__.py b/tests/fixtures/isaac_lab_cartpole/__init__.py new file mode 100644 index 000000000..93cde0fda --- /dev/null +++ b/tests/fixtures/isaac_lab_cartpole/__init__.py @@ -0,0 +1,3 @@ +"""Isaac Lab Cartpole migration fixture.""" + +from .task import FIXTURE_ENV_NAME as FIXTURE_ENV_NAME diff --git a/tests/fixtures/isaac_lab_cartpole/cartpole.xml b/tests/fixtures/isaac_lab_cartpole/cartpole.xml new file mode 100644 index 000000000..5a7d9e389 --- /dev/null +++ b/tests/fixtures/isaac_lab_cartpole/cartpole.xml @@ -0,0 +1,16 @@ + + diff --git a/tests/fixtures/isaac_lab_cartpole/conf/config.yaml b/tests/fixtures/isaac_lab_cartpole/conf/config.yaml new file mode 100644 index 000000000..202b6ab83 --- /dev/null +++ b/tests/fixtures/isaac_lab_cartpole/conf/config.yaml @@ -0,0 +1,117 @@ +# Derived from Isaac Lab b0542fe2d45bf91c4e1d9ef6952b9c709c80b4e8. +# Copyright (c) 2022-2026, The Isaac Lab Project Developers. +# Modified by UniLab for Hydra and the fixture-local MJCF adapter; BSD-3-Clause. +# Hydra is the fixture's only task-configuration owner. This declaration mirrors +# Isaac Lab's CartpoleEnvCfg term names and insertion order. +training: + task_name: IsaacLabCartpoleFixture + sim_backend: mujoco + +env: + scene: + model_file: tests/fixtures/isaac_lab_cartpole/cartpole.xml + entities: + robot: + root_body_name: cart + joint_names: [slider_to_cart, cart_to_pole] + actuator_names: [slider_to_cart] + sim_dt: 0.008333333333333333 + ctrl_dt: 0.016666666666666666 + max_episode_seconds: 5.0 + seed: 7 + observations: + policy: + _target_: unilab.managers.ObservationGroupCfg + enable_corruption: false + concatenate_terms: true + terms: + joint_pos_rel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_pos_rel + joint_vel_rel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_vel_rel + actions: + joint_effort: + _target_: tests.fixtures.isaac_lab_cartpole.task.JointEffortActionCfg + entity_name: robot + actuator_names: [slider_to_cart] + scale: 100.0 + events: + reset_cart_position: + _target_: unilab.managers.EventTermCfg + func: tests.fixtures.isaac_lab_cartpole.task.reset_joints_by_offset + mode: reset + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + joint_names: slider_to_cart + position_range: [-1.0, 1.0] + velocity_range: [-0.5, 0.5] + reset_pole_position: + _target_: unilab.managers.EventTermCfg + func: tests.fixtures.isaac_lab_cartpole.task.reset_joints_by_offset + mode: reset + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + joint_names: cart_to_pole + position_range: [-0.7853981633974483, 0.7853981633974483] + velocity_range: [-0.7853981633974483, 0.7853981633974483] + terminations: + time_out: + _target_: unilab.managers.TerminationTermCfg + func: unilab.envs.mdp.time_out + time_out: true + cart_out_of_bounds: + _target_: unilab.managers.TerminationTermCfg + func: tests.fixtures.isaac_lab_cartpole.task.joint_pos_out_of_manual_limit + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + joint_names: slider_to_cart + bounds: [-3.0, 3.0] + policy_observation_group: policy + critic_observation_group: null + scale_rewards_by_dt: true + +reward: + alive: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.mdp.is_alive + weight: 1.0 + terminating: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.mdp.is_terminated + weight: -2.0 + pole_pos: + _target_: unilab.managers.RewardTermCfg + func: tests.fixtures.isaac_lab_cartpole.task.joint_pos_target_l2 + weight: -1.0 + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + joint_names: cart_to_pole + target: 0.0 + cart_vel: + _target_: unilab.managers.RewardTermCfg + func: tests.fixtures.isaac_lab_cartpole.task.joint_vel_l1 + weight: -0.01 + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + joint_names: slider_to_cart + pole_vel: + _target_: unilab.managers.RewardTermCfg + func: tests.fixtures.isaac_lab_cartpole.task.joint_vel_l1 + weight: -0.005 + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + joint_names: cart_to_pole diff --git a/tests/fixtures/isaac_lab_cartpole/task.py b/tests/fixtures/isaac_lab_cartpole/task.py new file mode 100644 index 000000000..2e80e04c9 --- /dev/null +++ b/tests/fixtures/isaac_lab_cartpole/task.py @@ -0,0 +1,220 @@ +# Derived from Isaac Lab b0542fe2d45bf91c4e1d9ef6952b9c709c80b4e8, +# source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/cartpole. +# Copyright (c) 2022-2026, The Isaac Lab Project Developers. +# Modified by UniLab for NumPy and the fixture-local MJCF/entity adapter; BSD-3-Clause. +"""NumPy terms and adapters for the Isaac Lab Cartpole migration fixture.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from numbers import Real +from typing import TYPE_CHECKING, cast + +import numpy as np + +from unilab.base import registry +from unilab.envs import ManagerBasedRlEnvCfg, make_manager_based_rl_env +from unilab.managers import ActionTerm, ActionTermCfg +from unilab.managers.scene_entity_config import SceneEntityCfg + +if TYPE_CHECKING: + from unilab.base.entity import Entity + from unilab.managers._types import ManagerBasedRlEnv + + +FIXTURE_ENV_NAME = "IsaacLabCartpoleFixture" + + +def _finite_real(value: Real, *, label: str) -> float: + if isinstance(value, (bool, np.bool_)) or not isinstance(value, Real): + raise TypeError(f"{label} must be a real number") + result = float(value) + if not np.isfinite(result): + raise ValueError(f"{label} must be finite") + return result + + +def _range(value: tuple[float, float] | list[float], *, label: str) -> tuple[float, float]: + if isinstance(value, (str, bytes)) or not isinstance(value, (tuple, list)): + raise TypeError(f"{label} must be a two-value range") + if len(value) != 2: + raise ValueError(f"{label} must contain two values") + lower = _finite_real(value[0], label=f"{label}[0]") + upper = _finite_real(value[1], label=f"{label}[1]") + if lower > upper: + raise ValueError(f"{label} lower bound {lower} exceeds upper bound {upper}") + return lower, upper + + +@dataclass(kw_only=True) +class JointEffortActionCfg(ActionTermCfg): + """Fixture-local adapter for Isaac Lab's ``JointEffortActionCfg``.""" + + actuator_names: tuple[str, ...] | list[str] + scale: float = 1.0 + + def build(self, env: ManagerBasedRlEnv) -> JointEffortAction: + return JointEffortAction(self, env) + + +class JointEffortAction(ActionTerm): + """Scale policy actions and write entity-local actuator efforts.""" + + cfg: JointEffortActionCfg + + def __init__(self, cfg: JointEffortActionCfg, env: ManagerBasedRlEnv): + super().__init__(cfg, env) + if cfg.clip is not None: + raise NotImplementedError( + "IsaacLabCartpoleFixture JointEffortAction does not support clip" + ) + if isinstance(cfg.actuator_names, (str, bytes)) or not isinstance( + cfg.actuator_names, (tuple, list) + ): + raise TypeError("JointEffortActionCfg actuator_names must be an ordered sequence") + actuator_ids, actuator_names = self._entity.find_actuators( + cfg.actuator_names, + preserve_order=True, + ) + if not actuator_ids: + raise ValueError( + "JointEffortActionCfg actuator_names resolved no actuators; " + f"patterns={list(cfg.actuator_names)}" + ) + self._actuator_ids = np.asarray(actuator_ids, dtype=np.intp) + self._actuator_ids.setflags(write=False) + self._actuator_names = tuple(actuator_names) + self._scale = _finite_real(cfg.scale, label="JointEffortActionCfg scale") + self._raw_actions = np.zeros((self.num_envs, len(actuator_ids)), dtype=np.float32) + self._processed_actions = np.zeros_like(self._raw_actions) + + @property + def action_dim(self) -> int: + return self._raw_actions.shape[1] + + @property + def raw_action(self) -> np.ndarray: + return self._raw_actions + + def process_actions(self, actions: np.ndarray) -> None: + if not isinstance(actions, np.ndarray): + raise TypeError( + "IsaacLabCartpoleFixture JointEffortAction expected np.ndarray, " + f"received {type(actions).__name__}" + ) + if actions.shape != self._raw_actions.shape: + raise ValueError( + "IsaacLabCartpoleFixture JointEffortAction expected shape " + f"{self._raw_actions.shape}, received {actions.shape}" + ) + if not np.isfinite(actions).all(): + raise ValueError("IsaacLabCartpoleFixture JointEffortAction received NaN or Inf") + np.copyto(self._raw_actions, actions) + np.multiply(actions, self._scale, out=self._processed_actions) + + def apply_actions(self) -> None: + self._entity.data.write_ctrl( + self._processed_actions, + actuator_ids=self._actuator_ids, + ) + + def reset(self, env_ids: np.ndarray | slice | None = None) -> None: + ids = slice(None) if env_ids is None else env_ids + self._raw_actions[ids] = 0.0 + self._processed_actions[ids] = 0.0 + + +def reset_joints_by_offset( + env: ManagerBasedRlEnv, + env_ids: np.ndarray | None, + position_range: tuple[float, float] | list[float], + velocity_range: tuple[float, float] | list[float], + asset_cfg: SceneEntityCfg, +) -> None: + """Port Isaac Lab's joint-offset reset through the entity reset transaction.""" + if env_ids is None: + raise ValueError("reset_joints_by_offset requires concrete environment IDs") + position_lower, position_upper = _range(position_range, label="position_range") + velocity_lower, velocity_upper = _range(velocity_range, label="velocity_range") + asset = cast("Entity", env.scene[asset_cfg.name]) + joint_ids = asset_cfg.joint_ids + default_position = asset.data.default_joint_pos[env_ids][:, joint_ids] + default_velocity = asset.data.default_joint_vel[env_ids][:, joint_ids] + position = default_position + env.rng.uniform( + position_lower, + position_upper, + default_position.shape, + ) + velocity = default_velocity + env.rng.uniform( + velocity_lower, + velocity_upper, + default_velocity.shape, + ) + asset.write_joint_state_to_sim( + position.astype(default_position.dtype, copy=False), + velocity.astype(default_velocity.dtype, copy=False), + joint_ids=joint_ids, + env_ids=env_ids, + ) + + +def joint_pos_target_l2( + env: ManagerBasedRlEnv, + target: float, + asset_cfg: SceneEntityCfg, +) -> np.ndarray: + """Penalize wrapped joint-position deviation from a target value.""" + target_value = _finite_real(target, label="joint_pos_target_l2 target") + asset = cast("Entity", env.scene[asset_cfg.name]) + joint_pos = asset.data.joint_pos[:, asset_cfg.joint_ids] + wrapped = np.remainder(joint_pos + math.pi, 2.0 * math.pi) - math.pi + return np.sum(np.square(wrapped - target_value), axis=1) + + +def joint_vel_l1( + env: ManagerBasedRlEnv, + asset_cfg: SceneEntityCfg, +) -> np.ndarray: + """Penalize the absolute velocity of selected joints.""" + asset = cast("Entity", env.scene[asset_cfg.name]) + return np.sum(np.abs(asset.data.joint_vel[:, asset_cfg.joint_ids]), axis=1) + + +def joint_pos_out_of_manual_limit( + env: ManagerBasedRlEnv, + bounds: tuple[float, float] | list[float], + asset_cfg: SceneEntityCfg, +) -> np.ndarray: + """Terminate when a selected joint leaves the configured manual bounds.""" + lower, upper = _range(bounds, label="joint_pos_out_of_manual_limit bounds") + asset = cast("Entity", env.scene[asset_cfg.name]) + joint_pos = asset.data.joint_pos[:, asset_cfg.joint_ids] + return np.any((joint_pos < lower) | (joint_pos > upper), axis=1) + + +def register_fixture() -> None: + """Register the fixture without adding it to the production task package.""" + if registry.contains(FIXTURE_ENV_NAME): + return + registry.register_env_config(FIXTURE_ENV_NAME, ManagerBasedRlEnvCfg) + registry.register_env( + FIXTURE_ENV_NAME, + make_manager_based_rl_env, + sim_backend="mujoco", + ) + + +register_fixture() + + +__all__ = [ + "FIXTURE_ENV_NAME", + "JointEffortAction", + "JointEffortActionCfg", + "joint_pos_out_of_manual_limit", + "joint_pos_target_l2", + "joint_vel_l1", + "register_fixture", + "reset_joints_by_offset", +] diff --git a/tests/managers/test_isaac_lab_migration_fixture.py b/tests/managers/test_isaac_lab_migration_fixture.py new file mode 100644 index 000000000..eb9de528a --- /dev/null +++ b/tests/managers/test_isaac_lab_migration_fixture.py @@ -0,0 +1,188 @@ +"""End-to-end evidence for the Isaac Lab Manager-Based migration fixture.""" + +from __future__ import annotations + +from dataclasses import fields, is_dataclass +from pathlib import Path +from typing import Any + +import numpy as np +import pytest +from hydra import compose, initialize_config_dir +from hydra.core.global_hydra import GlobalHydra +from omegaconf import DictConfig, OmegaConf + +from tests.fixtures.isaac_lab_cartpole import FIXTURE_ENV_NAME +from tests.fixtures.isaac_lab_cartpole import task as fixture +from unilab.base import registry +from unilab.base.config_materialization import apply_cfg_overrides +from unilab.envs import ManagerBasedRlEnv, ManagerBasedRLEnvCfg, ManagerBasedRlEnvCfg +from unilab.training import BackendAdapter + +ROOT_DIR = Path(__file__).parents[2] +FIXTURE_DIR = ROOT_DIR / "tests" / "fixtures" / "isaac_lab_cartpole" + + +def _compose_fixture() -> DictConfig: + GlobalHydra.instance().clear() + with initialize_config_dir(config_dir=str(FIXTURE_DIR / "conf"), version_base="1.3"): + return compose("config") + + +def _materialize_fixture() -> tuple[DictConfig, ManagerBasedRlEnvCfg, dict[str, Any]]: + hydra_cfg = _compose_fixture() + override = BackendAdapter(hydra_cfg, root_dir=ROOT_DIR).build_task_env_cfg_override() + env_cfg = registry.materialize_env_config(FIXTURE_ENV_NAME) + assert isinstance(env_cfg, ManagerBasedRlEnvCfg) + apply_cfg_overrides(env_cfg, override) + env_cfg.validate() + return hydra_cfg, env_cfg, override + + +def _assert_plain(value: Any) -> None: + assert not OmegaConf.is_config(value) + if is_dataclass(value) and not isinstance(value, type): + for field in fields(value): + _assert_plain(getattr(value, field.name)) + elif isinstance(value, dict): + for key, item in value.items(): + _assert_plain(key) + _assert_plain(item) + elif isinstance(value, (tuple, list)): + for item in value: + _assert_plain(item) + + +def test_fixture_hydra_owner_materializes_source_order_as_plain_manager_cfg() -> None: + hydra_cfg, env_cfg, _ = _materialize_fixture() + + assert hydra_cfg.training.task_name == FIXTURE_ENV_NAME + assert hydra_cfg.training.sim_backend == "mujoco" + assert ManagerBasedRLEnvCfg is ManagerBasedRlEnvCfg + assert list(env_cfg.observations) == ["policy"] + assert list(env_cfg.observations["policy"].terms) == ["joint_pos_rel", "joint_vel_rel"] + assert list(env_cfg.actions) == ["joint_effort"] + assert list(env_cfg.events) == ["reset_cart_position", "reset_pole_position"] + assert list(env_cfg.rewards) == [ + "alive", + "terminating", + "pole_pos", + "cart_vel", + "pole_vel", + ] + assert list(env_cfg.terminations) == ["time_out", "cart_out_of_bounds"] + assert env_cfg.actions["joint_effort"].scale == pytest.approx(100.0) + assert env_cfg.sim_dt == pytest.approx(1.0 / 120.0) + assert env_cfg.ctrl_dt == pytest.approx(1.0 / 60.0) + assert env_cfg.max_episode_seconds == pytest.approx(5.0) + assert env_cfg.policy_observation_group == "policy" + assert env_cfg.critic_observation_group is None + _assert_plain(env_cfg) + + +def test_fixture_real_mujoco_reset_step_and_partial_reset() -> None: + _, _, override = _materialize_fixture() + env = registry.make( + FIXTURE_ENV_NAME, + sim_backend="mujoco", + env_cfg_override=override, + num_envs=8, + ) + assert isinstance(env, ManagerBasedRlEnv) + try: + state = env.init_state() + assert env.obs_groups_spec == {"obs": 4} + assert env.action_space.shape == (1,) + assert state.obs["obs"].shape == (8, 4) + assert np.isfinite(state.obs["obs"]).all() + + before = env.scene["robot"].data.joint_pos.copy() + reset_ids = np.asarray([1, 6], dtype=np.int32) + reset_obs, _ = env.reset(env_ids=reset_ids) + after = env.scene["robot"].data.joint_pos.copy() + assert reset_obs["obs"].shape == (2, 4) + np.testing.assert_array_equal(after[[0, 2, 3, 4, 5, 7]], before[[0, 2, 3, 4, 5, 7]]) + assert np.all(np.abs(after[reset_ids, 0]) <= 1.0) + assert np.all(np.abs(after[reset_ids, 1]) <= 0.25 * np.pi) + + state = env.step(np.zeros((8, 1), dtype=np.float32)) + assert state.obs["obs"].shape == (8, 4) + assert state.reward.shape == (8,) + assert state.terminated.shape == (8,) + assert state.truncated.shape == (8,) + assert np.isfinite(state.obs["obs"]).all() + assert np.isfinite(state.reward).all() + + robot = env.scene["robot"] + pole_pos = robot.data.joint_pos[:, 1] + expected = 1.0 - np.square(np.remainder(pole_pos + np.pi, 2.0 * np.pi) - np.pi) + expected -= 0.01 * np.abs(robot.data.joint_vel[:, 0]) + expected -= 0.005 * np.abs(robot.data.joint_vel[:, 1]) + np.testing.assert_allclose(state.reward, expected * env.step_dt, rtol=1e-5, atol=1e-6) + finally: + env.close() + + +def test_fixture_missing_actuator_fails_during_cold_path_binding() -> None: + _, _, override = _materialize_fixture() + override["actions"]["joint_effort"]["actuator_names"] = ["missing_actuator"] + + with pytest.raises( + ValueError, + match="Not all entity selector regular expressions matched.*missing_actuator", + ): + registry.make( + FIXTURE_ENV_NAME, + sim_backend="mujoco", + env_cfg_override=override, + num_envs=2, + ) + + +def test_fixture_stays_test_only_and_has_no_external_runtime_imports() -> None: + source = (FIXTURE_DIR / "task.py").read_text(encoding="utf-8") + executable = "\n".join( + line for line in source.splitlines() if not line.lstrip().startswith("#") + ) + for forbidden in ( + "import torch", + "from torch", + "import isaaclab", + "from isaaclab", + "unilab.ipc", + "unilab.training", + ): + assert forbidden not in executable + assert Path(fixture.__file__).is_relative_to(ROOT_DIR / "tests" / "fixtures") + assert "tests.fixtures" not in tuple(registry._DEFAULT_REGISTRY_PACKAGES) + assert set(registry._envs[FIXTURE_ENV_NAME].env_factory_dict) == {"mujoco"} + + +def test_fixture_local_term_math_matches_isaac_source() -> None: + class _Data: + joint_pos = np.asarray([[0.2, 3.5], [-0.3, -3.4]], dtype=np.float32) + joint_vel = np.asarray([[0.4, -0.7], [-0.2, 0.9]], dtype=np.float32) + + class _Entity: + data = _Data() + + class _Scene(dict): + pass + + env = type("FixtureMathEnv", (), {"scene": _Scene(robot=_Entity())})() + pole = fixture.SceneEntityCfg("robot", joint_ids=[1]) + slider = fixture.SceneEntityCfg("robot", joint_ids=[0]) + + expected_wrapped = np.remainder(_Data.joint_pos[:, 1] + np.pi, 2.0 * np.pi) - np.pi + np.testing.assert_allclose( + fixture.joint_pos_target_l2(env, target=0.0, asset_cfg=pole), + np.square(expected_wrapped), + ) + np.testing.assert_allclose( + fixture.joint_vel_l1(env, asset_cfg=pole), + np.abs(_Data.joint_vel[:, 1]), + ) + np.testing.assert_array_equal( + fixture.joint_pos_out_of_manual_limit(env, bounds=(-0.25, 0.25), asset_cfg=slider), + np.asarray([False, True]), + ) From c347da604d1cf9e09975c6130de0b461fbaf0779 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Wed, 19 Aug 2026 11:13:55 +0800 Subject: [PATCH 086/169] docs: align Isaac Lab manager migration guide [skip ci] (#1218) --- .../3-framework_migration/0-index.md | 2 +- .../3-framework_migration/1-from_isaac_lab.md | 240 ++++++++++++------ .../3-framework_migration/0-index.md | 2 +- .../3-framework_migration/1-from_isaac_lab.md | 195 +++++++++----- 4 files changed, 292 insertions(+), 147 deletions(-) diff --git a/docs/sphinx/source/en/3-deployment/3-framework_migration/0-index.md b/docs/sphinx/source/en/3-deployment/3-framework_migration/0-index.md index 85047eeb2..6d62a0de5 100644 --- a/docs/sphinx/source/en/3-deployment/3-framework_migration/0-index.md +++ b/docs/sphinx/source/en/3-deployment/3-framework_migration/0-index.md @@ -9,7 +9,7 @@ contract-driven layout. :::{grid-item-card} From Isaac Lab :link: 1-from_isaac_lab :link-type: doc -Map GPU-resident task structure to UniLab's CPU sim and learner split. +Keep Manager-Based terms while adapting Hydra config, NumPy execution, and scene access. ::: :::{grid-item-card} From Legged Gym diff --git a/docs/sphinx/source/en/3-deployment/3-framework_migration/1-from_isaac_lab.md b/docs/sphinx/source/en/3-deployment/3-framework_migration/1-from_isaac_lab.md index f2a4c6fbb..73168650b 100644 --- a/docs/sphinx/source/en/3-deployment/3-framework_migration/1-from_isaac_lab.md +++ b/docs/sphinx/source/en/3-deployment/3-framework_migration/1-from_isaac_lab.md @@ -1,93 +1,169 @@ # Migrating from Isaac Lab -If you have an Isaac Lab task you want to run in UniLab, this page tells -you what stays the same, what changes, and where the sharp edges are. - -## What stays the same - -- Gymnasium-style env interface (`reset`, `step`, `obs/reward/info`). -- Hydra-based configuration. Most of your existing YAML can be ported with - field-name remapping. -- The general idea of a "task" that composes scene + reward + DR + obs. -- PPO as the default algo — UniLab ships RSL-RL's PPO out of the box. +Port an Isaac Lab Manager-Based task to UniLab by keeping its manager and term +structure, then adapting configuration, numeric execution, and scene access at +their owner boundaries. Do not rewrite it as a monolithic `NpEnv` subclass. + +This is source-compatible migration, not a promise that an arbitrary Isaac Lab +task runs unchanged. The target path is: + +```text +Hydra owner YAML + -> plain ManagerBasedRlEnvCfg + -> Registry + make_manager_based_rl_env + -> ManagerBasedRlEnv on the NumPy/SimBackend runtime + -> NpEnvState for the existing training and IPC path +``` -## What changes +## Compatibility boundary ```{list-table} :header-rows: 1 -:widths: 30 35 35 - -* - Isaac Lab concept - - UniLab equivalent - - Notes -* - `DirectRLEnv` - - `unilab.base.np_env.NpEnv` - - UniLab obs is always a **dict**, not a tensor. -* - `RigidBody.cfg` - - Task-side asset import + scene composition - - See {doc}`../../4-developer_guide/1-architecture/4-scene_composition`. -* - GPU PhysX backend - - CPU MuJoCo / Motrix + GPU learner - - Architectural inversion — see below. -* - `RandomizationCfg` - - {doc}`../../4-developer_guide/2-contracts/4-dr_contract` - - UniLab DR runs in cold-path resampling only. -* - `RewardManager` chains - - Reward composition in env, plus - `unilab.training.reward` bookkeeping - - Reward terms still keyed for component-wise logging. -* - `EventCfg` event-driven hooks - - Phase + curriculum + DR providers - - Hooks are explicit, not implicit. +:widths: 28 24 48 + +* - Isaac Lab surface + - UniLab status + - Migration rule +* - Manager categories, term names and dictionary order + - Compatible + - Keep observation, action, event, reward, termination, command, and + curriculum terms in the same order. +* - Function/class terms and `func + params` + - Compatible + - Change imports to `unilab.managers`; keep term boundaries and partial + `reset(env_ids)` semantics. +* - `ManagerBasedRLEnv` / `ManagerBasedRLEnvCfg` + - Compatible spelling aliases + - The canonical UniLab names are `ManagerBasedRlEnv` and + `ManagerBasedRlEnvCfg`; the aliases point to the same implementation. +* - Tensor values and operations + - Adapted + - Replace `torch.Tensor` with `np.ndarray` and use vectorized NumPy. There is + no manager-facing device API. +* - Nested `@configclass` task configuration + - Adapted + - Move the complete task declaration to one Hydra owner YAML. `_target_` + selects concrete config dataclasses and dotted `func` values select terms. +* - `InteractiveSceneCfg`, USD, and PhysX views + - Adapted or unsupported + - Declare a task-owned `SceneCfg` and `EntityCfg`; access state and control + only through `SceneEntityCfg` and the public entity facade. Unsupported + capabilities raise during cold-path binding. +* - Omniverse, Isaac renderer, and Torch/PhysX mutation + - Unsupported + - UniLab does not install or silently emulate these runtimes. +``` + +The normative boundary is +{doc}`ADR-0006 `. Only +surfaces backed by registration, configuration, and tests should be described +as compatible. + +## Migration procedure + +### 1. Inventory the source task + +Pin the Isaac Lab revision and list the source manager groups, term names, term +order, parameters, observation dimensions, action dimensions, reset behavior, +and episode timing. Classify each dependency before writing code: + +- reuse an existing `unilab.managers` config or `unilab.envs.mdp` term; +- adapt a task-specific term from Torch to NumPy; +- stop if the term requires a capability absent from the public entity or + `SimBackend` contract. + +Do not probe backend objects with `getattr`/`hasattr`, return zeros, or route the +task back to a legacy environment. + +### 2. Port scene and assets on the cold path + +Replace Isaac Lab's USD/`InteractiveSceneCfg` declaration with a task-owned +`SceneCfg`. Declare every entity and selector needed by terms. The +`SceneEntityCfg` selector resolves names and regular expressions once during +materialization; reset and step reuse cached IDs and NumPy views. + +The Cartpole fixture uses a minimal task-owned MJCF asset. More complex assets +must follow +{doc}`scene composition <../../4-developer_guide/1-architecture/4-scene_composition>` +and the selected backend's formal capabilities. + +### 3. Port term code, not the manager structure + +Keep each function/class term and its parameters. Replace Torch types and +operators mechanically with NumPy, preserve batch shapes, and return one value +per environment where the source term does. Stateful terms resolve selectors +and allocate buffers in their constructor, then update only NumPy buffers on +the hot path. + +Python owns term implementations and reusable config dataclasses. It must not +hold a second task-specific list of enabled terms or default weights. + +### 4. Make Hydra the only task configuration owner + +Declare scene, timing, groups, terms, concrete config types, callables, +parameters, weights, and observation mapping in the owner YAML. For example: + +```yaml +env: + observations: + policy: + _target_: unilab.managers.ObservationGroupCfg + terms: + joint_pos_rel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_pos_rel + terminations: + time_out: + _target_: unilab.managers.TerminationTermCfg + func: unilab.envs.mdp.time_out + time_out: true + policy_observation_group: policy + critic_observation_group: null + +reward: + alive: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.mdp.is_alive + weight: 1.0 +``` + +Hydra composition materializes this declaration into plain typed config on the +cold path. Unknown fields, unresolved `_target_`/`func` references, and wrong +config types fail before reset or step. Direct Python config construction is +reserved for focused lower-level tests. + +### 5. Register one generic runtime path + +The task module registers `ManagerBasedRlEnvCfg` and +`make_manager_based_rl_env` for each backend that the repository actually +supports. Backend owner YAMLs carry backend identity and tuning. Users select +the composed owner through the normal CLI, for example: + +```bash +uv run train --algo ppo --task --sim mujoco ``` -## The architectural inversion - -Isaac Lab places the simulator on GPU and lets you batch thousands of -envs in PhysX. UniLab places the simulator on CPU (often multithread) and -batches across worker **processes**, sharing memory with a single GPU -learner. - -Implications: - -- **Per-env step time** in UniLab is comparable or worse than Isaac on a - single env. **Throughput** comes from process parallelism + asynchrony - (see `unilab.ipc.async_runner`). -- You can run on **MPS, ROCm, XPU** as the learner device — Isaac is - CUDA-only. -- **No GPU contention** between simulator and learner — your trainer's - memory usage is predictable. - -## Step-by-step migration - -1. **Audit observations.** Make sure every observation key is a vector - you can express without GPU PhysX queries. If not, add a state - estimator or move the query to cold path. -2. **Port the asset.** UniLab consumes MJCF as its source of truth. If - you have USD, convert to MJCF first. -3. **Port the env.** Subclass `unilab.base.np_env.NpEnv`. Move - reward computation into the env's `compute_reward()`. -4. **Port the YAML.** Map Isaac Lab's `EnvCfg` fields to UniLab task owner - YAML following the table in - {doc}`5-task_config_translation`. -5. **Port the reward.** Use the cookbook at - {doc}`6-reward_porting`. -6. **Validate.** Train a small run, compare reward curves against your - Isaac baseline. - -## What you'll miss (and how to compensate) - -- **Isaac Sim renderer.** Use Motrix's headless video export or build a - viser scene (`unilab.visualization.viser_scene`). -- **Per-env tensor obs.** UniLab gives you dict-of-arrays; wrap with your - own `obs_to_tensor` if you need a tensor. -- **Built-in GPU-side DR.** UniLab DR is CPU-side per process. For most - tasks this is plenty; for extreme parallelism use more worker - processes. +Do not add a task-specific training-script branch, environment factory, runner, +or IPC path. + +### 6. Validate near each adaptation + +Test Hydra composition and typed materialization, term order and math, selector +failure, observation/action shapes, partial reset, and at least one real +registered backend transition. Compare behavior with the pinned source task; +benchmark only after semantic migration is complete. + +## Repository evidence + +`tests/fixtures/isaac_lab_cartpole/` ports the Manager-Based Cartpole task from +Isaac Lab commit `b0542fe2d45bf91c4e1d9ef6952b9c709c80b4e8`. It preserves all +12 source term names and their order while adapting Torch to NumPy, nested +config objects to Hydra YAML, and the scene/action/reset boundaries to a +fixture-local MJCF implementation. It is test-only evidence, not a production +task or a blanket Isaac Lab support claim. ## See also -- {doc}`2-from_legged_gym` -- {doc}`3-from_rsl_rl` -- {doc}`5-task_config_translation` -- {doc}`6-reward_porting` +- {doc}`Manager-Based API <../../4-developer_guide/1-architecture/6-manager_based_api>` +- {doc}`Environment contract <../../4-developer_guide/2-contracts/1-env_contract>` +- {doc}`ADR-0006 ` diff --git a/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/0-index.md b/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/0-index.md index 3c60032d6..42d467c0c 100644 --- a/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/0-index.md +++ b/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/0-index.md @@ -8,7 +8,7 @@ :::{grid-item-card} 从 Isaac Lab 迁移 :link: 1-from_isaac_lab :link-type: doc -把 GPU 常驻的任务结构映射到 UniLab 的 CPU sim 与 learner 拆分。 +保留 Manager-Based term 结构，适配 Hydra 配置、NumPy 执行和场景访问。 ::: :::{grid-item-card} 从 Legged Gym 迁移 diff --git a/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/1-from_isaac_lab.md b/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/1-from_isaac_lab.md index 7eaaf7f6b..f6914d9ee 100644 --- a/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/1-from_isaac_lab.md +++ b/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/1-from_isaac_lab.md @@ -1,83 +1,152 @@ # 从 Isaac Lab 迁移 -如果你有一个想在 UniLab 中运行的 Isaac Lab 任务，本页会告诉你哪些保持不变、 -哪些会改变，以及锋利的边角在哪里。 +把 Isaac Lab Manager-Based task 迁入 UniLab 时，应保留 manager 与 term 结构，只在各自 +owner 边界适配配置、数值执行和场景访问；不要把 task 重写成单体 `NpEnv` 子类。 -## 哪些保持不变 +这是基于源码的兼容迁移，不代表任意 Isaac Lab task 都能不修改直接运行。目标路径是： -- Gymnasium 风格的 env 接口（`reset`、`step`、`obs/reward/info`）。 -- 基于 Hydra 的配置。你现有的大部分 YAML 可以通过字段名重映射来移植。 -- "任务"由 scene + reward + DR + obs 组合而成这一总体思路。 -- PPO 作为默认算法 —— UniLab 开箱即带 RSL-RL 的 PPO。 +```text +Hydra owner YAML + -> plain ManagerBasedRlEnvCfg + -> Registry + make_manager_based_rl_env + -> NumPy/SimBackend runtime 上的 ManagerBasedRlEnv + -> 交给现有 training 和 IPC 路径的 NpEnvState +``` -## 哪些会改变 +## 兼容边界 ```{list-table} :header-rows: 1 -:widths: 30 35 35 - -* - Isaac Lab 概念 - - UniLab 对应物 - - 备注 -* - `DirectRLEnv` - - `unilab.base.np_env.NpEnv` - - UniLab 的 obs 始终是 **dict**，而不是 tensor。 -* - `RigidBody.cfg` - - 任务侧的 asset 导入 + 场景组合 - - 参见 {doc}`../../4-developer_guide/1-architecture/4-scene_composition`。 -* - GPU PhysX 后端 - - CPU MuJoCo / Motrix + GPU learner - - 架构倒置 —— 见下文。 -* - `RandomizationCfg` - - {doc}`../../4-developer_guide/2-contracts/4-dr_contract` - - UniLab 的 DR 只在冷路径重采样中运行。 -* - `RewardManager` 链 - - env 中的 reward 组合，外加 - `unilab.training.reward` 记账 - - reward 项仍然以 key 标识，以便分量级别的日志记录。 -* - `EventCfg` 事件驱动钩子 - - Phase + curriculum + DR provider - - 钩子是显式的，而非隐式的。 +:widths: 28 24 48 + +* - Isaac Lab 表面 + - UniLab 状态 + - 迁移规则 +* - Manager 类别、term 名称和字典顺序 + - Compatible + - 保持 observation、action、event、reward、termination、command 与 + curriculum term 的顺序。 +* - Function/class term 与 `func + params` + - Compatible + - 把 import 改为 `unilab.managers`；保留 term 边界和局部 + `reset(env_ids)` 语义。 +* - `ManagerBasedRLEnv` / `ManagerBasedRLEnvCfg` + - Compatible 拼写 alias + - UniLab canonical 名称是 `ManagerBasedRlEnv` 与 `ManagerBasedRlEnvCfg`；alias + 指向同一份实现。 +* - Tensor 数值与运算 + - Adapted + - 把 `torch.Tensor` 换成 `np.ndarray` 并使用向量化 NumPy；manager-facing + API 没有 device 接口。 +* - 嵌套 `@configclass` task 配置 + - Adapted + - 把完整 task 声明迁入唯一 Hydra owner YAML；用 `_target_` 选择具体 config + dataclass，用 dotted `func` 选择 term。 +* - `InteractiveSceneCfg`、USD 与 PhysX view + - Adapted 或 Unsupported + - 声明 task-owned `SceneCfg` 与 `EntityCfg`；状态和控制只通过 + `SceneEntityCfg` 与公共 entity facade 访问。不支持的能力在冷路径绑定时报错。 +* - Omniverse、Isaac renderer 与 Torch/PhysX mutation + - Unsupported + - UniLab 不安装这些 runtime，也不提供静默模拟或回退。 +``` + +规范边界见 {doc}`ADR-0006 `。 +只有已经被 registry、配置和测试覆盖的表面才能声明为 Compatible。 + +## 迁移步骤 + +### 1. 盘点来源 task + +固定 Isaac Lab revision，并列出来源 manager group、term 名称与顺序、参数、observation +维度、action 维度、reset 行为和 episode timing。写代码前逐项分类： + +- 复用已有 `unilab.managers` config 或 `unilab.envs.mdp` term； +- 把 task-specific term 从 Torch 适配为 NumPy； +- 如果 term 依赖公共 entity 或 `SimBackend` contract 尚未提供的能力，立即停止。 + +不能用 `getattr`/`hasattr` 探测 backend 对象、返回零，或把 task 路由回 legacy env。 + +### 2. 在冷路径迁移 scene 与 asset + +用 task-owned `SceneCfg` 代替 Isaac Lab 的 USD/`InteractiveSceneCfg` 声明，显式声明 term +需要的每个 entity 与 selector。`SceneEntityCfg` 在 materialization 时只解析一次名称和 +正则表达式；reset/step 复用缓存 ID 与 NumPy view。 + +Cartpole fixture 使用最小 task-owned MJCF。更复杂的 asset 必须遵守 +{doc}`场景组合 <../../4-developer_guide/1-architecture/4-scene_composition>`，并只使用所选 +backend 的正式能力。 + +### 3. 迁移 term 代码，不改 manager 结构 + +保留每个 function/class term 及其参数，机械地把 Torch 类型与运算改为 NumPy，保持 batch +shape，并在来源 term 返回每环境数值时继续返回每环境数值。Stateful term 在构造时解析 +selector、分配 buffer，热路径只更新 NumPy buffer。 + +Python 只拥有 term 实现和可复用 config dataclass，不能再保存一份 task-specific term +启停清单或默认 weight。 + +### 4. 让 Hydra 成为唯一 task 配置 owner + +在 owner YAML 中声明 scene、timing、group、term、具体 config 类型、callable、参数、 +weight 和 observation mapping。例如： + +```yaml +env: + observations: + policy: + _target_: unilab.managers.ObservationGroupCfg + terms: + joint_pos_rel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_pos_rel + terminations: + time_out: + _target_: unilab.managers.TerminationTermCfg + func: unilab.envs.mdp.time_out + time_out: true + policy_observation_group: policy + critic_observation_group: null + +reward: + alive: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.mdp.is_alive + weight: 1.0 ``` -## 架构倒置 +Hydra compose 在冷路径把这份声明物化为 plain typed config。未知字段、无法解析的 +`_target_`/`func` 和错误 config 类型都会在 reset/step 之前报错。直接用 Python 构造 +config 只用于 focused 底层测试。 -Isaac Lab 把模拟器放在 GPU 上，让你在 PhysX 中批处理数千个 env。UniLab 把 -模拟器放在 CPU 上（通常是多线程），并跨 worker **进程**做批处理，与单个 GPU -learner 共享内存。 +### 5. 只注册一条通用 runtime 路径 -由此带来的影响： +Task module 为仓库已经实际支持的每个 backend 注册 `ManagerBasedRlEnvCfg` 与 +`make_manager_based_rl_env`。Backend owner YAML 只承载 backend 身份与 tuning。用户通过 +标准 CLI 选择 compose owner，例如： + +```bash +uv run train --algo ppo --task --sim mujoco +``` -- 在单个 env 上，UniLab 的**每个 env 步进时间**与 Isaac 相当甚至更差。 - **吞吐量**来自进程并行 + 异步（参见 `unilab.ipc.async_runner`）。 -- 你可以用 **MPS、ROCm、XPU** 作为 learner 设备 —— Isaac 仅支持 CUDA。 -- 模拟器与 learner 之间**不存在 GPU 争用** —— 你的 trainer 内存占用是可预测的。 +不要增加 task-specific 训练脚本分支、env factory、runner 或 IPC 路径。 -## 逐步迁移 +### 6. 在适配风险附近验证 -1. **审查观测。** 确保每个观测 key 都是一个无需 GPU PhysX 查询即可表达的向量。 - 如果不是，就添加一个状态估计器，或把该查询移到冷路径。 -2. **移植 asset。** UniLab 以 MJCF 作为唯一真实来源（source of truth）。如果你 - 有 USD，先转换为 MJCF。 -3. **移植 env。** 继承 `unilab.base.np_env.NpEnv`。把 reward 计算移进 env 的 - `compute_reward()`。 -4. **移植 YAML。** 按照 {doc}`5-task_config_translation` 中的表格，把 Isaac Lab - 的 `EnvCfg` 字段映射到 UniLab 任务 owner YAML。 -5. **移植 reward。** 使用 {doc}`6-reward_porting` 中的食谱。 -6. **验证。** 训练一个小规模运行，把 reward 曲线与你的 Isaac 基线对比。 +测试 Hydra compose 与 typed materialization、term 顺序与数学、selector 失败、 +observation/action shape、局部 reset，以及至少一个真实已注册 backend 的 transition。行为 +应与固定来源 task 对比；完成语义迁移后再做性能 benchmark。 -## 你会失去什么（以及如何弥补） +## 仓库证据 -- **Isaac Sim 渲染器。** 使用 Motrix 的无头视频导出，或构建一个 viser 场景 - （`unilab.visualization.viser_scene`）。 -- **每个 env 的 tensor obs。** UniLab 给你的是 dict-of-arrays；如果你需要 tensor， - 用你自己的 `obs_to_tensor` 包一层。 -- **内置的 GPU 侧 DR。** UniLab 的 DR 是 CPU 侧、按进程进行的。对大多数任务来说 - 这已经足够；对于极端并行，使用更多的 worker 进程。 +`tests/fixtures/isaac_lab_cartpole/` 迁移了 Isaac Lab commit +`b0542fe2d45bf91c4e1d9ef6952b9c709c80b4e8` 的 Manager-Based Cartpole task。它保留 +全部 12 个来源 term 的名称和顺序，同时把 Torch 适配为 NumPy、嵌套 config object 适配 +为 Hydra YAML，并用 fixture-local MJCF 实现 scene/action/reset 边界。这只是 test-only +证据，不是 production task 或 Isaac Lab 全量支持声明。 ## 另请参阅 -- {doc}`2-from_legged_gym` -- {doc}`3-from_rsl_rl` -- {doc}`5-task_config_translation` -- {doc}`6-reward_porting` +- {doc}`Manager-Based API <../../4-developer_guide/1-architecture/6-manager_based_api>` +- {doc}`Env contract <../../4-developer_guide/2-contracts/1-env_contract>` +- {doc}`ADR-0006 ` From 7eedf7cedcbf5c4c84a7322845d705c61414120c Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Wed, 19 Aug 2026 11:35:56 +0800 Subject: [PATCH 087/169] test: add pinned mjlab migration fixture [skip ci] (#1220) --- tests/fixtures/cartpole_manager_adapters.py | 160 ++++++++++++++++++ tests/fixtures/isaac_lab_cartpole/task.py | 149 ++-------------- tests/fixtures/mjlab_cartpole/README.md | 18 ++ tests/fixtures/mjlab_cartpole/__init__.py | 3 + tests/fixtures/mjlab_cartpole/cartpole.xml | 31 ++++ .../fixtures/mjlab_cartpole/conf/config.yaml | 109 ++++++++++++ tests/fixtures/mjlab_cartpole/task.py | 83 +++++++++ .../managers/test_mjlab_migration_fixture.py | 150 ++++++++++++++++ 8 files changed, 564 insertions(+), 139 deletions(-) create mode 100644 tests/fixtures/cartpole_manager_adapters.py create mode 100644 tests/fixtures/mjlab_cartpole/README.md create mode 100644 tests/fixtures/mjlab_cartpole/__init__.py create mode 100644 tests/fixtures/mjlab_cartpole/cartpole.xml create mode 100644 tests/fixtures/mjlab_cartpole/conf/config.yaml create mode 100644 tests/fixtures/mjlab_cartpole/task.py create mode 100644 tests/managers/test_mjlab_migration_fixture.py diff --git a/tests/fixtures/cartpole_manager_adapters.py b/tests/fixtures/cartpole_manager_adapters.py new file mode 100644 index 000000000..9e49e0b67 --- /dev/null +++ b/tests/fixtures/cartpole_manager_adapters.py @@ -0,0 +1,160 @@ +# Derived from Isaac Lab b0542fe2d45bf91c4e1d9ef6952b9c709c80b4e8, +# source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/cartpole. +# Copyright (c) 2022-2026, The Isaac Lab Project Developers. +# Modified by UniLab as test-only NumPy/entity adapters; BSD-3-Clause. +"""Test-only action and reset adapters shared by pinned Cartpole fixtures.""" + +from __future__ import annotations + +from dataclasses import dataclass +from numbers import Real +from typing import TYPE_CHECKING, cast + +import numpy as np + +from unilab.managers import ActionTerm, ActionTermCfg +from unilab.managers.scene_entity_config import SceneEntityCfg + +if TYPE_CHECKING: + from unilab.base.entity import Entity + from unilab.managers._types import ManagerBasedRlEnv + + +def finite_real(value: Real, *, label: str) -> float: + if isinstance(value, (bool, np.bool_)) or not isinstance(value, Real): + raise TypeError(f"{label} must be a real number") + result = float(value) + if not np.isfinite(result): + raise ValueError(f"{label} must be finite") + return result + + +def numeric_range(value: tuple[float, float] | list[float], *, label: str) -> tuple[float, float]: + if isinstance(value, (str, bytes)) or not isinstance(value, (tuple, list)): + raise TypeError(f"{label} must be a two-value range") + if len(value) != 2: + raise ValueError(f"{label} must contain two values") + lower = finite_real(value[0], label=f"{label}[0]") + upper = finite_real(value[1], label=f"{label}[1]") + if lower > upper: + raise ValueError(f"{label} lower bound {lower} exceeds upper bound {upper}") + return lower, upper + + +@dataclass(kw_only=True) +class JointEffortActionCfg(ActionTermCfg): + """Fixture-only adapter for the community ``JointEffortActionCfg`` surface.""" + + actuator_names: tuple[str, ...] | list[str] + scale: float = 1.0 + + def build(self, env: ManagerBasedRlEnv) -> JointEffortAction: + return JointEffortAction(self, env) + + +class JointEffortAction(ActionTerm): + """Scale policy actions and write entity-local actuator efforts.""" + + cfg: JointEffortActionCfg + + def __init__(self, cfg: JointEffortActionCfg, env: ManagerBasedRlEnv): + super().__init__(cfg, env) + if cfg.clip is not None: + raise NotImplementedError("Cartpole fixture JointEffortAction does not support clip") + if isinstance(cfg.actuator_names, (str, bytes)) or not isinstance( + cfg.actuator_names, (tuple, list) + ): + raise TypeError("JointEffortActionCfg actuator_names must be an ordered sequence") + actuator_ids, actuator_names = self._entity.find_actuators( + cfg.actuator_names, + preserve_order=True, + ) + if not actuator_ids: + raise ValueError( + "JointEffortActionCfg actuator_names resolved no actuators; " + f"patterns={list(cfg.actuator_names)}" + ) + self._actuator_ids = np.asarray(actuator_ids, dtype=np.intp) + self._actuator_ids.setflags(write=False) + self._actuator_names = tuple(actuator_names) + self._scale = finite_real(cfg.scale, label="JointEffortActionCfg scale") + self._raw_actions = np.zeros((self.num_envs, len(actuator_ids)), dtype=np.float32) + self._processed_actions = np.zeros_like(self._raw_actions) + + @property + def action_dim(self) -> int: + return self._raw_actions.shape[1] + + @property + def raw_action(self) -> np.ndarray: + return self._raw_actions + + def process_actions(self, actions: np.ndarray) -> None: + if not isinstance(actions, np.ndarray): + raise TypeError( + "Cartpole fixture JointEffortAction expected np.ndarray, " + f"received {type(actions).__name__}" + ) + if actions.shape != self._raw_actions.shape: + raise ValueError( + "Cartpole fixture JointEffortAction expected shape " + f"{self._raw_actions.shape}, received {actions.shape}" + ) + if not np.isfinite(actions).all(): + raise ValueError("Cartpole fixture JointEffortAction received NaN or Inf") + np.copyto(self._raw_actions, actions) + np.multiply(actions, self._scale, out=self._processed_actions) + + def apply_actions(self) -> None: + self._entity.data.write_ctrl( + self._processed_actions, + actuator_ids=self._actuator_ids, + ) + + def reset(self, env_ids: np.ndarray | slice | None = None) -> None: + ids = slice(None) if env_ids is None else env_ids + self._raw_actions[ids] = 0.0 + self._processed_actions[ids] = 0.0 + + +def reset_joints_by_offset( + env: ManagerBasedRlEnv, + env_ids: np.ndarray | None, + position_range: tuple[float, float] | list[float], + velocity_range: tuple[float, float] | list[float], + asset_cfg: SceneEntityCfg, +) -> None: + """Write uniformly offset joint defaults through the reset transaction.""" + if env_ids is None: + raise ValueError("reset_joints_by_offset requires concrete environment IDs") + position_lower, position_upper = numeric_range(position_range, label="position_range") + velocity_lower, velocity_upper = numeric_range(velocity_range, label="velocity_range") + asset = cast("Entity", env.scene[asset_cfg.name]) + joint_ids = asset_cfg.joint_ids + default_position = asset.data.default_joint_pos[env_ids][:, joint_ids] + default_velocity = asset.data.default_joint_vel[env_ids][:, joint_ids] + position = default_position + env.rng.uniform( + position_lower, + position_upper, + default_position.shape, + ) + velocity = default_velocity + env.rng.uniform( + velocity_lower, + velocity_upper, + default_velocity.shape, + ) + asset.write_joint_state_to_sim( + position.astype(default_position.dtype, copy=False), + velocity.astype(default_velocity.dtype, copy=False), + joint_ids=joint_ids, + env_ids=env_ids, + ) + + +__all__ = [ + "JointEffortAction", + "JointEffortActionCfg", + "finite_real", + "numeric_range", + "reset_joints_by_offset", +] diff --git a/tests/fixtures/isaac_lab_cartpole/task.py b/tests/fixtures/isaac_lab_cartpole/task.py index 2e80e04c9..ffa912b43 100644 --- a/tests/fixtures/isaac_lab_cartpole/task.py +++ b/tests/fixtures/isaac_lab_cartpole/task.py @@ -2,20 +2,24 @@ # source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/cartpole. # Copyright (c) 2022-2026, The Isaac Lab Project Developers. # Modified by UniLab for NumPy and the fixture-local MJCF/entity adapter; BSD-3-Clause. -"""NumPy terms and adapters for the Isaac Lab Cartpole migration fixture.""" +"""NumPy terms for the Isaac Lab Cartpole migration fixture.""" from __future__ import annotations import math -from dataclasses import dataclass -from numbers import Real from typing import TYPE_CHECKING, cast import numpy as np +from tests.fixtures.cartpole_manager_adapters import ( + JointEffortAction, + JointEffortActionCfg, + finite_real, + numeric_range, + reset_joints_by_offset, +) from unilab.base import registry from unilab.envs import ManagerBasedRlEnvCfg, make_manager_based_rl_env -from unilab.managers import ActionTerm, ActionTermCfg from unilab.managers.scene_entity_config import SceneEntityCfg if TYPE_CHECKING: @@ -26,146 +30,13 @@ FIXTURE_ENV_NAME = "IsaacLabCartpoleFixture" -def _finite_real(value: Real, *, label: str) -> float: - if isinstance(value, (bool, np.bool_)) or not isinstance(value, Real): - raise TypeError(f"{label} must be a real number") - result = float(value) - if not np.isfinite(result): - raise ValueError(f"{label} must be finite") - return result - - -def _range(value: tuple[float, float] | list[float], *, label: str) -> tuple[float, float]: - if isinstance(value, (str, bytes)) or not isinstance(value, (tuple, list)): - raise TypeError(f"{label} must be a two-value range") - if len(value) != 2: - raise ValueError(f"{label} must contain two values") - lower = _finite_real(value[0], label=f"{label}[0]") - upper = _finite_real(value[1], label=f"{label}[1]") - if lower > upper: - raise ValueError(f"{label} lower bound {lower} exceeds upper bound {upper}") - return lower, upper - - -@dataclass(kw_only=True) -class JointEffortActionCfg(ActionTermCfg): - """Fixture-local adapter for Isaac Lab's ``JointEffortActionCfg``.""" - - actuator_names: tuple[str, ...] | list[str] - scale: float = 1.0 - - def build(self, env: ManagerBasedRlEnv) -> JointEffortAction: - return JointEffortAction(self, env) - - -class JointEffortAction(ActionTerm): - """Scale policy actions and write entity-local actuator efforts.""" - - cfg: JointEffortActionCfg - - def __init__(self, cfg: JointEffortActionCfg, env: ManagerBasedRlEnv): - super().__init__(cfg, env) - if cfg.clip is not None: - raise NotImplementedError( - "IsaacLabCartpoleFixture JointEffortAction does not support clip" - ) - if isinstance(cfg.actuator_names, (str, bytes)) or not isinstance( - cfg.actuator_names, (tuple, list) - ): - raise TypeError("JointEffortActionCfg actuator_names must be an ordered sequence") - actuator_ids, actuator_names = self._entity.find_actuators( - cfg.actuator_names, - preserve_order=True, - ) - if not actuator_ids: - raise ValueError( - "JointEffortActionCfg actuator_names resolved no actuators; " - f"patterns={list(cfg.actuator_names)}" - ) - self._actuator_ids = np.asarray(actuator_ids, dtype=np.intp) - self._actuator_ids.setflags(write=False) - self._actuator_names = tuple(actuator_names) - self._scale = _finite_real(cfg.scale, label="JointEffortActionCfg scale") - self._raw_actions = np.zeros((self.num_envs, len(actuator_ids)), dtype=np.float32) - self._processed_actions = np.zeros_like(self._raw_actions) - - @property - def action_dim(self) -> int: - return self._raw_actions.shape[1] - - @property - def raw_action(self) -> np.ndarray: - return self._raw_actions - - def process_actions(self, actions: np.ndarray) -> None: - if not isinstance(actions, np.ndarray): - raise TypeError( - "IsaacLabCartpoleFixture JointEffortAction expected np.ndarray, " - f"received {type(actions).__name__}" - ) - if actions.shape != self._raw_actions.shape: - raise ValueError( - "IsaacLabCartpoleFixture JointEffortAction expected shape " - f"{self._raw_actions.shape}, received {actions.shape}" - ) - if not np.isfinite(actions).all(): - raise ValueError("IsaacLabCartpoleFixture JointEffortAction received NaN or Inf") - np.copyto(self._raw_actions, actions) - np.multiply(actions, self._scale, out=self._processed_actions) - - def apply_actions(self) -> None: - self._entity.data.write_ctrl( - self._processed_actions, - actuator_ids=self._actuator_ids, - ) - - def reset(self, env_ids: np.ndarray | slice | None = None) -> None: - ids = slice(None) if env_ids is None else env_ids - self._raw_actions[ids] = 0.0 - self._processed_actions[ids] = 0.0 - - -def reset_joints_by_offset( - env: ManagerBasedRlEnv, - env_ids: np.ndarray | None, - position_range: tuple[float, float] | list[float], - velocity_range: tuple[float, float] | list[float], - asset_cfg: SceneEntityCfg, -) -> None: - """Port Isaac Lab's joint-offset reset through the entity reset transaction.""" - if env_ids is None: - raise ValueError("reset_joints_by_offset requires concrete environment IDs") - position_lower, position_upper = _range(position_range, label="position_range") - velocity_lower, velocity_upper = _range(velocity_range, label="velocity_range") - asset = cast("Entity", env.scene[asset_cfg.name]) - joint_ids = asset_cfg.joint_ids - default_position = asset.data.default_joint_pos[env_ids][:, joint_ids] - default_velocity = asset.data.default_joint_vel[env_ids][:, joint_ids] - position = default_position + env.rng.uniform( - position_lower, - position_upper, - default_position.shape, - ) - velocity = default_velocity + env.rng.uniform( - velocity_lower, - velocity_upper, - default_velocity.shape, - ) - asset.write_joint_state_to_sim( - position.astype(default_position.dtype, copy=False), - velocity.astype(default_velocity.dtype, copy=False), - joint_ids=joint_ids, - env_ids=env_ids, - ) - - def joint_pos_target_l2( env: ManagerBasedRlEnv, target: float, asset_cfg: SceneEntityCfg, ) -> np.ndarray: """Penalize wrapped joint-position deviation from a target value.""" - target_value = _finite_real(target, label="joint_pos_target_l2 target") + target_value = finite_real(target, label="joint_pos_target_l2 target") asset = cast("Entity", env.scene[asset_cfg.name]) joint_pos = asset.data.joint_pos[:, asset_cfg.joint_ids] wrapped = np.remainder(joint_pos + math.pi, 2.0 * math.pi) - math.pi @@ -187,7 +58,7 @@ def joint_pos_out_of_manual_limit( asset_cfg: SceneEntityCfg, ) -> np.ndarray: """Terminate when a selected joint leaves the configured manual bounds.""" - lower, upper = _range(bounds, label="joint_pos_out_of_manual_limit bounds") + lower, upper = numeric_range(bounds, label="joint_pos_out_of_manual_limit bounds") asset = cast("Entity", env.scene[asset_cfg.name]) joint_pos = asset.data.joint_pos[:, asset_cfg.joint_ids] return np.any((joint_pos < lower) | (joint_pos > upper), axis=1) diff --git a/tests/fixtures/mjlab_cartpole/README.md b/tests/fixtures/mjlab_cartpole/README.md new file mode 100644 index 000000000..55fe82ea6 --- /dev/null +++ b/tests/fixtures/mjlab_cartpole/README.md @@ -0,0 +1,18 @@ +# mjlab Cartpole migration fixture + +This test-only fixture is derived from mjlab v1.6.0 commit +`0fb8a681136be94ffc636a3dd423cabb97d91f10`, specifically +`src/mjlab/tasks/cartpole/cartpole_env_cfg.py` and `cartpole.xml`. The source is +Apache-2.0; provenance is retained in the derived files. + +| Surface | Status | Migration delta | +| --- | --- | --- | +| Manager dictionaries, term names/order, `func + params` | Compatible | Imports change from `mjlab` to `unilab`; the Balance task keeps all source entries and ordering. | +| Term math and buffers | Adapted | `torch.Tensor` and Torch ops become `np.ndarray` and NumPy ops. | +| Config container | Adapted | The Python config factory becomes one Hydra owner YAML, materialized as a plain `ManagerBasedRlEnvCfg`. | +| Scene/simulation | Adapted | mjlab scene/sim/viewer objects become task-owned MJCF plus `SceneCfg`/`EntityCfg`; the source contact-disable setting is embedded and visual materials are inlined for replicated scenes. | +| Joint effort and reset mutation | Adapted, fixture-local | Shared test-only adapters write through the entity control/reset contracts; they are not public built-ins. | +| mjlab runner, viewer, Torch/Warp runtime | Unsupported | No dependency or fallback is provided. | + +This fixture is evidence for #1042, not a production registration or a claim +that arbitrary mjlab tasks run unchanged. diff --git a/tests/fixtures/mjlab_cartpole/__init__.py b/tests/fixtures/mjlab_cartpole/__init__.py new file mode 100644 index 000000000..4a64a6eab --- /dev/null +++ b/tests/fixtures/mjlab_cartpole/__init__.py @@ -0,0 +1,3 @@ +"""Pinned mjlab Cartpole migration fixture.""" + +from tests.fixtures.mjlab_cartpole.task import FIXTURE_ENV_NAME as FIXTURE_ENV_NAME diff --git a/tests/fixtures/mjlab_cartpole/cartpole.xml b/tests/fixtures/mjlab_cartpole/cartpole.xml new file mode 100644 index 000000000..0164d63ac --- /dev/null +++ b/tests/fixtures/mjlab_cartpole/cartpole.xml @@ -0,0 +1,31 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/tests/fixtures/mjlab_cartpole/conf/config.yaml b/tests/fixtures/mjlab_cartpole/conf/config.yaml new file mode 100644 index 000000000..0facd98aa --- /dev/null +++ b/tests/fixtures/mjlab_cartpole/conf/config.yaml @@ -0,0 +1,109 @@ +# Derived from mujocolab/mjlab 0fb8a681, +# src/mjlab/tasks/cartpole/cartpole_env_cfg.py; Apache-2.0. +# Hydra is the fixture's only task-configuration owner. +training: + task_name: MjlabCartpoleBalanceFixture + sim_backend: mujoco + +env: + scene: + model_file: tests/fixtures/mjlab_cartpole/cartpole.xml + entities: + cartpole: + root_body_name: cart + joint_names: [slider, hinge_1] + actuator_names: [slide] + sim_dt: 0.01 + ctrl_dt: 0.05 + max_episode_seconds: 50.0 + seed: 13 + observations: + actor: + _target_: unilab.managers.ObservationGroupCfg + enable_corruption: true + terms: + cart_pos: &cart_pos + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_pos_rel + params: + asset_cfg: &cart_cfg + _target_: unilab.managers.SceneEntityCfg + name: cartpole + joint_names: [slider] + pole_angle: &pole_angle + _target_: unilab.managers.ObservationTermCfg + func: tests.fixtures.mjlab_cartpole.task.pole_angle_cos_sin + params: + asset_cfg: &hinge_cfg + _target_: unilab.managers.SceneEntityCfg + name: cartpole + joint_names: [hinge_1] + cart_vel: &cart_vel + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_vel_rel + params: + asset_cfg: *cart_cfg + pole_vel: &pole_vel + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_vel_rel + params: + asset_cfg: *hinge_cfg + critic: + _target_: unilab.managers.ObservationGroupCfg + terms: + cart_pos: *cart_pos + pole_angle: *pole_angle + cart_vel: *cart_vel + pole_vel: *pole_vel + actions: + effort: + _target_: tests.fixtures.cartpole_manager_adapters.JointEffortActionCfg + entity_name: cartpole + actuator_names: [slide] + scale: 1.0 + events: + reset_slider: + _target_: unilab.managers.EventTermCfg + func: tests.fixtures.cartpole_manager_adapters.reset_joints_by_offset + mode: reset + params: + position_range: [-0.1, 0.1] + velocity_range: [-0.01, 0.01] + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: cartpole + joint_names: [slider] + reset_hinge: + _target_: unilab.managers.EventTermCfg + func: tests.fixtures.cartpole_manager_adapters.reset_joints_by_offset + mode: reset + params: + position_range: [-0.034, 0.034] + velocity_range: [-0.01, 0.01] + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: cartpole + joint_names: [hinge_1] + terminations: + time_out: + _target_: unilab.managers.TerminationTermCfg + func: unilab.envs.mdp.time_out + time_out: true + policy_observation_group: actor + critic_observation_group: critic + scale_rewards_by_dt: true + +reward: + smooth_reward: + _target_: unilab.managers.RewardTermCfg + func: tests.fixtures.mjlab_cartpole.task.cartpole_smooth_reward + weight: 1.0 + params: + cart_cfg: + _target_: unilab.managers.SceneEntityCfg + name: cartpole + joint_names: [slider] + hinge_cfg: + _target_: unilab.managers.SceneEntityCfg + name: cartpole + joint_names: [hinge_1] diff --git a/tests/fixtures/mjlab_cartpole/task.py b/tests/fixtures/mjlab_cartpole/task.py new file mode 100644 index 000000000..47a79365d --- /dev/null +++ b/tests/fixtures/mjlab_cartpole/task.py @@ -0,0 +1,83 @@ +# Derived from mujocolab/mjlab v1.6.0 (0fb8a681), +# src/mjlab/tasks/cartpole/cartpole_env_cfg.py. +# Copyright 2025, The mjlab Developers. +# Modified by UniLab for NumPy and fixture-only Hydra/entity adapters; Apache-2.0. +"""NumPy terms for the pinned mjlab Cartpole Balance fixture.""" + +from __future__ import annotations + +import math +from typing import TYPE_CHECKING, cast + +import numpy as np + +from unilab.base import registry +from unilab.envs import ManagerBasedRlEnvCfg, make_manager_based_rl_env +from unilab.managers import SceneEntityCfg + +if TYPE_CHECKING: + from unilab.base.entity import Entity + from unilab.managers._types import ManagerBasedRlEnv + + +FIXTURE_ENV_NAME = "MjlabCartpoleBalanceFixture" +_GAUSSIAN_SCALE = math.sqrt(-2.0 * math.log(0.1)) +_QUADRATIC_SCALE = math.sqrt(1.0 - 0.1) + + +def pole_angle_cos_sin(env: ManagerBasedRlEnv, asset_cfg: SceneEntityCfg) -> np.ndarray: + """Return cosine and sine of the selected pole angle.""" + asset = cast("Entity", env.scene[asset_cfg.name]) + angle = asset.data.joint_pos[:, asset_cfg.joint_ids] + return np.concatenate((np.cos(angle), np.sin(angle)), axis=-1) + + +def _gaussian_tolerance(x: np.ndarray, margin: float) -> np.ndarray: + if margin == 0.0: + return (x == 0.0).astype(np.float32) + scaled = x / margin * _GAUSSIAN_SCALE + return np.exp(-0.5 * np.square(scaled)) + + +def _quadratic_tolerance(x: np.ndarray, margin: float) -> np.ndarray: + if margin == 0.0: + return (x == 0.0).astype(np.float32) + scaled = x / margin * _QUADRATIC_SCALE + return np.maximum(1.0 - np.square(scaled), 0.0) + + +def cartpole_smooth_reward( + env: ManagerBasedRlEnv, + cart_cfg: SceneEntityCfg, + hinge_cfg: SceneEntityCfg, +) -> np.ndarray: + """Port mjlab's dm_control-style smooth Cartpole reward to NumPy.""" + asset = cast("Entity", env.scene[cart_cfg.name]) + hinge_angle = asset.data.joint_pos[:, hinge_cfg.joint_ids].squeeze(-1) + upright = (np.cos(hinge_angle) + 1.0) / 2.0 + cart_pos = asset.data.joint_pos[:, cart_cfg.joint_ids].squeeze(-1) + centered = (1.0 + _gaussian_tolerance(cart_pos, margin=2.0)) / 2.0 + control = env.action_manager.action.squeeze(-1) + small_control = (4.0 + _quadratic_tolerance(control, margin=1.0)) / 5.0 + hinge_vel = asset.data.joint_vel[:, hinge_cfg.joint_ids].squeeze(-1) + small_velocity = (1.0 + _gaussian_tolerance(hinge_vel, margin=5.0)) / 2.0 + return upright * centered * small_control * small_velocity + + +def register_fixture() -> None: + """Register this fixture without adding it to production bootstrap.""" + if registry.contains(FIXTURE_ENV_NAME): + return + registry.register_env_config(FIXTURE_ENV_NAME, ManagerBasedRlEnvCfg) + registry.register_env(FIXTURE_ENV_NAME, make_manager_based_rl_env, sim_backend="mujoco") + + +register_fixture() + + +__all__ = [ + "FIXTURE_ENV_NAME", + "cartpole_smooth_reward", + "pole_angle_cos_sin", + "register_fixture", +] diff --git a/tests/managers/test_mjlab_migration_fixture.py b/tests/managers/test_mjlab_migration_fixture.py new file mode 100644 index 000000000..3b9a6e243 --- /dev/null +++ b/tests/managers/test_mjlab_migration_fixture.py @@ -0,0 +1,150 @@ +"""End-to-end evidence for the pinned mjlab task migration fixture.""" + +from __future__ import annotations + +from dataclasses import fields, is_dataclass +from pathlib import Path +from typing import Any + +import numpy as np +import pytest +from hydra import compose, initialize_config_dir +from hydra.core.global_hydra import GlobalHydra +from omegaconf import DictConfig, OmegaConf + +from tests.fixtures.mjlab_cartpole import FIXTURE_ENV_NAME +from unilab.base import registry +from unilab.base.config_materialization import apply_cfg_overrides +from unilab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg +from unilab.training import BackendAdapter + +ROOT_DIR = Path(__file__).parents[2] +FIXTURE_DIR = ROOT_DIR / "tests" / "fixtures" / "mjlab_cartpole" + + +def _materialize() -> tuple[ManagerBasedRlEnvCfg, dict[str, Any]]: + GlobalHydra.instance().clear() + with initialize_config_dir(config_dir=str(FIXTURE_DIR / "conf"), version_base="1.3"): + hydra_cfg: DictConfig = compose("config") + override = BackendAdapter(hydra_cfg, root_dir=ROOT_DIR).build_task_env_cfg_override() + env_cfg = registry.materialize_env_config(FIXTURE_ENV_NAME) + assert isinstance(env_cfg, ManagerBasedRlEnvCfg) + apply_cfg_overrides(env_cfg, override) + env_cfg.validate() + return env_cfg, override + + +def _assert_plain(value: Any) -> None: + assert not OmegaConf.is_config(value) + if is_dataclass(value) and not isinstance(value, type): + for field in fields(value): + _assert_plain(getattr(value, field.name)) + elif isinstance(value, dict): + for item in value.values(): + _assert_plain(item) + elif isinstance(value, (tuple, list)): + for item in value: + _assert_plain(item) + + +def test_mjlab_fixture_hydra_materializes_source_structure() -> None: + cfg, _ = _materialize() + + assert list(cfg.observations) == ["actor", "critic"] + expected_obs = ["cart_pos", "pole_angle", "cart_vel", "pole_vel"] + assert list(cfg.observations["actor"].terms) == expected_obs + assert list(cfg.observations["critic"].terms) == expected_obs + assert cfg.observations["actor"].enable_corruption is True + assert cfg.observations["critic"].enable_corruption is False + assert list(cfg.actions) == ["effort"] + assert list(cfg.events) == ["reset_slider", "reset_hinge"] + assert list(cfg.rewards) == ["smooth_reward"] + assert list(cfg.terminations) == ["time_out"] + assert cfg.policy_observation_group == "actor" + assert cfg.critic_observation_group == "critic" + assert cfg.sim_dt == pytest.approx(0.01) + assert cfg.ctrl_dt == pytest.approx(0.05) + assert cfg.max_episode_seconds == pytest.approx(50.0) + assert cfg.scale_rewards_by_dt is True + _assert_plain(cfg) + + +def test_mjlab_fixture_real_mujoco_reset_step_and_reward() -> None: + _, override = _materialize() + env = registry.make( + FIXTURE_ENV_NAME, + sim_backend="mujoco", + env_cfg_override=override, + num_envs=8, + ) + assert isinstance(env, ManagerBasedRlEnv) + try: + state = env.init_state() + assert env.obs_groups_spec == {"obs": 5, "critic": 5} + assert state.obs["obs"].shape == state.obs["critic"].shape == (8, 5) + assert env.action_space.shape == (1,) + + before = env.scene["cartpole"].data.joint_pos.copy() + ids = np.asarray([1, 6], dtype=np.int32) + reset_obs, _ = env.reset(env_ids=ids) + after = env.scene["cartpole"].data.joint_pos.copy() + np.testing.assert_array_equal(after[[0, 2, 3, 4, 5, 7]], before[[0, 2, 3, 4, 5, 7]]) + assert reset_obs["obs"].shape == reset_obs["critic"].shape == (2, 5) + assert np.all(np.abs(after[ids, 0]) <= 0.1) + assert np.all(np.abs(after[ids, 1]) <= 0.034) + assert np.all(np.abs(env.scene["cartpole"].data.joint_vel[ids]) <= 0.01) + + actions = np.full((8, 1), 0.25, dtype=np.float32) + state = env.step(actions) + entity = env.scene["cartpole"] + hinge = entity.data.joint_pos[:, 1] + cart = entity.data.joint_pos[:, 0] + hinge_vel = entity.data.joint_vel[:, 1] + gaussian_scale = np.sqrt(-2.0 * np.log(0.1)) + quadratic_scale = np.sqrt(0.9) + expected = (np.cos(hinge) + 1.0) / 2.0 + expected *= (1.0 + np.exp(-0.5 * np.square(cart / 2.0 * gaussian_scale))) / 2.0 + expected *= (4.0 + np.maximum(1.0 - np.square(0.25 * quadratic_scale), 0.0)) / 5.0 + expected *= (1.0 + np.exp(-0.5 * np.square(hinge_vel / 5.0 * gaussian_scale))) / 2.0 + np.testing.assert_allclose(state.reward, expected * 0.05, rtol=1e-5, atol=1e-6) + finally: + env.close() + + +def test_mjlab_fixture_missing_actuator_fails_on_cold_path() -> None: + _, override = _materialize() + override["actions"]["effort"]["actuator_names"] = ["missing_actuator"] + with pytest.raises(ValueError, match="regular expressions matched.*missing_actuator"): + registry.make( + FIXTURE_ENV_NAME, + sim_backend="mujoco", + env_cfg_override=override, + num_envs=2, + ) + + +def test_mjlab_fixture_is_pinned_test_only_numpy_code() -> None: + task_source = (FIXTURE_DIR / "task.py").read_text(encoding="utf-8") + helper_source = (ROOT_DIR / "tests/fixtures/cartpole_manager_adapters.py").read_text( + encoding="utf-8" + ) + executable = "\n".join( + line + for line in (task_source + helper_source).splitlines() + if not line.lstrip().startswith("#") + ) + for forbidden in ( + "import torch", + "from torch", + "import mjlab", + "from mjlab", + "unilab.algos", + "unilab.ipc", + "unilab.training", + ): + assert forbidden not in executable + assert "0fb8a681136be94ffc636a3dd423cabb97d91f10" in (FIXTURE_DIR / "README.md").read_text( + encoding="utf-8" + ) + assert "tests.fixtures" not in tuple(registry._DEFAULT_REGISTRY_PACKAGES) + assert set(registry._envs[FIXTURE_ENV_NAME].env_factory_dict) == {"mujoco"} From 7794ec85689566b7262c6af6a9cafce46f69f3c8 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Thu, 20 Aug 2026 02:26:51 +0800 Subject: [PATCH 088/169] chore: codify task migration matrix and compatibility boundary (#1224) * chore(tasks): codify manager migration matrix * refactor(tasks): add internal legacy factory seam * refactor(tasks): freeze custom compatibility boundary * test(tasks): isolate production migration audit [skip ci] --- src/unilab/tasks/compatibility.py | 132 ++++++++++++ .../tasks/locomotion/go2_arm/manip_loco.py | 16 +- .../manipulation/sharpa_inhand/grasp_gen.py | 14 +- .../manipulation/sharpa_inhand/rotation.py | 15 +- src/unilab/tasks/migration_matrix.py | 167 +++++++++++++++ tests/tasks/test_legacy_task_compatibility.py | 201 ++++++++++++++++++ tests/tasks/test_migration_matrix.py | 44 ++++ 7 files changed, 581 insertions(+), 8 deletions(-) create mode 100644 src/unilab/tasks/compatibility.py create mode 100644 src/unilab/tasks/migration_matrix.py create mode 100644 tests/tasks/test_legacy_task_compatibility.py create mode 100644 tests/tasks/test_migration_matrix.py diff --git a/src/unilab/tasks/compatibility.py b/src/unilab/tasks/compatibility.py new file mode 100644 index 000000000..a33106b20 --- /dev/null +++ b/src/unilab/tasks/compatibility.py @@ -0,0 +1,132 @@ +"""Internal, cold-path compatibility seam for legacy task factories. + +This module is intentionally task-owned and is not part of the base registry +contract. It only admits legacy factories that already use UniLab's +``EnvCfg -> NpEnv`` lifecycle; it does not provide a fallback runtime. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum +from typing import Generic, Protocol, TypeVar + +from unilab.base.base import ABEnv, EnvCfg +from unilab.base.np_env import NpEnv + + +class CompatibilityStatus(str, Enum): + """Documented outcome for one legacy compatibility boundary.""" + + ADAPTED = "Adapted" + UNSUPPORTED = "Unsupported" + + +@dataclass(frozen=True) +class LegacyTaskCompatibility: + """Immutable compatibility evidence attached to one task-family seam.""" + + task_family: str + status: CompatibilityStatus + reason: str + + def __post_init__(self) -> None: + if not self.task_family.strip(): + raise ValueError("legacy compatibility task_family must be non-empty") + if not self.reason.strip(): + raise ValueError("legacy compatibility reason must be non-empty") + + +TCfg_contra = TypeVar("TCfg_contra", bound=EnvCfg, contravariant=True) + + +class LegacyEnvFactory(Protocol[TCfg_contra]): + """Existing registry-shaped legacy factory admitted by this seam.""" + + def __call__( + self, + cfg: TCfg_contra, + *, + num_envs: int = 1, + backend_type: str = "mujoco", + ) -> ABEnv: ... + + +@dataclass(frozen=True) +class LegacyFactoryAdapter(Generic[TCfg_contra]): + """Validate one legacy factory at the existing env construction boundary.""" + + factory: LegacyEnvFactory[TCfg_contra] + compatibility: LegacyTaskCompatibility + + def __post_init__(self) -> None: + if self.compatibility.status is not CompatibilityStatus.ADAPTED: + raise ValueError("LegacyFactoryAdapter compatibility status must be Adapted") + + def __call__( + self, + cfg: TCfg_contra, + *, + num_envs: int = 1, + backend_type: str = "mujoco", + ) -> NpEnv: + family = self.compatibility.task_family + if not isinstance(cfg, EnvCfg): + raise TypeError( + f"Legacy task family '{family}' expected EnvCfg, received {type(cfg).__name__}" + ) + + env = self.factory(cfg, num_envs=num_envs, backend_type=backend_type) + if not isinstance(env, ABEnv): + raise TypeError( + f"Legacy task family '{family}' factory returned {type(env).__name__}, " + "expected ABEnv" + ) + if not isinstance(env, NpEnv): + raise TypeError( + f"Legacy task family '{family}' compatibility is Unsupported: " + f"{type(env).__name__} does not use the NpEnv lifecycle" + ) + return env + + +def adapt_legacy_factory( + factory: LegacyEnvFactory[TCfg_contra], + *, + task_family: str, + reason: str, +) -> LegacyFactoryAdapter[TCfg_contra]: + """Mark and wrap an existing ``EnvCfg -> NpEnv`` task factory. + + The wrapper runs only while the registry constructs an environment. It + forwards the registry's fixed arguments exactly once and rejects any + other config or runtime shape instead of probing or falling back. + """ + + if not callable(factory): + raise TypeError(f"legacy task family '{task_family}' factory must be callable") + compatibility = LegacyTaskCompatibility( + task_family=task_family, + status=CompatibilityStatus.ADAPTED, + reason=reason, + ) + return LegacyFactoryAdapter(factory=factory, compatibility=compatibility) + + +def unsupported_legacy_task(*, task_family: str, reason: str) -> LegacyTaskCompatibility: + """Record an explicit unsupported surface without creating a factory.""" + + return LegacyTaskCompatibility( + task_family=task_family, + status=CompatibilityStatus.UNSUPPORTED, + reason=reason, + ) + + +__all__ = [ + "CompatibilityStatus", + "LegacyFactoryAdapter", + "LegacyTaskCompatibility", + "adapt_legacy_factory", + "unsupported_legacy_task", +] diff --git a/src/unilab/tasks/locomotion/go2_arm/manip_loco.py b/src/unilab/tasks/locomotion/go2_arm/manip_loco.py index ec961600d..1f3c35a24 100644 --- a/src/unilab/tasks/locomotion/go2_arm/manip_loco.py +++ b/src/unilab/tasks/locomotion/go2_arm/manip_loco.py @@ -12,6 +12,7 @@ from unilab.base.scene import SceneCfg from unilab.dr.types import ResetPlan from unilab.dtype_config import get_global_dtype +from unilab.tasks.compatibility import adapt_legacy_factory from unilab.tasks.locomotion.common import rewards from unilab.tasks.locomotion.common.commands import Commands from unilab.tasks.locomotion.common.domain_rand import DomainRandConfig @@ -264,9 +265,6 @@ def _compute_reset_obs( return env._update_history(actor_raw, env_ids=env_ids, critic_raw_obs=critic_raw) # type: ignore[no-any-return] -@registry.env("Go2ArmManipLoco", sim_backend="motrix") -@registry.env("Go2ArmManipLoco", sim_backend="drake") -@registry.env("Go2ArmManipLoco", sim_backend="mujoco") class Go2ArmManipLocoEnv(Go2ArmBaseEnv): _cfg: Go2ArmManipLocoCfg # pyright: ignore[reportIncompatibleVariableOverride] @@ -1027,3 +1025,15 @@ def _reward_arm_collision(self, _ctx: RewardContext) -> np.ndarray: for name in self._ARM_TOUCH_SENSORS: total += self._backend.get_sensor_data(name)[:, 0] return total + + +_GO2_ARM_COMPAT_FACTORY = adapt_legacy_factory( + Go2ArmManipLocoEnv, + task_family="Go2ArmManipLoco", + reason=( + "custom IK/Jacobian, end-effector goals, and observation history remain " + "task-owned until formal Manager-Based terms exist" + ), +) +for _backend_type in ("mujoco", "motrix", "drake"): + registry.register_env("Go2ArmManipLoco", _GO2_ARM_COMPAT_FACTORY, _backend_type) diff --git a/src/unilab/tasks/manipulation/sharpa_inhand/grasp_gen.py b/src/unilab/tasks/manipulation/sharpa_inhand/grasp_gen.py index 4a6583451..77c602f77 100644 --- a/src/unilab/tasks/manipulation/sharpa_inhand/grasp_gen.py +++ b/src/unilab/tasks/manipulation/sharpa_inhand/grasp_gen.py @@ -12,6 +12,7 @@ from unilab.base.run_control import RunComplete from unilab.dr import ResetPlan from unilab.dr.dr_utils import build_common_reset_randomization +from unilab.tasks.compatibility import adapt_legacy_factory from unilab.tasks.manipulation.sharpa_inhand.base import ( SOURCE_DEFAULT_HAND_JOINT_POS_DEG, SharpaDomainRandConfig, @@ -138,8 +139,6 @@ def build_reset_plan(self, env: Any, env_ids: np.ndarray) -> ResetPlan: ) -@registry.env("SharpaInhandRotationGrasp", sim_backend="mujoco") -@registry.env("SharpaInhandRotationGrasp", sim_backend="motrix") class SharpaInhandRotationGraspEnv(SharpaInhandRotationEnv): _cfg: SharpaInhandRotationGraspCfg # pyright: ignore[reportIncompatibleVariableOverride] _MATERIALIZE_ROTATION_GRASP_CACHE = False @@ -390,3 +389,14 @@ def update_state(self, state: NpEnvState) -> NpEnvState: SharpaWaveGraspCfg = SharpaInhandGraspEnvCfg + +_SHARPA_GRASP_COMPAT_FACTORY = adapt_legacy_factory( + SharpaInhandRotationGraspEnv, + task_family="Sharpa", + reason=( + "tactile grasp validation and cache-collection completion remain task-owned " + "until formal Manager-Based capabilities exist" + ), +) +for _backend_type in ("mujoco", "motrix"): + registry.register_env("SharpaInhandRotationGrasp", _SHARPA_GRASP_COMPAT_FACTORY, _backend_type) diff --git a/src/unilab/tasks/manipulation/sharpa_inhand/rotation.py b/src/unilab/tasks/manipulation/sharpa_inhand/rotation.py index 8a6c5ca37..9ff635f40 100644 --- a/src/unilab/tasks/manipulation/sharpa_inhand/rotation.py +++ b/src/unilab/tasks/manipulation/sharpa_inhand/rotation.py @@ -30,6 +30,7 @@ ResetRandomizationPayload, ) from unilab.dtype_config import get_global_dtype +from unilab.tasks.compatibility import adapt_legacy_factory from unilab.tasks.manipulation.sharpa_inhand.base import ( SharpaInhandBaseCfg, SharpaInhandBaseEnv, @@ -440,9 +441,6 @@ def build_interval_randomization_plan( ) -@registry.env("SharpaInhandRotation", sim_backend="drake") -@registry.env("SharpaInhandRotation", sim_backend="mujoco") -@registry.env("SharpaInhandRotation", sim_backend="motrix") class SharpaInhandRotationEnv(SharpaInhandBaseEnv): _cfg: SharpaInhandRotationCfg # pyright: ignore[reportIncompatibleVariableOverride] _reward_cfg: RewardConfig @@ -1536,3 +1534,14 @@ def update_state(self, state: NpEnvState) -> NpEnvState: SharpaWaveRewardConfig = RewardConfig SharpaWaveRotationCfg = SharpaInhandRotationCfg + +_SHARPA_ROTATION_COMPAT_FACTORY = adapt_legacy_factory( + SharpaInhandRotationEnv, + task_family="Sharpa", + reason=( + "tactile/contact latency, object variants, and grasp-cache state remain " + "task-owned until formal Manager-Based capabilities exist" + ), +) +for _backend_type in ("mujoco", "motrix", "drake"): + registry.register_env("SharpaInhandRotation", _SHARPA_ROTATION_COMPAT_FACTORY, _backend_type) diff --git a/src/unilab/tasks/migration_matrix.py b/src/unilab/tasks/migration_matrix.py new file mode 100644 index 000000000..db061874c --- /dev/null +++ b/src/unilab/tasks/migration_matrix.py @@ -0,0 +1,167 @@ +"""Production task migration status and closeout ownership. + +The matrix is deliberately small and explicit. It is an audit boundary for +the grouped #1042 migration work; it does not provide a second task runtime or +translate task configuration. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Literal + +MigrationStatus = Literal["Compatible", "Adapted"] +MigrationTarget = Literal["complete", "mba", "compatibility"] + + +@dataclass(frozen=True) +class TaskMigrationRecord: + task_name: str + family: str + status: MigrationStatus + target: MigrationTarget + rationale: str + next_step: str + + +_MBA_TASKS = frozenset( + { + "A2JoystickFlat", + "AllegroInhandRotation", + "AllegroInhandRotationGrasp", + "Go1JoystickFlat", + "Go2FootStand", + "Go2JoystickFlat", + "Go2WJoystickFlat", + "StewartBalance", + } +) + +_ROUGH_TASKS = frozenset( + { + "Go1JoystickRough", + "Go2JoystickRough", + "Go2WJoystickRough", + } +) + +_G1_LOCOMOTION_TASKS = frozenset( + { + "G1WalkFlat", + "G1WalkRough", + "G1Walk23DofFlat", + "G1Walk23DofRough", + } +) + +_CUSTOM_COMPAT_TASKS = frozenset( + { + "Go2ArmManipLoco", + "SharpaInhandRotation", + "SharpaInhandRotationGrasp", + } +) + +_MOTION_TASKS = frozenset( + { + "G1BoxTracking", + "G1BoxTracking23Dof", + "G1ClimbTracking", + "G1ClimbTracking23Dof", + "G1FlipTracking", + "G1FlipTracking23Dof", + "G1FlipTrackingSAC", + "G1FlipTrackingSAC23Dof", + "G1MotionTracking", + "G1MotionTracking23Dof", + "G1MotionTracking23DofDeploy", + "G1MotionTrackingDeploy", + "G1MotionTrackingSAC", + "G1MotionTrackingSAC23Dof", + "G1WallFlipTracking", + "G1WallFlipTracking23Dof", + "G1WallFlipTrackingSAC", + "G1WallFlipTrackingSAC23Dof", + "G1WBTObs", + "G1WBTObs23Dof", + "X2WallFlipTracking", + } +) + +PRODUCTION_TASK_NAMES = frozenset( + _MBA_TASKS | _ROUGH_TASKS | _G1_LOCOMOTION_TASKS | _CUSTOM_COMPAT_TASKS | _MOTION_TASKS +) + + +def migration_record(task_name: str) -> TaskMigrationRecord: + """Return the closeout status for one registered production task. + + Unknown names fail closed so adding a production registration requires an + explicit migration decision and cannot silently escape the audit. + """ + + if task_name in _MBA_TASKS: + return TaskMigrationRecord( + task_name, + "manager_based", + "Compatible", + "complete", + "Hydra owner YAML materializes the canonical NumPy Manager-Based runtime.", + "Keep the manager contract and regression evidence current.", + ) + if task_name in _ROUGH_TASKS: + return TaskMigrationRecord( + task_name, + "quadruped_rough", + "Adapted", + "mba", + "Terrain and height-scan terms depend on the pending raycaster capability boundary.", + "Migrate as one rough-family PR; use the compatibility seam only if a new public capability is required.", + ) + if task_name in _G1_LOCOMOTION_TASKS: + return TaskMigrationRecord( + task_name, + "g1_locomotion", + "Adapted", + "mba", + "The locomotion equations are reusable, but the 29/23-DoF sensor and gait surface is not yet manager-owned.", + "Migrate flat and rough variants together and delete the legacy owner.", + ) + if task_name in _CUSTOM_COMPAT_TASKS: + family = "go2_arm" if task_name == "Go2ArmManipLoco" else "sharpa" + return TaskMigrationRecord( + task_name, + family, + "Adapted", + "compatibility", + "Custom IK/history or tactile/contact/cache behavior is retained behind one frozen adapter.", + "Keep Hydra/Registry ownership single; migrate only when the formal capability exists.", + ) + if task_name in _MOTION_TASKS: + return TaskMigrationRecord( + task_name, + "motion_tracking", + "Adapted", + "mba", + "Stateful motion loading and profile-specific tracking terms need a grouped manager port.", + "Migrate the shared engine and all profiles together; stop on new backend contracts.", + ) + raise KeyError(f"Task '{task_name}' has no #1042 migration-matrix entry") + + +def migration_records( + task_names: list[str] | tuple[str, ...] | set[str], +) -> tuple[TaskMigrationRecord, ...]: + """Return records in deterministic task-name order.""" + + return tuple(migration_record(name) for name in sorted(task_names)) + + +__all__ = [ + "MigrationStatus", + "MigrationTarget", + "PRODUCTION_TASK_NAMES", + "TaskMigrationRecord", + "migration_record", + "migration_records", +] diff --git a/tests/tasks/test_legacy_task_compatibility.py b/tests/tasks/test_legacy_task_compatibility.py new file mode 100644 index 000000000..eebabe275 --- /dev/null +++ b/tests/tasks/test_legacy_task_compatibility.py @@ -0,0 +1,201 @@ +"""Focused contract tests for the internal legacy-task compatibility seam.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import gymnasium as gym +import numpy as np +import pytest + +from unilab.base import registry +from unilab.base.base import ABEnv, EnvCfg +from unilab.base.np_env import NpEnv, NpEnvState +from unilab.tasks.compatibility import ( + CompatibilityStatus, + LegacyFactoryAdapter, + adapt_legacy_factory, + unsupported_legacy_task, +) + + +@dataclass +class _Cfg(EnvCfg): + pass + + +class _PlainABEnv(ABEnv): + @property + def num_envs(self) -> int: + return 1 + + @property + def cfg(self) -> EnvCfg: + return _Cfg() + + @property + def observation_space(self) -> gym.Space: + return gym.spaces.Box(-np.inf, np.inf, shape=(1,), dtype=np.float32) + + @property + def action_space(self) -> gym.Space: + return gym.spaces.Box(-1.0, 1.0, shape=(1,), dtype=np.float32) + + @property + def obs_groups_spec(self) -> dict[str, int]: + return {"obs": 1} + + @property + def state(self) -> None: + return None + + def init_state(self) -> None: + return None + + def step(self, actions: np.ndarray) -> None: + return None + + def close(self) -> None: + return None + + +class _NpEnv(NpEnv): + @property + def action_space(self) -> gym.Space: + return gym.spaces.Box(-1.0, 1.0, shape=(1,), dtype=np.float32) + + @property + def obs_groups_spec(self) -> dict[str, int]: + return {"obs": 1} + + def apply_action(self, actions: np.ndarray, state: NpEnvState) -> np.ndarray: + return actions + + def update_state(self, state: NpEnvState) -> NpEnvState: + return state + + +def _uninitialized_np_env() -> _NpEnv: + return object.__new__(_NpEnv) + + +def test_adapter_records_fixed_adapted_metadata_and_forwards_registry_arguments() -> None: + received: list[tuple[EnvCfg, int, str]] = [] + expected = _uninitialized_np_env() + + def factory(cfg: EnvCfg, *, num_envs: int, backend_type: str) -> ABEnv: + received.append((cfg, num_envs, backend_type)) + return expected + + adapter = adapt_legacy_factory( + factory, + task_family="Go2ArmManipLoco", + reason="existing task owner already constructs an NpEnv", + ) + cfg = _Cfg() + + assert adapter(cfg, num_envs=4, backend_type="motrix") is expected + assert received == [(cfg, 4, "motrix")] + assert adapter.compatibility.task_family == "Go2ArmManipLoco" + assert adapter.compatibility.status is CompatibilityStatus.ADAPTED + assert adapter.compatibility.reason == "existing task owner already constructs an NpEnv" + + +def test_adapter_rejects_non_env_cfg_before_calling_factory() -> None: + called = False + + def factory(cfg: EnvCfg, *, num_envs: int, backend_type: str) -> ABEnv: + nonlocal called + called = True + return _uninitialized_np_env() + + adapter = adapt_legacy_factory(factory, task_family="Sharpa", reason="migration seam") + + with pytest.raises(TypeError, match=r"Sharpa.*expected EnvCfg.*dict"): + adapter({}) # type: ignore[arg-type] + + assert called is False + + +@pytest.mark.parametrize( + ("result", "match"), + ( + (object(), r"Go2ArmManipLoco.*object.*expected ABEnv"), + (_PlainABEnv(), r"Go2ArmManipLoco.*Unsupported.*_PlainABEnv.*NpEnv"), + ), +) +def test_adapter_rejects_factories_outside_the_np_env_lifecycle( + result: object, + match: str, +) -> None: + def factory(cfg: EnvCfg, *, num_envs: int, backend_type: str) -> object: + return result + + adapter = adapt_legacy_factory( + factory, # type: ignore[arg-type] + task_family="Go2ArmManipLoco", + reason="migration seam", + ) + + with pytest.raises(TypeError, match=match): + adapter(_Cfg()) + + +def test_factory_exception_propagates_without_fallback() -> None: + failure = RuntimeError("owner factory failed") + + def factory(cfg: EnvCfg, *, num_envs: int, backend_type: str) -> ABEnv: + raise failure + + adapter = adapt_legacy_factory(factory, task_family="Sharpa", reason="migration seam") + + with pytest.raises(RuntimeError) as exc_info: + adapter(_Cfg()) + + assert exc_info.value is failure + + +def test_unsupported_metadata_is_explicit_and_does_not_create_a_factory() -> None: + compatibility = unsupported_legacy_task( + task_family="Sharpa foreign lifecycle", + reason="only the existing NpEnv lifecycle is admitted", + ) + + assert compatibility.status is CompatibilityStatus.UNSUPPORTED + assert compatibility.reason == "only the existing NpEnv lifecycle is admitted" + + with pytest.raises(ValueError, match="status must be Adapted"): + LegacyFactoryAdapter(lambda cfg, **kwargs: _uninitialized_np_env(), compatibility) + + +@pytest.mark.parametrize(("task_family", "reason"), (("", "reason"), ("Sharpa", ""))) +def test_compatibility_metadata_requires_stable_family_and_reason( + task_family: str, + reason: str, +) -> None: + with pytest.raises(ValueError, match="must be non-empty"): + unsupported_legacy_task(task_family=task_family, reason=reason) + + +@pytest.mark.parametrize( + ("task_name", "family", "backends"), + ( + ("Go2ArmManipLoco", "Go2ArmManipLoco", {"mujoco", "motrix", "drake"}), + ("SharpaInhandRotation", "Sharpa", {"mujoco", "motrix", "drake"}), + ("SharpaInhandRotationGrasp", "Sharpa", {"mujoco", "motrix"}), + ), +) +def test_approved_production_families_are_registered_through_the_frozen_seam( + task_name: str, + family: str, + backends: set[str], +) -> None: + registry.ensure_registries() + factories = registry._envs[task_name].env_factory_dict + + assert set(factories) == backends + assert all(isinstance(factory, LegacyFactoryAdapter) for factory in factories.values()) + assert {factory.compatibility.task_family for factory in factories.values()} == {family} + assert {factory.compatibility.status for factory in factories.values()} == { + CompatibilityStatus.ADAPTED + } diff --git a/tests/tasks/test_migration_matrix.py b/tests/tasks/test_migration_matrix.py new file mode 100644 index 000000000..d58dc1b7d --- /dev/null +++ b/tests/tasks/test_migration_matrix.py @@ -0,0 +1,44 @@ +from __future__ import annotations + +import pytest + +from unilab.base import registry +from unilab.tasks.migration_matrix import ( + PRODUCTION_TASK_NAMES, + migration_record, + migration_records, +) + + +def test_registered_tasks_have_explicit_migration_records() -> None: + registry.ensure_registries() + registered = registry.list_registered_envs() + records = migration_records(set(PRODUCTION_TASK_NAMES)) + + assert PRODUCTION_TASK_NAMES <= registered.keys() + assert {record.task_name for record in records} == set(PRODUCTION_TASK_NAMES) + assert len(records) == 39 + assert sum(record.status == "Compatible" for record in records) == 8 + assert sum(record.target == "compatibility" for record in records) == 3 + + +@pytest.mark.parametrize( + ("task_name", "family", "target"), + [ + ("Go2ArmManipLoco", "go2_arm", "compatibility"), + ("SharpaInhandRotation", "sharpa", "compatibility"), + ("G1MotionTracking", "motion_tracking", "mba"), + ("G1WalkRough", "g1_locomotion", "mba"), + ("Go2JoystickRough", "quadruped_rough", "mba"), + ], +) +def test_matrix_records_high_risk_families(task_name: str, family: str, target: str) -> None: + record = migration_record(task_name) + assert record.family == family + assert record.target == target + assert record.status == "Adapted" + + +def test_unknown_task_fails_closed() -> None: + with pytest.raises(KeyError, match="no #1042 migration-matrix entry"): + migration_record("NewTaskWithoutDecision") From 20162bb10f26d177cd82899bed5ac21c03fa122d Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Thu, 20 Aug 2026 04:55:47 +0800 Subject: [PATCH 089/169] feat: migrate quadruped rough family to Manager-Based runtime Local validation: make test-all passed on head 6892d05b (2045 passed, 31 skipped, 272 deselected, 1 xfailed); Ruff, Mypy, Pyright, coverage, and benchmark smoke passed. No remote checks are configured for the dev base branch. --- conf/ppo/task/go1_joystick_rough/motrix.yaml | 115 +-- conf/ppo/task/go1_joystick_rough/mujoco.yaml | 131 +-- conf/ppo/task/go2_joystick_rough/motrix.yaml | 114 +-- conf/ppo/task/go2_joystick_rough/mujoco.yaml | 131 +-- conf/ppo/task/go2w_joystick_rough/motrix.yaml | 98 +-- conf/ppo/task/go2w_joystick_rough/mujoco.yaml | 119 ++- .../task/quadruped_joystick_rough/base.yaml | 166 ++++ .../task/quadruped_joystick_rough/go2w.yaml | 150 ++++ .../quadruped_joystick_rough/quadruped.yaml | 109 +++ .../benchmark/benchmark_drake_performance.py | 25 +- scripts/benchmark/core/task_names.py | 15 +- scripts/benchmark/env/benchmark_env_step.py | 58 +- src/unilab/base/backend/mujoco/backend.py | 38 +- .../locomotion/common/rough_manager_terms.py | 572 +++++++++++++ src/unilab/tasks/locomotion/go1/__init__.py | 24 +- src/unilab/tasks/locomotion/go1/base.py | 40 - src/unilab/tasks/locomotion/go1/joystick.py | 262 ------ src/unilab/tasks/locomotion/go1/rough.py | 786 ----------------- src/unilab/tasks/locomotion/go2/__init__.py | 26 +- src/unilab/tasks/locomotion/go2/base.py | 54 -- src/unilab/tasks/locomotion/go2/joystick.py | 329 -------- src/unilab/tasks/locomotion/go2/rough.py | 793 ------------------ src/unilab/tasks/locomotion/go2w/__init__.py | 24 +- src/unilab/tasks/locomotion/go2w/base.py | 95 +-- src/unilab/tasks/locomotion/go2w/joystick.py | 651 -------------- src/unilab/tasks/locomotion/go2w/rough.py | 380 --------- src/unilab/tasks/migration_matrix.py | 8 +- tests/base/test_backend_conformance.py | 4 +- .../test_env_step_config_contract.py | 36 +- tests/config/test_config_system.py | 46 - tests/config/test_locomotion_params.py | 216 ----- .../go1/test_manager_based_flat_cfg.py | 10 +- .../locomotion/go2w/test_go2w_height_scan.py | 77 -- .../go2w/test_go2w_manager_based_flat_cfg.py | 10 +- .../go2w/test_go2w_motor_control.py | 299 ------- .../test_go2_joystick_stand_still.py | 29 - .../locomotion/test_go2_rough_height_scan.py | 77 -- .../envs/locomotion/test_go2_terrain_spawn.py | 529 +++++------- tests/envs/test_env_configs.py | 2 - tests/envs/test_go1_obs_noise.py | 77 -- tests/envs/test_go2_obs_noise.py | 74 -- tests/tasks/test_migration_matrix.py | 20 +- 42 files changed, 1500 insertions(+), 5319 deletions(-) create mode 100644 conf/ppo/task/quadruped_joystick_rough/base.yaml create mode 100644 conf/ppo/task/quadruped_joystick_rough/go2w.yaml create mode 100644 conf/ppo/task/quadruped_joystick_rough/quadruped.yaml create mode 100644 src/unilab/tasks/locomotion/common/rough_manager_terms.py delete mode 100644 src/unilab/tasks/locomotion/go1/base.py delete mode 100644 src/unilab/tasks/locomotion/go1/joystick.py delete mode 100644 src/unilab/tasks/locomotion/go1/rough.py delete mode 100644 src/unilab/tasks/locomotion/go2/base.py delete mode 100644 src/unilab/tasks/locomotion/go2/joystick.py delete mode 100644 src/unilab/tasks/locomotion/go2/rough.py delete mode 100644 src/unilab/tasks/locomotion/go2w/joystick.py delete mode 100644 src/unilab/tasks/locomotion/go2w/rough.py delete mode 100644 tests/envs/locomotion/go2w/test_go2w_height_scan.py delete mode 100644 tests/envs/locomotion/go2w/test_go2w_motor_control.py delete mode 100644 tests/envs/locomotion/test_go2_joystick_stand_still.py delete mode 100644 tests/envs/locomotion/test_go2_rough_height_scan.py delete mode 100644 tests/envs/test_go1_obs_noise.py delete mode 100644 tests/envs/test_go2_obs_noise.py diff --git a/conf/ppo/task/go1_joystick_rough/motrix.yaml b/conf/ppo/task/go1_joystick_rough/motrix.yaml index 21be8340a..0af29a93b 100644 --- a/conf/ppo/task/go1_joystick_rough/motrix.yaml +++ b/conf/ppo/task/go1_joystick_rough/motrix.yaml @@ -1,121 +1,18 @@ # @package _global_ +defaults: + - /task/go1_joystick_rough/mujoco + - _self_ + training: task_name: Go1JoystickRough sim_backend: motrix - play_steps: 500 - play_env_num: 16 - cam_tracking: true - cam_tracking_env_idx: 0 - cam_tracking_extra_envs: 9 - -interactive: - action_mode: policy - policy_obs_mode: auto - camera_follow_body: true - use_env_visual_model: false algo: num_envs: 2048 - num_steps_per_env: 24 - max_iterations: 1000 - empirical_normalization: false - obs_groups: - actor: - - actor - critic: - - critic - policy: - init_noise_std: 1.0 - algorithm: - learning_rate: 1.0e-3 - entropy_coef: 1.0e-2 env: render_offset_mode: zero - control_config: - action_scale: 0.25 - hip_action_scale: 0.125 - non_hip_action_scale: 0.25 - clip_actions: 100.0 - commands: - vel_limit: [[-1.0, -1.0, -1.0], [1.0, 1.0, 1.0]] - resampling_time: 10.0 - heading_command: true - heading_range: [-3.141592653589793, 3.141592653589793] - rel_standing_envs: 0.1 - terrain_curriculum: - enabled: false scene: model_file: src/unilab/assets/robots/go1/go1.xml - fragment_files: - - src/unilab/assets/robots/go1/locomotion_task.xml - terrain: - hfield_name: terrain_hfield - geom_name: floor - generator: - seed: 42 - curriculum: false - size: [8.0, 8.0] - num_rows: 6 - num_cols: 6 - border_width: 20.0 - terrain_scan: - enabled: true - geom_name: floor - termination_config: - terrain_out_of_bounds: true - terrain_distance_buffer: 3.0 - domain_rand: - randomize_base_mass: true - added_mass_range: [-1.0, 3.0] - random_com: true - randomize_kp: true - kp_multiplier_range: [0.5, 2.0] - randomize_kd: true - kd_multiplier_range: [0.5, 2.0] - push_robots: true - push_interval: 625 - max_force: [1.0, 1.0, 0.5] - -reward: - scales: - lin_vel_z: -2.0 - ang_vel_xy: -0.05 - joint_torques_l2: -2.5e-5 - joint_acc_l2: -2.5e-7 - joint_power: -2.0e-5 - stand_still: -2.0 - hip_pos: -0.5 - joint_pos_penalty: -1.0 - joint_mirror: -0.05 - action_rate: -0.01 - undesired_contacts: -1.0 - contact_forces: -1.5e-4 - tracking_lin_vel: 3.0 - tracking_ang_vel: 1.5 - feet_air_time: 0.5 - feet_air_time_variance: -1.0 - feet_contact_without_cmd: 0.1 - feet_slide: -0.1 - feet_height_body: -5.0 - feet_gait: 0.5 - upward: 1.0 - tracking_sigma: 0.25 - base_height_target: 0.33 - stand_still_command_threshold: 0.1 - joint_pos_penalty_stand_still_scale: 5.0 - joint_pos_penalty_velocity_threshold: 0.5 - joint_pos_penalty_command_threshold: 0.1 - contact_threshold: 1.0 - contact_forces_threshold: 100.0 - feet_air_time_threshold: 0.5 - feet_height_body_target: -0.2 - feet_height_body_tanh_mult: 2.0 - feet_gait_std: 0.7071067811865476 - feet_gait_max_err: 0.2 - feet_gait_velocity_threshold: 0.5 - feet_gait_command_threshold: 0.1 -play_profile: - enabled: true - env: - render_spacing: 2.0 + events: + push_robot: null diff --git a/conf/ppo/task/go1_joystick_rough/mujoco.yaml b/conf/ppo/task/go1_joystick_rough/mujoco.yaml index 04a9dc179..db1c3e920 100644 --- a/conf/ppo/task/go1_joystick_rough/mujoco.yaml +++ b/conf/ppo/task/go1_joystick_rough/mujoco.yaml @@ -1,4 +1,8 @@ # @package _global_ +defaults: + - /task/quadruped_joystick_rough/quadruped + - _self_ + training: task_name: Go1JoystickRough sim_backend: mujoco @@ -20,103 +24,46 @@ algo: num_steps_per_env: 24 max_iterations: 1000 empirical_normalization: false - obs_groups: - actor: - - actor - critic: - - critic - policy: - init_noise_std: 1.0 - algorithm: - learning_rate: 1.0e-3 - entropy_coef: 1.0e-2 + obs_groups: {actor: [actor], critic: [critic]} + policy: {init_noise_std: 1.0} + algorithm: {learning_rate: 1.0e-3, entropy_coef: 1.0e-2} env: sim_dt: 0.005 - control_config: - action_scale: 0.25 - hip_action_scale: 0.125 - non_hip_action_scale: 0.25 - clip_actions: 100.0 - commands: - vel_limit: [[-1.0, -1.0, -1.0], [1.0, 1.0, 1.0]] - resampling_time: 10.0 - heading_command: true - heading_range: [-3.141592653589793, 3.141592653589793] - rel_standing_envs: 0.1 - terrain_curriculum: - enabled: false scene: model_file: src/unilab/assets/robots/go1/go1_mujoco.xml - fragment_files: - - src/unilab/assets/robots/go1/locomotion_task.xml - terrain: - hfield_name: terrain_hfield - geom_name: floor - generator: - seed: 42 - curriculum: false - size: [8.0, 8.0] - num_rows: 6 - num_cols: 6 - border_width: 20.0 - terrain_scan: - enabled: true - geom_name: floor - termination_config: - terrain_out_of_bounds: true - terrain_distance_buffer: 3.0 - domain_rand: - randomize_base_mass: true - added_mass_range: [-1.0, 3.0] - random_com: true - randomize_kp: true - kp_multiplier_range: [0.5, 2.0] - randomize_kd: true - kd_multiplier_range: [0.5, 2.0] - push_robots: true - push_interval: 625 - max_force: [1.0, 1.0, 0.5] + fragment_files: [src/unilab/assets/robots/go1/locomotion_task.xml] + entities: + robot: + root_body_name: trunk + joint_names: + - FR_hip_joint + - FR_thigh_joint + - FR_calf_joint + - FL_hip_joint + - FL_thigh_joint + - FL_calf_joint + - RR_hip_joint + - RR_thigh_joint + - RR_calf_joint + - RL_hip_joint + - RL_thigh_joint + - RL_calf_joint + actuator_names: + - FR_hip + - FR_thigh + - FR_calf + - FL_hip + - FL_thigh + - FL_calf + - RR_hip + - RR_thigh + - RR_calf + - RL_hip + - RL_thigh + - RL_calf + body_names: [trunk] -reward: - scales: - lin_vel_z: -2.0 - ang_vel_xy: -0.05 - joint_torques_l2: -2.5e-5 - joint_acc_l2: -2.5e-7 - joint_power: -2.0e-5 - stand_still: -2.0 - hip_pos: -0.5 - joint_pos_penalty: -1.0 - joint_mirror: -0.05 - action_rate: -0.01 - undesired_contacts: -1.0 - contact_forces: -1.5e-4 - tracking_lin_vel: 3.0 - tracking_ang_vel: 1.5 - feet_air_time: 0.5 - feet_air_time_variance: -1.0 - feet_contact_without_cmd: 0.1 - feet_slide: -0.1 - feet_height_body: -5.0 - feet_gait: 0.5 - upward: 1.0 - tracking_sigma: 0.25 - base_height_target: 0.33 - stand_still_command_threshold: 0.1 - joint_pos_penalty_stand_still_scale: 5.0 - joint_pos_penalty_velocity_threshold: 0.5 - joint_pos_penalty_command_threshold: 0.1 - contact_threshold: 1.0 - contact_forces_threshold: 100.0 - feet_air_time_threshold: 0.5 - feet_height_body_target: -0.2 - feet_height_body_tanh_mult: 2.0 - feet_gait_std: 0.7071067811865476 - feet_gait_max_err: 0.2 - feet_gait_velocity_threshold: 0.5 - feet_gait_command_threshold: 0.1 play_profile: enabled: true - env: - render_spacing: 2.0 + env: {render_spacing: 2.0} diff --git a/conf/ppo/task/go2_joystick_rough/motrix.yaml b/conf/ppo/task/go2_joystick_rough/motrix.yaml index 46992b1f0..a35a5cbc5 100644 --- a/conf/ppo/task/go2_joystick_rough/motrix.yaml +++ b/conf/ppo/task/go2_joystick_rough/motrix.yaml @@ -1,120 +1,18 @@ # @package _global_ +defaults: + - /task/go2_joystick_rough/mujoco + - _self_ + training: task_name: Go2JoystickRough sim_backend: motrix - play_steps: 500 - play_env_num: 16 - cam_tracking: true - cam_tracking_env_idx: 0 - cam_tracking_extra_envs: 9 - -interactive: - action_mode: policy - policy_obs_mode: auto - camera_follow_body: true - use_env_visual_model: false algo: num_envs: 4096 - num_steps_per_env: 24 - max_iterations: 1500 - empirical_normalization: false - obs_groups: - actor: - - actor - critic: - - critic - policy: - init_noise_std: 1.0 - algorithm: - learning_rate: 1.0e-3 - entropy_coef: 1.0e-2 env: render_offset_mode: zero - control_config: - action_scale: 0.25 - hip_action_scale: 0.125 - non_hip_action_scale: 0.25 - clip_actions: 100.0 - commands: - vel_limit: [[-1.0, -1.0, -1.0], [1.0, 1.0, 1.0]] - resampling_time: 10.0 - heading_command: true - heading_range: [-3.141592653589793, 3.141592653589793] - rel_standing_envs: 0.1 - terrain_curriculum: - enabled: false scene: model_file: src/unilab/assets/robots/go2/go2.xml - fragment_files: - - src/unilab/assets/robots/go2/locomotion_task.xml - terrain: - hfield_name: terrain_hfield - geom_name: floor - generator: - seed: 42 - curriculum: false - size: [8.0, 8.0] - num_rows: 6 - num_cols: 6 - border_width: 20.0 - terrain_scan: - enabled: true - geom_name: floor - termination_config: - terrain_out_of_bounds: true - terrain_distance_buffer: 3.0 - domain_rand: - randomize_base_mass: true - added_mass_range: [-1.0, 3.0] - random_com: true - randomize_kp: true - kp_multiplier_range: [0.5, 2.0] - randomize_kd: true - kd_multiplier_range: [0.5, 2.0] - push_robots: true - push_interval: 625 - max_force: [1.0, 1.0, 0.5] -reward: - scales: - lin_vel_z: -2.0 - ang_vel_xy: -0.05 - joint_torques_l2: -2.5e-5 - joint_acc_l2: -2.5e-7 - joint_power: -2.0e-5 - stand_still: -2.0 - hip_pos: -0.5 - joint_pos_penalty: -1.0 - joint_mirror: -0.05 - action_rate: -0.01 - undesired_contacts: -1.0 - contact_forces: -1.5e-4 - tracking_lin_vel: 3.0 - tracking_ang_vel: 1.5 - feet_air_time: 0.5 - feet_air_time_variance: -1.0 - feet_contact_without_cmd: 0.1 - feet_slide: -0.1 - feet_height_body: -5.0 - feet_gait: 0.5 - upward: 1.0 - tracking_sigma: 0.25 - base_height_target: 0.3 - stand_still_command_threshold: 0.1 - joint_pos_penalty_stand_still_scale: 5.0 - joint_pos_penalty_velocity_threshold: 0.5 - joint_pos_penalty_command_threshold: 0.1 - contact_threshold: 1.0 - contact_forces_threshold: 100.0 - feet_air_time_threshold: 0.5 - feet_height_body_target: -0.2 - feet_height_body_tanh_mult: 2.0 - feet_gait_std: 0.7071067811865476 - feet_gait_max_err: 0.2 - feet_gait_velocity_threshold: 0.5 - feet_gait_command_threshold: 0.1 -play_profile: - enabled: true - env: - render_spacing: 2.0 + events: + push_robot: null diff --git a/conf/ppo/task/go2_joystick_rough/mujoco.yaml b/conf/ppo/task/go2_joystick_rough/mujoco.yaml index ea665a5fa..5d3622fdc 100644 --- a/conf/ppo/task/go2_joystick_rough/mujoco.yaml +++ b/conf/ppo/task/go2_joystick_rough/mujoco.yaml @@ -1,4 +1,8 @@ # @package _global_ +defaults: + - /task/quadruped_joystick_rough/quadruped + - _self_ + training: task_name: Go2JoystickRough sim_backend: mujoco @@ -20,103 +24,46 @@ algo: num_steps_per_env: 24 max_iterations: 1500 empirical_normalization: false - obs_groups: - actor: - - actor - critic: - - critic - policy: - init_noise_std: 1.0 - algorithm: - learning_rate: 1.0e-3 - entropy_coef: 1.0e-2 + obs_groups: {actor: [actor], critic: [critic]} + policy: {init_noise_std: 1.0} + algorithm: {learning_rate: 1.0e-3, entropy_coef: 1.0e-2} env: sim_dt: 0.002 - control_config: - action_scale: 0.25 - hip_action_scale: 0.125 - non_hip_action_scale: 0.25 - clip_actions: 100.0 - commands: - vel_limit: [[-1.0, -1.0, -1.0], [1.0, 1.0, 1.0]] - resampling_time: 10.0 - heading_command: true - heading_range: [-3.141592653589793, 3.141592653589793] - rel_standing_envs: 0.1 - terrain_curriculum: - enabled: false scene: model_file: src/unilab/assets/robots/go2/go2_mujoco.xml - fragment_files: - - src/unilab/assets/robots/go2/locomotion_task.xml - terrain: - hfield_name: terrain_hfield - geom_name: floor - generator: - seed: 42 - curriculum: false - size: [8.0, 8.0] - num_rows: 6 - num_cols: 6 - border_width: 20.0 - terrain_scan: - enabled: true - geom_name: floor - termination_config: - terrain_out_of_bounds: true - terrain_distance_buffer: 3.0 - domain_rand: - randomize_base_mass: true - added_mass_range: [-1.0, 3.0] - random_com: true - randomize_kp: true - kp_multiplier_range: [0.5, 2.0] - randomize_kd: true - kd_multiplier_range: [0.5, 2.0] - push_robots: true - push_interval: 625 - max_force: [1.0, 1.0, 0.5] + fragment_files: [src/unilab/assets/robots/go2/locomotion_task.xml] + entities: + robot: + root_body_name: base + joint_names: + - FR_hip_joint + - FR_thigh_joint + - FR_calf_joint + - FL_hip_joint + - FL_thigh_joint + - FL_calf_joint + - RR_hip_joint + - RR_thigh_joint + - RR_calf_joint + - RL_hip_joint + - RL_thigh_joint + - RL_calf_joint + actuator_names: + - FR_hip + - FR_thigh + - FR_calf + - FL_hip + - FL_thigh + - FL_calf + - RR_hip + - RR_thigh + - RR_calf + - RL_hip + - RL_thigh + - RL_calf + body_names: [base] -reward: - scales: - lin_vel_z: -2.0 - ang_vel_xy: -0.05 - joint_torques_l2: -2.5e-5 - joint_acc_l2: -2.5e-7 - joint_power: -2.0e-5 - stand_still: -2.0 - hip_pos: -0.5 - joint_pos_penalty: -1.0 - joint_mirror: -0.05 - action_rate: -0.01 - undesired_contacts: -1.0 - contact_forces: -1.5e-4 - tracking_lin_vel: 3.0 - tracking_ang_vel: 1.5 - feet_air_time: 0.5 - feet_air_time_variance: -1.0 - feet_contact_without_cmd: 0.1 - feet_slide: -0.1 - feet_height_body: -5.0 - feet_gait: 0.5 - upward: 1.0 - tracking_sigma: 0.25 - base_height_target: 0.3 - stand_still_command_threshold: 0.1 - joint_pos_penalty_stand_still_scale: 5.0 - joint_pos_penalty_velocity_threshold: 0.5 - joint_pos_penalty_command_threshold: 0.1 - contact_threshold: 1.0 - contact_forces_threshold: 100.0 - feet_air_time_threshold: 0.5 - feet_height_body_target: -0.2 - feet_height_body_tanh_mult: 2.0 - feet_gait_std: 0.7071067811865476 - feet_gait_max_err: 0.2 - feet_gait_velocity_threshold: 0.5 - feet_gait_command_threshold: 0.1 play_profile: enabled: true - env: - render_spacing: 2.0 + env: {render_spacing: 2.0} diff --git a/conf/ppo/task/go2w_joystick_rough/motrix.yaml b/conf/ppo/task/go2w_joystick_rough/motrix.yaml index d15620ddb..0f4fcbb21 100644 --- a/conf/ppo/task/go2w_joystick_rough/motrix.yaml +++ b/conf/ppo/task/go2w_joystick_rough/motrix.yaml @@ -1,99 +1,19 @@ # @package _global_ +defaults: + - /task/go2w_joystick_rough/mujoco + - _self_ + training: task_name: Go2WJoystickRough sim_backend: motrix - play_steps: 500 - play_env_num: 16 - cam_tracking: true - cam_tracking_env_idx: 0 - cam_tracking_extra_envs: 9 - -interactive: - action_mode: policy - policy_obs_mode: auto - camera_follow_body: true - use_env_visual_model: false -algo: - num_envs: 2048 - num_steps_per_env: 24 - max_iterations: 1200 - empirical_normalization: false - obs_groups: - actor: - - actor - critic: - - critic env: render_offset_mode: zero scene: model_file: src/unilab/assets/robots/go2w/go2w.xml - fragment_files: - - src/unilab/assets/robots/go2w/locomotion_task.xml - terrain: - hfield_name: terrain_hfield - geom_name: floor - generator: - seed: 42 - curriculum: false - size: [8.0, 8.0] - num_rows: 6 - num_cols: 6 - border_width: 20.0 - commands: - vel_limit: [[-1.0, -1.0, -1.0], [1.0, 1.0, 1.0]] - resampling_time: 10.0 - heading_command: true - heading_range: [-3.141592653589793, 3.141592653589793] - rel_standing_envs: 0.1 - control_config: - action_scale: 0.25 - hip_action_scale: 0.125 - wheel_action_scale: 5.0 - wheel_Kd: 0.5 - clip_actions: 100.0 - simulate_action_latency: false - terrain_scan: - enabled: true - hfield_name: terrain_hfield - geom_name: floor - termination_config: - terrain_out_of_bounds: true - terrain_distance_buffer: 3.0 - domain_rand: - randomize_base_mass: true - added_mass_range: [-1.0, 3.0] - random_com: true - com_offset_x: [-0.05, 0.05] - randomize_kp: true - kp_multiplier_range: [0.5, 1.0] - randomize_kd: true - kd_multiplier_range: [0.5, 1.0] - push_robots: true - push_interval: 625 - max_force: [1.0, 1.0, 0.5] - push_body_name: base_link + events: + push_robot: null + reward: - scales: - tracking_lin_vel: 3.0 - tracking_ang_vel: 1.5 - lin_vel_z: -2.0 - ang_vel_xy: -0.05 - orientation: -2.0 - joint_torques_l2: -2.5e-5 - joint_acc_l2: -2.5e-7 - joint_acc_wheel_l2: -2.5e-9 - joint_power: -2.0e-5 - action_rate: -0.01 - stand_still: -2.0 - hip_pos: -0.5 - joint_pos_penalty: -1.0 - joint_mirror: -0.05 - upward: 1.0 - tracking_sigma: 0.25 - base_height_target: 0.4 - only_positive_rewards: false -play_profile: - enabled: true - env: - render_spacing: 2.0 + hip_pos: + weight: -0.5 diff --git a/conf/ppo/task/go2w_joystick_rough/mujoco.yaml b/conf/ppo/task/go2w_joystick_rough/mujoco.yaml index 5bcc31038..f0c77718c 100644 --- a/conf/ppo/task/go2w_joystick_rough/mujoco.yaml +++ b/conf/ppo/task/go2w_joystick_rough/mujoco.yaml @@ -1,4 +1,8 @@ # @package _global_ +defaults: + - /task/quadruped_joystick_rough/go2w + - _self_ + training: task_name: Go2WJoystickRough sim_backend: mujoco @@ -20,80 +24,55 @@ algo: num_steps_per_env: 24 max_iterations: 1200 empirical_normalization: false - obs_groups: - actor: - - actor - critic: - - critic + obs_groups: {actor: [actor], critic: [critic]} + env: + sim_dt: 0.005 scene: model_file: src/unilab/assets/robots/go2w/go2w_mujoco.xml - fragment_files: - - src/unilab/assets/robots/go2w/locomotion_task.xml + fragment_files: [src/unilab/assets/robots/go2w/locomotion_task.xml] terrain: - hfield_name: terrain_hfield - geom_name: floor generator: - seed: 42 - curriculum: false - size: [8.0, 8.0] - num_rows: 6 - num_cols: 6 - border_width: 20.0 - commands: - vel_limit: [[-1.0, -1.0, -1.0], [1.0, 1.0, 1.0]] - resampling_time: 10.0 - heading_command: true - heading_range: [-3.141592653589793, 3.141592653589793] - rel_standing_envs: 0.1 - control_config: - action_scale: 0.25 - hip_action_scale: 0.125 - wheel_action_scale: 5.0 - wheel_Kd: 0.5 - clip_actions: 100.0 - simulate_action_latency: false - terrain_scan: - enabled: true - hfield_name: terrain_hfield - geom_name: floor - termination_config: - terrain_out_of_bounds: true - terrain_distance_buffer: 3.0 - domain_rand: - randomize_base_mass: true - added_mass_range: [-1.0, 3.0] - random_com: true - com_offset_x: [-0.05, 0.05] - randomize_kp: true - kp_multiplier_range: [0.5, 1.0] - randomize_kd: true - kd_multiplier_range: [0.5, 1.0] - push_robots: true - push_interval: 625 - max_force: [1.0, 1.0, 0.5] - push_body_name: base_link -reward: - scales: - tracking_lin_vel: 3.0 - tracking_ang_vel: 1.5 - lin_vel_z: -2.0 - ang_vel_xy: -0.05 - orientation: -2.0 - joint_torques_l2: -2.5e-5 - joint_acc_l2: -2.5e-7 - joint_acc_wheel_l2: -2.5e-9 - joint_power: -2.0e-5 - action_rate: -0.01 - stand_still: -2.0 - hip_pos: -2.0 - joint_pos_penalty: -1.0 - joint_mirror: -0.05 - upward: 1.0 - tracking_sigma: 0.25 - base_height_target: 0.4 - only_positive_rewards: false + horizontal_scale: 0.1 + entities: + robot: + root_body_name: base_link + joint_names: + - FR_hip_joint + - FR_thigh_joint + - FR_calf_joint + - FL_hip_joint + - FL_thigh_joint + - FL_calf_joint + - RR_hip_joint + - RR_thigh_joint + - RR_calf_joint + - RL_hip_joint + - RL_thigh_joint + - RL_calf_joint + - FR_wheel_joint + - FL_wheel_joint + - RR_wheel_joint + - RL_wheel_joint + actuator_names: + - FR_hip + - FR_thigh + - FR_calf + - FL_hip + - FL_thigh + - FL_calf + - RR_hip + - RR_thigh + - RR_calf + - RL_hip + - RL_thigh + - RL_calf + - FR_wheel + - FL_wheel + - RR_wheel + - RL_wheel + body_names: [base_link] + play_profile: enabled: true - env: - render_spacing: 2.0 + env: {render_spacing: 2.0} diff --git a/conf/ppo/task/quadruped_joystick_rough/base.yaml b/conf/ppo/task/quadruped_joystick_rough/base.yaml new file mode 100644 index 000000000..0c1f99f94 --- /dev/null +++ b/conf/ppo/task/quadruped_joystick_rough/base.yaml @@ -0,0 +1,166 @@ +# @package _global_ +# Shared terrain/reset/command owner for all production rough quadrupeds. +env: + scene: + default_keyframe_name: home + terrain: + hfield_name: terrain_hfield + geom_name: floor + generator: + _target_: unilab.tasks.locomotion.common.rough_manager_terms.QuadrupedRoughTerrainCfg + seed: 42 + curriculum: false + size: [8.0, 8.0] + num_rows: 6 + num_cols: 6 + border_width: 20.0 + ctrl_dt: 0.02 + max_episode_seconds: 20.0 + commands: + twist: + _target_: unilab.tasks.locomotion.common.rough_manager_terms.RoughVelocityCommandCfg + entity_name: robot + resampling_time_range: [10.0, 10.0] + heading_command: true + heading_control_stiffness: 0.5 + rel_standing_envs: 0.1 + rel_heading_envs: 1.0 + rel_world_envs: 0.0 + rel_forward_envs: 0.0 + init_velocity_prob: 0.0 + planar_dead_zone: 0.08 + ranges: + lin_vel_x: [-1.0, 1.0] + lin_vel_y: [-1.0, 1.0] + ang_vel_z: [-1.0, 1.0] + heading: [-3.141592653589793, 3.141592653589793] + events: + reset_scene_to_default: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.reset_scene_to_default + mode: reset + terrain_root_state: + _target_: unilab.managers.EventTermCfg + func: unilab.tasks.locomotion.common.rough_manager_terms.RoughTerrainReset + mode: reset + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + pose_range: + x: [-0.5, 0.5] + y: [-0.5, 0.5] + z: [0.25, 0.5] + roll: [-3.14, 3.14] + pitch: [-3.14, 3.14] + yaw: [-3.14, 3.14] + velocity_range: + x: [-0.5, 0.5] + y: [-0.5, 0.5] + z: [-0.5, 0.5] + roll: [-0.5, 0.5] + pitch: [-0.5, 0.5] + yaw: [-0.5, 0.5] + promote_frac: 0.5 + demote_frac: 0.25 + cycle_top_frac: 0.5 + spawn_height_margin: 0.05 + base_mass: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.randomize_rigid_body_mass + mode: reset + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + body_names: ".*" + mass_distribution_params: [-1.0, 3.0] + operation: add + recompute_inertia: false + base_com: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.randomize_rigid_body_com + mode: reset + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + body_names: ".*" + com_range: + x: [-0.05, 0.05] + y: [0.0, 0.0] + z: [0.0, 0.0] + pd_gains: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.pd_gains + mode: reset + params: + kp_range: [17.5, 70.0] + kd_range: [0.25, 1.0] + operation: abs + push_robot: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.push_by_setting_velocity + mode: interval + interval_range_s: [12.5, 12.5] + is_global_time: true + params: + velocity_range: + x: [-1.0, 1.0] + y: [-1.0, 1.0] + z: [-0.5, 0.5] + roll: [0.0, 0.0] + pitch: [0.0, 0.0] + yaw: [0.0, 0.0] + terminations: + time_out: + _target_: unilab.managers.TerminationTermCfg + func: unilab.envs.mdp.time_out + time_out: true + terrain_out_of_bounds: + _target_: unilab.managers.TerminationTermCfg + func: unilab.tasks.locomotion.common.rough_manager_terms.RoughTerrainOutOfBounds + time_out: true + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + distance_buffer: 3.0 + curriculum: + terrain_levels: + _target_: unilab.managers.CurriculumTermCfg + func: unilab.tasks.locomotion.common.rough_manager_terms.RoughTerrainCurriculum + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + policy_observation_group: policy + critic_observation_group: critic + +reward: + tracking_lin_vel: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel_xy_exp + weight: 3.0 + params: + std: 0.5 + command_name: twist + tracking_ang_vel: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel_z_exp + weight: 1.5 + params: + std: 0.5 + command_name: twist + lin_vel_z: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z_l2 + weight: -2.0 + ang_vel_xy: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy_l2 + weight: -0.05 + action_rate: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.mdp.action_rate_l2 + weight: -0.01 diff --git a/conf/ppo/task/quadruped_joystick_rough/go2w.yaml b/conf/ppo/task/quadruped_joystick_rough/go2w.yaml new file mode 100644 index 000000000..e8d7a6564 --- /dev/null +++ b/conf/ppo/task/quadruped_joystick_rough/go2w.yaml @@ -0,0 +1,150 @@ +# @package _global_ +defaults: + - /task/quadruped_joystick_rough/base + - _self_ + +env: + observations: + policy: + _target_: unilab.managers.ObservationGroupCfg + terms: + base_ang_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.builtin_sensor + params: {sensor_name: gyro} + scale: 0.25 + projected_gravity: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.projected_gravity_from_sensor + params: {sensor_name: upvector} + command: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.generated_commands + params: {command_name: twist} + leg_joint_pos: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_pos_rel + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + joint_names: ".*_(hip|thigh|calf)_joint" + joint_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_vel_rel + scale: 0.05 + actions: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.last_action + params: {action_name: motor} + critic: + _target_: unilab.managers.ObservationGroupCfg + terms: + base_lin_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.builtin_sensor + params: {sensor_name: local_linvel} + base_ang_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.builtin_sensor + params: {sensor_name: gyro} + projected_gravity: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.projected_gravity_from_sensor + params: {sensor_name: upvector} + command: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.generated_commands + params: {command_name: twist} + leg_joint_pos: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_pos_rel + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + joint_names: ".*_(hip|thigh|calf)_joint" + joint_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_vel_rel + actions: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.last_action + params: {action_name: motor} + height_scan: + _target_: unilab.managers.ObservationTermCfg + func: unilab.tasks.locomotion.common.rough_manager_terms.RoughHeightScan + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + geom_name: floor + vertical_offset: 0.5 + scale: 5.0 + actions: + motor: + _target_: unilab.tasks.locomotion.go2w.manager_terms.Go2WMixedActionCfg + entity_name: robot + actuator_names: [".*"] + leg_action_scale: 0.25 + hip_action_scale: 0.125 + wheel_action_scale: 5.0 + leg_kp: 35.0 + leg_kd: 0.5 + wheel_kd: 0.5 + clip_actions: 100.0 + simulate_action_latency: false + events: + pd_gains: null + motor_gains: + _target_: unilab.managers.EventTermCfg + func: unilab.tasks.locomotion.go2w.manager_terms.randomize_motor_gains + mode: reset + params: + action_name: motor + kp_multiplier_range: [0.5, 1.0] + kd_multiplier_range: [0.5, 1.0] + +reward: + orientation: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.mdp.flat_orientation_l2 + weight: -2.0 + motor_torque: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.go2w.manager_terms.motor_torque_l2 + weight: -2.5e-5 + params: {action_name: motor} + stand_still: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.stand_still_l1 + weight: -2.0 + params: + command_name: twist + command_threshold: 0.1 + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + joint_names: ".*_(hip|thigh|calf)_joint" + hip_pos: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.rough_manager_terms.joint_deviation_l2 + weight: -2.0 + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + joint_names: ".*_hip_joint" + joint_pos_penalty: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.rough_manager_terms.joint_deviation_l2 + weight: -1.0 + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + joint_names: ".*_(hip|thigh|calf)_joint" + upward: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.go2w.manager_terms.upward_l2 + weight: 1.0 diff --git a/conf/ppo/task/quadruped_joystick_rough/quadruped.yaml b/conf/ppo/task/quadruped_joystick_rough/quadruped.yaml new file mode 100644 index 000000000..51fe14ed7 --- /dev/null +++ b/conf/ppo/task/quadruped_joystick_rough/quadruped.yaml @@ -0,0 +1,109 @@ +# @package _global_ +defaults: + - /task/quadruped_joystick_rough/base + - _self_ + +env: + observations: + policy: + _target_: unilab.managers.ObservationGroupCfg + terms: + base_ang_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.builtin_sensor + params: {sensor_name: gyro} + scale: 0.25 + projected_gravity: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.projected_gravity_from_sensor + params: {sensor_name: upvector} + command: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.generated_commands + params: {command_name: twist} + joint_pos: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_pos_rel + joint_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_vel_rel + scale: 0.05 + actions: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.last_action + params: {action_name: joint_pos} + critic: + _target_: unilab.managers.ObservationGroupCfg + terms: + base_lin_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.builtin_sensor + params: {sensor_name: local_linvel} + base_ang_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.builtin_sensor + params: {sensor_name: gyro} + projected_gravity: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.projected_gravity_from_sensor + params: {sensor_name: upvector} + command: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.generated_commands + params: {command_name: twist} + joint_pos: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_pos_rel + joint_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_vel_rel + actions: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.last_action + params: {action_name: joint_pos} + height_scan: + _target_: unilab.managers.ObservationTermCfg + func: unilab.tasks.locomotion.common.rough_manager_terms.RoughHeightScan + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + geom_name: floor + vertical_offset: 0.5 + scale: 5.0 + actions: + joint_pos: + _target_: unilab.tasks.locomotion.common.rough_manager_terms.RoughJointPositionActionCfg + entity_name: robot + actuator_names: [".*"] + scale: + ".*_hip_joint": 0.125 + ".*_(thigh|calf)_joint": 0.25 + use_default_offset: true + clip_actions: 100.0 + +reward: + stand_still: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.manager_terms.stand_still_l1 + weight: -2.0 + params: + command_name: twist + command_threshold: 0.1 + hip_pos: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.rough_manager_terms.joint_deviation_l2 + weight: -0.5 + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + joint_names: ".*_hip_joint" + joint_pos_penalty: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.common.rough_manager_terms.joint_deviation_l2 + weight: -1.0 + upward: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.go2w.manager_terms.upward_l2 + weight: 1.0 diff --git a/scripts/benchmark/benchmark_drake_performance.py b/scripts/benchmark/benchmark_drake_performance.py index 78bddfe80..e4a9c10be 100644 --- a/scripts/benchmark/benchmark_drake_performance.py +++ b/scripts/benchmark/benchmark_drake_performance.py @@ -44,7 +44,7 @@ def _install_import_paths(drakeuni_src: Path | None) -> None: @dataclass(frozen=True) class TaskSpec: env_cfg_factory: Callable[[], Any] - env_cls_factory: Callable[[], type] + env_cls_factory: Callable[[], Callable[..., Any]] @dataclass @@ -101,24 +101,19 @@ def wrapped(*args: Any, **kwargs: Any) -> Any: def _task_specs() -> dict[str, TaskSpec]: def go1_cfg() -> Any: - from unilab.tasks.locomotion.go1.joystick import Go1JoystickCfg + from unilab.envs import ManagerBasedRlEnvCfg - return Go1JoystickCfg() + return ManagerBasedRlEnvCfg() - def go1_env() -> type: - from unilab.tasks.locomotion.go1.joystick import Go1WalkTask + def manager_env() -> Callable[..., Any]: + from unilab.envs import make_manager_based_rl_env - return Go1WalkTask + return make_manager_based_rl_env def go2_cfg() -> Any: - from unilab.tasks.locomotion.go2.joystick import Go2JoystickCfg + from unilab.envs import ManagerBasedRlEnvCfg - return Go2JoystickCfg() - - def go2_env() -> type: - from unilab.tasks.locomotion.go2.joystick import Go2WalkTask - - return Go2WalkTask + return ManagerBasedRlEnvCfg() def g1_tracking_cfg() -> Any: from unilab.tasks.motion_tracking.g1.tracking import G1MotionTrackingEnvCfg @@ -132,8 +127,8 @@ def g1_tracking_env() -> type: return { "g1_motion_tracking": TaskSpec(g1_tracking_cfg, g1_tracking_env), - "go1_joystick_flat": TaskSpec(go1_cfg, go1_env), - "go2_joystick_flat": TaskSpec(go2_cfg, go2_env), + "go1_joystick_flat": TaskSpec(go1_cfg, manager_env), + "go2_joystick_flat": TaskSpec(go2_cfg, manager_env), } diff --git a/scripts/benchmark/core/task_names.py b/scripts/benchmark/core/task_names.py index 409ca6137..e7b1c09d2 100644 --- a/scripts/benchmark/core/task_names.py +++ b/scripts/benchmark/core/task_names.py @@ -2,9 +2,8 @@ from dataclasses import dataclass +from unilab.envs import ManagerBasedRlEnvCfg from unilab.tasks.locomotion.g1.joystick import G1WalkFlatCfg -from unilab.tasks.locomotion.go1.joystick import Go1JoystickCfg -from unilab.tasks.locomotion.go2.joystick import Go2JoystickCfg from unilab.tasks.manipulation.sharpa_inhand.rotation import SharpaInhandRotationCfg @@ -14,6 +13,7 @@ class LocomotionTaskSpec: env_task_name: str display_name: str config_cls: type + model_file: str | None = None _TASK_SPECS = { @@ -21,13 +21,15 @@ class LocomotionTaskSpec: owner_task_id="go1_joystick_flat", env_task_name="Go1JoystickFlat", display_name="go1_joystick_flat", - config_cls=Go1JoystickCfg, + config_cls=ManagerBasedRlEnvCfg, + model_file="src/unilab/assets/robots/go1/scene_flat.xml", ), "go2_joystick_flat": LocomotionTaskSpec( owner_task_id="go2_joystick_flat", env_task_name="Go2JoystickFlat", display_name="go2_joystick_flat", - config_cls=Go2JoystickCfg, + config_cls=ManagerBasedRlEnvCfg, + model_file="src/unilab/assets/robots/go2/scene_flat.xml", ), "g1_walk_flat": LocomotionTaskSpec( owner_task_id="g1_walk_flat", @@ -75,7 +77,10 @@ def locomotion_task_spec(task_name: str) -> LocomotionTaskSpec: def locomotion_task_model_file(task_name: str) -> str: - cfg = locomotion_task_spec(task_name).config_cls() + spec = locomotion_task_spec(task_name) + if spec.model_file is not None: + return spec.model_file + cfg = spec.config_cls() scene = getattr(cfg, "scene", None) model_file = getattr(scene, "model_file", None) if model_file: diff --git a/scripts/benchmark/env/benchmark_env_step.py b/scripts/benchmark/env/benchmark_env_step.py index e1ec4e312..7fea985f8 100644 --- a/scripts/benchmark/env/benchmark_env_step.py +++ b/scripts/benchmark/env/benchmark_env_step.py @@ -276,39 +276,29 @@ def _materialize_sharpa_motrix_scene() -> str: def _go1_cfg(backend: str, config_overrides: list[str]) -> Any: - from unilab.tasks.locomotion.go1.joystick import Go1JoystickCfg + from unilab.envs import ManagerBasedRlEnvCfg - return _ppo_owner_yaml_cfg("go1_joystick_flat", backend, Go1JoystickCfg, config_overrides) + return _ppo_owner_yaml_cfg("go1_joystick_flat", backend, ManagerBasedRlEnvCfg, config_overrides) -def _go1_env_cls() -> type: - from unilab.tasks.locomotion.go1.joystick import Go1WalkTask +def _manager_env_cls() -> Callable[..., Any]: + from unilab.envs import make_manager_based_rl_env - return Go1WalkTask + return make_manager_based_rl_env def _go2_cfg(backend: str, config_overrides: list[str]) -> Any: - from unilab.tasks.locomotion.go2.joystick import Go2JoystickCfg - - return _ppo_owner_yaml_cfg("go2_joystick_flat", backend, Go2JoystickCfg, config_overrides) - - -def _go2_env_cls() -> type: - from unilab.tasks.locomotion.go2.joystick import Go2WalkTask + from unilab.envs import ManagerBasedRlEnvCfg - return Go2WalkTask + return _ppo_owner_yaml_cfg("go2_joystick_flat", backend, ManagerBasedRlEnvCfg, config_overrides) def _go2_rough_cfg(backend: str, config_overrides: list[str]) -> Any: - from unilab.tasks.locomotion.go2.rough import Go2JoystickRoughCfg - - return _ppo_owner_yaml_cfg("go2_joystick_rough", backend, Go2JoystickRoughCfg, config_overrides) - - -def _go2_rough_env_cls() -> type: - from unilab.tasks.locomotion.go2.rough import Go2JoystickRoughEnv + from unilab.envs import ManagerBasedRlEnvCfg - return Go2JoystickRoughEnv + return _ppo_owner_yaml_cfg( + "go2_joystick_rough", backend, ManagerBasedRlEnvCfg, config_overrides + ) def _go2w_cfg(backend: str, config_overrides: list[str]) -> Any: @@ -320,23 +310,15 @@ def _go2w_cfg(backend: str, config_overrides: list[str]) -> Any: def _go2w_rough_cfg(backend: str, config_overrides: list[str]) -> Any: - from unilab.tasks.locomotion.go2w.rough import Go2WJoystickRoughCfg + from unilab.envs import ManagerBasedRlEnvCfg return _ppo_owner_yaml_cfg( - "go2w_joystick_rough", backend, Go2WJoystickRoughCfg, config_overrides + "go2w_joystick_rough", backend, ManagerBasedRlEnvCfg, config_overrides ) def _go2w_env_cls() -> Callable[..., Any]: - from unilab.envs import make_manager_based_rl_env - - return make_manager_based_rl_env - - -def _go2w_rough_env_cls() -> type: - from unilab.tasks.locomotion.go2w.rough import Go2WJoystickRoughEnv - - return Go2WJoystickRoughEnv + return _manager_env_cls() def _g1_flat_cfg(backend: str, config_overrides: list[str]) -> Any: @@ -440,22 +422,22 @@ def _sharpa_inhand_env_cls() -> type: task_id="go1_joystick_flat", env_name="Go1JoystickFlat", cfg_factory=_go1_cfg, - env_cls_factory=_go1_env_cls, + env_cls_factory=_manager_env_cls, backends=("mujoco", "motrix", "mjwarp"), ), "go2": TaskConfig( task_id="go2_joystick_flat", env_name="Go2JoystickFlat", cfg_factory=_go2_cfg, - env_cls_factory=_go2_env_cls, + env_cls_factory=_manager_env_cls, backends=("mujoco", "motrix", "mjwarp"), ), "go2_rough": TaskConfig( task_id="go2_joystick_rough", env_name="Go2JoystickRough", cfg_factory=_go2_rough_cfg, - env_cls_factory=_go2_rough_env_cls, - backends=("mujoco", "motrix", "mjwarp"), + env_cls_factory=_manager_env_cls, + backends=("mujoco", "motrix"), ), "go2w": TaskConfig( task_id="go2w_joystick_flat", @@ -468,8 +450,8 @@ def _sharpa_inhand_env_cls() -> type: task_id="go2w_joystick_rough", env_name="Go2WJoystickRough", cfg_factory=_go2w_rough_cfg, - env_cls_factory=_go2w_rough_env_cls, - backends=("mujoco", "motrix", "mjwarp"), + env_cls_factory=_manager_env_cls, + backends=("mujoco", "motrix"), ), "g1": TaskConfig( task_id="g1_walk_flat", diff --git a/src/unilab/base/backend/mujoco/backend.py b/src/unilab/base/backend/mujoco/backend.py index 1decc5389..6069a6fe0 100644 --- a/src/unilab/base/backend/mujoco/backend.py +++ b/src/unilab/base/backend/mujoco/backend.py @@ -62,17 +62,26 @@ class _MuJoCoHeightScanner(BackendHeightScanner): def scan(self) -> np.ndarray: pool = self.backend._pool - if pool is None: - raise RuntimeError("MuJoCo backend pool must be materialized before hfield scanning") - - heights = pool.sample_hfield_height( - self.backend._physics_state, - hfield_geom_id=self.hfield_geom_id, - offsets=self.offsets, - frame_body_id=self.frame_body_id, - alignment=self.alignment, - output=self.output, - ) + transient_pool = pool is None + if transient_pool: + # ObservationManager evaluates terms once to infer their dimensions + # before startup randomization and the backend's formal materialize + # phase. Use a real, short-lived pool for that cold-path read; the + # scanner automatically switches to the final pool afterwards. + pool = self.backend._build_pool() + assert pool is not None + try: + heights = pool.sample_hfield_height( + self.backend._physics_state, + hfield_geom_id=self.hfield_geom_id, + offsets=self.offsets, + frame_body_id=self.frame_body_id, + alignment=self.alignment, + output=self.output, + ) + finally: + if transient_pool: + pool.close() return np.asarray(heights, dtype=self.backend._np_dtype) @@ -242,7 +251,7 @@ def _build_mujoco_scene_context(scene: SceneCfg) -> _MuJoCoSceneContext: output_dir = tempfile.TemporaryDirectory(prefix="unilab_scene_") try: - model, terrain_origins, terrain_surface_sampler = materialize_mujoco_hfield_attached_scene( + _, terrain_origins, terrain_surface_sampler = materialize_mujoco_hfield_attached_scene( model_file=scene.model_file, terrain_cfg=scene.terrain.generator, output_dir=output_dir.name, @@ -256,7 +265,10 @@ def _build_mujoco_scene_context(scene: SceneCfg) -> _MuJoCoSceneContext: raise return _MuJoCoSceneContext( - model_source=model, + # The materializer already writes the complete composed scene. Keep XML + # as the physics source so manager-requested body sensors can be injected + # before compilation just like they are for static scenes. + model_source=os.path.join(output_dir.name, "scene.xml"), model_file=scene.model_file, visual_model_file=os.path.join(output_dir.name, "scene.xml"), artifacts_dir=output_dir.name, diff --git a/src/unilab/tasks/locomotion/common/rough_manager_terms.py b/src/unilab/tasks/locomotion/common/rough_manager_terms.py new file mode 100644 index 000000000..25ed1091f --- /dev/null +++ b/src/unilab/tasks/locomotion/common/rough_manager_terms.py @@ -0,0 +1,572 @@ +"""Manager-Based terrain terms shared by the production rough quadrupeds.""" + +from __future__ import annotations + +from collections.abc import Callable, Sequence +from dataclasses import dataclass, field +from numbers import Real +from typing import TYPE_CHECKING, Any, cast +from weakref import WeakKeyDictionary + +import numpy as np + +from unilab.base.backend.base import BackendTerrainSpawnData +from unilab.dtype_config import get_global_dtype +from unilab.envs.mdp.actions.actions import JointPositionAction, JointPositionActionCfg +from unilab.envs.mdp.commands.velocity_command import ( + UniformVelocityCommand, + UniformVelocityCommandCfg, +) +from unilab.managers.manager_base import ManagerTermBase, ManagerTermBaseCfg +from unilab.managers.scene_entity_config import SceneEntityCfg +from unilab.tasks.locomotion.common.height_scan import ( + DEFAULT_SCAN_POINTS_X, + DEFAULT_SCAN_POINTS_Y, + height_scan_offsets, +) +from unilab.tasks.locomotion.common.terrain_spawn import ( + TerrainCurriculumCfg, + TerrainSpawnManager, +) +from unilab.terrains import ( + SubTerrainCfg, + TerrainGeneratorCfg, + flat, + hf_pyramid_slope, + hf_pyramid_slope_inv, + pyramid_stairs, + pyramid_stairs_inv, + random_rough, + wave_terrain, +) +from unilab.utils.rotation import np_quat_from_euler_xyz, np_quat_mul + +if TYPE_CHECKING: + from unilab.base.backend.base import BackendHeightScanner + from unilab.base.entity import Entity + from unilab.envs.manager_based_rl_env import ManagerBasedRlEnv as RoughManagerBasedRlEnv + from unilab.managers._types import ManagerBasedRlEnv + + +_DEFAULT_ASSET_CFG = SceneEntityCfg("robot") +_POSE_AXES = ("x", "y", "z", "roll", "pitch", "yaw") +_VELOCITY_AXES = ("x", "y", "z", "roll", "pitch", "yaw") + + +def _rough_sub_terrains() -> dict[str, SubTerrainCfg]: + return { + "flat": flat(proportion=0.0), + "pyramid_stairs": pyramid_stairs( + proportion=0.1, + step_height_range=(0.025, 0.10), + step_width=0.4, + platform_width=3.0, + border_width=0.2, + ), + "pyramid_stairs_inv": pyramid_stairs_inv( + proportion=0.1, + step_height_range=(0.025, 0.10), + step_width=0.4, + platform_width=3.0, + border_width=0.2, + ), + "hf_pyramid_slope": hf_pyramid_slope( + proportion=0.2, + slope_range=(0.0, 0.3), + platform_width=2.0, + border_width=0.2, + ), + "hf_pyramid_slope_inv": hf_pyramid_slope_inv( + proportion=0.2, + slope_range=(0.0, 0.3), + platform_width=2.0, + border_width=0.2, + ), + "random_rough": random_rough( + proportion=0.3, + noise_range=(0.01, 0.06), + noise_step=0.01, + border_width=0.2, + ), + "wave_terrain": wave_terrain( + proportion=0.3, + amplitude_range=(0.0, 0.12), + num_waves=4, + border_width=0.2, + ), + } + + +@dataclass(kw_only=True) +class QuadrupedRoughTerrainCfg(TerrainGeneratorCfg): + """Shared seven-terrain production generator for Go1, Go2, and Go2W.""" + + seed: int | None = 42 + curriculum: bool = False + size: tuple[float, float] = (8.0, 8.0) + horizontal_scale: float = 0.2 + vertical_scale: float = 0.005 + border_width: float = 20.0 + num_rows: int = 6 + num_cols: int = 6 + add_lights: bool = True + sub_terrains: dict[str, SubTerrainCfg] = field(default_factory=_rough_sub_terrains) + + +def _real( + value: Any, + *, + label: str, + minimum: float | None = None, + strict_minimum: bool = False, +) -> float: + if isinstance(value, (bool, np.bool_)) or not isinstance(value, Real): + raise TypeError(f"{label} must be a real number") + result = float(value) + if not np.isfinite(result): + raise ValueError(f"{label} must be finite") + if minimum is not None and (result <= minimum if strict_minimum else result < minimum): + relation = "greater than" if strict_minimum else "at least" + raise ValueError(f"{label} must be {relation} {minimum}") + return result + + +def _pair(value: Any, *, label: str) -> tuple[float, float]: + if isinstance(value, (str, bytes)) or not isinstance(value, (tuple, list)): + raise TypeError(f"{label} must be a two-value range") + if len(value) != 2: + raise ValueError(f"{label} must contain two values") + lower = _real(value[0], label=f"{label}[0]") + upper = _real(value[1], label=f"{label}[1]") + if lower > upper: + raise ValueError(f"{label} lower bound {lower} exceeds upper bound {upper}") + return lower, upper + + +def _ranges(value: Any, axes: Sequence[str], *, label: str) -> dict[str, tuple[float, float]]: + if not isinstance(value, dict): + raise TypeError(f"{label} must be a mapping") + if set(value) != set(axes): + raise ValueError(f"{label} must declare exactly {list(axes)}, got {sorted(value)}") + return {axis: _pair(value[axis], label=f"{label}.{axis}") for axis in axes} + + +def _env_ids(env: ManagerBasedRlEnv, env_ids: np.ndarray | slice | None) -> np.ndarray: + if env_ids is None: + return np.arange(env.num_envs, dtype=np.int32) + if isinstance(env_ids, slice): + return np.arange(env.num_envs, dtype=np.int32)[env_ids] + raw = np.asarray(env_ids) + if ( + raw.ndim != 1 + or not np.issubdtype(raw.dtype, np.integer) + or np.issubdtype(raw.dtype, np.bool_) + ): + raise TypeError(f"rough terrain env_ids must be a 1-D integer array, got {raw.dtype}") + ids = np.asarray(raw, dtype=np.int32) + if np.any(ids < 0) or np.any(ids >= env.num_envs): + raise IndexError(f"rough terrain env_ids out of range: {ids.tolist()}") + if np.unique(ids).size != ids.size: + raise ValueError(f"rough terrain env_ids contain duplicates: {ids.tolist()}") + return ids + + +def _terrain_generator(env: RoughManagerBasedRlEnv) -> TerrainGeneratorCfg: + scene = env._cfg.scene + terrain = None if scene is None else scene.terrain + generator = None if terrain is None else terrain.generator + if not isinstance(generator, TerrainGeneratorCfg): + raise TypeError("rough manager terms require SceneCfg.terrain.generator") + return generator + + +def _strict_height_sampler( + sample_height: Callable[[np.ndarray], np.ndarray], +) -> Callable[[np.ndarray], np.ndarray]: + def sample(xy: np.ndarray) -> np.ndarray: + points = np.asarray(xy, dtype=np.float64) + if points.ndim < 1 or points.shape[-1] != 2: + raise ValueError(f"terrain sample xy must have shape (..., 2), got {points.shape}") + if not np.isfinite(points).all(): + raise ValueError("terrain sample xy contains NaN or Inf") + heights = np.asarray(sample_height(points), dtype=np.float64) + expected = points.shape[:-1] + if heights.shape != expected: + raise ValueError( + f"terrain sample_height returned shape {heights.shape}, expected {expected}" + ) + if not np.isfinite(heights).all(): + raise ValueError("terrain sample_height returned NaN or Inf") + return heights + + return sample + + +@dataclass(frozen=True) +class _RoughTerrainContext: + spawn_manager: TerrainSpawnManager + generator: TerrainGeneratorCfg + + +_TERRAIN_CONTEXTS: WeakKeyDictionary[Any, _RoughTerrainContext] = WeakKeyDictionary() + + +def _materialize_terrain_context( + env: RoughManagerBasedRlEnv, + *, + promote_frac: float, + demote_frac: float, + cycle_top_frac: float, + spawn_height_margin: float, +) -> _RoughTerrainContext: + existing = _TERRAIN_CONTEXTS.get(env) + if existing is not None: + return existing + spawn_data = env._backend.get_terrain_spawn_data() + if not isinstance(spawn_data, BackendTerrainSpawnData): + raise NotImplementedError( + "rough terrain reset requires SimBackend.get_terrain_spawn_data()" + ) + if spawn_data.sample_height is None: + raise NotImplementedError("rough terrain reset requires terrain sample_height") + generator = _terrain_generator(env) + curriculum_cfg = TerrainCurriculumCfg( + enabled=bool(generator.curriculum), + promote_frac=promote_frac, + demote_frac=demote_frac, + cycle_top_frac=cycle_top_frac, + spawn_height_margin=spawn_height_margin, + seed=env._cfg.seed, + ) + context = _RoughTerrainContext( + spawn_manager=TerrainSpawnManager( + env.num_envs, + spawn_data.terrain_origins, + cell_size=float(generator.size[0]), + cfg=curriculum_cfg, + sample_height=_strict_height_sampler(spawn_data.sample_height), + ), + generator=generator, + ) + _TERRAIN_CONTEXTS[env] = context + return context + + +class RoughTerrainReset(ManagerTermBase): + """Stage a terrain-aware randomized root state in the reset transaction.""" + + def __init__(self, cfg: ManagerTermBaseCfg, env: RoughManagerBasedRlEnv): + super().__init__(env) + self._asset_cfg = cast(SceneEntityCfg, cfg.params.get("asset_cfg", _DEFAULT_ASSET_CFG)) + self._asset = cast("Entity", env.scene[self._asset_cfg.name]) + self._pose_range = _ranges( + cfg.params.get("pose_range"), _POSE_AXES, label="RoughTerrainReset pose_range" + ) + self._velocity_range = _ranges( + cfg.params.get("velocity_range"), + _VELOCITY_AXES, + label="RoughTerrainReset velocity_range", + ) + promote = _real( + cfg.params.get("promote_frac", 0.5), + label="RoughTerrainReset promote_frac", + minimum=0.0, + ) + demote = _real( + cfg.params.get("demote_frac", 0.25), + label="RoughTerrainReset demote_frac", + minimum=0.0, + ) + cycle = _real( + cfg.params.get("cycle_top_frac", 0.5), + label="RoughTerrainReset cycle_top_frac", + minimum=0.0, + ) + margin = _real( + cfg.params.get("spawn_height_margin", 0.05), + label="RoughTerrainReset spawn_height_margin", + minimum=0.0, + ) + self._context = _materialize_terrain_context( + env, + promote_frac=promote, + demote_frac=demote, + cycle_top_frac=cycle, + spawn_height_margin=margin, + ) + default = self._asset.data.default_root_state + expected = (env.num_envs, 13) + if default.shape != expected or not np.isfinite(default).all(): + raise ValueError( + f"RoughTerrainReset default root state must be finite {expected}, got {default.shape}" + ) + + @property + def spawn_manager(self) -> TerrainSpawnManager: + return self._context.spawn_manager + + def __call__( + self, + env: ManagerBasedRlEnv, + env_ids: np.ndarray | slice | None, + **params: Any, + ) -> None: + del params + ids = _env_ids(env, env_ids) + count = len(ids) + root_state = self._asset.data.default_root_state[ids].copy() + rng = env.rng + root_state[:, 0] += rng.uniform(*self._pose_range["x"], size=count) + root_state[:, 1] += rng.uniform(*self._pose_range["y"], size=count) + root_state[:, 2] += rng.uniform(*self._pose_range["z"], size=count) + roll = rng.uniform(*self._pose_range["roll"], size=count) + pitch = rng.uniform(*self._pose_range["pitch"], size=count) + yaw = rng.uniform(*self._pose_range["yaw"], size=count) + root_state[:, 3:7] = np_quat_mul( + root_state[:, 3:7], np_quat_from_euler_xyz(roll, pitch, yaw) + ) + for column, axis in enumerate(_VELOCITY_AXES, start=7): + root_state[:, column] = rng.uniform(*self._velocity_range[axis], size=count) + root_state[:, :3] = self._context.spawn_manager.apply_spawn( + ids, + root_state[:, :3], + yaw=yaw, + ) + if not np.isfinite(root_state).all(): + raise ValueError("RoughTerrainReset produced NaN or Inf") + self._asset.write_root_state_to_sim(root_state, env_ids=ids) + self._context.spawn_manager.record_episode_start(ids, root_state[:, :3]) + + +class RoughTerrainCurriculum(ManagerTermBase): + """Settle completed episodes before the following terrain reset selects a cell.""" + + def __init__(self, cfg: ManagerTermBaseCfg, env: RoughManagerBasedRlEnv): + super().__init__(env) + asset_cfg = cast(SceneEntityCfg, cfg.params.get("asset_cfg", _DEFAULT_ASSET_CFG)) + self._asset = cast("Entity", env.scene[asset_cfg.name]) + context = _TERRAIN_CONTEXTS.get(env) + if context is None: + raise RuntimeError("RoughTerrainCurriculum requires RoughTerrainReset") + self._spawn_manager = context.spawn_manager + + def __call__( + self, + env: RoughManagerBasedRlEnv, + env_ids: np.ndarray | slice | None, + **params: Any, + ) -> dict[str, float]: + del params + ids = _env_ids(env, env_ids) + done_ids = ids[env.reset_buf[ids]] + root_pos = self._asset.data.root_link_pos_w + expected = (env.num_envs, 3) + if root_pos.shape != expected or not np.isfinite(root_pos).all(): + raise ValueError( + f"RoughTerrainCurriculum root position must be finite {expected}, got {root_pos.shape}" + ) + return self._spawn_manager.update_on_done(done_ids, root_pos[done_ids]) + + +class RoughTerrainOutOfBounds(ManagerTermBase): + """Cached terrain-footprint truncation term.""" + + def __init__(self, cfg: ManagerTermBaseCfg, env: RoughManagerBasedRlEnv): + super().__init__(env) + asset_cfg = cast(SceneEntityCfg, cfg.params.get("asset_cfg", _DEFAULT_ASSET_CFG)) + self._asset = cast("Entity", env.scene[asset_cfg.name]) + self._distance_buffer = _real( + cfg.params.get("distance_buffer", 3.0), + label="RoughTerrainOutOfBounds distance_buffer", + minimum=0.0, + ) + generator = _terrain_generator(env) + num_cols = len(generator.sub_terrains) if generator.curriculum else generator.num_cols + self._half_width = 0.5 * ( + generator.num_rows * float(generator.size[0]) + 2.0 * generator.border_width + ) + self._half_height = 0.5 * ( + num_cols * float(generator.size[1]) + 2.0 * generator.border_width + ) + if self._distance_buffer >= min(self._half_width, self._half_height): + raise ValueError("RoughTerrainOutOfBounds distance_buffer consumes the terrain map") + + def __call__(self, env: ManagerBasedRlEnv, **params: Any) -> np.ndarray: + del params + root_pos = self._asset.data.root_link_pos_w + expected = (env.num_envs, 3) + if root_pos.shape != expected or not np.isfinite(root_pos).all(): + raise ValueError( + f"RoughTerrainOutOfBounds root position must be finite {expected}, got {root_pos.shape}" + ) + x_out = np.abs(root_pos[:, 0]) > self._half_width - self._distance_buffer + y_out = np.abs(root_pos[:, 1]) > self._half_height - self._distance_buffer + return np.asarray(x_out | y_out, dtype=np.bool_) + + +class RoughHeightScan(ManagerTermBase): + """Strict cached yaw-aligned height scan in the legacy critic format.""" + + def __init__(self, cfg: ManagerTermBaseCfg, env: RoughManagerBasedRlEnv): + super().__init__(env) + asset_cfg = cast(SceneEntityCfg, cfg.params.get("asset_cfg", _DEFAULT_ASSET_CFG)) + self._asset = cast("Entity", env.scene[asset_cfg.name]) + geom_name = cfg.params.get("geom_name", "floor") + base_body_name = cfg.params.get("base_body_name") + if not isinstance(geom_name, str) or not geom_name: + raise ValueError("RoughHeightScan geom_name must be non-empty") + if base_body_name is None: + scene = env._cfg.scene + if scene is None or asset_cfg.name not in scene.entities: + raise ValueError( + f"RoughHeightScan scene entity '{asset_cfg.name}' is not configured" + ) + base_body_name = scene.entities[asset_cfg.name].root_body_name + if not isinstance(base_body_name, str) or not base_body_name: + raise ValueError("RoughHeightScan requires base_body_name or an entity root_body_name") + points_x = cfg.params.get("measured_points_x", DEFAULT_SCAN_POINTS_X) + points_y = cfg.params.get("measured_points_y", DEFAULT_SCAN_POINTS_Y) + if isinstance(points_x, (str, bytes)) or not isinstance(points_x, Sequence): + raise TypeError("RoughHeightScan measured_points_x must be a sequence") + if isinstance(points_y, (str, bytes)) or not isinstance(points_y, Sequence): + raise TypeError("RoughHeightScan measured_points_y must be a sequence") + offsets = height_scan_offsets(points_x, points_y) + if offsets.shape[0] == 0 or not np.isfinite(offsets).all(): + raise ValueError("RoughHeightScan measured points must be finite and non-empty") + self._num_points = int(offsets.shape[0]) + self._vertical_offset = _real( + cfg.params.get("vertical_offset", 0.5), label="RoughHeightScan vertical_offset" + ) + self._scale = _real( + cfg.params.get("scale", 5.0), + label="RoughHeightScan scale", + minimum=0.0, + ) + geom_id = env._backend.get_geom_id(geom_name) + frame_body_id = env._backend.get_body_id(base_body_name) + self._scanner: BackendHeightScanner = env._backend.create_hfield_scanner( + hfield_geom_id=geom_id, + offsets=offsets, + frame_body_id=frame_body_id, + alignment="yaw", + output="height", + ) + + def __call__(self, env: ManagerBasedRlEnv, **params: Any) -> np.ndarray: + del params + raw = np.asarray(self._scanner.scan()) + expected = (env.num_envs, self._num_points) + if raw.shape != expected: + raise ValueError( + f"RoughHeightScan scanner returned shape {raw.shape}, expected {expected}" + ) + if not np.issubdtype(raw.dtype, np.number): + raise TypeError(f"RoughHeightScan scanner returned non-numeric dtype {raw.dtype}") + if not np.isfinite(raw).all(): + raise ValueError("RoughHeightScan scanner returned NaN or Inf") + base_pos = self._asset.data.root_link_pos_w + if base_pos.shape != (env.num_envs, 3): + raise ValueError( + f"RoughHeightScan root position has shape {base_pos.shape}, expected ({env.num_envs}, 3)" + ) + if not np.isfinite(base_pos).all(): + raise ValueError("RoughHeightScan root position contains NaN or Inf") + value = np.clip(base_pos[:, 2:3] - self._vertical_offset - raw, -1.0, 1.0) + return np.asarray(value * self._scale, dtype=get_global_dtype()) + + +@dataclass(kw_only=True) +class RoughJointPositionActionCfg(JointPositionActionCfg): + """Joint-position action with legacy raw-action clipping.""" + + clip_actions: float = 100.0 + + def build(self, env: ManagerBasedRlEnv) -> RoughJointPositionAction: + return RoughJointPositionAction(self, env) + + +class RoughJointPositionAction(JointPositionAction): + cfg: RoughJointPositionActionCfg # pyright: ignore[reportIncompatibleVariableOverride] + + def __init__(self, cfg: RoughJointPositionActionCfg, env: ManagerBasedRlEnv): + self._clip_actions = _real( + cfg.clip_actions, + label="RoughJointPositionActionCfg clip_actions", + minimum=0.0, + strict_minimum=True, + ) + super().__init__(cfg, env) + self._clipped_input = np.empty_like(self.raw_action) + + def process_actions(self, actions: np.ndarray) -> None: + if not isinstance(actions, np.ndarray): + raise TypeError( + f"RoughJointPositionAction expected np.ndarray, got {type(actions).__name__}" + ) + if actions.shape != self._clipped_input.shape: + raise ValueError( + f"RoughJointPositionAction expected shape {self._clipped_input.shape}, got {actions.shape}" + ) + np.clip(actions, -self._clip_actions, self._clip_actions, out=self._clipped_input) + super().process_actions(self._clipped_input) + + +@dataclass(kw_only=True) +class RoughVelocityCommandCfg(UniformVelocityCommandCfg): + """Rough-task velocity command with a planar-norm dead zone.""" + + planar_dead_zone: float = 0.08 + + def build(self, env: ManagerBasedRlEnv) -> RoughVelocityCommand: + return RoughVelocityCommand(self, env) + + +class RoughVelocityCommand(UniformVelocityCommand): + cfg: RoughVelocityCommandCfg # pyright: ignore[reportIncompatibleVariableOverride] + + def __init__(self, cfg: RoughVelocityCommandCfg, env: ManagerBasedRlEnv): + self._planar_dead_zone = _real( + cfg.planar_dead_zone, + label="RoughVelocityCommandCfg planar_dead_zone", + minimum=0.0, + ) + if cfg.heading_command and not np.isclose(cfg.rel_heading_envs, 1.0): + raise ValueError( + "RoughVelocityCommandCfg heading_command requires rel_heading_envs=1.0" + ) + super().__init__(cfg, env) + + def _resample_command(self, env_ids: np.ndarray) -> None: + super()._resample_command(env_ids) + planar = self.vel_command_b[env_ids, :2] + moving = np.linalg.norm(planar, axis=1) > self._planar_dead_zone + self.vel_command_b[env_ids, :2] = planar * moving[:, None] + + +def joint_deviation_l2( + env: ManagerBasedRlEnv, + asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG, +) -> np.ndarray: + """Squared deviation of selected joints from the keyframe default.""" + asset = cast("Entity", env.scene[asset_cfg.name]) + delta = ( + asset.data.joint_pos[:, asset_cfg.joint_ids] + - asset.data.default_joint_pos[:, asset_cfg.joint_ids] + ) + if delta.ndim != 2 or delta.shape[0] != env.num_envs or not np.isfinite(delta).all(): + raise ValueError("joint_deviation_l2 received invalid joint state") + return np.asarray(np.sum(np.square(delta), axis=1), dtype=get_global_dtype()) + + +__all__ = [ + "QuadrupedRoughTerrainCfg", + "RoughHeightScan", + "RoughJointPositionAction", + "RoughJointPositionActionCfg", + "RoughTerrainCurriculum", + "RoughTerrainOutOfBounds", + "RoughTerrainReset", + "RoughVelocityCommand", + "RoughVelocityCommandCfg", + "joint_deviation_l2", +] diff --git a/src/unilab/tasks/locomotion/go1/__init__.py b/src/unilab/tasks/locomotion/go1/__init__.py index 4f26efac3..3c5413d39 100644 --- a/src/unilab/tasks/locomotion/go1/__init__.py +++ b/src/unilab/tasks/locomotion/go1/__init__.py @@ -1,9 +1,15 @@ -from .joystick import Go1JoystickCfg, Go1WalkTask -from .rough import Go1JoystickRoughCfg, Go1JoystickRoughEnv - -__all__ = [ - "Go1JoystickCfg", - "Go1JoystickRoughCfg", - "Go1JoystickRoughEnv", - "Go1WalkTask", -] +"""Hydra-owned Manager-Based Go1 production registrations.""" + +from unilab.base import registry +from unilab.envs import ManagerBasedRlEnvCfg, make_manager_based_rl_env + +registry.register_env_config("Go1JoystickFlat", ManagerBasedRlEnvCfg) +registry.register_env("Go1JoystickFlat", make_manager_based_rl_env, sim_backend="mujoco") +registry.register_env("Go1JoystickFlat", make_manager_based_rl_env, sim_backend="motrix") +registry.register_env("Go1JoystickFlat", make_manager_based_rl_env, sim_backend="drake") + +registry.register_env_config("Go1JoystickRough", ManagerBasedRlEnvCfg) +registry.register_env("Go1JoystickRough", make_manager_based_rl_env, sim_backend="mujoco") +registry.register_env("Go1JoystickRough", make_manager_based_rl_env, sim_backend="motrix") + +__all__: list[str] = [] diff --git a/src/unilab/tasks/locomotion/go1/base.py b/src/unilab/tasks/locomotion/go1/base.py deleted file mode 100644 index 8b82c6681..000000000 --- a/src/unilab/tasks/locomotion/go1/base.py +++ /dev/null @@ -1,40 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass, field - -from unilab.tasks.locomotion.common.base import ( - BaseNoiseConfig, - LocomotionBaseCfg, - LocomotionBaseEnv, - PdControlConfig, -) - - -@dataclass -class NoiseConfig(BaseNoiseConfig): - pass - - -@dataclass -class ControlConfig(PdControlConfig): - pass - - -@dataclass -class Asset: - base_name = "trunk" - foot_name = "foot" - ground = "floor" - - -@dataclass -class Go1BaseCfg(LocomotionBaseCfg): - noise_config: NoiseConfig = field(default_factory=NoiseConfig) # type: ignore[assignment] - control_config: ControlConfig = field(default_factory=ControlConfig) # type: ignore[assignment] - asset: Asset = field(default_factory=Asset) - sim_dt: float = 0.01 - ctrl_dt: float = 0.02 - - -class Go1BaseEnv(LocomotionBaseEnv): - _cfg: Go1BaseCfg # pyright: ignore[reportIncompatibleVariableOverride] diff --git a/src/unilab/tasks/locomotion/go1/joystick.py b/src/unilab/tasks/locomotion/go1/joystick.py deleted file mode 100644 index 8df9d23d3..000000000 --- a/src/unilab/tasks/locomotion/go1/joystick.py +++ /dev/null @@ -1,262 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass, field -from typing import Any - -import numpy as np - -from unilab.assets import ASSETS_ROOT_PATH -from unilab.base import registry -from unilab.base.backend import create_backend, env_backend_kwargs -from unilab.base.np_env import NpEnvState -from unilab.base.scene import SceneCfg -from unilab.dtype_config import get_global_dtype -from unilab.envs import ManagerBasedRlEnvCfg, make_manager_based_rl_env -from unilab.tasks.locomotion.common import rewards -from unilab.tasks.locomotion.common.commands import Commands -from unilab.tasks.locomotion.common.domain_rand import DomainRandConfig -from unilab.tasks.locomotion.common.dr_provider import LocomotionDRProvider -from unilab.tasks.locomotion.common.rewards import RewardContext -from unilab.tasks.locomotion.common.terrain_spawn import ( - TerrainCurriculumCfg, - TerrainSpawnManager, -) -from unilab.tasks.locomotion.go1.base import Go1BaseCfg, Go1BaseEnv - - -@dataclass -class InitState: - pos = [0.0, 0.0, 0.45] - - -@dataclass -class RewardConfig: - scales: dict[str, float] - tracking_sigma: float - base_height_target: float - - -@dataclass -class JoystickSensor: - local_linvel = "local_linvel" - gyro = "gyro" - feet_force = ["FL_foot_contact", "FR_foot_contact", "RL_foot_contact", "RR_foot_contact"] - feet_pos = ["FL_pos", "FR_pos", "RL_pos", "RR_pos"] - - -@dataclass -class Go1JoystickCfg(Go1BaseCfg): - scene: SceneCfg = field( # pyright: ignore[reportIncompatibleVariableOverride] - default_factory=lambda: SceneCfg( - model_file=str(ASSETS_ROOT_PATH / "robots" / "go1" / "scene_flat.xml") - ) - ) - max_episode_seconds: float = 20.0 # pyright: ignore[reportIncompatibleVariableOverride] - init_state: InitState = field(default_factory=InitState) - commands: Commands = field(default_factory=Commands) - reward_config: RewardConfig | None = None - sensor: JoystickSensor = field(default_factory=JoystickSensor) # type: ignore[assignment] - domain_rand: DomainRandConfig = field( - default_factory=lambda: DomainRandConfig( - randomize_base_mass=True, - random_com=True, - push_robots=True, - ) - ) - - -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]: - return env._compute_obs( # type: ignore[no-any-return] - info_updates, linvel, gyro, gravity, dof_pos, dof_vel, env.feet_phase[env_ids] - ) - - -class Go1WalkTask(Go1BaseEnv): - _cfg: Go1JoystickCfg # pyright: ignore[reportIncompatibleVariableOverride] - - def __init__(self, cfg: Go1JoystickCfg, num_envs=1, backend_type="mujoco"): - if cfg.reward_config is None: - raise ValueError("reward_config must be provided via Hydra configuration") - self._scene_terrain_origins: np.ndarray | None = None - scene_cfg = cfg.scene - terrain_generator = scene_cfg.terrain.generator if scene_cfg.terrain is not None else None - backend = create_backend( - backend_type, - cfg.scene, - num_envs, - cfg.sim_dt, - base_name=cfg.asset.base_name, - push_body_name=cfg.domain_rand.push_body_name, - position_actuator_gains={"kp": cfg.control_config.Kp, "kd": cfg.control_config.Kd}, - **env_backend_kwargs(cfg), - ) - terrain_spawn_data = backend.get_terrain_spawn_data() - if terrain_spawn_data is not None: - self._scene_terrain_origins = terrain_spawn_data.terrain_origins - super().__init__(cfg, backend, num_envs) - self._enable_reward_log = True - self._reward_cfg = cfg.reward_config - self._init_reward_functions() - if terrain_spawn_data is not None and terrain_generator is not None: - self._spawn = TerrainSpawnManager( - num_envs, - terrain_spawn_data.terrain_origins, - cell_size=float(terrain_generator.size[0]), - cfg=getattr(cfg, "terrain_curriculum", TerrainCurriculumCfg()), - sample_height=terrain_spawn_data.sample_height, - ) - self.phase = np.zeros((num_envs,), dtype=np.float32) - self.feet_phase = np.zeros((num_envs, len(cfg.sensor.feet_force)), dtype=np.float32) - self.gait_frequency = 2 - self.feet_force = np.zeros((num_envs, len(cfg.sensor.feet_force), 3), dtype=np.float32) - self._init_domain_randomization(Go1JoystickDomainRandomizationProvider()) - self.feet_pos = np.zeros((num_envs, len(cfg.sensor.feet_pos), 3), dtype=np.float32) - - @property - def obs_groups_spec(self) -> dict[str, int]: - # gyro(3) + gravity(3) + diff(12) + dof_vel(12) + action(12) + cmd(3) + phase(4) = 49 - return {"obs": 49, "critic": 52} - - def _init_reward_functions(self): - self._reward_fns: dict[str, Any] = { - "tracking_lin_vel": rewards.tracking_lin_vel, - "tracking_ang_vel": rewards.tracking_ang_vel, - "lin_vel_z": rewards.lin_vel_z, - "ang_vel_xy": rewards.ang_vel_xy, - "base_height": rewards.base_height, - "action_rate": rewards.action_rate, - "action_smooth": rewards.action_smooth, - "similar_to_default": rewards.similar_to_default, - "contact": self._reward_contact, - "swing_feet_z": self._reward_swing_feet_z, - } - - def update_state(self, state: NpEnvState) -> NpEnvState: - self.phase = np.fmod(self.phase + self._cfg.ctrl_dt * self.gait_frequency, 1.0) - self.feet_phase[:, 0] = self.phase - self.feet_phase[:, 3] = self.phase - - self.feet_phase[:, 1] = (self.phase + 0.5) % 1 - self.feet_phase[:, 2] = (self.phase + 0.5) % 1 - - linvel = self.get_local_linvel() - gyro = self.get_gyro() - gravity = self._backend.get_sensor_data("upvector") - dof_pos = self.get_dof_pos() - dof_vel = self.get_dof_vel() - self.feet_force[:, :, :] = 0 - for i in range(len(self._cfg.sensor.feet_force)): - self.feet_force[:, i, :] = self._backend.get_sensor_data(self._cfg.sensor.feet_force[i]) - for i in range(len(self._cfg.sensor.feet_pos)): - self.feet_pos[:, i, :] = self._backend.get_sensor_data(self._cfg.sensor.feet_pos[i]) - terminated = gravity[:, 2] <= 0.5 - reward = self._compute_reward(state.info, linvel, gyro, dof_pos) - obs = self._compute_obs( - state.info, linvel, gyro, gravity, dof_pos, dof_vel, self.feet_phase - ) - return state.replace(obs=obs, reward=reward, terminated=terminated) - - def _compute_obs( - self, info: dict, linvel, gyro, gravity, dof_pos, dof_vel, feet_phase - ) -> dict[str, np.ndarray]: - noise_cfg = self._cfg.noise_config - diff = dof_pos - self.default_angles - noisy_gyro = self._obs_noise(gyro, noise_cfg.scale_gyro) - noisy_gravity = self._obs_noise(gravity, noise_cfg.scale_gravity) - noisy_diff = self._obs_noise(diff, noise_cfg.scale_joint_angle) - noisy_dof_vel = self._obs_noise(dof_vel, noise_cfg.scale_joint_vel) - command = info["commands"] - last_actions = info.get("current_actions", np.zeros_like(diff)) - obs = np.concatenate( - [ - noisy_gyro, - -noisy_gravity, - noisy_diff, - noisy_dof_vel, - last_actions, - command, - feet_phase, - ], - axis=1, - dtype=get_global_dtype(), - ) - critic = np.concatenate( - [gyro, -gravity, diff, dof_vel, last_actions, command, feet_phase, linvel], - axis=1, - dtype=get_global_dtype(), - ) - return {"obs": obs, "critic": critic} - - def _compute_reward(self, info: dict, linvel, gyro, dof_pos) -> np.ndarray: - cfg = self._reward_cfg - ctx = RewardContext( - info=info, - linvel=linvel, - gyro=gyro, - dof_pos=dof_pos, - num_envs=self._num_envs, - default_angles=self.default_angles, - tracking_sigma=cfg.tracking_sigma, - base_height_target=cfg.base_height_target, - base_height=self._backend.get_base_pos()[:, 2], - ) - return rewards.run_reward_dispatch( - scales=cfg.scales, - fns=self._reward_fns, - ctx=ctx, - info=info, - enable_log=self._enable_reward_log, - ctrl_dt=self._cfg.ctrl_dt, - ) - - def _reward_contact(self, ctx: RewardContext) -> np.ndarray: - contact = self.feet_force[:, :, 2] > 0.1 - res = np.zeros(self.num_envs, dtype=np.float32) - for i in range(len(self._cfg.sensor.feet_force)): - is_contact = (self.feet_phase[:, i] < 0.6) | (self.gait_frequency < 1.0e-8) - res += ~(contact[:, i] ^ is_contact) - return res - - def _reward_swing_feet_z(self, ctx: RewardContext) -> np.ndarray: - is_swing = self.feet_phase >= 0.6 - target_height = 0.1 - height_error = np.square(self.feet_pos[:, :, 2] - target_height) - swing_rew = np.exp(-height_error / 0.01) * is_swing - reward: np.ndarray = np.sum(swing_rew, axis=1) / len(self._cfg.sensor.feet_pos) - return reward - - def apply_action(self, actions: np.ndarray, state: NpEnvState) -> np.ndarray: - previous_current = state.info.get("current_actions", np.zeros_like(actions)) - previous_last = state.info.get("last_actions", np.zeros_like(actions)) - state.info["previous_actions"] = previous_last - state.info["last_actions"] = previous_current - state.info["current_actions"] = actions - exec_actions = ( - state.info["last_actions"] - if self._cfg.control_config.simulate_action_latency - else actions - ) - ctrl: np.ndarray = ( - exec_actions * self._cfg.control_config.action_scale + self.default_angles - ) - return ctrl - - -# Go1JoystickCfg and Go1WalkTask remain as the rough-task bridge and as inputs -# for pre-migration A/B benchmark adapters. Neither owns the flat production -# identity, which is Hydra-owned and uses the generic Manager-Based factories. -registry.register_env_config("Go1JoystickFlat", ManagerBasedRlEnvCfg) -registry.register_env("Go1JoystickFlat", make_manager_based_rl_env, sim_backend="mujoco") -registry.register_env("Go1JoystickFlat", make_manager_based_rl_env, sim_backend="motrix") -registry.register_env("Go1JoystickFlat", make_manager_based_rl_env, sim_backend="drake") diff --git a/src/unilab/tasks/locomotion/go1/rough.py b/src/unilab/tasks/locomotion/go1/rough.py deleted file mode 100644 index 3a11d17c9..000000000 --- a/src/unilab/tasks/locomotion/go1/rough.py +++ /dev/null @@ -1,786 +0,0 @@ -"""Go1 joystick rough-terrain task.""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from typing import Any - -import numpy as np - -from unilab.assets import ASSETS_ROOT_PATH -from unilab.base import registry -from unilab.base.np_env import NpEnvState -from unilab.base.scene import SceneCfg, TerrainSceneCfg -from unilab.dr import DomainRandomizationManager, ResetPlan -from unilab.dr.dr_utils import build_common_reset_randomization, zero_actions -from unilab.dtype_config import get_global_dtype -from unilab.tasks.locomotion.common import rewards -from unilab.tasks.locomotion.common.commands import ( - Commands, - apply_heading_yaw_feedback, - sample_heading_commands, - zero_small_xy_commands, -) -from unilab.tasks.locomotion.common.domain_rand import DomainRandConfig -from unilab.tasks.locomotion.common.height_scan import ( - HeightScanConfig, - base_height_from_scan, - height_scan_obs, - init_height_scan_sensor, - raw_height_scan_obs, - terrain_out_of_bounds, -) -from unilab.tasks.locomotion.common.rewards import RewardContext -from unilab.tasks.locomotion.common.terrain_spawn import ( - TerrainCurriculumCfg, -) -from unilab.tasks.locomotion.go1.base import ControlConfig -from unilab.tasks.locomotion.go1.joystick import ( - Go1JoystickCfg, - Go1JoystickDomainRandomizationProvider, - Go1WalkTask, - JoystickSensor, - RewardConfig, -) -from unilab.terrains import ( - SubTerrainCfg, - TerrainGeneratorCfg, - flat, - hf_pyramid_slope, - hf_pyramid_slope_inv, - pyramid_stairs, - pyramid_stairs_inv, - random_rough, - wave_terrain, -) -from unilab.utils.rotation import ( - np_quat_apply_inverse, - np_quat_from_euler_xyz, - np_quat_mul, -) - -# pyright: reportIncompatibleVariableOverride=false, reportAttributeAccessIssue=false, reportCallIssue=false - - -GO1_HIP_INDICES = np.asarray([0, 3, 6, 9], dtype=np.int32) -GO1_FRONT_LEFT = 0 -GO1_FRONT_RIGHT = 1 -GO1_REAR_LEFT = 2 -GO1_REAR_RIGHT = 3 - - -@dataclass -class RoughControlConfig(ControlConfig): - hip_action_scale: float = 0.125 - non_hip_action_scale: float = 0.25 - clip_actions: float = 100.0 - - -@dataclass -class Go1RoughDomainRandConfig(DomainRandConfig): - randomize_kp: bool = True - kp_multiplier_range: list[float] = field(default_factory=lambda: [0.9, 1.1]) - randomize_kd: bool = True - kd_multiplier_range: list[float] = field(default_factory=lambda: [0.9, 1.1]) - - -@dataclass -class RoughCommands(Commands): - vel_limit: list[list[float]] = field( - default_factory=lambda: [[-1.0, -1.0, -1.0], [1.0, 1.0, 1.0]] - ) - resampling_time: float = 10.0 - heading_command: bool = True - heading_range: list[float] = field(default_factory=lambda: [-np.pi, np.pi]) - - -@dataclass -class RoughRewardConfig(RewardConfig): - stand_still_command_threshold: float = 0.1 - joint_pos_penalty_stand_still_scale: float = 5.0 - joint_pos_penalty_velocity_threshold: float = 0.5 - joint_pos_penalty_command_threshold: float = 0.1 - contact_threshold: float = 1.0 - contact_forces_threshold: float = 100.0 - feet_air_time_threshold: float = 0.5 - feet_height_body_target: float = -0.2 - feet_height_body_tanh_mult: float = 2.0 - feet_gait_std: float = float(np.sqrt(0.5)) - feet_gait_max_err: float = 0.2 - feet_gait_velocity_threshold: float = 0.5 - feet_gait_command_threshold: float = 0.1 - - -@dataclass -class RoughJoystickSensor(JoystickSensor): - feet_vel = ["FL_vel", "FR_vel", "RL_vel", "RR_vel"] - undesired_contact = [ - "base1_contact", - "base2_contact", - "base3_contact", - "FL_hip_contact", - "FR_hip_contact", - "RL_hip_contact", - "RR_hip_contact", - "FL_thigh_contact", - "FR_thigh_contact", - "RL_thigh_contact", - "RR_thigh_contact", - "FL_calf_contact1", - "FR_calf_contact1", - "RL_calf_contact1", - "RR_calf_contact1", - "FL_calf_contact2", - "FR_calf_contact2", - "RL_calf_contact2", - "RR_calf_contact2", - ] - - -@dataclass -class RoughTerminationConfig: - terrain_out_of_bounds: bool = True - terrain_distance_buffer: float = 3.0 - - -@dataclass(kw_only=True) -class Go1RoughTerrainCfg(TerrainGeneratorCfg): - size: tuple[float, float] = (8.0, 8.0) - num_rows: int = 6 - num_cols: int = 6 - border_width: float = 1.0 - add_lights: bool = True - horizontal_scale: float = 0.2 - - sub_terrains: dict[str, SubTerrainCfg] = field( - default_factory=lambda: { - "flat": flat(proportion=0.0), - "pyramid_stairs": pyramid_stairs( - proportion=0.1, - step_height_range=(0.025, 0.10), - step_width=0.4, - platform_width=3.0, - border_width=0.2, - ), - "pyramid_stairs_inv": pyramid_stairs_inv( - proportion=0.1, - step_height_range=(0.025, 0.10), - step_width=0.4, - platform_width=3.0, - border_width=0.2, - ), - "hf_pyramid_slope": hf_pyramid_slope( - proportion=0.2, - slope_range=(0.0, 0.3), - platform_width=2.0, - border_width=0.2, - ), - "hf_pyramid_slope_inv": hf_pyramid_slope_inv( - proportion=0.2, - slope_range=(0.0, 0.3), - platform_width=2.0, - border_width=0.2, - ), - "random_rough": random_rough( - proportion=0.3, - noise_range=(0.01, 0.06), - noise_step=0.01, - border_width=0.2, - ), - "wave_terrain": wave_terrain( - proportion=0.3, - amplitude_range=(0.0, 0.12), - num_waves=4, - border_width=0.2, - ), - } - ) - - -@registry.envcfg("Go1JoystickRough") -@dataclass -class Go1JoystickRoughCfg(Go1JoystickCfg): - scene: SceneCfg = field( - default_factory=lambda: SceneCfg( - model_file=str(ASSETS_ROOT_PATH / "robots" / "go1" / "go1.xml"), - fragment_files=[ - str(ASSETS_ROOT_PATH / "robots" / "go1" / "locomotion_task.xml"), - ], - terrain=TerrainSceneCfg( - generator=Go1RoughTerrainCfg(), - hfield_name="terrain_hfield", - geom_name="floor", - ), - ) - ) - control_config: RoughControlConfig = field(default_factory=RoughControlConfig) - commands: RoughCommands = field(default_factory=RoughCommands) - terrain_scan: HeightScanConfig = field(default_factory=HeightScanConfig) - termination_config: RoughTerminationConfig = field(default_factory=RoughTerminationConfig) - terrain_curriculum: TerrainCurriculumCfg = field(default_factory=TerrainCurriculumCfg) - sensor: RoughJoystickSensor = field(default_factory=RoughJoystickSensor) - domain_rand: Go1RoughDomainRandConfig = field(default_factory=Go1RoughDomainRandConfig) - reward_config: RoughRewardConfig | None = None - - -class Go1JoystickRoughDomainRandomizationProvider(Go1JoystickDomainRandomizationProvider): - def _sample_commands(self, env: Any, num_reset: int) -> np.ndarray: - commands = super()._sample_commands(env, num_reset) - zero_small_xy_commands(commands, threshold=0.08) - standing_prob = env.cfg.commands.rel_standing_envs - if standing_prob > 0.0: - standing = np.random.uniform(size=(num_reset,)) < min(standing_prob, 1.0) - commands[standing] = 0.0 - if env.cfg.commands.heading_command: - commands[:, 2] = 0.0 - return commands - - 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)) - qpos[:, 2] += np.random.uniform(0.25, 0.5, (num_reset,)) - qpos[:, 0:3] += env._spawn.origins_for(env_ids) - roll = np.random.uniform(-3.14, 3.14, (num_reset,)) - pitch = np.random.uniform(-3.14, 3.14, (num_reset,)) - yaw = np.random.uniform(-3.14, 3.14, (num_reset,)) - qpos[:, 3:7] = np_quat_mul(qpos[:, 3:7], np_quat_from_euler_xyz(roll, pitch, yaw)) - qvel[:, 0:6] = np.asarray( - np.random.uniform(-0.5, 0.5, size=(num_reset, 6)), dtype=get_global_dtype() - ) - commands = self._sample_commands(env, num_reset) - info_updates: dict[str, Any] = { - "commands": commands, - "current_actions": zero_actions(num_reset, env._num_action), - "last_actions": zero_actions(num_reset, env._num_action), - "qacc": np.zeros((num_reset, env._num_action), dtype=get_global_dtype()), - "torques": np.zeros((num_reset, env._num_action), dtype=get_global_dtype()), - } - if env.cfg.commands.heading_command: - info_updates["heading_commands"] = sample_heading_commands(env, num_reset) - env._spawn.record_episode_start(env_ids, qpos[:, 0:3]) - return ResetPlan( - env_ids=env_ids, - qpos=qpos, - qvel=qvel, - info_updates=info_updates, - randomization=build_common_reset_randomization(env, num_reset), - ) - - -@registry.env("Go1JoystickRough", sim_backend="mujoco") -class Go1JoystickRoughEnv(Go1WalkTask): - _cfg: Go1JoystickRoughCfg - _reward_cfg: RoughRewardConfig - _height_scan_dim: int = 0 - - def __init__(self, cfg: Go1JoystickRoughCfg, num_envs=1, backend_type="mujoco"): - super().__init__(cfg, num_envs=num_envs, backend_type=backend_type) - self._dr_manager = DomainRandomizationManager( - self, Go1JoystickRoughDomainRandomizationProvider() - ) - - self._last_dof_vel_for_acc = np.zeros( - (num_envs, self._num_action), dtype=get_global_dtype() - ) - joint_range = self._backend.get_joint_range() - self._joint_range = ( - np.asarray(joint_range, dtype=get_global_dtype()) if joint_range is not None else None - ) - - self._action_scale = np.full( - (self._num_action,), - float(cfg.control_config.non_hip_action_scale), - dtype=get_global_dtype(), - ) - self._action_scale[GO1_HIP_INDICES] = float(cfg.control_config.hip_action_scale) - - self.feet_vel = np.zeros((num_envs, len(cfg.sensor.feet_vel), 3), dtype=np.float32) - self._last_foot_contact = np.zeros((num_envs, len(cfg.sensor.feet_force)), dtype=bool) - self._current_air_time = np.zeros((num_envs, len(cfg.sensor.feet_force)), dtype=np.float32) - self._current_contact_time = np.zeros( - (num_envs, len(cfg.sensor.feet_force)), dtype=np.float32 - ) - self._last_air_time = np.zeros((num_envs, len(cfg.sensor.feet_force)), dtype=np.float32) - self._last_contact_time = np.zeros((num_envs, len(cfg.sensor.feet_force)), dtype=np.float32) - self._first_foot_contact = np.zeros((num_envs, len(cfg.sensor.feet_force)), dtype=bool) - - init_height_scan_sensor(self, cfg.terrain_scan, cfg.asset.base_name) - - @property - def obs_groups_spec(self) -> dict[str, int]: - # Match the rough format: policy obs = 45, critic = 48 + height_scan. - return {"obs": 45, "critic": 48 + self._height_scan_dim} - - def reset(self, env_indices: np.ndarray) -> tuple[dict[str, np.ndarray], dict]: - env_ids = np.asarray(env_indices, dtype=np.int32) - obs, info = super().reset(env_ids) - dof_vel = self.get_dof_vel() - if dof_vel.shape[0] == self._num_envs: - self._last_dof_vel_for_acc[env_ids] = dof_vel[env_ids] - self._reset_contact_timers(env_ids) - return obs, info - - def _upright_scale(self, gravity: np.ndarray | None) -> np.ndarray: - return rewards.upright_scale(gravity, self._num_envs) - - def _init_reward_functions(self): - scale_gravity = self._upright_scale - - def gated(fn): - return lambda ctx: fn(ctx) * scale_gravity(ctx.gravity) - - def _joint_pos_penalty(ctx: RewardContext) -> np.ndarray: - cfg = self._reward_cfg - return rewards.joint_pos_penalty( - ctx, - stand_still_scale=cfg.joint_pos_penalty_stand_still_scale, - velocity_threshold=cfg.joint_pos_penalty_velocity_threshold, - command_threshold=cfg.joint_pos_penalty_command_threshold, - ) * scale_gravity(ctx.gravity) - - def _stand_still(ctx: RewardContext) -> np.ndarray: - return rewards.stand_still( - ctx, command_threshold=self._reward_cfg.stand_still_command_threshold - ) * scale_gravity(ctx.gravity) - - self._reward_fns: dict[str, Any] = { - "tracking_lin_vel": gated(rewards.tracking_lin_vel), - "tracking_ang_vel": gated(rewards.tracking_ang_vel), - "lin_vel_z": gated(rewards.lin_vel_z), - "ang_vel_xy": gated(rewards.ang_vel_xy), - "dof_torques_l2": gated(rewards.dof_torques_l2), - "joint_torques_l2": gated(rewards.dof_torques_l2), - "dof_acc_l2": gated(rewards.dof_acc_l2), - "joint_acc_l2": gated(rewards.dof_acc_l2), - "joint_power": gated(rewards.joint_power), - "stand_still": _stand_still, - "hip_pos": self._reward_hip_pos, - "joint_pos_penalty": _joint_pos_penalty, - "joint_mirror": self._reward_joint_mirror, - "action_rate": rewards.action_rate, - "action_rate_l2": rewards.action_rate, - "undesired_contacts": self._reward_undesired_contacts, - "contact_forces": self._reward_contact_forces, - "feet_air_time": self._reward_feet_air_time, - "feet_air_time_variance": self._reward_feet_air_time_variance, - "feet_contact_without_cmd": self._reward_feet_contact_without_cmd, - "feet_slide": self._reward_feet_slide, - "feet_height_body": self._reward_feet_height_body, - "feet_gait": self._reward_feet_gait, - "upward": rewards.upward, - } - - def apply_action(self, actions: np.ndarray, state: NpEnvState) -> np.ndarray: - clipped = np.asarray( - np.clip( - actions, - -float(self._cfg.control_config.clip_actions), - float(self._cfg.control_config.clip_actions), - ), - dtype=get_global_dtype(), - ) - state.info["last_actions"] = state.info.get("current_actions", np.zeros_like(clipped)) - state.info["current_actions"] = clipped - exec_actions = ( - state.info["last_actions"] - if self._cfg.control_config.simulate_action_latency - else clipped - ) - return np.asarray( - exec_actions * self._action_scale + self.default_angles, dtype=get_global_dtype() - ) - - def update_state(self, state: NpEnvState) -> NpEnvState: - self._update_commands(state.info) - self.phase = np.fmod(self.phase + self._cfg.ctrl_dt * self.gait_frequency, 1.0) - self.feet_phase[:, 0] = self.phase - self.feet_phase[:, 3] = self.phase - self.feet_phase[:, 1] = (self.phase + 0.5) % 1 - self.feet_phase[:, 2] = (self.phase + 0.5) % 1 - - linvel = self.get_local_linvel() - gyro = self.get_gyro() - gravity = self._backend.get_sensor_data("upvector") - dof_pos = self.get_dof_pos() - dof_vel = self.get_dof_vel() - - self.feet_force[:, :, :] = 0 - for i in range(len(self._cfg.sensor.feet_force)): - self.feet_force[:, i, :] = self._backend.get_sensor_data(self._cfg.sensor.feet_force[i]) - for i in range(len(self._cfg.sensor.feet_pos)): - self.feet_pos[:, i, :] = self._backend.get_sensor_data(self._cfg.sensor.feet_pos[i]) - for i in range(len(self._cfg.sensor.feet_vel)): - self.feet_vel[:, i, :] = self._backend.get_sensor_data(self._cfg.sensor.feet_vel[i]) - - self._update_contact_timers(self._foot_contact_mask()) - state.info["qacc"] = self._estimate_dof_acc(dof_vel) - state.info["torques"] = self._estimate_pd_torques(state.info, dof_pos, dof_vel) - - terminated = self._compute_terminated(gravity) - reward = self._compute_rough_reward(state.info, linvel, gyro, gravity, dof_pos, dof_vel) - obs = self._compute_obs( - state.info, linvel, gyro, gravity, dof_pos, dof_vel, self.feet_phase - ) - state = state.replace(obs=obs, reward=reward, terminated=terminated) - - done = state.terminated | state.truncated - if np.any(done): - done_indices = np.where(done)[0] - stats = self._spawn.update_on_done( - done_indices, self._backend.get_base_pos()[done_indices] - ) - if stats: - if "log" not in state.info: - state.info["log"] = {} - for k, v in stats.items(): - state.info["log"][f"terrain_curriculum/{k}"] = float(v) - return state - - def _compute_obs( - self, - info: dict, - linvel: np.ndarray, - gyro: np.ndarray, - gravity: np.ndarray, - dof_pos: np.ndarray, - dof_vel: np.ndarray, - feet_phase: np.ndarray, - ) -> dict[str, np.ndarray]: - del feet_phase - noise_cfg = self._cfg.noise_config - diff = dof_pos - self.default_angles - policy_gyro = self._obs_noise(gyro, noise_cfg.scale_gyro) * 0.25 - policy_gravity = self._obs_noise(-gravity, noise_cfg.scale_gravity) - policy_diff = self._obs_noise(diff, noise_cfg.scale_joint_angle) - policy_dof_vel = self._obs_noise(dof_vel, noise_cfg.scale_joint_vel) * 0.05 - last_actions = info.get("current_actions", np.zeros_like(diff)) - commands = info["commands"] - obs = np.concatenate( - [policy_gyro, policy_gravity, commands, policy_diff, policy_dof_vel, last_actions], - axis=1, - dtype=get_global_dtype(), - ) - critic_base = np.concatenate( - [linvel, gyro, -gravity, commands, diff, dof_vel, last_actions], - axis=1, - dtype=get_global_dtype(), - ) - critic = np.concatenate( - [critic_base, height_scan_obs(self, self._cfg.terrain_scan, critic_base.shape[0])], - axis=1, - dtype=get_global_dtype(), - ) - return {"obs": obs, "critic": critic} - - def _compute_rough_reward( - self, - info: dict, - linvel: np.ndarray, - gyro: np.ndarray, - gravity: np.ndarray, - dof_pos: np.ndarray, - dof_vel: np.ndarray, - ) -> np.ndarray: - cfg = self._reward_cfg - ctx = RewardContext( - info=info, - linvel=linvel, - gyro=gyro, - dof_pos=dof_pos, - num_envs=self._num_envs, - default_angles=self.default_angles, - tracking_sigma=cfg.tracking_sigma, - base_height_target=cfg.base_height_target, - base_height=base_height_from_scan(self, self._num_envs), - gravity=gravity, - dof_vel=dof_vel, - joint_range=self._joint_range, - ) - return rewards.run_reward_dispatch( - scales=cfg.scales, - fns=self._reward_fns, - ctx=ctx, - info=info, - enable_log=self._enable_reward_log, - ctrl_dt=self._cfg.ctrl_dt, - ) - - def _compute_terminated(self, gravity: np.ndarray) -> np.ndarray: - del gravity - return np.zeros((self._num_envs,), dtype=bool) - - def _compute_truncated(self, state: NpEnvState) -> np.ndarray: - truncated = super()._compute_truncated(state) - if self._cfg.termination_config.terrain_out_of_bounds: - terrain_scene = self._cfg.scene.terrain - terrain_cfg = terrain_scene.generator if terrain_scene is not None else None - np.logical_or( - truncated, - terrain_out_of_bounds( - self, - terrain_cfg, - float(self._cfg.termination_config.terrain_distance_buffer), - ), - out=truncated, - ) - return truncated - - def _reward_base_height_values(self, num_obs: int | None = None) -> np.ndarray: - return base_height_from_scan(self, num_obs) - - def _raw_height_scan_obs(self, num_obs: int) -> tuple[np.ndarray | None, np.ndarray | None]: - return raw_height_scan_obs(self, num_obs) - - def _estimate_dof_acc(self, dof_vel: np.ndarray) -> np.ndarray: - qacc = np.asarray((dof_vel - self._last_dof_vel_for_acc) / self._cfg.ctrl_dt) - self._last_dof_vel_for_acc[:] = dof_vel - return np.asarray(qacc, dtype=get_global_dtype()) - - def _estimate_pd_torques( - self, info: dict, dof_pos: np.ndarray, dof_vel: np.ndarray - ) -> np.ndarray: - actions = np.asarray( - info.get("current_actions", np.zeros((dof_pos.shape[0], self._num_action))), - dtype=get_global_dtype(), - ) - if self._cfg.control_config.simulate_action_latency: - actions = np.asarray(info.get("last_actions", actions), dtype=get_global_dtype()) - targets = actions * self._action_scale + self.default_angles - torques = ( - float(self._cfg.control_config.Kp) * (targets - dof_pos) - - float(self._cfg.control_config.Kd) * dof_vel - ) - return np.asarray(torques, dtype=get_global_dtype()) - - def _update_commands(self, info: dict) -> None: - commands_arr = np.asarray(info["commands"], dtype=get_global_dtype()) - resampling_time = float(self._cfg.commands.resampling_time) - if resampling_time > 0.0: - interval_steps = max(int(round(resampling_time / self._cfg.ctrl_dt)), 1) - steps = np.asarray(info.get("steps", np.zeros((self._num_envs,), dtype=np.uint32))) - resample_mask = (steps > 0) & ((steps % interval_steps) == 0) - if np.any(resample_mask): - num_resample = int(np.count_nonzero(resample_mask)) - low = np.asarray(self._cfg.commands.vel_limit[0], dtype=get_global_dtype()) - high = np.asarray(self._cfg.commands.vel_limit[1], dtype=get_global_dtype()) - sampled = np.random.uniform(low=low, high=high, size=(num_resample, 3)).astype( - get_global_dtype() - ) - zero_small_xy_commands(sampled, threshold=0.08) - commands_arr[resample_mask] = sampled - if self._cfg.commands.heading_command: - heading_commands = self._ensure_heading_commands(info, commands_arr.shape[0]) - heading_commands[resample_mask] = sample_heading_commands(self, num_resample) - info["heading_commands"] = heading_commands - - if self._cfg.commands.heading_command: - heading_commands = self._ensure_heading_commands(info, commands_arr.shape[0]) - base_quat = np.asarray(self._backend.get_base_quat(), dtype=get_global_dtype()) - if base_quat.shape[0] == commands_arr.shape[0]: - apply_heading_yaw_feedback(commands_arr, base_quat, heading_commands, stiffness=0.5) - info["commands"] = commands_arr - - def _ensure_heading_commands(self, info: dict, num_obs: int) -> np.ndarray: - heading_commands = info.get("heading_commands") - if heading_commands is None or np.asarray(heading_commands).shape != (num_obs,): - heading_commands = sample_heading_commands(self, num_obs) - heading_commands = np.asarray(heading_commands, dtype=get_global_dtype()) - info["heading_commands"] = heading_commands - return heading_commands - - def _foot_contact_mask(self) -> np.ndarray: - contact_force = np.linalg.norm(self.feet_force, axis=2) - return np.asarray(contact_force > self._reward_cfg.contact_threshold, dtype=bool) - - def _reset_contact_timers(self, env_ids: np.ndarray) -> None: - self._current_air_time[env_ids] = 0.0 - self._current_contact_time[env_ids] = 0.0 - self._last_air_time[env_ids] = 0.0 - self._last_contact_time[env_ids] = 0.0 - self._first_foot_contact[env_ids] = False - self._last_foot_contact[env_ids] = self._foot_contact_mask()[env_ids] - - def _update_contact_timers(self, contact: np.ndarray) -> None: - first_contact = contact & ~self._last_foot_contact - first_air = ~contact & self._last_foot_contact - self._first_foot_contact[:] = first_contact - self._last_air_time[first_contact] = self._current_air_time[first_contact] - self._last_contact_time[first_air] = self._current_contact_time[first_air] - self._current_air_time[contact] = 0.0 - self._current_air_time[~contact] += self._cfg.ctrl_dt - self._current_contact_time[~contact] = 0.0 - self._current_contact_time[contact] += self._cfg.ctrl_dt - self._last_foot_contact[:] = contact - - def _relative_foot_vel_body(self) -> np.ndarray: - base_quat = np.asarray(self._backend.get_base_quat(), dtype=get_global_dtype()) - base_linvel = np.asarray( - self._backend.get_sensor_data("global_linvel"), dtype=get_global_dtype() - ) - relative_vel = self.feet_vel - base_linvel[:, None, :] - flat = relative_vel.reshape(self._num_envs * relative_vel.shape[1], 3) - quat = np.repeat(base_quat, relative_vel.shape[1], axis=0) - return np_quat_apply_inverse(quat, flat).reshape(relative_vel.shape) - - def _relative_foot_pos_body(self) -> np.ndarray: - base_quat = np.asarray(self._backend.get_base_quat(), dtype=get_global_dtype()) - base_pos = np.asarray(self._backend.get_base_pos(), dtype=get_global_dtype()) - relative_pos = self.feet_pos - base_pos[:, None, :] - flat = relative_pos.reshape(self._num_envs * relative_pos.shape[1], 3) - quat = np.repeat(base_quat, relative_pos.shape[1], axis=0) - return np_quat_apply_inverse(quat, flat).reshape(relative_pos.shape) - - def _reward_hip_pos(self, ctx: RewardContext) -> np.ndarray: - diff = ctx.dof_pos[:, GO1_HIP_INDICES] - self.default_angles[GO1_HIP_INDICES] - return np.asarray( - np.sum(np.square(diff), axis=1) * self._upright_scale(ctx.gravity), - dtype=get_global_dtype(), - ) - - def _reward_joint_mirror(self, ctx: RewardContext) -> np.ndarray: - fr_rl = ctx.dof_pos[:, 0:3] - ctx.dof_pos[:, 9:12] - fl_rr = ctx.dof_pos[:, 3:6] - ctx.dof_pos[:, 6:9] - mirror = 0.5 * (np.sum(np.square(fr_rl), axis=1) + np.sum(np.square(fl_rr), axis=1)) - return np.asarray(mirror * self._upright_scale(ctx.gravity), dtype=get_global_dtype()) - - def _reward_undesired_contacts(self, ctx: RewardContext) -> np.ndarray: - contacts = [ - _force_norm_columns( - np.asarray(self._backend.get_sensor_data(name), dtype=get_global_dtype()), - ctx.num_envs, - ) - for name in self._cfg.sensor.undesired_contact - ] - if not contacts: - return np.zeros((ctx.num_envs,), dtype=get_global_dtype()) - contact_force = np.concatenate(contacts, axis=1) - contact_count = np.sum(contact_force > self._reward_cfg.contact_threshold, axis=1) - return np.asarray( - contact_count * self._upright_scale(ctx.gravity), dtype=get_global_dtype() - ) - - def _reward_contact_forces(self, ctx: RewardContext) -> np.ndarray: - force_norm = np.linalg.norm(self.feet_force, axis=2) - force_clip = np.minimum(force_norm, 1500.0) - violation = np.clip(force_clip - self._reward_cfg.contact_forces_threshold, 0.0, None) - return np.asarray( - np.sum(violation, axis=1) * self._upright_scale(ctx.gravity), - dtype=get_global_dtype(), - ) - - def _reward_feet_air_time(self, ctx: RewardContext) -> np.ndarray: - cfg = self._reward_cfg - reward = np.sum( - (self._last_air_time - cfg.feet_air_time_threshold) * self._first_foot_contact, - axis=1, - ) - moving = np.linalg.norm(ctx.info["commands"], axis=1) > 0.1 - return np.asarray( - reward * moving * self._upright_scale(ctx.gravity), dtype=get_global_dtype() - ) - - def _reward_feet_air_time_variance(self, ctx: RewardContext) -> np.ndarray: - air_var = np.var(np.clip(self._last_air_time, 0.0, 0.5), axis=1) - contact_var = np.var(np.clip(self._last_contact_time, 0.0, 0.5), axis=1) - return np.asarray( - (air_var + contact_var) * self._upright_scale(ctx.gravity), dtype=get_global_dtype() - ) - - def _reward_feet_contact_without_cmd(self, ctx: RewardContext) -> np.ndarray: - reward = np.sum(self._first_foot_contact, axis=1) - stopped = np.linalg.norm(ctx.info["commands"], axis=1) < 0.1 - return np.asarray( - reward * stopped * self._upright_scale(ctx.gravity), dtype=get_global_dtype() - ) - - def _reward_feet_slide(self, ctx: RewardContext) -> np.ndarray: - foot_vel_body = self._relative_foot_vel_body() - lateral_vel = np.linalg.norm(foot_vel_body[:, :, :2], axis=2) - reward = np.sum(lateral_vel * self._foot_contact_mask(), axis=1) - return np.asarray(reward * self._upright_scale(ctx.gravity), dtype=get_global_dtype()) - - def _reward_feet_height_body(self, ctx: RewardContext) -> np.ndarray: - cfg = self._reward_cfg - foot_pos_body = self._relative_foot_pos_body() - foot_vel_body = self._relative_foot_vel_body() - z_error = np.square(foot_pos_body[:, :, 2] - cfg.feet_height_body_target) - velocity_tanh = np.tanh( - cfg.feet_height_body_tanh_mult * np.linalg.norm(foot_vel_body[:, :, :2], axis=2) - ) - moving = np.linalg.norm(ctx.info["commands"], axis=1) > 0.1 - reward = np.sum(z_error * velocity_tanh, axis=1) - return np.asarray( - reward * moving * self._upright_scale(ctx.gravity), dtype=get_global_dtype() - ) - - def _reward_feet_gait(self, ctx: RewardContext) -> np.ndarray: - cfg = self._reward_cfg - command_norm = np.linalg.norm(ctx.info["commands"], axis=1) - body_vel = np.linalg.norm(ctx.linvel[:, :2], axis=1) - enabled = (command_norm > cfg.feet_gait_command_threshold) | ( - body_vel > cfg.feet_gait_velocity_threshold - ) - air = self._current_air_time - contact = self._current_contact_time - sync_fl_rr = _gait_sync_reward( - air, contact, GO1_FRONT_LEFT, GO1_REAR_RIGHT, cfg.feet_gait_std, cfg.feet_gait_max_err - ) - sync_fr_rl = _gait_sync_reward( - air, contact, GO1_FRONT_RIGHT, GO1_REAR_LEFT, cfg.feet_gait_std, cfg.feet_gait_max_err - ) - async_fl_fr = _gait_async_reward( - air, contact, GO1_FRONT_LEFT, GO1_FRONT_RIGHT, cfg.feet_gait_std, cfg.feet_gait_max_err - ) - async_rr_rl = _gait_async_reward( - air, contact, GO1_REAR_RIGHT, GO1_REAR_LEFT, cfg.feet_gait_std, cfg.feet_gait_max_err - ) - async_fl_rl = _gait_async_reward( - air, contact, GO1_FRONT_LEFT, GO1_REAR_LEFT, cfg.feet_gait_std, cfg.feet_gait_max_err - ) - async_fr_rr = _gait_async_reward( - air, contact, GO1_FRONT_RIGHT, GO1_REAR_RIGHT, cfg.feet_gait_std, cfg.feet_gait_max_err - ) - reward = sync_fl_rr * sync_fr_rl * async_fl_fr * async_rr_rl * async_fl_rl * async_fr_rr - return np.asarray( - reward * enabled * self._upright_scale(ctx.gravity), dtype=get_global_dtype() - ) - - -def _force_norm_columns(force: np.ndarray, num_envs: int) -> np.ndarray: - force = np.asarray(force, dtype=get_global_dtype()).reshape(num_envs, -1) - if force.shape[1] == 0: - return force - if force.shape[1] % 3 == 0: - return np.linalg.norm(force.reshape(num_envs, -1, 3), axis=2) - return np.abs(force) - - -def _gait_sync_reward( - air: np.ndarray, - contact: np.ndarray, - foot_0: int, - foot_1: int, - std: float, - max_err: float, -) -> np.ndarray: - se_air = np.clip(np.square(air[:, foot_0] - air[:, foot_1]), 0.0, max_err**2) - se_contact = np.clip(np.square(contact[:, foot_0] - contact[:, foot_1]), 0.0, max_err**2) - return np.exp(-(se_air + se_contact) / std) - - -def _gait_async_reward( - air: np.ndarray, - contact: np.ndarray, - foot_0: int, - foot_1: int, - std: float, - max_err: float, -) -> np.ndarray: - se_act_0 = np.clip(np.square(air[:, foot_0] - contact[:, foot_1]), 0.0, max_err**2) - se_act_1 = np.clip(np.square(contact[:, foot_0] - air[:, foot_1]), 0.0, max_err**2) - return np.exp(-(se_act_0 + se_act_1) / std) - - -registry.register_env("Go1JoystickRough", Go1JoystickRoughEnv, sim_backend="motrix") diff --git a/src/unilab/tasks/locomotion/go2/__init__.py b/src/unilab/tasks/locomotion/go2/__init__.py index 41578562f..75357ce80 100644 --- a/src/unilab/tasks/locomotion/go2/__init__.py +++ b/src/unilab/tasks/locomotion/go2/__init__.py @@ -1,11 +1,17 @@ +"""Hydra-owned Manager-Based Go2 production registrations.""" + +from unilab.base import registry +from unilab.envs import ManagerBasedRlEnvCfg, make_manager_based_rl_env + from . import footstand as footstand -from .joystick import Go2JoystickCfg, Go2WalkTask -from .rough import Go2JoystickRoughCfg, Go2JoystickRoughEnv - -__all__ = [ - "Go2JoystickCfg", - "Go2JoystickRoughCfg", - "Go2JoystickRoughEnv", - "Go2WalkTask", - "footstand", -] + +registry.register_env_config("Go2JoystickFlat", ManagerBasedRlEnvCfg) +registry.register_env("Go2JoystickFlat", make_manager_based_rl_env, sim_backend="mujoco") +registry.register_env("Go2JoystickFlat", make_manager_based_rl_env, sim_backend="motrix") +registry.register_env("Go2JoystickFlat", make_manager_based_rl_env, sim_backend="drake") + +registry.register_env_config("Go2JoystickRough", ManagerBasedRlEnvCfg) +registry.register_env("Go2JoystickRough", make_manager_based_rl_env, sim_backend="mujoco") +registry.register_env("Go2JoystickRough", make_manager_based_rl_env, sim_backend="motrix") + +__all__ = ["footstand"] diff --git a/src/unilab/tasks/locomotion/go2/base.py b/src/unilab/tasks/locomotion/go2/base.py deleted file mode 100644 index d3704c0a0..000000000 --- a/src/unilab/tasks/locomotion/go2/base.py +++ /dev/null @@ -1,54 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass, field - -import numpy as np - -from unilab.tasks.locomotion.common.base import ( - BaseNoiseConfig, - LocomotionBaseCfg, - LocomotionBaseEnv, - PdControlConfig, -) - - -@dataclass -class NoiseConfig(BaseNoiseConfig): - pass - - -@dataclass -class ControlConfig(PdControlConfig): - pass - - -@dataclass -class Asset: - base_name = "base" - foot_name = "foot" - ground = "floor" - - -@dataclass -class Go2BaseCfg(LocomotionBaseCfg): - noise_config: NoiseConfig = field(default_factory=NoiseConfig) # type: ignore[assignment] - control_config: ControlConfig = field(default_factory=ControlConfig) # type: ignore[assignment] - asset: Asset = field(default_factory=Asset) - sim_dt: float = 0.01 - ctrl_dt: float = 0.02 - - -class Go2BaseEnv(LocomotionBaseEnv): - _cfg: Go2BaseCfg # pyright: ignore[reportIncompatibleVariableOverride] - - def get_foot_pos(self) -> np.ndarray: - """Get foot positions. Returns shape (num_envs, 4, 3)""" - foot_names = ["FL_pos", "FR_pos", "RL_pos", "RR_pos"] - foot_pos = [self._backend.get_sensor_data(name) for name in foot_names] - return np.stack(foot_pos, axis=1) - - def get_foot_contact(self) -> np.ndarray: - """Get foot contact forces. Returns shape (num_envs, 4)""" - contact_names = ["FL_foot_contact", "FR_foot_contact", "RL_foot_contact", "RR_foot_contact"] - contacts = [self._backend.get_sensor_data(name)[:, 0] for name in contact_names] - return np.stack(contacts, axis=1) diff --git a/src/unilab/tasks/locomotion/go2/joystick.py b/src/unilab/tasks/locomotion/go2/joystick.py deleted file mode 100644 index 8009226ca..000000000 --- a/src/unilab/tasks/locomotion/go2/joystick.py +++ /dev/null @@ -1,329 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass, field -from typing import Any, cast - -import numpy as np - -from unilab.assets import ASSETS_ROOT_PATH -from unilab.base import registry -from unilab.base.backend import create_backend, env_backend_kwargs -from unilab.base.np_env import NpEnvState -from unilab.base.scene import SceneCfg -from unilab.dtype_config import get_global_dtype -from unilab.envs.manager_based_rl_env import ( - ManagerBasedRlEnvCfg, - make_manager_based_rl_env, -) -from unilab.tasks.locomotion.common import rewards -from unilab.tasks.locomotion.common.base import Sensor -from unilab.tasks.locomotion.common.commands import Commands -from unilab.tasks.locomotion.common.domain_rand import DomainRandConfig -from unilab.tasks.locomotion.common.dr_provider import LocomotionDRProvider -from unilab.tasks.locomotion.common.rewards import RewardContext -from unilab.tasks.locomotion.common.terrain_spawn import ( - TerrainCurriculumCfg, - TerrainSpawnManager, -) -from unilab.tasks.locomotion.go2.base import Go2BaseCfg, Go2BaseEnv - - -@dataclass -class InitState: - pos = [0.0, 0.0, 0.42] - - -@dataclass -class Go2DomainRandConfig(DomainRandConfig): - randomize_kp: bool = True - kp_multiplier_range: list[float] = field(default_factory=lambda: [0.9, 1.1]) - - randomize_kd: bool = True - kd_multiplier_range: list[float] = field(default_factory=lambda: [0.9, 1.1]) - - -@dataclass -class RewardConfig: - scales: dict[str, float] - tracking_sigma: float - base_height_target: float - target_foot_height: float = 0.1 - - -@dataclass -class JoystickSensor(Sensor): - local_linvel = "local_linvel" - gyro = "gyro" - feet_force = ["FL_foot_contact", "FR_foot_contact", "RL_foot_contact", "RR_foot_contact"] - feet_pos = ["FL_pos", "FR_pos", "RL_pos", "RR_pos"] - - -@dataclass -class Go2JoystickCfg(Go2BaseCfg): - scene: SceneCfg = field( # pyright: ignore[reportIncompatibleVariableOverride] - default_factory=lambda: SceneCfg( - model_file=str(ASSETS_ROOT_PATH / "robots" / "go2" / "scene_flat.xml") - ) - ) - max_episode_seconds: float = 20.0 # pyright: ignore[reportIncompatibleVariableOverride] - init_state: InitState = field(default_factory=InitState) - commands: Commands = field(default_factory=Commands) - reward_config: RewardConfig | None = None - sensor: JoystickSensor = field( # pyright: ignore[reportIncompatibleVariableOverride] - default_factory=JoystickSensor - ) - domain_rand: Go2DomainRandConfig = field(default_factory=Go2DomainRandConfig) - terrain_curriculum: TerrainCurriculumCfg = field(default_factory=TerrainCurriculumCfg) - - -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]: - return cast( - dict[str, np.ndarray], - env._compute_obs( - info_updates, linvel, gyro, gravity, dof_pos, dof_vel, env.feet_phase[env_ids] - ), - ) - - -class Go2WalkTask(Go2BaseEnv): - _cfg: Go2JoystickCfg # pyright: ignore[reportIncompatibleVariableOverride] - - def __init__(self, cfg: Go2JoystickCfg, num_envs=1, backend_type="mujoco"): - if cfg.reward_config is None: - raise ValueError("reward_config must be provided via Hydra configuration") - - self._scene_terrain_origins: np.ndarray | None = None - scene_cfg = cfg.scene - terrain_generator = scene_cfg.terrain.generator if scene_cfg.terrain is not None else None - - backend = create_backend( - backend_type, - cfg.scene, - num_envs, - cfg.sim_dt, - base_name=cfg.asset.base_name, - push_body_name=cfg.domain_rand.push_body_name, - position_actuator_gains=cfg.control_config.position_gains(), - **env_backend_kwargs(cfg), - ) - terrain_spawn_data = backend.get_terrain_spawn_data() - self._terrain_surface_sample_height = ( - None if terrain_spawn_data is None else terrain_spawn_data.sample_height - ) - if terrain_spawn_data is not None: - self._scene_terrain_origins = terrain_spawn_data.terrain_origins - super().__init__(cfg, backend, num_envs) - self._enable_reward_log = True - self._reward_cfg = cfg.reward_config - self._init_reward_functions() - self._init_domain_randomization(self._make_dr_provider()) - if terrain_spawn_data is not None and terrain_generator is not None: - self._spawn = TerrainSpawnManager( - num_envs, - terrain_spawn_data.terrain_origins, - cell_size=float(terrain_generator.size[0]), - cfg=cfg.terrain_curriculum, - sample_height=self._terrain_surface_sample_height, - ) - self.phase = np.zeros((num_envs,), dtype=np.float32) - self.feet_phase = np.zeros((num_envs, len(cfg.sensor.feet_force)), dtype=np.float32) - self.gait_frequency = 2 - self.feet_force = np.zeros((num_envs, len(cfg.sensor.feet_force), 3), dtype=np.float32) - self.feet_pos = np.zeros((num_envs, len(cfg.sensor.feet_pos), 3), dtype=np.float32) - - def _make_dr_provider(self) -> LocomotionDRProvider: - """Domain-randomization provider for this task. Subclasses override to - inject robot-specific behaviour (e.g. A2's per-joint base gains).""" - return Go2JoystickDomainRandomizationProvider() - - def _update_commands(self, info: dict) -> None: - """Mid-episode command resample hook. No-op for Go2 flat; subclasses - (e.g. A2) override to draw standing-aware commands at intervals.""" - - def get_playback_model(self, env_index: int | None = None) -> Any: - return super().get_playback_model(env_index) - - @property - def obs_groups_spec(self) -> dict[str, int]: - # gyro(3) + gravity(3) + diff(12) + dof_vel(12) + action(12) + cmd(3) + phase(4) = 49 - return {"obs": 49, "critic": 52} - - def _init_reward_functions(self): - self._reward_fns: dict[str, Any] = { - "tracking_lin_vel": rewards.tracking_lin_vel, - "tracking_ang_vel": rewards.tracking_ang_vel, - "lin_vel_z": rewards.lin_vel_z, - "ang_vel_xy": rewards.ang_vel_xy, - "base_height": rewards.base_height, - "action_rate": rewards.action_rate, - "similar_to_default": rewards.similar_to_default, - "alive": rewards.alive, - "swing_feet_z": self._reward_swing_feet_z, - "contact": self._reward_contact, - "foot_drag": self._reward_foot_drag, - } - - def _advance_phase(self, phase: np.ndarray) -> np.ndarray: - """Advance the gait phase clock one control step. - - Subclasses override to modulate the advance (e.g. A2 freezes standing - envs). The base advance is unconditional, matching the Go2 flat gait.""" - return np.fmod(phase + self._cfg.ctrl_dt * self.gait_frequency, 1.0) - - def update_state(self, state: NpEnvState) -> NpEnvState: - self._update_commands(state.info) - self.phase = self._advance_phase(self.phase) - self.feet_phase[:, 0] = self.phase - self.feet_phase[:, 3] = self.phase - - self.feet_phase[:, 1] = (self.phase + 0.5) % 1 - self.feet_phase[:, 2] = (self.phase + 0.5) % 1 - - linvel = self.get_local_linvel() - gyro = self.get_gyro() - gravity = self._backend.get_sensor_data("upvector") - dof_pos = self.get_dof_pos() - dof_vel = self.get_dof_vel() - self.feet_force[:, :, :] = 0 - for i in range(len(self._cfg.sensor.feet_force)): - contact = self._backend.get_sensor_data(self._cfg.sensor.feet_force[i]) - if contact.shape[1] == 1: - self.feet_force[:, i, 2] = contact[:, 0] - elif contact.shape[1] == 3: - self.feet_force[:, i, :] = contact - else: - raise ValueError( - "foot contact sensor must return either scalar found flags " - f"or 3D force vectors, got {contact.shape}" - ) - for i in range(len(self._cfg.sensor.feet_pos)): - self.feet_pos[:, i, :] = self._backend.get_sensor_data(self._cfg.sensor.feet_pos[i]) - terminated = gravity[:, 2] <= 0.5 - reward = self._compute_reward(state.info, linvel, gyro, dof_pos) - obs = self._compute_obs( - state.info, linvel, gyro, gravity, dof_pos, dof_vel, self.feet_phase - ) - state = state.replace(obs=obs, reward=reward, terminated=terminated) - done = state.terminated | state.truncated - if np.any(done): - done_indices = np.where(done)[0] - stats = self._spawn.update_on_done( - done_indices, self._backend.get_base_pos()[done_indices] - ) - if stats: - if "log" not in state.info: - state.info["log"] = {} - for k, v in stats.items(): - state.info["log"][f"terrain_curriculum/{k}"] = float(v) - return state - - def _compute_obs( - self, info: dict, linvel, gyro, gravity, dof_pos, dof_vel, feet_phase - ) -> dict[str, np.ndarray]: - noise_cfg = self._cfg.noise_config - diff = dof_pos - self.default_angles - noisy_gyro = self._obs_noise(gyro, noise_cfg.scale_gyro) - noisy_gravity = self._obs_noise(gravity, noise_cfg.scale_gravity) - noisy_diff = self._obs_noise(diff, noise_cfg.scale_joint_angle) - noisy_dof_vel = self._obs_noise(dof_vel, noise_cfg.scale_joint_vel) - command = info["commands"] - last_actions = info.get("current_actions", np.zeros_like(diff)) - obs = np.concatenate( - [ - noisy_gyro, - -noisy_gravity, - noisy_diff, - noisy_dof_vel, - last_actions, - command, - feet_phase, - ], - axis=1, - dtype=get_global_dtype(), - ) - critic = np.concatenate( - [gyro, -gravity, diff, dof_vel, last_actions, command, feet_phase, linvel], - axis=1, - dtype=get_global_dtype(), - ) - return {"obs": obs, "critic": critic} - - def _compute_reward(self, info: dict, linvel, gyro, dof_pos) -> np.ndarray: - cfg = self._reward_cfg - ctx = RewardContext( - info=info, - linvel=linvel, - gyro=gyro, - dof_pos=dof_pos, - num_envs=self._num_envs, - default_angles=self.default_angles, - tracking_sigma=cfg.tracking_sigma, - base_height_target=cfg.base_height_target, - base_height=self._reward_base_height_values(), - ) - return rewards.run_reward_dispatch( - scales=cfg.scales, - fns=self._reward_fns, - ctx=ctx, - info=info, - enable_log=self._enable_reward_log, - ctrl_dt=self._cfg.ctrl_dt, - ) - - # ── reward functions (robot-specific) ──────────────────────────── - - def _reward_base_height_values(self) -> np.ndarray: - base_pos = np.asarray(self._backend.get_base_pos(), dtype=get_global_dtype()) - sample_height = self._terrain_surface_sample_height - if sample_height is None: - return np.asarray(base_pos[:, 2], dtype=get_global_dtype()) - - surface = np.asarray(sample_height(base_pos[:, :2]), dtype=get_global_dtype()) - return np.asarray(base_pos[:, 2] - surface, dtype=get_global_dtype()) - - def _reward_swing_feet_z(self, ctx: RewardContext) -> np.ndarray: - is_swing = self.feet_phase >= 0.6 - target_height = 0.1 - height_error = np.square(self.feet_pos[:, :, 2] - target_height) - swing_rew = np.exp(-height_error / 0.01) * is_swing - reward: np.ndarray = np.sum(swing_rew, axis=1) / len(self._cfg.sensor.feet_pos) - return reward - - def _reward_foot_drag(self, ctx: RewardContext) -> np.ndarray: - foot_pos = self.get_foot_pos() - foot_heights = foot_pos[..., 2] - foot_contact = self.get_foot_contact() - is_swing = foot_contact < 0.5 - safe_height = self._reward_cfg.target_foot_height / 2.0 - height_error = np.clip(safe_height - foot_heights, 0.0, None) - error = np.square(height_error) * is_swing - drag_penalty: np.ndarray = np.sum(error, axis=1) - return drag_penalty - - def _reward_contact(self, ctx: RewardContext) -> np.ndarray: - contact = self.feet_force[:, :, 2] > 0.1 - res = np.zeros(self._num_envs, dtype=np.float32) - for i in range(len(self._cfg.sensor.feet_force)): - is_contact = (self.feet_phase[:, i] < 0.6) | (self.gait_frequency < 1.0e-8) - res += (contact[:, i] == is_contact).astype(np.float32) - return res / len(self._cfg.sensor.feet_force) - - -# Go2JoystickCfg and Go2WalkTask remain as implementation bases for legacy rough -# and A2 tasks. The flat production identity is Hydra-owned and uses only the -# generic Manager-Based config/runtime factories. -registry.register_env_config("Go2JoystickFlat", ManagerBasedRlEnvCfg) -registry.register_env("Go2JoystickFlat", make_manager_based_rl_env, sim_backend="mujoco") -registry.register_env("Go2JoystickFlat", make_manager_based_rl_env, sim_backend="motrix") -registry.register_env("Go2JoystickFlat", make_manager_based_rl_env, sim_backend="drake") diff --git a/src/unilab/tasks/locomotion/go2/rough.py b/src/unilab/tasks/locomotion/go2/rough.py deleted file mode 100644 index d8a2dedf7..000000000 --- a/src/unilab/tasks/locomotion/go2/rough.py +++ /dev/null @@ -1,793 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass, field -from typing import Any - -import numpy as np - -from unilab.assets import ASSETS_ROOT_PATH -from unilab.base import registry -from unilab.base.np_env import NpEnvState -from unilab.base.scene import SceneCfg, TerrainSceneCfg -from unilab.dr import DomainRandomizationManager, ResetPlan -from unilab.dr.dr_utils import build_common_reset_randomization, zero_actions -from unilab.dtype_config import get_global_dtype -from unilab.tasks.locomotion.common import rewards -from unilab.tasks.locomotion.common.commands import ( - apply_heading_yaw_feedback, - sample_heading_commands, - zero_small_xy_commands, -) -from unilab.tasks.locomotion.common.height_scan import ( - DEFAULT_SCAN_POINTS_X, - DEFAULT_SCAN_POINTS_Y, - HeightScanConfig, - base_height_from_scan, - height_scan_obs, - init_height_scan_sensor, - raw_height_scan_obs, - terrain_out_of_bounds, -) -from unilab.tasks.locomotion.common.rewards import RewardContext -from unilab.tasks.locomotion.go2.base import ControlConfig -from unilab.tasks.locomotion.go2.joystick import ( - Commands, - Go2JoystickCfg, - Go2JoystickDomainRandomizationProvider, - Go2WalkTask, - JoystickSensor, - RewardConfig, -) -from unilab.terrains import ( - SubTerrainCfg, - TerrainGeneratorCfg, - flat, - hf_pyramid_slope, - hf_pyramid_slope_inv, - pyramid_stairs, - pyramid_stairs_inv, - random_rough, - wave_terrain, -) -from unilab.utils.rotation import ( - np_quat_apply_inverse, - np_quat_from_euler_xyz, - np_quat_mul, -) - -# pyright: reportIncompatibleVariableOverride=false, reportAttributeAccessIssue=false, reportCallIssue=false - -GO2_HIP_INDICES = np.asarray([0, 3, 6, 9], dtype=np.int32) -GO2_ACTUATOR_TO_DOF_INDICES = np.asarray([3, 4, 5, 0, 1, 2, 9, 10, 11, 6, 7, 8], dtype=np.int32) -GO2_FRONT_LEFT = 0 -GO2_FRONT_RIGHT = 1 -GO2_REAR_LEFT = 2 -GO2_REAR_RIGHT = 3 - - -@dataclass -class RoughControlConfig(ControlConfig): - hip_action_scale: float = 0.125 - non_hip_action_scale: float = 0.25 - clip_actions: float = 100.0 - - -@dataclass -class RoughCommands(Commands): - vel_limit: list[list[float]] = field( - default_factory=lambda: [[-1.0, -1.0, -1.0], [1.0, 1.0, 1.0]] - ) - resampling_time: float = 10.0 - heading_command: bool = True - heading_range: list[float] = field(default_factory=lambda: [-np.pi, np.pi]) - - -@dataclass -class RoughRewardConfig(RewardConfig): - stand_still_command_threshold: float = 0.1 - joint_pos_penalty_stand_still_scale: float = 5.0 - joint_pos_penalty_velocity_threshold: float = 0.5 - joint_pos_penalty_command_threshold: float = 0.1 - contact_threshold: float = 1.0 - contact_forces_threshold: float = 100.0 - feet_air_time_threshold: float = 0.5 - feet_height_body_target: float = -0.2 - feet_height_body_tanh_mult: float = 2.0 - feet_gait_std: float = np.sqrt(0.5) - feet_gait_max_err: float = 0.2 - feet_gait_velocity_threshold: float = 0.5 - feet_gait_command_threshold: float = 0.1 - - -@dataclass -class RoughJoystickSensor(JoystickSensor): - feet_vel = ["FL_vel", "FR_vel", "RL_vel", "RR_vel"] - undesired_contact = [ - "base1_contact", - "base2_contact", - "base3_contact", - "FL_hip_contact", - "FR_hip_contact", - "RL_hip_contact", - "RR_hip_contact", - "FL_thigh_contact", - "FR_thigh_contact", - "RL_thigh_contact", - "RR_thigh_contact", - "FL_calf_contact1", - "FR_calf_contact1", - "RL_calf_contact1", - "RR_calf_contact1", - "FL_calf_contact2", - "FR_calf_contact2", - "RL_calf_contact2", - "RR_calf_contact2", - ] - - -@dataclass -class RoughTerminationConfig: - terrain_out_of_bounds: bool = True - terrain_distance_buffer: float = 3.0 - - -@dataclass(kw_only=True) -class Go2RoughTerrainCfg(TerrainGeneratorCfg): - size: tuple[float, float] = (8.0, 8.0) - num_rows: int = 6 - num_cols: int = 6 - border_width: float = 1.0 - add_lights: bool = True - horizontal_scale: float = 0.2 - - sub_terrains: dict[str, SubTerrainCfg] = field( - default_factory=lambda: { - "flat": flat(proportion=0.0), - "pyramid_stairs": pyramid_stairs( - proportion=0.1, - step_height_range=(0.025, 0.10), - step_width=0.4, - platform_width=3.0, - border_width=0.2, - ), - "pyramid_stairs_inv": pyramid_stairs_inv( - proportion=0.1, - step_height_range=(0.025, 0.10), - step_width=0.4, - platform_width=3.0, - border_width=0.2, - ), - "hf_pyramid_slope": hf_pyramid_slope( - proportion=0.2, - slope_range=(0.0, 0.3), - platform_width=2.0, - border_width=0.2, - ), - "hf_pyramid_slope_inv": hf_pyramid_slope_inv( - proportion=0.2, - slope_range=(0.0, 0.3), - platform_width=2.0, - border_width=0.2, - ), - "random_rough": random_rough( - proportion=0.3, - noise_range=(0.01, 0.06), - noise_step=0.01, - border_width=0.2, - ), - "wave_terrain": wave_terrain( - proportion=0.3, - amplitude_range=(0.0, 0.12), - num_waves=4, - border_width=0.2, - ), - } - ) - - -@registry.envcfg("Go2JoystickRough") -@dataclass -class Go2JoystickRoughCfg(Go2JoystickCfg): - scene: SceneCfg = field( - default_factory=lambda: SceneCfg( - model_file=str(ASSETS_ROOT_PATH / "robots" / "go2" / "go2.xml"), - fragment_files=[ - str(ASSETS_ROOT_PATH / "robots" / "go2" / "locomotion_task.xml"), - ], - terrain=TerrainSceneCfg( - generator=Go2RoughTerrainCfg(), - hfield_name="terrain_hfield", - geom_name="floor", - ), - ) - ) - control_config: RoughControlConfig = field(default_factory=RoughControlConfig) - commands: RoughCommands = field(default_factory=RoughCommands) - terrain_scan: HeightScanConfig = field(default_factory=HeightScanConfig) - termination_config: RoughTerminationConfig = field(default_factory=RoughTerminationConfig) - sensor: RoughJoystickSensor = field(default_factory=RoughJoystickSensor) - reward_config: RoughRewardConfig | None = None - - -class Go2JoystickRoughDomainRandomizationProvider(Go2JoystickDomainRandomizationProvider): - def _sample_commands(self, env: Any, num_reset: int) -> np.ndarray: - commands = super()._sample_commands(env, num_reset) - zero_small_xy_commands(commands, threshold=0.08) - standing_prob = env.cfg.commands.rel_standing_envs - if standing_prob > 0.0: - standing = np.random.uniform(size=(num_reset,)) < min(standing_prob, 1.0) - commands[standing] = 0.0 - if env.cfg.commands.heading_command: - commands[:, 2] = 0.0 - return commands - - 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)) - qpos[:, 2] += np.random.uniform(0.25, 0.5, (num_reset,)) - qpos[:, 0:3] += env._spawn.origins_for(env_ids) - roll = np.random.uniform(-3.14, 3.14, (num_reset,)) - pitch = np.random.uniform(-3.14, 3.14, (num_reset,)) - yaw = np.random.uniform(-3.14, 3.14, (num_reset,)) - qpos[:, 3:7] = np_quat_mul(qpos[:, 3:7], np_quat_from_euler_xyz(roll, pitch, yaw)) - qvel[:, 0:6] = np.asarray( - np.random.uniform(-0.5, 0.5, size=(num_reset, 6)), dtype=get_global_dtype() - ) - commands = self._sample_commands(env, num_reset) - info_updates: dict[str, Any] = { - "commands": commands, - "current_actions": zero_actions(num_reset, env._num_action), - "last_actions": zero_actions(num_reset, env._num_action), - "qacc": np.zeros((num_reset, env._num_action), dtype=get_global_dtype()), - "torques": np.zeros((num_reset, env._num_action), dtype=get_global_dtype()), - } - if env.cfg.commands.heading_command: - info_updates["heading_commands"] = sample_heading_commands(env, num_reset) - env._spawn.record_episode_start(env_ids, qpos[:, 0:3]) - return ResetPlan( - env_ids=env_ids, - qpos=qpos, - qvel=qvel, - info_updates=info_updates, - randomization=build_common_reset_randomization(env, num_reset), - ) - - -@registry.env("Go2JoystickRough", sim_backend="mujoco") -class Go2JoystickRoughEnv(Go2WalkTask): - _cfg: Go2JoystickRoughCfg - _reward_cfg: RoughRewardConfig - - def __init__(self, cfg: Go2JoystickRoughCfg, num_envs=1, backend_type="mujoco"): - self._height_scan_dim = len(cfg.terrain_scan.measured_points_x) * len( - cfg.terrain_scan.measured_points_y - ) - super().__init__(cfg, num_envs=num_envs, backend_type=backend_type) - self._dr_manager = DomainRandomizationManager( - self, Go2JoystickRoughDomainRandomizationProvider() - ) - self._last_dof_vel_for_acc = np.zeros( - (num_envs, self._num_action), dtype=get_global_dtype() - ) - self._action_scale = np.full( - (self._num_action,), - float(cfg.control_config.non_hip_action_scale), - dtype=get_global_dtype(), - ) - self._action_scale[GO2_HIP_INDICES] = float(cfg.control_config.hip_action_scale) - self._default_angles_actuator = self.default_angles[GO2_ACTUATOR_TO_DOF_INDICES] - joint_range = self._backend.get_joint_range() - self._joint_range = ( - np.asarray(joint_range, dtype=get_global_dtype()) if joint_range is not None else None - ) - self.feet_vel = np.zeros((num_envs, len(cfg.sensor.feet_vel), 3), dtype=np.float32) - self._last_foot_contact = np.zeros((num_envs, len(cfg.sensor.feet_force)), dtype=bool) - self._current_air_time = np.zeros((num_envs, len(cfg.sensor.feet_force)), dtype=np.float32) - self._current_contact_time = np.zeros( - (num_envs, len(cfg.sensor.feet_force)), dtype=np.float32 - ) - self._last_air_time = np.zeros((num_envs, len(cfg.sensor.feet_force)), dtype=np.float32) - self._last_contact_time = np.zeros((num_envs, len(cfg.sensor.feet_force)), dtype=np.float32) - self._first_foot_contact = np.zeros((num_envs, len(cfg.sensor.feet_force)), dtype=bool) - init_height_scan_sensor(self, cfg.terrain_scan, cfg.asset.base_name) - - @property - def obs_groups_spec(self) -> dict[str, int]: - return {"obs": 45, "critic": 48 + self._height_scan_dim} - - def reset(self, env_indices: np.ndarray) -> tuple[dict[str, np.ndarray], dict]: - env_ids = np.asarray(env_indices, dtype=np.int32) - obs, info = super().reset(env_ids) - dof_vel = self.get_dof_vel() - if dof_vel.shape[0] == self._num_envs: - self._last_dof_vel_for_acc[env_ids] = dof_vel[env_ids] - self._reset_contact_timers(env_ids) - return obs, info - - def _init_reward_functions(self): - scale_gravity = self._upright_scale # local alias for lambda capture - - def gated(fn): - return lambda ctx: fn(ctx) * scale_gravity(ctx.gravity) - - # joint_pos_penalty needs its three thresholds from the reward config - def _joint_pos_penalty(ctx: RewardContext) -> np.ndarray: - cfg = self._reward_cfg - return rewards.joint_pos_penalty( - ctx, - stand_still_scale=cfg.joint_pos_penalty_stand_still_scale, - velocity_threshold=cfg.joint_pos_penalty_velocity_threshold, - command_threshold=cfg.joint_pos_penalty_command_threshold, - ) * scale_gravity(ctx.gravity) - - def _stand_still(ctx: RewardContext) -> np.ndarray: - return rewards.stand_still( - ctx, command_threshold=self._reward_cfg.stand_still_command_threshold - ) * scale_gravity(ctx.gravity) - - self._reward_fns: dict[str, Any] = { - "tracking_lin_vel": gated(rewards.tracking_lin_vel), - "tracking_ang_vel": gated(rewards.tracking_ang_vel), - "lin_vel_z": gated(rewards.lin_vel_z), - "ang_vel_xy": gated(rewards.ang_vel_xy), - "dof_torques_l2": gated(rewards.dof_torques_l2), - "joint_torques_l2": gated(rewards.dof_torques_l2), - "dof_acc_l2": gated(rewards.dof_acc_l2), - "joint_acc_l2": gated(rewards.dof_acc_l2), - "joint_power": gated(rewards.joint_power), - "stand_still": _stand_still, - "hip_pos": self._reward_hip_pos, - "joint_pos_penalty": _joint_pos_penalty, - "joint_mirror": self._reward_joint_mirror, - "action_rate": rewards.action_rate, - "action_rate_l2": rewards.action_rate, - "undesired_contacts": self._reward_undesired_contacts, - "contact_forces": self._reward_contact_forces, - "feet_air_time": self._reward_feet_air_time, - "feet_air_time_variance": self._reward_feet_air_time_variance, - "feet_contact_without_cmd": self._reward_feet_contact_without_cmd, - "feet_slide": self._reward_feet_slide, - "feet_height_body": self._reward_feet_height_body, - "feet_gait": self._reward_feet_gait, - "upward": rewards.upward, - } - - def apply_action(self, actions: np.ndarray, state: NpEnvState) -> np.ndarray: - clipped_actions = np.asarray( - np.clip( - actions, - -float(self._cfg.control_config.clip_actions), - float(self._cfg.control_config.clip_actions), - ), - dtype=get_global_dtype(), - ) - state.info["last_actions"] = state.info.get( - "current_actions", np.zeros_like(clipped_actions) - ) - state.info["current_actions"] = clipped_actions - exec_actions = ( - state.info["last_actions"] - if self._cfg.control_config.simulate_action_latency - else clipped_actions - ) - return np.asarray( - exec_actions * self._action_scale + self._default_angles_actuator, - dtype=get_global_dtype(), - ) - - def update_state(self, state: NpEnvState) -> NpEnvState: - self._update_commands(state.info) - self.phase = np.fmod(self.phase + self._cfg.ctrl_dt * self.gait_frequency, 1.0) - self.feet_phase[:, 0] = self.phase - self.feet_phase[:, 3] = self.phase - self.feet_phase[:, 1] = (self.phase + 0.5) % 1 - self.feet_phase[:, 2] = (self.phase + 0.5) % 1 - - linvel = self.get_local_linvel() - gyro = self.get_gyro() - gravity = self._backend.get_sensor_data("upvector") - dof_pos = self.get_dof_pos() - dof_vel = self.get_dof_vel() - self.feet_force[:, :, :] = 0 - for i in range(len(self._cfg.sensor.feet_force)): - self.feet_force[:, i, :] = self._backend.get_sensor_data(self._cfg.sensor.feet_force[i]) - for i in range(len(self._cfg.sensor.feet_pos)): - self.feet_pos[:, i, :] = self._backend.get_sensor_data(self._cfg.sensor.feet_pos[i]) - for i in range(len(self._cfg.sensor.feet_vel)): - self.feet_vel[:, i, :] = self._backend.get_sensor_data(self._cfg.sensor.feet_vel[i]) - self._update_contact_timers(self._foot_contact_mask()) - state.info["qacc"] = self._estimate_dof_acc(dof_vel) - state.info["torques"] = self._estimate_pd_torques(state.info, dof_pos, dof_vel) - terminated = self._compute_terminated(gravity) - reward = self._compute_rough_reward(state.info, linvel, gyro, gravity, dof_pos, dof_vel) - obs = self._compute_obs( - state.info, linvel, gyro, gravity, dof_pos, dof_vel, self.feet_phase - ) - state = state.replace(obs=obs, reward=reward, terminated=terminated) - done = state.terminated | state.truncated - if np.any(done): - done_indices = np.where(done)[0] - stats = self._spawn.update_on_done( - done_indices, self._backend.get_base_pos()[done_indices] - ) - if stats: - if "log" not in state.info: - state.info["log"] = {} - for k, v in stats.items(): - state.info["log"][f"terrain_curriculum/{k}"] = float(v) - return state - - def _compute_obs( - self, - info: dict, - linvel: np.ndarray, - gyro: np.ndarray, - gravity: np.ndarray, - dof_pos: np.ndarray, - dof_vel: np.ndarray, - feet_phase: np.ndarray, - ) -> dict[str, np.ndarray]: - del feet_phase - noise_cfg = self._cfg.noise_config - diff = dof_pos - self.default_angles - policy_gyro = self._obs_noise(gyro, noise_cfg.scale_gyro) * 0.25 - policy_gravity = self._obs_noise(-gravity, noise_cfg.scale_gravity) - policy_diff = self._obs_noise(diff, noise_cfg.scale_joint_angle) - policy_dof_vel = self._obs_noise(dof_vel, noise_cfg.scale_joint_vel) * 0.05 - last_actions = info.get("current_actions", np.zeros_like(diff)) - commands = info["commands"] - obs = np.concatenate( - [policy_gyro, policy_gravity, commands, policy_diff, policy_dof_vel, last_actions], - axis=1, - dtype=get_global_dtype(), - ) - critic_base = np.concatenate( - [linvel, gyro, -gravity, commands, diff, dof_vel, last_actions], - axis=1, - dtype=get_global_dtype(), - ) - critic = np.concatenate( - [critic_base, height_scan_obs(self, self._cfg.terrain_scan, critic_base.shape[0])], - axis=1, - dtype=get_global_dtype(), - ) - return {"obs": obs, "critic": critic} - - def _compute_rough_reward( - self, - info: dict, - linvel: np.ndarray, - gyro: np.ndarray, - gravity: np.ndarray, - dof_pos: np.ndarray, - dof_vel: np.ndarray, - ) -> np.ndarray: - cfg = self._reward_cfg - ctx = RewardContext( - info=info, - linvel=linvel, - gyro=gyro, - dof_pos=dof_pos, - num_envs=self._num_envs, - default_angles=self.default_angles, - tracking_sigma=cfg.tracking_sigma, - base_height_target=cfg.base_height_target, - base_height=base_height_from_scan(self, self._num_envs), - gravity=gravity, - dof_vel=dof_vel, - joint_range=self._joint_range, - ) - return rewards.run_reward_dispatch( - scales=cfg.scales, - fns=self._reward_fns, - ctx=ctx, - info=info, - enable_log=self._enable_reward_log, - ctrl_dt=self._cfg.ctrl_dt, - ) - - def _compute_terminated(self, gravity: np.ndarray) -> np.ndarray: - # return gravity[:, 2] <= 0.5 - del gravity - return np.zeros((self._num_envs,), dtype=bool) - - def _compute_truncated(self, state: NpEnvState) -> np.ndarray: - truncated = super()._compute_truncated(state) - if self._cfg.termination_config.terrain_out_of_bounds: - terrain_scene = self._cfg.scene.terrain - terrain_cfg = terrain_scene.generator if terrain_scene is not None else None - np.logical_or( - truncated, - terrain_out_of_bounds( - self, - terrain_cfg, - float(self._cfg.termination_config.terrain_distance_buffer), - ), - out=truncated, - ) - return truncated - - def _reward_base_height_values(self, num_obs: int | None = None) -> np.ndarray: - return base_height_from_scan(self, num_obs) - - def _raw_height_scan_obs(self, num_obs: int) -> tuple[np.ndarray | None, np.ndarray | None]: - return raw_height_scan_obs(self, num_obs) - - def _estimate_dof_acc(self, dof_vel: np.ndarray) -> np.ndarray: - qacc = np.asarray((dof_vel - self._last_dof_vel_for_acc) / self._cfg.ctrl_dt) - self._last_dof_vel_for_acc[:] = dof_vel - return np.asarray(qacc, dtype=get_global_dtype()) - - def _estimate_pd_torques( - self, info: dict, dof_pos: np.ndarray, dof_vel: np.ndarray - ) -> np.ndarray: - actions = np.asarray( - info.get("current_actions", np.zeros((dof_pos.shape[0], self._num_action))), - dtype=get_global_dtype(), - ) - if self._cfg.control_config.simulate_action_latency: - actions = np.asarray(info.get("last_actions", actions), dtype=get_global_dtype()) - targets_actuator = actions * self._action_scale + self._default_angles_actuator - targets_dof = np.empty_like(targets_actuator) - targets_dof[:, GO2_ACTUATOR_TO_DOF_INDICES] = targets_actuator - torques = ( - float(self._cfg.control_config.Kp) * (targets_dof - dof_pos) - - float(self._cfg.control_config.Kd) * dof_vel - ) - return np.asarray(torques, dtype=get_global_dtype()) - - def _update_commands(self, info: dict) -> None: - commands_arr = np.asarray(info["commands"], dtype=get_global_dtype()) - resampling_time = float(self._cfg.commands.resampling_time) - if resampling_time > 0.0: - interval_steps = max(int(round(resampling_time / self._cfg.ctrl_dt)), 1) - steps = np.asarray(info.get("steps", np.zeros((self._num_envs,), dtype=np.uint32))) - resample_mask = (steps > 0) & ((steps % interval_steps) == 0) - if np.any(resample_mask): - num_resample = int(np.count_nonzero(resample_mask)) - low = np.asarray(self._cfg.commands.vel_limit[0], dtype=get_global_dtype()) - high = np.asarray(self._cfg.commands.vel_limit[1], dtype=get_global_dtype()) - sampled = np.random.uniform(low=low, high=high, size=(num_resample, 3)).astype( - get_global_dtype() - ) - zero_small_xy_commands(sampled, threshold=0.08) - commands_arr[resample_mask] = sampled - if self._cfg.commands.heading_command: - heading_commands = self._ensure_heading_commands(info, commands_arr.shape[0]) - heading_commands[resample_mask] = sample_heading_commands(self, num_resample) - info["heading_commands"] = heading_commands - - if self._cfg.commands.heading_command: - heading_commands = self._ensure_heading_commands(info, commands_arr.shape[0]) - base_quat = np.asarray(self._backend.get_base_quat(), dtype=get_global_dtype()) - if base_quat.shape[0] == commands_arr.shape[0]: - apply_heading_yaw_feedback(commands_arr, base_quat, heading_commands, stiffness=0.5) - info["commands"] = commands_arr - - def _ensure_heading_commands(self, info: dict, num_obs: int) -> np.ndarray: - heading_commands = info.get("heading_commands") - if heading_commands is None or np.asarray(heading_commands).shape != (num_obs,): - heading_commands = sample_heading_commands(self, num_obs) - heading_commands = np.asarray(heading_commands, dtype=get_global_dtype()) - info["heading_commands"] = heading_commands - return heading_commands - - def _foot_contact_mask(self) -> np.ndarray: - contact_force = np.linalg.norm(self.feet_force, axis=2) - return np.asarray(contact_force > self._reward_cfg.contact_threshold, dtype=bool) - - def _reset_contact_timers(self, env_ids: np.ndarray) -> None: - self._current_air_time[env_ids] = 0.0 - self._current_contact_time[env_ids] = 0.0 - self._last_air_time[env_ids] = 0.0 - self._last_contact_time[env_ids] = 0.0 - self._first_foot_contact[env_ids] = False - self._last_foot_contact[env_ids] = self._foot_contact_mask()[env_ids] - - def _update_contact_timers(self, contact: np.ndarray) -> None: - first_contact = contact & ~self._last_foot_contact - first_air = ~contact & self._last_foot_contact - self._first_foot_contact[:] = first_contact - self._last_air_time[first_contact] = self._current_air_time[first_contact] - self._last_contact_time[first_air] = self._current_contact_time[first_air] - self._current_air_time[contact] = 0.0 - self._current_air_time[~contact] += self._cfg.ctrl_dt - self._current_contact_time[~contact] = 0.0 - self._current_contact_time[contact] += self._cfg.ctrl_dt - self._last_foot_contact[:] = contact - - def _upright_scale(self, gravity: np.ndarray | None) -> np.ndarray: - return rewards.upright_scale(gravity, self._num_envs) - - # ── reward functions that need backend / env state (kept as methods) ──── - - def _reward_hip_pos(self, ctx: RewardContext) -> np.ndarray: - diff = ctx.dof_pos[:, GO2_HIP_INDICES] - self.default_angles[GO2_HIP_INDICES] - return np.asarray( - np.sum(np.square(diff), axis=1) * self._upright_scale(ctx.gravity), - dtype=get_global_dtype(), - ) - - def _reward_joint_mirror(self, ctx: RewardContext) -> np.ndarray: - fr_rl = ctx.dof_pos[:, 0:3] - ctx.dof_pos[:, 9:12] - fl_rr = ctx.dof_pos[:, 3:6] - ctx.dof_pos[:, 6:9] - mirror = 0.5 * (np.sum(np.square(fr_rl), axis=1) + np.sum(np.square(fl_rr), axis=1)) - return np.asarray(mirror * self._upright_scale(ctx.gravity), dtype=get_global_dtype()) - - def _reward_undesired_contacts(self, ctx: RewardContext) -> np.ndarray: - contacts = [ - _force_norm_columns( - np.asarray(self._backend.get_sensor_data(name), dtype=get_global_dtype()), - ctx.num_envs, - ) - for name in self._cfg.sensor.undesired_contact - ] - if not contacts: - return np.zeros((ctx.num_envs,), dtype=get_global_dtype()) - contact_force = np.concatenate(contacts, axis=1) - contact_count = np.sum(contact_force > self._reward_cfg.contact_threshold, axis=1) - return np.asarray( - contact_count * self._upright_scale(ctx.gravity), dtype=get_global_dtype() - ) - - def _reward_contact_forces(self, ctx: RewardContext) -> np.ndarray: - force_norm = np.linalg.norm(self.feet_force, axis=2) - force_clip = np.minimum(force_norm, 1500.0) - violation = np.clip(force_clip - self._reward_cfg.contact_forces_threshold, 0.0, None) - return np.asarray( - np.sum(violation, axis=1) * self._upright_scale(ctx.gravity), - dtype=get_global_dtype(), - ) - - def _reward_feet_air_time(self, ctx: RewardContext) -> np.ndarray: - cfg = self._reward_cfg - reward = np.sum( - (self._last_air_time - cfg.feet_air_time_threshold) * self._first_foot_contact, - axis=1, - ) - moving = np.linalg.norm(ctx.info["commands"], axis=1) > 0.1 - return np.asarray( - reward * moving * self._upright_scale(ctx.gravity), dtype=get_global_dtype() - ) - - def _reward_feet_air_time_variance(self, ctx: RewardContext) -> np.ndarray: - air_var = np.var(np.clip(self._last_air_time, 0.0, 0.5), axis=1) - contact_var = np.var(np.clip(self._last_contact_time, 0.0, 0.5), axis=1) - return np.asarray( - (air_var + contact_var) * self._upright_scale(ctx.gravity), dtype=get_global_dtype() - ) - - def _reward_feet_contact_without_cmd(self, ctx: RewardContext) -> np.ndarray: - reward = np.sum(self._first_foot_contact, axis=1) - stopped = np.linalg.norm(ctx.info["commands"], axis=1) < 0.1 - return np.asarray( - reward * stopped * self._upright_scale(ctx.gravity), dtype=get_global_dtype() - ) - - def _relative_foot_vel_body(self) -> np.ndarray: - base_quat = np.asarray(self._backend.get_base_quat(), dtype=get_global_dtype()) - base_linvel = np.asarray( - self._backend.get_sensor_data("global_linvel"), dtype=get_global_dtype() - ) - relative_vel = self.feet_vel - base_linvel[:, None, :] - flat = relative_vel.reshape(self._num_envs * relative_vel.shape[1], 3) - quat = np.repeat(base_quat, relative_vel.shape[1], axis=0) - return np_quat_apply_inverse(quat, flat).reshape(relative_vel.shape) - - def _relative_foot_pos_body(self) -> np.ndarray: - base_quat = np.asarray(self._backend.get_base_quat(), dtype=get_global_dtype()) - base_pos = np.asarray(self._backend.get_base_pos(), dtype=get_global_dtype()) - relative_pos = self.feet_pos - base_pos[:, None, :] - flat = relative_pos.reshape(self._num_envs * relative_pos.shape[1], 3) - quat = np.repeat(base_quat, relative_pos.shape[1], axis=0) - return np_quat_apply_inverse(quat, flat).reshape(relative_pos.shape) - - def _reward_feet_slide(self, ctx: RewardContext) -> np.ndarray: - foot_vel_body = self._relative_foot_vel_body() - lateral_vel = np.linalg.norm(foot_vel_body[:, :, :2], axis=2) - reward = np.sum(lateral_vel * self._foot_contact_mask(), axis=1) - return np.asarray(reward * self._upright_scale(ctx.gravity), dtype=get_global_dtype()) - - def _reward_feet_height_body(self, ctx: RewardContext) -> np.ndarray: - cfg = self._reward_cfg - foot_pos_body = self._relative_foot_pos_body() - foot_vel_body = self._relative_foot_vel_body() - z_error = np.square(foot_pos_body[:, :, 2] - cfg.feet_height_body_target) - velocity_tanh = np.tanh( - cfg.feet_height_body_tanh_mult * np.linalg.norm(foot_vel_body[:, :, :2], axis=2) - ) - moving = np.linalg.norm(ctx.info["commands"], axis=1) > 0.1 - reward = np.sum(z_error * velocity_tanh, axis=1) - return np.asarray( - reward * moving * self._upright_scale(ctx.gravity), dtype=get_global_dtype() - ) - - def _reward_feet_gait(self, ctx: RewardContext) -> np.ndarray: - cfg = self._reward_cfg - command_norm = np.linalg.norm(ctx.info["commands"], axis=1) - body_vel = np.linalg.norm(ctx.linvel[:, :2], axis=1) - enabled = (command_norm > cfg.feet_gait_command_threshold) | ( - body_vel > cfg.feet_gait_velocity_threshold - ) - air = self._current_air_time - contact = self._current_contact_time - sync_fl_rr = _gait_sync_reward( - air, contact, GO2_FRONT_LEFT, GO2_REAR_RIGHT, cfg.feet_gait_std, cfg.feet_gait_max_err - ) - sync_fr_rl = _gait_sync_reward( - air, contact, GO2_FRONT_RIGHT, GO2_REAR_LEFT, cfg.feet_gait_std, cfg.feet_gait_max_err - ) - async_fl_fr = _gait_async_reward( - air, contact, GO2_FRONT_LEFT, GO2_FRONT_RIGHT, cfg.feet_gait_std, cfg.feet_gait_max_err - ) - async_rr_rl = _gait_async_reward( - air, contact, GO2_REAR_RIGHT, GO2_REAR_LEFT, cfg.feet_gait_std, cfg.feet_gait_max_err - ) - async_fl_rl = _gait_async_reward( - air, contact, GO2_FRONT_LEFT, GO2_REAR_LEFT, cfg.feet_gait_std, cfg.feet_gait_max_err - ) - async_fr_rr = _gait_async_reward( - air, contact, GO2_FRONT_RIGHT, GO2_REAR_RIGHT, cfg.feet_gait_std, cfg.feet_gait_max_err - ) - reward = sync_fl_rr * sync_fr_rl * async_fl_fr * async_rr_rl * async_fl_rl * async_fr_rr - return np.asarray( - reward * enabled * self._upright_scale(ctx.gravity), dtype=get_global_dtype() - ) - - -def _force_norm_columns(force: np.ndarray, num_envs: int) -> np.ndarray: - force = np.asarray(force, dtype=get_global_dtype()).reshape(num_envs, -1) - if force.shape[1] == 0: - return force - if force.shape[1] % 3 == 0: - return np.linalg.norm(force.reshape(num_envs, -1, 3), axis=2) - return np.abs(force) - - -def _gait_sync_reward( - air: np.ndarray, - contact: np.ndarray, - foot_0: int, - foot_1: int, - std: float, - max_err: float, -) -> np.ndarray: - se_air = np.clip(np.square(air[:, foot_0] - air[:, foot_1]), 0.0, max_err**2) - se_contact = np.clip(np.square(contact[:, foot_0] - contact[:, foot_1]), 0.0, max_err**2) - return np.exp(-(se_air + se_contact) / std) - - -def _gait_async_reward( - air: np.ndarray, - contact: np.ndarray, - foot_0: int, - foot_1: int, - std: float, - max_err: float, -) -> np.ndarray: - se_act_0 = np.clip(np.square(air[:, foot_0] - contact[:, foot_1]), 0.0, max_err**2) - se_act_1 = np.clip(np.square(contact[:, foot_0] - air[:, foot_1]), 0.0, max_err**2) - return np.exp(-(se_act_0 + se_act_1) / std) - - -# Backwards-compat aliases for any callers that imported the unused defaults. -__all__ = [ - "DEFAULT_SCAN_POINTS_X", - "DEFAULT_SCAN_POINTS_Y", - "GO2_ACTUATOR_TO_DOF_INDICES", - "GO2_HIP_INDICES", - "Go2JoystickRoughCfg", - "Go2JoystickRoughDomainRandomizationProvider", - "Go2JoystickRoughEnv", - "Go2RoughTerrainCfg", - "RoughCommands", - "RoughControlConfig", - "RoughJoystickSensor", - "RoughRewardConfig", - "RoughTerminationConfig", -] - - -registry.register_env("Go2JoystickRough", Go2JoystickRoughEnv, sim_backend="motrix") diff --git a/src/unilab/tasks/locomotion/go2w/__init__.py b/src/unilab/tasks/locomotion/go2w/__init__.py index d82303136..2c908918d 100644 --- a/src/unilab/tasks/locomotion/go2w/__init__.py +++ b/src/unilab/tasks/locomotion/go2w/__init__.py @@ -1,9 +1,15 @@ -from .joystick import Go2WJoystickCfg, Go2WJoystickEnv -from .rough import Go2WJoystickRoughCfg, Go2WJoystickRoughEnv - -__all__ = [ - "Go2WJoystickCfg", - "Go2WJoystickEnv", - "Go2WJoystickRoughCfg", - "Go2WJoystickRoughEnv", -] +"""Hydra-owned Manager-Based Go2W production registrations.""" + +from unilab.base import registry +from unilab.envs import ManagerBasedRlEnvCfg, make_manager_based_rl_env + +registry.register_env_config("Go2WJoystickFlat", ManagerBasedRlEnvCfg) +registry.register_env("Go2WJoystickFlat", make_manager_based_rl_env, sim_backend="mujoco") +registry.register_env("Go2WJoystickFlat", make_manager_based_rl_env, sim_backend="motrix") +registry.register_env("Go2WJoystickFlat", make_manager_based_rl_env, sim_backend="drake") + +registry.register_env_config("Go2WJoystickRough", ManagerBasedRlEnvCfg) +registry.register_env("Go2WJoystickRough", make_manager_based_rl_env, sim_backend="mujoco") +registry.register_env("Go2WJoystickRough", make_manager_based_rl_env, sim_backend="motrix") + +__all__: list[str] = [] diff --git a/src/unilab/tasks/locomotion/go2w/base.py b/src/unilab/tasks/locomotion/go2w/base.py index 7ef1f4def..d4912a879 100644 --- a/src/unilab/tasks/locomotion/go2w/base.py +++ b/src/unilab/tasks/locomotion/go2w/base.py @@ -1,17 +1,7 @@ from __future__ import annotations -from dataclasses import dataclass, field - -import gymnasium as gym import numpy as np -from unilab.tasks.locomotion.common.base import ( - BaseNoiseConfig, - LocomotionBaseCfg, - LocomotionBaseEnv, - PdControlConfig, -) - LEG_JOINT_SENSOR_PREFIXES: tuple[str, ...] = ( "FR_hip", "FR_thigh", @@ -33,62 +23,6 @@ NUM_WHEEL_ACTIONS = len(WHEEL_JOINT_SENSOR_PREFIXES) NUM_GO2W_ACTIONS = len(JOINT_SENSOR_PREFIXES) -DEFAULT_LEG_ANGLES = np.asarray( - [ - 0.0, - 0.8, - -1.5, - 0.0, - 0.8, - -1.5, - 0.0, - 0.8, - -1.5, - 0.0, - 0.8, - -1.5, - ], - dtype=np.float64, -) -DEFAULT_GO2W_ANGLES = np.concatenate( - [DEFAULT_LEG_ANGLES, np.zeros((NUM_WHEEL_ACTIONS,), dtype=np.float64)] -) - - -@dataclass -class NoiseConfig(BaseNoiseConfig): - scale_wheel_vel: float = 0.5 - - -@dataclass -class ControlConfig(PdControlConfig): - action_scale: float = 0.25 - hip_action_scale: float | None = None - wheel_action_scale: float = 10.0 - wheel_Kd: float = 0.5 # noqa: N815 - Hydra config key kept for compatibility. - clip_actions: float = 1.0 - - -@dataclass -class Asset: - base_name = "base_link" - ground = "floor" - - -@dataclass -class Go2WBaseCfg(LocomotionBaseCfg): - noise_config: NoiseConfig = field(default_factory=NoiseConfig) # type: ignore[assignment] - control_config: ControlConfig = field(default_factory=ControlConfig) # type: ignore[assignment] - asset: Asset = field(default_factory=Asset) - sim_dt: float = 0.005 - ctrl_dt: float = 0.02 - - -def stack_joint_sensors(backend, suffix: str, *, dtype: np.dtype | type) -> np.ndarray: - names = tuple(f"{prefix}_{suffix}" for prefix in JOINT_SENSOR_PREFIXES) - values = backend.get_sensor_data_batch(names) - return np.asarray(values.reshape(values.shape[0], -1)[:, :NUM_GO2W_ACTIONS], dtype=dtype) - def compute_go2w_motor_ctrl( policy_ctrl: np.ndarray, @@ -116,23 +50,12 @@ def compute_go2w_motor_ctrl( return out -class Go2WBaseEnv(LocomotionBaseEnv): - _cfg: Go2WBaseCfg # pyright: ignore[reportIncompatibleVariableOverride] - - def _init_action_space(self) -> None: - self._action_space = gym.spaces.Box( - low=-1.0, - high=1.0, - shape=(NUM_GO2W_ACTIONS,), - dtype=np.float32, - ) - - def _init_buffers(self) -> None: - super()._init_buffers() - self.default_angles = np.asarray(DEFAULT_GO2W_ANGLES, dtype=self.default_angles.dtype) - - def get_dof_pos(self) -> np.ndarray: - return stack_joint_sensors(self._backend, "pos", dtype=self.default_angles.dtype) - - def get_dof_vel(self) -> np.ndarray: - return stack_joint_sensors(self._backend, "vel", dtype=self.default_angles.dtype) +__all__ = [ + "JOINT_SENSOR_PREFIXES", + "LEG_JOINT_SENSOR_PREFIXES", + "NUM_GO2W_ACTIONS", + "NUM_LEG_ACTIONS", + "NUM_WHEEL_ACTIONS", + "WHEEL_JOINT_SENSOR_PREFIXES", + "compute_go2w_motor_ctrl", +] diff --git a/src/unilab/tasks/locomotion/go2w/joystick.py b/src/unilab/tasks/locomotion/go2w/joystick.py deleted file mode 100644 index 9a61a063f..000000000 --- a/src/unilab/tasks/locomotion/go2w/joystick.py +++ /dev/null @@ -1,651 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass, field -from typing import Any, cast - -import numpy as np - -from unilab.assets import ASSETS_ROOT_PATH -from unilab.base import registry -from unilab.base.backend import create_backend, env_backend_kwargs -from unilab.base.np_env import NpEnvState -from unilab.base.scene import SceneCfg -from unilab.dr import DomainRandomizationCapabilities, ResetPlan, ResetRandomizationPayload -from unilab.dr.dr_utils import ( - build_interval_push_plan, - validate_interval_push_support, - zero_actions, -) -from unilab.dtype_config import get_global_dtype -from unilab.envs import ManagerBasedRlEnvCfg, make_manager_based_rl_env -from unilab.tasks.locomotion.common import rewards -from unilab.tasks.locomotion.common.commands import ( - Commands, - apply_heading_yaw_feedback, - zero_small_xy_commands, -) -from unilab.tasks.locomotion.common.commands import ( - sample_heading_commands as sample_go2w_heading_commands, -) -from unilab.tasks.locomotion.common.domain_rand import DomainRandConfig -from unilab.tasks.locomotion.common.dr_provider import LocomotionDRProvider -from unilab.tasks.locomotion.common.rewards import RewardContext -from unilab.tasks.locomotion.go2w.base import ( - DEFAULT_GO2W_ANGLES, - NUM_GO2W_ACTIONS, - NUM_LEG_ACTIONS, - NUM_WHEEL_ACTIONS, - Go2WBaseCfg, - Go2WBaseEnv, - compute_go2w_motor_ctrl, - stack_joint_sensors, -) -from unilab.utils.rotation import ( - np_quat_mul, - np_yaw_to_quat, -) - -GO2W_HIP_INDICES = np.asarray([0, 3, 6, 9], dtype=np.int32) - - -@dataclass -class InitState: - pos = [0.0, 0.0, 0.42] - - -@dataclass -class Go2WDomainRandConfig(DomainRandConfig): - randomize_init_yaw: bool = True - init_z_range: list[float] = field(default_factory=lambda: [0.0, 0.2]) - init_roll_range: list[float] = field(default_factory=lambda: [0.0, 0.0]) - init_pitch_range: list[float] = field(default_factory=lambda: [0.0, 0.0]) - init_yaw_range: list[float] = field(default_factory=lambda: [-np.pi, np.pi]) - - randomize_kp: bool = True - kp_multiplier_range: list[float] = field(default_factory=lambda: [0.9, 1.1]) - - randomize_kd: bool = True - kd_multiplier_range: list[float] = field(default_factory=lambda: [0.9, 1.1]) - - -@dataclass -class RewardConfig: - scales: dict[str, float] - tracking_sigma: float - base_height_target: float - only_positive_rewards: bool = False - joint_pos_penalty_stand_still_scale: float = 5.0 - joint_pos_penalty_velocity_threshold: float = 0.5 - joint_pos_penalty_command_threshold: float = 0.1 - - -@dataclass -class JoystickSensor: - local_linvel = "local_linvel" - gyro = "gyro" - gravity = "upvector" - - -@dataclass -class Go2WJoystickCfg(Go2WBaseCfg): - scene: SceneCfg = field( # pyright: ignore[reportIncompatibleVariableOverride] - default_factory=lambda: SceneCfg( - model_file=str(ASSETS_ROOT_PATH / "robots" / "go2w" / "scene_flat.xml") - ) - ) - max_episode_seconds: float = 20.0 # pyright: ignore[reportIncompatibleVariableOverride] - init_state: InitState = field(default_factory=InitState) - commands: Commands = field(default_factory=Commands) - reward_config: RewardConfig | None = None - sensor: JoystickSensor = field(default_factory=JoystickSensor) # type: ignore[assignment] - domain_rand: Go2WDomainRandConfig = field(default_factory=Go2WDomainRandConfig) - - -def build_go2w_backend_reset_randomization( - env: Any, num_reset: int -) -> ResetRandomizationPayload | None: - """Build reset DR payloads that are valid for a motor-actuator Go2W model. - - kp/kd are intentionally excluded here. Go2W samples them through the same - config path as Go2, but applies them inside its owner pre-step motor control. - """ - domain_rand = getattr(env.cfg, "domain_rand", None) - if domain_rand is None: - return None - - payload = ResetRandomizationPayload() - if getattr(domain_rand, "randomize_base_mass", False): - low, high = domain_rand.added_mass_range - payload.base_mass_delta = np.random.uniform(low, high, size=(num_reset,)) - - if getattr(domain_rand, "random_com", False): - low, high = domain_rand.com_offset_x - base_com_offset = np.zeros((num_reset, 3), dtype=np.float64) - base_com_offset[:, 0] = np.random.uniform(low, high, size=(num_reset,)) - payload.base_com_offset = base_com_offset - - if getattr(domain_rand, "randomize_gravity", False): - gravity_range = np.asarray(domain_rand.gravity_range, dtype=np.float64) - if gravity_range.shape != (2, 3): - raise ValueError( - f"domain_rand.gravity_range must have shape (2, 3), got {gravity_range.shape}" - ) - low = np.minimum(gravity_range[0], gravity_range[1]) - high = np.maximum(gravity_range[0], gravity_range[1]) - payload.gravity = np.random.uniform(low=low, high=high, size=(num_reset, 3)) - - return None if payload.is_empty() else payload - - -def sample_go2w_reset_yaw(domain_rand: Go2WDomainRandConfig, num_reset: int) -> np.ndarray: - if not domain_rand.randomize_init_yaw: - return np.zeros((num_reset,), dtype=get_global_dtype()) - - yaw_range = np.asarray(domain_rand.init_yaw_range, dtype=np.float64) - if yaw_range.shape != (2,): - raise ValueError(f"domain_rand.init_yaw_range must have shape (2,), got {yaw_range.shape}") - low, high = float(np.min(yaw_range)), float(np.max(yaw_range)) - return np.asarray(np.random.uniform(low, high, size=(num_reset,)), dtype=get_global_dtype()) - - -class Go2WJoystickDomainRandomizationProvider(LocomotionDRProvider): - def validate(self, env: Any, capabilities: DomainRandomizationCapabilities) -> None: - payload = build_go2w_backend_reset_randomization(env, num_reset=1) - if payload is not None: - unsupported = capabilities.get_unsupported_reset_terms(payload.requested_terms()) - if unsupported: - names = ", ".join(sorted(unsupported)) - raise NotImplementedError( - f"{env._backend.backend_type} backend does not support Go2W reset randomization terms: {names}" - ) - validate_interval_push_support(env, capabilities) - - def build_interval_randomization_plan(self, env: Any, step_counter: int): - return build_interval_push_plan(env, step_counter) - - 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)) - qpos[:, 0:3] += env._spawn.origins_for(env_ids) - yaw = sample_go2w_reset_yaw(env.cfg.domain_rand, num_reset) - qpos[:, 3:7] = np_quat_mul(qpos[:, 3:7], np_yaw_to_quat(yaw)) - qvel[:, 0:6] = np.asarray( - np.random.uniform(-0.5, 0.5, size=(num_reset, 6)), dtype=get_global_dtype() - ) - - motor_kp, motor_kd = env.sample_reset_motor_gains(num_reset) - env.set_motor_gains(env_ids, motor_kp, motor_kd) - - commands = self._sample_commands(env, num_reset) - info_updates: dict[str, Any] = { - "commands": commands, - "current_actions": zero_actions(num_reset, env._num_action), - "last_actions": zero_actions(num_reset, env._num_action), - "motor_kp": motor_kp.astype(get_global_dtype()), - "motor_kd": motor_kd.astype(get_global_dtype()), - "torques": np.zeros((num_reset, env._num_action), dtype=get_global_dtype()), - } - if getattr(env.cfg.commands, "heading_command", False): - info_updates["heading_commands"] = sample_go2w_heading_commands(env, num_reset) - return ResetPlan( - env_ids=env_ids, - qpos=qpos, - qvel=qvel, - info_updates=info_updates, - randomization=build_go2w_backend_reset_randomization(env, num_reset), - ) - - 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]: - del env_ids - return cast( - dict[str, np.ndarray], - env._compute_obs(info_updates, linvel, gyro, gravity, dof_pos, dof_vel), - ) - - def _sample_commands(self, env: Any, num_reset: int) -> np.ndarray: - commands = super()._sample_commands(env, num_reset) - zero_small_xy_commands(commands) - standing_prob = float(getattr(env.cfg.commands, "rel_standing_envs", 0.0)) - if standing_prob > 0.0: - standing = np.random.uniform(size=(num_reset,)) < min(standing_prob, 1.0) - commands[standing] = 0.0 - if getattr(env.cfg.commands, "heading_command", False): - commands[:, 2] = 0.0 - return commands - - -class Go2WJoystickEnv(Go2WBaseEnv): - _cfg: Go2WJoystickCfg # pyright: ignore[reportIncompatibleVariableOverride] - - def __init__(self, cfg: Go2WJoystickCfg, num_envs=1, backend_type="mujoco"): - if cfg.reward_config is None: - raise ValueError("reward_config must be provided via Hydra configuration") - backend = create_backend( - backend_type, - cfg.scene, - num_envs, - cfg.sim_dt, - base_name=cfg.asset.base_name, - push_body_name=cfg.domain_rand.push_body_name, - **env_backend_kwargs(cfg), - ) - super().__init__(cfg, backend, num_envs) - self._np_dtype = get_global_dtype() - self._leg_action_scale = self._build_leg_action_scale() - self._reward_cfg = cfg.reward_config - self._enable_reward_log = True - ctrl_range = np.asarray(self._backend.get_actuator_ctrl_range(), dtype=np.float64) - self._validate_motor_control_contract(ctrl_range, num_envs) - self._ctrl_lower = ctrl_range[:, 0].astype(self._np_dtype) - self._ctrl_upper = ctrl_range[:, 1].astype(self._np_dtype) - joint_range = self._backend.get_joint_range() - self._leg_joint_range = ( - np.asarray(joint_range[:NUM_LEG_ACTIONS], dtype=get_global_dtype()) - if joint_range is not None - else None - ) - self._base_motor_kp = np.full((NUM_LEG_ACTIONS,), cfg.control_config.Kp, dtype=np.float64) - self._base_motor_kd = np.full((NUM_LEG_ACTIONS,), cfg.control_config.Kd, dtype=np.float64) - self._base_wheel_kd = np.full( - (NUM_WHEEL_ACTIONS,), cfg.control_config.wheel_Kd, dtype=np.float64 - ) - self._motor_kp = np.broadcast_to(self._base_motor_kp, (num_envs, NUM_LEG_ACTIONS)).copy() - self._motor_kd = np.broadcast_to(self._base_motor_kd, (num_envs, NUM_LEG_ACTIONS)).copy() - self._wheel_kd = np.broadcast_to(self._base_wheel_kd, (num_envs, NUM_WHEEL_ACTIONS)).copy() - self._last_motor_ctrl = np.zeros((num_envs, NUM_GO2W_ACTIONS), dtype=self._np_dtype) - self._last_dof_vel_for_acc = np.zeros( - (num_envs, NUM_GO2W_ACTIONS), dtype=get_global_dtype() - ) - self._backend.set_pre_step_control(self._pre_step_motor_control) - self._init_reward_functions() - self._init_domain_randomization(Go2WJoystickDomainRandomizationProvider()) - - @property - def obs_groups_spec(self) -> dict[str, int]: - return {"obs": 53, "critic": 72} - - def reset(self, env_indices: np.ndarray) -> tuple[dict[str, np.ndarray], dict]: - env_ids = np.asarray(env_indices, dtype=np.int32) - obs, info = super().reset(env_ids) - dof_vel = self.get_dof_vel() - if dof_vel.shape[0] == self._num_envs: - self._last_dof_vel_for_acc[env_ids] = dof_vel[env_ids] - return obs, info - - def _validate_motor_control_contract(self, ctrl_range: np.ndarray, num_envs: int) -> None: - if self._backend.num_actuators != NUM_GO2W_ACTIONS: - raise ValueError( - f"Go2W requires {NUM_GO2W_ACTIONS} motor actuators, got {self._backend.num_actuators}" - ) - if ctrl_range.shape != (NUM_GO2W_ACTIONS, 2): - raise ValueError( - f"Go2W actuator ctrl_range must have shape ({NUM_GO2W_ACTIONS}, 2), " - f"got {ctrl_range.shape}" - ) - pos = stack_joint_sensors(self._backend, "pos", dtype=self.default_angles.dtype) - vel = stack_joint_sensors(self._backend, "vel", dtype=self.default_angles.dtype) - expected_shape = (num_envs, NUM_GO2W_ACTIONS) - if pos.shape != expected_shape: - raise ValueError(f"Go2W joint position sensor stack must have shape {expected_shape}") - if vel.shape != expected_shape: - raise ValueError(f"Go2W joint velocity sensor stack must have shape {expected_shape}") - - def _build_leg_action_scale(self) -> np.ndarray: - scale = np.full( - (NUM_LEG_ACTIONS,), - float(self._cfg.control_config.action_scale), - dtype=self._np_dtype, - ) - hip_action_scale = self._cfg.control_config.hip_action_scale - if hip_action_scale is not None: - scale[GO2W_HIP_INDICES] = float(hip_action_scale) - return scale - - def _init_reward_functions(self) -> None: - self._reward_fns: dict[str, Any] = { - "tracking_lin_vel": rewards.tracking_lin_vel, - "tracking_ang_vel": rewards.tracking_ang_vel, - "lin_vel_z": rewards.lin_vel_z, - "ang_vel_xy": rewards.ang_vel_xy, - "base_height": rewards.base_height, - "action_rate": rewards.action_rate, - "similar_to_default": rewards.similar_to_default, - "orientation": rewards.orientation, - "torques": self._reward_torques_l2, - "joint_torques_l2": self._reward_joint_torques_l2, - "energy": rewards.energy, - "dof_vel": self._reward_dof_vel, - "dof_acc": self._reward_dof_acc, - "joint_acc_l2": self._reward_dof_acc, - "wheel_acc": self._reward_wheel_acc, - "joint_acc_wheel_l2": self._reward_wheel_acc, - "stand_still": self._reward_stand_still, - "hip_pos": self._reward_hip_pos, - "dof_error": self._reward_dof_error, - "joint_pos_penalty": self._reward_joint_pos_penalty, - "joint_power": self._reward_joint_power, - "joint_mirror": self._reward_joint_mirror, - "alive": rewards.alive, - "upward": rewards.upward, - "wheel_vel": self._reward_wheel_vel, - } - - def sample_reset_motor_gains(self, num_reset: int) -> tuple[np.ndarray, np.ndarray]: - kp = np.broadcast_to(self._base_motor_kp, (num_reset, NUM_LEG_ACTIONS)).copy() - kd = np.broadcast_to(self._base_motor_kd, (num_reset, NUM_LEG_ACTIONS)).copy() - domain_rand = self._cfg.domain_rand - if domain_rand.randomize_kp: - low, high = domain_rand.kp_multiplier_range - kp *= np.random.uniform(low, high, size=(num_reset, 1)) - if domain_rand.randomize_kd: - low, high = domain_rand.kd_multiplier_range - kd *= np.random.uniform(low, high, size=(num_reset, 1)) - return kp, kd - - def set_motor_gains(self, env_ids: np.ndarray, kp: np.ndarray, kd: np.ndarray) -> None: - self._motor_kp[env_ids] = np.asarray(kp, dtype=np.float64) - self._motor_kd[env_ids] = np.asarray(kd, dtype=np.float64) - - def apply_action(self, actions: np.ndarray, state: NpEnvState) -> np.ndarray: - clipped_actions = np.asarray( - np.clip( - actions, - -self._cfg.control_config.clip_actions, - self._cfg.control_config.clip_actions, - ), - dtype=self._np_dtype, - ) - state.info["last_actions"] = state.info.get( - "current_actions", np.zeros_like(clipped_actions) - ) - state.info["current_actions"] = clipped_actions - exec_actions = ( - state.info["last_actions"] - if self._cfg.control_config.simulate_action_latency - else clipped_actions - ) - - leg_targets = ( - exec_actions[:, :NUM_LEG_ACTIONS] * self._leg_action_scale - + self.default_angles[:NUM_LEG_ACTIONS] - ) - wheel_velocity_targets = ( - exec_actions[:, NUM_LEG_ACTIONS:] * self._cfg.control_config.wheel_action_scale - ) - return np.concatenate([leg_targets, wheel_velocity_targets], axis=1, dtype=self._np_dtype) - - def _pre_step_motor_control(self, backend: Any, policy_ctrl: np.ndarray) -> np.ndarray: - joint_pos = stack_joint_sensors(backend, "pos", dtype=self.default_angles.dtype) - joint_vel = stack_joint_sensors(backend, "vel", dtype=self.default_angles.dtype) - motor_ctrl = compute_go2w_motor_ctrl( - policy_ctrl, - joint_pos, - joint_vel, - self._motor_kp, - self._motor_kd, - self._wheel_kd, - self._ctrl_lower, - self._ctrl_upper, - self._last_motor_ctrl, - ) - return motor_ctrl - - def update_state(self, state: NpEnvState) -> NpEnvState: - self._update_commands(state.info) - linvel = self.get_local_linvel() - gyro = self.get_gyro() - gravity = self._backend.get_sensor_data(self._cfg.sensor.gravity) - dof_pos = self.get_dof_pos() - dof_vel = self.get_dof_vel() - state.info["torques"] = self._last_motor_ctrl.copy() - state.info["qacc"] = self._estimate_dof_acc(dof_vel) - terminated = self._compute_terminated(gravity) - reward = self._compute_reward(state.info, linvel, gyro, gravity, dof_pos, dof_vel) - obs = self._compute_obs(state.info, linvel, gyro, gravity, dof_pos, dof_vel) - return state.replace(obs=obs, reward=reward, terminated=terminated) - - def _compute_terminated(self, gravity: np.ndarray) -> np.ndarray: - return gravity[:, 2] <= 0.5 - - def _compute_obs( - self, - info: dict, - linvel: np.ndarray, - gyro: np.ndarray, - gravity: np.ndarray, - dof_pos: np.ndarray, - dof_vel: np.ndarray, - ) -> dict[str, np.ndarray]: - noise_cfg = self._cfg.noise_config - leg_diff = dof_pos[:, :NUM_LEG_ACTIONS] - self.default_angles[:NUM_LEG_ACTIONS] - leg_vel = dof_vel[:, :NUM_LEG_ACTIONS] - wheel_vel = dof_vel[:, NUM_LEG_ACTIONS:] - noisy_gyro = self._obs_noise(gyro, noise_cfg.scale_gyro) - noisy_gravity = self._obs_noise(gravity, noise_cfg.scale_gravity) - noisy_leg_diff = self._obs_noise(leg_diff, noise_cfg.scale_joint_angle) - noisy_leg_vel = self._obs_noise(leg_vel, noise_cfg.scale_joint_vel) - noisy_wheel_vel = self._obs_noise(wheel_vel, noise_cfg.scale_wheel_vel) - num_obs = gyro.shape[0] - last_actions = info.get("current_actions", np.zeros((num_obs, self._num_action))) - motor_ctrl = info.get("torques", np.zeros((num_obs, self._num_action), dtype=dof_pos.dtype)) - - obs = np.concatenate( - [ - noisy_gyro, - -noisy_gravity, - noisy_leg_diff, - noisy_leg_vel, - noisy_wheel_vel, - last_actions, - info["commands"], - ], - axis=1, - dtype=get_global_dtype(), - ) - critic = np.concatenate( - [ - gyro, - -gravity, - leg_diff, - leg_vel, - wheel_vel, - last_actions, - info["commands"], - linvel, - motor_ctrl, - ], - axis=1, - dtype=get_global_dtype(), - ) - return {"obs": obs, "critic": critic} - - def _compute_reward(self, info: dict, linvel, gyro, gravity, dof_pos, dof_vel) -> np.ndarray: - dtype = get_global_dtype() - num_obs = linvel.shape[0] - ctx = RewardContext( - info=info, - linvel=linvel, - gyro=gyro, - dof_pos=dof_pos[:, :NUM_LEG_ACTIONS], - dof_vel=dof_vel, - num_envs=num_obs, - default_angles=DEFAULT_GO2W_ANGLES[:NUM_LEG_ACTIONS].astype(dtype), - tracking_sigma=self._reward_cfg.tracking_sigma, - base_height_target=self._reward_cfg.base_height_target, - base_height=self._reward_base_height_values(num_obs), - gravity=gravity, - joint_range=self._leg_joint_range, - ) - return rewards.run_reward_dispatch( - scales=self._reward_cfg.scales, - fns=self._reward_fns, - ctx=ctx, - info=info, - enable_log=self._enable_reward_log, - ctrl_dt=self._cfg.ctrl_dt, - only_positive=self._reward_cfg.only_positive_rewards, - ) - - def _update_commands(self, info: dict) -> None: - commands = info.get("commands") - if commands is None: - return - - commands_arr = np.asarray(commands, dtype=get_global_dtype()) - resampling_time = float(getattr(self._cfg.commands, "resampling_time", 0.0)) - if resampling_time > 0.0: - interval_steps = max(int(round(resampling_time / self._cfg.ctrl_dt)), 1) - steps = np.asarray(info.get("steps", np.zeros((self._num_envs,), dtype=np.uint32))) - resample_mask = (steps > 0) & ((steps % interval_steps) == 0) - if np.any(resample_mask): - num_resample = int(np.count_nonzero(resample_mask)) - low = np.asarray(self._cfg.commands.vel_limit[0], dtype=get_global_dtype()) - high = np.asarray(self._cfg.commands.vel_limit[1], dtype=get_global_dtype()) - sampled = np.random.uniform(low=low, high=high, size=(num_resample, 3)).astype( - get_global_dtype() - ) - zero_small_xy_commands(sampled) - standing_prob = float(getattr(self._cfg.commands, "rel_standing_envs", 0.0)) - if standing_prob > 0.0: - standing = np.random.uniform(size=(num_resample,)) < min(standing_prob, 1.0) - sampled[standing] = 0.0 - commands_arr[resample_mask] = sampled - if getattr(self._cfg.commands, "heading_command", False): - heading_commands = self._ensure_heading_commands(info, commands_arr.shape[0]) - heading_commands[resample_mask] = sample_go2w_heading_commands( - self, num_resample - ) - info["heading_commands"] = heading_commands - - if getattr(self._cfg.commands, "heading_command", False): - heading_commands = self._ensure_heading_commands(info, commands_arr.shape[0]) - base_quat = np.asarray(self._backend.get_base_quat(), dtype=get_global_dtype()) - if base_quat.shape[0] == commands_arr.shape[0]: - stiffness = float(getattr(self._cfg.commands, "heading_control_stiffness", 0.5)) - apply_heading_yaw_feedback( - commands_arr, base_quat, heading_commands, stiffness=stiffness - ) - info["commands"] = commands_arr - - def _ensure_heading_commands(self, info: dict, num_obs: int) -> np.ndarray: - heading_commands = info.get("heading_commands") - if heading_commands is None or np.asarray(heading_commands).shape != (num_obs,): - heading_commands = sample_go2w_heading_commands(self, num_obs) - info["heading_commands"] = heading_commands - heading_commands = np.asarray(heading_commands, dtype=get_global_dtype()) - info["heading_commands"] = heading_commands - return heading_commands - - def _estimate_dof_acc(self, dof_vel: np.ndarray) -> np.ndarray: - qacc = np.asarray((dof_vel - self._last_dof_vel_for_acc) / self._cfg.ctrl_dt) - self._last_dof_vel_for_acc[:] = dof_vel - return np.asarray(qacc, dtype=get_global_dtype()) - - def _reward_base_height_values(self, num_obs: int) -> np.ndarray: - base_pos = np.asarray(self._backend.get_base_pos(), dtype=get_global_dtype()) - if base_pos.shape[0] != num_obs: - return np.zeros((num_obs,), dtype=get_global_dtype()) - return np.asarray(base_pos[:, 2], dtype=get_global_dtype()) - - def _reward_wheel_vel(self, ctx: RewardContext) -> np.ndarray: - assert ctx.dof_vel is not None - wheel_vel = ctx.dof_vel[:, NUM_LEG_ACTIONS:] - return np.asarray(np.sum(np.square(wheel_vel), axis=1), dtype=get_global_dtype()) - - def _reward_torques_l2(self, ctx: RewardContext) -> np.ndarray: - torques = np.asarray( - ctx.info.get("torques", np.zeros((ctx.num_envs, self._num_action))), - dtype=get_global_dtype(), - ) - return np.asarray(np.sum(np.square(torques), axis=1), dtype=get_global_dtype()) - - def _reward_joint_torques_l2(self, ctx: RewardContext) -> np.ndarray: - torques = np.asarray( - ctx.info.get("torques", np.zeros((ctx.num_envs, self._num_action))), - dtype=get_global_dtype(), - ) - return np.asarray( - np.sum(np.square(torques[:, :NUM_LEG_ACTIONS]), axis=1), - dtype=get_global_dtype(), - ) - - def _reward_dof_vel(self, ctx: RewardContext) -> np.ndarray: - assert ctx.dof_vel is not None - return np.asarray( - np.sum(np.square(ctx.dof_vel[:, :NUM_LEG_ACTIONS]), axis=1), - dtype=get_global_dtype(), - ) - - def _reward_dof_acc(self, ctx: RewardContext) -> np.ndarray: - qacc = np.asarray( - ctx.info.get("qacc", np.zeros((ctx.num_envs, NUM_GO2W_ACTIONS))), - dtype=get_global_dtype(), - ) - return np.asarray(np.sum(np.square(qacc[:, :NUM_LEG_ACTIONS]), axis=1), dtype=qacc.dtype) - - def _reward_wheel_acc(self, ctx: RewardContext) -> np.ndarray: - qacc = np.asarray( - ctx.info.get("qacc", np.zeros((ctx.num_envs, NUM_GO2W_ACTIONS))), - dtype=get_global_dtype(), - ) - return np.asarray(np.sum(np.square(qacc[:, NUM_LEG_ACTIONS:]), axis=1), dtype=qacc.dtype) - - def _reward_stand_still(self, ctx: RewardContext) -> np.ndarray: - commands = ctx.info["commands"] - stopped = np.linalg.norm(commands[:, :2], axis=1) < 0.1 - dof_error = np.sum(np.abs(ctx.dof_pos - DEFAULT_GO2W_ANGLES[:NUM_LEG_ACTIONS]), axis=1) - return np.asarray(dof_error * stopped, dtype=get_global_dtype()) - - def _reward_hip_pos(self, ctx: RewardContext) -> np.ndarray: - diff = ctx.dof_pos[:, GO2W_HIP_INDICES] - DEFAULT_GO2W_ANGLES[GO2W_HIP_INDICES] - return np.asarray(np.sum(np.square(diff), axis=1), dtype=get_global_dtype()) - - def _reward_dof_error(self, ctx: RewardContext) -> np.ndarray: - diff = ctx.dof_pos - DEFAULT_GO2W_ANGLES[:NUM_LEG_ACTIONS] - return np.asarray(np.sum(np.square(diff), axis=1), dtype=get_global_dtype()) - - def _reward_joint_pos_penalty(self, ctx: RewardContext) -> np.ndarray: - return rewards.joint_pos_penalty( - ctx, - stand_still_scale=self._reward_cfg.joint_pos_penalty_stand_still_scale, - velocity_threshold=self._reward_cfg.joint_pos_penalty_velocity_threshold, - command_threshold=self._reward_cfg.joint_pos_penalty_command_threshold, - ) - - def _reward_joint_power(self, ctx: RewardContext) -> np.ndarray: - assert ctx.dof_vel is not None - torques = np.asarray( - ctx.info.get("torques", np.zeros((ctx.num_envs, self._num_action))), - dtype=get_global_dtype(), - ) - return np.asarray( - np.sum(np.abs(ctx.dof_vel[:, :NUM_LEG_ACTIONS] * torques[:, :NUM_LEG_ACTIONS]), axis=1), - dtype=get_global_dtype(), - ) - - def _reward_joint_mirror(self, ctx: RewardContext) -> np.ndarray: - fr_rl = ctx.dof_pos[:, 0:3] - ctx.dof_pos[:, 9:12] - fl_rr = ctx.dof_pos[:, 3:6] - ctx.dof_pos[:, 6:9] - mirror = 0.5 * (np.sum(np.square(fr_rl), axis=1) + np.sum(np.square(fl_rr), axis=1)) - return np.asarray(mirror, dtype=get_global_dtype()) - - -# Go2WJoystickCfg and Go2WJoystickEnv remain solely as the rough-task bridge. -# The flat production identity is Hydra-owned and uses the generic Manager-Based -# factory; the bridge is deleted with the rough-task migration. -registry.register_env_config("Go2WJoystickFlat", ManagerBasedRlEnvCfg) -registry.register_env("Go2WJoystickFlat", make_manager_based_rl_env, sim_backend="mujoco") -registry.register_env("Go2WJoystickFlat", make_manager_based_rl_env, sim_backend="motrix") -registry.register_env("Go2WJoystickFlat", make_manager_based_rl_env, sim_backend="drake") diff --git a/src/unilab/tasks/locomotion/go2w/rough.py b/src/unilab/tasks/locomotion/go2w/rough.py deleted file mode 100644 index 5c9151ba4..000000000 --- a/src/unilab/tasks/locomotion/go2w/rough.py +++ /dev/null @@ -1,380 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass, field -from typing import Any - -import numpy as np - -from unilab.assets import ASSETS_ROOT_PATH -from unilab.base import registry -from unilab.base.np_env import NpEnvState -from unilab.base.scene import SceneCfg, TerrainSceneCfg -from unilab.dr import DomainRandomizationManager, ResetPlan -from unilab.dr.dr_utils import zero_actions -from unilab.dtype_config import get_global_dtype -from unilab.tasks.locomotion.common import rewards -from unilab.tasks.locomotion.common.commands import ( - Commands, - apply_heading_yaw_feedback, - zero_small_xy_commands, -) -from unilab.tasks.locomotion.common.height_scan import ( - HeightScanConfig, - base_height_from_scan, - height_scan_obs, - init_height_scan_sensor, - raw_height_scan_obs, - terrain_out_of_bounds, -) -from unilab.tasks.locomotion.common.rewards import RewardContext -from unilab.tasks.locomotion.common.terrain_spawn import ( - TerrainCurriculumCfg, - TerrainSpawnManager, -) -from unilab.tasks.locomotion.go2w.base import NUM_GO2W_ACTIONS, NUM_LEG_ACTIONS -from unilab.tasks.locomotion.go2w.joystick import ( - Go2WJoystickCfg, - Go2WJoystickDomainRandomizationProvider, - Go2WJoystickEnv, - build_go2w_backend_reset_randomization, - sample_go2w_heading_commands, -) -from unilab.terrains import ( - SubTerrainCfg, - TerrainGeneratorCfg, - flat, - hf_pyramid_slope, - hf_pyramid_slope_inv, - pyramid_stairs, - pyramid_stairs_inv, - random_rough, - wave_terrain, -) -from unilab.utils.rotation import ( - np_quat_from_euler_xyz, - np_quat_mul, -) - -# pyright: reportIncompatibleVariableOverride=false, reportAttributeAccessIssue=false, reportCallIssue=false - - -@dataclass -class Go2WRoughCommands(Commands): - vel_limit: list[list[float]] = field( - default_factory=lambda: [[-1.0, -1.0, -1.0], [1.0, 1.0, 1.0]] - ) - resampling_time: float = 10.0 - heading_command: bool = True - heading_range: list[float] = field(default_factory=lambda: [-np.pi, np.pi]) - - -@dataclass -class RoughTerminationConfig: - terrain_out_of_bounds: bool = True - terrain_distance_buffer: float = 3.0 - - -@dataclass(kw_only=True) -class Go2WRoughTerrainCfg(TerrainGeneratorCfg): - size: tuple[float, float] = (8.0, 8.0) - num_rows: int = 6 - num_cols: int = 6 - border_width: float = 1.0 - add_lights: bool = True - horizontal_scale: float = 0.1 - - sub_terrains: dict[str, SubTerrainCfg] = field( - default_factory=lambda: { - "flat": flat(proportion=0.0), - "pyramid_stairs": pyramid_stairs( - proportion=0.1, - step_height_range=(0.025, 0.10), - step_width=0.4, - platform_width=3.0, - border_width=0.2, - ), - "pyramid_stairs_inv": pyramid_stairs_inv( - proportion=0.1, - step_height_range=(0.025, 0.10), - step_width=0.4, - platform_width=3.0, - border_width=0.2, - ), - "hf_pyramid_slope": hf_pyramid_slope( - proportion=0.2, - slope_range=(0.0, 0.3), - platform_width=2.0, - border_width=0.2, - ), - "hf_pyramid_slope_inv": hf_pyramid_slope_inv( - proportion=0.2, - slope_range=(0.0, 0.3), - platform_width=2.0, - border_width=0.2, - ), - "random_rough": random_rough( - proportion=0.3, - noise_range=(0.01, 0.06), - noise_step=0.01, - border_width=0.2, - ), - "wave_terrain": wave_terrain( - proportion=0.3, - amplitude_range=(0.0, 0.12), - num_waves=4, - border_width=0.2, - ), - } - ) - - -@registry.envcfg("Go2WJoystickRough") -@dataclass -class Go2WJoystickRoughCfg(Go2WJoystickCfg): - """Go2W rough terrain task with procedurally generated sub-terrains.""" - - scene: SceneCfg = field( - default_factory=lambda: SceneCfg( - model_file=str(ASSETS_ROOT_PATH / "robots" / "go2w" / "go2w.xml"), - fragment_files=[ - str(ASSETS_ROOT_PATH / "robots" / "go2w" / "locomotion_task.xml"), - ], - terrain=TerrainSceneCfg( - generator=Go2WRoughTerrainCfg(), - hfield_name="terrain_hfield", - geom_name="floor", - ), - ) - ) - commands: Go2WRoughCommands = field(default_factory=Go2WRoughCommands) - terrain_scan: HeightScanConfig = field(default_factory=HeightScanConfig) - termination_config: RoughTerminationConfig = field(default_factory=RoughTerminationConfig) - terrain_curriculum: TerrainCurriculumCfg = field(default_factory=TerrainCurriculumCfg) - - -class Go2WJoystickRoughDomainRandomizationProvider(Go2WJoystickDomainRandomizationProvider): - def _sample_commands(self, env: Any, num_reset: int) -> np.ndarray: - commands = super()._sample_commands(env, num_reset) - zero_small_xy_commands(commands, threshold=0.08) - standing_prob = env.cfg.commands.rel_standing_envs - if standing_prob > 0.0: - standing = np.random.uniform(size=(num_reset,)) < min(standing_prob, 1.0) - commands[standing] = 0.0 - if env.cfg.commands.heading_command: - commands[:, 2] = 0.0 - return commands - - 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)) - qpos[:, 2] += np.random.uniform(0.25, 0.5, (num_reset,)) - qpos[:, 0:3] += env._spawn.origins_for(env_ids) - roll = np.random.uniform(-3.14, 3.14, (num_reset,)) - pitch = np.random.uniform(-3.14, 3.14, (num_reset,)) - yaw = np.random.uniform(-3.14, 3.14, (num_reset,)) - qpos[:, 3:7] = np_quat_mul(qpos[:, 3:7], np_quat_from_euler_xyz(roll, pitch, yaw)) - qvel[:, 0:6] = np.asarray( - np.random.uniform(-0.5, 0.5, size=(num_reset, 6)), dtype=get_global_dtype() - ) - - motor_kp, motor_kd = env.sample_reset_motor_gains(num_reset) - env.set_motor_gains(env_ids, motor_kp, motor_kd) - commands = self._sample_commands(env, num_reset) - info_updates: dict[str, Any] = { - "commands": commands, - "current_actions": zero_actions(num_reset, env._num_action), - "last_actions": zero_actions(num_reset, env._num_action), - "motor_kp": motor_kp.astype(get_global_dtype()), - "motor_kd": motor_kd.astype(get_global_dtype()), - "torques": np.zeros((num_reset, env._num_action), dtype=get_global_dtype()), - } - if getattr(env.cfg.commands, "heading_command", False): - info_updates["heading_commands"] = sample_go2w_heading_commands(env, num_reset) - return ResetPlan( - env_ids=env_ids, - qpos=qpos, - qvel=qvel, - info_updates=info_updates, - randomization=build_go2w_backend_reset_randomization(env, num_reset), - ) - - -@registry.env("Go2WJoystickRough", sim_backend="mujoco") -class Go2WJoystickRoughEnv(Go2WJoystickEnv): - _cfg: Go2WJoystickRoughCfg - _height_scan_dim: int = 0 - - def __init__(self, cfg: Go2WJoystickRoughCfg, num_envs=1, backend_type="mujoco"): - super().__init__(cfg, num_envs=num_envs, backend_type=backend_type) - terrain_spawn_data = self._backend.get_terrain_spawn_data() - terrain_generator = cfg.scene.terrain.generator if cfg.scene.terrain is not None else None - if terrain_spawn_data is not None and terrain_generator is not None: - self._spawn = TerrainSpawnManager( - num_envs, - terrain_spawn_data.terrain_origins, - cell_size=float(terrain_generator.size[0]), - cfg=cfg.terrain_curriculum, - sample_height=terrain_spawn_data.sample_height, - ) - self._dr_manager = DomainRandomizationManager( - self, Go2WJoystickRoughDomainRandomizationProvider() - ) - init_height_scan_sensor(self, cfg.terrain_scan, cfg.asset.base_name) - - @property - def obs_groups_spec(self) -> dict[str, int]: - return {"obs": 53, "critic": 56 + self._height_scan_dim} - - def _init_reward_functions(self) -> None: - def gated(fn): - return lambda ctx: fn(ctx) * self._upright_scale(ctx.gravity) - - def _joint_pos_penalty(ctx: RewardContext) -> np.ndarray: - return self._reward_joint_pos_penalty(ctx) * self._upright_scale(ctx.gravity) - - def _stand_still(ctx: RewardContext) -> np.ndarray: - return self._reward_stand_still(ctx) * self._upright_scale(ctx.gravity) - - self._reward_fns = { - "tracking_lin_vel": gated(rewards.tracking_lin_vel), - "tracking_ang_vel": gated(rewards.tracking_ang_vel), - "lin_vel_z": gated(rewards.lin_vel_z), - "ang_vel_xy": gated(rewards.ang_vel_xy), - "base_height": gated(rewards.base_height), - "orientation": gated(rewards.orientation), - "similar_to_default": gated(rewards.similar_to_default), - "torques": gated(self._reward_torques_l2), - "joint_torques_l2": gated(self._reward_joint_torques_l2), - "energy": gated(rewards.energy), - "dof_vel": gated(self._reward_dof_vel), - "dof_acc": gated(self._reward_dof_acc), - "joint_acc_l2": gated(self._reward_dof_acc), - "wheel_acc": gated(self._reward_wheel_acc), - "joint_acc_wheel_l2": gated(self._reward_wheel_acc), - "stand_still": _stand_still, - "hip_pos": gated(self._reward_hip_pos), - "dof_error": gated(self._reward_dof_error), - "joint_pos_penalty": _joint_pos_penalty, - "joint_power": gated(self._reward_joint_power), - "joint_mirror": gated(self._reward_joint_mirror), - "alive": rewards.alive, - "upward": rewards.upward, - "wheel_vel": gated(self._reward_wheel_vel), - "action_rate": rewards.action_rate, - } - - def _upright_scale(self, gravity: np.ndarray | None) -> np.ndarray: - return rewards.upright_scale(gravity, self._num_envs) - - def _compute_obs( - self, - info: dict, - linvel: np.ndarray, - gyro: np.ndarray, - gravity: np.ndarray, - dof_pos: np.ndarray, - dof_vel: np.ndarray, - ) -> dict[str, np.ndarray]: - noise_cfg = self._cfg.noise_config - leg_diff = dof_pos[:, :NUM_LEG_ACTIONS] - self.default_angles[:NUM_LEG_ACTIONS] - policy_gyro = self._obs_noise(gyro, noise_cfg.scale_gyro) * 0.25 - policy_gravity = self._obs_noise(-gravity, noise_cfg.scale_gravity) - policy_leg_diff = self._obs_noise(leg_diff, noise_cfg.scale_joint_angle) - policy_dof_vel = self._obs_noise(dof_vel, noise_cfg.scale_joint_vel) * 0.05 - num_obs = gyro.shape[0] - last_actions = info.get( - "current_actions", np.zeros((num_obs, NUM_GO2W_ACTIONS), dtype=dof_pos.dtype) - ) - commands = info["commands"] - - obs = np.concatenate( - [ - policy_gyro, - policy_gravity, - commands, - policy_leg_diff, - policy_dof_vel, - last_actions, - ], - axis=1, - dtype=get_global_dtype(), - ) - critic_base = np.concatenate( - [linvel, gyro, -gravity, commands, leg_diff, dof_vel, last_actions], - axis=1, - dtype=get_global_dtype(), - ) - critic = np.concatenate( - [critic_base, height_scan_obs(self, self._cfg.terrain_scan, num_obs)], - axis=1, - dtype=get_global_dtype(), - ) - return {"obs": obs, "critic": critic} - - def _reward_base_height_values(self, num_obs: int) -> np.ndarray: - height = base_height_from_scan(self, num_obs) - if height.shape[0] != num_obs: - return super()._reward_base_height_values(num_obs) - return height - - def _update_commands(self, info: dict) -> None: - commands = info.get("commands") - if commands is None: - return - - commands_arr = np.asarray(commands, dtype=get_global_dtype()) - resampling_time = float(self._cfg.commands.resampling_time) - if resampling_time > 0.0: - interval_steps = max(int(round(resampling_time / self._cfg.ctrl_dt)), 1) - steps = np.asarray(info.get("steps", np.zeros((self._num_envs,), dtype=np.uint32))) - resample_mask = (steps > 0) & ((steps % interval_steps) == 0) - if np.any(resample_mask): - num_resample = int(np.count_nonzero(resample_mask)) - low = np.asarray(self._cfg.commands.vel_limit[0], dtype=get_global_dtype()) - high = np.asarray(self._cfg.commands.vel_limit[1], dtype=get_global_dtype()) - sampled = np.random.uniform(low=low, high=high, size=(num_resample, 3)).astype( - get_global_dtype() - ) - zero_small_xy_commands(commands, threshold=0.08) - commands_arr[resample_mask] = sampled - if self._cfg.commands.heading_command: - heading_commands = self._ensure_heading_commands(info, commands_arr.shape[0]) - heading_commands[resample_mask] = sample_go2w_heading_commands( - self, num_resample - ) - info["heading_commands"] = heading_commands - - if self._cfg.commands.heading_command: - heading_commands = self._ensure_heading_commands(info, commands_arr.shape[0]) - base_quat = np.asarray(self._backend.get_base_quat(), dtype=get_global_dtype()) - if base_quat.shape[0] == commands_arr.shape[0]: - apply_heading_yaw_feedback(commands_arr, base_quat, heading_commands, stiffness=0.5) - info["commands"] = commands_arr - - def _compute_terminated(self, gravity: np.ndarray) -> np.ndarray: - del gravity - return np.zeros((self._num_envs,), dtype=bool) - - def _raw_height_scan_obs(self, num_obs: int) -> tuple[np.ndarray | None, np.ndarray | None]: - return raw_height_scan_obs(self, num_obs) - - def _compute_truncated(self, state: NpEnvState) -> np.ndarray: - truncated = super()._compute_truncated(state) - if self._cfg.termination_config.terrain_out_of_bounds: - terrain_scene = self._cfg.scene.terrain - terrain_cfg = terrain_scene.generator if terrain_scene is not None else None - np.logical_or( - truncated, - terrain_out_of_bounds( - self, - terrain_cfg, - float(self._cfg.termination_config.terrain_distance_buffer), - ), - out=truncated, - ) - return truncated - - -registry.register_env("Go2WJoystickRough", Go2WJoystickRoughEnv, sim_backend="motrix") diff --git a/src/unilab/tasks/migration_matrix.py b/src/unilab/tasks/migration_matrix.py index db061874c..350d3f946 100644 --- a/src/unilab/tasks/migration_matrix.py +++ b/src/unilab/tasks/migration_matrix.py @@ -113,10 +113,10 @@ def migration_record(task_name: str) -> TaskMigrationRecord: return TaskMigrationRecord( task_name, "quadruped_rough", - "Adapted", - "mba", - "Terrain and height-scan terms depend on the pending raycaster capability boundary.", - "Migrate as one rough-family PR; use the compatibility seam only if a new public capability is required.", + "Compatible", + "complete", + "Hydra owners materialize shared terrain, height-scan, reset, and curriculum manager terms on the canonical runtime.", + "Keep the shared rough-family contract and both backend owners in sync.", ) if task_name in _G1_LOCOMOTION_TASKS: return TaskMigrationRecord( diff --git a/tests/base/test_backend_conformance.py b/tests/base/test_backend_conformance.py index c9c8f06c3..6a9474fec 100644 --- a/tests/base/test_backend_conformance.py +++ b/tests/base/test_backend_conformance.py @@ -33,10 +33,8 @@ SRC_ROOT / "unilab" / "tasks", ) _TERRAIN_CONSUMER_PATHS = ( + Path("locomotion/common/rough_manager_terms.py"), Path("locomotion/common/terrain_spawn.py"), - Path("locomotion/go1/joystick.py"), - Path("locomotion/go2/joystick.py"), - Path("locomotion/go2w/rough.py"), ) NUM_ENVS = 2 diff --git a/tests/benchmark/test_env_step_config_contract.py b/tests/benchmark/test_env_step_config_contract.py index b0d33e862..720094c5a 100644 --- a/tests/benchmark/test_env_step_config_contract.py +++ b/tests/benchmark/test_env_step_config_contract.py @@ -1,13 +1,11 @@ from __future__ import annotations -from pathlib import Path - import pytest from hydra.errors import ConfigCompositionException -from omegaconf import OmegaConf from scripts.benchmark.env import benchmark_env_step as bench from unilab.envs import ManagerBasedRlEnvCfg, make_manager_based_rl_env +from unilab.tasks.locomotion.common.rough_manager_terms import QuadrupedRoughTerrainCfg def test_go2w_flat_benchmark_uses_production_manager_owner() -> None: @@ -22,32 +20,28 @@ def test_go2w_flat_benchmark_uses_production_manager_owner() -> None: def test_go2w_rough_cfg_matches_ppo_owner_yaml() -> None: cfg = bench.TASK_CONFIGS["go2w_rough"].build_cfg("mujoco") - owner_path = ( - Path(bench.ROOT_DIR) / "conf" / "ppo" / "task" / "go2w_joystick_rough" / "mujoco.yaml" - ) - owner_cfg = OmegaConf.load(owner_path) - assert cfg.reward_config.scales == OmegaConf.to_container( - owner_cfg.reward.scales, - resolve=True, - ) - assert cfg.reward_config.tracking_sigma == owner_cfg.reward.tracking_sigma - assert cfg.reward_config.base_height_target == owner_cfg.reward.base_height_target - assert cfg.scene.model_file == owner_cfg.env.scene.model_file + assert isinstance(cfg, ManagerBasedRlEnvCfg) + assert cfg.scene is not None + assert cfg.scene.terrain is not None + assert isinstance(cfg.scene.terrain.generator, QuadrupedRoughTerrainCfg) + assert cfg.scene.model_file.endswith("go2w_mujoco.xml") + assert cfg.actions["motor"].wheel_action_scale == pytest.approx(5.0) + assert cfg.rewards["tracking_lin_vel"].weight == pytest.approx(3.0) + assert bench.TASK_CONFIGS["go2w_rough"].env_cls_factory() is make_manager_based_rl_env def test_env_and_reward_overrides_use_hydra_composition() -> None: cfg = bench.TASK_CONFIGS["go2w_rough"].build_cfg( "mujoco", [ - "env.control_config.action_scale=0.125", - "reward.tracking_sigma=${reward.base_height_target}", + "env.actions.motor.leg_action_scale=0.125", + "reward.tracking_lin_vel.weight=2.25", ], ) - assert cfg.control_config.action_scale == pytest.approx(0.125) - assert cfg.reward_config.tracking_sigma == pytest.approx(0.4) - assert isinstance(cfg.reward_config.tracking_sigma, float) + assert cfg.actions["motor"].leg_action_scale == pytest.approx(0.125) + assert cfg.rewards["tracking_lin_vel"].weight == pytest.approx(2.25) def test_unknown_env_override_fails_in_hydra() -> None: @@ -70,8 +64,8 @@ def test_training_sim_backend_override_is_rejected(override: str) -> None: def test_only_owner_config_overrides_are_forwarded() -> None: overrides = [ "task=go2w_joystick_rough/mujoco", - "env.control_config.action_scale=0.125", - "+reward.scales.custom_term=1.0", + "env.actions.motor.leg_action_scale=0.125", + "reward.tracking_lin_vel.weight=2.25", ] assert bench._owner_config_overrides(overrides) == overrides[1:] diff --git a/tests/config/test_config_system.py b/tests/config/test_config_system.py index 1fa7f78d8..9548dc595 100644 --- a/tests/config/test_config_system.py +++ b/tests/config/test_config_system.py @@ -481,52 +481,6 @@ def test_ppo_go2w_motrix_uses_motor_owner_scene_path(): assert cfg.reward.torques.weight < 0.0 -def test_ppo_go2w_rough_mujoco_uses_terrain_generator(): - cfg = _compose("ppo", overrides=["task=go2w_joystick_rough/mujoco"]) - - assert cfg.training.task_name == "Go2WJoystickRough" - assert cfg.training.sim_backend == "mujoco" - assert str(cfg.env.scene.model_file).endswith("src/unilab/assets/robots/go2w/go2w_mujoco.xml") - assert cfg.env.scene.terrain.hfield_name == "terrain_hfield" - assert cfg.env.scene.terrain.geom_name == "floor" - assert cfg.env.terrain_scan.hfield_name == "terrain_hfield" - assert cfg.env.terrain_scan.geom_name == "floor" - assert cfg.env.commands.resampling_time == pytest.approx(10.0) - assert cfg.env.commands.heading_command is True - assert cfg.env.commands.vel_limit == [[-1.0, -1.0, -1.0], [1.0, 1.0, 1.0]] - assert cfg.env.commands.heading_range == pytest.approx([-3.141592653589793, 3.141592653589793]) - assert cfg.env.control_config.clip_actions == pytest.approx(100.0) - assert cfg.env.control_config.action_scale == pytest.approx(0.25) - assert cfg.env.control_config.hip_action_scale == pytest.approx(0.125) - assert cfg.env.control_config.wheel_action_scale == pytest.approx(5.0) - assert cfg.env.domain_rand.randomize_kp is True - assert cfg.env.domain_rand.randomize_kd is True - assert cfg.env.domain_rand.kp_multiplier_range == [0.5, 1.0] - assert cfg.reward.scales.tracking_lin_vel == pytest.approx(3.0) - assert cfg.reward.scales.hip_pos == pytest.approx(-2.0) - assert cfg.reward.scales.joint_mirror == pytest.approx(-0.05) - assert cfg.reward.only_positive_rewards is False - assert cfg.algo.max_iterations == 1200 - - -def test_ppo_go2w_rough_motrix_uses_yaw_reset_and_strong_control(): - cfg = _compose("ppo", overrides=["task=go2w_joystick_rough/motrix"]) - - assert cfg.training.task_name == "Go2WJoystickRough" - assert cfg.training.sim_backend == "motrix" - assert cfg.env.commands.vel_limit == [[-1.0, -1.0, -1.0], [1.0, 1.0, 1.0]] - assert cfg.env.commands.heading_range == pytest.approx([-3.141592653589793, 3.141592653589793]) - assert cfg.env.control_config.action_scale == pytest.approx(0.25) - assert cfg.env.control_config.hip_action_scale == pytest.approx(0.125) - assert cfg.env.control_config.wheel_action_scale == pytest.approx(5.0) - assert cfg.env.domain_rand.randomize_kp is True - assert cfg.env.domain_rand.randomize_kd is True - assert cfg.reward.scales.orientation == pytest.approx(-2.0) - assert cfg.reward.scales.hip_pos == pytest.approx(-0.5) - assert cfg.reward.scales.upward == pytest.approx(1.0) - assert cfg.algo.max_iterations == 1200 - - def test_offpolicy_g1_walk_flat_motrix_preserves_backend_env_overrides(): cfg = _compose("offpolicy", overrides=["algo=sac", "task=sac/g1_walk_flat/motrix"]) diff --git a/tests/config/test_locomotion_params.py b/tests/config/test_locomotion_params.py index 62a03f468..b9c9ff793 100644 --- a/tests/config/test_locomotion_params.py +++ b/tests/config/test_locomotion_params.py @@ -224,135 +224,6 @@ def test_offpolicy_flashsac_go2_task_overrides(): assert cfg.env.actions.joint_pos.scale == pytest.approx(0.4) -def test_go2_joystick_rough_uses_terrain_generator(): - from unilab.assets import ASSETS_ROOT_PATH - from unilab.base.scene import SceneCfg, TerrainSceneCfg - from unilab.tasks.locomotion.go2.rough import Go2JoystickRoughCfg - from unilab.terrains import TerrainGeneratorCfg - - cfg = Go2JoystickRoughCfg() - assert isinstance(cfg.scene, SceneCfg) - assert isinstance(cfg.scene.terrain, TerrainSceneCfg) - assert cfg.scene.model_file.endswith("go2.xml") - assert cfg.scene.fragment_files == [ - str(ASSETS_ROOT_PATH / "robots" / "go2" / "locomotion_task.xml") - ] - assert isinstance(cfg.scene.terrain.generator, TerrainGeneratorCfg) - assert cfg.scene.terrain.hfield_name == "terrain_hfield" - assert cfg.scene.terrain.geom_name == "floor" - assert len(cfg.scene.terrain.generator.sub_terrains) == 7 - - -def test_go2_joystick_rough_terrain_cfg_is_independent_per_instance(): - """Confirm rough terrain cfg defaults are not shared across instances.""" - from unilab.tasks.locomotion.go2.rough import Go2JoystickRoughCfg - - a = Go2JoystickRoughCfg() - b = Go2JoystickRoughCfg() - b.scene.terrain.generator.num_rows = 3 - assert a.scene.terrain.generator is not b.scene.terrain.generator - a.scene.terrain.generator.num_rows = 4 - assert b.scene.terrain.generator.num_rows == 3 - - -def test_go2_joystick_rough_playback_model_uses_backend_scene(tmp_path): - """Offline playback / video rendering must reuse the backend-compiled scene model.""" - import mujoco - - from unilab.tasks.locomotion.go2.joystick import RewardConfig - from unilab.tasks.locomotion.go2.rough import Go2JoystickRoughCfg, Go2JoystickRoughEnv - from unilab.visualization.playback import _resolve_render_play_model_files - - cfg = Go2JoystickRoughCfg( - reward_config=RewardConfig(scales={}, tracking_sigma=0.25, base_height_target=0.3) - ) - cfg.scene.terrain.generator.num_rows = 2 - cfg.scene.terrain.generator.num_cols = 2 - cfg.scene.terrain.generator.border_width = 0.0 - cfg.scene.terrain.generator.add_lights = False - cfg.scene.terrain.generator.seed = 0 - - env = Go2JoystickRoughEnv(cfg, num_envs=2, backend_type="mujoco") - try: - playback_model = env.get_playback_model(0) - assert isinstance(playback_model, mujoco.MjModel) - assert env._backend.terrain_origins is not None - assert (Path(env._backend.scene_artifacts_dir) / "hfields" / "hfield.png").is_file() - assert Path(env._backend.scene_visual_model_file).is_file() - assert mujoco.mj_name2id(playback_model, mujoco.mjtObj.mjOBJ_HFIELD, "terrain_hfield") >= 0 - assert mujoco.mj_name2id(playback_model, mujoco.mjtObj.mjOBJ_GEOM, "floor") >= 0 - assert mujoco.mj_name2id(playback_model, mujoco.mjtObj.mjOBJ_SENSOR, "FL_foot_contact") >= 0 - model_file = _resolve_render_play_model_files(env, num_envs=2, tmp_dir=tmp_path) - assert isinstance(model_file, str) - assert model_file.endswith(".mjb") - assert Path(model_file).is_file() - rendered_model = mujoco.MjModel.from_binary_path(model_file) - assert mujoco.mj_name2id(rendered_model, mujoco.mjtObj.mjOBJ_HFIELD, "terrain_hfield") >= 0 - assert mujoco.mj_name2id(rendered_model, mujoco.mjtObj.mjOBJ_GEOM, "floor") >= 0 - assert rendered_model.ngeom > playback_model.ngeom - finally: - env.close() - - -def test_go2_joystick_flat_no_terrain_materialized(): - """Flat task keeps the static scene source and has no terrain origins.""" - from unilab.tasks.locomotion.go2.joystick import ( - Go2JoystickCfg, - Go2WalkTask, - RewardConfig, - ) - - cfg = Go2JoystickCfg( - reward_config=RewardConfig(scales={}, tracking_sigma=0.25, base_height_target=0.3) - ) - env = Go2WalkTask(cfg, num_envs=4, backend_type="mujoco") - try: - assert env._backend.scene_model_file == cfg.scene.model_file - assert env._backend.terrain_origins is None - assert env._backend.scene_artifacts_dir is None - finally: - env.close() - - -def test_ppo_go2_joystick_rough_task_compose(): - from hydra import compose, initialize_config_dir - from hydra.core.global_hydra import GlobalHydra - - GlobalHydra.instance().clear() - with initialize_config_dir(config_dir=str(CONF_DIR / "ppo"), version_base="1.3"): - cfg = compose("config", overrides=["task=go2_joystick_rough/mujoco"]) - assert cfg.training.task_name == "Go2JoystickRough" - assert cfg.training.sim_backend == "mujoco" - - -def test_ppo_go2_joystick_rough_motrix_task_compose(): - from hydra import compose, initialize_config_dir - from hydra.core.global_hydra import GlobalHydra - - GlobalHydra.instance().clear() - with initialize_config_dir(config_dir=str(CONF_DIR / "ppo"), version_base="1.3"): - cfg = compose("config", overrides=["task=go2_joystick_rough/motrix"]) - assert cfg.training.task_name == "Go2JoystickRough" - assert cfg.training.sim_backend == "motrix" - assert cfg.algo.num_envs == 4096 - assert cfg.algo.max_iterations == 1500 - assert cfg.env.render_offset_mode == "zero" - assert cfg.env.scene.model_file.endswith("go2.xml") - assert cfg.env.scene.terrain.generator.num_rows == 6 - assert cfg.env.scene.terrain.generator.num_cols == 6 - assert cfg.env.terrain_scan.enabled is True - assert cfg.reward.scales.tracking_lin_vel == pytest.approx(3.0) - assert "base_height" not in cfg.reward.scales - assert "swing_feet_z" not in cfg.reward.scales - - -def test_go2_joystick_rough_motrix_registers_rough_env(): - from unilab.base import registry - from unilab.tasks.locomotion.go2.rough import Go2JoystickRoughEnv - - assert registry._envs["Go2JoystickRough"].env_factory_dict["motrix"] is Go2JoystickRoughEnv - - def test_offpolicy_g1_rough_terrain_task_overrides(): from hydra import compose, initialize_config_dir from hydra.core.global_hydra import GlobalHydra @@ -654,90 +525,3 @@ def test_ppo_x2_wall_flip_tracking(): assert cfg.reward.scales.motion_ee_body_pos_z == pytest.approx(2.0) assert cfg.reward.scales.action_rate_l2 == pytest.approx(-0.005) assert cfg.reward.scales.undesired_contacts == pytest.approx(-0.1) - - -# --------------------------------------------------------------------------- -# Issue #197 DoD: rough terrain profile params overridable via Hydra -# --------------------------------------------------------------------------- - - -def test_apply_cfg_overrides_deep_merges_dataclass_field(): - """registry.apply_cfg_overrides must deep-merge into existing dataclass - instances rather than re-instantiating them, so partial overrides like - `scene.terrain.generator.num_rows=4` keep `sub_terrains` and other defaults.""" - from unilab.base.registry import apply_cfg_overrides - from unilab.tasks.locomotion.go2.rough import Go2JoystickRoughCfg - - cfg = Go2JoystickRoughCfg() - cfg.scene.terrain.generator.num_cols = 3 - cfg.scene.terrain.generator.border_width = 2.5 - cfg.scene.terrain.generator.sub_terrains = { - "test_flat": cfg.scene.terrain.generator.sub_terrains["flat"] - } - cfg.scene.terrain.generator.add_lights = False - apply_cfg_overrides( - cfg, - {"scene": {"terrain": {"generator": {"num_rows": 4, "seed": 42, "curriculum": True}}}}, - ) - - # Overridden fields take effect. - assert cfg.scene.terrain.generator.num_rows == 4 - assert cfg.scene.terrain.generator.seed == 42 - assert cfg.scene.terrain.generator.curriculum is True - # Non-overridden fields preserve the pre-existing instance state. - assert cfg.scene.terrain.generator.num_cols == 3 - assert cfg.scene.terrain.generator.border_width == pytest.approx(2.5) - assert list(cfg.scene.terrain.generator.sub_terrains) == ["test_flat"] - assert cfg.scene.terrain.generator.add_lights is False - - -def test_ppo_go2_joystick_rough_hydra_terrain_override(): - """Issue #197 DoD: rough terrain profile parameters must be overridable - via Hydra command-line. Composes the resolved config and feeds it through - the same BackendAdapter -> registry.apply_cfg_overrides path the trainer - uses.""" - from hydra import compose, initialize_config_dir - from hydra.core.global_hydra import GlobalHydra - - from unilab.base.registry import apply_cfg_overrides - from unilab.tasks.locomotion.go2.rough import Go2JoystickRoughCfg - from unilab.training.backend_adapter import BackendAdapter - - GlobalHydra.instance().clear() - with initialize_config_dir(config_dir=str(CONF_DIR / "ppo"), version_base="1.3"): - cfg = compose( - "config", - overrides=[ - "task=go2_joystick_rough/mujoco", - "env.scene.terrain.generator.num_rows=4", - "env.scene.terrain.generator.num_cols=6", - "env.scene.terrain.generator.seed=42", - "env.scene.terrain.generator.curriculum=true", - ], - ) - - # Yaml exposes the overridable schema (struct-mode acceptance). - assert cfg.env.scene.terrain.hfield_name == "terrain_hfield" - assert cfg.env.scene.terrain.geom_name == "floor" - assert cfg.env.scene.terrain.generator.num_rows == 4 - assert cfg.env.scene.terrain.generator.num_cols == 6 - assert cfg.env.scene.terrain.generator.seed == 42 - assert cfg.env.scene.terrain.generator.curriculum is True - - # End-to-end: the override dict produced by the adapter must, after the - # registry's deep-merge, leave Go2JoystickRoughCfg in a coherent state — - # overridden fields applied, untouched dataclass defaults preserved. - adapter = BackendAdapter(cfg, root_dir=Path.cwd()) - env_cfg_override = adapter.build_task_env_cfg_override() - assert "scene" in env_cfg_override - assert env_cfg_override["scene"]["terrain"]["generator"]["num_rows"] == 4 - - env_cfg = Go2JoystickRoughCfg() - apply_cfg_overrides(env_cfg, env_cfg_override) - - assert env_cfg.scene.terrain.generator.num_rows == 4 - assert env_cfg.scene.terrain.generator.num_cols == 6 - assert env_cfg.scene.terrain.generator.seed == 42 - assert env_cfg.scene.terrain.generator.curriculum is True - # sub_terrains is not in the yaml schema, so its Python default survives. - assert len(env_cfg.scene.terrain.generator.sub_terrains) == 7 diff --git a/tests/envs/locomotion/go1/test_manager_based_flat_cfg.py b/tests/envs/locomotion/go1/test_manager_based_flat_cfg.py index 3b071ddd9..f393eee42 100644 --- a/tests/envs/locomotion/go1/test_manager_based_flat_cfg.py +++ b/tests/envs/locomotion/go1/test_manager_based_flat_cfg.py @@ -263,17 +263,17 @@ def test_go1_flat_owner_materializes_complete_plain_manager_cfg( _assert_no_omegaconf(env_cfg) -def test_go1_flat_registry_is_manager_only_and_rough_owns_the_legacy_bridge() -> None: +def test_go1_flat_and_rough_registries_are_manager_only() -> None: registry.ensure_registries() - from unilab.tasks.locomotion.go1.joystick import Go1JoystickCfg, Go1WalkTask - from unilab.tasks.locomotion.go1.rough import Go1JoystickRoughCfg, Go1JoystickRoughEnv assert registry.list_registered_envs()["Go1JoystickFlat"] == { "config_factory": "ManagerBasedRlEnvCfg", "available_backends": ["mujoco", "motrix", "drake"], } - assert Go1JoystickRoughCfg.__bases__ == (Go1JoystickCfg,) - assert Go1JoystickRoughEnv.__bases__ == (Go1WalkTask,) + assert registry.list_registered_envs()["Go1JoystickRough"] == { + "config_factory": "ManagerBasedRlEnvCfg", + "available_backends": ["mujoco", "motrix"], + } for legacy_override in ( {"reward_config": {}}, diff --git a/tests/envs/locomotion/go2w/test_go2w_height_scan.py b/tests/envs/locomotion/go2w/test_go2w_height_scan.py deleted file mode 100644 index c9e1a2c87..000000000 --- a/tests/envs/locomotion/go2w/test_go2w_height_scan.py +++ /dev/null @@ -1,77 +0,0 @@ -from __future__ import annotations - -from typing import Any - -import numpy as np -import pytest - -pytest.importorskip("mujoco", reason="mujoco not installed") - -try: - from mujoco_uni.batch_env import BatchEnvPool -except Exception: - pytest.skip( - "mujoco_uni.batch_env not available (platform/libstdc++ issue)", - allow_module_level=True, - ) - -if not hasattr(BatchEnvPool, "sample_hfield_height"): - pytest.skip( - "BatchEnvPool.sample_hfield_height requires a mujoco-uni-runtime build that ships it", - allow_module_level=True, - ) - -from unilab.tasks.locomotion.go2w.rough import Go2WJoystickRoughEnv - - -def test_go2w_rough_height_scan_uses_backend_native_sampling() -> None: - class FakeHeightScanner: - def __init__(self, heights: np.ndarray) -> None: - self.calls = 0 - self.heights = heights - - def scan(self) -> np.ndarray: - self.calls += 1 - return self.heights - - class FakeBackend: - def __init__(self) -> None: - self.scanner_calls: list[dict[str, Any]] = [] - self.base_pos = np.asarray([[0.0, 0.0, 0.6], [1.0, 0.0, 0.7]], dtype=np.float32) - self.heights = np.asarray([[0.1, 0.2], [0.3, 0.4]], dtype=np.float32) - self.scanner = FakeHeightScanner(self.heights) - - def get_base_pos(self) -> np.ndarray: - return self.base_pos - - def create_hfield_scanner(self, **kwargs: Any) -> FakeHeightScanner: - self.scanner_calls.append(kwargs) - return self.scanner - - env = object.__new__(Go2WJoystickRoughEnv) - fake_backend = FakeBackend() - env._backend = fake_backend - env._height_scan_dim = 2 - env._height_scan_hfield_geom_id = 7 - env._height_scan_frame_body_id = 3 - env._height_scan_offsets = np.asarray([[0.0, 0.0], [0.1, -0.1]], dtype=np.float64) - env._height_scan_sensor = fake_backend.create_hfield_scanner( - hfield_geom_id=env._height_scan_hfield_geom_id, - offsets=env._height_scan_offsets, - frame_body_id=env._height_scan_frame_body_id, - alignment="yaw", - output="height", - ) - - raw_heights, base_pos = env._raw_height_scan_obs(num_obs=2) - - np.testing.assert_array_equal(raw_heights, fake_backend.heights) - np.testing.assert_array_equal(base_pos, fake_backend.base_pos) - assert fake_backend.scanner.calls == 1 - assert len(fake_backend.scanner_calls) == 1 - call = fake_backend.scanner_calls[0] - assert call["hfield_geom_id"] == 7 - assert call["frame_body_id"] == 3 - assert call["alignment"] == "yaw" - assert call["output"] == "height" - np.testing.assert_array_equal(call["offsets"], env._height_scan_offsets) diff --git a/tests/envs/locomotion/go2w/test_go2w_manager_based_flat_cfg.py b/tests/envs/locomotion/go2w/test_go2w_manager_based_flat_cfg.py index 9e33af9ff..2c80f5cb3 100644 --- a/tests/envs/locomotion/go2w/test_go2w_manager_based_flat_cfg.py +++ b/tests/envs/locomotion/go2w/test_go2w_manager_based_flat_cfg.py @@ -187,17 +187,17 @@ def test_go2w_sac_drake_owner_uses_the_same_manager_contract() -> None: env_cfg.validate() -def test_go2w_flat_registry_is_manager_only_and_rough_owns_legacy_bridge() -> None: +def test_go2w_flat_and_rough_registries_are_manager_only() -> None: registry.ensure_registries() - from unilab.tasks.locomotion.go2w.joystick import Go2WJoystickCfg, Go2WJoystickEnv - from unilab.tasks.locomotion.go2w.rough import Go2WJoystickRoughCfg, Go2WJoystickRoughEnv assert registry.list_registered_envs()["Go2WJoystickFlat"] == { "config_factory": "ManagerBasedRlEnvCfg", "available_backends": ["mujoco", "motrix", "drake"], } - assert Go2WJoystickRoughCfg.__bases__ == (Go2WJoystickCfg,) - assert Go2WJoystickRoughEnv.__bases__ == (Go2WJoystickEnv,) + assert registry.list_registered_envs()["Go2WJoystickRough"] == { + "config_factory": "ManagerBasedRlEnvCfg", + "available_backends": ["mujoco", "motrix"], + } for legacy_override in ( {"reward_config": {}}, {"domain_rand": {"randomize_kp": False}}, diff --git a/tests/envs/locomotion/go2w/test_go2w_motor_control.py b/tests/envs/locomotion/go2w/test_go2w_motor_control.py deleted file mode 100644 index fbd49b175..000000000 --- a/tests/envs/locomotion/go2w/test_go2w_motor_control.py +++ /dev/null @@ -1,299 +0,0 @@ -from __future__ import annotations - -from types import SimpleNamespace -from typing import Any, cast - -import numpy as np -import pytest - -from unilab.base.np_env import NpEnvState -from unilab.tasks.locomotion.go2w.base import ( - DEFAULT_GO2W_ANGLES, - JOINT_SENSOR_PREFIXES, - NUM_GO2W_ACTIONS, - NUM_LEG_ACTIONS, - NUM_WHEEL_ACTIONS, - Go2WBaseEnv, - compute_go2w_motor_ctrl, -) -from unilab.tasks.locomotion.go2w.joystick import ( - Go2WJoystickCfg, - Go2WJoystickDomainRandomizationProvider, - Go2WJoystickEnv, - RewardConfig, - build_go2w_backend_reset_randomization, -) - - -def _reward_config() -> RewardConfig: - return RewardConfig(scales={"alive": 1.0}, tracking_sigma=0.25, base_height_target=0.3) - - -class _ConcreteGo2WBaseEnv(Go2WBaseEnv): - def update_state(self, state: NpEnvState) -> NpEnvState: - return state - - -def test_compute_go2w_motor_ctrl_converts_legs_and_tracks_wheel_velocity() -> None: - policy_ctrl = np.zeros((1, NUM_GO2W_ACTIONS), dtype=np.float64) - policy_ctrl[:, :NUM_LEG_ACTIONS] = 0.5 - policy_ctrl[:, NUM_LEG_ACTIONS:] = np.array([[2.0, -2.0, 20.0, -20.0]]) - joint_pos = np.zeros_like(policy_ctrl) - joint_vel = np.ones_like(policy_ctrl) * 0.1 - leg_kp = np.ones((1, NUM_LEG_ACTIONS), dtype=np.float64) * 10.0 - leg_kd = np.ones((1, NUM_LEG_ACTIONS), dtype=np.float64) * 2.0 - wheel_kd = np.ones((1, NUM_WHEEL_ACTIONS), dtype=np.float64) * 2.0 - ctrl_range = np.tile(np.array([-15.0, 15.0]), (NUM_GO2W_ACTIONS, 1)) - out = np.zeros_like(policy_ctrl) - - result = compute_go2w_motor_ctrl( - policy_ctrl, - joint_pos, - joint_vel, - leg_kp, - leg_kd, - wheel_kd, - ctrl_range[:, 0], - ctrl_range[:, 1], - out, - ) - - assert result is out - np.testing.assert_allclose(result[:, :NUM_LEG_ACTIONS], 4.8) - np.testing.assert_allclose(result[:, NUM_LEG_ACTIONS:], [[3.8, -4.2, 15.0, -15.0]]) - - -def test_go2w_joint_state_reads_named_sensors() -> None: - calls: list[str] = [] - - class FakeBackend: - def get_sensor_data(self, name: str) -> np.ndarray: - calls.append(name) - idx = len(calls) - return np.array([[float(idx)], [float(idx + 100)]], dtype=np.float32) - - def get_sensor_data_batch(self, names) -> np.ndarray: - return np.concatenate([self.get_sensor_data(name) for name in names], axis=1) - - env = cast(Any, object.__new__(_ConcreteGo2WBaseEnv)) - env._backend = FakeBackend() - env.default_angles = DEFAULT_GO2W_ANGLES.astype(np.float32) - - pos = env.get_dof_pos() - vel = env.get_dof_vel() - - assert pos.shape == (2, NUM_GO2W_ACTIONS) - assert vel.shape == (2, NUM_GO2W_ACTIONS) - assert calls[:NUM_GO2W_ACTIONS] == [f"{prefix}_pos" for prefix in JOINT_SENSOR_PREFIXES] - assert calls[NUM_GO2W_ACTIONS:] == [f"{prefix}_vel" for prefix in JOINT_SENSOR_PREFIXES] - - -def test_go2w_backend_reset_randomization_excludes_kp_kd_payload() -> None: - cfg = Go2WJoystickCfg() - cfg.domain_rand.randomize_base_mass = True - cfg.domain_rand.random_com = True - cfg.domain_rand.randomize_gravity = True - cfg.domain_rand.randomize_kp = True - cfg.domain_rand.randomize_kd = True - - payload = build_go2w_backend_reset_randomization(SimpleNamespace(cfg=cfg), num_reset=3) - - assert payload is not None - assert payload.kp is None - assert payload.kd is None - assert payload.base_mass_delta is not None - assert payload.base_com_offset is not None - assert payload.gravity is not None - - -def test_go2w_reset_plan_can_disable_initial_yaw_randomization() -> None: - from unilab.tasks.locomotion.common.terrain_spawn import BaseSpawnManager - - cfg = Go2WJoystickCfg(reward_config=_reward_config()) - cfg.domain_rand.randomize_init_yaw = False - env = SimpleNamespace( - cfg=cfg, - _spawn=BaseSpawnManager(), - _init_qpos=np.concatenate( - [np.array([0.0, 0.0, 0.42, 1.0, 0.0, 0.0, 0.0]), np.zeros(NUM_GO2W_ACTIONS)] - ), - _init_qvel=np.zeros(6 + NUM_GO2W_ACTIONS), - _num_action=NUM_GO2W_ACTIONS, - sample_reset_motor_gains=lambda num_reset: ( - np.ones((num_reset, NUM_LEG_ACTIONS)), - np.ones((num_reset, NUM_LEG_ACTIONS)), - ), - set_motor_gains=lambda env_ids, motor_kp, motor_kd: None, - ) - - plan = Go2WJoystickDomainRandomizationProvider().build_reset_plan( - env, np.array([0, 1], dtype=np.int32) - ) - - expected = np.tile(np.array([[1.0, 0.0, 0.0, 0.0]]), (2, 1)) - np.testing.assert_allclose(plan.qpos[:, 3:7], expected) - - -def test_go2w_init_does_not_pass_position_actuator_gains(monkeypatch: pytest.MonkeyPatch) -> None: - from unilab.tasks.locomotion.go2w import joystick as go2w_module - - captured: dict[str, Any] = {} - - class FakeBackend: - backend_type = "mujoco" - num_actuators = NUM_GO2W_ACTIONS - - def get_actuator_ctrl_range(self) -> np.ndarray: - return np.tile(np.array([-15.0, 15.0]), (NUM_GO2W_ACTIONS, 1)) - - def get_joint_range(self) -> np.ndarray: - return np.tile(np.array([-3.14, 3.14]), (NUM_GO2W_ACTIONS, 1)) - - def get_sensor_data(self, name: str) -> np.ndarray: - if name.endswith("_pos") or name.endswith("_vel"): - return np.zeros((2, 1), dtype=np.float32) - raise KeyError(name) - - def get_sensor_data_batch(self, names) -> np.ndarray: - return np.concatenate([self.get_sensor_data(name) for name in names], axis=1) - - def set_pre_step_control(self, fn) -> None: - captured["pre_step_control"] = fn - - def fake_create_backend(*args, **kwargs): - captured["args"] = args - captured["kwargs"] = kwargs - return FakeBackend() - - def fake_base_init(self, cfg, backend, num_envs): - self._cfg = cfg - self._backend = backend - self._num_envs = num_envs - self._num_action = NUM_GO2W_ACTIONS - self.default_angles = DEFAULT_GO2W_ANGLES.astype(np.float32) - - monkeypatch.setattr(go2w_module, "create_backend", fake_create_backend) - monkeypatch.setattr(Go2WBaseEnv, "__init__", fake_base_init) - monkeypatch.setattr(Go2WJoystickEnv, "_init_domain_randomization", lambda self, provider: None) - - cfg = Go2WJoystickCfg(reward_config=_reward_config()) - Go2WJoystickEnv(cfg, num_envs=2, backend_type="mujoco") - - assert captured["args"][0] == "mujoco" - assert "position_actuator_gains" not in captured["kwargs"] - assert callable(captured["pre_step_control"]) - - -def test_go2w_apply_action_maps_legs_to_targets_and_wheels_to_velocity_targets() -> None: - env = cast(Any, object.__new__(Go2WJoystickEnv)) - env._cfg = Go2WJoystickCfg(reward_config=_reward_config()) - env._np_dtype = np.float32 - env._num_envs = 1 - env._num_action = NUM_GO2W_ACTIONS - env.default_angles = DEFAULT_GO2W_ANGLES.astype(np.float32) - env._leg_action_scale = env._build_leg_action_scale() - state = NpEnvState( - obs={}, - reward=np.zeros((1,), dtype=np.float32), - terminated=np.zeros((1,), dtype=bool), - truncated=np.zeros((1,), dtype=bool), - info={}, - ) - action = np.ones((1, NUM_GO2W_ACTIONS), dtype=np.float32) - - ctrl = env.apply_action(action, state) - - expected_leg_targets = (DEFAULT_GO2W_ANGLES[:NUM_LEG_ACTIONS] + 0.25).reshape(1, -1) - np.testing.assert_allclose(ctrl[:, :NUM_LEG_ACTIONS], expected_leg_targets) - np.testing.assert_allclose(ctrl[:, NUM_LEG_ACTIONS:], 10.0) - np.testing.assert_allclose(state.info["current_actions"], action) - - -def test_go2w_compute_obs_accepts_reset_subset_info_shapes() -> None: - env = cast(Any, object.__new__(Go2WJoystickEnv)) - env._cfg = Go2WJoystickCfg(reward_config=_reward_config()) - env._num_envs = 2 - env._num_action = NUM_GO2W_ACTIONS - env.default_angles = DEFAULT_GO2W_ANGLES.astype(np.float32) - info = { - "current_actions": np.zeros((1, NUM_GO2W_ACTIONS), dtype=np.float32), - "commands": np.zeros((1, 3), dtype=np.float32), - "torques": np.ones((1, NUM_GO2W_ACTIONS), dtype=np.float32), - } - linvel = np.zeros((1, 3), dtype=np.float32) - gyro = np.zeros((1, 3), dtype=np.float32) - gravity = np.zeros((1, 3), dtype=np.float32) - dof_pos = DEFAULT_GO2W_ANGLES.reshape(1, -1).astype(np.float32) - dof_vel = np.zeros((1, NUM_GO2W_ACTIONS), dtype=np.float32) - - obs = env._compute_obs(info, linvel, gyro, gravity, dof_pos, dof_vel) - - assert obs["obs"].shape == (1, env.obs_groups_spec["obs"]) - assert obs["critic"].shape == (1, env.obs_groups_spec["critic"]) - np.testing.assert_allclose(obs["critic"][:, -NUM_GO2W_ACTIONS:], 1.0) - - -def test_go2w_heading_command_updates_yaw_rate_from_heading_error() -> None: - class FakeBackend: - def get_base_quat(self) -> np.ndarray: - return np.array([[1.0, 0.0, 0.0, 0.0]], dtype=np.float32) - - env = cast(Any, object.__new__(Go2WJoystickEnv)) - env._cfg = Go2WJoystickCfg(reward_config=_reward_config()) - env._cfg.commands.heading_command = True - env._num_envs = 1 - env._backend = FakeBackend() - info = { - "commands": np.zeros((1, 3), dtype=np.float32), - "heading_commands": np.array([np.pi / 2.0], dtype=np.float32), - "steps": np.zeros((1,), dtype=np.uint32), - } - - env._update_commands(info) - - np.testing.assert_allclose(info["commands"][:, 2], [0.25 * np.pi], rtol=1e-6) - - -def test_go2w_pre_step_motor_control_reads_from_passed_backend() -> None: - calls: list[str] = [] - - class FakeBackend: - def get_sensor_data(self, name: str) -> np.ndarray: - calls.append(name) - if name.endswith("_pos"): - return np.zeros((1, 1), dtype=np.float32) - if name.endswith("_vel"): - return np.ones((1, 1), dtype=np.float32) * 0.1 - raise KeyError(name) - - def get_sensor_data_batch(self, names) -> np.ndarray: - return np.concatenate([self.get_sensor_data(name) for name in names], axis=1) - - class PoisonBackend: - def get_sensor_data(self, name: str) -> np.ndarray: - raise AssertionError(f"unexpected env backend read: {name}") - - def get_sensor_data_batch(self, names) -> np.ndarray: - raise AssertionError(f"unexpected env backend read: {names}") - - env = cast(Any, object.__new__(Go2WJoystickEnv)) - env._backend = PoisonBackend() - env.default_angles = DEFAULT_GO2W_ANGLES.astype(np.float32) - env._np_dtype = np.float32 - env._motor_kp = np.ones((1, NUM_LEG_ACTIONS), dtype=np.float64) * 10.0 - env._motor_kd = np.ones((1, NUM_LEG_ACTIONS), dtype=np.float64) * 2.0 - env._wheel_kd = np.ones((1, NUM_WHEEL_ACTIONS), dtype=np.float64) * 2.0 - ctrl_range = np.tile(np.array([-15.0, 15.0]), (NUM_GO2W_ACTIONS, 1)) - env._ctrl_lower = ctrl_range[:, 0].astype(np.float32) - env._ctrl_upper = ctrl_range[:, 1].astype(np.float32) - env._last_motor_ctrl = np.zeros((1, NUM_GO2W_ACTIONS), dtype=np.float32) - policy_ctrl = np.zeros((1, NUM_GO2W_ACTIONS), dtype=np.float32) - policy_ctrl[:, :NUM_LEG_ACTIONS] = 0.5 - policy_ctrl[:, NUM_LEG_ACTIONS:] = np.array([[1.0, -1.0, 2.0, -2.0]], dtype=np.float32) - - motor_ctrl = env._pre_step_motor_control(FakeBackend(), policy_ctrl) - - assert calls[:NUM_GO2W_ACTIONS] == [f"{prefix}_pos" for prefix in JOINT_SENSOR_PREFIXES] - assert calls[NUM_GO2W_ACTIONS:] == [f"{prefix}_vel" for prefix in JOINT_SENSOR_PREFIXES] - np.testing.assert_allclose(motor_ctrl[:, :NUM_LEG_ACTIONS], 4.8) - np.testing.assert_allclose(motor_ctrl[:, NUM_LEG_ACTIONS:], [[1.8, -2.2, 3.8, -4.2]]) diff --git a/tests/envs/locomotion/test_go2_joystick_stand_still.py b/tests/envs/locomotion/test_go2_joystick_stand_still.py deleted file mode 100644 index 26850eddf..000000000 --- a/tests/envs/locomotion/test_go2_joystick_stand_still.py +++ /dev/null @@ -1,29 +0,0 @@ -"""Go2 legacy-base regressions around the A2 Manager-Based migration.""" - -from __future__ import annotations - -from types import SimpleNamespace - -import numpy as np - - -def test_go2_advance_phase_is_unconditional(): - """Go2WalkTask advances the gait clock every step regardless of command — - the A2 standing freeze must not have leaked into the Go2 owner.""" - from unilab.tasks.locomotion.go2.joystick import Go2WalkTask - - stub = SimpleNamespace(_cfg=SimpleNamespace(ctrl_dt=0.02), gait_frequency=2.0) - phase = np.array([0.3, 0.3]) - out = Go2WalkTask._advance_phase(stub, phase) - expected = np.fmod(phase + 0.02 * 2.0, 1.0) - np.testing.assert_allclose(out, expected) - - -def test_go2_reward_config_has_no_command_threshold(): - """A2's threshold stays in Hydra Manager terms, not the shared legacy config.""" - import dataclasses - - from unilab.tasks.locomotion.go2.joystick import RewardConfig - - names = {f.name for f in dataclasses.fields(RewardConfig)} - assert "command_threshold" not in names diff --git a/tests/envs/locomotion/test_go2_rough_height_scan.py b/tests/envs/locomotion/test_go2_rough_height_scan.py deleted file mode 100644 index 8776f78f3..000000000 --- a/tests/envs/locomotion/test_go2_rough_height_scan.py +++ /dev/null @@ -1,77 +0,0 @@ -from __future__ import annotations - -from typing import Any - -import numpy as np - -from unilab.tasks.locomotion.common.height_scan import height_scan_offsets as _height_scan_offsets -from unilab.tasks.locomotion.go2.rough import Go2JoystickRoughEnv - - -def test_go2_rough_height_scan_uses_backend_native_sampling() -> None: - class FakeHeightScanner: - def __init__(self, heights: np.ndarray) -> None: - self.calls = 0 - self.heights = heights - - def scan(self) -> np.ndarray: - self.calls += 1 - return self.heights - - class FakeBackend: - def __init__(self) -> None: - self.scanner_calls: list[dict[str, Any]] = [] - self.base_pos = np.asarray([[0.0, 0.0, 0.6], [1.0, 0.0, 0.7]], dtype=np.float32) - self.heights = np.asarray([[0.1, 0.2], [0.3, 0.4]], dtype=np.float32) - self.scanner = FakeHeightScanner(self.heights) - - def get_base_pos(self) -> np.ndarray: - return self.base_pos - - def create_hfield_scanner(self, **kwargs: Any) -> FakeHeightScanner: - self.scanner_calls.append(kwargs) - return self.scanner - - env = object.__new__(Go2JoystickRoughEnv) - fake_backend = FakeBackend() - env._backend = fake_backend - env._height_scan_dim = 2 - env._height_scan_hfield_geom_id = 7 - env._height_scan_frame_body_id = 3 - env._height_scan_offsets = np.asarray([[0.0, 0.0], [0.1, -0.1]], dtype=np.float64) - env._height_scan_sensor = fake_backend.create_hfield_scanner( - hfield_geom_id=env._height_scan_hfield_geom_id, - offsets=env._height_scan_offsets, - frame_body_id=env._height_scan_frame_body_id, - alignment="yaw", - output="height", - ) - - raw_heights, base_pos = env._raw_height_scan_obs(num_obs=2) - - np.testing.assert_array_equal(raw_heights, fake_backend.heights) - np.testing.assert_array_equal(base_pos, fake_backend.base_pos) - assert fake_backend.scanner.calls == 1 - assert len(fake_backend.scanner_calls) == 1 - call = fake_backend.scanner_calls[0] - assert call["hfield_geom_id"] == 7 - assert call["frame_body_id"] == 3 - assert call["alignment"] == "yaw" - assert call["output"] == "height" - np.testing.assert_array_equal(call["offsets"], env._height_scan_offsets) - - -def test_go2_rough_height_scan_offsets_are_grid_ordered() -> None: - offsets = _height_scan_offsets(points_x=[-0.1, 0.2], points_y=[-0.3, 0.4]) - - expected = np.asarray( - [ - [-0.1, -0.3], - [-0.1, 0.4], - [0.2, -0.3], - [0.2, 0.4], - ], - dtype=np.float64, - ) - np.testing.assert_array_equal(offsets, expected) - assert offsets.flags.c_contiguous diff --git a/tests/envs/locomotion/test_go2_terrain_spawn.py b/tests/envs/locomotion/test_go2_terrain_spawn.py index b977e8243..ec9dd468e 100644 --- a/tests/envs/locomotion/test_go2_terrain_spawn.py +++ b/tests/envs/locomotion/test_go2_terrain_spawn.py @@ -1,365 +1,224 @@ -"""Integration tests for Go2 + TerrainSpawnManager.""" +"""Production contracts for the Manager-Based quadruped rough family.""" from __future__ import annotations +import inspect +from pathlib import Path from types import SimpleNamespace +from typing import Any, cast import numpy as np import pytest - -from unilab.tasks.locomotion.common.terrain_spawn import ( - TerrainCurriculumCfg, - TerrainSpawnManager, +from hydra import compose, initialize_config_dir +from hydra.core.global_hydra import GlobalHydra +from omegaconf import DictConfig + +from unilab.base import registry +from unilab.base.config_materialization import apply_cfg_overrides +from unilab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg +from unilab.tasks.locomotion.common.rough_manager_terms import ( + QuadrupedRoughTerrainCfg, + RoughHeightScan, + RoughJointPositionAction, + RoughTerrainCurriculum, + RoughTerrainOutOfBounds, + RoughTerrainReset, ) - - -def _class_path(obj) -> str: - cls = type(obj) - return f"{cls.__module__}.{cls.__name__}" - - -def _configure_small_terrain(cfg, *, seed: int = 0) -> None: - assert cfg.scene.terrain is not None - generator = cfg.scene.terrain.generator - assert generator is not None - generator.num_rows = 3 - generator.num_cols = 3 - generator.border_width = 0.0 - generator.add_lights = False - generator.seed = seed - - -def _rough_cfg(*, curriculum_enabled: bool = False, seed: int = 0): - from unilab.tasks.locomotion.go2.rough import Go2JoystickRoughCfg, RoughRewardConfig - - cfg = Go2JoystickRoughCfg( - reward_config=RoughRewardConfig(scales={}, tracking_sigma=0.25, base_height_target=0.3) +from unilab.tasks.locomotion.go2w.manager_terms import Go2WMixedAction +from unilab.training.backend_adapter import BackendAdapter +from unilab.training.sim2sim import extract_contract_snapshot + +ROOT_DIR = Path(__file__).parents[3] +CONF_DIR = ROOT_DIR / "conf" / "ppo" + +_OWNER_CASES = tuple( + pytest.param(task_id, task_name, backend, *dims, id=f"{task_id.split('_')[0]}-{backend}") + for task_id, task_name, dims in ( + ("go1_joystick_rough", "Go1JoystickRough", (45, 235, 12)), + ("go2_joystick_rough", "Go2JoystickRough", (45, 235, 12)), + ("go2w_joystick_rough", "Go2WJoystickRough", (53, 243, 16)), ) - _configure_small_terrain(cfg, seed=seed) - cfg.terrain_curriculum = TerrainCurriculumCfg(enabled=curriculum_enabled, seed=seed) - return cfg - - -def test_terrain_spawn_attached_when_rough(): - from unilab.tasks.locomotion.go2.rough import Go2JoystickRoughEnv - - cfg = _rough_cfg() - env = Go2JoystickRoughEnv(cfg, num_envs=4, backend_type="mujoco") - try: - terrain_data = env._backend.get_terrain_spawn_data() - assert terrain_data is not None - assert terrain_data.sample_height is not None - assert terrain_data.terrain_origins.shape == (3, 3, 3) - assert not terrain_data.terrain_origins.flags.writeable - assert _class_path(env._spawn) == ( - "unilab.tasks.locomotion.common.terrain_spawn.TerrainSpawnManager" - ) - assert env._scene_terrain_origins is not None - assert env._scene_terrain_origins.shape == (3, 3, 3) - finally: - env.close() - - -def test_terrain_spawn_attached_when_rough_motrix(): - pytest.importorskip("motrixsim") - - from unilab.tasks.locomotion.go2.joystick import Go2WalkTask - - cfg = _rough_cfg() - cfg.domain_rand.randomize_kp = False - cfg.domain_rand.randomize_kd = False - env = Go2WalkTask(cfg, num_envs=2, backend_type="motrix") - try: - terrain_data = env._backend.get_terrain_spawn_data() - assert terrain_data is not None - assert terrain_data.sample_height is not None - assert terrain_data.terrain_origins.shape == (3, 3, 3) - assert not terrain_data.terrain_origins.flags.writeable - assert _class_path(env._spawn) == ( - "unilab.tasks.locomotion.common.terrain_spawn.TerrainSpawnManager" - ) - assert env._scene_terrain_origins is not None - assert env._scene_terrain_origins.shape == (3, 3, 3) - state = env.init_state() - assert state.obs["obs"].shape == (2, 49) - finally: - env.close() + for backend in ("mujoco", "motrix") +) -def test_go1_rough_initialization_and_reset_use_backend_terrain_contract(): - from unilab.tasks.locomotion.go1.rough import ( - Go1JoystickRoughCfg, - Go1JoystickRoughEnv, - RoughRewardConfig, +def _compose(task_id: str, backend: str, extra: tuple[str, ...] = ()) -> DictConfig: + GlobalHydra.instance().clear() + with initialize_config_dir(config_dir=str(CONF_DIR), version_base="1.3"): + return compose("config", overrides=[f"task={task_id}/{backend}", *extra]) + + +def _materialize( + task_id: str, + task_name: str, + backend: str, + extra: tuple[str, ...] = (), +) -> tuple[DictConfig, ManagerBasedRlEnvCfg, dict[str, Any]]: + hydra_cfg = _compose(task_id, backend, extra) + override = BackendAdapter(hydra_cfg, root_dir=ROOT_DIR).build_task_env_cfg_override() + env_cfg = registry.materialize_env_config(task_name) + assert isinstance(env_cfg, ManagerBasedRlEnvCfg) + apply_cfg_overrides(env_cfg, override) + env_cfg.validate() + return hydra_cfg, env_cfg, override + + +@pytest.mark.parametrize("task_id,task_name,backend,policy_dim,critic_dim,action_dim", _OWNER_CASES) +def test_rough_registry_executes_real_backend_contract( + task_id: str, + task_name: str, + backend: str, + policy_dim: int, + critic_dim: int, + action_dim: int, +) -> None: + if backend == "motrix": + pytest.importorskip("motrixsim") + registry.ensure_registries() + hydra_cfg, env_cfg, override = _materialize(task_id, task_name, backend) + assert (hydra_cfg.training.task_name, hydra_cfg.training.sim_backend) == (task_name, backend) + assert hydra_cfg.env.commands.twist.planar_dead_zone == pytest.approx(0.08) + override["commands"]["twist"]["ranges"]["lin_vel_x"] = [0.04, 0.04] + override["commands"]["twist"]["ranges"]["lin_vel_y"] = [0.0, 0.0] + assert env_cfg.scene is not None and env_cfg.scene.terrain is not None + terrain = env_cfg.scene.terrain.generator + assert isinstance(terrain, QuadrupedRoughTerrainCfg) + assert (terrain.num_rows, terrain.num_cols, len(terrain.sub_terrains)) == (6, 6, 7) + assert terrain.horizontal_scale == pytest.approx(0.1 if "go2w" in task_id else 0.2) + assert registry.list_registered_envs()[task_name]["config_factory"] == "ManagerBasedRlEnvCfg" + env = cast( + ManagerBasedRlEnv, + registry.make( + task_name, + sim_backend=backend, + env_cfg_override=override, + num_envs=2, + ), ) - - cfg = Go1JoystickRoughCfg( - reward_config=RoughRewardConfig(scales={}, tracking_sigma=0.25, base_height_target=0.3) - ) - _configure_small_terrain(cfg) - cfg.terrain_scan.enabled = False - cfg.domain_rand.randomize_kp = False - cfg.domain_rand.randomize_kd = False - - env = Go1JoystickRoughEnv(cfg, num_envs=2, backend_type="mujoco") try: - terrain_data = env._backend.get_terrain_spawn_data() - assert terrain_data is not None - assert terrain_data.sample_height is not None - assert env._scene_terrain_origins is terrain_data.terrain_origins - assert env._spawn._sample_height is terrain_data.sample_height - state = env.init_state() - assert state.obs["obs"].shape == (2, 45) - finally: - env.close() - - -def test_go2w_rough_initialization_and_reset_use_backend_terrain_contract(): - from unilab.tasks.locomotion.go2w.joystick import RewardConfig - from unilab.tasks.locomotion.go2w.rough import Go2WJoystickRoughCfg, Go2WJoystickRoughEnv + obs, _ = env.reset(seed=7) + np.testing.assert_allclose(env.command_manager.get_command("twist")[:, :2], 0.0) + assert obs["obs"].shape == (2, policy_dim) + assert obs["critic"].shape == (2, critic_dim) + assert env.obs_groups_spec == {"obs": policy_dim, "critic": critic_dim} + assert env.action_space.shape == (action_dim,) + np.testing.assert_array_equal(env.scene.env_origins, np.zeros((2, 3))) - cfg = Go2WJoystickRoughCfg( - reward_config=RewardConfig(scales={}, tracking_sigma=0.25, base_height_target=0.3) - ) - _configure_small_terrain(cfg) - cfg.terrain_scan.enabled = False - cfg.domain_rand.randomize_kp = False - cfg.domain_rand.randomize_kd = False - - env = Go2WJoystickRoughEnv(cfg, num_envs=2, backend_type="mujoco") - try: terrain_data = env._backend.get_terrain_spawn_data() assert terrain_data is not None assert terrain_data.sample_height is not None - np.testing.assert_array_equal(env._spawn._terrain_origins, terrain_data.terrain_origins) - assert env._spawn._sample_height is terrain_data.sample_height - state = env.init_state() - assert state.obs["obs"].shape == (2, 53) + assert terrain_data.terrain_origins.shape == (6, 6, 3) + reset_term = env.event_manager.get_term_cfg("terrain_root_state").func + assert isinstance(reset_term, RoughTerrainReset) + spawn = reset_term.spawn_manager + ids = np.arange(2, dtype=np.int32) + base_pos = env.scene["robot"].data.root_link_pos_w + assigned_origins = spawn.origins_for(ids) + assert np.all(np.abs(base_pos[:, :2] - assigned_origins[:, :2]) <= 0.5) + surface_height = terrain_data.sample_height(base_pos[:, :2]) + assert np.all(base_pos[:, 2] - surface_height > 0.25) + + action = env.action_manager.get_term("motor" if "go2w" in task_id else "joint_pos") + saturated = np.full((2, action_dim), 200.0, dtype=np.float32) + action.process_actions(saturated) + np.testing.assert_allclose(action.raw_action, 100.0) + if isinstance(action, RoughJointPositionAction): + scale = np.asarray(action.scale)[0] + for name, value in zip(action.target_names, scale, strict=True): + expected = 0.125 if "_hip_joint" in name else 0.25 + assert value == pytest.approx(expected) + assert action.cfg.clip_actions == pytest.approx(100.0) + else: + assert isinstance(action, Go2WMixedAction) + assert action.cfg.leg_action_scale == pytest.approx(0.25) + assert action.cfg.hip_action_scale == pytest.approx(0.125) + assert action.cfg.wheel_action_scale == pytest.approx(5.0) + assert action.cfg.clip_actions == pytest.approx(100.0) + + state = env.step(np.zeros((2, action_dim), dtype=np.float32)) + assert state.obs["obs"].shape == (2, policy_dim) + assert state.obs["critic"].shape == (2, critic_dim) + for value in (*state.obs.values(), state.reward): + assert np.isfinite(value).all() finally: env.close() -def test_terrain_spawn_samples_height_after_xy_jitter(): - class FakeSurface: - def sample_height(self, xy): - xy = np.asarray(xy, dtype=np.float64) - return xy[:, 0] * 0.5 + xy[:, 1] * 0.25 - - origins = np.zeros((1, 1, 3), dtype=np.float64) - cfg = TerrainCurriculumCfg(spawn_height_margin=0.05) - surface = FakeSurface() - sample_height = surface.sample_height - spawn = TerrainSpawnManager( - 1, - origins, - cell_size=1.0, - cfg=cfg, - sample_height=sample_height, - ) - assert spawn._sample_height is sample_height - - qpos_xyz = np.asarray([[0.2, 0.4, 0.42]], dtype=np.float64) - spawned = spawn.apply_spawn(np.asarray([0], dtype=np.int32), qpos_xyz) - - assert spawned[0, 0] == pytest.approx(0.2) - assert spawned[0, 1] == pytest.approx(0.4) - assert spawned[0, 2] == pytest.approx(0.2 * 0.5 + 0.4 * 0.25 + 0.42 + 0.05) - - -def test_terrain_spawn_rejects_non_callable_height_sampler_on_init(): - with pytest.raises(TypeError, match="sample_height must be callable"): - TerrainSpawnManager( - 1, - np.zeros((1, 1, 3), dtype=np.float64), - cell_size=1.0, - cfg=TerrainCurriculumCfg(), - sample_height=object(), # type: ignore[arg-type] - ) - - -def test_go2_rough_base_height_reward_uses_terrain_relative_height(): - from unilab.tasks.locomotion.go2.joystick import Go2WalkTask - - class FakeBackend: - def get_base_pos(self): - return np.asarray( - [ - [1.0, 2.0, 1.25], - [-1.0, 0.5, -0.2], - ], - dtype=np.float32, - ) - - class FakeSurface: - def sample_height(self, xy): - xy = np.asarray(xy, dtype=np.float64) - return np.asarray([0.75, -0.55], dtype=np.float64) - - surface = FakeSurface() - env = Go2WalkTask.__new__(Go2WalkTask) - env._backend = FakeBackend() - env._terrain_surface_sample_height = surface.sample_height - - np.testing.assert_allclose(env._reward_base_height_values(), np.asarray([0.5, 0.35])) - - -def test_go2_rough_pd_torque_estimate_returns_dof_order(): - from unilab.tasks.locomotion.go2.rough import ( - GO2_ACTUATOR_TO_DOF_INDICES, - Go2JoystickRoughEnv, - ) - - env = Go2JoystickRoughEnv.__new__(Go2JoystickRoughEnv) - env._num_action = 12 - env._cfg = SimpleNamespace( - control_config=SimpleNamespace(Kp=2.0, Kd=0.5, simulate_action_latency=False) - ) - env._action_scale = np.ones((12,), dtype=np.float32) - env.default_angles = np.arange(12, dtype=np.float32) - env._default_angles_actuator = env.default_angles[GO2_ACTUATOR_TO_DOF_INDICES] - dof_pos = np.zeros((1, 12), dtype=np.float32) - dof_vel = np.zeros((1, 12), dtype=np.float32) - info = {"current_actions": np.zeros((1, 12), dtype=np.float32)} - - torques = env._estimate_pd_torques(info, dof_pos, dof_vel) - - np.testing.assert_allclose(torques, 2.0 * env.default_angles[None, :]) - - -def test_default_spawn_used_when_flat(): - from unilab.tasks.locomotion.go2.joystick import ( - Go2JoystickCfg, - Go2WalkTask, - RewardConfig, +@pytest.mark.parametrize( + ("task_id", "task_name"), + ( + ("go1_joystick_rough", "Go1JoystickRough"), + ("go2_joystick_rough", "Go2JoystickRough"), + ("go2w_joystick_rough", "Go2WJoystickRough"), + ), +) +def test_rough_sim2sim_snapshot_matches_across_backends(task_id: str, task_name: str) -> None: + registry.ensure_registries() + mujoco_cfg, _, _ = _materialize(task_id, task_name, "mujoco") + motrix_cfg, _, _ = _materialize(task_id, task_name, "motrix") + assert extract_contract_snapshot(mujoco_cfg) == extract_contract_snapshot(motrix_cfg) + + +def test_curriculum_updates_before_next_spawn_and_oob_fails_closed() -> None: + registry.ensure_registries() + extra = ( + "env.scene.terrain.generator.curriculum=true", + "env.scene.terrain.generator.num_rows=3", ) - - cfg = Go2JoystickCfg( - reward_config=RewardConfig(scales={}, tracking_sigma=0.25, base_height_target=0.3) + _, _, override = _materialize("go2_joystick_rough", "Go2JoystickRough", "mujoco", extra) + env = cast( + ManagerBasedRlEnv, + registry.make( + "Go2JoystickRough", + sim_backend="mujoco", + env_cfg_override=override, + num_envs=2, + ), ) - env = Go2WalkTask(cfg, num_envs=4, backend_type="mujoco") try: - assert env._backend.get_terrain_spawn_data() is None - assert _class_path(env._spawn) == ( - "unilab.tasks.locomotion.common.terrain_spawn.BaseSpawnManager" + env.reset(seed=7) + reset_term = env.event_manager.get_term_cfg("terrain_root_state").func + assert isinstance(reset_term, RoughTerrainReset) + spawn = reset_term.spawn_manager + np.testing.assert_array_equal(spawn.levels, np.zeros(2, dtype=np.int32)) + + current = env.scene["robot"].data.root_link_pos_w.copy() + spawn._episode_start_xyz[0, :2] = current[0, :2] - np.asarray([5.0, 0.0]) + env.reset_buf[0] = True + _, info = env.reset(env_ids=np.asarray([0], dtype=np.int32)) + + assert spawn.levels[0] == 1 + assert info["log"]["Curriculum/terrain_levels/num_promoted"] == 1 + next_pos = env.scene["robot"].data.root_link_pos_w[0] + next_origin = spawn.origins_for(np.asarray([0], dtype=np.int32))[0] + assert np.all(np.abs(next_pos[:2] - next_origin[:2]) <= 0.5) + np.testing.assert_allclose(spawn._episode_start_xyz[0], next_pos) + + oob = env.termination_manager.get_term_cfg("terrain_out_of_bounds").func + assert isinstance(oob, RoughTerrainOutOfBounds) + oob._asset = SimpleNamespace( + data=SimpleNamespace( + root_link_pos_w=np.asarray([[0.0, 0.0, 0.5], [oob._half_width + 1.0, 0.0, 0.5]]) + ) ) - assert env._scene_terrain_origins is None - # Origins are zeros (flat scene needs no spread; per-env xy jitter still applies). - np.testing.assert_array_equal(env._spawn.origins_for(np.arange(4)), np.zeros((4, 3))) - finally: - env.close() - - -def test_curriculum_disabled_distributes_levels_uniformly(): - from unilab.tasks.locomotion.go2.rough import Go2JoystickRoughEnv - - cfg = _rough_cfg(curriculum_enabled=False, seed=0) - env = Go2JoystickRoughEnv(cfg, num_envs=64, backend_type="mujoco") - try: - sm = env._spawn - assert sm is not None - assert sm.levels.min() == 0 - assert sm.levels.max() == 2 - assert sm.type_cols.min() >= 0 - assert sm.type_cols.max() <= 2 - finally: - env.close() - - -def test_curriculum_enabled_levels_start_at_zero(): - from unilab.tasks.locomotion.go2.rough import Go2JoystickRoughEnv - - cfg = _rough_cfg(curriculum_enabled=True, seed=0) - env = Go2JoystickRoughEnv(cfg, num_envs=8, backend_type="mujoco") - try: - sm = env._spawn - assert sm is not None - assert np.all(sm.levels == 0) - finally: - env.close() - - -def test_reset_qpos_xy_matches_terrain_origins(): - from unilab.tasks.locomotion.go2.rough import Go2JoystickRoughEnv - - cfg = _rough_cfg(curriculum_enabled=False, seed=0) - env = Go2JoystickRoughEnv(cfg, num_envs=4, backend_type="mujoco") - try: - sm = env._spawn - assert sm is not None - env.init_state() - base_pos = env._backend.get_base_pos() - rows = sm.levels - cols = sm.type_cols - expected_xy = sm._terrain_origins[rows, cols, :2] - # Reset adds a uniform [-0.5, 0.5] xy jitter on top of the spawn xy. - diff = base_pos[:, :2] - expected_xy - assert np.all(np.abs(diff) < 0.6) - finally: - env.close() - - -def test_rough_reset_spawns_above_sampled_terrain(): - from unilab.tasks.locomotion.go2.rough import Go2JoystickRoughEnv - - cfg = _rough_cfg(curriculum_enabled=False, seed=0) - env = Go2JoystickRoughEnv(cfg, num_envs=64, backend_type="mujoco") - try: - env.init_state() - raw_heights, _ = env._raw_height_scan_obs(env.num_envs) - assert raw_heights is not None - center_heights = raw_heights[:, raw_heights.shape[1] // 2] - clearance = env._backend.get_base_pos()[:, 2] - center_heights - assert np.all(clearance > 0.25) + np.testing.assert_array_equal(oob(env), np.asarray([False, True])) finally: env.close() -def test_curriculum_logs_appear_after_done(): - from unilab.tasks.locomotion.go2.rough import Go2JoystickRoughEnv - - cfg = _rough_cfg(curriculum_enabled=True, seed=0) - env = Go2JoystickRoughEnv(cfg, num_envs=4, backend_type="mujoco") - try: - state = env.init_state() - env.apply_action(np.zeros((4, 12), dtype=np.float32), state) - state.truncated[:] = True - out = env.update_state(state) - log = out.info.get("log", {}) - for key in ( - "terrain_curriculum/mean_level", - "terrain_curriculum/max_level", - "terrain_curriculum/mean_walked", - "terrain_curriculum/num_promoted", - "terrain_curriculum/num_demoted", - "terrain_curriculum/num_skipped", +def test_rough_hot_paths_use_only_cached_runtime_objects() -> None: + for term in ( + RoughTerrainReset, + RoughTerrainCurriculum, + RoughTerrainOutOfBounds, + RoughHeightScan, + ): + source = inspect.getsource(term.__call__) + for forbidden in ( + "ASSETS_ROOT_PATH", + "model_file", + "getattr(", + "hasattr(", + "._backend", ): - assert key in log - finally: - env.close() - - -@pytest.mark.parametrize("preset", ["flat", "rough"]) -def test_episode_start_recorded_after_reset(preset): - from unilab.tasks.locomotion.go2.joystick import ( - Go2JoystickCfg, - Go2WalkTask, - RewardConfig, - ) - from unilab.tasks.locomotion.go2.rough import Go2JoystickRoughEnv - - if preset == "flat": - cfg = Go2JoystickCfg( - reward_config=RewardConfig(scales={}, tracking_sigma=0.25, base_height_target=0.3) - ) - env_cls = Go2WalkTask - else: - cfg = _rough_cfg(curriculum_enabled=False, seed=0) - env_cls = Go2JoystickRoughEnv - env = env_cls(cfg, num_envs=4, backend_type="mujoco") - try: - env.init_state() - sm = env._spawn - if _class_path(sm) == "unilab.tasks.locomotion.common.terrain_spawn.TerrainSpawnManager": - assert np.all(sm._has_started) - finally: - env.close() + assert forbidden not in source diff --git a/tests/envs/test_env_configs.py b/tests/envs/test_env_configs.py index 7c9279249..aa470d69d 100644 --- a/tests/envs/test_env_configs.py +++ b/tests/envs/test_env_configs.py @@ -2140,8 +2140,6 @@ def test_g1_motion_tracking_clip_end_does_not_override_true_termination(): # Environments that don't need special config overrides _STANDARD_ENVS = [ - "Go1JoystickRough", - "Go2WJoystickRough", "G1WalkFlat", "G1WalkRough", "AllegroInhandRotation", diff --git a/tests/envs/test_go1_obs_noise.py b/tests/envs/test_go1_obs_noise.py deleted file mode 100644 index a70eed7ae..000000000 --- a/tests/envs/test_go1_obs_noise.py +++ /dev/null @@ -1,77 +0,0 @@ -"""Tests for Go1 per-step observation noise.""" - -from __future__ import annotations - -import numpy as np - -from unilab.tasks.locomotion.go1.base import Go1BaseCfg, Go1BaseEnv, NoiseConfig - - -class _ConcreteGo1Env(Go1BaseEnv): - """Minimal concrete subclass — only needed to satisfy the ABC.""" - - def update_state(self, state): - raise NotImplementedError - - -def _make_env(level: float) -> Go1BaseEnv: - cfg = Go1BaseCfg(noise_config=NoiseConfig(level=level)) - env = object.__new__(_ConcreteGo1Env) - env._cfg = cfg - return env - - -class TestObsNoise: - def test_noise_applied_when_level_positive(self): - env = _make_env(level=1.0) - data = np.ones((4, 10), dtype=np.float32) - cfg = env._cfg.noise_config - - results = [env._obs_noise(data.copy(), cfg.scale_joint_angle) for _ in range(5)] - # At least one result should differ from the original - assert any(not np.allclose(r, data) for r in results) - - def test_no_noise_when_level_zero(self): - env = _make_env(level=0.0) - data = np.ones((4, 10), dtype=np.float32) - cfg = env._cfg.noise_config - - result = env._obs_noise(data.copy(), cfg.scale_joint_angle) - np.testing.assert_array_equal(result, data) - - def test_noise_bounded_by_level_times_scale(self): - env = _make_env(level=1.0) - data = np.zeros((128, 29), dtype=np.float32) - scale = 0.2 - result = env._obs_noise(data.copy(), scale) - # uniform[-1,1] * 1.0 * 0.2 => bounded by [-0.2, 0.2] - assert np.all(result >= -scale) - assert np.all(result <= scale) - - def test_noise_scales_with_level(self): - env_half = _make_env(level=0.5) - env_full = _make_env(level=1.0) - data = np.zeros((1024, 10), dtype=np.float32) - scale = 1.0 - - np.random.seed(0) - r_half = env_half._obs_noise(data.copy(), scale) - np.random.seed(0) - r_full = env_full._obs_noise(data.copy(), scale) - - # Same random seed: full-level noise should be exactly 2x half-level noise - np.testing.assert_allclose(r_full, r_half * 2.0) - - def test_noise_preserves_dtype(self): - for dt in [np.float32, np.float64]: - env = _make_env(level=1.0) - data = np.ones((4, 5), dtype=dt) - result = env._obs_noise(data, 0.1) - assert result.dtype == dt - - def test_noise_preserves_shape(self): - env = _make_env(level=1.0) - for shape in [(1, 3), (64, 29), (1024, 10)]: - data = np.zeros(shape, dtype=np.float32) - result = env._obs_noise(data, 0.1) - assert result.shape == shape diff --git a/tests/envs/test_go2_obs_noise.py b/tests/envs/test_go2_obs_noise.py deleted file mode 100644 index 0eb1840a2..000000000 --- a/tests/envs/test_go2_obs_noise.py +++ /dev/null @@ -1,74 +0,0 @@ -"""Tests for Go2 per-step observation noise.""" - -from __future__ import annotations - -import numpy as np - -from unilab.tasks.locomotion.go2.base import Go2BaseCfg, Go2BaseEnv, NoiseConfig - - -class _ConcreteGo2Env(Go2BaseEnv): - """Minimal concrete subclass — only needed to satisfy the ABC.""" - - def update_state(self, state): - raise NotImplementedError - - -def _make_env(level: float) -> Go2BaseEnv: - cfg = Go2BaseCfg(noise_config=NoiseConfig(level=level)) - env = object.__new__(_ConcreteGo2Env) - env._cfg = cfg - return env - - -class TestObsNoise: - def test_noise_applied_when_level_positive(self): - env = _make_env(level=1.0) - data = np.ones((4, 10), dtype=np.float32) - cfg = env._cfg.noise_config - - results = [env._obs_noise(data.copy(), cfg.scale_joint_angle) for _ in range(5)] - assert any(not np.allclose(r, data) for r in results) - - def test_no_noise_when_level_zero(self): - env = _make_env(level=0.0) - data = np.ones((4, 10), dtype=np.float32) - cfg = env._cfg.noise_config - - result = env._obs_noise(data.copy(), cfg.scale_joint_angle) - np.testing.assert_array_equal(result, data) - - def test_noise_bounded_by_level_times_scale(self): - env = _make_env(level=1.0) - data = np.zeros((128, 29), dtype=np.float32) - scale = 0.2 - result = env._obs_noise(data.copy(), scale) - assert np.all(result >= -scale) - assert np.all(result <= scale) - - def test_noise_scales_with_level(self): - env_half = _make_env(level=0.5) - env_full = _make_env(level=1.0) - data = np.zeros((1024, 10), dtype=np.float32) - scale = 1.0 - - np.random.seed(0) - r_half = env_half._obs_noise(data.copy(), scale) - np.random.seed(0) - r_full = env_full._obs_noise(data.copy(), scale) - - np.testing.assert_allclose(r_full, r_half * 2.0) - - def test_noise_preserves_dtype(self): - for dt in [np.float32, np.float64]: - env = _make_env(level=1.0) - data = np.ones((4, 5), dtype=dt) - result = env._obs_noise(data, 0.1) - assert result.dtype == dt - - def test_noise_preserves_shape(self): - env = _make_env(level=1.0) - for shape in [(1, 3), (64, 29), (1024, 10)]: - data = np.zeros(shape, dtype=np.float32) - result = env._obs_noise(data, 0.1) - assert result.shape == shape diff --git a/tests/tasks/test_migration_matrix.py b/tests/tasks/test_migration_matrix.py index d58dc1b7d..00d0dcfe1 100644 --- a/tests/tasks/test_migration_matrix.py +++ b/tests/tasks/test_migration_matrix.py @@ -18,25 +18,27 @@ def test_registered_tasks_have_explicit_migration_records() -> None: assert PRODUCTION_TASK_NAMES <= registered.keys() assert {record.task_name for record in records} == set(PRODUCTION_TASK_NAMES) assert len(records) == 39 - assert sum(record.status == "Compatible" for record in records) == 8 + assert sum(record.status == "Compatible" for record in records) == 11 assert sum(record.target == "compatibility" for record in records) == 3 @pytest.mark.parametrize( - ("task_name", "family", "target"), + ("task_name", "family", "target", "status"), [ - ("Go2ArmManipLoco", "go2_arm", "compatibility"), - ("SharpaInhandRotation", "sharpa", "compatibility"), - ("G1MotionTracking", "motion_tracking", "mba"), - ("G1WalkRough", "g1_locomotion", "mba"), - ("Go2JoystickRough", "quadruped_rough", "mba"), + ("Go2ArmManipLoco", "go2_arm", "compatibility", "Adapted"), + ("SharpaInhandRotation", "sharpa", "compatibility", "Adapted"), + ("G1MotionTracking", "motion_tracking", "mba", "Adapted"), + ("G1WalkRough", "g1_locomotion", "mba", "Adapted"), + ("Go2JoystickRough", "quadruped_rough", "complete", "Compatible"), ], ) -def test_matrix_records_high_risk_families(task_name: str, family: str, target: str) -> None: +def test_matrix_records_high_risk_families( + task_name: str, family: str, target: str, status: str +) -> None: record = migration_record(task_name) assert record.family == family assert record.target == target - assert record.status == "Adapted" + assert record.status == status def test_unknown_task_fails_closed() -> None: From d5bec570957e74d50be0f2b5963582071d21c281 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Thu, 20 Aug 2026 12:39:07 +0800 Subject: [PATCH 090/169] feat: infer manager term _target_ from type annotation (#1231) * feat: infer manager term _target_ from type annotation Manager mappings whose value type is a single concrete config dataclass (observations/groups, events, rewards, terminations, curriculum, metrics, recorders) no longer require each entry to repeat '_target_'; the materialization engine infers it from the field type annotation. actions/commands keep explicit '_target_' because their base configs are abstract. Bare class names now resolve against 'unilab.managers.' (e.g. SceneEntityCfg, UniformNoiseCfg), and the noise cfg family is re-exported from 'unilab.managers'. Sweeps the 23 migrated MBA task YAMLs (ppo/appo/offpolicy) to drop the redundant declarations and syncs the IsaacLab migration guide examples. * test: skip cpu affinity wiring tests without os.sched_getaffinity The module-level _AVAILABLE_CPUS lookup crashed collection on macOS, which lacks os.sched_getaffinity. Skip the module on such platforms, matching the existing mujoco_uni availability guards. --- conf/appo/task/allegro_inhand/base.yaml | 13 --- conf/appo/task/go1_joystick_flat/base.yaml | 38 +-------- conf/appo/task/go1_joystick_flat/motrix.yaml | 1 - conf/appo/task/go2_joystick_flat/base.yaml | 32 ------- .../flashsac/go2_joystick_flat/mujoco.yaml | 16 ++-- .../task/go1_joystick_flat/base.yaml | 38 +-------- .../task/go2_joystick_flat/base.yaml | 31 ------- .../task/go2w_joystick_flat/base.yaml | 51 ++--------- .../task/sac/go2_footstand/base.yaml | 16 +--- conf/offpolicy/task/stewart_balance/base.yaml | 10 --- conf/ppo/task/a2_joystick_flat/base.yaml | 51 ++--------- conf/ppo/task/allegro_inhand/base.yaml | 13 --- .../ppo/task/allegro_inhand_grasp/motrix.yaml | 6 -- .../ppo/task/allegro_inhand_grasp/mujoco.yaml | 6 -- conf/ppo/task/go1_joystick_flat/base.yaml | 38 +-------- conf/ppo/task/go2_footstand/base.yaml | 16 +--- conf/ppo/task/go2_footstand/mujoco.yaml | 12 +-- conf/ppo/task/go2_joystick_flat/base.yaml | 31 ------- conf/ppo/task/go2w_joystick_flat/base.yaml | 51 ++--------- .../task/quadruped_joystick_rough/base.yaml | 24 ++---- .../task/quadruped_joystick_rough/go2w.yaml | 35 ++------ .../quadruped_joystick_rough/quadruped.yaml | 24 +----- conf/ppo/task/stewart_balance/base.yaml | 10 --- .../3-framework_migration/1-from_isaac_lab.md | 13 ++- .../3-framework_migration/1-from_isaac_lab.md | 10 ++- src/unilab/base/config_materialization.py | 35 +++++++- src/unilab/managers/__init__.py | 8 ++ .../test_mujoco_cpu_affinity_wiring.py | 7 ++ tests/base/test_manager_config_overlay.py | 85 ++++++++++++++++++- tests/config/test_config_system.py | 1 - 30 files changed, 200 insertions(+), 522 deletions(-) diff --git a/conf/appo/task/allegro_inhand/base.yaml b/conf/appo/task/allegro_inhand/base.yaml index 7ee1f9deb..a36f24c9f 100644 --- a/conf/appo/task/allegro_inhand/base.yaml +++ b/conf/appo/task/allegro_inhand/base.yaml @@ -48,12 +48,10 @@ env: max_episode_seconds: 20.0 observations: policy: - _target_: unilab.managers.ObservationGroupCfg history_length: 3 flatten_history_dim: true terms: rotation: - _target_: unilab.managers.ObservationTermCfg func: unilab.tasks.manipulation.allegro_inhand.manager_terms.AllegroRotationObservation params: entity_name: robot @@ -70,11 +68,9 @@ env: raw_action_clip: [-1.0, 1.0] events: reset_scene_to_default: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.reset_scene_to_default mode: reset reset_hand_ball: - _target_: unilab.managers.EventTermCfg func: unilab.tasks.manipulation.allegro_inhand.manager_terms.AllegroHandBallReset mode: reset params: @@ -86,7 +82,6 @@ env: ball_velocity_noise: 0.0 ball_z_offset: 0.0 pd_gains: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.pd_gains mode: reset params: @@ -95,14 +90,12 @@ env: operation: abs terminations: dropped: - _target_: unilab.managers.TerminationTermCfg func: unilab.tasks.manipulation.allegro_inhand.manager_terms.AllegroDropTermination params: observation_group: policy observation_term: rotation minimum_ball_height: 0.125 time_out: - _target_: unilab.managers.TerminationTermCfg func: unilab.envs.mdp.time_out time_out: true scale_rewards_by_dt: true @@ -111,7 +104,6 @@ env: reward: rotate: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.manipulation.allegro_inhand.manager_terms.AllegroRotateReward weight: 1.25 params: @@ -120,31 +112,26 @@ reward: clip_min: -0.5 clip_max: 0.5 obj_linvel: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.manipulation.allegro_inhand.manager_terms.object_linear_velocity_l1 weight: -0.3 params: state_term_name: dropped pose_diff: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.manipulation.allegro_inhand.manager_terms.hand_pose_deviation_l2 weight: -0.3 params: state_term_name: dropped torque: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.manipulation.allegro_inhand.manager_terms.estimated_torque_l2 weight: -0.1 params: state_term_name: dropped work: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.manipulation.allegro_inhand.manager_terms.estimated_work_l2 weight: -2.0 params: state_term_name: dropped drop: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.manipulation.allegro_inhand.manager_terms.dropped weight: 0.0 params: diff --git a/conf/appo/task/go1_joystick_flat/base.yaml b/conf/appo/task/go1_joystick_flat/base.yaml index 07c3f2990..8a2f3bf57 100644 --- a/conf/appo/task/go1_joystick_flat/base.yaml +++ b/conf/appo/task/go1_joystick_flat/base.yaml @@ -40,71 +40,54 @@ env: max_episode_seconds: 20.0 observations: policy: - _target_: unilab.managers.ObservationGroupCfg terms: base_ang_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.builtin_sensor params: sensor_name: gyro projected_gravity: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.projected_gravity_from_sensor params: sensor_name: upvector joint_pos: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_pos_rel joint_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_vel_rel actions: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.last_action command: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.generated_commands params: command_name: twist gait_phase: - _target_: unilab.managers.ObservationTermCfg func: unilab.tasks.locomotion.common.manager_terms.quadruped_gait_phase params: frequency: 2.0 critic: - _target_: unilab.managers.ObservationGroupCfg terms: base_ang_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.builtin_sensor params: sensor_name: gyro projected_gravity: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.projected_gravity_from_sensor params: sensor_name: upvector joint_pos: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_pos_rel joint_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_vel_rel actions: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.last_action command: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.generated_commands params: command_name: twist gait_phase: - _target_: unilab.managers.ObservationTermCfg func: unilab.tasks.locomotion.common.manager_terms.quadruped_gait_phase params: frequency: 2.0 base_lin_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.builtin_sensor params: sensor_name: local_linvel @@ -133,11 +116,9 @@ env: ang_vel_z: [-0.8, 0.8] events: reset_scene_to_default: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.reset_scene_to_default mode: reset reset_root_state_uniform: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.reset_root_state_uniform mode: reset params: @@ -156,24 +137,22 @@ env: pitch: [-0.5, 0.5] yaw: [-0.5, 0.5] base_mass: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.randomize_rigid_body_mass mode: reset params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot body_names: trunk mass_distribution_params: [-1.5, 1.5] operation: add recompute_inertia: false base_com: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.randomize_rigid_body_com mode: reset params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot body_names: trunk com_range: @@ -181,7 +160,6 @@ env: y: [0.0, 0.0] z: [0.0, 0.0] pd_gains: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.pd_gains mode: reset params: @@ -189,7 +167,6 @@ env: kd_range: [0.5, 0.5] operation: abs push_robot: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.push_by_setting_velocity mode: interval interval_range_s: [15.0, 15.0] @@ -204,11 +181,9 @@ env: yaw: [0.0, 0.0] terminations: time_out: - _target_: unilab.managers.TerminationTermCfg func: unilab.envs.mdp.time_out time_out: true bad_orientation: - _target_: unilab.managers.TerminationTermCfg func: unilab.envs.mdp.bad_orientation params: limit_angle: 1.0471975511965976 @@ -217,43 +192,35 @@ env: reward: tracking_lin_vel: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel_xy_exp weight: 1.0 params: std: 0.5 command_name: twist tracking_ang_vel: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel_z_exp weight: 0.2 params: std: 0.5 command_name: twist lin_vel_z: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z_l2 weight: -5.0 ang_vel_xy: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy_l2 weight: -0.1 base_height: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.base_height_l2 weight: -100.0 params: target_height: 0.3 action_rate: - _target_: unilab.managers.RewardTermCfg func: unilab.envs.mdp.action_rate_l2 weight: -0.005 similar_to_default: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.joint_deviation_l1 weight: -0.1 contact: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.feet_phase_contact # Legacy Go1 sums four matching feet while this community term returns their mean. weight: 0.96 @@ -267,7 +234,6 @@ reward: contact_threshold: 0.1 stance_threshold: 0.6 swing_feet_z: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.feet_phase_swing_height weight: 4.0 params: diff --git a/conf/appo/task/go1_joystick_flat/motrix.yaml b/conf/appo/task/go1_joystick_flat/motrix.yaml index 09b622ce6..d4bb53521 100644 --- a/conf/appo/task/go1_joystick_flat/motrix.yaml +++ b/conf/appo/task/go1_joystick_flat/motrix.yaml @@ -30,7 +30,6 @@ reward: action_rate: weight: -0.015 action_smooth: - _target_: unilab.managers.RewardTermCfg func: unilab.envs.mdp.action_acc_l2 weight: -0.01 similar_to_default: diff --git a/conf/appo/task/go2_joystick_flat/base.yaml b/conf/appo/task/go2_joystick_flat/base.yaml index 288c022c7..be7364f78 100644 --- a/conf/appo/task/go2_joystick_flat/base.yaml +++ b/conf/appo/task/go2_joystick_flat/base.yaml @@ -39,71 +39,54 @@ env: max_episode_seconds: 20.0 observations: policy: - _target_: unilab.managers.ObservationGroupCfg terms: base_ang_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.builtin_sensor params: sensor_name: gyro projected_gravity: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.projected_gravity_from_sensor params: sensor_name: upvector joint_pos: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_pos_rel joint_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_vel_rel actions: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.last_action command: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.generated_commands params: command_name: twist gait_phase: - _target_: unilab.managers.ObservationTermCfg func: unilab.tasks.locomotion.common.manager_terms.quadruped_gait_phase params: frequency: 2.0 critic: - _target_: unilab.managers.ObservationGroupCfg terms: base_ang_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.builtin_sensor params: sensor_name: gyro projected_gravity: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.projected_gravity_from_sensor params: sensor_name: upvector joint_pos: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_pos_rel joint_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_vel_rel actions: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.last_action command: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.generated_commands params: command_name: twist gait_phase: - _target_: unilab.managers.ObservationTermCfg func: unilab.tasks.locomotion.common.manager_terms.quadruped_gait_phase params: frequency: 2.0 base_lin_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.builtin_sensor params: sensor_name: local_linvel @@ -132,11 +115,9 @@ env: ang_vel_z: [-0.8, 0.8] events: reset_scene_to_default: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.reset_scene_to_default mode: reset reset_root_state_uniform: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.reset_root_state_uniform mode: reset params: @@ -155,7 +136,6 @@ env: pitch: [-0.5, 0.5] yaw: [-0.5, 0.5] pd_gains: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.pd_gains mode: reset params: @@ -164,11 +144,9 @@ env: operation: abs terminations: time_out: - _target_: unilab.managers.TerminationTermCfg func: unilab.envs.mdp.time_out time_out: true bad_orientation: - _target_: unilab.managers.TerminationTermCfg func: unilab.envs.mdp.bad_orientation params: limit_angle: 1.0471975511965976 @@ -177,47 +155,38 @@ env: reward: tracking_lin_vel: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel_xy_exp weight: 1.0 params: std: 0.5 command_name: twist tracking_ang_vel: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel_z_exp weight: 0.2 params: std: 0.5 command_name: twist lin_vel_z: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z_l2 weight: -5.0 ang_vel_xy: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy_l2 weight: -0.1 base_height: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.base_height_l2 weight: -100.0 params: target_height: 0.3 action_rate: - _target_: unilab.managers.RewardTermCfg func: unilab.envs.mdp.action_rate_l2 weight: -0.005 similar_to_default: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.joint_deviation_l1 weight: -0.1 alive: - _target_: unilab.managers.RewardTermCfg func: unilab.envs.mdp.is_alive weight: 0.0 contact: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.feet_phase_contact weight: 0.24 params: @@ -230,7 +199,6 @@ reward: contact_threshold: 0.1 stance_threshold: 0.6 swing_feet_z: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.feet_phase_swing_height weight: 4.0 params: diff --git a/conf/offpolicy/task/flashsac/go2_joystick_flat/mujoco.yaml b/conf/offpolicy/task/flashsac/go2_joystick_flat/mujoco.yaml index ba0ddaa30..7fbfdec78 100644 --- a/conf/offpolicy/task/flashsac/go2_joystick_flat/mujoco.yaml +++ b/conf/offpolicy/task/flashsac/go2_joystick_flat/mujoco.yaml @@ -29,12 +29,12 @@ env: terms: joint_pos: noise: - _target_: unilab.managers._noise.UniformNoiseCfg + _target_: UniformNoiseCfg n_min: -0.01 n_max: 0.01 joint_vel: noise: - _target_: unilab.managers._noise.UniformNoiseCfg + _target_: UniformNoiseCfg n_min: -0.1 n_max: 0.1 critic: @@ -42,40 +42,37 @@ env: terms: joint_pos: noise: - _target_: unilab.managers._noise.UniformNoiseCfg + _target_: UniformNoiseCfg n_min: -0.01 n_max: 0.01 joint_vel: noise: - _target_: unilab.managers._noise.UniformNoiseCfg + _target_: UniformNoiseCfg n_min: -0.1 n_max: 0.1 events: randomize_rigid_body_mass: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.randomize_rigid_body_mass mode: reset params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot body_names: base mass_distribution_params: [-1.5, 1.5] operation: add recompute_inertia: false randomize_rigid_body_com: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.randomize_rigid_body_com mode: reset params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot body_names: base com_range: x: [-0.05, 0.05] randomize_physics_scene_gravity: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.randomize_physics_scene_gravity mode: reset params: @@ -84,7 +81,6 @@ env: - [0.0, 0.0, -9.81] operation: abs push_by_setting_velocity: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.push_by_setting_velocity mode: interval interval_range_s: [15.0, 15.0] diff --git a/conf/offpolicy/task/go1_joystick_flat/base.yaml b/conf/offpolicy/task/go1_joystick_flat/base.yaml index 07c3f2990..8a2f3bf57 100644 --- a/conf/offpolicy/task/go1_joystick_flat/base.yaml +++ b/conf/offpolicy/task/go1_joystick_flat/base.yaml @@ -40,71 +40,54 @@ env: max_episode_seconds: 20.0 observations: policy: - _target_: unilab.managers.ObservationGroupCfg terms: base_ang_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.builtin_sensor params: sensor_name: gyro projected_gravity: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.projected_gravity_from_sensor params: sensor_name: upvector joint_pos: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_pos_rel joint_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_vel_rel actions: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.last_action command: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.generated_commands params: command_name: twist gait_phase: - _target_: unilab.managers.ObservationTermCfg func: unilab.tasks.locomotion.common.manager_terms.quadruped_gait_phase params: frequency: 2.0 critic: - _target_: unilab.managers.ObservationGroupCfg terms: base_ang_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.builtin_sensor params: sensor_name: gyro projected_gravity: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.projected_gravity_from_sensor params: sensor_name: upvector joint_pos: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_pos_rel joint_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_vel_rel actions: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.last_action command: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.generated_commands params: command_name: twist gait_phase: - _target_: unilab.managers.ObservationTermCfg func: unilab.tasks.locomotion.common.manager_terms.quadruped_gait_phase params: frequency: 2.0 base_lin_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.builtin_sensor params: sensor_name: local_linvel @@ -133,11 +116,9 @@ env: ang_vel_z: [-0.8, 0.8] events: reset_scene_to_default: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.reset_scene_to_default mode: reset reset_root_state_uniform: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.reset_root_state_uniform mode: reset params: @@ -156,24 +137,22 @@ env: pitch: [-0.5, 0.5] yaw: [-0.5, 0.5] base_mass: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.randomize_rigid_body_mass mode: reset params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot body_names: trunk mass_distribution_params: [-1.5, 1.5] operation: add recompute_inertia: false base_com: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.randomize_rigid_body_com mode: reset params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot body_names: trunk com_range: @@ -181,7 +160,6 @@ env: y: [0.0, 0.0] z: [0.0, 0.0] pd_gains: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.pd_gains mode: reset params: @@ -189,7 +167,6 @@ env: kd_range: [0.5, 0.5] operation: abs push_robot: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.push_by_setting_velocity mode: interval interval_range_s: [15.0, 15.0] @@ -204,11 +181,9 @@ env: yaw: [0.0, 0.0] terminations: time_out: - _target_: unilab.managers.TerminationTermCfg func: unilab.envs.mdp.time_out time_out: true bad_orientation: - _target_: unilab.managers.TerminationTermCfg func: unilab.envs.mdp.bad_orientation params: limit_angle: 1.0471975511965976 @@ -217,43 +192,35 @@ env: reward: tracking_lin_vel: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel_xy_exp weight: 1.0 params: std: 0.5 command_name: twist tracking_ang_vel: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel_z_exp weight: 0.2 params: std: 0.5 command_name: twist lin_vel_z: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z_l2 weight: -5.0 ang_vel_xy: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy_l2 weight: -0.1 base_height: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.base_height_l2 weight: -100.0 params: target_height: 0.3 action_rate: - _target_: unilab.managers.RewardTermCfg func: unilab.envs.mdp.action_rate_l2 weight: -0.005 similar_to_default: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.joint_deviation_l1 weight: -0.1 contact: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.feet_phase_contact # Legacy Go1 sums four matching feet while this community term returns their mean. weight: 0.96 @@ -267,7 +234,6 @@ reward: contact_threshold: 0.1 stance_threshold: 0.6 swing_feet_z: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.feet_phase_swing_height weight: 4.0 params: diff --git a/conf/offpolicy/task/go2_joystick_flat/base.yaml b/conf/offpolicy/task/go2_joystick_flat/base.yaml index 639d29849..409129c16 100644 --- a/conf/offpolicy/task/go2_joystick_flat/base.yaml +++ b/conf/offpolicy/task/go2_joystick_flat/base.yaml @@ -39,71 +39,54 @@ env: max_episode_seconds: 20.0 observations: policy: - _target_: unilab.managers.ObservationGroupCfg terms: base_ang_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.builtin_sensor params: sensor_name: gyro projected_gravity: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.projected_gravity_from_sensor params: sensor_name: upvector joint_pos: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_pos_rel joint_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_vel_rel actions: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.last_action command: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.generated_commands params: command_name: twist gait_phase: - _target_: unilab.managers.ObservationTermCfg func: unilab.tasks.locomotion.common.manager_terms.quadruped_gait_phase params: frequency: 2.0 critic: - _target_: unilab.managers.ObservationGroupCfg terms: base_ang_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.builtin_sensor params: sensor_name: gyro projected_gravity: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.projected_gravity_from_sensor params: sensor_name: upvector joint_pos: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_pos_rel joint_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_vel_rel actions: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.last_action command: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.generated_commands params: command_name: twist gait_phase: - _target_: unilab.managers.ObservationTermCfg func: unilab.tasks.locomotion.common.manager_terms.quadruped_gait_phase params: frequency: 2.0 base_lin_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.builtin_sensor params: sensor_name: local_linvel @@ -132,11 +115,9 @@ env: ang_vel_z: [-0.8, 0.8] events: reset_scene_to_default: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.reset_scene_to_default mode: reset reset_root_state_uniform: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.reset_root_state_uniform mode: reset params: @@ -155,7 +136,6 @@ env: pitch: [-0.5, 0.5] yaw: [-0.5, 0.5] pd_gains: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.pd_gains mode: reset params: @@ -164,11 +144,9 @@ env: operation: abs terminations: time_out: - _target_: unilab.managers.TerminationTermCfg func: unilab.envs.mdp.time_out time_out: true bad_orientation: - _target_: unilab.managers.TerminationTermCfg func: unilab.envs.mdp.bad_orientation params: limit_angle: 1.0471975511965976 @@ -177,43 +155,35 @@ env: reward: tracking_lin_vel: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel_xy_exp weight: 1.0 params: std: 0.5 command_name: twist tracking_ang_vel: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel_z_exp weight: 0.2 params: std: 0.5 command_name: twist lin_vel_z: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z_l2 weight: -5.0 ang_vel_xy: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy_l2 weight: -0.1 base_height: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.base_height_l2 weight: -100.0 params: target_height: 0.3 action_rate: - _target_: unilab.managers.RewardTermCfg func: unilab.envs.mdp.action_rate_l2 weight: -0.005 similar_to_default: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.joint_deviation_l1 weight: -0.1 contact: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.feet_phase_contact weight: 0.24 params: @@ -226,7 +196,6 @@ reward: contact_threshold: 0.1 stance_threshold: 0.6 swing_feet_z: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.feet_phase_swing_height weight: 4.0 params: diff --git a/conf/offpolicy/task/go2w_joystick_flat/base.yaml b/conf/offpolicy/task/go2w_joystick_flat/base.yaml index 0f13a638c..3d057e1ee 100644 --- a/conf/offpolicy/task/go2w_joystick_flat/base.yaml +++ b/conf/offpolicy/task/go2w_joystick_flat/base.yaml @@ -48,106 +48,88 @@ env: max_episode_seconds: 20.0 observations: policy: - _target_: unilab.managers.ObservationGroupCfg terms: base_ang_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.builtin_sensor params: sensor_name: gyro projected_gravity: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.projected_gravity_from_sensor params: sensor_name: upvector leg_joint_pos: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_pos_rel params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot joint_names: ".*_(hip|thigh|calf)_joint" leg_joint_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_vel_rel params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot joint_names: ".*_(hip|thigh|calf)_joint" wheel_joint_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_vel_rel params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot joint_names: ".*_wheel_joint" actions: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.last_action params: action_name: motor command: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.generated_commands params: command_name: twist critic: - _target_: unilab.managers.ObservationGroupCfg terms: base_ang_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.builtin_sensor params: sensor_name: gyro projected_gravity: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.projected_gravity_from_sensor params: sensor_name: upvector leg_joint_pos: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_pos_rel params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot joint_names: ".*_(hip|thigh|calf)_joint" leg_joint_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_vel_rel params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot joint_names: ".*_(hip|thigh|calf)_joint" wheel_joint_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_vel_rel params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot joint_names: ".*_wheel_joint" actions: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.last_action params: action_name: motor command: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.generated_commands params: command_name: twist base_lin_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.builtin_sensor params: sensor_name: local_linvel motor_torque: - _target_: unilab.managers.ObservationTermCfg func: unilab.tasks.locomotion.go2w.manager_terms.motor_torque params: action_name: motor @@ -182,11 +164,9 @@ env: ang_vel_z: [-1.0, 1.0] events: reset_scene_to_default: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.reset_scene_to_default mode: reset reset_root_state_uniform: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.reset_root_state_uniform mode: reset params: @@ -205,7 +185,6 @@ env: pitch: [-0.5, 0.5] yaw: [-0.5, 0.5] motor_gains: - _target_: unilab.managers.EventTermCfg func: unilab.tasks.locomotion.go2w.manager_terms.randomize_motor_gains mode: reset params: @@ -214,11 +193,9 @@ env: kd_multiplier_range: [1.0, 1.0] terminations: time_out: - _target_: unilab.managers.TerminationTermCfg func: unilab.envs.mdp.time_out time_out: true bad_orientation: - _target_: unilab.managers.TerminationTermCfg func: unilab.envs.mdp.bad_orientation params: limit_angle: 1.0471975511965976 @@ -227,72 +204,60 @@ env: reward: tracking_lin_vel: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel_xy_exp weight: 1.0 params: std: 0.5 command_name: twist tracking_ang_vel: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel_z_exp weight: 0.75 params: std: 0.5 command_name: twist lin_vel_z: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z_l2 weight: -5.0 ang_vel_xy: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy_l2 weight: -0.1 base_height: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.base_height_l2 weight: -100.0 params: target_height: 0.4 orientation: - _target_: unilab.managers.RewardTermCfg func: unilab.envs.mdp.flat_orientation_l2 weight: -2.0 action_rate: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.go2w.manager_terms.clipped_action_rate_l2 weight: -0.005 params: action_name: motor similar_to_default: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.joint_deviation_l1 weight: -0.5 params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot joint_names: ".*_(hip|thigh|calf)_joint" torques: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.go2w.manager_terms.motor_torque_l2 weight: -0.0002 params: action_name: motor wheel_vel: - _target_: unilab.managers.RewardTermCfg func: unilab.envs.mdp.joint_vel_l2 weight: 0.0 params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot joint_names: ".*_wheel_joint" alive: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.go2w.manager_terms.constant_alive weight: 0.5 upward: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.go2w.manager_terms.upward_l2 weight: 1.0 diff --git a/conf/offpolicy/task/sac/go2_footstand/base.yaml b/conf/offpolicy/task/sac/go2_footstand/base.yaml index 814e6d7a2..08095758c 100644 --- a/conf/offpolicy/task/sac/go2_footstand/base.yaml +++ b/conf/offpolicy/task/sac/go2_footstand/base.yaml @@ -55,30 +55,25 @@ env: adaptive_chunk_size: false observations: policy: - _target_: unilab.managers.ObservationGroupCfg enable_corruption: true terms: frame: - _target_: unilab.managers.ObservationTermCfg func: unilab.tasks.locomotion.go2.footstand.frame_observation params: action_name: joint_pos noise: - _target_: unilab.managers._noise.UniformNoiseCfg + _target_: UniformNoiseCfg n_min: [-0.1, -0.1, -0.1, -0.2, -0.2, -0.2, -0.05, -0.05, -0.05, -0.01, -0.01, -0.01, -0.01, -0.01, -0.01, -0.01, -0.01, -0.01, -0.01, -0.01, -0.01, -1.5, -1.5, -1.5, -1.5, -1.5, -1.5, -1.5, -1.5, -1.5, -1.5, -1.5, -1.5, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] n_max: [0.1, 0.1, 0.1, 0.2, 0.2, 0.2, 0.05, 0.05, 0.05, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] history_length: 15 critic: - _target_: unilab.managers.ObservationGroupCfg terms: frame: - _target_: unilab.managers.ObservationTermCfg func: unilab.tasks.locomotion.go2.footstand.frame_observation params: action_name: joint_pos history_length: 15 privileged: - _target_: unilab.managers.ObservationTermCfg func: unilab.tasks.locomotion.go2.footstand.privileged_observation params: action_name: joint_pos @@ -132,11 +127,9 @@ env: simulate_action_latency: false events: reset_scene_to_default: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.reset_scene_to_default mode: reset reset_root_state_uniform: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.reset_root_state_uniform mode: reset params: @@ -155,17 +148,15 @@ env: pitch: [-0.5, 0.5] yaw: [-0.5, 0.5] reset_joints: - _target_: unilab.managers.EventTermCfg func: unilab.tasks.locomotion.go2.footstand.FootstandJointReset mode: reset params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot joint_names: ".*" position_offset_range: [-0.05, 0.05] pd_gains: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.pd_gains mode: reset params: @@ -178,7 +169,6 @@ env: joint_armature: null terminations: footstand: - _target_: unilab.managers.TerminationTermCfg func: unilab.tasks.locomotion.go2.footstand.FootstandTermination params: action_name: joint_pos @@ -187,7 +177,6 @@ env: orientation_threshold: 0.2 energy_threshold: 200.0 time_out: - _target_: unilab.managers.TerminationTermCfg func: unilab.envs.mdp.time_out time_out: true policy_observation_group: policy @@ -196,7 +185,6 @@ env: reward: footstand: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.go2.footstand.FootstandReward weight: 1.0 params: diff --git a/conf/offpolicy/task/stewart_balance/base.yaml b/conf/offpolicy/task/stewart_balance/base.yaml index 5dfdf4279..e38b87656 100644 --- a/conf/offpolicy/task/stewart_balance/base.yaml +++ b/conf/offpolicy/task/stewart_balance/base.yaml @@ -29,10 +29,8 @@ env: render_spacing: 4.5 observations: policy: - _target_: unilab.managers.ObservationGroupCfg terms: balance: - _target_: unilab.managers.ObservationTermCfg func: unilab.tasks.manipulation.stewart.balance.StewartObservation params: entity_name: stewart @@ -63,11 +61,9 @@ env: center_control_min_gain: 0.15 events: reset_scene_to_default: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.reset_scene_to_default mode: reset reset_ball: - _target_: unilab.managers.EventTermCfg func: unilab.tasks.manipulation.stewart.balance.StewartBallReset mode: reset params: @@ -77,7 +73,6 @@ env: ball_home_z: 1.2 terminations: balance_state: - _target_: unilab.managers.TerminationTermCfg func: unilab.tasks.manipulation.stewart.balance.StewartBalanceState params: observation_group: policy @@ -92,7 +87,6 @@ env: zero_vel_thresh: 0.07 still_steps_needed: 5 time_out: - _target_: unilab.managers.TerminationTermCfg func: unilab.envs.mdp.time_out time_out: true scale_rewards_by_dt: false @@ -101,25 +95,21 @@ env: reward: center: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.manipulation.stewart.balance.center_reward weight: 0.7 params: state_term_name: balance_state progress: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.manipulation.stewart.balance.progress_reward weight: 0.6 params: state_term_name: balance_state still: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.manipulation.stewart.balance.still_reward weight: 3.0 params: state_term_name: balance_state fall: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.manipulation.stewart.balance.fall_reward weight: -6.0 params: diff --git a/conf/ppo/task/a2_joystick_flat/base.yaml b/conf/ppo/task/a2_joystick_flat/base.yaml index 8ac59f4dc..87203212a 100644 --- a/conf/ppo/task/a2_joystick_flat/base.yaml +++ b/conf/ppo/task/a2_joystick_flat/base.yaml @@ -41,75 +41,58 @@ env: max_episode_seconds: 20.0 observations: policy: - _target_: unilab.managers.ObservationGroupCfg terms: base_ang_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.builtin_sensor params: sensor_name: gyro projected_gravity: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.projected_gravity_from_sensor params: sensor_name: upvector joint_pos: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_pos_rel joint_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_vel_rel actions: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.last_action command: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.generated_commands params: command_name: twist gait_phase: - _target_: unilab.managers.ObservationTermCfg func: unilab.tasks.locomotion.common.manager_terms.quadruped_gait_phase params: frequency: 2.0 command_name: twist command_threshold: 0.1 critic: - _target_: unilab.managers.ObservationGroupCfg terms: base_ang_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.builtin_sensor params: sensor_name: gyro projected_gravity: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.projected_gravity_from_sensor params: sensor_name: upvector joint_pos: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_pos_rel joint_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_vel_rel actions: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.last_action command: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.generated_commands params: command_name: twist gait_phase: - _target_: unilab.managers.ObservationTermCfg func: unilab.tasks.locomotion.common.manager_terms.quadruped_gait_phase params: frequency: 2.0 command_name: twist command_threshold: 0.1 base_lin_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.builtin_sensor params: sensor_name: local_linvel @@ -138,11 +121,9 @@ env: ang_vel_z: [-0.8, 0.8] events: reset_scene_to_default: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.reset_scene_to_default mode: reset reset_root_state_uniform: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.reset_root_state_uniform mode: reset params: @@ -161,24 +142,22 @@ env: pitch: [-0.5, 0.5] yaw: [-0.5, 0.5] base_mass: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.randomize_rigid_body_mass mode: reset params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot body_names: base_link mass_distribution_params: [0.0, 8.0] operation: add recompute_inertia: false base_com: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.randomize_rigid_body_com mode: reset params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot body_names: base_link com_range: @@ -186,12 +165,11 @@ env: y: [-0.08, 0.08] z: [-0.08, 0.08] foot_friction: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.geom_friction mode: reset params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot geom_names: floor ranges: [0.3, 1.6] @@ -199,30 +177,27 @@ env: axes: [0] shared_random: true joint_armature: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.joint_armature mode: reset params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot joint_names: ".*" ranges: [0.9, 1.1] operation: scale pd_gains: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.pd_gains mode: reset params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot actuator_names: ".*" kp_range: [0.9, 1.1] kd_range: [0.9, 1.1] operation: scale push_robot: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.push_by_setting_velocity mode: interval interval_range_s: [8.0, 8.0] @@ -237,11 +212,9 @@ env: yaw: [0.0, 0.0] terminations: time_out: - _target_: unilab.managers.TerminationTermCfg func: unilab.envs.mdp.time_out time_out: true bad_orientation: - _target_: unilab.managers.TerminationTermCfg func: unilab.envs.mdp.bad_orientation params: limit_angle: 1.0471975511965976 @@ -250,43 +223,35 @@ env: reward: tracking_lin_vel: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel_xy_exp weight: 1.0 params: std: 0.5 command_name: twist tracking_ang_vel: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel_z_exp weight: 0.4 params: std: 0.5 command_name: twist lin_vel_z: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z_l2 weight: -5.0 ang_vel_xy: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy_l2 weight: -0.1 base_height: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.base_height_l2 weight: -100.0 params: target_height: 0.4 action_rate: - _target_: unilab.managers.RewardTermCfg func: unilab.envs.mdp.action_rate_l2 weight: -0.02 similar_to_default: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.joint_deviation_l1 weight: -0.25 contact: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.feet_phase_contact weight: 0.5 params: @@ -297,7 +262,6 @@ reward: contact_threshold: 0.1 stance_threshold: 0.6 swing_feet_z: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.feet_phase_swing_height weight: 4.0 params: @@ -309,23 +273,20 @@ reward: kernel: 0.01 swing_start: 0.6 stand_still: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.stand_still_l1 weight: -4.0 params: command_name: twist command_threshold: 0.1 hip_deviation: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.joint_deviation_l1 weight: -1.0 params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot joint_names: ".*_hip_joint" stand_feet_air: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.feet_air_while_standing weight: -1.0 params: diff --git a/conf/ppo/task/allegro_inhand/base.yaml b/conf/ppo/task/allegro_inhand/base.yaml index 6c7851171..afafae38c 100644 --- a/conf/ppo/task/allegro_inhand/base.yaml +++ b/conf/ppo/task/allegro_inhand/base.yaml @@ -49,12 +49,10 @@ env: max_episode_seconds: 20.0 observations: policy: - _target_: unilab.managers.ObservationGroupCfg history_length: 3 flatten_history_dim: true terms: rotation: - _target_: unilab.managers.ObservationTermCfg func: unilab.tasks.manipulation.allegro_inhand.manager_terms.AllegroRotationObservation params: entity_name: robot @@ -71,11 +69,9 @@ env: raw_action_clip: [-1.0, 1.0] events: reset_scene_to_default: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.reset_scene_to_default mode: reset reset_hand_ball: - _target_: unilab.managers.EventTermCfg func: unilab.tasks.manipulation.allegro_inhand.manager_terms.AllegroHandBallReset mode: reset params: @@ -87,7 +83,6 @@ env: ball_velocity_noise: 0.0 ball_z_offset: 0.0 pd_gains: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.pd_gains mode: reset params: @@ -96,14 +91,12 @@ env: operation: abs terminations: dropped: - _target_: unilab.managers.TerminationTermCfg func: unilab.tasks.manipulation.allegro_inhand.manager_terms.AllegroDropTermination params: observation_group: policy observation_term: rotation minimum_ball_height: 0.125 time_out: - _target_: unilab.managers.TerminationTermCfg func: unilab.envs.mdp.time_out time_out: true scale_rewards_by_dt: true @@ -112,7 +105,6 @@ env: reward: rotate: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.manipulation.allegro_inhand.manager_terms.AllegroRotateReward weight: 1.25 params: @@ -121,31 +113,26 @@ reward: clip_min: -0.5 clip_max: 0.5 obj_linvel: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.manipulation.allegro_inhand.manager_terms.object_linear_velocity_l1 weight: -0.3 params: state_term_name: dropped pose_diff: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.manipulation.allegro_inhand.manager_terms.hand_pose_deviation_l2 weight: -0.3 params: state_term_name: dropped torque: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.manipulation.allegro_inhand.manager_terms.estimated_torque_l2 weight: -0.1 params: state_term_name: dropped work: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.manipulation.allegro_inhand.manager_terms.estimated_work_l2 weight: -2.0 params: state_term_name: dropped drop: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.manipulation.allegro_inhand.manager_terms.dropped weight: 0.0 params: diff --git a/conf/ppo/task/allegro_inhand_grasp/motrix.yaml b/conf/ppo/task/allegro_inhand_grasp/motrix.yaml index 0006f9b43..e75de1bd9 100644 --- a/conf/ppo/task/allegro_inhand_grasp/motrix.yaml +++ b/conf/ppo/task/allegro_inhand_grasp/motrix.yaml @@ -48,7 +48,6 @@ env: ball_z_offset: 0.0 terminations: invalid_grasp: - _target_: unilab.managers.TerminationTermCfg func: unilab.tasks.manipulation.allegro_inhand.grasp_gen.AllegroGraspQualityTermination params: entity_name: robot @@ -62,24 +61,19 @@ env: enabled: true metrics: fingertips_close: - _target_: unilab.managers.MetricsTermCfg func: unilab.tasks.manipulation.allegro_inhand.grasp_gen.AllegroGraspQualityMetric params: {quality_term_name: invalid_grasp, condition: fingertips_close} enough_contacts: - _target_: unilab.managers.MetricsTermCfg func: unilab.tasks.manipulation.allegro_inhand.grasp_gen.AllegroGraspQualityMetric params: {quality_term_name: invalid_grasp, condition: enough_contacts} ball_held: - _target_: unilab.managers.MetricsTermCfg func: unilab.tasks.manipulation.allegro_inhand.grasp_gen.AllegroGraspQualityMetric params: {quality_term_name: invalid_grasp, condition: ball_held} valid: - _target_: unilab.managers.MetricsTermCfg func: unilab.tasks.manipulation.allegro_inhand.grasp_gen.AllegroGraspQualityMetric params: {quality_term_name: invalid_grasp, condition: valid} recorders: grasp_cache: - _target_: unilab.managers.RecorderTermCfg func: unilab.tasks.manipulation.allegro_inhand.grasp_gen.AllegroGraspRecorder params: quality_term_name: invalid_grasp diff --git a/conf/ppo/task/allegro_inhand_grasp/mujoco.yaml b/conf/ppo/task/allegro_inhand_grasp/mujoco.yaml index 9d16b9809..693937e3d 100644 --- a/conf/ppo/task/allegro_inhand_grasp/mujoco.yaml +++ b/conf/ppo/task/allegro_inhand_grasp/mujoco.yaml @@ -52,7 +52,6 @@ env: ball_z_offset: 0.0 terminations: invalid_grasp: - _target_: unilab.managers.TerminationTermCfg func: unilab.tasks.manipulation.allegro_inhand.grasp_gen.AllegroGraspQualityTermination params: entity_name: robot @@ -66,24 +65,19 @@ env: enabled: true metrics: fingertips_close: - _target_: unilab.managers.MetricsTermCfg func: unilab.tasks.manipulation.allegro_inhand.grasp_gen.AllegroGraspQualityMetric params: {quality_term_name: invalid_grasp, condition: fingertips_close} enough_contacts: - _target_: unilab.managers.MetricsTermCfg func: unilab.tasks.manipulation.allegro_inhand.grasp_gen.AllegroGraspQualityMetric params: {quality_term_name: invalid_grasp, condition: enough_contacts} ball_held: - _target_: unilab.managers.MetricsTermCfg func: unilab.tasks.manipulation.allegro_inhand.grasp_gen.AllegroGraspQualityMetric params: {quality_term_name: invalid_grasp, condition: ball_held} valid: - _target_: unilab.managers.MetricsTermCfg func: unilab.tasks.manipulation.allegro_inhand.grasp_gen.AllegroGraspQualityMetric params: {quality_term_name: invalid_grasp, condition: valid} recorders: grasp_cache: - _target_: unilab.managers.RecorderTermCfg func: unilab.tasks.manipulation.allegro_inhand.grasp_gen.AllegroGraspRecorder params: quality_term_name: invalid_grasp diff --git a/conf/ppo/task/go1_joystick_flat/base.yaml b/conf/ppo/task/go1_joystick_flat/base.yaml index 07c3f2990..8a2f3bf57 100644 --- a/conf/ppo/task/go1_joystick_flat/base.yaml +++ b/conf/ppo/task/go1_joystick_flat/base.yaml @@ -40,71 +40,54 @@ env: max_episode_seconds: 20.0 observations: policy: - _target_: unilab.managers.ObservationGroupCfg terms: base_ang_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.builtin_sensor params: sensor_name: gyro projected_gravity: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.projected_gravity_from_sensor params: sensor_name: upvector joint_pos: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_pos_rel joint_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_vel_rel actions: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.last_action command: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.generated_commands params: command_name: twist gait_phase: - _target_: unilab.managers.ObservationTermCfg func: unilab.tasks.locomotion.common.manager_terms.quadruped_gait_phase params: frequency: 2.0 critic: - _target_: unilab.managers.ObservationGroupCfg terms: base_ang_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.builtin_sensor params: sensor_name: gyro projected_gravity: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.projected_gravity_from_sensor params: sensor_name: upvector joint_pos: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_pos_rel joint_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_vel_rel actions: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.last_action command: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.generated_commands params: command_name: twist gait_phase: - _target_: unilab.managers.ObservationTermCfg func: unilab.tasks.locomotion.common.manager_terms.quadruped_gait_phase params: frequency: 2.0 base_lin_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.builtin_sensor params: sensor_name: local_linvel @@ -133,11 +116,9 @@ env: ang_vel_z: [-0.8, 0.8] events: reset_scene_to_default: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.reset_scene_to_default mode: reset reset_root_state_uniform: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.reset_root_state_uniform mode: reset params: @@ -156,24 +137,22 @@ env: pitch: [-0.5, 0.5] yaw: [-0.5, 0.5] base_mass: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.randomize_rigid_body_mass mode: reset params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot body_names: trunk mass_distribution_params: [-1.5, 1.5] operation: add recompute_inertia: false base_com: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.randomize_rigid_body_com mode: reset params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot body_names: trunk com_range: @@ -181,7 +160,6 @@ env: y: [0.0, 0.0] z: [0.0, 0.0] pd_gains: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.pd_gains mode: reset params: @@ -189,7 +167,6 @@ env: kd_range: [0.5, 0.5] operation: abs push_robot: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.push_by_setting_velocity mode: interval interval_range_s: [15.0, 15.0] @@ -204,11 +181,9 @@ env: yaw: [0.0, 0.0] terminations: time_out: - _target_: unilab.managers.TerminationTermCfg func: unilab.envs.mdp.time_out time_out: true bad_orientation: - _target_: unilab.managers.TerminationTermCfg func: unilab.envs.mdp.bad_orientation params: limit_angle: 1.0471975511965976 @@ -217,43 +192,35 @@ env: reward: tracking_lin_vel: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel_xy_exp weight: 1.0 params: std: 0.5 command_name: twist tracking_ang_vel: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel_z_exp weight: 0.2 params: std: 0.5 command_name: twist lin_vel_z: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z_l2 weight: -5.0 ang_vel_xy: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy_l2 weight: -0.1 base_height: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.base_height_l2 weight: -100.0 params: target_height: 0.3 action_rate: - _target_: unilab.managers.RewardTermCfg func: unilab.envs.mdp.action_rate_l2 weight: -0.005 similar_to_default: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.joint_deviation_l1 weight: -0.1 contact: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.feet_phase_contact # Legacy Go1 sums four matching feet while this community term returns their mean. weight: 0.96 @@ -267,7 +234,6 @@ reward: contact_threshold: 0.1 stance_threshold: 0.6 swing_feet_z: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.feet_phase_swing_height weight: 4.0 params: diff --git a/conf/ppo/task/go2_footstand/base.yaml b/conf/ppo/task/go2_footstand/base.yaml index 8b7abd5d2..5478fc5b5 100644 --- a/conf/ppo/task/go2_footstand/base.yaml +++ b/conf/ppo/task/go2_footstand/base.yaml @@ -55,30 +55,25 @@ env: adaptive_chunk_size: false observations: policy: - _target_: unilab.managers.ObservationGroupCfg enable_corruption: true terms: frame: - _target_: unilab.managers.ObservationTermCfg func: unilab.tasks.locomotion.go2.footstand.frame_observation params: action_name: joint_pos noise: - _target_: unilab.managers._noise.UniformNoiseCfg + _target_: UniformNoiseCfg n_min: [-0.1, -0.1, -0.1, -0.2, -0.2, -0.2, -0.05, -0.05, -0.05, -0.01, -0.01, -0.01, -0.01, -0.01, -0.01, -0.01, -0.01, -0.01, -0.01, -0.01, -0.01, -1.5, -1.5, -1.5, -1.5, -1.5, -1.5, -1.5, -1.5, -1.5, -1.5, -1.5, -1.5, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] n_max: [0.1, 0.1, 0.1, 0.2, 0.2, 0.2, 0.05, 0.05, 0.05, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] history_length: 15 critic: - _target_: unilab.managers.ObservationGroupCfg terms: frame: - _target_: unilab.managers.ObservationTermCfg func: unilab.tasks.locomotion.go2.footstand.frame_observation params: action_name: joint_pos history_length: 15 privileged: - _target_: unilab.managers.ObservationTermCfg func: unilab.tasks.locomotion.go2.footstand.privileged_observation params: action_name: joint_pos @@ -132,11 +127,9 @@ env: simulate_action_latency: false events: reset_scene_to_default: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.reset_scene_to_default mode: reset reset_root_state_uniform: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.reset_root_state_uniform mode: reset params: @@ -155,17 +148,15 @@ env: pitch: [-0.5, 0.5] yaw: [-0.5, 0.5] reset_joints: - _target_: unilab.managers.EventTermCfg func: unilab.tasks.locomotion.go2.footstand.FootstandJointReset mode: reset params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot joint_names: ".*" position_offset_range: [-0.05, 0.05] pd_gains: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.pd_gains mode: reset params: @@ -178,7 +169,6 @@ env: joint_armature: null terminations: footstand: - _target_: unilab.managers.TerminationTermCfg func: unilab.tasks.locomotion.go2.footstand.FootstandTermination params: action_name: joint_pos @@ -187,7 +177,6 @@ env: orientation_threshold: 0.2 energy_threshold: 200.0 time_out: - _target_: unilab.managers.TerminationTermCfg func: unilab.envs.mdp.time_out time_out: true policy_observation_group: policy @@ -196,7 +185,6 @@ env: reward: footstand: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.go2.footstand.FootstandReward weight: 1.0 params: diff --git a/conf/ppo/task/go2_footstand/mujoco.yaml b/conf/ppo/task/go2_footstand/mujoco.yaml index 6e8e80c94..8a4b428cc 100644 --- a/conf/ppo/task/go2_footstand/mujoco.yaml +++ b/conf/ppo/task/go2_footstand/mujoco.yaml @@ -10,35 +10,32 @@ training: env: events: floor_friction: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.geom_friction mode: reset params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot geom_names: floor ranges: [0.4, 1.0] operation: abs link_mass: - _target_: unilab.managers.EventTermCfg func: unilab.tasks.locomotion.go2.footstand.FootstandMassRandomization mode: reset params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot body_names: ".*" torso_body_name: base link_mass_scale_range: [0.9, 1.1] torso_added_mass_range: [-1.0, 1.0] torso_com: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.randomize_rigid_body_com mode: reset params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot body_names: base com_range: @@ -46,12 +43,11 @@ env: y: [-0.05, 0.05] z: [-0.05, 0.05] joint_armature: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.joint_armature mode: reset params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot joint_names: ".*" ranges: [1.0, 1.05] diff --git a/conf/ppo/task/go2_joystick_flat/base.yaml b/conf/ppo/task/go2_joystick_flat/base.yaml index 639d29849..409129c16 100644 --- a/conf/ppo/task/go2_joystick_flat/base.yaml +++ b/conf/ppo/task/go2_joystick_flat/base.yaml @@ -39,71 +39,54 @@ env: max_episode_seconds: 20.0 observations: policy: - _target_: unilab.managers.ObservationGroupCfg terms: base_ang_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.builtin_sensor params: sensor_name: gyro projected_gravity: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.projected_gravity_from_sensor params: sensor_name: upvector joint_pos: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_pos_rel joint_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_vel_rel actions: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.last_action command: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.generated_commands params: command_name: twist gait_phase: - _target_: unilab.managers.ObservationTermCfg func: unilab.tasks.locomotion.common.manager_terms.quadruped_gait_phase params: frequency: 2.0 critic: - _target_: unilab.managers.ObservationGroupCfg terms: base_ang_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.builtin_sensor params: sensor_name: gyro projected_gravity: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.projected_gravity_from_sensor params: sensor_name: upvector joint_pos: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_pos_rel joint_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_vel_rel actions: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.last_action command: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.generated_commands params: command_name: twist gait_phase: - _target_: unilab.managers.ObservationTermCfg func: unilab.tasks.locomotion.common.manager_terms.quadruped_gait_phase params: frequency: 2.0 base_lin_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.builtin_sensor params: sensor_name: local_linvel @@ -132,11 +115,9 @@ env: ang_vel_z: [-0.8, 0.8] events: reset_scene_to_default: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.reset_scene_to_default mode: reset reset_root_state_uniform: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.reset_root_state_uniform mode: reset params: @@ -155,7 +136,6 @@ env: pitch: [-0.5, 0.5] yaw: [-0.5, 0.5] pd_gains: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.pd_gains mode: reset params: @@ -164,11 +144,9 @@ env: operation: abs terminations: time_out: - _target_: unilab.managers.TerminationTermCfg func: unilab.envs.mdp.time_out time_out: true bad_orientation: - _target_: unilab.managers.TerminationTermCfg func: unilab.envs.mdp.bad_orientation params: limit_angle: 1.0471975511965976 @@ -177,43 +155,35 @@ env: reward: tracking_lin_vel: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel_xy_exp weight: 1.0 params: std: 0.5 command_name: twist tracking_ang_vel: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel_z_exp weight: 0.2 params: std: 0.5 command_name: twist lin_vel_z: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z_l2 weight: -5.0 ang_vel_xy: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy_l2 weight: -0.1 base_height: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.base_height_l2 weight: -100.0 params: target_height: 0.3 action_rate: - _target_: unilab.managers.RewardTermCfg func: unilab.envs.mdp.action_rate_l2 weight: -0.005 similar_to_default: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.joint_deviation_l1 weight: -0.1 contact: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.feet_phase_contact weight: 0.24 params: @@ -226,7 +196,6 @@ reward: contact_threshold: 0.1 stance_threshold: 0.6 swing_feet_z: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.feet_phase_swing_height weight: 4.0 params: diff --git a/conf/ppo/task/go2w_joystick_flat/base.yaml b/conf/ppo/task/go2w_joystick_flat/base.yaml index 0f13a638c..3d057e1ee 100644 --- a/conf/ppo/task/go2w_joystick_flat/base.yaml +++ b/conf/ppo/task/go2w_joystick_flat/base.yaml @@ -48,106 +48,88 @@ env: max_episode_seconds: 20.0 observations: policy: - _target_: unilab.managers.ObservationGroupCfg terms: base_ang_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.builtin_sensor params: sensor_name: gyro projected_gravity: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.projected_gravity_from_sensor params: sensor_name: upvector leg_joint_pos: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_pos_rel params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot joint_names: ".*_(hip|thigh|calf)_joint" leg_joint_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_vel_rel params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot joint_names: ".*_(hip|thigh|calf)_joint" wheel_joint_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_vel_rel params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot joint_names: ".*_wheel_joint" actions: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.last_action params: action_name: motor command: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.generated_commands params: command_name: twist critic: - _target_: unilab.managers.ObservationGroupCfg terms: base_ang_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.builtin_sensor params: sensor_name: gyro projected_gravity: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.projected_gravity_from_sensor params: sensor_name: upvector leg_joint_pos: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_pos_rel params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot joint_names: ".*_(hip|thigh|calf)_joint" leg_joint_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_vel_rel params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot joint_names: ".*_(hip|thigh|calf)_joint" wheel_joint_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_vel_rel params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot joint_names: ".*_wheel_joint" actions: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.last_action params: action_name: motor command: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.generated_commands params: command_name: twist base_lin_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.builtin_sensor params: sensor_name: local_linvel motor_torque: - _target_: unilab.managers.ObservationTermCfg func: unilab.tasks.locomotion.go2w.manager_terms.motor_torque params: action_name: motor @@ -182,11 +164,9 @@ env: ang_vel_z: [-1.0, 1.0] events: reset_scene_to_default: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.reset_scene_to_default mode: reset reset_root_state_uniform: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.reset_root_state_uniform mode: reset params: @@ -205,7 +185,6 @@ env: pitch: [-0.5, 0.5] yaw: [-0.5, 0.5] motor_gains: - _target_: unilab.managers.EventTermCfg func: unilab.tasks.locomotion.go2w.manager_terms.randomize_motor_gains mode: reset params: @@ -214,11 +193,9 @@ env: kd_multiplier_range: [1.0, 1.0] terminations: time_out: - _target_: unilab.managers.TerminationTermCfg func: unilab.envs.mdp.time_out time_out: true bad_orientation: - _target_: unilab.managers.TerminationTermCfg func: unilab.envs.mdp.bad_orientation params: limit_angle: 1.0471975511965976 @@ -227,72 +204,60 @@ env: reward: tracking_lin_vel: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel_xy_exp weight: 1.0 params: std: 0.5 command_name: twist tracking_ang_vel: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel_z_exp weight: 0.75 params: std: 0.5 command_name: twist lin_vel_z: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z_l2 weight: -5.0 ang_vel_xy: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy_l2 weight: -0.1 base_height: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.base_height_l2 weight: -100.0 params: target_height: 0.4 orientation: - _target_: unilab.managers.RewardTermCfg func: unilab.envs.mdp.flat_orientation_l2 weight: -2.0 action_rate: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.go2w.manager_terms.clipped_action_rate_l2 weight: -0.005 params: action_name: motor similar_to_default: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.joint_deviation_l1 weight: -0.5 params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot joint_names: ".*_(hip|thigh|calf)_joint" torques: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.go2w.manager_terms.motor_torque_l2 weight: -0.0002 params: action_name: motor wheel_vel: - _target_: unilab.managers.RewardTermCfg func: unilab.envs.mdp.joint_vel_l2 weight: 0.0 params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot joint_names: ".*_wheel_joint" alive: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.go2w.manager_terms.constant_alive weight: 0.5 upward: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.go2w.manager_terms.upward_l2 weight: 1.0 diff --git a/conf/ppo/task/quadruped_joystick_rough/base.yaml b/conf/ppo/task/quadruped_joystick_rough/base.yaml index 0c1f99f94..c916ccd63 100644 --- a/conf/ppo/task/quadruped_joystick_rough/base.yaml +++ b/conf/ppo/task/quadruped_joystick_rough/base.yaml @@ -36,16 +36,14 @@ env: heading: [-3.141592653589793, 3.141592653589793] events: reset_scene_to_default: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.reset_scene_to_default mode: reset terrain_root_state: - _target_: unilab.managers.EventTermCfg func: unilab.tasks.locomotion.common.rough_manager_terms.RoughTerrainReset mode: reset params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot pose_range: x: [-0.5, 0.5] @@ -66,24 +64,22 @@ env: cycle_top_frac: 0.5 spawn_height_margin: 0.05 base_mass: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.randomize_rigid_body_mass mode: reset params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot body_names: ".*" mass_distribution_params: [-1.0, 3.0] operation: add recompute_inertia: false base_com: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.randomize_rigid_body_com mode: reset params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot body_names: ".*" com_range: @@ -91,7 +87,6 @@ env: y: [0.0, 0.0] z: [0.0, 0.0] pd_gains: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.pd_gains mode: reset params: @@ -99,7 +94,6 @@ env: kd_range: [0.25, 1.0] operation: abs push_robot: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.push_by_setting_velocity mode: interval interval_range_s: [12.5, 12.5] @@ -114,53 +108,45 @@ env: yaw: [0.0, 0.0] terminations: time_out: - _target_: unilab.managers.TerminationTermCfg func: unilab.envs.mdp.time_out time_out: true terrain_out_of_bounds: - _target_: unilab.managers.TerminationTermCfg func: unilab.tasks.locomotion.common.rough_manager_terms.RoughTerrainOutOfBounds time_out: true params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot distance_buffer: 3.0 curriculum: terrain_levels: - _target_: unilab.managers.CurriculumTermCfg func: unilab.tasks.locomotion.common.rough_manager_terms.RoughTerrainCurriculum params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot policy_observation_group: policy critic_observation_group: critic reward: tracking_lin_vel: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel_xy_exp weight: 3.0 params: std: 0.5 command_name: twist tracking_ang_vel: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel_z_exp weight: 1.5 params: std: 0.5 command_name: twist lin_vel_z: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z_l2 weight: -2.0 ang_vel_xy: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy_l2 weight: -0.05 action_rate: - _target_: unilab.managers.RewardTermCfg func: unilab.envs.mdp.action_rate_l2 weight: -0.01 diff --git a/conf/ppo/task/quadruped_joystick_rough/go2w.yaml b/conf/ppo/task/quadruped_joystick_rough/go2w.yaml index e8d7a6564..746943680 100644 --- a/conf/ppo/task/quadruped_joystick_rough/go2w.yaml +++ b/conf/ppo/task/quadruped_joystick_rough/go2w.yaml @@ -6,77 +6,61 @@ defaults: env: observations: policy: - _target_: unilab.managers.ObservationGroupCfg terms: base_ang_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.builtin_sensor params: {sensor_name: gyro} scale: 0.25 projected_gravity: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.projected_gravity_from_sensor params: {sensor_name: upvector} command: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.generated_commands params: {command_name: twist} leg_joint_pos: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_pos_rel params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot joint_names: ".*_(hip|thigh|calf)_joint" joint_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_vel_rel scale: 0.05 actions: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.last_action params: {action_name: motor} critic: - _target_: unilab.managers.ObservationGroupCfg terms: base_lin_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.builtin_sensor params: {sensor_name: local_linvel} base_ang_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.builtin_sensor params: {sensor_name: gyro} projected_gravity: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.projected_gravity_from_sensor params: {sensor_name: upvector} command: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.generated_commands params: {command_name: twist} leg_joint_pos: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_pos_rel params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot joint_names: ".*_(hip|thigh|calf)_joint" joint_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_vel_rel actions: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.last_action params: {action_name: motor} height_scan: - _target_: unilab.managers.ObservationTermCfg func: unilab.tasks.locomotion.common.rough_manager_terms.RoughHeightScan params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot geom_name: floor vertical_offset: 0.5 @@ -97,7 +81,6 @@ env: events: pd_gains: null motor_gains: - _target_: unilab.managers.EventTermCfg func: unilab.tasks.locomotion.go2w.manager_terms.randomize_motor_gains mode: reset params: @@ -107,44 +90,38 @@ env: reward: orientation: - _target_: unilab.managers.RewardTermCfg func: unilab.envs.mdp.flat_orientation_l2 weight: -2.0 motor_torque: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.go2w.manager_terms.motor_torque_l2 weight: -2.5e-5 params: {action_name: motor} stand_still: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.stand_still_l1 weight: -2.0 params: command_name: twist command_threshold: 0.1 asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot joint_names: ".*_(hip|thigh|calf)_joint" hip_pos: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.rough_manager_terms.joint_deviation_l2 weight: -2.0 params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot joint_names: ".*_hip_joint" joint_pos_penalty: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.rough_manager_terms.joint_deviation_l2 weight: -1.0 params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot joint_names: ".*_(hip|thigh|calf)_joint" upward: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.go2w.manager_terms.upward_l2 weight: 1.0 diff --git a/conf/ppo/task/quadruped_joystick_rough/quadruped.yaml b/conf/ppo/task/quadruped_joystick_rough/quadruped.yaml index 51fe14ed7..a464a98c8 100644 --- a/conf/ppo/task/quadruped_joystick_rough/quadruped.yaml +++ b/conf/ppo/task/quadruped_joystick_rough/quadruped.yaml @@ -6,67 +6,51 @@ defaults: env: observations: policy: - _target_: unilab.managers.ObservationGroupCfg terms: base_ang_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.builtin_sensor params: {sensor_name: gyro} scale: 0.25 projected_gravity: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.projected_gravity_from_sensor params: {sensor_name: upvector} command: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.generated_commands params: {command_name: twist} joint_pos: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_pos_rel joint_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_vel_rel scale: 0.05 actions: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.last_action params: {action_name: joint_pos} critic: - _target_: unilab.managers.ObservationGroupCfg terms: base_lin_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.builtin_sensor params: {sensor_name: local_linvel} base_ang_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.builtin_sensor params: {sensor_name: gyro} projected_gravity: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.projected_gravity_from_sensor params: {sensor_name: upvector} command: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.generated_commands params: {command_name: twist} joint_pos: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_pos_rel joint_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_vel_rel actions: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.last_action params: {action_name: joint_pos} height_scan: - _target_: unilab.managers.ObservationTermCfg func: unilab.tasks.locomotion.common.rough_manager_terms.RoughHeightScan params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot geom_name: floor vertical_offset: 0.5 @@ -84,26 +68,22 @@ env: reward: stand_still: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.manager_terms.stand_still_l1 weight: -2.0 params: command_name: twist command_threshold: 0.1 hip_pos: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.rough_manager_terms.joint_deviation_l2 weight: -0.5 params: asset_cfg: - _target_: unilab.managers.SceneEntityCfg + _target_: SceneEntityCfg name: robot joint_names: ".*_hip_joint" joint_pos_penalty: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.common.rough_manager_terms.joint_deviation_l2 weight: -1.0 upward: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.go2w.manager_terms.upward_l2 weight: 1.0 diff --git a/conf/ppo/task/stewart_balance/base.yaml b/conf/ppo/task/stewart_balance/base.yaml index 87db4efc5..ef8a4c49a 100644 --- a/conf/ppo/task/stewart_balance/base.yaml +++ b/conf/ppo/task/stewart_balance/base.yaml @@ -30,10 +30,8 @@ env: render_spacing: 4.5 observations: policy: - _target_: unilab.managers.ObservationGroupCfg terms: balance: - _target_: unilab.managers.ObservationTermCfg func: unilab.tasks.manipulation.stewart.balance.StewartObservation params: entity_name: stewart @@ -64,11 +62,9 @@ env: center_control_min_gain: 0.15 events: reset_scene_to_default: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.reset_scene_to_default mode: reset reset_ball: - _target_: unilab.managers.EventTermCfg func: unilab.tasks.manipulation.stewart.balance.StewartBallReset mode: reset params: @@ -78,7 +74,6 @@ env: ball_home_z: 1.2 terminations: balance_state: - _target_: unilab.managers.TerminationTermCfg func: unilab.tasks.manipulation.stewart.balance.StewartBalanceState params: observation_group: policy @@ -93,7 +88,6 @@ env: zero_vel_thresh: 0.07 still_steps_needed: 5 time_out: - _target_: unilab.managers.TerminationTermCfg func: unilab.envs.mdp.time_out time_out: true # Legacy Stewart rewards were discrete per-control-step values, not rates. @@ -103,25 +97,21 @@ env: reward: center: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.manipulation.stewart.balance.center_reward weight: 0.7 params: state_term_name: balance_state progress: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.manipulation.stewart.balance.progress_reward weight: 0.6 params: state_term_name: balance_state still: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.manipulation.stewart.balance.still_reward weight: 3.0 params: state_term_name: balance_state fall: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.manipulation.stewart.balance.fall_reward weight: -6.0 params: diff --git a/docs/sphinx/source/en/3-deployment/3-framework_migration/1-from_isaac_lab.md b/docs/sphinx/source/en/3-deployment/3-framework_migration/1-from_isaac_lab.md index 73168650b..7585d81de 100644 --- a/docs/sphinx/source/en/3-deployment/3-framework_migration/1-from_isaac_lab.md +++ b/docs/sphinx/source/en/3-deployment/3-framework_migration/1-from_isaac_lab.md @@ -107,14 +107,11 @@ parameters, weights, and observation mapping in the owner YAML. For example: env: observations: policy: - _target_: unilab.managers.ObservationGroupCfg terms: joint_pos_rel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_pos_rel terminations: time_out: - _target_: unilab.managers.TerminationTermCfg func: unilab.envs.mdp.time_out time_out: true policy_observation_group: policy @@ -122,11 +119,19 @@ env: reward: alive: - _target_: unilab.managers.RewardTermCfg func: unilab.envs.mdp.is_alive weight: 1.0 ``` +Manager mappings whose value type is a single concrete config dataclass +(observations / events / rewards / terminations / curriculum / metrics / +recorders) may omit `_target_`; materialization infers it from the field type +annotation. `actions` / `commands` have abstract base configs, so they must +still declare a concrete `_target_` (for example +`unilab.envs.mdp.JointPositionActionCfg`). Config classes under +`unilab.managers.` (such as `SceneEntityCfg`) may be referenced by their bare +class name. + Hydra composition materializes this declaration into plain typed config on the cold path. Unknown fields, unresolved `_target_`/`func` references, and wrong config types fail before reset or step. Direct Python config construction is diff --git a/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/1-from_isaac_lab.md b/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/1-from_isaac_lab.md index f6914d9ee..ea55f2f58 100644 --- a/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/1-from_isaac_lab.md +++ b/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/1-from_isaac_lab.md @@ -95,14 +95,11 @@ weight 和 observation mapping。例如： env: observations: policy: - _target_: unilab.managers.ObservationGroupCfg terms: joint_pos_rel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_pos_rel terminations: time_out: - _target_: unilab.managers.TerminationTermCfg func: unilab.envs.mdp.time_out time_out: true policy_observation_group: policy @@ -110,11 +107,16 @@ env: reward: alive: - _target_: unilab.managers.RewardTermCfg func: unilab.envs.mdp.is_alive weight: 1.0 ``` +值类型唯一且具体的 manager mapping（observations / events / rewards / +terminations / curriculum / metrics / recorders）可以省略 `_target_`，物化时按字段 +类型注解推断；`actions` / `commands` 的基类是抽象的，必须显式声明具体 `_target_` +（如 `unilab.envs.mdp.JointPositionActionCfg`）。`unilab.managers.` 下的 config 类 +（如 `SceneEntityCfg`）可以直接写裸类名。 + Hydra compose 在冷路径把这份声明物化为 plain typed config。未知字段、无法解析的 `_target_`/`func` 和错误 config 类型都会在 reset/step 之前报错。直接用 Python 构造 config 只用于 focused 底层测试。 diff --git a/src/unilab/base/config_materialization.py b/src/unilab/base/config_materialization.py index 7e9e23517..1fe694952 100644 --- a/src/unilab/base/config_materialization.py +++ b/src/unilab/base/config_materialization.py @@ -66,6 +66,14 @@ def _dict_value_type(annotation: Any) -> Any: def _resolve(reference: Any, *, path: str) -> Any: if not isinstance(reference, str) or not reference.strip(): raise TypeError(f"Config field '{path}' must be a non-empty dotted string") + if "." not in reference: + try: + return get_object(f"unilab.managers.{reference}") + except Exception as exc: + raise ValueError( + f"Config field '{path}' could not resolve short reference {reference!r} " + f"(tried 'unilab.managers.{reference}'): {exc}" + ) from exc try: return get_object(reference) except Exception as exc: @@ -74,6 +82,13 @@ def _resolve(reference: Any, *, path: str) -> Any: ) from exc +def _inferable_target(annotation: Any) -> type[Any] | None: + candidates = _dataclass_types(annotation) + if len(candidates) == 1 and not inspect.isabstract(candidates[0]): + return candidates[0] + return None + + def _resolve_target( reference: Any, *, @@ -155,7 +170,7 @@ def _prepare_manager_mapping(value: Any, *, annotation: Any, path: str) -> dict[ dict(entry), expected=value_type, path=entry_path, - require_target=True, + require_target=_inferable_target(value_type) is None, ) return result @@ -185,6 +200,7 @@ def _prepare_dataclass( reference = _target_path(target) else: target = _resolve_target(reference, expected=expected, path=path) + reference = _target_path(target) fields = {field.name: field for field in dataclasses.fields(target) if field.init} unknown = [name for name in values if name not in fields] @@ -219,8 +235,16 @@ def _prepare_dataclass( return prepared -def _materialize_entry(value: Mapping[str, Any], *, expected: Any, path: str) -> Any: - prepared = _prepare_dataclass(value, expected=expected, path=path, require_target=True) +def _materialize_entry( + value: Mapping[str, Any], + *, + expected: Any, + path: str, + require_target: bool = True, +) -> Any: + prepared = _prepare_dataclass( + value, expected=expected, path=path, require_target=require_target + ) try: result = instantiate(prepared, _convert_="all") except Exception as exc: @@ -299,6 +323,11 @@ def _apply_manager_mapping( else: current = existing.get(term_name, _MISSING) if current is _MISSING or current is None: + if _inferable_target(value_type) is not None: + existing[term_name] = _materialize_entry( + value, expected=value_type, path=path, require_target=False + ) + continue raise ValueError( f"Config field '{path}' is a new Manager-Based entry and must declare " f"'{HYDRA_TARGET_KEY}'" diff --git a/src/unilab/managers/__init__.py b/src/unilab/managers/__init__.py index be3224090..ec031811c 100644 --- a/src/unilab/managers/__init__.py +++ b/src/unilab/managers/__init__.py @@ -3,6 +3,14 @@ # Modified by UniLab for NumPy and UniLab contracts; licensed under Apache-2.0. """Environment managers.""" +from unilab.managers._noise.noise_cfg import ConstantNoiseCfg as ConstantNoiseCfg +from unilab.managers._noise.noise_cfg import GaussianNoiseCfg as GaussianNoiseCfg +from unilab.managers._noise.noise_cfg import NoiseCfg as NoiseCfg +from unilab.managers._noise.noise_cfg import NoiseModelCfg as NoiseModelCfg +from unilab.managers._noise.noise_cfg import ( + NoiseModelWithAdditiveBiasCfg as NoiseModelWithAdditiveBiasCfg, +) +from unilab.managers._noise.noise_cfg import UniformNoiseCfg as UniformNoiseCfg from unilab.managers.action_manager import ActionManager as ActionManager from unilab.managers.action_manager import ActionTerm as ActionTerm from unilab.managers.action_manager import ActionTermCfg as ActionTermCfg diff --git a/tests/base/backend/test_mujoco_cpu_affinity_wiring.py b/tests/base/backend/test_mujoco_cpu_affinity_wiring.py index 7dad28f1f..8be9e1b0c 100644 --- a/tests/base/backend/test_mujoco_cpu_affinity_wiring.py +++ b/tests/base/backend/test_mujoco_cpu_affinity_wiring.py @@ -34,6 +34,13 @@ _MODEL_FILE = str(ASSETS_ROOT_PATH / "robots" / "go2_arm" / "scene_flat.xml") _NUM_ENVS = 4 + +if not hasattr(os, "sched_getaffinity"): + pytest.skip( + "os.sched_getaffinity unavailable on this platform (e.g. macOS)", + allow_module_level=True, + ) + _AVAILABLE_CPUS = sorted(os.sched_getaffinity(0)) diff --git a/tests/base/test_manager_config_overlay.py b/tests/base/test_manager_config_overlay.py index 443ce5ec2..8722c09bc 100644 --- a/tests/base/test_manager_config_overlay.py +++ b/tests/base/test_manager_config_overlay.py @@ -280,8 +280,7 @@ def test_hydra_mapping_fully_materializes_empty_manager_config() -> None: @pytest.mark.parametrize( ("overrides", "match"), [ - ({"rewards": {"missing": {"weight": 1.0}}}, "missing.*_target_"), - ({"rewards": {"disabled": {"weight": 1.0}}}, "disabled.*_target_"), + ({"actions": {"missing": {"scale": 0.25}}}, "missing.*_target_"), ({"rewards": {"tracking": _second_term}}, "tracking.*field mapping"), ({"rewards": []}, "rewards.*mapping"), ({"rewards": {"tracking": {"func": _second_term}}}, "func.*typed term"), @@ -292,6 +291,88 @@ def test_manager_mapping_overlay_fails_closed(overrides: dict, match: str) -> No apply_cfg_overrides(_manager_cfg(), overrides) +def test_manager_mapping_overlay_infers_concrete_term_target() -> None: + cfg = _manager_cfg() + + apply_cfg_overrides( + cfg, + { + "rewards": { + "missing": { + "func": "unilab.envs.mdp.is_alive", + "weight": 0.5, + }, + "disabled": { + "func": "unilab.envs.mdp.is_alive", + "weight": 0.2, + }, + }, + "observations": { + "critic": { + "terms": { + "joint_vel": {"func": "unilab.envs.mdp.joint_vel_rel"}, + } + } + }, + }, + ) + + missing = cfg.rewards["missing"] + assert isinstance(missing, RewardTermCfg) + assert callable(missing.func) + assert missing.weight == pytest.approx(0.5) + disabled = cfg.rewards["disabled"] + assert isinstance(disabled, RewardTermCfg) + assert disabled.weight == pytest.approx(0.2) + assert list(cfg.rewards) == ["tracking", "alive", "disabled", "missing"] + critic = cfg.observations["critic"] + assert isinstance(critic, ObservationGroupCfg) + joint_vel = critic.terms["joint_vel"] + assert isinstance(joint_vel, ObservationTermCfg) + assert callable(joint_vel.func) + + +def test_hydra_materialization_resolves_managers_short_name() -> None: + cfg = ManagerBasedRlEnvCfg() + + apply_cfg_overrides( + cfg, + { + "events": { + "base_mass": { + "func": "unilab.envs.mdp.is_alive", + "mode": "reset", + "params": { + "asset_cfg": {"_target_": "SceneEntityCfg", "name": "robot"}, + }, + } + } + }, + ) + + base_mass = cfg.events["base_mass"] + assert isinstance(base_mass, EventTermCfg) + asset_cfg = base_mass.params["asset_cfg"] + assert isinstance(asset_cfg, SceneEntityCfg) + assert asset_cfg.name == "robot" + + +def test_hydra_materialization_rejects_unknown_short_name() -> None: + with pytest.raises(ValueError, match="could not resolve"): + apply_cfg_overrides( + ManagerBasedRlEnvCfg(), + { + "rewards": { + "term": { + "_target_": "NoSuchCfg", + "func": "unilab.envs.mdp.is_alive", + "weight": 1.0, + } + } + }, + ) + + @pytest.mark.parametrize( ("entry", "match"), [ diff --git a/tests/config/test_config_system.py b/tests/config/test_config_system.py index 9548dc595..32baef7fd 100644 --- a/tests/config/test_config_system.py +++ b/tests/config/test_config_system.py @@ -72,7 +72,6 @@ def _assert_reward_populated(cfg, label: str): assert active_terms, f"{label} Manager-Based reward terms must be non-empty" for term_name, term in active_terms.items(): assert isinstance(term, dict), f"{label} reward.{term_name} must be a mapping" - assert "_target_" in term, f"{label} reward.{term_name} must declare _target_" assert "func" in term, f"{label} reward.{term_name} must declare func" assert "weight" in term, f"{label} reward.{term_name} must declare weight" From 05c6d2947440ebb06e5c5c5d8ba1a655eaa0f991 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Thu, 20 Aug 2026 14:18:00 +0800 Subject: [PATCH 091/169] feat: migrate G1 locomotion family to Manager-Based runtime (#1232) --- AGENTS.md | 2 +- conf/appo/task/g1_23dof_walk_flat/base.yaml | 66 ++ conf/appo/task/g1_23dof_walk_flat/mujoco.yaml | 35 +- conf/appo/task/g1_walk_flat/base.yaml | 270 +++++ conf/appo/task/g1_walk_flat/mujoco.yaml | 35 +- .../flashsac/g1_23dof_walk_flat/motrix.yaml | 76 +- .../flashsac/g1_23dof_walk_flat/mujoco.yaml | 56 +- .../task/flashsac/g1_walk_flat/mjwarp.yaml | 70 +- .../task/flashsac/g1_walk_flat/motrix.yaml | 79 +- .../task/flashsac/g1_walk_flat/mujoco.yaml | 58 +- .../task/g1_23dof_walk_flat/base.yaml | 66 ++ conf/offpolicy/task/g1_walk_flat/base.yaml | 285 +++++ .../task/sac/g1_23dof_walk_flat/motrix.yaml | 63 +- .../task/sac/g1_23dof_walk_flat/mujoco.yaml | 46 +- .../task/sac/g1_23dof_walk_rough/motrix.yaml | 65 +- .../task/sac/g1_23dof_walk_rough/mujoco.yaml | 47 +- .../task/sac/g1_walk_flat/mjwarp.yaml | 58 +- .../task/sac/g1_walk_flat/motrix.yaml | 71 +- .../task/sac/g1_walk_flat/mujoco.yaml | 47 +- .../task/sac/g1_walk_rough/motrix.yaml | 65 +- .../task/sac/g1_walk_rough/mujoco.yaml | 48 +- .../task/td3/g1_23dof_walk_flat/mujoco.yaml | 46 +- .../task/td3/g1_walk_flat/mujoco.yaml | 47 +- conf/ppo/task/g1_23dof_walk_flat/base.yaml | 66 ++ conf/ppo/task/g1_23dof_walk_flat/motrix.yaml | 147 ++- conf/ppo/task/g1_23dof_walk_flat/mujoco.yaml | 35 +- conf/ppo/task/g1_23dof_walk_rough/mujoco.yaml | 52 +- conf/ppo/task/g1_walk_flat/base.yaml | 270 +++++ conf/ppo/task/g1_walk_flat/mjwarp.yaml | 47 +- conf/ppo/task/g1_walk_flat/motrix.yaml | 147 ++- conf/ppo/task/g1_walk_flat/mujoco.yaml | 37 +- .../5-domain_randomization/0-index.md | 16 +- .../5-domain_randomization/1-configuration.md | 6 +- .../5-domain_randomization/0-index.md | 16 +- .../5-domain_randomization/1-configuration.md | 6 +- .../9-sim2sim_contract_status.md | 7 +- scripts/benchmark/core/task_names.py | 4 +- scripts/benchmark/env/benchmark_env_step.py | 14 +- scripts/benchmark/torch_env/walk_flat.py | 3 +- src/unilab/tasks/locomotion/g1/__init__.py | 39 +- src/unilab/tasks/locomotion/g1/joystick.py | 717 ------------ .../tasks/locomotion/g1/manager_terms.py | 1021 +++++++++++++++++ src/unilab/tasks/migration_matrix.py | 8 +- .../tasks/motion_tracking/common/config.py | 2 +- .../common/g1_base.py} | 7 +- .../tasks/motion_tracking/common/tracking.py | 4 +- .../tasks/motion_tracking/g1/tracking_obs.py | 2 +- .../tasks/motion_tracking/x2/flip_tracking.py | 2 +- tests/algos/test_rsl_rl_runner.py | 28 +- tests/base/test_reward_override.py | 40 +- ...st_offpolicy_collector_active_benchmark.py | 11 +- tests/config/test_config_system.py | 56 +- tests/config/test_locomotion_params.py | 30 +- tests/config/test_reward_injection.py | 61 +- tests/conftest.py | 121 -- .../locomotion/g1/test_g1_owner_contract.py | 900 +++++++++++---- .../envs/locomotion/g1/test_manager_terms.py | 313 +++++ .../locomotion/g1/test_symmetry_contract.py | 83 +- tests/envs/test_env_configs.py | 233 +--- tests/envs/test_g1_obs_noise.py | 109 -- .../test_reward_injection_integration.py | 10 +- tests/scripts/test_train_scripts.py | 56 +- tests/tasks/test_migration_matrix.py | 4 +- tests/training/test_sim2sim_resolver.py | 4 +- tests/training/test_training_helpers.py | 5 +- 65 files changed, 3865 insertions(+), 2575 deletions(-) create mode 100644 conf/appo/task/g1_23dof_walk_flat/base.yaml create mode 100644 conf/appo/task/g1_walk_flat/base.yaml create mode 100644 conf/offpolicy/task/g1_23dof_walk_flat/base.yaml create mode 100644 conf/offpolicy/task/g1_walk_flat/base.yaml create mode 100644 conf/ppo/task/g1_23dof_walk_flat/base.yaml create mode 100644 conf/ppo/task/g1_walk_flat/base.yaml delete mode 100644 src/unilab/tasks/locomotion/g1/joystick.py create mode 100644 src/unilab/tasks/locomotion/g1/manager_terms.py rename src/unilab/tasks/{locomotion/g1/base.py => motion_tracking/common/g1_base.py} (86%) create mode 100644 tests/envs/locomotion/g1/test_manager_terms.py delete mode 100644 tests/envs/test_g1_obs_noise.py diff --git a/AGENTS.md b/AGENTS.md index 4dd206c14..77dd3986f 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -44,7 +44,7 @@ UniLab 是一个 **高性能、模块化、contract 驱动** 的 RL infrastructu - **WARNING_LIST**：`reward.*`、`env.control_config.simulate_action_latency`、`env.ctrl_dt`。 - **ALLOWLIST**（自由覆盖）：`training.sim_backend`、`env.scene`、`training.play_steps`、`env.domain_rand`、`env.noise_config`、`env.commands.vel_limit`。 -训练时 `ExperimentTracker.start()` 把上述字段写入 `run_config.json` 的 `contract_snapshot`（不改 checkpoint 格式，旧 run 无 snapshot 时 fallback + warning）；五个 play 入口在建 env 前调用 `resolve_sim2sim_config` 校验，并用 `policy_load_dim_guard` 包裹 checkpoint 加载以把维度不匹配的隐晦报错重抛为显式诊断。设 `training.sim2sim_strict=false` 可把 DENYLIST 差异降级为 warning（默认 `true`）。DENYLIST 字段在每个后端 owner 配置中显式声明并保持跨后端一致（范例：`conf/ppo/task/g1_walk_flat/{mujoco,motrix}.yaml`）；跨后端契约审计见 `scripts/audit_sim2sim_contracts.py`。 +训练时 `ExperimentTracker.start()` 把上述字段写入 `run_config.json` 的 `contract_snapshot`（不改 checkpoint 格式，旧 run 无 snapshot 时 fallback + warning）；五个 play 入口在建 env 前调用 `resolve_sim2sim_config` 校验，并用 `policy_load_dim_guard` 包裹 checkpoint 加载以把维度不匹配的隐晦报错重抛为显式诊断。设 `training.sim2sim_strict=false` 可把 DENYLIST 差异降级为 warning（默认 `true`）。DENYLIST 字段在共享 base owner 与后端 owner 配置中显式声明并保持跨后端一致（范例：`conf/ppo/task/g1_walk_flat/{base,mujoco,motrix}.yaml`）；跨后端契约审计见 `scripts/audit_sim2sim_contracts.py`。 ## Pointers diff --git a/conf/appo/task/g1_23dof_walk_flat/base.yaml b/conf/appo/task/g1_23dof_walk_flat/base.yaml new file mode 100644 index 000000000..66c3da41e --- /dev/null +++ b/conf/appo/task/g1_23dof_walk_flat/base.yaml @@ -0,0 +1,66 @@ +# @package _global_ +# Canonical G1 23-DoF walk Manager-Based task declaration (APPO owners; mirrors the PPO base). +# Inherits the 29-DoF flat contract and swaps the scene to the 23-DoF model +# (no waist roll/pitch, no wrist pitch/yaw) with the 23-entry pose weights. +defaults: + - /task/g1_walk_flat/base + - _self_ + +env: + scene: + model_file: src/unilab/assets/robots/g1/scene_flat_23dof.xml + entities: + robot: + joint_names: + - left_hip_pitch_joint + - left_hip_roll_joint + - left_hip_yaw_joint + - left_knee_joint + - left_ankle_pitch_joint + - left_ankle_roll_joint + - right_hip_pitch_joint + - right_hip_roll_joint + - right_hip_yaw_joint + - right_knee_joint + - right_ankle_pitch_joint + - right_ankle_roll_joint + - waist_yaw_joint + - left_shoulder_pitch_joint + - left_shoulder_roll_joint + - left_shoulder_yaw_joint + - left_elbow_joint + - left_wrist_roll_joint + - right_shoulder_pitch_joint + - right_shoulder_roll_joint + - right_shoulder_yaw_joint + - right_elbow_joint + - right_wrist_roll_joint + actuator_names: + - left_hip_pitch_joint + - left_hip_roll_joint + - left_hip_yaw_joint + - left_knee_joint + - left_ankle_pitch_joint + - left_ankle_roll_joint + - right_hip_pitch_joint + - right_hip_roll_joint + - right_hip_yaw_joint + - right_knee_joint + - right_ankle_pitch_joint + - right_ankle_roll_joint + - waist_yaw_joint + - left_shoulder_pitch_joint + - left_shoulder_roll_joint + - left_shoulder_yaw_joint + - left_elbow_joint + - left_wrist_roll_joint + - right_shoulder_pitch_joint + - right_shoulder_roll_joint + - right_shoulder_yaw_joint + - right_elbow_joint + - right_wrist_roll_joint + +reward: + pose: + params: + pose_weights: [0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0] diff --git a/conf/appo/task/g1_23dof_walk_flat/mujoco.yaml b/conf/appo/task/g1_23dof_walk_flat/mujoco.yaml index e68cd0af1..aa34c3339 100644 --- a/conf/appo/task/g1_23dof_walk_flat/mujoco.yaml +++ b/conf/appo/task/g1_23dof_walk_flat/mujoco.yaml @@ -1,36 +1,13 @@ # @package _global_ +# MuJoCo APPO owner: inherits the shared 23-DoF flat Manager-Based contract +# from base.yaml and only carries backend/algo identity. +defaults: + - /task/g1_23dof_walk_flat/base + - _self_ + training: task_name: G1Walk23DofFlat sim_backend: mujoco algo: max_iterations: 500 save_interval: 100 -env: - control_config: - action_scale: 0.25 - curriculum: - enabled: false - noise_config: - level: 1.0 - scale_joint_angle: 0.01 - scale_joint_vel: 1.5 - scale_gyro: 0.2 -reward: - scales: - tracking_lin_vel: 2.0 - tracking_ang_vel: 0.2 - feet_phase: 1.0 - lin_vel_z: -1.0 - ang_vel_xy: -0.25 - base_height: -500.0 - orientation: -5.0 - action_rate: -0.01 - pose: -0.1 - tracking_sigma: 0.25 - gait_frequency: 1.5 - feet_phase_swing_height: 0.09 - feet_phase_tracking_sigma: 0.008 - base_height_target: 0.754 - min_base_height: 0.55 - max_tilt_deg: 25.0 - pose_weights: [0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0] diff --git a/conf/appo/task/g1_walk_flat/base.yaml b/conf/appo/task/g1_walk_flat/base.yaml new file mode 100644 index 000000000..8c74c34c1 --- /dev/null +++ b/conf/appo/task/g1_walk_flat/base.yaml @@ -0,0 +1,270 @@ +# @package _global_ +# Canonical G1 29-DoF walk Manager-Based task declaration (APPO owners; mirrors the PPO base). +# Backend owner leaves inherit this file and only override backend/algo tuning +# or explicitly disabled terms. Observation scaling follows the legacy profile +# (unit scales); the walk profile lives in the offpolicy base owners. +env: + scene: + model_file: src/unilab/assets/robots/g1/scene_flat.xml + default_keyframe_name: stand + entities: + robot: + root_body_name: pelvis + joint_names: + - left_hip_pitch_joint + - left_hip_roll_joint + - left_hip_yaw_joint + - left_knee_joint + - left_ankle_pitch_joint + - left_ankle_roll_joint + - right_hip_pitch_joint + - right_hip_roll_joint + - right_hip_yaw_joint + - right_knee_joint + - right_ankle_pitch_joint + - right_ankle_roll_joint + - waist_yaw_joint + - waist_roll_joint + - waist_pitch_joint + - left_shoulder_pitch_joint + - left_shoulder_roll_joint + - left_shoulder_yaw_joint + - left_elbow_joint + - left_wrist_roll_joint + - left_wrist_pitch_joint + - left_wrist_yaw_joint + - right_shoulder_pitch_joint + - right_shoulder_roll_joint + - right_shoulder_yaw_joint + - right_elbow_joint + - right_wrist_roll_joint + - right_wrist_pitch_joint + - right_wrist_yaw_joint + actuator_names: + - left_hip_pitch_joint + - left_hip_roll_joint + - left_hip_yaw_joint + - left_knee_joint + - left_ankle_pitch_joint + - left_ankle_roll_joint + - right_hip_pitch_joint + - right_hip_roll_joint + - right_hip_yaw_joint + - right_knee_joint + - right_ankle_pitch_joint + - right_ankle_roll_joint + - waist_yaw_joint + - waist_roll_joint + - waist_pitch_joint + - left_shoulder_pitch_joint + - left_shoulder_roll_joint + - left_shoulder_yaw_joint + - left_elbow_joint + - left_wrist_roll_joint + - left_wrist_pitch_joint + - left_wrist_yaw_joint + - right_shoulder_pitch_joint + - right_shoulder_roll_joint + - right_shoulder_yaw_joint + - right_elbow_joint + - right_wrist_roll_joint + - right_wrist_pitch_joint + - right_wrist_yaw_joint + body_names: [pelvis] + sim_dt: 0.006666666666666667 + ctrl_dt: 0.02 + max_episode_seconds: 20.0 + observations: + policy: + _target_: unilab.managers.ObservationGroupCfg + terms: + base_ang_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.builtin_sensor + params: {sensor_name: torso_gyro} + projected_gravity: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.projected_gravity_from_sensor + params: {sensor_name: torso_upvector} + joint_pos: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_pos_rel + joint_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_vel_rel + actions: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.last_action + params: {action_name: joint_pos} + command: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.generated_commands + params: {command_name: twist} + gait_phase: + _target_: unilab.managers.ObservationTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.G1GaitPhase + params: + frequency: 1.5 + init_mode: offset_phase + critic: + _target_: unilab.managers.ObservationGroupCfg + terms: + base_ang_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.builtin_sensor + params: {sensor_name: torso_gyro} + projected_gravity: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.projected_gravity_from_sensor + params: {sensor_name: torso_upvector} + joint_pos: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_pos_rel + joint_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_vel_rel + actions: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.last_action + params: {action_name: joint_pos} + command: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.generated_commands + params: {command_name: twist} + gait_phase: + _target_: unilab.managers.ObservationTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.G1GaitPhase + params: + frequency: 1.5 + init_mode: offset_phase + base_lin_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.builtin_sensor + params: {sensor_name: pelvis_local_linvel} + actions: + joint_pos: + _target_: unilab.envs.mdp.JointPositionActionCfg + entity_name: robot + actuator_names: [".*"] + scale: 0.25 + use_default_offset: true + commands: + twist: + _target_: unilab.tasks.locomotion.g1.manager_terms.G1VelocityCommandCfg + entity_name: robot + resampling_time_range: [20.0, 20.0] + heading_command: false + heading_control_stiffness: 0.5 + rel_standing_envs: 0.0 + rel_heading_envs: 0.0 + rel_world_envs: 0.0 + rel_forward_envs: 0.0 + init_velocity_prob: 0.0 + planar_dead_zone: 0.2 + ranges: + lin_vel_x: [-0.6, 1.0] + lin_vel_y: [-0.4, 0.4] + ang_vel_z: [-0.8, 0.8] + events: + reset_scene_to_default: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.reset_scene_to_default + mode: reset + reset_root_state_uniform: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.reset_root_state_uniform + mode: reset + params: + pose_range: + x: [-0.5, 0.5] + y: [-0.5, 0.5] + z: [0.0, 0.0] + roll: [0.0, 0.0] + pitch: [0.0, 0.0] + yaw: [-3.141592653589793, 3.141592653589793] + velocity_range: + x: [-0.5, 0.5] + y: [-0.5, 0.5] + z: [-0.5, 0.5] + roll: [-0.5, 0.5] + pitch: [-0.5, 0.5] + yaw: [-0.5, 0.5] + pd_gains: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.pd_gains + mode: reset + params: + kp_range: [0.9, 1.1] + kd_range: [0.9, 1.1] + operation: scale + terminations: + time_out: + _target_: unilab.managers.TerminationTermCfg + func: unilab.envs.mdp.time_out + time_out: true + tilt: + _target_: unilab.managers.TerminationTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.g1_tilt_exceeded + params: + max_tilt_deg: 25.0 + base_height: + _target_: unilab.managers.TerminationTermCfg + func: unilab.envs.mdp.root_height_below_minimum + params: + minimum_height: 0.55 + policy_observation_group: policy + critic_observation_group: critic + +reward: + tracking_lin_vel: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.track_lin_vel + weight: 2.0 + params: + tracking_sigma: 0.25 + command_name: twist + tracking_ang_vel: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.track_ang_vel + weight: 0.2 + params: + tracking_sigma: 0.25 + command_name: twist + feet_phase: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.feet_phase + weight: 1.0 + params: + frequency: 1.5 + swing_height: 0.09 + tracking_sigma: 0.008 + min_forward_speed: 0.0 + command_name: twist + lin_vel_z: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.lin_vel_z + weight: -1.0 + ang_vel_xy: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.ang_vel_xy + weight: -0.25 + base_height: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.base_height + weight: -500.0 + params: + target_height: 0.754 + orientation: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.orientation + weight: -5.0 + action_rate: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.mdp.action_rate_l2 + weight: -0.01 + pose: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.weighted_pose + weight: -0.1 + params: + pose_weights: [0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0] diff --git a/conf/appo/task/g1_walk_flat/mujoco.yaml b/conf/appo/task/g1_walk_flat/mujoco.yaml index 0fce8bf11..6ee7b66ed 100644 --- a/conf/appo/task/g1_walk_flat/mujoco.yaml +++ b/conf/appo/task/g1_walk_flat/mujoco.yaml @@ -1,36 +1,13 @@ # @package _global_ +# MuJoCo APPO owner: inherits the shared 29-DoF flat Manager-Based contract +# from base.yaml and only carries backend/algo identity. +defaults: + - /task/g1_walk_flat/base + - _self_ + training: task_name: G1WalkFlat sim_backend: mujoco algo: max_iterations: 500 save_interval: 100 -env: - control_config: - action_scale: 0.25 - curriculum: - enabled: false - noise_config: - level: 1.0 - scale_joint_angle: 0.01 - scale_joint_vel: 1.5 - scale_gyro: 0.2 -reward: - scales: - tracking_lin_vel: 2.0 - tracking_ang_vel: 0.2 - feet_phase: 1.0 - lin_vel_z: -1.0 - ang_vel_xy: -0.25 - base_height: -500.0 - orientation: -5.0 - action_rate: -0.01 - pose: -0.1 - tracking_sigma: 0.25 - gait_frequency: 1.5 - feet_phase_swing_height: 0.09 - feet_phase_tracking_sigma: 0.008 - base_height_target: 0.754 - min_base_height: 0.55 - max_tilt_deg: 25.0 - pose_weights: [0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0] diff --git a/conf/offpolicy/task/flashsac/g1_23dof_walk_flat/motrix.yaml b/conf/offpolicy/task/flashsac/g1_23dof_walk_flat/motrix.yaml index d1698fb4c..487f03787 100644 --- a/conf/offpolicy/task/flashsac/g1_23dof_walk_flat/motrix.yaml +++ b/conf/offpolicy/task/flashsac/g1_23dof_walk_flat/motrix.yaml @@ -1,8 +1,10 @@ # @package _global_ -# Motrix owner for FlashSAC G1 23-DoF walk flat. -# Mirrors 29-DoF flashsac/g1_walk_flat/motrix.yaml: -# - Keeps the mujoco owner's FlashSAC algo identity -# - Adopts the Motrix-direction env + reward tuning (kp/kd rand off, retuned shaping) +# FlashSAC Motrix 23-DoF owner: mirrors flashsac/g1_walk_flat/motrix.yaml with +# the 23-DoF contract (algo identity kept, Motrix-direction retuning). +defaults: + - /task/g1_23dof_walk_flat/base + - _self_ + training: task_name: G1Walk23DofFlat sim_backend: motrix @@ -15,45 +17,29 @@ algo: replay_buffer_n: 256 tau: 0.05 env: - control_config: - action_scale: 1.0 - gait_phase_init_mode: "offset_phase" - reset_base_qvel_limit: 0.5 - domain_rand: - randomize_kp: false - randomize_kd: false - curriculum: - enabled: true - initial_scale: 0.5 - min_scale: 0.5 - max_scale: 1.0 - level_down_threshold: 150.0 - level_up_threshold: 750.0 - degree: 0.001 - noise_config: - level: 0.0 - scale_gyro: 0.0 - scale_gravity: 0.0 - scale_joint_angle: 0.01 - scale_joint_vel: 0.1 - scale_linvel: 0.0 + events: + # Legacy Motrix owners disable kp/kd randomization. + pd_gains: null reward: - scales: - tracking_lin_vel: 2.2 - tracking_ang_vel: 1.8 - penalty_ang_vel_xy: -1.2 - penalty_orientation: -12.0 - penalty_action_rate: -2.5 - pose: -0.6 - penalty_feet_ori: -5.0 - feet_phase: 6.0 - alive: 12.0 - tracking_sigma: 0.25 - base_height_target: 0.754 - min_base_height: 0.3 - max_tilt_deg: 65.0 - gait_frequency: 1.5 - feet_phase_swing_height: 0.09 - feet_phase_tracking_sigma: 0.008 - close_feet_threshold: 0.15 - pose_weights: [0.01, 2.0, 5.0, 0.01, 5.0, 5.0, 0.01, 2.0, 5.0, 0.01, 5.0, 5.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0] + tracking_lin_vel: + weight: 2.2 + tracking_ang_vel: + weight: 1.8 + penalty_ang_vel_xy: + weight: -1.2 + penalty_orientation: + weight: -12.0 + penalty_action_rate: + weight: -2.5 + pose: + weight: -0.6 + params: + pose_weights: [0.01, 2.0, 5.0, 0.01, 5.0, 5.0, 0.01, 2.0, 5.0, 0.01, 5.0, 5.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0] + penalty_feet_ori: + weight: -5.0 + feet_phase: + weight: 6.0 + params: + tracking_sigma: 0.008 + alive: + weight: 12.0 diff --git a/conf/offpolicy/task/flashsac/g1_23dof_walk_flat/mujoco.yaml b/conf/offpolicy/task/flashsac/g1_23dof_walk_flat/mujoco.yaml index 34e63d24d..4045dd65a 100644 --- a/conf/offpolicy/task/flashsac/g1_23dof_walk_flat/mujoco.yaml +++ b/conf/offpolicy/task/flashsac/g1_23dof_walk_flat/mujoco.yaml @@ -1,4 +1,10 @@ # @package _global_ +# FlashSAC MuJoCo 23-DoF owner: 23-DoF offpolicy contract plus the FlashSAC +# algo identity and its reward retuning. +defaults: + - /task/g1_23dof_walk_flat/base + - _self_ + training: task_name: G1Walk23DofFlat sim_backend: mujoco @@ -8,46 +14,16 @@ algo: max_iterations: 5000 save_interval: 1000 updates_per_step: 8 - #use_symmetry: true replay_buffer_n: 256 tau: 0.05 -env: - control_config: - action_scale: 1.0 - gait_phase_init_mode: "offset_phase" - reset_base_qvel_limit: 0.5 - curriculum: - enabled: true - initial_scale: 0.5 - min_scale: 0.5 - max_scale: 1.0 - level_down_threshold: 150.0 - level_up_threshold: 750.0 - degree: 0.001 - noise_config: - level: 0.0 - scale_gyro: 0.0 - scale_gravity: 0.0 - scale_joint_angle: 0.01 - scale_joint_vel: 0.1 - scale_linvel: 0.0 reward: - scales: - tracking_lin_vel: 2.0 - tracking_ang_vel: 1.5 - penalty_ang_vel_xy: -1.0 - penalty_orientation: -10.0 - penalty_action_rate: -5.0 - pose: -0.5 - penalty_feet_ori: -25.0 - feet_phase: 5.0 - alive: 10.0 - tracking_sigma: 0.25 - base_height_target: 0.754 - min_base_height: 0.3 - max_tilt_deg: 65.0 - gait_frequency: 1.5 - feet_phase_swing_height: 0.09 - feet_phase_tracking_sigma: 0.005 - close_feet_threshold: 0.15 - pose_weights: [0.01, 2.0, 5.0, 0.01, 5.0, 5.0, 0.01, 2.0, 5.0, 0.01, 5.0, 5.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0] + penalty_action_rate: + weight: -5.0 + penalty_feet_ori: + weight: -25.0 + feet_phase: + params: + tracking_sigma: 0.005 + pose: + params: + pose_weights: [0.01, 2.0, 5.0, 0.01, 5.0, 5.0, 0.01, 2.0, 5.0, 0.01, 5.0, 5.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0] diff --git a/conf/offpolicy/task/flashsac/g1_walk_flat/mjwarp.yaml b/conf/offpolicy/task/flashsac/g1_walk_flat/mjwarp.yaml index 4a7602acb..35efa5faa 100644 --- a/conf/offpolicy/task/flashsac/g1_walk_flat/mjwarp.yaml +++ b/conf/offpolicy/task/flashsac/g1_walk_flat/mjwarp.yaml @@ -1,8 +1,12 @@ # @package _global_ -# Configured-only mjwarp owner for FlashSAC G1 walk flat. Mirrors the mujoco -# owner's algo / env / reward identity; mjwarp-specific host-adapter settings -# follow conf/offpolicy/task/sac/g1_walk_flat/mjwarp.yaml. Offline record -# reuses MuJoCo rendering; native playback and device-resident runtime are absent. +# Configured-only FlashSAC mjwarp owner. Mirrors the MuJoCo owner's algo / +# env / reward identity; mjwarp-specific host-adapter settings follow +# conf/offpolicy/task/sac/g1_walk_flat/mjwarp.yaml. Offline record reuses +# MuJoCo rendering; native playback and device-resident runtime are absent. +defaults: + - /task/g1_walk_flat/base + - _self_ + training: task_name: G1WalkFlat sim_backend: mjwarp @@ -13,55 +17,23 @@ algo: max_iterations: 5000 save_interval: 1000 updates_per_step: 8 - #use_symmetry: true replay_buffer_n: 256 tau: 0.05 env: mjwarp_nconmax: 128 mjwarp_njmax: 256 render_spacing: 2.0 - control_config: - action_scale: 1.0 - gait_phase_init_mode: "offset_phase" - reset_base_qvel_limit: 0.5 - domain_rand: - randomize_kp: false - randomize_kd: false - randomize_dof_armature: false - randomize_body_gravity_compensation: false - curriculum: - enabled: true - initial_scale: 0.5 - min_scale: 0.5 - max_scale: 1.0 - level_down_threshold: 150.0 - level_up_threshold: 750.0 - degree: 0.001 - noise_config: - level: 0.0 - scale_gyro: 0.0 - scale_gravity: 0.0 - scale_joint_angle: 0.01 - scale_joint_vel: 0.1 - scale_linvel: 0.0 - seed: null + events: + # Legacy mjwarp owners disable kp/kd and armature randomization. + pd_gains: null reward: - scales: - tracking_lin_vel: 2.0 - tracking_ang_vel: 1.5 - penalty_ang_vel_xy: -1.0 - penalty_orientation: -10.0 - penalty_action_rate: -5.0 - pose: -0.5 - penalty_feet_ori: -25.0 - feet_phase: 5.0 - alive: 10.0 - tracking_sigma: 0.25 - base_height_target: 0.754 - min_base_height: 0.3 - max_tilt_deg: 65.0 - gait_frequency: 1.5 - feet_phase_swing_height: 0.09 - feet_phase_tracking_sigma: 0.005 - close_feet_threshold: 0.15 - pose_weights: [0.01, 2.0, 5.0, 0.01, 5.0, 5.0, 0.01, 2.0, 5.0, 0.01, 5.0, 5.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0] + penalty_action_rate: + weight: -5.0 + penalty_feet_ori: + weight: -25.0 + feet_phase: + params: + tracking_sigma: 0.005 + pose: + params: + pose_weights: [0.01, 2.0, 5.0, 0.01, 5.0, 5.0, 0.01, 2.0, 5.0, 0.01, 5.0, 5.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0] diff --git a/conf/offpolicy/task/flashsac/g1_walk_flat/motrix.yaml b/conf/offpolicy/task/flashsac/g1_walk_flat/motrix.yaml index 5f61c51b6..c8d8e6b8d 100644 --- a/conf/offpolicy/task/flashsac/g1_walk_flat/motrix.yaml +++ b/conf/offpolicy/task/flashsac/g1_walk_flat/motrix.yaml @@ -1,10 +1,10 @@ # @package _global_ -# Motrix owner for FlashSAC G1 walk flat. -# Keeps the mujoco owner's FlashSAC algo identity (num_envs/updates_per_step/ -# replay_buffer_n/tau plus the distributional critic from conf/offpolicy/algo/flashsac.yaml) -# while adopting the Motrix-direction env + reward tuning used by the SAC motrix owner -# (kp/kd randomization off, retuned reward shaping, tighter feet-phase sigma). -# control_config.action_scale stays 1.0 to keep sim2sim DENYLIST parity with mujoco. +# FlashSAC Motrix owner: keeps the MuJoCo owner's FlashSAC algo identity while +# adopting the Motrix-direction reward retuning; kp/kd randomization disabled. +defaults: + - /task/g1_walk_flat/base + - _self_ + training: task_name: G1WalkFlat sim_backend: motrix @@ -17,46 +17,29 @@ algo: replay_buffer_n: 256 tau: 0.05 env: - control_config: - action_scale: 1.0 - gait_phase_init_mode: "offset_phase" - reset_base_qvel_limit: 0.5 - domain_rand: - randomize_kp: false - randomize_kd: false - curriculum: - enabled: true - initial_scale: 0.5 - min_scale: 0.5 - max_scale: 1.0 - level_down_threshold: 150.0 - level_up_threshold: 750.0 - degree: 0.001 - noise_config: - level: 0.0 - scale_gyro: 0.0 - scale_gravity: 0.0 - scale_joint_angle: 0.01 - scale_joint_vel: 0.1 - scale_linvel: 0.0 - seed: null + events: + # Legacy Motrix owners disable kp/kd randomization. + pd_gains: null reward: - scales: - tracking_lin_vel: 2.2 - tracking_ang_vel: 1.8 - penalty_ang_vel_xy: -1.2 - penalty_orientation: -12.0 - penalty_action_rate: -2.5 - pose: -0.6 - penalty_feet_ori: -5.0 - feet_phase: 6.0 - alive: 12.0 - tracking_sigma: 0.25 - base_height_target: 0.754 - min_base_height: 0.3 - max_tilt_deg: 65.0 - gait_frequency: 1.5 - feet_phase_swing_height: 0.09 - feet_phase_tracking_sigma: 0.008 - close_feet_threshold: 0.15 - pose_weights: [0.01, 2.0, 5.0, 0.01, 5.0, 5.0, 0.01, 2.0, 5.0, 0.01, 5.0, 5.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0] + tracking_lin_vel: + weight: 2.2 + tracking_ang_vel: + weight: 1.8 + penalty_ang_vel_xy: + weight: -1.2 + penalty_orientation: + weight: -12.0 + penalty_action_rate: + weight: -2.5 + pose: + weight: -0.6 + params: + pose_weights: [0.01, 2.0, 5.0, 0.01, 5.0, 5.0, 0.01, 2.0, 5.0, 0.01, 5.0, 5.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0] + penalty_feet_ori: + weight: -5.0 + feet_phase: + weight: 6.0 + params: + tracking_sigma: 0.008 + alive: + weight: 12.0 diff --git a/conf/offpolicy/task/flashsac/g1_walk_flat/mujoco.yaml b/conf/offpolicy/task/flashsac/g1_walk_flat/mujoco.yaml index 35624c3b6..23d43c5e9 100644 --- a/conf/offpolicy/task/flashsac/g1_walk_flat/mujoco.yaml +++ b/conf/offpolicy/task/flashsac/g1_walk_flat/mujoco.yaml @@ -1,4 +1,11 @@ # @package _global_ +# FlashSAC MuJoCo owner: inherits the 29-DoF offpolicy Manager-Based contract +# and carries the FlashSAC algo identity plus its reward retuning +# (stiffer action-rate / feet-orientation penalties, tighter feet-phase sigma). +defaults: + - /task/g1_walk_flat/base + - _self_ + training: task_name: G1WalkFlat sim_backend: mujoco @@ -8,47 +15,16 @@ algo: max_iterations: 5000 save_interval: 1000 updates_per_step: 8 - #use_symmetry: true replay_buffer_n: 256 tau: 0.05 -env: - control_config: - action_scale: 1.0 - gait_phase_init_mode: "offset_phase" - reset_base_qvel_limit: 0.5 - curriculum: - enabled: true - initial_scale: 0.5 - min_scale: 0.5 - max_scale: 1.0 - level_down_threshold: 150.0 - level_up_threshold: 750.0 - degree: 0.001 - noise_config: - level: 0.0 - scale_gyro: 0.0 - scale_gravity: 0.0 - scale_joint_angle: 0.01 - scale_joint_vel: 0.1 - scale_linvel: 0.0 - seed: null reward: - scales: - tracking_lin_vel: 2.0 - tracking_ang_vel: 1.5 - penalty_ang_vel_xy: -1.0 - penalty_orientation: -10.0 - penalty_action_rate: -5.0 - pose: -0.5 - penalty_feet_ori: -25.0 - feet_phase: 5.0 - alive: 10.0 - tracking_sigma: 0.25 - base_height_target: 0.754 - min_base_height: 0.3 - max_tilt_deg: 65.0 - gait_frequency: 1.5 - feet_phase_swing_height: 0.09 - feet_phase_tracking_sigma: 0.005 - close_feet_threshold: 0.15 - pose_weights: [0.01, 2.0, 5.0, 0.01, 5.0, 5.0, 0.01, 2.0, 5.0, 0.01, 5.0, 5.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0] + penalty_action_rate: + weight: -5.0 + penalty_feet_ori: + weight: -25.0 + feet_phase: + params: + tracking_sigma: 0.005 + pose: + params: + pose_weights: [0.01, 2.0, 5.0, 0.01, 5.0, 5.0, 0.01, 2.0, 5.0, 0.01, 5.0, 5.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0] diff --git a/conf/offpolicy/task/g1_23dof_walk_flat/base.yaml b/conf/offpolicy/task/g1_23dof_walk_flat/base.yaml new file mode 100644 index 000000000..3d36ad65b --- /dev/null +++ b/conf/offpolicy/task/g1_23dof_walk_flat/base.yaml @@ -0,0 +1,66 @@ +# @package _global_ +# Canonical G1 23-DoF walk Manager-Based task declaration (offpolicy owners). +# Inherits the 29-DoF offpolicy contract and swaps the scene to the 23-DoF +# model (no waist roll/pitch, no wrist pitch/yaw) with 23-entry pose weights. +defaults: + - /task/g1_walk_flat/base + - _self_ + +env: + scene: + model_file: src/unilab/assets/robots/g1/scene_flat_23dof.xml + entities: + robot: + joint_names: + - left_hip_pitch_joint + - left_hip_roll_joint + - left_hip_yaw_joint + - left_knee_joint + - left_ankle_pitch_joint + - left_ankle_roll_joint + - right_hip_pitch_joint + - right_hip_roll_joint + - right_hip_yaw_joint + - right_knee_joint + - right_ankle_pitch_joint + - right_ankle_roll_joint + - waist_yaw_joint + - left_shoulder_pitch_joint + - left_shoulder_roll_joint + - left_shoulder_yaw_joint + - left_elbow_joint + - left_wrist_roll_joint + - right_shoulder_pitch_joint + - right_shoulder_roll_joint + - right_shoulder_yaw_joint + - right_elbow_joint + - right_wrist_roll_joint + actuator_names: + - left_hip_pitch_joint + - left_hip_roll_joint + - left_hip_yaw_joint + - left_knee_joint + - left_ankle_pitch_joint + - left_ankle_roll_joint + - right_hip_pitch_joint + - right_hip_roll_joint + - right_hip_yaw_joint + - right_knee_joint + - right_ankle_pitch_joint + - right_ankle_roll_joint + - waist_yaw_joint + - left_shoulder_pitch_joint + - left_shoulder_roll_joint + - left_shoulder_yaw_joint + - left_elbow_joint + - left_wrist_roll_joint + - right_shoulder_pitch_joint + - right_shoulder_roll_joint + - right_shoulder_yaw_joint + - right_elbow_joint + - right_wrist_roll_joint + +reward: + pose: + params: + pose_weights: [0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0] diff --git a/conf/offpolicy/task/g1_walk_flat/base.yaml b/conf/offpolicy/task/g1_walk_flat/base.yaml new file mode 100644 index 000000000..82fde7853 --- /dev/null +++ b/conf/offpolicy/task/g1_walk_flat/base.yaml @@ -0,0 +1,285 @@ +# @package _global_ +# Canonical G1 29-DoF walk Manager-Based task declaration (offpolicy owners). +# Backend owner leaves inherit this file and only override backend/algo tuning +# or explicitly disabled terms. Observation scaling follows the walk profile +# (gyro x0.25, joint velocity x0.05, critic linear velocity x2.0); every +# offpolicy owner carries the penalty curriculum. +env: + scene: + model_file: src/unilab/assets/robots/g1/scene_flat.xml + default_keyframe_name: stand + entities: + robot: + root_body_name: pelvis + joint_names: + - left_hip_pitch_joint + - left_hip_roll_joint + - left_hip_yaw_joint + - left_knee_joint + - left_ankle_pitch_joint + - left_ankle_roll_joint + - right_hip_pitch_joint + - right_hip_roll_joint + - right_hip_yaw_joint + - right_knee_joint + - right_ankle_pitch_joint + - right_ankle_roll_joint + - waist_yaw_joint + - waist_roll_joint + - waist_pitch_joint + - left_shoulder_pitch_joint + - left_shoulder_roll_joint + - left_shoulder_yaw_joint + - left_elbow_joint + - left_wrist_roll_joint + - left_wrist_pitch_joint + - left_wrist_yaw_joint + - right_shoulder_pitch_joint + - right_shoulder_roll_joint + - right_shoulder_yaw_joint + - right_elbow_joint + - right_wrist_roll_joint + - right_wrist_pitch_joint + - right_wrist_yaw_joint + actuator_names: + - left_hip_pitch_joint + - left_hip_roll_joint + - left_hip_yaw_joint + - left_knee_joint + - left_ankle_pitch_joint + - left_ankle_roll_joint + - right_hip_pitch_joint + - right_hip_roll_joint + - right_hip_yaw_joint + - right_knee_joint + - right_ankle_pitch_joint + - right_ankle_roll_joint + - waist_yaw_joint + - waist_roll_joint + - waist_pitch_joint + - left_shoulder_pitch_joint + - left_shoulder_roll_joint + - left_shoulder_yaw_joint + - left_elbow_joint + - left_wrist_roll_joint + - left_wrist_pitch_joint + - left_wrist_yaw_joint + - right_shoulder_pitch_joint + - right_shoulder_roll_joint + - right_shoulder_yaw_joint + - right_elbow_joint + - right_wrist_roll_joint + - right_wrist_pitch_joint + - right_wrist_yaw_joint + body_names: [pelvis] + sim_dt: 0.006666666666666667 + ctrl_dt: 0.02 + max_episode_seconds: 20.0 + observations: + policy: + _target_: unilab.managers.ObservationGroupCfg + terms: + base_ang_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.builtin_sensor + params: {sensor_name: torso_gyro} + scale: 0.25 + projected_gravity: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.projected_gravity_from_sensor + params: {sensor_name: torso_upvector} + joint_pos: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_pos_rel + joint_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_vel_rel + scale: 0.05 + actions: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.last_action + params: {action_name: joint_pos} + command: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.generated_commands + params: {command_name: twist} + gait_phase: + _target_: unilab.managers.ObservationTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.G1GaitPhase + params: + frequency: 1.5 + init_mode: offset_phase + critic: + _target_: unilab.managers.ObservationGroupCfg + terms: + base_ang_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.builtin_sensor + params: {sensor_name: torso_gyro} + scale: 0.25 + projected_gravity: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.projected_gravity_from_sensor + params: {sensor_name: torso_upvector} + joint_pos: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_pos_rel + joint_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_vel_rel + scale: 0.05 + actions: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.last_action + params: {action_name: joint_pos} + command: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.generated_commands + params: {command_name: twist} + gait_phase: + _target_: unilab.managers.ObservationTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.G1GaitPhase + params: + frequency: 1.5 + init_mode: offset_phase + base_lin_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.builtin_sensor + params: {sensor_name: pelvis_local_linvel} + scale: 2.0 + actions: + joint_pos: + _target_: unilab.envs.mdp.JointPositionActionCfg + entity_name: robot + actuator_names: [".*"] + scale: 1.0 + use_default_offset: true + commands: + twist: + _target_: unilab.tasks.locomotion.g1.manager_terms.G1VelocityCommandCfg + entity_name: robot + resampling_time_range: [20.0, 20.0] + heading_command: false + heading_control_stiffness: 0.5 + rel_standing_envs: 0.0 + rel_heading_envs: 0.0 + rel_world_envs: 0.0 + rel_forward_envs: 0.0 + init_velocity_prob: 0.0 + planar_dead_zone: 0.2 + ranges: + lin_vel_x: [-0.6, 1.0] + lin_vel_y: [-0.4, 0.4] + ang_vel_z: [-0.8, 0.8] + events: + reset_scene_to_default: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.reset_scene_to_default + mode: reset + reset_root_state_uniform: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.reset_root_state_uniform + mode: reset + params: + pose_range: + x: [-0.5, 0.5] + y: [-0.5, 0.5] + z: [0.0, 0.0] + roll: [0.0, 0.0] + pitch: [0.0, 0.0] + yaw: [-3.141592653589793, 3.141592653589793] + velocity_range: + x: [-0.5, 0.5] + y: [-0.5, 0.5] + z: [-0.5, 0.5] + roll: [-0.5, 0.5] + pitch: [-0.5, 0.5] + yaw: [-0.5, 0.5] + pd_gains: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.pd_gains + mode: reset + params: + kp_range: [0.9, 1.1] + kd_range: [0.9, 1.1] + operation: scale + terminations: + time_out: + _target_: unilab.managers.TerminationTermCfg + func: unilab.envs.mdp.time_out + time_out: true + tilt: + _target_: unilab.managers.TerminationTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.g1_tilt_exceeded + params: + max_tilt_deg: 65.0 + base_height: + _target_: unilab.managers.TerminationTermCfg + func: unilab.envs.mdp.root_height_below_minimum + params: + minimum_height: 0.3 + curriculum: + penalty_scaling: + _target_: unilab.managers.CurriculumTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.G1PenaltyCurriculum + params: + initial_scale: 0.5 + min_scale: 0.5 + max_scale: 1.0 + level_down_threshold: 150.0 + level_up_threshold: 750.0 + degree: 0.001 + policy_observation_group: policy + critic_observation_group: critic + +reward: + tracking_lin_vel: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.track_lin_vel + weight: 2.0 + params: + tracking_sigma: 0.25 + command_name: twist + tracking_ang_vel: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.track_ang_vel + weight: 1.5 + params: + tracking_sigma: 0.25 + command_name: twist + penalty_ang_vel_xy: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.ang_vel_xy + weight: -1.0 + penalty_orientation: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.orientation + weight: -10.0 + penalty_action_rate: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.mdp.action_rate_l2 + weight: -4.0 + pose: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.weighted_pose + weight: -0.5 + params: + pose_weights: [0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0] + penalty_feet_ori: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.penalty_feet_ori + weight: -20.0 + feet_phase: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.feet_phase + weight: 5.0 + params: + frequency: 1.5 + swing_height: 0.09 + tracking_sigma: 0.04 + min_forward_speed: 0.0 + command_name: twist + alive: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.alive + weight: 10.0 diff --git a/conf/offpolicy/task/sac/g1_23dof_walk_flat/motrix.yaml b/conf/offpolicy/task/sac/g1_23dof_walk_flat/motrix.yaml index d67013177..bddbcf5ef 100644 --- a/conf/offpolicy/task/sac/g1_23dof_walk_flat/motrix.yaml +++ b/conf/offpolicy/task/sac/g1_23dof_walk_flat/motrix.yaml @@ -1,4 +1,10 @@ # @package _global_ +# SAC Motrix 23-DoF owner: keeps DENYLIST parity with the MuJoCo owner and +# retunes reward shaping / disables kp/kd randomization for Motrix. +defaults: + - /task/g1_23dof_walk_flat/base + - _self_ + training: task_name: G1Walk23DofFlat sim_backend: motrix @@ -13,38 +19,27 @@ algo: alpha_init: 0.001 target_entropy_ratio: 0.0 env: - control_config: - action_scale: 1.0 - gait_phase_init_mode: "offset_phase" - reset_base_qvel_limit: 0.5 - domain_rand: - randomize_kp: false - randomize_kd: false - curriculum: - enabled: true - initial_scale: 0.5 - min_scale: 0.5 - max_scale: 1.0 - level_down_threshold: 150.0 - level_up_threshold: 750.0 - degree: 0.001 + events: + # Legacy Motrix owners disable kp/kd randomization. + pd_gains: null reward: - scales: - tracking_lin_vel: 2.2 - tracking_ang_vel: 1.8 - penalty_ang_vel_xy: -1.2 - penalty_orientation: -12.0 - penalty_action_rate: -2.5 - pose: -0.6 - penalty_feet_ori: -5.0 - feet_phase: 6.0 - alive: 12.0 - tracking_sigma: 0.25 - base_height_target: 0.754 - min_base_height: 0.3 - max_tilt_deg: 65.0 - gait_frequency: 1.5 - feet_phase_swing_height: 0.09 - feet_phase_tracking_sigma: 0.008 - close_feet_threshold: 0.15 - pose_weights: [0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0] + tracking_lin_vel: + weight: 2.2 + tracking_ang_vel: + weight: 1.8 + penalty_ang_vel_xy: + weight: -1.2 + penalty_orientation: + weight: -12.0 + penalty_action_rate: + weight: -2.5 + pose: + weight: -0.6 + penalty_feet_ori: + weight: -5.0 + feet_phase: + weight: 6.0 + params: + tracking_sigma: 0.008 + alive: + weight: 12.0 diff --git a/conf/offpolicy/task/sac/g1_23dof_walk_flat/mujoco.yaml b/conf/offpolicy/task/sac/g1_23dof_walk_flat/mujoco.yaml index 1f2ca273c..324fdb985 100644 --- a/conf/offpolicy/task/sac/g1_23dof_walk_flat/mujoco.yaml +++ b/conf/offpolicy/task/sac/g1_23dof_walk_flat/mujoco.yaml @@ -1,4 +1,10 @@ # @package _global_ +# SAC MuJoCo 23-DoF owner: inherits the 23-DoF offpolicy Manager-Based +# contract and only carries backend/algo identity. +defaults: + - /task/g1_23dof_walk_flat/base + - _self_ + training: task_name: G1Walk23DofFlat sim_backend: mujoco @@ -12,43 +18,3 @@ algo: algo_params: alpha_init: 0.001 target_entropy_ratio: 0.0 -env: - control_config: - action_scale: 1.0 - gait_phase_init_mode: "offset_phase" - reset_base_qvel_limit: 0.5 - curriculum: - enabled: true - initial_scale: 0.5 - min_scale: 0.5 - max_scale: 1.0 - level_down_threshold: 150.0 - level_up_threshold: 750.0 - degree: 0.001 - noise_config: - level: 1.0 - scale_gyro: 0.0 - scale_gravity: 0.0 - scale_joint_angle: 0.01 - scale_joint_vel: 0.1 - scale_linvel: 0.0 -reward: - scales: - tracking_lin_vel: 2.0 - tracking_ang_vel: 1.5 - penalty_ang_vel_xy: -1.0 - penalty_orientation: -10.0 - penalty_action_rate: -4.0 - pose: -0.5 - penalty_feet_ori: -20.0 - feet_phase: 5.0 - alive: 10.0 - tracking_sigma: 0.25 - base_height_target: 0.754 - min_base_height: 0.3 - max_tilt_deg: 65.0 - gait_frequency: 1.5 - feet_phase_swing_height: 0.09 - feet_phase_tracking_sigma: 0.04 - close_feet_threshold: 0.15 - pose_weights: [0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0] diff --git a/conf/offpolicy/task/sac/g1_23dof_walk_rough/motrix.yaml b/conf/offpolicy/task/sac/g1_23dof_walk_rough/motrix.yaml index 844fb2f2a..f1c9af30d 100644 --- a/conf/offpolicy/task/sac/g1_23dof_walk_rough/motrix.yaml +++ b/conf/offpolicy/task/sac/g1_23dof_walk_rough/motrix.yaml @@ -1,4 +1,10 @@ # @package _global_ +# SAC Motrix 23-DoF rough owner: static-hfield rough scene, Motrix +# sim_dt=0.01, Motrix-direction reward retuning; kp/kd randomization disabled. +defaults: + - /task/g1_23dof_walk_flat/base + - _self_ + training: task_name: G1Walk23DofRough sim_backend: motrix @@ -14,38 +20,29 @@ algo: target_entropy_ratio: 0.0 env: sim_dt: 0.01 - control_config: - action_scale: 1.0 - gait_phase_init_mode: "offset_phase" - reset_base_qvel_limit: 0.5 - domain_rand: - randomize_kp: false - randomize_kd: false - curriculum: - enabled: true - initial_scale: 0.5 - min_scale: 0.5 - max_scale: 1.0 - level_down_threshold: 150.0 - level_up_threshold: 750.0 - degree: 0.001 + scene: + model_file: src/unilab/assets/robots/g1/scene_rough_23dof.xml + events: + # Legacy Motrix owners disable kp/kd randomization. + pd_gains: null reward: - scales: - tracking_lin_vel: 2.2 - tracking_ang_vel: 1.8 - penalty_ang_vel_xy: -1.2 - penalty_orientation: -12.0 - penalty_action_rate: -2.5 - pose: -0.6 - penalty_feet_ori: -5.0 - feet_phase: 6.0 - alive: 12.0 - tracking_sigma: 0.25 - base_height_target: 0.754 - min_base_height: 0.3 - max_tilt_deg: 65.0 - gait_frequency: 1.5 - feet_phase_swing_height: 0.09 - feet_phase_tracking_sigma: 0.008 - close_feet_threshold: 0.15 - pose_weights: [0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0] + tracking_lin_vel: + weight: 2.2 + tracking_ang_vel: + weight: 1.8 + penalty_ang_vel_xy: + weight: -1.2 + penalty_orientation: + weight: -12.0 + penalty_action_rate: + weight: -2.5 + pose: + weight: -0.6 + penalty_feet_ori: + weight: -5.0 + feet_phase: + weight: 6.0 + params: + tracking_sigma: 0.008 + alive: + weight: 12.0 diff --git a/conf/offpolicy/task/sac/g1_23dof_walk_rough/mujoco.yaml b/conf/offpolicy/task/sac/g1_23dof_walk_rough/mujoco.yaml index 16689327b..4c8ab2dad 100644 --- a/conf/offpolicy/task/sac/g1_23dof_walk_rough/mujoco.yaml +++ b/conf/offpolicy/task/sac/g1_23dof_walk_rough/mujoco.yaml @@ -1,4 +1,10 @@ # @package _global_ +# SAC MuJoCo 23-DoF rough owner: 23-DoF offpolicy contract plus the +# static-hfield rough scene. +defaults: + - /task/g1_23dof_walk_flat/base + - _self_ + training: task_name: G1Walk23DofRough sim_backend: mujoco @@ -13,42 +19,5 @@ algo: alpha_init: 0.001 target_entropy_ratio: 0.0 env: - control_config: - action_scale: 1.0 - gait_phase_init_mode: "offset_phase" - reset_base_qvel_limit: 0.5 - curriculum: - enabled: true - initial_scale: 0.5 - min_scale: 0.5 - max_scale: 1.0 - level_down_threshold: 150.0 - level_up_threshold: 750.0 - degree: 0.001 - noise_config: - level: 1.0 - scale_gyro: 0.0 - scale_gravity: 0.0 - scale_joint_angle: 0.01 - scale_joint_vel: 0.1 - scale_linvel: 0.0 -reward: - scales: - tracking_lin_vel: 2.0 - tracking_ang_vel: 1.5 - penalty_ang_vel_xy: -1.0 - penalty_orientation: -10.0 - penalty_action_rate: -4.0 - pose: -0.5 - penalty_feet_ori: -20.0 - feet_phase: 5.0 - alive: 10.0 - tracking_sigma: 0.25 - base_height_target: 0.754 - min_base_height: 0.3 - max_tilt_deg: 65.0 - gait_frequency: 1.5 - feet_phase_swing_height: 0.09 - feet_phase_tracking_sigma: 0.04 - close_feet_threshold: 0.15 - pose_weights: [0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0] + scene: + model_file: src/unilab/assets/robots/g1/scene_rough_23dof.xml diff --git a/conf/offpolicy/task/sac/g1_walk_flat/mjwarp.yaml b/conf/offpolicy/task/sac/g1_walk_flat/mjwarp.yaml index dc9e23014..c079cddaf 100644 --- a/conf/offpolicy/task/sac/g1_walk_flat/mjwarp.yaml +++ b/conf/offpolicy/task/sac/g1_walk_flat/mjwarp.yaml @@ -1,6 +1,12 @@ # @package _global_ -# Configured-only SAC owner for the mjwarp host adapter. Offline record reuses -# MuJoCo rendering; native playback and device-resident runtime are absent. +# Configured-only SAC mjwarp owner for the unified host contract adapter. Keeps +# DENYLIST parity with the MuJoCo owner plus the mjwarp capacity knobs; legacy +# kp/kd randomization is disabled. Offline record reuses MuJoCo rendering; +# native playback and device-resident runtime are intentionally absent. +defaults: + - /task/g1_walk_flat/base + - _self_ + training: task_name: G1WalkFlat sim_backend: mjwarp @@ -19,48 +25,6 @@ env: mjwarp_nconmax: 128 mjwarp_njmax: 256 render_spacing: 2.0 - control_config: - action_scale: 1.0 - gait_phase_init_mode: "offset_phase" - reset_base_qvel_limit: 0.5 - domain_rand: - randomize_kp: false - randomize_kd: false - randomize_dof_armature: false - randomize_body_gravity_compensation: false - curriculum: - enabled: true - initial_scale: 0.5 - min_scale: 0.5 - max_scale: 1.0 - level_down_threshold: 150.0 - level_up_threshold: 750.0 - degree: 0.001 - noise_config: - level: 0.0 - scale_gyro: 0.0 - scale_gravity: 0.0 - scale_joint_angle: 0.01 - scale_joint_vel: 0.1 - scale_linvel: 0.0 - seed: null -reward: - scales: - tracking_lin_vel: 2.0 - tracking_ang_vel: 1.5 - penalty_ang_vel_xy: -1.0 - penalty_orientation: -10.0 - penalty_action_rate: -4.0 - pose: -0.5 - penalty_feet_ori: -20.0 - feet_phase: 5.0 - alive: 10.0 - tracking_sigma: 0.25 - base_height_target: 0.754 - min_base_height: 0.3 - max_tilt_deg: 65.0 - gait_frequency: 1.5 - feet_phase_swing_height: 0.09 - feet_phase_tracking_sigma: 0.04 - close_feet_threshold: 0.15 - pose_weights: [0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0] + events: + # Legacy mjwarp owners disable kp/kd and armature randomization. + pd_gains: null diff --git a/conf/offpolicy/task/sac/g1_walk_flat/motrix.yaml b/conf/offpolicy/task/sac/g1_walk_flat/motrix.yaml index a193eeb51..120a865c1 100644 --- a/conf/offpolicy/task/sac/g1_walk_flat/motrix.yaml +++ b/conf/offpolicy/task/sac/g1_walk_flat/motrix.yaml @@ -1,4 +1,10 @@ # @package _global_ +# SAC Motrix owner: keeps DENYLIST parity with the MuJoCo owner and retunes +# reward shaping / disables kp/kd randomization for the Motrix direction. +defaults: + - /task/g1_walk_flat/base + - _self_ + training: task_name: G1WalkFlat sim_backend: motrix @@ -13,46 +19,27 @@ algo: alpha_init: 0.001 target_entropy_ratio: 0.0 env: - control_config: - action_scale: 1.0 - gait_phase_init_mode: "offset_phase" - reset_base_qvel_limit: 0.5 - domain_rand: - randomize_kp: false - randomize_kd: false - curriculum: - enabled: true - initial_scale: 0.5 - min_scale: 0.5 - max_scale: 1.0 - level_down_threshold: 150.0 - level_up_threshold: 750.0 - degree: 0.001 - noise_config: - level: 0.0 - scale_gyro: 0.0 - scale_gravity: 0.0 - scale_joint_angle: 0.01 - scale_joint_vel: 0.1 - scale_linvel: 0.0 - seed: null + events: + # Legacy Motrix owners disable kp/kd randomization. + pd_gains: null reward: - scales: - tracking_lin_vel: 2.2 - tracking_ang_vel: 1.8 - penalty_ang_vel_xy: -1.2 - penalty_orientation: -12.0 - penalty_action_rate: -2.5 - pose: -0.6 - penalty_feet_ori: -5.0 - feet_phase: 6.0 - alive: 12.0 - tracking_sigma: 0.25 - base_height_target: 0.754 - min_base_height: 0.3 - max_tilt_deg: 65.0 - gait_frequency: 1.5 - feet_phase_swing_height: 0.09 - feet_phase_tracking_sigma: 0.008 - close_feet_threshold: 0.15 - pose_weights: [0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0] + tracking_lin_vel: + weight: 2.2 + tracking_ang_vel: + weight: 1.8 + penalty_ang_vel_xy: + weight: -1.2 + penalty_orientation: + weight: -12.0 + penalty_action_rate: + weight: -2.5 + pose: + weight: -0.6 + penalty_feet_ori: + weight: -5.0 + feet_phase: + weight: 6.0 + params: + tracking_sigma: 0.008 + alive: + weight: 12.0 diff --git a/conf/offpolicy/task/sac/g1_walk_flat/mujoco.yaml b/conf/offpolicy/task/sac/g1_walk_flat/mujoco.yaml index 0b0e2d4bd..2e3ffe39a 100644 --- a/conf/offpolicy/task/sac/g1_walk_flat/mujoco.yaml +++ b/conf/offpolicy/task/sac/g1_walk_flat/mujoco.yaml @@ -1,4 +1,10 @@ # @package _global_ +# SAC MuJoCo owner: inherits the 29-DoF offpolicy Manager-Based contract from +# the shared base and only carries backend/algo identity. +defaults: + - /task/g1_walk_flat/base + - _self_ + training: task_name: G1WalkFlat sim_backend: mujoco @@ -12,44 +18,3 @@ algo: algo_params: alpha_init: 0.001 target_entropy_ratio: 0.0 -env: - control_config: - action_scale: 1.0 - gait_phase_init_mode: "offset_phase" - reset_base_qvel_limit: 0.5 - curriculum: - enabled: true - initial_scale: 0.5 - min_scale: 0.5 - max_scale: 1.0 - level_down_threshold: 150.0 - level_up_threshold: 750.0 - degree: 0.001 - noise_config: - level: 1.0 - scale_gyro: 0.0 - scale_gravity: 0.0 - scale_joint_angle: 0.01 - scale_joint_vel: 0.1 - scale_linvel: 0.0 - seed: null -reward: - scales: - tracking_lin_vel: 2.0 - tracking_ang_vel: 1.5 - penalty_ang_vel_xy: -1.0 - penalty_orientation: -10.0 - penalty_action_rate: -4.0 - pose: -0.5 - penalty_feet_ori: -20.0 - feet_phase: 5.0 - alive: 10.0 - tracking_sigma: 0.25 - base_height_target: 0.754 - min_base_height: 0.3 - max_tilt_deg: 65.0 - gait_frequency: 1.5 - feet_phase_swing_height: 0.09 - feet_phase_tracking_sigma: 0.04 - close_feet_threshold: 0.15 - pose_weights: [0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0] diff --git a/conf/offpolicy/task/sac/g1_walk_rough/motrix.yaml b/conf/offpolicy/task/sac/g1_walk_rough/motrix.yaml index 4144b1713..b34f0c7cd 100644 --- a/conf/offpolicy/task/sac/g1_walk_rough/motrix.yaml +++ b/conf/offpolicy/task/sac/g1_walk_rough/motrix.yaml @@ -1,4 +1,10 @@ # @package _global_ +# SAC Motrix rough owner: static-hfield rough scene, Motrix sim_dt=0.01, and +# the Motrix-direction reward retuning; kp/kd randomization stays disabled. +defaults: + - /task/g1_walk_flat/base + - _self_ + training: task_name: G1WalkRough sim_backend: motrix @@ -14,38 +20,29 @@ algo: target_entropy_ratio: 0.0 env: sim_dt: 0.01 - control_config: - action_scale: 1.0 - gait_phase_init_mode: "offset_phase" - reset_base_qvel_limit: 0.5 - domain_rand: - randomize_kp: false - randomize_kd: false - curriculum: - enabled: true - initial_scale: 0.5 - min_scale: 0.5 - max_scale: 1.0 - level_down_threshold: 150.0 - level_up_threshold: 750.0 - degree: 0.001 + scene: + model_file: src/unilab/assets/robots/g1/scene_rough.xml + events: + # Legacy Motrix owners disable kp/kd randomization. + pd_gains: null reward: - scales: - tracking_lin_vel: 2.2 - tracking_ang_vel: 1.8 - penalty_ang_vel_xy: -1.2 - penalty_orientation: -12.0 - penalty_action_rate: -2.5 - pose: -0.6 - penalty_feet_ori: -5.0 - feet_phase: 6.0 - alive: 12.0 - tracking_sigma: 0.25 - base_height_target: 0.754 - min_base_height: 0.3 - max_tilt_deg: 65.0 - gait_frequency: 1.5 - feet_phase_swing_height: 0.09 - feet_phase_tracking_sigma: 0.008 - close_feet_threshold: 0.15 - pose_weights: [0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0] + tracking_lin_vel: + weight: 2.2 + tracking_ang_vel: + weight: 1.8 + penalty_ang_vel_xy: + weight: -1.2 + penalty_orientation: + weight: -12.0 + penalty_action_rate: + weight: -2.5 + pose: + weight: -0.6 + penalty_feet_ori: + weight: -5.0 + feet_phase: + weight: 6.0 + params: + tracking_sigma: 0.008 + alive: + weight: 12.0 diff --git a/conf/offpolicy/task/sac/g1_walk_rough/mujoco.yaml b/conf/offpolicy/task/sac/g1_walk_rough/mujoco.yaml index bd9a1282a..8d34219db 100644 --- a/conf/offpolicy/task/sac/g1_walk_rough/mujoco.yaml +++ b/conf/offpolicy/task/sac/g1_walk_rough/mujoco.yaml @@ -1,4 +1,11 @@ # @package _global_ +# SAC MuJoCo rough owner: inherits the 29-DoF offpolicy Manager-Based contract +# and swaps the scene to the static-hfield rough XML (no height-scan +# observation and no terrain curriculum, matching the legacy rough task). +defaults: + - /task/g1_walk_flat/base + - _self_ + training: task_name: G1WalkRough sim_backend: mujoco @@ -13,42 +20,5 @@ algo: alpha_init: 0.001 target_entropy_ratio: 0.0 env: - control_config: - action_scale: 1.0 - gait_phase_init_mode: "offset_phase" - reset_base_qvel_limit: 0.5 - curriculum: - enabled: true - initial_scale: 0.5 - min_scale: 0.5 - max_scale: 1.0 - level_down_threshold: 150.0 - level_up_threshold: 750.0 - degree: 0.001 - noise_config: - level: 1.0 - scale_gyro: 0.0 - scale_gravity: 0.0 - scale_joint_angle: 0.01 - scale_joint_vel: 0.1 - scale_linvel: 0.0 -reward: - scales: - tracking_lin_vel: 2.0 - tracking_ang_vel: 1.5 - penalty_ang_vel_xy: -1.0 - penalty_orientation: -10.0 - penalty_action_rate: -4.0 - pose: -0.5 - penalty_feet_ori: -20.0 - feet_phase: 5.0 - alive: 10.0 - tracking_sigma: 0.25 - base_height_target: 0.754 - min_base_height: 0.3 - max_tilt_deg: 65.0 - gait_frequency: 1.5 - feet_phase_swing_height: 0.09 - feet_phase_tracking_sigma: 0.04 - close_feet_threshold: 0.15 - pose_weights: [0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0] + scene: + model_file: src/unilab/assets/robots/g1/scene_rough.xml diff --git a/conf/offpolicy/task/td3/g1_23dof_walk_flat/mujoco.yaml b/conf/offpolicy/task/td3/g1_23dof_walk_flat/mujoco.yaml index 850f155e8..12fd2d9c4 100644 --- a/conf/offpolicy/task/td3/g1_23dof_walk_flat/mujoco.yaml +++ b/conf/offpolicy/task/td3/g1_23dof_walk_flat/mujoco.yaml @@ -1,46 +1,12 @@ # @package _global_ +# TD3 MuJoCo 23-DoF owner: inherits the 23-DoF offpolicy Manager-Based +# contract and only carries backend/algo identity. +defaults: + - /task/g1_23dof_walk_flat/base + - _self_ + training: task_name: G1Walk23DofFlat sim_backend: mujoco algo: max_iterations: 100000 -env: - control_config: - action_scale: 1.0 - gait_phase_init_mode: "offset_phase" - reset_base_qvel_limit: 0.5 - curriculum: - enabled: true - initial_scale: 0.5 - min_scale: 0.5 - max_scale: 1.0 - level_down_threshold: 150.0 - level_up_threshold: 750.0 - degree: 0.001 - noise_config: - level: 1.0 - scale_gyro: 0.0 - scale_gravity: 0.0 - scale_joint_angle: 0.01 - scale_joint_vel: 0.1 - scale_linvel: 0.0 -reward: - scales: - tracking_lin_vel: 2.0 - tracking_ang_vel: 1.5 - penalty_ang_vel_xy: -1.0 - penalty_orientation: -10.0 - penalty_action_rate: -4.0 - pose: -0.5 - penalty_feet_ori: -20.0 - feet_phase: 5.0 - alive: 10.0 - tracking_sigma: 0.25 - base_height_target: 0.754 - min_base_height: 0.3 - max_tilt_deg: 65.0 - gait_frequency: 1.5 - feet_phase_swing_height: 0.09 - feet_phase_tracking_sigma: 0.04 - close_feet_threshold: 0.15 - pose_weights: [0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0] diff --git a/conf/offpolicy/task/td3/g1_walk_flat/mujoco.yaml b/conf/offpolicy/task/td3/g1_walk_flat/mujoco.yaml index 54f87699c..53dfcff6d 100644 --- a/conf/offpolicy/task/td3/g1_walk_flat/mujoco.yaml +++ b/conf/offpolicy/task/td3/g1_walk_flat/mujoco.yaml @@ -1,47 +1,12 @@ # @package _global_ +# TD3 MuJoCo owner: inherits the 29-DoF offpolicy Manager-Based contract and +# only carries backend/algo identity. +defaults: + - /task/g1_walk_flat/base + - _self_ + training: task_name: G1WalkFlat sim_backend: mujoco algo: max_iterations: 100000 -env: - control_config: - action_scale: 1.0 - gait_phase_init_mode: "offset_phase" - reset_base_qvel_limit: 0.5 - curriculum: - enabled: true - initial_scale: 0.5 - min_scale: 0.5 - max_scale: 1.0 - level_down_threshold: 150.0 - level_up_threshold: 750.0 - degree: 0.001 - noise_config: - level: 1.0 - scale_gyro: 0.0 - scale_gravity: 0.0 - scale_joint_angle: 0.01 - scale_joint_vel: 0.1 - scale_linvel: 0.0 - seed: null -reward: - scales: - tracking_lin_vel: 2.0 - tracking_ang_vel: 1.5 - penalty_ang_vel_xy: -1.0 - penalty_orientation: -10.0 - penalty_action_rate: -4.0 - pose: -0.5 - penalty_feet_ori: -20.0 - feet_phase: 5.0 - alive: 10.0 - tracking_sigma: 0.25 - base_height_target: 0.754 - min_base_height: 0.3 - max_tilt_deg: 65.0 - gait_frequency: 1.5 - feet_phase_swing_height: 0.09 - feet_phase_tracking_sigma: 0.04 - close_feet_threshold: 0.15 - pose_weights: [0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0] diff --git a/conf/ppo/task/g1_23dof_walk_flat/base.yaml b/conf/ppo/task/g1_23dof_walk_flat/base.yaml new file mode 100644 index 000000000..7489083e8 --- /dev/null +++ b/conf/ppo/task/g1_23dof_walk_flat/base.yaml @@ -0,0 +1,66 @@ +# @package _global_ +# Canonical G1 23-DoF walk Manager-Based task declaration (PPO/APPO owners). +# Inherits the 29-DoF flat contract and swaps the scene to the 23-DoF model +# (no waist roll/pitch, no wrist pitch/yaw) with the 23-entry pose weights. +defaults: + - /task/g1_walk_flat/base + - _self_ + +env: + scene: + model_file: src/unilab/assets/robots/g1/scene_flat_23dof.xml + entities: + robot: + joint_names: + - left_hip_pitch_joint + - left_hip_roll_joint + - left_hip_yaw_joint + - left_knee_joint + - left_ankle_pitch_joint + - left_ankle_roll_joint + - right_hip_pitch_joint + - right_hip_roll_joint + - right_hip_yaw_joint + - right_knee_joint + - right_ankle_pitch_joint + - right_ankle_roll_joint + - waist_yaw_joint + - left_shoulder_pitch_joint + - left_shoulder_roll_joint + - left_shoulder_yaw_joint + - left_elbow_joint + - left_wrist_roll_joint + - right_shoulder_pitch_joint + - right_shoulder_roll_joint + - right_shoulder_yaw_joint + - right_elbow_joint + - right_wrist_roll_joint + actuator_names: + - left_hip_pitch_joint + - left_hip_roll_joint + - left_hip_yaw_joint + - left_knee_joint + - left_ankle_pitch_joint + - left_ankle_roll_joint + - right_hip_pitch_joint + - right_hip_roll_joint + - right_hip_yaw_joint + - right_knee_joint + - right_ankle_pitch_joint + - right_ankle_roll_joint + - waist_yaw_joint + - left_shoulder_pitch_joint + - left_shoulder_roll_joint + - left_shoulder_yaw_joint + - left_elbow_joint + - left_wrist_roll_joint + - right_shoulder_pitch_joint + - right_shoulder_roll_joint + - right_shoulder_yaw_joint + - right_elbow_joint + - right_wrist_roll_joint + +reward: + pose: + params: + pose_weights: [0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0] diff --git a/conf/ppo/task/g1_23dof_walk_flat/motrix.yaml b/conf/ppo/task/g1_23dof_walk_flat/motrix.yaml index ab1c8a94f..615cb7ef4 100644 --- a/conf/ppo/task/g1_23dof_walk_flat/motrix.yaml +++ b/conf/ppo/task/g1_23dof_walk_flat/motrix.yaml @@ -1,4 +1,11 @@ # @package _global_ +# Motrix owner: inherits the shared 23-DoF flat Manager-Based contract from +# base.yaml, then overrides contract fields for Motrix-specific tuning +# (intentionally non-transferable from MuJoCo; drop overrides to restore parity). +defaults: + - /task/g1_23dof_walk_flat/base + - _self_ + training: task_name: G1Walk23DofFlat sim_backend: motrix @@ -17,48 +24,102 @@ algo: learning_rate: 3.0e-4 entropy_coef: 5.0e-3 env: - domain_rand: - randomize_kp: false - randomize_kd: false - control_config: - action_scale: 0.5 + actions: + joint_pos: + scale: 0.5 commands: - vel_limit: - - [0.4, 0.0, 0.0] - - [0.7, 0.0, 0.0] - gait_phase_init_mode: offset_phase - reset_base_qvel_limit: 0.05 - curriculum: - enabled: false - noise_config: - level: 1.0 - scale_joint_angle: 0.01 - scale_joint_vel: 1.5 - scale_gyro: 0.2 + twist: + ranges: + lin_vel_x: [0.4, 0.7] + lin_vel_y: [0.0, 0.0] + ang_vel_z: [0.0, 0.0] + events: + # Legacy Motrix owners disable kp/kd randomization. + pd_gains: null + reset_root_state_uniform: + params: + velocity_range: + x: [-0.05, 0.05] + y: [-0.05, 0.05] + z: [-0.05, 0.05] + roll: [-0.05, 0.05] + pitch: [-0.05, 0.05] + yaw: [-0.05, 0.05] + terminations: + tilt: + params: + max_tilt_deg: 35.0 + base_height: + params: + minimum_height: 0.5 reward: - scales: - tracking_lin_vel: 2.0 - tracking_ang_vel: 0.25 - forward_progress: 0.0 - under_speed: -0.2 - upper_body_pose: -0.05 - penalty_feet_ori: 0.0 - feet_phase: 1.2 - feet_phase_contrast: 1.5 - feet_phase_contact: 1.0 - feet_double_stance: -1.0 - lin_vel_z: -1.0 - ang_vel_xy: -0.2 - base_height: -120.0 - orientation: -2.5 - action_rate: -0.005 - pose: -0.05 - pose_weights: [0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0] - tracking_sigma: 0.25 - gait_frequency: 1.5 - feet_phase_swing_height: 0.09 - feet_phase_tracking_sigma: 0.008 - base_height_target: 0.765 - min_forward_speed_for_gait_reward: 0.05 - min_base_height: 0.5 - max_tilt_deg: 35.0 + tracking_ang_vel: + weight: 0.25 + forward_progress: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.forward_progress + weight: 0.0 + params: + command_name: twist + under_speed: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.under_speed + weight: -0.2 + params: + command_name: twist + upper_body_pose: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.upper_body_pose + weight: -0.05 + params: + pose_weights: [0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0] + penalty_feet_ori: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.penalty_feet_ori + weight: 0.0 + feet_phase: + weight: 1.2 + params: + min_forward_speed: 0.05 + feet_phase_contrast: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.feet_phase_contrast + weight: 1.5 + params: + frequency: 1.5 + swing_height: 0.09 + tracking_sigma: 0.008 + min_forward_speed: 0.05 + command_name: twist + feet_phase_contact: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.feet_phase_contact + weight: 1.0 + params: + frequency: 1.5 + swing_height: 0.09 + tracking_sigma: 0.008 + min_forward_speed: 0.05 + command_name: twist + feet_double_stance: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.feet_double_stance + weight: -1.0 + params: + frequency: 1.5 + swing_height: 0.09 + tracking_sigma: 0.008 + min_forward_speed: 0.05 + command_name: twist + ang_vel_xy: + weight: -0.2 + base_height: + weight: -120.0 + params: + target_height: 0.765 + orientation: + weight: -2.5 + action_rate: + weight: -0.005 + pose: + weight: -0.05 diff --git a/conf/ppo/task/g1_23dof_walk_flat/mujoco.yaml b/conf/ppo/task/g1_23dof_walk_flat/mujoco.yaml index 0397b885d..e8e5776ed 100644 --- a/conf/ppo/task/g1_23dof_walk_flat/mujoco.yaml +++ b/conf/ppo/task/g1_23dof_walk_flat/mujoco.yaml @@ -1,4 +1,10 @@ # @package _global_ +# MuJoCo owner: inherits the shared 23-DoF flat Manager-Based contract from +# base.yaml and only carries backend/algo identity. +defaults: + - /task/g1_23dof_walk_flat/base + - _self_ + training: task_name: G1Walk23DofFlat sim_backend: mujoco @@ -8,32 +14,3 @@ algo: obs_groups: actor: - actor -env: - control_config: - action_scale: 0.25 - curriculum: - enabled: false - noise_config: - level: 1.0 - scale_joint_angle: 0.01 - scale_joint_vel: 1.5 - scale_gyro: 0.2 -reward: - scales: - tracking_lin_vel: 2.0 - tracking_ang_vel: 0.2 - feet_phase: 1.0 - lin_vel_z: -1.0 - ang_vel_xy: -0.25 - base_height: -500.0 - orientation: -5.0 - action_rate: -0.01 - pose: -0.1 - tracking_sigma: 0.25 - gait_frequency: 1.5 - feet_phase_swing_height: 0.09 - feet_phase_tracking_sigma: 0.008 - base_height_target: 0.754 - min_base_height: 0.55 - max_tilt_deg: 25.0 - pose_weights: [0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0] diff --git a/conf/ppo/task/g1_23dof_walk_rough/mujoco.yaml b/conf/ppo/task/g1_23dof_walk_rough/mujoco.yaml index ea82c6792..2a960832a 100644 --- a/conf/ppo/task/g1_23dof_walk_rough/mujoco.yaml +++ b/conf/ppo/task/g1_23dof_walk_rough/mujoco.yaml @@ -1,4 +1,11 @@ # @package _global_ +# MuJoCo 23-DoF rough owner: inherits the 23-DoF flat Manager-Based contract, +# swaps the scene to the static-hfield rough XML, and enables the penalty +# curriculum (the only PPO walk owner that carries one). +defaults: + - /task/g1_23dof_walk_flat/base + - _self_ + training: task_name: G1Walk23DofRough sim_backend: mujoco @@ -9,37 +16,16 @@ algo: actor: - actor env: - control_config: - action_scale: 0.25 + scene: + model_file: src/unilab/assets/robots/g1/scene_rough_23dof.xml curriculum: - enabled: true - initial_scale: 0.5 - min_scale: 0.5 - max_scale: 1.0 - level_down_threshold: 150.0 - level_up_threshold: 750.0 - degree: 0.001 - noise_config: - level: 1.0 - scale_joint_angle: 0.01 - scale_joint_vel: 1.5 - scale_gyro: 0.2 -reward: - scales: - tracking_lin_vel: 2.0 - tracking_ang_vel: 0.2 - feet_phase: 1.0 - lin_vel_z: -1.0 - ang_vel_xy: -0.25 - base_height: -500.0 - orientation: -5.0 - action_rate: -0.01 - pose: -0.1 - tracking_sigma: 0.25 - gait_frequency: 1.5 - feet_phase_swing_height: 0.09 - feet_phase_tracking_sigma: 0.008 - base_height_target: 0.754 - min_base_height: 0.55 - max_tilt_deg: 25.0 - pose_weights: [0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0] + penalty_scaling: + _target_: unilab.managers.CurriculumTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.G1PenaltyCurriculum + params: + initial_scale: 0.5 + min_scale: 0.5 + max_scale: 1.0 + level_down_threshold: 150.0 + level_up_threshold: 750.0 + degree: 0.001 diff --git a/conf/ppo/task/g1_walk_flat/base.yaml b/conf/ppo/task/g1_walk_flat/base.yaml new file mode 100644 index 000000000..ae8130aaf --- /dev/null +++ b/conf/ppo/task/g1_walk_flat/base.yaml @@ -0,0 +1,270 @@ +# @package _global_ +# Canonical G1 29-DoF walk Manager-Based task declaration (PPO/APPO owners). +# Backend owner leaves inherit this file and only override backend/algo tuning +# or explicitly disabled terms. Observation scaling follows the legacy profile +# (unit scales); the walk profile lives in the offpolicy base owners. +env: + scene: + model_file: src/unilab/assets/robots/g1/scene_flat.xml + default_keyframe_name: stand + entities: + robot: + root_body_name: pelvis + joint_names: + - left_hip_pitch_joint + - left_hip_roll_joint + - left_hip_yaw_joint + - left_knee_joint + - left_ankle_pitch_joint + - left_ankle_roll_joint + - right_hip_pitch_joint + - right_hip_roll_joint + - right_hip_yaw_joint + - right_knee_joint + - right_ankle_pitch_joint + - right_ankle_roll_joint + - waist_yaw_joint + - waist_roll_joint + - waist_pitch_joint + - left_shoulder_pitch_joint + - left_shoulder_roll_joint + - left_shoulder_yaw_joint + - left_elbow_joint + - left_wrist_roll_joint + - left_wrist_pitch_joint + - left_wrist_yaw_joint + - right_shoulder_pitch_joint + - right_shoulder_roll_joint + - right_shoulder_yaw_joint + - right_elbow_joint + - right_wrist_roll_joint + - right_wrist_pitch_joint + - right_wrist_yaw_joint + actuator_names: + - left_hip_pitch_joint + - left_hip_roll_joint + - left_hip_yaw_joint + - left_knee_joint + - left_ankle_pitch_joint + - left_ankle_roll_joint + - right_hip_pitch_joint + - right_hip_roll_joint + - right_hip_yaw_joint + - right_knee_joint + - right_ankle_pitch_joint + - right_ankle_roll_joint + - waist_yaw_joint + - waist_roll_joint + - waist_pitch_joint + - left_shoulder_pitch_joint + - left_shoulder_roll_joint + - left_shoulder_yaw_joint + - left_elbow_joint + - left_wrist_roll_joint + - left_wrist_pitch_joint + - left_wrist_yaw_joint + - right_shoulder_pitch_joint + - right_shoulder_roll_joint + - right_shoulder_yaw_joint + - right_elbow_joint + - right_wrist_roll_joint + - right_wrist_pitch_joint + - right_wrist_yaw_joint + body_names: [pelvis] + sim_dt: 0.006666666666666667 + ctrl_dt: 0.02 + max_episode_seconds: 20.0 + observations: + policy: + _target_: unilab.managers.ObservationGroupCfg + terms: + base_ang_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.builtin_sensor + params: {sensor_name: torso_gyro} + projected_gravity: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.projected_gravity_from_sensor + params: {sensor_name: torso_upvector} + joint_pos: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_pos_rel + joint_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_vel_rel + actions: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.last_action + params: {action_name: joint_pos} + command: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.generated_commands + params: {command_name: twist} + gait_phase: + _target_: unilab.managers.ObservationTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.G1GaitPhase + params: + frequency: 1.5 + init_mode: offset_phase + critic: + _target_: unilab.managers.ObservationGroupCfg + terms: + base_ang_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.builtin_sensor + params: {sensor_name: torso_gyro} + projected_gravity: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.projected_gravity_from_sensor + params: {sensor_name: torso_upvector} + joint_pos: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_pos_rel + joint_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.joint_vel_rel + actions: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.last_action + params: {action_name: joint_pos} + command: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.generated_commands + params: {command_name: twist} + gait_phase: + _target_: unilab.managers.ObservationTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.G1GaitPhase + params: + frequency: 1.5 + init_mode: offset_phase + base_lin_vel: + _target_: unilab.managers.ObservationTermCfg + func: unilab.envs.mdp.builtin_sensor + params: {sensor_name: pelvis_local_linvel} + actions: + joint_pos: + _target_: unilab.envs.mdp.JointPositionActionCfg + entity_name: robot + actuator_names: [".*"] + scale: 0.25 + use_default_offset: true + commands: + twist: + _target_: unilab.tasks.locomotion.g1.manager_terms.G1VelocityCommandCfg + entity_name: robot + resampling_time_range: [20.0, 20.0] + heading_command: false + heading_control_stiffness: 0.5 + rel_standing_envs: 0.0 + rel_heading_envs: 0.0 + rel_world_envs: 0.0 + rel_forward_envs: 0.0 + init_velocity_prob: 0.0 + planar_dead_zone: 0.2 + ranges: + lin_vel_x: [-0.6, 1.0] + lin_vel_y: [-0.4, 0.4] + ang_vel_z: [-0.8, 0.8] + events: + reset_scene_to_default: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.reset_scene_to_default + mode: reset + reset_root_state_uniform: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.reset_root_state_uniform + mode: reset + params: + pose_range: + x: [-0.5, 0.5] + y: [-0.5, 0.5] + z: [0.0, 0.0] + roll: [0.0, 0.0] + pitch: [0.0, 0.0] + yaw: [-3.141592653589793, 3.141592653589793] + velocity_range: + x: [-0.5, 0.5] + y: [-0.5, 0.5] + z: [-0.5, 0.5] + roll: [-0.5, 0.5] + pitch: [-0.5, 0.5] + yaw: [-0.5, 0.5] + pd_gains: + _target_: unilab.managers.EventTermCfg + func: unilab.envs.mdp.pd_gains + mode: reset + params: + kp_range: [0.9, 1.1] + kd_range: [0.9, 1.1] + operation: scale + terminations: + time_out: + _target_: unilab.managers.TerminationTermCfg + func: unilab.envs.mdp.time_out + time_out: true + tilt: + _target_: unilab.managers.TerminationTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.g1_tilt_exceeded + params: + max_tilt_deg: 25.0 + base_height: + _target_: unilab.managers.TerminationTermCfg + func: unilab.envs.mdp.root_height_below_minimum + params: + minimum_height: 0.55 + policy_observation_group: policy + critic_observation_group: critic + +reward: + tracking_lin_vel: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.track_lin_vel + weight: 2.0 + params: + tracking_sigma: 0.25 + command_name: twist + tracking_ang_vel: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.track_ang_vel + weight: 0.2 + params: + tracking_sigma: 0.25 + command_name: twist + feet_phase: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.feet_phase + weight: 1.0 + params: + frequency: 1.5 + swing_height: 0.09 + tracking_sigma: 0.008 + min_forward_speed: 0.0 + command_name: twist + lin_vel_z: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.lin_vel_z + weight: -1.0 + ang_vel_xy: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.ang_vel_xy + weight: -0.25 + base_height: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.base_height + weight: -500.0 + params: + target_height: 0.754 + orientation: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.orientation + weight: -5.0 + action_rate: + _target_: unilab.managers.RewardTermCfg + func: unilab.envs.mdp.action_rate_l2 + weight: -0.01 + pose: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.weighted_pose + weight: -0.1 + params: + pose_weights: [0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0] diff --git a/conf/ppo/task/g1_walk_flat/mjwarp.yaml b/conf/ppo/task/g1_walk_flat/mjwarp.yaml index 2d8c3a9d7..647fc6d00 100644 --- a/conf/ppo/task/g1_walk_flat/mjwarp.yaml +++ b/conf/ppo/task/g1_walk_flat/mjwarp.yaml @@ -1,7 +1,12 @@ # @package _global_ -# Configured-only mjwarp owner for the unified host contract adapter. Offline -# record reuses MuJoCo rendering; native playback and device-resident runtime -# routing are intentionally absent. +# Configured-only mjwarp owner for the unified host contract adapter. Keeps +# DENYLIST parity with the MuJoCo owner; legacy kp/kd randomization is disabled. +# Offline record reuses MuJoCo rendering; native playback and device-resident +# runtime routing are intentionally absent. +defaults: + - /task/g1_walk_flat/base + - _self_ + training: task_name: G1WalkFlat sim_backend: mjwarp @@ -19,39 +24,9 @@ algo: env: mjwarp_nconmax: 128 mjwarp_njmax: 256 - domain_rand: - randomize_kp: false - randomize_kd: false - randomize_dof_armature: false - randomize_body_gravity_compensation: false - control_config: - action_scale: 0.25 - curriculum: - enabled: false - noise_config: - level: 0.0 - scale_joint_angle: 0.01 - scale_joint_vel: 1.5 - scale_gyro: 0.2 -reward: - scales: - tracking_lin_vel: 2.0 - tracking_ang_vel: 0.2 - feet_phase: 1.0 - lin_vel_z: -1.0 - ang_vel_xy: -0.25 - base_height: -500.0 - orientation: -5.0 - action_rate: -0.01 - pose: -0.1 - tracking_sigma: 0.25 - gait_frequency: 1.5 - feet_phase_swing_height: 0.09 - feet_phase_tracking_sigma: 0.008 - base_height_target: 0.754 - min_base_height: 0.55 - max_tilt_deg: 25.0 - pose_weights: [0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0] + events: + # Legacy mjwarp owners disable kp/kd and armature randomization. + pd_gains: null play_profile: enabled: true env: diff --git a/conf/ppo/task/g1_walk_flat/motrix.yaml b/conf/ppo/task/g1_walk_flat/motrix.yaml index 88c971936..5690fa578 100644 --- a/conf/ppo/task/g1_walk_flat/motrix.yaml +++ b/conf/ppo/task/g1_walk_flat/motrix.yaml @@ -1,7 +1,11 @@ # @package _global_ -# Standalone Motrix owner config: carries the shared contract inline, then -# overrides contract fields for Motrix-specific tuning +# Motrix owner: inherits the shared 29-DoF flat Manager-Based contract from +# base.yaml, then overrides contract fields for Motrix-specific tuning # (intentionally non-transferable from MuJoCo; drop overrides to restore parity). +defaults: + - /task/g1_walk_flat/base + - _self_ + training: task_name: G1WalkFlat sim_backend: motrix @@ -22,50 +26,105 @@ algo: learning_rate: 3.0e-4 entropy_coef: 5.0e-3 env: - domain_rand: - randomize_kp: false - randomize_kd: false - control_config: - action_scale: 0.5 + actions: + joint_pos: + scale: 0.5 commands: - vel_limit: - - [0.4, 0.0, 0.0] - - [0.7, 0.0, 0.0] - gait_phase_init_mode: offset_phase - reset_base_qvel_limit: 0.05 - curriculum: - enabled: false - noise_config: - level: 1.0 - scale_joint_angle: 0.01 - scale_joint_vel: 1.5 - scale_gyro: 0.2 + twist: + ranges: + lin_vel_x: [0.4, 0.7] + lin_vel_y: [0.0, 0.0] + ang_vel_z: [0.0, 0.0] + events: + # Legacy Motrix owners disable kp/kd randomization. + pd_gains: null + reset_root_state_uniform: + params: + velocity_range: + x: [-0.05, 0.05] + y: [-0.05, 0.05] + z: [-0.05, 0.05] + roll: [-0.05, 0.05] + pitch: [-0.05, 0.05] + yaw: [-0.05, 0.05] + terminations: + tilt: + params: + max_tilt_deg: 35.0 + base_height: + params: + minimum_height: 0.5 reward: - scales: - tracking_lin_vel: 2.0 - tracking_ang_vel: 0.25 - forward_progress: 0.0 - under_speed: -0.2 - upper_body_pose: -0.05 - penalty_feet_ori: 0.0 - feet_phase: 1.2 - feet_phase_contrast: 1.5 - feet_phase_contact: 1.0 - feet_double_stance: -1.0 - lin_vel_z: -1.0 - ang_vel_xy: -0.2 - base_height: -120.0 - orientation: -2.5 - action_rate: -0.005 - pose: -0.05 - tracking_sigma: 0.25 - gait_frequency: 1.5 - feet_phase_swing_height: 0.09 - feet_phase_tracking_sigma: 0.008 - base_height_target: 0.765 - min_forward_speed_for_gait_reward: 0.05 - min_base_height: 0.5 - max_tilt_deg: 35.0 + tracking_ang_vel: + weight: 0.25 + forward_progress: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.forward_progress + weight: 0.0 + params: + command_name: twist + under_speed: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.under_speed + weight: -0.2 + params: + command_name: twist + upper_body_pose: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.upper_body_pose + weight: -0.05 + params: + pose_weights: [0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0] + penalty_feet_ori: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.penalty_feet_ori + weight: 0.0 + feet_phase: + weight: 1.2 + params: + min_forward_speed: 0.05 + feet_phase_contrast: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.feet_phase_contrast + weight: 1.5 + params: + frequency: 1.5 + swing_height: 0.09 + tracking_sigma: 0.008 + min_forward_speed: 0.05 + command_name: twist + feet_phase_contact: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.feet_phase_contact + weight: 1.0 + params: + frequency: 1.5 + swing_height: 0.09 + tracking_sigma: 0.008 + min_forward_speed: 0.05 + command_name: twist + feet_double_stance: + _target_: unilab.managers.RewardTermCfg + func: unilab.tasks.locomotion.g1.manager_terms.feet_double_stance + weight: -1.0 + params: + frequency: 1.5 + swing_height: 0.09 + tracking_sigma: 0.008 + min_forward_speed: 0.05 + command_name: twist + ang_vel_xy: + weight: -0.2 + base_height: + weight: -120.0 + params: + target_height: 0.765 + orientation: + weight: -2.5 + action_rate: + weight: -0.005 + pose: + weight: -0.05 play_profile: enabled: true env: diff --git a/conf/ppo/task/g1_walk_flat/mujoco.yaml b/conf/ppo/task/g1_walk_flat/mujoco.yaml index 843d7b85f..a900f8d8d 100644 --- a/conf/ppo/task/g1_walk_flat/mujoco.yaml +++ b/conf/ppo/task/g1_walk_flat/mujoco.yaml @@ -1,6 +1,10 @@ # @package _global_ -# Standalone MuJoCo owner config: carries the shared cross-backend contract inline -# (formerly base.yaml), plus backend-specific tuning. +# MuJoCo owner: inherits the shared 29-DoF flat Manager-Based contract from +# base.yaml and only carries backend/algo identity. +defaults: + - /task/g1_walk_flat/base + - _self_ + training: task_name: G1WalkFlat sim_backend: mujoco @@ -14,35 +18,6 @@ algo: policy: actor_hidden_dims: [512, 256, 128] critic_hidden_dims: [512, 256, 128] -env: - control_config: - action_scale: 0.25 - curriculum: - enabled: false - noise_config: - level: 1.0 - scale_joint_angle: 0.01 - scale_joint_vel: 1.5 - scale_gyro: 0.2 -reward: - scales: - tracking_lin_vel: 2.0 - tracking_ang_vel: 0.2 - feet_phase: 1.0 - lin_vel_z: -1.0 - ang_vel_xy: -0.25 - base_height: -500.0 - orientation: -5.0 - action_rate: -0.01 - pose: -0.1 - tracking_sigma: 0.25 - gait_frequency: 1.5 - feet_phase_swing_height: 0.09 - feet_phase_tracking_sigma: 0.008 - base_height_target: 0.754 - min_base_height: 0.55 - max_tilt_deg: 25.0 - pose_weights: [0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0] play_profile: enabled: true env: diff --git a/docs/sphinx/source/en/2-user_guide/5-domain_randomization/0-index.md b/docs/sphinx/source/en/2-user_guide/5-domain_randomization/0-index.md index 8d50e3800..c76b3a2fb 100644 --- a/docs/sphinx/source/en/2-user_guide/5-domain_randomization/0-index.md +++ b/docs/sphinx/source/en/2-user_guide/5-domain_randomization/0-index.md @@ -18,7 +18,7 @@ These three paths correspond to three lifecycle classes: ## Status Conclusions 1. All tasks currently wired to a DR provider use the unified DR entry point; no task bypasses `DomainRandomizationManager` to run a separate DR flow inside `reset()`. -2. They are all roughly structured: task files define a `domain_rand` config dataclass, a `DomainRandomizationProvider`, and a `ResetPlan`; `G1WalkFlat` reuses `G1Walk`'s provider. +2. They are all roughly structured: task files define a `domain_rand` config dataclass, a `DomainRandomizationProvider`, and a `ResetPlan`; `G1Walk*` instead declares DR through Hydra `EventTermCfg` Manager-Based reset terms. 3. What is "unified" today is mainly the entry point and execution flow, not every randomization item itself. The shared helper `build_common_reset_randomization()` currently generates `base_mass_delta`, `base_com_offset`, `gravity`, `kp`, `kd`; the shared interval helper currently only generates push. 4. `ResetRandomizationPayload` can already express `gravity`, `body_iquat`, `body_inertia`, `kp`, `kd`, and `MuJoCoBackend` has declared support. Whether these are actually used still depends on whether the task provider samples and dispatches them. 5. `MotrixBackend` currently supports `base_mass_delta`, `base_com_offset`, `kp`, `kd`, and interval push; and it requires all model actuators to be position actuators during initialization. @@ -30,8 +30,8 @@ These three paths correspond to three lifecycle classes: | --- | --- | --- | --- | --- | --- | | `Go1JoystickFlat` | Yes | Yes: `Domain_Rand + Provider + ResetPlan` | task state sampling + common payload | push | `go1/joystick.py` | | `Go2JoystickFlat` | Yes | Yes: `Domain_Rand + Provider + ResetPlan` | task state sampling + common payload | push | `go2/joystick.py` | -| `G1WalkFlat` | Yes | Yes: `Domain_Rand + Provider + ResetPlan` | task state sampling + common payload | push | `g1/joystick.py` | -| `G1WalkRough` | Yes | Yes: reuses `G1WalkDomainRandomizationProvider` | task state sampling + common payload | push | `g1/joystick.py` | +| `G1WalkFlat` | Yes | Yes: Hydra `EventTermCfg` + Manager-Based reset terms | root-state reset + kp/kd via `pd_gains` | none | `g1/manager_terms.py` | +| `G1WalkRough` | Yes | Yes: same Manager-Based event terms as `G1WalkFlat` | root-state reset + kp/kd via `pd_gains` | none | `g1/manager_terms.py` | | `G1MotionTracking` | Yes | Yes: `Domain_Rand + Provider + ResetPlan` | extensive task-specific reset sampling + common payload | push | `motion_tracking/g1/tracking.py` | | `AllegroInhandRotation` | Yes | Yes: Hydra `EventTermCfg` + Manager-Based reset term | entity-scoped hand/ball reset | none | `allegro_inhand/manager_terms.py` | | `AllegroInhandRotationGrasp` | Yes | Yes: reuses the rotation reset event + `RecorderTermCfg` | noisy hand reset + grasp collection | none | `allegro_inhand/grasp_gen.py` | @@ -44,8 +44,8 @@ These three paths correspond to three lifecycle classes: | --- | --- | --- | --- | | `Go1JoystickFlat` | base xy; base yaw; base qvel; command sampling; `current_actions/last_actions` zeroed; optional `base_mass_delta`; optional `base_com_offset`; optional `gravity` | `push_robots` | `base_mass_delta`, `base_com_offset`, and push enabled by default; `gravity` disabled by default | | `Go2JoystickFlat` | base xy; base yaw; base qvel; command sampling; `current_actions/last_actions` zeroed; kp/kd randomization (enabled by default); optional `base_mass_delta`; optional `base_com_offset`; optional `gravity` | `push_robots` | kp/kd enabled by default; common payload and push disabled by default | -| `G1WalkFlat` | base xy; base yaw; base qvel sampled by `reset_base_qvel_limit`; command sampling; `gait_phase` sampling; `current_actions/last_actions` zeroed; kp/kd randomization (enabled by default); optional `base_mass_delta`; optional `base_com_offset`; optional `gravity` | `push_robots` | kp/kd enabled by default; common payload and push disabled by default | -| `G1WalkRough` | Same as `G1WalkFlat`, directly reuses the same provider | `push_robots` | kp/kd enabled by default; common payload and push disabled by default | +| `G1WalkFlat` | base xy/yaw and base qvel via `reset_root_state_uniform`; command sampling with a planar dead zone; `gait_phase` sampling; kp/kd randomization via `pd_gains` | none | kp/kd enabled on mujoco owners by default; disabled on motrix/mjwarp owners | +| `G1WalkRough` | Same as `G1WalkFlat` (shared owner bases, rough scene) | none | Same defaults as `G1WalkFlat` | | `G1MotionTracking` | motion frame sampling; root pose perturbation `x/y/z/roll/pitch/yaw`; root velocity perturbation `x/y/z/roll/pitch/yaw`; joint position noise; under MuJoCo clipped by joint range; `current_actions/last_actions` zeroed; optional `base_mass_delta`; optional `base_com_offset`; optional `gravity` | `push_robots` | `pose_randomization`, `velocity_randomization`, `joint_position_range` have non-zero perturbations by default; common payload and push disabled by default | | `AllegroInhandRotation` | Entity-scoped hand/ball reset; an explicitly configured grasp cache is sampled, otherwise `null` explicitly selects the model home pose; optional `joint_noise`, `ball_velocity_noise`, and `ball_z_offset` | none | owner YAML explicitly selects the home pose and zero reset noise; a configured missing or malformed cache fails closed | | `AllegroInhandRotationGrasp` | Reuses the rotation reset with `joint_noise=0.25`; Manager-Based termination checks fingertip distance, contact count, and ball height; recorder stores successful timeout rows | none | generates the 50k-row Allegro grasp cache and raises `RunComplete` after a successful save | @@ -134,7 +134,7 @@ Field semantics: If you only want to randomize the magnitude while keeping the vertical-down direction, only open up the `z` component: ```bash -uv run train --algo ppo --task g1_walk_flat --sim mujoco \ +uv run train --algo ppo --task g1_motion_tracking --sim mujoco \ env.domain_rand.randomize_gravity=true \ 'env.domain_rand.gravity_range=[[0.0,0.0,-10.5],[0.0,0.0,-8.5]]' ``` @@ -142,7 +142,7 @@ uv run train --algo ppo --task g1_walk_flat --sim mujoco \ If you want to randomize both direction and magnitude, open up `x/y/z`: ```bash -uv run train --algo ppo --task g1_walk_flat --sim mujoco \ +uv run train --algo ppo --task g1_motion_tracking --sim mujoco \ env.domain_rand.randomize_gravity=true \ 'env.domain_rand.gravity_range=[[-0.3,-0.3,-10.5],[0.3,0.3,-8.5]]' ``` @@ -174,7 +174,7 @@ env: - `push_body_name`: the target body / link to apply the force to. Defaults to `null`, meaning the backend's `base_name` is used. ```bash -uv run train --algo ppo --task g1_walk_flat --sim mujoco \ +uv run train --algo ppo --task g1_motion_tracking --sim mujoco \ env.domain_rand.push_robots=true \ env.domain_rand.push_interval=500 \ 'env.domain_rand.max_force=[20.0,20.0,5.0]' \ diff --git a/docs/sphinx/source/en/2-user_guide/5-domain_randomization/1-configuration.md b/docs/sphinx/source/en/2-user_guide/5-domain_randomization/1-configuration.md index bc1ee6ae0..f2d91d87b 100644 --- a/docs/sphinx/source/en/2-user_guide/5-domain_randomization/1-configuration.md +++ b/docs/sphinx/source/en/2-user_guide/5-domain_randomization/1-configuration.md @@ -5,7 +5,7 @@ under `env.domain_rand`. Use `--task` and `--sim` to select backend-specific behavior first, then override fields inside that selected owner. ```bash -uv run train --algo ppo --task g1_walk_flat --sim mujoco \ +uv run train --algo ppo --task g1_motion_tracking --sim mujoco \ env.domain_rand.randomize_gravity=true \ 'env.domain_rand.gravity_range=[[0.0,0.0,-10.5],[0.0,0.0,-8.5]]' ``` @@ -31,7 +31,7 @@ Use `--sim mujoco` when enabling gravity reset randomization; Motrix does not advertise the same gravity capability in the current backend. ```bash -uv run train --algo ppo --task g1_walk_flat --sim mujoco \ +uv run train --algo ppo --task g1_motion_tracking --sim mujoco \ env.domain_rand.randomize_gravity=true \ 'env.domain_rand.gravity_range=[[0.0,0.0,-10.5],[0.0,0.0,-8.5]]' ``` @@ -39,7 +39,7 @@ uv run train --algo ppo --task g1_walk_flat --sim mujoco \ ## Interval Push ```bash -uv run train --algo ppo --task g1_walk_flat --sim mujoco \ +uv run train --algo ppo --task g1_motion_tracking --sim mujoco \ env.domain_rand.push_robots=true \ env.domain_rand.push_interval=500 \ 'env.domain_rand.max_force=[20.0,20.0,5.0]' diff --git a/docs/sphinx/source/zh_CN/2-user_guide/5-domain_randomization/0-index.md b/docs/sphinx/source/zh_CN/2-user_guide/5-domain_randomization/0-index.md index e8dd9c389..a75dbbd25 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/5-domain_randomization/0-index.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/5-domain_randomization/0-index.md @@ -18,7 +18,7 @@ ## 状态结论 1. 当前所有接入 DR provider 的任务都使用统一的 DR 入口点；没有任何任务绕开 `DomainRandomizationManager` 在 `reset()` 内部运行单独的 DR 流程。 -2. 它们的结构都大致相同：task 文件定义一个 `domain_rand` 配置 dataclass、一个 `DomainRandomizationProvider` 和一个 `ResetPlan`；`G1WalkFlat` 复用 `G1Walk` 的 provider。 +2. 它们的结构都大致相同：task 文件定义一个 `domain_rand` 配置 dataclass、一个 `DomainRandomizationProvider` 和一个 `ResetPlan`；`G1Walk*` 则改为通过 Hydra `EventTermCfg` Manager-Based reset term 声明 DR。 3. 今天所"统一"的主要是入口点和执行流程，而不是每一个随机化项本身。共享辅助函数 `build_common_reset_randomization()` 目前生成 `base_mass_delta`、`base_com_offset`、`gravity`、`kp`、`kd`；共享的 interval 辅助函数目前只生成 push。 4. `ResetRandomizationPayload` 已经可以表达 `gravity`、`body_iquat`、`body_inertia`、`kp`、`kd`，并且 `MuJoCoBackend` 已声明支持。这些是否实际被使用，仍取决于 task provider 是否对它们进行采样和 dispatch。 5. `MotrixBackend` 目前支持 `base_mass_delta`、`base_com_offset`、`kp`、`kd` 和 interval push；并且它要求在初始化期间所有模型 actuator 都是 position actuator。 @@ -30,8 +30,8 @@ | --- | --- | --- | --- | --- | --- | | `Go1JoystickFlat` | 是 | 是：`Domain_Rand + Provider + ResetPlan` | task 状态采样 + common payload | push | `go1/joystick.py` | | `Go2JoystickFlat` | 是 | 是：`Domain_Rand + Provider + ResetPlan` | task 状态采样 + common payload | push | `go2/joystick.py` | -| `G1WalkFlat` | 是 | 是：`Domain_Rand + Provider + ResetPlan` | task 状态采样 + common payload | push | `g1/joystick.py` | -| `G1WalkRough` | 是 | 是：复用 `G1WalkDomainRandomizationProvider` | task 状态采样 + common payload | push | `g1/joystick.py` | +| `G1WalkFlat` | 是 | 是：Hydra `EventTermCfg` + Manager-Based reset term | root-state reset + 经 `pd_gains` 的 kp/kd | 无 | `g1/manager_terms.py` | +| `G1WalkRough` | 是 | 是：与 `G1WalkFlat` 相同的 Manager-Based event term | root-state reset + 经 `pd_gains` 的 kp/kd | 无 | `g1/manager_terms.py` | | `G1MotionTracking` | 是 | 是：`Domain_Rand + Provider + ResetPlan` | 大量 task 专属的 reset 采样 + common payload | push | `motion_tracking/g1/tracking.py` | | `AllegroInhandRotation` | 是 | 是：Hydra `EventTermCfg` + Manager-Based reset term | entity 范围的手/球 reset | 无 | `allegro_inhand/manager_terms.py` | | `AllegroInhandRotationGrasp` | 是 | 是：复用 rotation reset event + `RecorderTermCfg` | 带噪声的手部 reset + grasp 收集 | 无 | `allegro_inhand/grasp_gen.py` | @@ -44,8 +44,8 @@ | --- | --- | --- | --- | | `Go1JoystickFlat` | base xy；base yaw；base qvel；command 采样；`current_actions/last_actions` 清零；可选 `base_mass_delta`；可选 `base_com_offset`；可选 `gravity` | `push_robots` | `base_mass_delta`、`base_com_offset` 和 push 默认启用；`gravity` 默认禁用 | | `Go2JoystickFlat` | base xy；base yaw；base qvel；command 采样；`current_actions/last_actions` 清零；kp/kd 随机化（默认启用）；可选 `base_mass_delta`；可选 `base_com_offset`；可选 `gravity` | `push_robots` | kp/kd 默认启用；common payload 和 push 默认禁用 | -| `G1WalkFlat` | base xy；base yaw；由 `reset_base_qvel_limit` 采样的 base qvel；command 采样；`gait_phase` 采样；`current_actions/last_actions` 清零；kp/kd 随机化（默认启用）；可选 `base_mass_delta`；可选 `base_com_offset`；可选 `gravity` | `push_robots` | kp/kd 默认启用；common payload 和 push 默认禁用 | -| `G1WalkRough` | 与 `G1WalkFlat` 相同，直接复用同一 provider | `push_robots` | kp/kd 默认启用；common payload 和 push 默认禁用 | +| `G1WalkFlat` | 经 `reset_root_state_uniform` 的 base xy/yaw 与 base qvel；带平面死区的 command 采样；`gait_phase` 采样；经 `pd_gains` 的 kp/kd 随机化 | 无 | mujoco owner 默认启用 kp/kd；motrix/mjwarp owner 默认禁用 | +| `G1WalkRough` | 与 `G1WalkFlat` 相同（共享 owner base，rough 场景） | 无 | 与 `G1WalkFlat` 相同的默认值 | | `G1MotionTracking` | 动作帧采样；root 位姿扰动 `x/y/z/roll/pitch/yaw`；root 速度扰动 `x/y/z/roll/pitch/yaw`；关节位置噪声；在 MuJoCo 下被关节范围 clip；`current_actions/last_actions` 清零；可选 `base_mass_delta`；可选 `base_com_offset`；可选 `gravity` | `push_robots` | `pose_randomization`、`velocity_randomization`、`joint_position_range` 默认有非零扰动；common payload 和 push 默认禁用 | | `AllegroInhandRotation` | entity 范围的手/球 reset；显式配置 grasp cache 时进行采样，否则以 `null` 显式选择模型 home pose；可选 `joint_noise`、`ball_velocity_noise` 与 `ball_z_offset` | 无 | owner YAML 显式选择 home pose 与零 reset 噪声；配置的 cache 缺失或格式错误时 fail-closed | | `AllegroInhandRotationGrasp` | 复用 rotation reset 并设置 `joint_noise=0.25`；Manager-Based termination 检查指尖距离、接触数和球高度；recorder 保存成功 timeout rows | 无 | 生成 5 万行 Allegro grasp cache，成功保存后抛出 `RunComplete` | @@ -134,7 +134,7 @@ env: 如果你只想随机化大小而保持竖直向下的方向，只开放 `z` 分量： ```bash -uv run train --algo ppo --task g1_walk_flat --sim mujoco \ +uv run train --algo ppo --task g1_motion_tracking --sim mujoco \ env.domain_rand.randomize_gravity=true \ 'env.domain_rand.gravity_range=[[0.0,0.0,-10.5],[0.0,0.0,-8.5]]' ``` @@ -142,7 +142,7 @@ uv run train --algo ppo --task g1_walk_flat --sim mujoco \ 如果你想同时随机化方向和大小，开放 `x/y/z`： ```bash -uv run train --algo ppo --task g1_walk_flat --sim mujoco \ +uv run train --algo ppo --task g1_motion_tracking --sim mujoco \ env.domain_rand.randomize_gravity=true \ 'env.domain_rand.gravity_range=[[-0.3,-0.3,-10.5],[0.3,0.3,-8.5]]' ``` @@ -174,7 +174,7 @@ env: - `push_body_name`：施加力的目标 body / link。默认为 `null`，表示使用后端的 `base_name`。 ```bash -uv run train --algo ppo --task g1_walk_flat --sim mujoco \ +uv run train --algo ppo --task g1_motion_tracking --sim mujoco \ env.domain_rand.push_robots=true \ env.domain_rand.push_interval=500 \ 'env.domain_rand.max_force=[20.0,20.0,5.0]' \ diff --git a/docs/sphinx/source/zh_CN/2-user_guide/5-domain_randomization/1-configuration.md b/docs/sphinx/source/zh_CN/2-user_guide/5-domain_randomization/1-configuration.md index d83415575..6bcda32d0 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/5-domain_randomization/1-configuration.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/5-domain_randomization/1-configuration.md @@ -5,7 +5,7 @@ 然后在所选的 owner 内部 override 字段。 ```bash -uv run train --algo ppo --task g1_walk_flat --sim mujoco \ +uv run train --algo ppo --task g1_motion_tracking --sim mujoco \ env.domain_rand.randomize_gravity=true \ 'env.domain_rand.gravity_range=[[0.0,0.0,-10.5],[0.0,0.0,-8.5]]' ``` @@ -28,7 +28,7 @@ uv run train --algo ppo --task g1_walk_flat --sim mujoco \ 未提供相同的 gravity 能力。 ```bash -uv run train --algo ppo --task g1_walk_flat --sim mujoco \ +uv run train --algo ppo --task g1_motion_tracking --sim mujoco \ env.domain_rand.randomize_gravity=true \ 'env.domain_rand.gravity_range=[[0.0,0.0,-10.5],[0.0,0.0,-8.5]]' ``` @@ -36,7 +36,7 @@ uv run train --algo ppo --task g1_walk_flat --sim mujoco \ ## Interval Push ```bash -uv run train --algo ppo --task g1_walk_flat --sim mujoco \ +uv run train --algo ppo --task g1_motion_tracking --sim mujoco \ env.domain_rand.push_robots=true \ env.domain_rand.push_interval=500 \ 'env.domain_rand.max_force=[20.0,20.0,5.0]' diff --git a/docs/sphinx/source/zh_CN/4-developer_guide/9-sim2sim_contract_status.md b/docs/sphinx/source/zh_CN/4-developer_guide/9-sim2sim_contract_status.md index f981f54c5..f5d31b2bb 100644 --- a/docs/sphinx/source/zh_CN/4-developer_guide/9-sim2sim_contract_status.md +++ b/docs/sphinx/source/zh_CN/4-developer_guide/9-sim2sim_contract_status.md @@ -31,7 +31,7 @@ uv run scripts/audit_sim2sim_contracts.py | allegro_inhand · allegro_inhand_grasp · g1_climb_tracking · g1_motion_tracking · g1_wall_flip_tracking · go1_joystick_rough · go2_arm_manip_loco · go2_footstand · go2_handstand · go2_joystick_flat · go2_joystick_rough · go2w_joystick_flat · go2w_joystick_rough · sharpa_inhand · sharpa_inhand_grasp | ✅ | 无 | | g1_box_tracking | ❌ | `empirical_normalization` false↔true；`obs_groups` critic 组差异 | | g1_flip_tracking | ❌ | `empirical_normalization` true↔false；`obs_groups`；`action_scale` 29 维↔默认 0.25；`sampling_mode` 两后端运行时同为 `start`（无害） | -| g1_walk_flat | ❌ | `action_scale` 0.25↔0.5；`empirical_normalization` false↔true；`obs_groups` | +| g1_walk_flat | ❌ | `env.actions.joint_pos.scale` 0.25↔0.5；`empirical_normalization` false↔true；`obs_groups` | | go1_joystick_flat | ❌ | `empirical_normalization` false↔true | | g1_motion_tracking_deploy | ⚪ | 仅 mujoco | @@ -69,6 +69,7 @@ uv run scripts/audit_sim2sim_contracts.py | `action_scale` | **不可** | 改值即改训练动力学，必须 owner 决策 + 重训 | | `empirical_normalization` | **不可** | 改变网络结构，必须重训 | -试点示例：`conf/ppo/task/g1_walk_flat/{mujoco,motrix}.yaml`。每个后端 owner 自包含完整契约， -`motrix.yaml` 为单后端调参 override 了若干契约字段——这种 override +试点示例：`conf/ppo/task/g1_walk_flat/{base,mujoco,motrix}.yaml`。后端 owner 通过 Hydra +defaults 继承共享 base owner 的完整契约，`motrix.yaml` 为单后端调参 override +了若干契约字段——这种 override 即令该 task 在该后端不可 sim2sim 迁移，去掉 override 即可恢复。 diff --git a/scripts/benchmark/core/task_names.py b/scripts/benchmark/core/task_names.py index e7b1c09d2..2bc007acc 100644 --- a/scripts/benchmark/core/task_names.py +++ b/scripts/benchmark/core/task_names.py @@ -3,7 +3,6 @@ from dataclasses import dataclass from unilab.envs import ManagerBasedRlEnvCfg -from unilab.tasks.locomotion.g1.joystick import G1WalkFlatCfg from unilab.tasks.manipulation.sharpa_inhand.rotation import SharpaInhandRotationCfg @@ -35,7 +34,8 @@ class LocomotionTaskSpec: owner_task_id="g1_walk_flat", env_task_name="G1WalkFlat", display_name="g1_walk_flat", - config_cls=G1WalkFlatCfg, + config_cls=ManagerBasedRlEnvCfg, + model_file="src/unilab/assets/robots/g1/scene_flat.xml", ), "sharpa_inhand": LocomotionTaskSpec( owner_task_id="sharpa_inhand", diff --git a/scripts/benchmark/env/benchmark_env_step.py b/scripts/benchmark/env/benchmark_env_step.py index 7fea985f8..f8422efab 100644 --- a/scripts/benchmark/env/benchmark_env_step.py +++ b/scripts/benchmark/env/benchmark_env_step.py @@ -78,7 +78,7 @@ def _load_helper_module(module_name: str, relative_path: str): def _install_mjwarp_patch() -> bool: """Route ``backend_type == "mjwarp"`` to ``scripts/benchmark/mjwarp`` via factory patch. - Must run before any task env module (e.g. ``unilab.tasks.locomotion.g1.joystick``) + Must run before any task env module (e.g. ``unilab.tasks.locomotion.g1``) is imported, because those modules bind ``create_backend`` at module load time via ``from unilab.base.backend import create_backend``. @@ -322,15 +322,15 @@ def _go2w_env_cls() -> Callable[..., Any]: def _g1_flat_cfg(backend: str, config_overrides: list[str]) -> Any: - from unilab.tasks.locomotion.g1.joystick import G1WalkFlatCfg + from unilab.envs import ManagerBasedRlEnvCfg - return _ppo_owner_yaml_cfg("g1_walk_flat", backend, G1WalkFlatCfg, config_overrides) + return _ppo_owner_yaml_cfg("g1_walk_flat", backend, ManagerBasedRlEnvCfg, config_overrides) def _g1_rough_cfg(backend: str, config_overrides: list[str]) -> Any: - from unilab.tasks.locomotion.g1.joystick import G1WalkRoughCfg + from unilab.envs import ManagerBasedRlEnvCfg - return _sac_owner_yaml_cfg("g1_walk_rough", backend, G1WalkRoughCfg, config_overrides) + return _sac_owner_yaml_cfg("g1_walk_rough", backend, ManagerBasedRlEnvCfg, config_overrides) def _g1_motion_tracking_cfg(backend: str, config_overrides: list[str]) -> Any: @@ -400,9 +400,9 @@ def _ensure_sharpa_benchmark_grasp_cache(cfg: Any, _: str) -> None: def _g1_walk_env_cls() -> type: - from unilab.tasks.locomotion.g1.joystick import G1WalkEnv + from unilab.tasks.locomotion.g1 import make_g1_walk_env - return G1WalkEnv + return make_g1_walk_env def _g1_motion_tracking_env_cls() -> type: diff --git a/scripts/benchmark/torch_env/walk_flat.py b/scripts/benchmark/torch_env/walk_flat.py index 12d02b6e6..3d2ccc538 100644 --- a/scripts/benchmark/torch_env/walk_flat.py +++ b/scripts/benchmark/torch_env/walk_flat.py @@ -3,7 +3,8 @@ Faithful xp-port of the NumPy computation in the collector-timed sections of `uv run train --algo sac --task g1_walk_flat --sim mujoco` (num_envs=2048): -- `G1WalkEnv.update_state` (src/unilab/tasks/locomotion/g1/joystick.py): +- the legacy `G1WalkEnv.update_state` computation (pre-Manager-Based migration; + src/unilab/tasks/locomotion/g1/manager_terms.py now owns the same math): termination, `_compute_reward` (9 active terms under the SAC scales incl. per-term logging every 4 steps), `_compute_obs` (noise + concat, walk profile), and the done-triggered curriculum bookkeeping. diff --git a/src/unilab/tasks/locomotion/g1/__init__.py b/src/unilab/tasks/locomotion/g1/__init__.py index fbd4f6660..4b77e4553 100644 --- a/src/unilab/tasks/locomotion/g1/__init__.py +++ b/src/unilab/tasks/locomotion/g1/__init__.py @@ -1,19 +1,30 @@ -from .joystick import ( - G1WalkControlConfig, - G1WalkEnv, - G1WalkEnvCfg, - G1WalkFlatCfg, - G1WalkRewardConfig, - G1WalkRoughCfg, -) +"""Hydra-owned Manager-Based G1 locomotion production registrations.""" + +from unilab.base import registry +from unilab.envs import ManagerBasedRlEnvCfg + +from .manager_terms import G1WalkManagerBasedEnv, make_g1_walk_env from .symmetry import G1SymmetryAugmentation +registry.register_env_config("G1WalkFlat", ManagerBasedRlEnvCfg) +registry.register_env("G1WalkFlat", make_g1_walk_env, sim_backend="mujoco") +registry.register_env("G1WalkFlat", make_g1_walk_env, sim_backend="mjwarp") +registry.register_env("G1WalkFlat", make_g1_walk_env, sim_backend="motrix") + +registry.register_env_config("G1WalkRough", ManagerBasedRlEnvCfg) +registry.register_env("G1WalkRough", make_g1_walk_env, sim_backend="mujoco") +registry.register_env("G1WalkRough", make_g1_walk_env, sim_backend="motrix") + +registry.register_env_config("G1Walk23DofFlat", ManagerBasedRlEnvCfg) +registry.register_env("G1Walk23DofFlat", make_g1_walk_env, sim_backend="mujoco") +registry.register_env("G1Walk23DofFlat", make_g1_walk_env, sim_backend="motrix") + +registry.register_env_config("G1Walk23DofRough", ManagerBasedRlEnvCfg) +registry.register_env("G1Walk23DofRough", make_g1_walk_env, sim_backend="mujoco") +registry.register_env("G1Walk23DofRough", make_g1_walk_env, sim_backend="motrix") + __all__ = [ "G1SymmetryAugmentation", - "G1WalkControlConfig", - "G1WalkEnv", - "G1WalkEnvCfg", - "G1WalkFlatCfg", - "G1WalkRewardConfig", - "G1WalkRoughCfg", + "G1WalkManagerBasedEnv", + "make_g1_walk_env", ] diff --git a/src/unilab/tasks/locomotion/g1/joystick.py b/src/unilab/tasks/locomotion/g1/joystick.py deleted file mode 100644 index 0191ca44b..000000000 --- a/src/unilab/tasks/locomotion/g1/joystick.py +++ /dev/null @@ -1,717 +0,0 @@ -"""G1 joystick locomotion environments.""" - -from __future__ import annotations - -import math -from dataclasses import dataclass, field -from typing import Any - -import numpy as np - -from unilab.assets import ASSETS_ROOT_PATH -from unilab.base import registry -from unilab.base.augmentation import SymmetryObsLayout -from unilab.base.backend import create_backend, env_backend_kwargs -from unilab.base.curriculum import EpisodeLengthTracker, PenaltyCurriculum -from unilab.base.np_env import NpEnvState -from unilab.base.scene import SceneCfg -from unilab.dtype_config import get_global_dtype -from unilab.tasks.locomotion.common import rewards -from unilab.tasks.locomotion.common.commands import ( - Commands, - sample_heading_commands, - zero_small_xy_commands, -) -from unilab.tasks.locomotion.common.domain_rand import DomainRandConfig -from unilab.tasks.locomotion.common.dr_provider import LocomotionDRProvider -from unilab.tasks.locomotion.common.rewards import RewardContext - -from .base import G1BaseCfg, G1BaseEnv - - -@dataclass -class G1DomainRandConfig(DomainRandConfig): - randomize_kp: bool = True - kp_multiplier_range: list[float] = field(default_factory=lambda: [0.9, 1.1]) - - randomize_kd: bool = True - kd_multiplier_range: list[float] = field(default_factory=lambda: [0.9, 1.1]) - - randomize_body_gravity_compensation: bool = False - body_gravity_compensation_range: list[float] = field(default_factory=lambda: [-0.1, 0.3]) - body_gravity_compensation_bodies: list[str] = field( - default_factory=lambda: ["pelvis", "torso_link"] - ) - - -@dataclass -class InitState: - pos = [0.0, 0.0, 0.754] - - -def build_upper_body_pose_weights(pose_weights: list[float]) -> np.ndarray: - weights = np.asarray(pose_weights, dtype=get_global_dtype()).copy() - weights[:12] = 0.0 - return np.asarray(weights, dtype=get_global_dtype()) - - -def compute_feet_phase_height_targets( - gait_phase: np.ndarray, swing_height: float -) -> tuple[np.ndarray, np.ndarray]: - def cubic_bezier_height(phi: np.ndarray, swing_height: float) -> np.ndarray: - phi_normalized = np.fmod(phi + np.pi, 2 * np.pi) - np.pi - x = (phi_normalized + np.pi) / (2 * np.pi) - - def cubic_bezier_interpolation( - y_start: np.ndarray, y_end: np.ndarray, t: np.ndarray - ) -> np.ndarray: - y_diff = y_end - y_start - bezier = t**3 + 3 * (t**2 * (1 - t)) - return np.asarray(y_start + y_diff * bezier, dtype=get_global_dtype()) - - stance = cubic_bezier_interpolation(np.zeros_like(x), np.full_like(x, swing_height), 2 * x) - swing = cubic_bezier_interpolation( - np.full_like(x, swing_height), np.zeros_like(x), 2 * x - 1 - ) - return np.where(x <= 0.5, stance, swing) - - left_target = cubic_bezier_height(gait_phase[:, 0], swing_height) - right_target = cubic_bezier_height(gait_phase[:, 1], swing_height) - return left_target, right_target - - -LEFT_FOOT_CONTACT_SENSORS = [f"left_foot_contact_{i}" for i in range(4)] -RIGHT_FOOT_CONTACT_SENSORS = [f"right_foot_contact_{i}" for i in range(4)] - - -def _scalarize_sensor_values(sensor_values: np.ndarray) -> np.ndarray: - sensor_array = np.asarray(sensor_values, dtype=get_global_dtype()) - if sensor_array.ndim == 1: - return sensor_array - if sensor_array.ndim == 2 and sensor_array.shape[1] == 1: - return sensor_array[:, 0] - raise ValueError(f"Expected scalar sensor values, got shape {sensor_array.shape}") - - -def compute_aggregated_foot_contact(backend: Any, sensor_names: list[str]) -> np.ndarray: - contacts = [_scalarize_sensor_values(backend.get_sensor_data(name)) for name in sensor_names] - return np.asarray(np.any(np.stack(contacts, axis=1) > 0.5, axis=1), dtype=np.bool_) - - -def compute_feet_phase_contact_targets( - gait_phase: np.ndarray, swing_height: float -) -> tuple[np.ndarray, np.ndarray]: - left_target, right_target = compute_feet_phase_height_targets(gait_phase, swing_height) - contact_height_threshold = swing_height * 0.5 - return left_target <= contact_height_threshold, right_target <= contact_height_threshold - - -def compute_forward_speed_gate(linvel: np.ndarray, min_forward_speed: float) -> np.ndarray: - forward_speed = np.maximum(linvel[:, 0], 0.0) - return np.asarray(forward_speed >= min_forward_speed, dtype=get_global_dtype()) - - -def compute_forward_command_mask(commands: np.ndarray) -> np.ndarray: - return np.asarray(np.maximum(commands[:, 0], 0.0) > 1.0e-6, dtype=get_global_dtype()) - - -@dataclass -class G1RewardConfig: - scales: dict[str, float] - tracking_sigma: float - gait_frequency: float - feet_phase_swing_height: float - feet_phase_tracking_sigma: float - base_height_target: float - min_base_height: float - max_tilt_deg: float - min_forward_speed_for_gait_reward: float = 0.0 - close_feet_threshold: float = 0.15 - pose_weights: list[float] = field( - default_factory=lambda: [ - 0.01, - 1.0, - 5.0, - 0.01, - 5.0, - 5.0, - 0.01, - 1.0, - 5.0, - 0.01, - 5.0, - 5.0, - 50.0, - 50.0, - 50.0, - 50.0, - 50.0, - 50.0, - 50.0, - 50.0, - 50.0, - 50.0, - 50.0, - 50.0, - 50.0, - 50.0, - 50.0, - 50.0, - 50.0, - ] - ) - - -@dataclass -class CurriculumConfig: - enabled: bool = False - initial_scale: float = 0.5 - min_scale: float = 0.5 - max_scale: float = 1.0 - level_down_threshold: float = 150.0 - level_up_threshold: float = 750.0 - degree: float = 0.001 - - -@dataclass -class G1WalkEnvCfg(G1BaseCfg): - scene: SceneCfg = field( # pyright: ignore[reportIncompatibleVariableOverride] - default_factory=lambda: SceneCfg( - model_file=str(ASSETS_ROOT_PATH / "robots" / "g1" / "scene_flat.xml") - ) - ) - max_episode_seconds: float | None = 20.0 - init_state: InitState = field(default_factory=InitState) - commands: Commands = field(default_factory=Commands) - reward_config: G1RewardConfig | None = None - domain_rand: G1DomainRandConfig = field(default_factory=G1DomainRandConfig) - gait_phase_init_mode: str = "offset_phase" - reset_base_qvel_limit: float = 0.5 - curriculum: CurriculumConfig = field(default_factory=CurriculumConfig) - - -class G1WalkDomainRandomizationProvider(LocomotionDRProvider): - def __init__( - self, - *, - base_kp: np.ndarray | None = None, - base_kd: np.ndarray | None = None, - base_dof_armature: np.ndarray | None = None, - ): - self._base_kp = base_kp - self._base_kd = base_kd - self._base_dof_armature = base_dof_armature - - def _get_base_actuator_gains(self, env: Any) -> tuple[np.ndarray | None, np.ndarray | None]: - return self._base_kp, self._base_kd - - def _get_reset_randomization_baselines( - self, env: Any - ) -> tuple[np.ndarray | None, np.ndarray | None, int | None, np.ndarray | None]: - return None, None, None, self._base_dof_armature - - def _get_qvel_limit(self, env: Any) -> float: - return float(env.cfg.reset_base_qvel_limit) - - def _build_extra_info_updates(self, env: Any, num_reset: int) -> dict[str, np.ndarray]: - updates = {"gait_phase": self._sample_gait_phase(env, num_reset)} - if getattr(env.cfg.commands, "heading_command", False): - updates["heading_commands"] = sample_heading_commands(env, num_reset) - return updates - - def _sample_commands(self, env: Any, num_reset: int) -> np.ndarray: - commands = super()._sample_commands(env, num_reset) - zero_small_xy_commands(commands) - standing_prob = float(getattr(env.cfg.commands, "rel_standing_envs", 0.0)) - if standing_prob > 0.0: - standing = np.random.uniform(size=(num_reset,)) < min(standing_prob, 1.0) - commands[standing] = 0.0 - if getattr(env.cfg.commands, "heading_command", False): - commands[:, 2] = 0.0 - return commands - - def _sample_gait_phase(self, env: Any, num_reset: int) -> np.ndarray: - mode = env.cfg.gait_phase_init_mode - if mode == "independent": - left = np.random.uniform(0.0, 2.0 * np.pi, size=(num_reset,)) - right = np.random.uniform(0.0, 2.0 * np.pi, size=(num_reset,)) - return np.asarray(np.column_stack([left, right]), dtype=get_global_dtype()) - - 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 _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]: - return env._compute_obs(info_updates, linvel, gyro, gravity, dof_pos, dof_vel) # type: ignore[no-any-return] - - -class G1WalkEnv(G1BaseEnv): - _cfg: G1WalkEnvCfg # pyright: ignore[reportIncompatibleVariableOverride] - _reward_cfg: Any - - def __init__(self, cfg: G1WalkEnvCfg, num_envs=1, backend_type="mujoco"): - if cfg.reward_config is None: - raise ValueError("reward_config must be provided via Hydra configuration") - backend = create_backend( - backend_type, - cfg.scene, - num_envs, - cfg.sim_dt, - base_name=cfg.asset.base_name, - push_body_name=cfg.domain_rand.push_body_name, - **env_backend_kwargs(cfg), - ) - super().__init__(cfg, backend, num_envs) - self._enable_reward_log = True - self._reward_cfg = cfg.reward_config - - self._gait_phase_delta = float( - 2.0 * math.pi * self._reward_cfg.gait_frequency * cfg.ctrl_dt - ) - self._pose_weights = np.array(self._reward_cfg.pose_weights, dtype=get_global_dtype()) - if self._pose_weights.shape[0] != self._num_action: - raise ValueError("pose_weights length mismatch") - self._upper_body_pose_weights = build_upper_body_pose_weights(self._reward_cfg.pose_weights) - self._episode_tracker: EpisodeLengthTracker | None = None - self._penalty_curriculum: PenaltyCurriculum | None = None - if cfg.curriculum.enabled: - self._episode_tracker = EpisodeLengthTracker(num_envs) - self._penalty_curriculum = PenaltyCurriculum( - self, - enabled=True, - initial_scale=cfg.curriculum.initial_scale, - min_scale=cfg.curriculum.min_scale, - max_scale=cfg.curriculum.max_scale, - level_down_threshold=cfg.curriculum.level_down_threshold, - level_up_threshold=cfg.curriculum.level_up_threshold, - degree=cfg.curriculum.degree, - ) - - self._init_reward_functions() - base_kp: np.ndarray | None = None - base_kd: np.ndarray | None = None - if cfg.domain_rand.randomize_kp or cfg.domain_rand.randomize_kd: - base_kp, base_kd = backend.get_actuator_gains() - base_dof_armature = ( - backend.get_dof_armature() if cfg.domain_rand.randomize_dof_armature else None - ) - dr_provider = G1WalkDomainRandomizationProvider( - base_kp=base_kp, - base_kd=base_kd, - base_dof_armature=base_dof_armature, - ) - self._init_domain_randomization(dr_provider) - - @property - def obs_groups_spec(self) -> dict[str, int]: - # gyro(3) + gravity(3) + diff(29) + dof_vel(29) + action(29) + cmd(3) + phase(2) = 98 - return {"obs": 98, "critic": 101} - - def _init_reward_functions(self): - self._reward_fns: dict[str, Any] = { - "tracking_lin_vel": rewards.tracking_lin_vel, - "tracking_ang_vel": rewards.tracking_ang_vel, - "forward_progress": rewards.forward_progress, - "under_speed": rewards.under_speed, - "lin_vel_z": rewards.lin_vel_z, - "orientation": rewards.orientation, - "penalty_orientation": rewards.orientation, - "ang_vel_xy": rewards.ang_vel_xy, - "penalty_ang_vel_xy": rewards.ang_vel_xy, - "action_rate": rewards.action_rate, - "penalty_action_rate": rewards.action_rate, - "base_height": rewards.base_height, - "pose": rewards.weighted_pose, - "upper_body_pose": self._reward_upper_body_pose, - "penalty_close_feet_xy": self._reward_close_feet_xy, - "penalty_feet_ori": self._reward_feet_ori, - "feet_phase": self._reward_feet_phase, - "feet_phase_contrast": self._reward_feet_phase_contrast, - "feet_phase_contact": self._reward_feet_phase_contact, - "feet_double_stance": self._reward_feet_double_stance, - "feet_air_time": self._reward_feet_air_time, - "alive": rewards.alive, - } - - def _terrain_relative_base_height(self) -> np.ndarray: - return np.asarray(self._backend.get_base_pos()[:, 2], dtype=get_global_dtype()) - - def update_state(self, state: NpEnvState) -> NpEnvState: - linvel = self.get_local_linvel() - gyro = self.get_gyro() - gravity = self._backend.get_sensor_data(self._cfg.sensor.upvector) - dof_pos = self.get_dof_pos() - dof_vel = self.get_dof_vel() - - max_tilt_rad = np.deg2rad(self._reward_cfg.max_tilt_deg) - tilt = np.arccos(np.clip(gravity[:, 2], -1, 1)) - terminated = np.logical_or( - tilt > max_tilt_rad, - self._terrain_relative_base_height() < self._reward_cfg.min_base_height, - ) - reward = self._compute_reward(state.info, linvel, gyro, gravity, dof_pos, dof_vel) - obs = self._compute_obs(state.info, linvel, gyro, gravity, dof_pos, dof_vel) - - state = state.replace(obs=obs, reward=reward, terminated=terminated) - - done = state.terminated | state.truncated - if self._episode_tracker is None or self._penalty_curriculum is None or not np.any(done): - return state - - done_indices = np.where(done)[0] - episode_lengths = state.info["steps"][done_indices] + 1 - self._episode_tracker.update(episode_lengths) - self._penalty_curriculum.update(self._episode_tracker.average_length) - - if "log" not in state.info: - state.info["log"] = {} - state.info["log"]["curriculum/average_episode_length"] = float( - self._episode_tracker.average_length - ) - state.info["log"]["curriculum/penalty_scale"] = float( - self._penalty_curriculum.current_scale - ) - return state - - def _compute_obs( - self, info: dict, linvel, gyro, gravity, dof_pos, dof_vel - ) -> dict[str, np.ndarray]: - noise_cfg = self._cfg.noise_config - diff = dof_pos - self.default_angles - command = info["commands"] - last_actions = info.get("current_actions", np.zeros_like(diff)) - gait_phase = info.get("gait_phase", np.zeros((self._num_envs, 2), dtype=get_global_dtype())) - walk_profile = self._uses_walk_observation_profile() - - noisy_gyro = self._obs_noise(gyro, noise_cfg.scale_gyro) - noisy_gravity = self._obs_noise(gravity, noise_cfg.scale_gravity) - noisy_diff = self._obs_noise(diff, noise_cfg.scale_joint_angle) - noisy_dof_vel = self._obs_noise(dof_vel, noise_cfg.scale_joint_vel) - actor_gyro_scale = 0.25 if walk_profile else 1.0 - actor_dof_vel_scale = 0.05 if walk_profile else 1.0 - - actor = np.concatenate( - [ - noisy_gyro * actor_gyro_scale, - -noisy_gravity, - noisy_diff, - noisy_dof_vel * actor_dof_vel_scale, - last_actions, - command, - gait_phase, - ], - axis=1, - dtype=get_global_dtype(), - ) - - critic_gyro_scale = 0.25 if walk_profile else 1.0 - critic_dof_vel_scale = 0.05 if walk_profile else 1.0 - critic_linvel_scale = 2.0 if walk_profile else 1.0 - critic_base = np.concatenate( - [ - gyro * critic_gyro_scale, - -gravity, - diff, - dof_vel * critic_dof_vel_scale, - last_actions, - command, - gait_phase, - ], - axis=1, - dtype=get_global_dtype(), - ) - critic = np.concatenate( - [ - critic_base, - np.asarray(linvel * critic_linvel_scale, dtype=get_global_dtype()), - ], - axis=1, - dtype=get_global_dtype(), - ) - - return {"obs": actor, "critic": critic} - - def _uses_walk_observation_profile(self) -> bool: - scales = getattr(getattr(self, "_reward_cfg", None), "scales", None) - if scales is None: - reward_cfg = getattr(self._cfg, "reward_config", None) - scales = getattr(reward_cfg, "scales", None) - - if scales is not None: - if any( - key in scales - for key in ( - "penalty_orientation", - "penalty_ang_vel_xy", - "penalty_action_rate", - "alive", - ) - ): - return True - if any(key in scales for key in ("orientation", "ang_vel_xy", "action_rate")): - return False - - curriculum = getattr(self._cfg, "curriculum", None) - return bool(curriculum is not None and curriculum.enabled) - - def _actor_symmetry_obs_layout(self) -> SymmetryObsLayout: - return ( - ("gyro", 3), - ("gravity", 3), - ("dof_pos", self._num_action), - ("dof_vel", self._num_action), - ("actions", self._num_action), - ("command", 3), - ("gait_phase", 2), - ) - - def get_symmetry_obs_layouts(self) -> dict[str, SymmetryObsLayout]: - actor_layout = self._actor_symmetry_obs_layout() - return { - "obs": actor_layout, - "critic": (*actor_layout, ("linvel", 3)), - } - - def build_symmetry_augmentation(self, *, device: str): - try: - actuator_names = self._backend.get_actuator_names() - except NotImplementedError: - return None - from unilab.tasks.locomotion.g1.symmetry import G1SymmetryAugmentation - - return G1SymmetryAugmentation( - actuator_names, - self.get_symmetry_obs_layouts(), - device=device, - ) - - def _build_reward_context( - self, info: dict, linvel, gyro, gravity, dof_pos, dof_vel - ) -> RewardContext: - return RewardContext( - info=info, - linvel=linvel, - gyro=gyro, - dof_pos=dof_pos, - num_envs=self._num_envs, - default_angles=self.default_angles, - tracking_sigma=self._reward_cfg.tracking_sigma, - base_height_target=self._reward_cfg.base_height_target, - base_height=self._backend.get_base_pos()[:, 2], - gravity=gravity, - dof_vel=dof_vel, - pose_weights=self._pose_weights, - ) - - def _compute_reward(self, info: dict, linvel, gyro, gravity, dof_pos, dof_vel) -> np.ndarray: - cfg = self._reward_cfg - ctx = self._build_reward_context(info, linvel, gyro, gravity, dof_pos, dof_vel) - return rewards.run_reward_dispatch( - scales=cfg.scales, - fns=self._reward_fns, - ctx=ctx, - info=info, - enable_log=self._enable_reward_log, - ctrl_dt=self._cfg.ctrl_dt, - ) - - def _reward_feet_phase(self, ctx: RewardContext): - """Reward gait phase tracking by encouraging the expected swing-foot height.""" - left_foot = self._backend.get_sensor_data("left_foot_pos") - right_foot = self._backend.get_sensor_data("right_foot_pos") - gait_phase = ctx.info.get( - "gait_phase", np.zeros((self._num_envs, 2), dtype=get_global_dtype()) - ) - swing_height = self._reward_cfg.feet_phase_swing_height - left_target, right_target = compute_feet_phase_height_targets(gait_phase, swing_height) - left_error = np.square(left_foot[:, 2] - left_target) - right_error = np.square(right_foot[:, 2] - right_target) - reward = np.exp(-(left_error + right_error) / self._reward_cfg.feet_phase_tracking_sigma) - return np.asarray(reward * self._gait_reward_gate(ctx.linvel), dtype=get_global_dtype()) - - def _gait_reward_gate(self, linvel: np.ndarray) -> np.ndarray: - min_forward_speed = getattr(self._reward_cfg, "min_forward_speed_for_gait_reward", 0.0) - return compute_forward_speed_gate(linvel, min_forward_speed) - - def _reward_feet_phase_contrast(self, ctx: RewardContext): - left_foot = self._backend.get_sensor_data("left_foot_pos") - right_foot = self._backend.get_sensor_data("right_foot_pos") - gait_phase = ctx.info.get( - "gait_phase", np.zeros((self._num_envs, 2), dtype=get_global_dtype()) - ) - swing_height = self._reward_cfg.feet_phase_swing_height - left_target, right_target = compute_feet_phase_height_targets(gait_phase, swing_height) - actual_delta = left_foot[:, 2] - right_foot[:, 2] - target_delta = left_target - right_target - error = np.square(actual_delta - target_delta) - reward = np.exp(-error / self._reward_cfg.feet_phase_tracking_sigma) - return np.asarray(reward * self._gait_reward_gate(ctx.linvel), dtype=get_global_dtype()) - - def _reward_feet_phase_contact(self, ctx: RewardContext): - gait_phase = ctx.info.get( - "gait_phase", np.zeros((self._num_envs, 2), dtype=get_global_dtype()) - ) - swing_height = self._reward_cfg.feet_phase_swing_height - left_target_contact, right_target_contact = compute_feet_phase_contact_targets( - gait_phase, swing_height - ) - left_contact = compute_aggregated_foot_contact(self._backend, LEFT_FOOT_CONTACT_SENSORS) - right_contact = compute_aggregated_foot_contact(self._backend, RIGHT_FOOT_CONTACT_SENSORS) - left_match = np.asarray(left_contact == left_target_contact, dtype=get_global_dtype()) - right_match = np.asarray(right_contact == right_target_contact, dtype=get_global_dtype()) - reward = np.asarray(0.5 * (left_match + right_match), dtype=get_global_dtype()) - return np.asarray(reward * self._gait_reward_gate(ctx.linvel), dtype=get_global_dtype()) - - def _reward_feet_double_stance(self, ctx: RewardContext): - commands = ctx.info.get("commands", np.zeros((self._num_envs, 3), dtype=get_global_dtype())) - left_contact = compute_aggregated_foot_contact(self._backend, LEFT_FOOT_CONTACT_SENSORS) - right_contact = compute_aggregated_foot_contact(self._backend, RIGHT_FOOT_CONTACT_SENSORS) - double_stance = np.asarray( - np.logical_and(left_contact, right_contact), dtype=get_global_dtype() - ) - return np.asarray( - double_stance * compute_forward_command_mask(commands), dtype=get_global_dtype() - ) - - def _reward_feet_ori(self, ctx: RewardContext): - left_foot_quat = self._backend.get_sensor_data("left_foot_quat") - right_foot_quat = self._backend.get_sensor_data("right_foot_quat") - return ( - np.square(left_foot_quat[:, 1]) - + np.square(left_foot_quat[:, 2]) - + np.square(right_foot_quat[:, 1]) - + np.square(right_foot_quat[:, 2]) - ) - - def _reward_close_feet_xy(self, ctx: RewardContext): - left_foot = self._backend.get_sensor_data("left_foot_pos") - right_foot = self._backend.get_sensor_data("right_foot_pos") - feet_dist = np.linalg.norm(left_foot[:, :2] - right_foot[:, :2], axis=1) - return np.where( - feet_dist < self._reward_cfg.close_feet_threshold, - np.square(feet_dist - self._reward_cfg.close_feet_threshold), - 0.0, - ) - - def _reward_feet_air_time(self, ctx: RewardContext): - air_time = ctx.info.get( - "feet_air_time", np.zeros((self._num_envs, 2), dtype=get_global_dtype()) - ) - in_range = (air_time > 0.05) & (air_time < 0.5) - return np.sum(in_range.astype(float), axis=1) - - def _reward_upper_body_pose(self, ctx: RewardContext): - diff = ctx.dof_pos - self.default_angles - return np.asarray( - np.sum(self._upper_body_pose_weights * np.square(diff), axis=1), - dtype=get_global_dtype(), - ) - - def apply_action(self, actions: np.ndarray, state: NpEnvState) -> np.ndarray: - state.info["last_actions"] = state.info.get("current_actions", np.zeros_like(actions)) - state.info["current_actions"] = actions - - gait_phase = state.info.get( - "gait_phase", np.zeros((self._num_envs, 2), dtype=get_global_dtype()) - ) - gait_phase[:, 0] = (gait_phase[:, 0] + self._gait_phase_delta) % (2 * np.pi) - gait_phase[:, 1] = (gait_phase[:, 1] + self._gait_phase_delta) % (2 * np.pi) - state.info["gait_phase"] = gait_phase - - ctrl: np.ndarray = actions * self._cfg.control_config.action_scale + self.default_angles - return ctrl - - -def _walk_curriculum() -> CurriculumConfig: - return CurriculumConfig( - enabled=True, - initial_scale=0.5, - min_scale=0.5, - max_scale=1.0, - level_down_threshold=150.0, - level_up_threshold=750.0, - degree=0.001, - ) - - -@dataclass -class G1WalkControlConfig: - action_scale: float = 1.0 - simulate_action_latency: bool = False - - -@dataclass -class G1WalkRewardConfig(G1RewardConfig): - """Align reward weights with holosoma G1 walking.""" - - -@registry.envcfg("G1WalkFlat") -@dataclass -class G1WalkFlatCfg(G1WalkEnvCfg): - reward_config: G1WalkRewardConfig | None = None # pyright: ignore[reportIncompatibleVariableOverride] - scene: SceneCfg = field( - default_factory=lambda: SceneCfg( - model_file=str(ASSETS_ROOT_PATH / "robots" / "g1" / "scene_flat.xml") - ) - ) - control_config: G1WalkControlConfig = field(default_factory=G1WalkControlConfig) # type: ignore[assignment] - curriculum: CurriculumConfig = field(default_factory=_walk_curriculum) - - -@registry.envcfg("G1WalkRough") -@dataclass -class G1WalkRoughCfg(G1WalkFlatCfg): - scene: SceneCfg = field( - default_factory=lambda: SceneCfg( - model_file=str(ASSETS_ROOT_PATH / "robots" / "g1" / "scene_rough.xml") - ) - ) - - -registry.register_env("G1WalkFlat", G1WalkEnv, sim_backend="mujoco") -registry.register_env("G1WalkFlat", G1WalkEnv, sim_backend="mjwarp") -registry.register_env("G1WalkFlat", G1WalkEnv, sim_backend="motrix") -registry.register_env("G1WalkRough", G1WalkEnv, sim_backend="mujoco") -registry.register_env("G1WalkRough", G1WalkEnv, sim_backend="motrix") - - -class G1Walk23DofEnv(G1WalkEnv): - @property - def obs_groups_spec(self) -> dict[str, int]: - # gyro(3) + gravity(3) + diff(23) + dof_vel(23) + action(23) + cmd(3) + phase(2) = 80 - return {"obs": 80, "critic": 83} - - -@registry.envcfg("G1Walk23DofFlat") -@dataclass -class G1Walk23DofFlatCfg(G1WalkFlatCfg): - scene: SceneCfg = field( - default_factory=lambda: SceneCfg( - model_file=str(ASSETS_ROOT_PATH / "robots" / "g1" / "scene_flat_23dof.xml") - ) - ) - - -@registry.envcfg("G1Walk23DofRough") -@dataclass -class G1Walk23DofRoughCfg(G1Walk23DofFlatCfg): - scene: SceneCfg = field( - default_factory=lambda: SceneCfg( - model_file=str(ASSETS_ROOT_PATH / "robots" / "g1" / "scene_rough_23dof.xml") - ) - ) - - -registry.register_env("G1Walk23DofFlat", G1Walk23DofEnv, sim_backend="mujoco") -registry.register_env("G1Walk23DofFlat", G1Walk23DofEnv, sim_backend="motrix") -registry.register_env("G1Walk23DofRough", G1Walk23DofEnv, sim_backend="mujoco") -registry.register_env("G1Walk23DofRough", G1Walk23DofEnv, sim_backend="motrix") diff --git a/src/unilab/tasks/locomotion/g1/manager_terms.py b/src/unilab/tasks/locomotion/g1/manager_terms.py new file mode 100644 index 000000000..e8756af14 --- /dev/null +++ b/src/unilab/tasks/locomotion/g1/manager_terms.py @@ -0,0 +1,1021 @@ +"""Task-owned Manager-Based terms for G1 biped locomotion. + +The equations come from UniLab's legacy G1Walk joystick tasks. The adaptation +uses community ``func + params`` terms, NumPy, cached cold-path sensor bindings, +and the base-owned entity facade; hot paths never parse assets or probe backend +privates. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, ClassVar, Protocol, cast +from weakref import WeakKeyDictionary + +import numpy as np + +from unilab.base.augmentation import SymmetryAugmentation, SymmetryObsLayout +from unilab.base.backend import create_backend, env_backend_kwargs +from unilab.base.curriculum import EpisodeLengthTracker +from unilab.dtype_config import get_global_dtype +from unilab.envs.manager_based_rl_env import ( + ManagerBasedRlEnv as _ConcreteManagerBasedRlEnv, +) +from unilab.envs.manager_based_rl_env import ( + ManagerBasedRlEnvCfg, + _resolve_backend_entity_contract, +) +from unilab.envs.mdp.commands.velocity_command import ( + UniformVelocityCommand, + UniformVelocityCommandCfg, +) +from unilab.managers.manager_base import ManagerTermBase, ManagerTermBaseCfg +from unilab.managers.scene_entity_config import SceneEntityCfg + +if TYPE_CHECKING: + from unilab.base.entity import Entity + from unilab.managers._types import ManagerBasedRlEnv, ManagerSensorView + + class _RewardTermCfgView(Protocol): + weight: float + + class _RewardManagerView(Protocol): + @property + def active_terms(self) -> list[str]: ... + + def get_term_cfg(self, term_name: str) -> _RewardTermCfgView: ... + + class _G1Env(ManagerBasedRlEnv, Protocol): + @property + def common_step_counter(self) -> int: ... + + @property + def reset_buf(self) -> np.ndarray: ... + + @property + def reward_manager(self) -> _RewardManagerView: ... + + +_DEFAULT_ASSET_CFG = SceneEntityCfg("robot") +_LEFT_FOOT_CONTACT_SENSORS = tuple(f"left_foot_contact_{index}" for index in range(4)) +_RIGHT_FOOT_CONTACT_SENSORS = tuple(f"right_foot_contact_{index}" for index in range(4)) +_FOOT_CONTACT_SENSORS = _LEFT_FOOT_CONTACT_SENSORS + _RIGHT_FOOT_CONTACT_SENSORS +_FOOT_POS_SENSORS = ("left_foot_pos", "right_foot_pos") +_FOOT_QUAT_SENSORS = ("left_foot_quat", "right_foot_quat") +_GAIT_INIT_MODES = ("offset_phase", "independent") + + +def _real( + term: str, + name: str, + value: Any, + *, + minimum: float | None = None, + strict_minimum: bool = False, +) -> float: + if isinstance(value, (bool, np.bool_)) or not isinstance(value, (int, float, np.number)): + raise TypeError(f"{term} {name} must be a real number") + result = float(value) + if not np.isfinite(result): + raise ValueError(f"{term} {name} must be finite") + if minimum is not None and (result <= minimum if strict_minimum else result < minimum): + relation = "greater than" if strict_minimum else "at least" + raise ValueError(f"{term} {name} must be {relation} {minimum}") + return result + + +def _weights(term: str, name: str, value: Any) -> np.ndarray: + if isinstance(value, (str, bytes)) or not isinstance(value, (tuple, list, np.ndarray)): + raise TypeError(f"{term} {name} must be a sequence of real numbers") + result = np.asarray( + [_real(term, f"{name}[{index}]", item) for index, item in enumerate(value)], + dtype=get_global_dtype(), + ) + if result.ndim != 1 or result.shape[0] == 0: + raise ValueError(f"{term} {name} must be a non-empty one-dimensional sequence") + return result + + +def _state(term: str, capability: str, value: Any, shape: tuple[int, ...]) -> np.ndarray: + if not isinstance(value, np.ndarray): + raise TypeError(f"{term} {capability} must be an np.ndarray") + if value.shape != shape: + raise ValueError(f"{term} {capability} must have shape {shape}, got {value.shape}") + if not np.isfinite(value).all(): + env_ids = np.flatnonzero(~np.isfinite(value).reshape(shape[0], -1).all(axis=1)).tolist() + raise ValueError(f"{term} {capability} contains NaN or Inf for environments {env_ids[:10]}") + return value + + +def _command(env: ManagerBasedRlEnv, term: str, command_name: str) -> np.ndarray: + if not isinstance(command_name, str) or not command_name: + raise ValueError(f"{term} command_name must be a non-empty string") + try: + command = env.command_manager.get_command(command_name) + except KeyError as exc: + raise KeyError(f"{term} command capability '{command_name}' is unavailable") from exc + if command is None: + raise KeyError(f"{term} command capability '{command_name}' is unavailable") + return _state(term, f"command '{command_name}'", command, (env.num_envs, 3)) + + +def _asset(env: ManagerBasedRlEnv, asset_cfg: SceneEntityCfg) -> Entity: + return cast("Entity", env.scene[asset_cfg.name]) + + +class _SensorTerm(ManagerTermBase): + """Cold-path named-sensor binding shared by G1 manager terms.""" + + _allowed_params: ClassVar[frozenset[str]] = frozenset() + + def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv): + super().__init__(env) + unexpected = set(cfg.params) - self._allowed_params + if unexpected: + raise TypeError(f"{self.name} received unsupported parameters: {sorted(unexpected)}") + + def _bind(self, sensor_names: tuple[str, ...]) -> ManagerSensorView: + try: + return self._env.scene.bind_sensor_data(sensor_names) + except (KeyError, TypeError, ValueError, NotImplementedError) as exc: + raise type(exc)( + f"Manager term '{self.name}' named-sensor capability could not be " + f"materialized for {sensor_names}: {exc}" + ) from exc + + @staticmethod + def _read(view: ManagerSensorView, term: str) -> np.ndarray: + try: + return view.read() + except (KeyError, TypeError, ValueError, NotImplementedError) as exc: + raise type(exc)( + f"Manager term '{term}' named-sensor capability failed on " + f"backend '{view.backend_type}': {exc}" + ) from exc + + +# --------------------------------------------------------------------------- +# Gait phase state (shared by the observation term and the gait reward terms) +# --------------------------------------------------------------------------- + + +def compute_feet_phase_height_targets( + gait_phase: np.ndarray, swing_height: float +) -> tuple[np.ndarray, np.ndarray]: + """Cubic-Bézier per-foot height targets, ported from the legacy G1 task.""" + + def cubic_bezier_height(phi: np.ndarray, swing_height: float) -> np.ndarray: + phi_normalized = np.fmod(phi + np.pi, 2 * np.pi) - np.pi + x = (phi_normalized + np.pi) / (2 * np.pi) + + def cubic_bezier_interpolation( + y_start: np.ndarray, y_end: np.ndarray, t: np.ndarray + ) -> np.ndarray: + y_diff = y_end - y_start + bezier = t**3 + 3 * (t**2 * (1 - t)) + return np.asarray(y_start + y_diff * bezier, dtype=get_global_dtype()) + + stance = cubic_bezier_interpolation(np.zeros_like(x), np.full_like(x, swing_height), 2 * x) + swing = cubic_bezier_interpolation( + np.full_like(x, swing_height), np.zeros_like(x), 2 * x - 1 + ) + return np.where(x <= 0.5, stance, swing) + + left_target = cubic_bezier_height(gait_phase[:, 0], swing_height) + right_target = cubic_bezier_height(gait_phase[:, 1], swing_height) + return left_target, right_target + + +def compute_feet_phase_contact_targets( + gait_phase: np.ndarray, swing_height: float +) -> tuple[np.ndarray, np.ndarray]: + """Expected per-foot contact flags derived from the Bézier height targets.""" + left_target, right_target = compute_feet_phase_height_targets(gait_phase, swing_height) + contact_height_threshold = swing_height * 0.5 + return left_target <= contact_height_threshold, right_target <= contact_height_threshold + + +@dataclass +class _G1GaitContext: + phase: np.ndarray # (num_envs, 2), radians in [0, 2*pi) + delta: float # 2*pi*frequency*ctrl_dt + frequency: float + init_mode: str + last_counter: int + + +_GAIT_CONTEXTS: WeakKeyDictionary[Any, _G1GaitContext] = WeakKeyDictionary() + + +def _gait_context(env: _G1Env, term: str, frequency: float, init_mode: str) -> _G1GaitContext: + context = _GAIT_CONTEXTS.get(env) + if context is None: + delta = float( + 2.0 + * math.pi + * frequency + * _real(term, "step_dt", env.step_dt, minimum=0.0, strict_minimum=True) + ) + context = _G1GaitContext( + phase=np.zeros((env.num_envs, 2), dtype=get_global_dtype()), + delta=delta, + frequency=frequency, + init_mode=init_mode, + last_counter=int(env.common_step_counter), + ) + _GAIT_CONTEXTS[env] = context + return context + if context.frequency != frequency or context.init_mode != init_mode: + raise ValueError( + f"{term} gait terms on one env must share frequency and init_mode; " + f"got frequency {context.frequency} vs {frequency}, " + f"init_mode {context.init_mode!r} vs {init_mode!r}" + ) + return context + + +def _advance_gait(env: _G1Env, context: _G1GaitContext) -> np.ndarray: + counter = env.common_step_counter + if isinstance(counter, (bool, np.bool_)) or not isinstance(counter, (int, np.integer)): + raise TypeError("G1 gait terms require an integer common_step_counter") + counter = int(counter) + if counter < context.last_counter: + raise ValueError("G1 gait terms common_step_counter cannot move backwards") + two_pi = 2.0 * np.pi + for _ in range(counter - context.last_counter): + context.phase = np.asarray( + np.fmod(context.phase + context.delta, two_pi), dtype=get_global_dtype() + ) + context.last_counter = counter + return context.phase + + +def _resample_gait(env: _G1Env, context: _G1GaitContext, env_ids: np.ndarray) -> None: + ids = np.asarray(env_ids, dtype=np.intp).reshape(-1) + count = len(ids) + if count == 0: + return + if context.init_mode == "independent": + left = env.rng.uniform(0.0, 2.0 * np.pi, size=(count,)) + right = env.rng.uniform(0.0, 2.0 * np.pi, size=(count,)) + else: + left = env.rng.uniform(0.0, 2.0 * np.pi, size=(count,)) + right = left + np.pi + context.phase[ids] = np.column_stack([left, right]).astype(get_global_dtype(), copy=False) + + +class G1GaitPhase(ManagerTermBase): + """Two-foot gait phase observation in radians, with per-reset phase sampling.""" + + _allowed_params: ClassVar[frozenset[str]] = frozenset({"frequency", "init_mode"}) + + def __init__(self, cfg: ManagerTermBaseCfg, env: _G1Env): + super().__init__(env) + unexpected = set(cfg.params) - self._allowed_params + if unexpected: + raise TypeError(f"{self.name} received unsupported parameters: {sorted(unexpected)}") + frequency = _real(self.name, "frequency", cfg.params.get("frequency", 1.5), minimum=0.0) + init_mode = cfg.params.get("init_mode", "offset_phase") + if init_mode not in _GAIT_INIT_MODES: + raise ValueError(f"{self.name} init_mode must be one of {_GAIT_INIT_MODES}") + self._context = _gait_context(env, self.name, frequency, init_mode) + + def __call__(self, env: _G1Env, **params: Any) -> np.ndarray: + del params + return _advance_gait(env, self._context).copy() + + def reset(self, env_ids: np.ndarray | slice | None = None) -> None: + if env_ids is None: + ids = np.arange(self.num_envs, dtype=np.intp) + elif isinstance(env_ids, slice): + ids = np.arange(self.num_envs, dtype=np.intp)[env_ids] + else: + ids = np.asarray(env_ids, dtype=np.intp).reshape(-1) + _resample_gait(cast("_G1Env", self._env), self._context, ids) + + +class _GaitRewardTerm(_SensorTerm): + """Shared gait-context, forward-speed gate, and foot-position binding.""" + + _allowed_params = frozenset( + { + "frequency", + "init_mode", + "swing_height", + "tracking_sigma", + "min_forward_speed", + "command_name", + } + ) + + def __init__(self, cfg: ManagerTermBaseCfg, env: _G1Env): + super().__init__(cfg, env) + frequency = _real(self.name, "frequency", cfg.params.get("frequency", 1.5), minimum=0.0) + init_mode = cfg.params.get("init_mode", "offset_phase") + if init_mode not in _GAIT_INIT_MODES: + raise ValueError(f"{self.name} init_mode must be one of {_GAIT_INIT_MODES}") + self._swing_height = _real( + self.name, "swing_height", cfg.params.get("swing_height", 0.09), minimum=0.0 + ) + self._tracking_sigma = _real( + self.name, + "tracking_sigma", + cfg.params.get("tracking_sigma", 0.008), + minimum=0.0, + strict_minimum=True, + ) + self._min_forward_speed = _real( + self.name, + "min_forward_speed", + cfg.params.get("min_forward_speed", 0.0), + minimum=0.0, + ) + command_name = cfg.params.get("command_name", "twist") + if not isinstance(command_name, str) or not command_name: + raise ValueError(f"{self.name} command_name must be a non-empty string") + self._command_name = command_name + self._context = _gait_context(env, self.name, frequency, init_mode) + self._feet_pos = self._bind(_FOOT_POS_SENSORS) + if self._feet_pos.dimensions != (3, 3): + raise ValueError( + f"{self.name} foot position sensors must expose 3-D xyz; received " + f"{self._feet_pos.dimensions} on backend '{self._feet_pos.backend_type}'" + ) + self._linvel = self._bind(("pelvis_local_linvel",)) + + def _targets(self, env: _G1Env) -> tuple[np.ndarray, np.ndarray]: + phase = _advance_gait(env, self._context) + return compute_feet_phase_height_targets(phase, self._swing_height) + + def _foot_heights(self) -> tuple[np.ndarray, np.ndarray]: + values = _state( + self.name, "foot position", self._read(self._feet_pos, self.name), (self.num_envs, 6) + ) + return values[:, 2], values[:, 5] + + def _gate(self, env: ManagerBasedRlEnv) -> np.ndarray: + linvel = _state( + self.name, + "pelvis linear velocity", + self._read(self._linvel, self.name), + (env.num_envs, 3), + ) + forward_speed = np.maximum(linvel[:, 0], 0.0) + return np.asarray(forward_speed >= self._min_forward_speed, dtype=get_global_dtype()) + + +class feet_phase(_GaitRewardTerm): + """Reward gait phase tracking by encouraging the expected swing-foot height.""" + + def __call__(self, env: _G1Env, **params: Any) -> np.ndarray: + del params + left_target, right_target = self._targets(env) + left_height, right_height = self._foot_heights() + error = np.square(left_height - left_target) + np.square(right_height - right_target) + reward = np.exp(-error / self._tracking_sigma) + return np.asarray(reward * self._gate(env), dtype=get_global_dtype()) + + +class feet_phase_contrast(_GaitRewardTerm): + """Reward left/right foot-height contrast against the gait-phase targets.""" + + def __call__(self, env: _G1Env, **params: Any) -> np.ndarray: + del params + left_target, right_target = self._targets(env) + left_height, right_height = self._foot_heights() + error = np.square((left_height - right_height) - (left_target - right_target)) + reward = np.exp(-error / self._tracking_sigma) + return np.asarray(reward * self._gate(env), dtype=get_global_dtype()) + + +class _FootContactTerm(_GaitRewardTerm): + """Adds the aggregated per-foot contact binding shared by contact gait terms.""" + + def __init__(self, cfg: ManagerTermBaseCfg, env: _G1Env): + super().__init__(cfg, env) + self._contacts = self._bind(_FOOT_CONTACT_SENSORS) + if self._contacts.dimensions != (1,) * 8: + raise ValueError( + f"{self.name} foot contact sensors must each expose 1-D found; received " + f"{self._contacts.dimensions} on backend '{self._contacts.backend_type}'" + ) + + def _contact_pair(self) -> tuple[np.ndarray, np.ndarray]: + values = _state( + self.name, "foot contact", self._read(self._contacts, self.name), (self.num_envs, 8) + ) + left = np.any(values[:, :4] > 0.5, axis=1) + right = np.any(values[:, 4:] > 0.5, axis=1) + return left, right + + +class feet_phase_contact(_FootContactTerm): + """Reward foot contact matching the expected stance phase of the gait.""" + + def __call__(self, env: _G1Env, **params: Any) -> np.ndarray: + del params + phase = _advance_gait(env, self._context) + left_target, right_target = compute_feet_phase_contact_targets(phase, self._swing_height) + left_contact, right_contact = self._contact_pair() + left_match = np.asarray(left_contact == left_target, dtype=get_global_dtype()) + right_match = np.asarray(right_contact == right_target, dtype=get_global_dtype()) + reward = np.asarray(0.5 * (left_match + right_match), dtype=get_global_dtype()) + return np.asarray(reward * self._gate(env), dtype=get_global_dtype()) + + +class feet_double_stance(_FootContactTerm): + """Penalize double-stance contact while a forward command is active.""" + + def __call__(self, env: ManagerBasedRlEnv, **params: Any) -> np.ndarray: + del params + command = _command(env, self.name, self._command_name) + left_contact, right_contact = self._contact_pair() + double_stance = np.asarray( + np.logical_and(left_contact, right_contact), dtype=get_global_dtype() + ) + forward_mask = np.asarray(np.maximum(command[:, 0], 0.0) > 1.0e-6, dtype=get_global_dtype()) + return np.asarray(double_stance * forward_mask, dtype=get_global_dtype()) + + +class feet_air_time(_FootContactTerm): + """Count feet whose current air time sits inside the rewarded window.""" + + def __init__(self, cfg: ManagerTermBaseCfg, env: _G1Env): + super().__init__(cfg, env) + self._air_time = np.zeros((env.num_envs, 2), dtype=get_global_dtype()) + + def reset(self, env_ids: np.ndarray | slice | None = None) -> None: + self._air_time[env_ids if env_ids is not None else slice(None)] = 0.0 + + def __call__(self, env: ManagerBasedRlEnv, **params: Any) -> np.ndarray: + del params + left_contact, right_contact = self._contact_pair() + contact = np.column_stack([left_contact, right_contact]) + step_dt = _real(self.name, "step_dt", env.step_dt, minimum=0.0, strict_minimum=True) + self._air_time = np.where(contact, 0.0, self._air_time + step_dt).astype( + get_global_dtype(), copy=False + ) + in_range = (self._air_time > 0.05) & (self._air_time < 0.5) + return np.asarray( + np.sum(in_range.astype(get_global_dtype()), axis=1), dtype=get_global_dtype() + ) + + +# --------------------------------------------------------------------------- +# Velocity / orientation reward terms (sensor-bound, legacy equations) +# --------------------------------------------------------------------------- + + +class _LinVelTerm(_SensorTerm): + def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv): + super().__init__(cfg, env) + self._linvel = self._bind(("pelvis_local_linvel",)) + + def _read_linvel(self, env: ManagerBasedRlEnv) -> np.ndarray: + return _state( + self.name, + "pelvis linear velocity", + self._read(self._linvel, self.name), + (env.num_envs, 3), + ) + + +class _GyroTerm(_SensorTerm): + def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv): + super().__init__(cfg, env) + self._gyro = self._bind(("torso_gyro",)) + + def _read_gyro(self, env: ManagerBasedRlEnv) -> np.ndarray: + return _state(self.name, "torso gyro", self._read(self._gyro, self.name), (env.num_envs, 3)) + + +class _UpvectorTerm(_SensorTerm): + def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv): + super().__init__(cfg, env) + self._upvector = self._bind(("torso_upvector",)) + + def _read_upvector(self, env: ManagerBasedRlEnv) -> np.ndarray: + return _state( + self.name, "torso upvector", self._read(self._upvector, self.name), (env.num_envs, 3) + ) + + +class track_lin_vel(_LinVelTerm): + """Exponential reward for tracking commanded xy linear velocity.""" + + _allowed_params = frozenset({"tracking_sigma", "command_name"}) + + def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv): + super().__init__(cfg, env) + self._sigma = _real( + self.name, + "tracking_sigma", + cfg.params.get("tracking_sigma", 0.25), + minimum=0.0, + strict_minimum=True, + ) + command_name = cfg.params.get("command_name", "twist") + if not isinstance(command_name, str) or not command_name: + raise ValueError(f"{self.name} command_name must be a non-empty string") + self._command_name = command_name + + def __call__(self, env: ManagerBasedRlEnv, **params: Any) -> np.ndarray: + del params + linvel = self._read_linvel(env) + command = _command(env, self.name, self._command_name) + error = np.sum(np.square(command[:, :2] - linvel[:, :2]), axis=1) + return np.asarray(np.exp(-error / self._sigma), dtype=get_global_dtype()) + + +class track_ang_vel(_GyroTerm): + """Exponential reward for tracking commanded yaw angular velocity.""" + + _allowed_params = frozenset({"tracking_sigma", "command_name"}) + + def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv): + super().__init__(cfg, env) + self._sigma = _real( + self.name, + "tracking_sigma", + cfg.params.get("tracking_sigma", 0.25), + minimum=0.0, + strict_minimum=True, + ) + command_name = cfg.params.get("command_name", "twist") + if not isinstance(command_name, str) or not command_name: + raise ValueError(f"{self.name} command_name must be a non-empty string") + self._command_name = command_name + + def __call__(self, env: ManagerBasedRlEnv, **params: Any) -> np.ndarray: + del params + gyro = self._read_gyro(env) + command = _command(env, self.name, self._command_name) + error = np.square(command[:, 2] - gyro[:, 2]) + return np.asarray(np.exp(-error / self._sigma), dtype=get_global_dtype()) + + +class forward_progress(_LinVelTerm): + """Reward forward progress relative to commanded speed.""" + + _allowed_params = frozenset({"command_name"}) + + def __call__(self, env: ManagerBasedRlEnv, **params: Any) -> np.ndarray: + command = _command(env, self.name, _term_command_name(self.name, params)) + linvel = self._read_linvel(env) + commanded_speed = np.maximum(command[:, 0], 1e-6) + forward_speed = np.maximum(linvel[:, 0], 0.0) + return np.asarray( + np.minimum(forward_speed / commanded_speed, 1.0), dtype=get_global_dtype() + ) + + +def _term_command_name(term: str, params: dict[str, Any]) -> str: + name = params.get("command_name", "twist") + if not isinstance(name, str) or not name: + raise ValueError(f"{term} command_name must be a non-empty string") + return name + + +class under_speed(_LinVelTerm): + """Penalty for being below commanded forward speed.""" + + _allowed_params = frozenset({"command_name"}) + + def __call__(self, env: ManagerBasedRlEnv, **params: Any) -> np.ndarray: + command = _command(env, self.name, _term_command_name(self.name, params)) + linvel = self._read_linvel(env) + commanded_speed = np.maximum(command[:, 0], 1e-6) + forward_speed = np.maximum(linvel[:, 0], 0.0) + gap = np.maximum(command[:, 0] - forward_speed, 0.0) + return np.asarray(gap / commanded_speed, dtype=get_global_dtype()) + + +class lin_vel_z(_LinVelTerm): + """Penalty for vertical (z) linear velocity.""" + + def __call__(self, env: ManagerBasedRlEnv, **params: Any) -> np.ndarray: + del params + linvel = self._read_linvel(env) + return np.asarray(np.square(linvel[:, 2]), dtype=get_global_dtype()) + + +class ang_vel_xy(_GyroTerm): + """Penalty for roll/pitch angular velocity.""" + + def __call__(self, env: ManagerBasedRlEnv, **params: Any) -> np.ndarray: + del params + gyro = self._read_gyro(env) + return np.asarray(np.sum(np.square(gyro[:, :2]), axis=1), dtype=get_global_dtype()) + + +class orientation(_UpvectorTerm): + """Penalty for deviation from upright orientation (roll/pitch).""" + + def __call__(self, env: ManagerBasedRlEnv, **params: Any) -> np.ndarray: + del params + upvector = self._read_upvector(env) + return np.asarray( + np.square(upvector[:, 0]) + np.square(upvector[:, 1]), dtype=get_global_dtype() + ) + + +class g1_tilt_exceeded(_UpvectorTerm): + """Terminate when the base tilt from upright exceeds ``max_tilt_deg``.""" + + _allowed_params = frozenset({"max_tilt_deg"}) + + def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv): + super().__init__(cfg, env) + max_tilt_deg = _real(self.name, "max_tilt_deg", cfg.params.get("max_tilt_deg"), minimum=0.0) + self._max_tilt_rad = math.radians(max_tilt_deg) + + def __call__(self, env: ManagerBasedRlEnv, **params: Any) -> np.ndarray: + del params + upvector = self._read_upvector(env) + tilt = np.arccos(np.clip(upvector[:, 2], -1.0, 1.0)) + return np.asarray(tilt > self._max_tilt_rad, dtype=np.bool_) + + +class penalty_feet_ori(_SensorTerm): + """Penalty for non-flat foot orientations (roll/pitch quaternion rows).""" + + def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv): + super().__init__(cfg, env) + self._quats = self._bind(_FOOT_QUAT_SENSORS) + if self._quats.dimensions != (4, 4): + raise ValueError( + f"{self.name} foot quaternion sensors must expose 4-D quats; received " + f"{self._quats.dimensions} on backend '{self._quats.backend_type}'" + ) + + def __call__(self, env: ManagerBasedRlEnv, **params: Any) -> np.ndarray: + del params + values = _state( + self.name, "foot quaternion", self._read(self._quats, self.name), (env.num_envs, 8) + ) + return np.asarray( + np.square(values[:, 1]) + + np.square(values[:, 2]) + + np.square(values[:, 5]) + + np.square(values[:, 6]), + dtype=get_global_dtype(), + ) + + +class penalty_close_feet_xy(_SensorTerm): + """Penalty for feet closer than ``threshold`` in the horizontal plane.""" + + _allowed_params = frozenset({"threshold"}) + + def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv): + super().__init__(cfg, env) + self._threshold = _real( + self.name, "threshold", cfg.params.get("threshold", 0.15), minimum=0.0 + ) + self._feet_pos = self._bind(_FOOT_POS_SENSORS) + if self._feet_pos.dimensions != (3, 3): + raise ValueError( + f"{self.name} foot position sensors must expose 3-D xyz; received " + f"{self._feet_pos.dimensions} on backend '{self._feet_pos.backend_type}'" + ) + + def __call__(self, env: ManagerBasedRlEnv, **params: Any) -> np.ndarray: + del params + values = _state( + self.name, "foot position", self._read(self._feet_pos, self.name), (env.num_envs, 6) + ) + feet_dist = np.linalg.norm(values[:, :2] - values[:, 3:5], axis=1) + return np.asarray( + np.where( + feet_dist < self._threshold, + np.square(feet_dist - self._threshold), + 0.0, + ), + dtype=get_global_dtype(), + ) + + +# --------------------------------------------------------------------------- +# Entity-based reward terms (legacy equations, no sensor binding) +# --------------------------------------------------------------------------- + + +def base_height( + env: ManagerBasedRlEnv, + target_height: float, + asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG, +) -> np.ndarray: + """Penalty for base height deviation from target (world frame).""" + target = _real("base_height", "target_height", target_height) + position = _state( + "base_height", + "root position", + _asset(env, asset_cfg).data.root_link_pos_w, + (env.num_envs, 3), + ) + return np.asarray(np.square(position[:, 2] - target), dtype=get_global_dtype()) + + +def weighted_pose( + env: ManagerBasedRlEnv, + pose_weights: Any, + asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG, +) -> np.ndarray: + """Weighted L2 penalty for joint position deviation from the default pose.""" + asset = _asset(env, asset_cfg) + position = asset.data.joint_pos[:, asset_cfg.joint_ids] + default = asset.data.default_joint_pos[:, asset_cfg.joint_ids] + weights = _weights("weighted_pose", "pose_weights", pose_weights) + if weights.shape[0] != position.shape[1]: + raise ValueError( + f"weighted_pose pose_weights length {weights.shape[0]} does not match " + f"joint count {position.shape[1]}" + ) + diff = _state("weighted_pose", "joint position", position, position.shape) - _state( + "weighted_pose", "default joint position", default, position.shape + ) + return np.asarray(np.sum(weights * np.square(diff), axis=1), dtype=get_global_dtype()) + + +def upper_body_pose( + env: ManagerBasedRlEnv, + pose_weights: Any, + asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG, +) -> np.ndarray: + """Weighted L2 pose penalty with the twelve leg joints zeroed out.""" + weights = _weights("upper_body_pose", "pose_weights", pose_weights) + if weights.shape[0] < 12: + raise ValueError("upper_body_pose pose_weights must cover at least the twelve leg joints") + weights = weights.copy() + weights[:12] = 0.0 + return weighted_pose(env, weights, asset_cfg=asset_cfg) + + +def alive(env: ManagerBasedRlEnv) -> np.ndarray: + """Constant reward for staying alive (unconditional, as in the legacy task).""" + return np.ones((env.num_envs,), dtype=get_global_dtype()) + + +# --------------------------------------------------------------------------- +# Velocity command term (legacy dead zone + standing zeroing semantics) +# --------------------------------------------------------------------------- + + +@dataclass(kw_only=True) +class G1VelocityCommandCfg(UniformVelocityCommandCfg): + """G1 velocity command with the legacy planar-norm dead zone.""" + + planar_dead_zone: float = 0.2 + + def build(self, env: ManagerBasedRlEnv) -> G1VelocityCommand: + return G1VelocityCommand(self, env) + + +class G1VelocityCommand(UniformVelocityCommand): + cfg: G1VelocityCommandCfg # pyright: ignore[reportIncompatibleVariableOverride] + + def __init__(self, cfg: G1VelocityCommandCfg, env: ManagerBasedRlEnv): + self._planar_dead_zone = _real( + "G1VelocityCommand", "planar_dead_zone", cfg.planar_dead_zone, minimum=0.0 + ) + if cfg.heading_command: + raise NotImplementedError( + "G1VelocityCommand capability 'heading command' is unavailable in the " + "Manager-Based runtime; the legacy G1 heading-channel zeroing has no " + "production owner and fails closed instead of falling back" + ) + super().__init__(cfg, env) + + def _resample_command(self, env_ids: np.ndarray) -> None: + super()._resample_command(env_ids) + planar = self.vel_command_b[env_ids, :2] + moving = np.linalg.norm(planar, axis=1) > self._planar_dead_zone + self.vel_command_b[env_ids, :2] = planar * moving[:, None] + + +# --------------------------------------------------------------------------- +# Penalty curriculum (ports EpisodeLengthTracker + PenaltyCurriculum semantics) +# --------------------------------------------------------------------------- + + +class G1PenaltyCurriculum(ManagerTermBase): + """Scale negative-weight reward terms by average episode length. + + Ports the legacy G1 penalty curriculum: penalty weights start at + ``initial_scale`` of their configured value and relax toward ``max_scale`` + as the tracked average episode length crosses the configured thresholds. + """ + + _allowed_params: ClassVar[frozenset[str]] = frozenset( + { + "initial_scale", + "min_scale", + "max_scale", + "level_down_threshold", + "level_up_threshold", + "degree", + } + ) + + def __init__(self, cfg: ManagerTermBaseCfg, env: _G1Env): + super().__init__(env) + unexpected = set(cfg.params) - self._allowed_params + if unexpected: + raise TypeError(f"{self.name} received unsupported parameters: {sorted(unexpected)}") + self._min_scale = _real( + self.name, "min_scale", cfg.params.get("min_scale", 0.5), minimum=0.0 + ) + self._max_scale = _real( + self.name, + "max_scale", + cfg.params.get("max_scale", 1.0), + minimum=self._min_scale, + ) + self._current_scale = _real( + self.name, + "initial_scale", + cfg.params.get("initial_scale", 0.5), + minimum=self._min_scale, + ) + if self._current_scale > self._max_scale: + raise ValueError(f"{self.name} initial_scale must be at most max_scale") + self._level_down_threshold = _real( + self.name, + "level_down_threshold", + cfg.params.get("level_down_threshold", 150.0), + minimum=0.0, + ) + self._level_up_threshold = _real( + self.name, + "level_up_threshold", + cfg.params.get("level_up_threshold", 750.0), + minimum=0.0, + ) + self._degree = _real(self.name, "degree", cfg.params.get("degree", 0.001), minimum=0.0) + self._tracker = EpisodeLengthTracker(env.num_envs) + self._original_weights: dict[str, float] = {} + for name in env.reward_manager.active_terms: + weight = float(env.reward_manager.get_term_cfg(name).weight) + if weight < 0.0: + self._original_weights[name] = weight + self._apply_scale() + + def _apply_scale(self) -> None: + env = cast("_G1Env", self._env) + for name, original in self._original_weights.items(): + env.reward_manager.get_term_cfg(name).weight = original * self._current_scale + + def __call__( + self, + env: _G1Env, + env_ids: np.ndarray | slice | None, + **params: Any, + ) -> dict[str, float]: + del params + ids = ( + np.arange(env.num_envs, dtype=np.intp) + if env_ids is None + else np.arange(env.num_envs, dtype=np.intp)[env_ids] + if isinstance(env_ids, slice) + else np.asarray(env_ids, dtype=np.intp).reshape(-1) + ) + done_ids = ids[env.reset_buf[ids]] + if len(done_ids) > 0: + self._tracker.update(env.episode_length_buf[done_ids].astype(np.float64)) + average = self._tracker.average_length + if average < self._level_down_threshold: + self._current_scale *= 1.0 - self._degree + elif average > self._level_up_threshold: + self._current_scale *= 1.0 + self._degree + self._current_scale = float( + np.clip(self._current_scale, self._min_scale, self._max_scale) + ) + self._apply_scale() + return { + "average_episode_length": float(self._tracker.average_length), + "penalty_scale": float(self._current_scale), + } + + +# --------------------------------------------------------------------------- +# G1 Manager-Based env: adds the public symmetry hook on the single lifecycle +# --------------------------------------------------------------------------- + +_SYMMETRY_TERM_SEGMENTS = { + "base_ang_vel": "gyro", + "projected_gravity": "gravity", + "joint_pos": "dof_pos", + "joint_vel": "dof_vel", + "actions": "actions", + "command": "command", + "gait_phase": "gait_phase", + "base_lin_vel": "linvel", +} + + +class G1WalkManagerBasedEnv(_ConcreteManagerBasedRlEnv): + """Manager-Based G1 walk runtime with the legacy symmetry-augmentation hook.""" + + def get_symmetry_obs_layouts(self) -> dict[str, SymmetryObsLayout]: + """Derive symmetry segment layouts from the observation manager terms.""" + mapping = {"obs": self._cfg.policy_observation_group} + if self._cfg.critic_observation_group is not None: + mapping["critic"] = self._cfg.critic_observation_group + layouts: dict[str, SymmetryObsLayout] = {} + for output_name, group_name in mapping.items(): + term_names = self.observation_manager.active_terms[group_name] + term_dims = self.observation_manager.group_obs_term_dim[group_name] + segments: list[tuple[str, int]] = [] + for term_name, dim in zip(term_names, term_dims, strict=True): + segment = _SYMMETRY_TERM_SEGMENTS.get(term_name) + if segment is None: + raise KeyError( + f"G1 symmetry layout has no segment mapping for observation term " + f"'{term_name}' in group '{group_name}'" + ) + segments.append((segment, int(np.prod(dim)))) + layouts[output_name] = tuple(segments) + return layouts + + def build_symmetry_augmentation(self, *, device: str) -> SymmetryAugmentation | None: + try: + actuator_names = self._backend.get_actuator_names() + except NotImplementedError: + return None + from unilab.tasks.locomotion.g1.symmetry import G1SymmetryAugmentation + + return G1SymmetryAugmentation( + actuator_names, + self.get_symmetry_obs_layouts(), + device=device, + ) + + +def make_g1_walk_env( + cfg: ManagerBasedRlEnvCfg, + num_envs: int = 1, + backend_type: str = "mujoco", +) -> G1WalkManagerBasedEnv: + """Construct the Registry-owned G1 Manager-Based production runtime.""" + if not isinstance(cfg, ManagerBasedRlEnvCfg): + raise TypeError( + f"make_g1_walk_env expected ManagerBasedRlEnvCfg, received {type(cfg).__name__}" + ) + if isinstance(num_envs, bool) or not isinstance(num_envs, int) or num_envs <= 0: + raise ValueError(f"make_g1_walk_env num_envs must be a positive integer, got {num_envs!r}") + if not isinstance(backend_type, str) or not backend_type: + raise ValueError( + f"make_g1_walk_env backend_type must be a non-empty string, got {backend_type!r}" + ) + + cfg.validate() + assert cfg.scene is not None + base_name, body_state_requested = _resolve_backend_entity_contract(cfg) + backend_kwargs = env_backend_kwargs(cfg) + backend_kwargs["base_name"] = base_name + + backend = create_backend( + backend_type, + cfg.scene, + num_envs, + cfg.sim_dt, + body_state_required=body_state_requested, + **backend_kwargs, + ) + try: + return G1WalkManagerBasedEnv(cfg, backend, num_envs) + except Exception: + backend.cleanup_scene_assets() + raise + + +__all__ = [ + "G1GaitPhase", + "G1PenaltyCurriculum", + "G1VelocityCommand", + "G1VelocityCommandCfg", + "G1WalkManagerBasedEnv", + "alive", + "ang_vel_xy", + "base_height", + "compute_feet_phase_contact_targets", + "compute_feet_phase_height_targets", + "feet_air_time", + "feet_double_stance", + "feet_phase", + "feet_phase_contact", + "feet_phase_contrast", + "forward_progress", + "g1_tilt_exceeded", + "lin_vel_z", + "make_g1_walk_env", + "orientation", + "penalty_close_feet_xy", + "penalty_feet_ori", + "track_ang_vel", + "track_lin_vel", + "under_speed", + "upper_body_pose", + "weighted_pose", +] diff --git a/src/unilab/tasks/migration_matrix.py b/src/unilab/tasks/migration_matrix.py index 350d3f946..ccde6acdb 100644 --- a/src/unilab/tasks/migration_matrix.py +++ b/src/unilab/tasks/migration_matrix.py @@ -122,10 +122,10 @@ def migration_record(task_name: str) -> TaskMigrationRecord: return TaskMigrationRecord( task_name, "g1_locomotion", - "Adapted", - "mba", - "The locomotion equations are reusable, but the 29/23-DoF sensor and gait surface is not yet manager-owned.", - "Migrate flat and rough variants together and delete the legacy owner.", + "Compatible", + "complete", + "Hydra owners materialize biped gait, sensor, command, and penalty-curriculum manager terms on the canonical runtime.", + "Keep the manager contract and regression evidence current.", ) if task_name in _CUSTOM_COMPAT_TASKS: family = "go2_arm" if task_name == "Go2ArmManipLoco" else "sharpa" diff --git a/src/unilab/tasks/motion_tracking/common/config.py b/src/unilab/tasks/motion_tracking/common/config.py index fb9d5419e..9f1aaee34 100644 --- a/src/unilab/tasks/motion_tracking/common/config.py +++ b/src/unilab/tasks/motion_tracking/common/config.py @@ -13,7 +13,7 @@ from unilab.assets import ASSETS_ROOT_PATH from unilab.base.scene import SceneCfg -from unilab.tasks.locomotion.g1.base import G1BaseCfg +from unilab.tasks.motion_tracking.common.g1_base import G1BaseCfg from .rewards import RewardConfig diff --git a/src/unilab/tasks/locomotion/g1/base.py b/src/unilab/tasks/motion_tracking/common/g1_base.py similarity index 86% rename from src/unilab/tasks/locomotion/g1/base.py rename to src/unilab/tasks/motion_tracking/common/g1_base.py index 9c25bacf3..f52afb0e0 100644 --- a/src/unilab/tasks/locomotion/g1/base.py +++ b/src/unilab/tasks/motion_tracking/common/g1_base.py @@ -1,4 +1,9 @@ -"""Shared runtime and configuration for G1 task owners.""" +"""Shared G1 runtime and configuration base for motion-tracking tasks. + +Relocated from ``unilab.tasks.locomotion.g1.base`` when the G1 locomotion +family migrated to the Manager-Based runtime; motion tracking is the sole +remaining consumer of the legacy class-based G1 base env. +""" from __future__ import annotations diff --git a/src/unilab/tasks/motion_tracking/common/tracking.py b/src/unilab/tasks/motion_tracking/common/tracking.py index 7da3516a6..99be8a52c 100644 --- a/src/unilab/tasks/motion_tracking/common/tracking.py +++ b/src/unilab/tasks/motion_tracking/common/tracking.py @@ -1,7 +1,7 @@ """Robot-agnostic motion-tracking engine. Holds :class:`MotionTrackingEnv` (the imitation engine, inheriting the shared -``G1BaseEnv`` locomotion base) and :class:`MotionTrackingDeployEnv` (the +``G1BaseEnv`` locomotion base (``motion_tracking.common.g1_base``)) and :class:`MotionTrackingDeployEnv` (the unitree_rl_lab mimic actor variant). Per-concern math lives in the owner modules (``rewards`` / ``observations`` / ``terminations`` / ``transforms`` / ``reset`` / ``domain_randomization``); the engine keeps only the @@ -17,7 +17,7 @@ from unilab.base.backend import create_backend, env_backend_kwargs from unilab.base.np_env import NpEnvState from unilab.dtype_config import get_global_dtype -from unilab.tasks.locomotion.g1.base import G1BaseEnv +from unilab.tasks.motion_tracking.common.g1_base import G1BaseEnv from . import observations from .config import MotionTrackingCfg, MotionTrackingDeployEnvCfg diff --git a/src/unilab/tasks/motion_tracking/g1/tracking_obs.py b/src/unilab/tasks/motion_tracking/g1/tracking_obs.py index ea194f3c0..d6c53abd9 100644 --- a/src/unilab/tasks/motion_tracking/g1/tracking_obs.py +++ b/src/unilab/tasks/motion_tracking/g1/tracking_obs.py @@ -43,7 +43,7 @@ ) from unilab.dr.types import RESET_TERM_GEOM_FRICTION from unilab.dtype_config import get_global_dtype -from unilab.tasks.locomotion.g1.base import NoiseConfig +from unilab.tasks.motion_tracking.common.g1_base import NoiseConfig from ..common.rewards import RewardContext from .tracking import ( diff --git a/src/unilab/tasks/motion_tracking/x2/flip_tracking.py b/src/unilab/tasks/motion_tracking/x2/flip_tracking.py index 8c0a5340f..a72b9e0fc 100644 --- a/src/unilab/tasks/motion_tracking/x2/flip_tracking.py +++ b/src/unilab/tasks/motion_tracking/x2/flip_tracking.py @@ -8,7 +8,7 @@ from unilab.assets import ASSETS_ROOT_PATH from unilab.base import registry from unilab.base.scene import SceneCfg -from unilab.tasks.locomotion.g1.base import Sensor +from unilab.tasks.motion_tracking.common.g1_base import Sensor from ..common.config import ( PoseRandomization, diff --git a/tests/algos/test_rsl_rl_runner.py b/tests/algos/test_rsl_rl_runner.py index 3585d577f..cfd834250 100644 --- a/tests/algos/test_rsl_rl_runner.py +++ b/tests/algos/test_rsl_rl_runner.py @@ -118,33 +118,23 @@ def get_privileged_observations(self): ) def test_rsl_rl_ppo_one_iteration( env_name: str, - default_g1_reward_config, ): """RSL-RL PPO can complete 1 training iteration on a real env.""" from rsl_rl.runners import OnPolicyRunner + num_envs = 256 + root_dir = Path(__file__).parents[2] if "Go2" in env_name: - num_envs = 256 - root_dir = Path(__file__).parents[2] - GlobalHydra.instance().clear() - with initialize_config_dir(config_dir=str(root_dir / "conf" / "ppo"), version_base="1.3"): - hydra_cfg = compose("config", overrides=["task=go2_joystick_flat/mujoco"]) - env_cfg_override = BackendAdapter( - hydra_cfg, root_dir=root_dir - ).build_task_env_cfg_override() + task = "go2_joystick_flat/mujoco" elif "G1" in env_name: - reward_cfg = default_g1_reward_config - num_envs = 256 - env_cfg_override = {"reward_config": reward_cfg} + task = "g1_walk_flat/mujoco" else: num_envs = 128 - root_dir = Path(__file__).parents[2] - GlobalHydra.instance().clear() - with initialize_config_dir(config_dir=str(root_dir / "conf" / "ppo"), version_base="1.3"): - hydra_cfg = compose("config", overrides=["task=allegro_inhand/mujoco"]) - env_cfg_override = BackendAdapter( - hydra_cfg, root_dir=root_dir - ).build_task_env_cfg_override() + task = "allegro_inhand/mujoco" + GlobalHydra.instance().clear() + with initialize_config_dir(config_dir=str(root_dir / "conf" / "ppo"), version_base="1.3"): + hydra_cfg = compose("config", overrides=[f"task={task}"]) + env_cfg_override = BackendAdapter(hydra_cfg, root_dir=root_dir).build_task_env_cfg_override() env = registry.make( env_name, diff --git a/tests/base/test_reward_override.py b/tests/base/test_reward_override.py index 3eca81e67..8bbce5283 100644 --- a/tests/base/test_reward_override.py +++ b/tests/base/test_reward_override.py @@ -1,29 +1,32 @@ """Test reward config override through registry.""" +from pathlib import Path from typing import Any, cast +from hydra import compose, initialize_config_dir +from hydra.core.global_hydra import GlobalHydra + from unilab.base import registry from unilab.base.registry import ensure_registries +from unilab.training.backend_adapter import BackendAdapter + +ROOT_DIR = Path(__file__).parents[2] def test_reward_override_g1(): - """Test G1 reward config override.""" + """Test G1 manager reward override through the registry.""" ensure_registries() - from unilab.tasks.locomotion.g1.joystick import G1WalkRewardConfig - - override_config = G1WalkRewardConfig( - scales={"tracking_lin_vel": 888.0, "alive": 20.0}, - tracking_sigma=0.3, - base_height_target=0.8, - min_base_height=0.3, - max_tilt_deg=65.0, - gait_frequency=1.5, - feet_phase_swing_height=0.09, - feet_phase_tracking_sigma=0.008, - close_feet_threshold=0.15, - pose_weights=[0.01] * 29, - ) + GlobalHydra.instance().clear() + with initialize_config_dir(config_dir=str(ROOT_DIR / "conf" / "ppo"), version_base="1.3"): + cfg = compose("config", overrides=["task=g1_walk_flat/mujoco"]) + env_cfg_override = BackendAdapter(cfg, root_dir=ROOT_DIR).build_task_env_cfg_override() + env_cfg_override["rewards"]["tracking_lin_vel"]["weight"] = 888.0 + env_cfg_override["rewards"]["alive"] = { + "_target_": "unilab.managers.RewardTermCfg", + "func": "unilab.tasks.locomotion.g1.manager_terms.alive", + "weight": 20.0, + } env = cast( Any, @@ -31,10 +34,11 @@ def test_reward_override_g1(): "G1WalkFlat", num_envs=1, sim_backend="mujoco", - env_cfg_override={"reward_config": override_config}, + env_cfg_override=env_cfg_override, ), ) - assert env._cfg.reward_config.scales["tracking_lin_vel"] == 888.0 - assert env._cfg.reward_config.scales["alive"] == 20.0 + assert env._cfg.rewards["tracking_lin_vel"].weight == 888.0 + assert env._cfg.rewards["alive"].weight == 20.0 + assert env.reward_manager.get_term_cfg("tracking_lin_vel").weight == 888.0 env.close() diff --git a/tests/benchmark/test_offpolicy_collector_active_benchmark.py b/tests/benchmark/test_offpolicy_collector_active_benchmark.py index d88e31c4c..25665b2c8 100644 --- a/tests/benchmark/test_offpolicy_collector_active_benchmark.py +++ b/tests/benchmark/test_offpolicy_collector_active_benchmark.py @@ -135,14 +135,9 @@ def test_auto_discovery_supports_motrixsim_alias() -> None: def test_noise_seed_override_composes_for_target_g1_profiles() -> None: - for spec in ( - ("sac", "g1_motion_tracking", "mujoco"), - ("sac", "g1_walk_flat", "mujoco"), - ("sac", "g1_walk_flat", "motrix"), - ("flashsac", "g1_walk_flat", "mujoco"), - ("flashsac", "g1_walk_flat", "motrix"), - ("td3", "g1_walk_flat", "mujoco"), - ): + # Manager-Based G1 walk owners no longer carry env.noise_config; only the + # legacy motion-tracking owner still composes the seed override. + for spec in (("sac", "g1_motion_tracking", "mujoco"),): cfg = bench._compose_offpolicy_cfg( *spec, extra_overrides=["env.noise_config.seed=123"], diff --git a/tests/config/test_config_system.py b/tests/config/test_config_system.py index 32baef7fd..710de281e 100644 --- a/tests/config/test_config_system.py +++ b/tests/config/test_config_system.py @@ -215,9 +215,8 @@ def test_offpolicy_g1_walk_flat_motrix_sac_preserves_backend_overrides(): assert cfg.algo.num_envs == 2048 assert cfg.algo.max_iterations == 5000 - assert cfg.reward.scales.tracking_lin_vel == pytest.approx(2.2) - assert cfg.env.domain_rand.randomize_kp is False - assert cfg.env.domain_rand.randomize_kd is False + assert cfg.reward.tracking_lin_vel.weight == pytest.approx(2.2) + assert cfg.env.events.pd_gains is None def test_offpolicy_g1_walk_flat_mujoco_td3_uses_td3_task_owner(): @@ -229,8 +228,8 @@ def test_offpolicy_g1_walk_flat_mujoco_td3_uses_td3_task_owner(): assert cfg.algo.tau == pytest.approx(0.1) assert cfg.algo.actor_hidden_dim == 512 assert cfg.algo.critic_hidden_dim == 1024 - assert cfg.reward.scales.tracking_lin_vel == pytest.approx(2.0) - assert cfg.env.control_config.action_scale == pytest.approx(1.0) + assert cfg.reward.tracking_lin_vel.weight == pytest.approx(2.0) + assert cfg.env.actions.joint_pos.scale == pytest.approx(1.0) def test_offpolicy_td3_go2_joystick_flat_motrix_composes(): @@ -280,14 +279,13 @@ def test_offpolicy_g1_walk_flat_mjwarp_owner_preserves_sac_contract(): assert mjwarp_cfg.training.play_render_mode == "record" assert mjwarp_cfg.algo.num_envs == mujoco_cfg.algo.num_envs assert mjwarp_cfg.algo.use_symmetry is mujoco_cfg.algo.use_symmetry is True - assert mjwarp_cfg.env.control_config.action_scale == pytest.approx( - mujoco_cfg.env.control_config.action_scale + assert mjwarp_cfg.env.actions.joint_pos.scale == pytest.approx( + mujoco_cfg.env.actions.joint_pos.scale ) assert mjwarp_cfg.env.mjwarp_nconmax == 128 assert mjwarp_cfg.env.mjwarp_njmax == 256 assert mjwarp_cfg.env.render_spacing == pytest.approx(2.0) - assert mjwarp_cfg.env.domain_rand.randomize_kp is False - assert mjwarp_cfg.env.domain_rand.randomize_kd is False + assert mjwarp_cfg.env.events.pd_gains is None assert OmegaConf.to_container(mjwarp_cfg.reward, resolve=True) == OmegaConf.to_container( mujoco_cfg.reward, resolve=True ) @@ -312,24 +310,24 @@ def test_ppo_g1_backend_specific_hyperparams_remain_separate(): assert motrix_cfg.algo.empirical_normalization is True assert motrix_cfg.algo.obs_groups.actor == ["policy"] assert OmegaConf.select(motrix_cfg, "env.motrix_max_iterations") is None - assert motrix_cfg.env.control_config.action_scale == pytest.approx(0.5) - assert motrix_cfg.env.commands.vel_limit == [[0.4, 0.0, 0.0], [0.7, 0.0, 0.0]] - assert motrix_cfg.env.gait_phase_init_mode == "offset_phase" - assert motrix_cfg.reward.scales.tracking_lin_vel == pytest.approx(2.0) - assert motrix_cfg.reward.scales.tracking_ang_vel == pytest.approx(0.25) - assert motrix_cfg.reward.scales.forward_progress == pytest.approx(0.0) - assert motrix_cfg.reward.scales.under_speed == pytest.approx(-0.2) - assert motrix_cfg.reward.scales.penalty_feet_ori == pytest.approx(0.0) - assert motrix_cfg.reward.scales.feet_phase == pytest.approx(1.2) - assert motrix_cfg.reward.scales.feet_phase_contrast == pytest.approx(1.5) - assert motrix_cfg.reward.scales.feet_phase_contact == pytest.approx(1.0) - assert motrix_cfg.reward.scales.feet_double_stance == pytest.approx(-1.0) - assert motrix_cfg.reward.scales.base_height == pytest.approx(-120.0) - assert motrix_cfg.reward.scales.pose == pytest.approx(-0.05) - assert motrix_cfg.reward.base_height_target == pytest.approx(0.765) - assert motrix_cfg.reward.min_forward_speed_for_gait_reward == pytest.approx(0.05) - assert motrix_cfg.reward.min_base_height == pytest.approx(0.5) - assert motrix_cfg.reward.max_tilt_deg == pytest.approx(35.0) + assert motrix_cfg.env.actions.joint_pos.scale == pytest.approx(0.5) + assert motrix_cfg.env.commands.twist.ranges.lin_vel_x == [0.4, 0.7] + assert motrix_cfg.env.observations.policy.terms.gait_phase.params.init_mode == "offset_phase" + assert motrix_cfg.reward.tracking_lin_vel.weight == pytest.approx(2.0) + assert motrix_cfg.reward.tracking_ang_vel.weight == pytest.approx(0.25) + assert motrix_cfg.reward.forward_progress.weight == pytest.approx(0.0) + assert motrix_cfg.reward.under_speed.weight == pytest.approx(-0.2) + assert motrix_cfg.reward.penalty_feet_ori.weight == pytest.approx(0.0) + assert motrix_cfg.reward.feet_phase.weight == pytest.approx(1.2) + assert motrix_cfg.reward.feet_phase_contrast.weight == pytest.approx(1.5) + assert motrix_cfg.reward.feet_phase_contact.weight == pytest.approx(1.0) + assert motrix_cfg.reward.feet_double_stance.weight == pytest.approx(-1.0) + assert motrix_cfg.reward.base_height.weight == pytest.approx(-120.0) + assert motrix_cfg.reward.pose.weight == pytest.approx(-0.05) + assert motrix_cfg.reward.base_height.params.target_height == pytest.approx(0.765) + assert motrix_cfg.reward.feet_phase.params.min_forward_speed == pytest.approx(0.05) + assert motrix_cfg.env.terminations.base_height.params.minimum_height == pytest.approx(0.5) + assert motrix_cfg.env.terminations.tilt.params.max_tilt_deg == pytest.approx(35.0) @pytest.mark.parametrize( @@ -486,8 +484,8 @@ def test_offpolicy_g1_walk_flat_motrix_preserves_backend_env_overrides(): assert cfg.training.sim_backend == "motrix" assert cfg.algo.num_envs == 2048 assert cfg.algo.max_iterations == 5000 - assert cfg.env.domain_rand.randomize_kp is False - assert cfg.env.domain_rand.randomize_kd is False + assert cfg.env.events.pd_gains is None + assert cfg.reward.tracking_lin_vel.weight == pytest.approx(2.2) def test_offpolicy_flashsac_go2_joystick_mujoco_enables_full_dr_stack(): diff --git a/tests/config/test_locomotion_params.py b/tests/config/test_locomotion_params.py index b9c9ff793..14e2e4f25 100644 --- a/tests/config/test_locomotion_params.py +++ b/tests/config/test_locomotion_params.py @@ -3,7 +3,6 @@ from __future__ import annotations from pathlib import Path -from typing import Any, cast import pytest @@ -151,9 +150,9 @@ def test_offpolicy_sac_g1_task_overrides(): assert cfg.algo.algo_params.use_compile is True assert cfg.training.task_name == "G1WalkFlat" - assert cfg.env.control_config.action_scale == pytest.approx(1.0) - assert cfg.env.gait_phase_init_mode == "offset_phase" - assert cfg.env.reset_base_qvel_limit == pytest.approx(0.5) + assert cfg.env.actions.joint_pos.scale == pytest.approx(1.0) + assert cfg.env.observations.policy.terms.gait_phase.params.init_mode == "offset_phase" + assert cfg.env.events.reset_root_state_uniform.params.velocity_range.x == [-0.5, 0.5] def test_offpolicy_td3_defaults(): @@ -181,7 +180,7 @@ def test_offpolicy_td3_g1_task_overrides(): cfg = compose("config", overrides=["algo=td3", "task=td3/g1_walk_flat/mujoco"]) assert cfg.training.task_name == "G1WalkFlat" assert cfg.algo.max_iterations == 100000 - assert cfg.env.control_config.action_scale == pytest.approx(1.0) + assert cfg.env.actions.joint_pos.scale == pytest.approx(1.0) def test_offpolicy_flashsac_g1_task_overrides(): @@ -228,8 +227,6 @@ def test_offpolicy_g1_rough_terrain_task_overrides(): from hydra import compose, initialize_config_dir from hydra.core.global_hydra import GlobalHydra - from unilab.tasks.locomotion.g1.joystick import G1WalkRoughCfg - GlobalHydra.instance().clear() with initialize_config_dir(config_dir=str(CONF_DIR / "offpolicy"), version_base="1.3"): cfg = compose( @@ -239,23 +236,24 @@ def test_offpolicy_g1_rough_terrain_task_overrides(): assert cfg.algo.algo == "sac" assert cfg.training.task_name == "G1WalkRough" assert cfg.training.sim_backend == "mujoco" - assert G1WalkRoughCfg().scene.model_file.endswith("scene_rough.xml") + assert cfg.env.scene.model_file.endswith("scene_rough.xml") def test_g1_task_owner_yamls_preserve_legacy_and_walk_observation_profiles(): from hydra import compose, initialize_config_dir from hydra.core.global_hydra import GlobalHydra - from unilab.tasks.locomotion.g1.joystick import G1WalkEnv - def uses_walk_profile(config_group: str, overrides: list[str]) -> bool: GlobalHydra.instance().clear() with initialize_config_dir(config_dir=str(CONF_DIR / config_group), version_base="1.3"): cfg = compose("config", overrides=overrides) - env = cast(Any, object.__new__(G1WalkEnv)) - env._cfg = cfg.env - env._reward_cfg = cfg.reward - return bool(env._uses_walk_observation_profile()) + gyro_scale = cfg.env.observations.policy.terms.base_ang_vel.get("scale") + if gyro_scale is None: + return False + assert gyro_scale == pytest.approx(0.25) + assert cfg.env.observations.policy.terms.joint_vel.scale == pytest.approx(0.05) + assert cfg.env.observations.critic.terms.base_lin_vel.scale == pytest.approx(2.0) + return True assert uses_walk_profile("ppo", ["task=g1_walk_flat/mujoco"]) is False assert uses_walk_profile("appo", ["task=g1_walk_flat/mujoco"]) is False @@ -297,7 +295,7 @@ def test_appo_g1_task_overrides(): assert cfg.algo.save_interval == 100 assert cfg.training.task_name == "G1WalkFlat" assert "obs_profile" not in cfg.env - assert cfg.env.curriculum.enabled is False + assert "curriculum" not in cfg.env # --------------------------------------------------------------------------- @@ -344,7 +342,7 @@ def test_ppo_g1_num_envs(): assert cfg.algo.max_iterations == 2200 assert cfg.training.task_name == "G1WalkFlat" assert "obs_profile" not in cfg.env - assert cfg.env.curriculum.enabled is False + assert "curriculum" not in cfg.env def test_ppo_go2_num_envs(): diff --git a/tests/config/test_reward_injection.py b/tests/config/test_reward_injection.py index 0c9e5e8fc..80f8e132a 100644 --- a/tests/config/test_reward_injection.py +++ b/tests/config/test_reward_injection.py @@ -12,9 +12,9 @@ def test_reward_config_loading_g1(): with initialize(config_path="../../conf/offpolicy", version_base="1.3"): cfg = compose(config_name="config", overrides=["task=sac/g1_walk_flat/mujoco"]) assert hasattr(cfg, "reward") - assert cfg.reward.scales.tracking_lin_vel == 2.0 - assert cfg.reward.scales.alive == 10.0 - assert cfg.reward.base_height_target == 0.754 + assert cfg.reward.tracking_lin_vel.weight == 2.0 + assert cfg.reward.alive.weight == 10.0 + assert cfg.reward.feet_phase.params.swing_height == 0.09 def test_reward_config_loading_g1_motrix(): @@ -22,8 +22,8 @@ def test_reward_config_loading_g1_motrix(): with initialize(config_path="../../conf/offpolicy", version_base="1.3"): cfg = compose(config_name="config", overrides=["task=sac/g1_walk_flat/motrix"]) assert hasattr(cfg, "reward") - assert cfg.reward.scales.tracking_lin_vel == 2.2 - assert cfg.reward.scales.alive == 12.0 + assert cfg.reward.tracking_lin_vel.weight == 2.2 + assert cfg.reward.alive.weight == 12.0 def test_resolve_reward_dict_reads_task_reward(): @@ -44,34 +44,35 @@ def test_resolve_reward_dict_reads_task_reward(): def test_reward_config_conversion(): - """Test reward config converts to dataclasses via registry.""" + """Test reward config materializes into manager reward terms via registry.""" from unilab.base import registry + from unilab.base.config_materialization import apply_cfg_overrides from unilab.base.registry import ensure_registries + from unilab.envs import ManagerBasedRlEnvCfg ensure_registries() - # Test G1 walk config - registry auto-converts dict to G1WalkRewardConfig - g1_dict = { - "scales": {"tracking_lin_vel": 2.0, "alive": 10.0}, - "tracking_sigma": 0.25, - "base_height_target": 0.754, - "gait_frequency": 1.5, - "feet_phase_swing_height": 0.09, - "feet_phase_tracking_sigma": 0.008, - "min_base_height": 0.3, - "max_tilt_deg": 65.0, - "close_feet_threshold": 0.15, - "pose_weights": [0.01] * 29, - } - env = cast( - Any, - registry.make( - "G1WalkFlat", - num_envs=1, - sim_backend="mujoco", - env_cfg_override={"reward_config": g1_dict}, - ), + env_cfg = registry.materialize_env_config("G1WalkFlat") + assert isinstance(env_cfg, ManagerBasedRlEnvCfg) + apply_cfg_overrides( + env_cfg, + { + "rewards": { + "tracking_lin_vel": { + "_target_": "unilab.managers.RewardTermCfg", + "func": "unilab.tasks.locomotion.g1.manager_terms.track_lin_vel", + "weight": 2.0, + "params": {"tracking_sigma": 0.25, "command_name": "twist"}, + }, + "alive": { + "_target_": "unilab.managers.RewardTermCfg", + "func": "unilab.tasks.locomotion.g1.manager_terms.alive", + "weight": 10.0, + }, + } + }, ) - assert hasattr(env._cfg.reward_config, "scales") - assert env._cfg.reward_config.scales["tracking_lin_vel"] == 2.0 - env.close() + assert env_cfg.rewards["tracking_lin_vel"].weight == 2.0 + assert env_cfg.rewards["tracking_lin_vel"].params["tracking_sigma"] == 0.25 + assert env_cfg.rewards["alive"].weight == 10.0 + assert OmegaConf.is_config(env_cfg.rewards["tracking_lin_vel"].func) is False diff --git a/tests/conftest.py b/tests/conftest.py index 7e5dc819e..f2a9abda2 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -163,124 +163,3 @@ def default_go2_reward_config(): "tracking_sigma": 0.25, "base_height_target": 0.3, } - - -@pytest.fixture -def default_g1_reward_config(): - """Default reward config for G1 testing.""" - return { - "scales": { - "tracking_lin_vel": 2.0, - "tracking_ang_vel": 0.25, - "forward_progress": 0.0, - "under_speed": -0.2, - "upper_body_pose": -0.05, - "penalty_feet_ori": 0.0, - "feet_phase": 1.0, - "feet_phase_contrast": 1.0, - "feet_phase_contact": 0.5, - "feet_double_stance": -0.5, - "lin_vel_z": -1.0, - "ang_vel_xy": -0.2, - "base_height": -120.0, - "orientation": -2.5, - "action_rate": -0.005, - "pose": -0.05, - }, - "tracking_sigma": 0.25, - "gait_frequency": 1.5, - "feet_phase_swing_height": 0.09, - "feet_phase_tracking_sigma": 0.008, - "base_height_target": 0.765, - "min_forward_speed_for_gait_reward": 0.05, - "min_base_height": 0.5, - "max_tilt_deg": 35.0, - "pose_weights": [ - 0.01, - 1.0, - 5.0, - 0.01, - 5.0, - 5.0, - 0.01, - 1.0, - 5.0, - 0.01, - 5.0, - 5.0, - 50.0, - 50.0, - 50.0, - 50.0, - 50.0, - 50.0, - 50.0, - 50.0, - 50.0, - 50.0, - 50.0, - 50.0, - 50.0, - 50.0, - 50.0, - 50.0, - 50.0, - ], - } - - -@pytest.fixture -def default_g1_walk_flat_reward_config(): - """Default reward config for G1 SAC testing.""" - return { - "scales": { - "tracking_lin_vel": 2.0, - "tracking_ang_vel": 1.5, - "penalty_ang_vel_xy": -1.0, - "penalty_orientation": -10.0, - "penalty_action_rate": -2.0, - "pose": -0.5, - "penalty_feet_ori": -25.0, - "feet_phase": 5.0, - "alive": 10.0, - }, - "tracking_sigma": 0.25, - "base_height_target": 0.754, - "min_base_height": 0.3, - "max_tilt_deg": 65.0, - "gait_frequency": 1.5, - "feet_phase_swing_height": 0.09, - "feet_phase_tracking_sigma": 0.008, - "close_feet_threshold": 0.15, - "pose_weights": [ - 0.01, - 1.0, - 5.0, - 0.01, - 5.0, - 5.0, - 0.01, - 1.0, - 5.0, - 0.01, - 5.0, - 5.0, - 50.0, - 50.0, - 50.0, - 50.0, - 50.0, - 50.0, - 50.0, - 50.0, - 50.0, - 50.0, - 50.0, - 50.0, - 50.0, - 50.0, - 50.0, - 50.0, - 50.0, - ], - } diff --git a/tests/envs/locomotion/g1/test_g1_owner_contract.py b/tests/envs/locomotion/g1/test_g1_owner_contract.py index 2b22c9d97..2956d198c 100644 --- a/tests/envs/locomotion/g1/test_g1_owner_contract.py +++ b/tests/envs/locomotion/g1/test_g1_owner_contract.py @@ -1,266 +1,678 @@ +"""Hydra-owned production contract for the G1 walk Manager-Based tasks.""" + from __future__ import annotations +from collections.abc import Mapping, Sequence +from dataclasses import fields, is_dataclass from pathlib import Path -from typing import Any, cast +from typing import Any import numpy as np import pytest from hydra import compose, initialize_config_dir from hydra.core.global_hydra import GlobalHydra -from omegaconf import OmegaConf +from omegaconf import DictConfig, OmegaConf from unilab.base import registry -from unilab.base.registry import ensure_registries +from unilab.base.config_materialization import apply_cfg_overrides +from unilab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg, mdp +from unilab.tasks.locomotion.g1 import manager_terms as g1_terms from unilab.training.backend_adapter import BackendAdapter ROOT_DIR = Path(__file__).parents[4] CONF_DIR = ROOT_DIR / "conf" - -_G1_OWNER_CASES = [ - { - "id": "ppo_mujoco", - "config_group": "ppo", - "overrides": ["task=g1_walk_flat/mujoco"], - "task_name": "G1WalkFlat", - "backend": "mujoco", - "profile": "legacy", - "action_scale": 0.25, - "curriculum_enabled": False, - }, - { - "id": "ppo_motrix", - "config_group": "ppo", - "overrides": ["task=g1_walk_flat/motrix"], - "task_name": "G1WalkFlat", - "backend": "motrix", - "profile": "legacy", - "action_scale": 0.5, - "curriculum_enabled": False, - }, - { - "id": "appo_mujoco", - "config_group": "appo", - "overrides": ["task=g1_walk_flat/mujoco"], - "task_name": "G1WalkFlat", - "backend": "mujoco", - "profile": "legacy", - "action_scale": 0.25, - "curriculum_enabled": False, - }, - { - "id": "sac_mujoco", - "config_group": "offpolicy", - "overrides": ["algo=sac", "task=sac/g1_walk_flat/mujoco"], - "task_name": "G1WalkFlat", - "backend": "mujoco", - "profile": "walk", - "action_scale": 1.0, - "curriculum_enabled": True, - }, - { - "id": "sac_motrix", - "config_group": "offpolicy", - "overrides": ["algo=sac", "task=sac/g1_walk_flat/motrix"], - "task_name": "G1WalkFlat", - "backend": "motrix", - "profile": "walk", - "action_scale": 1.0, - "curriculum_enabled": True, - }, - { - "id": "ppo_mjwarp", - "config_group": "ppo", - "overrides": ["task=g1_walk_flat/mjwarp"], - "task_name": "G1WalkFlat", - "backend": "mjwarp", - "profile": "legacy", - "action_scale": 0.25, - "curriculum_enabled": False, - }, - { - "id": "sac_mjwarp", - "config_group": "offpolicy", - "overrides": ["algo=sac", "task=sac/g1_walk_flat/mjwarp"], - "task_name": "G1WalkFlat", - "backend": "mjwarp", - "profile": "walk", - "action_scale": 1.0, - "curriculum_enabled": True, - }, - { - "id": "sac_rough", - "config_group": "offpolicy", - "overrides": ["algo=sac", "task=sac/g1_walk_rough/mujoco"], - "task_name": "G1WalkRough", - "backend": "mujoco", - "profile": "walk", - "action_scale": 1.0, - "curriculum_enabled": True, - "model_suffix": "scene_rough.xml", - }, - { - "id": "td3_mujoco", - "config_group": "offpolicy", - "overrides": ["algo=td3", "task=td3/g1_walk_flat/mujoco"], - "task_name": "G1WalkFlat", - "backend": "mujoco", - "profile": "walk", - "action_scale": 1.0, - "curriculum_enabled": True, - }, - { - "id": "flashsac_walk_mujoco", - "config_group": "offpolicy", - "overrides": ["algo=flashsac", "task=flashsac/g1_walk_flat/mujoco"], - "task_name": "G1WalkFlat", - "backend": "mujoco", - "profile": "walk", - "action_scale": 1.0, - "curriculum_enabled": True, - }, -] - - -def _compose_cfg(config_group: str, overrides: list[str]): +_RESET_EVENTS = ("reset_scene_to_default", "reset_root_state_uniform") +_PPO_REWARDS = ( + "tracking_lin_vel", + "tracking_ang_vel", + "feet_phase", + "lin_vel_z", + "ang_vel_xy", + "base_height", + "orientation", + "action_rate", + "pose", +) +_MOTRIX_EXTRA_REWARDS = ( + "forward_progress", + "under_speed", + "upper_body_pose", + "penalty_feet_ori", + "feet_phase_contrast", + "feet_phase_contact", + "feet_double_stance", +) +_OFFPOLICY_REWARDS = ( + "tracking_lin_vel", + "tracking_ang_vel", + "penalty_ang_vel_xy", + "penalty_orientation", + "penalty_action_rate", + "pose", + "penalty_feet_ori", + "feet_phase", + "alive", +) + +_OBSERVATION_TERMS = ( + "base_ang_vel", + "projected_gravity", + "joint_pos", + "joint_vel", + "actions", + "command", + "gait_phase", +) + +_POSE_WEIGHTS_29 = [0.01, 1.0, 5.0, 0.01, 5.0, 5.0] * 2 + [50.0] * 17 +_POSE_WEIGHTS_23 = [0.01, 1.0, 5.0, 0.01, 5.0, 5.0] * 2 + [50.0] * 11 + +_OWNER_CASES = ( + pytest.param( + "ppo", + ("task=g1_walk_flat/mujoco",), + "G1WalkFlat", + "mujoco", + 29, + 0.25, + "scene_flat.xml", + _PPO_REWARDS, + (*_RESET_EVENTS, "pd_gains"), + False, + id="ppo-mujoco", + ), + pytest.param( + "ppo", + ("task=g1_walk_flat/motrix",), + "G1WalkFlat", + "motrix", + 29, + 0.5, + "scene_flat.xml", + (*_PPO_REWARDS, *_MOTRIX_EXTRA_REWARDS), + _RESET_EVENTS, + False, + id="ppo-motrix", + ), + pytest.param( + "ppo", + ("task=g1_walk_flat/mjwarp",), + "G1WalkFlat", + "mjwarp", + 29, + 0.25, + "scene_flat.xml", + _PPO_REWARDS, + _RESET_EVENTS, + False, + id="ppo-mjwarp", + ), + pytest.param( + "ppo", + ("task=g1_23dof_walk_flat/mujoco",), + "G1Walk23DofFlat", + "mujoco", + 23, + 0.25, + "scene_flat_23dof.xml", + _PPO_REWARDS, + (*_RESET_EVENTS, "pd_gains"), + False, + id="ppo-23dof-mujoco", + ), + pytest.param( + "ppo", + ("task=g1_23dof_walk_flat/motrix",), + "G1Walk23DofFlat", + "motrix", + 23, + 0.5, + "scene_flat_23dof.xml", + (*_PPO_REWARDS, *_MOTRIX_EXTRA_REWARDS), + _RESET_EVENTS, + False, + id="ppo-23dof-motrix", + ), + pytest.param( + "ppo", + ("task=g1_23dof_walk_rough/mujoco",), + "G1Walk23DofRough", + "mujoco", + 23, + 0.25, + "scene_rough_23dof.xml", + _PPO_REWARDS, + (*_RESET_EVENTS, "pd_gains"), + True, + id="ppo-23dof-rough-mujoco", + ), + pytest.param( + "appo", + ("task=g1_walk_flat/mujoco",), + "G1WalkFlat", + "mujoco", + 29, + 0.25, + "scene_flat.xml", + _PPO_REWARDS, + (*_RESET_EVENTS, "pd_gains"), + False, + id="appo-mujoco", + ), + pytest.param( + "appo", + ("task=g1_23dof_walk_flat/mujoco",), + "G1Walk23DofFlat", + "mujoco", + 23, + 0.25, + "scene_flat_23dof.xml", + _PPO_REWARDS, + (*_RESET_EVENTS, "pd_gains"), + False, + id="appo-23dof-mujoco", + ), + pytest.param( + "offpolicy", + ("algo=sac", "task=sac/g1_walk_flat/mujoco"), + "G1WalkFlat", + "mujoco", + 29, + 1.0, + "scene_flat.xml", + _OFFPOLICY_REWARDS, + (*_RESET_EVENTS, "pd_gains"), + True, + id="sac-mujoco", + ), + pytest.param( + "offpolicy", + ("algo=sac", "task=sac/g1_walk_flat/motrix"), + "G1WalkFlat", + "motrix", + 29, + 1.0, + "scene_flat.xml", + _OFFPOLICY_REWARDS, + _RESET_EVENTS, + True, + id="sac-motrix", + ), + pytest.param( + "offpolicy", + ("algo=sac", "task=sac/g1_walk_flat/mjwarp"), + "G1WalkFlat", + "mjwarp", + 29, + 1.0, + "scene_flat.xml", + _OFFPOLICY_REWARDS, + _RESET_EVENTS, + True, + id="sac-mjwarp", + ), + pytest.param( + "offpolicy", + ("algo=sac", "task=sac/g1_walk_rough/mujoco"), + "G1WalkRough", + "mujoco", + 29, + 1.0, + "scene_rough.xml", + _OFFPOLICY_REWARDS, + (*_RESET_EVENTS, "pd_gains"), + True, + id="sac-rough-mujoco", + ), + pytest.param( + "offpolicy", + ("algo=sac", "task=sac/g1_walk_rough/motrix"), + "G1WalkRough", + "motrix", + 29, + 1.0, + "scene_rough.xml", + _OFFPOLICY_REWARDS, + _RESET_EVENTS, + True, + id="sac-rough-motrix", + ), + pytest.param( + "offpolicy", + ("algo=sac", "task=sac/g1_23dof_walk_flat/mujoco"), + "G1Walk23DofFlat", + "mujoco", + 23, + 1.0, + "scene_flat_23dof.xml", + _OFFPOLICY_REWARDS, + (*_RESET_EVENTS, "pd_gains"), + True, + id="sac-23dof-mujoco", + ), + pytest.param( + "offpolicy", + ("algo=sac", "task=sac/g1_23dof_walk_rough/motrix"), + "G1Walk23DofRough", + "motrix", + 23, + 1.0, + "scene_rough_23dof.xml", + _OFFPOLICY_REWARDS, + _RESET_EVENTS, + True, + id="sac-23dof-rough-motrix", + ), + pytest.param( + "offpolicy", + ("algo=td3", "task=td3/g1_walk_flat/mujoco"), + "G1WalkFlat", + "mujoco", + 29, + 1.0, + "scene_flat.xml", + _OFFPOLICY_REWARDS, + (*_RESET_EVENTS, "pd_gains"), + True, + id="td3-mujoco", + ), + pytest.param( + "offpolicy", + ("algo=flashsac", "task=flashsac/g1_walk_flat/mujoco"), + "G1WalkFlat", + "mujoco", + 29, + 1.0, + "scene_flat.xml", + _OFFPOLICY_REWARDS, + (*_RESET_EVENTS, "pd_gains"), + True, + id="flashsac-mujoco", + ), +) + +_WALK_PROFILE_IDS = { + "sac-mujoco", + "sac-motrix", + "sac-mjwarp", + "sac-rough-mujoco", + "sac-rough-motrix", + "sac-23dof-mujoco", + "sac-23dof-rough-motrix", + "td3-mujoco", + "flashsac-mujoco", +} + + +def _compose(config_group: str, overrides: Sequence[str]) -> DictConfig: GlobalHydra.instance().clear() with initialize_config_dir(config_dir=str(CONF_DIR / config_group), version_base="1.3"): - return compose("config", overrides=overrides) - - -def _materialize_env_cfg(cfg: Any): - from unilab.tasks.locomotion.g1.joystick import G1WalkFlatCfg, G1WalkRoughCfg - - env_cfg_cls = G1WalkRoughCfg if cfg.training.task_name == "G1WalkRough" else G1WalkFlatCfg - return OmegaConf.merge(OmegaConf.structured(env_cfg_cls()), cfg.env) - - -def _build_probe_env(cfg: Any): - from unilab.tasks.locomotion.g1.joystick import G1WalkEnv - - env = cast(Any, object.__new__(G1WalkEnv)) - env._num_envs = 1 - env._cfg = _materialize_env_cfg(cfg) - env._reward_cfg = cfg.reward - env.default_angles = np.zeros((1, 29), dtype=np.float32) - env._obs_noise = lambda data, scale: np.asarray(data + 100.0, dtype=np.float32) - return env - - -def _compute_probe_obs(cfg: Any) -> dict[str, np.ndarray]: - env = _build_probe_env(cfg) - return cast( - dict[str, np.ndarray], - env._compute_obs( - { - "commands": np.array([[0.7, 0.0, 0.2]], dtype=np.float32), - "current_actions": np.zeros((1, 29), dtype=np.float32), - "gait_phase": np.array([[0.3, 3.4]], dtype=np.float32), - }, - linvel=np.array([[1.0, 2.0, 3.0]], dtype=np.float32), - gyro=np.array([[4.0, 5.0, 6.0]], dtype=np.float32), - gravity=np.array([[0.1, 0.2, 0.9]], dtype=np.float32), - dof_pos=np.zeros((1, 29), dtype=np.float32), - dof_vel=np.array([np.arange(7.0, 36.0, dtype=np.float32)], dtype=np.float32), - ), + return compose("config", overrides=list(overrides)) + + +def _materialize( + config_group: str, overrides: Sequence[str], task_name: str +) -> tuple[DictConfig, ManagerBasedRlEnvCfg, dict[str, Any]]: + hydra_cfg = _compose(config_group, overrides) + env_override = BackendAdapter(hydra_cfg, root_dir=ROOT_DIR).build_task_env_cfg_override() + env_cfg = registry.materialize_env_config(task_name) + assert isinstance(env_cfg, ManagerBasedRlEnvCfg) + apply_cfg_overrides(env_cfg, env_override) + env_cfg.validate() + return hydra_cfg, env_cfg, env_override + + +def _assert_no_omegaconf(value: Any) -> None: + assert not OmegaConf.is_config(value) + if is_dataclass(value) and not isinstance(value, type): + for item in fields(value): + _assert_no_omegaconf(getattr(value, item.name)) + elif isinstance(value, Mapping): + for key, item in value.items(): + _assert_no_omegaconf(key) + _assert_no_omegaconf(item) + elif isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + for item in value: + _assert_no_omegaconf(item) + + +@pytest.mark.parametrize( + "config_group,overrides,task_name,backend,num_dof,action_scale,model_suffix," + "expected_rewards,expected_events,has_curriculum", + _OWNER_CASES, + ids=[case.id for case in _OWNER_CASES], +) +def test_g1_owner_materializes_complete_plain_manager_cfg( + config_group: str, + overrides: tuple[str, ...], + task_name: str, + backend: str, + num_dof: int, + action_scale: float, + model_suffix: str, + expected_rewards: tuple[str, ...], + expected_events: tuple[str, ...], + has_curriculum: bool, +) -> None: + registry.ensure_registries() + hydra_cfg, env_cfg, _ = _materialize(config_group, overrides, task_name) + case_id = next(case.id for case in _OWNER_CASES if case.values[1] == overrides) + + assert hydra_cfg.training.task_name == task_name + assert hydra_cfg.training.sim_backend == backend + assert env_cfg.ctrl_dt == pytest.approx(0.02) + assert env_cfg.max_episode_seconds == pytest.approx(20.0) + assert env_cfg.policy_observation_group == "policy" + assert env_cfg.critic_observation_group == "critic" + assert env_cfg.scale_rewards_by_dt is True + + assert env_cfg.scene is not None + assert env_cfg.scene.model_file.endswith(f"robots/g1/{model_suffix}") + assert env_cfg.scene.default_keyframe_name == "stand" + robot = env_cfg.scene.entities["robot"] + assert robot.root_body_name == "pelvis" + assert len(robot.joint_names or ()) == num_dof + assert len(robot.actuator_names or ()) == num_dof + assert robot.body_names == ["pelvis"] + + assert list(env_cfg.observations) == ["policy", "critic"] + assert list(env_cfg.observations["policy"].terms) == list(_OBSERVATION_TERMS) + assert list(env_cfg.observations["critic"].terms) == [*_OBSERVATION_TERMS, "base_lin_vel"] + + # Observation scaling profiles are explicit per-owner term scales. + policy_terms = env_cfg.observations["policy"].terms + critic_terms = env_cfg.observations["critic"].terms + walk_profile = case_id in _WALK_PROFILE_IDS + expected_gyro = 0.25 if walk_profile else None + expected_joint_vel = 0.05 if walk_profile else None + expected_linvel = 2.0 if walk_profile else None + for terms in (policy_terms, critic_terms): + assert terms["base_ang_vel"].scale == expected_gyro + assert terms["joint_vel"].scale == expected_joint_vel + assert terms["projected_gravity"].scale is None + assert terms["joint_pos"].scale is None + assert terms["actions"].scale is None + assert terms["command"].scale is None + assert terms["gait_phase"].scale is None + assert terms["gait_phase"].params["frequency"] == pytest.approx(1.5) + assert critic_terms["base_lin_vel"].scale == expected_linvel + + assert list(env_cfg.actions) == ["joint_pos"] + assert env_cfg.actions["joint_pos"].scale == pytest.approx(action_scale) + assert env_cfg.actions["joint_pos"].use_default_offset is True + + assert list(env_cfg.terminations) == ["time_out", "tilt", "base_height"] + assert env_cfg.terminations["time_out"].time_out is True + assert env_cfg.terminations["tilt"].func is g1_terms.g1_tilt_exceeded + assert env_cfg.terminations["base_height"].func is mdp.root_height_below_minimum + + assert tuple(name for name, term in env_cfg.events.items() if term is not None) == ( + expected_events ) - - -@pytest.mark.parametrize("case", _G1_OWNER_CASES, ids=[case["id"] for case in _G1_OWNER_CASES]) -def test_g1_owner_yaml_regression_contract(case: dict[str, Any]): - from unilab.tasks.locomotion.g1.joystick import G1WalkEnv - - cfg = _compose_cfg(case["config_group"], case["overrides"]) - full_env_cfg = _materialize_env_cfg(cfg) - env = _build_probe_env(cfg) - env_cfg_override = BackendAdapter( - cfg, root_dir=ROOT_DIR, algo_name=cfg.algo.algo if "algo" in cfg.algo else None - ).build_task_env_cfg_override() - - assert cfg.training.task_name == case["task_name"] - assert cfg.training.sim_backend == case["backend"] - assert full_env_cfg.control_config.action_scale == pytest.approx(case["action_scale"]) - assert full_env_cfg.curriculum.enabled is case["curriculum_enabled"] - assert env._uses_walk_observation_profile() is (case["profile"] == "walk") - assert ( - registry._envs[cfg.training.task_name].env_factory_dict[cfg.training.sim_backend] - is G1WalkEnv + assert tuple(name for name, term in env_cfg.rewards.items() if term is not None) == ( + expected_rewards ) - - if "model_suffix" in case: - assert full_env_cfg.scene.model_file.endswith(case["model_suffix"]) - - reward_config = OmegaConf.to_container(cfg.reward, resolve=True) - assert env_cfg_override["reward_config"] == reward_config - env_override = cast(dict[str, Any], OmegaConf.to_container(cfg.env, resolve=True)) - for key, value in env_override.items(): - assert env_cfg_override[key] == value - - env._reward_fns = {} - env._init_reward_functions() - for reward_name in cfg.reward.scales.keys(): - assert reward_name in env._reward_fns - - -@pytest.mark.parametrize("case", _G1_OWNER_CASES, ids=[case["id"] for case in _G1_OWNER_CASES]) -def test_g1_owner_yaml_observation_profiles_match_expected_family(case: dict[str, Any]): - cfg = _compose_cfg(case["config_group"], case["overrides"]) - obs = _compute_probe_obs(cfg) - - if case["profile"] == "legacy": - np.testing.assert_allclose(obs["obs"][:, :3], [[104.0, 105.0, 106.0]]) - np.testing.assert_allclose(obs["obs"][:, 35:37], [[107.0, 108.0]]) - np.testing.assert_allclose(obs["critic"][:, :3], [[4.0, 5.0, 6.0]]) - np.testing.assert_allclose(obs["critic"][:, 35:37], [[7.0, 8.0]]) - np.testing.assert_allclose(obs["critic"][:, 98:101], [[1.0, 2.0, 3.0]]) + if has_curriculum: + assert list(env_cfg.curriculum) == ["penalty_scaling"] + assert env_cfg.curriculum["penalty_scaling"].func is g1_terms.G1PenaltyCurriculum else: - np.testing.assert_allclose(obs["obs"][:, :3], [[26.0, 26.25, 26.5]]) - np.testing.assert_allclose(obs["obs"][:, 35:37], [[5.35, 5.4]]) - np.testing.assert_allclose(obs["critic"][:, :3], [[1.0, 1.25, 1.5]]) - np.testing.assert_allclose(obs["critic"][:, 35:37], [[0.35, 0.4]]) - np.testing.assert_allclose(obs["critic"][:, 98:101], [[2.0, 4.0, 6.0]]) - - -def test_g1_observation_profile_selection_prefers_reward_family_over_curriculum_flag(): - from unilab.tasks.locomotion.g1.joystick import G1WalkEnv - - env = cast(Any, object.__new__(G1WalkEnv)) - - env._cfg = cast( - Any, - type( - "Cfg", - (), - {"curriculum": type("Curriculum", (), {"enabled": True})(), "reward_config": None}, - )(), - ) - env._reward_cfg = cast( - Any, - type("RewardCfg", (), {"scales": {"orientation": -2.5, "ang_vel_xy": -0.2}})(), - ) - assert env._uses_walk_observation_profile() is False - - env._cfg = cast( - Any, - type( - "Cfg", - (), - {"curriculum": type("Curriculum", (), {"enabled": False})(), "reward_config": None}, - )(), + assert not env_cfg.curriculum + + command = env_cfg.commands["twist"] + assert isinstance(command, g1_terms.G1VelocityCommandCfg) + assert command.planar_dead_zone == pytest.approx(0.2) + assert command.resampling_time_range == [20.0, 20.0] + if backend == "motrix" and config_group == "ppo": + assert tuple(command.ranges.lin_vel_x) == (0.4, 0.7) + assert tuple(command.ranges.lin_vel_y) == (0.0, 0.0) + else: + assert tuple(command.ranges.lin_vel_x) == (-0.6, 1.0) + assert tuple(command.ranges.lin_vel_y) == (-0.4, 0.4) + assert tuple(command.ranges.ang_vel_z) == (-0.8, 0.8) + + if backend == "mjwarp": + assert env_cfg.mjwarp_nconmax == 128 + assert env_cfg.mjwarp_njmax == 256 + + pose = env_cfg.rewards["pose"] + expected_weights = _POSE_WEIGHTS_29 if num_dof == 29 else _POSE_WEIGHTS_23 + if case_id == "flashsac-mujoco": + expected_weights = [2.0 if i in (1, 7) else w for i, w in enumerate(expected_weights)] + assert list(pose.params["pose_weights"]) == pytest.approx(expected_weights) + + for manager_name in ("observations", "events", "rewards", "terminations", "curriculum"): + for term in getattr(env_cfg, manager_name).values(): + if term is None: + continue + nested_terms = term.terms.values() if manager_name == "observations" else (term,) + for nested in nested_terms: + if nested is None: + continue + module = nested.func.__module__ + assert ".backend." not in module + assert not any(name in module for name in (".mujoco", ".motrix", ".mjwarp")) + + _assert_no_omegaconf(env_cfg) + + +def test_g1_walk_registries_are_manager_only() -> None: + registry.ensure_registries() + metadata = registry.list_registered_envs() + + assert metadata["G1WalkFlat"] == { + "config_factory": "ManagerBasedRlEnvCfg", + "available_backends": ["mujoco", "mjwarp", "motrix"], + } + assert metadata["G1WalkRough"] == { + "config_factory": "ManagerBasedRlEnvCfg", + "available_backends": ["mujoco", "motrix"], + } + assert metadata["G1Walk23DofFlat"] == { + "config_factory": "ManagerBasedRlEnvCfg", + "available_backends": ["mujoco", "motrix"], + } + assert metadata["G1Walk23DofRough"] == { + "config_factory": "ManagerBasedRlEnvCfg", + "available_backends": ["mujoco", "motrix"], + } + + for legacy_override in ( + {"reward_config": {}}, + {"domain_rand": {"randomize_kp": True}}, + {"control_config": {"action_scale": 0.25}}, + {"gait_phase_init_mode": "offset_phase"}, + {"reset_base_qvel_limit": 0.5}, + {"noise_config": {"level": 1.0}}, + ): + with pytest.raises(ValueError, match="has no attribute"): + apply_cfg_overrides(ManagerBasedRlEnvCfg(), legacy_override) + + +@pytest.mark.parametrize( + ("config_group", "overrides", "task_name", "backend", "num_dof", "obs_dim", "critic_dim"), + ( + pytest.param( + "ppo", + ("task=g1_walk_flat/mujoco",), + "G1WalkFlat", + "mujoco", + 29, + 98, + 101, + id="ppo-mujoco", + ), + pytest.param( + "ppo", + ("task=g1_walk_flat/motrix",), + "G1WalkFlat", + "motrix", + 29, + 98, + 101, + id="ppo-motrix", + ), + pytest.param( + "offpolicy", + ("algo=sac", "task=sac/g1_walk_flat/mujoco"), + "G1WalkFlat", + "mujoco", + 29, + 98, + 101, + id="sac-mujoco", + ), + pytest.param( + "ppo", + ("task=g1_23dof_walk_flat/mujoco",), + "G1Walk23DofFlat", + "mujoco", + 23, + 80, + 83, + id="ppo-23dof-mujoco", + ), + pytest.param( + "ppo", + ("task=g1_23dof_walk_rough/mujoco",), + "G1Walk23DofRough", + "mujoco", + 23, + 80, + 83, + id="ppo-23dof-rough-mujoco", + ), + ), +) +def test_g1_registry_executes_real_manager_runtime( + config_group: str, + overrides: tuple[str, ...], + task_name: str, + backend: str, + num_dof: int, + obs_dim: int, + critic_dim: int, +) -> None: + registry.ensure_registries() + _, env_cfg, env_override = _materialize(config_group, overrides, task_name) + try: + env = registry.make( + task_name, + sim_backend=backend, + env_cfg_override=env_override, + num_envs=2, + ) + except ImportError as exc: + pytest.skip(f"{backend} runtime unavailable: {exc}") + + try: + assert isinstance(env, ManagerBasedRlEnv) + assert env.obs_groups_spec == {"obs": obs_dim, "critic": critic_dim} + assert env.action_space.shape == (num_dof,) + action = env.action_manager.get_term("joint_pos") + assert len(action.target_names) == num_dof + + obs, info = env.reset(seed=7) + assert {name: value.shape for name, value in obs.items()} == { + "obs": (2, obs_dim), + "critic": (2, critic_dim), + } + assert isinstance(info, dict) + for _ in range(5): + state = env.step(np.zeros((2, num_dof), dtype=np.float32)) + for value in (*state.obs.values(), state.reward): + assert isinstance(value, np.ndarray) + assert np.isfinite(value).all() + + # The command and gait-phase segments pin the legacy obs layout tail. + command = env.command_manager.get_command("twist") + np.testing.assert_allclose( + state.obs["obs"][:, obs_dim - 5 : obs_dim - 2], command, rtol=0.0, atol=1.0e-6 + ) + finally: + env.close() + + +def test_g1_walk_profile_runtime_obs_scaling_matches_legacy_layout() -> None: + """Walk-profile owners scale gyro x0.25, dof_vel x0.05, critic linvel x2.0.""" + registry.ensure_registries() + _, _, env_override = _materialize( + "offpolicy", ("algo=sac", "task=sac/g1_walk_flat/mujoco"), "G1WalkFlat" ) - env._reward_cfg = cast( - Any, - type( - "RewardCfg", - (), - {"scales": {"penalty_orientation": -10.0, "penalty_ang_vel_xy": -1.0, "alive": 10.0}}, - )(), + try: + env = registry.make( + "G1WalkFlat", sim_backend="mujoco", env_cfg_override=env_override, num_envs=2 + ) + except ImportError as exc: + pytest.skip(f"mujoco runtime unavailable: {exc}") + + try: + env.reset(seed=3) + state = env.step(np.zeros((2, 29), dtype=np.float32)) + gyro = env._backend.get_sensor_data("torso_gyro") + upvector = env._backend.get_sensor_data("torso_upvector") + dof_vel = env._backend.get_dof_vel() + linvel = env._backend.get_sensor_data("pelvis_local_linvel") + np.testing.assert_allclose(state.obs["obs"][:, :3], 0.25 * gyro, rtol=0.0, atol=1.0e-6) + np.testing.assert_allclose(state.obs["obs"][:, 3:6], -upvector, rtol=0.0, atol=1.0e-6) + np.testing.assert_allclose( + state.obs["obs"][:, 35:64], 0.05 * dof_vel, rtol=0.0, atol=1.0e-6 + ) + np.testing.assert_allclose( + state.obs["critic"][:, 98:101], 2.0 * linvel, rtol=0.0, atol=1.0e-6 + ) + np.testing.assert_allclose(state.obs["critic"][:, :3], 0.25 * gyro, rtol=0.0, atol=1.0e-6) + np.testing.assert_allclose( + state.obs["obs"][:, 96:98], state.obs["critic"][:, 96:98], rtol=0.0, atol=0.0 + ) + finally: + env.close() + + +def test_g1_legacy_profile_runtime_obs_scaling_matches_legacy_layout() -> None: + """Legacy-profile owners keep unit scaling on every observation segment.""" + registry.ensure_registries() + _, _, env_override = _materialize("ppo", ("task=g1_walk_flat/mujoco",), "G1WalkFlat") + try: + env = registry.make( + "G1WalkFlat", sim_backend="mujoco", env_cfg_override=env_override, num_envs=2 + ) + except ImportError as exc: + pytest.skip(f"mujoco runtime unavailable: {exc}") + + try: + env.reset(seed=3) + state = env.step(np.zeros((2, 29), dtype=np.float32)) + gyro = env._backend.get_sensor_data("torso_gyro") + dof_vel = env._backend.get_dof_vel() + linvel = env._backend.get_sensor_data("pelvis_local_linvel") + np.testing.assert_allclose(state.obs["obs"][:, :3], gyro, rtol=0.0, atol=1.0e-6) + np.testing.assert_allclose(state.obs["obs"][:, 35:64], dof_vel, rtol=0.0, atol=1.0e-6) + np.testing.assert_allclose(state.obs["critic"][:, 98:101], linvel, rtol=0.0, atol=1.0e-6) + finally: + env.close() + + +def test_g1_penalty_curriculum_scales_negative_weights_from_start() -> None: + registry.ensure_registries() + _, _, env_override = _materialize( + "offpolicy", ("algo=sac", "task=sac/g1_walk_flat/mujoco"), "G1WalkFlat" ) - assert env._uses_walk_observation_profile() is True - - -def test_g1_walk_tasks_are_registered(): - ensure_registries() - - assert registry.contains("G1WalkFlat") - assert registry.contains("G1WalkRough") + try: + env = registry.make( + "G1WalkFlat", sim_backend="mujoco", env_cfg_override=env_override, num_envs=2 + ) + except ImportError as exc: + pytest.skip(f"mujoco runtime unavailable: {exc}") + + try: + assert env.curriculum_manager.active_terms == ["penalty_scaling"] + # initial_scale=0.5 halves every negative weight from construction. + assert env.reward_manager.get_term_cfg("penalty_orientation").weight == pytest.approx(-5.0) + assert env.reward_manager.get_term_cfg("penalty_action_rate").weight == pytest.approx(-2.0) + assert env.reward_manager.get_term_cfg("pose").weight == pytest.approx(-0.25) + # Positive weights stay untouched. + assert env.reward_manager.get_term_cfg("alive").weight == pytest.approx(10.0) + assert env.reward_manager.get_term_cfg("feet_phase").weight == pytest.approx(5.0) + finally: + env.close() diff --git a/tests/envs/locomotion/g1/test_manager_terms.py b/tests/envs/locomotion/g1/test_manager_terms.py new file mode 100644 index 000000000..b7d65c494 --- /dev/null +++ b/tests/envs/locomotion/g1/test_manager_terms.py @@ -0,0 +1,313 @@ +"""Unit contracts for the G1 Manager-Based terms (fake-env, no simulator).""" + +from __future__ import annotations + +import inspect +from types import SimpleNamespace +from typing import Any, cast + +import numpy as np +import pytest + +from unilab.managers import ObservationTermCfg, RewardTermCfg +from unilab.tasks.locomotion.g1 import manager_terms as g1_terms +from unilab.tasks.locomotion.g1.manager_terms import ( + G1GaitPhase, + G1PenaltyCurriculum, + G1VelocityCommandCfg, + compute_feet_phase_contact_targets, + compute_feet_phase_height_targets, +) + + +def _fake_scene(sensor_data: dict[str, np.ndarray]) -> Any: + class _Scene(dict): + def bind_sensor_data(self, names): + arrays = [sensor_data[name] for name in names] + concatenated = np.concatenate(arrays, axis=1) + dimensions = tuple(array.shape[1] for array in arrays) + return SimpleNamespace( + dimensions=dimensions, + backend_type="fake", + read=lambda: concatenated, + ) + + return _Scene() + + +class _FakeEnv: + """Weakref-able stand-in for ManagerBasedRlEnv in term unit tests.""" + + def __init__( + self, + sensor_data: dict[str, np.ndarray], + *, + num_envs: int = 2, + counter: int = 0, + command: np.ndarray | None = None, + ) -> None: + commands = {"twist": command if command is not None else np.zeros((num_envs, 3))} + self.num_envs = num_envs + self.common_step_counter = counter + self.step_dt = 0.02 + self.rng = np.random.default_rng(0) + self.scene = _fake_scene(sensor_data) + self.command_manager = SimpleNamespace(get_command=lambda name: commands[name]) + + +def _fake_env( + sensor_data: dict[str, np.ndarray], + *, + num_envs: int = 2, + counter: int = 0, + command: np.ndarray | None = None, +) -> Any: + return _FakeEnv(sensor_data, num_envs=num_envs, counter=counter, command=command) + + +def test_gait_phase_advances_with_counter_and_resamples_per_init_mode(): + env = _fake_env({}, num_envs=4) + term = G1GaitPhase( + ObservationTermCfg( + func=G1GaitPhase, params={"frequency": 1.5, "init_mode": "offset_phase"} + ), + cast(Any, env), + ) + + np.testing.assert_array_equal(term(env), np.zeros((4, 2))) + term.reset(np.arange(4, dtype=np.int32)) + phase = term(env) + np.testing.assert_allclose(phase[:, 1] - phase[:, 0], np.pi, rtol=1.0e-6) + assert np.all(phase[:, 0] >= 0.0) and np.all(phase[:, 0] < 2.0 * np.pi) + + env.common_step_counter = 1 + advanced = term(env) + delta = 2.0 * np.pi * 1.5 * 0.02 + np.testing.assert_allclose(advanced, np.fmod(phase + delta, 2.0 * np.pi), rtol=1.0e-6) + + # Idempotent per counter: reading twice does not advance twice. + np.testing.assert_array_equal(term(env), advanced) + + env.common_step_counter = 0 + with pytest.raises(ValueError, match="cannot move backwards"): + term(env) + + +def test_gait_phase_independent_mode_samples_feet_independently(): + env = _fake_env({}, num_envs=64) + term = G1GaitPhase( + ObservationTermCfg(func=G1GaitPhase, params={"frequency": 1.5, "init_mode": "independent"}), + cast(Any, env), + ) + term.reset(np.arange(64, dtype=np.int32)) + phase = term(env) + assert not np.allclose(phase[:, 1] - phase[:, 0], np.pi) + + +def test_bezier_targets_match_legacy_reference_values(): + phase = np.array([[0.0, np.pi], [np.pi / 2.0, 3.0 * np.pi / 2.0]]) + left, right = compute_feet_phase_height_targets(phase, 0.09) + # phi=0 -> x=0.5 boundary -> stance peak; phi=pi -> x=0 -> zero; + # phi=pi/2 -> x=0.75 -> swing midpoint 0.045. + np.testing.assert_allclose(left, [0.09, 0.045], atol=1.0e-7) + np.testing.assert_allclose(right, [0.0, 0.045], atol=1.0e-7) + left_contact, right_contact = compute_feet_phase_contact_targets(phase, 0.09) + np.testing.assert_array_equal(left_contact, [False, True]) + np.testing.assert_array_equal(right_contact, [True, True]) + + +def test_feet_phase_reward_is_gated_by_forward_speed(): + sensors = { + "left_foot_pos": np.zeros((2, 3), dtype=np.float32), + "right_foot_pos": np.zeros((2, 3), dtype=np.float32), + "pelvis_local_linvel": np.array([[0.01, 0.0, 0.0], [0.10, 0.0, 0.0]], dtype=np.float32), + } + env = _fake_env(sensors) + term = g1_terms.feet_phase( + RewardTermCfg( + func=g1_terms.feet_phase, + weight=1.0, + params={ + "frequency": 1.5, + "swing_height": 0.09, + "tracking_sigma": 0.008, + "min_forward_speed": 0.05, + "command_name": "twist", + }, + ), + cast(Any, env), + ) + + reward = term(env) + + assert reward[0] == pytest.approx(0.0) + assert reward[1] > 0.0 + + +def test_feet_double_stance_masks_on_forward_command(): + sensors = { + "left_foot_pos": np.zeros((2, 3), dtype=np.float32), + "right_foot_pos": np.zeros((2, 3), dtype=np.float32), + "pelvis_local_linvel": np.zeros((2, 3), dtype=np.float32), + **{f"left_foot_contact_{i}": np.ones((2, 1)) for i in range(4)}, + **{f"right_foot_contact_{i}": np.ones((2, 1)) for i in range(4)}, + } + command = np.array([[0.5, 0.0, 0.0], [0.0, 0.0, 0.0]]) + env = _fake_env(sensors, command=command) + term = g1_terms.feet_double_stance( + RewardTermCfg( + func=g1_terms.feet_double_stance, + weight=-1.0, + params={"frequency": 1.5, "command_name": "twist"}, + ), + cast(Any, env), + ) + + np.testing.assert_array_equal(term(env), [1.0, 0.0]) + + +def _curriculum_env(weights: dict[str, float], num_envs: int = 4) -> Any: + cfgs = {name: SimpleNamespace(weight=value) for name, value in weights.items()} + return SimpleNamespace( + num_envs=num_envs, + reward_manager=SimpleNamespace( + active_terms=list(weights), + get_term_cfg=lambda name: cfgs[name], + ), + reset_buf=np.zeros(num_envs, dtype=np.bool_), + episode_length_buf=np.zeros(num_envs, dtype=np.int64), + rng=np.random.default_rng(0), + ) + + +def test_penalty_curriculum_scales_only_negative_weights_and_tracks_episodes(): + env = _curriculum_env({"pose": -0.5, "alive": 10.0, "penalty_orientation": -10.0}) + term = G1PenaltyCurriculum( + RewardTermCfg( + func=G1PenaltyCurriculum, + weight=1.0, + params={"initial_scale": 0.5, "min_scale": 0.5, "max_scale": 1.0}, + ), + cast(Any, env), + ) + + assert env.reward_manager.get_term_cfg("pose").weight == pytest.approx(-0.25) + assert env.reward_manager.get_term_cfg("penalty_orientation").weight == pytest.approx(-5.0) + assert env.reward_manager.get_term_cfg("alive").weight == pytest.approx(10.0) + + # Short episodes (< level_down_threshold=150 default) shrink the scale, + # clamped at min_scale. + env.reset_buf[:] = True + env.episode_length_buf[:] = 10 + state = term(cast(Any, env), np.arange(4, dtype=np.int32)) + assert state["average_episode_length"] == pytest.approx(10.0) + assert state["penalty_scale"] == pytest.approx(0.5) + assert env.reward_manager.get_term_cfg("pose").weight == pytest.approx(-0.25) + + # Long episodes (> level_up_threshold=750 default) relax the scale. + env.episode_length_buf[:] = 1000 + state = term(cast(Any, env), np.arange(4, dtype=np.int32)) + assert state["penalty_scale"] == pytest.approx(0.5 * (1.0 + 0.001)) + assert env.reward_manager.get_term_cfg("pose").weight == pytest.approx( + -0.5 * 0.5 * (1.0 + 0.001) + ) + + +def test_penalty_curriculum_shrinks_scale_below_initial_when_min_allows(): + env = _curriculum_env({"pose": -0.5}) + term = G1PenaltyCurriculum( + RewardTermCfg( + func=G1PenaltyCurriculum, + weight=1.0, + params={"initial_scale": 0.5, "min_scale": 0.0, "max_scale": 1.0}, + ), + cast(Any, env), + ) + + env.reset_buf[:] = True + env.episode_length_buf[:] = 10 + state = term(cast(Any, env), np.arange(4, dtype=np.int32)) + assert state["penalty_scale"] == pytest.approx(0.5 * (1.0 - 0.001)) + assert env.reward_manager.get_term_cfg("pose").weight == pytest.approx( + -0.5 * 0.5 * (1.0 - 0.001) + ) + + +def test_velocity_command_dead_zone_zeroes_small_planar_commands(): + env = SimpleNamespace( + num_envs=64, + rng=np.random.default_rng(0), + step_dt=0.02, + scene={ + "robot": SimpleNamespace( + data=SimpleNamespace( + root_link_lin_vel_b=np.zeros((64, 3)), + root_link_ang_vel_b=np.zeros((64, 3)), + heading_w=np.zeros(64), + ) + ) + }, + ) + cfg = G1VelocityCommandCfg( + entity_name="robot", + resampling_time_range=(20.0, 20.0), + ranges=G1VelocityCommandCfg.Ranges( + lin_vel_x=(-0.15, 0.15), + lin_vel_y=(-0.15, 0.15), + ang_vel_z=(0.0, 0.0), + ), + ) + term = cfg.build(cast(Any, env)) + term._resample_command(np.arange(64, dtype=np.int32)) + + planar_norm = np.linalg.norm(term.vel_command_b[:, :2], axis=1) + assert np.all((planar_norm == 0.0) | (planar_norm > 0.2)) + assert np.any(planar_norm == 0.0) + + +def test_velocity_command_fails_closed_on_heading_command(): + cfg = G1VelocityCommandCfg( + entity_name="robot", + resampling_time_range=(20.0, 20.0), + heading_command=True, + ranges=G1VelocityCommandCfg.Ranges( + lin_vel_x=(0.0, 1.0), + lin_vel_y=(0.0, 0.0), + ang_vel_z=(0.0, 0.0), + heading=(-3.14, 3.14), + ), + ) + env = SimpleNamespace(num_envs=1, rng=np.random.default_rng(0), step_dt=0.02, scene={}) + with pytest.raises(NotImplementedError, match="heading command"): + cfg.build(cast(Any, env)) + + +def test_g1_hot_paths_use_only_cached_runtime_objects(): + for term in ( + G1GaitPhase, + g1_terms.feet_phase, + g1_terms.feet_phase_contrast, + g1_terms.feet_phase_contact, + g1_terms.feet_double_stance, + g1_terms.feet_air_time, + g1_terms.track_lin_vel, + g1_terms.track_ang_vel, + g1_terms.forward_progress, + g1_terms.under_speed, + g1_terms.lin_vel_z, + g1_terms.ang_vel_xy, + g1_terms.orientation, + g1_terms.g1_tilt_exceeded, + g1_terms.penalty_feet_ori, + g1_terms.penalty_close_feet_xy, + g1_terms.G1PenaltyCurriculum, + ): + source = inspect.getsource(term.__call__) + for forbidden in ( + "ASSETS_ROOT_PATH", + "model_file", + "getattr(", + "hasattr(", + "._backend", + ): + assert forbidden not in source, f"{term.__name__} hot path references {forbidden}" diff --git a/tests/envs/locomotion/g1/test_symmetry_contract.py b/tests/envs/locomotion/g1/test_symmetry_contract.py index 3594e94ca..61bb42a87 100644 --- a/tests/envs/locomotion/g1/test_symmetry_contract.py +++ b/tests/envs/locomotion/g1/test_symmetry_contract.py @@ -1,74 +1,87 @@ +"""G1 symmetry contract on the Manager-Based runtime.""" + from __future__ import annotations +from pathlib import Path from typing import Any, cast import pytest import torch +from hydra import compose, initialize_config_dir +from hydra.core.global_hydra import GlobalHydra from unilab.base import registry from unilab.base.registry import ensure_registries -from unilab.tasks.locomotion.g1.joystick import G1WalkRewardConfig +from unilab.training.backend_adapter import BackendAdapter pytest.importorskip("mujoco", reason="mujoco is required for G1 symmetry contract tests") +ROOT_DIR = Path(__file__).parents[4] +CONF_DIR = ROOT_DIR / "conf" + -def _reward_config() -> G1WalkRewardConfig: - return G1WalkRewardConfig( - scales={"tracking_lin_vel": 2.0, "alive": 10.0}, - tracking_sigma=0.25, - base_height_target=0.754, - min_base_height=0.3, - max_tilt_deg=65.0, - gait_frequency=1.5, - feet_phase_swing_height=0.09, - feet_phase_tracking_sigma=0.04, - close_feet_threshold=0.15, - pose_weights=[0.01] * 29, +def _make_env(task_name: str = "G1WalkFlat", num_envs: int = 1) -> Any: + owner = "g1_23dof_walk_flat/mujoco" if "23Dof" in task_name else "g1_walk_flat/mujoco" + GlobalHydra.instance().clear() + with initialize_config_dir(config_dir=str(CONF_DIR / "ppo"), version_base="1.3"): + cfg = compose("config", overrides=[f"task={owner}"]) + env_cfg_override = BackendAdapter(cfg, root_dir=ROOT_DIR).build_task_env_cfg_override() + return registry.make( + task_name, + num_envs=num_envs, + sim_backend="mujoco", + env_cfg_override=env_cfg_override, ) def test_g1_walk_flat_symmetry_contract_matches_obs_groups(): ensure_registries() - env = cast( - Any, - registry.make( - "G1WalkFlat", - num_envs=1, - sim_backend="mujoco", - env_cfg_override={"reward_config": _reward_config()}, - ), - ) + env = cast(Any, _make_env()) try: layouts = env.get_symmetry_obs_layouts() assert set(layouts) == {"obs", "critic"} + assert [name for name, _ in layouts["obs"]] == [ + "gyro", + "gravity", + "dof_pos", + "dof_vel", + "actions", + "command", + "gait_phase", + ] + assert [name for name, _ in layouts["critic"]][-1] == "linvel" for group_name, layout in layouts.items(): assert sum(dim for _, dim in layout) == env.obs_groups_spec[group_name] finally: env.close() -def test_g1_walk_flat_symmetry_can_augment_critic_group(): +@pytest.mark.parametrize( + ("task_name", "obs_dim", "critic_dim", "action_dim"), + [ + ("G1WalkFlat", 98, 101, 29), + ("G1Walk23DofFlat", 80, 83, 23), + ], +) +def test_g1_walk_symmetry_can_augment_critic_group( + task_name: str, obs_dim: int, critic_dim: int, action_dim: int +): ensure_registries() - env = cast( - Any, - registry.make( - "G1WalkFlat", - num_envs=1, - sim_backend="mujoco", - env_cfg_override={"reward_config": _reward_config()}, - ), - ) + env = cast(Any, _make_env(task_name)) try: augmentation = env.build_symmetry_augmentation(device="cpu") assert augmentation is not None - action_dim = env.action_space.shape[0] + assert env.action_space.shape[0] == action_dim obs = torch.zeros((1, env.obs_groups_spec["obs"])) critic = torch.zeros((1, env.obs_groups_spec["critic"])) actions = torch.zeros((1, action_dim)) + assert env.obs_groups_spec["obs"] == obs_dim + assert env.obs_groups_spec["critic"] == critic_dim + actor_aug, action_aug = augmentation.augment_obs_and_actions(obs, actions, obs_group="obs") critic_aug, critic_action_aug = augmentation.augment_obs_and_actions( critic, @@ -78,8 +91,8 @@ def test_g1_walk_flat_symmetry_can_augment_critic_group(): actor_obs_aug = augmentation.augment_obs(obs, obs_group="obs") critic_obs_aug = augmentation.augment_obs(critic, obs_group="critic") - assert actor_aug.shape == (2, env.obs_groups_spec["obs"]) - assert critic_aug.shape == (2, env.obs_groups_spec["critic"]) + assert actor_aug.shape == (2, obs_dim) + assert critic_aug.shape == (2, critic_dim) assert action_aug.shape == (2, action_dim) assert critic_action_aug.shape == (2, action_dim) assert torch.equal(actor_obs_aug, actor_aug) diff --git a/tests/envs/test_env_configs.py b/tests/envs/test_env_configs.py index aa470d69d..d4a1ec3b7 100644 --- a/tests/envs/test_env_configs.py +++ b/tests/envs/test_env_configs.py @@ -50,6 +50,26 @@ def _allegro_manager_override( ).build_task_env_cfg_override() +def _g1_manager_override(task: str = "g1_walk_flat") -> dict[str, Any]: + from hydra import compose, initialize_config_dir + + from unilab.training.backend_adapter import BackendAdapter + + repo_root = Path(__file__).parents[2] + if task == "g1_walk_rough": + # There is no ppo g1_walk_rough owner; use the SAC owner instead. + with initialize_config_dir( + config_dir=str(repo_root / "conf" / "offpolicy"), version_base="1.3" + ): + cfg = compose("config", overrides=["algo=sac", f"task=sac/{task}/mujoco"]) + return BackendAdapter( + cfg, root_dir=repo_root, algo_name="sac" + ).build_task_env_cfg_override() + with initialize_config_dir(config_dir=str(repo_root / "conf" / "ppo"), version_base="1.3"): + cfg = compose("config", overrides=[f"task={task}/mujoco"]) + return BackendAdapter(cfg, root_dir=repo_root, algo_name="ppo").build_task_env_cfg_override() + + # --------------------------------------------------------------------------- # Non-slow: config attribute completeness (no env.step(), no MuJoCo sim) # --------------------------------------------------------------------------- @@ -102,201 +122,14 @@ def blocked_import(name, globals=None, locals=None, fromlist=(), level=0): assert result.returncode == 0, result.stderr or result.stdout -def test_g1_walk_env_cfg_obs_groups_spec(): - """G1WalkEnv must declare obs_groups_spec with actor and critic groups.""" - from unilab.tasks.locomotion.g1.joystick import G1WalkEnvCfg - - cfg = G1WalkEnvCfg() - assert not hasattr(cfg, "obs_config"), "obs_config should have been removed" - - -def test_g1_walk_flat_cfg_no_obs_config(): - """G1WalkFlatCfg should no longer have obs_config after dict obs refactor.""" - from unilab.tasks.locomotion.g1.joystick import G1WalkFlatCfg - - cfg = G1WalkFlatCfg() - assert not hasattr(cfg, "obs_config"), ( - "obs_config should have been removed in the dict obs refactor" - ) - - -def test_g1_walk_flat_cfg_has_domain_rand_for_motrix(): - from unilab.tasks.locomotion.g1.joystick import G1WalkFlatCfg - - cfg = G1WalkFlatCfg() - assert hasattr(cfg, "domain_rand") - assert hasattr(cfg, "gait_phase_init_mode") - assert hasattr(cfg, "reset_base_qvel_limit") - assert cfg.domain_rand.randomize_base_mass is False - assert cfg.domain_rand.random_com is False - assert cfg.domain_rand.randomize_gravity is False - assert cfg.domain_rand.push_robots is False - - -def test_g1_walk_flat_cfg_defaults_match_walk_profile(): - from unilab.tasks.locomotion.g1.joystick import G1WalkFlatCfg - - cfg = G1WalkFlatCfg() - assert not hasattr(cfg, "obs_profile") - assert cfg.curriculum.enabled is True - - -def test_g1_walk_tasks_register_to_algorithm_agnostic_env_base(): +def test_g1_walk_tasks_register_to_manager_based_env(): from unilab.base import registry - from unilab.tasks.locomotion.g1.joystick import G1WalkEnv, G1WalkRewardConfig - - env = cast( - Any, - registry.make( - "G1WalkFlat", - num_envs=1, - sim_backend="mujoco", - env_cfg_override={ - "reward_config": G1WalkRewardConfig( - scales={"tracking_lin_vel": 2.0, "alive": 10.0}, - tracking_sigma=0.25, - base_height_target=0.754, - min_base_height=0.3, - max_tilt_deg=65.0, - gait_frequency=1.5, - feet_phase_swing_height=0.09, - feet_phase_tracking_sigma=0.04, - close_feet_threshold=0.15, - pose_weights=[0.01] * 29, - ) - }, - ), - ) - try: - assert env.__class__ is G1WalkEnv - finally: - env.close() - - -def test_g1_walk_flat_observation_construction_is_hardcoded_for_legacy_and_walk_modes(): - from unilab.tasks.locomotion.g1.joystick import G1WalkEnv - - class NoiseCfg: - level = 0.0 - scale_gyro = 0.0 - scale_gravity = 0.0 - scale_joint_angle = 0.0 - scale_joint_vel = 0.0 - scale_linvel = 0.0 - - def compute_obs(curriculum_enabled: bool) -> dict[str, np.ndarray]: - env = cast(Any, object.__new__(G1WalkEnv)) - env._num_envs = 1 - env.default_angles = np.array([[0.5, -0.5]], dtype=np.float32) - env._cfg = type( - "Cfg", - (), - { - "noise_config": NoiseCfg(), - "curriculum": type("Curriculum", (), {"enabled": curriculum_enabled})(), - }, - )() - env._obs_noise = lambda data, scale: data + 100.0 - info = { - "commands": np.array([[0.7, 0.0, 0.2]], dtype=np.float32), - "current_actions": np.array([[0.1, -0.2]], dtype=np.float32), - "gait_phase": np.array([[0.3, 3.4]], dtype=np.float32), - } - return cast( - dict[str, np.ndarray], - env._compute_obs( - info, - linvel=np.array([[1.0, 2.0, 3.0]], dtype=np.float32), - gyro=np.array([[4.0, 5.0, 6.0]], dtype=np.float32), - gravity=np.array([[0.1, 0.2, 0.9]], dtype=np.float32), - dof_pos=np.array([[0.6, -0.3]], dtype=np.float32), - dof_vel=np.array([[7.0, 8.0]], dtype=np.float32), - ), - ) - - legacy = compute_obs(curriculum_enabled=False) - walk = compute_obs(curriculum_enabled=True) - - np.testing.assert_allclose(legacy["obs"][:, :3], [[104.0, 105.0, 106.0]]) - np.testing.assert_allclose(legacy["obs"][:, 8:10], [[107.0, 108.0]]) - np.testing.assert_allclose(legacy["critic"][:, :3], [[4.0, 5.0, 6.0]]) - np.testing.assert_allclose(legacy["critic"][:, 17:20], [[1.0, 2.0, 3.0]]) - - np.testing.assert_allclose(walk["obs"][:, :3], [[26.0, 26.25, 26.5]]) - np.testing.assert_allclose(walk["obs"][:, 8:10], [[5.35, 5.4]]) - np.testing.assert_allclose(walk["critic"][:, :3], [[1.0, 1.25, 1.5]]) - np.testing.assert_allclose(walk["critic"][:, 8:10], [[0.35, 0.4]]) - np.testing.assert_allclose(walk["critic"][:, 17:20], [[2.0, 4.0, 6.0]]) - - -def test_g1_walk_env_obs_groups_spec_dims(): - """obs_groups_spec total dim must match what _compute_obs actually produces. - - G1WalkEnv._compute_obs outputs (G1 has 29 DoF): - actor: gyro(3) + gravity(3) + diff(29) + dof_vel(29) - + last_actions(29) + command(3) + gait_phase(2) = 98 - critic: actor(98) + linvel(3) = 101 - """ - from unilab.tasks.locomotion.g1.joystick import G1WalkEnv - - # obs_groups_spec is a @property; access via descriptor protocol - spec = G1WalkEnv.obs_groups_spec.fget(None) # type: ignore[union-attr] - assert spec is not None - assert spec["obs"] == 98 - assert spec["critic"] == 101 - - -def test_g1_walk_env_reward_dispatch_restores_motrix_terms(): - from unilab.tasks.locomotion.g1.joystick import G1WalkEnv - - env = cast(Any, object.__new__(G1WalkEnv)) - env._reward_fns = {} - env._init_reward_functions() - - assert "penalty_feet_ori" in env._reward_fns - assert "feet_phase_contrast" in env._reward_fns - assert "feet_phase_contact" in env._reward_fns - assert "feet_double_stance" in env._reward_fns - - -def test_g1_walk_env_feet_phase_reward_is_gated_by_forward_speed(): - from unilab.tasks.locomotion.common.rewards import RewardContext - from unilab.tasks.locomotion.g1.joystick import G1WalkEnv - - class FakeBackend: - def get_sensor_data(self, name: str) -> np.ndarray: - if name == "left_foot_pos": - return np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], dtype=np.float32) - if name == "right_foot_pos": - return np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], dtype=np.float32) - raise KeyError(name) - - env = cast(Any, object.__new__(G1WalkEnv)) - env._backend = FakeBackend() - env._num_envs = 2 - env._reward_cfg = type( - "RewardCfg", - (), - { - "feet_phase_swing_height": 0.09, - "feet_phase_tracking_sigma": 0.008, - "min_forward_speed_for_gait_reward": 0.05, - }, - )() - - ctx = RewardContext( - info={"gait_phase": np.zeros((2, 2), dtype=np.float32), "commands": np.zeros((2, 3))}, - linvel=np.array([[0.01, 0.0, 0.0], [0.10, 0.0, 0.0]], dtype=np.float32), - gyro=np.zeros((2, 3), dtype=np.float32), - dof_pos=np.zeros((2, 29), dtype=np.float32), - num_envs=2, - ) - - reward = env._reward_feet_phase(ctx) - - assert reward[0] == pytest.approx(0.0) - assert reward[1] > 0.0 + ensure_registries() + metadata = registry.list_registered_envs() + assert metadata["G1WalkFlat"]["config_factory"] == "ManagerBasedRlEnvCfg" + assert metadata["G1WalkFlat"]["available_backends"] == ["mujoco", "mjwarp", "motrix"] + assert metadata["G1WalkRough"]["available_backends"] == ["mujoco", "motrix"] def test_g1_walk_flat_assets_define_contact_sensors_for_gait_rewards(): @@ -1480,8 +1313,8 @@ def test_g1_motion_tracking_cfg_preserves_legacy_defaults(): def test_g1_motion_tracking_init_delegates_motion_body_ids_to_backend(monkeypatch): - from unilab.tasks.locomotion.g1.base import G1BaseEnv from unilab.tasks.motion_tracking.common import tracking as tracking_module + from unilab.tasks.motion_tracking.common.g1_base import G1BaseEnv from unilab.tasks.motion_tracking.g1.tracking import ( G1MotionTrackingCfg, G1MotionTrackingEnv, @@ -2152,8 +1985,6 @@ def test_env_reset_and_step( env_name: str, default_go1_reward_config, default_go2_reward_config, - default_g1_reward_config, - default_g1_walk_flat_reward_config, ): """Every registered env must be constructible, resetable, and steppable. @@ -2166,16 +1997,16 @@ def test_env_reset_and_step( ensure_registries() from unilab.base import registry - # Provide reward_config for envs that require it via Hydra + # Provide config overrides for envs that require them via Hydra env_cfg_override = None if "Go1" in env_name: env_cfg_override = {"reward_config": default_go1_reward_config} elif "Go2" in env_name: env_cfg_override = {"reward_config": default_go2_reward_config} - elif "G1Walk" in env_name: - env_cfg_override = {"reward_config": default_g1_walk_flat_reward_config} - elif "G1" in env_name: - env_cfg_override = {"reward_config": default_g1_reward_config} + elif env_name == "G1WalkFlat": + env_cfg_override = _g1_manager_override("g1_walk_flat") + elif env_name == "G1WalkRough": + env_cfg_override = _g1_manager_override("g1_walk_rough") elif env_name == "AllegroInhandRotation": env_cfg_override = _allegro_manager_override() elif env_name == "AllegroInhandRotationGrasp": diff --git a/tests/envs/test_g1_obs_noise.py b/tests/envs/test_g1_obs_noise.py deleted file mode 100644 index 2dbb1c4c4..000000000 --- a/tests/envs/test_g1_obs_noise.py +++ /dev/null @@ -1,109 +0,0 @@ -"""Tests for G1 per-step observation noise.""" - -from __future__ import annotations - -import numpy as np - -from unilab.tasks.locomotion.g1.base import G1BaseCfg, G1BaseEnv, NoiseConfig - - -class _ConcreteG1Env(G1BaseEnv): - """Minimal concrete subclass — only needed to satisfy the ABC.""" - - def update_state(self, state): - raise NotImplementedError - - -def _make_env(level: float, *, seed: int | None = None) -> G1BaseEnv: - cfg = G1BaseCfg(noise_config=NoiseConfig(level=level, seed=seed)) - env = object.__new__(_ConcreteG1Env) - env._cfg = cfg - return env - - -class TestObsNoise: - def test_noise_applied_when_level_positive(self): - env = _make_env(level=1.0) - data = np.ones((4, 10), dtype=np.float32) - cfg = env._cfg.noise_config - - results = [env._obs_noise(data.copy(), cfg.scale_joint_angle) for _ in range(5)] - # At least one result should differ from the original - assert any(not np.allclose(r, data) for r in results) - - def test_no_noise_when_level_zero(self): - env = _make_env(level=0.0) - data = np.ones((4, 10), dtype=np.float32) - cfg = env._cfg.noise_config - - result = env._obs_noise(data, cfg.scale_joint_angle) - assert result is data - np.testing.assert_array_equal(result, data) - - def test_noise_bounded_by_level_times_scale(self): - env = _make_env(level=1.0) - data = np.zeros((128, 29), dtype=np.float32) - scale = 0.2 - result = env._obs_noise(data.copy(), scale) - # uniform[-1,1] * 1.0 * 0.2 => bounded by [-0.2, 0.2] - assert np.all(result >= -scale) - assert np.all(result <= scale) - - def test_noise_scales_with_level(self): - env_half = _make_env(level=0.5, seed=123) - env_full = _make_env(level=1.0, seed=123) - data = np.zeros((1024, 10), dtype=np.float32) - scale = 1.0 - - r_half = env_half._obs_noise(data.copy(), scale) - r_full = env_full._obs_noise(data.copy(), scale) - - np.testing.assert_allclose(r_full, r_half * 2.0) - - def test_configured_seed_is_reproducible_across_fresh_envs(self): - env_a = _make_env(level=1.0, seed=11) - env_b = _make_env(level=1.0, seed=11) - data = np.zeros((32, 10), dtype=np.float32) - - first_a = env_a._obs_noise(data.copy(), 0.25) - first_b = env_b._obs_noise(data.copy(), 0.25) - second_a = env_a._obs_noise(data.copy(), 0.25) - - np.testing.assert_allclose(first_a, first_b) - assert not np.allclose(first_a, second_a) - - def test_seed_observation_noise_resets_stream(self): - env = _make_env(level=1.0) - data = np.zeros((32, 10), dtype=np.float32) - - env.seed_observation_noise(17) - first = env._obs_noise(data.copy(), 0.25) - env.seed_observation_noise(17) - replayed = env._obs_noise(data.copy(), 0.25) - - np.testing.assert_allclose(first, replayed) - - def test_seed_observation_noise_overrides_configured_seed(self): - env = _make_env(level=1.0, seed=11) - replay = _make_env(level=1.0, seed=99) - data = np.zeros((32, 10), dtype=np.float32) - - env.seed_observation_noise(99) - result = env._obs_noise(data.copy(), 0.25) - expected = replay._obs_noise(data.copy(), 0.25) - - np.testing.assert_allclose(result, expected) - - def test_noise_preserves_dtype(self): - for dt in [np.float32, np.float64]: - env = _make_env(level=1.0) - data = np.ones((4, 5), dtype=dt) - result = env._obs_noise(data, 0.1) - assert result.dtype == dt - - def test_noise_preserves_shape(self): - env = _make_env(level=1.0) - for shape in [(1, 3), (64, 29), (1024, 10)]: - data = np.zeros(shape, dtype=np.float32) - result = env._obs_noise(data, 0.1) - assert result.shape == shape diff --git a/tests/integration/test_reward_injection_integration.py b/tests/integration/test_reward_injection_integration.py index d884a2eb3..dee025824 100644 --- a/tests/integration/test_reward_injection_integration.py +++ b/tests/integration/test_reward_injection_integration.py @@ -21,9 +21,9 @@ def test_reward_injection_in_offpolicy_env_override(): env_cfg_override = build_offpolicy_env_cfg_override("sac", cfg) assert env_cfg_override is not None - assert "reward_config" in env_cfg_override + assert "rewards" in env_cfg_override - # Verify reward config dict has correct values - reward_dict = env_cfg_override["reward_config"] - assert reward_dict["scales"]["tracking_lin_vel"] == 2.0 - assert reward_dict["scales"]["alive"] == 10.0 + # Verify reward term mapping has correct values + reward_dict = env_cfg_override["rewards"] + assert reward_dict["tracking_lin_vel"]["weight"] == 2.0 + assert reward_dict["alive"]["weight"] == 10.0 diff --git a/tests/scripts/test_train_scripts.py b/tests/scripts/test_train_scripts.py index 8d2bd510a..ef523fa1b 100644 --- a/tests/scripts/test_train_scripts.py +++ b/tests/scripts/test_train_scripts.py @@ -605,14 +605,13 @@ def test_offpolicy_g1_walk_flat_motrix_resolved_algo_matches_task_owner(): assert cfg.algo.use_symmetry is False -def test_offpolicy_g1_walk_flat_env_cfg_override_has_reward_and_domain_rand(): +def test_offpolicy_g1_walk_flat_env_cfg_override_has_rewards_and_events(): cfg = _offpolicy_cfg(["task=sac/g1_walk_flat/motrix"]) env_cfg_override = _offpolicy().build_offpolicy_env_cfg_override("sac", cfg) - assert env_cfg_override["reward_config"]["scales"]["tracking_lin_vel"] == pytest.approx(2.2) - assert env_cfg_override["domain_rand"]["randomize_kp"] is False - assert env_cfg_override["domain_rand"]["randomize_kd"] is False + assert env_cfg_override["rewards"]["tracking_lin_vel"]["weight"] == pytest.approx(2.2) + assert env_cfg_override["events"]["pd_gains"] is None def test_offpolicy_g1_walk_flat_backend_scoped_use_symmetry(): @@ -661,14 +660,14 @@ def test_ppo_g1_env_preset_has_env_overrides(): cfg = _ppo_cfg(["task=g1_walk_flat/motrix"]) assert OmegaConf.select(cfg, "env.motrix_max_iterations") is None - assert cfg.env.control_config.action_scale == pytest.approx(0.5) - assert cfg.env.commands.vel_limit == [[0.4, 0.0, 0.0], [0.7, 0.0, 0.0]] - assert cfg.env.gait_phase_init_mode == "offset_phase" - assert cfg.env.reset_base_qvel_limit == pytest.approx(0.05) - assert cfg.reward.scales.feet_phase_contrast == pytest.approx(1.5) - assert cfg.reward.scales.feet_phase_contact == pytest.approx(1.0) - assert cfg.reward.scales.feet_double_stance == pytest.approx(-1.0) - assert cfg.reward.min_forward_speed_for_gait_reward == pytest.approx(0.05) + assert cfg.env.actions.joint_pos.scale == pytest.approx(0.5) + assert cfg.env.commands.twist.ranges.lin_vel_x == [0.4, 0.7] + assert cfg.env.observations.policy.terms.gait_phase.params.init_mode == "offset_phase" + assert cfg.env.events.reset_root_state_uniform.params.velocity_range.x == [-0.05, 0.05] + assert cfg.reward.feet_phase_contrast.weight == pytest.approx(1.5) + assert cfg.reward.feet_phase_contact.weight == pytest.approx(1.0) + assert cfg.reward.feet_double_stance.weight == pytest.approx(-1.0) + assert cfg.reward.feet_phase.params.min_forward_speed == pytest.approx(0.05) def test_ppo_task_go2_aligns_mujoco_with_motrix_defaults(): @@ -744,19 +743,21 @@ def test_build_ppo_env_cfg_override_g1_motrix( env_cfg_override = mod.build_ppo_env_cfg_override(cfg) # env_cfg_override has reward + env preset fields (flat, matching env cfg structure) - assert env_cfg_override["reward_config"]["scales"]["upper_body_pose"] == pytest.approx(-0.05) - assert env_cfg_override["reward_config"]["scales"]["penalty_feet_ori"] == pytest.approx(0.0) - assert env_cfg_override["reward_config"]["scales"]["feet_phase_contrast"] == pytest.approx(1.5) - assert env_cfg_override["reward_config"]["scales"]["feet_phase_contact"] == pytest.approx(1.0) - assert env_cfg_override["reward_config"]["scales"]["feet_double_stance"] == pytest.approx(-1.0) - assert env_cfg_override["reward_config"]["min_forward_speed_for_gait_reward"] == pytest.approx( - 0.05 - ) + assert env_cfg_override["rewards"]["upper_body_pose"]["weight"] == pytest.approx(-0.05) + assert env_cfg_override["rewards"]["penalty_feet_ori"]["weight"] == pytest.approx(0.0) + assert env_cfg_override["rewards"]["feet_phase_contrast"]["weight"] == pytest.approx(1.5) + assert env_cfg_override["rewards"]["feet_phase_contact"]["weight"] == pytest.approx(1.0) + assert env_cfg_override["rewards"]["feet_double_stance"]["weight"] == pytest.approx(-1.0) + assert env_cfg_override["rewards"]["feet_phase"]["params"][ + "min_forward_speed" + ] == pytest.approx(0.05) assert "motrix_max_iterations" not in env_cfg_override - assert env_cfg_override["control_config"]["action_scale"] == pytest.approx(0.5) - assert env_cfg_override["commands"]["vel_limit"] == [[0.4, 0.0, 0.0], [0.7, 0.0, 0.0]] - assert env_cfg_override["gait_phase_init_mode"] == "offset_phase" - assert env_cfg_override["reset_base_qvel_limit"] == pytest.approx(0.05) + assert env_cfg_override["actions"]["joint_pos"]["scale"] == pytest.approx(0.5) + assert env_cfg_override["commands"]["twist"]["ranges"]["lin_vel_x"] == [0.4, 0.7] + assert env_cfg_override["events"]["pd_gains"] is None + assert env_cfg_override["events"]["reset_root_state_uniform"]["params"]["velocity_range"][ + "x" + ] == [-0.05, 0.05] def test_build_ppo_env_cfg_override_carries_motrix_max_iterations_override( @@ -782,14 +783,13 @@ def test_build_ppo_env_cfg_override_carries_post_step_forward_sensor_override( assert env_cfg_override["post_step_forward_sensor"] is value -def test_offpolicy_g1_walk_flat_motrix_env_cfg_override_has_domain_rand(): +def test_offpolicy_g1_walk_flat_motrix_env_cfg_override_disables_pd_gains(): cfg = _offpolicy_cfg(["algo=sac", "task=sac/g1_walk_flat/motrix"]) env_cfg_override = _offpolicy().build_offpolicy_env_cfg_override("sac", cfg) - assert env_cfg_override["domain_rand"]["randomize_kp"] is False - assert env_cfg_override["domain_rand"]["randomize_kd"] is False - assert env_cfg_override["reward_config"]["scales"]["tracking_lin_vel"] == pytest.approx(2.2) + assert env_cfg_override["events"]["pd_gains"] is None + assert env_cfg_override["rewards"]["tracking_lin_vel"]["weight"] == pytest.approx(2.2) def test_build_ppo_env_cfg_override_applies_go2_motrix_reward( diff --git a/tests/tasks/test_migration_matrix.py b/tests/tasks/test_migration_matrix.py index 00d0dcfe1..10b06b87a 100644 --- a/tests/tasks/test_migration_matrix.py +++ b/tests/tasks/test_migration_matrix.py @@ -18,7 +18,7 @@ def test_registered_tasks_have_explicit_migration_records() -> None: assert PRODUCTION_TASK_NAMES <= registered.keys() assert {record.task_name for record in records} == set(PRODUCTION_TASK_NAMES) assert len(records) == 39 - assert sum(record.status == "Compatible" for record in records) == 11 + assert sum(record.status == "Compatible" for record in records) == 15 assert sum(record.target == "compatibility" for record in records) == 3 @@ -28,7 +28,7 @@ def test_registered_tasks_have_explicit_migration_records() -> None: ("Go2ArmManipLoco", "go2_arm", "compatibility", "Adapted"), ("SharpaInhandRotation", "sharpa", "compatibility", "Adapted"), ("G1MotionTracking", "motion_tracking", "mba", "Adapted"), - ("G1WalkRough", "g1_locomotion", "mba", "Adapted"), + ("G1WalkRough", "g1_locomotion", "complete", "Compatible"), ("Go2JoystickRough", "quadruped_rough", "complete", "Compatible"), ], ) diff --git a/tests/training/test_sim2sim_resolver.py b/tests/training/test_sim2sim_resolver.py index fc971bbba..b9530a3d9 100644 --- a/tests/training/test_sim2sim_resolver.py +++ b/tests/training/test_sim2sim_resolver.py @@ -421,9 +421,9 @@ def _compose_task(task: str) -> Any: def test_g1_walk_flat_mujoco_inherits_base_contract(): - # The MuJoCo owner carries the full contract in its standalone owner config. + # The MuJoCo owner inherits the full contract from the shared base owner. mujoco = _compose_task("g1_walk_flat/mujoco") - assert OmegaConf.select(mujoco, "env.control_config.action_scale") == 0.25 + assert OmegaConf.select(mujoco, "env.actions.joint_pos.scale") == 0.25 assert OmegaConf.select(mujoco, "algo.empirical_normalization") is False assert OmegaConf.select(mujoco, "algo.obs_groups.actor") == ["actor"] diff --git a/tests/training/test_training_helpers.py b/tests/training/test_training_helpers.py index 84abff4a8..e2eee8462 100644 --- a/tests/training/test_training_helpers.py +++ b/tests/training/test_training_helpers.py @@ -381,9 +381,8 @@ def test_backend_adapter_env_cfg_override_for_motrix_sac_g1_walk_flat(): env_cfg_override = adapter.build_task_env_cfg_override() # env_cfg_override has reward + env preset fields - assert env_cfg_override["reward_config"]["scales"]["tracking_lin_vel"] == pytest.approx(2.2) - assert env_cfg_override["domain_rand"]["randomize_kp"] is False - assert env_cfg_override["domain_rand"]["randomize_kd"] is False + assert env_cfg_override["rewards"]["tracking_lin_vel"]["weight"] == pytest.approx(2.2) + assert env_cfg_override["events"]["pd_gains"] is None # algo values come straight from YAML compose — no mutation, matches task owner values assert cfg.algo.num_envs == 2048 assert cfg.algo.max_iterations == 5000 From 6ac9bec3b686715d6064b8fc319671c9a9a929bb Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Thu, 20 Aug 2026 14:39:12 +0800 Subject: [PATCH 092/169] chore: drop inferable manager term _target_ from G1 owner configs (#1233) --- conf/appo/task/g1_walk_flat/base.yaml | 32 ------------------ conf/offpolicy/task/g1_walk_flat/base.yaml | 33 ------------------- conf/ppo/task/g1_23dof_walk_flat/motrix.yaml | 7 ---- conf/ppo/task/g1_23dof_walk_rough/mujoco.yaml | 1 - conf/ppo/task/g1_walk_flat/base.yaml | 32 ------------------ conf/ppo/task/g1_walk_flat/motrix.yaml | 7 ---- scripts/benchmark/env/benchmark_env_step.py | 17 ++++++++++ .../test_env_step_config_contract.py | 5 +++ 8 files changed, 22 insertions(+), 112 deletions(-) diff --git a/conf/appo/task/g1_walk_flat/base.yaml b/conf/appo/task/g1_walk_flat/base.yaml index 8c74c34c1..bd6890ac7 100644 --- a/conf/appo/task/g1_walk_flat/base.yaml +++ b/conf/appo/task/g1_walk_flat/base.yaml @@ -76,69 +76,52 @@ env: max_episode_seconds: 20.0 observations: policy: - _target_: unilab.managers.ObservationGroupCfg terms: base_ang_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.builtin_sensor params: {sensor_name: torso_gyro} projected_gravity: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.projected_gravity_from_sensor params: {sensor_name: torso_upvector} joint_pos: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_pos_rel joint_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_vel_rel actions: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.last_action params: {action_name: joint_pos} command: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.generated_commands params: {command_name: twist} gait_phase: - _target_: unilab.managers.ObservationTermCfg func: unilab.tasks.locomotion.g1.manager_terms.G1GaitPhase params: frequency: 1.5 init_mode: offset_phase critic: - _target_: unilab.managers.ObservationGroupCfg terms: base_ang_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.builtin_sensor params: {sensor_name: torso_gyro} projected_gravity: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.projected_gravity_from_sensor params: {sensor_name: torso_upvector} joint_pos: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_pos_rel joint_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_vel_rel actions: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.last_action params: {action_name: joint_pos} command: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.generated_commands params: {command_name: twist} gait_phase: - _target_: unilab.managers.ObservationTermCfg func: unilab.tasks.locomotion.g1.manager_terms.G1GaitPhase params: frequency: 1.5 init_mode: offset_phase base_lin_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.builtin_sensor params: {sensor_name: pelvis_local_linvel} actions: @@ -167,11 +150,9 @@ env: ang_vel_z: [-0.8, 0.8] events: reset_scene_to_default: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.reset_scene_to_default mode: reset reset_root_state_uniform: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.reset_root_state_uniform mode: reset params: @@ -190,7 +171,6 @@ env: pitch: [-0.5, 0.5] yaw: [-0.5, 0.5] pd_gains: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.pd_gains mode: reset params: @@ -199,16 +179,13 @@ env: operation: scale terminations: time_out: - _target_: unilab.managers.TerminationTermCfg func: unilab.envs.mdp.time_out time_out: true tilt: - _target_: unilab.managers.TerminationTermCfg func: unilab.tasks.locomotion.g1.manager_terms.g1_tilt_exceeded params: max_tilt_deg: 25.0 base_height: - _target_: unilab.managers.TerminationTermCfg func: unilab.envs.mdp.root_height_below_minimum params: minimum_height: 0.55 @@ -217,21 +194,18 @@ env: reward: tracking_lin_vel: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.g1.manager_terms.track_lin_vel weight: 2.0 params: tracking_sigma: 0.25 command_name: twist tracking_ang_vel: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.g1.manager_terms.track_ang_vel weight: 0.2 params: tracking_sigma: 0.25 command_name: twist feet_phase: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.g1.manager_terms.feet_phase weight: 1.0 params: @@ -241,29 +215,23 @@ reward: min_forward_speed: 0.0 command_name: twist lin_vel_z: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.g1.manager_terms.lin_vel_z weight: -1.0 ang_vel_xy: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.g1.manager_terms.ang_vel_xy weight: -0.25 base_height: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.g1.manager_terms.base_height weight: -500.0 params: target_height: 0.754 orientation: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.g1.manager_terms.orientation weight: -5.0 action_rate: - _target_: unilab.managers.RewardTermCfg func: unilab.envs.mdp.action_rate_l2 weight: -0.01 pose: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.g1.manager_terms.weighted_pose weight: -0.1 params: diff --git a/conf/offpolicy/task/g1_walk_flat/base.yaml b/conf/offpolicy/task/g1_walk_flat/base.yaml index 82fde7853..2c25972d4 100644 --- a/conf/offpolicy/task/g1_walk_flat/base.yaml +++ b/conf/offpolicy/task/g1_walk_flat/base.yaml @@ -77,73 +77,56 @@ env: max_episode_seconds: 20.0 observations: policy: - _target_: unilab.managers.ObservationGroupCfg terms: base_ang_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.builtin_sensor params: {sensor_name: torso_gyro} scale: 0.25 projected_gravity: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.projected_gravity_from_sensor params: {sensor_name: torso_upvector} joint_pos: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_pos_rel joint_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_vel_rel scale: 0.05 actions: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.last_action params: {action_name: joint_pos} command: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.generated_commands params: {command_name: twist} gait_phase: - _target_: unilab.managers.ObservationTermCfg func: unilab.tasks.locomotion.g1.manager_terms.G1GaitPhase params: frequency: 1.5 init_mode: offset_phase critic: - _target_: unilab.managers.ObservationGroupCfg terms: base_ang_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.builtin_sensor params: {sensor_name: torso_gyro} scale: 0.25 projected_gravity: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.projected_gravity_from_sensor params: {sensor_name: torso_upvector} joint_pos: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_pos_rel joint_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_vel_rel scale: 0.05 actions: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.last_action params: {action_name: joint_pos} command: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.generated_commands params: {command_name: twist} gait_phase: - _target_: unilab.managers.ObservationTermCfg func: unilab.tasks.locomotion.g1.manager_terms.G1GaitPhase params: frequency: 1.5 init_mode: offset_phase base_lin_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.builtin_sensor params: {sensor_name: pelvis_local_linvel} scale: 2.0 @@ -173,11 +156,9 @@ env: ang_vel_z: [-0.8, 0.8] events: reset_scene_to_default: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.reset_scene_to_default mode: reset reset_root_state_uniform: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.reset_root_state_uniform mode: reset params: @@ -196,7 +177,6 @@ env: pitch: [-0.5, 0.5] yaw: [-0.5, 0.5] pd_gains: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.pd_gains mode: reset params: @@ -205,22 +185,18 @@ env: operation: scale terminations: time_out: - _target_: unilab.managers.TerminationTermCfg func: unilab.envs.mdp.time_out time_out: true tilt: - _target_: unilab.managers.TerminationTermCfg func: unilab.tasks.locomotion.g1.manager_terms.g1_tilt_exceeded params: max_tilt_deg: 65.0 base_height: - _target_: unilab.managers.TerminationTermCfg func: unilab.envs.mdp.root_height_below_minimum params: minimum_height: 0.3 curriculum: penalty_scaling: - _target_: unilab.managers.CurriculumTermCfg func: unilab.tasks.locomotion.g1.manager_terms.G1PenaltyCurriculum params: initial_scale: 0.5 @@ -234,43 +210,35 @@ env: reward: tracking_lin_vel: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.g1.manager_terms.track_lin_vel weight: 2.0 params: tracking_sigma: 0.25 command_name: twist tracking_ang_vel: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.g1.manager_terms.track_ang_vel weight: 1.5 params: tracking_sigma: 0.25 command_name: twist penalty_ang_vel_xy: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.g1.manager_terms.ang_vel_xy weight: -1.0 penalty_orientation: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.g1.manager_terms.orientation weight: -10.0 penalty_action_rate: - _target_: unilab.managers.RewardTermCfg func: unilab.envs.mdp.action_rate_l2 weight: -4.0 pose: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.g1.manager_terms.weighted_pose weight: -0.5 params: pose_weights: [0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0] penalty_feet_ori: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.g1.manager_terms.penalty_feet_ori weight: -20.0 feet_phase: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.g1.manager_terms.feet_phase weight: 5.0 params: @@ -280,6 +248,5 @@ reward: min_forward_speed: 0.0 command_name: twist alive: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.g1.manager_terms.alive weight: 10.0 diff --git a/conf/ppo/task/g1_23dof_walk_flat/motrix.yaml b/conf/ppo/task/g1_23dof_walk_flat/motrix.yaml index 615cb7ef4..d8fd852d0 100644 --- a/conf/ppo/task/g1_23dof_walk_flat/motrix.yaml +++ b/conf/ppo/task/g1_23dof_walk_flat/motrix.yaml @@ -56,25 +56,21 @@ reward: tracking_ang_vel: weight: 0.25 forward_progress: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.g1.manager_terms.forward_progress weight: 0.0 params: command_name: twist under_speed: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.g1.manager_terms.under_speed weight: -0.2 params: command_name: twist upper_body_pose: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.g1.manager_terms.upper_body_pose weight: -0.05 params: pose_weights: [0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0] penalty_feet_ori: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.g1.manager_terms.penalty_feet_ori weight: 0.0 feet_phase: @@ -82,7 +78,6 @@ reward: params: min_forward_speed: 0.05 feet_phase_contrast: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.g1.manager_terms.feet_phase_contrast weight: 1.5 params: @@ -92,7 +87,6 @@ reward: min_forward_speed: 0.05 command_name: twist feet_phase_contact: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.g1.manager_terms.feet_phase_contact weight: 1.0 params: @@ -102,7 +96,6 @@ reward: min_forward_speed: 0.05 command_name: twist feet_double_stance: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.g1.manager_terms.feet_double_stance weight: -1.0 params: diff --git a/conf/ppo/task/g1_23dof_walk_rough/mujoco.yaml b/conf/ppo/task/g1_23dof_walk_rough/mujoco.yaml index 2a960832a..cd7700207 100644 --- a/conf/ppo/task/g1_23dof_walk_rough/mujoco.yaml +++ b/conf/ppo/task/g1_23dof_walk_rough/mujoco.yaml @@ -20,7 +20,6 @@ env: model_file: src/unilab/assets/robots/g1/scene_rough_23dof.xml curriculum: penalty_scaling: - _target_: unilab.managers.CurriculumTermCfg func: unilab.tasks.locomotion.g1.manager_terms.G1PenaltyCurriculum params: initial_scale: 0.5 diff --git a/conf/ppo/task/g1_walk_flat/base.yaml b/conf/ppo/task/g1_walk_flat/base.yaml index ae8130aaf..c8466f9a2 100644 --- a/conf/ppo/task/g1_walk_flat/base.yaml +++ b/conf/ppo/task/g1_walk_flat/base.yaml @@ -76,69 +76,52 @@ env: max_episode_seconds: 20.0 observations: policy: - _target_: unilab.managers.ObservationGroupCfg terms: base_ang_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.builtin_sensor params: {sensor_name: torso_gyro} projected_gravity: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.projected_gravity_from_sensor params: {sensor_name: torso_upvector} joint_pos: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_pos_rel joint_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_vel_rel actions: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.last_action params: {action_name: joint_pos} command: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.generated_commands params: {command_name: twist} gait_phase: - _target_: unilab.managers.ObservationTermCfg func: unilab.tasks.locomotion.g1.manager_terms.G1GaitPhase params: frequency: 1.5 init_mode: offset_phase critic: - _target_: unilab.managers.ObservationGroupCfg terms: base_ang_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.builtin_sensor params: {sensor_name: torso_gyro} projected_gravity: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.projected_gravity_from_sensor params: {sensor_name: torso_upvector} joint_pos: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_pos_rel joint_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.joint_vel_rel actions: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.last_action params: {action_name: joint_pos} command: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.generated_commands params: {command_name: twist} gait_phase: - _target_: unilab.managers.ObservationTermCfg func: unilab.tasks.locomotion.g1.manager_terms.G1GaitPhase params: frequency: 1.5 init_mode: offset_phase base_lin_vel: - _target_: unilab.managers.ObservationTermCfg func: unilab.envs.mdp.builtin_sensor params: {sensor_name: pelvis_local_linvel} actions: @@ -167,11 +150,9 @@ env: ang_vel_z: [-0.8, 0.8] events: reset_scene_to_default: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.reset_scene_to_default mode: reset reset_root_state_uniform: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.reset_root_state_uniform mode: reset params: @@ -190,7 +171,6 @@ env: pitch: [-0.5, 0.5] yaw: [-0.5, 0.5] pd_gains: - _target_: unilab.managers.EventTermCfg func: unilab.envs.mdp.pd_gains mode: reset params: @@ -199,16 +179,13 @@ env: operation: scale terminations: time_out: - _target_: unilab.managers.TerminationTermCfg func: unilab.envs.mdp.time_out time_out: true tilt: - _target_: unilab.managers.TerminationTermCfg func: unilab.tasks.locomotion.g1.manager_terms.g1_tilt_exceeded params: max_tilt_deg: 25.0 base_height: - _target_: unilab.managers.TerminationTermCfg func: unilab.envs.mdp.root_height_below_minimum params: minimum_height: 0.55 @@ -217,21 +194,18 @@ env: reward: tracking_lin_vel: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.g1.manager_terms.track_lin_vel weight: 2.0 params: tracking_sigma: 0.25 command_name: twist tracking_ang_vel: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.g1.manager_terms.track_ang_vel weight: 0.2 params: tracking_sigma: 0.25 command_name: twist feet_phase: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.g1.manager_terms.feet_phase weight: 1.0 params: @@ -241,29 +215,23 @@ reward: min_forward_speed: 0.0 command_name: twist lin_vel_z: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.g1.manager_terms.lin_vel_z weight: -1.0 ang_vel_xy: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.g1.manager_terms.ang_vel_xy weight: -0.25 base_height: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.g1.manager_terms.base_height weight: -500.0 params: target_height: 0.754 orientation: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.g1.manager_terms.orientation weight: -5.0 action_rate: - _target_: unilab.managers.RewardTermCfg func: unilab.envs.mdp.action_rate_l2 weight: -0.01 pose: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.g1.manager_terms.weighted_pose weight: -0.1 params: diff --git a/conf/ppo/task/g1_walk_flat/motrix.yaml b/conf/ppo/task/g1_walk_flat/motrix.yaml index 5690fa578..b93cc9269 100644 --- a/conf/ppo/task/g1_walk_flat/motrix.yaml +++ b/conf/ppo/task/g1_walk_flat/motrix.yaml @@ -58,25 +58,21 @@ reward: tracking_ang_vel: weight: 0.25 forward_progress: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.g1.manager_terms.forward_progress weight: 0.0 params: command_name: twist under_speed: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.g1.manager_terms.under_speed weight: -0.2 params: command_name: twist upper_body_pose: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.g1.manager_terms.upper_body_pose weight: -0.05 params: pose_weights: [0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0] penalty_feet_ori: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.g1.manager_terms.penalty_feet_ori weight: 0.0 feet_phase: @@ -84,7 +80,6 @@ reward: params: min_forward_speed: 0.05 feet_phase_contrast: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.g1.manager_terms.feet_phase_contrast weight: 1.5 params: @@ -94,7 +89,6 @@ reward: min_forward_speed: 0.05 command_name: twist feet_phase_contact: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.g1.manager_terms.feet_phase_contact weight: 1.0 params: @@ -104,7 +98,6 @@ reward: min_forward_speed: 0.05 command_name: twist feet_double_stance: - _target_: unilab.managers.RewardTermCfg func: unilab.tasks.locomotion.g1.manager_terms.feet_double_stance weight: -1.0 params: diff --git a/scripts/benchmark/env/benchmark_env_step.py b/scripts/benchmark/env/benchmark_env_step.py index f8422efab..1445b83e4 100644 --- a/scripts/benchmark/env/benchmark_env_step.py +++ b/scripts/benchmark/env/benchmark_env_step.py @@ -116,9 +116,26 @@ def _patched_create_backend(backend_type, scene, num_envs, sim_dt, **kwargs): _ub_backend.create_backend = _patched_create_backend _ub_backend._mjwarp_patched = True + _ub_backend._mjwarp_orig_create_backend = _orig_create_backend return True +def _uninstall_mjwarp_patch() -> None: + """Undo the import-time factory patch installed by :func:`_install_mjwarp_patch`. + + Importing this module inside a pytest session would otherwise leak the + mjwarp rerouting into unrelated tests that exercise the real factory. + """ + import unilab.base.backend as _ub_backend + + if not getattr(_ub_backend, "_mjwarp_patched", False): + return + orig = getattr(_ub_backend, "_mjwarp_orig_create_backend", None) + if orig is not None: + _ub_backend.create_backend = orig + _ub_backend._mjwarp_patched = False + + MJWARP_AVAILABLE = _install_mjwarp_patch() BACKENDS = ["mujoco", "motrix", "mjwarp"] diff --git a/tests/benchmark/test_env_step_config_contract.py b/tests/benchmark/test_env_step_config_contract.py index 720094c5a..f7bbe7c64 100644 --- a/tests/benchmark/test_env_step_config_contract.py +++ b/tests/benchmark/test_env_step_config_contract.py @@ -7,6 +7,11 @@ from unilab.envs import ManagerBasedRlEnvCfg, make_manager_based_rl_env from unilab.tasks.locomotion.common.rough_manager_terms import QuadrupedRoughTerrainCfg +# Importing benchmark_env_step installs a process-wide create_backend patch for +# the benchmark script. These tests only build configs, so undo the patch to +# keep the global factory pristine for the rest of the pytest session. +bench._uninstall_mjwarp_patch() + def test_go2w_flat_benchmark_uses_production_manager_owner() -> None: cfg = bench.TASK_CONFIGS["go2w"].build_cfg("mujoco") From 331caa64b9053ee97ee582fb96a0f6de8d567047 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Thu, 20 Aug 2026 15:48:42 +0800 Subject: [PATCH 093/169] refactor: split unilab.tools into owner layers and scripts/tools (#1234) - scripts/tools/import_robot.py: one-off URDF->MJCF asset authoring (hardcodes REPO_ROOT, drops unilab-import-robot console entry) - scripts/motion/bones_seed_csv.py: CSV contract helper co-located with its only consumers - src/unilab/utils/nan_viz.py: viewer next to its nan_guard dump producer - src/unilab/assets/pull.py: thin CLI over assets.hub - src/unilab/base/backend/mujoco/export_scene.py: MuJoCo-only export - src/unilab/visualization/teaser.py: motrix teaser scene renderer - src/unilab/cli_completion.py: shell completion sibling of cli.py - update console entries (pyproject.toml, pyproject.rocm.toml), demo.py, shell completion fallbacks, tests, docs; drop api_reference/tools page --- docs/sphinx/source/api_reference/index.md | 7 ------- docs/sphinx/source/api_reference/tools/index.md | 15 --------------- .../en/2-user_guide/7-tooling/3-nan_visualizer.md | 2 +- .../en/2-user_guide/7-tooling/4-scene_export.md | 2 +- .../en/2-user_guide/7-tooling/5-robot_import.md | 4 ++-- .../2-user_guide/7-tooling/3-nan_visualizer.md | 2 +- .../2-user_guide/7-tooling/4-scene_export.md | 2 +- .../2-user_guide/7-tooling/5-robot_import.md | 4 ++-- pyproject.rocm.toml | 11 +++++------ pyproject.toml | 11 +++++------ scripts/completions/unilab.bash | 2 +- scripts/completions/unilab.zsh | 2 +- .../tools => scripts/motion}/bones_seed_csv.py | 0 scripts/motion/bones_seed_csv_to_npz.py | 10 +++++----- scripts/motion/replay_bones_seed_csv.py | 6 +++--- {src/unilab => scripts}/tools/import_robot.py | 4 ++-- .../{tools/pull_assets.py => assets/pull.py} | 0 .../backend/mujoco}/export_scene.py | 0 .../{tools/completion.py => cli_completion.py} | 0 src/unilab/demo.py | 4 ++-- src/unilab/tools/__init__.py | 1 - src/unilab/{tools/viz_nan.py => utils/nan_viz.py} | 0 .../render_teaser.py => visualization/teaser.py} | 0 tests/scripts/test_mujoco_only_tooling_markers.py | 2 +- tests/test_bones_seed_csv.py | 5 ++--- tests/test_cli.py | 8 ++++---- tests/test_completion.py | 2 +- tests/test_export_scene.py | 2 +- tests/test_import_robot.py | 2 +- tests/test_nan_guard.py | 2 +- tests/test_render_teaser.py | 4 ++-- tests/utils/test_utils_package_policy.py | 1 + 32 files changed, 46 insertions(+), 71 deletions(-) delete mode 100644 docs/sphinx/source/api_reference/tools/index.md rename {src/unilab/tools => scripts/motion}/bones_seed_csv.py (100%) rename {src/unilab => scripts}/tools/import_robot.py (99%) rename src/unilab/{tools/pull_assets.py => assets/pull.py} (100%) rename src/unilab/{tools => base/backend/mujoco}/export_scene.py (100%) rename src/unilab/{tools/completion.py => cli_completion.py} (100%) delete mode 100644 src/unilab/tools/__init__.py rename src/unilab/{tools/viz_nan.py => utils/nan_viz.py} (100%) rename src/unilab/{tools/render_teaser.py => visualization/teaser.py} (100%) diff --git a/docs/sphinx/source/api_reference/index.md b/docs/sphinx/source/api_reference/index.md index fdbed0fb2..194fa57e8 100644 --- a/docs/sphinx/source/api_reference/index.md +++ b/docs/sphinx/source/api_reference/index.md @@ -97,12 +97,6 @@ Procedural and heightfield terrain generators. Scene rendering and viser bridges. ::: -:::{grid-item-card} 🔧 `unilab.tools` -:link: tools/index -:link-type: doc -Scene export, NaN visualizer, ONNX export. -::: - :::{grid-item-card} 🧰 `unilab.utils` :link: utils/index :link-type: doc @@ -151,7 +145,6 @@ backend/index dr/index terrains/index visualization/index -tools/index utils/index logging/index ``` diff --git a/docs/sphinx/source/api_reference/tools/index.md b/docs/sphinx/source/api_reference/tools/index.md deleted file mode 100644 index a8e1918a5..000000000 --- a/docs/sphinx/source/api_reference/tools/index.md +++ /dev/null @@ -1,15 +0,0 @@ -# `unilab.tools` — CLI Tools - -Console-script entrypoints registered in `pyproject.toml`: - -- `unilab-viz-nan` — interactive NaN trace viewer for failed runs. -- `unilab-export-scene` — dump the resolved scene of a task to disk. - -```{eval-rst} -.. autosummary:: - :toctree: _autosummary - :template: autosummary/module.rst - :recursive: - - unilab.tools -``` diff --git a/docs/sphinx/source/en/2-user_guide/7-tooling/3-nan_visualizer.md b/docs/sphinx/source/en/2-user_guide/7-tooling/3-nan_visualizer.md index 53495316c..bab4bef06 100644 --- a/docs/sphinx/source/en/2-user_guide/7-tooling/3-nan_visualizer.md +++ b/docs/sphinx/source/en/2-user_guide/7-tooling/3-nan_visualizer.md @@ -12,7 +12,7 @@ uv run train --algo ppo --task go2_joystick_flat --sim mujoco \ training.nan_guard.output_dir=/tmp/unilab/nan_dumps ``` -The viewer implementation is `src/unilab/tools/viz_nan.py`, registered as the +The viewer implementation is `src/unilab/utils/nan_viz.py`, registered as the `unilab-viz-nan` console entry. It replays a dump path and lets you select the environment index. Dump format and round-trip loading are covered by `tests/test_nan_guard.py`. diff --git a/docs/sphinx/source/en/2-user_guide/7-tooling/4-scene_export.md b/docs/sphinx/source/en/2-user_guide/7-tooling/4-scene_export.md index f92fa64b7..dfeea30da 100644 --- a/docs/sphinx/source/en/2-user_guide/7-tooling/4-scene_export.md +++ b/docs/sphinx/source/en/2-user_guide/7-tooling/4-scene_export.md @@ -1,6 +1,6 @@ # Scene Export -Scene export is implemented by `src/unilab/tools/export_scene.py` and registered +Scene export is implemented by `src/unilab/base/backend/mujoco/export_scene.py` and registered as the `unilab-export-scene` console entry in `pyproject.toml`. It accepts a MuJoCo XML or MJB model path, writes `scene.xml`, copies mesh assets when they are discoverable, and can create a zip archive. diff --git a/docs/sphinx/source/en/2-user_guide/7-tooling/5-robot_import.md b/docs/sphinx/source/en/2-user_guide/7-tooling/5-robot_import.md index bb6569642..888be3ff6 100644 --- a/docs/sphinx/source/en/2-user_guide/7-tooling/5-robot_import.md +++ b/docs/sphinx/source/en/2-user_guide/7-tooling/5-robot_import.md @@ -30,7 +30,7 @@ Prefer MuJoCo/MJCF `.xml`, copied according to the contract above. If the source is URDF-only, convert it with the repository script: ```bash -uv run unilab-import-robot [robot_name] +uv run scripts/tools/import_robot.py [robot_name] ``` ```{important} @@ -72,7 +72,7 @@ When checking `home`, confirm at least: ## Output Artifacts -After running `uv run unilab-import-robot [robot_name]`, the script +After running `uv run scripts/tools/import_robot.py [robot_name]`, the script generates: - `src/unilab/assets/robots//assets/`: converted and organized mesh diff --git a/docs/sphinx/source/zh_CN/2-user_guide/7-tooling/3-nan_visualizer.md b/docs/sphinx/source/zh_CN/2-user_guide/7-tooling/3-nan_visualizer.md index 7fe13f910..cd442ddd7 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/7-tooling/3-nan_visualizer.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/7-tooling/3-nan_visualizer.md @@ -8,4 +8,4 @@ uv run train --algo ppo --task go2_joystick_flat --sim mujoco \ training.nan_guard.output_dir=/tmp/unilab/nan_dumps ``` -viewer 的实现是 `src/unilab/tools/viz_nan.py`，注册为 `unilab-viz-nan` 控制台入口。它会回放一个 dump 路径，并让你选择环境索引。dump 格式和往返加载由 `tests/test_nan_guard.py` 覆盖。 +viewer 的实现是 `src/unilab/utils/nan_viz.py`，注册为 `unilab-viz-nan` 控制台入口。它会回放一个 dump 路径，并让你选择环境索引。dump 格式和往返加载由 `tests/test_nan_guard.py` 覆盖。 diff --git a/docs/sphinx/source/zh_CN/2-user_guide/7-tooling/4-scene_export.md b/docs/sphinx/source/zh_CN/2-user_guide/7-tooling/4-scene_export.md index 4d785c0c3..fff4b9de8 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/7-tooling/4-scene_export.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/7-tooling/4-scene_export.md @@ -1,6 +1,6 @@ # 场景导出 -场景导出由 `src/unilab/tools/export_scene.py` 实现，并在 `pyproject.toml` 中注册为 `unilab-export-scene` 控制台入口。它接受一个 MuJoCo XML 或 MJB 模型路径，写出 `scene.xml`，在能够发现 mesh asset 时复制它们，并且可以创建一个 zip 归档。 +场景导出由 `src/unilab/base/backend/mujoco/export_scene.py` 实现，并在 `pyproject.toml` 中注册为 `unilab-export-scene` 控制台入口。它接受一个 MuJoCo XML 或 MJB 模型路径，写出 `scene.xml`，在能够发现 mesh asset 时复制它们，并且可以创建一个 zip 归档。 对于 task 级别的实例化检查，请使用从 registry 和 owner config 构造 env 的脚本： diff --git a/docs/sphinx/source/zh_CN/2-user_guide/7-tooling/5-robot_import.md b/docs/sphinx/source/zh_CN/2-user_guide/7-tooling/5-robot_import.md index 08a48e67b..cd3787a30 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/7-tooling/5-robot_import.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/7-tooling/5-robot_import.md @@ -25,7 +25,7 @@ task/reward/env 语义。 如果只有 URDF，使用仓库自带脚本进行转换： ```bash -uv run unilab-import-robot [robot_name] +uv run scripts/tools/import_robot.py [robot_name] ``` ```{important} @@ -58,7 +58,7 @@ visual mesh 作为 collision mesh，尽量把碰撞体简化为 box / capsule / ## 输出产物 -使用 `uv run unilab-import-robot [robot_name]` 转换后，会在仓库内生成： +使用 `uv run scripts/tools/import_robot.py [robot_name]` 转换后，会在仓库内生成： - `src/unilab/assets/robots//assets/`：转换并整理后的 mesh 资产。 - `src/unilab/assets/robots//.xml`：机器人 MJCF 描述，只包含机器人 diff --git a/pyproject.rocm.toml b/pyproject.rocm.toml index 5108794fd..3cff64e32 100644 --- a/pyproject.rocm.toml +++ b/pyproject.rocm.toml @@ -40,12 +40,11 @@ dependencies = [ train = "unilab.cli:train_main" eval = "unilab.cli:eval_main" demo = "unilab.cli:demo_main" -unilab-complete = "unilab.tools.completion:main" -unilab-viz-nan = "unilab.tools.viz_nan:main" -unilab-export-scene = "unilab.tools.export_scene:main" -unilab-render-teaser = "unilab.tools.render_teaser:main" -unilab-import-robot = "unilab.tools.import_robot:main" -unilab-pull-assets = "unilab.tools.pull_assets:main" +unilab-complete = "unilab.cli_completion:main" +unilab-viz-nan = "unilab.utils.nan_viz:main" +unilab-export-scene = "unilab.base.backend.mujoco.export_scene:main" +unilab-render-teaser = "unilab.visualization.teaser:main" +unilab-pull-assets = "unilab.assets.pull:main" [project.optional-dependencies] mujoco = [ diff --git a/pyproject.toml b/pyproject.toml index 7b56b062f..df809ab77 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,12 +41,11 @@ dependencies = [ train = "unilab.cli:train_main" eval = "unilab.cli:eval_main" demo = "unilab.cli:demo_main" -unilab-complete = "unilab.tools.completion:main" -unilab-viz-nan = "unilab.tools.viz_nan:main" -unilab-export-scene = "unilab.tools.export_scene:main" -unilab-render-teaser = "unilab.tools.render_teaser:main" -unilab-import-robot = "unilab.tools.import_robot:main" -unilab-pull-assets = "unilab.tools.pull_assets:main" +unilab-complete = "unilab.cli_completion:main" +unilab-viz-nan = "unilab.utils.nan_viz:main" +unilab-export-scene = "unilab.base.backend.mujoco.export_scene:main" +unilab-render-teaser = "unilab.visualization.teaser:main" +unilab-pull-assets = "unilab.assets.pull:main" [project.optional-dependencies] mujoco = [ diff --git a/scripts/completions/unilab.bash b/scripts/completions/unilab.bash index 024a569b8..b54177023 100644 --- a/scripts/completions/unilab.bash +++ b/scripts/completions/unilab.bash @@ -14,7 +14,7 @@ _unilab_uv_complete() { local candidates if ! mapfile -t candidates < <( uv run --no-sync unilab-complete --cword "$COMP_CWORD" -- "${COMP_WORDS[@]}" 2>/dev/null \ - || PYTHONPATH="$repo_root/src${PYTHONPATH:+:$PYTHONPATH}" uv run --no-sync python -m unilab.tools.completion --cword "$COMP_CWORD" -- "${COMP_WORDS[@]}" 2>/dev/null + || PYTHONPATH="$repo_root/src${PYTHONPATH:+:$PYTHONPATH}" uv run --no-sync python -m unilab.cli_completion --cword "$COMP_CWORD" -- "${COMP_WORDS[@]}" 2>/dev/null ); then return 0 fi diff --git a/scripts/completions/unilab.zsh b/scripts/completions/unilab.zsh index 3432470bf..53189f294 100644 --- a/scripts/completions/unilab.zsh +++ b/scripts/completions/unilab.zsh @@ -13,7 +13,7 @@ _unilab_uv_complete() { local output output="$( uv run --no-sync unilab-complete --cword "$((CURRENT - 1))" -- "${words[@]}" 2>/dev/null \ - || PYTHONPATH="$repo_root/src${PYTHONPATH:+:$PYTHONPATH}" uv run --no-sync python -m unilab.tools.completion --cword "$((CURRENT - 1))" -- "${words[@]}" 2>/dev/null + || PYTHONPATH="$repo_root/src${PYTHONPATH:+:$PYTHONPATH}" uv run --no-sync python -m unilab.cli_completion --cword "$((CURRENT - 1))" -- "${words[@]}" 2>/dev/null )" || return 0 if [[ -z "$output" ]]; then diff --git a/src/unilab/tools/bones_seed_csv.py b/scripts/motion/bones_seed_csv.py similarity index 100% rename from src/unilab/tools/bones_seed_csv.py rename to scripts/motion/bones_seed_csv.py diff --git a/scripts/motion/bones_seed_csv_to_npz.py b/scripts/motion/bones_seed_csv_to_npz.py index ed21defdf..163d9d533 100644 --- a/scripts/motion/bones_seed_csv_to_npz.py +++ b/scripts/motion/bones_seed_csv_to_npz.py @@ -37,17 +37,17 @@ import mujoco import numpy as np -from tqdm import tqdm - -from unilab.assets import ASSETS_ROOT_PATH -from unilab.base.backend.mujoco.xml import inject_mujoco_tracking_sensors -from unilab.tools.bones_seed_csv import ( +from scripts.motion.bones_seed_csv import ( ROOT_COLUMNS, euler_deg_to_quat_wxyz, load_header, parse_joint_names, resolve_input_files, ) +from tqdm import tqdm + +from unilab.assets import ASSETS_ROOT_PATH +from unilab.base.backend.mujoco.xml import inject_mujoco_tracking_sensors from unilab.utils.rotation import np_quat_angular_velocity, np_quat_ensure_continuity DEFAULT_INPUT = "src/unilab/assets/motions/g1/flip" diff --git a/scripts/motion/replay_bones_seed_csv.py b/scripts/motion/replay_bones_seed_csv.py index c791ff8f6..553230026 100644 --- a/scripts/motion/replay_bones_seed_csv.py +++ b/scripts/motion/replay_bones_seed_csv.py @@ -36,9 +36,7 @@ import mujoco import mujoco.viewer import numpy as np - -from unilab.assets import ASSETS_ROOT_PATH -from unilab.tools.bones_seed_csv import ( +from scripts.motion.bones_seed_csv import ( ROOT_COLUMNS, euler_deg_to_quat_wxyz, load_header, @@ -46,6 +44,8 @@ resolve_input_files, ) +from unilab.assets import ASSETS_ROOT_PATH + DEFAULT_INPUT = "src/unilab/assets/motions/g1/flip" diff --git a/src/unilab/tools/import_robot.py b/scripts/tools/import_robot.py similarity index 99% rename from src/unilab/tools/import_robot.py rename to scripts/tools/import_robot.py index 5a46ffe21..e5fa0549e 100644 --- a/src/unilab/tools/import_robot.py +++ b/scripts/tools/import_robot.py @@ -2,7 +2,7 @@ """Convert a URDF robot to a UniLab robot MJCF asset directory. Usage: - uv run unilab-import-robot [robot_name] + uv run scripts/tools/import_robot.py [robot_name] """ @@ -16,7 +16,7 @@ from pathlib import Path from typing import Any, Iterable, Sequence, cast -REPO_ROOT = Path(__file__).resolve().parents[3] +REPO_ROOT = Path(__file__).resolve().parents[2] ROBOT_ASSET_ROOT = REPO_ROOT / "src" / "unilab" / "assets" / "robots" TEMP_MESH_PREFIX = "meshes/meshes/" DEFAULT_MATERIAL = "default_material" diff --git a/src/unilab/tools/pull_assets.py b/src/unilab/assets/pull.py similarity index 100% rename from src/unilab/tools/pull_assets.py rename to src/unilab/assets/pull.py diff --git a/src/unilab/tools/export_scene.py b/src/unilab/base/backend/mujoco/export_scene.py similarity index 100% rename from src/unilab/tools/export_scene.py rename to src/unilab/base/backend/mujoco/export_scene.py diff --git a/src/unilab/tools/completion.py b/src/unilab/cli_completion.py similarity index 100% rename from src/unilab/tools/completion.py rename to src/unilab/cli_completion.py diff --git a/src/unilab/demo.py b/src/unilab/demo.py index e3f31721b..ef90c211f 100644 --- a/src/unilab/demo.py +++ b/src/unilab/demo.py @@ -238,13 +238,13 @@ def _run_teaser_demo() -> int: if platform.system() == "Darwin" and Path(sys.executable).name != "mxpython": command = [ _mxpython_executable(), - str(_repo_root() / "src" / "unilab" / "tools" / "render_teaser.py"), + str(_repo_root() / "src" / "unilab" / "visualization" / "teaser.py"), ] env = os.environ.copy() env["UV_PROJECT_ENVIRONMENT"] = str(_repo_root() / ".venv") return subprocess.run(command, check=False, env=env).returncode - from unilab.tools.render_teaser import main as render_teaser_main + from unilab.visualization.teaser import main as render_teaser_main render_teaser_main() return 0 diff --git a/src/unilab/tools/__init__.py b/src/unilab/tools/__init__.py deleted file mode 100644 index 5961d279a..000000000 --- a/src/unilab/tools/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Debug and export tools for UniLab.""" diff --git a/src/unilab/tools/viz_nan.py b/src/unilab/utils/nan_viz.py similarity index 100% rename from src/unilab/tools/viz_nan.py rename to src/unilab/utils/nan_viz.py diff --git a/src/unilab/tools/render_teaser.py b/src/unilab/visualization/teaser.py similarity index 100% rename from src/unilab/tools/render_teaser.py rename to src/unilab/visualization/teaser.py diff --git a/tests/scripts/test_mujoco_only_tooling_markers.py b/tests/scripts/test_mujoco_only_tooling_markers.py index f94878b7c..98bc53c00 100644 --- a/tests/scripts/test_mujoco_only_tooling_markers.py +++ b/tests/scripts/test_mujoco_only_tooling_markers.py @@ -11,7 +11,7 @@ def test_mujoco_only_tooling_is_marked_explicitly(): root / "scripts" / "motion" / "replay_npz.py", root / "scripts" / "motion" / "bones_seed_csv_to_npz.py", root / "scripts" / "motion" / "replay_bones_seed_csv.py", - root / "src" / "unilab" / "tools" / "bones_seed_csv.py", + root / "scripts" / "motion" / "bones_seed_csv.py", root / "src" / "unilab" / "visualization" / "render_many.py", ] diff --git a/tests/test_bones_seed_csv.py b/tests/test_bones_seed_csv.py index cf13cf488..cda4b625f 100644 --- a/tests/test_bones_seed_csv.py +++ b/tests/test_bones_seed_csv.py @@ -8,9 +8,8 @@ import numpy as np import pytest - -from unilab.tools import bones_seed_csv -from unilab.tools.bones_seed_csv import ( +from scripts.motion import bones_seed_csv +from scripts.motion.bones_seed_csv import ( ROOT_COLUMNS, euler_deg_to_quat_wxyz, load_header, diff --git a/tests/test_cli.py b/tests/test_cli.py index 420935ee2..0bdb7b72c 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -601,7 +601,7 @@ def test_demo_teaser_run_demo_invokes_render_teaser_main( def fake_render_teaser_main() -> None: called.append("rendered") - import unilab.tools.render_teaser as render_teaser_module + import unilab.visualization.teaser as render_teaser_module monkeypatch.setattr(render_teaser_module, "main", fake_render_teaser_main) @@ -633,7 +633,7 @@ def fake_run(command: list[str], *, check: bool, env: dict[str, str]) -> SimpleN def fail_render_teaser_main() -> None: raise AssertionError("macOS teaser must route through mxpython") - import unilab.tools.render_teaser as render_teaser_module + import unilab.visualization.teaser as render_teaser_module monkeypatch.setattr(render_teaser_module, "main", fail_render_teaser_main) @@ -643,7 +643,7 @@ def fail_render_teaser_main() -> None: command, env = calls[0] assert command == [ "/tmp/unilab/.venv/bin/mxpython", - str(demo._repo_root() / "src" / "unilab" / "tools" / "render_teaser.py"), + str(demo._repo_root() / "src" / "unilab" / "visualization" / "teaser.py"), ] assert env["UV_PROJECT_ENVIRONMENT"] == str(demo._repo_root() / ".venv") @@ -657,7 +657,7 @@ def test_demo_main_teaser_dispatches_to_render_teaser( def fake_render_teaser_main() -> None: called.append("rendered") - import unilab.tools.render_teaser as render_teaser_module + import unilab.visualization.teaser as render_teaser_module monkeypatch.setattr(render_teaser_module, "main", fake_render_teaser_main) rc = cli.demo_main(["teaser"]) diff --git a/tests/test_completion.py b/tests/test_completion.py index 26ccf2c25..87cbccb63 100644 --- a/tests/test_completion.py +++ b/tests/test_completion.py @@ -2,7 +2,7 @@ from pathlib import Path -from unilab.tools.completion import ( +from unilab.cli_completion import ( COMPLETION_BLOCK_END, COMPLETION_BLOCK_START, build_metadata, diff --git a/tests/test_export_scene.py b/tests/test_export_scene.py index 7691f65fe..f7e38b50a 100644 --- a/tests/test_export_scene.py +++ b/tests/test_export_scene.py @@ -9,7 +9,7 @@ pytest.importorskip("mujoco") -from unilab.tools.export_scene import export_scene, main +from unilab.base.backend.mujoco.export_scene import export_scene, main MINIMAL_XML = """\ diff --git a/tests/test_import_robot.py b/tests/test_import_robot.py index 9aaa48337..afd8865cd 100644 --- a/tests/test_import_robot.py +++ b/tests/test_import_robot.py @@ -9,7 +9,7 @@ def _load_script() -> Any: - return importlib.import_module("unilab.tools.import_robot") + return importlib.import_module("scripts.tools.import_robot") def test_robot_name_defaults_to_urdf_stem(tmp_path: Path) -> None: diff --git a/tests/test_nan_guard.py b/tests/test_nan_guard.py index 2a491a20a..0d94952d2 100644 --- a/tests/test_nan_guard.py +++ b/tests/test_nan_guard.py @@ -7,8 +7,8 @@ import numpy as np import pytest -from unilab.tools.viz_nan import load_dump from unilab.utils.nan_guard import NanGuard, NanGuardCfg +from unilab.utils.nan_viz import load_dump NUM_ENVS = 4 OBS_DIM = 10 diff --git a/tests/test_render_teaser.py b/tests/test_render_teaser.py index 30d204d11..d21500d21 100644 --- a/tests/test_render_teaser.py +++ b/tests/test_render_teaser.py @@ -1,11 +1,11 @@ -"""Tests for the teaser scene renderer (``unilab.tools.render_teaser``).""" +"""Tests for the teaser scene renderer (``unilab.visualization.teaser``).""" from __future__ import annotations import inspect from unittest.mock import MagicMock, patch -from unilab.tools import render_teaser +from unilab.visualization import teaser as render_teaser def test_set_teaser_system_camera_view_uses_packaged_view(): diff --git a/tests/utils/test_utils_package_policy.py b/tests/utils/test_utils_package_policy.py index a126dddb1..8030ed238 100644 --- a/tests/utils/test_utils_package_policy.py +++ b/tests/utils/test_utils_package_policy.py @@ -9,6 +9,7 @@ "device", "geometry", "nan_guard", + "nan_viz", "rotation", "support_matrix", "tensor", From 9b0c1aac232b7029b73b7549604de6eea494461e Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Thu, 20 Aug 2026 16:17:10 +0800 Subject: [PATCH 094/169] feat: migrate motion tracking core to manager api (#1235) --- .../task/g1_23dof_motion_tracking/motrix.yaml | 29 +- .../task/g1_23dof_motion_tracking/mujoco.yaml | 91 +- conf/appo/task/g1_motion_tracking/motrix.yaml | 30 +- conf/appo/task/g1_motion_tracking/mujoco.yaml | 239 +++++- .../sac/g1_23dof_motion_tracking/motrix.yaml | 12 +- .../sac/g1_23dof_motion_tracking/mujoco.yaml | 139 ++-- .../task/sac/g1_motion_tracking/motrix.yaml | 12 +- .../task/sac/g1_motion_tracking/mujoco.yaml | 271 +++++- .../task/g1_23dof_motion_tracking/motrix.yaml | 71 +- .../task/g1_23dof_motion_tracking/mujoco.yaml | 93 ++- .../motrix.yaml | 123 +-- .../mujoco.yaml | 152 ++-- conf/ppo/task/g1_motion_tracking/motrix.yaml | 72 +- conf/ppo/task/g1_motion_tracking/mujoco.yaml | 243 +++++- .../g1_motion_tracking_deploy/motrix.yaml | 131 +-- .../g1_motion_tracking_deploy/mujoco.yaml | 163 ++-- .../benchmark/benchmark_drake_performance.py | 23 +- scripts/benchmark/env/benchmark_env_step.py | 14 +- src/unilab/base/backend/motrix/backend.py | 15 +- src/unilab/base/entity.py | 60 ++ src/unilab/envs/manager_based_rl_env.py | 13 +- src/unilab/managers/_types.py | 8 + src/unilab/managers/command_manager.py | 10 + .../tasks/manipulation/stewart/balance.py | 4 +- src/unilab/tasks/migration_matrix.py | 33 +- .../motion_tracking/common/manager_terms.py | 776 ++++++++++++++++++ .../motion_tracking/common/motion_loader.py | 53 +- .../tasks/motion_tracking/g1/tracking.py | 31 +- .../tasks/motion_tracking/g1/tracking_sac.py | 13 +- src/unilab/training/sim2sim.py | 35 +- tests/base/test_entity_facade.py | 12 + tests/base/test_sim_backend.py | 5 +- tests/base/test_sim_backend_smoke.py | 4 +- ...st_offpolicy_collector_active_benchmark.py | 8 +- tests/config/test_locomotion_params.py | 40 +- tests/envs/mdp/test_events.py | 3 + tests/envs/mdp/test_joint_position_action.py | 3 + tests/envs/mdp/test_observations.py | 3 + tests/envs/test_env_configs.py | 299 +++++-- tests/envs/test_manager_based_rl_env.py | 45 + tests/envs/test_motion_loader.py | 24 + tests/scripts/test_train_scripts.py | 55 +- tests/tasks/test_migration_matrix.py | 4 +- tests/training/test_sim2sim_resolver.py | 26 +- tests/training/test_training_helpers.py | 11 +- 45 files changed, 2584 insertions(+), 917 deletions(-) create mode 100644 src/unilab/tasks/motion_tracking/common/manager_terms.py diff --git a/conf/appo/task/g1_23dof_motion_tracking/motrix.yaml b/conf/appo/task/g1_23dof_motion_tracking/motrix.yaml index aa74c1a66..9d8ecf00d 100644 --- a/conf/appo/task/g1_23dof_motion_tracking/motrix.yaml +++ b/conf/appo/task/g1_23dof_motion_tracking/motrix.yaml @@ -1,29 +1,8 @@ # @package _global_ +defaults: + - /task/g1_23dof_motion_tracking/mujoco + - _self_ + training: task_name: G1MotionTracking23Dof sim_backend: motrix - play_steps: 1000 -algo: - num_envs: 1024 - max_iterations: 5000 - save_interval: 500 -reward: - scales: - motion_global_root_pos: 0.5 - motion_global_root_ori: 0.5 - motion_body_pos: 1.0 - motion_body_ori: 1.0 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_joint_pos: 0.0 - motion_joint_vel: 0.0 - action_rate_l2: -0.1 - joint_limit: -10.0 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 diff --git a/conf/appo/task/g1_23dof_motion_tracking/mujoco.yaml b/conf/appo/task/g1_23dof_motion_tracking/mujoco.yaml index 3b31907a5..8605feb92 100644 --- a/conf/appo/task/g1_23dof_motion_tracking/mujoco.yaml +++ b/conf/appo/task/g1_23dof_motion_tracking/mujoco.yaml @@ -1,32 +1,67 @@ # @package _global_ +defaults: + - /task/g1_motion_tracking/mujoco + - _self_ + training: task_name: G1MotionTracking23Dof sim_backend: mujoco - play_steps: 1000 -algo: - num_envs: 1024 - max_iterations: 5000 - save_interval: 500 - algorithm: - adaptive_kl_factor: 2.0 - adaptive_lr_factor: 1.5 -reward: - scales: - motion_global_root_pos: 0.5 - motion_global_root_ori: 0.5 - motion_body_pos: 1.0 - motion_body_ori: 1.0 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_joint_pos: 0.0 - motion_joint_vel: 0.0 - action_rate_l2: -0.1 - joint_limit: -10.0 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 + +env: + scene: + model_file: src/unilab/assets/robots/g1/scene_flat_23dof.xml + entities: + robot: + joint_names: &g1_23dof_joints + - left_hip_pitch_joint + - left_hip_roll_joint + - left_hip_yaw_joint + - left_knee_joint + - left_ankle_pitch_joint + - left_ankle_roll_joint + - right_hip_pitch_joint + - right_hip_roll_joint + - right_hip_yaw_joint + - right_knee_joint + - right_ankle_pitch_joint + - right_ankle_roll_joint + - waist_yaw_joint + - left_shoulder_pitch_joint + - left_shoulder_roll_joint + - left_shoulder_yaw_joint + - left_elbow_joint + - left_wrist_roll_joint + - right_shoulder_pitch_joint + - right_shoulder_roll_joint + - right_shoulder_yaw_joint + - right_elbow_joint + - right_wrist_roll_joint + actuator_names: *g1_23dof_joints + body_names: &tracked_bodies_23dof + - pelvis + - left_hip_roll_link + - left_knee_link + - left_ankle_roll_link + - right_hip_roll_link + - right_knee_link + - right_ankle_roll_link + - torso_link + - left_shoulder_roll_link + - left_elbow_link + - left_wrist_roll_rubber_hand + - right_shoulder_roll_link + - right_elbow_link + - right_wrist_roll_rubber_hand + commands: + motion: + params: + motion_file: motions/g1/dance1_subject2_part_23dof.npz + body_names: *tracked_bodies_23dof + terminations: + ee_body_pos: + params: + body_names: + - left_ankle_roll_link + - right_ankle_roll_link + - left_wrist_roll_rubber_hand + - right_wrist_roll_rubber_hand diff --git a/conf/appo/task/g1_motion_tracking/motrix.yaml b/conf/appo/task/g1_motion_tracking/motrix.yaml index eabea1ad6..22bbdadfa 100644 --- a/conf/appo/task/g1_motion_tracking/motrix.yaml +++ b/conf/appo/task/g1_motion_tracking/motrix.yaml @@ -1,30 +1,8 @@ # @package _global_ +defaults: + - /task/g1_motion_tracking/mujoco + - _self_ + training: task_name: G1MotionTracking sim_backend: motrix - play_steps: 1000 -algo: - num_envs: 1024 - max_iterations: 5000 - save_interval: 500 -env: -reward: - scales: - motion_global_root_pos: 0.5 - motion_global_root_ori: 0.5 - motion_body_pos: 1.0 - motion_body_ori: 1.0 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_joint_pos: 0.0 - motion_joint_vel: 0.0 - action_rate_l2: -0.1 - joint_limit: -10.0 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 diff --git a/conf/appo/task/g1_motion_tracking/mujoco.yaml b/conf/appo/task/g1_motion_tracking/mujoco.yaml index ec978a19c..b1898ef6d 100644 --- a/conf/appo/task/g1_motion_tracking/mujoco.yaml +++ b/conf/appo/task/g1_motion_tracking/mujoco.yaml @@ -3,6 +3,7 @@ training: task_name: G1MotionTracking sim_backend: mujoco play_steps: 1000 + algo: num_envs: 1024 max_iterations: 5000 @@ -10,24 +11,224 @@ algo: algorithm: adaptive_kl_factor: 2.0 adaptive_lr_factor: 1.5 + env: + scene: + model_file: src/unilab/assets/robots/g1/scene_flat.xml + default_keyframe_name: stand + entities: + robot: + root_body_name: pelvis + joint_names: &g1_joints + - left_hip_pitch_joint + - left_hip_roll_joint + - left_hip_yaw_joint + - left_knee_joint + - left_ankle_pitch_joint + - left_ankle_roll_joint + - right_hip_pitch_joint + - right_hip_roll_joint + - right_hip_yaw_joint + - right_knee_joint + - right_ankle_pitch_joint + - right_ankle_roll_joint + - waist_yaw_joint + - waist_roll_joint + - waist_pitch_joint + - left_shoulder_pitch_joint + - left_shoulder_roll_joint + - left_shoulder_yaw_joint + - left_elbow_joint + - left_wrist_roll_joint + - left_wrist_pitch_joint + - left_wrist_yaw_joint + - right_shoulder_pitch_joint + - right_shoulder_roll_joint + - right_shoulder_yaw_joint + - right_elbow_joint + - right_wrist_roll_joint + - right_wrist_pitch_joint + - right_wrist_yaw_joint + actuator_names: *g1_joints + body_names: &tracked_bodies + - pelvis + - left_hip_roll_link + - left_knee_link + - left_ankle_roll_link + - right_hip_roll_link + - right_knee_link + - right_ankle_roll_link + - torso_link + - left_shoulder_roll_link + - left_elbow_link + - left_wrist_yaw_link + - right_shoulder_roll_link + - right_elbow_link + - right_wrist_yaw_link + geom_names: + - left_foot1_collision + - left_foot2_collision + - left_foot3_collision + - left_foot4_collision + - left_foot5_collision + - left_foot6_collision + - left_foot7_collision + - right_foot1_collision + - right_foot2_collision + - right_foot3_collision + - right_foot4_collision + - right_foot5_collision + - right_foot6_collision + - right_foot7_collision + sim_dt: 0.006666666666666667 + ctrl_dt: 0.02 + max_episode_seconds: 10.0 + observations: + actor: + terms: + command: &command_obs + func: unilab.envs.mdp.generated_commands + params: {command_name: motion} + motion_anchor_pos_b: &anchor_pos_obs + func: unilab.tasks.motion_tracking.common.manager_terms.motion_anchor_pos_b + params: {command_name: motion} + motion_anchor_ori_b: &anchor_ori_obs + func: unilab.tasks.motion_tracking.common.manager_terms.motion_anchor_ori_b + params: {command_name: motion} + base_lin_vel: &base_lin_vel_obs + func: unilab.envs.mdp.builtin_sensor + params: {sensor_name: pelvis_local_linvel} + base_ang_vel: &base_ang_vel_obs + func: unilab.envs.mdp.builtin_sensor + params: {sensor_name: torso_gyro} + joint_pos: &joint_pos_obs + func: unilab.tasks.motion_tracking.common.manager_terms.motion_joint_pos_rel + params: {command_name: motion} + joint_vel: &joint_vel_obs + func: unilab.envs.mdp.joint_vel_rel + actions: &actions_obs + func: unilab.envs.mdp.last_action + critic: + terms: + command: *command_obs + motion_anchor_pos_b: *anchor_pos_obs + motion_anchor_ori_b: *anchor_ori_obs + base_lin_vel: *base_lin_vel_obs + base_ang_vel: *base_ang_vel_obs + joint_pos: *joint_pos_obs + joint_vel: *joint_vel_obs + actions: *actions_obs + body_pos: + func: unilab.tasks.motion_tracking.common.manager_terms.robot_body_pos_b + params: {command_name: motion} + body_ori: + func: unilab.tasks.motion_tracking.common.manager_terms.robot_body_ori_b + params: {command_name: motion} + actions: + joint_pos: + _target_: unilab.tasks.motion_tracking.common.manager_terms.MotionJointPositionActionCfg + entity_name: robot + actuator_names: [".*"] + scale: 0.25 + use_default_offset: true + command_name: motion + commands: + motion: + _target_: unilab.tasks.motion_tracking.common.manager_terms.MotionCommandCfg + entity_name: robot + resampling_time_range: [1.0e9, 1.0e9] + params: + motion_file: motions/g1/dance1_subject2_part.npz + anchor_body_name: torso_link + body_names: *tracked_bodies + sampling_mode: adaptive + sampling_start_ratio: 0.0 + truncate_on_clip_end: false + pose_range: + x: [-0.05, 0.05] + y: [-0.05, 0.05] + z: [-0.01, 0.01] + roll: [-0.1, 0.1] + pitch: [-0.1, 0.1] + yaw: [-0.2, 0.2] + velocity_range: + x: [-0.5, 0.5] + y: [-0.5, 0.5] + z: [-0.2, 0.2] + roll: [-0.52, 0.52] + pitch: [-0.52, 0.52] + yaw: [-0.78, 0.78] + joint_position_range: [-0.1, 0.1] + joint_default_position_range: [0.0, 0.0] + terminations: + time_out: + func: unilab.envs.mdp.time_out + time_out: true + anchor_pos: + func: unilab.tasks.motion_tracking.common.manager_terms.bad_anchor_pos_z_only + params: {command_name: motion, threshold: 0.25} + anchor_ori: + func: unilab.tasks.motion_tracking.common.manager_terms.bad_anchor_ori + params: + command_name: motion + threshold: 0.8 + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + ee_body_pos: + func: unilab.tasks.motion_tracking.common.manager_terms.bad_motion_body_pos_z_only + params: + command_name: motion + threshold: 0.25 + body_names: + - left_ankle_roll_link + - right_ankle_roll_link + - left_wrist_yaw_link + - right_wrist_yaw_link + policy_observation_group: actor + critic_observation_group: critic + reward: - scales: - motion_global_root_pos: 0.5 - motion_global_root_ori: 0.5 - motion_body_pos: 1.0 - motion_body_ori: 1.0 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_joint_pos: 0.0 - motion_joint_vel: 0.0 - action_rate_l2: -0.1 - joint_limit: -10.0 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 + motion_global_root_pos: + func: unilab.tasks.motion_tracking.common.manager_terms.motion_global_anchor_position_error_exp + weight: 0.5 + params: {command_name: motion, std: 0.3} + motion_global_root_ori: + func: unilab.tasks.motion_tracking.common.manager_terms.motion_global_anchor_orientation_error_exp + weight: 0.5 + params: {command_name: motion, std: 0.4} + motion_body_pos: + func: unilab.tasks.motion_tracking.common.manager_terms.motion_relative_body_position_error_exp + weight: 1.0 + params: {command_name: motion, std: 0.3} + motion_body_ori: + func: unilab.tasks.motion_tracking.common.manager_terms.motion_relative_body_orientation_error_exp + weight: 1.0 + params: {command_name: motion, std: 0.4} + motion_body_lin_vel: + func: unilab.tasks.motion_tracking.common.manager_terms.motion_global_body_linear_velocity_error_exp + weight: 1.0 + params: {command_name: motion, std: 1.0} + motion_body_ang_vel: + func: unilab.tasks.motion_tracking.common.manager_terms.motion_global_body_angular_velocity_error_exp + weight: 1.0 + params: {command_name: motion, std: 3.14} + motion_joint_pos: + func: unilab.tasks.motion_tracking.common.manager_terms.motion_joint_position_error_exp + weight: 0.0 + params: {command_name: motion, std: 0.2} + motion_joint_vel: + func: unilab.tasks.motion_tracking.common.manager_terms.motion_joint_velocity_error_exp + weight: 0.0 + params: {command_name: motion, std: 1.0} + action_rate_l2: + func: unilab.envs.mdp.action_rate_l2 + weight: -0.1 + joint_limit: + func: unilab.tasks.motion_tracking.common.manager_terms.joint_pos_limits + weight: -10.0 + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + joint_names: ".*" diff --git a/conf/offpolicy/task/sac/g1_23dof_motion_tracking/motrix.yaml b/conf/offpolicy/task/sac/g1_23dof_motion_tracking/motrix.yaml index e042ee799..814b02db3 100644 --- a/conf/offpolicy/task/sac/g1_23dof_motion_tracking/motrix.yaml +++ b/conf/offpolicy/task/sac/g1_23dof_motion_tracking/motrix.yaml @@ -1,8 +1,5 @@ # @package _global_ -# G1 23-DoF Motion Tracking SAC — Motrix variant for sim2sim eval. -# Inherits the mujoco training config in full and only switches the rendering -# backend so checkpoints trained on mujoco can be replayed via motrix's native -# renderer (`eval --sim motrix`). Training on motrix is not the intended path. +# Motrix is the sim2sim eval owner for MuJoCo-trained WBT checkpoints. defaults: - /task/sac/g1_23dof_motion_tracking/mujoco - _self_ @@ -10,10 +7,3 @@ defaults: training: task_name: G1MotionTrackingSAC23Dof sim_backend: motrix -env: - # motrix backend's kp/kd override path is broken on column slices, and DR - # is not desirable during deterministic sim2sim eval anyway. Match the - # `g1_walk_flat/motrix.yaml` convention by switching them off. - domain_rand: - randomize_kp: false - randomize_kd: false diff --git a/conf/offpolicy/task/sac/g1_23dof_motion_tracking/mujoco.yaml b/conf/offpolicy/task/sac/g1_23dof_motion_tracking/mujoco.yaml index a68d07fae..28a2cbe17 100644 --- a/conf/offpolicy/task/sac/g1_23dof_motion_tracking/mujoco.yaml +++ b/conf/offpolicy/task/sac/g1_23dof_motion_tracking/mujoco.yaml @@ -1,73 +1,76 @@ # @package _global_ +defaults: + - /task/sac/g1_motion_tracking/mujoco + - _self_ + training: task_name: G1MotionTrackingSAC23Dof sim_backend: mujoco -algo: - num_envs: 2048 - max_iterations: 25000 - save_interval: 1000 - gamma: 0.99 - tau: 0.05 - num_atoms: 501 - updates_per_step: 4 - policy_frequency: 2 - use_symmetry: false - algo_params: - alpha_init: 0.1 - target_entropy_ratio: 0.5 - max_grad_norm: 10.0 + env: - control_config: - action_scale: - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.07450087032950714 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.07450087032950714 - anchor_pos_z_threshold: 0.5 - ee_body_pos_z_threshold: 0.5 - truncate_on_clip_end: true - noise_config: - level: 1.0 - scale_joint_angle: 0.01 - scale_joint_vel: 1.5 - scale_gyro: 0.2 -reward: - scales: - motion_global_root_pos: 1.0 - motion_global_root_ori: 0.5 - motion_body_pos: 2.0 - motion_body_ori: 1.0 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_joint_pos: 0.0 - motion_joint_vel: 0.0 - action_rate_l2: -0.1 - joint_limit: -2.0 - undesired_contacts: -0.1 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 + scene: + model_file: src/unilab/assets/robots/g1/scene_flat_23dof.xml + entities: + robot: + joint_names: &g1_23dof_joints + - left_hip_pitch_joint + - left_hip_roll_joint + - left_hip_yaw_joint + - left_knee_joint + - left_ankle_pitch_joint + - left_ankle_roll_joint + - right_hip_pitch_joint + - right_hip_roll_joint + - right_hip_yaw_joint + - right_knee_joint + - right_ankle_pitch_joint + - right_ankle_roll_joint + - waist_yaw_joint + - left_shoulder_pitch_joint + - left_shoulder_roll_joint + - left_shoulder_yaw_joint + - left_elbow_joint + - left_wrist_roll_joint + - right_shoulder_pitch_joint + - right_shoulder_roll_joint + - right_shoulder_yaw_joint + - right_elbow_joint + - right_wrist_roll_joint + actuator_names: *g1_23dof_joints + body_names: &tracked_bodies_23dof + - pelvis + - left_hip_roll_link + - left_knee_link + - left_ankle_roll_link + - right_hip_roll_link + - right_knee_link + - right_ankle_roll_link + - torso_link + - left_shoulder_roll_link + - left_elbow_link + - left_wrist_roll_rubber_hand + - right_shoulder_roll_link + - right_elbow_link + - right_wrist_roll_rubber_hand + actions: + joint_pos: + scale: + ".*_(hip_pitch|hip_yaw)_joint": 0.5475464629911068 + ".*_(hip_roll|knee)_joint": 0.35066146637882434 + ".*_ankle_(pitch|roll)_joint": 0.43857731392336724 + "waist_yaw_joint": 0.5475464629911068 + ".*_(shoulder_(pitch|roll|yaw)|elbow)_joint": 0.43857731392336724 + ".*_wrist_roll_joint": 0.07450087032950714 + commands: + motion: + params: + motion_file: motions/g1/dance1_subject2_part_23dof.npz + body_names: *tracked_bodies_23dof + terminations: + ee_body_pos: + params: + body_names: + - left_ankle_roll_link + - right_ankle_roll_link + - left_wrist_roll_rubber_hand + - right_wrist_roll_rubber_hand diff --git a/conf/offpolicy/task/sac/g1_motion_tracking/motrix.yaml b/conf/offpolicy/task/sac/g1_motion_tracking/motrix.yaml index eddb1839f..4777d3098 100644 --- a/conf/offpolicy/task/sac/g1_motion_tracking/motrix.yaml +++ b/conf/offpolicy/task/sac/g1_motion_tracking/motrix.yaml @@ -1,8 +1,5 @@ # @package _global_ -# G1 Whole-Body Tracking (WBT) FastSAC — Motrix variant for sim2sim eval. -# Inherits the mujoco training config in full and only switches the rendering -# backend so checkpoints trained on mujoco can be replayed via motrix's native -# renderer (`eval --sim motrix`). Training on motrix is not the intended path. +# Motrix is the sim2sim eval owner for MuJoCo-trained WBT checkpoints. defaults: - /task/sac/g1_motion_tracking/mujoco - _self_ @@ -10,10 +7,3 @@ defaults: training: task_name: G1MotionTrackingSAC sim_backend: motrix -env: - # motrix backend's kp/kd override path is broken on column slices, and DR - # is not desirable during deterministic sim2sim eval anyway. Match the - # `g1_walk_flat/motrix.yaml` convention by switching them off. - domain_rand: - randomize_kp: false - randomize_kd: false diff --git a/conf/offpolicy/task/sac/g1_motion_tracking/mujoco.yaml b/conf/offpolicy/task/sac/g1_motion_tracking/mujoco.yaml index 598feccbb..c5d4506a1 100644 --- a/conf/offpolicy/task/sac/g1_motion_tracking/mujoco.yaml +++ b/conf/offpolicy/task/sac/g1_motion_tracking/mujoco.yaml @@ -4,11 +4,11 @@ training: task_name: G1MotionTrackingSAC sim_backend: mujoco + algo: num_envs: 2048 max_iterations: 25000 save_interval: 1000 - # --- holosoma WBT-specific overrides (vs sac.yaml defaults) --- gamma: 0.99 tau: 0.05 num_atoms: 501 @@ -19,28 +19,251 @@ algo: alpha_init: 0.1 target_entropy_ratio: 0.5 max_grad_norm: 10.0 + env: - control_config: - action_scale: 2.0 - anchor_pos_z_threshold: 0.5 - ee_body_pos_z_threshold: 0.5 - truncate_on_clip_end: true - noise_config: - level: 1.0 - scale_joint_angle: 0.01 - scale_joint_vel: 1.5 - scale_gyro: 0.2 - seed: null + seed: null + scene: + model_file: src/unilab/assets/robots/g1/scene_flat.xml + default_keyframe_name: stand + entities: + robot: + root_body_name: pelvis + joint_names: &g1_joints + - left_hip_pitch_joint + - left_hip_roll_joint + - left_hip_yaw_joint + - left_knee_joint + - left_ankle_pitch_joint + - left_ankle_roll_joint + - right_hip_pitch_joint + - right_hip_roll_joint + - right_hip_yaw_joint + - right_knee_joint + - right_ankle_pitch_joint + - right_ankle_roll_joint + - waist_yaw_joint + - waist_roll_joint + - waist_pitch_joint + - left_shoulder_pitch_joint + - left_shoulder_roll_joint + - left_shoulder_yaw_joint + - left_elbow_joint + - left_wrist_roll_joint + - left_wrist_pitch_joint + - left_wrist_yaw_joint + - right_shoulder_pitch_joint + - right_shoulder_roll_joint + - right_shoulder_yaw_joint + - right_elbow_joint + - right_wrist_roll_joint + - right_wrist_pitch_joint + - right_wrist_yaw_joint + actuator_names: *g1_joints + body_names: &tracked_bodies + - pelvis + - left_hip_roll_link + - left_knee_link + - left_ankle_roll_link + - right_hip_roll_link + - right_knee_link + - right_ankle_roll_link + - torso_link + - left_shoulder_roll_link + - left_elbow_link + - left_wrist_yaw_link + - right_shoulder_roll_link + - right_elbow_link + - right_wrist_yaw_link + sim_dt: 0.006666666666666667 + ctrl_dt: 0.02 + max_episode_seconds: 10.0 + observations: + actor: + enable_corruption: true + terms: + command: &command_obs + func: unilab.envs.mdp.generated_commands + params: {command_name: motion} + motion_anchor_pos_b: &anchor_pos_obs + func: unilab.tasks.motion_tracking.common.manager_terms.motion_anchor_pos_b + params: {command_name: motion} + motion_anchor_ori_b: &anchor_ori_obs + func: unilab.tasks.motion_tracking.common.manager_terms.motion_anchor_ori_b + params: {command_name: motion} + base_lin_vel: &base_lin_vel_obs + func: unilab.envs.mdp.builtin_sensor + params: {sensor_name: pelvis_local_linvel} + noise: + _target_: unilab.managers._noise.UniformNoiseCfg + n_min: -0.1 + n_max: 0.1 + base_ang_vel: &base_ang_vel_obs + func: unilab.envs.mdp.builtin_sensor + params: {sensor_name: torso_gyro} + noise: + _target_: unilab.managers._noise.UniformNoiseCfg + n_min: -0.2 + n_max: 0.2 + joint_pos: &joint_pos_obs + func: unilab.tasks.motion_tracking.common.manager_terms.motion_joint_pos_rel + params: {command_name: motion} + noise: + _target_: unilab.managers._noise.UniformNoiseCfg + n_min: -0.01 + n_max: 0.01 + joint_vel: &joint_vel_obs + func: unilab.envs.mdp.joint_vel_rel + noise: + _target_: unilab.managers._noise.UniformNoiseCfg + n_min: -1.5 + n_max: 1.5 + actions: &actions_obs + func: unilab.envs.mdp.last_action + critic: + terms: + command: *command_obs + motion_anchor_pos_b: *anchor_pos_obs + motion_anchor_ori_b: *anchor_ori_obs + base_lin_vel: *base_lin_vel_obs + base_ang_vel: *base_ang_vel_obs + joint_pos: *joint_pos_obs + joint_vel: *joint_vel_obs + actions: *actions_obs + body_pos: + func: unilab.tasks.motion_tracking.common.manager_terms.robot_body_pos_b + params: {command_name: motion} + body_ori: + func: unilab.tasks.motion_tracking.common.manager_terms.robot_body_ori_b + params: {command_name: motion} + sac_base_lin_vel: + func: unilab.envs.mdp.builtin_sensor + params: {sensor_name: pelvis_local_linvel} + actions: + joint_pos: + _target_: unilab.tasks.motion_tracking.common.manager_terms.MotionJointPositionActionCfg + entity_name: robot + actuator_names: [".*"] + scale: 2.0 + use_default_offset: true + command_name: motion + commands: + motion: + _target_: unilab.tasks.motion_tracking.common.manager_terms.MotionCommandCfg + entity_name: robot + resampling_time_range: [1.0e9, 1.0e9] + params: + motion_file: motions/g1/dance1_subject2_part.npz + anchor_body_name: torso_link + body_names: *tracked_bodies + sampling_mode: adaptive + sampling_start_ratio: 0.0 + truncate_on_clip_end: true + pose_range: + x: [-0.05, 0.05] + y: [-0.05, 0.05] + z: [-0.01, 0.01] + roll: [-0.1, 0.1] + pitch: [-0.1, 0.1] + yaw: [-0.2, 0.2] + velocity_range: + x: [-0.5, 0.5] + y: [-0.5, 0.5] + z: [-0.2, 0.2] + roll: [-0.52, 0.52] + pitch: [-0.52, 0.52] + yaw: [-0.78, 0.78] + joint_position_range: [-0.1, 0.1] + joint_default_position_range: [0.0, 0.0] + terminations: + time_out: + func: unilab.envs.mdp.time_out + time_out: true + motion_clip_end: + func: unilab.tasks.motion_tracking.common.manager_terms.motion_clip_end + time_out: true + params: {command_name: motion} + anchor_pos: + func: unilab.tasks.motion_tracking.common.manager_terms.bad_anchor_pos_z_only + params: {command_name: motion, threshold: 0.5} + anchor_ori: + func: unilab.tasks.motion_tracking.common.manager_terms.bad_anchor_ori + params: + command_name: motion + threshold: 0.8 + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + ee_body_pos: + func: unilab.tasks.motion_tracking.common.manager_terms.bad_motion_body_pos_z_only + params: + command_name: motion + threshold: 0.5 + body_names: + - left_ankle_roll_link + - right_ankle_roll_link + - left_wrist_yaw_link + - right_wrist_yaw_link + policy_observation_group: actor + critic_observation_group: critic + reward: - scales: - motion_global_root_pos: 1.0 - motion_global_root_ori: 0.5 - motion_body_pos: 2.0 - motion_body_ori: 1.0 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_joint_pos: 0.0 - motion_joint_vel: 0.0 - action_rate_l2: -0.1 - joint_limit: -2.0 - undesired_contacts: -0.1 + motion_global_root_pos: + func: unilab.tasks.motion_tracking.common.manager_terms.motion_global_anchor_position_error_exp + weight: 1.0 + params: {command_name: motion, std: 0.3} + motion_global_root_ori: + func: unilab.tasks.motion_tracking.common.manager_terms.motion_global_anchor_orientation_error_exp + weight: 0.5 + params: {command_name: motion, std: 0.4} + motion_body_pos: + func: unilab.tasks.motion_tracking.common.manager_terms.motion_relative_body_position_error_exp + weight: 2.0 + params: {command_name: motion, std: 0.3} + motion_body_ori: + func: unilab.tasks.motion_tracking.common.manager_terms.motion_relative_body_orientation_error_exp + weight: 1.0 + params: {command_name: motion, std: 0.4} + motion_body_lin_vel: + func: unilab.tasks.motion_tracking.common.manager_terms.motion_global_body_linear_velocity_error_exp + weight: 1.0 + params: {command_name: motion, std: 1.0} + motion_body_ang_vel: + func: unilab.tasks.motion_tracking.common.manager_terms.motion_global_body_angular_velocity_error_exp + weight: 1.0 + params: {command_name: motion, std: 3.14} + motion_joint_pos: + func: unilab.tasks.motion_tracking.common.manager_terms.motion_joint_position_error_exp + weight: 0.0 + params: {command_name: motion, std: 0.2} + motion_joint_vel: + func: unilab.tasks.motion_tracking.common.manager_terms.motion_joint_velocity_error_exp + weight: 0.0 + params: {command_name: motion, std: 1.0} + action_rate_l2: + func: unilab.envs.mdp.action_rate_l2 + weight: -0.1 + joint_limit: + func: unilab.tasks.motion_tracking.common.manager_terms.joint_pos_limits + weight: -2.0 + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + joint_names: ".*" + undesired_contacts: + func: unilab.tasks.motion_tracking.common.manager_terms.undesired_body_contacts + weight: -0.1 + params: + command_name: motion + threshold: 0.05 + body_names: + - pelvis + - left_hip_roll_link + - left_knee_link + - right_hip_roll_link + - right_knee_link + - torso_link + - left_shoulder_roll_link + - left_elbow_link + - right_shoulder_roll_link + - right_elbow_link diff --git a/conf/ppo/task/g1_23dof_motion_tracking/motrix.yaml b/conf/ppo/task/g1_23dof_motion_tracking/motrix.yaml index 6bcc0a74f..2141c58cf 100644 --- a/conf/ppo/task/g1_23dof_motion_tracking/motrix.yaml +++ b/conf/ppo/task/g1_23dof_motion_tracking/motrix.yaml @@ -1,52 +1,37 @@ # @package _global_ +defaults: + - /task/g1_23dof_motion_tracking/mujoco + - _self_ + training: task_name: G1MotionTracking23Dof sim_backend: motrix play_env_num: 16 - play_steps: 1000 -algo: - num_envs: 1024 - max_iterations: 15000 - save_interval: 500 - obs_groups: - actor: - - actor - algorithm: - entropy_coef: 0.005 - noise_config: - level: 1.0 - scale_joint_angle: 0.01 - scale_joint_vel: 1.5 - scale_gyro: 0.2 + +reward: + motion_global_root_pos: + weight: 1.0 + action_rate_l2: + weight: -0.05 + undesired_contacts: + func: unilab.tasks.motion_tracking.common.manager_terms.undesired_body_contacts + weight: -0.1 + params: + command_name: motion + threshold: 0.05 + body_names: + - pelvis + - left_hip_roll_link + - left_knee_link + - right_hip_roll_link + - right_knee_link + - torso_link + - left_shoulder_roll_link + - left_elbow_link + - right_shoulder_roll_link + - right_elbow_link + play_profile: enabled: true env: render_spacing: 2.5 - scene: - enabled: true - source_model_file: src/unilab/assets/robots/g1/scene_flat_23dof.xml - ground_texture_file: src/unilab/assets/robots/g1/textures/floor.png - skybox_rgb1: [0.90, 0.90, 0.91] - skybox_rgb2: [0.68, 0.68, 0.70] - ground_texrepeat: [0.25, 0.25] -reward: - scales: - motion_global_root_pos: 1.0 - motion_global_root_ori: 0.5 - motion_body_pos: 1.0 - motion_body_ori: 1.0 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_joint_pos: 0.0 - motion_joint_vel: 0.0 - action_rate_l2: -0.05 - joint_limit: -10.0 - undesired_contacts: -0.1 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 \ No newline at end of file diff --git a/conf/ppo/task/g1_23dof_motion_tracking/mujoco.yaml b/conf/ppo/task/g1_23dof_motion_tracking/mujoco.yaml index 3a411cd18..8605feb92 100644 --- a/conf/ppo/task/g1_23dof_motion_tracking/mujoco.yaml +++ b/conf/ppo/task/g1_23dof_motion_tracking/mujoco.yaml @@ -1,34 +1,67 @@ # @package _global_ +defaults: + - /task/g1_motion_tracking/mujoco + - _self_ + training: task_name: G1MotionTracking23Dof sim_backend: mujoco - play_steps: 1000 -algo: - num_envs: 1024 - max_iterations: 15000 - save_interval: 500 - obs_groups: - actor: - - actor - algorithm: - entropy_coef: 0.005 -reward: - scales: - motion_global_root_pos: 0.5 - motion_global_root_ori: 0.5 - motion_body_pos: 1.0 - motion_body_ori: 1.0 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_joint_pos: 0.0 - motion_joint_vel: 0.0 - action_rate_l2: -0.1 - joint_limit: -10.0 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 \ No newline at end of file + +env: + scene: + model_file: src/unilab/assets/robots/g1/scene_flat_23dof.xml + entities: + robot: + joint_names: &g1_23dof_joints + - left_hip_pitch_joint + - left_hip_roll_joint + - left_hip_yaw_joint + - left_knee_joint + - left_ankle_pitch_joint + - left_ankle_roll_joint + - right_hip_pitch_joint + - right_hip_roll_joint + - right_hip_yaw_joint + - right_knee_joint + - right_ankle_pitch_joint + - right_ankle_roll_joint + - waist_yaw_joint + - left_shoulder_pitch_joint + - left_shoulder_roll_joint + - left_shoulder_yaw_joint + - left_elbow_joint + - left_wrist_roll_joint + - right_shoulder_pitch_joint + - right_shoulder_roll_joint + - right_shoulder_yaw_joint + - right_elbow_joint + - right_wrist_roll_joint + actuator_names: *g1_23dof_joints + body_names: &tracked_bodies_23dof + - pelvis + - left_hip_roll_link + - left_knee_link + - left_ankle_roll_link + - right_hip_roll_link + - right_knee_link + - right_ankle_roll_link + - torso_link + - left_shoulder_roll_link + - left_elbow_link + - left_wrist_roll_rubber_hand + - right_shoulder_roll_link + - right_elbow_link + - right_wrist_roll_rubber_hand + commands: + motion: + params: + motion_file: motions/g1/dance1_subject2_part_23dof.npz + body_names: *tracked_bodies_23dof + terminations: + ee_body_pos: + params: + body_names: + - left_ankle_roll_link + - right_ankle_roll_link + - left_wrist_roll_rubber_hand + - right_wrist_roll_rubber_hand diff --git a/conf/ppo/task/g1_23dof_motion_tracking_deploy/motrix.yaml b/conf/ppo/task/g1_23dof_motion_tracking_deploy/motrix.yaml index a0359647e..46d82769f 100644 --- a/conf/ppo/task/g1_23dof_motion_tracking_deploy/motrix.yaml +++ b/conf/ppo/task/g1_23dof_motion_tracking_deploy/motrix.yaml @@ -1,101 +1,42 @@ # @package _global_ +defaults: + - /task/g1_23dof_motion_tracking_deploy/mujoco + - _self_ + training: task_name: G1MotionTracking23DofDeploy sim_backend: motrix play_env_num: 16 - play_steps: 1000 -algo: - num_envs: 1024 - max_iterations: 15000 - save_interval: 500 - obs_groups: - actor: - - actor - algorithm: - entropy_coef: 0.005 - noise_config: - level: 1.0 - scale_joint_angle: 0.01 - scale_joint_vel: 1.5 - scale_gyro: 0.2 + env: - sim_dt: 0.005 - sensor: - local_linvel: pelvis_local_linvel - gyro: pelvis_gyro - upvector: pelvis_upvector - domain_rand: - random_com: true - com_offset_x: [-0.025, 0.025] - com_offset_y: [-0.05, 0.05] - com_offset_z: [-0.05, 0.05] - randomize_base_mass: true - added_mass_range: [-1.5, 1.5] - push_robots: true - push_interval: 750 - max_force: [1.0, 1.0, 0.5] - randomize_joint_default_pos: true - joint_default_pos_range: [-0.01, 0.01] - noise_config: - scale_joint_angle: 0.01 - scale_joint_vel: 0.5 - scale_gyro: 0.2 - scale_linvel: 0.5 - scale_gravity: 0.05 - control_config: - action_scale: - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 + events: + foot_friction: null + push_robot: null + +reward: + motion_global_root_pos: + weight: 1.0 + action_rate_l2: + weight: -0.05 + undesired_contacts: + func: unilab.tasks.motion_tracking.common.manager_terms.undesired_body_contacts + weight: -0.1 + params: + command_name: motion + threshold: 0.05 + body_names: + - pelvis + - left_hip_roll_link + - left_knee_link + - right_hip_roll_link + - right_knee_link + - torso_link + - left_shoulder_roll_link + - left_elbow_link + - right_shoulder_roll_link + - right_elbow_link + play_profile: enabled: true env: render_spacing: 2.5 - scene: - enabled: true - source_model_file: src/unilab/assets/robots/g1/scene_flat_23dof.xml - ground_texture_file: src/unilab/assets/robots/g1/textures/floor.png - skybox_rgb1: [0.90, 0.90, 0.91] - skybox_rgb2: [0.68, 0.68, 0.70] - ground_texrepeat: [0.25, 0.25] -reward: - scales: - motion_global_root_pos: 1.0 - motion_global_root_ori: 0.5 - motion_body_pos: 1.0 - motion_body_ori: 1.0 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_joint_pos: 0.0 - motion_joint_vel: 0.0 - action_rate_l2: -0.05 - joint_limit: -10.0 - undesired_contacts: -0.1 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 diff --git a/conf/ppo/task/g1_23dof_motion_tracking_deploy/mujoco.yaml b/conf/ppo/task/g1_23dof_motion_tracking_deploy/mujoco.yaml index 4c725f29d..0e4667f90 100644 --- a/conf/ppo/task/g1_23dof_motion_tracking_deploy/mujoco.yaml +++ b/conf/ppo/task/g1_23dof_motion_tracking_deploy/mujoco.yaml @@ -1,85 +1,77 @@ # @package _global_ +defaults: + - /task/g1_motion_tracking_deploy/mujoco + - _self_ + training: task_name: G1MotionTracking23DofDeploy sim_backend: mujoco - play_steps: 1000 -algo: - num_envs: 1024 - max_iterations: 15000 - save_interval: 500 - obs_groups: - actor: - - actor - algorithm: - entropy_coef: 0.005 + +_g1_23dof_deploy_action_scale: + ".*_(hip_pitch|hip_yaw)_joint": 0.5475464629911068 + ".*_(hip_roll|knee)_joint": 0.35066146637882434 + ".*_ankle_(pitch|roll)_joint": 0.43857731392336724 + "waist_yaw_joint": 0.5475464629911068 + ".*_(shoulder_(pitch|roll|yaw)|elbow|wrist_roll)_joint": 0.43857731392336724 + env: - sim_dt: 0.005 - sensor: - local_linvel: pelvis_local_linvel - gyro: pelvis_gyro - upvector: pelvis_upvector - domain_rand: - random_com: true - com_offset_x: [-0.025, 0.025] - com_offset_y: [-0.05, 0.05] - com_offset_z: [-0.05, 0.05] - randomize_base_mass: true - added_mass_range: [-1.5, 1.5] - push_robots: true - push_interval: 750 - max_force: [1.0, 1.0, 0.5] - randomize_geom_friction: true - friction_range: [0.3, 1.2] - randomize_joint_default_pos: true - joint_default_pos_range: [-0.01, 0.01] - noise_config: - scale_joint_angle: 0.01 - scale_joint_vel: 0.5 - scale_gyro: 0.2 - scale_linvel: 0.5 - scale_gravity: 0.05 - control_config: - action_scale: - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 -reward: - scales: - motion_global_root_pos: 0.5 - motion_global_root_ori: 0.5 - motion_body_pos: 1.0 - motion_body_ori: 1.0 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_joint_pos: 0.0 - motion_joint_vel: 0.0 - action_rate_l2: -0.1 - joint_limit: -10.0 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 + scene: + model_file: src/unilab/assets/robots/g1/scene_flat_23dof.xml + entities: + robot: + joint_names: &g1_23dof_joints + - left_hip_pitch_joint + - left_hip_roll_joint + - left_hip_yaw_joint + - left_knee_joint + - left_ankle_pitch_joint + - left_ankle_roll_joint + - right_hip_pitch_joint + - right_hip_roll_joint + - right_hip_yaw_joint + - right_knee_joint + - right_ankle_pitch_joint + - right_ankle_roll_joint + - waist_yaw_joint + - left_shoulder_pitch_joint + - left_shoulder_roll_joint + - left_shoulder_yaw_joint + - left_elbow_joint + - left_wrist_roll_joint + - right_shoulder_pitch_joint + - right_shoulder_roll_joint + - right_shoulder_yaw_joint + - right_elbow_joint + - right_wrist_roll_joint + actuator_names: *g1_23dof_joints + body_names: &tracked_bodies_23dof + - pelvis + - left_hip_roll_link + - left_knee_link + - left_ankle_roll_link + - right_hip_roll_link + - right_knee_link + - right_ankle_roll_link + - torso_link + - left_shoulder_roll_link + - left_elbow_link + - left_wrist_roll_rubber_hand + - right_shoulder_roll_link + - right_elbow_link + - right_wrist_roll_rubber_hand + actions: + joint_pos: + scale: ${_g1_23dof_deploy_action_scale} + commands: + motion: + params: + motion_file: motions/g1/dance1_subject2_part_23dof.npz + body_names: *tracked_bodies_23dof + terminations: + ee_body_pos: + params: + body_names: + - left_ankle_roll_link + - right_ankle_roll_link + - left_wrist_roll_rubber_hand + - right_wrist_roll_rubber_hand diff --git a/conf/ppo/task/g1_motion_tracking/motrix.yaml b/conf/ppo/task/g1_motion_tracking/motrix.yaml index b4e531856..57f4e35be 100644 --- a/conf/ppo/task/g1_motion_tracking/motrix.yaml +++ b/conf/ppo/task/g1_motion_tracking/motrix.yaml @@ -1,53 +1,37 @@ # @package _global_ +defaults: + - /task/g1_motion_tracking/mujoco + - _self_ + training: task_name: G1MotionTracking sim_backend: motrix play_env_num: 16 - play_steps: 1000 -algo: - num_envs: 1024 - max_iterations: 15000 - save_interval: 500 - obs_groups: - actor: - - actor - algorithm: - entropy_coef: 0.005 - noise_config: - level: 1.0 - scale_joint_angle: 0.01 - scale_joint_vel: 1.5 - scale_gyro: 0.2 -env: + +reward: + motion_global_root_pos: + weight: 1.0 + action_rate_l2: + weight: -0.05 + undesired_contacts: + func: unilab.tasks.motion_tracking.common.manager_terms.undesired_body_contacts + weight: -0.1 + params: + command_name: motion + threshold: 0.05 + body_names: + - pelvis + - left_hip_roll_link + - left_knee_link + - right_hip_roll_link + - right_knee_link + - torso_link + - left_shoulder_roll_link + - left_elbow_link + - right_shoulder_roll_link + - right_elbow_link + play_profile: enabled: true env: render_spacing: 2.5 - scene: - enabled: true - source_model_file: src/unilab/assets/robots/g1/scene_flat.xml - ground_texture_file: src/unilab/assets/robots/g1/textures/floor.png - skybox_rgb1: [0.90, 0.90, 0.91] - skybox_rgb2: [0.68, 0.68, 0.70] - ground_texrepeat: [0.25, 0.25] -reward: - scales: - motion_global_root_pos: 1.0 - motion_global_root_ori: 0.5 - motion_body_pos: 1.0 - motion_body_ori: 1.0 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_joint_pos: 0.0 - motion_joint_vel: 0.0 - action_rate_l2: -0.05 - joint_limit: -10.0 - undesired_contacts: -0.1 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 diff --git a/conf/ppo/task/g1_motion_tracking/mujoco.yaml b/conf/ppo/task/g1_motion_tracking/mujoco.yaml index 54ffea9fd..343c454e0 100644 --- a/conf/ppo/task/g1_motion_tracking/mujoco.yaml +++ b/conf/ppo/task/g1_motion_tracking/mujoco.yaml @@ -3,36 +3,237 @@ training: task_name: G1MotionTracking sim_backend: mujoco play_steps: 1000 + algo: num_envs: 1024 max_iterations: 15000 save_interval: 500 obs_groups: - actor: - - actor + actor: [actor] algorithm: entropy_coef: 0.005 + env: + scene: + model_file: src/unilab/assets/robots/g1/scene_flat.xml + default_keyframe_name: stand + entities: + robot: + root_body_name: pelvis + joint_names: &g1_joints + - left_hip_pitch_joint + - left_hip_roll_joint + - left_hip_yaw_joint + - left_knee_joint + - left_ankle_pitch_joint + - left_ankle_roll_joint + - right_hip_pitch_joint + - right_hip_roll_joint + - right_hip_yaw_joint + - right_knee_joint + - right_ankle_pitch_joint + - right_ankle_roll_joint + - waist_yaw_joint + - waist_roll_joint + - waist_pitch_joint + - left_shoulder_pitch_joint + - left_shoulder_roll_joint + - left_shoulder_yaw_joint + - left_elbow_joint + - left_wrist_roll_joint + - left_wrist_pitch_joint + - left_wrist_yaw_joint + - right_shoulder_pitch_joint + - right_shoulder_roll_joint + - right_shoulder_yaw_joint + - right_elbow_joint + - right_wrist_roll_joint + - right_wrist_pitch_joint + - right_wrist_yaw_joint + actuator_names: *g1_joints + body_names: &tracked_bodies + - pelvis + - left_hip_roll_link + - left_knee_link + - left_ankle_roll_link + - right_hip_roll_link + - right_knee_link + - right_ankle_roll_link + - torso_link + - left_shoulder_roll_link + - left_elbow_link + - left_wrist_yaw_link + - right_shoulder_roll_link + - right_elbow_link + - right_wrist_yaw_link + geom_names: + - left_foot1_collision + - left_foot2_collision + - left_foot3_collision + - left_foot4_collision + - left_foot5_collision + - left_foot6_collision + - left_foot7_collision + - right_foot1_collision + - right_foot2_collision + - right_foot3_collision + - right_foot4_collision + - right_foot5_collision + - right_foot6_collision + - right_foot7_collision + sim_dt: 0.006666666666666667 + ctrl_dt: 0.02 + max_episode_seconds: 10.0 + observations: + actor: + terms: + command: &command_obs + func: unilab.envs.mdp.generated_commands + params: {command_name: motion} + motion_anchor_pos_b: &anchor_pos_obs + func: unilab.tasks.motion_tracking.common.manager_terms.motion_anchor_pos_b + params: {command_name: motion} + motion_anchor_ori_b: &anchor_ori_obs + func: unilab.tasks.motion_tracking.common.manager_terms.motion_anchor_ori_b + params: {command_name: motion} + base_lin_vel: &base_lin_vel_obs + func: unilab.envs.mdp.builtin_sensor + params: {sensor_name: pelvis_local_linvel} + base_ang_vel: &base_ang_vel_obs + func: unilab.envs.mdp.builtin_sensor + params: {sensor_name: torso_gyro} + joint_pos: &joint_pos_obs + func: unilab.tasks.motion_tracking.common.manager_terms.motion_joint_pos_rel + params: {command_name: motion} + joint_vel: &joint_vel_obs + func: unilab.envs.mdp.joint_vel_rel + actions: &actions_obs + func: unilab.envs.mdp.last_action + critic: + terms: + command: *command_obs + motion_anchor_pos_b: *anchor_pos_obs + motion_anchor_ori_b: *anchor_ori_obs + base_lin_vel: *base_lin_vel_obs + base_ang_vel: *base_ang_vel_obs + joint_pos: *joint_pos_obs + joint_vel: *joint_vel_obs + actions: *actions_obs + body_pos: + func: unilab.tasks.motion_tracking.common.manager_terms.robot_body_pos_b + params: {command_name: motion} + body_ori: + func: unilab.tasks.motion_tracking.common.manager_terms.robot_body_ori_b + params: {command_name: motion} + actions: + joint_pos: + _target_: unilab.tasks.motion_tracking.common.manager_terms.MotionJointPositionActionCfg + entity_name: robot + actuator_names: [".*"] + scale: 0.25 + use_default_offset: true + command_name: motion + commands: + motion: + _target_: unilab.tasks.motion_tracking.common.manager_terms.MotionCommandCfg + entity_name: robot + resampling_time_range: [1.0e9, 1.0e9] + params: + motion_file: motions/g1/dance1_subject2_part.npz + anchor_body_name: torso_link + body_names: *tracked_bodies + sampling_mode: adaptive + sampling_start_ratio: 0.0 + truncate_on_clip_end: false + pose_range: + x: [-0.05, 0.05] + y: [-0.05, 0.05] + z: [-0.01, 0.01] + roll: [-0.1, 0.1] + pitch: [-0.1, 0.1] + yaw: [-0.2, 0.2] + velocity_range: + x: [-0.5, 0.5] + y: [-0.5, 0.5] + z: [-0.2, 0.2] + roll: [-0.52, 0.52] + pitch: [-0.52, 0.52] + yaw: [-0.78, 0.78] + joint_position_range: [-0.1, 0.1] + joint_default_position_range: [0.0, 0.0] + terminations: + time_out: + func: unilab.envs.mdp.time_out + time_out: true + anchor_pos: + func: unilab.tasks.motion_tracking.common.manager_terms.bad_anchor_pos_z_only + params: {command_name: motion, threshold: 0.25} + anchor_ori: + func: unilab.tasks.motion_tracking.common.manager_terms.bad_anchor_ori + params: + command_name: motion + threshold: 0.8 + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + ee_body_pos: + func: unilab.tasks.motion_tracking.common.manager_terms.bad_motion_body_pos_z_only + params: + command_name: motion + threshold: 0.25 + body_names: + - left_ankle_roll_link + - right_ankle_roll_link + - left_wrist_yaw_link + - right_wrist_yaw_link + policy_observation_group: actor + critic_observation_group: critic + reward: - scales: - motion_global_root_pos: 0.5 - motion_global_root_ori: 0.5 - motion_body_pos: 1.0 - motion_body_ori: 1.0 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_joint_pos: 0.0 - motion_joint_vel: 0.0 - action_rate_l2: -0.1 - joint_limit: -10.0 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 + motion_global_root_pos: + func: unilab.tasks.motion_tracking.common.manager_terms.motion_global_anchor_position_error_exp + weight: 0.5 + params: {command_name: motion, std: 0.3} + motion_global_root_ori: + func: unilab.tasks.motion_tracking.common.manager_terms.motion_global_anchor_orientation_error_exp + weight: 0.5 + params: {command_name: motion, std: 0.4} + motion_body_pos: + func: unilab.tasks.motion_tracking.common.manager_terms.motion_relative_body_position_error_exp + weight: 1.0 + params: {command_name: motion, std: 0.3} + motion_body_ori: + func: unilab.tasks.motion_tracking.common.manager_terms.motion_relative_body_orientation_error_exp + weight: 1.0 + params: {command_name: motion, std: 0.4} + motion_body_lin_vel: + func: unilab.tasks.motion_tracking.common.manager_terms.motion_global_body_linear_velocity_error_exp + weight: 1.0 + params: {command_name: motion, std: 1.0} + motion_body_ang_vel: + func: unilab.tasks.motion_tracking.common.manager_terms.motion_global_body_angular_velocity_error_exp + weight: 1.0 + params: {command_name: motion, std: 3.14} + motion_joint_pos: + func: unilab.tasks.motion_tracking.common.manager_terms.motion_joint_position_error_exp + weight: 0.0 + params: {command_name: motion, std: 0.2} + motion_joint_vel: + func: unilab.tasks.motion_tracking.common.manager_terms.motion_joint_velocity_error_exp + weight: 0.0 + params: {command_name: motion, std: 1.0} + action_rate_l2: + func: unilab.envs.mdp.action_rate_l2 + weight: -0.1 + joint_limit: + func: unilab.tasks.motion_tracking.common.manager_terms.joint_pos_limits + weight: -10.0 + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + joint_names: ".*" + play_profile: enabled: true env: diff --git a/conf/ppo/task/g1_motion_tracking_deploy/motrix.yaml b/conf/ppo/task/g1_motion_tracking_deploy/motrix.yaml index 6ebc34b25..9d2fe0e6e 100644 --- a/conf/ppo/task/g1_motion_tracking_deploy/motrix.yaml +++ b/conf/ppo/task/g1_motion_tracking_deploy/motrix.yaml @@ -1,107 +1,42 @@ # @package _global_ +defaults: + - /task/g1_motion_tracking_deploy/mujoco + - _self_ + training: task_name: G1MotionTrackingDeploy sim_backend: motrix play_env_num: 16 - play_steps: 1000 -algo: - num_envs: 1024 - max_iterations: 15000 - save_interval: 500 - obs_groups: - actor: - - actor - algorithm: - entropy_coef: 0.005 - noise_config: - level: 1.0 - scale_joint_angle: 0.01 - scale_joint_vel: 1.5 - scale_gyro: 0.2 + +env: + events: + foot_friction: null + push_robot: null + +reward: + motion_global_root_pos: + weight: 1.0 + action_rate_l2: + weight: -0.05 + undesired_contacts: + func: unilab.tasks.motion_tracking.common.manager_terms.undesired_body_contacts + weight: -0.1 + params: + command_name: motion + threshold: 0.05 + body_names: + - pelvis + - left_hip_roll_link + - left_knee_link + - right_hip_roll_link + - right_knee_link + - torso_link + - left_shoulder_roll_link + - left_elbow_link + - right_shoulder_roll_link + - right_elbow_link + play_profile: enabled: true env: render_spacing: 2.5 - scene: - enabled: true - source_model_file: src/unilab/assets/robots/g1/scene_flat.xml - ground_texture_file: src/unilab/assets/robots/g1/textures/floor.png - skybox_rgb1: [0.90, 0.90, 0.91] - skybox_rgb2: [0.68, 0.68, 0.70] - ground_texrepeat: [0.25, 0.25] -env: - sim_dt: 0.005 - sensor: - local_linvel: pelvis_local_linvel - gyro: pelvis_gyro - upvector: pelvis_upvector - domain_rand: - random_com: true - com_offset_x: [-0.025, 0.025] - com_offset_y: [-0.05, 0.05] - com_offset_z: [-0.05, 0.05] - randomize_base_mass: true - added_mass_range: [-1.5, 1.5] - push_robots: true - push_interval: 750 - max_force: [1.0, 1.0, 0.5] - randomize_joint_default_pos: true - joint_default_pos_range: [-0.01, 0.01] - # randomize_geom_friction omitted: Motrix does not reliably support it - noise_config: - scale_joint_angle: 0.01 - scale_joint_vel: 0.5 - scale_gyro: 0.2 - scale_linvel: 0.5 - scale_gravity: 0.05 - control_config: - action_scale: - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.07450087032950714 - - 0.07450087032950714 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.07450087032950714 - - 0.07450087032950714 -reward: - scales: - motion_global_root_pos: 0.5 - motion_global_root_ori: 0.5 - motion_body_pos: 1.0 - motion_body_ori: 1.0 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_joint_pos: 0.0 - motion_joint_vel: 0.0 - action_rate_l2: -0.1 - joint_limit: -10.0 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 diff --git a/conf/ppo/task/g1_motion_tracking_deploy/mujoco.yaml b/conf/ppo/task/g1_motion_tracking_deploy/mujoco.yaml index 39a10a9ca..6b8d5ef89 100644 --- a/conf/ppo/task/g1_motion_tracking_deploy/mujoco.yaml +++ b/conf/ppo/task/g1_motion_tracking_deploy/mujoco.yaml @@ -1,91 +1,84 @@ # @package _global_ +defaults: + - /task/g1_motion_tracking/mujoco + - _self_ + training: task_name: G1MotionTrackingDeploy sim_backend: mujoco - play_steps: 1000 -algo: - num_envs: 1024 - max_iterations: 15000 - save_interval: 500 - obs_groups: - actor: - - actor - algorithm: - entropy_coef: 0.005 + env: sim_dt: 0.005 - sensor: - local_linvel: pelvis_local_linvel - gyro: pelvis_gyro - upvector: pelvis_upvector - domain_rand: - random_com: true - com_offset_x: [-0.025, 0.025] - com_offset_y: [-0.05, 0.05] - com_offset_z: [-0.05, 0.05] - randomize_base_mass: true - added_mass_range: [-1.5, 1.5] - push_robots: true - push_interval: 750 - max_force: [1.0, 1.0, 0.5] - randomize_geom_friction: true - friction_range: [0.3, 1.2] - randomize_joint_default_pos: true - joint_default_pos_range: [-0.01, 0.01] - noise_config: - scale_joint_angle: 0.01 - scale_joint_vel: 0.5 - scale_gyro: 0.2 - scale_linvel: 0.5 - scale_gravity: 0.05 - control_config: - action_scale: - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.07450087032950714 - - 0.07450087032950714 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.07450087032950714 - - 0.07450087032950714 -reward: - scales: - motion_global_root_pos: 0.5 - motion_global_root_ori: 0.5 - motion_body_pos: 1.0 - motion_body_ori: 1.0 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_joint_pos: 0.0 - motion_joint_vel: 0.0 - action_rate_l2: -0.1 - joint_limit: -10.0 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 + observations: + actor: + terms: + motion_anchor_pos_b: null + base_lin_vel: null + base_ang_vel: + params: {sensor_name: pelvis_gyro} + critic: + terms: + base_ang_vel: + params: {sensor_name: pelvis_gyro} + actions: + joint_pos: + scale: + ".*_(hip_pitch|hip_yaw)_joint": 0.5475464629911068 + ".*_(hip_roll|knee)_joint": 0.35066146637882434 + ".*_ankle_(pitch|roll)_joint": 0.43857731392336724 + "waist_yaw_joint": 0.5475464629911068 + "waist_(roll|pitch)_joint": 0.43857731392336724 + ".*_(shoulder_(pitch|roll|yaw)|elbow|wrist_roll)_joint": 0.43857731392336724 + ".*_wrist_(pitch|yaw)_joint": 0.07450087032950714 + commands: + motion: + params: + joint_default_position_range: [-0.01, 0.01] + events: + base_mass: + func: unilab.envs.mdp.randomize_rigid_body_mass + mode: reset + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + body_names: pelvis + mass_distribution_params: [-1.5, 1.5] + operation: add + recompute_inertia: false + base_com: + func: unilab.envs.mdp.randomize_rigid_body_com + mode: reset + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + body_names: pelvis + com_range: + x: [-0.025, 0.025] + y: [-0.05, 0.05] + z: [-0.05, 0.05] + foot_friction: + func: unilab.envs.mdp.geom_friction + mode: reset + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + geom_names: ".*" + ranges: [0.3, 1.2] + operation: abs + shared_random: true + push_robot: + func: unilab.envs.mdp.push_by_setting_velocity + mode: interval + interval_range_s: [15.0, 15.0] + is_global_time: true + params: + velocity_range: + x: [-1.0, 1.0] + y: [-1.0, 1.0] + z: [-0.5, 0.5] + roll: [0.0, 0.0] + pitch: [0.0, 0.0] + yaw: [0.0, 0.0] diff --git a/scripts/benchmark/benchmark_drake_performance.py b/scripts/benchmark/benchmark_drake_performance.py index e4a9c10be..01429e921 100644 --- a/scripts/benchmark/benchmark_drake_performance.py +++ b/scripts/benchmark/benchmark_drake_performance.py @@ -1,10 +1,5 @@ #!/usr/bin/env python3 -"""Profile Drake vs MuJoCo env-step performance on selected UniLab tasks. - -This benchmark is intentionally task-level rather than raw-simulator-level. It -keeps UniLab's reset, observation, sensor-view, and body-query paths in the -loop so G1 motion tracking can expose the expensive integration points. -""" +"""Profile Drake vs MuJoCo env-step performance on selected UniLab tasks.""" from __future__ import annotations @@ -115,18 +110,7 @@ def go2_cfg() -> Any: return ManagerBasedRlEnvCfg() - def g1_tracking_cfg() -> Any: - from unilab.tasks.motion_tracking.g1.tracking import G1MotionTrackingEnvCfg - - return G1MotionTrackingEnvCfg() - - def g1_tracking_env() -> type: - from unilab.tasks.motion_tracking.g1.tracking import G1MotionTrackingEnv - - return G1MotionTrackingEnv - return { - "g1_motion_tracking": TaskSpec(g1_tracking_cfg, g1_tracking_env), "go1_joystick_flat": TaskSpec(go1_cfg, manager_env), "go2_joystick_flat": TaskSpec(go2_cfg, manager_env), } @@ -336,10 +320,7 @@ def main() -> None: parser.add_argument( "--tasks", default="go1_joystick_flat,go2_joystick_flat", - help=( - "Comma-separated task ids. Defaults stay within committed Drake task configs; " - "pass g1_motion_tracking explicitly when its Drake config is available." - ), + help="Comma-separated task ids with committed Drake YAML owners.", ) parser.add_argument("--backends", default="drake,mujoco", help="Comma-separated backends.") parser.add_argument("--num-envs", default="64,256,1024", help="Comma-separated env counts.") diff --git a/scripts/benchmark/env/benchmark_env_step.py b/scripts/benchmark/env/benchmark_env_step.py index 1445b83e4..940d0f8ef 100644 --- a/scripts/benchmark/env/benchmark_env_step.py +++ b/scripts/benchmark/env/benchmark_env_step.py @@ -351,12 +351,12 @@ def _g1_rough_cfg(backend: str, config_overrides: list[str]) -> Any: def _g1_motion_tracking_cfg(backend: str, config_overrides: list[str]) -> Any: - from unilab.tasks.motion_tracking.g1.tracking import G1MotionTrackingEnvCfg + from unilab.envs import ManagerBasedRlEnvCfg return _ppo_owner_yaml_cfg( "g1_motion_tracking", backend, - G1MotionTrackingEnvCfg, + ManagerBasedRlEnvCfg, config_overrides, ) @@ -422,12 +422,6 @@ def _g1_walk_env_cls() -> type: return make_g1_walk_env -def _g1_motion_tracking_env_cls() -> type: - from unilab.tasks.motion_tracking.g1.tracking import G1MotionTrackingEnv - - return G1MotionTrackingEnv - - def _sharpa_inhand_env_cls() -> type: from unilab.tasks.manipulation.sharpa_inhand.rotation import SharpaInhandRotationEnv @@ -489,8 +483,8 @@ def _sharpa_inhand_env_cls() -> type: task_id="g1_motion_tracking", env_name="G1MotionTracking", cfg_factory=_g1_motion_tracking_cfg, - env_cls_factory=_g1_motion_tracking_env_cls, - backends=("mujoco", "motrix", "mjwarp"), + env_cls_factory=_manager_env_cls, + backends=("mujoco", "motrix"), ), "sharpa_inhand": TaskConfig( task_id="sharpa_inhand", diff --git a/src/unilab/base/backend/motrix/backend.py b/src/unilab/base/backend/motrix/backend.py index 16672e3c9..53c02ec97 100644 --- a/src/unilab/base/backend/motrix/backend.py +++ b/src/unilab/base/backend/motrix/backend.py @@ -638,7 +638,20 @@ def get_gravity(self) -> np.ndarray: return np.asarray(self._model.options.gravity, dtype=np.float64).copy() def get_joint_range(self) -> np.ndarray | None: - return None + """Return single-DoF joint limits in backend DOF order. + + Motrix stores the model-wide limits as a ``(2, num_dof)`` table, + whereas the UniLab backend contract exposes the MuJoCo-shaped + ``(num_dof, 2)`` table. This is materialized once by ``Entity`` and + never queried from a task hot path. + """ + raw_limits = np.asarray(self._model.joint_limits, dtype=self._np_dtype) + if raw_limits.ndim != 2 or raw_limits.shape != (2, self.num_dof_vel): + raise ValueError( + "Motrix joint limits must have shape (2, num_dof); " + f"received {raw_limits.shape} for {self.num_dof_vel} DOFs" + ) + return np.array(raw_limits.T, copy=True) # ------------------------------------------------------------------ # # Simulation control # diff --git a/src/unilab/base/entity.py b/src/unilab/base/entity.py index 8ccbc5eea..d9593b0f0 100644 --- a/src/unilab/base/entity.py +++ b/src/unilab/base/entity.py @@ -146,6 +146,7 @@ def __init__( default_root_state_error: str | None, default_joint_pos: np.ndarray | None, default_joint_vel: np.ndarray | None, + soft_joint_pos_limits: np.ndarray | None, gravity_vec_w: np.ndarray | None, body_ids: np.ndarray | None, actuator_ids: np.ndarray | None, @@ -164,6 +165,7 @@ def __init__( self._default_root_state_error = default_root_state_error self._default_joint_pos = default_joint_pos self._default_joint_vel = default_joint_vel + self._soft_joint_pos_limits = soft_joint_pos_limits self._gravity_vec_w = gravity_vec_w self._encoder_bias = ( None @@ -224,6 +226,11 @@ def projected_gravity_b(self) -> np.ndarray: gravity = self._require(self._gravity_vec_w, "projected gravity") return np_quat_apply_inverse(self.root_link_quat_w, gravity) + @property + def gravity_vec_w(self) -> np.ndarray: + """Read-only world-frame unit gravity vector for every environment.""" + return self._require(self._gravity_vec_w, "world-frame gravity") + @property def root_link_pose_w(self) -> np.ndarray: return np.concatenate((self.root_link_pos_w, self.root_link_quat_w), axis=-1) @@ -268,6 +275,11 @@ def default_joint_vel(self) -> np.ndarray: """Read-only zero default velocities from the UniLab reset contract.""" return self._require(self._default_joint_vel, "default joint velocity") + @property + def soft_joint_pos_limits(self) -> np.ndarray: + """Read-only joint position limits in the declared entity joint order.""" + return self._require(self._soft_joint_pos_limits, "joint position limits") + @property def encoder_bias(self) -> np.ndarray: """Mutable per-environment joint encoder bias used by position actions.""" @@ -427,6 +439,7 @@ def __init__( self._reset_joint_qpos_ids: np.ndarray | None = None self._reset_joint_qvel_ids: np.ndarray | None = None self._joint_model_dof_ids: np.ndarray | None = None + self._motion_body_ids: np.ndarray | None = None self._joint_names = _normalize_names(name, "joint", cfg.joint_names) self._body_names = _normalize_names(name, "body", cfg.body_names) @@ -493,6 +506,7 @@ def __init__( default_qpos, ) default_joint_vel = self._materialize_default_joint_vel(backend, joint_vel_ids) + soft_joint_pos_limits = self._materialize_soft_joint_pos_limits(backend, joint_pos_ids) gravity_vec_w = self._materialize_gravity_vector(backend, root_body_ids) actuator_ctrl_range = self._materialize_actuator_ctrl_range(backend, actuator_ids) ( @@ -521,6 +535,7 @@ def __init__( default_root_state_error=self._reset_root_layout_error, default_joint_pos=default_joint_pos, default_joint_vel=default_joint_vel, + soft_joint_pos_limits=soft_joint_pos_limits, gravity_vec_w=gravity_vec_w, body_ids=body_ids, actuator_ids=actuator_ids, @@ -530,6 +545,22 @@ def __init__( backend_type=self._backend_type, ) + @property + def motion_body_ids(self) -> np.ndarray: + """Motion-dataset body columns for the declared entity body order.""" + if self._body_names is None: + raise self._capability_error( + "motion body IDs", + "body_names were not declared in EntityCfg", + ) + if self._motion_body_ids is None: + self._motion_body_ids = self._resolve_ids( + "motion body", + self._body_names, + self._backend.get_motion_body_ids, + ) + return self._motion_body_ids + def _capability_error(self, capability: str, detail: str) -> NotImplementedError: return NotImplementedError( f"Entity '{self.name}' capability '{capability}' is unavailable on " @@ -723,6 +754,35 @@ def _materialize_default_joint_pos( materialized.setflags(write=False) return materialized + def _materialize_soft_joint_pos_limits( + self, + backend: SimBackend, + joint_pos_ids: np.ndarray | None, + ) -> np.ndarray | None: + if joint_pos_ids is None: + return None + try: + raw_ranges = backend.get_joint_range() + except (AttributeError, NotImplementedError) as exc: + raise self._capability_error("joint position limits", str(exc)) from exc + if raw_ranges is None: + return None + ranges = np.asarray(raw_ranges) + if ranges.ndim != 2 or ranges.shape[1] != 2: + raise ValueError( + f"Entity '{self.name}' capability 'joint position limits' on backend " + f"'{self._backend_type}' returned shape {ranges.shape}; expected (num_dof, 2)" + ) + if joint_pos_ids.size and int(np.max(joint_pos_ids)) >= ranges.shape[0]: + raise ValueError( + f"Entity '{self.name}' capability 'joint position limits' resolved index " + f"{int(np.max(joint_pos_ids))}, but backend '{self._backend_type}' returned " + f"only {ranges.shape[0]} rows" + ) + selected = np.array(ranges[_as_column_index(joint_pos_ids)], copy=True) + selected.setflags(write=False) + return selected + def _materialize_root_state( self, backend: SimBackend, diff --git a/src/unilab/envs/manager_based_rl_env.py b/src/unilab/envs/manager_based_rl_env.py index 21d8edd0a..01c4133af 100644 --- a/src/unilab/envs/manager_based_rl_env.py +++ b/src/unilab/envs/manager_based_rl_env.py @@ -243,6 +243,8 @@ def __init__(self, cfg: ManagerBasedRlEnvCfg, backend: SimBackend, num_envs: int self._command_dt = np.zeros(num_envs, dtype=get_global_dtype()) self._no_truncation = np.zeros(num_envs, dtype=np.bool_) self._manual_reset_pending = np.zeros(num_envs, dtype=np.bool_) + self._all_env_ids = np.arange(num_envs, dtype=np.int32) + self._all_env_ids.setflags(write=False) self._has_transition = False self._uses_pre_step_control = False @@ -459,7 +461,9 @@ def update_state(self, state: NpEnvState) -> NpEnvState: self._command_dt.fill(self.step_dt) self._command_dt[self.reset_buf] = 0.0 - self.command_manager.compute(dt=self._command_dt) + with self._reset_state.scoped(self._all_env_ids): + self.command_manager.compute(dt=self._command_dt) + self.command_manager.post_compute() manager_obs = self.observation_manager.compute(update_history=True) self.obs_buf = self._map_observations(manager_obs) self._has_transition = True @@ -504,13 +508,14 @@ def reset( log: dict[str, Any] = {} self.curriculum_manager.compute(env_ids=ids) - if "reset" in self.event_manager.available_modes: - with self._reset_state.scoped(ids): + with self._reset_state.scoped(ids): + if "reset" in self.event_manager.available_modes: self.event_manager.apply( mode="reset", env_ids=ids, global_env_step_count=self.step_counter, ) + log.update(self.command_manager.reset(ids)) for manager in ( self.observation_manager, @@ -518,7 +523,6 @@ def reset( self.reward_manager, self.metrics_manager, self.curriculum_manager, - self.command_manager, self.event_manager, self.termination_manager, ): @@ -531,6 +535,7 @@ def reset( self._state.info["steps"][ids] = 0 self.command_manager.compute(dt=0.0, env_ids=ids) + self.command_manager.post_compute() manager_obs = self.observation_manager.compute(update_history=True, env_ids=ids) mapped_obs = self._map_observations(manager_obs) reset_obs = {name: values[ids].copy() for name, values in mapped_obs.items()} diff --git a/src/unilab/managers/_types.py b/src/unilab/managers/_types.py index 6971fcf6c..d59d21f5f 100644 --- a/src/unilab/managers/_types.py +++ b/src/unilab/managers/_types.py @@ -196,6 +196,8 @@ def get_term(self, name: str) -> ManagerActionTerm: ... class ManagerCommandManager(Protocol): def get_command(self, name: str) -> np.ndarray | None: ... + def get_term(self, name: str) -> Any: ... + class ManagerTerminationManager(Protocol): @property @@ -236,6 +238,12 @@ def termination_manager(self) -> ManagerTerminationManager: ... @property def episode_length_buf(self) -> np.ndarray: ... + @property + def reset_buf(self) -> np.ndarray: ... + + @property + def common_step_counter(self) -> int: ... + @property def max_episode_length(self) -> int: ... diff --git a/src/unilab/managers/command_manager.py b/src/unilab/managers/command_manager.py index 2b20323ee..81451bb2a 100644 --- a/src/unilab/managers/command_manager.py +++ b/src/unilab/managers/command_manager.py @@ -175,6 +175,9 @@ def _update_command(self, env_ids: np.ndarray | None) -> None: """ raise NotImplementedError + def post_compute(self) -> None: + """Refresh state that depends on committed command-side simulation writes.""" + class CommandManager(ManagerBase): """Manages command generation for the environment. @@ -235,6 +238,10 @@ def compute(self, dt: float | np.ndarray, env_ids: np.ndarray | None = None) -> term.compute(dt, env_ids) self._validate_command(name, term.command) + def post_compute(self) -> None: + for term in self._terms.values(): + term.post_compute() + def get_command(self, name: str) -> np.ndarray: return self._validate_command(name, self._terms[name].command) @@ -303,6 +310,9 @@ def reset(self, env_ids: np.ndarray | None = None) -> dict[str, np.ndarray]: def compute(self, dt: float | np.ndarray, env_ids: np.ndarray | None = None) -> None: pass + def post_compute(self) -> None: + pass + def get_command(self, name: str) -> None: return None diff --git a/src/unilab/tasks/manipulation/stewart/balance.py b/src/unilab/tasks/manipulation/stewart/balance.py index 03bea58ed..51faa06d9 100644 --- a/src/unilab/tasks/manipulation/stewart/balance.py +++ b/src/unilab/tasks/manipulation/stewart/balance.py @@ -33,7 +33,9 @@ from unilab.managers.termination_manager import TerminationManager class _StewartEnv(ManagerBasedRlEnv, Protocol): - common_step_counter: int + @property + def common_step_counter(self) -> int: ... + observation_manager: ObservationManager diff --git a/src/unilab/tasks/migration_matrix.py b/src/unilab/tasks/migration_matrix.py index ccde6acdb..8691ae85e 100644 --- a/src/unilab/tasks/migration_matrix.py +++ b/src/unilab/tasks/migration_matrix.py @@ -62,6 +62,17 @@ class TaskMigrationRecord: } ) +_MOTION_CORE_TASKS = frozenset( + { + "G1MotionTracking", + "G1MotionTracking23Dof", + "G1MotionTracking23DofDeploy", + "G1MotionTrackingDeploy", + "G1MotionTrackingSAC", + "G1MotionTrackingSAC23Dof", + } +) + _MOTION_TASKS = frozenset( { "G1BoxTracking", @@ -72,12 +83,6 @@ class TaskMigrationRecord: "G1FlipTracking23Dof", "G1FlipTrackingSAC", "G1FlipTrackingSAC23Dof", - "G1MotionTracking", - "G1MotionTracking23Dof", - "G1MotionTracking23DofDeploy", - "G1MotionTrackingDeploy", - "G1MotionTrackingSAC", - "G1MotionTrackingSAC23Dof", "G1WallFlipTracking", "G1WallFlipTracking23Dof", "G1WallFlipTrackingSAC", @@ -89,7 +94,12 @@ class TaskMigrationRecord: ) PRODUCTION_TASK_NAMES = frozenset( - _MBA_TASKS | _ROUGH_TASKS | _G1_LOCOMOTION_TASKS | _CUSTOM_COMPAT_TASKS | _MOTION_TASKS + _MBA_TASKS + | _ROUGH_TASKS + | _G1_LOCOMOTION_TASKS + | _CUSTOM_COMPAT_TASKS + | _MOTION_CORE_TASKS + | _MOTION_TASKS ) @@ -137,6 +147,15 @@ def migration_record(task_name: str) -> TaskMigrationRecord: "Custom IK/history or tactile/contact/cache behavior is retained behind one frozen adapter.", "Keep Hydra/Registry ownership single; migrate only when the formal capability exists.", ) + if task_name in _MOTION_CORE_TASKS: + return TaskMigrationRecord( + task_name, + "motion_tracking", + "Compatible", + "complete", + "Hydra owner YAML materializes task-owned NumPy motion manager terms on the canonical runtime.", + "Keep PPO, APPO, and SAC owners aligned with the shared motion manager contract.", + ) if task_name in _MOTION_TASKS: return TaskMigrationRecord( task_name, diff --git a/src/unilab/tasks/motion_tracking/common/manager_terms.py b/src/unilab/tasks/motion_tracking/common/manager_terms.py new file mode 100644 index 000000000..b520a1053 --- /dev/null +++ b/src/unilab/tasks/motion_tracking/common/manager_terms.py @@ -0,0 +1,776 @@ +"""Manager-native NumPy terms for motion tracking.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Literal, cast + +import numpy as np + +from unilab.envs.mdp.actions import JointPositionAction, JointPositionActionCfg +from unilab.managers import CommandTerm, CommandTermCfg, ManagerTermBase, ManagerTermBaseCfg +from unilab.managers.scene_entity_config import SceneEntityCfg +from unilab.utils.geometry import ( + np_write_relative_anchor_transform_pos_rot6d, +) +from unilab.utils.rotation import ( + np_quat_apply_inverse, + np_quat_error_magnitude_squared_batched, + np_quat_from_euler_xyz, + np_quat_mul, +) + +from .motion_loader import MotionData, MotionLoader, MotionSampler +from .observations import write_body_ori6_in_anchor_frame, write_body_pos_in_anchor_frame +from .transforms import update_relative_transforms + +if TYPE_CHECKING: + from unilab.base.entity import Entity + from unilab.managers._types import ManagerBasedRlEnv + + +SamplingMode = Literal["start", "clip_start", "uniform", "adaptive", "mixed"] +_RANGE_KEYS = ("x", "y", "z", "roll", "pitch", "yaw") +_DEFAULT_ASSET_CFG = SceneEntityCfg("robot") + + +def _range_matrix(value: dict[str, tuple[float, float]], *, name: str) -> np.ndarray: + unknown = sorted(set(value) - set(_RANGE_KEYS)) + if unknown: + raise ValueError(f"{name} has unknown axes {unknown}") + try: + ranges = np.asarray([value.get(key, (0.0, 0.0)) for key in _RANGE_KEYS], dtype=np.float64) + except (TypeError, ValueError) as exc: + raise TypeError(f"{name} must map axes to numeric (min, max) pairs") from exc + if ranges.shape != (6, 2) or not np.isfinite(ranges).all(): + raise ValueError(f"{name} must contain six finite (min, max) pairs") + if np.any(ranges[:, 0] > ranges[:, 1]): + raise ValueError(f"{name} contains a minimum greater than its maximum") + ranges.setflags(write=False) + return ranges + + +def _pair(value: tuple[float, float], *, name: str) -> tuple[float, float]: + try: + values = np.asarray(value, dtype=np.float64) + except (TypeError, ValueError) as exc: + raise TypeError(f"{name} must be a numeric (min, max) pair") from exc + if values.shape != (2,) or not np.isfinite(values).all(): + raise ValueError(f"{name} must be a finite (min, max) pair") + lower, upper = float(values[0]), float(values[1]) + if lower > upper: + raise ValueError(f"{name} minimum {lower} exceeds maximum {upper}") + return lower, upper + + +@dataclass +class MotionCommandParamsCfg: + """Hydra-owned motion data and sampling parameters.""" + + motion_file: str | list[str] + anchor_body_name: str + body_names: tuple[str, ...] | list[str] + sampling_mode: SamplingMode = "adaptive" + sampling_start_ratio: float = 0.0 + truncate_on_clip_end: bool = False + pose_range: dict[str, tuple[float, float]] = field(default_factory=dict) + velocity_range: dict[str, tuple[float, float]] = field(default_factory=dict) + joint_position_range: tuple[float, float] = (-0.1, 0.1) + joint_default_position_range: tuple[float, float] = (0.0, 0.0) + adaptive_lambda: float = 0.8 + adaptive_kernel_size: int = 1 + adaptive_uniform_ratio: float = 0.1 + adaptive_alpha: float = 0.001 + + +@dataclass(kw_only=True) +class MotionCommandCfg(CommandTermCfg): + """Community-shaped motion command with Hydra-owned nested parameters.""" + + entity_name: str + params: MotionCommandParamsCfg + + def build(self, env: ManagerBasedRlEnv) -> MotionCommand: + return MotionCommand(self, env) + + @property + def motion_file(self) -> str | list[str]: + return self.params.motion_file + + @property + def anchor_body_name(self) -> str: + return self.params.anchor_body_name + + @property + def body_names(self) -> tuple[str, ...]: + return tuple(self.params.body_names) + + @property + def sampling_mode(self) -> SamplingMode: + return self.params.sampling_mode + + +class MotionCommand(CommandTerm): + """Motion reference command on UniLab's NumPy/entity runtime.""" + + cfg: MotionCommandCfg + + def __init__(self, cfg: MotionCommandCfg, env: ManagerBasedRlEnv): + self._validate_cfg(cfg) + super().__init__(cfg, env) + self.robot = cast("Entity", env.scene[cfg.entity_name]) + body_ids, body_names = self.robot.find_bodies(cfg.body_names, preserve_order=True) + if tuple(body_names) != cfg.body_names: + raise ValueError( + f"MotionCommand body order {tuple(body_names)} does not match {cfg.body_names}" + ) + self._robot_body_ids = np.asarray(body_ids, dtype=np.intp) + self._robot_body_ids.setflags(write=False) + motion_body_ids = self.robot.motion_body_ids[self._robot_body_ids] + self.motion = MotionLoader(cfg.motion_file, body_indices=motion_body_ids) + if self.motion.num_joints != len(self.robot.joint_names): + raise ValueError( + f"MotionCommand motion joint width {self.motion.num_joints} does not match " + f"entity '{self.robot.name}' joint width {len(self.robot.joint_names)}" + ) + if self.motion.num_bodies != len(cfg.body_names): + raise ValueError( + f"MotionCommand motion body width {self.motion.num_bodies} does not match " + f"configured body width {len(cfg.body_names)}" + ) + + self.anchor_body_idx = cfg.body_names.index(cfg.anchor_body_name) + self.sampler = MotionSampler( + self.motion, + mode=cfg.params.sampling_mode, + num_envs=self.num_envs, + adaptive_lambda=cfg.params.adaptive_lambda, + adaptive_kernel_size=cfg.params.adaptive_kernel_size, + adaptive_uniform_ratio=cfg.params.adaptive_uniform_ratio, + adaptive_alpha=cfg.params.adaptive_alpha, + start_ratio=cfg.params.sampling_start_ratio, + rng=env.rng, + ) + self._pose_range = _range_matrix(cfg.params.pose_range, name="MotionCommand pose_range") + self._velocity_range = _range_matrix( + cfg.params.velocity_range, name="MotionCommand velocity_range" + ) + self._joint_position_range = _pair( + cfg.params.joint_position_range, + name="MotionCommand joint_position_range", + ) + self._joint_default_position_range = _pair( + cfg.params.joint_default_position_range, + name="MotionCommand joint_default_position_range", + ) + + num_bodies = len(cfg.body_names) + num_joints = self.motion.num_joints + dtype = self.motion.joint_pos.dtype + self.time_steps = self.sampler.current_frames + self._motion_data = self.motion.make_motion_data_buffer(self.num_envs) + self._command = np.empty((self.num_envs, num_joints * 2), dtype=dtype) + self._body_pos_w = np.empty((self.num_envs, num_bodies, 3), dtype=dtype) + self.body_pos_relative_w = np.empty_like(self._body_pos_w) + self.body_quat_relative_w = np.empty((self.num_envs, num_bodies, 4), dtype=dtype) + self.motion_anchor_pos_b = np.empty((self.num_envs, 3), dtype=dtype) + self.motion_anchor_ori_b = np.empty((self.num_envs, 6), dtype=dtype) + self.robot_body_pos_b = np.empty_like(self._body_pos_w) + self.robot_body_ori_b = np.empty((self.num_envs, num_bodies, 6), dtype=dtype) + self.joint_default_bias = np.zeros((self.num_envs, num_joints), dtype=dtype) + self._delta_pos_w = np.empty((self.num_envs, 3), dtype=dtype) + self._delta_ori_w = np.empty((self.num_envs, 4), dtype=dtype) + self._body_vec_error = np.empty_like(self._body_pos_w) + self._env_error = np.empty(self.num_envs, dtype=dtype) + self._reward_term = np.empty(self.num_envs, dtype=dtype) + self._robot_cache_step = -1 + self._robot_body_pos_w = np.empty_like(self._body_pos_w) + self._robot_body_quat_w = np.empty((self.num_envs, num_bodies, 4), dtype=dtype) + self._robot_body_lin_vel_w = np.empty_like(self._body_pos_w) + self._robot_body_ang_vel_w = np.empty_like(self._body_pos_w) + + for name in ( + "error_anchor_pos", + "error_anchor_rot", + "error_anchor_lin_vel", + "error_anchor_ang_vel", + "error_body_pos", + "error_body_rot", + "error_body_lin_vel", + "error_body_ang_vel", + "error_joint_pos", + "error_joint_vel", + "sampling_entropy", + "sampling_top1_prob", + "sampling_top1_bin", + ): + self.metrics[name] = np.zeros(self.num_envs, dtype=dtype) + self._refresh_motion() + self._refresh_robot_state(force=True) + self._refresh_relative_state() + + @staticmethod + def _validate_cfg(cfg: MotionCommandCfg) -> None: + if not isinstance(cfg.entity_name, str) or not cfg.entity_name: + raise ValueError("MotionCommandCfg entity_name must be non-empty") + if not isinstance(cfg.params, MotionCommandParamsCfg): + raise TypeError("MotionCommandCfg params must be MotionCommandParamsCfg") + if not cfg.motion_file: + raise ValueError("MotionCommandCfg motion_file must be configured") + if not cfg.anchor_body_name or cfg.anchor_body_name not in cfg.body_names: + raise ValueError("MotionCommandCfg anchor_body_name must occur in body_names") + if len(set(cfg.body_names)) != len(cfg.body_names): + raise ValueError("MotionCommandCfg body_names must be unique") + if cfg.sampling_mode not in ("start", "clip_start", "uniform", "adaptive", "mixed"): + raise ValueError( + f"MotionCommandCfg has unsupported sampling_mode {cfg.sampling_mode!r}" + ) + if not 0.0 <= cfg.params.sampling_start_ratio <= 1.0: + raise ValueError("MotionCommandCfg sampling_start_ratio must be within [0, 1]") + if not isinstance(cfg.params.truncate_on_clip_end, bool): + raise TypeError("MotionCommandCfg truncate_on_clip_end must be bool") + + @property + def command(self) -> np.ndarray: + return self._command + + @property + def joint_pos(self) -> np.ndarray: + return self._motion_data.joint_pos + + @property + def joint_vel(self) -> np.ndarray: + return self._motion_data.joint_vel + + @property + def body_pos_w(self) -> np.ndarray: + return self._body_pos_w + + @property + def body_quat_w(self) -> np.ndarray: + return self._motion_data.body_quat_w + + @property + def body_lin_vel_w(self) -> np.ndarray: + return self._motion_data.body_lin_vel_w + + @property + def body_ang_vel_w(self) -> np.ndarray: + return self._motion_data.body_ang_vel_w + + @property + def anchor_pos_w(self) -> np.ndarray: + return self._body_pos_w[:, self.anchor_body_idx] + + @property + def anchor_quat_w(self) -> np.ndarray: + return self._motion_data.body_quat_w[:, self.anchor_body_idx] + + @property + def anchor_lin_vel_w(self) -> np.ndarray: + return self._motion_data.body_lin_vel_w[:, self.anchor_body_idx] + + @property + def anchor_ang_vel_w(self) -> np.ndarray: + return self._motion_data.body_ang_vel_w[:, self.anchor_body_idx] + + @property + def robot_joint_pos(self) -> np.ndarray: + return self.robot.data.joint_pos + + @property + def robot_joint_vel(self) -> np.ndarray: + return self.robot.data.joint_vel + + @property + def robot_body_pos_w(self) -> np.ndarray: + self._refresh_robot_state() + return self._robot_body_pos_w + + @property + def robot_body_quat_w(self) -> np.ndarray: + self._refresh_robot_state() + return self._robot_body_quat_w + + @property + def robot_body_lin_vel_w(self) -> np.ndarray: + self._refresh_robot_state() + return self._robot_body_lin_vel_w + + @property + def robot_body_ang_vel_w(self) -> np.ndarray: + self._refresh_robot_state() + return self._robot_body_ang_vel_w + + @property + def robot_anchor_pos_w(self) -> np.ndarray: + return self.robot_body_pos_w[:, self.anchor_body_idx] + + @property + def robot_anchor_quat_w(self) -> np.ndarray: + return self.robot_body_quat_w[:, self.anchor_body_idx] + + @property + def robot_anchor_lin_vel_w(self) -> np.ndarray: + return self.robot_body_lin_vel_w[:, self.anchor_body_idx] + + @property + def robot_anchor_ang_vel_w(self) -> np.ndarray: + return self.robot_body_ang_vel_w[:, self.anchor_body_idx] + + def reset(self, env_ids: np.ndarray | slice | None) -> dict[str, float]: + ids = ( + np.arange(self.num_envs, dtype=np.int32) + if env_ids is None + else np.arange(self.num_envs, dtype=np.int32)[env_ids] + if isinstance(env_ids, slice) + else env_ids + ) + lower, upper = self._joint_default_position_range + self.joint_default_bias[ids] = self._env.rng.uniform( + lower, upper, size=(len(ids), self.motion.num_joints) + ) + return super().reset(ids) + + def _refresh_motion(self) -> None: + self.motion.get_motion_at_frame(self.time_steps, out=self._motion_data) + np.add( + self._motion_data.body_pos_w, + self._env.scene.env_origins[:, None, :], + out=self._body_pos_w, + ) + width = self.motion.num_joints + self._command[:, :width] = self._motion_data.joint_pos + self._command[:, width:] = self._motion_data.joint_vel + + def _refresh_robot_state(self, *, force: bool = False) -> None: + step = self._env.common_step_counter + if not force and self._robot_cache_step == step: + return + body_index = self._robot_body_ids + self._robot_body_pos_w[:] = self.robot.data.body_link_pos_w[:, body_index] + self._robot_body_quat_w[:] = self.robot.data.body_link_quat_w[:, body_index] + self._robot_body_lin_vel_w[:] = self.robot.data.body_link_lin_vel_w[:, body_index] + self._robot_body_ang_vel_w[:] = self.robot.data.body_link_ang_vel_w[:, body_index] + self._robot_cache_step = step + + def _refresh_relative_state(self) -> None: + update_relative_transforms( + self, + self._motion_data, + self._robot_body_pos_w, + self._robot_body_quat_w, + ) + np_write_relative_anchor_transform_pos_rot6d( + self.robot_anchor_pos_w, + self.robot_anchor_quat_w, + self.anchor_pos_w, + self.anchor_quat_w, + self.motion_anchor_pos_b, + self.motion_anchor_ori_b, + ) + write_body_pos_in_anchor_frame( + self.robot_anchor_pos_w, + self.robot_anchor_quat_w, + self._robot_body_pos_w, + self.robot_body_pos_b, + body_vec_error=self._body_vec_error, + ) + write_body_ori6_in_anchor_frame( + self.robot_anchor_quat_w, + self._robot_body_quat_w, + self.robot_body_ori_b, + ) + + def _update_metrics(self) -> None: + self.metrics["error_anchor_pos"][:] = np.linalg.norm( + self.anchor_pos_w - self.robot_anchor_pos_w, axis=-1 + ) + self.metrics["error_anchor_rot"][:] = np.sqrt( + np_quat_error_magnitude_squared_batched(self.anchor_quat_w, self.robot_anchor_quat_w) + ) + self.metrics["error_anchor_lin_vel"][:] = np.linalg.norm( + self.anchor_lin_vel_w - self.robot_anchor_lin_vel_w, axis=-1 + ) + self.metrics["error_anchor_ang_vel"][:] = np.linalg.norm( + self.anchor_ang_vel_w - self.robot_anchor_ang_vel_w, axis=-1 + ) + self.metrics["error_body_pos"][:] = np.linalg.norm( + self.body_pos_relative_w - self.robot_body_pos_w, axis=-1 + ).mean(axis=-1) + self.metrics["error_body_rot"][:] = np.sqrt( + np_quat_error_magnitude_squared_batched( + self.body_quat_relative_w, self.robot_body_quat_w + ) + ).mean(axis=-1) + self.metrics["error_body_lin_vel"][:] = np.linalg.norm( + self.body_lin_vel_w - self.robot_body_lin_vel_w, axis=-1 + ).mean(axis=-1) + self.metrics["error_body_ang_vel"][:] = np.linalg.norm( + self.body_ang_vel_w - self.robot_body_ang_vel_w, axis=-1 + ).mean(axis=-1) + self.metrics["error_joint_pos"][:] = np.linalg.norm( + self.joint_pos - self.robot_joint_pos, axis=-1 + ) + self.metrics["error_joint_vel"][:] = np.linalg.norm( + self.joint_vel - self.robot_joint_vel, axis=-1 + ) + self.metrics["sampling_entropy"].fill(self.sampler.sampling_entropy) + self.metrics["sampling_top1_prob"].fill(self.sampler.sampling_top1_prob) + self.metrics["sampling_top1_bin"].fill(self.sampler.sampling_top1_bin) + + def _resample_command(self, env_ids: np.ndarray) -> None: + frames = self.sampler.sample_frames(env_ids) + motion = self.motion.get_motion_at_frame(frames) + count = len(env_ids) + pose = self._env.rng.uniform( + self._pose_range[:, 0], self._pose_range[:, 1], size=(count, 6) + ) + velocity = self._env.rng.uniform( + self._velocity_range[:, 0], self._velocity_range[:, 1], size=(count, 6) + ) + root_pos = motion.body_pos_w[:, 0].copy() + root_pos += self._env.scene.env_origins[env_ids] + root_pos += pose[:, :3] + root_quat = np_quat_mul( + np_quat_from_euler_xyz(pose[:, 3], pose[:, 4], pose[:, 5]), + motion.body_quat_w[:, 0], + ) + root_lin_vel = motion.body_lin_vel_w[:, 0] + velocity[:, :3] + root_ang_vel = motion.body_ang_vel_w[:, 0] + velocity[:, 3:] + joint_pos = motion.joint_pos.copy() + joint_pos += self._env.rng.uniform( + *self._joint_position_range, + size=joint_pos.shape, + ) + limits = self.robot.data.soft_joint_pos_limits + np.clip(joint_pos, limits[:, 0], limits[:, 1], out=joint_pos) + self.robot.write_joint_state_to_sim(joint_pos, motion.joint_vel, env_ids=env_ids) + root_state = np.concatenate((root_pos, root_quat, root_lin_vel, root_ang_vel), axis=-1) + self.robot.write_root_state_to_sim(root_state, env_ids=env_ids) + + def _update_command(self, env_ids: np.ndarray | None) -> None: + if env_ids is not None: + self._refresh_motion() + return + self.sampler.update_failure_stats(self._env.termination_manager.terminated) + active_ids = np.flatnonzero(~self._env.reset_buf).astype(np.int32, copy=False) + wrap_ids = self.sampler.step(active_ids) + if len(wrap_ids) and not self.cfg.params.truncate_on_clip_end: + self._resample_command(wrap_ids) + self._refresh_motion() + + def post_compute(self) -> None: + self._refresh_robot_state(force=True) + self._refresh_relative_state() + + +@dataclass(kw_only=True) +class MotionJointPositionActionCfg(JointPositionActionCfg): + command_name: str = "motion" + + def build(self, env: ManagerBasedRlEnv) -> MotionJointPositionAction: + return MotionJointPositionAction(self, env) + + +class MotionJointPositionAction(JointPositionAction): + cfg: MotionJointPositionActionCfg # pyright: ignore[reportIncompatibleVariableOverride] + + def __init__(self, cfg: MotionJointPositionActionCfg, env: ManagerBasedRlEnv): + super().__init__(cfg, env) + self._motion_command = _command(env, cfg.command_name) + + def apply_actions(self) -> None: + encoder_bias = self._entity.data.encoder_bias[:, self._target_ids] + np.add( + self._processed_actions, + self._motion_command.joint_default_bias[:, self._target_ids], + out=self._target, + ) + self._target -= encoder_bias + self._entity.set_joint_position_target(self._target, joint_ids=self._target_ids) + + +def _command(env: ManagerBasedRlEnv, command_name: str) -> MotionCommand: + try: + command = env.command_manager.get_term(command_name) + except KeyError as exc: + raise KeyError(f"Motion command term '{command_name}' not found") from exc + if not isinstance(command, MotionCommand): + raise TypeError( + f"Command term '{command_name}' is {type(command).__name__}, expected MotionCommand" + ) + return command + + +def motion_anchor_pos_b(env: ManagerBasedRlEnv, command_name: str) -> np.ndarray: + return _command(env, command_name).motion_anchor_pos_b + + +def motion_anchor_ori_b(env: ManagerBasedRlEnv, command_name: str) -> np.ndarray: + return _command(env, command_name).motion_anchor_ori_b + + +def robot_body_pos_b(env: ManagerBasedRlEnv, command_name: str) -> np.ndarray: + command = _command(env, command_name) + return command.robot_body_pos_b.reshape(env.num_envs, -1) + + +def robot_body_ori_b(env: ManagerBasedRlEnv, command_name: str) -> np.ndarray: + command = _command(env, command_name) + return command.robot_body_ori_b.reshape(env.num_envs, -1) + + +def motion_joint_pos_rel(env: ManagerBasedRlEnv, command_name: str) -> np.ndarray: + command = _command(env, command_name) + return ( + command.robot_joint_pos - command.robot.data.default_joint_pos - command.joint_default_bias + ) + + +def _positive_std(value: float, *, term_name: str) -> float: + if isinstance(value, bool) or not isinstance(value, (int, float, np.number)): + raise TypeError(f"{term_name} std must be a real number") + result = float(value) + if not math.isfinite(result) or result <= 0.0: + raise ValueError(f"{term_name} std must be finite and positive") + return result + + +def motion_global_anchor_position_error_exp( + env: ManagerBasedRlEnv, command_name: str, std: float +) -> np.ndarray: + command = _command(env, command_name) + error = np.sum(np.square(command.anchor_pos_w - command.robot_anchor_pos_w), axis=-1) + return np.exp(-error / _positive_std(std, term_name="motion anchor position") ** 2) + + +def motion_global_anchor_orientation_error_exp( + env: ManagerBasedRlEnv, command_name: str, std: float +) -> np.ndarray: + command = _command(env, command_name) + error = np_quat_error_magnitude_squared_batched( + command.anchor_quat_w, command.robot_anchor_quat_w + ) + return np.exp(-error / _positive_std(std, term_name="motion anchor orientation") ** 2) + + +class _BodyTerm(ManagerTermBase): + def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv): + super().__init__(env) + command_name = cfg.params.get("command_name") + if not isinstance(command_name, str) or not command_name: + raise ValueError(f"{type(self).__name__} requires a non-empty command_name") + self._command_name = command_name + command = _command(env, command_name) + body_names = cfg.params.get("body_names") + if body_names is None: + self._body_ids = slice(None) + else: + requested = tuple(body_names) + missing = [name for name in requested if name not in command.cfg.body_names] + if missing: + raise ValueError( + f"Body names {missing} are not tracked by command '{command_name}'" + ) + self._body_ids = np.asarray( + [command.cfg.body_names.index(name) for name in requested], dtype=np.intp + ) + + def _validate(self, command_name: str, std: float) -> tuple[MotionCommand, float]: + if command_name != self._command_name: + raise ValueError( + f"{type(self).__name__} was bound to '{self._command_name}', got '{command_name}'" + ) + return _command(self._env, command_name), _positive_std(std, term_name=type(self).__name__) + + +class motion_relative_body_position_error_exp(_BodyTerm): + def __call__( + self, + env: ManagerBasedRlEnv, + command_name: str, + std: float, + body_names: tuple[str, ...] | None = None, + ) -> np.ndarray: + del env, body_names + command, scale = self._validate(command_name, std) + error = np.sum( + np.square( + command.body_pos_relative_w[:, self._body_ids] + - command.robot_body_pos_w[:, self._body_ids] + ), + axis=-1, + ) + return np.exp(-error.mean(axis=-1) / scale**2) + + +class motion_relative_body_orientation_error_exp(_BodyTerm): + def __call__( + self, + env: ManagerBasedRlEnv, + command_name: str, + std: float, + body_names: tuple[str, ...] | None = None, + ) -> np.ndarray: + del env, body_names + command, scale = self._validate(command_name, std) + error = np_quat_error_magnitude_squared_batched( + command.body_quat_relative_w[:, self._body_ids], + command.robot_body_quat_w[:, self._body_ids], + ) + return np.exp(-error.mean(axis=-1) / scale**2) + + +class motion_global_body_linear_velocity_error_exp(_BodyTerm): + def __call__( + self, + env: ManagerBasedRlEnv, + command_name: str, + std: float, + body_names: tuple[str, ...] | None = None, + ) -> np.ndarray: + del env, body_names + command, scale = self._validate(command_name, std) + error = np.sum( + np.square( + command.body_lin_vel_w[:, self._body_ids] + - command.robot_body_lin_vel_w[:, self._body_ids] + ), + axis=-1, + ) + return np.exp(-error.mean(axis=-1) / scale**2) + + +class motion_global_body_angular_velocity_error_exp(_BodyTerm): + def __call__( + self, + env: ManagerBasedRlEnv, + command_name: str, + std: float, + body_names: tuple[str, ...] | None = None, + ) -> np.ndarray: + del env, body_names + command, scale = self._validate(command_name, std) + error = np.sum( + np.square( + command.body_ang_vel_w[:, self._body_ids] + - command.robot_body_ang_vel_w[:, self._body_ids] + ), + axis=-1, + ) + return np.exp(-error.mean(axis=-1) / scale**2) + + +def motion_joint_position_error_exp( + env: ManagerBasedRlEnv, command_name: str, std: float +) -> np.ndarray: + command = _command(env, command_name) + error = np.square(command.joint_pos - command.robot_joint_pos).mean(axis=-1) + return np.exp(-error / _positive_std(std, term_name="motion joint position") ** 2) + + +def motion_joint_velocity_error_exp( + env: ManagerBasedRlEnv, command_name: str, std: float +) -> np.ndarray: + command = _command(env, command_name) + error = np.square(command.joint_vel - command.robot_joint_vel).mean(axis=-1) + return np.exp(-error / _positive_std(std, term_name="motion joint velocity") ** 2) + + +def joint_pos_limits( + env: ManagerBasedRlEnv, + asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG, +) -> np.ndarray: + """Penalize selected joint-limit violations through the entity facade.""" + asset = cast("Entity", env.scene[asset_cfg.name]) + joint_pos = asset.data.joint_pos[:, asset_cfg.joint_ids] + limits = asset.data.soft_joint_pos_limits[asset_cfg.joint_ids] + lower_error = np.maximum(limits[:, 0] - joint_pos, 0.0) + upper_error = np.maximum(joint_pos - limits[:, 1], 0.0) + return np.sum(np.square(lower_error + upper_error), axis=-1) + + +class undesired_body_contacts(_BodyTerm): + def __call__( + self, + env: ManagerBasedRlEnv, + command_name: str, + threshold: float, + body_names: tuple[str, ...] | None = None, + ) -> np.ndarray: + del env, body_names + command = _command(self._env, command_name) + return np.sum(command.robot_body_pos_w[:, self._body_ids, 2] < threshold, axis=-1) + + +def bad_anchor_pos_z_only( + env: ManagerBasedRlEnv, command_name: str, threshold: float +) -> np.ndarray: + command = _command(env, command_name) + return np.abs(command.anchor_pos_w[:, 2] - command.robot_anchor_pos_w[:, 2]) > threshold + + +def bad_anchor_ori( + env: ManagerBasedRlEnv, + command_name: str, + threshold: float, + asset_cfg: SceneEntityCfg | None = None, +) -> np.ndarray: + command = _command(env, command_name) + asset = command.robot if asset_cfg is None else cast("Entity", env.scene[asset_cfg.name]) + gravity_vec_w = asset.data.gravity_vec_w + motion_z = np_quat_apply_inverse(command.anchor_quat_w, gravity_vec_w)[:, 2] + robot_z = np_quat_apply_inverse(command.robot_anchor_quat_w, gravity_vec_w)[:, 2] + return np.abs(motion_z - robot_z) > threshold + + +class bad_motion_body_pos_z_only(_BodyTerm): + def __call__( + self, + env: ManagerBasedRlEnv, + command_name: str, + threshold: float, + body_names: tuple[str, ...] | None = None, + ) -> np.ndarray: + del env, body_names + command = _command(self._env, command_name) + error = np.abs( + command.body_pos_relative_w[:, self._body_ids, 2] + - command.robot_body_pos_w[:, self._body_ids, 2] + ) + return np.any(error > threshold, axis=-1) + + +def motion_clip_end(env: ManagerBasedRlEnv, command_name: str) -> np.ndarray: + command = _command(env, command_name) + return command.time_steps >= command.sampler.current_clip_end_frames + + +__all__ = [ + "MotionCommand", + "MotionCommandCfg", + "MotionCommandParamsCfg", + "MotionJointPositionAction", + "MotionJointPositionActionCfg", + "bad_anchor_ori", + "bad_anchor_pos_z_only", + "bad_motion_body_pos_z_only", + "joint_pos_limits", + "motion_anchor_ori_b", + "motion_anchor_pos_b", + "motion_clip_end", + "motion_global_anchor_orientation_error_exp", + "motion_global_anchor_position_error_exp", + "motion_global_body_angular_velocity_error_exp", + "motion_global_body_linear_velocity_error_exp", + "motion_joint_pos_rel", + "motion_joint_position_error_exp", + "motion_joint_velocity_error_exp", + "motion_relative_body_orientation_error_exp", + "motion_relative_body_position_error_exp", + "robot_body_ori_b", + "robot_body_pos_b", + "undesired_body_contacts", +] diff --git a/src/unilab/tasks/motion_tracking/common/motion_loader.py b/src/unilab/tasks/motion_tracking/common/motion_loader.py index b11b497aa..62e438dd2 100644 --- a/src/unilab/tasks/motion_tracking/common/motion_loader.py +++ b/src/unilab/tasks/motion_tracking/common/motion_loader.py @@ -214,6 +214,7 @@ def __init__( adaptive_uniform_ratio: float = 0.1, adaptive_alpha: float = 0.001, start_ratio: float = 0.0, + rng: np.random.Generator | None = None, ): """Initialize motion sampler. @@ -236,6 +237,7 @@ def __init__( self.mode = mode self.num_envs = num_envs self.start_ratio = start_ratio + self.rng = rng # Current frame indices for each environment self.current_frames = np.zeros(num_envs, dtype=np.int32) @@ -306,16 +308,28 @@ def _sample_clip_start(self, env_ids: np.ndarray) -> np.ndarray: if self.motion_loader.num_clips == 1: frames = np.zeros(len(env_ids), dtype=np.int32) else: - clip_indices = np.random.randint( - 0, self.motion_loader.num_clips, len(env_ids), dtype=np.int32 - ) + if self.rng is None: + clip_indices = np.random.randint( + 0, self.motion_loader.num_clips, len(env_ids), dtype=np.int32 + ) + else: + clip_indices = self.rng.integers( + 0, self.motion_loader.num_clips, len(env_ids), dtype=np.int32 + ) frames = np.asarray(self.motion_loader.clip_offsets[clip_indices], dtype=np.int32) self._set_sampled_frames(env_ids, frames) return frames def _sample_uniform(self, env_ids: np.ndarray) -> np.ndarray: """Sample uniformly across motion.""" - frames = np.random.randint(0, self.motion_loader.num_frames, len(env_ids), dtype=np.int32) + if self.rng is None: + frames = np.random.randint( + 0, self.motion_loader.num_frames, len(env_ids), dtype=np.int32 + ) + else: + frames = self.rng.integers( + 0, self.motion_loader.num_frames, len(env_ids), dtype=np.int32 + ) self._set_sampled_frames(env_ids, frames) # Update metrics @@ -334,11 +348,16 @@ def _sample_mixed(self, env_ids: np.ndarray) -> np.ndarray: uniform RSI's whole-clip coverage everywhere else. """ n = len(env_ids) - use_start = np.random.random(n) < self.start_ratio + random_values = np.random.random(n) if self.rng is None else self.rng.random(n) + if self.rng is None: + uniform_frames = np.random.randint(0, self.motion_loader.num_frames, n) + else: + uniform_frames = self.rng.integers(0, self.motion_loader.num_frames, n) + use_start = random_values < self.start_ratio frames = np.where( use_start, 0, - np.random.randint(0, self.motion_loader.num_frames, n), + uniform_frames, ).astype(np.int32) self._set_sampled_frames(env_ids, frames) @@ -374,10 +393,18 @@ def _sample_adaptive(self, env_ids: np.ndarray) -> np.ndarray: sampling_probs = sampling_probs / sampling_probs.sum() # Sample bins - sampled_bins = np.random.choice(self.bin_count, size=len(env_ids), p=sampling_probs) + sampled_bins = ( + np.random.choice(self.bin_count, size=len(env_ids), p=sampling_probs) + if self.rng is None + else self.rng.choice(self.bin_count, size=len(env_ids), p=sampling_probs) + ) # Add random offset within bin - bin_offsets = np.random.uniform(0.0, 1.0, len(env_ids)) + bin_offsets = ( + np.random.uniform(0.0, 1.0, len(env_ids)) + if self.rng is None + else self.rng.uniform(0.0, 1.0, len(env_ids)) + ) frames = ( (sampled_bins + bin_offsets) / self.bin_count * (self.motion_loader.num_frames - 1) ).astype(np.int32) @@ -437,12 +464,14 @@ def _set_sampled_frames(self, env_ids: np.ndarray, frames: np.ndarray) -> None: self.current_clip_indices[env_ids] = clip_indices self.current_clip_end_frames[env_ids] = self.motion_loader.clip_end_frames[clip_indices] - def step(self): - """Advance all frames by one step.""" - self.current_frames += 1 + def step(self, env_ids: np.ndarray | None = None) -> np.ndarray: + """Advance selected frames by one step and return their clip-end rows.""" + ids = np.arange(self.num_envs, dtype=np.int32) if env_ids is None else env_ids + self.current_frames[ids] += 1 # Find environments that reached the end of their current clip. - np.greater(self.current_frames, self.current_clip_end_frames, out=self._done_mask) + self._done_mask.fill(False) + self._done_mask[ids] = self.current_frames[ids] > self.current_clip_end_frames[ids] return np.flatnonzero(self._done_mask) def get_current_motion(self, out: MotionData | None = None) -> MotionData: diff --git a/src/unilab/tasks/motion_tracking/g1/tracking.py b/src/unilab/tasks/motion_tracking/g1/tracking.py index 58cb5d7a4..de3618803 100644 --- a/src/unilab/tasks/motion_tracking/g1/tracking.py +++ b/src/unilab/tasks/motion_tracking/g1/tracking.py @@ -14,6 +14,7 @@ from unilab.assets import ASSETS_ROOT_PATH from unilab.base import registry from unilab.base.scene import SceneCfg +from unilab.envs import ManagerBasedRlEnvCfg, make_manager_based_rl_env from ..common.config import ( Domain_Rand, @@ -41,7 +42,6 @@ _build_motion_reference_state = build_motion_reference_state -@registry.envcfg("G1MotionTracking") @dataclass class G1MotionTrackingEnvCfg(MotionTrackingCfg): """Registered configuration for G1 motion tracking.""" @@ -49,7 +49,6 @@ class G1MotionTrackingEnvCfg(MotionTrackingCfg): pass -@registry.envcfg("G1MotionTrackingDeploy") @dataclass class G1MotionTrackingDeployEnvCfg(MotionTrackingDeployEnvCfg): """Registered deploy configuration for G1 motion tracking.""" @@ -91,44 +90,40 @@ class G1MotionTracking23DofCfg(G1MotionTrackingCfg): ) -@registry.envcfg("G1MotionTracking23Dof") @dataclass class G1MotionTracking23DofEnvCfg(G1MotionTracking23DofCfg): pass -@registry.envcfg("G1MotionTracking23DofDeploy") @dataclass class G1MotionTracking23DofDeployEnvCfg(G1MotionTracking23DofCfg): pass -@registry.env("G1MotionTracking", sim_backend="drake") -@registry.env("G1MotionTracking", sim_backend="mujoco") -@registry.env("G1MotionTracking", sim_backend="motrix") class G1MotionTrackingEnv(MotionTrackingEnv): """G1 Motion Tracking Environment.""" _cfg: MotionTrackingCfg -@registry.env("G1MotionTrackingDeploy", sim_backend="mujoco") -@registry.env("G1MotionTrackingDeploy", sim_backend="motrix") class G1MotionTrackingDeployEnv(MotionTrackingDeployEnv): """Deploy-oriented G1 motion tracking env with unitree_rl_lab mimic actor inputs.""" _cfg: MotionTrackingDeployEnvCfg -# --- 23-DoF env registrations (same env classes, 23-DoF configs) --- -registry.register_env("G1MotionTracking23Dof", G1MotionTrackingEnv, sim_backend="mujoco") -registry.register_env("G1MotionTracking23Dof", G1MotionTrackingEnv, sim_backend="motrix") -registry.register_env( - "G1MotionTracking23DofDeploy", G1MotionTrackingDeployEnv, sim_backend="mujoco" -) -registry.register_env( - "G1MotionTracking23DofDeploy", G1MotionTrackingDeployEnv, sim_backend="motrix" -) +# The legacy classes above remain explicit consumers for profiles that are not +# part of #1227. The four core identities have one Hydra-owned manager config +# and one generic runtime factory instead of inheriting those classes. +for _task_name in ( + "G1MotionTracking", + "G1MotionTrackingDeploy", + "G1MotionTracking23Dof", + "G1MotionTracking23DofDeploy", +): + registry.register_env_config(_task_name, ManagerBasedRlEnvCfg) + registry.register_env(_task_name, make_manager_based_rl_env, sim_backend="mujoco") + registry.register_env(_task_name, make_manager_based_rl_env, sim_backend="motrix") __all__ = [ diff --git a/src/unilab/tasks/motion_tracking/g1/tracking_sac.py b/src/unilab/tasks/motion_tracking/g1/tracking_sac.py index 1654838f4..a7e3c5e36 100644 --- a/src/unilab/tasks/motion_tracking/g1/tracking_sac.py +++ b/src/unilab/tasks/motion_tracking/g1/tracking_sac.py @@ -15,6 +15,7 @@ from unilab.base import registry from unilab.dtype_config import get_global_dtype +from unilab.envs import ManagerBasedRlEnvCfg, make_manager_based_rl_env from .tracking import ( G1MotionTracking23DofCfg, @@ -23,14 +24,11 @@ ) -@registry.envcfg("G1MotionTrackingSAC") @dataclass class G1MotionTrackingSACCfg(G1MotionTrackingCfg): """Config for SAC-based motion tracking (identical fields, separate registry entry).""" -@registry.env("G1MotionTrackingSAC", sim_backend="mujoco") -@registry.env("G1MotionTrackingSAC", sim_backend="motrix") class G1MotionTrackingSACEnv(G1MotionTrackingEnv): """G1 Motion Tracking environment for FastSAC training. @@ -75,13 +73,16 @@ def _compute_obs( return obs -@registry.envcfg("G1MotionTrackingSAC23Dof") @dataclass class G1MotionTrackingSAC23DofCfg(G1MotionTracking23DofCfg): pass -@registry.env("G1MotionTrackingSAC23Dof", sim_backend="mujoco") -@registry.env("G1MotionTrackingSAC23Dof", sim_backend="motrix") class G1MotionTrackingSAC23DofEnv(G1MotionTrackingSACEnv): _cfg: G1MotionTrackingSAC23DofCfg # pyright: ignore[reportIncompatibleVariableOverride] + + +for _task_name in ("G1MotionTrackingSAC", "G1MotionTrackingSAC23Dof"): + registry.register_env_config(_task_name, ManagerBasedRlEnvCfg) + registry.register_env(_task_name, make_manager_based_rl_env, sim_backend="mujoco") + registry.register_env(_task_name, make_manager_based_rl_env, sim_backend="motrix") diff --git a/src/unilab/training/sim2sim.py b/src/unilab/training/sim2sim.py index d7e35426d..1489c84d3 100644 --- a/src/unilab/training/sim2sim.py +++ b/src/unilab/training/sim2sim.py @@ -44,19 +44,34 @@ class CrossBackendIncompatibleError(RuntimeError): "algo.policy.critic_hidden_dims", "algo.empirical_normalization", "algo.obs_normalization", - "env.sampling_mode", + "env.commands.motion.params.sampling_mode", ] SNAPSHOT_FIELDS: list[str] = DENYLIST + WARNING_LIST ENV_STRUCTURAL_DENYLIST: list[str] = [path for path in DENYLIST if path.startswith("env.")] +LEGACY_PATH_ALIASES: dict[str, str] = { + "env.sampling_mode": "env.commands.motion.params.sampling_mode", +} +_CANONICAL_PATH_FALLBACKS: dict[str, str] = { + canonical: legacy for legacy, canonical in LEGACY_PATH_ALIASES.items() +} + def _select(cfg: Any, path: str) -> Any: """Return the effective value at a dotted path (or ``None`` if absent).""" return OmegaConf.select(cfg, path) +def _select_contract(cfg: Any, path: str) -> Any: + value = _select(cfg, path) + if value is not None: + return value + legacy_path = _CANONICAL_PATH_FALLBACKS.get(path) + return None if legacy_path is None else _select(cfg, legacy_path) + + def _to_plain(value: Any) -> Any: if OmegaConf.is_config(value): return OmegaConf.to_container(value, resolve=True) @@ -68,7 +83,7 @@ def extract_contract_snapshot(full_cfg: DictConfig) -> dict[str, Any]: cfg: Any = full_cfg if OmegaConf.is_config(full_cfg) else OmegaConf.create(full_cfg) snapshot: dict[str, Any] = {} for path in SNAPSHOT_FIELDS: - value = _select(cfg, path) + value = _select_contract(cfg, path) if value is None: continue snapshot[path] = _to_plain(value) @@ -162,8 +177,13 @@ def resolve_sim2sim_config( return target_cfg denials: list[str] = [] - for path, source_value in snapshot.items(): - target_value = _select(target_cfg, path) + canonical_snapshot: dict[str, Any] = {} + for raw_path, source_value in snapshot.items(): + path = LEGACY_PATH_ALIASES.get(raw_path, raw_path) + canonical_snapshot[path] = source_value + + for path, source_value in canonical_snapshot.items(): + target_value = _select_contract(target_cfg, path) if target_value is None: if path in ENV_STRUCTURAL_DENYLIST: denials.append(_asymmetric_line(path, source_value, source_present=True)) @@ -177,10 +197,11 @@ def resolve_sim2sim_config( print(f"[sim2sim] WARNING override {line}") for path in ENV_STRUCTURAL_DENYLIST: - if path in snapshot: + if path in canonical_snapshot: continue - if _select(target_cfg, path) is not None: - denials.append(_asymmetric_line(path, _select(target_cfg, path), source_present=False)) + target_value = _select_contract(target_cfg, path) + if target_value is not None: + denials.append(_asymmetric_line(path, target_value, source_present=False)) if denials: message = ( diff --git a/tests/base/test_entity_facade.py b/tests/base/test_entity_facade.py index ba1f9ce34..857f1314e 100644 --- a/tests/base/test_entity_facade.py +++ b/tests/base/test_entity_facade.py @@ -59,6 +59,10 @@ def __init__(self, backend_type: str, *, unsupported: frozenset[str] = frozenset dtype=np.float32, ) self.set_state_calls: list[tuple[np.ndarray, np.ndarray, np.ndarray]] = [] + self.joint_range = np.asarray( + [[-1.0, 1.0], [-2.0, 2.0], [-3.0, 3.0], [-4.0, 4.0], [-5.0, 5.0]], + dtype=np.float32, + ) def _check(self, capability: str) -> None: self.calls[capability] += 1 @@ -133,6 +137,10 @@ def get_dof_vel(self) -> np.ndarray: self._check("joint velocity state") return self.dof_vel + def get_joint_range(self) -> np.ndarray: + self._check("joint position limits") + return self.joint_range + def get_body_pos_w(self, ids: np.ndarray) -> np.ndarray: self._check("body position state") return self.body_pos[:, ids] @@ -186,6 +194,8 @@ def test_backend_profiles_materialize_identical_local_entity_contract(backend_ty np.testing.assert_array_equal(robot.data.joint_pos, backend.dof_pos[:, [4, 2]]) np.testing.assert_array_equal(robot.data.joint_vel, backend.dof_vel[:, [4, 2]]) np.testing.assert_array_equal(robot.data.default_joint_vel, 0.0) + np.testing.assert_array_equal(robot.data.soft_joint_pos_limits, backend.joint_range[[4, 2]]) + np.testing.assert_array_equal(robot.data.gravity_vec_w, [[0.0, 0.0, -1.0]] * 3) np.testing.assert_array_equal(robot.data.joint_pos_biased, backend.dof_pos[:, [4, 2]]) np.testing.assert_array_equal(robot.data.body_link_pos_w, backend.body_pos[:, [7, 4]]) np.testing.assert_array_equal(robot.data.root_link_pos_w, backend.body_pos[:, 4]) @@ -198,6 +208,8 @@ def test_backend_profiles_materialize_identical_local_entity_contract(backend_ty np.tile([1.0, 2.0, 3.0, 1.0, 0.0, 0.0, 0.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0], (3, 1)), ) assert not robot.data.default_root_state.flags.writeable + assert not robot.data.soft_joint_pos_limits.flags.writeable + assert not robot.data.gravity_vec_w.flags.writeable np.testing.assert_array_equal( robot.data.actuator_ctrl_range, np.arange(10, dtype=np.float32).reshape(5, 2)[[4, 2]], diff --git a/tests/base/test_sim_backend.py b/tests/base/test_sim_backend.py index 88eb1a8e3..17996fe11 100644 --- a/tests/base/test_sim_backend.py +++ b/tests/base/test_sim_backend.py @@ -1221,7 +1221,10 @@ def test_get_body_ids_missing(self, bkd): bkd.get_body_ids(["nonexistent_body_xyz"]) def test_get_joint_range(self, bkd): - assert bkd.get_joint_range() is None + jr = bkd.get_joint_range() + assert jr is not None + assert jr.shape == (bkd.num_dof_vel, 2) + assert np.all(jr[:, 0] <= jr[:, 1]) # --------------------------------------------------------------------------- diff --git a/tests/base/test_sim_backend_smoke.py b/tests/base/test_sim_backend_smoke.py index ee3b9fb6f..cf6ec6e21 100644 --- a/tests/base/test_sim_backend_smoke.py +++ b/tests/base/test_sim_backend_smoke.py @@ -535,7 +535,9 @@ def test_motrix_model_properties_smoke(): ctrl_range = bkd.get_actuator_ctrl_range() _shape(ctrl_range, bkd.num_actuators, 2) assert bkd.get_default_qpos().ndim == 1 - assert bkd.get_joint_range() is None + joint_range = bkd.get_joint_range() + assert joint_range is not None + assert joint_range.shape == (bkd.num_dof_vel, 2) def test_motrix_copy_body_state_matches_split_queries(): diff --git a/tests/benchmark/test_offpolicy_collector_active_benchmark.py b/tests/benchmark/test_offpolicy_collector_active_benchmark.py index 25665b2c8..0b2054452 100644 --- a/tests/benchmark/test_offpolicy_collector_active_benchmark.py +++ b/tests/benchmark/test_offpolicy_collector_active_benchmark.py @@ -135,15 +135,15 @@ def test_auto_discovery_supports_motrixsim_alias() -> None: def test_noise_seed_override_composes_for_target_g1_profiles() -> None: - # Manager-Based G1 walk owners no longer carry env.noise_config; only the - # legacy motion-tracking owner still composes the seed override. + # Manager-Based owners seed their NumPy runtime directly rather than + # carrying the legacy observation-noise config. for spec in (("sac", "g1_motion_tracking", "mujoco"),): cfg = bench._compose_offpolicy_cfg( *spec, - extra_overrides=["env.noise_config.seed=123"], + extra_overrides=["env.seed=123"], ) - assert cfg.env.noise_config.seed == 123 + assert cfg.env.seed == 123 def test_stats_reports_distribution() -> None: diff --git a/tests/config/test_locomotion_params.py b/tests/config/test_locomotion_params.py index 14e2e4f25..cffa86828 100644 --- a/tests/config/test_locomotion_params.py +++ b/tests/config/test_locomotion_params.py @@ -2,6 +2,7 @@ from __future__ import annotations +import re from pathlib import Path import pytest @@ -39,6 +40,7 @@ 0.07450087032950714, 0.07450087032950714, ] +G1_23DOF_BEYONDMIMIC_ACTION_SCALE = G1_BEYONDMIMIC_ACTION_SCALE[:13] + [0.43857731392336724] * 10 X2_ACTION_SCALE = [0.25] * 29 @@ -411,8 +413,42 @@ def test_ppo_g1_motion_tracking_deploy(): assert cfg.algo.max_iterations == 15000 assert cfg.algo.algorithm.entropy_coef == pytest.approx(0.005) assert cfg.env.sim_dt == pytest.approx(0.005) - assert cfg.env.sensor.gyro == "pelvis_gyro" - assert list(cfg.env.control_config.action_scale) == pytest.approx(G1_BEYONDMIMIC_ACTION_SCALE) + assert cfg.env.observations.actor.terms.base_ang_vel.params.sensor_name == "pelvis_gyro" + assert cfg.env.actions.joint_pos.scale[".*_(hip_pitch|hip_yaw)_joint"] == pytest.approx( + G1_BEYONDMIMIC_ACTION_SCALE[0] + ) + assert cfg.env.actions.joint_pos.scale[".*_wrist_(pitch|yaw)_joint"] == pytest.approx( + G1_BEYONDMIMIC_ACTION_SCALE[20] + ) + + +@pytest.mark.parametrize( + ("task", "expected"), + [ + ("g1_motion_tracking_deploy", G1_BEYONDMIMIC_ACTION_SCALE), + ("g1_23dof_motion_tracking_deploy", G1_23DOF_BEYONDMIMIC_ACTION_SCALE), + ], +) +def test_ppo_g1_motion_tracking_deploy_action_scale_expands_in_joint_order( + task: str, + expected: list[float], +) -> None: + from hydra import compose, initialize_config_dir + from hydra.core.global_hydra import GlobalHydra + + GlobalHydra.instance().clear() + with initialize_config_dir(config_dir=str(CONF_DIR / "ppo"), version_base="1.3"): + cfg = compose("config", overrides=[f"task={task}/mujoco"]) + + scales = cfg.env.actions.joint_pos.scale + resolved: list[float] = [] + for joint_name in cfg.env.scene.entities.robot.joint_names: + matches = [ + float(value) for pattern, value in scales.items() if re.fullmatch(pattern, joint_name) + ] + assert len(matches) == 1, f"{joint_name} matched {len(matches)} action-scale patterns" + resolved.append(matches[0]) + assert resolved == pytest.approx(expected) def test_ppo_g1_box_tracking(): diff --git a/tests/envs/mdp/test_events.py b/tests/envs/mdp/test_events.py index 3b2a3ccae..9d5468c6a 100644 --- a/tests/envs/mdp/test_events.py +++ b/tests/envs/mdp/test_events.py @@ -215,6 +215,9 @@ def get_dof_vel(self) -> np.ndarray: def get_default_dof_pos(self) -> np.ndarray: return np.zeros(3) + def get_joint_range(self) -> np.ndarray: + return np.tile([-1.0, 1.0], (3, 1)) + def get_joint_dof_indices(self, names) -> np.ndarray: table = {"j0": 6, "j1": 7, "j2": 8} return np.asarray([table[name] for name in names], dtype=np.int32) diff --git a/tests/envs/mdp/test_joint_position_action.py b/tests/envs/mdp/test_joint_position_action.py index dbba0a161..c6b82a148 100644 --- a/tests/envs/mdp/test_joint_position_action.py +++ b/tests/envs/mdp/test_joint_position_action.py @@ -60,6 +60,9 @@ def get_dof_vel(self) -> np.ndarray: def get_default_dof_pos(self) -> np.ndarray: return np.asarray([0.1, 0.2, 0.3], dtype=np.float32) + def get_joint_range(self) -> np.ndarray: + return np.tile(np.asarray([[-1.0, 1.0]], dtype=np.float32), (3, 1)) + def _action( **overrides, diff --git a/tests/envs/mdp/test_observations.py b/tests/envs/mdp/test_observations.py index 83092c79f..e5dfd3b2a 100644 --- a/tests/envs/mdp/test_observations.py +++ b/tests/envs/mdp/test_observations.py @@ -73,6 +73,9 @@ def get_dof_vel(self) -> np.ndarray: def get_default_dof_pos(self) -> np.ndarray: return np.asarray([0.1, 0.2, 0.3], dtype=np.float32) + def get_joint_range(self) -> np.ndarray: + return np.tile(np.asarray([[-1.0, 1.0]], dtype=np.float32), (3, 1)) + def get_body_pos_w(self, ids: np.ndarray) -> np.ndarray: return self.body_pos[:, ids] diff --git a/tests/envs/test_env_configs.py b/tests/envs/test_env_configs.py index d4a1ec3b7..f8915967b 100644 --- a/tests/envs/test_env_configs.py +++ b/tests/envs/test_env_configs.py @@ -70,6 +70,33 @@ def _g1_manager_override(task: str = "g1_walk_flat") -> dict[str, Any]: return BackendAdapter(cfg, root_dir=repo_root, algo_name="ppo").build_task_env_cfg_override() +def _motion_manager_override( + task: str, + backend: str, + *, + config_root: str = "ppo", +) -> tuple[str, dict[str, Any]]: + from hydra import compose, initialize_config_dir + from hydra.core.global_hydra import GlobalHydra + + from unilab.training.backend_adapter import BackendAdapter + + repo_root = Path(__file__).parents[2] + GlobalHydra.instance().clear() + with initialize_config_dir( + config_dir=str(repo_root / "conf" / config_root), version_base="1.3" + ): + overrides = [f"task={task}/{backend}"] + if config_root == "offpolicy": + overrides.insert(0, "algo=sac") + cfg = compose("config", overrides=overrides) + return str(cfg.training.task_name), BackendAdapter( + cfg, + root_dir=repo_root, + algo_name="sac" if config_root == "offpolicy" else config_root, + ).build_task_env_cfg_override() + + # --------------------------------------------------------------------------- # Non-slow: config attribute completeness (no env.step(), no MuJoCo sim) # --------------------------------------------------------------------------- @@ -2154,103 +2181,221 @@ def test_allegro_grasp_manager_runtime_uses_zero_increment_action(sim_backend: s env.close() -@pytest.mark.parametrize("sim_backend", ["mujoco", "motrix"]) -def test_g1_motion_tracking_reset_and_step(sim_backend: str): - """G1MotionTracking needs a motion_file — skip if not available.""" - ensure_registries() +_MOTION_CORE_RUNTIME_CASES = ( + pytest.param("ppo", "g1_motion_tracking", "G1MotionTracking", 160, 286, 29, False), + pytest.param( + "ppo", + "g1_motion_tracking_deploy", + "G1MotionTrackingDeploy", + 154, + 286, + 29, + False, + ), + pytest.param("ppo", "g1_23dof_motion_tracking", "G1MotionTracking23Dof", 130, 256, 23, False), + pytest.param( + "ppo", + "g1_23dof_motion_tracking_deploy", + "G1MotionTracking23DofDeploy", + 124, + 256, + 23, + False, + ), + pytest.param("appo", "g1_motion_tracking", "G1MotionTracking", 160, 286, 29, False), + pytest.param("appo", "g1_23dof_motion_tracking", "G1MotionTracking23Dof", 130, 256, 23, False), + pytest.param( + "offpolicy", + "sac/g1_motion_tracking", + "G1MotionTrackingSAC", + 160, + 289, + 29, + True, + ), + pytest.param( + "offpolicy", + "sac/g1_23dof_motion_tracking", + "G1MotionTrackingSAC23Dof", + 130, + 259, + 23, + True, + ), +) + + +def test_g1_motion_core_registrations_are_manager_only() -> None: from unilab.base import registry - if sim_backend == "mujoco": - _require_mujoco_runtime() - else: - pytest.importorskip("motrixsim") + ensure_registries() + metadata = registry.list_registered_envs() + for task_name in ( + "G1MotionTracking", + "G1MotionTrackingDeploy", + "G1MotionTracking23Dof", + "G1MotionTracking23DofDeploy", + "G1MotionTrackingSAC", + "G1MotionTrackingSAC23Dof", + ): + assert metadata[task_name] == { + "config_factory": "ManagerBasedRlEnvCfg", + "available_backends": ["mujoco", "motrix"], + } - # Look for any motion file in the expected location - motion_dir = Path(__file__).parents[2] / "src" / "unilab" / "assets" / "motions" / "g1" - if not motion_dir.exists(): - pytest.skip(f"Motion data directory not found: {motion_dir}") - npz_files = list(motion_dir.glob("*.npz")) - if not npz_files: - pytest.skip(f"No .npz motion files in {motion_dir}") +def test_g1_motion_manager_ppo_wraps_only_active_rows_in_one_state_commit( + monkeypatch: pytest.MonkeyPatch, +) -> None: + ensure_registries() + _require_mujoco_runtime() + from unilab.base import registry - # Filter out 23-DoF motion files — G1MotionTracking uses 29-DoF config - non_23dof = [f for f in npz_files if "_23dof" not in f.name] - if not non_23dof: - pytest.skip("No non-23-DoF motion files available for 29-DoF config") - motion_file = str(non_23dof[0]) - env = cast( - Any, - registry.make( - "G1MotionTracking", - num_envs=2, - sim_backend=sim_backend, - env_cfg_override={"motion_file": motion_file}, - ), + _, override = _motion_manager_override("g1_motion_tracking", "mujoco") + env = registry.make( + "G1MotionTracking", + num_envs=2, + sim_backend="mujoco", + env_cfg_override=override, ) try: - spec = env.obs_groups_spec - assert isinstance(spec, dict) - assert "obs" in spec - assert "critic" in spec - obs_shape = env.observation_space.shape - assert obs_shape is not None - assert sum(spec.values()) == obs_shape[0] - - state = env.init_state() - assert isinstance(state.obs, dict) - for key, dim in spec.items(): - assert state.obs[key].shape == (2, dim) + env.init_state() + command = env.command_manager.get_term("motion") + command.time_steps[:] = command.sampler.current_clip_end_frames + env.reset_buf[:] = [True, False] + + set_state_env_ids: list[np.ndarray] = [] + original_set_state = env._backend.set_state + + def record_set_state( + env_ids: np.ndarray, + qpos: np.ndarray, + qvel: np.ndarray, + *, + randomization: Any = None, + ) -> Any: + set_state_env_ids.append(env_ids.copy()) + return original_set_state( + env_ids, + qpos, + qvel, + randomization=randomization, + ) - action_shape = env.action_space.shape - assert action_shape is not None - actions = np.zeros((2, action_shape[0])) - state = env.step(actions) - assert isinstance(state.obs, dict) - assert state.reward.shape == (2,) - assert state.terminated.shape == (2,) - assert state.truncated.shape == (2,) + monkeypatch.setattr(env._backend, "set_state", record_set_state) + all_ids = np.arange(env.num_envs, dtype=np.int32) + with env._reset_state.scoped(all_ids): + env.command_manager.compute(dt=0.0) + env.command_manager.post_compute() + + assert len(set_state_env_ids) == 1 + np.testing.assert_array_equal(set_state_env_ids[0], [1]) + assert command.time_steps[0] == command.sampler.current_clip_end_frames[0] + assert command.time_steps[1] <= command.sampler.current_clip_end_frames[1] + expected_motion = command.motion.get_motion_at_frame(command.time_steps) + np.testing.assert_array_equal(command.joint_pos, expected_motion.joint_pos) + np.testing.assert_array_equal( + command._robot_body_pos_w, + command.robot.data.body_link_pos_w[:, command._robot_body_ids], + ) + assert command._robot_cache_step == env.common_step_counter finally: env.close() -def test_g1_motion_tracking_deploy_reset_and_step_mujoco(): - """Deploy env keeps motion-tracking behavior but exposes unitree mimic actor inputs.""" +def test_g1_motion_manager_sac_clip_end_is_truncation() -> None: ensure_registries() _require_mujoco_runtime() from unilab.base import registry - motion_dir = Path(__file__).parents[2] / "src" / "unilab" / "assets" / "motions" / "g1" - if not motion_dir.exists(): - pytest.skip(f"Motion data directory not found: {motion_dir}") + _, override = _motion_manager_override( + "sac/g1_motion_tracking", + "mujoco", + config_root="offpolicy", + ) + override["auto_reset"] = False + env = registry.make( + "G1MotionTrackingSAC", + num_envs=2, + sim_backend="mujoco", + env_cfg_override=override, + ) + try: + env.init_state() + command = env.command_manager.get_term("motion") + command.time_steps[:] = command.sampler.current_clip_end_frames + + state = env.step(np.zeros((2, 29), dtype=np.float32)) - npz_files = list(motion_dir.glob("*.npz")) - if not npz_files: - pytest.skip(f"No .npz motion files in {motion_dir}") + np.testing.assert_array_equal(state.terminated, [False, False]) + np.testing.assert_array_equal(state.truncated, [True, True]) + np.testing.assert_array_equal(command.time_steps, command.sampler.current_clip_end_frames) + finally: + env.close() - # Filter out 23-DoF motion files — G1MotionTrackingDeploy uses 29-DoF config - non_23dof = [f for f in npz_files if "_23dof" not in f.name] - if not non_23dof: - pytest.skip("No non-23-DoF motion files available for 29-DoF deploy config") - env = cast( - Any, - registry.make( - "G1MotionTrackingDeploy", - num_envs=2, - sim_backend="mujoco", - env_cfg_override={"motion_file": str(non_23dof[0])}, - ), + +@pytest.mark.parametrize( + ("config_root", "task", "identity", "actor_dim", "critic_dim", "action_dim", "truncate"), + _MOTION_CORE_RUNTIME_CASES, +) +@pytest.mark.parametrize("sim_backend", ["mujoco", "motrix"]) +def test_g1_motion_core_manager_reset_and_step( + config_root: str, + task: str, + identity: str, + actor_dim: int, + critic_dim: int, + action_dim: int, + truncate: bool, + sim_backend: str, +) -> None: + ensure_registries() + from unilab.base import registry + from unilab.envs import ManagerBasedRlEnv + + if sim_backend == "mujoco": + _require_mujoco_runtime() + else: + pytest.importorskip("motrixsim") + + task_name, override = _motion_manager_override( + task, + sim_backend, + config_root=config_root, ) + assert task_name == identity + env = registry.make( + identity, + num_envs=2, + sim_backend=sim_backend, + env_cfg_override=override, + ) + assert isinstance(env, ManagerBasedRlEnv) try: - assert env.obs_groups_spec == {"obs": 154, "critic": 286} + assert env.obs_groups_spec == {"obs": actor_dim, "critic": critic_dim} + assert env.action_space.shape == (action_dim,) + command = env.command_manager.get_term("motion") + assert command.cfg.params.truncate_on_clip_end is truncate + if sim_backend == "motrix" and "deploy" in task: + assert env._cfg.events["foot_friction"] is None + assert env._cfg.events["push_robot"] is None + if sim_backend == "motrix" and config_root in {"ppo", "appo"}: + root_pos = env._cfg.rewards["motion_global_root_pos"] + action_rate = env._cfg.rewards["action_rate_l2"] + assert root_pos is not None + assert action_rate is not None + expected_weights = (1.0, -0.05) if config_root == "ppo" else (0.5, -0.1) + assert (root_pos.weight, action_rate.weight) == pytest.approx(expected_weights) + state = env.init_state() - assert state.obs["obs"].shape == (2, 154) - assert state.obs["critic"].shape == (2, 286) - - action_shape = env.action_space.shape - assert action_shape is not None - state = env.step(np.zeros((2, action_shape[0]))) - assert state.obs["obs"].shape == (2, 154) - assert state.obs["critic"].shape == (2, 286) + assert state.obs["obs"].shape == (2, actor_dim) + assert state.obs["critic"].shape == (2, critic_dim) + + state = env.step(np.zeros((2, action_dim), dtype=np.float32)) + assert state.reward.shape == (2,) + assert state.terminated.shape == (2,) + assert state.truncated.shape == (2,) + assert np.isfinite(state.reward).all() + assert all(np.isfinite(values).all() for values in state.obs.values()) finally: env.close() diff --git a/tests/envs/test_manager_based_rl_env.py b/tests/envs/test_manager_based_rl_env.py index 0cbf39fd7..351d58326 100644 --- a/tests/envs/test_manager_based_rl_env.py +++ b/tests/envs/test_manager_based_rl_env.py @@ -140,6 +140,9 @@ def get_default_dof_pos(self) -> np.ndarray: def get_dof_vel(self) -> np.ndarray: return np.zeros((self.num_envs, 1), dtype=np.float32) + def get_joint_range(self) -> np.ndarray: + return np.array([[-1.0, 1.0]], dtype=np.float32) + def set_state( self, env_ids: np.ndarray, @@ -234,6 +237,22 @@ def _update_command(self, env_ids: np.ndarray | None) -> None: self._env.command_update_ids.append(ids) +@dataclass(kw_only=True) +class _StateWritingCommandCfg(CommandTermCfg): + def build(self, env) -> CommandTerm: + return _StateWritingCommand(self, env) + + +class _StateWritingCommand(_Command): + def _resample_command(self, env_ids: np.ndarray) -> None: + self._command[env_ids, 0] = 0.75 + self._env.scene["robot"].write_joint_state_to_sim( + np.full((len(env_ids), 1), 0.75, dtype=np.float32), + np.full((len(env_ids), 1), -0.75, dtype=np.float32), + env_ids=env_ids, + ) + + class _Recorder(RecorderTerm): def record_pre_reset(self, env_ids: np.ndarray) -> None: self._env.trace.append(("pre_reset", env_ids.tolist())) @@ -997,6 +1016,32 @@ def test_reset_events_compose_then_commit_default_state_once() -> None: assert backend.joint_layout_calls == 2 +def test_reset_event_and_command_state_writes_share_one_commit() -> None: + cfg = _make_cfg(include_optional_managers=False) + cfg.events = { + "default_first": EventTermCfg(func=mdp.reset_scene_to_default, mode="reset"), + "joint_state": EventTermCfg(func=_write_reset_joint_state, mode="reset"), + } + cfg.commands = {"state_writer": _StateWritingCommandCfg(resampling_time_range=(1.0, 1.0))} + cfg.scene.entities["robot"] = EntityCfg( + joint_names=("joint",), + actuator_names=("motor",), + ) + backend = _ResetBackend(2) + env = _TestEnv(cfg, cast(SimBackend, backend), 2) + + env.reset() + + assert len(backend.set_state_calls) == 1 + ids, qpos, qvel = backend.set_state_calls[0] + np.testing.assert_array_equal(ids, [0, 1]) + np.testing.assert_array_equal( + qpos, + [[0.0, 0.0, 0.5, 0.75], [0.0, 0.0, 0.5, 0.75]], + ) + np.testing.assert_array_equal(qvel, [[0.0, 0.0, -0.75], [0.0, 0.0, -0.75]]) + + def test_pure_reset_event_does_not_request_backend_state_capability() -> None: env, backend = _make_env() diff --git a/tests/envs/test_motion_loader.py b/tests/envs/test_motion_loader.py index b7961cc1b..1139b8ba0 100644 --- a/tests/envs/test_motion_loader.py +++ b/tests/envs/test_motion_loader.py @@ -184,6 +184,30 @@ def test_motion_sampler_step_respects_current_clip_end(tmp_path): np.testing.assert_array_equal(sampler.current_frames, np.array([2, 4], dtype=np.int32)) +def test_motion_sampler_uses_env_owned_rng_and_steps_only_selected_rows(tmp_path): + motion = tmp_path / "motion.npz" + _write_motion_npz(motion, base_value=0.0, num_frames=8) + loader = MotionLoader(str(motion)) + env_ids = np.array([0, 2], dtype=np.int32) + sampler = MotionSampler( + loader, + mode="uniform", + num_envs=3, + rng=np.random.default_rng(17), + ) + expected_rng = np.random.default_rng(17) + + frames = sampler.sample_frames(env_ids) + + np.testing.assert_array_equal(frames, expected_rng.integers(0, 8, 2, dtype=np.int32)) + untouched = int(sampler.current_frames[1]) + sampler.current_clip_end_frames[:] = 7 + done = sampler.step(np.array([2], dtype=np.int32)) + assert done.size == 0 + assert sampler.current_frames[1] == untouched + assert sampler.current_frames[2] == frames[1] + 1 + + def test_box_motion_loader_reads_object_state_and_trims_robot_joints(tmp_path): from unilab.tasks.motion_tracking.g1.motion_box_loader import BoxMotionLoader diff --git a/tests/scripts/test_train_scripts.py b/tests/scripts/test_train_scripts.py index ef523fa1b..63baace03 100644 --- a/tests/scripts/test_train_scripts.py +++ b/tests/scripts/test_train_scripts.py @@ -1275,12 +1275,12 @@ def test_g1_motion_tracking_ppo_motrix_prefers_backend_specific_reward( mod = _train_rsl_rl(monkeypatch) cfg = _ppo_cfg(["task=g1_motion_tracking/motrix"]) - assert cfg.reward.scales.motion_body_pos == pytest.approx(1.0) - cfg.reward.scales.motion_body_pos = 1.25 + assert cfg.reward.motion_body_pos.weight == pytest.approx(1.0) + cfg.reward.motion_body_pos.weight = 1.25 env_cfg_override = mod.build_ppo_env_cfg_override(cfg) - assert env_cfg_override["reward_config"]["scales"]["motion_body_pos"] == pytest.approx(1.25) + assert env_cfg_override["rewards"]["motion_body_pos"]["weight"] == pytest.approx(1.25) def test_build_ppo_play_env_cfg_override_applies_g1_motion_tracking_play_profile( @@ -1293,15 +1293,16 @@ def test_build_ppo_play_env_cfg_override_applies_g1_motion_tracking_play_profile monkeypatch.setattr( mod, "materialize_scene_visual_override", - lambda source_model_file, **kwargs: "/tmp/g1_motion_tracking_play_scene.xml", + lambda *_args, **_kwargs: pytest.fail("manager scene must not be replaced"), ) env_cfg_override = mod.build_ppo_play_env_cfg_override(cfg) assert cfg.training.play_env_num == 16 assert env_cfg_override["render_spacing"] == pytest.approx(2.5) - assert env_cfg_override["scene"].model_file == "/tmp/g1_motion_tracking_play_scene.xml" - assert env_cfg_override["reward_config"]["scales"]["motion_body_pos"] == pytest.approx(1.0) + assert env_cfg_override["scene"]["model_file"].endswith("robots/g1/scene_flat.xml") + assert "robot" in env_cfg_override["scene"]["entities"] + assert env_cfg_override["rewards"]["motion_body_pos"]["weight"] == pytest.approx(1.0) def test_build_ppo_play_env_cfg_override_respects_cli_play_env_override( @@ -1328,27 +1329,21 @@ def test_build_ppo_play_env_cfg_override_respects_cli_play_env_override( assert env_cfg_override["render_spacing"] == pytest.approx(2.5) -def test_build_ppo_play_env_cfg_override_resolves_relative_ground_texture( +def test_build_ppo_play_env_cfg_override_keeps_task_owned_manager_scene( monkeypatch: pytest.MonkeyPatch, ): mod = _train_rsl_rl(monkeypatch) cfg = _ppo_cfg(["task=g1_motion_tracking/motrix", "training.play_only=true"]) - cfg.play_profile.scene.ground_texture_file = "src/unilab/assets/robots/g1/textures/floor.png" - - captured = {} - - def _fake_materialize(source_model_file, **kwargs): - captured["source_model_file"] = source_model_file - captured.update(kwargs) - return "/tmp/g1_motion_tracking_play_scene.xml" - - monkeypatch.setattr(mod, "materialize_scene_visual_override", _fake_materialize) + monkeypatch.setattr( + mod, + "materialize_scene_visual_override", + lambda *_args, **_kwargs: pytest.fail("manager scene must not be replaced"), + ) - mod.build_ppo_play_env_cfg_override(cfg) + env_cfg_override = mod.build_ppo_play_env_cfg_override(cfg) - assert captured["ground_texture_file"] == str( - mod.ROOT_DIR / "src/unilab/assets/robots/g1/textures/floor.png" - ) + assert env_cfg_override["scene"]["entities"]["robot"]["root_body_name"] == "pelvis" + assert env_cfg_override["scene"]["entities"]["robot"]["joint_names"] def test_go2_arm_manip_loco_motrix_eval_uses_visual_floor( @@ -1472,28 +1467,28 @@ def step(self, actions): def test_g1_motion_tracking_appo_reward_extraction_prefers_backend_specific_reward(): - from unilab.training.reward import extract_reward_config + from unilab.training import BackendAdapter cfg = _appo_cfg(["task=g1_motion_tracking/motrix"]) - assert cfg.reward.scales.motion_body_pos == pytest.approx(1.0) - cfg.reward.scales.motion_body_pos = 1.5 + assert cfg.reward.motion_body_pos.weight == pytest.approx(1.0) + cfg.reward.motion_body_pos.weight = 1.5 - env_cfg_override = extract_reward_config(cfg) + env_cfg_override = BackendAdapter(cfg, root_dir=_SRC_DIR.parent).build_task_env_cfg_override() - assert env_cfg_override["reward_config"]["scales"]["motion_body_pos"] == pytest.approx(1.5) + assert env_cfg_override["rewards"]["motion_body_pos"]["weight"] == pytest.approx(1.5) def test_g1_motion_tracking_ppo_task_exposes_final_reward(): cfg = _ppo_cfg(["task=g1_motion_tracking/motrix"]) - assert cfg.reward.scales.motion_body_pos == pytest.approx(1.0) + assert cfg.reward.motion_body_pos.weight == pytest.approx(1.0) def test_g1_motion_tracking_appo_task_exposes_final_reward(): cfg = _appo_cfg(["task=g1_motion_tracking/motrix"]) - assert cfg.reward.scales.motion_body_pos == pytest.approx(1.0) + assert cfg.reward.motion_body_pos.weight == pytest.approx(1.0) def test_sharpa_appo_motrix_owner_uses_backend_specific_overrides(): @@ -1519,7 +1514,7 @@ def test_build_appo_runner_kwargs_forwards_sim_backend(): runner_kwargs = mod.build_appo_runner_kwargs( cfg, - env_cfg_override={"reward_config": {"scales": {}}}, + env_cfg_override={"rewards": {}}, collector_device="cpu", ) @@ -1530,7 +1525,7 @@ def test_build_appo_runner_kwargs_forwards_sim_backend(): assert runner_kwargs["steps_per_env"] == cfg.algo.steps_per_env assert "num_workers" not in runner_kwargs assert "num_collectors" not in runner_kwargs - assert runner_kwargs["env_cfg_overrides"]["reward_config"]["scales"] == {} + assert runner_kwargs["env_cfg_overrides"]["rewards"] == {} def test_run_motrix_play_loop_runs_without_physics_state(): diff --git a/tests/tasks/test_migration_matrix.py b/tests/tasks/test_migration_matrix.py index 10b06b87a..77bfb065a 100644 --- a/tests/tasks/test_migration_matrix.py +++ b/tests/tasks/test_migration_matrix.py @@ -18,7 +18,7 @@ def test_registered_tasks_have_explicit_migration_records() -> None: assert PRODUCTION_TASK_NAMES <= registered.keys() assert {record.task_name for record in records} == set(PRODUCTION_TASK_NAMES) assert len(records) == 39 - assert sum(record.status == "Compatible" for record in records) == 15 + assert sum(record.status == "Compatible" for record in records) == 21 assert sum(record.target == "compatibility" for record in records) == 3 @@ -27,7 +27,7 @@ def test_registered_tasks_have_explicit_migration_records() -> None: [ ("Go2ArmManipLoco", "go2_arm", "compatibility", "Adapted"), ("SharpaInhandRotation", "sharpa", "compatibility", "Adapted"), - ("G1MotionTracking", "motion_tracking", "mba", "Adapted"), + ("G1MotionTracking", "motion_tracking", "complete", "Compatible"), ("G1WalkRough", "g1_locomotion", "complete", "Compatible"), ("Go2JoystickRough", "quadruped_rough", "complete", "Compatible"), ], diff --git a/tests/training/test_sim2sim_resolver.py b/tests/training/test_sim2sim_resolver.py index b9530a3d9..a7710ce47 100644 --- a/tests/training/test_sim2sim_resolver.py +++ b/tests/training/test_sim2sim_resolver.py @@ -115,7 +115,7 @@ def test_extract_snapshot_includes_only_present_contract_fields(): assert "training.sim_backend" not in snapshot # ...and absent fields are omitted (never stored as None). assert "algo.obs_normalization" not in snapshot - assert "env.sampling_mode" not in snapshot + assert "env.commands.motion.params.sampling_mode" not in snapshot assert "reward.base_height_target" not in snapshot @@ -272,7 +272,7 @@ def test_env_structural_denylist_is_the_env_subset(): "env.actions", "env.policy_observation_group", "env.critic_observation_group", - "env.sampling_mode", + "env.commands.motion.params.sampling_mode", ] assert set(ENV_STRUCTURAL_DENYLIST) <= set(DENYLIST) @@ -294,7 +294,10 @@ def test_env_field_present_in_target_absent_in_source_raises(tmp_path): # the target sets it explicitly. Still unverifiable -> fail closed. _write_sidecar(tmp_path, {"algo.empirical_normalization": False}) target = OmegaConf.create( - {"algo": {"empirical_normalization": False}, "env": {"sampling_mode": "adaptive"}} + { + "algo": {"empirical_normalization": False}, + "env": {"commands": {"motion": {"params": {"sampling_mode": "adaptive"}}}}, + } ) with pytest.raises(CrossBackendIncompatibleError) as excinfo: resolve_sim2sim_config(tmp_path, target) @@ -303,6 +306,23 @@ def test_env_field_present_in_target_absent_in_source_raises(tmp_path): assert "source=" in msg +def test_legacy_sampling_mode_snapshot_resolves_against_manager_path(tmp_path) -> None: + _write_sidecar(tmp_path, {"env.sampling_mode": "adaptive"}) + target = OmegaConf.create( + {"env": {"commands": {"motion": {"params": {"sampling_mode": "adaptive"}}}}} + ) + + assert resolve_sim2sim_config(tmp_path, target) is target + + +def test_legacy_config_sampling_mode_is_snapshotted_under_canonical_path() -> None: + cfg = OmegaConf.create({"env": {"sampling_mode": "adaptive"}}) + + assert extract_contract_snapshot(cfg) == { + "env.commands.motion.params.sampling_mode": "adaptive" + } + + def test_env_field_symmetric_absence_does_not_raise(tmp_path): # Neither side sets the env structural field -> both use the same env default -> ok. _write_sidecar(tmp_path, {"algo.empirical_normalization": False}) diff --git a/tests/training/test_training_helpers.py b/tests/training/test_training_helpers.py index e2eee8462..8ff811352 100644 --- a/tests/training/test_training_helpers.py +++ b/tests/training/test_training_helpers.py @@ -388,7 +388,7 @@ def test_backend_adapter_env_cfg_override_for_motrix_sac_g1_walk_flat(): assert cfg.algo.max_iterations == 5000 -def test_backend_adapter_builds_play_scene_override(): +def test_backend_adapter_keeps_motion_manager_scene_during_play(): cfg = _ppo_cfg(["task=g1_motion_tracking/motrix", "training.play_only=true"]) assert cfg.training.play_env_num == 16 captured: dict[str, object] = {} @@ -396,7 +396,7 @@ def test_backend_adapter_builds_play_scene_override(): def _fake_materializer(source_model_file: str, **kwargs) -> str: captured["source_model_file"] = source_model_file captured.update(kwargs) - return "/tmp/g1_motion_tracking_play_scene.xml" + pytest.fail("manager scene must not be replaced during play") env_cfg_override = BackendAdapter( cfg, @@ -407,10 +407,9 @@ def _fake_materializer(source_model_file: str, **kwargs) -> str: assert cfg.training.play_env_num == 16 assert env_cfg_override["render_spacing"] == pytest.approx(2.5) - assert env_cfg_override["scene"].model_file == "/tmp/g1_motion_tracking_play_scene.xml" - assert captured["ground_texture_file"] == str( - _ROOT_DIR / "src/unilab/assets/robots/g1/textures/floor.png" - ) + assert env_cfg_override["scene"]["model_file"].endswith("robots/g1/scene_flat.xml") + assert "robot" in env_cfg_override["scene"]["entities"] + assert captured == {} def test_render_play_mode_uses_env_interactive_contract(): From 8d58deb26539563d1b39d64bd4717bcce111e3fb Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Thu, 20 Aug 2026 19:37:21 +0800 Subject: [PATCH 095/169] feat: migrate remaining motion profiles to manager api (#1236) --- .../task/g1_23dof_climb_tracking/motrix.yaml | 65 +- .../task/g1_23dof_climb_tracking/mujoco.yaml | 126 +- .../task/g1_23dof_flip_tracking/motrix.yaml | 82 +- .../task/g1_23dof_flip_tracking/mujoco.yaml | 135 +- .../g1_23dof_wall_flip_tracking/motrix.yaml | 57 +- .../g1_23dof_wall_flip_tracking/mujoco.yaml | 84 +- conf/appo/task/g1_climb_tracking/motrix.yaml | 71 +- conf/appo/task/g1_climb_tracking/mujoco.yaml | 134 +- conf/appo/task/g1_flip_tracking/motrix.yaml | 92 +- conf/appo/task/g1_flip_tracking/mujoco.yaml | 141 +- .../task/g1_wall_flip_tracking/motrix.yaml | 58 +- .../task/g1_wall_flip_tracking/mujoco.yaml | 93 +- .../sac/g1_23dof_flip_tracking/mujoco.yaml | 138 +- .../g1_23dof_wall_flip_tracking/mujoco.yaml | 98 +- .../task/sac/g1_23dof_wbt_obs/mujoco.yaml | 224 ++- .../task/sac/g1_flip_tracking/mujoco.yaml | 146 +- .../sac/g1_wall_flip_tracking/mujoco.yaml | 104 +- .../offpolicy/task/sac/g1_wbt_obs/mujoco.yaml | 223 ++- .../task/g1_23dof_box_tracking/motrix.yaml | 59 +- .../task/g1_23dof_box_tracking/mujoco.yaml | 98 +- .../task/g1_23dof_climb_tracking/motrix.yaml | 74 +- .../task/g1_23dof_climb_tracking/mujoco.yaml | 133 +- .../task/g1_23dof_flip_tracking/motrix.yaml | 49 +- .../task/g1_23dof_flip_tracking/mujoco.yaml | 144 +- .../g1_23dof_wall_flip_tracking/motrix.yaml | 74 +- .../g1_23dof_wall_flip_tracking/mujoco.yaml | 88 +- conf/ppo/task/g1_box_tracking/motrix.yaml | 59 +- conf/ppo/task/g1_box_tracking/mujoco.yaml | 95 +- conf/ppo/task/g1_climb_tracking/motrix.yaml | 84 +- conf/ppo/task/g1_climb_tracking/mujoco.yaml | 138 +- conf/ppo/task/g1_flip_tracking/motrix.yaml | 53 +- conf/ppo/task/g1_flip_tracking/mujoco.yaml | 149 +- .../task/g1_wall_flip_tracking/motrix.yaml | 80 +- .../task/g1_wall_flip_tracking/mujoco.yaml | 94 +- .../task/x2_wall_flip_tracking/motrix.yaml | 89 +- .../task/x2_wall_flip_tracking/mujoco.yaml | 254 +++- .../2-user_guide/4-tasks/2-motion_tracking.md | 74 +- .../5-domain_randomization/0-index.md | 10 +- .../2-writing_providers.md | 6 +- .../3-deployment/1-sim_to_real/1-overview.md | 8 +- .../1-sim_to_real/8-latency_budget.md | 15 +- .../2-sim_to_sim/7-config_guard.md | 2 +- .../2-contracts/4-dr_contract.md | 4 +- .../en/4-developer_guide/7-motion_assets.md | 9 +- .../2-user_guide/4-tasks/2-motion_tracking.md | 72 +- .../5-domain_randomization/0-index.md | 10 +- .../2-writing_providers.md | 6 +- .../3-deployment/1-sim_to_real/1-overview.md | 6 +- .../1-sim_to_real/8-latency_budget.md | 13 +- .../2-contracts/4-dr_contract.md | 4 +- .../4-developer_guide/7-motion_assets.md | 7 +- .../benchmark/torch_env/motion_tracking.py | 27 +- scripts/benchmark/torch_env/run_benchmark.py | 8 +- scripts/deploy/export_deploy_config.py | 22 +- src/unilab/envs/manager_based_rl_env.py | 14 +- src/unilab/tasks/migration_matrix.py | 8 +- src/unilab/tasks/motion_tracking/__init__.py | 100 +- .../tasks/motion_tracking/common/__init__.py | 8 +- .../tasks/motion_tracking/common/config.py | 168 --- .../common/domain_randomization.py | 186 --- .../tasks/motion_tracking/common/g1_base.py | 70 - .../motion_tracking/common/manager_terms.py | 82 +- .../motion_tracking/common/observations.py | 219 +-- .../tasks/motion_tracking/common/reset.py | 83 -- .../tasks/motion_tracking/common/rewards.py | 386 ------ .../motion_tracking/common/terminations.py | 69 - .../tasks/motion_tracking/common/tracking.py | 539 -------- .../tasks/motion_tracking/g1/__init__.py | 127 +- .../tasks/motion_tracking/g1/box_tracking.py | 412 ------ .../tasks/motion_tracking/g1/flip_tracking.py | 264 ---- .../motion_tracking/g1/flip_tracking_sac.py | 85 -- .../tasks/motion_tracking/g1/manager_terms.py | 327 +++++ .../motion_tracking/g1/motion_box_loader.py | 20 + .../tasks/motion_tracking/g1/tracking.py | 152 -- .../tasks/motion_tracking/g1/tracking_obs.py | 517 ------- .../tasks/motion_tracking/g1/tracking_sac.py | 88 -- .../tasks/motion_tracking/x2/__init__.py | 38 +- .../tasks/motion_tracking/x2/flip_tracking.py | 135 -- src/unilab/training/backend_adapter.py | 35 +- .../test_mujoco_scene_context_visual.py | 18 +- tests/config/test_locomotion_params.py | 89 +- tests/envs/test_env_configs.py | 1225 +---------------- tests/envs/test_manager_based_rl_env.py | 12 +- tests/envs/test_motion_profiles.py | 530 +++++++ tests/envs/test_motion_tracking_rewards.py | 155 --- tests/scripts/test_obs_alignment_g1_wbt.py | 21 +- tests/scripts/test_train_script_configs.py | 2 +- tests/tasks/test_migration_matrix.py | 4 +- tests/training/test_training_helpers.py | 26 + 89 files changed, 3083 insertions(+), 7319 deletions(-) delete mode 100644 src/unilab/tasks/motion_tracking/common/config.py delete mode 100644 src/unilab/tasks/motion_tracking/common/domain_randomization.py delete mode 100644 src/unilab/tasks/motion_tracking/common/g1_base.py delete mode 100644 src/unilab/tasks/motion_tracking/common/reset.py delete mode 100644 src/unilab/tasks/motion_tracking/common/rewards.py delete mode 100644 src/unilab/tasks/motion_tracking/common/terminations.py delete mode 100644 src/unilab/tasks/motion_tracking/common/tracking.py delete mode 100644 src/unilab/tasks/motion_tracking/g1/box_tracking.py delete mode 100644 src/unilab/tasks/motion_tracking/g1/flip_tracking.py delete mode 100644 src/unilab/tasks/motion_tracking/g1/flip_tracking_sac.py create mode 100644 src/unilab/tasks/motion_tracking/g1/manager_terms.py delete mode 100644 src/unilab/tasks/motion_tracking/g1/tracking.py delete mode 100644 src/unilab/tasks/motion_tracking/g1/tracking_obs.py delete mode 100644 src/unilab/tasks/motion_tracking/g1/tracking_sac.py delete mode 100644 src/unilab/tasks/motion_tracking/x2/flip_tracking.py create mode 100644 tests/envs/test_motion_profiles.py delete mode 100644 tests/envs/test_motion_tracking_rewards.py diff --git a/conf/appo/task/g1_23dof_climb_tracking/motrix.yaml b/conf/appo/task/g1_23dof_climb_tracking/motrix.yaml index 7844addb7..ae708bf1d 100644 --- a/conf/appo/task/g1_23dof_climb_tracking/motrix.yaml +++ b/conf/appo/task/g1_23dof_climb_tracking/motrix.yaml @@ -1,65 +1,8 @@ # @package _global_ +defaults: + - /task/g1_23dof_climb_tracking/mujoco + - _self_ + training: task_name: G1ClimbTracking23Dof sim_backend: motrix - play_steps: 1000 -algo: - num_envs: 1024 - max_iterations: 20000 - save_interval: 500 -env: - sampling_mode: adaptive - truncate_on_clip_end: false - sim_dt: 0.005 - control_config: - action_scale: - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - anchor_pos_z_threshold: 0.3 - ee_body_pos_z_threshold: 0.3 - terminate_on_undesired_contacts: true - noise_config: - level: 0.0 -reward: - scales: - motion_global_root_pos: 0.5 - motion_global_root_ori: 0.5 - motion_body_pos: 2.0 - motion_body_ori: 1.5 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_ee_body_pos_z: 2.0 - motion_joint_pos: 0.5 - motion_joint_vel: 0.25 - action_rate_l2: -0.005 - joint_limit: -10.0 - undesired_contacts: -0.1 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 diff --git a/conf/appo/task/g1_23dof_climb_tracking/mujoco.yaml b/conf/appo/task/g1_23dof_climb_tracking/mujoco.yaml index 21380ecb4..3586affd7 100644 --- a/conf/appo/task/g1_23dof_climb_tracking/mujoco.yaml +++ b/conf/appo/task/g1_23dof_climb_tracking/mujoco.yaml @@ -1,65 +1,85 @@ # @package _global_ +defaults: + - /task/g1_23dof_motion_tracking/mujoco + - _self_ + training: task_name: G1ClimbTracking23Dof sim_backend: mujoco play_steps: 1000 + algo: num_envs: 1024 max_iterations: 20000 save_interval: 500 + algorithm: + adaptive_kl_factor: 1.2 + adaptive_lr_factor: 1.1 + env: - sampling_mode: adaptive - truncate_on_clip_end: false + scene: + model_file: src/unilab/assets/robots/g1/scene_climb_20_z_scale_1_23dof.xml sim_dt: 0.005 - control_config: - action_scale: - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - anchor_pos_z_threshold: 0.3 - ee_body_pos_z_threshold: 0.3 - terminate_on_undesired_contacts: true - noise_config: - level: 0.0 + max_episode_seconds: 15.0 + actions: + joint_pos: + scale: + ".*_(hip_pitch|hip_yaw)_joint": 0.5475464629911068 + ".*_(hip_roll|knee)_joint": 0.35066146637882434 + ".*_ankle_(pitch|roll)_joint": 0.43857731392336724 + "waist_yaw_joint": 0.5475464629911068 + ".*_(shoulder_(pitch|roll|yaw)|elbow|wrist_roll)_joint": 0.43857731392336724 + commands: + motion: + params: + motion_file: motions/g1/climb_20_z_scale_1.0_23dof.npz + sampling_mode: adaptive + truncate_on_clip_end: false + terminations: + anchor_pos: + params: {command_name: motion, threshold: 0.3} + ee_body_pos: + params: + command_name: motion + threshold: 0.3 + body_names: &ee_bodies + - left_ankle_roll_link + - right_ankle_roll_link + - left_wrist_roll_rubber_hand + - right_wrist_roll_rubber_hand + undesired_contacts: + func: unilab.tasks.motion_tracking.common.manager_terms.bad_undesired_body_contacts + params: + command_name: motion + threshold: 0.05 + body_names: &undesired_bodies + - pelvis + - left_hip_roll_link + - left_knee_link + - right_hip_roll_link + - right_knee_link + - torso_link + - left_shoulder_roll_link + - left_elbow_link + - right_shoulder_roll_link + - right_elbow_link + reward: - scales: - motion_global_root_pos: 0.5 - motion_global_root_ori: 0.5 - motion_body_pos: 2.0 - motion_body_ori: 1.5 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_ee_body_pos_z: 2.0 - motion_joint_pos: 0.5 - motion_joint_vel: 0.25 - action_rate_l2: -0.005 - joint_limit: -10.0 - undesired_contacts: -0.1 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 + motion_body_pos: + weight: 2.0 + motion_body_ori: + weight: 1.5 + motion_ee_body_pos_z: + func: unilab.tasks.motion_tracking.common.manager_terms.motion_relative_body_position_z_error_exp + weight: 2.0 + params: {command_name: motion, std: 0.3, body_names: *ee_bodies} + motion_joint_pos: + weight: 0.5 + motion_joint_vel: + weight: 0.25 + action_rate_l2: + weight: -0.005 + undesired_contacts: + func: unilab.tasks.motion_tracking.common.manager_terms.undesired_body_contacts + weight: -0.1 + params: {command_name: motion, threshold: 0.05, body_names: *undesired_bodies} diff --git a/conf/appo/task/g1_23dof_flip_tracking/motrix.yaml b/conf/appo/task/g1_23dof_flip_tracking/motrix.yaml index 05098dac7..a600df907 100644 --- a/conf/appo/task/g1_23dof_flip_tracking/motrix.yaml +++ b/conf/appo/task/g1_23dof_flip_tracking/motrix.yaml @@ -1,86 +1,14 @@ # @package _global_ +defaults: + - /task/g1_23dof_flip_tracking/mujoco + - _self_ + training: task_name: G1FlipTracking23Dof sim_backend: motrix - play_steps: 1000 + algo: - num_envs: 1024 - steps_per_env: 24 - max_iterations: 3500 - save_interval: 500 algorithm: - num_learning_epochs: 10 - num_mini_batches: 8 - clip_param: 0.2 - gamma: 0.99 - lam: 0.95 - value_loss_coef: 1.0 - entropy_coef: 0.005 - learning_rate: 1.0e-3 - max_grad_norm: 1.0 - use_clipped_value_loss: true - schedule: adaptive - desired_kl: 0.01 adaptive_kl_factor: 2.0 adaptive_lr_factor: 1.5 - optimizer: adam - tau: 1.0 - target_update_freq: 1 - vtrace_clip_rho: 1.0 - vtrace_clip_c: 1.0 enable_compile: true -env: - sampling_mode: start - truncate_on_clip_end: true - sim_dt: 0.005 - control_config: - action_scale: - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - anchor_pos_z_threshold: 0.5 - ee_body_pos_z_threshold: 0.5 - terminate_on_undesired_contacts: true - noise_config: - level: 0.0 -reward: - scales: - motion_global_root_pos: 0.5 - motion_global_root_ori: 0.5 - motion_body_pos: 2.0 - motion_body_ori: 1.5 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_ee_body_pos_z: 2.0 - motion_joint_pos: 0.0 - motion_joint_vel: 0.0 - action_rate_l2: -0.005 - joint_limit: -10.0 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 diff --git a/conf/appo/task/g1_23dof_flip_tracking/mujoco.yaml b/conf/appo/task/g1_23dof_flip_tracking/mujoco.yaml index 6676d9360..6798916ca 100644 --- a/conf/appo/task/g1_23dof_flip_tracking/mujoco.yaml +++ b/conf/appo/task/g1_23dof_flip_tracking/mujoco.yaml @@ -1,8 +1,13 @@ # @package _global_ +defaults: + - /task/g1_23dof_motion_tracking/mujoco + - _self_ + training: task_name: G1FlipTracking23Dof sim_backend: mujoco play_steps: 1000 + algo: num_envs: 1024 steps_per_env: 24 @@ -13,58 +18,84 @@ algo: num_learning_epochs: 10 num_mini_batches: 8 desired_kl: 0.01 + adaptive_kl_factor: 1.2 + adaptive_lr_factor: 1.1 + env: - sampling_mode: start - truncate_on_clip_end: true + scene: + model_file: src/unilab/assets/robots/g1/scene_flat_23dof.xml sim_dt: 0.005 - control_config: - action_scale: - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - anchor_pos_z_threshold: 0.5 - ee_body_pos_z_threshold: 0.5 - terminate_on_undesired_contacts: true - noise_config: - level: 0.0 + actions: + joint_pos: + scale: + ".*_(hip_pitch|hip_yaw)_joint": 0.5475464629911068 + ".*_(hip_roll|knee)_joint": 0.35066146637882434 + ".*_ankle_(pitch|roll)_joint": 0.43857731392336724 + "waist_yaw_joint": 0.5475464629911068 + ".*_(shoulder_(pitch|roll|yaw)|elbow|wrist_roll)_joint": 0.43857731392336724 + commands: + motion: + params: + motion_file: motions/g1/flip_360_001__A304_23dof.npz + sampling_mode: start + truncate_on_clip_end: true + pose_range: &zero_pose + x: [0.0, 0.0] + y: [0.0, 0.0] + z: [0.0, 0.0] + roll: [0.0, 0.0] + pitch: [0.0, 0.0] + yaw: [0.0, 0.0] + velocity_range: *zero_pose + joint_position_range: [0.0, 0.0] + terminations: + motion_clip_end: + func: unilab.tasks.motion_tracking.common.manager_terms.motion_clip_end + time_out: true + params: {command_name: motion} + anchor_pos: + params: {command_name: motion, threshold: 0.5} + anchor_ori: + params: + command_name: motion + threshold: 1.0e9 + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + ee_body_pos: + params: + command_name: motion + threshold: 0.5 + body_names: &ee_bodies + - left_ankle_roll_link + - right_ankle_roll_link + - left_wrist_roll_rubber_hand + - right_wrist_roll_rubber_hand + undesired_contacts: + func: unilab.tasks.motion_tracking.common.manager_terms.bad_undesired_body_contacts + params: + command_name: motion + threshold: 0.05 + body_names: + - pelvis + - left_hip_roll_link + - left_knee_link + - right_hip_roll_link + - right_knee_link + - torso_link + - left_shoulder_roll_link + - left_elbow_link + - right_shoulder_roll_link + - right_elbow_link + reward: - scales: - motion_global_root_pos: 0.5 - motion_global_root_ori: 0.5 - motion_body_pos: 2.0 - motion_body_ori: 1.5 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_ee_body_pos_z: 2.0 - motion_joint_pos: 0.0 - motion_joint_vel: 0.0 - action_rate_l2: -0.005 - joint_limit: -10.0 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 + motion_body_pos: + weight: 2.0 + motion_body_ori: + weight: 1.5 + motion_ee_body_pos_z: + func: unilab.tasks.motion_tracking.common.manager_terms.motion_relative_body_position_z_error_exp + weight: 2.0 + params: {command_name: motion, std: 0.3, body_names: *ee_bodies} + action_rate_l2: + weight: -0.005 diff --git a/conf/appo/task/g1_23dof_wall_flip_tracking/motrix.yaml b/conf/appo/task/g1_23dof_wall_flip_tracking/motrix.yaml index 3b543c9c4..3a43778c1 100644 --- a/conf/appo/task/g1_23dof_wall_flip_tracking/motrix.yaml +++ b/conf/appo/task/g1_23dof_wall_flip_tracking/motrix.yaml @@ -1,50 +1,33 @@ # @package _global_ +defaults: + - /task/g1_23dof_wall_flip_tracking/mujoco + - _self_ + training: task_name: G1WallFlipTracking23Dof sim_backend: motrix - play_steps: 1000 + replay_queue_size: null + algo: - num_envs: 1024 + steps_per_env: 24 max_iterations: 5000 - save_interval: 500 algorithm: num_learning_epochs: 5 num_mini_batches: 4 - clip_param: 0.2 - gamma: 0.99 - lam: 0.95 - value_loss_coef: 1.0 entropy_coef: 0.01 - learning_rate: 1.0e-3 - max_grad_norm: 1.0 - use_clipped_value_loss: true - schedule: adaptive desired_kl: 0.01 - adaptive_kl_factor: 2.0 - adaptive_lr_factor: 1.5 - optimizer: adam - tau: 1.0 - target_update_freq: 1 - vtrace_clip_rho: 1.0 - vtrace_clip_c: 1.0 enable_compile: true + +env: + actions: + joint_pos: + scale: 0.25 + reward: - scales: - motion_global_root_pos: 0.5 - motion_global_root_ori: 0.5 - motion_body_pos: 1.0 - motion_body_ori: 1.0 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_joint_pos: 0.0 - motion_joint_vel: 0.0 - action_rate_l2: -0.1 - joint_limit: -10.0 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 + motion_body_pos: + weight: 1.0 + motion_body_ori: + weight: 1.0 + motion_ee_body_pos_z: null + action_rate_l2: + weight: -0.1 diff --git a/conf/appo/task/g1_23dof_wall_flip_tracking/mujoco.yaml b/conf/appo/task/g1_23dof_wall_flip_tracking/mujoco.yaml index 0decec32d..a140546c1 100644 --- a/conf/appo/task/g1_23dof_wall_flip_tracking/mujoco.yaml +++ b/conf/appo/task/g1_23dof_wall_flip_tracking/mujoco.yaml @@ -1,87 +1,27 @@ # @package _global_ +defaults: + - /task/g1_23dof_flip_tracking/mujoco + - _self_ + training: task_name: G1WallFlipTracking23Dof sim_backend: mujoco - play_steps: 1000 replay_queue_size: 5 + algo: - num_envs: 1024 steps_per_env: 20 max_iterations: 7000 - save_interval: 500 algorithm: num_learning_epochs: 6 - num_mini_batches: 8 - clip_param: 0.2 - gamma: 0.99 - lam: 0.95 - value_loss_coef: 1.0 - entropy_coef: 0.005 - learning_rate: 1.0e-3 - max_grad_norm: 1.0 - use_clipped_value_loss: true - schedule: adaptive desired_kl: 0.008 adaptive_kl_factor: 2.0 adaptive_lr_factor: 1.5 - optimizer: adam - tau: 1.0 - target_update_freq: 1 - vtrace_clip_rho: 1.0 - vtrace_clip_c: 1.0 enable_compile: true + env: - sampling_mode: start - truncate_on_clip_end: true - sim_dt: 0.005 - control_config: - action_scale: - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - anchor_pos_z_threshold: 0.5 - ee_body_pos_z_threshold: 0.5 - terminate_on_undesired_contacts: true - noise_config: - level: 0.0 -reward: - scales: - motion_global_root_pos: 0.5 - motion_global_root_ori: 0.5 - motion_body_pos: 2.0 - motion_body_ori: 1.5 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_ee_body_pos_z: 2.0 - motion_joint_pos: 0.0 - motion_joint_vel: 0.0 - action_rate_l2: -0.005 - joint_limit: -10.0 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 + scene: + model_file: src/unilab/assets/robots/g1/scene_flat_23dof_with_wall.xml + commands: + motion: + params: + motion_file: motions/g1/flip_from_wall_104__A304_23dof.npz diff --git a/conf/appo/task/g1_climb_tracking/motrix.yaml b/conf/appo/task/g1_climb_tracking/motrix.yaml index ec5684fe2..01a98604a 100644 --- a/conf/appo/task/g1_climb_tracking/motrix.yaml +++ b/conf/appo/task/g1_climb_tracking/motrix.yaml @@ -1,71 +1,8 @@ # @package _global_ +defaults: + - /task/g1_climb_tracking/mujoco + - _self_ + training: task_name: G1ClimbTracking sim_backend: motrix - play_steps: 1000 -algo: - num_envs: 1024 - max_iterations: 20000 - save_interval: 500 -env: - sampling_mode: adaptive - truncate_on_clip_end: false - sim_dt: 0.005 - control_config: - action_scale: - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.07450087032950714 - - 0.07450087032950714 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.07450087032950714 - - 0.07450087032950714 - anchor_pos_z_threshold: 0.3 - ee_body_pos_z_threshold: 0.3 - terminate_on_undesired_contacts: true - noise_config: - level: 0.0 -reward: - scales: - motion_global_root_pos: 0.5 - motion_global_root_ori: 0.5 - motion_body_pos: 2.0 - motion_body_ori: 1.5 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_ee_body_pos_z: 2.0 - motion_joint_pos: 0.5 - motion_joint_vel: 0.25 - action_rate_l2: -0.005 - joint_limit: -10.0 - undesired_contacts: -0.1 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 diff --git a/conf/appo/task/g1_climb_tracking/mujoco.yaml b/conf/appo/task/g1_climb_tracking/mujoco.yaml index bff7b13c8..8d15fff57 100644 --- a/conf/appo/task/g1_climb_tracking/mujoco.yaml +++ b/conf/appo/task/g1_climb_tracking/mujoco.yaml @@ -1,71 +1,87 @@ # @package _global_ +defaults: + - /task/g1_motion_tracking/mujoco + - _self_ + training: task_name: G1ClimbTracking sim_backend: mujoco play_steps: 1000 + algo: num_envs: 1024 max_iterations: 20000 save_interval: 500 + algorithm: + adaptive_kl_factor: 1.2 + adaptive_lr_factor: 1.1 + env: - sampling_mode: adaptive - truncate_on_clip_end: false + scene: + model_file: src/unilab/assets/robots/g1/scene_climb_20_z_scale_1.xml sim_dt: 0.005 - control_config: - action_scale: - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.07450087032950714 - - 0.07450087032950714 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.07450087032950714 - - 0.07450087032950714 - anchor_pos_z_threshold: 0.3 - ee_body_pos_z_threshold: 0.3 - terminate_on_undesired_contacts: true - noise_config: - level: 0.0 + max_episode_seconds: 15.0 + actions: + joint_pos: + scale: + ".*_(hip_pitch|hip_yaw)_joint": 0.5475464629911068 + ".*_(hip_roll|knee)_joint": 0.35066146637882434 + ".*_ankle_(pitch|roll)_joint": 0.43857731392336724 + "waist_yaw_joint": 0.5475464629911068 + "waist_(roll|pitch)_joint": 0.43857731392336724 + ".*_(shoulder_(pitch|roll|yaw)|elbow|wrist_roll)_joint": 0.43857731392336724 + ".*_wrist_(pitch|yaw)_joint": 0.07450087032950714 + commands: + motion: + params: + motion_file: motions/g1/climb_20_z_scale_1.0.npz + sampling_mode: adaptive + truncate_on_clip_end: false + terminations: + anchor_pos: + params: {command_name: motion, threshold: 0.3} + ee_body_pos: + params: + command_name: motion + threshold: 0.3 + body_names: &ee_bodies + - left_ankle_roll_link + - right_ankle_roll_link + - left_wrist_yaw_link + - right_wrist_yaw_link + undesired_contacts: + func: unilab.tasks.motion_tracking.common.manager_terms.bad_undesired_body_contacts + params: + command_name: motion + threshold: 0.05 + body_names: &undesired_bodies + - pelvis + - left_hip_roll_link + - left_knee_link + - right_hip_roll_link + - right_knee_link + - torso_link + - left_shoulder_roll_link + - left_elbow_link + - right_shoulder_roll_link + - right_elbow_link + reward: - scales: - motion_global_root_pos: 0.5 - motion_global_root_ori: 0.5 - motion_body_pos: 2.0 - motion_body_ori: 1.5 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_ee_body_pos_z: 2.0 - motion_joint_pos: 0.5 - motion_joint_vel: 0.25 - action_rate_l2: -0.005 - joint_limit: -10.0 - undesired_contacts: -0.1 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 + motion_body_pos: + weight: 2.0 + motion_body_ori: + weight: 1.5 + motion_ee_body_pos_z: + func: unilab.tasks.motion_tracking.common.manager_terms.motion_relative_body_position_z_error_exp + weight: 2.0 + params: {command_name: motion, std: 0.3, body_names: *ee_bodies} + motion_joint_pos: + weight: 0.5 + motion_joint_vel: + weight: 0.25 + action_rate_l2: + weight: -0.005 + undesired_contacts: + func: unilab.tasks.motion_tracking.common.manager_terms.undesired_body_contacts + weight: -0.1 + params: {command_name: motion, threshold: 0.05, body_names: *undesired_bodies} diff --git a/conf/appo/task/g1_flip_tracking/motrix.yaml b/conf/appo/task/g1_flip_tracking/motrix.yaml index 9b2fd0c57..d5f5e8b76 100644 --- a/conf/appo/task/g1_flip_tracking/motrix.yaml +++ b/conf/appo/task/g1_flip_tracking/motrix.yaml @@ -1,92 +1,8 @@ # @package _global_ +defaults: + - /task/g1_flip_tracking/mujoco + - _self_ + training: task_name: G1FlipTracking sim_backend: motrix - play_steps: 1000 -algo: - num_envs: 1024 - steps_per_env: 24 - max_iterations: 3500 - save_interval: 500 - algorithm: - num_learning_epochs: 10 - num_mini_batches: 8 - clip_param: 0.2 - gamma: 0.99 - lam: 0.95 - value_loss_coef: 1.0 - entropy_coef: 0.005 - learning_rate: 1.0e-3 - max_grad_norm: 1.0 - use_clipped_value_loss: true - schedule: adaptive - desired_kl: 0.01 - adaptive_kl_factor: 2.0 - adaptive_lr_factor: 1.5 - optimizer: adam - tau: 1.0 - target_update_freq: 1 - vtrace_clip_rho: 1.0 - vtrace_clip_c: 1.0 - enable_compile: true -env: - sampling_mode: start - truncate_on_clip_end: true - sim_dt: 0.005 - control_config: - action_scale: - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.07450087032950714 - - 0.07450087032950714 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.07450087032950714 - - 0.07450087032950714 - anchor_pos_z_threshold: 0.5 - ee_body_pos_z_threshold: 0.5 - terminate_on_undesired_contacts: true - noise_config: - level: 0.0 -reward: - scales: - motion_global_root_pos: 0.5 - motion_global_root_ori: 0.5 - motion_body_pos: 2.0 - motion_body_ori: 1.5 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_ee_body_pos_z: 2.0 - motion_joint_pos: 0.0 - motion_joint_vel: 0.0 - action_rate_l2: -0.005 - joint_limit: -10.0 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 diff --git a/conf/appo/task/g1_flip_tracking/mujoco.yaml b/conf/appo/task/g1_flip_tracking/mujoco.yaml index 628bdbb12..b8f05ad61 100644 --- a/conf/appo/task/g1_flip_tracking/mujoco.yaml +++ b/conf/appo/task/g1_flip_tracking/mujoco.yaml @@ -1,8 +1,13 @@ # @package _global_ +defaults: + - /task/g1_motion_tracking/mujoco + - _self_ + training: task_name: G1FlipTracking sim_backend: mujoco play_steps: 1000 + algo: num_envs: 1024 steps_per_env: 24 @@ -29,64 +34,84 @@ algo: vtrace_clip_rho: 1.0 vtrace_clip_c: 1.0 enable_compile: true + env: - sampling_mode: start - truncate_on_clip_end: true + scene: + model_file: src/unilab/assets/robots/g1/scene_flat.xml sim_dt: 0.005 - control_config: - action_scale: - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.07450087032950714 - - 0.07450087032950714 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.07450087032950714 - - 0.07450087032950714 - anchor_pos_z_threshold: 0.5 - ee_body_pos_z_threshold: 0.5 - terminate_on_undesired_contacts: true - noise_config: - level: 0.0 + actions: + joint_pos: + scale: + ".*_(hip_pitch|hip_yaw)_joint": 0.5475464629911068 + ".*_(hip_roll|knee)_joint": 0.35066146637882434 + ".*_ankle_(pitch|roll)_joint": 0.43857731392336724 + "waist_yaw_joint": 0.5475464629911068 + "waist_(roll|pitch)_joint": 0.43857731392336724 + ".*_(shoulder_(pitch|roll|yaw)|elbow|wrist_roll)_joint": 0.43857731392336724 + ".*_wrist_(pitch|yaw)_joint": 0.07450087032950714 + commands: + motion: + params: + motion_file: motions/g1/flip_360_001__A304.npz + sampling_mode: start + truncate_on_clip_end: true + pose_range: &zero_pose + x: [0.0, 0.0] + y: [0.0, 0.0] + z: [0.0, 0.0] + roll: [0.0, 0.0] + pitch: [0.0, 0.0] + yaw: [0.0, 0.0] + velocity_range: *zero_pose + joint_position_range: [0.0, 0.0] + terminations: + motion_clip_end: + func: unilab.tasks.motion_tracking.common.manager_terms.motion_clip_end + time_out: true + params: {command_name: motion} + anchor_pos: + params: {command_name: motion, threshold: 0.5} + anchor_ori: + params: + command_name: motion + threshold: 1.0e9 + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + ee_body_pos: + params: + command_name: motion + threshold: 0.5 + body_names: &ee_bodies + - left_ankle_roll_link + - right_ankle_roll_link + - left_wrist_yaw_link + - right_wrist_yaw_link + undesired_contacts: + func: unilab.tasks.motion_tracking.common.manager_terms.bad_undesired_body_contacts + params: + command_name: motion + threshold: 0.05 + body_names: + - pelvis + - left_hip_roll_link + - left_knee_link + - right_hip_roll_link + - right_knee_link + - torso_link + - left_shoulder_roll_link + - left_elbow_link + - right_shoulder_roll_link + - right_elbow_link + reward: - scales: - motion_global_root_pos: 0.5 - motion_global_root_ori: 0.5 - motion_body_pos: 2.0 - motion_body_ori: 1.5 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_ee_body_pos_z: 2.0 - motion_joint_pos: 0.0 - motion_joint_vel: 0.0 - action_rate_l2: -0.005 - joint_limit: -10.0 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 \ No newline at end of file + motion_body_pos: + weight: 2.0 + motion_body_ori: + weight: 1.5 + motion_ee_body_pos_z: + func: unilab.tasks.motion_tracking.common.manager_terms.motion_relative_body_position_z_error_exp + weight: 2.0 + params: {command_name: motion, std: 0.3, body_names: *ee_bodies} + action_rate_l2: + weight: -0.005 diff --git a/conf/appo/task/g1_wall_flip_tracking/motrix.yaml b/conf/appo/task/g1_wall_flip_tracking/motrix.yaml index 83cca7b39..3ce1d8cc3 100644 --- a/conf/appo/task/g1_wall_flip_tracking/motrix.yaml +++ b/conf/appo/task/g1_wall_flip_tracking/motrix.yaml @@ -1,50 +1,32 @@ # @package _global_ +defaults: + - /task/g1_wall_flip_tracking/mujoco + - _self_ + training: task_name: G1WallFlipTracking sim_backend: motrix - play_steps: 1000 + replay_queue_size: null + algo: - num_envs: 1024 + steps_per_env: 24 max_iterations: 5000 - save_interval: 500 algorithm: num_learning_epochs: 5 num_mini_batches: 4 - clip_param: 0.2 - gamma: 0.99 - lam: 0.95 - value_loss_coef: 1.0 entropy_coef: 0.01 - learning_rate: 1.0e-3 - max_grad_norm: 1.0 - use_clipped_value_loss: true - schedule: adaptive desired_kl: 0.01 - adaptive_kl_factor: 2.0 - adaptive_lr_factor: 1.5 - optimizer: adam - tau: 1.0 - target_update_freq: 1 - vtrace_clip_rho: 1.0 - vtrace_clip_c: 1.0 - enable_compile: true + +env: + actions: + joint_pos: + scale: 0.25 + reward: - scales: - motion_global_root_pos: 0.5 - motion_global_root_ori: 0.5 - motion_body_pos: 1.0 - motion_body_ori: 1.0 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_joint_pos: 0.0 - motion_joint_vel: 0.0 - action_rate_l2: -0.1 - joint_limit: -10.0 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 + motion_body_pos: + weight: 1.0 + motion_body_ori: + weight: 1.0 + motion_ee_body_pos_z: null + action_rate_l2: + weight: -0.1 diff --git a/conf/appo/task/g1_wall_flip_tracking/mujoco.yaml b/conf/appo/task/g1_wall_flip_tracking/mujoco.yaml index 0083729c0..d177ef8b4 100644 --- a/conf/appo/task/g1_wall_flip_tracking/mujoco.yaml +++ b/conf/appo/task/g1_wall_flip_tracking/mujoco.yaml @@ -1,93 +1,24 @@ # @package _global_ +defaults: + - /task/g1_flip_tracking/mujoco + - _self_ + training: task_name: G1WallFlipTracking sim_backend: mujoco - play_steps: 1000 replay_queue_size: 5 + algo: - num_envs: 1024 steps_per_env: 20 max_iterations: 7000 - save_interval: 500 algorithm: num_learning_epochs: 6 - num_mini_batches: 8 - clip_param: 0.2 - gamma: 0.99 - lam: 0.95 - value_loss_coef: 1.0 - entropy_coef: 0.005 - learning_rate: 1.0e-3 - max_grad_norm: 1.0 - use_clipped_value_loss: true - schedule: adaptive desired_kl: 0.008 - adaptive_kl_factor: 2.0 - adaptive_lr_factor: 1.5 - optimizer: adam - tau: 1.0 - target_update_freq: 1 - vtrace_clip_rho: 1.0 - vtrace_clip_c: 1.0 - enable_compile: true + env: - sampling_mode: start - truncate_on_clip_end: true - sim_dt: 0.005 - control_config: - action_scale: - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.07450087032950714 - - 0.07450087032950714 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.07450087032950714 - - 0.07450087032950714 - anchor_pos_z_threshold: 0.5 - ee_body_pos_z_threshold: 0.5 - terminate_on_undesired_contacts: true - noise_config: - level: 0.0 -reward: - scales: - motion_global_root_pos: 0.5 - motion_global_root_ori: 0.5 - motion_body_pos: 2.0 - motion_body_ori: 1.5 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_ee_body_pos_z: 2.0 - motion_joint_pos: 0.0 - motion_joint_vel: 0.0 - action_rate_l2: -0.005 - joint_limit: -10.0 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 \ No newline at end of file + scene: + model_file: src/unilab/assets/robots/g1/scene_flat_with_wall.xml + commands: + motion: + params: + motion_file: motions/g1/flip_from_wall_104__A304.npz diff --git a/conf/offpolicy/task/sac/g1_23dof_flip_tracking/mujoco.yaml b/conf/offpolicy/task/sac/g1_23dof_flip_tracking/mujoco.yaml index f48e6108a..0369aaeec 100644 --- a/conf/offpolicy/task/sac/g1_23dof_flip_tracking/mujoco.yaml +++ b/conf/offpolicy/task/sac/g1_23dof_flip_tracking/mujoco.yaml @@ -1,8 +1,13 @@ # @package _global_ +defaults: + - /task/sac/g1_23dof_motion_tracking/mujoco + - _self_ + training: task_name: G1FlipTrackingSAC23Dof sim_backend: mujoco play_steps: 1000 + algo: num_envs: 4096 max_iterations: 25000 @@ -17,60 +22,85 @@ algo: alpha_init: 0.005 target_entropy_ratio: 0.05 max_grad_norm: 10.0 + env: - sampling_mode: mixed - sampling_start_ratio: 0.1 - truncate_on_clip_end: true + scene: + model_file: src/unilab/assets/robots/g1/scene_flat_23dof.xml sim_dt: 0.005 - control_config: - action_scale: - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - anchor_pos_z_threshold: 0.5 - ee_body_pos_z_threshold: 0.5 - terminate_on_undesired_contacts: true - noise_config: - level: 0.0 + actions: + joint_pos: + scale: + ".*_(hip_pitch|hip_yaw)_joint": 0.5475464629911068 + ".*_(hip_roll|knee)_joint": 0.35066146637882434 + ".*_ankle_(pitch|roll)_joint": 0.43857731392336724 + "waist_yaw_joint": 0.5475464629911068 + ".*_wrist_roll_joint": 0.43857731392336724 + commands: + motion: + params: + motion_file: motions/g1/flip_360_001__A304_23dof.npz + sampling_mode: mixed + sampling_start_ratio: 0.1 + truncate_on_clip_end: true + pose_range: &zero_pose + x: [0.0, 0.0] + y: [0.0, 0.0] + z: [0.0, 0.0] + roll: [0.0, 0.0] + pitch: [0.0, 0.0] + yaw: [0.0, 0.0] + velocity_range: *zero_pose + joint_position_range: [0.0, 0.0] + terminations: + anchor_pos: + params: {command_name: motion, threshold: 0.5} + anchor_ori: + params: + command_name: motion + threshold: 1.0e9 + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + ee_body_pos: + params: + command_name: motion + threshold: 0.5 + body_names: &ee_bodies + - left_ankle_roll_link + - right_ankle_roll_link + - left_wrist_roll_rubber_hand + - right_wrist_roll_rubber_hand + undesired_contacts: + func: unilab.tasks.motion_tracking.common.manager_terms.bad_undesired_body_contacts + params: + command_name: motion + threshold: 0.05 + body_names: &undesired_bodies + - pelvis + - left_hip_roll_link + - left_knee_link + - right_hip_roll_link + - right_knee_link + - torso_link + - left_shoulder_roll_link + - left_elbow_link + - right_shoulder_roll_link + - right_elbow_link + reward: - scales: - motion_global_root_pos: 0.5 - motion_global_root_ori: 0.5 - motion_body_pos: 2.0 - motion_body_ori: 1.5 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_ee_body_pos_z: 2.0 - motion_joint_pos: 0.0 - motion_joint_vel: 0.0 - action_rate_l2: -0.005 - joint_limit: -10.0 - undesired_contacts: -0.1 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 + motion_global_root_pos: + weight: 0.5 + motion_body_ori: + weight: 1.5 + motion_ee_body_pos_z: + func: unilab.tasks.motion_tracking.common.manager_terms.motion_relative_body_position_z_error_exp + weight: 2.0 + params: {command_name: motion, std: 0.3, body_names: *ee_bodies} + action_rate_l2: + weight: -0.005 + joint_limit: + weight: -10.0 + undesired_contacts: + func: unilab.tasks.motion_tracking.common.manager_terms.undesired_body_contacts + weight: -0.1 + params: {command_name: motion, threshold: 0.05, body_names: *undesired_bodies} diff --git a/conf/offpolicy/task/sac/g1_23dof_wall_flip_tracking/mujoco.yaml b/conf/offpolicy/task/sac/g1_23dof_wall_flip_tracking/mujoco.yaml index f67d34d2e..242052e0e 100644 --- a/conf/offpolicy/task/sac/g1_23dof_wall_flip_tracking/mujoco.yaml +++ b/conf/offpolicy/task/sac/g1_23dof_wall_flip_tracking/mujoco.yaml @@ -1,75 +1,41 @@ # @package _global_ +defaults: + - /task/sac/g1_23dof_flip_tracking/mujoco + - _self_ + training: task_name: G1WallFlipTrackingSAC23Dof sim_backend: mujoco - play_steps: 1000 + algo: - num_envs: 4096 - max_iterations: 25000 - save_interval: 1000 - gamma: 0.99 - tau: 0.05 - num_atoms: 501 - updates_per_step: 4 - policy_frequency: 2 - use_symmetry: false algo_params: - alpha_init: 0.005 target_entropy_ratio: 0.0 - max_grad_norm: 10.0 + env: - sampling_mode: uniform - truncate_on_clip_end: true - sim_dt: 0.005 - control_config: - action_scale: - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - anchor_pos_z_threshold: 1000000000.0 - ee_body_pos_z_threshold: 1000000000.0 - terminate_on_undesired_contacts: false - noise_config: - level: 0.0 + scene: + model_file: src/unilab/assets/robots/g1/scene_flat_23dof_with_wall.xml + commands: + motion: + params: + motion_file: motions/g1/flip_from_wall_104__A304_23dof.npz + sampling_mode: uniform + sampling_start_ratio: 0.0 + terminations: + anchor_pos: + params: {command_name: motion, threshold: 1.0e9} + ee_body_pos: + params: + command_name: motion + threshold: 1.0e9 + body_names: + - left_ankle_roll_link + - right_ankle_roll_link + - left_wrist_roll_rubber_hand + - right_wrist_roll_rubber_hand + undesired_contacts: null + reward: - scales: - motion_global_root_pos: 0.5 - motion_global_root_ori: 0.5 - motion_body_pos: 2.0 - motion_body_ori: 1.5 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_ee_body_pos_z: 2.0 - motion_joint_pos: 0.5 - motion_joint_vel: 0.25 - action_rate_l2: -0.005 - joint_limit: -10.0 - undesired_contacts: -0.1 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 + motion_joint_pos: + weight: 0.5 + motion_joint_vel: + weight: 0.25 diff --git a/conf/offpolicy/task/sac/g1_23dof_wbt_obs/mujoco.yaml b/conf/offpolicy/task/sac/g1_23dof_wbt_obs/mujoco.yaml index d0623bad2..2d311a755 100644 --- a/conf/offpolicy/task/sac/g1_23dof_wbt_obs/mujoco.yaml +++ b/conf/offpolicy/task/sac/g1_23dof_wbt_obs/mujoco.yaml @@ -1,7 +1,12 @@ # @package _global_ +defaults: + - /task/sac/g1_23dof_motion_tracking/mujoco + - _self_ + training: task_name: G1WBTObs23Dof sim_backend: mujoco + algo: num_envs: 4096 max_iterations: 140000 @@ -16,64 +21,167 @@ algo: alpha_init: 0.1 target_entropy_ratio: 0.5 max_grad_norm: 10.0 + env: sim_dt: 0.005 - sensor: - local_linvel: pelvis_local_linvel - gyro: pelvis_gyro - upvector: pelvis_upvector - control_config: - action_scale: 2.0 - simulate_action_latency: true - anchor_pos_z_threshold: 0.40 - ee_body_pos_z_threshold: 0.5 - truncate_on_clip_end: true - noise_config: - level: 1.0 - scale_joint_angle: 0.01 - scale_joint_vel: 0.5 - scale_gyro: 0.2 - enable_zero_linvel: true - enable_zero_anchor_pos: true - enable_anchor_ori_noise: true - scale_anchor_ori: 0.05 - obs_history_length: 5 - domain_rand: - randomize_base_mass: true - added_mass_range: [-1.0, 1.0] - random_com: true - com_offset_x: [-0.05, 0.05] - randomize_com_y: true - com_offset_y: [-0.05, 0.05] - randomize_com_z: true - com_offset_z: [-0.05, 0.05] - randomize_gravity: false - gravity_range: [[0.0, 0.0, -9.81], [0.0, 0.0, -9.81]] - push_robots: true - push_interval: 200 - max_force: [300.0, 300.0, 120.0] - push_body_name: null - randomize_kp: true - kp_multiplier_range: [0.9, 1.1] - randomize_kd: true - kd_multiplier_range: [0.85, 1.15] - randomize_geom_friction: true - friction_range: [0.3, 1.2] - friction_geom_pattern: "^(left|right)_foot[1-7]_collision$" - enable_encoder_bias: true - encoder_bias_range: [-0.01, 0.01] + scene: + entities: + robot: + geom_names: + - left_foot1_collision + - left_foot2_collision + - left_foot3_collision + - left_foot4_collision + - left_foot5_collision + - left_foot6_collision + - left_foot7_collision + - right_foot1_collision + - right_foot2_collision + - right_foot3_collision + - right_foot4_collision + - right_foot5_collision + - right_foot6_collision + - right_foot7_collision + observations: + actor: + terms: + motion_anchor_pos_b: null + motion_anchor_ori_b: + func: unilab.tasks.motion_tracking.common.manager_terms.motion_anchor_ori_b + params: {command_name: motion} + noise: + _target_: unilab.managers._noise.UniformNoiseCfg + n_min: -0.05 + n_max: 0.05 + base_lin_vel: null + base_ang_vel: + func: unilab.envs.mdp.builtin_sensor + params: {sensor_name: pelvis_gyro} + noise: + _target_: unilab.managers._noise.UniformNoiseCfg + n_min: -0.2 + n_max: 0.2 + history_length: 5 + joint_pos: + func: unilab.tasks.motion_tracking.common.manager_terms.motion_joint_pos_rel_biased + params: {command_name: motion} + noise: + _target_: unilab.managers._noise.UniformNoiseCfg + n_min: -0.01 + n_max: 0.01 + history_length: 5 + joint_vel: + func: unilab.envs.mdp.joint_vel_rel + noise: + _target_: unilab.managers._noise.UniformNoiseCfg + n_min: -0.5 + n_max: 0.5 + history_length: 5 + actions: + func: unilab.envs.mdp.last_action + history_length: 5 + critic: + terms: + base_ang_vel: + params: {sensor_name: pelvis_gyro} + actions: + joint_pos: + scale: 2.0 + simulate_action_latency: true + terminations: + anchor_pos: + params: {command_name: motion, threshold: 0.4} + ee_body_pos: + params: + command_name: motion + threshold: 0.5 + body_names: + - left_ankle_roll_link + - right_ankle_roll_link + - left_wrist_roll_rubber_hand + - right_wrist_roll_rubber_hand + events: + base_mass: + func: unilab.envs.mdp.randomize_rigid_body_mass + mode: reset + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + body_names: pelvis + mass_distribution_params: [-1.0, 1.0] + operation: add + recompute_inertia: false + base_com: + func: unilab.envs.mdp.randomize_rigid_body_com + mode: reset + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + body_names: pelvis + com_range: + x: [-0.05, 0.05] + y: [-0.05, 0.05] + z: [-0.05, 0.05] + pd_gains: + func: unilab.envs.mdp.pd_gains + mode: reset + params: + kp_range: [0.9, 1.1] + kd_range: [0.85, 1.15] + operation: scale + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + actuator_names: ".*" + foot_friction: + func: unilab.envs.mdp.geom_friction + mode: reset + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + geom_names: "^(left|right)_foot[1-7]_collision$" + ranges: [0.3, 1.2] + operation: abs + shared_random: true + encoder_bias: + func: unilab.tasks.motion_tracking.g1.manager_terms.randomize_encoder_bias + mode: reset + params: + bias_range: [-0.01, 0.01] + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + joint_names: ".*" + push_robot: + func: unilab.envs.mdp.push_by_setting_velocity + mode: interval + interval_range_s: [4.0, 4.0] + is_global_time: true + params: + velocity_range: + x: [-1.0, 1.0] + y: [-1.0, 1.0] + z: [-0.5, 0.5] + roll: [0.0, 0.0] + pitch: [0.0, 0.0] + yaw: [0.0, 0.0] + reward: - scales: - motion_global_root_pos: 1.0 - motion_global_root_ori: 1.0 - motion_body_pos: 1.0 - motion_body_ori: 1.0 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_joint_pos: 0.0 - motion_joint_vel: 0.0 - action_rate_l2: -0.1 - joint_limit: -5.0 - undesired_contacts: -0.1 - joint_acc_l2: -2.5e-7 - joint_torque_l2: -1e-5 + motion_global_root_ori: + weight: 1.0 + motion_body_pos: + weight: 1.0 + action_rate_l2: + weight: -0.1 + joint_limit: + weight: -5.0 + joint_acc_l2: + func: unilab.tasks.motion_tracking.g1.manager_terms.joint_acc_l2 + weight: -2.5e-7 + joint_torque_l2: + func: unilab.tasks.motion_tracking.g1.manager_terms.joint_torque_l2 + weight: -1.0e-5 + params: {action_name: joint_pos} diff --git a/conf/offpolicy/task/sac/g1_flip_tracking/mujoco.yaml b/conf/offpolicy/task/sac/g1_flip_tracking/mujoco.yaml index 48a956b4f..4fca11f73 100644 --- a/conf/offpolicy/task/sac/g1_flip_tracking/mujoco.yaml +++ b/conf/offpolicy/task/sac/g1_flip_tracking/mujoco.yaml @@ -1,8 +1,13 @@ # @package _global_ +defaults: + - /task/sac/g1_motion_tracking/mujoco + - _self_ + training: task_name: G1FlipTrackingSAC sim_backend: mujoco play_steps: 1000 + algo: num_envs: 4096 max_iterations: 25000 @@ -17,66 +22,87 @@ algo: alpha_init: 0.005 target_entropy_ratio: 0.05 max_grad_norm: 10.0 + env: - sampling_mode: mixed - sampling_start_ratio: 0.1 - truncate_on_clip_end: true + scene: + model_file: src/unilab/assets/robots/g1/scene_flat.xml sim_dt: 0.005 - control_config: - action_scale: - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.07450087032950714 - - 0.07450087032950714 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.07450087032950714 - - 0.07450087032950714 - anchor_pos_z_threshold: 0.5 - ee_body_pos_z_threshold: 0.5 - terminate_on_undesired_contacts: true - noise_config: - level: 0.0 + actions: + joint_pos: + scale: + ".*_(hip_pitch|hip_yaw)_joint": 0.5475464629911068 + ".*_(hip_roll|knee)_joint": 0.35066146637882434 + ".*_ankle_(pitch|roll)_joint": 0.43857731392336724 + "waist_yaw_joint": 0.5475464629911068 + "waist_(roll|pitch)_joint": 0.43857731392336724 + ".*_(shoulder_(pitch|roll|yaw)|elbow|wrist_roll)_joint": 0.43857731392336724 + ".*_wrist_(pitch|yaw)_joint": 0.07450087032950714 + commands: + motion: + params: + motion_file: motions/g1/flip_360_001__A304.npz + sampling_mode: mixed + sampling_start_ratio: 0.1 + truncate_on_clip_end: true + pose_range: &zero_pose + x: [0.0, 0.0] + y: [0.0, 0.0] + z: [0.0, 0.0] + roll: [0.0, 0.0] + pitch: [0.0, 0.0] + yaw: [0.0, 0.0] + velocity_range: *zero_pose + joint_position_range: [0.0, 0.0] + terminations: + anchor_pos: + params: {command_name: motion, threshold: 0.5} + anchor_ori: + params: + command_name: motion + threshold: 1.0e9 + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + ee_body_pos: + params: + command_name: motion + threshold: 0.5 + body_names: &ee_bodies + - left_ankle_roll_link + - right_ankle_roll_link + - left_wrist_yaw_link + - right_wrist_yaw_link + undesired_contacts: + func: unilab.tasks.motion_tracking.common.manager_terms.bad_undesired_body_contacts + params: + command_name: motion + threshold: 0.05 + body_names: &undesired_bodies + - pelvis + - left_hip_roll_link + - left_knee_link + - right_hip_roll_link + - right_knee_link + - torso_link + - left_shoulder_roll_link + - left_elbow_link + - right_shoulder_roll_link + - right_elbow_link + reward: - scales: - motion_global_root_pos: 0.5 - motion_global_root_ori: 0.5 - motion_body_pos: 2.0 - motion_body_ori: 1.5 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_ee_body_pos_z: 2.0 - motion_joint_pos: 0.0 - motion_joint_vel: 0.0 - action_rate_l2: -0.005 - joint_limit: -10.0 - undesired_contacts: -0.1 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 + motion_global_root_pos: + weight: 0.5 + motion_body_ori: + weight: 1.5 + motion_ee_body_pos_z: + func: unilab.tasks.motion_tracking.common.manager_terms.motion_relative_body_position_z_error_exp + weight: 2.0 + params: {command_name: motion, std: 0.3, body_names: *ee_bodies} + action_rate_l2: + weight: -0.005 + joint_limit: + weight: -10.0 + undesired_contacts: + func: unilab.tasks.motion_tracking.common.manager_terms.undesired_body_contacts + weight: -0.1 + params: {command_name: motion, threshold: 0.05, body_names: *undesired_bodies} diff --git a/conf/offpolicy/task/sac/g1_wall_flip_tracking/mujoco.yaml b/conf/offpolicy/task/sac/g1_wall_flip_tracking/mujoco.yaml index da48802f1..3b831e946 100644 --- a/conf/offpolicy/task/sac/g1_wall_flip_tracking/mujoco.yaml +++ b/conf/offpolicy/task/sac/g1_wall_flip_tracking/mujoco.yaml @@ -1,81 +1,41 @@ # @package _global_ +defaults: + - /task/sac/g1_flip_tracking/mujoco + - _self_ + training: task_name: G1WallFlipTrackingSAC sim_backend: mujoco - play_steps: 1000 + algo: - num_envs: 4096 - max_iterations: 25000 - save_interval: 1000 - gamma: 0.99 - tau: 0.05 - num_atoms: 501 - updates_per_step: 4 - policy_frequency: 2 - use_symmetry: false algo_params: - alpha_init: 0.005 target_entropy_ratio: 0.0 - max_grad_norm: 10.0 + env: - sampling_mode: uniform - truncate_on_clip_end: true - sim_dt: 0.005 - control_config: - action_scale: - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.07450087032950714 - - 0.07450087032950714 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.07450087032950714 - - 0.07450087032950714 - anchor_pos_z_threshold: 1000000000.0 - ee_body_pos_z_threshold: 1000000000.0 - terminate_on_undesired_contacts: false - noise_config: - level: 0.0 + scene: + model_file: src/unilab/assets/robots/g1/scene_flat_with_wall.xml + commands: + motion: + params: + motion_file: motions/g1/flip_from_wall_104__A304.npz + sampling_mode: uniform + sampling_start_ratio: 0.0 + terminations: + anchor_pos: + params: {command_name: motion, threshold: 1.0e9} + ee_body_pos: + params: + command_name: motion + threshold: 1.0e9 + body_names: + - left_ankle_roll_link + - right_ankle_roll_link + - left_wrist_yaw_link + - right_wrist_yaw_link + undesired_contacts: null + reward: - scales: - motion_global_root_pos: 0.5 - motion_global_root_ori: 0.5 - motion_body_pos: 2.0 - motion_body_ori: 1.5 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_ee_body_pos_z: 2.0 - motion_joint_pos: 0.5 - motion_joint_vel: 0.25 - action_rate_l2: -0.005 - joint_limit: -10.0 - undesired_contacts: -0.1 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 + motion_joint_pos: + weight: 0.5 + motion_joint_vel: + weight: 0.25 diff --git a/conf/offpolicy/task/sac/g1_wbt_obs/mujoco.yaml b/conf/offpolicy/task/sac/g1_wbt_obs/mujoco.yaml index 5fd91fc07..6c37235f0 100644 --- a/conf/offpolicy/task/sac/g1_wbt_obs/mujoco.yaml +++ b/conf/offpolicy/task/sac/g1_wbt_obs/mujoco.yaml @@ -1,7 +1,12 @@ # @package _global_ +defaults: + - /task/sac/g1_motion_tracking/mujoco + - _self_ + training: task_name: G1WBTObs sim_backend: mujoco + algo: num_envs: 4096 max_iterations: 140000 @@ -16,64 +21,166 @@ algo: alpha_init: 0.1 target_entropy_ratio: 0.5 max_grad_norm: 10.0 + env: sim_dt: 0.005 - sensor: - local_linvel: pelvis_local_linvel - gyro: pelvis_gyro - upvector: pelvis_upvector - control_config: - action_scale: 2.0 - simulate_action_latency: true - anchor_pos_z_threshold: 0.40 - ee_body_pos_z_threshold: 0.5 - truncate_on_clip_end: true - noise_config: - level: 1.0 - scale_joint_angle: 0.01 - scale_joint_vel: 0.5 - scale_gyro: 0.2 - enable_zero_linvel: true - enable_zero_anchor_pos: true - enable_anchor_ori_noise: true - scale_anchor_ori: 0.05 - obs_history_length: 5 - domain_rand: - randomize_base_mass: true - added_mass_range: [-1.0, 1.0] - random_com: true - com_offset_x: [-0.05, 0.05] - randomize_com_y: true - com_offset_y: [-0.05, 0.05] - randomize_com_z: true - com_offset_z: [-0.05, 0.05] - randomize_gravity: false - gravity_range: [[0.0, 0.0, -9.81], [0.0, 0.0, -9.81]] - push_robots: true - push_interval: 200 - max_force: [300.0, 300.0, 120.0] - push_body_name: null - randomize_kp: true - kp_multiplier_range: [0.9, 1.1] - randomize_kd: true - kd_multiplier_range: [0.85, 1.15] - randomize_geom_friction: true - friction_range: [0.3, 1.2] - friction_geom_pattern: "^(left|right)_foot[1-7]_collision$" - enable_encoder_bias: true - encoder_bias_range: [-0.01, 0.01] + scene: + entities: + robot: + geom_names: + - left_foot1_collision + - left_foot2_collision + - left_foot3_collision + - left_foot4_collision + - left_foot5_collision + - left_foot6_collision + - left_foot7_collision + - right_foot1_collision + - right_foot2_collision + - right_foot3_collision + - right_foot4_collision + - right_foot5_collision + - right_foot6_collision + - right_foot7_collision + observations: + actor: + terms: + motion_anchor_pos_b: null + motion_anchor_ori_b: + func: unilab.tasks.motion_tracking.common.manager_terms.motion_anchor_ori_b + params: {command_name: motion} + noise: + _target_: unilab.managers._noise.UniformNoiseCfg + n_min: -0.05 + n_max: 0.05 + base_lin_vel: null + base_ang_vel: + func: unilab.envs.mdp.builtin_sensor + params: {sensor_name: pelvis_gyro} + noise: + _target_: unilab.managers._noise.UniformNoiseCfg + n_min: -0.2 + n_max: 0.2 + history_length: 5 + joint_pos: + func: unilab.tasks.motion_tracking.common.manager_terms.motion_joint_pos_rel_biased + params: {command_name: motion} + noise: + _target_: unilab.managers._noise.UniformNoiseCfg + n_min: -0.01 + n_max: 0.01 + history_length: 5 + joint_vel: + func: unilab.envs.mdp.joint_vel_rel + noise: + _target_: unilab.managers._noise.UniformNoiseCfg + n_min: -0.5 + n_max: 0.5 + history_length: 5 + actions: + func: unilab.envs.mdp.last_action + history_length: 5 + critic: + terms: + base_ang_vel: + params: {sensor_name: pelvis_gyro} + actions: + joint_pos: + simulate_action_latency: true + terminations: + anchor_pos: + params: {command_name: motion, threshold: 0.4} + ee_body_pos: + params: + command_name: motion + threshold: 0.5 + body_names: + - left_ankle_roll_link + - right_ankle_roll_link + - left_wrist_yaw_link + - right_wrist_yaw_link + events: + base_mass: + func: unilab.envs.mdp.randomize_rigid_body_mass + mode: reset + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + body_names: pelvis + mass_distribution_params: [-1.0, 1.0] + operation: add + recompute_inertia: false + base_com: + func: unilab.envs.mdp.randomize_rigid_body_com + mode: reset + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + body_names: pelvis + com_range: + x: [-0.05, 0.05] + y: [-0.05, 0.05] + z: [-0.05, 0.05] + pd_gains: + func: unilab.envs.mdp.pd_gains + mode: reset + params: + kp_range: [0.9, 1.1] + kd_range: [0.85, 1.15] + operation: scale + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + actuator_names: ".*" + foot_friction: + func: unilab.envs.mdp.geom_friction + mode: reset + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + geom_names: "^(left|right)_foot[1-7]_collision$" + ranges: [0.3, 1.2] + operation: abs + shared_random: true + encoder_bias: + func: unilab.tasks.motion_tracking.g1.manager_terms.randomize_encoder_bias + mode: reset + params: + bias_range: [-0.01, 0.01] + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + joint_names: ".*" + push_robot: + func: unilab.envs.mdp.push_by_setting_velocity + mode: interval + interval_range_s: [4.0, 4.0] + is_global_time: true + params: + velocity_range: + x: [-1.0, 1.0] + y: [-1.0, 1.0] + z: [-0.5, 0.5] + roll: [0.0, 0.0] + pitch: [0.0, 0.0] + yaw: [0.0, 0.0] + reward: - scales: - motion_global_root_pos: 1.0 - motion_global_root_ori: 1.0 - motion_body_pos: 1.0 - motion_body_ori: 1.0 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_joint_pos: 0.0 - motion_joint_vel: 0.0 - action_rate_l2: -0.1 - joint_limit: -5.0 - undesired_contacts: -0.1 - joint_acc_l2: -2.5e-7 - joint_torque_l2: -1e-5 + motion_global_root_ori: + weight: 1.0 + motion_body_pos: + weight: 1.0 + action_rate_l2: + weight: -0.1 + joint_limit: + weight: -5.0 + joint_acc_l2: + func: unilab.tasks.motion_tracking.g1.manager_terms.joint_acc_l2 + weight: -2.5e-7 + joint_torque_l2: + func: unilab.tasks.motion_tracking.g1.manager_terms.joint_torque_l2 + weight: -1.0e-5 + params: {action_name: joint_pos} diff --git a/conf/ppo/task/g1_23dof_box_tracking/motrix.yaml b/conf/ppo/task/g1_23dof_box_tracking/motrix.yaml index e453b0721..c9456f1e4 100644 --- a/conf/ppo/task/g1_23dof_box_tracking/motrix.yaml +++ b/conf/ppo/task/g1_23dof_box_tracking/motrix.yaml @@ -1,27 +1,37 @@ # @package _global_ +defaults: + - /task/g1_23dof_box_tracking/mujoco + - _self_ + training: task_name: G1BoxTracking23Dof sim_backend: motrix play_env_num: 16 - play_steps: 1000 + algo: - num_envs: 1024 max_iterations: 40000 - save_interval: 500 empirical_normalization: true obs_groups: - actor: - - actor - critic: - - critic + actor: [actor] + critic: [critic] algorithm: entropy_coef: 0.002 desired_kl: 0.01 - noise_config: - level: 1.0 - scale_joint_angle: 0.01 - scale_joint_vel: 1.5 - scale_gyro: 0.2 + +reward: + motion_global_root_pos: + weight: 1.0 + motion_body_ori: + weight: 1.5 + motion_body_ang_vel: + weight: 1.5 + object_global_ref_position_error_exp: + weight: 4.0 + params: {command_name: motion, std: 0.12} + object_global_ref_orientation_error_exp: + weight: 3.0 + params: {command_name: motion, std: 0.2} + play_profile: enabled: true env: @@ -33,28 +43,3 @@ play_profile: skybox_rgb1: [0.90, 0.90, 0.91] skybox_rgb2: [0.68, 0.68, 0.70] ground_texrepeat: [0.25, 0.25] -reward: - scales: - motion_global_root_pos: 1.0 - motion_global_root_ori: 0.5 - motion_body_pos: 1.0 - motion_body_ori: 1.5 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.5 - motion_joint_pos: 0.0 - motion_joint_vel: 0.0 - action_rate_l2: -0.1 - joint_limit: -10.0 - undesired_contacts: -0.1 - object_global_ref_position_error_exp: 4.0 - object_global_ref_orientation_error_exp: 3.0 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 - std_object_pos: 0.12 - std_object_ori: 0.2 diff --git a/conf/ppo/task/g1_23dof_box_tracking/mujoco.yaml b/conf/ppo/task/g1_23dof_box_tracking/mujoco.yaml index 4c7b186f3..d835ffe06 100644 --- a/conf/ppo/task/g1_23dof_box_tracking/mujoco.yaml +++ b/conf/ppo/task/g1_23dof_box_tracking/mujoco.yaml @@ -1,44 +1,84 @@ # @package _global_ +defaults: + - /task/g1_23dof_motion_tracking/mujoco + - _self_ + training: task_name: G1BoxTracking23Dof sim_backend: mujoco play_steps: 1000 + algo: num_envs: 1024 max_iterations: 30000 save_interval: 500 obs_groups: - actor: - - actor + actor: [actor] algorithm: entropy_coef: 0.005 + +play_profile: + enabled: false + env: null + env: + scene: + model_file: src/unilab/assets/robots/g1/scene_flat_23dof_with_largebox.xml + entities: + object: + root_body_name: largebox sim_dt: 0.005 - sensor: - gyro: pelvis_gyro - upvector: pelvis_upvector + observations: + actor: + terms: + motion_anchor_pos_b: null + base_lin_vel: null + base_ang_vel: + params: {sensor_name: pelvis_gyro} + critic: + terms: + base_ang_vel: + params: {sensor_name: pelvis_gyro} + object_state: + func: unilab.tasks.motion_tracking.g1.manager_terms.object_state_b + params: {command_name: motion} + commands: + motion: + _target_: unilab.tasks.motion_tracking.g1.manager_terms.BoxMotionCommandCfg + object_entity_name: object + params: + motion_file: motions/g1/sub3_largebox_003_boxconverted_23dof.npz + terminations: + object_pos: + func: unilab.tasks.motion_tracking.g1.manager_terms.bad_object_position + params: {command_name: motion, threshold: 0.25} + object_ori: + func: unilab.tasks.motion_tracking.g1.manager_terms.bad_object_orientation + params: {command_name: motion, threshold: 0.8} + reward: - scales: - motion_global_root_pos: 0.5 - motion_global_root_ori: 0.5 - motion_body_pos: 1.0 - motion_body_ori: 1.0 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_joint_pos: 0.0 - motion_joint_vel: 0.0 - action_rate_l2: -0.1 - joint_limit: -10.0 - undesired_contacts: -0.1 - object_global_ref_position_error_exp: 2.0 - object_global_ref_orientation_error_exp: 2.0 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 - std_object_pos: 0.2 - std_object_ori: 0.3 + undesired_contacts: + func: unilab.tasks.motion_tracking.common.manager_terms.undesired_body_contacts + weight: -0.1 + params: + command_name: motion + threshold: 0.05 + body_names: + - pelvis + - left_hip_roll_link + - left_knee_link + - right_hip_roll_link + - right_knee_link + - torso_link + - left_shoulder_roll_link + - left_elbow_link + - right_shoulder_roll_link + - right_elbow_link + object_global_ref_position_error_exp: + func: unilab.tasks.motion_tracking.g1.manager_terms.object_global_position_error_exp + weight: 2.0 + params: {command_name: motion, std: 0.2} + object_global_ref_orientation_error_exp: + func: unilab.tasks.motion_tracking.g1.manager_terms.object_global_orientation_error_exp + weight: 2.0 + params: {command_name: motion, std: 0.3} diff --git a/conf/ppo/task/g1_23dof_climb_tracking/motrix.yaml b/conf/ppo/task/g1_23dof_climb_tracking/motrix.yaml index f554a13d7..ae708bf1d 100644 --- a/conf/ppo/task/g1_23dof_climb_tracking/motrix.yaml +++ b/conf/ppo/task/g1_23dof_climb_tracking/motrix.yaml @@ -1,74 +1,8 @@ # @package _global_ +defaults: + - /task/g1_23dof_climb_tracking/mujoco + - _self_ + training: task_name: G1ClimbTracking23Dof sim_backend: motrix - play_steps: 1000 -algo: - num_envs: 1024 - max_iterations: 20000 - save_interval: 500 - empirical_normalization: true - obs_groups: - actor: - - actor - critic: - - critic - algorithm: - entropy_coef: 0.005 - desired_kl: 0.01 -env: - sampling_mode: adaptive - truncate_on_clip_end: false - sim_dt: 0.005 - control_config: - action_scale: - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - anchor_pos_z_threshold: 0.3 - ee_body_pos_z_threshold: 0.3 - terminate_on_undesired_contacts: true - noise_config: - level: 0.0 -reward: - scales: - motion_global_root_pos: 0.5 - motion_global_root_ori: 0.5 - motion_body_pos: 2.0 - motion_body_ori: 1.5 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_ee_body_pos_z: 2.0 - motion_joint_pos: 0.5 - motion_joint_vel: 0.25 - action_rate_l2: -0.005 - joint_limit: -10.0 - undesired_contacts: -0.1 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 diff --git a/conf/ppo/task/g1_23dof_climb_tracking/mujoco.yaml b/conf/ppo/task/g1_23dof_climb_tracking/mujoco.yaml index 34fc4ae59..99706fa98 100644 --- a/conf/ppo/task/g1_23dof_climb_tracking/mujoco.yaml +++ b/conf/ppo/task/g1_23dof_climb_tracking/mujoco.yaml @@ -1,74 +1,93 @@ # @package _global_ +defaults: + - /task/g1_23dof_motion_tracking/mujoco + - _self_ + training: task_name: G1ClimbTracking23Dof sim_backend: mujoco play_steps: 1000 + algo: num_envs: 1024 max_iterations: 20000 save_interval: 500 empirical_normalization: true obs_groups: - actor: - - actor - critic: - - critic + actor: [actor] + critic: [critic] algorithm: entropy_coef: 0.005 desired_kl: 0.01 + +play_profile: + enabled: false + env: null + env: - sampling_mode: adaptive - truncate_on_clip_end: false + scene: + model_file: src/unilab/assets/robots/g1/scene_climb_20_z_scale_1_23dof.xml sim_dt: 0.005 - control_config: - action_scale: - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - anchor_pos_z_threshold: 0.3 - ee_body_pos_z_threshold: 0.3 - terminate_on_undesired_contacts: true - noise_config: - level: 0.0 + max_episode_seconds: 15.0 + actions: + joint_pos: + scale: + ".*_(hip_pitch|hip_yaw)_joint": 0.5475464629911068 + ".*_(hip_roll|knee)_joint": 0.35066146637882434 + ".*_ankle_(pitch|roll)_joint": 0.43857731392336724 + "waist_yaw_joint": 0.5475464629911068 + ".*_(shoulder_(pitch|roll|yaw)|elbow|wrist_roll)_joint": 0.43857731392336724 + commands: + motion: + params: + motion_file: motions/g1/climb_20_z_scale_1.0_23dof.npz + sampling_mode: adaptive + truncate_on_clip_end: false + terminations: + anchor_pos: + params: {command_name: motion, threshold: 0.3} + ee_body_pos: + params: + command_name: motion + threshold: 0.3 + body_names: &ee_bodies + - left_ankle_roll_link + - right_ankle_roll_link + - left_wrist_roll_rubber_hand + - right_wrist_roll_rubber_hand + undesired_contacts: + func: unilab.tasks.motion_tracking.common.manager_terms.bad_undesired_body_contacts + params: + command_name: motion + threshold: 0.05 + body_names: &undesired_bodies + - pelvis + - left_hip_roll_link + - left_knee_link + - right_hip_roll_link + - right_knee_link + - torso_link + - left_shoulder_roll_link + - left_elbow_link + - right_shoulder_roll_link + - right_elbow_link + reward: - scales: - motion_global_root_pos: 0.5 - motion_global_root_ori: 0.5 - motion_body_pos: 2.0 - motion_body_ori: 1.5 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_ee_body_pos_z: 2.0 - motion_joint_pos: 0.5 - motion_joint_vel: 0.25 - action_rate_l2: -0.005 - joint_limit: -10.0 - undesired_contacts: -0.1 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 + motion_body_pos: + weight: 2.0 + motion_body_ori: + weight: 1.5 + motion_ee_body_pos_z: + func: unilab.tasks.motion_tracking.common.manager_terms.motion_relative_body_position_z_error_exp + weight: 2.0 + params: {command_name: motion, std: 0.3, body_names: *ee_bodies} + motion_joint_pos: + weight: 0.5 + motion_joint_vel: + weight: 0.25 + action_rate_l2: + weight: -0.005 + undesired_contacts: + func: unilab.tasks.motion_tracking.common.manager_terms.undesired_body_contacts + weight: -0.1 + params: {command_name: motion, threshold: 0.05, body_names: *undesired_bodies} diff --git a/conf/ppo/task/g1_23dof_flip_tracking/motrix.yaml b/conf/ppo/task/g1_23dof_flip_tracking/motrix.yaml index fff96e018..fe5afd7e6 100644 --- a/conf/ppo/task/g1_23dof_flip_tracking/motrix.yaml +++ b/conf/ppo/task/g1_23dof_flip_tracking/motrix.yaml @@ -1,34 +1,29 @@ # @package _global_ +defaults: + - /task/g1_23dof_flip_tracking/mujoco + - _self_ + training: task_name: G1FlipTracking23Dof sim_backend: motrix - play_steps: 1000 + algo: - num_envs: 1024 max_iterations: 30000 - save_interval: 500 - obs_groups: - actor: - - actor - algorithm: - entropy_coef: 0.005 + empirical_normalization: false + +env: + actions: + joint_pos: + scale: 0.25 + reward: - scales: - motion_global_root_pos: 1.0 - motion_global_root_ori: 0.5 - motion_body_pos: 1.0 - motion_body_ori: 1.0 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_joint_pos: 0.0 - motion_joint_vel: 0.0 - action_rate_l2: -0.05 - joint_limit: -10.0 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 + motion_global_root_pos: + weight: 1.0 + motion_body_pos: + weight: 1.0 + motion_body_ori: + weight: 1.0 + motion_ee_body_pos_z: null + action_rate_l2: + weight: -0.05 + undesired_contacts: null diff --git a/conf/ppo/task/g1_23dof_flip_tracking/mujoco.yaml b/conf/ppo/task/g1_23dof_flip_tracking/mujoco.yaml index 8e62a80ff..5cec48e23 100644 --- a/conf/ppo/task/g1_23dof_flip_tracking/mujoco.yaml +++ b/conf/ppo/task/g1_23dof_flip_tracking/mujoco.yaml @@ -1,74 +1,104 @@ # @package _global_ +defaults: + - /task/g1_23dof_motion_tracking/mujoco + - _self_ + training: task_name: G1FlipTracking23Dof sim_backend: mujoco play_steps: 1000 + algo: num_envs: 1024 max_iterations: 20000 save_interval: 500 empirical_normalization: true obs_groups: - actor: - - actor - critic: - - critic + actor: [actor] + critic: [critic] algorithm: entropy_coef: 0.005 desired_kl: 0.01 + +play_profile: + enabled: false + env: null + env: - sampling_mode: start - truncate_on_clip_end: false + scene: + model_file: src/unilab/assets/robots/g1/scene_flat_23dof.xml sim_dt: 0.005 - control_config: - action_scale: - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - anchor_pos_z_threshold: 0.5 - ee_body_pos_z_threshold: 0.5 - terminate_on_undesired_contacts: true - noise_config: - level: 0.0 + actions: + joint_pos: + scale: + ".*_(hip_pitch|hip_yaw)_joint": 0.5475464629911068 + ".*_(hip_roll|knee)_joint": 0.35066146637882434 + ".*_ankle_(pitch|roll)_joint": 0.43857731392336724 + "waist_yaw_joint": 0.5475464629911068 + ".*_(shoulder_(pitch|roll|yaw)|elbow|wrist_roll)_joint": 0.43857731392336724 + commands: + motion: + params: + motion_file: motions/g1/flip_360_001__A304_23dof.npz + sampling_mode: start + truncate_on_clip_end: false + pose_range: &zero_pose + x: [0.0, 0.0] + y: [0.0, 0.0] + z: [0.0, 0.0] + roll: [0.0, 0.0] + pitch: [0.0, 0.0] + yaw: [0.0, 0.0] + velocity_range: *zero_pose + joint_position_range: [0.0, 0.0] + terminations: + anchor_pos: + params: {command_name: motion, threshold: 0.5} + anchor_ori: + params: + command_name: motion + threshold: 1.0e9 + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + ee_body_pos: + params: + command_name: motion + threshold: 0.5 + body_names: &ee_bodies + - left_ankle_roll_link + - right_ankle_roll_link + - left_wrist_roll_rubber_hand + - right_wrist_roll_rubber_hand + undesired_contacts: + func: unilab.tasks.motion_tracking.common.manager_terms.bad_undesired_body_contacts + params: + command_name: motion + threshold: 0.05 + body_names: &undesired_bodies + - pelvis + - left_hip_roll_link + - left_knee_link + - right_hip_roll_link + - right_knee_link + - torso_link + - left_shoulder_roll_link + - left_elbow_link + - right_shoulder_roll_link + - right_elbow_link + reward: - scales: - motion_global_root_pos: 0.5 - motion_global_root_ori: 0.5 - motion_body_pos: 2.0 - motion_body_ori: 1.5 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_ee_body_pos_z: 2.0 - motion_joint_pos: 0.0 - motion_joint_vel: 0.0 - action_rate_l2: -0.005 - joint_limit: -10.0 - undesired_contacts: -0.1 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 + motion_body_pos: + weight: 2.0 + motion_body_ori: + weight: 1.5 + motion_ee_body_pos_z: + func: unilab.tasks.motion_tracking.common.manager_terms.motion_relative_body_position_z_error_exp + weight: 2.0 + params: {command_name: motion, std: 0.3, body_names: *ee_bodies} + action_rate_l2: + weight: -0.005 + undesired_contacts: + func: unilab.tasks.motion_tracking.common.manager_terms.undesired_body_contacts + weight: -0.1 + params: {command_name: motion, threshold: 0.05, body_names: *undesired_bodies} diff --git a/conf/ppo/task/g1_23dof_wall_flip_tracking/motrix.yaml b/conf/ppo/task/g1_23dof_wall_flip_tracking/motrix.yaml index 04c69a001..ca2f0e7be 100644 --- a/conf/ppo/task/g1_23dof_wall_flip_tracking/motrix.yaml +++ b/conf/ppo/task/g1_23dof_wall_flip_tracking/motrix.yaml @@ -1,58 +1,20 @@ # @package _global_ +defaults: + - /task/g1_23dof_wall_flip_tracking/mujoco + - _self_ + training: task_name: G1WallFlipTracking23Dof sim_backend: motrix play_env_num: 16 - play_steps: 1000 render_spacing: 3.0 + algo: - num_envs: 1024 max_iterations: 12000 - save_interval: 500 - empirical_normalization: true - obs_groups: - actor: - - actor - critic: - - critic - algorithm: - entropy_coef: 0.005 - desired_kl: 0.01 + env: motrix_max_iterations: 3 - sampling_mode: start - truncate_on_clip_end: false - sim_dt: 0.005 - control_config: - action_scale: - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - anchor_pos_z_threshold: 0.5 - ee_body_pos_z_threshold: 0.5 - terminate_on_undesired_contacts: true - noise_config: - level: 0.0 + play_profile: enabled: true env: @@ -64,25 +26,3 @@ play_profile: skybox_rgb1: [0.90, 0.90, 0.91] skybox_rgb2: [0.68, 0.68, 0.70] ground_texrepeat: [0.25, 0.25] -reward: - scales: - motion_global_root_pos: 0.5 - motion_global_root_ori: 0.5 - motion_body_pos: 2.0 - motion_body_ori: 1.5 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_ee_body_pos_z: 2.0 - motion_joint_pos: 0.5 - motion_joint_vel: 0.25 - action_rate_l2: -0.005 - joint_limit: -10.0 - undesired_contacts: -0.1 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 diff --git a/conf/ppo/task/g1_23dof_wall_flip_tracking/mujoco.yaml b/conf/ppo/task/g1_23dof_wall_flip_tracking/mujoco.yaml index 7a49d091d..3f2c7ec81 100644 --- a/conf/ppo/task/g1_23dof_wall_flip_tracking/mujoco.yaml +++ b/conf/ppo/task/g1_23dof_wall_flip_tracking/mujoco.yaml @@ -1,74 +1,26 @@ # @package _global_ +defaults: + - /task/g1_23dof_flip_tracking/mujoco + - _self_ + training: task_name: G1WallFlipTracking23Dof sim_backend: mujoco - play_steps: 1000 -algo: - num_envs: 1024 - max_iterations: 20000 - save_interval: 500 - empirical_normalization: true - obs_groups: - actor: - - actor - critic: - - critic - algorithm: - entropy_coef: 0.005 - desired_kl: 0.01 + +play_profile: + enabled: false + env: null + env: - sampling_mode: start - truncate_on_clip_end: false - sim_dt: 0.005 - control_config: - action_scale: - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - anchor_pos_z_threshold: 0.5 - ee_body_pos_z_threshold: 0.5 - terminate_on_undesired_contacts: true - noise_config: - level: 0.0 + scene: + model_file: src/unilab/assets/robots/g1/scene_flat_23dof_with_wall.xml + commands: + motion: + params: + motion_file: motions/g1/flip_from_wall_104__A304_23dof.npz + reward: - scales: - motion_global_root_pos: 0.5 - motion_global_root_ori: 0.5 - motion_body_pos: 2.0 - motion_body_ori: 1.5 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_ee_body_pos_z: 2.0 - motion_joint_pos: 0.5 - motion_joint_vel: 0.25 - action_rate_l2: -0.005 - joint_limit: -10.0 - undesired_contacts: -0.1 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 + motion_joint_pos: + weight: 0.5 + motion_joint_vel: + weight: 0.25 diff --git a/conf/ppo/task/g1_box_tracking/motrix.yaml b/conf/ppo/task/g1_box_tracking/motrix.yaml index 863ac1215..ba9ad2abc 100644 --- a/conf/ppo/task/g1_box_tracking/motrix.yaml +++ b/conf/ppo/task/g1_box_tracking/motrix.yaml @@ -1,27 +1,37 @@ # @package _global_ +defaults: + - /task/g1_box_tracking/mujoco + - _self_ + training: task_name: G1BoxTracking sim_backend: motrix play_env_num: 16 - play_steps: 1000 + algo: - num_envs: 1024 max_iterations: 40000 - save_interval: 500 empirical_normalization: true obs_groups: - actor: - - actor - critic: - - critic + actor: [actor] + critic: [critic] algorithm: entropy_coef: 0.002 desired_kl: 0.01 - noise_config: - level: 1.0 - scale_joint_angle: 0.01 - scale_joint_vel: 1.5 - scale_gyro: 0.2 + +reward: + motion_global_root_pos: + weight: 1.0 + motion_body_ori: + weight: 1.5 + motion_body_ang_vel: + weight: 1.5 + object_global_ref_position_error_exp: + weight: 4.0 + params: {command_name: motion, std: 0.12} + object_global_ref_orientation_error_exp: + weight: 3.0 + params: {command_name: motion, std: 0.2} + play_profile: enabled: true env: @@ -33,28 +43,3 @@ play_profile: skybox_rgb1: [0.90, 0.90, 0.91] skybox_rgb2: [0.68, 0.68, 0.70] ground_texrepeat: [0.25, 0.25] -reward: - scales: - motion_global_root_pos: 1.0 - motion_global_root_ori: 0.5 - motion_body_pos: 1.0 - motion_body_ori: 1.5 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.5 - motion_joint_pos: 0.0 - motion_joint_vel: 0.0 - action_rate_l2: -0.1 - joint_limit: -10.0 - undesired_contacts: -0.1 - object_global_ref_position_error_exp: 4.0 - object_global_ref_orientation_error_exp: 3.0 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 - std_object_pos: 0.12 - std_object_ori: 0.2 diff --git a/conf/ppo/task/g1_box_tracking/mujoco.yaml b/conf/ppo/task/g1_box_tracking/mujoco.yaml index 278a8b8f3..73d278bf9 100644 --- a/conf/ppo/task/g1_box_tracking/mujoco.yaml +++ b/conf/ppo/task/g1_box_tracking/mujoco.yaml @@ -1,47 +1,84 @@ # @package _global_ +defaults: + - /task/g1_motion_tracking/mujoco + - _self_ + training: task_name: G1BoxTracking sim_backend: mujoco play_steps: 1000 + algo: num_envs: 1024 max_iterations: 30000 save_interval: 500 obs_groups: - actor: - - actor + actor: [actor] algorithm: entropy_coef: 0.005 + env: + scene: + model_file: src/unilab/assets/robots/g1/scene_flat_with_largebox.xml + entities: + object: + root_body_name: largebox sim_dt: 0.005 - sensor: - gyro: pelvis_gyro - upvector: pelvis_upvector + observations: + actor: + terms: + motion_anchor_pos_b: null + base_lin_vel: null + base_ang_vel: + params: {sensor_name: pelvis_gyro} + critic: + terms: + base_ang_vel: + params: {sensor_name: pelvis_gyro} + object_state: + func: unilab.tasks.motion_tracking.g1.manager_terms.object_state_b + params: {command_name: motion} + commands: + motion: + _target_: unilab.tasks.motion_tracking.g1.manager_terms.BoxMotionCommandCfg + object_entity_name: object + params: + motion_file: motions/g1/sub3_largebox_003_boxconverted.npz + terminations: + object_pos: + func: unilab.tasks.motion_tracking.g1.manager_terms.bad_object_position + params: {command_name: motion, threshold: 0.25} + object_ori: + func: unilab.tasks.motion_tracking.g1.manager_terms.bad_object_orientation + params: {command_name: motion, threshold: 0.8} + reward: - scales: - motion_global_root_pos: 0.5 - motion_global_root_ori: 0.5 - motion_body_pos: 1.0 - motion_body_ori: 1.0 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_joint_pos: 0.0 - motion_joint_vel: 0.0 - action_rate_l2: -0.1 - joint_limit: -10.0 - undesired_contacts: -0.1 - object_global_ref_position_error_exp: 2.0 - object_global_ref_orientation_error_exp: 2.0 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 - std_object_pos: 0.2 - std_object_ori: 0.3 + undesired_contacts: + func: unilab.tasks.motion_tracking.common.manager_terms.undesired_body_contacts + weight: -0.1 + params: + command_name: motion + threshold: 0.05 + body_names: + - pelvis + - left_hip_roll_link + - left_knee_link + - right_hip_roll_link + - right_knee_link + - torso_link + - left_shoulder_roll_link + - left_elbow_link + - right_shoulder_roll_link + - right_elbow_link + object_global_ref_position_error_exp: + func: unilab.tasks.motion_tracking.g1.manager_terms.object_global_position_error_exp + weight: 2.0 + params: {command_name: motion, std: 0.2} + object_global_ref_orientation_error_exp: + func: unilab.tasks.motion_tracking.g1.manager_terms.object_global_orientation_error_exp + weight: 2.0 + params: {command_name: motion, std: 0.3} + play_profile: enabled: true env: diff --git a/conf/ppo/task/g1_climb_tracking/motrix.yaml b/conf/ppo/task/g1_climb_tracking/motrix.yaml index 0c30db38e..01a98604a 100644 --- a/conf/ppo/task/g1_climb_tracking/motrix.yaml +++ b/conf/ppo/task/g1_climb_tracking/motrix.yaml @@ -1,84 +1,8 @@ # @package _global_ +defaults: + - /task/g1_climb_tracking/mujoco + - _self_ + training: task_name: G1ClimbTracking sim_backend: motrix - play_steps: 1000 -algo: - num_envs: 1024 - max_iterations: 20000 - save_interval: 500 - empirical_normalization: true - obs_groups: - actor: - - actor - critic: - - critic - algorithm: - entropy_coef: 0.005 - desired_kl: 0.01 -env: - sampling_mode: adaptive - truncate_on_clip_end: false - sim_dt: 0.005 - control_config: - action_scale: - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.07450087032950714 - - 0.07450087032950714 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.07450087032950714 - - 0.07450087032950714 - anchor_pos_z_threshold: 0.3 - ee_body_pos_z_threshold: 0.3 - terminate_on_undesired_contacts: true - noise_config: - level: 0.0 -reward: - scales: - motion_global_root_pos: 0.5 - motion_global_root_ori: 0.5 - motion_body_pos: 2.0 - motion_body_ori: 1.5 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_ee_body_pos_z: 2.0 - motion_joint_pos: 0.5 - motion_joint_vel: 0.25 - action_rate_l2: -0.005 - joint_limit: -10.0 - undesired_contacts: -0.1 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 -play_profile: - enabled: true - env: - render_spacing: 2.0 diff --git a/conf/ppo/task/g1_climb_tracking/mujoco.yaml b/conf/ppo/task/g1_climb_tracking/mujoco.yaml index b7111e85a..fd7826eb4 100644 --- a/conf/ppo/task/g1_climb_tracking/mujoco.yaml +++ b/conf/ppo/task/g1_climb_tracking/mujoco.yaml @@ -1,83 +1,95 @@ # @package _global_ +defaults: + - /task/g1_motion_tracking/mujoco + - _self_ + training: task_name: G1ClimbTracking sim_backend: mujoco play_steps: 1000 + algo: num_envs: 1024 max_iterations: 20000 save_interval: 500 empirical_normalization: true obs_groups: - actor: - - actor - critic: - - critic + actor: [actor] + critic: [critic] algorithm: entropy_coef: 0.005 desired_kl: 0.01 + env: - sampling_mode: adaptive - truncate_on_clip_end: false + scene: + model_file: src/unilab/assets/robots/g1/scene_climb_20_z_scale_1.xml sim_dt: 0.005 - control_config: - action_scale: - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.07450087032950714 - - 0.07450087032950714 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.07450087032950714 - - 0.07450087032950714 - anchor_pos_z_threshold: 0.3 - ee_body_pos_z_threshold: 0.3 - terminate_on_undesired_contacts: true - noise_config: - level: 0.0 + max_episode_seconds: 15.0 + actions: + joint_pos: + scale: + ".*_(hip_pitch|hip_yaw)_joint": 0.5475464629911068 + ".*_(hip_roll|knee)_joint": 0.35066146637882434 + ".*_ankle_(pitch|roll)_joint": 0.43857731392336724 + "waist_yaw_joint": 0.5475464629911068 + "waist_(roll|pitch)_joint": 0.43857731392336724 + ".*_(shoulder_(pitch|roll|yaw)|elbow|wrist_roll)_joint": 0.43857731392336724 + ".*_wrist_(pitch|yaw)_joint": 0.07450087032950714 + commands: + motion: + params: + motion_file: motions/g1/climb_20_z_scale_1.0.npz + sampling_mode: adaptive + truncate_on_clip_end: false + terminations: + anchor_pos: + params: {command_name: motion, threshold: 0.3} + ee_body_pos: + params: + command_name: motion + threshold: 0.3 + body_names: &ee_bodies + - left_ankle_roll_link + - right_ankle_roll_link + - left_wrist_yaw_link + - right_wrist_yaw_link + undesired_contacts: + func: unilab.tasks.motion_tracking.common.manager_terms.bad_undesired_body_contacts + params: + command_name: motion + threshold: 0.05 + body_names: &undesired_bodies + - pelvis + - left_hip_roll_link + - left_knee_link + - right_hip_roll_link + - right_knee_link + - torso_link + - left_shoulder_roll_link + - left_elbow_link + - right_shoulder_roll_link + - right_elbow_link + reward: - scales: - motion_global_root_pos: 0.5 - motion_global_root_ori: 0.5 - motion_body_pos: 2.0 - motion_body_ori: 1.5 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_ee_body_pos_z: 2.0 - motion_joint_pos: 0.5 - motion_joint_vel: 0.25 - action_rate_l2: -0.005 - joint_limit: -10.0 - undesired_contacts: -0.1 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 + motion_body_pos: + weight: 2.0 + motion_body_ori: + weight: 1.5 + motion_ee_body_pos_z: + func: unilab.tasks.motion_tracking.common.manager_terms.motion_relative_body_position_z_error_exp + weight: 2.0 + params: {command_name: motion, std: 0.3, body_names: *ee_bodies} + motion_joint_pos: + weight: 0.5 + motion_joint_vel: + weight: 0.25 + action_rate_l2: + weight: -0.005 + undesired_contacts: + func: unilab.tasks.motion_tracking.common.manager_terms.undesired_body_contacts + weight: -0.1 + params: {command_name: motion, threshold: 0.05, body_names: *undesired_bodies} + play_profile: enabled: true env: diff --git a/conf/ppo/task/g1_flip_tracking/motrix.yaml b/conf/ppo/task/g1_flip_tracking/motrix.yaml index 3435bd57a..5948e0166 100644 --- a/conf/ppo/task/g1_flip_tracking/motrix.yaml +++ b/conf/ppo/task/g1_flip_tracking/motrix.yaml @@ -1,38 +1,29 @@ # @package _global_ +defaults: + - /task/g1_flip_tracking/mujoco + - _self_ + training: task_name: G1FlipTracking sim_backend: motrix - play_steps: 1000 + algo: - num_envs: 1024 max_iterations: 30000 - save_interval: 500 - obs_groups: - actor: - - actor - algorithm: - entropy_coef: 0.005 + empirical_normalization: false + +env: + actions: + joint_pos: + scale: 0.25 + reward: - scales: - motion_global_root_pos: 1.0 - motion_global_root_ori: 0.5 - motion_body_pos: 1.0 - motion_body_ori: 1.0 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_joint_pos: 0.0 - motion_joint_vel: 0.0 - action_rate_l2: -0.05 - joint_limit: -10.0 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 -play_profile: - enabled: true - env: - render_spacing: 2.0 + motion_global_root_pos: + weight: 1.0 + motion_body_pos: + weight: 1.0 + motion_body_ori: + weight: 1.0 + motion_ee_body_pos_z: null + action_rate_l2: + weight: -0.05 + undesired_contacts: null diff --git a/conf/ppo/task/g1_flip_tracking/mujoco.yaml b/conf/ppo/task/g1_flip_tracking/mujoco.yaml index 1edd403da..eea70d92f 100644 --- a/conf/ppo/task/g1_flip_tracking/mujoco.yaml +++ b/conf/ppo/task/g1_flip_tracking/mujoco.yaml @@ -1,83 +1,106 @@ # @package _global_ +defaults: + - /task/g1_motion_tracking/mujoco + - _self_ + training: task_name: G1FlipTracking sim_backend: mujoco play_steps: 1000 + algo: num_envs: 1024 max_iterations: 20000 save_interval: 500 empirical_normalization: true obs_groups: - actor: - - actor - critic: - - critic + actor: [actor] + critic: [critic] algorithm: entropy_coef: 0.005 desired_kl: 0.01 + env: - sampling_mode: start - truncate_on_clip_end: false + scene: + model_file: src/unilab/assets/robots/g1/scene_flat.xml sim_dt: 0.005 - control_config: - action_scale: - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.07450087032950714 - - 0.07450087032950714 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.07450087032950714 - - 0.07450087032950714 - anchor_pos_z_threshold: 0.5 - ee_body_pos_z_threshold: 0.5 - terminate_on_undesired_contacts: true - noise_config: - level: 0.0 + actions: + joint_pos: + scale: + ".*_(hip_pitch|hip_yaw)_joint": 0.5475464629911068 + ".*_(hip_roll|knee)_joint": 0.35066146637882434 + ".*_ankle_(pitch|roll)_joint": 0.43857731392336724 + "waist_yaw_joint": 0.5475464629911068 + "waist_(roll|pitch)_joint": 0.43857731392336724 + ".*_(shoulder_(pitch|roll|yaw)|elbow|wrist_roll)_joint": 0.43857731392336724 + ".*_wrist_(pitch|yaw)_joint": 0.07450087032950714 + commands: + motion: + params: + motion_file: motions/g1/flip_360_001__A304.npz + sampling_mode: start + truncate_on_clip_end: false + pose_range: &zero_pose + x: [0.0, 0.0] + y: [0.0, 0.0] + z: [0.0, 0.0] + roll: [0.0, 0.0] + pitch: [0.0, 0.0] + yaw: [0.0, 0.0] + velocity_range: *zero_pose + joint_position_range: [0.0, 0.0] + terminations: + anchor_pos: + params: {command_name: motion, threshold: 0.5} + anchor_ori: + params: + command_name: motion + threshold: 1.0e9 + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + ee_body_pos: + params: + command_name: motion + threshold: 0.5 + body_names: &ee_bodies + - left_ankle_roll_link + - right_ankle_roll_link + - left_wrist_yaw_link + - right_wrist_yaw_link + undesired_contacts: + func: unilab.tasks.motion_tracking.common.manager_terms.bad_undesired_body_contacts + params: + command_name: motion + threshold: 0.05 + body_names: &undesired_bodies + - pelvis + - left_hip_roll_link + - left_knee_link + - right_hip_roll_link + - right_knee_link + - torso_link + - left_shoulder_roll_link + - left_elbow_link + - right_shoulder_roll_link + - right_elbow_link + reward: - scales: - motion_global_root_pos: 0.5 - motion_global_root_ori: 0.5 - motion_body_pos: 2.0 - motion_body_ori: 1.5 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_ee_body_pos_z: 2.0 - motion_joint_pos: 0.0 - motion_joint_vel: 0.0 - action_rate_l2: -0.005 - joint_limit: -10.0 - undesired_contacts: -0.1 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 + motion_body_pos: + weight: 2.0 + motion_body_ori: + weight: 1.5 + motion_ee_body_pos_z: + func: unilab.tasks.motion_tracking.common.manager_terms.motion_relative_body_position_z_error_exp + weight: 2.0 + params: {command_name: motion, std: 0.3, body_names: *ee_bodies} + action_rate_l2: + weight: -0.005 + undesired_contacts: + func: unilab.tasks.motion_tracking.common.manager_terms.undesired_body_contacts + weight: -0.1 + params: {command_name: motion, threshold: 0.05, body_names: *undesired_bodies} + play_profile: enabled: true env: diff --git a/conf/ppo/task/g1_wall_flip_tracking/motrix.yaml b/conf/ppo/task/g1_wall_flip_tracking/motrix.yaml index c595203a6..8fd67bfc7 100644 --- a/conf/ppo/task/g1_wall_flip_tracking/motrix.yaml +++ b/conf/ppo/task/g1_wall_flip_tracking/motrix.yaml @@ -1,64 +1,20 @@ # @package _global_ +defaults: + - /task/g1_wall_flip_tracking/mujoco + - _self_ + training: task_name: G1WallFlipTracking sim_backend: motrix play_env_num: 16 - play_steps: 1000 render_spacing: 3.0 + algo: - num_envs: 1024 max_iterations: 12000 - save_interval: 500 - empirical_normalization: true - obs_groups: - actor: - - actor - critic: - - critic - algorithm: - entropy_coef: 0.005 - desired_kl: 0.01 + env: motrix_max_iterations: 3 - sampling_mode: start - truncate_on_clip_end: false - sim_dt: 0.005 - control_config: - action_scale: - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.07450087032950714 - - 0.07450087032950714 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.07450087032950714 - - 0.07450087032950714 - anchor_pos_z_threshold: 0.5 - ee_body_pos_z_threshold: 0.5 - terminate_on_undesired_contacts: true - noise_config: - level: 0.0 + play_profile: enabled: true env: @@ -70,25 +26,3 @@ play_profile: skybox_rgb1: [0.90, 0.90, 0.91] skybox_rgb2: [0.68, 0.68, 0.70] ground_texrepeat: [0.25, 0.25] -reward: - scales: - motion_global_root_pos: 0.5 - motion_global_root_ori: 0.5 - motion_body_pos: 2.0 - motion_body_ori: 1.5 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_ee_body_pos_z: 2.0 - motion_joint_pos: 0.5 - motion_joint_vel: 0.25 - action_rate_l2: -0.005 - joint_limit: -10.0 - undesired_contacts: -0.1 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 diff --git a/conf/ppo/task/g1_wall_flip_tracking/mujoco.yaml b/conf/ppo/task/g1_wall_flip_tracking/mujoco.yaml index e11d9b05b..082c23fd9 100644 --- a/conf/ppo/task/g1_wall_flip_tracking/mujoco.yaml +++ b/conf/ppo/task/g1_wall_flip_tracking/mujoco.yaml @@ -1,84 +1,22 @@ # @package _global_ +defaults: + - /task/g1_flip_tracking/mujoco + - _self_ + training: task_name: G1WallFlipTracking sim_backend: mujoco - play_steps: 1000 -algo: - num_envs: 1024 - max_iterations: 20000 - save_interval: 500 - empirical_normalization: true - obs_groups: - actor: - - actor - critic: - - critic - algorithm: - entropy_coef: 0.005 - desired_kl: 0.01 + env: - sampling_mode: start - truncate_on_clip_end: false - sim_dt: 0.005 - control_config: - action_scale: - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.5475464629911068 - - 0.35066146637882434 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.5475464629911068 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.07450087032950714 - - 0.07450087032950714 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.43857731392336724 - - 0.07450087032950714 - - 0.07450087032950714 - anchor_pos_z_threshold: 0.5 - ee_body_pos_z_threshold: 0.5 - terminate_on_undesired_contacts: true - noise_config: - level: 0.0 + scene: + model_file: src/unilab/assets/robots/g1/scene_flat_with_wall.xml + commands: + motion: + params: + motion_file: motions/g1/flip_from_wall_104__A304.npz + reward: - scales: - motion_global_root_pos: 0.5 - motion_global_root_ori: 0.5 - motion_body_pos: 2.0 - motion_body_ori: 1.5 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_ee_body_pos_z: 2.0 - motion_joint_pos: 0.5 - motion_joint_vel: 0.25 - action_rate_l2: -0.005 - joint_limit: -10.0 - undesired_contacts: -0.1 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 -play_profile: - enabled: true - env: - render_spacing: 2.0 + motion_joint_pos: + weight: 0.5 + motion_joint_vel: + weight: 0.25 diff --git a/conf/ppo/task/x2_wall_flip_tracking/motrix.yaml b/conf/ppo/task/x2_wall_flip_tracking/motrix.yaml index d2a45c89e..9b0aecd0a 100644 --- a/conf/ppo/task/x2_wall_flip_tracking/motrix.yaml +++ b/conf/ppo/task/x2_wall_flip_tracking/motrix.yaml @@ -1,70 +1,15 @@ # @package _global_ +defaults: + - /task/x2_wall_flip_tracking/mujoco + - _self_ + training: task_name: X2WallFlipTracking sim_backend: motrix - play_steps: 300 # 6s play video @ ctrl_dt=0.02 (fps 50) - play_env_num: 16 - render_spacing: 3.0 - cam_distance: 14.0 - cam_azimuth: 225.0 - cam_elevation: -18.0 - cam_lookat: [4.5, 4.5, 1.0] -interactive: - action_mode: policy -algo: - num_envs: 1024 - max_iterations: 9500 - save_interval: 500 - empirical_normalization: true - obs_groups: - actor: - - actor - critic: - - critic - algorithm: - entropy_coef: 0.005 - desired_kl: 0.01 + env: motrix_max_iterations: 3 - sampling_mode: start - truncate_on_clip_end: false - sim_dt: 0.005 - control_config: - action_scale: - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - anchor_pos_z_threshold: 0.5 - ee_body_pos_z_threshold: 0.5 - terminate_on_undesired_contacts: true - noise_config: - level: 0.0 + play_profile: enabled: true env: @@ -76,25 +21,3 @@ play_profile: skybox_rgb1: [0.90, 0.90, 0.91] skybox_rgb2: [0.68, 0.68, 0.70] ground_texrepeat: [0.25, 0.25] -reward: - scales: - motion_global_root_pos: 0.5 - motion_global_root_ori: 0.5 - motion_body_pos: 2.0 - motion_body_ori: 1.5 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_ee_body_pos_z: 2.0 - motion_joint_pos: 0.5 - motion_joint_vel: 0.25 - action_rate_l2: -0.005 - joint_limit: -10.0 - undesired_contacts: -0.1 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 diff --git a/conf/ppo/task/x2_wall_flip_tracking/mujoco.yaml b/conf/ppo/task/x2_wall_flip_tracking/mujoco.yaml index 83511b770..50922321b 100644 --- a/conf/ppo/task/x2_wall_flip_tracking/mujoco.yaml +++ b/conf/ppo/task/x2_wall_flip_tracking/mujoco.yaml @@ -1,94 +1,208 @@ # @package _global_ +defaults: + - /task/g1_motion_tracking/mujoco + - _self_ + training: task_name: X2WallFlipTracking sim_backend: mujoco - play_steps: 300 # 6s play video @ ctrl_dt=0.02 (fps 50) - # Offline-render (play) only — does not affect the training loop. 16 envs are - # laid out on a 4x4 grid; render_spacing must exceed the per-env wall reach - # (~2.14m toward -Y) so neighbouring cells don't overlap. The oblique - # azimuth (225) views the grid corner-on with each robot in front of its - # wall (azimuth 90 would put the walls between camera and robots); lookat is - # the grid centre (offsets span 0..9m in X/Y at spacing 3.0). + play_steps: 300 play_env_num: 16 render_spacing: 3.0 cam_distance: 14.0 cam_azimuth: 225.0 cam_elevation: -18.0 cam_lookat: [4.5, 4.5, 1.0] + interactive: action_mode: policy + algo: num_envs: 1024 max_iterations: 9500 save_interval: 500 empirical_normalization: true obs_groups: - actor: - - actor - critic: - - critic + actor: [actor] + critic: [critic] algorithm: entropy_coef: 0.005 desired_kl: 0.01 + +play_profile: + enabled: false + env: null + env: - sampling_mode: start - truncate_on_clip_end: false + scene: + model_file: src/unilab/assets/robots/x2/scene_flat_with_wall.xml + visual_model_file: src/unilab/assets/robots/x2/scene_flat_with_wall_visual.xml + default_keyframe_name: home + entities: + robot: + root_body_name: pelvis + joint_names: &x2_joints + - left_hip_pitch_joint + - left_hip_roll_joint + - left_hip_yaw_joint + - left_knee_joint + - left_ankle_pitch_joint + - left_ankle_roll_joint + - right_hip_pitch_joint + - right_hip_roll_joint + - right_hip_yaw_joint + - right_knee_joint + - right_ankle_pitch_joint + - right_ankle_roll_joint + - waist_yaw_joint + - waist_pitch_joint + - waist_roll_joint + - left_shoulder_pitch_joint + - left_shoulder_roll_joint + - left_shoulder_yaw_joint + - left_elbow_joint + - left_wrist_yaw_joint + - left_wrist_pitch_joint + - left_wrist_roll_joint + - right_shoulder_pitch_joint + - right_shoulder_roll_joint + - right_shoulder_yaw_joint + - right_elbow_joint + - right_wrist_yaw_joint + - right_wrist_pitch_joint + - right_wrist_roll_joint + actuator_names: *x2_joints + geom_names: null + body_names: &tracked_bodies + - pelvis + - left_hip_pitch_link + - left_hip_roll_link + - left_hip_yaw_link + - left_knee_link + - left_ankle_pitch_link + - left_ankle_roll_link + - right_hip_pitch_link + - right_hip_roll_link + - right_hip_yaw_link + - right_knee_link + - right_ankle_pitch_link + - right_ankle_roll_link + - waist_yaw_link + - waist_pitch_link + - torso_link + - left_shoulder_pitch_link + - left_shoulder_roll_link + - left_shoulder_yaw_link + - left_elbow_link + - left_wrist_yaw_link + - left_wrist_pitch_link + - left_wrist_roll_link + - right_shoulder_pitch_link + - right_shoulder_roll_link + - right_shoulder_yaw_link + - right_elbow_link + - right_wrist_yaw_link + - right_wrist_pitch_link + - right_wrist_roll_link sim_dt: 0.005 - control_config: - action_scale: - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - - 0.25 - anchor_pos_z_threshold: 0.5 - ee_body_pos_z_threshold: 0.5 - terminate_on_undesired_contacts: true - noise_config: - level: 0.0 + observations: + actor: + terms: + motion_anchor_pos_b: null + base_lin_vel: null + base_ang_vel: + params: {sensor_name: body-angular-velocity} + critic: + terms: + base_lin_vel: + params: {sensor_name: body-linear-vel} + base_ang_vel: + params: {sensor_name: body-angular-velocity} + commands: + motion: + params: + motion_file: motions/x2/tictacflip_6-3_g1format.npz + anchor_body_name: torso_link + body_names: *tracked_bodies + sampling_mode: start + truncate_on_clip_end: false + pose_range: &zero_pose + x: [0.0, 0.0] + y: [0.0, 0.0] + z: [0.0, 0.0] + roll: [0.0, 0.0] + pitch: [0.0, 0.0] + yaw: [0.0, 0.0] + velocity_range: *zero_pose + joint_position_range: [0.0, 0.0] + terminations: + anchor_pos: + params: {command_name: motion, threshold: 0.5} + anchor_ori: + params: + command_name: motion + threshold: 1.0e9 + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + ee_body_pos: + params: + command_name: motion + threshold: 0.5 + body_names: &ee_bodies + - left_ankle_roll_link + - right_ankle_roll_link + - left_wrist_yaw_link + - right_wrist_yaw_link + undesired_contacts: + func: unilab.tasks.motion_tracking.common.manager_terms.bad_undesired_body_contacts + params: + command_name: motion + threshold: 0.05 + body_names: &undesired_bodies + - pelvis + - left_hip_pitch_link + - left_hip_roll_link + - left_hip_yaw_link + - left_knee_link + - left_ankle_pitch_link + - right_hip_pitch_link + - right_hip_roll_link + - right_hip_yaw_link + - right_knee_link + - right_ankle_pitch_link + - waist_yaw_link + - waist_pitch_link + - torso_link + - left_shoulder_pitch_link + - left_shoulder_roll_link + - left_shoulder_yaw_link + - left_elbow_link + - left_wrist_pitch_link + - left_wrist_roll_link + - right_shoulder_pitch_link + - right_shoulder_roll_link + - right_shoulder_yaw_link + - right_elbow_link + - right_wrist_pitch_link + - right_wrist_roll_link + reward: - scales: - motion_global_root_pos: 0.5 - motion_global_root_ori: 0.5 - motion_body_pos: 2.0 - motion_body_ori: 1.5 - motion_body_lin_vel: 1.0 - motion_body_ang_vel: 1.0 - motion_ee_body_pos_z: 2.0 - motion_joint_pos: 0.5 - motion_joint_vel: 0.25 - action_rate_l2: -0.005 - joint_limit: -10.0 - undesired_contacts: -0.1 - std_root_pos: 0.3 - std_root_ori: 0.4 - std_body_pos: 0.3 - std_body_ori: 0.4 - std_body_lin_vel: 1.0 - std_body_ang_vel: 3.14 - std_joint_pos: 0.2 - std_joint_vel: 1.0 + motion_body_pos: + weight: 2.0 + motion_body_ori: + weight: 1.5 + motion_ee_body_pos_z: + func: unilab.tasks.motion_tracking.common.manager_terms.motion_relative_body_position_z_error_exp + weight: 2.0 + params: {command_name: motion, std: 0.3, body_names: *ee_bodies} + motion_joint_pos: + weight: 0.5 + motion_joint_vel: + weight: 0.25 + action_rate_l2: + weight: -0.005 + undesired_contacts: + func: unilab.tasks.motion_tracking.common.manager_terms.undesired_body_contacts + weight: -0.1 + params: {command_name: motion, threshold: 0.05, body_names: *undesired_bodies} diff --git a/docs/sphinx/source/en/2-user_guide/4-tasks/2-motion_tracking.md b/docs/sphinx/source/en/2-user_guide/4-tasks/2-motion_tracking.md index 3245e8c45..150be56ff 100644 --- a/docs/sphinx/source/en/2-user_guide/4-tasks/2-motion_tracking.md +++ b/docs/sphinx/source/en/2-user_guide/4-tasks/2-motion_tracking.md @@ -13,28 +13,31 @@ off-policy paths. ## Task Owners -Each task ships a default motion clip defined in the env config dataclass: +Each task ships a default motion clip in its Hydra task-owner YAML. Hydra is the +configuration entry point; the selected owner is materialized into the shared +`ManagerBasedRlEnvCfg` and then consumed by the NumPy Manager-Based runtime. | CLI Task | Registered Env | Default Motion | Owner Evidence | | --- | --- | --- | --- | | `g1_motion_tracking` | `G1MotionTracking` | `dance1_subject2_part.npz` | `conf/ppo/task/g1_motion_tracking/`, `conf/appo/task/g1_motion_tracking/` | | `g1_flip_tracking` | `G1FlipTracking` | `flip_360_001__A304.npz` | `conf/ppo/task/g1_flip_tracking/`, `conf/appo/task/g1_flip_tracking/` | | `g1_wall_flip_tracking` | `G1WallFlipTracking` | `flip_from_wall_104__A304.npz` | `conf/ppo/task/g1_wall_flip_tracking/`, `conf/appo/task/g1_wall_flip_tracking/` | -| `x2_wall_flip_tracking` | `X2WallFlipTracking` | `tictacflip_6-3_g1format.npz` | `conf/ppo/task/x2_wall_flip_tracking/` (MuJoCo only) | -| `g1_climb_tracking` | G1 climb tracking env | clip from env config | `conf/ppo/task/g1_climb_tracking/`, `conf/appo/task/g1_climb_tracking/` | -| `g1_box_tracking` | G1 box tracking env | clip from env config | `conf/ppo/task/g1_box_tracking/` | -| `g1_wbt_obs` | `G1MotionTrackingSAC` | shared with `g1_motion_tracking` | `conf/offpolicy/task/sac/g1_wbt_obs/mujoco.yaml` | +| `x2_wall_flip_tracking` | `X2WallFlipTracking` | `tictacflip_6-3_g1format.npz` | `conf/ppo/task/x2_wall_flip_tracking/` | +| `g1_climb_tracking` | `G1ClimbTracking` | `climb_20_z_scale_1.0.npz` | `conf/ppo/task/g1_climb_tracking/`, `conf/appo/task/g1_climb_tracking/` | +| `g1_box_tracking` | `G1BoxTracking` | `sub3_largebox_003_boxconverted.npz` | `conf/ppo/task/g1_box_tracking/` | +| `g1_wbt_obs` | `G1WBTObs` | `dance1_subject2_part.npz` | `conf/offpolicy/task/sac/g1_wbt_obs/mujoco.yaml` | -The defaults are set in code: `dance1_subject2_part.npz` -(`g1/tracking.py`), `flip_360_001__A304.npz` and `flip_from_wall_104__A304.npz` -(`g1/flip_tracking.py`), and `tictacflip_6-3_g1format.npz` (`x2/flip_tracking.py`). +The 23-DoF task-owner directories select their matching 23-DoF scene, motion, +entity, and action declarations. Profile differences remain in Hydra. The G1 +identities use the shared manager factory; X2 adds only a cold-path mesh resolver +before delegating to that factory. ## PPO And APPO PPO owner iteration budgets (the `--sim mujoco` owner YAMLs): `g1_motion_tracking` runs `algo.max_iterations=15000`; `g1_flip_tracking` and `g1_wall_flip_tracking` -run `20000`; the MuJoCo-only `x2_wall_flip_tracking` runs `9500`. (The Motrix -owner YAML for `g1_flip_tracking` raises this to `30000`.) +run `20000`; `x2_wall_flip_tracking` runs `9500`. (The Motrix owner YAML for +`g1_flip_tracking` raises this to `30000`.) ```bash uv run train --algo ppo --task g1_motion_tracking --sim mujoco @@ -57,11 +60,13 @@ uv run train --algo sac --task g1_motion_tracking --sim mujoco training.use_amp= uv run train --algo sac --task g1_wbt_obs --sim mujoco training.use_amp=true ``` -The `g1_wbt_obs` owner is the deploy-aligned off-policy observation profile: a -pelvis IMU state (`pelvis_local_linvel` / `pelvis_gyro` / `pelvis_upvector`) plus -per-term observation history (`noise_config.obs_history_length: 5`), byte-aligned -with the deploy-time `ObservationManager`. Deploy tooling lives under -`scripts/deploy/`, and the observation alignment is cross-checked by +The `g1_wbt_obs` owner is the deploy-aligned off-policy observation profile. Its +actor keeps the command and anchor-orientation terms at one step while the +`base_ang_vel`, `joint_pos`, `joint_vel`, and `actions` terms declare +`history_length: 5`. `ObservationManager` owns and flattens those per-term +histories; the actor uses the configured encoder-biased joint-position term while +the critic keeps the clean term. Deploy tooling lives under `scripts/deploy/`, and +the observation alignment is cross-checked by `tests/scripts/test_obs_alignment_g1_wbt.py`. When a Motrix sim2sim replay needs a checkpoint from another log root, pass the absolute path through `uv run eval`: @@ -72,16 +77,20 @@ uv run eval --algo sac --task g1_motion_tracking --sim motrix \ ## Motion Files -Motion NPZ files are read through `env.motion_file`, which also accepts a list of +Motion NPZ files are selected through +`env.commands.motion.params.motion_file`, which accepts one path or a list of paths. A standard clip must contain the seven keys `fps`, `joint_pos`, `joint_vel`, `body_pos_w`, `body_quat_w`, `body_lin_vel_w`, and `body_ang_vel_w` (validated in `common/motion_loader.py`): ```yaml env: - motion_file: - - src/unilab/assets/motions/g1/dance1_subject2_part.npz - - src/unilab/assets/motions/g1/walk1_subject5_from_csv.npz + commands: + motion: + params: + motion_file: + - motions/g1/dance1_subject2_part.npz + - motions/g1/walk1_subject5_from_csv.npz ``` Conversion and inspection helpers are in `scripts/motion/`: @@ -115,21 +124,22 @@ randomization so the precise clip start state is reused: ```bash CUDA_VISIBLE_DEVICES=1 uv run train --algo sac --task g1_motion_tracking --sim mujoco \ training.use_amp=true algo.seed=1 \ - +env.motion_file=src/unilab/assets/motions/g1/motion_crawl_slope_uni.npz \ - +env.scene.model_file=src/unilab/assets/robots/g1/scene_crawl_slope.xml \ - +env.sampling_mode=start \ - env.truncate_on_clip_end=true \ - +env.max_episode_seconds=20.0 \ - '+env.pose_randomization={x:[0,0],y:[0,0],z:[0,0],roll:[0,0],pitch:[0,0],yaw:[0,0]}' \ - '+env.velocity_randomization={x:[0,0],y:[0,0],z:[0,0],roll:[0,0],pitch:[0,0],yaw:[0,0]}' \ - '+env.joint_position_range=[0,0]' + env.commands.motion.params.motion_file=motions/g1/motion_crawl_slope_uni.npz \ + env.scene.model_file=src/unilab/assets/robots/g1/scene_crawl_slope.xml \ + env.commands.motion.params.sampling_mode=start \ + env.commands.motion.params.truncate_on_clip_end=true \ + env.max_episode_seconds=20.0 \ + 'env.commands.motion.params.pose_range={x:[0,0],y:[0,0],z:[0,0],roll:[0,0],pitch:[0,0],yaw:[0,0]}' \ + 'env.commands.motion.params.velocity_range={x:[0,0],y:[0,0],z:[0,0],roll:[0,0],pitch:[0,0],yaw:[0,0]}' \ + 'env.commands.motion.params.joint_position_range=[0,0]' ``` -Key overrides: `env.motion_file` selects the crawl-slope clip; -`env.scene.model_file` switches to the slope scene (`scene_crawl_slope.xml` exists -under `src/unilab/assets/robots/g1/`); `sampling_mode=start` plus -`truncate_on_clip_end=true` starts from the clip beginning and truncates there; and -zeroing the randomization ranges reuses the exact clip initial state. +Key overrides: `env.commands.motion.params.motion_file` selects the crawl-slope +clip; `env.scene.model_file` switches to the slope scene +(`scene_crawl_slope.xml` exists under `src/unilab/assets/robots/g1/`); +`sampling_mode=start` plus `truncate_on_clip_end=true` starts from the clip +beginning and truncates there; and zeroing the command reset ranges reuses the +exact clip initial state. ## Interactive Debugging diff --git a/docs/sphinx/source/en/2-user_guide/5-domain_randomization/0-index.md b/docs/sphinx/source/en/2-user_guide/5-domain_randomization/0-index.md index c76b3a2fb..fe98d3b9b 100644 --- a/docs/sphinx/source/en/2-user_guide/5-domain_randomization/0-index.md +++ b/docs/sphinx/source/en/2-user_guide/5-domain_randomization/0-index.md @@ -18,7 +18,7 @@ These three paths correspond to three lifecycle classes: ## Status Conclusions 1. All tasks currently wired to a DR provider use the unified DR entry point; no task bypasses `DomainRandomizationManager` to run a separate DR flow inside `reset()`. -2. They are all roughly structured: task files define a `domain_rand` config dataclass, a `DomainRandomizationProvider`, and a `ResetPlan`; `G1Walk*` instead declares DR through Hydra `EventTermCfg` Manager-Based reset terms. +2. They are all roughly structured: legacy task owners define a `domain_rand` config dataclass, a `DomainRandomizationProvider`, and a `ResetPlan`; Manager-Based owners declare reset behavior through Hydra command/event terms. G1 motion reset perturbations belong to `MotionCommandCfg`, while WBT adds `EventTermCfg` reset and interval terms. 3. What is "unified" today is mainly the entry point and execution flow, not every randomization item itself. The shared helper `build_common_reset_randomization()` currently generates `base_mass_delta`, `base_com_offset`, `gravity`, `kp`, `kd`; the shared interval helper currently only generates push. 4. `ResetRandomizationPayload` can already express `gravity`, `body_iquat`, `body_inertia`, `kp`, `kd`, and `MuJoCoBackend` has declared support. Whether these are actually used still depends on whether the task provider samples and dispatches them. 5. `MotrixBackend` currently supports `base_mass_delta`, `base_com_offset`, `kp`, `kd`, and interval push; and it requires all model actuators to be position actuators during initialization. @@ -32,7 +32,8 @@ These three paths correspond to three lifecycle classes: | `Go2JoystickFlat` | Yes | Yes: `Domain_Rand + Provider + ResetPlan` | task state sampling + common payload | push | `go2/joystick.py` | | `G1WalkFlat` | Yes | Yes: Hydra `EventTermCfg` + Manager-Based reset terms | root-state reset + kp/kd via `pd_gains` | none | `g1/manager_terms.py` | | `G1WalkRough` | Yes | Yes: same Manager-Based event terms as `G1WalkFlat` | root-state reset + kp/kd via `pd_gains` | none | `g1/manager_terms.py` | -| `G1MotionTracking` | Yes | Yes: `Domain_Rand + Provider + ResetPlan` | extensive task-specific reset sampling + common payload | push | `motion_tracking/g1/tracking.py` | +| `G1MotionTracking` | Yes | Yes: Hydra `MotionCommandCfg` + Manager-Based command reset | motion frame, root pose/velocity, and joint-position sampling | none | `motion_tracking/common/manager_terms.py` | +| `G1WBTObs` | Yes | Yes: same motion command + Hydra `EventTermCfg` | motion reset plus mass/COM/PD/friction/encoder-bias events | interval velocity kick | `motion_tracking/g1/manager_terms.py` | | `AllegroInhandRotation` | Yes | Yes: Hydra `EventTermCfg` + Manager-Based reset term | entity-scoped hand/ball reset | none | `allegro_inhand/manager_terms.py` | | `AllegroInhandRotationGrasp` | Yes | Yes: reuses the rotation reset event + `RecorderTermCfg` | noisy hand reset + grasp collection | none | `allegro_inhand/grasp_gen.py` | | `SharpaInhandRotation` | Yes | Yes: `InitRandomizationPlan + ResetPlan + IntervalRandomizationPlan` | grasp cache sampling + common payload | object `body_force` | `sharpa_inhand/rotation.py` | @@ -46,7 +47,8 @@ These three paths correspond to three lifecycle classes: | `Go2JoystickFlat` | base xy; base yaw; base qvel; command sampling; `current_actions/last_actions` zeroed; kp/kd randomization (enabled by default); optional `base_mass_delta`; optional `base_com_offset`; optional `gravity` | `push_robots` | kp/kd enabled by default; common payload and push disabled by default | | `G1WalkFlat` | base xy/yaw and base qvel via `reset_root_state_uniform`; command sampling with a planar dead zone; `gait_phase` sampling; kp/kd randomization via `pd_gains` | none | kp/kd enabled on mujoco owners by default; disabled on motrix/mjwarp owners | | `G1WalkRough` | Same as `G1WalkFlat` (shared owner bases, rough scene) | none | Same defaults as `G1WalkFlat` | -| `G1MotionTracking` | motion frame sampling; root pose perturbation `x/y/z/roll/pitch/yaw`; root velocity perturbation `x/y/z/roll/pitch/yaw`; joint position noise; under MuJoCo clipped by joint range; `current_actions/last_actions` zeroed; optional `base_mass_delta`; optional `base_com_offset`; optional `gravity` | `push_robots` | `pose_randomization`, `velocity_randomization`, `joint_position_range` have non-zero perturbations by default; common payload and push disabled by default | +| `G1MotionTracking` | Motion-command frame sampling; root pose perturbation `x/y/z/roll/pitch/yaw`; root velocity perturbation `x/y/z/roll/pitch/yaw`; joint-position noise clipped through the public entity soft limits; action-manager state reset | none | `pose_range`, `velocity_range`, and `joint_position_range` have non-zero perturbations in the base owner | +| `G1WBTObs` | Same motion reset plus base mass, base COM, PD gain, foot friction, and encoder-bias event terms | `push_by_setting_velocity` | The WBT owner explicitly enables all listed event terms; unsupported capabilities raise rather than fall back | | `AllegroInhandRotation` | Entity-scoped hand/ball reset; an explicitly configured grasp cache is sampled, otherwise `null` explicitly selects the model home pose; optional `joint_noise`, `ball_velocity_noise`, and `ball_z_offset` | none | owner YAML explicitly selects the home pose and zero reset noise; a configured missing or malformed cache fails closed | | `AllegroInhandRotationGrasp` | Reuses the rotation reset with `joint_noise=0.25`; Manager-Based termination checks fingertip distance, contact count, and ball height; recorder stores successful timeout rows | none | generates the 50k-row Allegro grasp cache and raises `RunComplete` after a successful save | | `SharpaInhandRotation` | grasp cache bucketed sampling by `scale_ids`; object pose / quat reset; optional common reset randomization payload (incl. `gravity`) | object `body_force` direct force disturbance | `domain_rand.scale_list` defaults come from the owner YAML; under MuJoCo, object geom scale is materialized during init; common payload disabled by default; object force enabled by default via the Sharpa owner YAML | @@ -74,7 +76,7 @@ So from an execution-path perspective, the tasks are already unified. This means: - Although locomotion tasks all go through the unified entry point, their base xy, yaw, qvel, command, and gait phase are still sampled directly inside each provider -- `G1MotionTracking`'s pose / velocity / joint noise is also task-specific logic +- `G1MotionTracking`'s pose / velocity / joint noise is owned by its manager command - Allegro's grasp / object initial state sampling is entirely task-specific logic - Sharpa's `geom_size` scale is init-lifecycle model materialization and is not part of the reset common payload 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 4f277d09d..b63fe6bd9 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 @@ -27,10 +27,8 @@ The shared types live in `src/unilab/dr/types.py`, and the manager lives in Representative provider implementations are in: -- `src/unilab/tasks/locomotion/go1/joystick.py` -- `src/unilab/tasks/locomotion/g1/joystick.py` -- `src/unilab/tasks/motion_tracking/g1/tracking.py` -- `src/unilab/tasks/locomotion/go2/footstand.py` +- `src/unilab/tasks/locomotion/common/dr_provider.py` +- `src/unilab/tasks/locomotion/go2_arm/manip_loco.py` - `src/unilab/tasks/manipulation/sharpa_inhand/rotation.py` Developer contract details are in diff --git a/docs/sphinx/source/en/3-deployment/1-sim_to_real/1-overview.md b/docs/sphinx/source/en/3-deployment/1-sim_to_real/1-overview.md index 38a1bc8b9..b1d3d14d4 100644 --- a/docs/sphinx/source/en/3-deployment/1-sim_to_real/1-overview.md +++ b/docs/sphinx/source/en/3-deployment/1-sim_to_real/1-overview.md @@ -70,10 +70,10 @@ flowchart LR runtime (units, frame, filter cutoffs). Log the first deploy-side observation window and compare it with a sim rollout built from the same owner YAML. -- **Action latency.** Some task configs expose one-step delayed action - execution through `control_config.simulate_action_latency`. Measure the - deploy loop and make the training owner match that contract before a - hardware run. See {doc}`8-latency_budget`. +- **Action latency.** Some task owners expose one-step delayed action execution + through a control config or Manager-Based action term. Measure the deploy loop + and make the training owner match that contract before a hardware run. See + {doc}`8-latency_budget`. - **Friction / damping mismatch.** Especially for in-hand manipulation. Sweep friction in DR; cross-check via {doc}`../2-sim_to_sim/3-contact_and_friction_alignment`. - **Reset transients.** Sim resets to a stable pose; deployment starts from a diff --git a/docs/sphinx/source/en/3-deployment/1-sim_to_real/8-latency_budget.md b/docs/sphinx/source/en/3-deployment/1-sim_to_real/8-latency_budget.md index 7e3177b7f..6e546e0d4 100644 --- a/docs/sphinx/source/en/3-deployment/1-sim_to_real/8-latency_budget.md +++ b/docs/sphinx/source/en/3-deployment/1-sim_to_real/8-latency_budget.md @@ -8,21 +8,22 @@ budgets as robot-specific measurements, not UniLab defaults. | Surface | Repo evidence | What it covers | | --- | --- | --- | -| One-step action delay | `control_config.simulate_action_latency` in locomotion and G1 motion-tracking envs | Executes the previous action instead of the current action. | -| G1 WBT observation history | `noise_config.obs_history_length` and `scripts/deploy/export_deploy_config.py` | Exports per-term `obs_layout` history for `gyro`, `joint_pos_rel`, `dof_vel`, and `last_actions`. | +| One-step action delay | Manager action term `simulate_action_latency` declarations in task owners | Executes the previous action instead of the current action. | +| G1 WBT observation history | Per-term `history_length` in `conf/offpolicy/task/sac/g1_wbt_obs/mujoco.yaml` and `scripts/deploy/export_deploy_config.py` | Exports per-term `obs_layout` history for `gyro`, `joint_pos_rel`, `dof_vel`, and `last_actions`. | | Sharpa tactile contact latency | `domain_rand.contact_latency` in Sharpa in-hand configs | Keeps previous tactile contact values for sampled contact channels. | | Deploy-side ONNX contract check | `scripts/deploy/sim_prototype.py` | Validates `obs_layout`, `obs_dim`, ONNX input width, clipping, and EMA action smoothing for the G1 WBT path. | ## Action Latency -For tasks that expose `control_config.simulate_action_latency`, the env applies -`last_actions` when the flag is enabled. Keep this in the selected task owner -YAML instead of adding deploy-only behavior later. +For Manager-Based tasks that enable action latency, the action manager applies +the previous action when the flag is enabled. Keep this in the selected task +owner YAML instead of adding deploy-only behavior later. ```yaml env: - control_config: - simulate_action_latency: true + actions: + joint_pos: + simulate_action_latency: true ``` The checked-in G1 WBT owner enables this flag in diff --git a/docs/sphinx/source/en/3-deployment/2-sim_to_sim/7-config_guard.md b/docs/sphinx/source/en/3-deployment/2-sim_to_sim/7-config_guard.md index eaa655e34..af3a5db91 100644 --- a/docs/sphinx/source/en/3-deployment/2-sim_to_sim/7-config_guard.md +++ b/docs/sphinx/source/en/3-deployment/2-sim_to_sim/7-config_guard.md @@ -50,4 +50,4 @@ guard still remains active. - {doc}`1-backend_swap` - {doc}`4-reward_parity` -- {doc}`../../4-developer_guide/9-sim2sim_contract_status` +- {doc}`/zh_CN/4-developer_guide/9-sim2sim_contract_status` 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 94ddb57d4..644ecd7ae 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 @@ -100,6 +100,6 @@ payloads. - DR types: `src/unilab/dr/types.py` - DR manager: `src/unilab/dr/manager.py` - Backend interface: `src/unilab/base/backend/base.py` -- Example providers: `src/unilab/tasks/locomotion/g1/joystick.py`, - `src/unilab/tasks/motion_tracking/g1/tracking.py`, +- Example providers: `src/unilab/tasks/locomotion/common/dr_provider.py`, + `src/unilab/tasks/locomotion/go2_arm/manip_loco.py`, `src/unilab/tasks/manipulation/sharpa_inhand/rotation.py` diff --git a/docs/sphinx/source/en/4-developer_guide/7-motion_assets.md b/docs/sphinx/source/en/4-developer_guide/7-motion_assets.md index c32b02dfb..de3a966b7 100644 --- a/docs/sphinx/source/en/4-developer_guide/7-motion_assets.md +++ b/docs/sphinx/source/en/4-developer_guide/7-motion_assets.md @@ -120,7 +120,8 @@ To add a new robot's meshes: - Hot paths (`step` / `reset`) never trigger any file download or parsing. - `ASSETS_ROOT_PATH` is unchanged, so the download target matches the original local path exactly. -- Robot meshes use the same directory resolver (`resolve_robot_asset_dir`), - integrated at `X2WallFlipTrackingEnv.__init__` in - `src/unilab/tasks/motion_tracking/x2/flip_tracking.py`, and exposed as the - `unilab-pull-assets` CLI. +- Robot meshes use the same directory resolver (`resolve_robot_asset_dir`). The + thin `make_x2_wall_flip_env` factory in + `src/unilab/tasks/motion_tracking/x2/__init__.py` resolves them once before + delegating to the shared manager environment factory. The resolver is also + exposed through the `unilab-pull-assets` CLI. diff --git a/docs/sphinx/source/zh_CN/2-user_guide/4-tasks/2-motion_tracking.md b/docs/sphinx/source/zh_CN/2-user_guide/4-tasks/2-motion_tracking.md index 8812fd3c4..0c2d74903 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/4-tasks/2-motion_tracking.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/4-tasks/2-motion_tracking.md @@ -11,28 +11,30 @@ G1 动作追踪任务位于 `src/unilab/tasks/motion_tracking/` 下，并通过 ## Task Owners -每个 task 在 env 配置 dataclass 中定义了默认 motion 片段： +每个 task 都在 Hydra task owner YAML 中定义默认 motion 片段。Hydra 是唯一配置入口； +选中的 owner 会被物化为共享的 `ManagerBasedRlEnvCfg`，再由 NumPy Manager-Based +runtime 执行。 | CLI Task | Registered Env | 默认 motion | Owner Evidence | | --- | --- | --- | --- | | `g1_motion_tracking` | `G1MotionTracking` | `dance1_subject2_part.npz` | `conf/ppo/task/g1_motion_tracking/`, `conf/appo/task/g1_motion_tracking/` | | `g1_flip_tracking` | `G1FlipTracking` | `flip_360_001__A304.npz` | `conf/ppo/task/g1_flip_tracking/`, `conf/appo/task/g1_flip_tracking/` | | `g1_wall_flip_tracking` | `G1WallFlipTracking` | `flip_from_wall_104__A304.npz` | `conf/ppo/task/g1_wall_flip_tracking/`, `conf/appo/task/g1_wall_flip_tracking/` | -| `x2_wall_flip_tracking` | `X2WallFlipTracking` | `tictacflip_6-3_g1format.npz` | `conf/ppo/task/x2_wall_flip_tracking/`（仅 MuJoCo） | -| `g1_climb_tracking` | G1 climb tracking env | 由 env 配置给出 | `conf/ppo/task/g1_climb_tracking/`, `conf/appo/task/g1_climb_tracking/` | -| `g1_box_tracking` | G1 box tracking env | 由 env 配置给出 | `conf/ppo/task/g1_box_tracking/` | -| `g1_wbt_obs` | `G1MotionTrackingSAC` | 与 `g1_motion_tracking` 共用 | `conf/offpolicy/task/sac/g1_wbt_obs/mujoco.yaml` | +| `x2_wall_flip_tracking` | `X2WallFlipTracking` | `tictacflip_6-3_g1format.npz` | `conf/ppo/task/x2_wall_flip_tracking/` | +| `g1_climb_tracking` | `G1ClimbTracking` | `climb_20_z_scale_1.0.npz` | `conf/ppo/task/g1_climb_tracking/`, `conf/appo/task/g1_climb_tracking/` | +| `g1_box_tracking` | `G1BoxTracking` | `sub3_largebox_003_boxconverted.npz` | `conf/ppo/task/g1_box_tracking/` | +| `g1_wbt_obs` | `G1WBTObs` | `dance1_subject2_part.npz` | `conf/offpolicy/task/sac/g1_wbt_obs/mujoco.yaml` | -默认值在代码中设定：`dance1_subject2_part.npz`（`g1/tracking.py`）， -`flip_360_001__A304.npz` 与 `flip_from_wall_104__A304.npz`（`g1/flip_tracking.py`）， -以及 `tictacflip_6-3_g1format.npz`（`x2/flip_tracking.py`）。 +23-DoF task owner 目录选择对应的 23-DoF 场景、motion、entity 与 action 声明。 +profile 差异全部留在 Hydra 中。G1 identity 使用共享 manager factory；X2 只在委托给 +该 factory 前增加一层冷路径 mesh resolver。 ## PPO 与 APPO PPO owner 迭代预算（`--sim mujoco` owner YAML）：`g1_motion_tracking` 为 `algo.max_iterations=15000`；`g1_flip_tracking` 和 `g1_wall_flip_tracking` 为 -`20000`；仅 MuJoCo 的 `x2_wall_flip_tracking` 为 `9500`。（`g1_flip_tracking` 的 -Motrix owner YAML 将其提到 `30000`。） +`20000`；`x2_wall_flip_tracking` 为 `9500`。（`g1_flip_tracking` 的 Motrix owner +YAML 将其提到 `30000`。） ```bash uv run train --algo ppo --task g1_motion_tracking --sim mujoco @@ -55,12 +57,13 @@ uv run train --algo sac --task g1_motion_tracking --sim mujoco training.use_amp= uv run train --algo sac --task g1_wbt_obs --sim mujoco training.use_amp=true ``` -`g1_wbt_obs` owner 是与部署对齐的 off-policy 观测配置：pelvis IMU 状态 -（`pelvis_local_linvel` / `pelvis_gyro` / `pelvis_upvector`）加上 per-term 历史观测 -（`noise_config.obs_history_length: 5`），与部署侧的 `ObservationManager` 按字节对齐。 -部署工具在 `scripts/deploy/`，观测对齐由 `tests/scripts/test_obs_alignment_g1_wbt.py` -交叉校验。当 Motrix sim2sim 回放需要引用其他日志根目录下的 checkpoint 时，用 -`uv run eval` 透传绝对路径： +`g1_wbt_obs` owner 是与部署对齐的 off-policy 观测配置。actor 的 command 与 anchor +orientation term 保持单步，`base_ang_vel`、`joint_pos`、`joint_vel` 和 `actions` term +分别声明 `history_length: 5`。这些逐项历史由 `ObservationManager` 维护并展开；actor +使用配置中的 encoder-biased joint-position term，critic 则保留 clean term。部署工具在 +`scripts/deploy/`，观测对齐由 `tests/scripts/test_obs_alignment_g1_wbt.py` 交叉校验。 +当 Motrix sim2sim 回放需要引用其他日志根目录下的 checkpoint 时，用 `uv run eval` +透传绝对路径： ```bash uv run eval --algo sac --task g1_motion_tracking --sim motrix \ @@ -69,15 +72,19 @@ uv run eval --algo sac --task g1_motion_tracking --sim motrix \ ## 动作文件 -动作 NPZ 文件通过 `env.motion_file` 读取，也支持路径列表。标准片段必须包含七个 key： +动作 NPZ 文件通过 `env.commands.motion.params.motion_file` 选择，既可传单个路径， +也可传路径列表。标准片段必须包含七个 key： `fps`、`joint_pos`、`joint_vel`、`body_pos_w`、`body_quat_w`、`body_lin_vel_w`、 `body_ang_vel_w`（在 `common/motion_loader.py` 中校验）： ```yaml env: - motion_file: - - src/unilab/assets/motions/g1/dance1_subject2_part.npz - - src/unilab/assets/motions/g1/walk1_subject5_from_csv.npz + commands: + motion: + params: + motion_file: + - motions/g1/dance1_subject2_part.npz + - motions/g1/walk1_subject5_from_csv.npz ``` 转换与检查辅助工具在 `scripts/motion/` 中： @@ -109,20 +116,21 @@ uv run scripts/motion/replay_npz.py \ ```bash CUDA_VISIBLE_DEVICES=1 uv run train --algo sac --task g1_motion_tracking --sim mujoco \ training.use_amp=true algo.seed=1 \ - +env.motion_file=src/unilab/assets/motions/g1/motion_crawl_slope_uni.npz \ - +env.scene.model_file=src/unilab/assets/robots/g1/scene_crawl_slope.xml \ - +env.sampling_mode=start \ - env.truncate_on_clip_end=true \ - +env.max_episode_seconds=20.0 \ - '+env.pose_randomization={x:[0,0],y:[0,0],z:[0,0],roll:[0,0],pitch:[0,0],yaw:[0,0]}' \ - '+env.velocity_randomization={x:[0,0],y:[0,0],z:[0,0],roll:[0,0],pitch:[0,0],yaw:[0,0]}' \ - '+env.joint_position_range=[0,0]' + env.commands.motion.params.motion_file=motions/g1/motion_crawl_slope_uni.npz \ + env.scene.model_file=src/unilab/assets/robots/g1/scene_crawl_slope.xml \ + env.commands.motion.params.sampling_mode=start \ + env.commands.motion.params.truncate_on_clip_end=true \ + env.max_episode_seconds=20.0 \ + 'env.commands.motion.params.pose_range={x:[0,0],y:[0,0],z:[0,0],roll:[0,0],pitch:[0,0],yaw:[0,0]}' \ + 'env.commands.motion.params.velocity_range={x:[0,0],y:[0,0],z:[0,0],roll:[0,0],pitch:[0,0],yaw:[0,0]}' \ + 'env.commands.motion.params.joint_position_range=[0,0]' ``` -关键覆写：`env.motion_file` 切爬坡动作；`env.scene.model_file` 切斜坡场景 -（`scene_crawl_slope.xml` 在 `src/unilab/assets/robots/g1/` 下）；`sampling_mode=start` -加 `truncate_on_clip_end=true` 从 clip 起点出发并在结尾截断；randomization 范围全置零 -复用 motion 精确初始状态。 +关键覆写：`env.commands.motion.params.motion_file` 切换爬坡动作； +`env.scene.model_file` 切换斜坡场景（`scene_crawl_slope.xml` 在 +`src/unilab/assets/robots/g1/` 下）；`sampling_mode=start` 加 +`truncate_on_clip_end=true` 从 clip 起点出发并在结尾截断；command reset 范围全置零 +即可复用 motion 的精确初始状态。 ## 交互式调试 diff --git a/docs/sphinx/source/zh_CN/2-user_guide/5-domain_randomization/0-index.md b/docs/sphinx/source/zh_CN/2-user_guide/5-domain_randomization/0-index.md index a75dbbd25..46d6fb1d7 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/5-domain_randomization/0-index.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/5-domain_randomization/0-index.md @@ -18,7 +18,7 @@ ## 状态结论 1. 当前所有接入 DR provider 的任务都使用统一的 DR 入口点；没有任何任务绕开 `DomainRandomizationManager` 在 `reset()` 内部运行单独的 DR 流程。 -2. 它们的结构都大致相同：task 文件定义一个 `domain_rand` 配置 dataclass、一个 `DomainRandomizationProvider` 和一个 `ResetPlan`；`G1Walk*` 则改为通过 Hydra `EventTermCfg` Manager-Based reset term 声明 DR。 +2. 它们的结构都大致相同：legacy task owner 定义 `domain_rand` 配置 dataclass、`DomainRandomizationProvider` 和 `ResetPlan`；Manager-Based owner 则通过 Hydra command/event term 声明 reset 行为。G1 motion reset 扰动归 `MotionCommandCfg` 所有，WBT 另加 `EventTermCfg` reset 与 interval term。 3. 今天所"统一"的主要是入口点和执行流程，而不是每一个随机化项本身。共享辅助函数 `build_common_reset_randomization()` 目前生成 `base_mass_delta`、`base_com_offset`、`gravity`、`kp`、`kd`；共享的 interval 辅助函数目前只生成 push。 4. `ResetRandomizationPayload` 已经可以表达 `gravity`、`body_iquat`、`body_inertia`、`kp`、`kd`，并且 `MuJoCoBackend` 已声明支持。这些是否实际被使用，仍取决于 task provider 是否对它们进行采样和 dispatch。 5. `MotrixBackend` 目前支持 `base_mass_delta`、`base_com_offset`、`kp`、`kd` 和 interval push；并且它要求在初始化期间所有模型 actuator 都是 position actuator。 @@ -32,7 +32,8 @@ | `Go2JoystickFlat` | 是 | 是：`Domain_Rand + Provider + ResetPlan` | task 状态采样 + common payload | push | `go2/joystick.py` | | `G1WalkFlat` | 是 | 是：Hydra `EventTermCfg` + Manager-Based reset term | root-state reset + 经 `pd_gains` 的 kp/kd | 无 | `g1/manager_terms.py` | | `G1WalkRough` | 是 | 是：与 `G1WalkFlat` 相同的 Manager-Based event term | root-state reset + 经 `pd_gains` 的 kp/kd | 无 | `g1/manager_terms.py` | -| `G1MotionTracking` | 是 | 是：`Domain_Rand + Provider + ResetPlan` | 大量 task 专属的 reset 采样 + common payload | push | `motion_tracking/g1/tracking.py` | +| `G1MotionTracking` | 是 | 是：Hydra `MotionCommandCfg` + Manager-Based command reset | motion frame、root pose/velocity 与 joint-position 采样 | 无 | `motion_tracking/common/manager_terms.py` | +| `G1WBTObs` | 是 | 是：同一 motion command + Hydra `EventTermCfg` | motion reset 加 mass/COM/PD/friction/encoder-bias event | interval velocity kick | `motion_tracking/g1/manager_terms.py` | | `AllegroInhandRotation` | 是 | 是：Hydra `EventTermCfg` + Manager-Based reset term | entity 范围的手/球 reset | 无 | `allegro_inhand/manager_terms.py` | | `AllegroInhandRotationGrasp` | 是 | 是：复用 rotation reset event + `RecorderTermCfg` | 带噪声的手部 reset + grasp 收集 | 无 | `allegro_inhand/grasp_gen.py` | | `SharpaInhandRotation` | 是 | 是：`InitRandomizationPlan + ResetPlan + IntervalRandomizationPlan` | grasp cache 采样 + common payload | 物体 `body_force` | `sharpa_inhand/rotation.py` | @@ -46,7 +47,8 @@ | `Go2JoystickFlat` | base xy；base yaw；base qvel；command 采样；`current_actions/last_actions` 清零；kp/kd 随机化（默认启用）；可选 `base_mass_delta`；可选 `base_com_offset`；可选 `gravity` | `push_robots` | kp/kd 默认启用；common payload 和 push 默认禁用 | | `G1WalkFlat` | 经 `reset_root_state_uniform` 的 base xy/yaw 与 base qvel；带平面死区的 command 采样；`gait_phase` 采样；经 `pd_gains` 的 kp/kd 随机化 | 无 | mujoco owner 默认启用 kp/kd；motrix/mjwarp owner 默认禁用 | | `G1WalkRough` | 与 `G1WalkFlat` 相同（共享 owner base，rough 场景） | 无 | 与 `G1WalkFlat` 相同的默认值 | -| `G1MotionTracking` | 动作帧采样；root 位姿扰动 `x/y/z/roll/pitch/yaw`；root 速度扰动 `x/y/z/roll/pitch/yaw`；关节位置噪声；在 MuJoCo 下被关节范围 clip；`current_actions/last_actions` 清零；可选 `base_mass_delta`；可选 `base_com_offset`；可选 `gravity` | `push_robots` | `pose_randomization`、`velocity_randomization`、`joint_position_range` 默认有非零扰动；common payload 和 push 默认禁用 | +| `G1MotionTracking` | Motion-command frame 采样；root 位姿扰动 `x/y/z/roll/pitch/yaw`；root 速度扰动 `x/y/z/roll/pitch/yaw`；通过 public entity soft limit clip 的关节位置噪声；action-manager 状态 reset | 无 | base owner 中 `pose_range`、`velocity_range` 与 `joint_position_range` 默认有非零扰动 | +| `G1WBTObs` | 同一 motion reset 加 base mass、base COM、PD gain、足端摩擦和 encoder-bias event term | `push_by_setting_velocity` | WBT owner 显式启用上述全部 event term；能力不支持时直接报错，不回退 | | `AllegroInhandRotation` | entity 范围的手/球 reset；显式配置 grasp cache 时进行采样，否则以 `null` 显式选择模型 home pose；可选 `joint_noise`、`ball_velocity_noise` 与 `ball_z_offset` | 无 | owner YAML 显式选择 home pose 与零 reset 噪声；配置的 cache 缺失或格式错误时 fail-closed | | `AllegroInhandRotationGrasp` | 复用 rotation reset 并设置 `joint_noise=0.25`；Manager-Based termination 检查指尖距离、接触数和球高度；recorder 保存成功 timeout rows | 无 | 生成 5 万行 Allegro grasp cache，成功保存后抛出 `RunComplete` | | `SharpaInhandRotation` | grasp cache 按 `scale_ids` 分桶采样；物体位姿 / quat reset；可选 common reset 随机化 payload（含 `gravity`） | 物体 `body_force` 直接力扰动 | `domain_rand.scale_list` 默认值来自 owner YAML；在 MuJoCo 下，物体 geom 缩放在 init 期间 materialize；common payload 默认禁用；物体 force 通过 Sharpa owner YAML 默认启用 | @@ -74,7 +76,7 @@ 这意味着： - 尽管运动控制任务都走统一入口点，但它们的 base xy、yaw、qvel、command 和 gait phase 仍然直接在各自的 provider 内部采样 -- `G1MotionTracking` 的 pose / velocity / joint 噪声也是 task 专属逻辑 +- `G1MotionTracking` 的 pose / velocity / joint 噪声由其 manager command 所有 - Allegro 的 grasp / 物体初始状态采样完全是 task 专属逻辑 - Sharpa 的 `geom_size` 缩放是 init 生命周期的模型 materialization，不属于 reset common payload 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 6f36f1c30..e05b06cde 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 @@ -26,10 +26,8 @@ 具有代表性的 provider 实现位于： -- `src/unilab/tasks/locomotion/go1/joystick.py` -- `src/unilab/tasks/locomotion/g1/joystick.py` -- `src/unilab/tasks/motion_tracking/g1/tracking.py` -- `src/unilab/tasks/locomotion/go2/footstand.py` +- `src/unilab/tasks/locomotion/common/dr_provider.py` +- `src/unilab/tasks/locomotion/go2_arm/manip_loco.py` - `src/unilab/tasks/manipulation/sharpa_inhand/rotation.py` 开发者 contract 详情见 diff --git a/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/1-overview.md b/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/1-overview.md index 223c74a92..f385074e4 100644 --- a/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/1-overview.md +++ b/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/1-overview.md @@ -60,9 +60,9 @@ flowchart LR - **观测漂移。** 仿真与部署运行时之间的传感器预处理不同（单位、坐标系、滤波截止频率）。 记录第一段部署侧观测窗口，并与用同一份 owner YAML 构建的仿真回合作对比。 -- **动作延迟。** 一些任务配置通过 `control_config.simulate_action_latency` 暴露单步 - 延迟的动作执行。测量部署回路，并在硬件运行前让训练 owner 匹配该契约。见 - {doc}`8-latency_budget`。 +- **动作延迟。** 一些 task owner 通过 control config 或 Manager-Based action term + 暴露单步延迟的动作执行。测量部署回路，并在硬件运行前让训练 owner 匹配该契约。 + 见 {doc}`8-latency_budget`。 - **摩擦 / 阻尼不匹配。** 尤其对于手内操作。在 DR 中扫动摩擦；通过 {doc}`../2-sim_to_sim/3-contact_and_friction_alignment` 交叉核对。 - **复位瞬态。** 仿真复位到一个稳定姿态；部署则从一个控制器状态开始。安全层必须在 diff --git a/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/8-latency_budget.md b/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/8-latency_budget.md index 0d972f6f6..916b08856 100644 --- a/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/8-latency_budget.md +++ b/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/8-latency_budget.md @@ -7,20 +7,21 @@ | 面 | 仓库证据 | 它覆盖什么 | | --- | --- | --- | -| 单步动作延迟 | locomotion 与 G1 运动跟踪环境中的 `control_config.simulate_action_latency` | 执行上一步动作而非当前动作。 | -| G1 WBT 观测历史 | `noise_config.obs_history_length` 与 `scripts/deploy/export_deploy_config.py` | 为 `gyro`、`joint_pos_rel`、`dof_vel` 与 `last_actions` 导出逐项的 `obs_layout` 历史。 | +| 单步动作延迟 | task owner 中 Manager action term 的 `simulate_action_latency` 声明 | 执行上一步动作而非当前动作。 | +| G1 WBT 观测历史 | `conf/offpolicy/task/sac/g1_wbt_obs/mujoco.yaml` 中逐 term 的 `history_length` 与 `scripts/deploy/export_deploy_config.py` | 为 `gyro`、`joint_pos_rel`、`dof_vel` 与 `last_actions` 导出逐项的 `obs_layout` 历史。 | | Sharpa 触觉接触延迟 | Sharpa 手内配置中的 `domain_rand.contact_latency` | 为采样到的接触通道保留上一步的触觉接触值。 | | 部署侧 ONNX 契约检查 | `scripts/deploy/sim_prototype.py` | 为 G1 WBT 路径校验 `obs_layout`、`obs_dim`、ONNX 输入宽度、钳制以及 EMA 动作平滑。 | ## 动作延迟 -对于暴露 `control_config.simulate_action_latency` 的任务，当该开关启用时，环境会应用 -`last_actions`。把它保留在所选的任务 owner YAML 中，而不要事后添加仅部署的行为。 +对于启用 action latency 的 Manager-Based 任务，action manager 会在该开关开启时应用 +上一步 action。把它保留在所选的 task owner YAML 中，而不要事后添加仅部署的行为。 ```yaml env: - control_config: - simulate_action_latency: true + actions: + joint_pos: + simulate_action_latency: true ``` 已签入的 G1 WBT owner 在 `conf/offpolicy/task/sac/g1_wbt_obs/mujoco.yaml` 中启用了 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 ef5f820f4..1fc3f0f03 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 @@ -90,6 +90,6 @@ actuator 的机制泄漏到共享 payload 里。 - DR 类型：`src/unilab/dr/types.py` - DR manager：`src/unilab/dr/manager.py` - Backend 接口：`src/unilab/base/backend/base.py` -- 示例 provider：`src/unilab/tasks/locomotion/g1/joystick.py`、 - `src/unilab/tasks/motion_tracking/g1/tracking.py`、 +- 示例 provider：`src/unilab/tasks/locomotion/common/dr_provider.py`、 + `src/unilab/tasks/locomotion/go2_arm/manip_loco.py`、 `src/unilab/tasks/manipulation/sharpa_inhand/rotation.py` diff --git a/docs/sphinx/source/zh_CN/4-developer_guide/7-motion_assets.md b/docs/sphinx/source/zh_CN/4-developer_guide/7-motion_assets.md index 5db21042e..7e04ba0b4 100644 --- a/docs/sphinx/source/zh_CN/4-developer_guide/7-motion_assets.md +++ b/docs/sphinx/source/zh_CN/4-developer_guide/7-motion_assets.md @@ -107,6 +107,7 @@ uv run unilab-pull-assets --robot x2 `MotionLoader.__init__`，在冷路径上调用一次 resolver。 - 热路径（`step` / `reset`）**不会**触发任何文件下载或解析。 - `ASSETS_ROOT_PATH` 定义不变，下载落盘位置与原始本地路径完全一致。 -- 机器人网格使用同一目录 resolver（`resolve_robot_asset_dir`），集成点为 - `src/unilab/tasks/motion_tracking/x2/flip_tracking.py` 中的 - `X2WallFlipTrackingEnv.__init__`，并通过 `unilab-pull-assets` CLI 暴露。 +- 机器人网格使用同一目录 resolver（`resolve_robot_asset_dir`）。 + `src/unilab/tasks/motion_tracking/x2/__init__.py` 中的薄 + `make_x2_wall_flip_env` factory 会先在冷路径解析一次，再委托给共享 manager env + factory；同一 resolver 也通过 `unilab-pull-assets` CLI 暴露。 diff --git a/scripts/benchmark/torch_env/motion_tracking.py b/scripts/benchmark/torch_env/motion_tracking.py index 7d7d07bf9..be0bae6cb 100644 --- a/scripts/benchmark/torch_env/motion_tracking.py +++ b/scripts/benchmark/torch_env/motion_tracking.py @@ -1,30 +1,23 @@ -"""G1MotionTrackingSAC (SAC/mujoco) update_state / reset_done workload. +"""G1MotionTrackingSAC (SAC/MuJoCo) numeric manager workload. -Faithful xp-port of the NumPy computation in the collector-timed sections of +Synthetic xp-port of the NumPy kernels in the collector-timed sections of `uv run train --algo sac --task g1_motion_tracking --sim mujoco` (num_envs=2048, 29-dof, 14 tracked bodies): -- `MotionTrackingEnv.update_state` - (src/unilab/tasks/motion_tracking/common/tracking.py): - motion gather, relative transforms (transforms.py), terminations - (terminations.py), 9 active reward terms (rewards.py, incl. per-term logging - every 4 steps), observation build (observations.py, actor 160 / critic 289 - with the SAC +3 linvel tail), adaptive motion-sampler bookkeeping - (motion_loader.py). -- `MotionTrackingDomainRandomizationProvider.build_reset_plan` / - `build_reset_observation` + `build_motion_reference_state` (reset.py). -- `NpEnv._reset_done_envs` scatter/gather. +- `MotionCommand` gather, relative transforms, termination/reward terms, + observation-group assembly (actor 160 / critic 289), and adaptive sampler + bookkeeping from the Manager-Based motion runtime. +- Motion-command reset-state construction and `NpEnv._reset_done_envs` + scatter/gather. Excluded (identical across variants, not NumPy/Torch env math): `backend.step` physics, `backend.set_state`, sensor/body-state reads (replaced by persistent arrays), and the adaptive-sampler entropy metrics (3 scalar reductions per reset). -The real `build_motion_reference_state` samples pose/velocity randomization -with a per-element Python loop (num_reset x 6 draws, twice). The NumPy -workload reproduces that faithfully; pass ``vectorized_reset_rng=True`` for a -column-wise vectorized NumPy draw (used for cross-backend RNG replay during -validation, which is also the only mode Torch implements). +Pass ``vectorized_reset_rng=True`` for a column-wise vectorized NumPy draw; +validation uses that mode for cross-backend RNG replay, and Torch implements +that mode only. """ from __future__ import annotations diff --git a/scripts/benchmark/torch_env/run_benchmark.py b/scripts/benchmark/torch_env/run_benchmark.py index a9e3dae71..91d21a3c7 100644 --- a/scripts/benchmark/torch_env/run_benchmark.py +++ b/scripts/benchmark/torch_env/run_benchmark.py @@ -1,9 +1,9 @@ #!/usr/bin/env python3 -"""NumPy vs Torch comparison for the collector-timed env computation. +"""NumPy vs Torch comparison for collector-timed numeric manager kernels. -Reproduces the `update_state` and `reset_done` sections of `NpEnv.step` (the -`env_step_update_state_ms` / `env_step_reset_done_ms` collector metrics) for the -two SAC/mujoco tasks, at identical scale and computation items: +Reproduces the numeric work represented by the `env_step_update_state_ms` and +`env_step_reset_done_ms` collector metrics for two SAC/MuJoCo tasks, at +identical scale and computation items: - g1_walk_flat (num_envs=2048, 29-dof, obs 98 / critic 101) - g1_motion_tracking (num_envs=2048, 29-dof, 14 bodies, obs 160 / critic 289) diff --git a/scripts/deploy/export_deploy_config.py b/scripts/deploy/export_deploy_config.py index b692e992e..dc6bbddf5 100644 --- a/scripts/deploy/export_deploy_config.py +++ b/scripts/deploy/export_deploy_config.py @@ -1,12 +1,12 @@ #!/usr/bin/env python3 """Export deploy_config.yaml for the C++ G1-29DOF WBT deployment side. -Reads g1.xml + scene_flat.xml + tracking.py defaults to emit a single yaml +Reads the G1 scene plus the Manager-Based WBT owner contract to emit a single yaml that the deploy framework (~/deploy_ws/unitree_rl_lab/.../State_WBT) can load to drive the actor at runtime. This file is the SINGLE SOURCE OF TRUTH for the actor obs schema: - * Training side (tracking.py) assembles obs in the order documented here. + * Training side (ObservationManager) assembles terms in the order documented here. * Deploy side (State_WBT.cpp + ObservationManager) reads obs_layout from this yaml and assembles in matching order with per-term history buffers. * Alignment test (tests/test_obs_alignment_g1_wbt.py) verifies both code @@ -64,9 +64,9 @@ KEYFRAME_NAME = "stand" ROOT_QPOS_DIM = 7 # free joint: xyz + quat(wxyz) -# Default obs history length — matches g1_wbt_obs/mujoco.yaml's -# `noise_config.obs_history_length`. Override via --obs-history-length when -# exporting for other training profiles (e.g. g1_motion_tracking/mujoco.yaml uses 1). +# Default obs history length — matches the per-term ``history_length`` declarations +# in g1_wbt_obs/mujoco.yaml. Override via --obs-history-length when exporting for +# other training profiles (e.g. g1_motion_tracking/mujoco.yaml uses one step). DEFAULT_OBS_HISTORY_LENGTH = 5 @@ -77,7 +77,7 @@ def _round_list(arr, ndigits=6): def _build_obs_layout( num_action: int, hist_len: int, enable_zero_anchor_pos: bool, enable_zero_linvel: bool ): - """Build obs_layout in the exact order tracking.py:_compute_obs assembles. + """Build obs_layout in the exact order ObservationManager assembles it. Returns (layout_list, total_obs_dim). Order = single-step refs first, then per-term proprio history blocks, @@ -162,8 +162,8 @@ def main(): "--obs-history-length", type=int, default=DEFAULT_OBS_HISTORY_LENGTH, - help="Proprio history length H. Must match training-side " - "noise_config.obs_history_length. Default 5 = current " + help="Proprio history length H. Must match training-side per-term " + "history_length. Default 5 = current " "g1_wbt_obs/mujoco.yaml. Set 1 for the legacy 154-d schema.", ) ap.add_argument( @@ -171,14 +171,14 @@ def main(): action="store_true", default=True, help="Drop motion_anchor_pos_b from actor obs (mjlab parity). " - "Matches g1_wbt_obs/mujoco.yaml's noise_config flag.", + "Matches the null actor term in g1_wbt_obs/mujoco.yaml.", ) ap.add_argument( "--enable-zero-linvel", action="store_true", default=True, help="Drop base_lin_vel from actor obs (mjlab parity). " - "Matches g1_wbt_obs/mujoco.yaml's noise_config flag.", + "Matches the null actor term in g1_wbt_obs/mujoco.yaml.", ) args = ap.parse_args() @@ -299,7 +299,7 @@ def main(): # ---- obs layout (SINGLE SOURCE OF TRUTH for both ends) ---- # Each entry: name, dim (per-step), history_length, source. # Total obs_dim = sum(dim * history_length). - # Order MUST match tracking.py:_compute_obs assembly order. + # Order MUST match the WBT owner term order consumed by ObservationManager. # State_WBT.cpp:build_env_cfg translates names via its alias table. "obs_layout": obs_layout, } diff --git a/src/unilab/envs/manager_based_rl_env.py b/src/unilab/envs/manager_based_rl_env.py index 01c4133af..a5b615126 100644 --- a/src/unilab/envs/manager_based_rl_env.py +++ b/src/unilab/envs/manager_based_rl_env.py @@ -174,13 +174,19 @@ def _resolve_backend_entity_contract(cfg: ManagerBasedRlEnvCfg) -> tuple[str, bo if entity_cfg.body_names is not None: body_state_requested = True - if len(root_entities) != 1: + if not root_entities: + raise ValueError( + "ManagerBasedRlEnv factory requires at least one scene entity with an explicit " + "root_body_name" + ) + primary = next((item for item in root_entities if item[0] == "robot"), None) + if primary is None and len(root_entities) != 1: declared = [name for name, _ in root_entities] raise ValueError( - "ManagerBasedRlEnv factory requires exactly one scene entity with an explicit " - f"root_body_name; found {len(root_entities)} root entities {declared}" + "ManagerBasedRlEnv factory requires a conventional 'robot' root entity when " + f"multiple floating entities are declared; found {declared}" ) - return root_entities[0][1], body_state_requested + return (primary or root_entities[0])[1], body_state_requested class ManagerBasedRlEnv(NpEnv): diff --git a/src/unilab/tasks/migration_matrix.py b/src/unilab/tasks/migration_matrix.py index 8691ae85e..511fe2db7 100644 --- a/src/unilab/tasks/migration_matrix.py +++ b/src/unilab/tasks/migration_matrix.py @@ -160,10 +160,10 @@ def migration_record(task_name: str) -> TaskMigrationRecord: return TaskMigrationRecord( task_name, "motion_tracking", - "Adapted", - "mba", - "Stateful motion loading and profile-specific tracking terms need a grouped manager port.", - "Migrate the shared engine and all profiles together; stop on new backend contracts.", + "Compatible", + "complete", + "Hydra profile owners specialize the shared NumPy motion managers without a legacy runtime.", + "Keep profile scene, motion, observation, reward, and termination declarations aligned.", ) raise KeyError(f"Task '{task_name}' has no #1042 migration-matrix entry") diff --git a/src/unilab/tasks/motion_tracking/__init__.py b/src/unilab/tasks/motion_tracking/__init__.py index c121df023..4e23b0f14 100644 --- a/src/unilab/tasks/motion_tracking/__init__.py +++ b/src/unilab/tasks/motion_tracking/__init__.py @@ -1,99 +1,5 @@ -"""Motion-tracking task packages.""" +"""Hydra-owned motion tracking tasks on the NumPy Manager-Based runtime.""" -from .g1 import ( - BoxMotionData, - BoxMotionLoader, - G1BoxTracking23DofCfg, - G1BoxTracking23DofEnvCfg, - G1BoxTrackingCfg, - G1BoxTrackingEnv, - G1BoxTrackingEnvCfg, - G1ClimbTrackingCfg, - G1ClimbTrackingEnv, - G1ClimbTrackingEnvCfg, - G1FlipTracking23DofCfg, - G1FlipTracking23DofEnvCfg, - G1FlipTrackingCfg, - G1FlipTrackingEnv, - G1FlipTrackingEnvCfg, - G1FlipTrackingSAC23DofCfg, - G1FlipTrackingSAC23DofEnv, - G1FlipTrackingSACCfg, - G1FlipTrackingSACEnv, - G1MotionTracking23DofCfg, - G1MotionTracking23DofDeployEnvCfg, - G1MotionTracking23DofEnvCfg, - G1MotionTrackingCfg, - G1MotionTrackingEnv, - G1MotionTrackingEnvCfg, - G1MotionTrackingSAC23DofCfg, - G1MotionTrackingSAC23DofEnv, - G1MotionTrackingSACCfg, - G1MotionTrackingSACEnv, - G1WallFlipTracking23DofCfg, - G1WallFlipTracking23DofEnvCfg, - G1WallFlipTrackingCfg, - G1WallFlipTrackingEnv, - G1WallFlipTrackingEnvCfg, - G1WallFlipTrackingSAC23DofCfg, - G1WallFlipTrackingSAC23DofEnv, - G1WallFlipTrackingSACCfg, - G1WallFlipTrackingSACEnv, - G1WBTObs23DofCfg, - G1WBTObsCfg, - G1WBTObsEnv, -) -from .x2 import ( - X2MotionTrackingCfg, - X2WallFlipTrackingCfg, - X2WallFlipTrackingEnv, - X2WallFlipTrackingEnvCfg, -) +from .g1.motion_box_loader import BoxMotionData, BoxMotionLoader -__all__ = [ - "BoxMotionData", - "BoxMotionLoader", - "G1BoxTracking23DofCfg", - "G1BoxTracking23DofEnvCfg", - "G1BoxTrackingCfg", - "G1BoxTrackingEnv", - "G1BoxTrackingEnvCfg", - "G1ClimbTrackingCfg", - "G1ClimbTrackingEnv", - "G1ClimbTrackingEnvCfg", - "G1FlipTracking23DofCfg", - "G1FlipTracking23DofEnvCfg", - "G1FlipTrackingCfg", - "G1FlipTrackingEnv", - "G1FlipTrackingEnvCfg", - "G1FlipTrackingSAC23DofCfg", - "G1FlipTrackingSAC23DofEnv", - "G1FlipTrackingSACCfg", - "G1FlipTrackingSACEnv", - "G1MotionTracking23DofCfg", - "G1MotionTracking23DofDeployEnvCfg", - "G1MotionTracking23DofEnvCfg", - "G1MotionTrackingCfg", - "G1MotionTrackingEnv", - "G1MotionTrackingEnvCfg", - "G1MotionTrackingSAC23DofCfg", - "G1MotionTrackingSAC23DofEnv", - "G1MotionTrackingSACCfg", - "G1MotionTrackingSACEnv", - "G1WallFlipTracking23DofCfg", - "G1WallFlipTracking23DofEnvCfg", - "G1WallFlipTrackingCfg", - "G1WallFlipTrackingEnv", - "G1WallFlipTrackingEnvCfg", - "G1WallFlipTrackingSAC23DofCfg", - "G1WallFlipTrackingSAC23DofEnv", - "G1WallFlipTrackingSACCfg", - "G1WallFlipTrackingSACEnv", - "G1WBTObs23DofCfg", - "G1WBTObsCfg", - "G1WBTObsEnv", - "X2MotionTrackingCfg", - "X2WallFlipTrackingCfg", - "X2WallFlipTrackingEnv", - "X2WallFlipTrackingEnvCfg", -] +__all__ = ["BoxMotionData", "BoxMotionLoader"] diff --git a/src/unilab/tasks/motion_tracking/common/__init__.py b/src/unilab/tasks/motion_tracking/common/__init__.py index 941751b15..fa134d4b7 100644 --- a/src/unilab/tasks/motion_tracking/common/__init__.py +++ b/src/unilab/tasks/motion_tracking/common/__init__.py @@ -1,7 +1 @@ -"""Shared motion-tracking task engine and owner modules. - -This package holds the motion-tracking task engine and its per-concern owner -modules: config, rewards, observations, terminations, transforms, reset, -domain randomization, and motion loading. Per-robot profiles live under -``g1/`` and ``x2/`` and contain robot-specific defaults and registry leaves. -""" +"""Shared NumPy motion command, loader, and numeric manager helpers.""" diff --git a/src/unilab/tasks/motion_tracking/common/config.py b/src/unilab/tasks/motion_tracking/common/config.py deleted file mode 100644 index 9f1aaee34..000000000 --- a/src/unilab/tasks/motion_tracking/common/config.py +++ /dev/null @@ -1,168 +0,0 @@ -"""Configuration dataclasses for the robot-agnostic motion-tracking engine. - -The default field values are kept at the historical G1 profile so that the -registered ``G1MotionTracking`` / ``G1MotionTrackingDeploy`` configs remain -field-identical to the pre-refactor monolith. Per-robot profiles (g1/x2) -override these via subclasses. -""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from typing import Literal - -from unilab.assets import ASSETS_ROOT_PATH -from unilab.base.scene import SceneCfg -from unilab.tasks.motion_tracking.common.g1_base import G1BaseCfg - -from .rewards import RewardConfig - - -@dataclass -class PoseRandomization: - """Pose randomization ranges for reset.""" - - x: tuple[float, float] = (-0.05, 0.05) - y: tuple[float, float] = (-0.05, 0.05) - z: tuple[float, float] = (-0.01, 0.01) - roll: tuple[float, float] = (-0.1, 0.1) - pitch: tuple[float, float] = (-0.1, 0.1) - yaw: tuple[float, float] = (-0.2, 0.2) - - -@dataclass -class VelocityRandomization: - """Velocity randomization ranges for reset.""" - - x: tuple[float, float] = (-0.5, 0.5) - y: tuple[float, float] = (-0.5, 0.5) - z: tuple[float, float] = (-0.2, 0.2) - roll: tuple[float, float] = (-0.52, 0.52) - pitch: tuple[float, float] = (-0.52, 0.52) - yaw: tuple[float, float] = (-0.78, 0.78) - - -@dataclass -class DomainRand: - """Domain randomization config required by motrix backend hooks.""" - - randomize_base_mass: bool = False - added_mass_range: list[float] = field(default_factory=lambda: [-1.5, 1.5]) - - random_com: bool = False - com_offset_x: list[float] = field(default_factory=lambda: [-0.05, 0.05]) - com_offset_y: list[float] = field(default_factory=lambda: [-0.05, 0.05]) - com_offset_z: list[float] = field(default_factory=lambda: [-0.05, 0.05]) - - randomize_gravity: bool = False - gravity_range: list[list[float]] = field( - default_factory=lambda: [[0.0, 0.0, -9.81], [0.0, 0.0, -9.81]] - ) - - push_robots: bool = False - push_interval: int = 750 - max_force: list[float] = field(default_factory=lambda: [1.0, 1.0, 0.5]) - push_body_name: str | None = None - - randomize_kp: bool = False - kp_multiplier_range: list[float] = field(default_factory=lambda: [0.9, 1.1]) - - randomize_kd: bool = False - kd_multiplier_range: list[float] = field(default_factory=lambda: [0.9, 1.1]) - - randomize_geom_friction: bool = False - friction_range: list[float] = field(default_factory=lambda: [0.3, 1.2]) - friction_geom_pattern: str = r"^(left|right)_foot[1-7]_collision$" - - randomize_joint_default_pos: bool = False - joint_default_pos_range: list[float] = field(default_factory=lambda: [-0.01, 0.01]) - - -# Backward-compatible alias for the historical (snake_case) class name. -Domain_Rand = DomainRand - - -def _zero_pose_randomization() -> PoseRandomization: - return PoseRandomization( - x=(0.0, 0.0), - y=(0.0, 0.0), - z=(0.0, 0.0), - roll=(0.0, 0.0), - pitch=(0.0, 0.0), - yaw=(0.0, 0.0), - ) - - -def _zero_velocity_randomization() -> VelocityRandomization: - return VelocityRandomization( - x=(0.0, 0.0), - y=(0.0, 0.0), - z=(0.0, 0.0), - roll=(0.0, 0.0), - pitch=(0.0, 0.0), - yaw=(0.0, 0.0), - ) - - -@dataclass -class MotionTrackingCfg(G1BaseCfg): - """Configuration for the motion tracking environment.""" - - scene: SceneCfg = field( # pyright: ignore[reportIncompatibleVariableOverride] - default_factory=lambda: SceneCfg( - model_file=str(ASSETS_ROOT_PATH / "robots" / "g1" / "scene_flat.xml") - ) - ) - # Kept at the historical single-clip default for backward compatibility. - motion_file: str | list[str] = str( - ASSETS_ROOT_PATH / "motions" / "g1" / "dance1_subject2_part.npz" - ) - # motion_file: str | list[str] = str(ASSETS_ROOT_PATH / "motions" / "g1" / "gangnam_style.npz") - # motion_file: str | list[str] = str(ASSETS_ROOT_PATH / "motions" / "g1" / "fight1_subject5_from_csv.npz") #LAFAN - # motion_file: str | list[str] = str(ASSETS_ROOT_PATH / "motions" / "g1" / "dance_basic_slide_180_R_loop_001__A322_M.npz") #LAFAN - # motion_file: str | list[str] = str(ASSETS_ROOT_PATH / "motions" / "g1" / "playing_violin_R_003__A327_from_csv.npz") #Seed - anchor_body_name: str = "torso_link" - body_names: tuple[str, ...] = ( - "pelvis", - "left_hip_roll_link", - "left_knee_link", - "left_ankle_roll_link", - "right_hip_roll_link", - "right_knee_link", - "right_ankle_roll_link", - "torso_link", - "left_shoulder_roll_link", - "left_elbow_link", - "left_wrist_yaw_link", - "right_shoulder_roll_link", - "right_elbow_link", - "right_wrist_yaw_link", - ) - sampling_mode: Literal["start", "clip_start", "uniform", "adaptive", "mixed"] = "adaptive" - sampling_start_ratio: float = 0.0 - truncate_on_clip_end: bool = False - max_episode_seconds: float = 10.0 # pyright: ignore[reportIncompatibleVariableOverride] - reward_config: RewardConfig = field(default_factory=RewardConfig) - pose_randomization: PoseRandomization = field(default_factory=PoseRandomization) - velocity_randomization: VelocityRandomization = field(default_factory=VelocityRandomization) - domain_rand: DomainRand = field(default_factory=DomainRand) - joint_position_range: tuple[float, float] = (-0.1, 0.1) - # Termination thresholds - anchor_pos_z_threshold: float = 0.25 - anchor_ori_threshold: float = 0.8 - ee_body_pos_z_threshold: float = 0.25 - ee_body_names: tuple[str, ...] = ( - "left_ankle_roll_link", - "right_ankle_roll_link", - "left_wrist_yaw_link", - "right_wrist_yaw_link", - ) - undesired_contact_z_threshold: float = 0.05 - terminate_on_undesired_contacts: bool = False - - -@dataclass -class MotionTrackingDeployEnvCfg(MotionTrackingCfg): - """Base deploy configuration for motion tracking.""" - - pass diff --git a/src/unilab/tasks/motion_tracking/common/domain_randomization.py b/src/unilab/tasks/motion_tracking/common/domain_randomization.py deleted file mode 100644 index aa74e50e8..000000000 --- a/src/unilab/tasks/motion_tracking/common/domain_randomization.py +++ /dev/null @@ -1,186 +0,0 @@ -"""Domain-randomization reset/observation provider for motion tracking.""" - -from __future__ import annotations - -import time -from typing import Any, cast - -import numpy as np - -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.dr.types import RESET_TERM_GEOM_FRICTION, ResetRandomizationPayload -from unilab.dtype_config import get_global_dtype - -from .reset import build_motion_reference_state - - -class MotionTrackingDomainRandomizationProvider(DomainRandomizationProvider): - def __init__( - self, - *, - base_kp: np.ndarray | None = None, - base_kd: np.ndarray | None = None, - base_geom_friction: np.ndarray | None = None, - foot_geom_ids: np.ndarray | None = None, - ) -> None: - self._base_kp = base_kp - self._base_kd = base_kd - self._base_geom_friction = base_geom_friction - self._foot_geom_ids = foot_geom_ids - self._last_reset_observation_timing_ms: dict[str, float] = {} - - @property - def last_reset_observation_timing_ms(self) -> dict[str, float]: - return dict(self._last_reset_observation_timing_ms) - - def validate(self, env: Any, capabilities: DomainRandomizationCapabilities) -> None: - validate_common_reset_randomization( - env, capabilities, base_kp=self._base_kp, base_kd=self._base_kd - ) - validate_interval_push_support(env, capabilities) - if getattr(env.cfg.domain_rand, "randomize_geom_friction", False): - if not capabilities.supports_reset_term(RESET_TERM_GEOM_FRICTION): - raise NotImplementedError( - f"{env._backend.backend_type} backend does not support " - "geom-friction reset randomization" - ) - if ( - self._base_geom_friction is None - or self._foot_geom_ids is None - or self._foot_geom_ids.size == 0 - ): - raise ValueError("randomize_geom_friction=True but provider has no foot geom IDs") - - def build_interval_randomization_plan( - self, env: Any, step_counter: int - ) -> IntervalRandomizationPlan | None: - return build_interval_push_plan(env, step_counter) - - def build_reset_plan(self, env: Any, env_ids: np.ndarray) -> ResetPlan: - num_reset = len(env_ids) - motion_frames = env.motion_sampler.sample_frames(env_ids) - motion_data = env.motion_loader.get_motion_at_frame(motion_frames) - qpos, qvel = build_motion_reference_state(env, env_ids, motion_data) - - info_updates = { - "current_actions": zero_actions(num_reset, env._num_action), - "last_actions": zero_actions(num_reset, env._num_action), - } - randomization = build_common_reset_randomization( - env, num_reset, base_kp=self._base_kp, base_kd=self._base_kd - ) - - dr_cfg = env.cfg.domain_rand - if getattr(dr_cfg, "randomize_geom_friction", False): - assert self._base_geom_friction is not None - assert self._foot_geom_ids is not None - payload = randomization or ResetRandomizationPayload() - low, high = dr_cfg.friction_range - scale = np.random.uniform(low, high, size=(num_reset, 1)).astype(np.float64) - geom_friction = np.broadcast_to( - self._base_geom_friction, - (num_reset, *self._base_geom_friction.shape), - ).copy() - geom_friction[:, self._foot_geom_ids, 0] = scale * np.ones( - (1, self._foot_geom_ids.size) - ) - payload.geom_friction = geom_friction - randomization = payload - - if getattr(dr_cfg, "randomize_joint_default_pos", False): - low, high = dr_cfg.joint_default_pos_range - info_updates["default_dof_pos_bias"] = np.random.uniform( - low, high, size=(num_reset, env._num_action) - ).astype(get_global_dtype()) - - return ResetPlan( - env_ids=env_ids, - qpos=qpos, - qvel=qvel, - info_updates=info_updates, - randomization=randomization, - ) - - def build_reset_observation( - self, env: Any, env_ids: np.ndarray, info_updates: dict[str, Any] - ) -> dict[str, np.ndarray]: - obs_t0 = time.perf_counter() - - t0 = time.perf_counter() - motion_data = env.motion_loader.get_motion_at_frame( - env.motion_sampler.current_frames[env_ids] - ) - motion_ms = (time.perf_counter() - t0) * 1000.0 - - t0 = time.perf_counter() - linvel = env.get_local_linvel()[env_ids] - linvel_ms = (time.perf_counter() - t0) * 1000.0 - - t0 = time.perf_counter() - gyro = env.get_gyro()[env_ids] - gyro_ms = (time.perf_counter() - t0) * 1000.0 - - t0 = time.perf_counter() - dof_pos = env.get_dof_pos()[env_ids] - dof_pos_ms = (time.perf_counter() - t0) * 1000.0 - - t0 = time.perf_counter() - dof_vel = env.get_dof_vel()[env_ids] - dof_vel_ms = (time.perf_counter() - t0) * 1000.0 - - t0 = time.perf_counter() - robot_body_pos_w, robot_body_quat_w = env._backend.get_body_pose_w_rows( - env_ids, env.body_ids - ) - body_pose_ms = (time.perf_counter() - t0) * 1000.0 - - obs_info = dict(info_updates) - default_dof_pos_bias = info_updates.get("default_dof_pos_bias") - if isinstance(default_dof_pos_bias, np.ndarray): - obs_info["default_dof_pos_bias"] = default_dof_pos_bias - obs_info["env_ids"] = env_ids - - t0 = time.perf_counter() - obs = cast( - dict[str, np.ndarray], - env._compute_obs( - obs_info, - motion_data, - linvel, - gyro, - dof_pos, - dof_vel, - robot_body_pos_w, - robot_body_quat_w, - ), - ) - compute_obs_ms = (time.perf_counter() - t0) * 1000.0 - - total_ms = (time.perf_counter() - obs_t0) * 1000.0 - getters_ms = linvel_ms + gyro_ms + dof_pos_ms + dof_vel_ms + body_pose_ms - self._last_reset_observation_timing_ms = { - "dr_reset_observation_getters_ms": getters_ms, - "dr_reset_obs_get_motion_ms": motion_ms, - "dr_reset_obs_get_local_linvel_ms": linvel_ms, - "dr_reset_obs_get_gyro_ms": gyro_ms, - "dr_reset_obs_get_dof_pos_ms": dof_pos_ms, - "dr_reset_obs_get_dof_vel_ms": dof_vel_ms, - "dr_reset_obs_get_body_pose_ms": body_pose_ms, - "dr_reset_observation_compute_obs_ms": compute_obs_ms, - "dr_reset_observation_internal_gap_ms": ( - total_ms - motion_ms - getters_ms - compute_obs_ms - ), - } - return obs diff --git a/src/unilab/tasks/motion_tracking/common/g1_base.py b/src/unilab/tasks/motion_tracking/common/g1_base.py deleted file mode 100644 index f52afb0e0..000000000 --- a/src/unilab/tasks/motion_tracking/common/g1_base.py +++ /dev/null @@ -1,70 +0,0 @@ -"""Shared G1 runtime and configuration base for motion-tracking tasks. - -Relocated from ``unilab.tasks.locomotion.g1.base`` when the G1 locomotion -family migrated to the Manager-Based runtime; motion tracking is the sole -remaining consumer of the legacy class-based G1 base env. -""" - -from __future__ import annotations - -from dataclasses import dataclass, field - -import numpy as np - -from unilab.tasks.locomotion.common.base import ( - BaseNoiseConfig, - ControlConfigBase, - LocomotionBaseCfg, - LocomotionBaseEnv, -) -from unilab.tasks.locomotion.common.base import ( - Sensor as LocomotionSensor, -) - - -@dataclass -class NoiseConfig(BaseNoiseConfig): - scale_joint_angle: float = 0.02 - scale_joint_vel: float = 0.3 - scale_gyro: float = 0.1 - - -@dataclass -class ControlConfig(ControlConfigBase): - action_scale: float | np.ndarray = 0.25 # type: ignore[assignment] - - -@dataclass -class Sensor(LocomotionSensor): - local_linvel: str = "pelvis_local_linvel" - gyro: str = "torso_gyro" - upvector: str = "torso_upvector" - - -@dataclass -class Asset: - base_name = "pelvis" - foot_name = "ankle_roll_link" - ground = "floor" - - -@dataclass -class G1BaseCfg(LocomotionBaseCfg): - noise_config: NoiseConfig = field(default_factory=NoiseConfig) # type: ignore[assignment] - control_config: ControlConfig = field(default_factory=ControlConfig) # type: ignore[assignment] - sensor: Sensor = field( # pyright: ignore[reportIncompatibleVariableOverride] - default_factory=Sensor - ) - asset: Asset = field(default_factory=Asset) - sim_dt: float = 0.02 / 3.0 - ctrl_dt: float = 0.02 - - -class G1BaseEnv(LocomotionBaseEnv): - _cfg: G1BaseCfg # pyright: ignore[reportIncompatibleVariableOverride] - _keyframe_name = "stand" - _use_global_dtype = False - - def _obs_noise(self, data: np.ndarray, scale: float) -> np.ndarray: - """Same as base, but coerces back to ``data.dtype`` (G1 runs in float32).""" - return np.asarray(super()._obs_noise(data, scale), dtype=data.dtype) diff --git a/src/unilab/tasks/motion_tracking/common/manager_terms.py b/src/unilab/tasks/motion_tracking/common/manager_terms.py index b520a1053..9dc3b6b13 100644 --- a/src/unilab/tasks/motion_tracking/common/manager_terms.py +++ b/src/unilab/tasks/motion_tracking/common/manager_terms.py @@ -128,7 +128,7 @@ def __init__(self, cfg: MotionCommandCfg, env: ManagerBasedRlEnv): self._robot_body_ids = np.asarray(body_ids, dtype=np.intp) self._robot_body_ids.setflags(write=False) motion_body_ids = self.robot.motion_body_ids[self._robot_body_ids] - self.motion = MotionLoader(cfg.motion_file, body_indices=motion_body_ids) + self.motion = self._make_motion_loader(cfg.motion_file, motion_body_ids) if self.motion.num_joints != len(self.robot.joint_names): raise ValueError( f"MotionCommand motion joint width {self.motion.num_joints} does not match " @@ -210,6 +210,14 @@ def __init__(self, cfg: MotionCommandCfg, env: ManagerBasedRlEnv): self._refresh_robot_state(force=True) self._refresh_relative_state() + def _make_motion_loader( + self, + motion_file: str | list[str], + body_indices: np.ndarray, + ) -> MotionLoader: + """Materialize the profile-owned motion loader on the cold path.""" + return MotionLoader(motion_file, body_indices=body_indices) + @staticmethod def _validate_cfg(cfg: MotionCommandCfg) -> None: if not isinstance(cfg.entity_name, str) or not cfg.entity_name: @@ -469,6 +477,7 @@ def post_compute(self) -> None: @dataclass(kw_only=True) class MotionJointPositionActionCfg(JointPositionActionCfg): command_name: str = "motion" + simulate_action_latency: bool = False def build(self, env: ManagerBasedRlEnv) -> MotionJointPositionAction: return MotionJointPositionAction(self, env) @@ -478,8 +487,36 @@ class MotionJointPositionAction(JointPositionAction): cfg: MotionJointPositionActionCfg # pyright: ignore[reportIncompatibleVariableOverride] def __init__(self, cfg: MotionJointPositionActionCfg, env: ManagerBasedRlEnv): + if not isinstance(cfg.simulate_action_latency, bool): + raise TypeError("MotionJointPositionActionCfg simulate_action_latency must be bool") super().__init__(cfg, env) self._motion_command = _command(env, cfg.command_name) + self._previous_raw_actions = np.zeros_like(self._raw_actions) + + @property + def target(self) -> np.ndarray: + """Most recently applied physical joint target in entity joint order.""" + return self._target + + def process_actions(self, actions: np.ndarray) -> None: + self._previous_raw_actions[:] = self._raw_actions + super().process_actions(actions) + if not self.cfg.simulate_action_latency: + return + np.multiply(self._previous_raw_actions, self._scale, out=self._processed_actions) + np.add(self._processed_actions, self._offset, out=self._processed_actions) + if self._clip is not None: + np.clip( + self._processed_actions, + self._clip[..., 0], + self._clip[..., 1], + out=self._processed_actions, + ) + + def reset(self, env_ids: np.ndarray | slice | None = None) -> None: + super().reset(env_ids) + ids = slice(None) if env_ids is None else env_ids + self._previous_raw_actions[ids] = 0.0 def apply_actions(self) -> None: encoder_bias = self._entity.data.encoder_bias[:, self._target_ids] @@ -529,6 +566,16 @@ def motion_joint_pos_rel(env: ManagerBasedRlEnv, command_name: str) -> np.ndarra ) +def motion_joint_pos_rel_biased(env: ManagerBasedRlEnv, command_name: str) -> np.ndarray: + """Joint position relative to the episode default, including encoder bias.""" + command = _command(env, command_name) + return ( + command.robot.data.joint_pos_biased + - command.robot.data.default_joint_pos + - command.joint_default_bias + ) + + def _positive_std(value: float, *, term_name: str) -> float: if isinstance(value, bool) or not isinstance(value, (int, float, np.number)): raise TypeError(f"{term_name} std must be a real number") @@ -663,6 +710,23 @@ def __call__( return np.exp(-error.mean(axis=-1) / scale**2) +class motion_relative_body_position_z_error_exp(_BodyTerm): + def __call__( + self, + env: ManagerBasedRlEnv, + command_name: str, + std: float, + body_names: tuple[str, ...] | None = None, + ) -> np.ndarray: + del env, body_names + command, scale = self._validate(command_name, std) + error = np.square( + command.body_pos_relative_w[:, self._body_ids, 2] + - command.robot_body_pos_w[:, self._body_ids, 2] + ) + return np.exp(-error.mean(axis=-1) / scale**2) + + def motion_joint_position_error_exp( env: ManagerBasedRlEnv, command_name: str, std: float ) -> np.ndarray: @@ -743,6 +807,19 @@ def __call__( return np.any(error > threshold, axis=-1) +class bad_undesired_body_contacts(_BodyTerm): + def __call__( + self, + env: ManagerBasedRlEnv, + command_name: str, + threshold: float, + body_names: tuple[str, ...] | None = None, + ) -> np.ndarray: + del env, body_names + command = _command(self._env, command_name) + return np.any(command.robot_body_pos_w[:, self._body_ids, 2] < threshold, axis=-1) + + def motion_clip_end(env: ManagerBasedRlEnv, command_name: str) -> np.ndarray: command = _command(env, command_name) return command.time_steps >= command.sampler.current_clip_end_frames @@ -757,6 +834,7 @@ def motion_clip_end(env: ManagerBasedRlEnv, command_name: str) -> np.ndarray: "bad_anchor_ori", "bad_anchor_pos_z_only", "bad_motion_body_pos_z_only", + "bad_undesired_body_contacts", "joint_pos_limits", "motion_anchor_ori_b", "motion_anchor_pos_b", @@ -766,10 +844,12 @@ def motion_clip_end(env: ManagerBasedRlEnv, command_name: str) -> np.ndarray: "motion_global_body_angular_velocity_error_exp", "motion_global_body_linear_velocity_error_exp", "motion_joint_pos_rel", + "motion_joint_pos_rel_biased", "motion_joint_position_error_exp", "motion_joint_velocity_error_exp", "motion_relative_body_orientation_error_exp", "motion_relative_body_position_error_exp", + "motion_relative_body_position_z_error_exp", "robot_body_ori_b", "robot_body_pos_b", "undesired_body_contacts", diff --git a/src/unilab/tasks/motion_tracking/common/observations.py b/src/unilab/tasks/motion_tracking/common/observations.py index ba6f6203f..6332c4ac5 100644 --- a/src/unilab/tasks/motion_tracking/common/observations.py +++ b/src/unilab/tasks/motion_tracking/common/observations.py @@ -1,108 +1,9 @@ -"""Shared observation construction for motion tracking. - -Holds the robot-agnostic observation builders. The environment classes keep a -thin polymorphic method surface (``_compute_obs`` / ``_build_actor_obs`` / -``_write_body_*``) that delegates here so subclasses can still override obs -layout (SAC critic tail, box object obs, deploy mimic actor) while the core -math lives in one place. -""" +"""In-place body-frame writers shared by motion manager terms.""" from __future__ import annotations -from typing import Any - import numpy as np -from unilab.dtype_config import get_global_dtype -from unilab.utils.geometry import np_write_relative_anchor_transform_pos_rot6d - - -def actor_obs_dim(n: int) -> int: - return 3 + 6 + 3 + 3 + n * 5 - - -def critic_base_obs_dim(n: int) -> int: - return 3 + 6 + 3 + 3 + n * 5 - - -def mimic_actor_obs_dim(n: int) -> int: - # unitree_rl_lab mimic deploy actor input: - # motion_command(2n), motion_anchor_ori_b(6), gyro(3), joints, actions. - return 6 + 3 + n * 5 - - -def obs_groups_spec(env: Any) -> dict[str, int]: - # Actor: command(2n) + motion_anchor_pos_b(3) + motion_anchor_ori_b(6) - # + linvel(3) + gyro(3) + joint_pos(n) + joint_vel(n) + actions(n) - # Critic mirrors BeyondMimic physical terms without actor observation noise: - # command, motion anchor, robot body pos/ori, linvel, gyro, joints, actions. - n = env._num_action - actor_width = getattr(env, "_actor_obs_width", env._actor_obs_dim(n)) - critic_width = getattr( - env, - "_critic_obs_width", - env._critic_base_obs_dim(n) + len(env._cfg.body_names) * 9, - ) - return {"obs": actor_width, "critic": critic_width} - - -def build_actor_obs( - *, - actor_obs_dim: int, - command: np.ndarray, - motion_anchor_pos_b: np.ndarray, - motion_anchor_ori_b: np.ndarray, - noisy_linvel: np.ndarray, - noisy_gyro: np.ndarray, - noisy_joint_pos_rel: np.ndarray, - noisy_dof_vel: np.ndarray, - last_actions: np.ndarray, -) -> np.ndarray: - num_envs = command.shape[0] - n_action = noisy_joint_pos_rel.shape[1] - actor_obs = np.empty((num_envs, actor_obs_dim), dtype=get_global_dtype()) - offset = 0 - actor_obs[:, offset : offset + command.shape[1]] = command - offset += command.shape[1] - actor_obs[:, offset : offset + 3] = motion_anchor_pos_b - offset += 3 - actor_obs[:, offset : offset + 6] = motion_anchor_ori_b - offset += 6 - actor_obs[:, offset : offset + 3] = noisy_linvel - offset += 3 - actor_obs[:, offset : offset + 3] = noisy_gyro - offset += 3 - actor_obs[:, offset : offset + n_action] = noisy_joint_pos_rel - offset += n_action - actor_obs[:, offset : offset + n_action] = noisy_dof_vel - offset += n_action - actor_obs[:, offset : offset + n_action] = last_actions - return actor_obs - - -def build_mimic_actor_obs( - *, - command: np.ndarray, - motion_anchor_ori_b: np.ndarray, - noisy_gyro: np.ndarray, - noisy_joint_pos_rel: np.ndarray, - noisy_dof_vel: np.ndarray, - last_actions: np.ndarray, -) -> np.ndarray: - """unitree_rl_lab mimic deploy actor layout: 2n + 6 + 3 + n + n + n.""" - return np.concatenate( - [ - command, - motion_anchor_ori_b, - noisy_gyro, - noisy_joint_pos_rel, - noisy_dof_vel, - last_actions, - ], - axis=1, - dtype=get_global_dtype(), - ) - def write_body_pos_in_anchor_frame( anchor_pos: np.ndarray, @@ -117,9 +18,8 @@ def write_body_pos_in_anchor_frame( ay = anchor_quat[:, None, 2] az = anchor_quat[:, None, 3] - num_envs, n_body = body_pos.shape[:2] - rel_pos = body_vec_error[:num_envs, :n_body] - + num_envs, num_bodies = body_pos.shape[:2] + rel_pos = body_vec_error[:num_envs, :num_bodies] vx = rel_pos[..., 0] vy = rel_pos[..., 1] vz = rel_pos[..., 2] @@ -163,115 +63,4 @@ def write_body_ori6_in_anchor_frame( out[..., 5] = 2 * (ry * rz + rw * rx) -def compute_obs( - env: Any, - info: dict, - motion_data: Any, - linvel: np.ndarray, - gyro: np.ndarray, - dof_pos: np.ndarray, - dof_vel: np.ndarray, - robot_body_pos_w: np.ndarray, - robot_body_quat_w: np.ndarray, -) -> dict[str, np.ndarray]: - """Compute observations as dict with actor and critic groups.""" - num_envs = linvel.shape[0] - dtype = get_global_dtype() - n_action = dof_pos.shape[1] - n_body = env._n_motion_bodies - - # Get anchor states - anchor_pos_w = motion_data.body_pos_w[:, env.anchor_body_idx] - anchor_quat_w = motion_data.body_quat_w[:, env.anchor_body_idx] - robot_anchor_pos_w = robot_body_pos_w[:, env.anchor_body_idx] - robot_anchor_quat_w = robot_body_quat_w[:, env.anchor_body_idx] - - # Motion anchor pose in robot frame - if num_envs == env._num_envs: - motion_anchor_pos_b = env._motion_anchor_pos_b - motion_anchor_ori_b = env._motion_anchor_ori_b - joint_pos_rel = env._joint_pos_rel - zero_actions = env._zero_actions - else: - motion_anchor_pos_b = np.empty((num_envs, 3), dtype=dtype) - motion_anchor_ori_b = np.empty((num_envs, 6), dtype=dtype) - joint_pos_rel = np.empty((num_envs, n_action), dtype=dtype) - zero_actions = np.zeros((num_envs, n_action), dtype=dtype) - np_write_relative_anchor_transform_pos_rot6d( - robot_anchor_pos_w, - robot_anchor_quat_w, - anchor_pos_w, - anchor_quat_w, - motion_anchor_pos_b, - motion_anchor_ori_b, - ) - - # Joint positions and velocities - bias = info.get("default_dof_pos_bias") - effective_default = env.default_angles + bias if bias is not None else env.default_angles - np.subtract(dof_pos, effective_default, out=joint_pos_rel) - last_actions = info.get("current_actions") - if not isinstance(last_actions, np.ndarray): - last_actions = zero_actions - - if num_envs == env._num_envs: - command = env._motion_command - else: - command = np.empty((num_envs, n_action * 2), dtype=dtype) - command[:, :n_action] = motion_data.joint_pos - command[:, n_action : n_action * 2] = motion_data.joint_vel - - # Per-step observation noise on sensor channels (actor only). - # Critic uses the clean originals — asymmetric actor–critic contract. - noise_cfg = env._cfg.noise_config - noise_enabled = noise_cfg.level > 0.0 - if noise_enabled: - linvel_actor = env._obs_noise(linvel, noise_cfg.scale_linvel) - gyro_actor = env._obs_noise(gyro, noise_cfg.scale_gyro) - joint_pos_actor = env._obs_noise(joint_pos_rel, noise_cfg.scale_joint_angle) - dof_vel_actor = env._obs_noise(dof_vel, noise_cfg.scale_joint_vel) - else: - linvel_actor = linvel - gyro_actor = gyro - joint_pos_actor = joint_pos_rel - dof_vel_actor = dof_vel - - # Actor observations (noisy proprioception) - actor_obs = env._build_actor_obs( - command=command, - motion_anchor_pos_b=motion_anchor_pos_b, - motion_anchor_ori_b=motion_anchor_ori_b, - noisy_linvel=linvel_actor, - noisy_gyro=gyro_actor, - noisy_joint_pos_rel=joint_pos_actor, - noisy_dof_vel=dof_vel_actor, - last_actions=last_actions, - ) - - # Critic observations (clean proprioception + privileged body transforms) - critic_obs = np.empty((num_envs, env._critic_obs_width), dtype=dtype) - offset = 0 - critic_obs[:, offset : offset + command.shape[1]] = command - offset += command.shape[1] - critic_obs[:, offset : offset + 3] = motion_anchor_pos_b - offset += 3 - critic_obs[:, offset : offset + 6] = motion_anchor_ori_b - offset += 6 - critic_obs[:, offset : offset + 3] = linvel - offset += 3 - critic_obs[:, offset : offset + 3] = gyro - offset += 3 - critic_obs[:, offset : offset + n_action] = joint_pos_rel - offset += n_action - critic_obs[:, offset : offset + n_action] = dof_vel - offset += n_action - critic_obs[:, offset : offset + n_action] = last_actions - offset += n_action - robot_body_pos_b = critic_obs[:, offset : offset + n_body * 3].reshape(num_envs, n_body, 3) - env._write_body_pos_in_anchor_frame( - robot_anchor_pos_w, robot_anchor_quat_w, robot_body_pos_w, robot_body_pos_b - ) - offset += n_body * 3 - robot_body_ori_b = critic_obs[:, offset : offset + n_body * 6].reshape(num_envs, n_body, 6) - env._write_body_ori6_in_anchor_frame(robot_anchor_quat_w, robot_body_quat_w, robot_body_ori_b) - return {"obs": actor_obs, "critic": critic_obs} +__all__ = ["write_body_ori6_in_anchor_frame", "write_body_pos_in_anchor_frame"] diff --git a/src/unilab/tasks/motion_tracking/common/reset.py b/src/unilab/tasks/motion_tracking/common/reset.py deleted file mode 100644 index 4959bfc47..000000000 --- a/src/unilab/tasks/motion_tracking/common/reset.py +++ /dev/null @@ -1,83 +0,0 @@ -"""Shared reset-state construction for motion tracking.""" - -from __future__ import annotations - -from typing import Any - -import numpy as np - -from unilab.dtype_config import get_global_dtype -from unilab.utils.geometry import np_sample_uniform -from unilab.utils.rotation import np_quat_apply, np_quat_from_euler_xyz, np_quat_inv, np_quat_mul - -from .motion_loader import MotionData - - -def build_motion_reference_state( - env: Any, env_ids: np.ndarray, motion_data: MotionData -) -> tuple[np.ndarray, np.ndarray]: - dtype = get_global_dtype() - num_reset = len(env_ids) - - root_pos = motion_data.body_pos_w[:, 0].copy() - root_ori = motion_data.body_quat_w[:, 0].copy() - root_lin_vel = motion_data.body_lin_vel_w[:, 0].copy() - root_ang_vel = motion_data.body_ang_vel_w[:, 0].copy() - joint_pos = motion_data.joint_pos.copy() - joint_vel = motion_data.joint_vel.copy() - - pose_rand = env.cfg.pose_randomization - pose_ranges = [ - (pose_rand.x[0], pose_rand.x[1]), - (pose_rand.y[0], pose_rand.y[1]), - (pose_rand.z[0], pose_rand.z[1]), - (pose_rand.roll[0], pose_rand.roll[1]), - (pose_rand.pitch[0], pose_rand.pitch[1]), - (pose_rand.yaw[0], pose_rand.yaw[1]), - ] - pose_samples = np.array( - [[np.random.uniform(low, high) for low, high in pose_ranges] for _ in range(num_reset)], - dtype=dtype, - ) - root_pos += pose_samples[:, 0:3] - root_ori = np_quat_mul( - np_quat_from_euler_xyz(pose_samples[:, 3], pose_samples[:, 4], pose_samples[:, 5]), - root_ori, - ) - - vel_rand = env.cfg.velocity_randomization - vel_ranges = [ - (vel_rand.x[0], vel_rand.x[1]), - (vel_rand.y[0], vel_rand.y[1]), - (vel_rand.z[0], vel_rand.z[1]), - (vel_rand.roll[0], vel_rand.roll[1]), - (vel_rand.pitch[0], vel_rand.pitch[1]), - (vel_rand.yaw[0], vel_rand.yaw[1]), - ] - vel_samples = np.array( - [[np.random.uniform(low, high) for low, high in vel_ranges] for _ in range(num_reset)], - dtype=dtype, - ) - root_lin_vel += vel_samples[:, :3] - root_ang_vel += vel_samples[:, 3:] - - joint_pos += np_sample_uniform( - env.cfg.joint_position_range[0], - env.cfg.joint_position_range[1], - joint_pos.shape, - dtype=np.float32, - ) - joint_range = env._get_joint_range() - if joint_range is not None: - joint_pos = np.clip(joint_pos, joint_range[:, 0], joint_range[:, 1]) - - qpos = np.tile(env._init_qpos, (num_reset, 1)) - qvel = np.tile(env._init_qvel, (num_reset, 1)) - qpos[:, 0:3] = root_pos - qpos[:, 3:7] = root_ori - qpos[:, 7:] = joint_pos - - qvel[:, 0:3] = root_lin_vel - qvel[:, 3:6] = np_quat_apply(np_quat_inv(root_ori), root_ang_vel) - qvel[:, 6:] = joint_vel - return qpos, qvel diff --git a/src/unilab/tasks/motion_tracking/common/rewards.py b/src/unilab/tasks/motion_tracking/common/rewards.py deleted file mode 100644 index 36a169c02..000000000 --- a/src/unilab/tasks/motion_tracking/common/rewards.py +++ /dev/null @@ -1,386 +0,0 @@ -"""Shared reward configuration and functions for motion tracking. - -Reward terms are plain module-level callables ``fn(ctx: RewardContext) -> np.ndarray`` -mirroring :mod:`unilab.tasks.locomotion.common.rewards`. Robot-specific terms that -live on env subclasses (box-object / joint-effort terms) are stored in the same -``_reward_fns`` dispatch table as bound methods and are called with the same -``ctx`` argument. - -``RewardContext`` intentionally carries the environment's *preallocated scratch -buffers* (``body_vec_error``, ``joint_error``, ... ``undesired_contact_mask``). -These buffers are env-owned so the hot path runs with zero per-step allocations; -the term functions write into them in place. The op order, constants, clip bounds, -and in-place ``out=`` usage are load-bearing for stable numeric behavior. -""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from typing import Any, Callable, Mapping - -import numpy as np - - -@dataclass -class RewardConfig: - """Reward configuration for motion tracking.""" - - scales: dict[str, float] = field( - default_factory=lambda: { - "motion_global_root_pos": 0.5, - "motion_global_root_ori": 0.5, - "motion_body_pos": 1.0, - "motion_body_ori": 1.0, - "motion_body_lin_vel": 1.0, - "motion_body_ang_vel": 1.0, - "motion_ee_body_pos_z": 0.0, - "motion_joint_pos": 0.0, - "motion_joint_vel": 0.0, - "action_rate_l2": -0.1, - "joint_limit": -10.0, - } - ) - # Standard deviations for exponential rewards - std_root_pos: float = 0.3 - std_root_ori: float = 0.4 - std_body_pos: float = 0.3 - std_body_ori: float = 0.4 - std_body_lin_vel: float = 1.0 - std_body_ang_vel: float = 3.14 - std_joint_pos: float = 0.2 - std_joint_vel: float = 1.0 - - -@dataclass -class RewardContext: - """Bundle of everything the motion-tracking reward functions may read. - - Built once per ``_compute_reward`` call. The array fields prefixed as - buffers (``env_error``, ``reward_term``, ``weighted_reward``, ...) are the - environment's preallocated scratch arrays; reward terms write into them in - place to keep the hot path allocation-free and numerically stable. - """ - - # ── semantic inputs ────────────────────────────────────────────── - info: dict - motion_data: Any = None - robot_body_pos_w: np.ndarray | None = None - robot_body_quat_w: np.ndarray | None = None - robot_body_lin_vel_w: np.ndarray | None = None - robot_body_ang_vel_w: np.ndarray | None = None - ref_body_pos_w: np.ndarray | None = None # env.body_pos_relative_w - ref_body_quat_w: np.ndarray | None = None # env.body_quat_relative_w - dof_pos: np.ndarray | None = None - dof_vel: np.ndarray | None = None - - # ── config-derived scalars / indices ──────────────────────────── - reward_config: Any = None - anchor_body_idx: int = 0 - ee_body_indices: np.ndarray | None = None - undesired_contact_body_indices: np.ndarray | None = None - joint_lower: np.ndarray | None = None - joint_upper: np.ndarray | None = None - undesired_contact_z_threshold: float = 0.0 - num_envs: int = 0 - - # ── env-owned scratch buffers (zero-alloc hot path) ────────────── - body_vec_error: np.ndarray | None = None - joint_error: np.ndarray | None = None - joint_error_upper: np.ndarray | None = None - env_error: np.ndarray | None = None - env_error2: np.ndarray | None = None - reward_term: np.ndarray | None = None - weighted_reward: np.ndarray | None = None - quat_error_w: np.ndarray | None = None - quat_error_x: np.ndarray | None = None - ee_pos_error_z: np.ndarray | None = None - undesired_contact_mask: np.ndarray | None = None - - -# ── buffered math helpers ──────────────────────────────────────────── - - -def _required_array(value: np.ndarray | None, name: str) -> np.ndarray: - if value is None: - raise RuntimeError(f"RewardContext.{name} is required for this reward term") - return value - - -def _mean_body_xyz_squared_error( - ctx: RewardContext, reference: np.ndarray, actual: np.ndarray -) -> np.ndarray: - vec_error = _required_array(ctx.body_vec_error, "body_vec_error") - env_error = _required_array(ctx.env_error, "env_error") - tmp_error = _required_array(ctx.reward_term, "reward_term") - np.subtract(reference[..., 0], actual[..., 0], out=vec_error[..., 0]) - np.square(vec_error[..., 0], out=vec_error[..., 0]) - np.sum(vec_error[..., 0], axis=1, out=env_error) - np.subtract(reference[..., 1], actual[..., 1], out=vec_error[..., 1]) - np.square(vec_error[..., 1], out=vec_error[..., 1]) - np.sum(vec_error[..., 1], axis=1, out=tmp_error) - env_error += tmp_error - np.subtract(reference[..., 2], actual[..., 2], out=vec_error[..., 2]) - np.square(vec_error[..., 2], out=vec_error[..., 2]) - np.sum(vec_error[..., 2], axis=1, out=tmp_error) - env_error += tmp_error - env_error /= reference.shape[1] - return env_error - - -def _quat_error_magnitude_squared_body( - ctx: RewardContext, q1: np.ndarray, q2: np.ndarray -) -> np.ndarray: - rel_w = _required_array(ctx.quat_error_w, "quat_error_w") - rel_x = _required_array(ctx.quat_error_x, "quat_error_x") - # Motion/backend quaternions are unit quaternions, so the relative - # rotation angle only needs abs(dot(q1, q2)). - np.multiply(q1[..., 0], q2[..., 0], out=rel_w) - np.multiply(q1[..., 1], q2[..., 1], out=rel_x) - rel_w += rel_x - np.multiply(q1[..., 2], q2[..., 2], out=rel_x) - rel_w += rel_x - np.multiply(q1[..., 3], q2[..., 3], out=rel_x) - rel_w += rel_x - np.abs(rel_w, out=rel_w) - np.clip(rel_w, 0.0, 1.0, out=rel_w) - np.arccos(rel_w, out=rel_x) - rel_x *= 2.0 - np.square(rel_x, out=rel_x) - return rel_x - - -def _exp_reward_from_error(ctx: RewardContext, error: np.ndarray, std: float) -> np.ndarray: - out = _required_array(ctx.reward_term, "reward_term") - np.divide(error, -(std**2), out=out) - np.exp(out, out=out) - return out - - -# ── reward terms ───────────────────────────────────────────────────── - - -def motion_global_root_pos(ctx: RewardContext) -> np.ndarray: - motion_data = ctx.motion_data - robot_body_pos_w = _required_array(ctx.robot_body_pos_w, "robot_body_pos_w") - anchor_pos_w = motion_data.body_pos_w[:, ctx.anchor_body_idx] - robot_anchor_pos_w = robot_body_pos_w[:, ctx.anchor_body_idx] - error = _required_array(ctx.env_error, "env_error") - reward_term = _required_array(ctx.reward_term, "reward_term") - np.subtract(anchor_pos_w[:, 0], robot_anchor_pos_w[:, 0], out=error) - np.square(error, out=error) - np.subtract(anchor_pos_w[:, 1], robot_anchor_pos_w[:, 1], out=reward_term) - np.square(reward_term, out=reward_term) - error += reward_term - np.subtract(anchor_pos_w[:, 2], robot_anchor_pos_w[:, 2], out=reward_term) - np.square(reward_term, out=reward_term) - error += reward_term - return _exp_reward_from_error(ctx, error, ctx.reward_config.std_root_pos) - - -def motion_global_root_ori(ctx: RewardContext) -> np.ndarray: - motion_data = ctx.motion_data - robot_body_quat_w = _required_array(ctx.robot_body_quat_w, "robot_body_quat_w") - anchor_quat_w = motion_data.body_quat_w[:, ctx.anchor_body_idx] - robot_anchor_quat_w = robot_body_quat_w[:, ctx.anchor_body_idx] - env_error = _required_array(ctx.env_error, "env_error") - reward_term = _required_array(ctx.reward_term, "reward_term") - np.multiply(anchor_quat_w[:, 0], robot_anchor_quat_w[:, 0], out=env_error) - np.multiply(anchor_quat_w[:, 1], robot_anchor_quat_w[:, 1], out=reward_term) - env_error += reward_term - np.multiply(anchor_quat_w[:, 2], robot_anchor_quat_w[:, 2], out=reward_term) - env_error += reward_term - np.multiply(anchor_quat_w[:, 3], robot_anchor_quat_w[:, 3], out=reward_term) - env_error += reward_term - np.abs(env_error, out=env_error) - np.clip(env_error, 0.0, 1.0, out=env_error) - np.arccos(env_error, out=env_error) - env_error *= 2.0 - np.square(env_error, out=env_error) - return _exp_reward_from_error(ctx, env_error, ctx.reward_config.std_root_ori) - - -def motion_body_pos(ctx: RewardContext) -> np.ndarray: - ref_body_pos_w = _required_array(ctx.ref_body_pos_w, "ref_body_pos_w") - robot_body_pos_w = _required_array(ctx.robot_body_pos_w, "robot_body_pos_w") - error = _mean_body_xyz_squared_error(ctx, ref_body_pos_w, robot_body_pos_w) - return _exp_reward_from_error(ctx, error, ctx.reward_config.std_body_pos) - - -def motion_body_ori(ctx: RewardContext) -> np.ndarray: - ref_body_quat_w = _required_array(ctx.ref_body_quat_w, "ref_body_quat_w") - robot_body_quat_w = _required_array(ctx.robot_body_quat_w, "robot_body_quat_w") - error = _quat_error_magnitude_squared_body(ctx, ref_body_quat_w, robot_body_quat_w) - env_error = _required_array(ctx.env_error, "env_error") - np.sum(error, axis=-1, out=env_error) - env_error /= error.shape[1] - return _exp_reward_from_error(ctx, env_error, ctx.reward_config.std_body_ori) - - -def motion_body_lin_vel(ctx: RewardContext) -> np.ndarray: - motion_data = ctx.motion_data - robot_body_lin_vel_w = _required_array(ctx.robot_body_lin_vel_w, "robot_body_lin_vel_w") - error = _mean_body_xyz_squared_error(ctx, motion_data.body_lin_vel_w, robot_body_lin_vel_w) - return _exp_reward_from_error(ctx, error, ctx.reward_config.std_body_lin_vel) - - -def motion_body_ang_vel(ctx: RewardContext) -> np.ndarray: - motion_data = ctx.motion_data - robot_body_ang_vel_w = _required_array(ctx.robot_body_ang_vel_w, "robot_body_ang_vel_w") - error = _mean_body_xyz_squared_error(ctx, motion_data.body_ang_vel_w, robot_body_ang_vel_w) - return _exp_reward_from_error(ctx, error, ctx.reward_config.std_body_ang_vel) - - -def motion_ee_body_pos_z(ctx: RewardContext) -> np.ndarray: - ref_body_pos_w = _required_array(ctx.ref_body_pos_w, "ref_body_pos_w") - robot_body_pos_w = _required_array(ctx.robot_body_pos_w, "robot_body_pos_w") - ee_body_indices = _required_array(ctx.ee_body_indices, "ee_body_indices") - ee_pos_error_z = _required_array(ctx.ee_pos_error_z, "ee_pos_error_z") - env_error = _required_array(ctx.env_error, "env_error") - np.subtract( - ref_body_pos_w[:, ee_body_indices, 2], - robot_body_pos_w[:, ee_body_indices, 2], - out=ee_pos_error_z, - ) - np.square(ee_pos_error_z, out=ee_pos_error_z) - np.sum(ee_pos_error_z, axis=-1, out=env_error) - env_error /= ee_pos_error_z.shape[1] - return _exp_reward_from_error(ctx, env_error, ctx.reward_config.std_body_pos) - - -def motion_joint_pos(ctx: RewardContext) -> np.ndarray: - motion_data = ctx.motion_data - dof_pos = _required_array(ctx.dof_pos, "dof_pos") - joint_error = _required_array(ctx.joint_error, "joint_error") - env_error = _required_array(ctx.env_error, "env_error") - np.subtract(motion_data.joint_pos, dof_pos, out=joint_error) - np.square(joint_error, out=joint_error) - np.sum(joint_error, axis=1, out=env_error) - env_error /= dof_pos.shape[1] - return _exp_reward_from_error(ctx, env_error, ctx.reward_config.std_joint_pos) - - -def motion_joint_vel(ctx: RewardContext) -> np.ndarray: - motion_data = ctx.motion_data - dof_vel = _required_array(ctx.dof_vel, "dof_vel") - joint_error = _required_array(ctx.joint_error, "joint_error") - env_error = _required_array(ctx.env_error, "env_error") - np.subtract(motion_data.joint_vel, dof_vel, out=joint_error) - np.square(joint_error, out=joint_error) - np.sum(joint_error, axis=1, out=env_error) - env_error /= dof_vel.shape[1] - return _exp_reward_from_error(ctx, env_error, ctx.reward_config.std_joint_vel) - - -def undesired_contacts(ctx: RewardContext) -> np.ndarray: - robot_body_pos_w = _required_array(ctx.robot_body_pos_w, "robot_body_pos_w") - undesired_contact_body_indices = _required_array( - ctx.undesired_contact_body_indices, "undesired_contact_body_indices" - ) - undesired_contact_mask = _required_array(ctx.undesired_contact_mask, "undesired_contact_mask") - env_error = _required_array(ctx.env_error, "env_error") - body_z = robot_body_pos_w[:, undesired_contact_body_indices, 2] - np.less( - body_z, - ctx.undesired_contact_z_threshold, - out=undesired_contact_mask, - ) - np.sum(undesired_contact_mask, axis=-1, out=env_error) - return env_error - - -def action_rate_l2(ctx: RewardContext) -> np.ndarray: - info = ctx.info - joint_error = _required_array(ctx.joint_error, "joint_error") - env_error = _required_array(ctx.env_error, "env_error") - np.subtract(info["current_actions"], info["last_actions"], out=joint_error) - np.square(joint_error, out=joint_error) - np.sum(joint_error, axis=1, out=env_error) - return env_error - - -def joint_limit(ctx: RewardContext) -> np.ndarray: - dof_pos = _required_array(ctx.dof_pos, "dof_pos") - reward_term = _required_array(ctx.reward_term, "reward_term") - joint_error = _required_array(ctx.joint_error, "joint_error") - joint_error_upper = _required_array(ctx.joint_error_upper, "joint_error_upper") - lower = ctx.joint_lower - upper = ctx.joint_upper - if lower is None or upper is None: - reward_term.fill(0.0) - return reward_term - - # Compute violation - np.subtract(lower, dof_pos, out=joint_error) - np.maximum(joint_error, 0, out=joint_error) - np.subtract(dof_pos, upper, out=joint_error_upper) - np.maximum(joint_error_upper, 0, out=joint_error_upper) - joint_error += joint_error_upper - np.square(joint_error, out=joint_error) - np.sum(joint_error, axis=1, out=reward_term) - return reward_term - - -def build_reward_functions() -> dict[str, Callable[[RewardContext], np.ndarray]]: - """Return the robot-agnostic reward-term dispatch table. - - Keys define the canonical reward-term dispatch order. - """ - return { - "motion_global_root_pos": motion_global_root_pos, - "motion_global_root_ori": motion_global_root_ori, - "motion_body_pos": motion_body_pos, - "motion_body_ori": motion_body_ori, - "motion_body_lin_vel": motion_body_lin_vel, - "motion_body_ang_vel": motion_body_ang_vel, - "motion_ee_body_pos_z": motion_ee_body_pos_z, - "motion_joint_pos": motion_joint_pos, - "motion_joint_vel": motion_joint_vel, - "action_rate_l2": action_rate_l2, - "joint_limit": joint_limit, - "undesired_contacts": undesired_contacts, - } - - -def compute_reward( - ctx: RewardContext, - *, - active_reward_fns: Mapping[str, Callable[[RewardContext], np.ndarray]], - all_reward_fns: Mapping[str, Callable[[RewardContext], np.ndarray]], - scales: Mapping[str, float], - ctrl_dt: float, - enable_log: bool, -) -> np.ndarray: - """Reduce ``scales × fns(ctx)`` into the per-env reward (in place, zero-alloc). - - Uses ``ctx.env_error2`` as the reward accumulator and ``ctx.weighted_reward`` - as the per-term scratch. Logs per-term - means into ``ctx.info["log"]`` every 4th step, then scales by ``ctrl_dt``. - """ - reward = _required_array(ctx.env_error2, "env_error2") - reward.fill(0.0) - - info = ctx.info - step_count = info.get("steps") - should_log = enable_log and ( - int(step_count[0]) % 4 == 0 if isinstance(step_count, np.ndarray) else True - ) - log = {} if should_log else info.get("log", {}) - - for name, scale in scales.items(): - if scale == 0: - continue - reward_fn = active_reward_fns.get(name) - if reward_fn is None: - if should_log and name in all_reward_fns: - log[f"reward/{name}"] = 0.0 - continue - rew = reward_fn(ctx) - weighted_rew = _required_array(ctx.weighted_reward, "weighted_reward") - np.multiply(rew, scale, out=weighted_rew) - reward += weighted_rew - if should_log: - log[f"reward/{name}"] = float(np.sum(weighted_rew) / weighted_rew.size) - - info["log"] = log - reward *= ctrl_dt - return reward diff --git a/src/unilab/tasks/motion_tracking/common/terminations.py b/src/unilab/tasks/motion_tracking/common/terminations.py deleted file mode 100644 index 6a9ece326..000000000 --- a/src/unilab/tasks/motion_tracking/common/terminations.py +++ /dev/null @@ -1,69 +0,0 @@ -"""Shared termination computation for motion tracking.""" - -from __future__ import annotations - -from typing import Any - -import numpy as np - -from unilab.utils.geometry import np_gravity_z_in_body_from_quat - - -def compute_terminations( - env: Any, - motion_data: Any, - robot_body_pos_w: np.ndarray, - robot_body_quat_w: np.ndarray, -) -> np.ndarray: - """Compute termination conditions (writes into ``env._terminated``).""" - terminated = env._terminated - terminated.fill(False) - - # Anchor position error (Z-axis only) - anchor_pos_w = motion_data.body_pos_w[:, env.anchor_body_idx] - robot_anchor_pos_w = robot_body_pos_w[:, env.anchor_body_idx] - np.subtract(anchor_pos_w[:, 2], robot_anchor_pos_w[:, 2], out=env._env_error) - np.abs(env._env_error, out=env._env_error) - np.greater(env._env_error, env._cfg.anchor_pos_z_threshold, out=env._env_bool) - terminated |= env._env_bool - - # Anchor orientation error (gravity direction). The gravity-z difference - # is bounded by 2 for unit quaternions, so huge thresholds disable this - # termination without doing the per-step math. - if env._cfg.anchor_ori_threshold < 2.0: - anchor_quat_w = motion_data.body_quat_w[:, env.anchor_body_idx] - robot_anchor_quat_w = robot_body_quat_w[:, env.anchor_body_idx] - motion_gravity_z_b = np_gravity_z_in_body_from_quat(anchor_quat_w) - robot_gravity_z_b = np_gravity_z_in_body_from_quat(robot_anchor_quat_w) - np.subtract(motion_gravity_z_b, robot_gravity_z_b, out=env._env_error) - np.abs(env._env_error, out=env._env_error) - np.greater(env._env_error, env._cfg.anchor_ori_threshold, out=env._env_bool) - terminated |= env._env_bool - - # End-effector position error (Z-axis only) - if env._has_ee_body_indices: - np.subtract( - env.body_pos_relative_w[:, env.ee_body_indices, 2], - robot_body_pos_w[:, env.ee_body_indices, 2], - out=env._ee_pos_error_z, - ) - np.abs(env._ee_pos_error_z, out=env._ee_pos_error_z) - np.greater( - env._ee_pos_error_z, - env._cfg.ee_body_pos_z_threshold, - out=env._ee_terminated, - ) - np.logical_or.reduce(env._ee_terminated, axis=1, out=env._env_bool) - terminated |= env._env_bool - - if env._cfg.terminate_on_undesired_contacts and env._has_undesired_contact_body_indices: - body_z = robot_body_pos_w[:, env.undesired_contact_body_indices, 2] - np.less( - body_z, - env._cfg.undesired_contact_z_threshold, - out=env._undesired_contact_mask, - ) - np.logical_or.reduce(env._undesired_contact_mask, axis=-1, out=env._env_bool) - terminated |= env._env_bool - - return terminated diff --git a/src/unilab/tasks/motion_tracking/common/tracking.py b/src/unilab/tasks/motion_tracking/common/tracking.py deleted file mode 100644 index 99be8a52c..000000000 --- a/src/unilab/tasks/motion_tracking/common/tracking.py +++ /dev/null @@ -1,539 +0,0 @@ -"""Robot-agnostic motion-tracking engine. - -Holds :class:`MotionTrackingEnv` (the imitation engine, inheriting the shared -``G1BaseEnv`` locomotion base (``motion_tracking.common.g1_base``)) and :class:`MotionTrackingDeployEnv` (the -unitree_rl_lab mimic actor variant). Per-concern math lives in the owner -modules (``rewards`` / ``observations`` / ``terminations`` / ``transforms`` / -``reset`` / ``domain_randomization``); the engine keeps only the -thin polymorphic method surface and per-step orchestration. -""" - -from __future__ import annotations - -from typing import Any - -import numpy as np - -from unilab.base.backend import create_backend, env_backend_kwargs -from unilab.base.np_env import NpEnvState -from unilab.dtype_config import get_global_dtype -from unilab.tasks.motion_tracking.common.g1_base import G1BaseEnv - -from . import observations -from .config import MotionTrackingCfg, MotionTrackingDeployEnvCfg -from .domain_randomization import MotionTrackingDomainRandomizationProvider -from .motion_loader import MotionData, MotionLoader, MotionSampler -from .reset import build_motion_reference_state -from .rewards import ( - RewardContext, - build_reward_functions, - compute_reward, -) -from .terminations import compute_terminations -from .transforms import update_relative_transforms - - -class MotionTrackingEnv(G1BaseEnv): - """Motion Tracking Environment (robot-agnostic imitation engine).""" - - _cfg: MotionTrackingCfg # pyright: ignore[reportIncompatibleVariableOverride] - - def __init__(self, cfg: MotionTrackingCfg, num_envs=1, backend_type="mujoco"): - if not cfg.motion_file: - raise ValueError("motion_file must be specified in config") - - backend = create_backend( - backend_type, - cfg.scene, - num_envs, - cfg.sim_dt, - base_name=cfg.asset.base_name, - push_body_name=cfg.domain_rand.push_body_name, - add_body_sensors=True, - **env_backend_kwargs(cfg), - ) - super().__init__(cfg, backend, num_envs) - - # Resolve body IDs for backend querying and motion-file indexing. - self.body_ids = self._backend.get_body_ids(cfg.body_names) - motion_body_ids = self._backend.get_motion_body_ids(cfg.body_names) - - self.anchor_body_idx = cfg.body_names.index(cfg.anchor_body_name) - - # Get end-effector body indices for termination - self.ee_body_indices = np.array( - [cfg.body_names.index(name) for name in cfg.ee_body_names], dtype=np.int32 - ) - self._has_ee_body_indices = bool(self.ee_body_indices.size) - - # Get non-EE body indices for undesired contact penalty - ee_set = set(cfg.ee_body_names) - self.undesired_contact_body_indices = np.array( - [i for i, name in enumerate(cfg.body_names) if name not in ee_set], - dtype=np.int32, - ) - self._has_undesired_contact_body_indices = bool(self.undesired_contact_body_indices.size) - - # Load motion data - self.motion_loader = MotionLoader(cfg.motion_file, body_indices=motion_body_ids) - self.motion_sampler = MotionSampler( - self.motion_loader, - mode=cfg.sampling_mode, - num_envs=num_envs, - start_ratio=cfg.sampling_start_ratio, - ) - needs_kp_kd = cfg.domain_rand.randomize_kp or cfg.domain_rand.randomize_kd - needs_friction = getattr(cfg.domain_rand, "randomize_geom_friction", False) - base_kp = base_kd = None - if needs_kp_kd: - base_kp, base_kd = backend.get_actuator_gains() - base_geom_friction = None - foot_geom_ids = None - if needs_friction: - import re as _re - - base_geom_friction = backend.get_geom_friction() - geom_names = backend.get_geom_names() - pattern = _re.compile(cfg.domain_rand.friction_geom_pattern) - foot_geom_ids = np.asarray( - [i for i, name in enumerate(geom_names) if name and pattern.match(name)], - dtype=np.int64, - ) - if foot_geom_ids.size == 0: - raise ValueError( - "friction_geom_pattern " - f"'{cfg.domain_rand.friction_geom_pattern}' did not match any geom" - ) - dr_provider = MotionTrackingDomainRandomizationProvider( - base_kp=base_kp, - base_kd=base_kd, - base_geom_friction=base_geom_friction, - foot_geom_ids=foot_geom_ids, - ) - self._init_domain_randomization(dr_provider) - - dtype = get_global_dtype() - n_body = len(cfg.body_names) - self._n_motion_bodies = n_body - self._actor_obs_width = self._actor_obs_dim(self._num_action) - self._critic_base_obs_width = self._critic_base_obs_dim(self._num_action) - self._critic_obs_width = self._critic_base_obs_width + n_body * 9 - self._copy_body_state_w = self._backend.copy_body_state_w - - # Buffers for relative body transforms - self.body_pos_relative_w = np.zeros((num_envs, n_body, 3), dtype=dtype) - self.body_quat_relative_w = np.zeros((num_envs, n_body, 4), dtype=dtype) - self.body_quat_relative_w[:, :, 0] = 1.0 # Initialize to identity quaternion - self._motion_data_buffer = ( - self.motion_loader.make_motion_data_buffer(num_envs) - if hasattr(self.motion_loader, "make_motion_data_buffer") - else None - ) - self._zero_actions = np.zeros((num_envs, self._num_action), dtype=dtype) - self._joint_range = self._backend.get_joint_range() - if self._joint_range is not None: - self._joint_range = np.asarray(self._joint_range, dtype=dtype) - self._joint_lower = self._joint_range[:, 0] - self._joint_upper = self._joint_range[:, 1] - else: - self._joint_lower = None - self._joint_upper = None - self._delta_pos_w = np.empty((num_envs, 3), dtype=dtype) - self._delta_ori_w = np.empty((num_envs, 4), dtype=dtype) - self._motion_anchor_pos_b = np.empty((num_envs, 3), dtype=dtype) - self._motion_anchor_ori_b = np.empty((num_envs, 6), dtype=dtype) - self._motion_command = np.empty((num_envs, self._num_action * 2), dtype=dtype) - self._joint_pos_rel = np.empty((num_envs, self._num_action), dtype=dtype) - self._robot_body_pos_w = np.empty((num_envs, n_body, 3), dtype=dtype) - self._robot_body_quat_w = np.empty((num_envs, n_body, 4), dtype=dtype) - self._robot_body_lin_vel_w = np.empty((num_envs, n_body, 3), dtype=dtype) - self._robot_body_ang_vel_w = np.empty((num_envs, n_body, 3), dtype=dtype) - self._quat_error_w = np.empty((num_envs, n_body), dtype=dtype) - self._quat_error_x = np.empty((num_envs, n_body), dtype=dtype) - self._body_vec_error = np.empty((num_envs, n_body, 3), dtype=dtype) - self._body_vec_tmp = np.empty((num_envs, n_body, 3), dtype=dtype) - self._joint_error = np.empty((num_envs, self._num_action), dtype=dtype) - self._joint_error_upper = np.empty((num_envs, self._num_action), dtype=dtype) - self._env_error = np.empty((num_envs,), dtype=dtype) - self._env_error2 = np.empty((num_envs,), dtype=dtype) - self._reward_term = np.empty((num_envs,), dtype=dtype) - self._weighted_reward = np.empty((num_envs,), dtype=dtype) - self._terminated = np.empty((num_envs,), dtype=bool) - self._env_bool = np.empty((num_envs,), dtype=bool) - self._ee_pos_error_z = np.empty((num_envs, self.ee_body_indices.size), dtype=dtype) - self._ee_terminated = np.empty((num_envs, self.ee_body_indices.size), dtype=bool) - self._undesired_contact_mask = np.empty( - (num_envs, self.undesired_contact_body_indices.size), dtype=bool - ) - - self._enable_reward_log = True - self._init_reward_functions() - self._active_reward_fns = { - name: reward_fn - for name, reward_fn in self._reward_fns.items() - if self._reward_term_is_active(name) - } - self._clip_end_truncated = np.zeros((num_envs,), dtype=bool) - - def _effective_default_angles(self, env_ids: np.ndarray | None = None) -> np.ndarray: - """Return default_angles with per-episode joint-default-pos bias applied.""" - state = getattr(self, "_state", None) - if state is not None: - bias = state.info.get("default_dof_pos_bias") - if bias is not None: - if env_ids is not None: - return self.default_angles + bias[env_ids] - return self.default_angles + bias - return self.default_angles - - def apply_action(self, actions: np.ndarray, state: NpEnvState) -> np.ndarray: - state.info["last_actions"] = state.info.get("current_actions", np.zeros_like(actions)) - state.info["current_actions"] = actions - exec_actions = ( - state.info["last_actions"] - if self._cfg.control_config.simulate_action_latency - else actions - ) - bias = state.info.get("default_dof_pos_bias") - base = self.default_angles + bias if bias is not None else self.default_angles - ctrl: np.ndarray = exec_actions * self._cfg.control_config.action_scale + base - return ctrl - - def _resample_reference_state(self, env_ids: np.ndarray) -> None: - motion_frames = self.motion_sampler.sample_frames(env_ids) - motion_data = self.motion_loader.get_motion_at_frame(motion_frames) - qpos, qvel = build_motion_reference_state(self, env_ids, motion_data) - self._backend.set_state(env_ids, qpos, qvel) - - def _refresh_observation_rows( - self, obs: dict[str, np.ndarray], info: dict, env_ids: np.ndarray - ) -> None: - motion_data = self.motion_loader.get_motion_at_frame( - self.motion_sampler.current_frames[env_ids] - ) - row_ids = np.asarray(env_ids, dtype=np.intp) - linvel = self._backend.get_sensor_data_rows(self._cfg.sensor.local_linvel, row_ids) - gyro = self._backend.get_sensor_data_rows(self._cfg.sensor.gyro, row_ids) - dof_pos = self.get_dof_pos()[row_ids] - dof_vel = self.get_dof_vel()[row_ids] - robot_body_pos_w, robot_body_quat_w = self._backend.get_body_pose_w_rows( - row_ids, self.body_ids - ) - - obs_info: dict[str, Any] = {} - current_actions = info.get("current_actions") - if isinstance(current_actions, np.ndarray): - obs_info["current_actions"] = current_actions[env_ids] - obs_info["env_ids"] = env_ids - - refreshed_obs = self._compute_obs( - obs_info, - motion_data, - linvel, - gyro, - dof_pos, - dof_vel, - robot_body_pos_w, - robot_body_quat_w, - ) - for key, value in refreshed_obs.items(): - if value.shape[0] == len(env_ids): - obs[key][env_ids] = value - else: - obs[key][env_ids] = value[env_ids] - - def _get_body_pose_w(self) -> tuple[np.ndarray, np.ndarray]: - return self._backend.get_body_pose_w(self.body_ids) - - def _get_body_state_w(self) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: - copy_body_state_w = self._copy_body_state_w - if copy_body_state_w is not None: - return copy_body_state_w( - self.body_ids, - self._robot_body_pos_w, - self._robot_body_quat_w, - self._robot_body_lin_vel_w, - self._robot_body_ang_vel_w, - ) - robot_body_pos_w, robot_body_quat_w = self._get_body_pose_w() - robot_body_lin_vel_w, robot_body_ang_vel_w = self._backend.get_body_vel_w(self.body_ids) - return ( - robot_body_pos_w, - robot_body_quat_w, - robot_body_lin_vel_w, - robot_body_ang_vel_w, - ) - - def _get_joint_range(self) -> np.ndarray | None: - return self._joint_range - - def _get_current_motion(self) -> MotionData: - if self._motion_data_buffer is None: - return self.motion_sampler.get_current_motion() - return self.motion_sampler.get_current_motion(self._motion_data_buffer) - - @property - def obs_groups_spec(self) -> dict[str, int]: - return observations.obs_groups_spec(self) - - def _actor_obs_dim(self, n: int) -> int: - return observations.actor_obs_dim(n) - - def _critic_base_obs_dim(self, n: int) -> int: - return observations.critic_base_obs_dim(n) - - def _build_actor_obs( - self, - *, - command: np.ndarray, - motion_anchor_pos_b: np.ndarray, - motion_anchor_ori_b: np.ndarray, - noisy_linvel: np.ndarray, - noisy_gyro: np.ndarray, - noisy_joint_pos_rel: np.ndarray, - noisy_dof_vel: np.ndarray, - last_actions: np.ndarray, - ) -> np.ndarray: - n_action = noisy_joint_pos_rel.shape[1] - return observations.build_actor_obs( - actor_obs_dim=self._actor_obs_dim(n_action), - command=command, - motion_anchor_pos_b=motion_anchor_pos_b, - motion_anchor_ori_b=motion_anchor_ori_b, - noisy_linvel=noisy_linvel, - noisy_gyro=noisy_gyro, - noisy_joint_pos_rel=noisy_joint_pos_rel, - noisy_dof_vel=noisy_dof_vel, - last_actions=last_actions, - ) - - def _init_reward_functions(self): - self._reward_fns = build_reward_functions() - - def _reward_term_is_active(self, name: str) -> bool: - if name == "joint_limit": - return self._joint_lower is not None and self._joint_upper is not None - if name == "undesired_contacts": - return self._has_undesired_contact_body_indices - if name == "motion_ee_body_pos_z": - return self._has_ee_body_indices - return True - - def update_state(self, state: NpEnvState) -> NpEnvState: - self._clip_end_truncated.fill(False) - - # Get current motion data - motion_data = self._get_current_motion() - - # Get robot state - linvel = self.get_local_linvel() - gyro = self.get_gyro() - dof_pos = self.get_dof_pos() - dof_vel = self.get_dof_vel() - - # Get body states - ( - robot_body_pos_w, - robot_body_quat_w, - robot_body_lin_vel_w, - robot_body_ang_vel_w, - ) = self._get_body_state_w() - - # Compute relative body transforms (for observations and rewards) - self._update_relative_transforms(motion_data, robot_body_pos_w, robot_body_quat_w) - - # Compute terminations - terminated = self._compute_terminations(motion_data, robot_body_pos_w, robot_body_quat_w) - - # Compute reward - reward = self._compute_reward( - state.info, - motion_data, - robot_body_pos_w, - robot_body_quat_w, - robot_body_lin_vel_w, - robot_body_ang_vel_w, - dof_pos, - dof_vel, - ) - - # Compute observations - obs = self._compute_obs( - state.info, - motion_data, - linvel, - gyro, - dof_pos, - dof_vel, - robot_body_pos_w, - robot_body_quat_w, - ) - - # Update failure statistics for adaptive sampling - self.motion_sampler.update_failure_stats(terminated) - - # Advance motion frames - done_env_ids = self.motion_sampler.step() - if len(done_env_ids) > 0: - if self._cfg.truncate_on_clip_end: - self._clip_end_truncated[done_env_ids] = True - else: - # Match BeyondMimic: clip boundaries are command resampling points, not - # episode boundaries; sync the simulated robot to the new reference. - resample_env_ids = done_env_ids[~terminated[done_env_ids]] - if len(resample_env_ids) > 0: - self._resample_reference_state(resample_env_ids) - self._refresh_observation_rows(obs, state.info, resample_env_ids) - - return state.replace(obs=obs, reward=reward, terminated=terminated) - - def _compute_truncated(self, state: NpEnvState) -> np.ndarray: - truncated = super()._compute_truncated(state) - clip_end_only = getattr(self, "_env_bool", None) - if clip_end_only is None or clip_end_only.shape != (self._num_envs,): - clip_end_only = np.empty((self._num_envs,), dtype=bool) - self._env_bool = clip_end_only - np.logical_not(state.terminated, out=clip_end_only) - np.logical_and(self._clip_end_truncated, clip_end_only, out=clip_end_only) - np.logical_or(truncated, clip_end_only, out=truncated) - return truncated - - def _update_relative_transforms( - self, motion_data, robot_body_pos_w: np.ndarray, robot_body_quat_w: np.ndarray - ): - """Update relative body transforms for tracking.""" - update_relative_transforms(self, motion_data, robot_body_pos_w, robot_body_quat_w) - - def _compute_terminations( - self, - motion_data, - robot_body_pos_w: np.ndarray, - robot_body_quat_w: np.ndarray, - ) -> np.ndarray: - """Compute termination conditions.""" - return compute_terminations(self, motion_data, robot_body_pos_w, robot_body_quat_w) - - def _write_body_pos_in_anchor_frame( - self, - anchor_pos: np.ndarray, - anchor_quat: np.ndarray, - body_pos: np.ndarray, - out: np.ndarray, - ) -> None: - observations.write_body_pos_in_anchor_frame( - anchor_pos, anchor_quat, body_pos, out, body_vec_error=self._body_vec_error - ) - - def _write_body_ori6_in_anchor_frame( - self, - anchor_quat: np.ndarray, - body_quat: np.ndarray, - out: np.ndarray, - ) -> None: - observations.write_body_ori6_in_anchor_frame(anchor_quat, body_quat, out) - - def _compute_obs( - self, - info: dict, - motion_data, - linvel: np.ndarray, - gyro: np.ndarray, - dof_pos: np.ndarray, - dof_vel: np.ndarray, - robot_body_pos_w: np.ndarray, - robot_body_quat_w: np.ndarray, - ) -> dict[str, np.ndarray]: - """Compute observations as dict with actor and critic groups.""" - return observations.compute_obs( - self, - info, - motion_data, - linvel, - gyro, - dof_pos, - dof_vel, - robot_body_pos_w, - robot_body_quat_w, - ) - - def _compute_reward( - self, - info: dict, - motion_data, - robot_body_pos_w: np.ndarray, - robot_body_quat_w: np.ndarray, - robot_body_lin_vel_w: np.ndarray, - robot_body_ang_vel_w: np.ndarray, - dof_pos: np.ndarray, - dof_vel: np.ndarray, - ) -> np.ndarray: - """Compute reward.""" - ctx = RewardContext( - info=info, - motion_data=motion_data, - robot_body_pos_w=robot_body_pos_w, - robot_body_quat_w=robot_body_quat_w, - robot_body_lin_vel_w=robot_body_lin_vel_w, - robot_body_ang_vel_w=robot_body_ang_vel_w, - ref_body_pos_w=self.body_pos_relative_w, - ref_body_quat_w=self.body_quat_relative_w, - dof_pos=dof_pos, - dof_vel=dof_vel, - reward_config=self._cfg.reward_config, - anchor_body_idx=self.anchor_body_idx, - ee_body_indices=self.ee_body_indices, - undesired_contact_body_indices=self.undesired_contact_body_indices, - joint_lower=self._joint_lower, - joint_upper=self._joint_upper, - undesired_contact_z_threshold=self._cfg.undesired_contact_z_threshold, - num_envs=self._num_envs, - body_vec_error=self._body_vec_error, - joint_error=self._joint_error, - joint_error_upper=self._joint_error_upper, - env_error=self._env_error, - env_error2=self._env_error2, - reward_term=self._reward_term, - weighted_reward=self._weighted_reward, - quat_error_w=self._quat_error_w, - quat_error_x=self._quat_error_x, - ee_pos_error_z=self._ee_pos_error_z, - undesired_contact_mask=self._undesired_contact_mask, - ) - return compute_reward( - ctx, - active_reward_fns=self._active_reward_fns, - all_reward_fns=self._reward_fns, - scales=self._cfg.reward_config.scales, - ctrl_dt=self._cfg.ctrl_dt, - enable_log=self._enable_reward_log, - ) - - -class MotionTrackingDeployEnv(MotionTrackingEnv): - """Deploy-oriented motion tracking env with unitree_rl_lab mimic actor inputs.""" - - _cfg: MotionTrackingDeployEnvCfg # pyright: ignore[reportIncompatibleVariableOverride] - - def _actor_obs_dim(self, n: int) -> int: - return observations.mimic_actor_obs_dim(n) - - def _build_actor_obs( - self, - *, - command: np.ndarray, - motion_anchor_pos_b: np.ndarray, - motion_anchor_ori_b: np.ndarray, - noisy_linvel: np.ndarray, - noisy_gyro: np.ndarray, - noisy_joint_pos_rel: np.ndarray, - noisy_dof_vel: np.ndarray, - last_actions: np.ndarray, - ) -> np.ndarray: - return observations.build_mimic_actor_obs( - command=command, - motion_anchor_ori_b=motion_anchor_ori_b, - noisy_gyro=noisy_gyro, - noisy_joint_pos_rel=noisy_joint_pos_rel, - noisy_dof_vel=noisy_dof_vel, - last_actions=last_actions, - ) diff --git a/src/unilab/tasks/motion_tracking/g1/__init__.py b/src/unilab/tasks/motion_tracking/g1/__init__.py index 117e74579..040166e9b 100644 --- a/src/unilab/tasks/motion_tracking/g1/__init__.py +++ b/src/unilab/tasks/motion_tracking/g1/__init__.py @@ -1,98 +1,37 @@ -"""G1 motion-tracking tasks.""" +"""G1 motion profiles on the shared NumPy Manager-Based runtime.""" + +from unilab.base import registry +from unilab.envs import ManagerBasedRlEnvCfg, make_manager_based_rl_env -from .box_tracking import ( - G1BoxTracking23DofCfg, - G1BoxTracking23DofEnvCfg, - G1BoxTrackingCfg, - G1BoxTrackingEnv, - G1BoxTrackingEnvCfg, -) -from .flip_tracking import ( - G1ClimbTrackingCfg, - G1ClimbTrackingEnv, - G1ClimbTrackingEnvCfg, - G1FlipTracking23DofCfg, - G1FlipTracking23DofEnvCfg, - G1FlipTrackingCfg, - G1FlipTrackingEnv, - G1FlipTrackingEnvCfg, - G1WallFlipTracking23DofCfg, - G1WallFlipTracking23DofEnvCfg, - G1WallFlipTrackingCfg, - G1WallFlipTrackingEnv, - G1WallFlipTrackingEnvCfg, -) -from .flip_tracking_sac import ( - G1FlipTrackingSAC23DofCfg, - G1FlipTrackingSAC23DofEnv, - G1FlipTrackingSACCfg, - G1FlipTrackingSACEnv, - G1WallFlipTrackingSAC23DofCfg, - G1WallFlipTrackingSAC23DofEnv, - G1WallFlipTrackingSACCfg, - G1WallFlipTrackingSACEnv, -) from .motion_box_loader import BoxMotionData, BoxMotionLoader -from .tracking import ( - G1MotionTracking23DofCfg, - G1MotionTracking23DofDeployEnvCfg, - G1MotionTracking23DofEnvCfg, - G1MotionTrackingCfg, - G1MotionTrackingDeployEnv, - G1MotionTrackingDeployEnvCfg, - G1MotionTrackingEnv, - G1MotionTrackingEnvCfg, -) -from .tracking_obs import G1WBTObs23DofCfg, G1WBTObsCfg, G1WBTObsEnv -from .tracking_sac import ( - G1MotionTrackingSAC23DofCfg, - G1MotionTrackingSAC23DofEnv, - G1MotionTrackingSACCfg, - G1MotionTrackingSACEnv, + +G1_MOTION_TASKS = ( + "G1MotionTracking", + "G1MotionTrackingDeploy", + "G1MotionTracking23Dof", + "G1MotionTracking23DofDeploy", + "G1MotionTrackingSAC", + "G1MotionTrackingSAC23Dof", + "G1BoxTracking", + "G1BoxTracking23Dof", + "G1ClimbTracking", + "G1ClimbTracking23Dof", + "G1FlipTracking", + "G1FlipTracking23Dof", + "G1FlipTrackingSAC", + "G1FlipTrackingSAC23Dof", + "G1WallFlipTracking", + "G1WallFlipTracking23Dof", + "G1WallFlipTrackingSAC", + "G1WallFlipTrackingSAC23Dof", + "G1WBTObs", + "G1WBTObs23Dof", ) -__all__ = [ - "BoxMotionData", - "BoxMotionLoader", - "G1BoxTracking23DofCfg", - "G1BoxTracking23DofEnvCfg", - "G1BoxTrackingCfg", - "G1BoxTrackingEnv", - "G1BoxTrackingEnvCfg", - "G1ClimbTrackingCfg", - "G1ClimbTrackingEnv", - "G1ClimbTrackingEnvCfg", - "G1FlipTracking23DofCfg", - "G1FlipTracking23DofEnvCfg", - "G1FlipTrackingCfg", - "G1FlipTrackingEnv", - "G1FlipTrackingEnvCfg", - "G1FlipTrackingSAC23DofCfg", - "G1FlipTrackingSAC23DofEnv", - "G1FlipTrackingSACCfg", - "G1FlipTrackingSACEnv", - "G1MotionTracking23DofCfg", - "G1MotionTracking23DofDeployEnvCfg", - "G1MotionTracking23DofEnvCfg", - "G1MotionTrackingCfg", - "G1MotionTrackingDeployEnv", - "G1MotionTrackingDeployEnvCfg", - "G1MotionTrackingEnv", - "G1MotionTrackingEnvCfg", - "G1MotionTrackingSAC23DofCfg", - "G1MotionTrackingSAC23DofEnv", - "G1MotionTrackingSACCfg", - "G1MotionTrackingSACEnv", - "G1WBTObs23DofCfg", - "G1WBTObsCfg", - "G1WBTObsEnv", - "G1WallFlipTracking23DofCfg", - "G1WallFlipTracking23DofEnvCfg", - "G1WallFlipTrackingCfg", - "G1WallFlipTrackingEnv", - "G1WallFlipTrackingEnvCfg", - "G1WallFlipTrackingSAC23DofCfg", - "G1WallFlipTrackingSAC23DofEnv", - "G1WallFlipTrackingSACCfg", - "G1WallFlipTrackingSACEnv", -] +for _task_name in G1_MOTION_TASKS: + registry.register_env_config(_task_name, ManagerBasedRlEnvCfg) + registry.register_env(_task_name, make_manager_based_rl_env, sim_backend="mujoco") + registry.register_env(_task_name, make_manager_based_rl_env, sim_backend="motrix") + + +__all__ = ["BoxMotionData", "BoxMotionLoader", "G1_MOTION_TASKS"] diff --git a/src/unilab/tasks/motion_tracking/g1/box_tracking.py b/src/unilab/tasks/motion_tracking/g1/box_tracking.py deleted file mode 100644 index 64ef22763..000000000 --- a/src/unilab/tasks/motion_tracking/g1/box_tracking.py +++ /dev/null @@ -1,412 +0,0 @@ -"""G1 box tracking environment with object-aware motion imitation.""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from typing import Any, cast - -import numpy as np - -from unilab.assets import ASSETS_ROOT_PATH -from unilab.base import registry -from unilab.base.scene import SceneCfg -from unilab.dr import DomainRandomizationManager, ResetPlan -from unilab.dr.dr_utils import build_common_reset_randomization, zero_actions -from unilab.dtype_config import get_global_dtype -from unilab.utils.geometry import np_sample_uniform -from unilab.utils.rotation import ( - np_matrix_from_quat, - np_quat_apply, - np_quat_error_magnitude, - np_quat_from_euler_xyz, - np_quat_inv, - np_quat_mul, - np_subtract_frame_transforms, -) - -from ..common.rewards import RewardContext -from .motion_box_loader import BoxMotionData, BoxMotionLoader -from .tracking import ( - G1MotionTrackingCfg, - G1MotionTrackingDomainRandomizationProvider, - G1MotionTrackingEnv, - RewardConfig, -) - - -@dataclass -class BoxRewardConfig(RewardConfig): - """Reward config extended with object-tracking terms.""" - - scales: dict[str, float] = field( - default_factory=lambda: { - **RewardConfig().scales, - "undesired_contacts": -0.1, - "object_global_ref_position_error_exp": 1.0, - "object_global_ref_orientation_error_exp": 1.0, - } - ) - std_object_pos: float = 0.3 - std_object_ori: float = 0.4 - - -@dataclass -class G1BoxTrackingCfg(G1MotionTrackingCfg): - """Configuration for the G1 large-box tracking task.""" - - scene: SceneCfg = field( - default_factory=lambda: SceneCfg( - model_file=str(ASSETS_ROOT_PATH / "robots" / "g1" / "scene_flat_with_largebox.xml") - ) - ) - motion_file: str | list[str] = str( - ASSETS_ROOT_PATH / "motions" / "g1" / "sub3_largebox_003_boxconverted.npz" - ) - object_body_name: str = "largebox" - object_pos_threshold: float = 0.25 - object_ori_threshold: float = 0.8 - reward_config: BoxRewardConfig = field( # pyright: ignore[reportIncompatibleVariableOverride] - default_factory=BoxRewardConfig - ) - - -@registry.envcfg("G1BoxTracking") -@dataclass -class G1BoxTrackingEnvCfg(G1BoxTrackingCfg): - """Registered config for G1 box tracking.""" - - pass - - -@dataclass -class G1BoxTracking23DofCfg(G1BoxTrackingCfg): - scene: SceneCfg = field( - default_factory=lambda: SceneCfg( - model_file=str( - ASSETS_ROOT_PATH / "robots" / "g1" / "scene_flat_23dof_with_largebox.xml" - ) - ) - ) - motion_file: str | list[str] = str( - ASSETS_ROOT_PATH / "motions" / "g1" / "sub3_largebox_003_boxconverted_23dof.npz" - ) - body_names: tuple[str, ...] = ( - "pelvis", - "left_hip_roll_link", - "left_knee_link", - "left_ankle_roll_link", - "right_hip_roll_link", - "right_knee_link", - "right_ankle_roll_link", - "torso_link", - "left_shoulder_roll_link", - "left_elbow_link", - "left_wrist_roll_rubber_hand", - "right_shoulder_roll_link", - "right_elbow_link", - "right_wrist_roll_rubber_hand", - ) - ee_body_names: tuple[str, ...] = ( - "left_ankle_roll_link", - "right_ankle_roll_link", - "left_wrist_roll_rubber_hand", - "right_wrist_roll_rubber_hand", - ) - - -@registry.envcfg("G1BoxTracking23Dof") -@dataclass -class G1BoxTracking23DofEnvCfg(G1BoxTracking23DofCfg): - pass - - -def _build_box_motion_reference_state( - env: Any, env_ids: np.ndarray, motion_data: BoxMotionData -) -> tuple[np.ndarray, np.ndarray]: - dtype = get_global_dtype() - num_reset = len(env_ids) - - root_pos = motion_data.body_pos_w[:, 0].copy() - root_ori = motion_data.body_quat_w[:, 0].copy() - root_lin_vel = motion_data.body_lin_vel_w[:, 0].copy() - root_ang_vel = motion_data.body_ang_vel_w[:, 0].copy() - joint_pos = motion_data.joint_pos.copy() - joint_vel = motion_data.joint_vel.copy() - - pose_rand = env.cfg.pose_randomization - pose_ranges = [ - (pose_rand.x[0], pose_rand.x[1]), - (pose_rand.y[0], pose_rand.y[1]), - (pose_rand.z[0], pose_rand.z[1]), - (pose_rand.roll[0], pose_rand.roll[1]), - (pose_rand.pitch[0], pose_rand.pitch[1]), - (pose_rand.yaw[0], pose_rand.yaw[1]), - ] - pose_samples = np.array( - [[np.random.uniform(low, high) for low, high in pose_ranges] for _ in range(num_reset)], - dtype=dtype, - ) - root_pos += pose_samples[:, 0:3] - root_ori = np_quat_mul( - np_quat_from_euler_xyz(pose_samples[:, 3], pose_samples[:, 4], pose_samples[:, 5]), - root_ori, - ) - - vel_rand = env.cfg.velocity_randomization - vel_ranges = [ - (vel_rand.x[0], vel_rand.x[1]), - (vel_rand.y[0], vel_rand.y[1]), - (vel_rand.z[0], vel_rand.z[1]), - (vel_rand.roll[0], vel_rand.roll[1]), - (vel_rand.pitch[0], vel_rand.pitch[1]), - (vel_rand.yaw[0], vel_rand.yaw[1]), - ] - vel_samples = np.array( - [[np.random.uniform(low, high) for low, high in vel_ranges] for _ in range(num_reset)], - dtype=dtype, - ) - root_lin_vel += vel_samples[:, :3] - root_ang_vel += vel_samples[:, 3:] - - joint_pos += np_sample_uniform( - env.cfg.joint_position_range[0], - env.cfg.joint_position_range[1], - joint_pos.shape, - dtype=np.float32, - ) - joint_range = env._get_joint_range() - if joint_range is not None: - joint_pos = np.clip(joint_pos, joint_range[:, 0], joint_range[:, 1]) - - qpos = np.tile(env._init_qpos, (num_reset, 1)) - qvel = np.tile(env._init_qvel, (num_reset, 1)) - - qpos[:, 0:3] = root_pos - qpos[:, 3:7] = root_ori - qpos[:, 7 : 7 + joint_pos.shape[1]] = joint_pos - - qvel[:, 0:3] = root_lin_vel - qvel[:, 3:6] = np_quat_apply(np_quat_inv(root_ori), root_ang_vel) - qvel[:, 6 : 6 + joint_vel.shape[1]] = joint_vel - - if motion_data.object_pos_w is not None: - qpos[:, env._obj_pos_slice] = motion_data.object_pos_w - qpos[:, env._obj_quat_slice] = motion_data.object_quat_w - qvel[:, env._obj_lin_vel_slice] = motion_data.object_lin_vel_w - qvel[:, env._obj_ang_vel_slice] = motion_data.object_ang_vel_w - - return qpos, qvel - - -class G1BoxTrackingDomainRandomizationProvider(G1MotionTrackingDomainRandomizationProvider): - """Reset provider that restores both robot and object state from motion data.""" - - def build_reset_plan(self, env: Any, env_ids: np.ndarray) -> ResetPlan: - num_reset = len(env_ids) - motion_frames = env.motion_sampler.sample_frames(env_ids) - motion_data = cast(BoxMotionData, env.motion_loader.get_motion_at_frame(motion_frames)) - qpos, qvel = _build_box_motion_reference_state(env, env_ids, motion_data) - - info_updates = { - "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, base_kp=self._base_kp, base_kd=self._base_kd - ), - ) - - -@registry.env("G1BoxTracking", sim_backend="mujoco") -@registry.env("G1BoxTracking", sim_backend="motrix") -@registry.env("G1BoxTracking23Dof", sim_backend="mujoco") -@registry.env("G1BoxTracking23Dof", sim_backend="motrix") -class G1BoxTrackingEnv(G1MotionTrackingEnv): - """Motion tracking env extended with large-box state and rewards.""" - - _cfg: G1BoxTrackingCfg # pyright: ignore[reportIncompatibleVariableOverride] - - def __init__(self, cfg: G1BoxTrackingCfg, num_envs=1, backend_type="mujoco"): - super().__init__(cfg, num_envs, backend_type) - - motion_body_ids = self._backend.get_motion_body_ids(cfg.body_names) - self.motion_loader = BoxMotionLoader(cfg.motion_file, body_indices=motion_body_ids) - self.motion_sampler = type(self.motion_sampler)( - self.motion_loader, mode=cfg.sampling_mode, num_envs=num_envs - ) - - if cfg.domain_rand.randomize_kp or cfg.domain_rand.randomize_kd: - base_kp, base_kd = self._backend.get_actuator_gains() - dr_provider = G1BoxTrackingDomainRandomizationProvider(base_kp=base_kp, base_kd=base_kd) - else: - dr_provider = G1BoxTrackingDomainRandomizationProvider() - # Parent init already applied init randomization and materialized the backend. - # Box tracking only needs to swap in a box-aware reset/obs provider for future resets. - self._dr_manager = DomainRandomizationManager(self, dr_provider) - - self._object_body_ids = self._backend.get_body_ids([cfg.object_body_name]) - - nq = self._init_qpos.shape[0] - self._obj_pos_slice = slice(nq - 7, nq - 4) - self._obj_quat_slice = slice(nq - 4, nq) - nv = self._init_qvel.shape[0] - self._obj_lin_vel_slice = slice(nv - 6, nv - 3) - self._obj_ang_vel_slice = slice(nv - 3, nv) - - if not self.motion_loader.has_object: - raise ValueError( - f"Motion file '{cfg.motion_file}' does not contain object data. " - "Expected keys: object_pos_w, object_quat_w, object_lin_vel_w, object_ang_vel_w" - ) - - def _get_joint_range(self) -> np.ndarray | None: - joint_range = super()._get_joint_range() - if joint_range is not None and joint_range.shape[0] > self.motion_loader.num_joints: - joint_range = joint_range[: self.motion_loader.num_joints] - return joint_range - - def _resample_reference_state(self, env_ids: np.ndarray) -> None: - motion_frames = self.motion_sampler.sample_frames(env_ids) - motion_data = cast(BoxMotionData, self.motion_loader.get_motion_at_frame(motion_frames)) - qpos, qvel = _build_box_motion_reference_state(self, env_ids, motion_data) - self._backend.set_state(env_ids, qpos, qvel) - - def get_dof_pos(self) -> np.ndarray: - dof_pos = super().get_dof_pos() - return dof_pos[:, : self.motion_loader.num_joints] - - def get_dof_vel(self) -> np.ndarray: - dof_vel = super().get_dof_vel() - return dof_vel[:, : self.motion_loader.num_joints] - - @property - def obs_groups_spec(self) -> dict[str, int]: - spec = super().obs_groups_spec - return {**spec, "critic": spec["critic"] + 12} - - def _actor_obs_dim(self, n: int) -> int: - return 6 + 3 + n * 5 - - def _build_actor_obs( - self, - *, - command: np.ndarray, - motion_anchor_pos_b: np.ndarray, - motion_anchor_ori_b: np.ndarray, - noisy_linvel: np.ndarray, - noisy_gyro: np.ndarray, - noisy_joint_pos_rel: np.ndarray, - noisy_dof_vel: np.ndarray, - last_actions: np.ndarray, - ) -> np.ndarray: - return np.concatenate( - [ - command, - motion_anchor_ori_b, - noisy_gyro, - noisy_joint_pos_rel, - noisy_dof_vel, - last_actions, - ], - axis=1, - dtype=get_global_dtype(), - ) - - def _init_reward_functions(self): - super()._init_reward_functions() - self._reward_fns["object_global_ref_position_error_exp"] = self._reward_object_position - self._reward_fns["object_global_ref_orientation_error_exp"] = ( - self._reward_object_orientation - ) - - def _compute_terminations( - self, - motion_data: BoxMotionData, - robot_body_pos_w: np.ndarray, - robot_body_quat_w: np.ndarray, - ) -> np.ndarray: - terminated = super()._compute_terminations(motion_data, robot_body_pos_w, robot_body_quat_w) - - if motion_data.object_pos_w is not None: - obj_pos_w = self._backend.get_body_pos_w(self._object_body_ids)[:, 0, :] - obj_pos_error = np.linalg.norm(obj_pos_w - motion_data.object_pos_w, axis=-1) - terminated |= obj_pos_error > self._cfg.object_pos_threshold - - if motion_data.object_quat_w is not None: - obj_quat_w = self._backend.get_body_quat_w(self._object_body_ids)[:, 0, :] - obj_ori_error = np_quat_error_magnitude(obj_quat_w, motion_data.object_quat_w) - terminated |= obj_ori_error > self._cfg.object_ori_threshold - - return terminated - - def _compute_obs( - self, - info: dict, - motion_data: BoxMotionData, - linvel: np.ndarray, - gyro: np.ndarray, - dof_pos: np.ndarray, - dof_vel: np.ndarray, - robot_body_pos_w: np.ndarray, - robot_body_quat_w: np.ndarray, - ) -> dict[str, np.ndarray]: - obs = super()._compute_obs( - info, motion_data, linvel, gyro, dof_pos, dof_vel, robot_body_pos_w, robot_body_quat_w - ) - - env_ids = info.get("env_ids") - if isinstance(env_ids, np.ndarray): - obj_pos_w = self._backend.get_body_pos_w(self._object_body_ids)[env_ids, 0, :] - obj_quat_w = self._backend.get_body_quat_w(self._object_body_ids)[env_ids, 0, :] - obj_lin_vel_w = self._backend.get_body_lin_vel_w(self._object_body_ids)[env_ids, 0, :] - else: - num_envs = linvel.shape[0] - obj_pos_w = self._backend.get_body_pos_w(self._object_body_ids)[:num_envs, 0, :] - obj_quat_w = self._backend.get_body_quat_w(self._object_body_ids)[:num_envs, 0, :] - obj_lin_vel_w = self._backend.get_body_lin_vel_w(self._object_body_ids)[:num_envs, 0, :] - - anchor_pos_w = robot_body_pos_w[:, self.anchor_body_idx] - anchor_quat_w = robot_body_quat_w[:, self.anchor_body_idx] - - obj_pos_b, obj_ori_rel = np_subtract_frame_transforms( - anchor_pos_w, anchor_quat_w, obj_pos_w, obj_quat_w - ) - obj_ori_mat = np_matrix_from_quat(obj_ori_rel) - num_envs = linvel.shape[0] - obj_ori_b = obj_ori_mat[:, :, :2].reshape(num_envs, 6) - obj_lin_vel_b = np_quat_apply(np_quat_inv(anchor_quat_w), obj_lin_vel_w) - - object_obs = np.concatenate( - [obj_pos_b, obj_ori_b, obj_lin_vel_b], - axis=1, - dtype=get_global_dtype(), - ) - obs["critic"] = np.concatenate( - [obs["critic"], object_obs], axis=1, dtype=get_global_dtype() - ) - return obs - - def _reward_object_position(self, ctx: RewardContext) -> np.ndarray: - motion_data: BoxMotionData = ctx.motion_data - if motion_data.object_pos_w is None: - return np.zeros((self._num_envs,), dtype=get_global_dtype()) - obj_pos_w = self._backend.get_body_pos_w(self._object_body_ids)[:, 0, :] - error = np.sum(np.square(obj_pos_w - motion_data.object_pos_w), axis=-1) - return np.asarray( - np.exp(-error / self._cfg.reward_config.std_object_pos**2), dtype=get_global_dtype() - ) - - def _reward_object_orientation(self, ctx: RewardContext) -> np.ndarray: - motion_data: BoxMotionData = ctx.motion_data - if motion_data.object_quat_w is None: - return np.zeros((self._num_envs,), dtype=get_global_dtype()) - obj_quat_w = self._backend.get_body_quat_w(self._object_body_ids)[:, 0, :] - error = np_quat_error_magnitude(obj_quat_w, motion_data.object_quat_w) ** 2 - return np.asarray( - np.exp(-error / self._cfg.reward_config.std_object_ori**2), dtype=get_global_dtype() - ) diff --git a/src/unilab/tasks/motion_tracking/g1/flip_tracking.py b/src/unilab/tasks/motion_tracking/g1/flip_tracking.py deleted file mode 100644 index 2bda6a55b..000000000 --- a/src/unilab/tasks/motion_tracking/g1/flip_tracking.py +++ /dev/null @@ -1,264 +0,0 @@ -"""Flip-specialized G1 motion tracking environment. - -This keeps the generic G1MotionTracking defaults backward-compatible while -providing a dedicated registry task for flip-focused datasets/profiles. -""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from typing import Literal - -from unilab.assets import ASSETS_ROOT_PATH -from unilab.base import registry -from unilab.base.scene import SceneCfg - -from ..common.config import ( - _zero_pose_randomization, - _zero_velocity_randomization, -) -from .tracking import ( - G1MotionTrackingCfg, - G1MotionTrackingEnv, - PoseRandomization, - VelocityRandomization, -) - - -@dataclass -class G1FlipTrackingCfg(G1MotionTrackingCfg): - """Config profile for flip tracking clips.""" - - scene: SceneCfg = field( - default_factory=lambda: SceneCfg( - model_file=str(ASSETS_ROOT_PATH / "robots" / "g1" / "scene_flat.xml") - ) - ) - motion_file: str | list[str] = str( - ASSETS_ROOT_PATH / "motions" / "g1" / "flip_360_001__A304.npz" - ) - pose_randomization: PoseRandomization = field(default_factory=_zero_pose_randomization) - velocity_randomization: VelocityRandomization = field( - default_factory=_zero_velocity_randomization - ) - joint_position_range: tuple[float, float] = (0.0, 0.0) - sampling_mode: Literal["start", "clip_start", "uniform", "adaptive", "mixed"] = "start" - terminate_on_undesired_contacts: bool = True - # Some flip clips include large anchor orientation deviations. - anchor_ori_threshold: float = 1e9 - - -@registry.envcfg("G1FlipTracking") -@dataclass -class G1FlipTrackingEnvCfg(G1FlipTrackingCfg): - """Registered configuration for G1 flip tracking.""" - - pass - - -@registry.env("G1FlipTracking", sim_backend="mujoco") -@registry.env("G1FlipTracking", sim_backend="motrix") -class G1FlipTrackingEnv(G1MotionTrackingEnv): - """G1 flip-tracking environment implementation.""" - - _cfg: G1FlipTrackingCfg # pyright: ignore[reportIncompatibleVariableOverride] - - -@dataclass -class G1WallFlipTrackingCfg(G1FlipTrackingCfg): - """Config profile for wall-assisted G1 flip tracking clips.""" - - scene: SceneCfg = field( - default_factory=lambda: SceneCfg( - model_file=str(ASSETS_ROOT_PATH / "robots" / "g1" / "scene_flat_with_wall.xml") - ) - ) - motion_file: str | list[str] = str( - ASSETS_ROOT_PATH / "motions" / "g1" / "flip_from_wall_104__A304.npz" - ) - sampling_mode: Literal["start", "clip_start", "uniform", "adaptive", "mixed"] = "adaptive" - anchor_pos_z_threshold: float = 0.5 - ee_body_pos_z_threshold: float = 0.5 - - -@registry.envcfg("G1WallFlipTracking") -@dataclass -class G1WallFlipTrackingEnvCfg(G1WallFlipTrackingCfg): - """Registered configuration for G1 wall flip tracking.""" - - pass - - -@registry.env("G1WallFlipTracking", sim_backend="mujoco") -@registry.env("G1WallFlipTracking", sim_backend="motrix") -class G1WallFlipTrackingEnv(G1MotionTrackingEnv): - """G1 wall flip-tracking environment implementation.""" - - _cfg: G1WallFlipTrackingCfg # pyright: ignore[reportIncompatibleVariableOverride] - - -@dataclass -class G1ClimbTrackingCfg(G1MotionTrackingCfg): - """Config profile for the climb_20_z_scale_1 motion clip.""" - - scene: SceneCfg = field( - default_factory=lambda: SceneCfg( - model_file=str(ASSETS_ROOT_PATH / "robots" / "g1" / "scene_climb_20_z_scale_1.xml") - ) - ) - motion_file: str | list[str] = str( - ASSETS_ROOT_PATH / "motions" / "g1" / "climb_20_z_scale_1.0.npz" - ) - max_episode_seconds: float = 15.0 - anchor_pos_z_threshold: float = 0.5 - ee_body_pos_z_threshold: float = 0.5 - - -@registry.envcfg("G1ClimbTracking") -@dataclass -class G1ClimbTrackingEnvCfg(G1ClimbTrackingCfg): - """Registered configuration for G1 box-climb motion tracking.""" - - pass - - -@registry.env("G1ClimbTracking", sim_backend="mujoco") -@registry.env("G1ClimbTracking", sim_backend="motrix") -class G1ClimbTrackingEnv(G1MotionTrackingEnv): - """G1 climb-tracking environment implementation.""" - - _cfg: G1ClimbTrackingCfg # pyright: ignore[reportIncompatibleVariableOverride] - - -@dataclass -class G1FlipTracking23DofCfg(G1FlipTrackingCfg): - scene: SceneCfg = field( - default_factory=lambda: SceneCfg( - model_file=str(ASSETS_ROOT_PATH / "robots" / "g1" / "scene_flat_23dof.xml") - ) - ) - motion_file: str | list[str] = str( - ASSETS_ROOT_PATH / "motions" / "g1" / "flip_360_001__A304_23dof.npz" - ) - body_names: tuple[str, ...] = ( - "pelvis", - "left_hip_roll_link", - "left_knee_link", - "left_ankle_roll_link", - "right_hip_roll_link", - "right_knee_link", - "right_ankle_roll_link", - "torso_link", - "left_shoulder_roll_link", - "left_elbow_link", - "left_wrist_roll_rubber_hand", - "right_shoulder_roll_link", - "right_elbow_link", - "right_wrist_roll_rubber_hand", - ) - ee_body_names: tuple[str, ...] = ( - "left_ankle_roll_link", - "right_ankle_roll_link", - "left_wrist_roll_rubber_hand", - "right_wrist_roll_rubber_hand", - ) - - -@registry.envcfg("G1FlipTracking23Dof") -@dataclass -class G1FlipTracking23DofEnvCfg(G1FlipTracking23DofCfg): - pass - - -@dataclass -class G1WallFlipTracking23DofCfg(G1WallFlipTrackingCfg): - scene: SceneCfg = field( - default_factory=lambda: SceneCfg( - model_file=str(ASSETS_ROOT_PATH / "robots" / "g1" / "scene_flat_23dof_with_wall.xml") - ) - ) - motion_file: str | list[str] = str( - ASSETS_ROOT_PATH / "motions" / "g1" / "flip_from_wall_104__A304_23dof.npz" - ) - body_names: tuple[str, ...] = ( - "pelvis", - "left_hip_roll_link", - "left_knee_link", - "left_ankle_roll_link", - "right_hip_roll_link", - "right_knee_link", - "right_ankle_roll_link", - "torso_link", - "left_shoulder_roll_link", - "left_elbow_link", - "left_wrist_roll_rubber_hand", - "right_shoulder_roll_link", - "right_elbow_link", - "right_wrist_roll_rubber_hand", - ) - ee_body_names: tuple[str, ...] = ( - "left_ankle_roll_link", - "right_ankle_roll_link", - "left_wrist_roll_rubber_hand", - "right_wrist_roll_rubber_hand", - ) - - -@registry.envcfg("G1WallFlipTracking23Dof") -@dataclass -class G1WallFlipTracking23DofEnvCfg(G1WallFlipTracking23DofCfg): - pass - - -@dataclass -class G1ClimbTracking23DofCfg(G1ClimbTrackingCfg): - """23-DoF config for the climb_20_z_scale_1 motion clip.""" - - scene: SceneCfg = field( - default_factory=lambda: SceneCfg( - model_file=str( - ASSETS_ROOT_PATH / "robots" / "g1" / "scene_climb_20_z_scale_1_23dof.xml" - ) - ) - ) - motion_file: str | list[str] = str( - ASSETS_ROOT_PATH / "motions" / "g1" / "climb_20_z_scale_1.0_23dof.npz" - ) - body_names: tuple[str, ...] = ( - "pelvis", - "left_hip_roll_link", - "left_knee_link", - "left_ankle_roll_link", - "right_hip_roll_link", - "right_knee_link", - "right_ankle_roll_link", - "torso_link", - "left_shoulder_roll_link", - "left_elbow_link", - "left_wrist_roll_rubber_hand", - "right_shoulder_roll_link", - "right_elbow_link", - "right_wrist_roll_rubber_hand", - ) - ee_body_names: tuple[str, ...] = ( - "left_ankle_roll_link", - "right_ankle_roll_link", - "left_wrist_roll_rubber_hand", - "right_wrist_roll_rubber_hand", - ) - - -@registry.envcfg("G1ClimbTracking23Dof") -@dataclass -class G1ClimbTracking23DofEnvCfg(G1ClimbTracking23DofCfg): - """Registered 23-DoF configuration for G1 box-climb motion tracking.""" - - pass - - -registry.register_env("G1ClimbTracking23Dof", G1ClimbTrackingEnv, sim_backend="mujoco") -registry.register_env("G1ClimbTracking23Dof", G1ClimbTrackingEnv, sim_backend="motrix") -registry.register_env("G1FlipTracking23Dof", G1FlipTrackingEnv, sim_backend="mujoco") -registry.register_env("G1FlipTracking23Dof", G1FlipTrackingEnv, sim_backend="motrix") -registry.register_env("G1WallFlipTracking23Dof", G1WallFlipTrackingEnv, sim_backend="mujoco") -registry.register_env("G1WallFlipTracking23Dof", G1WallFlipTrackingEnv, sim_backend="motrix") diff --git a/src/unilab/tasks/motion_tracking/g1/flip_tracking_sac.py b/src/unilab/tasks/motion_tracking/g1/flip_tracking_sac.py deleted file mode 100644 index f3ec25e5d..000000000 --- a/src/unilab/tasks/motion_tracking/g1/flip_tracking_sac.py +++ /dev/null @@ -1,85 +0,0 @@ -"""SAC variants of the G1 flip / wall-flip tracking tasks. - -These pair the flip-specialized configs (scene / motion / zeroed reset -randomization, from :mod:`flip_tracking`) with the asymmetric actor-critic -observation layout of :class:`G1MotionTrackingSACEnv` (``base_lin_vel`` -appended to the critic obs). - -The flip behavior lives entirely in the cfg dataclasses and the SAC obs -augmentation lives entirely in the env class, so the two compose by simple -inheritance — no obs override is duplicated here. Registered under separate -names so the PPO flip pipeline (``G1FlipTracking`` / ``G1WallFlipTracking``) -is untouched, and sit at the same level as ``G1WBTObs``. -""" - -from __future__ import annotations - -from dataclasses import dataclass - -from unilab.base import registry - -from .flip_tracking import ( - G1FlipTracking23DofEnvCfg, - G1FlipTrackingCfg, - G1WallFlipTracking23DofEnvCfg, - G1WallFlipTrackingCfg, -) -from .tracking_sac import G1MotionTrackingSACEnv - - -@registry.envcfg("G1FlipTrackingSAC") -@dataclass -class G1FlipTrackingSACCfg(G1FlipTrackingCfg): - """Flip tracking cfg for SAC (identical fields, separate registry entry).""" - - -@registry.env("G1FlipTrackingSAC", sim_backend="mujoco") -@registry.env("G1FlipTrackingSAC", sim_backend="motrix") -class G1FlipTrackingSACEnv(G1MotionTrackingSACEnv): - """Flip tracking env with the SAC asymmetric actor-critic obs layout. - - The motrix backend is registered for sim2sim eval/playback only. - """ - - _cfg: G1FlipTrackingSACCfg # pyright: ignore[reportIncompatibleVariableOverride] - - -@registry.envcfg("G1WallFlipTrackingSAC") -@dataclass -class G1WallFlipTrackingSACCfg(G1WallFlipTrackingCfg): - """Wall-flip tracking cfg for SAC (identical fields, separate registry entry).""" - - -@registry.env("G1WallFlipTrackingSAC", sim_backend="mujoco") -@registry.env("G1WallFlipTrackingSAC", sim_backend="motrix") -class G1WallFlipTrackingSACEnv(G1MotionTrackingSACEnv): - """Wall-flip tracking env with the SAC asymmetric actor-critic obs layout. - - The motrix backend is registered for sim2sim eval/playback only. - """ - - _cfg: G1WallFlipTrackingSACCfg # pyright: ignore[reportIncompatibleVariableOverride] - - -@registry.envcfg("G1FlipTrackingSAC23Dof") -@dataclass -class G1FlipTrackingSAC23DofCfg(G1FlipTracking23DofEnvCfg): - pass - - -@registry.env("G1FlipTrackingSAC23Dof", sim_backend="mujoco") -@registry.env("G1FlipTrackingSAC23Dof", sim_backend="motrix") -class G1FlipTrackingSAC23DofEnv(G1MotionTrackingSACEnv): - _cfg: G1FlipTrackingSAC23DofCfg # pyright: ignore[reportIncompatibleVariableOverride] - - -@registry.envcfg("G1WallFlipTrackingSAC23Dof") -@dataclass -class G1WallFlipTrackingSAC23DofCfg(G1WallFlipTracking23DofEnvCfg): - pass - - -@registry.env("G1WallFlipTrackingSAC23Dof", sim_backend="mujoco") -@registry.env("G1WallFlipTrackingSAC23Dof", sim_backend="motrix") -class G1WallFlipTrackingSAC23DofEnv(G1MotionTrackingSACEnv): - _cfg: G1WallFlipTrackingSAC23DofCfg # pyright: ignore[reportIncompatibleVariableOverride] diff --git a/src/unilab/tasks/motion_tracking/g1/manager_terms.py b/src/unilab/tasks/motion_tracking/g1/manager_terms.py new file mode 100644 index 000000000..09ddc07f9 --- /dev/null +++ b/src/unilab/tasks/motion_tracking/g1/manager_terms.py @@ -0,0 +1,327 @@ +"""G1 profile-specific NumPy manager terms for motion tracking.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, cast + +import numpy as np + +from unilab.managers import ManagerTermBase, ManagerTermBaseCfg +from unilab.managers.scene_entity_config import SceneEntityCfg +from unilab.tasks.motion_tracking.common.manager_terms import ( + MotionCommand, + MotionCommandCfg, + MotionJointPositionAction, +) +from unilab.utils.geometry import np_write_relative_anchor_transform_pos_rot6d +from unilab.utils.rotation import ( + np_quat_apply_inverse, + np_quat_error_magnitude_squared_batched, +) + +from .motion_box_loader import BoxMotionData, BoxMotionLoader + +if TYPE_CHECKING: + from unilab.base.entity import Entity + from unilab.managers._types import ManagerBasedRlEnv + + +_DEFAULT_ASSET_CFG = SceneEntityCfg("robot") + + +@dataclass(kw_only=True) +class BoxMotionCommandCfg(MotionCommandCfg): + """Motion command extended with one independently simulated object.""" + + object_entity_name: str + + def build(self, env: ManagerBasedRlEnv) -> BoxMotionCommand: + return BoxMotionCommand(self, env) + + +class BoxMotionCommand(MotionCommand): + cfg: BoxMotionCommandCfg # pyright: ignore[reportIncompatibleVariableOverride] + + def __init__(self, cfg: BoxMotionCommandCfg, env: ManagerBasedRlEnv): + if not isinstance(cfg.object_entity_name, str) or not cfg.object_entity_name: + raise ValueError("BoxMotionCommandCfg object_entity_name must be non-empty") + self.object = cast("Entity", env.scene[cfg.object_entity_name]) + self._object_pos_w = np.empty((env.num_envs, 3), dtype=np.float32) + self._object_obs_b = np.empty((env.num_envs, 12), dtype=np.float32) + super().__init__(cfg, env) + if not isinstance(self._motion_data, BoxMotionData): + raise TypeError("BoxMotionCommand requires BoxMotionData") + self._refresh_object_state() + + def _make_motion_loader( + self, + motion_file: str | list[str], + body_indices: np.ndarray, + ) -> BoxMotionLoader: + return BoxMotionLoader(motion_file, body_indices=body_indices) + + @property + def box_motion(self) -> BoxMotionData: + return cast(BoxMotionData, self._motion_data) + + @property + def object_pos_w(self) -> np.ndarray: + return self._object_pos_w + + @property + def object_quat_w(self) -> np.ndarray: + value = self.box_motion.object_quat_w + if value is None: + raise RuntimeError("Box motion object quaternion was not materialized") + return value + + @property + def object_state_b(self) -> np.ndarray: + return self._object_obs_b + + def _refresh_motion(self) -> None: + super()._refresh_motion() + value = self.box_motion.object_pos_w + if value is None: + raise RuntimeError("Box motion object position was not materialized") + np.add(value, self._env.scene.env_origins, out=self._object_pos_w) + + def _resample_command(self, env_ids: np.ndarray) -> None: + super()._resample_command(env_ids) + motion = cast(BoxMotionData, self.motion.get_motion_at_frame(self.time_steps[env_ids])) + values = ( + motion.object_pos_w, + motion.object_quat_w, + motion.object_lin_vel_w, + motion.object_ang_vel_w, + ) + if any(value is None for value in values): + raise RuntimeError("Box motion reset requires complete object state") + object_pos = cast(np.ndarray, motion.object_pos_w).copy() + object_pos += self._env.scene.env_origins[env_ids] + object_state = np.concatenate( + ( + object_pos, + cast(np.ndarray, motion.object_quat_w), + cast(np.ndarray, motion.object_lin_vel_w), + cast(np.ndarray, motion.object_ang_vel_w), + ), + axis=-1, + ) + self.object.write_root_state_to_sim(object_state, env_ids=env_ids) + + def _refresh_object_state(self) -> None: + np_write_relative_anchor_transform_pos_rot6d( + self.robot_anchor_pos_w, + self.robot_anchor_quat_w, + self.object.data.root_link_pos_w, + self.object.data.root_link_quat_w, + self._object_obs_b[:, :3], + self._object_obs_b[:, 3:9], + ) + self._object_obs_b[:, 9:12] = np_quat_apply_inverse( + self.robot_anchor_quat_w, + self.object.data.root_link_lin_vel_w, + ) + + def post_compute(self) -> None: + super().post_compute() + self._refresh_object_state() + + +def _box_command(env: ManagerBasedRlEnv, command_name: str) -> BoxMotionCommand: + try: + command = env.command_manager.get_term(command_name) + except KeyError as exc: + raise KeyError(f"Box motion command term '{command_name}' not found") from exc + if not isinstance(command, BoxMotionCommand): + raise TypeError( + f"Command term '{command_name}' is {type(command).__name__}, expected BoxMotionCommand" + ) + return command + + +def object_state_b(env: ManagerBasedRlEnv, command_name: str) -> np.ndarray: + return _box_command(env, command_name).object_state_b + + +def object_global_position_error_exp( + env: ManagerBasedRlEnv, + command_name: str, + std: float, +) -> np.ndarray: + command = _box_command(env, command_name) + if not np.isfinite(std) or std <= 0.0: + raise ValueError("object position std must be finite and positive") + error = np.sum( + np.square(command.object_pos_w - command.object.data.root_link_pos_w), + axis=-1, + ) + return np.exp(-error / float(std) ** 2) + + +def object_global_orientation_error_exp( + env: ManagerBasedRlEnv, + command_name: str, + std: float, +) -> np.ndarray: + command = _box_command(env, command_name) + if not np.isfinite(std) or std <= 0.0: + raise ValueError("object orientation std must be finite and positive") + error = np_quat_error_magnitude_squared_batched( + command.object_quat_w, + command.object.data.root_link_quat_w, + ) + return np.exp(-error / float(std) ** 2) + + +def bad_object_position( + env: ManagerBasedRlEnv, + command_name: str, + threshold: float, +) -> np.ndarray: + command = _box_command(env, command_name) + error = np.linalg.norm(command.object_pos_w - command.object.data.root_link_pos_w, axis=-1) + return error > threshold + + +def bad_object_orientation( + env: ManagerBasedRlEnv, + command_name: str, + threshold: float, +) -> np.ndarray: + command = _box_command(env, command_name) + error = np.sqrt( + np_quat_error_magnitude_squared_batched( + command.object_quat_w, + command.object.data.root_link_quat_w, + ) + ) + return error > threshold + + +class randomize_encoder_bias(ManagerTermBase): + """Per-reset encoder calibration error through the public Entity data surface.""" + + def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv): + super().__init__(env) + asset_cfg = cfg.params.get("asset_cfg", _DEFAULT_ASSET_CFG) + if not isinstance(asset_cfg, SceneEntityCfg): + raise TypeError("randomize_encoder_bias asset_cfg must be SceneEntityCfg") + bias_range = np.asarray(cfg.params.get("bias_range"), dtype=np.float64) + if bias_range.shape != (2,) or not np.isfinite(bias_range).all(): + raise ValueError("randomize_encoder_bias bias_range must be a finite pair") + if bias_range[0] > bias_range[1]: + raise ValueError("randomize_encoder_bias minimum exceeds maximum") + self._range = (float(bias_range[0]), float(bias_range[1])) + self._entity = cast("Entity", env.scene[asset_cfg.name]) + raw_ids = np.arange(self._entity.num_joints)[asset_cfg.joint_ids] + self._joint_ids = np.asarray(raw_ids, dtype=np.intp) + + def __call__( + self, + env: ManagerBasedRlEnv, + env_ids: np.ndarray | None, + bias_range: tuple[float, float], + asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG, + ) -> None: + del bias_range, asset_cfg + ids = np.arange(env.num_envs, dtype=np.int32) if env_ids is None else env_ids + samples = env.rng.uniform( + *self._range, + size=(len(ids), len(self._joint_ids)), + ) + self._entity.data.encoder_bias[np.ix_(ids, self._joint_ids)] = samples + + +class joint_acc_l2(ManagerTermBase): + """Squared finite-difference joint acceleration with reset-aware state.""" + + def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv): + super().__init__(env) + asset_cfg = cfg.params.get("asset_cfg", _DEFAULT_ASSET_CFG) + if not isinstance(asset_cfg, SceneEntityCfg): + raise TypeError("joint_acc_l2 asset_cfg must be SceneEntityCfg") + self._entity = cast("Entity", env.scene[asset_cfg.name]) + self._joint_ids = np.arange(self._entity.num_joints, dtype=np.intp)[asset_cfg.joint_ids] + self._previous = self._entity.data.joint_vel[:, self._joint_ids].copy() + + def reset(self, env_ids: np.ndarray | slice | None) -> None: + ids = np.arange(self.num_envs, dtype=np.intp) + if env_ids is not None: + ids = ids[env_ids] + self._previous[ids] = self._entity.data.joint_vel[np.ix_(ids, self._joint_ids)] + + def __call__( + self, + env: ManagerBasedRlEnv, + asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG, + ) -> np.ndarray: + del asset_cfg + velocity = self._entity.data.joint_vel[:, self._joint_ids] + acceleration = (velocity - self._previous) / env.step_dt + self._previous[:] = velocity + return np.sum(np.square(acceleration), axis=-1) + + +class joint_torque_l2(ManagerTermBase): + """Position-controller torque estimate using cold-path actuator gain binding.""" + + def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv): + super().__init__(env) + asset_cfg = cfg.params.get("asset_cfg", _DEFAULT_ASSET_CFG) + if not isinstance(asset_cfg, SceneEntityCfg): + raise TypeError("joint_torque_l2 asset_cfg must be SceneEntityCfg") + action_name = cfg.params.get("action_name", "joint_pos") + if not isinstance(action_name, str) or not action_name: + raise ValueError("joint_torque_l2 action_name must be non-empty") + action = env.action_manager.get_term(action_name) + if not isinstance(action, MotionJointPositionAction): + raise TypeError("joint_torque_l2 requires MotionJointPositionAction") + self._action = action + self._entity = cast("Entity", env.scene[asset_cfg.name]) + actuator_ids, kp, kd = self._entity.bind_actuator_gain_write( + asset_cfg.actuator_ids, + term_name="joint_torque_l2", + ) + selected_names = tuple(self._entity.actuator_names[int(index)] for index in actuator_ids) + if selected_names != tuple(action.target_names): + raise ValueError( + "joint_torque_l2 actuator order does not match the action target order: " + f"{selected_names} != {tuple(action.target_names)}" + ) + self._kp = kp + self._kd = kd + self._joint_ids = action.target_ids + self._torque = np.empty_like(action.target) + + def __call__( + self, + env: ManagerBasedRlEnv, + action_name: str = "joint_pos", + asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG, + ) -> np.ndarray: + del env, action_name, asset_cfg + np.subtract( + self._action.target, + self._entity.data.joint_pos[:, self._joint_ids], + out=self._torque, + ) + self._torque *= self._kp + self._torque -= self._kd * self._entity.data.joint_vel[:, self._joint_ids] + return np.sum(np.square(self._torque), axis=-1) + + +__all__ = [ + "BoxMotionCommand", + "BoxMotionCommandCfg", + "bad_object_orientation", + "bad_object_position", + "joint_acc_l2", + "joint_torque_l2", + "object_global_orientation_error_exp", + "object_global_position_error_exp", + "object_state_b", + "randomize_encoder_bias", +] diff --git a/src/unilab/tasks/motion_tracking/g1/motion_box_loader.py b/src/unilab/tasks/motion_tracking/g1/motion_box_loader.py index 795f13274..a44557d58 100644 --- a/src/unilab/tasks/motion_tracking/g1/motion_box_loader.py +++ b/src/unilab/tasks/motion_tracking/g1/motion_box_loader.py @@ -92,6 +92,26 @@ def __init__(self, motion_file, body_indices=None): self.joint_pos = self.joint_pos[:, :n_robot_joints] self.joint_vel = self.joint_vel[:, :n_robot_joints] + def make_motion_data_buffer(self, num_frames: int) -> BoxMotionData: + base = super().make_motion_data_buffer(num_frames) + if not self.has_object: + raise ValueError( + "BoxMotionLoader requires object_pos_w, object_quat_w, " + "object_lin_vel_w, and object_ang_vel_w" + ) + return BoxMotionData( + joint_pos=base.joint_pos, + joint_vel=base.joint_vel, + body_pos_w=base.body_pos_w, + body_quat_w=base.body_quat_w, + body_lin_vel_w=base.body_lin_vel_w, + body_ang_vel_w=base.body_ang_vel_w, + object_pos_w=np.empty((num_frames, 3), dtype=self.object_pos_w.dtype), + object_quat_w=np.empty((num_frames, 4), dtype=self.object_quat_w.dtype), + object_lin_vel_w=np.empty((num_frames, 3), dtype=self.object_lin_vel_w.dtype), + object_ang_vel_w=np.empty((num_frames, 3), dtype=self.object_ang_vel_w.dtype), + ) + def get_motion_at_frame( self, frame_idx: np.ndarray, out: MotionData | None = None ) -> BoxMotionData: diff --git a/src/unilab/tasks/motion_tracking/g1/tracking.py b/src/unilab/tasks/motion_tracking/g1/tracking.py deleted file mode 100644 index de3618803..000000000 --- a/src/unilab/tasks/motion_tracking/g1/tracking.py +++ /dev/null @@ -1,152 +0,0 @@ -"""G1 Motion Tracking profiles — thin registry subclasses over the shared engine. - -The robot-agnostic engine and owner modules live in -:mod:`unilab.tasks.motion_tracking.common`. This module keeps the G1 registry -entries (``G1MotionTracking`` / ``G1MotionTrackingDeploy``) and re-exports the -historical ``G1*`` / ``Domain_Rand`` / ``_build_motion_reference_state`` symbol -names so existing subclasses and tests keep importing them from ``.tracking``. -""" - -from __future__ import annotations - -from dataclasses import dataclass, field - -from unilab.assets import ASSETS_ROOT_PATH -from unilab.base import registry -from unilab.base.scene import SceneCfg -from unilab.envs import ManagerBasedRlEnvCfg, make_manager_based_rl_env - -from ..common.config import ( - Domain_Rand, - DomainRand, - MotionTrackingCfg, - MotionTrackingDeployEnvCfg, - PoseRandomization, - VelocityRandomization, - _zero_pose_randomization, - _zero_velocity_randomization, -) -from ..common.domain_randomization import ( - MotionTrackingDomainRandomizationProvider, -) -from ..common.reset import build_motion_reference_state -from ..common.rewards import RewardConfig -from ..common.tracking import ( - MotionTrackingDeployEnv, - MotionTrackingEnv, -) - -# ── backward-compatible aliases (historical G1* symbol names) ──────── -G1MotionTrackingCfg = MotionTrackingCfg -G1MotionTrackingDomainRandomizationProvider = MotionTrackingDomainRandomizationProvider -_build_motion_reference_state = build_motion_reference_state - - -@dataclass -class G1MotionTrackingEnvCfg(MotionTrackingCfg): - """Registered configuration for G1 motion tracking.""" - - pass - - -@dataclass -class G1MotionTrackingDeployEnvCfg(MotionTrackingDeployEnvCfg): - """Registered deploy configuration for G1 motion tracking.""" - - pass - - -@dataclass -class G1MotionTracking23DofCfg(G1MotionTrackingCfg): - motion_file: str | list[str] = str( - ASSETS_ROOT_PATH / "motions" / "g1" / "dance1_subject2_part_23dof.npz" - ) - scene: SceneCfg = field( - default_factory=lambda: SceneCfg( - model_file=str(ASSETS_ROOT_PATH / "robots" / "g1" / "scene_flat_23dof.xml") - ) - ) - body_names: tuple[str, ...] = ( - "pelvis", - "left_hip_roll_link", - "left_knee_link", - "left_ankle_roll_link", - "right_hip_roll_link", - "right_knee_link", - "right_ankle_roll_link", - "torso_link", - "left_shoulder_roll_link", - "left_elbow_link", - "left_wrist_roll_rubber_hand", - "right_shoulder_roll_link", - "right_elbow_link", - "right_wrist_roll_rubber_hand", - ) - ee_body_names: tuple[str, ...] = ( - "left_ankle_roll_link", - "right_ankle_roll_link", - "left_wrist_roll_rubber_hand", - "right_wrist_roll_rubber_hand", - ) - - -@dataclass -class G1MotionTracking23DofEnvCfg(G1MotionTracking23DofCfg): - pass - - -@dataclass -class G1MotionTracking23DofDeployEnvCfg(G1MotionTracking23DofCfg): - pass - - -class G1MotionTrackingEnv(MotionTrackingEnv): - """G1 Motion Tracking Environment.""" - - _cfg: MotionTrackingCfg - - -class G1MotionTrackingDeployEnv(MotionTrackingDeployEnv): - """Deploy-oriented G1 motion tracking env with unitree_rl_lab mimic actor inputs.""" - - _cfg: MotionTrackingDeployEnvCfg - - -# The legacy classes above remain explicit consumers for profiles that are not -# part of #1227. The four core identities have one Hydra-owned manager config -# and one generic runtime factory instead of inheriting those classes. -for _task_name in ( - "G1MotionTracking", - "G1MotionTrackingDeploy", - "G1MotionTracking23Dof", - "G1MotionTracking23DofDeploy", -): - registry.register_env_config(_task_name, ManagerBasedRlEnvCfg) - registry.register_env(_task_name, make_manager_based_rl_env, sim_backend="mujoco") - registry.register_env(_task_name, make_manager_based_rl_env, sim_backend="motrix") - - -__all__ = [ - "DomainRand", - "Domain_Rand", - "G1MotionTracking23DofCfg", - "G1MotionTracking23DofDeployEnvCfg", - "G1MotionTracking23DofEnvCfg", - "G1MotionTrackingCfg", - "G1MotionTrackingDeployEnv", - "G1MotionTrackingDeployEnvCfg", - "G1MotionTrackingDomainRandomizationProvider", - "G1MotionTrackingEnv", - "G1MotionTrackingEnvCfg", - "MotionTrackingCfg", - "MotionTrackingDeployEnv", - "MotionTrackingDeployEnvCfg", - "MotionTrackingDomainRandomizationProvider", - "MotionTrackingEnv", - "PoseRandomization", - "RewardConfig", - "VelocityRandomization", - "_build_motion_reference_state", - "_zero_pose_randomization", - "_zero_velocity_randomization", -] diff --git a/src/unilab/tasks/motion_tracking/g1/tracking_obs.py b/src/unilab/tasks/motion_tracking/g1/tracking_obs.py deleted file mode 100644 index d6c53abd9..000000000 --- a/src/unilab/tasks/motion_tracking/g1/tracking_obs.py +++ /dev/null @@ -1,517 +0,0 @@ -"""G1 Whole-Body Tracking — sim2real-oriented SAC variant (task ``G1WBTObs``). - -This module registers a strict subclass of :class:`G1MotionTrackingSACEnv` that -adds the training-pipeline pieces needed for ONNX-on-real-G1 deployment: - -* drop deploy-unavailable channels from the actor obs - (``base_lin_vel``, ``motion_anchor_pos_b``); -* per-step uniform noise on ``motion_anchor_ori_b`` (actor only); -* proprio observation history (``gyro`` / ``joint_pos_rel`` / ``dof_vel`` / - ``last_actions``) flattened oldest-first per term, matching the deploy-side - ``ObservationManager`` when ``use_gym_history=false``; -* per-episode encoder bias on ``joint_pos_rel`` (actor only); -* per-episode foot-geom friction sampled across regex-matched geoms; -* per-episode y / z COM offsets layered on top of the existing x offset; -* ``joint_acc_l2`` and ``joint_torque_l2`` reward terms. - -All extensions are gated by flags on :class:`G1WBTObsCfg`; the bases -(``G1MotionTrackingSACCfg`` / ``G1MotionTrackingSACEnv`` / -``G1MotionTrackingEnv``) are untouched. Switch the pelvis IMU via the -yaml ``env.sensor.gyro``/``env.sensor.upvector``/``env.sensor.local_linvel`` -fields (no XML duplication required — ``g1.xml`` already exposes both IMUs). -""" - -from __future__ import annotations - -import re -from dataclasses import dataclass, field -from typing import Any - -import numpy as np - -from unilab.assets import ASSETS_ROOT_PATH -from unilab.base import registry -from unilab.base.scene import SceneCfg -from unilab.dr import ( - DomainRandomizationCapabilities, - ResetPlan, - ResetRandomizationPayload, -) -from unilab.dr.dr_utils import ( - build_common_reset_randomization, - zero_actions, -) -from unilab.dr.types import RESET_TERM_GEOM_FRICTION -from unilab.dtype_config import get_global_dtype -from unilab.tasks.motion_tracking.common.g1_base import NoiseConfig - -from ..common.rewards import RewardContext -from .tracking import ( - Domain_Rand, - G1MotionTrackingDomainRandomizationProvider, - _build_motion_reference_state, -) -from .tracking_sac import G1MotionTrackingSACCfg, G1MotionTrackingSACEnv - -# --------------------------------------------------------------------------- # -# Config extensions -# --------------------------------------------------------------------------- # - - -@dataclass -class ObsNoiseConfig(NoiseConfig): - """Actor obs masking flags + proprio history depth. - - Defaults preserve baseline behaviour so the parent ``NoiseConfig`` remains - a drop-in replacement; ``G1WBTObs`` flips the flags via its task yaml. - """ - - # Drop ``base_lin_vel`` from actor obs (G1 has no on-robot linvel sensor). - enable_zero_linvel: bool = False - # Drop ``motion_anchor_pos_b`` from actor obs (no torso-pose estimator). - enable_zero_anchor_pos: bool = False - # Per-step uniform noise on ``motion_anchor_ori_b`` (actor only). - enable_anchor_ori_noise: bool = False - scale_anchor_ori: float = 0.05 - # When > 1, proprio terms (gyro / joint_pos_rel / dof_vel / last_actions) - # are flattened oldest-first as an H-step history block. Reference terms - # stay single-step. Critic stays single-step. Mirrors deploy-side - # ``ObservationManager`` with ``use_gym_history=false``. - obs_history_length: int = 1 - - -@dataclass -class ObsDomainRand(Domain_Rand): - """y / z COM offsets, per-episode encoder bias, foot-geom friction.""" - - randomize_com_y: bool = False - com_offset_y: list[float] = field(default_factory=lambda: [-0.05, 0.05]) - randomize_com_z: bool = False - com_offset_z: list[float] = field(default_factory=lambda: [-0.05, 0.05]) - - # Per-episode additive bias on actor's joint_pos channel. - enable_encoder_bias: bool = False - encoder_bias_range: list[float] = field(default_factory=lambda: [-0.01, 0.01]) - - # Per-reset foot-geom friction. ``shared_random=True`` — a single scalar - # is broadcast across all foot geoms of one env, applied to the - # sliding-friction column. Matches mjlab. - randomize_geom_friction: bool = False - friction_range: list[float] = field(default_factory=lambda: [0.3, 1.2]) - friction_geom_pattern: str = r"^(left|right)_foot[1-7]_collision$" - - -@registry.envcfg("G1WBTObs") -@dataclass -class G1WBTObsCfg(G1MotionTrackingSACCfg): - """SAC whole-body tracking with sim2real obs / DR / reward extensions.""" - - noise_config: ObsNoiseConfig = field(default_factory=ObsNoiseConfig) # type: ignore[assignment] - domain_rand: ObsDomainRand = field(default_factory=ObsDomainRand) # type: ignore[assignment] - - -# --------------------------------------------------------------------------- # -# DR provider extension -# --------------------------------------------------------------------------- # - - -class G1WBTObsDomainRandomizationProvider(G1MotionTrackingDomainRandomizationProvider): - """Extends the SAC tracking DR provider with encoder bias, foot-geom - friction, y/z COM offsets, and post-reset ``prev_dof_vel`` seeding.""" - - def __init__( - self, - *, - base_kp: np.ndarray | None = None, - base_kd: np.ndarray | None = None, - base_geom_friction: np.ndarray | None = None, - foot_geom_ids: np.ndarray | None = None, - ) -> None: - super().__init__(base_kp=base_kp, base_kd=base_kd) - self._base_geom_friction = base_geom_friction - self._foot_geom_ids = foot_geom_ids - - def validate(self, env: Any, capabilities: DomainRandomizationCapabilities) -> None: - super().validate(env, capabilities) - if not getattr(env.cfg.domain_rand, "randomize_geom_friction", False): - return - if not capabilities.supports_reset_term(RESET_TERM_GEOM_FRICTION): - raise NotImplementedError( - f"{env._backend.backend_type} backend does not support " - "geom-friction reset randomization" - ) - if ( - self._base_geom_friction is None - or self._foot_geom_ids is None - or self._foot_geom_ids.size == 0 - ): - raise ValueError("randomize_geom_friction=True but provider has no foot geom IDs") - - def build_reset_plan(self, env: Any, env_ids: np.ndarray) -> ResetPlan: - num_reset = len(env_ids) - motion_frames = env.motion_sampler.sample_frames(env_ids) - motion_data = env.motion_loader.get_motion_at_frame(motion_frames) - qpos, qvel = _build_motion_reference_state(env, env_ids, motion_data) - - info_updates: dict[str, Any] = { - "current_actions": zero_actions(num_reset, env._num_action), - "last_actions": zero_actions(num_reset, env._num_action), - # Seed prev_dof_vel with the post-reset joint velocity so the first - # joint_acc_l2 sample is physically meaningful (Δv from the new - # starting velocity, not a spurious step from pre-termination). - "prev_dof_vel": qvel[:, 6:].astype(get_global_dtype()), - } - - dr_cfg = env.cfg.domain_rand - if getattr(dr_cfg, "enable_encoder_bias", False): - low, high = dr_cfg.encoder_bias_range - info_updates["joint_pos_obs_bias"] = np.random.uniform( - low, high, size=(num_reset, env._num_action) - ).astype(get_global_dtype()) - - randomization = build_common_reset_randomization( - env, num_reset, base_kp=self._base_kp, base_kd=self._base_kd - ) - - # Foot-geom friction. - if getattr(dr_cfg, "randomize_geom_friction", False): - assert self._base_geom_friction is not None - assert self._foot_geom_ids is not None - payload = randomization or ResetRandomizationPayload() - low, high = dr_cfg.friction_range - scale = np.random.uniform(low, high, size=(num_reset, 1)).astype(np.float64) - geom_friction = np.broadcast_to( - self._base_geom_friction, - (num_reset, *self._base_geom_friction.shape), - ).copy() - geom_friction[:, self._foot_geom_ids, 0] = scale * np.ones( - (1, self._foot_geom_ids.size) - ) - payload.geom_friction = geom_friction - randomization = payload - - # y / z COM offsets, layered on top of parent's x-only common build. - has_com_y = getattr(dr_cfg, "randomize_com_y", False) - has_com_z = getattr(dr_cfg, "randomize_com_z", False) - if has_com_y or has_com_z: - payload = randomization or ResetRandomizationPayload() - com_offset = payload.base_com_offset - if com_offset is None: - com_offset = np.zeros((num_reset, 3), dtype=np.float64) - if has_com_y: - low, high = dr_cfg.com_offset_y - com_offset[:, 1] = np.random.uniform(low, high, size=(num_reset,)) - if has_com_z: - low, high = dr_cfg.com_offset_z - com_offset[:, 2] = np.random.uniform(low, high, size=(num_reset,)) - payload.base_com_offset = com_offset - randomization = payload - - return ResetPlan( - env_ids=env_ids, - qpos=qpos, - qvel=qvel, - info_updates=info_updates, - randomization=randomization, - ) - - -# --------------------------------------------------------------------------- # -# Env -# --------------------------------------------------------------------------- # - - -@registry.env("G1WBTObs", sim_backend="mujoco") -@registry.env("G1WBTObs", sim_backend="motrix") -class G1WBTObsEnv(G1MotionTrackingSACEnv): - """SAC WBT with deploy-aligned obs, proprio history, and extra DR/rewards. - - All extensions live in this subclass — base classes are untouched. Flags - on ``G1WBTObsCfg`` are toggled from the task yaml. - """ - - _cfg: G1WBTObsCfg # pyright: ignore[reportIncompatibleVariableOverride] - - def __init__(self, cfg: G1WBTObsCfg, num_envs: int = 1, backend_type: str = "mujoco"): - super().__init__(cfg, num_envs=num_envs, backend_type=backend_type) - - # Cache base actuator gains for joint_torque_l2. - # Position-control torque approx: τ ≈ kp·(target_q − q) − kd·qd. - # DR (kp/kd ±10–15%) leaves small error vs true per-env torque, but - # the gradient direction (penalise large action / large Δq) is preserved. - base_kp, base_kd = self._backend.get_actuator_gains() - self._base_kp = np.asarray(base_kp, dtype=get_global_dtype()) - self._base_kd = np.asarray(base_kd, dtype=get_global_dtype()) - - # Proprio history buffers — per-term, oldest-first. Allocated only when - # H > 1 so H = 1 is zero-overhead. - H = max(1, int(cfg.noise_config.obs_history_length)) - self._hist_len = H - self._hist_buf: dict[str, np.ndarray] | None = None - if H > 1: - n = self._num_action - dtype = get_global_dtype() - self._hist_buf = { - "gyro": np.zeros((num_envs, H, 3), dtype=dtype), - "joint_pos_rel": np.zeros((num_envs, H, n), dtype=dtype), - "dof_vel": np.zeros((num_envs, H, n), dtype=dtype), - "last_actions": np.zeros((num_envs, H, n), dtype=dtype), - } - # Plumbs ``info`` from ``_compute_obs`` down to ``_build_actor_obs`` - # without changing the base-class hook signature. - self._obs_compute_info: dict | None = None - - # Swap to the extended DR provider whenever an extended flag is on. - dr_cfg = cfg.domain_rand - needs_extended = ( - getattr(dr_cfg, "enable_encoder_bias", False) - or getattr(dr_cfg, "randomize_geom_friction", False) - or getattr(dr_cfg, "randomize_com_y", False) - or getattr(dr_cfg, "randomize_com_z", False) - ) - if needs_extended: - kp = self._base_kp if (dr_cfg.randomize_kp or dr_cfg.randomize_kd) else None - kd = self._base_kd if (dr_cfg.randomize_kp or dr_cfg.randomize_kd) else None - base_geom_friction = None - foot_geom_ids = None - if dr_cfg.randomize_geom_friction: - base_geom_friction = self._backend.get_geom_friction() - geom_names = self._backend.get_geom_names() - pattern = re.compile(dr_cfg.friction_geom_pattern) - foot_geom_ids = np.asarray( - [i for i, name in enumerate(geom_names) if name and pattern.match(name)], - dtype=np.int64, - ) - if foot_geom_ids.size == 0: - raise ValueError( - "friction_geom_pattern " - f"'{dr_cfg.friction_geom_pattern}' did not match any geom" - ) - extended_provider = G1WBTObsDomainRandomizationProvider( - base_kp=kp, - base_kd=kd, - base_geom_friction=base_geom_friction, - foot_geom_ids=foot_geom_ids, - ) - # Swap the per-reset DR provider directly. ``_init_domain_randomization`` - # cannot be called twice — it materializes the backend at the end and - # MuJoCo's pool raises on a second materialize. The parent's call - # already (a) ran init randomization and (b) materialized; we only - # need the new provider's ``build_reset_plan`` for per-episode DR. - from unilab.dr import DomainRandomizationManager - - self._dr_manager = DomainRandomizationManager(self, extended_provider) - - # ------------------------------------------------------------------ # - # Rewards - # ------------------------------------------------------------------ # - - def _init_reward_functions(self) -> None: - super()._init_reward_functions() - self._reward_fns["joint_acc_l2"] = self._reward_joint_acc_l2 - self._reward_fns["joint_torque_l2"] = self._reward_joint_torque_l2 - - def _reward_joint_acc_l2(self, ctx: RewardContext) -> np.ndarray: - dof_vel = ctx.dof_vel - prev_dof_vel = ctx.info.get("prev_dof_vel") - if dof_vel is None: - raise RuntimeError("RewardContext.dof_vel is required for joint_acc_l2") - if prev_dof_vel is None or prev_dof_vel.shape != dof_vel.shape: - return np.zeros((self._num_envs,), dtype=get_global_dtype()) - joint_acc = (dof_vel - prev_dof_vel) / self._cfg.ctrl_dt - return np.asarray(np.sum(np.square(joint_acc), axis=1), dtype=get_global_dtype()) - - def _reward_joint_torque_l2(self, ctx: RewardContext) -> np.ndarray: - dof_pos = ctx.dof_pos - dof_vel = ctx.dof_vel - last_actions = ctx.info.get("last_actions") - if dof_pos is None or dof_vel is None: - raise RuntimeError("RewardContext.dof_pos and dof_vel are required for joint_torque_l2") - if last_actions is None: - return np.zeros((self._num_envs,), dtype=get_global_dtype()) - target_q = ( - last_actions * self._cfg.control_config.action_scale + self._effective_default_angles() - ) - torque = self._base_kp * (target_q - dof_pos) - self._base_kd * dof_vel - return np.asarray(np.sum(np.square(torque), axis=1), dtype=get_global_dtype()) - - # ------------------------------------------------------------------ # - # Obs - # ------------------------------------------------------------------ # - - def _actor_obs_dim(self, n: int) -> int: - nc = self._cfg.noise_config - H = max(1, int(nc.obs_history_length)) - single_step = 2 * n + 6 # command(2n) + anchor_ori(6) - if not nc.enable_zero_anchor_pos: - single_step += 3 - if not nc.enable_zero_linvel: - single_step += 3 - proprio_step = 3 + 3 * n # gyro + joint_pos_rel + dof_vel + last_actions - return single_step + H * proprio_step - - def _compute_obs( - self, - info: dict, - motion_data: Any, - linvel: np.ndarray, - gyro: np.ndarray, - dof_pos: np.ndarray, - dof_vel: np.ndarray, - robot_body_pos_w: np.ndarray, - robot_body_quat_w: np.ndarray, - ) -> dict[str, np.ndarray]: - # Stash so the overridden ``_build_actor_obs`` (called inside super) - # can read env_ids / joint_pos_obs_bias without a signature change. - self._obs_compute_info = info - try: - obs = super()._compute_obs( - info, - motion_data, - linvel, - gyro, - dof_pos, - dof_vel, - robot_body_pos_w, - robot_body_quat_w, - ) - finally: - self._obs_compute_info = None - - # Cache for next-step joint_acc_l2. The reset path overwrites this - # via the DR provider's ``prev_dof_vel`` info_update. - info["prev_dof_vel"] = dof_vel.copy() - return obs - - def _build_actor_obs( - self, - *, - command: np.ndarray, - motion_anchor_pos_b: np.ndarray, - motion_anchor_ori_b: np.ndarray, - noisy_linvel: np.ndarray, - noisy_gyro: np.ndarray, - noisy_joint_pos_rel: np.ndarray, - noisy_dof_vel: np.ndarray, - last_actions: np.ndarray, - ) -> np.ndarray: - info = self._obs_compute_info or {} - # Reset path is signalled by ``env_ids`` in obs_info (set by parent's - # ``_refresh_observation_rows`` and the DR provider's - # ``build_reset_observation``). In that case fill history slots; in - # the per-step path we push (oldest out, current in). - env_ids = info.get("env_ids") - is_reset = env_ids is not None - - nc = self._cfg.noise_config - - # Per-episode encoder bias on actor's joint_pos channel. - bias = info.get("joint_pos_obs_bias") - if bias is not None and bias.shape == noisy_joint_pos_rel.shape: - noisy_joint_pos_rel = np.asarray( - noisy_joint_pos_rel + bias, dtype=noisy_joint_pos_rel.dtype - ) - - # Per-step anchor_ori noise (actor only). - actor_anchor_ori_b = motion_anchor_ori_b - if nc.enable_anchor_ori_noise: - actor_anchor_ori_b = self._obs_noise(motion_anchor_ori_b, nc.scale_anchor_ori) - - # Single-step reference terms, dropping deploy-unavailable channels. - actor_terms: list[np.ndarray] = [command] - if not nc.enable_zero_anchor_pos: - actor_terms.append(motion_anchor_pos_b) - actor_terms.append(actor_anchor_ori_b) - if not nc.enable_zero_linvel: - actor_terms.append(noisy_linvel) - - # Proprio history (or single-step pass-through when H = 1). - if self._hist_buf is not None: - components = { - "gyro": noisy_gyro, - "joint_pos_rel": noisy_joint_pos_rel, - "dof_vel": noisy_dof_vel, - "last_actions": last_actions, - } - if is_reset: - self._fill_obs_history(env_ids, components) - else: - self._push_obs_history(env_ids, components) - sel = slice(None) if env_ids is None else env_ids - for key in ("gyro", "joint_pos_rel", "dof_vel", "last_actions"): - buf = self._hist_buf[key][sel] # (n_e, H, D) - actor_terms.append(buf.reshape(buf.shape[0], -1)) - else: - actor_terms.extend([noisy_gyro, noisy_joint_pos_rel, noisy_dof_vel, last_actions]) - - return np.concatenate(actor_terms, axis=1, dtype=get_global_dtype()) - - # ------------------------------------------------------------------ # - # Proprio history buffer maintenance. - # Mirrors deploy ``ObservationManager`` / ``ObservationTermCfg``: - # * On reset: fill all H slots with the current value (matches - # ``ObservationTermCfg::reset`` which calls ``add()`` H times). - # * On step: pop oldest, push current at end. - # * Read order is oldest-first, so ``flatten(buf[env, :, :])`` yields - # ``[t-H+1, t-H+2, ..., t]`` — matches deploy - # ``ObservationTermCfg::get`` (deque front-to-back). - # ------------------------------------------------------------------ # - - def _push_obs_history( - self, env_ids: np.ndarray | None, components: dict[str, np.ndarray] - ) -> None: - if self._hist_buf is None: - return - sel = slice(None) if env_ids is None else env_ids - for key, val in components.items(): - buf = self._hist_buf[key] - buf[sel, :-1] = buf[sel, 1:] - buf[sel, -1] = val - - def _fill_obs_history( - self, env_ids: np.ndarray | None, components: dict[str, np.ndarray] - ) -> None: - if self._hist_buf is None: - return - sel = slice(None) if env_ids is None else env_ids - for key, val in components.items(): - self._hist_buf[key][sel, :] = val[:, None, :] - - -@registry.envcfg("G1WBTObs23Dof") -@dataclass -class G1WBTObs23DofCfg(G1WBTObsCfg): - motion_file: str | list[str] = str( - ASSETS_ROOT_PATH / "motions" / "g1" / "dance1_subject2_part_23dof.npz" - ) - scene: SceneCfg = field( - default_factory=lambda: SceneCfg( - model_file=str(ASSETS_ROOT_PATH / "robots" / "g1" / "scene_flat_23dof.xml") - ) - ) - body_names: tuple[str, ...] = ( - "pelvis", - "left_hip_roll_link", - "left_knee_link", - "left_ankle_roll_link", - "right_hip_roll_link", - "right_knee_link", - "right_ankle_roll_link", - "torso_link", - "left_shoulder_roll_link", - "left_elbow_link", - "left_wrist_roll_rubber_hand", - "right_shoulder_roll_link", - "right_elbow_link", - "right_wrist_roll_rubber_hand", - ) - ee_body_names: tuple[str, ...] = ( - "left_ankle_roll_link", - "right_ankle_roll_link", - "left_wrist_roll_rubber_hand", - "right_wrist_roll_rubber_hand", - ) - - -registry.register_env("G1WBTObs23Dof", G1WBTObsEnv, sim_backend="mujoco") -registry.register_env("G1WBTObs23Dof", G1WBTObsEnv, sim_backend="motrix") diff --git a/src/unilab/tasks/motion_tracking/g1/tracking_sac.py b/src/unilab/tasks/motion_tracking/g1/tracking_sac.py deleted file mode 100644 index a7e3c5e36..000000000 --- a/src/unilab/tasks/motion_tracking/g1/tracking_sac.py +++ /dev/null @@ -1,88 +0,0 @@ -"""G1 Motion Tracking SAC Environment — thin SAC wrapper over G1MotionTrackingEnv. - -Differences from the PPO base: -- Critic observations additionally include ``base_lin_vel`` (3 dims), - matching holosoma's asymmetric actor-critic design for WBT. -- Registered under a separate name so it can be paired with FastSAC - configs without affecting the PPO motion-tracking pipeline. -""" - -from __future__ import annotations - -from dataclasses import dataclass - -import numpy as np - -from unilab.base import registry -from unilab.dtype_config import get_global_dtype -from unilab.envs import ManagerBasedRlEnvCfg, make_manager_based_rl_env - -from .tracking import ( - G1MotionTracking23DofCfg, - G1MotionTrackingCfg, - G1MotionTrackingEnv, -) - - -@dataclass -class G1MotionTrackingSACCfg(G1MotionTrackingCfg): - """Config for SAC-based motion tracking (identical fields, separate registry entry).""" - - -class G1MotionTrackingSACEnv(G1MotionTrackingEnv): - """G1 Motion Tracking environment for FastSAC training. - - Extends the PPO motion-tracking environment with ``base_lin_vel`` - appended to the critic observation, matching holosoma's asymmetric - actor-critic WBT design. - - The motrix backend is registered for sim2sim eval/playback only — checkpoints - trained on mujoco can be replayed via motrix's native renderer through - ``eval --sim motrix``. - """ - - @property - def obs_groups_spec(self) -> dict[str, int]: - spec = super().obs_groups_spec - # Append base_lin_vel (3) to critic observations. - return {**spec, "critic": spec["critic"] + 3} - - def _compute_obs( - self, - info: dict, - motion_data, - linvel: np.ndarray, - gyro: np.ndarray, - dof_pos: np.ndarray, - dof_vel: np.ndarray, - robot_body_pos_w: np.ndarray, - robot_body_quat_w: np.ndarray, - ) -> dict[str, np.ndarray]: - obs = super()._compute_obs( # pyright: ignore[reportAttributeAccessIssue] - info, - motion_data, - linvel, - gyro, - dof_pos, - dof_vel, - robot_body_pos_w, - robot_body_quat_w, - ) - # Append base_lin_vel to critic observations. - obs["critic"] = np.concatenate([obs["critic"], linvel], axis=1, dtype=get_global_dtype()) # type: ignore[call-overload] - return obs - - -@dataclass -class G1MotionTrackingSAC23DofCfg(G1MotionTracking23DofCfg): - pass - - -class G1MotionTrackingSAC23DofEnv(G1MotionTrackingSACEnv): - _cfg: G1MotionTrackingSAC23DofCfg # pyright: ignore[reportIncompatibleVariableOverride] - - -for _task_name in ("G1MotionTrackingSAC", "G1MotionTrackingSAC23Dof"): - registry.register_env_config(_task_name, ManagerBasedRlEnvCfg) - registry.register_env(_task_name, make_manager_based_rl_env, sim_backend="mujoco") - registry.register_env(_task_name, make_manager_based_rl_env, sim_backend="motrix") diff --git a/src/unilab/tasks/motion_tracking/x2/__init__.py b/src/unilab/tasks/motion_tracking/x2/__init__.py index c553d13e2..c8612d854 100644 --- a/src/unilab/tasks/motion_tracking/x2/__init__.py +++ b/src/unilab/tasks/motion_tracking/x2/__init__.py @@ -1,15 +1,23 @@ -"""Motion tracking environments for AgiBot X2.""" - -from .flip_tracking import ( - X2MotionTrackingCfg, - X2WallFlipTrackingCfg, - X2WallFlipTrackingEnv, - X2WallFlipTrackingEnvCfg, -) - -__all__ = [ - "X2MotionTrackingCfg", - "X2WallFlipTrackingCfg", - "X2WallFlipTrackingEnv", - "X2WallFlipTrackingEnvCfg", -] +"""AgiBot X2 motion profiles on the shared NumPy Manager-Based runtime.""" + +from unilab.assets.hub import resolve_robot_asset_dir +from unilab.base import registry +from unilab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg, make_manager_based_rl_env + + +def make_x2_wall_flip_env( + cfg: ManagerBasedRlEnvCfg, + num_envs: int = 1, + backend_type: str = "mujoco", +) -> ManagerBasedRlEnv: + """Resolve untracked X2 meshes before backend scene materialization.""" + resolve_robot_asset_dir("robots/x2/meshes", marker="pelvis.STL") + return make_manager_based_rl_env(cfg, num_envs=num_envs, backend_type=backend_type) + + +registry.register_env_config("X2WallFlipTracking", ManagerBasedRlEnvCfg) +registry.register_env("X2WallFlipTracking", make_x2_wall_flip_env, sim_backend="mujoco") +registry.register_env("X2WallFlipTracking", make_x2_wall_flip_env, sim_backend="motrix") + + +__all__ = ["make_x2_wall_flip_env"] diff --git a/src/unilab/tasks/motion_tracking/x2/flip_tracking.py b/src/unilab/tasks/motion_tracking/x2/flip_tracking.py deleted file mode 100644 index a72b9e0fc..000000000 --- a/src/unilab/tasks/motion_tracking/x2/flip_tracking.py +++ /dev/null @@ -1,135 +0,0 @@ -"""X2 motion-tracking profiles backed by the shared humanoid tracker.""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from typing import Literal - -from unilab.assets import ASSETS_ROOT_PATH -from unilab.base import registry -from unilab.base.scene import SceneCfg -from unilab.tasks.motion_tracking.common.g1_base import Sensor - -from ..common.config import ( - PoseRandomization, - VelocityRandomization, - _zero_pose_randomization, - _zero_velocity_randomization, -) -from ..common.tracking import ( - MotionTrackingDeployEnv, - MotionTrackingDeployEnvCfg, -) - - -@dataclass -class X2MotionTrackingCfg(MotionTrackingDeployEnvCfg): - """Base X2 motion-tracking config profile.""" - - scene: SceneCfg = field( - default_factory=lambda: SceneCfg( - model_file=str(ASSETS_ROOT_PATH / "robots" / "x2" / "scene_flat.xml") - ) - ) - motion_file: str | list[str] = str( - ASSETS_ROOT_PATH / "motions" / "x2" / "tictacflip_6-3_g1format.npz" - ) - sensor: Sensor = field( - default_factory=lambda: Sensor( - local_linvel="body-linear-vel", - gyro="body-angular-velocity", - upvector="body-orientation", - ) - ) - anchor_body_name: str = "torso_link" - body_names: tuple[str, ...] = ( - "pelvis", - "left_hip_pitch_link", - "left_hip_roll_link", - "left_hip_yaw_link", - "left_knee_link", - "left_ankle_pitch_link", - "left_ankle_roll_link", - "right_hip_pitch_link", - "right_hip_roll_link", - "right_hip_yaw_link", - "right_knee_link", - "right_ankle_pitch_link", - "right_ankle_roll_link", - "waist_yaw_link", - "waist_pitch_link", - "torso_link", - "left_shoulder_pitch_link", - "left_shoulder_roll_link", - "left_shoulder_yaw_link", - "left_elbow_link", - "left_wrist_yaw_link", - "left_wrist_pitch_link", - "left_wrist_roll_link", - "right_shoulder_pitch_link", - "right_shoulder_roll_link", - "right_shoulder_yaw_link", - "right_elbow_link", - "right_wrist_yaw_link", - "right_wrist_pitch_link", - "right_wrist_roll_link", - ) - ee_body_names: tuple[str, ...] = ( - "left_ankle_roll_link", - "right_ankle_roll_link", - "left_wrist_yaw_link", - "right_wrist_yaw_link", - ) - - -@dataclass -class X2WallFlipTrackingCfg(X2MotionTrackingCfg): - """Config profile for wall-assisted X2 flip tracking.""" - - scene: SceneCfg = field( - default_factory=lambda: SceneCfg( - model_file=str(ASSETS_ROOT_PATH / "robots" / "x2" / "scene_flat_with_wall.xml"), - # Render-only twin: wall as a worldbody geom so the offline grid - # renderer replicates it under every env cell (matches G1's per-env - # wall view). Physics still uses the wall in model_file. - visual_model_file=str( - ASSETS_ROOT_PATH / "robots" / "x2" / "scene_flat_with_wall_visual.xml" - ), - ) - ) - pose_randomization: PoseRandomization = field(default_factory=_zero_pose_randomization) - velocity_randomization: VelocityRandomization = field( - default_factory=_zero_velocity_randomization - ) - joint_position_range: tuple[float, float] = (0.0, 0.0) - sampling_mode: Literal["start", "clip_start", "uniform", "adaptive", "mixed"] = "adaptive" - truncate_on_clip_end: bool = False - terminate_on_undesired_contacts: bool = True - anchor_ori_threshold: float = 1e9 - anchor_pos_z_threshold: float = 0.5 - ee_body_pos_z_threshold: float = 0.5 - - -@registry.envcfg("X2WallFlipTracking") -@dataclass -class X2WallFlipTrackingEnvCfg(X2WallFlipTrackingCfg): - """Registered configuration for X2 wall flip tracking.""" - - pass - - -@registry.env("X2WallFlipTracking", sim_backend="mujoco") -@registry.env("X2WallFlipTracking", sim_backend="motrix") -class X2WallFlipTrackingEnv(MotionTrackingDeployEnv): - """X2 wall flip-tracking environment implementation.""" - - _cfg: X2WallFlipTrackingCfg # pyright: ignore[reportIncompatibleVariableOverride] - _keyframe_name = "home" - - def __init__(self, cfg: X2WallFlipTrackingCfg, num_envs: int = 1, backend_type: str = "mujoco"): - # X2 meshes are hosted on Hugging Face, not committed to git. Ensure they - # exist locally before the scene XML is parsed by the simulation backend. - from unilab.assets.hub import resolve_robot_asset_dir - - resolve_robot_asset_dir("robots/x2/meshes", marker="pelvis.STL") - super().__init__(cfg, num_envs=num_envs, backend_type=backend_type) diff --git a/src/unilab/training/backend_adapter.py b/src/unilab/training/backend_adapter.py index fc16ce948..f26aed062 100644 --- a/src/unilab/training/backend_adapter.py +++ b/src/unilab/training/backend_adapter.py @@ -2,6 +2,7 @@ from __future__ import annotations +from dataclasses import replace from pathlib import Path from typing import Any, Callable @@ -69,19 +70,29 @@ def build_play_env_cfg_override(self) -> dict[str, Any]: if not source_model_file: raise ValueError("play_profile.scene.source_model_file must be configured") - env_cfg_override["scene"] = SceneCfg( - model_file=self.scene_materializer( - self._resolve_root_relative_path(str(source_model_file)), - ground_texture_file=( - self._resolve_root_relative_path(str(scene_override.ground_texture_file)) - if getattr(scene_override, "ground_texture_file", None) - else None - ), - ground_texrepeat=getattr(scene_override, "ground_texrepeat", None), - skybox_rgb1=getattr(scene_override, "skybox_rgb1", None), - skybox_rgb2=getattr(scene_override, "skybox_rgb2", None), - ) + materialized_model_file = self.scene_materializer( + self._resolve_root_relative_path(str(source_model_file)), + ground_texture_file=( + self._resolve_root_relative_path(str(scene_override.ground_texture_file)) + if getattr(scene_override, "ground_texture_file", None) + else None + ), + ground_texrepeat=getattr(scene_override, "ground_texrepeat", None), + skybox_rgb1=getattr(scene_override, "skybox_rgb1", None), + skybox_rgb2=getattr(scene_override, "skybox_rgb2", None), ) + scene = env_cfg_override.get("scene") + if scene is None: + env_cfg_override["scene"] = SceneCfg(model_file=materialized_model_file) + elif isinstance(scene, SceneCfg): + env_cfg_override["scene"] = replace(scene, model_file=materialized_model_file) + elif isinstance(scene, dict): + env_cfg_override["scene"] = {**scene, "model_file": materialized_model_file} + else: + raise TypeError( + "play_profile.scene can only override a missing, SceneCfg, or mapping scene; " + f"got {type(scene).__name__}" + ) return env_cfg_override def _apply_env_profile(self, env_cfg_override: dict[str, Any], env_profile: Any) -> None: diff --git a/tests/base/backend/test_mujoco_scene_context_visual.py b/tests/base/backend/test_mujoco_scene_context_visual.py index 2aa276f88..b8af7ca43 100644 --- a/tests/base/backend/test_mujoco_scene_context_visual.py +++ b/tests/base/backend/test_mujoco_scene_context_visual.py @@ -35,9 +35,21 @@ def test_visual_model_file_override_is_render_only() -> None: def test_x2_wall_flip_wires_render_only_visual_twin() -> None: - from unilab.tasks.motion_tracking.x2.flip_tracking import X2WallFlipTrackingEnvCfg - - cfg = X2WallFlipTrackingEnvCfg() + from hydra import compose, initialize_config_dir + + from unilab.base import registry + from unilab.base.config_materialization import apply_cfg_overrides + from unilab.training.backend_adapter import BackendAdapter + + repo_root = Path(__file__).parents[3] + with initialize_config_dir(config_dir=str(repo_root / "conf" / "ppo"), version_base="1.3"): + owner = compose("config", overrides=["task=x2_wall_flip_tracking/mujoco"]) + registry.ensure_registries() + cfg = registry.materialize_env_config("X2WallFlipTracking") + apply_cfg_overrides( + cfg, + BackendAdapter(owner, root_dir=repo_root, algo_name="ppo").build_task_env_cfg_override(), + ) # Physics = trained wall; render = worldbody-geom twin. assert cfg.scene.model_file.endswith("scene_flat_with_wall.xml") diff --git a/tests/config/test_locomotion_params.py b/tests/config/test_locomotion_params.py index cffa86828..ae2e50f3e 100644 --- a/tests/config/test_locomotion_params.py +++ b/tests/config/test_locomotion_params.py @@ -461,10 +461,12 @@ def test_ppo_g1_box_tracking(): assert cfg.training.task_name == "G1BoxTracking" assert cfg.algo.max_iterations == 30000 assert cfg.algo.algorithm.entropy_coef == pytest.approx(0.005) - assert cfg.reward.scales.object_global_ref_position_error_exp == pytest.approx(2.0) - assert cfg.reward.scales.object_global_ref_orientation_error_exp == pytest.approx(2.0) - assert cfg.reward.std_object_pos == pytest.approx(0.2) - assert cfg.reward.std_object_ori == pytest.approx(0.3) + assert cfg.env.scene.entities.object.root_body_name == "largebox" + assert cfg.env.commands.motion.object_entity_name == "object" + assert cfg.reward.object_global_ref_position_error_exp.weight == pytest.approx(2.0) + assert cfg.reward.object_global_ref_orientation_error_exp.weight == pytest.approx(2.0) + assert cfg.reward.object_global_ref_position_error_exp.params.std == pytest.approx(0.2) + assert cfg.reward.object_global_ref_orientation_error_exp.params.std == pytest.approx(0.3) def test_ppo_g1_flip_tracking(): @@ -481,19 +483,20 @@ def test_ppo_g1_flip_tracking(): assert cfg.algo.obs_groups.critic == ["critic"] assert cfg.algo.algorithm.entropy_coef == pytest.approx(0.005) assert cfg.algo.algorithm.desired_kl == pytest.approx(0.01) - assert cfg.env.sampling_mode == "start" - assert cfg.env.truncate_on_clip_end is False + assert cfg.env.commands.motion.params.sampling_mode == "start" + assert cfg.env.commands.motion.params.truncate_on_clip_end is False assert cfg.env.sim_dt == pytest.approx(0.005) - assert list(cfg.env.control_config.action_scale) == pytest.approx(G1_BEYONDMIMIC_ACTION_SCALE) - assert cfg.env.anchor_pos_z_threshold == pytest.approx(0.5) - assert cfg.env.ee_body_pos_z_threshold == pytest.approx(0.5) - assert cfg.env.terminate_on_undesired_contacts is True - assert cfg.env.noise_config.level == pytest.approx(0.0) - assert cfg.reward.scales.motion_body_pos == pytest.approx(2.0) - assert cfg.reward.scales.motion_body_ori == pytest.approx(1.5) - assert cfg.reward.scales.motion_ee_body_pos_z == pytest.approx(2.0) - assert cfg.reward.scales.action_rate_l2 == pytest.approx(-0.005) - assert cfg.reward.scales.undesired_contacts == pytest.approx(-0.1) + assert cfg.env.actions.joint_pos.scale[".*_(hip_pitch|hip_yaw)_joint"] == pytest.approx( + G1_BEYONDMIMIC_ACTION_SCALE[0] + ) + assert cfg.env.terminations.anchor_pos.params.threshold == pytest.approx(0.5) + assert cfg.env.terminations.ee_body_pos.params.threshold == pytest.approx(0.5) + assert cfg.env.terminations.undesired_contacts is not None + assert cfg.reward.motion_body_pos.weight == pytest.approx(2.0) + assert cfg.reward.motion_body_ori.weight == pytest.approx(1.5) + assert cfg.reward.motion_ee_body_pos_z.weight == pytest.approx(2.0) + assert cfg.reward.action_rate_l2.weight == pytest.approx(-0.005) + assert cfg.reward.undesired_contacts.weight == pytest.approx(-0.1) def test_ppo_g1_wall_flip_tracking(): @@ -510,21 +513,20 @@ def test_ppo_g1_wall_flip_tracking(): assert cfg.algo.obs_groups.critic == ["critic"] assert cfg.algo.algorithm.entropy_coef == pytest.approx(0.005) assert cfg.algo.algorithm.desired_kl == pytest.approx(0.01) - assert cfg.env.sampling_mode == "start" - assert cfg.env.truncate_on_clip_end is False + assert cfg.env.commands.motion.params.sampling_mode == "start" + assert cfg.env.commands.motion.params.truncate_on_clip_end is False assert cfg.env.sim_dt == pytest.approx(0.005) - assert list(cfg.env.control_config.action_scale) == pytest.approx(G1_BEYONDMIMIC_ACTION_SCALE) - assert cfg.env.anchor_pos_z_threshold == pytest.approx(0.5) - assert cfg.env.ee_body_pos_z_threshold == pytest.approx(0.5) - assert cfg.env.terminate_on_undesired_contacts is True - assert cfg.env.noise_config.level == pytest.approx(0.0) - assert cfg.reward.scales.motion_joint_pos == pytest.approx(0.5) - assert cfg.reward.scales.motion_joint_vel == pytest.approx(0.25) - assert cfg.reward.scales.motion_body_pos == pytest.approx(2.0) - assert cfg.reward.scales.motion_body_ori == pytest.approx(1.5) - assert cfg.reward.scales.motion_ee_body_pos_z == pytest.approx(2.0) - assert cfg.reward.scales.action_rate_l2 == pytest.approx(-0.005) - assert cfg.reward.scales.undesired_contacts == pytest.approx(-0.1) + assert cfg.env.actions.joint_pos.scale[".*_(hip_pitch|hip_yaw)_joint"] == pytest.approx( + G1_BEYONDMIMIC_ACTION_SCALE[0] + ) + assert cfg.env.scene.model_file.endswith("scene_flat_with_wall.xml") + assert cfg.reward.motion_joint_pos.weight == pytest.approx(0.5) + assert cfg.reward.motion_joint_vel.weight == pytest.approx(0.25) + assert cfg.reward.motion_body_pos.weight == pytest.approx(2.0) + assert cfg.reward.motion_body_ori.weight == pytest.approx(1.5) + assert cfg.reward.motion_ee_body_pos_z.weight == pytest.approx(2.0) + assert cfg.reward.action_rate_l2.weight == pytest.approx(-0.005) + assert cfg.reward.undesired_contacts.weight == pytest.approx(-0.1) def test_ppo_x2_wall_flip_tracking(): @@ -544,18 +546,17 @@ def test_ppo_x2_wall_flip_tracking(): assert cfg.algo.algorithm.desired_kl == pytest.approx(0.01) # Interactive playback defaults to policy mode for this task. assert cfg.interactive.action_mode == "policy" - assert cfg.env.sampling_mode == "start" - assert cfg.env.truncate_on_clip_end is False + assert cfg.env.commands.motion.params.sampling_mode == "start" + assert cfg.env.commands.motion.params.truncate_on_clip_end is False assert cfg.env.sim_dt == pytest.approx(0.005) - assert list(cfg.env.control_config.action_scale) == pytest.approx(X2_ACTION_SCALE) - assert cfg.env.anchor_pos_z_threshold == pytest.approx(0.5) - assert cfg.env.ee_body_pos_z_threshold == pytest.approx(0.5) - assert cfg.env.terminate_on_undesired_contacts is True - assert cfg.env.noise_config.level == pytest.approx(0.0) - assert cfg.reward.scales.motion_joint_pos == pytest.approx(0.5) - assert cfg.reward.scales.motion_joint_vel == pytest.approx(0.25) - assert cfg.reward.scales.motion_body_pos == pytest.approx(2.0) - assert cfg.reward.scales.motion_body_ori == pytest.approx(1.5) - assert cfg.reward.scales.motion_ee_body_pos_z == pytest.approx(2.0) - assert cfg.reward.scales.action_rate_l2 == pytest.approx(-0.005) - assert cfg.reward.scales.undesired_contacts == pytest.approx(-0.1) + assert cfg.env.actions.joint_pos.scale == pytest.approx(X2_ACTION_SCALE[0]) + assert cfg.env.terminations.anchor_pos.params.threshold == pytest.approx(0.5) + assert cfg.env.terminations.ee_body_pos.params.threshold == pytest.approx(0.5) + assert cfg.env.terminations.undesired_contacts is not None + assert cfg.reward.motion_joint_pos.weight == pytest.approx(0.5) + assert cfg.reward.motion_joint_vel.weight == pytest.approx(0.25) + assert cfg.reward.motion_body_pos.weight == pytest.approx(2.0) + assert cfg.reward.motion_body_ori.weight == pytest.approx(1.5) + assert cfg.reward.motion_ee_body_pos_z.weight == pytest.approx(2.0) + assert cfg.reward.action_rate_l2.weight == pytest.approx(-0.005) + assert cfg.reward.undesired_contacts.weight == pytest.approx(-0.1) diff --git a/tests/envs/test_env_configs.py b/tests/envs/test_env_configs.py index f8915967b..0f6710f0d 100644 --- a/tests/envs/test_env_configs.py +++ b/tests/envs/test_env_configs.py @@ -119,11 +119,6 @@ def blocked_import(name, globals=None, locals=None, fromlist=(), level=0): from unilab.base import registry from unilab.base.backend import create_backend - from unilab.tasks.motion_tracking.g1.tracking import ( - G1MotionTrackingCfg, - G1MotionTrackingDeployEnvCfg, - ) - from unilab.tasks.motion_tracking.x2 import X2WallFlipTrackingCfg from unilab.base.registry import ensure_registries ensure_registries() @@ -132,9 +127,9 @@ def blocked_import(name, globals=None, locals=None, fromlist=(), level=0): assert registry.contains("G1MotionTrackingDeploy") assert registry.contains("X2WallFlipTracking") assert registry.contains("AllegroInhandRotation") - G1MotionTrackingCfg() - G1MotionTrackingDeployEnvCfg() - X2WallFlipTrackingCfg() + metadata = registry.list_registered_envs() + assert metadata["G1MotionTracking"]["config_factory"] == "ManagerBasedRlEnvCfg" + assert metadata["X2WallFlipTracking"]["config_factory"] == "ManagerBasedRlEnvCfg" """ ) @@ -542,893 +537,6 @@ def test_allegro_grasp_recorder_close_autosaves_and_io_failure_is_fail_closed( assert failed.cache_saved is False -def test_g1_motion_tracking_uses_combined_body_pose_query(): - """G1MotionTracking should query pos/quat via the stable combined backend API.""" - from unilab.tasks.motion_tracking.g1.tracking import G1MotionTrackingEnv - - class FakeBackend: - def __init__(self) -> None: - self.calls: list[np.ndarray] = [] - - def get_body_pose_w(self, body_ids: np.ndarray) -> tuple[np.ndarray, np.ndarray]: - self.calls.append(body_ids.copy()) - return np.ones((2, len(body_ids), 3)), np.ones((2, len(body_ids), 4)) - - env = cast(Any, object.__new__(G1MotionTrackingEnv)) - env._backend = FakeBackend() - env.body_ids = np.array([1, 3], dtype=np.int32) - - pos_w, quat_w = env._get_body_pose_w() - - assert pos_w.shape == (2, 2, 3) - assert quat_w.shape == (2, 2, 4) - assert len(env._backend.calls) == 1 - np.testing.assert_array_equal(env._backend.calls[0], np.array([1, 3], dtype=np.int32)) - - -def test_g1_motion_tracking_reset_observation_uses_sparse_body_pose_rows(): - from unilab.tasks.motion_tracking.common.motion_loader import MotionData - from unilab.tasks.motion_tracking.g1.tracking import ( - G1MotionTrackingDomainRandomizationProvider, - ) - - class FakeBackend: - def __init__(self) -> None: - self.row_calls: list[tuple[np.ndarray, np.ndarray]] = [] - - def get_body_pose_w_rows( - self, env_ids: np.ndarray, body_ids: np.ndarray - ) -> tuple[np.ndarray, np.ndarray]: - self.row_calls.append((env_ids.copy(), body_ids.copy())) - rows = len(env_ids) - bodies = len(body_ids) - return ( - np.full((rows, bodies, 3), 2.0, dtype=np.float32), - np.full((rows, bodies, 4), 3.0, dtype=np.float32), - ) - - def get_body_pose_w(self, body_ids: np.ndarray) -> tuple[np.ndarray, np.ndarray]: - raise AssertionError("reset observation should use sparse body pose rows") - - class FakeMotionLoader: - def get_motion_at_frame(self, frames: np.ndarray) -> MotionData: - rows = len(frames) - return MotionData( - joint_pos=np.zeros((rows, 2), dtype=np.float32), - joint_vel=np.zeros((rows, 2), dtype=np.float32), - body_pos_w=np.zeros((rows, 2, 3), dtype=np.float32), - body_quat_w=np.zeros((rows, 2, 4), dtype=np.float32), - body_lin_vel_w=np.zeros((rows, 2, 3), dtype=np.float32), - body_ang_vel_w=np.zeros((rows, 2, 3), dtype=np.float32), - ) - - class FakeMotionSampler: - current_frames = np.array([10, 11, 12, 13], dtype=np.int32) - - env = SimpleNamespace( - _backend=FakeBackend(), - body_ids=np.array([1, 3], dtype=np.int32), - motion_loader=FakeMotionLoader(), - motion_sampler=FakeMotionSampler(), - get_local_linvel=lambda: np.zeros((4, 3), dtype=np.float32), - get_gyro=lambda: np.zeros((4, 3), dtype=np.float32), - get_dof_pos=lambda: np.zeros((4, 2), dtype=np.float32), - get_dof_vel=lambda: np.zeros((4, 2), dtype=np.float32), - ) - - captured: dict[str, np.ndarray] = {} - - def compute_obs( - obs_info, - motion_data, - linvel, - gyro, - dof_pos, - dof_vel, - robot_body_pos_w, - robot_body_quat_w, - ): - del obs_info, motion_data, linvel, gyro, dof_pos, dof_vel - captured["robot_body_pos_w"] = robot_body_pos_w - captured["robot_body_quat_w"] = robot_body_quat_w - return {"obs": np.zeros((2, 1), dtype=np.float32)} - - env._compute_obs = compute_obs - provider = G1MotionTrackingDomainRandomizationProvider() - env_ids = np.array([1, 3], dtype=np.int32) - - obs = provider.build_reset_observation(env, env_ids, {}) - - assert obs["obs"].shape == (2, 1) - assert len(env._backend.row_calls) == 1 - np.testing.assert_array_equal(env._backend.row_calls[0][0], env_ids) - np.testing.assert_array_equal(env._backend.row_calls[0][1], env.body_ids) - assert captured["robot_body_pos_w"].shape == (2, 2, 3) - assert captured["robot_body_quat_w"].shape == (2, 2, 4) - - -def _compute_g1_motion_tracking_obs_stub(env_cls: type): - from unilab.tasks.motion_tracking.common.motion_loader import MotionData - - env = cast(Any, object.__new__(env_cls)) - env._num_envs = 1 - env._num_action = 2 - env._n_motion_bodies = 2 - env._critic_obs_width = env._critic_base_obs_dim(env._num_action) + env._n_motion_bodies * 9 - env._cfg = SimpleNamespace( - noise_config=SimpleNamespace( - level=1.0, - scale_linvel=1.0, - scale_gyro=1.0, - scale_joint_angle=1.0, - scale_joint_vel=1.0, - ), - body_names=("pelvis", "torso_link"), - ) - env.default_angles = np.array([[0.5, -0.5]], dtype=np.float32) - env.anchor_body_idx = 0 - env._motion_anchor_pos_b = np.empty((1, 3), dtype=np.float32) - env._motion_anchor_ori_b = np.empty((1, 6), dtype=np.float32) - env._motion_command = np.empty((1, 4), dtype=np.float32) - env._joint_pos_rel = np.empty((1, 2), dtype=np.float32) - env._body_vec_error = np.empty((1, 2, 3), dtype=np.float32) - env._body_vec_tmp = np.empty((1, 2, 3), dtype=np.float32) - env._quat_error_w = np.empty((1, 2), dtype=np.float32) - env._quat_error_x = np.empty((1, 2), dtype=np.float32) - env._zero_actions = np.zeros((1, 2), dtype=np.float32) - env._obs_noise = lambda data, scale: np.asarray(data + 100.0, dtype=np.float32) - - motion_data = MotionData( - joint_pos=np.array([[0.1, 0.2]], dtype=np.float32), - joint_vel=np.array([[0.3, 0.4]], dtype=np.float32), - body_pos_w=np.zeros((1, 2, 3), dtype=np.float32), - body_quat_w=np.tile(np.array([[[1.0, 0.0, 0.0, 0.0]]], dtype=np.float32), (1, 2, 1)), - body_lin_vel_w=np.zeros((1, 2, 3), dtype=np.float32), - body_ang_vel_w=np.zeros((1, 2, 3), dtype=np.float32), - ) - linvel = np.array([[1.0, 2.0, 3.0]], dtype=np.float32) - gyro = np.array([[4.0, 5.0, 6.0]], dtype=np.float32) - dof_pos = np.array([[0.7, -0.2]], dtype=np.float32) - dof_vel = np.array([[7.0, 8.0]], dtype=np.float32) - robot_body_pos_w = np.array([[[0.0, 0.0, 0.0], [0.2, 0.0, 0.1]]], dtype=np.float32) - robot_body_quat_w = np.tile(np.array([[[1.0, 0.0, 0.0, 0.0]]], dtype=np.float32), (1, 2, 1)) - info = {"current_actions": np.array([[0.1, -0.2]], dtype=np.float32)} - - obs = env._compute_obs( - info, - motion_data, - linvel, - gyro, - dof_pos, - dof_vel, - robot_body_pos_w, - robot_body_quat_w, - ) - return env, obs, motion_data, linvel, gyro, dof_pos, dof_vel, info - - -def test_g1_motion_tracking_critic_uses_clean_beyondmimic_aligned_terms(): - from unilab.tasks.motion_tracking.g1.tracking import G1MotionTrackingEnv - - env, obs, motion_data, linvel, gyro, dof_pos, dof_vel, info = ( - _compute_g1_motion_tracking_obs_stub(G1MotionTrackingEnv) - ) - - assert env.obs_groups_spec == {"obs": 25, "critic": 43} - assert obs["obs"].shape == (1, 25) - np.testing.assert_allclose(obs["obs"][:, 13:16], linvel + 100.0) - np.testing.assert_allclose(obs["obs"][:, 16:19], gyro + 100.0) - np.testing.assert_allclose(obs["obs"][:, 19:21], dof_pos - env.default_angles + 100.0) - np.testing.assert_allclose(obs["obs"][:, 21:23], dof_vel + 100.0) - np.testing.assert_allclose(obs["obs"][:, 23:25], info["current_actions"]) - - command_dim = motion_data.joint_pos.shape[1] + motion_data.joint_vel.shape[1] - anchor_dim = 3 + 6 - clean_proprio_start = command_dim + anchor_dim - np.testing.assert_allclose( - obs["critic"][:, clean_proprio_start : clean_proprio_start + 3], linvel - ) - np.testing.assert_allclose( - obs["critic"][:, clean_proprio_start + 3 : clean_proprio_start + 6], gyro - ) - np.testing.assert_allclose( - obs["critic"][:, clean_proprio_start + 6 : clean_proprio_start + 8], - dof_pos - env.default_angles, - ) - np.testing.assert_allclose( - obs["critic"][:, clean_proprio_start + 8 : clean_proprio_start + 10], dof_vel - ) - np.testing.assert_allclose( - obs["critic"][:, clean_proprio_start + 10 : clean_proprio_start + 12], - info["current_actions"], - ) - - -def test_g1_motion_tracking_anchor_frame_writers_match_reference(): - from unilab.tasks.motion_tracking.g1.tracking import G1MotionTrackingEnv - from unilab.utils.rotation import ( - np_matrix_from_quat, - np_quat_apply, - np_quat_inv, - np_quat_mul, - ) - - rng = np.random.default_rng(123) - num_envs = 4 - num_bodies = 5 - dtype = np.float64 - - def random_quat(shape: tuple[int, ...]) -> np.ndarray: - quat = rng.normal(size=(*shape, 4)).astype(dtype) - quat /= np.linalg.norm(quat, axis=-1, keepdims=True) - return quat - - anchor_pos = rng.normal(size=(num_envs, 3)).astype(dtype) - anchor_quat = random_quat((num_envs,)) - body_pos = rng.normal(size=(num_envs, num_bodies, 3)).astype(dtype) - body_quat = random_quat((num_envs, num_bodies)) - - env = cast(Any, object.__new__(G1MotionTrackingEnv)) - env._body_vec_error = np.empty((num_envs, num_bodies, 3), dtype=dtype) - env._body_vec_tmp = np.empty((num_envs, num_bodies, 3), dtype=dtype) - env._quat_error_w = np.empty((num_envs, num_bodies), dtype=dtype) - env._quat_error_x = np.empty((num_envs, num_bodies), dtype=dtype) - - pos_out = np.empty((num_envs, num_bodies, 3), dtype=dtype) - ori_out = np.empty((num_envs, num_bodies, 6), dtype=dtype) - - env._write_body_pos_in_anchor_frame(anchor_pos, anchor_quat, body_pos, pos_out) - env._write_body_ori6_in_anchor_frame(anchor_quat, body_quat, ori_out) - - anchor_quat_inv = np_quat_inv(anchor_quat) - tiled_anchor_quat_inv = np.repeat(anchor_quat_inv, num_bodies, axis=0) - rel_pos_flat = (body_pos - anchor_pos[:, None, :]).reshape(num_envs * num_bodies, 3) - ref_pos = np_quat_apply(tiled_anchor_quat_inv, rel_pos_flat).reshape(num_envs, num_bodies, 3) - ref_quat = np_quat_mul( - tiled_anchor_quat_inv, - body_quat.reshape(num_envs * num_bodies, 4), - ) - ref_ori = np_matrix_from_quat(ref_quat)[:, :, :2].reshape(num_envs, num_bodies, 6) - - np.testing.assert_allclose(pos_out, ref_pos, rtol=1e-12, atol=1e-12) - np.testing.assert_allclose(ori_out, ref_ori, rtol=1e-12, atol=1e-12) - - -def test_g1_motion_tracking_relative_transform_fast_path_matches_reference(): - from unilab.tasks.motion_tracking.g1.tracking import G1MotionTrackingEnv - from unilab.utils.rotation import np_quat_apply, np_quat_inv, np_quat_mul, np_yaw_quat - - rng = np.random.default_rng(321) - num_envs = 4 - num_bodies = 5 - anchor_idx = 2 - dtype = np.float64 - - def random_quat(shape: tuple[int, ...]) -> np.ndarray: - quat = rng.normal(size=(*shape, 4)).astype(dtype) - quat /= np.linalg.norm(quat, axis=-1, keepdims=True) - return quat - - motion_data = SimpleNamespace( - body_pos_w=rng.normal(size=(num_envs, num_bodies, 3)).astype(dtype), - body_quat_w=random_quat((num_envs, num_bodies)), - ) - robot_body_pos_w = rng.normal(size=(num_envs, num_bodies, 3)).astype(dtype) - robot_body_quat_w = random_quat((num_envs, num_bodies)) - - env = cast(Any, object.__new__(G1MotionTrackingEnv)) - env.anchor_body_idx = anchor_idx - env.body_pos_relative_w = np.empty((num_envs, num_bodies, 3), dtype=dtype) - env.body_quat_relative_w = np.empty((num_envs, num_bodies, 4), dtype=dtype) - env._delta_pos_w = np.empty((num_envs, 3), dtype=dtype) - env._delta_ori_w = np.empty((num_envs, 4), dtype=dtype) - env._body_vec_error = np.empty((num_envs, num_bodies, 3), dtype=dtype) - env._env_error = np.empty((num_envs,), dtype=dtype) - env._reward_term = np.empty((num_envs,), dtype=dtype) - - env._update_relative_transforms(motion_data, robot_body_pos_w, robot_body_quat_w) - - anchor_pos_w = motion_data.body_pos_w[:, anchor_idx] - anchor_quat_w = motion_data.body_quat_w[:, anchor_idx] - robot_anchor_pos_w = robot_body_pos_w[:, anchor_idx] - robot_anchor_quat_w = robot_body_quat_w[:, anchor_idx] - delta_pos_w = robot_anchor_pos_w.copy() - delta_pos_w[:, 2] = anchor_pos_w[:, 2] - delta_ori_w = np_yaw_quat(np_quat_mul(robot_anchor_quat_w, np_quat_inv(anchor_quat_w))) - delta_ori_tiled = np.tile(delta_ori_w, (1, num_bodies)).reshape(num_envs * num_bodies, 4) - expected_quat = np_quat_mul( - delta_ori_tiled, - motion_data.body_quat_w.reshape(num_envs * num_bodies, 4), - ).reshape(num_envs, num_bodies, 4) - rel_pos_flat = (motion_data.body_pos_w - anchor_pos_w[:, None, :]).reshape( - num_envs * num_bodies, 3 - ) - expected_pos = delta_pos_w[:, None, :] + np_quat_apply(delta_ori_tiled, rel_pos_flat).reshape( - num_envs, num_bodies, 3 - ) - - np.testing.assert_allclose(env.body_pos_relative_w, expected_pos, rtol=1e-12, atol=1e-12) - np.testing.assert_allclose(env.body_quat_relative_w, expected_quat, rtol=1e-12, atol=1e-12) - - -def test_g1_motion_tracking_reward_fast_path_matches_reference(): - from unilab.tasks.motion_tracking.common.motion_loader import MotionData - from unilab.tasks.motion_tracking.g1.tracking import G1MotionTrackingEnv, RewardConfig - from unilab.utils.rotation import np_quat_error_magnitude - - rng = np.random.default_rng(456) - num_envs = 3 - num_bodies = 4 - num_actions = 2 - anchor_idx = 1 - ee_indices = np.array([2, 3], dtype=np.int32) - undesired_indices = np.array([0, 1], dtype=np.int32) - dtype = np.float64 - - def random_quat(shape: tuple[int, ...]) -> np.ndarray: - quat = rng.normal(size=(*shape, 4)).astype(dtype) - quat /= np.linalg.norm(quat, axis=-1, keepdims=True) - return quat - - reward_config = RewardConfig() - reward_config.scales = { - "motion_global_root_pos": 0.5, - "motion_global_root_ori": 0.25, - "motion_body_pos": 1.0, - "motion_body_ori": 0.75, - "motion_body_lin_vel": 0.4, - "motion_body_ang_vel": 0.3, - "motion_ee_body_pos_z": 0.2, - "motion_joint_pos": 0.6, - "motion_joint_vel": 0.7, - "action_rate_l2": -0.1, - "joint_limit": -0.2, - "undesired_contacts": -0.3, - } - ctrl_dt = 0.02 - contact_threshold = 0.05 - - motion_data = MotionData( - joint_pos=rng.normal(size=(num_envs, num_actions)).astype(dtype), - joint_vel=rng.normal(size=(num_envs, num_actions)).astype(dtype), - body_pos_w=rng.normal(size=(num_envs, num_bodies, 3)).astype(dtype), - body_quat_w=random_quat((num_envs, num_bodies)), - body_lin_vel_w=rng.normal(size=(num_envs, num_bodies, 3)).astype(dtype), - body_ang_vel_w=rng.normal(size=(num_envs, num_bodies, 3)).astype(dtype), - ) - robot_body_pos_w = rng.normal(size=(num_envs, num_bodies, 3)).astype(dtype) - robot_body_pos_w[:, undesired_indices, 2] = np.array( - [[0.01, 0.10], [0.20, 0.02], [0.07, 0.03]], dtype=dtype - ) - robot_body_quat_w = random_quat((num_envs, num_bodies)) - robot_body_lin_vel_w = rng.normal(size=(num_envs, num_bodies, 3)).astype(dtype) - robot_body_ang_vel_w = rng.normal(size=(num_envs, num_bodies, 3)).astype(dtype) - dof_pos = rng.normal(size=(num_envs, num_actions)).astype(dtype) - dof_vel = rng.normal(size=(num_envs, num_actions)).astype(dtype) - current_actions = rng.normal(size=(num_envs, num_actions)).astype(dtype) - last_actions = rng.normal(size=(num_envs, num_actions)).astype(dtype) - body_pos_relative_w = rng.normal(size=(num_envs, num_bodies, 3)).astype(dtype) - body_quat_relative_w = random_quat((num_envs, num_bodies)) - joint_lower = np.array([-0.2, -0.1], dtype=dtype) - joint_upper = np.array([0.1, 0.2], dtype=dtype) - - env = cast(Any, object.__new__(G1MotionTrackingEnv)) - env._num_envs = num_envs - env.anchor_body_idx = anchor_idx - env.ee_body_indices = ee_indices - env.undesired_contact_body_indices = undesired_indices - env._has_ee_body_indices = True - env._has_undesired_contact_body_indices = True - env._cfg = SimpleNamespace( - reward_config=reward_config, - ctrl_dt=ctrl_dt, - undesired_contact_z_threshold=contact_threshold, - ) - env.body_pos_relative_w = body_pos_relative_w.copy() - env.body_quat_relative_w = body_quat_relative_w.copy() - env._joint_lower = joint_lower - env._joint_upper = joint_upper - env._body_vec_error = np.empty((num_envs, num_bodies, 3), dtype=dtype) - env._joint_error = np.empty((num_envs, num_actions), dtype=dtype) - env._joint_error_upper = np.empty((num_envs, num_actions), dtype=dtype) - env._env_error = np.empty((num_envs,), dtype=dtype) - env._env_error2 = np.empty((num_envs,), dtype=dtype) - env._reward_term = np.empty((num_envs,), dtype=dtype) - env._weighted_reward = np.empty((num_envs,), dtype=dtype) - env._quat_error_w = np.empty((num_envs, num_bodies), dtype=dtype) - env._quat_error_x = np.empty((num_envs, num_bodies), dtype=dtype) - env._ee_pos_error_z = np.empty((num_envs, ee_indices.size), dtype=dtype) - env._undesired_contact_mask = np.empty((num_envs, undesired_indices.size), dtype=bool) - env._enable_reward_log = False - env._init_reward_functions() - env._active_reward_fns = { - name: fn for name, fn in env._reward_fns.items() if env._reward_term_is_active(name) - } - - info = { - "current_actions": current_actions, - "last_actions": last_actions, - "steps": np.zeros((num_envs,), dtype=np.uint32), - } - actual = env._compute_reward( - info, - motion_data, - robot_body_pos_w, - robot_body_quat_w, - robot_body_lin_vel_w, - robot_body_ang_vel_w, - dof_pos, - dof_vel, - ).copy() - - cfg = reward_config - expected = np.zeros((num_envs,), dtype=dtype) - root_pos_error = np.sum( - np.square(motion_data.body_pos_w[:, anchor_idx] - robot_body_pos_w[:, anchor_idx]), - axis=-1, - ) - expected += cfg.scales["motion_global_root_pos"] * np.exp(-root_pos_error / cfg.std_root_pos**2) - root_ori_error = ( - np_quat_error_magnitude( - motion_data.body_quat_w[:, anchor_idx], - robot_body_quat_w[:, anchor_idx], - ) - ** 2 - ) - expected += cfg.scales["motion_global_root_ori"] * np.exp(-root_ori_error / cfg.std_root_ori**2) - body_pos_error = np.sum(np.square(body_pos_relative_w - robot_body_pos_w), axis=-1) - expected += cfg.scales["motion_body_pos"] * np.exp( - -body_pos_error.mean(-1) / cfg.std_body_pos**2 - ) - body_ori_error = np_quat_error_magnitude( - body_quat_relative_w.reshape(num_envs * num_bodies, 4), - robot_body_quat_w.reshape(num_envs * num_bodies, 4), - ).reshape(num_envs, num_bodies) - expected += cfg.scales["motion_body_ori"] * np.exp( - -np.square(body_ori_error).mean(-1) / cfg.std_body_ori**2 - ) - body_lin_error = np.sum(np.square(motion_data.body_lin_vel_w - robot_body_lin_vel_w), axis=-1) - expected += cfg.scales["motion_body_lin_vel"] * np.exp( - -body_lin_error.mean(-1) / cfg.std_body_lin_vel**2 - ) - body_ang_error = np.sum(np.square(motion_data.body_ang_vel_w - robot_body_ang_vel_w), axis=-1) - expected += cfg.scales["motion_body_ang_vel"] * np.exp( - -body_ang_error.mean(-1) / cfg.std_body_ang_vel**2 - ) - ee_error = np.square(body_pos_relative_w[:, ee_indices, 2] - robot_body_pos_w[:, ee_indices, 2]) - expected += cfg.scales["motion_ee_body_pos_z"] * np.exp( - -ee_error.mean(-1) / cfg.std_body_pos**2 - ) - joint_pos_error = np.mean(np.square(motion_data.joint_pos - dof_pos), axis=-1) - expected += cfg.scales["motion_joint_pos"] * np.exp(-joint_pos_error / cfg.std_joint_pos**2) - joint_vel_error = np.mean(np.square(motion_data.joint_vel - dof_vel), axis=-1) - expected += cfg.scales["motion_joint_vel"] * np.exp(-joint_vel_error / cfg.std_joint_vel**2) - expected += cfg.scales["action_rate_l2"] * np.sum( - np.square(current_actions - last_actions), axis=1 - ) - lower_violation = np.maximum(0, joint_lower - dof_pos) - upper_violation = np.maximum(0, dof_pos - joint_upper) - expected += cfg.scales["joint_limit"] * np.sum( - np.square(lower_violation + upper_violation), axis=1 - ) - expected += cfg.scales["undesired_contacts"] * np.sum( - robot_body_pos_w[:, undesired_indices, 2] < contact_threshold, - axis=-1, - ) - expected *= ctrl_dt - - np.testing.assert_allclose(actual, expected, rtol=1e-12, atol=1e-12) - - -def test_g1_motion_tracking_deploy_actor_matches_unitree_mimic_terms(): - from unilab.tasks.motion_tracking.g1.tracking import G1MotionTrackingDeployEnv - - env, obs, _motion_data, _linvel, gyro, dof_pos, dof_vel, info = ( - _compute_g1_motion_tracking_obs_stub(G1MotionTrackingDeployEnv) - ) - - assert env.obs_groups_spec == {"obs": 19, "critic": 43} - assert obs["obs"].shape == (1, 19) - np.testing.assert_allclose(obs["obs"][:, 10:13], gyro + 100.0) - np.testing.assert_allclose(obs["obs"][:, 13:15], dof_pos - env.default_angles + 100.0) - np.testing.assert_allclose(obs["obs"][:, 15:17], dof_vel + 100.0) - np.testing.assert_allclose(obs["obs"][:, 17:19], info["current_actions"]) - - -def test_g1_box_tracking_cfg_uses_largebox_scene_and_motion_defaults(): - from unilab.tasks.motion_tracking.g1.box_tracking import BoxRewardConfig, G1BoxTrackingCfg - - cfg = G1BoxTrackingCfg() - - assert cfg.scene.model_file.endswith("scene_flat_with_largebox.xml") - assert str(cfg.motion_file).endswith("sub3_largebox_003_boxconverted.npz") - assert cfg.object_body_name == "largebox" - assert cfg.object_pos_threshold == pytest.approx(0.25) - assert cfg.object_ori_threshold == pytest.approx(0.8) - assert isinstance(cfg.reward_config, BoxRewardConfig) - assert cfg.reward_config.scales["object_global_ref_position_error_exp"] == pytest.approx(1.0) - assert cfg.reward_config.scales["object_global_ref_orientation_error_exp"] == pytest.approx(1.0) - - -def test_g1_box_tracking_is_exported_from_g1_and_motion_tracking_packages(): - from unilab.tasks.motion_tracking import ( - G1BoxTrackingCfg as TopLevelCfg, - ) - from unilab.tasks.motion_tracking import ( - G1BoxTrackingEnv as TopLevelEnv, - ) - from unilab.tasks.motion_tracking import ( - G1BoxTrackingEnvCfg as TopLevelEnvCfg, - ) - from unilab.tasks.motion_tracking.g1 import ( - G1BoxTrackingCfg as G1PkgCfg, - ) - from unilab.tasks.motion_tracking.g1 import ( - G1BoxTrackingEnv as G1PkgEnv, - ) - from unilab.tasks.motion_tracking.g1 import ( - G1BoxTrackingEnvCfg as G1PkgEnvCfg, - ) - - assert TopLevelCfg is G1PkgCfg - assert TopLevelEnv is G1PkgEnv - assert TopLevelEnvCfg is G1PkgEnvCfg - - -def _compute_g1_box_tracking_obs_stub(): - from unilab.tasks.motion_tracking.g1.box_tracking import G1BoxTrackingEnv - from unilab.tasks.motion_tracking.g1.motion_box_loader import BoxMotionData - - env = cast(Any, object.__new__(G1BoxTrackingEnv)) - env._num_envs = 1 - env._num_action = 2 - env._n_motion_bodies = 2 - env._critic_obs_width = env._critic_base_obs_dim(env._num_action) + env._n_motion_bodies * 9 - env._cfg = SimpleNamespace( - noise_config=SimpleNamespace( - level=1.0, - scale_linvel=1.0, - scale_gyro=1.0, - scale_joint_angle=1.0, - scale_joint_vel=1.0, - ), - body_names=("pelvis", "torso_link"), - ) - env.default_angles = np.array([[0.5, -0.5]], dtype=np.float32) - env.anchor_body_idx = 0 - env._object_body_ids = np.array([7], dtype=np.int32) - env._motion_anchor_pos_b = np.empty((1, 3), dtype=np.float32) - env._motion_anchor_ori_b = np.empty((1, 6), dtype=np.float32) - env._motion_command = np.empty((1, 4), dtype=np.float32) - env._joint_pos_rel = np.empty((1, 2), dtype=np.float32) - env._body_vec_error = np.empty((1, 2, 3), dtype=np.float32) - env._zero_actions = np.zeros((1, 2), dtype=np.float32) - env._obs_noise = lambda data, scale: np.asarray(data + 100.0, dtype=np.float32) - - class FakeBackend: - def get_body_pos_w(self, body_ids: np.ndarray) -> np.ndarray: - np.testing.assert_array_equal(body_ids, np.array([7], dtype=np.int32)) - return np.array([[[1.0, 2.0, 3.0]]], dtype=np.float32) - - def get_body_quat_w(self, body_ids: np.ndarray) -> np.ndarray: - np.testing.assert_array_equal(body_ids, np.array([7], dtype=np.int32)) - return np.array([[[1.0, 0.0, 0.0, 0.0]]], dtype=np.float32) - - def get_body_lin_vel_w(self, body_ids: np.ndarray) -> np.ndarray: - np.testing.assert_array_equal(body_ids, np.array([7], dtype=np.int32)) - return np.array([[[4.0, 5.0, 6.0]]], dtype=np.float32) - - env._backend = FakeBackend() - - motion_data = BoxMotionData( - joint_pos=np.array([[0.1, 0.2]], dtype=np.float32), - joint_vel=np.array([[0.3, 0.4]], dtype=np.float32), - body_pos_w=np.zeros((1, 2, 3), dtype=np.float32), - body_quat_w=np.tile(np.array([[[1.0, 0.0, 0.0, 0.0]]], dtype=np.float32), (1, 2, 1)), - body_lin_vel_w=np.zeros((1, 2, 3), dtype=np.float32), - body_ang_vel_w=np.zeros((1, 2, 3), dtype=np.float32), - object_pos_w=np.array([[1.0, 2.0, 3.0]], dtype=np.float32), - object_quat_w=np.array([[1.0, 0.0, 0.0, 0.0]], dtype=np.float32), - object_lin_vel_w=np.array([[4.0, 5.0, 6.0]], dtype=np.float32), - object_ang_vel_w=np.array([[7.0, 8.0, 9.0]], dtype=np.float32), - ) - linvel = np.array([[1.0, 2.0, 3.0]], dtype=np.float32) - gyro = np.array([[4.0, 5.0, 6.0]], dtype=np.float32) - dof_pos = np.array([[0.7, -0.2]], dtype=np.float32) - dof_vel = np.array([[7.0, 8.0]], dtype=np.float32) - robot_body_pos_w = np.array([[[0.0, 0.0, 0.0], [0.2, 0.0, 0.1]]], dtype=np.float32) - robot_body_quat_w = np.tile(np.array([[[1.0, 0.0, 0.0, 0.0]]], dtype=np.float32), (1, 2, 1)) - info = {"current_actions": np.array([[0.1, -0.2]], dtype=np.float32)} - - obs = env._compute_obs( - info, - motion_data, - linvel, - gyro, - dof_pos, - dof_vel, - robot_body_pos_w, - robot_body_quat_w, - ) - return env, obs, gyro, dof_pos, dof_vel, info - - -def test_g1_box_tracking_actor_matches_deploy_and_critic_adds_object_state(): - env, obs, gyro, dof_pos, dof_vel, info = _compute_g1_box_tracking_obs_stub() - - assert env.obs_groups_spec == {"obs": 19, "critic": 55} - assert obs["obs"].shape == (1, 19) - np.testing.assert_allclose(obs["obs"][:, 10:13], gyro + 100.0) - np.testing.assert_allclose(obs["obs"][:, 13:15], dof_pos - env.default_angles + 100.0) - np.testing.assert_allclose(obs["obs"][:, 15:17], dof_vel + 100.0) - np.testing.assert_allclose(obs["obs"][:, 17:19], info["current_actions"]) - np.testing.assert_allclose( - obs["critic"][:, -12:], - np.array([[1.0, 2.0, 3.0, 1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 4.0, 5.0, 6.0]], dtype=np.float32), - ) - - -def test_g1_box_tracking_critic_object_state_respects_subset_env_order(): - from unilab.tasks.motion_tracking.g1.box_tracking import G1BoxTrackingEnv - from unilab.tasks.motion_tracking.g1.motion_box_loader import BoxMotionData - - env = cast(Any, object.__new__(G1BoxTrackingEnv)) - env._num_envs = 4 - env._num_action = 2 - env._n_motion_bodies = 2 - env._critic_obs_width = env._critic_base_obs_dim(env._num_action) + env._n_motion_bodies * 9 - env.anchor_body_idx = 0 - env._object_body_ids = np.array([7], dtype=np.int32) - env._cfg = SimpleNamespace( - noise_config=SimpleNamespace( - level=0.0, - scale_linvel=0.0, - scale_gyro=0.0, - scale_joint_angle=0.0, - scale_joint_vel=0.0, - ), - body_names=("pelvis", "torso_link"), - ) - env.default_angles = np.zeros((2,), dtype=np.float32) - env._body_vec_error = np.empty((4, 2, 3), dtype=np.float32) - env._obs_noise = lambda data, scale: np.asarray(data, dtype=np.float32) - - class FakeBackend: - def get_body_pos_w(self, body_ids: np.ndarray) -> np.ndarray: - np.testing.assert_array_equal(body_ids, np.array([7], dtype=np.int32)) - return np.array( - [ - [[10.0, 0.0, 0.0]], - [[20.0, 0.0, 0.0]], - [[30.0, 0.0, 0.0]], - [[40.0, 0.0, 0.0]], - ], - dtype=np.float32, - ) - - def get_body_quat_w(self, body_ids: np.ndarray) -> np.ndarray: - np.testing.assert_array_equal(body_ids, np.array([7], dtype=np.int32)) - return np.tile(np.array([[[1.0, 0.0, 0.0, 0.0]]], dtype=np.float32), (4, 1, 1)) - - def get_body_lin_vel_w(self, body_ids: np.ndarray) -> np.ndarray: - np.testing.assert_array_equal(body_ids, np.array([7], dtype=np.int32)) - return np.array( - [ - [[1.0, 0.0, 0.0]], - [[2.0, 0.0, 0.0]], - [[3.0, 0.0, 0.0]], - [[4.0, 0.0, 0.0]], - ], - dtype=np.float32, - ) - - env._backend = FakeBackend() - - motion_data = BoxMotionData( - joint_pos=np.zeros((2, 2), dtype=np.float32), - joint_vel=np.zeros((2, 2), dtype=np.float32), - body_pos_w=np.zeros((2, 2, 3), dtype=np.float32), - body_quat_w=np.tile(np.array([[[1.0, 0.0, 0.0, 0.0]]], dtype=np.float32), (2, 2, 1)), - body_lin_vel_w=np.zeros((2, 2, 3), dtype=np.float32), - body_ang_vel_w=np.zeros((2, 2, 3), dtype=np.float32), - object_pos_w=np.zeros((2, 3), dtype=np.float32), - object_quat_w=np.tile(np.array([[1.0, 0.0, 0.0, 0.0]], dtype=np.float32), (2, 1)), - object_lin_vel_w=np.zeros((2, 3), dtype=np.float32), - object_ang_vel_w=np.zeros((2, 3), dtype=np.float32), - ) - - linvel = np.zeros((2, 3), dtype=np.float32) - gyro = np.zeros((2, 3), dtype=np.float32) - dof_pos = np.zeros((2, 2), dtype=np.float32) - dof_vel = np.zeros((2, 2), dtype=np.float32) - robot_body_pos_w = np.zeros((2, 2, 3), dtype=np.float32) - robot_body_quat_w = np.tile(np.array([[[1.0, 0.0, 0.0, 0.0]]], dtype=np.float32), (2, 2, 1)) - - obs = env._compute_obs( - { - "env_ids": np.array([2, 0], dtype=np.int32), - "current_actions": np.zeros((2, 2), dtype=np.float32), - }, - motion_data, - linvel, - gyro, - dof_pos, - dof_vel, - robot_body_pos_w, - robot_body_quat_w, - ) - - np.testing.assert_allclose( - obs["critic"][:, -12:], - np.array( - [ - [30.0, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 3.0, 0.0, 0.0], - [10.0, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 0.0, 0.0], - ], - dtype=np.float32, - ), - ) - - -def test_g1_motion_tracking_can_terminate_on_undesired_contacts(): - from unilab.tasks.motion_tracking.g1.tracking import G1MotionTrackingEnv - - env = cast(Any, object.__new__(G1MotionTrackingEnv)) - env._num_envs = 2 - env.anchor_body_idx = 0 - env.ee_body_indices = np.array([1], dtype=np.int32) - env._has_ee_body_indices = True - env.undesired_contact_body_indices = np.array([2], dtype=np.int32) - env._has_undesired_contact_body_indices = True - env._terminated = np.empty((2,), dtype=bool) - env._env_bool = np.empty((2,), dtype=bool) - env._env_error = np.empty((2,), dtype=np.float32) - env._ee_pos_error_z = np.empty((2, 1), dtype=np.float32) - env._ee_terminated = np.empty((2, 1), dtype=bool) - env._undesired_contact_mask = np.empty((2, 1), dtype=bool) - env.body_pos_relative_w = np.array( - [ - [[0.0, 0.0, 1.0], [0.0, 0.0, 0.8], [0.0, 0.0, 0.8]], - [[0.0, 0.0, 1.0], [0.0, 0.0, 0.8], [0.0, 0.0, 0.8]], - ], - dtype=np.float32, - ) - env._cfg = SimpleNamespace( - anchor_pos_z_threshold=0.5, - anchor_ori_threshold=1e9, - ee_body_pos_z_threshold=0.5, - terminate_on_undesired_contacts=True, - undesired_contact_z_threshold=0.05, - ) - quat = np.tile(np.array([[[1.0, 0.0, 0.0, 0.0]]], dtype=np.float32), (2, 3, 1)) - motion_data = SimpleNamespace(body_pos_w=env.body_pos_relative_w.copy(), body_quat_w=quat) - robot_body_pos_w = env.body_pos_relative_w.copy() - robot_body_pos_w[0, 2, 2] = 0.04 - robot_body_pos_w[1, 2, 2] = 0.10 - - terminated = env._compute_terminations(motion_data, robot_body_pos_w, quat) - np.testing.assert_array_equal(terminated, np.array([True, False])) - - env._cfg.terminate_on_undesired_contacts = False - terminated_without_contact = env._compute_terminations(motion_data, robot_body_pos_w, quat) - np.testing.assert_array_equal(terminated_without_contact, np.array([False, False])) - - -def test_g1_motion_tracking_cfg_has_domain_rand_for_motrix(): - from unilab.tasks.motion_tracking.g1.tracking import G1MotionTrackingCfg - - cfg = G1MotionTrackingCfg() - assert hasattr(cfg, "domain_rand") - assert cfg.domain_rand.randomize_base_mass is False - assert cfg.domain_rand.random_com is False - assert cfg.domain_rand.randomize_gravity is False - assert cfg.domain_rand.push_robots is False - - -def test_g1_motion_tracking_cfg_preserves_legacy_defaults(): - from unilab.tasks.motion_tracking.g1.tracking import G1MotionTrackingCfg - - cfg = G1MotionTrackingCfg() - - assert str(cfg.motion_file).endswith("dance1_subject2_part.npz") - assert cfg.pose_randomization.x == (-0.05, 0.05) - assert cfg.velocity_randomization.x == (-0.5, 0.5) - assert cfg.joint_position_range == (-0.1, 0.1) - assert cfg.anchor_ori_threshold == pytest.approx(0.8) - assert cfg.sampling_mode == "adaptive" - assert cfg.truncate_on_clip_end is False - - -def test_g1_motion_tracking_init_delegates_motion_body_ids_to_backend(monkeypatch): - from unilab.tasks.motion_tracking.common import tracking as tracking_module - from unilab.tasks.motion_tracking.common.g1_base import G1BaseEnv - from unilab.tasks.motion_tracking.g1.tracking import ( - G1MotionTrackingCfg, - G1MotionTrackingEnv, - ) - - calls: dict[str, Any] = {} - - class FakeBackend: - def get_body_ids(self, names: tuple[str, ...]) -> np.ndarray: - calls["body_ids_names"] = names - return np.array([10, 11], dtype=np.int32) - - def get_motion_body_ids(self, names: tuple[str, ...]) -> np.ndarray: - calls["motion_body_ids_names"] = names - return np.array([1, 2], dtype=np.int32) - - def copy_body_state_w( - self, - body_ids: np.ndarray, - out_pos: np.ndarray, - out_quat: np.ndarray, - out_lin_vel: np.ndarray, - out_ang_vel: np.ndarray, - ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: - return out_pos, out_quat, out_lin_vel, out_ang_vel - - def get_joint_range(self) -> None: - return None - - def fake_base_init(self, cfg, backend, num_envs): - self._cfg = cfg - self._backend = backend - self._num_envs = num_envs - self._num_action = 2 - self._init_qpos = np.zeros((9,), dtype=np.float32) - self._init_qvel = np.zeros((8,), dtype=np.float32) - - class FakeMotionLoader: - def __init__(self, motion_file: str, body_indices: np.ndarray): - calls["motion_loader"] = (motion_file, body_indices.copy()) - - class FakeMotionSampler: - def __init__( - self, - motion_loader: Any, - mode: str, - num_envs: int, - start_ratio: float = 0.0, - ): - calls["motion_sampler"] = (motion_loader, mode, num_envs, start_ratio) - - fake_backend = FakeBackend() - monkeypatch.setattr(tracking_module, "create_backend", lambda *args, **kwargs: fake_backend) - monkeypatch.setattr(G1BaseEnv, "__init__", fake_base_init) - monkeypatch.setattr(tracking_module, "MotionLoader", FakeMotionLoader) - monkeypatch.setattr(tracking_module, "MotionSampler", FakeMotionSampler) - monkeypatch.setattr( - G1MotionTrackingEnv, - "_init_domain_randomization", - lambda self, provider: calls.setdefault("dr_provider", provider.__class__.__name__), - ) - monkeypatch.setattr( - G1MotionTrackingEnv, - "_init_reward_functions", - lambda self: ( - calls.setdefault("reward_init", True), - setattr(self, "_reward_fns", {}), - ), - ) - - cfg = G1MotionTrackingCfg( - motion_file="dummy_motion.npz", - body_names=("pelvis", "torso_link"), - ee_body_names=("torso_link",), - ) - env = cast(Any, G1MotionTrackingEnv)(cfg, num_envs=4, backend_type="motrix") - - np.testing.assert_array_equal(env.body_ids, np.array([10, 11], dtype=np.int32)) - assert calls["body_ids_names"] == cfg.body_names - assert calls["motion_body_ids_names"] == cfg.body_names - assert calls["motion_loader"][0] == "dummy_motion.npz" - np.testing.assert_array_equal(calls["motion_loader"][1], np.array([1, 2], dtype=np.int32)) - assert calls["motion_sampler"][1:] == ("adaptive", 4, cfg.sampling_start_ratio) - assert calls["dr_provider"] == "MotionTrackingDomainRandomizationProvider" - assert calls["reward_init"] is True - - def _patch_sharpa_rotation_constructor( monkeypatch: pytest.MonkeyPatch, initialized_providers: list[Any], @@ -1667,333 +775,6 @@ def test_sharpa_run_complete_reports_effective_collection_target() -> None: } -def test_g1_flip_tracking_cfg_uses_flip_profile(): - from unilab.tasks.motion_tracking.g1.flip_tracking import G1FlipTrackingCfg - - cfg = G1FlipTrackingCfg() - - assert cfg.scene.model_file.endswith("scene_flat.xml") - assert str(cfg.motion_file).endswith("flip_360_001__A304.npz") - assert cfg.pose_randomization.x == (0.0, 0.0) - assert cfg.velocity_randomization.x == (0.0, 0.0) - assert cfg.joint_position_range == (0.0, 0.0) - assert cfg.truncate_on_clip_end is False - assert cfg.anchor_ori_threshold == pytest.approx(1e9) - assert cfg.terminate_on_undesired_contacts is True - assert cfg.sampling_mode == "start" - - -def test_g1_wall_flip_tracking_cfg_uses_wall_flip_profile(): - from unilab.tasks.motion_tracking.g1.flip_tracking import G1WallFlipTrackingCfg - - cfg = G1WallFlipTrackingCfg() - - assert cfg.scene.model_file.endswith("scene_flat_with_wall.xml") - assert str(cfg.motion_file).endswith("flip_from_wall_104__A304.npz") - assert cfg.pose_randomization.x == (0.0, 0.0) - assert cfg.velocity_randomization.x == (0.0, 0.0) - assert cfg.joint_position_range == (0.0, 0.0) - assert cfg.truncate_on_clip_end is False - assert cfg.anchor_ori_threshold == pytest.approx(1e9) - assert cfg.anchor_pos_z_threshold == pytest.approx(0.5) - assert cfg.ee_body_pos_z_threshold == pytest.approx(0.5) - assert cfg.terminate_on_undesired_contacts is True - assert cfg.sampling_mode == "adaptive" - - -def test_x2_wall_flip_tracking_cfg_uses_x2_wall_flip_profile(): - from unilab.tasks.motion_tracking.x2 import X2WallFlipTrackingCfg - - cfg = X2WallFlipTrackingCfg() - - assert cfg.scene.model_file.endswith("scene_flat_with_wall.xml") - assert str(cfg.motion_file).endswith("tictacflip_6-3_g1format.npz") - assert cfg.sensor.local_linvel == "body-linear-vel" - assert cfg.sensor.gyro == "body-angular-velocity" - assert cfg.anchor_body_name == "torso_link" - assert cfg.body_names[0] == "pelvis" - assert cfg.body_names[-1] == "right_wrist_roll_link" - assert len(cfg.body_names) == 30 - assert cfg.pose_randomization.x == (0.0, 0.0) - assert cfg.velocity_randomization.x == (0.0, 0.0) - assert cfg.joint_position_range == (0.0, 0.0) - assert cfg.truncate_on_clip_end is False - assert cfg.anchor_ori_threshold == pytest.approx(1e9) - assert cfg.anchor_pos_z_threshold == pytest.approx(0.5) - assert cfg.ee_body_pos_z_threshold == pytest.approx(0.5) - assert cfg.terminate_on_undesired_contacts is True - assert cfg.sampling_mode == "adaptive" - - -def test_g1_motion_tracking_apply_action_accepts_per_joint_action_scale(): - from unilab.tasks.motion_tracking.g1.tracking import G1MotionTrackingEnv - - env = cast(Any, object.__new__(G1MotionTrackingEnv)) - env.default_angles = np.array([0.5, -0.5, 1.0], dtype=np.float32) - env._cfg = SimpleNamespace( - control_config=SimpleNamespace( - action_scale=[0.1, 0.2, 0.3], - simulate_action_latency=False, - ) - ) - state = SimpleNamespace(info={}) - actions = np.array([[1.0, -1.0, 0.5]], dtype=np.float32) - - ctrl = env.apply_action(actions, state) - - np.testing.assert_allclose(ctrl, np.array([[0.6, -0.7, 1.15]], dtype=np.float32)) - np.testing.assert_array_equal(state.info["current_actions"], actions) - - -def _make_g1_motion_tracking_clip_end_stub( - *, - truncate_on_clip_end: bool, - terminated: np.ndarray | None = None, - step_env_ids: np.ndarray | None = None, -): - from unilab.base.np_env import NpEnvState - from unilab.tasks.motion_tracking.common.motion_loader import MotionData - from unilab.tasks.motion_tracking.g1.tracking import G1MotionTrackingEnv - - class FakeBackend: - def __init__(self) -> None: - self.set_state_calls: list[tuple[np.ndarray, np.ndarray, np.ndarray]] = [] - - def get_body_lin_vel_w(self, body_ids: np.ndarray) -> np.ndarray: - return np.zeros((2, len(body_ids), 3), dtype=np.float32) - - def get_body_ang_vel_w(self, body_ids: np.ndarray) -> np.ndarray: - return np.zeros((2, len(body_ids), 3), dtype=np.float32) - - def get_body_pos_w(self, body_ids: np.ndarray) -> np.ndarray: - return np.zeros((2, len(body_ids), 3), dtype=np.float32) - - def get_body_quat_w(self, body_ids: np.ndarray) -> np.ndarray: - return np.tile( - np.array([[[1.0, 0.0, 0.0, 0.0]]], dtype=np.float32), - (2, len(body_ids), 1), - ) - - def get_body_pose_w(self, body_ids: np.ndarray) -> tuple[np.ndarray, np.ndarray]: - return self.get_body_pos_w(body_ids), self.get_body_quat_w(body_ids) - - def get_body_pose_w_rows( - self, env_ids: np.ndarray, body_ids: np.ndarray - ) -> tuple[np.ndarray, np.ndarray]: - rows = np.asarray(env_ids, dtype=np.intp) - return self.get_body_pos_w(body_ids)[rows], self.get_body_quat_w(body_ids)[rows] - - def get_sensor_data_rows(self, name: str, env_ids: np.ndarray) -> np.ndarray: - del name - return np.zeros((len(env_ids), 3), dtype=np.float32) - - def get_body_vel_w(self, body_ids: np.ndarray) -> tuple[np.ndarray, np.ndarray]: - return self.get_body_lin_vel_w(body_ids), self.get_body_ang_vel_w(body_ids) - - def set_state(self, env_ids: np.ndarray, qpos: np.ndarray, qvel: np.ndarray) -> None: - self.set_state_calls.append((env_ids.copy(), qpos.copy(), qvel.copy())) - - class FakeMotionLoader: - def get_motion_at_frame(self, frames: np.ndarray) -> MotionData: - frame_values = frames.astype(np.float32)[:, None] - return MotionData( - joint_pos=np.repeat(frame_values, 2, axis=1), - joint_vel=np.repeat(frame_values + 10.0, 2, axis=1), - body_pos_w=np.pad(frame_values[:, None, :], ((0, 0), (0, 0), (0, 2))), - body_quat_w=np.tile( - np.array([[[1.0, 0.0, 0.0, 0.0]]], dtype=np.float32), (len(frames), 1, 1) - ), - body_lin_vel_w=np.zeros((len(frames), 1, 3), dtype=np.float32), - body_ang_vel_w=np.zeros((len(frames), 1, 3), dtype=np.float32), - ) - - class FakeSampler: - def __init__(self) -> None: - self.failure_updates: list[np.ndarray] = [] - self.sampled_env_ids: list[np.ndarray] = [] - self.current_frames = np.zeros((2,), dtype=np.int32) - self._after_step = False - - def get_current_motion(self) -> MotionData: - if self._after_step and np.any(self.current_frames == 99): - raise AssertionError("queried all current motion while a clip-end frame is invalid") - return MotionData( - joint_pos=np.zeros((2, 2), dtype=np.float32), - joint_vel=np.zeros((2, 2), dtype=np.float32), - body_pos_w=np.zeros((2, 1, 3), dtype=np.float32), - body_quat_w=np.tile( - np.array([[[1.0, 0.0, 0.0, 0.0]]], dtype=np.float32), (2, 1, 1) - ), - body_lin_vel_w=np.zeros((2, 1, 3), dtype=np.float32), - body_ang_vel_w=np.zeros((2, 1, 3), dtype=np.float32), - ) - - def update_failure_stats(self, terminated: np.ndarray) -> None: - self.failure_updates.append(terminated.copy()) - - def step(self) -> np.ndarray: - env_ids = np.array([1], dtype=np.int32) if step_env_ids is None else step_env_ids.copy() - self.current_frames[env_ids] = 99 - self._after_step = True - return env_ids - - def sample_frames(self, env_ids: np.ndarray) -> np.ndarray: - self.sampled_env_ids.append(env_ids.copy()) - self.current_frames[env_ids] = 7 - return np.full(len(env_ids), 7, dtype=np.int32) - - env = cast(Any, object.__new__(G1MotionTrackingEnv)) - env._num_envs = 2 - zero_pose = SimpleNamespace( - x=(0.0, 0.0), - y=(0.0, 0.0), - z=(0.0, 0.0), - roll=(0.0, 0.0), - pitch=(0.0, 0.0), - yaw=(0.0, 0.0), - ) - env._cfg = SimpleNamespace( - max_episode_steps=None, - truncate_on_clip_end=truncate_on_clip_end, - sensor=SimpleNamespace(local_linvel="local_linvel", gyro="gyro"), - pose_randomization=zero_pose, - velocity_randomization=zero_pose, - joint_position_range=(0.0, 0.0), - ) - env.body_ids = np.array([0], dtype=np.int32) - env._backend = FakeBackend() - env.motion_sampler = FakeSampler() - env.motion_loader = FakeMotionLoader() - env._motion_data_buffer = None - env._copy_body_state_w = None - env._clip_end_truncated = np.zeros((2,), dtype=bool) - env._env_bool = np.empty((2,), dtype=bool) - env._init_qpos = np.zeros((9,), dtype=np.float32) - env._init_qvel = np.zeros((8,), dtype=np.float32) - env.get_local_linvel = lambda: np.zeros((2, 3), dtype=np.float32) - env.get_gyro = lambda: np.zeros((2, 3), dtype=np.float32) - env.get_dof_pos = lambda: np.zeros((2, 2), dtype=np.float32) - env.get_dof_vel = lambda: np.zeros((2, 2), dtype=np.float32) - env._get_joint_range = lambda: None - env._get_body_pose_w = lambda: ( - np.zeros((2, 1, 3), dtype=np.float32), - np.tile(np.array([[[1.0, 0.0, 0.0, 0.0]]], dtype=np.float32), (2, 1, 1)), - ) - env._update_relative_transforms = lambda *args: None - env._compute_terminations = lambda *args: ( - np.zeros((2,), dtype=bool) if terminated is None else terminated.copy() - ) - env._compute_reward = lambda *args: np.zeros((2,), dtype=np.float32) - env._compute_obs = lambda *args: { - "obs": np.zeros((2, 1), dtype=np.float32), - "critic": np.zeros((2, 2), dtype=np.float32), - } - - state = NpEnvState( - obs={ - "obs": np.zeros((2, 1), dtype=np.float32), - "critic": np.zeros((2, 2), dtype=np.float32), - }, - reward=np.zeros((2,), dtype=np.float32), - terminated=np.zeros((2,), dtype=bool), - truncated=np.zeros((2,), dtype=bool), - info={"steps": np.zeros((2,), dtype=np.uint32)}, - ) - - return env, state - - -def test_g1_motion_tracking_clip_end_resamples_by_default_without_truncation(): - env, state = _make_g1_motion_tracking_clip_end_stub(truncate_on_clip_end=False) - - next_state = env.update_state(state) - truncated = env._compute_truncated(next_state) - - np.testing.assert_array_equal(next_state.terminated, np.array([False, False])) - np.testing.assert_array_equal(truncated, np.array([False, False])) - np.testing.assert_array_equal(env.motion_sampler.sampled_env_ids[0], np.array([1])) - assert len(env._backend.set_state_calls) == 1 - set_state_env_ids, qpos, qvel = env._backend.set_state_calls[0] - np.testing.assert_array_equal(set_state_env_ids, np.array([1], dtype=np.int32)) - np.testing.assert_array_equal(qpos[:, 0], np.array([7.0], dtype=np.float32)) - np.testing.assert_array_equal(qpos[:, 7:], np.array([[7.0, 7.0]], dtype=np.float32)) - np.testing.assert_array_equal(qvel[:, 6:], np.array([[17.0, 17.0]], dtype=np.float32)) - np.testing.assert_array_equal( - env.motion_sampler.failure_updates[0], np.array([False, False], dtype=bool) - ) - - -def test_g1_motion_tracking_clip_end_truncates_when_config_enabled(): - env, state = _make_g1_motion_tracking_clip_end_stub(truncate_on_clip_end=True) - - next_state = env.update_state(state) - truncated = env._compute_truncated(next_state) - - np.testing.assert_array_equal(next_state.terminated, np.array([False, False])) - np.testing.assert_array_equal(truncated, np.array([False, True])) - assert env.motion_sampler.sampled_env_ids == [] - assert env._backend.set_state_calls == [] - np.testing.assert_array_equal( - env.motion_sampler.failure_updates[0], np.array([False, False], dtype=bool) - ) - - -def test_g1_motion_tracking_clip_end_resample_skips_terminated_envs(): - env, state = _make_g1_motion_tracking_clip_end_stub( - truncate_on_clip_end=False, - terminated=np.array([False, True], dtype=bool), - ) - - next_state = env.update_state(state) - truncated = env._compute_truncated(next_state) - - np.testing.assert_array_equal(next_state.terminated, np.array([False, True])) - np.testing.assert_array_equal(truncated, np.array([False, False])) - assert env.motion_sampler.sampled_env_ids == [] - assert env._backend.set_state_calls == [] - - -def test_g1_motion_tracking_clip_end_resample_keeps_terminated_final_obs_valid(): - env, state = _make_g1_motion_tracking_clip_end_stub( - truncate_on_clip_end=False, - terminated=np.array([False, True], dtype=bool), - step_env_ids=np.array([0, 1], dtype=np.int32), - ) - - next_state = env.update_state(state) - truncated = env._compute_truncated(next_state) - - np.testing.assert_array_equal(next_state.terminated, np.array([False, True])) - np.testing.assert_array_equal(truncated, np.array([False, False])) - np.testing.assert_array_equal(env.motion_sampler.sampled_env_ids[0], np.array([0])) - assert len(env._backend.set_state_calls) == 1 - set_state_env_ids, qpos, qvel = env._backend.set_state_calls[0] - np.testing.assert_array_equal(set_state_env_ids, np.array([0], dtype=np.int32)) - np.testing.assert_array_equal(qpos[:, 0], np.array([7.0], dtype=np.float32)) - np.testing.assert_array_equal(qvel[:, 6:], np.array([[17.0, 17.0]], dtype=np.float32)) - - -def test_g1_motion_tracking_clip_end_does_not_override_true_termination(): - from unilab.base.np_env import NpEnvState - from unilab.tasks.motion_tracking.g1.tracking import G1MotionTrackingEnv - - env = cast(Any, object.__new__(G1MotionTrackingEnv)) - env._num_envs = 2 - env._cfg = type("Cfg", (), {"max_episode_steps": None})() - env._clip_end_truncated = np.array([False, True], dtype=bool) - - state = NpEnvState( - obs={}, - reward=np.zeros((2,), dtype=np.float32), - terminated=np.array([False, True], dtype=bool), - truncated=np.zeros((2,), dtype=bool), - info={"steps": np.zeros((2,), dtype=np.uint32)}, - ) - - truncated = env._compute_truncated(state) - np.testing.assert_array_equal(truncated, np.array([False, False])) - - # --------------------------------------------------------------------------- # Fast env/backend smoke tests # --------------------------------------------------------------------------- diff --git a/tests/envs/test_manager_based_rl_env.py b/tests/envs/test_manager_based_rl_env.py index 351d58326..8f1fb1ec9 100644 --- a/tests/envs/test_manager_based_rl_env.py +++ b/tests/envs/test_manager_based_rl_env.py @@ -418,6 +418,7 @@ def test_generic_factory_routes_only_public_backend_contract( root_body_name="base", actuator_names=("motor",), ) + cfg.scene.entities["payload"] = EntityCfg(root_body_name="box") backend = _FakeBackend(3) constructed: dict[str, Any] = {} @@ -470,17 +471,20 @@ def fake_make_env( @pytest.mark.parametrize( ("entities", "match"), [ - ({"robot": EntityCfg(actuator_names=("motor",))}, "found 0 root entities"), + ( + {"robot": EntityCfg(actuator_names=("motor",))}, + "at least one scene entity with an explicit root_body_name", + ), ( { - "robot": EntityCfg(root_body_name="base", actuator_names=("motor",)), "payload": EntityCfg(root_body_name="box"), + "tool": EntityCfg(root_body_name="tool"), }, - "found 2 root entities.*robot.*payload", + "conventional 'robot' root entity.*payload.*tool", ), ], ) -def test_generic_factory_requires_exactly_one_explicit_root_entity( +def test_generic_factory_rejects_missing_or_ambiguous_root_entity( monkeypatch: pytest.MonkeyPatch, entities: dict[str, EntityCfg], match: str, diff --git a/tests/envs/test_motion_profiles.py b/tests/envs/test_motion_profiles.py new file mode 100644 index 000000000..04bd383a3 --- /dev/null +++ b/tests/envs/test_motion_profiles.py @@ -0,0 +1,530 @@ +from __future__ import annotations + +from pathlib import Path +from types import SimpleNamespace +from typing import Any, cast + +import numpy as np +import pytest +from hydra import compose, initialize_config_dir +from hydra.core.global_hydra import GlobalHydra + +from unilab.base import registry +from unilab.base.config_materialization import apply_cfg_overrides +from unilab.envs import ManagerBasedRlEnvCfg +from unilab.envs.mdp.actions import JointPositionAction +from unilab.training.backend_adapter import BackendAdapter + +_ROOT = Path(__file__).parents[2] + +_PROFILE_IDENTITIES = ( + "G1BoxTracking", + "G1BoxTracking23Dof", + "G1ClimbTracking", + "G1ClimbTracking23Dof", + "G1FlipTracking", + "G1FlipTracking23Dof", + "G1FlipTrackingSAC", + "G1FlipTrackingSAC23Dof", + "G1WallFlipTracking", + "G1WallFlipTracking23Dof", + "G1WallFlipTrackingSAC", + "G1WallFlipTrackingSAC23Dof", + "G1WBTObs", + "G1WBTObs23Dof", + "X2WallFlipTracking", +) + +_PPO_PROFILES = ( + ( + "g1_box_tracking", + "G1BoxTracking", + "scene_flat_with_largebox.xml", + "sub3_largebox_003_boxconverted.npz", + ), + ( + "g1_23dof_box_tracking", + "G1BoxTracking23Dof", + "scene_flat_23dof_with_largebox.xml", + "sub3_largebox_003_boxconverted_23dof.npz", + ), + ( + "g1_climb_tracking", + "G1ClimbTracking", + "scene_climb_20_z_scale_1.xml", + "climb_20_z_scale_1.0.npz", + ), + ( + "g1_23dof_climb_tracking", + "G1ClimbTracking23Dof", + "scene_climb_20_z_scale_1_23dof.xml", + "climb_20_z_scale_1.0_23dof.npz", + ), + ("g1_flip_tracking", "G1FlipTracking", "scene_flat.xml", "flip_360_001__A304.npz"), + ( + "g1_23dof_flip_tracking", + "G1FlipTracking23Dof", + "scene_flat_23dof.xml", + "flip_360_001__A304_23dof.npz", + ), + ( + "g1_wall_flip_tracking", + "G1WallFlipTracking", + "scene_flat_with_wall.xml", + "flip_from_wall_104__A304.npz", + ), + ( + "g1_23dof_wall_flip_tracking", + "G1WallFlipTracking23Dof", + "scene_flat_23dof_with_wall.xml", + "flip_from_wall_104__A304_23dof.npz", + ), + ( + "x2_wall_flip_tracking", + "X2WallFlipTracking", + "scene_flat_with_wall.xml", + "tictacflip_6-3_g1format.npz", + ), +) + +_APPO_PROFILES = tuple( + profile + for profile in _PPO_PROFILES + if "box" not in profile[0] and not profile[0].startswith("x2") +) + +_SAC_PROFILES = ( + ("sac/g1_flip_tracking", "G1FlipTrackingSAC", "scene_flat.xml", "flip_360_001__A304.npz"), + ( + "sac/g1_23dof_flip_tracking", + "G1FlipTrackingSAC23Dof", + "scene_flat_23dof.xml", + "flip_360_001__A304_23dof.npz", + ), + ( + "sac/g1_wall_flip_tracking", + "G1WallFlipTrackingSAC", + "scene_flat_with_wall.xml", + "flip_from_wall_104__A304.npz", + ), + ( + "sac/g1_23dof_wall_flip_tracking", + "G1WallFlipTrackingSAC23Dof", + "scene_flat_23dof_with_wall.xml", + "flip_from_wall_104__A304_23dof.npz", + ), + ("sac/g1_wbt_obs", "G1WBTObs", "scene_flat.xml", "dance1_subject2_part.npz"), + ( + "sac/g1_23dof_wbt_obs", + "G1WBTObs23Dof", + "scene_flat_23dof.xml", + "dance1_subject2_part_23dof.npz", + ), +) + +_OWNER_CASES = ( + tuple( + ("ppo", task, backend, identity, scene, motion) + for task, identity, scene, motion in _PPO_PROFILES + for backend in ("mujoco", "motrix") + ) + + tuple( + ("appo", task, backend, identity, scene, motion) + for task, identity, scene, motion in _APPO_PROFILES + for backend in ("mujoco", "motrix") + ) + + tuple( + ("offpolicy", task, "mujoco", identity, scene, motion) + for task, identity, scene, motion in _SAC_PROFILES + ) +) + +_LEGACY_G1_ACTION_SCALE = ( + (r".*_(hip_pitch|hip_yaw)_joint", 0.5475464629911068), + (r".*_(hip_roll|knee)_joint", 0.35066146637882434), + (r".*_ankle_(pitch|roll)_joint", 0.43857731392336724), + (r"waist_yaw_joint", 0.5475464629911068), + (r"waist_(roll|pitch)_joint", 0.43857731392336724), + (r".*_(shoulder_(pitch|roll|yaw)|elbow|wrist_roll)_joint", 0.43857731392336724), + (r".*_wrist_(pitch|yaw)_joint", 0.07450087032950714), +) + +_LEGACY_SCALAR_ACTION_SCALE = { + ("ppo", "g1_flip_tracking", "motrix"): 0.25, + ("ppo", "g1_23dof_flip_tracking", "motrix"): 0.25, + ("appo", "g1_wall_flip_tracking", "motrix"): 0.25, + ("appo", "g1_23dof_wall_flip_tracking", "motrix"): 0.25, +} + + +def _compose_owner(config_root: str, task: str, backend: str) -> Any: + GlobalHydra.instance().clear() + with initialize_config_dir(config_dir=str(_ROOT / "conf" / config_root), version_base="1.3"): + overrides = [f"task={task}/{backend}"] + if config_root == "offpolicy": + overrides.insert(0, "algo=sac") + return compose("config", overrides=overrides) + + +def _materialize_profile( + config_root: str, + task: str, + backend: str, + identity: str, +) -> tuple[Any, ManagerBasedRlEnvCfg, dict[str, Any]]: + registry.ensure_registries() + owner = _compose_owner(config_root, task, backend) + cfg = registry.materialize_env_config(identity) + assert isinstance(cfg, ManagerBasedRlEnvCfg) + override = BackendAdapter( + owner, + root_dir=_ROOT, + algo_name="sac" if config_root == "offpolicy" else config_root, + ).build_task_env_cfg_override() + apply_cfg_overrides( + cfg, + override, + ) + return owner, cfg, override + + +def _resolved_action_scale(cfg: ManagerBasedRlEnvCfg) -> tuple[list[str], np.ndarray]: + joint_names = list(cfg.scene.entities["robot"].joint_names or ()) + entity = SimpleNamespace( + data=SimpleNamespace( + default_joint_pos=np.zeros((1, len(joint_names)), dtype=np.float32), + ), + find_joints_by_actuator_names=lambda actuator_names: ( + list(range(len(joint_names))), + joint_names, + ), + ) + env = SimpleNamespace(num_envs=1, scene={"robot": entity}) + action = JointPositionAction(cfg.actions["joint_pos"], cast(Any, env)) + scale = action.scale + if np.isscalar(scale): + scale = np.full(len(joint_names), scale, dtype=np.float32) + else: + scale = np.asarray(scale[0]) + return joint_names, scale + + +def _legacy_g1_action_scale(joint_names: list[str]) -> np.ndarray: + import re + + expected = [] + for joint_name in joint_names: + matches = [ + value for pattern, value in _LEGACY_G1_ACTION_SCALE if re.fullmatch(pattern, joint_name) + ] + assert len(matches) == 1, (joint_name, matches) + expected.append(matches[0]) + return np.asarray(expected, dtype=np.float32) + + +@pytest.mark.parametrize( + ("config_root", "task", "backend", "identity", "scene_file", "motion_file"), + _OWNER_CASES, +) +def test_motion_profile_owner_composes_to_manager_runtime( + config_root: str, + task: str, + backend: str, + identity: str, + scene_file: str, + motion_file: str, +) -> None: + registry.ensure_registries() + owner, cfg, _ = _materialize_profile(config_root, task, backend, identity) + + assert owner.training.task_name == identity + assert owner.training.sim_backend == backend + assert cfg.scene.model_file.endswith(scene_file) + assert str(cfg.commands["motion"].params.motion_file).endswith(motion_file) + assert list(cfg.actions) == ["joint_pos"] + assert cfg.policy_observation_group == "actor" + assert cfg.critic_observation_group == "critic" + + +@pytest.mark.parametrize( + ("config_root", "task", "backend", "identity", "scene_file", "motion_file"), + _OWNER_CASES, +) +def test_motion_profile_action_scale_matches_legacy_runtime( + config_root: str, + task: str, + backend: str, + identity: str, + scene_file: str, + motion_file: str, +) -> None: + del scene_file, motion_file + _, cfg, _ = _materialize_profile(config_root, task, backend, identity) + joint_names, actual = _resolved_action_scale(cfg) + + scalar = _LEGACY_SCALAR_ACTION_SCALE.get((config_root, task, backend)) + if "box_tracking" in task or task == "x2_wall_flip_tracking": + scalar = 0.25 + elif "wbt_obs" in task: + scalar = 2.0 + expected = ( + np.full(len(joint_names), scalar, dtype=np.float32) + if scalar is not None + else _legacy_g1_action_scale(joint_names) + ) + + np.testing.assert_allclose(actual, expected) + + +@pytest.mark.parametrize( + ( + "task", + "backend", + "adaptive_kl_factor", + "adaptive_lr_factor", + "enable_compile", + "steps_per_env", + "replay_queue_size", + ), + ( + ("g1_climb_tracking", "mujoco", 1.2, 1.1, False, 24, None), + ("g1_climb_tracking", "motrix", 1.2, 1.1, False, 24, None), + ("g1_23dof_climb_tracking", "mujoco", 1.2, 1.1, False, 24, None), + ("g1_23dof_climb_tracking", "motrix", 1.2, 1.1, False, 24, None), + ("g1_23dof_flip_tracking", "mujoco", 1.2, 1.1, False, 24, None), + ("g1_23dof_flip_tracking", "motrix", 2.0, 1.5, True, 24, None), + ("g1_wall_flip_tracking", "motrix", 2.0, 1.5, True, 24, None), + ("g1_23dof_wall_flip_tracking", "mujoco", 2.0, 1.5, True, 20, 5), + ("g1_23dof_wall_flip_tracking", "motrix", 2.0, 1.5, True, 24, None), + ), +) +def test_appo_profiles_preserve_training_owner_contract( + task: str, + backend: str, + adaptive_kl_factor: float, + adaptive_lr_factor: float, + enable_compile: bool, + steps_per_env: int, + replay_queue_size: int | None, +) -> None: + owner = _compose_owner("appo", task, backend) + + assert owner.algo.algorithm.adaptive_kl_factor == pytest.approx(adaptive_kl_factor) + assert owner.algo.algorithm.adaptive_lr_factor == pytest.approx(adaptive_lr_factor) + assert owner.algo.algorithm.enable_compile is enable_compile + assert owner.algo.steps_per_env == steps_per_env + assert owner.training.replay_queue_size == replay_queue_size + + +@pytest.mark.parametrize("task", ("g1_flip_tracking", "g1_23dof_flip_tracking")) +def test_ppo_motrix_flip_profiles_keep_actor_normalization_disabled(task: str) -> None: + owner = _compose_owner("ppo", task, "motrix") + + assert owner.algo.empirical_normalization is False + assert owner.algo.obs_groups.actor == ["actor"] + assert owner.algo.obs_groups.critic == ["critic"] + + +@pytest.mark.parametrize( + "task", + ( + "g1_23dof_box_tracking", + "g1_23dof_climb_tracking", + "g1_23dof_flip_tracking", + "g1_23dof_wall_flip_tracking", + "x2_wall_flip_tracking", + ), +) +def test_ppo_profiles_without_legacy_play_overrides_stay_disabled(task: str) -> None: + owner = _compose_owner("ppo", task, "mujoco") + + assert owner.play_profile.enabled is False + assert owner.play_profile.env is None + + +@pytest.mark.parametrize("task", ("g1_box_tracking", "g1_23dof_box_tracking")) +def test_box_motrix_drops_unconsumed_algorithm_noise_config(task: str) -> None: + owner = _compose_owner("ppo", task, "motrix") + + assert "noise_config" not in owner.algo + + +def test_all_motion_profiles_have_one_manager_factory_and_both_backends() -> None: + registry.ensure_registries() + metadata = registry.list_registered_envs() + + for identity in _PROFILE_IDENTITIES: + assert metadata[identity] == { + "config_factory": "ManagerBasedRlEnvCfg", + "available_backends": ["mujoco", "motrix"], + } + + +def test_box_wall_wbt_and_x2_profiles_keep_only_owner_differences() -> None: + from unilab.tasks.motion_tracking.g1.manager_terms import BoxMotionCommandCfg + + _, box, _ = _materialize_profile("ppo", "g1_box_tracking", "mujoco", "G1BoxTracking") + _, flip, _ = _materialize_profile( + "offpolicy", "sac/g1_flip_tracking", "mujoco", "G1FlipTrackingSAC" + ) + _, wall, _ = _materialize_profile( + "offpolicy", + "sac/g1_wall_flip_tracking", + "mujoco", + "G1WallFlipTrackingSAC", + ) + _, wbt, _ = _materialize_profile("offpolicy", "sac/g1_wbt_obs", "mujoco", "G1WBTObs") + _, x2, _ = _materialize_profile("ppo", "x2_wall_flip_tracking", "mujoco", "X2WallFlipTracking") + + assert set(box.scene.entities) == {"robot", "object"} + assert isinstance(box.commands["motion"], BoxMotionCommandCfg) + assert box.observations["actor"].terms["motion_anchor_pos_b"] is None + assert box.observations["actor"].terms["base_ang_vel"].params == {"sensor_name": "pelvis_gyro"} + assert box.observations["critic"].terms["base_ang_vel"].params == {"sensor_name": "pelvis_gyro"} + assert box.observations["critic"].terms["object_state"] is not None + + assert flip.commands["motion"].params.sampling_mode == "mixed" + assert flip.commands["motion"].params.sampling_start_ratio == pytest.approx(0.1) + assert wall.commands["motion"].params.sampling_mode == "uniform" + assert wall.terminations["undesired_contacts"] is None + assert wall.terminations["anchor_pos"].params["threshold"] == pytest.approx(1.0e9) + + actor_terms = wbt.observations["actor"].terms + critic_terms = wbt.observations["critic"].terms + assert actor_terms["motion_anchor_pos_b"] is None + assert actor_terms["base_lin_vel"] is None + assert [ + actor_terms[name].history_length + for name in ("base_ang_vel", "joint_pos", "joint_vel", "actions") + ] == [5, 5, 5, 5] + assert actor_terms["joint_pos"].func.__name__ == "motion_joint_pos_rel_biased" + assert critic_terms["joint_pos"].func.__name__ == "motion_joint_pos_rel" + assert wbt.actions["joint_pos"].simulate_action_latency is True + assert list(wbt.events) == [ + "base_mass", + "base_com", + "pd_gains", + "foot_friction", + "encoder_bias", + "push_robot", + ] + + assert len(x2.scene.entities["robot"].joint_names) == 29 + assert x2.scene.entities["robot"].geom_names is None + assert x2.scene.visual_model_file is not None + assert x2.observations["actor"].terms["motion_anchor_pos_b"] is None + assert x2.observations["actor"].terms["base_lin_vel"] is None + + +def test_manager_factory_selects_robot_from_multiple_floating_entities() -> None: + from unilab.base.entity import EntityCfg + from unilab.base.scene import SceneCfg + from unilab.envs.manager_based_rl_env import _resolve_backend_entity_contract + + cfg = ManagerBasedRlEnvCfg( + scene=SceneCfg( + model_file="scene.xml", + entities={ + "robot": EntityCfg(root_body_name="pelvis", body_names=("pelvis",)), + "object": EntityCfg(root_body_name="largebox"), + }, + ) + ) + + assert _resolve_backend_entity_contract(cfg) == ("pelvis", True) + + cfg.scene.entities = { + "first": EntityCfg(root_body_name="first"), + "second": EntityCfg(root_body_name="second"), + } + with pytest.raises(ValueError, match="conventional 'robot' root entity"): + _resolve_backend_entity_contract(cfg) + + +def test_x2_factory_resolves_meshes_only_before_manager_construction( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from unilab.tasks.motion_tracking import x2 + + calls: list[tuple[str, str]] = [] + sentinel = object() + monkeypatch.setattr( + x2, + "resolve_robot_asset_dir", + lambda path, marker: calls.append((path, marker)), + ) + monkeypatch.setattr(x2, "make_manager_based_rl_env", lambda *args, **kwargs: sentinel) + + result = x2.make_x2_wall_flip_env( + ManagerBasedRlEnvCfg(), + num_envs=3, + backend_type="motrix", + ) + + assert result is sentinel + assert calls == [("robots/x2/meshes", "pelvis.STL")] + + +def test_joint_acc_reset_updates_selected_rows_without_pairwise_indexing() -> None: + from unilab.managers import ManagerTermBaseCfg, SceneEntityCfg + from unilab.tasks.motion_tracking.g1.manager_terms import joint_acc_l2 + + velocity = np.arange(12, dtype=np.float32).reshape(4, 3) + entity = SimpleNamespace(num_joints=3, data=SimpleNamespace(joint_vel=velocity)) + env = SimpleNamespace(num_envs=4, step_dt=0.02, scene={"robot": entity}) + cfg = ManagerTermBaseCfg( + func=joint_acc_l2, + params={"asset_cfg": SceneEntityCfg("robot", joint_ids=[0, 2])}, + ) + term = joint_acc_l2(cfg, cast(Any, env)) + velocity[[1, 3]] += 10.0 + + term.reset(np.array([1, 3], dtype=np.int32)) + + np.testing.assert_array_equal(term._previous[[1, 3]], velocity[np.ix_([1, 3], [0, 2])]) + np.testing.assert_array_equal(term._previous[[0, 2]], np.array([[0.0, 2.0], [6.0, 8.0]])) + + +@pytest.mark.parametrize( + ("task", "identity", "backend", "actor_dim", "critic_dim", "action_dim"), + ( + ("g1_box_tracking", "G1BoxTracking", "mujoco", 154, 298, 29), + ("g1_box_tracking", "G1BoxTracking", "motrix", 154, 298, 29), + ("x2_wall_flip_tracking", "X2WallFlipTracking", "mujoco", 154, 430, 29), + ("x2_wall_flip_tracking", "X2WallFlipTracking", "motrix", 154, 430, 29), + ("sac/g1_wbt_obs", "G1WBTObs", "mujoco", 514, 289, 29), + ("sac/g1_23dof_wbt_obs", "G1WBTObs23Dof", "mujoco", 412, 259, 23), + ), +) +def test_representative_motion_profiles_reset_and_step( + task: str, + identity: str, + backend: str, + actor_dim: int, + critic_dim: int, + action_dim: int, +) -> None: + if backend == "mujoco": + pytest.importorskip("mujoco") + try: + from mujoco_uni.batch_env import BatchEnvPool as _ # noqa: F401 + except Exception: + pytest.skip("mujoco_uni.batch_env not available") + else: + pytest.importorskip("motrixsim") + + config_root = "offpolicy" if task.startswith("sac/") else "ppo" + _, _, override = _materialize_profile(config_root, task, backend, identity) + env = registry.make(identity, num_envs=2, sim_backend=backend, env_cfg_override=override) + try: + initial = env.init_state() + assert env.obs_groups_spec == {"obs": actor_dim, "critic": critic_dim} + assert initial.obs["obs"].shape == (2, actor_dim) + assert initial.obs["critic"].shape == (2, critic_dim) + + state = env.step(np.zeros((2, action_dim), dtype=np.float32)) + assert state.reward.shape == (2,) + assert np.isfinite(state.reward).all() + assert all(np.isfinite(value).all() for value in state.obs.values()) + finally: + env.close() diff --git a/tests/envs/test_motion_tracking_rewards.py b/tests/envs/test_motion_tracking_rewards.py deleted file mode 100644 index 5c6783cef..000000000 --- a/tests/envs/test_motion_tracking_rewards.py +++ /dev/null @@ -1,155 +0,0 @@ -"""Guard tests for the extracted motion-tracking reward module. - -These lock the ``common.rewards`` term functions and dispatch against -hand-computed values on a small synthetic ``RewardContext``. -""" - -from __future__ import annotations - -from types import SimpleNamespace - -import numpy as np - -from unilab.tasks.motion_tracking.common import rewards -from unilab.tasks.motion_tracking.common.rewards import RewardConfig, RewardContext - - -def _make_ctx(*, scales: dict[str, float] | None = None) -> RewardContext: - """Build a small deterministic RewardContext (2 envs, 3 bodies, 2 joints).""" - rng = np.random.default_rng(7) - num_envs, n_body, n_action = 2, 3, 2 - dtype = np.float64 - - def rq(shape: tuple[int, ...]) -> np.ndarray: - q = rng.normal(size=(*shape, 4)).astype(dtype) - q /= np.linalg.norm(q, axis=-1, keepdims=True) - return q - - reward_config = RewardConfig() - if scales is not None: - reward_config.scales = scales - - motion_data = SimpleNamespace( - joint_pos=rng.normal(size=(num_envs, n_action)).astype(dtype), - joint_vel=rng.normal(size=(num_envs, n_action)).astype(dtype), - body_pos_w=rng.normal(size=(num_envs, n_body, 3)).astype(dtype), - body_quat_w=rq((num_envs, n_body)), - body_lin_vel_w=rng.normal(size=(num_envs, n_body, 3)).astype(dtype), - body_ang_vel_w=rng.normal(size=(num_envs, n_body, 3)).astype(dtype), - ) - - robot_body_pos_w = rng.normal(size=(num_envs, n_body, 3)).astype(dtype) - # Force a deterministic contact pattern for the undesired_contacts term. - robot_body_pos_w[:, :, 2] = np.array([[0.01, 0.20, 0.03], [0.10, 0.02, 0.30]], dtype=dtype) - undesired_idx = np.array([0, 1], dtype=np.int32) - - current_actions = rng.normal(size=(num_envs, n_action)).astype(dtype) - last_actions = rng.normal(size=(num_envs, n_action)).astype(dtype) - - return RewardContext( - info={ - "current_actions": current_actions, - "last_actions": last_actions, - "steps": np.zeros((num_envs,), dtype=np.uint32), - }, - motion_data=motion_data, - robot_body_pos_w=robot_body_pos_w, - robot_body_quat_w=rq((num_envs, n_body)), - robot_body_lin_vel_w=rng.normal(size=(num_envs, n_body, 3)).astype(dtype), - robot_body_ang_vel_w=rng.normal(size=(num_envs, n_body, 3)).astype(dtype), - ref_body_pos_w=rng.normal(size=(num_envs, n_body, 3)).astype(dtype), - ref_body_quat_w=rq((num_envs, n_body)), - dof_pos=rng.normal(size=(num_envs, n_action)).astype(dtype), - dof_vel=rng.normal(size=(num_envs, n_action)).astype(dtype), - reward_config=reward_config, - anchor_body_idx=0, - ee_body_indices=np.array([2], dtype=np.int32), - undesired_contact_body_indices=undesired_idx, - joint_lower=np.array([-0.5, -0.5], dtype=dtype), - joint_upper=np.array([0.5, 0.5], dtype=dtype), - undesired_contact_z_threshold=0.05, - num_envs=num_envs, - body_vec_error=np.empty((num_envs, n_body, 3), dtype=dtype), - joint_error=np.empty((num_envs, n_action), dtype=dtype), - joint_error_upper=np.empty((num_envs, n_action), dtype=dtype), - env_error=np.empty((num_envs,), dtype=dtype), - env_error2=np.empty((num_envs,), dtype=dtype), - reward_term=np.empty((num_envs,), dtype=dtype), - weighted_reward=np.empty((num_envs,), dtype=dtype), - quat_error_w=np.empty((num_envs, n_body), dtype=dtype), - quat_error_x=np.empty((num_envs, n_body), dtype=dtype), - ee_pos_error_z=np.empty((num_envs, 1), dtype=dtype), - undesired_contact_mask=np.empty((num_envs, undesired_idx.size), dtype=bool), - ) - - -def test_build_reward_functions_contains_all_canonical_terms(): - assert set(rewards.build_reward_functions()) == { - "motion_global_root_pos", - "motion_global_root_ori", - "motion_body_pos", - "motion_body_ori", - "motion_body_lin_vel", - "motion_body_ang_vel", - "motion_ee_body_pos_z", - "motion_joint_pos", - "motion_joint_vel", - "action_rate_l2", - "joint_limit", - "undesired_contacts", - } - - -def test_motion_joint_pos_term_matches_hand_computed(): - ctx = _make_ctx() - out = rewards.motion_joint_pos(ctx).copy() - error = np.mean(np.square(ctx.motion_data.joint_pos - ctx.dof_pos), axis=1) - expected = np.exp(-error / ctx.reward_config.std_joint_pos**2) - np.testing.assert_allclose(out, expected, rtol=1e-12, atol=1e-12) - - -def test_action_rate_l2_term_matches_hand_computed(): - ctx = _make_ctx() - out = rewards.action_rate_l2(ctx).copy() - expected = np.sum(np.square(ctx.info["current_actions"] - ctx.info["last_actions"]), axis=1) - np.testing.assert_allclose(out, expected, rtol=1e-12, atol=1e-12) - - -def test_undesired_contacts_term_matches_hand_computed(): - ctx = _make_ctx() - out = rewards.undesired_contacts(ctx).copy() - idx = ctx.undesired_contact_body_indices - expected = np.sum( - ctx.robot_body_pos_w[:, idx, 2] < ctx.undesired_contact_z_threshold, axis=1 - ).astype(np.float64) - np.testing.assert_allclose(out, expected, rtol=1e-12, atol=1e-12) - - -def test_compute_reward_matches_hand_computed_weighted_sum(): - scales = { - "motion_joint_pos": 1.0, - "action_rate_l2": -0.1, - "undesired_contacts": -0.5, - # zero-weighted terms are skipped by compute_reward - "motion_body_pos": 0.0, - } - ctrl_dt = 0.02 - fns = rewards.build_reward_functions() - - # Reference terms computed on an untouched ctx (fresh buffers per call). - joint_pos_ref = rewards.motion_joint_pos(_make_ctx()).copy() - action_rate_ref = rewards.action_rate_l2(_make_ctx()).copy() - undesired_ref = rewards.undesired_contacts(_make_ctx()).copy() - - ctx = _make_ctx(scales=scales) - reward = rewards.compute_reward( - ctx, - active_reward_fns=fns, - all_reward_fns=fns, - scales=scales, - ctrl_dt=ctrl_dt, - enable_log=False, - ).copy() - - expected = (1.0 * joint_pos_ref - 0.1 * action_rate_ref - 0.5 * undesired_ref) * ctrl_dt - np.testing.assert_allclose(reward, expected, rtol=1e-12, atol=1e-12) diff --git a/tests/scripts/test_obs_alignment_g1_wbt.py b/tests/scripts/test_obs_alignment_g1_wbt.py index 93b5323b8..29b7c3f79 100644 --- a/tests/scripts/test_obs_alignment_g1_wbt.py +++ b/tests/scripts/test_obs_alignment_g1_wbt.py @@ -4,9 +4,8 @@ pipeline produces byte-identical actor obs at every step. Three independent implementations are exercised against the SAME inputs: - 1. Training side — tracking_obs.py's _push_obs_history / - _fill_obs_history + actor obs assembly in - _build_actor_obs (replicated below in numpy). + 1. Training side — ObservationManager's per-term CircularBuffer history + and term-major actor assembly (replicated below in NumPy). 2. Schema side — sim_prototype.ObsAssembler driven by deploy_config.yaml. 3. Deploy side — observation_manager.h::ObservationTermCfg semantics replicated in Python (oldest-first deque per term, @@ -208,13 +207,11 @@ def test_history_terms_carry_oldest_first(self, deploy_cfg, rng): class TestTrainingAssemblerVsDeploy: - """Replicate training-side history maintenance (np ring buffer) and compare. + """Replicate training-side ObservationManager history and compare. - This replicates the exact buffer logic from tracking_obs.py: - * _fill_obs_history: buf[:, :] = val[:, None, :] - * _push_obs_history: buf[:, :-1] = buf[:, 1:]; buf[:, -1] = val - and the actor obs assembly: refs first, then per-term history blocks - (gyro, joint_pos_rel, dof_vel, last_actions) flattened (n_env, H*dim). + The manager resets each term history by filling it with the current sample, + then shifts oldest-first and appends each new value. Concatenation is + term-major, matching the deploy-side observation layout. """ @staticmethod @@ -236,7 +233,7 @@ def test_training_path_matches_deploy(self, deploy_cfg, rng): n = deploy_cfg["action_dim"] H = 5 - # Initial buffer of zeros (matches tracking_obs.py allocation). + # Initial buffers before ObservationManager.reset fills active rows. buf = { "gyro": np.zeros((n_env, H, 3), dtype=np.float32), "joint_pos_rel": np.zeros((n_env, H, n), dtype=np.float32), @@ -247,7 +244,7 @@ def test_training_path_matches_deploy(self, deploy_cfg, rng): all_segments = [_random_segments(rng) for _ in range(15)] deploy_seq = _deploy_compute_group(deploy_cfg["obs_layout"], all_segments) - # Step 0: reset (matches tracking_obs.py is_reset=True path). + # Step 0: reset fills every history slot with the current value. s0 = all_segments[0] for key in ("gyro", "joint_pos_rel", "dof_vel", "last_actions"): buf[key][:, :, :] = s0[key][None, None, :] @@ -259,7 +256,7 @@ def test_training_path_matches_deploy(self, deploy_cfg, rng): train_obs0 = self._training_actor_obs(buf, refs0, s0, n_env) np.testing.assert_array_equal(train_obs0[0], deploy_seq[0]) - # Steps 1..N-1: push (matches tracking_obs.py is_reset=False path). + # Steps 1..N-1: evict oldest and append current. for k in range(1, len(all_segments)): sk = all_segments[k] for key in ("gyro", "joint_pos_rel", "dof_vel", "last_actions"): diff --git a/tests/scripts/test_train_script_configs.py b/tests/scripts/test_train_script_configs.py index ba0ea43d8..08ff973d1 100644 --- a/tests/scripts/test_train_script_configs.py +++ b/tests/scripts/test_train_script_configs.py @@ -80,7 +80,7 @@ def _appo_motion_file_overrides(task: str, tmp_path: Path) -> list[str]: return [] motion_file = tmp_path / f"{task.split('/', 1)[0]}_smoke_motion.npz" _write_g1_motion_smoke_npz(motion_file) - return [f"++env.motion_file={motion_file}"] + return [f"env.commands.motion.params.motion_file={motion_file}"] def test_appo_mujoco_smoke_tasks_have_owner_configs(): diff --git a/tests/tasks/test_migration_matrix.py b/tests/tasks/test_migration_matrix.py index 77bfb065a..1da07a41c 100644 --- a/tests/tasks/test_migration_matrix.py +++ b/tests/tasks/test_migration_matrix.py @@ -18,7 +18,7 @@ def test_registered_tasks_have_explicit_migration_records() -> None: assert PRODUCTION_TASK_NAMES <= registered.keys() assert {record.task_name for record in records} == set(PRODUCTION_TASK_NAMES) assert len(records) == 39 - assert sum(record.status == "Compatible" for record in records) == 21 + assert sum(record.status == "Compatible" for record in records) == 36 assert sum(record.target == "compatibility" for record in records) == 3 @@ -28,6 +28,8 @@ def test_registered_tasks_have_explicit_migration_records() -> None: ("Go2ArmManipLoco", "go2_arm", "compatibility", "Adapted"), ("SharpaInhandRotation", "sharpa", "compatibility", "Adapted"), ("G1MotionTracking", "motion_tracking", "complete", "Compatible"), + ("G1WBTObs", "motion_tracking", "complete", "Compatible"), + ("X2WallFlipTracking", "motion_tracking", "complete", "Compatible"), ("G1WalkRough", "g1_locomotion", "complete", "Compatible"), ("Go2JoystickRough", "quadruped_rough", "complete", "Compatible"), ], diff --git a/tests/training/test_training_helpers.py b/tests/training/test_training_helpers.py index 8ff811352..4e12db895 100644 --- a/tests/training/test_training_helpers.py +++ b/tests/training/test_training_helpers.py @@ -412,6 +412,32 @@ def _fake_materializer(source_model_file: str, **kwargs) -> str: assert captured == {} +def test_backend_adapter_materializes_visuals_without_dropping_manager_entities(): + cfg = _ppo_cfg(["task=g1_box_tracking/motrix", "training.play_only=true"]) + captured: dict[str, object] = {} + + def _fake_materializer(source_model_file: str, **kwargs) -> str: + captured["source_model_file"] = source_model_file + captured.update(kwargs) + return "/tmp/materialized_box.xml" + + env_cfg_override = BackendAdapter( + cfg, + root_dir=_ROOT_DIR, + algo_name="ppo", + scene_materializer=_fake_materializer, + ).build_play_env_cfg_override() + + scene = env_cfg_override["scene"] + assert scene["model_file"] == "/tmp/materialized_box.xml" + assert scene["default_keyframe_name"] == "stand" + assert set(scene["entities"]) == {"robot", "object"} + assert scene["entities"]["object"]["root_body_name"] == "largebox" + assert str(captured["source_model_file"]).endswith( + "src/unilab/assets/robots/g1/scene_flat_with_largebox.xml" + ) + + def test_render_play_mode_uses_env_interactive_contract(): class FakeEnv: def __init__(self): From 6f72bc26e7107ee281d900eb0aaa02ff98231b81 Mon Sep 17 00:00:00 2001 From: tatp-yf Date: Fri, 21 Aug 2026 01:01:22 +0800 Subject: [PATCH 096/169] refactor: remove torch layer from unilab.algos package Move src/unilab/algos/torch/* up to src/unilab/algos/ and update all references across src, scripts, tests, conf, docs, pyproject and .gitignore. Fix _REPO_ROOT parents index in hora/distill_config.py after the move, and fold api_reference algos/torch.md into index.md. --- .gitignore | 2 +- conf/appo/task/sharpa_inhand/mujoco_hora.yaml | 6 +-- .../student_model/hora_actor.yaml | 2 +- conf/hora_distill/student_model/hora_sac.yaml | 2 +- .../task/sac/sharpa_inhand/mujoco_hora.yaml | 2 +- conf/ppo/config.yaml | 2 +- conf/ppo/task/sharpa_inhand/mujoco_hora.yaml | 8 ++-- .../source/api_reference/algos/index.md | 29 ++++++++---- .../source/api_reference/algos/torch.md | 29 ------------ .../en/2-user_guide/2-algorithms/1-ppo.md | 2 +- .../en/2-user_guide/2-algorithms/2-appo.md | 2 +- .../2-user_guide/2-algorithms/5-flash_sac.md | 2 +- .../en/2-user_guide/2-algorithms/7-hora.md | 4 +- .../2-from_legged_gym.md | 4 +- .../3-framework_migration/3-from_rsl_rl.md | 4 +- .../3-framework_migration/4-from_skrl.md | 2 +- .../1-architecture/2-runtime_model.md | 4 +- .../3-extending/3-new_algorithm.md | 2 +- .../zh_CN/2-user_guide/2-algorithms/1-ppo.md | 2 +- .../zh_CN/2-user_guide/2-algorithms/2-appo.md | 2 +- .../2-user_guide/2-algorithms/5-flash_sac.md | 2 +- .../zh_CN/2-user_guide/2-algorithms/7-hora.md | 4 +- .../2-from_legged_gym.md | 4 +- .../3-framework_migration/3-from_rsl_rl.md | 4 +- .../3-framework_migration/4-from_skrl.md | 2 +- .../1-architecture/2-runtime_model.md | 4 +- .../3-extending/3-new_algorithm.md | 2 +- pyproject.rocm.toml | 6 +-- pyproject.toml | 4 +- scripts/play_interactive.py | 2 +- scripts/train_appo.py | 2 +- scripts/train_him_ppo.py | 2 +- scripts/train_hora_distill.py | 16 +++---- scripts/train_offpolicy.py | 10 ++--- scripts/train_rsl_rl.py | 2 +- src/unilab/algos/{torch => appo}/__init__.py | 0 src/unilab/algos/{torch => }/appo/learner.py | 2 +- src/unilab/algos/{torch => }/appo/runner.py | 6 +-- src/unilab/algos/{torch => }/appo/runtime.py | 2 +- src/unilab/algos/{torch => }/appo/staging.py | 0 src/unilab/algos/{torch => }/appo/worker.py | 2 +- src/unilab/algos/common/__init__.py | 18 ++++++++ .../algos/{torch => }/common/actor_factory.py | 8 ++-- .../{torch => }/common/collector_timing.py | 0 .../algos/{torch => }/common/compile.py | 0 src/unilab/algos/{torch => }/common/device.py | 0 .../algos/{torch => }/common/networks.py | 0 .../algos/{torch => }/common/normalization.py | 0 .../algos/{torch => }/common/stability.py | 0 .../algos/{torch => }/fast_sac/__init__.py | 0 .../{torch => }/fast_sac/double_buffer.py | 6 +-- .../algos/{torch => }/fast_sac/learner.py | 4 +- .../algos/{torch => }/fast_sac/runner.py | 4 +- .../algos/{torch => }/fast_td3/__init__.py | 0 .../{torch => }/fast_td3/double_buffer.py | 6 +-- .../algos/{torch => }/fast_td3/learner.py | 6 +-- src/unilab/algos/flash_sac/__init__.py | 12 +++++ .../{torch => }/flash_sac/double_buffer.py | 4 +- .../algos/{torch => }/flash_sac/layers.py | 0 .../algos/{torch => }/flash_sac/learner.py | 8 ++-- .../algos/{torch => }/flash_sac/network.py | 2 +- .../algos/{torch => }/flash_sac/runner.py | 4 +- .../algos/{torch => }/flash_sac/update.py | 0 src/unilab/algos/him_ppo/__init__.py | 11 +++++ .../algos/{torch => }/him_ppo/actor_critic.py | 2 +- .../algos/{torch => }/him_ppo/algorithm.py | 4 +- .../algos/{torch => }/him_ppo/estimator.py | 0 .../algos/{torch => }/him_ppo/runner.py | 4 +- .../algos/{torch => }/him_ppo/storage.py | 0 src/unilab/algos/{torch => }/hora/__init__.py | 0 src/unilab/algos/{torch => }/hora/appo.py | 4 +- .../algos/{torch => }/hora/appo_learner.py | 4 +- .../algos/{torch => }/hora/appo_runner.py | 12 ++--- .../algos/{torch => }/hora/appo_worker.py | 8 ++-- src/unilab/algos/{torch => }/hora/distill.py | 8 ++-- .../algos/{torch => }/hora/distill_config.py | 2 +- src/unilab/algos/{torch => }/hora/models.py | 0 .../algos/{torch => }/hora/observations.py | 0 src/unilab/algos/{torch => }/hora/ppo.py | 4 +- src/unilab/algos/{torch => }/hora/rsl_rl.py | 0 .../algos/{torch => }/hora/rsl_rl_compat.py | 0 src/unilab/algos/{torch => }/hora/runtime.py | 0 src/unilab/algos/{torch => }/hora/sac.py | 6 +-- .../algos/{torch => }/hora/sac_learner.py | 4 +- .../algos/{torch => }/hora/sac_models.py | 0 .../algos/{torch => }/offpolicy/__init__.py | 4 +- .../offpolicy/double_buffer_runner.py | 6 +-- .../algos/{torch => }/offpolicy/runner.py | 2 +- .../algos/{torch => }/offpolicy/runtime.py | 0 .../{torch => }/offpolicy/thread_budget.py | 0 .../algos/{torch => }/offpolicy/worker.py | 6 +-- src/unilab/algos/{torch => }/rsl_rl_ppo.py | 2 +- .../algos/{torch => }/rsl_rl_runtime.py | 0 src/unilab/algos/torch/appo/__init__.py | 0 src/unilab/algos/torch/common/__init__.py | 18 -------- src/unilab/algos/torch/flash_sac/__init__.py | 12 ----- src/unilab/algos/torch/him_ppo/__init__.py | 11 ----- src/unilab/structured_configs.py | 2 +- src/unilab/training/offpolicy.py | 8 ++-- .../visualization/interactive_playback.py | 22 +++++----- tests/algos/test_appo_learner.py | 4 +- tests/algos/test_appo_learner_metrics.py | 2 +- tests/algos/test_appo_runner.py | 2 +- tests/algos/test_appo_runner_unit.py | 4 +- tests/algos/test_appo_staging.py | 2 +- tests/algos/test_appo_worker.py | 2 +- tests/algos/test_fast_sac_compile.py | 2 +- .../algos/test_fast_sac_symmetry_contract.py | 24 +++++----- tests/algos/test_fast_td3_learner.py | 2 +- tests/algos/test_flash_sac_learner.py | 4 +- tests/algos/test_him_ppo_runner.py | 8 ++-- tests/algos/test_hora_contract.py | 44 +++++++++---------- tests/algos/test_hora_distill_config.py | 6 +-- tests/algos/test_hora_imports.py | 20 ++++----- .../test_offpolicy_bootstrap_contract.py | 4 +- .../test_offpolicy_double_buffer_runner.py | 28 ++++++------ tests/algos/test_offpolicy_dp_sync.py | 24 +++++----- tests/algos/test_offpolicy_runner_unit.py | 16 +++---- tests/algos/test_offpolicy_runtime.py | 8 ++-- tests/algos/test_offpolicy_thread_budget.py | 2 +- tests/algos/test_offpolicy_worker.py | 4 +- tests/algos/test_rsl_rl_ppo.py | 4 +- tests/algos/test_torch_compile_helper.py | 2 +- tests/config/test_locomotion_params.py | 2 +- .../ipc/test_replay_pipeline_gpu_resident.py | 2 +- tests/nan_injection/proto_him_ppo_inject.py | 2 +- tests/nan_injection/stage2_nan_inject.py | 2 +- tests/nan_injection/stage3_nan_inject.py | 8 ++-- tests/scripts/test_repo_hygiene.py | 4 +- tests/scripts/test_train_scripts.py | 20 ++++----- tests/utils/test_algo_utils.py | 2 +- tests/utils/test_utils_package_policy.py | 6 +-- .../test_interactive_playback.py | 12 ++--- 133 files changed, 349 insertions(+), 367 deletions(-) delete mode 100644 docs/sphinx/source/api_reference/algos/torch.md rename src/unilab/algos/{torch => appo}/__init__.py (100%) rename src/unilab/algos/{torch => }/appo/learner.py (99%) rename src/unilab/algos/{torch => }/appo/runner.py (99%) rename src/unilab/algos/{torch => }/appo/runtime.py (97%) rename src/unilab/algos/{torch => }/appo/staging.py (100%) rename src/unilab/algos/{torch => }/appo/worker.py (99%) create mode 100644 src/unilab/algos/common/__init__.py rename src/unilab/algos/{torch => }/common/actor_factory.py (88%) rename src/unilab/algos/{torch => }/common/collector_timing.py (100%) rename src/unilab/algos/{torch => }/common/compile.py (100%) rename src/unilab/algos/{torch => }/common/device.py (100%) rename src/unilab/algos/{torch => }/common/networks.py (100%) rename src/unilab/algos/{torch => }/common/normalization.py (100%) rename src/unilab/algos/{torch => }/common/stability.py (100%) rename src/unilab/algos/{torch => }/fast_sac/__init__.py (100%) rename src/unilab/algos/{torch => }/fast_sac/double_buffer.py (96%) rename src/unilab/algos/{torch => }/fast_sac/learner.py (99%) rename src/unilab/algos/{torch => }/fast_sac/runner.py (97%) rename src/unilab/algos/{torch => }/fast_td3/__init__.py (100%) rename src/unilab/algos/{torch => }/fast_td3/double_buffer.py (93%) rename src/unilab/algos/{torch => }/fast_td3/learner.py (98%) create mode 100644 src/unilab/algos/flash_sac/__init__.py rename src/unilab/algos/{torch => }/flash_sac/double_buffer.py (97%) rename src/unilab/algos/{torch => }/flash_sac/layers.py (100%) rename src/unilab/algos/{torch => }/flash_sac/learner.py (99%) rename src/unilab/algos/{torch => }/flash_sac/network.py (99%) rename src/unilab/algos/{torch => }/flash_sac/runner.py (97%) rename src/unilab/algos/{torch => }/flash_sac/update.py (100%) create mode 100644 src/unilab/algos/him_ppo/__init__.py rename src/unilab/algos/{torch => }/him_ppo/actor_critic.py (98%) rename src/unilab/algos/{torch => }/him_ppo/algorithm.py (98%) rename src/unilab/algos/{torch => }/him_ppo/estimator.py (100%) rename src/unilab/algos/{torch => }/him_ppo/runner.py (99%) rename src/unilab/algos/{torch => }/him_ppo/storage.py (100%) rename src/unilab/algos/{torch => }/hora/__init__.py (100%) rename src/unilab/algos/{torch => }/hora/appo.py (98%) rename src/unilab/algos/{torch => }/hora/appo_learner.py (97%) rename src/unilab/algos/{torch => }/hora/appo_runner.py (97%) rename src/unilab/algos/{torch => }/hora/appo_worker.py (98%) rename src/unilab/algos/{torch => }/hora/distill.py (98%) rename src/unilab/algos/{torch => }/hora/distill_config.py (99%) rename src/unilab/algos/{torch => }/hora/models.py (100%) rename src/unilab/algos/{torch => }/hora/observations.py (100%) rename src/unilab/algos/{torch => }/hora/ppo.py (98%) rename src/unilab/algos/{torch => }/hora/rsl_rl.py (100%) rename src/unilab/algos/{torch => }/hora/rsl_rl_compat.py (100%) rename src/unilab/algos/{torch => }/hora/runtime.py (100%) rename src/unilab/algos/{torch => }/hora/sac.py (88%) rename src/unilab/algos/{torch => }/hora/sac_learner.py (97%) rename src/unilab/algos/{torch => }/hora/sac_models.py (100%) rename src/unilab/algos/{torch => }/offpolicy/__init__.py (56%) rename src/unilab/algos/{torch => }/offpolicy/double_buffer_runner.py (99%) rename src/unilab/algos/{torch => }/offpolicy/runner.py (99%) rename src/unilab/algos/{torch => }/offpolicy/runtime.py (100%) rename src/unilab/algos/{torch => }/offpolicy/thread_budget.py (100%) rename src/unilab/algos/{torch => }/offpolicy/worker.py (98%) rename src/unilab/algos/{torch => }/rsl_rl_ppo.py (99%) rename src/unilab/algos/{torch => }/rsl_rl_runtime.py (100%) delete mode 100644 src/unilab/algos/torch/appo/__init__.py delete mode 100644 src/unilab/algos/torch/common/__init__.py delete mode 100644 src/unilab/algos/torch/flash_sac/__init__.py delete mode 100644 src/unilab/algos/torch/him_ppo/__init__.py diff --git a/.gitignore b/.gitignore index d1de5a693..0e0c0777f 100644 --- a/.gitignore +++ b/.gitignore @@ -58,7 +58,7 @@ run_summary.json src/unilab/assets/checkpoints/ scripts/benchmark/outputs/ -src/unilab/algos/torch/rsl_rl +src/unilab/algos/rsl_rl third-party temp/ diff --git a/conf/appo/task/sharpa_inhand/mujoco_hora.yaml b/conf/appo/task/sharpa_inhand/mujoco_hora.yaml index 47aa3947a..5bdcab1b5 100644 --- a/conf/appo/task/sharpa_inhand/mujoco_hora.yaml +++ b/conf/appo/task/sharpa_inhand/mujoco_hora.yaml @@ -23,7 +23,7 @@ interactive: algo: algo_log_name: hora_appo runtime_impl: hora_appo - runtime_resolver: unilab.algos.torch.hora.appo:resolve_hora_appo_runtime + runtime_resolver: unilab.algos.hora.appo:resolve_hora_appo_runtime num_envs: 2048 steps_per_env: 8 max_iterations: 305 @@ -38,11 +38,11 @@ algo: actor: 0 priv_info: 0 actor: - class_name: unilab.algos.torch.hora:HoraActorModel + class_name: unilab.algos.hora:HoraActorModel priv_info_embed_dim: 9 priv_mlp_hidden_dims: [256, 128, 9] critic: - class_name: unilab.algos.torch.hora:HoraCriticModel + class_name: unilab.algos.hora:HoraCriticModel priv_info_embed_dim: 9 priv_mlp_hidden_dims: [256, 128, 9] algorithm: diff --git a/conf/hora_distill/student_model/hora_actor.yaml b/conf/hora_distill/student_model/hora_actor.yaml index 4aa117957..05c17bc83 100644 --- a/conf/hora_distill/student_model/hora_actor.yaml +++ b/conf/hora_distill/student_model/hora_actor.yaml @@ -1,6 +1,6 @@ # Teacher -> student `algo.model` mapping for HoraActorModel teachers (PPO/APPO). # -# Loaded by unilab.algos.torch.hora.distill_config and merged next to the +# Loaded by unilab.algos.hora.distill_config and merged next to the # Hydra-composed teacher owner config mounted at `teacher_owner`, so every # student field interpolates directly from the teacher owner YAML. The teacher # owner config stays the single source of truth for these hyperparameters; diff --git a/conf/hora_distill/student_model/hora_sac.yaml b/conf/hora_distill/student_model/hora_sac.yaml index a462608f0..c9f2e9a2b 100644 --- a/conf/hora_distill/student_model/hora_sac.yaml +++ b/conf/hora_distill/student_model/hora_sac.yaml @@ -1,6 +1,6 @@ # Teacher -> student `algo.model` mapping for hora_sac teachers (offpolicy SAC). # -# Loaded by unilab.algos.torch.hora.distill_config and merged next to the +# Loaded by unilab.algos.hora.distill_config and merged next to the # Hydra-composed teacher owner config mounted at `teacher_owner`. Values come # from the teacher owner YAML; the `oc.select` fallback after the comma only # applies when the teacher owner config does not define the field at all. diff --git a/conf/offpolicy/task/sac/sharpa_inhand/mujoco_hora.yaml b/conf/offpolicy/task/sac/sharpa_inhand/mujoco_hora.yaml index aa69445f8..cbe8210bd 100644 --- a/conf/offpolicy/task/sac/sharpa_inhand/mujoco_hora.yaml +++ b/conf/offpolicy/task/sac/sharpa_inhand/mujoco_hora.yaml @@ -24,7 +24,7 @@ interactive: algo: algo_log_name: hora_sac runtime_impl: hora_sac - runtime_resolver: unilab.algos.torch.hora.sac:resolve_hora_sac_runtime + runtime_resolver: unilab.algos.hora.sac:resolve_hora_sac_runtime num_envs: 1024 batch_size: 2048 replay_buffer_n: 1280 diff --git a/conf/ppo/config.yaml b/conf/ppo/config.yaml index a5210ce18..2954e01ae 100644 --- a/conf/ppo/config.yaml +++ b/conf/ppo/config.yaml @@ -28,7 +28,7 @@ algo: activation: elu class_name: ActorCritic algorithm: - class_name: unilab.algos.torch.rsl_rl_ppo:FinalObservationAwarePPO + class_name: unilab.algos.rsl_rl_ppo:FinalObservationAwarePPO value_loss_coef: 1.0 use_clipped_value_loss: true clip_param: 0.2 diff --git a/conf/ppo/task/sharpa_inhand/mujoco_hora.yaml b/conf/ppo/task/sharpa_inhand/mujoco_hora.yaml index 7937b66ee..7a7f8c13d 100644 --- a/conf/ppo/task/sharpa_inhand/mujoco_hora.yaml +++ b/conf/ppo/task/sharpa_inhand/mujoco_hora.yaml @@ -14,12 +14,12 @@ interactive: algo: algo_log_name: hora_ppo runtime_impl: hora_ppo - runtime_resolver: unilab.algos.torch.hora.rsl_rl:resolve_hora_ppo_runtime + runtime_resolver: unilab.algos.hora.rsl_rl:resolve_hora_ppo_runtime obs_groups: actor: [actor] critic: [actor] actor: - class_name: unilab.algos.torch.hora:HoraActorModel + class_name: unilab.algos.hora:HoraActorModel hidden_dims: [512, 256, 128] activation: elu obs_normalization: true @@ -30,14 +30,14 @@ algo: init_std: 1.0 std_type: scalar critic: - class_name: unilab.algos.torch.hora:HoraCriticModel + class_name: unilab.algos.hora:HoraCriticModel hidden_dims: [512, 256, 128] activation: elu obs_normalization: true priv_info_embed_dim: 9 priv_mlp_hidden_dims: [256, 128, 9] algorithm: - class_name: unilab.algos.torch.hora:HoraPPO + class_name: unilab.algos.hora:HoraPPO env: obs: diff --git a/docs/sphinx/source/api_reference/algos/index.md b/docs/sphinx/source/api_reference/algos/index.md index 41c42cb8e..94a4ff6ab 100644 --- a/docs/sphinx/source/api_reference/algos/index.md +++ b/docs/sphinx/source/api_reference/algos/index.md @@ -1,22 +1,35 @@ # `unilab.algos` — Learning Algorithms -- **`unilab.algos.torch`** — PPO (RSL-RL), APPO, FastSAC, FastTD3, FlashSAC, +- **`unilab.algos`** — PPO (RSL-RL), APPO, FastSAC, FastTD3, FlashSAC, HIM-PPO, HORA + distillation, generic off-policy runner. All trainers conform to a single runner contract — see {doc}`../../en/4-developer_guide/2-contracts/5-runner_lifecycle`. -```{toctree} -:maxdepth: 2 - -torch -``` - ```{eval-rst} .. autosummary:: :toctree: _autosummary :template: autosummary/module.rst :recursive: - unilab.algos + unilab.algos.common + unilab.algos.appo + unilab.algos.fast_sac + unilab.algos.fast_td3 + unilab.algos.flash_sac + unilab.algos.him_ppo + unilab.algos.hora + unilab.algos.offpolicy +``` + +## Standalone PPO entrypoints + +```{eval-rst} +.. automodule:: unilab.algos.rsl_rl_ppo + :members: +``` + +```{eval-rst} +.. automodule:: unilab.algos.rsl_rl_runtime + :members: ``` diff --git a/docs/sphinx/source/api_reference/algos/torch.md b/docs/sphinx/source/api_reference/algos/torch.md deleted file mode 100644 index 347557974..000000000 --- a/docs/sphinx/source/api_reference/algos/torch.md +++ /dev/null @@ -1,29 +0,0 @@ -# `unilab.algos.torch` - -```{eval-rst} -.. autosummary:: - :toctree: _autosummary - :template: autosummary/module.rst - :recursive: - - unilab.algos.torch.common - unilab.algos.torch.appo - unilab.algos.torch.fast_sac - unilab.algos.torch.fast_td3 - unilab.algos.torch.flash_sac - unilab.algos.torch.him_ppo - unilab.algos.torch.hora - unilab.algos.torch.offpolicy -``` - -## Standalone PPO entrypoints - -```{eval-rst} -.. automodule:: unilab.algos.torch.rsl_rl_ppo - :members: -``` - -```{eval-rst} -.. automodule:: unilab.algos.torch.rsl_rl_runtime - :members: -``` diff --git a/docs/sphinx/source/en/2-user_guide/2-algorithms/1-ppo.md b/docs/sphinx/source/en/2-user_guide/2-algorithms/1-ppo.md index e52e20183..bbbca179e 100644 --- a/docs/sphinx/source/en/2-user_guide/2-algorithms/1-ppo.md +++ b/docs/sphinx/source/en/2-user_guide/2-algorithms/1-ppo.md @@ -2,7 +2,7 @@ PPO is the default synchronous on-policy training path. It uses `scripts/train_rsl_rl.py`, composes from `conf/ppo/config.yaml`, and runs the -RSL-RL adapter code in `src/unilab/algos/torch/rsl_rl_ppo.py` and +RSL-RL adapter code in `src/unilab/algos/rsl_rl_ppo.py` and `src/unilab/training/rsl_rl.py`. ## Quick Start diff --git a/docs/sphinx/source/en/2-user_guide/2-algorithms/2-appo.md b/docs/sphinx/source/en/2-user_guide/2-algorithms/2-appo.md index aea0cc9c9..f095dab65 100644 --- a/docs/sphinx/source/en/2-user_guide/2-algorithms/2-appo.md +++ b/docs/sphinx/source/en/2-user_guide/2-algorithms/2-appo.md @@ -1,7 +1,7 @@ # APPO APPO is UniLab's asynchronous PPO path. It uses `scripts/train_appo.py`, -`conf/appo/config.yaml`, and the runtime under `src/unilab/algos/torch/appo/`. +`conf/appo/config.yaml`, and the runtime under `src/unilab/algos/appo/`. The config exposes `algo.steps_per_env`, `training.collector_device`, and `training.replay_queue_size`; the algorithm config includes V-trace clipping fields. diff --git a/docs/sphinx/source/en/2-user_guide/2-algorithms/5-flash_sac.md b/docs/sphinx/source/en/2-user_guide/2-algorithms/5-flash_sac.md index f1ce54b40..d1ef9c75f 100644 --- a/docs/sphinx/source/en/2-user_guide/2-algorithms/5-flash_sac.md +++ b/docs/sphinx/source/en/2-user_guide/2-algorithms/5-flash_sac.md @@ -3,7 +3,7 @@ FlashSAC is the third algorithm on the shared off-policy entrypoint. Select it with `--algo flashsac`; defaults live in `conf/offpolicy/algo/flashsac.yaml`, and the implementation lives under -`src/unilab/algos/torch/flash_sac/`. +`src/unilab/algos/flash_sac/`. It shares the off-policy training script with SAC and TD3, but does not use the same default networks: the actor uses a block-based structure and the critic diff --git a/docs/sphinx/source/en/2-user_guide/2-algorithms/7-hora.md b/docs/sphinx/source/en/2-user_guide/2-algorithms/7-hora.md index 808ccf28b..d320f633e 100644 --- a/docs/sphinx/source/en/2-user_guide/2-algorithms/7-hora.md +++ b/docs/sphinx/source/en/2-user_guide/2-algorithms/7-hora.md @@ -13,7 +13,7 @@ uv run train --algo appo --task sharpa_inhand --sim mujoco --profile hora traini ``` The HORA PPO owner sets `algo.algo_log_name=hora_ppo` and resolves the runtime -through `unilab.algos.torch.hora.rsl_rl:resolve_hora_ppo_runtime`. The APPO +through `unilab.algos.hora.rsl_rl:resolve_hora_ppo_runtime`. The APPO variant sets `algo.algo_log_name=hora_appo`. ## Student Distillation @@ -24,5 +24,5 @@ CLI does not currently declare a separate HORA distillation `--algo` route, so the public CLI examples on this page stay on the teacher path above. Teacher checkpoint resolution is implemented in -`src/unilab/algos/torch/hora/distill_config.py`. The student log family is +`src/unilab/algos/hora/distill_config.py`. The student log family is `hora_distill`. diff --git a/docs/sphinx/source/en/3-deployment/3-framework_migration/2-from_legged_gym.md b/docs/sphinx/source/en/3-deployment/3-framework_migration/2-from_legged_gym.md index 02117cfe0..16d82862b 100644 --- a/docs/sphinx/source/en/3-deployment/3-framework_migration/2-from_legged_gym.md +++ b/docs/sphinx/source/en/3-deployment/3-framework_migration/2-from_legged_gym.md @@ -14,7 +14,7 @@ mostly mechanical. | `_reward_*` methods | env's `compute_reward()` + reward term registry | | `command_ranges` | task owner YAML's `commands` block | | Terrain curriculum | {doc}`../../2-user_guide/6-terrain/1-procedural` | -| RSL-RL PPO | `unilab.algos.torch.rsl_rl_ppo` | +| RSL-RL PPO | `unilab.algos.rsl_rl_ppo` | ## What's new @@ -22,7 +22,7 @@ mostly mechanical. + Motrix. Pick one (or both) before porting; see {doc}`../2-sim_to_sim/1-backend_swap`. - **Async collection.** Legged Gym collects on-GPU synchronously; UniLab's - APPO (`unilab.algos.torch.appo`) decouples collectors from + APPO (`unilab.algos.appo`) decouples collectors from learner. If wall-clock matters, port to APPO once your reward parity is established. - **Hardware deployment.** Legged Gym → real-world deployment is a diff --git a/docs/sphinx/source/en/3-deployment/3-framework_migration/3-from_rsl_rl.md b/docs/sphinx/source/en/3-deployment/3-framework_migration/3-from_rsl_rl.md index c8d2c5f9f..a104ceb59 100644 --- a/docs/sphinx/source/en/3-deployment/3-framework_migration/3-from_rsl_rl.md +++ b/docs/sphinx/source/en/3-deployment/3-framework_migration/3-from_rsl_rl.md @@ -1,7 +1,7 @@ # Migrating from RSL-RL You're already using RSL-RL standalone? Good news: UniLab ships RSL-RL PPO -as one of its supported algorithms (`unilab.algos.torch.rsl_rl_ppo`) +as one of its supported algorithms (`unilab.algos.rsl_rl_ppo`) and it's nearly drop-in. ## What you gain by moving inside UniLab @@ -13,7 +13,7 @@ and it's nearly drop-in. backend / task / algo selection. No more bespoke train scripts per robot. 3. **Async runner.** Wrap RSL-RL PPO inside - `unilab.algos.torch.appo` for higher throughput on machines + `unilab.algos.appo` for higher throughput on machines with many CPU cores. 4. **Deployment story.** ONNX export with the right wrapper, safety layer documentation, and the diff --git a/docs/sphinx/source/en/3-deployment/3-framework_migration/4-from_skrl.md b/docs/sphinx/source/en/3-deployment/3-framework_migration/4-from_skrl.md index 55d748e1e..c4d7616e6 100644 --- a/docs/sphinx/source/en/3-deployment/3-framework_migration/4-from_skrl.md +++ b/docs/sphinx/source/en/3-deployment/3-framework_migration/4-from_skrl.md @@ -8,7 +8,7 @@ deployment path. | skrl | UniLab | |---|---| -| `Agent` (PPO, SAC, …) | `unilab.algos.torch.*` | +| `Agent` (PPO, SAC, …) | `unilab.algos.*` | | `RolloutMemory` | `unilab.ipc.rollout_ring_buffer` | | `ReplayMemory` | `unilab.ipc.replay_buffer` | | `Trainer` | `unilab.training.run` | diff --git a/docs/sphinx/source/en/4-developer_guide/1-architecture/2-runtime_model.md b/docs/sphinx/source/en/4-developer_guide/1-architecture/2-runtime_model.md index b2989ba31..423775d28 100644 --- a/docs/sphinx/source/en/4-developer_guide/1-architecture/2-runtime_model.md +++ b/docs/sphinx/source/en/4-developer_guide/1-architecture/2-runtime_model.md @@ -41,8 +41,8 @@ CPU physics env loop -> shared IPC buffer -> learner ## Evidence In Repo - PPO entrypoint: `scripts/train_rsl_rl.py` -- APPO runner: `src/unilab/algos/torch/appo/runner.py` -- Off-policy runner: `src/unilab/algos/torch/offpolicy/double_buffer_runner.py` +- APPO runner: `src/unilab/algos/appo/runner.py` +- Off-policy runner: `src/unilab/algos/offpolicy/double_buffer_runner.py` - IPC primitives: `src/unilab/ipc/async_runner.py`, `src/unilab/ipc/rollout_ring_buffer.py`, `src/unilab/ipc/replay_buffer.py`, `src/unilab/ipc/weight_sync.py` diff --git a/docs/sphinx/source/en/4-developer_guide/3-extending/3-new_algorithm.md b/docs/sphinx/source/en/4-developer_guide/3-extending/3-new_algorithm.md index 8d15edeec..ab2411298 100644 --- a/docs/sphinx/source/en/4-developer_guide/3-extending/3-new_algorithm.md +++ b/docs/sphinx/source/en/4-developer_guide/3-extending/3-new_algorithm.md @@ -41,4 +41,4 @@ Algorithm work must preserve the env, config, and runner contracts. Start with - Structured config dataclasses: `src/unilab/structured_configs.py` - Training helpers: `src/unilab/training/common.py`, `src/unilab/training/run.py` -- Existing algorithm packages: `src/unilab/algos/torch/` +- Existing algorithm packages: `src/unilab/algos/` diff --git a/docs/sphinx/source/zh_CN/2-user_guide/2-algorithms/1-ppo.md b/docs/sphinx/source/zh_CN/2-user_guide/2-algorithms/1-ppo.md index 5a900bfc1..61afe0bc7 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/2-algorithms/1-ppo.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/2-algorithms/1-ppo.md @@ -1,7 +1,7 @@ # PPO PPO 是默认的同步 on-policy 训练路径。它使用 `scripts/train_rsl_rl.py`，从 -`conf/ppo/config.yaml` 组合配置，并运行 `src/unilab/algos/torch/rsl_rl_ppo.py` +`conf/ppo/config.yaml` 组合配置，并运行 `src/unilab/algos/rsl_rl_ppo.py` 和 `src/unilab/training/rsl_rl.py` 中的 RSL-RL 适配代码。 ## 快速开始 diff --git a/docs/sphinx/source/zh_CN/2-user_guide/2-algorithms/2-appo.md b/docs/sphinx/source/zh_CN/2-user_guide/2-algorithms/2-appo.md index d15b568b5..251e352c1 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/2-algorithms/2-appo.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/2-algorithms/2-appo.md @@ -1,7 +1,7 @@ # APPO APPO 是 UniLab 的异步 PPO 路径。它使用 `scripts/train_appo.py`、 -`conf/appo/config.yaml` 以及 `src/unilab/algos/torch/appo/` 下的运行时。该配置暴露 +`conf/appo/config.yaml` 以及 `src/unilab/algos/appo/` 下的运行时。该配置暴露 了 `algo.steps_per_env`、`training.collector_device` 和 `training.replay_queue_size`；算法配置中包含 V-trace 裁剪字段。 diff --git a/docs/sphinx/source/zh_CN/2-user_guide/2-algorithms/5-flash_sac.md b/docs/sphinx/source/zh_CN/2-user_guide/2-algorithms/5-flash_sac.md index e1f8c5911..24e83ca82 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/2-algorithms/5-flash_sac.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/2-algorithms/5-flash_sac.md @@ -2,7 +2,7 @@ FlashSAC 是共享 off-policy 入口上的第三个算法。使用 `--algo flashsac` 选择它；默认 值位于 `conf/offpolicy/algo/flashsac.yaml`，实现位于 -`src/unilab/algos/torch/flash_sac/` 下。 +`src/unilab/algos/flash_sac/` 下。 它与 SAC、TD3 共用 off-policy 训练脚本，但默认网络并不相同：actor 使用 block-based 结构，critic 使用 distributional（categorical）Q 变体。 diff --git a/docs/sphinx/source/zh_CN/2-user_guide/2-algorithms/7-hora.md b/docs/sphinx/source/zh_CN/2-user_guide/2-algorithms/7-hora.md index bf623244b..1cddce6ed 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/2-algorithms/7-hora.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/2-algorithms/7-hora.md @@ -13,7 +13,7 @@ uv run train --algo appo --task sharpa_inhand --sim mujoco --profile hora traini ``` HORA PPO owner 设置 `algo.algo_log_name=hora_ppo`，并通过 -`unilab.algos.torch.hora.rsl_rl:resolve_hora_ppo_runtime` 解析运行时。APPO 变体设置 +`unilab.algos.hora.rsl_rl:resolve_hora_ppo_runtime` 解析运行时。APPO 变体设置 `algo.algo_log_name=hora_appo`。 ## Student 蒸馏 @@ -22,5 +22,5 @@ student 蒸馏由 `scripts/train_hora_distill.py` 实现，并由 `conf/hora_distill/task/sharpa_inhand/mujoco.yaml` 配置。顶层 CLI 目前没有声明独立的 HORA 蒸馏 `--algo` 路由，因此本页的公开 CLI 示例仍保持在上面的 teacher 路径上。 -teacher 检查点的解析在 `src/unilab/algos/torch/hora/distill_config.py` 中实现。 +teacher 检查点的解析在 `src/unilab/algos/hora/distill_config.py` 中实现。 student 日志族为 `hora_distill`。 diff --git a/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/2-from_legged_gym.md b/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/2-from_legged_gym.md index 026a674b9..b4e4b775e 100644 --- a/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/2-from_legged_gym.md +++ b/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/2-from_legged_gym.md @@ -13,7 +13,7 @@ Legged Gym 曾是那套 GPU 常驻的 PPO 模板，教会了整个领域如何 | `_reward_*` 方法 | env 的 `compute_reward()` + reward 项 registry | | `command_ranges` | 任务 owner YAML 的 `commands` 块 | | 地形课程 | {doc}`../../2-user_guide/6-terrain/1-procedural` | -| RSL-RL PPO | `unilab.algos.torch.rsl_rl_ppo` | +| RSL-RL PPO | `unilab.algos.rsl_rl_ppo` | ## 有哪些新东西 @@ -21,7 +21,7 @@ Legged Gym 曾是那套 GPU 常驻的 PPO 模板，教会了整个领域如何 在移植之前先选一个（或两个都选）；参见 {doc}`../2-sim_to_sim/1-backend_swap`。 - **异步采集。** Legged Gym 在 GPU 上同步采集；UniLab 的 - APPO（`unilab.algos.torch.appo`）把 collector 与 learner 解耦。如果你在意 + APPO（`unilab.algos.appo`）把 collector 与 learner 解耦。如果你在意 wall-clock 时间，在建立起 reward 一致性之后，就移植到 APPO。 - **硬件部署。** Legged Gym → 真实世界部署，是各实验室各自手工搭建的流程。UniLab 把 {doc}`../1-sim_to_real/1-overview` 流水线作为一等公民产物提供给你。 diff --git a/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/3-from_rsl_rl.md b/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/3-from_rsl_rl.md index a00068b26..a011bb933 100644 --- a/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/3-from_rsl_rl.md +++ b/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/3-from_rsl_rl.md @@ -1,7 +1,7 @@ # 从 RSL-RL 迁移 你已经在独立使用 RSL-RL 了？好消息：UniLab 把 RSL-RL PPO 作为其受支持算法之一 -（`unilab.algos.torch.rsl_rl_ppo`）提供，而且几乎是即插即用的。 +（`unilab.algos.rsl_rl_ppo`）提供，而且几乎是即插即用的。 ## 迁移进 UniLab 后你能获得什么 @@ -10,7 +10,7 @@ 更不容易出错。 2. **任务 owner。** 基于 Hydra 的配置组合，外加 registry 驱动的 backend / task / algo 选择。不再需要为每种机器人编写定制的训练脚本。 -3. **异步 runner。** 把 RSL-RL PPO 包进 `unilab.algos.torch.appo`，在拥有许多 +3. **异步 runner。** 把 RSL-RL PPO 包进 `unilab.algos.appo`，在拥有许多 CPU 核心的机器上获得更高吞吐量。 4. **部署流程。** 配合正确 wrapper 的 ONNX 导出、安全层文档，以及 {doc}`../1-sim_to_real/1-overview` 流水线。 diff --git a/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/4-from_skrl.md b/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/4-from_skrl.md index 01415f9ba..0eaf6de38 100644 --- a/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/4-from_skrl.md +++ b/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/4-from_skrl.md @@ -7,7 +7,7 @@ skrl 的强项在于算法广度。UniLab 专注于一组精选算法（PPO、SA | skrl | UniLab | |---|---| -| `Agent`（PPO、SAC……） | `unilab.algos.torch.*` | +| `Agent`（PPO、SAC……） | `unilab.algos.*` | | `RolloutMemory` | `unilab.ipc.rollout_ring_buffer` | | `ReplayMemory` | `unilab.ipc.replay_buffer` | | `Trainer` | `unilab.training.run` | diff --git a/docs/sphinx/source/zh_CN/4-developer_guide/1-architecture/2-runtime_model.md b/docs/sphinx/source/zh_CN/4-developer_guide/1-architecture/2-runtime_model.md index ca937096b..349b3431c 100644 --- a/docs/sphinx/source/zh_CN/4-developer_guide/1-architecture/2-runtime_model.md +++ b/docs/sphinx/source/zh_CN/4-developer_guide/1-architecture/2-runtime_model.md @@ -50,8 +50,8 @@ CPU physics env loop -> shared IPC buffer -> learner ## 仓库中的证据 - PPO 入口：`scripts/train_rsl_rl.py` -- APPO runner：`src/unilab/algos/torch/appo/runner.py` -- Off-policy runner：`src/unilab/algos/torch/offpolicy/double_buffer_runner.py` +- APPO runner：`src/unilab/algos/appo/runner.py` +- Off-policy runner：`src/unilab/algos/offpolicy/double_buffer_runner.py` - IPC 原语：`src/unilab/ipc/async_runner.py`、 `src/unilab/ipc/rollout_ring_buffer.py`、`src/unilab/ipc/replay_buffer.py`、 `src/unilab/ipc/weight_sync.py` diff --git a/docs/sphinx/source/zh_CN/4-developer_guide/3-extending/3-new_algorithm.md b/docs/sphinx/source/zh_CN/4-developer_guide/3-extending/3-new_algorithm.md index 7e90fc4b9..f693b2d61 100644 --- a/docs/sphinx/source/zh_CN/4-developer_guide/3-extending/3-new_algorithm.md +++ b/docs/sphinx/source/zh_CN/4-developer_guide/3-extending/3-new_algorithm.md @@ -41,4 +41,4 @@ - 结构化 config dataclass：`src/unilab/structured_configs.py` - 训练辅助工具：`src/unilab/training/common.py`、 `src/unilab/training/run.py` -- 现有算法包：`src/unilab/algos/torch/` +- 现有算法包：`src/unilab/algos/` diff --git a/pyproject.rocm.toml b/pyproject.rocm.toml index 3cff64e32..b0051db67 100644 --- a/pyproject.rocm.toml +++ b/pyproject.rocm.toml @@ -127,7 +127,7 @@ warn_unused_configs = true ignore_missing_imports = true no_site_packages = true exclude = [ - "src/unilab/algos/torch/rsl_rl/", + "src/unilab/algos/rsl_rl/", ] [tool.pytest.ini_options] @@ -150,8 +150,8 @@ venvPath = "." venv = ".venv" include = ["src/unilab"] exclude = [ - "src/unilab/algos/torch/rsl_rl/", - "src/unilab/algos/torch/common/ane_*", + "src/unilab/algos/rsl_rl/", + "src/unilab/algos/common/ane_*", "src/unilab/base/backend/", "src/unilab/envs/", "src/unilab/terrains/", diff --git a/pyproject.toml b/pyproject.toml index df809ab77..dfae55b18 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -139,7 +139,7 @@ warn_unused_configs = true ignore_missing_imports = true no_site_packages = true exclude = [ - "src/unilab/algos/torch/rsl_rl/", # vendored third-party, tracked in .gitignore + "src/unilab/algos/rsl_rl/", # vendored third-party, tracked in .gitignore ] [tool.pytest.ini_options] @@ -162,7 +162,7 @@ venvPath = "." venv = ".venv" include = ["src/unilab"] exclude = [ - "src/unilab/algos/torch/rsl_rl/", # vendored third-party + "src/unilab/algos/rsl_rl/", # vendored third-party "src/unilab/base/backend/", # mujoco-uni-runtime stubs mismatch; optional backends "src/unilab/envs/", # mujoco-uni-runtime internal API, stubs mismatch "src/unilab/terrains/", # mujoco-uni-runtime MjSpec API, stubs mismatch diff --git a/scripts/play_interactive.py b/scripts/play_interactive.py index c0bff8328..c2e8efcb5 100644 --- a/scripts/play_interactive.py +++ b/scripts/play_interactive.py @@ -1050,7 +1050,7 @@ def _create_env(num_envs: int): if algo == "ppo": wrapper_cls = RslRlVecEnvWrapper if cfg is not None: - from unilab.algos.torch.rsl_rl_runtime import resolve_rsl_rl_ppo_runtime + from unilab.algos.rsl_rl_runtime import resolve_rsl_rl_ppo_runtime wrapper_cls = resolve_rsl_rl_ppo_runtime( _algo_config_dict(cfg), diff --git a/scripts/train_appo.py b/scripts/train_appo.py index 93e193429..f3d3daf9d 100644 --- a/scripts/train_appo.py +++ b/scripts/train_appo.py @@ -16,7 +16,7 @@ ROOT_DIR = Path(__file__).parent.parent sys.path.append(str(ROOT_DIR)) -from unilab.algos.torch.appo.runtime import resolve_appo_runtime +from unilab.algos.appo.runtime import resolve_appo_runtime from unilab.training import ( BackendAdapter, apply_configured_training_seed, diff --git a/scripts/train_him_ppo.py b/scripts/train_him_ppo.py index e104aa42f..559f7ad47 100644 --- a/scripts/train_him_ppo.py +++ b/scripts/train_him_ppo.py @@ -18,7 +18,7 @@ if str(ROOT_DIR) not in sys.path: sys.path.insert(0, str(ROOT_DIR)) -from unilab.algos.torch.him_ppo.runner import HIMOnPolicyRunner +from unilab.algos.him_ppo.runner import HIMOnPolicyRunner from unilab.base.backend import materialize_scene_visual_override from unilab.training import ( BackendAdapter, diff --git a/scripts/train_hora_distill.py b/scripts/train_hora_distill.py index d289424e9..b16383398 100644 --- a/scripts/train_hora_distill.py +++ b/scripts/train_hora_distill.py @@ -15,29 +15,29 @@ if str(ROOT_DIR) not in sys.path: sys.path.insert(0, str(ROOT_DIR)) -from unilab.algos.torch.hora import HoraDistillationTrainer -from unilab.algos.torch.hora.distill import ( +from unilab.algos.hora import HoraDistillationTrainer +from unilab.algos.hora.distill import ( build_student_actor_and_normalizer, cfg_with_checkpoint_runtime, load_distilled_checkpoint, student_policy, ) -from unilab.algos.torch.hora.distill_config import ( +from unilab.algos.hora.distill_config import ( apply_teacher_defaults as _apply_teacher_defaults, ) -from unilab.algos.torch.hora.distill_config import ( +from unilab.algos.hora.distill_config import ( get_teacher_owner_spec as _get_teacher_owner_spec, ) -from unilab.algos.torch.hora.distill_config import ( +from unilab.algos.hora.distill_config import ( resolve_teacher_checkpoint_path as _resolve_teacher_checkpoint_path, ) -from unilab.algos.torch.hora.distill_config import ( +from unilab.algos.hora.distill_config import ( resolved_distill_runtime_cfg as _resolved_distill_runtime_cfg, ) -from unilab.algos.torch.hora.distill_config import ( +from unilab.algos.hora.distill_config import ( teacher_run_metadata as _teacher_run_metadata, ) -from unilab.algos.torch.hora.rsl_rl import HoraRslRlVecEnvWrapper as RslRlVecEnvWrapper +from unilab.algos.hora.rsl_rl import HoraRslRlVecEnvWrapper as RslRlVecEnvWrapper from unilab.base.backend import materialize_scene_visual_override from unilab.training import ( BackendAdapter, diff --git a/scripts/train_offpolicy.py b/scripts/train_offpolicy.py index e490d64d8..5604056a2 100644 --- a/scripts/train_offpolicy.py +++ b/scripts/train_offpolicy.py @@ -92,7 +92,7 @@ def build_offpolicy_env_cfg_override(algo_name: str, cfg: DictConfig) -> dict[st def build_runner(algo_name: str, cfg: DictConfig, log_dir: str | None = None): """Build algorithm runner from unified Hydra config.""" env_cfg_override = build_offpolicy_env_cfg_override(algo_name, cfg) - from unilab.algos.torch.offpolicy.thread_budget import ( + from unilab.algos.offpolicy.thread_budget import ( apply_torch_thread_runtime, resolve_torch_thread_runtime, ) @@ -165,7 +165,7 @@ def build_runner(algo_name: str, cfg: DictConfig, log_dir: str | None = None): replay_device = require_offpolicy_replay_device(rank_device) if algo_name == "sac": - from unilab.algos.torch.fast_sac.double_buffer import ( + from unilab.algos.fast_sac.double_buffer import ( build_sac_double_buffer_runner, ) @@ -181,7 +181,7 @@ def build_runner(algo_name: str, cfg: DictConfig, log_dir: str | None = None): ) if algo_name == "td3": - from unilab.algos.torch.fast_td3.double_buffer import ( + from unilab.algos.fast_td3.double_buffer import ( build_td3_double_buffer_runner, ) @@ -197,7 +197,7 @@ def build_runner(algo_name: str, cfg: DictConfig, log_dir: str | None = None): ) if algo_name == "flashsac": - from unilab.algos.torch.flash_sac.double_buffer import ( + from unilab.algos.flash_sac.double_buffer import ( build_flashsac_double_buffer_runner, ) @@ -220,7 +220,7 @@ def play_offpolicy(algo_name: str, cfg: DictConfig) -> str | None: import numpy as np import torch - from unilab.algos.torch.offpolicy.worker import resolve_offpolicy_actor_priv_info + from unilab.algos.offpolicy.worker import resolve_offpolicy_actor_priv_info load_path, load_path_dir = resolve_checkpoint_path( ROOT_DIR, diff --git a/scripts/train_rsl_rl.py b/scripts/train_rsl_rl.py index feea13c9c..bdadf732b 100644 --- a/scripts/train_rsl_rl.py +++ b/scripts/train_rsl_rl.py @@ -17,7 +17,7 @@ if str(ROOT_DIR) not in sys.path: sys.path.insert(0, str(ROOT_DIR)) -from unilab.algos.torch.rsl_rl_runtime import resolve_rsl_rl_ppo_runtime +from unilab.algos.rsl_rl_runtime import resolve_rsl_rl_ppo_runtime from unilab.base.backend import RenderClosedError, materialize_scene_visual_override from unilab.base.run_control import RunComplete from unilab.ipc.dp_launcher import ( diff --git a/src/unilab/algos/torch/__init__.py b/src/unilab/algos/appo/__init__.py similarity index 100% rename from src/unilab/algos/torch/__init__.py rename to src/unilab/algos/appo/__init__.py diff --git a/src/unilab/algos/torch/appo/learner.py b/src/unilab/algos/appo/learner.py similarity index 99% rename from src/unilab/algos/torch/appo/learner.py rename to src/unilab/algos/appo/learner.py index 966ad2aaf..71305b917 100644 --- a/src/unilab/algos/torch/appo/learner.py +++ b/src/unilab/algos/appo/learner.py @@ -21,7 +21,7 @@ from rsl_rl.utils import resolve_optimizer from tensordict import TensorDict -from unilab.algos.torch.common.compile import get_torch_compile_for_cuda +from unilab.algos.common.compile import get_torch_compile_for_cuda _LOG_2_PI = math.log(2.0 * math.pi) _NORMAL_ENTROPY_OFFSET = 0.5 * (1.0 + _LOG_2_PI) diff --git a/src/unilab/algos/torch/appo/runner.py b/src/unilab/algos/appo/runner.py similarity index 99% rename from src/unilab/algos/torch/appo/runner.py rename to src/unilab/algos/appo/runner.py index 22800aaff..851d25290 100644 --- a/src/unilab/algos/torch/appo/runner.py +++ b/src/unilab/algos/appo/runner.py @@ -17,9 +17,9 @@ import torch from rsl_rl.utils import resolve_callable -from unilab.algos.torch.appo.learner import APPOLearner -from unilab.algos.torch.appo.staging import RolloutStagingPool -from unilab.algos.torch.appo.worker import appo_collector_fn +from unilab.algos.appo.learner import APPOLearner +from unilab.algos.appo.staging import RolloutStagingPool +from unilab.algos.appo.worker import appo_collector_fn from unilab.ipc import AsyncRunner, RolloutRingBuffer, SharedWeightSync from unilab.logging import OffPolicyLogger from unilab.training.seed import apply_training_seed, derive_worker_seed diff --git a/src/unilab/algos/torch/appo/runtime.py b/src/unilab/algos/appo/runtime.py similarity index 97% rename from src/unilab/algos/torch/appo/runtime.py rename to src/unilab/algos/appo/runtime.py index ad5b58bd8..d48131f18 100644 --- a/src/unilab/algos/torch/appo/runtime.py +++ b/src/unilab/algos/appo/runtime.py @@ -46,7 +46,7 @@ def resolve_appo_runtime( """ runtime_resolver = rl_cfg.get("runtime_resolver") if runtime_resolver in (None, ""): - from unilab.algos.torch.appo.runner import APPORunner + from unilab.algos.appo.runner import APPORunner return APPORuntime(runner_cls=APPORunner, play_fn=default_play_fn) diff --git a/src/unilab/algos/torch/appo/staging.py b/src/unilab/algos/appo/staging.py similarity index 100% rename from src/unilab/algos/torch/appo/staging.py rename to src/unilab/algos/appo/staging.py diff --git a/src/unilab/algos/torch/appo/worker.py b/src/unilab/algos/appo/worker.py similarity index 99% rename from src/unilab/algos/torch/appo/worker.py rename to src/unilab/algos/appo/worker.py index 3058c3bca..f819031d6 100644 --- a/src/unilab/algos/torch/appo/worker.py +++ b/src/unilab/algos/appo/worker.py @@ -16,7 +16,7 @@ import torch from rsl_rl.utils import resolve_callable -from unilab.algos.torch.common.collector_timing import extract_env_step_breakdown_timing_ms +from unilab.algos.common.collector_timing import extract_env_step_breakdown_timing_ms from unilab.base.final_observation import resolve_terminal_observation_contract from unilab.base.observations import split_obs_dict from unilab.base.registry import ensure_registries diff --git a/src/unilab/algos/common/__init__.py b/src/unilab/algos/common/__init__.py new file mode 100644 index 000000000..8863fd9f4 --- /dev/null +++ b/src/unilab/algos/common/__init__.py @@ -0,0 +1,18 @@ +from unilab.algos.common.actor_factory import build_actor +from unilab.algos.common.device import get_env_dims +from unilab.algos.common.networks import Critic, DistributionalQNetwork +from unilab.algos.common.normalization import EmpiricalNormalization +from unilab.algos.common.stability import check_nan_loss, clip_gradients, safe_tensor +from unilab.base.registry import ensure_registries + +__all__ = [ + "EmpiricalNormalization", + "DistributionalQNetwork", + "Critic", + "get_env_dims", + "check_nan_loss", + "clip_gradients", + "safe_tensor", + "ensure_registries", + "build_actor", +] diff --git a/src/unilab/algos/torch/common/actor_factory.py b/src/unilab/algos/common/actor_factory.py similarity index 88% rename from src/unilab/algos/torch/common/actor_factory.py rename to src/unilab/algos/common/actor_factory.py index af1656a54..2c6eead1a 100644 --- a/src/unilab/algos/torch/common/actor_factory.py +++ b/src/unilab/algos/common/actor_factory.py @@ -21,7 +21,7 @@ def build_actor( ): """Build the correct actor model based on algorithm type.""" if algo_type == "sac": - from unilab.algos.torch.fast_sac.learner import SACActor + from unilab.algos.fast_sac.learner import SACActor return SACActor( obs_dim=obs_dim, @@ -33,7 +33,7 @@ def build_actor( if algo_type == "hora_sac": if priv_info_dim is None: raise ValueError("build_actor(algo_type='hora_sac') requires priv_info_dim.") - from unilab.algos.torch.hora.sac_models import HoraSACActor + from unilab.algos.hora.sac_models import HoraSACActor return HoraSACActor( obs_dim=obs_dim, @@ -46,7 +46,7 @@ def build_actor( device=device, ) if algo_type == "td3": - from unilab.algos.torch.fast_td3.learner import TD3Actor + from unilab.algos.fast_td3.learner import TD3Actor return TD3Actor( obs_dim=obs_dim, @@ -59,7 +59,7 @@ def build_actor( device=device, ) if algo_type == "flashsac": - from unilab.algos.torch.flash_sac.network import FlashSACActor + from unilab.algos.flash_sac.network import FlashSACActor return FlashSACActor( num_blocks=actor_num_blocks, diff --git a/src/unilab/algos/torch/common/collector_timing.py b/src/unilab/algos/common/collector_timing.py similarity index 100% rename from src/unilab/algos/torch/common/collector_timing.py rename to src/unilab/algos/common/collector_timing.py diff --git a/src/unilab/algos/torch/common/compile.py b/src/unilab/algos/common/compile.py similarity index 100% rename from src/unilab/algos/torch/common/compile.py rename to src/unilab/algos/common/compile.py diff --git a/src/unilab/algos/torch/common/device.py b/src/unilab/algos/common/device.py similarity index 100% rename from src/unilab/algos/torch/common/device.py rename to src/unilab/algos/common/device.py diff --git a/src/unilab/algos/torch/common/networks.py b/src/unilab/algos/common/networks.py similarity index 100% rename from src/unilab/algos/torch/common/networks.py rename to src/unilab/algos/common/networks.py diff --git a/src/unilab/algos/torch/common/normalization.py b/src/unilab/algos/common/normalization.py similarity index 100% rename from src/unilab/algos/torch/common/normalization.py rename to src/unilab/algos/common/normalization.py diff --git a/src/unilab/algos/torch/common/stability.py b/src/unilab/algos/common/stability.py similarity index 100% rename from src/unilab/algos/torch/common/stability.py rename to src/unilab/algos/common/stability.py diff --git a/src/unilab/algos/torch/fast_sac/__init__.py b/src/unilab/algos/fast_sac/__init__.py similarity index 100% rename from src/unilab/algos/torch/fast_sac/__init__.py rename to src/unilab/algos/fast_sac/__init__.py diff --git a/src/unilab/algos/torch/fast_sac/double_buffer.py b/src/unilab/algos/fast_sac/double_buffer.py similarity index 96% rename from src/unilab/algos/torch/fast_sac/double_buffer.py rename to src/unilab/algos/fast_sac/double_buffer.py index 88b813366..f18da10ef 100644 --- a/src/unilab/algos/torch/fast_sac/double_buffer.py +++ b/src/unilab/algos/fast_sac/double_buffer.py @@ -6,9 +6,9 @@ from omegaconf import DictConfig, OmegaConf -from unilab.algos.torch.fast_sac.learner import FastSACLearner -from unilab.algos.torch.offpolicy.double_buffer_runner import DoubleBufferOffPolicyRunner -from unilab.algos.torch.offpolicy.runtime import resolve_custom_offpolicy_runtime +from unilab.algos.fast_sac.learner import FastSACLearner +from unilab.algos.offpolicy.double_buffer_runner import DoubleBufferOffPolicyRunner +from unilab.algos.offpolicy.runtime import resolve_custom_offpolicy_runtime from unilab.base.np_env import NpEnv from unilab.training import create_env, ensure_registries from unilab.utils.nan_guard import NanGuardCfg diff --git a/src/unilab/algos/torch/fast_sac/learner.py b/src/unilab/algos/fast_sac/learner.py similarity index 99% rename from src/unilab/algos/torch/fast_sac/learner.py rename to src/unilab/algos/fast_sac/learner.py index 5d714a37c..5d281b3d3 100644 --- a/src/unilab/algos/torch/fast_sac/learner.py +++ b/src/unilab/algos/fast_sac/learner.py @@ -20,8 +20,8 @@ import torch.nn.functional as F import torch.optim as optim -from unilab.algos.torch.common.compile import get_torch_compile_for_cuda -from unilab.algos.torch.common.normalization import EmpiricalNormalization +from unilab.algos.common.compile import get_torch_compile_for_cuda +from unilab.algos.common.normalization import EmpiricalNormalization from unilab.base.augmentation import SymmetryAugmentation diff --git a/src/unilab/algos/torch/fast_sac/runner.py b/src/unilab/algos/fast_sac/runner.py similarity index 97% rename from src/unilab/algos/torch/fast_sac/runner.py rename to src/unilab/algos/fast_sac/runner.py index f6a72a8c8..5e185650a 100644 --- a/src/unilab/algos/torch/fast_sac/runner.py +++ b/src/unilab/algos/fast_sac/runner.py @@ -3,8 +3,8 @@ import logging from typing import Any -from unilab.algos.torch.fast_sac.learner import FastSACLearner -from unilab.algos.torch.offpolicy.double_buffer_runner import DoubleBufferOffPolicyRunner +from unilab.algos.fast_sac.learner import FastSACLearner +from unilab.algos.offpolicy.double_buffer_runner import DoubleBufferOffPolicyRunner from unilab.ipc.replay_pipelines.gpu_resident import require_offpolicy_replay_device from unilab.utils.device import get_default_device diff --git a/src/unilab/algos/torch/fast_td3/__init__.py b/src/unilab/algos/fast_td3/__init__.py similarity index 100% rename from src/unilab/algos/torch/fast_td3/__init__.py rename to src/unilab/algos/fast_td3/__init__.py diff --git a/src/unilab/algos/torch/fast_td3/double_buffer.py b/src/unilab/algos/fast_td3/double_buffer.py similarity index 93% rename from src/unilab/algos/torch/fast_td3/double_buffer.py rename to src/unilab/algos/fast_td3/double_buffer.py index 34f5df0bc..c260e1c2a 100644 --- a/src/unilab/algos/torch/fast_td3/double_buffer.py +++ b/src/unilab/algos/fast_td3/double_buffer.py @@ -6,9 +6,9 @@ from omegaconf import DictConfig -from unilab.algos.torch.common.device import get_env_dims -from unilab.algos.torch.fast_td3.learner import FastTD3Learner -from unilab.algos.torch.offpolicy.double_buffer_runner import DoubleBufferOffPolicyRunner +from unilab.algos.common.device import get_env_dims +from unilab.algos.fast_td3.learner import FastTD3Learner +from unilab.algos.offpolicy.double_buffer_runner import DoubleBufferOffPolicyRunner from unilab.utils.nan_guard import NanGuardCfg if TYPE_CHECKING: diff --git a/src/unilab/algos/torch/fast_td3/learner.py b/src/unilab/algos/fast_td3/learner.py similarity index 98% rename from src/unilab/algos/torch/fast_td3/learner.py rename to src/unilab/algos/fast_td3/learner.py index 772cce22e..6d7804714 100644 --- a/src/unilab/algos/torch/fast_td3/learner.py +++ b/src/unilab/algos/fast_td3/learner.py @@ -23,9 +23,9 @@ import torch.nn.functional as F import torch.optim as optim -from unilab.algos.torch.common.networks import Critic -from unilab.algos.torch.common.normalization import EmpiricalNormalization -from unilab.algos.torch.common.stability import check_nan_loss, clip_gradients +from unilab.algos.common.networks import Critic +from unilab.algos.common.normalization import EmpiricalNormalization +from unilab.algos.common.stability import check_nan_loss, clip_gradients # --------------------------------------------------------------------------- # Actor (deterministic, ReLU, per-env noise) diff --git a/src/unilab/algos/flash_sac/__init__.py b/src/unilab/algos/flash_sac/__init__.py new file mode 100644 index 000000000..7cadda2ed --- /dev/null +++ b/src/unilab/algos/flash_sac/__init__.py @@ -0,0 +1,12 @@ +"""FlashSAC algorithm package.""" + +from unilab.algos.flash_sac.learner import FlashSACLearner +from unilab.algos.flash_sac.network import FlashSACActor, FlashSACDoubleCritic +from unilab.algos.flash_sac.runner import FlashSACRunner + +__all__ = [ + "FlashSACActor", + "FlashSACDoubleCritic", + "FlashSACLearner", + "FlashSACRunner", +] diff --git a/src/unilab/algos/torch/flash_sac/double_buffer.py b/src/unilab/algos/flash_sac/double_buffer.py similarity index 97% rename from src/unilab/algos/torch/flash_sac/double_buffer.py rename to src/unilab/algos/flash_sac/double_buffer.py index 26d21869b..9a2239f09 100644 --- a/src/unilab/algos/torch/flash_sac/double_buffer.py +++ b/src/unilab/algos/flash_sac/double_buffer.py @@ -6,8 +6,8 @@ from omegaconf import DictConfig -from unilab.algos.torch.flash_sac.learner import FlashSACLearner -from unilab.algos.torch.offpolicy.double_buffer_runner import DoubleBufferOffPolicyRunner +from unilab.algos.flash_sac.learner import FlashSACLearner +from unilab.algos.offpolicy.double_buffer_runner import DoubleBufferOffPolicyRunner from unilab.ipc.replay_pipelines.gpu_resident import require_offpolicy_replay_device from unilab.training import create_env, ensure_registries from unilab.training.seed import apply_training_seed diff --git a/src/unilab/algos/torch/flash_sac/layers.py b/src/unilab/algos/flash_sac/layers.py similarity index 100% rename from src/unilab/algos/torch/flash_sac/layers.py rename to src/unilab/algos/flash_sac/layers.py diff --git a/src/unilab/algos/torch/flash_sac/learner.py b/src/unilab/algos/flash_sac/learner.py similarity index 99% rename from src/unilab/algos/torch/flash_sac/learner.py rename to src/unilab/algos/flash_sac/learner.py index 3f0285838..dc150f478 100644 --- a/src/unilab/algos/torch/flash_sac/learner.py +++ b/src/unilab/algos/flash_sac/learner.py @@ -11,14 +11,14 @@ import torch.nn as nn import torch.optim as optim -from unilab.algos.torch.common.compile import get_torch_compile_for_cuda -from unilab.algos.torch.common.normalization import EmpiricalNormalization -from unilab.algos.torch.flash_sac.network import ( +from unilab.algos.common.compile import get_torch_compile_for_cuda +from unilab.algos.common.normalization import EmpiricalNormalization +from unilab.algos.flash_sac.network import ( FlashSACActor, FlashSACDoubleCritic, FlashSACTemperature, ) -from unilab.algos.torch.flash_sac.update import ( +from unilab.algos.flash_sac.update import ( build_lr_lambda, resolve_target_entropy, select_min_q_log_probs, diff --git a/src/unilab/algos/torch/flash_sac/network.py b/src/unilab/algos/flash_sac/network.py similarity index 99% rename from src/unilab/algos/torch/flash_sac/network.py rename to src/unilab/algos/flash_sac/network.py index e810513fd..37acff9bf 100644 --- a/src/unilab/algos/torch/flash_sac/network.py +++ b/src/unilab/algos/flash_sac/network.py @@ -8,7 +8,7 @@ import torch import torch.nn as nn -from unilab.algos.torch.flash_sac.layers import ( +from unilab.algos.flash_sac.layers import ( EnsembleCategoricalValue, EnsembleFlashSACBlock, EnsembleFlashSACEmbedder, diff --git a/src/unilab/algos/torch/flash_sac/runner.py b/src/unilab/algos/flash_sac/runner.py similarity index 97% rename from src/unilab/algos/torch/flash_sac/runner.py rename to src/unilab/algos/flash_sac/runner.py index 729f49de9..877711314 100644 --- a/src/unilab/algos/torch/flash_sac/runner.py +++ b/src/unilab/algos/flash_sac/runner.py @@ -4,8 +4,8 @@ from typing import Any -from unilab.algos.torch.flash_sac.learner import FlashSACLearner -from unilab.algos.torch.offpolicy.double_buffer_runner import DoubleBufferOffPolicyRunner +from unilab.algos.flash_sac.learner import FlashSACLearner +from unilab.algos.offpolicy.double_buffer_runner import DoubleBufferOffPolicyRunner from unilab.ipc.replay_pipelines.gpu_resident import require_offpolicy_replay_device from unilab.utils.device import get_default_device diff --git a/src/unilab/algos/torch/flash_sac/update.py b/src/unilab/algos/flash_sac/update.py similarity index 100% rename from src/unilab/algos/torch/flash_sac/update.py rename to src/unilab/algos/flash_sac/update.py diff --git a/src/unilab/algos/him_ppo/__init__.py b/src/unilab/algos/him_ppo/__init__.py new file mode 100644 index 000000000..2ecfec5cd --- /dev/null +++ b/src/unilab/algos/him_ppo/__init__.py @@ -0,0 +1,11 @@ +from unilab.algos.him_ppo.actor_critic import HIMActorCritic +from unilab.algos.him_ppo.algorithm import HIMPPO +from unilab.algos.him_ppo.estimator import HIMEstimator +from unilab.algos.him_ppo.storage import HIMRolloutStorage + +__all__ = [ + "HIMActorCritic", + "HIMPPO", + "HIMEstimator", + "HIMRolloutStorage", +] diff --git a/src/unilab/algos/torch/him_ppo/actor_critic.py b/src/unilab/algos/him_ppo/actor_critic.py similarity index 98% rename from src/unilab/algos/torch/him_ppo/actor_critic.py rename to src/unilab/algos/him_ppo/actor_critic.py index 47ac66e01..dc62a4074 100644 --- a/src/unilab/algos/torch/him_ppo/actor_critic.py +++ b/src/unilab/algos/him_ppo/actor_critic.py @@ -8,7 +8,7 @@ import torch.nn as nn from torch.distributions import Normal -from unilab.algos.torch.him_ppo.estimator import HIMEstimator, get_activation +from unilab.algos.him_ppo.estimator import HIMEstimator, get_activation class HIMActorCritic(nn.Module): diff --git a/src/unilab/algos/torch/him_ppo/algorithm.py b/src/unilab/algos/him_ppo/algorithm.py similarity index 98% rename from src/unilab/algos/torch/him_ppo/algorithm.py rename to src/unilab/algos/him_ppo/algorithm.py index eba081cf0..5892479c5 100644 --- a/src/unilab/algos/torch/him_ppo/algorithm.py +++ b/src/unilab/algos/him_ppo/algorithm.py @@ -11,8 +11,8 @@ import torch.optim as optim from tensordict import TensorDict -from unilab.algos.torch.him_ppo.actor_critic import HIMActorCritic -from unilab.algos.torch.him_ppo.storage import HIMRolloutStorage +from unilab.algos.him_ppo.actor_critic import HIMActorCritic +from unilab.algos.him_ppo.storage import HIMRolloutStorage class HIMPPO: diff --git a/src/unilab/algos/torch/him_ppo/estimator.py b/src/unilab/algos/him_ppo/estimator.py similarity index 100% rename from src/unilab/algos/torch/him_ppo/estimator.py rename to src/unilab/algos/him_ppo/estimator.py diff --git a/src/unilab/algos/torch/him_ppo/runner.py b/src/unilab/algos/him_ppo/runner.py similarity index 99% rename from src/unilab/algos/torch/him_ppo/runner.py rename to src/unilab/algos/him_ppo/runner.py index bab234d42..6a8d7ab16 100644 --- a/src/unilab/algos/torch/him_ppo/runner.py +++ b/src/unilab/algos/him_ppo/runner.py @@ -12,8 +12,8 @@ import torch -from unilab.algos.torch.him_ppo.actor_critic import HIMActorCritic -from unilab.algos.torch.him_ppo.algorithm import HIMPPO +from unilab.algos.him_ppo.actor_critic import HIMActorCritic +from unilab.algos.him_ppo.algorithm import HIMPPO logger = logging.getLogger(__name__) diff --git a/src/unilab/algos/torch/him_ppo/storage.py b/src/unilab/algos/him_ppo/storage.py similarity index 100% rename from src/unilab/algos/torch/him_ppo/storage.py rename to src/unilab/algos/him_ppo/storage.py diff --git a/src/unilab/algos/torch/hora/__init__.py b/src/unilab/algos/hora/__init__.py similarity index 100% rename from src/unilab/algos/torch/hora/__init__.py rename to src/unilab/algos/hora/__init__.py diff --git a/src/unilab/algos/torch/hora/appo.py b/src/unilab/algos/hora/appo.py similarity index 98% rename from src/unilab/algos/torch/hora/appo.py rename to src/unilab/algos/hora/appo.py index e8b9835f1..7b2e930cc 100644 --- a/src/unilab/algos/torch/hora/appo.py +++ b/src/unilab/algos/hora/appo.py @@ -11,8 +11,8 @@ import torch from omegaconf import DictConfig -from unilab.algos.torch.hora.appo_runner import HoraAPPORunner -from unilab.algos.torch.hora.rsl_rl_compat import ( +from unilab.algos.hora.appo_runner import HoraAPPORunner +from unilab.algos.hora.rsl_rl_compat import ( convert_config_v3_to_v4, is_rsl_rl_v4, is_rsl_rl_v5, diff --git a/src/unilab/algos/torch/hora/appo_learner.py b/src/unilab/algos/hora/appo_learner.py similarity index 97% rename from src/unilab/algos/torch/hora/appo_learner.py rename to src/unilab/algos/hora/appo_learner.py index ed9578fed..be77ab29c 100644 --- a/src/unilab/algos/torch/hora/appo_learner.py +++ b/src/unilab/algos/hora/appo_learner.py @@ -7,13 +7,13 @@ import torch from tensordict import TensorDict -from unilab.algos.torch.appo.learner import ( +from unilab.algos.appo.learner import ( APPOLearner, _distribution_std, _sample_tensor_for_metric, vtrace_advantages, ) -from unilab.algos.torch.hora.models import HoraActorModel, HoraCriticModel +from unilab.algos.hora.models import HoraActorModel, HoraCriticModel def _build_hora_obs_td( diff --git a/src/unilab/algos/torch/hora/appo_runner.py b/src/unilab/algos/hora/appo_runner.py similarity index 97% rename from src/unilab/algos/torch/hora/appo_runner.py rename to src/unilab/algos/hora/appo_runner.py index d3d0d40c6..c780a43fd 100644 --- a/src/unilab/algos/torch/hora/appo_runner.py +++ b/src/unilab/algos/hora/appo_runner.py @@ -13,16 +13,16 @@ import torch from rsl_rl.utils import resolve_callable -from unilab.algos.torch.appo.runner import ( +from unilab.algos.appo.runner import ( APPORunner, _optimizer_lr_from_state, _sync_resume_target_actor, ) -from unilab.algos.torch.appo.staging import RolloutStagingPool -from unilab.algos.torch.hora.appo_learner import HoraAPPOLearner -from unilab.algos.torch.hora.appo_worker import hora_appo_collector_fn -from unilab.algos.torch.hora.models import build_hora_shared_actor_critic -from unilab.algos.torch.hora.rsl_rl_compat import ( +from unilab.algos.appo.staging import RolloutStagingPool +from unilab.algos.hora.appo_learner import HoraAPPOLearner +from unilab.algos.hora.appo_worker import hora_appo_collector_fn +from unilab.algos.hora.models import build_hora_shared_actor_critic +from unilab.algos.hora.rsl_rl_compat import ( convert_config_v3_to_v4, is_rsl_rl_v4, is_rsl_rl_v5, diff --git a/src/unilab/algos/torch/hora/appo_worker.py b/src/unilab/algos/hora/appo_worker.py similarity index 98% rename from src/unilab/algos/torch/hora/appo_worker.py rename to src/unilab/algos/hora/appo_worker.py index b303452aa..725edd9fb 100644 --- a/src/unilab/algos/torch/hora/appo_worker.py +++ b/src/unilab/algos/hora/appo_worker.py @@ -12,11 +12,11 @@ import torch from rsl_rl.utils import resolve_callable -from unilab.algos.torch.appo.worker import ( +from unilab.algos.appo.worker import ( compute_rollout_active_steps_per_sec, put_latest_metrics, ) -from unilab.algos.torch.common.collector_timing import extract_env_step_breakdown_timing_ms +from unilab.algos.common.collector_timing import extract_env_step_breakdown_timing_ms from unilab.base.final_observation import resolve_terminal_observation_contract from unilab.base.registry import ensure_registries from unilab.training.seed import apply_training_seed @@ -89,8 +89,8 @@ def hora_appo_collector_fn( from tensordict import TensorDict - from unilab.algos.torch.hora.models import build_hora_shared_actor_critic - from unilab.algos.torch.hora.rsl_rl_compat import ( + from unilab.algos.hora.models import build_hora_shared_actor_critic + from unilab.algos.hora.rsl_rl_compat import ( convert_config_v3_to_v4, is_rsl_rl_v4, is_rsl_rl_v5, diff --git a/src/unilab/algos/torch/hora/distill.py b/src/unilab/algos/hora/distill.py similarity index 98% rename from src/unilab/algos/torch/hora/distill.py rename to src/unilab/algos/hora/distill.py index d98bd93ea..b1f1a84be 100644 --- a/src/unilab/algos/torch/hora/distill.py +++ b/src/unilab/algos/hora/distill.py @@ -13,14 +13,14 @@ from omegaconf import DictConfig, OmegaConf from tensordict import TensorDict -from unilab.algos.torch.common.normalization import EmpiricalNormalization -from unilab.algos.torch.hora.models import ( +from unilab.algos.common.normalization import EmpiricalNormalization +from unilab.algos.hora.models import ( HoraActorModel, HoraCoreOutput, HoraSharedActorCritic, ProprioAdaptTConv, ) -from unilab.algos.torch.hora.sac_models import HoraSACActor +from unilab.algos.hora.sac_models import HoraSACActor class HoraSACDistillShared(nn.Module): @@ -306,7 +306,7 @@ def cfg_with_checkpoint_runtime(cfg: DictConfig, checkpoint: dict[str, Any]) -> Config using the current owner env/reward settings and checkpoint model construction fields. """ - from unilab.algos.torch.hora.distill_config import apply_teacher_defaults + from unilab.algos.hora.distill_config import apply_teacher_defaults cfg_with_owner_defaults = apply_teacher_defaults(cfg) cfg_clone = OmegaConf.create(OmegaConf.to_container(cfg_with_owner_defaults, resolve=False)) diff --git a/src/unilab/algos/torch/hora/distill_config.py b/src/unilab/algos/hora/distill_config.py similarity index 99% rename from src/unilab/algos/torch/hora/distill_config.py rename to src/unilab/algos/hora/distill_config.py index 7e7b4c9fb..bede00873 100644 --- a/src/unilab/algos/torch/hora/distill_config.py +++ b/src/unilab/algos/hora/distill_config.py @@ -12,7 +12,7 @@ from unilab.training.run import resolve_task_checkpoint_path -_REPO_ROOT = Path(__file__).resolve().parents[5] +_REPO_ROOT = Path(__file__).resolve().parents[4] # Teacher owner configs are Hydra-composed from their family config tree. # SAC teachers live in the shared offpolicy tree behind the `algo` group. diff --git a/src/unilab/algos/torch/hora/models.py b/src/unilab/algos/hora/models.py similarity index 100% rename from src/unilab/algos/torch/hora/models.py rename to src/unilab/algos/hora/models.py diff --git a/src/unilab/algos/torch/hora/observations.py b/src/unilab/algos/hora/observations.py similarity index 100% rename from src/unilab/algos/torch/hora/observations.py rename to src/unilab/algos/hora/observations.py diff --git a/src/unilab/algos/torch/hora/ppo.py b/src/unilab/algos/hora/ppo.py similarity index 98% rename from src/unilab/algos/torch/hora/ppo.py rename to src/unilab/algos/hora/ppo.py index 899de6fb0..56c501da8 100644 --- a/src/unilab/algos/torch/hora/ppo.py +++ b/src/unilab/algos/hora/ppo.py @@ -14,8 +14,8 @@ from rsl_rl.utils import resolve_obs_groups, resolve_optimizer from tensordict import TensorDict -from unilab.algos.torch.hora.models import HoraActorModel, HoraCriticModel, HoraSharedActorCritic -from unilab.algos.torch.rsl_rl_ppo import FinalObservationAwarePPO +from unilab.algos.hora.models import HoraActorModel, HoraCriticModel, HoraSharedActorCritic +from unilab.algos.rsl_rl_ppo import FinalObservationAwarePPO logger = logging.getLogger(__name__) diff --git a/src/unilab/algos/torch/hora/rsl_rl.py b/src/unilab/algos/hora/rsl_rl.py similarity index 100% rename from src/unilab/algos/torch/hora/rsl_rl.py rename to src/unilab/algos/hora/rsl_rl.py diff --git a/src/unilab/algos/torch/hora/rsl_rl_compat.py b/src/unilab/algos/hora/rsl_rl_compat.py similarity index 100% rename from src/unilab/algos/torch/hora/rsl_rl_compat.py rename to src/unilab/algos/hora/rsl_rl_compat.py diff --git a/src/unilab/algos/torch/hora/runtime.py b/src/unilab/algos/hora/runtime.py similarity index 100% rename from src/unilab/algos/torch/hora/runtime.py rename to src/unilab/algos/hora/runtime.py diff --git a/src/unilab/algos/torch/hora/sac.py b/src/unilab/algos/hora/sac.py similarity index 88% rename from src/unilab/algos/torch/hora/sac.py rename to src/unilab/algos/hora/sac.py index 4d0524716..ec8b807eb 100644 --- a/src/unilab/algos/torch/hora/sac.py +++ b/src/unilab/algos/hora/sac.py @@ -5,9 +5,9 @@ from dataclasses import dataclass, field from typing import Any -from unilab.algos.torch.hora.runtime import HORA_SAC_RUNTIME_IMPL, is_hora_sac_runtime -from unilab.algos.torch.hora.sac_learner import HoraSACLearner -from unilab.algos.torch.offpolicy.runtime import OffPolicyRuntime +from unilab.algos.hora.runtime import HORA_SAC_RUNTIME_IMPL, is_hora_sac_runtime +from unilab.algos.hora.sac_learner import HoraSACLearner +from unilab.algos.offpolicy.runtime import OffPolicyRuntime @dataclass(frozen=True) diff --git a/src/unilab/algos/torch/hora/sac_learner.py b/src/unilab/algos/hora/sac_learner.py similarity index 97% rename from src/unilab/algos/torch/hora/sac_learner.py rename to src/unilab/algos/hora/sac_learner.py index 029783ec7..ce90021fe 100644 --- a/src/unilab/algos/torch/hora/sac_learner.py +++ b/src/unilab/algos/hora/sac_learner.py @@ -8,8 +8,8 @@ import torch import torch.optim as optim -from unilab.algos.torch.fast_sac.learner import FastSACLearner -from unilab.algos.torch.hora.sac_models import HoraSACActor +from unilab.algos.fast_sac.learner import FastSACLearner +from unilab.algos.hora.sac_models import HoraSACActor def derive_priv_info_from_critic_obs( diff --git a/src/unilab/algos/torch/hora/sac_models.py b/src/unilab/algos/hora/sac_models.py similarity index 100% rename from src/unilab/algos/torch/hora/sac_models.py rename to src/unilab/algos/hora/sac_models.py diff --git a/src/unilab/algos/torch/offpolicy/__init__.py b/src/unilab/algos/offpolicy/__init__.py similarity index 56% rename from src/unilab/algos/torch/offpolicy/__init__.py rename to src/unilab/algos/offpolicy/__init__.py index 60749cd92..d578222bd 100644 --- a/src/unilab/algos/torch/offpolicy/__init__.py +++ b/src/unilab/algos/offpolicy/__init__.py @@ -1,7 +1,7 @@ """Off-policy RL unified infrastructure.""" -from unilab.algos.torch.offpolicy.runner import OffPolicyRunner -from unilab.algos.torch.offpolicy.worker import off_policy_collector_fn +from unilab.algos.offpolicy.runner import OffPolicyRunner +from unilab.algos.offpolicy.worker import off_policy_collector_fn from unilab.logging import OffPolicyLogger __all__ = [ diff --git a/src/unilab/algos/torch/offpolicy/double_buffer_runner.py b/src/unilab/algos/offpolicy/double_buffer_runner.py similarity index 99% rename from src/unilab/algos/torch/offpolicy/double_buffer_runner.py rename to src/unilab/algos/offpolicy/double_buffer_runner.py index 114f9d927..851461da3 100644 --- a/src/unilab/algos/torch/offpolicy/double_buffer_runner.py +++ b/src/unilab/algos/offpolicy/double_buffer_runner.py @@ -16,17 +16,17 @@ if TYPE_CHECKING: from unilab.ipc.dp_sync import DpParameterSync -from unilab.algos.torch.offpolicy.runner import ( +from unilab.algos.offpolicy.runner import ( OffPolicyRunner, build_offpolicy_sample_info, build_reward_comparison_metrics, replay_buffer_ready_for_learning, ) -from unilab.algos.torch.offpolicy.thread_budget import ( +from unilab.algos.offpolicy.thread_budget import ( format_torch_thread_runtime, torch_thread_env, ) -from unilab.algos.torch.offpolicy.worker import off_policy_collector_fn, sample_offpolicy_actions +from unilab.algos.offpolicy.worker import off_policy_collector_fn, sample_offpolicy_actions from unilab.ipc.async_runner import _SPAWN_CTX from unilab.ipc.inference_slot import SharedInferenceSlot from unilab.ipc.replay_buffer import DEFAULT_REPLAY_INGRESS_DEPTH, ReplayBuffer diff --git a/src/unilab/algos/torch/offpolicy/runner.py b/src/unilab/algos/offpolicy/runner.py similarity index 99% rename from src/unilab/algos/torch/offpolicy/runner.py rename to src/unilab/algos/offpolicy/runner.py index 09fbed4bd..ec45d7c96 100644 --- a/src/unilab/algos/torch/offpolicy/runner.py +++ b/src/unilab/algos/offpolicy/runner.py @@ -6,7 +6,7 @@ from collections import deque from typing import Any -from unilab.algos.torch.common.device import get_env_dims +from unilab.algos.common.device import get_env_dims from unilab.ipc.async_runner import AsyncRunner from unilab.logging import OffPolicyLogger from unilab.training.seed import apply_training_seed diff --git a/src/unilab/algos/torch/offpolicy/runtime.py b/src/unilab/algos/offpolicy/runtime.py similarity index 100% rename from src/unilab/algos/torch/offpolicy/runtime.py rename to src/unilab/algos/offpolicy/runtime.py diff --git a/src/unilab/algos/torch/offpolicy/thread_budget.py b/src/unilab/algos/offpolicy/thread_budget.py similarity index 100% rename from src/unilab/algos/torch/offpolicy/thread_budget.py rename to src/unilab/algos/offpolicy/thread_budget.py diff --git a/src/unilab/algos/torch/offpolicy/worker.py b/src/unilab/algos/offpolicy/worker.py similarity index 98% rename from src/unilab/algos/torch/offpolicy/worker.py rename to src/unilab/algos/offpolicy/worker.py index f00fcc16b..9c176fa4a 100644 --- a/src/unilab/algos/torch/offpolicy/worker.py +++ b/src/unilab/algos/offpolicy/worker.py @@ -8,8 +8,8 @@ import numpy as np import torch -from unilab.algos.torch.common.collector_timing import extract_env_step_breakdown_timing_ms -from unilab.algos.torch.offpolicy.thread_budget import apply_torch_thread_runtime +from unilab.algos.common.collector_timing import extract_env_step_breakdown_timing_ms +from unilab.algos.offpolicy.thread_budget import apply_torch_thread_runtime from unilab.base.final_observation import resolve_terminal_observation_contract from unilab.base.observations import split_obs_dict from unilab.base.registry import ensure_registries @@ -64,7 +64,7 @@ def resolve_offpolicy_actor_priv_info( if algo_type != "hora_sac": return None - from unilab.algos.torch.hora.observations import split_hora_obs_with_priv_info + from unilab.algos.hora.observations import split_hora_obs_with_priv_info _, _, priv_info_np = split_hora_obs_with_priv_info( {"obs": obs_np, "critic": critic_np}, diff --git a/src/unilab/algos/torch/rsl_rl_ppo.py b/src/unilab/algos/rsl_rl_ppo.py similarity index 99% rename from src/unilab/algos/torch/rsl_rl_ppo.py rename to src/unilab/algos/rsl_rl_ppo.py index 67ca1b17d..5f8836db5 100644 --- a/src/unilab/algos/torch/rsl_rl_ppo.py +++ b/src/unilab/algos/rsl_rl_ppo.py @@ -7,7 +7,7 @@ from rsl_rl.algorithms import PPO from tensordict import TensorDict -from unilab.algos.torch.common.compile import get_torch_compile_for_cuda +from unilab.algos.common.compile import get_torch_compile_for_cuda _LOG_2_PI = math.log(2.0 * math.pi) _NORMAL_ENTROPY_OFFSET = 0.5 * (1.0 + _LOG_2_PI) diff --git a/src/unilab/algos/torch/rsl_rl_runtime.py b/src/unilab/algos/rsl_rl_runtime.py similarity index 100% rename from src/unilab/algos/torch/rsl_rl_runtime.py rename to src/unilab/algos/rsl_rl_runtime.py diff --git a/src/unilab/algos/torch/appo/__init__.py b/src/unilab/algos/torch/appo/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/src/unilab/algos/torch/common/__init__.py b/src/unilab/algos/torch/common/__init__.py deleted file mode 100644 index 75b190f0b..000000000 --- a/src/unilab/algos/torch/common/__init__.py +++ /dev/null @@ -1,18 +0,0 @@ -from unilab.algos.torch.common.actor_factory import build_actor -from unilab.algos.torch.common.device import get_env_dims -from unilab.algos.torch.common.networks import Critic, DistributionalQNetwork -from unilab.algos.torch.common.normalization import EmpiricalNormalization -from unilab.algos.torch.common.stability import check_nan_loss, clip_gradients, safe_tensor -from unilab.base.registry import ensure_registries - -__all__ = [ - "EmpiricalNormalization", - "DistributionalQNetwork", - "Critic", - "get_env_dims", - "check_nan_loss", - "clip_gradients", - "safe_tensor", - "ensure_registries", - "build_actor", -] diff --git a/src/unilab/algos/torch/flash_sac/__init__.py b/src/unilab/algos/torch/flash_sac/__init__.py deleted file mode 100644 index 2219b855b..000000000 --- a/src/unilab/algos/torch/flash_sac/__init__.py +++ /dev/null @@ -1,12 +0,0 @@ -"""FlashSAC algorithm package.""" - -from unilab.algos.torch.flash_sac.learner import FlashSACLearner -from unilab.algos.torch.flash_sac.network import FlashSACActor, FlashSACDoubleCritic -from unilab.algos.torch.flash_sac.runner import FlashSACRunner - -__all__ = [ - "FlashSACActor", - "FlashSACDoubleCritic", - "FlashSACLearner", - "FlashSACRunner", -] diff --git a/src/unilab/algos/torch/him_ppo/__init__.py b/src/unilab/algos/torch/him_ppo/__init__.py deleted file mode 100644 index 225f04629..000000000 --- a/src/unilab/algos/torch/him_ppo/__init__.py +++ /dev/null @@ -1,11 +0,0 @@ -from unilab.algos.torch.him_ppo.actor_critic import HIMActorCritic -from unilab.algos.torch.him_ppo.algorithm import HIMPPO -from unilab.algos.torch.him_ppo.estimator import HIMEstimator -from unilab.algos.torch.him_ppo.storage import HIMRolloutStorage - -__all__ = [ - "HIMActorCritic", - "HIMPPO", - "HIMEstimator", - "HIMRolloutStorage", -] diff --git a/src/unilab/structured_configs.py b/src/unilab/structured_configs.py index 48e59a805..4b25f1e0e 100644 --- a/src/unilab/structured_configs.py +++ b/src/unilab/structured_configs.py @@ -238,7 +238,7 @@ class PPOPolicyConfig: @dataclass class PPOAlgorithmConfig: - class_name: str = "unilab.algos.torch.rsl_rl_ppo:FinalObservationAwarePPO" + class_name: str = "unilab.algos.rsl_rl_ppo:FinalObservationAwarePPO" value_loss_coef: float = 1.0 use_clipped_value_loss: bool = True clip_param: float = 0.2 diff --git a/src/unilab/training/offpolicy.py b/src/unilab/training/offpolicy.py index 33854813e..36e0ee0b3 100644 --- a/src/unilab/training/offpolicy.py +++ b/src/unilab/training/offpolicy.py @@ -64,7 +64,7 @@ def resolve_play_actor_spec( if algo_name != "sac": return algo_name, {} - from unilab.algos.torch.offpolicy.runtime import resolve_custom_offpolicy_runtime + from unilab.algos.offpolicy.runtime import resolve_custom_offpolicy_runtime rl_cfg = cast(dict[str, Any], OmegaConf.to_container(cfg.algo, resolve=True)) custom_runtime = resolve_custom_offpolicy_runtime(rl_cfg) @@ -91,7 +91,7 @@ def build_play_actor( """Build the policy actor selected by an off-policy owner config.""" import torch - from unilab.algos.torch.common.actor_factory import build_actor + from unilab.algos.common.actor_factory import build_actor actor_algo_type, actor_kwargs = resolve_play_actor_spec( algo_name, @@ -111,7 +111,7 @@ def build_play_actor( **actor_kwargs, ) elif algo_name == "td3": - from unilab.algos.torch.fast_td3.learner import EmpiricalNormalization, TD3Actor + from unilab.algos.fast_td3.learner import EmpiricalNormalization, TD3Actor actor = TD3Actor( obs_dim, @@ -138,7 +138,7 @@ def build_play_actor( actor_noise_zeta_max=cfg.algo.algo_params.actor_noise_zeta_max, ) if cfg.algo.obs_normalization: - from unilab.algos.torch.common.normalization import EmpiricalNormalization + from unilab.algos.common.normalization import EmpiricalNormalization normalizer = EmpiricalNormalization(shape=obs_dim, device=device) else: diff --git a/src/unilab/visualization/interactive_playback.py b/src/unilab/visualization/interactive_playback.py index 5c319fd68..b76b8d7c0 100644 --- a/src/unilab/visualization/interactive_playback.py +++ b/src/unilab/visualization/interactive_playback.py @@ -536,9 +536,9 @@ def _build_appo_actor( rl_cfg_dict = deepcopy(rl_cfg) if is_hora: - from unilab.algos.torch.hora.appo import _update_hora_obs_groups - from unilab.algos.torch.hora.models import build_hora_shared_actor_critic - from unilab.algos.torch.hora.rsl_rl_compat import ( + from unilab.algos.hora.appo import _update_hora_obs_groups + from unilab.algos.hora.models import build_hora_shared_actor_critic + from unilab.algos.hora.rsl_rl_compat import ( convert_config_v3_to_v4, is_rsl_rl_v4, is_rsl_rl_v5, @@ -635,13 +635,13 @@ def create_appo_playback_session( if env is None: raise RuntimeError("Playback env factory did not return an environment.") - from unilab.algos.torch.hora.runtime import is_hora_appo_runtime + from unilab.algos.hora.runtime import is_hora_appo_runtime is_hora = is_hora_appo_runtime(rl_cfg) selected_wrapper_cls = wrapper_cls policy_obs_mode = playback_cfg.policy_obs_mode if is_hora: - from unilab.algos.torch.hora.rsl_rl import HoraRslRlVecEnvWrapper + from unilab.algos.hora.rsl_rl import HoraRslRlVecEnvWrapper selected_wrapper_cls = HoraRslRlVecEnvWrapper policy_obs_mode = "actor" @@ -710,8 +710,8 @@ def create_sac_playback_session( import os - from unilab.algos.torch.common.actor_factory import build_actor - from unilab.algos.torch.offpolicy.worker import resolve_offpolicy_actor_priv_info + from unilab.algos.common.actor_factory import build_actor + from unilab.algos.offpolicy.worker import resolve_offpolicy_actor_priv_info from unilab.training.offpolicy import ( default_device, extract_play_obs, @@ -749,7 +749,7 @@ def create_sac_playback_session( checkpoint_path: str | None = None normalizer = None if bool(getattr(cfg.algo, "obs_normalization", False)): - from unilab.algos.torch.common.normalization import EmpiricalNormalization + from unilab.algos.common.normalization import EmpiricalNormalization normalizer = EmpiricalNormalization(shape=obs_dim, device=device_name) if playback_cfg.action_mode == "policy": @@ -813,14 +813,14 @@ def create_sac_playback_session( def _default_hora_distill_playback_deps(root_dir: str | Path) -> dict[str, Any]: - from unilab.algos.torch.hora.distill import ( + from unilab.algos.hora.distill import ( build_student_actor_and_normalizer, cfg_with_checkpoint_runtime, load_distilled_checkpoint, student_policy, ) - from unilab.algos.torch.hora.distill_config import apply_teacher_defaults - from unilab.algos.torch.hora.rsl_rl import HoraRslRlVecEnvWrapper + from unilab.algos.hora.distill_config import apply_teacher_defaults + from unilab.algos.hora.rsl_rl import HoraRslRlVecEnvWrapper from unilab.base.backend import materialize_scene_visual_override from unilab.training import ( BackendAdapter, diff --git a/tests/algos/test_appo_learner.py b/tests/algos/test_appo_learner.py index 1e481aa08..b1311d89f 100644 --- a/tests/algos/test_appo_learner.py +++ b/tests/algos/test_appo_learner.py @@ -4,8 +4,8 @@ import torch -from unilab.algos.torch.appo.learner import APPOLearner -from unilab.algos.torch.hora.appo_learner import HoraAPPOLearner +from unilab.algos.appo.learner import APPOLearner +from unilab.algos.hora.appo_learner import HoraAPPOLearner def test_appo_learner_compile_targets_minibatch_loss(monkeypatch) -> None: diff --git a/tests/algos/test_appo_learner_metrics.py b/tests/algos/test_appo_learner_metrics.py index c1a37a4e7..177c7a196 100644 --- a/tests/algos/test_appo_learner_metrics.py +++ b/tests/algos/test_appo_learner_metrics.py @@ -4,7 +4,7 @@ import torch import torch.nn as nn -from unilab.algos.torch.appo.learner import APPOLearner +from unilab.algos.appo.learner import APPOLearner class _Actor(nn.Module): diff --git a/tests/algos/test_appo_runner.py b/tests/algos/test_appo_runner.py index 4c3ceb4e1..17813ec33 100644 --- a/tests/algos/test_appo_runner.py +++ b/tests/algos/test_appo_runner.py @@ -15,7 +15,7 @@ pytest.importorskip("mujoco") -from unilab.algos.torch.appo.runner import APPORunner +from unilab.algos.appo.runner import APPORunner from unilab.structured_configs import APPOConfig from unilab.training.backend_adapter import BackendAdapter diff --git a/tests/algos/test_appo_runner_unit.py b/tests/algos/test_appo_runner_unit.py index b6df4c55c..c4b895ad6 100644 --- a/tests/algos/test_appo_runner_unit.py +++ b/tests/algos/test_appo_runner_unit.py @@ -6,8 +6,8 @@ import pytest import torch -import unilab.algos.torch.appo.runner as appo_runner_module -from unilab.algos.torch.appo.runner import APPORunner +import unilab.algos.appo.runner as appo_runner_module +from unilab.algos.appo.runner import APPORunner @pytest.fixture(autouse=True) diff --git a/tests/algos/test_appo_staging.py b/tests/algos/test_appo_staging.py index ab0f34c8f..1fdb095a5 100644 --- a/tests/algos/test_appo_staging.py +++ b/tests/algos/test_appo_staging.py @@ -4,7 +4,7 @@ import pytest import torch -from unilab.algos.torch.appo.staging import RolloutStagingPool +from unilab.algos.appo.staging import RolloutStagingPool _NUM_ENVS = 2 _NUM_STEPS = 3 diff --git a/tests/algos/test_appo_worker.py b/tests/algos/test_appo_worker.py index 3a64b9363..6218d58f3 100644 --- a/tests/algos/test_appo_worker.py +++ b/tests/algos/test_appo_worker.py @@ -4,7 +4,7 @@ import numpy as np -from unilab.algos.torch.appo.worker import ( +from unilab.algos.appo.worker import ( compute_rollout_active_steps_per_sec, compute_timeout_bootstrap_correction, put_latest_metrics, diff --git a/tests/algos/test_fast_sac_compile.py b/tests/algos/test_fast_sac_compile.py index 298a6ff85..8bac88556 100644 --- a/tests/algos/test_fast_sac_compile.py +++ b/tests/algos/test_fast_sac_compile.py @@ -6,7 +6,7 @@ import pytest import torch -from unilab.algos.torch.fast_sac.learner import FastSACLearner, SACActor +from unilab.algos.fast_sac.learner import FastSACLearner, SACActor def _small_fast_sac_learner(*, use_autotune: bool = True) -> FastSACLearner: diff --git a/tests/algos/test_fast_sac_symmetry_contract.py b/tests/algos/test_fast_sac_symmetry_contract.py index f02eb7848..5143a360d 100644 --- a/tests/algos/test_fast_sac_symmetry_contract.py +++ b/tests/algos/test_fast_sac_symmetry_contract.py @@ -8,7 +8,7 @@ import pytest import torch -import unilab.algos.torch.fast_sac.learner as learner_module +import unilab.algos.fast_sac.learner as learner_module class _FakeSymmetryAugmentation: @@ -62,9 +62,9 @@ def test_fast_sac_runner_uses_env_owned_symmetry_contract( monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, ): - import unilab.algos.torch.fast_sac.runner as runner_module - import unilab.algos.torch.offpolicy.double_buffer_runner as device_runner_module - from unilab.algos.torch.fast_sac.runner import FastSACRunner + import unilab.algos.fast_sac.runner as runner_module + import unilab.algos.offpolicy.double_buffer_runner as device_runner_module + from unilab.algos.fast_sac.runner import FastSACRunner from unilab.base import registry augmentation = _FakeSymmetryAugmentation() @@ -79,7 +79,7 @@ def test_fast_sac_runner_uses_env_owned_symmetry_contract( lambda value: value, ) - with caplog.at_level(logging.INFO, logger="unilab.algos.torch.fast_sac.runner"): + with caplog.at_level(logging.INFO, logger="unilab.algos.fast_sac.runner"): runner = FastSACRunner( env_name="FakeEnv", device="cpu", @@ -101,9 +101,9 @@ def test_fast_sac_runner_uses_env_owned_symmetry_contract( def test_fast_sac_runner_skips_symmetry_builder_when_disabled(monkeypatch: pytest.MonkeyPatch): - import unilab.algos.torch.fast_sac.runner as runner_module - import unilab.algos.torch.offpolicy.double_buffer_runner as device_runner_module - from unilab.algos.torch.fast_sac.runner import FastSACRunner + import unilab.algos.fast_sac.runner as runner_module + import unilab.algos.offpolicy.double_buffer_runner as device_runner_module + from unilab.algos.fast_sac.runner import FastSACRunner from unilab.base import registry fake_env = _FakeEnv(_FakeSymmetryAugmentation()) @@ -141,7 +141,7 @@ def _unexpected_builder(*args, **kwargs): def test_fast_sac_learner_rejects_symmetry_without_augmentation(): - from unilab.algos.torch.fast_sac.learner import FastSACLearner + from unilab.algos.fast_sac.learner import FastSACLearner with pytest.raises( ValueError, @@ -157,7 +157,7 @@ def test_fast_sac_learner_rejects_symmetry_without_augmentation(): def test_fast_sac_obs_normalization_uses_local_moments_and_round_trips(): - from unilab.algos.torch.fast_sac.learner import FastSACLearner + from unilab.algos.fast_sac.learner import FastSACLearner learner = FastSACLearner( obs_dim=2, @@ -199,7 +199,7 @@ def test_fast_sac_obs_normalization_uses_local_moments_and_round_trips(): def test_fast_sac_symmetry_augmentation_emits_fine_grained_nvtx_ranges( monkeypatch: pytest.MonkeyPatch, ) -> None: - from unilab.algos.torch.fast_sac.learner import FastSACLearner + from unilab.algos.fast_sac.learner import FastSACLearner seen_ranges: list[str] = [] @@ -275,7 +275,7 @@ def fake_actor_loss_tensors(*args, **kwargs): def test_fast_sac_symmetry_uses_obs_only_augmentation_for_obs_only_paths( monkeypatch: pytest.MonkeyPatch, ) -> None: - from unilab.algos.torch.fast_sac.learner import FastSACLearner + from unilab.algos.fast_sac.learner import FastSACLearner symmetry = _FakeSymmetryAugmentation() learner = FastSACLearner( diff --git a/tests/algos/test_fast_td3_learner.py b/tests/algos/test_fast_td3_learner.py index 9d0423431..993f17b41 100644 --- a/tests/algos/test_fast_td3_learner.py +++ b/tests/algos/test_fast_td3_learner.py @@ -5,7 +5,7 @@ import pytest import torch -from unilab.algos.torch.fast_td3.learner import FastTD3Learner, TD3Actor +from unilab.algos.fast_td3.learner import FastTD3Learner, TD3Actor # --------------------------------------------------------------------------- # TD3Actor diff --git a/tests/algos/test_flash_sac_learner.py b/tests/algos/test_flash_sac_learner.py index 422a8091d..58ddbf7d5 100644 --- a/tests/algos/test_flash_sac_learner.py +++ b/tests/algos/test_flash_sac_learner.py @@ -8,8 +8,8 @@ import pytest import torch -from unilab.algos.torch.flash_sac.learner import FlashSACLearner, RewardNormalizer -from unilab.algos.torch.flash_sac.update import compute_categorical_td_target +from unilab.algos.flash_sac.learner import FlashSACLearner, RewardNormalizer +from unilab.algos.flash_sac.update import compute_categorical_td_target def _make_batch(batch_size: int = 32) -> dict[str, torch.Tensor]: diff --git a/tests/algos/test_him_ppo_runner.py b/tests/algos/test_him_ppo_runner.py index f64a23bb9..e036a164b 100644 --- a/tests/algos/test_him_ppo_runner.py +++ b/tests/algos/test_him_ppo_runner.py @@ -11,7 +11,7 @@ def test_him_iteration_progress_is_one_equivalent_multiline_log_record( caplog: pytest.LogCaptureFixture, ) -> None: - from unilab.algos.torch.him_ppo.runner import HIMOnPolicyRunner + from unilab.algos.him_ppo.runner import HIMOnPolicyRunner runner = cast(Any, HIMOnPolicyRunner.__new__(HIMOnPolicyRunner)) runner.logger = SimpleNamespace( @@ -19,7 +19,7 @@ def test_him_iteration_progress_is_one_equivalent_multiline_log_record( lenbuffer=deque([10.0, 14.0]), ) - with caplog.at_level(logging.INFO, logger="unilab.algos.torch.him_ppo.runner"): + with caplog.at_level(logging.INFO, logger="unilab.algos.him_ppo.runner"): runner._log_iter( it=1, tot=3, @@ -31,9 +31,7 @@ def test_him_iteration_progress_is_one_equivalent_multiline_log_record( infos={"log": {"reward/feet": 1.25}}, ) - records = [ - record for record in caplog.records if record.name == "unilab.algos.torch.him_ppo.runner" - ] + records = [record for record in caplog.records if record.name == "unilab.algos.him_ppo.runner"] assert len(records) == 1 lines = records[0].getMessage().splitlines() assert lines[0] == "-" * 80 diff --git a/tests/algos/test_hora_contract.py b/tests/algos/test_hora_contract.py index 8b794bc0b..1b8bbc6a3 100644 --- a/tests/algos/test_hora_contract.py +++ b/tests/algos/test_hora_contract.py @@ -15,7 +15,7 @@ def test_hora_ppo_logs_when_symmetry_is_logging_only( caplog: pytest.LogCaptureFixture, ) -> None: - from unilab.algos.torch.hora.ppo import HoraPPO + from unilab.algos.hora.ppo import HoraPPO actor = torch.nn.Linear(2, 2) critic = torch.nn.Linear(2, 1) @@ -27,7 +27,7 @@ def test_hora_ppo_logs_when_symmetry_is_logging_only( "data_augmentation_func": lambda *_args: None, } - with caplog.at_level(logging.WARNING, logger="unilab.algos.torch.hora.ppo"): + with caplog.at_level(logging.WARNING, logger="unilab.algos.hora.ppo"): HoraPPO( cast(Any, actor), cast(Any, critic), @@ -39,7 +39,7 @@ def test_hora_ppo_logs_when_symmetry_is_logging_only( def test_hora_sac_actor_shapes_and_stable_module_names() -> None: - from unilab.algos.torch.hora.sac_models import HoraSACActor + from unilab.algos.hora.sac_models import HoraSACActor actor = HoraSACActor( obs_dim=5, @@ -65,7 +65,7 @@ def test_hora_sac_actor_shapes_and_stable_module_names() -> None: def test_hora_sac_learner_derives_priv_info_from_critic_contract() -> None: - from unilab.algos.torch.hora.sac_learner import derive_priv_info_from_critic_obs + from unilab.algos.hora.sac_learner import derive_priv_info_from_critic_obs actor_obs = torch.zeros((4, 5), dtype=torch.float32) priv_info = torch.arange(12, dtype=torch.float32).reshape(4, 3) @@ -81,7 +81,7 @@ def test_hora_sac_learner_derives_priv_info_from_critic_contract() -> None: def test_hora_sac_learner_updates_with_privileged_tail() -> None: - from unilab.algos.torch.hora.sac_learner import HoraSACLearner + from unilab.algos.hora.sac_learner import HoraSACLearner torch.manual_seed(23) learner = HoraSACLearner( @@ -123,7 +123,7 @@ def test_hora_sac_learner_updates_with_privileged_tail() -> None: def test_hora_sac_disables_cuda_graph_critic_path() -> None: - from unilab.algos.torch.hora.sac_learner import HoraSACLearner + from unilab.algos.hora.sac_learner import HoraSACLearner learner = HoraSACLearner( obs_dim=5, @@ -144,7 +144,7 @@ def test_hora_sac_disables_cuda_graph_critic_path() -> None: def test_hora_sac_distilled_student_forward_does_not_require_priv_info() -> None: - from unilab.algos.torch.hora.distill import HoraSACDistillActor, HoraSACDistillShared + from unilab.algos.hora.distill import HoraSACDistillActor, HoraSACDistillShared shared = HoraSACDistillShared( obs_dim=12, @@ -175,12 +175,12 @@ def test_hora_sac_distilled_student_forward_does_not_require_priv_info() -> None def test_hora_sac_distill_loads_teacher_actor_weights(tmp_path) -> None: - from unilab.algos.torch.hora.distill import ( + from unilab.algos.hora.distill import ( HoraSACDistillActor, HoraSACDistillShared, load_teacher_actor_weights, ) - from unilab.algos.torch.hora.sac_models import HoraSACActor + from unilab.algos.hora.sac_models import HoraSACActor teacher = HoraSACActor( obs_dim=12, @@ -226,7 +226,7 @@ def test_hora_sac_distill_loads_teacher_actor_weights(tmp_path) -> None: def test_hora_rsl_wrapper_uses_explicit_np_env_state_contract() -> None: """HORA wrapper must not probe required NpEnvState fields dynamically.""" - from unilab.algos.torch.hora.rsl_rl import HoraRslRlVecEnvWrapper + from unilab.algos.hora.rsl_rl import HoraRslRlVecEnvWrapper source = textwrap.dedent(inspect.getsource(HoraRslRlVecEnvWrapper.step)) tree = ast.parse(source) @@ -245,7 +245,7 @@ def test_hora_rsl_wrapper_uses_explicit_np_env_state_contract() -> None: def test_hora_appo_learner_derives_priv_info_from_critic_contract() -> None: - from unilab.algos.torch.hora.appo_learner import _derive_priv_info_from_critic + from unilab.algos.hora.appo_learner import _derive_priv_info_from_critic actor_obs = torch.zeros((2, 3, 4), dtype=torch.float32) priv_info = torch.arange(12, dtype=torch.float32).reshape(2, 3, 2) @@ -261,8 +261,8 @@ def test_hora_appo_learner_derives_priv_info_from_critic_contract() -> None: def _make_hora_appo_learner(**algorithm_overrides): - from unilab.algos.torch.hora.appo_learner import HoraAPPOLearner - from unilab.algos.torch.hora.models import ( + from unilab.algos.hora.appo_learner import HoraAPPOLearner + from unilab.algos.hora.models import ( HoraActorModel, HoraCriticModel, HoraSharedActorCritic, @@ -326,7 +326,7 @@ def test_hora_appo_minibatch_tensor_path_matches_tensordict_forward() -> None: def test_hora_appo_runner_builds_shared_actor_critic_core() -> None: - from unilab.algos.torch.hora.appo_runner import HoraAPPORunner + from unilab.algos.hora.appo_runner import HoraAPPORunner runner = HoraAPPORunner.__new__(HoraAPPORunner) runner.num_envs = 4 @@ -341,13 +341,13 @@ def test_hora_appo_runner_builds_shared_actor_critic_core() -> None: "critic": {"actor": 5, "priv_info": 2}, }, "actor": { - "class_name": "unilab.algos.torch.hora:HoraActorModel", + "class_name": "unilab.algos.hora:HoraActorModel", "hidden_dims": [8], "priv_info_embed_dim": 2, "priv_mlp_hidden_dims": [4, 2], }, "critic": { - "class_name": "unilab.algos.torch.hora:HoraCriticModel", + "class_name": "unilab.algos.hora:HoraCriticModel", "priv_info_embed_dim": 2, "priv_mlp_hidden_dims": [4, 2], }, @@ -363,7 +363,7 @@ def test_hora_appo_runner_builds_shared_actor_critic_core() -> None: def test_hora_appo_worker_builds_shared_actor_critic_core() -> None: - from unilab.algos.torch.hora.appo_worker import hora_appo_collector_fn + from unilab.algos.hora.appo_worker import hora_appo_collector_fn source = textwrap.dedent(inspect.getsource(hora_appo_collector_fn)) tree = ast.parse(source) @@ -381,7 +381,7 @@ def test_hora_appo_worker_builds_shared_actor_critic_core() -> None: def test_hora_appo_play_builds_explicit_shared_actor_core() -> None: - from unilab.algos.torch.hora.appo import play_hora_appo + from unilab.algos.hora.appo import play_hora_appo source = textwrap.dedent(inspect.getsource(play_hora_appo)) tree = ast.parse(source) @@ -405,7 +405,7 @@ def _patch_hora_appo_play_fakes(monkeypatch: pytest.MonkeyPatch, *, actor_cls: t import numpy as np import rsl_rl.utils as rsl_rl_utils - import unilab.algos.torch.hora.appo as hora_appo + import unilab.algos.hora.appo as hora_appo fake_env = SimpleNamespace( obs_groups_spec={"obs": 3, "critic": 5}, @@ -466,7 +466,7 @@ def test_hora_appo_play_runs_sim2sim_preflight_before_checkpoint_load( monkeypatch: pytest.MonkeyPatch, tmp_path, ) -> None: - import unilab.algos.torch.hora.appo as hora_appo + import unilab.algos.hora.appo as hora_appo checkpoint = tmp_path / "model_10.pt" torch.save({"actor": {"weight": torch.tensor(1.0)}}, checkpoint) @@ -506,7 +506,7 @@ def test_hora_appo_play_dim_mismatch_reraises_explicit_sim2sim_diagnostic( tmp_path, capsys, ) -> None: - import unilab.algos.torch.hora.appo as hora_appo + import unilab.algos.hora.appo as hora_appo from unilab.training.sim2sim import CrossBackendIncompatibleError checkpoint = tmp_path / "model_10.pt" @@ -534,7 +534,7 @@ def load_state_dict(self, state_dict, strict=True): def test_hora_appo_resume_rejects_inconsistent_shared_checkpoint() -> None: - from unilab.algos.torch.hora.appo_runner import _validate_hora_shared_checkpoint + from unilab.algos.hora.appo_runner import _validate_hora_shared_checkpoint learner = _make_hora_appo_learner() joint_checkpoint = { diff --git a/tests/algos/test_hora_distill_config.py b/tests/algos/test_hora_distill_config.py index 457e7a2bf..f8f57c86e 100644 --- a/tests/algos/test_hora_distill_config.py +++ b/tests/algos/test_hora_distill_config.py @@ -17,7 +17,7 @@ from omegaconf import OmegaConf from omegaconf.errors import InterpolationResolutionError -from unilab.algos.torch.hora import distill_config +from unilab.algos.hora import distill_config _REPO_ROOT = Path(__file__).resolve().parents[2] @@ -182,7 +182,7 @@ def test_hora_actor_mapping_strips_distribution_class_name() -> None: { "algo": { "actor": { - "class_name": "unilab.algos.torch.hora:HoraActorModel", + "class_name": "unilab.algos.hora:HoraActorModel", "hidden_dims": [64, 32], "activation": "relu", "obs_normalization": False, @@ -215,7 +215,7 @@ def test_hora_actor_mapping_fails_closed_when_teacher_field_is_missing() -> None { "algo": { "actor": { - "class_name": "unilab.algos.torch.hora:HoraActorModel", + "class_name": "unilab.algos.hora:HoraActorModel", "activation": "elu", "obs_normalization": True, "priv_info_embed_dim": 9, diff --git a/tests/algos/test_hora_imports.py b/tests/algos/test_hora_imports.py index f5cfcadee..045680f6d 100644 --- a/tests/algos/test_hora_imports.py +++ b/tests/algos/test_hora_imports.py @@ -8,31 +8,31 @@ def test_hora_package_import_keeps_appo_lazy() -> None: - sys.modules.pop("unilab.algos.torch.hora", None) - sys.modules.pop("unilab.algos.torch.hora.appo", None) + sys.modules.pop("unilab.algos.hora", None) + sys.modules.pop("unilab.algos.hora.appo", None) - importlib.import_module("unilab.algos.torch.hora") + importlib.import_module("unilab.algos.hora") - assert "unilab.algos.torch.hora.appo" not in sys.modules + assert "unilab.algos.hora.appo" not in sys.modules def test_resolve_callable_loads_hora_ppo_from_package_export() -> None: - resolved = resolve_callable("unilab.algos.torch.hora:HoraPPO") + resolved = resolve_callable("unilab.algos.hora:HoraPPO") - from unilab.algos.torch.hora.ppo import HoraPPO + from unilab.algos.hora.ppo import HoraPPO assert resolved is HoraPPO def test_rsl_rl_runtime_resolver_loads_hora_wrapper_from_owner_marker() -> None: - from unilab.algos.torch.hora.rsl_rl import HoraRslRlVecEnvWrapper - from unilab.algos.torch.rsl_rl_runtime import resolve_rsl_rl_ppo_runtime + from unilab.algos.hora.rsl_rl import HoraRslRlVecEnvWrapper + from unilab.algos.rsl_rl_runtime import resolve_rsl_rl_ppo_runtime from unilab.training.rsl_rl import RslRlVecEnvWrapper runtime = resolve_rsl_rl_ppo_runtime( { "runtime_impl": "hora_ppo", - "runtime_resolver": "unilab.algos.torch.hora.rsl_rl:resolve_hora_ppo_runtime", + "runtime_resolver": "unilab.algos.hora.rsl_rl:resolve_hora_ppo_runtime", }, default_wrapper_cls=RslRlVecEnvWrapper, ) @@ -41,7 +41,7 @@ def test_rsl_rl_runtime_resolver_loads_hora_wrapper_from_owner_marker() -> None: def test_rsl_rl_runtime_resolver_rejects_unresolved_custom_runtime() -> None: - from unilab.algos.torch.rsl_rl_runtime import resolve_rsl_rl_ppo_runtime + from unilab.algos.rsl_rl_runtime import resolve_rsl_rl_ppo_runtime from unilab.training.rsl_rl import RslRlVecEnvWrapper with pytest.raises(ValueError, match="runtime_impl='hora_ppo'.*runtime_resolver"): diff --git a/tests/algos/test_offpolicy_bootstrap_contract.py b/tests/algos/test_offpolicy_bootstrap_contract.py index c6a8c0ea2..98145cd67 100644 --- a/tests/algos/test_offpolicy_bootstrap_contract.py +++ b/tests/algos/test_offpolicy_bootstrap_contract.py @@ -2,8 +2,8 @@ import torch -from unilab.algos.torch.fast_sac.learner import FastSACLearner -from unilab.algos.torch.fast_td3.learner import FastTD3Learner +from unilab.algos.fast_sac.learner import FastSACLearner +from unilab.algos.fast_td3.learner import FastTD3Learner class _CaptureSacTargetCritic(torch.nn.Module): diff --git a/tests/algos/test_offpolicy_double_buffer_runner.py b/tests/algos/test_offpolicy_double_buffer_runner.py index 9998a33a9..c28334025 100644 --- a/tests/algos/test_offpolicy_double_buffer_runner.py +++ b/tests/algos/test_offpolicy_double_buffer_runner.py @@ -147,15 +147,15 @@ def reject_env(*args, **kwargs): raise AssertionError("unsupported replay device must fail before env creation") if algo == "sac": - import unilab.algos.torch.fast_sac.double_buffer as owner_module + import unilab.algos.fast_sac.double_buffer as owner_module monkeypatch.setattr(owner_module, "create_env", reject_env) elif algo == "td3": - import unilab.algos.torch.fast_td3.double_buffer as owner_module + import unilab.algos.fast_td3.double_buffer as owner_module monkeypatch.setattr(owner_module, "get_env_dims", reject_env) else: - import unilab.algos.torch.flash_sac.double_buffer as owner_module + import unilab.algos.flash_sac.double_buffer as owner_module monkeypatch.setattr(owner_module, "create_env", reject_env) with pytest.raises(ValueError, match="training.devices entries"): @@ -167,7 +167,7 @@ def test_sac_dispatch_constructs_unique_runner(monkeypatch: pytest.MonkeyPatch): module = _offpolicy() cfg = _offpolicy_cfg(["algo=sac", "algo.use_symmetry=false"]) - import unilab.algos.torch.fast_sac.double_buffer as owner_module + import unilab.algos.fast_sac.double_buffer as owner_module monkeypatch.setattr(owner_module, "ensure_registries", lambda: None) monkeypatch.setattr(owner_module, "create_env", lambda *args, **kwargs: _FakeEnv()) @@ -222,8 +222,8 @@ def test_sac_dispatch_constructs_unique_runner(monkeypatch: pytest.MonkeyPatch): def test_sac_owner_custom_runtime_can_override_base_learner_kwargs( monkeypatch: pytest.MonkeyPatch, ): - from unilab.algos.torch.fast_sac import double_buffer as owner_module - from unilab.algos.torch.offpolicy.runtime import OffPolicyRuntime + from unilab.algos.fast_sac import double_buffer as owner_module + from unilab.algos.offpolicy.runtime import OffPolicyRuntime cfg = _offpolicy_cfg(["algo=sac", "algo.use_symmetry=false"]) custom_runtime = OffPolicyRuntime( @@ -256,8 +256,8 @@ def test_sac_owner_custom_runtime_can_override_base_learner_kwargs( def test_sac_owner_rejects_custom_runtime_without_symmetry_support( monkeypatch: pytest.MonkeyPatch, ): - from unilab.algos.torch.fast_sac import double_buffer as owner_module - from unilab.algos.torch.offpolicy.runtime import OffPolicyRuntime + from unilab.algos.fast_sac import double_buffer as owner_module + from unilab.algos.offpolicy.runtime import OffPolicyRuntime cfg = _offpolicy_cfg(["algo=sac", "algo.use_symmetry=true"]) monkeypatch.setattr(owner_module, "ensure_registries", lambda: None) @@ -280,7 +280,7 @@ def test_sac_owner_rejects_custom_runtime_without_symmetry_support( def test_sac_owner_preserves_symmetry_batch_and_learner_contract( monkeypatch: pytest.MonkeyPatch, ): - import unilab.algos.torch.fast_sac.double_buffer as owner_module + import unilab.algos.fast_sac.double_buffer as owner_module cfg = _offpolicy_cfg(["algo=sac", "algo.use_symmetry=true"]) symmetry = MagicMock(batch_multiplier=4) @@ -321,7 +321,7 @@ def test_sac_owner_preserves_symmetry_validation( symmetry: MagicMock | None, match: str, ): - import unilab.algos.torch.fast_sac.double_buffer as owner_module + import unilab.algos.fast_sac.double_buffer as owner_module cfg = _offpolicy_cfg(["algo=sac", "algo.use_symmetry=true", f"algo.batch_size={batch_size}"]) @@ -346,7 +346,7 @@ def test_td3_dispatch_constructs_unique_runner(monkeypatch: pytest.MonkeyPatch): module = _offpolicy() cfg = _offpolicy_cfg(["algo=td3"]) - import unilab.algos.torch.fast_td3.double_buffer as owner_module + import unilab.algos.fast_td3.double_buffer as owner_module monkeypatch.setattr(owner_module, "get_env_dims", lambda *args, **kwargs: (4, 2, 6)) monkeypatch.setattr(owner_module, "FastTD3Learner", _FakeLearner) @@ -418,7 +418,7 @@ def test_flashsac_dispatch_constructs_unique_runner(monkeypatch: pytest.MonkeyPa module = _offpolicy() cfg = _offpolicy_cfg(["algo=flashsac"]) - import unilab.algos.torch.flash_sac.double_buffer as flash_module + import unilab.algos.flash_sac.double_buffer as flash_module monkeypatch.setattr(flash_module, "ensure_registries", lambda: None) monkeypatch.setattr(flash_module, "create_env", lambda *args, **kwargs: _FakeEnv()) @@ -443,7 +443,7 @@ def test_flashsac_n_step_is_rejected(): def _bare_runner(): - from unilab.algos.torch.offpolicy.double_buffer_runner import DoubleBufferOffPolicyRunner + from unilab.algos.offpolicy.double_buffer_runner import DoubleBufferOffPolicyRunner return object.__new__(DoubleBufferOffPolicyRunner) @@ -476,7 +476,7 @@ def fake_create_env(*args, **kwargs): monkeypatch.setattr(module.os, "cpu_count", lambda: cpu_count) - import unilab.algos.torch.fast_sac.double_buffer as owner_module + import unilab.algos.fast_sac.double_buffer as owner_module monkeypatch.setattr(owner_module, "ensure_registries", lambda: None) monkeypatch.setattr(owner_module, "create_env", fake_create_env) diff --git a/tests/algos/test_offpolicy_dp_sync.py b/tests/algos/test_offpolicy_dp_sync.py index b5266af47..5021ee671 100644 --- a/tests/algos/test_offpolicy_dp_sync.py +++ b/tests/algos/test_offpolicy_dp_sync.py @@ -19,7 +19,7 @@ def _bare_runner(): - from unilab.algos.torch.offpolicy.double_buffer_runner import DoubleBufferOffPolicyRunner + from unilab.algos.offpolicy.double_buffer_runner import DoubleBufferOffPolicyRunner return object.__new__(DoubleBufferOffPolicyRunner) @@ -162,7 +162,7 @@ def test_close_closes_dp_sync_idempotently(): def test_close_restores_terminal_and_ipc_before_destroying_process_group(monkeypatch): - from unilab.algos.torch.offpolicy.runner import OffPolicyRunner + from unilab.algos.offpolicy.runner import OffPolicyRunner events: list[str] = [] @@ -185,7 +185,7 @@ def close(self) -> None: def test_close_still_destroys_process_group_when_local_cleanup_fails(monkeypatch): - from unilab.algos.torch.offpolicy.runner import OffPolicyRunner + from unilab.algos.offpolicy.runner import OffPolicyRunner dp_sync = _FakeDpSync() runner = _runner_with(_SyncLearner(), dp_sync) @@ -335,7 +335,7 @@ def log_save(self, path: str) -> None: def test_learn_source_orders_sync_around_collector_and_logging(): """Startup broadcast precedes collection; timing is consumed after updates.""" - from unilab.algos.torch.offpolicy.double_buffer_runner import DoubleBufferOffPolicyRunner + from unilab.algos.offpolicy.double_buffer_runner import DoubleBufferOffPolicyRunner source = inspect.getsource(DoubleBufferOffPolicyRunner.learn) assert source.index("self._dp_init_broadcast()") < source.index("self._start_collector(") @@ -350,7 +350,7 @@ def test_learn_source_orders_sync_around_collector_and_logging(): def test_fast_sac_initial_sync_tensors_return_live_references(): - from unilab.algos.torch.fast_sac.learner import FastSACLearner + from unilab.algos.fast_sac.learner import FastSACLearner learner = FastSACLearner( obs_dim=4, @@ -389,7 +389,7 @@ def test_fast_sac_initial_sync_tensors_return_live_references(): def test_fast_sac_syncs_each_optimizer_gradient_before_step(): - from unilab.algos.torch.fast_sac.learner import FastSACLearner + from unilab.algos.fast_sac.learner import FastSACLearner learner = FastSACLearner( obs_dim=4, @@ -446,7 +446,7 @@ def record_gradients(parameters) -> None: def test_fast_sac_gradient_sync_preserves_cuda_graph_capture(): - from unilab.algos.torch.fast_sac.learner import FastSACLearner + from unilab.algos.fast_sac.learner import FastSACLearner learner = FastSACLearner( obs_dim=4, @@ -482,7 +482,7 @@ def _build_sac_runner_with_dp_fakes(monkeypatch: pytest.MonkeyPatch, overrides: cfg = _offpolicy_cfg(overrides) monkeypatch.setattr(module.os, "cpu_count", lambda: 128) - import unilab.algos.torch.fast_sac.double_buffer as owner_module + import unilab.algos.fast_sac.double_buffer as owner_module class _Learner: def __init__(self, *args, **kwargs): @@ -556,7 +556,7 @@ def test_build_runner_multi_gpu_rank0_requires_log_dir(monkeypatch: pytest.Monke def test_flash_sac_initial_sync_tensors_return_live_references(): - from unilab.algos.torch.flash_sac.learner import FlashSACLearner + from unilab.algos.flash_sac.learner import FlashSACLearner learner = FlashSACLearner( obs_dim=4, @@ -603,7 +603,7 @@ def test_flash_sac_initial_sync_tensors_return_live_references(): def test_flash_sac_syncs_each_optimizer_gradient_before_step(): - from unilab.algos.torch.flash_sac.learner import FlashSACLearner + from unilab.algos.flash_sac.learner import FlashSACLearner learner = FlashSACLearner( obs_dim=4, @@ -650,7 +650,7 @@ def record_gradients(parameters) -> None: def test_flash_sac_gradient_sync_preserves_cuda_graph_and_cpu_fallback_updates(): - from unilab.algos.torch.flash_sac.learner import FlashSACLearner + from unilab.algos.flash_sac.learner import FlashSACLearner learner = FlashSACLearner( obs_dim=4, @@ -706,7 +706,7 @@ def _build_flashsac_runner_with_dp_fakes(monkeypatch: pytest.MonkeyPatch, overri cfg = _offpolicy_cfg(overrides) monkeypatch.setattr(module.os, "cpu_count", lambda: 128) - import unilab.algos.torch.flash_sac.double_buffer as flash_module + import unilab.algos.flash_sac.double_buffer as flash_module monkeypatch.setattr(flash_module, "ensure_registries", lambda: None) monkeypatch.setattr(flash_module, "create_env", lambda *args, **kwargs: _FakeEnv()) diff --git a/tests/algos/test_offpolicy_runner_unit.py b/tests/algos/test_offpolicy_runner_unit.py index 87d43b9c3..eac5ce832 100644 --- a/tests/algos/test_offpolicy_runner_unit.py +++ b/tests/algos/test_offpolicy_runner_unit.py @@ -11,13 +11,13 @@ import pytest import torch -import unilab.algos.torch.offpolicy.double_buffer_runner as device_runner_module -import unilab.algos.torch.offpolicy.runner as runner_module -from unilab.algos.torch.offpolicy.double_buffer_runner import ( +import unilab.algos.offpolicy.double_buffer_runner as device_runner_module +import unilab.algos.offpolicy.runner as runner_module +from unilab.algos.offpolicy.double_buffer_runner import ( _LearnerInferenceScheduler, algo_display_name, ) -from unilab.algos.torch.offpolicy.runner import ( +from unilab.algos.offpolicy.runner import ( OffPolicyRunner, build_offpolicy_sample_info, compute_train_start_threshold, @@ -575,15 +575,15 @@ def test_drain_metrics_propagates_collector_error(): @pytest.mark.parametrize("algo_type", ["sac", "td3", "flashsac"]) def test_learner_inference_matches_existing_actor_exploration(algo_type: str) -> None: if algo_type == "sac": - from unilab.algos.torch.fast_sac.learner import SACActor + from unilab.algos.fast_sac.learner import SACActor actor = SACActor(3, 2, hidden_dim=8, use_layer_norm=False) elif algo_type == "flashsac": - from unilab.algos.torch.flash_sac.network import FlashSACActor + from unilab.algos.flash_sac.network import FlashSACActor actor = FlashSACActor(num_blocks=1, input_dim=3, hidden_dim=8, action_dim=2) else: - from unilab.algos.torch.fast_td3.learner import TD3Actor + from unilab.algos.fast_td3.learner import TD3Actor actor = TD3Actor(3, 2, num_envs=2, init_scale=0.01, hidden_dim=8) expected_actor = copy.deepcopy(actor) @@ -624,7 +624,7 @@ def test_learner_inference_matches_existing_actor_exploration(algo_type: str) -> def test_hora_learner_inference_uses_privileged_context() -> None: - from unilab.algos.torch.hora.sac_models import HoraSACActor + from unilab.algos.hora.sac_models import HoraSACActor actor = HoraSACActor( obs_dim=3, diff --git a/tests/algos/test_offpolicy_runtime.py b/tests/algos/test_offpolicy_runtime.py index 437d78435..d399bf87b 100644 --- a/tests/algos/test_offpolicy_runtime.py +++ b/tests/algos/test_offpolicy_runtime.py @@ -4,7 +4,7 @@ def test_offpolicy_runtime_defaults_to_standard_sac_overrides() -> None: - from unilab.algos.torch.offpolicy.runtime import OffPolicyRuntime + from unilab.algos.offpolicy.runtime import OffPolicyRuntime runtime = OffPolicyRuntime() @@ -15,7 +15,7 @@ def test_offpolicy_runtime_defaults_to_standard_sac_overrides() -> None: def test_hora_sac_runtime_builds_privileged_actor_kwargs() -> None: - from unilab.algos.torch.hora.sac import resolve_hora_sac_runtime + from unilab.algos.hora.sac import resolve_hora_sac_runtime runtime = resolve_hora_sac_runtime( { @@ -38,7 +38,7 @@ def test_hora_sac_runtime_builds_privileged_actor_kwargs() -> None: def test_hora_sac_runtime_requires_critic_tail() -> None: - from unilab.algos.torch.hora.sac import resolve_hora_sac_runtime + from unilab.algos.hora.sac import resolve_hora_sac_runtime runtime = resolve_hora_sac_runtime({"runtime_impl": "hora_sac"}) @@ -48,7 +48,7 @@ def test_hora_sac_runtime_requires_critic_tail() -> None: def test_offpolicy_runtime_rejects_marker_without_resolver() -> None: - from unilab.algos.torch.offpolicy.runtime import resolve_custom_offpolicy_runtime + from unilab.algos.offpolicy.runtime import resolve_custom_offpolicy_runtime with pytest.raises(ValueError, match="runtime_impl='hora_sac'.*runtime_resolver"): resolve_custom_offpolicy_runtime({"runtime_impl": "hora_sac"}) diff --git a/tests/algos/test_offpolicy_thread_budget.py b/tests/algos/test_offpolicy_thread_budget.py index 15403ae6b..13509af8b 100644 --- a/tests/algos/test_offpolicy_thread_budget.py +++ b/tests/algos/test_offpolicy_thread_budget.py @@ -4,7 +4,7 @@ import pytest -from unilab.algos.torch.offpolicy.thread_budget import ( +from unilab.algos.offpolicy.thread_budget import ( apply_torch_thread_runtime, format_torch_thread_runtime, resolve_torch_thread_runtime, diff --git a/tests/algos/test_offpolicy_worker.py b/tests/algos/test_offpolicy_worker.py index d6b54368e..17b27352f 100644 --- a/tests/algos/test_offpolicy_worker.py +++ b/tests/algos/test_offpolicy_worker.py @@ -6,8 +6,8 @@ import pytest import torch -from unilab.algos.torch.common.collector_timing import extract_env_step_breakdown_timing_ms -from unilab.algos.torch.offpolicy.worker import ( +from unilab.algos.common.collector_timing import extract_env_step_breakdown_timing_ms +from unilab.algos.offpolicy.worker import ( _publish_inference_tick, _wait_for_inference_tick, compute_collector_active_steps_per_sec, diff --git a/tests/algos/test_rsl_rl_ppo.py b/tests/algos/test_rsl_rl_ppo.py index d93d0d968..078226c3e 100644 --- a/tests/algos/test_rsl_rl_ppo.py +++ b/tests/algos/test_rsl_rl_ppo.py @@ -9,7 +9,7 @@ from omegaconf import OmegaConf from tensordict import TensorDict -from unilab.algos.torch.rsl_rl_ppo import FinalObservationAwarePPO +from unilab.algos.rsl_rl_ppo import FinalObservationAwarePPO from unilab.training.rsl_rl import ( RslRlVecEnvWrapper, apply_rsl_rl_rank_seed, @@ -225,7 +225,7 @@ def test_normalize_ppo_train_cfg_preserves_unilab_runtime_flags() -> None: train_cfg = normalize_ppo_train_cfg( { "algorithm": { - "class_name": "unilab.algos.torch.rsl_rl_ppo:FinalObservationAwarePPO", + "class_name": "unilab.algos.rsl_rl_ppo:FinalObservationAwarePPO", "enable_compile": True, "target_kl_stop": None, }, diff --git a/tests/algos/test_torch_compile_helper.py b/tests/algos/test_torch_compile_helper.py index d4b1eb30f..599e185e5 100644 --- a/tests/algos/test_torch_compile_helper.py +++ b/tests/algos/test_torch_compile_helper.py @@ -6,7 +6,7 @@ import torch -from unilab.algos.torch.common import compile as compile_helper +from unilab.algos.common import compile as compile_helper def _without_triton(name: str, *args: Any, **kwargs: Any) -> Any: diff --git a/tests/config/test_locomotion_params.py b/tests/config/test_locomotion_params.py index ae2e50f3e..34b01053a 100644 --- a/tests/config/test_locomotion_params.py +++ b/tests/config/test_locomotion_params.py @@ -96,7 +96,7 @@ def test_ppo_config_defaults(): assert cfg.algo == "ppo" assert cfg.max_iterations == 101 assert cfg.algorithm.clip_param == 0.2 - assert cfg.algorithm.class_name == "unilab.algos.torch.rsl_rl_ppo:FinalObservationAwarePPO" + assert cfg.algorithm.class_name == "unilab.algos.rsl_rl_ppo:FinalObservationAwarePPO" assert cfg.algorithm.enable_compile is True assert cfg.policy.class_name == "ActorCritic" diff --git a/tests/ipc/test_replay_pipeline_gpu_resident.py b/tests/ipc/test_replay_pipeline_gpu_resident.py index 862d379ba..43bbdc224 100644 --- a/tests/ipc/test_replay_pipeline_gpu_resident.py +++ b/tests/ipc/test_replay_pipeline_gpu_resident.py @@ -151,7 +151,7 @@ def test_invalid_pack_layout_rejected(self): GPUResidentReplayPipeline(rb, device="cpu", sample_count=8, pack_layout="bogus") def test_runner_rejects_non_accelerator_before_base_initialization(self): - from unilab.algos.torch.offpolicy.double_buffer_runner import ( + from unilab.algos.offpolicy.double_buffer_runner import ( DoubleBufferOffPolicyRunner, ) diff --git a/tests/nan_injection/proto_him_ppo_inject.py b/tests/nan_injection/proto_him_ppo_inject.py index ffb28bb50..885cb3b78 100644 --- a/tests/nan_injection/proto_him_ppo_inject.py +++ b/tests/nan_injection/proto_him_ppo_inject.py @@ -34,7 +34,7 @@ sys.path.insert(0, str(ROOT_DIR)) sys.path.insert(0, str(ROOT_DIR / "src")) -from unilab.algos.torch.him_ppo.runner import HIMOnPolicyRunner # noqa: E402 +from unilab.algos.him_ppo.runner import HIMOnPolicyRunner # noqa: E402 from unilab.base.backend.mujoco.xml import materialize_scene_visual_override # noqa: E402 from unilab.training import BackendAdapter, create_env, ensure_registries # noqa: E402 from unilab.training.rsl_rl import RslRlVecEnvWrapper # noqa: E402 diff --git a/tests/nan_injection/stage2_nan_inject.py b/tests/nan_injection/stage2_nan_inject.py index 06736e5a8..bb01202e9 100644 --- a/tests/nan_injection/stage2_nan_inject.py +++ b/tests/nan_injection/stage2_nan_inject.py @@ -67,7 +67,7 @@ def _attach_guard(env, output_dir: Path) -> NanGuard: def build_him_ppo(num_envs: int, log_dir: Path, output_dir: Path): - from unilab.algos.torch.him_ppo.runner import HIMOnPolicyRunner + from unilab.algos.him_ppo.runner import HIMOnPolicyRunner ensure_registries() with initialize_config_dir(config_dir=str(ROOT_DIR / "conf/ppo_him"), version_base=None): diff --git a/tests/nan_injection/stage3_nan_inject.py b/tests/nan_injection/stage3_nan_inject.py index b23860558..5d86b1858 100644 --- a/tests/nan_injection/stage3_nan_inject.py +++ b/tests/nan_injection/stage3_nan_inject.py @@ -193,8 +193,8 @@ def cleanup(self): def _check_double_buffer_runner_wires_nan_guard(): """Verify DoubleBufferOffPolicyRunner passes nan_guard_cfg to collector.""" - import unilab.algos.torch.offpolicy.double_buffer_runner as db_mod - import unilab.algos.torch.offpolicy.runner as runner_mod + import unilab.algos.offpolicy.double_buffer_runner as db_mod + import unilab.algos.offpolicy.runner as runner_mod with ( patch.object(db_mod, "ReplayBuffer", _FakeReplayBuffer), @@ -244,8 +244,8 @@ def capture_start_collector(*, target_fn, kwargs): def _check_appo_runner_wires_nan_guard(): """Verify APPORunner passes nan_guard_cfg to collector.""" - import unilab.algos.torch.appo.runner as appo_mod - from unilab.algos.torch.appo.runner import APPORunner + import unilab.algos.appo.runner as appo_mod + from unilab.algos.appo.runner import APPORunner def fake_detect_dims(self): self.critic_dim = 4 diff --git a/tests/scripts/test_repo_hygiene.py b/tests/scripts/test_repo_hygiene.py index c12632aec..61409f0f5 100644 --- a/tests/scripts/test_repo_hygiene.py +++ b/tests/scripts/test_repo_hygiene.py @@ -5,8 +5,8 @@ from tests.scripts import repo_hygiene_checks REMOVED_ORPHAN_MODULES = ( - "src/unilab/algos/torch/common/base_collector.py", - "src/unilab/algos/torch/fast_td3/runner.py", + "src/unilab/algos/common/base_collector.py", + "src/unilab/algos/fast_td3/runner.py", ) diff --git a/tests/scripts/test_train_scripts.py b/tests/scripts/test_train_scripts.py index 63baace03..9af4db514 100644 --- a/tests/scripts/test_train_scripts.py +++ b/tests/scripts/test_train_scripts.py @@ -407,7 +407,7 @@ def test_hora_distill_runtime_checkpoint_records_model_only(): def test_hora_distill_checkpoint_runtime_only_restores_model_structure(): - from unilab.algos.torch.hora.distill import cfg_with_checkpoint_runtime + from unilab.algos.hora.distill import cfg_with_checkpoint_runtime cfg = _hora_distill_cfg(["task=sharpa_inhand/mujoco_nodr"]) checkpoint = { @@ -471,8 +471,8 @@ def test_hora_distill_checkpoint_runtime_only_overrides_model_side( teacher_algo_family: str, checkpoint_model: dict[str, Any], ): - from unilab.algos.torch.hora import distill_config - from unilab.algos.torch.hora.distill import cfg_with_checkpoint_runtime + from unilab.algos.hora import distill_config + from unilab.algos.hora.distill import cfg_with_checkpoint_runtime owner_cfg = OmegaConf.create( { @@ -1967,8 +1967,8 @@ def test_offpolicy_play_actor_spec_keeps_standard_sac_and_flashsac(): def test_offpolicy_build_play_actor_preserves_flashsac_model_kwargs( monkeypatch: pytest.MonkeyPatch, ): - import unilab.algos.torch.common.actor_factory as actor_factory - import unilab.algos.torch.common.normalization as normalization + import unilab.algos.common.actor_factory as actor_factory + import unilab.algos.common.normalization as normalization from unilab.training.offpolicy import build_play_actor captured: dict[str, Any] = {} @@ -2024,7 +2024,7 @@ def test_offpolicy_build_play_actor_restores_td3_state_and_normalizer( ): import torch - import unilab.algos.torch.fast_td3.learner as learner_module + import unilab.algos.fast_td3.learner as learner_module from unilab.training.offpolicy import build_play_actor, load_play_actor captured: dict[str, Any] = {} @@ -2185,7 +2185,7 @@ def run_playback_mode(self, **kwargs): ), ) - import unilab.algos.torch.common.actor_factory as actor_factory + import unilab.algos.common.actor_factory as actor_factory monkeypatch.setattr(actor_factory, "build_actor", lambda *args, **kwargs: FakeActor()) @@ -2308,7 +2308,7 @@ def run_playback_mode(self, **kwargs): lambda *args, **kwargs: (str(checkpoint), str(run_dir)), ) - import unilab.algos.torch.common.actor_factory as actor_factory + import unilab.algos.common.actor_factory as actor_factory def fake_build_actor(algo_type, obs_dim, action_dim, hidden_dim, use_layer_norm, device, **kw): captured["build_actor"] = (algo_type, obs_dim, action_dim, kw) @@ -2547,7 +2547,7 @@ def reset(self, env_indices): def test_play_wrapper_preserves_hora_priv_info_and_proprio_history(): import numpy as np - from unilab.algos.torch.hora.rsl_rl import HoraRslRlVecEnvWrapper + from unilab.algos.hora.rsl_rl import HoraRslRlVecEnvWrapper class FakeEnv: def __init__(self): @@ -2668,7 +2668,7 @@ def step(self, actions): def test_play_wrapper_timeout_bootstrap_preserves_hora_priv_info(): import torch - from unilab.algos.torch.hora.rsl_rl import HoraRslRlVecEnvWrapper + from unilab.algos.hora.rsl_rl import HoraRslRlVecEnvWrapper class FakeEnv: def __init__(self): diff --git a/tests/utils/test_algo_utils.py b/tests/utils/test_algo_utils.py index bb4ba181b..950b98914 100644 --- a/tests/utils/test_algo_utils.py +++ b/tests/utils/test_algo_utils.py @@ -8,7 +8,7 @@ import pytest -from unilab.algos.torch.common.actor_factory import build_actor +from unilab.algos.common.actor_factory import build_actor from unilab.base.registry import ensure_registries diff --git a/tests/utils/test_utils_package_policy.py b/tests/utils/test_utils_package_policy.py index 8030ed238..c6990686c 100644 --- a/tests/utils/test_utils_package_policy.py +++ b/tests/utils/test_utils_package_policy.py @@ -35,8 +35,8 @@ "xml_utils", } REMOVED_OWNER_ALIASES = { - "unilab.algos.torch.offpolicy.logging", - "unilab.algos.torch.common.tensor", + "unilab.algos.offpolicy.logging", + "unilab.algos.common.tensor", } @@ -69,7 +69,7 @@ def test_removed_owner_aliases_are_not_importable() -> None: def test_algos_torch_common_no_longer_reexports_utils_primitives() -> None: - common = importlib.import_module("unilab.algos.torch.common") + common = importlib.import_module("unilab.algos.common") assert "get_default_device" not in common.__all__ assert "to_numpy" not in common.__all__ assert "to_torch" not in common.__all__ diff --git a/tests/visualization/test_interactive_playback.py b/tests/visualization/test_interactive_playback.py index fbc0ec523..da62e95d7 100644 --- a/tests/visualization/test_interactive_playback.py +++ b/tests/visualization/test_interactive_playback.py @@ -527,7 +527,7 @@ def test_sac_playback_session_runs_sim2sim_preflight( ) -> None: from omegaconf import OmegaConf - import unilab.algos.torch.common.actor_factory as actor_factory + import unilab.algos.common.actor_factory as actor_factory import unilab.training.offpolicy as offpolicy_play import unilab.training.run as training_run import unilab.visualization.interactive_playback as interactive_playback @@ -667,8 +667,8 @@ def test_appo_hora_playback_session_uses_hora_wrapper_and_actor_checkpoint( from omegaconf import OmegaConf from tensordict import TensorDict - import unilab.algos.torch.hora.models as hora_models - import unilab.algos.torch.hora.rsl_rl as hora_rsl + import unilab.algos.hora.models as hora_models + import unilab.algos.hora.rsl_rl as hora_rsl checkpoint = tmp_path / "model_10.pt" torch.save({"actor": {"weight": torch.tensor(1.0)}}, checkpoint) @@ -782,7 +782,7 @@ def test_sac_hora_playback_session_updates_priv_info_after_reset_and_step( ) -> None: from omegaconf import OmegaConf - import unilab.algos.torch.common.actor_factory as actor_factory + import unilab.algos.common.actor_factory as actor_factory import unilab.training.offpolicy as offpolicy_play import unilab.training.run as training_run @@ -915,8 +915,8 @@ def test_hora_distill_playback_session_loads_stage2_checkpoint_and_student_polic from omegaconf import OmegaConf from tensordict import TensorDict - import unilab.algos.torch.hora.distill as distill - import unilab.algos.torch.hora.rsl_rl as hora_rsl + import unilab.algos.hora.distill as distill + import unilab.algos.hora.rsl_rl as hora_rsl import unilab.training as training checkpoint = tmp_path / "hora_stage2_last.pt" From 3be1fc7f647898ec19e3f8cc394ae107e04ddc53 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Fri, 21 Aug 2026 02:07:01 +0800 Subject: [PATCH 097/169] chore: close the #1042 registry, documentation, and performance acceptance (#1237) - Pin the post-migration production registry with closeout boundary tests: exact registry/matrix equality, frozen LegacyFactoryAdapter seam limited to the 3 approved families, canonical manager-runtime factories for all other registrations, and removed legacy unilab.envs.* package guards. - Delete unused legacy code: PenaltyCurriculum, four dead command helpers, PdControlConfig, 18 dead shared reward functions plus run_reward_dispatch, and dead Go1/Go2 reward-config test fixtures. - Document the final task status (36 Compatible / 3 Adapted, 39 tasks / 86 backend registrations) and the two approved factory wrappers, refresh stale domain-randomization pages to the manager-based events path, and realign the English support-matrix mirror. Closes #1228 --- .../5-domain_randomization/0-index.md | 77 +++--- .../5-domain_randomization/1-configuration.md | 37 ++- .../2-writing_providers.md | 13 +- .../3-framework_migration/1-from_isaac_lab.md | 24 ++ .../source/en/5-reference/5-support_matrix.md | 148 ++++++++---- .../5-domain_randomization/0-index.md | 66 +++--- .../5-domain_randomization/1-configuration.md | 29 ++- .../2-writing_providers.md | 11 +- .../3-framework_migration/1-from_isaac_lab.md | 22 ++ src/unilab/base/curriculum.py | 68 ------ .../tasks/locomotion/common/__init__.py | 14 +- src/unilab/tasks/locomotion/common/base.py | 17 -- .../tasks/locomotion/common/commands.py | 42 ---- src/unilab/tasks/locomotion/common/rewards.py | 224 ------------------ .../tasks/locomotion/g1/manager_terms.py | 2 +- tests/conftest.py | 39 --- tests/envs/test_env_configs.py | 12 +- .../test_production_registry_closeout.py | 167 +++++++++++++ 18 files changed, 469 insertions(+), 543 deletions(-) create mode 100644 tests/tasks/test_production_registry_closeout.py diff --git a/docs/sphinx/source/en/2-user_guide/5-domain_randomization/0-index.md b/docs/sphinx/source/en/2-user_guide/5-domain_randomization/0-index.md index fe98d3b9b..5caf9e859 100644 --- a/docs/sphinx/source/en/2-user_guide/5-domain_randomization/0-index.md +++ b/docs/sphinx/source/en/2-user_guide/5-domain_randomization/0-index.md @@ -1,9 +1,14 @@ # Domain Randomization -This page only describes the current status of tasks in the repo that are already registered and already wired to a DR provider. All conclusions come from the code; nothing is inferred from design intent. +This page only describes the current domain randomization status of registered tasks in the repo. All conclusions come from the code; nothing is inferred from design intent. -The current unified entry point lives in `NpEnv._init_domain_randomization()` and `DomainRandomizationManager`: +Two DR declaration paths exist today: + +- **Manager-Based (Compatible) tasks**: reset / interval randomization is declared through Hydra `events:` manager terms in the owner YAML; reset-lifecycle events sample at reset, interval-lifecycle events perturb between steps. See the `events:` block of `conf/ppo/task/go1_joystick_flat/base.yaml` for an example. +- **Legacy provider path**: only the 3 Adapted families (`sharpa_inhand` / `sharpa_inhand_grasp` / `go2_arm_manip_loco`, including their appo / hora / ppo_him owners) still declare `env.domain_rand.*` configuration through a `DomainRandomizationProvider` + `DomainRandomizationManager`. + +The unified entry point of the legacy provider path lives in `NpEnv._init_domain_randomization()` and `DomainRandomizationManager`: - init path: the task provider produces an `InitRandomizationPlan`; the manager calls the backend's `apply_init_randomization(...)` during env initialization - reset path: the task provider produces a `ResetPlan`; the manager validates capability and then calls the backend's `set_state(..., randomization=...)` @@ -17,34 +22,35 @@ These three paths correspond to three lifecycle classes: ## Status Conclusions -1. All tasks currently wired to a DR provider use the unified DR entry point; no task bypasses `DomainRandomizationManager` to run a separate DR flow inside `reset()`. -2. They are all roughly structured: legacy task owners define a `domain_rand` config dataclass, a `DomainRandomizationProvider`, and a `ResetPlan`; Manager-Based owners declare reset behavior through Hydra command/event terms. G1 motion reset perturbations belong to `MotionCommandCfg`, while WBT adds `EventTermCfg` reset and interval terms. -3. What is "unified" today is mainly the entry point and execution flow, not every randomization item itself. The shared helper `build_common_reset_randomization()` currently generates `base_mass_delta`, `base_com_offset`, `gravity`, `kp`, `kd`; the shared interval helper currently only generates push. +1. Manager-Based tasks do not register a DR provider; their reset/interval randomization consists of `events:` manager terms in the owner YAML, executed uniformly by the manager lifecycle. Only the frozen compatibility factories of the Adapted families still go through the `DomainRandomizationManager` unified entry point. +2. Adapted-family owners define a `domain_rand` config dataclass, a `DomainRandomizationProvider`, and a `ResetPlan`; Manager-Based owners declare reset behavior through Hydra command/event terms. G1 motion reset perturbations belong to `MotionCommandCfg`, while WBT adds `EventTermCfg` reset and interval terms. +3. What is "unified" today is mainly the entry point and execution flow, not every randomization item itself. The legacy path's shared helper `build_common_reset_randomization()` currently generates `base_mass_delta`, `base_com_offset`, `gravity`, `kp`, `kd`; the shared interval helper currently only generates push. 4. `ResetRandomizationPayload` can already express `gravity`, `body_iquat`, `body_inertia`, `kp`, `kd`, and `MuJoCoBackend` has declared support. Whether these are actually used still depends on whether the task provider samples and dispatches them. 5. `MotrixBackend` currently supports `base_mass_delta`, `base_com_offset`, `kp`, `kd`, and interval push; and it requires all model actuators to be position actuators during initialization. 6. `geom_size` is not a reset-lifecycle field; Sharpa-hand object geom scale is handled by init-lifecycle model materialization. ## Uniformity Assessment Table -| Task | Uses unified DR entry? | Structured form? | reset form | interval form | Code | +| Task | Declaration path | Structured form? | reset form | interval form | Code | | --- | --- | --- | --- | --- | --- | -| `Go1JoystickFlat` | Yes | Yes: `Domain_Rand + Provider + ResetPlan` | task state sampling + common payload | push | `go1/joystick.py` | -| `Go2JoystickFlat` | Yes | Yes: `Domain_Rand + Provider + ResetPlan` | task state sampling + common payload | push | `go2/joystick.py` | -| `G1WalkFlat` | Yes | Yes: Hydra `EventTermCfg` + Manager-Based reset terms | root-state reset + kp/kd via `pd_gains` | none | `g1/manager_terms.py` | -| `G1WalkRough` | Yes | Yes: same Manager-Based event terms as `G1WalkFlat` | root-state reset + kp/kd via `pd_gains` | none | `g1/manager_terms.py` | -| `G1MotionTracking` | Yes | Yes: Hydra `MotionCommandCfg` + Manager-Based command reset | motion frame, root pose/velocity, and joint-position sampling | none | `motion_tracking/common/manager_terms.py` | -| `G1WBTObs` | Yes | Yes: same motion command + Hydra `EventTermCfg` | motion reset plus mass/COM/PD/friction/encoder-bias events | interval velocity kick | `motion_tracking/g1/manager_terms.py` | -| `AllegroInhandRotation` | Yes | Yes: Hydra `EventTermCfg` + Manager-Based reset term | entity-scoped hand/ball reset | none | `allegro_inhand/manager_terms.py` | -| `AllegroInhandRotationGrasp` | Yes | Yes: reuses the rotation reset event + `RecorderTermCfg` | noisy hand reset + grasp collection | none | `allegro_inhand/grasp_gen.py` | -| `SharpaInhandRotation` | Yes | Yes: `InitRandomizationPlan + ResetPlan + IntervalRandomizationPlan` | grasp cache sampling + common payload | object `body_force` | `sharpa_inhand/rotation.py` | -| `SharpaInhandRotationGrasp` | Yes | Yes: reuses the Sharpa rotation provider and overrides reset sampling | grasp collection reset + common payload | none | `sharpa_inhand/grasp_gen.py` | +| `Go1JoystickFlat` | Hydra `events:` terms | Yes: owner YAML declares reset/interval events | root-state reset + base mass/COM + `pd_gains` | `push_by_setting_velocity` event | `conf/ppo/task/go1_joystick_flat/base.yaml` | +| `Go2JoystickFlat` | Hydra `events:` terms | Yes: owner YAML declares reset events | root-state reset + `pd_gains` kp/kd | none | `conf/ppo/task/go2_joystick_flat/base.yaml` | +| `G1WalkFlat` | Hydra `events:` terms | Yes: Hydra `EventTermCfg` + Manager-Based reset terms | root-state reset + kp/kd via `pd_gains` | none | `g1/manager_terms.py` | +| `G1WalkRough` | Hydra `events:` terms | Yes: same Manager-Based event terms as `G1WalkFlat` | root-state reset + kp/kd via `pd_gains` | none | `g1/manager_terms.py` | +| `G1MotionTracking` | Hydra command term | Yes: Hydra `MotionCommandCfg` + Manager-Based command reset | motion frame, root pose/velocity, and joint-position sampling | none | `motion_tracking/common/manager_terms.py` | +| `G1WBTObs` | Hydra `events:` terms | Yes: same motion command + Hydra `EventTermCfg` | motion reset plus mass/COM/PD/friction/encoder-bias events | interval velocity kick | `motion_tracking/g1/manager_terms.py` | +| `AllegroInhandRotation` | Hydra `events:` terms | Yes: Hydra `EventTermCfg` + Manager-Based reset term | entity-scoped hand/ball reset | none | `allegro_inhand/manager_terms.py` | +| `AllegroInhandRotationGrasp` | Hydra `events:` terms | Yes: reuses the rotation reset event + `RecorderTermCfg` | noisy hand reset + grasp collection | none | `allegro_inhand/grasp_gen.py` | +| `SharpaInhandRotation` | legacy provider | Yes: `InitRandomizationPlan + ResetPlan + IntervalRandomizationPlan` | grasp cache sampling + common payload | object `body_force` | `sharpa_inhand/rotation.py` | +| `SharpaInhandRotationGrasp` | legacy provider | Yes: reuses the Sharpa rotation provider and overrides reset sampling | grasp collection reset + common payload | none | `sharpa_inhand/grasp_gen.py` | +| `Go2ArmManipLoco` | legacy provider | Yes: `DomainRandConfig + LocomotionDRProvider subclass + ResetPlan` | task state sampling + common payload | push | `go2_arm/manip_loco.py` | ## Per-task Domain Randomization List | Task | Currently implemented reset domain randomization | Currently implemented interval domain randomization | Default state | | --- | --- | --- | --- | -| `Go1JoystickFlat` | base xy; base yaw; base qvel; command sampling; `current_actions/last_actions` zeroed; optional `base_mass_delta`; optional `base_com_offset`; optional `gravity` | `push_robots` | `base_mass_delta`, `base_com_offset`, and push enabled by default; `gravity` disabled by default | -| `Go2JoystickFlat` | base xy; base yaw; base qvel; command sampling; `current_actions/last_actions` zeroed; kp/kd randomization (enabled by default); optional `base_mass_delta`; optional `base_com_offset`; optional `gravity` | `push_robots` | kp/kd enabled by default; common payload and push disabled by default | +| `Go1JoystickFlat` | base xy/yaw and base qvel via `reset_root_state_uniform`; command sampling (`UniformVelocityCommandCfg`); base mass via `randomize_rigid_body_mass`; base COM via `randomize_rigid_body_com`; kp/kd via `pd_gains` | `push_by_setting_velocity` interval event | all listed event terms are declared and enabled by default in `conf/ppo/task/go1_joystick_flat/base.yaml` | +| `Go2JoystickFlat` | base xy/yaw and base qvel via `reset_root_state_uniform`; command sampling; kp/kd via `pd_gains` | none | event terms declared and enabled by default in `conf/ppo/task/go2_joystick_flat/base.yaml` | | `G1WalkFlat` | base xy/yaw and base qvel via `reset_root_state_uniform`; command sampling with a planar dead zone; `gait_phase` sampling; kp/kd randomization via `pd_gains` | none | kp/kd enabled on mujoco owners by default; disabled on motrix/mjwarp owners | | `G1WalkRough` | Same as `G1WalkFlat` (shared owner bases, rough scene) | none | Same defaults as `G1WalkFlat` | | `G1MotionTracking` | Motion-command frame sampling; root pose perturbation `x/y/z/roll/pitch/yaw`; root velocity perturbation `x/y/z/roll/pitch/yaw`; joint-position noise clipped through the public entity soft limits; action-manager state reset | none | `pose_range`, `velocity_range`, and `joint_position_range` have non-zero perturbations in the base owner | @@ -56,26 +62,30 @@ These three paths correspond to three lifecycle classes: ## Current Unified DR Capabilities and Boundaries -### 1. Unified Entry Point Is Complete +### 1. The Legacy Provider Entry Point Is Unified -The unified entry point is guaranteed by `NpEnv` and `DomainRandomizationManager`: +The unified entry point of the legacy provider path is guaranteed by `NpEnv` +and `DomainRandomizationManager`: - Tasks only need to register a provider - The manager uniformly performs capability validation - The backend is uniformly responsible for actually applying the randomization payload -So from an execution-path perspective, the tasks are already unified. +So from an execution-path perspective, the Adapted families still on this path +are unified; Manager-Based tasks instead execute the `events:` terms declared +in the owner YAML through the manager lifecycle. ### 2. The Shared Helpers Are Still Narrow -`dr_utils.py` currently has only two classes of shared helpers: +The legacy path's `dr_utils.py` currently has only two classes of shared helpers: - reset common payload: `base_mass_delta`, `base_com_offset`, `gravity`, `kp`, `kd` - interval common payload: push This means: -- Although locomotion tasks all go through the unified entry point, their base xy, yaw, qvel, command, and gait phase are still sampled directly inside each provider +- The go2_arm / sharpa families still on the legacy provider path sample their + task-specific state directly inside each provider - `G1MotionTracking`'s pose / velocity / joint noise is owned by its manager command - Allegro's grasp / object initial state sampling is entirely task-specific logic - Sharpa's `geom_size` scale is init-lifecycle model materialization and is not part of the reset common payload @@ -116,7 +126,10 @@ But on the task side, the current reality is: not every provider constructs thes - Lifecycle: only sampled and written at reset; the env retains that gravity until the next reset re-samples it. - Backend: currently in UniLab, only the MuJoCo backend declares support for this reset term; the Motrix backend does not. Some tasks filter it by capability and skip it; others raise an error in the validate stage. -The config entry is under each task's `env.domain_rand`: +The config entry exists only under `env.domain_rand` of the Adapted-family +owners still on the legacy provider path (`sharpa_inhand_grasp`, +`go2_arm_manip_loco`, and their hora / appo / ppo_him variants); Manager-Based +tasks have no `env.domain_rand`: ```yaml env: @@ -136,7 +149,7 @@ Field semantics: If you only want to randomize the magnitude while keeping the vertical-down direction, only open up the `z` component: ```bash -uv run train --algo ppo --task g1_motion_tracking --sim mujoco \ +uv run train --algo ppo --task sharpa_inhand_grasp --sim mujoco \ env.domain_rand.randomize_gravity=true \ 'env.domain_rand.gravity_range=[[0.0,0.0,-10.5],[0.0,0.0,-8.5]]' ``` @@ -144,7 +157,7 @@ uv run train --algo ppo --task g1_motion_tracking --sim mujoco \ If you want to randomize both direction and magnitude, open up `x/y/z`: ```bash -uv run train --algo ppo --task g1_motion_tracking --sim mujoco \ +uv run train --algo ppo --task sharpa_inhand_grasp --sim mujoco \ env.domain_rand.randomize_gravity=true \ 'env.domain_rand.gravity_range=[[-0.3,-0.3,-10.5],[0.3,0.3,-8.5]]' ``` @@ -159,7 +172,13 @@ Notes: ## Interval push Usage -Tasks supporting interval push configure it under `env.domain_rand`: +The `env.domain_rand.push_robots` family of fields exists only in the go2_arm +Adapted-family owners (`conf/ppo/task/go2_arm_manip_loco/mujoco.yaml` etc.); +Manager-Based tasks declare push through a `push_by_setting_velocity` interval +event term instead (for example `conf/ppo/task/go1_joystick_flat/base.yaml` and +`conf/ppo/task/quadruped_joystick_rough/base.yaml`). + +The go2_arm owners configure push under `env.domain_rand`: ```yaml env: @@ -176,11 +195,11 @@ env: - `push_body_name`: the target body / link to apply the force to. Defaults to `null`, meaning the backend's `base_name` is used. ```bash -uv run train --algo ppo --task g1_motion_tracking --sim mujoco \ +uv run train --algo ppo --task go2_arm_manip_loco --sim mujoco \ env.domain_rand.push_robots=true \ env.domain_rand.push_interval=500 \ 'env.domain_rand.max_force=[20.0,20.0,5.0]' \ - env.domain_rand.push_body_name=torso_link + env.domain_rand.push_body_name=base ``` Notes: diff --git a/docs/sphinx/source/en/2-user_guide/5-domain_randomization/1-configuration.md b/docs/sphinx/source/en/2-user_guide/5-domain_randomization/1-configuration.md index f2d91d87b..ca555110b 100644 --- a/docs/sphinx/source/en/2-user_guide/5-domain_randomization/1-configuration.md +++ b/docs/sphinx/source/en/2-user_guide/5-domain_randomization/1-configuration.md @@ -1,11 +1,19 @@ # Configuration -Domain randomization is configured inside the selected task owner YAML, usually -under `env.domain_rand`. Use `--task` and `--sim` to select backend-specific -behavior first, then override fields inside that selected owner. +Domain randomization is configured inside the selected task owner YAML. Use +`--task` and `--sim` to select backend-specific behavior first, then override +fields inside that selected owner. + +Two declaration paths exist today: + +- Manager-Based (Compatible) tasks declare reset / interval randomization + through Hydra `events:` manager terms in the owner YAML, for example + `conf/ppo/task/go1_joystick_flat/base.yaml`. +- Only the Adapted families (sharpa / go2_arm and their hora / appo / ppo_him + owners) still configure legacy provider fields under `env.domain_rand`. ```bash -uv run train --algo ppo --task g1_motion_tracking --sim mujoco \ +uv run train --algo ppo --task sharpa_inhand_grasp --sim mujoco \ env.domain_rand.randomize_gravity=true \ 'env.domain_rand.gravity_range=[[0.0,0.0,-10.5],[0.0,0.0,-8.5]]' ``` @@ -21,25 +29,30 @@ Common lifecycle boundaries: The detailed task status and field semantics are in {doc}`0-index`. Domain randomization is split by lifecycle: init, reset, and interval. The -manager path is `src/unilab/dr/manager.py`; task providers live near the env -owners, and backend capabilities are declared through +legacy path's manager is `src/unilab/dr/manager.py`; task providers live near +the env owners, and backend capabilities are declared through `src/unilab/base/backend/base.py`. ## Reset Gravity Use `--sim mujoco` when enabling gravity reset randomization; Motrix does not -advertise the same gravity capability in the current backend. +advertise the same gravity capability in the current backend. This item is only +available on the legacy provider path (Adapted-family owners). ```bash -uv run train --algo ppo --task g1_motion_tracking --sim mujoco \ +uv run train --algo ppo --task sharpa_inhand_grasp --sim mujoco \ env.domain_rand.randomize_gravity=true \ 'env.domain_rand.gravity_range=[[0.0,0.0,-10.5],[0.0,0.0,-8.5]]' ``` ## Interval Push +Manager-Based tasks declare push through a `push_by_setting_velocity` interval +event term; `env.domain_rand.push_robots` is only available on the go2_arm +Adapted-family owners. + ```bash -uv run train --algo ppo --task g1_motion_tracking --sim mujoco \ +uv run train --algo ppo --task go2_arm_manip_loco --sim mujoco \ env.domain_rand.push_robots=true \ env.domain_rand.push_interval=500 \ 'env.domain_rand.max_force=[20.0,20.0,5.0]' @@ -48,8 +61,10 @@ uv run train --algo ppo --task g1_motion_tracking --sim mujoco \ ## Owner-Local Defaults Keep ranges in the task owner YAML when they are part of the task contract. For -example, `conf/ppo/task/go2_joystick_rough/mujoco.yaml` enables base mass, -center-of-mass, kp/kd, and push randomization, while +example, the rough quadruped family's base mass, center-of-mass, kp/kd, and +push randomization are declared as event terms in the shared base +`conf/ppo/task/quadruped_joystick_rough/base.yaml` (the `go2_joystick_rough` +backend owners compose it through Hydra defaults), while `conf/ppo/task/sharpa_inhand/mujoco.yaml` configures object scale, friction, and force disturbance for Sharpa. 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 b63fe6bd9..061ab04e6 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 @@ -1,5 +1,12 @@ # Writing Providers +This page describes the legacy provider path: only the 3 Adapted families +(`sharpa_inhand` / `sharpa_inhand_grasp` / `go2_arm_manip_loco`) still declare +domain randomization through a task-level `DomainRandomizationProvider`. +Migrated Manager-Based tasks do not write providers; they declare randomization +through Hydra `events:` manager terms in the owner YAML (see {doc}`0-index` +and {doc}`1-configuration`). + Task-level domain randomization providers live with the task env owner. They sample task-specific state and return plans consumed by `DomainRandomizationManager`. @@ -25,9 +32,11 @@ The shared types live in `src/unilab/dr/types.py`, and the manager lives in ## Evidence -Representative provider implementations are in: +Representative provider implementations are in (all on the Adapted-family +compatibility path): -- `src/unilab/tasks/locomotion/common/dr_provider.py` +- `src/unilab/tasks/locomotion/common/dr_provider.py` (`LocomotionDRProvider`, + used by the go2_arm family) - `src/unilab/tasks/locomotion/go2_arm/manip_loco.py` - `src/unilab/tasks/manipulation/sharpa_inhand/rotation.py` diff --git a/docs/sphinx/source/en/3-deployment/3-framework_migration/1-from_isaac_lab.md b/docs/sphinx/source/en/3-deployment/3-framework_migration/1-from_isaac_lab.md index 7585d81de..2604f3bd6 100644 --- a/docs/sphinx/source/en/3-deployment/3-framework_migration/1-from_isaac_lab.md +++ b/docs/sphinx/source/en/3-deployment/3-framework_migration/1-from_isaac_lab.md @@ -151,6 +151,15 @@ uv run train --algo ppo --task --sim mujoco Do not add a task-specific training-script branch, environment factory, runner, or IPC path. +Two maintainer-approved factory wrappers are the only registered exceptions to +the generic-factory rule: `make_g1_walk_env` +(`src/unilab/tasks/locomotion/g1/manager_terms.py`) constructs a +`G1WalkManagerBasedEnv` subclass that provides the `NpEnv` +`build_symmetry_augmentation` symmetry hook, and `make_x2_wall_flip_env` +(`src/unilab/tasks/motion_tracking/x2/__init__.py`) resolves untracked X2 +meshes on the cold path before delegating to `make_manager_based_rl_env`. +Every other Compatible task registers `make_manager_based_rl_env` directly. + ### 6. Validate near each adaptation Test Hydra composition and typed materialization, term order and math, selector @@ -158,6 +167,21 @@ failure, observation/action shapes, partial reset, and at least one real registered backend transition. Compare behavior with the pinned source task; benchmark only after semantic migration is complete. +## Final task status + +The #1042 migration closeout covers 39 production tasks and 86 task/backend +registrations. The fail-closed source of truth is +`src/unilab/tasks/migration_matrix.py`: `migration_record()` raises `KeyError` +for a production task name with no entry, so adding a production registration +requires an explicit migration decision. + +- 36 tasks are **Compatible** (`target=complete`): the Hydra owner YAML + materializes the canonical NumPy Manager-Based runtime. +- 3 tasks are **Adapted** (`target=compatibility`): `Go2ArmManipLoco`, + `SharpaInhandRotation`, and `SharpaInhandRotationGrasp` keep custom + IK/history or tactile/contact/cache behavior behind one frozen compatibility + factory each; they migrate only when the formal capability exists. + ## Repository evidence `tests/fixtures/isaac_lab_cartpole/` ports the Manager-Based Cartpole task from diff --git a/docs/sphinx/source/en/5-reference/5-support_matrix.md b/docs/sphinx/source/en/5-reference/5-support_matrix.md index bae44e4f4..c317d8bb0 100644 --- a/docs/sphinx/source/en/5-reference/5-support_matrix.md +++ b/docs/sphinx/source/en/5-reference/5-support_matrix.md @@ -3,12 +3,16 @@ This matrix is generated conceptually from registry entries, owner YAMLs, and tests. The generator implementation is `src/unilab/utils/support_matrix.py`; the write target for the generated block is currently the Chinese reference page -`docs/sphinx/source/zh_CN/5-reference/5-support_matrix.md`. +`docs/sphinx/source/zh_CN/5-reference/5-support_matrix.md`. This English page +mirrors that generated content. ## Backend Selection Rules - The default backend is `mujoco`. - Switch to Motrix with `--sim motrix` on the unified CLI. +- `--sim mjwarp` currently maps only to the `g1_walk_flat` host adapter; PPO + (torch) and SAC (torch) are Tested, other entrypoints follow the matrix + below, and using it requires installing the `mjwarp` extra. - `--algo`, `--task`, and `--sim` jointly select the owner YAML. - Do not treat `training.sim_backend` as a standalone backend switch. @@ -17,70 +21,120 @@ write target for the generated block is currently the Chinese reference page - `mujoco`: `--render-mode auto` exports `play_video.mp4`. - `motrix`: `--render-mode auto` opens an interactive renderer window; it does not record a video and is not bound by `play_steps`. -- `--render-mode record`: both backends record a video only. +- `mjwarp`: only supports explicit, finite-step `record`, rendered offline + through the task owner's MuJoCo visual model; `auto`, interactive, and + native renderers are not supported. +- `--render-mode record`: MuJoCo, mjwarp, and Motrix all record a video only. - `--render-mode none`: no playback. ## Evidence Grades | Grade | Repository Evidence | | --- | --- | -| `Registered` | The env/backend pair appears after `registry.ensure_registries()`. | -| `Configured` | A matching owner YAML exists under `conf/ppo/task`, `conf/appo/task`, or `conf/offpolicy/task`. | -| `Tested` | Automated tests cover the entrypoint/task-owner/backend combination through config compose or runtime smoke. | +| `Registered` | The env/backend pair exists in `registry.list_registered_envs()` after `ensure_registries()`. | +| `Configured` | A matching owner YAML exists under `conf/{ppo,appo,offpolicy}/task/...`. | +| `Tested` | Automated tests under `tests/` cover the entrypoint/task-owner/backend combination, or an explicit maintainer full-training validation with near-risk automated tests exists. `Tested` here does not mean the default recommended path. | | `Benchmarked` | A checked-in benchmark manifest exists for the combination. | | `Recommended` | Explicit recommendation metadata exists in the repo. | -The current generator reports no checked-in benchmark manifest and no separate -recommendation metadata, so rows do not auto-promote to `Benchmarked` or +`Tested` only describes existing automated coverage or explicit maintainer +training validation; it does not imply the combination has all the backend +capabilities of the same-named MuJoCo owner. For example, a phase-1 Motrix +owner may only cover training smoke and an explicitly enabled DR subset. + +`mjwarp` only supports the `g1_walk_flat` host adapter. The PPO (torch) and SAC +(torch) owners have completed training validation and have backend, contract, +and playback automated coverage, so they are marked `Tested`. mjwarp playback +only supports explicit, finite-step `record` and reuses the MuJoCo offline +renderer; it does not support `auto`, interactive, or native playback. A +`Registered` mark on other entrypoints only denotes env/backend registry +identity, not support for the corresponding algorithm, terrain, full DR, or +production training. + +No checked-in benchmark manifest bound to these combinations has been detected, +so rows do not auto-promote to `Benchmarked`. There is also no separate +recommendation metadata in the repo, so rows do not auto-promote to `Recommended`. ## Entrypoint x Task Owner -| Entrypoint | Task owner | MuJoCo | Motrix | -| --- | --- | --- | --- | -| PPO (torch) | `go1_joystick_flat` | Tested | Tested | -| PPO (torch) | `go2_joystick_flat` | Tested | Tested | -| PPO (torch) | `go2_joystick_rough` | Tested | Tested | -| PPO (torch) | `g1_walk_flat` | Tested | Tested | -| PPO (torch) | `g1_motion_tracking` | Tested | Tested | -| PPO (torch) | `g1_flip_tracking` | Tested | Tested | -| PPO (torch) | `g1_wall_flip_tracking` | Tested | Tested | -| PPO (torch) | `allegro_inhand` | Tested | Tested | -| PPO (torch) | `sharpa_inhand` | Tested | Tested | -| PPO (torch) | `sharpa_inhand_grasp` | Tested | Tested | -| PPO (torch) | `allegro_inhand_grasp` | Tested | Tested | -| PPO (torch) | `g1_box_tracking` | Tested | Tested | -| PPO (torch) | `g1_climb_tracking` | Tested | Tested | -| PPO (torch) | `g1_motion_tracking_deploy` | Tested | Registered | -| PPO (torch) | `go1_joystick_rough` | Tested | Tested | -| PPO (torch) | `go2_arm_manip_loco` | Tested | - | -| PPO (torch) | `go2_footstand` | Tested | Tested | -| PPO (torch) | `go2w_joystick_flat` | Tested | Tested | -| PPO (torch) | `go2w_joystick_rough` | Tested | Tested | -| APPO (torch) | `go1_joystick_flat` | Tested | Registered | -| APPO (torch) | `go2_joystick_flat` | Tested | Registered | -| APPO (torch) | `g1_walk_flat` | Tested | Registered | -| APPO (torch) | `g1_motion_tracking` | Tested | Tested | -| APPO (torch) | `g1_flip_tracking` | Tested | Tested | -| APPO (torch) | `g1_wall_flip_tracking` | Tested | Tested | -| APPO (torch) | `allegro_inhand` | Tested | Tested | -| APPO (torch) | `sharpa_inhand` | Tested | Registered | -| APPO (torch) | `g1_climb_tracking` | Tested | Tested | -| SAC (torch) | `g1_walk_flat` | Tested | Tested | -| SAC (torch) | `g1_walk_rough` | Tested | Tested | -| SAC (torch) | `g1_motion_tracking` | Tested | Tested | -| SAC (torch) | `g1_wbt_obs` | Tested | Registered | -| TD3 (torch) | `go1_joystick_flat` | Registered | Tested | -| TD3 (torch) | `go2_joystick_flat` | Registered | Tested | -| TD3 (torch) | `g1_walk_flat` | Tested | Registered | -| FlashSAC (torch) | `go2_joystick_flat` | Tested | Registered | -| FlashSAC (torch) | `g1_walk_flat` | Tested | Registered | +| Entrypoint | Task owner | MuJoCo | mjwarp | Motrix | +| --- | --- | --- | --- | --- | +| PPO (torch) | `go1_joystick_flat` (Go1 joystick) | Tested | - | Tested | +| PPO (torch) | `go2_joystick_flat` (Go2 joystick) | Tested | - | Tested | +| PPO (torch) | `go2_joystick_rough` (Go2 joystick rough) | Tested | - | Tested | +| PPO (torch) | `g1_walk_flat` (G1 walk flat) | Tested | Tested | Tested | +| PPO (torch) | `g1_motion_tracking` (G1 motion tracking) | Tested | - | Tested | +| PPO (torch) | `g1_flip_tracking` (G1 flip tracking) | Tested | - | Tested | +| PPO (torch) | `g1_wall_flip_tracking` (G1 wall flip tracking) | Tested | - | Tested | +| PPO (torch) | `x2_wall_flip_tracking` (X2 wall flip tracking) | Tested | - | Tested | +| PPO (torch) | `allegro_inhand` (Allegro in-hand) | Tested | - | Tested | +| PPO (torch) | `sharpa_inhand` (Sharpa in-hand) | Tested | - | Tested | +| PPO (torch) | `sharpa_inhand_grasp` (Sharpa in-hand grasp) | Tested | - | Tested | +| PPO (torch) | `a2_joystick_flat` (a2 joystick flat) | Tested | - | - | +| PPO (torch) | `allegro_inhand_grasp` (allegro inhand grasp) | Tested | - | Tested | +| PPO (torch) | `g1_23dof_box_tracking` (g1 23dof box tracking) | Tested | - | Tested | +| PPO (torch) | `g1_23dof_climb_tracking` (g1 23dof climb tracking) | Tested | - | Tested | +| PPO (torch) | `g1_23dof_flip_tracking` (g1 23dof flip tracking) | Tested | - | Tested | +| PPO (torch) | `g1_23dof_motion_tracking` (g1 23dof motion tracking) | Tested | - | Tested | +| PPO (torch) | `g1_23dof_motion_tracking_deploy` (g1 23dof motion tracking deploy) | Tested | - | Tested | +| PPO (torch) | `g1_23dof_walk_flat` (g1 23dof walk flat) | Tested | - | Tested | +| PPO (torch) | `g1_23dof_walk_rough` (g1 23dof walk rough) | Tested | - | Registered | +| PPO (torch) | `g1_23dof_wall_flip_tracking` (g1 23dof wall flip tracking) | Tested | - | Tested | +| PPO (torch) | `g1_box_tracking` (g1 box tracking) | Tested | - | Tested | +| PPO (torch) | `g1_climb_tracking` (g1 climb tracking) | Tested | - | Tested | +| PPO (torch) | `g1_motion_tracking_deploy` (g1 motion tracking deploy) | Tested | - | Tested | +| PPO (torch) | `go1_joystick_rough` (go1 joystick rough) | Tested | - | Tested | +| PPO (torch) | `go2_arm_manip_loco` (go2 arm manip loco) | Tested | - | Tested | +| PPO (torch) | `go2_footstand` (go2 footstand) | Tested | - | Tested | +| PPO (torch) | `go2w_joystick_flat` (go2w joystick flat) | Tested | - | Tested | +| PPO (torch) | `go2w_joystick_rough` (go2w joystick rough) | Tested | - | Tested | +| PPO (torch) | `stewart_balance` (stewart balance) | Tested | - | Tested | +| APPO (torch) | `go1_joystick_flat` (Go1 joystick) | Tested | - | Tested | +| APPO (torch) | `go2_joystick_flat` (Go2 joystick) | Tested | - | Tested | +| APPO (torch) | `g1_walk_flat` (G1 walk flat) | Tested | Registered | Registered | +| APPO (torch) | `g1_motion_tracking` (G1 motion tracking) | Tested | - | Tested | +| APPO (torch) | `g1_flip_tracking` (G1 flip tracking) | Tested | - | Tested | +| APPO (torch) | `g1_wall_flip_tracking` (G1 wall flip tracking) | Tested | - | Tested | +| APPO (torch) | `allegro_inhand` (Allegro in-hand) | Tested | - | Tested | +| APPO (torch) | `sharpa_inhand` (Sharpa in-hand) | Tested | - | Tested | +| APPO (torch) | `g1_23dof_climb_tracking` (g1 23dof climb tracking) | Tested | - | Tested | +| APPO (torch) | `g1_23dof_flip_tracking` (g1 23dof flip tracking) | Tested | - | Tested | +| APPO (torch) | `g1_23dof_motion_tracking` (g1 23dof motion tracking) | Tested | - | Tested | +| APPO (torch) | `g1_23dof_walk_flat` (g1 23dof walk flat) | Tested | - | Registered | +| APPO (torch) | `g1_23dof_wall_flip_tracking` (g1 23dof wall flip tracking) | Tested | - | Tested | +| APPO (torch) | `g1_climb_tracking` (g1 climb tracking) | Tested | - | Tested | +| SAC (torch) | `g1_walk_flat` (G1 walk flat) | Tested | Tested | Tested | +| SAC (torch) | `g1_walk_rough` (G1 walk rough) | Tested | - | Tested | +| SAC (torch) | `g1_motion_tracking` (G1 motion tracking) | Tested | - | Tested | +| SAC (torch) | `g1_flip_tracking` (G1 flip tracking) | Tested | - | Registered | +| SAC (torch) | `g1_wall_flip_tracking` (G1 wall flip tracking) | Tested | - | Registered | +| SAC (torch) | `g1_23dof_flip_tracking` (g1 23dof flip tracking) | Tested | - | Registered | +| SAC (torch) | `g1_23dof_motion_tracking` (g1 23dof motion tracking) | Tested | - | Tested | +| SAC (torch) | `g1_23dof_walk_flat` (g1 23dof walk flat) | Tested | - | Tested | +| SAC (torch) | `g1_23dof_walk_rough` (g1 23dof walk rough) | Tested | - | Tested | +| SAC (torch) | `g1_23dof_wall_flip_tracking` (g1 23dof wall flip tracking) | Tested | - | Registered | +| SAC (torch) | `g1_23dof_wbt_obs` (g1 23dof wbt obs) | Tested | - | Registered | +| SAC (torch) | `g1_wbt_obs` (g1 wbt obs) | Tested | - | Registered | +| TD3 (torch) | `go1_joystick_flat` (Go1 joystick) | Registered | - | Tested | +| TD3 (torch) | `go2_joystick_flat` (Go2 joystick) | Registered | - | Tested | +| TD3 (torch) | `g1_walk_flat` (G1 walk flat) | Tested | Registered | Registered | +| TD3 (torch) | `g1_23dof_walk_flat` (g1 23dof walk flat) | Tested | - | Registered | +| FlashSAC (torch) | `go2_joystick_flat` (Go2 joystick) | Tested | - | Registered | +| FlashSAC (torch) | `g1_walk_flat` (G1 walk flat) | Tested | Configured | Tested | +| FlashSAC (torch) | `g1_23dof_walk_flat` (g1 23dof walk flat) | Tested | - | Tested | ## Source Index -- Registry bootstrap: `src/unilab/envs/**` registrations via +- Registry bootstrap: `src/unilab/envs/**` decorators via `unilab.base.registry.ensure_registries()`. - Owner YAML scan: `conf/ppo/task/**`, `conf/appo/task/**`, `conf/offpolicy/task/**`. - Generic compose coverage: `tests/config/test_config_system.py::test_supported_task_composes`. +- Validated mjwarp entrypoints are explicitly recorded in + `_MAINTAINER_VALIDATED_MJWARP_ENTRYPOINT_TASKS`; near-risk coverage lives in + `tests/base/test_mjwarp_backend.py`, + `tests/base/test_backend_conformance.py`, + `tests/base/test_mjwarp_differential.py`, and + `tests/base/test_mjwarp_playback.py`. diff --git a/docs/sphinx/source/zh_CN/2-user_guide/5-domain_randomization/0-index.md b/docs/sphinx/source/zh_CN/2-user_guide/5-domain_randomization/0-index.md index 46d6fb1d7..82bfcd39a 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/5-domain_randomization/0-index.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/5-domain_randomization/0-index.md @@ -1,9 +1,14 @@ # 域随机化 -本页仅描述仓库中那些已经注册、且已经接入 DR provider 的任务的当前状态。所有结论都来自代码；不从设计意图推断任何内容。 +本页仅描述仓库中已注册任务的域随机化现状。所有结论都来自代码；不从设计意图推断任何内容。 -当前统一的入口点位于 `NpEnv._init_domain_randomization()` 和 `DomainRandomizationManager`： +当前存在两条 DR 声明路径： + +- **Manager-Based（Compatible）任务**：reset / interval 随机化通过 owner YAML 中的 Hydra `events:` manager term 声明；reset 生命周期的 event 在 reset 时采样，interval 生命周期的 event 在 step 之间施加扰动。例如 `conf/ppo/task/go1_joystick_flat/base.yaml` 的 `events:` 段。 +- **legacy provider 路径**：只有 3 个 Adapted family（`sharpa_inhand` / `sharpa_inhand_grasp` / `go2_arm_manip_loco`，含 appo / hora / ppo_him owner）仍通过 `DomainRandomizationProvider` + `DomainRandomizationManager` 声明 `env.domain_rand.*` 配置。 + +legacy provider 路径的统一入口点位于 `NpEnv._init_domain_randomization()` 和 `DomainRandomizationManager`： - init 路径：task provider 产生一个 `InitRandomizationPlan`；manager 在 env 初始化期间调用后端的 `apply_init_randomization(...)` - reset 路径：task provider 产生一个 `ResetPlan`；manager 验证能力，然后调用后端的 `set_state(..., randomization=...)` @@ -17,34 +22,35 @@ ## 状态结论 -1. 当前所有接入 DR provider 的任务都使用统一的 DR 入口点；没有任何任务绕开 `DomainRandomizationManager` 在 `reset()` 内部运行单独的 DR 流程。 -2. 它们的结构都大致相同：legacy task owner 定义 `domain_rand` 配置 dataclass、`DomainRandomizationProvider` 和 `ResetPlan`；Manager-Based owner 则通过 Hydra command/event term 声明 reset 行为。G1 motion reset 扰动归 `MotionCommandCfg` 所有，WBT 另加 `EventTermCfg` reset 与 interval term。 -3. 今天所"统一"的主要是入口点和执行流程，而不是每一个随机化项本身。共享辅助函数 `build_common_reset_randomization()` 目前生成 `base_mass_delta`、`base_com_offset`、`gravity`、`kp`、`kd`；共享的 interval 辅助函数目前只生成 push。 +1. Manager-Based 任务不注册 DR provider；它们的 reset/interval 随机化是 owner YAML 中的 `events:` manager term，由 manager 生命周期统一执行。只有 Adapted family 的冻结兼容工厂仍走 `DomainRandomizationManager` 统一入口。 +2. Adapted family owner 定义 `domain_rand` 配置 dataclass、`DomainRandomizationProvider` 和 `ResetPlan`；Manager-Based owner 则通过 Hydra command/event term 声明 reset 行为。G1 motion reset 扰动归 `MotionCommandCfg` 所有，WBT 另加 `EventTermCfg` reset 与 interval term。 +3. 今天所"统一"的主要是入口点和执行流程，而不是每一个随机化项本身。legacy 路径的共享辅助函数 `build_common_reset_randomization()` 目前生成 `base_mass_delta`、`base_com_offset`、`gravity`、`kp`、`kd`；共享的 interval 辅助函数目前只生成 push。 4. `ResetRandomizationPayload` 已经可以表达 `gravity`、`body_iquat`、`body_inertia`、`kp`、`kd`，并且 `MuJoCoBackend` 已声明支持。这些是否实际被使用，仍取决于 task provider 是否对它们进行采样和 dispatch。 5. `MotrixBackend` 目前支持 `base_mass_delta`、`base_com_offset`、`kp`、`kd` 和 interval push；并且它要求在初始化期间所有模型 actuator 都是 position actuator。 6. `geom_size` 不是 reset 生命周期字段；Sharpa 手物体的 geom 缩放由 init 生命周期的模型 materialization 处理。 ## 统一性评估表 -| Task | 使用统一 DR 入口？ | 结构化形式？ | reset 形式 | interval 形式 | Code | +| Task | 声明路径 | 结构化形式？ | reset 形式 | interval 形式 | Code | | --- | --- | --- | --- | --- | --- | -| `Go1JoystickFlat` | 是 | 是：`Domain_Rand + Provider + ResetPlan` | task 状态采样 + common payload | push | `go1/joystick.py` | -| `Go2JoystickFlat` | 是 | 是：`Domain_Rand + Provider + ResetPlan` | task 状态采样 + common payload | push | `go2/joystick.py` | -| `G1WalkFlat` | 是 | 是：Hydra `EventTermCfg` + Manager-Based reset term | root-state reset + 经 `pd_gains` 的 kp/kd | 无 | `g1/manager_terms.py` | -| `G1WalkRough` | 是 | 是：与 `G1WalkFlat` 相同的 Manager-Based event term | root-state reset + 经 `pd_gains` 的 kp/kd | 无 | `g1/manager_terms.py` | -| `G1MotionTracking` | 是 | 是：Hydra `MotionCommandCfg` + Manager-Based command reset | motion frame、root pose/velocity 与 joint-position 采样 | 无 | `motion_tracking/common/manager_terms.py` | -| `G1WBTObs` | 是 | 是：同一 motion command + Hydra `EventTermCfg` | motion reset 加 mass/COM/PD/friction/encoder-bias event | interval velocity kick | `motion_tracking/g1/manager_terms.py` | -| `AllegroInhandRotation` | 是 | 是：Hydra `EventTermCfg` + Manager-Based reset term | entity 范围的手/球 reset | 无 | `allegro_inhand/manager_terms.py` | -| `AllegroInhandRotationGrasp` | 是 | 是：复用 rotation reset event + `RecorderTermCfg` | 带噪声的手部 reset + grasp 收集 | 无 | `allegro_inhand/grasp_gen.py` | -| `SharpaInhandRotation` | 是 | 是：`InitRandomizationPlan + ResetPlan + IntervalRandomizationPlan` | grasp cache 采样 + common payload | 物体 `body_force` | `sharpa_inhand/rotation.py` | -| `SharpaInhandRotationGrasp` | 是 | 是：复用 Sharpa rotation provider 并 override reset 采样 | grasp 收集 reset + common payload | 无 | `sharpa_inhand/grasp_gen.py` | +| `Go1JoystickFlat` | Hydra `events:` term | 是：owner YAML 声明 reset/interval event | root-state reset + base mass/COM + `pd_gains` | `push_by_setting_velocity` event | `conf/ppo/task/go1_joystick_flat/base.yaml` | +| `Go2JoystickFlat` | Hydra `events:` term | 是：owner YAML 声明 reset event | root-state reset + `pd_gains` kp/kd | 无 | `conf/ppo/task/go2_joystick_flat/base.yaml` | +| `G1WalkFlat` | Hydra `events:` term | 是：Hydra `EventTermCfg` + Manager-Based reset term | root-state reset + 经 `pd_gains` 的 kp/kd | 无 | `g1/manager_terms.py` | +| `G1WalkRough` | Hydra `events:` term | 是：与 `G1WalkFlat` 相同的 Manager-Based event term | root-state reset + 经 `pd_gains` 的 kp/kd | 无 | `g1/manager_terms.py` | +| `G1MotionTracking` | Hydra command term | 是：Hydra `MotionCommandCfg` + Manager-Based command reset | motion frame、root pose/velocity 与 joint-position 采样 | 无 | `motion_tracking/common/manager_terms.py` | +| `G1WBTObs` | Hydra `events:` term | 是：同一 motion command + Hydra `EventTermCfg` | motion reset 加 mass/COM/PD/friction/encoder-bias event | interval velocity kick | `motion_tracking/g1/manager_terms.py` | +| `AllegroInhandRotation` | Hydra `events:` term | 是：Hydra `EventTermCfg` + Manager-Based reset term | entity 范围的手/球 reset | 无 | `allegro_inhand/manager_terms.py` | +| `AllegroInhandRotationGrasp` | Hydra `events:` term | 是：复用 rotation reset event + `RecorderTermCfg` | 带噪声的手部 reset + grasp 收集 | 无 | `allegro_inhand/grasp_gen.py` | +| `SharpaInhandRotation` | legacy provider | 是：`InitRandomizationPlan + ResetPlan + IntervalRandomizationPlan` | grasp cache 采样 + common payload | 物体 `body_force` | `sharpa_inhand/rotation.py` | +| `SharpaInhandRotationGrasp` | legacy provider | 是：复用 Sharpa rotation provider 并 override reset 采样 | grasp 收集 reset + common payload | 无 | `sharpa_inhand/grasp_gen.py` | +| `Go2ArmManipLoco` | legacy provider | 是：`DomainRandConfig + LocomotionDRProvider 子类 + ResetPlan` | task 状态采样 + common payload | push | `go2_arm/manip_loco.py` | ## 各任务域随机化清单 | Task | 当前已实现的 reset 域随机化 | 当前已实现的 interval 域随机化 | 默认状态 | | --- | --- | --- | --- | -| `Go1JoystickFlat` | base xy；base yaw；base qvel；command 采样；`current_actions/last_actions` 清零；可选 `base_mass_delta`；可选 `base_com_offset`；可选 `gravity` | `push_robots` | `base_mass_delta`、`base_com_offset` 和 push 默认启用；`gravity` 默认禁用 | -| `Go2JoystickFlat` | base xy；base yaw；base qvel；command 采样；`current_actions/last_actions` 清零；kp/kd 随机化（默认启用）；可选 `base_mass_delta`；可选 `base_com_offset`；可选 `gravity` | `push_robots` | kp/kd 默认启用；common payload 和 push 默认禁用 | +| `Go1JoystickFlat` | 经 `reset_root_state_uniform` 的 base xy/yaw 与 base qvel；command 采样（`UniformVelocityCommandCfg`）；经 `randomize_rigid_body_mass` 的 base mass；经 `randomize_rigid_body_com` 的 base COM；经 `pd_gains` 的 kp/kd | `push_by_setting_velocity` interval event | 上述 event term 全部在 `conf/ppo/task/go1_joystick_flat/base.yaml` 中默认声明并启用 | +| `Go2JoystickFlat` | 经 `reset_root_state_uniform` 的 base xy/yaw 与 base qvel；command 采样；经 `pd_gains` 的 kp/kd | 无 | event term 在 `conf/ppo/task/go2_joystick_flat/base.yaml` 中默认声明并启用 | | `G1WalkFlat` | 经 `reset_root_state_uniform` 的 base xy/yaw 与 base qvel；带平面死区的 command 采样；`gait_phase` 采样；经 `pd_gains` 的 kp/kd 随机化 | 无 | mujoco owner 默认启用 kp/kd；motrix/mjwarp owner 默认禁用 | | `G1WalkRough` | 与 `G1WalkFlat` 相同（共享 owner base，rough 场景） | 无 | 与 `G1WalkFlat` 相同的默认值 | | `G1MotionTracking` | Motion-command frame 采样；root 位姿扰动 `x/y/z/roll/pitch/yaw`；root 速度扰动 `x/y/z/roll/pitch/yaw`；通过 public entity soft limit clip 的关节位置噪声；action-manager 状态 reset | 无 | base owner 中 `pose_range`、`velocity_range` 与 `joint_position_range` 默认有非零扰动 | @@ -56,26 +62,26 @@ ## 当前统一 DR 的能力与边界 -### 1. 统一入口点是完整的 +### 1. legacy provider 入口是统一的 -统一入口点由 `NpEnv` 和 `DomainRandomizationManager` 保证： +legacy provider 路径的统一入口点由 `NpEnv` 和 `DomainRandomizationManager` 保证： - 任务只需注册一个 provider - manager 统一执行能力验证 - 后端统一负责实际施加随机化 payload -因此从执行路径的角度看，这些任务已经是统一的。 +因此从执行路径的角度看，仍走该路径的 Adapted family 是统一的；Manager-Based 任务则由 manager 生命周期统一执行 owner YAML 声明的 `events:` term。 ### 2. 共享辅助函数仍然较窄 -`dr_utils.py` 目前只有两类共享辅助函数： +legacy 路径的 `dr_utils.py` 目前只有两类共享辅助函数： - reset common payload：`base_mass_delta`、`base_com_offset`、`gravity`、`kp`、`kd` - interval common payload：push 这意味着： -- 尽管运动控制任务都走统一入口点，但它们的 base xy、yaw、qvel、command 和 gait phase 仍然直接在各自的 provider 内部采样 +- 仍走 legacy provider 的 go2_arm / sharpa family，其 task 专属状态仍直接在各自的 provider 内部采样 - `G1MotionTracking` 的 pose / velocity / joint 噪声由其 manager command 所有 - Allegro 的 grasp / 物体初始状态采样完全是 task 专属逻辑 - Sharpa 的 `geom_size` 缩放是 init 生命周期的模型 materialization，不属于 reset common payload @@ -116,7 +122,7 @@ - 生命周期：仅在 reset 时采样和写入；env 会保留该重力，直到下一次 reset 重新采样。 - 后端：当前在 UniLab 中，只有 MuJoCo 后端声明支持该 reset 项；Motrix 后端不支持。一些任务按能力过滤并跳过它；另一些任务在 validate 阶段抛出错误。 -配置入口在每个任务的 `env.domain_rand` 下： +配置入口仅在仍走 legacy provider 路径的 Adapted family owner 的 `env.domain_rand` 下（`sharpa_inhand_grasp`、`go2_arm_manip_loco` 及对应 hora / appo / ppo_him 变体）；Manager-Based 任务没有 `env.domain_rand`： ```yaml env: @@ -136,7 +142,7 @@ env: 如果你只想随机化大小而保持竖直向下的方向，只开放 `z` 分量： ```bash -uv run train --algo ppo --task g1_motion_tracking --sim mujoco \ +uv run train --algo ppo --task sharpa_inhand_grasp --sim mujoco \ env.domain_rand.randomize_gravity=true \ 'env.domain_rand.gravity_range=[[0.0,0.0,-10.5],[0.0,0.0,-8.5]]' ``` @@ -144,7 +150,7 @@ uv run train --algo ppo --task g1_motion_tracking --sim mujoco \ 如果你想同时随机化方向和大小，开放 `x/y/z`： ```bash -uv run train --algo ppo --task g1_motion_tracking --sim mujoco \ +uv run train --algo ppo --task sharpa_inhand_grasp --sim mujoco \ env.domain_rand.randomize_gravity=true \ 'env.domain_rand.gravity_range=[[-0.3,-0.3,-10.5],[0.3,0.3,-8.5]]' ``` @@ -159,7 +165,9 @@ uv run train --algo ppo --task g1_motion_tracking --sim mujoco \ ## Interval push 用法 -支持 interval push 的任务在 `env.domain_rand` 下配置它： +`env.domain_rand.push_robots` 系列字段只存在于 go2_arm Adapted family 的 owner（`conf/ppo/task/go2_arm_manip_loco/mujoco.yaml` 等）；Manager-Based 任务改用 `push_by_setting_velocity` interval event term 声明 push（例如 `conf/ppo/task/go1_joystick_flat/base.yaml` 和 `conf/ppo/task/quadruped_joystick_rough/base.yaml`）。 + +go2_arm owner 在 `env.domain_rand` 下配置 push： ```yaml env: @@ -176,11 +184,11 @@ env: - `push_body_name`：施加力的目标 body / link。默认为 `null`，表示使用后端的 `base_name`。 ```bash -uv run train --algo ppo --task g1_motion_tracking --sim mujoco \ +uv run train --algo ppo --task go2_arm_manip_loco --sim mujoco \ env.domain_rand.push_robots=true \ env.domain_rand.push_interval=500 \ 'env.domain_rand.max_force=[20.0,20.0,5.0]' \ - env.domain_rand.push_body_name=torso_link + env.domain_rand.push_body_name=base ``` 说明： diff --git a/docs/sphinx/source/zh_CN/2-user_guide/5-domain_randomization/1-configuration.md b/docs/sphinx/source/zh_CN/2-user_guide/5-domain_randomization/1-configuration.md index 6bcda32d0..2a3033b33 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/5-domain_randomization/1-configuration.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/5-domain_randomization/1-configuration.md @@ -1,11 +1,17 @@ # 配置 -域随机化在所选的 task owner YAML 内部配置，通常位于 -`env.domain_rand` 下。先使用 `--task` 和 `--sim` 选择后端专属行为， +域随机化在所选的 task owner YAML 内部配置。先使用 `--task` 和 `--sim` 选择后端专属行为， 然后在所选的 owner 内部 override 字段。 +当前有两条声明路径： + +- Manager-Based（Compatible）任务通过 owner YAML 的 `events:` manager term 声明 + reset / interval 随机化，例如 `conf/ppo/task/go1_joystick_flat/base.yaml`。 +- 只有 Adapted family（sharpa / go2_arm 及对应 hora / appo / ppo_him owner）仍在 + `env.domain_rand` 下配置 legacy provider 字段。 + ```bash -uv run train --algo ppo --task g1_motion_tracking --sim mujoco \ +uv run train --algo ppo --task sharpa_inhand_grasp --sim mujoco \ env.domain_rand.randomize_gravity=true \ 'env.domain_rand.gravity_range=[[0.0,0.0,-10.5],[0.0,0.0,-8.5]]' ``` @@ -18,25 +24,29 @@ uv run train --algo ppo --task g1_motion_tracking --sim mujoco \ 详细的任务状态和字段语义见 {doc}`0-index`。 -域随机化按生命周期划分：init、reset 和 interval。manager 路径是 +域随机化按生命周期划分：init、reset 和 interval。legacy 路径的 manager 位于 `src/unilab/dr/manager.py`；task provider 位于 env owner 附近， 后端能力通过 `src/unilab/base/backend/base.py` 声明。 ## Reset Gravity 在启用 gravity reset 随机化时使用 `--sim mujoco`；Motrix 在当前后端中 -未提供相同的 gravity 能力。 +未提供相同的 gravity 能力。该项只在 legacy provider 路径（Adapted family owner） +上可用。 ```bash -uv run train --algo ppo --task g1_motion_tracking --sim mujoco \ +uv run train --algo ppo --task sharpa_inhand_grasp --sim mujoco \ env.domain_rand.randomize_gravity=true \ 'env.domain_rand.gravity_range=[[0.0,0.0,-10.5],[0.0,0.0,-8.5]]' ``` ## Interval Push +Manager-Based 任务通过 `push_by_setting_velocity` interval event term 声明 push； +`env.domain_rand.push_robots` 只在 go2_arm Adapted family owner 上可用。 + ```bash -uv run train --algo ppo --task g1_motion_tracking --sim mujoco \ +uv run train --algo ppo --task go2_arm_manip_loco --sim mujoco \ env.domain_rand.push_robots=true \ env.domain_rand.push_interval=500 \ 'env.domain_rand.max_force=[20.0,20.0,5.0]' @@ -45,8 +55,9 @@ uv run train --algo ppo --task g1_motion_tracking --sim mujoco \ ## Owner 本地默认值 当取值范围是任务 contract 的一部分时，将其保留在 task owner YAML 中。例如， -`conf/ppo/task/go2_joystick_rough/mujoco.yaml` 启用了 base mass、 -质心、kp/kd 和 push 随机化，而 +rough 四足家族的 base mass、质心、kp/kd 和 push 随机化作为 event term 声明在共享 base +`conf/ppo/task/quadruped_joystick_rough/base.yaml`（`go2_joystick_rough` 的 backend +owner 通过 Hydra defaults 组合它），而 `conf/ppo/task/sharpa_inhand/mujoco.yaml` 为 Sharpa 配置了物体缩放、摩擦和 力扰动。 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 e05b06cde..901600ee3 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 @@ -1,5 +1,11 @@ # 编写 Provider +本页描述 legacy provider 路径：只有 3 个 Adapted family（`sharpa_inhand` / +`sharpa_inhand_grasp` / `go2_arm_manip_loco`）仍通过任务级 +`DomainRandomizationProvider` 声明域随机化。已迁移的 Manager-Based 任务不写 +provider；它们在 owner YAML 中通过 Hydra `events:` manager term 声明随机化（见 +{doc}`0-index` 与 {doc}`1-configuration`）。 + 任务级域随机化 provider 与 task env owner 放在一起。它们采样任务专属的 状态，并返回由 `DomainRandomizationManager` 消费的 plan。 @@ -24,9 +30,10 @@ ## 证据 -具有代表性的 provider 实现位于： +具有代表性的 provider 实现位于（全部属于 Adapted family 的兼容路径）： -- `src/unilab/tasks/locomotion/common/dr_provider.py` +- `src/unilab/tasks/locomotion/common/dr_provider.py`（`LocomotionDRProvider`， + 由 go2_arm family 使用） - `src/unilab/tasks/locomotion/go2_arm/manip_loco.py` - `src/unilab/tasks/manipulation/sharpa_inhand/rotation.py` diff --git a/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/1-from_isaac_lab.md b/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/1-from_isaac_lab.md index ea55f2f58..876fbf489 100644 --- a/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/1-from_isaac_lab.md +++ b/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/1-from_isaac_lab.md @@ -133,12 +133,34 @@ uv run train --algo ppo --task --sim mujoco 不要增加 task-specific 训练脚本分支、env factory、runner 或 IPC 路径。 +generic factory 规则只有两个 maintainer 批准的已注册例外： +`make_g1_walk_env`（`src/unilab/tasks/locomotion/g1/manager_terms.py`）构造 +`G1WalkManagerBasedEnv` 子类，提供 `NpEnv` 的 `build_symmetry_augmentation` +对称性 hook；`make_x2_wall_flip_env` +（`src/unilab/tasks/motion_tracking/x2/__init__.py`）在冷路径解析未跟踪的 X2 +mesh，然后委托给 `make_manager_based_rl_env`。其余所有 Compatible task 都直接注册 +`make_manager_based_rl_env`。 + ### 6. 在适配风险附近验证 测试 Hydra compose 与 typed materialization、term 顺序与数学、selector 失败、 observation/action shape、局部 reset，以及至少一个真实已注册 backend 的 transition。行为 应与固定来源 task 对比；完成语义迁移后再做性能 benchmark。 +## 任务迁移最终状态 + +#1042 迁移收尾覆盖 39 个 production task、86 个 task/backend 注册。fail-closed 的 +source of truth 是 `src/unilab/tasks/migration_matrix.py`：`migration_record()` 对没有 +entry 的 production task 名称抛出 `KeyError`，因此新增 production 注册必须显式做出 +迁移决策。 + +- 36 个 task 为 **Compatible**（`target=complete`）：Hydra owner YAML 物化 canonical + NumPy Manager-Based runtime。 +- 3 个 task 为 **Adapted**（`target=compatibility`）：`Go2ArmManipLoco`、 + `SharpaInhandRotation` 和 `SharpaInhandRotationGrasp` 各自把自定义 IK/history 或 + tactile/contact/cache 行为保留在一个冻结的兼容 factory 后面；只有当正式能力存在时 + 才迁移。 + ## 仓库证据 `tests/fixtures/isaac_lab_cartpole/` 迁移了 Isaac Lab commit diff --git a/src/unilab/base/curriculum.py b/src/unilab/base/curriculum.py index 245f702b2..2b87f3968 100644 --- a/src/unilab/base/curriculum.py +++ b/src/unilab/base/curriculum.py @@ -2,8 +2,6 @@ from __future__ import annotations -from typing import Any - import numpy as np @@ -22,69 +20,3 @@ def update(self, episode_lengths: np.ndarray) -> None: current_avg = float(np.mean(episode_lengths)) weight = min(len(episode_lengths) / self.window_size, 1.0) self.average_length = self.average_length * (1 - weight) + current_avg * weight - - -class PenaltyCurriculum: - """Adaptive penalty scaling based on episode length.""" - - def __init__( - self, - env: Any, - enabled: bool = True, - initial_scale: float = 0.5, - min_scale: float = 0.5, - max_scale: float = 1.0, - level_down_threshold: float = 150.0, - level_up_threshold: float = 750.0, - degree: float = 0.001, - ): - self.env = env - self.enabled = enabled - self.current_scale = initial_scale - self.min_scale = min_scale - self.max_scale = max_scale - self.level_down_threshold = level_down_threshold - self.level_up_threshold = level_up_threshold - self.degree = degree - - # Store original penalty weights - self.penalty_names: list[str] = [] - self.original_weights: dict[str, float] = {} - - if enabled: - self._identify_penalties() - self._apply_initial_scale() - - def _identify_penalties(self) -> None: - """Identify penalty rewards (negative scales).""" - for name, scale in self.env.cfg.reward_config.scales.items(): - if scale < 0: - self.penalty_names.append(name) - self.original_weights[name] = scale - - def _apply_initial_scale(self) -> None: - """Apply initial penalty scaling.""" - for name in self.penalty_names: - self.env.cfg.reward_config.scales[name] = ( - self.original_weights[name] * self.current_scale - ) - - def update(self, average_episode_length: float) -> None: - """Update penalty scale based on average episode length.""" - if not self.enabled: - return - - # Adjust scale - if average_episode_length < self.level_down_threshold: - self.current_scale *= 1.0 - self.degree - elif average_episode_length > self.level_up_threshold: - self.current_scale *= 1.0 + self.degree - - # Clamp - self.current_scale = float(np.clip(self.current_scale, self.min_scale, self.max_scale)) - - # Apply to all penalty rewards - for name in self.penalty_names: - self.env.cfg.reward_config.scales[name] = ( - self.original_weights[name] * self.current_scale - ) diff --git a/src/unilab/tasks/locomotion/common/__init__.py b/src/unilab/tasks/locomotion/common/__init__.py index cc78aac1f..67b9c7141 100644 --- a/src/unilab/tasks/locomotion/common/__init__.py +++ b/src/unilab/tasks/locomotion/common/__init__.py @@ -5,16 +5,9 @@ ControlConfigBase, LocomotionBaseCfg, LocomotionBaseEnv, - PdControlConfig, Sensor, ) -from .commands import ( - Commands, - apply_heading_yaw_feedback, - sample_heading_commands, - sample_velocity_commands, - zero_small_xy_commands, -) +from .commands import Commands from .domain_rand import DomainRandConfig from .dr_provider import LocomotionDRProvider from .height_scan import ( @@ -35,11 +28,6 @@ "LocomotionDRProvider", "LocomotionBaseCfg", "LocomotionBaseEnv", - "PdControlConfig", "RewardContext", "Sensor", - "apply_heading_yaw_feedback", - "sample_heading_commands", - "sample_velocity_commands", - "zero_small_xy_commands", ] diff --git a/src/unilab/tasks/locomotion/common/base.py b/src/unilab/tasks/locomotion/common/base.py index 68d82754b..34508bfbf 100644 --- a/src/unilab/tasks/locomotion/common/base.py +++ b/src/unilab/tasks/locomotion/common/base.py @@ -27,23 +27,6 @@ class ControlConfigBase: simulate_action_latency: bool = False -@dataclass -class PdControlConfig(ControlConfigBase): - """``ControlConfigBase`` plus shared PD-actuator gains (Go1/Go2/Go2W defaults).""" - - Kp: float = 35.0 - Kd: float = 0.5 - - def position_gains(self) -> dict[str, float | np.ndarray]: - """Gains forwarded to the backend's ``position_actuator_gains``. - - Default applies the scalar ``Kp``/``Kd`` uniformly to every actuator. - Subclasses with per-joint gains (e.g. A2's stronger calf) override this - to return per-actuator arrays in model/actuator order. - """ - return {"kp": self.Kp, "kd": self.Kd} - - @dataclass class BaseNoiseConfig: level: float = 0.0 diff --git a/src/unilab/tasks/locomotion/common/commands.py b/src/unilab/tasks/locomotion/common/commands.py index 08b85fcaa..92a63ff41 100644 --- a/src/unilab/tasks/locomotion/common/commands.py +++ b/src/unilab/tasks/locomotion/common/commands.py @@ -3,12 +3,6 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import Any - -import numpy as np - -from unilab.dtype_config import get_global_dtype -from unilab.utils.rotation import np_wrap_to_pi, np_yaw_from_quat @dataclass @@ -24,39 +18,3 @@ class Commands: heading_range: list[float] = field(default_factory=lambda: [-3.14, 3.14]) heading_control_stiffness: float = 0.5 rel_standing_envs: float = 0.0 - - -def sample_velocity_commands( - rng: np.random.Generator, num_samples: int, low: np.ndarray, high: np.ndarray -) -> np.ndarray: - return np.asarray( - rng.uniform(low=low, high=high, size=(num_samples, 3)), dtype=get_global_dtype() - ) - - -def zero_small_xy_commands(commands: np.ndarray, *, threshold: float = 0.2) -> None: - """Zero ``commands[:, :2]`` in-place wherever its norm is below ``threshold``.""" - moving = np.linalg.norm(commands[:, :2], axis=1) > threshold - commands[:, :2] *= moving[:, None] - - -def sample_heading_commands(env: Any, num_samples: int) -> np.ndarray: - """Uniformly sample heading targets from ``env.cfg.commands.heading_range``.""" - heading_range = np.asarray(env.cfg.commands.heading_range, dtype=get_global_dtype()) - if heading_range.shape != (2,): - raise ValueError(f"commands.heading_range must have shape (2,), got {heading_range.shape}") - low, high = float(np.min(heading_range)), float(np.max(heading_range)) - return np.asarray(np.random.uniform(low, high, size=(num_samples,)), dtype=get_global_dtype()) - - -def apply_heading_yaw_feedback( - commands: np.ndarray, - base_quat: np.ndarray, - heading_commands: np.ndarray, - *, - stiffness: float, - clip: float = 2.0, -) -> None: - """In-place P-control on heading error → ``commands[:, 2]`` (yaw rate).""" - heading = np_yaw_from_quat(base_quat) - commands[:, 2] = np.clip(stiffness * np_wrap_to_pi(heading_commands - heading), -clip, clip) diff --git a/src/unilab/tasks/locomotion/common/rewards.py b/src/unilab/tasks/locomotion/common/rewards.py index a147a52be..bd127ed66 100644 --- a/src/unilab/tasks/locomotion/common/rewards.py +++ b/src/unilab/tasks/locomotion/common/rewards.py @@ -9,9 +9,7 @@ from __future__ import annotations -from collections.abc import Callable, Mapping from dataclasses import dataclass, field -from typing import Any import numpy as np @@ -67,23 +65,6 @@ def tracking_ang_vel(ctx: RewardContext) -> np.ndarray: return np.exp(-ang_vel_error / ctx.tracking_sigma) # type: ignore[no-any-return] -def forward_progress(ctx: RewardContext) -> np.ndarray: - """Reward for forward progress relative to commanded speed.""" - commands = ctx.info["commands"] - commanded_speed = np.maximum(commands[:, 0], 1e-6) - forward_speed = np.maximum(ctx.linvel[:, 0], 0.0) - return np.asarray(np.minimum(forward_speed / commanded_speed, 1.0), dtype=get_global_dtype()) - - -def under_speed(ctx: RewardContext) -> np.ndarray: - """Penalty for being below commanded forward speed.""" - commands = ctx.info["commands"] - commanded_speed = np.maximum(commands[:, 0], 1e-6) - forward_speed = np.maximum(ctx.linvel[:, 0], 0.0) - gap = np.maximum(commands[:, 0] - forward_speed, 0.0) - return np.asarray(gap / commanded_speed, dtype=get_global_dtype()) - - # ── velocity / orientation penalties ───────────────────────────────── @@ -97,13 +78,6 @@ def ang_vel_xy(ctx: RewardContext) -> np.ndarray: return np.sum(np.square(ctx.gyro[:, :2]), axis=1) # type: ignore[no-any-return] -def orientation(ctx: RewardContext) -> np.ndarray: - """Penalty for deviation from upright orientation (roll/pitch).""" - g = ctx.gravity - assert g is not None - return np.square(g[:, 0]) + np.square(g[:, 1]) # type: ignore[no-any-return] - - def roll(ctx: RewardContext) -> np.ndarray: """Penalty for deviation from roll orientation.""" g = ctx.gravity @@ -111,14 +85,6 @@ def roll(ctx: RewardContext) -> np.ndarray: return np.square(g[:, 0]) # type: ignore[no-any-return] -def upright(ctx: RewardContext) -> np.ndarray: - """Exponential reward for upright orientation.""" - g = ctx.gravity - assert g is not None - xy_squared = np.sum(np.square(g[:, :2]), axis=1) - return np.exp(-xy_squared / 0.25) # type: ignore[no-any-return] - - # ── height / pose penalties ────────────────────────────────────────── @@ -149,14 +115,6 @@ def action_rate(ctx: RewardContext) -> np.ndarray: return np.sum(np.square(current - last), axis=1) # type: ignore[no-any-return] -def action_smooth(ctx: RewardContext) -> np.ndarray: - """Penalty for second-order action changes between timesteps.""" - current = ctx.info["current_actions"] - last = ctx.info["last_actions"] - previous = ctx.info.get("previous_actions", last) - return np.sum(np.square(current - 2.0 * last + previous), axis=1) # type: ignore[no-any-return] - - # ── effort penalties ───────────────────────────────────────────────── @@ -190,185 +148,3 @@ def dof_acc(ctx: RewardContext) -> np.ndarray: def alive(ctx: RewardContext) -> np.ndarray: """Constant reward for staying alive.""" return np.ones((ctx.num_envs,), dtype=get_global_dtype()) - - -# ── quadruped-rough helpers / penalties ────────────────────────────── - - -def upright_scale(gravity: np.ndarray | None, num_envs: int) -> np.ndarray: - """Scalar gate in [0, 1] from the body-up projection of gravity. - - Used by quadruped rough tasks to suppress reward / penalty bookkeeping - while the robot is tipping over. Returns 1.0 when the body is upright - (gravity[:, 2] >= 0.7) and 0.0 when fully tipped. - """ - if gravity is None: - return np.ones((num_envs,), dtype=get_global_dtype()) - return np.asarray(np.clip(gravity[:, 2], 0.0, 0.7) / 0.7, dtype=get_global_dtype()) - - -def dof_torques_l2(ctx: RewardContext) -> np.ndarray: - """Penalty for joint torque magnitude (L2).""" - torques = np.asarray( - ctx.info.get("torques", np.zeros((ctx.num_envs, ctx.dof_pos.shape[1]))), - dtype=get_global_dtype(), - ) - return np.asarray(np.sum(np.square(torques), axis=1), dtype=get_global_dtype()) - - -def dof_acc_l2(ctx: RewardContext) -> np.ndarray: - """Penalty for joint acceleration magnitude (L2).""" - qacc = np.asarray( - ctx.info.get("qacc", np.zeros((ctx.num_envs, ctx.dof_pos.shape[1]))), - dtype=get_global_dtype(), - ) - return np.asarray(np.sum(np.square(qacc), axis=1), dtype=get_global_dtype()) - - -def joint_pos_limits(ctx: RewardContext) -> np.ndarray: - """Penalty for joint position over/under-shoot relative to backend limits.""" - if ctx.joint_range is None: - return np.zeros((ctx.num_envs,), dtype=get_global_dtype()) - lower = ctx.joint_range[:, 0] - upper = ctx.joint_range[:, 1] - low_error = np.clip(lower - ctx.dof_pos, 0.0, None) - high_error = np.clip(ctx.dof_pos - upper, 0.0, None) - return np.asarray(np.sum(low_error + high_error, axis=1), dtype=get_global_dtype()) - - -def joint_power(ctx: RewardContext) -> np.ndarray: - """Penalty for joint mechanical power (|tau * dq|).""" - assert ctx.dof_vel is not None - torques = np.asarray( - ctx.info.get("torques", np.zeros((ctx.num_envs, ctx.dof_pos.shape[1]))), - dtype=get_global_dtype(), - ) - return np.asarray(np.sum(np.abs(ctx.dof_vel * torques), axis=1), dtype=get_global_dtype()) - - -def stand_still(ctx: RewardContext, command_threshold: float = 0.1) -> np.ndarray: - """Penalty for joint deviation from default while command norm is below threshold.""" - stopped = np.linalg.norm(ctx.info["commands"], axis=1) < command_threshold - dof_error = np.sum(np.abs(ctx.dof_pos - ctx.default_angles), axis=1) - return np.asarray(dof_error * stopped, dtype=get_global_dtype()) - - -def joint_pos_penalty( - ctx: RewardContext, - *, - stand_still_scale: float = 5.0, - velocity_threshold: float = 0.5, - command_threshold: float = 0.1, -) -> np.ndarray: - """Penalty for joint deviation that switches scale based on command/body motion.""" - command_norm = np.linalg.norm(ctx.info["commands"], axis=1) - body_vel = np.linalg.norm(ctx.linvel[:, :2], axis=1) - running_error = np.linalg.norm(ctx.dof_pos - ctx.default_angles, axis=1) - moving = (command_norm > command_threshold) | (body_vel > velocity_threshold) - return np.asarray( - np.where(moving, running_error, stand_still_scale * running_error), - dtype=get_global_dtype(), - ) - - -def upward(ctx: RewardContext) -> np.ndarray: - """Reward favouring an upright body (no Go2 upright gate).""" - assert ctx.gravity is not None - return np.asarray(np.square(1.0 + ctx.gravity[:, 2]), dtype=get_global_dtype()) - - -# ── biped-style rewards ───────────────────── - - -def track_lin_vel_xy_yaw_frame_exp(ctx: RewardContext) -> np.ndarray: - """Exponential tracking of xy linear velocity in the gravity-aligned yaw frame. - - Requires ``ctx.linvel_yaw`` (base linvel rotated into yaw frame). - """ - linvel = ctx.linvel_yaw if ctx.linvel_yaw is not None else ctx.linvel - commands = ctx.info["commands"] - lin_vel_error = np.sum(np.square(commands[:, :2] - linvel[:, :2]), axis=1) - return np.asarray(np.exp(-lin_vel_error / ctx.tracking_sigma), dtype=get_global_dtype()) - - -def track_ang_vel_z_world_exp(ctx: RewardContext) -> np.ndarray: - """Exponential tracking of yaw angular velocity (world frame).""" - commands = ctx.info["commands"] - ang_vel_error = np.square(commands[:, 2] - ctx.gyro[:, 2]) - return np.asarray(np.exp(-ang_vel_error / ctx.tracking_sigma), dtype=get_global_dtype()) - - -def feet_air_time_positive_biped( - ctx: RewardContext, - *, - threshold: float = 0.4, - command_threshold: float = 0.1, -) -> np.ndarray: - """Biped foot air-time reward: only rewards single-stance phase. - - Reads ``ctx.info`` keys ``current_air_time``, ``current_contact_time`` (each - shape (N, 2)); the environment populates them per step. - """ - air = np.asarray( - ctx.info.get("current_air_time", np.zeros((ctx.num_envs, 2))), dtype=get_global_dtype() - ) - contact = np.asarray( - ctx.info.get("current_contact_time", np.zeros((ctx.num_envs, 2))), dtype=get_global_dtype() - ) - in_contact = contact > 0.0 - in_mode_time = np.where(in_contact, contact, air) - single_stance = np.sum(in_contact.astype(np.int32), axis=1) == 1 - masked = np.where(single_stance[:, None], in_mode_time, 0.0) - reward = np.min(masked, axis=1) - reward = np.clip(reward, None, threshold) - moving = np.linalg.norm(ctx.info["commands"][:, :2], axis=1) > command_threshold - return np.asarray(reward * moving, dtype=get_global_dtype()) - - -def joint_deviation_l1(ctx: RewardContext, joint_indices: np.ndarray | None = None) -> np.ndarray: - """L1 penalty for joints deviating from their default positions.""" - diff = ctx.dof_pos - ctx.default_angles - if joint_indices is not None: - diff = diff[:, joint_indices] - return np.asarray(np.sum(np.abs(diff), axis=1), dtype=get_global_dtype()) - - -# ── reward dispatch ────────────────────────────────────────────────── - - -def run_reward_dispatch( - *, - scales: Mapping[str, float], - fns: Mapping[str, Callable[[RewardContext], np.ndarray]], - ctx: RewardContext, - info: dict[str, Any], - enable_log: bool, - ctrl_dt: float, - log_every_n_steps: int = 4, - only_positive: bool = False, -) -> np.ndarray: - """Standard ``scales × fns(ctx)`` reduction shared by all locomotion envs. - - - Writes per-reward means into ``info["log"]`` when ``enable_log`` and the - ``steps[0]`` cadence matches ``log_every_n_steps``. - - Returns ``reward * ctrl_dt`` (with optional positive clamp). - """ - dtype = get_global_dtype() - reward = np.zeros((ctx.num_envs,), dtype=dtype) - step_count = info.get("steps", np.zeros((ctx.num_envs,), dtype=np.uint32)) - should_log = enable_log and (int(step_count[0]) % log_every_n_steps == 0) - log = {} if should_log else info.get("log", {}) - - for name, scale in scales.items(): - if scale == 0 or name not in fns: - continue - rew = fns[name](ctx) - weighted_rew = rew * scale - reward += weighted_rew - if should_log: - log[f"reward/{name}"] = float(np.mean(weighted_rew)) - - info["log"] = log - if only_positive: - np.maximum(reward, 0.0, out=reward) - return reward * ctrl_dt diff --git a/src/unilab/tasks/locomotion/g1/manager_terms.py b/src/unilab/tasks/locomotion/g1/manager_terms.py index e8756af14..cba15584d 100644 --- a/src/unilab/tasks/locomotion/g1/manager_terms.py +++ b/src/unilab/tasks/locomotion/g1/manager_terms.py @@ -795,7 +795,7 @@ def _resample_command(self, env_ids: np.ndarray) -> None: # --------------------------------------------------------------------------- -# Penalty curriculum (ports EpisodeLengthTracker + PenaltyCurriculum semantics) +# Penalty curriculum (ports the legacy G1 penalty-curriculum semantics) # --------------------------------------------------------------------------- diff --git a/tests/conftest.py b/tests/conftest.py index f2a9abda2..5fbc4ecc4 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -124,42 +124,3 @@ def tiny_weight_shapes(): @pytest.fixture def mock_env_name() -> str: return _DUMMY_ENV_NAME - - -@pytest.fixture -def default_go1_reward_config(): - """Default reward config for Go1 testing.""" - return { - "scales": { - "tracking_lin_vel": 1.0, - "tracking_ang_vel": 0.2, - "lin_vel_z": -5.0, - "ang_vel_xy": -0.1, - "base_height": -100.0, - "action_rate": -0.005, - "similar_to_default": -0.1, - "contact": 0.24, - }, - "tracking_sigma": 0.25, - "base_height_target": 0.3, - } - - -@pytest.fixture -def default_go2_reward_config(): - """Default reward config for Go2 testing.""" - return { - "scales": { - "tracking_lin_vel": 1.0, - "tracking_ang_vel": 0.2, - "lin_vel_z": -5.0, - "ang_vel_xy": -0.1, - "base_height": -100.0, - "action_rate": -0.005, - "similar_to_default": -0.1, - "contact": 0.24, - "swing_feet_z": 4.0, - }, - "tracking_sigma": 0.25, - "base_height_target": 0.3, - } diff --git a/tests/envs/test_env_configs.py b/tests/envs/test_env_configs.py index 0f6710f0d..a0a7ee3d6 100644 --- a/tests/envs/test_env_configs.py +++ b/tests/envs/test_env_configs.py @@ -789,11 +789,7 @@ def test_sharpa_run_complete_reports_effective_collection_target() -> None: @pytest.mark.parametrize("env_name", _STANDARD_ENVS) -def test_env_reset_and_step( - env_name: str, - default_go1_reward_config, - default_go2_reward_config, -): +def test_env_reset_and_step(env_name: str): """Every registered env must be constructible, resetable, and steppable. Verifies: @@ -807,11 +803,7 @@ def test_env_reset_and_step( # Provide config overrides for envs that require them via Hydra env_cfg_override = None - if "Go1" in env_name: - env_cfg_override = {"reward_config": default_go1_reward_config} - elif "Go2" in env_name: - env_cfg_override = {"reward_config": default_go2_reward_config} - elif env_name == "G1WalkFlat": + if env_name == "G1WalkFlat": env_cfg_override = _g1_manager_override("g1_walk_flat") elif env_name == "G1WalkRough": env_cfg_override = _g1_manager_override("g1_walk_rough") diff --git a/tests/tasks/test_production_registry_closeout.py b/tests/tasks/test_production_registry_closeout.py new file mode 100644 index 000000000..28c53a389 --- /dev/null +++ b/tests/tasks/test_production_registry_closeout.py @@ -0,0 +1,167 @@ +"""Repository boundary tests closing out the Manager-Based migration (#1042). + +Pin the post-migration production registry so no legacy fallback or dual +registration can come back: + +- the production registry matches the #1042 migration matrix exactly, +- only the three approved families use the frozen LegacyFactoryAdapter seam, +- every other registered factory is one of the canonical manager-runtime + callables (generic factory plus the two maintainer-approved wrappers), +- the deleted ``unilab.envs.{locomotion,manipulation,motion_tracking}`` + packages stay removed (``unilab.envs`` itself remains: it owns the + manager-based runtime such as ``manager_based_rl_env`` and ``mdp``). + +Scope note: the registry has no unregister API and no provenance tracking, and +the pytest session pollutes it with fixture-only envs (``DummyFlatTest`` via +``UNILAB_EXTRA_REGISTRY_PACKAGES``, the cartpole fixtures reusing +``ManagerBasedRlEnvCfg``/``make_manager_based_rl_env``). The registry snapshot +is therefore taken in a fresh subprocess with that env var scrubbed +(``tests/base/test_backend_imports.py`` idiom), so only the production +``unilab.tasks`` bootstrap contributes registrations. +""" + +from __future__ import annotations + +import importlib.util +import json +import os +import subprocess +import sys +import textwrap + +from unilab.tasks.migration_matrix import PRODUCTION_TASK_NAMES + +APPROVED_COMPATIBILITY_TASKS = { + "Go2ArmManipLoco", + "SharpaInhandRotation", + "SharpaInhandRotationGrasp", +} + +ADAPTER_FACTORY = ("unilab.tasks.compatibility", "LegacyFactoryAdapter") + +CANONICAL_MANAGER_RUNTIME_FACTORIES = ( + ("unilab.envs.manager_based_rl_env", "make_manager_based_rl_env"), + # Approved wrapper: G1WalkManagerBasedEnv subclass providing the NpEnv + # build_symmetry_augmentation hook. + ("unilab.tasks.locomotion.g1.manager_terms", "make_g1_walk_env"), + # Approved wrapper: cold-path untracked X2 mesh resolution before + # delegating to the generic factory. + ("unilab.tasks.motion_tracking.x2", "make_x2_wall_flip_env"), +) + +REMOVED_LEGACY_ENV_PACKAGES = ( + "unilab.envs.locomotion", + "unilab.envs.manipulation", + "unilab.envs.motion_tracking", +) + +_SNAPSHOT_CODE = textwrap.dedent( + """ + import json + + from unilab.base import registry + + registry.ensure_registries() + snapshot = { + name: { + backend: [ + getattr(factory, "__module__", None) or type(factory).__module__, + getattr(factory, "__qualname__", None) or type(factory).__qualname__, + ] + for backend, factory in meta.env_factory_dict.items() + } + for name, meta in registry._envs.items() + } + print(json.dumps(snapshot)) + """ +) + +_snapshot_cache: dict[str, dict[str, tuple[str, str]]] | None = None + + +def _production_factories() -> dict[str, dict[str, tuple[str, str]]]: + """Snapshot the production registry in a clean interpreter. + + Returns ``{task: {backend: (factory_module, factory_qualname)}}``. The + subprocess scrubs ``UNILAB_EXTRA_REGISTRY_PACKAGES`` so fixture-only test + envs injected by ``tests/conftest.py`` cannot leak into the snapshot. + """ + global _snapshot_cache + if _snapshot_cache is None: + env = { + key: value + for key, value in os.environ.items() + if key != "UNILAB_EXTRA_REGISTRY_PACKAGES" + } + result = subprocess.run( + [sys.executable, "-c", _SNAPSHOT_CODE], + check=True, + capture_output=True, + text=True, + env=env, + ) + raw = json.loads(result.stdout.strip().splitlines()[-1]) + _snapshot_cache = { + name: {backend: tuple(factory) for backend, factory in backends.items()} + for name, backends in raw.items() + } + return _snapshot_cache + + +def test_production_registry_matches_migration_matrix_exactly() -> None: + factories = _production_factories() + + assert set(factories) == set(PRODUCTION_TASK_NAMES), ( + "production registry must match the #1042 migration matrix exactly: " + f"missing={sorted(set(PRODUCTION_TASK_NAMES) - set(factories))}, " + f"stray={sorted(set(factories) - set(PRODUCTION_TASK_NAMES))}" + ) + empty = sorted(name for name, backends in factories.items() if not backends) + assert empty == [], f"registered tasks without any backend: {empty}" + + +def test_only_approved_families_use_the_frozen_compatibility_seam() -> None: + factories = _production_factories() + adapter_tasks = { + task_name + for task_name, backends in factories.items() + if any(factory == ADAPTER_FACTORY for factory in backends.values()) + } + + assert adapter_tasks == APPROVED_COMPATIBILITY_TASKS, ( + "LegacyFactoryAdapter is a frozen seam: only the approved families may use it, " + f"unexpected={sorted(adapter_tasks - APPROVED_COMPATIBILITY_TASKS)}, " + f"missing={sorted(APPROVED_COMPATIBILITY_TASKS - adapter_tasks)}" + ) + + offenders = [ + f"{task_name}/{backend_type}: {factory[1]}" + for task_name in sorted(APPROVED_COMPATIBILITY_TASKS) + for backend_type, factory in factories[task_name].items() + if factory != ADAPTER_FACTORY + ] + assert offenders == [], ( + "approved compatibility families must route every backend through " + f"LegacyFactoryAdapter: {offenders}" + ) + + +def test_all_other_factories_are_the_canonical_manager_runtime_factories() -> None: + factories = _production_factories() + + offenders = [ + f"{task_name}/{backend_type}: {factory[0]}.{factory[1]}" + for task_name, backends in sorted(factories.items()) + for backend_type, factory in sorted(backends.items()) + if factory != ADAPTER_FACTORY and factory not in CANONICAL_MANAGER_RUNTIME_FACTORIES + ] + assert offenders == [], ( + "non-adapter factories must be one of the canonical manager-runtime " + f"factories {[qualname for _, qualname in CANONICAL_MANAGER_RUNTIME_FACTORIES]}: " + f"{offenders}" + ) + + +def test_removed_legacy_env_packages_stay_removed() -> None: + for module_name in REMOVED_LEGACY_ENV_PACKAGES: + assert importlib.util.find_spec(module_name) is None, module_name From 3ea5905d87697dc29a8aa49e15e63e266859ada8 Mon Sep 17 00:00:00 2001 From: TATP-233 Date: Fri, 21 Aug 2026 13:55:24 +0800 Subject: [PATCH 098/169] fix: restore legacy-tuned g1 walk training behavior after MBA migration Three regressions found by a full effective-config audit (main vs branch, 63 task leaves x 5 algo groups, constructed envs on both refs): - offpolicy g1 walk penalty curriculum: the legacy PenaltyCurriculum halved the shared override dict in place on each env construction (two probe envs + the spawned collector in every offpolicy runner), so collectors actually trained at 1/8 initial / 1/4 cap of the YAML penalty weights. The manager runtime isolates each env, so declare the tuned effective schedule explicitly: initial_scale/min_scale 0.125, max_scale 0.25. On-policy owners (single env construction) keep 0.5 -> 1.0, matching their legacy baseline. - per-term reward logs: RewardManager gains step_reward_extras() emitting reward/ entries (weighted, pre-dt per-step mean, matching the legacy log format), merged into info["log"] by ManagerBasedRlEnv.update_state so the offpolicy/rsl_rl/him_ppo logging pipelines pick them up again. - g1 walk observation noise: restore the legacy actor-only uniform noise (ppo/appo: gyro 0.2, gravity 0.05, joint pos 0.01, joint vel 1.5; offpolicy: joint pos 0.01, joint vel 0.1) via per-term UniformNoiseCfg with enable_corruption on the policy group only. Regression coverage: repeated env construction never mutates the shared override and yields identical effective weights; per-owner penalty schedule pinning for every curriculum-carrying leaf; per-owner obs noise assertions; step_reward_extras semantics; per-term log keys present after step. The quadruped rough family's reduced reward set (12 legacy terms without a manager port) is intentionally kept as-is and recorded in the migration matrix. Validation: make test-all passed on the final tree (ruff, mypy, pyright, full pytest, benchmark smoke). --- conf/appo/task/g1_walk_flat/base.yaml | 24 +++ conf/offpolicy/task/g1_walk_flat/base.yaml | 26 +++- conf/ppo/task/g1_walk_flat/base.yaml | 24 +++ src/unilab/envs/manager_based_rl_env.py | 1 + src/unilab/managers/reward_manager.py | 17 +++ src/unilab/tasks/migration_matrix.py | 8 +- .../locomotion/g1/test_g1_owner_contract.py | 144 +++++++++++++++++- .../envs/locomotion/g1/test_manager_terms.py | 41 +++++ tests/managers/test_core_managers.py | 20 +++ 9 files changed, 293 insertions(+), 12 deletions(-) diff --git a/conf/appo/task/g1_walk_flat/base.yaml b/conf/appo/task/g1_walk_flat/base.yaml index bd6890ac7..b9df403e8 100644 --- a/conf/appo/task/g1_walk_flat/base.yaml +++ b/conf/appo/task/g1_walk_flat/base.yaml @@ -76,17 +76,41 @@ env: max_episode_seconds: 20.0 observations: policy: + # Observation noise matches the legacy noise_config (level=1.0, actor-only + # since the critic reads clean observations): gyro +/-0.2, gravity +/-0.05, + # joint pos +/-0.01, joint vel +/-1.5, applied before term scaling. + enable_corruption: true terms: base_ang_vel: func: unilab.envs.mdp.builtin_sensor params: {sensor_name: torso_gyro} + noise: + _target_: unilab.managers._noise.UniformNoiseCfg + n_min: -0.2 + n_max: 0.2 + operation: add projected_gravity: func: unilab.envs.mdp.projected_gravity_from_sensor params: {sensor_name: torso_upvector} + noise: + _target_: unilab.managers._noise.UniformNoiseCfg + n_min: -0.05 + n_max: 0.05 + operation: add joint_pos: func: unilab.envs.mdp.joint_pos_rel + noise: + _target_: unilab.managers._noise.UniformNoiseCfg + n_min: -0.01 + n_max: 0.01 + operation: add joint_vel: func: unilab.envs.mdp.joint_vel_rel + noise: + _target_: unilab.managers._noise.UniformNoiseCfg + n_min: -1.5 + n_max: 1.5 + operation: add actions: func: unilab.envs.mdp.last_action params: {action_name: joint_pos} diff --git a/conf/offpolicy/task/g1_walk_flat/base.yaml b/conf/offpolicy/task/g1_walk_flat/base.yaml index 2c25972d4..e61646cb8 100644 --- a/conf/offpolicy/task/g1_walk_flat/base.yaml +++ b/conf/offpolicy/task/g1_walk_flat/base.yaml @@ -77,6 +77,10 @@ env: max_episode_seconds: 20.0 observations: policy: + # Observation noise matches the legacy offpolicy noise_config (level=1.0, + # actor-only; gyro/gravity/linvel scales were 0.0 there): joint pos + # +/-0.01 and joint vel +/-0.1, applied before term scaling. + enable_corruption: true terms: base_ang_vel: func: unilab.envs.mdp.builtin_sensor @@ -87,9 +91,19 @@ env: params: {sensor_name: torso_upvector} joint_pos: func: unilab.envs.mdp.joint_pos_rel + noise: + _target_: unilab.managers._noise.UniformNoiseCfg + n_min: -0.01 + n_max: 0.01 + operation: add joint_vel: func: unilab.envs.mdp.joint_vel_rel scale: 0.05 + noise: + _target_: unilab.managers._noise.UniformNoiseCfg + n_min: -0.1 + n_max: 0.1 + operation: add actions: func: unilab.envs.mdp.last_action params: {action_name: joint_pos} @@ -198,10 +212,16 @@ env: curriculum: penalty_scaling: func: unilab.tasks.locomotion.g1.manager_terms.G1PenaltyCurriculum + # Effective schedule matches the tuned legacy baseline: the legacy env + # halved the shared override dict once per env construction (two probe + # envs + the collector in every offpolicy runner), so collectors actually + # trained at 1/8 initial / 1/4 cap of these YAML weights. The manager + # runtime isolates each env, so the tuned effective range is declared + # explicitly here. params: - initial_scale: 0.5 - min_scale: 0.5 - max_scale: 1.0 + initial_scale: 0.125 + min_scale: 0.125 + max_scale: 0.25 level_down_threshold: 150.0 level_up_threshold: 750.0 degree: 0.001 diff --git a/conf/ppo/task/g1_walk_flat/base.yaml b/conf/ppo/task/g1_walk_flat/base.yaml index c8466f9a2..4559ea737 100644 --- a/conf/ppo/task/g1_walk_flat/base.yaml +++ b/conf/ppo/task/g1_walk_flat/base.yaml @@ -76,17 +76,41 @@ env: max_episode_seconds: 20.0 observations: policy: + # Observation noise matches the legacy noise_config (level=1.0, actor-only + # since the critic reads clean observations): gyro +/-0.2, gravity +/-0.05, + # joint pos +/-0.01, joint vel +/-1.5, applied before term scaling. + enable_corruption: true terms: base_ang_vel: func: unilab.envs.mdp.builtin_sensor params: {sensor_name: torso_gyro} + noise: + _target_: unilab.managers._noise.UniformNoiseCfg + n_min: -0.2 + n_max: 0.2 + operation: add projected_gravity: func: unilab.envs.mdp.projected_gravity_from_sensor params: {sensor_name: torso_upvector} + noise: + _target_: unilab.managers._noise.UniformNoiseCfg + n_min: -0.05 + n_max: 0.05 + operation: add joint_pos: func: unilab.envs.mdp.joint_pos_rel + noise: + _target_: unilab.managers._noise.UniformNoiseCfg + n_min: -0.01 + n_max: 0.01 + operation: add joint_vel: func: unilab.envs.mdp.joint_vel_rel + noise: + _target_: unilab.managers._noise.UniformNoiseCfg + n_min: -1.5 + n_max: 1.5 + operation: add actions: func: unilab.envs.mdp.last_action params: {action_name: joint_pos} diff --git a/src/unilab/envs/manager_based_rl_env.py b/src/unilab/envs/manager_based_rl_env.py index a5b615126..c7fd4b19b 100644 --- a/src/unilab/envs/manager_based_rl_env.py +++ b/src/unilab/envs/manager_based_rl_env.py @@ -456,6 +456,7 @@ def update_state(self, state: NpEnvState) -> NpEnvState: np.logical_or(self.reset_terminated, self.reset_time_outs, out=self.reset_buf) self.reward_buf = self.reward_manager.compute(dt=self.step_dt) + log.update(self.reward_manager.step_reward_extras()) if self._cfg.sim_substeps == 1: self.metrics_manager.compute_substep() self.metrics_manager.compute() diff --git a/src/unilab/managers/reward_manager.py b/src/unilab/managers/reward_manager.py index 67b2a5a1c..262d64164 100644 --- a/src/unilab/managers/reward_manager.py +++ b/src/unilab/managers/reward_manager.py @@ -48,6 +48,11 @@ class RewardManager(ManagerBase): Regardless of the scaling setting: - ``_step_reward`` (via ``get_active_iterable_terms()``) always contains the unscaled reward rate (raw_value * weight) + + ``step_reward_extras()`` exposes the latest ``compute()`` call's per-term + means as ``reward/`` log entries (weighted, pre-dt rate), matching + the legacy envs' per-step reward log contract consumed by training + runners. """ _env: ManagerBasedRlEnv @@ -127,6 +132,18 @@ def compute(self, dt: float) -> np.ndarray: self._step_reward[:, term_idx] = value / scale return self._reward_buf + def step_reward_extras(self) -> dict[str, float]: + """Per-term log entries of the latest ``compute()`` call. + + Returns ``reward/`` -> mean weighted reward rate across envs + (raw_value * weight, before dt scaling), mirroring the legacy envs' + per-step reward log format. + """ + return { + f"reward/{name}": float(np.mean(self._step_reward[:, term_idx])) + for term_idx, name in enumerate(self._term_names) + } + def get_active_iterable_terms(self, env_idx: int) -> list[tuple[str, list[float]]]: terms = [] for idx, name in enumerate(self._term_names): diff --git a/src/unilab/tasks/migration_matrix.py b/src/unilab/tasks/migration_matrix.py index 511fe2db7..2238f4ecc 100644 --- a/src/unilab/tasks/migration_matrix.py +++ b/src/unilab/tasks/migration_matrix.py @@ -126,7 +126,13 @@ def migration_record(task_name: str) -> TaskMigrationRecord: "Compatible", "complete", "Hydra owners materialize shared terrain, height-scan, reset, and curriculum manager terms on the canonical runtime.", - "Keep the shared rough-family contract and both backend owners in sync.", + "Keep the shared rough-family contract and both backend owners in sync. " + "Known intentional divergence from the legacy rough env (recorded 2026-08): " + "the legacy reward terms feet_gait, feet_air_time(+variance), " + "feet_contact_without_cmd, feet_height_body, feet_slide, contact_forces, " + "undesired_contacts, joint_mirror, joint_power, joint_torques_l2, " + "joint_acc_l2(+wheel) have no manager port and are not part of the " + "manager-based rough reward set.", ) if task_name in _G1_LOCOMOTION_TASKS: return TaskMigrationRecord( diff --git a/tests/envs/locomotion/g1/test_g1_owner_contract.py b/tests/envs/locomotion/g1/test_g1_owner_contract.py index 2956d198c..0b04cebc7 100644 --- a/tests/envs/locomotion/g1/test_g1_owner_contract.py +++ b/tests/envs/locomotion/g1/test_g1_owner_contract.py @@ -3,6 +3,7 @@ from __future__ import annotations from collections.abc import Mapping, Sequence +from copy import deepcopy from dataclasses import fields, is_dataclass from pathlib import Path from typing import Any @@ -398,6 +399,30 @@ def test_g1_owner_materializes_complete_plain_manager_cfg( assert terms["gait_phase"].params["frequency"] == pytest.approx(1.5) assert critic_terms["base_lin_vel"].scale == expected_linvel + # Observation noise reproduces the legacy noise_config: actor-only (the + # critic reads clean observations), applied before term scaling. + assert env_cfg.observations["policy"].enable_corruption is True + assert env_cfg.observations["critic"].enable_corruption is False + expected_noise = ( + {"joint_pos": 0.01, "joint_vel": 0.1} + if walk_profile + else { + "base_ang_vel": 0.2, + "projected_gravity": 0.05, + "joint_pos": 0.01, + "joint_vel": 1.5, + } + ) + for name, term in policy_terms.items(): + if name in expected_noise: + assert term.noise is not None + assert term.noise.n_max == pytest.approx(expected_noise[name]) + assert term.noise.n_min == pytest.approx(-expected_noise[name]) + else: + assert term.noise is None + for term in critic_terms.values(): + assert term.noise is None + assert list(env_cfg.actions) == ["joint_pos"] assert env_cfg.actions["joint_pos"].scale == pytest.approx(action_scale) assert env_cfg.actions["joint_pos"].use_default_offset is True @@ -599,6 +624,8 @@ def test_g1_walk_profile_runtime_obs_scaling_matches_legacy_layout() -> None: _, _, env_override = _materialize( "offpolicy", ("algo=sac", "task=sac/g1_walk_flat/mujoco"), "G1WalkFlat" ) + # Exact comparison against raw sensor reads requires clean observations. + env_override["observations"]["policy"]["enable_corruption"] = False try: env = registry.make( "G1WalkFlat", sim_backend="mujoco", env_cfg_override=env_override, num_envs=2 @@ -633,6 +660,8 @@ def test_g1_legacy_profile_runtime_obs_scaling_matches_legacy_layout() -> None: """Legacy-profile owners keep unit scaling on every observation segment.""" registry.ensure_registries() _, _, env_override = _materialize("ppo", ("task=g1_walk_flat/mujoco",), "G1WalkFlat") + # Exact comparison against raw sensor reads requires clean observations. + env_override["observations"]["policy"]["enable_corruption"] = False try: env = registry.make( "G1WalkFlat", sim_backend="mujoco", env_cfg_override=env_override, num_envs=2 @@ -658,21 +687,120 @@ def test_g1_penalty_curriculum_scales_negative_weights_from_start() -> None: _, _, env_override = _materialize( "offpolicy", ("algo=sac", "task=sac/g1_walk_flat/mujoco"), "G1WalkFlat" ) + override_snapshot = deepcopy(env_override) try: env = registry.make( "G1WalkFlat", sim_backend="mujoco", env_cfg_override=env_override, num_envs=2 ) + # Repeated construction from the same override must not drift: the + # legacy env halved the shared override dict on every construction. + env_repeat = registry.make( + "G1WalkFlat", sim_backend="mujoco", env_cfg_override=env_override, num_envs=2 + ) except ImportError as exc: pytest.skip(f"mujoco runtime unavailable: {exc}") try: - assert env.curriculum_manager.active_terms == ["penalty_scaling"] - # initial_scale=0.5 halves every negative weight from construction. - assert env.reward_manager.get_term_cfg("penalty_orientation").weight == pytest.approx(-5.0) - assert env.reward_manager.get_term_cfg("penalty_action_rate").weight == pytest.approx(-2.0) - assert env.reward_manager.get_term_cfg("pose").weight == pytest.approx(-0.25) - # Positive weights stay untouched. - assert env.reward_manager.get_term_cfg("alive").weight == pytest.approx(10.0) - assert env.reward_manager.get_term_cfg("feet_phase").weight == pytest.approx(5.0) + for built in (env, env_repeat): + assert built.curriculum_manager.active_terms == ["penalty_scaling"] + # initial_scale=0.125 scales every negative weight from construction, + # matching the tuned legacy effective schedule (1/8 initial, 1/4 cap). + assert built.reward_manager.get_term_cfg("penalty_orientation").weight == pytest.approx( + -1.25 + ) + assert built.reward_manager.get_term_cfg("penalty_action_rate").weight == pytest.approx( + -0.5 + ) + assert built.reward_manager.get_term_cfg("pose").weight == pytest.approx(-0.0625) + # Positive weights stay untouched. + assert built.reward_manager.get_term_cfg("alive").weight == pytest.approx(10.0) + assert built.reward_manager.get_term_cfg("feet_phase").weight == pytest.approx(5.0) + # The shared override dict is never mutated in place. + assert env_override == override_snapshot + + state = env.step(np.zeros((2, 29), dtype=np.float32)) + log = state.info["log"] + for name in _OFFPOLICY_REWARDS: + assert f"reward/{name}" in log + assert log["reward/penalty_action_rate"] == pytest.approx(0.0) finally: env.close() + env_repeat.close() + + +# Every task owner carrying the penalty curriculum, with the schedule that +# reproduces its tuned legacy baseline. Legacy offpolicy runners built three +# envs per training run (two probe envs + the spawned collector) and the +# legacy PenaltyCurriculum halved the shared override dict in place on each +# construction, so collectors effectively trained at 1/8 initial / 1/4 cap of +# the YAML weights. The on-policy runners built a single env, so their +# effective schedule was the declared 0.5 -> 1.0. The manager runtime isolates +# each env, so these params are now the single source of truth. +_PENALTY_CURRICULUM_CASES = ( + pytest.param( + "offpolicy", + ("algo=sac", "task=sac/g1_walk_flat/mujoco"), + "G1WalkFlat", + id="sac-walk-flat", + ), + pytest.param( + "offpolicy", + ("algo=sac", "task=sac/g1_walk_rough/mujoco"), + "G1WalkRough", + id="sac-walk-rough", + ), + pytest.param( + "offpolicy", + ("algo=sac", "task=sac/g1_23dof_walk_flat/mujoco"), + "G1Walk23DofFlat", + id="sac-23dof-walk-flat", + ), + pytest.param( + "offpolicy", + ("algo=sac", "task=sac/g1_23dof_walk_rough/mujoco"), + "G1Walk23DofRough", + id="sac-23dof-walk-rough", + ), + pytest.param( + "offpolicy", + ("algo=td3", "task=td3/g1_walk_flat/mujoco"), + "G1WalkFlat", + id="td3-walk-flat", + ), + pytest.param( + "offpolicy", + ("algo=flashsac", "task=flashsac/g1_walk_flat/mujoco"), + "G1WalkFlat", + id="flashsac-walk-flat", + ), +) + +_OFFPOLICY_ALIGNED_SCHEDULE = {"initial_scale": 0.125, "min_scale": 0.125, "max_scale": 0.25} + + +@pytest.mark.parametrize( + "config_group,overrides,task_name", + _PENALTY_CURRICULUM_CASES, + ids=[case.id for case in _PENALTY_CURRICULUM_CASES], +) +def test_offpolicy_penalty_curriculum_matches_legacy_effective_schedule( + config_group: str, overrides: tuple[str, ...], task_name: str +) -> None: + _, env_cfg, _ = _materialize(config_group, overrides, task_name) + params = env_cfg.curriculum["penalty_scaling"].params + for key, expected in _OFFPOLICY_ALIGNED_SCHEDULE.items(): + assert params[key] == pytest.approx(expected) + # Thresholds and annealing rate stay as tuned. + assert params["level_down_threshold"] == pytest.approx(150.0) + assert params["level_up_threshold"] == pytest.approx(750.0) + assert params["degree"] == pytest.approx(0.001) + + +def test_ppo_penalty_curriculum_matches_legacy_effective_schedule() -> None: + # The on-policy runner builds a single env per training run, so the legacy + # effective schedule equals the declared 0.5 -> 1.0 range. + _, env_cfg, _ = _materialize("ppo", ("task=g1_23dof_walk_rough/mujoco",), "G1Walk23DofRough") + params = env_cfg.curriculum["penalty_scaling"].params + assert params["initial_scale"] == pytest.approx(0.5) + assert params["min_scale"] == pytest.approx(0.5) + assert params["max_scale"] == pytest.approx(1.0) diff --git a/tests/envs/locomotion/g1/test_manager_terms.py b/tests/envs/locomotion/g1/test_manager_terms.py index b7d65c494..52f50747b 100644 --- a/tests/envs/locomotion/g1/test_manager_terms.py +++ b/tests/envs/locomotion/g1/test_manager_terms.py @@ -213,6 +213,47 @@ def test_penalty_curriculum_scales_only_negative_weights_and_tracks_episodes(): ) +def test_penalty_curriculum_repeated_construction_never_mutates_source_cfg(): + """Regression guard for the legacy shared-override mutation. + + The legacy PenaltyCurriculum halved the shared override dict in place on + every env construction (two probe envs + the collector in each offpolicy + runner), so collectors silently trained at 1/8 of the configured penalty + weights. The manager runtime must isolate each construction from the + source config so repeated env builds keep identical effective weights. + """ + from unilab.managers import RewardManager + + source_cfg = { + "pose": RewardTermCfg(func=lambda env: np.zeros(env.num_envs), weight=-0.5), + "alive": RewardTermCfg(func=lambda env: np.zeros(env.num_envs), weight=10.0), + } + effective_weights: list[float] = [] + for _ in range(3): # legacy offpolicy runners built probe + probe + collector + reward_manager = RewardManager(source_cfg, cast(Any, SimpleNamespace(num_envs=4))) + env = SimpleNamespace( + num_envs=4, + reward_manager=reward_manager, + reset_buf=np.zeros(4, dtype=np.bool_), + episode_length_buf=np.zeros(4, dtype=np.int64), + rng=np.random.default_rng(0), + ) + G1PenaltyCurriculum( + RewardTermCfg( + func=G1PenaltyCurriculum, + weight=1.0, + params={"initial_scale": 0.125, "min_scale": 0.125, "max_scale": 0.25}, + ), + cast(Any, env), + ) + effective_weights.append(reward_manager.get_term_cfg("pose").weight) + + assert source_cfg["pose"].weight == -0.5 + assert source_cfg["alive"].weight == 10.0 + for weight in effective_weights: + assert weight == pytest.approx(-0.0625) + + def test_penalty_curriculum_shrinks_scale_below_initial_when_min_allows(): env = _curriculum_env({"pose": -0.5}) term = G1PenaltyCurriculum( diff --git a/tests/managers/test_core_managers.py b/tests/managers/test_core_managers.py index 653362c16..2fe0fb5db 100644 --- a/tests/managers/test_core_managers.py +++ b/tests/managers/test_core_managers.py @@ -144,6 +144,26 @@ def test_reward_dt_scaling_reset_and_config_immutability(fake_env: FakeEnv) -> N assert isinstance(manager.get_term_cfg("stateful").func, StatefulReward) +def test_reward_step_extras_report_per_term_weighted_rates(fake_env: FakeEnv) -> None: + def ones(env: FakeEnv) -> np.ndarray: + return np.ones(env.num_envs, dtype=np.float32) + + cfg = { + "pos": RewardTermCfg(func=ones, weight=2.0), + "neg": RewardTermCfg(func=lambda env: env.value, weight=-0.5), + "zero": RewardTermCfg(func=ones, weight=0.0), + } + manager = RewardManager(cfg, fake_env) + manager.compute(dt=0.25) + + extras = manager.step_reward_extras() + assert set(extras) == {"reward/pos", "reward/neg", "reward/zero"} + # Weighted reward rate (raw_value * weight), not scaled by dt. + assert extras["reward/pos"] == pytest.approx(2.0) + assert extras["reward/neg"] == pytest.approx(float(np.mean(fake_env.value)) * -0.5) + assert extras["reward/zero"] == 0.0 + + @pytest.mark.parametrize("bad", [np.nan, np.inf, -np.inf]) def test_reward_nonfinite_is_an_error(fake_env: FakeEnv, bad: float) -> None: def reward(env: FakeEnv) -> np.ndarray: From 34acc0d95acbe4698a8392a71930adbd9c98838d Mon Sep 17 00:00:00 2001 From: Xiaofan Li <103843797+LeeLeno@users.noreply.github.com> Date: Fri, 21 Aug 2026 15:14:39 +0800 Subject: [PATCH 099/169] fix:robot xml of all (#1238) --- src/unilab/assets/robots/a2/scene_flat.xml | 15 ++++++++------- .../assets/robots/allegro_hand/scene.xml | 15 ++++++++------- .../robots/g1/assets/scene_crawl_slope.xml | 9 +++++---- .../robots/g1/scene_climb_20_z_scale_1.xml | 11 ++++++----- .../g1/scene_climb_20_z_scale_1_23dof.xml | 11 ++++++----- .../assets/robots/g1/scene_crawl_slope.xml | 11 ++++++----- src/unilab/assets/robots/g1/scene_flat.xml | 13 +++++++------ .../assets/robots/g1/scene_flat_23dof.xml | 13 +++++++------ .../g1/scene_flat_23dof_with_largebox.xml | 13 +++++++------ .../robots/g1/scene_flat_23dof_with_wall.xml | 13 +++++++------ .../robots/g1/scene_flat_with_largebox.xml | 13 +++++++------ .../assets/robots/g1/scene_flat_with_wall.xml | 17 +++++++++-------- src/unilab/assets/robots/go1/scene_flat.xml | 17 +++++++++-------- src/unilab/assets/robots/go2/scene_flat.xml | 17 +++++++++-------- .../assets/robots/go2_arm/scene_flat.xml | 17 +++++++++-------- src/unilab/assets/robots/go2w/scene_flat.xml | 15 ++++++++------- src/unilab/assets/robots/sharpa_wave/scene.xml | 17 +++++++++-------- .../assets/robots/stewart/motphys-ground.png | Bin 56655 -> 0 bytes src/unilab/assets/robots/stewart/scene.xml | 17 ++++++++++++----- src/unilab/assets/robots/x2/scene_flat.xml | 13 +++++++------ .../assets/robots/x2/scene_flat_with_wall.xml | 13 +++++++------ .../robots/x2/scene_flat_with_wall_visual.xml | 13 +++++++------ 22 files changed, 160 insertions(+), 133 deletions(-) delete mode 100644 src/unilab/assets/robots/stewart/motphys-ground.png diff --git a/src/unilab/assets/robots/a2/scene_flat.xml b/src/unilab/assets/robots/a2/scene_flat.xml index e378a4668..11e1fc5f4 100644 --- a/src/unilab/assets/robots/a2/scene_flat.xml +++ b/src/unilab/assets/robots/a2/scene_flat.xml @@ -4,20 +4,21 @@ - - + + + - - - + + + - + - - - + + + - + diff --git a/src/unilab/assets/robots/g1/assets/scene_crawl_slope.xml b/src/unilab/assets/robots/g1/assets/scene_crawl_slope.xml index dceb86081..94bdaada3 100644 --- a/src/unilab/assets/robots/g1/assets/scene_crawl_slope.xml +++ b/src/unilab/assets/robots/g1/assets/scene_crawl_slope.xml @@ -5,15 +5,16 @@ - + + - - + - + - - + + + - - - + + + - + diff --git a/src/unilab/assets/robots/sharpa_wave/scene.xml b/src/unilab/assets/robots/sharpa_wave/scene.xml index 42bca1199..3b95c2cb0 100644 --- a/src/unilab/assets/robots/sharpa_wave/scene.xml +++ b/src/unilab/assets/robots/sharpa_wave/scene.xml @@ -8,22 +8,23 @@ - - + + + - - - - + + + + - + diff --git a/src/unilab/assets/robots/stewart/motphys-ground.png b/src/unilab/assets/robots/stewart/motphys-ground.png deleted file mode 100644 index 04967dcfbf9b80824210939ffd81c69eb3007e85..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 56655 zcmeFacU05c);1bNK}BG13l^H&hK*vOR~yoc^csqQgetwq1}KPZL_k15Y9x?EO6U*} z1?e?}8boRc5LyT&A>Rt;-RGV=#yjpG_pfh^V@8A|{K{H$%{A9EpZTo#T37o4*P+vg zAP@+bn(DoW5Xb@W@&M!jJNQ+E=oJUQ4tl7XdO;vMWY)huzep!B;KNhi_f5R@-5tFB zY&`8Detv$UPHxU#b~YaNqVArKDU0%_A&@f=wR^V>{6Ef*ntE6mr_Ioyi&y`K9FbE# z@^kO=Y)M#troqF!{7M7Ms6NY4g+CP({(QSGIDTNeo~Lwdx?bmP;W_$^@smn}<0l@p zz0t5lBGC`?R3nEjyvtOLJac<~`b zKI7l9na&(r`O3r0uLdpt=TVp8fTI5MhV{E{XHbaV$y|am4Y(slXRUs|rga1nsO_^K zYiTPo3@PT|FjkOTm$oyGV?g885!7@#HD9)euYl_Dv~r!AYRcI9WjsfABH%U&OnX!W zTVD-x2g;Na5s9EqJL`~g9AgN9(#n<^ImKI2%tfhp+02=-otUCQX->$=%`s*tQwM5+ zLpXK{;yd9>7XP}y6mtW6WANk743i|}ZTkt)->_$HC@2v_Vs635^o@g*ncv4kaI8;> zI0kNv&}gK|1lKZA3TAAY@o;>0&NEPFE{!qV-w7Uscu3EC9)Unk7MWg?V`2vY@W!RPec8}(OiyY3bfR__z_w=lZU+AMp#mjQ+;`W;}OrI)&qG1;-AE(EGfdaYP zH;kgf#;z+L{Kzfmb*k=VL*Xh#oOM>F$xeyz#22{@0k-!~zrg>)wx=O`{^K3w7G&># zzI_bY_n&XOc75;PpM2W$$G=~L7s_mB{`2iqcIE$k3*mbD|9boT#=lMd+qwR8O8j3t zovblb2)CJRYJY!!!bj~0`c5v}+|3*mN~B^`FrbFWMk;^MKw*F=HoMzO^rQ|E%Te5?ane>xOqAn796ZO z=wkG)ueUcVEscm=H&xTuj{}X{k1@Z-o2n3`Et*JHZY(Y@&yJ6`ZwT_OZQKXxOR*3| zv?hx25m&`_zK^vm(uCTcvYLM@d&$g%!WidK|2FY1dwkp`g!a`he`p3YQq-#Yv$w2> znkcU36cI6>mHELnZnc1VbJKKM<;lsm=`V|Gm{wAB*qwA$$9;HIUh7Qsw5cW>?&e`SA<OvGqEMBL#~Y$hVTUPd7vqwXllxyE6~5Brh`OI;9(nKVp?!M|o!~Y$l=~l5{@})t zN1=?l*FW4?SA*dHh*rU3u{qD`_cUzhA`srwZAp=90W&G?w9 zIF6Hw`=_U;OTK*hL?jZCCCsTrsO0>JH19APc9{iAA^xAn*&rY+uyb(m6yO#XT-)2* zd41f^7iHw+T=4PpGY&A|u(PqLy|M?w^PKavp=hG(z8@I9AQSqEnbm(RdF>u?}^jbn>X%%R|_pq?YQd1pInP~>=bneEZdV7C< zdUV4{OY==)NmA7pap)KMdYj_7T-=HinKQ>@cYXQ!>HEYJ%S!LOY*{J+ZPk0{g^OE6 z13CG_htK|1gvi66S=WM`JO}x72YB4Yix*pw3sMpic?MTZ!}J8I!@@S-@!oDLz+jw% zf~ss6cHxpP8VsJs39+lZQHk`C_^U#&AUUI)Sb+jE0ryZc1!WR+^ypDT8QtAW>zv&S zx$?-(O}l&6$h6bp#wjgL&0F_QZ+Oe1@c8)1lK9RUY>){K6CHE6sjGnkh_tILNVZCyT*_^ zeFbtRwd1k(nw6>32Cb*JH(xL=?5V#zIw^@1xmH%no6t$fUSD62EHSH2<-;dEr2559y6(#hV*R1=o0^WcPwdOuUjPkw>NaU`tCAy4>74J=qIC zJUR5kLe#PRq;1*o)Uv;Ch+?ej$(?KEql=cePlJV^wYAmE($w&pps~E5zGzC^@I-9W zg3%X+&uM9CAsl}1uL#CmJ$C3aZNE^%AgQO8b?R_Dd=4%`of_hhp#(ELzPI%L#oKff3R) zxsF2q?wpa3k+CXlhpU-cSb%r>UaQ0;NYP-y#fyzslWX05eO0V5##WdzTi5Gpv%Er#ufu5K2l*b4{)pec^vGX4jo}uKnGSttP)6J-0IRkdTnmf{LiO z9oh|w{SD+m+Nhs~pt-yt0Z!*(7XGA8G{umq6kIKNm4LDAotn-3k$~Ua3sR>-KaKv|a>rLigvb+-ClPY&EIVs7=U&O`V-OH<(ot?vb zaYUwpHu59Y-J9I++E5xwpAgDNl$9S<4tni!Z><%&)3cN&w6n>>#r6oqa1c@@t}AB5 z#~-)9(RiDaXCv5(N~Jk_ER|PQRjmvBU`@VVOaVfOq01_CD9>?>_b?KJaqlEIKYExb zYK45!zF{~7h&n?f9Y*OZ{1v(_#ZHj`g(t=sgLYkrIV z_3Ov}!_#d|G3mM*8d2G)J_S=#_FzrXBv^oXn4FS=@9j0mOCU8bW`V`#{#m|Uj`JgK zWdNVy@p$|4gSFMwE@nu($Fn8$`+^q-G{6}A_#w~`q@5;j=%JU%x)~kbX&;3dO6Q_vPhjY{HL? zVJKD`Z|OTLyoxY;H#aMxkPW@=C{Yl4z)}wS)LUrRv^L5)CPKRPG6~F(;W9UJKgv8P zEqS1IdOFuFG(`WL-{M@wYUkWYg;x+!N?yKTxD0pswCx^M=+@-C+`^!DXNmm;|lu36jL*^g4yQCJDwEtT*95jRxls*?=4THg;h}Rc|g}E^4r)9A-&2*5H z)U3Apt!>Ug-(xp%RW(`wlABU+%$;Wry{uV@v=4T%pe;uKtNj?u8sz5YCUUJpAbRda ztI$dIOG4}u6BB0UH;IQwTca^<{68!%q-v8GO7z?5!7EOtDa@FtduLCds?f(uWWWxD z-8n5J#A@RcK0fm4Q>XGwrlU0-5HDxGe!XK#RsxtFI0%R}#MPOmDbG`B)rr|2(k^Ca zyEXzQeEzofpuBkv2FPzQIBRflF!Iqct@H0da`4$d9MFIeezyok2mpmnPEMwxiTN-g z;NY#T8a&%00|OiV>G?L!&KjJYoQ4Y%XWm!pzSGP_^0qC1c9G2Q?(UA%8R#XIh(aq1 z3bz{(b>#YufpW#2jjRwc$?vvL7>j-s!}DxSIT|Ohl|@AqU)D$fFfhD2w6yePytu-m z!}}=i@A#?XuCAxj56cFu^r1V$8pvQ&Le1m%XIBi$Fc#B8q026;B(LuF(9lp1%Dk+v zug?Gr(iW85Sn%qXPB5mUwT$Tc&+G0sHn*W;8^aT z0iXvM{lehT^uWN%{GJCF+S=O<3sCpu(ICK+9hI1Q1^p`%k(u2LF^Ieeh4{I$vdV|P zvFsjFi^D;mDHorOZGYbRmdJ-sy6LFC*n5ou_tS|Lt`t`w7MOS_k_uP+f z;9&)5bPL6OD1NWf9|wL+d0t7nHew8>J^*v?wk@Di-1Nt8ZbZYgwliPq?+u!sj4Z6F zL6Jz81(ZTFVIiQ2KRu+^_CPo;C$>7-+R~M?oDc}a_zDe=*?ILLvV5;nXPTk8uyxcNT7C;o9j$ANL!v z;@9}V4Q_1=TwKp;Nh&aE04B=H%E}d<=jp+V;ri<8>d>EF>fgOvxiy7<-iZPj-c=;q z`pX8!RWV&kgeiWjWP6R-j&X>=O}HrXWpILv;WJA!Lw2jm!r?(USd3!Q?; z4~5;5kg;c>$}zu$3b2NB|FZQdmYN3ZnwHkw{Gb=lOe%MT<-$T$FAmXrxuL)vz&I4j zbw6f1xR$|FzPGfj3|wT+LFCT6>6w|no}MCd4S`6s`H*yjHwpmRW-^s`l%PEgnV1g= zLS48fDw^QEG6usokAJzGK6W`t*Cmh{9nE$B?4SGhKFiz&8LINcUyx7xVB2EI3||q_H*4=v&?<8XLBY!@%$!3v1X+u%#s#%=dXrnJ4A2Aj2Hgr z>Nb89PxW{q6RUYFzJKQ#2mx8W2%QwsjrX<)UufJsH=-!1M6xXNKEd_P_g~2$KlHth z8+#inboz7;h_feEh%s-8gfPzUVvpN8+xhzMdRJJ;`pyrhP{8D)QTPxwYyR@GL5m;X z@{^3gh`2}&m!XA^90Y3-%*93XX79~#PEXq6=x#a<5nt|(5IW^>I6oY~c`!%k5K7xg z^XpsYNGDfky`Ugn+T!sv+{*?s2;YR}IAN~gO5cTKC-_WfIw3+R)MQ0k=vC81<)+Va z&hcQyba}xT7s(XQis~ClC7InKIla>etbX46gv3MupU$ay%38_>uFayiGrK?AK%S1e zjjdL|JY+3lu$5+qVVN&~H)msOYAn~-sVyz)-`OAly#Wgs7=bGmylZ~cjhF%#la%cD z+poLt#R9<7w?^c7A_(e-Ag}J7ZE@g?X>}{3BsF`vwo5xdZ5Y>ebVTfZ_VlAR0WDDn z(80pOXTaQP^8ESoPX2C2EI&TyZe)r`NZ`d57FxyCPAr!L#Cl8GWaF5>2z71k3O_fW zgNliTg)WGi^=HB#d$+C`glcGzYkTKdgZg@i4WX0X?=_Pi*EBJ>x-~NYv6jA~(VM(* z?wlWW?Ph1+)ZUIPK-W-LTE+6WbvSur%q=bPC(_|)JpOMRxs(+1u7?T1LGYIAXO&+ae^pLc4#(V5Ljw7Rkt4lVb>9juTQk{9^z(a)0c(2i}uE{sc?}`IivW#N)|C{cE`p4@MHK;HZm?lz>G#z?+(<98}6f7b@A=*$;JOvAB=& zFOLU0mnWL=-A)#`(YD#y$9jW3va+&Za7C&{3umZchP8KgCb_-ApwXZ6^YcY{+~ixx zR3wh6T91tf30OsiY41C*lv|jarvOw!qImp8KIS9xh^=~c+j zdWUMg(q<3o0wS?C6X!SnWmzRuygZ2O-09JW7xgz7v&;V8bA5l0Z@n0GGd|Omsr>1F ztm-{Lx9#VFxVX82%xux@1?Xm7Tvn~)#UK#1TO5ixw*Gv550E&zy@ZOWc_SOaD`6zXm|5!7e)Czff&Akw_Y<2m;O;xP|Ux1sB+ z>3LZV(!Ue{@Bkdcm5fTc9BVVBxPUZe0xkqVf#pl$COGbJu(hA{z5; zf`b8s8_g%3JhbdDAK3DUn>XtFEluL+Ux;X|IQJP6 znJh;7se0xQ@nB2PqP)KP_u_y_Q4T>;?3!ah@>}UganmCa_`2!zkJVN&^qsBzg9C$u zdAg7B;#Wce+^e^a;D76mTj(b=@0cN-N2>0E)cV&HAq%FtUv=RD% zUJEA%%m{Jm@V2#gkT#aL+K4}C+iQ*KHF|ybdVw9V1gl;-TyWEpqxd6T?#z>&sc^T> z_LUlW=2uQLE7R^4&+(rhc9RUOCF$Vsb61$VyYVFh>HPZH*~R%iK`e$|Q){5R8(mu3 zZTruBTKGffVWn|smKEkUT9e{evYD|m5s7KEmh zaRr_M7hhoh-JH{4B_-ugZ0TYe4bO(?&3X2$KC^|o;L^r3*OBUec&0*y^2~N^Z4HTX zr8#U`ZEbHKk5)Mzb$szAU^95577krp+6)FrFWhTf_p#sl(;C3|@g#9fA`JdejuwYNp8=~jNXz20uSA)VNo zQ(F4zPR3VPclQSxu+=lbJn!Fg6Wz-4v$8CsqBQnX)Xx8EakyDDs17yrmgb+EyCPDT zlb^2<7k4zAmi;hKq4;1)ncriFT#s-Z3e}$3GIc8c$luu@SGpD(5LW>|8~eUNk3RIK zOGGQ`Ab-v zU#vpK5)Q{Xp=u3a6xzVg<=YinynT(0(fSV>C!UCi9Je}FIJ_Vm8f@L^?2xj|!iQ&L zL4oH9AfEsX+^=N9T9uU*szS&Kxfc;Ig*VBdcvZPl&k-)%;XL)N`DhTkiPGHSj~`{) z%VPo+&Vg5$fSX8itK&UaWb-j|lp^MC^~kqUDOR4{!%qC< zwJUt}S$p@a9u63EDYb?>B=NNO^MmX!v->N>CZ*kINx`AKqmP!Vo>%V zXwvE7j12<7b_T1ft5`vm(E%b6J&bt(O@9)-0cex&-|wr*o!avzvpFaAmX8BR_ORo$g*kpTzU!W;QCBhGs3AT^0;!2|n}wTju_n0ORlvl=@FC8~OO*I+!PVyG=H54qix+TwP(9CLE=m2aYI2A}^w9KRMP>(_ z5xU6;#g^>r8{z-?QVH~hI|oq3mbDh?!bE3MKWg|8`>=$#v@SU#SmS{s;;n#OE{XKj z2@i-mpXqs10I#mtJnUEwI{_^i7#VX>obz~Z%-asK54SQtw-v|Yq{@`I3X7ba1UP-r zzuH6je#IW1skhAzUaD8*1+%aXZ{KpG3a*`~PE*gpR#fDHE+}A9><@^yM2^$8?XhND z%+S_HH_WW8dREfIA7jm?yk)06q^CS(raa0gJ!K|+tANSuu?YzZe}cyABvkUSgfxL$ zq@KAV$^+5FUQn^b#pl*qwYRnXm{{(&)T*nkwQ_YS;M1<2o_5g6TJd{+@Qwa_{Moqc z+>kGFEzBn#ve9gihxTr-jN=+iLt*zL8a4(EDytv+B<5F#9VvdZPbD}Mx3#r>OA5e{ zwC(CGyy%;y(`T>l+@< zKldQ!Gf0DNqaK4sD@y<>o-vA}SAiTpUY#@wt#M9Hs1*#}Ep58$X+%9?FMP57HA?KU zuW!GH{%TaG5}~m%+(ypQl9g3F{@!Ts8#v@}Rq0f2y8}N&H|7*XU8IyOB4aL>4Wn{%9atXHSK`;js&A%QL}c!(yPcetCVLDlcpZ4;o{oQ$>Az zPa`~IzT4&s%WZ8(w;p_dC0%OKUFoP=%IKqq77+-Ob$y-=m-_S2X)@NNg}+ykR^I+l ztE01%a_Z;x>xvE=Bcxu0=62|Gk|SV>LM6kSK;9bt=FM!)7!KhVx~VaWgX3&nSzsdi z3@ZVTzZzi+*mJcEIO*4goYuo7=mapVT=WCa_ zi5SKHARL_1wNH{&Tf?tYqxr&5)y?UhKhhmxtcawF`vxplI8J%Sr&}(5dwn!KulJcc zkV0G^dD%I$8kA6S>x!3Cp-EqVzr#>jyF+ht90K|9EFgcK0|V73ChUBY!dRoo5d(Ru zGzeHZz(b-U^b!DBJT;YL#xr4FV3ZGf&=S~hTll%GtnBN#y>GZu1#eU<%F1q`lrZS< zL%HLL*>+8Dn-@gUrH??Sy|aSZc_WT zb*g``*fHwR$x9DF_Xh?Tm957>$@%Gv2au<}a}8SB+I@j1w{IvZVcOfjro>{qh(yt{ zVf_`aT|gW;EPdbggZ-x58=sko@bE`vC(w7x0p!fdfpeE6ofMLC>Xc_NvjF2~6~vHW z{5pZ_0`AwkHF*69D3=%ODCIv+0j8?6ITEcx9gh&Y0(B#<$a*ZyVJqdT#e*FN#Iwx) za_mTp@6?&b1^Rw~8yz~xKFp=TcmBvhfNqrNUeYcJJ4z-GnUX1%Y-}%Ny{1174a;3g z`|}KMl#`oVX{IX}il0lDtA`EG0dy1j-Iim>@nW3QtJwii|KQ?DkL@+BNVeU!0pm#^ zAfCj@l7J-91wa&Zi+>x8urQ-wonYJUCMJ5CerC#Z+rK8L?_1D&ycrxrc+*Sj`}y{XbxfAFI8f5 z{Hy5(J77VRrS+Nae~~@FqW2d-HO8{m+jQJxykrOu)+nA3qXGZ}t_lTzE@Q5U040za z&x4*0i*2wTbOGT*WN71pFBuGVUW@;my;V7>N8G1w1ob=Q+1gCcuW&`#;fPN zmo)XeUIc1{9UY$06}$d=~l$2#tPdtq*F zqOx+gb#!7}O6RIY+J{CvC$G2`y(&pheuaN`L)FG{tDu>Tm7ANxBxI(_c=d}ztYMz& zW1Ag*B}GM<_cE7M=y@6X$2U(MW`A*r9T$>(22h(pLFU85!;x;T&gY6Us2(PO0=O^Y zp#4S%NIF5Ik}E?Ok^#d4!o%IGv9DMH5LS+CdwMohm$jy%(Tka_M-{h6iVgF^M3Xvk z#e*x(*(n12mn4tVZx40aa1=Z80MK&Y$Ge1&9rCh)Z;sOGh}x}7iJP0$Yh%U`H`K zDBb^)5K&tj>gqCZPSl!{GvwD&Db)pjV`7_WlN!j$OIRum{VIW!GiIj$O@3K}9=%#^Mi406vaY6A5F&>7B+Tn8jb?Fp;k45Sx4 zaabC{?A}_;dG-BGda=nhpQJt3K$cQa#4$Xq*Biv`47>ze#+Ody2{S4w)RShw&;0S( zjF45wL;0dI$B~LW9HBVZ(bKl|<-TOf#IQ!vL;zjXyprQ$i8)>0i zdGFFvK!mM!D;}JrZPRX+B+G`{>KbD#Mf@X)1)9xWUAlSl2^#sy0;HCpAU-*kb$p#9oQ}li7}l?U+HB|e$qF-TYu@#7 z&6bvF+gHk`BjaEuLc!YFhd_w`I=`|4}}@w2L}y%yjY^nET=uk#>`CARa7E~lL`ukOM#$RY-+z%ZIwRarc=N2vY;4*YxzN# zAWPifGe1!GY#qpiOnqC*8h}$dd3xSyXmET~jHeG@s~XKfqeVqTMHd5Va5&tY-O^^8 zR1OBZt**2&&J*c&^T_MRBaiM@_huJ7Yp{{G0X&WR#W=&)1-MNr;q*!A(m3tf1Ax@C zvunhZfnee`r<>3KsG7d!qxLx$sDSA9oK)>*v~{%cT1ns9jHKPdE;-p8k^H$ntUfK05RTC~#O| zXmwjDDJjky?M#k$&F_Kc%6rb>iSWC8RDQm>8*Nd4!zds7xqXx2Db2qU7!vNZ&59GO zLii%xD|Z3txk5ZPi=5tvGXI?T8E)Ic1}s6E%xfXUgUm0G$pi08ao?#3zKB={)V8SQ zEqF}`BP5>RR_ibi4|#pTT~k9qesMBS|3XgDT3`*RP6plf3~*8dmY$mx*Vfo52Y84| z{pxh6#%e;7rJ7Whm^W;#Oa`_4)$lPo+!|C?jnNj}Jn)SMWNyD$f3J=bsvIZL{k6*ZX^tH?E zs2>~^E9bYa*^~rkhoU6f4^Bl(O1^<&0Na3^wyRm)WZ3V&@ln>^;sxc2O3F5%HCcCQ z@95wiB{nuz0yXU8a@}5=cd;z^DLoH{{ML;d55W@s85D~{EPwVfZhJ}dpFYVyKeZ`kTsOFQz>!S`IzBoy?)t$CnHnpS7)l6xq&QwaZyl9^H^Fns9^(#Wi!3T8z3(j zATPIvjT@m4Up}EE3cIEP?xM(~bB1hA8Dh&grl?zwux3CrdhAfxq5f|{BR`kxBkbhn z!y_UjK+3tsP*51XAjW_AF!#yOvAE{!osAtKe1AVO>p@H&SdzgbB1+};y@$>B?wEIr zfU1OZs2kBcV|L?pgU}~dN^y*ZZsTH<@~GmlUt~V1K0xa1u>WW&g(AYW0r_g4%QqPN zBODaEfl@3$-fGar7eo{wANy`w2zpjkSsB?<+tt^X4JODlD!^B9adKYw5WRm@h0K)H zj=BN%5)pYiDi(PA;NU4h5r9HSGs*jhZQ-SObx^T9Uen(V%2DE%r-z5-NR#s8;5EW1V>q_vKF z%x#{q>Jk=$Tl6siUqtd=PN|Cl8pQjKX3l_!961nba3wzUgCucWTzeRb1kf5(?n*h0 zRj>VkVL>`tmN}7^1Hvy@vMP3Ny zL!ZOLwVgskwR`n@4ap6*g^#VRzkKcmyn@P?;f3oWmvuWq`g}$3Z(V0l^5Es2)}j;- z&fD7)N=&W+^5eten2?C0{{;V&$6>%tPfXbB}n0NwQou%_~aCB@5tUU<^4ygAYf9QDw zI)7>Gnc)60HIcVwhvIc{wXq=Lq{;^r{W@w4m>j9GV0|_=whKypkDa<{WB!)h51g=* zEj}xgYN?6}B@|!nkJ1^&IKl_v)>Lf-kaKHNt_p(;IXpakjUiPF^%SVLy(J|@5(`~9 zN~P*?up%Xg-J?g@{dfI;MZbM}gGT5Gp9L$gDq>y6G3evyVppC%qD7W#aoPWtw6wqV z;3<$L9Y4%Y8`Ty+30#P6TSymdUt0kDg+vYJ;~Fz&Zg#fo=TB?HBCqmMH!iN!&|j@^ za2M;QXAThm9{o-NTRz^*x5P7RP_5qA;F+@W>Z;fFm7^d-_NIS(QnS_*!9Nm(1$PKoKt(6EwyE;#WY}sR8EQ%>3#Iw2+^>s5k&yz`M|Yi_ZNM&o`|jZ^Av=g+prCoE zWZ#urX%38yCMDKxh}?_9*plkZXsS<6FDbHwSzbKTF*jRyt}cz$EV5u0Jm*51i>%K{5ToE3Qjq=FSM;Y>wym_(rf8xIlBaFwraHA z?<}?yEdya1XmS543u30!^}IMbd^hrYDM;thpp1U+tR6N!1wdjHy1MjkHIvo@s?>OxWtMnMljRC&ZD%Lz-60)i3v`QSaFl)7 zu!OY#$Wof;*&-3%C?Km8$_c#i1)05sZcrcRI;-6W% zXx?`~5wf!bWz>u-a6br$IyAOTq^-4F#T=`yOjmE`_86XoQQyA+r3+{& zd{VI^Yi9vqyns9*N8c&urjvb-lWBnwut&v5?SK$PY$-c;K`^WMWStG%omKM7$K#Jr zd9^Wk%$%L0_Fl-%WiEzs7JXi#A!T%vp3-as$xxmgWg6^*;0@-ZwaSL5)RUU-Wyf3;>yD#YezM03{8O=K!KlUXZ6=0d!Z~ z=cmj#KlcMui5^|Pg@tX*QDZf^IMLNgm+dKPAZ=s>x7ChRaXdRmk~6}$Giyvx{m8yo zyPG78jT0*?E0H~+*hVO7EuqlFf~7+B5ApE$G?%i-8+9MR1wb;PVJwGIz7S9&C?Mxs z(BGgO3szku^^e-`;~v0fpO7J-*n}-B%dZKcH81MaT6N88-mA?`K)UTv0MatZQK3)J zLvi#_xZx1Ge*@J&SCRs%BtXHj6bQHs{8s4u;RUowIBUt1AG~yk+Och{P-={!~tC|-GixL0YbJH2%#YljsEK_TD= z$(uCU90xK*tRQPYk}fZ}N8L*>|3Ddit*2G2mE_$YD7|~|r?0>pp6^Cy`{F^E0zdfl zlMA?goghCq0;CH*ld3o{i&`i&LZIruPI(So2k5)@KH~+U;>a4%T|Nz01{=seY4fBT zxw}K_$?ebT_Posvtyv#(Zh3T*hv`x_IO!n^9LZVw%I-DVe7wDZ1SsdHPuXzD)2uuw zScO}rR=|;1vq!~)YOA{!1;+3C^t|48|BD4yv6Zvrj|025lm-%z^lB?8;OS|AiTIaM z`A<5dSS!)LAL}vye#v_A|9f{eC>a_2qXhuh`M-8gsxqsWtj7*~eeq%=dP{L)gVK?D zYWPi|*7d&W`(5uo94mNb9h-ef6{VUSdFGSWprj#67Iyl={NGm$_3josj+3~4U5`valLEQdi1X%-NaS28N(#T#>L_c)#_;{smxA4f zTcK`q!5S)CE7RC&Xw#%(3BfXI8XqRlS;~DLoHtZ1fRFK8L?e-0WYPTF0b4@X>=)ZA zU@8+s-qhnht6Bs<1&%<9a3=+wdkWu{Y$%~1Q9&20!+h~N!8=C}+cc9at}LyAYY@CX z(R>**waIcJ0K;BXu;ESUE_JQB?pKNoy$soCOD7sT8|nB-9wKt31x|+rt`!RPZDF0~ zG8Tj=toEOSGOQ%pKMPh<(3SSBT$Y{BtcI(Y9t!jyQZQ+LX*=uvwJYri^5i^eD>bGX zr?_0m;{#fTI>mrOmQJK^J=W3*^kPp~SrAUkSJ>tzfQ(BweKZu#YFFsaT(_+Mq<_d- zhdFs2CFKiFT&PH2zzD-PY%P7J!UJb};aybz*%~o8ql!AO14Ysp)6-j%3Vz<;P}k6m zJh704Z9~JUDCVwXfgt_#e_lW!dqL*;ugC8XckS@oAS^Tdc81^1@H;a6o(Q|M0Q~qJ z6uanj^<_hSQBbBU9OVL;+emgx#w<}R+PGqioNi&mKRJ+wEyz3%He_n-W@1ib z5*Kz4B*5Xd(;Xl4x8ck(@{6({OLmQTl&08k{E3wIvf_cZ6@M+Hkm!~9$X7@fVny(e&A=Od{%fw3Ct!_e#c zhE#KNG!1*|%NrkO0>~|fT8-oYer1>Vj4XAl4%K(9APu2slraL*8)k+q6}l|=>&3Cq zdFFIJ#v6Ay;mIPM{tGMC_CKNHDde9-0CM2W|9l4d{o$@9ej9{khTqQc+ZoR6n&7uH z{9}gSk%482-;rV0AO2gMK|Fx;>$?JnvoybV#r{CbE^*nOs%RcW6I%4=S{o~O%}2s0 zSDWSlkH%sm1*H9}kx0=tDF9xct)7_9N6vM1$6=je1Aw?J%RMDxMHz-?Ukq2|glHs% zO#EVsA{;vV&`jB`g_Hg#%GYqr`g0gRz*|yg3*!t8eHQ8a0f4OA1mApONF8c6rL|h0 zTIFp2McqZd+-3hvP9_g;I>cjPI$GlT(mseuOPr%CMh|&K+D*&o5lBKVbYO+*ZgL!d-Nzcy(5SA zh?ngtZu@IGJ{jMe%Jtmu1tKjVLoj5X18Xy?sZhe0U_r^C$^5Z=aNdlcyl$CGMaC|) zr3s-Ugf>;UB92mx(}O;-AUSk%y4S~Z-J8)}0SWe z_=R9~lQxBcm!&aC8d#rJA8K>Z;%*W^zoQS`8lvSR)wfvRS;Bhk+4ryC5Xc|O|GfJ5 zhs;*u=(iuHT~GL#oo$Pi!M@3#b_~Fyj=*Y$wC$v?LLT#)4#m_f> z`^Kr~ezEbS)+^6tk!9;Mw_jfRxP5MT# zc5#y94W`SDO{LCxKOjv#^Yq`%_?;rNO!3TPd?$=YlEwSc+G0TH@_cs$~BOG zHb|_4KM#rYac+V=^8Y;U^a+JEXc7%TpQ2Alc_Lod(g zgB513uAYQMo`%Yqp%S9tCT;Vv6B2Nw`Q5v0*wA2~-U;Qgpq(yFk!x!(??<1*Ew>vo z-=<2N0?jW{yULPJs^DL@*6i#bu_$Bbki!PLWP|o$sDU-fioqpFtK+P}+kIf5Y>sZh zYj-IH+KArCYhARRN@plzyGzrxf{yZ8ji$JE2IsxH)+0-c6LZqWsX^A{$V)_OSj zq?kknWwUbf*K!~MpD3xoy;id=v||psGp9f7{+Wx+cyRNY13_sWZ>KYl#M?FaOl)Iz z+W`-zlVp8u(QfoQEOi~#(Q17iFbo{_d_8A(Il z%UfIDT!ma%O;p;RdVvS~PX)}!=3;Jsc`5KBH4p}R{Mh^3+h_LWlzJunyhhVpkrtU8 zPIF`3(b;Qt)sgLK%xGW;*DI!to}3>u{~}ER>Mn@;0Z#G%8H<2r1o}*H5KWC;iG6F* zT?DfQBZc*EXfuf$E=_ee{v>(SZ&JiR)Z?O}?(;e)Pg>c7z^g$hL|GfGw zAKG=x-%6ofGyLyLq2B^CV4U9`^4mjxPlVqSVRsgQAO9_(@}#tv-J(Gm5huljHjW{m15K>IN~8@)-`l{-q`BxMJn8DkhKIoG@(!Xewj+jFp1SSx|MmgNK6J+;j`s z?&$0v##ENDpeSQCDQ~OM(vP5wu=ER^TF-5IR!eIL5-cD!JRq<7xnNk33!Dv2NvHG+ zKr2D%yLM+PhS2r1?num2sZ!_|?cM^>>yu(|Z|PRcjN%}GUbNcvb|L%@8#8ZMsR&hS z8Dz@3!1gMX*3aN!)}Z}|8@CqAYP)t1WloIsor|FqIAbca6>AtFEV%0GhBhDD?yfPU zA|ao8oj}`q2f*G6@h(|L#;FAuq3fwsyLE+ZYJvcp#_ZwhmD7cG)=nrR(6%90V9KbB z`aMDH30*Q?3u8rg)-pSOW=5&nO>duF@a|YVF2|mbH7d)HC6$A3tITYwVUEcX)sM*pnRWM}5`e`uSak9DKR3uhfmgL~&j)Sx7ofSv)=g-pV3W;Si0|OsEwAK9aX1w*J3w5N4Eo`mE0fS2CLPGT9 zWb>65zBA2;lxRK6&d;xWb)M|Kst^p?y5+7Uuk%F=n}D&V!3vM1*)GXCC5uFxkf!ZH zwaaeK&aBuLcCI2U>#e_ymAjQw@}xjO^l7ypM;{1@3G!T((9ObK7ZJ5G@?E1$n-S!t zT?Rfi<*%RcR!(~rq)SXT%*IOL4EHDO&D*bhO3Lwl5Zh6{nv}@P*vsl)>S2d%U{SHV z;nbZR`|+~1kTrT0x$yR(4rTyIlCux3tgY=sb^itD>)c0b>^^^=T|gwiz*9ehmlH(Y z0+?9!21y%3!r@#tC|x#$xuh}AOK?(-Zr^i0_PhP04~xmDQZD(!J=k~LG!-N!oc#P& zO;ZW%OY1L>#k_%&-n`1l`C?Rn;A=xII8%q0;Tjzp5dHM4C(L;!w zkZxKMoJF6A6OBkst@EZo1<7oXe0JXWqyaK{eZ9T0Z~}T}<0VGky{7c`!7kMyCEmZdk3aeR3H94p$`-l|9{V-?oCommq;Gi?)};fQa!*hBskiLa>r7 zu^!A!Rw=pa=jT^iT`gKyS4+)FZQ$uRme%ER*U8L#+1vbuje$@cNU*!Cyg{zK^l#!Z z)R&F2ncWfZ!-cW-262oC(vY& z`{$6H`Stdls7skaUR9QX%bRx_{W3h?*>6KtxQ~ccnlY3v)=P%caL;q^*1}oTyga9< zC;f79@4Z`$?9=~%)A^}NcYCL9B9~S#3PdB&j>P5U5Z_pdaTl4P>H!9bsHOve4+nA* zpm0JEzrj2!xc4`Z2J{6*#4J`qLRQtv5WAeb9Br|MHT=(81#pxGqr!+YeTvi?v}+e8 zEsaRDQ5D0vo#pE(b0Mxg5ZFxb%v9j_cOM-9)i>L+fmH`tY3WSQcL6)L&ehe=bc}iO zxg_C13I@?t8EhUu1A;yjIBMw2PgpcHGcyL=rBC+Wy7Z_^?;p2BI(Q-U1|{pgxpe3X7Z|;@P41Q&XnZU-;h_Rm4N~XOnA|%lluo< zlj>JNIisUv>Apal5e#=#m3f~EGSj~Js-x7RK#B*w)CyRk>Bm&Y(ncbjxb;0mSga0Z z-MiV*bOn?xOslJdWM{TYtzpif+|lt#M-+sof~jA>-v7E7a74p*6XZ7%+U;Bhfc_P3 zn0LKn15s2ABOcy{l3iL)Qz3SC`WdYOqrhtY+hfq5kJ9wr1Ul6GXsAI+?K~N86IY8| zH|i{wL`uTSV*xM&#wwv>ztDoEDD+dKE0=;3f>(};o->L{TccxazLn-W)zeR&rp zMLq&4vE;3(nBdR|@CTQd6S7$Ja#OKv0Kicixm4p@Z50e{qE-XE3_+R-D2bYs$uS~V zf{r0;>$Vo&I{Fl8oL>=w=l|TOE5y!D9j2&eAVR2wL(^dPpZ4oZ!@+%n)1L(+8t=Rg z86Xm!lXa16L0S8YA3(iI})JGAAKr8 z#5(r<3=Oo@>Q>-JdOACw0}ca*9SBPBNpOT&j2j0u%}u6UCE3044AC&dz(A%d2UOg2 zgqAzrKWn;`KyZ#}8BZY2#E>i%Oi%JhCt8| zUP(G2*l{YY3mxa)f2+9d)|qy9e^85uOG-FUij?SU+F)jaq73#*W9W+jrR=zq5$r>= zv$KwF1#+?!IN)0bC!uR=Ym--q!VfQ;$Kgal0o{jyhDqw?>e}yleIl~u_G69A$1Q+@ zeA_kw2adg^Fj)-}k(-_F=XCCXIuN#^Vy*k`+D*983mD8!>(LQ8mHgOtJHXBuu2i+p z83qNl+#_Dhc=jr2Cy^(0W`0`bbVRhqZGa#cI7+{UNptDgsZ){>eWw!pveb^IgB<(z z``eJ=0n7==99+l%+}&%#(BBsBoaOU6n?j-N#^&T`8oNwqeFBMnpGP`Jl<|1G?ackx zV6l}uZ^}1Ju7QBubHeDzNGd#EkS&SV5lU#ak8hu_jn$H?w(`RmS*Pizq#OaU3b7z@ ziJeI<&_MaDtYwZuX-ZVEiH{H{y3FB@gfN;3> z9eeiD_f4jxxPidUM=-Ay6~LTvAWIp;_E1pd;RuYz_2AAo1+BKRqf`?co4T8x??H3~ z@D&_%P_6uw5h*eO>Xh(tc{edVJ!cwGVV90&Tt1OsQ1Ex}bQ|+XfYOKu_pfbchY&+$ zEeMMI($amhLz{`K@&;u0qWkJ{ig0W4%XQ{vR#q`^6L5lEm#*v%5@%$xveC{*z)cuf zsTLDV(gA%2EBM#&hB1k1t->&E!z(7LcoD7Lzp??%wGRQMDg}U0Kpl%AAOgs@)YCN{ z0pIeOlp;=k!z554Tdd2&*Iwy_)3bzXl8l@Z|D90oixWr*0(>lzC331snR#sUEipWCJ!H;1KGZrm1$|ON@sb}C;Jj}Q`7F%Rb$c`7`HiC&<+loAf|FspYur7 z**|vusx5IEj{YCtj=FTL>4C%vkZ;;=h586e%w}eWXZi9r){9L+0sx`)ffy+QyAic9 zMhr%Pdq_y3#PR~@+*gwID#@vd)?Jg{2un}GfiK=N3g&QBv_9Xthas8b;^Iy|KBZsj zp#S!M%;1QIj{&UMrKPtQGCHu6%rOH~;sxRg?LiR@PYu<_7p&nL~JbThrB8;A1v?(;Z>_TKPzBf*^Po2rYtx*z&?t48a1srdY9FGZS(f+ z%qsu2!m(l2bX<<+Zq$pEq%$=mJw2e-_a%b}z6cDe-Jy3kUmFO+a6)CsrRwf@_DJxu zZc3eJ>XR_I#UC29%EIK2_-atVcy~RqcF)86g_f1CE`6OaUjhI2+P7U89}Tc(DwSRn zruy~7#*K|+?t|Q-wN#MSYVrEzD0dX4LV#f$&98#h;*J3Pj6+u{K}`JH9aAf-f!YUC8ptPu!xrok5DWd&3oxA9y6org zPosyf<7-0}(syJk^4p+k2{3v_G&%^*^|sW}Q;V$*2^ieweJAdVAD8(G1Lh@GzKH@Dy^X4H8q11#ci~(7NdEz0Mz^)|csnK*rC?__@ zLevrHd}UJdWzmmHz_ErWT?V6|mv+wk3$5v)I-<%c;ms<5M>k(l*C!dWE)RZ$EEQ4Q zD0*(9f~p_CBQPVh-nk2ZL>61k-DTqx z&CcKlL_A&@ROSWRfwK0u5`PiMw>+J=OM7|PW(pl$>X{xCq)ufF-zT8bbOBhrOPQCI_U^r}Fb{>Fr4Xu~iJ z=NH+vww)+G@Zz%d3oXw*-|yRnrV%*c&0q_OcOOQf(RkR=!{aKzJ(!v$NeiBg0IO?m zbw)4E#(By3sbYnPz^SI-N+1YxgRL{3kWhn-)EQCCxQ z>xJx>b8)RFs>u>Kf(Qui1lOj6IGvAz(ZE+NC4Lj({ef8r55U?C8wZ^^Sw$Di(dq%v zsRM2P=$A$X)&R%?S9a|C1c)2~{#)B!X?}J#UhMcJ^6#C&c6Mg1wo8}M@#7HeQmE9p zzUkN#ck$L}YS5^1x4VajJ{NzmRjB|>n)V%}UpL*{GzDBaW)UYTQ?tr*>_TS5QdH2| zyo;^R3Wk~lQ#i3*G_eERprZWIBL+O92rdE2*QQAIAtZE=o4Pod_A7#!WNk!* z76jJNeFsjqJG&stxD~B)eyy(MSOa04I5yRv!xo7+A2x?T-z)9?B7hhAc`hd&85e7&;DujFC;C zA&9e;Dlio@H$Dnn`)gvg)Ej{I1*@2>RFT~kyKTV0 z3>bfSV^cAdSzD{J&S1Vua{;Rb3hWyjRbz+V-~8t!*l|#40%5Z@`^c^XlnmG~?d{FS z$M@}&JA=6gwBVL;!l3zLIs@hvqR)xR%aRfWiV$s$%h8!)W^_AYb@DPuz}l730)}}G z9Wv0=Y?0?JTr}00;hizJuo6AD_Q_VNDrge09%L#&|0EJMTADdahfGILv!mXSTq8is z0o0a2*Of9i9n&{$e15y&zTY71aD;HM(sJ?Kgm*MhDwvz=$sLRi`x037phOJPT?qhN zD=EQ2O6l#uF~hug(B-Z>MP45Uxh&u1@4#-Gezm3&w3`146;@e!x3TNv3@vZcf`LTR zot?a^L$F}7Dn<-#Y;1a@l3@=9MqqDaL^L}~1NIBpB2lP?WMGNw*jTe5M9clNOHXLu z#!A}(Mu%j`RiC}xmpy9NVTPe!IhOu-L@5#oE&%XwrDfs8qOMP$;+06Ma&p<_<*Ub| zXszQdf+wSe;W{^M@t)egubbaB(GCUWX8hGBa~6C@3OT1#MZ{ z1!pvMg!YSQWz1E-2t}5Pg}P+&fVr+3v0ytVb4u zy*>-%!1ejN#&+sJTnd@Pi_+4~hff2@0%4zQtV|z|YH^jD`I;4}0O5Paa=)!q9(_$< zs;4nRke-~9r2ISTMwz*^ zvcchiYiQCw7~jEPU~HbJmEUUH`|)Eg$YR}8jA+aN)(cEoo`Jq0{XKPPN6MQ0AM86o zy7AdfUK}AC0?PJV#h^ae_-1Z(wG%{kVL#W<(D=5vpT!A@c1R2YJ4II2gHK%m#Pho* zxpZEGk}EqKDmASWZd&5pcbFLXy`hPmxQ9ts4%nccy@tU4F*nmvA}25?sI10v^OCMY z$}wbqqjUjOaw zC=QR+iI3bM=`&9&sN$ z7+DU%Eii|H@72&x?YMR8)&*+lB)U;fB2im)vxm#BWN78v;=ac6$UTwZ&VyOlAVmo9 zgzN-CZ7P1PaBwrks?J&|q3#5PGZ!a;XfPhW4Mopcc91pJSbZ3vltGitWPfVUA+W?N zY_#+hjP`*eKc+#15;NJwYd`*^=@k;zplX_}^R{B%EuV3`1ruIprlS)0`o_`pl12w! z<#@JE&NRZeZ?A7Y{|nqSs5)4>%}J(x4OKd4=C!=?wP_C^Aj8deTzPZ}1e^~Oe}~;m zv~r#G5}q-rlpTngKNWW>HPFr9KgWp;HSC@UN}OhX^oEA>kPF}$NQrF8)a2=d*pow^ zWwU?RuVBD{>Nxkx4KhPt7hlv^k7^ew@^J#pXKtRm>+A2kAA)Z5=fTH^QYG%DHZKf4 z(d)d7b5{OVS08ZG1TsO5H>eN?DoqJ)*mzwX6`@ug@bnl&!K zG5H=0kuEO^sHn90AnJ266Tddn*S7R{%v`rpko)oa{93KWs{d&@mF%W0b#ZZt18yZY zznl(EpG`qh5TR+tCXXHN1Z;eBGTVj%$1dEP1*x4%@=j@+X6=0SyEUn?UCi>-RQJaG z?XBMycNl>Ckq-RLZ73wfwq9Q;BX5SzZ_s8a>^^c67wYNMx|~tQ{P-&(WGm-~D{%s> zZ!M!W&s%MMM$>nsmfA#n1mF7laa>gP9NOBua4&bHQNxWyV~*Z!{DO`dBY{D5QH7iC zU`l~{9zz_A$iwjHxQydo_Rdb zlXbpFE~BG4K1vJ#&WR#$-NmLLlblfm)Hp{Bc*BBmkFX(`P^N(7j?Q$rWWYGk?NKN# zb+&5&j)a)sMVg#|gid z2jFXC_Z)6gVU?H5k>;9tXEUDN`EnV26ES0@kD!l_Ehu1gyU1Y{scueAdA4XBWRfQU z6IX;s+1$9;7u)v&-bdGU*_lK^%H z)NWmFf1suq_cb|}O@7{nf0P@)()Bn~gh}=uMOC4=9IWt>?!F-luX&sNPb09FCr-u8 z-Jy`A8)p=O2=Z6tPF6MS?b{#iSRv*XRk-6oQOLG65_IfhazA) zik(Zdw6=a;s#-)E2D;zOVnF|hHAW{QOKyyC1Bb>$ z>bOI5ia$A#M$ zzk8Fr6HrZ?5BEgzw`oJ@(cIi3yFPWbM-IGAWGey0ZmZ4?DxU7HtB1OrHIqhiO-MsM zTRj?i5oH z42Rm<0dZXekt=m2B|qMW56wOAus0IAU89B==#-R{MC%fgD+nCJZ_uaoI{R|;^{#>w zWyj$Ty4k_^l#??SmAL1-KS?KhTK+Q>Rv|viPE8%r;F=JmPYm4;-^SZBHDdE!s`)&j zH$?4xLHzFeBDsa(Tg4E#@Yp+PXBjR9JYmB}V~7@vz~rGhx#4g6IXDYN+^XW6#2<8n z(?3$6ctsf@SF49cg19Q=d(^e}ebVXTQ3(nNcs*5i!^8mg%2wM&*GJ%oXh&-AY~y=J zmbhD^a#xujk%(O%>L2MrNH58@?XVZBxj1jI3|9j18`X8J>f1rqGCX}AsvDPeU)*Gc z4OLKcrGM|PSDVNircC!{MS^cE^>SATM6^vDz1iCIQ-gu@4QIJR!+qBxn+$zNjd4v2 z^4I1b{Ro7}KeR4O1?S9$!XAlhz0L@)AUGR=r#Ys@^gQe1h*DGHmw#CwgtxA|@_J)t zZm#d^*IVms%t-K~bcDRqH^GgKCC*&4z}TduY~x0D?S(6zXZey}1dM-91+pxXMeEy% zx=cSg4I0BBs|rzLERz}H=BCa*8LF+lM^L}X-^b^p#WCd|!7Ckg+S{M;7K7~mHdK}r zm)!jh_m>S_3nP1X3QA2=xVzvQ35E>+N${y;WL}dD{}_CkY!FCu}(R}*S+xZGGQko6xyA9R%qis zI9w2Fo*C|T5kj^yE5;|!y}i4@w!=Ov<+?9YgF_$WaP9Go)rsV%qBg>1Z?2)U@0r`3 zI;*G2>zx;b)4pmm*Feh**#tmbI5D#)H@5Xbq|%sY-U=(Lk7vgNtEB+viqA-G5!7e1 zvIKz_k1VcJX&$vGq~xA{6~G8k(8Ij74mtz++E+{}N#>HkkHO=clYYD6#suG$wnkhv zxR@8~r1=|DK5>_dM|0TT4xUknaLxm0f}gRth$URV&isIK$hh(k2Lyoq)yFkVKeCqK8wg4O|>#ydHZ~{&&{EEsD$az0VLmTuR4V#2Fsby|t zsBGj!mzLnImYMJ=3zR=mUnEmi?!E$_Fv0!vU$CdvZY-zqDnChlfjq&^AtN&60elvF zAM!-+Iz2$J3pQ8fdsr;{DvgVxl{`DB+Z^Tsrv=g5&h;MvFMW$p(JeHmHW zr#F6xAz4uPJUlF%)QU3$izzN{WUZ2)x~jjhX5w1}0|Bhg9oy0enp>LLi&TT;XuZRs zPIp>QUz>t(w>`B5oJPji#JX-Mb=s8}*(7H`N$Ln=j4RYY#@W(qAH01*(e&eslgsP3 z_07u44)i42$Psr0W5)Vb$cLeFMxzT}(IqTvgme0ga;_307f(I0 zSIoNO(@dQ-5_{DBA8Ch17Sp!m@Z9`-Zh@Psbko4|#aA~^n<_9^74RllAiYAT2w019cs0eLX0*bw{^Fw)K8&dLV&IWE zIx~w%&`U~TI;uPu?+bP6IyyUFS@vEAa%T}~v=IueE;hj#T%pXHTbELg%geJj2Nr`s z=w|Wz3=?Udx*q1{L?UH&HrIK)*m@+CrHo$b(X`XIg&3|g1(Vxu>S{bbMY)i%VUEtC zY?>96sS?oFHS=^|v>zU&ivUPTM`-&m+n`i)&~4xDhrsYFgef;Sm-1>3(jP#!Faj=@ zI5^{mAVLjm6mmZh2LpM~Fs~Fm`QB+Cqfr|u36sg>;|Y`S>`+fd(z+mhfAWgZscS}2 zK=(P9Lg+e2xq0J;fqtrzL^U2@@g^qzR?WAH4>{ZtEnl?Y%oOvCwJjsz=Burr8cIpU z>GE0^Y3`^FgpG`cpMT)nH`^Y$u2dDmmoH(Uug!fh(t$h$HPw<50M^_ZfAQi)_lFOE z!PQ*A+yn?7LPo0&y4rO6s(NNIwze)I;Q(x&F*jj_e~}%yz6}L5?%T_(T~NFOT(0$XL8fnV8GOtm@4(%M#CiPX6A3!a zWfcdFl&lpxMpsw2>Ix{wTqqXFv7*7DaDiMFm?&NE`w%kaW@cW5ilki`l`4>H@#P)C zc{7Pm5Op*~P6GGcQ^p2WfAva3R&4Xya|Q6X{&Aooc1-|o7G zm){gONq&wj{@g5BL~rFEUA2VP3P7@lM__W;ztq=f2{l%fx(1% z*y}XDZ!v4&2+yB}6Ll3}u}wdn{Wv`BT)SrRw0u>Lwn}qy^QCpq(9g?Bn_HVp3`Rii zCrLVoBJ)0U(x>X1PZh+rzQzuMS{Pp5jjxe2gCi15a$(Uy)nT}pumOjc2j^b6y1Py_ zM`6Hnx0_m-5mJWyd6b`c03;eVbxyu^ zSQ1f&&kyAeG68tkWCgfDdfxJKgU@G7{^~4!eHy+QQoiByS_Y0o1HFs71#NX z`R(^*PdU<@^LDKdt3I(9AWD(n`{T?JuLWROoV!W#=3CU>kM_Uws z9^;)T-WwsHq0z3*N^7-Ukz{R3su!iHps#qOs@fRaNJ3!>c=Xp^U)G4sjIV;nRgWtg z+lGP!5vwnBU9x4h20X-30Mr1TL8aV5u5B`O!*b4T&e_Ey_)sTkucs89etg(X>>+WF>`V8QTAn3k782-pg*B~fsF{4v_aoV zQ3?0@UJ=S%>e^_M3^iLF4)yh&>t0)VHDWz}k;AWg(NZi<3)<)i6G59B4Qx;fdS%;O zJ=u1*4vpS9L#VrgBnDfWv2E$3fdT7P2d5HJDIRa-Gw)Gsj@=DuKZv%6M<@JD2q)qC z%r7>wp`H#+riC)iU76bS2F^5q`-MQ?L#wmfrirqY0(9y#*@`e7K(?Z#z4+{kta4}H ztbpAWu)9LNKSjZ}>e~hI-l5{G(>ht6xT{_`M@bD5*!udum?WAx^D>YGN$(=|GTQP+WYcG5$-7;GtOs8QB_=tIG2*+sG!7BBG0Ft^4iU*WivS zYgr6x;^rsFb`*vLmVNWqf2rgYfPWdq$jdv0P%Z?AjG#0#6&<>HNg#OP+V%&{NnG#8 zo)v~{IJK%kbErkz*x4sfpF(&QbfAmT(ySrkO!caSY~GPTH;4nE4?)4~>)Ns#PE=H# z9)Cg}@jOvQWzE#x0+|q`0-39}r|~S(472fdS#l*)fR)>G5n5PdyG%x=rveW{A1OOT zrm7{7sDw2Jr5Lt7&{K9{%+oZZ+rCxnRpZ#RE{B(odmtVh^E@P>hg7l?+i~zbIV14- z>G!;VJ*|x108oxW!;OXd6Htge#m5Qc(9Cc2GE*_iY%q-t1?gc0g~0=yyovq+@Zp9I zId_sCc`OIw0pg)v)CG+1gew5BtImH|0~KNDfjp$P*=+eAy#VElbY8j#ayToufec#TJdgxi-bntW_SP2+z<+-anaqnBx(_d4pb-*~pi)UXZMIVCF7o_BhuGU(%lYmqbsRl3 zTPTc5hnIeBN`g)yNMgStLc@u%@pDKN0ANfqxw3#UMufh6CHq86Fjo8b~UYkr~ zX~hCl{hQ}ms1N}7HNN$Ad&tPcJrfBPQ`2)sKF(EBG>D|C{i)OQyjSIl4RbOx zu}f_K?!e6-#n5ab2JQruAAw9~D|PeGOQsv}`$j)$+xU*}8$Yle#m>QTWEc0TiT>$6 z=n=qfy`P1=2c>~QrwB}$q)Q)-CLXWuKpbCNkt`_S;K^3vMft+U$KI*O>ZPe%>p^;<6zaw*ay7kFI`_J@HNuz6;2(Ca5fL<#P^d1X zwKyg>Sehmb-#2RbRDkqNxw~uJzP(RYSAZq5L}x&M7*rO_$*E&LRnTFlRt$T%xVqNZUO{WYo=%y=ZGL$G8N#*NH7iCRYwc?5(r(QO zm~dn&x}vZPK3;8EaR3#j8O?v>+~Wg;$d-LuMzojbUS4m%qbT>(IaRTqnY`5b{8p4?0X?fMcM$RcHL8ZLvc8p z!U_C&=fFU$dt?G|i9-YwgTi^C zo-L7BpxvwtJqui_n@Kg+`s!}+k{3ayCks)=w*k$e<@S_0l^chHTWeLCgm>v+M?`( ztgNgmT^Vaxq%?+SrK~|~4|&cHTB3B^E2_T+6&PbWQ;)xWnu=~xo8YR-5mrvX7K6bvW4jEptk@BxJ^GjpGpf*E`is($gT^Rx&b8BkO%Ry!8e zp-s@q)AJ1_YN+woU$!wQ#5xGpO1=gGT~Or*DFSH~;*E}hlH_!X|L8eG0eHaF8oHi# z7Be&k9Ik~ocj}+OlU-OK9aqjNDl$=JVbX|#%B)Ag>%rNu4BgybGoc5+{)_`^AwWs@ z+T^D1gy&gv30?YC9UYzAx+0qjcQQh*CooLKxM{1a*y`%Z%aOam=L2J7Ze<0zEo_kU zJLt+emm&eEHKUhz>&HB&pml#~9jlFIVIjlQ+uR)3GXr!fq4T`-hU1gSxdsS!j_;%d z;H}{gWt(d-&SBsCoTnRX)Qxck*}Nk~Rq&SU!EVq7+1uHvM_`=gcfkGEzZSL9Suu*E?S}h9=mptI)({VF6_|s|a zZ}nfolv3N?D)WKir%+i;Ds5rxqH^Gmh094t!JvR`@R=^q7H2WwQKUIi+8X9XhbkgDg%THVQ#b<=&SGI$8I7&LvBIdowqC#UufQj*_~dA4@=L+efv5e}Z{ zLWV~Qc~YZkGVNaT=2~1lJ9;I>5kX zLb(Y!fdcL3F@Gm2;m18cPe7!%yxu=D@@F;Ky5tVx6eS#j6>}`cKq19_#lTt{D#RBr zo`f`6_PEzMa1Gx;3-_^@XRE8P>^jT%z@+qid9)8|SEcy_gJN0d_&Cv%)uTOG6%tNP zcK~AVGB`MLRUyGpBq7jg-ou{{9U>oa%Ye10$~dWvUb+O_*1=)+186GUyU(2?`U%oj zb)e`t!9~Tn%$$R+(su7)QP)Y}49Xn34IQeK+PV%~sp$^+(SBZ!nhXNL2m&BcK4es? zZJpGLh&fL?{b;mRr>Xu^fT4FtN3!HIV46^MqeZL*0oDU((3vrw;c0pS4JaZij}~z@ zj4Rs@q_HbHpZ=}juc5hK59X?TxSM4mMM%|c&MlBWS0Shc{cJIpZ9{vYx&>!YTU`fA zri~a(ikq3ZAA`!G~>^UW;9%=!hj zo-x8Wy`S)MzHy_D&srtTM2Y{JP#nX)7NRDf<%av;6E7-+x|_1vRX1C(sn;aHN{+>p z@bI5A4_x;GEE+`e9LRrg%ID$5n}V@gg8S|qC}maB86gm&!=a_ZODRBaRNW+o=FZi# zzDB5S+_c(Ac-;~$aFDzSA+N%nua`Q~oS(o4rv8k|D+y|L!C&&4}aNQ5{T(W3?|jyr&!OHEB>E0QcYS;noG z+c;WfwjW&P&5+(^Z9@oaAVTSxxu^E!bFcIK&Y{yih~>gH*h)duG}O(6dw6`uW{_E| zjvUVRSVmcyFT_Q>_990PA3pc=Po#7vdlAeUaq#_nUVN8SNk{mjUN5;B1ZT0*S}zJ~3FyruKf z)-`m$hw!qWKqmil5bMBk_#u%me$`-WUiCC0uSTJK_abcp)l`{(;BH3wQ2=+V5kc~# zFVBfaHuX+rZ2erA!W6g((V{vg`8i@-3^~&a3AVy;8#H(5lktbEsh&}_ckB{ zC%e*!KR*2G9ST2N(~Ddi1hxf8^{vtpQXSA1o=DPZ+B#Z@GJn`N5j#3*1W9(pZDZaG z`ObuwIsOKNAr2=~Y#MI(_skfJ=FOgfUUePFoyFQpY1vCpGq&!uD7<1MPz+R?m$ z45YQdZ1C3GB0m6+{6FQkkiVvHM?$QB|5Dp?^xuDmKZJ3s{nx*u#KZo#znY`|%Nv8+ zi2CO#{_AHP|MOHe@|~X_{Bx*}{%dM%{qN`G`hRvg^$n>qRCmX&j7*#{`}BbgrNM0? zS5~?=mT+reYPwc~Yo}viTwlA${@2B0=Y7lRU?|h`pnaMa;0%;nnLB>zJ&oIvNBY)f z?WE7q)^^Cqd?rtgp>)MBGfZBi8Oee-%m0EVIVHR>nzTbfWn7bgz7W{(f-q5w@cGqQ?^s*}$`Eebq}%XgDMg z^>RLjIOn?J8J5jWdh8S;cp`czk6NQVf?P9oqeH7{-zkn&Uo0*z&OR&#n$&~h_QbwU z+a#!eXR7_iCFD@{$Kcg;GA@qor3~AZ$~DLH>le;bEqRPEU0b*9=Hz=#Kl}NF*Pv|3 zbM_LoQ_-p7q6wYOuHXLJ5wuQ~s=q8}(@l_^*9wHXE-@6X<%cPo(nY;yt&Q`-6+6$) z{>HFj4UKMHGrcd7E1`3e_Q1bIM&|dei}s_>oa-q1tDiJK9(p@WN=L=QKq;noy?Eb7 zrv)?-!n4ULA$k0QVs^1;@^8(>IQakQ9%fQmlXXtVOEcrxwlNqH9`eY51T6$yUAH{X zT9;St7oPZpx+uU}6je?#H^y{zu$obJFD%$kC3I3wj^cCESh6g8)bU5Q8yFO7_djmJ zKZE~&!_a?WIZ)kieJh+f`!v+S&#yv%#p$4Nz_$IY^AC$hqC?=(fiM>@o3n;C(_SvB17?WK(;2t&&lv}GHgu*= z@N+W!oDBbi7J`}sRo9;rhYdV#jxyD%Q_a80;nG1_O#WP2JAJM@u*Ev55PLJ~i4XQ# zO7^50j$=n74|{~{Aa9;{qMLglXX(gdTpy6=K!#rTtZ31X`|jf~3BGg%YjZ0(rj z)V%feCW@Z|VJ@(-{>ld9JO_^QwpCYFWxxeqJUCP0Ela1$WBc-WrmIe}iQZj)LMpv! zRbWwF1)uKxVC>D)v=@cgihnXHFzoeBE~i>+s3VHbuGceuYhUjh?;9ml=sS>y(%Dkd`K??PA?+Tw|S)*i~;V5H_N%5=l9nxHtQQm1z^^e9Icsoqa zj8|Bi5JSIUk^7?^q|*}Up$0}q0`!Gc3QfsvNBiB)q>Hk$$`-_Vl2uymTsZvP40<^u zVrJ!&s}&FyrL=tT2}yymP4C@uFOSq6<(lpuG- zV7>`ir#oON;@j#AQ!HF)*y3i5hNdrBnJ0nf5}a>Bv@9zjo$>tR2SFee*O2Z^KM+?I zL@dvqtXIz%0%(e5SWBvcDRSv=TcZGT(r~K|R1dn06rJ-=G9xB78CiLdt}ft6%)irI zzumMtutu-z`rK!^u@d(D8#jw>kc(U_zKXn2YX>i|@s`8Okhj@<|J4w6lnz-p zxVXuK^@9j!ES$Xdl#VR#2-Vp(a^EXlg_M6 zeK~kFl}4(>VikORu-x2qj-q>hiknM74oSW{_e6S!M19O{aH=STzLHwW2M1Q@tZL4g zS^RY8%?B3$ntovw4ZV^mOO?tDg(Lc%^l_h{ z$)Dd9YG?+DvbJpsCU zv#x(xk{^esj#}B&&a7_vsrIAt&LKB5RPHRx`cBBn%iKJtU+4KcWc6cT*+j1{t)PfZ z0RI>-Nv7YxdoB4`h0aw|s&M|Y;2^t&+$pYQe(8eoyH_-iOAXzIY7hJZes7mPqhJK@ z;(}iKjz_MpnyRQ;(8NriC;11>jg%E2PA**l{kC>(+=d8u;JY^I;fi>}6?H63a#9LT zpB#9R?3FnqW!g#dUZzz_j@TF(k;DIu6ZH6r^+0Cl|1Z+)|1QV=WK(PzrT;6>y*_lX z8bhq2C_Nn?9kPy-JEQWUx2IP@;kE#u1)pL3A6L62M5IpNFSHl^v9r?e&v$!=A6c26 ze|B&1N=8%Og=PhnQezRLp=7y{ApNJ2pL?60`S|OpEC;%I$sJU|G*X8#{$5BW{AWTC zd1OCJ1zQvM=VbUf8MY<@^5uWvWMJpy4=gCaRynw5?|5*Y7?e&dln^Jrc5)a%}Wx$CadhL zN%T%98aro6*i>2@8Fg*Cg}U(sNIWQexv@s73ii#*&Iy+-v8zH}r)pz{Eyr?mDj9nF zcls#r44lJx)D;zVS><xM2%dF`T=m3bk)+MnR9uRMalnkz=u)1McL6FcdXF3fon!a}y%T)cv2HBJo5g__`( z)*QC(XTImQFY|pCjCwi$@R?C~X<4@0%%`uF*eXHP*_GJ~Ze%F}N%*s{UaECIDQP-z zz3F6C>M!^;5{c?8jVyjL9>=K~ z`ey6WvnvmyYJ37RvcnsFt+%m8M?%&D(5p@BgLhfz&O`4wsA$g`=fW~&%mW(8lHl(fX|^#4S!VIMfdqwl^0Rr}*Vh|T{4*TuiPP=4-% zTdV8mBSwbd=VbUf8Gc%Zp9^7Y6~LFDRuLJA{{yR74=$Rc48?1w>#p(EcK}N(FX##A zUF7+U8YBp0_Pj2XRCKr@UXdya#ao-9LVl@nEhxeJ%5=1K%=F||u0Pqa=^;8tQI@Me^<3z78QGGCf$OZ_G5KG3 abslVM9}iL`zdu24K}%ilQsIRg_x~4Cxjq~K diff --git a/src/unilab/assets/robots/stewart/scene.xml b/src/unilab/assets/robots/stewart/scene.xml index 2887522a6..a3a8b7f41 100644 --- a/src/unilab/assets/robots/stewart/scene.xml +++ b/src/unilab/assets/robots/stewart/scene.xml @@ -1,12 +1,19 @@