diff --git a/docs/sphinx/source/adr/ADR-0000-index.md b/docs/sphinx/source/adr/ADR-0000-index.md index 9bab70067..ba198abf5 100644 --- a/docs/sphinx/source/adr/ADR-0000-index.md +++ b/docs/sphinx/source/adr/ADR-0000-index.md @@ -32,6 +32,7 @@ orphan: true | [ADR-0018 Collector Tensor Device Sole Placement Knob](ADR-0018-collector-tensor-device-sole-placement-knob.md) | Training config / off-policy placement | Proposed | | [ADR-0019 CUDA Learner Forbids CPU Carrier](ADR-0019-cuda-learner-forbids-cpu-carrier.md) | Training config / carrier pairing contract | Proposed | | [ADR-0020 FlashSAC DP Graph Ownership And Wire Precision](ADR-0020-flashsac-dp-graph-and-wire-precision.md) | RL algorithm / distributed runtime | Proposed | +| [ADR-0021 Genesis Device-Resident DR Lifecycle](ADR-0021-genesis-device-resident-dr-lifecycle.md) | Backend DR / Manager reset plane | Accepted | | [ADR-0231 Integration-only K=1 MJWarp Collector CUDA Graph](ADR-0231-k1-mjwarp-collector-cuda-graph-integration-only.md) | Collector CUDA graph / historical pre-registration | Rejected / no-go at M2 | | [ADR-0238 State-Update-Only CUDA Graph Boundary](ADR-0238-state-update-only-cuda-graph-boundary.md) | Manager CUDA graph / historical pre-registration | Rejected / no-go | diff --git a/docs/sphinx/source/adr/ADR-0021-genesis-device-resident-dr-lifecycle.md b/docs/sphinx/source/adr/ADR-0021-genesis-device-resident-dr-lifecycle.md new file mode 100644 index 000000000..b5d6a164e --- /dev/null +++ b/docs/sphinx/source/adr/ADR-0021-genesis-device-resident-dr-lifecycle.md @@ -0,0 +1,107 @@ +--- +orphan: true +--- + +# ADR-0021 Genesis Device-Resident DR Lifecycle + +语言: 简体中文 + +- Status: Accepted +- Date: 2026-10-26 +- Owners: Simulation backend / Manager runtime maintainers +- Supersedes: None +- Superseded by: None +- Roadmap: [Issue #2520](https://github.com/Motphys/UniLab/issues/2520); + FlashMimic subset extension: + [Issue #2521](https://github.com/Motphys/UniLab/issues/2521) + +## Context + +Genesis previously carried two runtime-DR profiles. The whole-MJCF exact-CUDA lane exposed a +narrow device-resident packed reset contract, while portable `entity_assets` scenes retained +host-plane model-field setters and a portable-only wrench path. Capability wording also mixed +payload plane with lifecycle support: portable host setters could look like Manager runtime DR +even though the Manager reset transaction requires a device-resident packed commit. + +This crossed the public `SimBackend` DR capability boundary and the Manager reset/interval +execution contract. UniSim issue #2520 therefore approved one adapter-level lifecycle decision +before this integration; the authoritative adapter ADR and implementation live in `unisim-core` +(`docs/en/adr-genesis-dr-lifecycle.md`). + +## Decision + +UniLab consumes the unified contract as device-only Genesis runtime DR: + +- Genesis declares reset terms only in `supported_device_reset_terms`, and interval terms only + in `supported_device_interval_terms`. The host term sets are empty, including for portable + scenes and sessions without the exact CUDA/AMDGPU lifecycle. +- Manager capability negotiation selects the plane from the backend-bound tensor device. It never + treats default host capability wording as a fallback for a device-resident runtime. +- The packed reset device is declared before cold event-term binds, so owner terms negotiate + against the same plane that will commit after backend materialization. +- Reset DR is limited to the audited device lifecycle (`body_mass`, `base_mass_delta`, + `body_ipos`, `base_com_offset`, `kp`, `kd`); interval DR is limited to `body_force`, + `body_torque`, world-frame root linear velocity delta, and root-local angular velocity delta. + Unsupported, host-plane, mixed-plane, and CPU/host payloads fail closed with the term and plane + named. +- Genesis owner YAML enables `pd_gains` only because that device-plane packed lane declares + `kp`/`kd`. It does not re-enable broad portable host setters, geometry, gravity, DOF, or the + opaque legacy push descriptor. Startup body mass and the six-axis velocity-delta event use + their audited public device lifecycle; absolute `geom_friction` remains intentionally + fail-closed because upstream Genesis has no equivalent per-environment absolute friction API. +- Event operands and interval body IDs are lifted to the backend-bound device before entering the + public interval plan. A host RNG or NumPy input at a Manager leaf is a cold compatibility edge, + not permission to execute runtime DR on CPU. + +## Stable Contracts + +- `DomainRandomizationCapabilities` plane sets are authoritative; aggregate projections keep + `dr.reset.*`/`dr.interval.*` and `dr.device.reset.*`/`dr.device.interval.*` independent. +- The Manager packed reset transaction declares its device before reset event construction. +- Genesis interval dispatch uses public `IntervalTermOp` values only; backend IDs are Torch int64 + on the backend device and wrench payloads are float32 on that same device. +- Genesis owners cannot declare backend capability by comments or backend-name branches; owner + YAML may only enable terms whose plane the current backend declares. + +## Alternatives Considered + +- **Default to the historical host plane before packed-device declaration:** rejected. It made a + portable host setter look valid during cold binds and could accept a term that the eventual + device commit must reject. +- **Convert host operands inside each backend adapter:** rejected as the Manager contract. It + reintroduces hidden hot-path H2D detours and plane ambiguity at every adapter boundary. +- **Keep `events.pd_gains: null` until every Genesis reset term is device-capable:** rejected. + `kp`/`kd` have one audited lifecycle; unrelated unsupported terms already fail closed and need + no owner-level blanket disable. + +## Consequences + +- Genesis SAC and PPO walk owners have the same PD-gain DR default as their MuJoCo counterparts. +- The SAC G1 motion-tracking owner carries the FlashMimic-compatible DR subset: startup + pelvis/torso mass scaling without inertia recomputation, torso COM and full-actuator PD gain + reset writes, and six-axis velocity kicks. Its foot-friction event stays disabled. +- Portable Genesis scenes have a smaller reset contract, and removed portable host runtime paths + have no compatibility alias in UniLab. +- Owners must sample device-resident operands on the Manager RNG/backend device. CPU-only Genesis + sessions do not advertise runtime DR. +- Adding a Genesis reset or interval term requires first extending the audited device lifecycle + and capability declaration in `unisim-core`, then updating owners/docs here. + +## Evidence In Repo + +- `src/unilab/base/entity.py` — interval plane negotiation and device-resident interval ops. +- `src/unilab/envs/manager_based_rl_env.py` — packed reset device declaration before event binds. +- `src/unilab/base/reset_state.py` — plane-aware capability checks. +- `src/unilab/conf/{ppo,sac}/task/g1_walk_flat/genesis.yaml` — walk owner parity. +- `src/unilab/conf/sac/task/g1_motion_tracking/genesis.yaml` — FlashMimic DR subset owner. +- `tests/base/test_genesis_backend.py`; `tests/envs/mdp/test_events.py`; + `tests/envs/locomotion/g1/test_g1_owner_contract.py`. +- Generated backend support evidence: + `docs/sphinx/source/{en,zh_CN}/5-reference/5-support_matrix.md`. +- Adapter decision: `unisim-core` `docs/en/adr-genesis-dr-lifecycle.md`. + +## Related Documents + +- {doc}`ADR Index ` +- {doc}`ADR-0012 Sole Tensor Manager And Scoped Backends ` +- {doc}`Genesis Backend ` diff --git a/docs/sphinx/source/changelog.md b/docs/sphinx/source/changelog.md index 9e3c2bdfe..d94d9efa9 100644 --- a/docs/sphinx/source/changelog.md +++ b/docs/sphinx/source/changelog.md @@ -13,6 +13,72 @@ UniLab 遵循[语义化版本](https://semver.org/)。本共享页面以中英 PyPI 版本变更与未发布变更；发布日期采用 PyPI 上传日期。完整提交历史请参阅 [UniLab 仓库](https://github.com/unilabsim/UniLab)。 +## Unreleased / 未发布 + +### Added / 新增 + +- The Manager tensor ray observation now consumes UniSim's full-scene public + `RayCollisionDescription`, so mesh/hfield materialization no longer requires + a backend-specific `uni_ray` builder, and entity-local packed pose reads are + aggregated onto the public full-scene body axis. `RayPatternCfg` supports + ring, parallel-grid, and pinhole-camera ray patterns; body, yaw-only, and + world alignment; body or public-site attachment with a local sensor offset; + `min_distance`; public geom/tag/body/ancestor filters; optional normal and + hit-point outputs in addition to distance; and reset-boundary per-row geometry + synchronization. The optional `ray` extra now resolves the published + `uniray[warp]==0.1.0` distribution in both CUDA and ROCm lockfiles + ([UniSim #472](https://github.com/unilabsim/unisim/issues/472)). + Manager tensor ray observation 现在消费 UniSim 公开 full-scene + `RayCollisionDescription`，mesh/hfield materialization 不再要求 backend-specific + `uni_ray` builder，并将 entity-local packed pose read 聚合到公开 full-scene body + 轴。`RayPatternCfg` 支持 ring、平行 grid 与 pinhole-camera ray pattern， + body/yaw-only/world alignment、body/public-site attachment 与本地 sensor + offset、`min_distance`、公开 geom/tag/body/ancestor filter、可选 normal 与 + hit-point 输出，以及 reset 边界的逐行 geometry 同步。可选 `ray` extra + 在 CUDA 与 ROCm lockfile 中均解析已发布的 `uniray[warp]==0.1.0` 发行版 + （[UniSim #472](https://github.com/unilabsim/unisim/issues/472)）。 + The first facade is a minimal grouped ray sensor, not a camera: intrinsics, + image projection, noise, history, and scheduling remain application logic. + Full-scene pose aggregation covers entity-owned bodies plus the identity + world root; uncovered non-root bodies fail closed. MuJoCo and MJWarp + fixed-variant collision descriptors now publish per-row geometry through the + capability added by [UniSim #479](https://github.com/unilabsim/unisim/issues/479); + Motrix source mesh/hfield scenes follow the UniSim fail-closed boundary. + The real end-to-end fixture covers + primitive geometry, static filtering, minimum distance, attachment + transformation, site attachment, and normals; focused fake-caster tests cover + grid and pinhole ray geometry, yaw/world alignment, hit-point packing, and + geometry synchronization; mesh/hfield + numerical parity is owned by the UniSim and `uni_ray` contract suites. + 首个 facade 是最小 grouped ray sensor，不是 camera：intrinsics、图像投影、 + noise、history 与 scheduling 仍属于应用逻辑。full-scene pose 聚合覆盖 + entity-owned body 与 identity world root；未覆盖的非 root body fail closed。 + MuJoCo 与 MJWarp fixed-variant collision scene 以及 Motrix source + mesh/hfield scene 遵循 UniSim fail-closed 边界；fixed-variant collision + descriptor 通过 [UniSim #479](https://github.com/unilabsim/unisim/issues/479) + 新增的能力发布逐行 geometry。真实端到端 fixture 覆盖 + primitive geometry、static filtering、minimum distance、attachment + transformation、site attachment 与 normal；grid/pinhole ray 几何、yaw/world alignment、hit-point packing 与 geometry synchronization 由 focused fake-caster test 覆盖；mesh/hfield 数值 parity 由 UniSim 与 `uni_ray` contract suite 负责。 + +### Changed / 变更 + +- Genesis runtime domain randomization now follows the unified device-resident + `unisim-core` lifecycle. Manager interval negotiation selects the + backend-bound device plane, wrench payloads and body IDs enter the interval + plan on that device, and the packed reset device is declared before cold + event-term binds. The Genesis G1 walk SAC owner re-enables `pd_gains` with + the PPO owner because Genesis declares device-plane packed reset `kp`/`kd`; + host-plane runtime DR, push, velocity-delta, and other undeclared terms + still fail closed + ([#2520](https://github.com/Motphys/UniLab/issues/2520)). + Genesis 运行时域随机化现在遵循 `unisim-core` 统一的 device-resident 生命 + 周期。Manager interval 协商选择 backend 绑定的 device plane；wrench payload + 与 body ID 在该设备上进入 interval plan；packed reset 设备会在冷路径 + event-term bind 之前声明。Genesis G1 walk SAC owner 与 PPO owner 一致重新 + 启用 `pd_gains`，因为 Genesis 声明了 device-plane packed reset 的 + `kp`/`kd`；host-plane 运行时 DR、push、velocity-delta 以及其他未声明 term + 仍会 fail closed（[#2520](https://github.com/Motphys/UniLab/issues/2520)）。 + ## 1.5.1 (2026-10-23) ### Added / 新增 diff --git a/docs/sphinx/source/en/1-getting_started/2-installation.md b/docs/sphinx/source/en/1-getting_started/2-installation.md index e94f8fcc8..86fb7356e 100644 --- a/docs/sphinx/source/en/1-getting_started/2-installation.md +++ b/docs/sphinx/source/en/1-getting_started/2-installation.md @@ -250,8 +250,9 @@ ROCm notes: explicit `training.collector_tensor_device=cuda` forces the same routing for other learner devices. Direct environment construction remains CPU. Host-bridge physics stays on CPU; packed transfers connect it to the GPU - observations, actions, rewards, and resets. CUDA-only physics backends remain - unsupported on ROCm. See + observations, actions, rewards, and resets. Genesis additionally supports its + existing device-resident direct lifecycle with HIP/ROCm Torch and exact + `gs.amdgpu`; other CUDA-only physics backends remain unsupported on ROCm. See {doc}`/adr/ADR-0012-sole-tensor-manager-and-scoped-backends`. - When installing from PyPI instead of a source checkout, `make sync-rocm` does not apply. Install the torch build validated by the repository from the diff --git a/docs/sphinx/source/en/2-user_guide/1-training/4-tensor_runtime.md b/docs/sphinx/source/en/2-user_guide/1-training/4-tensor_runtime.md index 5328788aa..f9e555efe 100644 --- a/docs/sphinx/source/en/2-user_guide/1-training/4-tensor_runtime.md +++ b/docs/sphinx/source/en/2-user_guide/1-training/4-tensor_runtime.md @@ -58,9 +58,10 @@ actions use the CPU inference ring; in the default-legal pairings the learner itself is CPU or Apple MPS — see the pairing contract below for CUDA learners), and `cuda` derives the CUDA ring. `cuda` follows Torch device-type semantics and covers both NVIDIA CUDA and AMD ROCm. It does not relax separate -platform gates: CUDA-only tensor backends still fail closed on ROCm/HIP, and -`training.cuda_process_sharing` remains available only to the NVIDIA-specific -CUDA MPS daemon and native H2D implementation. +platform gates: CUDA-only tensor backends still fail closed on ROCm/HIP, while +Genesis separately admits only its reviewed HIP/ROCm + exact `gs.amdgpu` +profile. `training.cuda_process_sharing` remains available only to the +NVIDIA-specific CUDA MPS daemon and native H2D implementation. Removed keys fail closed with a migration error: use `training.collector_tensor_device` instead of `training.sim_carrier_device` diff --git a/docs/sphinx/source/en/2-user_guide/3-backends/5-genesis.md b/docs/sphinx/source/en/2-user_guide/3-backends/5-genesis.md index 84ecba2ab..7d9ce4eb2 100644 --- a/docs/sphinx/source/en/2-user_guide/3-backends/5-genesis.md +++ b/docs/sphinx/source/en/2-user_guide/3-backends/5-genesis.md @@ -52,14 +52,24 @@ compensates on the materialize cold path — hot paths never parse XML: reproduction of it. - **Actuators**: `` actuators import losslessly and are driven with `control_dofs_position`; kp/kd reset randomization is supported - (the DR capability set declares only the per-env round-trip-measured terms: - body mass, base mass, kp, kd, plus interval body force). + on the device plane (the DR capability sets declare only the measured device + reset terms: body mass, base-mass delta, body CoM / base-CoM delta, kp, + kd; the device interval set additionally covers body force, body torque, + and world-frame root linear plus root-local angular velocity deltas). ## Prerequisites -- Linux x86_64 with an NVIDIA GPU and driver. Only the `gs.gpu` lane is - validated by the feasibility probe; the CPU backend is not a validated - support lane. +- Linux x86_64 with a supported GPU runtime for the Manager tensor runtime. + UniSim's reviewed Genesis tensor profile requires Genesis zero-copy and one + exact runtime pairing: non-HIP Torch CUDA with `gs.cuda`, or HIP/ROCm Torch + with `gs.amdgpu`. Both negotiate `DEVICE_RESIDENT`, in-process, `DIRECT` and + retain the public `cuda` Torch device namespace. +- Genesis 1.3.3 itself natively supports `gs.metal` on Apple Silicon and + `gs.amdgpu` on Linux HIP, as shown by the direct-engine G1 throughput work + (UniLab PR #2518 / discussion #2517). Metal remains outside the UniSim + Genesis adapter tensor lifecycle. On ROCm, `gs.cuda` is rejected as a runtime + mismatch; on non-HIP CUDA, `gs.amdgpu` is likewise rejected. CPU, portable + Genesis, Metal, MPS, and substitute backends are never silent fallbacks. - The repository's pinned torch window (torch 2.8 on x86_64): the IMUSensor + contact-force sensor combination crashed under torch 2.7 in the probe (REPORT §3.4/§8). @@ -145,8 +155,13 @@ degrading: - **Absolute geom-friction DR**: upstream Genesis only offers a per-env friction *ratio* API, so `geom_friction` reset randomization and `get_geom_friction` fail closed. -- **Interval push / body-velocity-delta DR**: rejected at construction/plan - time (`push_body_name`, push perturbation). +- **Legacy interval push term**: the opaque `push` descriptor is rejected at + construction/plan time. Use the six-axis velocity-delta event, which commits + world-frame linear and root-local angular deltas through Genesis public + state without touching qpos, controls, solver state, or unrelated rows. +- **Host-plane runtime DR**: Genesis accepts reset and interval payloads only + on its backend-bound CUDA/AMDGPU device. Host payloads fail closed with the + term and plane named; there is no NumPy or host-bridge fallback. - **Repeated `gs.init`** in one process (see above). ## Cross-Backend Migration (sim2sim) 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 d96c46718..06f64d0a4 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 @@ -40,7 +40,7 @@ These three paths correspond to three lifecycle classes: | Task | Currently implemented reset domain randomization | Currently implemented interval domain randomization | Default state | | --- | --- | --- | --- | | `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 `src/unilab/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 | +| `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 the mujoco and genesis owners by default; disabled on motrix/mjwarp owners | | `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 | ## Current DR Capabilities and Boundaries 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 e0fcfef3a..96f21f0dd 100644 --- a/docs/sphinx/source/en/5-reference/5-support_matrix.md +++ b/docs/sphinx/source/en/5-reference/5-support_matrix.md @@ -81,7 +81,7 @@ This table is derived from UniSim's SDK-free public static inventory. It describ | `mjwarp` | Device-resident / in-process / direct | CUDA | Required for the entire tensor lifecycle | Supported: Linux CUDA only | Unsupported; no CPU, MPS, or ROCm fallback | Unsupported; no CPU, MPS, or ROCm fallback | In-process; no external Python worker | Unsupported | Unsupported | Unsupported | Unsupported | | `newton` | Device-resident / in-process / direct | CUDA | Required for the entire tensor lifecycle | Supported: Linux CUDA only | Unsupported; no CPU, MPS, or ROCm fallback | Unsupported; no CPU, MPS, or ROCm fallback | In-process; no external Python worker | Exact | Unsupported | Unsupported | Unsupported | | `superdex` | Host bridge / in-process / host bridge | CPU / CUDA | Required only when the learner requests CUDA state/control buffers | Supported: CPU-authoritative physics with optional CUDA Torch buffers | CPU-authoritative host bridge only; no CUDA physics claim | CPU-authoritative host bridge with explicit ROCm Torch buffers (`cuda`) | In-process; no external Python worker | Unsupported | Unsupported | Unsupported | Exact | -| `genesis` | Device-resident / in-process / direct | CUDA | Required for the entire tensor lifecycle | Supported: Linux CUDA only | Unsupported; no CPU, MPS, or ROCm fallback | Unsupported; no CPU, MPS, or ROCm fallback | In-process; no external Python worker | Exact | Unsupported | Unsupported | Unsupported | +| `genesis` | Device-resident / in-process / direct | CUDA | Required for the entire tensor lifecycle | Supported: Linux CUDA only | Unsupported UniSim tensor profile; Genesis engine natively supports Metal, with no CPU or substitute-backend fallback | Supported: Linux ROCm with HIP Torch and exact gs.amdgpu (cuda); no CPU or substitute-backend fallback | In-process; no external Python worker | Exact | Unsupported | Unsupported | Unsupported | | `isaacgym` | Device-resident / external worker / cuda ipc | CUDA | Required for the entire tensor lifecycle | Supported: Linux CUDA only | Unsupported; no CPU, MPS, or ROCm fallback | Unsupported; no CPU, MPS, or ROCm fallback | Dedicated external Python 3.8 worker; host Python paths are not inherited | Exact | Unsupported | Unsupported | Unsupported | | `isaacsim` | Device-resident / external worker / cuda ipc | CUDA | Required for the entire tensor lifecycle | Supported: Linux CUDA only | Unsupported; no CPU, MPS, or ROCm fallback | Unsupported; no CPU, MPS, or ROCm fallback | Dedicated external Python 3.11 worker; host Python paths are not inherited | Exact | Unsupported | Unsupported | Unsupported | @@ -116,7 +116,8 @@ Pairing contract (issue #2483, ADR-0019): with a CUDA learner on a non-ROCm Torc CUDA-only tensor backends are rejected before backend construction on CPU, MPS, macOS, ROCm/HIP, unavailable CUDA, malformed CUDA ordinals, and out-of-range -device requests. They do not fall back to NumPy or a host bridge. +device requests. Genesis instead accepts its declared HIP/ROCm + `gs.amdgpu` +profile. Unsupported requests do not fall back to NumPy or a host bridge. Check the Torch runtime and visible ordinals in the same process environment: @@ -124,17 +125,19 @@ Check the Torch runtime and visible ordinals in the same process environment: uv run python -c "import torch; print(torch.__version__, torch.version.cuda, torch.version.hip, torch.cuda.is_available(), torch.cuda.device_count(), torch.cuda.current_device())" ``` -`torch.version.hip` must be `None` for the CUDA-only matrix. If CUDA is -unavailable, check the NVIDIA driver/container mismatch before changing task -configuration. If `CUDA_VISIBLE_DEVICES` is set, backend ordinals address that -remapped namespace, not host-global physical indices. - -On macOS and ROCm, use a CPU-authoritative host-bridge backend. On ROCm, -Manager/TorchEnv uses the current GPU whenever the learner device is CUDA — -the default `training.collector_tensor_device=auto` follows the learner device — -or when training explicitly sets `training.collector_tensor_device=cuda`, provided -the selected HOST_BRIDGE backend declares and accepts the routed `cuda` -`env.collector_tensor_device`; CPU learners keep CPU, torch.mps learners retain the -owner/default CPU placement, and direct environment construction remains CPU. +`torch.version.hip` must be `None` for CUDA-only backends and the Genesis +`gs.cuda` pairing; the Genesis `gs.amdgpu` pairing requires a HIP build. If the +GPU runtime is unavailable, check the driver/container mismatch before changing +task configuration. If `CUDA_VISIBLE_DEVICES` is set, backend ordinals address +that remapped namespace, not host-global physical indices. + +On macOS, use a CPU-authoritative host-bridge backend. ROCm supports both host +bridges and the Genesis exact-`gs.amdgpu` device-resident profile. For a host +bridge, Manager/TorchEnv uses the current GPU whenever the learner device is +CUDA — the default `training.collector_tensor_device=auto` follows the learner +device — or when training explicitly sets `training.collector_tensor_device=cuda`, +provided the selected backend declares and accepts the routed `cuda` +`env.collector_tensor_device`; CPU learners keep CPU, torch.mps learners retain +the owner/default CPU placement, and direct environment construction remains CPU. GPU Torch buffers on a host bridge do not imply GPU physics or a device-resident backend lifecycle. diff --git a/docs/sphinx/source/zh_CN/1-getting_started/2-installation.md b/docs/sphinx/source/zh_CN/1-getting_started/2-installation.md index 83a833826..b49262085 100644 --- a/docs/sphinx/source/zh_CN/1-getting_started/2-installation.md +++ b/docs/sphinx/source/zh_CN/1-getting_started/2-installation.md @@ -225,7 +225,9 @@ ROCm 说明： 设备，显式 `training.collector_tensor_device=cuda` 也可对其他 learner 设备强制同样的路由。 直接构造环境仍为 CPU。host-bridge 物理仿真仍在 CPU 上，通过 packed 传输连接 - GPU observation、action、reward 与 reset。CUDA-only 物理后端仍不支持 ROCm。见 + GPU observation、action、reward 与 reset。Genesis 还支持 HIP/ROCm Torch + + exact `gs.amdgpu` 的既有 device-resident direct 生命周期；其他 CUDA-only + 物理后端仍不支持 ROCm。见 {doc}`/adr/ADR-0012-sole-tensor-manager-and-scoped-backends`。 - 从 PyPI 安装（不克隆仓库）时，`make sync-rocm` 不适用；先从 PyTorch ROCm 索引 安装仓库验证过的 torch build，再安装 `unilab`。发布的依赖范围是 diff --git a/docs/sphinx/source/zh_CN/2-user_guide/1-training/4-tensor_runtime.md b/docs/sphinx/source/zh_CN/2-user_guide/1-training/4-tensor_runtime.md index 5b976d0f3..60194ab17 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/1-training/4-tensor_runtime.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/1-training/4-tensor_runtime.md @@ -54,7 +54,8 @@ off-policy inference transport 是纯派生事实，不再是配置键： ring；默认合法配对中 learner 是 CPU 或 Apple MPS，CUDA learner 的配对约束 见下文的配对契约），`cuda` 推导 CUDA ring。`cuda` 遵循 Torch device type 语义，覆盖 NVIDIA CUDA 与 AMD ROCm。它不会放宽独立 -平台 gate：CUDA-only tensor backend 在 ROCm/HIP 上仍 fail closed， +平台 gate：CUDA-only tensor backend 在 ROCm/HIP 上仍 fail closed；Genesis +仅通过独立审查的 HIP/ROCm + exact `gs.amdgpu` profile 支持 ROCm。 `training.cuda_process_sharing` 仍只开放给 NVIDIA 专属的 CUDA MPS daemon 与 native H2D 实现。 diff --git a/docs/sphinx/source/zh_CN/2-user_guide/3-backends/5-genesis.md b/docs/sphinx/source/zh_CN/2-user_guide/3-backends/5-genesis.md index 0934207d5..1a221e59d 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/3-backends/5-genesis.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/3-backends/5-genesis.md @@ -47,14 +47,25 @@ materialize 冷路径补偿——热路径从不解析 XML： contact sensor 变为 per-link net contact force 阈值（1 N）——这是对 geom 对 `found` 语义的近似，不是复现。 - **执行器**：`` 执行器无损导入，用 - `control_dofs_position` 驱动；kp/kd reset 随机化受支持（DR 能力集只 - 声明 per-env 往返实测的项：body mass、base mass、kp、kd，以及 - interval body force）。 + `control_dofs_position` 驱动；kp/kd reset 随机化在 device plane 受支持 + （DR 能力集只声明实测的 device reset 项：body mass、kp、kd，以及 + base-mass delta、body CoM / base-CoM delta、kp、kd；device interval 集 + 另覆盖 body force、body torque，以及世界系 root 线速度与 root 本体系 + 角速度 delta）。 ## 前置条件 -- Linux x86_64，装有 NVIDIA GPU 与驱动。可行性探针只验证了 `gs.gpu` - 通道；CPU backend 不是已验证的支持通道。 +- Manager tensor 运行时要求 Linux x86_64 与受支持的 GPU runtime。UniSim 已 + 审查的 Genesis tensor profile 要求 Genesis zero-copy 及以下一种 exact + runtime 配对：非 HIP Torch CUDA + `gs.cuda`，或 HIP/ROCm Torch + + `gs.amdgpu`。两者均协商 `DEVICE_RESIDENT`、in-process、`DIRECT`，并保留 + 公开 `cuda` Torch device namespace。 +- Genesis 1.3.3 引擎本身在 Apple Silicon 上原生支持 `gs.metal`，在 + Linux HIP 上原生支持 `gs.amdgpu`；direct-engine G1 throughput 工作 + （UniLab PR #2518 / discussion #2517）已证明这两条引擎 lane。Metal 仍在 + UniSim Genesis adapter tensor 生命周期之外。ROCm 上的 `gs.cuda` 与非 HIP + CUDA 上的 `gs.amdgpu` 都因 runtime 不匹配而被拒绝；CPU、portable Genesis、 + Metal、MPS 或替代后端不会成为静默 fallback。 - 仓库钉定的 torch 窗口（x86_64 上为 torch 2.8）：探针中 IMUSensor + contact-force 传感器组合在 torch 2.7 下崩溃（REPORT §3.4/§8）。 @@ -130,8 +141,13 @@ Genesis 的 init/destroy 循环每轮泄漏 200–450 MB host RSS（REPORT flat owner YAML。 - **绝对值 geom 摩擦 DR**：上游 Genesis 只有 per-env 摩擦*比例* API， 因此 `geom_friction` reset 随机化与 `get_geom_friction` fail-closed。 -- **interval push / body-velocity-delta DR**：在构造/plan 时拒绝 - （`push_body_name`、push perturbation）。 +- **legacy interval push term**：不透明 `push` descriptor 在构造/plan 时 + 拒绝。请使用六轴 velocity-delta event；它通过 Genesis public state 提交 + 世界系线速度与 root 本体系角速度 delta，且不修改 qpos、control、solver + state 或未选中的 rows。 +- **host-plane 运行时 DR**：Genesis 只接受位于其 backend 绑定 + CUDA/AMDGPU device 上的 reset 与 interval payload。host payload 会带上 + term 与 plane fail-closed；没有 NumPy 或 host-bridge 回退。 - **同进程重复 `gs.init`**（见上文）。 ## 跨后端迁移（sim2sim） 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 272f03158..79c6f15db 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 @@ -36,7 +36,7 @@ Manager-Based event term 是唯一 DR 声明路径： | Task | 当前已实现的 reset 域随机化 | 当前已实现的 interval 域随机化 | 默认状态 | | --- | --- | --- | --- | | `Go2JoystickFlat` | 经 `reset_root_state_uniform` 的 base xy/yaw 与 base qvel；command 采样；经 `pd_gains` 的 kp/kd | 无 | event term 在 `src/unilab/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 默认禁用 | +| `G1WalkFlat` | 经 `reset_root_state_uniform` 的 base xy/yaw 与 base qvel；带平面死区的 command 采样；`gait_phase` 采样；经 `pd_gains` 的 kp/kd 随机化 | 无 | mujoco 与 genesis owner 默认启用 kp/kd；motrix/mjwarp owner 默认禁用 | | `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` 默认有非零扰动 | ## 当前 DR 的能力与边界 diff --git a/docs/sphinx/source/zh_CN/5-reference/5-support_matrix.md b/docs/sphinx/source/zh_CN/5-reference/5-support_matrix.md index fdb0f83ac..c39629725 100644 --- a/docs/sphinx/source/zh_CN/5-reference/5-support_matrix.md +++ b/docs/sphinx/source/zh_CN/5-reference/5-support_matrix.md @@ -73,7 +73,7 @@ uv run scripts/generate_support_matrix.py --write | `mjwarp` | Device-resident / in-process / direct | CUDA | Required for the entire tensor lifecycle | Supported: Linux CUDA only | Unsupported; no CPU, MPS, or ROCm fallback | Unsupported; no CPU, MPS, or ROCm fallback | In-process; no external Python worker | 不支持 | 不支持 | 不支持 | 不支持 | | `newton` | Device-resident / in-process / direct | CUDA | Required for the entire tensor lifecycle | Supported: Linux CUDA only | Unsupported; no CPU, MPS, or ROCm fallback | Unsupported; no CPU, MPS, or ROCm fallback | In-process; no external Python worker | 支持 | 不支持 | 不支持 | 不支持 | | `superdex` | Host bridge / in-process / host bridge | CPU / CUDA | Required only when the learner requests CUDA state/control buffers | Supported: CPU-authoritative physics with optional CUDA Torch buffers | CPU-authoritative host bridge only; no CUDA physics claim | CPU-authoritative host bridge with explicit ROCm Torch buffers (`cuda`) | In-process; no external Python worker | 不支持 | 不支持 | 不支持 | 支持 | -| `genesis` | Device-resident / in-process / direct | CUDA | Required for the entire tensor lifecycle | Supported: Linux CUDA only | Unsupported; no CPU, MPS, or ROCm fallback | Unsupported; no CPU, MPS, or ROCm fallback | In-process; no external Python worker | 支持 | 不支持 | 不支持 | 不支持 | +| `genesis` | Device-resident / in-process / direct | CUDA | Required for the entire tensor lifecycle | Supported: Linux CUDA only | Unsupported UniSim tensor profile; Genesis engine natively supports Metal, with no CPU or substitute-backend fallback | Supported: Linux ROCm with HIP Torch and exact gs.amdgpu (cuda); no CPU or substitute-backend fallback | In-process; no external Python worker | 支持 | 不支持 | 不支持 | 不支持 | | `isaacgym` | Device-resident / external worker / cuda ipc | CUDA | Required for the entire tensor lifecycle | Supported: Linux CUDA only | Unsupported; no CPU, MPS, or ROCm fallback | Unsupported; no CPU, MPS, or ROCm fallback | Dedicated external Python 3.8 worker; host Python paths are not inherited | 支持 | 不支持 | 不支持 | 不支持 | | `isaacsim` | Device-resident / external worker / cuda ipc | CUDA | Required for the entire tensor lifecycle | Supported: Linux CUDA only | Unsupported; no CPU, MPS, or ROCm fallback | Unsupported; no CPU, MPS, or ROCm fallback | Dedicated external Python 3.11 worker; host Python paths are not inherited | 支持 | 不支持 | 不支持 | 不支持 | @@ -107,7 +107,8 @@ uv run scripts/generate_support_matrix.py --write ## 平台故障排查 CUDA-only tensor 后端在 CPU、MPS、macOS、ROCm/HIP、CUDA 不可用、CUDA 序号 -非法以及设备越界时，都会在后端构造前失败；不会回退到 NumPy 或 host bridge。 +非法以及设备越界时，都会在后端构造前失败；Genesis 则接受其声明的 +HIP/ROCm + `gs.amdgpu` profile。不支持的请求不会回退到 NumPy 或 host bridge。 请在同一个进程环境中检查 Torch runtime 与可见序号： @@ -115,15 +116,16 @@ CUDA-only tensor 后端在 CPU、MPS、macOS、ROCm/HIP、CUDA 不可用、CUDA uv run python -c "import torch; print(torch.__version__, torch.version.cuda, torch.version.hip, torch.cuda.is_available(), torch.cuda.device_count(), torch.cuda.current_device())" ``` -CUDA-only 矩阵要求 `torch.version.hip` 为 `None`。若 CUDA 不可用，先检查 -NVIDIA 驱动与容器运行时不匹配，再调整任务配置。设置了 -`CUDA_VISIBLE_DEVICES` 时，后端序号指向重映射后的命名空间，而不是宿主机 -全局物理索引。 +CUDA-only 后端与 Genesis `gs.cuda` 配对要求 `torch.version.hip` 为 `None`； +Genesis `gs.amdgpu` 配对则要求 HIP build。若 GPU runtime 不可用，先检查 +驱动与容器运行时不匹配，再调整任务配置。设置了 `CUDA_VISIBLE_DEVICES` 时， +后端序号指向重映射后的命名空间，而不是宿主机全局物理索引。 -在 macOS 与 ROCm 上请使用 CPU-authoritative host-bridge 后端。在 ROCm 上， -只要 learner 设备为 CUDA——默认的 `training.collector_tensor_device=auto` 会跟随 +在 macOS 上请使用 CPU-authoritative host-bridge 后端。ROCm 支持 host bridge +与 Genesis exact-`gs.amdgpu` device-resident profile。对于 host bridge，只要 +learner 设备为 CUDA——默认的 `training.collector_tensor_device=auto` 会跟随 learner 设备——或训练显式设置 `training.collector_tensor_device=cuda`，且所选 -HOST_BRIDGE 后端声明并接受被路由的 `cuda` `env.collector_tensor_device`， -Manager/TorchEnv 就使用当前 GPU；CPU learner 保持 CPU，torch.mps learner -保留 owner/default CPU 放置，直接构造环境仍为 CPU。host bridge 使用 GPU -Torch buffer 不代表 GPU physics，也不代表 device-resident 后端生命周期。 +后端声明并接受被路由的 `cuda` `env.collector_tensor_device`，Manager/TorchEnv +就使用当前 GPU；CPU learner 保持 CPU，torch.mps learner 保留 owner/default +CPU 放置，直接构造环境仍为 CPU。host bridge 使用 GPU Torch buffer 不代表 +GPU physics，也不代表 device-resident 后端生命周期。 diff --git a/pyproject.rocm.toml b/pyproject.rocm.toml index 6ff905e47..6b91cd1e6 100644 --- a/pyproject.rocm.toml +++ b/pyproject.rocm.toml @@ -23,7 +23,7 @@ classifiers = [ requires-python = ">=3.10,<3.14" dependencies = [ "numpy", - "unisim-core==2.1.1", + "unisim-core==2.2.0", "prettytable>=3.10", "torch==2.14.0", "triton-rocm==3.8.0 ; sys_platform == 'linux' and platform_machine == 'x86_64'", @@ -72,7 +72,9 @@ mujoco = [ "mujoco~=3.11.0", "mjbatch-uni~=0.2.5", ] -motrix = ["unisim-core[motrix]==2.1.1"] +motrix = ["unisim-core[motrix]==2.2.0"] +# Keep the optional tensor-ray extra aligned with pyproject.toml. +ray = ["uniray[warp]==0.1.0"] [dependency-groups] # The RL stack (APPO/off-policy/multi-GPU) is merged into unilab itself diff --git a/pyproject.toml b/pyproject.toml index c2e991376..75901cbca 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,7 +33,7 @@ classifiers = [ requires-python = ">=3.10,<3.14" dependencies = [ "numpy", - "unisim-core==2.1.1", + "unisim-core==2.2.0", "prettytable>=3.10", # A range (not an exact pin) lets ROCm users substitute a ROCm torch # build; uv.lock pins the tested CUDA builds via tool.uv.sources. @@ -106,13 +106,17 @@ newton = [ "pyglet>=2.1.6,<3", "imgui-bundle>=1.92.0", ] -motrix = ["unisim-core[motrix]==2.1.1"] +motrix = ["unisim-core[motrix]==2.2.0"] genesis = ["genesis-world==1.3.3"] # Wheels are CPython 3.12/3.13 Linux x86_64 only; the marker keeps uv lock # resolvable for the other required-environments. superdex = [ - "unisim-core[superdex]==2.1.1 ; python_version >= '3.12' and sys_platform == 'linux' and platform_machine == 'x86_64'", + "unisim-core[superdex]==2.2.0 ; python_version >= '3.12' and sys_platform == 'linux' and platform_machine == 'x86_64'", ] +# Tensor ray perception is optional and is discovered through UniSim's public +# caster factory. The PyPI distribution is `uniray`; the import namespace is +# `uni_ray`. +ray = ["uniray[warp]==0.1.0"] [dependency-groups] # The RL stack (APPO/off-policy/multi-GPU) is merged into unilab itself diff --git a/src/unilab/base/backend_factory.py b/src/unilab/base/backend_factory.py index f7655c367..892e82e51 100644 --- a/src/unilab/base/backend_factory.py +++ b/src/unilab/base/backend_factory.py @@ -157,7 +157,9 @@ def _validate_cuda_only_backend_platform(backend_type: str, kwargs: dict[str, An "macOS has no supported CUDA runtime", "Use a CPU-authoritative host-bridge backend or a Linux CUDA runtime.", ) - if runtime_state["torch_hip"] not in (None, "None", ""): + if runtime_state["torch_hip"] not in (None, "None", "") and not ( + profile.rocm_tensor_profile.startswith("Supported:") + ): raise _fail_closed_cuda_only( backend_type, requested_device, diff --git a/src/unilab/base/entity.py b/src/unilab/base/entity.py index 854ca79f9..407cce505 100644 --- a/src/unilab/base/entity.py +++ b/src/unilab/base/entity.py @@ -7,6 +7,7 @@ from __future__ import annotations +import inspect import re from collections.abc import Callable, Iterator, Mapping, Sequence from contextlib import contextmanager @@ -41,6 +42,7 @@ RESET_TERM_GEOM_FRICTION, RESET_TERM_KD, RESET_TERM_KP, + DataPlaneName, IntervalRandomizationPlan, IntervalTermOp, ) @@ -205,11 +207,26 @@ def __post_init__(self) -> None: class EntityTensorBodyStateView: """Entity-ordered world-frame Torch body state for one read phase.""" + body_names: tuple[str, ...] + pos_w: torch.Tensor | None + quat_w: torch.Tensor | None + lin_vel_w: torch.Tensor | None + ang_vel_w: torch.Tensor | None + + +@dataclass(frozen=True) +class SceneTensorPoseView: + """Ray-shaped full-batch Torch pose backed by packed scene tensor views. + + This is deliberately ``RayPoseTensorView``-shaped without importing the + caster package. The packet-backed tensors may be row-strided column views; + the public caster accepts such packed-strided Torch views without requiring + global contiguity. + """ + body_names: tuple[str, ...] pos_w: torch.Tensor quat_w: torch.Tensor - lin_vel_w: torch.Tensor - ang_vel_w: torch.Tensor _TENSOR_BODY_SENSOR_FIELDS: tuple[tuple[str, str], ...] = ( @@ -218,23 +235,51 @@ class EntityTensorBodyStateView: ("lin_vel_w", "track_linvel_w_"), ("ang_vel_w", "track_angvel_w_"), ) +_TENSOR_POSE_BODY_FIELDS = ("pos_w", "quat_w") +_TENSOR_BODY_FIELDS = frozenset((*_TENSOR_POSE_BODY_FIELDS, "lin_vel_w", "ang_vel_w")) +_TENSOR_BODY_FIELD_ORDER = ("pos_w", "quat_w", "lin_vel_w", "ang_vel_w") +_TENSOR_BODY_PACKET_FIELDS = { + "pos_w": "body_pos_w", + "quat_w": "body_quat_w", + "lin_vel_w": "body_lin_vel_w", + "ang_vel_w": "body_ang_vel_w", +} +_TENSOR_PACKET_FIELD_WIDTHS = { + "body_pos_w": 3, + "body_quat_w": 4, + "body_lin_vel_w": 3, + "body_ang_vel_w": 3, +} +_PACKET_FIELD_BODY_FIELDS = { + packet_name: field for field, packet_name in _TENSOR_BODY_PACKET_FIELDS.items() +} _TRACKED_BODY_SENSOR_PREFIXES = {field: prefix for field, prefix in _TENSOR_BODY_SENSOR_FIELDS} _TENSOR_SENSOR_WIDTHS = {"pos_w": 3, "quat_w": 4, "lin_vel_w": 3, "ang_vel_w": 3} +def _packet_body_name(entity_name: str, body_name: str) -> str: + return f"{entity_name}/{body_name}" + + @dataclass(frozen=True) class SceneTensorReadSpec: """Declare one entity's named state reads for the scene tensor phase. ``sensor_names`` are explicit backend names. ``body_names`` are entity-local - names; the scene owner expands them into the canonical ``track_*`` sensor - contract. Requests are explicit so the packed packet never grows from an - implicit EntityCfg namespace that is only needed by a cold-path consumer. + names; ``body_fields`` selects their public semantic blocks. The legacy + default is the exact four-field body-state set; pose-only ray consumers + select ``(pos_w, quat_w)`` so velocities never enter the packed transfer. """ entity: str sensor_names: tuple[str, ...] = () body_names: tuple[str, ...] = () + body_fields: tuple[str, ...] = ( + "pos_w", + "quat_w", + "lin_vel_w", + "ang_vel_w", + ) def __post_init__(self) -> None: if not isinstance(self.entity, str) or not self.entity: @@ -243,8 +288,11 @@ def __post_init__(self) -> None: raise TypeError("Scene tensor read sensor_names must be a sequence of strings") if isinstance(self.body_names, (str, bytes)): raise TypeError("Scene tensor read body_names must be a sequence of strings") + if isinstance(self.body_fields, (str, bytes)): + raise TypeError("Scene tensor read body_fields must be a sequence of strings") sensors = tuple(self.sensor_names) bodies = tuple(self.body_names) + fields = tuple(self.body_fields) if not sensors and not bodies: # Generic joint observations request the entity's packed qpos/qvel # packet without named sensors or body-projected state. This is a @@ -255,8 +303,19 @@ def __post_init__(self) -> None: raise TypeError(f"Scene tensor read {label} names must be non-empty strings") if len(set(names)) != len(names): raise ValueError(f"Scene tensor read {label} names must be unique: {names}") + if bodies and not fields: + raise ValueError("Scene tensor read body_fields cannot be empty when bodies are set") + unknown = [field for field in fields if field not in _TENSOR_BODY_FIELDS] + if unknown: + raise ValueError( + f"Scene tensor read body_fields are unknown: {unknown}; " + f"expected={sorted(_TENSOR_BODY_FIELDS)}" + ) + if len(set(fields)) != len(fields): + raise ValueError(f"Scene tensor read body_fields must be unique: {fields}") object.__setattr__(self, "sensor_names", sensors) object.__setattr__(self, "body_names", bodies) + object.__setattr__(self, "body_fields", fields) class SceneTensorReadPlan: @@ -277,6 +336,11 @@ def __init__( entities: Mapping[str, "Entity"], sensor_names: Mapping[str, tuple[str, ...]], body_names: Mapping[str, tuple[str, ...]], + body_fields: Mapping[str, tuple[str, ...]], + packet_names: tuple[str, ...], + packet_sensor_owners: Mapping[str, str], + packet_body_axes: Mapping[str, tuple[str, ...]], + packet_body_blocks: bool, host_plan: HostBridgeTransferPlan | None, ) -> None: self._scene = scene @@ -284,20 +348,24 @@ def __init__( self._entities = MappingProxyType(dict(entities)) self._sensor_names = MappingProxyType(dict(sensor_names)) self._body_names = MappingProxyType(dict(body_names)) + self._body_fields = MappingProxyType(dict(body_fields)) self._host_plan_value = host_plan - self._packet_names = self._aggregate_packet_names() + self._packet_names = packet_names + owners = dict(packet_sensor_owners) + if not packet_body_blocks: + for entity_name, body_names_value in self._body_names.items(): + for body_name in body_names_value: + for field, prefix in _TENSOR_BODY_SENSOR_FIELDS: + if field not in self._body_fields[entity_name]: + continue + name = prefix + body_name + owners.setdefault(name, entity_name) + self._packet_sensor_owners = MappingProxyType(owners) + self._packet_body_axes = MappingProxyType(dict(packet_body_axes)) + self._packet_body_blocks = packet_body_blocks self._aggregate_body_views: dict[str, EntityTensorBodyStateView] = {} self._joint_column_indices: dict[str, tuple[torch.Tensor, torch.Tensor]] = {} - self._packet_sensor_owners: dict[str, str] = {} - self._packet_sensor_owners.update( - (name, entity_name) - for entity_name, names in self._sensor_names.items() - for name in names - ) - for entity_name, body_names_value in self._body_names.items(): - for body_name in body_names_value: - for _, prefix in _TENSOR_BODY_SENSOR_FIELDS: - self._packet_sensor_owners.setdefault(prefix + body_name, entity_name) + self._packet_body_index_cache: dict[tuple[str, tuple[str, ...]], torch.Tensor] = {} self._packet: Mapping[str, torch.Tensor] = MappingProxyType({}) self._refreshed = False self._closed = False @@ -322,6 +390,16 @@ def sensor_names(self) -> Mapping[str, tuple[str, ...]]: def body_names(self) -> Mapping[str, tuple[str, ...]]: return self._body_names + @property + def body_fields(self) -> Mapping[str, tuple[str, ...]]: + """Return the semantic body blocks compiled for each entity.""" + return self._body_fields + + @property + def packet_body_blocks(self) -> bool: + """Return whether the backend publishes native packed body blocks.""" + return self._packet_body_blocks + @property def host_plan(self) -> HostBridgeTransferPlan | None: return self._host_plan_value @@ -336,11 +414,53 @@ def packet_names(self) -> tuple[str, ...]: """Return the immutable backend sensor names compiled into this plan.""" return self._packet_names + def ray_pose_view( + self, entity: "str | Entity", body_names: Sequence[str] | None = None + ) -> SceneTensorPoseView: + """Return a ray-shaped full-batch pose view backed by this phase packet. + The tensors are persistent packed views in the requested body order. + They are valid until the next phase refresh, invalidation, or close. + """ + self._require_open() + owner = self._require_entity(entity) + requested = self._compiled_body_request(owner, body_names) + fields = self._packet_body_fields( + owner, + requested, + required_fields=_TENSOR_POSE_BODY_FIELDS, + capability="ray pose tensor read", + ) + return SceneTensorPoseView( + body_names=requested, + pos_w=fields["pos_w"], + quat_w=fields["quat_w"], + ) + + def packet_body_axis(self, packet_field: str) -> tuple[str, ...]: + """Return immutable aggregate body identity for one packed body field.""" + try: + return self._packet_body_axes[packet_field] + except KeyError as exc: + raise KeyError( + f"Scene tensor body field '{packet_field}' was not compiled into the " + f"packed read; compiled={list(self._packet_body_axes)}" + ) from exc + def refresh(self) -> None: """Publish one full-batch packet for the current read phase.""" self._require_not_closed() if self._host_plan_value is not None: - packet = self._read_packet(self._host_plan_value.read_state_sensors) + if self._packet_body_blocks: + + def reader() -> dict[str, torch.Tensor]: + return self._read_host_body_blocks(selected=False) + + else: + reader = cast( + Callable[[], dict[str, torch.Tensor]], + self._host_plan_value.read_state_sensors, + ) + packet = self._read_packet(reader) else: packet = self._read_device_resident_packet() self._publish_packet(packet) @@ -354,7 +474,17 @@ def refresh_selected(self) -> None: """ self._require_not_closed() if self._host_plan_value is not None: - packet = self._read_packet(self._host_plan_value.read_selected_state_sensors) + if self._packet_body_blocks: + + def reader() -> dict[str, torch.Tensor]: + return self._read_host_body_blocks(selected=True) + + else: + reader = cast( + Callable[[], dict[str, torch.Tensor]], + self._host_plan_value.read_selected_state_sensors, + ) + packet = self._read_packet(reader) else: packet = self._read_device_resident_packet() self._publish_packet(packet) @@ -370,6 +500,22 @@ def _read_device_resident_packet(self) -> dict[str, torch.Tensor]: """ return self._read_packet(None) + def _read_host_body_blocks(self, *, selected: bool) -> dict[str, torch.Tensor]: + """Expose packed body blocks as views of the single host-bridge packet.""" + host_plan = self._host_plan_value + assert host_plan is not None + source = ( + host_plan.read_selected_state_sensors() if selected else host_plan.read_state_sensors() + ) + packet = dict(source) + poses = host_plan.get_body_pose_views() + for field in _TENSOR_POSE_BODY_FIELDS: + packet_name = _TENSOR_BODY_PACKET_FIELDS[field] + value = getattr(poses, field) + if value is not None: + packet[packet_name] = value.reshape(value.shape[0], -1) + return packet + def invalidate(self) -> None: """Drop phase values after an in-phase simulation mutation.""" self._require_not_closed() @@ -500,14 +646,9 @@ def body_tensor_view( self, entity: "str | Entity", body_names: Sequence[str] | None = None ) -> EntityTensorBodyStateView: """Return entity-local body state from canonical phase sensors.""" + self._require_open() owner = self._require_entity(entity) - if owner._body_names is None: - raise owner._capability_error( - "scene body tensor state", "body_names were not declared in EntityCfg" - ) - requested = self._scene._normalize_tensor_body_request( - owner, tuple(body_names) if body_names is not None else owner._body_names - ) + requested = self._compiled_body_request(owner, body_names) missing = [name for name in requested if name not in set(self._body_names[owner.name])] if missing: raise ValueError( @@ -519,8 +660,31 @@ def body_tensor_view( if aggregate is not None and aggregate.body_names == requested: return aggregate + compiled_fields = self._body_fields[owner.name] + packet_fields_available = all( + _TENSOR_BODY_PACKET_FIELDS[field] in self._packet for field in compiled_fields + ) + if packet_fields_available: + fields = self._packet_body_fields( + owner, + requested, + required_fields=compiled_fields, + capability="aggregate body tensor read", + ) + view = EntityTensorBodyStateView( + body_names=requested, + pos_w=fields.get("pos_w"), + quat_w=fields.get("quat_w"), + lin_vel_w=fields.get("lin_vel_w"), + ang_vel_w=fields.get("ang_vel_w"), + ) + if requested == self._body_names[owner.name]: + self._aggregate_body_views[owner.name] = view + return view capabilities = self._scene._backend.get_tensor_capabilities() - if getattr(capabilities, "tracked_body_views", False): + if getattr(capabilities, "tracked_body_views", False) and set(compiled_fields) == set( + _TENSOR_BODY_FIELDS + ): try: # A single requested body needs no four-block column projection; # the per-body canonical sensors are already the exact packed @@ -553,8 +717,9 @@ def body_tensor_view( self._aggregate_body_views[owner.name] = view return view - fields: dict[str, torch.Tensor] = {} - for field, prefix in _TENSOR_BODY_SENSOR_FIELDS: + stacked_fields: dict[str, torch.Tensor] = {} + for field in compiled_fields: + prefix = _TRACKED_BODY_SENSOR_PREFIXES[field] stacked = [ Entity._validate_tensor_sensor( self._packet[prefix + body_name], @@ -568,14 +733,122 @@ def body_tensor_view( ) for body_name in requested ] - fields[field] = torch.stack(stacked, dim=1) + stacked_fields[field] = torch.stack(stacked, dim=1) return EntityTensorBodyStateView( body_names=requested, - pos_w=fields["pos_w"], - quat_w=fields["quat_w"], - lin_vel_w=fields["lin_vel_w"], - ang_vel_w=fields["ang_vel_w"], + pos_w=stacked_fields["pos_w"], + quat_w=stacked_fields["quat_w"], + lin_vel_w=stacked_fields["lin_vel_w"], + ang_vel_w=stacked_fields["ang_vel_w"], + ) + + def _compiled_body_request( + self, owner: "Entity", body_names: Sequence[str] | None + ) -> tuple[str, ...]: + if owner._body_names is None: + raise owner._capability_error( + "scene body tensor state", "body_names were not declared in EntityCfg" + ) + return self._scene._normalize_tensor_body_request( + owner, tuple(body_names) if body_names is not None else owner._body_names + ) + + def _packet_body_fields( + self, + owner: "Entity", + requested: tuple[str, ...], + *, + required_fields: tuple[str, ...], + capability: str, + ) -> dict[str, torch.Tensor]: + fields: dict[str, torch.Tensor] = {} + missing: list[str] = [] + compiled_fields = self._body_fields[owner.name] + for field in required_fields: + if field not in compiled_fields: + missing.append(field) + if missing: + raise ValueError( + f"Entity '{owner.name}' {capability} fields {missing} were not compiled " + f"into the packed read; compiled={list(compiled_fields)}" + ) + compiled_bodies = self._body_names[owner.name] + if not self._packet_body_blocks: + canonical_fields: dict[str, torch.Tensor] = {} + for field, prefix in _TENSOR_BODY_SENSOR_FIELDS: + if field not in required_fields: + continue + values = tuple(self._packet.get(prefix + name) for name in requested) + if any(value is None for value in values): + missing.append(field) + continue + canonical_fields[field] = torch.stack( + tuple(value for value in values if value is not None), dim=1 + ) + if missing: + raise ValueError( + f"Entity '{owner.name}' {capability} fields {missing} were not compiled " + f"into the packed read; compiled={list(compiled_fields)}" + ) + return canonical_fields + for field in required_fields: + packet_name = _TENSOR_BODY_PACKET_FIELDS[field] + raw = self._packet.get(packet_name) + if raw is None: + missing.append(field) + continue + axis = self._packet_body_axes.get(packet_name, ()) + if not axis: + missing.append(field) + continue + width = _TENSOR_PACKET_FIELD_WIDTHS[packet_name] + expected = (self._scene._backend.num_envs, len(axis) * width) + if tuple(raw.shape) != expected: + raise ValueError( + f"Entity '{owner.name}' {capability} field '{field}' has shape " + f"{tuple(raw.shape)}; expected {expected}" + ) + if raw.dtype != torch.float32 or raw.device != self._device: + raise TypeError( + f"Entity '{owner.name}' {capability} field '{field}' must be float32 " + f"on {self._device}; got {raw.dtype} on {raw.device}" + ) + # ``unflatten`` preserves a view for row-strided packet slices. + # ``reshape`` would silently materialize a contiguous copy when the + # backend publishes one packed body block beside unrelated columns. + value = raw.unflatten(1, (len(axis), width)) + if self._host_plan_value is not None: + owner_axis = compiled_bodies + requested_axis = requested + else: + owner_axis = tuple(_packet_body_name(owner.name, name) for name in compiled_bodies) + requested_axis = tuple(_packet_body_name(owner.name, name) for name in requested) + if requested_axis != owner_axis or owner_axis != axis: + rows = self._packet_body_indices(requested_axis, axis) + value = value.index_select(1, rows) + fields[field] = value + if missing: + raise ValueError( + f"Entity '{owner.name}' {capability} fields {missing} were not compiled " + f"into the packed read; compiled={list(self._body_fields.get(owner.name, ()))}" + ) + return fields + + def _packet_body_indices( + self, requested: tuple[str, ...], compiled: tuple[str, ...] + ) -> torch.Tensor: + key: tuple[str, tuple[str, ...]] = ("body", compiled) + cached = self._packet_body_index_cache.get(key) + if cached is not None: + return cached + positions = {name: index for index, name in enumerate(compiled)} + value = torch.tensor( + [positions[name] for name in requested], + dtype=torch.int64, + device=self._device, ) + self._packet_body_index_cache[key] = value + return value def _validated_aggregate_body_view( self, owner: "Entity", views: Any @@ -632,18 +905,32 @@ def close(self) -> None: self._host_plan_value.close() finally: self._aggregate_body_views.clear() + self._packet_body_index_cache.clear() self._packet = MappingProxyType({}) self._refreshed = False self._closed = True - def _aggregate_packet_names(self) -> tuple[str, ...]: - names: list[str] = [] - for entity_names in self._sensor_names.values(): - names.extend(entity_names) - for entity_names in self._body_names.values(): - for body_name in entity_names: - names.extend(prefix + body_name for _, prefix in _TENSOR_BODY_SENSOR_FIELDS) - return tuple(dict.fromkeys(names)) + @staticmethod + def _aggregate_packet_body_axes( + body_names: Mapping[str, tuple[str, ...]], + body_fields: Mapping[str, set[str] | tuple[str, ...]], + ) -> dict[str, tuple[str, ...]]: + axes: dict[str, tuple[str, ...]] = {} + for entity_name, names in body_names.items(): + if not names: + continue + for field in body_fields[entity_name]: + packet_name = _TENSOR_BODY_PACKET_FIELDS[field] + axis = axes.setdefault(packet_name, ()) + axes[packet_name] = ( + *axis, + *( + _packet_body_name(entity_name, body_name) + for body_name in names + if _packet_body_name(entity_name, body_name) not in axis + ), + ) + return axes def _read_packet( self, reader: Callable[[], Mapping[str, Any]] | None @@ -663,9 +950,23 @@ def _read_packet( # stable adapter-owned storage here; only a fresh call proves # that storage reflects the current post-step/post-reset state. raw = dict(backend.get_state_views(("qpos", "qvel"), device=self._device)) + body_packet_names = tuple( + name for name in self._packet_names if name not in self._packet_sensor_owners + ) + body_view: Any = None + if body_packet_names: + body_names = tuple( + name.partition("/")[2] if "/" in name else name + for name in self._packet_body_axes[body_packet_names[0]] + ) + body_view = backend.get_tracked_body_views(body_names, device=self._device) for name in self._packet_names: - view = backend.get_sensor_view(name, device=self._device) - raw[name] = view + if name in self._packet_sensor_owners: + raw[name] = backend.get_sensor_view(name, device=self._device) + else: + raw[name] = getattr(body_view, _PACKET_FIELD_BODY_FIELDS[name]).reshape( + backend.num_envs, -1 + ) except (AttributeError, KeyError, TypeError, ValueError, NotImplementedError) as exc: raise type(exc)( "Manager scene device-resident tensor read on backend " @@ -706,17 +1007,38 @@ def _publish_packet(self, packet: dict[str, torch.Tensor]) -> None: f"got {value.dtype} on {value.device}" ) for name in self._packet_names: - owner = self._packet_sensor_owners[name] - packet[name] = Entity._validate_tensor_sensor( - packet[name], - entity_name=owner, - backend_type=backend.backend_type, - sensor_name=name, - expected_width=None, - num_envs=num_envs, - device=self._device, - finite=False, + value = packet[name] + if name in self._packet_sensor_owners: + packet[name] = Entity._validate_tensor_sensor( + value, + entity_name=self._packet_sensor_owners[name], + backend_type=backend.backend_type, + sensor_name=name, + expected_width=None, + num_envs=num_envs, + device=self._device, + finite=False, + ) + continue + if not isinstance(value, torch.Tensor): + raise TypeError( + f"Scene tensor body field '{name}' on backend " + f"'{backend.backend_type}' is {type(value).__name__}, expected torch.Tensor" + ) + expected = ( + num_envs, + len(self._packet_body_axes[name]) * _TENSOR_PACKET_FIELD_WIDTHS[name], ) + if tuple(value.shape) != expected: + raise ValueError( + f"Scene tensor body field '{name}' on backend " + f"'{backend.backend_type}' has shape {tuple(value.shape)}; expected {expected}" + ) + if value.dtype != torch.float32 or value.device != self._device: + raise TypeError( + f"Scene tensor body field '{name}' must be float32 on {self._device}; " + f"got {value.dtype} on {value.device}" + ) self._packet = MappingProxyType(packet) self._aggregate_body_views.clear() self._refreshed = True @@ -1436,6 +1758,8 @@ def __init__( self.root_body_name = cfg.root_body_name self._backend_type = backend.backend_type self._backend = backend + self._interval_operand_device_value: torch.device | None = None + self._interval_root_body_ids_value: torch.Tensor | None = None self._body_tensor_layouts: dict[tuple[str, ...], np.ndarray] = {} self._reset_state = reset_state self._reset_root_layout: RootStateLayout | None = None @@ -3453,7 +3777,9 @@ def _bind_root_velocity_delta(self, *, angular: bool, term_name: str) -> None: if angular else INTERVAL_TERM_BODY_LINEAR_VELOCITY_DELTA ) - supported = capabilities.supports_interval_term(term) + supported = capabilities.supports_interval_term( + term, plane=self._interval_capability_plane() + ) if not supported: raise self._capability_error( "interval root velocity delta", @@ -3528,9 +3854,19 @@ def apply_root_velocity_delta_to_sim( ) from exc def _interval_root_body_ids(self) -> torch.Tensor: - """Return the immutable root body IDs as a torch-first host int64 tensor.""" + """Return immutable root body IDs on the backend-bound interval plane.""" assert self._root_body_ids is not None - return torch.tensor(self._root_body_ids, dtype=torch.int64) + device = self._interval_operand_device() + cached = self._interval_root_body_ids_value + if cached is not None and cached.device == device: + return cached + cached = torch.as_tensor( + np.asarray(self._root_body_ids).copy(), + dtype=torch.int64, + device=device, + ) + self._interval_root_body_ids_value = cached + return cached def _validate_root_velocity_delta( self, @@ -3559,13 +3895,19 @@ def _validate_root_velocity_delta( raise ValueError( f"EventManager term '{term_name}' root velocity delta must be finite floating data" ) + device = self._interval_operand_device() + if values.device != device: + raise ValueError( + f"EventManager term '{term_name}' root velocity delta must live on " + f"{device}; got {values.device}" + ) assert self._root_body_ids is not None delta = torch.zeros( (self._backend.num_envs, len(self._root_body_ids), 3), - dtype=values.dtype, - device=values.device, + dtype=torch.float32, + device=device, ) - delta[ids.to(device=values.device, dtype=torch.int64), 0, :] = values + delta[ids.to(device=device, dtype=torch.int64), 0, :] = values.to(dtype=torch.float32) return delta def bind_body_wrench( @@ -3593,18 +3935,38 @@ def bind_body_wrench( capabilities = self._backend.get_dr_capabilities() except (AttributeError, NotImplementedError) as exc: raise self._capability_error("interval body wrench", str(exc)) from exc - if not capabilities.supports_interval_term(INTERVAL_TERM_BODY_FORCE): + plane = self._interval_capability_plane() + if not capabilities.supports_interval_term(INTERVAL_TERM_BODY_FORCE, plane=plane): raise self._capability_error( "interval body wrench", f"EventManager term '{term_name}' requested an unsupported backend capability", ) - if torque and not capabilities.supports_interval_term(INTERVAL_TERM_BODY_TORQUE): + if torque and not capabilities.supports_interval_term( + INTERVAL_TERM_BODY_TORQUE, plane=plane + ): raise self._capability_error( "interval body wrench", f"EventManager term '{term_name}' requested an unsupported backend " "capability: interval body torque", ) - return torch.as_tensor(local_ids), self._body_ids[local_ids] + return torch.as_tensor(local_ids, device=self._interval_operand_device()), self._body_ids[ + local_ids + ] + + def _interval_operand_device(self) -> torch.device: + """Return the backend-bound device for direct interval operands.""" + + if self._interval_operand_device_value is not None: + return self._interval_operand_device_value + views = self._backend.get_state_views(("qpos",), device=None) + device = views["qpos"].device + self._interval_operand_device_value = device + return device + + def _interval_capability_plane(self) -> DataPlaneName: + """Return the capability plane used by this entity's interval payloads.""" + + return "device" if self._interval_operand_device().type != "cpu" else "host" def apply_body_wrench_to_sim( self, @@ -3637,7 +3999,12 @@ def apply_body_wrench_to_sim( term_name=term_name, ) ops = [] - torch_body_ids = torch.tensor(np.asarray(body_ids), dtype=torch.int64) + device = self._interval_operand_device() + torch_body_ids = torch.tensor(np.asarray(body_ids), dtype=torch.int64, device=device) + if force_values is not None: + force_values = force_values.to(device=device, dtype=torch.float32) + if torque_values is not None: + torque_values = torque_values.to(device=device, dtype=torch.float32) if force_values is not None: ops.append( IntervalTermOp(INTERVAL_TERM_BODY_FORCE, force_values, body_ids=torch_body_ids) @@ -4224,6 +4591,7 @@ def compile_tensor_reads( entities: dict[str, Entity] = {} per_entity_sensors: dict[str, tuple[str, ...]] = {} per_entity_bodies: dict[str, tuple[str, ...]] = {} + per_entity_body_fields: dict[str, set[str]] = {} request_owners: dict[str, str] = {} for spec in normalized: try: @@ -4236,6 +4604,7 @@ def compile_tensor_reads( entities[spec.entity] = entity sensors = per_entity_sensors.setdefault(spec.entity, ()) bodies = per_entity_bodies.setdefault(spec.entity, ()) + fields = per_entity_body_fields.setdefault(spec.entity, set()) for name in spec.sensor_names: owner = request_owners.get(name) if owner is not None and owner != spec.entity: @@ -4246,6 +4615,8 @@ def compile_tensor_reads( request_owners[name] = spec.entity per_entity_sensors[spec.entity] = (*sensors, *spec.sensor_names) per_entity_bodies[spec.entity] = (*bodies, *spec.body_names) + if spec.body_names: + fields.update(spec.body_fields) expanded_bodies: dict[str, tuple[str, ...]] = {} for entity_name, body_names in per_entity_bodies.items(): @@ -4265,17 +4636,68 @@ def compile_tensor_reads( capabilities = self._backend.get_tensor_capabilities() self._validate_tensor_read_capabilities(capabilities, resolved_device) - sensor_names = self._aggregate_tensor_read_names( - entities, per_entity_sensors, expanded_bodies - ) - host_plan: HostBridgeTransferPlan | None = None - if capabilities.execution is TensorExecution.HOST_BRIDGE: - spec = TensorIOSpec( - state_fields=("qpos", "qvel"), - sensor_names=sensor_names, - device=resolved_device, + host_body_blocks = False + host_sensor_packet_names: tuple[str, ...] | None = None + sensor_packet_names = self._aggregate_tensor_read_names(tuple(entities), per_entity_sensors) + requested_body_fields = any(expanded_bodies.values()) + packet_body_blocks = False + if capabilities.execution is TensorExecution.DEVICE_RESIDENT and requested_body_fields: + packet_body_blocks = bool(getattr(capabilities, "tracked_body_views", False)) + packet_names = self._aggregate_tensor_read_names( + tuple(entities), per_entity_sensors, expanded_bodies, per_entity_body_fields + ) + packet_body_axes: dict[str, tuple[str, ...]] = {} + packet_body_names = tuple( + _TENSOR_BODY_PACKET_FIELDS[field] + for field in _TENSOR_BODY_FIELD_ORDER + if any(field in fields for fields in per_entity_body_fields.values()) + ) + if packet_body_blocks: + packet_body_axes = SceneTensorReadPlan._aggregate_packet_body_axes( + expanded_bodies, per_entity_body_fields ) + packet_names = (*sensor_packet_names, *packet_body_names) + if capabilities.execution is TensorExecution.HOST_BRIDGE: + host_body_blocks = {"body_names", "body_fields"}.issubset( + inspect.signature(TensorIOSpec).parameters + ) + native_body_names: tuple[str, ...] = () + host_pose_only_blocks = host_body_blocks and all( + not (set(fields) - set(_TENSOR_POSE_BODY_FIELDS)) + for fields in per_entity_body_fields.values() + ) + if host_pose_only_blocks and any(expanded_bodies.values()): + packet_body_blocks = True + packet_names = (*sensor_packet_names, *packet_body_names) + host_sensor_packet_names = sensor_packet_names + # Host-bridge body blocks use the backend's native public body + # namespace. Entity ownership remains available through the + # per-entity read-plan mappings below. + native_body_names = tuple( + dict.fromkeys( + body_name + for body_names_value in expanded_bodies.values() + for body_name in body_names_value + ) + ) + for packet_name in packet_body_names: + packet_body_axes[packet_name] = native_body_names + tensor_io_options: dict[str, Any] = { + "state_fields": ("qpos", "qvel"), + "sensor_names": ( + sensor_packet_names if host_sensor_packet_names is not None else packet_names + ), + "device": resolved_device, + } + if host_pose_only_blocks and any(expanded_bodies.values()): + tensor_io_options.update( + { + "body_names": native_body_names, + "body_fields": ("pos_w", "quat_w"), + } + ) + spec = TensorIOSpec(**tensor_io_options) try: host_plan = self._backend.bind_host_bridge_io(spec) except (AttributeError, KeyError, TypeError, ValueError, NotImplementedError) as exc: @@ -4290,6 +4712,18 @@ def compile_tensor_reads( entities=entities, sensor_names=per_entity_sensors, body_names=expanded_bodies, + body_fields={ + name: tuple(field for field in _TENSOR_BODY_FIELD_ORDER if field in fields) + for name, fields in per_entity_body_fields.items() + }, + packet_names=packet_names, + packet_sensor_owners={ + name: entity_name + for entity_name, names in per_entity_sensors.items() + for name in names + }, + packet_body_axes=packet_body_axes, + packet_body_blocks=packet_body_blocks, host_plan=host_plan, ) @@ -4343,9 +4777,10 @@ def _validate_tensor_read_capabilities( @staticmethod def _aggregate_tensor_read_names( - entities: Mapping[str, "Entity"], + entity_names: Sequence[str], sensor_names: Mapping[str, tuple[str, ...]], - body_names: Mapping[str, tuple[str, ...]], + body_names: Mapping[str, tuple[str, ...]] | None = None, + body_fields: Mapping[str, set[str] | tuple[str, ...]] | None = None, ) -> tuple[str, ...]: aggregate: list[str] = [] @@ -4355,13 +4790,32 @@ def add(name: str) -> None: aggregate.append(name) seen: set[str] = set() - for entity_name in entities: + for entity_name in entity_names: for name in sensor_names[entity_name]: add(name) - for entity_name in entities: - for body_name in body_names[entity_name]: - for _, prefix in _TENSOR_BODY_SENSOR_FIELDS: - add(prefix + body_name) + if body_names is not None and body_fields is not None: + for entity_name in entity_names: + fields = body_fields[entity_name] + if entity_name not in body_names: + continue + for body_name in body_names[entity_name]: + for field, prefix in _TENSOR_BODY_SENSOR_FIELDS: + if field in fields: + add(prefix + body_name) + return tuple(aggregate) + + @staticmethod + def _aggregate_tensor_read_body_names( + body_names: Mapping[str, tuple[str, ...]], + ) -> tuple[str, ...]: + aggregate: list[str] = [] + for entity_name, names in body_names.items(): + if not names: + continue + for name in names: + packet_name = _packet_body_name(entity_name, name) + if packet_name not in aggregate: + aggregate.append(packet_name) return tuple(aggregate) @contextmanager @@ -4466,5 +4920,6 @@ def __len__(self) -> int: "EntityTensorSensorViews", "EntityTensorStateView", "SceneTensorReadPlan", + "SceneTensorPoseView", "SceneTensorReadSpec", ] diff --git a/src/unilab/base/reset_state.py b/src/unilab/base/reset_state.py index d472a98b4..865336071 100644 --- a/src/unilab/base/reset_state.py +++ b/src/unilab/base/reset_state.py @@ -1812,9 +1812,9 @@ def _packed_reset_plane(self) -> DataPlaneName: """Data plane of the packed composer for capability checks. Torch tensors on ``torch.cpu`` form the host plane; any other device - forms the device plane. Before the packed device is declared (cold-path - binds run before ``declare_packed_reset_device``) the check falls back - to the host plane, matching the historical declaration semantics. + forms the device plane. The Manager runtime declares this device before + cold-path event binds; standalone transactions must do the same before + constructing terms so capability checks match the eventual commit plane. """ device = self._packed_reset_device if device is None or device.type == "cpu": diff --git a/src/unilab/base/torch_env.py b/src/unilab/base/torch_env.py index 35f939a96..1a977bb4b 100644 --- a/src/unilab/base/torch_env.py +++ b/src/unilab/base/torch_env.py @@ -69,10 +69,9 @@ def _cpu_time() -> float: # Backends whose ``step_tensor`` publishes an event-backed asynchronous step # receipt (a unisim ``CompletionToken`` carried in ``info["completion"]``). -# Synchronous producers — host-bridge backends and device-resident backends -# that synchronize before returning (for example genesis) — publish no -# receipt, so consumers must accept a missing receipt without waiting. -_EVENT_BACKED_STEP_RECEIPT_BACKENDS = frozenset({"mjwarp"}) +# Event-backed Genesis and MJWarp CUDA paths publish a receipt per step; +# synchronous host-bridge producers publish none. +_EVENT_BACKED_STEP_RECEIPT_BACKENDS = frozenset({"genesis", "mjwarp"}) class TorchEnv(ABEnv): diff --git a/src/unilab/conf/ppo/task/g1_walk_flat/genesis.yaml b/src/unilab/conf/ppo/task/g1_walk_flat/genesis.yaml index 9af74d93b..e7d8ce7cd 100644 --- a/src/unilab/conf/ppo/task/g1_walk_flat/genesis.yaml +++ b/src/unilab/conf/ppo/task/g1_walk_flat/genesis.yaml @@ -9,8 +9,8 @@ # stay at Genesis defaults: the MJCF option block does not declare them and # MuJoCo's implicit defaults (PGS solver, pyramidal cone) have no Genesis # equivalent. Keyframes and sensors are scanned from the MJCF once on the -# materialize cold path; kp/kd reset randomization stays enabled because the -# backend declares the measured RESET_TERM_KP/KD DR terms (REPORT §5.7). +# materialize cold path; the device-plane packed reset lane declares kp/kd, so +# PD-gain randomization follows the shared owner baseline (REPORT #1372 §5.7). # Playback uses Genesis native rendering: with a display, play_render_mode=auto # opens the interactive viewer; headless hosts fall back to offscreen camera # recording. diff --git a/src/unilab/conf/sac/task/g1_motion_tracking/genesis.yaml b/src/unilab/conf/sac/task/g1_motion_tracking/genesis.yaml index c9881063f..07b43b69a 100644 --- a/src/unilab/conf/sac/task/g1_motion_tracking/genesis.yaml +++ b/src/unilab/conf/sac/task/g1_motion_tracking/genesis.yaml @@ -6,11 +6,13 @@ # flows through the existing base.yaml sim_dt -> ctrl_dt chain, and the # integrator through the env field below. Requires unisim-core>=1.5.1 for the # motion body id capability (unilabsim/unisim#137), plus the contact-sensor -# scan fix for found+mindist geom pairs (unilabsim/unisim#276). Under the -# genesis-world==1.3.3 scene loading pinned here, the G1 scene declares no -# model-field reset terms and the backend exposes no interval velocity-delta -# capability, so base_com/foot_friction/push_robot stay disabled; encoder_bias -# is a host-side observation bias and stays enabled. Playback uses Genesis +# scan fix for found+mindist geom pairs (unilabsim/unisim#276), and the +# device-resident Genesis DR subset from issue #2521. The owner enables +# startup pelvis/torso mass scaling, inherited torso COM reset, reset PD-gain +# scaling, and the six-axis interval velocity kick. Absolute geom friction +# remains intentionally fail-closed because Genesis has no equivalent +# per-environment friction API; encoder_bias is a host-side observation bias +# and stays enabled. Playback uses Genesis # native rendering: with a display, play_render_mode=auto opens the # interactive viewer; headless hosts fall back to offscreen camera recording. defaults: @@ -34,14 +36,39 @@ env: scene: entities: robot: - # The genesis backend does not expose portable geom names; the - # foot-geom declaration only feeds the disabled foot_friction DR - # binding. geom_names: null events: - # With the genesis-world==1.3.3 scene loading pinned above, the G1 scene - # declares no model-field reset terms and the backend exposes no interval - # velocity-delta capability; keep the owner fail-closed. - base_com: null + body_mass: + func: unilab.envs.mdp.randomize_rigid_body_mass + mode: startup + params: + asset_cfg: + _target_: unilab.managers.SceneEntityCfg + name: robot + body_names: [pelvis, torso_link] + mass_distribution_params: [0.9, 1.1] + operation: scale + distribution: uniform + recompute_inertia: false foot_friction: null - push_robot: null + pd_gains: + func: unilab.envs.mdp.pd_gains + mode: reset + params: + kp_range: [0.9, 1.1] + kd_range: [0.9, 1.1] + distribution: uniform + operation: scale + push_robot: + func: unilab.envs.mdp.push_by_setting_velocity + mode: interval + interval_range_s: [1.0, 3.0] + is_global_time: true + params: + 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] diff --git a/src/unilab/conf/sac/task/g1_walk_flat/genesis.yaml b/src/unilab/conf/sac/task/g1_walk_flat/genesis.yaml index d0ee834b4..5d53a6f51 100644 --- a/src/unilab/conf/sac/task/g1_walk_flat/genesis.yaml +++ b/src/unilab/conf/sac/task/g1_walk_flat/genesis.yaml @@ -8,10 +8,9 @@ # sim_dt -> ctrl_dt chain, and the integrator through the EnvCfg genesis # fields below. Solver/constraint-cone/iterations stay at Genesis defaults: # the MJCF option block does not declare them and MuJoCo's implicit defaults -# (PGS solver, pyramidal cone) have no Genesis equivalent. Unlike the -# shelved Isaac owner (events.pd_gains: null), the public Genesis tensor selected -# reset does not support randomization, so unsupported host reset terms are -# explicitly disabled rather than falling back. Playback uses Genesis native +# (PGS solver, pyramidal cone) have no Genesis equivalent. The device-plane +# packed reset lane declares kp/kd, so PD-gain randomization follows the shared +# owner baseline. Playback uses Genesis native # rendering: with a display, play_render_mode=auto opens the interactive # viewer; headless hosts fall back to offscreen camera recording. defaults: @@ -76,8 +75,6 @@ env: func: unilab.envs.mdp.skip_scene_reset reset_root_state_uniform: func: unilab.envs.mdp.reset_root_state_uniform_fused - # Genesis' public tensor mask reset does not support randomization. - pd_gains: null terminations: tilt: null base_height: null diff --git a/src/unilab/envs/manager_based_rl_env.py b/src/unilab/envs/manager_based_rl_env.py index 9bd1ae60b..5253fe4d0 100644 --- a/src/unilab/envs/manager_based_rl_env.py +++ b/src/unilab/envs/manager_based_rl_env.py @@ -398,7 +398,10 @@ def __init__(self, cfg: ManagerBasedRlEnvCfg, backend: SimBackend, num_envs: int self._reset_mask_buffer = ResetMaskBuffer(num_envs, self.device) self.event_term_mask_buffer = ResetMaskBuffer(num_envs, self.device) self._has_transition = False - + # Cold-path binds happen while managers construct. Declare the Manager + # carrier's device now so reset-capability negotiation checks the same + # plane that the packed reset transaction will use after materialize. + self._reset_state.declare_packed_reset_device(self.device) self._load_managers() self._mapped_obs_dims = self._validate_observation_mapping() self._validate_substep_capabilities() @@ -427,12 +430,14 @@ def __init__(self, cfg: ManagerBasedRlEnvCfg, backend: SimBackend, num_envs: int env_ids=self._all_env_rows if self.event_manager.uses_startup_rows else None, ) self._compile_tensor_read_plan() + self.observation_manager.bind_read_phase() self.command_manager.bind_read_phase() self._validate_manager_tensor_runtime() def _compile_tensor_read_plan(self) -> None: """Compile the scene's only packed tensor read phase.""" specs: list[SceneTensorReadSpec] = [] + specs.extend(self._observation_ray_read_specs()) specs.extend(self._action_tensor_read_specs()) specs.extend(self._observation_tensor_read_specs()) specs.extend(self._task_tensor_state_read_specs()) @@ -507,6 +512,10 @@ def _observation_tensor_read_specs(self) -> list[SceneTensorReadSpec]: ) ) term = term_cfg.func + if getattr(term, "ray_pattern", None) is not None: + # Ray terms already declared a pose-only body request above. + # They must not also add the legacy four-field body-state set. + continue body_names = term_cfg.params.get("tensor_body_names") if body_names is None: body_names = getattr(term, "tensor_body_names", None) @@ -532,6 +541,41 @@ def _observation_tensor_read_specs(self) -> list[SceneTensorReadSpec]: specs.append(SceneTensorReadSpec(entity=entity_name, body_names=tuple(body_names))) return specs + def _observation_ray_read_specs(self) -> list[SceneTensorReadSpec]: + """Collect pose-only blocks needed by tensor ray observations.""" + specs: list[SceneTensorReadSpec] = [] + requested_entities: set[str] = set() + for group_name, terms in self.observation_manager.active_terms.items(): + for name in terms: + term_cfg = self.observation_manager.get_term_cfg(group_name, name) + ray_pattern = getattr(term_cfg.func, "ray_pattern", None) + if ray_pattern is None: + continue + if not isinstance(ray_pattern, mdp.RayPatternCfg): + raise TypeError( + "ManagerBasedRlEnv ray declaration for observation term " + f"'{group_name}/{name}' must be RayPatternCfg; got " + f"{type(ray_pattern).__name__}" + ) + for entity_name, entity in self.scene.entities.items(): + if entity_name in requested_entities: + continue + try: + body_names = tuple(entity.body_names) + except NotImplementedError: + continue + if not body_names: + continue + requested_entities.add(entity_name) + specs.append( + SceneTensorReadSpec( + entity=entity_name, + body_names=body_names, + body_fields=("pos_w", "quat_w"), + ) + ) + return specs + def _manager_term_tensor_read_specs(self) -> list[SceneTensorReadSpec]: specs: list[SceneTensorReadSpec] = [] narrow = ( @@ -1725,6 +1769,7 @@ def import_training_state(self, state: Mapping[str, Any]) -> None: self._sim_step_counter = self.step_counter * self._cfg.sim_substeps def close(self) -> None: + self.observation_manager.close() read_plan = self.scene.read_plan if read_plan is not None: read_plan.close() diff --git a/src/unilab/envs/mdp/__init__.py b/src/unilab/envs/mdp/__init__.py index a7e460a0b..9ef797a7e 100644 --- a/src/unilab/envs/mdp/__init__.py +++ b/src/unilab/envs/mdp/__init__.py @@ -58,6 +58,18 @@ from unilab.envs.mdp.observations import ( projected_gravity_imu_misaligned as projected_gravity_imu_misaligned, ) +from unilab.envs.mdp.ray_observation import RayAlignment as RayAlignment +from unilab.envs.mdp.ray_observation import ( + RayCameraPatternCfg as RayCameraPatternCfg, +) +from unilab.envs.mdp.ray_observation import ( + RayGridPatternCfg as RayGridPatternCfg, +) +from unilab.envs.mdp.ray_observation import RayPatternCfg as RayPatternCfg +from unilab.envs.mdp.ray_observation import ( + RayRingPatternCfg as RayRingPatternCfg, +) +from unilab.envs.mdp.ray_observation import RayScan as RayScan from unilab.envs.mdp.recorders import LifecycleCounterRecorder as LifecycleCounterRecorder 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 @@ -96,6 +108,12 @@ "UniformVelocityCommandCfg", "UniformPoseCommand", "UniformPoseCommandCfg", + "RayPatternCfg", + "RayAlignment", + "RayCameraPatternCfg", + "RayGridPatternCfg", + "RayRingPatternCfg", + "RayScan", "action_acc_l2", "action_rate_l2", "apply_body_impulse", diff --git a/src/unilab/envs/mdp/ray_observation.py b/src/unilab/envs/mdp/ray_observation.py new file mode 100644 index 000000000..5f2caf4a3 --- /dev/null +++ b/src/unilab/envs/mdp/ray_observation.py @@ -0,0 +1,826 @@ +"""Tensor-native ray-pattern observation terms for the Manager runtime.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Literal, Sequence + +import torch +from unisim.backend.base import torch_device_matches +from unisim.ray_query import ( + RayGeometryTensorView, + RayPoseTensorView, + RayTraceFilter, + RayTraceOutputs, +) + +from unilab.managers.manager_base import ManagerTermBase, ManagerTermBaseCfg + +if TYPE_CHECKING: + from unilab.managers.types import ManagerBasedRlEnv + + +def _real_sequence(value: object, name: str, width: int) -> tuple[float, ...]: + """Validate and normalize a finite real sequence used by ray patterns.""" + + if isinstance(value, (str, bytes)) or not isinstance(value, Sequence): + raise TypeError(f"{name} must be a sequence of {width} real numbers") + values = tuple(value) + if len(values) != width or any( + isinstance(item, bool) or not isinstance(item, (int, float)) for item in values + ): + raise TypeError(f"{name} must be a sequence of {width} real numbers") + if any(not math.isfinite(float(item)) for item in values): + raise ValueError(f"{name} must be finite") + return tuple(float(item) for item in values) + + +@dataclass(frozen=True) +class RayRingPatternCfg: + """Rays on a body-local yaw arc with one shared pitch.""" + + num_rays: int + start: float + stop: float + pitch: float = 0.0 + + def __post_init__(self) -> None: + if isinstance(self.num_rays, bool) or not isinstance(self.num_rays, int): + raise TypeError("RayRingPatternCfg num_rays must be an integer") + if self.num_rays <= 0: + raise ValueError("RayRingPatternCfg num_rays must be positive") + for name, value in (("start", self.start), ("stop", self.stop), ("pitch", self.pitch)): + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise TypeError(f"RayRingPatternCfg {name} must be a real number") + if not math.isfinite(float(value)): + raise ValueError(f"RayRingPatternCfg {name} must be finite") + + @property + def ray_count(self) -> int: + return self.num_rays + + def origins(self, device: torch.device) -> torch.Tensor: + return torch.zeros((self.ray_count, 3), dtype=torch.float32, device=device) + + def directions(self, device: torch.device) -> torch.Tensor: + angles = torch.linspace( + self.start, self.stop, self.num_rays, dtype=torch.float32, device=device + ) + pitch = float(self.pitch) + return torch.stack( + ( + torch.cos(angles) * math.cos(pitch), + torch.sin(angles) * math.cos(pitch), + torch.full_like(angles, math.sin(pitch)), + ), + dim=-1, + ).contiguous() + + +@dataclass(frozen=True) +class RayGridPatternCfg: + """Parallel rays on an inclusive two-dimensional local grid.""" + + size: Sequence[float] + resolution: float | Sequence[float] + offset: Sequence[float] = (0.0, 0.0) + direction: Sequence[float] = (0.0, 0.0, -1.0) + + def __post_init__(self) -> None: + size = _real_sequence(self.size, "RayGridPatternCfg size", 2) + if any(value <= 0.0 for value in size): + raise ValueError("RayGridPatternCfg size components must be positive") + object.__setattr__(self, "size", size) + + if isinstance(self.resolution, bool) or not isinstance( + self.resolution, (int, float, Sequence) + ): + raise TypeError( + "RayGridPatternCfg resolution must be a real number or two real numbers" + ) + if isinstance(self.resolution, (str, bytes)): + raise TypeError( + "RayGridPatternCfg resolution must be a real number or two real numbers" + ) + if isinstance(self.resolution, Sequence): + resolution = _real_sequence(self.resolution, "RayGridPatternCfg resolution", 2) + else: + value = float(self.resolution) + if not math.isfinite(value): + raise ValueError("RayGridPatternCfg resolution must be finite") + resolution = (value, value) + if any(value <= 0.0 for value in resolution): + raise ValueError("RayGridPatternCfg resolution components must be positive") + object.__setattr__(self, "resolution", resolution) + + object.__setattr__( + self, + "offset", + _real_sequence(self.offset, "RayGridPatternCfg offset", 2), + ) + direction = _real_sequence(self.direction, "RayGridPatternCfg direction", 3) + if torch.linalg.vector_norm(torch.tensor(direction)).item() == 0.0: + raise ValueError("RayGridPatternCfg direction must be nonzero") + object.__setattr__(self, "direction", direction) + + @property + def _counts(self) -> tuple[int, int]: + resolution = self.resolution + if not isinstance(resolution, Sequence): + resolution = (resolution, resolution) + return tuple( # type: ignore[return-value] + int(math.floor(axis / step + 1.0e-9)) + 1 + for axis, step in zip(self.size, resolution, strict=True) + ) + + @property + def ray_count(self) -> int: + count_x, count_y = self._counts + return count_x * count_y + + def origins(self, device: torch.device) -> torch.Tensor: + count_x, count_y = self._counts + x = torch.linspace(-self.size[0] / 2.0, self.size[0] / 2.0, count_x, device=device) + y = torch.linspace(-self.size[1] / 2.0, self.size[1] / 2.0, count_y, device=device) + grid_x, grid_y = torch.meshgrid(x, y, indexing="xy") + origins = torch.zeros((self.ray_count, 3), dtype=torch.float32, device=device) + origins[:, 0] = (grid_x.flatten() + self.offset[0]).float() + origins[:, 1] = (grid_y.flatten() + self.offset[1]).float() + return origins.contiguous() + + def directions(self, device: torch.device) -> torch.Tensor: + direction = torch.tensor(self.direction, dtype=torch.float32, device=device) + direction = direction / torch.linalg.vector_norm(direction) + return direction.expand(self.ray_count, 3).clone().contiguous() + + +@dataclass(frozen=True) +class RayCameraPatternCfg: + """Unit ray directions for a pinhole camera in an explicit local frame.""" + + width: int + height: int + fovy: float + forward: Sequence[float] = (1.0, 0.0, 0.0) + up: Sequence[float] = (0.0, 0.0, 1.0) + + def __post_init__(self) -> None: + if isinstance(self.width, bool) or not isinstance(self.width, int): + raise TypeError("RayCameraPatternCfg width must be an integer") + if isinstance(self.height, bool) or not isinstance(self.height, int): + raise TypeError("RayCameraPatternCfg height must be an integer") + if self.width <= 0 or self.height <= 0: + raise ValueError("RayCameraPatternCfg width and height must be positive") + if isinstance(self.fovy, bool) or not isinstance(self.fovy, (int, float)): + raise TypeError("RayCameraPatternCfg fovy must be a real number") + if not math.isfinite(float(self.fovy)) or not 0.0 < float(self.fovy) < 180.0: + raise ValueError("RayCameraPatternCfg fovy must be in (0, 180) degrees") + forward = _real_sequence(self.forward, "RayCameraPatternCfg forward", 3) + up = _real_sequence(self.up, "RayCameraPatternCfg up", 3) + forward_tensor = torch.tensor(forward, dtype=torch.float64) + up_tensor = torch.tensor(up, dtype=torch.float64) + if ( + torch.linalg.vector_norm(forward_tensor).item() == 0.0 + or torch.linalg.vector_norm(up_tensor).item() == 0.0 + or abs(float(torch.dot(forward_tensor, up_tensor).item())) > 1.0e-6 + ): + raise ValueError("RayCameraPatternCfg forward and up must be nonzero and orthogonal") + object.__setattr__(self, "forward", forward) + object.__setattr__(self, "up", up) + + @property + def ray_count(self) -> int: + return self.width * self.height + + def origins(self, device: torch.device) -> torch.Tensor: + return torch.zeros((self.ray_count, 3), dtype=torch.float32, device=device) + + def directions(self, device: torch.device) -> torch.Tensor: + forward = torch.tensor(self.forward, dtype=torch.float32, device=device) + up = torch.tensor(self.up, dtype=torch.float32, device=device) + right = torch.linalg.cross(forward, up) + right = right / torch.linalg.vector_norm(right) + focal = (self.height / 2.0) / math.tan(math.radians(self.fovy) / 2.0) + x = ((torch.arange(self.width, device=device) + 0.5) - self.width / 2.0) / focal + y = ((torch.arange(self.height, device=device) + 0.5) - self.height / 2.0) / focal + grid_x, grid_y = torch.meshgrid(x, y, indexing="xy") + directions = ( + forward[None, None, :] + + grid_x[..., None] * right[None, None, :] + - grid_y[..., None] * up[None, None, :] + ) + directions = directions / torch.linalg.vector_norm(directions, dim=-1, keepdim=True) + return directions.reshape(-1, 3).contiguous() + + +RayPattern = RayRingPatternCfg | RayGridPatternCfg | RayCameraPatternCfg +RayAlignment = Literal["body", "yaw", "world"] + + +@dataclass(frozen=True) +class RayPatternCfg: + """Declare one observation-owned ray pattern and its attachment frame. + + ``caster`` is a public UniSim ray-caster factory name, not an imported + plugin class. ``attachment_body_name`` is a public full-scene body identity + from ``RayCollisionDescription.scene.body_names``. The local pattern and + sensor offset are transformed by the selected attachment alignment. + """ + + entity_name: str + pattern: RayPattern + max_distance: float + attachment_body_name: str | None = None + attachment_site_name: str | None = None + alignment: RayAlignment = "body" + caster: str = "uni_ray" + sensor_offset: Sequence[float] = (0.0, 0.0, 0.0) + min_distance: float = 0.0 + include_normal: bool = False + include_hit_point: bool = False + sync_geometry: bool = False + allowed_geom_names: Sequence[str] | None = None + allowed_tags: Sequence[str] | None = None + excluded_body_names: Sequence[str] | None = None + excluded_ancestor_body_names: Sequence[str] | None = None + caster_options: dict[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + if not isinstance(self.entity_name, str) or not self.entity_name: + raise TypeError("RayPatternCfg entity_name must be a non-empty string") + if not isinstance( + self.pattern, + (RayRingPatternCfg, RayGridPatternCfg, RayCameraPatternCfg), + ): + raise TypeError( + "RayPatternCfg pattern must be RayRingPatternCfg, RayGridPatternCfg, " + "or RayCameraPatternCfg" + ) + if self.alignment not in ("body", "yaw", "world"): + raise ValueError("RayPatternCfg alignment must be one of 'body', 'yaw', or 'world'") + for name, value in ( + ("max_distance", self.max_distance), + ("min_distance", self.min_distance), + ): + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise TypeError(f"RayPatternCfg {name} must be a real number") + if not math.isfinite(float(value)): + raise ValueError(f"RayPatternCfg {name} must be finite") + if self.max_distance <= 0.0: + raise ValueError("RayPatternCfg max_distance must be positive") + if self.min_distance < 0.0 or self.min_distance >= self.max_distance: + raise ValueError( + "RayPatternCfg min_distance must satisfy 0 <= min_distance < max_distance" + ) + if (self.attachment_body_name is None) == (self.attachment_site_name is None): + raise ValueError( + "RayPatternCfg requires exactly one of attachment_body_name or attachment_site_name" + ) + for attachment_value in ( + self.attachment_body_name, + self.attachment_site_name, + ): + if attachment_value is not None and ( + not isinstance(attachment_value, str) or not attachment_value + ): + raise TypeError("RayPatternCfg attachment names must be non-empty strings") + if isinstance(self.sensor_offset, (str, bytes)) or not isinstance( + self.sensor_offset, Sequence + ): + raise TypeError("RayPatternCfg sensor_offset must be a sequence of three real numbers") + object.__setattr__( + self, + "sensor_offset", + _real_sequence(self.sensor_offset, "RayPatternCfg sensor_offset", 3), + ) + if not isinstance(self.include_normal, bool): + raise TypeError("RayPatternCfg include_normal must be a boolean") + if not isinstance(self.include_hit_point, bool): + raise TypeError("RayPatternCfg include_hit_point must be a boolean") + if not isinstance(self.sync_geometry, bool): + raise TypeError("RayPatternCfg sync_geometry must be a boolean") + if not isinstance(self.caster, str) or not self.caster: + raise TypeError("RayPatternCfg caster must be a non-empty factory name") + if not isinstance(self.caster_options, dict): + raise TypeError("RayPatternCfg caster_options must be a mapping") + if "collision" in self.caster_options: + raise ValueError( + "RayPatternCfg caster_options cannot override the public collision descriptor" + ) + for name in ( + "allowed_geom_names", + "allowed_tags", + "excluded_body_names", + "excluded_ancestor_body_names", + ): + value = getattr(self, name) + if value is None: + continue + if isinstance(value, (str, bytes)) or not isinstance(value, Sequence): + raise TypeError(f"RayPatternCfg {name} must be a string sequence or None") + values = tuple(value) + if any(not isinstance(item, str) or not item for item in values): + raise TypeError(f"RayPatternCfg {name} must contain non-empty strings") + if len(set(values)) != len(values): + raise ValueError(f"RayPatternCfg {name} must be unique: {values}") + object.__setattr__(self, name, values) + + +class RayScan(ManagerTermBase): + """Bind one tensor caster and publish a ray observation carrier.""" + + def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv): + super().__init__(env) + pattern = cfg.params.get("ray_pattern") + if not isinstance(pattern, RayPatternCfg): + raise TypeError( + "RayScan requires params['ray_pattern'] to be a RayPatternCfg; " + f"got {type(pattern).__name__}" + ) + self.ray_pattern = pattern + output_width = 1 + if pattern.include_normal: + output_width += 3 + if pattern.include_hit_point: + output_width += 3 + self.observation_width = pattern.pattern.ray_count * output_width + self._device = torch.device(env.device) + self._entity_name = pattern.entity_name + self._caster: Any = None + self._bound = False + self._closed = False + self._distance: torch.Tensor | None = None + self._normal: torch.Tensor | None = None + self._hit_point: torch.Tensor | None = None + self._observation: torch.Tensor | None = None + self._ray_filter: RayTraceFilter | None = None + self._outputs = RayTraceOutputs() + self._attachment_body_index: int | None = None + self._attachment_body_name: str | None = None + self._attachment_site_local_pos: torch.Tensor | None = None + self._attachment_site_local_quat: torch.Tensor | None = None + self._local_origins: torch.Tensor | None = None + self._local_directions: torch.Tensor | None = None + self._ray_origins: torch.Tensor | None = None + self._ray_directions: torch.Tensor | None = None + self._pose_sources: dict[str, tuple[tuple[str, ...], torch.Tensor]] = {} + self._pose_body_names: tuple[str, ...] = () + self._pose_pos: torch.Tensor | None = None + self._pose_quat: torch.Tensor | None = None + + def bind_read_phase(self) -> None: + """Materialize and bind the caster after packed pose reads compile.""" + if self._bound or self._closed: + return + read_plan = getattr(self._env.scene, "read_plan", None) + if read_plan is None: + raise NotImplementedError("RayScan requires the Manager packed Torch read plan") + try: + collision = self._env.backend.get_ray_collision_description() + scene = collision.scene + attachment_body_name = self.ray_pattern.attachment_body_name + if self.ray_pattern.attachment_site_name is not None: + (site_frame,) = self._env.backend.get_ray_site_frames( + (self.ray_pattern.attachment_site_name,) + ) + attachment_body_name = site_frame.body_name + self._attachment_site_local_pos = torch.tensor( + site_frame.pos, dtype=torch.float32, device=self._device + ) + self._attachment_site_local_quat = torch.tensor( + site_frame.quat, dtype=torch.float32, device=self._device + ) + assert attachment_body_name is not None + self._attachment_body_name = attachment_body_name + if attachment_body_name not in scene.body_names: + raise KeyError( + f"ray pattern attachment body {attachment_body_name!r} is not in " + f"{scene.body_names}" + ) + self._attachment_body_index = tuple(scene.body_names).index(attachment_body_name) + self._pose_sources = self._resolve_pose_sources(collision) + self._validate_pose_coverage(collision) + self._pose_body_names = tuple(scene.body_names) + self._ray_filter = self._resolve_ray_filter(collision) + caster = _create_ray_caster( + self.ray_pattern.caster, + num_envs=self._env.num_envs, + num_rays=self.ray_pattern.pattern.ray_count, + device=self._device, + collision=collision, + **self.ray_pattern.caster_options, + ) + try: + capabilities = caster.get_ray_capabilities() + if not ( + capabilities.supports_tensor_pose_sync + and capabilities.supports_tensor_rays + and capabilities.supports_device_output + and capabilities.stream_event_ownership == "caller" + and torch_device_matches(capabilities.torch_devices, self._device) + ): + raise NotImplementedError( + f"RayScan caster '{self.ray_pattern.caster}' does not support the " + f"tensor pose/ray/output profile on {self._device}" + ) + if self.ray_pattern.min_distance != 0.0 and not capabilities.supports_min_distance: + raise NotImplementedError( + f"RayScan caster '{self.ray_pattern.caster}' does not support min_distance" + ) + if self._ray_filter.has_constraints() and not capabilities.supports_static_filter: + raise NotImplementedError( + f"RayScan caster '{self.ray_pattern.caster}' does not support " + "static ray filters" + ) + if self.ray_pattern.include_normal and not capabilities.supports_tensor_normal: + raise NotImplementedError( + f"RayScan caster '{self.ray_pattern.caster}' does not support tensor normals" + ) + if self.ray_pattern.include_hit_point and not capabilities.supports_hit_point: + raise NotImplementedError( + f"RayScan caster '{self.ray_pattern.caster}' does not support " + "tensor hit points" + ) + if self.ray_pattern.sync_geometry and not ( + capabilities.supports_geometry_sync + and capabilities.supports_tensor_geometry_sync + ): + raise NotImplementedError( + f"RayScan caster '{self.ray_pattern.caster}' does not support " + "tensor geometry sync" + ) + caster.materialize(scene) + caster.bind_pose_layout(tuple(scene.body_names)) + except Exception: + caster.close() + raise + except (AttributeError, KeyError, TypeError, ValueError, NotImplementedError) as exc: + raise type(exc)( + f"RayScan caster '{self.ray_pattern.caster}' read-phase binding failed: {exc}" + ) from exc + self._caster = caster + self._outputs = RayTraceOutputs( + normal=self.ray_pattern.include_normal, + hit_point=self.ray_pattern.include_hit_point, + ) + self._distance = torch.empty( + (self._env.num_envs, self.ray_pattern.pattern.ray_count), + dtype=torch.float32, + device=self._device, + ) + if self.ray_pattern.include_normal or self.ray_pattern.include_hit_point: + output_width = 1 + if self.ray_pattern.include_normal: + output_width += 3 + self._normal = torch.empty( + (self._env.num_envs, self.ray_pattern.pattern.ray_count, 3), + dtype=torch.float32, + device=self._device, + ) + if self.ray_pattern.include_hit_point: + output_width += 3 + self._hit_point = torch.empty( + (self._env.num_envs, self.ray_pattern.pattern.ray_count, 3), + dtype=torch.float32, + device=self._device, + ) + self._observation = torch.empty( + (self._env.num_envs, self.ray_pattern.pattern.ray_count, output_width), + dtype=torch.float32, + device=self._device, + ) + self._local_origins = self.ray_pattern.pattern.origins(self._device) + self._local_origins = self._local_origins + torch.tensor( + self.ray_pattern.sensor_offset, + dtype=torch.float32, + device=self._device, + ) + self._local_directions = self.ray_pattern.pattern.directions(self._device) + self._ray_origins = torch.empty( + (self._env.num_envs, self.ray_pattern.pattern.ray_count, 3), + dtype=torch.float32, + device=self._device, + ) + self._ray_directions = torch.empty_like(self._ray_origins) + self._pose_pos = torch.zeros( + (self._env.num_envs, len(scene.body_names), 3), + dtype=torch.float32, + device=self._device, + ) + self._pose_quat = torch.zeros( + (self._env.num_envs, len(scene.body_names), 4), + dtype=torch.float32, + device=self._device, + ) + self._pose_quat[..., 0] = 1.0 + if self.ray_pattern.sync_geometry: + self._sync_geometry() + self._bound = True + + def __call__( + self, + env: ManagerBasedRlEnv, + ray_pattern: RayPatternCfg | None = None, + ) -> torch.Tensor: + if ray_pattern is not None and ray_pattern is not self.ray_pattern: + raise ValueError( + "RayScan was bound to a different ray_pattern; manager term parameters " + "must remain immutable" + ) + if not self._bound or self._caster is None or self._distance is None: + raise NotImplementedError("RayScan must bind its read phase before compute") + read_plan = getattr(env.scene, "read_plan", None) + if read_plan is None or not read_plan.ready: + raise RuntimeError("RayScan requires a refreshed packed read phase") + pose = self._refresh_full_pose(read_plan) + try: + self._refresh_attachment_rays(pose) + self._caster.update_pose_tensor(pose) + result = self._caster.trace_tensor( + self._ray_origins, + self._ray_directions, + self.ray_pattern.max_distance, + outputs=self._outputs, + min_distance=self.ray_pattern.min_distance, + ray_filter=self._ray_filter, + ) + except (AttributeError, KeyError, TypeError, ValueError, NotImplementedError) as exc: + raise type(exc)(f"RayScan tensor trace failed: {exc}") from exc + if self._observation is None: + self._distance.copy_(result.distance) + return self._distance + output_slice = 1 + self._observation[..., 0].copy_(result.distance) + if self._normal is not None: + self._normal.copy_(result.normal) + self._observation[..., output_slice : output_slice + 3].copy_(self._normal) + output_slice += 3 + if self._hit_point is not None: + self._hit_point.copy_(result.hit_point) + self._observation[..., output_slice : output_slice + 3].copy_(self._hit_point) + return self._observation.reshape(self._env.num_envs, -1) + + def _refresh_attachment_rays(self, pose: Any) -> None: + """Transform the local ray pattern by the selected attachment frame.""" + + assert self._attachment_body_index is not None + assert self._local_origins is not None + assert self._local_directions is not None + assert self._ray_origins is not None + assert self._ray_directions is not None + index = self._attachment_body_index + position = pose.pos_w[:, index, :].unsqueeze(1) + quaternion = pose.quat_w[:, index, :] + if self._attachment_site_local_pos is not None: + assert self._attachment_site_local_quat is not None + quaternion = _quat_multiply_wxyz( + quaternion, self._attachment_site_local_quat.expand_as(quaternion) + ) + position = position + _quat_rotate_wxyz( + pose.quat_w[:, index, :], self._attachment_site_local_pos.expand(1, 3) + )[:, 0, :].unsqueeze(1) + rotation = _attachment_rotation_wxyz(quaternion, self.ray_pattern.alignment) + self._ray_origins.copy_( + position + torch.einsum("bij,rj->bri", rotation, self._local_origins) + ) + self._ray_directions.copy_(torch.einsum("bij,rj->bri", rotation, self._local_directions)) + + def _resolve_pose_sources( + self, collision: Any + ) -> dict[str, tuple[tuple[str, ...], torch.Tensor]]: + """Map public full-scene body identities to entity-local pose reads.""" + + scene = collision.scene + body_names = tuple(scene.body_names) + body_ids = self._env.backend.get_body_ids(body_names) + collision_index_by_id = { + int(body_id): index for index, body_id in enumerate(body_ids.tolist()) + } + sources: dict[str, tuple[tuple[str, ...], torch.Tensor]] = {} + entities = getattr(self._env.scene, "entities", {}) + for entity_name, entity in entities.items(): + try: + local_names = tuple(entity.body_names) + except NotImplementedError: + continue + if not local_names: + continue + local_ids = self._env.backend.get_body_ids(local_names) + selected_names: list[str] = [] + destinations: list[int] = [] + for local_name, local_id in zip(local_names, local_ids.tolist(), strict=True): + collision_index = collision_index_by_id.get(int(local_id)) + if collision_index is None: + continue + selected_names.append(local_name) + destinations.append(collision_index) + if not selected_names: + continue + sources[entity_name] = ( + tuple(selected_names), + torch.tensor(destinations, dtype=torch.long, device=self._device), + ) + return sources + + def _validate_pose_coverage(self, collision: Any) -> None: + """Require entity reads for every non-root full-scene body.""" + + scene = collision.scene + source = self._pose_sources.get(self._entity_name) + attachment_name = self._attachment_body_name + assert attachment_name is not None + if source is None or attachment_name not in source[0]: + raise KeyError( + f"ray pattern attachment body {attachment_name!r} is not owned by " + f"entity {self._entity_name!r}" + ) + covered = { + body_name for local_names, _ in self._pose_sources.values() for body_name in local_names + } + uncovered = set(scene.body_names) - covered + root_name = scene.body_names[0] + if uncovered and ( + uncovered != {root_name} + or (collision.body_parent_ids is not None and int(collision.body_parent_ids[0]) != -1) + ): + raise NotImplementedError( + "RayScan full-scene pose aggregation requires every non-root body to " + "belong to a scene entity with a packed body read; uncovered bodies: " + f"{sorted(uncovered)}" + ) + + def _refresh_full_pose(self, read_plan: Any) -> RayPoseTensorView: + """Assemble the full public body-axis pose from entity-local read views.""" + + assert self._pose_body_names + assert self._pose_pos is not None + assert self._pose_quat is not None + self._pose_pos.zero_() + self._pose_quat.zero_() + self._pose_quat[..., 0] = 1.0 + for entity_name, (local_names, destinations) in self._pose_sources.items(): + source = read_plan.ray_pose_view(entity_name, local_names) + self._pose_pos[:, destinations, :] = source.pos_w + self._pose_quat[:, destinations, :] = source.quat_w + return RayPoseTensorView( + body_names=self._pose_body_names, + pos_w=self._pose_pos, + quat_w=self._pose_quat, + ) + + def _resolve_ray_filter(self, collision: Any) -> RayTraceFilter: + """Resolve public names/tags to the compact static filter on the cold path.""" + + scene = collision.scene + geom_names = tuple(collision.geom_names) + body_names = tuple(scene.body_names) + allowed_geom_ids = None + if self.ray_pattern.allowed_geom_names is not None: + unknown = set(self.ray_pattern.allowed_geom_names) - set(geom_names) + if unknown: + raise KeyError(f"ray pattern geom names are missing: {sorted(unknown)}") + allowed_geom_ids = tuple( + geom_names.index(name) for name in self.ray_pattern.allowed_geom_names + ) + excluded_body_ids = None + if self.ray_pattern.excluded_body_names is not None: + unknown = set(self.ray_pattern.excluded_body_names) - set(body_names) + if unknown: + raise KeyError(f"ray pattern body names are missing: {sorted(unknown)}") + excluded_body_ids = tuple( + body_names.index(name) for name in self.ray_pattern.excluded_body_names + ) + excluded_ancestor_body_ids = None + if self.ray_pattern.excluded_ancestor_body_names is not None: + unknown = set(self.ray_pattern.excluded_ancestor_body_names) - set(body_names) + if unknown: + raise KeyError(f"ray pattern ancestor body names are missing: {sorted(unknown)}") + excluded_ancestor_body_ids = tuple( + body_names.index(name) for name in self.ray_pattern.excluded_ancestor_body_names + ) + ray_filter = RayTraceFilter( + allowed_geom_ids=allowed_geom_ids, + allowed_tags=self.ray_pattern.allowed_tags, + excluded_body_ids=excluded_body_ids, + excluded_ancestor_body_ids=excluded_ancestor_body_ids, + ) + ray_filter.geom_mask(collision) + return ray_filter + + def close(self) -> None: + """Release caster resources before the scene read plan closes.""" + self._closed = True + caster, self._caster = self._caster, None + if caster is not None: + caster.close() + + def reset(self, env_ids: torch.Tensor | slice | None = None) -> None: + """Retain caster state; observations refresh from the next read phase.""" + del env_ids + if self._bound and self.ray_pattern.sync_geometry: + self._sync_geometry() + + def _sync_geometry(self) -> None: + """Publish backend-owned per-row geometry to the caster on reset paths.""" + + if self._caster is None: + return + state = self._env.backend.get_ray_geometry_state() + self._caster.update_geometry_tensor( + RayGeometryTensorView( + geom_sizes=torch.tensor(state.geom_sizes, dtype=torch.float32, device=self._device), + geom_local_pos=torch.tensor( + state.geom_local_pos, dtype=torch.float32, device=self._device + ), + geom_local_quat=torch.tensor( + state.geom_local_quat, dtype=torch.float32, device=self._device + ), + geom_enabled=torch.tensor(state.geom_enabled, device=self._device), + geom_mesh_ids=torch.tensor( + state.geom_mesh_ids, dtype=torch.int32, device=self._device + ), + ) + ) + + +__all__ = [ + "RayAlignment", + "RayCameraPatternCfg", + "RayGridPatternCfg", + "RayPatternCfg", + "RayRingPatternCfg", + "RayScan", +] + + +def _attachment_rotation_wxyz(quaternion: torch.Tensor, alignment: RayAlignment) -> torch.Tensor: + """Return one ``[batch, 3, 3]`` attachment rotation matrix.""" + + batch = quaternion.shape[0] + identity = torch.eye(3, dtype=quaternion.dtype, device=quaternion.device) + if alignment == "world": + return identity.unsqueeze(0).expand(batch, 3, 3) + + basis = torch.eye(3, dtype=quaternion.dtype, device=quaternion.device) + rotated_basis = _quat_rotate_wxyz(quaternion, basis) + if alignment == "body": + return rotated_basis.transpose(1, 2) + + forward = rotated_basis[:, 0, :].clone() + forward[:, 2] = 0.0 + forward_norm = torch.linalg.vector_norm(forward, dim=1) + singular = forward_norm < 1.0e-3 + left = rotated_basis[:, 1, :].clone() + left[:, 2] = 0.0 + left_norm = torch.linalg.vector_norm(left, dim=1).clamp(min=1.0e-6) + left = left / left_norm.unsqueeze(1) + fallback = torch.stack((left[:, 1], -left[:, 0], torch.zeros_like(left[:, 0])), dim=1) + forward = torch.where( + singular.unsqueeze(1), + fallback, + forward / forward_norm.clamp(min=1.0e-6).unsqueeze(1), + ) + up = torch.tensor((0.0, 0.0, 1.0), dtype=quaternion.dtype, device=quaternion.device).expand( + batch, 3 + ) + left = torch.stack((-forward[:, 1], forward[:, 0], torch.zeros_like(forward[:, 0])), dim=1) + return torch.stack((forward, left, up), dim=2) + + +def _quat_multiply_wxyz(first: torch.Tensor, second: torch.Tensor) -> torch.Tensor: + """Compose batched ``wxyz`` quaternions with matching leading dimensions.""" + + w1, x1, y1, z1 = first.unbind(dim=-1) + w2, x2, y2, z2 = second.unbind(dim=-1) + return torch.stack( + ( + w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2, + w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2, + w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2, + w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2, + ), + dim=-1, + ) + + +def _quat_rotate_wxyz(quaternion: torch.Tensor, vector: torch.Tensor) -> torch.Tensor: + """Rotate shared local vectors by full-batch public ``wxyz`` quaternions.""" + + qvec = quaternion[:, None, 1:] + scalar = quaternion[:, None, 0:1] + cross = 2.0 * torch.linalg.cross(qvec, vector[None, :, :]) + return vector[None, :, :] + scalar * cross + torch.linalg.cross(qvec, cross) + + +def _create_ray_caster(caster: str, **options: Any) -> Any: + """Lazily call the installed UniSim public caster factory.""" + import unisim + + factory = getattr(unisim, "create_ray_caster", None) + if not callable(factory): + raise NotImplementedError( + "installed unisim-core does not expose create_ray_caster; use the " + "ray-local optional dependency profile" + ) + return factory(caster, **options) diff --git a/src/unilab/managers/observation_manager.py b/src/unilab/managers/observation_manager.py index 37222b72f..fbae777fa 100644 --- a/src/unilab/managers/observation_manager.py +++ b/src/unilab/managers/observation_manager.py @@ -362,6 +362,23 @@ def reset(self, env_ids: torch.Tensor | slice | None = None) -> dict[str, float] mod.reset(reset_mask) return {} + def bind_read_phase(self) -> None: + """Bind class terms after the scene's packed tensor read plan exists.""" + for group_cfg in self._group_obs_class_term_cfgs.values(): + for term_cfg in group_cfg: + binder = getattr(term_cfg.func, "bind_read_phase", None) + if not callable(binder): + continue + binder() + + def close(self) -> None: + """Forward close to class terms that own external resources.""" + for group_cfg in self._group_obs_class_term_cfgs.values(): + for term_cfg in group_cfg: + closer = getattr(term_cfg.func, "close", None) + if callable(closer): + closer() + def _check_and_handle_nans( self, tensor: torch.Tensor, @@ -787,6 +804,16 @@ def _observations_to_tensor_boundary( return values raise TypeError("ObservationManager results must be torch.Tensor carriers") + def _deferred_term_width(self, term_cfg: ObservationTermCfg) -> int: + """Resolve the fixed width of a term that binds after the read phase.""" + width = getattr(term_cfg.func, "observation_width", None) + if isinstance(width, bool) or not isinstance(width, int) or width <= 0: + raise TypeError( + f"ObservationManager term '{term_cfg.func}' with deferred read binding " + "must declare a positive integer observation_width" + ) + return width + 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() @@ -878,30 +905,34 @@ def _prepare_terms(self) -> None: 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, torch.Tensor): - raise TypeError( - f"ObservationManager term '{group_name}/{term_name}' returned " - f"{type(initial_obs).__name__}, expected torch.Tensor." - ) - if initial_obs.dtype != torch.float32 or initial_obs.device != self._device: - raise TypeError( - f"ObservationManager term '{group_name}/{term_name}' must return " - f"float32 observations on {self._device}; got " - f"{initial_obs.dtype} on {initial_obs.device}." - ) - 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) + deferred_bind = callable(getattr(term_cfg.func, "bind_read_phase", None)) + if not deferred_bind: + initial_obs = term_cfg.func(self._env, **term_cfg.params) + if not isinstance(initial_obs, torch.Tensor): + raise TypeError( + f"ObservationManager term '{group_name}/{term_name}' returned " + f"{type(initial_obs).__name__}, expected torch.Tensor." + ) + if initial_obs.dtype != torch.float32 or initial_obs.device != self._device: + raise TypeError( + f"ObservationManager term '{group_name}/{term_name}' must return " + f"float32 observations on {self._device}; got " + f"{initial_obs.dtype} on {initial_obs.device}." + ) + 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) + else: + obs_dims = (self.num_envs, self._deferred_term_width(term_cfg)) if term_cfg.scale is not None: scale = torch.as_tensor(term_cfg.scale, dtype=torch.float32) self._scale_tensors[id(term_cfg)] = ( - scale.broadcast_to(initial_obs.shape).clone().to(device=self._device) + scale.broadcast_to(obs_dims).clone().to(device=self._device) ) if term_cfg.noise is not None and isinstance( diff --git a/src/unilab/managers/types.py b/src/unilab/managers/types.py index 55c0a8610..421bb5165 100644 --- a/src/unilab/managers/types.py +++ b/src/unilab/managers/types.py @@ -292,6 +292,9 @@ def action_manager(self) -> ManagerActionManager: ... @property def command_manager(self) -> ManagerCommandManager: ... + @property + def observation_manager(self) -> Any: ... + @property def event_manager(self) -> ManagerEventManager: ... diff --git a/tests/base/genesis_fake_runtime.py b/tests/base/genesis_fake_runtime.py index 64ee4e8e0..199efaa45 100644 --- a/tests/base/genesis_fake_runtime.py +++ b/tests/base/genesis_fake_runtime.py @@ -108,6 +108,7 @@ def __init__(self, morph): ] self.built_envs = 0 self.control_calls: list[np.ndarray] = [] + self.velocity_calls: list[tuple[np.ndarray, list[int], list[int]]] = [] self.step_count = 0 def get_link(self, name): @@ -206,12 +207,36 @@ def set_qpos( def set_dofs_velocity( self, velocity=None, dofs_idx_local=None, envs_idx=None, *, skip_forward=False ): - rows = self._rows(envs_idx) + # Genesis 1.3.3 accepts int64 rows with a selected-column payload. A + # bool mask means the payload is full-width even when columns narrow. + raw_rows = self._rows(envs_idx) + bool_mask = envs_idx is not None and np.asarray(envs_idx.dtype == torch.bool).item() + dofs = list(range(NV)) if dofs_idx_local is None else list(dofs_idx_local) value = torch.as_tensor(velocity, dtype=torch.float32).cpu() - if rows is None: - self.qvel = value.clone() + if raw_rows is None: + submitted = value + rows = None + elif bool_mask: + submitted = value + rows = list(np.flatnonzero(np.asarray(envs_idx.detach().cpu()))) + else: + rows = list(raw_rows) + submitted = value if tuple(value.shape) == (len(rows), len(dofs)) else value[-1:] + self.velocity_calls.append((submitted.numpy().copy(), dofs, [-1] if rows is None else rows)) + device = self.qvel.device + row_indices = ( + None if rows is None else torch.as_tensor(rows, dtype=torch.int64, device=device) + ) + columns = torch.as_tensor(dofs, dtype=torch.int64, device=device) + if row_indices is None: + staged = self.qvel.clone() else: - self.qvel[rows] = value.clone() + staged = self.qvel.index_select(0, row_indices) + staged_values = submitted.to(device=device, dtype=self.qvel.dtype) + if tuple(staged_values.shape) != tuple(staged[:, columns].shape): + staged_values = staged_values[-1:, :].expand(staged[:, columns].shape).clone() + staged = staged.index_copy(1, columns, staged_values) + self.qvel = staged if row_indices is None else self.qvel.index_copy(0, row_indices, staged) def control_dofs_position(self, position, dofs_idx_local=None, envs_idx=None): value = torch.as_tensor(position, dtype=torch.float32).cpu() @@ -259,11 +284,20 @@ def step(self): class _FakeRigidSolver: def __init__(self): - self.external_forces: list[tuple[np.ndarray, list[int]]] = [] + self.external_forces: list[tuple[np.ndarray, list[int], list[int] | None]] = [] + self.external_torques: list[tuple[np.ndarray, list[int], list[int] | None]] = [] def apply_links_external_force(self, force, links_idx=None, envs_idx=None, **kwargs): value = torch.as_tensor(force, dtype=torch.float32).cpu().numpy().copy() - self.external_forces.append((value, list(links_idx))) + self.external_forces.append( + (value, list(links_idx), None if envs_idx is None else list(envs_idx)) + ) + + def apply_links_external_torque(self, torque, links_idx=None, envs_idx=None, **kwargs): + value = torch.as_tensor(torque, dtype=torch.float32).cpu().numpy().copy() + self.external_torques.append( + (value, list(links_idx), None if envs_idx is None else list(envs_idx)) + ) class _FakeCamera: diff --git a/tests/base/test_cuda_backend_platform_preflight.py b/tests/base/test_cuda_backend_platform_preflight.py index 22835a84c..0b0ad854a 100644 --- a/tests/base/test_cuda_backend_platform_preflight.py +++ b/tests/base/test_cuda_backend_platform_preflight.py @@ -69,7 +69,10 @@ def test_host_bridge_backends_do_not_require_torch_cuda( ("newton", {"newton_device": "cpu"}, "valid", "not supported"), ("newton", {"newton_device": "mps"}, "valid", "not supported"), ("mjwarp", {}, "darwin", "macOS"), - ("genesis", {"genesis_device_id": 0}, "hip", "ROCm/HIP"), + ("mjwarp", {}, "hip", "ROCm/HIP"), + ("newton", {"newton_device": "cuda:0"}, "hip", "ROCm/HIP"), + ("isaacgym", {"isaacgym_device_id": 0}, "hip", "ROCm/HIP"), + ("isaacsim", {"isaacsim_device_id": 0}, "hip", "ROCm/HIP"), ("mjwarp", {}, "unavailable", "unavailable"), ("newton", {"newton_device": "cuda:2"}, "valid", "out of range"), ("newton", {"newton_device": "cuda:not-an-index"}, "valid", "not supported"), @@ -121,6 +124,25 @@ def test_cuda_only_backend_can_construct_with_valid_cuda_device( assert fake_cuda == [] +def test_genesis_backend_can_construct_with_valid_rocm_device( + fake_cuda: list[int], + backend_calls: list[dict[str, object]], + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(torch.version, "hip", "7.2", raising=False) + + backend_factory.create_backend( + "genesis", + _Scene(), + 1, + 0.005, + genesis_device_id=0, + ) + + assert len(backend_calls) == 1 + assert fake_cuda == [0] + + @pytest.mark.parametrize("learner_device", ["cpu", "mps"]) def test_cuda_only_process_devices_reject_cpu_learner(learner_device: str) -> None: with pytest.raises(ValueError, match="requires a CUDA process device"): diff --git a/tests/base/test_entity_tensor_read_plan.py b/tests/base/test_entity_tensor_read_plan.py index 550b87742..79b0ddce3 100644 --- a/tests/base/test_entity_tensor_read_plan.py +++ b/tests/base/test_entity_tensor_read_plan.py @@ -2,7 +2,8 @@ from __future__ import annotations -from types import MethodType +from dataclasses import dataclass, field +from types import MethodType, SimpleNamespace from typing import cast import numpy as np @@ -36,6 +37,7 @@ def __init__(self, *, device_resident: bool = False) -> None: self.sensor_view_calls = 0 self.aggregate_body_view_calls = 0 self.closes = 0 + self.io_specs: list[TensorIOSpec] = [] self.qpos = torch.tensor([[0.1], [0.3]], dtype=torch.float32) self.qvel = torch.tensor([[10.0], [20.0]], dtype=torch.float32) self.sensors = { @@ -49,6 +51,12 @@ def __init__(self, *, device_resident: bool = False) -> None: "track_linvel_w_hip": torch.zeros((2, 3), dtype=torch.float32), "track_angvel_w_hip": torch.ones((2, 3), dtype=torch.float32), } + self.body_blocks = { + "body_pos_w": self.sensors["track_pos_w_hip"].reshape(2, 3).contiguous(), + "body_quat_w": self.sensors["track_quat_w_hip"].reshape(2, 4).contiguous(), + "body_lin_vel_w": self.sensors["track_linvel_w_hip"].reshape(2, 3).contiguous(), + "body_ang_vel_w": self.sensors["track_angvel_w_hip"].reshape(2, 3).contiguous(), + } self._selected_rows = torch.empty((0,), dtype=torch.int64) def get_tensor_capabilities(self) -> TensorLifecycleCapabilities: @@ -97,6 +105,7 @@ def get_sensor_view(self, name: str, device: str | torch.device = "cpu") -> torc def get_tracked_body_views(self, body_names=None, device=None) -> TrackedBodyStateViews: self.aggregate_body_view_calls += 1 names = ("hip",) if body_names is None else tuple(body_names) + names = tuple(name.rsplit("/", 1)[-1] for name in names) if names != ("hip",): raise ValueError(f"unknown tracked bodies: {names}") target = torch.device(device) if device is not None else torch.device("cpu") @@ -183,6 +192,8 @@ def get_joint_dof_vel_indices(self, names) -> np.ndarray: def bind_host_bridge_io(self, spec: TensorIOSpec): backend = self + if getattr(spec, "body_names", ()): + backend.io_specs.append(spec) class _Plan(HostBridgeTransferPlan): def __init__(self, plan_spec: TensorIOSpec) -> None: @@ -208,6 +219,13 @@ def read_selected_state_sensors(self): backend.selected_reads += 1 return backend._packet() + def get_body_pose_views(self): + return SimpleNamespace( + body_names=("robot/hip",), + pos_w=backend.body_blocks["body_pos_w"], + quat_w=backend.body_blocks["body_quat_w"], + ) + def write_control(self, ctrl) -> None: return None @@ -226,6 +244,8 @@ def close(self) -> None: def _packet(self) -> dict[str, torch.Tensor]: packet = self.get_state_views(("qpos", "qvel"), device="cpu") packet.update(self.sensors) + if self.device_resident: + return packet return packet @@ -242,6 +262,30 @@ def _scene(backend: _SceneTensorBackend) -> EntityScene: ) +def _use_body_capable_tensor_spec(monkeypatch: pytest.MonkeyPatch) -> None: + """Expose a UniSim-2.2-shaped public TensorIOSpec for packed-block tests.""" + + @dataclass(frozen=True) + class _BodyCapableTensorIOSpec(TensorIOSpec): + body_names: tuple[str, ...] = () + body_fields: tuple[str, ...] = () + device: object | None = field(default=None) + + monkeypatch.setattr("unilab.base.entity.TensorIOSpec", _BodyCapableTensorIOSpec) + + +def _use_legacy_tensor_spec(monkeypatch: pytest.MonkeyPatch) -> None: + """Expose a pre-2.2 TensorIOSpec without packed body-block parameters.""" + + @dataclass(frozen=True) + class _LegacyTensorIOSpec: + state_fields: tuple[str, ...] + sensor_names: tuple[str, ...] + device: object | None = None + + monkeypatch.setattr("unilab.base.entity.TensorIOSpec", _LegacyTensorIOSpec) + + def _specs() -> tuple[SceneTensorReadSpec, ...]: return ( SceneTensorReadSpec(entity="robot", sensor_names=("imu_gyro",)), @@ -249,8 +293,11 @@ def _specs() -> tuple[SceneTensorReadSpec, ...]: ) -def test_host_bridge_scene_plan_packs_all_reads_into_one_phase_transfer() -> None: +def test_host_bridge_scene_plan_packs_all_reads_into_one_phase_transfer( + monkeypatch: pytest.MonkeyPatch, +) -> None: backend = _SceneTensorBackend() + _use_body_capable_tensor_spec(monkeypatch) scene = _scene(backend) plan = scene.compile_tensor_reads("cpu", _specs()) @@ -276,6 +323,32 @@ def test_host_bridge_scene_plan_packs_all_reads_into_one_phase_transfer() -> Non } plan.close() + + +def test_scene_plan_pose_reads_fall_back_to_canonical_sensors_on_older_host_contracts( + monkeypatch: pytest.MonkeyPatch, +) -> None: + backend = _SceneTensorBackend() + scene = _scene(backend) + _use_legacy_tensor_spec(monkeypatch) + plan = scene.compile_tensor_reads( + "cpu", + ( + SceneTensorReadSpec( + entity="robot", body_names=("hip",), body_fields=("pos_w", "quat_w") + ), + ), + ) + + assert plan.packet_body_blocks is False + assert plan.packet_names == ("track_pos_w_hip", "track_quat_w_hip") + assert backend.io_specs == [] + plan.refresh() + pose = plan.ray_pose_view("robot") + assert pose.body_names == ("hip",) + torch.testing.assert_close(pose.pos_w[:, 0], backend.sensors["track_pos_w_hip"]) + torch.testing.assert_close(pose.quat_w[:, 0], backend.sensors["track_quat_w_hip"]) + plan.close() assert backend.closes == 1 @@ -331,7 +404,7 @@ def test_device_resident_scene_plan_refreshes_every_sensor_boundary() -> None: torch.testing.assert_close(second, torch.full((2, 3), 7.0)) # Every compiled sensor crosses its public backend boundary on every phase. # Genesis/Newton derived sensors are stable outputs, not native-live arrays. - assert backend.sensor_view_calls == calls_after_first_refresh + len(plan.packet_names) + assert backend.sensor_view_calls == calls_after_first_refresh + len(plan.sensor_names["robot"]) assert plan.transfer_stats == {} assert plan.last_timing == {} plan.close() @@ -350,7 +423,9 @@ def test_device_resident_selected_reset_refreshes_body_sensor_boundary() -> None body = plan.body_tensor_view("robot") torch.testing.assert_close(body.pos_w[:, 0], backend.sensors["track_pos_w_hip"]) - assert backend.sensor_view_calls == calls_after_selected_refresh + len(plan.packet_names) + assert backend.sensor_view_calls == calls_after_selected_refresh + len( + plan.sensor_names["robot"] + ) plan.close() @@ -385,6 +460,105 @@ def counting_qpos_indices(names): plan.close() +def test_scene_plan_pose_reads_use_packed_pose_blocks_without_velocities( + monkeypatch: pytest.MonkeyPatch, +) -> None: + backend = _SceneTensorBackend() + _use_body_capable_tensor_spec(monkeypatch) + scene = _scene(backend) + plan = scene.compile_tensor_reads( + "cpu", + ( + SceneTensorReadSpec( + entity="robot", body_names=("hip",), body_fields=("pos_w", "quat_w") + ), + ), + ) + + assert plan.body_fields == {"robot": ("pos_w", "quat_w")} + assert plan.packet_names == ("body_pos_w", "body_quat_w") + if backend.io_specs: + assert backend.io_specs[0].body_names == ("hip",) + assert backend.io_specs[0].body_fields == ("pos_w", "quat_w") + assert plan.packet_body_axis("body_pos_w") == ("hip",) + plan.refresh() + + pose = plan.ray_pose_view("robot") + assert pose.body_names == ("hip",) + assert pose.pos_w.shape == (2, 1, 3) + assert pose.quat_w.shape == (2, 1, 4) + torch.testing.assert_close(pose.pos_w[:, 0], backend.sensors["track_pos_w_hip"]) + torch.testing.assert_close(pose.quat_w[:, 0], backend.sensors["track_quat_w_hip"]) + body = plan.body_tensor_view("robot") + assert body.pos_w is not None and body.quat_w is not None + assert body.lin_vel_w is None and body.ang_vel_w is None + plan.close() + + +def test_scene_plan_ray_pose_view_preserves_row_strided_packet_slices( + monkeypatch: pytest.MonkeyPatch, +) -> None: + backend = _SceneTensorBackend() + _use_body_capable_tensor_spec(monkeypatch) + scene = _scene(backend) + plan = scene.compile_tensor_reads( + "cpu", + ( + SceneTensorReadSpec( + entity="robot", body_names=("hip",), body_fields=("pos_w", "quat_w") + ), + ), + ) + + storage = torch.zeros((2, 12), dtype=torch.float32) + pos = torch.arange(6, dtype=torch.float32).reshape(2, 3) + quat = torch.tensor(((1.0, 0.0, 0.0, 0.0), (1.0, 0.0, 0.0, 0.0)), dtype=torch.float32) + storage[:, 3:6] = pos + storage[:, 7:11] = quat + backend.body_blocks["body_pos_w"] = storage[:, 3:6] + backend.body_blocks["body_quat_w"] = storage[:, 7:11] + plan.refresh() + pose = plan.ray_pose_view("robot") + + assert not pose.pos_w.is_contiguous() + assert not pose.quat_w.is_contiguous() + assert pose.pos_w.data_ptr() == storage[:, 3:6].data_ptr() + assert pose.quat_w.data_ptr() == storage[:, 7:11].data_ptr() + assert pose.pos_w.stride() == (12, 3, 1) + assert pose.quat_w.stride() == (12, 4, 1) + storage[0, 3] = 42.0 + storage[1, 7] = -1.0 + assert pose.pos_w[0, 0, 0].item() == 42.0 + assert pose.quat_w[1, 0, 0].item() == -1.0 + plan.close() + + +def test_scene_plan_validates_body_field_requests_and_pose_reads( + monkeypatch: pytest.MonkeyPatch, +) -> None: + backend = _SceneTensorBackend() + _use_body_capable_tensor_spec(monkeypatch) + scene = _scene(backend) + + with pytest.raises(ValueError, match="body_fields are unknown"): + SceneTensorReadSpec(entity="robot", body_names=("hip",), body_fields=("pos_w", "bogus")) + with pytest.raises(ValueError, match="body_fields must be unique"): + SceneTensorReadSpec(entity="robot", body_names=("hip",), body_fields=("pos_w", "pos_w")) + with pytest.raises(ValueError, match="body_fields cannot be empty"): + SceneTensorReadSpec(entity="robot", body_names=("hip",), body_fields=()) + + plan = scene.compile_tensor_reads( + "cpu", + (SceneTensorReadSpec(entity="robot", body_names=("hip",), body_fields=("pos_w",)),), + ) + with pytest.raises(RuntimeError, match="must be refreshed"): + plan.ray_pose_view("robot") + plan.refresh() + with pytest.raises(ValueError, match=r"fields \['quat_w'\] were not compiled"): + plan.ray_pose_view("robot") + plan.close() + + def test_scene_plan_rejects_reads_before_refresh_and_after_close() -> None: backend = _SceneTensorBackend() plan = _scene(backend).compile_tensor_reads("cpu", _specs()) diff --git a/tests/base/test_genesis_backend.py b/tests/base/test_genesis_backend.py index 1d0563221..24584fda8 100644 --- a/tests/base/test_genesis_backend.py +++ b/tests/base/test_genesis_backend.py @@ -31,8 +31,10 @@ ) from unisim.backend.genesis.materialization import preserve_torch_globals from unisim.dr.types import ( + INTERVAL_TERM_BODY_ANGULAR_VELOCITY_DELTA, INTERVAL_TERM_BODY_FORCE, INTERVAL_TERM_BODY_LINEAR_VELOCITY_DELTA, + INTERVAL_TERM_BODY_TORQUE, INTERVAL_TERM_PUSH, IntervalRandomizationPlan, IntervalTermOp, @@ -115,20 +117,50 @@ def fake_genesis(monkeypatch: pytest.MonkeyPatch): if not torch.cuda.is_available(): factories = { name: getattr(torch, name) - for name in ("arange", "empty", "full", "ones", "tensor", "zeros") + for name in ( + "arange", + "asarray", + "as_tensor", + "empty", + "empty_like", + "full", + "full_like", + "ones", + "ones_like", + "tensor", + "zeros", + "zeros_like", + ) } def cuda_free_factory(original): def factory(*args, **kwargs): device = kwargs.get("device") - if device is not None and torch.device(device).type == "cuda": + requested_cuda = device is not None and torch.device(device).type == "cuda" + if requested_cuda: kwargs["device"] = "cpu" - return original(*args, **kwargs) + result = original(*args, **kwargs) + source = next((value for value in args if isinstance(value, torch.Tensor)), None) + if requested_cuda or getattr(source, "_fake_genesis_cuda", False): + result._fake_genesis_cuda = True + return result return factory for name, original in factories.items(): monkeypatch.setattr(torch, name, cuda_free_factory(original)) + + class FakeCudaEvent: + def record(self, stream: object) -> None: + del stream + + def synchronize(self) -> None: + return None + + def wait(self, stream: object) -> None: + del stream + + monkeypatch.setattr(torch.cuda, "Event", lambda *args, **kwargs: FakeCudaEvent()) monkeypatch.setattr( torch.cuda, "current_stream", @@ -144,6 +176,39 @@ def factory(*args, **kwargs): and bool(getattr(self._gs, "use_zerocopy", False)) ), ) + monkeypatch.setattr( + GenesisBackend, + "_device_dr_supported", + lambda self: self._tensor_lifecycle_supported(), + ) + original_interval_operand_plane = GenesisBackend._interval_operand_plane + + def fake_interval_operand_plane(self, op): + operands = tuple( + value + for value in (op.payload, op.body_ids) + if value is not None and isinstance(value, torch.Tensor) + ) + planes = { + "device" + if getattr(value, "_fake_genesis_cuda", False) or value.device.type != "cpu" + else "host" + for value in operands + } + if len(planes) > 1: + raise ValueError( + "genesis runtime interval randomization operands must share one " + f"data plane; term {op.term!r} mixes host and device tensors" + ) + if planes: + return next(iter(planes)) + return original_interval_operand_plane(self, op) + + monkeypatch.setattr( + GenesisBackend, + "_interval_operand_plane", + fake_interval_operand_plane, + ) def tensor_device(self, requested=None): target = self._device if requested is None else torch.device(requested) @@ -155,6 +220,9 @@ def tensor_device(self, requested=None): fake = make_fake_genesis() fake.backend = fake.gpu fake.cuda = fake.gpu + fake.amdgpu = fake.gpu + if torch.version.hip is not None: + fake.backend = fake.amdgpu fake.use_zerocopy = True original_add_entity = fake.Scene.add_entity device_getters = ( @@ -527,10 +595,34 @@ def test_dr_capabilities_and_tensor_reset_randomization_boundary( ) -> None: backend = _backend(tiny_model_file) caps = backend.get_dr_capabilities() - assert caps.supported_reset_terms == {"body_mass", "base_mass_delta", "kp", "kd"} - assert caps.supports_interval_term(INTERVAL_TERM_BODY_FORCE) is True + assert caps.supported_reset_terms == frozenset() + assert caps.supported_interval_terms == frozenset() + assert caps.supported_device_reset_terms == { + "base_com_offset", + "base_mass_delta", + "body_ipos", + "body_mass", + "kp", + "kd", + } + assert caps.supported_device_interval_terms == { + "body_force", + "body_torque", + "body_linear_velocity_delta", + "body_angular_velocity_delta", + } + assert caps.supports_interval_term(INTERVAL_TERM_BODY_FORCE, plane="device") is True + assert caps.supports_interval_term(INTERVAL_TERM_BODY_TORQUE, plane="device") is True assert caps.supports_interval_term(INTERVAL_TERM_PUSH) is False - assert caps.supports_interval_term(INTERVAL_TERM_BODY_LINEAR_VELOCITY_DELTA) is False + assert caps.supports_interval_term(INTERVAL_TERM_PUSH, plane="device") is False + assert ( + caps.supports_interval_term(INTERVAL_TERM_BODY_LINEAR_VELOCITY_DELTA, plane="device") + is True + ) + assert ( + caps.supports_interval_term(INTERVAL_TERM_BODY_ANGULAR_VELOCITY_DELTA, plane="device") + is True + ) rows = torch.tensor([1, 3], dtype=torch.int64, device="cuda") qpos, qvel = _stand_state(backend, 2) @@ -566,29 +658,102 @@ def test_dr_capabilities_and_tensor_reset_randomization_boundary( def test_interval_randomization_and_body_force(fake_genesis, tiny_model_file: str) -> None: backend = _backend(tiny_model_file) solver = backend._scene.sim.rigid_solver - force = torch.ones((4, 2, 3), dtype=torch.float32) + device = torch.device("cuda") if not torch.cuda.is_available() else backend._device + force = torch.ones((4, 2, 3), dtype=torch.float32, device=device) + torque = torch.full_like(force, 0.25) backend.apply_interval_randomization( IntervalRandomizationPlan( ops=( IntervalTermOp( INTERVAL_TERM_BODY_FORCE, force, - body_ids=torch.as_tensor([1, 3], dtype=torch.int32), + body_ids=torch.as_tensor([1, 3], dtype=torch.int64, device=device), ), - ) + IntervalTermOp( + INTERVAL_TERM_BODY_TORQUE, + torque, + body_ids=torch.as_tensor([1, 3], dtype=torch.int64, device=device), + ), + ), ) ) - assert len(solver.external_forces) == 2 - np.testing.assert_allclose(solver.external_forces[0][0], np.ones((4, 3))) - assert solver.external_forces[0][1] == [1] - assert solver.external_forces[1][1] == [3] + # The single-entity lane submits one solver force call per selected body; + # torque rides the same API as a zero force plus the torque payload. + assert len(solver.external_forces) == 4 + assert len(solver.external_torques) == 0 + for offset, force_call in enumerate(solver.external_forces): + expected_magnitude = 1.0 if offset < 2 else 0.0 + np.testing.assert_allclose(force_call[0], np.full((4, 3), expected_magnitude)) + assert force_call[1] == [1 if offset % 2 == 0 else 3] + assert force_call[2] == [0, 1, 2, 3] - with pytest.raises(NotImplementedError, match="does not support interval term 'push'"): + with pytest.raises( + NotImplementedError, + match="genesis runtime interval randomization does not support host-plane term: body_force", + ): backend.apply_interval_randomization( - IntervalRandomizationPlan(ops=(IntervalTermOp(INTERVAL_TERM_PUSH, torch.ones((1, 3))),)) + IntervalRandomizationPlan( + ops=( + IntervalTermOp( + INTERVAL_TERM_BODY_FORCE, + torch.zeros((4, 1, 3), dtype=torch.float32), + body_ids=torch.zeros((1,), dtype=torch.int64), + ), + ) + ) ) + + with pytest.raises(NotImplementedError, match="does not support device-plane term: push"): + backend.apply_interval_randomization( + IntervalRandomizationPlan( + ops=(IntervalTermOp(INTERVAL_TERM_PUSH, torch.ones((1, 3), device=device)),) + ) + ) + linear = torch.zeros((4, 1, 3), dtype=torch.float32, device=device) + linear[1, 0] = torch.asarray((0.25, -0.5, 0.75), device=device) + angular = torch.zeros_like(linear) + angular[1, 0] = torch.asarray((-0.1, 0.2, -0.3), device=device) + entity = backend._scene.entities[0] + selected_linear = torch.zeros((entity.built_envs, 3), dtype=torch.float32) + selected_linear[1] = linear[1, 0].cpu() + selected_angular = torch.zeros_like(selected_linear) + selected_angular[1] = angular[1, 0].cpu() + velocity_before = entity.qvel.clone() + qpos_before = entity.qpos.clone() + linear_ids = torch.as_tensor([1], dtype=torch.int64, device=device) + backend.apply_interval_randomization( + IntervalRandomizationPlan( + ops=( + IntervalTermOp( + INTERVAL_TERM_BODY_LINEAR_VELOCITY_DELTA, + linear, + body_ids=linear_ids, + ), + IntervalTermOp( + INTERVAL_TERM_BODY_ANGULAR_VELOCITY_DELTA, + angular, + body_ids=linear_ids, + ), + ) + ) + ) + # Genesis submits selected rows through column-scoped public state; qpos, + # unrelated rows, and pending wrenches remain untouched. + torch.testing.assert_close(entity.qpos, qpos_before) + torch.testing.assert_close(entity.qvel[0], velocity_before[0]) + torch.testing.assert_close(entity.qvel[2:], velocity_before[2:]) + torch.testing.assert_close(entity.qvel[1, 0:3], velocity_before[1, 0:3] + selected_linear[1]) + torch.testing.assert_close(entity.qvel[1, 3:6], velocity_before[1, 3:6] + selected_angular[1]) + assert len(entity.velocity_calls) == 2 + np.testing.assert_allclose(entity.velocity_calls[0][0], selected_linear[1:2]) + assert entity.velocity_calls[0][1] == [0, 1, 2] + np.testing.assert_allclose(entity.velocity_calls[1][0], selected_angular[1:2]) + assert entity.velocity_calls[1][1] == [3, 4, 5] + for velocity_call in entity.velocity_calls: + assert velocity_call[2] == [1] with pytest.raises( - NotImplementedError, match="does not support interval term 'body_linear_velocity_delta'" + ValueError, + match="term 'body_linear_velocity_delta' mixes host and device tensors", ): backend.apply_interval_randomization( IntervalRandomizationPlan( @@ -596,17 +761,19 @@ def test_interval_randomization_and_body_force(fake_genesis, tiny_model_file: st IntervalTermOp( INTERVAL_TERM_BODY_LINEAR_VELOCITY_DELTA, torch.zeros((4, 1, 3), dtype=torch.float32), - body_ids=torch.as_tensor([1], dtype=torch.int32), + body_ids=torch.as_tensor([1], dtype=torch.int64, device=device), ), ) ) ) - with pytest.raises(ValueError, match="requires body_ids"): + with pytest.raises(ValueError, match="body force must have shape"): backend.apply_interval_randomization( IntervalRandomizationPlan( ops=( IntervalTermOp( - INTERVAL_TERM_BODY_FORCE, torch.zeros((4, 1, 3), dtype=torch.float32) + INTERVAL_TERM_BODY_FORCE, + torch.zeros((4, 2, 3), dtype=torch.float32, device=device), + body_ids=linear_ids, ), ) ) diff --git a/tests/base/test_torch_env.py b/tests/base/test_torch_env.py index 42b469203..14dae3169 100644 --- a/tests/base/test_torch_env.py +++ b/tests/base/test_torch_env.py @@ -654,10 +654,13 @@ def test_produces_completion_receipts_follows_backend_step_receipt_contract() -> # Synchronous host-bridge producers (mujoco/motrix/superdex) emit no # receipt, so the env declares the relaxed sync semantic to unilab.rl. assert _StubTorchEnv().produces_completion_receipts is False - # The event-backed async mjwarp CUDA path publishes a receipt per step. + # The event-backed async Genesis and MJWarp CUDA paths publish a receipt + # on every step. backend = _backend() backend.backend_type = "mjwarp" assert _StubTorchEnv(backend=backend).produces_completion_receipts is True + backend.backend_type = "genesis" + assert _StubTorchEnv(backend=backend).produces_completion_receipts is True def test_backend_reset_timing_hook_is_merged() -> None: diff --git a/tests/config/test_g1_motion_tracking_owners.py b/tests/config/test_g1_motion_tracking_owners.py index 095181904..d25d288b6 100644 --- a/tests/config/test_g1_motion_tracking_owners.py +++ b/tests/config/test_g1_motion_tracking_owners.py @@ -172,12 +172,27 @@ def test_sac_g1_motion_tracking_genesis_inherits_mujoco_parity() -> None: assert cfg.training.play_render_mode == "auto" assert cfg.env.genesis_device_id == 0 assert cfg.env.genesis_integrator == "implicitfast" - # Genesis legacy scenes declare no model-field reset terms and no interval - # velocity delta; only the host-side encoder_bias bias stays enabled. - assert set(cfg.env.events) == {"base_com", "encoder_bias", "foot_friction", "push_robot"} + # Issue #2521 enables the FlashMimic-compatible device-resident subset: + # startup mass, inherited COM/PD reset writes, and the six-axis kick. + # Genesis geom friction remains intentionally fail-closed. + assert set(cfg.env.events) == { + "base_com", + "body_mass", + "encoder_bias", + "foot_friction", + "pd_gains", + "push_robot", + } assert cfg.env.events.foot_friction is None - assert cfg.env.events.base_com is None - assert cfg.env.events.push_robot is None + assert cfg.env.events.body_mass.mode == "startup" + assert cfg.env.events.body_mass.params.asset_cfg.body_names == ["pelvis", "torso_link"] + assert cfg.env.events.body_mass.params.mass_distribution_params == [0.9, 1.1] + assert cfg.env.events.body_mass.params.recompute_inertia is False + assert cfg.env.events.pd_gains.mode == "reset" + assert cfg.env.events.pd_gains.params.kp_range == [0.9, 1.1] + assert cfg.env.events.pd_gains.params.kd_range == [0.9, 1.1] + assert cfg.env.events.push_robot.params.velocity_range.roll == [-0.52, 0.52] + assert cfg.env.events.push_robot.params.velocity_range.yaw == [-0.78, 0.78] assert cfg.env.events.encoder_bias is not None assert cfg.env.scene.entities.robot.geom_names is None # Algo block inherits the MuJoCo owner verbatim (DENYLIST parity). diff --git a/tests/envs/locomotion/g1/test_g1_owner_contract.py b/tests/envs/locomotion/g1/test_g1_owner_contract.py index 48317a2ab..6a6830507 100644 --- a/tests/envs/locomotion/g1/test_g1_owner_contract.py +++ b/tests/envs/locomotion/g1/test_g1_owner_contract.py @@ -174,7 +174,8 @@ 1.0, "scene_flat.xml", _OFFPOLICY_REWARDS, - _RESET_EVENTS, + # Genesis commits reset randomization through its device-plane packed lane. + (*_RESET_EVENTS, "pd_gains"), True, id="sac-genesis", ), diff --git a/tests/envs/mdp/test_events.py b/tests/envs/mdp/test_events.py index 182f58aac..117fa65cf 100644 --- a/tests/envs/mdp/test_events.py +++ b/tests/envs/mdp/test_events.py @@ -1107,6 +1107,38 @@ def test_apply_body_impulse_offset_generates_cross_torque() -> None: ) +@requires_cuda +def test_apply_body_impulse_cuda_payload_stays_on_declared_device_plane() -> None: + env, backend, _ = _tensor_dr_env(device_reset_randomization=True, device=torch.device("cuda")) + backend.interval_force_supported = True + backend.interval_torque_supported = True + manager = EventManager( + { + "impulse": EventTermCfg( + func=mdp.apply_body_impulse, + mode="step", + params={ + "force_range": (1.0, 1.0), + "torque_range": (0.25, 0.25), + "duration_s": (1.0, 1.0), + "cooldown_s": (0.0, 0.0), + }, + ) + }, + env, + ) + + manager.apply(mode="step", dt=0.02) + + assert len(backend.interval_plans) == 1 + plan = backend.interval_plans[0] + for term in (INTERVAL_TERM_BODY_FORCE, INTERVAL_TERM_BODY_TORQUE): + op = _interval_op(plan, term) + assert op is not None + assert op.payload.device.type == "cuda", term + assert op.body_ids.device.type == "cuda", term + + def test_apply_body_impulse_reset_clears_active_wrench() -> None: env, backend, _ = _transaction_env(rng_seed=43, interval_force_supported=True) manager = EventManager( @@ -1292,6 +1324,37 @@ def test_velocity_push_dispatches_angular_delta_when_supported() -> None: ) +@requires_cuda +def test_velocity_push_cuda_root_body_ids_match_payload_plane() -> None: + """Root interval IDs must match the sampled payload plane exactly.""" + + env, backend, _ = _tensor_dr_env( + device_reset_randomization=True, + device=torch.device("cuda"), + ) + backend.interval_velocity_supported = True + 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)}}, + ) + }, + env, + ) + manager._interval_term_time_left[0][:] = 0.0 + + manager.apply(mode="interval", dt=0.1) + + assert len(backend.interval_plans) == 1 + velocity_op = _interval_op(backend.interval_plans[0], INTERVAL_TERM_BODY_LINEAR_VELOCITY_DELTA) + assert velocity_op is not None + assert velocity_op.payload.device == torch.device("cuda", torch.cuda.current_device()) + assert velocity_op.body_ids.device == velocity_op.payload.device + + @pytest.mark.parametrize( ( "velocity_supported", @@ -1673,11 +1736,17 @@ def __init__( self.device_reset_randomization = device_reset_randomization self.device_plane = device_plane self.tensor_reset_calls: list[tuple[Any, Any, Any, Any]] = [] + view_device = torch.device("cuda" if device_plane else "cpu") + self.view_device = view_device self.body_tensor_views = { - "pos_w": torch.zeros((self.num_envs, 1, 3), dtype=torch.float32), - "quat_w": torch.zeros((self.num_envs, 1, 4), dtype=torch.float32), - "lin_vel_w": torch.zeros((self.num_envs, 1, 3), dtype=torch.float32), - "ang_vel_w": torch.zeros((self.num_envs, 1, 3), dtype=torch.float32), + "pos_w": torch.zeros((self.num_envs, 1, 3), dtype=torch.float32, device=view_device), + "quat_w": torch.zeros((self.num_envs, 1, 4), dtype=torch.float32, device=view_device), + "lin_vel_w": torch.zeros( + (self.num_envs, 1, 3), dtype=torch.float32, device=view_device + ), + "ang_vel_w": torch.zeros( + (self.num_envs, 1, 3), dtype=torch.float32, device=view_device + ), } self.body_tensor_views["quat_w"][..., 0] = 1.0 @@ -1713,6 +1782,7 @@ def get_dr_capabilities(self) -> DomainRandomizationCapabilities: supported_reset_terms=capabilities.supported_reset_terms, supported_interval_terms=capabilities.supported_interval_terms, supported_device_reset_terms=capabilities.supported_reset_terms, + supported_device_interval_terms=capabilities.supported_interval_terms, ) def set_state_tensor(self, env_indices, qpos, qvel, randomization=None) -> dict: @@ -1724,9 +1794,10 @@ def get_tracked_body_views(self, body_names=None, device=None): return SimpleNamespace(body_names=tuple(body_names), **self.body_tensor_views) def get_state_views(self, names, device=None): + target = self.view_device if device is None else torch.device(device) return { - "qpos": torch.zeros((self.num_envs, 3), dtype=torch.float32), - "qvel": torch.zeros((self.num_envs, 3), dtype=torch.float32), + "qpos": torch.zeros((self.num_envs, 3), dtype=torch.float32, device=target), + "qvel": torch.zeros((self.num_envs, 3), dtype=torch.float32, device=target), } def get_sensor_view(self, name, device="cpu"): diff --git a/tests/envs/mdp/test_ray_observation.py b/tests/envs/mdp/test_ray_observation.py new file mode 100644 index 000000000..6bcd08fe0 --- /dev/null +++ b/tests/envs/mdp/test_ray_observation.py @@ -0,0 +1,851 @@ +"""Focused tests for the tensor-native ray observation integration.""" + +from __future__ import annotations + +import sys +from dataclasses import dataclass, replace +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import numpy as np +import pytest +import torch +from unisim.errors import UnsupportedCapabilityError +from unisim.ray_query import ( + RayCollisionDescription, + RayGeometryState, + RaySceneDescription, + RaySiteFrameDescription, + RayTraceOutputs, +) + +from unilab.envs.mdp.ray_observation import ( + RayCameraPatternCfg, + RayGridPatternCfg, + RayPatternCfg, + RayRingPatternCfg, + RayScan, +) +from unilab.managers.manager_base import ManagerTermBaseCfg + +_REPO_ROOT = Path(__file__).resolve().parents[3] + + +class _Block: + """Minimal RayTraceTensorResult-shaped caster output.""" + + def __init__(self, values: torch.Tensor): + self.distance = values + self.normal = torch.zeros(*values.shape, 3, dtype=torch.float32) + self.hit_point = torch.ones(*values.shape, 3, dtype=torch.float32) + + +@dataclass +class _Capabilities: + supports_tensor_pose_sync: bool = True + supports_tensor_rays: bool = True + supports_device_output: bool = True + stream_event_ownership: str | None = "caller" + torch_devices: tuple[str, ...] = ("cpu",) + supports_min_distance: bool = True + supports_static_filter: bool = True + supports_tensor_normal: bool = True + supports_hit_point: bool = True + supports_geometry_sync: bool = True + supports_tensor_geometry_sync: bool = True + + +class _FakeCaster: + caster_type = "fake-ray" + + def __init__(self, *, num_rays: int = 4, values: torch.Tensor | None = None): + self.values = values + self.expected_num_rays = num_rays + self.scene: Any = None + self.materialized = False + self.layout: tuple[str, ...] | None = None + self.poses: list[Any] = [] + self.traces = 0 + self.geometry_updates: list[Any] = [] + self.last_query: tuple[float, object, RayTraceOutputs] | None = None + self.closed = 0 + + @property + def num_envs(self) -> int: + return 2 + + @property + def num_rays(self) -> int: + return self.expected_num_rays + + def get_ray_capabilities(self) -> _Capabilities: + return _Capabilities() + + def materialize(self, scene: Any) -> None: + assert scene is self.scene + self.materialized = True + + def bind_pose_layout(self, body_names) -> None: + self.layout = tuple(body_names) + + def update_pose_tensor(self, pose: Any) -> None: + self.poses.append(pose) + + def update_geometry_tensor(self, geometry: Any) -> None: + self.geometry_updates.append(geometry) + + def trace_tensor( + self, + origins, + directions, + max_distance, + env_ids=None, + outputs=None, + *, + min_distance=0.0, + ray_filter=None, + ): + self.traces += 1 + self.last_inputs = (origins, directions, max_distance) + self.last_query = (min_distance, ray_filter, outputs) + if origins.shape != directions.shape: + raise ValueError( + f"origin and direction tensor batches must have the same shape: " + f"{tuple(origins.shape)} != {tuple(directions.shape)}" + ) + if origins.ndim not in (2, 3) or origins.shape[-1] != 3: + raise ValueError( + "ray tensors must have shape [num_rays, 3] or " + f"[num_envs, num_rays, 3]: {tuple(origins.shape)}" + ) + if origins.ndim == 2 and origins.shape[0] != self.num_rays: + raise ValueError(f"expected {self.num_rays} rays, got {origins.shape[0]}") + if origins.ndim == 3 and origins.shape[:2] != (self.num_envs, self.num_rays): + raise ValueError( + f"expected {self.num_envs} x {self.num_rays} rays, got {tuple(origins.shape)}" + ) + if not origins.is_contiguous() or not directions.is_contiguous(): + raise ValueError("shared ray tensors must be contiguous") + values = self.values + if values is None: + values = torch.arange(2 * self.expected_num_rays, dtype=torch.float32).reshape( + 2, self.expected_num_rays + ) + return _Block(values) + + def close(self) -> None: + self.closed += 1 + + +def _RayScene() -> RaySceneDescription: + return RaySceneDescription( + num_bodies=2, + geom_types=("plane", "sphere"), + geom_sizes=np.array([[0.0, 0.0, 0.0], [0.5, 0.0, 0.0]], dtype=np.float64), + geom_local_pos=np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]], dtype=np.float64), + geom_local_quat=np.array([[1.0, 0.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0]], dtype=np.float64), + geom_body_ids=np.array([0, 1], dtype=np.intp), + body_names=("world", "base"), + ) + + +class _ReadPlan: + device = torch.device("cpu") + ready = True + + def __init__(self, pos: torch.Tensor, quat: torch.Tensor): + self.pos_w = pos + self.quat_w = quat + + def ray_pose_view(self, entity_name: str, body_names) -> Any: + self.entity_name = entity_name + self.body_names = tuple(body_names) + return SimpleNamespace(body_names=self.body_names, pos_w=self.pos_w, quat_w=self.quat_w) + + +class _Backend: + backend_type = "fake-ray-backend" + num_envs = 2 + + def __init__(self, scene: Any | None = None, collision: Any | None = None): + self.scene = scene if scene is not None else _RayScene() + self.collision = ( + collision if collision is not None else RayCollisionDescription(scene=self.scene) + ) + + def get_ray_scene_description(self, body_names=None) -> Any: + if tuple(body_names or ()) != self.scene.body_names: + raise KeyError(tuple(body_names)) + return self.scene + + def get_ray_collision_description(self, body_names=None) -> Any: + if body_names is not None and tuple(body_names) != self.scene.body_names: + raise KeyError(tuple(body_names)) + return self.collision + + def get_body_ids(self, body_names) -> np.ndarray: + names = tuple(self.scene.body_names) + return np.asarray([names.index(name) for name in body_names], dtype=np.int64) + + def get_ray_site_frames(self, site_names) -> tuple[RaySiteFrameDescription, ...]: + names = tuple(site_names) + if names != ("foot",): + raise KeyError(names) + return ( + RaySiteFrameDescription( + site_name="foot", + body_name="base", + pos=np.array([0.1, 0.2, 0.3]), + quat=np.array([1.0, 0.0, 0.0, 0.0]), + ), + ) + + def get_ray_geometry_state(self) -> RayGeometryState: + return RayGeometryState( + num_envs=2, + geom_names=("ground", "ball"), + geom_sizes=np.broadcast_to( + np.array([[0.0, 0.0, 0.0], [0.5, 0.0, 0.0]]), (2, 2, 3) + ).copy(), + geom_local_pos=np.zeros((2, 2, 3)), + geom_local_quat=np.broadcast_to(np.array([1.0, 0.0, 0.0, 0.0]), (2, 2, 4)).copy(), + geom_enabled=np.ones((2, 2), dtype=np.bool_), + geom_mesh_ids=np.full((2, 2), -1, dtype=np.int32), + ) + + +class _Env: + num_envs = 2 + device = torch.device("cpu") + backend = _Backend() + + def __init__(self, read_plan: Any, caster: _FakeCaster, scene: Any, collision: Any): + self.scene = SimpleNamespace( + read_plan=read_plan, + entities={"robot": SimpleNamespace(body_names=("base",))}, + ) + self.backend = _Backend(scene, collision) + self._caster = caster + + @property + def caster(self) -> _FakeCaster: + return self._caster + + +def _bind( + monkeypatch, + values: torch.Tensor | None = None, + pose_pos: torch.Tensor | None = None, + pose_quat: torch.Tensor | None = None, + pattern: RayPatternCfg | None = None, + **pattern_kwargs: Any, +) -> tuple[_FakeCaster, RayScan]: + scene = _RayScene() + collision = RayCollisionDescription( + scene=scene, + geom_names=("ground", "ball"), + geom_tags=(("terrain",), ("robot",)), + body_parent_ids=np.array([-1, 0], dtype=np.intp), + ) + if pattern is None: + pattern = RayPatternCfg( + entity_name="robot", + attachment_body_name="base", + pattern=RayRingPatternCfg(num_rays=4, start=0.0, stop=1.5), + max_distance=3.0, + ) + pattern = replace(pattern, **pattern_kwargs) + caster = _FakeCaster(num_rays=pattern.pattern.ray_count, values=values) + if pose_pos is None: + pose_pos = torch.zeros((2, 1, 3), dtype=torch.float32) + if pose_quat is None: + pose_quat = torch.zeros((2, 1, 4), dtype=torch.float32) + pose_quat[..., 0] = 1.0 + read_plan = _ReadPlan(pose_pos, pose_quat) + env = _Env(read_plan, caster, scene, collision) + term = RayScan(ManagerTermBaseCfg(func=RayScan, params={"ray_pattern": pattern}), env) + monkeypatch.setattr( + "unilab.envs.mdp.ray_observation._create_ray_caster", + lambda caster_name, **kwargs: ( + setattr(caster, "scene", kwargs["collision"].scene), + setattr(caster, "expected_num_rays", kwargs["num_rays"]), + caster, + )[2], + ) + term.bind_read_phase() + return caster, term + + +def test_ray_pattern_config_validates_public_names_and_values() -> None: + assert ( + RayPatternCfg( + entity_name="robot", + attachment_body_name="base", + pattern=RayRingPatternCfg(num_rays=4, start=0.0, stop=1.0), + max_distance=2, + ).attachment_body_name + == "base" + ) + for kwargs in ( + {"entity_name": ""}, + {"attachment_body_name": None}, + {"num_rays": 0}, + {"max_distance": 0.0}, + {"min_distance": -0.1}, + {"min_distance": 1.0}, + {"sensor_offset": "0 0 0"}, + {"sensor_offset": (0.0, 0.0)}, + {"attachment_body_name": ""}, + {"caster_options": {"collision": object()}}, + {"caster": ""}, + ): + base = dict( + entity_name="robot", + attachment_body_name="base", + pattern=RayRingPatternCfg(num_rays=4, start=0.0, stop=1.0), + max_distance=1.0, + ) + base.update(kwargs) + with pytest.raises((TypeError, ValueError)): + RayPatternCfg(**base) + + +def test_ray_pattern_configs_validate_geometry_values() -> None: + for make_pattern in ( + lambda: RayRingPatternCfg(num_rays=0, start=0.0, stop=1.0), + lambda: RayRingPatternCfg(num_rays=2, start=0.0, stop=float("nan")), + lambda: RayGridPatternCfg(size=(0.0, 1.0), resolution=0.1), + lambda: RayGridPatternCfg(size=(1.0, 1.0), resolution=0.0), + lambda: RayGridPatternCfg(size=(1.0, 1.0), resolution=0.1, direction=(0.0, 0.0, 0.0)), + lambda: RayCameraPatternCfg(width=0, height=2, fovy=90.0), + lambda: RayCameraPatternCfg(width=2, height=2, fovy=180.0), + lambda: RayCameraPatternCfg( + width=2, height=2, fovy=90.0, forward=(1.0, 0.0, 0.0), up=(1.0, 0.0, 0.0) + ), + ): + with pytest.raises((TypeError, ValueError)): + make_pattern() + with pytest.raises(ValueError): + RayPatternCfg( + entity_name="robot", + attachment_body_name="base", + pattern=RayRingPatternCfg(num_rays=2, start=0.0, stop=1.0), + max_distance=1.0, + alignment="attachment", + ) + + +def test_ray_scan_binds_and_traces_without_importing_uni_ray(monkeypatch) -> None: + class _Blocker: + def find_module(self, fullname, path=None): + if fullname == "uni_ray" or fullname.startswith("uni_ray."): + return self + return None + + def load_module(self, fullname): + raise AssertionError(f"uni_ray was imported: {fullname}") + + sys.meta_path.insert(0, _Blocker()) + try: + caster, term = _bind(monkeypatch) + result = term(term._env) + + assert caster.materialized + assert caster.layout == ("world", "base") + assert term.observation_width == 4 + assert caster.poses[-1].body_names == ("world", "base") + assert result.shape == (2, 4) + assert result.dtype == torch.float32 + assert result.device == torch.device("cpu") + term.close() + assert caster.closed == 1 + term.close() + assert caster.closed == 1 + finally: + del sys.meta_path[0] + + +def test_ray_scan_copies_reused_caster_output_into_term_owned_carrier(monkeypatch) -> None: + source = torch.arange(8, dtype=torch.float32).reshape(2, 4) + caster, term = _bind(monkeypatch, source) + + first = term(term._env) + second = term(term._env) + + assert first is second + torch.testing.assert_close(first, source) + source[0, 0] = 99.0 + third = term(term._env) + assert third[0, 0].item() == source[0, 0].item() + term.close() + + +def test_ray_scan_uses_public_collision_descriptor_and_rich_query(monkeypatch) -> None: + source = torch.arange(8, dtype=torch.float32).reshape(2, 4) + caster, term = _bind( + monkeypatch, + values=source, + min_distance=0.2, + include_normal=True, + allowed_geom_names=("ball",), + allowed_tags=("robot",), + ) + result = term(term._env) + collision = term._env.backend.collision + packed = result.reshape(2, -1, 4) + + assert term.observation_width == 16 + assert result.shape == (2, term.observation_width) + assert caster.scene is collision.scene + assert caster.last_query is not None + min_distance, ray_filter, outputs = caster.last_query + assert min_distance == 0.2 + assert ray_filter is not None + assert ray_filter.geom_mask(collision).tolist() == [False, True] + assert outputs.normal + torch.testing.assert_close(packed[..., 0], source) + torch.testing.assert_close(packed[..., 1:4], torch.zeros((2, 4, 3))) + term.close() + + +def test_ray_scan_excludes_attachment_body_through_public_filter(monkeypatch) -> None: + caster, term = _bind(monkeypatch, excluded_body_names=("base",)) + term(term._env) + collision = term._env.backend.collision + + assert caster.last_query is not None + ray_filter = caster.last_query[1] + assert ray_filter is not None + assert ray_filter.geom_mask(collision).tolist() == [True, False] + term.close() + + +def test_ray_scan_fails_closed_on_unsupported_rich_capabilities(monkeypatch) -> None: + caster = _FakeCaster() + read_plan = _ReadPlan( + torch.zeros((2, 1, 3), dtype=torch.float32), + torch.zeros((2, 1, 4), dtype=torch.float32), + ) + scene = _RayScene() + collision = RayCollisionDescription( + scene=scene, + geom_names=("ground", "ball"), + geom_tags=(("terrain",), ("robot",)), + body_parent_ids=np.array([-1, 0], dtype=np.intp), + ) + env = _Env(read_plan, caster, scene, collision) + pattern = RayPatternCfg( + entity_name="robot", + attachment_body_name="base", + pattern=RayRingPatternCfg(num_rays=2, start=0.0, stop=1.0), + max_distance=2.0, + min_distance=0.1, + ) + term = RayScan(ManagerTermBaseCfg(func=RayScan, params={"ray_pattern": pattern}), env) + caster.get_ray_capabilities = lambda: _Capabilities(supports_min_distance=False) + monkeypatch.setattr( + "unilab.envs.mdp.ray_observation._create_ray_caster", + lambda caster_name, **kwargs: caster, + ) + with pytest.raises(NotImplementedError, match="min_distance"): + term.bind_read_phase() + assert not caster.materialized + + +def test_ray_scan_fails_closed_on_unsupported_filter_and_normal(monkeypatch) -> None: + def bind(**pattern_kwargs): + caster = _FakeCaster() + read_plan = _ReadPlan( + torch.zeros((2, 1, 3), dtype=torch.float32), + torch.zeros((2, 1, 4), dtype=torch.float32), + ) + scene = _RayScene() + collision = RayCollisionDescription( + scene=scene, + geom_names=("ground", "ball"), + geom_tags=(("terrain",), ("robot",)), + body_parent_ids=np.array([-1, 0], dtype=np.intp), + ) + env = _Env(read_plan, caster, scene, collision) + values = { + "entity_name": "robot", + "attachment_body_name": "base", + "pattern": RayRingPatternCfg(num_rays=2, start=0.0, stop=1.0), + "max_distance": 2.0, + } + values.update(pattern_kwargs) + term = RayScan( + ManagerTermBaseCfg(func=RayScan, params={"ray_pattern": RayPatternCfg(**values)}), + env, + ) + monkeypatch.setattr( + "unilab.envs.mdp.ray_observation._create_ray_caster", + lambda caster_name, **kwargs: caster, + ) + return caster, term + + caster, term = bind(excluded_body_names=("base",)) + caster.get_ray_capabilities = lambda: _Capabilities(supports_static_filter=False) + with pytest.raises(NotImplementedError, match="static ray filters"): + term.bind_read_phase() + + caster, term = bind(include_normal=True) + caster.get_ray_capabilities = lambda: _Capabilities(supports_tensor_normal=False) + with pytest.raises(NotImplementedError, match="tensor normals"): + term.bind_read_phase() + + caster, term = bind(include_hit_point=True) + caster.get_ray_capabilities = lambda: _Capabilities(supports_hit_point=False) + with pytest.raises(NotImplementedError, match="tensor hit points"): + term.bind_read_phase() + + caster, term = bind(sync_geometry=True) + caster.get_ray_capabilities = lambda: _Capabilities(supports_tensor_geometry_sync=False) + with pytest.raises(NotImplementedError, match="tensor geometry sync"): + term.bind_read_phase() + + +def test_ray_scan_maps_ancestor_exclusion_from_public_names(monkeypatch) -> None: + caster, term = _bind(monkeypatch, excluded_ancestor_body_names=("world",)) + term(term._env) + collision = term._env.backend.collision + + assert caster.last_query is not None + ray_filter = caster.last_query[1] + assert ray_filter is not None + assert ray_filter.geom_mask(collision).tolist() == [False, False] + term.close() + + +def test_ray_scan_fails_closed_on_missing_ancestry(monkeypatch) -> None: + caster = _FakeCaster() + read_plan = _ReadPlan( + torch.zeros((2, 1, 3), dtype=torch.float32), + torch.zeros((2, 1, 4), dtype=torch.float32), + ) + scene = _RayScene() + collision = RayCollisionDescription( + scene=scene, + geom_names=("ground", "ball"), + geom_tags=(("terrain",), ("robot",)), + ) + env = _Env(read_plan, caster, scene, collision) + pattern = RayPatternCfg( + entity_name="robot", + attachment_body_name="base", + pattern=RayRingPatternCfg(num_rays=2, start=0.0, stop=1.0), + max_distance=2.0, + excluded_ancestor_body_names=("world",), + ) + term = RayScan(ManagerTermBaseCfg(func=RayScan, params={"ray_pattern": pattern}), env) + with pytest.raises(UnsupportedCapabilityError, match="ancestor"): + term.bind_read_phase() + + +def test_ray_scan_fails_closed_on_uncovered_non_root_body(monkeypatch) -> None: + scene = RaySceneDescription( + num_bodies=3, + geom_types=("plane", "sphere", "box"), + geom_sizes=np.array([[0.0, 0.0, 0.0], [0.5, 0.0, 0.0], [0.2, 0.2, 0.2]], dtype=np.float64), + geom_local_pos=np.zeros((3, 3), dtype=np.float64), + geom_local_quat=np.tile(np.array([1.0, 0.0, 0.0, 0.0]), (3, 1)), + geom_body_ids=np.array([0, 1, 2], dtype=np.intp), + body_names=("world", "base", "static_object"), + ) + collision = RayCollisionDescription( + scene=scene, + geom_names=("ground", "ball", "object"), + geom_tags=(("terrain",), ("robot",), ("terrain",)), + body_parent_ids=np.array([-1, 0, 0], dtype=np.intp), + ) + caster = _FakeCaster() + read_plan = _ReadPlan( + torch.zeros((2, 1, 3), dtype=torch.float32), + torch.zeros((2, 1, 4), dtype=torch.float32), + ) + env = _Env(read_plan, caster, scene, collision) + pattern = RayPatternCfg( + entity_name="robot", + attachment_body_name="base", + pattern=RayRingPatternCfg(num_rays=2, start=0.0, stop=1.0), + max_distance=2.0, + ) + term = RayScan(ManagerTermBaseCfg(func=RayScan, params={"ray_pattern": pattern}), env) + with pytest.raises(NotImplementedError, match="uncovered bodies"): + term.bind_read_phase() + + +def test_ray_scan_publishes_matching_attachment_ray_batches(monkeypatch) -> None: + caster, term = _bind(monkeypatch) + + term(term._env) + + origins, directions, max_distance = caster.last_inputs + assert origins.shape == directions.shape == (2, 4, 3) + assert origins.is_contiguous() and directions.is_contiguous() + assert origins.dtype == directions.dtype == torch.float32 + assert origins.device == directions.device == torch.device("cpu") + torch.testing.assert_close( + origins, + torch.tensor( + [[[0.0, 0.0, 0.0]] * 4] * 2, + dtype=torch.float32, + ), + ) + angles = torch.linspace( + term.ray_pattern.pattern.start, + term.ray_pattern.pattern.stop, + term.ray_pattern.pattern.num_rays, + dtype=torch.float32, + ) + expected_directions = torch.stack( + (torch.cos(angles), torch.sin(angles), torch.zeros_like(angles)), dim=-1 + ).repeat(2, 1, 1) + torch.testing.assert_close(directions, expected_directions) + assert max_distance == term.ray_pattern.max_distance + term.close() + + +def test_ray_pattern_follows_body_attachment_pose(monkeypatch) -> None: + pose_pos = torch.tensor([[[0.1, 0.2, 0.3]], [[-0.4, 0.5, -0.2]]], dtype=torch.float32) + pose_quat = torch.zeros((2, 1, 4), dtype=torch.float32) + pose_quat[..., 0] = 0.5**0.5 + pose_quat[..., 3] = 0.5**0.5 + caster, term = _bind( + monkeypatch, + pose_pos=pose_pos, + pose_quat=pose_quat, + sensor_offset=(0.2, 0.0, 0.1), + ) + + term(term._env) + origins, directions, _ = caster.last_inputs + pose = caster.poses[-1] + assert pose.body_names == ("world", "base") + assert torch.count_nonzero(pose.pos_w[:, 0, :]).item() == 0 + torch.testing.assert_close(pose.pos_w[:, 1, :], pose_pos[:, 0, :]) + rotation = torch.tensor( + [[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]], + dtype=torch.float32, + ) + expected_origin = pose_pos + torch.tensor([0.0, 0.2, 0.1], dtype=torch.float32)[None, None, :] + angles = torch.linspace(0.0, 1.5, 4, dtype=torch.float32) + local_directions = torch.stack( + (torch.cos(angles), torch.sin(angles), torch.zeros_like(angles)), dim=-1 + ) + expected_directions = local_directions[None, :, :] @ rotation.T + torch.testing.assert_close(origins, expected_origin.repeat(1, 4, 1)) + torch.testing.assert_close(directions, expected_directions.repeat(2, 1, 1)) + term.close() + + +def test_grid_pattern_uses_yaw_alignment(monkeypatch) -> None: + pose_quat = torch.zeros((2, 1, 4), dtype=torch.float32) + pose_quat[..., 0] = 0.5**0.5 + pose_quat[..., 3] = 0.5**0.5 + caster, term = _bind( + monkeypatch, + pose_quat=pose_quat, + pattern=RayPatternCfg( + entity_name="robot", + attachment_body_name="base", + pattern=RayGridPatternCfg(size=(0.2, 0.1), resolution=0.1, offset=(0.05, 0.0)), + alignment="yaw", + max_distance=4.0, + ), + ) + + term(term._env) + origins, directions, _ = caster.last_inputs + assert origins.shape == directions.shape == (2, 6, 3) + local_origins = torch.tensor( + [ + [-0.05, -0.05, 0.0], + [0.05, -0.05, 0.0], + [0.15, -0.05, 0.0], + [-0.05, 0.05, 0.0], + [0.05, 0.05, 0.0], + [0.15, 0.05, 0.0], + ], + dtype=torch.float32, + ) + yaw_rotation = torch.tensor( + [[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]], + dtype=torch.float32, + ) + expected_origins = local_origins[None, :, :] @ yaw_rotation.T + torch.testing.assert_close(origins, expected_origins.repeat(2, 1, 1)) + torch.testing.assert_close( + directions, torch.zeros_like(directions) - torch.tensor((0.0, 0.0, 1.0)) + ) + term.close() + + +def test_world_alignment_ignores_attachment_rotation(monkeypatch) -> None: + pose_pos = torch.tensor([[[0.3, 0.2, 0.1]]] * 2, dtype=torch.float32) + pose_quat = torch.zeros((2, 1, 4), dtype=torch.float32) + pose_quat[..., 0] = 0.5**0.5 + pose_quat[..., 3] = 0.5**0.5 + caster, term = _bind( + monkeypatch, + pose_pos=pose_pos, + pose_quat=pose_quat, + pattern=RayPatternCfg( + entity_name="robot", + attachment_body_name="base", + pattern=RayRingPatternCfg(num_rays=1, start=0.0, stop=0.0), + alignment="world", + max_distance=1.0, + sensor_offset=(0.2, 0.0, 0.0), + ), + ) + + term(term._env) + origins, directions, _ = caster.last_inputs + torch.testing.assert_close(origins, torch.tensor([[[0.5, 0.2, 0.1]]] * 2, dtype=torch.float32)) + torch.testing.assert_close( + directions, torch.tensor([[[1.0, 0.0, 0.0]]] * 2, dtype=torch.float32) + ) + term.close() + + +def test_site_attachment_composes_public_site_frame(monkeypatch) -> None: + pose_pos = torch.tensor([[[0.4, 0.2, 0.1]]] * 2, dtype=torch.float32) + pose_quat = torch.zeros((2, 1, 4), dtype=torch.float32) + pose_quat[..., 0] = 0.5**0.5 + pose_quat[..., 3] = 0.5**0.5 + caster, term = _bind( + monkeypatch, + pose_pos=pose_pos, + pose_quat=pose_quat, + pattern=RayPatternCfg( + entity_name="robot", + attachment_site_name="foot", + pattern=RayRingPatternCfg(num_rays=1, start=0.0, stop=0.0), + max_distance=2.0, + sensor_offset=(0.03, 0.0, 0.0), + ), + ) + + term(term._env) + origins, directions, _ = caster.last_inputs + expected_origin = torch.tensor( + [[[0.2, 0.33, 0.4]]] * 2, + dtype=torch.float32, + ) + expected_direction = torch.tensor([[[0.0, 1.0, 0.0]]] * 2, dtype=torch.float32) + torch.testing.assert_close(origins, expected_origin) + torch.testing.assert_close(directions, expected_direction) + term.close() + + +def test_ray_scan_syncs_backend_geometry_on_bind_and_reset(monkeypatch) -> None: + caster, term = _bind(monkeypatch, sync_geometry=True) + + assert len(caster.geometry_updates) == 1 + geometry = caster.geometry_updates[0] + assert geometry.geom_sizes.shape == (2, 2, 3) + assert geometry.geom_enabled.shape == (2, 2) + assert geometry.geom_mesh_ids.dtype == torch.int32 + + term.reset() + assert len(caster.geometry_updates) == 2 + term.close() + + +def test_camera_pattern_generates_normalized_pinhole_rays() -> None: + pattern = RayCameraPatternCfg(width=3, height=2, fovy=90.0) + directions = pattern.directions(torch.device("cpu")) + + assert pattern.ray_count == 6 + assert directions.shape == (6, 3) + torch.testing.assert_close( + torch.linalg.vector_norm(directions, dim=-1), torch.ones(6), atol=1e-6, rtol=0.0 + ) + torch.testing.assert_close( + directions[1], + torch.tensor((2.0 / 5.0**0.5, 0.0, 1.0 / 5.0**0.5), dtype=torch.float32), + atol=1e-6, + rtol=0.0, + ) + + +def test_ray_scan_publishes_optional_hit_point(monkeypatch) -> None: + source = torch.arange(8, dtype=torch.float32).reshape(2, 4) + caster, term = _bind( + monkeypatch, + values=source, + include_normal=True, + include_hit_point=True, + ) + result = term(term._env) + + assert term.observation_width == 28 + assert result.shape == (2, 28) + assert caster.last_query is not None + outputs = caster.last_query[2] + assert outputs.normal and outputs.hit_point + packed = result.reshape(2, 4, 7) + torch.testing.assert_close(packed[..., 0], source) + torch.testing.assert_close(packed[..., 1:4], torch.zeros((2, 4, 3))) + torch.testing.assert_close(packed[..., 4:7], torch.ones((2, 4, 3))) + term.close() + + +def test_ray_scan_fails_closed_before_binding_and_on_unsupported_capability(monkeypatch) -> None: + scene = _RayScene() + caster = _FakeCaster() + read_plan = _ReadPlan( + torch.zeros((2, 1, 3), dtype=torch.float32), + torch.zeros((2, 1, 4), dtype=torch.float32), + ) + collision = RayCollisionDescription( + scene=scene, + geom_names=("ground", "ball"), + geom_tags=(("terrain",), ("robot",)), + body_parent_ids=np.array([-1, 0], dtype=np.intp), + ) + env = _Env(read_plan, caster, scene, collision) + pattern = RayPatternCfg( + entity_name="robot", + attachment_body_name="base", + pattern=RayRingPatternCfg(num_rays=2, start=0.0, stop=1.0), + max_distance=1.0, + ) + term = RayScan(ManagerTermBaseCfg(func=RayScan, params={"ray_pattern": pattern}), env) + with pytest.raises(NotImplementedError, match="bind its read phase"): + term(env) + + caster.get_ray_capabilities = lambda: _Capabilities(supports_tensor_rays=False) + monkeypatch.setattr( + "unilab.envs.mdp.ray_observation._create_ray_caster", + lambda caster_name, **kwargs: caster, + ) + with pytest.raises(NotImplementedError, match="tensor pose/ray/output profile"): + term.bind_read_phase() + assert not caster.materialized + assert caster.layout is None + + +def test_ray_scan_uses_only_public_backend_and_factory_boundaries(monkeypatch) -> None: + caster, term = _bind(monkeypatch) + env = term._env + + collision = env.backend.get_ray_collision_description() + assert collision.scene is env.backend.scene + assert collision is env.backend.collision + assert term.ray_pattern.caster == "uni_ray" + source = (_REPO_ROOT / "src/unilab/envs/mdp/ray_observation.py").read_text(encoding="utf-8") + for banned in ( + "import uni_ray", + "MjModel", + "mjwarp.Model", + "MjBatchCollision", + "build_collision_description", + ): + assert banned not in source, banned + assert not any( + path.name.startswith("_") and path.is_file() for path in Path(__file__).parent.rglob("*.py") + ) + term.close() diff --git a/tests/envs/test_env_configs.py b/tests/envs/test_env_configs.py index a53e611c2..639953a96 100644 --- a/tests/envs/test_env_configs.py +++ b/tests/envs/test_env_configs.py @@ -637,6 +637,41 @@ def test_flashsac_g1_motion_mujoco_tensor_command_model_field_dr_roll_out( env._backend.close() +def test_sac_g1_motion_genesis_flashmimic_dr_roll_out() -> None: + """The SAC Genesis owner exercises the FlashMimic device-DR subset.""" + ensure_registries() + _require_genesis_runtime() + from unilab.base import registry + from unilab.envs import ManagerBasedRlEnv + + _, override = _motion_manager_override( + "g1_motion_tracking", + "genesis", + config_root="sac", + ) + env = registry.make( + "G1MotionTracking", + num_envs=2, + sim_backend="genesis", + env_cfg_override=override, + ) + assert isinstance(env, ManagerBasedRlEnv) + try: + names = env.event_manager._mode_term_names + assert set(names["startup"]) == {"body_mass"} + assert set(names["reset"]) >= {"base_com", "encoder_bias", "pd_gains"} + assert set(names["interval"]) == {"push_robot"} + obs, _ = env.reset(seed=2521) + assert all(torch.isfinite(values).all() for values in obs.values()) + for _ in range(3): + state = env.step(torch.zeros((2, 29), dtype=torch.float32, device=env.device)) + assert torch.isfinite(state.reward).all() + assert all(torch.isfinite(values).all() for values in state.obs.values()) + finally: + env.close() + env._backend.close() + + def test_flashsac_g1_motion_genesis_manager_tensor_command_roll_out() -> None: """The canonical FlashSAC Genesis owner exercises the Manager tensor path.""" ensure_registries() diff --git a/tests/envs/test_manager_based_rl_env.py b/tests/envs/test_manager_based_rl_env.py index 25382bb2c..6ba330f36 100644 --- a/tests/envs/test_manager_based_rl_env.py +++ b/tests/envs/test_manager_based_rl_env.py @@ -1580,6 +1580,110 @@ def test_mujoco_host_bridge_packed_reset_lifecycle( backend.close() +def test_mujoco_tensor_ray_observation_uses_public_rich_query_contract( + tmp_path: Path, +) -> None: + """Exercise the real MuJoCo/UniSim/uni_ray perception path end to end.""" + + pytest.importorskip("uni_ray", reason="local uni_ray is required") + pytest.importorskip("warp", reason="warp-lang is required by uni_ray") + + def build(pattern: mdp.RayPatternCfg): + model = tmp_path / "ray-rich-query.xml" + model.write_text( + "" + "" + "" + "" + "" + ) + scene = SceneCfg( + model_file=str(model), + entities={ + "robot": EntityCfg( + root_body_name="base", + body_names=("base",), + joint_names=("joint",), + actuator_names=("motor",), + ) + }, + ) + cfg = ManagerBasedRlEnvCfg( + scene=scene, + sim_dt=0.01, + ctrl_dt=0.02, + max_episode_seconds=0.04, + observations={ + "actor": ObservationGroupCfg( + terms={ + "rays": ObservationTermCfg( + func=mdp.RayScan, params={"ray_pattern": pattern} + ) + } + ) + }, + actions={ + "effort": mdp.JointEffortActionCfg(entity_name="robot", actuator_names=(".*",)) + }, + events={"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 + ) + return ManagerBasedRlEnv(cfg, backend, 2), backend + + def ray_pattern(**kwargs): + values = { + "entity_name": "robot", + "attachment_body_name": "base", + "pattern": mdp.RayRingPatternCfg( + num_rays=2, + start=0.0, + stop=0.0, + pitch=-0.5 * 3.141592653589793, + ), + "max_distance": 3.0, + "sensor_offset": (0.0, 0.0, 0.2), + "include_normal": True, + "sync_geometry": True, + } + values.update(kwargs) + return mdp.RayPatternCfg(**values) + + near_env, near_backend = build(ray_pattern()) + valid_env, valid_backend = build(ray_pattern(min_distance=0.15, excluded_body_names=("base",))) + try: + near = near_env.init_state().obs["obs"].reshape(2, -1, 4) + valid = valid_env.init_state().obs["obs"].reshape(2, -1, 4) + torch.testing.assert_close(near[..., 0], torch.full((2, 2), 0.1), atol=2e-5, rtol=0.0) + expected_normal = torch.tensor([0.0, 0.0, 1.0]).repeat(2, 2, 1) + torch.testing.assert_close(near[..., 1:4], expected_normal, atol=2e-5, rtol=0.0) + torch.testing.assert_close(valid[..., 0], torch.full((2, 2), 1.2), atol=2e-5, rtol=0.0) + torch.testing.assert_close(valid[..., 1:4], expected_normal, atol=2e-5, rtol=0.0) + + site_env, site_backend = build( + ray_pattern( + attachment_body_name=None, + attachment_site_name="sensor", + sensor_offset=(0.0, 0.0, 0.0), + sync_geometry=False, + ) + ) + site = site_env.init_state().obs["obs"].reshape(2, -1, 4) + torch.testing.assert_close(site[..., 0], torch.full((2, 2), 0.1), atol=2e-5, rtol=0.0) + finally: + site_env.close() + site_backend.close() + near_env.close() + valid_env.close() + near_backend.close() + valid_backend.close() + + def _make_state_env( *, commands: dict[str, CommandTermCfg | None] | None = None, diff --git a/tests/managers/test_device_contracts.py b/tests/managers/test_device_contracts.py index dfdde5429..f8c760ead 100644 --- a/tests/managers/test_device_contracts.py +++ b/tests/managers/test_device_contracts.py @@ -16,6 +16,7 @@ TensorExecution, TensorLifecycleCapabilities, TensorProcessTopology, + TrackedBodyStateViews, ) from unilab.base.entity import EntityCfg, EntityScene, SceneTensorReadSpec @@ -308,6 +309,18 @@ def get_state_views(self, names, device=None): def get_sensor_view(self, name, device=None): return torch.as_tensor(self.sensor_values[name], dtype=torch.float32, device=device) + def get_tracked_body_views(self, body_names, device=None): + names = tuple(name.rsplit("/", 1)[-1] for name in body_names) + if names != ("base",): + raise KeyError(names) + return TrackedBodyStateViews( + body_names=names, + pos_w=torch.as_tensor(self.body_pos, dtype=torch.float32, device=device), + quat_w=torch.as_tensor(self.body_quat, dtype=torch.float32, device=device), + lin_vel_w=torch.as_tensor(self.body_lin_vel_w, dtype=torch.float32, device=device), + ang_vel_w=torch.as_tensor(self.body_ang_vel_w, dtype=torch.float32, device=device), + ) + 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) diff --git a/tests/scripts/test_rocm_dependency_profile.py b/tests/scripts/test_rocm_dependency_profile.py index bfc1a7c27..7724ae775 100644 --- a/tests/scripts/test_rocm_dependency_profile.py +++ b/tests/scripts/test_rocm_dependency_profile.py @@ -24,6 +24,10 @@ def _project_scripts(path: Path) -> dict[str, str]: return dict(project["project"]["scripts"]) +def _project(path: Path) -> dict: + return tomllib.loads(path.read_text(encoding="utf-8")) + + def test_rocm_profile_preserves_non_substituted_runtime_dependencies() -> None: default_dependencies = _dependencies(ROOT / "pyproject.toml") rocm_dependencies = _dependencies(ROOT / "pyproject.rocm.toml") @@ -55,3 +59,15 @@ def test_rocm_lock_contains_profile_runtime_dependencies() -> None: assert "prettytable" in packages assert "prettytable" in root_dependencies assert "numba" not in root_dependencies + + +def test_ray_optional_dependencies_use_published_packages() -> None: + for path in (ROOT / "pyproject.toml", ROOT / "pyproject.rocm.toml"): + project = _project(path) + extras = project["project"]["optional-dependencies"] + assert extras["ray"] == ["uniray[warp]==0.1.0"] + assert "uniray" not in _dependencies(path) + + sources = project["tool"]["uv"]["sources"] + assert "unisim-core" not in sources + assert "uniray" not in sources diff --git a/uv.lock b/uv.lock index 3402bb6c2..acb7e352b 100644 --- a/uv.lock +++ b/uv.lock @@ -5217,6 +5217,9 @@ newton = [ { name = "pyglet" }, { name = "warp-lang" }, ] +ray = [ + { name = "uniray", extra = ["warp"] }, +] superdex = [ { name = "unisim-core", extra = ["superdex"], marker = "python_full_version >= '3.12' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, ] @@ -5268,15 +5271,16 @@ requires-dist = [ { name = "tqdm" }, { name = "trimesh", specifier = ">=3.21.7" }, { name = "typing-extensions" }, - { name = "unisim-core", specifier = "==2.1.1" }, - { name = "unisim-core", extras = ["motrix"], marker = "extra == 'motrix'", specifier = "==2.1.1" }, - { name = "unisim-core", extras = ["superdex"], marker = "python_full_version >= '3.12' and platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'superdex'", specifier = "==2.1.1" }, + { name = "uniray", extras = ["warp"], marker = "extra == 'ray'", specifier = "==0.1.0" }, + { name = "unisim-core", specifier = "==2.2.0" }, + { name = "unisim-core", extras = ["motrix"], marker = "extra == 'motrix'", specifier = "==2.2.0" }, + { name = "unisim-core", extras = ["superdex"], marker = "python_full_version >= '3.12' and platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'superdex'", specifier = "==2.2.0" }, { name = "viser", specifier = ">=1.0.26" }, { name = "wandb", specifier = ">=0.22.3" }, { name = "warp-lang", marker = "extra == 'mjwarp'", specifier = "==1.16.0" }, { name = "warp-lang", marker = "extra == 'newton'", specifier = "==1.16.0" }, ] -provides-extras = ["drake", "mujoco", "mjwarp", "newton", "motrix", "genesis", "superdex"] +provides-extras = ["drake", "mujoco", "mjwarp", "newton", "motrix", "genesis", "superdex", "ray"] [package.metadata.requires-dev] dev = [ @@ -5287,17 +5291,36 @@ dev = [ { name = "ruff" }, ] +[[package]] +name = "uniray" +version = "0.1.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "unisim-core" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/12/d4/71015d68bae08d996f0c7d19e17516cb012e43c5f480fc56a9c1fe9ab4ba/uniray-0.1.0.tar.gz", hash = "sha256:2bbc87029615d6a4efa7315a12615f217d5219f53e0627a8ba8ac93b5ea045ee", size = 32970, upload-time = "2026-10-26T12:03:45.548Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7a/79/4c27367b240a58610657673f9e83499ed4ed5295dad0d17d07ac551d897c/uniray-0.1.0-py3-none-any.whl", hash = "sha256:b4bbfd5e3befb4f6327593172f4ed8e583e49086e437926b80cdf81738f7c88b", size = 36533, upload-time = "2026-10-26T12:03:40.835Z" }, +] + +[package.optional-dependencies] +warp = [ + { name = "warp-lang" }, +] + [[package]] name = "unisim-core" -version = "2.1.1" +version = "2.2.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/ae/dc/5c8b56a04996cc3f47767ac4c3c665c0eae131a188f1554525ff7514a9ed/unisim_core-2.1.1.tar.gz", hash = "sha256:f3c9bc2497b3ef85ec984cf38f8a8ce5cfd24b900db2734685575c5bb8cef2ed", size = 699580, upload-time = "2026-10-22T10:26:44.017Z" } +sdist = { url = "https://files.pythonhosted.org/packages/f6/85/7c0c271839c31e1042e9a3cefc0b29abf344e46709a91f958e9d45de526c/unisim_core-2.2.0.tar.gz", hash = "sha256:4097bf45bd18b1bbdf124e3f662665d24e0e050e0d0d71866d87a72d88a4bfd6", size = 735592, upload-time = "2026-10-26T11:33:02.698Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/31/99/b532c452a8544eb508ff3b6009146aa24f44e09c5c43e63aaf46153b8490/unisim_core-2.1.1-py3-none-any.whl", hash = "sha256:fee2dcecb8ceb60ee2532a9e7d15f66daf1c1213d990fe4a607aef92130f6bf2", size = 746595, upload-time = "2026-10-22T10:26:42.255Z" }, + { url = "https://files.pythonhosted.org/packages/02/df/aceda11969a7c2b4db39c6899ec9990eda3f10dab14ed42ab3822a6fb8dd/unisim_core-2.2.0-py3-none-any.whl", hash = "sha256:82d72bca8f7e925cd218174d128d6e483fd227e49d3409b4e5c24720528c83f3", size = 783936, upload-time = "2026-10-26T11:33:00.742Z" }, ] [package.optional-dependencies] diff --git a/uv.rocm.lock b/uv.rocm.lock index fcde2bf27..3a8f4153c 100644 --- a/uv.rocm.lock +++ b/uv.rocm.lock @@ -3756,6 +3756,9 @@ mujoco = [ { name = "mjbatch-uni" }, { name = "mujoco" }, ] +ray = [ + { name = "uniray", extra = ["warp"] }, +] [package.dev-dependencies] dev = [ @@ -3797,12 +3800,13 @@ requires-dist = [ { name = "trimesh", specifier = ">=3.21.7" }, { name = "triton-rocm", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'", specifier = "==3.8.0", index = "https://download.pytorch.org/whl/rocm7.2" }, { name = "typing-extensions" }, - { name = "unisim-core", specifier = "==2.1.1" }, - { name = "unisim-core", extras = ["motrix"], marker = "extra == 'motrix'", specifier = "==2.1.1" }, + { name = "uniray", extras = ["warp"], marker = "extra == 'ray'", specifier = "==0.1.0" }, + { name = "unisim-core", specifier = "==2.2.0" }, + { name = "unisim-core", extras = ["motrix"], marker = "extra == 'motrix'", specifier = "==2.2.0" }, { name = "viser", specifier = ">=1.0.26" }, { name = "wandb", specifier = ">=0.22.3" }, ] -provides-extras = ["drake", "mujoco", "motrix"] +provides-extras = ["drake", "mujoco", "motrix", "ray"] [package.metadata.requires-dev] dev = [ @@ -3813,17 +3817,36 @@ dev = [ { name = "ruff" }, ] +[[package]] +name = "uniray" +version = "0.1.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "unisim-core" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/12/d4/71015d68bae08d996f0c7d19e17516cb012e43c5f480fc56a9c1fe9ab4ba/uniray-0.1.0.tar.gz", hash = "sha256:2bbc87029615d6a4efa7315a12615f217d5219f53e0627a8ba8ac93b5ea045ee", size = 32970, upload-time = "2026-10-26T12:03:45.548Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7a/79/4c27367b240a58610657673f9e83499ed4ed5295dad0d17d07ac551d897c/uniray-0.1.0-py3-none-any.whl", hash = "sha256:b4bbfd5e3befb4f6327593172f4ed8e583e49086e437926b80cdf81738f7c88b", size = 36533, upload-time = "2026-10-26T12:03:40.835Z" }, +] + +[package.optional-dependencies] +warp = [ + { name = "warp-lang" }, +] + [[package]] name = "unisim-core" -version = "2.1.1" +version = "2.2.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/ae/dc/5c8b56a04996cc3f47767ac4c3c665c0eae131a188f1554525ff7514a9ed/unisim_core-2.1.1.tar.gz", hash = "sha256:f3c9bc2497b3ef85ec984cf38f8a8ce5cfd24b900db2734685575c5bb8cef2ed", size = 699580, upload-time = "2026-10-22T10:26:44.017Z" } +sdist = { url = "https://files.pythonhosted.org/packages/f6/85/7c0c271839c31e1042e9a3cefc0b29abf344e46709a91f958e9d45de526c/unisim_core-2.2.0.tar.gz", hash = "sha256:4097bf45bd18b1bbdf124e3f662665d24e0e050e0d0d71866d87a72d88a4bfd6", size = 735592, upload-time = "2026-10-26T11:33:02.698Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/31/99/b532c452a8544eb508ff3b6009146aa24f44e09c5c43e63aaf46153b8490/unisim_core-2.1.1-py3-none-any.whl", hash = "sha256:fee2dcecb8ceb60ee2532a9e7d15f66daf1c1213d990fe4a607aef92130f6bf2", size = 746595, upload-time = "2026-10-22T10:26:42.255Z" }, + { url = "https://files.pythonhosted.org/packages/02/df/aceda11969a7c2b4db39c6899ec9990eda3f10dab14ed42ab3822a6fb8dd/unisim_core-2.2.0-py3-none-any.whl", hash = "sha256:82d72bca8f7e925cd218174d128d6e483fd227e49d3409b4e5c24720528c83f3", size = 783936, upload-time = "2026-10-26T11:33:00.742Z" }, ] [package.optional-dependencies] @@ -3932,6 +3955,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c4/e8/76836b75d401ff5912aaf513176e64557ceaec4c4946bfd38a698ff84d48/wandb-0.25.1-py3-none-win_arm64.whl", hash = "sha256:cc7c34b70cf4b7be4d395541e82e325fd9d2be978d62c9ec01f1a7141523b6bb", size = 22080774, upload-time = "2026-03-10T23:51:40.196Z" }, ] +[[package]] +name = "warp-lang" +version = "1.18.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/59/ca/0f88f27e9367b55df2c87fd6dbb96f15871ab7fcca14dd9ec92ec4688a30/warp_lang-1.18.0-py3-none-macosx_11_0_arm64.whl", hash = "sha256:3d0a12f411555c08015f84060187a0466800d2c837eb8800977b9604be438c21", size = 27819487, upload-time = "2026-10-05T00:22:32.346Z" }, + { url = "https://files.pythonhosted.org/packages/13/69/5fc4388c6af4c78cd247776bb6928cabe83d893896444dd2127ccbe8b2ca/warp_lang-1.18.0-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:090962bfae7dd22ade0d8e0258f41635e0e58980f15bcc9e7a7d312819dfe440", size = 148811186, upload-time = "2026-10-05T00:23:01.197Z" }, + { url = "https://files.pythonhosted.org/packages/be/41/6e255904a10cd1b3b1a61cc9ac21551812ffc55b1c68324de4c0c24a6660/warp_lang-1.18.0-py3-none-manylinux_2_34_aarch64.whl", hash = "sha256:d76ab8c7d4f4897405fa0a744d38812fc44cfc8162c40e658dc9ca7790b6a645", size = 147348850, upload-time = "2026-10-05T00:23:31.015Z" }, + { url = "https://files.pythonhosted.org/packages/3e/75/888dfe20326fea1a3b46b68a239ccd5db717bb039d511ab6fca7c77e6fea/warp_lang-1.18.0-py3-none-win_amd64.whl", hash = "sha256:1769e55746739f2e244a2b412123003c462ec3d977dfb5f420f924a8553d30aa", size = 134607061, upload-time = "2026-10-05T00:23:56.792Z" }, +] + [[package]] name = "wcwidth" version = "0.6.0"