diff --git a/docs/sphinx/source/api_reference/algos/index.md b/docs/sphinx/source/api_reference/algos/index.md index 720af6e13..328053439 100644 --- a/docs/sphinx/source/api_reference/algos/index.md +++ b/docs/sphinx/source/api_reference/algos/index.md @@ -8,7 +8,6 @@ published on PyPI; issue #1480): - `uni_rl.algos.appo` — APPO runner, learner, staging, worker - `uni_rl.algos.fast_sac` / `uni_rl.algos.fast_td3` / `uni_rl.algos.flash_sac` — off-policy learners and runners - `uni_rl.offpolicy` — generic off-policy runner, worker, thread budget -- `uni_rl.algos.hora` — HORA models, trainers, and distillation - `uni_rl.algos.common` — shared actor factory, networks, normalization, compile helpers UniLab keeps the training *entrypoints* (`src/unilab/scripts/train_*.py`), diff --git a/docs/sphinx/source/api_reference/tasks/index.md b/docs/sphinx/source/api_reference/tasks/index.md index eb4ed03d1..1b2a6fa2b 100644 --- a/docs/sphinx/source/api_reference/tasks/index.md +++ b/docs/sphinx/source/api_reference/tasks/index.md @@ -3,7 +3,7 @@ Concrete RL tasks split by family: - **locomotion** — A2, Go1, Go2, Go2w, Go2 + Airbot, and Unitree G1 -- **manipulation** — Allegro / Sharpa in-hand cube and Stewart balance +- **manipulation** — Allegro in-hand cube and Stewart balance - **motion_tracking** — G1 and X2 whole-body motion tracking Every task is registered into the task `Registry` so it can be selected via diff --git a/docs/sphinx/source/api_reference/tasks/manipulation.md b/docs/sphinx/source/api_reference/tasks/manipulation.md index 328e820de..ebf881e31 100644 --- a/docs/sphinx/source/api_reference/tasks/manipulation.md +++ b/docs/sphinx/source/api_reference/tasks/manipulation.md @@ -7,6 +7,5 @@ :recursive: unilab.tasks.manipulation.allegro_inhand - unilab.tasks.manipulation.sharpa_inhand unilab.tasks.manipulation.stewart ``` diff --git a/docs/sphinx/source/en/0-index.md b/docs/sphinx/source/en/0-index.md index a122605fe..5d52c8084 100644 --- a/docs/sphinx/source/en/0-index.md +++ b/docs/sphinx/source/en/0-index.md @@ -108,7 +108,7 @@ capability requirements. :::{grid-item-card} Pick an algorithm :link: 2-user_guide/2-algorithms/0-index :link-type: doc -Compare PPO, APPO, SAC, TD3, FlashSAC, and HORA entrypoints. +Compare PPO, APPO, SAC, TD3, and FlashSAC entrypoints. ::: :::{grid-item-card} Deploy or switch sims @@ -158,7 +158,6 @@ committed benchmark manifest or separate recommendation metadata. | Go2W joystick | PPO | PPO owner YAMLs exist for MuJoCo and Motrix flat/rough variants under `src/unilab/conf/ppo/task/go2w_joystick_*`. | | G1 locomotion / tracking | PPO, APPO, SAC, TD3 | PPO, APPO, and SAC include committed MuJoCo and Motrix owner YAMLs for G1 tasks; TD3 has a `g1_walk_flat` MuJoCo owner. | | Allegro in-hand | PPO, APPO | PPO and APPO have committed MuJoCo and Motrix owner YAMLs for Allegro in-hand tasks. | -| Sharpa in-hand | PPO, APPO HORA teacher, HORA distillation | Sharpa owner YAMLs are committed for PPO/APPO teacher paths; student distillation uses `src/unilab/conf/hora_distill/task/sharpa_inhand/mujoco.yaml`. | ```{toctree} :hidden: 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 a49bbea33..7ef2c90a3 100644 --- a/docs/sphinx/source/en/1-getting_started/2-installation.md +++ b/docs/sphinx/source/en/1-getting_started/2-installation.md @@ -134,9 +134,8 @@ ROCm/XPU profiles, prefer the uv paths above. Robot meshes and textures are intentionally excluded from the wheel and downloaded on the cold path from the `unilabsim/unilab-robots` dataset. Ensure the installed package location is writable, or pre-fetch assets with `uv run unilab-pull-assets` from a -source checkout. The isaacgym / isaacsim backends and the HORA multi-GPU -submission path still assume a source checkout; use their dedicated setup pages -below. +source checkout. The isaacgym / isaacsim backends still assume a source +checkout; use their dedicated setup pages below. ## Runtime Assets diff --git a/docs/sphinx/source/en/1-getting_started/3-evaluation_and_playback.md b/docs/sphinx/source/en/1-getting_started/3-evaluation_and_playback.md index 43e187961..c004fdb81 100644 --- a/docs/sphinx/source/en/1-getting_started/3-evaluation_and_playback.md +++ b/docs/sphinx/source/en/1-getting_started/3-evaluation_and_playback.md @@ -39,7 +39,7 @@ Use `uv run eval` for regular evaluation and video export. When you need a live `src/unilab/scripts/play_interactive.py` script. `src/unilab/scripts/play_interactive.py` is the general MuJoCo viewer entrypoint for PPO, -APPO, SAC, FlashSAC, and HORA distill policies. It uses `--algo / --task / +APPO, SAC, and FlashSAC policies. It uses `--algo / --task / --sim` to select the algorithm and owner config. The viewer is always `mujoco.viewer`; `--sim` only selects which config to read. diff --git a/docs/sphinx/source/en/1-getting_started/4-project_structure.md b/docs/sphinx/source/en/1-getting_started/4-project_structure.md index 0114b6602..dbfc042ab 100644 --- a/docs/sphinx/source/en/1-getting_started/4-project_structure.md +++ b/docs/sphinx/source/en/1-getting_started/4-project_structure.md @@ -10,7 +10,7 @@ changing behavior. | `src/unilab/conf/` | Hydra roots and task owner YAMLs. The top-level CLI exposes backend selection as `--task` plus `--sim`, then composes the matching owner YAML. | | `src/unilab/base/` | Registry, env state, scene, and backend contracts. | | `src/unilab/envs/` | Task env implementations and task-specific reset, reward, observation, and DR logic. | -| `uni_rl` (unilab-rl repo) | PPO, APPO, off-policy, and HORA algorithm code. | +| `uni_rl` (unilab-rl repo) | PPO, APPO, and off-policy algorithm code. | | `uni_rl.ipc` (unilab-rl repo) | Shared-memory and async runner primitives. | | `src/unilab/training/` | Shared training helpers for logging, playback, seed handling, and config guards. | | `src/unilab/visualization/` | Playback, rendering, NaN inspection, and scene/export utilities. | @@ -27,7 +27,6 @@ The main config roots are: - `src/unilab/conf/sac/config.yaml`, `src/unilab/conf/td3/config.yaml`, and `src/unilab/conf/flashsac/config.yaml` for SAC, TD3, and FlashSAC, each with its algorithm hyperparameters inlined. -- `src/unilab/conf/hora_distill/config.yaml` for HORA student distillation. Task owner YAMLs are the backend identity. Examples: diff --git a/docs/sphinx/source/en/2-user_guide/0-index.md b/docs/sphinx/source/en/2-user_guide/0-index.md index c3cdd8b36..b99c7a074 100644 --- a/docs/sphinx/source/en/2-user_guide/0-index.md +++ b/docs/sphinx/source/en/2-user_guide/0-index.md @@ -15,7 +15,7 @@ CLI routes, Hydra owner YAMLs, logs, checkpoints, and Docker. :::{grid-item-card} Algorithms :link: 2-algorithms/0-index :link-type: doc -Compare PPO, APPO, SAC, TD3, FlashSAC, and HORA. +Compare PPO, APPO, SAC, TD3, and FlashSAC. ::: :::{grid-item-card} Backends @@ -42,12 +42,6 @@ Configure reset, init, and interval randomization through task owner configs. Export ONNX, inspect NaNs, send W&B logs, and export scenes. ::: -:::{grid-item-card} Manipulation Notes -:link: 8-manipulation/0-index -:link-type: doc -Task-specific Allegro, Sharpa, and Go2+Airbot notes. -::: - :::: ```{toctree} @@ -61,5 +55,4 @@ Task-specific Allegro, Sharpa, and Go2+Airbot notes. 5-domain_randomization/0-index 6-terrain/0-index 7-tooling/0-index -8-manipulation/0-index ``` diff --git a/docs/sphinx/source/en/2-user_guide/1-training/2-hydra_config.md b/docs/sphinx/source/en/2-user_guide/1-training/2-hydra_config.md index 365a19165..ea9617419 100644 --- a/docs/sphinx/source/en/2-user_guide/1-training/2-hydra_config.md +++ b/docs/sphinx/source/en/2-user_guide/1-training/2-hydra_config.md @@ -10,7 +10,6 @@ identity of the task, backend, reward, scene, and task-specific runtime fields. | PPO | `src/unilab/conf/ppo/task//.yaml` | | APPO | `src/unilab/conf/appo/task//.yaml` | | SAC / TD3 / FlashSAC | `src/unilab/conf//task//.yaml` | -| HORA distillation | `src/unilab/conf/hora_distill/task//.yaml` | Examples: diff --git a/docs/sphinx/source/en/2-user_guide/1-training/3-logging.md b/docs/sphinx/source/en/2-user_guide/1-training/3-logging.md index 49422b03f..a59971009 100644 --- a/docs/sphinx/source/en/2-user_guide/1-training/3-logging.md +++ b/docs/sphinx/source/en/2-user_guide/1-training/3-logging.md @@ -74,7 +74,7 @@ Normally the rows sum to `Iter Wall`, so their displayed percentages sum to abou hides applicable phases below a 1% threshold: a zero row is kept so it can be matched directly with TensorBoard / W&B. Algorithm-specific phases neither occupy terminal rows nor get persisted for other algorithms; for example, `Replay Stage` and -`Weight Publish` exist only for APPO / HORA-APPO. +`Weight Publish` exist only for APPO. ### Learner Main Timeline diff --git a/docs/sphinx/source/en/2-user_guide/2-algorithms/0-index.md b/docs/sphinx/source/en/2-user_guide/2-algorithms/0-index.md index debdb9e2f..8fbb1da48 100644 --- a/docs/sphinx/source/en/2-user_guide/2-algorithms/0-index.md +++ b/docs/sphinx/source/en/2-user_guide/2-algorithms/0-index.md @@ -11,7 +11,6 @@ lives, and which command shape selects it. For general flags, see | SAC | off-policy | `src/unilab/scripts/train_sac.py` | `src/unilab/conf/sac/config.yaml` | | TD3 | off-policy | `src/unilab/scripts/train_td3.py` | `src/unilab/conf/td3/config.yaml` | | FlashSAC | off-policy | `src/unilab/scripts/train_flashsac.py` | `src/unilab/conf/flashsac/config.yaml` | -| HORA | teacher/student distillation path | `scripts/train_hora_distill.py` | `src/unilab/conf/hora_distill/config.yaml` | ```{toctree} :hidden: @@ -21,5 +20,4 @@ lives, and which command shape selects it. For general flags, see 3-sac 4-td3 5-flash_sac -7-hora ``` diff --git a/docs/sphinx/source/en/2-user_guide/2-algorithms/7-hora.md b/docs/sphinx/source/en/2-user_guide/2-algorithms/7-hora.md deleted file mode 100644 index a2e11c19e..000000000 --- a/docs/sphinx/source/en/2-user_guide/2-algorithms/7-hora.md +++ /dev/null @@ -1,28 +0,0 @@ -# HORA - -The committed HORA path is the Sharpa in-hand teacher/student flow. Teacher -owners live under the PPO and APPO task trees through the `7-hora` profile for -`sharpa_inhand`; student distillation uses `scripts/train_hora_distill.py` and -`src/unilab/conf/hora_distill/task/sharpa_inhand/mujoco.yaml`. - -## Teacher - -```bash -uv run train --algo ppo --task sharpa_inhand --sim mujoco --profile hora -uv run train --algo appo --task sharpa_inhand --sim mujoco --profile hora training.no_play=true -``` - -The HORA PPO owner sets `algo.algo_log_name=hora_ppo` and resolves the runtime -through `uni_rl.algos.hora.rsl_rl:resolve_hora_ppo_runtime`. The APPO -variant sets `algo.algo_log_name=hora_appo`. - -## Student Distillation - -Student distillation is implemented by `scripts/train_hora_distill.py` and -configured by `src/unilab/conf/hora_distill/task/sharpa_inhand/mujoco.yaml`. The top-level -CLI does not currently declare a separate HORA distillation `--algo` route, so -the public CLI examples on this page stay on the teacher path above. - -Teacher checkpoint resolution is implemented in -`src/unilab/training/hora_distill_config.py`. The student log family is -`hora_distill`. diff --git a/docs/sphinx/source/en/2-user_guide/4-tasks/0-index.md b/docs/sphinx/source/en/2-user_guide/4-tasks/0-index.md index 83fe54bc6..a211f93e8 100644 --- a/docs/sphinx/source/en/2-user_guide/4-tasks/0-index.md +++ b/docs/sphinx/source/en/2-user_guide/4-tasks/0-index.md @@ -21,7 +21,7 @@ G1 motion tracking, flips, climbs, wall flips, and box tracking. :::{grid-item-card} Manipulation :link: 3-manipulation :link-type: doc -Allegro and Sharpa in-hand rotation and grasp generation. +Allegro in-hand rotation and grasp generation. ::: :::: diff --git a/docs/sphinx/source/en/2-user_guide/4-tasks/3-manipulation.md b/docs/sphinx/source/en/2-user_guide/4-tasks/3-manipulation.md index b9e55f62d..42423650f 100644 --- a/docs/sphinx/source/en/2-user_guide/4-tasks/3-manipulation.md +++ b/docs/sphinx/source/en/2-user_guide/4-tasks/3-manipulation.md @@ -5,19 +5,12 @@ Manipulation tasks live in `src/unilab/tasks/manipulation/`. ## In-Hand - `allegro_inhand` and `allegro_inhand_grasp` have MuJoCo and Motrix PPO owners. -- `sharpa_inhand`, `sharpa_inhand_grasp`, and the `hora` profile for - `sharpa_inhand` are MuJoCo owner paths in the current configs. ```bash uv run train --algo ppo --task allegro_inhand --sim mujoco uv run train --algo ppo --task allegro_inhand --sim motrix training.no_play=true -uv run train --algo ppo --task sharpa_inhand --sim mujoco --profile hora training.no_play=true ``` -HORA student distillation is configured by -`src/unilab/conf/hora_distill/task/sharpa_inhand/mujoco.yaml`; it is not currently exposed -as a separate top-level CLI route. - ## Platform Balancing `stewart_balance` is a 6-DOF parallel (Stewart) platform that balances a free @@ -34,5 +27,3 @@ training-stable under load. ```bash uv run train --algo ppo --task stewart_balance --sim motrix training.no_play=true ``` - -See {doc}`../8-manipulation/1-dexterous_inhand` for in-hand task notes. 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 171c2ccda..b559859ae 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 @@ -6,7 +6,7 @@ This page only describes the current domain randomization status of registered t Two DR declaration paths exist today: - **Manager-Based (Compatible) tasks**: reset / interval randomization is declared through Hydra `events:` manager terms in the owner YAML; reset-lifecycle events sample at reset, interval-lifecycle events perturb between steps. See the `events:` block of `src/unilab/conf/ppo/task/go1_joystick_flat/base.yaml` for an example. -- **Legacy provider path**: only the 2 Adapted families (`sharpa_inhand` / `sharpa_inhand_grasp`, including their appo / hora owners) still declare `env.domain_rand.*` configuration through a `DomainRandomizationProvider` + `DomainRandomizationManager`. +- **Task-level provider path**: custom tasks (including tasks hosted in external repos) may declare `env.domain_rand.*` configuration through a `DomainRandomizationProvider` + `DomainRandomizationManager`. No in-repo task currently uses this path. The unified entry point of the legacy provider path lives in `NpEnv._init_domain_randomization()` and `DomainRandomizationManager`: @@ -16,18 +16,18 @@ The unified entry point of the legacy provider path lives in `NpEnv._init_domain These three paths correspond to three lifecycle classes: -- **init-lifecycle DR**: items that change the model identity or model geometry; can only take effect during env/backend initialization and materialization, e.g. Sharpa-hand object `geom_size` scaling. +- **init-lifecycle DR**: items that change the model identity or model geometry; can only take effect during env/backend initialization and materialization, e.g. object `geom_size` scaling via model variants. - **reset-lifecycle DR**: items that do not change model identity, only change parameters or reset state within the same model, e.g. `base_mass_delta`, `base_com_offset`, `gravity`, `kp`, `kd`. - **interval-lifecycle DR**: external perturbations between steps, e.g. push. ## Status Conclusions -1. Manager-Based tasks do not register a DR provider; their reset/interval randomization consists of `events:` manager terms in the owner YAML, executed uniformly by the manager lifecycle. Only the frozen compatibility factories of the Adapted families still go through the `DomainRandomizationManager` unified entry point. -2. Adapted-family owners define a `domain_rand` config dataclass, a `DomainRandomizationProvider`, and a `ResetPlan`; Manager-Based owners declare reset behavior through Hydra command/event terms. G1 motion reset perturbations belong to `MotionCommandCfg`, while WBT adds `EventTermCfg` reset and interval terms. +1. Manager-Based tasks do not register a DR provider; their reset/interval randomization consists of `events:` manager terms in the owner YAML, executed uniformly by the manager lifecycle. Custom tasks on the provider path instead go through the `DomainRandomizationManager` unified entry point. +2. Provider-path owners define a `domain_rand` config dataclass, a `DomainRandomizationProvider`, and a `ResetPlan`; Manager-Based owners declare reset behavior through Hydra command/event terms. G1 motion reset perturbations belong to `MotionCommandCfg`, while WBT adds `EventTermCfg` reset and interval terms. 3. What is "unified" today is mainly the entry point and execution flow, not every randomization item itself. The legacy path's shared helper `build_common_reset_randomization()` currently generates `base_mass_delta`, `base_com_offset`, `gravity`, `kp`, `kd`. 4. `ResetRandomizationPayload` can already express `gravity`, `body_iquat`, `body_inertia`, `kp`, `kd`, and `MuJoCoBackend` has declared support. Whether these are actually used still depends on whether the task provider samples and dispatches them. 5. `MotrixBackend` currently supports `base_mass_delta`, `base_com_offset`, `kp`, `kd`, and interval push; and it requires all model actuators to be position actuators during initialization. -6. `geom_size` is not a reset-lifecycle field; Sharpa-hand object geom scale is handled by init-lifecycle model materialization. +6. `geom_size` is not a reset-lifecycle field; object geom scale is handled by init-lifecycle model materialization. ## Uniformity Assessment Table @@ -41,8 +41,6 @@ These three paths correspond to three lifecycle classes: | `G1WBTObs` | Hydra `events:` terms | Yes: same motion command + Hydra `EventTermCfg` | motion reset plus mass/COM/PD/friction/encoder-bias events | interval velocity kick | `motion_tracking/g1/manager_terms.py` | | `AllegroInhandRotation` | Hydra `events:` terms | Yes: Hydra `EventTermCfg` + Manager-Based reset term | entity-scoped hand/ball reset | none | `allegro_inhand/manager_terms.py` | | `AllegroInhandRotationGrasp` | Hydra `events:` terms | Yes: reuses the rotation reset event + `RecorderTermCfg` | noisy hand reset + grasp collection | none | `allegro_inhand/grasp_gen.py` | -| `SharpaInhandRotation` | legacy provider | Yes: `InitRandomizationPlan + ResetPlan + IntervalRandomizationPlan` | grasp cache sampling + common payload | object `body_force` | `sharpa_inhand/rotation.py` | -| `SharpaInhandRotationGrasp` | legacy provider | Yes: reuses the Sharpa rotation provider and overrides reset sampling | grasp collection reset + common payload | none | `sharpa_inhand/grasp_gen.py` | ## Per-task Domain Randomization List @@ -56,8 +54,6 @@ These three paths correspond to three lifecycle classes: | `G1WBTObs` | Same motion reset plus base mass, base COM, PD gain, foot friction, and encoder-bias event terms | `push_by_setting_velocity` | The WBT owner explicitly enables all listed event terms; unsupported capabilities raise rather than fall back | | `AllegroInhandRotation` | Entity-scoped hand/ball reset; an explicitly configured grasp cache is sampled, otherwise `null` explicitly selects the model home pose; optional `joint_noise`, `ball_velocity_noise`, and `ball_z_offset` | none | owner YAML explicitly selects the home pose and zero reset noise; a configured missing or malformed cache fails closed | | `AllegroInhandRotationGrasp` | Reuses the rotation reset with `joint_noise=0.25`; Manager-Based termination checks fingertip distance, contact count, and ball height; recorder stores successful timeout rows | none | generates the 50k-row Allegro grasp cache and raises `RunComplete` after a successful save | -| `SharpaInhandRotation` | grasp cache bucketed sampling by `scale_ids`; object pose / quat reset; optional common reset randomization payload (incl. `gravity`) | object `body_force` direct force disturbance | `domain_rand.scale_list` defaults come from the owner YAML; under MuJoCo, object geom scale is materialized during init; common payload disabled by default; object force enabled by default via the Sharpa owner YAML | -| `SharpaInhandRotationGrasp` | hand pose reset; object pose / quat reset; collects successful grasps and stores them bucketed by `scale_ids`; optional `base_mass_delta`; optional `base_com_offset`; optional `gravity` | none | Used by default to generate the Sharpa grasp cache; cache filename includes the single scale value; common payload disabled by default | ## Current Unified DR Capabilities and Boundaries @@ -70,9 +66,9 @@ and `DomainRandomizationManager`: - The manager uniformly performs capability validation - The backend is uniformly responsible for actually applying the randomization payload -So from an execution-path perspective, the Adapted families still on this path -are unified; Manager-Based tasks instead execute the `events:` terms declared -in the owner YAML through the manager lifecycle. +So from an execution-path perspective, provider-path tasks are unified; +Manager-Based tasks instead execute the `events:` terms declared in the owner +YAML through the manager lifecycle. ### 2. The Shared Helpers Are Still Narrow @@ -82,11 +78,10 @@ The legacy path's `dr_utils.py` builds and validates common reset payloads: This means: -- The Sharpa families still on the legacy provider path sample their - task-specific state directly inside each provider +- Provider-path tasks sample their task-specific state directly inside each provider - `G1MotionTracking`'s pose / velocity / joint noise is owned by its manager command - Allegro's grasp / object initial state sampling is entirely task-specific logic -- Sharpa's `geom_size` scale is init-lifecycle model materialization and is not part of the reset common payload +- `geom_size` scale is init-lifecycle model materialization and is not part of the reset common payload So today's "uniformity" is more about the contract and the calling convention than "all tasks share the same set of randomization-item schemas". @@ -110,7 +105,7 @@ Backend capability today: Notes: - The current `IntervalRandomizationPlan` supports `push_perturbation_limit`, `body_linear_velocity_delta`, `body_angular_velocity_delta`, `body_force`, and `body_torque`; among these, `body_force`/`body_torque` express hot-path direct external-wrench perturbations without exposing the backend-private `xfrc_applied` details. -- The current MuJoCo backend's interval push and interval body force are both dispatched through `xfrc_applied`; the Sharpa-hand object disturbance has been switched to direct force disturbance. +- The current MuJoCo backend's interval push and interval body force are both dispatched through `xfrc_applied`. - The Motrix backend currently still does not support direct body-force disturbance, so such owner configs must continue to be explicitly disabled. But on the task side, the current reality is: not every provider constructs these fields. The backend contract is the capability boundary; whether the task config and provider dispatch a payload is what determines whether a given task actually enables the corresponding DR item. @@ -124,9 +119,8 @@ But on the task side, the current reality is: not every provider constructs thes - Lifecycle: only sampled and written at reset; the env retains that gravity until the next reset re-samples it. - Backend: currently in UniLab, only the MuJoCo backend declares support for this reset term; the Motrix backend does not. Some tasks filter it by capability and skip it; others raise an error in the validate stage. -The config entry lives under `env.domain_rand` in Sharpa owners on the legacy -provider path, such as `sharpa_inhand_grasp`; Manager-Based tasks have no -`env.domain_rand`: +The config entry lives under `env.domain_rand` in provider-path task owners; +Manager-Based tasks have no `env.domain_rand`: ```yaml env: @@ -143,21 +137,7 @@ Field semantics: - `gravity_range`: a `(2, 3)`-shaped per-dimension sampling range; the first and second rows give the upper and lower bounds of each component. - On each reset, each dimension is uniformly sampled within `[min(row0, row1), max(row0, row1)]`. The direction is not automatically normalized, and the gravity norm is not fixed. -If you only want to randomize the magnitude while keeping the vertical-down direction, only open up the `z` component: - -```bash -uv run train --algo ppo --task sharpa_inhand_grasp --sim mujoco \ - env.domain_rand.randomize_gravity=true \ - 'env.domain_rand.gravity_range=[[0.0,0.0,-10.5],[0.0,0.0,-8.5]]' -``` - -If you want to randomize both direction and magnitude, open up `x/y/z`: - -```bash -uv run train --algo ppo --task sharpa_inhand_grasp --sim mujoco \ - env.domain_rand.randomize_gravity=true \ - 'env.domain_rand.gravity_range=[[-0.3,-0.3,-10.5],[0.3,0.3,-8.5]]' -``` +If you only want to randomize the magnitude while keeping the vertical-down direction, only open up the `z` component; to randomize both direction and magnitude, open up `x/y/z`. Enable it from the CLI with `env.domain_rand.randomize_gravity=true` and a `env.domain_rand.gravity_range=[...]` override on a provider-path task owner. Notes: @@ -198,100 +178,9 @@ The reason is that `geom_size` changes model geometry and model identity; the co This boundary exists to honor the cold-path asset/model-metadata access principle: `step()`, `reset()`, and hot-path DR do not parse XML, do not read assets, and do not branch at runtime based on asset metadata. -## Sharpa-hand Object Geom Scale Usage - -Sharpa-hand is the current example task for `geom_size` init-lifecycle DR in the repo. Related task configs: - -- `src/unilab/conf/ppo/task/sharpa_inhand/mujoco.yaml` -- `src/unilab/conf/ppo/task/sharpa_inhand_grasp/mujoco.yaml` - -### 1. Config Entry - -Sharpa's scale configuration lives in the env owner YAML at `env.domain_rand.scale_list`: - -```yaml -env: - object_body_name: object - object_geom_name: object - domain_rand: - scale_list: [0.5, 0.6, 0.7, 0.8] -``` - -Field semantics: - -- `object_body_name`: the object body name; used to locate the object body during reset / observation, not the target field for scaling. -- `object_geom_name`: the MuJoCo geom name to scale; defaults to `object`. -- `domain_rand.scale_list`: explicit scale list; each value must be greater than 0. -- The order of `domain_rand.scale_list` is the `scale_id` order. -- The length of `domain_rand.scale_list` is the number of model variants. - -Each env is statically assigned a `scale_id`. The current assignment rule is contiguous bucket assignment; when `algo.num_envs` is not divisible by `num_scales`, the first few scale buckets get one extra env: - -```bash -uv run train --algo ppo --task sharpa_inhand --sim mujoco 'env.domain_rand.scale_list=[0.5,0.6,0.7,0.8]' algo.num_envs=4096 -``` - -If `algo.num_envs=4096` and `num_scales=4`, then every 1024 envs use the same scale bucket. - -### 2. MuJoCo Materialization Behavior - -How the MuJoCo backend applies it: - -1. The env/provider builds `ModelVariantSpec` based on `scale_list` during init. -2. The backend uses `MjSpec` to read the model and modify the `size` of the geom corresponding to `object_geom_name`. -3. Each scale compiles a scale-specific `MjModel`. -4. The first time a physics pool is needed, the env-to-model assignment is expanded into a model sequence of length `num_envs`, and then `BatchEnvPool` is constructed. - -Therefore, `domain_rand.scale_list` only takes effect during env/backend initialization. Changing `env.domain_rand.scale_list` after env creation does not change the already-materialized model pool. - -This flow has three important boundaries: - -- `BatchEnvPool` is lazily constructed; the normal path does not first construct a pool for the default model and then rebuild it for `scale_list`. -- Compilation of multiple model variants is done in chunks using process-based parallelism; do not compile in a Python thread, and do not serially compile `num_envs` models in an upper-level for loop. -- Workers compile variants with `MjSpec` and save `.mjb`; the parent process only loads `MjModel.from_binary_path(...)` by `.mjb` path. Do not transmit modified model objects or model bytes back via IPC. - -### 3. Grasp Cache and Scale Buckets - -The Sharpa rotation task samples from multiple single-scale grasp caches by `scale_ids`: - -- The cache filename defaults are jointly determined by `grasp_cache_path` and a single scale value. -- `scale_list: [0.5, 0.6, 0.7, 0.8]` by default corresponds to `caches/sharpa_grasp_linspace_0.5.npy`, `caches/sharpa_grasp_linspace_0.6.npy`, `caches/sharpa_grasp_linspace_0.7.npy`, `caches/sharpa_grasp_linspace_0.8.npy`. -- At rotation startup all cache files for `scale_list` are checked; if any is missing, it errors. -- Each scale bucket only samples from the cache file of its own scale, avoiding mixing grasp initial states across different object scales. - -Default-scale caches are hosted on Hugging Face (`unilabsim/unilab-caches`) and are downloaded automatically into `src/unilab/assets/caches/` on first rotation training, so no manual collection is needed for the standard `scale_list`. - -To collect caches for custom scales not on HF — or to regenerate them locally — run the grasp collection task once per scale. Generated files land under `src/unilab/assets/caches/`, the same location HF downloads to, so subsequent rotation training auto-resolves them. Regeneration is **slow**. - -The helper script collects each scale sequentially: - -```bash -./scripts/sharpa_collect_grasps.sh 0.5 0.6 0.7 0.8 -``` - -Equivalent per-scale invocations: `uv run train --algo ppo --task sharpa_inhand_grasp --sim mujoco 'env.domain_rand.scale_list=[0.5]' algo.num_envs=4096` (repeat for `[0.6]`, `[0.7]`, `[0.8]`, …). - -Then train rotation with the same `scale_list`: - -```bash -uv run train --algo ppo --task sharpa_inhand --sim mujoco 'env.domain_rand.scale_list=[0.5,0.6,0.7,0.8]' algo.num_envs=4096 -``` - -### 4. Boundaries and Caveats - -- `geom_size` is not a reset DR field and must not be written into `ResetPlan.randomization`. -- `BatchEnvPool.reset(..., randomization=...)` currently does not support `geom_size`. -- `geom_size` scale is only materialized under the MuJoCo backend; the Motrix backend currently does not produce a multi-model pool from `scale_list`. -- The length of `scale_list` is the number of model variants, not the number of resamples per reset. -- Each env's `scale_id` is statically assigned during init and does not change at reset. -- When scaling out, scale out the number of model variants; do not compile one model per env according to `num_envs`. Multiple envs share the same `MjModel` corresponding to the same scale bucket. -- The hot path must not read XML, parse assets, or use `getattr` / `hasattr` to probe backend-private capability to decide scaling behavior. -- When extending to other shape DR, prefer to reuse the init-lifecycle contract; do not stuff shape fields into the reset payload. - ## Related Tasks - {doc}`G1 Motion Tracking <../4-tasks/2-motion_tracking>`: confirm motion assets and replay first before enabling DR. -- {doc}`Sharpa In-Hand <../8-manipulation/1-dexterous_inhand>`: the scale / grasp-cache / DR boundary is sensitive. - {doc}`Go2 Rough Terrain <../4-tasks/1-locomotion>`: common items are mass, COM, friction, and push. For configuration examples, see {doc}`1-configuration`. For the developer diff --git a/docs/sphinx/source/en/2-user_guide/5-domain_randomization/1-configuration.md b/docs/sphinx/source/en/2-user_guide/5-domain_randomization/1-configuration.md index 78dca0951..e0f2d0874 100644 --- a/docs/sphinx/source/en/2-user_guide/5-domain_randomization/1-configuration.md +++ b/docs/sphinx/source/en/2-user_guide/5-domain_randomization/1-configuration.md @@ -9,14 +9,8 @@ Two declaration paths exist today: - Manager-Based (Compatible) tasks declare reset / interval randomization through Hydra `events:` manager terms in the owner YAML, for example `src/unilab/conf/ppo/task/go1_joystick_flat/base.yaml`. -- The Sharpa Adapted families and their hora / appo owners configure legacy - provider fields under `env.domain_rand`. - -```bash -uv run train --algo ppo --task sharpa_inhand_grasp --sim mujoco \ - env.domain_rand.randomize_gravity=true \ - 'env.domain_rand.gravity_range=[[0.0,0.0,-10.5],[0.0,0.0,-8.5]]' -``` +- Tasks may also attach a task-level provider and configure legacy provider + fields under `env.domain_rand`; no in-repo task currently uses this path. Common lifecycle boundaries: @@ -37,13 +31,8 @@ the env owners, and backend capabilities are declared through Use `--sim mujoco` when enabling gravity reset randomization; Motrix does not advertise the same gravity capability in the current backend. This item is only -available on the legacy provider path (Adapted-family owners). - -```bash -uv run train --algo ppo --task sharpa_inhand_grasp --sim mujoco \ - env.domain_rand.randomize_gravity=true \ - 'env.domain_rand.gravity_range=[[0.0,0.0,-10.5],[0.0,0.0,-8.5]]' -``` +available on the task-level provider path (`env.domain_rand.randomize_gravity` +and `env.domain_rand.gravity_range`), which no in-repo task currently uses. ## Interval Push @@ -62,8 +51,6 @@ Keep ranges in the task owner YAML when they are part of the task contract. For example, the rough quadruped family's base mass, center-of-mass, kp/kd, and push randomization are declared as event terms in the shared base `src/unilab/conf/ppo/task/quadruped_joystick_rough/base.yaml` (the `go2_joystick_rough` -backend owners compose it through Hydra defaults), while -`src/unilab/conf/ppo/task/sharpa_inhand/mujoco.yaml` configures object scale, friction, and -force disturbance for Sharpa. +backend owners compose it through Hydra defaults). For the full current inventory, see {doc}`0-index`. diff --git a/docs/sphinx/source/en/2-user_guide/5-domain_randomization/2-writing_providers.md b/docs/sphinx/source/en/2-user_guide/5-domain_randomization/2-writing_providers.md index 1c35d55bc..3fea0fbab 100644 --- a/docs/sphinx/source/en/2-user_guide/5-domain_randomization/2-writing_providers.md +++ b/docs/sphinx/source/en/2-user_guide/5-domain_randomization/2-writing_providers.md @@ -1,11 +1,10 @@ # Writing Providers -This page describes the legacy provider path: only the 2 Adapted families -(`sharpa_inhand` / `sharpa_inhand_grasp`) still declare -domain randomization through a task-level `DomainRandomizationProvider`. -Migrated Manager-Based tasks do not write providers; they declare randomization -through Hydra `events:` manager terms in the owner YAML (see {doc}`0-index` -and {doc}`1-configuration`). +This page describes the task-level provider path: custom tasks (including +tasks hosted in external repos) may declare domain randomization through a +task-level `DomainRandomizationProvider`. Manager-Based tasks do not write +providers; they declare randomization through Hydra `events:` manager terms in +the owner YAML (see {doc}`0-index` and {doc}`1-configuration`). Task-level domain randomization providers live with the task env owner. They sample task-specific state and return plans consumed by @@ -59,10 +58,10 @@ manager lives in `src/unilab/dr/manager.py`. ## Evidence -Representative provider implementations are in (all on the Adapted-family -compatibility path): +The provider interface and manager live in: -- `src/unilab/tasks/manipulation/sharpa_inhand/rotation.py` +- `src/unilab/dr/provider.py` +- `src/unilab/dr/manager.py` Developer contract details are in {doc}`../../4-developer_guide/2-contracts/4-dr_contract`. diff --git a/docs/sphinx/source/en/2-user_guide/8-manipulation/0-index.md b/docs/sphinx/source/en/2-user_guide/8-manipulation/0-index.md deleted file mode 100644 index f11b3938d..000000000 --- a/docs/sphinx/source/en/2-user_guide/8-manipulation/0-index.md +++ /dev/null @@ -1,21 +0,0 @@ -# Manipulation Notes - -Task-specific notes for the manipulation paths that need more operational -detail than the category overview in {doc}`../4-tasks/3-manipulation`. - -::::{grid} 1 1 2 2 -:gutter: 3 - -:::{grid-item-card} Dexterous in-hand -:link: 1-dexterous_inhand -:link-type: doc -Allegro and Sharpa owner YAMLs, grasp caches, and train commands. -::: - -:::: - -```{toctree} -:hidden: - -1-dexterous_inhand -``` diff --git a/docs/sphinx/source/en/2-user_guide/8-manipulation/1-dexterous_inhand.md b/docs/sphinx/source/en/2-user_guide/8-manipulation/1-dexterous_inhand.md deleted file mode 100644 index b3a2af32f..000000000 --- a/docs/sphinx/source/en/2-user_guide/8-manipulation/1-dexterous_inhand.md +++ /dev/null @@ -1,145 +0,0 @@ -# Dexterous In-Hand Manipulation - -This page covers the checked-in Allegro and Sharpa in-hand manipulation paths. -Select backends with `--task` and `--sim`; do not override -`training.sim_backend` alone. The owner YAMLs remain the internal evidence for -which combinations are configured. - -## Allegro - -Allegro rotation uses the registered env `AllegroInhandRotation`. The rotation -owner is `allegro_inhand`, and grasp-cache generation uses -`allegro_inhand_grasp`. - -Owner evidence: - -- `src/unilab/conf/ppo/task/allegro_inhand/mujoco.yaml` -- `src/unilab/conf/ppo/task/allegro_inhand/motrix.yaml` -- `src/unilab/conf/ppo/task/allegro_inhand_grasp/mujoco.yaml` -- `src/unilab/conf/ppo/task/allegro_inhand_grasp/motrix.yaml` -- `src/unilab/conf/appo/task/allegro_inhand/mujoco.yaml` -- `src/unilab/conf/appo/task/allegro_inhand/motrix.yaml` - -The typical flow is two stages: first generate a grasp cache, then train the -rotation policy. - -```bash -uv run train --algo ppo --task allegro_inhand_grasp --sim mujoco training.no_play=true -uv run train --algo ppo --task allegro_inhand --sim mujoco training.no_play=true -``` - -Motrix owner YAMLs also exist for the PPO Allegro paths: - -```bash -uv run train --algo ppo --task allegro_inhand_grasp --sim motrix training.no_play=true -uv run train --algo ppo --task allegro_inhand --sim motrix training.no_play=true -``` - -The rotation owner defaults to the grasp cache at `caches/allegro_grasp_50k.npy`. -To use a custom cache, override `env.grasp_cache_path`: - -```bash -uv run train --algo ppo --task allegro_inhand --sim mujoco \ - env.grasp_cache_path=caches/my_allegro_grasp.npy -``` - -Replay a trained checkpoint with `eval` (`--load-run -1` picks the latest run): - -```bash -uv run eval --algo ppo --task allegro_inhand --sim mujoco --load-run -1 -uv run eval --algo appo --task allegro_inhand --sim mujoco --load-run -1 -``` - -## Sharpa - -Sharpa rotation uses the registered env `SharpaInhandRotation`. Current checked -in training paths are MuJoCo owner paths. - -Owner evidence: - -- `src/unilab/conf/ppo/task/sharpa_inhand/mujoco.yaml` -- `src/unilab/conf/ppo/task/sharpa_inhand/mujoco_hora.yaml` -- `src/unilab/conf/ppo/task/sharpa_inhand_grasp/mujoco.yaml` -- `src/unilab/conf/appo/task/sharpa_inhand/mujoco.yaml` -- `src/unilab/conf/appo/task/sharpa_inhand/mujoco_hora.yaml` -- `src/unilab/conf/hora_distill/task/sharpa_inhand/mujoco.yaml` - -The full HORA path is three stages: - -1. Generate the grasp cache. -2. Train the teacher policy. -3. Distill a student policy when needed. - -The full HORA teacher/student path is MuJoCo-owner-primary. The Motrix path -currently covers only phase-1 PPO rotation and grasp-cache collection; it is not -a full HORA capability-equivalent path. - -### Grasp cache and scale - -The default caches are hosted on Hugging Face (`unilabsim/unilab-caches`) and are -downloaded automatically into `src/unilab/assets/caches/` on first training, so -no manual step is needed. - -To collect caches for custom scales not on HF — or to regenerate them locally — -run the grasp task once per scale (cache files are named `_.npy`). -Generated files land under `src/unilab/assets/caches/`, the same location HF -downloads to, so subsequent training auto-resolves them without further -configuration. Regeneration is **slow**. - -The helper script collects each scale sequentially: - -```bash -bash scripts/sharpa_collect_grasps.sh 0.8 1.0 1.2 -``` - -Equivalent per-scale invocations: `uv run train --algo ppo --task sharpa_inhand_grasp --sim mujoco 'env.domain_rand.scale_list=[0.8]' training.no_play=true` (repeat for `[1.0]`, `[1.2]`, …). - -Motrix can also collect a grasp cache (phase-1 scope only): - -```bash -uv run train --algo ppo --task sharpa_inhand_grasp --sim motrix \ - 'env.domain_rand.scale_list=[1.0]' \ - env.grasp_collection_target=1000 \ - training.no_play=true -``` - -To use a custom cache prefix, override `env.grasp_cache_path`: - -```bash -uv run train --algo ppo --task sharpa_inhand --sim mujoco \ - env.grasp_cache_path=caches/my_sharpa_grasp_cache -``` - -### Teacher and student - -Train the HORA teacher with the `hora` profile (PPO or APPO): - -```bash -uv run train --algo ppo --task sharpa_inhand --sim mujoco --profile hora training.no_play=true -uv run train --algo appo --task sharpa_inhand --sim mujoco --profile hora training.no_play=true -``` - -Replay a teacher run with `eval --profile hora --load-run -1`: - -```bash -uv run eval --algo ppo --task sharpa_inhand --sim mujoco --profile hora --load-run -1 -uv run eval --algo appo --task sharpa_inhand --sim mujoco --profile hora --load-run -1 -``` - -Student distillation is configured by -`src/unilab/conf/hora_distill/task/sharpa_inhand/mujoco.yaml` and implemented by -`scripts/train_hora_distill.py`; the top-level CLI does not currently expose a -separate HORA distillation route (it is not in the CLI `SUPPORTED_ALGOS`). To -distill from an APPO teacher, set `teacher.algo_family=appo` in that low-level -config. - -Common log directories: - -- `logs/hora_ppo/SharpaInhandRotation/` -- `logs/hora_appo/SharpaInhandRotation/` -- `logs/hora_distill/SharpaInhandRotation/` - -The scale / grasp-cache / DR boundary is sensitive here; see -{doc}`../5-domain_randomization/0-index` for the lifecycle rules. - -For the category-level task page, see {doc}`../4-tasks/3-manipulation`. diff --git a/docs/sphinx/source/en/3-deployment/1-sim_to_real/4-allegro_inhand.md b/docs/sphinx/source/en/3-deployment/1-sim_to_real/4-allegro_inhand.md index 484f9febb..af15f12e1 100644 --- a/docs/sphinx/source/en/3-deployment/1-sim_to_real/4-allegro_inhand.md +++ b/docs/sphinx/source/en/3-deployment/1-sim_to_real/4-allegro_inhand.md @@ -1,7 +1,7 @@ -# Allegro / Sharpa In-Hand Manipulation Deployment +# Allegro In-Hand Manipulation Deployment -Cube reorientation on a 16-DoF Allegro hand or 17-DoF Sharpa hand. UniLab -trains these tactile-free — observations are joint state + cube pose only. +Cube reorientation on a 16-DoF Allegro hand. UniLab +trains it tactile-free — observations are joint state + cube pose only. ## What makes this hard @@ -26,10 +26,10 @@ bring-up; see {doc}`6-domain_randomization`. - Dim - Source on hardware * - Joint positions - - 16 (Allegro) / 17 (Sharpa) + - 16 - encoder * - Joint velocities - - 16 / 17 + - 16 - encoder differentiated, low-pass * - Cube pose (world) - 7 @@ -41,7 +41,7 @@ bring-up; see {doc}`6-domain_randomization`. - 4 - command * - Previous action - - 16 / 17 + - 16 - last policy output ``` @@ -55,7 +55,7 @@ agree on the observation timing before hardware deployment. See ## Grasp generator -Both `4-allegro_inhand` and `sharpa_inhand` envs ship a **grasp generator** +The `allegro_inhand` env ships a **grasp generator** that samples plausible initial hand configurations. The hardware-side equivalent is the operator placing the cube in the hand — verify your distribution of starting configurations matches the trained env's grasp @@ -69,8 +69,7 @@ to the grasp generator**, retrain, and try again. The manipulation envs map policy actions to joint position targets through their task control config. Allegro owns this declaration in -`src/unilab/conf/ppo/task/allegro_inhand/base.yaml` and its Manager-Based action term; -Sharpa currently owns it in `src/unilab/tasks/manipulation/sharpa_inhand/base.py`. +`src/unilab/conf/ppo/task/allegro_inhand/base.yaml` and its Manager-Based action term. The deploy controller must use the same joint order, action scale, and limit policy. @@ -85,4 +84,3 @@ alert the operator. - {doc}`5-onnx_runtime` - {doc}`6-domain_randomization` -- {doc}`../../2-user_guide/8-manipulation/1-dexterous_inhand` diff --git a/docs/sphinx/source/en/3-deployment/1-sim_to_real/6-domain_randomization.md b/docs/sphinx/source/en/3-deployment/1-sim_to_real/6-domain_randomization.md index 2618f0e9f..bb49eb95f 100644 --- a/docs/sphinx/source/en/3-deployment/1-sim_to_real/6-domain_randomization.md +++ b/docs/sphinx/source/en/3-deployment/1-sim_to_real/6-domain_randomization.md @@ -46,11 +46,10 @@ Manager-Based tasks declare reset and interval randomization through `env.events` in their owner YAML, executed by the manager lifecycle. See `src/unilab/conf/ppo/task/quadruped_joystick_rough/base.yaml`. -The Sharpa Adapted tasks still attach a task provider to -`src/unilab/dr/manager.py`. The current example is -`SharpaInhandRotationDRProvider` in -`src/unilab/tasks/manipulation/sharpa_inhand/rotation.py`. The capability -boundary for both paths is described in +Tasks may also attach a task-level provider (the +`DomainRandomizationProvider` interface in `src/unilab/dr/provider.py`) to +`src/unilab/dr/manager.py`; no in-repo task currently uses this path. The +capability boundary for both paths is described in {doc}`../../4-developer_guide/2-contracts/4-dr_contract`. ## Recipe: starting ranges diff --git a/docs/sphinx/source/en/3-deployment/1-sim_to_real/8-latency_budget.md b/docs/sphinx/source/en/3-deployment/1-sim_to_real/8-latency_budget.md index ef32bda90..9e94878fe 100644 --- a/docs/sphinx/source/en/3-deployment/1-sim_to_real/8-latency_budget.md +++ b/docs/sphinx/source/en/3-deployment/1-sim_to_real/8-latency_budget.md @@ -10,7 +10,6 @@ budgets as robot-specific measurements, not UniLab defaults. | --- | --- | --- | | One-step action delay | Manager action term `simulate_action_latency` declarations in task owners | Executes the previous action instead of the current action. | | G1 WBT observation history | Per-term `history_length` in `src/unilab/conf/sac/task/g1_wbt_obs/mujoco.yaml` | Per-term history for `base_ang_vel`, `joint_pos`, `joint_vel`, and `actions`. | -| Sharpa tactile contact latency | `domain_rand.contact_latency` in Sharpa in-hand configs | Keeps previous tactile contact values for sampled contact channels. | | Obs history ordering guard | `tests/scripts/test_obs_alignment_g1_wbt.py` | Asserts per-term oldest-first flatten for the G1 WBT actor obs. | ## Action Latency diff --git a/docs/sphinx/source/en/3-deployment/1-sim_to_real/9-troubleshooting.md b/docs/sphinx/source/en/3-deployment/1-sim_to_real/9-troubleshooting.md index 5020148c1..89bd1d9dd 100644 --- a/docs/sphinx/source/en/3-deployment/1-sim_to_real/9-troubleshooting.md +++ b/docs/sphinx/source/en/3-deployment/1-sim_to_real/9-troubleshooting.md @@ -35,7 +35,7 @@ Almost always one of: against the owner's `env.observations.actor.terms` — term order first, then per-term history ordering. -## Cube drops in Allegro / Sharpa inhand +## Cube drops in Allegro inhand | Likely cause | Check | Fix | |---|---|---| diff --git a/docs/sphinx/source/en/3-deployment/3-framework_migration/1-from_isaac_lab.md b/docs/sphinx/source/en/3-deployment/3-framework_migration/1-from_isaac_lab.md index 48d208c80..7ff1eaf92 100644 --- a/docs/sphinx/source/en/3-deployment/3-framework_migration/1-from_isaac_lab.md +++ b/docs/sphinx/source/en/3-deployment/3-framework_migration/1-from_isaac_lab.md @@ -177,9 +177,6 @@ requires an explicit migration decision. - 36 tasks are **Compatible** (`target=complete`): the Hydra owner YAML materializes the canonical NumPy Manager-Based runtime. -- 2 tasks are **Adapted** (`target=compatibility`): `SharpaInhandRotation` - and `SharpaInhandRotationGrasp` keep tactile/contact/cache behavior behind one - frozen compatibility factory each; they migrate only when the formal capability exists. ## Repository evidence diff --git a/docs/sphinx/source/en/4-developer_guide/0-index.md b/docs/sphinx/source/en/4-developer_guide/0-index.md index 274606565..7e7dbfbe0 100644 --- a/docs/sphinx/source/en/4-developer_guide/0-index.md +++ b/docs/sphinx/source/en/4-developer_guide/0-index.md @@ -93,5 +93,4 @@ Extend terrain generation while keeping asset access on cold paths. 5-contributing_workflow 6-agent_quick_reference 7-motion_assets -8-motrix_contact_sensor ``` diff --git a/docs/sphinx/source/en/4-developer_guide/2-contracts/3-task_owner.md b/docs/sphinx/source/en/4-developer_guide/2-contracts/3-task_owner.md index 619447775..60b7d45db 100644 --- a/docs/sphinx/source/en/4-developer_guide/2-contracts/3-task_owner.md +++ b/docs/sphinx/source/en/4-developer_guide/2-contracts/3-task_owner.md @@ -10,8 +10,6 @@ contract is recorded in `src/unilab/conf/{ppo,appo}/task//.yaml`. - Off-policy algorithms (SAC / TD3 / FlashSAC) each have their own config tree: `src/unilab/conf//task//.yaml`. -- `src/unilab/conf/hora_distill/` follows the same owner-YAML identity rule for its - supported tasks. ## Required Semantics diff --git a/docs/sphinx/source/en/4-developer_guide/2-contracts/4-dr_contract.md b/docs/sphinx/source/en/4-developer_guide/2-contracts/4-dr_contract.md index 895a2654e..349c8bf12 100644 --- a/docs/sphinx/source/en/4-developer_guide/2-contracts/4-dr_contract.md +++ b/docs/sphinx/source/en/4-developer_guide/2-contracts/4-dr_contract.md @@ -136,4 +136,4 @@ payloads. `src/unilab/dr/__init__.py` - DR manager: `src/unilab/dr/manager.py` - Backend interface: `unisim.backend.base` -- Example provider: `src/unilab/tasks/manipulation/sharpa_inhand/rotation.py` +- Provider interface: `src/unilab/dr/provider.py` diff --git a/docs/sphinx/source/en/4-developer_guide/3-extending/3-new_algorithm.md b/docs/sphinx/source/en/4-developer_guide/3-extending/3-new_algorithm.md index 38e2c4601..029b52223 100644 --- a/docs/sphinx/source/en/4-developer_guide/3-extending/3-new_algorithm.md +++ b/docs/sphinx/source/en/4-developer_guide/3-extending/3-new_algorithm.md @@ -18,7 +18,8 @@ Change no code — only adjust the Hydra config under `src/unilab/conf//`: `src/unilab/conf//task//.yaml`; - Swapping policy / algorithm implementation classes works through the `class_name` dotted path in the owner YAML (in-repo example: - `uni_rl.algos.hora:HoraActorModel`), with no new code path required. + `uni_rl.algos.rsl_rl_ppo:FinalObservationAwarePPO` in + `src/unilab/conf/ppo/config.yaml`), with no new code path required. ### 2. `runtime_resolver`: Algorithm Code In Your Own Repository @@ -43,11 +44,6 @@ Contract: `uni_rl.algos.rsl_rl_runtime`); the dotted path may point at any importable module. -In-repo example: `src/unilab/conf/appo/task/sharpa_inhand/mujoco_hora.yaml` -points at `unilab.scripts.play_hora_appo:resolve_hora_appo_runtime`; the HORA -SAC / PPO variants point at `uni_rl.algos.hora.sac:resolve_hora_sac_runtime` -and `uni_rl.algos.hora.rsl_rl:resolve_hora_ppo_runtime` respectively. - ### 3. Fork unilab_rl: Modify `uni_rl/algos/` Only when you need to change the shared runner / learner / collector @@ -79,9 +75,8 @@ convention-discovered). Notes: -- Config trees that have a conf directory but no entrypoint script (such as - `hora_distill`) are not routable — they are not standalone - CLI algos. +- Config trees that have a conf directory but no entrypoint script are not + routable — they are not standalone CLI algos. - The special script-name mappings for built-in algorithms are preserved: `ppo` → `train_rsl_rl.py`, `appo` → `train_appo.py`. - The dataclasses in `src/unilab/structured_configs.py` are an **optional** diff --git a/docs/sphinx/source/en/4-developer_guide/6-agent_quick_reference.md b/docs/sphinx/source/en/4-developer_guide/6-agent_quick_reference.md index d9b19f5b3..6f9e242ab 100644 --- a/docs/sphinx/source/en/4-developer_guide/6-agent_quick_reference.md +++ b/docs/sphinx/source/en/4-developer_guide/6-agent_quick_reference.md @@ -13,7 +13,6 @@ repo facts. - APPO entrypoint: `src/unilab/scripts/train_appo.py` - SAC / TD3 / FlashSAC entrypoints: `src/unilab/scripts/train_sac.py` / `src/unilab/scripts/train_td3.py` / `src/unilab/scripts/train_flashsac.py` -- HORA distillation entrypoint: `scripts/train_hora_distill.py` ## Contracts To Keep In Mind diff --git a/docs/sphinx/source/en/4-developer_guide/7-motion_assets.md b/docs/sphinx/source/en/4-developer_guide/7-motion_assets.md index e32ed5838..f0368e4ed 100644 --- a/docs/sphinx/source/en/4-developer_guide/7-motion_assets.md +++ b/docs/sphinx/source/en/4-developer_guide/7-motion_assets.md @@ -88,7 +88,7 @@ Robot binary meshes and textures (for example `.STL`, `.obj`, and `.png`) are externalized the same way, on the Hugging Face dataset repo [unilabsim/unilab-robots](https://huggingface.co/datasets/unilabsim/unilab-robots). The registered robots are a2, allegro_hand, g1, go2, -sharpa_wave, and x2 (`ROBOT_ASSET_SPECS` in `src/unilab/assets/hub.py`). +and x2 (`ROBOT_ASSET_SPECS` in `src/unilab/assets/hub.py`). Their mesh/texture directories download lazily on first use and land under their original paths (for example `src/unilab/assets/robots/g1/assets/` and `robots/g1/textures/` for G1), so the original relative XML paths remain diff --git a/docs/sphinx/source/en/4-developer_guide/8-motrix_contact_sensor.md b/docs/sphinx/source/en/4-developer_guide/8-motrix_contact_sensor.md deleted file mode 100644 index 1cc419a7d..000000000 --- a/docs/sphinx/source/en/4-developer_guide/8-motrix_contact_sensor.md +++ /dev/null @@ -1,139 +0,0 @@ -# Motrix Contact Sensor Notes - -## Background - -The tactile observation for the Sharpa in-hand rotation task depends on -contact-sensor data (fingertip contact force against the object). The code -works under the MuJoCo backend, but Motrix returns contact-sensor data in a -different format, so the read path needs to account for both shapes. - -## Contact Sensor Configuration - -The fingertip contact sensors are defined in the robot scene XML -`src/unilab/assets/robots/sharpa_wave/right_sharpa_wave.xml`: - -```xml - - -``` - -Parameters: - -- `num="1"` — report at most one contact point. -- `data="force"` — return only the force data. -- `reduce="netforce"` — reduce to a single net contact force in the global - frame. - -## Return-Format Difference Between Backends - -### MuJoCo - -Returns a `(num_envs, 3)` force vector in the global frame: - -``` -[[fx, fy, fz], # env 0 - [fx, fy, fz], # env 1 - ...] -``` - -### Motrix - -Returns a `(num_envs, 1 + num * stride)` flat array whose **first element is -the actual contact count**: - -``` -[[count, fx, fy, fz], # env 0, shape = (4,) because num=1, stride=3 - [count, fx, fy, fz], # env 1 - ...] -``` - -For `reduce="netforce"` + `data="force"` + `num=1`: - -- shape = `(num_envs, 4)` -- `[0]` = contact count (0 or 1) -- `[1:4]` = net force vector in the global frame (`netforce` uses global - coordinates) - -### General multi-contact layout (Motrix) - -With `num=4, data="force pos normal tangent"` and no reduce: - -``` -shape = (num_envs, 1 + 4 * 12) = (num_envs, 49) - -[count, - f1_normal, f1_tangent0, f1_tangent1, # contact 1 force (contact frame) - f1_x, f1_y, f1_z, # contact 1 position - f1_nx, f1_ny, f1_nz, # contact 1 normal - f1_tx, f1_ty, f1_tz, # contact 1 tangent - f2_normal, f2_tangent0, f2_tangent1, # contact 2 ... - ... - padding_zeros] # zero-padded up to num contacts -``` - -Note: without a reduce mode the force is expressed in the **contact frame** -(one normal scalar plus two tangent scalars), not a global `xyz` vector. Only -`reduce="netforce"` returns a global-frame force vector. - -## Motrix `reduce` Mode Reference - -| Reduce mode | Frame | Returns | -| --- | --- | --- | -| `netforce` | global | single reduced net force vector | -| `maxforce` | contact | the contact point with the largest force | -| `mindist` | contact | the contact point with the shallowest penetration | -| none | contact | the first `num` contact points | - -Only `netforce` returns a global-frame force vector; the other modes return -contact-frame data (one normal scalar plus two tangent scalars). - -## Why a Single Norm Branch Is Not Enough - -The env reads tactile force through `_read_tactile_force()` → -`_extract_sensor_scalar()` in -`src/unilab/tasks/manipulation/sharpa_inhand/base.py`. That helper currently -collapses any `(N, >=3)` array with `np.linalg.norm(data[:, :3], axis=1)`. - -If the env still routes both backend shapes through that one branch, the -MuJoCo `(N, 3)` case is correct (`norm` of the real force vector), but the -Motrix `(N, 4)` case would be wrong: `data[:, :3]` would pick up -`[count, fx, fy]` — treating the contact count as a force component and -dropping `fz`. The fix is not to special-case shapes inside the env, but to -move the per-backend knowledge behind a backend method. - -## Recommended Contract: a Backend Method for Force Magnitude - -Add `get_contact_force_magnitude(sensor_name) -> np.ndarray` to the -`SimBackend` interface (`unisim.backend.base`), returning a -`(num_envs,)` scalar magnitude. Each backend implements it according to its -own data layout: - -- **MuJoCo** (`unisim.backend.mujoco.backend`): take the norm of - the 3D force vector returned by `get_sensor_data(name)`. -- **Motrix** (`unisim.backend.motrix.backend`): interpret the - layout by reduce mode: - - `reduce="netforce"`: take `[1:4]`, then norm. - - no reduce, multiple contacts: sum the per-contact forces, then norm. - - `reduce="maxforce"`: take the strongest contact. - -The env's `_read_tactile_force()` then routes contact sensors through -`get_contact_force_magnitude()`, while ordinary scalar sensors keep using -`_extract_sensor_scalar()`. - -This keeps the env backend-agnostic and aligns with the high-risk-area -invariant that the env layer only calls methods declared on `SimBackend` — -no feature leakage into the env, and a new backend only has to implement the -interface method. `get_sensor_data` is already a declared `SimBackend` -method, so the proposal extends the same boundary rather than reaching into a -backend subclass. - -## Related Files - -| File | Role | -| --- | --- | -| `src/unilab/tasks/manipulation/sharpa_inhand/base.py` | `_extract_sensor_scalar()`, `_read_tactile_force()` | -| `src/unilab/tasks/manipulation/sharpa_inhand/rotation.py` | reward computation, virtual torque | -| `src/unilab/assets/robots/sharpa_wave/right_sharpa_wave.xml` | contact-sensor XML definitions | -| `unisim.backend.motrix.backend` | Motrix `get_sensor_data()` | -| `unisim.backend.mujoco.backend` | MuJoCo `get_sensor_data()` | -| `unisim.backend.base` | `SimBackend` interface | 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 07e743a2e..75fd336a6 100644 --- a/docs/sphinx/source/en/5-reference/5-support_matrix.md +++ b/docs/sphinx/source/en/5-reference/5-support_matrix.md @@ -87,8 +87,6 @@ rendering/playback paths remain unsupported. | PPO (torch) | `g1_wall_flip_tracking` (G1 wall flip tracking) | Tested | - | Tested | - | - | | PPO (torch) | `x2_wall_flip_tracking` (X2 wall flip tracking) | Tested | - | Tested | - | - | | PPO (torch) | `allegro_inhand` (Allegro in-hand) | Tested | - | Tested | - | - | -| PPO (torch) | `sharpa_inhand` (Sharpa in-hand) | Tested | - | Tested | - | - | -| PPO (torch) | `sharpa_inhand_grasp` (Sharpa in-hand grasp) | Tested | - | Tested | - | - | | PPO (torch) | `a2_joystick_flat` (a2 joystick flat) | Tested | - | - | - | - | | PPO (torch) | `allegro_inhand_grasp` (allegro inhand grasp) | Tested | - | Tested | - | - | | PPO (torch) | `g1_23dof_box_tracking` (g1 23dof box tracking) | Tested | - | Tested | - | - | @@ -114,7 +112,6 @@ rendering/playback paths remain unsupported. | APPO (torch) | `g1_flip_tracking` (G1 flip tracking) | Tested | - | Tested | - | - | | APPO (torch) | `g1_wall_flip_tracking` (G1 wall flip tracking) | Tested | - | Tested | - | - | | APPO (torch) | `allegro_inhand` (Allegro in-hand) | Tested | - | Tested | - | - | -| APPO (torch) | `sharpa_inhand` (Sharpa in-hand) | Tested | - | Tested | - | - | | APPO (torch) | `g1_23dof_climb_tracking` (g1 23dof climb tracking) | Tested | - | Tested | - | - | | APPO (torch) | `g1_23dof_flip_tracking` (g1 23dof flip tracking) | Tested | - | Tested | - | - | | APPO (torch) | `g1_23dof_motion_tracking` (g1 23dof motion tracking) | Tested | - | Tested | - | - | diff --git a/docs/sphinx/source/zh_CN/0-index.md b/docs/sphinx/source/zh_CN/0-index.md index c8d9fd6a5..3c55e8318 100644 --- a/docs/sphinx/source/zh_CN/0-index.md +++ b/docs/sphinx/source/zh_CN/0-index.md @@ -101,7 +101,7 @@ uv run train --algo ppo --task go2_joystick_flat --sim motrix \ :::{grid-item-card} 挑选算法 :link: 2-user_guide/2-algorithms/0-index :link-type: doc -对比 PPO、APPO、SAC、TD3、FlashSAC 和 HORA 的入口。 +对比 PPO、APPO、SAC、TD3 和 FlashSAC 的入口。 ::: :::{grid-item-card} 部署或切换仿真 @@ -151,7 +151,6 @@ recommendation 元数据。 | Go2W joystick | PPO | `src/unilab/conf/ppo/task/go2w_joystick_*` 下存在 MuJoCo 与 Motrix flat/rough 变体的 PPO owner YAML。 | | G1 locomotion / tracking | PPO、APPO、SAC、TD3 | PPO、APPO、SAC 都为 G1 任务提供了已提交的 MuJoCo 与 Motrix owner YAML；TD3 有一个 `g1_walk_flat` 的 MuJoCo owner。 | | Allegro in-hand | PPO、APPO | PPO 和 APPO 为 Allegro in-hand 任务提供了已提交的 MuJoCo 与 Motrix owner YAML。 | -| Sharpa in-hand | PPO、APPO HORA teacher、HORA distillation | Sharpa owner YAML 为 PPO/APPO teacher 路径已提交；student distillation 使用 `src/unilab/conf/hora_distill/task/sharpa_inhand/mujoco.yaml`。 | ```{toctree} :hidden: 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 acc1a780d..f2a9685b1 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 @@ -121,7 +121,7 @@ editable install 会指向源码 checkout；常规安装会把 package 和任务 MJWarp、Genesis、平台相关 torch index、ROCm / XPU profile 请优先使用上面的 uv 路径。机器人 mesh 和纹理不会打进 wheel，而是在 cold path 从 `unilabsim/unilab-robots` 数据集下载。请确保安装位置可写，或从源码 checkout 使用 -`uv run unilab-pull-assets` 预拉取。isaacgym / isaacsim 后端和 HORA 多卡提交路径仍假设 +`uv run unilab-pull-assets` 预拉取。isaacgym / isaacsim 后端仍假设 源码 checkout；外部后端请使用下方专用安装页。 ## 运行时 Asset diff --git a/docs/sphinx/source/zh_CN/1-getting_started/3-evaluation_and_playback.md b/docs/sphinx/source/zh_CN/1-getting_started/3-evaluation_and_playback.md index 76de27894..266859c37 100644 --- a/docs/sphinx/source/zh_CN/1-getting_started/3-evaluation_and_playback.md +++ b/docs/sphinx/source/zh_CN/1-getting_started/3-evaluation_and_playback.md @@ -37,7 +37,7 @@ uv run demo dance 调试策略时，可以使用低层脚本 `src/unilab/scripts/play_interactive.py`。 `src/unilab/scripts/play_interactive.py` 是通用 MuJoCo viewer 入口，适合 PPO、APPO、 -SAC、FlashSAC 和 HORA distill 的策略可视化。它使用 `--algo / --task / --sim` +SAC 和 FlashSAC 的策略可视化。它使用 `--algo / --task / --sim` 选择算法和 owner config；无论 `--sim` 选择 MuJoCo 还是 Motrix，窗口都使用 `mujoco.viewer` 可视化，`--sim` 只决定读取哪份配置。 diff --git a/docs/sphinx/source/zh_CN/1-getting_started/4-project_structure.md b/docs/sphinx/source/zh_CN/1-getting_started/4-project_structure.md index 50d6eadd1..f13283913 100644 --- a/docs/sphinx/source/zh_CN/1-getting_started/4-project_structure.md +++ b/docs/sphinx/source/zh_CN/1-getting_started/4-project_structure.md @@ -8,7 +8,7 @@ UniLab 将运行时 contract、配置、训练脚本和文档分置于不同的 | `src/unilab/conf/` | Hydra 根配置和任务 owner YAML。顶层 CLI 将后端选择暴露为 `--task` 加 `--sim`，然后组合出匹配的 owner YAML。 | | `src/unilab/base/` | Registry、env state、scene 以及 backend contract。 | | `src/unilab/envs/` | 任务 env 实现，以及任务专属的 reset、reward、observation 和 DR 逻辑。 | -| `uni_rl` (unilab-rl repo) | PPO、APPO、off-policy 和 HORA 算法代码。 | +| `uni_rl` (unilab-rl repo) | PPO、APPO 和 off-policy 算法代码。 | | `uni_rl.ipc` (unilab-rl repo) | 共享内存与异步 runner 原语。 | | `src/unilab/training/` | 共享的训练辅助工具，用于日志、回放、种子处理和配置守卫（config guard）。 | | `src/unilab/visualization/` | 回放、渲染、NaN 检查以及 scene/export 工具。 | @@ -24,7 +24,6 @@ UniLab 将运行时 contract、配置、训练脚本和文档分置于不同的 - `src/unilab/conf/appo/config.yaml`，用于 APPO。 - `src/unilab/conf/sac/config.yaml`、`src/unilab/conf/td3/config.yaml` 和 `src/unilab/conf/flashsac/config.yaml`， 分别用于 SAC、TD3 和 FlashSAC，算法超参数内联在各自的 config.yaml 中。 -- `src/unilab/conf/hora_distill/config.yaml`，用于 HORA student distillation。 任务 owner YAML 即后端身份。示例： diff --git a/docs/sphinx/source/zh_CN/2-user_guide/0-index.md b/docs/sphinx/source/zh_CN/2-user_guide/0-index.md index 376723d59..a6eb6abb4 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/0-index.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/0-index.md @@ -15,7 +15,7 @@ CLI 路由、Hydra owner YAML、日志、检查点与 Docker。 :::{grid-item-card} 算法 :link: 2-algorithms/0-index :link-type: doc -对比 PPO、APPO、SAC、TD3、FlashSAC 与 HORA。 +对比 PPO、APPO、SAC、TD3 与 FlashSAC。 ::: :::{grid-item-card} 后端 @@ -42,12 +42,6 @@ CLI 路由、Hydra owner YAML、日志、检查点与 Docker。 导出 ONNX、检查 NaN、发送 W&B 日志以及导出场景。 ::: -:::{grid-item-card} 操作说明 -:link: 8-manipulation/0-index -:link-type: doc -针对 Allegro、Sharpa 与 Go2+Airbot 的任务专属说明。 -::: - :::: ```{toctree} @@ -61,5 +55,4 @@ CLI 路由、Hydra owner YAML、日志、检查点与 Docker。 5-domain_randomization/0-index 6-terrain/0-index 7-tooling/0-index -8-manipulation/0-index ``` diff --git a/docs/sphinx/source/zh_CN/2-user_guide/1-training/2-hydra_config.md b/docs/sphinx/source/zh_CN/2-user_guide/1-training/2-hydra_config.md index 567db3862..3d216c6ab 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/1-training/2-hydra_config.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/1-training/2-hydra_config.md @@ -10,7 +10,6 @@ reward、scene 以及 task 专属运行时字段的身份标识。 | PPO | `src/unilab/conf/ppo/task//.yaml` | | APPO | `src/unilab/conf/appo/task//.yaml` | | SAC / TD3 / FlashSAC | `src/unilab/conf//task//.yaml` | -| HORA 蒸馏 | `src/unilab/conf/hora_distill/task//.yaml` | 示例： diff --git a/docs/sphinx/source/zh_CN/2-user_guide/1-training/3-logging.md b/docs/sphinx/source/zh_CN/2-user_guide/1-training/3-logging.md index 27ee341ac..eb1d9f300 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/1-training/3-logging.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/1-training/3-logging.md @@ -64,7 +64,7 @@ rate 求和，因此两者都是整个训练任务的总吞吐。标题中的 `A 等于 `Iter Wall`，每行百分比合计约为 100%（逐行取整可能略有误差）。 终端不再按 1% 阈值隐藏适用于当前算法的阶段，0 ms 也保留，便于与 TensorBoard / W&B 逐项对应。算法专有阶段不会跨算法占位或落盘；例如 `Replay Stage` 与 `Weight Publish` 只在 -APPO / HORA-APPO 出现。 +APPO 出现。 ### Learner 主时间线 diff --git a/docs/sphinx/source/zh_CN/2-user_guide/2-algorithms/0-index.md b/docs/sphinx/source/zh_CN/2-user_guide/2-algorithms/0-index.md index d0fe592ab..183eaf0cb 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/2-algorithms/0-index.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/2-algorithms/0-index.md @@ -10,7 +10,6 @@ | SAC | off-policy | `src/unilab/scripts/train_sac.py` | `src/unilab/conf/sac/config.yaml` | | TD3 | off-policy | `src/unilab/scripts/train_td3.py` | `src/unilab/conf/td3/config.yaml` | | FlashSAC | off-policy | `src/unilab/scripts/train_flashsac.py` | `src/unilab/conf/flashsac/config.yaml` | -| HORA | teacher/student 蒸馏路径 | `scripts/train_hora_distill.py` | `src/unilab/conf/hora_distill/config.yaml` | ```{toctree} :hidden: @@ -20,5 +19,4 @@ 3-sac 4-td3 5-flash_sac -7-hora ``` diff --git a/docs/sphinx/source/zh_CN/2-user_guide/2-algorithms/7-hora.md b/docs/sphinx/source/zh_CN/2-user_guide/2-algorithms/7-hora.md deleted file mode 100644 index 3741e6b30..000000000 --- a/docs/sphinx/source/zh_CN/2-user_guide/2-algorithms/7-hora.md +++ /dev/null @@ -1,26 +0,0 @@ -# HORA - -已提交的 HORA 路径是 Sharpa 手内（in-hand）teacher/student 流程。teacher owner 位 -于 PPO 与 APPO 的 task 树下，通过 `sharpa_inhand` 的 `7-hora` profile 选择；student -蒸馏使用 `scripts/train_hora_distill.py` 和 -`src/unilab/conf/hora_distill/task/sharpa_inhand/mujoco.yaml`。 - -## Teacher - -```bash -uv run train --algo ppo --task sharpa_inhand --sim mujoco --profile hora -uv run train --algo appo --task sharpa_inhand --sim mujoco --profile hora training.no_play=true -``` - -HORA PPO owner 设置 `algo.algo_log_name=hora_ppo`，并通过 -`uni_rl.algos.hora.rsl_rl:resolve_hora_ppo_runtime` 解析运行时。APPO 变体设置 -`algo.algo_log_name=hora_appo`。 - -## Student 蒸馏 - -student 蒸馏由 `scripts/train_hora_distill.py` 实现，并由 -`src/unilab/conf/hora_distill/task/sharpa_inhand/mujoco.yaml` 配置。顶层 CLI 目前没有声明独立的 -HORA 蒸馏 `--algo` 路由，因此本页的公开 CLI 示例仍保持在上面的 teacher 路径上。 - -teacher 检查点的解析在 `src/unilab/training/hora_distill_config.py` 中实现。 -student 日志族为 `hora_distill`。 diff --git a/docs/sphinx/source/zh_CN/2-user_guide/4-tasks/0-index.md b/docs/sphinx/source/zh_CN/2-user_guide/4-tasks/0-index.md index c97e5b68d..e1265c0db 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/4-tasks/0-index.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/4-tasks/0-index.md @@ -21,7 +21,7 @@ G1 动作追踪、翻转、攀爬、墙面翻转和箱体追踪。 :::{grid-item-card} 操作 :link: 3-manipulation :link-type: doc -Allegro 和 Sharpa 手内旋转与抓取生成。 +Allegro 手内旋转与抓取生成。 ::: :::: diff --git a/docs/sphinx/source/zh_CN/2-user_guide/4-tasks/3-manipulation.md b/docs/sphinx/source/zh_CN/2-user_guide/4-tasks/3-manipulation.md index 1dac402ba..09c631511 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/4-tasks/3-manipulation.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/4-tasks/3-manipulation.md @@ -5,19 +5,12 @@ ## 手内操作 - `allegro_inhand` 和 `allegro_inhand_grasp` 拥有 MuJoCo 和 Motrix PPO owner。 -- `sharpa_inhand`、`sharpa_inhand_grasp` 以及 `sharpa_inhand` 的 `hora` - 配置在当前 config 中都是 MuJoCo owner 路径。 ```bash uv run train --algo ppo --task allegro_inhand --sim mujoco uv run train --algo ppo --task allegro_inhand --sim motrix training.no_play=true -uv run train --algo ppo --task sharpa_inhand --sim mujoco --profile hora training.no_play=true ``` -HORA student 蒸馏由 -`src/unilab/conf/hora_distill/task/sharpa_inhand/mujoco.yaml` 配置；它当前未作为 -单独的顶层 CLI 路线暴露。 - ## 平台平衡 `stewart_balance` 是一个 6 自由度并联（Stewart）平台，用于把一个自由小球 @@ -31,5 +24,3 @@ HORA student 蒸馏由 ```bash uv run train --algo ppo --task stewart_balance --sim motrix training.no_play=true ``` - -手内操作任务说明见 {doc}`../8-manipulation/1-dexterous_inhand`。 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 1ee39f40f..55c07ff78 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 @@ -6,7 +6,7 @@ 当前存在两条 DR 声明路径： - **Manager-Based（Compatible）任务**：reset / interval 随机化通过 owner YAML 中的 Hydra `events:` manager term 声明；reset 生命周期的 event 在 reset 时采样，interval 生命周期的 event 在 step 之间施加扰动。例如 `src/unilab/conf/ppo/task/go1_joystick_flat/base.yaml` 的 `events:` 段。 -- **legacy provider 路径**：只有 2 个 Adapted family（`sharpa_inhand` / `sharpa_inhand_grasp`，含 appo / hora owner）仍通过 `DomainRandomizationProvider` + `DomainRandomizationManager` 声明 `env.domain_rand.*` 配置。 +- **任务级 provider 路径**：自定义任务（包括托管在外部仓库中的任务）可以通过 `DomainRandomizationProvider` + `DomainRandomizationManager` 声明 `env.domain_rand.*` 配置。当前仓库内没有任务使用该路径。 legacy provider 路径的统一入口点位于 `NpEnv._init_domain_randomization()` 和 `DomainRandomizationManager`： @@ -16,18 +16,18 @@ legacy provider 路径的统一入口点位于 `NpEnv._init_domain_randomization 这三条路径对应三个生命周期类别： -- **init 生命周期 DR**：改变模型 identity 或模型几何的项；只能在 env/backend 初始化和 materialization 期间生效，例如 Sharpa 手物体的 `geom_size` 缩放。 +- **init 生命周期 DR**：改变模型 identity 或模型几何的项；只能在 env/backend 初始化和 materialization 期间生效，例如通过模型变体进行的物体 `geom_size` 缩放。 - **reset 生命周期 DR**：不改变模型 identity，只在同一模型内改变参数或 reset 状态的项，例如 `base_mass_delta`、`base_com_offset`、`gravity`、`kp`、`kd`。 - **interval 生命周期 DR**：step 之间的外部扰动，例如 push。 ## 状态结论 -1. Manager-Based 任务不注册 DR provider；它们的 reset/interval 随机化是 owner YAML 中的 `events:` manager term，由 manager 生命周期统一执行。只有 Adapted family 的冻结兼容工厂仍走 `DomainRandomizationManager` 统一入口。 -2. Adapted family owner 定义 `domain_rand` 配置 dataclass、`DomainRandomizationProvider` 和 `ResetPlan`；Manager-Based owner 则通过 Hydra command/event term 声明 reset 行为。G1 motion reset 扰动归 `MotionCommandCfg` 所有，WBT 另加 `EventTermCfg` reset 与 interval term。 +1. Manager-Based 任务不注册 DR provider；它们的 reset/interval 随机化是 owner YAML 中的 `events:` manager term，由 manager 生命周期统一执行。走 provider 路径的自定义任务则经过 `DomainRandomizationManager` 统一入口。 +2. provider 路径的 owner 定义 `domain_rand` 配置 dataclass、`DomainRandomizationProvider` 和 `ResetPlan`；Manager-Based owner 则通过 Hydra command/event term 声明 reset 行为。G1 motion reset 扰动归 `MotionCommandCfg` 所有，WBT 另加 `EventTermCfg` reset 与 interval term。 3. 今天所"统一"的主要是入口点和执行流程，而不是每一个随机化项本身。legacy 路径的共享辅助函数 `build_common_reset_randomization()` 目前生成 `base_mass_delta`、`base_com_offset`、`gravity`、`kp`、`kd`。 4. `ResetRandomizationPayload` 已经可以表达 `gravity`、`body_iquat`、`body_inertia`、`kp`、`kd`，并且 `MuJoCoBackend` 已声明支持。这些是否实际被使用，仍取决于 task provider 是否对它们进行采样和 dispatch。 5. `MotrixBackend` 目前支持 `base_mass_delta`、`base_com_offset`、`kp`、`kd` 和 interval push；并且它要求在初始化期间所有模型 actuator 都是 position actuator。 -6. `geom_size` 不是 reset 生命周期字段；Sharpa 手物体的 geom 缩放由 init 生命周期的模型 materialization 处理。 +6. `geom_size` 不是 reset 生命周期字段；物体 geom 缩放由 init 生命周期的模型 materialization 处理。 ## 统一性评估表 @@ -41,8 +41,6 @@ legacy provider 路径的统一入口点位于 `NpEnv._init_domain_randomization | `G1WBTObs` | Hydra `events:` term | 是：同一 motion command + Hydra `EventTermCfg` | motion reset 加 mass/COM/PD/friction/encoder-bias event | interval velocity kick | `motion_tracking/g1/manager_terms.py` | | `AllegroInhandRotation` | Hydra `events:` term | 是：Hydra `EventTermCfg` + Manager-Based reset term | entity 范围的手/球 reset | 无 | `allegro_inhand/manager_terms.py` | | `AllegroInhandRotationGrasp` | Hydra `events:` term | 是：复用 rotation reset event + `RecorderTermCfg` | 带噪声的手部 reset + grasp 收集 | 无 | `allegro_inhand/grasp_gen.py` | -| `SharpaInhandRotation` | legacy provider | 是：`InitRandomizationPlan + ResetPlan + IntervalRandomizationPlan` | grasp cache 采样 + common payload | 物体 `body_force` | `sharpa_inhand/rotation.py` | -| `SharpaInhandRotationGrasp` | legacy provider | 是：复用 Sharpa rotation provider 并 override reset 采样 | grasp 收集 reset + common payload | 无 | `sharpa_inhand/grasp_gen.py` | ## 各任务域随机化清单 @@ -56,8 +54,6 @@ legacy provider 路径的统一入口点位于 `NpEnv._init_domain_randomization | `G1WBTObs` | 同一 motion reset 加 base mass、base COM、PD gain、足端摩擦和 encoder-bias event term | `push_by_setting_velocity` | WBT owner 显式启用上述全部 event term；能力不支持时直接报错，不回退 | | `AllegroInhandRotation` | entity 范围的手/球 reset；显式配置 grasp cache 时进行采样，否则以 `null` 显式选择模型 home pose；可选 `joint_noise`、`ball_velocity_noise` 与 `ball_z_offset` | 无 | owner YAML 显式选择 home pose 与零 reset 噪声；配置的 cache 缺失或格式错误时 fail-closed | | `AllegroInhandRotationGrasp` | 复用 rotation reset 并设置 `joint_noise=0.25`；Manager-Based termination 检查指尖距离、接触数和球高度；recorder 保存成功 timeout rows | 无 | 生成 5 万行 Allegro grasp cache，成功保存后抛出 `RunComplete` | -| `SharpaInhandRotation` | grasp cache 按 `scale_ids` 分桶采样；物体位姿 / quat reset；可选 common reset 随机化 payload（含 `gravity`） | 物体 `body_force` 直接力扰动 | `domain_rand.scale_list` 默认值来自 owner YAML；在 MuJoCo 下，物体 geom 缩放在 init 期间 materialize；common payload 默认禁用；物体 force 通过 Sharpa owner YAML 默认启用 | -| `SharpaInhandRotationGrasp` | 手部位姿 reset；物体位姿 / quat reset；收集成功的 grasp 并按 `scale_ids` 分桶存储；可选 `base_mass_delta`；可选 `base_com_offset`；可选 `gravity` | 无 | 默认用于生成 Sharpa grasp cache；cache 文件名包含单个 scale 值；common payload 默认禁用 | ## 当前统一 DR 的能力与边界 @@ -69,7 +65,7 @@ legacy provider 路径的统一入口点由 `NpEnv` 和 `DomainRandomizationMana - manager 统一执行能力验证 - 后端统一负责实际施加随机化 payload -因此从执行路径的角度看，仍走该路径的 Adapted family 是统一的；Manager-Based 任务则由 manager 生命周期统一执行 owner YAML 声明的 `events:` term。 +因此从执行路径的角度看，provider 路径的任务是统一的；Manager-Based 任务则由 manager 生命周期统一执行 owner YAML 声明的 `events:` term。 ### 2. 共享辅助函数仍然较窄 @@ -79,10 +75,10 @@ legacy 路径的 `dr_utils.py` 构造并校验通用 reset payload： 这意味着： -- 仍走 legacy provider 的 sharpa family，其 task 专属状态仍直接在各自的 provider 内部采样 +- provider 路径的任务直接在各自的 provider 内部采样 task 专属状态 - `G1MotionTracking` 的 pose / velocity / joint 噪声由其 manager command 所有 - Allegro 的 grasp / 物体初始状态采样完全是 task 专属逻辑 -- Sharpa 的 `geom_size` 缩放是 init 生命周期的模型 materialization，不属于 reset common payload +- `geom_size` 缩放是 init 生命周期的模型 materialization，不属于 reset common payload 所以今天的"统一性"更多是关于 contract 和调用约定，而不是"所有任务共享同一套随机化项 schema"。 @@ -106,7 +102,7 @@ legacy 路径的 `dr_utils.py` 构造并校验通用 reset payload： 说明： - 当前的 `IntervalRandomizationPlan` 支持 `push_perturbation_limit`、`body_linear_velocity_delta`、`body_angular_velocity_delta`、`body_force` 和 `body_torque`；其中 `body_force`/`body_torque` 表达热路径上的直接外力/力矩扰动，而不暴露后端私有的 `xfrc_applied` 细节。 -- 当前 MuJoCo 后端的 interval push 和 interval body force 都通过 `xfrc_applied` dispatch；Sharpa 手物体扰动已切换为直接力扰动。 +- 当前 MuJoCo 后端的 interval push 和 interval body force 都通过 `xfrc_applied` dispatch。 - Motrix 后端目前仍不支持直接 body-force 扰动，因此这类 owner 配置必须继续显式禁用。 但在任务侧，当前的现实是：并非每个 provider 都构造这些字段。后端 contract 是能力边界；task 配置和 provider 是否 dispatch 一个 payload，才决定了某个任务是否实际启用对应的 DR 项。 @@ -120,7 +116,7 @@ legacy 路径的 `dr_utils.py` 构造并校验通用 reset payload： - 生命周期：仅在 reset 时采样和写入；env 会保留该重力，直到下一次 reset 重新采样。 - 后端：当前在 UniLab 中，只有 MuJoCo 后端声明支持该 reset 项；Motrix 后端不支持。一些任务按能力过滤并跳过它；另一些任务在 validate 阶段抛出错误。 -配置入口位于采用 legacy provider 路径的 Sharpa owner（如 `sharpa_inhand_grasp`）的 `env.domain_rand` 下；Manager-Based 任务没有 `env.domain_rand`： +配置入口位于 provider 路径任务 owner 的 `env.domain_rand` 下；Manager-Based 任务没有 `env.domain_rand`： ```yaml env: @@ -137,21 +133,7 @@ env: - `gravity_range`：一个形状为 `(2, 3)` 的逐维采样范围；第一行和第二行给出每个分量的上界和下界。 - 在每次 reset 时，每个维度在 `[min(row0, row1), max(row0, row1)]` 内均匀采样。方向不会自动归一化，重力范数也不固定。 -如果你只想随机化大小而保持竖直向下的方向，只开放 `z` 分量： - -```bash -uv run train --algo ppo --task sharpa_inhand_grasp --sim mujoco \ - env.domain_rand.randomize_gravity=true \ - 'env.domain_rand.gravity_range=[[0.0,0.0,-10.5],[0.0,0.0,-8.5]]' -``` - -如果你想同时随机化方向和大小，开放 `x/y/z`： - -```bash -uv run train --algo ppo --task sharpa_inhand_grasp --sim mujoco \ - env.domain_rand.randomize_gravity=true \ - 'env.domain_rand.gravity_range=[[-0.3,-0.3,-10.5],[0.3,0.3,-8.5]]' -``` +如果你只想随机化大小而保持竖直向下的方向，只开放 `z` 分量；如果想同时随机化方向和大小，开放 `x/y/z`。在 provider 路径的任务 owner 上，可通过 CLI 以 `env.domain_rand.randomize_gravity=true` 与 `env.domain_rand.gravity_range=[...]` override 启用。 说明： @@ -192,100 +174,9 @@ uv run train --algo ppo --task go1_joystick_flat --sim mujoco \ 这条边界存在的目的是遵循冷路径 asset/model-metadata 访问原则：`step()`、`reset()` 和热路径 DR 不解析 XML、不读取 asset，也不在运行时基于 asset 元数据进行分支。 -## Sharpa 手物体 Geom 缩放用法 - -Sharpa 手是仓库中当前 `geom_size` init 生命周期 DR 的示例任务。相关任务配置： - -- `src/unilab/conf/ppo/task/sharpa_inhand/mujoco.yaml` -- `src/unilab/conf/ppo/task/sharpa_inhand_grasp/mujoco.yaml` - -### 1. 配置入口 - -Sharpa 的 scale 配置位于 env owner YAML 的 `env.domain_rand.scale_list`： - -```yaml -env: - object_body_name: object - object_geom_name: object - domain_rand: - scale_list: [0.5, 0.6, 0.7, 0.8] -``` - -字段语义： - -- `object_body_name`：物体 body name；用于在 reset / observation 期间定位物体 body，而不是缩放的目标字段。 -- `object_geom_name`：要缩放的 MuJoCo geom name；默认为 `object`。 -- `domain_rand.scale_list`：显式的 scale 列表；每个值必须大于 0。 -- `domain_rand.scale_list` 的顺序就是 `scale_id` 的顺序。 -- `domain_rand.scale_list` 的长度就是模型变体的数量。 - -每个 env 被静态分配一个 `scale_id`。当前的分配规则是连续分桶分配；当 `algo.num_envs` 不能被 `num_scales` 整除时，前几个 scale 桶各多分到一个 env： - -```bash -uv run train --algo ppo --task sharpa_inhand --sim mujoco 'env.domain_rand.scale_list=[0.5,0.6,0.7,0.8]' algo.num_envs=4096 -``` - -如果 `algo.num_envs=4096` 且 `num_scales=4`，那么每 1024 个 env 使用同一个 scale 桶。 - -### 2. MuJoCo Materialization 行为 - -MuJoCo 后端如何应用它： - -1. env/provider 在 init 期间基于 `scale_list` 构建 `ModelVariantSpec`。 -2. 后端使用 `MjSpec` 读取模型，并修改与 `object_geom_name` 对应的 geom 的 `size`。 -3. 每个 scale 编译一个 scale 专属的 `MjModel`。 -4. 在首次需要物理 pool 时，将 env 到模型的分配展开为长度为 `num_envs` 的模型序列，然后构造 `BatchEnvPool`。 - -因此，`domain_rand.scale_list` 仅在 env/backend 初始化期间生效。在 env 创建之后修改 `env.domain_rand.scale_list` 不会改变已经 materialize 的模型 pool。 - -该流程有三条重要边界： - -- `BatchEnvPool` 是惰性构造的；正常路径不会先为默认模型构造一个 pool，再为 `scale_list` 重建它。 -- 多个模型变体的编译使用基于进程的并行分块完成；不要在 Python 线程中编译，也不要在上层 for 循环中串行编译 `num_envs` 个模型。 -- worker 使用 `MjSpec` 编译变体并保存 `.mjb`；父进程仅按 `.mjb` 路径加载 `MjModel.from_binary_path(...)`。不要通过 IPC 传回修改后的模型对象或模型字节。 - -### 3. Grasp Cache 与 Scale 桶 - -Sharpa rotation 任务按 `scale_ids` 从多个单 scale 的 grasp cache 中采样： - -- cache 文件名默认由 `grasp_cache_path` 和单个 scale 值共同决定。 -- `scale_list: [0.5, 0.6, 0.7, 0.8]` 默认对应 `caches/sharpa_grasp_linspace_0.5.npy`、`caches/sharpa_grasp_linspace_0.6.npy`、`caches/sharpa_grasp_linspace_0.7.npy`、`caches/sharpa_grasp_linspace_0.8.npy`。 -- 在 rotation 启动时会检查 `scale_list` 的所有 cache 文件；如果有任何缺失，则报错。 -- 每个 scale 桶只从其自身 scale 的 cache 文件中采样，避免在不同物体 scale 之间混用 grasp 初始状态。 - -默认 scale 的 cache 托管在 Hugging Face (`unilabsim/unilab-caches`)，首次 rotation 训练时会自动下载到 `src/unilab/assets/caches/`，标准 `scale_list` 无需手动采集。 - -如需为 HF 上没有的自定义 scale 采集 cache，或在本地重新生成，可按每个 scale 分别运行 grasp 收集任务。生成的文件会落到 `src/unilab/assets/caches/`，与 HF 下载位置一致，后续 rotation 训练能直接命中；**但重新生成耗时较长**。 - -辅助脚本按顺序采集每个 scale： - -```bash -./scripts/sharpa_collect_grasps.sh 0.5 0.6 0.7 0.8 -``` - -等价的逐 scale 调用：`uv run train --algo ppo --task sharpa_inhand_grasp --sim mujoco 'env.domain_rand.scale_list=[0.5]' algo.num_envs=4096`（对 `[0.6]`、`[0.7]`、`[0.8]` 等同理重复）。 - -然后用相同的 `scale_list` 训练 rotation： - -```bash -uv run train --algo ppo --task sharpa_inhand --sim mujoco 'env.domain_rand.scale_list=[0.5,0.6,0.7,0.8]' algo.num_envs=4096 -``` - -### 4. 边界与注意事项 - -- `geom_size` 不是 reset DR 字段，且不得写入 `ResetPlan.randomization`。 -- `BatchEnvPool.reset(..., randomization=...)` 目前不支持 `geom_size`。 -- `geom_size` 缩放仅在 MuJoCo 后端下 materialize；Motrix 后端目前不会从 `scale_list` 产生多模型 pool。 -- `scale_list` 的长度是模型变体的数量，而不是每次 reset 的重采样次数。 -- 每个 env 的 `scale_id` 在 init 期间静态分配，且在 reset 时不变。 -- 在扩容时，扩容模型变体的数量；不要按 `num_envs` 为每个 env 编译一个模型。多个 env 共享与同一 scale 桶对应的同一个 `MjModel`。 -- 热路径不得读取 XML、解析 asset，也不得使用 `getattr` / `hasattr` 探测后端私有能力来决定缩放行为。 -- 在扩展到其他 shape DR 时，优先复用 init 生命周期 contract；不要把 shape 字段塞进 reset payload。 - ## 相关任务 - {doc}`G1 Motion Tracking <../4-tasks/2-motion_tracking>`：开启 DR 前先确认 motion 资产和 replay。 -- {doc}`Sharpa Inhand <../8-manipulation/1-dexterous_inhand>`：scale / grasp cache / DR 边界更敏感。 - {doc}`Go2 Rough Terrain <../4-tasks/1-locomotion>`：常见的是 mass、COM、friction、push。 有关配置示例，请参阅 {doc}`1-configuration`。有关开发者 diff --git a/docs/sphinx/source/zh_CN/2-user_guide/5-domain_randomization/1-configuration.md b/docs/sphinx/source/zh_CN/2-user_guide/5-domain_randomization/1-configuration.md index c15ef3cb7..eb135258c 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/5-domain_randomization/1-configuration.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/5-domain_randomization/1-configuration.md @@ -7,14 +7,8 @@ - Manager-Based（Compatible）任务通过 owner YAML 的 `events:` manager term 声明 reset / interval 随机化，例如 `src/unilab/conf/ppo/task/go1_joystick_flat/base.yaml`。 -- 只有 Adapted family（sharpa 及对应 hora / appo owner）仍在 - `env.domain_rand` 下配置 legacy provider 字段。 - -```bash -uv run train --algo ppo --task sharpa_inhand_grasp --sim mujoco \ - env.domain_rand.randomize_gravity=true \ - 'env.domain_rand.gravity_range=[[0.0,0.0,-10.5],[0.0,0.0,-8.5]]' -``` +- 任务也可以挂载任务级 provider，并在 `env.domain_rand` 下配置 legacy provider + 字段；目前仓内没有任务使用这条路径。 常见的生命周期边界： @@ -31,14 +25,9 @@ uv run train --algo ppo --task sharpa_inhand_grasp --sim mujoco \ ## Reset Gravity 在启用 gravity reset 随机化时使用 `--sim mujoco`；Motrix 在当前后端中 -未提供相同的 gravity 能力。该项只在 legacy provider 路径（Adapted family owner） -上可用。 - -```bash -uv run train --algo ppo --task sharpa_inhand_grasp --sim mujoco \ - env.domain_rand.randomize_gravity=true \ - 'env.domain_rand.gravity_range=[[0.0,0.0,-10.5],[0.0,0.0,-8.5]]' -``` +未提供相同的 gravity 能力。该项只在任务级 provider 路径 +（`env.domain_rand.randomize_gravity` 与 `env.domain_rand.gravity_range`）上可用， +目前仓内没有任务使用这条路径。 ## Interval Push @@ -56,8 +45,6 @@ uv run train --algo ppo --task go1_joystick_flat --sim mujoco \ 当取值范围是任务 contract 的一部分时，将其保留在 task owner YAML 中。例如， rough 四足家族的 base mass、质心、kp/kd 和 push 随机化作为 event term 声明在共享 base `src/unilab/conf/ppo/task/quadruped_joystick_rough/base.yaml`（`go2_joystick_rough` 的 backend -owner 通过 Hydra defaults 组合它），而 -`src/unilab/conf/ppo/task/sharpa_inhand/mujoco.yaml` 为 Sharpa 配置了物体缩放、摩擦和 -力扰动。 +owner 通过 Hydra defaults 组合它）。 完整的当前清单见 {doc}`0-index`。 diff --git a/docs/sphinx/source/zh_CN/2-user_guide/5-domain_randomization/2-writing_providers.md b/docs/sphinx/source/zh_CN/2-user_guide/5-domain_randomization/2-writing_providers.md index 902d0f065..a6095beec 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/5-domain_randomization/2-writing_providers.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/5-domain_randomization/2-writing_providers.md @@ -1,10 +1,9 @@ # 编写 Provider -本页描述 legacy provider 路径：只有 2 个 Adapted family（`sharpa_inhand` / -`sharpa_inhand_grasp`）仍通过任务级 -`DomainRandomizationProvider` 声明域随机化。已迁移的 Manager-Based 任务不写 -provider；它们在 owner YAML 中通过 Hydra `events:` manager term 声明随机化（见 -{doc}`0-index` 与 {doc}`1-configuration`）。 +本页描述任务级 provider 路径：自定义任务（包括托管在外部仓库中的任务） +可以通过任务级 `DomainRandomizationProvider` 声明域随机化。Manager-Based +任务不写 provider；它们在 owner YAML 中通过 Hydra `events:` manager term +声明随机化（见 {doc}`0-index` 与 {doc}`1-configuration`）。 任务级域随机化 provider 与 task env owner 放在一起。它们采样任务专属的 状态，并返回由 `DomainRandomizationManager` 消费的 plan。 @@ -56,9 +55,10 @@ def build_interval_randomization_plan(self, env, step_counter): ## 证据 -具有代表性的 provider 实现位于（全部属于 Adapted family 的兼容路径）： +provider 接口与 manager 位于： -- `src/unilab/tasks/manipulation/sharpa_inhand/rotation.py` +- `src/unilab/dr/provider.py` +- `src/unilab/dr/manager.py` 开发者 contract 详情见 {doc}`../../4-developer_guide/2-contracts/4-dr_contract`。 diff --git a/docs/sphinx/source/zh_CN/2-user_guide/8-manipulation/0-index.md b/docs/sphinx/source/zh_CN/2-user_guide/8-manipulation/0-index.md deleted file mode 100644 index 444082b7e..000000000 --- a/docs/sphinx/source/zh_CN/2-user_guide/8-manipulation/0-index.md +++ /dev/null @@ -1,20 +0,0 @@ -# 操作说明 - -针对操作（manipulation）路径的任务专属说明，提供比 {doc}`../4-tasks/3-manipulation` 中分类概览更多的运维细节。 - -::::{grid} 1 1 2 2 -:gutter: 3 - -:::{grid-item-card} 灵巧手内操作 -:link: 1-dexterous_inhand -:link-type: doc -Allegro 和 Sharpa 的 owner YAML、抓取缓存以及训练命令。 -::: - -:::: - -```{toctree} -:hidden: - -1-dexterous_inhand -``` diff --git a/docs/sphinx/source/zh_CN/2-user_guide/8-manipulation/1-dexterous_inhand.md b/docs/sphinx/source/zh_CN/2-user_guide/8-manipulation/1-dexterous_inhand.md deleted file mode 100644 index 241d699d4..000000000 --- a/docs/sphinx/source/zh_CN/2-user_guide/8-manipulation/1-dexterous_inhand.md +++ /dev/null @@ -1,123 +0,0 @@ -# 灵巧手内操作 - -本页介绍已提交的 Allegro 和 Sharpa 手内操作路径。通过 `--task` 和 `--sim` 选择后端；不要单独覆盖 `training.sim_backend`。owner YAML 始终是哪些组合被配置的内部证据。 - -## Allegro - -Allegro 旋转使用已注册的 env `AllegroInhandRotation`。旋转 owner 是 `allegro_inhand`，抓取缓存生成使用 `allegro_inhand_grasp`。 - -Owner 证据： - -- `src/unilab/conf/ppo/task/allegro_inhand/mujoco.yaml` -- `src/unilab/conf/ppo/task/allegro_inhand/motrix.yaml` -- `src/unilab/conf/ppo/task/allegro_inhand_grasp/mujoco.yaml` -- `src/unilab/conf/ppo/task/allegro_inhand_grasp/motrix.yaml` -- `src/unilab/conf/appo/task/allegro_inhand/mujoco.yaml` -- `src/unilab/conf/appo/task/allegro_inhand/motrix.yaml` - -典型流程分两个阶段：先生成抓取缓存，然后训练旋转策略。 - -```bash -uv run train --algo ppo --task allegro_inhand_grasp --sim mujoco training.no_play=true -uv run train --algo ppo --task allegro_inhand --sim mujoco training.no_play=true -``` - -PPO Allegro 路径也存在 Motrix owner YAML： - -```bash -uv run train --algo ppo --task allegro_inhand_grasp --sim motrix training.no_play=true -uv run train --algo ppo --task allegro_inhand --sim motrix training.no_play=true -``` - -旋转 owner 默认使用 `caches/allegro_grasp_50k.npy` 处的抓取缓存。要使用自定义缓存，override `env.grasp_cache_path`： - -```bash -uv run train --algo ppo --task allegro_inhand --sim mujoco \ - env.grasp_cache_path=caches/my_allegro_grasp.npy -``` - -用 `eval` 回放已训练的 checkpoint（`--load-run -1` 取最新的 run）： - -```bash -uv run eval --algo ppo --task allegro_inhand --sim mujoco --load-run -1 -uv run eval --algo appo --task allegro_inhand --sim mujoco --load-run -1 -``` - -## Sharpa - -Sharpa 旋转使用已注册的 env `SharpaInhandRotation`。当前已提交的训练路径是 MuJoCo owner 路径。 - -Owner 证据： - -- `src/unilab/conf/ppo/task/sharpa_inhand/mujoco.yaml` -- `src/unilab/conf/ppo/task/sharpa_inhand/mujoco_hora.yaml` -- `src/unilab/conf/ppo/task/sharpa_inhand_grasp/mujoco.yaml` -- `src/unilab/conf/appo/task/sharpa_inhand/mujoco.yaml` -- `src/unilab/conf/appo/task/sharpa_inhand/mujoco_hora.yaml` -- `src/unilab/conf/hora_distill/task/sharpa_inhand/mujoco.yaml` - -完整 HORA 流程分三个阶段： - -1. 生成 grasp cache。 -2. 训练 teacher policy。 -3. 需要时再蒸馏出 student policy。 - -完整 HORA teacher / student 流程以 MuJoCo owner 为主。Motrix 路径当前只承担 phase-1 PPO rotation 和 grasp cache 采集，不是完整 HORA 能力等价路径。 - -### Grasp cache 与 scale - -默认 cache 托管在 Hugging Face (`unilabsim/unilab-caches`)，首次训练时自动下载到 `src/unilab/assets/caches/`，无需手动操作。 - -如需为 HF 上没有的自定义 scale 采集 cache，或在本地重新生成，可按每个 scale 分别运行 grasp 任务（cache 文件命名为 `_.npy`）。生成的文件会落到 `src/unilab/assets/caches/`，与 HF 下载位置一致，后续训练能直接命中，**无需额外配置；但重新生成耗时较长**。 - -辅助脚本会按顺序采集每个 scale： - -```bash -bash scripts/sharpa_collect_grasps.sh 0.8 1.0 1.2 -``` - -等价的逐 scale 调用：`uv run train --algo ppo --task sharpa_inhand_grasp --sim mujoco 'env.domain_rand.scale_list=[0.8]' training.no_play=true`（对 `[1.0]`、`[1.2]` 等同理重复）。 - -Motrix 也可以采集 grasp cache（仅 phase-1 范围）： - -```bash -uv run train --algo ppo --task sharpa_inhand_grasp --sim motrix \ - 'env.domain_rand.scale_list=[1.0]' \ - env.grasp_collection_target=1000 \ - training.no_play=true -``` - -要使用自定义 cache 前缀，override `env.grasp_cache_path`： - -```bash -uv run train --algo ppo --task sharpa_inhand --sim mujoco \ - env.grasp_cache_path=caches/my_sharpa_grasp_cache -``` - -### Teacher 与 student - -用 `hora` profile 训练 HORA teacher（PPO 或 APPO）： - -```bash -uv run train --algo ppo --task sharpa_inhand --sim mujoco --profile hora training.no_play=true -uv run train --algo appo --task sharpa_inhand --sim mujoco --profile hora training.no_play=true -``` - -用 `eval --profile hora --load-run -1` 回放 teacher run： - -```bash -uv run eval --algo ppo --task sharpa_inhand --sim mujoco --profile hora --load-run -1 -uv run eval --algo appo --task sharpa_inhand --sim mujoco --profile hora --load-run -1 -``` - -Student 蒸馏由 `src/unilab/conf/hora_distill/task/sharpa_inhand/mujoco.yaml` 配置，由 `scripts/train_hora_distill.py` 实现；顶层 CLI 目前没有暴露单独的 HORA 蒸馏路由（不在 CLI `SUPPORTED_ALGOS` 中）。需要从 APPO teacher 蒸馏时，在该低层配置中设置 `teacher.algo_family=appo`。 - -常见日志目录： - -- `logs/hora_ppo/SharpaInhandRotation/` -- `logs/hora_appo/SharpaInhandRotation/` -- `logs/hora_distill/SharpaInhandRotation/` - -这里的 scale / grasp cache / DR 边界比较敏感；生命周期规则参见 {doc}`../5-domain_randomization/0-index`。 - -关于分类级别的任务页面，参见 {doc}`../4-tasks/3-manipulation`。 diff --git a/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/4-allegro_inhand.md b/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/4-allegro_inhand.md index 0dadcc93e..894e74fa1 100644 --- a/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/4-allegro_inhand.md +++ b/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/4-allegro_inhand.md @@ -1,7 +1,7 @@ -# Allegro / Sharpa 手内操作部署 +# Allegro 手内操作部署 -在 16 自由度的 Allegro 手或 17 自由度的 Sharpa 手上进行方块重定向。UniLab 以无触觉 -方式训练它们 —— 观测仅为关节状态 + 方块位姿。 +在 16 自由度的 Allegro 手上进行方块重定向。UniLab 以无触觉 +方式训练 —— 观测仅为关节状态 + 方块位姿。 ## 难点在哪 @@ -24,10 +24,10 @@ - 维度 - 硬件上的来源 * - 关节位置 - - 16（Allegro）/ 17（Sharpa） + - 16 - 编码器 * - 关节速度 - - 16 / 17 + - 16 - 编码器差分，低通滤波 * - 方块位姿（世界系） - 7 @@ -39,7 +39,7 @@ - 4 - 指令 * - 上一步动作 - - 16 / 17 + - 16 - 上一次策略输出 ``` @@ -52,7 +52,7 @@ owner 与部署运行时在观测时序上达成一致。见 ## 抓取生成器 -`4-allegro_inhand` 与 `sharpa_inhand` 两个环境都自带一个**抓取生成器**，用于采样 +`allegro_inhand` 环境自带一个**抓取生成器**，用于采样 合理的初始手部构型。硬件侧的等价物是操作员把方块放到手里 —— 请核实你的起始构型分布 与训练环境的抓取生成器输出相匹配（参见 `unilab.tasks.manipulation.allegro_inhand.grasp_gen`）。 @@ -63,8 +63,7 @@ owner 与部署运行时在观测时序上达成一致。见 ## 动作接口 操作类环境通过任务控制配置把策略动作映射为关节位置目标。Allegro 的声明由 -`src/unilab/conf/ppo/task/allegro_inhand/base.yaml` 与其 Manager-Based action term 持有； -Sharpa 当前仍由 `src/unilab/tasks/manipulation/sharpa_inhand/base.py` 持有。 +`src/unilab/conf/ppo/task/allegro_inhand/base.yaml` 与其 Manager-Based action term 持有。 部署控制器必须使用相同的关节顺序、动作缩放与限位策略。 ## 失败恢复 @@ -76,4 +75,3 @@ Sharpa 当前仍由 `src/unilab/tasks/manipulation/sharpa_inhand/base.py` 持有 - {doc}`5-onnx_runtime` - {doc}`6-domain_randomization` -- {doc}`../../2-user_guide/8-manipulation/1-dexterous_inhand` diff --git a/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/6-domain_randomization.md b/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/6-domain_randomization.md index 0d70c9adb..ab9a8d321 100644 --- a/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/6-domain_randomization.md +++ b/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/6-domain_randomization.md @@ -44,9 +44,9 @@ Manager-Based 任务在 owner YAML 的 `env.events` 中声明 reset 与 interval 随机化，由 manager 生命周期执行。示例见 `src/unilab/conf/ppo/task/quadruped_joystick_rough/base.yaml`。 -Sharpa Adapted 任务仍通过任务 provider 接入 `src/unilab/dr/manager.py`； -现有实现为 `src/unilab/tasks/manipulation/sharpa_inhand/rotation.py` 中的 -`SharpaInhandRotationDRProvider`。两条路径的能力边界见 +任务也可以通过任务级 provider（`src/unilab/dr/provider.py` 中的 +`DomainRandomizationProvider` 接口）接入 `src/unilab/dr/manager.py`；目前仓内 +没有任务使用这条路径。两条路径的能力边界见 {doc}`../../4-developer_guide/2-contracts/4-dr_contract`。 ## 配方：起始范围 diff --git a/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/8-latency_budget.md b/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/8-latency_budget.md index 74affe62f..4b9611f94 100644 --- a/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/8-latency_budget.md +++ b/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/8-latency_budget.md @@ -9,7 +9,6 @@ | --- | --- | --- | | 单步动作延迟 | task owner 中 Manager action term 的 `simulate_action_latency` 声明 | 执行上一步动作而非当前动作。 | | G1 WBT 观测历史 | `src/unilab/conf/sac/task/g1_wbt_obs/mujoco.yaml` 中逐 term 的 `history_length` | 为 `base_ang_vel`、`joint_pos`、`joint_vel` 与 `actions` 提供逐项历史。 | -| Sharpa 触觉接触延迟 | Sharpa 手内配置中的 `domain_rand.contact_latency` | 为采样到的接触通道保留上一步的触觉接触值。 | | 观测历史顺序守护 | `tests/scripts/test_obs_alignment_g1_wbt.py` | 断言 G1 WBT actor 观测按逐项最旧优先展平。 | ## 动作延迟 diff --git a/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/9-troubleshooting.md b/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/9-troubleshooting.md index b0d832e3b..33a9f1077 100644 --- a/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/9-troubleshooting.md +++ b/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/9-troubleshooting.md @@ -31,7 +31,7 @@ 3. **观测布局不匹配。** 将硬件回路装配出的内容与 owner 的 `env.observations.actor.terms` 比较——先比分项顺序，再比逐项历史顺序。 -## Allegro / Sharpa 手内操作中方块掉落 +## Allegro 手内操作中方块掉落 | 可能原因 | 检查 | 修复 | |---|---|---| diff --git a/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/1-from_isaac_lab.md b/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/1-from_isaac_lab.md index 795433b2f..27459eba9 100644 --- a/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/1-from_isaac_lab.md +++ b/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/1-from_isaac_lab.md @@ -156,9 +156,6 @@ entry 的 production task 名称抛出 `KeyError`，因此新增 production 注 - 36 个 task 为 **Compatible**（`target=complete`）：Hydra owner YAML 物化 canonical NumPy Manager-Based runtime。 -- 2 个 task 为 **Adapted**（`target=compatibility`）：`SharpaInhandRotation` 和 - `SharpaInhandRotationGrasp` 各自把 tactile/contact/cache 行为保留在一个冻结的 - 兼容 factory 后面；只有当正式能力存在时才迁移。 ## 仓库证据 diff --git a/docs/sphinx/source/zh_CN/4-developer_guide/0-index.md b/docs/sphinx/source/zh_CN/4-developer_guide/0-index.md index e90d5a9e3..421d20157 100644 --- a/docs/sphinx/source/zh_CN/4-developer_guide/0-index.md +++ b/docs/sphinx/source/zh_CN/4-developer_guide/0-index.md @@ -93,6 +93,5 @@ DR provider 的 init、reset、interval 与后端能力边界。 5-contributing_workflow 6-agent_quick_reference 7-motion_assets -8-motrix_contact_sensor 9-sim2sim_contract_status ``` diff --git a/docs/sphinx/source/zh_CN/4-developer_guide/2-contracts/3-task_owner.md b/docs/sphinx/source/zh_CN/4-developer_guide/2-contracts/3-task_owner.md index 6782cb94a..992471786 100644 --- a/docs/sphinx/source/zh_CN/4-developer_guide/2-contracts/3-task_owner.md +++ b/docs/sphinx/source/zh_CN/4-developer_guide/2-contracts/3-task_owner.md @@ -9,7 +9,6 @@ `src/unilab/conf/{ppo,appo}/task//.yaml`。 - Off-policy 算法（SAC / TD3 / FlashSAC）各自有独立的配置树： `src/unilab/conf//task//.yaml`。 -- `src/unilab/conf/hora_distill/` 对其支持的任务遵循相同的 owner YAML identity 规则。 ## 必需语义 diff --git a/docs/sphinx/source/zh_CN/4-developer_guide/2-contracts/4-dr_contract.md b/docs/sphinx/source/zh_CN/4-developer_guide/2-contracts/4-dr_contract.md index 205df9727..74d080650 100644 --- a/docs/sphinx/source/zh_CN/4-developer_guide/2-contracts/4-dr_contract.md +++ b/docs/sphinx/source/zh_CN/4-developer_guide/2-contracts/4-dr_contract.md @@ -123,4 +123,4 @@ actuator 的机制泄漏到共享 payload 里。 `src/unilab/dr/__init__.py` 再导出 - DR manager：`src/unilab/dr/manager.py` - Backend 接口：`unisim.backend.base` -- 示例 provider：`src/unilab/tasks/manipulation/sharpa_inhand/rotation.py` +- Provider 接口：`src/unilab/dr/provider.py` diff --git a/docs/sphinx/source/zh_CN/4-developer_guide/3-extending/3-new_algorithm.md b/docs/sphinx/source/zh_CN/4-developer_guide/3-extending/3-new_algorithm.md index c179841b3..af4576425 100644 --- a/docs/sphinx/source/zh_CN/4-developer_guide/3-extending/3-new_algorithm.md +++ b/docs/sphinx/source/zh_CN/4-developer_guide/3-extending/3-new_algorithm.md @@ -16,7 +16,8 @@ - 每个 task×backend 组合对应一个 owner YAML： `src/unilab/conf//task//.yaml`； - 替换策略 / 算法实现类可以走 owner YAML 中的 `class_name` dotted path - （现网实例：`uni_rl.algos.hora:HoraActorModel`），无需新增代码路径。 + （仓内实例：`src/unilab/conf/ppo/config.yaml` 中的 + `uni_rl.algos.rsl_rl_ppo:FinalObservationAwarePPO`），无需新增代码路径。 ### 2. `runtime_resolver`：算法代码放在自己的仓库 @@ -38,11 +39,6 @@ algo: - 解析发生在 uni_rl 侧（`uni_rl.algos.appo.runtime` / `uni_rl.algos.rsl_rl_runtime`），dotted path 可以指向任何可 import 的模块。 -仓库内实例：`src/unilab/conf/appo/task/sharpa_inhand/mujoco_hora.yaml` 指向 -`unilab.scripts.play_hora_appo:resolve_hora_appo_runtime`；HORA 的 SAC / PPO -变体分别指向 `uni_rl.algos.hora.sac:resolve_hora_sac_runtime` 与 -`uni_rl.algos.hora.rsl_rl:resolve_hora_ppo_runtime`。 - ### 3. fork unilab_rl：改 `uni_rl/algos/` 只有当需要改动 runner / learner / collector 的共享实现（例如新的 IPC @@ -68,8 +64,8 @@ fail-closed，报错信息列出全部可用 algo（内置 + 约定发现的） 注意： -- 只有 conf 目录而没有入口脚本的 config 树（如 `hora_distill`） - 不可路由——它们不是独立的 CLI algo。 +- 只有 conf 目录而没有入口脚本的 config 树不可路由——它们不是独立的 + CLI algo。 - 内置算法的特殊脚本名映射保留不变：`ppo` → `train_rsl_rl.py`、 `appo` → `train_appo.py`。 - `src/unilab/structured_configs.py` 的 dataclass 是**可选**的约定俗成镜像， diff --git a/docs/sphinx/source/zh_CN/4-developer_guide/6-agent_quick_reference.md b/docs/sphinx/source/zh_CN/4-developer_guide/6-agent_quick_reference.md index 56ab753c2..58448a82e 100644 --- a/docs/sphinx/source/zh_CN/4-developer_guide/6-agent_quick_reference.md +++ b/docs/sphinx/source/zh_CN/4-developer_guide/6-agent_quick_reference.md @@ -12,7 +12,6 @@ - APPO 入口：`src/unilab/scripts/train_appo.py` - SAC / TD3 / FlashSAC 入口：`src/unilab/scripts/train_sac.py` / `src/unilab/scripts/train_td3.py` / `src/unilab/scripts/train_flashsac.py` -- HORA 蒸馏入口：`scripts/train_hora_distill.py` ## 需要记住的契约 diff --git a/docs/sphinx/source/zh_CN/4-developer_guide/7-motion_assets.md b/docs/sphinx/source/zh_CN/4-developer_guide/7-motion_assets.md index 76a51e4ed..a391d4c6d 100644 --- a/docs/sphinx/source/zh_CN/4-developer_guide/7-motion_assets.md +++ b/docs/sphinx/source/zh_CN/4-developer_guide/7-motion_assets.md @@ -80,7 +80,7 @@ env: 机器人二进制网格和纹理（例如 `.STL`、`.obj`、`.png`）采用相同方式外置， 托管在 Hugging Face 数据集仓库 [unilabsim/unilab-robots](https://huggingface.co/datasets/unilabsim/unilab-robots)。 -已注册的机器人为 a2、allegro_hand、g1、go2、sharpa_wave、 +已注册的机器人为 a2、allegro_hand、g1、go2、 x2（见 `src/unilab/assets/hub.py` 的 `ROBOT_ASSET_SPECS`）。它们的 mesh/纹理目录在首次使用时按需下载，落盘到原始路径（例如 G1 的 `src/unilab/assets/robots/g1/assets/` 与 `robots/g1/textures/`），因此 XML 中的 diff --git a/docs/sphinx/source/zh_CN/4-developer_guide/8-motrix_contact_sensor.md b/docs/sphinx/source/zh_CN/4-developer_guide/8-motrix_contact_sensor.md deleted file mode 100644 index 38c487adf..000000000 --- a/docs/sphinx/source/zh_CN/4-developer_guide/8-motrix_contact_sensor.md +++ /dev/null @@ -1,112 +0,0 @@ -# Motrix Contact Sensor 适配笔记 - -## 背景 - -Sharpa Inhand Rotation 任务的触觉观测依赖 contact sensor 数据（指尖与物体的接触力）。该代码在 MuJoCo 后端下工作正常，但 Motrix 后端的 contact sensor 返回格式不同，读取路径需要同时兼容两种形状。 - -## Contact Sensor 配置 - -指尖 contact sensor 定义在机器人场景 XML -`src/unilab/assets/robots/sharpa_wave/right_sharpa_wave.xml` 中： - -```xml - - -``` - -参数说明： - -- `num="1"` — 最多报告 1 个接触点 -- `data="force"` — 只返回力数据 -- `reduce="netforce"` — 合成为单个接触点的合力，全局坐标系 - -## 两个后端的返回格式差异 - -### MuJoCo - -返回 `(num_envs, 3)` 的力向量，全局坐标系： - -``` -[[fx, fy, fz], # env 0 - [fx, fy, fz], # env 1 - ...] -``` - -### Motrix - -返回 `(num_envs, 1 + num * stride)` 的扁平数组，**首个元素为实际接触点数**： - -``` -[[count, fx, fy, fz], # env 0, shape = (4,) 因为 num=1, stride=3 - [count, fx, fy, fz], # env 1 - ...] -``` - -对于 `reduce="netforce"` + `data="force"` + `num=1`： - -- shape = `(num_envs, 4)` -- `[0]` = 接触点数（0 或 1） -- `[1:4]` = 全局坐标系的合力向量（`netforce` 模式使用全局坐标） - -### 多接触点的通用格式（Motrix） - -当 `num=4, data="force pos normal tangent"` 且无 reduce 时： - -``` -shape = (num_envs, 1 + 4 * 12) = (num_envs, 49) - -[count, - f1_normal, f1_tangent0, f1_tangent1, # contact 1 force (contact 坐标系) - f1_x, f1_y, f1_z, # contact 1 position - f1_nx, f1_ny, f1_nz, # contact 1 normal - f1_tx, f1_ty, f1_tz, # contact 1 tangent - f2_normal, f2_tangent0, f2_tangent1, # contact 2 ... - ... - padding_zeros] # 不足 num 个的用 0 填充 -``` - -注意：无 reduce 时力是 **contact 坐标系**（法向标量 + 两个切向标量），不是全局 xyz 向量。只有 `reduce="netforce"` 返回全局坐标系的力向量。 - -## Motrix `reduce` 模式参考 - -| reduce 模式 | 坐标系 | 返回内容 | -| --- | --- | --- | -| `netforce` | 全局 | 单个合成的合力向量 | -| `maxforce` | contact | 力最大的那个接触点 | -| `mindist` | contact | 穿透最浅的那个接触点 | -| 无 | contact | 前 `num` 个接触点 | - -其中只有 `netforce` 返回全局坐标系力向量，其他模式返回 contact 坐标系数据（法向标量 + 两个切向标量）。 - -## 单一 norm 分支为何不够 - -env 通过 `src/unilab/tasks/manipulation/sharpa_inhand/base.py` 中的 -`_read_tactile_force()` → `_extract_sensor_scalar()` 读取触觉力。该 helper 目前对任意 `(N, >=3)` 数组都用 `np.linalg.norm(data[:, :3], axis=1)` 折叠。 - -如果 env 仍把两种后端形状都走这一个分支，MuJoCo 的 `(N, 3)` 是正确的（对真实力向量取 norm），但 Motrix 的 `(N, 4)` 会出错：`data[:, :3]` 取到的是 `[count, fx, fy]`——把接触点数当成了力分量，并且漏掉了 `fz`。正确做法不是在 env 里按形状特判，而是把每个后端的布局知识下沉到 backend 方法。 - -## 建议的契约：用 backend 方法返回力大小 - -在 `SimBackend` 接口（`unisim.backend.base`）增加 -`get_contact_force_magnitude(sensor_name) -> np.ndarray`，返回 `(num_envs,)` 的标量力大小。每个 backend 按自己的数据布局实现： - -- **MuJoCo**（`unisim.backend.mujoco.backend`）：对 `get_sensor_data(name)` 返回的 3D 力向量取范数。 -- **Motrix**（`unisim.backend.motrix.backend`）：按 reduce 模式解释布局： - - `reduce="netforce"`：取 `[1:4]` 后取 norm。 - - 无 reduce 多 contact：对各接触力求和后取 norm。 - - `reduce="maxforce"`：取力最大的接触点。 - -env 层的 `_read_tactile_force()` 对 contact sensor 走 `get_contact_force_magnitude()` 路径，普通标量 sensor 仍走 `_extract_sensor_scalar()`。 - -这样 env 代码保持 backend 无关，并与高风险区的不变量一致：env 层只能调用 `SimBackend` 上已声明的方法——不向 env 泄漏后端功能，新增 backend 只需实现该接口方法。`get_sensor_data` 已是 `SimBackend` 上声明的方法，因此该提案是在同一边界上扩展，而不是伸手去调用 backend 子类。 - -## 相关文件 - -| 文件 | 说明 | -| --- | --- | -| `src/unilab/tasks/manipulation/sharpa_inhand/base.py` | `_extract_sensor_scalar()`, `_read_tactile_force()` | -| `src/unilab/tasks/manipulation/sharpa_inhand/rotation.py` | reward 计算，virtual torque | -| `src/unilab/assets/robots/sharpa_wave/right_sharpa_wave.xml` | contact sensor XML 定义 | -| `unisim.backend.motrix.backend` | Motrix `get_sensor_data()` | -| `unisim.backend.mujoco.backend` | MuJoCo `get_sensor_data()` | -| `unisim.backend.base` | `SimBackend` 接口 | diff --git a/docs/sphinx/source/zh_CN/4-developer_guide/9-sim2sim_contract_status.md b/docs/sphinx/source/zh_CN/4-developer_guide/9-sim2sim_contract_status.md index 23383cecc..4c61f1c78 100644 --- a/docs/sphinx/source/zh_CN/4-developer_guide/9-sim2sim_contract_status.md +++ b/docs/sphinx/source/zh_CN/4-developer_guide/9-sim2sim_contract_status.md @@ -28,7 +28,7 @@ uv run scripts/audit_sim2sim_contracts.py | Task | 判定 | 分歧 | |---|---|---| -| allegro_inhand · allegro_inhand_grasp · g1_climb_tracking · g1_motion_tracking · g1_wall_flip_tracking · go1_joystick_rough · go2_footstand · go2_handstand · go2_joystick_flat · go2_joystick_rough · go2w_joystick_flat · go2w_joystick_rough · sharpa_inhand · sharpa_inhand_grasp | ✅ | 无 | +| allegro_inhand · allegro_inhand_grasp · g1_climb_tracking · g1_motion_tracking · g1_wall_flip_tracking · go1_joystick_rough · go2_footstand · go2_handstand · go2_joystick_flat · go2_joystick_rough · go2w_joystick_flat · go2w_joystick_rough | ✅ | 无 | | g1_box_tracking | ❌ | `empirical_normalization` false↔true；`obs_groups` critic 组差异 | | g1_flip_tracking | ❌ | `empirical_normalization` true↔false；`obs_groups`；`action_scale` 29 维↔默认 0.25；`sampling_mode` 两后端运行时同为 `start`（无害） | | g1_walk_flat | ❌ | `env.actions.joint_pos.scale` 0.25↔0.5；`empirical_normalization` false↔true；`obs_groups` | @@ -46,9 +46,8 @@ uv run scripts/audit_sim2sim_contracts.py ## 其它配置树 -`src/unilab/conf/sac/task`、`src/unilab/conf/td3/task`、`src/unilab/conf/flashsac/task`、 -`src/unilab/conf/hora_distill/task` 均无 mujoco↔motrix -配对，sim2sim 不适用。 +`src/unilab/conf/sac/task`、`src/unilab/conf/td3/task`、`src/unilab/conf/flashsac/task` +均无 mujoco↔motrix 配对，sim2sim 不适用。 ## 字段语义速查 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 301fdb91d..91bc45919 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 @@ -70,8 +70,6 @@ uv run scripts/generate_support_matrix.py --write | PPO (torch) | `g1_wall_flip_tracking` (G1 wall flip tracking) | Tested | - | Tested | - | - | - | - | | PPO (torch) | `x2_wall_flip_tracking` (X2 wall flip tracking) | Tested | - | Tested | - | - | - | - | | PPO (torch) | `allegro_inhand` (Allegro in-hand) | Tested | - | Tested | - | - | - | - | -| PPO (torch) | `sharpa_inhand` (Sharpa in-hand) | Tested | - | Tested | - | - | - | - | -| PPO (torch) | `sharpa_inhand_grasp` (Sharpa in-hand grasp) | Tested | - | Tested | - | - | - | - | | PPO (torch) | `a2_joystick_flat` (a2 joystick flat) | Tested | - | - | - | - | - | - | | PPO (torch) | `allegro_inhand_grasp` (allegro inhand grasp) | Tested | - | Tested | - | - | - | - | | PPO (torch) | `g1_23dof_box_tracking` (g1 23dof box tracking) | Tested | - | Tested | - | - | - | - | @@ -97,7 +95,6 @@ uv run scripts/generate_support_matrix.py --write | APPO (torch) | `g1_flip_tracking` (G1 flip tracking) | Tested | - | Tested | - | - | - | - | | APPO (torch) | `g1_wall_flip_tracking` (G1 wall flip tracking) | Tested | - | Tested | - | - | - | - | | APPO (torch) | `allegro_inhand` (Allegro in-hand) | Tested | - | Tested | - | - | - | - | -| APPO (torch) | `sharpa_inhand` (Sharpa in-hand) | Tested | - | Tested | - | - | - | - | | APPO (torch) | `g1_23dof_climb_tracking` (g1 23dof climb tracking) | Tested | - | Tested | - | - | - | - | | APPO (torch) | `g1_23dof_flip_tracking` (g1 23dof flip tracking) | Tested | - | Tested | - | - | - | - | | APPO (torch) | `g1_23dof_motion_tracking` (g1 23dof motion tracking) | Tested | - | Tested | - | - | - | - | diff --git a/pyproject.toml b/pyproject.toml index 5fd2f7331..9860eab6a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,7 +15,6 @@ source-exclude = [ "/src/unilab/assets/robots/g1/textures", "/src/unilab/assets/robots/go1/assets", "/src/unilab/assets/robots/go2/assets", - "/src/unilab/assets/robots/sharpa_wave/meshes", "/src/unilab/assets/robots/x2/meshes", ] @@ -43,7 +42,7 @@ dependencies = [ # Physics implementations are provided by the independently released # unisim-core package from the production PyPI index. "unisim-core>=1.1.4", - # RL algorithms and async runtimes (PPO/APPO/SAC/TD3/HORA runners, + # RL algorithms and async runtimes (PPO/APPO/SAC/TD3 runners, # collectors, IPC, logging) live in the independently released uni-rl # package (distribution name ``unilab-rl``), consumed via the injected # env contract (uni_rl.env_contract.EnvFactory). Published on PyPI. diff --git a/scripts/benchmark/core/task_names.py b/scripts/benchmark/core/task_names.py index 2bc007acc..88e0f15e3 100644 --- a/scripts/benchmark/core/task_names.py +++ b/scripts/benchmark/core/task_names.py @@ -3,7 +3,6 @@ from dataclasses import dataclass from unilab.envs import ManagerBasedRlEnvCfg -from unilab.tasks.manipulation.sharpa_inhand.rotation import SharpaInhandRotationCfg @dataclass(frozen=True) @@ -37,12 +36,6 @@ class LocomotionTaskSpec: config_cls=ManagerBasedRlEnvCfg, model_file="src/unilab/assets/robots/g1/scene_flat.xml", ), - "sharpa_inhand": LocomotionTaskSpec( - owner_task_id="sharpa_inhand", - env_task_name="SharpaInhandRotation", - display_name="sharpa_inhand", - config_cls=SharpaInhandRotationCfg, - ), } _TASK_ALIASES = {spec.env_task_name: spec.owner_task_id for spec in _TASK_SPECS.values()} _TASK_ALIASES.update({f"task={task_id}/mujoco": task_id for task_id in _TASK_SPECS}) diff --git a/scripts/benchmark/env/benchmark_env_step.py b/scripts/benchmark/env/benchmark_env_step.py index cfee0385f..d9d807600 100644 --- a/scripts/benchmark/env/benchmark_env_step.py +++ b/scripts/benchmark/env/benchmark_env_step.py @@ -223,32 +223,6 @@ def _materialize_g1_rough_benchmark_scene() -> str: return str(output_path) -def _materialize_sharpa_motrix_scene() -> str: - import xml.etree.ElementTree as ET - - source_dir = ROOT_DIR / "src" / "unilab" / "assets" / "robots" / "sharpa_wave" - output_dir = Path("/tmp/unilab_benchmark_sharpa_scene") - output_dir.mkdir(parents=True, exist_ok=True) - - robot_tree = ET.parse(source_dir / "right_sharpa_wave.xml") - robot_root = robot_tree.getroot() - compiler = robot_root.find("compiler") - if compiler is not None: - compiler.set("meshdir", str((source_dir / "meshes").resolve())) - - contact = robot_root.find("contact") - if contact is not None: - for exclude in list(contact.findall("exclude")): - if exclude.get("body1") == exclude.get("body2"): - contact.remove(exclude) - robot_tree.write(output_dir / "right_sharpa_wave.xml") - - scene_tree = ET.parse(source_dir / "scene.xml") - output_path = output_dir / "scene.xml" - scene_tree.write(output_path) - return str(output_path) - - def _go1_cfg(backend: str, config_overrides: list[str]) -> Any: from unilab.envs import ManagerBasedRlEnvCfg @@ -318,75 +292,12 @@ def _g1_motion_tracking_cfg(backend: str, config_overrides: list[str]) -> Any: ) -def _sharpa_inhand_cfg(backend: str, config_overrides: list[str]) -> Any: - from hydra import compose, initialize_config_dir - from hydra.core.global_hydra import GlobalHydra - - from unilab.base.config_adapter import BackendAdapter - from unilab.base.registry import apply_cfg_overrides - from unilab.tasks.manipulation.sharpa_inhand.rotation import SharpaInhandRotationCfg - - yaml_backend = _hydra_yaml_backend(backend) - GlobalHydra.instance().clear() - with initialize_config_dir( - config_dir=str(ROOT_DIR / "src" / "unilab" / "conf" / "ppo"), version_base="1.3" - ): - owner_cfg = compose( - config_name="config", - overrides=[ - f"task=sharpa_inhand/{yaml_backend}", - "env.grasp_cache_path=/tmp/unilab_benchmark_sharpa_grasp", - *config_overrides, - "hydra.run.dir=.", - "hydra.output_subdir=null", - "hydra/job_logging=disabled", - "hydra/hydra_logging=disabled", - ], - ) - - env_cfg_override = BackendAdapter( - owner_cfg, - root_dir=ROOT_DIR, - algo_name="ppo", - ).build_task_env_cfg_override() - - cfg = SharpaInhandRotationCfg() - apply_cfg_overrides(cfg, env_cfg_override) - return cfg - - -def _ensure_sharpa_benchmark_grasp_cache(cfg: Any, _: str) -> None: - from unilab.tasks.manipulation.sharpa_inhand.base import ( - SOURCE_DEFAULT_HAND_JOINT_POS_DEG, - resolve_grasp_cache_file, - ) - - if not str(cfg.grasp_cache_path).startswith("/tmp/unilab_benchmark_sharpa_grasp"): - return - - hand_qpos = np.deg2rad(np.asarray(SOURCE_DEFAULT_HAND_JOINT_POS_DEG, dtype=np.float64)) - object_height = 0.5 * (float(cfg.reset_height_lower) + float(cfg.reset_height_upper)) - object_pose = np.asarray([0.0, 0.0, object_height, 1.0, 0.0, 0.0, 0.0], dtype=np.float64) - cache_row = np.concatenate([hand_qpos, object_pose], axis=0) - - for scale_value in np.asarray(cfg.domain_rand.scale_list, dtype=np.float64): - cache_file = resolve_grasp_cache_file(cfg.grasp_cache_path, float(scale_value)) - cache_file.parent.mkdir(parents=True, exist_ok=True) - np.save(cache_file, cache_row[None, :]) - - def _g1_walk_env_cls() -> type: from unilab.tasks.locomotion.g1 import make_g1_walk_env return make_g1_walk_env -def _sharpa_inhand_env_cls() -> type: - from unilab.tasks.manipulation.sharpa_inhand.rotation import SharpaInhandRotationEnv - - return SharpaInhandRotationEnv - - TASK_CONFIGS: dict[str, TaskConfig] = { "go1": TaskConfig( task_id="go1_joystick_flat", @@ -445,14 +356,6 @@ def _sharpa_inhand_env_cls() -> type: env_cls_factory=_manager_env_cls, backends=("mujoco", "motrix"), ), - "sharpa_inhand": TaskConfig( - task_id="sharpa_inhand", - env_name="SharpaInhandRotation", - cfg_factory=_sharpa_inhand_cfg, - env_cls_factory=_sharpa_inhand_env_cls, - cfg_finalizer=_ensure_sharpa_benchmark_grasp_cache, - backends=("mujoco", "motrix", "mjwarp"), - ), } # Default benchmark parameters @@ -470,7 +373,6 @@ def _sharpa_inhand_env_cls() -> type: "go2w": "#72B7B2", "go2w_rough": "#499894", "go2w_rough_tiles": "#499894", - "sharpa_inhand": "#D37295", } BACKEND_STYLES = { "mujoco": {"marker": "o", "linestyle": "-", "hatch": "//"}, @@ -818,8 +720,6 @@ def _short_task_label(task_name: str) -> str: name = task_name.lower() if "motiontracking" in name: return "g1_mt" - if name.startswith("sharpainhand"): - return "sharpa_inhand" if "rough" in name and name.startswith("g1"): return "g1_rough" if name.startswith("go2w") and "roughtiles" in name: diff --git a/scripts/benchmark/env/benchmark_sharpa_init_dr_construct.py b/scripts/benchmark/env/benchmark_sharpa_init_dr_construct.py deleted file mode 100644 index 4008d0d83..000000000 --- a/scripts/benchmark/env/benchmark_sharpa_init_dr_construct.py +++ /dev/null @@ -1,612 +0,0 @@ -#!/usr/bin/env python3 -"""Benchmark Sharpa-hand init-DR construction cost vs variant count. - -What this benchmark measures: -- `construct_only`: time spent in `create_env(...)` only. -- `construct_plus_pool`: time spent in `create_env(...)` plus the first lazy - `BatchEnvPool` materialization via `env.init_state()`. - -Why variant count matters: -- Sharpa init-DR compiles `variant_count` scale-specific MuJoCo models. -- UniLab then expands env-to-variant assignments into a per-env model sequence - before constructing `BatchEnvPool`. - -Relevant `mujoco-uni-runtime` constraint from source: -- `BatchEnvPool(model=...)` accepts either - - one `MjModel`, or - - a sequence of `MjModel` with length `1` or `nbatch`. - -See: -- `mujoco_uni/runtime/batch.py` (re-exported as `mujoco_uni.batch_env`) -- `mujoco_uni/native/batch_env.cc` - -Usage: - uv run scripts/benchmark/env/benchmark_sharpa_init_dr_construct.py - uv run scripts/benchmark/env/benchmark_sharpa_init_dr_construct.py --env-nums 256,512,1024 - uv run scripts/benchmark/env/benchmark_sharpa_init_dr_construct.py --variant-counts 1,2,4,8 - uv run scripts/benchmark/env/benchmark_sharpa_init_dr_construct.py --measure construct_plus_pool -""" - -from __future__ import annotations - -import argparse -import gc -import json -import os -import sys -import tempfile -import time -from contextlib import contextmanager -from dataclasses import asdict, dataclass -from datetime import datetime, timezone -from pathlib import Path -from statistics import mean, median, pstdev -from typing import Any, Iterator - -import numpy as np -from hydra import compose, initialize_config_dir -from hydra.core.global_hydra import GlobalHydra - -ROOT_DIR = Path(__file__).resolve().parents[3] -if str(ROOT_DIR) not in sys.path: - sys.path.append(str(ROOT_DIR)) - -os.environ.setdefault("MPLCONFIGDIR", str(Path(tempfile.gettempdir()) / "unilab-matplotlib")) - -plt: Any | None = None -try: - import matplotlib as _matplotlib - - _matplotlib.use("Agg") - import matplotlib.pyplot as _plt - - plt = _plt -except Exception: - plt = None - -from scripts.benchmark.core.device_info import get_device_info_dict, get_device_info_line - -DEFAULT_ENV_NUMS = [2**power for power in range(8, 14)] -DEFAULT_VARIANT_COUNTS = [128, 256] -DEFAULT_SCALE_LOWER = 0.5 -DEFAULT_SCALE_UPPER = 0.8 -DEFAULT_OUTPUT_DIR = ROOT_DIR / "scripts" / "benchmark" / "outputs" / "sharpa_init_dr_construct" - - -@dataclass -class ConstructRecord: - measure: str - mode: str - variant_count: int - num_envs: int - scale_list: list[float] - repeats: int - samples_sec: list[float] - mean_sec: float - median_sec: float - std_sec: float - min_sec: float - max_sec: float - init_randomization_applied: bool - model_variant_count: int - pool_built: bool - - -def _parse_env_nums(value: str | None) -> list[int]: - if value: - env_nums = [int(part.strip()) for part in value.split(",") if part.strip()] - else: - env_nums = list(DEFAULT_ENV_NUMS) - if not env_nums: - raise ValueError("env nums cannot be empty") - if any(num_envs <= 0 for num_envs in env_nums): - raise ValueError(f"env nums must be positive, got {env_nums}") - return env_nums - - -def _parse_variant_counts(value: str | None) -> list[int]: - if value: - counts = [int(part.strip()) for part in value.split(",") if part.strip()] - else: - counts = list(DEFAULT_VARIANT_COUNTS) - if not counts: - raise ValueError("variant counts cannot be empty") - if any(count <= 0 for count in counts): - raise ValueError(f"variant counts must be positive, got {counts}") - deduped: list[int] = [] - for count in counts: - if count not in deduped: - deduped.append(count) - return deduped - - -def _compose_cfg(task: str, *, lower: float, upper: float, variant_count: int): - config_dir = str(ROOT_DIR / "src" / "unilab" / "conf" / "ppo") - scale_list = np.linspace(lower, upper, variant_count, dtype=np.float64) - scale_override = ",".join(f"{float(scale):g}" for scale in scale_list) - overrides = [ - f"task={task}", - f"env.domain_rand.scale_list=[{scale_override}]", - "hydra.run.dir=.", - "hydra.output_subdir=null", - "hydra/job_logging=disabled", - "hydra/hydra_logging=disabled", - ] - - GlobalHydra.instance().clear() - with initialize_config_dir(config_dir=config_dir, version_base="1.3"): - return compose(config_name="config", overrides=overrides) - - -@contextmanager -def _init_dr_mode(enabled: bool) -> Iterator[None]: - from unilab.tasks.manipulation.sharpa_inhand.rotation import SharpaInhandRotationDRProvider - - original = SharpaInhandRotationDRProvider.build_init_randomization_plan - if enabled: - yield - return - - def disabled_build_init_randomization_plan(self: Any, env: Any) -> None: - del self, env - return None - - SharpaInhandRotationDRProvider.build_init_randomization_plan = ( - disabled_build_init_randomization_plan - ) - try: - yield - finally: - SharpaInhandRotationDRProvider.build_init_randomization_plan = original - - -@contextmanager -def _synthetic_grasp_cache_mode(enabled: bool) -> Iterator[None]: - from unilab.tasks.manipulation.sharpa_inhand.rotation import SharpaInhandRotationDRProvider - - original = SharpaInhandRotationDRProvider._load_grasp_cache - if not enabled: - yield - return - - def synthetic_load_grasp_cache(self: Any, env: Any) -> np.ndarray: - del self - num_scales = int(env._num_scales) - hand_qpos = np.asarray(env.default_angles, dtype=np.float64) - object_height = 0.5 * ( - float(env.cfg.reset_height_lower) + float(env.cfg.reset_height_upper) - ) - object_pose = np.array([0.0, 0.0, object_height, 1.0, 0.0, 0.0, 0.0], dtype=np.float64) - row = np.concatenate([hand_qpos, object_pose], axis=0) - return np.repeat(row[None, :], num_scales, axis=0) - - SharpaInhandRotationDRProvider._load_grasp_cache = synthetic_load_grasp_cache - try: - yield - finally: - SharpaInhandRotationDRProvider._load_grasp_cache = original - - -def _cleanup_env(env: Any) -> None: - backend = getattr(env, "_backend", None) - pool = getattr(backend, "_pool", None) - if pool is not None: - pool.close() - if backend is not None: - backend._pool = None - close = getattr(env, "close", None) - if close is not None: - close() - del env - gc.collect() - - -def _construct_once( - cfg: Any, - *, - num_envs: int, - init_dr_enabled: bool, - force_pool: bool, -) -> tuple[float, dict[str, Any]]: - from unilab.base.config_adapter import BackendAdapter, create_env - from unilab.training import ensure_registries - - ensure_registries() - - adapter = BackendAdapter(cfg, root_dir=ROOT_DIR, algo_name="ppo") - env_cfg_override = adapter.build_task_env_cfg_override() - task_name = str(cfg.training.task_name) - sim_backend = str(cfg.training.sim_backend) - - with _init_dr_mode(init_dr_enabled), _synthetic_grasp_cache_mode(force_pool): - t0 = time.perf_counter() - env = create_env( - cfg, - num_envs=num_envs, - env_cfg_override=env_cfg_override, - sim_backend=sim_backend, - task_name=task_name, - ) - if force_pool: - env.init_state() - elapsed = time.perf_counter() - t0 - - backend = getattr(env, "_backend", None) - meta = { - "init_randomization_applied": bool(getattr(env, "_init_randomization_applied", False)), - "model_variant_count": len(getattr(backend, "_model_variants", ())), - "pool_built": getattr(backend, "_pool", None) is not None, - } - _cleanup_env(env) - return elapsed, meta - - -def _summarize_record( - *, - measure: str, - mode: str, - variant_count: int, - num_envs: int, - lower: float, - upper: float, - samples: list[float], - meta: dict[str, Any], -) -> ConstructRecord: - return ConstructRecord( - measure=measure, - mode=mode, - variant_count=variant_count, - num_envs=num_envs, - scale_list=np.linspace(lower, upper, variant_count, dtype=np.float64).tolist(), - repeats=len(samples), - samples_sec=[float(sample) for sample in samples], - mean_sec=float(mean(samples)), - median_sec=float(median(samples)), - std_sec=float(pstdev(samples)) if len(samples) > 1 else 0.0, - min_sec=float(min(samples)), - max_sec=float(max(samples)), - init_randomization_applied=bool(meta["init_randomization_applied"]), - model_variant_count=int(meta["model_variant_count"]), - pool_built=bool(meta["pool_built"]), - ) - - -def run_benchmark( - *, - task: str, - measure: str, - env_nums: list[int], - variant_counts: list[int], - scale_lower: float, - scale_upper: float, - repeats: int, - warmup: int, -) -> list[ConstructRecord]: - if repeats <= 0: - raise ValueError(f"repeats must be positive, got {repeats}") - if warmup < 0: - raise ValueError(f"warmup must be non-negative, got {warmup}") - if scale_lower <= 0.0 or scale_upper <= 0.0: - raise ValueError( - f"scale bounds must be positive, got lower={scale_lower}, upper={scale_upper}" - ) - - for variant_count in variant_counts: - indivisible = [num_envs for num_envs in env_nums if num_envs % variant_count != 0] - if indivisible: - raise ValueError( - f"All num_envs must be divisible by variant_count={variant_count}, got {indivisible}" - ) - - force_pool = measure == "construct_plus_pool" - records: list[ConstructRecord] = [] - baseline_cfg = _compose_cfg(task, lower=scale_lower, upper=scale_lower, variant_count=1) - variant_cfgs = { - variant_count: _compose_cfg( - task, - lower=scale_lower, - upper=scale_upper, - variant_count=variant_count, - ) - for variant_count in variant_counts - } - - for num_envs in env_nums: - print(f"\nnum_envs={num_envs}", flush=True) - - for warmup_idx in range(warmup): - elapsed, _ = _construct_once( - baseline_cfg, - num_envs=num_envs, - init_dr_enabled=False, - force_pool=force_pool, - ) - print( - f" [init_dr_off] warmup {warmup_idx + 1}/{warmup}: {elapsed:.3f}s", - flush=True, - ) - - off_samples: list[float] = [] - off_meta: dict[str, Any] = {} - for repeat_idx in range(repeats): - elapsed, off_meta = _construct_once( - baseline_cfg, - num_envs=num_envs, - init_dr_enabled=False, - force_pool=force_pool, - ) - off_samples.append(elapsed) - print( - f" [init_dr_off] repeat {repeat_idx + 1}/{repeats}: {elapsed:.3f}s", - flush=True, - ) - records.append( - _summarize_record( - measure=measure, - mode="init_dr_off", - variant_count=1, - num_envs=num_envs, - lower=scale_lower, - upper=scale_lower, - samples=off_samples, - meta=off_meta, - ) - ) - - for variant_count in variant_counts: - cfg = variant_cfgs[variant_count] - for warmup_idx in range(warmup): - elapsed, _ = _construct_once( - cfg, - num_envs=num_envs, - init_dr_enabled=True, - force_pool=force_pool, - ) - print( - f" [init_dr_on:v{variant_count}] warmup {warmup_idx + 1}/{warmup}: {elapsed:.3f}s", - flush=True, - ) - - on_samples: list[float] = [] - on_meta: dict[str, Any] = {} - for repeat_idx in range(repeats): - elapsed, on_meta = _construct_once( - cfg, - num_envs=num_envs, - init_dr_enabled=True, - force_pool=force_pool, - ) - on_samples.append(elapsed) - print( - f" [init_dr_on:v{variant_count}] repeat {repeat_idx + 1}/{repeats}: {elapsed:.3f}s", - flush=True, - ) - records.append( - _summarize_record( - measure=measure, - mode="init_dr_on", - variant_count=variant_count, - num_envs=num_envs, - lower=scale_lower, - upper=scale_upper, - samples=on_samples, - meta=on_meta, - ) - ) - - return records - - -def _record_map(records: list[ConstructRecord]) -> dict[tuple[str, int, int], ConstructRecord]: - return {(record.mode, record.variant_count, record.num_envs): record for record in records} - - -def print_table(records: list[ConstructRecord]) -> None: - env_nums = sorted({record.num_envs for record in records}) - variant_counts = sorted( - {record.variant_count for record in records if record.mode == "init_dr_on"} - ) - by_key = _record_map(records) - - print() - header = f"{'num_envs':>8} | {'off_mean(s)':>11}" - for variant_count in variant_counts: - header += f" | {f'on_v{variant_count}(s)':>11}" - print(header) - print("-" * len(header)) - for num_envs in env_nums: - off = by_key[("init_dr_off", 1, num_envs)] - row = f"{num_envs:8d} | {off.mean_sec:11.3f}" - for variant_count in variant_counts: - on = by_key[("init_dr_on", variant_count, num_envs)] - row += f" | {on.mean_sec:11.3f}" - print(row) - - print() - print( - f"{'num_envs':>8} | {'variant':>7} | {'delta(s)':>9} | {'ratio':>7} | {'variants(on)':>12}" - ) - print("-" * 64) - for num_envs in env_nums: - off = by_key[("init_dr_off", 1, num_envs)] - for variant_count in variant_counts: - on = by_key[("init_dr_on", variant_count, num_envs)] - delta = on.mean_sec - off.mean_sec - ratio = on.mean_sec / max(off.mean_sec, 1e-12) - print( - f"{num_envs:8d} | {variant_count:7d} | {delta:9.3f} | {ratio:7.2f} | {on.model_variant_count:12d}" - ) - - -def save_json(path: Path, records: list[ConstructRecord], meta: dict[str, Any]) -> None: - path.parent.mkdir(parents=True, exist_ok=True) - payload = { - "meta": { - "timestamp_utc": datetime.now(timezone.utc).isoformat(), - "device_info": get_device_info_dict(), - **meta, - }, - "results": [asdict(record) for record in records], - } - path.write_text(json.dumps(payload, indent=2), encoding="utf-8") - print(f"Saved: {path.resolve()}") - - -def save_plot(path: Path, records: list[ConstructRecord]) -> bool: - if plt is None or not records: - return False - - path.parent.mkdir(parents=True, exist_ok=True) - env_nums = sorted({record.num_envs for record in records}) - variant_counts = sorted( - {record.variant_count for record in records if record.mode == "init_dr_on"} - ) - - fig, (ax_env, ax_variant) = plt.subplots(1, 2, figsize=(15, 5)) - cmap = plt.get_cmap("tab10") - - off_subset = sorted( - [record for record in records if record.mode == "init_dr_off"], - key=lambda record: record.num_envs, - ) - ax_env.plot( - [record.num_envs for record in off_subset], - [record.mean_sec for record in off_subset], - marker="o", - color="#2563eb", - label="init DR off", - ) - - for color_idx, variant_count in enumerate(variant_counts): - subset = sorted( - [ - record - for record in records - if record.mode == "init_dr_on" and record.variant_count == variant_count - ], - key=lambda record: record.num_envs, - ) - ax_env.plot( - [record.num_envs for record in subset], - [record.mean_sec for record in subset], - marker="o", - color=cmap(color_idx % 10), - label=f"init DR on (variants={variant_count})", - ) - - for color_idx, num_envs in enumerate(env_nums): - subset = sorted( - [ - record - for record in records - if record.mode == "init_dr_on" and record.num_envs == num_envs - ], - key=lambda record: record.variant_count, - ) - ax_variant.plot( - [record.variant_count for record in subset], - [record.mean_sec for record in subset], - marker="s", - color=cmap(color_idx % 10), - label=f"num_envs={num_envs}", - ) - - ax_env.set_xscale("log", base=2) - ax_env.set_yscale("log") - ax_env.set_xticks(env_nums) - ax_env.set_xticklabels([str(num_envs) for num_envs in env_nums], rotation=30) - ax_env.set_xlabel("num_envs") - ax_env.set_ylabel("construction time mean (s)") - ax_env.set_title("Construction Time vs num_envs") - ax_env.grid(True, alpha=0.3) - ax_env.legend(fontsize=8) - - ax_variant.set_xscale("log", base=2) - ax_variant.set_yscale("log") - ax_variant.set_xticks(variant_counts) - ax_variant.set_xticklabels([str(variant_count) for variant_count in variant_counts]) - ax_variant.set_xlabel("variant_count") - ax_variant.set_ylabel("construction time mean (s)") - ax_variant.set_title("Construction Time vs variant_count") - ax_variant.grid(True, alpha=0.3) - ax_variant.legend(fontsize=8) - - measure = records[0].measure - title = f"Sharpa-hand init-DR construction benchmark ({measure})" - device_info = get_device_info_line() - fig.suptitle(f"{title}\n{device_info}" if device_info else title) - fig.tight_layout() - fig.savefig(path, dpi=180) - plt.close(fig) - print(f"Saved: {path.resolve()}") - return True - - -def main() -> None: - parser = argparse.ArgumentParser( - description="Benchmark Sharpa-hand construction time vs init-DR variant count" - ) - parser.add_argument("--task", type=str, default="sharpa_inhand/mujoco") - parser.add_argument("--measure", type=str, default="construct_only") - parser.add_argument("--env-nums", type=str, default=None, help="Comma-separated env counts") - parser.add_argument("--variant-counts", type=str, default=None) - parser.add_argument("--scale-lower", type=float, default=DEFAULT_SCALE_LOWER) - parser.add_argument("--scale-upper", type=float, default=DEFAULT_SCALE_UPPER) - parser.add_argument("--repeats", type=int, default=1) - parser.add_argument("--warmup", type=int, default=0) - parser.add_argument( - "--out-json", - type=Path, - default=DEFAULT_OUTPUT_DIR / "results.json", - ) - parser.add_argument( - "--out-png", - type=Path, - default=DEFAULT_OUTPUT_DIR / "construct_time.png", - ) - args = parser.parse_args() - - if args.measure not in {"construct_only", "construct_plus_pool"}: - raise ValueError( - f"--measure must be one of {{'construct_only','construct_plus_pool'}}, got {args.measure}" - ) - - env_nums = _parse_env_nums(args.env_nums) - variant_counts = _parse_variant_counts(args.variant_counts) - - records = run_benchmark( - task=args.task, - measure=args.measure, - env_nums=env_nums, - variant_counts=variant_counts, - scale_lower=float(args.scale_lower), - scale_upper=float(args.scale_upper), - repeats=args.repeats, - warmup=args.warmup, - ) - print_table(records) - save_json( - args.out_json, - records, - meta={ - "task": args.task, - "measure": args.measure, - "env_nums": env_nums, - "variant_counts": variant_counts, - "scale_lower": float(args.scale_lower), - "scale_upper": float(args.scale_upper), - "repeats": args.repeats, - "warmup": args.warmup, - "batch_env_contract": { - "source_py": "mujoco_uni/runtime/batch.py (re-exported as mujoco_uni.batch_env)", - "source_cc": "mujoco_uni/native/batch_env.cc", - "accepted_model_arity": "single MjModel, or sequence length 1 or nbatch", - }, - "init_dr_off_mode": "local monkeypatch: provider returns no InitRandomizationPlan", - }, - ) - save_plot(args.out_png, records) - - -if __name__ == "__main__": - main() diff --git a/scripts/benchmark/physics/benchmark_physics_step_isaacgym.py b/scripts/benchmark/physics/benchmark_physics_step_isaacgym.py index 23e45ab3a..57de6bb90 100644 --- a/scripts/benchmark/physics/benchmark_physics_step_isaacgym.py +++ b/scripts/benchmark/physics/benchmark_physics_step_isaacgym.py @@ -128,32 +128,20 @@ class BenchRecord: asset_file="g1_description/g1_29dof_rev_1_0.urdf", initial_height=0.78, ), - "sharpa_inhand": TaskSpec( - owner_task_id="sharpa_inhand", - display_name="sharpa_inhand", - asset_root=DEFAULT_MODELS_ROOT, - asset_file="right_sharpa_wave/right_sharpa_wave.urdf", - initial_height=0.30, - ), } TASK_ALIASES = { "Go1JoystickFlat": "go1_joystick_flat", "Go2JoystickFlat": "go2_joystick_flat", "G1WalkFlat": "g1_walk_flat", - "SharpaInhandRotation": "sharpa_inhand", "task=go1_joystick_flat/isaacgym": "go1_joystick_flat", "task=go2_joystick_flat/isaacgym": "go2_joystick_flat", "task=g1_walk_flat/isaacgym": "g1_walk_flat", - "task=sharpa_inhand/isaacgym": "sharpa_inhand", "go1_joystick_flat/isaacgym": "go1_joystick_flat", "go2_joystick_flat/isaacgym": "go2_joystick_flat", "g1_walk_flat/isaacgym": "g1_walk_flat", - "sharpa_inhand/isaacgym": "sharpa_inhand", "go1": "go1_joystick_flat", "go2": "go2_joystick_flat", "g1": "g1_walk_flat", - "sharpa": "sharpa_inhand", - "sharpahand": "sharpa_inhand", } DEFAULT_TASK_IDS = list(TASK_SPECS.keys()) DEFAULT_BATCH_SIZES = [2**k for k in range(8, 15)] # 256 .. 16384 @@ -449,7 +437,7 @@ def main(): "--models-root", type=str, default=str(DEFAULT_MODELS_ROOT), - help="Root containing go1_description/go2_description/g1_description/sharpa_wave URDFs", + help="Root containing go1_description/go2_description/g1_description URDFs", ) parser.add_argument("--tasks", type=str, default=",".join(DEFAULT_TASK_IDS)) parser.add_argument( diff --git a/scripts/benchmark/physics/benchmark_physics_step_isaacsim.py b/scripts/benchmark/physics/benchmark_physics_step_isaacsim.py index 14d12e75d..8c4852791 100644 --- a/scripts/benchmark/physics/benchmark_physics_step_isaacsim.py +++ b/scripts/benchmark/physics/benchmark_physics_step_isaacsim.py @@ -96,24 +96,14 @@ class BenchRecord: articulation_root_prim="pelvis", initial_height=0.78, ), - "sharpa_inhand": TaskSpec( - owner_task_id="sharpa_inhand", - display_name="sharpa_inhand", - usd_file="sharpa_wave/right_sharpa_wave.usda", - articulation_root_prim="root_joint", - initial_height=0.30, - ), } TASK_ALIASES = { "Go1JoystickFlat": "go1_joystick_flat", "Go2JoystickFlat": "go2_joystick_flat", "G1WalkFlat": "g1_walk_flat", - "SharpaInhandRotation": "sharpa_inhand", "go1": "go1_joystick_flat", "go2": "go2_joystick_flat", "g1": "g1_walk_flat", - "sharpa": "sharpa_inhand", - "sharpahand": "sharpa_inhand", **{f"task={task}/isaacsim": task for task in TASK_SPECS}, **{f"{task}/isaacsim": task for task in TASK_SPECS}, } @@ -578,7 +568,7 @@ def main() -> None: "--usd-root", type=str, default=str(DEFAULT_USD_ROOT), - help="Directory containing go1/go2/g1/sharpa USD assets", + help="Directory containing go1/go2/g1 USD assets", ) parser.add_argument("--tasks", type=str, default=",".join(DEFAULT_TASK_IDS)) parser.add_argument( diff --git a/scripts/benchmark/physics/benchmark_physics_step_mujoco_warp.py b/scripts/benchmark/physics/benchmark_physics_step_mujoco_warp.py index 5bb973b69..699829b37 100644 --- a/scripts/benchmark/physics/benchmark_physics_step_mujoco_warp.py +++ b/scripts/benchmark/physics/benchmark_physics_step_mujoco_warp.py @@ -84,7 +84,6 @@ class BenchRecord: "go1_joystick_flat": 100, "go2_joystick_flat": 100, "g1_walk_flat": 256, - "sharpa_inhand": 128, } diff --git a/scripts/benchmark/rl/benchmark_offpolicy_collector_active.py b/scripts/benchmark/rl/benchmark_offpolicy_collector_active.py index bd79d178b..0c3ba0bb1 100644 --- a/scripts/benchmark/rl/benchmark_offpolicy_collector_active.py +++ b/scripts/benchmark/rl/benchmark_offpolicy_collector_active.py @@ -1624,8 +1624,7 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: default="default", help=( "'default', 'auto', or comma-separated // cases. " - "Default covers SAC G1 motion tracking and FlashSAC G1 walk flat on the selected backend. " - "Run Sharpa explicitly with --cases sac/sharpa_inhand/mujoco_hora." + "Default covers SAC G1 motion tracking and FlashSAC G1 walk flat on the selected backend." ), ) parser.add_argument( diff --git a/scripts/benchmark/rl/benchmark_sac_replay_buffer_sampling.py b/scripts/benchmark/rl/benchmark_sac_replay_buffer_sampling.py index 51364774e..18630486f 100644 --- a/scripts/benchmark/rl/benchmark_sac_replay_buffer_sampling.py +++ b/scripts/benchmark/rl/benchmark_sac_replay_buffer_sampling.py @@ -68,7 +68,7 @@ DEFAULT_TASK = "g1_motion_tracking" DEFAULT_SIM = "motrix" LEGACY_DEFAULT_TASK = "g1_walk_flat" -KNOWN_SIMS = ("mujoco", "motrix", "mujoco_hora") +KNOWN_SIMS = ("mujoco", "motrix") DEFAULT_GPU_COUNTS = "1,2,4,8" DEFAULT_CAPACITY_MULTIPLIERS = "0.25,0.5,1.0" FLOAT_BYTES = 4 diff --git a/scripts/sharpa_collect_grasps.sh b/scripts/sharpa_collect_grasps.sh deleted file mode 100755 index 87c9ff6f4..000000000 --- a/scripts/sharpa_collect_grasps.sh +++ /dev/null @@ -1,42 +0,0 @@ -#!/usr/bin/env bash - -set -euo pipefail - -if [ "$#" -lt 1 ]; then - echo "Usage: $0 [scale2 ...]" - echo "Environment:" - echo " SHARPA_GRASP_BACKEND=mujoco|motrix default: mujoco" - echo " SHARPA_GRASP_CACHE_PATH= optional env.grasp_cache_path override" - echo " SHARPA_GRASP_TARGET= optional env.grasp_collection_target override" - echo " SHARPA_GRASP_NUM_ENVS= optional algo.num_envs override" - exit 1 -fi - -backend="${SHARPA_GRASP_BACKEND:-mujoco}" -case "${backend}" in - mujoco|motrix) ;; - *) - echo "Unsupported SHARPA_GRASP_BACKEND=${backend}; expected mujoco or motrix" - exit 1 - ;; -esac - -extra_args=(training.no_play=true) -if [ -n "${SHARPA_GRASP_CACHE_PATH:-}" ]; then - extra_args+=("env.grasp_cache_path=${SHARPA_GRASP_CACHE_PATH}") -fi -if [ -n "${SHARPA_GRASP_TARGET:-}" ]; then - extra_args+=("env.grasp_collection_target=${SHARPA_GRASP_TARGET}") -fi -if [ -n "${SHARPA_GRASP_NUM_ENVS:-}" ]; then - extra_args+=("algo.num_envs=${SHARPA_GRASP_NUM_ENVS}") -fi - -for scale in "$@"; do - echo "[sharpa_collect_grasps] collecting backend=${backend} scale=${scale}" - - uv run src/unilab/scripts/train_rsl_rl.py \ - "task=sharpa_inhand_grasp/${backend}" \ - "env.domain_rand.scale_list=[${scale}]" \ - "${extra_args[@]}" -done diff --git a/scripts/tools/support_matrix.py b/scripts/tools/support_matrix.py index 68a76a8bd..14d66b90a 100644 --- a/scripts/tools/support_matrix.py +++ b/scripts/tools/support_matrix.py @@ -89,8 +89,6 @@ "x2_wall_flip_tracking": 8, "allegro_inhand": 9, "allegro_sac": 10, - "sharpa_inhand": 11, - "sharpa_inhand_grasp": 12, } _TASK_LABELS = { "go1_joystick_flat": "Go1 joystick", @@ -104,8 +102,6 @@ "x2_wall_flip_tracking": "X2 wall flip tracking", "allegro_inhand": "Allegro in-hand", "allegro_sac": "Allegro SAC in-hand", - "sharpa_inhand": "Sharpa in-hand", - "sharpa_inhand_grasp": "Sharpa in-hand grasp", } diff --git a/scripts/train_hora_distill.py b/scripts/train_hora_distill.py deleted file mode 100644 index 147b60677..000000000 --- a/scripts/train_hora_distill.py +++ /dev/null @@ -1,258 +0,0 @@ -import datetime -import sys -from pathlib import Path -from typing import Any, cast - -import hydra -import torch -from omegaconf import DictConfig, OmegaConf - -ROOT_DIR = Path(__file__).parent.parent -SRC_DIR = ROOT_DIR / "src" -if str(SRC_DIR) not in sys.path: - sys.path.insert(0, str(SRC_DIR)) -if str(ROOT_DIR) not in sys.path: - sys.path.insert(0, str(ROOT_DIR)) - -from uni_rl.algos.hora import HoraDistillationTrainer -from uni_rl.algos.hora.distill import ( - build_student_actor_and_normalizer, - cfg_with_checkpoint_runtime, - load_distilled_checkpoint, - student_policy, -) -from uni_rl.algos.hora.rsl_rl import HoraRslRlVecEnvWrapper as RslRlVecEnvWrapper -from unisim.backend.base import log_playback_plan -from unisim.backend.mujoco.xml import materialize_scene_visual_override - -from unilab.base.config_adapter import ( - BackendAdapter, - create_env, -) -from unilab.training import ( - ensure_registries, - format_hora_stage2_checkpoint_error, - get_log_root, - resolve_hora_stage2_checkpoint_path, - setup_logger, - should_run_playback, -) -from unilab.training.experiment import get_device_info_dict, write_run_config_snapshot -from unilab.training.hora_distill_config import ( - apply_teacher_defaults as _apply_teacher_defaults, -) -from unilab.training.hora_distill_config import ( - get_teacher_owner_spec as _get_teacher_owner_spec, -) -from unilab.training.hora_distill_config import ( - resolve_teacher_checkpoint_path as _resolve_teacher_checkpoint_path, -) -from unilab.training.hora_distill_config import ( - resolved_distill_runtime_cfg as _resolved_distill_runtime_cfg, -) -from unilab.training.hora_distill_config import ( - teacher_run_metadata as _teacher_run_metadata, -) - - -def _write_distill_run_config( - log_dir: Path, - *, - cfg: DictConfig, - teacher_metadata: dict[str, Any], -) -> None: - """Persist distillation run config plus teacher provenance near the checkpoints. - - Args: - log_dir: Run directory where the metadata file should be written. - cfg: Resolved distillation config for this run. - teacher_metadata: Explicit teacher provenance dictionary for this run. - - Returns: - None. Writes `distill_run_config.json` into `log_dir`. - """ - write_run_config_snapshot( - log_dir, - run_metadata={ - "algo": "hora_distill", - "task": str(OmegaConf.select(cfg, "training.task_name")), - "sim_backend": str(OmegaConf.select(cfg, "training.sim_backend")), - "log_dir": str(log_dir), - "hardware": get_device_info_dict(), - "teacher": teacher_metadata, - }, - full_cfg=cfg, - filename="distill_run_config.json", - trailing_newline=True, - ) - - -def _build_env_cfg_override(cfg: DictConfig) -> dict[str, Any]: - adapter = BackendAdapter( - cfg, - root_dir=ROOT_DIR, - algo_name="hora_distill", - scene_materializer=materialize_scene_visual_override, - ) - return cast(dict[str, Any], adapter.build_task_env_cfg_override()) - - -def _build_play_env_cfg_override(cfg: DictConfig) -> dict[str, Any]: - adapter = BackendAdapter( - cfg, - root_dir=ROOT_DIR, - algo_name="hora_distill", - scene_materializer=materialize_scene_visual_override, - ) - return cast(dict[str, Any], adapter.build_play_env_cfg_override()) - - -def _play_camera_kwargs(cfg: DictConfig) -> dict[str, Any]: - camera_kwargs = { - "cam_tracking": getattr(cfg.training, "cam_tracking", False), - "cam_tracking_env_idx": getattr(cfg.training, "cam_tracking_env_idx", 0), - "cam_tracking_extra_envs": getattr(cfg.training, "cam_tracking_extra_envs", 2), - } - for key in ("cam_distance", "cam_elevation", "cam_azimuth", "cam_lookat"): - value = getattr(cfg.training, key, None) - if value is not None: - camera_kwargs[key] = value - return camera_kwargs - - -def play_hora_distill(cfg: DictConfig, device: str) -> str | None: - task_log_root = get_log_root(ROOT_DIR, cfg) / str(cfg.training.task_name) - load_path, load_path_dir = resolve_hora_stage2_checkpoint_path(cfg, root_dir=ROOT_DIR) - if load_path is None or load_path_dir is None or not load_path.exists(): - print( - format_hora_stage2_checkpoint_error( - cfg, - task_log_root=task_log_root, - load_path=load_path, - load_path_dir=load_path_dir, - ) - ) - return None - - print(f"Loading distilled model: {load_path}") - checkpoint = torch.load(load_path, map_location="cpu", weights_only=False) - if "model_state_dict" not in checkpoint: - print( - f"Checkpoint at {load_path} is not a HORA distillation checkpoint " - f"(found keys: {set(checkpoint.keys())}). Aborting play." - ) - return None - - # uni_rl's cfg_with_checkpoint_runtime no longer composes teacher defaults; - # the caller owns that composition (issue #1480). - cfg = cfg_with_checkpoint_runtime(_apply_teacher_defaults(cfg), checkpoint) - env = create_env( - cfg, - num_envs=int(cfg.training.play_env_num), - env_cfg_override=_build_play_env_cfg_override(cfg), - ) - wrapped_env = RslRlVecEnvWrapper(env, device=device, policy_obs_mode="actor") - torch_device = torch.device(device) - actor, hist_normalizer = build_student_actor_and_normalizer( - wrapped_env, - cfg, - device=torch_device, - ) - load_distilled_checkpoint(actor, hist_normalizer, load_path, device=torch_device) - actor.eval() - hist_normalizer.eval() - - with torch.inference_mode(): - play_video_path = env.run_playback_mode( - play_render_mode=getattr(cfg.training, "play_render_mode", "auto"), - play_steps=getattr(cfg.training, "play_steps", None), - output_video=Path(load_path_dir) / "play_video_stage2.mp4", - render_spacing=float( - getattr(cfg.training, "render_spacing", getattr(env.cfg, "render_spacing", 1.0)) - ), - initialize=lambda: wrapped_env.reset()[0], - step=lambda obs: wrapped_env.step( - student_policy(actor, hist_normalizer, obs, device=torch_device) - )[0], - camera_kwargs=_play_camera_kwargs(cfg), - on_plan=log_playback_plan, - ) - print("Done.") - return play_video_path - - -@hydra.main(version_base="1.3", config_path="../src/unilab/conf/hora_distill", config_name="config") -def main(cfg: DictConfig) -> None: - ensure_registries() - - if cfg.training.device: - device = str(cfg.training.device) - elif torch.cuda.is_available(): - device = "cuda" - elif torch.backends.mps.is_available(): - device = "mps" - else: - device = "cpu" - - if should_run_playback( - play_only=cfg.training.play_only, - no_play=True, - play_render_mode=getattr(cfg.training, "play_render_mode", "auto"), - ): - play_hora_distill(cfg, device) - return - - cfg = _apply_teacher_defaults(cfg) - teacher_algo_family, teacher_task = _get_teacher_owner_spec(cfg) - if teacher_algo_family is None or teacher_task is None: - raise ValueError("HORA distillation requires teacher.algo_family and teacher.task.") - - teacher_checkpoint, _ = _resolve_teacher_checkpoint_path(cfg) - if teacher_checkpoint is None: - raise FileNotFoundError( - "Could not resolve HORA teacher checkpoint. " - f"teacher.algo_family={teacher_algo_family!r} " - f"teacher.task={teacher_task!r}. " - "Set algo.load_run and optionally algo.checkpoint." - ) - - teacher_metadata = _teacher_run_metadata( - cfg, - teacher_algo_family=teacher_algo_family, - teacher_checkpoint=teacher_checkpoint, - ) - timestamp = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S") - log_root = Path(cfg.training.log_dir) if cfg.training.log_dir else get_log_root(ROOT_DIR, cfg) - run_name = f"{timestamp}_{cfg.training.sim_backend}_{teacher_metadata['run_slug']}" - log_dir = log_root / str(cfg.training.task_name) / run_name - logger = setup_logger(log_dir, "hora_distill", echo=str(cfg.training.logger) != "no_print") - _write_distill_run_config(log_dir, cfg=cfg, teacher_metadata=teacher_metadata) - logger.info( - "teacher_algo=%s teacher_task=%s teacher_checkpoint=%s", - teacher_metadata["algo_family"], - teacher_metadata["task"], - teacher_metadata["checkpoint_path"], - ) - - env = create_env( - cfg, - num_envs=int(cfg.algo.num_envs), - env_cfg_override=_build_env_cfg_override(cfg), - ) - wrapped_env = RslRlVecEnvWrapper(env, device=device, policy_obs_mode="actor") - trainer = HoraDistillationTrainer( - wrapped_env, - cfg, - device=device, - log_dir=log_dir, - teacher_checkpoint=teacher_checkpoint, - teacher_algo_family=teacher_algo_family, - teacher_metadata=teacher_metadata, - distill_runtime_cfg=_resolved_distill_runtime_cfg(cfg), - logger=logger, - ) - trainer.train() - - -if __name__ == "__main__": - main() diff --git a/src/unilab/assets/hub.py b/src/unilab/assets/hub.py index 5cd89264b..89651f0fc 100644 --- a/src/unilab/assets/hub.py +++ b/src/unilab/assets/hub.py @@ -24,7 +24,6 @@ logger = logging.getLogger(__name__) _HF_MOTIONS_REPO_ID = "unilabsim/unilab-motions" -_HF_CACHES_REPO_ID = "unilabsim/unilab-caches" _HF_SCENES_REPO_ID = "unilabsim/unilab-scenes" _HF_CHECKPOINTS_REPO_ID = "unilabsim/unilab-checkpoints" _HF_ROBOTS_REPO_ID = "unilabsim/unilab-robots" @@ -48,7 +47,6 @@ "go2": (("robots/go2/assets", "base_0.obj", "**/*", "asset"),), # go2w points its meshdir at ``../go2/assets``. "go2w": (("robots/go2/assets", "base_0.obj", "**/*", "asset"),), - "sharpa_wave": (("robots/sharpa_wave/meshes", "DP_HB1_4F.STL", "*.STL", "STL"),), "x2": (("robots/x2/meshes", "pelvis.STL", "*.STL", "STL"),), } @@ -72,31 +70,6 @@ def resolve_motion_files( return [_resolve_single(p, repo_id=_HF_MOTIONS_REPO_ID) for p in motion_file] -def resolve_grasp_cache_files( - cache_file: str | Sequence[str], - *, - show_progress: bool = False, -) -> str | list[str]: - """Ensure grasp cache file(s) exist locally, downloading from HF if needed. - - Args: - cache_file: Absolute path or ``ASSETS_ROOT_PATH``-relative path - (single string or sequence of strings). - show_progress: Whether Hugging Face downloads may render progress bars. - - Returns: - Resolved absolute path(s) guaranteed to exist on disk. - A single string input returns a single string; a sequence input - returns a list of strings. - """ - if isinstance(cache_file, str): - return _resolve_single(cache_file, repo_id=_HF_CACHES_REPO_ID, show_progress=show_progress) - return [ - _resolve_single(p, repo_id=_HF_CACHES_REPO_ID, show_progress=show_progress) - for p in cache_file - ] - - def resolve_checkpoint_file( checkpoint_file: str | Sequence[str], ) -> str | list[str]: diff --git a/src/unilab/assets/robots/sharpa_wave/right_sharpa_wave.xml b/src/unilab/assets/robots/sharpa_wave/right_sharpa_wave.xml deleted file mode 100644 index 96775ecfe..000000000 --- a/src/unilab/assets/robots/sharpa_wave/right_sharpa_wave.xml +++ /dev/null @@ -1,263 +0,0 @@ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - diff --git a/src/unilab/assets/robots/sharpa_wave/scene.xml b/src/unilab/assets/robots/sharpa_wave/scene.xml deleted file mode 100644 index 3b95c2cb0..000000000 --- a/src/unilab/assets/robots/sharpa_wave/scene.xml +++ /dev/null @@ -1,59 +0,0 @@ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - diff --git a/src/unilab/cli.py b/src/unilab/cli.py index 2b52ad2c7..a1aa82b86 100644 --- a/src/unilab/cli.py +++ b/src/unilab/cli.py @@ -95,7 +95,7 @@ def _check_profile(profile: str | None) -> None: return if TASK_NAME_PATTERN.fullmatch(profile) is None: raise SystemExit( - "--profile must be a task owner variant such as `hora`; " + "--profile must be a task owner variant such as `nodr`; " "do not include slashes, dots, or path separators." ) @@ -250,7 +250,7 @@ def available_algos(root: Path | None = None) -> tuple[str, ...]: A custom algo ``X`` is routable when both ``conf/X/config.yaml`` and ``scripts/train_X.py`` exist under the package root. Config trees without - an entrypoint script (e.g. ``hora_distill``) are not routable. + an entrypoint script are not routable. """ selected_root = root or package_root() discovered: list[str] = [] diff --git a/src/unilab/conf/appo/task/sharpa_inhand/motrix.yaml b/src/unilab/conf/appo/task/sharpa_inhand/motrix.yaml deleted file mode 100644 index 51fa7985f..000000000 --- a/src/unilab/conf/appo/task/sharpa_inhand/motrix.yaml +++ /dev/null @@ -1,18 +0,0 @@ -# @package _global_ -defaults: - - /task/sharpa_inhand/mujoco - - _self_ - -training: - task_name: SharpaInhandRotation - sim_backend: motrix - -algo: - num_envs: 2048 - -env: - sim_dt: 0.01 - domain_rand: - randomize_gravity: true - randomize_gravity_direction: false - randomize_pd_gains: true diff --git a/src/unilab/conf/appo/task/sharpa_inhand/mujoco.yaml b/src/unilab/conf/appo/task/sharpa_inhand/mujoco.yaml deleted file mode 100644 index 53ce00191..000000000 --- a/src/unilab/conf/appo/task/sharpa_inhand/mujoco.yaml +++ /dev/null @@ -1,119 +0,0 @@ -# @package _global_ -# Base Sharpa APPO MuJoCo owner config. The HORA variant inherits this file -# so backend support and shared hyperparameters stay visible at the backend -# owner layer rather than being duplicated in a HORA-only variant. - -training: - task_name: SharpaInhandRotation - sim_backend: mujoco - play_steps: 200 - render_spacing: 0.5 - cam_distance: 1.5 - cam_lookat: [0.75, 0.75, 0.4] - cam_elevation: -20.0 - -algo: - num_envs: 16384 - steps_per_env: 8 - max_iterations: 301 - save_interval: 50 - actor: - hidden_dims: [512, 256, 128] - activation: elu - obs_normalization: true - distribution_cfg: - class_name: rsl_rl.modules.distribution.GaussianDistribution - init_std: 1.0 - std_type: scalar - critic: - hidden_dims: [512, 256, 128] - activation: elu - obs_normalization: true - algorithm: - value_loss_coef: 4.0 - entropy_coef: 0.01 - learning_rate: 0.001 - desired_kl: 0.02 - adaptive_kl_factor: 1.2 - adaptive_lr_factor: 1.1 - num_learning_epochs: 5 - num_mini_batches: 4 - clip_param: 0.2 - gamma: 0.99 - lam: 0.95 - -reward: - scales: - rotate: 2.5 - obj_linvel: -0.3 - pose_diff: -0.4 - torque: -0.1 - work: -0.5 - object_pos: 0.003 - angvel_clip_min: -0.5 - angvel_clip_max: 0.5 - -env: - post_step_forward_sensor: true - zero_action_test_mode: false - clip_obs: 5.0 - clip_actions: 1.0 - reset_height_lower: 0.59906 - reset_height_upper: 0.63906 - reset_angle_diff: 0.7853981633974483 - rot_axis: [0.0, 0.0, 1.0] - grasp_cache_path: caches/sharpa_grasp_linspace - sensor: - tactile_force_sensor_names: - - contact_right_thumb_elastomer_force - - contact_right_index_elastomer_force - - contact_right_middle_elastomer_force - - contact_right_ring_elastomer_force - - contact_right_pinky_elastomer_force - disable_tactile_ids: [] - use_default_object_pose_for_object_pos_anchor: true - obs: - observation_mode: flattened - enable_tactile: true - binary_contact: false - enable_contact_pos: false - contact_smooth: 0.5 - contact_threshold: 0.05 - tactile_force_clip_max: 4.0 - priv_info: - include_friction_scale: true - include_gravity_direction: false - control_config: - action_scale: 0.041666666666666664 - p_gain: 1.0 # use the XML value instead of this value by default - d_gain: 0.1 - torque_control: false # can only be false - dof_limits_scale: 0.9 - domain_rand: - scale_list: [0.8, 0.9, 1.0, 1.1, 1.2, 1.3, 1.4, 1.5] - randomize_gravity_direction: true - gravity_direction_magnitude: 9.81 - randomize_pd_gains: true - randomize_p_gain_scale_lower: 0.5 - randomize_p_gain_scale_upper: 2.0 - randomize_d_gain_scale_lower: 0.5 - randomize_d_gain_scale_upper: 2.0 - randomize_friction: true - randomize_friction_scale_lower: 0.75 - randomize_friction_scale_upper: 1.25 - elastomer_base_friction: 2.0 - metal_base_friction: 1.0 - object_base_friction: 2.0 - randomize_com: true - randomize_com_lower: -0.01 - randomize_com_upper: 0.01 - randomize_mass: true - randomize_mass_lower: 0.01 - randomize_mass_upper: 0.25 - force_scale: 2.0 - random_force_prob_scalar: 0.25 - force_decay: 0.9 - force_decay_interval: 0.08 - joint_noise_scale: 0.02 - contact_latency: 0.005 - contact_sensor_noise: 0.01 diff --git a/src/unilab/conf/appo/task/sharpa_inhand/mujoco_hora.yaml b/src/unilab/conf/appo/task/sharpa_inhand/mujoco_hora.yaml deleted file mode 100644 index 7a76aaec5..000000000 --- a/src/unilab/conf/appo/task/sharpa_inhand/mujoco_hora.yaml +++ /dev/null @@ -1,64 +0,0 @@ -# @package _global_ -# HORA Sharpa APPO variant. Inherit the shared MuJoCo owner so backend support -# and shared hyperparameters stay aligned with the non-HORA Sharpa APPO config. -defaults: - - /task/sharpa_inhand/mujoco - - _self_ - -training: - replay_queue_size: 8 - cam_distance: 0.75 - cam_lookat: [0.0, 0.0, 0.62] - cam_elevation: -25.0 - cam_azimuth: 45.0 - -interactive: - action_mode: policy - policy_obs_mode: actor - camera_distance: 0.75 - camera_elevation: -25.0 - camera_azimuth: 45.0 - use_env_visual_model: true - -algo: - algo_log_name: hora_appo - runtime_impl: hora_appo - runtime_resolver: unilab.scripts.play_hora_appo:resolve_hora_appo_runtime - num_envs: 2048 - steps_per_env: 8 - max_iterations: 305 - save_interval: 51 - obs_groups: - # Keep grouped keys explicit in the owner YAML; runtime support for these - # grouped observations lands in the next implementation step. - actor: - actor: 0 - priv_info: 0 - critic: - actor: 0 - priv_info: 0 - actor: - class_name: uni_rl.algos.hora:HoraActorModel - priv_info_embed_dim: 9 - priv_mlp_hidden_dims: [256, 128, 9] - critic: - class_name: uni_rl.algos.hora:HoraCriticModel - priv_info_embed_dim: 9 - priv_mlp_hidden_dims: [256, 128, 9] - algorithm: - learning_rate: 0.001 - desired_kl: 0.04 - adaptive_kl_factor: 1.2 - adaptive_lr_factor: 1.1 - -env: - use_default_object_pose_for_object_pos_anchor: true - obs: - observation_mode: separated - domain_rand: - randomize_friction: true - randomize_friction_scale_lower: 0.75 - randomize_friction_scale_upper: 1.25 - elastomer_base_friction: 2.0 - metal_base_friction: 1.0 - object_base_friction: 2.0 diff --git a/src/unilab/conf/hora_distill/config.yaml b/src/unilab/conf/hora_distill/config.yaml deleted file mode 100644 index 384b20c5b..000000000 --- a/src/unilab/conf/hora_distill/config.yaml +++ /dev/null @@ -1,77 +0,0 @@ -defaults: - - _self_ - - task: sharpa_inhand/mujoco - -algo: - algo_log_name: hora_distill - seed: 1 - num_envs: 4096 - max_agent_steps: 1000000000 - save_interval_steps: 100000000 - log_interval_steps: 32768 - learning_rate: 3.0e-4 - load_run: "-1" - checkpoint: -1 - model: {} - -training: - task_name: SharpaInhandRotation - device: null - logger: tensorboard - sim_backend: mujoco - play_only: false - play_render_mode: auto - play_env_num: 16 - play_steps: 200 - render_spacing: 1.0 - cam_distance: 6.0 - cam_elevation: -20.0 - cam_azimuth: 90.0 - cam_lookat: null - cam_tracking: false - cam_tracking_env_idx: 0 - cam_tracking_extra_envs: 2 - log_root: null - log_dir: null - -interactive: - action_mode: zero - policy_obs_mode: auto - show_target_bodies: false - show_reward_debug: false - target_show_axes: false - target_body_names: "" - target_max_bodies: 0 - target_marker_radius: 0.02 - target_axis_length: 0.08 - target_marker_alpha: 0.75 - reward_debug_show_velocity: false - reward_debug_lin_vel_scale: 0.08 - reward_debug_ang_vel_scale: 0.05 - reward_debug_show_connectors: false - reward_debug_show_global_anchor: false - camera_follow_body: true - camera_focus_body_name: "" - camera_height_offset: 0.15 - camera_distance: null - camera_elevation: null - camera_azimuth: null - use_env_visual_model: true - speed: 1.0 - start_paused: false - keyboard: false - keyboard_step_lin: 0.1 - keyboard_step_ang: 0.2 - -env: - post_step_forward_sensor: false - -hydra: - run: - dir: . - output_subdir: null - job: - chdir: false - job_logging: - root: - handlers: [console] diff --git a/src/unilab/conf/hora_distill/student_model/hora_actor.yaml b/src/unilab/conf/hora_distill/student_model/hora_actor.yaml deleted file mode 100644 index 5931860cc..000000000 --- a/src/unilab/conf/hora_distill/student_model/hora_actor.yaml +++ /dev/null @@ -1,17 +0,0 @@ -# Teacher -> student `algo.model` mapping for HoraActorModel teachers (PPO/APPO). -# -# Loaded by unilab.training.hora_distill_config and merged next to the -# Hydra-composed teacher owner config mounted at `teacher_owner`, so every -# student field interpolates directly from the teacher owner YAML. The teacher -# owner config stays the single source of truth for these hyperparameters; -# missing teacher fields fail at resolution time instead of falling back to -# Python-side defaults. -model: - hidden_dims: ${teacher_owner.algo.actor.hidden_dims} - activation: ${teacher_owner.algo.actor.activation} - obs_normalization: ${teacher_owner.algo.actor.obs_normalization} - priv_info_embed_dim: ${teacher_owner.algo.actor.priv_info_embed_dim} - priv_mlp_hidden_dims: ${teacher_owner.algo.actor.priv_mlp_hidden_dims} - # The student re-binds its own distribution class; distill_config strips - # `class_name` from the resolved dict before exposing `algo.model`. - distribution_cfg: ${teacher_owner.algo.actor.distribution_cfg} diff --git a/src/unilab/conf/hora_distill/student_model/hora_sac.yaml b/src/unilab/conf/hora_distill/student_model/hora_sac.yaml deleted file mode 100644 index 06aeaa93f..000000000 --- a/src/unilab/conf/hora_distill/student_model/hora_sac.yaml +++ /dev/null @@ -1,12 +0,0 @@ -# Teacher -> student `algo.model` mapping for hora_sac teachers (off-policy SAC). -# -# Loaded by unilab.training.hora_distill_config and merged next to the -# Hydra-composed teacher owner config mounted at `teacher_owner`. Values come -# from the teacher owner YAML; the `oc.select` fallback after the comma only -# applies when the teacher owner config does not define the field at all. -model: - teacher_arch: hora_sac - actor_hidden_dim: ${oc.select:teacher_owner.algo.actor_hidden_dim,512} - use_layer_norm: ${oc.select:teacher_owner.algo.use_layer_norm,true} - priv_info_embed_dim: ${oc.select:teacher_owner.algo.actor.priv_info_embed_dim,9} - priv_mlp_hidden_dims: ${oc.select:teacher_owner.algo.actor.priv_mlp_hidden_dims,[256, 128, 9]} diff --git a/src/unilab/conf/hora_distill/task/sharpa_inhand/mujoco.yaml b/src/unilab/conf/hora_distill/task/sharpa_inhand/mujoco.yaml deleted file mode 100644 index ca633879d..000000000 --- a/src/unilab/conf/hora_distill/task/sharpa_inhand/mujoco.yaml +++ /dev/null @@ -1,30 +0,0 @@ -# @package _global_ -teacher: - algo_family: ppo - task: sharpa_inhand/mujoco_hora - -training: - task_name: SharpaInhandRotation - sim_backend: mujoco - render_spacing: 0.5 - cam_distance: 1.5 - cam_lookat: [0.75, 0.75, 0.4] - cam_elevation: -20.0 - cam_azimuth: 90.0 - -interactive: - action_mode: policy - policy_obs_mode: actor - camera_distance: 1.5 - camera_elevation: -20.0 - camera_azimuth: 90.0 - use_env_visual_model: true - -env: - post_step_forward_sensor: true - -algo: - algo_log_name: hora_distill - num_envs: 16384 - max_agent_steps: 100000000 - save_interval_steps: 10000000 diff --git a/src/unilab/conf/hora_distill/task/sharpa_inhand/mujoco_nodr.yaml b/src/unilab/conf/hora_distill/task/sharpa_inhand/mujoco_nodr.yaml deleted file mode 100644 index 0c9556cf9..000000000 --- a/src/unilab/conf/hora_distill/task/sharpa_inhand/mujoco_nodr.yaml +++ /dev/null @@ -1,72 +0,0 @@ -# @package _global_ -defaults: - - /task/sharpa_inhand/mujoco - - _self_ - -teacher: - algo_family: appo - task: sharpa_inhand/mujoco_hora - -training: - task_name: SharpaInhandRotation - sim_backend: mujoco - render_spacing: 0.5 - cam_distance: 1.5 - cam_lookat: [0.75, 0.75, 0.4] - cam_elevation: -20.0 - cam_azimuth: 90.0 - -interactive: - action_mode: policy - policy_obs_mode: actor - show_target_bodies: false - show_reward_debug: false - target_show_axes: false - target_body_names: "" - target_max_bodies: 0 - target_marker_radius: 0.02 - target_axis_length: 0.08 - target_marker_alpha: 0.75 - reward_debug_show_velocity: false - reward_debug_lin_vel_scale: 0.08 - reward_debug_ang_vel_scale: 0.05 - reward_debug_show_connectors: false - reward_debug_show_global_anchor: false - camera_follow_body: true - camera_focus_body_name: "" - camera_height_offset: 0.15 - camera_distance: 1.5 - camera_elevation: -20.0 - camera_azimuth: 90.0 - use_env_visual_model: true - speed: 1.0 - start_paused: false - keyboard: false - keyboard_step_lin: 0.1 - keyboard_step_ang: 0.2 - -env: - post_step_forward_sensor: true - domain_rand: - scale_list: [0.8, 0.9, 1.0, 1.1, 1.2, 1.3, 1.4, 1.5] - randomize_base_mass: false - added_mass_range: [0.0, 0.0] - random_com: false - com_offset_x: [0.0, 0.0] - randomize_gravity: false - randomize_gravity_direction: false - randomize_pd_gains: false - randomize_friction: false - randomize_com: false - randomize_mass: false - force_scale: 0.0 - random_force_prob_scalar: 0.0 - joint_noise_scale: 0.0 - contact_latency: 0.0 - contact_sensor_noise: 0.0 - -algo: - algo_log_name: hora_distill - num_envs: 16384 - max_agent_steps: 100000000 - save_interval_steps: 10000000 diff --git a/src/unilab/conf/ppo/task/sharpa_inhand/motrix.yaml b/src/unilab/conf/ppo/task/sharpa_inhand/motrix.yaml deleted file mode 100644 index 91cbf1436..000000000 --- a/src/unilab/conf/ppo/task/sharpa_inhand/motrix.yaml +++ /dev/null @@ -1,26 +0,0 @@ -# @package _global_ -defaults: - - /task/sharpa_inhand/mujoco - - _self_ - -training: - task_name: SharpaInhandRotation - sim_backend: motrix - -algo: - num_envs: 2048 - -env: - sim_dt: 0.01 - # Motrix phase-1 supports object scale variants, object mass/COM, geom friction, - # direct object-force perturbations, and position-actuator kp/kd reset - # randomization. MotrixSim gravity override now supports reset-time gravity - # direction randomization. - domain_rand: - randomize_gravity: false - randomize_gravity_direction: true - randomize_pd_gains: true -play_profile: - enabled: true - env: - render_spacing: 2.0 diff --git a/src/unilab/conf/ppo/task/sharpa_inhand/mujoco.yaml b/src/unilab/conf/ppo/task/sharpa_inhand/mujoco.yaml deleted file mode 100644 index 638b8aaa4..000000000 --- a/src/unilab/conf/ppo/task/sharpa_inhand/mujoco.yaml +++ /dev/null @@ -1,105 +0,0 @@ -# @package _global_ -training: - task_name: SharpaInhandRotation - sim_backend: mujoco - render_spacing: 0.5 - cam_distance: 1.5 - cam_lookat: [0.75, 0.75, 0.4] - cam_elevation: -20.0 -algo: - num_envs: 2048 - num_steps_per_env: 8 - max_iterations: 301 - empirical_normalization: true - policy: - actor_hidden_dims: [512, 256, 128] - critic_hidden_dims: [512, 256, 128] - algorithm: - value_loss_coef: 4.0 - entropy_coef: 0.01 - learning_rate: 0.001 - desired_kl: 0.02 - num_learning_epochs: 5 - num_mini_batches: 4 - clip_param: 0.2 - gamma: 0.99 - lam: 0.95 - save_interval: 50 -reward: - scales: - rotate: 2.5 - obj_linvel: -0.3 - pose_diff: -0.4 - torque: -0.1 - work: -0.5 - object_pos: 0.003 - angvel_clip_min: -0.5 - angvel_clip_max: 0.5 -env: - post_step_forward_sensor: true - zero_action_test_mode: false - clip_obs: 5.0 - clip_actions: 1.0 - reset_height_lower: 0.59906 # assume the hand base height is 0.5 - reset_height_upper: 0.63906 - reset_angle_diff: 0.7853981633974483 - rot_axis: [0.0, 0.0, 1.0] - grasp_cache_path: caches/sharpa_grasp_linspace - sensor: - tactile_force_sensor_names: - - contact_right_thumb_elastomer_force - - contact_right_index_elastomer_force - - contact_right_middle_elastomer_force - - contact_right_ring_elastomer_force - - contact_right_pinky_elastomer_force - disable_tactile_ids: [] - use_default_object_pose_for_object_pos_anchor: false - obs: - observation_mode: flattened # non-hora setting - enable_tactile: true - binary_contact: false - enable_contact_pos: false - contact_smooth: 0.5 - contact_threshold: 0.05 - tactile_force_clip_max: 4.0 - priv_info: - include_friction_scale: true - include_gravity_direction: false - control_config: - action_scale: 0.041666666666666664 - p_gain: 1.0 # use the XML value instead of this value by default - d_gain: 0.1 - torque_control: false # can only be false - dof_limits_scale: 0.9 # tighten the XML dof limits - domain_rand: - scale_list: [0.8, 0.9, 1.0, 1.1, 1.2, 1.3, 1.4, 1.5] # cylinder object scales - randomize_gravity_direction: true - gravity_direction_magnitude: 9.81 - randomize_pd_gains: true - randomize_p_gain_scale_lower: 0.5 - randomize_p_gain_scale_upper: 2.0 - randomize_d_gain_scale_lower: 0.5 - randomize_d_gain_scale_upper: 2.0 - randomize_friction: true - randomize_friction_scale_lower: 0.75 - randomize_friction_scale_upper: 1.25 - elastomer_base_friction: 2.0 - metal_base_friction: 1.0 - object_base_friction: 2.0 - randomize_com: true - randomize_com_lower: -0.01 - randomize_com_upper: 0.01 - randomize_mass: true - randomize_mass_lower: 0.01 - randomize_mass_upper: 0.25 - force_scale: 2.0 - random_force_prob_scalar: 0.25 - force_decay: 0.9 - force_decay_interval: 0.08 - joint_noise_scale: 0.02 - contact_latency: 0.005 - contact_sensor_noise: 0.01 -play_profile: - enabled: true - env: - render_spacing: 2.0 diff --git a/src/unilab/conf/ppo/task/sharpa_inhand/mujoco_hora.yaml b/src/unilab/conf/ppo/task/sharpa_inhand/mujoco_hora.yaml deleted file mode 100644 index a5660212f..000000000 --- a/src/unilab/conf/ppo/task/sharpa_inhand/mujoco_hora.yaml +++ /dev/null @@ -1,44 +0,0 @@ -# @package _global_ -defaults: - - /task/sharpa_inhand/mujoco - - _self_ - -interactive: - action_mode: policy - policy_obs_mode: actor - camera_distance: 1.5 - camera_elevation: -20.0 - camera_azimuth: 90.0 - use_env_visual_model: true - -algo: - algo_log_name: hora_ppo - runtime_impl: hora_ppo - runtime_resolver: uni_rl.algos.hora.rsl_rl:resolve_hora_ppo_runtime - obs_groups: - actor: [actor] - critic: [actor] - actor: - class_name: uni_rl.algos.hora:HoraActorModel - hidden_dims: [512, 256, 128] - activation: elu - obs_normalization: true - priv_info_embed_dim: 9 - priv_mlp_hidden_dims: [256, 128, 9] - distribution_cfg: - class_name: GaussianDistribution - init_std: 1.0 - std_type: scalar - critic: - class_name: uni_rl.algos.hora:HoraCriticModel - hidden_dims: [512, 256, 128] - activation: elu - obs_normalization: true - priv_info_embed_dim: 9 - priv_mlp_hidden_dims: [256, 128, 9] - algorithm: - class_name: uni_rl.algos.hora:HoraPPO - -env: - obs: - observation_mode: separated diff --git a/src/unilab/conf/ppo/task/sharpa_inhand_grasp/motrix.yaml b/src/unilab/conf/ppo/task/sharpa_inhand_grasp/motrix.yaml deleted file mode 100644 index 86168005b..000000000 --- a/src/unilab/conf/ppo/task/sharpa_inhand_grasp/motrix.yaml +++ /dev/null @@ -1,18 +0,0 @@ -# @package _global_ -defaults: - - /task/sharpa_inhand_grasp/mujoco - - _self_ - -training: - task_name: SharpaInhandRotationGrasp - sim_backend: motrix - -env: - sim_dt: 0.01 - domain_rand: - randomize_gravity: false - randomize_gravity_direction: false -play_profile: - enabled: true - env: - render_spacing: 2.0 diff --git a/src/unilab/conf/ppo/task/sharpa_inhand_grasp/mujoco.yaml b/src/unilab/conf/ppo/task/sharpa_inhand_grasp/mujoco.yaml deleted file mode 100644 index 871b4be0a..000000000 --- a/src/unilab/conf/ppo/task/sharpa_inhand_grasp/mujoco.yaml +++ /dev/null @@ -1,106 +0,0 @@ -# @package _global_ -training: - task_name: SharpaInhandRotationGrasp - sim_backend: mujoco - render_spacing: 0.5 - cam_distance: 1.5 - cam_lookat: [0.75, 0.75, 0.4] - cam_elevation: -20.0 -algo: - num_envs: 2048 - num_steps_per_env: 8 - max_iterations: 2290 - empirical_normalization: true - policy: - actor_hidden_dims: [512, 256, 128] - critic_hidden_dims: [512, 256, 128] - algorithm: - value_loss_coef: 4.0 - entropy_coef: 0.0 - learning_rate: 5.0e-3 - desired_kl: 0.02 - num_learning_epochs: 5 - num_mini_batches: 4 - clip_param: 0.2 - gamma: 0.99 - lam: 0.95 -reward: - scales: - rotate: 0.0 - obj_linvel: 0.0 - pose_diff: 0.0 - torque: 0.0 - work: 0.0 - object_pos: 0.0 - angvel_clip_min: -0.5 - angvel_clip_max: 0.5 -env: - post_step_forward_sensor: true - max_episode_seconds: 3.0 - domain_rand: - scale_list: [0.8] - randomize_gravity: false - gravity_range: - - [0.0, 0.0, -9.81] - - [0.0, 0.0, -9.81] - randomize_gravity_direction: false - gravity_direction_magnitude: 9.81 - randomize_pd_gains: false - randomize_p_gain_scale_lower: 0.5 - randomize_p_gain_scale_upper: 2.0 - randomize_d_gain_scale_lower: 0.5 - randomize_d_gain_scale_upper: 2.0 - randomize_friction: false - randomize_friction_scale_lower: 0.5 - randomize_friction_scale_upper: 2.0 - elastomer_base_friction: 0.8 - metal_base_friction: 0.1 - object_base_friction: 0.5 - randomize_com: false - randomize_com_lower: -0.01 - randomize_com_upper: 0.01 - randomize_mass: true - randomize_mass_lower: 0.05 - randomize_mass_upper: 0.051 - force_scale: 0.0 - random_force_prob_scalar: 0.0 - force_decay: 0.9 - force_decay_interval: 0.08 - joint_noise_scale: 0.02 - contact_latency: 0.005 - contact_sensor_noise: 0.01 - zero_action_test_mode: false - clip_obs: 5.0 - clip_actions: 1.0 - control_config: - action_scale: 0.041666666666666664 - p_gain: 1.0 - d_gain: 0.1 - torque_control: false - dof_limits_scale: 0.9 - reset_height_lower: 0.61406 - reset_height_upper: 0.62406 - reset_angle_diff: 0.5235987755982988 - rot_axis: [0.0, 0.0, 1.0] - grasp_cache_path: caches/sharpa_grasp_linspace - obs: - observation_mode: flattened - enable_tactile: true - binary_contact: false - enable_contact_pos: false - contact_smooth: 0.5 - contact_threshold: 0.05 - tactile_force_clip_max: 5.0 - sensor: - tactile_force_sensor_names: - - contact_right_thumb_elastomer_force - - contact_right_index_elastomer_force - - contact_right_middle_elastomer_force - - contact_right_ring_elastomer_force - - contact_right_pinky_elastomer_force - disable_tactile_ids: [] - grasp_collection_target: 10000 -play_profile: - enabled: true - env: - render_spacing: 2.0 diff --git a/src/unilab/conf/sac/task/sharpa_inhand/mujoco_hora.yaml b/src/unilab/conf/sac/task/sharpa_inhand/mujoco_hora.yaml deleted file mode 100644 index a0025de0e..000000000 --- a/src/unilab/conf/sac/task/sharpa_inhand/mujoco_hora.yaml +++ /dev/null @@ -1,120 +0,0 @@ -# @package _global_ -# HORA Sharpa SAC MuJoCo teacher-training owner config. - -training: - task_name: SharpaInhandRotation - sim_backend: mujoco - use_amp: true - env_steps_per_sync: 2 - no_play: true - play_steps: 200 - render_spacing: 0.5 - cam_distance: 1.5 - cam_lookat: [0.75, 0.75, 0.4] - cam_elevation: -20.0 - -interactive: - action_mode: policy - policy_obs_mode: actor - camera_distance: 1.5 - camera_elevation: -20.0 - camera_azimuth: 90.0 - use_env_visual_model: true - -algo: - algo_log_name: hora_sac - runtime_impl: hora_sac - runtime_resolver: uni_rl.algos.hora.sac:resolve_hora_sac_runtime - num_envs: 1024 - batch_size: 2048 - replay_buffer_n: 1280 - updates_per_step: 14 - learning_starts: 1 - policy_frequency: 2 - max_iterations: 5371 - save_interval: 896 - actor_lr: 4.5e-4 - critic_lr: 4.5e-4 - actor: - priv_info_embed_dim: 9 - priv_mlp_hidden_dims: [256, 128, 9] - algo_params: - alpha_lr: 4.5e-4 - amp_dtype: auto - use_compile: true - -reward: - scales: - rotate: 2.5 - obj_linvel: -0.3 - pose_diff: -0.4 - torque: -0.1 - work: -0.5 - object_pos: 0.003 - angvel_clip_min: -0.5 - angvel_clip_max: 0.5 - -env: - post_step_forward_sensor: true - zero_action_test_mode: false - clip_obs: 5.0 - clip_actions: 1.0 - reset_height_lower: 0.59906 - reset_height_upper: 0.63906 - reset_angle_diff: 0.7853981633974483 - rot_axis: [0.0, 0.0, 1.0] - grasp_cache_path: caches/sharpa_grasp_linspace - sensor: - tactile_force_sensor_names: - - contact_right_thumb_elastomer_force - - contact_right_index_elastomer_force - - contact_right_middle_elastomer_force - - contact_right_ring_elastomer_force - - contact_right_pinky_elastomer_force - disable_tactile_ids: [] - use_default_object_pose_for_object_pos_anchor: true - obs: - observation_mode: separated - enable_tactile: true - binary_contact: false - enable_contact_pos: false - contact_smooth: 0.5 - contact_threshold: 0.05 - tactile_force_clip_max: 4.0 - priv_info: - include_friction_scale: true - include_gravity_direction: false - control_config: - action_scale: 0.041666666666666664 - p_gain: 1.0 - d_gain: 0.1 - torque_control: false - dof_limits_scale: 0.9 - domain_rand: - scale_list: [0.8, 0.9, 1.0, 1.1, 1.2, 1.3, 1.4, 1.5] - randomize_gravity_direction: true - gravity_direction_magnitude: 9.81 - randomize_pd_gains: true - randomize_p_gain_scale_lower: 0.5 - randomize_p_gain_scale_upper: 2.0 - randomize_d_gain_scale_lower: 0.5 - randomize_d_gain_scale_upper: 2.0 - randomize_friction: true - randomize_friction_scale_lower: 0.75 - randomize_friction_scale_upper: 1.25 - elastomer_base_friction: 2.0 - metal_base_friction: 1.0 - object_base_friction: 2.0 - randomize_com: true - randomize_com_lower: -0.01 - randomize_com_upper: 0.01 - randomize_mass: true - randomize_mass_lower: 0.01 - randomize_mass_upper: 0.25 - force_scale: 2.0 - random_force_prob_scalar: 0.25 - force_decay: 0.9 - force_decay_interval: 0.08 - joint_noise_scale: 0.02 - contact_latency: 0.005 - contact_sensor_noise: 0.01 diff --git a/src/unilab/demo.py b/src/unilab/demo.py index 608126398..202092adf 100644 --- a/src/unilab/demo.py +++ b/src/unilab/demo.py @@ -28,25 +28,11 @@ class DemoSpec: "dance": DemoSpec(algo="ppo", task="g1_motion_tracking", sim="motrix", entry="eval"), "wallflip": DemoSpec(algo="ppo", task="g1_wall_flip_tracking", sim="motrix", entry="eval"), "boxtracking": DemoSpec(algo="ppo", task="g1_box_tracking", sim="motrix", entry="eval"), - "inhandgrasp": DemoSpec( - algo="hora_distill", - task="sharpa_inhand", - sim="mujoco_nodr", - entry="play_interactive", - ), - "sharpa_appo_student": DemoSpec( - algo="hora_distill", - task="sharpa_inhand", - sim="mujoco_nodr", - entry="play_interactive", - ), "teaser": DemoSpec(algo="", task="", sim="", entry="teaser"), } -_LOCAL_ONLY_CHECKPOINT_DEMOS = {"sharpa_appo_student"} -_DEMO_PLAY_INTERACTIVE_OVERRIDES: dict[str, tuple[str, ...]] = { - "inhandgrasp": ("interactive.camera_follow_body=false",), -} +_LOCAL_ONLY_CHECKPOINT_DEMOS: set[str] = set() +_DEMO_PLAY_INTERACTIVE_OVERRIDES: dict[str, tuple[str, ...]] = {} def _package_root() -> Path: diff --git a/src/unilab/scripts/play_hora_appo.py b/src/unilab/scripts/play_hora_appo.py deleted file mode 100644 index 26bb5f69c..000000000 --- a/src/unilab/scripts/play_hora_appo.py +++ /dev/null @@ -1,290 +0,0 @@ -"""HORA APPO checkpoint playback, owned by UniLab (issue #1480). - -The HORA APPO *training* runner lives in uni_rl (``uni_rl.algos.hora.appo_runner``); -the play-mode orchestration is UniLab-side business logic because it drives -UniLab's registry-backed env construction, sim2sim contract validation, and -backend playback plan logging. ``resolve_hora_appo_runtime`` is the owner -config's ``runtime_resolver`` entrypoint consumed by -``uni_rl.algos.appo.runtime.resolve_appo_runtime``. -""" - -from __future__ import annotations - -import os -from collections.abc import Callable -from copy import deepcopy -from dataclasses import dataclass -from typing import Any, cast - -import torch -from omegaconf import DictConfig, OmegaConf -from uni_rl.algos.hora.appo_runner import HoraAPPORunner -from uni_rl.algos.hora.models import build_hora_shared_actor_critic -from uni_rl.algos.hora.observations import ( - build_hora_actor_tensordict, - split_hora_obs_with_priv_info, -) -from uni_rl.algos.hora.rsl_rl_compat import ( - convert_config_v3_to_v4, - is_rsl_rl_v4, - is_rsl_rl_v5, -) -from uni_rl.algos.hora.runtime import is_hora_appo_runtime -from uni_rl.utils.observations import get_obs_dims -from unisim.backend.base import log_playback_plan - -from unilab.base.config_adapter import BackendAdapter, create_env -from unilab.base.process_device import ( - apply_backend_env_device_override, - configure_backend_process_device, -) -from unilab.utils.sim2sim import policy_load_dim_guard, resolve_sim2sim_config - - -@dataclass(frozen=True) -class HoraAPPORuntime: - """Resolved HORA APPO entrypoints used by the generic APPO script. - - Args: - runner_cls: Runner class used for HORA APPO training mode. - play_fn: Play-mode callable used for HORA APPO checkpoint playback. - - Returns: - Immutable entrypoint bundle consumed by generic APPO script assembly. - """ - - runner_cls: type[HoraAPPORunner] - play_fn: Callable[..., str | None] - - -def resolve_hora_appo_runtime(rl_cfg: dict[str, Any]) -> HoraAPPORuntime | None: - """Resolve HORA APPO entrypoints from an explicit runtime marker. - - Args: - rl_cfg: Resolved algorithm config dictionary from Hydra composition. - - Returns: - ``HoraAPPORuntime`` when the owner config selects HORA APPO, otherwise - ``None``. - """ - if not is_hora_appo_runtime(rl_cfg): - return None - return HoraAPPORuntime(runner_cls=HoraAPPORunner, play_fn=play_hora_appo) - - -def _update_hora_obs_groups( - rl_cfg: dict[str, Any], - *, - obs_dim: int, - priv_info_dim: int, -) -> None: - """Update grouped actor/critic dims for the HORA APPO runtime. - - Args: - rl_cfg: Mutable algorithm config dictionary to update in place. - obs_dim: Actor observation dimension reported by the env contract. - priv_info_dim: Privileged-info dimension reported by the env contract. - - Returns: - None. Mutates ``rl_cfg["obs_groups"]`` directly. - """ - obs_groups = rl_cfg.setdefault("obs_groups", {}) - actor_group = obs_groups.setdefault("actor", {}) - critic_group = obs_groups.setdefault("critic", {}) - if isinstance(actor_group, dict): - actor_group["actor"] = obs_dim - actor_group["priv_info"] = priv_info_dim - if isinstance(critic_group, dict): - critic_group["actor"] = obs_dim - critic_group["priv_info"] = priv_info_dim - - -def play_hora_appo( - cfg: DictConfig, - rl_cfg: dict[str, Any], - *, - root_dir, - resolve_checkpoint_path, -) -> str | None: - """Play HORA APPO checkpoints with grouped actor and privileged inputs.""" - import numpy as np - from rsl_rl.utils import resolve_callable - from tensordict import TensorDict - - device = cfg.training.device or ( - "cuda" - if torch.cuda.is_available() - else "mps" - if torch.backends.mps.is_available() - else "cpu" - ) - # Genesis owns a process-wide session and must select its CUDA device - # before the first backend construction. Keep the rank/device routing in - # the shared owner-layer helper so HORA playback follows the same contract - # as the generic APPO/PPO/off-policy play paths. A non-zero Genesis - # request pins CUDA_VISIBLE_DEVICES; the bound in-process device replaces - # the requested one for the policy and the env override. - sim_backend = str(OmegaConf.select(cfg, "training.sim_backend", default="mujoco")) - if str(device).strip().lower().startswith("cuda"): - bound_device = configure_backend_process_device(sim_backend, device) - if bound_device is not None: - device = bound_device - env_cfg_override = apply_backend_env_device_override( - BackendAdapter( - cfg, - root_dir=root_dir, - algo_name="appo", - ).build_task_env_cfg_override(), - sim_backend, - learner_device=device, - ) - print(f"Using device for play: {device}") - - env = cast( - Any, - create_env( - cfg, - num_envs=cfg.training.play_env_num, - env_cfg_override=env_cfg_override, - ), - ) - obs_dim, _ = get_obs_dims(env.obs_groups_spec) - if env.state is None: - env.init_state() - _, _, state_priv_info = split_hora_obs_with_priv_info( - env.state.obs, - env.state.info if env.state is not None else None, - ) - priv_info_dim = int(state_priv_info.shape[1]) if state_priv_info is not None else 0 - if priv_info_dim <= 0: - raise ValueError("HORA APPO play requires privileged info from the environment.") - - action_shape = env.action_space.shape - if action_shape is None: - raise ValueError("env.action_space.shape must be defined") - action_dim = int(action_shape[0]) - - rl_cfg_dict = dict(rl_cfg) - _update_hora_obs_groups(rl_cfg_dict, obs_dim=obs_dim, priv_info_dim=priv_info_dim) - - if is_rsl_rl_v5(): - pass - elif is_rsl_rl_v4(): - rl_cfg_dict = convert_config_v3_to_v4(rl_cfg_dict) - - obs_example = torch.zeros((cfg.training.play_env_num, obs_dim), device=device) - td_example = TensorDict( - { - "actor": obs_example, - "priv_info": torch.zeros((cfg.training.play_env_num, priv_info_dim), device=device), - }, - batch_size=cfg.training.play_env_num, - ) - - actor_cfg = deepcopy(rl_cfg_dict["actor"]) - actor_cls = resolve_callable(actor_cfg.pop("class_name")) - actor_cfg.pop("num_actions", None) - - critic_cfg = deepcopy(rl_cfg_dict.get("critic") or rl_cfg_dict.get("actor") or {}) - critic_cfg.pop("class_name", None) - critic_cfg.pop("num_actions", None) - critic_cfg.pop("distribution_cfg", None) - - shared_model = build_hora_shared_actor_critic( - obs_dim=obs_dim, - action_dim=action_dim, - priv_info_dim=priv_info_dim, - actor_cfg=actor_cfg, - critic_cfg=critic_cfg, - ).to(device) - actor = actor_cls( - td_example, - rl_cfg_dict["obs_groups"], - "actor", - action_dim, - shared_model=shared_model, - **actor_cfg, - ) - actor = actor.to(device) - actor.eval() - - load_path, load_path_dir = resolve_checkpoint_path(cfg) - if not load_path or not os.path.exists(load_path): - print(f"Could not find run to load. load_path={load_path}") - return None - - resolve_sim2sim_config( - load_path_dir, - cfg, - algo_name="appo", - strict=bool(getattr(cfg.training, "sim2sim_strict", True)), - ) - - print(f"Loading model: {load_path}") - checkpoint = torch.load(load_path, map_location=device, weights_only=True) - with policy_load_dim_guard(env_obs_dim=obs_dim, env_action_dim=action_dim, algo_name="appo"): - actor.load_state_dict(checkpoint["actor"]) - - current_priv_info: np.ndarray | None = None - - def initialize_play_obs() -> np.ndarray: - nonlocal current_priv_info - obs_out, info_out = env.reset(np.arange(cfg.training.play_env_num, dtype=np.int32)) - actor_obs, _, priv_info = split_hora_obs_with_priv_info(obs_out, info_out) - current_priv_info = priv_info.astype(np.float32) if priv_info is not None else None - return np.asarray(actor_obs, dtype=np.float32) - - def step_play_obs(obs_np: np.ndarray) -> np.ndarray: - nonlocal current_priv_info - if current_priv_info is None: - raise ValueError("HORA APPO play step is missing privileged info.") - td = build_hora_actor_tensordict( - obs_np, - priv_info=current_priv_info, - device=device, - batch_size=cfg.training.play_env_num, - ) - actions = actor(td).cpu().numpy().astype(np.float32) - state = env.step(actions) - actor_obs, _, priv_info = split_hora_obs_with_priv_info(state.obs, state.info) - current_priv_info = priv_info.astype(np.float32) if priv_info is not None else None - return np.asarray(actor_obs, dtype=np.float32) - - print("Collecting physics states...") - with torch.inference_mode(): - play_video_path = cast( - str | None, - env.run_playback_mode( - play_render_mode=getattr(cfg.training, "play_render_mode", "auto"), - play_steps=getattr(cfg.training, "play_steps", None), - output_video=os.path.join(load_path_dir, "play_video.mp4") - if load_path_dir - else None, - render_spacing=float( - getattr(cfg.training, "render_spacing", getattr(env.cfg, "render_spacing", 1.0)) - ), - initialize=initialize_play_obs, - step=step_play_obs, - camera_kwargs={ - "cam_distance": cfg.training.cam_distance, - "cam_elevation": cfg.training.cam_elevation, - "cam_azimuth": cfg.training.cam_azimuth, - "cam_lookat": getattr(cfg.training, "cam_lookat", None), - "cam_tracking": getattr(cfg.training, "cam_tracking", False), - "cam_tracking_env_idx": getattr(cfg.training, "cam_tracking_env_idx", 0), - "cam_tracking_extra_envs": getattr(cfg.training, "cam_tracking_extra_envs", 2), - }, - on_plan=log_playback_plan, - ), - ) - if play_video_path is not None: - print(f"Saving video to {play_video_path} ...") - print("Done.") - return play_video_path - - -__all__ = [ - "HoraAPPORuntime", - "play_hora_appo", - "resolve_hora_appo_runtime", -] diff --git a/src/unilab/scripts/play_interactive.py b/src/unilab/scripts/play_interactive.py index 0f6bd7cd3..5c2cbf8b2 100644 --- a/src/unilab/scripts/play_interactive.py +++ b/src/unilab/scripts/play_interactive.py @@ -64,7 +64,6 @@ ) from unilab.utils.rotation import np_matrix_from_quat from unilab.visualization.interactive_playback import ( - _HORA_DISTILL_CHECKPOINT_UNAVAILABLE, KeyboardCommander, PlaybackControls, PlayInteractiveArgs, @@ -73,7 +72,6 @@ build_play_backend_adapter, build_playback_config, create_appo_playback_session, - create_hora_distill_playback_session, create_rsl_rl_playback_session, create_sac_playback_session, infer_checkpoint_actor_input_dim, @@ -132,14 +130,13 @@ def _algo_config_dict(cfg: DictConfig | None) -> dict[str, Any]: return algo_config_dict(cfg) -SUPPORTED_INTERACTIVE_ALGOS = ("ppo", "appo", "sac", "td3", "flashsac", "hora_distill") +SUPPORTED_INTERACTIVE_ALGOS = ("ppo", "appo", "sac", "td3", "flashsac") _CONFIG_ROOT_BY_ALGO = { "ppo": "ppo", "appo": "appo", "sac": "sac", "td3": "td3", "flashsac": "flashsac", - "hora_distill": "hora_distill", } _OFFPOLICY_INTERACTIVE_ALGOS = {"sac", "td3", "flashsac"} @@ -1024,22 +1021,10 @@ def _create_env(num_envs: int): algo_name=algo, log=lambda message: print(f"[play_interactive] {message}"), ) - elif algo == "hora_distill": - if cfg is None: - raise ValueError( - "HORA distill interactive playback requires a composed Hydra config." - ) - session = create_hora_distill_playback_session( - playback_cfg=playback_cfg, - cfg=cfg, - root_dir=Path.cwd(), - device=device, - log=lambda message: print(f"[play_interactive] {message}"), - ) else: raise ValueError(f"Unsupported interactive playback algo: {algo}") except RuntimeError as exc: - if str(exc) in {_PLAYBACK_ENV_UNAVAILABLE, _HORA_DISTILL_CHECKPOINT_UNAVAILABLE}: + if str(exc) == _PLAYBACK_ENV_UNAVAILABLE: return raise playback_session = session[0] diff --git a/src/unilab/scripts/train_offpolicy.py b/src/unilab/scripts/train_offpolicy.py index d02efd727..aca2fb7ff 100644 --- a/src/unilab/scripts/train_offpolicy.py +++ b/src/unilab/scripts/train_offpolicy.py @@ -57,7 +57,6 @@ RslRlPlaybackConfig, create_sac_playback_session, default_device, - resolve_play_actor_spec, resolve_play_obs_dims, ) from unilab.visualization.interactive_playback import ( @@ -344,28 +343,12 @@ def play_offpolicy(algo_name: str, cfg: DictConfig) -> str | None: env = cast(Any, session.env) actor = session.actor normalizer = session.normalizer - actor_algo_type = session.actor_algo_type # Export actor to ONNX if load_path_dir is not None and bool(getattr(cfg.training, "export_onnx", True)): - obs_dim, critic_obs_dim = resolve_play_obs_dims(env.obs_groups_spec) - _, actor_kwargs = resolve_play_actor_spec( - algo_name, - cfg, - obs_dim=obs_dim, - critic_obs_dim=critic_obs_dim, - ) + obs_dim, _ = resolve_play_obs_dims(env.obs_groups_spec) onnx_path = os.path.join(load_path_dir, "policy.onnx") dummy_input = torch.randn(1, obs_dim, device=device) - dummy_priv_info = ( - torch.zeros( - (1, int(actor_kwargs["priv_info_dim"])), - device=device, - dtype=dummy_input.dtype, - ) - if actor_algo_type == "hora_sac" - else None - ) with torch.inference_mode(): if normalizer: dummy_input = normalizer(dummy_input, update=False) @@ -374,24 +357,15 @@ def play_offpolicy(algo_name: str, cfg: DictConfig) -> str | None: export_module = actor.as_export_module() else: export_module = actor - export_inputs = ( - (dummy_input, dummy_priv_info) if dummy_priv_info is not None else (dummy_input,) - ) - input_names = ["obs", "priv_info"] if dummy_priv_info is not None else ["obs"] + export_inputs = (dummy_input,) + input_names = ["obs"] export_policy_onnx(export_module, onnx_path, export_inputs, input_names=input_names) # Verify ONNX output matches PyTorch verify_input = torch.randn(1, obs_dim, device=device) with torch.inference_mode(): onnx_feed = normalizer(verify_input, update=False) if normalizer else verify_input - verify_priv_info = ( - torch.zeros((1, int(actor_kwargs["priv_info_dim"])), device=device) - if actor_algo_type == "hora_sac" - else None - ) - verify_inputs = ( - (onnx_feed, verify_priv_info) if verify_priv_info is not None else (onnx_feed,) - ) + verify_inputs = (onnx_feed,) verify_policy_onnx(export_module, onnx_path, verify_inputs, input_names=input_names) elif load_path_dir is not None: print("Skipping ONNX export because training.export_onnx=false.") diff --git a/src/unilab/tasks/__init__.py b/src/unilab/tasks/__init__.py index 6d59d3958..9122364d1 100644 --- a/src/unilab/tasks/__init__.py +++ b/src/unilab/tasks/__init__.py @@ -12,7 +12,6 @@ "unilab.tasks.locomotion.g1", "unilab.tasks.locomotion.a2", "unilab.tasks.manipulation.allegro_inhand", - "unilab.tasks.manipulation.sharpa_inhand", "unilab.tasks.manipulation.stewart", "unilab.tasks.motion_tracking.g1", "unilab.tasks.motion_tracking.x2", diff --git a/src/unilab/tasks/manipulation/sharpa_inhand/__init__.py b/src/unilab/tasks/manipulation/sharpa_inhand/__init__.py deleted file mode 100644 index 79e9b91b2..000000000 --- a/src/unilab/tasks/manipulation/sharpa_inhand/__init__.py +++ /dev/null @@ -1,19 +0,0 @@ -from . import ( - grasp_gen, # registers SharpaInhandRotationGrasp via @registry decorators - rotation, # registers SharpaInhandRotation via @registry decorators -) -from .grasp_gen import ( - SharpaInhandGraspEnvCfg, - SharpaInhandRotationGraspCfg, - SharpaInhandRotationGraspEnv, -) -from .rotation import RewardConfig, SharpaInhandRotationCfg, SharpaInhandRotationEnv - -__all__ = [ - "RewardConfig", - "SharpaInhandRotationCfg", - "SharpaInhandRotationEnv", - "SharpaInhandRotationGraspCfg", - "SharpaInhandGraspEnvCfg", - "SharpaInhandRotationGraspEnv", -] diff --git a/src/unilab/tasks/manipulation/sharpa_inhand/base.py b/src/unilab/tasks/manipulation/sharpa_inhand/base.py deleted file mode 100644 index 741cbe4e4..000000000 --- a/src/unilab/tasks/manipulation/sharpa_inhand/base.py +++ /dev/null @@ -1,670 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass, field -from pathlib import Path -from typing import Any, Sequence, cast - -import gymnasium as gym -import numpy as np -from unisim import SimBackend - -from unilab.assets import ASSETS_ROOT_PATH -from unilab.base.base import EnvCfg -from unilab.base.np_env import NpEnv, NpEnvState -from unilab.base.scene import SceneCfg -from unilab.dtype_config import get_global_dtype - -DEFAULT_ACTUATED_JOINT_NAMES: list[str] = [ - "right_thumb_CMC_FE", - "right_thumb_CMC_AA", - "right_thumb_MCP_FE", - "right_thumb_MCP_AA", - "right_thumb_IP", - "right_index_MCP_FE", - "right_index_MCP_AA", - "right_index_PIP", - "right_index_DIP", - "right_middle_MCP_FE", - "right_middle_MCP_AA", - "right_middle_PIP", - "right_middle_DIP", - "right_ring_MCP_FE", - "right_ring_MCP_AA", - "right_ring_PIP", - "right_ring_DIP", - "right_pinky_CMC", - "right_pinky_MCP_FE", - "right_pinky_MCP_AA", - "right_pinky_PIP", - "right_pinky_DIP", -] - -DEFAULT_FINGERTIP_BODY_NAMES: list[str] = [ - "right_thumb_DP", - "right_index_DP", - "right_middle_DP", - "right_ring_DP", - "right_pinky_DP", -] - -# Source parity anchor from sharpa-rl-lab: -# rl_isaaclab/tasks/inhand_rotate/sharpa_wave_env_cfg.py (hand init_state joint_pos). -SOURCE_DEFAULT_HAND_JOINT_POS_DEG: tuple[float, ...] = ( - 95.12771, - -3.11244, - 14.81626, - -1.03493, - 12.23986, - 65.21091, - 6.1133, - 15.58495, - 5.90325, - 31.74149, - -0.95812, - 41.88173, - 12.844, - 31.72383, - 9.84458, - 35.22366, - 18.02839, - 10.9712, - 68.30895, - 7.99151, - 5.89626, - 5.89875, -) - - -@dataclass -class SharpaControlConfig: - action_scale: float = 1.0 / 24.0 - # MuJoCo Sharpa loads PD defaults from XML actuator gains. These fields remain - # fallback defaults for backends that cannot expose actuator gains yet. - p_gain: float = 1.0 - d_gain: float = 0.1 - # MuJoCo Sharpa currently uses position actuators, so torque-control mode - # stays declared here for owner-config structure but is rejected at runtime. - torque_control: bool = False - dof_limits_scale: float = 0.9 - - -@dataclass -class SharpaSensorConfig: - tactile_force_sensor_names: list[str] = field(default_factory=list) - - -@dataclass -class SharpaObservationConfig: - observation_mode: str = "separated" - enable_tactile: bool = True - binary_contact: bool = False - enable_contact_pos: bool = False - contact_smooth: float = 0.5 - contact_threshold: float = 0.05 - tactile_force_clip_max: float = 4.0 - - -@dataclass -class SharpaPrivilegedInfoConfig: - include_friction_scale: bool = True - include_gravity_direction: bool = False - - -@dataclass -class SharpaDomainRandConfig: - scale_list: list[float] = field(default_factory=lambda: [0.5]) - randomize_base_mass: bool = False - added_mass_range: list[float] = field(default_factory=lambda: [0.0, 0.0]) - random_com: bool = False - com_offset_x: list[float] = field(default_factory=lambda: [0.0, 0.0]) - randomize_gravity: bool = False - gravity_range: list[list[float]] = field( - default_factory=lambda: [[0.0, 0.0, -9.81], [0.0, 0.0, -9.81]] - ) - randomize_gravity_direction: bool = False - gravity_direction_magnitude: float = 9.81 - randomize_pd_gains: bool = True - randomize_p_gain_scale_lower: float = 0.5 - randomize_p_gain_scale_upper: float = 2.0 - randomize_d_gain_scale_lower: float = 0.5 - randomize_d_gain_scale_upper: float = 2.0 - randomize_friction: bool = True - randomize_friction_scale_lower: float = 0.5 - randomize_friction_scale_upper: float = 2.0 - elastomer_base_friction: float = 1.6 - metal_base_friction: float = 0.2 - object_base_friction: float = 1.0 - randomize_com: bool = True - randomize_com_lower: float = -0.01 - randomize_com_upper: float = 0.01 - randomize_mass: bool = True - randomize_mass_lower: float = 0.01 - randomize_mass_upper: float = 0.25 - force_scale: float = 2.0 - random_force_prob_scalar: float = 0.25 - force_decay: float = 0.9 - force_decay_interval: float = 0.08 - joint_noise_scale: float = 0.02 - contact_latency: float = 0.005 - contact_sensor_noise: float = 0.01 - push_body_name: str | None = None - - -@dataclass -class SharpaInhandBaseCfg(EnvCfg): - scene: SceneCfg = field( # pyright: ignore[reportIncompatibleVariableOverride] - default_factory=lambda: SceneCfg( - model_file=str(ASSETS_ROOT_PATH / "robots" / "sharpa_wave" / "scene.xml") - ) - ) - max_episode_seconds: float = 20.0 # pyright: ignore[reportIncompatibleVariableOverride] - sim_dt: float = 1.0 / 240.0 - ctrl_dt: float = 12.0 / 240.0 - - action_space: int = 22 - observation_space: int = 192 - prop_hist_len: int = 30 - critic_info_dim: int = 8 - - clip_obs: float = 5.0 - clip_actions: float = 1.0 - - num_hand_dofs: int = 22 - frame_obs_dim: int = 64 - obs_lag_steps: int = 3 - obs_history_len: int = 80 - - base_name: str = "right_hand_C_MC" - object_body_name: str = "object" - object_geom_name: str = "object" - actuated_joint_names: list[str] = field( - default_factory=lambda: list(DEFAULT_ACTUATED_JOINT_NAMES) - ) - fingertip_body_names: list[str] = field( - default_factory=lambda: list(DEFAULT_FINGERTIP_BODY_NAMES) - ) - - control_config: SharpaControlConfig = field(default_factory=SharpaControlConfig) - sensor: SharpaSensorConfig = field(default_factory=SharpaSensorConfig) # type: ignore[assignment] - obs: SharpaObservationConfig = field(default_factory=SharpaObservationConfig) - priv_info: SharpaPrivilegedInfoConfig = field(default_factory=SharpaPrivilegedInfoConfig) - domain_rand: SharpaDomainRandConfig = field(default_factory=SharpaDomainRandConfig) - - reset_height_lower: float = 0.59906 - reset_height_upper: float = 0.63906 - reset_angle_diff: float = 45.0 / 180.0 * np.pi - - rot_axis: tuple[float, float, float] = (0.0, 0.0, 1.0) - - grasp_cache_path: str = str(ASSETS_ROOT_PATH / "caches" / "sharpa_grasp_linspace") - disable_tactile_ids: list[int] = field(default_factory=list) - # Match the reference Sharpa object-position reward/privileged-info anchor - # by using the fixed XML/default object pose instead of the sampled grasp reset. - use_default_object_pose_for_object_pos_anchor: bool = False - - debug_show_axes: bool = False - - -def format_scale_tag(scale_value: float) -> str: - """Convert one object scale into a stable cache filename tag. - - Args: - scale_value: Single object scale value. - - Returns: - Scale tag used in cache filenames. - """ - scale_value = float(scale_value) - if scale_value <= 0.0: - raise ValueError(f"scale values must be positive, got {scale_value}") - return f"{scale_value:g}" - - -def resolve_grasp_cache_file(grasp_cache_path: str, scale_value: float) -> Path: - """Resolve the grasp cache path for a single object scale. - - Args: - grasp_cache_path: Configured cache prefix or template path. - scale_value: Single object scale value for this cache file. - - Returns: - Cache path for that exact scale. - """ - scale_tag = format_scale_tag(scale_value) - if "{scale}" in grasp_cache_path: - return Path(grasp_cache_path.format(scale=scale_tag)) - - base = Path(grasp_cache_path) - if base.suffix == ".npy": - return base.with_name(f"{base.stem}_{scale_tag}{base.suffix}") - return Path(f"{grasp_cache_path}_{scale_tag}.npy") - - -def sample_scale_grasp_caches( - grasp_caches: Sequence[np.ndarray], - scale_ids: np.ndarray, -) -> np.ndarray: - """Sample one cached grasp per reset environment from per-scale cache files. - - Args: - grasp_caches: Cache arrays ordered the same way as env scale ids. - scale_ids: Scale-bucket assignment for each reset environment. - - Returns: - Cached grasp states with shape ``(num_envs, 29)``. - """ - num_envs = scale_ids.shape[0] - num_scales = len(grasp_caches) - if num_scales <= 0: - raise ValueError("grasp_caches must contain at least one scale bucket") - - sampled = np.zeros((num_envs, 29), dtype=np.float64) - for scale_idx, grasp_cache in enumerate(grasp_caches): - if grasp_cache.ndim != 2 or grasp_cache.shape[1] < 29: - raise ValueError(f"Expected cached grasp shape (?, 29), got {grasp_cache.shape}") - if grasp_cache.shape[0] == 0: - raise ValueError(f"grasp cache for scale id {scale_idx} is empty") - env_ids = np.flatnonzero(scale_ids == scale_idx) - if len(env_ids) == 0: - continue - sample_ids = np.random.randint(0, grasp_cache.shape[0], size=len(env_ids)) - sampled[env_ids] = grasp_cache[sample_ids] - return sampled - - -def repeat_obs_history(init_frame: np.ndarray, history_len: int) -> np.ndarray: - history = np.broadcast_to( - init_frame[:, None, :], (init_frame.shape[0], history_len, init_frame.shape[1]) - ).copy() - return np.asarray(history, dtype=init_frame.dtype) - - -class SharpaInhandBaseEnv(NpEnv): - _cfg: SharpaInhandBaseCfg - _default_p_gain: np.ndarray - _default_d_gain: np.ndarray - - def __init__(self, cfg: SharpaInhandBaseCfg, backend: SimBackend, num_envs: int = 1) -> None: - super().__init__(cfg, backend, num_envs) - - self._np_dtype = get_global_dtype() - - self._num_action = int(cfg.num_hand_dofs) - actuator_range = np.asarray(self._backend.get_actuator_ctrl_range(), dtype=self._np_dtype) - if actuator_range.shape[0] < self._num_action: - raise ValueError( - f"Model has {actuator_range.shape[0]} actuators, but Sharpa task needs {self._num_action}" - ) - - # Keep the raw XML actuator limits for observation normalization and any - # logic that needs the original backend contract. Target-position clipping - # uses a separate scaled limit pair to match Sharpa source behavior. - self._ctrl_lower = np.asarray(actuator_range[: self._num_action, 0], dtype=self._np_dtype) - self._ctrl_upper = np.asarray(actuator_range[: self._num_action, 1], dtype=self._np_dtype) - self._target_lower, self._target_upper = self._resolve_target_joint_limits( - self._ctrl_lower, - self._ctrl_upper, - cfg.control_config.dof_limits_scale, - ) - - self._init_qpos = self._resolve_init_qpos() - self._init_qvel = np.asarray(self._backend.get_init_qvel(), dtype=np.float64) - self.nq = int(self._init_qpos.shape[0]) - self.nv = int(self._init_qvel.shape[0]) - - if self.nq < self._num_action + 7: - raise ValueError( - f"Model qpos dim {self.nq} is too small for {self._num_action} hand DoFs + object pose" - ) - - self._obj_pos_slice = slice(self._num_action, self._num_action + 3) - self._obj_quat_slice = slice(self._num_action + 3, self._num_action + 7) - - source_default_angles = np.deg2rad( - np.asarray(SOURCE_DEFAULT_HAND_JOINT_POS_DEG, dtype=np.float64) - ) - if source_default_angles.shape[0] != self._num_action: - raise ValueError( - "Source default hand joint pose size mismatch: " - f"{source_default_angles.shape[0]} vs expected {self._num_action}" - ) - self.default_angles = np.asarray(source_default_angles, dtype=self._np_dtype) - - self._action_space = gym.spaces.Box( - low=-float(cfg.clip_actions), - high=float(cfg.clip_actions), - shape=(self._num_action,), - dtype=np.float32, - ) - - self._object_body_ids = self._backend.get_body_ids([cfg.object_body_name]) - self._fingertip_body_ids = self._backend.get_body_ids(cfg.fingertip_body_names) - self._object_geom_base_size = self._resolve_object_geom_base_size() - - self._num_tactile = len(cfg.fingertip_body_names) - self.last_contacts = np.zeros((num_envs, self._num_tactile), dtype=self._np_dtype) - self._prev_tactile_force = np.zeros((num_envs, self._num_tactile), dtype=self._np_dtype) - - self.object_default_pose = np.zeros((num_envs, 7), dtype=self._np_dtype) - - self.obs_buf_lag_history = np.zeros( - (num_envs, cfg.obs_history_len, cfg.frame_obs_dim), dtype=self._np_dtype - ) - self.proprio_hist_buf = np.zeros( - (num_envs, cfg.prop_hist_len, cfg.frame_obs_dim), dtype=self._np_dtype - ) - self.critic_info_buf = np.zeros((num_envs, cfg.critic_info_dim), dtype=self._np_dtype) - - self.scale_ids, self._num_scales, self._bucket_env = self._build_scale_ids( - num_envs, cfg.domain_rand.scale_list - ) - self.scale_values = self._build_scale_values(cfg.domain_rand.scale_list) - - @property - def action_space(self) -> gym.spaces.Box: - return self._action_space # type: ignore[no-any-return] - - def _resolve_target_joint_limits( - self, - raw_lower: np.ndarray, - raw_upper: np.ndarray, - scale: float, - ) -> tuple[np.ndarray, np.ndarray]: - """Build scaled target-position clipping limits from raw XML actuator bounds. - - Args: - raw_lower: Original lower control limits loaded from the backend XML. - raw_upper: Original upper control limits loaded from the backend XML. - scale: Multiplicative scale applied to both bound arrays. - - Returns: - Tuple of scaled ``(lower, upper)`` target-position limits. - """ - scale_value = float(scale) - if scale_value <= 0.0: - raise ValueError(f"dof_limits_scale must be positive, got {scale_value}") - target_lower = np.asarray(raw_lower, dtype=self._np_dtype) * scale_value - target_upper = np.asarray(raw_upper, dtype=self._np_dtype) * scale_value - if np.any(target_lower > target_upper): - raise ValueError("Scaled Sharpa target joint limits are invalid") - return target_lower.astype(self._np_dtype, copy=False), target_upper.astype( - self._np_dtype, - copy=False, - ) - - def _resolve_init_qpos(self) -> np.ndarray: - for key_name in ("home", "stand", "default"): - try: - return np.asarray(self._backend.get_keyframe_qpos(key_name), dtype=np.float64) - except Exception: - continue - - try: - return np.asarray(self._backend.get_default_qpos(), dtype=np.float64) - except NotImplementedError as exc: - raise ValueError("Could not resolve initial qpos from backend contract") from exc - - def _build_scale_ids( - self, num_envs: int, scale_list: Sequence[float] - ) -> tuple[np.ndarray, int, int]: - """Build deterministic near-even environment assignments for each scale. - - Args: - num_envs: Number of vectorized environments to assign. - scale_list: Explicit object scale values used by this env instance. - - Returns: - Tuple of scale id per environment, total scale count, and the minimum - number of environments assigned to any scale. - """ - if len(scale_list) == 0: - raise ValueError("scale_list must contain at least one scale") - scale_values = np.asarray(scale_list, dtype=np.float64) - if np.any(scale_values <= 0.0): - raise ValueError(f"scale_list values must be positive, got {list(scale_list)}") - num_scales = int(scale_values.shape[0]) - if num_scales <= 0: - raise ValueError(f"scale_list must contain at least one value, got {list(scale_list)}") - - bucket_env = num_envs // num_scales - remainder = num_envs % num_scales - # Assign the remainder to the lowest scale ids to keep bucket sizes within one env. - counts = np.full((num_scales,), bucket_env, dtype=np.int32) - counts[:remainder] += 1 - scale_ids = np.repeat(np.arange(num_scales, dtype=np.int32), counts) - return scale_ids, num_scales, bucket_env - - def _build_scale_values(self, scale_list: Sequence[float]) -> np.ndarray: - """Normalize configured scale values into a stable numpy array. - - Args: - scale_list: Explicit list of object scales. - - Returns: - Array of configured scale values in config order. - """ - scale_values = np.asarray(scale_list, dtype=np.float64) - if scale_values.ndim != 1 or scale_values.size == 0: - raise ValueError(f"scale_list must be a non-empty flat list, got {list(scale_list)}") - if np.any(scale_values <= 0.0): - raise ValueError(f"scale_list values must be positive, got {list(scale_list)}") - return scale_values - - def _resolve_object_geom_base_size(self) -> np.ndarray | None: - try: - return cast(np.ndarray, self._backend.get_geom_size(self._cfg.object_geom_name)) - except NotImplementedError: - return None - - def apply_action(self, actions: np.ndarray, state: NpEnvState) -> np.ndarray: - clipped_actions = np.clip(actions, -self._cfg.clip_actions, self._cfg.clip_actions) - clipped_actions = np.asarray(clipped_actions[:, : self._num_action], dtype=self._np_dtype) - - state.info["last_actions"] = state.info.get("current_actions", clipped_actions.copy()) - state.info["current_actions"] = clipped_actions - - prev_targets = state.info.get( - "prev_targets", - np.broadcast_to(self.default_angles, (self._num_envs, self._num_action)).copy(), - ) - targets = prev_targets + self._cfg.control_config.action_scale * clipped_actions - # Clip action targets by the scaled control range only. Observation - # normalization continues to use the raw XML actuator limits. - targets = np.clip(targets, self._target_lower, self._target_upper) - prev_targets = np.asarray(targets, dtype=self._np_dtype) - state.info["prev_targets"] = prev_targets - return prev_targets - - def get_hand_dof_pos(self) -> np.ndarray: - return np.asarray(self._backend.get_dof_pos()[:, : self._num_action], dtype=self._np_dtype) - - def get_hand_dof_vel(self) -> np.ndarray: - return np.asarray(self._backend.get_dof_vel()[:, : self._num_action], dtype=self._np_dtype) - - def get_fingertip_pos(self) -> np.ndarray: - return np.asarray( - self._backend.get_body_pos_w(self._fingertip_body_ids), dtype=self._np_dtype - ) - - def get_object_pos(self) -> np.ndarray: - return np.asarray( - self._backend.get_body_pos_w(self._object_body_ids)[:, 0, :], dtype=self._np_dtype - ) - - def get_object_quat(self) -> np.ndarray: - return np.asarray( - self._backend.get_body_quat_w(self._object_body_ids)[:, 0, :], dtype=self._np_dtype - ) - - def _extract_sensor_scalar(self, sensor_name: str) -> np.ndarray: - data = np.asarray(self._backend.get_sensor_data(sensor_name), dtype=self._np_dtype) - if data.ndim == 1: - return data - if data.ndim == 2 and data.shape[1] == 1: - return data[:, 0] - if data.ndim == 2 and data.shape[1] >= 3: - return np.asarray(np.linalg.norm(data[:, :3], axis=1), dtype=self._np_dtype) - flat = data.reshape(data.shape[0], -1) - return np.asarray(flat[:, 0], dtype=self._np_dtype) - - def _read_tactile_force(self) -> np.ndarray: - """Read per-finger tactile force magnitudes in configured sensor order. - - Args: - None. - - Returns: - Array of shape ``(num_envs, num_tactile)`` ordered exactly as - ``sensor.tactile_force_sensor_names``. - """ - tactile_force = np.zeros((self._num_envs, self._num_tactile), dtype=self._np_dtype) - if not self._cfg.sensor.tactile_force_sensor_names: - return tactile_force - - for sensor_id, sensor_name in enumerate( - self._cfg.sensor.tactile_force_sensor_names[: self._num_tactile] - ): - try: - tactile_force[:, sensor_id] = self._extract_sensor_scalar(sensor_name) - except Exception: - tactile_force[:, sensor_id] = 0.0 - return tactile_force - - def _clip_tactile_force(self, tactile_force: np.ndarray) -> np.ndarray: - """Clip raw tactile-force magnitudes before they enter observation smoothing. - - Args: - tactile_force: Raw per-finger tactile magnitudes with shape - ``(num_envs, num_tactile)``. - - Returns: - Clipped tactile-force array with the same shape. Non-positive clip values - disable this clamp so the caller can opt out explicitly. - """ - clip_max = float(self._cfg.obs.tactile_force_clip_max) - tactile_force = np.asarray(tactile_force, dtype=self._np_dtype) - if clip_max <= 0.0: - return tactile_force - return np.asarray(np.clip(tactile_force, 0.0, clip_max), dtype=self._np_dtype) - - def _clear_tactile_history(self, env_ids: np.ndarray | None = None) -> None: - """Clear tactile-output and raw-force history buffers. - - Args: - env_ids: Optional environment ids to clear. When ``None``, clear all envs. - - Returns: - None. Buffers are updated in place. - """ - if env_ids is None: - self.last_contacts.fill(0.0) - self._prev_tactile_force.fill(0.0) - return - - self.last_contacts[env_ids] = 0.0 - self._prev_tactile_force[env_ids] = 0.0 - - def _compute_tactile_observation(self) -> np.ndarray: - """Build tactile observations with source-equivalent smoothing and latency. - - Args: - None. - - Returns: - Tactile-force observation array with shape ``(num_envs, num_tactile)``. - """ - obs_cfg = self._cfg.obs - domain_rand = self._cfg.domain_rand - if not obs_cfg.enable_tactile: - self._clear_tactile_history() - return np.zeros((self._num_envs, self._num_tactile), dtype=self._np_dtype) - - current_force = SharpaInhandBaseEnv._clip_tactile_force(self, self._read_tactile_force()) - smooth_contact = ( - current_force * obs_cfg.contact_smooth - + self._prev_tactile_force * (1.0 - obs_cfg.contact_smooth) - ).astype(self._np_dtype) - self._prev_tactile_force[:] = current_force - - for disabled_id in self._cfg.disable_tactile_ids: - if 0 <= disabled_id < self._num_tactile: - smooth_contact[:, disabled_id] = 0.0 - - latency = np.where( - np.random.rand(self._num_envs, self._num_tactile) < domain_rand.contact_latency, - 1.0, - 0.0, - ).astype(self._np_dtype) - - if obs_cfg.binary_contact: - binary_contact = (smooth_contact > obs_cfg.contact_threshold).astype(self._np_dtype) - self.last_contacts = self.last_contacts * latency + binary_contact * (1.0 - latency) - noise_mask = ( - np.random.rand(self._num_envs, self._num_tactile) - >= domain_rand.contact_sensor_noise - ).astype(self._np_dtype) - return np.where( - self.last_contacts > 0.1, - noise_mask * self.last_contacts, - self.last_contacts, - ) - - self.last_contacts = self.last_contacts * latency + smooth_contact * (1.0 - latency) - return self.last_contacts.copy() - - def _compute_contact_positions(self, tactile: np.ndarray) -> np.ndarray: - del tactile - # TODO(sharpa_inhand): IsaacLab contact positions are defined through contact sensor - # frame transforms. Backend-level contact point parity is not available yet in the - # current UniLab sensor contract, so we keep this channel as zeros for now. - return np.zeros((self._num_envs, self._num_tactile * 3), dtype=self._np_dtype) - - def _normalize_joint_pos(self, dof_pos: np.ndarray) -> np.ndarray: - return np.asarray( - (2.0 * dof_pos - self._ctrl_upper - self._ctrl_lower) - / (self._ctrl_upper - self._ctrl_lower + 1.0e-8), - dtype=self._np_dtype, - ) - - def _sample_pd_scales(self, lower: float, upper: float, shape: tuple[int, int]) -> np.ndarray: - if lower > 1.0 or upper < 1.0: - raise ValueError("PD randomization scales must satisfy lower <= 1 <= upper") - small = np.random.uniform(lower, 1.0, size=shape) - large = np.random.uniform(1.0, upper, size=shape) - use_small = np.random.rand(*shape) > 0.5 - return np.where(use_small, small, large).astype(self._np_dtype) - - def _load_default_pd_gains(self) -> tuple[np.ndarray, np.ndarray]: - """Resolve the default per-DOF PD gains used as the randomization baseline. - - Args: - None. - - Returns: - Tuple of ``(p_gain, d_gain)`` arrays with shape ``(num_action,)``. - """ - try: - p_gain, d_gain = self._backend.get_actuator_gains() - return ( - np.asarray(p_gain[: self._num_action], dtype=self._np_dtype).copy(), - np.asarray(d_gain[: self._num_action], dtype=self._np_dtype).copy(), - ) - except NotImplementedError: - return ( - np.full((self._num_action,), self._cfg.control_config.p_gain, dtype=self._np_dtype), - np.full((self._num_action,), self._cfg.control_config.d_gain, dtype=self._np_dtype), - ) - - def _resolve_pd_gains(self, info: dict[str, Any]) -> tuple[np.ndarray, np.ndarray]: - p_gain = info.get( - "p_gain", - np.broadcast_to(self._default_p_gain, (self._num_envs, self._num_action)).copy(), - ) - d_gain = info.get( - "d_gain", - np.broadcast_to(self._default_d_gain, (self._num_envs, self._num_action)).copy(), - ) - return np.asarray(p_gain, dtype=self._np_dtype), np.asarray(d_gain, dtype=self._np_dtype) - - def _update_proprio_history(self, obs_history: np.ndarray) -> np.ndarray: - return np.asarray(obs_history[:, -self._cfg.prop_hist_len :], dtype=self._np_dtype) diff --git a/src/unilab/tasks/manipulation/sharpa_inhand/grasp_gen.py b/src/unilab/tasks/manipulation/sharpa_inhand/grasp_gen.py deleted file mode 100644 index 77c602f77..000000000 --- a/src/unilab/tasks/manipulation/sharpa_inhand/grasp_gen.py +++ /dev/null @@ -1,402 +0,0 @@ -from __future__ import annotations - -import logging -from dataclasses import dataclass, field -from typing import Any - -import numpy as np - -from unilab.assets import ASSETS_ROOT_PATH -from unilab.base import registry -from unilab.base.np_env import NpEnvState -from unilab.base.run_control import RunComplete -from unilab.dr import ResetPlan -from unilab.dr.dr_utils import build_common_reset_randomization -from unilab.tasks.compatibility import adapt_legacy_factory -from unilab.tasks.manipulation.sharpa_inhand.base import ( - SOURCE_DEFAULT_HAND_JOINT_POS_DEG, - SharpaDomainRandConfig, - resolve_grasp_cache_file, -) -from unilab.tasks.manipulation.sharpa_inhand.rotation import ( - RewardConfig, - SharpaInhandRotationCfg, - SharpaInhandRotationDRProvider, - SharpaInhandRotationEnv, -) -from unilab.utils.rotation import np_quat_error_magnitude - -logger = logging.getLogger(__name__) - - -def _default_sharpa_grasp_domain_rand() -> SharpaDomainRandConfig: - """Build the nested DR defaults used by Sharpa grasp collection. - - Returns: - Domain-randomization config matching the grasp-task owner defaults. - """ - return SharpaDomainRandConfig( - randomize_pd_gains=False, - randomize_friction=False, - randomize_com=False, - randomize_mass=True, - randomize_mass_lower=0.05, - randomize_mass_upper=0.051, - force_scale=0.0, - random_force_prob_scalar=0.0, - ) - - -@dataclass -class SharpaInhandRotationGraspCfg(SharpaInhandRotationCfg): - max_episode_seconds: float = 3.0 # 12.0 - - reset_height_lower: float = 0.61406 - reset_height_upper: float = 0.62406 - reset_angle_diff: float = 30.0 / 180.0 * np.pi - - grasp_cache_path: str = "" - domain_rand: SharpaDomainRandConfig = field(default_factory=_default_sharpa_grasp_domain_rand) - - reward_config: RewardConfig = field( # pyright: ignore[reportIncompatibleVariableOverride] - default_factory=lambda: RewardConfig( - scales={ - "rotate": 0.0, - "obj_linvel": 0.0, - "pose_diff": 0.0, - "torque": 0.0, - "work": 0.0, - "object_pos": 0.0, - } - ) - ) - - grasp_collection_target: int = 50_000 - grasp_auto_save: bool = True - - -@registry.envcfg("SharpaInhandRotationGrasp") -@dataclass -class SharpaInhandGraspEnvCfg(SharpaInhandRotationGraspCfg): - pass - - -class SharpaInhandGraspDRProvider(SharpaInhandRotationDRProvider): - def build_reset_plan(self, env: Any, env_ids: np.ndarray) -> ResetPlan: - # Keep original grasp task behavior: collect successful pre-reset states on each reset. - env._collect_successful_grasps(env_ids) - - num_reset = len(env_ids) - if num_reset == 0: - return ResetPlan( - env_ids=env_ids, - qpos=np.zeros((0, env.nq), dtype=np.float64), - qvel=np.zeros((0, env.nv), dtype=np.float64), - info_updates={}, - randomization=None, - ) - - rand = 2.0 * np.random.rand(num_reset, env._num_action) - 1.0 - hand_qpos = np.broadcast_to(env._grasp_default_angles, (num_reset, env._num_action)).copy() - hand_qpos += 0.15 * rand - hand_qpos = np.clip(hand_qpos, env._ctrl_lower, env._ctrl_upper) - - object_pos = np.broadcast_to(env._init_qpos[env._obj_pos_slice], (num_reset, 3)).copy() - object_quat = np.broadcast_to(env._init_qpos[env._obj_quat_slice], (num_reset, 4)).copy() - - qpos = np.zeros((num_reset, env.nq), dtype=np.float64) - qpos[:, : env._num_action] = hand_qpos - qpos[:, env._obj_pos_slice] = object_pos - qpos[:, env._obj_quat_slice] = object_quat - - qvel = np.zeros((num_reset, env.nv), dtype=np.float64) - p_gain, d_gain = self._sample_reset_pd_gains(env, num_reset, dtype=env._np_dtype) - - info_updates = self._build_info_updates( - env, - env_ids=env_ids, - hand_qpos=hand_qpos, - object_pos=object_pos, - object_quat=object_quat, - reset_height_lower=np.full((num_reset,), env.cfg.reset_height_lower, dtype=np.float64), - reset_height_upper=np.full((num_reset,), env.cfg.reset_height_upper, dtype=np.float64), - rot_axis=np.broadcast_to(env._rot_axis, (num_reset, 3)).astype(np.float64), - p_gain=p_gain, - d_gain=d_gain, - friction_scale=None, - randomized_mass=None, - randomized_com_offset=None, - gravity=None, - ) - env._clear_tactile_history(env_ids) - - return ResetPlan( - env_ids=env_ids, - qpos=qpos, - qvel=qvel, - info_updates=info_updates, - randomization=build_common_reset_randomization(env, num_reset), - ) - - -class SharpaInhandRotationGraspEnv(SharpaInhandRotationEnv): - _cfg: SharpaInhandRotationGraspCfg # pyright: ignore[reportIncompatibleVariableOverride] - _MATERIALIZE_ROTATION_GRASP_CACHE = False - - def __init__( - self, - cfg: SharpaInhandRotationGraspCfg, - num_envs: int = 1, - backend_type: str = "motrix", - ) -> None: - if cfg.domain_rand.randomize_gravity or cfg.domain_rand.randomize_gravity_direction: - raise ValueError( - "SharpaInhandRotationGrasp does not support gravity randomization; " - "disable env.domain_rand.randomize_gravity and " - "env.domain_rand.randomize_gravity_direction." - ) - super().__init__( - cfg, - num_envs=num_envs, - backend_type=backend_type, - dr_provider=SharpaInhandGraspDRProvider(), - ) - - self._saved_grasping_states: list[list[np.ndarray]] = [ - list() for _ in range(self._num_scales) - ] - if self._num_scales != 1: - raise ValueError( - "Sharpa grasp generation now collects exactly one object scale per run; " - f"got scale_list={list(self.scale_values)}" - ) - self._grasp_target_per_scale = max(1, int(cfg.grasp_collection_target)) - self._grasp_cache_saved = False - self._grasp_target_reached_notified = False - self._last_grasp_progress_step = -1 - self._last_grasp_progress_counts: tuple[int, ...] | None = None - self._grasp_default_angles = np.asarray( - np.deg2rad(np.asarray(SOURCE_DEFAULT_HAND_JOINT_POS_DEG, dtype=np.float64)), - dtype=self._np_dtype, - ) - if self._grasp_default_angles.shape[0] != self._num_action: - raise ValueError( - "Source grasp default angle count mismatch: " - f"{self._grasp_default_angles.shape[0]} vs expected {self._num_action}" - ) - if np.any(np.bincount(self.scale_ids, minlength=self._num_scales) == 0): - raise ValueError( - "Sharpa grasp generation requires at least one environment for the configured scale; " - f"got num_envs={num_envs}, num_scales={self._num_scales}" - ) - - def apply_action(self, actions: np.ndarray, state: NpEnvState) -> np.ndarray: - # Grasp-cache collection should not use policy/random actions. - # Keep controls fixed at reset targets by forcing zero action input. - zero_actions = np.zeros_like(actions, dtype=self._np_dtype) - return super().apply_action(zero_actions, state) - - def _total_saved_grasps(self) -> int: - return int(sum(len(bucket) for bucket in self._saved_grasping_states)) - - def _collection_target_reached(self) -> bool: - return all( - len(bucket) >= self._grasp_target_per_scale for bucket in self._saved_grasping_states - ) - - def _get_per_scale_grasp_counts(self) -> tuple[int, ...]: - """Return collected grasp counts for each scale bucket. - - Returns: - Tuple where index is scale id and value is collected grasp count. - """ - return tuple(len(bucket) for bucket in self._saved_grasping_states) - - def _maybe_log_grasp_progress(self, force: bool = False) -> None: - """Log runtime grasp-collection progress grouped by scale. - - Args: - force: When True, log even if throttling would normally skip. - """ - if self.state is None: - return - - counts = self._get_per_scale_grasp_counts() - step_info = self.state.info.get("steps") - step = int(step_info[0]) if isinstance(step_info, np.ndarray) and step_info.size > 0 else 0 - if not force: - if counts == self._last_grasp_progress_counts: - return - if self._last_grasp_progress_step >= 0 and step - self._last_grasp_progress_step < 32: - return - - total = int(sum(counts)) - per_scale = ", ".join( - f"scale={float(self.scale_values[i]):g}:{count}" for i, count in enumerate(counts) - ) - logger.info( - "[SharpaInhandRotationGrasp] grasp progress total=%d/%d, per_scale=[%s]", - total, - int(self._cfg.grasp_collection_target), - per_scale, - ) - self._last_grasp_progress_step = step - self._last_grasp_progress_counts = counts - - def _stop_collection(self) -> None: - if self._grasp_target_reached_notified: - return - if not self._collection_target_reached(): - return - - self._maybe_log_grasp_progress(force=True) - self._grasp_target_reached_notified = True - collected = self._total_saved_grasps() - target = int(self._cfg.grasp_collection_target) - logger.info( - "[SharpaInhandRotationGrasp] Grasp collection target reached " - "(saved=%d, configured_target=%d). Collection completed.", - collected, - target, - ) - - if self.state is not None: - log = self.state.info.get("log", {}) - log["grasp/target_reached"] = 1.0 - self.state.info["log"] = log - - raise RunComplete( - reason="grasp_collection_target_reached", - summary={ - "collected_grasps": int(collected), - "grasp_collection_target": int(self._grasp_target_per_scale), - "grasp_collection_counts_by_scale": self._get_per_scale_grasp_counts(), - }, - ) - - def _collect_successful_grasps(self, env_ids: np.ndarray) -> None: - if self.state is None or len(env_ids) == 0: - return - - success_mask = self.state.truncated[env_ids] & ~self.state.terminated[env_ids] - if not np.any(success_mask): - return - - success_env_ids = env_ids[np.flatnonzero(success_mask)] - hand_qpos = self.get_hand_dof_pos()[success_env_ids] - object_pos = self.get_object_pos()[success_env_ids] - object_quat = self.get_object_quat()[success_env_ids] - all_states = np.concatenate([hand_qpos, object_pos, object_quat], axis=1).astype(np.float32) - - saved_scale_ids = self.scale_ids[success_env_ids] - for i, scale_id in enumerate(saved_scale_ids): - bucket = self._saved_grasping_states[int(scale_id)] - if len(bucket) < self._grasp_target_per_scale: - bucket.append(all_states[i : i + 1]) - - self._maybe_log_grasp_progress() - if self._grasp_cache_saved: - return - - finished_scales = sum( - int(len(bucket) >= self._grasp_target_per_scale) - for bucket in self._saved_grasping_states - ) - if finished_scales < self._num_scales: - return - - if not self._cfg.grasp_auto_save: - self._grasp_cache_saved = True - self._stop_collection() - return - - output_file = resolve_grasp_cache_file( - self._cfg.grasp_cache_path or "caches/sharpa_grasp_linspace", - float(self.scale_values[0]), - ) - if not output_file.is_absolute(): - output_file = ASSETS_ROOT_PATH / output_file - output_file.parent.mkdir(parents=True, exist_ok=True) - save_data = np.concatenate(self._saved_grasping_states[0], axis=0)[ - : self._grasp_target_per_scale - ] - np.save(output_file, save_data) - - self._grasp_cache_saved = True - if self.state is not None: - log = self.state.info.get("log", {}) - log["grasp_cache/saved"] = 1.0 - log["grasp_cache/num_states"] = float(save_data.shape[0]) - self.state.info["log"] = log - - self._stop_collection() - - def _compute_reward( - self, - info: dict[str, Any], - dof_pos: np.ndarray, - dof_vel: np.ndarray, - object_pos: np.ndarray, - object_linvel: np.ndarray, - object_angvel: np.ndarray, - torques: np.ndarray, - ) -> np.ndarray: - del info, dof_pos, dof_vel, object_pos, object_linvel, object_angvel, torques - return np.zeros((self._num_envs,), dtype=self._np_dtype) - - def update_state(self, state: NpEnvState) -> NpEnvState: - next_state = super().update_state(state) - - fingertip_pos = self.get_fingertip_pos() - object_pos = self.get_object_pos() - object_quat = self.get_object_quat() - object_default_pose = np.asarray( - next_state.info.get( - "object_default_pose", np.zeros((self._num_envs, 7), dtype=self._np_dtype) - ), - dtype=self._np_dtype, - ) - - cond1 = np.all( - np.linalg.norm(fingertip_pos - object_pos[:, None, :], axis=-1) < 0.1, axis=1 - ) - tactile = np.asarray(self.last_contacts, dtype=self._np_dtype) - cond2 = np.sum(tactile > 0.5, axis=1) >= 3 - quat_error = np_quat_error_magnitude(object_default_pose[:, 3:7], object_quat) - cond3 = quat_error < self._cfg.reset_angle_diff - - grasp_valid = cond1 & cond2 & cond3 - terminated = np.asarray(next_state.terminated | (~grasp_valid), dtype=bool) - - reward = np.zeros((self._num_envs,), dtype=self._np_dtype) - - step_count = next_state.info.get("steps", np.zeros((self._num_envs,), dtype=np.uint32)) - should_log = self._enable_reward_log and (int(step_count[0]) % 4 == 0) - if should_log: - log = next_state.info.get("log", {}) - log["grasp/cond1"] = float(np.mean(cond1.astype(np.float32))) - log["grasp/cond2"] = float(np.mean(cond2.astype(np.float32))) - log["grasp/cond3"] = float(np.mean(cond3.astype(np.float32))) - log["grasp/valid"] = float(np.mean(grasp_valid.astype(np.float32))) - per_scale_counts = self._get_per_scale_grasp_counts() - log["grasp/target_cache_size"] = float(self._cfg.grasp_collection_target) - for scale_idx, count in enumerate(per_scale_counts): - scale_value = float(self.scale_values[scale_idx]) - log[f"grasp/cache_size_scale_{scale_value:g}"] = float(count) - next_state.info["log"] = log - - return next_state.replace(reward=reward, terminated=terminated) - - -SharpaWaveGraspCfg = SharpaInhandGraspEnvCfg - -_SHARPA_GRASP_COMPAT_FACTORY = adapt_legacy_factory( - SharpaInhandRotationGraspEnv, - task_family="Sharpa", - reason=( - "tactile grasp validation and cache-collection completion remain task-owned " - "until formal Manager-Based capabilities exist" - ), -) -for _backend_type in ("mujoco", "motrix"): - registry.register_env("SharpaInhandRotationGrasp", _SHARPA_GRASP_COMPAT_FACTORY, _backend_type) diff --git a/src/unilab/tasks/manipulation/sharpa_inhand/rotation.py b/src/unilab/tasks/manipulation/sharpa_inhand/rotation.py deleted file mode 100644 index e940b6351..000000000 --- a/src/unilab/tasks/manipulation/sharpa_inhand/rotation.py +++ /dev/null @@ -1,1554 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass, field -from pathlib import Path -from typing import Any, cast - -import numpy as np -from unisim.dr.types import ( - RESET_TERM_BODY_IPOS, - RESET_TERM_BODY_MASS, - RESET_TERM_GEOM_FRICTION, - RESET_TERM_GRAVITY, - RESET_TERM_KD, - RESET_TERM_KP, - ResetRandomizationPayload, -) - -from unilab.assets.hub import resolve_grasp_cache_files -from unilab.base import registry -from unilab.base.backend_factory import create_backend, env_backend_kwargs -from unilab.base.np_env import NpEnvState -from unilab.dr import ( - INTERVAL_TERM_BODY_FORCE, - DomainRandomizationCapabilities, - DomainRandomizationProvider, - GeomSizeOverride, - InitRandomizationPlan, - IntervalRandomizationPlan, - IntervalTermOp, - ModelVariantSpec, - ResetPlan, -) -from unilab.dr.dr_utils import build_common_reset_randomization, validate_common_reset_randomization -from unilab.dtype_config import get_global_dtype -from unilab.tasks.compatibility import adapt_legacy_factory -from unilab.tasks.manipulation.sharpa_inhand.base import ( - SharpaInhandBaseCfg, - SharpaInhandBaseEnv, - repeat_obs_history, - resolve_grasp_cache_file, - sample_scale_grasp_caches, -) -from unilab.utils.geometry import np_sample_uniform_quaternion -from unilab.utils.rotation import ( - np_quat_apply, - np_quat_conjugate, - np_quat_mul, - np_quat_to_axis_angle, -) - - -@dataclass -class RewardConfig: - scales: dict[str, float] = field( - default_factory=lambda: { - "rotate": 2.5, - "obj_linvel": -0.3, - "pose_diff": -0.4, - "torque": -0.1, - "work": -0.5, - "object_pos": 0.003, - } - ) - angvel_clip_min: float = -0.5 - angvel_clip_max: float = 0.5 - - -@registry.envcfg("SharpaInhandRotation") -@dataclass -class SharpaInhandRotationCfg(SharpaInhandBaseCfg): - critic_info_dim: int = 9 - reward_config: RewardConfig | None = None - zero_action_test_mode: bool = False - - -def sample_random_quaternion(num_envs: int) -> np.ndarray: - """Sample uniformly distributed random quaternions in wxyz convention. - - Args: - num_envs: Number of quaternions to sample. - - Returns: - Quaternion array with shape ``(num_envs, 4)`` in float64. - """ - return np_sample_uniform_quaternion(num_envs).astype(np.float64) - - -def _materialize_grasp_caches( - grasp_cache_path: str, - scale_values: np.ndarray, -) -> tuple[np.ndarray, ...]: - """Load one Sharpa grasp cache for each scale on the env construction path.""" - grasp_caches: list[np.ndarray] = [] - missing_files: list[str] = [] - for scale_value in np.asarray(scale_values, dtype=np.float64): - cache_file = resolve_grasp_cache_file(grasp_cache_path, float(scale_value)) - resolved = cast(str, resolve_grasp_cache_files(str(cache_file))) - cache_file = Path(resolved) - if not cache_file.exists(): - missing_files.append(str(cache_file)) - continue - grasp_caches.append(np.load(cache_file).astype(np.float64)) - - if missing_files: - missing = ", ".join(missing_files) - raise RuntimeError( - f"Missing Sharpa grasp cache file(s): {missing}\n" - "Generate them with:\n" - " bash scripts/sharpa_collect_grasps.sh 0.8 0.9 1.0 1.1 1.2 1.3 1.4 1.5 1.6\n" - "See docs/sphinx/source/zh_CN/user_guide/D-tasks/04-sharpa-inhand.md" - ) - - return tuple(grasp_caches) - - -class SharpaInhandRotationDRProvider(DomainRandomizationProvider): - def validate(self, env: Any, capabilities: DomainRandomizationCapabilities) -> None: - unsupported = validate_common_reset_randomization(env, capabilities) - domain_rand = getattr(env.cfg, "domain_rand", None) - if domain_rand is not None and getattr(domain_rand, "randomize_gravity_direction", False): - if getattr(domain_rand, "randomize_gravity", False): - raise ValueError( - "Use only one Sharpa gravity randomization mode: " - "domain_rand.randomize_gravity_direction or domain_rand.randomize_gravity" - ) - if not capabilities.supports_reset_term(RESET_TERM_GRAVITY): - unsupported = unsupported | frozenset({RESET_TERM_GRAVITY}) - if ( - domain_rand is not None - and domain_rand.force_scale > 0.0 - and not capabilities.supports_interval_term(INTERVAL_TERM_BODY_FORCE) - ): - raise NotImplementedError( - f"{env._backend.backend_type} backend does not support interval body force perturbation" - ) - if domain_rand is not None and domain_rand.randomize_pd_gains: - if not capabilities.supports_reset_term(RESET_TERM_KP): - unsupported = unsupported | frozenset({RESET_TERM_KP}) - if not capabilities.supports_reset_term(RESET_TERM_KD): - unsupported = unsupported | frozenset({RESET_TERM_KD}) - if ( - domain_rand is not None - and domain_rand.randomize_com - and not capabilities.supports_reset_term(RESET_TERM_BODY_IPOS) - ): - unsupported = unsupported | frozenset({RESET_TERM_BODY_IPOS}) - if ( - domain_rand is not None - and domain_rand.randomize_mass - and not capabilities.supports_reset_term(RESET_TERM_BODY_MASS) - ): - unsupported = unsupported | frozenset({RESET_TERM_BODY_MASS}) - if ( - domain_rand is not None - and domain_rand.randomize_friction - and not capabilities.supports_reset_term(RESET_TERM_GEOM_FRICTION) - ): - unsupported = unsupported | frozenset({RESET_TERM_GEOM_FRICTION}) - if unsupported: - names = ", ".join(sorted(unsupported)) - raise NotImplementedError( - f"{env._backend.backend_type} backend does not support reset randomization terms: {names}" - ) - - def build_init_randomization_plan(self, env: Any) -> InitRandomizationPlan | None: - base_size = getattr(env, "_object_geom_base_size", None) - - if base_size is None: - return None - - model_variants = tuple( - ModelVariantSpec( - geom_size_overrides=( - GeomSizeOverride( - geom_name=env.cfg.object_geom_name, - size=tuple(np.asarray(base_size * scale, dtype=np.float64)), - ), - ) - ) - for scale in np.asarray(env.scale_values, dtype=np.float64) - ) - return InitRandomizationPlan( - model_assignments=np.asarray(env.scale_ids, dtype=np.int32).copy(), - model_variants=model_variants, - ) - - def _sample_reset_pd_gains( - self, - env: Any, - num_reset: int, - *, - dtype: np.dtype[Any], - ) -> tuple[np.ndarray, np.ndarray]: - """Sample absolute reset-time PD gains from split-around-1 scale ranges. - - Args: - env: Sharpa rotation env instance. - num_reset: Number of environments being reset. - dtype: Output dtype for the returned gain arrays. - - Returns: - Tuple of absolute ``(p_gain, d_gain)`` arrays with shape - ``(num_reset, env._num_action)``. - """ - p_gain = np.broadcast_to(env._default_p_gain, (num_reset, env._num_action)).astype( - dtype, copy=True - ) - d_gain = np.broadcast_to(env._default_d_gain, (num_reset, env._num_action)).astype( - dtype, copy=True - ) - domain_rand = env.cfg.domain_rand - if domain_rand.randomize_pd_gains: - p_scale = env._sample_pd_scales( - domain_rand.randomize_p_gain_scale_lower, - domain_rand.randomize_p_gain_scale_upper, - shape=(num_reset, env._num_action), - ) - d_scale = env._sample_pd_scales( - domain_rand.randomize_d_gain_scale_lower, - domain_rand.randomize_d_gain_scale_upper, - shape=(num_reset, env._num_action), - ) - p_gain *= p_scale - d_gain *= d_scale - return p_gain, d_gain - - def _build_info_updates( - self, - env: Any, - env_ids: np.ndarray, - hand_qpos: np.ndarray, - object_pos: np.ndarray, - object_quat: np.ndarray, - reset_height_lower: np.ndarray, - reset_height_upper: np.ndarray, - rot_axis: np.ndarray, - p_gain: np.ndarray, - d_gain: np.ndarray, - friction_scale: np.ndarray | None, - randomized_mass: np.ndarray | None, - randomized_com_offset: np.ndarray | None, - gravity: np.ndarray | None, - ) -> dict[str, np.ndarray]: - num_reset = hand_qpos.shape[0] - dtype = get_global_dtype() - - critic_info = env._build_reset_critic_info( - num_reset, - env_ids, - friction_scale=friction_scale, - randomized_mass=randomized_mass, - randomized_com_offset=randomized_com_offset, - gravity=gravity, - ).astype(dtype) - - tactile = np.zeros((num_reset, env._num_tactile), dtype=dtype) - contact_pos = np.zeros((num_reset, env._num_tactile * 3), dtype=dtype) - hand_qpos_f = hand_qpos.astype(dtype) - targets = hand_qpos_f.copy() - object_pos_f = object_pos.astype(dtype) - init_frame = env._build_policy_frame( - dof_pos=hand_qpos_f, - targets=targets, - tactile=tactile, - contact_pos=contact_pos, - add_noise=True, - ) - critic_init_frame = env._build_policy_frame( - dof_pos=hand_qpos_f, - targets=targets, - tactile=tactile, - contact_pos=contact_pos, - add_noise=False, - ) - obs_lag_history = repeat_obs_history(init_frame, env.cfg.obs_history_len).astype(dtype) - critic_obs_lag_history = repeat_obs_history( - critic_init_frame, env.cfg.obs_history_len - ).astype(dtype) - - object_default_pose = np.concatenate( - [object_pos_f, object_quat.astype(dtype)], axis=1 - ).astype(dtype) - object_pos_anchor = env._build_object_pos_anchor(object_pos_f).astype(dtype) - critic_info = env._fill_critic_info( - critic_info=critic_info, - object_pos=object_pos_f, - object_pos_anchor=object_pos_anchor, - ) - - info_updates = { - "current_actions": np.zeros((num_reset, env._num_action), dtype=dtype), - "last_actions": np.zeros((num_reset, env._num_action), dtype=dtype), - "prev_targets": hand_qpos_f.copy(), - "init_pose": hand_qpos_f.copy(), - "prev_hand_pos": hand_qpos_f.copy(), - "prev_object_pos": object_pos.astype(dtype).copy(), - "prev_object_quat": object_quat.astype(dtype).copy(), - "object_default_pose": object_default_pose, - "object_pos_anchor": object_pos_anchor, - "reset_height_lower": reset_height_lower.astype(dtype), - "reset_height_upper": reset_height_upper.astype(dtype), - "rot_axis": rot_axis.astype(dtype), - "p_gain": p_gain, - "d_gain": d_gain, - "critic_info": critic_info, - "obs_lag_history": obs_lag_history, - "critic_obs_lag_history": critic_obs_lag_history, - "proprio_hist": env._update_proprio_history(obs_lag_history), - } - return info_updates - - def build_reset_plan(self, env: Any, env_ids: np.ndarray) -> ResetPlan: - num_reset = len(env_ids) - if num_reset == 0: - return ResetPlan( - env_ids=env_ids, - qpos=np.zeros((0, env.nq), dtype=np.float64), - qvel=np.zeros((0, env.nv), dtype=np.float64), - info_updates={}, - randomization=None, - ) - - friction_scale = env._sample_friction_scale(num_reset) - randomized_mass = env._sample_object_mass(num_reset) - randomized_com_offset = env._sample_object_com_offset(num_reset) - gravity = env._sample_reset_gravity(num_reset) - p_gain, d_gain = self._sample_reset_pd_gains(env, num_reset, dtype=get_global_dtype()) - grasp_cache = cast(tuple[np.ndarray, ...], env._grasp_cache) - sampled_pose = sample_scale_grasp_caches(grasp_cache, env.scale_ids[env_ids]) - - hand_qpos = sampled_pose[:, : env._num_action] - object_pos = sampled_pose[:, env._num_action : env._num_action + 3] - object_quat = sampled_pose[:, env._num_action + 3 : env._num_action + 7] - - rot_axis = np.broadcast_to(env._rot_axis, (num_reset, 3)).copy().astype(np.float64) - - qpos = np.zeros((num_reset, env.nq), dtype=np.float64) - qpos[:, : env._num_action] = hand_qpos - qpos[:, env._obj_pos_slice] = object_pos - qpos[:, env._obj_quat_slice] = object_quat - - qvel = np.zeros((num_reset, env.nv), dtype=np.float64) - - height_range = env.cfg.reset_height_upper - env.cfg.reset_height_lower - reset_height_lower = object_pos[:, 2] - 0.5 * height_range - reset_height_upper = object_pos[:, 2] + 0.5 * height_range - - info_updates = self._build_info_updates( - env, - env_ids=env_ids, - hand_qpos=hand_qpos, - object_pos=object_pos, - object_quat=object_quat, - reset_height_lower=reset_height_lower, - reset_height_upper=reset_height_upper, - rot_axis=rot_axis, - p_gain=p_gain, - d_gain=d_gain, - friction_scale=friction_scale, - randomized_mass=randomized_mass, - randomized_com_offset=randomized_com_offset, - gravity=gravity, - ) - # Match the source task by clearing any cached external object force on reset. - env._random_object_force[env_ids] = 0.0 - env._clear_tactile_history(env_ids) - - return ResetPlan( - env_ids=env_ids, - qpos=qpos, - qvel=qvel, - info_updates=info_updates, - randomization=env._build_reset_randomization( - num_reset, - p_gain=p_gain, - d_gain=d_gain, - friction_scale=friction_scale, - randomized_mass=randomized_mass, - randomized_com_offset=randomized_com_offset, - gravity=gravity, - ), - ) - - def build_reset_observation( - self, - env: Any, - env_ids: np.ndarray, - info_updates: dict[str, Any], - ) -> dict[str, np.ndarray]: - del env_ids - tactile, contact_pos = env._policy_frame_zeros(len(info_updates["prev_targets"])) - return cast( - dict[str, np.ndarray], - env._compute_obs_from_inputs( - info_updates, - dof_pos=np.asarray(info_updates["prev_targets"]), - object_pos=np.asarray(info_updates["prev_object_pos"]), - tactile=tactile, - contact_pos=contact_pos, - ), - ) - - def build_interval_randomization_plan( - self, - env: Any, - step_counter: int, - ) -> IntervalRandomizationPlan | None: - """Build Sharpa object-force perturbations for the upcoming control step. - - Args: - env: Sharpa rotation env instance. - step_counter: Global environment step counter. - - Returns: - Interval randomization plan carrying direct object-force perturbations, - or ``None`` when object-force injection is disabled. - """ - del step_counter - domain_rand = env.cfg.domain_rand - if domain_rand.force_scale <= 0.0: - return None - - decay = float( - np.power( - domain_rand.force_decay, - env.cfg.ctrl_dt / max(domain_rand.force_decay_interval, 1.0e-8), - ) - ) - env._random_object_force *= decay - - random_mask = np.random.rand(env._num_envs) < float(domain_rand.random_force_prob_scalar) - if np.any(random_mask): - object_mass = env._resolve_current_object_mass() - env._random_object_force[random_mask] = ( - np.random.randn(int(np.sum(random_mask)), 3).astype(np.float64) - * object_mass[random_mask, None] - * float(domain_rand.force_scale) - ) - - return IntervalRandomizationPlan( - ops=( - IntervalTermOp( - INTERVAL_TERM_BODY_FORCE, - env._random_object_force[:, None, :].copy(), - body_ids=np.asarray([env._object_body_id], dtype=np.int32), - ), - ), - ) - - -class SharpaInhandRotationEnv(SharpaInhandBaseEnv): - _cfg: SharpaInhandRotationCfg # pyright: ignore[reportIncompatibleVariableOverride] - _reward_cfg: RewardConfig - _MATERIALIZE_ROTATION_GRASP_CACHE = True - _OBS_MODE_ALIASES: dict[str, str] = { - "separated": "separated", - "flattened": "flattened", - } - _CRITIC_BASE_DIM_WITHOUT_OPTIONALS = 8 - - def __init__( - self, - cfg: SharpaInhandRotationCfg, - num_envs: int = 1, - backend_type: str = "motrix", - dr_provider: DomainRandomizationProvider | None = None, - ) -> None: - if cfg.reward_config is None: - raise ValueError("reward_config must be provided via Hydra configuration") - - backend = create_backend( - backend_type, - cfg.scene, - num_envs, - cfg.sim_dt, - base_name=cfg.base_name, - push_body_name=cfg.domain_rand.push_body_name, - add_body_sensors=True, - **env_backend_kwargs(cfg), - ) - super().__init__(cfg, backend, num_envs) - - self._observation_mode = self._resolve_observation_mode(cfg.obs.observation_mode) - expected_critic_info_dim = self._expected_critic_info_dim() - if cfg.critic_info_dim != expected_critic_info_dim: - cfg.critic_info_dim = expected_critic_info_dim - policy_frame_dim = self._policy_frame_dim() - self.obs_buf_lag_history = np.zeros( - (num_envs, cfg.obs_history_len, policy_frame_dim), dtype=self._np_dtype - ) - self.proprio_hist_buf = np.zeros( - (num_envs, cfg.prop_hist_len, policy_frame_dim), dtype=self._np_dtype - ) - self.critic_info_buf = np.zeros((num_envs, expected_critic_info_dim), dtype=self._np_dtype) - self._friction_geom_ids = self._resolve_friction_geom_ids() - self._base_geom_friction = self._resolve_base_geom_friction() - self._base_gravity = self._resolve_base_gravity() - self._object_body_id = self._resolve_object_body_id() - self._base_body_mass = self._resolve_base_body_mass() - self._base_body_ipos = self._resolve_base_body_ipos() - self._default_p_gain, self._default_d_gain = self._load_default_pd_gains() - self._random_object_force = np.zeros((num_envs, 3), dtype=np.float64) - - if cfg.control_config.torque_control: - raise NotImplementedError( - "Sharpa torque_control=True is not implemented with the current position-actuator XML setup. " - "Set env.control_config.torque_control=false. Virtual torques are still computed explicitly for reward terms." - ) - - self._reward_cfg = cfg.reward_config - self._zero_action_test_mode = bool(cfg.zero_action_test_mode) - self._enable_reward_log = True - self._grasp_cache: tuple[np.ndarray, ...] | None = None - - axis = np.asarray(cfg.rot_axis, dtype=self._np_dtype) - axis_norm = np.linalg.norm(axis) - if axis_norm <= 1.0e-8: - raise ValueError("rot_axis must be non-zero") - self._rot_axis = np.asarray(axis / axis_norm, dtype=self._np_dtype) - - provider = dr_provider if dr_provider is not None else SharpaInhandRotationDRProvider() - if self._MATERIALIZE_ROTATION_GRASP_CACHE: - self._grasp_cache = _materialize_grasp_caches(cfg.grasp_cache_path, self.scale_values) - self._init_domain_randomization(provider) - - def apply_action(self, actions: np.ndarray, state: NpEnvState) -> np.ndarray: - actions_np = np.asarray(actions, dtype=self._np_dtype) - if self._zero_action_test_mode: - actions_np = np.zeros_like(actions_np, dtype=self._np_dtype) - return super().apply_action(actions_np, state) - - def _scale_randomization_enabled(self) -> bool: - return self._num_scales > 1 or not np.allclose(self.scale_values, self.scale_values[0]) - - def _expected_critic_info_dim(self) -> int: - priv_info_cfg = self._cfg.priv_info - dim = self._CRITIC_BASE_DIM_WITHOUT_OPTIONALS - if priv_info_cfg.include_friction_scale: - dim += 1 - if priv_info_cfg.include_gravity_direction: - dim += 3 - return dim - - def _resolve_friction_geom_ids(self) -> dict[str, np.ndarray]: - """Resolve MuJoCo geom ids touched by Sharpa friction randomization. - - Args: - None. - - Returns: - Mapping with object, elastomer, and metal collision geom id arrays. - """ - try: - object_geom_id = self._backend.get_geom_id(self._cfg.object_geom_name) - base_body_id = self._backend.get_body_id(self._cfg.base_name) - hand_body_ids = set( - int(body_id) for body_id in self._backend.get_body_subtree_ids(base_body_id) - ) - geom_body_ids = self._backend.get_geom_body_ids() - geom_contype, geom_conaffinity = self._backend.get_geom_contact_masks() - geom_names = self._backend.get_geom_names() - except NotImplementedError: - empty = np.zeros((0,), dtype=np.int32) - return {"object": empty, "elastomer": empty, "metal": empty} - elastomer_ids: list[int] = [] - metal_ids: list[int] = [] - for geom_id, geom_name in enumerate(geom_names): - if int(geom_body_ids[geom_id]) not in hand_body_ids: - continue - if int(geom_contype[geom_id]) == 0 and int(geom_conaffinity[geom_id]) == 0: - continue - if "elastomer" in geom_name: - elastomer_ids.append(geom_id) - else: - metal_ids.append(geom_id) - - if not elastomer_ids: - raise ValueError("No Sharpa elastomer collision geoms found for friction randomization") - if not metal_ids: - raise ValueError("No Sharpa metal collision geoms found for friction randomization") - - return { - "object": np.asarray([object_geom_id], dtype=np.int32), - "elastomer": np.asarray(elastomer_ids, dtype=np.int32), - "metal": np.asarray(metal_ids, dtype=np.int32), - } - - def _resolve_object_body_id(self) -> int: - """Resolve the MuJoCo body id of the manipulated object. - - Args: - None. - - Returns: - The MuJoCo object body id, or -1 when the backend is not MuJoCo. - """ - if self._object_body_ids.size == 0: - return -1 - return int(self._object_body_ids[0]) - - def _resolve_current_object_mass(self) -> np.ndarray: - """Resolve the current object mass used by force perturbation sampling. - - Args: - None. - - Returns: - Array of shape ``(num_envs,)`` with current object masses. - """ - if self.state is not None: - critic_info = np.asarray( - self.state.info.get( - "critic_info", - np.zeros((self._num_envs, self._cfg.critic_info_dim), dtype=self._np_dtype), - ), - dtype=np.float64, - ) - mass = critic_info[:, self._critic_info_layout()["mass"]].reshape(self._num_envs) - if np.all(mass > 0.0): - return mass - if self._base_body_mass is None or self._object_body_id < 0: - raise ValueError("MuJoCo object body-mass cache is unavailable") - return np.full( - (self._num_envs,), - float(self._base_body_mass[self._object_body_id]), - dtype=np.float64, - ) - - def _resolve_base_geom_friction(self) -> np.ndarray | None: - """Cache MuJoCo model friction vectors used as torsional/rolling templates. - - Args: - None. - - Returns: - Full geom friction table, or None when the backend has no geom-friction hook. - """ - try: - return np.asarray(self._backend.get_geom_friction(), dtype=np.float64).copy() - except NotImplementedError: - return None - - def _resolve_base_gravity(self) -> np.ndarray: - """Cache the default gravity vector for privileged-info fallbacks. - - Args: - None. - - Returns: - Gravity vector with shape ``(3,)``. - """ - try: - return np.asarray(self._backend.get_gravity(), dtype=np.float64).copy() - except NotImplementedError: - pass - return np.asarray([0.0, 0.0, -9.81], dtype=np.float64) - - def _resolve_base_body_mass(self) -> np.ndarray | None: - """Cache the MuJoCo body-mass table used as reset randomization baseline. - - Args: - None. - - Returns: - Full body-mass table, or None when the backend is not MuJoCo. - """ - try: - return np.asarray(self._backend.get_body_mass(), dtype=np.float64).copy() - except NotImplementedError: - return None - - def _resolve_base_body_ipos(self) -> np.ndarray | None: - """Cache the MuJoCo inertial-position table used for object COM randomization. - - Args: - None. - - Returns: - Full body inertial-position table, or None when the backend is not MuJoCo. - """ - try: - return np.asarray(self._backend.get_body_ipos(), dtype=np.float64).copy() - except NotImplementedError: - return None - - def _sample_friction_scale(self, batch_size: int) -> np.ndarray | None: - """Sample one friction multiplier per reset environment. - - Args: - batch_size: Number of reset environments. - - Returns: - Shape (batch_size, 1) multipliers, or None when disabled. - """ - domain_rand = self._cfg.domain_rand - if not domain_rand.randomize_friction: - return None - return np.random.uniform( - domain_rand.randomize_friction_scale_lower, - domain_rand.randomize_friction_scale_upper, - size=(batch_size, 1), - ).astype(np.float64) - - def _sample_object_mass(self, batch_size: int) -> np.ndarray | None: - """Sample reset-time object masses. - - Args: - batch_size: Number of reset environments. - - Returns: - Shape (batch_size, 1) absolute object masses, or None when disabled. - """ - domain_rand = self._cfg.domain_rand - if not domain_rand.randomize_mass: - return None - return np.random.uniform( - domain_rand.randomize_mass_lower, - domain_rand.randomize_mass_upper, - size=(batch_size, 1), - ).astype(np.float64) - - def _sample_object_com_offset(self, batch_size: int) -> np.ndarray | None: - """Sample reset-time object COM offsets. - - Args: - batch_size: Number of reset environments. - - Returns: - Shape (batch_size, 3) COM offsets, or None when disabled. - """ - domain_rand = self._cfg.domain_rand - if not domain_rand.randomize_com: - return None - return np.random.uniform( - domain_rand.randomize_com_lower, - domain_rand.randomize_com_upper, - size=(batch_size, 3), - ).astype(np.float64) - - def _friction_profile(self, material: str, base_sliding_friction: float) -> np.ndarray: - """Build one MuJoCo friction vector while preserving XML friction ratios. - - Args: - material: One of object, elastomer, or metal. - base_sliding_friction: Sliding-friction coefficient before random scaling. - - Returns: - Shape (3,) MuJoCo friction vector. - """ - if self._base_geom_friction is None: - raise ValueError("MuJoCo base geom friction cache is unavailable") - - geom_ids = self._friction_geom_ids[material] - template = np.asarray(self._base_geom_friction[int(geom_ids[0])], dtype=np.float64) - sliding = float(template[0]) - if sliding <= 0.0: - raise ValueError(f"{material} base sliding friction must be positive, got {sliding}") - return np.asarray(template / sliding * base_sliding_friction, dtype=np.float64) - - def _sample_reset_gravity(self, batch_size: int) -> np.ndarray | None: - """Sample the exact reset-time gravity vector applied by Sharpa. - - Args: - batch_size: Number of reset environments. - - Returns: - Gravity array with shape ``(batch_size, 3)``, or ``None`` when reset-time - gravity randomization is disabled. - """ - gravity = self._build_gravity_direction_randomization(batch_size) - if gravity is not None: - return gravity - domain_rand = self._cfg.domain_rand - if not getattr(domain_rand, "randomize_gravity", False): - return None - gravity_range = np.asarray(domain_rand.gravity_range, dtype=np.float64) - if gravity_range.shape != (2, 3): - raise ValueError( - f"domain_rand.gravity_range must have shape (2, 3), got {gravity_range.shape}" - ) - low = np.minimum(gravity_range[0], gravity_range[1]) - high = np.maximum(gravity_range[0], gravity_range[1]) - return np.random.uniform(low=low, high=high, size=(batch_size, 3)).astype(np.float64) - - def _resolved_privileged_gravity( - self, - batch_size: int, - gravity: np.ndarray | None, - ) -> np.ndarray: - """Resolve the gravity vector written into privileged info for one reset batch. - - Args: - batch_size: Number of reset environments. - gravity: Sampled reset gravity, if reset-time randomization is active. - - Returns: - Gravity array with shape ``(batch_size, 3)``. - """ - if gravity is not None: - return np.asarray(gravity, dtype=self._np_dtype) - return np.broadcast_to(self._base_gravity, (batch_size, 3)).astype( - self._np_dtype, copy=True - ) - - def _build_friction_randomization( - self, - batch_size: int, - friction_scale: np.ndarray | None, - ) -> np.ndarray | None: - """Build the reset-time MuJoCo geom_friction randomization table. - - Args: - batch_size: Number of reset environments. - friction_scale: Shape (batch_size, 1) sampled friction multipliers. - - Returns: - Shape (batch_size, ngeom, 3) friction table, or None when disabled. - """ - if friction_scale is None: - return None - if self._base_geom_friction is None: - raise ValueError("MuJoCo base geom friction cache is unavailable") - - geom_friction = np.broadcast_to( - self._base_geom_friction, - (batch_size, *self._base_geom_friction.shape), - ).copy() - scale = np.asarray(friction_scale, dtype=np.float64).reshape(batch_size, 1, 1) - domain_rand = self._cfg.domain_rand - material_profiles = { - "object": self._friction_profile("object", domain_rand.object_base_friction), - "elastomer": self._friction_profile("elastomer", domain_rand.elastomer_base_friction), - "metal": self._friction_profile("metal", domain_rand.metal_base_friction), - } - for material, profile in material_profiles.items(): - geom_friction[:, self._friction_geom_ids[material], :] = scale * profile.reshape( - 1, 1, 3 - ) - return geom_friction - - def _build_object_mass_randomization( - self, - batch_size: int, - randomized_mass: np.ndarray | None, - ) -> np.ndarray | None: - """Build the reset-time MuJoCo body_mass table for object mass randomization. - - Args: - batch_size: Number of reset environments. - randomized_mass: Shape (batch_size, 1) sampled absolute object masses. - - Returns: - Shape (batch_size, nbody) mass table, or None when disabled. - """ - if randomized_mass is None: - return None - if self._base_body_mass is None or self._object_body_id < 0: - raise ValueError("MuJoCo base body-mass cache is unavailable") - body_mass = np.broadcast_to( - self._base_body_mass, (batch_size, self._base_body_mass.size) - ).copy() - body_mass[:, self._object_body_id] = np.asarray(randomized_mass, dtype=np.float64).reshape( - batch_size - ) - return body_mass - - def _build_object_com_randomization( - self, - batch_size: int, - randomized_com_offset: np.ndarray | None, - ) -> np.ndarray | None: - """Build the reset-time MuJoCo body_ipos table for object COM randomization. - - Args: - batch_size: Number of reset environments. - randomized_com_offset: Shape (batch_size, 3) sampled COM offsets. - - Returns: - Shape (batch_size, nbody, 3) inertial-position table, or None when disabled. - """ - if randomized_com_offset is None: - return None - if self._base_body_ipos is None or self._object_body_id < 0: - raise ValueError("MuJoCo base body-ipos cache is unavailable") - body_ipos = np.broadcast_to( - self._base_body_ipos, - (batch_size, *self._base_body_ipos.shape), - ).copy() - body_ipos[:, self._object_body_id, :] += np.asarray(randomized_com_offset, dtype=np.float64) - return body_ipos - - def _build_gravity_direction_randomization(self, batch_size: int) -> np.ndarray | None: - """Sample fixed-magnitude gravity vectors with randomized directions. - - Args: - batch_size: Number of reset environments. - - Returns: - Gravity vectors with shape ``(batch_size, 3)``, or None when disabled. - """ - domain_rand = self._cfg.domain_rand - if not getattr(domain_rand, "randomize_gravity_direction", False): - return None - magnitude = float(getattr(domain_rand, "gravity_direction_magnitude", 9.81)) - if magnitude <= 0.0: - raise ValueError(f"gravity_direction_magnitude must be positive, got {magnitude}") - gravity = np.zeros((batch_size, 3), dtype=np.float64) - gravity[:, 2] = -magnitude - random_quat = sample_random_quaternion(batch_size) - # A uniform quaternion and its inverse have the same distribution, so - # rotating gravity samples the same relative gravity directions induced - # by uniformly rotating the global hand/object/task frame. - return np.asarray(np_quat_apply(random_quat, gravity), dtype=np.float64) - - def _build_reset_randomization( - self, - batch_size: int, - *, - p_gain: np.ndarray | None, - d_gain: np.ndarray | None, - friction_scale: np.ndarray | None, - randomized_mass: np.ndarray | None, - randomized_com_offset: np.ndarray | None, - gravity: np.ndarray | None, - ) -> ResetRandomizationPayload | None: - """Build reset-randomization payloads owned by the Sharpa rotation env. - - Args: - batch_size: Number of reset environments. - friction_scale: Shape (batch_size, 1) sampled friction multipliers. - - Returns: - Reset-randomization payload, or None when no backend randomization is requested. - """ - payload = build_common_reset_randomization(self, batch_size) - if self._cfg.domain_rand.randomize_pd_gains: - if payload is None: - payload = ResetRandomizationPayload() - payload.kp = np.asarray(p_gain, dtype=np.float64) - payload.kd = np.asarray(d_gain, dtype=np.float64) - body_mass = self._build_object_mass_randomization(batch_size, randomized_mass) - body_ipos = self._build_object_com_randomization(batch_size, randomized_com_offset) - geom_friction = self._build_friction_randomization(batch_size, friction_scale) - if body_mass is not None: - if payload is None: - payload = ResetRandomizationPayload() - payload.body_mass = body_mass - if body_ipos is not None: - if payload is None: - payload = ResetRandomizationPayload() - payload.body_ipos = body_ipos - if geom_friction is not None: - if payload is None: - payload = ResetRandomizationPayload() - payload.geom_friction = geom_friction - if gravity is not None: - if payload is None: - payload = ResetRandomizationPayload() - payload.gravity = gravity - return payload - - def _critic_info_layout(self) -> dict[str, slice]: - """Describe the flat critic_info channel layout. - - Args: - None. - - Returns: - Mapping from logical field names to channel slices. - """ - offset = 0 - layout = {"object_pos_delta": slice(offset, offset + 3)} - offset += 3 - priv_info_cfg = self._cfg.priv_info - if priv_info_cfg.include_friction_scale: - layout["friction"] = slice(offset, offset + 1) - offset += 1 - layout["mass"] = slice(offset, offset + 1) - offset += 1 - layout["com"] = slice(offset, offset + 3) - offset += 3 - layout["scale"] = slice(offset, offset + 1) - offset += 1 - if priv_info_cfg.include_gravity_direction: - layout["gravity"] = slice(offset, offset + 3) - return layout - - def _assign_critic_info_field( - self, - critic_info: np.ndarray, - field_name: str, - values: np.ndarray, - ) -> None: - """Assign one critic_info field according to the declared layout. - - Args: - critic_info: Critic-info buffer to update in place. - field_name: Logical field name from the layout. - values: Batch-major values for the field. - - Returns: - None. The critic_info array is updated in place. - """ - field_slice = self._critic_info_layout().get(field_name) - if field_slice is None: - return - field_values = np.asarray(values, dtype=self._np_dtype).reshape(critic_info.shape[0], -1) - critic_info[:, field_slice] = field_values - - def _build_reset_critic_info( - self, - batch_size: int, - env_ids: np.ndarray, - *, - friction_scale: np.ndarray | None, - randomized_mass: np.ndarray | None, - randomized_com_offset: np.ndarray | None, - gravity: np.ndarray | None, - ) -> np.ndarray: - """Build reset-time critic_info for all randomized object properties. - - Args: - batch_size: Number of reset environments. - env_ids: Global environment ids for this reset batch. - - Returns: - Critic-info tensor with shape (batch_size, critic_info_dim). - """ - critic_info = np.zeros((batch_size, self._cfg.critic_info_dim), dtype=self._np_dtype) - - priv_info_cfg = self._cfg.priv_info - if priv_info_cfg.include_friction_scale: - self._assign_critic_info_field( - critic_info, - "friction", - ( - friction_scale - if friction_scale is not None - else np.ones((batch_size, 1), dtype=np.float64) - ), - ) - if randomized_mass is not None: - self._assign_critic_info_field( - critic_info, - "mass", - randomized_mass, - ) - if randomized_com_offset is not None: - self._assign_critic_info_field( - critic_info, - "com", - randomized_com_offset, - ) - if self._scale_randomization_enabled(): - self._assign_critic_info_field( - critic_info, - "scale", - self.scale_values[self.scale_ids[env_ids]].reshape(batch_size, 1), - ) - if priv_info_cfg.include_gravity_direction: - self._assign_critic_info_field( - critic_info, - "gravity", - self._resolved_privileged_gravity(batch_size, gravity), - ) - return critic_info - - def _policy_frame_dim(self) -> int: - obs_cfg = self._cfg.obs - dim = self._num_action + self._num_action - if obs_cfg.enable_tactile: - dim += self._num_tactile - if obs_cfg.enable_contact_pos: - dim += self._num_tactile * 3 - return dim - - def _policy_frame_zeros(self, batch_size: int) -> tuple[np.ndarray, np.ndarray]: - """Build zero-filled optional policy inputs for reset observations. - - Args: - batch_size: Number of environments in the batch. - - Returns: - Tuple of tactile and contact-position arrays sized for the current config. - """ - obs_cfg = self._cfg.obs - tactile_dim = self._num_tactile if obs_cfg.enable_tactile else 0 - contact_pos_dim = self._num_tactile * 3 if obs_cfg.enable_contact_pos else 0 - return ( - np.zeros((batch_size, tactile_dim), dtype=self._np_dtype), - np.zeros((batch_size, contact_pos_dim), dtype=self._np_dtype), - ) - - def _fixed_default_object_pose(self, batch_size: int) -> np.ndarray: - """Build the fixed default object pose from the backend init state. - - Args: - batch_size: Number of environments in the batch. - - Returns: - Array with shape ``(batch_size, 7)`` containing the default object - position and quaternion loaded from the backend/model init qpos. - """ - object_pos = np.broadcast_to( - np.asarray(self._init_qpos[self._obj_pos_slice], dtype=self._np_dtype), - (batch_size, 3), - ).copy() - object_quat = np.broadcast_to( - np.asarray(self._init_qpos[self._obj_quat_slice], dtype=self._np_dtype), - (batch_size, 4), - ).copy() - return np.concatenate([object_pos, object_quat], axis=1).astype(self._np_dtype) - - def _build_object_pos_anchor( - self, - object_pos: np.ndarray, - ) -> np.ndarray: - """Choose the reward/privileged-info object-position anchor. - - Args: - object_pos: Reset-time object positions with shape ``(batch_size, 3)``. - - Returns: - Array with shape ``(batch_size, 3)``. When the config flag is set, - the fixed XML/default object position is used; otherwise the sampled - reset position is returned to preserve the legacy UniLab behavior. - """ - object_pos = np.asarray(object_pos, dtype=self._np_dtype) - if self._cfg.use_default_object_pose_for_object_pos_anchor: - return self._fixed_default_object_pose(object_pos.shape[0])[:, 0:3] - return object_pos.copy() - - def _resolve_object_pos_anchor( - self, - info: dict[str, Any], - batch_size: int, - ) -> np.ndarray: - """Resolve the cached object-position anchor for one runtime batch. - - Args: - info: Mutable env-state info dictionary. - batch_size: Number of environments in the current batch. - - Returns: - Array with shape ``(batch_size, 3)`` used by the object-position - reward and privileged-info delta channels. - """ - cached_anchor = info.get("object_pos_anchor") - if cached_anchor is not None: - return np.asarray(cached_anchor, dtype=self._np_dtype) - - object_default_pose = info.get("object_default_pose") - if object_default_pose is not None: - return np.asarray(object_default_pose, dtype=self._np_dtype)[:, 0:3] - - if self._cfg.use_default_object_pose_for_object_pos_anchor: - return self._fixed_default_object_pose(batch_size)[:, 0:3] - return np.zeros((batch_size, 3), dtype=self._np_dtype) - - def _policy_frame_parts( - self, - dof_norm: np.ndarray, - targets: np.ndarray, - tactile: np.ndarray, - contact_pos: np.ndarray, - ) -> list[np.ndarray]: - """Collect policy-frame components according to the configured observation layout. - - Args: - dof_norm: Normalized hand joint positions. - targets: Hand joint targets. - tactile: Tactile features. - contact_pos: Contact-position features. - - Returns: - Ordered list of arrays that should be concatenated into the policy frame. - """ - parts = [dof_norm, targets] - obs_cfg = self._cfg.obs - if obs_cfg.enable_tactile: - parts.append(np.asarray(tactile, dtype=self._np_dtype)) - if obs_cfg.enable_contact_pos: - parts.append(np.asarray(contact_pos, dtype=self._np_dtype)) - return parts - - def _fill_critic_info( - self, - critic_info: np.ndarray, - object_pos: np.ndarray, - object_pos_anchor: np.ndarray, - ) -> np.ndarray: - """Populate privileged channels that are derived from runtime state. - - Args: - critic_info: Critic info buffer to update in place. - object_pos: Current object positions with shape (batch, 3). - object_pos_anchor: Object-position anchor with shape (batch, 3). - - Returns: - Updated critic info array with shape (batch, critic_info_dim). - """ - self._assign_critic_info_field( - critic_info, - "object_pos_delta", - object_pos - object_pos_anchor, - ) - return critic_info - - @classmethod - def _resolve_observation_mode(cls, observation_mode: str) -> str: - normalized_mode = str(observation_mode).strip().lower() - resolved_mode = cls._OBS_MODE_ALIASES.get(normalized_mode) - if resolved_mode is None: - supported_modes = "', '".join(sorted(cls._OBS_MODE_ALIASES)) - raise ValueError( - f"observation_mode must be one of '{supported_modes}', got {observation_mode!r}" - ) - return resolved_mode - - def _build_policy_frame( - self, - dof_pos: np.ndarray, - targets: np.ndarray, - tactile: np.ndarray, - contact_pos: np.ndarray, - *, - add_noise: bool = True, - ) -> np.ndarray: - dof_pos_f = np.asarray(dof_pos, dtype=self._np_dtype) - targets_f = np.asarray(targets, dtype=self._np_dtype) - - dof_norm = self._normalize_joint_pos(dof_pos_f) - joint_noise_scale = float(self._cfg.domain_rand.joint_noise_scale) - if add_noise and joint_noise_scale > 0.0: - dof_norm += ( - np.random.uniform(-1.0, 1.0, size=dof_norm.shape).astype(self._np_dtype) - * joint_noise_scale - ) - - frame = np.asarray( - np.concatenate( - self._policy_frame_parts( - dof_norm=dof_norm, - targets=targets_f, - tactile=tactile, - contact_pos=contact_pos, - ), - axis=1, - ), - dtype=self._np_dtype, - ) - return self._clip_observation_values(frame) - - def _clip_observation_values(self, values: np.ndarray) -> np.ndarray: - """Clamp observation tensors to a configurable absolute value bound. - - Args: - values: Observation-like array to clip. - - Returns: - Clipped array with the same shape. Non-positive ``clip_obs`` disables the - clamp so callers can opt out when needed. - """ - clip_max = float(getattr(self._cfg, "clip_obs", 5.0)) - values = np.asarray(values, dtype=self._np_dtype) - if clip_max <= 0.0: - return values - return np.asarray(np.clip(values, -clip_max, clip_max), dtype=self._np_dtype) - - @property - def obs_groups_spec(self) -> dict[str, int]: - policy_obs_dim = self._cfg.obs_lag_steps * self._policy_frame_dim() - if self._observation_mode == "flattened": - return {"obs": policy_obs_dim + self._cfg.critic_info_dim} - return {"obs": policy_obs_dim, "critic": policy_obs_dim + self._cfg.critic_info_dim} - - def _build_critic_info( - self, - info: dict[str, Any], - batch_size: int, - object_pos: np.ndarray, - ) -> np.ndarray: - """Build privileged critic info for the current batch. - - Args: - info: Mutable state info dictionary carrying reset-time caches. - batch_size: Number of environments in the current batch. - object_pos: Current object positions with shape (batch, 3). - - Returns: - Privileged info array with shape (batch, critic_info_dim). - """ - critic_info = np.asarray( - info.get( - "critic_info", - np.zeros((batch_size, self._cfg.critic_info_dim), dtype=self._np_dtype), - ), - dtype=self._np_dtype, - ) - - object_pos_anchor = self._resolve_object_pos_anchor(info, batch_size) - critic_info = self._fill_critic_info( - critic_info=critic_info, - object_pos=object_pos, - object_pos_anchor=object_pos_anchor, - ) - - info["critic_info"] = critic_info - return critic_info - - def _pack_observations( - self, - policy_obs: np.ndarray, - critic_info: np.ndarray, - critic_policy_obs: np.ndarray | None = None, - ) -> dict[str, np.ndarray]: - """Pack actor and privileged info into the env observation groups. - - Args: - policy_obs: Actor observation tensor with shape (batch, actor_dim). - critic_info: Privileged tensor with shape (batch, critic_info_dim). - - Returns: - Observation groups that satisfy the UniLab env contract. - """ - if self._observation_mode == "flattened": - flattened_obs = self._clip_observation_values( - np.concatenate([policy_obs, critic_info], axis=1).astype(self._np_dtype) - ) - return {"obs": flattened_obs} - - if critic_policy_obs is None: - critic_policy_obs = policy_obs - return { - "obs": self._clip_observation_values(policy_obs), - "critic": self._clip_observation_values( - np.concatenate([critic_policy_obs, critic_info], axis=1).astype(self._np_dtype) - ), - } - - def _compute_obs_from_inputs( - self, - info: dict[str, Any], - dof_pos: np.ndarray, - object_pos: np.ndarray, - tactile: np.ndarray, - contact_pos: np.ndarray, - ) -> dict[str, np.ndarray]: - targets = np.asarray(info.get("prev_targets", dof_pos), dtype=self._np_dtype) - frame = self._build_policy_frame( - dof_pos=dof_pos, - targets=targets, - tactile=tactile, - contact_pos=contact_pos, - add_noise=True, - ) - critic_frame = self._build_policy_frame( - dof_pos=dof_pos, - targets=targets, - tactile=tactile, - contact_pos=contact_pos, - add_noise=False, - ) - batch_size = int(frame.shape[0]) - - history = info.get("obs_lag_history") - critic_history = info.get("critic_obs_lag_history") - if history is None: - history = repeat_obs_history(frame, self._cfg.obs_history_len).astype(self._np_dtype) - else: - history = np.asarray(history, dtype=self._np_dtype) - history[:, :-1] = history[:, 1:] - history[:, -1] = frame - if critic_history is None: - critic_history = repeat_obs_history(critic_frame, self._cfg.obs_history_len).astype( - self._np_dtype - ) - else: - critic_history = np.asarray(critic_history, dtype=self._np_dtype) - critic_history[:, :-1] = critic_history[:, 1:] - critic_history[:, -1] = critic_frame - - info["obs_lag_history"] = history - info["critic_obs_lag_history"] = critic_history - info["proprio_hist"] = self._update_proprio_history(history) - - obs = np.asarray( - history[:, -self._cfg.obs_lag_steps :].reshape(batch_size, -1), - dtype=self._np_dtype, - ) - critic_obs = np.asarray( - critic_history[:, -self._cfg.obs_lag_steps :].reshape(batch_size, -1), - dtype=self._np_dtype, - ) - critic_info = self._build_critic_info(info, batch_size=batch_size, object_pos=object_pos) - return self._pack_observations(obs, critic_info, critic_policy_obs=critic_obs) - - def _compute_reward( - self, - info: dict[str, Any], - dof_pos: np.ndarray, - dof_vel: np.ndarray, - object_pos: np.ndarray, - object_linvel: np.ndarray, - object_angvel: np.ndarray, - torques: np.ndarray, - ) -> np.ndarray: - rot_axis = np.asarray( - info.get("rot_axis", np.broadcast_to(self._rot_axis, (self._num_envs, 3))), - dtype=self._np_dtype, - ) - rotate_reward = np.clip( - np.sum(object_angvel * rot_axis, axis=1), - self._reward_cfg.angvel_clip_min, - self._reward_cfg.angvel_clip_max, - ) - object_linvel_penalty = np.sum(np.abs(object_linvel), axis=1) - pos_diff_penalty = np.sum(np.square(dof_pos - self.default_angles), axis=1) - torque_penalty = np.sum(np.square(torques), axis=1) - work_penalty = np.square(np.sum(torques * dof_vel, axis=1)) - - object_pos_anchor = self._resolve_object_pos_anchor(info, object_pos.shape[0]) - object_pos_reward = 1.0 / (np.linalg.norm(object_pos - object_pos_anchor, axis=1) + 0.001) - - reward_terms: dict[str, np.ndarray] = { - "rotate": np.asarray(rotate_reward, dtype=self._np_dtype), - "obj_linvel": np.asarray(object_linvel_penalty, dtype=self._np_dtype), - "pose_diff": np.asarray(pos_diff_penalty, dtype=self._np_dtype), - "torque": np.asarray(torque_penalty, dtype=self._np_dtype), - "work": np.asarray(work_penalty, dtype=self._np_dtype), - "object_pos": np.asarray(object_pos_reward, dtype=self._np_dtype), - } - - reward = np.zeros((self._num_envs,), dtype=self._np_dtype) - step_count = info.get("steps", np.zeros((self._num_envs,), dtype=np.uint32)) - should_log = self._enable_reward_log and (int(step_count[0]) % 4 == 0) - log = {} if should_log else info.get("log", {}) - - for name, scale in self._reward_cfg.scales.items(): - if scale == 0.0 or name not in reward_terms: - continue - weighted = reward_terms[name] * scale - reward += weighted - if should_log: - log[f"reward/{name}"] = float(np.mean(weighted)) - - if should_log: - log["reward/total"] = float(np.mean(reward)) - info["log"] = log - - return np.asarray(reward, dtype=self._np_dtype) * self._cfg.ctrl_dt - - def update_state(self, state: NpEnvState) -> NpEnvState: - dof_pos = self.get_hand_dof_pos() - dof_vel = self.get_hand_dof_vel() - object_pos = self.get_object_pos() - object_quat = self.get_object_quat() - - prev_object_pos = np.asarray( - state.info.get("prev_object_pos", object_pos), dtype=self._np_dtype - ) - prev_object_quat = np.asarray( - state.info.get("prev_object_quat", object_quat), dtype=self._np_dtype - ) - - object_linvel = (object_pos - prev_object_pos) / self._cfg.ctrl_dt - object_angvel = ( - np_quat_to_axis_angle(np_quat_mul(object_quat, np_quat_conjugate(prev_object_quat))) - / self._cfg.ctrl_dt - ) - - targets = np.asarray( - state.info.get( - "prev_targets", - np.broadcast_to(self.default_angles, (self._num_envs, self._num_action)).copy(), - ), - dtype=self._np_dtype, - ) - p_gain, d_gain = self._resolve_pd_gains(state.info) - - # Explicit virtual torque used for reward parity with source Sharpa formulation. - virtual_torques = np.asarray( - p_gain * (targets - dof_pos) - d_gain * dof_vel, - dtype=self._np_dtype, - ) - - tactile = self._compute_tactile_observation() - contact_pos = self._compute_contact_positions(tactile) - - reward = self._compute_reward( - state.info, - dof_pos=dof_pos, - dof_vel=dof_vel, - object_pos=object_pos, - object_linvel=object_linvel, - object_angvel=object_angvel, - torques=virtual_torques, - ) - - reset_height_lower = np.asarray( - state.info.get( - "reset_height_lower", - np.full((self._num_envs,), self._cfg.reset_height_lower, dtype=self._np_dtype), - ), - dtype=self._np_dtype, - ) - reset_height_upper = np.asarray( - state.info.get( - "reset_height_upper", - np.full((self._num_envs,), self._cfg.reset_height_upper, dtype=self._np_dtype), - ), - dtype=self._np_dtype, - ) - terminated = (object_pos[:, 2] > reset_height_upper) | ( - object_pos[:, 2] < reset_height_lower - ) - - obs = self._compute_obs_from_inputs( - state.info, - dof_pos=dof_pos, - object_pos=object_pos, - tactile=tactile, - contact_pos=contact_pos, - ) - - state.info["prev_hand_pos"] = dof_pos.copy() - state.info["hand_dof_vel"] = dof_vel.copy() - state.info["prev_object_pos"] = object_pos.copy() - state.info["prev_object_quat"] = object_quat.copy() - state.info["torques"] = virtual_torques - state.info["virtual_torques"] = virtual_torques.copy() - state.info["object_linvel"] = object_linvel - state.info["object_angvel"] = object_angvel - - return state.replace( - obs=obs, - reward=reward, - terminated=np.asarray(terminated, dtype=bool), - ) - - -SharpaWaveRewardConfig = RewardConfig -SharpaWaveRotationCfg = SharpaInhandRotationCfg - -_SHARPA_ROTATION_COMPAT_FACTORY = adapt_legacy_factory( - SharpaInhandRotationEnv, - task_family="Sharpa", - reason=( - "tactile/contact latency, object variants, and grasp-cache state remain " - "task-owned until formal Manager-Based capabilities exist" - ), -) -for _backend_type in ("mujoco", "motrix", "drake"): - registry.register_env("SharpaInhandRotation", _SHARPA_ROTATION_COMPAT_FACTORY, _backend_type) diff --git a/src/unilab/tasks/migration_matrix.py b/src/unilab/tasks/migration_matrix.py index ef0af5569..fbd59793f 100644 --- a/src/unilab/tasks/migration_matrix.py +++ b/src/unilab/tasks/migration_matrix.py @@ -54,13 +54,6 @@ class TaskMigrationRecord: } ) -_CUSTOM_COMPAT_TASKS = frozenset( - { - "SharpaInhandRotation", - "SharpaInhandRotationGrasp", - } -) - _MOTION_CORE_TASKS = frozenset( { "G1MotionTracking", @@ -93,12 +86,7 @@ class TaskMigrationRecord: ) PRODUCTION_TASK_NAMES = frozenset( - _MBA_TASKS - | _ROUGH_TASKS - | _G1_LOCOMOTION_TASKS - | _CUSTOM_COMPAT_TASKS - | _MOTION_CORE_TASKS - | _MOTION_TASKS + _MBA_TASKS | _ROUGH_TASKS | _G1_LOCOMOTION_TASKS | _MOTION_CORE_TASKS | _MOTION_TASKS ) @@ -142,15 +130,6 @@ def migration_record(task_name: str) -> TaskMigrationRecord: "Hydra owners materialize biped gait, sensor, command, and penalty-curriculum manager terms on the canonical runtime.", "Keep the manager contract and regression evidence current.", ) - if task_name in _CUSTOM_COMPAT_TASKS: - return TaskMigrationRecord( - task_name, - "sharpa", - "Adapted", - "compatibility", - "Custom tactile/contact/cache behavior is retained behind one frozen adapter.", - "Keep Hydra/Registry ownership single; migrate only when the formal capability exists.", - ) if task_name in _MOTION_CORE_TASKS: return TaskMigrationRecord( task_name, diff --git a/src/unilab/training/__init__.py b/src/unilab/training/__init__.py index 5908aad06..7dda18b39 100644 --- a/src/unilab/training/__init__.py +++ b/src/unilab/training/__init__.py @@ -20,11 +20,9 @@ algo_config_dict, apply_env_nan_guard, build_run_dir_name, - format_hora_stage2_checkpoint_error, format_play_checkpoint_error, get_log_root, parse_checkpoint_path, - resolve_hora_stage2_checkpoint_path, resolve_nan_guard_cfg, should_run_playback, ) @@ -72,9 +70,7 @@ "apply_configured_training_seed", "apply_training_seed", "derive_worker_seed", - "format_hora_stage2_checkpoint_error", "resolve_appo_checkpoint_path", - "resolve_hora_stage2_checkpoint_path", "resolve_offpolicy_checkpoint_path", "resolve_training_seed", "setup_logger", diff --git a/src/unilab/training/hora_distill_config.py b/src/unilab/training/hora_distill_config.py deleted file mode 100644 index 5bb6dbd2a..000000000 --- a/src/unilab/training/hora_distill_config.py +++ /dev/null @@ -1,258 +0,0 @@ -"""HORA distillation config and teacher-owner resolution helpers.""" - -from __future__ import annotations - -import re -from pathlib import Path -from typing import Any, cast - -from hydra import compose, initialize_config_dir -from hydra.core.global_hydra import GlobalHydra -from omegaconf import DictConfig, OmegaConf - -from unilab.utils.checkpoint import resolve_task_checkpoint_path - -_PACKAGE_CONF_ROOT = Path(__file__).resolve().parents[1] / "conf" - -# Teacher owner configs are Hydra-composed from their family config tree. -# SAC teachers live in their own per-algo tree; there is no `algo` group anymore. -_TEACHER_TREE_BY_FAMILY = {"sac": "sac"} - -# Teacher -> student `algo.model` mappings expressed in YAML; see -# hora_distill/student_model/*.yaml in the packaged conf tree. The mapping files -# interpolate against the composed teacher owner config mounted at -# `teacher_owner`. -_STUDENT_MODEL_MAPPING_DIR = Path("hora_distill") / "student_model" - - -def _conf_root(root_dir: str | Path | None) -> Path: - """Root of the conf tree: ``root_dir/"conf"`` when given, else packaged conf.""" - return Path(root_dir) / "conf" if root_dir is not None else _PACKAGE_CONF_ROOT - - -def _logs_root(root_dir: str | Path | None) -> Path: - """Root for run artifacts: ``root_dir`` when given, else the caller's CWD.""" - return Path(root_dir) if root_dir is not None else Path.cwd() - - -def _load_yaml_config(path: Path) -> DictConfig: - loaded = OmegaConf.load(path) - if not isinstance(loaded, DictConfig): - raise TypeError(f"Expected DictConfig from {path}, got {type(loaded)!r}") - return loaded - - -def _sanitize_path_token(value: str, *, fallback: str) -> str: - sanitized = re.sub(r"[^A-Za-z0-9._-]+", "-", str(value)).strip("-._") - return sanitized or fallback - - -def load_teacher_owner_config( - algo_family: str, - task: str, - *, - root_dir: str | Path | None = None, -) -> DictConfig: - """Compose a HORA teacher owner config with standard Hydra semantics. - - Uses ``initialize_config_dir`` + ``compose`` (same pattern as - ``scripts/audit_sim2sim_contracts.py``), so package directives, nested - ``defaults`` lists, and interpolations in the teacher tree all resolve. - """ - algo_family = str(algo_family) - conf_dir = _conf_root(root_dir) / _TEACHER_TREE_BY_FAMILY.get(algo_family, algo_family) - overrides = [f"task={task}"] - GlobalHydra.instance().clear() - with initialize_config_dir(config_dir=str(conf_dir.absolute()), version_base="1.3"): - return compose("config", overrides=overrides) - - -def get_teacher_owner_spec(cfg: DictConfig) -> tuple[str | None, str | None]: - """Resolve the teacher algo family and task owner from distillation config.""" - algo_family = OmegaConf.select(cfg, "teacher.algo_family") - task = OmegaConf.select(cfg, "teacher.task") - if algo_family in (None, "") or task in (None, ""): - return None, None - return str(algo_family), str(task) - - -def _teacher_contract_mapping_name( - teacher_algo_family: str, - teacher_task: str, - teacher_cfg: DictConfig, -) -> str: - """Validate the teacher owner contract and select the YAML mapping name.""" - if teacher_algo_family == "sac": - runtime_impl = OmegaConf.select(teacher_cfg, "algo.runtime_impl") - if runtime_impl != "hora_sac": - raise ValueError( - "HORA distillation SAC teacher owner must select runtime_impl='hora_sac'. " - f"Got task={teacher_task} runtime_impl={runtime_impl!r}." - ) - return "hora_sac" - - actor_class_name = str(OmegaConf.select(teacher_cfg, "algo.actor.class_name") or "") - if "HoraActorModel" not in actor_class_name: - raise ValueError( - "HORA distillation teacher owner must resolve to HoraActorModel. " - f"Got algo_family={teacher_algo_family} task={teacher_task} " - f"actor.class_name={actor_class_name!r}." - ) - return "hora_actor" - - -def _student_model_defaults( - mapping_name: str, - teacher_cfg: DictConfig, - *, - root: str | Path | None = None, -) -> dict[str, Any]: - """Resolve the YAML-expressed teacher -> student model mapping. - - The mapping YAML is merged next to the composed teacher owner config - (mounted at ``teacher_owner``) so its interpolations resolve against the - Hydra-composed teacher hyperparameters. The mapping file owns all fallback - defaults; this function only mounts and resolves it. - """ - mapping_path = _conf_root(root) / _STUDENT_MODEL_MAPPING_DIR / f"{mapping_name}.yaml" - merged = OmegaConf.merge( - {"teacher_owner": teacher_cfg}, - _load_yaml_config(mapping_path), - ) - model_cfg = OmegaConf.to_container(OmegaConf.select(merged, "model"), resolve=True) - if not isinstance(model_cfg, dict): - raise TypeError( - f"Expected mapping 'model' dict from {mapping_path}, got {type(model_cfg)!r}" - ) - model_cfg = cast(dict[str, Any], model_cfg) - distribution_cfg = model_cfg.get("distribution_cfg") - if isinstance(distribution_cfg, dict): - # The student re-binds its own distribution class; only the teacher's - # distribution hyperparameters carry over, not its class binding. - distribution_cfg.pop("class_name", None) - return model_cfg - - -def teacher_default_cfg( - cfg: DictConfig, - *, - root_dir: str | Path | None = None, -) -> DictConfig: - """Build HORA student defaults from the selected teacher owner YAML.""" - teacher_algo_family, teacher_task = get_teacher_owner_spec(cfg) - if teacher_algo_family is None or teacher_task is None: - return OmegaConf.create() - - teacher_cfg = load_teacher_owner_config( - teacher_algo_family, - teacher_task, - root_dir=root_dir, - ) - mapping_name = _teacher_contract_mapping_name( - teacher_algo_family, - teacher_task, - teacher_cfg, - ) - model_cfg = _student_model_defaults(mapping_name, teacher_cfg, root=root_dir) - return OmegaConf.create( - { - "training": OmegaConf.select(teacher_cfg, "training"), - "reward": OmegaConf.select(teacher_cfg, "reward"), - "env": OmegaConf.select(teacher_cfg, "env"), - "algo": {"model": model_cfg}, - } - ) - - -def apply_teacher_defaults( - cfg: DictConfig, - *, - root_dir: str | Path | None = None, -) -> DictConfig: - """Merge teacher-owner defaults under the user distillation config.""" - return cast(DictConfig, OmegaConf.merge(teacher_default_cfg(cfg, root_dir=root_dir), cfg)) - - -def resolved_distill_runtime_cfg(cfg: DictConfig) -> DictConfig: - """Return checkpoint runtime fields needed to rebuild the student model. - - Stage-2 checkpoints intentionally do not persist owner runtime settings such - as env, reward, or domain randomization. Replay should use the currently - composed owner config for those fields. - """ - model_cfg = OmegaConf.select(cfg, "algo.model") - return OmegaConf.create( - { - "algo": { - "model": ( - OmegaConf.to_container(model_cfg, resolve=True) if model_cfg is not None else {} - ) - }, - } - ) - - -def teacher_run_metadata( - cfg: DictConfig, - *, - teacher_algo_family: str, - teacher_checkpoint: Path, - root_dir: str | Path | None = None, -) -> dict[str, Any]: - """Build explicit teacher provenance metadata for distillation outputs.""" - teacher_task = OmegaConf.select(cfg, "teacher.task") - root = _logs_root(root_dir).resolve() - checkpoint_path = teacher_checkpoint.resolve() - try: - checkpoint_display = str(checkpoint_path.relative_to(root)) - except ValueError: - checkpoint_display = str(checkpoint_path) - - checkpoint_name = checkpoint_path.name - return { - "algo_family": str(teacher_algo_family), - "task": None if teacher_task in (None, "") else str(teacher_task), - "checkpoint_path": checkpoint_display, - "checkpoint_name": checkpoint_name, - "checkpoint_stem": checkpoint_path.stem, - "run_name": checkpoint_path.parent.name, - "run_slug": f"teacher-{_sanitize_path_token(teacher_algo_family, fallback='teacher')}", - } - - -def resolve_teacher_checkpoint_path( - cfg: DictConfig, - *, - root_dir: str | Path | None = None, -) -> tuple[Path | None, Path | None]: - """Resolve the selected HORA teacher checkpoint through owner metadata.""" - teacher_algo_family, teacher_task = get_teacher_owner_spec(cfg) - if teacher_algo_family is None or teacher_task is None: - return None, None - - root = _logs_root(root_dir) - teacher_cfg = load_teacher_owner_config( - teacher_algo_family, - teacher_task, - root_dir=root_dir, - ) - teacher_task_name = OmegaConf.select(teacher_cfg, "training.task_name") - teacher_algo_log_name = OmegaConf.select(teacher_cfg, "algo.algo_log_name") - if teacher_task_name in (None, "") or teacher_algo_log_name in (None, ""): - raise ValueError( - "Teacher owner config must define training.task_name and algo.algo_log_name. " - f"Got algo_family={teacher_algo_family} task={teacher_task}." - ) - - selected_checkpoint = OmegaConf.select(cfg, "algo.checkpoint", default=-1) - return resolve_task_checkpoint_path( - root, - task_name=str(teacher_task_name), - load_run=str(OmegaConf.select(cfg, "algo.load_run", default="-1")), - algo_log_name=str(teacher_algo_log_name), - checkpoint=( - str(selected_checkpoint) if selected_checkpoint not in (None, "", -1, "-1") else None - ), - suffix=".pt", - log_root=OmegaConf.select(cfg, "training.log_root"), - ) diff --git a/src/unilab/training/run.py b/src/unilab/training/run.py index bc6eefc44..fa1f35799 100644 --- a/src/unilab/training/run.py +++ b/src/unilab/training/run.py @@ -154,81 +154,3 @@ def parse_checkpoint_path( suffix=suffix, log_root=OmegaConf.select(cfg, "training.log_root"), ) - - -def resolve_hora_stage2_checkpoint_path( - cfg: DictConfig, - *, - root_dir: str | Path, -) -> tuple[Path | None, Path | None]: - """Resolve the HORA stage-2 distillation checkpoint selected by ``cfg.algo``.""" - task_log_root = get_log_root(root_dir, cfg) / str(OmegaConf.select(cfg, "training.task_name")) - load_run = str(OmegaConf.select(cfg, "algo.load_run", default="-1")) - selected_checkpoint = OmegaConf.select(cfg, "algo.checkpoint", default=-1) - - run_dir: Path | None - if load_run == "-1": - run_dir = get_latest_run(task_log_root) - else: - candidate = Path(load_run) - if not candidate.exists(): - candidate = task_log_root / load_run - if candidate.is_file(): - return candidate, candidate.parent - run_dir = candidate if candidate.is_dir() else None - - if run_dir is None: - return None, None - - if selected_checkpoint not in (None, "", -1, "-1"): - checkpoint_name = ( - f"hora_stage2_{selected_checkpoint}.pt" - if str(selected_checkpoint).isdigit() - else str(selected_checkpoint) - ) - checkpoint_path = run_dir / checkpoint_name - return (checkpoint_path, run_dir) if checkpoint_path.exists() else (None, run_dir) - - last_path = run_dir / "hora_stage2_last.pt" - if last_path.exists(): - return last_path, run_dir - - numbered = [ - path for path in run_dir.glob("hora_stage2_*.pt") if path.stem.split("_")[-1].isdigit() - ] - if not numbered: - return None, run_dir - return max(numbered, key=lambda path: int(path.stem.split("_")[-1])), run_dir - - -def format_hora_stage2_checkpoint_error( - cfg: DictConfig, - *, - task_log_root: Path, - load_path: Path | None, - load_path_dir: Path | None, -) -> str: - """Build the user-facing diagnostic for an unresolvable stage-2 checkpoint.""" - selected_checkpoint = OmegaConf.select(cfg, "algo.checkpoint", default=-1) - checkpoint_hint = ( - f" algo.checkpoint={selected_checkpoint!r}" - if selected_checkpoint not in (None, "", -1, "-1") - else "" - ) - if load_path_dir is not None and load_path is None and checkpoint_hint: - reason = f"Requested stage-2 checkpoint was not found under resolved_run={load_path_dir}." - elif not task_log_root.exists(): - reason = "Task log root does not exist." - else: - latest_run = get_latest_run(task_log_root) - if latest_run is None: - reason = "No run directories were found under the task log root." - else: - reason = "Requested run or stage-2 checkpoint could not be resolved." - return ( - "Could not resolve a stage-2 HORA checkpoint for play mode. " - f"{reason} task={cfg.training.task_name} task_log_root={task_log_root} " - f"algo.load_run={cfg.algo.load_run!r}{checkpoint_hint}. " - "Use algo.load_run= and optionally " - "algo.checkpoint=." - ) diff --git a/src/unilab/visualization/interactive_playback.py b/src/unilab/visualization/interactive_playback.py index 515b632c9..1fe931a17 100644 --- a/src/unilab/visualization/interactive_playback.py +++ b/src/unilab/visualization/interactive_playback.py @@ -3,7 +3,7 @@ from __future__ import annotations import copy -from collections.abc import Callable, Mapping +from collections.abc import Callable from dataclasses import dataclass, field from pathlib import Path from typing import Any, ClassVar, Protocol, cast @@ -288,7 +288,6 @@ def __init__( normalizer: Any | None, num_envs: int, obs_extractor: Callable[[dict[str, np.ndarray]], np.ndarray], - priv_info_resolver: Callable[..., np.ndarray | None], ) -> None: self.env = env self.device = device @@ -298,9 +297,7 @@ def __init__( self.normalizer = normalizer self.num_envs = int(num_envs) self.obs_extractor = obs_extractor - self.priv_info_resolver = priv_info_resolver self.obs: np.ndarray | None = None - self.current_priv_info: np.ndarray | None = None self.step_count = 0 def reset(self) -> np.ndarray: @@ -310,9 +307,8 @@ def reset(self) -> np.ndarray: reset_result = self.env.reset(env_indices) if not isinstance(reset_result, tuple) or len(reset_result) != 2: raise ValueError(f"Unexpected env.reset return format: {type(reset_result)!r}") - obs_out, info_out = reset_result + obs_out, _ = reset_result self.obs = np.asarray(self.obs_extractor(obs_out), dtype=np.float32) - self.current_priv_info = self._resolve_priv_info(obs_out, info_out) self.step_count = 0 return self.obs @@ -320,7 +316,6 @@ def step_once(self) -> np.ndarray: actions = self._build_actions() state = self.env.step(actions) self.obs = np.asarray(self.obs_extractor(state.obs), dtype=np.float32) - self.current_priv_info = self._resolve_priv_info(state.obs, state.info) self.step_count += 1 return self.obs @@ -339,28 +334,6 @@ def info(self) -> dict[str, Any]: info = getattr(state, "info", None) return info if isinstance(info, dict) else {} - def _resolve_priv_info( - self, - obs_dict: dict[str, np.ndarray], - info: dict[str, Any] | None, - ) -> np.ndarray | None: - if self.actor_algo_type != "hora_sac": - return None - if self.action_mode != "policy" or self.actor is None: - return None - from uni_rl.utils.observations import split_obs_dict - - actor_obs_np, critic_np = split_obs_dict(obs_dict) - priv_info = self.priv_info_resolver( - algo_type=self.actor_algo_type, - obs_np=np.asarray(actor_obs_np, dtype=np.float32), - critic_np=np.asarray(critic_np, dtype=np.float32), - info=info, - ) - if priv_info is None: - raise ValueError("HORA-SAC interactive play step is missing privileged info.") - return np.asarray(priv_info, dtype=np.float32) - def _build_actions(self) -> np.ndarray: if self.obs is None: raise RuntimeError("Playback session must be reset before stepping.") @@ -372,17 +345,7 @@ def _build_actions(self) -> np.ndarray: obs_torch = obs_torch.float() if self.normalizer is not None: obs_torch = self.normalizer(obs_torch, update=False) - if self.actor_algo_type == "hora_sac": - if self.current_priv_info is None: - raise ValueError("HORA-SAC interactive play step is missing privileged info.") - priv_info_torch = torch.from_numpy(self.current_priv_info).to(self.device) - actions = self.actor.explore( - obs_torch, - priv_info_torch, - deterministic=True, - ) - else: - actions = self.actor.explore(obs_torch, deterministic=True) + actions = self.actor.explore(obs_torch, deterministic=True) return actions.detach().cpu().numpy().astype(np.float32) if self.action_mode == "random": return np.random.uniform( @@ -393,9 +356,6 @@ def _build_actions(self) -> np.ndarray: return np.zeros((self.num_envs, action_dim), dtype=np.float32) -_HORA_DISTILL_CHECKPOINT_UNAVAILABLE = "hora_distill_checkpoint_unavailable" - - def make_sim2sim_preflight( cfg: Any, *, @@ -658,7 +618,6 @@ def _build_appo_actor( cfg: Any, rl_cfg: dict[str, Any], device: str, - is_hora: bool, ) -> tuple[Any, int, int]: """Build the APPO actor and return it with ``(obs_dim, action_dim)``.""" from copy import deepcopy @@ -673,62 +632,6 @@ def _build_appo_actor( action_dim = int(action_shape[0]) rl_cfg_dict = deepcopy(rl_cfg) - if is_hora: - from uni_rl.algos.hora.models import build_hora_shared_actor_critic - from uni_rl.algos.hora.rsl_rl_compat import ( - convert_config_v3_to_v4, - is_rsl_rl_v4, - is_rsl_rl_v5, - ) - - from unilab.scripts.play_hora_appo import _update_hora_obs_groups - - obs_td = wrapped_env.get_observations() - num_envs = int(getattr(wrapped_env, "num_envs", getattr(env, "num_envs", 1))) - obs_dim = int(obs_td["actor"].shape[-1]) - priv_info_dim = int(obs_td["priv_info"].shape[-1]) - if priv_info_dim <= 0: - raise ValueError("HORA APPO interactive play requires privileged info.") - _update_hora_obs_groups(rl_cfg_dict, obs_dim=obs_dim, priv_info_dim=priv_info_dim) - if is_rsl_rl_v5(): - pass - elif is_rsl_rl_v4(): - rl_cfg_dict = convert_config_v3_to_v4(rl_cfg_dict) - - actor_cfg = deepcopy(rl_cfg_dict["actor"]) - actor_cls = resolve_callable(actor_cfg.pop("class_name")) - actor_cfg.pop("num_actions", None) - critic_cfg = deepcopy(rl_cfg_dict.get("critic") or rl_cfg_dict.get("actor") or {}) - critic_cfg.pop("class_name", None) - critic_cfg.pop("num_actions", None) - critic_cfg.pop("distribution_cfg", None) - shared_model = build_hora_shared_actor_critic( - obs_dim=obs_dim, - action_dim=action_dim, - priv_info_dim=priv_info_dim, - actor_cfg=actor_cfg, - critic_cfg=critic_cfg, - ).to(device) - td_example = TensorDict( - { - "actor": torch.zeros((num_envs, obs_dim), device=device), - "priv_info": torch.zeros( - (num_envs, priv_info_dim), - device=device, - ), - }, - batch_size=num_envs, - ) - actor = actor_cls( - td_example, - rl_cfg_dict["obs_groups"], - "actor", - action_dim, - shared_model=shared_model, - **actor_cfg, - ) - return actor.to(device).eval(), obs_dim, action_dim - obs_dim, critic_dim = get_obs_dims(env.obs_groups_spec) num_envs = int(getattr(wrapped_env, "num_envs", getattr(env, "num_envs", 1))) obs_groups = rl_cfg_dict.setdefault("obs_groups", {}) @@ -774,18 +677,8 @@ def create_appo_playback_session( if env is None: raise RuntimeError("Playback env factory did not return an environment.") - from uni_rl.algos.hora.runtime import is_hora_appo_runtime - - is_hora = is_hora_appo_runtime(rl_cfg) - selected_wrapper_cls = wrapper_cls policy_obs_mode = playback_cfg.policy_obs_mode - if is_hora: - from uni_rl.algos.hora.rsl_rl import HoraRslRlVecEnvWrapper - - selected_wrapper_cls = HoraRslRlVecEnvWrapper - policy_obs_mode = "actor" - - wrapped_env = selected_wrapper_cls(env, device=device_name, policy_obs_mode=policy_obs_mode) + wrapped_env = wrapper_cls(env, device=device_name, policy_obs_mode=policy_obs_mode) policy = None actor = None checkpoint_path: str | None = None @@ -809,7 +702,6 @@ def create_appo_playback_session( cfg=cfg, rl_cfg=rl_cfg, device=device_name, - is_hora=is_hora, ) checkpoint = torch.load(checkpoint_path, map_location=device_name, weights_only=True) with policy_load_dim_guard( @@ -1049,7 +941,6 @@ def create_sac_playback_session( import os from uni_rl.algos.common.actor_factory import build_actor - from uni_rl.offpolicy.worker import resolve_offpolicy_actor_priv_info from unilab.utils.checkpoint import resolve_offpolicy_checkpoint_path @@ -1135,169 +1026,12 @@ def create_sac_playback_session( normalizer=normalizer, num_envs=playback_cfg.num_envs, obs_extractor=extract_play_obs, - priv_info_resolver=resolve_offpolicy_actor_priv_info, ), "actor", checkpoint_path, ) -def _default_hora_distill_playback_deps(root_dir: str | Path) -> dict[str, Any]: - from uni_rl.algos.hora.distill import ( - build_student_actor_and_normalizer, - cfg_with_checkpoint_runtime, - load_distilled_checkpoint, - student_policy, - ) - from uni_rl.algos.hora.rsl_rl import HoraRslRlVecEnvWrapper - - from unilab.base.config_adapter import BackendAdapter, create_env - from unilab.training import ( - format_hora_stage2_checkpoint_error, - get_log_root, - resolve_hora_stage2_checkpoint_path, - ) - from unilab.training.hora_distill_config import apply_teacher_defaults - - return { - "apply_teacher_defaults": apply_teacher_defaults, - "build_play_env_cfg_override": lambda cfg: BackendAdapter( - cfg, - root_dir=root_dir, - algo_name="hora_distill", - scene_materializer=_scene_visual_materializer(), - ).build_play_env_cfg_override(), - "build_student_actor_and_normalizer": build_student_actor_and_normalizer, - "cfg_with_checkpoint_runtime": cfg_with_checkpoint_runtime, - "create_env": create_env, - "format_stage2_play_checkpoint_error": format_hora_stage2_checkpoint_error, - "get_log_root": get_log_root, - "load_distilled_checkpoint": load_distilled_checkpoint, - "resolve_stage2_checkpoint_path": lambda cfg: resolve_hora_stage2_checkpoint_path( - cfg, - root_dir=root_dir, - ), - "student_policy": student_policy, - "wrapper_cls": HoraRslRlVecEnvWrapper, - "checkpoint_reader": torch.load, - } - - -def create_hora_distill_playback_session( - *, - playback_cfg: RslRlPlaybackConfig, - cfg: Any, - root_dir: str | Path, - device: str | None, - deps: Mapping[str, Any] | None = None, - log: LogFn = print, -) -> tuple[RslRlPlaybackSession, str, str | None]: - """Create an interactive playback session for HORA stage-2 student checkpoints.""" - - resolved_deps = dict(_default_hora_distill_playback_deps(root_dir) if deps is None else deps) - device_name = select_torch_device() if device is None else str(device) - load_path, load_path_dir = resolved_deps["resolve_stage2_checkpoint_path"](cfg) - checkpoint_path = str(load_path) if load_path is not None else None - policy: Callable[[Any], Any] | None = None - - if playback_cfg.action_mode == "policy": - if load_path is None or load_path_dir is None or not Path(load_path).exists(): - task_log_root = resolved_deps["get_log_root"](Path(root_dir), cfg) / str( - cfg.training.task_name - ) - log( - resolved_deps["format_stage2_play_checkpoint_error"]( - cfg, - task_log_root=task_log_root, - load_path=load_path, - load_path_dir=load_path_dir, - ) - ) - log("WARNING: falling back to zero actions.") - runtime_cfg = resolved_deps["apply_teacher_defaults"](cfg) - else: - resolve_sim2sim_config( - load_path_dir, - cfg, - algo_name="hora_distill", - strict=bool(getattr(cfg.training, "sim2sim_strict", True)), - ) - log(f"Loading distilled checkpoint: {load_path}") - checkpoint = resolved_deps["checkpoint_reader"]( - load_path, map_location="cpu", weights_only=False - ) - if "model_state_dict" not in checkpoint: - raise ValueError( - f"Checkpoint at {load_path} is not a HORA distillation checkpoint " - f"(found keys: {set(checkpoint.keys())})." - ) - # uni_rl's cfg_with_checkpoint_runtime no longer composes teacher - # defaults; the caller owns that composition (issue #1480). - runtime_cfg = resolved_deps["cfg_with_checkpoint_runtime"]( - resolved_deps["apply_teacher_defaults"](cfg), checkpoint - ) - else: - runtime_cfg = resolved_deps["apply_teacher_defaults"](cfg) - - env_cfg_override = resolved_deps["build_play_env_cfg_override"](runtime_cfg) - create_env = resolved_deps["create_env"] - try: - env = create_env( - runtime_cfg, - num_envs=int(playback_cfg.num_envs), - env_cfg_override=env_cfg_override, - sim_backend="mujoco", - task_name=str(runtime_cfg.training.task_name), - ) - except TypeError: - if deps is None: - raise - env = create_env( - runtime_cfg, - num_envs=int(playback_cfg.num_envs), - env_cfg_override=env_cfg_override, - ) - if env is None: - raise RuntimeError("Playback env factory did not return an environment.") - - policy_obs_mode = "actor" - wrapper_cls = resolved_deps["wrapper_cls"] - wrapped_env = wrapper_cls(env, device=device_name, policy_obs_mode=policy_obs_mode) - torch_device = torch.device(device_name) - - if playback_cfg.action_mode == "policy" and load_path is not None and Path(load_path).exists(): - actor, hist_normalizer = resolved_deps["build_student_actor_and_normalizer"]( - wrapped_env, - runtime_cfg, - device=torch_device, - ) - with policy_load_dim_guard(algo_name="hora_distill"): - resolved_deps["load_distilled_checkpoint"]( - actor, - hist_normalizer, - load_path, - device=torch_device, - ) - actor.eval() - hist_normalizer.eval() - student_policy = resolved_deps["student_policy"] - - def policy(obs: Any) -> Any: - return student_policy(actor, hist_normalizer, obs, device=torch_device) - - log(f"Policy obs mode: {policy_obs_mode}") - log(f"Action mode: {playback_cfg.action_mode}") - session = RslRlPlaybackSession( - env=env, - wrapped_env=wrapped_env, - device=device_name, - action_mode=playback_cfg.action_mode, - policy=policy, - num_envs=playback_cfg.num_envs, - ) - return session, policy_obs_mode, checkpoint_path - - def prepare_motion_overlay_selection( env: Any, *, @@ -1364,7 +1098,6 @@ def prepare_motion_overlay_selection( "build_play_backend_adapter", "build_playback_config", "create_appo_playback_session", - "create_hora_distill_playback_session", "create_rsl_rl_playback_session", "create_sac_playback_session", "infer_checkpoint_actor_input_dim", diff --git a/tests/algos/test_hora_contract.py b/tests/algos/test_hora_contract.py deleted file mode 100644 index 0a834d763..000000000 --- a/tests/algos/test_hora_contract.py +++ /dev/null @@ -1,601 +0,0 @@ -from __future__ import annotations - -import ast -import copy -import inspect -import logging -import textwrap -from typing import Any, cast - -import pytest -import torch -from tensordict import TensorDict - - -def test_hora_ppo_logs_when_symmetry_is_logging_only( - caplog: pytest.LogCaptureFixture, -) -> None: - from uni_rl.algos.hora.ppo import HoraPPO - - actor = torch.nn.Linear(2, 2) - critic = torch.nn.Linear(2, 1) - actor.is_recurrent = False # type: ignore[attr-defined] - critic.is_recurrent = False # type: ignore[attr-defined] - symmetry_cfg = { - "use_data_augmentation": False, - "use_mirror_loss": False, - "data_augmentation_func": lambda *_args: None, - } - - with caplog.at_level(logging.WARNING, logger="uni_rl.algos.hora.ppo"): - HoraPPO( - cast(Any, actor), - cast(Any, critic), - cast(Any, object()), - symmetry_cfg=symmetry_cfg, - ) - - assert "Symmetry not used for learning. We will use it for logging instead." in caplog.text - - -def test_hora_sac_actor_shapes_and_stable_module_names() -> None: - from uni_rl.algos.hora.sac_models import HoraSACActor - - actor = HoraSACActor( - obs_dim=5, - priv_info_dim=3, - action_dim=2, - hidden_dim=16, - priv_info_embed_dim=4, - priv_mlp_hidden_dims=(8, 4), - use_layer_norm=False, - ) - - obs = torch.zeros(6, 5) - priv_info = torch.zeros(6, 3) - actions, log_probs, log_std = actor.get_actions_and_log_probs(obs, priv_info) - - assert hasattr(actor, "priv_encoder") - assert hasattr(actor, "actor_trunk") - assert hasattr(actor, "action_mean_head") - assert hasattr(actor, "action_logstd_head") - assert actions.shape == (6, 2) - assert log_probs.shape == (6,) - assert log_std.shape == (6, 2) - - -def test_hora_sac_learner_derives_priv_info_from_critic_contract() -> None: - from uni_rl.algos.hora.sac_learner import derive_priv_info_from_critic_obs - - actor_obs = torch.zeros((4, 5), dtype=torch.float32) - priv_info = torch.arange(12, dtype=torch.float32).reshape(4, 3) - critic_obs = torch.cat([actor_obs, priv_info], dim=-1) - - torch.testing.assert_close( - derive_priv_info_from_critic_obs(actor_obs, critic_obs, context="test"), - priv_info, - ) - - with pytest.raises(ValueError, match="privileged tail"): - derive_priv_info_from_critic_obs(actor_obs, actor_obs, context="test") - - -def test_hora_sac_learner_updates_with_privileged_tail() -> None: - from uni_rl.algos.hora.sac_learner import HoraSACLearner - - torch.manual_seed(23) - learner = HoraSACLearner( - obs_dim=5, - critic_obs_dim=8, - priv_info_dim=3, - action_dim=2, - device="cpu", - actor_hidden_dim=16, - critic_hidden_dim=16, - priv_info_embed_dim=4, - priv_mlp_hidden_dims=(8, 4), - num_atoms=11, - use_layer_norm=False, - actor_lr=1e-3, - critic_lr=1e-3, - alpha_lr=1e-3, - ) - obs = torch.randn(7, 5) - next_obs = torch.randn(7, 5) - priv_info = torch.randn(7, 3) - next_priv_info = torch.randn(7, 3) - batch = { - "obs": obs, - "critic": torch.cat([obs, priv_info], dim=-1), - "actions": torch.randn(7, 2).clamp(-0.5, 0.5), - "rewards": torch.randn(7), - "next_obs": next_obs, - "next_critic": torch.cat([next_obs, next_priv_info], dim=-1), - "dones": torch.zeros(7), - "truncated": torch.zeros(7), - } - - critic_metrics = learner.update_critic(batch) - actor_metrics = learner.update_actor(batch) - - assert torch.isfinite(torch.tensor(list(critic_metrics.values()))).all() - assert torch.isfinite(torch.tensor(list(actor_metrics.values()))).all() - - -def test_hora_sac_disables_cuda_graph_critic_path() -> None: - from uni_rl.algos.hora.sac_learner import HoraSACLearner - - learner = HoraSACLearner( - obs_dim=5, - critic_obs_dim=8, - priv_info_dim=3, - action_dim=2, - device="cpu", - actor_hidden_dim=16, - critic_hidden_dim=16, - num_atoms=11, - use_layer_norm=False, - use_cuda_graph_critic=True, - use_cuda_graph_actor=True, - ) - - assert not learner.use_cuda_graph_critic - assert not learner.use_cuda_graph_actor - - -def test_hora_sac_distilled_student_forward_does_not_require_priv_info() -> None: - from uni_rl.algos.hora.distill import HoraSACDistillActor, HoraSACDistillShared - - shared = HoraSACDistillShared( - obs_dim=12, - action_dim=4, - priv_info_dim=3, - hidden_dim=32, - priv_info_embed_dim=3, - priv_mlp_hidden_dims=(8, 3), - use_layer_norm=False, - proprio_hist_len=30, - proprio_frame_dim=2, - device="cpu", - ) - actor = HoraSACDistillActor(shared) - student_obs = TensorDict( - { - "actor": torch.zeros((5, 12), dtype=torch.float32), - "proprio_hist": torch.zeros((5, 30, 2), dtype=torch.float32), - }, - batch_size=[5], - ) - - actions = actor(student_obs) - - assert actions.shape == (5, 4) - with pytest.raises(ValueError, match="priv_info is required"): - shared.policy_mean(student_obs, prefer_student=True) - - -def test_hora_sac_distill_loads_teacher_actor_weights(tmp_path) -> None: - from uni_rl.algos.hora.distill import ( - HoraSACDistillActor, - HoraSACDistillShared, - load_teacher_actor_weights, - ) - from uni_rl.algos.hora.sac_models import HoraSACActor - - teacher = HoraSACActor( - obs_dim=12, - priv_info_dim=3, - action_dim=4, - hidden_dim=32, - priv_info_embed_dim=3, - priv_mlp_hidden_dims=(8, 3), - use_layer_norm=False, - ) - shared = HoraSACDistillShared( - obs_dim=12, - action_dim=4, - priv_info_dim=3, - hidden_dim=32, - priv_info_embed_dim=3, - priv_mlp_hidden_dims=(8, 3), - use_layer_norm=False, - proprio_hist_len=30, - proprio_frame_dim=2, - device="cpu", - ) - actor = HoraSACDistillActor(shared) - checkpoint = tmp_path / "model.pt" - torch.save({"actor": teacher.state_dict()}, checkpoint) - - load_teacher_actor_weights( - actor, - checkpoint, - teacher_algo_family="sac", - device=torch.device("cpu"), - ) - - torch.testing.assert_close( - actor.shared.action_mean_head.weight, - teacher.action_mean_head.weight, - ) - torch.testing.assert_close( - actor.shared.encode_privileged_info(torch.zeros(2, 3)), - teacher.encode_privileged_info(torch.zeros(2, 3)), - ) - - -def test_hora_rsl_wrapper_uses_explicit_np_env_state_contract() -> None: - """HORA wrapper must not probe required NpEnvState fields dynamically.""" - from uni_rl.algos.hora.rsl_rl import HoraRslRlVecEnvWrapper - - source = textwrap.dedent(inspect.getsource(HoraRslRlVecEnvWrapper.step)) - tree = ast.parse(source) - forbidden_calls: list[str] = [] - for node in ast.walk(tree): - if not isinstance(node, ast.Call) or not isinstance(node.func, ast.Name): - continue - if node.func.id not in {"getattr", "hasattr"}: - continue - if not node.args or not isinstance(node.args[0], ast.Name): - continue - if node.args[0].id == "state": - forbidden_calls.append(node.func.id) - - assert forbidden_calls == [] - - -def test_hora_appo_learner_derives_priv_info_from_critic_contract() -> None: - from uni_rl.algos.hora.appo_learner import _derive_priv_info_from_critic - - actor_obs = torch.zeros((2, 3, 4), dtype=torch.float32) - priv_info = torch.arange(12, dtype=torch.float32).reshape(2, 3, 2) - critic_obs = torch.cat([actor_obs, priv_info], dim=-1) - - torch.testing.assert_close( - _derive_priv_info_from_critic(actor_obs, critic_obs, context="test"), - priv_info, - ) - - with pytest.raises(ValueError, match="privileged tail"): - _derive_priv_info_from_critic(actor_obs, actor_obs, context="test") - - -def _make_hora_appo_learner(**algorithm_overrides): - from uni_rl.algos.hora.appo_learner import HoraAPPOLearner - from uni_rl.algos.hora.models import ( - HoraActorModel, - HoraCriticModel, - HoraSharedActorCritic, - ) - - obs = TensorDict( - { - "actor": torch.zeros(4, 5), - "priv_info": torch.zeros(4, 2), - }, - batch_size=4, - ) - shared = HoraSharedActorCritic( - obs_dim=5, - action_dim=3, - priv_info_dim=2, - priv_info_embed_dim=2, - actor_hidden_dims=(8,), - priv_mlp_hidden_dims=(4, 2), - ) - actor = HoraActorModel(obs, {}, "actor", 3, shared_model=shared) - critic = HoraCriticModel(obs, {}, "critic", 1, shared_model=shared) - kwargs = { - "actor": actor, - "critic": critic, - "num_learning_epochs": 1, - "num_mini_batches": 1, - "device": "cpu", - } - kwargs.update(algorithm_overrides) - return HoraAPPOLearner(**kwargs) - - -def test_hora_appo_learner_uses_one_shared_actor_critic_core() -> None: - learner = _make_hora_appo_learner() - - assert learner.actor.shared is learner.critic.shared - - -def test_hora_appo_minibatch_tensor_path_matches_tensordict_forward() -> None: - torch.manual_seed(29) - learner = _make_hora_appo_learner() - actor_obs = torch.randn(6, 5) - priv_info = torch.randn(6, 2) - critic_obs = torch.cat([actor_obs, priv_info], dim=-1) - obs_td = TensorDict( - {"actor": actor_obs, "priv_info": priv_info}, - batch_size=actor_obs.shape[0], - ) - - with torch.inference_mode(): - learner.actor(obs_td, stochastic_output=True) - expected_mean = learner.actor.output_mean.clone() - expected_std = learner.actor.output_std.clone() - expected_value = learner.critic(obs_td).squeeze(-1) - mean, std, value = learner._minibatch_policy_value(actor_obs, critic_obs) - - torch.testing.assert_close(mean, expected_mean) - torch.testing.assert_close(std, expected_std) - torch.testing.assert_close(value, expected_value) - - -def test_hora_appo_runner_builds_shared_actor_critic_core() -> None: - from uni_rl.algos.hora.appo_runner import HoraAPPORunner - - runner = HoraAPPORunner.__new__(HoraAPPORunner) - runner.num_envs = 4 - runner.obs_dim = 5 - runner.action_dim = 3 - runner.priv_info_dim = 2 - runner.device = "cpu" - runner.seed = None - runner.rl_cfg = { - "obs_groups": { - "actor": {"actor": 5, "priv_info": 2}, - "critic": {"actor": 5, "priv_info": 2}, - }, - "actor": { - "class_name": "uni_rl.algos.hora:HoraActorModel", - "hidden_dims": [8], - "priv_info_embed_dim": 2, - "priv_mlp_hidden_dims": [4, 2], - }, - "critic": { - "class_name": "uni_rl.algos.hora:HoraCriticModel", - "priv_info_embed_dim": 2, - "priv_mlp_hidden_dims": [4, 2], - }, - "algorithm": { - "num_learning_epochs": 1, - "num_mini_batches": 1, - }, - } - - learner = runner._build_learner() - - assert learner.actor.shared is learner.critic.shared - - -def test_hora_appo_worker_builds_shared_actor_critic_core() -> None: - from uni_rl.algos.hora.appo_worker import hora_appo_collector_fn - - source = textwrap.dedent(inspect.getsource(hora_appo_collector_fn)) - tree = ast.parse(source) - shared_model_keywords = [ - node - for node in ast.walk(tree) - if isinstance(node, ast.keyword) - and node.arg == "shared_model" - and isinstance(node.value, ast.Name) - and node.value.id == "shared_model" - ] - - assert "build_hora_shared_actor_critic" in source - assert len(shared_model_keywords) >= 2 - - -def test_hora_appo_play_builds_explicit_shared_actor_core() -> None: - from unilab.scripts.play_hora_appo import play_hora_appo - - source = textwrap.dedent(inspect.getsource(play_hora_appo)) - tree = ast.parse(source) - shared_model_keywords = [ - node - for node in ast.walk(tree) - if isinstance(node, ast.keyword) - and node.arg == "shared_model" - and isinstance(node.value, ast.Name) - and node.value.id == "shared_model" - ] - - assert "build_hora_shared_actor_critic" in source - assert len(shared_model_keywords) >= 1 - - -def _patch_hora_appo_play_fakes(monkeypatch: pytest.MonkeyPatch, *, actor_cls: type) -> None: - """Patch env/actor construction fakes for ``play_hora_appo`` contract tests.""" - from types import SimpleNamespace - - import numpy as np - import rsl_rl.utils as rsl_rl_utils - - import unilab.scripts.play_hora_appo as hora_appo - - fake_env = SimpleNamespace( - obs_groups_spec={"obs": 3, "critic": 5}, - action_space=SimpleNamespace(shape=(2,)), - state=SimpleNamespace( - obs={"obs": np.zeros((1, 3), dtype=np.float32)}, - info={}, - ), - cfg=SimpleNamespace(render_spacing=1.0), - run_playback_mode=lambda **kwargs: None, - ) - - class FakeBackendAdapter: - def __init__(self, *args, **kwargs): - pass - - def build_task_env_cfg_override(self): - return {} - - monkeypatch.setattr(hora_appo, "BackendAdapter", FakeBackendAdapter) - monkeypatch.setattr(hora_appo, "create_env", lambda *args, **kwargs: fake_env) - monkeypatch.setattr( - hora_appo, - "split_hora_obs_with_priv_info", - lambda obs, info: ( - np.zeros((1, 3), dtype=np.float32), - None, - np.zeros((1, 2), dtype=np.float32), - ), - ) - monkeypatch.setattr(hora_appo, "is_rsl_rl_v5", lambda: True) - monkeypatch.setattr( - hora_appo, - "build_hora_shared_actor_critic", - lambda **kwargs: torch.nn.Identity(), - ) - monkeypatch.setattr(rsl_rl_utils, "resolve_callable", lambda path: actor_cls) - - -def _hora_appo_play_cfg(): - from omegaconf import OmegaConf - - return OmegaConf.create( - { - "training": { - "task_name": "Task", - "device": "cpu", - "play_env_num": 1, - "cam_distance": 3.0, - "cam_elevation": -20.0, - "cam_azimuth": 45.0, - }, - } - ) - - -def test_hora_appo_play_runs_sim2sim_preflight_before_checkpoint_load( - monkeypatch: pytest.MonkeyPatch, - tmp_path, -) -> None: - import unilab.scripts.play_hora_appo as hora_appo - - checkpoint = tmp_path / "model_10.pt" - torch.save({"actor": {"weight": torch.tensor(1.0)}}, checkpoint) - events: list[str] = [] - preflight_calls: list[tuple] = [] - - class FakeActor(torch.nn.Module): - def __init__(self, *args, **kwargs): - super().__init__() - - def load_state_dict(self, state_dict, strict=True): - events.append("load") - return None - - _patch_hora_appo_play_fakes(monkeypatch, actor_cls=FakeActor) - - def fake_preflight(source_run_dir, target_cfg, *, algo_name, strict): - events.append("preflight") - preflight_calls.append((source_run_dir, algo_name, strict)) - return target_cfg - - monkeypatch.setattr(hora_appo, "resolve_sim2sim_config", fake_preflight) - - hora_appo.play_hora_appo( - _hora_appo_play_cfg(), - {"actor": {"class_name": "fake.Actor"}, "critic": {}}, - root_dir=tmp_path, - resolve_checkpoint_path=lambda current_cfg: (str(checkpoint), str(tmp_path)), - ) - - assert preflight_calls == [(str(tmp_path), "appo", True)] - assert events == ["preflight", "load"] - - -def test_hora_appo_play_dim_mismatch_reraises_explicit_sim2sim_diagnostic( - monkeypatch: pytest.MonkeyPatch, - tmp_path, - capsys, -) -> None: - import unilab.scripts.play_hora_appo as hora_appo - from unilab.utils.sim2sim import CrossBackendIncompatibleError - - checkpoint = tmp_path / "model_10.pt" - torch.save({"actor": {"weight": torch.tensor(1.0)}}, checkpoint) - - class MismatchActor(torch.nn.Module): - def __init__(self, *args, **kwargs): - super().__init__() - - def load_state_dict(self, state_dict, strict=True): - raise RuntimeError("size mismatch for actor.mlp.0.weight: copying a param ...") - - _patch_hora_appo_play_fakes(monkeypatch, actor_cls=MismatchActor) - - with pytest.raises(CrossBackendIncompatibleError, match="does not fit this play environment"): - hora_appo.play_hora_appo( - _hora_appo_play_cfg(), - {"actor": {"class_name": "fake.Actor"}, "critic": {}}, - root_dir=tmp_path, - resolve_checkpoint_path=lambda current_cfg: (str(checkpoint), str(tmp_path)), - ) - - # Old runs without a contract snapshot keep the fallback + warning semantics. - assert "no contract_snapshot" in capsys.readouterr().out - - -def test_hora_appo_resume_rejects_inconsistent_shared_checkpoint() -> None: - from uni_rl.algos.hora.appo_runner import _validate_hora_shared_checkpoint - - learner = _make_hora_appo_learner() - joint_checkpoint = { - "actor": copy.deepcopy(learner.actor.state_dict()), - "critic": copy.deepcopy(learner.critic.state_dict()), - } - - _validate_hora_shared_checkpoint(joint_checkpoint) - - bad_checkpoint = copy.deepcopy(joint_checkpoint) - shared_key = next( - key - for key, value in bad_checkpoint["critic"].items() - if key.startswith("shared.") and torch.is_floating_point(value) - ) - bad_checkpoint["critic"][shared_key] = bad_checkpoint["critic"][shared_key] + 1.0 - - with pytest.raises(ValueError, match="Invalid HORA APPO checkpoint"): - _validate_hora_shared_checkpoint(bad_checkpoint) - - -def test_hora_appo_combined_optimizer_has_unique_parameters() -> None: - learner = _make_hora_appo_learner() - - combined_ids = [id(param) for param in learner._combined_params] - optimizer_ids = [ - id(param) for group in learner.optimizer.param_groups for param in group["params"] - ] - - assert len(combined_ids) == len(set(combined_ids)) - assert optimizer_ids == combined_ids - assert len(optimizer_ids) == len(set(optimizer_ids)) - - -def test_hora_appo_update_uses_joint_shared_optimizer() -> None: - torch.manual_seed(13) - learner = _make_hora_appo_learner(learning_rate=1e-3) - observations = torch.randn(2, 3, 5) - priv_info = torch.randn(2, 3, 2) - last_obs = torch.randn(3, 5) - last_priv_info = torch.randn(3, 2) - batch = { - "observations": observations, - "critic": torch.cat([observations, priv_info], dim=-1), - "actions": torch.randn(2, 3, 3), - "actions_log_prob": torch.zeros(2, 3), - "rewards": torch.randn(2, 3), - "dones": torch.zeros(2, 3), - "last_obs": last_obs, - "last_critic": torch.cat([last_obs, last_priv_info], dim=-1), - } - - trunk_before = [param.detach().clone() for param in learner.actor.shared.trunk.parameters()] - - learner.process_batch(batch) - metrics = learner.update(batch) - - trunk_after = list(learner.actor.shared.trunk.parameters()) - - assert metrics["appo/updates_executed"] == pytest.approx(1.0) - assert any( - not torch.allclose(before, after) - for before, after in zip(trunk_before, trunk_after, strict=True) - ) diff --git a/tests/algos/test_hora_distill_config.py b/tests/algos/test_hora_distill_config.py deleted file mode 100644 index 7db6be786..000000000 --- a/tests/algos/test_hora_distill_config.py +++ /dev/null @@ -1,126 +0,0 @@ -"""Contract tests for HORA distill teacher-owner Hydra composition.""" - -from __future__ import annotations - -from pathlib import Path - -import pytest -from omegaconf import OmegaConf -from omegaconf.errors import InterpolationResolutionError - -from unilab.training import hora_distill_config as distill_config - -# distill_config's root_dir parameter expects the directory containing "conf"; -# after the packaging move that is the unilab package directory. -_PACKAGE_ROOT = Path(__file__).resolve().parents[2] / "src" / "unilab" - - -# --------------------------------------------------------------------------- -# Hydra composition capabilities used by teacher-owner configs -# --------------------------------------------------------------------------- - - -def _write(path: Path, content: str) -> None: - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text(content, encoding="utf-8") - - -def test_load_teacher_owner_config_supports_nested_defaults_packages_and_interpolation( - tmp_path: Path, -) -> None: - conf_dir = tmp_path / "conf" / "ppo" - _write( - conf_dir / "config.yaml", - "defaults:\n - _self_\n - task: group/owner\n - extra: packed\n\nroot_scalar: 3\n", - ) - # Package directive: file content lands under `extra`, not at the root. - _write(conf_dir / "extra" / "packed.yaml", "# @package extra\ninner: ${root_scalar}\n") - _write( - conf_dir / "task" / "group" / "owner.yaml", - "# @package _global_\ndefaults:\n - group/mid\n - _self_\n\nleaf: owner\n", - ) - # Nested defaults: an included group file with its own defaults list. - _write( - conf_dir / "task" / "group" / "mid.yaml", - "# @package _global_\ndefaults:\n - nested_leaf\n - _self_\n\nmid_value: 5\n", - ) - _write( - conf_dir / "task" / "group" / "nested_leaf.yaml", - "# @package _global_\nnested_value: 42\n", - ) - - cfg = distill_config.load_teacher_owner_config("ppo", "group/owner", root_dir=tmp_path) - - assert cfg.nested_value == 42 - assert cfg.mid_value == 5 - assert cfg.leaf == "owner" - assert cfg.extra.inner == 3 - - -def test_hora_sac_mapping_keeps_yaml_fallbacks_for_missing_teacher_fields(tmp_path: Path) -> None: - teacher_cfg = OmegaConf.create({"algo": {"runtime_impl": "hora_sac", "actor": {}}}) - - model_cfg = distill_config._student_model_defaults("hora_sac", teacher_cfg, root=_PACKAGE_ROOT) - - assert model_cfg == { - "teacher_arch": "hora_sac", - "actor_hidden_dim": 512, - "use_layer_norm": True, - "priv_info_embed_dim": 9, - "priv_mlp_hidden_dims": [256, 128, 9], - } - - -def test_hora_actor_mapping_strips_distribution_class_name() -> None: - teacher_cfg = OmegaConf.create( - { - "algo": { - "actor": { - "class_name": "uni_rl.algos.hora:HoraActorModel", - "hidden_dims": [64, 32], - "activation": "relu", - "obs_normalization": False, - "priv_info_embed_dim": 4, - "priv_mlp_hidden_dims": [16, 4], - "distribution_cfg": { - "class_name": "GaussianDistribution", - "init_std": 0.5, - "std_type": "scalar", - }, - } - } - } - ) - - model_cfg = distill_config._student_model_defaults( - "hora_actor", teacher_cfg, root=_PACKAGE_ROOT - ) - - assert model_cfg == { - "hidden_dims": [64, 32], - "activation": "relu", - "obs_normalization": False, - "priv_info_embed_dim": 4, - "priv_mlp_hidden_dims": [16, 4], - "distribution_cfg": {"init_std": 0.5, "std_type": "scalar"}, - } - - -def test_hora_actor_mapping_fails_closed_when_teacher_field_is_missing() -> None: - teacher_cfg = OmegaConf.create( - { - "algo": { - "actor": { - "class_name": "uni_rl.algos.hora:HoraActorModel", - "activation": "elu", - "obs_normalization": True, - "priv_info_embed_dim": 9, - "priv_mlp_hidden_dims": [256, 128, 9], - "distribution_cfg": {"init_std": 1.0, "std_type": "scalar"}, - } - } - } - ) - - with pytest.raises(InterpolationResolutionError): - distill_config._student_model_defaults("hora_actor", teacher_cfg, root=_PACKAGE_ROOT) diff --git a/tests/algos/test_offpolicy_double_buffer_runner.py b/tests/algos/test_offpolicy_double_buffer_runner.py index 1d072591c..2df95e9c9 100644 --- a/tests/algos/test_offpolicy_double_buffer_runner.py +++ b/tests/algos/test_offpolicy_double_buffer_runner.py @@ -94,12 +94,6 @@ def test_offpolicy_config_has_one_replay_path(): assert "replay_h2d_submitter" not in cfg.training -def test_hora_uses_same_learner_inference_path(): - cfg = _offpolicy_cfg(["task=sharpa_inhand/mujoco_hora"]) - assert "inference_owner" not in cfg.training - assert cfg.training.env_steps_per_sync == 2 - - @pytest.mark.parametrize( "override", [ diff --git a/tests/assets/test_hub.py b/tests/assets/test_hub.py index 78de7517a..3746812e7 100644 --- a/tests/assets/test_hub.py +++ b/tests/assets/test_hub.py @@ -10,7 +10,6 @@ from unilab.assets import ASSETS_ROOT_PATH from unilab.assets.hub import ( - resolve_grasp_cache_files, resolve_motion_files, resolve_robot_asset_dir, resolve_scene_dir, @@ -152,104 +151,6 @@ def test_resolve_relative_windows_path_uses_posix_hf_filename(): ) -# --------------------------------------------------------------------------- -# Grasp cache resolve — local fast path -# --------------------------------------------------------------------------- - - -def test_resolve_grasp_cache_returns_existing_path(tmp_path: Path): - npy = tmp_path / "cache.npy" - np.save(npy, np.array([1.0])) - assert resolve_grasp_cache_files(str(npy)) == str(npy) - - -def test_resolve_grasp_cache_returns_list_for_list_input(tmp_path: Path): - a = tmp_path / "a.npy" - b = tmp_path / "b.npy" - np.save(a, np.array([1.0])) - np.save(b, np.array([2.0])) - result = resolve_grasp_cache_files([str(a), str(b)]) - assert result == [str(a), str(b)] - - -# --------------------------------------------------------------------------- -# Grasp cache resolve — HF download path (mocked) -# --------------------------------------------------------------------------- - - -def test_resolve_grasp_cache_calls_hf_download_with_caches_repo(): - """Grasp cache resolver should use the unilab-caches repo, not unilab-motions.""" - missing = ASSETS_ROOT_PATH / "caches" / "__test_nonexistent__.npy" - assert not missing.exists() - - expected_relative = missing.relative_to(ASSETS_ROOT_PATH).as_posix() - - fake_download = MagicMock(return_value=str(missing)) - fake_module = MagicMock() - fake_module.hf_hub_download = fake_download - - with patch.dict("sys.modules", {"huggingface_hub": fake_module}): - result = resolve_grasp_cache_files(str(missing)) - - assert result == str(missing) - fake_download.assert_called_once() - kwargs = fake_download.call_args.kwargs - assert kwargs["repo_id"] == "unilabsim/unilab-caches" - assert kwargs["filename"] == expected_relative - assert kwargs["repo_type"] == "dataset" - assert kwargs["local_dir"] == str(ASSETS_ROOT_PATH) - assert kwargs["tqdm_class"] is not None - - -def test_resolve_grasp_cache_relative_path_uses_caches_repo(): - """A relative path triggers HF download with the caches repo.""" - rel = "caches/__test_nonexistent_rel__.npy" - local = ASSETS_ROOT_PATH / rel - assert not local.exists() - - fake_download = MagicMock(return_value=str(local)) - fake_module = MagicMock() - fake_module.hf_hub_download = fake_download - - with patch.dict("sys.modules", {"huggingface_hub": fake_module}): - result = resolve_grasp_cache_files(rel) - - assert result == str(local) - fake_download.assert_called_once() - kwargs = fake_download.call_args.kwargs - assert kwargs["repo_id"] == "unilabsim/unilab-caches" - assert kwargs["filename"] == rel - assert kwargs["repo_type"] == "dataset" - assert kwargs["local_dir"] == str(ASSETS_ROOT_PATH) - assert kwargs["tqdm_class"] is not None - - -def test_resolve_grasp_cache_can_disable_download_progress(): - rel = "caches/__test_nonexistent_quiet__.npy" - local = ASSETS_ROOT_PATH / rel - assert not local.exists() - - fake_download = MagicMock(return_value=str(local)) - fake_module = MagicMock() - fake_module.hf_hub_download = fake_download - - with patch.dict("sys.modules", {"huggingface_hub": fake_module}): - result = resolve_grasp_cache_files(rel, show_progress=False) - - assert result == str(local) - fake_download.assert_called_once() - kwargs = fake_download.call_args.kwargs - assert kwargs["repo_id"] == "unilabsim/unilab-caches" - assert kwargs["filename"] == rel - assert kwargs["repo_type"] == "dataset" - assert kwargs["local_dir"] == str(ASSETS_ROOT_PATH) - progress = kwargs["tqdm_class"](range(1)) - try: - assert progress.disable is True - finally: - progress.close() - - # Scene directory resolver # --------------------------------------------------------------------------- @@ -351,7 +252,6 @@ def test_robot_asset_specs_cover_hf_hosted_robots(): "go1", "go2", "go2w", - "sharpa_wave", "x2", } assert set(ROBOT_ASSET_SPECS) == expected @@ -397,13 +297,13 @@ def test_ensure_robot_assets_handles_absolute_and_windows_paths( hub.ensure_robot_assets_for_paths( [ "/home/user/project/src/unilab/assets/robots/go2/go2.xml", - r"src\unilab\assets\robots\sharpa_wave\right_sharpa_wave.xml", + r"src\unilab\assets\robots\x2\x2.xml", ] ) assert calls == [ ("robots/go2/assets", "base_0.obj"), - ("robots/sharpa_wave/meshes", "DP_HB1_4F.STL"), + ("robots/x2/meshes", "pelvis.STL"), ] diff --git a/tests/base/test_backend_imports.py b/tests/base/test_backend_imports.py index 33ed47ed9..23a545aff 100644 --- a/tests/base/test_backend_imports.py +++ b/tests/base/test_backend_imports.py @@ -11,7 +11,6 @@ _MATERIALIZER_CONSUMERS = ( "src/unilab/base/config_adapter.py", "src/unilab/scripts/train_rsl_rl.py", - "scripts/train_hora_distill.py", ) diff --git a/tests/base/test_sim_backend.py b/tests/base/test_sim_backend.py index 3421a718c..d2ac7d527 100644 --- a/tests/base/test_sim_backend.py +++ b/tests/base/test_sim_backend.py @@ -36,7 +36,6 @@ def _xml(robot: str, scene: str = "scene_flat.xml") -> str: _G1 = dict(model_file=_xml("g1"), base_name="pelvis") _ALLEGRO = dict(model_file=_xml("allegro_hand", "scene.xml"), base_name="palm") -_SHARPA = dict(model_file=_xml("sharpa_wave", "scene.xml"), base_name="right_hand_C_MC") NUM_ENVS = 2 SIM_DT = 0.005 @@ -166,11 +165,11 @@ def test_apply_init_randomization_sets_variants_before_materialization(self): from unisim.backend.mujoco.backend import MuJoCoBackend bkd = MuJoCoBackend( - SceneCfg(model_file=_SHARPA["model_file"]), 4, SIM_DT, base_name=_SHARPA["base_name"] + SceneCfg(model_file=_G1["model_file"]), 4, SIM_DT, base_name=_G1["base_name"] ) assert bkd._pool is None mujoco = _mujoco_module() - geom_id = mujoco.mj_name2id(bkd.model, mujoco.mjtObj.mjOBJ_GEOM, "object") + geom_id = mujoco.mj_name2id(bkd.model, mujoco.mjtObj.mjOBJ_GEOM, "floor") base_size = np.asarray(bkd.model.geom_size[geom_id], dtype=np.float64).copy() bkd.apply_init_randomization( @@ -178,10 +177,10 @@ def test_apply_init_randomization_sets_variants_before_materialization(self): model_assignments=np.array([0, 1, 0, 1], dtype=np.int32), model_variants=( ModelVariantSpec( - geom_size_overrides=(GeomSizeOverride("object", tuple(base_size * 0.5)),) + geom_size_overrides=(GeomSizeOverride("floor", tuple(base_size * 0.5)),) ), ModelVariantSpec( - geom_size_overrides=(GeomSizeOverride("object", tuple(base_size * 0.75)),) + geom_size_overrides=(GeomSizeOverride("floor", tuple(base_size * 0.75)),) ), ), ) @@ -202,10 +201,10 @@ def test_get_playback_model_returns_env_specific_variant(self): from unisim.backend.mujoco.backend import MuJoCoBackend bkd = MuJoCoBackend( - SceneCfg(model_file=_SHARPA["model_file"]), 4, SIM_DT, base_name=_SHARPA["base_name"] + SceneCfg(model_file=_G1["model_file"]), 4, SIM_DT, base_name=_G1["base_name"] ) mujoco = _mujoco_module() - geom_id = mujoco.mj_name2id(bkd.model, mujoco.mjtObj.mjOBJ_GEOM, "object") + geom_id = mujoco.mj_name2id(bkd.model, mujoco.mjtObj.mjOBJ_GEOM, "floor") base_size = np.asarray(bkd.model.geom_size[geom_id], dtype=np.float64).copy() bkd.apply_init_randomization( @@ -213,10 +212,10 @@ def test_get_playback_model_returns_env_specific_variant(self): model_assignments=np.array([0, 1, 0, 1], dtype=np.int32), model_variants=( ModelVariantSpec( - geom_size_overrides=(GeomSizeOverride("object", tuple(base_size * 0.5)),) + geom_size_overrides=(GeomSizeOverride("floor", tuple(base_size * 0.5)),) ), ModelVariantSpec( - geom_size_overrides=(GeomSizeOverride("object", tuple(base_size * 0.75)),) + geom_size_overrides=(GeomSizeOverride("floor", tuple(base_size * 0.75)),) ), ), ) diff --git a/tests/base/test_sim_backend_smoke.py b/tests/base/test_sim_backend_smoke.py index 40eeab175..51d79a94a 100644 --- a/tests/base/test_sim_backend_smoke.py +++ b/tests/base/test_sim_backend_smoke.py @@ -31,7 +31,6 @@ def _xml(robot: str, scene: str = "scene_flat.xml") -> str: _G1 = dict(model_file=_xml("g1"), base_name="pelvis") _ALLEGRO = dict(model_file=_xml("allegro_hand", "scene.xml"), base_name="palm") -_SHARPA = dict(model_file=_xml("sharpa_wave", "scene.xml"), base_name="right_hand_C_MC") NUM_ENVS = 2 SIM_DT = 0.005 @@ -543,16 +542,16 @@ def test_mujoco_metadata_getters_return_stable_copies(): from unilab.assets.hub import ensure_robot_assets_for_paths - ensure_robot_assets_for_paths([_SHARPA["model_file"]]) + ensure_robot_assets_for_paths([_G1["model_file"]]) bkd = MuJoCoBackend( - SceneCfg(model_file=_SHARPA["model_file"]), NUM_ENVS, SIM_DT, base_name=_SHARPA["base_name"] + SceneCfg(model_file=_G1["model_file"]), NUM_ENVS, SIM_DT, base_name=_G1["base_name"] ) model = bkd.model - object_geom_id = int(mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, "object")) - base_body_id = int(mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, _SHARPA["base_name"])) + floor_geom_id = int(mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, "floor")) + base_body_id = int(mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, _G1["base_name"])) - assert bkd.get_geom_id("object") == object_geom_id - assert bkd.get_body_id(_SHARPA["base_name"]) == base_body_id + assert bkd.get_geom_id("floor") == floor_geom_id + assert bkd.get_body_id(_G1["base_name"]) == base_body_id with pytest.raises(ValueError, match="Geom 'missing'"): bkd.get_geom_id("missing") with pytest.raises(ValueError, match="Body 'missing'"): @@ -564,17 +563,17 @@ def test_mujoco_metadata_getters_return_stable_copies(): default_qpos[0] += 1.0 assert not np.isclose(default_qpos[0], model.qpos0[0]) - geom_size = bkd.get_geom_size("object") + geom_size = bkd.get_geom_size("floor") _shape(geom_size, 3) - np.testing.assert_allclose(geom_size, model.geom_size[object_geom_id]) + np.testing.assert_allclose(geom_size, model.geom_size[floor_geom_id]) geom_size[0] += 1.0 - assert not np.isclose(geom_size[0], model.geom_size[object_geom_id, 0]) + assert not np.isclose(geom_size[0], model.geom_size[floor_geom_id, 0]) geom_body_ids = bkd.get_geom_body_ids() _shape(geom_body_ids, model.ngeom) np.testing.assert_array_equal(geom_body_ids, model.geom_bodyid) - geom_body_ids[object_geom_id] = -1 - assert int(model.geom_bodyid[object_geom_id]) != -1 + geom_body_ids[floor_geom_id] = -1 + assert int(model.geom_bodyid[floor_geom_id]) != -1 geom_contype, geom_conaffinity = bkd.get_geom_contact_masks() _shape(geom_contype, model.ngeom) @@ -584,14 +583,14 @@ def test_mujoco_metadata_getters_return_stable_copies(): geom_names = bkd.get_geom_names() assert len(geom_names) == model.ngeom - assert geom_names[object_geom_id] == "object" + assert geom_names[floor_geom_id] == "floor" assert base_body_id in set(int(body_id) for body_id in bkd.get_body_subtree_ids(base_body_id)) geom_friction = bkd.get_geom_friction() _shape(geom_friction, model.ngeom, 3) np.testing.assert_allclose(geom_friction, model.geom_friction) - geom_friction[object_geom_id, 0] += 1.0 - assert not np.isclose(geom_friction[object_geom_id, 0], model.geom_friction[object_geom_id, 0]) + geom_friction[floor_geom_id, 0] += 1.0 + assert not np.isclose(geom_friction[floor_geom_id, 0], model.geom_friction[floor_geom_id, 0]) gravity = bkd.get_gravity() _shape(gravity, 3) @@ -734,9 +733,9 @@ def test_motrix_default_qpos_uses_mujoco_quaternion_convention(): from unilab.assets.hub import ensure_robot_assets_for_paths - ensure_robot_assets_for_paths([_SHARPA["model_file"]]) + ensure_robot_assets_for_paths([_G1["model_file"]]) bkd = MotrixBackend( - SceneCfg(model_file=_SHARPA["model_file"]), NUM_ENVS, SIM_DT, base_name=_SHARPA["base_name"] + SceneCfg(model_file=_G1["model_file"]), NUM_ENVS, SIM_DT, base_name=_G1["base_name"] ) qpos = bkd.get_default_qpos() assert qpos.ndim == 1 @@ -755,9 +754,9 @@ def test_motrix_default_qpos_uses_mujoco_quaternion_convention(): np.broadcast_to(qpos, (NUM_ENVS, qpos.shape[0])).copy(), np.zeros((NUM_ENVS, bkd.model.num_dof_vel), dtype=np.float64), ) - object_body_id = bkd.get_body_id("object") + base_body_id = bkd.get_body_id(_G1["base_name"]) np.testing.assert_allclose( - np.abs(bkd.get_body_quat_w(np.asarray([object_body_id], dtype=np.int32))[:, 0, :]), + np.abs(bkd.get_body_quat_w(np.asarray([base_body_id], dtype=np.int32))[:, 0, :]), np.broadcast_to([1.0, 0.0, 0.0, 0.0], (NUM_ENVS, 4)), atol=1.0e-6, ) diff --git a/tests/benchmark/test_offpolicy_collector_active_benchmark.py b/tests/benchmark/test_offpolicy_collector_active_benchmark.py index 6daf14fdc..9edf25e15 100644 --- a/tests/benchmark/test_offpolicy_collector_active_benchmark.py +++ b/tests/benchmark/test_offpolicy_collector_active_benchmark.py @@ -70,7 +70,7 @@ def test_parse_case_requires_algo_task_sim() -> None: bench._parse_case("g1_walk_flat/mujoco") -def test_default_cases_cover_motrix_only_without_sharpa() -> None: +def test_default_cases_cover_motrix_only() -> None: specs = bench._resolve_case_specs( "default", algos_arg="sac,flashsac,td3", @@ -80,7 +80,6 @@ def test_default_cases_cover_motrix_only_without_sharpa() -> None: assert "sac/g1_motion_tracking/motrix" in specs assert "flashsac/g1_walk_flat/motrix" in specs assert "sac/g1_motion_tracking/mujoco" not in specs - assert "sac/sharpa_inhand/mujoco_hora" not in specs def test_all_backend_selection_expands_default_cases() -> None: diff --git a/tests/benchmark/test_sharpa_init_dr_benchmark.py b/tests/benchmark/test_sharpa_init_dr_benchmark.py deleted file mode 100644 index aa6fe87ab..000000000 --- a/tests/benchmark/test_sharpa_init_dr_benchmark.py +++ /dev/null @@ -1,19 +0,0 @@ -from __future__ import annotations - -import numpy as np -from scripts.benchmark.env import benchmark_sharpa_init_dr_construct as sharpa_benchmark - - -def test_sharpa_init_dr_benchmark_uses_owner_scale_list_key() -> None: - cfg = sharpa_benchmark._compose_cfg( - "sharpa_inhand/mujoco", - lower=0.5, - upper=0.8, - variant_count=4, - ) - - np.testing.assert_allclose( - np.asarray(cfg.env.domain_rand.scale_list, dtype=np.float64), - np.linspace(0.5, 0.8, 4, dtype=np.float64), - ) - assert "scale_list" not in cfg.env diff --git a/tests/config/test_config_system.py b/tests/config/test_config_system.py index 234ea5e77..6102c15ed 100644 --- a/tests/config/test_config_system.py +++ b/tests/config/test_config_system.py @@ -297,47 +297,25 @@ def test_ppo_g1_backend_specific_hyperparams_remain_separate(): ("flashsac", ["task=g1_walk_flat/mujoco"]), ], ) -def test_post_step_forward_sensor_defaults_false_outside_sharpa_mujoco( - algo_dir: str, overrides: list[str] -): +def test_post_step_forward_sensor_defaults_false(algo_dir: str, overrides: list[str]): cfg = _compose(algo_dir, overrides=overrides) assert cfg.env.post_step_forward_sensor is False -@pytest.mark.parametrize( - ("algo_dir", "overrides"), - [ - ("ppo", ["task=sharpa_inhand/mujoco"]), - ("ppo", ["task=sharpa_inhand/mujoco_hora"]), - ("ppo", ["task=sharpa_inhand_grasp/mujoco"]), - ("appo", ["task=sharpa_inhand/mujoco"]), - ("appo", ["task=sharpa_inhand/mujoco_hora"]), - ("sac", ["task=sharpa_inhand/mujoco_hora"]), - ("hora_distill", ["task=sharpa_inhand/mujoco"]), - ], -) -def test_post_step_forward_sensor_enabled_for_sharpa_mujoco(algo_dir: str, overrides: list[str]): - cfg = _compose(algo_dir, overrides=overrides) - - assert cfg.env.post_step_forward_sensor is True - - def test_mujoco_post_step_forward_sensor_can_be_overridden(): override_cfg = _compose( - "appo", - overrides=["task=sharpa_inhand/mujoco_hora", "env.post_step_forward_sensor=false"], + "ppo", + overrides=["task=g1_walk_flat/mujoco", "env.post_step_forward_sensor=true"], ) - assert override_cfg.env.post_step_forward_sensor is False + assert override_cfg.env.post_step_forward_sensor is True def test_appo_adaptive_lr_factors_are_overridden_only_by_dex_hand_owners(): g1_cfg = _compose("appo", overrides=["task=g1_walk_flat/mujoco"]) allegro_cfg = _compose("appo", overrides=["task=allegro_inhand/mujoco"]) allegro_motrix_cfg = _compose("appo", overrides=["task=allegro_inhand/motrix"]) - sharpa_cfg = _compose("appo", overrides=["task=sharpa_inhand/mujoco"]) - sharpa_hora_cfg = _compose("appo", overrides=["task=sharpa_inhand/mujoco_hora"]) assert g1_cfg.algo.algorithm.adaptive_kl_factor == pytest.approx(1.2) assert g1_cfg.algo.algorithm.adaptive_lr_factor == pytest.approx(1.1) @@ -345,10 +323,6 @@ def test_appo_adaptive_lr_factors_are_overridden_only_by_dex_hand_owners(): assert allegro_cfg.algo.algorithm.adaptive_lr_factor == pytest.approx(1.5) assert allegro_motrix_cfg.algo.algorithm.adaptive_kl_factor == pytest.approx(2.0) assert allegro_motrix_cfg.algo.algorithm.adaptive_lr_factor == pytest.approx(1.5) - assert sharpa_cfg.algo.algorithm.adaptive_kl_factor == pytest.approx(1.2) - assert sharpa_cfg.algo.algorithm.adaptive_lr_factor == pytest.approx(1.1) - assert sharpa_hora_cfg.algo.algorithm.adaptive_kl_factor == pytest.approx(1.2) - assert sharpa_hora_cfg.algo.algorithm.adaptive_lr_factor == pytest.approx(1.1) def test_ppo_go1_motrix_preserves_reward_and_algo_values(): diff --git a/tests/envs/test_env_configs.py b/tests/envs/test_env_configs.py index 186ed2a90..4f24ecd40 100644 --- a/tests/envs/test_env_configs.py +++ b/tests/envs/test_env_configs.py @@ -6,7 +6,6 @@ from __future__ import annotations -import logging import subprocess import sys import textwrap @@ -610,244 +609,6 @@ def test_allegro_grasp_recorder_close_autosaves_and_io_failure_is_fail_closed( assert failed.cache_saved is False -def _patch_sharpa_rotation_constructor( - monkeypatch: pytest.MonkeyPatch, - initialized_providers: list[Any], -) -> Any: - from unilab.tasks.manipulation.sharpa_inhand import rotation as sharpa_rotation_module - from unilab.tasks.manipulation.sharpa_inhand.base import SharpaInhandBaseEnv - - def fake_base_init(self, cfg, backend, num_envs): - self._cfg = cfg - self._backend = backend - self._num_envs = num_envs - self._np_dtype = np.float64 - self._num_action = 22 - self._num_tactile = 5 - self._num_scales = len(cfg.domain_rand.scale_list) - self.scale_values = np.asarray(cfg.domain_rand.scale_list, dtype=np.float64) - self.scale_ids = np.zeros((num_envs,), dtype=np.int32) - self._object_body_ids = np.zeros((0,), dtype=np.int32) - - def unsupported_backend_metadata(*args, **kwargs): - raise NotImplementedError("fake backend does not expose Sharpa MuJoCo metadata") - - monkeypatch.setattr( - sharpa_rotation_module, - "create_backend", - lambda *args, **kwargs: SimpleNamespace( - backend_type="motrix", - get_actuator_gains=lambda: ( - np.ones(22, dtype=np.float64), - np.ones(22, dtype=np.float64), - ), - get_geom_id=unsupported_backend_metadata, - get_body_id=unsupported_backend_metadata, - get_body_subtree_ids=unsupported_backend_metadata, - get_geom_body_ids=unsupported_backend_metadata, - get_geom_contact_masks=unsupported_backend_metadata, - get_geom_names=unsupported_backend_metadata, - get_geom_friction=unsupported_backend_metadata, - get_gravity=unsupported_backend_metadata, - get_body_mass=unsupported_backend_metadata, - get_body_ipos=unsupported_backend_metadata, - ), - ) - monkeypatch.setattr(SharpaInhandBaseEnv, "__init__", fake_base_init) - monkeypatch.setattr( - sharpa_rotation_module.SharpaInhandRotationEnv, - "_init_domain_randomization", - lambda self, provider: initialized_providers.append(provider), - ) - return sharpa_rotation_module - - -def test_sharpa_rotation_explicit_default_provider_materializes_cache(monkeypatch): - from unilab.tasks.manipulation.sharpa_inhand.rotation import ( - RewardConfig, - SharpaInhandRotationCfg, - SharpaInhandRotationDRProvider, - SharpaInhandRotationEnv, - ) - - initialized_providers: list[Any] = [] - rotation = _patch_sharpa_rotation_constructor(monkeypatch, initialized_providers) - materialize_calls: list[tuple[str, np.ndarray]] = [] - sentinel_cache = (np.zeros((1, 29), dtype=np.float64),) - - def materialize(path: str, scale_values: np.ndarray) -> tuple[np.ndarray, ...]: - materialize_calls.append((path, scale_values.copy())) - return sentinel_cache - - monkeypatch.setattr(rotation, "_materialize_grasp_caches", materialize) - cfg = SharpaInhandRotationCfg(reward_config=RewardConfig()) - provider = SharpaInhandRotationDRProvider() - - env = cast(Any, SharpaInhandRotationEnv)( - cfg, - num_envs=4, - backend_type="mujoco", - dr_provider=provider, - ) - - assert initialized_providers == [provider] - assert len(materialize_calls) == 1 - assert materialize_calls[0][0] == cfg.grasp_cache_path - np.testing.assert_array_equal(materialize_calls[0][1], env.scale_values) - assert env._grasp_cache is sentinel_cache - - -def test_sharpa_grasp_env_initializes_dr_once_with_grasp_provider(monkeypatch): - from unilab.tasks.manipulation.sharpa_inhand.grasp_gen import ( - SharpaInhandRotationGraspCfg, - SharpaInhandRotationGraspEnv, - ) - - initialized_providers: list[Any] = [] - rotation = _patch_sharpa_rotation_constructor(monkeypatch, initialized_providers) - - def fail_materialization(*_args: Any, **_kwargs: Any) -> None: - raise AssertionError("grasp generation must not materialize rotation caches") - - monkeypatch.setattr(rotation, "_materialize_grasp_caches", fail_materialization) - - cfg = SharpaInhandRotationGraspCfg() - assert cfg.domain_rand.randomize_pd_gains is False - assert cfg.domain_rand.randomize_friction is False - assert cfg.domain_rand.randomize_com is False - assert cfg.domain_rand.randomize_mass is True - assert cfg.domain_rand.randomize_mass_lower == pytest.approx(0.05) - assert cfg.domain_rand.randomize_mass_upper == pytest.approx(0.051) - assert cfg.domain_rand.force_scale == pytest.approx(0.0) - assert cfg.domain_rand.random_force_prob_scalar == pytest.approx(0.0) - assert cfg.domain_rand.joint_noise_scale == pytest.approx(0.02) - assert cfg.domain_rand.contact_latency == pytest.approx(0.005) - assert cfg.domain_rand.contact_sensor_noise == pytest.approx(0.01) - assert cfg.control_config.torque_control is False - assert cfg.control_config.dof_limits_scale == pytest.approx(0.9) - assert cfg.obs.enable_tactile is True - assert cfg.obs.binary_contact is False - assert cfg.obs.enable_contact_pos is False - assert cfg.obs.contact_smooth == pytest.approx(0.5) - assert cfg.obs.contact_threshold == pytest.approx(0.05) - assert cfg.obs.tactile_force_clip_max == pytest.approx(4.0) - assert cfg.priv_info.include_friction_scale is True - assert cfg.priv_info.include_gravity_direction is False - env = cast(Any, SharpaInhandRotationGraspEnv)(cfg, num_envs=4, backend_type="mujoco") - - assert [provider.__class__.__name__ for provider in initialized_providers] == [ - "SharpaInhandGraspDRProvider" - ] - assert len(env._saved_grasping_states) == env._num_scales - - -def test_sharpa_grasp_target_saves_cache_then_raises_run_complete( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, - caplog: pytest.LogCaptureFixture, -) -> None: - from unilab.base.run_control import RunComplete - from unilab.tasks.manipulation.sharpa_inhand.grasp_gen import ( - SharpaInhandRotationGraspCfg, - SharpaInhandRotationGraspEnv, - ) - - cache_path = tmp_path / "sharpa.npy" - env = cast(Any, object.__new__(SharpaInhandRotationGraspEnv)) - env._cfg = SharpaInhandRotationGraspCfg( - grasp_collection_target=1, - grasp_cache_path=str(cache_path), - grasp_auto_save=True, - ) - env._grasp_target_per_scale = 1 - env._num_scales = 1 - env.scale_values = np.asarray([0.8], dtype=np.float64) - env.scale_ids = np.asarray([0], dtype=np.int32) - env._saved_grasping_states = [[]] - env._grasp_cache_saved = False - env._grasp_target_reached_notified = False - env._last_grasp_progress_counts = () - env._last_grasp_progress_step = -1 - env._state = SimpleNamespace( - truncated=np.asarray([True]), - terminated=np.asarray([False]), - info={"steps": np.asarray([1], dtype=np.uint32), "log": {}}, - ) - env.get_hand_dof_pos = lambda: np.arange(22, dtype=np.float64).reshape(1, 22) - env.get_object_pos = lambda: np.asarray([[1.0, 2.0, 3.0]], dtype=np.float64) - env.get_object_quat = lambda: np.asarray([[1.0, 0.0, 0.0, 0.0]], dtype=np.float64) - - save_calls: list[Path] = [] - real_save = np.save - - def save_once(path: str | Path, values: np.ndarray) -> None: - save_calls.append(Path(path)) - real_save(path, values) - - monkeypatch.setattr( - "unilab.tasks.manipulation.sharpa_inhand.grasp_gen.np.save", - save_once, - ) - - with caplog.at_level( - logging.INFO, - logger="unilab.tasks.manipulation.sharpa_inhand.grasp_gen", - ): - with pytest.raises(RunComplete) as caught: - env._collect_successful_grasps(np.asarray([0], dtype=np.int32)) - - saved_path = tmp_path / "sharpa_0.8.npy" - saved = np.load(saved_path) - expected = np.concatenate( - [ - np.arange(22, dtype=np.float32), - np.asarray([1.0, 2.0, 3.0], dtype=np.float32), - np.asarray([1.0, 0.0, 0.0, 0.0], dtype=np.float32), - ] - )[None, :] - np.testing.assert_array_equal(saved, expected) - assert saved.dtype == np.float32 - assert save_calls == [saved_path] - assert "grasp progress total=1/1, per_scale=[scale=0.8:1]" in caplog.text - assert "target reached (saved=1, configured_target=1)" in caplog.text - assert env.state.info["log"]["grasp_cache/saved"] == 1.0 - assert env.state.info["log"]["grasp_cache/num_states"] == 1.0 - assert env.state.info["log"]["grasp/target_reached"] == 1.0 - assert dict(caught.value.summary) == { - "collected_grasps": 1, - "grasp_collection_target": 1, - "grasp_collection_counts_by_scale": (1,), - } - - env._collect_successful_grasps(np.asarray([0], dtype=np.int32)) - env._stop_collection() - assert save_calls == [saved_path] - - -def test_sharpa_run_complete_reports_effective_collection_target() -> None: - from unilab.base.run_control import RunComplete - from unilab.tasks.manipulation.sharpa_inhand.grasp_gen import ( - SharpaInhandRotationGraspCfg, - SharpaInhandRotationGraspEnv, - ) - - env = cast(Any, object.__new__(SharpaInhandRotationGraspEnv)) - env._cfg = SharpaInhandRotationGraspCfg(grasp_collection_target=0) - env._grasp_target_per_scale = 1 - env._saved_grasping_states = [[np.zeros((1, 29), dtype=np.float32)]] - env._grasp_target_reached_notified = False - env._state = None - - with pytest.raises(RunComplete) as caught: - env._stop_collection() - - assert dict(caught.value.summary) == { - "collected_grasps": 1, - "grasp_collection_target": 1, - "grasp_collection_counts_by_scale": (1,), - } - - # --------------------------------------------------------------------------- # Fast env/backend smoke tests # --------------------------------------------------------------------------- diff --git a/tests/envs/test_sharpa.py b/tests/envs/test_sharpa.py deleted file mode 100644 index 55af028ae..000000000 --- a/tests/envs/test_sharpa.py +++ /dev/null @@ -1,1061 +0,0 @@ -from __future__ import annotations - -from pathlib import Path -from tempfile import TemporaryDirectory -from types import SimpleNamespace -from typing import Any - -import numpy as np -import pytest -from hydra import compose, initialize_config_dir -from hydra.core.global_hydra import GlobalHydra -from omegaconf import OmegaConf - -from unilab.base.registry import ensure_registries -from unilab.tasks.manipulation.sharpa_inhand.base import ( - SOURCE_DEFAULT_HAND_JOINT_POS_DEG, - SharpaInhandBaseEnv, - resolve_grasp_cache_file, -) -from unilab.tasks.manipulation.sharpa_inhand.rotation import ( - SharpaInhandRotationDRProvider, - _materialize_grasp_caches, -) - -_CONF_DIR = Path(__file__).resolve().parents[2] / "src" / "unilab" / "conf" -_SRC_DIR = Path(__file__).resolve().parents[2] / "src" - - -def test_sharpa_env_uses_backend_contract_for_mujoco_metadata() -> None: - """Sharpa env code should not read MuJoCo model internals directly.""" - source = "\n".join( - (_SRC_DIR / "unilab" / "tasks" / "manipulation" / "sharpa_inhand" / path).read_text( - encoding="utf-8" - ) - for path in ("base.py", "rotation.py") - ) - - assert "import mujoco" not in source - assert "self._backend.model" not in source - assert "_backend.model" not in source - - -def test_sharpa_motrix_owner_cfg_enables_supported_dr_terms() -> None: - cfg, env_cfg_override = _compose_sharpa_motrix_owner_cfg(num_envs=2) - - assert cfg.training.task_name == "SharpaInhandRotation" - assert cfg.training.sim_backend == "motrix" - assert env_cfg_override["domain_rand"]["scale_list"] == [ - 0.8, - 0.9, - 1.0, - 1.1, - 1.2, - 1.3, - 1.4, - 1.5, - ] - assert env_cfg_override["domain_rand"]["randomize_pd_gains"] is True - assert env_cfg_override["domain_rand"]["randomize_friction"] is True - assert env_cfg_override["domain_rand"]["randomize_com"] is True - assert env_cfg_override["domain_rand"]["randomize_mass"] is True - assert env_cfg_override["domain_rand"]["randomize_gravity"] is False - assert env_cfg_override["domain_rand"]["randomize_gravity_direction"] is True - assert env_cfg_override["domain_rand"]["force_scale"] == pytest.approx(2.0) - - -def test_sharpa_registry_exposes_rotation_and_grasp_motrix() -> None: - ensure_registries() - - from unilab.base import registry - - registered = registry.list_registered_envs() - - assert set(registered["SharpaInhandRotation"]["available_backends"]) == { - "drake", - "mujoco", - "motrix", - } - assert set(registered["SharpaInhandRotationGrasp"]["available_backends"]) == { - "mujoco", - "motrix", - } - - -def test_sharpa_grasp_motrix_owner_cfg() -> None: - GlobalHydra.instance().clear() - with initialize_config_dir(config_dir=str(_CONF_DIR / "ppo"), version_base="1.3"): - cfg = compose( - "config", - overrides=[ - "task=sharpa_inhand_grasp/motrix", - "algo.num_envs=2", - ], - ) - - env_cfg_override = OmegaConf.to_container(cfg.env, resolve=True) - assert isinstance(env_cfg_override, dict) - - assert cfg.training.task_name == "SharpaInhandRotationGrasp" - assert cfg.training.sim_backend == "motrix" - assert env_cfg_override["domain_rand"]["scale_list"] == [0.8] - assert env_cfg_override["domain_rand"]["randomize_gravity"] is False - assert env_cfg_override["domain_rand"]["randomize_gravity_direction"] is False - assert env_cfg_override["domain_rand"]["randomize_pd_gains"] is False - assert env_cfg_override["domain_rand"]["randomize_friction"] is False - assert env_cfg_override["domain_rand"]["randomize_com"] is False - assert env_cfg_override["domain_rand"]["randomize_mass"] is True - assert env_cfg_override["domain_rand"]["force_scale"] == pytest.approx(0.0) - - -def _require_mujoco_runtime() -> None: - """Require the MuJoCo batch runtime used by Sharpa reset randomization tests. - - Args: - None. - - Returns: - None. The helper skips the test when MuJoCo batch runtime is unavailable. - """ - pytest.importorskip("mujoco", reason="mujoco not installed") - try: - from mujoco_uni.batch_env import BatchEnvPool as _ # noqa: F401 - except Exception: - pytest.skip( - "mujoco_uni.batch_env not available (platform/libstdc++ issue)", - allow_module_level=False, - ) - - -def _require_motrix_runtime() -> None: - pytest.importorskip("motrixsim", reason="motrixsim not installed") - - -def _compose_sharpa_mujoco_owner_cfg(num_envs: int) -> tuple[Any, dict[str, Any]]: - """Compose the Sharpa MuJoCo owner config used by the real training path. - - Args: - num_envs: Number of vectorized environments for the test env. - - Returns: - Tuple of the composed Hydra config and the env_cfg_override dict for registry.make(). - """ - GlobalHydra.instance().clear() - with initialize_config_dir(config_dir=str(_CONF_DIR / "ppo"), version_base="1.3"): - cfg = compose( - "config", - overrides=[ - "task=sharpa_inhand/mujoco", - f"algo.num_envs={num_envs}", - ], - ) - - env_cfg_override = OmegaConf.to_container(cfg.env, resolve=True) - assert isinstance(env_cfg_override, dict) - env_cfg_override["reward_config"] = OmegaConf.to_container(cfg.reward, resolve=True) - return cfg, env_cfg_override - - -def _compose_sharpa_motrix_owner_cfg(num_envs: int) -> tuple[Any, dict[str, Any]]: - GlobalHydra.instance().clear() - with initialize_config_dir(config_dir=str(_CONF_DIR / "ppo"), version_base="1.3"): - cfg = compose( - "config", - overrides=[ - "task=sharpa_inhand/motrix", - f"algo.num_envs={num_envs}", - ], - ) - - env_cfg_override = OmegaConf.to_container(cfg.env, resolve=True) - assert isinstance(env_cfg_override, dict) - env_cfg_override["reward_config"] = OmegaConf.to_container(cfg.reward, resolve=True) - return cfg, env_cfg_override - - -def _write_sharpa_grasp_cache( - cache_prefix: Path, - scale_values: list[float], - *, - rows: int = 8, -) -> None: - hand_qpos = np.deg2rad(np.asarray(SOURCE_DEFAULT_HAND_JOINT_POS_DEG, dtype=np.float64)) - object_pose = np.asarray([-0.09559, -0.00517, 0.61906, 1.0, 0.0, 0.0, 0.0]) - cache = np.broadcast_to(np.concatenate([hand_qpos, object_pose]), (rows, 29)).copy() - for scale_value in scale_values: - np.save(resolve_grasp_cache_file(str(cache_prefix), float(scale_value)), cache) - - -def test_sharpa_grasp_cache_is_materialized_once_and_reset_samples_memory( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -) -> None: - cache_prefix = tmp_path / "sharpa_grasp" - scale_values = np.asarray([0.8, 1.0], dtype=np.float64) - _write_sharpa_grasp_cache(cache_prefix, [0.8, 1.0]) - - from unilab.tasks.manipulation.sharpa_inhand import rotation - - path_resolve_calls: list[tuple[str, float]] = [] - hf_resolve_calls: list[str] = [] - load_calls: list[Path] = [] - original_resolve_grasp_cache_file = rotation.resolve_grasp_cache_file - original_load = rotation.np.load - - def resolve_path_once(path: str, scale: float) -> Path: - path_resolve_calls.append((path, scale)) - return original_resolve_grasp_cache_file(path, scale) - - def resolve_hf_once(path: str) -> str: - hf_resolve_calls.append(path) - return path - - def load_once(path: Path) -> np.ndarray: - load_calls.append(path) - return original_load(path) - - monkeypatch.setattr(rotation, "resolve_grasp_cache_file", resolve_path_once) - monkeypatch.setattr(rotation, "resolve_grasp_cache_files", resolve_hf_once) - monkeypatch.setattr(rotation.np, "load", load_once) - caches = _materialize_grasp_caches(str(cache_prefix), scale_values) - assert len(path_resolve_calls) == 2 - assert len(hf_resolve_calls) == 2 - assert len(load_calls) == 2 - - def fail_io(*_args: Any, **_kwargs: Any) -> None: - raise AssertionError("reset must not resolve or load grasp cache files") - - monkeypatch.setattr(rotation, "resolve_grasp_cache_file", fail_io) - monkeypatch.setattr(rotation, "resolve_grasp_cache_files", fail_io) - monkeypatch.setattr(rotation.np, "load", fail_io) - - cleared_env_ids: list[np.ndarray] = [] - env = SimpleNamespace( - _grasp_cache=caches, - scale_ids=np.asarray([0, 1], dtype=np.int32), - _num_action=22, - _obj_pos_slice=slice(22, 25), - _obj_quat_slice=slice(25, 29), - _rot_axis=np.asarray([0.0, 0.0, 1.0]), - _random_object_force=np.ones((2, 3), dtype=np.float64), - nq=29, - nv=28, - cfg=SimpleNamespace(reset_height_lower=0.61, reset_height_upper=0.63), - _sample_friction_scale=lambda _num_reset: None, - _sample_object_mass=lambda _num_reset: None, - _sample_object_com_offset=lambda _num_reset: None, - _sample_reset_gravity=lambda _num_reset: None, - _clear_tactile_history=lambda env_ids: cleared_env_ids.append(env_ids.copy()), - _build_reset_randomization=lambda *_args, **_kwargs: None, - ) - provider = SharpaInhandRotationDRProvider() - - def sample_reset_pd_gains( - _env: Any, num_reset: int, *, dtype: np.dtype[Any] - ) -> tuple[np.ndarray, np.ndarray]: - gains = np.zeros((num_reset, env._num_action), dtype=dtype) - return gains, gains.copy() - - monkeypatch.setattr(provider, "_sample_reset_pd_gains", sample_reset_pd_gains) - monkeypatch.setattr(provider, "_build_info_updates", lambda _env, **_kwargs: {}) - - env_ids = np.asarray([0, 1], dtype=np.int32) - for _ in range(2): - plan = provider.build_reset_plan(env, env_ids) - assert plan.qpos.shape == (2, 29) - assert plan.qvel.shape == (2, 28) - assert np.all(np.isfinite(plan.qpos)) - - assert len(cleared_env_ids) == 2 - - -def _build_fake_tactile_env( - sensor_data: dict[str, np.ndarray], - *, - enable_tactile: bool = True, - binary_contact: bool = False, - disable_tactile_ids: list[int] | None = None, - contact_smooth: float = 0.5, - contact_threshold: float = 0.05, - contact_latency: float = 0.0, - contact_sensor_noise: float = 0.01, - last_contacts: np.ndarray | None = None, - prev_tactile_force: np.ndarray | None = None, -) -> Any: - """Build a minimal fake Sharpa env for tactile-observation unit tests. - - Args: - sensor_data: Mapping from tactile sensor name to backend sensor array. - enable_tactile: Whether tactile observation is enabled. - binary_contact: Whether binary-contact mode is enabled. - disable_tactile_ids: Optional tactile ids to zero out. - contact_smooth: Smoothing weight applied to the latest raw force. - contact_threshold: Threshold used by binary-contact mode. - contact_latency: Bernoulli probability of keeping the previous tactile output. - contact_sensor_noise: Binary-contact sensor dropout probability. - last_contacts: Optional previous tactile output buffer. - prev_tactile_force: Optional previous raw tactile-force buffer. - - Returns: - Simple fake env exposing the fields used by ``_compute_tactile_observation``. - """ - tactile_names = [ - "contact_right_thumb_elastomer_force", - "contact_right_index_elastomer_force", - "contact_right_middle_elastomer_force", - "contact_right_ring_elastomer_force", - "contact_right_pinky_elastomer_force", - ] - num_envs = next(iter(sensor_data.values())).shape[0] - env = SimpleNamespace( - _num_envs=num_envs, - _num_tactile=len(tactile_names), - _np_dtype=np.float64, - _backend=SimpleNamespace(get_sensor_data=lambda name: sensor_data[name]), - _cfg=SimpleNamespace( - obs=SimpleNamespace( - enable_tactile=enable_tactile, - binary_contact=binary_contact, - contact_smooth=contact_smooth, - contact_threshold=contact_threshold, - tactile_force_clip_max=5.0, - ), - domain_rand=SimpleNamespace( - contact_latency=contact_latency, - contact_sensor_noise=contact_sensor_noise, - ), - disable_tactile_ids=list(disable_tactile_ids or []), - sensor=SimpleNamespace(tactile_force_sensor_names=tactile_names), - ), - last_contacts=np.zeros((num_envs, len(tactile_names)), dtype=np.float64) - if last_contacts is None - else np.asarray(last_contacts, dtype=np.float64).copy(), - _prev_tactile_force=np.zeros((num_envs, len(tactile_names)), dtype=np.float64) - if prev_tactile_force is None - else np.asarray(prev_tactile_force, dtype=np.float64).copy(), - ) - env._extract_sensor_scalar = lambda sensor_name: SharpaInhandBaseEnv._extract_sensor_scalar( - env, sensor_name - ) - env._read_tactile_force = lambda: SharpaInhandBaseEnv._read_tactile_force(env) - env._clear_tactile_history = lambda env_ids=None: SharpaInhandBaseEnv._clear_tactile_history( - env, env_ids - ) - return env - - -def test_sharpa_motrix_reset_and_step_smoke() -> None: - _require_motrix_runtime() - ensure_registries() - - from unilab.base import registry - - num_envs = 2 - _, env_cfg_override = _compose_sharpa_motrix_owner_cfg(num_envs) - with TemporaryDirectory() as tmp_dir: - cache_prefix = Path(tmp_dir) / "sharpa_grasp" - env_cfg_override["grasp_cache_path"] = str(cache_prefix) - _write_sharpa_grasp_cache( - cache_prefix, - [float(value) for value in env_cfg_override["domain_rand"]["scale_list"]], - ) - - env = registry.make( - "SharpaInhandRotation", - num_envs=num_envs, - sim_backend="motrix", - env_cfg_override=env_cfg_override, - ) - env_obj: Any = env - try: - obs, info = env.reset(np.arange(num_envs, dtype=np.int32)) - assert isinstance(obs, dict) - assert set(obs) == {"obs"} - assert obs["obs"].shape[0] == num_envs - assert info["p_gain"].shape == (num_envs, env_obj._num_action) - assert info["d_gain"].shape == (num_envs, env_obj._num_action) - - state = env.step(np.zeros((num_envs, env_obj._num_action), dtype=env_obj._np_dtype)) - - assert isinstance(state.obs, dict) - assert state.obs["obs"].shape[0] == num_envs - assert np.all(np.isfinite(state.reward)) - finally: - env.close() - - -def test_sharpa_grasp_motrix_reset_and_step_smoke() -> None: - _require_motrix_runtime() - ensure_registries() - - from unilab.base import registry - - num_envs = 2 - GlobalHydra.instance().clear() - with initialize_config_dir(config_dir=str(_CONF_DIR / "ppo"), version_base="1.3"): - cfg = compose( - "config", - overrides=[ - "task=sharpa_inhand_grasp/motrix", - f"algo.num_envs={num_envs}", - ], - ) - - env_cfg_override = OmegaConf.to_container(cfg.env, resolve=True) - assert isinstance(env_cfg_override, dict) - env_cfg_override["reward_config"] = OmegaConf.to_container(cfg.reward, resolve=True) - - env = registry.make( - "SharpaInhandRotationGrasp", - num_envs=num_envs, - sim_backend="motrix", - env_cfg_override=env_cfg_override, - ) - env_obj: Any = env - try: - obs, info = env.reset(np.arange(num_envs, dtype=np.int32)) - assert isinstance(obs, dict) - assert set(obs) == {"obs"} - assert obs["obs"].shape[0] == num_envs - assert info["p_gain"].shape == (num_envs, env_obj._num_action) - assert info["d_gain"].shape == (num_envs, env_obj._num_action) - np.testing.assert_allclose( - np.abs(info["object_default_pose"][:, 3:7]), - np.abs(env_obj.get_object_quat()), - atol=1.0e-6, - ) - - state = env.step(np.zeros((num_envs, env_obj._num_action), dtype=env_obj._np_dtype)) - - assert isinstance(state.obs, dict) - assert state.obs["obs"].shape[0] == num_envs - assert np.all(np.isfinite(state.reward)) - finally: - env.close() - - -def test_sharpa_motrix_reset_applies_scale_mass_and_com_randomization() -> None: - _require_motrix_runtime() - ensure_registries() - - from unilab.base import registry - - num_envs = 9 - _, env_cfg_override = _compose_sharpa_motrix_owner_cfg(num_envs) - with TemporaryDirectory() as tmp_dir: - cache_prefix = Path(tmp_dir) / "sharpa_grasp" - env_cfg_override["grasp_cache_path"] = str(cache_prefix) - _write_sharpa_grasp_cache( - cache_prefix, - [float(value) for value in env_cfg_override["domain_rand"]["scale_list"]], - ) - - env = registry.make( - "SharpaInhandRotation", - num_envs=num_envs, - sim_backend="motrix", - env_cfg_override=env_cfg_override, - ) - env_obj: Any = env - try: - _, info = env.reset(np.arange(num_envs, dtype=np.int32)) - backend: Any = env_obj._backend - object_link = backend.model.get_link("object") - object_geom = backend.model.get_geom("object") - - mass = np.asarray(object_link.get_mass_override(backend.data), dtype=np.float64) - com = np.asarray( - object_link.get_center_of_mass_override(backend.data), dtype=np.float64 - ) - size = np.asarray(object_geom.get_size_override(backend.data), dtype=np.float64) - geom_friction = np.stack( - [ - np.asarray(backend.model.get_geom(geom_id).get_friction_override(backend.data)) - for geom_id in range(int(backend.model.num_geoms)) - ], - axis=1, - ).astype(np.float64) - layout = env_obj._critic_info_layout() - critic_info = np.asarray(info["critic_info"], dtype=np.float64) - - np.testing.assert_allclose(mass, critic_info[:, layout["mass"]].reshape(num_envs)) - np.testing.assert_allclose(com, critic_info[:, layout["com"]], atol=1.0e-7) - expected_size = env_obj._object_geom_base_size[:2] * env_obj.scale_values[ - env_obj.scale_ids - ].reshape(num_envs, 1) - np.testing.assert_allclose(size, expected_size) - assert np.unique(np.round(mass, 6)).size > 1 - assert np.unique(np.round(com.reshape(-1), 6)).size > 1 - assert np.unique(np.round(size[:, 0], 6)).size > 1 - - friction_scale = critic_info[:, layout["friction"]].reshape(num_envs) - assert np.unique(np.round(friction_scale, 6)).size > 1 - assert np.all(friction_scale >= env_obj._cfg.domain_rand.randomize_friction_scale_lower) - assert np.all(friction_scale <= env_obj._cfg.domain_rand.randomize_friction_scale_upper) - for env_idx in range(num_envs): - scale = friction_scale[env_idx] - for material, base_friction in ( - ("object", env_obj._cfg.domain_rand.object_base_friction), - ("metal", env_obj._cfg.domain_rand.metal_base_friction), - ("elastomer", env_obj._cfg.domain_rand.elastomer_base_friction), - ): - actual = geom_friction[env_idx, env_obj._friction_geom_ids[material]] - expected = env_obj._friction_profile(material, base_friction) * scale - np.testing.assert_allclose(actual, np.broadcast_to(expected, actual.shape)) - finally: - env.close() - - -def test_sharpa_motrix_interval_force_writes_object_link_external_force() -> None: - _require_motrix_runtime() - ensure_registries() - - from unilab.base import registry - - num_envs = 4 - _, env_cfg_override = _compose_sharpa_motrix_owner_cfg(num_envs) - env_cfg_override["domain_rand"]["force_scale"] = 2.0 - env_cfg_override["domain_rand"]["random_force_prob_scalar"] = 1.0 - with TemporaryDirectory() as tmp_dir: - cache_prefix = Path(tmp_dir) / "sharpa_grasp" - env_cfg_override["grasp_cache_path"] = str(cache_prefix) - _write_sharpa_grasp_cache( - cache_prefix, - [float(value) for value in env_cfg_override["domain_rand"]["scale_list"]], - ) - - env = registry.make( - "SharpaInhandRotation", - num_envs=num_envs, - sim_backend="motrix", - env_cfg_override=env_cfg_override, - ) - env_obj: Any = env - try: - env.init_state() - backend: Any = env_obj._backend - body_id = int(env_obj._object_body_id) - - assert backend.get_dr_capabilities().supports_interval_body_force - env_obj._dr_manager.apply_interval_randomization_if_due(env_obj.step_counter) - - assert np.any(np.linalg.norm(env_obj._random_object_force, axis=1) > 0.0) - np.testing.assert_allclose( - backend._applied_body_forces[body_id], - env_obj._random_object_force, - atol=1.0e-6, - ) - finally: - env.close() - - -def test_sharpa_provider_builds_mujoco_init_geom_scale_plan() -> None: - env = SimpleNamespace( - _backend=SimpleNamespace(backend_type="mujoco"), - _object_geom_base_size=np.array([0.02, 0.016, 0.0], dtype=np.float64), - scale_values=np.array([0.5, 0.8], dtype=np.float64), - scale_ids=np.array([0, 0, 1, 1], dtype=np.int32), - cfg=SimpleNamespace(object_geom_name="object"), - ) - - plan = SharpaInhandRotationDRProvider().build_init_randomization_plan(env) - - assert plan is not None - np.testing.assert_array_equal( - plan.model_assignments, - np.array([0, 0, 1, 1], dtype=np.int32), - ) - assert len(plan.model_variants) == 2 - assert plan.model_variants[0].geom_size_overrides[0].geom_name == "object" - np.testing.assert_allclose( - plan.model_variants[0].geom_size_overrides[0].size, - [0.01, 0.008, 0.0], - ) - np.testing.assert_allclose( - plan.model_variants[1].geom_size_overrides[0].size, - [0.016, 0.0128, 0.0], - ) - - -def test_sharpa_provider_builds_motrix_init_geom_scale_plan() -> None: - env = SimpleNamespace( - _backend=SimpleNamespace(backend_type="motrix"), - _object_geom_base_size=np.array([0.02, 0.016, 0.0], dtype=np.float64), - scale_values=np.array([0.5, 0.8], dtype=np.float64), - scale_ids=np.array([0, 0, 1, 1], dtype=np.int32), - cfg=SimpleNamespace(object_geom_name="object"), - ) - - plan = SharpaInhandRotationDRProvider().build_init_randomization_plan(env) - - assert plan is not None - np.testing.assert_array_equal(plan.model_assignments, np.array([0, 0, 1, 1], dtype=np.int32)) - assert len(plan.model_variants) == 2 - - -def test_sharpa_tactile_force_matches_reference_smoothing_and_order() -> None: - """Verify tactile force uses reference sensor order and raw-force smoothing. - - Args: - None. - - Returns: - None. The assertions validate thumb→pinky ordering and 2-step smoothing. - """ - sensor_data = { - "contact_right_thumb_elastomer_force": np.array([[1.0, 0.0, 0.0]], dtype=np.float64), - "contact_right_index_elastomer_force": np.array([[2.0, 0.0, 0.0]], dtype=np.float64), - "contact_right_middle_elastomer_force": np.array([[3.0, 0.0, 0.0]], dtype=np.float64), - "contact_right_ring_elastomer_force": np.array([[4.0, 0.0, 0.0]], dtype=np.float64), - "contact_right_pinky_elastomer_force": np.array([[5.0, 0.0, 0.0]], dtype=np.float64), - } - env = _build_fake_tactile_env( - sensor_data, - contact_smooth=0.25, - prev_tactile_force=np.array([[10.0, 20.0, 30.0, 40.0, 50.0]], dtype=np.float64), - ) - - tactile = SharpaInhandBaseEnv._compute_tactile_observation(env) - - expected = np.array([[7.75, 15.5, 23.25, 31.0, 38.75]], dtype=np.float64) - np.testing.assert_allclose(tactile, expected) - np.testing.assert_allclose(env.last_contacts, expected) - np.testing.assert_allclose( - env._prev_tactile_force, - np.array([[1.0, 2.0, 3.0, 4.0, 5.0]], dtype=np.float64), - ) - - -def test_sharpa_tactile_binary_mode_matches_reference_latency_noise_and_disable( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """Verify binary tactile mode matches reference latency/noise semantics. - - Args: - monkeypatch: Pytest helper used to make the random masks deterministic. - - Returns: - None. The assertions validate binary thresholding, latency, and dropout. - """ - sensor_data = { - "contact_right_thumb_elastomer_force": np.array([[0.2, 0.0, 0.0]], dtype=np.float64), - "contact_right_index_elastomer_force": np.array([[0.01, 0.0, 0.0]], dtype=np.float64), - "contact_right_middle_elastomer_force": np.array([[0.5, 0.0, 0.0]], dtype=np.float64), - "contact_right_ring_elastomer_force": np.array([[0.3, 0.0, 0.0]], dtype=np.float64), - "contact_right_pinky_elastomer_force": np.array([[0.9, 0.0, 0.0]], dtype=np.float64), - } - env = _build_fake_tactile_env( - sensor_data, - binary_contact=True, - disable_tactile_ids=[4], - contact_smooth=1.0, - contact_threshold=0.05, - contact_latency=0.5, - contact_sensor_noise=0.01, - last_contacts=np.array([[0.4, 0.7, 0.0, 1.0, 0.2]], dtype=np.float64), - ) - - sampled_masks = iter( - [ - np.array([[0.9, 0.1, 0.9, 0.9, 0.1]], dtype=np.float64), - np.array([[0.2, 0.2, 0.005, 0.9, 0.2]], dtype=np.float64), - ] - ) - - monkeypatch.setattr(np.random, "rand", lambda *shape: next(sampled_masks).copy()) - - tactile = SharpaInhandBaseEnv._compute_tactile_observation(env) - - expected_last = np.array([[1.0, 0.7, 1.0, 1.0, 0.2]], dtype=np.float64) - expected_tactile = np.array([[1.0, 0.7, 0.0, 1.0, 0.2]], dtype=np.float64) - np.testing.assert_allclose(env.last_contacts, expected_last) - np.testing.assert_allclose(tactile, expected_tactile) - np.testing.assert_allclose( - env._prev_tactile_force, - np.array([[0.2, 0.01, 0.5, 0.3, 0.9]], dtype=np.float64), - ) - - -def test_sharpa_tactile_disabled_clears_history() -> None: - """Verify disabled tactile observation clears cached tactile history. - - Args: - None. - - Returns: - None. The assertions validate the disabled-tactile contract. - """ - sensor_data = { - "contact_right_thumb_elastomer_force": np.array([[1.0, 0.0, 0.0]], dtype=np.float64), - "contact_right_index_elastomer_force": np.array([[2.0, 0.0, 0.0]], dtype=np.float64), - "contact_right_middle_elastomer_force": np.array([[3.0, 0.0, 0.0]], dtype=np.float64), - "contact_right_ring_elastomer_force": np.array([[4.0, 0.0, 0.0]], dtype=np.float64), - "contact_right_pinky_elastomer_force": np.array([[5.0, 0.0, 0.0]], dtype=np.float64), - } - env = _build_fake_tactile_env( - sensor_data, - enable_tactile=False, - last_contacts=np.full((1, 5), 3.0, dtype=np.float64), - prev_tactile_force=np.full((1, 5), 4.0, dtype=np.float64), - ) - - tactile = SharpaInhandBaseEnv._compute_tactile_observation(env) - - np.testing.assert_allclose(tactile, np.zeros((1, 5), dtype=np.float64)) - np.testing.assert_allclose(env.last_contacts, np.zeros((1, 5), dtype=np.float64)) - np.testing.assert_allclose(env._prev_tactile_force, np.zeros((1, 5), dtype=np.float64)) - - -@pytest.mark.slow -def test_sharpa_mujoco_reset_applies_friction_randomization() -> None: - _require_mujoco_runtime() - ensure_registries() - - from unilab.base import registry - - num_envs = 4 - cfg, env_cfg_override = _compose_sharpa_mujoco_owner_cfg(num_envs) - with TemporaryDirectory() as tmp_dir: - cache_prefix = Path(tmp_dir) / "sharpa_grasp" - env_cfg_override["grasp_cache_path"] = str(cache_prefix) - for scale_value in env_cfg_override["domain_rand"]["scale_list"]: - cache_file = cache_prefix.parent / f"{cache_prefix.name}_{float(scale_value):g}.npy" - np.save(cache_file, np.zeros((8, 29), dtype=np.float32)) - - env = registry.make( - "SharpaInhandRotation", - num_envs=num_envs, - sim_backend="mujoco", - env_cfg_override=env_cfg_override, - ) - env_obj: Any = env - try: - env_ids = np.arange(num_envs, dtype=np.int32) - _, info = env_obj.reset(env_ids) - - backend: Any = env_obj._backend - pool = backend._pool - geom_friction = np.stack( - [pool.get_field(i, "geom_friction") for i in range(num_envs)], - axis=0, - ).reshape(num_envs, backend.model.ngeom, 3) - friction_scale = np.asarray(info["critic_info"][:, 3], dtype=np.float64) - - assert np.unique(np.round(friction_scale, 6)).size > 1 - assert np.all(friction_scale >= cfg.env.domain_rand.randomize_friction_scale_lower) - assert np.all(friction_scale <= cfg.env.domain_rand.randomize_friction_scale_upper) - - for env_idx in range(num_envs): - scale = friction_scale[env_idx] - for material, base_friction in ( - ("object", cfg.env.domain_rand.object_base_friction), - ("metal", cfg.env.domain_rand.metal_base_friction), - ("elastomer", cfg.env.domain_rand.elastomer_base_friction), - ): - actual = geom_friction[env_idx, env_obj._friction_geom_ids[material]] - expected = env_obj._friction_profile(material, base_friction) * scale - np.testing.assert_allclose(actual, np.broadcast_to(expected, actual.shape)) - finally: - pool = getattr(getattr(env_obj, "_backend", None), "_pool", None) - if pool is not None: - pool.close() - env_obj.close() - - -@pytest.mark.slow -def test_sharpa_mujoco_reset_randomizes_pd_gains_from_xml_defaults() -> None: - """Verify Sharpa reset PD gains scale MuJoCo XML actuator defaults per DOF. - - Args: - None. - - Returns: - None. The assertions validate info buffers and backend reset payloads. - """ - _require_mujoco_runtime() - ensure_registries() - - from unilab.base import registry - - num_envs = 4 - cfg, env_cfg_override = _compose_sharpa_mujoco_owner_cfg(num_envs) - with TemporaryDirectory() as tmp_dir: - cache_prefix = Path(tmp_dir) / "sharpa_grasp" - env_cfg_override["grasp_cache_path"] = str(cache_prefix) - for scale_value in env_cfg_override["domain_rand"]["scale_list"]: - cache_file = cache_prefix.parent / f"{cache_prefix.name}_{float(scale_value):g}.npy" - np.save(cache_file, np.zeros((8, 29), dtype=np.float32)) - - env = registry.make( - "SharpaInhandRotation", - num_envs=num_envs, - sim_backend="mujoco", - env_cfg_override=env_cfg_override, - ) - env_obj: Any = env - try: - env_ids = np.arange(num_envs, dtype=np.int32) - _, info = env_obj.reset(env_ids) - - backend: Any = env_obj._backend - pool = backend._pool - default_kp, default_kd = backend.get_actuator_gains() - default_kp = np.asarray(default_kp[: env_obj._num_action], dtype=np.float64) - default_kd = np.asarray(default_kd[: env_obj._num_action], dtype=np.float64) - info_kp = np.asarray(info["p_gain"], dtype=np.float64) - info_kd = np.asarray(info["d_gain"], dtype=np.float64) - pool_kp = np.stack([pool.get_field(i, "kp") for i in range(num_envs)], axis=0) - pool_kd = np.stack([pool.get_field(i, "kd") for i in range(num_envs)], axis=0) - - assert np.all(default_kp > 0.0) - assert np.all(default_kd > 0.0) - assert info_kp.shape == (num_envs, env_obj._num_action) - assert info_kd.shape == (num_envs, env_obj._num_action) - np.testing.assert_allclose(pool_kp[:, : env_obj._num_action], info_kp) - np.testing.assert_allclose(pool_kd[:, : env_obj._num_action], info_kd) - - kp_scale = info_kp / default_kp[None, :] - kd_scale = info_kd / default_kd[None, :] - assert np.unique(np.round(kp_scale.reshape(-1), 6)).size > 1 - assert np.unique(np.round(kd_scale.reshape(-1), 6)).size > 1 - assert np.all(kp_scale >= cfg.env.domain_rand.randomize_p_gain_scale_lower) - assert np.all(kp_scale <= cfg.env.domain_rand.randomize_p_gain_scale_upper) - assert np.all(kd_scale >= cfg.env.domain_rand.randomize_d_gain_scale_lower) - assert np.all(kd_scale <= cfg.env.domain_rand.randomize_d_gain_scale_upper) - finally: - pool = getattr(getattr(env_obj, "_backend", None), "_pool", None) - if pool is not None: - pool.close() - env_obj.close() - - -@pytest.mark.slow -def test_sharpa_mujoco_interval_force_disturbs_object_velocity() -> None: - """Verify Sharpa force randomization perturbs object linear velocity through DR. - - Args: - None. - - Returns: - None. The assertions validate the interval-randomization force path. - """ - _require_mujoco_runtime() - ensure_registries() - - from unilab.base import registry - - num_envs = 4 - _, env_cfg_override = _compose_sharpa_mujoco_owner_cfg(num_envs) - env_cfg_override["domain_rand"]["force_scale"] = 2.0 - env_cfg_override["domain_rand"]["random_force_prob_scalar"] = 1.0 - with TemporaryDirectory() as tmp_dir: - cache_prefix = Path(tmp_dir) / "sharpa_grasp" - env_cfg_override["grasp_cache_path"] = str(cache_prefix) - for scale_value in env_cfg_override["domain_rand"]["scale_list"]: - cache_file = cache_prefix.parent / f"{cache_prefix.name}_{float(scale_value):g}.npy" - np.save(cache_file, np.zeros((8, 29), dtype=np.float32)) - - env = registry.make( - "SharpaInhandRotation", - num_envs=num_envs, - sim_backend="mujoco", - env_cfg_override=env_cfg_override, - ) - env_obj: Any = env - try: - env.init_state() - assert env_obj._backend.get_dr_capabilities().supports_interval_body_force - before = env_obj._backend._physics_state.copy() - env_obj._dr_manager.apply_interval_randomization_if_due(env_obj.step_counter) - - body_id = int(env_obj._object_body_id) - force_slice = slice(6 * body_id, 6 * body_id + 3) - joint_adr = int(env_obj._backend.model.body_jntadr[body_id]) - dof_adr = int(env_obj._backend.model.jnt_dofadr[joint_adr]) - qvel_slice = slice( - env_obj._backend._idx_qvel + dof_adr, - env_obj._backend._idx_qvel + dof_adr + 3, - ) - - assert np.any(np.linalg.norm(env_obj._random_object_force, axis=1) > 0.0) - np.testing.assert_allclose( - env_obj._backend._pending_xfrc_applied[:, force_slice], - env_obj._random_object_force, - ) - env_obj._backend.step( - np.zeros((num_envs, env_obj._num_action), dtype=env_obj._np_dtype), - nsteps=1, - ) - after = env_obj._backend._physics_state - velocity_delta = after[:, qvel_slice] - before[:, qvel_slice] - assert np.any(np.linalg.norm(velocity_delta, axis=1) > 0.0) - finally: - pool = getattr(getattr(env_obj, "_backend", None), "_pool", None) - if pool is not None: - pool.close() - env_obj.close() - - -@pytest.mark.slow -def test_sharpa_mujoco_interval_force_plan_matches_decay_and_mass_scaled_resample( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """Verify Sharpa force DR matches decay plus mass-scaled Gaussian resampling. - - Args: - monkeypatch: Pytest helper used to make numpy random sampling deterministic. - - Returns: - None. The assertions validate the exact interval body-force payload. - """ - _require_mujoco_runtime() - ensure_registries() - - from unilab.base import registry - - num_envs = 4 - _, env_cfg_override = _compose_sharpa_mujoco_owner_cfg(num_envs) - env_cfg_override["domain_rand"]["force_scale"] = 2.0 - env_cfg_override["domain_rand"]["random_force_prob_scalar"] = 0.5 - env_cfg_override["domain_rand"]["randomize_mass"] = False - with TemporaryDirectory() as tmp_dir: - cache_prefix = Path(tmp_dir) / "sharpa_grasp" - env_cfg_override["grasp_cache_path"] = str(cache_prefix) - for scale_value in env_cfg_override["domain_rand"]["scale_list"]: - cache_file = cache_prefix.parent / f"{cache_prefix.name}_{float(scale_value):g}.npy" - np.save(cache_file, np.zeros((8, 29), dtype=np.float32)) - - env = registry.make( - "SharpaInhandRotation", - num_envs=num_envs, - sim_backend="mujoco", - env_cfg_override=env_cfg_override, - ) - env_obj: Any = env - try: - env.init_state() - env_obj._random_object_force[:] = np.array( - [ - [0.4, -0.2, 0.1], - [-0.3, 0.5, -0.7], - [0.8, -0.6, 0.2], - [0.1, 0.3, -0.4], - ], - dtype=np.float64, - ) - previous_force = env_obj._random_object_force.copy() - sampled_uniform = np.array([0.1, 0.9, 0.2, 0.8], dtype=np.float64) - sampled_gaussian = np.array( - [ - [1.0, -2.0, 0.5], - [-1.5, 0.25, 2.0], - ], - dtype=np.float64, - ) - - def _fake_rand(*shape: int) -> np.ndarray: - assert shape == (num_envs,) - return sampled_uniform.copy() - - def _fake_randn(*shape: int) -> np.ndarray: - assert shape == (2, 3) - return sampled_gaussian.copy() - - monkeypatch.setattr(np.random, "rand", _fake_rand) - monkeypatch.setattr(np.random, "randn", _fake_randn) - - decay = float( - np.power( - env_obj.cfg.domain_rand.force_decay, - env_obj.cfg.ctrl_dt / max(env_obj.cfg.domain_rand.force_decay_interval, 1.0e-8), - ) - ) - object_mass = env_obj._resolve_current_object_mass() - resample_mask = sampled_uniform < float( - env_obj.cfg.domain_rand.random_force_prob_scalar - ) - - plan = SharpaInhandRotationDRProvider().build_interval_randomization_plan( - env_obj, - step_counter=0, - ) - - assert plan is not None - assert plan.body_ids is not None - assert plan.body_force is not None - np.testing.assert_array_equal( - plan.body_ids, - np.asarray([env_obj._object_body_id], dtype=np.int32), - ) - - # Envs below the Bernoulli threshold get a fresh Gaussian force sample; - # the remaining envs keep the decayed previous force. - expected_force = previous_force * decay - expected_force[resample_mask] = ( - sampled_gaussian - * object_mass[resample_mask, None] - * float(env_obj.cfg.domain_rand.force_scale) - ) - np.testing.assert_allclose(env_obj._random_object_force, expected_force) - np.testing.assert_allclose(plan.body_force, expected_force[:, None, :]) - assert plan.body_linear_velocity_delta is None - finally: - pool = getattr(getattr(env_obj, "_backend", None), "_pool", None) - if pool is not None: - pool.close() - env_obj.close() - - -@pytest.mark.slow -def test_sharpa_mujoco_reset_applies_object_mass_and_com_randomization() -> None: - _require_mujoco_runtime() - ensure_registries() - - from unilab.base import registry - - num_envs = 4 - cfg, env_cfg_override = _compose_sharpa_mujoco_owner_cfg(num_envs) - with TemporaryDirectory() as tmp_dir: - cache_prefix = Path(tmp_dir) / "sharpa_grasp" - env_cfg_override["grasp_cache_path"] = str(cache_prefix) - for scale_value in env_cfg_override["domain_rand"]["scale_list"]: - cache_file = cache_prefix.parent / f"{cache_prefix.name}_{float(scale_value):g}.npy" - np.save(cache_file, np.zeros((8, 29), dtype=np.float32)) - - env = registry.make( - "SharpaInhandRotation", - num_envs=num_envs, - sim_backend="mujoco", - env_cfg_override=env_cfg_override, - ) - env_obj: Any = env - try: - env_ids = np.arange(num_envs, dtype=np.int32) - _, info = env_obj.reset(env_ids) - - backend: Any = env_obj._backend - pool = backend._pool - body_mass = np.stack([pool.get_field(i, "body_mass") for i in range(num_envs)], axis=0) - body_ipos = np.stack([pool.get_field(i, "body_ipos") for i in range(num_envs)], axis=0) - body_ipos = body_ipos.reshape(num_envs, backend.model.nbody, 3) - - object_body_id = int(env_obj._object_body_id) - randomized_mass = np.asarray(info["critic_info"][:, 4], dtype=np.float64) - randomized_com = np.asarray(info["critic_info"][:, 5:8], dtype=np.float64) - - assert np.unique(np.round(randomized_mass, 6)).size > 1 - assert np.unique(np.round(randomized_com.reshape(-1), 6)).size > 1 - assert np.all(randomized_mass >= cfg.env.domain_rand.randomize_mass_lower) - assert np.all(randomized_mass <= cfg.env.domain_rand.randomize_mass_upper) - assert np.all(randomized_com >= cfg.env.domain_rand.randomize_com_lower) - assert np.all(randomized_com <= cfg.env.domain_rand.randomize_com_upper) - - np.testing.assert_allclose(body_mass[:, object_body_id], randomized_mass) - np.testing.assert_allclose( - body_ipos[:, object_body_id, :], - env_obj._base_body_ipos[object_body_id][None, :] + randomized_com, - ) - finally: - pool = getattr(getattr(env_obj, "_backend", None), "_pool", None) - if pool is not None: - pool.close() - env_obj.close() diff --git a/tests/scripts/test_check_docs.py b/tests/scripts/test_check_docs.py index 33fa2cdc9..f0de99bca 100644 --- a/tests/scripts/test_check_docs.py +++ b/tests/scripts/test_check_docs.py @@ -2,8 +2,6 @@ from pathlib import Path -from omegaconf import OmegaConf - from tests.scripts import doc_checks @@ -15,30 +13,6 @@ def test_documentation_files_match_current_repo_contracts(): assert warnings == [] -def test_sharpa_domain_randomization_doc_matches_owner_config(): - root = Path(__file__).resolve().parents[2] - doc_path = ( - root - / "docs" - / "sphinx" - / "source" - / "zh_CN" - / "2-user_guide" - / "5-domain_randomization" - / "0-index.md" - ) - content = doc_path.read_text(encoding="utf-8") - - owner_cfg = OmegaConf.load( - root / "src" / "unilab" / "conf" / "ppo" / "task" / "sharpa_inhand" / "mujoco.yaml" - ) - - assert "Sharpa" in content - assert "`geom_size`" in content - if float(owner_cfg.env.domain_rand.force_scale) > 0.0: - assert "`gravity`" in content - - def test_check_training_entrypoint_semantics_flags_issue_204_patterns(): root = Path(__file__).resolve().parents[2] doc_path = root / "README.md" diff --git a/tests/scripts/test_support_matrix.py b/tests/scripts/test_support_matrix.py index 8f0864498..c074bb49e 100644 --- a/tests/scripts/test_support_matrix.py +++ b/tests/scripts/test_support_matrix.py @@ -162,16 +162,7 @@ def test_support_matrix_marks_appo_go1_backends_as_tested(): assert row.cells["motrix"].level == EvidenceLevel.TESTED -def test_support_matrix_marks_sharpa_motrix_phase1_support(): - row = _row("PPO (torch)", "sharpa_inhand") - - assert row.cells["mujoco"].level == EvidenceLevel.TESTED - assert row.cells["motrix"].level == EvidenceLevel.TESTED - - appo_row = _row("APPO (torch)", "sharpa_inhand") - - assert appo_row.cells["mujoco"].level == EvidenceLevel.TESTED - assert appo_row.cells["motrix"].level == EvidenceLevel.TESTED +def test_support_matrix_marks_allegro_appo_backends_as_tested(): allegro_appo_row = _row("APPO (torch)", "allegro_inhand") assert allegro_appo_row.cells["mujoco"].level == EvidenceLevel.TESTED diff --git a/tests/scripts/test_train_script_configs.py b/tests/scripts/test_train_script_configs.py index cbb68465b..229e1c9bd 100644 --- a/tests/scripts/test_train_script_configs.py +++ b/tests/scripts/test_train_script_configs.py @@ -18,23 +18,6 @@ ROOT_DIR = Path(__file__).resolve().parents[2] - -def _write_sharpa_smoke_cache(cache_prefix, scale_values: list[float]) -> None: - from unilab.tasks.manipulation.sharpa_inhand.base import ( - SOURCE_DEFAULT_HAND_JOINT_POS_DEG, - resolve_grasp_cache_file, - ) - - hand_qpos = np.deg2rad(np.asarray(SOURCE_DEFAULT_HAND_JOINT_POS_DEG, dtype=np.float64)) - object_pose = np.asarray([-0.09559, -0.00517, 0.61906, 1.0, 0.0, 0.0, 0.0]) - cache = np.broadcast_to(np.concatenate([hand_qpos, object_pose]), (32, 29)).copy() - for scale_value in scale_values: - np.save( - resolve_grasp_cache_file(str(cache_prefix), float(scale_value)), - cache.astype(np.float32), - ) - - APPO_MUJOCO_SMOKE_TASKS = [ "go1_joystick_flat/mujoco", "go2_joystick_flat/mujoco", @@ -140,43 +123,6 @@ def test_offpolicy_task_configs_load(algo, task): assert result.returncode == 0, f"Off-policy {algo} {task} failed:\n{result.stderr}" -@pytest.mark.slow -def test_ppo_sharpa_motrix_one_iteration_training_smoke(tmp_path): - """Sharpa Motrix owner can run a minimal RSL-RL learn loop.""" - pytest.importorskip("motrixsim", reason="motrixsim not installed") - cache_prefix = tmp_path / "sharpa_grasp" - _write_sharpa_smoke_cache( - cache_prefix, - [0.8, 0.9, 1.0, 1.1, 1.2, 1.3, 1.4, 1.5], - ) - - result = subprocess.run( - [ - sys.executable, - "src/unilab/scripts/train_rsl_rl.py", - "task=sharpa_inhand/motrix", - "algo.num_envs=16", - "algo.num_steps_per_env=2", - "algo.max_iterations=1", - "algo.save_interval=100", - "training.no_play=true", - f"training.log_root={tmp_path / 'logs'}", - f"env.grasp_cache_path={cache_prefix}", - ], - capture_output=True, - text=True, - timeout=180, - ) - - assert result.returncode == 0, ( - "Sharpa Motrix PPO one-iteration smoke failed:\n" - f"stdout:\n{result.stdout}\n" - f"stderr:\n{result.stderr}" - ) - assert "Learning iteration 0/1" in result.stdout - assert "reward/total" in result.stdout - - @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires >=2 CUDA devices") @pytest.mark.slow @pytest.mark.parametrize( diff --git a/tests/scripts/test_train_scripts.py b/tests/scripts/test_train_scripts.py index 9f4932ecb..1ed37c958 100644 --- a/tests/scripts/test_train_scripts.py +++ b/tests/scripts/test_train_scripts.py @@ -92,7 +92,6 @@ def test_analyze_offpolicy_trace_reports_training_e2e(tmp_path, capsys): # Helpers # --------------------------------------------------------------------------- - try: import mujoco # noqa: F401 @@ -100,7 +99,6 @@ def test_analyze_offpolicy_trace_reports_training_e2e(tmp_path, capsys): except ImportError: _HAS_MUJOCO = False - # --------------------------------------------------------------------------- # train_sac.py / train_td3.py / train_flashsac.py — Hydra config defaults # (composed from the per-algo trees conf/sac, conf/td3, conf/flashsac) @@ -129,20 +127,6 @@ def _appo_cfg(overrides=None): return compose("config", overrides=_normalize_overrides(overrides)) -def _hora_distill_cfg(overrides=None): - """Compose the HORA distillation Hydra config. - - Args: - overrides: Optional Hydra override strings to apply during composition. - - Returns: - The composed HORA distillation config. - """ - GlobalHydra.instance().clear() - with initialize_config_dir(config_dir=str(_CONF_DIR / "hora_distill"), version_base="1.3"): - return compose("config", overrides=overrides or []) - - def _train_rsl_rl(monkeypatch: pytest.MonkeyPatch): import types @@ -163,18 +147,6 @@ def _train_appo(): return _load_script("train_appo") -def _train_hora_distill(): - """Load the HORA distillation entrypoint module. - - Args: - None. - - Returns: - The loaded ``scripts/train_hora_distill.py`` module. - """ - return _load_script("train_hora_distill") - - def test_offpolicy_hydra_default_algo(): cfg = _offpolicy_cfg() assert cfg.algo.algo == "sac" @@ -329,266 +301,6 @@ def test_offpolicy_hydra_algo_td3(): assert cfg.algo.algo == "td3" -def test_hora_distill_run_config_records_hardware(tmp_path, monkeypatch): - mod = _train_hora_distill() - hardware = { - "platform": "test-platform", - "chip": "test-cpu", - "cpu_total_cores": "8", - "gpu_name": "test-gpu", - "memory": "32 GB", - } - monkeypatch.setattr(mod, "get_device_info_dict", lambda: hardware) - cfg = OmegaConf.create({"training": {"task_name": "Task", "sim_backend": "mujoco"}}) - - mod._write_distill_run_config( - tmp_path, - cfg=cfg, - teacher_metadata={"checkpoint_path": "teacher.pt"}, - ) - - payload = json.loads((tmp_path / "distill_run_config.json").read_text(encoding="utf-8")) - assert payload["run"]["hardware"] == hardware - - -def test_hora_distill_task_owner_overrides_root_config_defaults(): - mod = _train_hora_distill() - root_cfg = OmegaConf.load(_CONF_DIR / "hora_distill" / "config.yaml") - cfg = mod._apply_teacher_defaults(_hora_distill_cfg(["task=sharpa_inhand/mujoco"])) - - assert root_cfg.algo.num_envs == 4096 - assert root_cfg.algo.save_interval_steps == 100000000 - assert cfg.algo.num_envs == 16384 - assert cfg.algo.save_interval_steps == 10000000 - - -def test_hora_distill_sharpa_appo_student_owner_selects_nodr_demo_profile(): - mod = _train_hora_distill() - cfg = mod._apply_teacher_defaults(_hora_distill_cfg(["task=sharpa_inhand/mujoco_nodr"])) - - assert cfg.teacher.algo_family == "appo" - assert cfg.teacher.task == "sharpa_inhand/mujoco_hora" - assert cfg.training.task_name == "SharpaInhandRotation" - assert cfg.training.sim_backend == "mujoco" - assert cfg.interactive.action_mode == "policy" - assert cfg.interactive.policy_obs_mode == "actor" - assert cfg.env.post_step_forward_sensor is True - assert cfg.env.domain_rand.scale_list == [0.8, 0.9, 1.0, 1.1, 1.2, 1.3, 1.4, 1.5] - assert cfg.env.domain_rand.randomize_gravity is False - assert cfg.env.domain_rand.randomize_gravity_direction is False - assert cfg.env.domain_rand.randomize_pd_gains is False - assert cfg.env.domain_rand.randomize_friction is False - assert cfg.env.domain_rand.randomize_com is False - assert cfg.env.domain_rand.randomize_mass is False - assert cfg.env.domain_rand.force_scale == pytest.approx(0.0) - assert cfg.env.domain_rand.random_force_prob_scalar == pytest.approx(0.0) - assert cfg.env.domain_rand.joint_noise_scale == pytest.approx(0.0) - assert cfg.env.domain_rand.contact_latency == pytest.approx(0.0) - assert cfg.env.domain_rand.contact_sensor_noise == pytest.approx(0.0) - assert cfg.algo.model.priv_info_embed_dim == 9 - assert cfg.algo.model.priv_mlp_hidden_dims == [256, 128, 9] - - -def test_hora_distill_runtime_checkpoint_records_model_only(): - mod = _train_hora_distill() - cfg = OmegaConf.create( - { - "training": { - "task_name": "OwnerTask", - "sim_backend": "mujoco", - "cam_distance": 1.5, - }, - "env": { - "post_step_forward_sensor": True, - "domain_rand": {"force_scale": 1.2}, - }, - "reward": {"scales": {"rotate": 2.5}}, - "algo": {"model": {"hidden_dims": [512, 256, 128]}}, - } - ) - - runtime = OmegaConf.to_container(mod._resolved_distill_runtime_cfg(cfg), resolve=True) - - assert runtime == {"algo": {"model": {"hidden_dims": [512, 256, 128]}}} - - -def test_hora_distill_checkpoint_runtime_only_restores_model_structure(): - from uni_rl.algos.hora.distill import cfg_with_checkpoint_runtime - - from unilab.training.hora_distill_config import apply_teacher_defaults - - cfg = _hora_distill_cfg(["task=sharpa_inhand/mujoco_nodr"]) - checkpoint = { - "distill_runtime_cfg": { - "training": { - "task_name": "CheckpointTask", - "sim_backend": "motrix", - "render_spacing": 99.0, - }, - "reward": {"scales": {"rotate": 999.0}}, - "env": { - "post_step_forward_sensor": False, - "domain_rand": { - "scale_list": [9.9], - "randomize_mass": True, - "force_scale": 99.0, - }, - }, - "algo": { - "model": { - "hidden_dims": [32, 16], - "priv_info_embed_dim": 7, - "priv_mlp_hidden_dims": [11, 7], - } - }, - } - } - - restored = cfg_with_checkpoint_runtime(apply_teacher_defaults(cfg), checkpoint) - - assert restored.training.task_name == "SharpaInhandRotation" - assert restored.training.sim_backend == "mujoco" - assert restored.training.render_spacing == pytest.approx(0.5) - assert restored.reward.scales.rotate != pytest.approx(999.0) - assert restored.env.post_step_forward_sensor is True - assert restored.env.domain_rand.scale_list == [0.8, 0.9, 1.0, 1.1, 1.2, 1.3, 1.4, 1.5] - assert restored.env.domain_rand.randomize_mass is False - assert restored.env.domain_rand.force_scale == pytest.approx(0.0) - assert restored.algo.model.hidden_dims == [32, 16] - assert restored.algo.model.priv_info_embed_dim == 7 - assert restored.algo.model.priv_mlp_hidden_dims == [11, 7] - - -@pytest.mark.parametrize( - ("teacher_algo_family", "checkpoint_model"), - [ - ("ppo", {"hidden_dims": [512, 256, 128], "activation": "elu"}), - ("appo", {"hidden_dims": [512, 256, 128], "activation": "elu"}), - ( - "sac", - { - "teacher_arch": "hora_sac", - "actor_hidden_dim": 512, - "use_layer_norm": True, - }, - ), - ], -) -def test_hora_distill_checkpoint_runtime_only_overrides_model_side( - monkeypatch: pytest.MonkeyPatch, - teacher_algo_family: str, - checkpoint_model: dict[str, Any], -): - from uni_rl.algos.hora.distill import cfg_with_checkpoint_runtime - - from unilab.training import hora_distill_config as distill_config - - owner_cfg = OmegaConf.create( - { - "teacher": {"algo_family": teacher_algo_family}, - "training": { - "task_name": "OwnerTask", - "sim_backend": "mujoco", - "cam_distance": 1.5, - }, - "env": { - "post_step_forward_sensor": False, - "domain_rand": {"force_scale": 1.2, "randomize_mass": False}, - }, - "reward": {"scales": {"rotate": 2.5}}, - "algo": {"model": {"owner_model": True}}, - } - ) - checkpoint = { - "teacher_algo_family": teacher_algo_family, - "distill_runtime_cfg": { - "training": { - "task_name": "CheckpointTask", - "sim_backend": "mujoco", - "cam_distance": 9.0, - }, - "env": { - "post_step_forward_sensor": True, - "domain_rand": {"force_scale": 9.0, "randomize_mass": True}, - }, - "reward": {"scales": {"rotate": 99.0}}, - "algo": {"model": checkpoint_model}, - }, - } - - monkeypatch.setattr(distill_config, "apply_teacher_defaults", lambda cfg: owner_cfg) - - # uni_rl's cfg_with_checkpoint_runtime only restores model-side fields; - # composing teacher-owner defaults first is the caller's job (issue #1480). - effective_cfg = cfg_with_checkpoint_runtime( - distill_config.apply_teacher_defaults(OmegaConf.create({})), checkpoint - ) - - assert effective_cfg.training.task_name == "OwnerTask" - assert effective_cfg.training.cam_distance == pytest.approx(1.5) - assert effective_cfg.env.post_step_forward_sensor is False - assert effective_cfg.env.domain_rand.force_scale == pytest.approx(1.2) - assert effective_cfg.env.domain_rand.randomize_mass is False - assert effective_cfg.reward.scales.rotate == pytest.approx(2.5) - assert OmegaConf.to_container(effective_cfg.algo.model, resolve=True) == checkpoint_model - - -@pytest.mark.parametrize("teacher_algo_family", ["ppo", "appo", "sac"]) -def test_hora_distill_teacher_owner_defaults_support_ppo_appo_and_sac( - teacher_algo_family: str, -): - mod = _train_hora_distill() - teacher_task = "sharpa_inhand/mujoco_hora" - cfg = mod._apply_teacher_defaults( - _hora_distill_cfg( - [ - "task=sharpa_inhand/mujoco", - f"teacher.algo_family={teacher_algo_family}", - f"teacher.task={teacher_task}", - ] - ) - ) - - assert cfg.training.task_name == "SharpaInhandRotation" - assert cfg.training.sim_backend == "mujoco" - assert cfg.algo.model.priv_info_embed_dim == 9 - assert cfg.algo.model.priv_mlp_hidden_dims == [256, 128, 9] - if teacher_algo_family == "sac": - assert cfg.algo.model.teacher_arch - assert cfg.algo.model.actor_hidden_dim is not None - - -def test_hora_distill_sac_teacher_requires_hora_sac_runtime(): - mod = _train_hora_distill() - - with pytest.raises(ValueError, match="runtime_impl='hora_sac'"): - mod._apply_teacher_defaults( - _hora_distill_cfg( - [ - "task=sharpa_inhand/mujoco", - "teacher.algo_family=sac", - "teacher.task=g1_walk_flat/mujoco", - ] - ) - ) - - -@pytest.mark.parametrize("teacher_algo_family", ["ppo", "appo"]) -def test_hora_distill_teacher_run_slug_omits_teacher_run_name(teacher_algo_family: str): - mod = _train_hora_distill() - cfg = OmegaConf.create({"teacher": {"task": "sharpa_inhand/mujoco"}}) - teacher_checkpoint = Path("/tmp") / "2026-04-22_13-26-45_mujoco" / "model_10000.pt" - - metadata = mod._teacher_run_metadata( - cfg, - teacher_algo_family=teacher_algo_family, - teacher_checkpoint=teacher_checkpoint, - ) - - assert metadata["run_name"] == "2026-04-22_13-26-45_mujoco" - assert metadata["run_slug"] == f"teacher-{teacher_algo_family}" - - def test_offpolicy_go1_motrix_task_is_not_configured(): """SAC has no Go1 Motrix owner config; use PPO for Go1 joystick tasks.""" from hydra.errors import MissingConfigException @@ -934,7 +646,7 @@ def test_build_ppo_env_cfg_override_allegro_grasp_mujoco( monkeypatch: pytest.MonkeyPatch, ): mod = _train_rsl_rl(monkeypatch) - cfg = _ppo_cfg(["task=allegro_inhand_grasp/mujoco"]) + cfg = _ppo_cfg(["task=allegro_inhand_grasp/mujoco", "+env.grasp_collection_target=1"]) env_cfg_override = mod.build_ppo_env_cfg_override(cfg) @@ -978,46 +690,6 @@ def test_build_ppo_env_cfg_override_allegro_grasp_cli_override_wins( assert env_cfg_override["rewards"]["rotate"]["weight"] == pytest.approx(0.3) -def test_build_ppo_env_cfg_override_sharpa_grasp_cli_override_wins( - monkeypatch: pytest.MonkeyPatch, -): - mod = _train_rsl_rl(monkeypatch) - cfg = _ppo_cfg( - [ - "task=sharpa_inhand_grasp/mujoco", - "algo.max_iterations=1", - "env.grasp_collection_target=128", - "reward.scales.rotate=0.3", - ] - ) - - env_cfg_override = mod.build_ppo_env_cfg_override(cfg) - - assert cfg.algo.max_iterations == 1 - assert env_cfg_override["grasp_collection_target"] == 128 - assert env_cfg_override["reward_config"]["scales"]["rotate"] == pytest.approx(0.3) - - -def test_build_ppo_env_cfg_override_sharpa_grasp_motrix_owner( - monkeypatch: pytest.MonkeyPatch, -): - mod = _train_rsl_rl(monkeypatch) - cfg = _ppo_cfg( - [ - "task=sharpa_inhand_grasp/motrix", - "algo.max_iterations=1", - "env.grasp_collection_target=128", - ] - ) - - env_cfg_override = mod.build_ppo_env_cfg_override(cfg) - - assert cfg.training.task_name == "SharpaInhandRotationGrasp" - assert cfg.training.sim_backend == "motrix" - assert env_cfg_override["grasp_collection_target"] == 128 - assert env_cfg_override["domain_rand"]["scale_list"] == [0.8] - - @pytest.mark.parametrize("std_type", ["scalar", "log"]) @pytest.mark.parametrize("state_dependent", [False, True]) def test_rsl_action_std_logging_patch_delegates_with_detached_clone( @@ -1088,7 +760,7 @@ def _build_rsl_lifecycle_case( ) cfg = _ppo_cfg( [ - "task=sharpa_inhand_grasp/mujoco", + "task=allegro_inhand_grasp/mujoco", f"training.log_dir={tmp_path}", "training.logger=none", "training.nan_guard.enabled=false", @@ -1305,7 +977,7 @@ def test_train_rsl_rl_grasp_collection_rejects_multi_rank_before_launch( world_size: int, ) -> None: mod = _train_rsl_rl(monkeypatch) - cfg = _ppo_cfg(["task=sharpa_inhand_grasp/mujoco"]) + cfg = _ppo_cfg(["task=allegro_inhand_grasp/mujoco", "+env.grasp_collection_target=1"]) monkeypatch.setattr(mod, "resolve_dp_topology", lambda _devices: devices) monkeypatch.setattr(mod, "current_torch_distributed_rank", lambda: 0) monkeypatch.setattr(mod, "current_torch_distributed_local_rank", lambda: 0) @@ -1528,18 +1200,6 @@ def test_g1_motion_tracking_appo_task_exposes_final_reward(): assert cfg.reward.motion_body_pos.weight == pytest.approx(1.0) -def test_sharpa_appo_motrix_owner_uses_backend_specific_overrides(): - cfg = _appo_cfg(["task=sharpa_inhand/motrix"]) - - assert cfg.training.task_name == "SharpaInhandRotation" - assert cfg.training.sim_backend == "motrix" - assert cfg.algo.num_envs == 2048 - assert cfg.env.sim_dt == pytest.approx(0.01) - assert cfg.env.domain_rand.randomize_gravity is True - assert cfg.env.domain_rand.randomize_gravity_direction is False - assert cfg.env.domain_rand.randomize_pd_gains is True - - # --------------------------------------------------------------------------- # train_appo.py — motrix runner / play helpers # --------------------------------------------------------------------------- @@ -1965,26 +1625,6 @@ def test_offpolicy_extract_play_obs_uses_obs_group_only(): assert np.allclose(play_obs, 1.0) -def test_offpolicy_play_actor_spec_uses_hora_sac_runtime(): - from unilab.visualization.interactive_playback import resolve_play_actor_spec - - cfg = _offpolicy_cfg( - [ - "task=sharpa_inhand/mujoco_hora", - ] - ) - - actor_algo_type, actor_kwargs = resolve_play_actor_spec( - "sac", - cfg, - obs_dim=4, - critic_obs_dim=6, - ) - - assert actor_algo_type == "hora_sac" - assert actor_kwargs["priv_info_dim"] == 2 - - def test_offpolicy_play_actor_spec_keeps_standard_sac_and_flashsac(): from unilab.visualization.interactive_playback import resolve_play_actor_spec @@ -2256,142 +1896,6 @@ def run_playback_mode(self, **kwargs): assert not (run_dir / "policy.onnx").exists() -def test_play_offpolicy_uses_hora_sac_actor_and_priv_info( - monkeypatch: pytest.MonkeyPatch, tmp_path: Path -): - import torch - - mod = _offpolicy() - cfg = _offpolicy_cfg( - [ - "task=sharpa_inhand/mujoco_hora", - "training.play_only=true", - "training.play_render_mode=record", - "training.export_onnx=false", - "training.play_env_num=2", - ] - ) - run_dir = tmp_path / "run" - run_dir.mkdir() - checkpoint = run_dir / "model_5000.pt" - torch.save({"actor": {}}, checkpoint) - - captured: dict[str, Any] = {} - reset_priv = np.array([[4.0, 5.0], [6.0, 7.0]], dtype=np.float32) - step_priv = np.array([[8.0, 9.0], [10.0, 11.0]], dtype=np.float32) - - class FakeHoraActor: - def eval(self): - return self - - def load_state_dict(self, state_dict): - captured["loaded_state_dict"] = state_dict - - def explore(self, obs, priv_info, deterministic=True): - captured["obs_shape"] = tuple(obs.shape) - captured["priv_info"] = priv_info.detach().cpu().numpy() - captured["deterministic"] = deterministic - return torch.zeros((obs.shape[0], 2), dtype=obs.dtype, device=obs.device) - - class FakeEnv: - def __init__(self): - self.obs_groups_spec = {"obs": 3, "critic": 5} - self.action_space = type("ActionSpace", (), {"shape": (2,)})() - self.state = None - - def init_state(self): - self.state = type( - "State", - (), - { - "obs": { - "obs": np.zeros((cfg.training.play_env_num, 3), dtype=np.float32), - "critic": np.zeros((cfg.training.play_env_num, 5), dtype=np.float32), - }, - "info": {"critic_info": reset_priv}, - }, - )() - - def reset(self, env_ids): - batch = len(env_ids) - return ( - { - "obs": np.zeros((batch, 3), dtype=np.float32), - "critic": np.concatenate( - [np.zeros((batch, 3), dtype=np.float32), reset_priv], - axis=1, - ), - }, - {"critic_info": reset_priv}, - ) - - def step(self, actions): - batch = actions.shape[0] - captured["actions_shape"] = actions.shape - self.state = type( - "State", - (), - { - "obs": { - "obs": np.ones((batch, 3), dtype=np.float32), - "critic": np.concatenate( - [np.ones((batch, 3), dtype=np.float32), step_priv], - axis=1, - ), - }, - "info": {"critic_info": step_priv}, - }, - )() - return self.state - - def run_playback_mode(self, **kwargs): - init_obs = kwargs["initialize"]() - captured["init_obs_shape"] = init_obs.shape - next_obs = kwargs["step"](init_obs) - captured["next_obs_shape"] = next_obs.shape - return str(kwargs["output_video"]) - - monkeypatch.setattr(mod, "build_offpolicy_env_cfg_override", lambda algo_name, cfg: {}) - monkeypatch.setattr(mod, "default_device", lambda torch_module, preferred=None: "cpu") - monkeypatch.setattr(mod, "create_env", lambda *args, **kwargs: FakeEnv()) - monkeypatch.setattr( - mod, - "resolve_checkpoint_path", - lambda *args, **kwargs: (str(checkpoint), str(run_dir)), - ) - - import unilab.utils.checkpoint as checkpoint_utils - - monkeypatch.setattr( - checkpoint_utils, - "resolve_offpolicy_checkpoint_path", - lambda *args, **kwargs: (str(checkpoint), str(run_dir)), - ) - - import uni_rl.algos.common.actor_factory as actor_factory - - def fake_build_actor(algo_type, obs_dim, action_dim, hidden_dim, use_layer_norm, device, **kw): - captured["build_actor"] = (algo_type, obs_dim, action_dim, kw) - return FakeHoraActor() - - monkeypatch.setattr(actor_factory, "build_actor", fake_build_actor) - - result = mod.play_offpolicy("sac", cfg) - - assert result == str(run_dir / "play_video.mp4") - assert captured["build_actor"][0] == "hora_sac" - assert captured["build_actor"][1:3] == (3, 2) - assert captured["build_actor"][3]["priv_info_dim"] == 2 - assert captured["loaded_state_dict"] == {} - assert captured["actions_shape"] == (cfg.training.play_env_num, 2) - assert captured["init_obs_shape"] == (cfg.training.play_env_num, 3) - assert captured["next_obs_shape"] == (cfg.training.play_env_num, 3) - assert captured["obs_shape"] == (cfg.training.play_env_num, 3) - np.testing.assert_allclose(captured["priv_info"], reset_priv) - assert captured["deterministic"] is True - assert not (run_dir / "policy.onnx").exists() - - # --------------------------------------------------------------------------- # play_interactive.py — resolve_checkpoint() # --------------------------------------------------------------------------- @@ -2601,69 +2105,6 @@ def reset(self, env_indices): assert wrapper.num_privileged_obs == 4 -def test_play_wrapper_preserves_hora_priv_info_and_proprio_history(): - import numpy as np - from uni_rl.algos.hora.rsl_rl import HoraRslRlVecEnvWrapper - - class FakeEnv: - def __init__(self): - self.num_envs = 1 - self.state = type( - "State", - (), - { - "obs": { - "obs": np.array([[1.0, 2.0, 3.0]], dtype=np.float32), - "proprio": np.array([[6.0, 7.0]], dtype=np.float32), - "critic": np.array([[1.0, 2.0, 3.0, 4.0, 5.0]], dtype=np.float32), - }, - "info": { - "critic_info": np.array([[4.0, 5.0]], dtype=np.float32), - "proprio_hist": np.array( - [[[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]], - dtype=np.float32, - ), - }, - }, - )() - self.cfg = type("Cfg", (), {"max_episode_seconds": 10.0, "ctrl_dt": 0.02})() - self.observation_space = type("Space", (), {"shape": (5,)})() - self.action_space = type("Space", (), {"shape": (2,)})() - self.obs_groups_spec = {"obs": 3, "proprio": 2, "critic": 5} - - def init_state(self): - pass - - def reset(self, env_indices): - del env_indices - return ( - cast(dict[str, np.ndarray], getattr(self.state, "obs")), - cast(dict[str, np.ndarray], getattr(self.state, "info")), - ) - - assert "reset" not in HoraRslRlVecEnvWrapper.__dict__ - assert "get_observations" not in HoraRslRlVecEnvWrapper.__dict__ - - wrapper = HoraRslRlVecEnvWrapper(FakeEnv(), device="cpu", policy_obs_mode="flat") - reset_obs_td, reset_info = wrapper.reset() - current_obs_td = wrapper.get_observations() - - assert reset_info is wrapper.env.state.info - for obs_td in (reset_obs_td, current_obs_td): - np.testing.assert_allclose( - obs_td["policy"].cpu().numpy(), - np.array([[1.0, 2.0, 3.0, 6.0, 7.0]], dtype=np.float32), - ) - np.testing.assert_allclose( - obs_td["priv_info"].cpu().numpy(), - np.array([[4.0, 5.0]], dtype=np.float32), - ) - np.testing.assert_allclose( - obs_td["proprio_hist"].cpu().numpy(), - np.array([[[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]], dtype=np.float32), - ) - - def test_play_wrapper_step_exports_timeout_bootstrap_obs(): import torch from uni_rl.algos.rsl_rl import RslRlVecEnvWrapper @@ -2720,76 +2161,6 @@ def step(self, actions): ) -def test_play_wrapper_timeout_bootstrap_preserves_hora_priv_info(): - import torch - from uni_rl.algos.hora.rsl_rl import HoraRslRlVecEnvWrapper - - class FakeEnv: - def __init__(self): - self.num_envs = 1 - self.cfg = type("Cfg", (), {"max_episode_seconds": 10.0, "ctrl_dt": 0.02})() - self.observation_space = type("Space", (), {"shape": (5,)})() - self.action_space = type("Space", (), {"shape": (2,)})() - self.obs_groups_spec = {"obs": 3, "critic": 5} - self.state = type( - "State", - (), - { - "obs": { - "obs": np.zeros((1, 3), dtype=np.float32), - "critic": np.zeros((1, 5), dtype=np.float32), - }, - "info": { - "critic_info": np.zeros((1, 2), dtype=np.float32), - "proprio_hist": np.zeros((1, 2, 3), dtype=np.float32), - }, - }, - )() - - def init_state(self): - pass - - def reset(self, env_indices): - del env_indices - return cast(dict[str, np.ndarray], getattr(self.state, "obs")), cast( - dict[str, np.ndarray], getattr(self.state, "info") - ) - - def step(self, actions): - del actions - return type( - "StepState", - (), - { - "obs": {"obs": np.array([[1.0, 2.0, 3.0]], dtype=np.float32)}, - "reward": np.array([1.0], dtype=np.float32), - "terminated": np.array([True]), - "truncated": np.array([True]), - "final_observation": { - "obs": np.array([[7.0, 8.0, 9.0]], dtype=np.float32), - "critic": np.array([[7.0, 8.0, 9.0, 4.0, 5.0]], dtype=np.float32), - }, - "info": { - "final_observation": { - "obs": np.array([[7.0, 8.0, 9.0]], dtype=np.float32), - "critic": np.array([[7.0, 8.0, 9.0, 4.0, 5.0]], dtype=np.float32), - }, - "critic_info": np.array([[0.0, 0.0]], dtype=np.float32), - "proprio_hist": np.zeros((1, 2, 3), dtype=np.float32), - }, - }, - )() - - wrapper = HoraRslRlVecEnvWrapper(FakeEnv(), device="cpu", policy_obs_mode="flat") - - _, _, _, infos = wrapper.step(torch.zeros((1, 2))) - - np.testing.assert_allclose( - infos["time_out_bootstrap_obs"]["priv_info"].cpu().numpy(), - np.array([[4.0, 5.0]], dtype=np.float32), - ) - - # --------------------------------------------------------------------------- # Issue #168: Unified log directory and load_run resolution # --------------------------------------------------------------------------- @@ -3168,25 +2539,25 @@ def test_play_interactive_parses_explicit_cli(): mod = _play_interactive() parsed = mod._parse_interactive_cli( - ["--algo", "hora_distill", "--task", "sharpa_inhand", "--sim", "mujoco_nodr"] + ["--algo", "ppo", "--task", "go1_joystick_flat", "--sim", "mujoco"] ) - assert parsed.algo == "hora_distill" - assert parsed.task == "sharpa_inhand" - assert parsed.sim == "mujoco_nodr" - assert parsed.overrides == ["task=sharpa_inhand/mujoco_nodr"] + assert parsed.algo == "ppo" + assert parsed.task == "go1_joystick_flat" + assert parsed.sim == "mujoco" + assert parsed.overrides == ["task=go1_joystick_flat/mujoco"] -@pytest.mark.parametrize("algo", ["appo", "sac", "td3", "hora_distill"]) +@pytest.mark.parametrize("algo", ["appo", "sac", "td3"]) def test_play_interactive_parses_feature_algo_flags(algo: str): mod = _play_interactive() parsed = mod._parse_interactive_cli( - [f"--algo={algo}", "--task", "sharpa_inhand", "--sim", "mujoco_hora"] + [f"--algo={algo}", "--task", "allegro_inhand", "--sim", "mujoco"] ) assert parsed.algo == algo - assert parsed.overrides == ["task=sharpa_inhand/mujoco_hora"] + assert parsed.overrides == ["task=allegro_inhand/mujoco"] def test_play_interactive_cli_respects_owner_action_mode_and_user_override(): @@ -3232,20 +2603,14 @@ def test_play_interactive_dynamic_compose_supports_algo_roots(): mod = _play_interactive() ppo_cfg = mod._compose_interactive_config("ppo", ["task=go1_joystick_flat/mujoco"]) - appo_cfg = mod._compose_interactive_config("appo", ["task=sharpa_inhand/mujoco_hora"]) - sac_cfg = mod._compose_interactive_config("sac", ["task=sharpa_inhand/mujoco_hora"]) + appo_cfg = mod._compose_interactive_config("appo", ["task=allegro_inhand/mujoco"]) + sac_cfg = mod._compose_interactive_config("sac", ["task=g1_walk_flat/mujoco"]) td3_cfg = mod._compose_interactive_config("td3", ["task=g1_walk_flat/mujoco"]) - distill_cfg = mod._compose_interactive_config("hora_distill", ["task=sharpa_inhand/mujoco"]) assert ppo_cfg.algo.algo == "ppo" - assert appo_cfg.algo.runtime_impl == "hora_appo" - assert appo_cfg.interactive.action_mode == "policy" + assert appo_cfg.algo.algo == "appo" assert sac_cfg.algo.algo == "sac" - assert sac_cfg.algo.runtime_impl == "hora_sac" - assert sac_cfg.interactive.policy_obs_mode == "actor" assert td3_cfg.algo.algo == "td3" - assert distill_cfg.algo.algo_log_name == "hora_distill" - assert distill_cfg.interactive.action_mode == "policy" def test_play_interactive_sac_overrides_pass_through(): @@ -3253,11 +2618,11 @@ def test_play_interactive_sac_overrides_pass_through(): overrides = mod._normalize_interactive_overrides( "sac", - ["task=sharpa_inhand/mujoco_hora", "algo.load_run=my_run"], + ["task=g1_walk_flat/mujoco", "algo.load_run=my_run"], ) assert overrides == [ - "task=sharpa_inhand/mujoco_hora", + "task=g1_walk_flat/mujoco", "algo.load_run=my_run", ] diff --git a/tests/tasks/test_legacy_task_compatibility.py b/tests/tasks/test_legacy_task_compatibility.py index 68e2e0803..196862aea 100644 --- a/tests/tasks/test_legacy_task_compatibility.py +++ b/tests/tasks/test_legacy_task_compatibility.py @@ -8,7 +8,6 @@ import numpy as np import pytest -from unilab.base import registry from unilab.base.base import ABEnv, EnvCfg from unilab.base.np_env import NpEnv, NpEnvState from unilab.tasks.compatibility import ( @@ -109,9 +108,9 @@ def factory(cfg: EnvCfg, *, num_envs: int, backend_type: str) -> ABEnv: called = True return _uninitialized_np_env() - adapter = adapt_legacy_factory(factory, task_family="Sharpa", reason="migration seam") + adapter = adapt_legacy_factory(factory, task_family="CustomLegacyTask", reason="migration seam") - with pytest.raises(TypeError, match=r"Sharpa.*expected EnvCfg.*dict"): + with pytest.raises(TypeError, match=r"CustomLegacyTask.*expected EnvCfg.*dict"): adapter({}) # type: ignore[arg-type] assert called is False @@ -147,7 +146,7 @@ def test_factory_exception_propagates_without_fallback() -> None: def factory(cfg: EnvCfg, *, num_envs: int, backend_type: str) -> ABEnv: raise failure - adapter = adapt_legacy_factory(factory, task_family="Sharpa", reason="migration seam") + adapter = adapt_legacy_factory(factory, task_family="CustomLegacyTask", reason="migration seam") with pytest.raises(RuntimeError) as exc_info: adapter(_Cfg()) @@ -157,7 +156,7 @@ def factory(cfg: EnvCfg, *, num_envs: int, backend_type: str) -> ABEnv: def test_unsupported_metadata_is_explicit_and_does_not_create_a_factory() -> None: compatibility = unsupported_legacy_task( - task_family="Sharpa foreign lifecycle", + task_family="CustomLegacyTask foreign lifecycle", reason="only the existing NpEnv lifecycle is admitted", ) @@ -168,33 +167,10 @@ def test_unsupported_metadata_is_explicit_and_does_not_create_a_factory() -> Non LegacyFactoryAdapter(lambda cfg, **kwargs: _uninitialized_np_env(), compatibility) -@pytest.mark.parametrize(("task_family", "reason"), (("", "reason"), ("Sharpa", ""))) +@pytest.mark.parametrize(("task_family", "reason"), (("", "reason"), ("CustomLegacyTask", ""))) def test_compatibility_metadata_requires_stable_family_and_reason( task_family: str, reason: str, ) -> None: with pytest.raises(ValueError, match="must be non-empty"): unsupported_legacy_task(task_family=task_family, reason=reason) - - -@pytest.mark.parametrize( - ("task_name", "family", "backends"), - ( - ("SharpaInhandRotation", "Sharpa", {"mujoco", "motrix", "drake"}), - ("SharpaInhandRotationGrasp", "Sharpa", {"mujoco", "motrix"}), - ), -) -def test_approved_production_families_are_registered_through_the_frozen_seam( - task_name: str, - family: str, - backends: set[str], -) -> None: - registry.ensure_registries() - factories = registry._envs[task_name].env_factory_dict - - assert set(factories) == backends - assert all(isinstance(factory, LegacyFactoryAdapter) for factory in factories.values()) - assert {factory.compatibility.task_family for factory in factories.values()} == {family} - assert {factory.compatibility.status for factory in factories.values()} == { - CompatibilityStatus.ADAPTED - } diff --git a/tests/tasks/test_migration_matrix.py b/tests/tasks/test_migration_matrix.py index 849af3f45..64e72ee0e 100644 --- a/tests/tasks/test_migration_matrix.py +++ b/tests/tasks/test_migration_matrix.py @@ -22,7 +22,6 @@ def test_registered_tasks_have_explicit_migration_records() -> None: @pytest.mark.parametrize( ("task_name", "family", "target", "status"), [ - ("SharpaInhandRotation", "sharpa", "compatibility", "Adapted"), ("G1MotionTracking", "motion_tracking", "complete", "Compatible"), ("G1WBTObs", "motion_tracking", "complete", "Compatible"), ("X2WallFlipTracking", "motion_tracking", "complete", "Compatible"), diff --git a/tests/tasks/test_package_boundary.py b/tests/tasks/test_package_boundary.py index 34d7f5056..d5c7d08f9 100644 --- a/tests/tasks/test_package_boundary.py +++ b/tests/tasks/test_package_boundary.py @@ -19,7 +19,6 @@ "unilab.tasks.locomotion.g1", "unilab.tasks.locomotion.a2", "unilab.tasks.manipulation.allegro_inhand", - "unilab.tasks.manipulation.sharpa_inhand", "unilab.tasks.manipulation.stewart", "unilab.tasks.motion_tracking.g1", "unilab.tasks.motion_tracking.x2", diff --git a/tests/test_cli.py b/tests/test_cli.py index dae9f2ccb..9ef38b17a 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -108,9 +108,9 @@ def test_macos_motrix_finds_uv_venv_mxpython_when_not_on_path( def test_train_profile_routes_to_owner_variant(tmp_path: Path) -> None: (tmp_path / "scripts").mkdir(parents=True) (tmp_path / "scripts" / "train_rsl_rl.py").write_text("", encoding="utf-8") - owner_dir = tmp_path / "conf" / "ppo" / "task" / "sharpa_inhand" + owner_dir = tmp_path / "conf" / "ppo" / "task" / "go1_joystick_flat" owner_dir.mkdir(parents=True) - (owner_dir / "mujoco_hora.yaml").write_text( + (owner_dir / "mujoco_nodr.yaml").write_text( "training:\n sim_backend: mujoco\n", encoding="utf-8", ) @@ -118,16 +118,16 @@ def test_train_profile_routes_to_owner_variant(tmp_path: Path) -> None: command = cli.build_command( mode="train", algo="ppo", - task="sharpa_inhand", + task="go1_joystick_flat", sim="mujoco", - profile="hora", + profile="nodr", overrides=[], root=tmp_path, ) assert command[1:] == [ str(tmp_path / "scripts" / "train_rsl_rl.py"), - "task=sharpa_inhand/mujoco_hora", + "task=go1_joystick_flat/mujoco_nodr", ] @@ -264,7 +264,7 @@ def test_eval_mujoco_interactive_honors_profile_and_render_override( (scripts_dir / "play_interactive.py").write_text("", encoding="utf-8") owner_dir = tmp_path / "conf" / "td3" / "task" / "g1_walk_flat" owner_dir.mkdir(parents=True) - (owner_dir / "mujoco_hora.yaml").write_text( + (owner_dir / "mujoco_nodr.yaml").write_text( "training:\n sim_backend: mujoco\n", encoding="utf-8" ) monkeypatch.setattr( @@ -278,7 +278,7 @@ def test_eval_mujoco_interactive_honors_profile_and_render_override( algo="td3", task="g1_walk_flat", sim="mujoco", - profile="hora", + profile="nodr", overrides=["training.play_render_mode=interactive", "algo.num_envs=4096"], load_run="run_1", render_mode="record", @@ -293,7 +293,7 @@ def test_eval_mujoco_interactive_honors_profile_and_render_override( "--task", "g1_walk_flat", "--sim", - "mujoco_hora", + "mujoco_nodr", "interactive.action_mode=policy", "training.play_only=true", "algo.load_run=run_1", @@ -410,9 +410,9 @@ def test_eval_fallback_prefers_same_profile_owner( ) -> None: (tmp_path / "scripts").mkdir(parents=True) (tmp_path / "scripts" / "train_rsl_rl.py").write_text("", encoding="utf-8") - owner_dir = tmp_path / "conf" / "ppo" / "task" / "sharpa_inhand" + owner_dir = tmp_path / "conf" / "ppo" / "task" / "go1_joystick_flat" owner_dir.mkdir(parents=True) - (owner_dir / "mujoco_hora.yaml").write_text( + (owner_dir / "mujoco_nodr.yaml").write_text( "training:\n sim_backend: mujoco\n", encoding="utf-8" ) (owner_dir / "motrix.yaml").write_text("training:\n sim_backend: motrix\n", encoding="utf-8") @@ -422,15 +422,15 @@ def test_eval_fallback_prefers_same_profile_owner( command = cli.build_command( mode="eval", algo="ppo", - task="sharpa_inhand", + task="go1_joystick_flat", sim="motrix", - profile="hora", + profile="nodr", overrides=[], load_run="-1", root=tmp_path, ) - assert "task=sharpa_inhand/mujoco_hora" in command + assert "task=go1_joystick_flat/mujoco_nodr" in command assert "training.sim_backend=motrix" in command @@ -439,7 +439,7 @@ def test_eval_fallback_without_same_profile_sibling_fails( ) -> None: (tmp_path / "scripts").mkdir(parents=True) (tmp_path / "scripts" / "train_rsl_rl.py").write_text("", encoding="utf-8") - owner_dir = tmp_path / "conf" / "ppo" / "task" / "sharpa_inhand" + owner_dir = tmp_path / "conf" / "ppo" / "task" / "go1_joystick_flat" owner_dir.mkdir(parents=True) (owner_dir / "mujoco.yaml").write_text("training:\n sim_backend: mujoco\n", encoding="utf-8") _pretend_motrix_is_installed(monkeypatch) @@ -449,9 +449,9 @@ def test_eval_fallback_without_same_profile_sibling_fails( cli.build_command( mode="eval", algo="ppo", - task="sharpa_inhand", + task="go1_joystick_flat", sim="motrix", - profile="hora", + profile="nodr", overrides=[], load_run="-1", root=tmp_path, @@ -652,27 +652,23 @@ def _make_demo_checkout(root: Path, *, demo_name: str) -> None: ) +def _register_play_interactive_demo(monkeypatch: pytest.MonkeyPatch) -> str: + name = "playdemo" + monkeypatch.setitem( + demo.DEMO_REGISTRY, + name, + demo.DemoSpec(algo="ppo", task="g1_walk_flat", sim="mujoco", entry="play_interactive"), + ) + return name + + def test_demo_registry_contains_expected_entries() -> None: assert set(demo.DEMO_REGISTRY) == { "dance", "wallflip", "boxtracking", - "sharpa_appo_student", - "inhandgrasp", "teaser", } - assert demo.DEMO_REGISTRY["inhandgrasp"] == demo.DemoSpec( - algo="hora_distill", - task="sharpa_inhand", - sim="mujoco_nodr", - entry="play_interactive", - ) - assert demo.DEMO_REGISTRY["sharpa_appo_student"] == demo.DemoSpec( - algo="hora_distill", - task="sharpa_inhand", - sim="mujoco_nodr", - entry="play_interactive", - ) assert demo.DEMO_REGISTRY["teaser"].entry == "teaser" for name in ("dance", "wallflip", "boxtracking"): spec = demo.DEMO_REGISTRY[name] @@ -698,62 +694,18 @@ def test_demo_eval_entry_passes_checkpoint_as_load_run_override( assert f"algo.load_run={abs_pt}" in command -def test_demo_play_interactive_entry_assembles_inhandgrasp_command( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -) -> None: - _make_demo_checkout(tmp_path, demo_name="inhandgrasp") - monkeypatch.setattr(demo.platform, "system", lambda: "Linux") - abs_pt = str(tmp_path / "fake" / "model_0.pt") - command = demo.build_demo_command( - demo_name="inhandgrasp", - checkpoint_path=abs_pt, - device="cpu", - root=tmp_path, - ) - - assert command[0] == sys.executable - assert command[1] == str(tmp_path / "scripts" / "play_interactive.py") - assert command[2:4] == ["--algo", "hora_distill"] - assert command[4:8] == ["--task", "sharpa_inhand", "--sim", "mujoco_nodr"] - assert f"algo.load_run={abs_pt}" in command - assert "training.device=cpu" in command - assert "interactive.camera_follow_body=false" in command - - -def test_demo_play_interactive_hora_distill_nodr_command(tmp_path: Path) -> None: - _make_demo_checkout(tmp_path, demo_name="sharpa_appo_student") - abs_pt = str(tmp_path / "fake" / "model_0.pt") - - command = demo.build_demo_command( - demo_name="sharpa_appo_student", - checkpoint_path=abs_pt, - root=tmp_path, - ) - - assert command[1:] == [ - str(tmp_path / "scripts" / "play_interactive.py"), - "--algo", - "hora_distill", - "--task", - "sharpa_inhand", - "--sim", - "mujoco_nodr", - f"algo.load_run={abs_pt}", - ] - - def test_demo_play_interactive_sac_owner_path_uses_sac_tree(tmp_path: Path) -> None: (tmp_path / "scripts").mkdir(parents=True, exist_ok=True) (tmp_path / "scripts" / "play_interactive.py").write_text("", encoding="utf-8") - owner_dir = tmp_path / "conf" / "sac" / "task" / "sharpa_inhand" + owner_dir = tmp_path / "conf" / "sac" / "task" / "g1_walk_flat" owner_dir.mkdir(parents=True) - (owner_dir / "mujoco_hora.yaml").write_text( + (owner_dir / "mujoco_nodr.yaml").write_text( "training:\n sim_backend: mujoco\n", encoding="utf-8" ) spec = demo.DemoSpec( algo="sac", - task="sharpa_inhand", - sim="mujoco_hora", + task="g1_walk_flat", + sim="mujoco_nodr", entry="play_interactive", ) @@ -769,9 +721,9 @@ def test_demo_play_interactive_sac_owner_path_uses_sac_tree(tmp_path: Path) -> N "--algo", "sac", "--task", - "sharpa_inhand", + "g1_walk_flat", "--sim", - "mujoco_hora", + "mujoco_nodr", "algo.load_run=/tmp/model_0.pt", ] @@ -779,7 +731,8 @@ def test_demo_play_interactive_sac_owner_path_uses_sac_tree(tmp_path: Path) -> N def test_demo_play_interactive_linux_does_not_materialize_mjpython_app( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: - _make_demo_checkout(tmp_path, demo_name="inhandgrasp") + demo_name = _register_play_interactive_demo(monkeypatch) + _make_demo_checkout(tmp_path, demo_name=demo_name) monkeypatch.setattr(demo.platform, "system", lambda: "Linux") def fail_materialize() -> None: @@ -788,7 +741,7 @@ def fail_materialize() -> None: monkeypatch.setattr(demo, "_ensure_mujoco_mjpython_app", fail_materialize) command = demo.build_demo_command( - demo_name="inhandgrasp", + demo_name=demo_name, checkpoint_path="/tmp/fake/model_0.pt", root=tmp_path, ) @@ -799,7 +752,8 @@ def fail_materialize() -> None: def test_demo_play_interactive_uses_mjpython_on_macos( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: - _make_demo_checkout(tmp_path, demo_name="inhandgrasp") + demo_name = _register_play_interactive_demo(monkeypatch) + _make_demo_checkout(tmp_path, demo_name=demo_name) venv_bin = tmp_path / ".venv" / "bin" venv_bin.mkdir(parents=True) fake_python = venv_bin / "python" @@ -811,7 +765,7 @@ def test_demo_play_interactive_uses_mjpython_on_macos( monkeypatch.setattr(demo, "_ensure_mujoco_mjpython_app", lambda: None) command = demo.build_demo_command( - demo_name="inhandgrasp", + demo_name=demo_name, checkpoint_path="/tmp/fake/model_0.pt", root=tmp_path, ) @@ -824,13 +778,14 @@ def test_demo_play_interactive_checks_mujoco_mjpython_app_on_macos( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: calls: list[str] = [] - _make_demo_checkout(tmp_path, demo_name="inhandgrasp") + demo_name = _register_play_interactive_demo(monkeypatch) + _make_demo_checkout(tmp_path, demo_name=demo_name) monkeypatch.setattr(demo.platform, "system", lambda: "Darwin") monkeypatch.setattr(demo, "_ensure_mujoco_mjpython_app", lambda: calls.append("checked")) monkeypatch.setattr(demo, "_current_env_mjpython", lambda: "/tmp/mjpython") command = demo.build_demo_command( - demo_name="inhandgrasp", + demo_name=demo_name, checkpoint_path="/tmp/fake/model_0.pt", root=tmp_path, ) @@ -839,27 +794,33 @@ def test_demo_play_interactive_checks_mujoco_mjpython_app_on_macos( assert calls == ["checked"] -def test_demo_play_interactive_requires_owner_yaml(tmp_path: Path) -> None: +def test_demo_play_interactive_requires_owner_yaml( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + demo_name = _register_play_interactive_demo(monkeypatch) (tmp_path / "scripts").mkdir(parents=True) (tmp_path / "scripts" / "play_interactive.py").write_text("", encoding="utf-8") with pytest.raises(SystemExit, match="owner config"): demo.build_demo_command( - demo_name="inhandgrasp", + demo_name=demo_name, checkpoint_path="/tmp/fake/model_0.pt", root=tmp_path, ) -def test_demo_play_interactive_requires_script(tmp_path: Path) -> None: - spec = demo.DEMO_REGISTRY["inhandgrasp"] +def test_demo_play_interactive_requires_script( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + demo_name = _register_play_interactive_demo(monkeypatch) + spec = demo.DEMO_REGISTRY[demo_name] owner_dir = tmp_path / "conf" / spec.algo / "task" / spec.task owner_dir.mkdir(parents=True) (owner_dir / f"{spec.sim}.yaml").write_text("training:\n", encoding="utf-8") with pytest.raises(SystemExit, match="play_interactive.py"): demo.build_demo_command( - demo_name="inhandgrasp", + demo_name=demo_name, checkpoint_path="/tmp/fake/model_0.pt", root=tmp_path, ) @@ -891,13 +852,19 @@ def fail_resolve(_: str) -> str: raise AssertionError("local-only demo must not download from Hugging Face") monkeypatch.setattr(demo, "resolve_checkpoint_file", fail_resolve) + monkeypatch.setitem( + demo.DEMO_REGISTRY, + "localonly", + demo.DemoSpec(algo="ppo", task="g1_walk_flat", sim="mujoco", entry="play_interactive"), + ) + monkeypatch.setattr(demo, "_LOCAL_ONLY_CHECKPOINT_DEMOS", {"localonly"}) - rc = demo.run_demo(demo_name="sharpa_appo_student") + rc = demo.run_demo(demo_name="localonly") output = capsys.readouterr().out assert rc == 1 assert "Checkpoint not found" in output - assert "checkpoints/sharpa_appo_student/model_0.pt" in output.replace("\\", "/") + assert "checkpoints/localonly/model_0.pt" in output.replace("\\", "/") def test_demo_local_only_checkpoint_uses_existing_file( @@ -906,8 +873,14 @@ def test_demo_local_only_checkpoint_uses_existing_file( ) -> None: checkout = tmp_path / "checkout" assets = tmp_path / "assets" - checkpoint = assets / "checkpoints" / "sharpa_appo_student" / "model_0.pt" - _make_demo_checkout(checkout, demo_name="sharpa_appo_student") + checkpoint = assets / "checkpoints" / "localonly" / "model_0.pt" + monkeypatch.setitem( + demo.DEMO_REGISTRY, + "localonly", + demo.DemoSpec(algo="ppo", task="g1_walk_flat", sim="mujoco", entry="play_interactive"), + ) + monkeypatch.setattr(demo, "_LOCAL_ONLY_CHECKPOINT_DEMOS", {"localonly"}) + _make_demo_checkout(checkout, demo_name="localonly") checkpoint.parent.mkdir(parents=True) checkpoint.write_bytes(b"checkpoint") calls: list[list[str]] = [] @@ -929,7 +902,7 @@ def fake_run(command: list[str], *, check: bool, env: dict[str, str]) -> SimpleN monkeypatch.setattr(demo, "resolve_checkpoint_file", fail_resolve) monkeypatch.setattr(demo.subprocess, "run", fake_run) - rc = demo.run_demo(demo_name="sharpa_appo_student", device="cpu") + rc = demo.run_demo(demo_name="localonly", device="cpu") assert rc == 0 assert calls == [ @@ -937,11 +910,11 @@ def fake_run(command: list[str], *, check: bool, env: dict[str, str]) -> SimpleN sys.executable, str(checkout / "scripts" / "play_interactive.py"), "--algo", - "hora_distill", + "ppo", "--task", - "sharpa_inhand", + "g1_walk_flat", "--sim", - "mujoco_nodr", + "mujoco", f"algo.load_run={checkpoint}", "training.device=cpu", ] @@ -1225,13 +1198,6 @@ def test_unknown_algo_error_lists_builtin_and_discovered_algos(tmp_path: Path) - assert "dreamer" in message -def test_conf_tree_without_entrypoint_is_not_routable() -> None: - # hora_distill ships a conf tree (compose-only) but no train_hora_distill.py. - with pytest.raises(SystemExit, match="Unsupported algo='hora_distill'"): - cli.build_route("hora_distill", "sharpa_inhand", "mujoco") - assert "hora_distill" not in cli.available_algos() - - def test_malformed_algo_name_is_rejected_fail_closed(tmp_path: Path) -> None: with pytest.raises(SystemExit, match="Unsupported algo"): cli.build_route("../etc", "go2_joystick_flat", "motrix", root=tmp_path) diff --git a/tests/test_completion.py b/tests/test_completion.py index bc342a938..a42a1bfdb 100644 --- a/tests/test_completion.py +++ b/tests/test_completion.py @@ -33,12 +33,12 @@ def _write_completion_fixture(root: Path) -> None: task_name: Go1 sim_backend: mujoco """, - root / "conf" / "ppo" / "task" / "go1" / "mujoco_hora.yaml": """ + root / "conf" / "ppo" / "task" / "go1" / "mujoco_nodr.yaml": """ defaults: - /task/go1/mujoco - _self_ algo: - algo_log_name: hora_ppo + algo_log_name: nodr_ppo """, root / "conf" / "ppo" / "task" / "go1" / "motrix_lab.yaml": """ training: @@ -70,7 +70,7 @@ def _write_completion_fixture(root: Path) -> None: root / "logs" / "rsl_rl_ppo" / "Go1" / "2026-01-01_00-00-00_mujoco", root / "logs" / "rsl_rl_ppo" / "Go1" / "2026-01-02_00-00-00_mujoco", root / "logs" / "rsl_rl_ppo" / "Go2" / "2026-02-01_00-00-00_mujoco", - root / "logs" / "hora_ppo" / "Go1" / "2026-03-01_00-00-00_mujoco", + root / "logs" / "nodr_ppo" / "Go1" / "2026-03-01_00-00-00_mujoco", root / "custom_logs" / "Go3" / "2026-04-01_00-00-00_mujoco", ]: path.mkdir(parents=True) @@ -174,7 +174,7 @@ def test_eval_load_run_completion_respects_profile_log_name(tmp_path: Path) -> N "--sim", "mujoco", "--profile", - "hora", + "nodr", "--load-run", "", ], @@ -216,13 +216,13 @@ def test_train_profile_value_position_completes_profile_names(tmp_path: Path) -> ["uv", "run", "train", "--algo", "ppo", "--sim", "mujoco", "--profile", ""], 8, metadata, - ) == ["custom", "hora", "lab"] + ) == ["custom", "lab", "nodr"] choices = complete_words( - ["uv", "run", "train", "--algo", "ppo", "--sim", "mujoco", "--profile", "h"], + ["uv", "run", "train", "--algo", "ppo", "--sim", "mujoco", "--profile", "n"], 8, metadata, ) - assert choices == ["hora"] + assert choices == ["nodr"] assert "--algo" not in choices assert complete_words( [ @@ -240,7 +240,7 @@ def test_train_profile_value_position_completes_profile_names(tmp_path: Path) -> ], 10, metadata, - ) == ["hora"] + ) == ["nodr"] assert complete_words( [ "uv", @@ -259,10 +259,10 @@ def test_train_profile_value_position_completes_profile_names(tmp_path: Path) -> metadata, ) == ["lab"] assert complete_words( - ["uv", "run", "eval", "--algo", "ppo", "--sim", "mujoco", "--profile", "h"], + ["uv", "run", "eval", "--algo", "ppo", "--sim", "mujoco", "--profile", "n"], 8, metadata, - ) == ["hora"] + ) == ["nodr"] def test_task_completion_respects_selected_profile(tmp_path: Path) -> None: @@ -279,7 +279,7 @@ def test_task_completion_respects_selected_profile(tmp_path: Path) -> None: "--sim", "mujoco", "--profile", - "hora", + "nodr", "--task", "", ], @@ -307,8 +307,6 @@ def test_demo_positional_completes_all_demo_names(tmp_path: Path) -> None: assert choices == [ "boxtracking", "dance", - "inhandgrasp", - "sharpa_appo_student", "teaser", "wallflip", ] diff --git a/tests/test_sharpa.py b/tests/test_sharpa.py deleted file mode 100644 index 341a2d6ad..000000000 --- a/tests/test_sharpa.py +++ /dev/null @@ -1,97 +0,0 @@ -from __future__ import annotations - -from types import SimpleNamespace - -import numpy as np -import pytest - -from unilab.tasks.manipulation.sharpa_inhand.base import SharpaDomainRandConfig -from unilab.tasks.manipulation.sharpa_inhand.grasp_gen import ( - SharpaInhandRotationGraspCfg, - SharpaInhandRotationGraspEnv, -) -from unilab.tasks.manipulation.sharpa_inhand.rotation import ( - SharpaInhandRotationEnv, -) -from unilab.utils.rotation import np_quat_apply - - -def test_sharpa_gravity_direction_randomization_matches_rotated_gravity( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """Gravity-direction DR should rotate a fixed-magnitude downward vector. - - Args: - monkeypatch: Pytest helper used to replace quaternion sampling. - - Returns: - None. The assertions validate the exact gravity vectors produced. - """ - fixed_quat = np.asarray( - [ - [1.0, 0.0, 0.0, 0.0], - [0.0, 0.0, 1.0, 0.0], - ], - dtype=np.float64, - ) - - def _fake_sample_random_quaternion(num_envs: int) -> np.ndarray: - assert num_envs == 2 - return fixed_quat.copy() - - monkeypatch.setitem( - SharpaInhandRotationEnv._build_gravity_direction_randomization.__globals__, - "sample_random_quaternion", - _fake_sample_random_quaternion, - ) - - env = SimpleNamespace( - _cfg=SimpleNamespace( - domain_rand=SimpleNamespace( - randomize_gravity_direction=True, - gravity_direction_magnitude=9.81, - ) - ) - ) - - gravity = SharpaInhandRotationEnv._build_gravity_direction_randomization(env, batch_size=2) - - assert gravity is not None - expected = np_quat_apply( - fixed_quat, - np.asarray([[0.0, 0.0, -9.81]], dtype=np.float64), - ) - np.testing.assert_allclose(gravity, expected) - np.testing.assert_allclose(np.linalg.norm(gravity, axis=1), np.full((2,), 9.81)) - - -def test_sharpa_gravity_direction_randomization_disabled_returns_none() -> None: - env = SimpleNamespace( - _cfg=SimpleNamespace( - domain_rand=SimpleNamespace( - randomize_gravity_direction=False, - gravity_direction_magnitude=9.81, - ) - ) - ) - - gravity = SharpaInhandRotationEnv._build_gravity_direction_randomization(env, batch_size=3) - - assert gravity is None - - -def test_sharpa_grasp_env_rejects_gravity_randomization() -> None: - """Sharpa grasp collection should reject gravity DR explicitly. - - Args: - None. - - Returns: - None. The assertion validates the grasp-task contract. - """ - cfg = SharpaInhandRotationGraspCfg( - domain_rand=SharpaDomainRandConfig(randomize_gravity_direction=True) - ) - - with pytest.raises(ValueError, match="does not support gravity randomization"): - SharpaInhandRotationGraspEnv(cfg, num_envs=1, backend_type="mujoco") diff --git a/tests/training/test_seed_contract.py b/tests/training/test_seed_contract.py index 310624741..6d9fc5f19 100644 --- a/tests/training/test_seed_contract.py +++ b/tests/training/test_seed_contract.py @@ -64,9 +64,7 @@ def test_derive_worker_seed_is_deterministic_and_distinct_from_base_seed(): [ ("ppo", ["task=go1_joystick_flat/mujoco"]), ("ppo", ["task=go1_joystick_flat/mujoco", "algo.seed=41"]), - ("ppo", ["task=sharpa_inhand/mujoco_hora"]), ("appo", ["task=go1_joystick_flat/mujoco"]), - ("appo", ["task=sharpa_inhand/mujoco_hora"]), ("sac", ["task=g1_walk_flat/mujoco"]), ("td3", ["task=g1_walk_flat/mujoco"]), ], diff --git a/tests/training/test_training_helpers.py b/tests/training/test_training_helpers.py index 7ad205192..5536dea4c 100644 --- a/tests/training/test_training_helpers.py +++ b/tests/training/test_training_helpers.py @@ -17,10 +17,8 @@ from unilab.base.config_adapter import BackendAdapter from unilab.base.scene import SceneCfg from unilab.training import ( - format_hora_stage2_checkpoint_error, get_log_root, parse_checkpoint_path, - resolve_hora_stage2_checkpoint_path, ) from unilab.utils.checkpoint import ( get_entrypoint_log_root, @@ -136,81 +134,6 @@ def test_resolve_offpolicy_checkpoint_path_reads_repo_log_tree(tmp_path: Path): assert checkpoint_dir == str(run_dir) -def test_resolve_hora_stage2_checkpoint_path_prefers_last_then_numbered( - monkeypatch: pytest.MonkeyPatch, tmp_path: Path -): - monkeypatch.delenv("UNILAB_TEST_LOG_ROOT", raising=False) - cfg = OmegaConf.create( - { - "training": {"task_name": "Task", "log_root": None}, - "algo": {"algo_log_name": "hora_distill", "load_run": "-1", "checkpoint": -1}, - } - ) - run_dir = tmp_path / "logs" / "hora_distill" / "Task" / "2024-02-01_00-00-00_mujoco" - run_dir.mkdir(parents=True) - (run_dir / "hora_stage2_10.pt").write_bytes(b"") - (run_dir / "hora_stage2_20.pt").write_bytes(b"") - - checkpoint_path, checkpoint_dir = resolve_hora_stage2_checkpoint_path(cfg, root_dir=tmp_path) - - assert checkpoint_path == run_dir / "hora_stage2_20.pt" - assert checkpoint_dir == run_dir - - (run_dir / "hora_stage2_last.pt").write_bytes(b"") - checkpoint_path, checkpoint_dir = resolve_hora_stage2_checkpoint_path(cfg, root_dir=tmp_path) - - assert checkpoint_path == run_dir / "hora_stage2_last.pt" - assert checkpoint_dir == run_dir - - -def test_resolve_hora_stage2_checkpoint_path_reports_missing_explicit_checkpoint( - monkeypatch: pytest.MonkeyPatch, tmp_path: Path -): - monkeypatch.delenv("UNILAB_TEST_LOG_ROOT", raising=False) - cfg = OmegaConf.create( - { - "training": {"task_name": "Task", "log_root": None}, - "algo": {"algo_log_name": "hora_distill", "load_run": "-1", "checkpoint": 7}, - } - ) - run_dir = tmp_path / "logs" / "hora_distill" / "Task" / "2024-02-01_00-00-00_mujoco" - run_dir.mkdir(parents=True) - (run_dir / "hora_stage2_last.pt").write_bytes(b"") - - checkpoint_path, checkpoint_dir = resolve_hora_stage2_checkpoint_path(cfg, root_dir=tmp_path) - - assert checkpoint_path is None - assert checkpoint_dir == run_dir - - message = format_hora_stage2_checkpoint_error( - cfg, - task_log_root=run_dir.parent, - load_path=checkpoint_path, - load_path_dir=checkpoint_dir, - ) - - assert "Requested stage-2 checkpoint was not found" in message - assert "algo.checkpoint=7" in message - - -def test_format_hora_stage2_checkpoint_error_reports_missing_task_root(tmp_path: Path): - cfg = OmegaConf.create( - { - "training": {"task_name": "Task"}, - "algo": {"load_run": "-1", "checkpoint": -1}, - } - ) - - message = format_hora_stage2_checkpoint_error( - cfg, - task_log_root=tmp_path / "missing", - load_path=None, - load_path_dir=None, - ) - - assert "Task log root does not exist." in message - - def test_parse_checkpoint_path_uses_algo_log_name_from_cfg( monkeypatch: pytest.MonkeyPatch, tmp_path: Path ): diff --git a/tests/utils/test_xml_utils.py b/tests/utils/test_xml_utils.py index b6a987a53..e1a4354a3 100644 --- a/tests/utils/test_xml_utils.py +++ b/tests/utils/test_xml_utils.py @@ -55,10 +55,6 @@ def _go1_locomotion_task() -> str: return str(ASSETS_ROOT_PATH / "robots" / "go1" / "locomotion_task.xml") -def _sharpa_robot() -> str: - return str(ASSETS_ROOT_PATH / "robots" / "sharpa_wave" / "right_sharpa_wave.xml") - - def _go2w_locomotion_task() -> str: return str(ASSETS_ROOT_PATH / "robots" / "go2w" / "locomotion_task.xml") @@ -168,13 +164,6 @@ def test_materialize_scene_fragments_merges_static_scene_fragment(tmp_path) -> N os.remove(tmp_xml) -def test_sharpa_contact_excludes_do_not_reference_same_body() -> None: - root = ET.parse(_sharpa_robot()).getroot() - - for exclude in root.findall("./contact/exclude"): - assert exclude.get("body1") != exclude.get("body2") - - def test_materialize_mujoco_hfield_attached_scene_composes_robot_and_task_fragment( tmp_path, ) -> None: diff --git a/tests/visualization/test_interactive_playback.py b/tests/visualization/test_interactive_playback.py index e6c961c6f..bcf70013e 100644 --- a/tests/visualization/test_interactive_playback.py +++ b/tests/visualization/test_interactive_playback.py @@ -19,8 +19,6 @@ PlaybackControls, RslRlPlaybackConfig, RslRlPlaybackSession, - create_appo_playback_session, - create_hora_distill_playback_session, create_rsl_rl_playback_session, create_sac_playback_session, prepare_motion_overlay_selection, @@ -200,208 +198,6 @@ def test_create_rsl_rl_playback_session_rejects_missing_env() -> None: ) -def test_create_hora_distill_playback_session_loads_student_policy(tmp_path: Path) -> None: - from tensordict import TensorDict - - checkpoint_path = tmp_path / "hora_stage2_last.pt" - checkpoint_path.write_bytes(b"checkpoint") - runtime_cfg = SimpleNamespace(training=SimpleNamespace(task_name="RuntimeTask")) - cfg = SimpleNamespace( - training=SimpleNamespace(task_name="SharpaInhandRotation"), - algo=SimpleNamespace(load_run=str(checkpoint_path)), - ) - captured: dict[str, Any] = {} - - class FakeEnv: - num_actions = 2 - action_space = SimpleNamespace( - shape=(2,), - low=np.full((2,), -1.0), - high=np.full((2,), 1.0), - ) - state = SimpleNamespace(info={}) - - def get_physics_state_snapshot(self): - return np.zeros((1, 4), dtype=np.float32) - - class FakeWrapper: - num_actions = 2 - - def __init__(self, env: Any, *, device: str, policy_obs_mode: str): - self.env = env - captured["wrapper_env"] = env - captured["wrapper_device"] = device - captured["policy_obs_mode"] = policy_obs_mode - - def reset(self): - obs = TensorDict( - { - "actor": torch.ones((1, 3), dtype=torch.float32), - "proprio_hist": torch.ones((1, 2, 3), dtype=torch.float32), - }, - batch_size=[1], - ) - return obs, {} - - def step(self, actions): - captured["actions"] = actions - return self.reset()[0], torch.zeros((1,)), torch.zeros((1,), dtype=torch.bool), {} - - class FakeModule: - def __init__(self, name: str): - self.name = name - self.eval_calls = 0 - - def eval(self): - self.eval_calls += 1 - - actor = FakeModule("actor") - hist_normalizer = FakeModule("hist_normalizer") - - def fake_student_policy(actor_obj, hist_obj, obs, *, device): - captured["student_policy"] = (actor_obj, hist_obj, obs, device) - return torch.full((1, 2), 0.5, dtype=torch.float32) - - deps = { - "resolve_stage2_checkpoint_path": lambda cfg_obj: (checkpoint_path, checkpoint_path.parent), - "get_log_root": lambda root_dir, cfg_obj: tmp_path / "logs", - "format_stage2_play_checkpoint_error": lambda *args, **kwargs: "missing", - "checkpoint_reader": lambda path, *, map_location, weights_only: { - "model_state_dict": {}, - "distill_runtime_cfg": {"algo": {"model": {"hidden_dims": [8]}}}, - }, - "apply_teacher_defaults": lambda cfg_obj: cfg_obj, - "cfg_with_checkpoint_runtime": lambda cfg_obj, checkpoint: runtime_cfg, - "build_play_env_cfg_override": lambda cfg_obj: {"env": "override"}, - "create_env": lambda cfg_obj, *, num_envs, env_cfg_override: FakeEnv(), - "wrapper_cls": FakeWrapper, - "build_student_actor_and_normalizer": lambda wrapped_env, cfg_obj, *, device: ( - actor, - hist_normalizer, - ), - "load_distilled_checkpoint": lambda actor_obj, hist_obj, path, *, device: captured.update( - { - "loaded": (actor_obj, hist_obj, path, device), - } - ), - "student_policy": fake_student_policy, - } - - session, policy_obs_mode, checkpoint = create_hora_distill_playback_session( - playback_cfg=RslRlPlaybackConfig( - task="SharpaInhandRotation", - load_run=str(checkpoint_path), - checkpoint=None, - action_mode="policy", - policy_obs_mode="actor", - algo_log_name="hora_distill", - log_root=None, - num_envs=1, - ), - cfg=cfg, - root_dir=tmp_path, - device="cpu", - deps=deps, - log=lambda message: None, - ) - - assert session.env is captured["wrapper_env"] - assert policy_obs_mode == "actor" - assert checkpoint == str(checkpoint_path) - assert captured["policy_obs_mode"] == "actor" - assert captured["loaded"][0] is actor - assert captured["loaded"][1] is hist_normalizer - assert captured["loaded"][2] == checkpoint_path - assert actor.eval_calls == 1 - assert hist_normalizer.eval_calls == 1 - - session.reset() - session.step_once() - - assert torch.equal(captured["actions"], torch.full((1, 2), 0.5)) - assert captured["student_policy"][0] is actor - assert captured["student_policy"][1] is hist_normalizer - assert str(captured["student_policy"][3]) == "cpu" - - -def test_create_hora_distill_playback_session_missing_checkpoint_uses_zero_actions( - tmp_path: Path, -) -> None: - runtime_cfg = SimpleNamespace(training=SimpleNamespace(task_name="RuntimeTask")) - cfg = SimpleNamespace( - training=SimpleNamespace(task_name="SharpaInhandRotation"), - algo=SimpleNamespace(load_run="missing"), - ) - captured: dict[str, Any] = {} - messages: list[str] = [] - - class FakeEnv: - action_space = SimpleNamespace( - shape=(2,), - low=np.full((2,), -1.0), - high=np.full((2,), 1.0), - ) - state = SimpleNamespace(info={}) - - def get_physics_state_snapshot(self): - return np.zeros((1, 4), dtype=np.float32) - - class FakeWrapper: - def __init__(self, env: Any, *, device: str, policy_obs_mode: str): - self.env = env - captured["policy_obs_mode"] = policy_obs_mode - - def reset(self): - return "obs", {} - - def step(self, actions): - captured["actions"] = actions - return "obs", torch.zeros((1,)), torch.zeros((1,), dtype=torch.bool), {} - - deps = { - "resolve_stage2_checkpoint_path": lambda cfg_obj: (None, None), - "get_log_root": lambda root_dir, cfg_obj: tmp_path / "logs", - "format_stage2_play_checkpoint_error": lambda *args, **kwargs: "missing checkpoint", - "apply_teacher_defaults": lambda cfg_obj: runtime_cfg, - "build_play_env_cfg_override": lambda cfg_obj: {"env": "override"}, - "create_env": lambda cfg_obj, *, num_envs, env_cfg_override: FakeEnv(), - "wrapper_cls": FakeWrapper, - } - - session, policy_obs_mode, checkpoint = create_hora_distill_playback_session( - playback_cfg=RslRlPlaybackConfig( - task="SharpaInhandRotation", - load_run="missing", - checkpoint=None, - action_mode="policy", - policy_obs_mode="actor", - algo_log_name="hora_distill", - log_root=None, - num_envs=1, - ), - cfg=cfg, - root_dir=tmp_path, - device="cpu", - deps=deps, - log=messages.append, - ) - - assert policy_obs_mode == "actor" - assert checkpoint is None - assert captured["policy_obs_mode"] == "actor" - assert messages == [ - "missing checkpoint", - "WARNING: falling back to zero actions.", - "Policy obs mode: actor", - "Action mode: policy", - ] - - session.reset() - session.step_once() - - assert torch.equal(captured["actions"], torch.zeros((1, 2))) - - def _rsl_rl_session_test_env() -> SimpleNamespace: return SimpleNamespace( obs_groups_spec={"obs": 5}, @@ -645,388 +441,6 @@ def test_prepare_motion_overlay_selection_filters_body_names() -> None: assert messages == ["WARNING: body name not found in task body list: missing"] -def test_appo_hora_playback_session_uses_hora_wrapper_and_actor_checkpoint( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - import rsl_rl.utils as rsl_rl_utils - import uni_rl.algos.hora.models as hora_models - import uni_rl.algos.hora.rsl_rl as hora_rsl - from omegaconf import OmegaConf - from tensordict import TensorDict - - checkpoint = tmp_path / "model_10.pt" - torch.save({"actor": {"weight": torch.tensor(1.0)}}, checkpoint) - captured: dict[str, Any] = {} - - class FakeHoraWrapper: - def __init__(self, env, *, device, policy_obs_mode): - captured["wrapper_cls"] = "hora" - captured["policy_obs_mode"] = policy_obs_mode - self.env = env - self.device = device - self.num_envs = env.num_envs - - def get_observations(self): - return TensorDict( - { - "actor": torch.zeros((1, 3)), - "priv_info": torch.zeros((1, 2)), - "proprio_hist": torch.zeros((1, 4, 3)), - }, - batch_size=1, - ) - - def reset(self): - return self.get_observations(), {} - - def step(self, actions): - captured["step_actions"] = actions - return self.get_observations(), torch.zeros(1), torch.zeros(1).bool(), {} - - class FakeActor(torch.nn.Module): - def __init__(self, *args, **kwargs): - super().__init__() - captured["actor_kwargs"] = kwargs - - def load_state_dict(self, state_dict, strict=True): - captured["loaded_actor"] = state_dict - return None - - def forward(self, obs): - captured["policy_obs"] = obs - return torch.ones((1, 2), dtype=torch.float32) - - fake_env = SimpleNamespace( - num_envs=1, - obs_groups_spec={"obs": 3, "critic": 5}, - action_space=SimpleNamespace( - shape=(2,), - low=np.full((2,), -1.0), - high=np.full((2,), 1.0), - ), - get_physics_state_snapshot=lambda: np.zeros((1, 4), dtype=np.float32), - state=SimpleNamespace(info={}), - ) - cfg = OmegaConf.create( - { - "training": {"task_name": "Task", "play_env_num": 1, "log_root": None}, - "algo": { - "algo_log_name": "hora_appo", - "load_run": str(tmp_path), - "checkpoint": str(checkpoint), - "runtime_impl": "hora_appo", - "obs_groups": {"actor": {"actor": 0, "priv_info": 0}}, - "actor": {"class_name": "fake.Actor"}, - "critic": {}, - }, - } - ) - - monkeypatch.setattr(hora_rsl, "HoraRslRlVecEnvWrapper", FakeHoraWrapper) - monkeypatch.setattr(rsl_rl_utils, "resolve_callable", lambda path: FakeActor) - monkeypatch.setattr( - hora_models, - "build_hora_shared_actor_critic", - lambda **kwargs: torch.nn.Identity(), - ) - - session, policy_obs_mode, resolved_checkpoint = create_appo_playback_session( - playback_cfg=RslRlPlaybackConfig( - task="Task", - load_run="run", - checkpoint=None, - action_mode="policy", - policy_obs_mode="auto", - algo_log_name="hora_appo", - log_root=None, - ), - cfg=cfg, - rl_cfg=OmegaConf.to_container(cfg.algo, resolve=True), - env_factory=lambda num_envs: fake_env, - root_dir=tmp_path, - device="cpu", - wrapper_cls=object, - log=lambda message: None, - ) - - session.reset() - assert session.advance(PlaybackControls()) is True - assert isinstance(session.actor, FakeActor) - assert policy_obs_mode == "actor" - assert resolved_checkpoint == str(checkpoint) - assert captured["wrapper_cls"] == "hora" - assert captured["policy_obs_mode"] == "actor" - torch.testing.assert_close(captured["loaded_actor"]["weight"], torch.tensor(1.0)) - assert isinstance(captured["policy_obs"], TensorDict) - torch.testing.assert_close(captured["step_actions"], torch.ones((1, 2))) - - -def test_sac_hora_playback_session_updates_priv_info_after_reset_and_step( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - import uni_rl.algos.common.actor_factory as actor_factory - from omegaconf import OmegaConf - - import unilab.utils.checkpoint as checkpoint_utils - import unilab.visualization.interactive_playback as offpolicy_play - - checkpoint = tmp_path / "model_10.pt" - torch.save({"actor": {}}, checkpoint) - reset_priv = np.array([[4.0, 5.0]], dtype=np.float32) - step_priv = np.array([[8.0, 9.0]], dtype=np.float32) - captured: dict[str, Any] = {} - - class FakeActor: - def eval(self): - return self - - def load_state_dict(self, state_dict): - captured["loaded_actor"] = state_dict - - def explore(self, obs, priv_info, deterministic=True): - captured["reset_priv_info"] = priv_info.detach().cpu().numpy() - captured["deterministic"] = deterministic - return torch.zeros((obs.shape[0], 2), dtype=obs.dtype, device=obs.device) - - class FakeEnv: - num_envs = 1 - obs_groups_spec = {"obs": 3, "critic": 5} - action_space = SimpleNamespace( - shape=(2,), - low=np.full((2,), -1.0), - high=np.full((2,), 1.0), - ) - state = None - - def init_state(self): - self.state = SimpleNamespace(info={}) - - def reset(self, env_indices): - del env_indices - return ( - { - "obs": np.zeros((1, 3), dtype=np.float32), - "critic": np.concatenate( - [np.zeros((1, 3), dtype=np.float32), reset_priv], - axis=1, - ), - }, - {"critic_info": reset_priv}, - ) - - def step(self, actions): - captured["actions"] = actions - self.state = SimpleNamespace( - obs={ - "obs": np.ones((1, 3), dtype=np.float32), - "critic": np.concatenate( - [np.ones((1, 3), dtype=np.float32), step_priv], - axis=1, - ), - }, - info={"critic_info": step_priv}, - ) - return self.state - - def get_physics_state_snapshot(self): - return np.zeros((1, 4), dtype=np.float32) - - cfg = OmegaConf.create( - { - "training": {"task_name": "Task", "device": None}, - "algo": { - "algo_log_name": "hora_sac", - "load_run": "run", - "actor_hidden_dim": 16, - "use_layer_norm": False, - "runtime_impl": "hora_sac", - }, - } - ) - - monkeypatch.setattr( - offpolicy_play, - "default_device", - lambda torch_module, preferred=None: "cpu", - ) - monkeypatch.setattr(offpolicy_play, "resolve_play_obs_dims", lambda spec: (3, 5)) - monkeypatch.setattr( - offpolicy_play, - "resolve_play_actor_spec", - lambda algo_name, cfg, *, obs_dim, critic_obs_dim: ( - "hora_sac", - {"priv_info_dim": 2}, - ), - ) - monkeypatch.setattr( - checkpoint_utils, - "resolve_offpolicy_checkpoint_path", - lambda *args, **kwargs: (str(checkpoint), str(tmp_path)), - ) - monkeypatch.setattr(actor_factory, "build_actor", lambda *args, **kwargs: FakeActor()) - - session, policy_obs_mode, resolved_checkpoint = create_sac_playback_session( - playback_cfg=RslRlPlaybackConfig( - task="Task", - load_run="run", - checkpoint=None, - action_mode="policy", - policy_obs_mode="actor", - algo_log_name="hora_sac", - log_root=None, - ), - cfg=cfg, - env_factory=lambda num_envs: FakeEnv(), - root_dir=tmp_path, - device="cpu", - log=lambda message: None, - ) - - session.reset() - assert session.advance(PlaybackControls()) is True - assert policy_obs_mode == "actor" - assert resolved_checkpoint == str(checkpoint) - assert captured["loaded_actor"] == {} - np.testing.assert_allclose(captured["reset_priv_info"], reset_priv) - np.testing.assert_allclose(session.current_priv_info, step_priv) - assert captured["deterministic"] is True - - -def test_hora_distill_playback_session_loads_stage2_checkpoint_and_student_policy( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - import uni_rl.algos.hora.distill as distill - import uni_rl.algos.hora.rsl_rl as hora_rsl - from omegaconf import OmegaConf - from tensordict import TensorDict - - import unilab.base.config_adapter as config_adapter - import unilab.training as training - - checkpoint = tmp_path / "hora_stage2_last.pt" - torch.save({"model_state_dict": {}, "distill_runtime_cfg": {}}, checkpoint) - captured: dict[str, Any] = {} - - class FakeWrapper: - def __init__(self, env, *, device, policy_obs_mode): - self.env = env - self.device = device - self.num_envs = env.num_envs - self.num_actions = 2 - captured["policy_obs_mode"] = policy_obs_mode - - def get_observations(self): - return TensorDict( - { - "actor": torch.zeros((1, 3)), - "priv_info": torch.zeros((1, 2)), - "proprio_hist": torch.zeros((1, 4, 3)), - }, - batch_size=1, - ) - - def reset(self): - return self.get_observations(), {} - - def step(self, actions): - captured["actions"] = actions - return self.get_observations(), torch.zeros(1), torch.zeros(1).bool(), {} - - fake_env = SimpleNamespace( - num_envs=1, - action_space=SimpleNamespace( - shape=(2,), - low=np.full((2,), -1.0), - high=np.full((2,), 1.0), - ), - get_physics_state_snapshot=lambda: np.zeros((1, 4), dtype=np.float32), - state=SimpleNamespace(info={}), - ) - cfg = OmegaConf.create( - { - "training": {"task_name": "Task", "sim_backend": "mujoco", "log_root": None}, - "algo": {"algo_log_name": "hora_distill", "load_run": "run", "checkpoint": -1}, - } - ) - - monkeypatch.setattr( - training, - "resolve_hora_stage2_checkpoint_path", - lambda cfg, *, root_dir: (checkpoint, tmp_path), - ) - - def fake_cfg_with_checkpoint_runtime(cfg, checkpoint_payload): - captured["runtime_helper_checkpoint"] = checkpoint_payload - return cfg - - monkeypatch.setattr( - distill, - "cfg_with_checkpoint_runtime", - fake_cfg_with_checkpoint_runtime, - ) - - class FakeBackendAdapter: - def __init__(self, *args, **kwargs): - pass - - def build_play_env_cfg_override(self): - return {} - - monkeypatch.setattr(config_adapter, "BackendAdapter", FakeBackendAdapter) - monkeypatch.setattr( - distill, - "student_policy", - lambda actor, hist_normalizer, obs, *, device: torch.ones((1, 2)), - ) - monkeypatch.setattr(config_adapter, "create_env", lambda *args, **kwargs: fake_env) - monkeypatch.setattr(hora_rsl, "HoraRslRlVecEnvWrapper", FakeWrapper) - monkeypatch.setattr( - distill, - "build_student_actor_and_normalizer", - lambda wrapped_env, cfg, *, device: ( - SimpleNamespace(eval=lambda: None), - SimpleNamespace(eval=lambda: None), - ), - ) - - def fake_load_distilled_checkpoint(actor, hist_normalizer, checkpoint_path, *, device): - captured["loaded_checkpoint"] = checkpoint_path - return {} - - monkeypatch.setattr(distill, "load_distilled_checkpoint", fake_load_distilled_checkpoint) - - session, policy_obs_mode, resolved_checkpoint = create_hora_distill_playback_session( - playback_cfg=RslRlPlaybackConfig( - task="Task", - load_run="run", - checkpoint=None, - action_mode="policy", - policy_obs_mode="actor", - algo_log_name="hora_distill", - log_root=None, - ), - cfg=cfg, - root_dir=tmp_path, - device="cpu", - log=lambda message: None, - ) - - session.reset() - assert session.advance(PlaybackControls()) is True - assert policy_obs_mode == "actor" - assert resolved_checkpoint == str(checkpoint) - assert captured["runtime_helper_checkpoint"] == { - "model_state_dict": {}, - "distill_runtime_cfg": {}, - } - assert captured["policy_obs_mode"] == "actor" - assert captured["loaded_checkpoint"] == checkpoint - torch.testing.assert_close(captured["actions"], torch.ones((1, 2))) - - def test_infer_checkpoint_actor_input_dim_mlp_key(tmp_path: Path) -> None: from unilab.visualization.interactive_playback import infer_checkpoint_actor_input_dim