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..a8c388e841 100644 --- a/agent_context/topics/neural-dynamics/neural-dynamics.md +++ b/agent_context/topics/neural-dynamics/neural-dynamics.md @@ -57,8 +57,13 @@ 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. +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. ## 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/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 new file mode 100644 index 0000000000..0ca81ca2a0 --- /dev/null +++ b/examples/learning/rigidformer/README.md @@ -0,0 +1,228 @@ +# 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 +``` + +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. +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. + +## 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 +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..e44598c023 --- /dev/null +++ b/examples/learning/rigidformer/__main__.py @@ -0,0 +1,251 @@ +# ---------------------------------------------------------------------------- +# 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 RigidFormer assets, run an engine demo, evaluate and export weights.""" + +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("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) + 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 == "demo": + from ._demo import run_demo + + run_demo(args) + 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/_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() 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..2584b0e487 --- /dev/null +++ b/examples/learning/rigidformer/configs/sl.movi_a.yaml @@ -0,0 +1,15 @@ +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..f69d31b22b --- /dev/null +++ b/examples/learning/rigidformer/configs/sl.movi_b.yaml @@ -0,0 +1,15 @@ +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..f2644c6c56 --- /dev/null +++ b/examples/learning/rigidformer/configs/sl.movi_sphere.yaml @@ -0,0 +1,15 @@ +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..3e78cf1afe --- /dev/null +++ b/examples/learning/rigidformer/research/qualification.json @@ -0,0 +1,156 @@ +{ + "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": "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": "4e037cfb591fc73f9a47fd42d2927a7c9246583a78723b029edc54de99f4c2ed", + "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": "post-review-recovery/summary.json", + "compared_exactly": [ + "learner", + "training_states", + "sampler", + "rng" + ], + "exported_tensors": { + "ema": 485, + "raw": 485 + }, + "model_source_commit": "acb01d18e28832557dbd4bff39ede7b6d12bb02f", + "previous_artifact": "installed-final-recovery/summary.json", + "previous_sha256": "6619a190694b2a4e8bfb954a6ee4fa6c17703b0bd77f081d0398184aede3ddf7", + "sha256": "2daad678a5b633615a871db5b2a610bae55819650779a00567dee10273a4545b", + "stop_after": 2, + "updates": 4 + }, + "index_benchmark": { + "artifact": "qualified-index.json", + "sha256": "4b07998f2222853c8e789aaa849235067f0740bf7eeb83b62a0caf912436e5e5" + }, + "two_rank_nccl": { + "artifact": "post-review-two-rank/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, + "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, + "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, + "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." + } +} diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/data.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/data.py index 4e512ad991..ae0b3dc62e 100644 --- a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/data.py +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/data.py @@ -210,7 +210,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 @@ -277,14 +284,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, @@ -294,7 +309,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 @@ -308,7 +327,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 d3a85b3b74..cc4092ebb9 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"