From 121a1c83ae10a7457666897f240f7dd0321759c4 Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Sun, 27 Sep 2026 22:16:04 +0800 Subject: [PATCH] feat: G1 backflip whole-body tracking task under FastSAC MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add the `g1-wbt-backflip` manager-based environment tracking a bundled G1 backflip reference motion with the Motrix FastSAC recipe (60k iterations by default), plus documentation with hardware-attributed training evidence. Environment (motrix_envs): - G1 backflip preset with dt=0.005 (the launch impulse spans ~0.1 s of ground contact and is lost to a coarser contact solve), mixed frame sampling (20% frame-0 / 80% uniform with exact-velocity mid-air teleports) and hold-at-clip-end landing episodes - Terminations follow holosoma's BadTrackingZOnly shape: under the command's per-step re-anchoring, world-z is the only anchor axis with a physical reference (gravity), so bad_ref_z (0.5 m, the launch-forcing term) is the sole anchor-position check, backed by the tilt-only orientation check and body-z/NaN guards - Rewards: holosoma-aligned tracking terms plus a deliberately light action_rate_l2 and full-3D end-effector tracking for landing accuracy - wbt/g1.py split into the wbt/g1/ package (common base cfg, one file per task: largebox, dance, backflip) Trainer (motrix_rl): - resolve_trainer_topology now owns the whole device layout from the raw specs: learner devices, collector inference devices with owner co-location (generic "cuda" follows the owning learner's GPU for every learner count), env sharding, NUMA/CPU bindings and ring/weight transports. Workers receive explicit devices from the topology, eliminating the in-process torch.cuda.current_device() re-resolution that silently crossed GPUs when the startup backend probe (JAX/XLA) left the CUDA current device on the last probed GPU — restoring the CUDA-IPC transports (+11% iter/s) and fixing a latent parent/child device mismatch - runs annotate hardware provenance into metadata.json via system_metrics.capture_system_info: the async trainer records the devices it actually uses (with roles, from the resolved topology) after layout resolution and end-of-training performance (wall time, env steps, throughput) Performance snapshots (docs): - schema v2: the committed per-env JSON keeps the default profile inline and adds hardware provenance plus other_profiles preserved across machines; the SVG chart carries a centered hardware annotation above the legend. g1-wbt-backflip ships with a v2 snapshot generated from the 60k run (1x RTX 4090, EPYC 9J14, ~111k env-steps/s) Documentation: - WBT user guide: backflip demo video and poster, task table row with training curve, and the stale WbtManagerEnvCfg references renamed to WbtEnvCfg Tests: - test_all_envs smokes each registered environment in its own child process with a serial numba-cache warmup (hosting several live manager models in one process is not a supported runtime contract; see motrixsim issue #2171), a fixed worker pool and per-case timeouts --- .../task/g1-wbt-backflip/motrix.fastsac.yaml | 27 ++ docs/performance.yaml | 10 + docs/scripts/generate_performance_svg.py | 126 ++++++++- .../data/performance/g1-wbt-backflip.json | 3 + .../images/performance/g1-wbt-backflip.svg | 3 + .../_static/images/poster/g1-wbt-backflip.jpg | 3 + .../source/_static/videos/g1-wbt-backflip.mp4 | 3 + .../whole_body_tracking/adding_wbt_task.md | 19 +- .../envs/whole_body_tracking/env_design.md | 2 +- .../envs/whole_body_tracking/index.md | 29 +- .../envs/whole_body_tracking/motion_format.md | 2 +- .../whole_body_tracking/adding_wbt_task.md | 19 +- .../envs/whole_body_tracking/env_design.md | 2 +- .../envs/whole_body_tracking/index.md | 29 +- .../envs/whole_body_tracking/motion_format.md | 2 +- .../wbt/assets/motion/g1/backflip.npz | 3 + .../src/motrix_envs/locomotion/wbt/cfg.py | 1 - .../motrix_envs/locomotion/wbt/g1/__init__.py | 12 + .../motrix_envs/locomotion/wbt/g1/backflip.py | 165 +++++++++++ .../locomotion/wbt/{g1.py => g1/common.py} | 44 +-- .../motrix_envs/locomotion/wbt/g1/dance.py | 23 ++ .../motrix_envs/locomotion/wbt/g1/largebox.py | 23 ++ .../motrix_envs/locomotion/wbt/mdp/command.py | 17 +- .../motrix_envs/locomotion/wbt/mdp/reset.py | 5 + .../motrix_envs/locomotion/wbt/mdp/rewards.py | 37 +++ .../locomotion/wbt/mdp/terminations.py | 44 +++ motrix_envs/tests/test_package_boundary.py | 3 +- motrix_envs/tests/test_wbt_numba.py | 2 +- .../motrix_rl/fastsac/async_impl/topology.py | 87 +++++- .../src/motrix_rl/fastsac/async_impl/train.py | 62 ++-- .../motrix_rl/fastsac/async_impl/worker.py | 5 +- motrix_rl/src/motrix_rl/runs.py | 18 ++ motrix_rl/src/motrix_rl/system_metrics.py | 119 +++++++- motrix_rl/src/motrix_rl/utils.py | 12 + .../tests/test_doc_generate_video_script.py | 2 +- motrix_rl/tests/test_fastsac_async_multi.py | 82 ++++-- scripts/private/g1_flip_converter.py | 267 ++++++++++++++++++ test/test_all_envs.py | 89 ++++-- 38 files changed, 1223 insertions(+), 178 deletions(-) create mode 100644 configs/task/g1-wbt-backflip/motrix.fastsac.yaml create mode 100644 docs/source/_static/data/performance/g1-wbt-backflip.json create mode 100644 docs/source/_static/images/performance/g1-wbt-backflip.svg create mode 100644 docs/source/_static/images/poster/g1-wbt-backflip.jpg create mode 100644 docs/source/_static/videos/g1-wbt-backflip.mp4 create mode 100644 motrix_envs/src/motrix_envs/locomotion/wbt/assets/motion/g1/backflip.npz create mode 100644 motrix_envs/src/motrix_envs/locomotion/wbt/g1/__init__.py create mode 100644 motrix_envs/src/motrix_envs/locomotion/wbt/g1/backflip.py rename motrix_envs/src/motrix_envs/locomotion/wbt/{g1.py => g1/common.py} (67%) create mode 100644 motrix_envs/src/motrix_envs/locomotion/wbt/g1/dance.py create mode 100644 motrix_envs/src/motrix_envs/locomotion/wbt/g1/largebox.py create mode 100644 scripts/private/g1_flip_converter.py diff --git a/configs/task/g1-wbt-backflip/motrix.fastsac.yaml b/configs/task/g1-wbt-backflip/motrix.fastsac.yaml new file mode 100644 index 00000000..6cfc51e6 --- /dev/null +++ b/configs/task/g1-wbt-backflip/motrix.fastsac.yaml @@ -0,0 +1,27 @@ +# @package _global_ +# Task g1-wbt-backflip: manager-based G1 backflip tracking with Motrix FastSAC. +defaults: + - /algo_base@algo: motrix.fastsac + - _self_ +task: + env: g1-wbt-backflip + rllib: motrix + algo: fastsac +num_envs: 2048 +play_num_envs: 16 +seed: 1 +checkpoint: + interval: 1000 +algo: + agent: + num_updates: 4 + policy_frequency: 2 + gamma: 0.99 + tau: 0.05 + target_entropy_ratio: 0.5 + num_atoms: 501 + trainer: + num_learning_iterations: 60000 + async_options: + # Keep the WBT training UTD fixed at four updates per ingested batch. + utd_mode: strict diff --git a/docs/performance.yaml b/docs/performance.yaml index a3c033f2..524eff51 100644 --- a/docs/performance.yaml +++ b/docs/performance.yaml @@ -47,6 +47,16 @@ benchmarks: zh_CN: 惩罚缩放 default_seed: 1 benchmark_seeds: [1] + g1-wbt-backflip: + task: g1-wbt-backflip/motrix.fastsac + template: wbt + metric: rollout/mean_return + survival_metric: rollout/mean_ep_len + survival_metric_label: + en: Episode survival (%) + zh_CN: Episode 存活率（%） + episode_length_max: 500 + default_seed: 1 g1-wbt-dance: task: g1-wbt-dance/motrix.fastsac template: wbt diff --git a/docs/scripts/generate_performance_svg.py b/docs/scripts/generate_performance_svg.py index a33996a0..ff033c92 100644 --- a/docs/scripts/generate_performance_svg.py +++ b/docs/scripts/generate_performance_svg.py @@ -98,6 +98,7 @@ class RunSeries: values: np.ndarray curriculum_values: np.ndarray | None = None survival_values: np.ndarray | None = None + provenance: dict | None = None @dataclass(frozen=True) @@ -429,6 +430,7 @@ def select_run_series( values=scalar.values, curriculum_values=curriculum_values, survival_values=survival_values, + provenance=metadata, ) ], ) @@ -555,6 +557,7 @@ def _render_svg( *, has_curriculum: bool, has_survival: bool = False, + annotation: str | None = None, ) -> str: """Render the shared SVG layout for the configured template.""" if has_curriculum and has_survival: @@ -659,6 +662,15 @@ def secondary_y_coord(value: float) -> float: '', f"", ] + if annotation: + # Header metadata line centered above the legend: the bottom rows are + # reserved for the axis titles, and the earlier bottom-right footer + # placement sat too close to the canvas edge (a top-right placement + # collided with the centered legend at this font size). + lines.append( + f'{escape(annotation)}' + ) legend_entries = [("curve", spec.metric_label["en"])] if has_secondary: assert secondary_label is not None @@ -669,8 +681,8 @@ def secondary_y_coord(value: float) -> float: for (css_class, label), entry_width in zip(legend_entries, legend_widths): lines.extend( [ - f'', - f'{escape(label)}', + f'', + f'{escape(label)}', ] ) legend_x += entry_width + 28 @@ -753,23 +765,23 @@ def secondary_y_coord(value: float) -> float: return "\n".join(lines) + "\n" -def render_default_svg(spec: BenchmarkSpec, aggregate: Aggregate, run_count: int) -> str: +def render_default_svg(spec: BenchmarkSpec, aggregate: Aggregate, run_count: int, annotation: str | None = None) -> str: """Render the standard single-axis performance template.""" - return _render_svg(spec, aggregate, run_count, has_curriculum=False) + return _render_svg(spec, aggregate, run_count, has_curriculum=False, annotation=annotation) -def render_curriculum_svg(spec: BenchmarkSpec, aggregate: Aggregate, run_count: int) -> str: +def render_curriculum_svg(spec: BenchmarkSpec, aggregate: Aggregate, run_count: int, annotation: str | None = None) -> str: """Render performance and curriculum progress on independent y axes.""" if aggregate.curriculum_values is None or spec.curriculum_metric_label is None: raise RuntimeError(f"Curriculum template for {spec.benchmark_id!r} has no curriculum series") - return _render_svg(spec, aggregate, run_count, has_curriculum=True) + return _render_svg(spec, aggregate, run_count, has_curriculum=True, annotation=annotation) -def render_wbt_svg(spec: BenchmarkSpec, aggregate: Aggregate, run_count: int) -> str: +def render_wbt_svg(spec: BenchmarkSpec, aggregate: Aggregate, run_count: int, annotation: str | None = None) -> str: """Render WBT return and normalized episode survival on independent y axes.""" if aggregate.survival_values is None or spec.survival_metric_label is None: raise RuntimeError(f"WBT template for {spec.benchmark_id!r} has no survival series") - return _render_svg(spec, aggregate, run_count, has_curriculum=False, has_survival=True) + return _render_svg(spec, aggregate, run_count, has_curriculum=False, has_survival=True, annotation=annotation) SVG_RENDERERS = { @@ -779,18 +791,93 @@ def render_wbt_svg(spec: BenchmarkSpec, aggregate: Aggregate, run_count: int) -> } -def render_svg(spec: BenchmarkSpec, aggregate: Aggregate, run_count: int) -> str: +def render_svg(spec: BenchmarkSpec, aggregate: Aggregate, run_count: int, annotation: str | None = None) -> str: """Dispatch SVG rendering through the selected template.""" try: renderer = SVG_RENDERERS[spec.template] except KeyError as exc: raise RuntimeError(f"Unsupported performance template: {spec.template!r}") from exc - return renderer(spec, aggregate, run_count) + return renderer(spec, aggregate, run_count, annotation=annotation) + + +def _hardware_annotation(series: list[RunSeries]) -> str | None: + """Concise used-hardware line for the chart, from the runs' provenance.""" + + for item in series: + provenance = item.provenance + if not (isinstance(provenance, dict) and isinstance(provenance.get("system"), dict)): + continue + system = provenance["system"] + parts = [] + gpus = system.get("gpus_used") or [] + if gpus: + model = str(gpus[0].get("model", "GPU")).replace("NVIDIA GeForce ", "") + parts.append(f"{len(gpus)}× {model}") + cpu_model = str(system.get("cpu", {}).get("model") or "") + if cpu_model: + parts.append(cpu_model) + performance = provenance.get("performance") + if isinstance(performance, dict): + wall_minutes = performance.get("wall_time_s") + if isinstance(wall_minutes, (int, float)) and wall_minutes > 0: + parts.append(f"{wall_minutes / 60.0:.0f} min") + rate = performance.get("mean_env_steps_per_s") + if isinstance(rate, (int, float)) and rate > 0: + parts.append(f"{rate / 1000.0:.0f}k env-steps/s") + return " · ".join(parts) if parts else None + return None + +def _profile_slug(series: list[RunSeries]) -> str | None: + """Hardware slug of a series' provenance, or None when unrecorded.""" -def render_data(spec: BenchmarkSpec, aggregate: Aggregate, series: list[RunSeries]) -> str: - """Render a stable JSON sidecar containing chart provenance and aggregate values.""" + for item in series: + provenance = item.provenance + if isinstance(provenance, dict) and isinstance(provenance.get("system"), dict): + from motrix_rl.system_metrics import hardware_profile_slug + + return hardware_profile_slug(provenance["system"]) + return None + + +def render_data(spec: BenchmarkSpec, aggregate: Aggregate, series: list[RunSeries], data_path: Path | None = None) -> str: + """Render a stable JSON sidecar containing chart provenance and aggregate values. + + Schema v2 keeps the default profile inline at the top level (the shape + readers consume) and adds hardware provenance: ``default_profile`` names + the used-hardware slug of the inline points, ``provenance`` carries the + runs' system/performance records, and ``other_profiles`` preserves full + payloads generated on other machines so regenerating on one host never + drops another host's data. + """ + slug = _profile_slug(series) + other_profiles: dict[str, dict] = {} + if data_path is not None and data_path.is_file(): + try: + previous = json.loads(data_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + previous = None + if isinstance(previous, dict): + other_profiles.update(previous.get("other_profiles") or {}) + previous_slug = previous.get("default_profile") + if isinstance(previous_slug, str) and previous_slug != slug: + other_profiles[previous_slug] = { + key: previous[key] for key in ("benchmark", "task", "template", "metric", "metric_label", "points") if key in previous + } + if "provenance" in previous: + other_profiles[previous_slug]["provenance"] = previous["provenance"] + provenance_runs = [] + for item in series: + record = {} + if isinstance(item.provenance, dict): + for key in ("system", "performance"): + if item.provenance.get(key) is not None: + record[key] = item.provenance[key] + provenance_runs.append(record) payload = { + "schema_version": 2, + **({"default_profile": slug} if slug is not None else {}), + **({"other_profiles": other_profiles} if other_profiles else {}), "benchmark": spec.benchmark_id, "task": spec.task, "template": spec.template, @@ -806,6 +893,7 @@ def render_data(spec: BenchmarkSpec, aggregate: Aggregate, series: list[RunSerie **({"survival_metric_label": spec.survival_metric_label} if spec.survival_metric_label is not None else {}), **({"episode_length_max": spec.episode_length_max} if spec.episode_length_max is not None else {}), "runs": [item.run_dir.name for item in series], + **({"provenance": provenance_runs} if provenance_runs else {}), "points": [ { "environment_steps": float(environment_steps), @@ -904,6 +992,9 @@ def read_snapshot( runs = payload["runs"] if not points or not runs: raise RuntimeError(f"Performance snapshot is incomplete: {path}") + provenance_runs = payload.get("provenance") + if provenance_runs is not None and (not isinstance(provenance_runs, list) or len(provenance_runs) != len(runs)): + raise RuntimeError(f"Performance snapshot provenance does not match its runs: {path}") if all("environment_steps" in point for point in points): environment_steps = [point["environment_steps"] for point in points] elif all("step" in point for point in points): @@ -956,8 +1047,13 @@ def read_snapshot( environment_steps=np.asarray([], dtype=np.float64), elapsed_seconds=np.asarray([], dtype=np.float64), values=np.asarray([], dtype=np.float64), + provenance=( + {key: record[key] for key in ("system", "performance") if key in record} + if isinstance((record := provenance_runs[index]), dict) + else None + ), ) - for run in runs + for index, run in enumerate(runs) ] except (KeyError, TypeError, ValueError, json.JSONDecodeError) as exc: raise RuntimeError(f"Invalid performance snapshot: {path}") from exc @@ -1005,9 +1101,9 @@ def generate( if spec is not None and aggregate is not None: svg_path = STATIC_DIR / "images" / "performance" / f"{env_id}.svg" data_path = STATIC_DIR / "data" / "performance" / f"{env_id}.json" - if _update_text(svg_path, render_svg(spec, aggregate, len(series)), check=check): + if _update_text(svg_path, render_svg(spec, aggregate, len(series), annotation=_hardware_annotation(series)), check=check): stale.append(svg_path) - data_stale = _update_text(data_path, render_data(spec, aggregate, series), check=check) + data_stale = _update_text(data_path, render_data(spec, aggregate, series, data_path=data_path), check=check) if data_stale: stale.append(data_path) return stale diff --git a/docs/source/_static/data/performance/g1-wbt-backflip.json b/docs/source/_static/data/performance/g1-wbt-backflip.json new file mode 100644 index 00000000..b4facd73 --- /dev/null +++ b/docs/source/_static/data/performance/g1-wbt-backflip.json @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c003da366dcc1772fefdddf60caba49d0eee58390a2c44371a0ab3bbe8cd1c93 +size 74902 diff --git a/docs/source/_static/images/performance/g1-wbt-backflip.svg b/docs/source/_static/images/performance/g1-wbt-backflip.svg new file mode 100644 index 00000000..f633299c --- /dev/null +++ b/docs/source/_static/images/performance/g1-wbt-backflip.svg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:40fba37c78a845d25e383148ffc2e3bc8bfce2c00252c673b15bbed9a468ee3b +size 12543 diff --git a/docs/source/_static/images/poster/g1-wbt-backflip.jpg b/docs/source/_static/images/poster/g1-wbt-backflip.jpg new file mode 100644 index 00000000..970cbdab --- /dev/null +++ b/docs/source/_static/images/poster/g1-wbt-backflip.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b594105095cc9da1f9d76def8a41a808e250e13517e1f2b7ebecc3075ca3c921 +size 55227 diff --git a/docs/source/_static/videos/g1-wbt-backflip.mp4 b/docs/source/_static/videos/g1-wbt-backflip.mp4 new file mode 100644 index 00000000..cee98125 --- /dev/null +++ b/docs/source/_static/videos/g1-wbt-backflip.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8d28f390c17d4bd20d83e2252b5d4108727d075d450860629a7ebbef2acd53e8 +size 865945 diff --git a/docs/source/en/user_guide/envs/whole_body_tracking/adding_wbt_task.md b/docs/source/en/user_guide/envs/whole_body_tracking/adding_wbt_task.md index 3397eb2a..4a79a492 100644 --- a/docs/source/en/user_guide/envs/whole_body_tracking/adding_wbt_task.md +++ b/docs/source/en/user_guide/envs/whole_body_tracking/adding_wbt_task.md @@ -1,6 +1,6 @@ # Adding a WBT Training Task -The reusable unit of `ManagerEnv` is a complete `WbtManagerEnvCfg`. Adding a motion for an existing robot normally requires a motion +The reusable unit of `ManagerEnv` is a complete `WbtEnvCfg`. Adding a motion for an existing robot normally requires a motion file, an environment-config factory, Env registration, and matching Hydra Training Tasks. It does not require a copy of the environment implementation. This chapter adds G1 motion `dance1_subject1.npz` as Env ID `g1-wbt-dance1-subject1`. @@ -8,26 +8,21 @@ environment implementation. This chapter adds G1 motion `dance1_subject1.npz` as ## 1. Define the complete environment config Start from the existing WBT config subclass for the target robot. Edit -`motrix_envs/src/motrix_envs/locomotion/wbt/g1.py`: +`motrix_envs/src/motrix_envs/locomotion/wbt/g1/dance.py`: ```python -from pathlib import Path - from motrix_env_core import registry from motrix_env_core.manager import ManagerEnv -from motrix_envs.locomotion.wbt.g1 import G1WbtManagerCfg - - -_MOTION_DIR = Path(__file__).parent / "assets" / "motion" / "g1" +from motrix_envs.locomotion.wbt.g1.common import G1WbtEnvCfg, MOTION_DIR @registry.envcfg("g1-wbt-dance1-subject1") -def make_g129dof_wbt_dance1_subject1_cfg() -> G1WbtManagerCfg: - return G1WbtManagerCfg(motion_file=str(_MOTION_DIR / "dance1_subject1.npz")) +def make_g129dof_wbt_dance1_subject1_cfg() -> G1WbtEnvCfg: + return G1WbtEnvCfg(motion_file=str(MOTION_DIR / "dance1_subject1.npz")) ``` -`G1WbtManagerCfg` inherits `WbtManagerEnvCfg` and provides the G1 scene, tracked bodies, reference body, control scaling, +`G1WbtEnvCfg` inherits `WbtEnvCfg` and provides the G1 scene, tracked bodies, reference body, control scaling, rewards, and termination rules. When a new motion uses the same robot and tracking semantics, pass a different `motion_file` directly to the constructor. Do not copy `ManagerEnv` for each clip. @@ -77,7 +72,7 @@ effort as the largest absolute endpoint and uses it for position-target scaling. Existing robot config classes provide starting points: -- G1: `G1WbtManagerCfg(motion_file=...)` +- G1: `G1WbtEnvCfg(motion_file=...)` - Dex-EVT: `DexEvtWbtManagerCfg()` - K1: `K1WbtManagerCfg(commands=_k1_commands(...), rewards=...)` diff --git a/docs/source/en/user_guide/envs/whole_body_tracking/env_design.md b/docs/source/en/user_guide/envs/whole_body_tracking/env_design.md index 379749e3..327c5fc0 100644 --- a/docs/source/en/user_guide/envs/whole_body_tracking/env_design.md +++ b/docs/source/en/user_guide/envs/whole_body_tracking/env_design.md @@ -1,7 +1,7 @@ # Task Environment Design `ManagerEnv` advances one reference-motion frame per control step and compares the robot state with that frame's targets. -`WbtManagerEnvCfg` selects the motion, `tracked_body_names`, and `reference_body_name`; `scene.objs.robot` supplies the robot model, +`WbtEnvCfg` selects the motion, `tracked_body_names`, and `reference_body_name`; `scene.objs.robot` supplies the robot model, default key pose, base link, and actuators. Reference joint states are policy commands in the observation, but the action is still a position residual around the robot's default pose, and rewards primarily compare body poses and velocities. diff --git a/docs/source/en/user_guide/envs/whole_body_tracking/index.md b/docs/source/en/user_guide/envs/whole_body_tracking/index.md index 31bdddd8..9837b8af 100644 --- a/docs/source/en/user_guide/envs/whole_body_tracking/index.md +++ b/docs/source/en/user_guide/envs/whole_body_tracking/index.md @@ -3,14 +3,14 @@ `ManagerEnv` is MotrixLab's generic whole-body tracking (WBT) environment for humanoid robots. A policy follows a frame-by-frame reference motion under physics simulation while the task compares the global reference-body pose, relative poses of multiple body parts, body velocities, and joint feasibility. The `RobotCfg` and its assets own the robot model and -physical limits; `WbtManagerEnvCfg` selects the motion, tracked bodies, control scaling, rewards, and termination conditions. The +physical limits; `WbtEnvCfg` selects the motion, tracked bodies, control scaling, rewards, and termination conditions. The same environment implementation can therefore support different robots and motion clips. ## Demos -The following videos show Dex-EVT and Unitree G1 tracking dance motions, and Booster K1 tracking a free-kick motion. +The following videos show Dex-EVT and Unitree G1 tracking dance motions, Unitree G1 performing a backflip, and Booster K1 tracking a free-kick motion. -::::{grid} 1 1 2 3 +::::{grid} 1 1 2 2 :gutter: 2 2 2 2 :::{grid-item-card} Dex-EVT dance @@ -45,6 +45,22 @@ The following videos show Dex-EVT and Unitree G1 tracking dance motions, and Boo ::: +:::{grid-item-card} Unitree G1 backflip + +```{video} /_static/videos/g1-wbt-backflip.mp4 +:alt: 16 Unitree G1 humanoids performing backflips +:class: wbt-demo-video +:poster: /_static/images/poster/g1-wbt-backflip.jpg +:nocontrols: +:autoplay: +:playsinline: +:muted: +:loop: +:width: 100% +``` + +::: + :::{grid-item-card} Booster K1 free kick ```{video} /_static/videos/k1-wbt-freekick.mp4 @@ -84,11 +100,18 @@ thumbnail to open the full-size SVG. | Environment ID | Robot | Reference motion | Duration | Available training configs | Training curve | | --- | --- | --- | ---: | --- | --- | | `g1-29dof-wbt-largebox` | Unitree G1 29-DoF | `sub3_largebox_003.npz` | 6.50 s | `motrix.fastsac` | — | +| `g1-wbt-backflip` | Unitree G1 29-DoF | `backflip.npz` | 4.00 s | `motrix.fastsac` | | | `g1-wbt-dance` | Unitree G1 29-DoF | `dance1_subject2.npz` | 19.98 s | `motrix.fastsac` | | | `dex-evt-wbt-dance` | Dex-EVT | `dance1_easy.npz` | 39.72 s | `motrix.fastsac` | | | `k1-wbt-freekick` | Booster K1 | `freekick_shoot_arc_02.npz` | 2.50 s | `motrix.fastsac` | | ::: + + + Unitree G1 backflip WBT training curve +

Unitree G1 (g1-wbt-backflip) training curve

+ + Unitree G1 dance WBT training curve diff --git a/docs/source/en/user_guide/envs/whole_body_tracking/motion_format.md b/docs/source/en/user_guide/envs/whole_body_tracking/motion_format.md index 6c39ed0f..821b0248 100644 --- a/docs/source/en/user_guide/envs/whole_body_tracking/motion_format.md +++ b/docs/source/en/user_guide/envs/whole_body_tracking/motion_format.md @@ -52,7 +52,7 @@ Every quaternion in `body_quat_w` must be normalized. The loader's default norm | `clip_name` | Human-readable motion name | | `ext_*` | Extension arrays, exposed through `extensions` without the `ext_` prefix | -WBT training treats `WbtManagerEnvCfg.tracked_body_names`, `reference_body_name`, and the robot `base_link_name` as authoritative; +WBT training treats `WbtEnvCfg.tracked_body_names`, `reference_body_name`, and the robot `base_link_name` as authoritative; optional fields in the NPZ do not change task semantics. If `root_body_name` is absent, `replay.py` uses `body_names[0]`. ## Name binding and validation diff --git a/docs/source/zh_CN/user_guide/envs/whole_body_tracking/adding_wbt_task.md b/docs/source/zh_CN/user_guide/envs/whole_body_tracking/adding_wbt_task.md index ca37d889..594bcff2 100644 --- a/docs/source/zh_CN/user_guide/envs/whole_body_tracking/adding_wbt_task.md +++ b/docs/source/zh_CN/user_guide/envs/whole_body_tracking/adding_wbt_task.md @@ -1,32 +1,27 @@ # 新增 WBT 训练任务 -`ManagerEnv` 的复用单元是一份完整的 `WbtManagerEnvCfg`。为现有机器人增加一段动作时，通常只需新增 motion 文件、环境配置 +`ManagerEnv` 的复用单元是一份完整的 `WbtEnvCfg`。为现有机器人增加一段动作时，通常只需新增 motion 文件、环境配置 factory、Env 注册和对应的 Hydra Training Task，不需要复制环境实现。本章以 G1 的 `dance1_subject1.npz` 为例， 使用 Env ID `g1-wbt-dance1-subject1`。 ## 1. 定义完整环境配置 从目标机器人的 WBT 配置子类构造顶层配置。编辑 -`motrix_envs/src/motrix_envs/locomotion/wbt/g1.py`： +`motrix_envs/src/motrix_envs/locomotion/wbt/g1/dance.py`： ```python -from pathlib import Path - from motrix_env_core import registry from motrix_env_core.manager import ManagerEnv -from motrix_envs.locomotion.wbt.g1 import G1WbtManagerCfg - - -_MOTION_DIR = Path(__file__).parent / "assets" / "motion" / "g1" +from motrix_envs.locomotion.wbt.g1.common import G1WbtEnvCfg, MOTION_DIR @registry.envcfg("g1-wbt-dance1-subject1") -def make_g129dof_wbt_dance1_subject1_cfg() -> G1WbtManagerCfg: - return G1WbtManagerCfg(motion_file=str(_MOTION_DIR / "dance1_subject1.npz")) +def make_g129dof_wbt_dance1_subject1_cfg() -> G1WbtEnvCfg: + return G1WbtEnvCfg(motion_file=str(MOTION_DIR / "dance1_subject1.npz")) ``` -`G1WbtManagerCfg` 通过继承 `WbtManagerEnvCfg`，提供 G1 的机器人场景、tracked bodies、参考身体、控制缩放、奖励和 +`G1WbtEnvCfg` 通过继承 `WbtEnvCfg`，提供 G1 的机器人场景、tracked bodies、参考身体、控制缩放、奖励和 终止条件。新 motion 使用同一机器人和同一跟踪语义时，只需通过构造参数传入新的 `motion_file`。不要为每个 clip 复制一份 `ManagerEnv`。 @@ -74,7 +69,7 @@ range。 现有机器人配置类可作为起点： -- G1：`G1WbtManagerCfg(motion_file=...)` +- G1：`G1WbtEnvCfg(motion_file=...)` - Dex-EVT：`DexEvtWbtManagerCfg()` - K1：`K1WbtManagerCfg(commands=_k1_commands(...), rewards=...)` diff --git a/docs/source/zh_CN/user_guide/envs/whole_body_tracking/env_design.md b/docs/source/zh_CN/user_guide/envs/whole_body_tracking/env_design.md index 24c09d61..7b48079f 100644 --- a/docs/source/zh_CN/user_guide/envs/whole_body_tracking/env_design.md +++ b/docs/source/zh_CN/user_guide/envs/whole_body_tracking/env_design.md @@ -1,6 +1,6 @@ # 任务环境设计 -`ManagerEnv` 在每个控制周期推进一帧参考 motion，并将机器人状态与该帧目标进行比较。`WbtManagerEnvCfg` 指定 motion、 +`ManagerEnv` 在每个控制周期推进一帧参考 motion，并将机器人状态与该帧目标进行比较。`WbtEnvCfg` 指定 motion、 `tracked_body_names` 和 `reference_body_name`；`scene.objs.robot` 提供机器人模型、默认 key pose、基座 link 与 actuator。 参考关节状态是策略观察中的命令，但动作仍是相对于机器人默认姿态的位置残差，奖励主要根据 body 位姿和速度误差计算。 diff --git a/docs/source/zh_CN/user_guide/envs/whole_body_tracking/index.md b/docs/source/zh_CN/user_guide/envs/whole_body_tracking/index.md index 77db7939..2d361421 100644 --- a/docs/source/zh_CN/user_guide/envs/whole_body_tracking/index.md +++ b/docs/source/zh_CN/user_guide/envs/whole_body_tracking/index.md @@ -2,14 +2,14 @@ `ManagerEnv` 是 MotrixLab 面向人形机器人的通用全身动作跟踪（Whole-Body Tracking，WBT）环境。策略在物理仿真中 逐帧跟踪一段参考动作，任务同时约束参考身体的全局位姿、多个身体部位的相对位姿、身体速度和关节可行性。 -机器人模型与物理限制由 `RobotCfg` 及其资产提供；`WbtManagerEnvCfg` 选择 motion、跟踪身体、控制缩放、奖励和终止条件。 +机器人模型与物理限制由 `RobotCfg` 及其资产提供；`WbtEnvCfg` 选择 motion、跟踪身体、控制缩放、奖励和终止条件。 同一套环境实现因此可以支持不同机器人和不同动作片段。 ## 效果演示 -以下视频分别展示 Dex-EVT 和 Unitree G1 的舞蹈跟踪，以及 Booster K1 的任意球动作跟踪效果。 +以下视频分别展示 Dex-EVT 和 Unitree G1 的舞蹈跟踪、Unitree G1 的后空翻，以及 Booster K1 的任意球动作跟踪效果。 -::::{grid} 1 1 2 3 +::::{grid} 1 1 2 2 :gutter: 2 2 2 2 :::{grid-item-card} Dex-EVT 舞蹈 @@ -44,6 +44,22 @@ ::: +:::{grid-item-card} Unitree G1 后空翻 + +```{video} /_static/videos/g1-wbt-backflip.mp4 +:alt: 16 台 Unitree G1 人形机器人完成后空翻动作 +:class: wbt-demo-video +:poster: /_static/images/poster/g1-wbt-backflip.jpg +:nocontrols: +:autoplay: +:playsinline: +:muted: +:loop: +:width: 100% +``` + +::: + :::{grid-item-card} Booster K1 任意球 ```{video} /_static/videos/k1-wbt-freekick.mp4 @@ -82,11 +98,18 @@ adding_wbt_task | Env ID | 机器人 | 参考动作 | 时长 | 已提供的训练配置 | 训练曲线 | | --- | --- | --- | ---: | --- | --- | | `g1-29dof-wbt-largebox` | Unitree G1 29-DoF | `sub3_largebox_003.npz` | 6.50 s | `motrix.fastsac` | — | +| `g1-wbt-backflip` | Unitree G1 29-DoF | `backflip.npz` | 4.00 s | `motrix.fastsac` | | | `g1-wbt-dance` | Unitree G1 29-DoF | `dance1_subject2.npz` | 19.98 s | `motrix.fastsac` | | | `dex-evt-wbt-dance` | Dex-EVT | `dance1_easy.npz` | 39.72 s | `motrix.fastsac` | | | `k1-wbt-freekick` | Booster K1 | `freekick_shoot_arc_02.npz` | 2.50 s | `motrix.fastsac` | | ::: + + + Unitree G1 后空翻 WBT 训练曲线 +

Unitree G1（g1-wbt-backflip）训练曲线

+ + Unitree G1 舞蹈 WBT 训练曲线 diff --git a/docs/source/zh_CN/user_guide/envs/whole_body_tracking/motion_format.md b/docs/source/zh_CN/user_guide/envs/whole_body_tracking/motion_format.md index 2e6a2783..0b5a2ad5 100644 --- a/docs/source/zh_CN/user_guide/envs/whole_body_tracking/motion_format.md +++ b/docs/source/zh_CN/user_guide/envs/whole_body_tracking/motion_format.md @@ -51,7 +51,7 @@ $$ | `clip_name` | 人类可读的动作名称 | | `ext_*` | 扩展数组；loader 以去掉 `ext_` 的名称放入 `extensions` | -WBT 训练以 `WbtManagerEnvCfg.tracked_body_names`、`reference_body_name` 和机器人 `base_link_name` 为最终配置来源，不会 +WBT 训练以 `WbtEnvCfg.tracked_body_names`、`reference_body_name` 和机器人 `base_link_name` 为最终配置来源，不会 因为 NPZ 中存在同名可选字段而修改任务语义。`replay.py` 在 `root_body_name` 缺失时使用 `body_names[0]`。 ## 名称绑定与验证 diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/assets/motion/g1/backflip.npz b/motrix_envs/src/motrix_envs/locomotion/wbt/assets/motion/g1/backflip.npz new file mode 100644 index 00000000..64cfca2f --- /dev/null +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/assets/motion/g1/backflip.npz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5686c62541b8f2106cd5869f516c7f9fb801f59603b743a550473ad539127f1b +size 368092 diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/cfg.py b/motrix_envs/src/motrix_envs/locomotion/wbt/cfg.py index 6c5a5d8a..13825936 100644 --- a/motrix_envs/src/motrix_envs/locomotion/wbt/cfg.py +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/cfg.py @@ -142,7 +142,6 @@ class TerminationsCfg(ManagerTerminationsCfg): bad_dof_pos: BadDofPositionTerminationCfg = BadDofPositionTerminationCfg(threshold=0.5) bad_dof_vel: BadDofVelocityTerminationCfg = BadDofVelocityTerminationCfg(threshold=100.0) - @configclass class ObservationsCfg(ManagerObservationsCfg): """Typed observation groups for WBT.""" diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/g1/__init__.py b/motrix_envs/src/motrix_envs/locomotion/wbt/g1/__init__.py new file mode 100644 index 00000000..7bbc51c8 --- /dev/null +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/g1/__init__.py @@ -0,0 +1,12 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""Unitree G1 whole-body tracking tasks. + +Importing a task module registers its env config and env class; these +imports exist for that registration side effect. +""" + +from . import backflip, dance, largebox # noqa: F401 + +__all__ = ["backflip", "dance", "largebox"] diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/g1/backflip.py b/motrix_envs/src/motrix_envs/locomotion/wbt/g1/backflip.py new file mode 100644 index 00000000..511af235 --- /dev/null +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/g1/backflip.py @@ -0,0 +1,165 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""G1 backflip tracking task.""" + +from motrix_env_core import registry +from motrix_env_core.base import EnvCfg, SimCfg +from motrix_env_core.config import configclass +from motrix_env_core.config.scene import SystemCameraCfg +from motrix_env_core.manager import ManagerEnv +from motrix_env_core.mdp.rewards import ActionRateRewardCfg +from motrix_envs.config.scene import StandardSceneCfg, StandardSceneObjsCfg +from motrix_envs.locomotion.wbt.cfg import CommandsCfg, RewardsCfg, TerminationsCfg +from motrix_envs.locomotion.wbt.mdp.command import WbtMotionCommandCfg +from motrix_envs.locomotion.wbt.mdp.rewards import ( + EeBodyPosRewardCfg, + GlobalBodyAngularVelocityRewardCfg, + GlobalRefOrientationRewardCfg, + GlobalRefPositionRewardCfg, +) +from motrix_envs.locomotion.wbt.mdp.terminations import ( + BadBodyZTerminationCfg, + BadDofPositionTerminationCfg, + BadDofVelocityTerminationCfg, + BadRefOrientationTerminationCfg, + BadRefZTerminationCfg, +) +from motrix_envs.robot import UnitreeG129Dof + +from .common import G1WbtEnvCfg, MOTION_DIR + + +@configclass +class G1BackflipRewardsCfg(RewardsCfg): + """Backflip rewards: holosoma-aligned base plus a light action-rate term. + + zh_CN: 后空翻奖励：holosoma 对齐基础权重，外加轻量动作率项。 + + The light ``action_rate_l2`` keeps per-step action churn from + dominating while the position/orientation terms drive the flip. + """ + + # Deliberately light: heavier rates (or gating) suppress the rotation + # signal before the flip forms. + action_rate_l2: ActionRateRewardCfg = ActionRateRewardCfg(weight=-0.01) + + # Full-3D EE tracking: planar foot-placement error is diluted to 1/14 in + # the all-body mean, so landing accuracy needs its own gradient; the + # kernel sits comfortably inside the learning band at typical errors. + motion_ee_body_pos: EeBodyPosRewardCfg = EeBodyPosRewardCfg( + weight=3.0, + sigma=0.3, + body_names=( + "left_ankle_roll_link", + "right_ankle_roll_link", + "left_wrist_yaw_link", + "right_wrist_yaw_link", + ), + ) + + # Wide sigma on purpose: against the ~1.5-2 rad/s whole-body + # angular-velocity noise floor a tighter kernel saturates near zero and + # the rotation signal loses its gradient. + motion_global_body_ang_vel: GlobalBodyAngularVelocityRewardCfg = GlobalBodyAngularVelocityRewardCfg( + weight=1.0, + sigma=3.14, + ) + motion_global_ref_orientation_error_exp: GlobalRefOrientationRewardCfg = GlobalRefOrientationRewardCfg( + weight=1.0, + sigma=0.4, + ) + # Planar drift needs an unsaturated gradient: at meter-scale drift the + # stock sigma 0.3 kernel is vanishingly small. Widening to 1.0 keeps a + # learnable signal far out; near-field precision is handled by the + # relative-body terms. + motion_global_ref_position_error_exp: GlobalRefPositionRewardCfg = GlobalRefPositionRewardCfg( + weight=1.0, + sigma=1.0, + ) + + +@configclass(kw_only=True) +class G1BackflipWbtEnvCfg(G1WbtEnvCfg): + """Backflip tracking with the built-in G1 model. + + zh_CN: 后空翻跟踪：内置 G1 模型 + holosoma fastsac 配方（z 向跟踪终止）。 + Uses the stock menagerie-gain G1 unchanged — the flip is learnable with + the built-in actuator gains, whose large per-joint action scales + (scale = 0.25 * effort/kp) let the policy push targets far during the + violent launch. Termination is holosoma's ``BadTrackingZOnly`` shape: + under the command's per-step re-anchoring, world-z is the only anchor + axis with a physical reference (gravity), so ``bad_ref_z`` (0.5 m) is + the sole anchor-position check, backed by the tilt-only orientation + check and body-z/NaN guards. Rewards and reset are the holosoma + fastsac recipe unchanged (centered-noise reference-state teleport with + velocities, mixed frame sampling). + """ + + scene: StandardSceneCfg = StandardSceneCfg( + system_camera=SystemCameraCfg(distance=6.0, elevation=-20.0, azimuth=180.0), + objs=StandardSceneObjsCfg(robot=UnitreeG129Dof()), + ) + + # The launch impulse that creates angular momentum spans ~0.1 s of + # ground contact; a coarser timestep loses it to the contact solve and + # the flip stops being learnable, despite the higher collector + # throughput. + sim: SimCfg = SimCfg(dt=0.005, solver_iterations=3) + render_spacing = 1.5 + + commands: CommandsCfg = CommandsCfg( + motion=WbtMotionCommandCfg( + # Mixed sampling: 20% of episodes start at frame 0 (full skill + # from standing), 80% start uniformly anywhere — mid-air frames + # included. Frame-0-only start starves the flight window and the + # landing hold of episode-start coverage; adaptive + # failure-biased sampling is strictly harmful for this task. + # Mid-air resets are exact in velocity, so teleports into the + # flight window land on the reference ballistic trajectory; judge + # skill by frame-0-start episodes (play), since training-time + # means stay diluted by mid-clip starts. + adaptive_sampling_enabled=False, + start_at_timestep_zero_prob=0.2, + # Hold the final frame instead of wrap-rematerializing: an episode + # that reaches the clip end keeps tracking the landing hold until + # timeout, so "land and stay standing" is the terminal skill. + hold_at_clip_end=True, + ), + ) + + terminations: TerminationsCfg = TerminationsCfg( + # bad_ref_z is the launch-forcing term: standing through the flight + # window yields a pelvis-z error of ~0.44 against the 1.19 m apex, so + # standing is fatal and jumping is mandatory. Under the per-step + # re-anchoring semantics world-z is the only anchor axis with a + # physical reference (gravity), so it is also the only anchor-position + # check the preset enables. + bad_ref_z=BadRefZTerminationCfg(threshold=0.5), + bad_ref_ori=BadRefOrientationTerminationCfg(threshold=1.5), + bad_body_z=BadBodyZTerminationCfg( + threshold=0.5, + body_names=( + "left_ankle_roll_link", + "right_ankle_roll_link", + ), + ), + # Kept as NaN/divergence guards; they never fire on healthy rollouts. + bad_dof_pos=BadDofPositionTerminationCfg(threshold=0.5), + bad_dof_vel=BadDofVelocityTerminationCfg(threshold=100.0), + ) + + rewards: G1BackflipRewardsCfg = G1BackflipRewardsCfg() + + +@registry.envcfg("g1-wbt-backflip") +def make_g129dof_wbt_backflip_cfg() -> EnvCfg: + """Track the bundled G1 backflip reference motion. + + zh_CN: 让 Unitree G1 跟踪内置后空翻参考动作。 + """ + + return G1BackflipWbtEnvCfg(motion_file=str(MOTION_DIR / "backflip.npz")) + + +registry.env("g1-wbt-backflip")(ManagerEnv) diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/g1.py b/motrix_envs/src/motrix_envs/locomotion/wbt/g1/common.py similarity index 67% rename from motrix_envs/src/motrix_envs/locomotion/wbt/g1.py rename to motrix_envs/src/motrix_envs/locomotion/wbt/g1/common.py index 708d0c2f..f43c8c56 100644 --- a/motrix_envs/src/motrix_envs/locomotion/wbt/g1.py +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/g1/common.py @@ -1,30 +1,24 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 # SPDX-License-Identifier: Apache-2.0 -"""Unitree G1 whole-body tracking presets and registration.""" +"""Shared G1 whole-body tracking configuration.""" from dataclasses import InitVar from pathlib import Path -from motrix_env_core import registry -from motrix_env_core.base import EnvCfg, SimCfg +from motrix_env_core.base import SimCfg from motrix_env_core.config import configclass from motrix_env_core.config.scene import SystemCameraCfg -from motrix_env_core.manager import ManagerEnv from motrix_env_core.mdp.rewards import ActionRateRewardCfg from motrix_env_core.sim import BodyLinkNetContactForceQuery from motrix_envs.config.scene import StandardSceneCfg, StandardSceneObjsCfg from motrix_envs.locomotion.wbt.cfg import CommandsCfg, RewardsCfg, TerminationsCfg, WbtEnvCfg -from motrix_envs.locomotion.wbt.mdp.command import ( - WbtMotionCommandCfg, -) -from motrix_envs.locomotion.wbt.mdp.terminations import ( - BadBodyZTerminationCfg, -) +from motrix_envs.locomotion.wbt.mdp.command import WbtMotionCommandCfg +from motrix_envs.locomotion.wbt.mdp.terminations import BadBodyZTerminationCfg from motrix_envs.robot import UnitreeG129Dof -_MOTION_DIR = Path(__file__).parent / "assets" / "motion" / "g1" -_G1_TRACKED_BODY_NAMES = ( +MOTION_DIR = Path(__file__).resolve().parent.parent / "assets" / "motion" / "g1" +G1_TRACKED_BODY_NAMES = ( "pelvis", "left_hip_roll_link", "left_knee_link", @@ -70,7 +64,7 @@ def __post_init__(self, motion_file: str | None) -> None: super().__post_init__() if motion_file is not None: self.commands.motion.motion_file = motion_file - self._set_tracked_body_names(_G1_TRACKED_BODY_NAMES) + self._set_tracked_body_names(G1_TRACKED_BODY_NAMES) self.commands.motion.reference_body_name = "torso_link" self.queries.data["undesired_contact_forces"] = BodyLinkNetContactForceQuery( @@ -82,27 +76,3 @@ def __post_init__(self, motion_file: str | None) -> None: "right_ankle_roll_link", ), ) - - -@registry.envcfg("g1-29dof-wbt-largebox") -def make_g129dof_wbt_largebox_cfg() -> EnvCfg: - """Track a large-box carrying reference motion with Unitree G1. - - zh_CN: 让 Unitree G1 跟踪搬运大箱子的参考动作。 - """ - - return G1WbtEnvCfg(motion_file=str(_MOTION_DIR / "sub3_largebox_003.npz")) - - -@registry.envcfg("g1-wbt-dance") -def make_g129dof_wbt_dance_cfg() -> EnvCfg: - """Track the bundled G1 dance motion with the manager-based environment. - - zh_CN: 让 Unitree G1 跟踪内置舞蹈参考动作。 - """ - - return G1WbtEnvCfg(motion_file=str(_MOTION_DIR / "dance1_subject2.npz")) - - -registry.env("g1-29dof-wbt-largebox")(ManagerEnv) -registry.env("g1-wbt-dance")(ManagerEnv) diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/g1/dance.py b/motrix_envs/src/motrix_envs/locomotion/wbt/g1/dance.py new file mode 100644 index 00000000..a2f860b0 --- /dev/null +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/g1/dance.py @@ -0,0 +1,23 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""G1 dance tracking task.""" + +from motrix_env_core import registry +from motrix_env_core.base import EnvCfg +from motrix_env_core.manager import ManagerEnv + +from .common import G1WbtEnvCfg, MOTION_DIR + + +@registry.envcfg("g1-wbt-dance") +def make_g129dof_wbt_dance_cfg() -> EnvCfg: + """Track the bundled G1 dance motion with the manager-based environment. + + zh_CN: 让 Unitree G1 跟踪内置舞蹈参考动作。 + """ + + return G1WbtEnvCfg(motion_file=str(MOTION_DIR / "dance1_subject2.npz")) + + +registry.env("g1-wbt-dance")(ManagerEnv) diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/g1/largebox.py b/motrix_envs/src/motrix_envs/locomotion/wbt/g1/largebox.py new file mode 100644 index 00000000..8b9b232b --- /dev/null +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/g1/largebox.py @@ -0,0 +1,23 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""G1 large-box carrying tracking task.""" + +from motrix_env_core import registry +from motrix_env_core.base import EnvCfg +from motrix_env_core.manager import ManagerEnv + +from .common import G1WbtEnvCfg, MOTION_DIR + + +@registry.envcfg("g1-29dof-wbt-largebox") +def make_g129dof_wbt_largebox_cfg() -> EnvCfg: + """Track a large-box carrying reference motion with Unitree G1. + + zh_CN: 让 Unitree G1 跟踪搬运大箱子的参考动作。 + """ + + return G1WbtEnvCfg(motion_file=str(MOTION_DIR / "sub3_largebox_003.npz")) + + +registry.env("g1-29dof-wbt-largebox")(ManagerEnv) diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/command.py b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/command.py index 68320aec..ee54f95e 100644 --- a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/command.py +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/command.py @@ -29,7 +29,12 @@ @njit(inline="always") -def _sample_motion_step(rand, sampling_cdf, num_frames: np.int64, start_at_timestep_zero_prob: np.float32): +def _sample_motion_step( + rand, + sampling_cdf, + num_frames: np.int64, + start_at_timestep_zero_prob: np.float32, +): """Draw one start frame from the adaptive-bin CDF (or uniformly when disabled).""" unit = (rand.next_uniform() + np.float32(1.0)) * np.float32(0.5) if sampling_cdf.size == 0: @@ -203,7 +208,10 @@ def reset_env(self, ctx: ManagerContext) -> None: """Sample the starting frame for one reset environment lane.""" num_frames = self.clip.joint_pos.shape[0] self.steps[0] = _sample_motion_step( - ctx.rand, self.sampling_cdf, np.int64(num_frames), self.start_at_timestep_zero_prob + ctx.rand, + self.sampling_cdf, + np.int64(num_frames), + self.start_at_timestep_zero_prob, ) # ``clip_ended`` is intentionally left untouched: advance recomputes it # every transition, so a lane that just wrapped keeps its flag for the @@ -230,7 +238,10 @@ def advance(self, ctx: ManagerContext) -> None: # resample keeps steps valid and consistently distributed between # this kernel and the reset pipeline. self.steps[0] = _sample_motion_step( - ctx.rand, self.sampling_cdf, np.int64(num_frames), self.start_at_timestep_zero_prob + ctx.rand, + self.sampling_cdf, + np.int64(num_frames), + self.start_at_timestep_zero_prob, ) ctx.sim_reset_requested[0] = True diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/reset.py b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/reset.py index 4bb4d214..13321c7a 100644 --- a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/reset.py +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/reset.py @@ -40,6 +40,9 @@ def _reset_body_pos( motion: WbtMotionCommand = ctx.commands["motion"] position[0] = motion.clip.root_body_pos_w[motion.steps[0]] for index in range(3): + # next_uniform() returns [-1, 1), so this is centered uniform noise + # in [-scale, +scale) per axis. (An earlier "centering fix" subtracted + # 0.5 from an already-centered sample, producing a [-3s, +1s) bias.) position[0, index] += ctx.rand.next_uniform() * noise_scale[index] @@ -110,6 +113,7 @@ def _reset_body_lin_vel( motion: WbtMotionCommand = ctx.commands["motion"] linear_velocity[0] = motion.clip.root_body_lin_vel_w[motion.steps[0]] for index in range(3): + # Centered noise: see _reset_body_pos. linear_velocity[0, index] += ctx.rand.next_uniform() * noise_scale[index] @@ -140,6 +144,7 @@ def _reset_body_rot_vel( motion: WbtMotionCommand = ctx.commands["motion"] angular_velocity[0] = motion.clip.root_body_ang_vel_w[motion.steps[0]] for index in range(3): + # Centered noise: see _reset_body_pos. angular_velocity[0, index] += ctx.rand.next_uniform() * noise_scale[index] diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/rewards.py b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/rewards.py index c720f716..642db966 100644 --- a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/rewards.py +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/rewards.py @@ -73,6 +73,43 @@ def __call__(self, ctx) -> RewardTerm: return RewardTerm(relative_body_position_reward, np.float32(self.sigma)) +@dispatch +def ee_body_pos_reward(ctx: ManagerContext, body_indices: tuple[int, ...], sigma: np.float32) -> float: + tracked_body_pos = ctx.sim["tracked_body_pos"] + motion: WbtMotionCommand = ctx.commands["motion"] + error_sq = 0.0 + for body_id in body_indices: + diff = motion.target_body_position_relative[body_id] - tracked_body_pos[body_id] + error_sq += float(np.dot(diff, diff)) + return math.exp(-(error_sq / len(body_indices)) / (sigma * sigma)) + + +@configclass(kw_only=True) +class EeBodyPosRewardCfg(RewardTermCfg): + """Full-3D end-effector position tracking (ankles, wrists). + + zh_CN: 末端 body（踝、腕）的三维位置专项跟踪奖励。 + + Foot placement error lives mostly in the horizontal plane, and each foot + is 1/14 of the all-body relative-position mean — too diluted to shape + landing accuracy. This term tracks the full 3D end-effector error + directly, including the height component. + """ + + body_names: tuple[str, ...] = () + sigma: float + + def __call__(self, ctx) -> RewardTerm: + if not self.body_names: + raise ValueError( + "EeBodyPosRewardCfg requires a non-empty body_names: the kernel " + "divides the squared tracking error by the number of bodies." + ) + tracked_body_names = ctx.cfg.commands.motion.tracked_body_names + body_indices = tuple(tracked_body_names.index(name) for name in self.body_names) + return RewardTerm(ee_body_pos_reward, body_indices, np.float32(self.sigma)) + + @dispatch def relative_body_orientation_reward(ctx: ManagerContext, sigma: np.float32) -> float: tracked_body_quat = ctx.sim["tracked_body_quat"] diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/terminations.py b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/terminations.py index 7622e658..656e5d5b 100644 --- a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/terminations.py +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/terminations.py @@ -63,6 +63,50 @@ def __call__(self, ctx) -> TerminationTerm: ) +@dispatch +def bad_motion_body_position_termination( + ctx: ManagerContext, + body_indices: tuple[int, ...], + threshold: np.float32, +) -> bool: + """Holosoma-aligned tracked-body tracking termination. + + Terminates when the mean position error of the configured tracked bodies + (any subset of ``tracked_body_names``, not just end effectors) against + their reference targets (the yaw/height-anchored body targets the + relative rewards track) exceeds ``threshold``. + """ + tracked_body_pos = ctx.sim["tracked_body_pos"] + motion: WbtMotionCommand = ctx.commands["motion"] + count = len(body_indices) + if count == 0: + ctx.metrics["motion_body_pos_err"][0] = 0.0 + return False + error_sq_sum = 0.0 + for index in range(count): + body_id = body_indices[index] + diff = motion.target_body_position_relative[body_id] - tracked_body_pos[body_id] + error_sq_sum += float(np.dot(diff, diff)) + mean_err = math.sqrt(error_sq_sum / count) + ctx.metrics["motion_body_pos_err"][0] = mean_err + return mean_err > threshold + + +@configclass(kw_only=True) +class BadMotionBodyPositionTerminationCfg(_WbtTerminationCfg): + body_names: tuple[str, ...] = () + + def __call__(self, ctx) -> TerminationTerm: + tracked_body_names = ctx.cfg.commands.motion.tracked_body_names + body_indices = tuple(tracked_body_names.index(name) for name in self.body_names) + return TerminationTerm( + bad_motion_body_position_termination, + body_indices, + np.float32(self.threshold), + metric_names=("motion_body_pos_err",), + ) + + @dispatch def bad_body_z_termination( ctx: ManagerContext, diff --git a/motrix_envs/tests/test_package_boundary.py b/motrix_envs/tests/test_package_boundary.py index 3f2cf55b..47391598 100644 --- a/motrix_envs/tests/test_package_boundary.py +++ b/motrix_envs/tests/test_package_boundary.py @@ -10,7 +10,8 @@ from motrix_env_core.base import EnvCfg from motrix_env_core.config.scene import SceneCfg from motrix_envs.locomotion.wbt.dex_evt import DexEvtWbtEnvCfg -from motrix_envs.locomotion.wbt.g1 import G1WbtEnvCfg, make_g129dof_wbt_dance_cfg +from motrix_envs.locomotion.wbt.g1.common import G1WbtEnvCfg +from motrix_envs.locomotion.wbt.g1.dance import make_g129dof_wbt_dance_cfg from motrix_envs.locomotion.wbt.k1 import K1WbtEnvCfg diff --git a/motrix_envs/tests/test_wbt_numba.py b/motrix_envs/tests/test_wbt_numba.py index 35ce76a3..1fe2a3aa 100644 --- a/motrix_envs/tests/test_wbt_numba.py +++ b/motrix_envs/tests/test_wbt_numba.py @@ -29,7 +29,7 @@ WbtEnvCfg, ) from motrix_envs.locomotion.wbt.dex_evt import DexEvtWbtEnvCfg # noqa: E402 -from motrix_envs.locomotion.wbt.g1 import G1WbtEnvCfg # noqa: E402 +from motrix_envs.locomotion.wbt.g1.common import G1WbtEnvCfg # noqa: E402 from motrix_envs.locomotion.wbt.k1 import K1WbtEnvCfg # noqa: E402 from motrix_envs.locomotion.wbt.mdp.action import ( # noqa: E402 WbtJointPositionAction, diff --git a/motrix_rl/src/motrix_rl/fastsac/async_impl/topology.py b/motrix_rl/src/motrix_rl/fastsac/async_impl/topology.py index 06576387..74e9a60d 100644 --- a/motrix_rl/src/motrix_rl/fastsac/async_impl/topology.py +++ b/motrix_rl/src/motrix_rl/fastsac/async_impl/topology.py @@ -123,6 +123,61 @@ def _learner_node(device: torch.device, multi_node: bool) -> int | None: return numa.gpu_numa_node(index) +def resolve_collector_inference_devices( + collector_inference_device: str, + learner_devices: list[torch.device], + num_collectors: int, + num_learners: int, + default_device: torch.device, +) -> list[torch.device]: + """One explicit inference device per collector, co-located with its owner. + + The device layout is owned here rather than split between the trainer and + the workers: a generic ``cuda`` spec resolves to the owning learner's GPU + for EVERY learner count (a single learner included), so the collector/learner + pair shares one GPU and the CUDA-IPC ring/weight transports stay the + default. Explicit ``cpu`` or ``cuda:N`` specs pass through. The returned + devices are always explicit (indexed for CUDA), so the collector workers + never re-resolve a generic spec in-process against + ``torch.cuda.current_device()`` — that fallback silently crossed GPUs when + the current device drifted from the learner's. + + Pure device arithmetic (like :func:`same_cuda_device`): no CUDA context is + created, so the pre-spawn parent and CUDA-less unit tests can call it + freely. Runtime availability validation stays with the worker-side + :func:`motrix_rl.fastsac.async_impl.collector.resolve_collector_inference_device`. + """ + device = torch.device(collector_inference_device) + if device.type not in ("cpu", "cuda"): + raise ValueError( + f"collector_inference_device must be cpu or cuda, got '{collector_inference_device}'" + ) + if ( + device.type == "cuda" + and device.index is not None + and torch.cuda.is_available() + and device.index >= torch.cuda.device_count() + ): + raise RuntimeError( + f"collector_inference_device='{collector_inference_device}' selects CUDA device " + f"{device.index}, but only {torch.cuda.device_count()} device(s) are available" + ) + per_learner = num_collectors // num_learners + + def _owner(rank: int) -> torch.device: + return learner_devices[rank] if learner_devices else default_device + + devices = [] + for i in range(num_collectors): + owner = _owner(i // per_learner) + if device.type == "cuda" and device.index is None and owner.type == "cuda": + owner_index = owner.index if owner.index is not None else 0 + devices.append(torch.device("cuda", owner_index)) + else: + devices.append(device) + return devices + + def resolve_learner_devices( learner_device_specs: list[str] | None, num_learners: int, @@ -204,8 +259,8 @@ def resolve_trainer_topology( num_envs: int, num_collectors: int, num_learners: int, - learner_devices: list[torch.device], - collector_devices: list[torch.device], + learner_device_specs: list[str] | None, + collector_inference_device: str, default_device: torch.device, async_options: FastSacAsyncOptionsCfg, actor_param_numel: int, @@ -213,24 +268,30 @@ def resolve_trainer_topology( ) -> TrainerTopology: """Single computation API: derive the full trainer topology in one pass. - Combines env sharding (:func:`split_num_envs`), NUMA placement, CPU - bindings, and the transport decisions for transition rings - (:func:`ring_transport_is_ipc`) and weight channels - (:func:`use_ipc_weight_channel`) into one :class:`TrainerTopology`. + Takes the RAW device specs and owns the whole layout: learner devices + (:func:`resolve_learner_devices`), collector inference devices with + owner co-location (:func:`resolve_collector_inference_devices`), env + sharding (:func:`split_num_envs`), NUMA placement, CPU bindings, and the + transport decisions for transition rings (:func:`ring_transport_is_ipc`) + and weight channels (:func:`use_ipc_weight_channel`). Callers read the + complete layout — which GPU each worker runs on, which transport each + ring uses — from the returned :class:`TrainerTopology` and never resolve + a device themselves. ``learner_device_specs`` is the configured list + (``None`` replicates ``default_device`` per rank); an index-less + ``cuda`` ``default_device`` means device 0. ``actor_param_numel`` is the parent-computed actor parameter count that sizes the weight-transport threshold; ``cpus_per_collector`` optionally - chunks each binding base into per-collector slices. ``learner_devices`` - carries one indexed device per rank (empty for a single learner, whose - device comes from ``default_device`` — an index-less ``cuda`` means - device 0). ``collector_devices[i]`` is the inference device of collector - ``i``. + chunks each binding base into per-collector slices. """ if num_learners < 1 or num_collectors < 1: raise ValueError(f"invalid worker counts: {num_collectors=} {num_learners=}") if num_collectors % num_learners != 0: raise ValueError(f"num_collectors={num_collectors} must divide evenly across num_learners={num_learners}") - if len(collector_devices) != num_collectors: - raise ValueError(f"expected {num_collectors} collector devices, got {len(collector_devices)}") + + learner_devices = resolve_learner_devices(learner_device_specs, num_learners, default_device) + collector_devices = resolve_collector_inference_devices( + collector_inference_device, learner_devices, num_collectors, num_learners, default_device + ) env_shards = split_num_envs(num_envs, num_collectors) multi_node = len(numa.available_numa_nodes()) >= 2 diff --git a/motrix_rl/src/motrix_rl/fastsac/async_impl/train.py b/motrix_rl/src/motrix_rl/fastsac/async_impl/train.py index 408550f8..5de8d0a9 100644 --- a/motrix_rl/src/motrix_rl/fastsac/async_impl/train.py +++ b/motrix_rl/src/motrix_rl/fastsac/async_impl/train.py @@ -30,11 +30,12 @@ from motrix_env_core.renderer import RenderConfig from motrix_rl.console import TrainingPanelStats, emit_training_panel, open_training_live from motrix_rl.fastsac.agent import FastSacAgent -from motrix_rl.fastsac.async_impl.collector import resolve_collector_inference_device from motrix_rl.fastsac.async_impl.numa import spawn_placement from motrix_rl.fastsac.async_impl.panels import BootPanel from motrix_rl.fastsac.async_impl.stats import aggregate_collector_stats, nest_timing_path, timing_mean -from motrix_rl.fastsac.async_impl.topology import resolve_learner_devices, resolve_trainer_topology +from motrix_rl.fastsac.async_impl.topology import resolve_trainer_topology +from motrix_rl.runs import annotate_run +from motrix_rl.system_metrics import capture_system_info from motrix_rl.fastsac.async_impl.transport import Control, RingCursors, SharedTransitionRing from motrix_rl.fastsac.async_impl.transport.handshake import StartupHandshake from motrix_rl.fastsac.async_impl.transport.weight_channel import WeightChannelShared @@ -133,6 +134,7 @@ def train(self) -> None: async_options = cfg.trainer.async_options logging_interval = self._context.logging.interval save_interval = self._context.checkpoint.interval + train_started_at = time.perf_counter() if self._context.logging.backend != "tensorboard": raise ValueError("FastSAC supports only the 'tensorboard' logging backend.") @@ -145,7 +147,6 @@ def train(self) -> None: dims = (obs_dim, critic_obs_dim, act_dim) learner_device = self._device() - collector_device = resolve_collector_inference_device(async_options.collector_inference_device) num_collectors = async_options.num_collectors num_learners = async_options.num_learners @@ -161,23 +162,8 @@ def train(self) -> None: # its GPU's PCIe-local node and every collector follows its owning # learner, so a collector/learner pair never straddles a NUMA node. # Falls back to the OS default placement on single-node hosts. - learner_devices = resolve_learner_devices(async_options.learner_devices, num_learners, self._device()) cpus_per_collector = async_options.cpus_per_collector num_envs = self._context.num_envs - # Generic "cuda" collector inference resolves to the owning learner's - # GPU (num_collectors // num_learners collectors per learner); explicit - # specs pass through. The single-learner path keeps the in-process - # resolution (collector_device=None) byte-identical. - collector_device_specs = None - if num_learners > 1: - per_learner = num_collectors // num_learners - spec = async_options.collector_inference_device - collector_device_specs = [ - spec - if not (spec == "cuda" and learner_devices[i // per_learner].type == "cuda") - else f"cuda:{learner_devices[i // per_learner].index}" - for i in range(num_collectors) - ] # Shared-memory primitives allocated in the parent, inherited by children. # One SPSC ring + one weight channel per collector: every shared quantity # keeps exactly one producer and one consumer, so the lock-free @@ -192,28 +178,33 @@ def train(self) -> None: # Collectors are co-located with their owning learner's GPU by # default (collector_inference_device="cuda" resolves per owner), so # the IPC path is the default whenever both sides share that GPU. - if collector_device_specs is not None: - collector_devices = [torch.device(spec) for spec in collector_device_specs] - else: - collector_devices = [collector_device] * num_collectors - # One resolution pass derives the whole compute layout: env shards, - # NUMA placement, and per-collector transports (rings + weights). + # One resolution pass derives the whole compute layout — which GPU + # each worker runs on, env shards, NUMA placement, per-collector + # transports (rings + weights) — from the raw device specs; the + # workers read their explicit devices from the topology and never + # re-resolve a generic spec in-process. param_numel = actor_param_numel(cfg, dims, action_scale, action_bias) topology = resolve_trainer_topology( num_envs, num_collectors, num_learners, - learner_devices, - collector_devices, + async_options.learner_devices, + async_options.collector_inference_device, self._device(), async_options, param_numel, cpus_per_collector=cpus_per_collector, ) + learner_devices = [learner.device for learner in topology.learners] + collector_devices = [collector.device for collector in topology.collectors] numa_nodes = [collector.numa_node for collector in topology.collectors] learner_numa_nodes = [learner.numa_node for learner in topology.learners] env_shards = topology.env_shards ring_ipc = [collector.ring_ipc for collector in topology.collectors] + # Provenance for performance snapshots: record the devices this run + # actually uses (topology is the single source of truth), not the + # host inventory. + annotate_run(self._context.run_dir, system=capture_system_info(topology=topology)) rings: list[SharedTransitionRing | RingCursors] = [ RingCursors() if ipc @@ -274,11 +265,13 @@ def _drain_child_errors() -> list[tuple[str, str]]: print(traceback_text.rstrip()) return errors + unique_collector_devices = list(dict.fromkeys(str(d) for d in collector_devices)) + collector_banner = unique_collector_devices[0] if len(unique_collector_devices) == 1 else unique_collector_devices print( f"[motrix.fastsac async] collector/learner training '{self._env_name}' learner={learner_device} " f"learner_replicas={num_learners} " + (f"learner_devices={[str(d) for d in learner_devices]} " if num_learners > 1 else "") - + f"collector_env=cpu collector_inference={collector_device} num_collectors={num_collectors} " + + f"collector_env=cpu collector_inference={collector_banner} num_collectors={num_collectors} " f"numa_nodes={numa_nodes} learner_numa_nodes={learner_numa_nodes} " f"num_envs={num_envs} iters={num_iterations} " f"from={resume_step} utd_mode={async_options.utd_mode}" @@ -324,7 +317,7 @@ def _drain_child_errors() -> list[tuple[str, str]]: "panel_queue": panel_queue, "num_learners": num_learners, "rendezvous_file": rendezvous_file, - "learner_device": str(learner_devices[rank]) if learner_devices else None, + "learner_device": str(topology.learners[rank].device), "all_rings": rings, "weight_ipc": [ c.weight_ipc for c in topology.collectors[rank * per_learner : (rank + 1) * per_learner] @@ -358,7 +351,7 @@ def _drain_child_errors() -> list[tuple[str, str]]: "collector_id": i, "numa_node": numa_nodes[i], "cpus": topology.collectors[i].cpus, - "collector_device": collector_device_specs[i] if collector_device_specs is not None else None, + "collector_device": str(topology.collectors[i].device), "run_dir": str(self._context.run_dir), "handshake": handshake, }, @@ -760,6 +753,17 @@ def _shutdown(processes, grace_s): live.stop() except Exception: pass + wall_time_s = time.perf_counter() - train_started_at + annotate_run( + self._context.run_dir, + performance={ + "num_envs": num_envs, + "iterations": num_iterations, + "total_env_steps": int(control.global_step) * num_envs, + "wall_time_s": round(wall_time_s, 1), + "mean_env_steps_per_s": int(int(control.global_step) * num_envs / max(wall_time_s, 1e-9)), + }, + ) if tb_writer is not None: try: tb_writer.close() diff --git a/motrix_rl/src/motrix_rl/fastsac/async_impl/worker.py b/motrix_rl/src/motrix_rl/fastsac/async_impl/worker.py index eb68910c..7183caa6 100644 --- a/motrix_rl/src/motrix_rl/fastsac/async_impl/worker.py +++ b/motrix_rl/src/motrix_rl/fastsac/async_impl/worker.py @@ -344,8 +344,9 @@ def run_collector_process( apply_binding(role, numa_node, cpus or []) set_seed(seed) opts = cfg.trainer.async_options - # Multi-learner: the parent resolves the generic "cuda" spec to the - # owning learner's GPU; an explicit spec passes through unchanged. + # The parent's topology pass resolved every collector's inference + # device explicitly (generic "cuda" co-located with the owning + # learner's GPU), so no generic spec is re-resolved in-process here. if collector_device is not None: opts.collector_inference_device = collector_device _pin_worker_cpus(_resolve_cpu_set(opts.collector_cpu_cores, "collector_cpu_cores")) diff --git a/motrix_rl/src/motrix_rl/runs.py b/motrix_rl/src/motrix_rl/runs.py index f0e74066..a221f8b7 100644 --- a/motrix_rl/src/motrix_rl/runs.py +++ b/motrix_rl/src/motrix_rl/runs.py @@ -28,6 +28,8 @@ class RunMetadata: checkpoint_format: str sim: str | None = None motrixlab_version: str | None = None + system: dict | None = None + performance: dict | None = None @dataclass(frozen=True) @@ -126,6 +128,22 @@ def write_metadata(run_dir: str | Path, metadata: RunMetadata) -> Path: return metadata_path +def annotate_run(run_dir: str | Path, **fields: Any) -> Path: + """Merge annotation fields into an existing run's metadata.json. + + Used by trainers to attach facts that only become known after run + creation (resolved device layout, end-of-training performance). The + merge is a raw top-level dict update so future annotation keys do not + require a RunMetadata field; ``read_metadata`` ignores unknown keys. + """ + run_dir = Path(run_dir) + metadata_path = run_dir / METADATA_FILENAME + data = json.loads(metadata_path.read_text(encoding="utf-8")) + data.update(fields) + metadata_path.write_text(json.dumps(data, indent=2, sort_keys=True) + "\n", encoding="utf-8") + return metadata_path + + def read_metadata(run_dir: str | Path) -> RunMetadata: data = json.loads((Path(run_dir) / METADATA_FILENAME).read_text(encoding="utf-8")) # Metadata written by older versions may carry keys the current schema no diff --git a/motrix_rl/src/motrix_rl/system_metrics.py b/motrix_rl/src/motrix_rl/system_metrics.py index c1d17c8e..47a02b56 100644 --- a/motrix_rl/src/motrix_rl/system_metrics.py +++ b/motrix_rl/src/motrix_rl/system_metrics.py @@ -1,7 +1,9 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 # SPDX-License-Identifier: Apache-2.0 -"""Low-overhead host metrics sampled at training-panel refresh boundaries. +"""Host metrics and training-run hardware provenance. + +Low-overhead host metrics sampled at training-panel refresh boundaries. CPU samplers read Linux ``/proc`` interfaces and return ``None`` where they are unavailable, so panels degrade to ``n/a`` fields; memory sampling also @@ -13,7 +15,10 @@ import ctypes import os +import platform +import re import sys +from importlib import metadata as importlib_metadata from dataclasses import dataclass from pathlib import Path from typing import Any @@ -433,3 +438,115 @@ def _sysfs_gpu_busy_percent() -> list[float]: if 0.0 <= value <= 100.0: values.append(value) return values + + +# ---------------------------------------------------------------- provenance + + +def capture_system_info(topology: Any = None, fallback_device: Any = None) -> dict[str, Any]: + """Capture the system context of a training run. + + Records what a run actually used — the devices in the resolved trainer + topology with their roles — rather than what the host machine has, so + performance snapshots stay comparable across machines and layouts. + ``topology`` is the resolved async trainer layout (single source of + truth for which devices each worker uses); ``fallback_device`` covers + trainers without a topology, recording that single device. + """ + + roles: dict[int, set[str]] = {} + if topology is not None: + for learner in topology.learners: + _add_device_role(roles, learner.device, "learner") + for collector in topology.collectors: + _add_device_role(roles, collector.device, "collector-inference") + elif fallback_device is not None: + _add_device_role(roles, fallback_device, "trainer") + + gpus_used = [] + for index in sorted(roles): + gpus_used.append({"index": index, "model": _cuda_device_name(index), "roles": sorted(roles[index])}) + + system: dict[str, Any] = { + "platform": platform.platform(), + "cpu": { + "model": _cpu_model(), + "machine_cores": os.cpu_count() or 0, + }, + "gpus_used": gpus_used, + "software": _software_versions(), + } + if topology is not None: + system["cpu"]["bound_cores"] = _bound_cores(topology) + return system + + +def hardware_profile_slug(system: dict[str, Any]) -> str: + """Stable slug identifying the used-hardware profile, e.g. ``epyc-9004-rtx-4090-x1``.""" + + cpu = _slugify(str(system.get("cpu", {}).get("model", "cpu"))) + gpus = system.get("gpus_used") or [] + models = sorted(str(gpu.get("model", "gpu")).lower() for gpu in gpus) + if not models: + return f"{cpu}-cpuonly" + gpu_part = _slugify(models[0]) if len(set(models)) == 1 else _slugify("-".join(models)) + return f"{cpu}-{gpu_part}-x{len(models)}" + + +def _add_device_role(roles: dict[int, set[str]], device: Any, role: str) -> None: + if getattr(device, "type", None) != "cuda": + return + index = device.index if device.index is not None else 0 + roles.setdefault(index, set()).add(role) + + +def _cuda_device_name(index: int) -> str: + try: + import torch + + if torch.cuda.is_available(): + return torch.cuda.get_device_name(index) + except Exception: + pass + return f"cuda:{index}" + + +def _cpu_model() -> str: + try: + for line in Path("/proc/cpuinfo").read_text(encoding="utf-8").splitlines(): + if line.startswith("model name"): + return line.split(":", 1)[1].strip() + except OSError: + pass + return platform.processor() or "unknown" + + +def _bound_cores(topology: Any) -> dict[str, list[int]] | None: + bound: dict[str, list[int]] = {} + for rank, learner in enumerate(topology.learners): + if learner.cpus: + bound[f"learner{rank if len(topology.learners) > 1 else ''}"] = sorted(learner.cpus) + for collector in topology.collectors: + if collector.cpus: + bound[f"collector{collector.collector_id}"] = sorted(collector.cpus) + return bound or None + + +def _software_versions() -> dict[str, str]: + def _pkg(name: str) -> str | None: + try: + return importlib_metadata.version(name) + except importlib_metadata.PackageNotFoundError: + return None + + software = {"python": sys.version.split()[0]} + for package in ("torch", "motrixlab", "motrixsim", "motrix-env-core"): + resolved = _pkg(package) + if resolved is not None: + software[package] = resolved + return software + + +def _slugify(text: str) -> str: + slug = re.sub(r"[^a-z0-9]+", "-", text.lower()).strip("-") + return re.sub(r"-+", "-", slug)[:48] or "unknown" diff --git a/motrix_rl/src/motrix_rl/utils.py b/motrix_rl/src/motrix_rl/utils.py index 84072fe3..b7bd5d6d 100644 --- a/motrix_rl/src/motrix_rl/utils.py +++ b/motrix_rl/src/motrix_rl/utils.py @@ -28,11 +28,14 @@ def _check_gpu_available_for_torch(): def get_device_supports() -> DeviceSupports: supports = DeviceSupports() + torch_prior_device = None try: import torch # noqa: F401 supports.torch = True supports.torch_gpu = _check_gpu_available_for_torch() + if supports.torch_gpu: + torch_prior_device = torch.cuda.current_device() except ImportError: pass @@ -48,6 +51,15 @@ def get_device_supports() -> DeviceSupports: except ImportError: pass + # XLA's GPU probe enumerates every visible device and can leave the CUDA + # runtime's current device on the LAST probed GPU; torch lazily inherits + # that value, silently poisoning every later torch.cuda.current_device() + # consumer in this process. Restore the pre-probe device. + if torch_prior_device is not None: + import torch + + torch.cuda.set_device(torch_prior_device) + return supports diff --git a/motrix_rl/tests/test_doc_generate_video_script.py b/motrix_rl/tests/test_doc_generate_video_script.py index e49d8d79..f17dfdcd 100644 --- a/motrix_rl/tests/test_doc_generate_video_script.py +++ b/motrix_rl/tests/test_doc_generate_video_script.py @@ -9,7 +9,7 @@ import pytest from motrix_env_core.renderer import RenderConfig -from motrix_envs.locomotion.wbt.g1 import G1WbtEnvCfg +from motrix_envs.locomotion.wbt.g1.common import G1WbtEnvCfg from motrix_rl import checkpoints, runs GENERATE_VIDEO_SCRIPT = Path(__file__).resolve().parents[2] / "docs" / "scripts" / "generate_video.py" diff --git a/motrix_rl/tests/test_fastsac_async_multi.py b/motrix_rl/tests/test_fastsac_async_multi.py index 046a871d..43a28110 100644 --- a/motrix_rl/tests/test_fastsac_async_multi.py +++ b/motrix_rl/tests/test_fastsac_async_multi.py @@ -18,6 +18,7 @@ CollectorInfo, LearnerInfo, TrainerTopology, + resolve_collector_inference_devices, resolve_learner_devices, resolve_trainer_topology, ring_transport_is_ipc, @@ -244,21 +245,25 @@ def test_topology_pairs_collectors_with_owning_learner(monkeypatch) -> None: monkeypatch.setattr(numa, "available_numa_nodes", lambda: [0, 1]) monkeypatch.setattr(numa, "gpu_numa_node", lambda index: {0: 0, 1: 1}.get(index)) + monkeypatch.setattr("torch.cuda.device_count", lambda: 2) cuda = lambda n: torch.device("cuda", n) # noqa: E731 opts = SimpleNamespace(transition_ipc="auto", weight_ipc="auto", weight_ipc_min_bytes=0) - # 2 learners x 2 collectors: each collector sits on its owner's GPU-local - # node — the pair never straddles a NUMA boundary + # 2 learners x 2 collectors: a generic "cuda" collector spec co-locates + # each collector with its owner's GPU — the pair never straddles a NUMA + # boundary topo = resolve_trainer_topology( 8, 4, 2, - [cuda(0), cuda(1)], - [cuda(0), cuda(0), cuda(1), cuda(1)], + ["cuda:0", "cuda:1"], + "cuda", cuda(0), opts, actor_param_numel=0, ) + assert [str(learner.device) for learner in topo.learners] == ["cuda:0", "cuda:1"] + assert [str(collector.device) for collector in topo.collectors] == ["cuda:0", "cuda:0", "cuda:1", "cuda:1"] assert [learner.numa_node for learner in topo.learners] == [0, 1] assert [collector.numa_node for collector in topo.collectors] == [0, 0, 1, 1] assert topo.env_shards == [2, 2, 2, 2] @@ -266,7 +271,8 @@ def test_topology_pairs_collectors_with_owning_learner(monkeypatch) -> None: assert [collector.ring_ipc for collector in topo.collectors] == [True, True, True, True] # multi-collector single learner: every collector follows the one learner - topo = resolve_trainer_topology(8, 4, 1, [], [cuda(1)] * 4, cuda(1), opts, actor_param_numel=0) + topo = resolve_trainer_topology(8, 4, 1, None, "cuda", cuda(1), opts, actor_param_numel=0) + assert [str(collector.device) for collector in topo.collectors] == ["cuda:1"] * 4 assert [learner.numa_node for learner in topo.learners] == [1] assert [collector.numa_node for collector in topo.collectors] == [1, 1, 1, 1] @@ -284,6 +290,7 @@ def test_topology_chunks_cpus_by_node_local_ordinal(monkeypatch) -> None: monkeypatch.setattr(numa, "available_numa_nodes", lambda: [0, 1]) monkeypatch.setattr(numa, "gpu_numa_node", lambda index: {0: 0, 1: 1}.get(index)) monkeypatch.setattr(numa, "numa_node_cpus", lambda node: node_cpus[node]) + monkeypatch.setattr("torch.cuda.device_count", lambda: 2) cuda = lambda n: torch.device("cuda", n) # noqa: E731 opts = SimpleNamespace(transition_ipc="auto", weight_ipc="auto", weight_ipc_min_bytes=0) @@ -291,8 +298,8 @@ def test_topology_chunks_cpus_by_node_local_ordinal(monkeypatch) -> None: 8, 4, 2, - [cuda(0), cuda(1)], - [cuda(0), cuda(0), cuda(1), cuda(1)], + ["cuda:0", "cuda:1"], + "cuda", cuda(0), opts, actor_param_numel=0, @@ -316,38 +323,38 @@ def test_topology_unbound_without_gpu_locality(monkeypatch) -> None: cpu = torch.device("cpu") opts = SimpleNamespace(transition_ipc="auto", weight_ipc="auto", weight_ipc_min_bytes=0) - # single-node host: no binding anywhere + # single-node host, CPU collector inference: no binding anywhere monkeypatch.setattr(numa, "available_numa_nodes", lambda: [0]) - topo = resolve_trainer_topology(2, 2, 1, [], [cuda, cpu], cuda, opts, actor_param_numel=0) + topo = resolve_trainer_topology(2, 2, 1, None, "cpu", cuda, opts, actor_param_numel=0) assert [learner.numa_node for learner in topo.learners] == [None] assert [collector.numa_node for collector in topo.collectors] == [None, None] # multi-node host, CPU learner: no binding anywhere monkeypatch.setattr(numa, "available_numa_nodes", lambda: [0, 1]) - topo = resolve_trainer_topology(2, 2, 1, [], [cpu, cpu], cpu, opts, actor_param_numel=0) + topo = resolve_trainer_topology(2, 2, 1, None, "cpu", cpu, opts, actor_param_numel=0) assert [learner.numa_node for learner in topo.learners] == [None] assert [collector.numa_node for collector in topo.collectors] == [None, None] # multi-node host, unknown GPU locality: no binding anywhere monkeypatch.setattr(numa, "gpu_numa_node", lambda index: None) - topo = resolve_trainer_topology(2, 2, 1, [], [cuda, cuda], cuda, opts, actor_param_numel=0) + topo = resolve_trainer_topology(2, 2, 1, None, "cuda", cuda, opts, actor_param_numel=0) assert [learner.numa_node for learner in topo.learners] == [None] assert [collector.numa_node for collector in topo.collectors] == [None, None] # index-less cuda means device 0 monkeypatch.setattr(numa, "gpu_numa_node", lambda index: {0: 0}.get(index)) - topo = resolve_trainer_topology(2, 2, 1, [], [cuda, cuda], torch.device("cuda"), opts, actor_param_numel=0) + topo = resolve_trainer_topology(2, 2, 1, None, "cuda", torch.device("cuda"), opts, actor_param_numel=0) assert [learner.numa_node for learner in topo.learners] == [0] def test_topology_rejects_bad_counts() -> None: opts = SimpleNamespace(transition_ipc="auto", weight_ipc="auto", weight_ipc_min_bytes=0) with pytest.raises(ValueError, match="divide evenly"): - resolve_trainer_topology(3, 3, 2, [], [], torch.device("cpu"), opts, actor_param_numel=0) + resolve_trainer_topology(3, 3, 2, None, "cpu", torch.device("cpu"), opts, actor_param_numel=0) with pytest.raises(ValueError, match="invalid worker counts"): - resolve_trainer_topology(4, 0, 1, [], [], torch.device("cpu"), opts, actor_param_numel=0) - with pytest.raises(ValueError, match="collector devices"): - resolve_trainer_topology(4, 2, 1, [], [torch.device("cpu")], torch.device("cpu"), opts, actor_param_numel=0) + resolve_trainer_topology(4, 0, 1, None, "cpu", torch.device("cpu"), opts, actor_param_numel=0) + with pytest.raises(ValueError, match="collector_inference_device"): + resolve_trainer_topology(4, 2, 1, None, "meta", torch.device("cpu"), opts, actor_param_numel=0) def test_ring_slice_partitions_collectors_by_ownership() -> None: @@ -441,6 +448,49 @@ def _ipc_opts(mode: str = "auto") -> SimpleNamespace: return SimpleNamespace(transition_ipc=mode) +def test_resolve_collector_inference_devices_colocates_with_owner(monkeypatch) -> None: + """Generic "cuda" resolves to the owning learner's GPU for every learner count. + + Regression: the single-learner path used to resolve the generic spec + in-process via torch.cuda.current_device(), which silently crossed GPUs + (and therefore disabled the IPC transports) when the current device + drifted from the learner's. + """ + cuda = lambda n: torch.device("cuda", n) # noqa: E731 + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "device_count", lambda: 2) + + # single learner on an index-less cuda device: every collector gets the + # learner's GPU explicitly (index 0), never the ambient current device + devices = resolve_collector_inference_devices("cuda", [], 3, 1, torch.device("cuda")) + assert devices == [cuda(0)] * 3 + + # single learner on cuda:1: collectors follow it + devices = resolve_collector_inference_devices("cuda", [], 2, 1, cuda(1)) + assert devices == [cuda(1)] * 2 + + # multi-learner: collectors chunk to their owning learner's GPU + devices = resolve_collector_inference_devices("cuda", [cuda(0), cuda(1)], 4, 2, cuda(0)) + assert devices == [cuda(0), cuda(0), cuda(1), cuda(1)] + + # explicit specs pass through unchanged (even cross-GPU) + devices = resolve_collector_inference_devices("cuda:1", [], 2, 1, cuda(0)) + assert devices == [cuda(1)] * 2 + devices = resolve_collector_inference_devices("cpu", [cuda(0), cuda(1)], 4, 2, cuda(0)) + assert devices == [torch.device("cpu")] * 4 + + # the resolved co-located layout keeps the IPC ring the default + flags = ring_transport_is_ipc(_ipc_opts(), [], [cuda(0)] * 4, 4, 1, torch.device("cuda")) + assert flags == [True] * 4 + + # invalid specs are rejected + with pytest.raises(ValueError, match="collector_inference_device"): + resolve_collector_inference_devices("meta", [], 1, 1, cuda(0)) + monkeypatch.setattr(torch.cuda, "device_count", lambda: 1) + with pytest.raises(RuntimeError, match="only 1 device"): + resolve_collector_inference_devices("cuda:1", [], 1, 1, cuda(0)) + + def test_ring_transport_ipc_colocated_collectors() -> None: """Collectors co-located with their owning learner's GPU get IPC rings per rank.""" cuda = lambda n: torch.device("cuda", n) # noqa: E731 diff --git a/scripts/private/g1_flip_converter.py b/scripts/private/g1_flip_converter.py new file mode 100644 index 00000000..38d59086 --- /dev/null +++ b/scripts/private/g1_flip_converter.py @@ -0,0 +1,267 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""Convert mjbatch G1 backflip motion NPZ to MotrixLab NPZ v1. + +Source: https://github.com/kevinzakka/mjbatch ``examples/assets/flip.npz`` +(Apache-2.0; a mocap flip retargeted to the Unitree G1). The file ships +``joint_pos`` / ``body_*`` arrays plus ``fps`` but no ``joint_names`` / +``body_names`` and quaternions in ``wxyz`` order — the same situation as the +K1 MuJoCo clips, so this converter follows +:mod:`scripts.private.k1_mujoco_converter`: assign the menagerie G1 29-DOF +column order, remap joints by name into the training model's order, then +re-bake all body poses/velocities with the same :class:`UnitreeG129Dof` model +used by the WBT environment. + +The source is 50 fps, matching the WBT ``ctrl_dt`` of 0.02 s, so no resampling +is needed by default. mjbatch's ``g1_flip.py`` tracks only frames [0, 200) +("the ankles roll after 200"); trim with ``--end-sec 4.0`` for a clean clip. + +Column-order assumption is verified at conversion time: the FK-baked +``body_pos_w`` is compared against the source ``body_pos_w`` by body name, and +the conversion aborts if the error is large (wrong joint order assumption). +""" + +from __future__ import annotations + +from pathlib import Path + +import motrixsim as mtx +import numpy as np + +from motrix_env_core.config import configclass +from motrix_env_core.config.scene import RobotCfg, SceneCfg, SceneObjsCfg +from motrix_env_core.math import quaternion +from motrix_env_motrixsim.compiler import build_scene_model +from motrix_envs.motion.converters.lafan_converter import ( + _angular_velocity_w, + _normalize, + _resample, + _trim, +) +from motrix_envs.motion.schema import SCHEMA_VERSION, XYZW_FROM_WXYZ +from motrix_envs.robot import UnitreeG129Dof + +# Assumed joint column order of the mjbatch flip.npz: the MuJoCo Menagerie +# unitree_g1 29-DOF qpos order (verified indirectly by the FK cross-check). +_G1_JOINT_ORDER = ( + "left_hip_pitch_joint", + "left_hip_roll_joint", + "left_hip_yaw_joint", + "left_knee_joint", + "left_ankle_pitch_joint", + "left_ankle_roll_joint", + "right_hip_pitch_joint", + "right_hip_roll_joint", + "right_hip_yaw_joint", + "right_knee_joint", + "right_ankle_pitch_joint", + "right_ankle_roll_joint", + "waist_yaw_joint", + "waist_roll_joint", + "waist_pitch_joint", + "left_shoulder_pitch_joint", + "left_shoulder_roll_joint", + "left_shoulder_yaw_joint", + "left_elbow_joint", + "left_wrist_roll_joint", + "left_wrist_pitch_joint", + "left_wrist_yaw_joint", + "right_shoulder_pitch_joint", + "right_shoulder_roll_joint", + "right_shoulder_yaw_joint", + "right_elbow_joint", + "right_wrist_roll_joint", + "right_wrist_pitch_joint", + "right_wrist_yaw_joint", +) + +# Source body column order: menagerie G1 link order, pelvis first. +_G1_ROOT_BODY = "pelvis" +_G1_REFERENCE_BODY = "torso_link" + +_MAX_FK_MISMATCH = 5e-3 # m; baked vs source body position, per shared body +_MAX_ROOT_ANG_VEL_MISMATCH = 0.15 # rad/s; baked vs quat-finite-difference root angular velocity + + +@configclass +class _RobotOnlySceneObjsCfg(SceneObjsCfg): + robot: RobotCfg + + +def _build_default_model() -> mtx.SceneModel: + """Build the same G1 robot model used by the locomotion environments.""" + return build_scene_model(SceneCfg(objs=_RobotOnlySceneObjsCfg(robot=UnitreeG129Dof()))) + + +def convert_g1_flip( + input_path: str | Path, + output_path: str | Path, + *, + input_fps: float | None = None, + output_fps: float = 50.0, + start_sec: float = 0.0, + end_sec: float | None = None, + model_file: str | Path | None = None, +) -> dict[str, object]: + """Convert a mjbatch G1 flip NPZ to a MotrixLab motion NPZ v1.""" + input_path = Path(input_path).expanduser() + output_path = Path(output_path).expanduser() + if not input_path.exists(): + raise FileNotFoundError(f"Input motion file does not exist: {input_path}") + output_path.parent.mkdir(parents=True, exist_ok=True) + + required = {"fps", "joint_pos", "joint_vel", "body_pos_w", "body_quat_w"} + with np.load(input_path, allow_pickle=False) as data: + missing = required.difference(data.files) + if missing: + raise ValueError(f"mjbatch flip missing keys {sorted(missing)}: {input_path}") + stored_fps = float(np.asarray(data["fps"]).reshape(-1)[0]) + source_joint_pos = np.asarray(data["joint_pos"], dtype=np.float64) + source_body_pos = np.asarray(data["body_pos_w"], dtype=np.float64) + source_body_quat_wxyz = np.asarray(data["body_quat_w"], dtype=np.float64) + + if source_joint_pos.shape[1] != len(_G1_JOINT_ORDER): + raise ValueError(f"joint_pos has {source_joint_pos.shape[1]} columns, expected {len(_G1_JOINT_ORDER)}") + + model = mtx.load_model(str(Path(model_file).expanduser())) if model_file is not None else _build_default_model() + joint_names = [str(n) for n in model.joint_names] + if sorted(joint_names) != sorted(_G1_JOINT_ORDER): + raise ValueError( + "Model joints do not match the assumed G1 joint set.\n" + f" model: {joint_names}\n assumed: {list(_G1_JOINT_ORDER)}" + ) + + # Root pose from the source pelvis row (menagerie body 0). + root_pos = source_body_pos[:, 0, :] + root_quat = _normalize(source_body_quat_wxyz[:, 0, :][:, list(XYZW_FROM_WXYZ)]) + dof_src = source_joint_pos # already in _G1_JOINT_ORDER + + src_fps = stored_fps if input_fps is None else float(input_fps) + root_pos, root_quat, dof_src = _trim(root_pos, root_quat, dof_src, src_fps, start_sec, end_sec, input_path) + root_pos, root_quat, dof_src = _resample(root_pos, root_quat, dof_src, src_fps, output_fps) + dof = dof_src[:, [_G1_JOINT_ORDER.index(name) for name in joint_names]] + num_frames = root_pos.shape[0] + + dt = 1.0 / output_fps + root_lin_vel = np.gradient(root_pos, dt, axis=0) if num_frames > 1 else np.zeros_like(root_pos) + root_ang_vel_w = _angular_velocity_w(root_quat, dt) + dof_vel = np.gradient(dof, dt, axis=0) if num_frames > 1 else np.zeros_like(dof) + # MotrixSim free-joint qvel follows the MuJoCo convention: linear velocity + # in the world frame but ANGULAR velocity in the body frame. Feeding the + # world-frame angular velocity here makes FK re-bake every link's angular + # velocity as R @ omega_world (rotated by the instantaneous pose), which + # silently corrupts body_ang_vel_w while positions stay correct. + root_ang_vel_body = quaternion.rotate_inverse(root_quat, root_ang_vel_w) + + qpos = np.concatenate([root_pos, root_quat, dof], axis=1).astype(np.float32) + qvel = np.concatenate([root_lin_vel, root_ang_vel_body, dof_vel], axis=1).astype(np.float32) + if qpos.shape[1] != model.num_dof_pos or qvel.shape[1] != model.num_dof_vel: + raise ValueError( + f"qpos/qvel width ({qpos.shape[1]}/{qvel.shape[1]}) != model dof ({model.num_dof_pos}/{model.num_dof_vel})" + ) + + data = mtx.SceneData(model, batch=[num_frames]) + data.set_dof_pos(qpos, model) + data.set_dof_vel(qvel) + model.forward_kinematic(data) + + poses = np.asarray(model.get_link_poses(data), dtype=np.float32) + body_pos_w = poses[:, :, 0:3].copy() + body_names = [str(n) for n in model.link_names] + + # Cross-check the assumed joint order: FK must reproduce the source body + # positions (modulo the trim/resample boundary rows). + menagerie_bodies = [ + "pelvis", + "left_hip_pitch_link", + "left_hip_roll_link", + "left_hip_yaw_link", + "left_knee_link", + "left_ankle_pitch_link", + "left_ankle_roll_link", + "right_hip_pitch_link", + "right_hip_roll_link", + "right_hip_yaw_link", + "right_knee_link", + "right_ankle_pitch_link", + "right_ankle_roll_link", + "waist_yaw_link", + "waist_roll_link", + "torso_link", + "left_shoulder_pitch_link", + "left_shoulder_roll_link", + "left_shoulder_yaw_link", + "left_elbow_link", + "left_wrist_roll_link", + "left_wrist_pitch_link", + "left_wrist_yaw_link", + "right_shoulder_pitch_link", + "right_shoulder_roll_link", + "right_shoulder_yaw_link", + "right_elbow_link", + "right_wrist_roll_link", + "right_wrist_pitch_link", + "right_wrist_yaw_link", + ] + src_body_index = {name: i for i, name in enumerate(menagerie_bodies)} + n_check = min(num_frames, source_body_pos.shape[0]) + errors = [ + np.linalg.norm(body_pos_w[t, body_names.index(name)] - source_body_pos[t, src_body_index[name]]) + for t in range(n_check) + for name in menagerie_bodies + if name in body_names + ] + if not errors: + missing = [name for name in menagerie_bodies if name not in body_names] + raise ValueError( + "FK cross-check found none of the expected menagerie body names in " + f"the model links; unmatched names (first 5): {missing[:5]}" + ) + max_err = float(np.max(errors)) + if max_err > _MAX_FK_MISMATCH: + raise ValueError( + f"FK cross-check failed: max body position mismatch {max_err:.4f} m > " + f"{_MAX_FK_MISMATCH} m; the assumed joint column order is likely wrong." + ) + + # Velocity self-check: the baked root angular velocity must match the + # world-frame quaternion finite difference. A qvel frame-convention bug + # (world vs body angular velocity) corrupts velocities while leaving all + # position cross-checks green, so it needs its own guard. + body_ang_vel_w = np.asarray(model.get_link_angular_velocities(data), dtype=np.float32) + root_ang_err = float( + np.abs(body_ang_vel_w[:, body_names.index("pelvis")] - root_ang_vel_w.astype(np.float32)).mean() + ) + if root_ang_err > _MAX_ROOT_ANG_VEL_MISMATCH: + raise ValueError( + f"FK velocity self-check failed: mean root angular velocity mismatch {root_ang_err:.4f} rad/s > " + f"{_MAX_ROOT_ANG_VEL_MISMATCH} rad/s; the free-joint qvel angular frame convention is likely wrong." + ) + + output = { + "schema_version": np.int32(SCHEMA_VERSION), + "fps": np.int32(round(output_fps)), + "num_frames": np.int32(num_frames), + "joint_names": np.asarray(joint_names), + "body_names": np.asarray(body_names), + "joint_pos": qpos[:, 7:].copy(), + "joint_vel": qvel[:, 6:].copy(), + "body_pos_w": body_pos_w, + "body_quat_w": poses[:, :, 3:7].copy(), + "body_lin_vel_w": np.asarray(model.get_link_linear_velocities(data), dtype=np.float32), + "body_ang_vel_w": body_ang_vel_w, + "root_body_name": np.asarray(_G1_ROOT_BODY), + "reference_body_name": np.asarray(_G1_REFERENCE_BODY), + "clip_name": np.asarray(input_path.stem), + } + np.savez(output_path, **output) + + return { + "num_frames": num_frames, + "num_joints": len(joint_names), + "num_bodies": len(body_names), + "has_object": False, + "output_path": str(output_path), + "fk_max_mismatch_m": max_err, + } diff --git a/test/test_all_envs.py b/test/test_all_envs.py index 310a8a93..84fc94f4 100644 --- a/test/test_all_envs.py +++ b/test/test_all_envs.py @@ -1,40 +1,81 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 # SPDX-License-Identifier: Apache-2.0 -import numpy as np +"""All registered environments must build and step for a few iterations. + +Each environment is smoke-tested in its own child process: manager +environments own per-process simulation state, and hosting several live +manager models in one process is not a supported runtime contract, so the +contract under test is per-env build/step in isolation. + +The first whole-body-tracking case runs serially as a warmup: on a cold +numba kernel cache it pays the one-time kernel compilation cost, so the +parallel pool only ever hits the warm cache. +""" + +import concurrent.futures +import subprocess +import sys import motrix_envs # noqa: F401 registers built-in environments from motrix_env_core import registry +_RUNNER = """ +import sys -def test_all_demos(): - print("Start testing:") +import numpy as np - num_envs = [1, 2] +import motrix_envs # noqa: F401 registers built-in environments +from motrix_env_core import registry + +env = registry.make(sys.argv[1], num_envs=int(sys.argv[2])) +action_space = env.action_space +action = np.zeros((env.num_envs, *action_space.shape), dtype=action_space.dtype) +for _ in range(10): + env.step(action) +""" - all_envs = list(registry.list_registered_envs()) +# Fixed on purpose: the value only trades throughput for memory (four full +# interpreter imports), it must not change test behaviour across machines. +_MAX_WORKERS = 4 +_CASE_TIMEOUT_SECONDS = 1500 +_WARMUP_ENV = "g1-29dof-wbt-largebox" - total_count = 0 - failed_count = 0 - for num_env in num_envs: - for env_name in all_envs: - total_count += 1 - try: - # Create environment (manager envs inject the default backend) - env = registry.make(env_name, num_envs=num_env) +def _smoke_in_subprocess(case: tuple[str, int]) -> tuple[str, str | None]: + env_name, num_env = case + label = f"{env_name} (num_envs={num_env})" + try: + result = subprocess.run( + [sys.executable, "-c", _RUNNER, env_name, str(num_env)], + capture_output=True, + text=True, + timeout=_CASE_TIMEOUT_SECONDS, + ) + except subprocess.TimeoutExpired as exc: + tail = "\n".join((exc.stderr or "").strip().splitlines()[-5:]) + return (label, f"timed out after {_CASE_TIMEOUT_SECONDS}s\n{tail}") + if result.returncode != 0: + tail = "\n".join(result.stderr.strip().splitlines()[-5:]) + return (label, tail) + return (label, None) - action_space = env.action_space - action = np.zeros((env.num_envs, *action_space.shape), dtype=action_space.dtype) - for _ in range(10): - env.step(action) - print(f"{env_name} pass.") +def test_all_demos(): + all_envs = sorted(registry.list_registered_envs()) + cases = [(env_name, num_env) for num_env in (1, 2) for env_name in all_envs] + + failures = [] + warmup = [case for case in cases if case[0] == _WARMUP_ENV] + for case in warmup: + label, error = _smoke_in_subprocess(case) + if error is not None: + failures.append(f"{label}:\n{error}") + remaining = [case for case in cases if case not in warmup] - except Exception as e: - failed_count += 1 - print(f"{env_name} fail.") - print(e) + with concurrent.futures.ThreadPoolExecutor(max_workers=_MAX_WORKERS) as pool: + for label, error in pool.map(_smoke_in_subprocess, remaining): + if error is not None: + failures.append(f"{label}:\n{error}") - print(f"\nComplete {total_count} tests.\n{total_count - failed_count} cases pass.") - assert failed_count == 0, f"{failed_count} cases failed" + assert not failures, f"{len(failures)} env smoke cases failed:\n" + "\n".join(failures)