diff --git a/docs/sphinx/source/en/2-user_guide/4-tasks/2-motion_tracking.md b/docs/sphinx/source/en/2-user_guide/4-tasks/2-motion_tracking.md index 7b78c58a5..8023b95b9 100644 --- a/docs/sphinx/source/en/2-user_guide/4-tasks/2-motion_tracking.md +++ b/docs/sphinx/source/en/2-user_guide/4-tasks/2-motion_tracking.md @@ -59,11 +59,13 @@ uv run train --algo sac --task g1_wbt_obs --sim mujoco training.use_amp=true The `g1_wbt_obs` owner is the deploy-aligned off-policy observation profile: a pelvis IMU state (`pelvis_local_linvel` / `pelvis_gyro` / `pelvis_upvector`) plus -per-term observation history (`noise_config.obs_history_length: 5`), byte-aligned -with the deploy-time `ObservationManager`. Deploy tooling lives under -`scripts/deploy/`, and the observation alignment is cross-checked by -`tests/scripts/test_obs_alignment_g1_wbt.py`. When a Motrix sim2sim replay needs a -checkpoint from another log root, pass the absolute path through `uv run eval`: +per-term observation history (`noise_config.obs_history_length: 5`), flattened +oldest-first per term so a hardware runtime assembling per-term history reads the +same vector. That ordering is guarded by +`tests/scripts/test_obs_alignment_g1_wbt.py`; the hardware-side contract is +documented in the sim-to-real deployment guide. When a Motrix sim2sim replay +needs a checkpoint from another log root, pass the absolute path through +`uv run eval`: ```bash uv run eval --algo sac --task g1_motion_tracking --sim motrix \ diff --git a/docs/sphinx/source/en/3-deployment/1-sim_to_real/1-overview.md b/docs/sphinx/source/en/3-deployment/1-sim_to_real/1-overview.md index 38a1bc8b9..4a838a6fb 100644 --- a/docs/sphinx/source/en/3-deployment/1-sim_to_real/1-overview.md +++ b/docs/sphinx/source/en/3-deployment/1-sim_to_real/1-overview.md @@ -6,9 +6,8 @@ page in this section drills into one stage. ## What "sim-to-real" means in UniLab A deployable UniLab policy is the exported policy plus the exact observation -and action contracts used by the selected task owner. The G1 WBT helper path -materializes this as `policy.onnx`, `deploy_config.yaml`, and a motion binary; -other robots need an equivalent hardware-side runtime that: +and action contracts used by the selected task owner. UniLab ships the training +side and the ONNX export; every robot needs a hardware-side runtime that: 1. Reads sensors → assembles the **same observation vector** the policy saw in simulation. diff --git a/docs/sphinx/source/en/3-deployment/1-sim_to_real/2-g1_whole_body.md b/docs/sphinx/source/en/3-deployment/1-sim_to_real/2-g1_whole_body.md index e4ba8b088..3ea3dfd3c 100644 --- a/docs/sphinx/source/en/3-deployment/1-sim_to_real/2-g1_whole_body.md +++ b/docs/sphinx/source/en/3-deployment/1-sim_to_real/2-g1_whole_body.md @@ -2,20 +2,21 @@ ::::{admonition} Hardware target :class: note -Unitree G1 humanoid (29-DoF variant). Joints are assumed in the order exported -by `scripts/deploy/export_deploy_config.py` from -`src/unilab/assets/robots/g1/scene_flat.xml`; verify that order before -hardware bring-up. +Unitree G1 humanoid (29-DoF variant). Joint order comes from the task owner's +scene (`src/unilab/assets/robots/g1/scene_flat.xml`, actuator order); verify +that order against your SDK motor indices before hardware bring-up. :::: -This guide walks the **last mile** between a converged G1 motion-tracking -policy and a closed-loop run on the robot. +This guide covers the **observation and action contract** a G1 motion-tracking +policy expects on hardware. The repository does not ship a G1 deploy runtime — +you supply the hardware-side loop, and this page tells you what it must +reproduce. ## 0. Verify your sim-side checkpoint ```bash # Replay the policy headlessly and produce a video. -uv run eval --algo ppo --task g1_motion_tracking --sim motrix --load-run -1 \ +uv run eval --algo sac --task g1_wbt_obs --sim mujoco --load-run -1 \ --render-mode record ``` @@ -25,39 +26,51 @@ What to look for in the video: - Joint velocities and actions remain finite and within the expected range. - Contact timing looks consistent with the reference motion. -If any of those is off, fix the sim-side checkpoint or deploy contract before -hardware bring-up. +If any of those is off, fix the sim-side checkpoint before hardware bring-up. -## 1. Export +## 1. Pick the owner, then read its contract off the YAML -Use the training playback path to export `policy.onnx`, then export the G1 WBT -deploy config and motion binary with the committed deployment helpers: +Every field your hardware loop needs is declared in the task owner YAML. The +deploy-oriented G1 owners are: -```bash -uv run eval --algo ppo --task g1_motion_tracking --sim motrix --load-run -1 - -uv run scripts/deploy/export_deploy_config.py \ - --output logs/deploy/deploy_config.yaml - -uv run scripts/deploy/export_motion_bin.py \ - --output logs/deploy/dance1.bin +```{list-table} +:header-rows: 1 +:widths: 34 22 44 + +* - Owner + - Actor obs width + - Notes +* - `conf/offpolicy/task/sac/g1_wbt_obs/mujoco.yaml` + - 514 (H=5) + - Proprio history, no state estimation: drops `base_lin_vel` and + `motion_anchor_pos_b`, pelvis IMU. +* - `conf/ppo/task/g1_motion_tracking_deploy/mujoco.yaml` + - 154 (H=1) + - Single-step mimic actor layout, per-joint `action_scale` list. ``` -The deployment-side prototype consumes: +::::{admonition} Read the width off the env, not off this table +:class: warning +Actor obs width is a function of the owner's `noise_config` flags — see +`_actor_obs_dim` in `src/unilab/envs/motion_tracking/g1/tracking_obs.py` for +`g1_wbt_obs`, and `mimic_actor_obs_dim` in +`src/unilab/envs/motion_tracking/common/observations.py` for the deploy owner. +If the ONNX input width disagrees with what your hardware loop assembles, that +is a contract bug, not a hardware tuning problem. +:::: + +Export `policy.onnx` through the training playback path: -``` -runs// -└── policy.onnx -logs/deploy/ -├── deploy_config.yaml -└── dance1.bin +```bash +uv run eval --algo sac --task g1_wbt_obs --sim mujoco --load-run -1 ``` ## 2. Observation contract -For the committed G1 WBT deploy helper, the observation layout is exported into -`deploy_config.yaml` as `obs_layout`. `scripts/deploy/export_deploy_config.py` -is the source of truth for the segment order: +For `g1_wbt_obs`, the actor obs is assembled in this order (see +`_build_actor_obs` in `src/unilab/envs/motion_tracking/g1/tracking_obs.py`). +Single-step reference terms come first, then each proprio term's full history +flattened **oldest-first**: ```{list-table} :header-rows: 1 @@ -77,32 +90,39 @@ is the source of truth for the segment order: - anchor orientation term from the reference and robot torso frames * - `gyro` - 3 per history step - - IMU gyro term + - IMU gyro term (`env.sensor.gyro`, `pelvis_gyro` for this owner) * - `joint_pos_rel` - 29 per history step - - measured joint position minus `default_angles` + - measured joint position minus the `stand` keyframe joint angles * - `dof_vel` - 29 per history step - joint velocity term * - `last_actions` - - 29 + - 29 per history step - previous raw actor output ``` -The export script also records each segment's `history_length` and verifies the -total `obs_dim`. `scripts/deploy/sim_prototype.py` refuses to run when the ONNX -input width and `deploy_config.yaml` `obs_dim` disagree. +History depth `H` is `env.noise_config.obs_history_length` (5 for this owner). +Per-term oldest-first ordering is guarded by +`tests/scripts/test_obs_alignment_g1_wbt.py`; mirror that ordering on hardware +or the policy reads a permuted vector. ## 3. Actuator interface -The G1 deploy prototype maps actor output exactly as: -`action * action_scale + default_angles`, then clips to `joint_lower` / -`joint_upper` and applies EMA smoothing from `ema_alpha`. +Map actor output as `action * action_scale + default_angles`, then clamp to the +scene's joint range before the target reaches the motor driver. + +- `action_scale` is `env.control_config.action_scale` in the owner YAML. It may + be a **scalar** (2.0 for `g1_wbt_obs`) or a **per-joint list** (29 entries for + `g1_motion_tracking_deploy`). Reproduce the owner's form exactly — do not + average a list, take its first entry, or broadcast a scalar over a list owner. +- `default_angles` is the `stand` keyframe joint block of the owner's scene. +- Joint limits and gains come from the same scene XML (`jnt_range`, position + actuator `gainprm` / `biasprm`). -- Action = target joint position, **scaled** by the `action_scale` entry in - `deploy_config.yaml`. -- Clamp the target to the generated joint range before it reaches the motor - driver. +Training applies the target directly with no smoothing. If hardware jitter +forces you to add smoothing, verify the sim2sim impact first — every step of lag +pushes observations out of the training distribution. ## 4. Reference motion sync @@ -125,8 +145,7 @@ specifics: - Reject non-finite actions and shape mismatches before applying `action_scale`. -- Clamp generated targets with `joint_lower` / `joint_upper` from - `deploy_config.yaml`. +- Clamp generated targets with the joint range from the owner's scene XML. - Keep watchdog, pose monitor, and operator-stop thresholds in the deploy controller and test them independently of the policy. @@ -145,18 +164,9 @@ and `last_actions` wiring mistakes are easiest to catch. ## 7. What to log Log the **full observation vector**, **full action vector**, and **wall -clock** for every step. Before hardware bring-up, validate the same ONNX, -deploy config, and motion binary through the MuJoCo deployment prototype: - -```bash -uv run scripts/deploy/sim_prototype.py \ - --onnx runs//policy.onnx \ - --config logs/deploy/deploy_config.yaml \ - --motion logs/deploy/dance1.bin -``` - -A mismatch between the ONNX input width and `deploy_config.yaml` `obs_dim` is a -deployment contract bug, not a hardware tuning problem. +clock** for every step. Compare the first hardware observation window against a +sim episode built from the same owner YAML — that diff localizes unit, frame, +and ordering mistakes faster than any reward inspection. ## See also diff --git a/docs/sphinx/source/en/3-deployment/1-sim_to_real/5-onnx_runtime.md b/docs/sphinx/source/en/3-deployment/1-sim_to_real/5-onnx_runtime.md index 92c101742..71387f942 100644 --- a/docs/sphinx/source/en/3-deployment/1-sim_to_real/5-onnx_runtime.md +++ b/docs/sphinx/source/en/3-deployment/1-sim_to_real/5-onnx_runtime.md @@ -26,38 +26,21 @@ uv run eval --algo sac --task g1_walk_flat --sim mujoco --load-run -1 `uv run eval` sets playback mode and maps `--load-run` to the checkpoint selector used by the routed training script. The exported file is written into -the selected run directory. For deployment -prototypes, keep the exported `policy.onnx` together with the deploy-side -configuration and motion assets used by the runtime. +the selected run directory. For deployment, keep the exported `policy.onnx` +together with the task owner YAML it was trained from — that YAML is the +authority on the observation and action contract the runtime must reproduce. -## G1 Deployment Prototype +## Verifying the Exported Graph -The committed G1 WBT deployment helpers use these artifacts: +The playback path validates the exported graph against PyTorch before writing +it, so a successful export already establishes numerical parity. What it does +**not** establish is that your hardware-side loop assembles the same input +vector. Before hardware bring-up: -| Artifact | Producer | -| --- | --- | -| `policy.onnx` | Training playback export above. | -| `deploy_config.yaml` | `scripts/deploy/export_deploy_config.py`. | -| `dance1.bin` or another motion binary | `scripts/deploy/export_motion_bin.py`. | - -Example validation run: - -```bash -uv run scripts/deploy/export_deploy_config.py \ - --output logs/deploy/deploy_config.yaml - -uv run scripts/deploy/export_motion_bin.py \ - --output logs/deploy/dance1.bin - -uv run scripts/deploy/sim_prototype.py \ - --onnx runs//policy.onnx \ - --config logs/deploy/deploy_config.yaml \ - --motion logs/deploy/dance1.bin -``` - -`scripts/deploy/sim_prototype.py` checks that the ONNX input width matches the -`obs_dim` in `deploy_config.yaml` and then drives the policy in MuJoCo with the -same observation layout the deployment side expects. +- Read the actor obs width off the env (not off a doc table) and confirm it + matches the ONNX input width. +- Confirm your term order and per-term history ordering against the owner's + `_build_actor_obs`. For G1 whole-body tracking, see {doc}`2-g1_whole_body`. ## See Also diff --git a/docs/sphinx/source/en/3-deployment/1-sim_to_real/7-safety_layers.md b/docs/sphinx/source/en/3-deployment/1-sim_to_real/7-safety_layers.md index 9835121aa..6477f9cf3 100644 --- a/docs/sphinx/source/en/3-deployment/1-sim_to_real/7-safety_layers.md +++ b/docs/sphinx/source/en/3-deployment/1-sim_to_real/7-safety_layers.md @@ -43,28 +43,25 @@ flowchart LR ``` Keep the hard real-time safety checks in the deploy controller, not in the -training script. The repository's G1 helper path exports deploy config and runs -a MuJoCo prototype; it does not implement a production motor-driver safety -loop. +training script. The repository does not implement a production motor-driver +safety loop — that boundary is yours to build and test. ## What the policy assumes you've configured -The G1 deployment helper exports these fields into `deploy_config.yaml`: +The policy expects the action mapping and limits its training owner declared. +For the G1 WBT owner (`conf/offpolicy/task/sac/g1_wbt_obs/mujoco.yaml`): -```yaml -action_scale: 2.0 -ema_alpha: 1.0 -default_angles: [...] -joint_lower: [...] -joint_upper: [...] -kp: [...] -kd: [...] -``` +| Quantity | Authority | +| --- | --- | +| `action_scale` | `env.control_config.action_scale` (scalar `2.0` for this owner; other owners declare a per-joint list) | +| `default_angles` | `stand` keyframe joint block of the owner's scene XML | +| joint limits | `jnt_range` in the scene XML | +| `kp` / `kd` | position actuator `gainprm` / `biasprm` in the scene XML | -`scripts/deploy/sim_prototype.py` consumes the same fields and applies -`action * action_scale + default_angles`, joint clipping, and EMA smoothing. -Hardware controllers should consume generated config rather than hand-copying -joint ranges or gains. +Derive these from the owner YAML and its scene, and reproduce the owner's +`action_scale` **form** exactly — a scalar owner and a per-joint-list owner are +not interchangeable. Do not hand-copy joint ranges or gains into a second place +that can silently drift from the asset. ## Hand-off testing @@ -72,8 +69,8 @@ Before integrating policy → safety → motor, test the safety layer in isolation: 1. Inject a NaN action and verify the command is rejected. -2. Inject an out-of-range joint target and verify clamping uses - `joint_lower` / `joint_upper` from `deploy_config.yaml`. +2. Inject an out-of-range joint target and verify clamping uses the joint range + from the owner's scene XML. 3. Cut the policy feed mid-run and verify the controller enters its configured safe state. diff --git a/docs/sphinx/source/en/3-deployment/1-sim_to_real/8-latency_budget.md b/docs/sphinx/source/en/3-deployment/1-sim_to_real/8-latency_budget.md index 7e3177b7f..43c7fd5d8 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 @@ -9,9 +9,9 @@ budgets as robot-specific measurements, not UniLab defaults. | Surface | Repo evidence | What it covers | | --- | --- | --- | | One-step action delay | `control_config.simulate_action_latency` in locomotion and G1 motion-tracking envs | Executes the previous action instead of the current action. | -| G1 WBT observation history | `noise_config.obs_history_length` and `scripts/deploy/export_deploy_config.py` | Exports per-term `obs_layout` history for `gyro`, `joint_pos_rel`, `dof_vel`, and `last_actions`. | +| G1 WBT observation history | `noise_config.obs_history_length` in `conf/offpolicy/task/sac/g1_wbt_obs/mujoco.yaml` | Per-term history for `gyro`, `joint_pos_rel`, `dof_vel`, and `last_actions`. | | Sharpa tactile contact latency | `domain_rand.contact_latency` in Sharpa in-hand configs | Keeps previous tactile contact values for sampled contact channels. | -| Deploy-side ONNX contract check | `scripts/deploy/sim_prototype.py` | Validates `obs_layout`, `obs_dim`, ONNX input width, clipping, and EMA action smoothing for the G1 WBT path. | +| 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 @@ -30,23 +30,13 @@ The checked-in G1 WBT owner enables this flag in ## Observation Lag And History -The G1 WBT deployment helpers do not guess observation width. They export a -schema with `obs_layout`, per-term `history_length`, and `obs_dim`; the -prototype then assembles the same layout and refuses mismatches. +Observation width is a function of the owner's `noise_config`, not something a +hardware runtime may guess. For the G1 WBT owner, `obs_history_length: 5` gives +each proprioceptive term a 5-step history flattened oldest-first, while +reference terms stay single-step. See {doc}`2-g1_whole_body` for the full term +order. -```bash -uv run scripts/deploy/export_deploy_config.py \ - --output logs/deploy/deploy_config.yaml - -uv run scripts/deploy/sim_prototype.py \ - --onnx runs//policy.onnx \ - --config logs/deploy/deploy_config.yaml \ - --motion logs/deploy/dance1.bin -``` - -Do not lag command/reference terms unless the training owner did so. In the G1 -WBT schema, reference terms stay single-step while proprioceptive terms carry -history. +Do not lag command/reference terms unless the training owner did so. ## Deploy-Side Measurements @@ -58,10 +48,10 @@ Record these per policy tick in the hardware runtime: 4. actuator command send timestamp 5. the action vector before and after clamp / smoothing -Compare the observation vector against a sim rollout built from the same -`deploy_config.yaml`. If the measured pipeline needs filtering or buffering, -encode the matching behavior in the task owner and re-export the deployment -artifacts. +Compare the observation vector against a sim rollout built from the same task +owner YAML. If the measured pipeline needs filtering or buffering, encode the +matching behavior in the task owner and retrain, rather than adding it only on +the deploy side. ## Symptoms Of Mismatch 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 219afa77c..bbf4e607a 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 @@ -27,12 +27,13 @@ Almost always one of: order in your motor driver. Use `unilab-export-scene` to dump the training joint order. 2. **Action scale unit mismatch.** Policy outputs unscaled values; the - driver expects rad, but you fed it normalized [-1, 1]. Apply - the `action_scale` / default-angle convention from `deploy_config.yaml` - before sending targets to the driver. -3. **Observation layout mismatch.** Compare `deploy_config.yaml` `obs_layout` - against the training owner and validate it with `scripts/deploy/sim_prototype.py` - before running on hardware. + driver expects rad, but you fed it normalized [-1, 1]. Apply the + `env.control_config.action_scale` / default-angle convention from the + training owner YAML before sending targets to the driver, reproducing the + owner's scalar-or-list form exactly. +3. **Observation layout mismatch.** Compare what your hardware loop assembles + against the training owner's `_build_actor_obs` — term order first, then + per-term history ordering. ## Cube drops in Allegro / Sharpa inhand @@ -55,7 +56,7 @@ investigation tomorrow takes 30 minutes instead of 4 hours: - Full hardware trace (`obs / action / wall_clock` for the entire run). - Sim-side YAML used to train: `runs//config.yaml`. -- `policy.onnx` and, for the G1 WBT path, `deploy_config.yaml`. +- `policy.onnx` and the exact task owner YAML path it was trained from. - One sim rollout video using **the same** seed: `eval --seed --render-mode record`. - A `git diff` between the run's commit and `main` if any. diff --git a/docs/sphinx/source/zh_CN/2-user_guide/4-tasks/2-motion_tracking.md b/docs/sphinx/source/zh_CN/2-user_guide/4-tasks/2-motion_tracking.md index cf7a16812..c17efd7df 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/4-tasks/2-motion_tracking.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/4-tasks/2-motion_tracking.md @@ -57,10 +57,10 @@ uv run train --algo sac --task g1_wbt_obs --sim mujoco training.use_amp=true `g1_wbt_obs` owner 是与部署对齐的 off-policy 观测配置：pelvis IMU 状态 （`pelvis_local_linvel` / `pelvis_gyro` / `pelvis_upvector`）加上 per-term 历史观测 -（`noise_config.obs_history_length: 5`），与部署侧的 `ObservationManager` 按字节对齐。 -部署工具在 `scripts/deploy/`，观测对齐由 `tests/scripts/test_obs_alignment_g1_wbt.py` -交叉校验。当 Motrix sim2sim 回放需要引用其他日志根目录下的 checkpoint 时，用 -`uv run eval` 透传绝对路径： +（`noise_config.obs_history_length: 5`），逐项按最旧优先展平，使得按逐项历史装配的 +硬件运行时读到同一个向量。该顺序由 `tests/scripts/test_obs_alignment_g1_wbt.py` +守护；硬件侧契约见仿真到真机部署指南。当 Motrix sim2sim 回放需要引用其他日志根目录下的 +checkpoint 时，用 `uv run eval` 透传绝对路径： ```bash uv run eval --algo sac --task g1_motion_tracking --sim motrix \ diff --git a/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/1-overview.md b/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/1-overview.md index 223c74a92..1e9879e9c 100644 --- a/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/1-overview.md +++ b/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/1-overview.md @@ -5,8 +5,7 @@ ## "仿真到真机"在 UniLab 中的含义 一个可部署的 UniLab 策略，是导出的策略加上所选任务 owner 使用的那套精确的观测与 -动作契约。G1 WBT 辅助路径将其物化为 `policy.onnx`、`deploy_config.yaml` 以及一个 -运动二进制文件；其他机器人需要一个等价的硬件侧运行时，它需要： +动作契约。UniLab 提供训练侧与 ONNX 导出；每种机器人都需要一个硬件侧运行时，它需要： 1. 读取传感器 → 组装出策略在仿真中看到的**同一个观测向量**。 2. 通过一个支持所导出计算图的运行时来运行 `policy.onnx`。 diff --git a/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/2-g1_whole_body.md b/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/2-g1_whole_body.md index df79c491e..9402d128b 100644 --- a/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/2-g1_whole_body.md +++ b/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/2-g1_whole_body.md @@ -2,19 +2,19 @@ ::::{admonition} 硬件目标 :class: note -Unitree G1 人形机器人（29 自由度变体）。假定关节顺序与 -`scripts/deploy/export_deploy_config.py` 从 -`src/unilab/assets/robots/g1/scene_flat.xml` 导出的顺序一致；在硬件上机前请核对该 -顺序。 +Unitree G1 人形机器人（29 自由度变体）。关节顺序来自任务 owner 的场景 +（`src/unilab/assets/robots/g1/scene_flat.xml`，按 actuator 顺序）；在硬件上机前请 +先核对该顺序与你的 SDK 电机索引是否一致。 :::: -本指南讲解从一个收敛的 G1 运动跟踪策略到机器人上闭环运行之间的**最后一公里**。 +本指南说明 G1 运动跟踪策略在硬件上所期望的**观测与动作契约**。仓库不提供 G1 部署侧 +运行时——硬件侧回路由你实现，本页告诉你它必须复现哪些内容。 ## 0. 验证你的仿真侧检查点 ```bash # Replay the policy headlessly and produce a video. -uv run eval --algo ppo --task g1_motion_tracking --sim motrix --load-run -1 \ +uv run eval --algo sac --task g1_wbt_obs --sim mujoco --load-run -1 \ --render-mode record ``` @@ -24,37 +24,48 @@ uv run eval --algo ppo --task g1_motion_tracking --sim motrix --load-run -1 \ - 关节速度与动作保持有限且在预期范围内。 - 接触时序看起来与参考运动一致。 -如果其中任何一项不对，请在硬件上机前修复仿真侧检查点或部署契约。 +如果其中任何一项不对，请在硬件上机前修复仿真侧检查点。 -## 1. 导出 +## 1. 先选定 owner，再从 YAML 读取契约 -使用训练回放路径导出 `policy.onnx`，然后用已提交的部署辅助工具导出 G1 WBT 部署配置 -与运动二进制文件： +硬件回路需要的每个字段都在任务 owner YAML 中声明。面向部署的 G1 owner 有： -```bash -uv run eval --algo ppo --task g1_motion_tracking --sim motrix --load-run -1 - -uv run scripts/deploy/export_deploy_config.py \ - --output logs/deploy/deploy_config.yaml - -uv run scripts/deploy/export_motion_bin.py \ - --output logs/deploy/dance1.bin +```{list-table} +:header-rows: 1 +:widths: 34 22 44 + +* - Owner + - Actor 观测宽度 + - 说明 +* - `conf/offpolicy/task/sac/g1_wbt_obs/mujoco.yaml` + - 514（H=5） + - 带 proprio 历史、无状态估计：丢弃 `base_lin_vel` 与 + `motion_anchor_pos_b`，使用 pelvis IMU。 +* - `conf/ppo/task/g1_motion_tracking_deploy/mujoco.yaml` + - 154（H=1） + - 单步 mimic actor 布局，逐关节 `action_scale` 列表。 ``` -部署侧原型消费如下文件： +::::{admonition} 观测宽度应从 env 读取，而不是照抄本表 +:class: warning +Actor 观测宽度是 owner `noise_config` 各开关的函数——`g1_wbt_obs` 见 +`src/unilab/envs/motion_tracking/g1/tracking_obs.py` 的 `_actor_obs_dim`， +deploy owner 见 `src/unilab/envs/motion_tracking/common/observations.py` 的 +`mimic_actor_obs_dim`。如果 ONNX 输入宽度与硬件回路装配出的宽度不一致，那是契约 +bug，而不是硬件调参问题。 +:::: + +通过训练回放路径导出 `policy.onnx`： -``` -runs// -└── policy.onnx -logs/deploy/ -├── deploy_config.yaml -└── dance1.bin +```bash +uv run eval --algo sac --task g1_wbt_obs --sim mujoco --load-run -1 ``` ## 2. 观测契约 -对于已提交的 G1 WBT 部署辅助工具，观测布局会作为 `obs_layout` 导出到 -`deploy_config.yaml`。`scripts/deploy/export_deploy_config.py` 是分段顺序的权威来源： +对 `g1_wbt_obs`，actor 观测按如下顺序装配（见 +`src/unilab/envs/motion_tracking/g1/tracking_obs.py` 的 `_build_actor_obs`）。 +单步参考项在前，随后是各 proprio 项的完整历史，按**最旧优先**展平： ```{list-table} :header-rows: 1 @@ -74,30 +85,37 @@ logs/deploy/ - 来自参考帧与机器人躯干帧的锚点朝向项 * - `gyro` - 每个历史步 3 - - IMU 陀螺仪项 + - IMU 陀螺仪项（`env.sensor.gyro`，该 owner 为 `pelvis_gyro`） * - `joint_pos_rel` - 每个历史步 29 - - 测量到的关节位置减去 `default_angles` + - 测量到的关节位置减去 `stand` keyframe 的关节角 * - `dof_vel` - 每个历史步 29 - 关节速度项 * - `last_actions` - - 29 + - 每个历史步 29 - 上一步的原始 actor 输出 ``` -导出脚本还会记录每个分段的 `history_length` 并校验总的 `obs_dim`。当 ONNX 输入宽度 -与 `deploy_config.yaml` 的 `obs_dim` 不一致时，`scripts/deploy/sim_prototype.py` 会 -拒绝运行。 +历史深度 `H` 即 `env.noise_config.obs_history_length`（该 owner 为 5）。逐项的 +最旧优先顺序由 `tests/scripts/test_obs_alignment_g1_wbt.py` 守护；硬件侧必须镜像 +该顺序，否则策略读到的是被置换过的向量。 ## 3. 执行器接口 -G1 部署原型将 actor 输出严格映射为： -`action * action_scale + default_angles`，然后钳制到 `joint_lower` / -`joint_upper`，并应用来自 `ema_alpha` 的 EMA 平滑。 +将 actor 输出映射为 `action * action_scale + default_angles`，然后在目标到达电机 +驱动器之前钳制到场景的关节范围内。 + +- `action_scale` 即 owner YAML 中的 `env.control_config.action_scale`。它可能是 + **标量**（`g1_wbt_obs` 为 2.0），也可能是**逐关节列表** + （`g1_motion_tracking_deploy` 为 29 项）。必须原样复现 owner 的形态——不要对列表 + 取平均、取首项，也不要把标量广播到列表 owner 上。 +- `default_angles` 是 owner 场景中 `stand` keyframe 的关节段。 +- 关节限位与增益同样来自该场景 XML（`jnt_range`、position actuator 的 + `gainprm` / `biasprm`）。 -- 动作 = 目标关节位置，按 `deploy_config.yaml` 中的 `action_scale` 项**缩放**。 -- 在目标到达电机驱动器之前，将其钳制到生成的关节范围内。 +训练侧直接施加目标，不做平滑。如果硬件抖动迫使你加入平滑，请先验证其 sim2sim +影响——每一步滞后都会把观测推离训练分布。 ## 4. 参考运动同步 @@ -118,7 +136,7 @@ G1 部署原型将 actor 输出严格映射为： 硬件侧：标准结构见 {doc}`7-safety_layers`。G1 的具体事项： - 在应用 `action_scale` 之前拒绝非有限动作与形状不匹配。 -- 用 `deploy_config.yaml` 中的 `joint_lower` / `joint_upper` 钳制生成的目标。 +- 用 owner 场景 XML 中的关节范围钳制生成的目标。 - 把看门狗、姿态监控以及操作员停止阈值保留在部署控制器中，并独立于策略对它们进行 测试。 @@ -134,18 +152,9 @@ G1 部署原型将 actor 输出严格映射为： ## 7. 应记录什么 -为每一步记录**完整的观测向量**、**完整的动作向量**与**墙钟**。在硬件上机前，通过 -MuJoCo 部署原型用同一份 ONNX、部署配置与运动二进制进行验证： - -```bash -uv run scripts/deploy/sim_prototype.py \ - --onnx runs//policy.onnx \ - --config logs/deploy/deploy_config.yaml \ - --motion logs/deploy/dance1.bin -``` - -ONNX 输入宽度与 `deploy_config.yaml` 的 `obs_dim` 之间的不匹配是部署契约 bug，而不是 -硬件调参问题。 +为每一步记录**完整的观测向量**、**完整的动作向量**与**墙钟**。把第一段硬件观测窗口 +与用同一份 owner YAML 构建的仿真回合作对比——这个 diff 比任何奖励检查都更快定位 +单位、坐标系与顺序错误。 ## 另请参阅 diff --git a/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/5-onnx_runtime.md b/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/5-onnx_runtime.md index 198f5e902..1a3ebe27b 100644 --- a/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/5-onnx_runtime.md +++ b/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/5-onnx_runtime.md @@ -24,36 +24,18 @@ uv run eval --algo sac --task g1_walk_flat --sim mujoco --load-run -1 ``` `uv run eval` 设置回放模式，并把 `--load-run` 映射到所路由训练脚本使用的检查点 -选择器。导出的文件会写入所选的运行目录。对于部署原型，请把导出的 `policy.onnx` 与 -运行时所用的部署侧配置和运动资产放在一起。 +选择器。导出的文件会写入所选的运行目录。用于部署时，请把导出的 `policy.onnx` 与训练 +它所用的任务 owner YAML 放在一起——该 YAML 是运行时必须复现的观测与动作契约的权威 +来源。 -## G1 部署原型 +## 校验导出的计算图 -已提交的 G1 WBT 部署辅助工具使用如下产物： +回放路径在写出之前会把导出的计算图与 PyTorch 比对，因此导出成功本身已经建立了数值 +一致性。它**没有**建立的是：你的硬件侧回路装配出的输入向量是否相同。在硬件上机前： -| 产物 | 生产者 | -| --- | --- | -| `policy.onnx` | 上述训练回放导出。 | -| `deploy_config.yaml` | `scripts/deploy/export_deploy_config.py`。 | -| `dance1.bin` 或其他运动二进制 | `scripts/deploy/export_motion_bin.py`。 | - -验证运行示例： - -```bash -uv run scripts/deploy/export_deploy_config.py \ - --output logs/deploy/deploy_config.yaml - -uv run scripts/deploy/export_motion_bin.py \ - --output logs/deploy/dance1.bin - -uv run scripts/deploy/sim_prototype.py \ - --onnx runs//policy.onnx \ - --config logs/deploy/deploy_config.yaml \ - --motion logs/deploy/dance1.bin -``` - -`scripts/deploy/sim_prototype.py` 会检查 ONNX 输入宽度是否与 `deploy_config.yaml` 中的 -`obs_dim` 匹配，然后用部署侧期望的同一观测布局在 MuJoCo 中驱动策略。 +- 从 env 读取 actor 观测宽度（而不是照抄文档表格），确认它与 ONNX 输入宽度一致。 +- 对照 owner 的 `_build_actor_obs` 确认分项顺序与逐项历史顺序。G1 全身跟踪见 + {doc}`2-g1_whole_body`。 ## 另请参阅 diff --git a/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/7-safety_layers.md b/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/7-safety_layers.md index 0c5c25df0..0e1bbb7b3 100644 --- a/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/7-safety_layers.md +++ b/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/7-safety_layers.md @@ -40,34 +40,31 @@ flowchart LR OP -.->|E-stop| D ``` -把硬实时的安全检查放在部署控制器中，而不是训练脚本里。仓库的 G1 辅助路径导出部署 -配置并运行一个 MuJoCo 原型；它并不实现生产级的电机驱动器安全回路。 +把硬实时的安全检查放在部署控制器中，而不是训练脚本里。仓库不实现生产级的电机驱动器 +安全回路——该边界由你构建并测试。 ## 策略假定你已配置的内容 -G1 部署辅助工具会把这些字段导出到 `deploy_config.yaml`： +策略期望的是其训练 owner 所声明的动作映射与限位。对 G1 WBT owner +（`conf/offpolicy/task/sac/g1_wbt_obs/mujoco.yaml`）： -```yaml -action_scale: 2.0 -ema_alpha: 1.0 -default_angles: [...] -joint_lower: [...] -joint_upper: [...] -kp: [...] -kd: [...] -``` +| 量 | 权威来源 | +| --- | --- | +| `action_scale` | `env.control_config.action_scale`（该 owner 为标量 `2.0`；其他 owner 声明逐关节列表） | +| `default_angles` | owner 场景 XML 中 `stand` keyframe 的关节段 | +| 关节限位 | 场景 XML 中的 `jnt_range` | +| `kp` / `kd` | 场景 XML 中 position actuator 的 `gainprm` / `biasprm` | -`scripts/deploy/sim_prototype.py` 消费同样的字段，并应用 -`action * action_scale + default_angles`、关节钳制与 EMA 平滑。硬件控制器应当消费 -生成的配置，而不是手动复制关节范围或增益。 +请从 owner YAML 及其场景派生这些量，并原样复现 owner 的 `action_scale` **形态**—— +标量 owner 与逐关节列表 owner 不可互换。不要把关节范围或增益手工复制到第二处，那会 +与资产静默漂移。 ## 交接测试 在把 策略 → 安全层 → 电机 集成起来之前，先隔离测试安全层： 1. 注入一个 NaN 动作，验证该指令被拒绝。 -2. 注入一个超范围的关节目标，验证钳制使用了 `deploy_config.yaml` 中的 - `joint_lower` / `joint_upper`。 +2. 注入一个超范围的关节目标，验证钳制使用了 owner 场景 XML 中的关节范围。 3. 在运行途中切断策略输入，验证控制器进入其配置的安全状态。 ## 另请参阅 diff --git a/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/8-latency_budget.md b/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/8-latency_budget.md index 0d972f6f6..6f37f6179 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 @@ -8,9 +8,9 @@ | 面 | 仓库证据 | 它覆盖什么 | | --- | --- | --- | | 单步动作延迟 | locomotion 与 G1 运动跟踪环境中的 `control_config.simulate_action_latency` | 执行上一步动作而非当前动作。 | -| G1 WBT 观测历史 | `noise_config.obs_history_length` 与 `scripts/deploy/export_deploy_config.py` | 为 `gyro`、`joint_pos_rel`、`dof_vel` 与 `last_actions` 导出逐项的 `obs_layout` 历史。 | +| G1 WBT 观测历史 | `conf/offpolicy/task/sac/g1_wbt_obs/mujoco.yaml` 中的 `noise_config.obs_history_length` | 为 `gyro`、`joint_pos_rel`、`dof_vel` 与 `last_actions` 提供逐项历史。 | | Sharpa 触觉接触延迟 | Sharpa 手内配置中的 `domain_rand.contact_latency` | 为采样到的接触通道保留上一步的触觉接触值。 | -| 部署侧 ONNX 契约检查 | `scripts/deploy/sim_prototype.py` | 为 G1 WBT 路径校验 `obs_layout`、`obs_dim`、ONNX 输入宽度、钳制以及 EMA 动作平滑。 | +| 观测历史顺序守护 | `tests/scripts/test_obs_alignment_g1_wbt.py` | 断言 G1 WBT actor 观测按逐项最旧优先展平。 | ## 动作延迟 @@ -28,21 +28,11 @@ env: ## 观测滞后与历史 -G1 WBT 部署辅助工具不会猜测观测宽度。它们导出一份带有 `obs_layout`、逐项 -`history_length` 与 `obs_dim` 的模式；随后原型组装出同样的布局，并拒绝不匹配。 +观测宽度是 owner `noise_config` 的函数，硬件运行时不允许猜测。对 G1 WBT owner， +`obs_history_length: 5` 让每个本体感受项携带 5 步历史并按最旧优先展平，而参考项保持 +单步。完整的分项顺序见 {doc}`2-g1_whole_body`。 -```bash -uv run scripts/deploy/export_deploy_config.py \ - --output logs/deploy/deploy_config.yaml - -uv run scripts/deploy/sim_prototype.py \ - --onnx runs//policy.onnx \ - --config logs/deploy/deploy_config.yaml \ - --motion logs/deploy/dance1.bin -``` - -除非训练 owner 这样做了，否则不要让指令/参考项滞后。在 G1 WBT 模式中，参考项保持 -单步，而本体感受项携带历史。 +除非训练 owner 这样做了，否则不要让指令/参考项滞后。 ## 部署侧测量 @@ -54,8 +44,8 @@ uv run scripts/deploy/sim_prototype.py \ 4. 执行器指令的发送时间戳 5. 钳制 / 平滑前后的动作向量 -将观测向量与用同一份 `deploy_config.yaml` 构建的仿真回合作对比。如果实测管线需要 -滤波或缓冲，请把相匹配的行为编码到任务 owner 中，并重新导出部署产物。 +将观测向量与用同一份任务 owner YAML 构建的仿真回合作对比。如果实测管线需要滤波或 +缓冲，请把相匹配的行为编码到任务 owner 中并重新训练，而不是只在部署侧添加。 ## 不匹配的症状 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 f601b5560..12593a21a 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 @@ -25,10 +25,11 @@ 1. **关节顺序被调换。** 检查 `policy.onnx` 的输入宽度与你电机驱动器中的关节顺序。 用 `unilab-export-scene` 导出训练时的关节顺序。 2. **动作缩放单位不匹配。** 策略输出未缩放的值；驱动器期望的是弧度，而你喂给它的 - 是归一化的 [-1, 1]。在把目标发送给驱动器之前，应用 `deploy_config.yaml` 中的 - `action_scale` / 默认角度约定。 -3. **观测布局不匹配。** 将 `deploy_config.yaml` 的 `obs_layout` 与训练 owner 比较， - 并在硬件上运行前用 `scripts/deploy/sim_prototype.py` 进行验证。 + 是归一化的 [-1, 1]。在把目标发送给驱动器之前，应用训练 owner YAML 中的 + `env.control_config.action_scale` / 默认角度约定，并原样复现该 owner 的标量或 + 列表形态。 +3. **观测布局不匹配。** 将硬件回路装配出的内容与训练 owner 的 `_build_actor_obs` + 比较——先比分项顺序，再比逐项历史顺序。 ## Allegro / Sharpa 手内操作中方块掉落 @@ -51,7 +52,7 @@ - 完整的硬件轨迹（整段运行的 `obs / action / wall_clock`）。 - 用于训练的仿真侧 YAML：`runs//config.yaml`。 -- `policy.onnx`，以及对于 G1 WBT 路径的 `deploy_config.yaml`。 +- `policy.onnx`，以及训练它所用的确切任务 owner YAML 路径。 - 一段使用**同一**种子的仿真回合视频：`eval --seed --render-mode record`。 - 如果有的话，该运行所在 commit 与 `main` 之间的 `git diff`。 diff --git a/scripts/deploy/append_cooldown.py b/scripts/deploy/append_cooldown.py deleted file mode 100644 index 8ad20eb8d..000000000 --- a/scripts/deploy/append_cooldown.py +++ /dev/null @@ -1,299 +0,0 @@ -#!/usr/bin/env python3 -"""Append a dance-final-frame->FixStand cooldown suffix to a WBT motion bin. - -Why this exists ---------------- -Symmetric to prepend_warmup.py. The original dance bin's last frame is mid- -motion (joint vel ~18 rad/s L2, pelvis still descending at 0.3 m/s, joints -~42 deg from FixStand stand qpos). When State_WBT's time-end check fires -and the FSM hands off to State_FixStand, FixStand only does a joint-space PD -interpolation from current motor q -> stand qpos over 3 s; it has no balance -controller, so a robot left with angular momentum and an off-balance pose -collapses despite the PD targets pulling it back to stand. - -Fix: append N seconds of kinematically self-consistent interpolation frames -at the tail of the bin. The tracking policy decelerates joints, brings the -torso back to upright, and lands on FixStand stand qpos before the FSM exit -triggers — so the FixStand interpolation starts from an already-balanced -pose and just holds station. - -What this does NOT change -------------------------- -- Training pipeline (no retrain needed). -- State_WBT.cpp / deploy_config.yaml / FSM yaml time_end (leave null -> full - duration so the cooldown plays automatically). -- The original dance frames (they are preserved verbatim before the cooldown). - -Interpolation scheme — shared primitives from motion_primitives.py, run in -reverse direction: -- joint_pos (J=29): cubic Hermite per joint with boundaries - (orig.jp[-1], orig.jv[-1]) -> (default_angles, 0) - joint_vel is the analytic derivative. -- body_pos (B,3): cubic Hermite per axis with boundaries - (orig.bp[-1], orig.bv[-1]) -> (FK(stand), 0) - body_lin_vel is the analytic derivative. -- body_quat (B,4): SLERP along quintic smoothstep s(u) = 6u^5 - 15u^4 + 10u^3 - (s'(0) = s''(0) = s'(1) = s''(1) = 0). Shortest-path; near-parallel fallback. -- body_ang_vel: central difference over [orig.bq[-2:], cooldown_bq] so the - seam ang_vel is continuous with the original bin's final ang_vel. - -Use ---- - uv run scripts/deploy/append_cooldown.py \ - --input ../deploy_ws/assets/dance1.bin \ - --output ../deploy_ws/assets/dance1_cooldown.bin \ - --config ../deploy_ws/assets/deploy_config.yaml \ - --cooldown-sec 2.0 - -Validate in sim BEFORE swapping on the real robot: - uv run scripts/deploy/sim_prototype.py \ - --motion ../deploy_ws/assets/dance1_cooldown.bin \ - --init-mode stand --render -""" - -from __future__ import annotations - -import argparse -import sys -from pathlib import Path - -import numpy as np -import yaml -from motion_primitives import ( - compute_fixstand_body_states, - hermite, - load_motion_bin, - quat_seq_ang_vel, - save_motion_bin, - slerp_smoothstep, -) - -REPO_ROOT = Path(__file__).resolve().parents[2] -DEPLOY_WS = REPO_ROOT.parent / "deploy_ws" -DEFAULT_SCENE = REPO_ROOT / "src/unilab/assets/robots/g1/scene_flat.xml" -DEFAULT_CFG = DEPLOY_WS / "assets/deploy_config.yaml" -DEFAULT_IN = DEPLOY_WS / "assets/dance1.bin" -DEFAULT_OUT = DEPLOY_WS / "assets/dance1_cooldown.bin" -DEFAULT_COOLDOWN_SEC = 3.0 -DEFAULT_HOLD_SEC = 0.5 - - -# ---------------------------------------------------------------------------- -# Main -# ---------------------------------------------------------------------------- - - -def main() -> None: - ap = argparse.ArgumentParser( - description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter - ) - ap.add_argument("--input", type=Path, default=DEFAULT_IN) - ap.add_argument("--output", type=Path, default=DEFAULT_OUT) - ap.add_argument( - "--config", - type=Path, - default=DEFAULT_CFG, - help="deploy_config.yaml (for default_angles, tracked ids).", - ) - ap.add_argument( - "--scene", - type=Path, - default=DEFAULT_SCENE, - help="MuJoCo XML with 'stand' keyframe (for FixStand FK).", - ) - ap.add_argument( - "--cut-frame", - type=int, - default=None, - help="If set, truncate the input bin so that frame " - "CUT_FRAME becomes the new last frame (inclusive). " - "Frames after that are dropped, and the cooldown ramps " - "from this frame to default. Use this when the dance's " - "actual last frame has too much momentum for any post-" - "hoc interpolation to stabilize. Find a low-energy " - "frame (low joint vel, near-upright torso, near-zero " - "pelvis vertical vel) and cut there.", - ) - ap.add_argument( - "--cooldown-sec", - type=float, - default=DEFAULT_COOLDOWN_SEC, - help="Length of Hermite deceleration ramp [seconds]. " - "Lowers joint vel from orig.jv[-1] to 0 and joint pos " - "to default_angles over this interval.", - ) - ap.add_argument( - "--hold-sec", - type=float, - default=DEFAULT_HOLD_SEC, - help="Length of trailing static-hold segment [seconds]. " - "Appended AFTER the Hermite ramp; frames are pure " - "(default_angles, 0, FixStand FK body, 0). Lets the " - "tracking policy physically converge to stand before " - "the FSM timeout swaps to State_FixStand. Set to 0 to " - "disable. Recommended >= 0.3s.", - ) - args = ap.parse_args() - - with open(args.config) as f: - cfg = yaml.safe_load(f) - default_angles = np.asarray(cfg["default_angles"], dtype=np.float64) - tracked_ids = list(cfg["tracked_body_mujoco_ids"]) - - orig = load_motion_bin(args.input) - fps, J, B = orig["fps"], orig["nj"], orig["nb"] - if J != len(default_angles): - raise SystemExit(f"J mismatch: bin J={J}, default_angles len={len(default_angles)}") - if B != len(tracked_ids): - raise SystemExit(f"B mismatch: bin B={B}, tracked_body_mujoco_ids len={len(tracked_ids)}") - - if args.cut_frame is not None: - cf = int(args.cut_frame) - if cf < 1 or cf >= orig["nf"]: - raise SystemExit(f"cut-frame={cf} out of range [1, {orig['nf'] - 1}]") - # Keep frames [0, cf] (cf becomes the new last frame, inclusive). - keep = cf + 1 - print( - f"Truncating input: keep frames [0, {cf}] " - f"({orig['nf']} -> {keep} frames, {orig['nf'] / fps:.2f}s -> {keep / fps:.2f}s)", - file=sys.stderr, - ) - for k in ("jp", "jv", "bp", "bq", "bv", "bav"): - orig[k] = orig[k][:keep] - orig["nf"] = keep - # Report cut-point energy so the user sees what cool-down starts from. - a = cfg.get("anchor_body_idx_in_tracked", 7) - jv_l2 = float(np.linalg.norm(orig["jv"][-1])) - jp_dev = float(np.abs(orig["jp"][-1] - default_angles).max()) - bv_z = float(abs(orig["bv"][-1, a, 2])) - print( - f" cut-point energy: |jv|={jv_l2:.2f} rad/s jp_dev={jp_dev:.3f} rad " - f"|bv_z|={bv_z:.3f} m/s", - file=sys.stderr, - ) - - dt = 1.0 / fps - N = int(round(args.cooldown_sec * fps)) - if N < 2: - raise SystemExit( - f"cooldown-sec={args.cooldown_sec}s too short for fps={fps} (need >= 2 frames)" - ) - K = int(round(args.hold_sec * fps)) if args.hold_sec > 0.0 else 0 - if args.hold_sec > 0.0 and K < 1: - raise SystemExit( - f"hold-sec={args.hold_sec}s rounds to 0 frames at fps={fps}; pass 0 or >= {dt:.3f}s" - ) - T = N * dt # frame N of the cooldown == FixStand target exactly - - fk_jp, fk_bp, fk_bq = compute_fixstand_body_states(args.scene, tracked_ids) - - keyframe_diff = float(np.abs(fk_jp - default_angles).max()) - if keyframe_diff > 1e-3: - print( - f"WARN: 'stand' keyframe joint pos differs from deploy_config " - f"default_angles by {keyframe_diff:.4f} rad — using deploy_config " - "values for cooldown end (so command_joint_pos matches what " - "FixStand will hold on the real robot).", - file=sys.stderr, - ) - - # --- joint pos/vel: analytic Hermite, orig[-1] -> default_angles -------- - # ts[k] = (k+1)*dt so the first cooldown frame is one dt after orig[-1] - # and the last cooldown frame (k = N-1) lands exactly at t = T = N*dt, - # i.e. precisely on the (default_angles, 0) boundary. - ts = (np.arange(N) + 1) * dt - jp_c, jv_c = hermite(orig["jp"][-1], orig["jv"][-1], default_angles, np.zeros(J), ts, T) - - # --- body pos / lin_vel: analytic Hermite per axis ---------------------- - bp_c, bv_c = hermite(orig["bp"][-1], orig["bv"][-1], fk_bp, np.zeros((B, 3)), ts, T) - - # --- body quat: SLERP along quintic smoothstep -------------------------- - u = ts / T - bq_c = np.zeros((N, B, 4), dtype=np.float64) - for b in range(B): - bq_c[:, b, :] = slerp_smoothstep(orig["bq"][-1, b], fk_bq[b], u) - - # --- body ang_vel: central diff, with the seam differenced against the - # ORIGINAL bin's last two frames so the derivative is continuous across - # the dance/cooldown boundary ------------------------------------------- - bav_c = np.zeros((N, B, 3), dtype=np.float64) - for b in range(B): - ext = np.concatenate([orig["bq"][-2:, b, :], bq_c[:, b, :]], axis=0) - # extended sequence has 2+N frames; cooldown frames live at [2, 2+N). - bav_c[:, b, :] = quat_seq_ang_vel(ext, dt)[2:] - - # --- static hold segment: K frames all at (default, 0, FK, 0). Gives the - # tracking policy time to physically converge to the FixStand target before - # the FSM time-end check fires and hands off to State_FixStand. Without - # this segment, the cooldown's last (default, 0) frame is commanded for - # only one step_dt (~0.02s) before the swap, and any residual body momentum - # carried through the Hermite ramp ends up as the FixStand starting pose. - - if K > 0: - jp_h = np.broadcast_to(default_angles, (K, J)).copy() - jv_h = np.zeros((K, J), dtype=np.float64) - bp_h = np.broadcast_to(fk_bp, (K, B, 3)).copy() - bq_h = np.broadcast_to(fk_bq, (K, B, 4)).copy() - bv_h = np.zeros((K, B, 3), dtype=np.float64) - bav_h = np.zeros((K, B, 3), dtype=np.float64) - - # --- concatenate original + cooldown (+ hold) ------------------------- - parts_jp = [orig["jp"], jp_c] - parts_jv = [orig["jv"], jv_c] - parts_bp = [orig["bp"], bp_c] - parts_bq = [orig["bq"], bq_c] - parts_bv = [orig["bv"], bv_c] - parts_bav = [orig["bav"], bav_c] - if K > 0: - parts_jp.append(jp_h) - parts_jv.append(jv_h) - parts_bp.append(bp_h) - parts_bq.append(bq_h) - parts_bv.append(bv_h) - parts_bav.append(bav_h) - new_jp = np.concatenate(parts_jp, axis=0) - new_jv = np.concatenate(parts_jv, axis=0) - new_bp = np.concatenate(parts_bp, axis=0) - new_bq = np.concatenate(parts_bq, axis=0) - new_bv = np.concatenate(parts_bv, axis=0) - new_bav = np.concatenate(parts_bav, axis=0) - - args.output.parent.mkdir(parents=True, exist_ok=True) - save_motion_bin(args.output, fps, new_jp, new_jv, new_bp, new_bq, new_bv, new_bav) - - # --- report ------------------------------------------------------------ - # Seam = boundary between orig[-1] and cooldown[0], which are one dt apart. - seam_jp = float(np.abs(jp_c[0] - orig["jp"][-1]).max()) - seam_jv = float(np.abs(jv_c[0] - orig["jv"][-1]).max()) - seam_bp = float(np.abs(bp_c[0] - orig["bp"][-1]).max()) - seam_bv = float(np.abs(bv_c[0] - orig["bv"][-1]).max()) - seam_qang_deg = [] - for b in range(B): - d = float(np.clip(abs(np.dot(bq_c[0, b], orig["bq"][-1, b])), 0.0, 1.0)) - seam_qang_deg.append(np.degrees(2 * np.arccos(d))) - seam_qang_max = max(seam_qang_deg) - # End state must match FixStand boundary (this is the whole point). - end_jp_diff = float(np.abs(new_jp[-1] - default_angles).max()) - end_jv_l2 = float(np.linalg.norm(new_jv[-1])) - - print(f"Wrote {args.output}") - print(f" fps={fps} J={J} B={B}") - print( - f" frames: {orig['nf']} (orig) -> {new_jp.shape[0]} " - f"({orig['nf']} dance + {N} cooldown + {K} hold)" - ) - print( - f" duration: {orig['nf'] / fps:.2f}s -> {new_jp.shape[0] / fps:.2f}s " - f"(cooldown_sec={args.cooldown_sec}, hold_sec={args.hold_sec})" - ) - print(f" end (last frame) command_joint_pos == default_angles? max_diff={end_jp_diff:.6f} rad") - print(f" end (last frame) command_joint_vel L2 = {end_jv_l2:.6f} rad/s (should be 0)") - print(" seam check (orig last frame vs cooldown first frame; one-dt-step apart):") - print(f" |Δjoint_pos|_∞ = {seam_jp:.4f} rad") - print(f" |Δjoint_vel|_∞ = {seam_jv:.4f} rad/s") - print(f" |Δbody_pos|_∞ = {seam_bp:.4f} m") - print(f" |Δbody_lin_vel|_∞ = {seam_bv:.4f} m/s") - print(f" max body quat angle = {seam_qang_max:.4f} deg") - - -if __name__ == "__main__": - main() diff --git a/scripts/deploy/export_deploy_config.py b/scripts/deploy/export_deploy_config.py deleted file mode 100644 index b692e992e..000000000 --- a/scripts/deploy/export_deploy_config.py +++ /dev/null @@ -1,328 +0,0 @@ -#!/usr/bin/env python3 -"""Export deploy_config.yaml for the C++ G1-29DOF WBT deployment side. - -Reads g1.xml + scene_flat.xml + tracking.py defaults to emit a single yaml -that the deploy framework (~/deploy_ws/unitree_rl_lab/.../State_WBT) can load -to drive the actor at runtime. - -This file is the SINGLE SOURCE OF TRUTH for the actor obs schema: - * Training side (tracking.py) assembles obs in the order documented here. - * Deploy side (State_WBT.cpp + ObservationManager) reads obs_layout from - this yaml and assembles in matching order with per-term history buffers. - * Alignment test (tests/test_obs_alignment_g1_wbt.py) verifies both code - paths produce byte-identical obs from the same inputs. - -obs_layout schema v2 (per-term history_length): - Each entry: {name, dim, history_length, source} - * Reference terms (command_*, motion_anchor_*) use history_length=1 - (single-step, refs come fresh from the motion clip every tick). - * Proprio terms (gyro, joint_pos_rel, dof_vel, last_actions) carry - history_length=H (5 for the deploy profile), flattened oldest-first. - * Deploy framework reads via use_gym_history=false (group-by-term mode), - so each term independently flattens its full history → total obs_dim is - the sum of dim * history_length over all entries. -""" - -from __future__ import annotations - -import argparse -from pathlib import Path - -import mujoco -import yaml - -REPO_ROOT = Path(__file__).resolve().parents[2] -DEFAULT_SCENE = REPO_ROOT / "src/unilab/assets/robots/g1/scene_flat.xml" -DEFAULT_OUT = REPO_ROOT / "logs/deploy/deploy_config.yaml" - -TRACKED_BODY_NAMES = ( - "pelvis", - "left_hip_roll_link", - "left_knee_link", - "left_ankle_roll_link", - "right_hip_roll_link", - "right_knee_link", - "right_ankle_roll_link", - "torso_link", - "left_shoulder_roll_link", - "left_elbow_link", - "left_wrist_yaw_link", - "right_shoulder_roll_link", - "right_elbow_link", - "right_wrist_yaw_link", -) -ANCHOR_BODY_NAME = "torso_link" - -ACTION_SCALE = 2.0 -# EMA alpha for q_target smoothing on the deploy side. Training applies q_target -# directly with no smoothing; alpha=1.0 means deploy also applies directly (best -# for sim2sim correctness). Lower it (~0.5–0.8) only on real hardware if jitter -# requires smoothing — but every step of lag pushes obs out of training -# distribution, so verify sim2sim impact at the chosen alpha first. -EMA_ALPHA = 1.0 -CTRL_DT = 0.02 -KEYFRAME_NAME = "stand" -ROOT_QPOS_DIM = 7 # free joint: xyz + quat(wxyz) - -# Default obs history length — matches g1_wbt_obs/mujoco.yaml's -# `noise_config.obs_history_length`. Override via --obs-history-length when -# exporting for other training profiles (e.g. g1_motion_tracking/mujoco.yaml uses 1). -DEFAULT_OBS_HISTORY_LENGTH = 5 - - -def _round_list(arr, ndigits=6): - return [round(float(v), ndigits) for v in arr] - - -def _build_obs_layout( - num_action: int, hist_len: int, enable_zero_anchor_pos: bool, enable_zero_linvel: bool -): - """Build obs_layout in the exact order tracking.py:_compute_obs assembles. - - Returns (layout_list, total_obs_dim). - Order = single-step refs first, then per-term proprio history blocks, - matching the training-side actor obs concatenation order. - """ - layout = [ - # ---- single-step reference terms (history_length=1) ---- - { - "name": "command_joint_pos", - "dim": num_action, - "history_length": 1, - "source": "motion_ref_frame.joint_pos", - }, - { - "name": "command_joint_vel", - "dim": num_action, - "history_length": 1, - "source": "motion_ref_frame.joint_vel", - }, - ] - if not enable_zero_anchor_pos: - layout.append( - { - "name": "motion_anchor_pos_b", - "dim": 3, - "history_length": 1, - "source": "subtract_frame(robot_anchor_w, motion_anchor_w).pos", - } - ) - layout.append( - { - "name": "motion_anchor_ori_b", - "dim": 6, - "history_length": 1, - "source": "rotation_matrix(subtract_frame(...).quat)[:, :2].flatten()", - } - ) - if not enable_zero_linvel: - layout.append( - { - "name": "base_lin_vel", - "dim": 3, - "history_length": 1, - "source": "imu.local_linvel (state-estimated)", - } - ) - # ---- proprio terms with H-step oldest-first history ---- - layout.extend( - [ - {"name": "gyro", "dim": 3, "history_length": hist_len, "source": "imu.gyroscope"}, - { - "name": "joint_pos_rel", - "dim": num_action, - "history_length": hist_len, - "source": "dof_pos - default_angles", - }, - {"name": "dof_vel", "dim": num_action, "history_length": hist_len, "source": "dof_vel"}, - { - "name": "last_actions", - "dim": num_action, - "history_length": hist_len, - "source": "previous raw actor output", - }, - ] - ) - total = sum(seg["dim"] * seg["history_length"] for seg in layout) - return layout, total - - -def main(): - ap = argparse.ArgumentParser( - description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter - ) - ap.add_argument( - "--scene", - type=Path, - default=DEFAULT_SCENE, - help="MuJoCo scene file containing the 'stand' keyframe.", - ) - ap.add_argument("--output", "-o", type=Path, default=DEFAULT_OUT) - ap.add_argument( - "--obs-history-length", - type=int, - default=DEFAULT_OBS_HISTORY_LENGTH, - help="Proprio history length H. Must match training-side " - "noise_config.obs_history_length. Default 5 = current " - "g1_wbt_obs/mujoco.yaml. Set 1 for the legacy 154-d schema.", - ) - ap.add_argument( - "--enable-zero-anchor-pos", - action="store_true", - default=True, - help="Drop motion_anchor_pos_b from actor obs (mjlab parity). " - "Matches g1_wbt_obs/mujoco.yaml's noise_config flag.", - ) - ap.add_argument( - "--enable-zero-linvel", - action="store_true", - default=True, - help="Drop base_lin_vel from actor obs (mjlab parity). " - "Matches g1_wbt_obs/mujoco.yaml's noise_config flag.", - ) - args = ap.parse_args() - - if not args.scene.exists(): - raise SystemExit(f"Scene not found: {args.scene}") - - model = mujoco.MjModel.from_xml_path(str(args.scene)) - - if model.nu != 29: - raise SystemExit(f"Expected 29 actuators, got {model.nu}") - - # Joint names in actuator order (action[i] drives actuator i). - joint_names = [ - mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_ACTUATOR, i) for i in range(model.nu) - ] - - # kp from XML actuators. - kp = model.actuator_gainprm[:, 0].copy() # gainprm[0] is kp for position actuator - if not (kp > 0).all(): - raise SystemExit(f"kp parsing produced non-positive values: {kp}") - # kv recovery: for position actuator, biasprm = [0, -kp, -kv] - kv = -model.actuator_biasprm[:, 2].copy() - if not (kv > 0).all(): - raise SystemExit(f"kv parsing produced non-positive values: {kv}") - - # Joint limits: skip the floating-root joint (jnt 0). - if model.njnt < 1 + model.nu: - raise SystemExit(f"Insufficient joints: njnt={model.njnt}") - jnt_range = model.jnt_range[1 : 1 + model.nu].copy() - joint_lower = jnt_range[:, 0] - joint_upper = jnt_range[:, 1] - - # Force range from actuator forcerange. - force_range = model.actuator_forcerange.copy() - force_lower = force_range[:, 0] - force_upper = force_range[:, 1] - - # default_angles = stand keyframe qpos[7:36]. - if model.nkey == 0: - raise SystemExit(f"No keyframes in {args.scene}; expected '{KEYFRAME_NAME}'.") - key_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_KEY, KEYFRAME_NAME) - if key_id < 0: - names = [mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_KEY, i) for i in range(model.nkey)] - raise SystemExit(f"Keyframe '{KEYFRAME_NAME}' not found; have {names}") - stand_qpos = model.key_qpos[key_id] - if len(stand_qpos) != ROOT_QPOS_DIM + model.nu: - raise SystemExit(f"stand qpos len {len(stand_qpos)} != {ROOT_QPOS_DIM}+{model.nu}") - default_angles = stand_qpos[ROOT_QPOS_DIM:].copy() - - # 14 tracked body indices (in MuJoCo body-id space; deploy side won't use - # these directly but they are useful for debugging / cross-checks). - tracked_body_ids = [] - for nm in TRACKED_BODY_NAMES: - bid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, nm) - if bid < 0: - raise SystemExit(f"Tracked body '{nm}' missing from model") - tracked_body_ids.append(int(bid)) - anchor_body_idx_in_tracked = TRACKED_BODY_NAMES.index(ANCHOR_BODY_NAME) - - # Build the obs layout (single source of truth for actor obs assembly). - obs_layout, obs_dim = _build_obs_layout( - num_action=model.nu, - hist_len=args.obs_history_length, - enable_zero_anchor_pos=args.enable_zero_anchor_pos, - enable_zero_linvel=args.enable_zero_linvel, - ) - - cfg = { - # ---- meta ---- - # obs_dim = sum over obs_layout of dim * history_length. - # For the deploy profile (H=5, both zero flags ON) this is: - # command_joint_pos(29*1) + command_joint_vel(29*1) - # + motion_anchor_ori_b(6*1) + gyro(3*5) - # + joint_pos_rel(29*5) + dof_vel(29*5) + last_actions(29*5) = 514 - # Aligns with deploy-profile training run (g1_wbt_obs/mujoco.yaml) which - # drops motion_anchor_pos_b and base_lin_vel from actor obs to match - # Unitree's verified mjlab "No-State-Estimation" deploy yaml and - # adds BeyondMimic-style proprio history (default H=5). - "obs_dim": obs_dim, - "obs_history_length": args.obs_history_length, - # Deploy ObservationManager mode. False = group-by-term (each term's - # full history flattened oldest-first, then concat across terms). - # True = group-by-time-step (legacy gym style). Training side uses - # the false convention, so MUST stay false here for byte alignment. - "use_gym_history": False, - "action_dim": int(model.nu), - "ctrl_dt": CTRL_DT, - "action_scale": ACTION_SCALE, - "ema_alpha": EMA_ALPHA, - # ---- joint config (in MuJoCo actuator order; deploy side assumes - # this matches the SDK motor index 1:1 for G1-29DOF — verify per-motor - # before real-robot run) ---- - "joint_names": list(joint_names), - # Identity SDK mapping; replace with measured order if 1:1 assumption fails. - "joint_ids_map": list(range(model.nu)), - "default_angles": _round_list(default_angles), - "kp": _round_list(kp), - "kd": _round_list(kv), - "joint_lower": _round_list(joint_lower), - "joint_upper": _round_list(joint_upper), - "force_lower": _round_list(force_lower, 3), - "force_upper": _round_list(force_upper, 3), - # ---- motion / anchor ---- - "tracked_body_names": list(TRACKED_BODY_NAMES), - "tracked_body_mujoco_ids": tracked_body_ids, - "anchor_body_name": ANCHOR_BODY_NAME, - "anchor_body_idx_in_tracked": int(anchor_body_idx_in_tracked), - # ---- noise (NOT applied at deploy; documentation only — per-step - # uniform noise scales used during training, plus persistent encoder - # bias absorbed into joint_pos_rel) ---- - "training_noise_scales": { - "joint_angle": 0.01, - "joint_vel": 0.5, - "gyro": 0.2, - "anchor_ori": 0.05, - "joint_pos_encoder_bias_per_episode": 0.01, - }, - # ---- obs layout (SINGLE SOURCE OF TRUTH for both ends) ---- - # Each entry: name, dim (per-step), history_length, source. - # Total obs_dim = sum(dim * history_length). - # Order MUST match tracking.py:_compute_obs assembly order. - # State_WBT.cpp:build_env_cfg translates names via its alias table. - "obs_layout": obs_layout, - } - - args.output.parent.mkdir(parents=True, exist_ok=True) - with open(args.output, "w") as f: - yaml.safe_dump(cfg, f, sort_keys=False, default_flow_style=None, width=120) - - print(f"Wrote {args.output} ({args.output.stat().st_size} bytes)") - print( - f" joints: {model.nu}, tracked bodies: {len(TRACKED_BODY_NAMES)}, " - f"anchor='{ANCHOR_BODY_NAME}' (idx_in_tracked={anchor_body_idx_in_tracked})" - ) - print(f" obs_history_length={args.obs_history_length}, total obs_dim={obs_dim}") - print(" obs_layout segments:") - for seg in obs_layout: - print( - f" {seg['name']:24s} dim={seg['dim']:3d} H={seg['history_length']:1d} " - f"contrib={seg['dim'] * seg['history_length']}" - ) - print(f" default_angles[:6] = {_round_list(default_angles[:6], 3)}") - print(f" kp[:3] = {_round_list(kp[:3], 3)}, kd[:3] = {_round_list(kv[:3], 3)}") - - -if __name__ == "__main__": - main() diff --git a/scripts/deploy/export_motion_bin.py b/scripts/deploy/export_motion_bin.py deleted file mode 100644 index 27905845c..000000000 --- a/scripts/deploy/export_motion_bin.py +++ /dev/null @@ -1,179 +0,0 @@ -#!/usr/bin/env python3 -"""Export a tracked-motion NPZ to the flat binary format used by State_WBT. - -Layout (little-endian, contiguous): - header (16 bytes): - int32 fps - int32 num_frames - int32 num_joints (29 for G1-29DOF) - int32 num_bodies (14 tracked bodies) - data (all float32, frames-major): - joint_pos [num_frames][num_joints] - joint_vel [num_frames][num_joints] - body_pos_w [num_frames][num_bodies][3] - body_quat_w [num_frames][num_bodies][4] (wxyz) - body_lin_vel_w [num_frames][num_bodies][3] - body_ang_vel_w [num_frames][num_bodies][3] - -NPZ source layout (per src/unilab/envs/motion_tracking/g1/motion_loader.py): - - 'fps' (int) - - 'joint_pos' (N, 29) - - 'joint_vel' (N, 29) - - 'body_pos_w', 'body_quat_w', 'body_lin_vel_w', 'body_ang_vel_w' (N, 31, *) - Body axis is in MuJoCo body-id order; we slice down to 14 tracked bodies - using mj_name2id() so the deploy side does not need to repeat the lookup. -""" - -from __future__ import annotations - -import argparse -import struct -from pathlib import Path - -import mujoco -import numpy as np - -REPO_ROOT = Path(__file__).resolve().parents[2] -DEFAULT_NPZ = REPO_ROOT / "src/unilab/assets/motions/g1/dance1_subject2_part.npz" -DEFAULT_SCENE = REPO_ROOT / "src/unilab/assets/robots/g1/scene_flat.xml" -DEFAULT_OUT = REPO_ROOT / "logs/deploy/dance1.bin" - -TRACKED_BODY_NAMES = ( - "pelvis", - "left_hip_roll_link", - "left_knee_link", - "left_ankle_roll_link", - "right_hip_roll_link", - "right_knee_link", - "right_ankle_roll_link", - "torso_link", - "left_shoulder_roll_link", - "left_elbow_link", - "left_wrist_yaw_link", - "right_shoulder_roll_link", - "right_elbow_link", - "right_wrist_yaw_link", -) - - -def main(): - ap = argparse.ArgumentParser(description=__doc__) - ap.add_argument("--motion", type=Path, default=DEFAULT_NPZ) - ap.add_argument("--scene", type=Path, default=DEFAULT_SCENE) - ap.add_argument("--output", "-o", type=Path, default=DEFAULT_OUT) - ap.add_argument( - "--start-frame", type=int, default=0, help="First frame to include (inclusive). Default 0." - ) - end_group = ap.add_mutually_exclusive_group() - end_group.add_argument( - "--end-frame", - type=int, - default=None, - help="One past the last frame (exclusive). Default = NPZ frame count.", - ) - end_group.add_argument( - "--duration", - type=float, - default=None, - help="Seconds to keep starting at --start-frame. Resolved to frames via NPZ fps.", - ) - args = ap.parse_args() - - if not args.motion.exists(): - raise SystemExit(f"NPZ not found: {args.motion}") - if not args.scene.exists(): - raise SystemExit(f"Scene not found: {args.scene}") - - model = mujoco.MjModel.from_xml_path(str(args.scene)) - body_ids = [] - for nm in TRACKED_BODY_NAMES: - bid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, nm) - if bid < 0: - raise SystemExit(f"Tracked body '{nm}' missing from model") - body_ids.append(int(bid)) - - with np.load(args.motion) as data: - fps = int(np.asarray(data["fps"]).reshape(-1)[0]) - joint_pos = data["joint_pos"].astype(np.float32) - joint_vel = data["joint_vel"].astype(np.float32) - body_pos_w = data["body_pos_w"].astype(np.float32) - body_quat_w = data["body_quat_w"].astype(np.float32) - body_lin_vel_w = data["body_lin_vel_w"].astype(np.float32) - body_ang_vel_w = data["body_ang_vel_w"].astype(np.float32) - - total_frames, num_joints = joint_pos.shape - num_bodies = len(TRACKED_BODY_NAMES) - - if joint_pos.shape != joint_vel.shape: - raise SystemExit(f"joint_pos {joint_pos.shape} != joint_vel {joint_vel.shape}") - if num_joints != 29: - raise SystemExit(f"Expected 29 joints, got {num_joints}") - if body_pos_w.shape[0] != total_frames: - raise SystemExit(f"body_pos_w frames {body_pos_w.shape[0]} != {total_frames}") - - start = args.start_frame - if args.duration is not None: - if args.duration <= 0: - raise SystemExit(f"--duration must be > 0, got {args.duration}") - end = start + int(round(args.duration * fps)) - elif args.end_frame is not None: - end = args.end_frame - else: - end = total_frames - if not (0 <= start < end <= total_frames): - raise SystemExit(f"Invalid frame range [{start}, {end}); valid is [0, {total_frames}]") - if (start, end) != (0, total_frames): - joint_pos = joint_pos[start:end] - joint_vel = joint_vel[start:end] - body_pos_w = body_pos_w[start:end] - body_quat_w = body_quat_w[start:end] - body_lin_vel_w = body_lin_vel_w[start:end] - body_ang_vel_w = body_ang_vel_w[start:end] - num_frames = end - start - if body_pos_w.shape[1] < max(body_ids) + 1: - raise SystemExit( - f"NPZ body axis len {body_pos_w.shape[1]} insufficient for max body_id {max(body_ids)}" - ) - - # Slice 31-body axis down to 14 tracked bodies. - body_pos_w = body_pos_w[:, body_ids] - body_quat_w = body_quat_w[:, body_ids] - body_lin_vel_w = body_lin_vel_w[:, body_ids] - body_ang_vel_w = body_ang_vel_w[:, body_ids] - - # Sanity check quaternion is wxyz (norm ≈ 1 and first element typically positive - # for "uprightish" frames). This is a soft check. - quat_norms = np.linalg.norm(body_quat_w[0], axis=1) - if not np.allclose(quat_norms, 1.0, atol=1e-3): - print(f"WARN: frame-0 quat norms not unit: {quat_norms}") - - args.output.parent.mkdir(parents=True, exist_ok=True) - with open(args.output, "wb") as f: - f.write(struct.pack(" dict: - with open(path, "rb") as f: - fps, nf, nj, nb = struct.unpack(" None: - nf, nj = jp.shape - nb = bp.shape[1] - assert jv.shape == jp.shape, "jv shape mismatch" - assert bp.shape == (nf, nb, 3) and bq.shape == (nf, nb, 4), "body shape mismatch" - assert bv.shape == (nf, nb, 3) and bav.shape == (nf, nb, 3), "body vel shape mismatch" - with open(path, "wb") as f: - f.write(struct.pack(" tuple[np.ndarray, np.ndarray, np.ndarray]: - model = mujoco.MjModel.from_xml_path(str(scene)) - data = mujoco.MjData(model) - key_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_KEY, "stand") - if key_id < 0: - raise SystemExit(f"'stand' keyframe not found in {scene}") - mujoco.mj_resetDataKeyframe(model, data, key_id) - mujoco.mj_forward(model, data) - jp = np.asarray(data.qpos[7:], dtype=np.float64).copy() - bp = np.stack([np.asarray(data.xpos[i], dtype=np.float64).copy() for i in tracked_ids]) - bq = np.stack( - [ - np.asarray(data.xquat[i], dtype=np.float64).copy() # wxyz - for i in tracked_ids - ] - ) - return jp, bp, bq - - -# ---------------------------------------------------------------------------- -# Cubic Hermite — analytic position and velocity, kinematically consistent. -# Boundary conditions: p(0)=p0, p'(0)=v0, p(T)=p1, p'(T)=v1. -# Returns (p(t), p'(t)) where t may be an array with shape (N,) and p* may -# have additional trailing dims (broadcasting). -# ---------------------------------------------------------------------------- - - -def hermite(p0, v0, p1, v1, t, T): - s = t / T - s2, s3 = s * s, s * s * s - h00 = 2 * s3 - 3 * s2 + 1 - h10 = s3 - 2 * s2 + s - h01 = -2 * s3 + 3 * s2 - h11 = s3 - s2 - h00d = 6 * s2 - 6 * s - h10d = 3 * s2 - 4 * s + 1 - h01d = -6 * s2 + 6 * s - h11d = 3 * s2 - 2 * s - - # Reshape s-dim coefficients for broadcasting against p*'s trailing dims. - def _r(a): - return a.reshape(a.shape + (1,) * (np.ndim(p0))) - - p = _r(h00) * p0 + _r(h10) * T * v0 + _r(h01) * p1 + _r(h11) * T * v1 - pd = _r(h00d / T) * p0 + _r(h10d) * v0 + _r(h01d / T) * p1 + _r(h11d) * v1 - return p, pd - - -# ---------------------------------------------------------------------------- -# Quaternion SLERP along quintic smoothstep s(u) = 6u^5 - 15u^4 + 10u^3. -# u in [0, 1]. Operates on (4,) wxyz quaternions, returns (N, 4). -# ---------------------------------------------------------------------------- - - -def slerp_smoothstep(q0: np.ndarray, q1: np.ndarray, u: np.ndarray) -> np.ndarray: - q0 = q0 / np.linalg.norm(q0) - q1 = q1 / np.linalg.norm(q1) - dot = float(np.dot(q0, q1)) - if dot < 0.0: # shortest path - q1 = -q1 - dot = -dot - s = 6 * u**5 - 15 * u**4 + 10 * u**3 - if dot > 0.9995: # near-parallel: lerp + renormalise - out = (1 - s)[:, None] * q0 + s[:, None] * q1 - return out / np.linalg.norm(out, axis=-1, keepdims=True) - theta_0 = np.arccos(np.clip(dot, -1.0, 1.0)) - sin_theta_0 = np.sin(theta_0) - theta = theta_0 * s - a = np.cos(theta) - dot * np.sin(theta) / sin_theta_0 - b = np.sin(theta) / sin_theta_0 - return a[:, None] * q0 + b[:, None] * q1 - - -# ---------------------------------------------------------------------------- -# World-frame angular velocity from a quaternion sequence by finite difference. -# Δq = q[k+1] * q[k]^-1 ; ω_w ≈ 2 * Δq.xyz / dt (small-angle, shortest path). -# Central difference for interior, forward/backward at endpoints. -# ---------------------------------------------------------------------------- - - -def quat_seq_ang_vel(q_seq: np.ndarray, dt: float) -> np.ndarray: - n = q_seq.shape[0] - out = np.zeros((n, 3), dtype=np.float64) - - def diff(q_a, q_b, h): # ω over interval h, expressed in world - aw, ax, ay, az = q_a - bw, bx, by, bz = q_b - # Δq = q_b * q_a^{-1} - dw = bw * aw + bx * ax + by * ay + bz * az - dx = -bw * ax + bx * aw - by * az + bz * ay - dy = -bw * ay + bx * az + by * aw - bz * ax - dz = -bw * az - bx * ay + by * ax + bz * aw - if dw < 0.0: # shortest path - dw, dx, dy, dz = -dw, -dx, -dy, -dz - return np.array([2 * dx / h, 2 * dy / h, 2 * dz / h]) - - for i in range(n): - if i == 0: - out[i] = diff(q_seq[0], q_seq[1], dt) - elif i == n - 1: - out[i] = diff(q_seq[n - 2], q_seq[n - 1], dt) - else: - out[i] = diff(q_seq[i - 1], q_seq[i + 1], 2 * dt) - return out diff --git a/scripts/deploy/prepend_warmup.py b/scripts/deploy/prepend_warmup.py deleted file mode 100644 index 0d0c2c656..000000000 --- a/scripts/deploy/prepend_warmup.py +++ /dev/null @@ -1,215 +0,0 @@ -#!/usr/bin/env python3 -"""Prepend a stand->dance-frame-0 warmup prefix to a WBT motion bin. - -Why this exists ---------------- -The deploy-side FSM transitions FixStand -> State_WBT at t=0, but the original -dance bin's frame 0 is mid-dance (e.g. dance1.bin frame 0 has L_shoulder_roll -at +1.41 rad vs default_angles +0.20 — a ~70 deg jump). The training pipeline -uses Reference State Initialization (RSI): the env resets the robot AT the -motion's frame 0 pose, so the policy never sees "robot at default_angles, -commanded to track dance frame 0". That mismatch produces a transient action -burst at deploy time which exceeds the ankle's holding capacity, causing the -real robot's feet to slip for ~3 s before the policy stabilises. - -Fix: insert N seconds of kinematically self-consistent interpolation frames at -the head of the bin. After this the policy sees a slow, smooth tracking task -(ramping the command from FixStand to dance frame 0), which is in-distribution -for a tracking policy. The dance frames are then concatenated verbatim. - -What this does NOT change -------------------------- -- Training pipeline (no retrain needed). -- State_WBT.cpp / deploy_config.yaml / FSM yaml time_start. -- The original dance frames (they are appended unchanged after the warmup). - -Interpolation scheme --------------------- -- joint_pos (J=29): cubic Hermite per joint with boundaries - (default_angles, 0) -> (orig.jp[0], orig.jv[0]) - joint_vel is the analytic derivative of the polynomial (kinematically - consistent with joint_pos by construction). -- body_pos (B,3): cubic Hermite per axis with boundaries - (FK(stand), 0) -> (orig.bp[0], orig.bv[0]) - body_lin_vel is the analytic derivative. -- body_quat (B,4): SLERP along quintic smoothstep s(u) = 6u^5 - 15u^4 + 10u^3 - (s'(0) = s''(0) = s'(1) = s''(1) = 0). Shortest-path; near-parallel fallback. -- body_ang_vel: central difference on the SLERP path. The seam frame is - differenced against the original bin's frame 0 so the discrete derivative is - continuous across the seam. - -FixStand body states come from MuJoCo FK on the 'stand' keyframe in scene XML, -using the same tracked_body_mujoco_ids the C++ State_WBT consumes — so the -warmup's frame 0 (body_pos/body_quat) is exactly what default_angles produces. - -Use ---- - uv run scripts/deploy/prepend_warmup.py \ - --input ../deploy_ws/assets/dance1.bin \ - --output ../deploy_ws/assets/dance1_warmup.bin \ - --config ../deploy_ws/assets/deploy_config.yaml \ - --warmup-sec 1.5 - -Validate in sim BEFORE swapping on the real robot: - uv run scripts/deploy/sim_prototype.py \ - --motion ../deploy_ws/assets/dance1_warmup.bin \ - --init-mode stand --render -""" - -from __future__ import annotations - -import argparse -import sys -from pathlib import Path - -import numpy as np -import yaml -from motion_primitives import ( - compute_fixstand_body_states, - hermite, - load_motion_bin, - quat_seq_ang_vel, - save_motion_bin, - slerp_smoothstep, -) - -REPO_ROOT = Path(__file__).resolve().parents[2] -DEPLOY_WS = REPO_ROOT.parent / "deploy_ws" -DEFAULT_SCENE = REPO_ROOT / "src/unilab/assets/robots/g1/scene_flat.xml" -DEFAULT_CFG = DEPLOY_WS / "assets/deploy_config.yaml" -DEFAULT_IN = DEPLOY_WS / "assets/dance1.bin" -DEFAULT_OUT = DEPLOY_WS / "assets/dance1_warmup.bin" -DEFAULT_WARMUP_SEC = 1.5 - - -# ---------------------------------------------------------------------------- -# Main -# ---------------------------------------------------------------------------- - - -def main() -> None: - ap = argparse.ArgumentParser( - description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter - ) - ap.add_argument("--input", type=Path, default=DEFAULT_IN) - ap.add_argument("--output", type=Path, default=DEFAULT_OUT) - ap.add_argument( - "--config", - type=Path, - default=DEFAULT_CFG, - help="deploy_config.yaml (for default_angles, tracked ids).", - ) - ap.add_argument( - "--scene", - type=Path, - default=DEFAULT_SCENE, - help="MuJoCo XML with 'stand' keyframe (for FixStand FK).", - ) - ap.add_argument( - "--warmup-sec", - type=float, - default=DEFAULT_WARMUP_SEC, - help="Length of prepended warmup interval [seconds].", - ) - args = ap.parse_args() - - with open(args.config) as f: - cfg = yaml.safe_load(f) - default_angles = np.asarray(cfg["default_angles"], dtype=np.float64) - tracked_ids = list(cfg["tracked_body_mujoco_ids"]) - - orig = load_motion_bin(args.input) - fps, J, B = orig["fps"], orig["nj"], orig["nb"] - if J != len(default_angles): - raise SystemExit(f"J mismatch: bin J={J}, default_angles len={len(default_angles)}") - if B != len(tracked_ids): - raise SystemExit(f"B mismatch: bin B={B}, tracked_body_mujoco_ids len={len(tracked_ids)}") - - dt = 1.0 / fps - N = int(round(args.warmup_sec * fps)) - if N < 2: - raise SystemExit( - f"warmup-sec={args.warmup_sec}s too short for fps={fps} (need >= 2 frames)" - ) - T = N * dt # so frame N in the new bin == original frame 0 exactly - - fk_jp, fk_bp, fk_bq = compute_fixstand_body_states(args.scene, tracked_ids) - - keyframe_diff = float(np.abs(fk_jp - default_angles).max()) - if keyframe_diff > 1e-3: - print( - f"WARN: 'stand' keyframe joint pos differs from deploy_config " - f"default_angles by {keyframe_diff:.4f} rad — using deploy_config " - "values for warmup start (so command_joint_pos matches what " - "FixStand actually holds on the real robot).", - file=sys.stderr, - ) - - # --- joint pos/vel: analytic Hermite ----------------------------------- - ts = np.arange(N) * dt # (N,) - jp_w, jv_w = hermite(default_angles, np.zeros(J), orig["jp"][0], orig["jv"][0], ts, T) - - # --- body pos / lin_vel: analytic Hermite per axis --------------------- - bp_w, bv_w = hermite(fk_bp, np.zeros((B, 3)), orig["bp"][0], orig["bv"][0], ts, T) - - # --- body quat: SLERP along quintic smoothstep ------------------------- - u = ts / T - bq_w = np.zeros((N, B, 4), dtype=np.float64) - for b in range(B): - bq_w[:, b, :] = slerp_smoothstep(fk_bq[b], orig["bq"][0, b], u) - - # --- body ang_vel: central diff, with the seam differenced against the - # original bin's first two frames so the derivative is continuous across - # the warmup/dance boundary -------------------------------------------- - bav_w = np.zeros((N, B, 3), dtype=np.float64) - for b in range(B): - ext = np.concatenate([bq_w[:, b, :], orig["bq"][:2, b, :]], axis=0) - bav_w[:, b, :] = quat_seq_ang_vel(ext, dt)[:N] - - # --- concatenate warmup + original ------------------------------------- - new_jp = np.concatenate([jp_w, orig["jp"]], axis=0) - new_jv = np.concatenate([jv_w, orig["jv"]], axis=0) - new_bp = np.concatenate([bp_w, orig["bp"]], axis=0) - new_bq = np.concatenate([bq_w, orig["bq"]], axis=0) - new_bv = np.concatenate([bv_w, orig["bv"]], axis=0) - new_bav = np.concatenate([bav_w, orig["bav"]], axis=0) - - args.output.parent.mkdir(parents=True, exist_ok=True) - save_motion_bin(args.output, fps, new_jp, new_jv, new_bp, new_bq, new_bv, new_bav) - - # --- report ------------------------------------------------------------ - seam_jp = float(np.abs(jp_w[-1] - orig["jp"][0]).max()) - seam_jv = float(np.abs(jv_w[-1] - orig["jv"][0]).max()) - seam_bp = float(np.abs(bp_w[-1] - orig["bp"][0]).max()) - seam_bv = float(np.abs(bv_w[-1] - orig["bv"][0]).max()) - # Quaternion seam: angle between bq_w[-1] and orig.bq[0] for each body - seam_qang_deg = [] - for b in range(B): - d = float(np.clip(abs(np.dot(bq_w[-1, b], orig["bq"][0, b])), 0.0, 1.0)) - seam_qang_deg.append(np.degrees(2 * np.arccos(d))) - seam_qang_max = max(seam_qang_deg) - # Warmup-start command joint deviation from default (rate / sec) - init_jvel_l2 = float(np.linalg.norm(jv_w[0])) - - print(f"Wrote {args.output}") - print(f" fps={fps} J={J} B={B}") - print(f" frames: {orig['nf']} (orig) -> {new_jp.shape[0]} ({N} warmup + {orig['nf']} dance)") - print( - f" duration: {orig['nf'] / fps:.2f}s -> {new_jp.shape[0] / fps:.2f}s " - f"(warmup_sec={args.warmup_sec})" - ) - print( - f" init (frame 0) command_joint_pos == default_angles? " - f"max_diff={float(np.abs(new_jp[0] - default_angles).max()):.6f} rad" - ) - print(f" init (frame 0) command_joint_vel L2 = {init_jvel_l2:.6f} rad/s (should be 0)") - print(" seam check (warmup last frame vs original frame 0; one-dt-step apart):") - print(f" |Δjoint_pos|_∞ = {seam_jp:.4f} rad") - print(f" |Δjoint_vel|_∞ = {seam_jv:.4f} rad/s") - print(f" |Δbody_pos|_∞ = {seam_bp:.4f} m") - print(f" |Δbody_lin_vel|_∞ = {seam_bv:.4f} m/s") - print(f" max body quat angle = {seam_qang_max:.4f} deg") - - -if __name__ == "__main__": - main() diff --git a/scripts/deploy/sim_prototype.py b/scripts/deploy/sim_prototype.py deleted file mode 100644 index 7043b9bd2..000000000 --- a/scripts/deploy/sim_prototype.py +++ /dev/null @@ -1,520 +0,0 @@ -#!/usr/bin/env python3 -"""Python prototype of State_WBT — drives the ONNX policy in MuJoCo using the -exact obs assembly the C++ State_WBT will use, so we can validate the -deploy_config.yaml + dance1.bin against training-side expectations BEFORE -writing C++. - -Inputs (defaults match the artifacts produced by export_*.py): - ~/deploy_ws/assets/policy.onnx (optional — if missing or --no-onnx, the - prototype skips inference and only sanity- - checks obs assembly with default-angle ctrl) - ~/deploy_ws/assets/deploy_config.yaml - ~/deploy_ws/assets/dance1.bin - -What this verifies: - 1. Obs assembly matches training segment-for-segment, driven by - cfg["obs_layout"] (no hard-coded dimension) — so when training flips - enable_zero_linvel / enable_zero_anchor_pos the prototype follows. - 2. obs total length equals cfg["obs_dim"]; ONNX input width (when provided) - matches cfg["obs_dim"]. - 3. With an ONNX file: q_target = action*2.0 + default_angles + clip + EMA - produces motion that visually tracks the reference in MuJoCo. - -Differences from training-side eval (deliberate): - - obs construction is reimplemented in pure numpy here, NOT routed through - the env class — this is the same code path State_WBT will reproduce in C++. - - No noise injected on obs (matches deploy convention). - - Robot anchor pos in world is locked to the first motion frame's torso pos - (since real robot has no GPS/SLAM). -""" - -from __future__ import annotations - -import argparse -import struct -import sys -import time -from pathlib import Path - -import mujoco -import numpy as np -import yaml - -REPO_ROOT = Path(__file__).resolve().parents[2] -sys.path.insert(0, str(REPO_ROOT / "src")) -from unilab.utils.rotation import ( # noqa: E402 - np_matrix_from_quat, - np_subtract_frame_transforms, -) - -DEFAULT_ONNX = Path.home() / "deploy_ws/assets/policy.onnx" -DEFAULT_CFG = Path.home() / "deploy_ws/assets/deploy_config.yaml" -DEFAULT_BIN = Path.home() / "deploy_ws/assets/dance1.bin" -DEFAULT_SCENE = REPO_ROOT / "src/unilab/assets/robots/g1/scene_flat.xml" - - -def load_motion_bin(path: Path) -> dict: - with open(path, "rb") as f: - fps, nf, nj, nb = struct.unpack(" np.ndarray: - """Concatenate segments in the order dictated by cfg['obs_layout']. - - Every layout entry must resolve to a known segment with matching dim. - Raises SystemExit on any mismatch — the prototype refuses to fabricate - a vector that disagrees with the deploy contract. - - Single-step path only; history-aware assembly goes through ``ObsAssembler``. - """ - parts: list[np.ndarray] = [] - for seg in layout: - name = seg["name"] - dim = int(seg["dim"]) - if name not in segments: - raise SystemExit(f"obs_layout segment '{name}' has no value provider in prototype") - arr = np.asarray(segments[name], dtype=np.float32).reshape(-1) - if arr.size != dim: - raise SystemExit(f"obs segment '{name}': expected dim {dim}, got {arr.size}") - parts.append(arr) - obs = np.concatenate(parts, axis=0).astype(np.float32, copy=False) - if obs.size != obs_dim: - raise SystemExit(f"assembled obs dim {obs.size} != cfg.obs_dim {obs_dim}") - return obs - - -class ObsAssembler: - """Schema-driven actor obs assembler with per-term history buffers. - - Mirrors the deploy-side ObservationManager + ObservationTermCfg behaviour - byte-for-byte so sim_prototype validates the same obs vector State_WBT - will produce on the real robot: - - * Per layout term, owns a (H, dim) deque-like ring buffer (oldest at row 0). - * ``reset(segments)`` fills all H rows with the current value (matches - deploy ObservationTermCfg::reset which calls add() H times). - * ``step(segments)`` shifts oldest out, writes current at the last row - (matches deploy add()). - * ``assemble()`` concatenates each term's full history oldest-first, - then concatenates across terms in layout order (matches deploy - use_gym_history=false mode in ObservationManager::compute_group). - - History length per term is read from ``layout[i]['history_length']``, - defaulting to 1 (single-step / no history). Total assembled dim is - sum(dim * history_length) over all terms. - """ - - def __init__(self, cfg: dict) -> None: - self.obs_dim = int(cfg["obs_dim"]) - self.layout = cfg["obs_layout"] - if cfg.get("use_gym_history", False): - raise SystemExit( - "sim_prototype assumes use_gym_history=false (group-by-term flatten); " - "set it to false in deploy_config.yaml or extend ObsAssembler" - ) - self._buffers: dict[str, np.ndarray] = {} - self._hist_len: dict[str, int] = {} - layout_total = 0 - for seg in self.layout: - name = seg["name"] - dim = int(seg["dim"]) - h_len = int(seg.get("history_length", 1)) - if h_len < 1: - raise SystemExit(f"obs term '{name}' has invalid history_length={h_len}") - self._buffers[name] = np.zeros((h_len, dim), dtype=np.float32) - self._hist_len[name] = h_len - layout_total += dim * h_len - if layout_total != self.obs_dim: - raise SystemExit( - f"cfg internal inconsistency: sum(dim*history_length)={layout_total} " - f"!= obs_dim={self.obs_dim}" - ) - self._primed = False - - def reset(self, segments: dict[str, np.ndarray]) -> np.ndarray: - """Fill every history slot with the current segment values (deploy parity).""" - for seg in self.layout: - name = seg["name"] - self._set_validated(name, segments[name]) - self._buffers[name][:] = self._buffers[name][-1:, :] - self._primed = True - return self.assemble() - - def step(self, segments: dict[str, np.ndarray]) -> np.ndarray: - """Push current values; auto-primes on the first call.""" - if not self._primed: - return self.reset(segments) - for seg in self.layout: - name = seg["name"] - buf = self._buffers[name] - buf[:-1] = buf[1:] - self._set_validated(name, segments[name]) - return self.assemble() - - def assemble(self) -> np.ndarray: - parts: list[np.ndarray] = [] - for seg in self.layout: - name = seg["name"] - parts.append(self._buffers[name].reshape(-1)) - obs = np.concatenate(parts, axis=0).astype(np.float32, copy=False) - if obs.size != self.obs_dim: - raise SystemExit(f"assembled obs dim {obs.size} != cfg.obs_dim {self.obs_dim}") - return obs - - def _set_validated(self, name: str, value) -> None: - buf = self._buffers[name] - dim = buf.shape[1] - arr = np.asarray(value, dtype=np.float32).reshape(-1) - if arr.size != dim: - raise SystemExit(f"obs segment '{name}': expected dim {dim}, got {arr.size}") - buf[-1, :] = arr - - -def compute_obs_segments( - cfg: dict, - motion_frame: dict, - *, - robot_torso_pos_w: np.ndarray, - robot_torso_quat_w: np.ndarray, - gyro: np.ndarray, - dof_pos: np.ndarray, - dof_vel: np.ndarray, - last_actions: np.ndarray, -) -> dict[str, np.ndarray]: - """Compute the current-step value of every potential obs segment. - - Returns a dict keyed by segment name. Only the segments listed in - cfg['obs_layout'] are consumed downstream; unused keys are harmless. - - All inputs are batch-less (1D arrays). Each output segment is float32. - """ - default_angles = np.asarray(cfg["default_angles"], dtype=np.float32) - anchor_idx = int(cfg["anchor_body_idx_in_tracked"]) - - ref_joint_pos = motion_frame["joint_pos"] - ref_joint_vel = motion_frame["joint_vel"] - ref_torso_pos_w = motion_frame["body_pos_w"][anchor_idx] - ref_torso_quat_w = motion_frame["body_quat_w"][anchor_idx] - - pos_b, ori_q = np_subtract_frame_transforms( - robot_torso_pos_w[None, :], - robot_torso_quat_w[None, :], - ref_torso_pos_w[None, :], - ref_torso_quat_w[None, :], - ) - motion_anchor_pos_b = pos_b[0].astype(np.float32) - ori_R = np_matrix_from_quat(ori_q)[0] - motion_anchor_ori_b = ori_R[:, :2].reshape(6).astype(np.float32) - - # linvel: deploy default = zero (no real-robot sensor on G1). - linvel_strategy = str(cfg.get("linvel_strategy", "zero")) - if linvel_strategy != "zero": - raise SystemExit( - f"linvel_strategy='{linvel_strategy}' not supported in prototype; " - "G1 deploy must use 'zero'" - ) - base_lin_vel = np.zeros(3, dtype=np.float32) - - return { - "command_joint_pos": ref_joint_pos.astype(np.float32), - "command_joint_vel": ref_joint_vel.astype(np.float32), - "motion_anchor_pos_b": motion_anchor_pos_b, - "motion_anchor_ori_b": motion_anchor_ori_b, - # Two aliases so legacy schemas (which used 'linvel') still work. - "base_lin_vel": base_lin_vel, - "linvel": base_lin_vel, - "gyro": gyro.astype(np.float32), - "joint_pos_rel": (dof_pos - default_angles).astype(np.float32), - "dof_vel": dof_vel.astype(np.float32), - "last_actions": last_actions.astype(np.float32), - } - - -def main(): - ap = argparse.ArgumentParser(description=__doc__) - ap.add_argument( - "--onnx", - type=Path, - default=DEFAULT_ONNX, - help="ONNX policy. If absent or --no-onnx, prototype runs a " - "default-angle ctrl sanity check (obs assembly only).", - ) - ap.add_argument( - "--no-onnx", action="store_true", help="Skip ONNX inference even if the file exists." - ) - ap.add_argument("--config", type=Path, default=DEFAULT_CFG) - ap.add_argument("--motion", type=Path, default=DEFAULT_BIN) - ap.add_argument("--scene", type=Path, default=DEFAULT_SCENE) - ap.add_argument( - "--render", action="store_true", help="Open MuJoCo passive viewer (requires display)." - ) - ap.add_argument( - "--max-steps", type=int, default=0, help="0 = play to end of clip then loop once." - ) - ap.add_argument( - "--cheat-anchor", - action="store_true", - help="Use sim-true robot torso pos for anchor (debug; " - "would not be available on real robot).", - ) - ap.add_argument( - "--init-mode", - choices=["rsi", "stand"], - default="rsi", - help="rsi = teleport robot to motion frame 0 (training-time " - "reset condition; default). stand = leave robot in the " - "'stand' keyframe pose (== FixStand default_angles, " - "matching the deploy-time FSM transition).", - ) - args = ap.parse_args() - - with open(args.config) as f: - cfg = yaml.safe_load(f) - motion = load_motion_bin(args.motion) - - expected_dim = int(cfg["obs_dim"]) - layout_total = sum(int(s["dim"]) * int(s.get("history_length", 1)) for s in cfg["obs_layout"]) - if layout_total != expected_dim: - raise SystemExit( - f"cfg internal inconsistency: sum(obs_layout dim*history_length)={layout_total} " - f"!= obs_dim={expected_dim}" - ) - obs_assembler = ObsAssembler(cfg) - - use_onnx = (not args.no_onnx) and args.onnx.exists() - sess = None - inp_name = out_name = None - if use_onnx: - import onnxruntime as ort # local import: optional in sanity mode - - sess = ort.InferenceSession(str(args.onnx), providers=["CPUExecutionProvider"]) - inp_name = sess.get_inputs()[0].name - out_name = sess.get_outputs()[0].name - onnx_in_shape = sess.get_inputs()[0].shape - onnx_in_dim = int(onnx_in_shape[-1]) if isinstance(onnx_in_shape[-1], int) else -1 - print( - f"ONNX: input={inp_name} {onnx_in_shape}, output={out_name} " - f"{sess.get_outputs()[0].shape}" - ) - if onnx_in_dim != expected_dim: - raise SystemExit( - f"ONNX input dim {onnx_in_dim} != cfg.obs_dim {expected_dim}. " - "Retrain (or re-export ONNX) so the policy matches the deploy contract." - ) - else: - reason = "missing" if not args.onnx.exists() else "disabled (--no-onnx)" - print( - f"ONNX: {reason} — running OBS-ASSEMBLY SANITY CHECK only " - "(ctrl = default_angles, no policy in the loop)" - ) - - print( - f"obs_dim={expected_dim}, layout segments=[" - + ", ".join( - f"{s['name']}({s['dim']}×{int(s.get('history_length', 1))})" for s in cfg["obs_layout"] - ) - + "]" - ) - - model = mujoco.MjModel.from_xml_path(str(args.scene)) - data = mujoco.MjData(model) - ctrl_dt = float(cfg["ctrl_dt"]) - sim_dt = float(model.opt.timestep) - substeps = max(1, int(round(ctrl_dt / sim_dt))) - print(f"sim_dt={sim_dt:.5f}, ctrl_dt={ctrl_dt:.3f}, substeps/ctrl={substeps}") - - # Init pose: two modes. - # rsi — Reference State Initialization: teleport robot to motion frame - # 0's pose (matches the training env's reset condition). - # stand — leave robot in the "stand" keyframe pose (== FixStand - # default_angles), which is what the C++ FSM sees at the - # FixStand → State_WBT transition on the real robot. - # Use --init-mode stand to reproduce the deploy-time "first few seconds of - # slipping while the policy snaps the robot from default_angles to dance - # frame 0" scenario in simulation, without touching hardware. - key_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_KEY, "stand") - if key_id < 0: - raise SystemExit("'stand' keyframe not found") - mujoco.mj_resetDataKeyframe(model, data, key_id) - - pelvis_id_in_tracked = 0 # 'pelvis' is first in TRACKED_BODY_NAMES - if args.init_mode == "rsi": - data.qpos[0:3] = motion["body_pos_w"][0, pelvis_id_in_tracked] - data.qpos[3:7] = motion["body_quat_w"][0, pelvis_id_in_tracked] # wxyz - data.qpos[7:] = motion["joint_pos"][0] - data.qvel[0:3] = motion["body_lin_vel_w"][0, pelvis_id_in_tracked] - data.qvel[3:6] = motion["body_ang_vel_w"][0, pelvis_id_in_tracked] - data.qvel[6:] = motion["joint_vel"][0] - mujoco.mj_forward(model, data) - print(f"init_mode={args.init_mode}: base xyz={data.qpos[:3]}, base quat={data.qpos[3:7]}") - - default_angles = np.asarray(cfg["default_angles"], dtype=np.float32) - action_scale = float(cfg["action_scale"]) - ema_alpha = float(cfg["ema_alpha"]) - joint_lower = np.asarray(cfg["joint_lower"], dtype=np.float32) - joint_upper = np.asarray(cfg["joint_upper"], dtype=np.float32) - anchor_idx = int(cfg["anchor_body_idx_in_tracked"]) - anchor_body_name = cfg["anchor_body_name"] - anchor_body_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, anchor_body_name) - - gyro_sid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, "gyro") - if gyro_sid < 0: - raise SystemExit("'gyro' sensor not found in model") - gyro_adr = int(model.sensor_adr[gyro_sid]) - gyro_dim = int(model.sensor_dim[gyro_sid]) - if gyro_dim != 3: - raise SystemExit(f"gyro sensor has dim {gyro_dim}, expected 3") - - robot_anchor_pos_w_locked = motion["body_pos_w"][0, anchor_idx].astype(np.float32) - print(f"robot_anchor_pos_w (locked) = {robot_anchor_pos_w_locked}") - - last_actions = np.zeros(29, dtype=np.float32) - q_target_smoothed = default_angles.copy() - n_frames = motion["num_frames"] - if args.max_steps > 0: - total_steps = args.max_steps - elif use_onnx: - total_steps = n_frames - else: - total_steps = min(50, n_frames) # sanity mode: short run - - obs_norms = [] - action_amplitudes = [] - z_errors = [] - - viewer = None - if args.render: - from mujoco import viewer as mj_viewer - - viewer = mj_viewer.launch_passive(model, data) - - t_wall = time.time() - for step in range(total_steps): - frame_idx = step % n_frames - - robot_torso_quat_w = data.xquat[anchor_body_id].astype(np.float32) - gyro = data.sensordata[gyro_adr : gyro_adr + gyro_dim].astype(np.float32) - dof_pos = data.qpos[7:].astype(np.float32) - dof_vel = data.qvel[6:].astype(np.float32) - - motion_frame = { - "joint_pos": motion["joint_pos"][frame_idx], - "joint_vel": motion["joint_vel"][frame_idx], - "body_pos_w": motion["body_pos_w"][frame_idx], - "body_quat_w": motion["body_quat_w"][frame_idx], - } - if args.cheat_anchor: - robot_torso_pos_w_used = data.xpos[anchor_body_id].astype(np.float32) - else: - robot_torso_pos_w_used = robot_anchor_pos_w_locked - current_segments = compute_obs_segments( - cfg, - motion_frame, - robot_torso_pos_w=robot_torso_pos_w_used, - robot_torso_quat_w=robot_torso_quat_w, - gyro=gyro, - dof_pos=dof_pos, - dof_vel=dof_vel, - last_actions=last_actions, - ) - # Stateful: first call auto-resets (fills history with current segments, - # matching State_WBT.cpp's env_->reset() at FSM enter); subsequent calls - # push current and drop oldest (matching ObservationTermCfg::add). - obs = obs_assembler.step(current_segments) - if not np.all(np.isfinite(obs)): - raise SystemExit(f"non-finite obs at step {step}") - - if use_onnx: - action = sess.run([out_name], {inp_name: obs[None, :].astype(np.float32)})[0][0] - action = action.astype(np.float32) - last_actions = action.copy() - q_target = action * action_scale + default_angles - q_target = np.clip(q_target, joint_lower, joint_upper) - q_target_smoothed = ema_alpha * q_target + (1.0 - ema_alpha) * q_target_smoothed - else: - # Sanity mode: keep robot at default angles, freeze last_actions at 0. - q_target_smoothed = default_angles.copy() - - data.ctrl[:] = q_target_smoothed - - for _ in range(substeps): - mujoco.mj_step(model, data) - - obs_norms.append(float(np.linalg.norm(obs))) - action_amplitudes.append(float(np.max(np.abs(last_actions)))) - robot_z = float(data.xpos[anchor_body_id, 2]) - ref_z = float(motion["body_pos_w"][frame_idx, anchor_idx, 2]) - z_errors.append(abs(robot_z - ref_z)) - - if viewer is not None: - viewer.sync() - elapsed = time.time() - t_wall - target = (step + 1) * ctrl_dt - if elapsed < target: - time.sleep(target - elapsed) - - if step % 50 == 0: - print( - f"step {step:4d} frame={frame_idx:3d} " - f"obs_norm={obs_norms[-1]:7.2f} " - f"|action|={action_amplitudes[-1]:5.2f} " - f"z_err={z_errors[-1]:.3f}m " - f"q_target[:3]={q_target_smoothed[:3]}" - ) - - if not np.all(np.isfinite(data.qpos)): - print(f"!! NaN at step {step}, aborting") - break - if use_onnx and z_errors[-1] > 0.6: - print(f"!! z_err {z_errors[-1]:.2f}m exceeds 0.6m, robot likely fell at step {step}") - break - - if viewer is not None: - viewer.close() - - n = len(obs_norms) - print() - print(f"Ran {n} ctrl steps ({n * ctrl_dt:.2f}s of motion).") - print(f"obs_norm: mean={np.mean(obs_norms):.3f} max={np.max(obs_norms):.3f}") - print( - f"|action| max: mean={np.mean(action_amplitudes):.3f} max={np.max(action_amplitudes):.3f}" - ) - print(f"|torso z err|: mean={np.mean(z_errors):.4f}m max={np.max(z_errors):.4f}m") - - if not use_onnx: - # Sanity mode never claims tracking — only that obs assembly is finite & shaped. - print( - "SANITY OK — obs assembly produced finite vectors of expected width " - f"({expected_dim}); end-to-end tracking requires running with a " - f"matching ONNX (input dim {expected_dim})." - ) - elif np.max(z_errors) < 0.2 and n == total_steps: - print("PROTOTYPE OK — obs assembly + ONNX inference produces tracking behavior.") - elif n < total_steps: - print("WARNING: prototype aborted before clip end (see message above).") - else: - print("WARNING: large torso z error — obs assembly may be off, or policy weak.") - - -if __name__ == "__main__": - main() diff --git a/src/unilab/envs/motion_tracking/g1/tracking_obs.py b/src/unilab/envs/motion_tracking/g1/tracking_obs.py index 270c4bdc6..3ed23a902 100644 --- a/src/unilab/envs/motion_tracking/g1/tracking_obs.py +++ b/src/unilab/envs/motion_tracking/g1/tracking_obs.py @@ -1,26 +1,3 @@ -"""G1 Whole-Body Tracking — sim2real-oriented SAC variant (task ``G1WBTObs``). - -This module registers a strict subclass of :class:`G1MotionTrackingSACEnv` that -adds the training-pipeline pieces needed for ONNX-on-real-G1 deployment: - -* drop deploy-unavailable channels from the actor obs - (``base_lin_vel``, ``motion_anchor_pos_b``); -* per-step uniform noise on ``motion_anchor_ori_b`` (actor only); -* proprio observation history (``gyro`` / ``joint_pos_rel`` / ``dof_vel`` / - ``last_actions``) flattened oldest-first per term, matching the deploy-side - ``ObservationManager`` when ``use_gym_history=false``; -* per-episode encoder bias on ``joint_pos_rel`` (actor only); -* per-episode foot-geom friction sampled across regex-matched geoms; -* per-episode y / z COM offsets layered on top of the existing x offset; -* ``joint_acc_l2`` and ``joint_torque_l2`` reward terms. - -All extensions are gated by flags on :class:`G1WBTObsCfg`; the bases -(``G1MotionTrackingSACCfg`` / ``G1MotionTrackingSACEnv`` / -``G1MotionTrackingEnv``) are untouched. Switch the pelvis IMU via the -yaml ``env.sensor.gyro``/``env.sensor.upvector``/``env.sensor.local_linvel`` -fields (no XML duplication required — ``g1.xml`` already exposes both IMUs). -""" - from __future__ import annotations import re @@ -53,10 +30,6 @@ ) from .tracking_sac import G1MotionTrackingSACCfg, G1MotionTrackingSACEnv -# --------------------------------------------------------------------------- # -# Config extensions -# --------------------------------------------------------------------------- # - @dataclass class ObsNoiseConfig(NoiseConfig): @@ -66,17 +39,10 @@ class ObsNoiseConfig(NoiseConfig): a drop-in replacement; ``G1WBTObs`` flips the flags via its task yaml. """ - # Drop ``base_lin_vel`` from actor obs (G1 has no on-robot linvel sensor). enable_zero_linvel: bool = False - # Drop ``motion_anchor_pos_b`` from actor obs (no torso-pose estimator). enable_zero_anchor_pos: bool = False - # Per-step uniform noise on ``motion_anchor_ori_b`` (actor only). enable_anchor_ori_noise: bool = False scale_anchor_ori: float = 0.05 - # When > 1, proprio terms (gyro / joint_pos_rel / dof_vel / last_actions) - # are flattened oldest-first as an H-step history block. Reference terms - # stay single-step. Critic stays single-step. Mirrors deploy-side - # ``ObservationManager`` with ``use_gym_history=false``. obs_history_length: int = 1 @@ -89,13 +55,9 @@ class ObsDomainRand(Domain_Rand): randomize_com_z: bool = False com_offset_z: list[float] = field(default_factory=lambda: [-0.05, 0.05]) - # Per-episode additive bias on actor's joint_pos channel. enable_encoder_bias: bool = False encoder_bias_range: list[float] = field(default_factory=lambda: [-0.01, 0.01]) - # Per-reset foot-geom friction. ``shared_random=True`` — a single scalar - # is broadcast across all foot geoms of one env, applied to the - # sliding-friction column. Matches mjlab. randomize_geom_friction: bool = False friction_range: list[float] = field(default_factory=lambda: [0.3, 1.2]) friction_geom_pattern: str = r"^(left|right)_foot[1-7]_collision$" @@ -110,11 +72,6 @@ class G1WBTObsCfg(G1MotionTrackingSACCfg): domain_rand: ObsDomainRand = field(default_factory=ObsDomainRand) # type: ignore[assignment] -# --------------------------------------------------------------------------- # -# DR provider extension -# --------------------------------------------------------------------------- # - - class G1WBTObsDomainRandomizationProvider(G1MotionTrackingDomainRandomizationProvider): """Extends the SAC tracking DR provider with encoder bias, foot-geom friction, y/z COM offsets, and post-reset ``prev_dof_vel`` seeding.""" @@ -156,9 +113,6 @@ def build_reset_plan(self, env: Any, env_ids: np.ndarray) -> ResetPlan: info_updates: dict[str, Any] = { "current_actions": zero_actions(num_reset, env._num_action), "last_actions": zero_actions(num_reset, env._num_action), - # Seed prev_dof_vel with the post-reset joint velocity so the first - # joint_acc_l2 sample is physically meaningful (Δv from the new - # starting velocity, not a spurious step from pre-termination). "prev_dof_vel": qvel[:, 6:].astype(get_global_dtype()), } @@ -173,7 +127,6 @@ def build_reset_plan(self, env: Any, env_ids: np.ndarray) -> ResetPlan: env, num_reset, base_kp=self._base_kp, base_kd=self._base_kd ) - # Foot-geom friction. if getattr(dr_cfg, "randomize_geom_friction", False): assert self._base_geom_friction is not None assert self._foot_geom_ids is not None @@ -190,7 +143,6 @@ def build_reset_plan(self, env: Any, env_ids: np.ndarray) -> ResetPlan: payload.geom_friction = geom_friction randomization = payload - # y / z COM offsets, layered on top of parent's x-only common build. has_com_y = getattr(dr_cfg, "randomize_com_y", False) has_com_z = getattr(dr_cfg, "randomize_com_z", False) if has_com_y or has_com_z: @@ -216,11 +168,6 @@ def build_reset_plan(self, env: Any, env_ids: np.ndarray) -> ResetPlan: ) -# --------------------------------------------------------------------------- # -# Env -# --------------------------------------------------------------------------- # - - @registry.env("G1WBTObs", sim_backend="mujoco") @registry.env("G1WBTObs", sim_backend="motrix") class G1WBTObsEnv(G1MotionTrackingSACEnv): @@ -235,16 +182,10 @@ class G1WBTObsEnv(G1MotionTrackingSACEnv): def __init__(self, cfg: G1WBTObsCfg, num_envs: int = 1, backend_type: str = "mujoco"): super().__init__(cfg, num_envs=num_envs, backend_type=backend_type) - # Cache base actuator gains for joint_torque_l2. - # Position-control torque approx: τ ≈ kp·(target_q − q) − kd·qd. - # DR (kp/kd ±10–15%) leaves small error vs true per-env torque, but - # the gradient direction (penalise large action / large Δq) is preserved. base_kp, base_kd = self._backend.get_actuator_gains() self._base_kp = np.asarray(base_kp, dtype=get_global_dtype()) self._base_kd = np.asarray(base_kd, dtype=get_global_dtype()) - # Proprio history buffers — per-term, oldest-first. Allocated only when - # H > 1 so H = 1 is zero-overhead. H = max(1, int(cfg.noise_config.obs_history_length)) self._hist_len = H self._hist_buf: dict[str, np.ndarray] | None = None @@ -257,11 +198,8 @@ def __init__(self, cfg: G1WBTObsCfg, num_envs: int = 1, backend_type: str = "muj "dof_vel": np.zeros((num_envs, H, n), dtype=dtype), "last_actions": np.zeros((num_envs, H, n), dtype=dtype), } - # Plumbs ``info`` from ``_compute_obs`` down to ``_build_actor_obs`` - # without changing the base-class hook signature. self._obs_compute_info: dict | None = None - # Swap to the extended DR provider whenever an extended flag is on. dr_cfg = cfg.domain_rand needs_extended = ( getattr(dr_cfg, "enable_encoder_bias", False) @@ -293,19 +231,10 @@ def __init__(self, cfg: G1WBTObsCfg, num_envs: int = 1, backend_type: str = "muj base_geom_friction=base_geom_friction, foot_geom_ids=foot_geom_ids, ) - # Swap the per-reset DR provider directly. ``_init_domain_randomization`` - # cannot be called twice — it materializes the backend at the end and - # MuJoCo's pool raises on a second materialize. The parent's call - # already (a) ran init randomization and (b) materialized; we only - # need the new provider's ``build_reset_plan`` for per-episode DR. from unilab.dr import DomainRandomizationManager self._dr_manager = DomainRandomizationManager(self, extended_provider) - # ------------------------------------------------------------------ # - # Rewards - # ------------------------------------------------------------------ # - def _init_reward_functions(self) -> None: super()._init_reward_functions() self._reward_fns["joint_acc_l2"] = self._reward_joint_acc_l2 @@ -335,19 +264,15 @@ def _reward_joint_torque_l2(self, ctx: RewardContext) -> np.ndarray: torque = self._base_kp * (target_q - dof_pos) - self._base_kd * dof_vel return np.asarray(np.sum(np.square(torque), axis=1), dtype=get_global_dtype()) - # ------------------------------------------------------------------ # - # Obs - # ------------------------------------------------------------------ # - def _actor_obs_dim(self, n: int) -> int: nc = self._cfg.noise_config H = max(1, int(nc.obs_history_length)) - single_step = 2 * n + 6 # command(2n) + anchor_ori(6) + single_step = 2 * n + 6 if not nc.enable_zero_anchor_pos: single_step += 3 if not nc.enable_zero_linvel: single_step += 3 - proprio_step = 3 + 3 * n # gyro + joint_pos_rel + dof_vel + last_actions + proprio_step = 3 + 3 * n return single_step + H * proprio_step def _compute_obs( @@ -361,8 +286,6 @@ def _compute_obs( robot_body_pos_w: np.ndarray, robot_body_quat_w: np.ndarray, ) -> dict[str, np.ndarray]: - # Stash so the overridden ``_build_actor_obs`` (called inside super) - # can read env_ids / joint_pos_obs_bias without a signature change. self._obs_compute_info = info try: obs = super()._compute_obs( @@ -378,8 +301,6 @@ def _compute_obs( finally: self._obs_compute_info = None - # Cache for next-step joint_acc_l2. The reset path overwrites this - # via the DR provider's ``prev_dof_vel`` info_update. info["prev_dof_vel"] = dof_vel.copy() return obs @@ -396,28 +317,21 @@ def _build_actor_obs( last_actions: np.ndarray, ) -> np.ndarray: info = self._obs_compute_info or {} - # Reset path is signalled by ``env_ids`` in obs_info (set by parent's - # ``_refresh_observation_rows`` and the DR provider's - # ``build_reset_observation``). In that case fill history slots; in - # the per-step path we push (oldest out, current in). env_ids = info.get("env_ids") is_reset = env_ids is not None nc = self._cfg.noise_config - # Per-episode encoder bias on actor's joint_pos channel. bias = info.get("joint_pos_obs_bias") if bias is not None and bias.shape == noisy_joint_pos_rel.shape: noisy_joint_pos_rel = np.asarray( noisy_joint_pos_rel + bias, dtype=noisy_joint_pos_rel.dtype ) - # Per-step anchor_ori noise (actor only). actor_anchor_ori_b = motion_anchor_ori_b if nc.enable_anchor_ori_noise: actor_anchor_ori_b = self._obs_noise(motion_anchor_ori_b, nc.scale_anchor_ori) - # Single-step reference terms, dropping deploy-unavailable channels. actor_terms: list[np.ndarray] = [command] if not nc.enable_zero_anchor_pos: actor_terms.append(motion_anchor_pos_b) @@ -425,7 +339,6 @@ def _build_actor_obs( if not nc.enable_zero_linvel: actor_terms.append(noisy_linvel) - # Proprio history (or single-step pass-through when H = 1). if self._hist_buf is not None: components = { "gyro": noisy_gyro, @@ -439,24 +352,13 @@ def _build_actor_obs( self._push_obs_history(env_ids, components) sel = slice(None) if env_ids is None else env_ids for key in ("gyro", "joint_pos_rel", "dof_vel", "last_actions"): - buf = self._hist_buf[key][sel] # (n_e, H, D) + buf = self._hist_buf[key][sel] actor_terms.append(buf.reshape(buf.shape[0], -1)) else: actor_terms.extend([noisy_gyro, noisy_joint_pos_rel, noisy_dof_vel, last_actions]) return np.concatenate(actor_terms, axis=1, dtype=get_global_dtype()) - # ------------------------------------------------------------------ # - # Proprio history buffer maintenance. - # Mirrors deploy ``ObservationManager`` / ``ObservationTermCfg``: - # * On reset: fill all H slots with the current value (matches - # ``ObservationTermCfg::reset`` which calls ``add()`` H times). - # * On step: pop oldest, push current at end. - # * Read order is oldest-first, so ``flatten(buf[env, :, :])`` yields - # ``[t-H+1, t-H+2, ..., t]`` — matches deploy - # ``ObservationTermCfg::get`` (deque front-to-back). - # ------------------------------------------------------------------ # - def _push_obs_history( self, env_ids: np.ndarray | None, components: dict[str, np.ndarray] ) -> None: diff --git a/tests/scripts/test_obs_alignment_g1_wbt.py b/tests/scripts/test_obs_alignment_g1_wbt.py index 93b5323b8..454ea3fb6 100644 --- a/tests/scripts/test_obs_alignment_g1_wbt.py +++ b/tests/scripts/test_obs_alignment_g1_wbt.py @@ -1,56 +1,13 @@ -"""Cross-side obs alignment test for the G1 WBT Obs deploy chain. - -This is the load-bearing test that ensures the train -> export -> deploy -pipeline produces byte-identical actor obs at every step. Three independent -implementations are exercised against the SAME inputs: - - 1. Training side — tracking_obs.py's _push_obs_history / - _fill_obs_history + actor obs assembly in - _build_actor_obs (replicated below in numpy). - 2. Schema side — sim_prototype.ObsAssembler driven by deploy_config.yaml. - 3. Deploy side — observation_manager.h::ObservationTermCfg semantics - replicated in Python (oldest-first deque per term, - group-by-term flatten, use_gym_history=false). - -If any pair diverges, redeployment will silently fail at runtime — better to -catch it here than at the FSM transition with a robot on a rig. - -The test is hermetic: it does NOT spin up MuJoCo, does NOT load motion clips, -and does NOT depend on training infra. It synthesizes fixed random inputs, -runs the assembly on both sides, and asserts bit-for-bit equality. -""" - from __future__ import annotations -import sys from collections import deque -from pathlib import Path import numpy as np import pytest -REPO_ROOT = Path(__file__).resolve().parents[2] -SIM_PROTOTYPE = REPO_ROOT / "scripts" / "deploy" / "sim_prototype.py" - - -def _load_sim_prototype(): - """Import sim_prototype as a module (it's under scripts/, not src/).""" - import importlib.util - - spec = importlib.util.spec_from_file_location("sim_prototype", SIM_PROTOTYPE) - mod = importlib.util.module_from_spec(spec) - sys.modules["sim_prototype"] = mod - spec.loader.exec_module(mod) - return mod - - -# --------------------------------------------------------------------------- -# Reference deque (mirrors deploy ObservationTermCfg::reset / add / get). -# --------------------------------------------------------------------------- - class _DeployTerm: - """Bit-for-bit copy of ObservationTermCfg buffer semantics. + """Per-term history buffer with deque semantics. reset(obs): add(obs) H times (== fill). add(obs): push at back; if buffer > H, pop_front. @@ -78,20 +35,16 @@ def get(self) -> np.ndarray: def _deploy_compute_group(layout, current_segments_by_name): - """Mirror ObservationManager::compute_group with use_gym_history=False.""" + """Assemble per step: each term's full history oldest-first, then concat.""" terms = {} for seg in layout: terms[seg["name"]] = _DeployTerm(int(seg["dim"]), int(seg.get("history_length", 1))) - # Match deploy reset(): fill once with the FIRST step's value. for seg in layout: terms[seg["name"]].reset(current_segments_by_name[0][seg["name"]]) - # First "step" inside compute_group calls term.add() once more BEFORE get(); - # that final add corresponds to the current frame after the reset fill. out_per_step = [] - for step_idx, segments in enumerate(current_segments_by_name): + for segments in current_segments_by_name: for seg in layout: terms[seg["name"]].add(segments[seg["name"]]) - # Concat each term's full history oldest-first, then across terms. out = np.concatenate([terms[seg["name"]].get() for seg in layout], axis=0).astype( np.float32 ) @@ -99,12 +52,6 @@ def _deploy_compute_group(layout, current_segments_by_name): return out_per_step -# --------------------------------------------------------------------------- -# Schema fixture — mirrors what export_deploy_config.py writes for the -# deploy profile (H=5, both zero flags ON). -# --------------------------------------------------------------------------- - - @pytest.fixture def deploy_cfg(): n = 29 @@ -121,7 +68,6 @@ def deploy_cfg(): total = sum(s["dim"] * s["history_length"] for s in obs_layout) return { "obs_dim": total, - "use_gym_history": False, "action_dim": n, "obs_layout": obs_layout, } @@ -144,11 +90,6 @@ def _random_segments(rng, num_action=29): } -# --------------------------------------------------------------------------- -# Tests -# --------------------------------------------------------------------------- - - class TestObsDim: def test_deploy_profile_dim_is_514(self, deploy_cfg): assert deploy_cfg["obs_dim"] == 514 @@ -158,51 +99,24 @@ def test_layout_sum_matches_obs_dim(self, deploy_cfg): assert total == deploy_cfg["obs_dim"] -class TestSchemaAssemblerVsDeploy: - """sim_prototype.ObsAssembler vs the deploy-side ObservationTermCfg.""" - - def test_first_step_matches_deploy_reset_then_add(self, deploy_cfg, rng): - sp = _load_sim_prototype() - assembler = sp.ObsAssembler(deploy_cfg) - - seg0 = _random_segments(rng) - prototype = assembler.step(seg0) - - deploy = _deploy_compute_group(deploy_cfg["obs_layout"], [seg0])[0] - - np.testing.assert_array_equal(prototype, deploy) - - def test_multi_step_buffer_eviction_matches_deploy(self, deploy_cfg, rng): - sp = _load_sim_prototype() - assembler = sp.ObsAssembler(deploy_cfg) - - all_segments = [_random_segments(rng) for _ in range(20)] - prototype_seq = [assembler.step(s) for s in all_segments] - deploy_seq = _deploy_compute_group(deploy_cfg["obs_layout"], all_segments) - - for k, (p, d) in enumerate(zip(prototype_seq, deploy_seq)): - np.testing.assert_array_equal( - p, d, err_msg=f"sim_prototype <-> deploy mismatch at step {k}" - ) +class TestHistoryOrdering: + """History blocks must read oldest-first, not newest-first.""" def test_history_terms_carry_oldest_first(self, deploy_cfg, rng): - """Spot-check the gyro history block manually: at step k>=H, the 5*3 - gyro slot of obs should be [gyro_{k-H+1}, gyro_{k-H+2}, ..., gyro_k] - flattened, NOT the reverse.""" - sp = _load_sim_prototype() - assembler = sp.ObsAssembler(deploy_cfg) - + """At step k>=H the 5*3 gyro slot should be + [gyro_{k-H+1}, ..., gyro_k] flattened, NOT the reverse.""" + layout = deploy_cfg["obs_layout"] gyros = [np.array([float(k), 0.0, 0.0], dtype=np.float32) for k in range(10)] + all_segments = [] for k in range(10): seg = _random_segments(rng) seg["gyro"] = gyros[k] - obs = assembler.step(seg) + all_segments.append(seg) + + obs = _deploy_compute_group(layout, all_segments)[-1] - # Find the gyro block. layout order: cmd_jp, cmd_jv, anchor_ori, gyro,... - # offset = 29 + 29 + 6 = 64 offset = 29 + 29 + 6 gyro_block = obs[offset : offset + 3 * 5].reshape(5, 3) - # Newest at the END (idx 4) = gyros[9]; oldest = gyros[5]. expected = np.stack(gyros[5:10]) np.testing.assert_array_equal(gyro_block, expected) @@ -218,16 +132,14 @@ class TestTrainingAssemblerVsDeploy: """ @staticmethod - def _training_actor_obs(history_buf, current_refs, hist_components, num_envs): - # refs (single-step) + def _training_actor_obs(history_buf, current_refs, num_envs): parts = [ current_refs["command_joint_pos"], current_refs["command_joint_vel"], current_refs["motion_anchor_ori_b"], ] - # proprio history (oldest-first per term, then concat across terms) for key in ("gyro", "joint_pos_rel", "dof_vel", "last_actions"): - buf = history_buf[key] # (num_envs, H, D) + buf = history_buf[key] parts.append(buf.reshape(num_envs, -1)) return np.concatenate(parts, axis=1).astype(np.float32) @@ -236,7 +148,6 @@ def test_training_path_matches_deploy(self, deploy_cfg, rng): n = deploy_cfg["action_dim"] H = 5 - # Initial buffer of zeros (matches tracking_obs.py allocation). buf = { "gyro": np.zeros((n_env, H, 3), dtype=np.float32), "joint_pos_rel": np.zeros((n_env, H, n), dtype=np.float32), @@ -247,7 +158,6 @@ def test_training_path_matches_deploy(self, deploy_cfg, rng): all_segments = [_random_segments(rng) for _ in range(15)] deploy_seq = _deploy_compute_group(deploy_cfg["obs_layout"], all_segments) - # Step 0: reset (matches tracking_obs.py is_reset=True path). s0 = all_segments[0] for key in ("gyro", "joint_pos_rel", "dof_vel", "last_actions"): buf[key][:, :, :] = s0[key][None, None, :] @@ -256,10 +166,9 @@ def test_training_path_matches_deploy(self, deploy_cfg, rng): "command_joint_vel": s0["command_joint_vel"][None, :], "motion_anchor_ori_b": s0["motion_anchor_ori_b"][None, :], } - train_obs0 = self._training_actor_obs(buf, refs0, s0, n_env) + train_obs0 = self._training_actor_obs(buf, refs0, n_env) np.testing.assert_array_equal(train_obs0[0], deploy_seq[0]) - # Steps 1..N-1: push (matches tracking_obs.py is_reset=False path). for k in range(1, len(all_segments)): sk = all_segments[k] for key in ("gyro", "joint_pos_rel", "dof_vel", "last_actions"): @@ -270,19 +179,19 @@ def test_training_path_matches_deploy(self, deploy_cfg, rng): "command_joint_vel": sk["command_joint_vel"][None, :], "motion_anchor_ori_b": sk["motion_anchor_ori_b"][None, :], } - train_obs = self._training_actor_obs(buf, refs, sk, n_env) + train_obs = self._training_actor_obs(buf, refs, n_env) np.testing.assert_array_equal( train_obs[0], deploy_seq[k], err_msg=f"training <-> deploy mismatch at step {k}" ) class TestBackCompat: - """H=1 ('no history') must reproduce the pre-history 154-d path bit-exact.""" + """H=1 ('no history') must reproduce the plain-concat 154-d path bit-exact.""" @pytest.fixture - def legacy_cfg(self): + def legacy_layout(self): n = 29 - obs_layout = [ + return [ {"name": "command_joint_pos", "dim": n, "history_length": 1}, {"name": "command_joint_vel", "dim": n, "history_length": 1}, {"name": "motion_anchor_ori_b", "dim": 6, "history_length": 1}, @@ -291,22 +200,17 @@ def legacy_cfg(self): {"name": "dof_vel", "dim": n, "history_length": 1}, {"name": "last_actions", "dim": n, "history_length": 1}, ] - return { - "obs_dim": 154, - "use_gym_history": False, - "action_dim": n, - "obs_layout": obs_layout, - } - def test_legacy_obs_dim_154(self, legacy_cfg): - assert legacy_cfg["obs_dim"] == 154 + def test_legacy_obs_dim_154(self, legacy_layout): + total = sum(s["dim"] * s["history_length"] for s in legacy_layout) + assert total == 154 - def test_legacy_matches_simple_concat(self, legacy_cfg, rng): - sp = _load_sim_prototype() - assembler = sp.ObsAssembler(legacy_cfg) + def test_legacy_matches_simple_concat(self, legacy_layout, rng): + """With H=1 every term holds exactly the current value, so the + assembled vector is a plain concat in layout order.""" + all_segments = [_random_segments(rng) for _ in range(5)] + assembled = _deploy_compute_group(legacy_layout, all_segments) - for _ in range(5): - seg = _random_segments(rng) - obs = assembler.step(seg) - expected = np.concatenate([seg[s["name"]] for s in legacy_cfg["obs_layout"]]) + for seg, obs in zip(all_segments, assembled): + expected = np.concatenate([seg[s["name"]] for s in legacy_layout]) np.testing.assert_array_equal(obs, expected)