From 63f046ea1a5ae760e7ec706caaddd7ded2cfb51a Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Thu, 27 Aug 2026 18:27:00 +0800 Subject: [PATCH] Revert "cleanup: drop never-runnable G1 deploy scripts, bind docs to task owners(issue #1029)" --- .../2-user_guide/4-tasks/2-motion_tracking.md | 12 +- .../3-deployment/1-sim_to_real/1-overview.md | 5 +- .../1-sim_to_real/2-g1_whole_body.md | 124 ++--- .../1-sim_to_real/5-onnx_runtime.md | 41 +- .../1-sim_to_real/7-safety_layers.md | 35 +- .../1-sim_to_real/8-latency_budget.md | 34 +- .../1-sim_to_real/9-troubleshooting.md | 15 +- .../2-user_guide/4-tasks/2-motion_tracking.md | 8 +- .../3-deployment/1-sim_to_real/1-overview.md | 3 +- .../1-sim_to_real/2-g1_whole_body.md | 111 ++-- .../1-sim_to_real/5-onnx_runtime.md | 36 +- .../1-sim_to_real/7-safety_layers.md | 31 +- .../1-sim_to_real/8-latency_budget.md | 26 +- .../1-sim_to_real/9-troubleshooting.md | 11 +- scripts/deploy/append_cooldown.py | 299 ++++++++++ scripts/deploy/export_deploy_config.py | 328 +++++++++++ scripts/deploy/export_motion_bin.py | 179 ++++++ scripts/deploy/motion_primitives.py | 174 ++++++ scripts/deploy/prepend_warmup.py | 215 ++++++++ scripts/deploy/sim_prototype.py | 520 ++++++++++++++++++ .../envs/motion_tracking/g1/tracking_obs.py | 104 +++- tests/scripts/test_obs_alignment_g1_wbt.py | 154 +++++- 22 files changed, 2207 insertions(+), 258 deletions(-) create mode 100644 scripts/deploy/append_cooldown.py create mode 100644 scripts/deploy/export_deploy_config.py create mode 100644 scripts/deploy/export_motion_bin.py create mode 100644 scripts/deploy/motion_primitives.py create mode 100644 scripts/deploy/prepend_warmup.py create mode 100644 scripts/deploy/sim_prototype.py 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 8023b95b9..7b78c58a5 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,13 +59,11 @@ 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`), 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`: +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`: ```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 4a838a6fb..38a1bc8b9 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,8 +6,9 @@ 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. UniLab ships the training -side and the ONNX export; every robot needs a hardware-side runtime that: +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: 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 3ea3dfd3c..e4ba8b088 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,21 +2,20 @@ ::::{admonition} Hardware target :class: note -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. +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. :::: -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. +This guide walks the **last mile** between a converged G1 motion-tracking +policy and a closed-loop run on the robot. ## 0. Verify your sim-side checkpoint ```bash # Replay the policy headlessly and produce a video. -uv run eval --algo sac --task g1_wbt_obs --sim mujoco --load-run -1 \ +uv run eval --algo ppo --task g1_motion_tracking --sim motrix --load-run -1 \ --render-mode record ``` @@ -26,51 +25,39 @@ 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 before hardware bring-up. +If any of those is off, fix the sim-side checkpoint or deploy contract before +hardware bring-up. -## 1. Pick the owner, then read its contract off the YAML +## 1. Export -Every field your hardware loop needs is declared in the task owner YAML. The -deploy-oriented G1 owners are: +Use the training playback path to export `policy.onnx`, then export the G1 WBT +deploy config and motion binary with the committed deployment helpers: -```{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. -``` +```bash +uv run eval --algo ppo --task g1_motion_tracking --sim motrix --load-run -1 -::::{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. -:::: +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 +``` -Export `policy.onnx` through the training playback path: +The deployment-side prototype consumes: -```bash -uv run eval --algo sac --task g1_wbt_obs --sim mujoco --load-run -1 +``` +runs// +└── policy.onnx +logs/deploy/ +├── deploy_config.yaml +└── dance1.bin ``` ## 2. Observation contract -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**: +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: ```{list-table} :header-rows: 1 @@ -90,39 +77,32 @@ flattened **oldest-first**: - anchor orientation term from the reference and robot torso frames * - `gyro` - 3 per history step - - IMU gyro term (`env.sensor.gyro`, `pelvis_gyro` for this owner) + - IMU gyro term * - `joint_pos_rel` - 29 per history step - - measured joint position minus the `stand` keyframe joint angles + - measured joint position minus `default_angles` * - `dof_vel` - 29 per history step - joint velocity term * - `last_actions` - - 29 per history step + - 29 - previous raw actor output ``` -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. +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. ## 3. Actuator interface -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`). +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`. -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. +- 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. ## 4. Reference motion sync @@ -145,7 +125,8 @@ specifics: - Reject non-finite actions and shape mismatches before applying `action_scale`. -- Clamp generated targets with the joint range from the owner's scene XML. +- Clamp generated targets with `joint_lower` / `joint_upper` from + `deploy_config.yaml`. - Keep watchdog, pose monitor, and operator-stop thresholds in the deploy controller and test them independently of the policy. @@ -164,9 +145,18 @@ 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. 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. +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. ## 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 71387f942..92c101742 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,21 +26,38 @@ 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, 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. +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. -## Verifying the Exported Graph +## G1 Deployment Prototype -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: +The committed G1 WBT deployment helpers use these artifacts: -- 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`. +| 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. ## 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 6477f9cf3..9835121aa 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,25 +43,28 @@ flowchart LR ``` Keep the hard real-time safety checks in the deploy controller, not in the -training script. The repository does not implement a production motor-driver -safety loop — that boundary is yours to build and test. +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. ## What the policy assumes you've configured -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`): +The G1 deployment helper exports these fields into `deploy_config.yaml`: -| 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 | +```yaml +action_scale: 2.0 +ema_alpha: 1.0 +default_angles: [...] +joint_lower: [...] +joint_upper: [...] +kp: [...] +kd: [...] +``` -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. +`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. ## Hand-off testing @@ -69,8 +72,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 the joint range - from the owner's scene XML. +2. Inject an out-of-range joint target and verify clamping uses + `joint_lower` / `joint_upper` from `deploy_config.yaml`. 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 43c7fd5d8..7e3177b7f 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` in `conf/offpolicy/task/sac/g1_wbt_obs/mujoco.yaml` | Per-term history for `gyro`, `joint_pos_rel`, `dof_vel`, and `last_actions`. | +| 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`. | | Sharpa tactile contact latency | `domain_rand.contact_latency` in Sharpa in-hand configs | Keeps previous tactile contact values for sampled contact channels. | -| Obs history ordering guard | `tests/scripts/test_obs_alignment_g1_wbt.py` | Asserts per-term oldest-first flatten for the G1 WBT actor obs. | +| Deploy-side ONNX contract check | `scripts/deploy/sim_prototype.py` | Validates `obs_layout`, `obs_dim`, ONNX input width, clipping, and EMA action smoothing for the G1 WBT path. | ## Action Latency @@ -30,13 +30,23 @@ The checked-in G1 WBT owner enables this flag in ## Observation Lag And History -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. +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. -Do not lag command/reference terms unless the training owner did so. +```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. ## Deploy-Side Measurements @@ -48,10 +58,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 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. +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. ## 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 bbf4e607a..219afa77c 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,13 +27,12 @@ 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 - `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. + 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. ## Cube drops in Allegro / Sharpa inhand @@ -56,7 +55,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 the exact task owner YAML path it was trained from. +- `policy.onnx` and, for the G1 WBT path, `deploy_config.yaml`. - 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 c17efd7df..cf7a16812 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`），逐项按最旧优先展平，使得按逐项历史装配的 -硬件运行时读到同一个向量。该顺序由 `tests/scripts/test_obs_alignment_g1_wbt.py` -守护；硬件侧契约见仿真到真机部署指南。当 Motrix sim2sim 回放需要引用其他日志根目录下的 -checkpoint 时，用 `uv run eval` 透传绝对路径： +（`noise_config.obs_history_length: 5`），与部署侧的 `ObservationManager` 按字节对齐。 +部署工具在 `scripts/deploy/`，观测对齐由 `tests/scripts/test_obs_alignment_g1_wbt.py` +交叉校验。当 Motrix sim2sim 回放需要引用其他日志根目录下的 checkpoint 时，用 +`uv run eval` 透传绝对路径： ```bash uv run eval --algo sac --task g1_motion_tracking --sim motrix \ 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 1e9879e9c..223c74a92 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,7 +5,8 @@ ## "仿真到真机"在 UniLab 中的含义 一个可部署的 UniLab 策略，是导出的策略加上所选任务 owner 使用的那套精确的观测与 -动作契约。UniLab 提供训练侧与 ONNX 导出；每种机器人都需要一个硬件侧运行时，它需要： +动作契约。G1 WBT 辅助路径将其物化为 `policy.onnx`、`deploy_config.yaml` 以及一个 +运动二进制文件；其他机器人需要一个等价的硬件侧运行时，它需要： 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 9402d128b..df79c491e 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 自由度变体）。关节顺序来自任务 owner 的场景 -（`src/unilab/assets/robots/g1/scene_flat.xml`，按 actuator 顺序）；在硬件上机前请 -先核对该顺序与你的 SDK 电机索引是否一致。 +Unitree G1 人形机器人（29 自由度变体）。假定关节顺序与 +`scripts/deploy/export_deploy_config.py` 从 +`src/unilab/assets/robots/g1/scene_flat.xml` 导出的顺序一致；在硬件上机前请核对该 +顺序。 :::: -本指南说明 G1 运动跟踪策略在硬件上所期望的**观测与动作契约**。仓库不提供 G1 部署侧 -运行时——硬件侧回路由你实现，本页告诉你它必须复现哪些内容。 +本指南讲解从一个收敛的 G1 运动跟踪策略到机器人上闭环运行之间的**最后一公里**。 ## 0. 验证你的仿真侧检查点 ```bash # Replay the policy headlessly and produce a video. -uv run eval --algo sac --task g1_wbt_obs --sim mujoco --load-run -1 \ +uv run eval --algo ppo --task g1_motion_tracking --sim motrix --load-run -1 \ --render-mode record ``` @@ -24,48 +24,37 @@ uv run eval --algo sac --task g1_wbt_obs --sim mujoco --load-run -1 \ - 关节速度与动作保持有限且在预期范围内。 - 接触时序看起来与参考运动一致。 -如果其中任何一项不对，请在硬件上机前修复仿真侧检查点。 +如果其中任何一项不对，请在硬件上机前修复仿真侧检查点或部署契约。 -## 1. 先选定 owner，再从 YAML 读取契约 +## 1. 导出 -硬件回路需要的每个字段都在任务 owner YAML 中声明。面向部署的 G1 owner 有： +使用训练回放路径导出 `policy.onnx`，然后用已提交的部署辅助工具导出 G1 WBT 部署配置 +与运动二进制文件： -```{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` 列表。 -``` +```bash +uv run eval --algo ppo --task g1_motion_tracking --sim motrix --load-run -1 -::::{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，而不是硬件调参问题。 -:::: +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 +``` -通过训练回放路径导出 `policy.onnx`： +部署侧原型消费如下文件： -```bash -uv run eval --algo sac --task g1_wbt_obs --sim mujoco --load-run -1 +``` +runs// +└── policy.onnx +logs/deploy/ +├── deploy_config.yaml +└── dance1.bin ``` ## 2. 观测契约 -对 `g1_wbt_obs`，actor 观测按如下顺序装配（见 -`src/unilab/envs/motion_tracking/g1/tracking_obs.py` 的 `_build_actor_obs`）。 -单步参考项在前，随后是各 proprio 项的完整历史，按**最旧优先**展平： +对于已提交的 G1 WBT 部署辅助工具，观测布局会作为 `obs_layout` 导出到 +`deploy_config.yaml`。`scripts/deploy/export_deploy_config.py` 是分段顺序的权威来源： ```{list-table} :header-rows: 1 @@ -85,37 +74,30 @@ uv run eval --algo sac --task g1_wbt_obs --sim mujoco --load-run -1 - 来自参考帧与机器人躯干帧的锚点朝向项 * - `gyro` - 每个历史步 3 - - IMU 陀螺仪项（`env.sensor.gyro`，该 owner 为 `pelvis_gyro`） + - IMU 陀螺仪项 * - `joint_pos_rel` - 每个历史步 29 - - 测量到的关节位置减去 `stand` keyframe 的关节角 + - 测量到的关节位置减去 `default_angles` * - `dof_vel` - 每个历史步 29 - 关节速度项 * - `last_actions` - - 每个历史步 29 + - 29 - 上一步的原始 actor 输出 ``` -历史深度 `H` 即 `env.noise_config.obs_history_length`（该 owner 为 5）。逐项的 -最旧优先顺序由 `tests/scripts/test_obs_alignment_g1_wbt.py` 守护；硬件侧必须镜像 -该顺序，否则策略读到的是被置换过的向量。 +导出脚本还会记录每个分段的 `history_length` 并校验总的 `obs_dim`。当 ONNX 输入宽度 +与 `deploy_config.yaml` 的 `obs_dim` 不一致时，`scripts/deploy/sim_prototype.py` 会 +拒绝运行。 ## 3. 执行器接口 -将 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`）。 +G1 部署原型将 actor 输出严格映射为： +`action * action_scale + default_angles`，然后钳制到 `joint_lower` / +`joint_upper`，并应用来自 `ema_alpha` 的 EMA 平滑。 -训练侧直接施加目标，不做平滑。如果硬件抖动迫使你加入平滑，请先验证其 sim2sim -影响——每一步滞后都会把观测推离训练分布。 +- 动作 = 目标关节位置，按 `deploy_config.yaml` 中的 `action_scale` 项**缩放**。 +- 在目标到达电机驱动器之前，将其钳制到生成的关节范围内。 ## 4. 参考运动同步 @@ -136,7 +118,7 @@ uv run eval --algo sac --task g1_wbt_obs --sim mujoco --load-run -1 硬件侧：标准结构见 {doc}`7-safety_layers`。G1 的具体事项： - 在应用 `action_scale` 之前拒绝非有限动作与形状不匹配。 -- 用 owner 场景 XML 中的关节范围钳制生成的目标。 +- 用 `deploy_config.yaml` 中的 `joint_lower` / `joint_upper` 钳制生成的目标。 - 把看门狗、姿态监控以及操作员停止阈值保留在部署控制器中，并独立于策略对它们进行 测试。 @@ -152,9 +134,18 @@ uv run eval --algo sac --task g1_wbt_obs --sim mujoco --load-run -1 ## 7. 应记录什么 -为每一步记录**完整的观测向量**、**完整的动作向量**与**墙钟**。把第一段硬件观测窗口 -与用同一份 owner YAML 构建的仿真回合作对比——这个 diff 比任何奖励检查都更快定位 -单位、坐标系与顺序错误。 +为每一步记录**完整的观测向量**、**完整的动作向量**与**墙钟**。在硬件上机前，通过 +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，而不是 +硬件调参问题。 ## 另请参阅 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 1a3ebe27b..198f5e902 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,18 +24,36 @@ uv run eval --algo sac --task g1_walk_flat --sim mujoco --load-run -1 ``` `uv run eval` 设置回放模式，并把 `--load-run` 映射到所路由训练脚本使用的检查点 -选择器。导出的文件会写入所选的运行目录。用于部署时，请把导出的 `policy.onnx` 与训练 -它所用的任务 owner YAML 放在一起——该 YAML 是运行时必须复现的观测与动作契约的权威 -来源。 +选择器。导出的文件会写入所选的运行目录。对于部署原型，请把导出的 `policy.onnx` 与 +运行时所用的部署侧配置和运动资产放在一起。 -## 校验导出的计算图 +## G1 部署原型 -回放路径在写出之前会把导出的计算图与 PyTorch 比对，因此导出成功本身已经建立了数值 -一致性。它**没有**建立的是：你的硬件侧回路装配出的输入向量是否相同。在硬件上机前： +已提交的 G1 WBT 部署辅助工具使用如下产物： -- 从 env 读取 actor 观测宽度（而不是照抄文档表格），确认它与 ONNX 输入宽度一致。 -- 对照 owner 的 `_build_actor_obs` 确认分项顺序与逐项历史顺序。G1 全身跟踪见 - {doc}`2-g1_whole_body`。 +| 产物 | 生产者 | +| --- | --- | +| `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 中驱动策略。 ## 另请参阅 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 0e1bbb7b3..0c5c25df0 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,31 +40,34 @@ flowchart LR OP -.->|E-stop| D ``` -把硬实时的安全检查放在部署控制器中，而不是训练脚本里。仓库不实现生产级的电机驱动器 -安全回路——该边界由你构建并测试。 +把硬实时的安全检查放在部署控制器中，而不是训练脚本里。仓库的 G1 辅助路径导出部署 +配置并运行一个 MuJoCo 原型；它并不实现生产级的电机驱动器安全回路。 ## 策略假定你已配置的内容 -策略期望的是其训练 owner 所声明的动作映射与限位。对 G1 WBT owner -（`conf/offpolicy/task/sac/g1_wbt_obs/mujoco.yaml`）： +G1 部署辅助工具会把这些字段导出到 `deploy_config.yaml`： -| 量 | 权威来源 | -| --- | --- | -| `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` | +```yaml +action_scale: 2.0 +ema_alpha: 1.0 +default_angles: [...] +joint_lower: [...] +joint_upper: [...] +kp: [...] +kd: [...] +``` -请从 owner YAML 及其场景派生这些量，并原样复现 owner 的 `action_scale` **形态**—— -标量 owner 与逐关节列表 owner 不可互换。不要把关节范围或增益手工复制到第二处，那会 -与资产静默漂移。 +`scripts/deploy/sim_prototype.py` 消费同样的字段，并应用 +`action * action_scale + default_angles`、关节钳制与 EMA 平滑。硬件控制器应当消费 +生成的配置，而不是手动复制关节范围或增益。 ## 交接测试 在把 策略 → 安全层 → 电机 集成起来之前，先隔离测试安全层： 1. 注入一个 NaN 动作，验证该指令被拒绝。 -2. 注入一个超范围的关节目标，验证钳制使用了 owner 场景 XML 中的关节范围。 +2. 注入一个超范围的关节目标，验证钳制使用了 `deploy_config.yaml` 中的 + `joint_lower` / `joint_upper`。 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 6f37f6179..0d972f6f6 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 观测历史 | `conf/offpolicy/task/sac/g1_wbt_obs/mujoco.yaml` 中的 `noise_config.obs_history_length` | 为 `gyro`、`joint_pos_rel`、`dof_vel` 与 `last_actions` 提供逐项历史。 | +| G1 WBT 观测历史 | `noise_config.obs_history_length` 与 `scripts/deploy/export_deploy_config.py` | 为 `gyro`、`joint_pos_rel`、`dof_vel` 与 `last_actions` 导出逐项的 `obs_layout` 历史。 | | Sharpa 触觉接触延迟 | Sharpa 手内配置中的 `domain_rand.contact_latency` | 为采样到的接触通道保留上一步的触觉接触值。 | -| 观测历史顺序守护 | `tests/scripts/test_obs_alignment_g1_wbt.py` | 断言 G1 WBT actor 观测按逐项最旧优先展平。 | +| 部署侧 ONNX 契约检查 | `scripts/deploy/sim_prototype.py` | 为 G1 WBT 路径校验 `obs_layout`、`obs_dim`、ONNX 输入宽度、钳制以及 EMA 动作平滑。 | ## 动作延迟 @@ -28,11 +28,21 @@ env: ## 观测滞后与历史 -观测宽度是 owner `noise_config` 的函数，硬件运行时不允许猜测。对 G1 WBT owner， -`obs_history_length: 5` 让每个本体感受项携带 5 步历史并按最旧优先展平，而参考项保持 -单步。完整的分项顺序见 {doc}`2-g1_whole_body`。 +G1 WBT 部署辅助工具不会猜测观测宽度。它们导出一份带有 `obs_layout`、逐项 +`history_length` 与 `obs_dim` 的模式；随后原型组装出同样的布局，并拒绝不匹配。 -除非训练 owner 这样做了，否则不要让指令/参考项滞后。 +```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 模式中，参考项保持 +单步，而本体感受项携带历史。 ## 部署侧测量 @@ -44,8 +54,8 @@ env: 4. 执行器指令的发送时间戳 5. 钳制 / 平滑前后的动作向量 -将观测向量与用同一份任务 owner YAML 构建的仿真回合作对比。如果实测管线需要滤波或 -缓冲，请把相匹配的行为编码到任务 owner 中并重新训练，而不是只在部署侧添加。 +将观测向量与用同一份 `deploy_config.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 12593a21a..f601b5560 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,11 +25,10 @@ 1. **关节顺序被调换。** 检查 `policy.onnx` 的输入宽度与你电机驱动器中的关节顺序。 用 `unilab-export-scene` 导出训练时的关节顺序。 2. **动作缩放单位不匹配。** 策略输出未缩放的值；驱动器期望的是弧度，而你喂给它的 - 是归一化的 [-1, 1]。在把目标发送给驱动器之前，应用训练 owner YAML 中的 - `env.control_config.action_scale` / 默认角度约定，并原样复现该 owner 的标量或 - 列表形态。 -3. **观测布局不匹配。** 将硬件回路装配出的内容与训练 owner 的 `_build_actor_obs` - 比较——先比分项顺序，再比逐项历史顺序。 + 是归一化的 [-1, 1]。在把目标发送给驱动器之前，应用 `deploy_config.yaml` 中的 + `action_scale` / 默认角度约定。 +3. **观测布局不匹配。** 将 `deploy_config.yaml` 的 `obs_layout` 与训练 owner 比较， + 并在硬件上运行前用 `scripts/deploy/sim_prototype.py` 进行验证。 ## Allegro / Sharpa 手内操作中方块掉落 @@ -52,7 +51,7 @@ - 完整的硬件轨迹（整段运行的 `obs / action / wall_clock`）。 - 用于训练的仿真侧 YAML：`runs//config.yaml`。 -- `policy.onnx`，以及训练它所用的确切任务 owner YAML 路径。 +- `policy.onnx`，以及对于 G1 WBT 路径的 `deploy_config.yaml`。 - 一段使用**同一**种子的仿真回合视频：`eval --seed --render-mode record`。 - 如果有的话，该运行所在 commit 与 `main` 之间的 `git diff`。 diff --git a/scripts/deploy/append_cooldown.py b/scripts/deploy/append_cooldown.py new file mode 100644 index 000000000..8ad20eb8d --- /dev/null +++ b/scripts/deploy/append_cooldown.py @@ -0,0 +1,299 @@ +#!/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 new file mode 100644 index 000000000..b692e992e --- /dev/null +++ b/scripts/deploy/export_deploy_config.py @@ -0,0 +1,328 @@ +#!/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 new file mode 100644 index 000000000..27905845c --- /dev/null +++ b/scripts/deploy/export_motion_bin.py @@ -0,0 +1,179 @@ +#!/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 new file mode 100644 index 000000000..0d0c2c656 --- /dev/null +++ b/scripts/deploy/prepend_warmup.py @@ -0,0 +1,215 @@ +#!/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 new file mode 100644 index 000000000..7043b9bd2 --- /dev/null +++ b/scripts/deploy/sim_prototype.py @@ -0,0 +1,520 @@ +#!/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 3ed23a902..270c4bdc6 100644 --- a/src/unilab/envs/motion_tracking/g1/tracking_obs.py +++ b/src/unilab/envs/motion_tracking/g1/tracking_obs.py @@ -1,3 +1,26 @@ +"""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 @@ -30,6 +53,10 @@ ) from .tracking_sac import G1MotionTrackingSACCfg, G1MotionTrackingSACEnv +# --------------------------------------------------------------------------- # +# Config extensions +# --------------------------------------------------------------------------- # + @dataclass class ObsNoiseConfig(NoiseConfig): @@ -39,10 +66,17 @@ 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 @@ -55,9 +89,13 @@ 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$" @@ -72,6 +110,11 @@ 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.""" @@ -113,6 +156,9 @@ 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()), } @@ -127,6 +173,7 @@ 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 @@ -143,6 +190,7 @@ 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: @@ -168,6 +216,11 @@ 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): @@ -182,10 +235,16 @@ 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 @@ -198,8 +257,11 @@ 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) @@ -231,10 +293,19 @@ 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 @@ -264,15 +335,19 @@ 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 + single_step = 2 * n + 6 # command(2n) + anchor_ori(6) if not nc.enable_zero_anchor_pos: single_step += 3 if not nc.enable_zero_linvel: single_step += 3 - proprio_step = 3 + 3 * n + proprio_step = 3 + 3 * n # gyro + joint_pos_rel + dof_vel + last_actions return single_step + H * proprio_step def _compute_obs( @@ -286,6 +361,8 @@ 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( @@ -301,6 +378,8 @@ 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 @@ -317,21 +396,28 @@ 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) @@ -339,6 +425,7 @@ 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, @@ -352,13 +439,24 @@ 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] + buf = self._hist_buf[key][sel] # (n_e, H, D) actor_terms.append(buf.reshape(buf.shape[0], -1)) else: actor_terms.extend([noisy_gyro, noisy_joint_pos_rel, noisy_dof_vel, last_actions]) return np.concatenate(actor_terms, axis=1, dtype=get_global_dtype()) + # ------------------------------------------------------------------ # + # Proprio history buffer maintenance. + # Mirrors deploy ``ObservationManager`` / ``ObservationTermCfg``: + # * On reset: fill all H slots with the current value (matches + # ``ObservationTermCfg::reset`` which calls ``add()`` H times). + # * On step: pop oldest, push current at end. + # * Read order is oldest-first, so ``flatten(buf[env, :, :])`` yields + # ``[t-H+1, t-H+2, ..., t]`` — matches deploy + # ``ObservationTermCfg::get`` (deque front-to-back). + # ------------------------------------------------------------------ # + def _push_obs_history( self, env_ids: np.ndarray | None, components: dict[str, np.ndarray] ) -> None: diff --git a/tests/scripts/test_obs_alignment_g1_wbt.py b/tests/scripts/test_obs_alignment_g1_wbt.py index 454ea3fb6..93b5323b8 100644 --- a/tests/scripts/test_obs_alignment_g1_wbt.py +++ b/tests/scripts/test_obs_alignment_g1_wbt.py @@ -1,13 +1,56 @@ +"""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: - """Per-term history buffer with deque semantics. + """Bit-for-bit copy of ObservationTermCfg buffer semantics. reset(obs): add(obs) H times (== fill). add(obs): push at back; if buffer > H, pop_front. @@ -35,16 +78,20 @@ def get(self) -> np.ndarray: def _deploy_compute_group(layout, current_segments_by_name): - """Assemble per step: each term's full history oldest-first, then concat.""" + """Mirror ObservationManager::compute_group with use_gym_history=False.""" 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 segments in current_segments_by_name: + for step_idx, segments in enumerate(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 ) @@ -52,6 +99,12 @@ 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 @@ -68,6 +121,7 @@ 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, } @@ -90,6 +144,11 @@ 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 @@ -99,24 +158,51 @@ def test_layout_sum_matches_obs_dim(self, deploy_cfg): assert total == deploy_cfg["obs_dim"] -class TestHistoryOrdering: - """History blocks must read oldest-first, not newest-first.""" +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}" + ) def test_history_terms_carry_oldest_first(self, deploy_cfg, rng): - """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"] + """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) + 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] - all_segments.append(seg) - - obs = _deploy_compute_group(layout, all_segments)[-1] + obs = assembler.step(seg) + # 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) @@ -132,14 +218,16 @@ class TestTrainingAssemblerVsDeploy: """ @staticmethod - def _training_actor_obs(history_buf, current_refs, num_envs): + def _training_actor_obs(history_buf, current_refs, hist_components, num_envs): + # refs (single-step) 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] + buf = history_buf[key] # (num_envs, H, D) parts.append(buf.reshape(num_envs, -1)) return np.concatenate(parts, axis=1).astype(np.float32) @@ -148,6 +236,7 @@ def test_training_path_matches_deploy(self, deploy_cfg, rng): n = deploy_cfg["action_dim"] H = 5 + # Initial buffer of zeros (matches tracking_obs.py allocation). buf = { "gyro": np.zeros((n_env, H, 3), dtype=np.float32), "joint_pos_rel": np.zeros((n_env, H, n), dtype=np.float32), @@ -158,6 +247,7 @@ def test_training_path_matches_deploy(self, deploy_cfg, rng): all_segments = [_random_segments(rng) for _ in range(15)] deploy_seq = _deploy_compute_group(deploy_cfg["obs_layout"], all_segments) + # Step 0: reset (matches tracking_obs.py is_reset=True path). s0 = all_segments[0] for key in ("gyro", "joint_pos_rel", "dof_vel", "last_actions"): buf[key][:, :, :] = s0[key][None, None, :] @@ -166,9 +256,10 @@ 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, n_env) + train_obs0 = self._training_actor_obs(buf, refs0, s0, n_env) np.testing.assert_array_equal(train_obs0[0], deploy_seq[0]) + # Steps 1..N-1: push (matches tracking_obs.py is_reset=False path). for k in range(1, len(all_segments)): sk = all_segments[k] for key in ("gyro", "joint_pos_rel", "dof_vel", "last_actions"): @@ -179,19 +270,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, n_env) + train_obs = self._training_actor_obs(buf, refs, sk, 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 plain-concat 154-d path bit-exact.""" + """H=1 ('no history') must reproduce the pre-history 154-d path bit-exact.""" @pytest.fixture - def legacy_layout(self): + def legacy_cfg(self): n = 29 - return [ + obs_layout = [ {"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}, @@ -200,17 +291,22 @@ def legacy_layout(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_layout): - total = sum(s["dim"] * s["history_length"] for s in legacy_layout) - assert total == 154 + def test_legacy_obs_dim_154(self, legacy_cfg): + assert legacy_cfg["obs_dim"] == 154 - 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) + def test_legacy_matches_simple_concat(self, legacy_cfg, rng): + sp = _load_sim_prototype() + assembler = sp.ObsAssembler(legacy_cfg) - for seg, obs in zip(all_segments, assembled): - expected = np.concatenate([seg[s["name"]] for s in legacy_layout]) + 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"]]) np.testing.assert_array_equal(obs, expected)