From b235d211ef662e671589f182b39598a78708a33f Mon Sep 17 00:00:00 2001 From: yuecideng Date: Tue, 27 Oct 2026 13:15:29 +0000 Subject: [PATCH 1/4] feat(examples): qualify official RigidFormer workflows --- agent_context/MAP.yaml | 3 +- .../topics/neural-dynamics/neural-dynamics.md | 6 +- .../api_reference/model_rigidformer.rst | 7 +- embodichain/learning/sl/types.py | 6 +- examples/learning/rigidformer/README.md | 176 +++ examples/learning/rigidformer/__init__.py | 21 + examples/learning/rigidformer/__main__.py | 220 +++ examples/learning/rigidformer/_prepare.py | 295 ++++ .../learning/rigidformer/_support/__init__.py | 21 + .../learning/rigidformer/_support/_compare.py | 321 ++++ .../rigidformer/_support/_evaluator.py | 1336 +++++++++++++++++ .../rigidformer/_support/_evidence.py | 83 + .../rigidformer/_support/_recovery.py | 132 ++ .../rigidformer/configs/sl.movi_a.yaml | 16 + .../rigidformer/configs/sl.movi_b.yaml | 16 + .../rigidformer/configs/sl.movi_sphere.yaml | 16 + .../rigidformer/research/acceptance.md | 104 ++ .../rigidformer/research/qualification.json | 139 ++ .../dynamics/rigidformer/data.py | 33 +- .../dynamics/rigidformer/ema.py | 25 +- .../dynamics/rigidformer/model.py | 6 +- .../dynamics/rigidformer/supervised.py | 2 + tests/examples/test_rigidformer_showcase.py | 73 + 23 files changed, 3041 insertions(+), 16 deletions(-) create mode 100644 examples/learning/rigidformer/README.md create mode 100644 examples/learning/rigidformer/__init__.py create mode 100644 examples/learning/rigidformer/__main__.py create mode 100644 examples/learning/rigidformer/_prepare.py create mode 100644 examples/learning/rigidformer/_support/__init__.py create mode 100644 examples/learning/rigidformer/_support/_compare.py create mode 100644 examples/learning/rigidformer/_support/_evaluator.py create mode 100644 examples/learning/rigidformer/_support/_evidence.py create mode 100644 examples/learning/rigidformer/_support/_recovery.py create mode 100644 examples/learning/rigidformer/configs/sl.movi_a.yaml create mode 100644 examples/learning/rigidformer/configs/sl.movi_b.yaml create mode 100644 examples/learning/rigidformer/configs/sl.movi_sphere.yaml create mode 100644 examples/learning/rigidformer/research/acceptance.md create mode 100644 examples/learning/rigidformer/research/qualification.json create mode 100644 tests/examples/test_rigidformer_showcase.py diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index 39bed0f05d..b81e03e7fe 100644 --- a/agent_context/MAP.yaml +++ b/agent_context/MAP.yaml @@ -283,7 +283,8 @@ topics: - examples/learning/nerd/_support/ - examples/learning/nerd/configs/ - extensions/embodichain_models/embodichain_models/dynamics/rigidformer/ - watch_paths: [extensions/embodichain_models/, examples/learning/nerd/, + - examples/learning/rigidformer/ + watch_paths: [extensions/embodichain_models/, examples/learning/nerd/, examples/learning/rigidformer/, docs/source/features/neural_dynamics.md, docs/source/api_reference/model_catalogs.rst] related_topics: [supervised-learning, learning-layers, data-pipeline, data-assets] diff --git a/agent_context/topics/neural-dynamics/neural-dynamics.md b/agent_context/topics/neural-dynamics/neural-dynamics.md index 1fbe47dcd0..bf46041422 100644 --- a/agent_context/topics/neural-dynamics/neural-dynamics.md +++ b/agent_context/topics/neural-dynamics/neural-dynamics.md @@ -57,8 +57,10 @@ and at least two valid objects. Attention backend and source semantics belong to recipe/recovery identity. Source-owned epoch sampling does not claim the legacy research DataLoader's worker-RNG trajectory. -Numerical agreement, recovery and full benchmark/retraining qualification -remain separate evidence requirements. +The [showcase](../../../examples/learning/rigidformer/README.md) owns downloads, +conversion, fixed-weight evaluation and independent author-Trainer comparison. +Assets/results stay outside the wheel. Numerical agreement, recovery and full +benchmark/retraining qualification remain distinct evidence requirements. ## Resource and physical identity diff --git a/docs/source/api_reference/model_rigidformer.rst b/docs/source/api_reference/model_rigidformer.rst index 803c9730b1..6176d85b8d 100644 --- a/docs/source/api_reference/model_rigidformer.rst +++ b/docs/source/api_reference/model_rigidformer.rst @@ -146,5 +146,8 @@ source identity also binds splits, scaling, sampling and seed. EMA shadows remai FP32 and retain the author's warmup policy. Official weight import is distinct from native optimizer recovery; export selects EMA or raw weights explicitly. -Applications own acquisition and independent reference experiments. Short -numerical comparisons do not establish published accuracy. +The showcase owns asset preparation, the author evaluator and independent +reference comparison. Fixed-checkpoint test evaluation does not select a +checkpoint. New learning recipes should select using validation data. Reports +record source/data/checkpoint identities, attention backend and subset/full scope. +Short numerical comparisons do not establish published accuracy. diff --git a/embodichain/learning/sl/types.py b/embodichain/learning/sl/types.py index 1f451cd7bd..81ccb3f675 100644 --- a/embodichain/learning/sl/types.py +++ b/embodichain/learning/sl/types.py @@ -119,7 +119,11 @@ def on_update(self, learner: SupervisedLearner) -> None: ... def state_dict(self) -> Mapping[str, Any]: - """Return an owned serializable snapshot of the committed state.""" + """Return an owned serializable snapshot of the committed state. + + Returns: + State whose mutable values do not alias the live participant. + """ ... def validate_state_dict(self, state: Mapping[str, Any]) -> None: diff --git a/examples/learning/rigidformer/README.md b/examples/learning/rigidformer/README.md new file mode 100644 index 0000000000..c4ed3e0ee1 --- /dev/null +++ b/examples/learning/rigidformer/README.md @@ -0,0 +1,176 @@ +# RigidFormer author implementation + +This showcase integrates the author's RigidFormer implementation at +`160fe0852c51cffae7b7440452b58d0b49db6df6` from +[Frank-ZY-Dou/Dynamics-Modeling](https://github.com/Frank-ZY-Dou/Dynamics-Modeling/tree/160fe0852c51cffae7b7440452b58d0b49db6df6/RigidFormer). +The model lives at `embodichain_models.dynamics.rigidformer` and trains through +the common SL trainer. The community `rigidformer` package is not a runtime dependency. + +## Setup and ownership + +Use a Python 3.11 environment and a CUDA-capable Torch installation: + +```bash +conda activate py311 +python -m pip install 'embodichain[rigidformer]' +embodichain run-example learning/rigidformer --help +``` + +The standard native attention backend is Torch SDPA. The explicit `flash` +backend requires a Flash-Attention build compatible with Torch, CUDA and the +GPU. Reference evidence records the actual dependency/backend versions. +Imports, architecture construction and checkpoint inspection work on CPU; +geometry prediction uses the author's CUDA KNN path and needs at least two +valid objects per scene. + +Model/config/checkpoint/data/objective/recipe code is installable under the +model family. Downloading assets, conversion and reference experiments remain +showcase-owned. Private author-derived source retains MIT attribution in +`THIRD_PARTY_LICENSES`. Official scenes are CC BY 4.0; collision meshes and +derived geometry are CC BY-SA 4.0. Those licenses remain attached to prepared +data. Data, checkpoints, source checkouts and reports are not bundled in the wheel. + +## Prepare official assets + +```bash +embodichain run-example learning/rigidformer prepare --dataset movi_a +``` + +Data preparation is independent of pretrained weights. Download those separately +with `prepare --dataset movi_a --assets weights`, or select `--assets all`. +Training from scratch requires no pretrained checkpoint. + +The downloader pins dataset revision +`e78514b479b11acfa5d23d28f1211d7830646ff9` and weights revision +`8f2c0a3a157b64e4b310c64a01eee6e4e7e25af6`. It verifies published checksums. +Use `movi_sphere`, `movi_a`, or `movi_b`. Numeric scenes and collision meshes +are sufficient; videos are not downloaded. + +Prepared files use the author's packed layout, with aligned material points, +complete object identities (including recovered objects), physical properties, +reference anchors/FPS order, and a content-hashed manifest. Stored vertices are +**unscaled asset vertices transformed into world coordinates**. The size sidecar +applies each object's metadata size about its origin exactly once during reads. + +By default preparation requires all 960 training, 120 validation and 120 test +scenes. `--limit-per-split N` creates a clearly labelled debug subset. +`--local-files-only` requires already downloaded, verified assets. +Every output dataset needs a fresh directory. Failures after conversion starts +remain recorded; checksum preflight rejects bad assets before creating output. +Readers reject altered payloads and geometry omitted from the manifest. FPS and +anchor caches are committed inputs, rather than untracked files generated during +training. Debug preparation verifies collision meshes as well as scene shards. + +## Train with shared SL + +```bash +embodichain train-sl --config examples/learning/rigidformer/configs/sl.movi_a.yaml +``` + +Paths resolve from the caller's current working directory. The preset recipes +use eight-frame windows, four anchors, variable integration strides 1/5/10, +author augmentation, anchor-only before/after-Kabsch position and acceleration +losses, grouped AdamW, gradient clipping, epoch scheduling and FP32 EMA. + +The author model integrates with **model dt=1**. Physical time is +`step_size / 240` seconds. Loss scaling and conditioning use the recorded +step size; do not replace the model integration unit with seconds. +FP32 geometry augmentation precedes BF16 network autocast; integration and +rigid registration retain the author's numerical policy. + +Each source shuffles and shards its epoch traversal, drops an incomplete batch, +and stages the cursor until optimizer success. The source-sampling stream is +explicit and recoverable; the legacy research DataLoader's worker-RNG trajectory +is not claimed bitwise identical. Checkpoints preserve sampling, model/optimizer, +scheduler, EMA and per-rank state. A changed source/split/scale/recipe/world size +is rejected before restoration. + +For multiple GPUs, use the same launcher: + +```bash +OMP_NUM_THREADS=1 torchrun --standalone --nnodes=1 --nproc-per-node=4 \ + -m embodichain train-sl --config examples/learning/rigidformer/configs/sl.movi_a.yaml +``` + +Batch size is per rank, and learning rates are not scaled automatically. +The released MOVi-A/Sphere training layout uses four GPUs x 32 scenes; +MOVi-B uses six x 21. A smaller training run is useful for integration checks +but is not equivalent to that reference training budget. Native DDP preserves +the author's synchronized encoder statistics and uses a global valid-anchor +loss mean exactly once. + +Resume through the common entry point; change only the output directory when +continuing an older checkpoint into a fresh run destination: + +```bash +embodichain train-sl --config my-recipe.yaml --resume outputs/run/last.pt +``` + +## Export and evaluate + +```bash +embodichain run-example learning/rigidformer export \ + --checkpoint outputs/run/last.pt --weights ema --output outputs/model-ema.pt + +embodichain run-example learning/rigidformer evaluate \ + --checkpoint outputs/model-ema.pt \ + --data-root outputs/rigidformer/data/MOVi-A \ + --split validation --output outputs/rigidformer/eval-validation +``` + +Official weights can be passed directly to `evaluate`. Exported native weights +use the same named tensor/config layout; importing weights does not recover an +official optimizer process. EMA/raw selection is explicit, and an absent EMA +does not silently fall back to raw weights. + +The fixed-checkpoint evaluator uses four ground-truth warmup frames, step sizes +1/5/10, and the author's HopNet-aligned position/rotation metrics at physical +frames 25/50/75/100/200/300/400. Off-grid checkpoints are not interpolated. +Reports include corpus scope and dependency/source/checkpoint identities. +A debug-subset result does not establish complete benchmark accuracy. + +New native runs do not select a checkpoint from test scores. Use validation to +select, freeze the checkpoint, then evaluate test. The author's default research +training recipe monitors test metrics; that selection-protocol difference must +remain visible in any retraining comparison. + +## Independent reference comparison + +Supply an unmodified author checkout's `RigidFormer` directory at the pinned +revision. The runner verifies source hashes in `REFERENCE_SOURCES.json` before +importing the independent original model and Trainer. + +```bash +embodichain run-example learning/rigidformer compare \ + --reference /path/to/Dynamics-Modeling/RigidFormer \ + --checkpoint outputs/model-ema.pt \ + --data-root outputs/rigidformer/data/MOVi-A \ + --attention-backend flash --updates 3 --output outputs/rigidformer/compare +``` + +Comparison checks prediction, free rollout, losses, gradients and paired optimizer +updates on identical fixed inputs and randomness. `flash` executes the genuine +Flash-Attention package. Explicit `sdpa` uses a documented compatibility operator +for the reference's unmasked attention; its results do not qualify Flash execution. +Reports distinguish official released weights from other explicitly provided weights. +They also compare optimizer moments/groups, buffers, EMA and epoch schedules, +and record selected frames, source hashes and the executed numerical runtime. +Each paired update represents one fixed-batch comparison epoch; it is separate +from the source's normal training-epoch length. + +Check exact continuous versus resumed training using the same installable recipe: + +```bash +embodichain run-example learning/rigidformer check-recovery \ + --config my-four-update-recipe.yaml --stop-after 2 --output outputs/recovery-check +``` + +This creates fresh full/split/resumed runs, compares complete committed state +exactly, and verifies raw/EMA exports. The [acceptance protocol](research/acceptance.md) +also covers epoch/warmup boundaries, two-rank NCCL and installed-wheel checks. + +This is a numerical-integration check. Full published-accuracy reproduction +additionally requires frozen official weights/data/evaluation, or complete +retraining with the declared compute and selection protocol. No short-run +comparison is labelled a full training reproduction. + diff --git a/examples/learning/rigidformer/__init__.py b/examples/learning/rigidformer/__init__.py new file mode 100644 index 0000000000..2147beac27 --- /dev/null +++ b/examples/learning/rigidformer/__init__.py @@ -0,0 +1,21 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Author-based RigidFormer integration and frozen reference workflows.""" + +from __future__ import annotations + +__all__: list[str] = [] diff --git a/examples/learning/rigidformer/__main__.py b/examples/learning/rigidformer/__main__.py new file mode 100644 index 0000000000..361c3071c0 --- /dev/null +++ b/examples/learning/rigidformer/__main__.py @@ -0,0 +1,220 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Prepare official assets, evaluate inference weights, and export native runs.""" + +from __future__ import annotations +import argparse +from copy import deepcopy +from pathlib import Path +from types import SimpleNamespace +import json +import sys +from ._prepare import FOLDERS, download, prepare, sha256 + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + sub = parser.add_subparsers(dest="command", required=True) + p = sub.add_parser("prepare") + p.add_argument("--dataset", choices=tuple(FOLDERS), default="movi_b") + p.add_argument("--cache", type=Path, default=Path("outputs/rigidformer/assets")) + p.add_argument("--output", type=Path, default=Path("outputs/rigidformer/data")) + p.add_argument("--local-files-only", action="store_true") + p.add_argument("--assets", choices=("data", "weights", "all"), default="data") + p.add_argument( + "--limit-per-split", + type=int, + default=0, + help="Debug subset only; zero prepares the full frozen corpus", + ) + p = sub.add_parser("export") + p.add_argument("--checkpoint", type=Path, required=True) + p.add_argument("--output", type=Path, required=True) + p.add_argument("--weights", choices=("ema", "raw"), default="ema") + p = sub.add_parser("evaluate") + p.add_argument("--checkpoint", type=Path, required=True) + p.add_argument("--data-root", type=Path, required=True) + p.add_argument("--output", type=Path, required=True) + p.add_argument("--split", choices=("validation", "test"), default="validation") + p.add_argument("--device", default="cuda") + p.add_argument("--attention-backend", choices=("sdpa", "flash"), default="sdpa") + p.add_argument("--step-sizes", type=int, nargs="+", default=[1, 5, 10]) + p.add_argument("--num-frames", type=int, default=480) + p = sub.add_parser("compare") + p.add_argument("--checkpoint", type=Path, required=True) + p.add_argument( + "--reference", + type=Path, + required=True, + help="Pinned author checkout's RigidFormer directory", + ) + p.add_argument("--data-root", type=Path, required=True) + p.add_argument("--output", type=Path, required=True) + p.add_argument("--attention-backend", choices=("sdpa", "flash"), default="flash") + p.add_argument("--updates", type=int, default=3) + p.add_argument("--seed", type=int, default=42) + p = sub.add_parser("check-recovery") + p.add_argument("--config", type=Path, required=True) + p.add_argument("--output", type=Path, required=True) + p.add_argument("--stop-after", type=int, default=2) + args = parser.parse_args() + if args.command == "prepare": + weights = None + if args.local_files_only: + raw = args.cache / "dataset" + if args.assets in ("weights", "all"): + from ._prepare import WEIGHTS, WEIGHT_SHA256 + + weights = args.cache / "model" / WEIGHTS[args.dataset] + if ( + not weights.is_file() + or sha256(weights) != WEIGHT_SHA256[args.dataset] + ): + raise ValueError( + "A verified official checkpoint is required in the local cache" + ) + else: + raw, weights = download(args.cache, args.dataset, assets=args.assets) + root = None + if args.assets in ("data", "all"): + root = prepare( + raw, args.output, args.dataset, limit_per_split=args.limit_per_split + ) + print( + json.dumps( + { + "data_root": str(root) if root else None, + "checkpoint": str(weights) if weights else None, + "scope": "full" if args.limit_per_split == 0 else "debug-subset", + }, + indent=2, + ) + ) + elif args.command == "export": + from embodichain_models.dynamics.rigidformer.checkpoint import ( + export_rigidformer_checkpoint, + ) + + print( + export_rigidformer_checkpoint( + args.checkpoint, args.output, weights=args.weights + ) + ) + elif args.command == "check-recovery": + from ._support._recovery import check_recovery + + check_recovery(args) + elif args.command == "compare": + from ._support._compare import compare + + compare(args) + else: + evaluate(args) + + +def evaluate(args: argparse.Namespace) -> None: + from embodichain_models.dynamics.rigidformer import load_rigidformer_checkpoint + from embodichain_models.dynamics.rigidformer.data import _verified_manifest + from ._support._evidence import evidence + from ._support._evaluator import Evaluator + + if args.output.exists(): + raise FileExistsError(args.output) + args.output.mkdir(parents=True) + try: + root = args.data_root.resolve() + manifest = _verified_manifest(root) + model = load_rigidformer_checkpoint( + args.checkpoint, + device=args.device, + attention_backend=args.attention_backend, + ) + config = deepcopy(model.reference_config) + split = json.loads((root / "splits.json").read_text()) + key = "val_datasets" if args.split == "validation" else "test_datasets" + scenes = split[key] + if not scenes: + raise ValueError(f"No {args.split} scenes in the prepared corpus") + selected = { + **split, + "val_datasets": [], + "test_datasets": scenes, + "evaluated_partition": args.split, + } + splitfile = args.output / "evaluation-splits.json" + splitfile.write_text(json.dumps(selected, sort_keys=True)) + config["data"].update( + dataset_path=str(root), + split_file=str(splitfile.resolve()), + object_scale_file=str(root / "object_scales.json"), + eval_step_sizes=args.step_sizes, + ) + config["training"]["metrics_test"].update( + enable=True, + initial_eval=True, + print_every_n_epochs=1, + eval_num_frames=args.num_frames, + eval_warmup_frames=4, + ) + # A standalone, fixed-checkpoint evaluator: no training or checkpoint selection. + from torch.utils.tensorboard import SummaryWriter + + writer = SummaryWriter(log_dir=str(args.output / "tensorboard")) + trainer = SimpleNamespace( + model=model.network, + config=config, + device=args.device, + writer=writer, + log_dir=args.output, + ) + try: + metrics = Evaluator(trainer).evaluate_rollout_rmse(0) + finally: + writer.close() + # Author aggregates include NumPy scalars; persist portable JSON numbers. + metrics = json.loads( + json.dumps(metrics, default=lambda value: value.item(), allow_nan=False) + ) + record = { + **evidence(args.checkpoint, root, args.attention_backend), + "author_revision": __import__( + "embodichain_models.dynamics.rigidformer", + fromlist=["RIGIDFORMER_REVISION"], + ).RIGIDFORMER_REVISION, + "checkpoint_sha256": sha256(args.checkpoint), + "data_manifest_sha256": sha256(root / "manifest.json"), + "split": args.split, + "num_scenes": len(scenes), + "num_frames": args.num_frames, + "step_sizes": args.step_sizes, + "attention_backend": args.attention_backend, + "scope": manifest["scope"], + "metrics": metrics, + } + (args.output / "summary.json").write_text( + json.dumps(record, sort_keys=True, indent=2, allow_nan=False) + ) + print(json.dumps(record, indent=2, allow_nan=False)) + except BaseException as exc: + (args.output / "failure.json").write_text( + json.dumps({"type": type(exc).__name__, "message": str(exc)}) + ) + raise + + +if __name__ == "__main__": + main() diff --git a/examples/learning/rigidformer/_prepare.py b/examples/learning/rigidformer/_prepare.py new file mode 100644 index 0000000000..24f9822bfd --- /dev/null +++ b/examples/learning/rigidformer/_prepare.py @@ -0,0 +1,295 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Prepare pinned official RigidFormer assets and bounded packed scene files.""" + +from __future__ import annotations +from pathlib import Path +from typing import Any +import hashlib +import json +import os +import numpy as np + +DATA_REVISION = "e78514b479b11acfa5d23d28f1211d7830646ff9" +MODEL_REVISION = "8f2c0a3a157b64e4b310c64a01eee6e4e7e25af6" +WEIGHT_SHA256 = { + "movi_sphere": "ffb0efc6904b5a63504a5005bb6c4df2e79e8de20641dc205885593747fcf7d9", + "movi_a": "db2e560f132933e2169ec9f90330e6b0595b3c1656684ca7f31adf95b0bc3d19", + "movi_b": "7db1dacf8850939a6398028fea71d9d58e75eebd7fc5079e75bd8475e6acf1ec", +} +FOLDERS = {"movi_sphere": "MOVi-Sphere", "movi_a": "MOVi-A", "movi_b": "MOVi-B"} +WEIGHTS = { + "movi_sphere": "rigidformer_movisphere.pth", + "movi_a": "rigidformer_movia.pth", + "movi_b": "rigidformer_movib.pth", +} + + +def sha256(path: Path) -> str: + h = hashlib.sha256() + with path.open("rb") as f: + for block in iter(lambda: f.read(1 << 20), b""): + h.update(block) + return h.hexdigest() + + +def download( + cache: Path, dataset: str, *, assets: str = "data" +) -> tuple[Path, Path | None]: + from huggingface_hub import snapshot_download, hf_hub_download + + if assets not in ("data", "weights", "all"): + raise ValueError("assets must be data, weights or all") + raw = cache / "dataset" + if assets in ("data", "all"): + raw = Path( + snapshot_download( + "frankzydou/RigidFormer-MOVi", + repo_type="dataset", + revision=DATA_REVISION, + local_dir=raw, + allow_patterns=[ + "manifest.json", + f"data/{dataset}/*", + "assets/kubasic/*", + ], + ) + ) + weight = None + if assets in ("weights", "all"): + weight = Path( + hf_hub_download( + "frankzydou/RigidFormer", + WEIGHTS[dataset], + revision=MODEL_REVISION, + local_dir=cache / "model", + ) + ) + if sha256(weight) != WEIGHT_SHA256[dataset]: + raise ValueError("Official checkpoint checksum differs") + return raw, weight + + +def _rows(path: Path) -> list[dict[str, Any]]: + import pyarrow.parquet as pq + + columns = [ + "scene", + "split", + "num_frames", + "frame_rate", + "num_objects", + "shape", + "size", + "mass", + "friction", + "restitution", + "positions", + "quaternions", + "recovered", + ] + return pq.read_table(path, columns=columns).to_pylist() + + +def prepare(raw: Path, output: Path, dataset: str, *, limit_per_split: int = 0) -> Path: + """Write author-compatible packed vertices without storing videos.""" + import trimesh + from embodichain_models.dynamics.rigidformer._native.data.dataset import ( + MultiObjectDynamics, + ) + + if dataset not in FOLDERS or limit_per_split < 0: + raise ValueError("Unknown dataset or negative scene limit") + out = output / FOLDERS[dataset] + if out.exists(): + raise FileExistsError(f"Prepare into a fresh dataset directory: {out}") + original = json.loads((raw / "manifest.json").read_text()) + needed = { + k: v + for k, v in original["files"].items() + if k.startswith(f"data/{dataset}/") or k.startswith("assets/") + } + selected = sorted((raw / "data" / dataset).glob("*.parquet")) + if not selected: + raise FileNotFoundError("No official numeric scenes found") + selected_names = {str(p.relative_to(raw)) for p in selected} + for rel, info in needed.items(): + if not limit_per_split or rel.startswith("assets/") or rel in selected_names: + if not (raw / rel).is_file() or sha256(raw / rel) != info["sha256"]: + raise ValueError(f"Official raw data checksum differs: {rel}") + meshes = {} + out.mkdir(parents=True) + splits = { + "scene_set": "complete", + "train_datasets": [], + "val_datasets": [], + "test_datasets": [], + } + scales = { + "converted": FOLDERS[dataset], + "scales": {}, + "fingerprint": {}, + "radius0": {}, + } + totals = {"train": 0, "validation": 0, "test": 0} + files = {} + try: + for path in selected: + split = path.name.split("-")[0] + if split not in totals or ( + limit_per_split and totals[split] >= limit_per_split + ): + continue + for rec in _rows(path): + if rec["split"] != split: + raise ValueError("Scene split does not match its official shard") + if limit_per_split and totals[split] >= limit_per_split: + break + name = rec["scene"] + folder = out / name + folder.mkdir() + (folder / "anchors").mkdir() + (folder / "meshes").mkdir() + pos = np.asarray(rec["positions"], dtype=np.float64) + quat = np.asarray(rec["quaternions"], dtype=np.float64) + if ( + pos.shape != (rec["num_frames"], rec["num_objects"], 3) + or quat.shape != (*pos.shape[:2], 4) + or rec["frame_rate"] != 240 + ): + raise ValueError("Official pose layout or frame rate differs") + shapes = [] + for shape in rec["shape"]: + if shape not in meshes: + mesh = trimesh.load( + raw + / "assets" + / "kubasic" + / shape + / "collision_geometry.obj", + process=False, + ) + if not isinstance(mesh, trimesh.Trimesh): + raise ValueError("Expected one collision mesh") + meshes[shape] = mesh + shapes.append(meshes[shape]) + counts = np.array([len(m.vertices) for m in shapes], dtype=np.int64) + frames = np.lib.format.open_memmap( + folder / "frames.npy", + mode="w+", + dtype=np.float32, + shape=(len(pos), int(counts.sum()), 3), + ) + for t in range(len(pos)): + row = [] + for i, mesh in enumerate(shapes): + transform = trimesh.transformations.quaternion_matrix( + quat[t, i] + ) + transform[:3, 3] = pos[t, i] + row.append( + trimesh.transform_points(mesh.vertices, transform).astype( + np.float32 + ) + ) + frames[t] = np.concatenate(row, axis=0) + frames.flush() + np.save(folder / "frames_counts.npy", counts) + traj = np.concatenate((pos, quat), axis=-1).astype(np.float32)[None] + np.save(folder / "vis_traj.npy", traj) + physics = {} + radii = [] + offsets = np.r_[0, np.cumsum(counts)] + for i, mesh in enumerate(shapes): + physics[f"obj_{i}"] = { + k: float(rec[k][i]) + for k in ("mass", "friction", "restitution", "size") + } + np.save( + folder / "meshes" / f"obj_{i}_faces.npy", + np.asarray(mesh.faces, dtype=np.int32), + ) + geometry = np.asarray(frames[0, offsets[i] : offsets[i + 1]]) + order = MultiObjectDynamics._farthest_point_sample_np( + geometry, int(counts[i]), rule="centroid_farthest" + ) + np.save(folder / "anchors" / f"obj_{i}_fpsorder_cf.npy", order) + np.save( + folder / "anchors" / f"obj_{i}_anchors_fps4_cf.npy", order[:4] + ) + radii.append( + float( + np.linalg.norm( + geometry.astype(np.float64) - pos[0, i], axis=1 + ).max() + ) + ) + (folder / "physics_params.json").write_text( + json.dumps(physics, sort_keys=True, indent=2) + ) + (folder / "conversion.json").write_text( + json.dumps( + { + "num_frames": len(pos), + "frame_rate": 240, + "complete": True, + "vertices": "unscaled asset vertices; size sidecar required", + "source_revision": DATA_REVISION, + "recovered": rec["recovered"], + }, + sort_keys=True, + ) + ) + scales["scales"][name] = [float(v) for v in rec["size"]] + scales["fingerprint"][name] = pos[0, 0].tolist() + scales["radius0"][name] = radii + key = "val_datasets" if split == "validation" else f"{split}_datasets" + splits[key].append(name) + totals[split] += 1 + del frames + for f in sorted(folder.rglob("*")): + if f.is_file(): + files[str(f.relative_to(out))] = sha256(f) + if not totals["train"] or ( + not limit_per_split + and any(v != n for v, n in zip(totals.values(), (960, 120, 120))) + ): + raise ValueError(f"Official scene counts differ: {totals}") + (out / "splits.json").write_text(json.dumps(splits, sort_keys=True, indent=2)) + (out / "object_scales.json").write_text( + json.dumps(scales, sort_keys=True, indent=2) + ) + for f in ("splits.json", "object_scales.json"): + files[f] = sha256(out / f) + manifest = { + "format": "embodichain.rigidformer.packed.v1", + "dataset": dataset, + "data_revision": DATA_REVISION, + "scope": "full" if not limit_per_split else "debug-subset", + "scene_counts": totals, + "files": files, + "licenses": {"scenes": "CC-BY-4.0", "meshes": "CC-BY-SA-4.0"}, + } + tmp = out / "manifest.json.tmp" + tmp.write_text(json.dumps(manifest, sort_keys=True, indent=2)) + tmp.replace(out / "manifest.json") + except BaseException as exc: + (out / "failure.json").write_text( + json.dumps({"type": type(exc).__name__, "message": str(exc)}) + ) + raise + return out diff --git a/examples/learning/rigidformer/_support/__init__.py b/examples/learning/rigidformer/_support/__init__.py new file mode 100644 index 0000000000..e451e82ab9 --- /dev/null +++ b/examples/learning/rigidformer/_support/__init__.py @@ -0,0 +1,21 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Private showcase-owned reference validation.""" + +from __future__ import annotations + +__all__: list[str] = [] diff --git a/examples/learning/rigidformer/_support/_compare.py b/examples/learning/rigidformer/_support/_compare.py new file mode 100644 index 0000000000..3b1837ef41 --- /dev/null +++ b/examples/learning/rigidformer/_support/_compare.py @@ -0,0 +1,321 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Execute a frozen author implementation independently of the native model.""" + +from __future__ import annotations +import argparse +from copy import deepcopy +import importlib +from importlib.resources import files +import json +from pathlib import Path +import sys +import types +import random +import numpy as np +import torch +from .._prepare import sha256, WEIGHT_SHA256 +from embodichain_models.dynamics.rigidformer import ( + load_rigidformer_checkpoint, + RIGIDFORMER_REVISION, +) +from embodichain_models.dynamics.rigidformer.data import MOViSequenceSource +from embodichain_models.dynamics.rigidformer.supervised import ( + RigidFormerObjective, + rigidformer_batch, +) +from embodichain_models.dynamics.rigidformer.sl_recipe import _optimizer, _EpochSchedule +from embodichain_models.dynamics.rigidformer.ema import RigidFormerEMA +from embodichain.learning.sl import SupervisedLearner +from ._evidence import evidence + + +def _assert_tree(actual, expected) -> float: + """Assert complete numerical/state agreement and return its largest error.""" + if isinstance(actual, torch.Tensor): + torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-4) + return float((actual - expected).abs().max()) if actual.numel() else 0.0 + if isinstance(actual, dict): + assert actual.keys() == expected.keys() + return max( + (_assert_tree(v, expected[k]) for k, v in actual.items()), default=0.0 + ) + if isinstance(actual, (tuple, list)): + assert len(actual) == len(expected) + return max((_assert_tree(a, b) for a, b in zip(actual, expected)), default=0.0) + if isinstance(actual, float): + assert actual == expected + return 0.0 + assert actual == expected + return 0.0 + + +def _optimizer_state(optimizer, model) -> dict: + """Compare optimizer groups and moments by parameter name, not object ID.""" + names = {id(p): n.removeprefix("network.") for n, p in model.named_parameters()} + return { + "groups": [ + { + **{k: v for k, v in group.items() if k != "params"}, + "parameters": [names[id(p)] for p in group["params"]], + } + for group in optimizer.param_groups + ], + "moments": {names[id(p)]: state for p, state in optimizer.state.items()}, + } + + +def compare(args: argparse.Namespace) -> None: + if args.updates < 1: + raise ValueError("updates must be positive") + if args.output.exists(): + raise FileExistsError(args.output) + args.output.mkdir(parents=True) + report = { + **evidence(args.checkpoint, args.data_root, args.attention_backend), + "qualification": "numerical-parity-only", + "author_revision": RIGIDFORMER_REVISION, + "attention_backend": args.attention_backend, + "updates": args.updates, + "seed": args.seed, + "comparison_steps_per_epoch": 1, + } + try: + manifest = json.loads( + files("embodichain_models.dynamics.rigidformer") + .joinpath("REFERENCE_SOURCES.json") + .read_text() + ) + for path, expected in manifest["files"].items(): + if sha256(args.reference / path) != expected: + raise ValueError( + f"Reference source differs from pinned author revision: {path}" + ) + report["reference_source_sha256"] = manifest["files"] + if args.attention_backend == "sdpa": + # Explicit compatibility profile, recorded separately from real Flash-Attention. + # The author's model/trainer source is not rewritten. + module = types.ModuleType("flash_attn") + + def sdpa(q, k, v, dropout_p=0.0, softmax_scale=None, **kwargs): + return torch.nn.functional.scaled_dot_product_attention( + q.transpose(1, 2), + k.transpose(1, 2), + v.transpose(1, 2), + dropout_p=dropout_p, + scale=softmax_scale, + ).transpose(1, 2) + + module.flash_attn_func = sdpa + sys.modules["flash_attn"] = module + report["reference_attention_profile"] = "explicit-torch-sdpa-compatibility" + else: + import flash_attn + + report["reference_attention_profile"] = ( + "flash-attn-" + flash_attn.__version__ + ) + sys.path.insert(0, str(args.reference.resolve())) + original = importlib.import_module("src.train") + native = load_rigidformer_checkpoint( + args.checkpoint, device="cuda", attention_backend=args.attention_backend + ) + cfg = deepcopy(native.reference_config) + cfg["paths"] = { + "log_dir": str(args.output / "reference-logs"), + "checkpoint_dir": str(args.output / "reference-checkpoints"), + } + cfg["data"].update( + dataset_path=str(args.data_root), + split_file=str(args.data_root / "splits.json"), + object_scale_file=str(args.data_root / "object_scales.json"), + ) + cfg["training"]["checkpoint"]["save_video"] = False + cfg["training"]["compile"] = {"enable": False} + source = MOViSequenceSource( + args.data_root, + args.data_root / "splits.json", + args.data_root / "object_scales.json", + data_config=cfg["data"], + ) + random.seed(args.seed) + np.random.seed(args.seed) + torch.manual_seed(args.seed) + raw = source.collate([source.dataset[0]]) + report["input_scenes"] = raw["dataset_name"] + report["input_frame_ids"] = raw["frame_ids"].tolist() + raw = {k: v.to("cuda") for k, v in raw.items() if isinstance(v, torch.Tensor)} + reference = original.CUT3RDynamicsModel(deepcopy(cfg)).cuda().eval() + reference.load_state_dict(native.network.state_dict(), strict=True) + with torch.no_grad(): + expected = reference(raw) + got = native(raw) + prediction_errors = { + k: float((got[k] - v).abs().max()) + for k, v in expected.items() + if isinstance(v, torch.Tensor) + } + for key in prediction_errors: + torch.testing.assert_close(got[key], expected[key], atol=1e-5, rtol=1e-4) + warm = { + k: (v[:, :4] if k in ("obj_positions", "valid_vtx_mask") else v) + for k, v in raw.items() + } + trajectory_ref = reference.rollout( + warm, num_steps=12, warmup_frames=4, return_extras=False + ) + trajectory_native = native.rollout(warm, num_frames=12, warmup_frames=4) + torch.testing.assert_close( + trajectory_native["obj_positions"], + trajectory_ref["obj_positions"], + atol=1e-5, + rtol=1e-4, + ) + report["prediction_max_abs"] = prediction_errors + report["rollout_max_abs"] = float( + (trajectory_native["obj_positions"] - trajectory_ref["obj_positions"]) + .abs() + .max() + ) + # The real, unmodified author Trainer computes its own losses, optimizer + # groups, gradient clipping and EMA; the native side uses shared SL. + reference_trainer = original.Trainer( + deepcopy(cfg), reference, [raw], [], device=torch.device("cuda") + ) + native_optimizer = _optimizer(native, cfg["training"]["optimizer"]) + schedule = _EpochSchedule( + native_optimizer, + 1, + cfg["training"]["scheduler"], + int(cfg["training"]["max_epochs"]), + ) + average = RigidFormerEMA(native, float(cfg["training"]["ema"]["decay"])) + if reference_trainer.ema is None: + raise ValueError("The comparison requires the released EMA-enabled recipe") + _assert_tree( + _optimizer_state(native_optimizer, native), + _optimizer_state(reference_trainer.optimizer, reference), + ) + from embodichain_models.dynamics.rigidformer._native.data.gpu_augmentation import ( + GPUAugmentation, + ) + + objective = RigidFormerObjective( + { + name: cfg["training"]["loss_weights"][name] + for name in ( + "position_before_icp", + "position_after_icp", + "acceleration_before_icp", + "acceleration_after_icp", + ) + }, + time_step_scale=cfg["training"]["loss_weights"].get( + "if_time_step_scale", True + ), + augmentation=GPUAugmentation(cfg, torch.device("cuda")), + ) + learner = SupervisedLearner( + native, + objective, + native_optimizer, + max_grad_norm=cfg["training"]["gradient_clip"], + amp_dtype=torch.bfloat16, + ) + records = [] + for update in range(args.updates): + cpu = torch.get_rng_state() + cuda = torch.cuda.get_rng_state_all() + python_rng, numpy_rng = random.getstate(), np.random.get_state() + reference_trainer.epoch = update + expected_loss = reference_trainer.train_epoch()["loss"] + torch.set_rng_state(cpu) + torch.cuda.set_rng_state_all(cuda) + random.setstate(python_rng) + np.random.set_state(numpy_rng) + metrics = learner.update(rigidformer_batch(raw)) + errors, gradient_errors = [], [] + for name, parameter in native.network.named_parameters(): + target = dict(reference.named_parameters())[name] + torch.testing.assert_close(parameter, target, atol=1e-5, rtol=1e-4) + if parameter.grad is not None: + torch.testing.assert_close( + parameter.grad, target.grad, atol=1e-5, rtol=1e-4 + ) + gradient_errors.append( + float((parameter.grad - target.grad).abs().max()) + ) + else: + assert target.grad is None + errors.append(float((parameter - target).abs().max())) + average.on_update(learner) + buffer_error = _assert_tree( + dict(native.network.named_buffers()), dict(reference.named_buffers()) + ) + ema_error = _assert_tree( + {k.removeprefix("network."): v for k, v in average.shadow.items()}, + reference_trainer.ema.shadow, + ) + assert average.num_updates == reference_trainer.ema.num_updates + reference_trainer.scheduler.step() + schedule.on_update(learner) + optimizer_error = _assert_tree( + _optimizer_state(native_optimizer, native), + _optimizer_state(reference_trainer.optimizer, reference), + ) + scheduler_error = _assert_tree( + schedule.scheduler.state_dict(), + reference_trainer.scheduler.state_dict(), + ) + if abs(metrics["loss"] - expected_loss) > 1e-5 + 1e-4 * abs(expected_loss): + raise AssertionError("Paired training loss differs") + records.append( + { + "update": update + 1, + "reference_loss": expected_loss, + "sl_loss": metrics["loss"], + "parameter_max_abs": max(errors), + "gradient_max_abs": max(gradient_errors, default=0.0), + "buffer_max_abs": buffer_error, + "optimizer_max_abs": optimizer_error, + "ema_max_abs": ema_error, + "scheduler_max_abs": scheduler_error, + "learning_rates": [ + group["lr"] for group in native_optimizer.param_groups + ], + } + ) + report["training"] = records + reference_trainer.writer.close() + report["checkpoint_sha256"] = sha256(args.checkpoint) + report["checkpoint_origin"] = ( + "official-release" + if report["checkpoint_sha256"] in WEIGHT_SHA256.values() + else "other-explicit-weights" + ) + report["data_manifest_sha256"] = sha256(args.data_root / "manifest.json") + report["passed"] = True + (args.output / "summary.json").write_text( + json.dumps(report, indent=2, sort_keys=True, allow_nan=False) + ) + print(json.dumps(report, indent=2, allow_nan=False)) + except BaseException as exc: + report.update(passed=False, error_type=type(exc).__name__, error=str(exc)) + (args.output / "failure.json").write_text( + json.dumps(report, indent=2, sort_keys=True) + ) + raise diff --git a/examples/learning/rigidformer/_support/_evaluator.py b/examples/learning/rigidformer/_support/_evaluator.py new file mode 100644 index 0000000000..32e1e8f256 --- /dev/null +++ b/examples/learning/rigidformer/_support/_evaluator.py @@ -0,0 +1,1336 @@ +# Copyright (c) 2026 Zhiyang Dou. SPDX-License-Identifier: MIT +# Author reference evaluator, pinned at 160fe0852c51cffae7b7440452b58d0b49db6df6. +# Import mapping only; full attribution is in the model THIRD_PARTY_LICENSES. + +""" +Rollout evaluator computing RMSE metrics on test/validation sets. + +Metrics (HopNet-aligned): +- Position RMSE: area-weighted center of mass of the predicted mesh vs. the + ground-truth trajectory, in meters. +- Rotation error: quaternion angular distance in degrees, obtained by Kabsch + alignment of predicted vertices against the canonical (frame-0) mesh. +- Checkpoints: physical frames 25, 50, 75, 100, 200, 300, 400. + +Frame convention: +- Checkpoints use direct array indexing and include the ground-truth warmup + frames ("@25 steps" = array index 25), matching HopNet. The warmup length + is configured via eval_warmup_frames. +- Checkpoints are evaluated only on the step-size grid: physical frame f is + network step f / skip_step_size. A checkpoint between two predicted frames, + or at one of the given warm-up (context) frames, is not evaluated and is + reported as "/" (step size 10: @25 and @75). + +Degenerate meshes (zero total surface area) raise a detailed error rather +than being skipped or regularized with an epsilon. +""" + +from __future__ import annotations + +import os +import logging +import json +import time +from pathlib import Path +from typing import Dict, List, Optional, Tuple +import numpy as np +import torch +import torch.distributed as dist +from torch.nn.parallel import DistributedDataParallel as DDP +from torch.utils.data import DataLoader +from tqdm import tqdm +import trimesh +import roma # For rotation matrix to quaternion conversion +from embodichain_models.dynamics.rigidformer._native.data.scene_frames import load_frame + +logger = logging.getLogger(__name__) + + +def _volume_center_of_mass(V: np.ndarray, F: np.ndarray): + """Centre of mass of closed triangle meshes V (..., Nv, 3) with faces F (Nf, 3). + + Signed-tetrahedron decomposition about the origin: identical to trimesh's + ``Trimesh.center_mass`` (the same divergence-theorem integrals), vectorised over + leading dimensions. Returns (com (..., 3), volume (...)). + """ + v0 = V[..., F[:, 0], :] + v1 = V[..., F[:, 1], :] + v2 = V[..., F[:, 2], :] + det = np.einsum("...i,...i->...", v0, np.cross(v1, v2)) # 6 x signed tetra volume + vol = det.sum(axis=-1) / 6.0 + first_moment = (det[..., None] * (v0 + v1 + v2)).sum(axis=-2) / 24.0 + return first_moment / vol[..., None], vol + + +def _hopnet_object_poses( + pred_world: np.ndarray, ref_positions_list, faces_list, Nobj: int +): + """HopNet-aligned object poses from predicted vertices (utils/rollout.py::shape_matching). + + For each object: qi = canonical vertices - canonical centre of mass, pi = predicted vertices - + predicted centre of mass, Apq = mean_i pi qi^T, R = polar(Apq) = U V^T, and the object origin + t = com_pred - R @ com_canonical. Returns (positions (T, Nobj, 3), rotations list of (T, 3, 3)). + """ + T = pred_world.shape[0] + positions = np.zeros((T, Nobj, 3), dtype=np.float64) + rotations = [] + for obj_idx in range(Nobj): + C = np.asarray( + ref_positions_list[obj_idx], dtype=np.float64 + ) # (Nv, 3) canonical + F = np.asarray(faces_list[obj_idx]) + Nv = C.shape[0] + P = np.asarray(pred_world[:, obj_idx, :Nv, :], dtype=np.float64) # (T, Nv, 3) + com_c, vol_c = _volume_center_of_mass(C, F) + if not np.isfinite(vol_c) or abs(vol_c) < 1e-12: + raise ValueError( + f"HopNet position estimator: object {obj_idx} has a degenerate volume ({vol_c}); " + f"the collision mesh must be closed." + ) + com_w, _ = _volume_center_of_mass(P, F) # (T, 3) + q = C - com_c # (Nv, 3) + pmat = P - com_w[:, None, :] # (T, Nv, 3) + Apq = np.einsum("tni,nj->tij", pmat, q) / Nv # (T, 3, 3) + U, _, Vt = np.linalg.svd(Apq) + R = U @ Vt # polar decomposition (as HopNet) + positions[:, obj_idx, :] = com_w - np.einsum("tij,j->ti", R, com_c) + rotations.append(R) + return positions.astype(np.float32), rotations + + +class Evaluator: + """ + Evaluator for computing RMSE metrics on test/validation sets. + + Key features: + - Batch-wise evaluation (no drop_last, processes all samples) + - HopNet-aligned checkpoints (@25, @50, etc.) + - TensorBoard logging + - Cached dataset indexing (reuse across epochs) + """ + + def __init__(self, trainer): + """ + Initialize evaluator with reference to trainer. + + Args: + trainer: Main trainer instance (provides model, config, device, etc.) + """ + self.trainer = trainer + self.config = trainer.config + # data.object_scale_file: the canonical frame-0 meshes read below must carry the same scale as the model inputs + from embodichain_models.dynamics.rigidformer._native.data.object_scale import ( + from_config as _object_scale_from_config, + ) + + self._object_scale = _object_scale_from_config(self.config) + # data.partial_view: the inputs are partial point clouds of a simulated depth camera (src/data/partial_view.py); + # each observed object's rigid motion moves its full collision mesh before scoring, unobserved objects are + # left out of the metrics + _pv = self.config["data"].get("partial_view") or {} + self._partial_view = dict(_pv) if _pv.get("enable", False) else None + self._partial_counts = ( + 0, + 0, + ) # observed objects, all objects (last evaluated split and step size) + # Unwrap DDP wrapper to access custom methods (e.g., rollout) + self.model = ( + trainer.model.module if isinstance(trainer.model, DDP) else trainer.model + ) + self.device = trainer.device + self.writer = trainer.writer + self.log_dir = trainer.log_dir + self.epoch = 0 # Updated by trainer before calling evaluate + + # Metrics configuration + metrics_cfg = self.config["training"].get("metrics_test", {}) + self.metrics_enabled = metrics_cfg.get("enable", False) + self.print_every_n_epochs = metrics_cfg.get("print_every_n_epochs", 5) + + # Evaluation settings + self.eval_num_frames = metrics_cfg.get("eval_num_frames", 480) + self.eval_warmup_frames = metrics_cfg.get( + "eval_warmup_frames", 2 + ) # Ground-truth frames before autoregressive rollout + self.hopnet_checkpoints_physical = metrics_cfg.get( + "hopnet_checkpoints", [25, 50, 75, 100, 200, 300, 400] + ) + # Step sizes > 1 start from the model's own step-size-1 rollout instead of ground-truth warm-up frames + # (_warmup_from_step1): every step size then sees only the ground truth of the step-size-1 protocol + self.warmup_from_step1 = bool(metrics_cfg.get("warmup_from_step1", False)) + if self.warmup_from_step1 and self._partial_view is not None: + raise ValueError( + "metrics_test.warmup_from_step1 is not defined for partial-view inputs (the observation " + "frame differs between step sizes)" + ) + # Object-pose extraction from predicted vertices. + # 'area_centroid': area-weighted surface centroid vs. the ground-truth origin. + # 'hopnet': HopNet's shape matching: volume centre of mass + polar-decomposition + # rotation, object origin = com_pred - R @ com_canonical (exact for a perfect + # prediction, no constant offset). + self.position_estimator = str( + metrics_cfg.get("position_estimator", "area_centroid") + ) + if self.position_estimator not in ("area_centroid", "hopnet"): + raise ValueError( + f"metrics_test.position_estimator must be 'area_centroid' or 'hopnet'; got {self.position_estimator}" + ) + + # Cached datasets (per skip_step_size) + self._test_eval_datasets = {} + self._val_eval_datasets = {} + + # DDP info + if dist.is_available() and dist.is_initialized(): + self.rank = dist.get_rank() + self.world_size = dist.get_world_size() + else: + self.rank = 0 + self.world_size = 1 + + def _protocol_text(self, split: str, n_scenes: int) -> str: + """What every reported number of this evaluation depends on besides the checkpoint. A split file derived for + evaluating another partition through the test list (scripts/eval_cross_dataset.py --split val) names that + partition in 'evaluated_partition'.""" + d = self.config["data"] + with open(d["split_file"]) as f: + partition = json.load(f).get("evaluated_partition", split) + text = ( + f"position estimator {self.position_estimator}; scene set {d.get('scene_set', 'benchmark')} " + f"({os.path.basename(str(d.get('split_file')))}), {n_scenes} {partition} scenes; " + f"rotation error = RMS over objects, rotation MAE = mean absolute angle over objects" + ) + if self.warmup_from_step1: + text += ( + f"; step sizes > 1: warm-up frames s .. {self.eval_warmup_frames - 1}s predicted by the model at " + f"step size 1 from ground-truth frames 0-{self.eval_warmup_frames - 1} (only frame 0 is ground " + f"truth at step size s)" + ) + if self._partial_view is not None: + from src.data.partial_view import PartialView + + pv = PartialView( + {**self._partial_view, "observe_index": self.eval_warmup_frames - 1}, + Path(__file__).resolve().parents[2], + ) + seen, total = self._partial_counts + text += ( + f"; inputs: {pv.describe(train=False)}; metrics over the {seen} observed of {total} objects " + f"(full collision meshes moved by the predicted rigid motions)" + ) + return text + + def should_evaluate(self, epoch: int) -> bool: + """Check if RMSE evaluation should run this epoch.""" + if not self.metrics_enabled: + return False + # Evaluate at epoch 0 (initial baseline) and every N epochs thereafter + return (epoch % self.print_every_n_epochs) == 0 + + def evaluate_rollout_rmse(self, epoch: int): + """ + Main entry point for RMSE evaluation. + + Evaluates on validation and/or test sets based on config. + Multi-timescale aware: evaluates at each step_size in eval_step_sizes. + + Called by trainer every N epochs. + + Args: + epoch: Current training epoch + """ + if not self.should_evaluate(epoch): + return + + self.epoch = epoch + if self.rank == 0: + logger.info(f"{'='*80}") + logger.info(f"[RMSE Evaluation] Epoch {epoch}") + logger.info(f"{'='*80}") + + # Get split file + split_file = self.config["data"].get("split_file", "dataset_split.json") + if not os.path.exists(split_file): + if self.rank == 0: + logger.warning( + f"Split file not found: {split_file}, skipping RMSE evaluation" + ) + return + + with open(split_file, "r") as f: + split_info = json.load(f) + + val_datasets = split_info.get("val_datasets", []) + test_datasets = split_info.get("test_datasets", []) + + # Multi-timescale evaluation: loop over the configured eval_step_sizes, + # falling back to the single skip_step_size when unset + eval_step_sizes = self.config["data"].get( + "eval_step_sizes", [self.config["data"].get("skip_step_size", 1)] + ) + + if self.rank == 0: + logger.info( + f"[RMSE Evaluation] Evaluating at step_sizes: {eval_step_sizes}" + ) + + all_test_metrics = {} # Aggregate test metrics across all step_sizes + + for step_size in eval_step_sizes: + if self.rank == 0: + logger.info(f"\n{'='*80}") + logger.info(f"[RMSE Evaluation] Evaluating at step_size = {step_size}") + logger.info(f"{'='*80}") + + # Evaluate validation set at this step_size + enable_val = self.config["data"].get("enable_val_split", False) + if enable_val and len(val_datasets) > 0: + if self.rank == 0: + logger.info( + f"[RMSE Evaluation] Validation set: {len(val_datasets)} datasets (step_size={step_size})" + ) + self._evaluate_rmse_on_split( + val_datasets, split="val", step_size=step_size + ) + + # Evaluate test set at this step_size (always evaluate if exists) + if len(test_datasets) > 0: + if self.rank == 0: + logger.info( + f"[RMSE Evaluation] Test set: {len(test_datasets)} datasets (step_size={step_size})" + ) + test_metrics_ss = self._evaluate_rmse_on_split( + test_datasets, split="test", step_size=step_size + ) + + # Aggregate test metrics (already include step_size suffix from _evaluate_rmse_on_split) + for metric_name, value in test_metrics_ss.items(): + all_test_metrics[metric_name] = value + + if self.rank == 0: + logger.info( + f"\n[RMSE Evaluation] Epoch {epoch} completed for all step_sizes" + ) + + # Return aggregated test metrics for best checkpoint selection + return all_test_metrics + + def _evaluate_rmse_on_split( + self, dataset_names: List[str], split: str, step_size: int + ): + """ + Evaluate RMSE on a split (val or test) using batched processing. + + DDP-aware: Each rank processes a shard of data, then Rank 0 gathers results. + + Args: + dataset_names: List of dataset names to evaluate + split: 'val' or 'test' + step_size: Step size for frame sampling (1, 2, 5, 10, etc.) + """ + from embodichain_models.dynamics.rigidformer._native.data.dataset import ( + MultiObjectDynamics, + make_collate_fn, + ) + import torch.distributed as dist + + skip_step_size = step_size + + # Calculate sequence length for evaluation (network steps) + eval_sequence_length = self.eval_num_frames // skip_step_size + + # Only Rank 0 logs dataset creation info + if self.rank == 0: + logger.info( + f"[RMSE Dataset] Creating {split} dataset: " + f"physical_frames={self.eval_num_frames}, skip_step={skip_step_size}, " + f"sequence_length={eval_sequence_length}" + ) + + # Display the checkpoints: evaluated on the step-size grid after the warm-up frames, "/" otherwise + logger.info( + f"\n[Checkpoints] skip_step_size={skip_step_size}, warm-up frames 0 to {self.eval_warmup_frames - 1}" + ) + logger.info( + f"{'Checkpoint':^12} | {'Network Step':^13} | {'Rollout predictions':^20}" + ) + logger.info("-" * 52) + for physical_cp in self.hopnet_checkpoints_physical: + if physical_cp % skip_step_size: + logger.info( + f"{physical_cp:^12} | {'-':^13} | {'/ (not on the grid)':^20}" + ) + elif physical_cp // skip_step_size < self.eval_warmup_frames: + logger.info( + f"{physical_cp:^12} | {physical_cp // skip_step_size:^13} | {'/ (warm-up frame)':^20}" + ) + else: + network_step = physical_cp // skip_step_size + logger.info( + f"{physical_cp:^12} | {network_step:^13} | {network_step - self.eval_warmup_frames + 1:^20}" + ) + logger.info("") + + eval_dataset = self._get_eval_dataset( + split, skip_step_size, eval_sequence_length + ) + batch_size = self.config["training"].get("eval_batch_size", 8) + eval_loader = self._eval_loader(eval_dataset, batch_size) + # metrics_test.warmup_from_step1: for a step size s > 1 the warm-up frames s, 2s, ... are the model's own + # step-size-1 rollout from the step-size-1 warm-up (ground-truth frames 0 .. warm-up - 1), instead of + # ground-truth frames; every step size then uses exactly the ground truth of the step-size-1 protocol + step1_loader = None + if self.warmup_from_step1 and skip_step_size > 1: + step1_loader = self._eval_loader( + self._get_eval_dataset(split, 1, self.eval_num_frames), batch_size + ) + if len(step1_loader) != len(eval_loader): + raise ValueError( + f"warmup_from_step1: {len(step1_loader)} step-size-1 batches vs {len(eval_loader)}" + ) + + if self.rank == 0: + logger.info( + f"[RMSE Dataset] {split} set: {len(eval_dataset)} samples, " + f"batch_size={batch_size}, {len(eval_loader)} batches" + + ( + ", warm-up from the step-size-1 rollout" + if step1_loader is not None + else "" + ) + ) + + return self._evaluate_rmse_loop( + eval_dataset, + eval_loader, + step1_loader, + split, + skip_step_size, + eval_sequence_length, + ) + + def _get_eval_dataset( + self, split: str, skip_step_size: int, eval_sequence_length: int + ): + """The cached evaluation dataset of a split at a step size (one sample per scene, warm-up frames only).""" + from embodichain_models.dynamics.rigidformer._native.data.dataset import ( + MultiObjectDynamics, + ) + + # Create or reuse cached dataset + cache_key = f"_eval_dataset_{split}_step{skip_step_size}" + if cache_key not in ( + self._test_eval_datasets if split == "test" else self._val_eval_datasets + ): + if self.rank == 0: + logger.info(f"[RMSE Dataset] Creating new {split} dataset (first time)") + + # Load only the warmup frames; rollout() generates the remaining + # predictions autoregressively. + eval_load_frames = max( + self.eval_warmup_frames, 8 + ) # At least 8 frames for safety margin + + if self.rank == 0: + logger.info( + f"[RMSE Eval] Loading only {eval_load_frames} frames (not {eval_sequence_length})" + ) + + eval_dataset = MultiObjectDynamics( + data_path=self.config["data"]["dataset_path"], + sequence_length=eval_load_frames, # Load only the warmup frames + stride=999999, # Large stride yields one sample per dataset (no sliding windows) + skip_step_size=skip_step_size, + split=split, + split_file=self.config["data"].get("split_file", "dataset_split.json"), + max_vertices_per_object=self.config["data"].get( + "max_vertices_per_object", 1150 + ), + max_objects_per_frame=self.config["data"].get( + "max_objects_per_frame", 20 + ), + anchor_k=self.config["data"].get("anchor_k", 32), + if_ground_mesh=self.config["data"].get("if_ground_mesh", False), + ground_mesh_vertex_num=self.config["data"].get( + "ground_mesh_vertex_num", 968 + ), + if_debug_normalization_log=False, # Disable verbose logging during evaluation + random_step_size={ + "enable": False + }, # Disable random sampling for evaluation + augmentation={"enable": False}, # Disable augmentation for evaluation + anchor_fps_start=self.config["data"].get("anchor_fps_start", "index0"), + fps_order=self.config["data"].get("fps_order", False), + object_scale_file=self.config["data"].get("object_scale_file"), + # the same physics-parameter validation as the training data (missing entries are errors) + if_phys_param=self.config.get("model", {}).get("if_phys_param", False), + # partial inputs observed at the last warm-up frame + partial_view=( + {**self._partial_view, "observe_index": self.eval_warmup_frames - 1} + if self._partial_view is not None + else None + ), + ) + + # Cache the dataset + if split == "test": + self._test_eval_datasets[cache_key] = eval_dataset + else: + self._val_eval_datasets[cache_key] = eval_dataset + else: + if self.rank == 0: + logger.info(f"[RMSE Dataset] Reusing cached {split} dataset") + eval_dataset = ( + self._test_eval_datasets if split == "test" else self._val_eval_datasets + )[cache_key] + return eval_dataset + + def _eval_loader(self, eval_dataset, batch_size: int): + """Evaluation DataLoader: scenes in dataset order, every scene on exactly one rank.""" + from embodichain_models.dynamics.rigidformer._native.data.dataset import ( + make_collate_fn, + ) + + # Create DataLoader with DDP-aware sampler + collate_fn = make_collate_fn( + max_objects_per_frame=self.config["data"].get("max_objects_per_frame", 20), + max_vertices_per_object=self.config["data"].get( + "max_vertices_per_object", 1150 + ), + ) + + # DDP-aware: every scene on exactly one rank (strided shards without the padding a DistributedSampler adds, + # which would evaluate some scenes twice whenever the scene count is not a multiple of the rank count). The + # rollout contains no collective, so ranks may process different numbers of scenes. + if self.world_size > 1: + eval_loader = DataLoader( + torch.utils.data.Subset( + eval_dataset, + list(range(self.rank, len(eval_dataset), self.world_size)), + ), + batch_size=batch_size, + shuffle=False, + num_workers=0, # Avoid DDP deadlock: use main process only for evaluation + pin_memory=True, + drop_last=False, + collate_fn=collate_fn, + ) + else: + # Single GPU: no sampler needed + eval_loader = DataLoader( + eval_dataset, + batch_size=batch_size, + shuffle=False, + num_workers=0, # Single process loading for evaluation + pin_memory=True, + drop_last=False, + collate_fn=collate_fn, + ) + return eval_loader + + def _warmup_from_step1( + self, batch: dict, batch1: dict, skip_step_size: int + ) -> dict: + """Replace the warm-up frames s, 2s, ... of a step-size-s batch by the model's step-size-1 rollout of the same + scenes, started from the step-size-1 warm-up (ground-truth frames 0 .. warm-up - 1). + """ + W = self.eval_warmup_frames + if list(batch.get("dataset_name", [])) != list(batch1.get("dataset_name", [])): + raise ValueError( + f"warmup_from_step1: scene order differs between the step-size-{skip_step_size} and " + f"step-size-1 evaluation batches" + ) + frames = [k * skip_step_size for k in range(1, W)] + out = self.model.rollout( + initial_batch=batch1, + num_steps=frames[-1] + 1, + dt=1.0, + return_extras=False, + warmup_frames=W, + ) + pos = batch["obj_positions"].clone() + pos[:, 1:W] = out["obj_positions"][:, frames].to(pos.dtype) + return {**batch, "obj_positions": pos} + + def _evaluate_rmse_loop( + self, + eval_dataset, + eval_loader, + step1_loader, + split: str, + skip_step_size: int, + eval_sequence_length: int, + ): + import torch.distributed as dist + + step1_batches = iter(step1_loader) if step1_loader is not None else None + + # Synchronize all ranks before starting evaluation loop (avoid DDP deadlock) + if self.world_size > 1: + if self.rank == 0: + import sys + + print(f"[RMSE Eval] Rank 0: Entering barrier sync...", flush=True) + sys.stdout.flush() + dist.barrier() + if self.rank == 0: + print( + f"[RMSE Eval] Rank 0: Barrier sync complete, starting eval loop", + flush=True, + ) + + # Evaluation loop + self.model.eval() + + all_position_rmse = {} # checkpoint -> list of RMSE values + all_rotation_error = ( + {} + ) # checkpoint -> list of rotation errors (per-scene RMS over objects) + all_rotation_mae = ( + {} + ) # checkpoint -> list of per-scene mean absolute rotation errors over objects + checkpoint_mapping = ( + {} + ) # checkpoint -> (evaluated physical frame, warm-up frames) when not (cp, main warm-up) + not_evaluated = ( + {} + ) # checkpoint -> reason ("/"): off the step-size grid or at a warm-up (context) frame + total_samples = 0 + partial_seen, partial_total = ( + 0, + 0, + ) # data.partial_view: observed objects, all objects + + # Flash Attention handles dtype conversion inside the decoder; the rest + # of the model runs in FP32, so no global autocast is needed. + with torch.no_grad(): + import sys + + if self.rank == 0: + print(f"[RMSE Eval] Creating tqdm progress bar...", flush=True) + + pbar = tqdm( + eval_loader, desc=f"RMSE Eval ({split})", disable=(self.rank != 0) + ) + + if self.rank == 0: + print( + f"[RMSE Eval] tqdm created, entering for loop (num_batches={len(eval_loader)})...", + flush=True, + ) + sys.stdout.flush() + + for batch_idx, batch in enumerate(pbar): + import sys + + if self.rank == 0: + print( + f"[RMSE Eval] Batch {batch_idx+1}/{len(eval_loader)} starting...", + flush=True, + ) + sys.stdout.flush() + + # Move batch to device + batch = { + k: ( + v.to(self.device, non_blocking=True) + if isinstance(v, torch.Tensor) + else v + ) + for k, v in batch.items() + } + if step1_batches is not None: + batch1 = { + k: ( + v.to(self.device, non_blocking=True) + if isinstance(v, torch.Tensor) + else v + ) + for k, v in next(step1_batches).items() + } + batch = self._warmup_from_step1(batch, batch1, skip_step_size) + + # Get dataset names for this batch + batch_dataset_names = batch.get("dataset_name", None) + if batch_dataset_names is None: + logger.warning( + f"[RMSE] Batch {batch_idx} missing 'dataset_name', skipping" + ) + continue + + # Autoregressive rollout; return_extras=False skips storing + # intermediate tensors to reduce memory. + num_frames = eval_sequence_length + if self.rank == 0: + logger.info( + f"[RMSE Eval] Batch {batch_idx+1}: Starting rollout ({num_frames} frames, warmup={self.eval_warmup_frames})..." + ) + + import time + + rollout_start = time.time() + rollout_out = self.model.rollout( + initial_batch=batch, + num_steps=num_frames, + dt=1.0, + return_extras=False, + warmup_frames=self.eval_warmup_frames, + ) + rollout_time = time.time() - rollout_start + + if self.rank == 0: + logger.info( + f"[RMSE Eval] Batch {batch_idx+1}: Rollout done in {rollout_time:.1f}s" + ) + + # Extract predictions + pred_full = ( + rollout_out["obj_positions"].cpu().numpy() + ) # (B, T, Nobj, Nv, 3) + B_actual = pred_full.shape[0] + + # Process each sample in batch + for b_idx in range(B_actual): + dataset_name = ( + batch_dataset_names[b_idx] + if isinstance(batch_dataset_names, list) + else batch_dataset_names[b_idx].item() + ) + pred_single = pred_full[b_idx] # (T, Nobj, Nv, 3) + object_mask = None + if self._partial_view is not None: + pred_single, object_mask = self._full_from_partial( + pred_single, batch, b_idx + ) + partial_seen += int(object_mask.sum()) + partial_total += int(object_mask.size) + + # Compute RMSE for this sample + rmse_dict = self._compute_rmse_for_sample( + pred_single, + dataset_name, + eval_dataset, + skip_step_size, + self.hopnet_checkpoints_physical, + object_mask=object_mask, + ) + + # Aggregate results (position RMSE and rotation error) + for checkpoint, value in rmse_dict.items(): + if checkpoint == "step_size": + continue + + # A checkpoint off the step-size grid or at a warm-up (context) frame: not evaluated + if isinstance(checkpoint, str) and checkpoint.startswith( + "_na_" + ): + not_evaluated[int(checkpoint[len("_na_") :])] = value + continue + + # Mapping metadata: the frame / warm-up actually evaluated for a checkpoint + if isinstance(checkpoint, str) and checkpoint.startswith( + "_mapping_" + ): + checkpoint_mapping[int(checkpoint[len("_mapping_") :])] = ( + tuple(value) + ) + continue + + if isinstance(checkpoint, str) and checkpoint.startswith( + "rotation_mae_deg_" + ): + all_rotation_mae.setdefault( + int(checkpoint[len("rotation_mae_deg_") :]), [] + ).append(value) + continue + + # Check if this is a rotation error key (starts with 'rotation_deg_') + if isinstance(checkpoint, str) and checkpoint.startswith( + "rotation_deg_" + ): + # Extract physical frame number from key + physical_frame = int( + checkpoint.replace("rotation_deg_", "") + ) + if physical_frame not in all_rotation_error: + all_rotation_error[physical_frame] = [] + all_rotation_error[physical_frame].append(value) + else: + # Position RMSE, keyed by target checkpoint + if checkpoint not in all_position_rmse: + all_position_rmse[checkpoint] = [] + all_position_rmse[checkpoint].append(value) + + total_samples += 1 + + # Update progress bar + pbar.set_postfix({"samples": total_samples}) + + # Synchronize all ranks before gathering results (ensure all finished evaluation) + if self.world_size > 1: + dist.barrier() + + # DDP: Gather results from all ranks to Rank 0 + if self.world_size > 1: + # Convert dict of lists to list of dicts for gathering + position_rmse_list = [ + {"checkpoint": k, "values": v} for k, v in all_position_rmse.items() + ] + rotation_error_list = [ + {"checkpoint": k, "values": v} for k, v in all_rotation_error.items() + ] + + # Gather all results to Rank 0 + gathered_position = [None] * self.world_size + gathered_rotation = [None] * self.world_size + gathered_extra = [None] * self.world_size + + dist.all_gather_object(gathered_position, position_rmse_list) + dist.all_gather_object(gathered_rotation, rotation_error_list) + dist.all_gather_object( + gathered_extra, + { + "rotation_mae": all_rotation_mae, + "mapping": checkpoint_mapping, + "not_evaluated": not_evaluated, + "samples": total_samples, + "partial": (partial_seen, partial_total), + }, + ) + + if self.rank == 0: + # Merge results from all ranks + all_position_rmse = {} + all_rotation_error = {} + total_samples = 0 + + for rank_position in gathered_position: + for item in rank_position: + checkpoint = item["checkpoint"] + if checkpoint not in all_position_rmse: + all_position_rmse[checkpoint] = [] + all_position_rmse[checkpoint].extend(item["values"]) + + for rank_rotation in gathered_rotation: + for item in rank_rotation: + checkpoint = item["checkpoint"] + if checkpoint not in all_rotation_error: + all_rotation_error[checkpoint] = [] + all_rotation_error[checkpoint].extend(item["values"]) + + all_rotation_mae = {} + for _ex in gathered_extra: + for checkpoint, vals in _ex["rotation_mae"].items(): + all_rotation_mae.setdefault(checkpoint, []).extend(vals) + checkpoint_mapping.update(_ex["mapping"]) + not_evaluated.update(_ex["not_evaluated"]) + + # Count total samples across all ranks + total_samples = sum(_ex["samples"] for _ex in gathered_extra) + partial_seen = sum(_ex["partial"][0] for _ex in gathered_extra) + partial_total = sum(_ex["partial"][1] for _ex in gathered_extra) + self._partial_counts = (partial_seen, partial_total) + + # Only Rank 0 logs and writes to TensorBoard + if self.rank == 0: + if total_samples != len(eval_dataset): + raise RuntimeError( + f"[RMSE Eval] evaluated {total_samples} scenes, the {split} set has {len(eval_dataset)}" + ) + # Aggregate and log results + logger.info(f"[RMSE Results] {split} set ({total_samples} samples):") + + # Log position RMSE + logger.info(f"\n[Position RMSE] (actual physical frames evaluated)") + for physical_frame in sorted(all_position_rmse.keys()): + rmse_values = all_position_rmse[physical_frame] + mean_rmse = np.mean(rmse_values) + + logger.info( + f" Pos RMSE @ frame {physical_frame:3d}: {mean_rmse:.6f} m" + ) + + # Keyed by target checkpoint so all step sizes align in TensorBoard + self.writer.add_scalar( + f"metrics/p_rmse_@{physical_frame:03d}/s={skip_step_size}", + mean_rmse, + self.epoch, + ) + + # Log rotation error + if len(all_rotation_error) > 0: + logger.info(f"\n[Rotation Error] (actual physical frames evaluated)") + for physical_frame in sorted(all_rotation_error.keys()): + rot_values = all_rotation_error[physical_frame] + mean_rot = np.mean(rot_values) + + logger.info( + f" Rot Error @ frame {physical_frame:3d}: {mean_rot:.4f}°" + ) + + # Keyed by target checkpoint so all step sizes align in TensorBoard + self.writer.add_scalar( + f"metrics/r_error_@{physical_frame:03d}/s={skip_step_size}", + mean_rot, + self.epoch, + ) + # Rotation error as the mean absolute angle over objects (the statistic HopNet's reporting code prints; + # "Rot Error" above is the RMS over objects) + if len(all_rotation_mae) > 0: + logger.info( + f"\n[Rotation MAE] (mean absolute angle over objects, per scene; averaged over scenes)" + ) + for physical_frame in sorted(all_rotation_mae.keys()): + mean_mae = np.mean(all_rotation_mae[physical_frame]) + logger.info( + f" Rot MAE @ frame {physical_frame:3d}: {mean_mae:.4f}°" + ) + self.writer.add_scalar( + f"metrics/r_mae_@{physical_frame:03d}/s={skip_step_size}", + mean_mae, + self.epoch, + ) + _mapping_text = ", ".join( + f"@{cp} -> frame {fr} (warm-up {wu})" + for cp, (fr, wu) in sorted(checkpoint_mapping.items()) + ) + _na_text = ", ".join( + f"@{cp} ({why})" for cp, why in sorted(not_evaluated.items()) + ) + logger.info( + f"[Checkpoint mapping] step_size={skip_step_size}, rollout warm-up {self.eval_warmup_frames} frames; " + f"checkpoints evaluated elsewhere: {_mapping_text or 'none'}; " + f"not evaluated (/): {_na_text or 'none'}" + ) + logger.info(f"[Protocol] {self._protocol_text(split, total_samples)}") + + logger.info(f"\n[RMSE Results] {split} set completed") + + # Append results to eval.txt for post-hoc analysis + eval_txt_path = os.path.join(self.log_dir, "eval.txt") + with open(eval_txt_path, "a") as f: + import datetime + + f.write(f"\n{'='*80}\n") + f.write( + f"Epoch {self.epoch} | {split} | step_size={skip_step_size} | {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n" + ) + f.write(f"{'='*80}\n") + + def _rows(values, fmt, unit): + # one row per checkpoint, "/" for those not evaluated + for physical_frame in sorted(set(values) | set(not_evaluated)): + if physical_frame in values: + f.write( + f" frame {physical_frame:3d}: {np.mean(values[physical_frame]):{fmt}}{unit}\n" + ) + else: + f.write(f" frame {physical_frame:3d}: /\n") + + f.write(f"[Position RMSE]\n") + _rows(all_position_rmse, ".6f", " m") + if len(all_rotation_error) > 0: + f.write(f"[Rotation Error]\n") + _rows(all_rotation_error, ".4f", "°") + if len(all_rotation_mae) > 0: + f.write(f"[Rotation MAE]\n") + _rows(all_rotation_mae, ".4f", "°") + f.write( + f" mapping: warm-up {self.eval_warmup_frames}; " + + ( + ", ".join( + [ + f"@{cp} -> frame {fr} (warm-up {wu})" + for cp, (fr, wu) in sorted(checkpoint_mapping.items()) + ] + + [ + f"@{cp} / ({why})" + for cp, why in sorted(not_evaluated.items()) + ] + ) + or "all checkpoints at their own frame" + ) + + "\n" + ) + f.write(f" protocol: {self._protocol_text(split, total_samples)}\n") + f.write(f"\n") + logger.info(f"[Eval] Results appended to {eval_txt_path}") + + # Aggregated metrics for best-checkpoint selection, named to match + # the TensorBoard tags + metrics = {} + for physical_frame in sorted(all_position_rmse.keys()): + rmse_values = all_position_rmse[physical_frame] + mean_rmse = np.mean(rmse_values) + metrics[f"p_rmse_@{physical_frame:03d}/s={skip_step_size}"] = mean_rmse + + for physical_frame in sorted(all_rotation_error.keys()): + rot_values = all_rotation_error[physical_frame] + mean_rot = np.mean(rot_values) + metrics[f"r_error_@{physical_frame:03d}/s={skip_step_size}"] = mean_rot + for physical_frame in sorted(all_rotation_mae.keys()): + metrics[f"r_mae_@{physical_frame:03d}/s={skip_step_size}"] = np.mean( + all_rotation_mae[physical_frame] + ) + + return metrics + else: + # Non-rank-0: return empty dict + return {} + + def _full_from_partial(self, pred_single: np.ndarray, batch: dict, b_idx: int): + """Full-mesh prediction of a rollout from partial point clouds (data.partial_view). + + The rollout moves each observed object's points rigidly; Kabsch from the points at the observation frame (the + last warm-up frame, ground truth) to the points of every frame recovers that motion, which moves the object's + full collision mesh of the observation frame. Unobserved objects keep their observation-frame mesh as a + placeholder and are masked out of the metrics. Returns ((T, Nobj, Nv, 3), observed mask (Nobj,)). + """ + from src.data.partial_view import kabsch + + observed = np.asarray(batch["partial_obj_index"][b_idx], dtype=np.int64) + k = int(batch["partial_obs_index"][b_idx]) + full = batch["partial_full_obs"][b_idx] # (Nv_i, 3) per object of the scene + vmask = batch["valid_vtx_mask"][b_idx].detach().cpu().numpy() # (T0, Nobs, Np) + T = pred_single.shape[0] + nv = max(V.shape[0] for V in full) + out = np.zeros((T, len(full), nv, 3), dtype=np.float32) + for i, V in enumerate(full): + out[:, i, : V.shape[0]] = V + for j, i in enumerate(observed): + m = vmask[k, j] + R, t = kabsch(pred_single[k, j, m], pred_single[:, j, m]) + out[:, i, : full[i].shape[0]] = ( + np.einsum("tab,nb->tna", R, full[i].astype(np.float64)) + t[:, None, :] + ).astype(np.float32) + mask = np.zeros((len(full),), dtype=bool) + mask[observed] = True + return out, mask + + def _compute_rmse_for_sample( + self, + pred_single: np.ndarray, + dataset_name: str, + eval_dataset, + skip_step_size: int, + physical_checkpoints: List[int], + object_mask: Optional[np.ndarray] = None, + ) -> Dict[int, float]: + """ + Compute position RMSE and rotation error for one sample at the checkpoints. + + Args: + pred_single: Predicted positions (T, Nobj, Nv, 3) with the main warmup + dataset_name: Name of the dataset + eval_dataset: Dataset instance for normalization + skip_step_size: Frame subsampling factor + physical_checkpoints: Physical frame checkpoints (e.g., [25, 50, 75, ...]) + + Returns: + Dict keyed by target checkpoint (position RMSE), with additional + 'rotation_deg_' entries and '_na_' (the reason) for a + checkpoint that is not evaluated: off the step-size grid, or at a + warm-up (context) frame. + """ + # Load GT trajectory + dataset_path = Path(self.config["data"]["dataset_path"]) / dataset_name + traj_path = dataset_path / "vis_traj.npy" + + if not traj_path.exists(): + raise FileNotFoundError(f"Missing trajectory file: {traj_path}") + + traj_full = np.load(traj_path)[0] # (T_total, Nobj, 7) - remove batch dim + + # Subsample trajectory to match skip_step_size + traj_sampled = traj_full[::skip_step_size] # (T_sampled, Nobj, 7) + + T_eval = min(pred_single.shape[0], traj_sampled.shape[0]) + gt_com = traj_sampled[:T_eval, :, :3].astype(np.float32) # (T_eval, Nobj, 3) + + # Get object count + Nobj_gt = gt_com.shape[1] + + # Trim predictions to match GT object count + if pred_single.shape[1] > Nobj_gt: + pred_single = pred_single[:T_eval, :Nobj_gt] + else: + pred_single = pred_single[:T_eval] + + # Predictions are already in world coordinates (DISABLE_NORMALIZERS=1) + pred_world = pred_single # (T_eval, Nobj, Nv, 3) + + # Load faces for area-weighted COM computation + faces_list = [] + meshes_dir = dataset_path / "meshes" + for obj_idx in range(Nobj_gt): + faces_file = meshes_dir / f"obj_{obj_idx}_faces.npy" + if not faces_file.exists(): + raise FileNotFoundError(f"Missing faces file: {faces_file}") + faces = np.load(faces_file) + faces_list.append(faces) + + R_hopnet = None + if self.position_estimator == "hopnet": + # Canonical meshes (object frame, from the ground-truth frame-0 pose) for shape matching + _gt_quat_eval = traj_sampled[:T_eval, :Nobj_gt, 3:7].astype(np.float32) + _ref_list = [] + _frame0 = load_frame( + dataset_path, 0 + ) # either layout (src/data/scene_frames.py) + for obj_idx in range(Nobj_gt): + _p0 = dataset_path / "frame_00000" / f"obj_{obj_idx}_pos.npy" + _pos0 = _frame0[obj_idx] + if self._object_scale is not None: + _pos0 = self._object_scale.apply(_p0, _pos0) + _q0 = _gt_quat_eval[0, obj_idx] + _q0_xyzw = ( + torch.tensor([_q0[1], _q0[2], _q0[3], _q0[0]]).float().unsqueeze(0) + ) + _R0 = roma.unitquat_to_rotmat(_q0_xyzw).squeeze(0).numpy() + _ref_list.append((_pos0 - gt_com[0, obj_idx]) @ _R0) + pred_com, R_hopnet = _hopnet_object_poses( + pred_world, _ref_list, faces_list, Nobj_gt + ) + else: + # Vectorized area-weighted COM across all timesteps (avoids building a + # trimesh object per frame and per object) + pred_com = np.zeros((T_eval, Nobj_gt, 3), dtype=np.float32) + for obj_idx in range(Nobj_gt): + F = faces_list[obj_idx] # (Nf, 3) + # Process all timesteps for this object at once + V_all_t = pred_world[:, obj_idx, :, :] # (T_eval, Nv, 3) + + # Vectorized COM calculation across all timesteps + # Get triangle vertices: (T, Nf, 3, 3) + v0 = V_all_t[:, F[:, 0], :] # (T, Nf, 3) + v1 = V_all_t[:, F[:, 1], :] # (T, Nf, 3) + v2 = V_all_t[:, F[:, 2], :] # (T, Nf, 3) + + # Triangle areas: 0.5 * ||cross(v1-v0, v2-v0)|| + cross = np.cross(v1 - v0, v2 - v0) # (T, Nf, 3) + areas = 0.5 * np.linalg.norm(cross, axis=-1) # (T, Nf) + + # Triangle centroids: (v0 + v1 + v2) / 3 + centroids = (v0 + v1 + v2) / 3.0 # (T, Nf, 3) + + # Area-weighted COM: sum(area * centroid) / sum(area) + total_area = np.sum(areas, axis=1, keepdims=True) # (T, 1) + + # Degenerate mesh is a hard error + if np.any(total_area <= 0): + zero_frames = np.where(total_area.squeeze() <= 0)[0] + raise ValueError( + f"Degenerate mesh detected for object {obj_idx} in dataset {dataset_name}!\n" + f"Total area is zero or negative at frames: {zero_frames.tolist()}\n" + f"This indicates collapsed triangles or invalid mesh geometry.\n" + f"Check mesh file: meshes/obj_{obj_idx}_faces.npy" + ) + + weighted_sum = np.sum( + areas[:, :, np.newaxis] * centroids, axis=1 + ) # (T, 3) + pred_com[:, obj_idx, :] = weighted_sum / total_area # (T, 3) + + # Compute position error + pos_err = np.sum((pred_com - gt_com) ** 2, axis=-1) # (T_eval, Nobj) + + # Compute per-step RMSE (data.partial_view: over the observed objects) + if object_mask is not None: + pos_err = pos_err[:, np.asarray(object_mask, dtype=bool)] + position_rmse_all_steps = [] + for step in range(T_eval): + mse_step = np.mean(pos_err[step]) # Average over objects + rmse_step = np.sqrt(mse_step) + position_rmse_all_steps.append(rmse_step) + + # ============================================================ + # Rotation error (HopNet-aligned): GT vs. predicted quaternions + # ============================================================ + # traj_full shape: (T_total, Nobj, 7), where [:, :, 3:7] = quaternions WXYZ + gt_quaternions_full = traj_full[:, :, 3:7].astype( + np.float32 + ) # (T_total, Nobj, 4) WXYZ + + # Subsample GT quaternions to match evaluation frames + gt_quaternions_sampled = gt_quaternions_full[ + ::skip_step_size + ] # (T_sampled, Nobj, 4) + gt_quaternions_eval = gt_quaternions_sampled[ + :T_eval, :Nobj_gt + ] # (T_eval, Nobj_gt, 4) + + # Load canonical mesh for each object (inverse transform from frame 0) + # Canonical mesh = rotation-free reference at origin + ref_positions_list = [] + _frame0 = load_frame( + dataset_path, 0 + ) # either layout (src/data/scene_frames.py) + for obj_idx in range(Nobj_gt): + # Load frame 0 positions + _p0 = dataset_path / "frame_00000" / f"obj_{obj_idx}_pos.npy" + obj_pos_0 = _frame0[obj_idx] + if self._object_scale is not None: + obj_pos_0 = self._object_scale.apply(_p0, obj_pos_0) + + # Get frame 0 quaternion and COM from GT trajectory + quat_0 = gt_quaternions_eval[0, obj_idx] # (4,) WXYZ + com_0 = gt_com[0, obj_idx] # (3,) + + # Convert quaternion to rotation matrix (WXYZ format) + quat_0_torch = torch.from_numpy(quat_0).float().unsqueeze(0) # (1, 4) WXYZ + # roma expects XYZW, convert from WXYZ + quat_0_xyzw = ( + torch.tensor([quat_0[1], quat_0[2], quat_0[3], quat_0[0]]) + .float() + .unsqueeze(0) + ) # (1, 4) XYZW + R_0 = roma.unitquat_to_rotmat(quat_0_xyzw).squeeze(0).numpy() # (3, 3) + + # Inverse transform: canonical = R_0^T @ (pos_0 - com_0) + canonical_pos = (obj_pos_0 - com_0) @ R_0 # (Nv, 3) + ref_positions_list.append(canonical_pos) + + # Convert GT quaternions to torch + gt_quat_torch = ( + torch.from_numpy(gt_quaternions_eval).float().to(self.device) + ) # (T, Nobj, 4) WXYZ + + import math + + # Network steps needed for the requested checkpoints (on the step-size grid, after the warm-up frames) + checkpoint_network_steps = set() + for physical_cp in physical_checkpoints: + network_step = physical_cp // skip_step_size + if ( + physical_cp % skip_step_size == 0 + and self.eval_warmup_frames <= network_step < T_eval + ): + checkpoint_network_steps.add(network_step) + + # Pre-transfer reference positions to GPU (once, not per-frame) + ref_positions_gpu = [] + for obj_idx in range(Nobj_gt): + ref_canonical = ref_positions_list[obj_idx] + ref_positions_gpu.append( + torch.from_numpy(ref_canonical).float().to(self.device) + ) + + # Compute rotation error only at checkpoint timesteps + rotation_error_at_checkpoint = {} # network_step -> RMS over objects + rotation_mae_at_checkpoint = ( + {} + ) # network_step -> mean absolute angle over objects + + for network_step in sorted(checkpoint_network_steps): + rotation_errors_t = [] + + for obj_idx in range(Nobj_gt): + if object_mask is not None and not object_mask[obj_idx]: + continue + ref_pts = ref_positions_gpu[obj_idx] # Already on GPU + Nv_actual = ref_pts.shape[0] + if Nv_actual < 3: + continue + + # Get predicted vertices for this timestep + V_pred = pred_world[network_step, obj_idx] # (Nv_padded, 3) + pred_pts = torch.from_numpy(V_pred[:Nv_actual]).float().to(self.device) + + if R_hopnet is not None: + R_pred = ( + torch.from_numpy(R_hopnet[obj_idx][network_step]) + .float() + .to(self.device) + .unsqueeze(0) + ) + else: + # Center points + ref_centered = ref_pts - ref_pts.mean(dim=0, keepdim=True) + pred_centered = pred_pts - pred_pts.mean(dim=0, keepdim=True) + # Kabsch alignment + R_pred, _ = roma.rigid_points_registration( + ref_centered.unsqueeze(0), pred_centered.unsqueeze(0) + ) + + # Convert to quaternion and compute geodesic distance + quat_pred_xyzw = roma.rotmat_to_unitquat(R_pred.squeeze(0)) + quat_pred_wxyz = torch.stack( + [ + quat_pred_xyzw[3], + quat_pred_xyzw[0], + quat_pred_xyzw[1], + quat_pred_xyzw[2], + ] + ) + + quat_gt = gt_quat_torch[network_step, obj_idx] + quat_pred_conj = torch.stack( + [ + quat_pred_wxyz[0], + -quat_pred_wxyz[1], + -quat_pred_wxyz[2], + -quat_pred_wxyz[3], + ] + ) + + w0, x0, y0, z0 = quat_pred_conj + w1, x1, y1, z1 = quat_gt + q_err_x = w0 * x1 + x0 * w1 + y0 * z1 - z0 * y1 + q_err_y = w0 * y1 - x0 * z1 + y0 * w1 + z0 * x1 + q_err_z = w0 * z1 + x0 * y1 - y0 * x1 + z0 * w1 + + imag_norm = torch.sqrt(q_err_x**2 + q_err_y**2 + q_err_z**2) + angle_rad = 2.0 * torch.arcsin(torch.clamp(imag_norm, 0.0, 1.0)) + angle_deg = angle_rad * 180.0 / np.pi + rotation_errors_t.append(angle_deg.item()) + + if rotation_errors_t: + rotation_error_at_checkpoint[network_step] = np.sqrt( + np.mean(np.square(rotation_errors_t)) + ) + rotation_mae_at_checkpoint[network_step] = float( + np.mean(rotation_errors_t) + ) + else: + rotation_error_at_checkpoint[network_step] = 0.0 + rotation_mae_at_checkpoint[network_step] = 0.0 + + # Sparse per-step rotation errors; only checkpoint steps carry values + rotation_error_all_steps = [0.0] * T_eval + for network_step, rmse in rotation_error_at_checkpoint.items(): + rotation_error_all_steps[network_step] = rmse + + # ============================================================ + # Extract checkpoints for both position and rotation + # ============================================================ + import math + + rmse_dict = {"step_size": skip_step_size} + + # HopNet convention: direct indexing, so checkpoints include warmup frames. Only frames on the step-size grid + # are evaluated: a checkpoint between two predicted frames, or at a given warm-up (context) frame, is "/". + for physical_cp in physical_checkpoints: + if physical_cp % skip_step_size: + rmse_dict[f"_na_{physical_cp}"] = ( + f"not on the step-size-{skip_step_size} grid" + ) + continue + network_step = physical_cp // skip_step_size + if network_step < self.eval_warmup_frames: + rmse_dict[f"_na_{physical_cp}"] = "a warm-up frame" + continue + + if network_step < T_eval: + # Position RMSE, keyed by target checkpoint so all step sizes + # align to the same key in TensorBoard + rmse_dict[physical_cp] = position_rmse_all_steps[network_step] + + # Rotation error (degrees), keyed by target checkpoint + if network_step < len(rotation_error_all_steps): + rot_err_value = rotation_error_all_steps[network_step] + rmse_dict[f"rotation_deg_{physical_cp}"] = rot_err_value + rmse_dict[f"rotation_mae_deg_{physical_cp}"] = ( + rotation_mae_at_checkpoint.get(network_step, 0.0) + ) + + return rmse_dict diff --git a/examples/learning/rigidformer/_support/_evidence.py b/examples/learning/rigidformer/_support/_evidence.py new file mode 100644 index 0000000000..625ea8189e --- /dev/null +++ b/examples/learning/rigidformer/_support/_evidence.py @@ -0,0 +1,83 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Bind experiment reports to executed sources, dependencies and actual assets.""" + +from __future__ import annotations + +from importlib.metadata import PackageNotFoundError, version +from pathlib import Path +import ctypes +import json +import platform +from typing import Any + +import torch + +from .._prepare import sha256 + + +def evidence(checkpoint: Path, root: Path, backend: str) -> dict[str, Any]: + """Describe the executed installation and the immutable experiment inputs.""" + import embodichain.learning.sl as sl + import embodichain_models.dynamics.rigidformer as family + + dependencies = {} + for package in ( + "torch", + "einops", + "roma", + "numpy", + "pyarrow", + "trimesh", + "flash-attn", + ): + try: + dependencies[package] = version(package) + except PackageNotFoundError: + dependencies[package] = None + sources = {} + for label, directory in ( + ("rigidformer", Path(family.__file__).parent), + ("sl", Path(sl.__file__).parent), + ("showcase", Path(__file__).parents[1]), + ): + for path in sorted(directory.rglob("*")): + if path.is_file() and ( + path.suffix in (".py", ".yaml") + or (label == "rigidformer" and path.suffix == ".json") + ): + sources[f"{label}/{path.relative_to(directory)}"] = sha256(path) + corpus = json.loads((root / "manifest.json").read_text()) + nccl = None + if torch.distributed.is_nccl_available(): + loaded = ctypes.c_int() + ctypes.CDLL("libnccl.so.2").ncclGetVersion(ctypes.byref(loaded)) + nccl = [loaded.value // 10000, loaded.value % 10000 // 100, loaded.value % 100] + return { + "python": platform.python_version(), + "dependencies": dependencies, + "torch_cuda": torch.version.cuda, + "nccl_runtime": nccl, + "gpu": torch.cuda.get_device_name() if torch.cuda.is_available() else None, + "attention_backend": backend, + "executed_source_sha256": sources, + "checkpoint_sha256": sha256(checkpoint), + "data_manifest_sha256": sha256(root / "manifest.json"), + "dataset": corpus.get("dataset"), + "scope": corpus.get("scope"), + "scene_counts": corpus.get("scene_counts"), + } diff --git a/examples/learning/rigidformer/_support/_recovery.py b/examples/learning/rigidformer/_support/_recovery.py new file mode 100644 index 0000000000..d4e25cb0a9 --- /dev/null +++ b/examples/learning/rigidformer/_support/_recovery.py @@ -0,0 +1,132 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Compare complete continuous and resumed native training checkpoints.""" + +from __future__ import annotations + +import argparse +from copy import deepcopy +import json +from pathlib import Path +from typing import Any + +import numpy as np +import torch +import yaml + +from embodichain_models.dynamics.rigidformer import ( + build_rigidformer_trainer, + export_rigidformer_checkpoint, +) +from ._evidence import evidence + + +def _exact(actual: Any, expected: Any) -> None: + """Check every saved value, including optimizer, NumPy and device RNG state.""" + if isinstance(actual, torch.Tensor): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + elif isinstance(actual, np.ndarray): + np.testing.assert_array_equal(actual, expected) + elif isinstance(actual, dict): + assert actual.keys() == expected.keys() + for key in actual: + _exact(actual[key], expected[key]) + elif isinstance(actual, (tuple, list)): + assert len(actual) == len(expected) + for a, b in zip(actual, expected): + _exact(a, b) + else: + assert actual == expected + + +def check_recovery(args: argparse.Namespace) -> None: + """Run the public factory in three fresh destinations and compare full state.""" + if args.output.exists(): + raise FileExistsError(args.output) + config = yaml.safe_load(args.config.read_text()) + args.output.mkdir(parents=True) + try: + for name in ("full", "split", "resumed"): + current = deepcopy(config) + current.setdefault("trainer", {})["output_dir"] = str(args.output / name) + trainer = build_rigidformer_trainer(current) + if not 0 < args.stop_after < trainer.cfg.total_updates: + raise ValueError("stop_after must be between zero and total_updates") + if name == "resumed": + trainer.resume(args.output / "split/last.pt") + trainer.fit(stop_after=args.stop_after if name == "split" else None) + del trainer + # These are checkpoints produced by this process in fresh directories. + full = torch.load( + args.output / "full/last.pt", + map_location="cpu", + mmap=True, + weights_only=False, + ) + resumed = torch.load( + args.output / "resumed/last.pt", + map_location="cpu", + mmap=True, + weights_only=False, + ) + compared = ["learner", "training_states", "sampler", "rng"] + for key in compared: + _exact(full[key], resumed[key]) + tensors = {} + for selection in ("raw", "ema"): + path = export_rigidformer_checkpoint( + args.output / "resumed/last.pt", + args.output / f"{selection}.pt", + weights=selection, + ) + exported = torch.load( + path, map_location="cpu", weights_only=True, mmap=True + ) + selected = ( + resumed["learner"]["model"] + if selection == "raw" + else resumed["training_states"]["ema"]["state"]["shadow"] + ) + expected = { + key.removeprefix("network."): value.to( + resumed["learner"]["model"][key].dtype + ) + for key, value in selected.items() + } + _exact(exported["model_state_dict"], expected) + tensors[selection] = len(expected) + report = { + **evidence( + args.output / "resumed/last.pt", + Path(config["data"]["root"]), + config.get("model", {}).get("attention_backend", "sdpa"), + ), + "passed": True, + "updates": resumed["learner"]["num_updates"], + "stop_after": args.stop_after, + "compared_exactly": compared, + "exported_tensors": tensors, + } + (args.output / "summary.json").write_text( + json.dumps(report, indent=2, sort_keys=True, allow_nan=False) + ) + print(json.dumps(report, indent=2)) + except BaseException as exc: + (args.output / "failure.json").write_text( + json.dumps({"type": type(exc).__name__, "message": str(exc)}) + ) + raise diff --git a/examples/learning/rigidformer/configs/sl.movi_a.yaml b/examples/learning/rigidformer/configs/sl.movi_a.yaml new file mode 100644 index 0000000000..a1f15207b1 --- /dev/null +++ b/examples/learning/rigidformer/configs/sl.movi_a.yaml @@ -0,0 +1,16 @@ +factory: embodichain_models.dynamics.rigidformer.sl_recipe:build_rigidformer_trainer +model: + preset: movi_a + attention_backend: sdpa +precision: bf16 +data: + root: outputs/rigidformer/data/MOVi-A + split_file: outputs/rigidformer/data/MOVi-A/splits.json + scale_file: outputs/rigidformer/data/MOVi-A/object_scales.json +trainer: + output_dir: outputs/rigidformer/movi_a/sl + device: cuda + batch_size: 32 + seed: 42 + checkpoint_interval: 5000 + diff --git a/examples/learning/rigidformer/configs/sl.movi_b.yaml b/examples/learning/rigidformer/configs/sl.movi_b.yaml new file mode 100644 index 0000000000..9b368c194a --- /dev/null +++ b/examples/learning/rigidformer/configs/sl.movi_b.yaml @@ -0,0 +1,16 @@ +factory: embodichain_models.dynamics.rigidformer.sl_recipe:build_rigidformer_trainer +model: + preset: movi_b + attention_backend: sdpa +precision: bf16 +data: + root: outputs/rigidformer/data/MOVi-B + split_file: outputs/rigidformer/data/MOVi-B/splits.json + scale_file: outputs/rigidformer/data/MOVi-B/object_scales.json +trainer: + output_dir: outputs/rigidformer/movi_b/sl + device: cuda + batch_size: 21 + seed: 42 + checkpoint_interval: 5000 + diff --git a/examples/learning/rigidformer/configs/sl.movi_sphere.yaml b/examples/learning/rigidformer/configs/sl.movi_sphere.yaml new file mode 100644 index 0000000000..a87e51899b --- /dev/null +++ b/examples/learning/rigidformer/configs/sl.movi_sphere.yaml @@ -0,0 +1,16 @@ +factory: embodichain_models.dynamics.rigidformer.sl_recipe:build_rigidformer_trainer +model: + preset: movi_sphere + attention_backend: sdpa +precision: bf16 +data: + root: outputs/rigidformer/data/MOVi-Sphere + split_file: outputs/rigidformer/data/MOVi-Sphere/splits.json + scale_file: outputs/rigidformer/data/MOVi-Sphere/object_scales.json +trainer: + output_dir: outputs/rigidformer/movi_sphere/sl + device: cuda + batch_size: 32 + seed: 42 + checkpoint_interval: 5000 + diff --git a/examples/learning/rigidformer/research/acceptance.md b/examples/learning/rigidformer/research/acceptance.md new file mode 100644 index 0000000000..633e203a96 --- /dev/null +++ b/examples/learning/rigidformer/research/acceptance.md @@ -0,0 +1,104 @@ +# Native RigidFormer qualification + +This protocol qualifies the installable author model, prepared MOVi inputs and +shared SL training. Published accuracy and full-budget retraining are separate +experiments. A debug subset or a short paired run cannot satisfy either claim. + +## Frozen inputs + +- Author source: `Frank-ZY-Dou/Dynamics-Modeling`, revision + `160fe0852c51cffae7b7440452b58d0b49db6df6`. The comparison verifies all 23 + files in the model package's `REFERENCE_SOURCES.json` before importing the + independent original model, Trainer and EMA implementation. +- Data: `frankzydou/RigidFormer-MOVi` at + `e78514b479b11acfa5d23d28f1211d7830646ff9`. Use the scene's complete object + identities, collision mesh vertices, recorded size, four FPS anchors and + physical parameters. Preparation checks raw scene and mesh hashes, including + debug preparation. Model integration uses dt=1; physical time is stride/240 s. +- Weights: `frankzydou/RigidFormer` at + `8f2c0a3a157b64e4b310c64a01eee6e4e7e25af6`, with the checksum selected by + `_prepare.py`. Use Sphere weights on Sphere scenes for the evaluation smoke. + Cross-dataset runs can establish operator agreement, not dataset accuracy. + +Prepared readers verify actual payload hashes and manifest coverage, including +anchor/FPS caches. An omitted payload, changed split, changed size sidecar or +changed geometry is not an exact-resume-compatible source. + +## Required checks + +| Check | Acceptance | +|---|---| +| Independent author comparison | Prediction, free rollout, loss, clipped gradients, parameters, buffers, optimizer moments and EMA agree at atol=1e-5, rtol=1e-4 for at least three updates. Optimizer groups, scalar policy and scheduler state agree exactly. | +| Attention | Run the genuine Flash-Attention package for Flash qualification. SDPA compatibility of the reference is reported separately. | +| Full architecture recovery | Continuous four-update training equals two updates plus resume plus two updates, including model, optimizer, EMA, schedule, sampler and Python/NumPy/Torch/CUDA RNG. Raw and EMA exports retain every named tensor. | +| Epoch recovery | The small model/data fixture runs six updates over three epochs with two warmup epochs, resuming after update three. Every saved value must match exactly. | +| CUDA/NCCL | Two ranks execute SyncBatchNorm, unequal 3/2-object supervision and exact rank-local recovery. The weighted-update oracle uses a concatenated batch outside any process group, with stochastic dropout disabled on both sides. Recovery retains author dropout. | +| Fixed weights evaluation | Complete 480-frame rollout on the declared debug scenes, with four warmup frames and strides 1/5/10. Persist scope and physical-frame metrics without interpolating off-grid checkpoints. No accuracy threshold is assigned to a debug subset. | +| Packaging | Each dependent layer builds and installs its own wheel. Imports, configs and licenses resolve from the installed wheel without a repository PYTHONPATH. The final wheel includes runnable showcase resources, without downloaded assets. | + +The numerical comparison uses one fixed batch per comparison epoch to exercise +epoch scheduling on each update. Seed 42, selected scene/frame identities, +executed source hashes, dependency versions and asset hashes are recorded. +This does not reproduce the research DataLoader's worker-RNG trajectory or the +released four/six-GPU training budget. Native sources own an explicit recoverable +epoch traversal, and exact recovery requires the same world size and recipe. + +## Reproduce + +Install the optional dependencies and use a compatible CUDA/Torch/Flash/NCCL +stack. Put all downloaded assets, prepared data and reports outside the checkout. + +```bash +embodichain run-example learning/rigidformer prepare --dataset movi_sphere \ + --assets all --cache /path/to/assets --output /path/to/prepared --limit-per-split 2 + +embodichain run-example learning/rigidformer compare \ + --reference /path/to/Dynamics-Modeling/RigidFormer \ + --checkpoint /path/to/assets/model/rigidformer_movisphere.pth \ + --data-root /path/to/prepared/MOVi-Sphere --attention-backend flash \ + --seed 42 --updates 3 --output /path/to/fresh-comparison + +# Use the Sphere recipe with external data paths, batch_size=1, total_updates=4. +embodichain run-example learning/rigidformer check-recovery \ + --config /path/to/recovery.yaml --stop-after 2 --output /path/to/fresh-recovery + +python -m pytest tests/learning/sl tests/models/dynamics -m 'not gpu' -q +python -m pytest tests/models/dynamics/test_rigidformer_official.py \ + tests/models/dynamics/test_rigidformer_training.py --run-gpu -m gpu -q +CUDA_VISIBLE_DEVICES=0,1 python -m pytest \ + tests/models/dynamics/test_rigidformer_distributed.py --run-gpu -m gpu -q + +python -m scripts.benchmark.learning.rigidformer_index --output /path/to/index.json +``` + +The index benchmark measures CPU metadata construction at 1/32/960 scenes, +excluding data IO, epoch permutations, model work and GPU execution. It reports +median latency and separately measured Python allocations; these are not total +process-memory or training-throughput measurements. + +## Evidence + +The checked-in [qualification index](qualification.json) identifies the executed +reports and their hashes. Reports and binary inputs remain external artifacts; +historical reports are not relabeled after source or environment changes. +Every reported backend must be the backend actually executed. In particular, +Torch's compile-time NCCL version can differ from a process-local loaded runtime. + +The development Blackwell host required a locally built NCCL 2.29.7 (CUDA 12.8, +sm_120) for model-layout broadcasts; its inherited 2.26.2 runtime failed the same +broadcast in an independent module containing only the parameter/buffer shapes. +The qualification records this runtime override and source/binary identities. +No NCCL transport-throughput claim follows from these correctness checks. + +To reproduce that host's two-rank runtime, build outside the checkout and load +the library for the validation process. The recorded source commit is +`b91894bd5b190c874d98a017f93f5daa515b65d0`: + +```bash +git clone --depth 1 --branch v2.29.7-1 https://github.com/NVIDIA/nccl.git /path/to/nccl +git -C /path/to/nccl checkout --detach b91894bd5b190c874d98a017f93f5daa515b65d0 +make -C /path/to/nccl -j16 src.build NVCC_GENCODE='-gencode=arch=compute_120,code=sm_120' +LD_PRELOAD=/path/to/nccl/build/lib/libnccl.so.2 NCCL_IB_DISABLE=1 NCCL_P2P_DISABLE=1 \ + CUDA_VISIBLE_DEVICES=0,1 python -m pytest \ + tests/models/dynamics/test_rigidformer_distributed.py --run-gpu -m gpu -q +``` diff --git a/examples/learning/rigidformer/research/qualification.json b/examples/learning/rigidformer/research/qualification.json new file mode 100644 index 0000000000..de9a1731f3 --- /dev/null +++ b/examples/learning/rigidformer/research/qualification.json @@ -0,0 +1,139 @@ +{ + "author_revision": "160fe0852c51cffae7b7440452b58d0b49db6df6", + "base": "7cc143dd0c6160a3a5b726bb5c9b9982d86492e5", + "external_artifact_root": "/root/outputs/rigidformer-official-20261027", + "full_budget_retraining": "not performed", + "full_published_benchmark": "not performed", + "qualification": "integration-and-numerical-parity", + "reports": { + "author_flash": { + "all_recorded_max_abs_errors": 0.0, + "artifact": "installed-final-flash/summary.json", + "checkpoint_origin": "official-release", + "checkpoint_sha256": "ffb0efc6904b5a63504a5005bb6c4df2e79e8de20641dc205885593747fcf7d9", + "seed": 42, + "sha256": "bc33e89c7f29431a679007b7dd567aae26eac9494f56cd5a80c881a4e0a049ee", + "updates": 3 + }, + "fixed_weight_evaluation": { + "artifact": "installed-final-evaluation/summary.json", + "dataset": "movi_sphere", + "num_frames": 480, + "num_scenes": 2, + "scope": "debug-subset", + "sha256": "9151c178adf3193c7752253505b7205d60c5668f43bc71ad368572db47364267", + "split": "test", + "step_sizes": [ + 1, + 5, + 10 + ] + }, + "full_architecture_recovery": { + "artifact": "installed-final-recovery/summary.json", + "compared_exactly": [ + "learner", + "training_states", + "sampler", + "rng" + ], + "exported_tensors": { + "ema": 485, + "raw": 485 + }, + "sha256": "6619a190694b2a4e8bfb954a6ee4fa6c17703b0bd77f081d0398184aede3ddf7", + "stop_after": 2, + "updates": 4 + }, + "index_benchmark": { + "artifact": "qualified-index.json", + "sha256": "4b07998f2222853c8e789aaa849235067f0740bf7eeb83b62a0caf912436e5e5" + }, + "two_rank_nccl": { + "artifact": "final-gpu-tests/test_nccl_weighted_syncbn_and_0/summary.json", + "backend": "nccl", + "compared": [ + "model", + "optimizer", + "EMA", + "schedule", + "rank RNG", + "rank source cursor", + "rank buffers" + ], + "epochs": 3, + "loaded_nccl_version": 22907, + "recovery": "exact", + "sha256": "cbabbf2e951fa12a86baf2fc9ba60dcbcebbf2e3058731a69db43abfb1e36e36", + "split_after": 3, + "sync_batch_norm": true, + "updates": 6, + "warmup_epochs": 2, + "weighted_update": { + "global_loss": 0.11964317858219146, + "state_max_abs": 1.7881393432617188e-07, + "valid_objects_by_rank": [ + 3, + 2 + ] + }, + "world_size": 2 + }, + "wheel_layers": [ + { + "artifact": "layer1/validation/summary.json", + "sha256": "34bbe578633b63aea2d311c5e4c38eb59a1ff451645b9b157be59924ed3df79b" + }, + { + "artifact": "layer2/validation/summary.json", + "sha256": "c82784c3c951ded351356fc389791951f4d79bf0779cbcf21b14e7153b861b4c" + }, + { + "artifact": "layer3-final/verification/summary.json", + "sha256": "c2ceedab8d04ecaaf654e96e9f675f377e0a26f0d53625870131a3213e96e7bf" + } + ] + }, + "runtime": { + "cuda": "12.8", + "dependencies": { + "einops": "0.8.2", + "flash-attn": "2.8.3", + "numpy": "2.4.6", + "pyarrow": "24.0.0", + "roma": "1.5.4", + "torch": "2.7.0+cu128", + "trimesh": "4.12.2" + }, + "gpu": "NVIDIA RTX PRO 5000 72GB Blackwell", + "loaded_nccl": [ + 2, + 29, + 7 + ], + "nccl_architecture": "sm_120", + "nccl_library_sha256": "37ee50973f269610cfaae10760037cbaf13f7fd92762e3ccca80c9d2d075a97b", + "nccl_source_commit": "b91894bd5b190c874d98a017f93f5daa515b65d0", + "nccl_transport_settings": { + "NCCL_IB_DISABLE": "1", + "NCCL_P2P_DISABLE": "1" + }, + "python": "3.11.14" + }, + "validation": { + "api_exports": "3284/3284", + "context": "passed", + "cpu_log": { + "artifact": "final-cpu.log", + "sha256": "3d4c334b19fd38021ab54af3248d4350b3299f2d14c3aa74adb21e23c897e6e4" + }, + "cpu_tests": 149, + "focused_sphinx": "passed", + "gpu_log": { + "artifact": "final-gpu.log", + "sha256": "6e94ed43b02c14495e39a6d054d7474dd7c025f8d2e9c232443a5f70da02aab2" + }, + "gpu_tests": 6, + "test_basis": "Full behavioral run plus per-layer checks; subsequent production edits only clarify API docstrings." + } +} diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/data.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/data.py index 4a1551a8bc..007f534b75 100644 --- a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/data.py +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/data.py @@ -191,7 +191,14 @@ def _indices(self, epoch: int) -> torch.Tensor: return self._order def updates_per_epoch(self, batch_size: int) -> int: - """Return full-batch updates per rank, matching drop-last training.""" + """Return full-batch updates per rank, matching drop-last training. + + Args: + batch_size: Number of windows in each rank's optimizer update. + + Returns: + Number of complete rank-local batches in an epoch. + """ if type(batch_size) is not int or batch_size < 1 or self._per_rank < batch_size: raise ValueError("Every rank needs at least one complete batch") return self._per_rank // batch_size @@ -255,14 +262,22 @@ def sample(self, batch_size: int, *, generator: torch.Generator) -> SupervisedBa ) def on_update(self, learner: Any) -> None: - """Commit the staged window traversal after a successful optimizer step.""" + """Commit the staged window traversal after a successful optimizer step. + + Args: + learner: Host update context; progress comes from the staged batch. + """ if self._pending is None: raise RuntimeError("No geometry batch is pending") self.epoch, self.cursor = self._pending self._pending = None def state_dict(self) -> dict[str, Any]: - """Return the committed rank-local sampling boundary.""" + """Return the committed rank-local sampling boundary. + + Returns: + Owned source identity, epoch, cursor and rank/world-size values. + """ return { "source_id": self.source_id, "epoch": self.epoch, @@ -272,7 +287,11 @@ def state_dict(self) -> dict[str, Any]: } def validate_state_dict(self, state: Mapping[str, Any]) -> None: - """Reject incompatible geometry or traversal before training mutation.""" + """Reject incompatible geometry or traversal before training mutation. + + Args: + state: Candidate committed sampling snapshot. + """ if ( state.get("source_id") != self.source_id or state.get("rank") != self.rank @@ -286,7 +305,11 @@ def validate_state_dict(self, state: Mapping[str, Any]) -> None: self._verify_files() def load_state_dict(self, state: Mapping[str, Any]) -> None: - """Restore a prevalidated committed traversal and discard staged reads.""" + """Restore a prevalidated committed traversal and discard staged reads. + + Args: + state: Snapshot accepted by validate_state_dict. + """ self.epoch, self.cursor = state["epoch"], state["cursor"] self._pending = None self._order_epoch = -1 diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/ema.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/ema.py index 5420bca366..150e97dfd6 100644 --- a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/ema.py +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/ema.py @@ -48,7 +48,12 @@ def __init__(self, model: nn.Module, decay: float = 0.9999) -> None: @torch.no_grad() def on_update(self, learner: Any) -> None: - """Update the average once after a successful parameter update.""" + """Update the average once after a successful parameter update. + + Args: + learner: Successful host update context. Values are read from the + model supplied at construction. + """ self.num_updates += 1 decay = min(self.decay, (1 + self.num_updates) / (10 + self.num_updates)) floating, values = [], [] @@ -62,7 +67,11 @@ def on_update(self, learner: Any) -> None: torch._foreach_add_(floating, values, alpha=1 - decay) def state_dict(self) -> dict[str, Any]: - """Return owned FP32 shadows and the exact EMA update cursor.""" + """Return owned FP32 shadows and the exact EMA update cursor. + + Returns: + CPU shadow tensors, decay and committed update count. + """ return { "decay": self.decay, "num_updates": self.num_updates, @@ -70,7 +79,11 @@ def state_dict(self) -> dict[str, Any]: } def validate_state_dict(self, state: Mapping[str, Any]) -> None: - """Validate EMA policy and every shadow tensor before restoration.""" + """Validate EMA policy and every shadow tensor before restoration. + + Args: + state: Candidate snapshot; validation does not change live weights. + """ if ( state.get("decay") != self.decay or type(state.get("num_updates")) is not int @@ -86,7 +99,11 @@ def validate_state_dict(self, state: Mapping[str, Any]) -> None: @torch.no_grad() def load_state_dict(self, state: Mapping[str, Any]) -> None: - """Restore a prevalidated EMA without changing live model weights.""" + """Restore a prevalidated EMA without changing live model weights. + + Args: + state: Snapshot accepted by validate_state_dict. + """ self.num_updates = state["num_updates"] for key, value in self.shadow.items(): value.copy_(state["shadow"][key]) diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/model.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/model.py index 59c4baaa1e..529df0de44 100644 --- a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/model.py +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/model.py @@ -50,7 +50,11 @@ class RigidFormerCfg: max_batch_bytes: int = 512 * 1024 * 1024 def reference_config(self) -> dict[str, Any]: - """Return an owned author configuration without opening data assets.""" + """Return an owned author configuration without opening data assets. + + Returns: + Complete author configuration, including explicit architecture overrides. + """ self.validate() if self.preset not in ("movi_b", "movi_a", "movi_sphere"): raise ValueError("preset must be movi_b, movi_a, or movi_sphere") diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/supervised.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/supervised.py index 34693ac393..8bae1efefb 100644 --- a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/supervised.py +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/supervised.py @@ -57,6 +57,8 @@ class RigidFormerObjective: Args: weights: Position/acceleration before/after Kabsch weights. time_step_scale: Divide errors by step_size squared before smooth-L1. + augmentation: Optional author geometry augmentation, applied in FP32 + before the network's numerical context. """ def __init__( diff --git a/tests/examples/test_rigidformer_showcase.py b/tests/examples/test_rigidformer_showcase.py new file mode 100644 index 0000000000..7147d55d9f --- /dev/null +++ b/tests/examples/test_rigidformer_showcase.py @@ -0,0 +1,73 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Reject altered raw/prepared geometry before conversion or model loading.""" + +from __future__ import annotations + +import hashlib +import json +from types import SimpleNamespace + +import pytest + + +def test_debug_preparation_verifies_collision_meshes(tmp_path): + pytest.importorskip("trimesh") + from examples.learning.rigidformer._prepare import prepare + + raw = tmp_path / "raw" + scene = raw / "data/movi_a/train-00000.parquet" + mesh = raw / "assets/kubasic/cube/collision_geometry.obj" + for path in (scene, mesh): + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(b"original-payload") + manifest = { + "files": { + str(path.relative_to(raw)): { + "sha256": hashlib.sha256(path.read_bytes()).hexdigest() + } + for path in (scene, mesh) + } + } + (raw / "manifest.json").write_text(json.dumps(manifest)) + mesh.write_bytes(b"changed-mesh") + with pytest.raises(ValueError, match="raw data checksum differs: assets/"): + prepare(raw, tmp_path / "prepared", "movi_a", limit_per_split=1) + assert not (tmp_path / "prepared/MOVi-A").exists() + + +def test_evaluation_verifies_geometry_before_loading_weights(tmp_path): + from examples.learning.rigidformer.__main__ import evaluate + + root = tmp_path / "prepared" + root.mkdir() + data = root / "frames.npy" + data.write_bytes(b"committed-geometry") + (root / "manifest.json").write_text( + json.dumps( + {"files": {data.name: hashlib.sha256(data.read_bytes()).hexdigest()}} + ) + ) + data.write_bytes(b"altered-geometry") + output = tmp_path / "evaluation" + with pytest.raises(ValueError, match="geometry hash differs"): + evaluate( + SimpleNamespace( + data_root=root, output=output, checkpoint=tmp_path / "absent.pt" + ) + ) + assert json.loads((output / "failure.json").read_text())["type"] == "ValueError" From 59bb5e76e651d9820431c281a09db016d83a6d04 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Wed, 28 Oct 2026 01:34:45 +0000 Subject: [PATCH 2/4] docs(rigidformer): record post-review qualification --- .../rigidformer/research/qualification.json | 29 +++++++++++++++---- 1 file changed, 23 insertions(+), 6 deletions(-) diff --git a/examples/learning/rigidformer/research/qualification.json b/examples/learning/rigidformer/research/qualification.json index de9a1731f3..3e78cf1afe 100644 --- a/examples/learning/rigidformer/research/qualification.json +++ b/examples/learning/rigidformer/research/qualification.json @@ -8,11 +8,19 @@ "reports": { "author_flash": { "all_recorded_max_abs_errors": 0.0, - "artifact": "installed-final-flash/summary.json", + "artifact": "post-review-flash/summary.json", "checkpoint_origin": "official-release", "checkpoint_sha256": "ffb0efc6904b5a63504a5005bb6c4df2e79e8de20641dc205885593747fcf7d9", + "model_source_commit": "acb01d18e28832557dbd4bff39ede7b6d12bb02f", + "nccl_runtime": [ + 2, + 26, + 2 + ], + "previous_artifact": "installed-final-flash/summary.json", + "previous_sha256": "bc33e89c7f29431a679007b7dd567aae26eac9494f56cd5a80c881a4e0a049ee", "seed": 42, - "sha256": "bc33e89c7f29431a679007b7dd567aae26eac9494f56cd5a80c881a4e0a049ee", + "sha256": "4e037cfb591fc73f9a47fd42d2927a7c9246583a78723b029edc54de99f4c2ed", "updates": 3 }, "fixed_weight_evaluation": { @@ -30,7 +38,7 @@ ] }, "full_architecture_recovery": { - "artifact": "installed-final-recovery/summary.json", + "artifact": "post-review-recovery/summary.json", "compared_exactly": [ "learner", "training_states", @@ -41,7 +49,10 @@ "ema": 485, "raw": 485 }, - "sha256": "6619a190694b2a4e8bfb954a6ee4fa6c17703b0bd77f081d0398184aede3ddf7", + "model_source_commit": "acb01d18e28832557dbd4bff39ede7b6d12bb02f", + "previous_artifact": "installed-final-recovery/summary.json", + "previous_sha256": "6619a190694b2a4e8bfb954a6ee4fa6c17703b0bd77f081d0398184aede3ddf7", + "sha256": "2daad678a5b633615a871db5b2a610bae55819650779a00567dee10273a4545b", "stop_after": 2, "updates": 4 }, @@ -50,7 +61,7 @@ "sha256": "4b07998f2222853c8e789aaa849235067f0740bf7eeb83b62a0caf912436e5e5" }, "two_rank_nccl": { - "artifact": "final-gpu-tests/test_nccl_weighted_syncbn_and_0/summary.json", + "artifact": "post-review-two-rank/test_nccl_weighted_syncbn_and_0/summary.json", "backend": "nccl", "compared": [ "model", @@ -63,6 +74,9 @@ ], "epochs": 3, "loaded_nccl_version": 22907, + "model_source_commit": "acb01d18e28832557dbd4bff39ede7b6d12bb02f", + "previous_artifact": "final-gpu-tests/test_nccl_weighted_syncbn_and_0/summary.json", + "previous_sha256": "cbabbf2e951fa12a86baf2fc9ba60dcbcebbf2e3058731a69db43abfb1e36e36", "recovery": "exact", "sha256": "cbabbf2e951fa12a86baf2fc9ba60dcbcebbf2e3058731a69db43abfb1e36e36", "split_after": 3, @@ -134,6 +148,9 @@ "sha256": "6e94ed43b02c14495e39a6d054d7474dd7c025f8d2e9c232443a5f70da02aab2" }, "gpu_tests": 6, - "test_basis": "Full behavioral run plus per-layer checks; subsequent production edits only clarify API docstrings." + "post_review_cpu_tests": 86, + "post_review_gpu_tests": 3, + "post_review_note": "Flash comparison, full-model recovery and two-rank NCCL were rerun after the model review fixes; other indexed reports retain their original source hashes.", + "test_basis": "Initial full behavioral run and per-layer checks; post-review model fixes were checked separately." } } From 6a16643bf3621ffaa398c6475bd20c1e46daa55f Mon Sep 17 00:00:00 2001 From: yuecideng Date: Wed, 28 Oct 2026 14:40:54 +0000 Subject: [PATCH 3/4] feat(examples): compare RigidFormer rollout in the engine --- .../topics/neural-dynamics/neural-dynamics.md | 3 + examples/README.md | 6 + examples/learning/rigidformer/README.md | 53 ++ examples/learning/rigidformer/__main__.py | 33 +- examples/learning/rigidformer/_demo.py | 460 ++++++++++++++++++ 5 files changed, 554 insertions(+), 1 deletion(-) create mode 100644 examples/learning/rigidformer/_demo.py diff --git a/agent_context/topics/neural-dynamics/neural-dynamics.md b/agent_context/topics/neural-dynamics/neural-dynamics.md index bf46041422..a8c388e841 100644 --- a/agent_context/topics/neural-dynamics/neural-dynamics.md +++ b/agent_context/topics/neural-dynamics/neural-dynamics.md @@ -59,6 +59,9 @@ legacy research DataLoader's worker-RNG trajectory. The [showcase](../../../examples/learning/rigidformer/README.md) owns downloads, conversion, fixed-weight evaluation and independent author-Trainer comparison. +Its `demo` command in `_demo.py` compares live Newton spheres with render-only +model predictions after four observed frames, through a native window or an +engine-camera recording. It uses explicit prepared geometry and weight paths. Assets/results stay outside the wheel. Numerical agreement, recovery and full benchmark/retraining qualification remain distinct evidence requirements. diff --git a/examples/README.md b/examples/README.md index d720762027..31b1e903ba 100644 --- a/examples/README.md +++ b/examples/README.md @@ -8,6 +8,7 @@ embodichain run-example sim/motion/planners/curobo_planner --help embodichain run-example ods/simple_bc_demo/train --help embodichain run-example learning/nerd --help embodichain run-example learning/rwm --help +embodichain run-example learning/rigidformer demo --help embodichain run-example sim/atomic_action --skill place --variants 8 embodichain run-example sim/atomic_action/coordinated_pickment --help ``` @@ -48,6 +49,11 @@ state-transition fitting and state/action gradients. Its generated recipes use the installable `embodichain_models.dynamics.rwm` model family and shared SL trainer; its corpus generation and acceptance helpers remain private to the example. +The [RigidFormer showcase](learning/rigidformer/README.md) prepares official +weights and geometry, trains and evaluates the model, and runs an engine demo +with Newton spheres beside a free model rollout. The demo opens a native window +or records an engine-camera video with `--headless`. + The [Atomic Skill example catalog](sim/atomic_action/README.md) owns all per-skill, coordinated-manipulation and recovery demonstrations. The two source-checkout walkthroughs under `scripts/tutorials/atomic_action/` teach single-arm and dual-arm diff --git a/examples/learning/rigidformer/README.md b/examples/learning/rigidformer/README.md index c4ed3e0ee1..9db2013e90 100644 --- a/examples/learning/rigidformer/README.md +++ b/examples/learning/rigidformer/README.md @@ -16,6 +16,9 @@ python -m pip install 'embodichain[rigidformer]' embodichain run-example learning/rigidformer --help ``` +From a source checkout, use `python -m pip install -e '.[rigidformer]'` to install +that checkout and its optional model dependencies. + The standard native attention backend is Torch SDPA. The explicit `flash` backend requires a Flash-Attention build compatible with Torch, CUDA and the GPU. Reference evidence records the actual dependency/backend versions. @@ -61,6 +64,56 @@ Readers reject altered payloads and geometry omitted from the manifest. FPS and anchor caches are committed inputs, rather than untracked files generated during training. Debug preparation verifies collision meshes as well as scene shards. +## Engine demo + +Prepare the Sphere weights and a small geometry subset, then open the engine +window: + +```bash +embodichain run-example learning/rigidformer prepare \ + --dataset movi_sphere --assets all --limit-per-split 1 + +embodichain run-example learning/rigidformer demo +``` + +The left side contains dynamic Newton spheres. The right side contains engine +`VisualObject` spheres driven by the official RigidFormer model. Both sides use +the same geometry, masses, friction and restitution from a prepared test scene. +The model observes four sampled simulation frames, then predicts the remaining +trajectory in one free rollout. Predicted spheres have no colliders and provide +no feedback to the physical simulation. The window holds the final frame for +three seconds; Ctrl+C closes the demo early. + +To record on a machine without a display, pass `--headless`: + +```bash +embodichain run-example learning/rigidformer demo --headless \ + --output outputs/rigidformer/demo-headless +``` + +The output directory must be fresh. It contains engine-camera `rollout.mp4`, +`frame.png`, aligned `trajectory.npz`, and a `summary.json` with source hashes, +timing and center errors. Frames stream to the video writer as they are captured. +The defaults show 51 samples at stride 5 from 240 Hz physics, played at 12 fps. +`--frames`, `--stride`, `--fps` and `--hold-seconds` control the demonstration. +`--scene` selects a prepared test scene; `--object-ids` selects two or three +template objects. `--device cuda:1` selects a different visible GPU. + +Use existing resources or compatible exported inference weights explicitly: + +```bash +embodichain run-example learning/rigidformer demo --headless \ + --data-root /path/to/prepared/MOVi-Sphere \ + --checkpoint /path/to/rigidformer_movisphere.pth \ + --output outputs/rigidformer/my-demo +``` + +The scene uses Newton XPBD and newly chosen initial positions and velocities. +Its center errors describe this engine comparison. The fixed-checkpoint +`evaluate` command below reports the official MOVi data metrics. The preparation +limit bounds converted scenes; downloading may still acquire the dataset's +numeric shards and shared meshes. + ## Train with shared SL ```bash diff --git a/examples/learning/rigidformer/__main__.py b/examples/learning/rigidformer/__main__.py index 361c3071c0..e44598c023 100644 --- a/examples/learning/rigidformer/__main__.py +++ b/examples/learning/rigidformer/__main__.py @@ -14,7 +14,7 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Prepare official assets, evaluate inference weights, and export native runs.""" +"""Prepare RigidFormer assets, run an engine demo, evaluate and export weights.""" from __future__ import annotations import argparse @@ -41,6 +41,33 @@ def main() -> None: default=0, help="Debug subset only; zero prepares the full frozen corpus", ) + p = sub.add_parser("demo", help="Compare live Newton dynamics with RigidFormer") + p.add_argument( + "--checkpoint", + type=Path, + default=Path("outputs/rigidformer/assets/model/rigidformer_movisphere.pth"), + ) + p.add_argument( + "--data-root", type=Path, default=Path("outputs/rigidformer/data/MOVi-Sphere") + ) + p.add_argument("--output", type=Path, default=Path("outputs/rigidformer/demo")) + p.add_argument("--scene", help="Prepared test scene; default selects the first") + p.add_argument( + "--object-ids", type=int, nargs="+", help="Two or three template object IDs" + ) + p.add_argument("--device", default="cuda:0") + p.add_argument("--attention-backend", choices=("sdpa", "flash"), default="sdpa") + p.add_argument( + "--headless", action="store_true", help="Record without opening a window" + ) + p.add_argument("--frames", type=int, default=51) + p.add_argument("--stride", type=int, default=5) + p.add_argument( + "--fps", type=int, default=12, help="Window playback and video frame rate" + ) + p.add_argument( + "--hold-seconds", type=float, default=3.0, help="Final window hold duration" + ) p = sub.add_parser("export") p.add_argument("--checkpoint", type=Path, required=True) p.add_argument("--output", type=Path, required=True) @@ -104,6 +131,10 @@ def main() -> None: indent=2, ) ) + elif args.command == "demo": + from ._demo import run_demo + + run_demo(args) elif args.command == "export": from embodichain_models.dynamics.rigidformer.checkpoint import ( export_rigidformer_checkpoint, diff --git a/examples/learning/rigidformer/_demo.py b/examples/learning/rigidformer/_demo.py new file mode 100644 index 0000000000..64f4bce6d9 --- /dev/null +++ b/examples/learning/rigidformer/_demo.py @@ -0,0 +1,460 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Display Newton sphere dynamics beside a free RigidFormer rollout. + +Run: embodichain run-example learning/rigidformer demo --help +""" + +from __future__ import annotations + +import argparse +from dataclasses import dataclass +import json +from pathlib import Path +import time +from typing import TYPE_CHECKING, Any + +import numpy as np +import torch + +from ._prepare import sha256 + +if TYPE_CHECKING: + from embodichain.lab.sim import SimulationManager + from embodichain.lab.sim.objects import RigidObject, VisualObject + from embodichain.lab.sim.sensors import Camera + from embodichain_models.dynamics.rigidformer import RigidFormerModel + +__all__ = ["run_demo"] + +_PHYSICS_HZ = 240 +_WARMUP = 4 +_WIDTH, _HEIGHT = 960, 540 +_OFFSET = (5.0, 0.0, 0.0) +_EYE, _TARGET = (0.15, -10.0, 5.2), (0.15, 0.0, 1.6) +_COLORS = ((0.87, 0.16, 0.09, 1.0), (0.04, 0.50, 0.85, 1.0), (0.96, 0.65, 0.05, 1.0)) + + +@dataclass +class _Geometry: + scene: str + object_ids: list[int] + local_vertices: torch.Tensor + valid_vertices: torch.Tensor + anchors: torch.Tensor + vertex_rank: torch.Tensor + parameters: torch.Tensor + radii: torch.Tensor + + +def _load_geometry(args: argparse.Namespace, config: dict[str, Any]) -> _Geometry: + """Read committed sphere templates, anchors and physical properties.""" + from embodichain_models.dynamics.rigidformer._native.data.dataset import ( + MultiObjectDynamics, + ) + + data = config["data"] + dataset = MultiObjectDynamics( + str(args.data_root), + sequence_length=_WARMUP, + stride=999999, + split="test", + split_file=str(args.data_root / "splits.json"), + object_scale_file=str(args.data_root / "object_scales.json"), + anchor_k=data["anchor_k"], + anchor_fps_start=data.get("anchor_fps_start", "index0"), + fps_order=data.get("fps_order", False), + if_phys_param=True, + random_step_size={"enable": False}, + augmentation={}, + require_prepared_cache=True, + ) + selected = next( + ( + index + for index, item in enumerate(dataset.all_sequences) + if (args.scene is None or item["dataset"] == args.scene) + and len(dataset._mesh_nv_by_dataset[item["dataset"]]) >= 2 + ), + None, + ) + if selected is None: + raise ValueError("The selected test scene must contain at least two objects") + sample = dataset[selected] + scene = dataset.all_sequences[selected]["dataset"] + count = int(sample["valid_obj_mask"].sum()) + ids = args.object_ids + if ids is None: + ids = (np.arange(min(3, count)) * count // min(3, count)).tolist() + if ( + not 2 <= len(ids) <= 3 + or len(set(ids)) != len(ids) + or any(index < 0 or index >= count for index in ids) + ): + raise ValueError("--object-ids must select two or three distinct scene objects") + origins = torch.from_numpy( + np.load(args.data_root / scene / "vis_traj.npy")[0, 0, ids, :3].copy() + ).float() + mask = sample["valid_vtx_mask"][0, ids] + vertices = sample["ref_obj_positions"][ids] - origins[:, None] + vertices = torch.where(mask[..., None], vertices, 0.0) + radii = vertices.norm(dim=-1).amax(dim=-1) + return _Geometry( + scene, + list(ids), + vertices, + mask, + sample["anchor_idx"][ids].long(), + sample["vtx_rank"][ids].long(), + sample["phys_params"][ids], + radii, + ) + + +def _read_poses(bodies: list[RigidObject]) -> torch.Tensor: + """Snapshot public xyz/xyzw poses without retaining engine buffers.""" + return torch.cat([body.get_local_pose() for body in bodies]).detach().clone() + + +def _vertices(poses: torch.Tensor, geometry: _Geometry) -> torch.Tensor: + """Move the template vertices with the observed sphere poses.""" + from embodichain.utils.math import matrix_from_quat + + local = geometry.local_vertices.to(poses.device) + mask = geometry.valid_vertices.to(poses.device) + moved = torch.einsum("nij,npj->npi", matrix_from_quat(poses[:, 3:]), local) + return torch.where(mask[..., None], moved + poses[:, None, :3], 0.0) + + +def _predict( + model: RigidFormerModel, + warmup: torch.Tensor, + geometry: _Geometry, + args: argparse.Namespace, +) -> torch.Tensor: + """Infer the entire free rollout from the four engine observations once.""" + import roma + + device = warmup.device + mask = geometry.valid_vertices.to(device) + objects = len(geometry.object_ids) + batch = { + "obj_positions": warmup[None], + "valid_obj_mask": torch.ones(1, objects, dtype=torch.bool, device=device), + "valid_vtx_mask": mask[None, None].expand(1, _WARMUP, -1, -1), + "ref_obj_positions": warmup[0][None], + "anchor_idx": geometry.anchors[None].to(device), + "vtx_rank": geometry.vertex_rank[None].to(device), + "phys_params": geometry.parameters[None].to(device), + "skip_step_size": torch.tensor([float(args.stride)], device=device), + "dt": torch.ones(1, device=device), + "time_step_id": torch.zeros(1, dtype=torch.long, device=device), + } + with torch.inference_mode(): + predicted = model.rollout(batch, num_frames=args.frames, warmup_frames=_WARMUP)[ + "obj_positions" + ][0] + if not torch.isfinite(predicted).all(): + raise RuntimeError("RigidFormer produced a non-finite rollout") + # Registration recovers body origins even when the vertex centroid is + # offset. Sphere orientation does not affect its rendered appearance. + _, centers = roma.rigid_points_registration( + geometry.local_vertices.to(device)[None].expand_as(predicted), + predicted, + weights=mask[None].expand(predicted.shape[:-1]).float(), + ) + return centers + + +def _set_prediction(visuals: list[VisualObject], centers: torch.Tensor) -> None: + """Place render-only spheres beside the dynamic scene.""" + poses = torch.zeros(len(visuals), 7, device=centers.device) + poses[:, :3] = centers + centers.new_tensor(_OFFSET) + poses[:, 6] = 1 + for index in range(len(visuals)): + visuals[index].set_local_pose(poses[index : index + 1]) + + +def _capture( + sim: SimulationManager, camera: Camera, frame: int, stride: int +) -> np.ndarray: + """Read a synchronized engine image and label its two sides.""" + import cv2 + + sim.sync_render_state() + camera.update() + rgb = camera.get_data()["color"][0, ..., :3].cpu().numpy().copy() + if np.issubdtype(rgb.dtype, np.floating): + rgb = (np.clip(rgb, 0, 1) * 255).astype(np.uint8) + cv2.rectangle(rgb, (0, 0), (_WIDTH, 44), (16, 22, 30), -1) + cv2.putText( + rgb, + "NEWTON SIMULATION", + (90, 29), + cv2.FONT_HERSHEY_SIMPLEX, + 0.7, + (255, 255, 255), + 1, + ) + cv2.putText( + rgb, + "RIGIDFORMER PREDICTION", + (555, 29), + cv2.FONT_HERSHEY_SIMPLEX, + 0.7, + (255, 255, 255), + 1, + ) + phase = "observed warm-up" if frame < _WARMUP else "free model rollout" + cv2.rectangle(rgb, (0, _HEIGHT - 32), (_WIDTH, _HEIGHT), (16, 22, 30), -1) + cv2.putText( + rgb, + f"t={frame * stride / _PHYSICS_HZ:.3f}s | {phase}", + (16, _HEIGHT - 11), + cv2.FONT_HERSHEY_SIMPLEX, + 0.5, + (255, 255, 255), + 1, + ) + return rgb + + +def run_demo(args: argparse.Namespace) -> None: + """Run the engine comparison and save a video and aligned trajectories. + + Args: + args: Parsed ``learning/rigidformer demo`` command-line options. + """ + import imageio.v2 as imageio + from embodichain.lab.sim import ( + SimulationManager, + SimulationManagerCfg, + VisualMaterialCfg, + ) + from embodichain.lab.sim.cfg import ( + LightCfg, + MassPropertiesCfg, + NewtonPhysicsCfg, + NewtonRigidBodyMaterialCfg, + RenderCfg, + RigidBodyPhysicsCfg, + RigidObjectCfg, + VisualObjectCfg, + ) + from embodichain.lab.sim.sensors import CameraCfg + from embodichain.lab.sim.shapes import SphereCfg + from embodichain_models.dynamics.rigidformer import load_rigidformer_checkpoint + from embodichain_models.dynamics.rigidformer.data import _verified_manifest + + if not _WARMUP < args.frames <= 480 or args.stride < 1 or args.fps < 1: + raise ValueError("frames must be in [5, 480]; stride and fps must be positive") + if not np.isfinite(args.hold_seconds) or args.hold_seconds < 0: + raise ValueError("hold-seconds must be finite and nonnegative") + device = torch.device(args.device) + if device.type != "cuda" or not torch.cuda.is_available(): + raise ValueError("The engine/RigidFormer demo requires a CUDA device") + device = torch.device("cuda", device.index if device.index is not None else 0) + torch.cuda.set_device(device) + if args.output.exists(): + raise FileExistsError(args.output) + if not args.data_root.is_dir() or not args.checkpoint.is_file(): + raise FileNotFoundError( + "Prepare MOVi-Sphere data and weights first: embodichain run-example learning/rigidformer " + "prepare --dataset movi_sphere --assets all --limit-per-split 1; " + "or pass --data-root and --checkpoint" + ) + manifest = _verified_manifest(args.data_root) + if manifest.get("dataset") != "movi_sphere": + raise ValueError("This sphere demo requires prepared MOVi-Sphere geometry") + model = load_rigidformer_checkpoint( + args.checkpoint, device=device, attention_backend=args.attention_backend + ) + geometry = _load_geometry(args, model.reference_config) + args.output.mkdir(parents=True) + sim = None + writer = None + camera = None + window = None + bodies: list[RigidObject] = [] + visuals: list[VisualObject] = [] + try: + sim = SimulationManager( + SimulationManagerCfg( + headless=True, + device=str(device), + gpu_id=device.index, + width=_WIDTH, + height=_HEIGHT, + render_cfg=RenderCfg(renderer="hybrid"), + physics_cfg=NewtonPhysicsCfg( + physics_dt=1 / _PHYSICS_HZ, + num_substeps=2, + solver_cfg={"solver_type": "xpbd"}, + ), + ) + ) + sim.add_light( + LightCfg(uid="demo_light", intensity=100, radius=30, init_pos=(0, -4, 9)) + ) + count = len(geometry.object_ids) + starts = torch.tensor( + [[-3.35, -0.2, 2.0], [-1.35, -0.2, 2.0], [-2.35, 0.6, 3.6]] + )[:count] + velocities = torch.tensor([[1.4, 0, 0], [-1.4, 0, 0], [0, -0.2, -1.0]])[:count] + for index in range(count): + physical = geometry.parameters[index] + material = VisualMaterialCfg(base_color=list(_COLORS[index]), roughness=0.4) + shape = SphereCfg( + radius=float(geometry.radii[index]), visual_material=material + ) + bodies.append( + sim.add_rigid_object( + RigidObjectCfg( + uid=f"simulation_ball_{index}", + shape=shape, + init_pos=starts[index].tolist(), + body_type="dynamic", + attrs=RigidBodyPhysicsCfg( + mass_props=MassPropertiesCfg(mass=float(physical[0])), + material_props=NewtonRigidBodyMaterialCfg( + dynamic_friction=float(physical[1]), + restitution=float(physical[2]), + ), + ), + ) + ) + ) + visuals.append( + sim.add_visual_object( + VisualObjectCfg( + uid=f"prediction_ball_{index}", + shape=shape, + init_pos=(starts[index] + starts.new_tensor(_OFFSET)).tolist(), + ) + ) + ) + camera = sim.add_sensor( + CameraCfg( + uid="demo_camera", + width=_WIDTH, + height=_HEIGHT, + intrinsics=(850, 850, _WIDTH / 2, _HEIGHT / 2), + extrinsics=CameraCfg.ExtrinsicsCfg(eye=_EYE, target=_TARGET), + ) + ) + sim.prepare() + for index in range(count): + bodies[index].set_velocity(lin_vel=velocities[index : index + 1].to(device)) + if not args.headless: + if not sim.open_window(): + raise RuntimeError( + "Unable to open the engine window; use --headless to record" + ) + window = sim.get_world().get_windows() + window.set_look_at( + eye=np.array(_EYE), look_at=np.array(_TARGET), up=np.array([0, 0, 1]) + ) + window.set_window_title("Newton (left) | RigidFormer (right)") + writer = imageio.get_writer( + str(args.output / "rollout.mp4"), + fps=args.fps, + codec="libx264", + macro_block_size=None, + ) + observed, warmup = [], [] + predicted = None + print("Left: Newton dynamics. Right: RigidFormer after four observed frames.") + for frame in range(args.frames): + tick = time.perf_counter() + if frame: + # Publish only the sampled state after its matching model pose + # is written; intermediate physics steps remain unrendered. + for _ in range(args.stride): + sim.update(step=1, render_final_step=False) + poses = _read_poses(bodies) + observed.append(poses[:, :3]) + if frame < _WARMUP: + warmup.append(_vertices(poses, geometry)) + _set_prediction(visuals, poses[:, :3]) + else: + assert predicted is not None + _set_prediction(visuals, predicted[frame]) + image = _capture(sim, camera, frame, args.stride) + writer.append_data(image) + if frame == min(20, args.frames - 1): + imageio.imwrite(args.output / "frame.png", image) + if frame == _WARMUP - 1: + predicted = _predict(model, torch.stack(warmup), geometry, args) + if not args.headless: + time.sleep(max(0.0, 1 / args.fps - (time.perf_counter() - tick))) + writer.close() + writer = None + assert predicted is not None + simulated = torch.stack(observed).cpu().numpy() + predicted_centers = predicted.cpu().numpy() + errors = np.linalg.norm(simulated - predicted_centers, axis=-1) + np.savez_compressed( + args.output / "trajectory.npz", + simulated_centers=simulated, + predicted_centers=predicted_centers, + center_error_m=errors, + ) + summary = { + "physics_backend": "newton", + "solver": "xpbd", + "physics_step_seconds": 1 / _PHYSICS_HZ, + "scene": geometry.scene, + "object_ids": geometry.object_ids, + "frames": args.frames, + "warmup_frames": _WARMUP, + "model_stride": args.stride, + "playback_fps": args.fps, + "attention_backend": args.attention_backend, + "checkpoint_sha256": sha256(args.checkpoint), + "data_manifest_sha256": sha256(args.data_root / "manifest.json"), + "dataset_scope": manifest["scope"], + "model_rollout_calls": 1, + "sphere_radii_m": geometry.radii.tolist(), + "mass_friction_restitution": geometry.parameters.tolist(), + "initial_positions_m": starts.tolist(), + "initial_velocities_m_per_s": velocities.tolist(), + "elapsed_simulation_seconds": (args.frames - 1) * args.stride / _PHYSICS_HZ, + "mean_center_error_at_last_frame_m": float(errors[-1].mean()), + "max_center_error_m": float(errors.max()), + } + (args.output / "summary.json").write_text(json.dumps(summary, indent=2) + "\n") + print(json.dumps(summary, indent=2)) + print(f"Saved engine video: {args.output / 'rollout.mp4'}") + if not args.headless: + time.sleep(args.hold_seconds) + except Exception as error: + (args.output / "failure.json").write_text( + json.dumps({"type": type(error).__name__, "message": str(error)}) + "\n" + ) + raise + finally: + try: + if writer is not None: + writer.close() + finally: + if sim is not None: + sim.destroy(exit_process=False) + bodies.clear() + visuals.clear() + camera = None + window = None + SimulationManager.flush_cleanup_queue() From 7928471ec1c9a90ec874ddca613dc1b8f83b7091 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Wed, 28 Oct 2026 16:34:34 +0000 Subject: [PATCH 4/4] style(examples): trim RigidFormer file endings --- examples/learning/rigidformer/README.md | 1 - examples/learning/rigidformer/configs/sl.movi_a.yaml | 1 - examples/learning/rigidformer/configs/sl.movi_b.yaml | 1 - examples/learning/rigidformer/configs/sl.movi_sphere.yaml | 1 - 4 files changed, 4 deletions(-) diff --git a/examples/learning/rigidformer/README.md b/examples/learning/rigidformer/README.md index 9db2013e90..0ca81ca2a0 100644 --- a/examples/learning/rigidformer/README.md +++ b/examples/learning/rigidformer/README.md @@ -226,4 +226,3 @@ This is a numerical-integration check. Full published-accuracy reproduction additionally requires frozen official weights/data/evaluation, or complete retraining with the declared compute and selection protocol. No short-run comparison is labelled a full training reproduction. - diff --git a/examples/learning/rigidformer/configs/sl.movi_a.yaml b/examples/learning/rigidformer/configs/sl.movi_a.yaml index a1f15207b1..2584b0e487 100644 --- a/examples/learning/rigidformer/configs/sl.movi_a.yaml +++ b/examples/learning/rigidformer/configs/sl.movi_a.yaml @@ -13,4 +13,3 @@ trainer: batch_size: 32 seed: 42 checkpoint_interval: 5000 - diff --git a/examples/learning/rigidformer/configs/sl.movi_b.yaml b/examples/learning/rigidformer/configs/sl.movi_b.yaml index 9b368c194a..f69d31b22b 100644 --- a/examples/learning/rigidformer/configs/sl.movi_b.yaml +++ b/examples/learning/rigidformer/configs/sl.movi_b.yaml @@ -13,4 +13,3 @@ trainer: batch_size: 21 seed: 42 checkpoint_interval: 5000 - diff --git a/examples/learning/rigidformer/configs/sl.movi_sphere.yaml b/examples/learning/rigidformer/configs/sl.movi_sphere.yaml index a87e51899b..f2644c6c56 100644 --- a/examples/learning/rigidformer/configs/sl.movi_sphere.yaml +++ b/examples/learning/rigidformer/configs/sl.movi_sphere.yaml @@ -13,4 +13,3 @@ trainer: batch_size: 32 seed: 42 checkpoint_interval: 5000 -