diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index d01eba53f1..b3d4a4cf75 100644 --- a/agent_context/MAP.yaml +++ b/agent_context/MAP.yaml @@ -270,10 +270,11 @@ topics: status: active - id: neural-dynamics title: Neural Dynamics Models and Training - aliases: [neural dynamics, neural robot dynamics, NeRD, 神经动力学, 动力学模型训练, NeRD 示例] + aliases: [neural dynamics, neural robot dynamics, NeRD, 神经动力学, 动力学模型训练, NeRD 示例, RigidFormer, RigidFormer 示例] keywords: [NeRDStateCodec, NeRDModel, DynamicsGenerationCfg, DynamicsShardWriter, DynamicsWindowDataset, DynamicsTrainer, embodichain_models, - learning/nerd, nerd reference assets, + learning/nerd, nerd reference assets, RigidFormerCfg, RigidFormerModel, + RigidBodySequenceSource, build_rigidformer_trainer, learning/rigidformer, 官方训练精度对比] paths: [topics/neural-dynamics/neural-dynamics.md] source_of_truth: @@ -282,9 +283,13 @@ topics: - examples/learning/nerd/_resources.py - examples/learning/nerd/_support/ - examples/learning/nerd/configs/ + - extensions/embodichain_models/embodichain_models/dynamics/rigidformer/ + - examples/learning/rigidformer/ watch_paths: [extensions/embodichain_models/, examples/learning/nerd/, + examples/learning/rigidformer/, tests/learning/models/test_rigidformer.py, + tests/learning/models/test_rigidformer_data.py, tests/examples/rigidformer_reference/, docs/source/features/neural_dynamics.md, - docs/source/api_reference/model_catalogs.rst] + docs/source/api_reference/model_catalogs.rst, docs/source/api_reference/model_rigidformer.rst] related_topics: [supervised-learning, learning-layers, data-pipeline, data-assets] status: active - id: learning-layers diff --git a/agent_context/topics/data-pipeline/data-pipeline.md b/agent_context/topics/data-pipeline/data-pipeline.md index 612ba903f2..1ddb557d85 100644 --- a/agent_context/topics/data-pipeline/data-pipeline.md +++ b/agent_context/topics/data-pipeline/data-pipeline.md @@ -24,6 +24,7 @@ belongs to [data-assets](../data-assets/data-assets.md). | File-only validation and grouped quality splits | `embodichain/data_pipeline/datasets/inspection.py`; `embodichain dataset validate\|split` | | Durable commit diagnostics and metadata repair | `embodichain/data_pipeline/recording/journal.py`; `embodichain dataset recover` | | NeRD-specific HDF5 windows and teacher collection | `extensions/embodichain_models/embodichain_models/dynamics/nerd/data.py` and `examples/learning/nerd/_support/_generation.py`; contracts in [neural dynamics](../neural-dynamics/neural-dynamics.md) | +| RigidFormer aligned geometry windows | `extensions/embodichain_models/embodichain_models/dynamics/rigidformer/data.py`; model ownership in [neural dynamics](../neural-dynamics/neural-dynamics.md) | | Isolated LeRobot 0.6.1 language/depth conversion | `tools/lerobot_export/` (Python >=3.12, independent dependencies) | Read [online sampling](online-sampling.md) for worker states, refill, shared diff --git a/agent_context/topics/learning-layers/learning-layers.md b/agent_context/topics/learning-layers/learning-layers.md index e1d5ff3c23..600ed99a39 100644 --- a/agent_context/topics/learning-layers/learning-layers.md +++ b/agent_context/topics/learning-layers/learning-layers.md @@ -10,10 +10,10 @@ returns TCP poses and geometric Jacobians from one FK pass, retaining dJ/dq. `learning/layers/` owns `KinematicsLayer` and `DifferentiableIKLayer` Torch module interfaces. `learning/models/` owns general MLPs, feature normalization and configuration-driven MLP construction. These shared components do not import -RL, TensorDict or simulation. Concrete NeRD models, codecs and the shared-SL -recipe live in the bundled `embodichain_models.dynamics.nerd` extension; -teacher collection, historical training and research belong to -`examples/learning/nerd/`. See +RL, TensorDict or simulation. Concrete NeRD and RigidFormer models, model-specific codecs/data and shared-SL +recipes live in the bundled `embodichain_models.dynamics` families; +teacher collection, historical training and research belong to the corresponding +`examples/learning/` showcase. See [neural dynamics](../neural-dynamics/neural-dynamics.md) for model, SL integration, dataset and checkpoint ownership. RL policy interfaces, actor-critic composition and policy registration/builders belong to `learning/rl/policies/`; RL code consumes diff --git a/agent_context/topics/neural-dynamics/neural-dynamics.md b/agent_context/topics/neural-dynamics/neural-dynamics.md index 747cb0b625..8f0e2d7218 100644 --- a/agent_context/topics/neural-dynamics/neural-dynamics.md +++ b/agent_context/topics/neural-dynamics/neural-dynamics.md @@ -10,8 +10,8 @@ runtime registry. Shared MLPs and feature normalization stay in core The [model catalog contract](../../../extensions/embodichain_models/README.md) owns domain and public-import naming. Policy models predict actions; dynamics -models predict future state or geometry. RWM and RigidFormer follow-ups use -`dynamics.rwm` and `dynamics.rigidformer`; only `dynamics.nerd` is implemented here. +models predict future state or geometry. NeRD and RigidFormer are implemented +at `dynamics.nerd` and `dynamics.rigidformer`; the RWM follow-up uses `dynamics.rwm`. World-model naming does not introduce a parallel `world_models` domain. | Request | Owner | @@ -38,7 +38,26 @@ Scene API. Native learned-solver deployment and functional Gym integration are separate consumers. There are no NeRD-specific core `compute/dynamics`, `data_pipeline/dynamics`, `learning/dynamics` modules or core `dynamics` command. -## Resource and physical identity +## RigidFormer geometry model + +`extensions/embodichain_models/embodichain_models/dynamics/rigidformer/` owns +`model.py`, aligned CPU windows in `data.py`, objective/batch conversion in +`supervised.py` and the installable +`sl_recipe:build_rigidformer_trainer` factory. Package imports stay lazy; only +model construction loads the pinned optional upstream dependency. Core SL and +generic data-pipeline modules do not import this family. The model has no action +input (`action_conditioned=False`) and does not qualify action gradients or APG. + +`embodichain run-example learning/rigidformer` runs the private Box2D comparison +under the example's `_support/`. The [example guide](../../../examples/learning/rigidformer/README.md) +owns source/data provenance, finite-padding and pairwise-memory boundaries; the +[frozen protocol](../../../examples/learning/rigidformer/research/README.md) owns +reference seeds, budgets and acceptance gates. Existing tests in +`tests/learning/models/test_rigidformer*.py` cover the production geometry, +weighted loss/gradient, masking, budget and exact recovery contracts. New runs +bind executing source hashes; historical reports retain their original identity. + +## NeRD resource and physical identity `prepare --robot cartpole|ant --resource-root DIR` downloads the pinned dataset, portable model and raw reference assets. Revisions and SHA256 checks are owned by diff --git a/agent_context/topics/supervised-learning/supervised-learning.md b/agent_context/topics/supervised-learning/supervised-learning.md index 982c189bd7..ef9565e120 100644 --- a/agent_context/topics/supervised-learning/supervised-learning.md +++ b/agent_context/topics/supervised-learning/supervised-learning.md @@ -19,10 +19,9 @@ operations remain in `compute`; source storage belongs to data adapters. - `embodichain train-sl --config FILE [--resume FILE]` lazily resolves a `factory: module:function`. The factory receives the full model/data/objective/ optimizer/trainer/evaluation configuration and returns a `SupervisedTrainer`. -- The [NeRD recipe](../neural-dynamics/neural-dynamics.md) is an installable factory - at `embodichain_models.dynamics.nerd.sl_recipe:build_supervised_trainer`. - The core does not import the model extension or example helpers. The NeRD - showcase's default `train` command uses this same common recipe. +- [Model-specific factories](../neural-dynamics/neural-dynamics.md) are installable + from their `embodichain_models` family packages. The core does not import the + model extension or example helpers; model recipes return the common trainer. Objectives own loss masks, normalization, integration and multi-step graph construction. The optimizer does not infer action semantics or detach histories. diff --git a/docs/architecture/curated.json b/docs/architecture/curated.json index 4eeee8a4b8..a1f308db22 100644 --- a/docs/architecture/curated.json +++ b/docs/architecture/curated.json @@ -3342,13 +3342,13 @@ }, { "id": "neural-dynamics", - "label": "NeRD model extension", + "label": "Official dynamics models", "kind": "package", "topic_ids": [ "neural-dynamics" ], - "summary": "Bundles the causal NeRD network, physical codec, trajectory windows and supervised recipe in the main wheel.", - "details": "embodichain_models.dynamics.nerd owns the concrete model, portable checkpoints, physical representations, HDF5 windows and installable SL recipe. Generic MLP/normalization and supervised orchestration stay in core. The model extension imports neither the showcase nor simulation; collection, historical training and reference research belong to the example.", + "summary": "Bundles NeRD and optional RigidFormer models, model-specific data adapters and shared-SL recipes in the main wheel.", + "details": "embodichain_models.dynamics.nerd owns portable neural dynamics and HDF5 windows; dynamics.rigidformer owns pinned geometry prediction, aligned point windows and the valid-object weighted objective. Both provide installable shared-SL factories. Generic MLP/normalization and supervised orchestration stay in core. The model extension imports neither the showcase nor simulation; collection, historical training and reference research belong to the example.", "boundaries": [ "The core does not import concrete model implementations.", "Data/model/reference-asset payloads are external resources; numerical and physical identities remain explicit." diff --git a/docs/source/api_reference/model_catalogs.rst b/docs/source/api_reference/model_catalogs.rst index ba6c17cf3c..b0f66dd341 100644 --- a/docs/source/api_reference/model_catalogs.rst +++ b/docs/source/api_reference/model_catalogs.rst @@ -252,3 +252,15 @@ working directory and example-runner bootstrap. :members: sample .. autofunction:: build_supervised_trainer + +RigidFormer geometry prediction +------------------------------- + +The optional ``embodichain_models.dynamics.rigidformer`` family owns aligned +geometry windows, its pinned upstream adapter, valid-object weighted objective +and installable SL recipe. It predicts geometry without action inputs. + +.. toctree:: + :maxdepth: 1 + + model_rigidformer diff --git a/docs/source/api_reference/model_rigidformer.rst b/docs/source/api_reference/model_rigidformer.rst new file mode 100644 index 0000000000..3f6bca6dd7 --- /dev/null +++ b/docs/source/api_reference/model_rigidformer.rst @@ -0,0 +1,114 @@ +RigidFormer geometry dynamics +============================= + +The lazy optional adapter binds the pinned RigidFormer implementation to causal +point-cloud windows and shared supervised updates. Reference geometry and point +correspondence are explicit. The model has no action input and does not qualify +action gradients, APG or a native physics-solver interface. Install the +``rigidformer`` extra for this family. The Box2D comparison belongs to the +``learning/rigidformer`` example and retains its frozen protocol and provenance. + +Public entry points +------------------- + +.. currentmodule:: embodichain_models.dynamics.rigidformer + +.. autosummary:: + + RIGIDFORMER_REVISION + RigidFormerCfg + RigidFormerModel + RigidBodySequenceSource + RigidFormerObjective + rigidformer_batch + build_rigidformer_trainer + +.. autodata:: RIGIDFORMER_REVISION + +.. autoclass:: RigidFormerCfg + +.. autoclass:: RigidFormerModel + +.. autoclass:: RigidBodySequenceSource + +.. autoclass:: RigidFormerObjective + +.. autofunction:: rigidformer_batch + +.. autofunction:: build_rigidformer_trainer + +Aligned geometry windows +------------------------ + +The source owns a bounded CPU corpus and returns independent previous, current, +reference and next-frame point arrays without crossing episode boundaries. +The caller owns sampling RNG; material point identities and finite object/point +lengths remain explicit. + +.. currentmodule:: embodichain_models.dynamics.rigidformer.data + +.. autosummary:: + + RigidBodySequenceSource + +.. autoclass:: RigidBodySequenceSource + :members: get_batch, sample_batch + +Optional network adapter +------------------------ + +The optional upstream model loads at construction. Parameter keys and numerical +behavior follow the pinned upstream revision. The per-call budget bounds the +dense nearest-neighbour distance intermediate. Unequal point lengths require a +padded point axis of at most 128 for forward and rollout. + +.. currentmodule:: embodichain_models.dynamics.rigidformer.model + +.. autosummary:: + + RIGIDFORMER_REVISION + RigidFormerCfg + RigidFormerModel + +.. autodata:: RIGIDFORMER_REVISION + +.. autoclass:: RigidFormerCfg + :members: + +.. autoclass:: RigidFormerModel + :members: forward, rollout + +Objective and causal batch conversion +------------------------------------- + +The objective preserves the upstream position/acceleration losses. Effective +weight is the sum of valid objects, including when scenes have different object +counts. Targets remain separate from causal model inputs. + +.. currentmodule:: embodichain_models.dynamics.rigidformer.supervised + +.. autosummary:: + + RigidFormerObjective + rigidformer_batch + +.. autoclass:: RigidFormerObjective + :members: compute_loss + +.. autofunction:: rigidformer_batch + +Installable supervised recipe +----------------------------- + +Use ``factory: +embodichain_models.dynamics.rigidformer.sl_recipe:build_rigidformer_trainer`` +with ``embodichain train-sl``. The recipe samples fixed tensor-only archives and +returns the common trainer with source/content/configuration recovery identity. + +.. currentmodule:: embodichain_models.dynamics.rigidformer.sl_recipe + +.. autosummary:: + + build_rigidformer_trainer + +.. autofunction:: build_rigidformer_trainer diff --git a/examples/README.md b/examples/README.md index f252b0577d..342ed013e0 100644 --- a/examples/README.md +++ b/examples/README.md @@ -7,6 +7,7 @@ embodichain run-example --list 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/rigidformer --help embodichain run-example sim/atomic_action --skill place --variants 8 embodichain run-example sim/atomic_action/coordinated_pickment --help ``` @@ -42,3 +43,7 @@ The [Atomic Skill example catalog](sim/atomic_action/README.md) owns all per-ski coordinated-manipulation and recovery demonstrations. The two source-checkout walkthroughs under `scripts/tutorials/atomic_action/` teach single-arm and dual-arm API composition; they reuse the canonical example setup. + +The [RigidFormer showcase](learning/rigidformer/README.md) owns the pinned Box2D +reference comparison. Its installable `dynamics.rigidformer` model, geometry +source and objective use the shared SL trainer for fixed-corpus training. diff --git a/examples/learning/rigidformer/README.md b/examples/learning/rigidformer/README.md new file mode 100644 index 0000000000..c6511bef15 --- /dev/null +++ b/examples/learning/rigidformer/README.md @@ -0,0 +1,64 @@ +# RigidFormer geometry prediction + +The installable model family is `embodichain_models.dynamics.rigidformer`. It +wraps `lucidrains/rigidformer` at the pinned revision +`4ab4edac7271b49484748d6bf8086a497f13272d` and trains through the shared +`embodichain.learning.sl` contracts. Install `embodichain[rigidformer]` only when +using this model. The upstream model and Box2D reference recipe retain the +complete MIT license in the model package's `THIRD_PARTY_LICENSES`. + +The ordinary optional extra uses the public-index pin `rigidformer==0.2.0`, +without a Git dependency in EmbodiChain's published metadata: + +```bash +python -m pip install 'embodichain[rigidformer]' +``` + +All four runtime Python files in the [PyPI 0.2.0 wheel](https://pypi.org/project/rigidformer/0.2.0/) +were verified byte-for-byte against the frozen commit above. This establishes +that distribution's runtime-source equivalence; the version number alone does +not qualify a reference environment. For exact Git installation in an isolated +`py311` environment, use [requirements-reference.txt](requirements-reference.txt) +and the [reference setup](research/README.md). The acceptance runner still +checks the checkout and installed Python bytes against the pinned Git blobs. + +For a fixed tensor-only PyTorch geometry archive, use the installable factory: + +```bash +embodichain train-sl --config examples/learning/rigidformer/configs/sl.train.yaml +embodichain train-sl --config my-recipe.yaml --resume outputs/rigidformer/sl/last.pt +``` + +Set `data.path` in the configuration to your archive. Paths resolve from the +caller's working directory. Its fields are `positions[E,T,O,P,3]`, +`properties[E,O,A]` (or `[E,O,P,A]`) and positive `step_dt` in seconds. Optional +`episode_lengths`, `object_lens` and `object_point_lens` describe finite padding. +Point identities remain aligned across frames, and the owned previous-frame +geometry is the reference for each sampled window. Sampling preserves episode +boundaries and explicit resident/batch byte budgets. + +`RigidFormerModel.action_conditioned` is `False`. The adapter exposes future +geometry prediction and supervised parameter training. Action-gradient/APG and +native physics-solver deployment are outside its qualified capabilities. The +dense nearest-neighbour buffer is checked against `max_pairwise_bytes`; this +is a bound on that intermediate alone. Unequal point lengths require padded +`P <= 128`, including at rollout. Supplied point counts must be at least +`num_anchors`, including placeholder objects; fixed anchors index valid points. +The objective weights each scene contribution by its number of valid objects. + +The reference experiment is an example, with private support outside the +installable model package: + +```bash +embodichain run-example learning/rigidformer --help +embodichain run-example learning/rigidformer \ + --reference /tmp/rigidformer-reference \ + --output /tmp/rigidformer-acceptance --threads 4 +``` + +See the [frozen reference protocol](research/README.md) for source/data checks, +friction/collision modes, seeds, budgets and acceptance thresholds. Every new +manifest records the executing adapter, SL and runner source hashes. Historical +reports remain bound to their original source and data; relocation does not +change or requalify that evidence. Downloaded data, upstream checkouts, weights +and generated reports remain outside Git and the wheel. diff --git a/examples/learning/rigidformer/__init__.py b/examples/learning/rigidformer/__init__.py new file mode 100644 index 0000000000..5b8a0bd583 --- /dev/null +++ b/examples/learning/rigidformer/__init__.py @@ -0,0 +1,22 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + + +"""RigidFormer reference showcase.""" + +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..89af68409a --- /dev/null +++ b/examples/learning/rigidformer/__main__.py @@ -0,0 +1,26 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Run the pinned RigidFormer reference acceptance experiment.""" + +from __future__ import annotations + +__all__: list[str] = [] + +if __name__ == "__main__": + from ._support._acceptance import _main + + _main() diff --git a/examples/learning/rigidformer/_support/__init__.py b/examples/learning/rigidformer/_support/__init__.py new file mode 100644 index 0000000000..addf554b6c --- /dev/null +++ b/examples/learning/rigidformer/_support/__init__.py @@ -0,0 +1,22 @@ +# ---------------------------------------------------------------------------- +# 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 RigidFormer reference workflow support.""" + +from __future__ import annotations + +__all__: list[str] = [] diff --git a/examples/learning/rigidformer/_support/_acceptance.py b/examples/learning/rigidformer/_support/_acceptance.py new file mode 100644 index 0000000000..097a850b22 --- /dev/null +++ b/examples/learning/rigidformer/_support/_acceptance.py @@ -0,0 +1,626 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Pair the pinned Box2D reference recipe with the shared SL learner. + +The upstream examples are loaded from an explicitly supplied, revision-checked +checkout. They retain their own MIT license and are not copied into production. +""" + +from __future__ import annotations + +import argparse +from copy import deepcopy +import hashlib +import importlib.metadata +import importlib.util +import json +from pathlib import Path +import subprocess +import time +from typing import Any + +import numpy as np +import torch +from torch.utils.data.dataloader import default_collate + +from embodichain_models.dynamics.rigidformer.model import ( + RIGIDFORMER_REVISION, + RigidFormerCfg, + RigidFormerModel, +) +from embodichain.learning.sl import SupervisedLearner +from embodichain_models.dynamics.rigidformer.supervised import ( + RigidFormerObjective, + rigidformer_batch, +) + +__all__: list[str] = [] + +_NUMERICAL_ATOL = 1e-5 +_NUMERICAL_RTOL = 1e-4 + + +def _verify_sources(path: Path, installed: Path) -> dict[str, Any]: + """Check source bytes against Git objects, including a dirty checkout.""" + revision = subprocess.check_output( + ["git", "-C", str(path), "rev-parse", "HEAD"], text=True + ).strip() + if revision != RIGIDFORMER_REVISION: + raise ValueError(f"Expected reference revision {RIGIDFORMER_REVISION}") + package_files = subprocess.check_output( + [ + "git", + "-C", + str(path), + "ls-tree", + "-r", + "--name-only", + revision, + "rigidformer", + ], + text=True, + ).splitlines() + sources = ["examples/box2d_dataset.py", "examples/train_box2d.py"] + sources += [name for name in package_files if name.endswith(".py")] + verified = {} + for name in sources: + pinned = subprocess.check_output( + ["git", "-C", str(path), "show", f"{revision}:{name}"] + ) + if (path / name).read_bytes() != pinned: + raise ValueError(f"Reference source differs from pinned Git blob: {name}") + if name.startswith("rigidformer/"): + destination = installed / Path(name).relative_to("rigidformer") + if destination.read_bytes() != pinned: + raise ValueError( + f"Installed source differs from pinned Git blob: {name}" + ) + verified[name] = { + "git_blob": subprocess.check_output( + ["git", "-C", str(path), "rev-parse", f"{revision}:{name}"], text=True + ).strip(), + "sha256": hashlib.sha256(pinned).hexdigest(), + } + return verified + + +def _load_reference(path: Path) -> tuple[Any, dict[str, Any]]: + # Locate the package without importing it before its source is verified. + package = importlib.util.find_spec("rigidformer") + if package is None or package.origin is None: + raise ImportError("The optional rigidformer dependency is unavailable") + verified = _verify_sources(path, Path(package.origin).parent) + source = path / "examples" / "box2d_dataset.py" + spec = importlib.util.spec_from_file_location("rigidformer_box2d_reference", source) + if spec is None or spec.loader is None: + raise ValueError("Reference Box2D data module is unavailable") + module = importlib.util.module_from_spec(spec) + # Execute the checked bytes directly rather than accepting a cached bytecode file. + exec(compile(source.read_bytes(), str(source), "exec"), module.__dict__) + return module, verified + + +def _digest(values: Any) -> str: + digest = hashlib.sha256() + for value in values: + if isinstance(value, torch.Tensor): + value = value.detach().cpu().contiguous().numpy() + digest.update(np.asarray(value).tobytes()) + return digest.hexdigest() + + +def _write(path: Path, value: Any) -> None: + temporary = path.with_suffix(".tmp") + temporary.write_text(json.dumps(value, indent=2, allow_nan=False) + "\n") + temporary.replace(path) + + +def _sync(device: str) -> None: + if device.startswith("cuda"): + torch.cuda.synchronize(device) + + +def _sample(dataset: Any, size: int, device: str) -> dict[str, torch.Tensor]: + raw = default_collate( + [dataset[np.random.randint(len(dataset))] for _ in range(size)] + ) + return {key: value.to(device) for key, value in raw.items()} + + +def _rollout( + model: Any, trajectory: Any, *, collision: bool, device: str +) -> dict[str, Any]: + from rigidformer import RigidformerRolloutWrapper + + positions, props = trajectory + first = 10 if collision else 30 + frames = torch.from_numpy(positions[first : first + 10 * 13 : 10])[:13].to(device) + gt = frames.unsqueeze(1) + network = model.network if isinstance(model, RigidFormerModel) else model + with torch.no_grad(): + predicted = torch.stack( + RigidformerRolloutWrapper(network)( + delta_times=torch.full((1,), 1 / 6, device=device), + vertex_properties=torch.from_numpy(props).unsqueeze(0).to(device), + object_positions=[gt[:1, 0], gt[1:2, 0]], + num_steps=11, + ) + ) + constant_velocity = [gt[0], gt[1]] + for _ in range(11): + constant_velocity.append(2 * constant_velocity[-1] - constant_velocity[-2]) + constant = torch.stack(constant_velocity) + stay = gt[1].expand_as(gt) + centers = lambda value: value.mean(-2) + point_error = (predicted - gt).norm(dim=-1).mean(dim=(1, 2, 3)) + center_error = (centers(predicted) - centers(gt)).norm(dim=-1).mean(dim=(1, 2)) + constant_error = (centers(constant) - centers(gt)).norm(dim=-1).mean(dim=(1, 2)) + stay_error = (centers(stay) - centers(gt)).norm(dim=-1).mean(dim=(1, 2)) + rigidity = ( + ( + torch.cdist( + predicted[-1], + predicted[-1], + compute_mode="donot_use_mm_for_euclid_dist", + ) + - torch.cdist(gt[0], gt[0], compute_mode="donot_use_mm_for_euclid_dist") + ) + .abs() + .max() + ) + return { + "center_error_m": center_error.cpu().tolist(), + "point_error_m": point_error.cpu().tolist(), + "constant_velocity_center_error_m": constant_error.cpu().tolist(), + "stay_center_error_m": stay_error.cpu().tolist(), + "rigidity_max_distance_error_m": float(rigidity), + } + + +def _step_parity(dataset: Any, device: str) -> dict[str, Any]: + candidate = RigidFormerModel().to(device) + reference = deepcopy(candidate.network) + optimizer = torch.optim.AdamW(reference.parameters(), lr=3e-4, weight_decay=1e-2) + candidate_optimizer = torch.optim.AdamW( + candidate.parameters(), lr=3e-4, weight_decay=1e-2 + ) + batch = _sample(dataset, 8, device) + rng = torch.get_rng_state() + cuda_rng = torch.cuda.get_rng_state(device) if device.startswith("cuda") else None + loss, _ = reference(**batch) + loss.backward() + expected_gradients = { + name: parameter.grad.detach().clone() + for name, parameter in reference.named_parameters() + if parameter.grad is not None + } + torch.nn.utils.clip_grad_norm_(reference.parameters(), 1.0) + optimizer.step() + torch.set_rng_state(rng) + if cuda_rng is not None: + torch.cuda.set_rng_state(cuda_rng, device) + learner = SupervisedLearner( + candidate, RigidFormerObjective(), candidate_optimizer, max_grad_norm=1.0 + ) + # Capture unmodified gradients before the common clipping operation. + observed_gradients: dict[str, torch.Tensor] = {} + handles = [ + parameter.register_hook( + lambda gradient, name=name: observed_gradients.__setitem__( + name, gradient.detach().clone() + ) + ) + for name, parameter in candidate.network.named_parameters() + ] + metrics = learner.update(rigidformer_batch(batch)) + for handle in handles: + handle.remove() + gradient_match = all( + torch.allclose(observed_gradients[name], value, atol=1e-5, rtol=1e-4) + for name, value in expected_gradients.items() + ) + weight_match = all( + torch.allclose( + candidate.network.state_dict()[name], value, atol=1e-5, rtol=1e-4 + ) + for name, value in reference.state_dict().items() + ) + return { + "passed": bool( + gradient_match + and weight_match + and np.isclose(metrics["loss"], float(loss.detach()), atol=1e-5, rtol=1e-4) + ), + "reference_loss": float(loss), + "sl_loss": metrics["loss"], + "gradients_match": gradient_match, + "updated_parameters_match": weight_match, + "atol": 1e-5, + "rtol": 1e-4, + } + + +def _fit( + train: Any, + validation: Any, + *, + seed: int, + collision: bool, + steps: int, + device: str, + use_sl: bool, + output: Path, +) -> dict[str, Any]: + from rigidformer import Rigidformer + + np.random.seed(seed) + torch.manual_seed(seed) + cfg = RigidFormerCfg() + network_cfg = cfg.to_dict() + network_cfg.pop("max_pairwise_bytes") + model = ( + RigidFormerModel(cfg).to(device) + if use_sl + else Rigidformer(**network_cfg).to(device) + ) + optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1e-2) + learner = ( + SupervisedLearner(model, RigidFormerObjective(), optimizer, max_grad_norm=1.0) + if use_sl + else None + ) + losses, curve = [], [] + _sync(device) + started = time.monotonic() + for step in range(steps): + batch = _sample(train, 8, device) + if learner is not None: + metrics = learner.update(rigidformer_batch(batch)) + scalar_loss = metrics["loss"] + else: + loss, _ = model(**batch) + optimizer.zero_grad() + loss.backward() + torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) + optimizer.step() + scalar_loss = float(loss.detach()) + losses.append(scalar_loss) + if not np.isfinite(scalar_loss): + raise FloatingPointError("Reference recipe returned a non-finite loss") + if (step + 1) % 500 == 0: + model.eval() + measured = _rollout( + model, validation.trajectories[0], collision=collision, device=device + ) + curve.append({"step": step + 1, **measured}) + model.train() + print( + json.dumps( + { + "seed": seed, + "sl": use_sl, + "collision": collision, + "step": step + 1, + "loss": scalar_loss, + } + ), + flush=True, + ) + _sync(device) + training_seconds = time.monotonic() - started + model.eval() + single_step: dict[str, list[float]] = { + key: [] for key in ("model", "constant_velocity", "stay") + } + # Preserve the reference 32 sampled single-step windows and RNG consumption. + with torch.no_grad(): + for _ in range(32): + batch = _sample(validation, 1, device) + target = batch.pop("object_pos_next") + prediction = model(**batch).object_pos_next + for key, value in ( + ("model", prediction), + ( + "constant_velocity", + 2 * batch["object_pos"] - batch["object_pos_prev"], + ), + ("stay", batch["object_pos"]), + ): + single_step[key].append(float((value - target).norm(dim=-1).mean())) + reports = [ + _rollout(model, row, collision=collision, device=device) + for row in validation.trajectories + ] + network = model.network if isinstance(model, RigidFormerModel) else model + checkpoint = output / ("sl.pt" if use_sl else "reference.pt") + torch.save( + { + "model": network.state_dict(), + "optimizer": optimizer.state_dict(), + "seed": seed, + "steps": steps, + "config": network_cfg, + }, + checkpoint, + ) + report = { + "seed": seed, + "mode": "collision" if collision else "friction", + "engine": "sl" if use_sl else "reference", + "steps": steps, + "training_seconds": training_seconds, + "losses": losses, + "checkpoint_sha256": hashlib.sha256(checkpoint.read_bytes()).hexdigest(), + "single_step_mean_point_distance_m": { + key: float(np.mean(value)) for key, value in single_step.items() + }, + "reference_trajectory": reports[0], + "validation_trajectories": reports, + "mean_rollout_endpoint_center_error_m": float( + np.mean([row["center_error_m"][-1] for row in reports]) + ), + "mean_rollout_endpoint_point_error_m": float( + np.mean([row["point_error_m"][-1] for row in reports]) + ), + "curve": curve, + } + _write(output / ("sl.json" if use_sl else "reference.json"), report) + return report + + +def _gate(reference: list[float], candidate: list[float]) -> dict[str, Any]: + baseline, actual = np.asarray(reference), np.asarray(candidate) + absolute = bool((baseline <= 1e-5).any()) + differences = actual - baseline if absolute else actual / baseline - 1 + draws = np.random.default_rng(0).integers( + 0, len(differences), size=(10000, len(differences)) + ) + upper = float(np.percentile(differences[draws].mean(-1), 95)) + threshold = 1e-5 if absolute else 0.05 + return { + "status": ( + "pending" + if len(reference) < 5 + else "passed" if np.isfinite(upper) and upper <= threshold else "failed" + ), + "one_sided_95_upper": upper, + "threshold": threshold, + "difference_mode": "absolute" if absolute else "relative", + "reference": reference, + "candidate": candidate, + } + + +def _loss_curve_gate( + reference: list[float], candidate: list[float], *, steps: int +) -> dict[str, Any]: + """Require every paired optimization loss to match at the stated tolerance.""" + baseline = np.asarray(reference, dtype=np.float64) + actual = np.asarray(candidate, dtype=np.float64) + complete = baseline.shape == actual.shape == (steps,) + finite = bool(np.isfinite(baseline).all() and np.isfinite(actual).all()) + matching = ( + np.isclose(actual, baseline, atol=_NUMERICAL_ATOL, rtol=_NUMERICAL_RTOL) + if complete and finite + else np.zeros(steps, dtype=bool) + ) + return { + "passed": bool(complete and finite and matching.all()), + "expected_steps": steps, + "reference_steps": len(reference), + "candidate_steps": len(candidate), + "finite": finite, + "mismatched_steps": int((~matching).sum()), + "first_mismatched_step": ( + int(np.flatnonzero(~matching)[0]) + 1 if (~matching).any() else None + ), + "max_absolute_difference": ( + float(np.abs(actual - baseline).max()) if complete and finite else None + ), + "atol": _NUMERICAL_ATOL, + "rtol": _NUMERICAL_RTOL, + } + + +def _candidate_sources() -> dict[str, str]: + """Bind newly generated evidence to the executing model and training bytes.""" + from importlib import import_module + + modules = ( + "embodichain_models.dynamics.rigidformer.model", + "embodichain_models.dynamics.rigidformer.data", + "embodichain_models.dynamics.rigidformer.supervised", + "embodichain_models.dynamics.rigidformer.sl_recipe", + "embodichain.learning.sl.learner", + "embodichain.learning.sl.trainer", + "embodichain.learning.sl._artifacts", + "embodichain.learning.sl.types", + ) + sources = { + name: hashlib.sha256( + Path(import_module(name).__file__).read_bytes() + ).hexdigest() + for name in modules + } + sources["reference_runner"] = hashlib.sha256( + Path(__file__).read_bytes() + ).hexdigest() + return sources + + +def _run(args: argparse.Namespace) -> None: + torch.set_num_threads(args.threads) + reference, verified = _load_reference(args.reference) + versions = { + name: importlib.metadata.version(name) + for name in ( + "torch", + "numpy", + "rigidformer", + "roma", + "einx", + "torch-einops-utils", + "x-mlps-pytorch", + "box2d-py", + ) + } + manifest = { + "reference_revision": RIGIDFORMER_REVISION, + "candidate_sources_sha256": _candidate_sources(), + "reference_dataset_sha256": verified["examples/box2d_dataset.py"]["sha256"], + "verified_sources": verified, + "device": args.device, + "threads": args.threads, + "versions": versions, + "steps": args.steps, + "seeds": args.seeds, + "modes": args.modes, + "batch_size": 8, + "dt": 1 / 6, + "objects": 3, + "points_per_object": 48, + "candidate_is_upstream_model": True, + } + _write(args.output / "manifest.json", manifest) + summary: dict[str, Any] = {"passed": False, "modes": {}} + for mode in args.modes: + collision = mode == "collision" + train = reference.Box2DDataset( + [reference.trajectory(seed, collision) for seed in range(64)], collision + ) + validation = reference.Box2DDataset( + [reference.trajectory(seed, collision) for seed in range(1000, 1016)], + collision, + ) + corpus_hash = _digest( + [ + value + for row in (*train.trajectories, *validation.trajectories) + for value in row + ] + ) + paired: list[dict[str, Any]] = [] + for seed in args.seeds: + output = args.output / mode / str(seed) + output.mkdir(parents=True) + np.random.seed(seed) + torch.manual_seed(seed) + parity = _step_parity(train, args.device) + _write(output / "step-parity.json", parity) + baseline = _fit( + train, + validation, + seed=seed, + collision=collision, + steps=args.steps, + device=args.device, + use_sl=False, + output=output, + ) + candidate = _fit( + train, + validation, + seed=seed, + collision=collision, + steps=args.steps, + device=args.device, + use_sl=True, + output=output, + ) + paired.append( + { + "reference": baseline, + "candidate": candidate, + "parity": parity, + "loss_curve": _loss_curve_gate( + baseline["losses"], candidate["losses"], steps=args.steps + ), + } + ) + _write( + args.output / "progress.json", + { + "mode": mode, + "completed_seeds": [row["reference"]["seed"] for row in paired], + }, + ) + gates = { + key: _gate( + [row["reference"][key] for row in paired], + [row["candidate"][key] for row in paired], + ) + for key in ( + "mean_rollout_endpoint_center_error_m", + "mean_rollout_endpoint_point_error_m", + ) + } + summary["modes"][mode] = { + "corpus_sha256": corpus_hash, + "gates": gates, + "step_parity_passed": all(row["parity"]["passed"] for row in paired), + "loss_curves": { + str(row["reference"]["seed"]): row["loss_curve"] for row in paired + }, + } + _write(args.output / "summary.json", summary) + summary["passed"] = args.steps >= 1000 and all( + row["step_parity_passed"] + and all(gate["passed"] for gate in row["loss_curves"].values()) + and all(gate["status"] == "passed" for gate in row["gates"].values()) + for row in summary["modes"].values() + ) + _write(args.output / "summary.json", summary) + if not summary["passed"]: + raise SystemExit("Acceptance incomplete or failed; inspect summary.json") + + +def _main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--reference", required=True, type=Path) + parser.add_argument("--output", required=True, type=Path) + parser.add_argument("--steps", type=int, default=1000) + parser.add_argument("--seeds", type=int, nargs="+", default=list(range(42, 47))) + parser.add_argument( + "--modes", + choices=("friction", "collision"), + nargs="+", + default=["friction", "collision"], + ) + parser.add_argument("--device", default="cpu") + parser.add_argument("--threads", type=int, default=4) + args = parser.parse_args() + if ( + args.steps < 1 + or args.threads < 1 + or len(set(args.seeds)) != len(args.seeds) + or len(set(args.modes)) != len(args.modes) + ): + parser.error("Use positive steps/threads and distinct seeds/modes") + # Admission stays outside the failure guard: never write into an existing run. + args.output.mkdir(parents=True, exist_ok=False) + try: + _run(args) + except BaseException as exc: + _write( + args.output / "failure.json", + {"type": type(exc).__name__, "message": str(exc)}, + ) + raise + + +if __name__ == "__main__": + _main() diff --git a/examples/learning/rigidformer/configs/sl.train.yaml b/examples/learning/rigidformer/configs/sl.train.yaml new file mode 100644 index 0000000000..97ab2626c6 --- /dev/null +++ b/examples/learning/rigidformer/configs/sl.train.yaml @@ -0,0 +1,26 @@ +# Fixed CPU geometry archive; paths resolve from the caller's working directory. +factory: embodichain_models.dynamics.rigidformer.sl_recipe:build_rigidformer_trainer +data: + path: geometry.pt + max_resident_bytes: 536870912 + max_batch_bytes: 67108864 +model: + dim: 128 + dim_head: 32 + heads: 4 + num_register_tokens: 4 + object_self_attn_depth: 3 + anchor_cross_attn_depth: 3 + object_hidden_layers: [0, 1, 3] + num_anchors: 4 + vertex_properties_dim: 3 + max_pairwise_bytes: 268435456 +optimizer: + lr: 0.0003 + weight_decay: 0.01 +trainer: + output_dir: outputs/rigidformer/sl + total_updates: 1000 + batch_size: 8 + seed: 42 + device: cpu diff --git a/examples/learning/rigidformer/requirements-reference.txt b/examples/learning/rigidformer/requirements-reference.txt new file mode 100644 index 0000000000..e195227623 --- /dev/null +++ b/examples/learning/rigidformer/requirements-reference.txt @@ -0,0 +1,4 @@ +# Install the ordinary optional dependencies first, inside an isolated py311 venv. +# Then force this exact upstream source without replacing the inherited stack: +# python -m pip install --force-reinstall --no-deps -r requirements-reference.txt +rigidformer @ git+https://github.com/lucidrains/rigidformer.git@4ab4edac7271b49484748d6bf8086a497f13272d diff --git a/examples/learning/rigidformer/research/README.md b/examples/learning/rigidformer/research/README.md new file mode 100644 index 0000000000..edb6faa00b --- /dev/null +++ b/examples/learning/rigidformer/research/README.md @@ -0,0 +1,105 @@ +# RigidFormer reference acceptance + +This runner compares the shared SL learner with the unmodified optimization +recipe from `lucidrains/rigidformer` revision +`4ab4edac7271b49484748d6bf8086a497f13272d`. The external model and Box2D +generator retain their MIT license; the runner verifies the checkout revision, +the Box2D dataset and training recipe, and every model-package Python file +against the pinned Git blobs. Installed package Python files must match those +same blobs before the optional package is imported. A dirty checkout with the +correct HEAD is rejected. No upstream source or +generated dataset/checkpoint is vendored here. + +Use the `py311` environment with the optional pinned RigidFormer dependency, +`box2d-py` and its build prerequisite SWIG. Install optional dependencies in an +isolated environment inheriting `py311` if the main environment is shared. +The ordinary `embodichain[rigidformer]` extra resolves PyPI `rigidformer==0.2.0`; +its four runtime Python files match the frozen commit. The exact Git-source +requirement remains example-owned in `requirements-reference.txt`, outside +published dependency metadata. From a source checkout, install the ordinary +dependencies first, then replace only the optional upstream package at the +exact commit. The force reinstall prevents a same-version index installation +from standing in for this explicit Git-source installation: + +```bash +conda run -n py311 python -m venv --system-site-packages /tmp/rigidformer-reference-env +/tmp/rigidformer-reference-env/bin/python -m pip install 'embodichain[rigidformer]' +/tmp/rigidformer-reference-env/bin/python -m pip install --force-reinstall --no-deps \ + -r examples/learning/rigidformer/requirements-reference.txt +git clone https://github.com/lucidrains/rigidformer.git /tmp/rigidformer-reference +git -C /tmp/rigidformer-reference checkout 4ab4edac7271b49484748d6bf8086a497f13272d +/tmp/rigidformer-reference-env/bin/python -m embodichain run-example learning/rigidformer \ + --reference /tmp/rigidformer-reference \ + --output /tmp/rigidformer-acceptance --threads 4 +``` + +The full default qualification trains two modes (friction and collision), five +seeds (42–46), and two engines (reference and shared SL), each for 1,000 updates. +It keeps the reference's 64 training trajectories, 16 held-out trajectories, +three objects, 48 material points per object, 1/6-second model interval, B8, +dim128 model and AdamW settings. Both engines consume identical initialization, +sampling and FPS randomness. Periodic evaluation remains at updates 500/1,000 +and preserves the reference RNG consumption. + +The report includes a loss/gradient/update comparison, the reference's 32 +single-step samples, its first held-out trajectory and all 16 held-out +trajectories. Rollouts contain two initial frames followed by 11 predicted +frames. The single-step metric is mean Euclidean point distance, despite the +upstream example labeling it RMSE. Constant-velocity and stationary baselines, +point/center errors and intra-object rigidity errors are retained. The one-sided +95% seed-paired bootstrap upper bound allows at most 5% relative degradation; +near-zero references use an absolute tolerance of 1e-5. Numerical comparisons +use atol1e-5/rtol1e-4. Every paired training loss must also satisfy +`abs(sl_loss - reference_loss) <= 1e-5 + 1e-4 * abs(reference_loss)` at every +update; both complete curves must contain the requested number of finite +values. A matching final loss alone cannot qualify a run. Fewer than five +seeds or fewer than 1,000 updates report pending acceptance and exit nonzero. + +Outputs include a frozen manifest, source/data hashes, dependency versions, +per-seed immutable results/checkpoints, curves and `summary.json`. Once a fresh +output directory is admitted, setup, dependency, revision, training and failed +or pending qualification failures are retained in `failure.json`. An existing +output directory is rejected without modifying its artifacts. A completed +reference comparison establishes +adapter/trainer agreement, not paper accuracy, action gradients or physics +solver qualification. The reference's collision limitations remain visible. +`training_seconds` includes the two periodic reference-trajectory evaluations; +the 32-window and full-panel final evaluations are outside that interval. + +For a user-supplied fixed geometry corpus, the generic training entry point uses +`factory: embodichain_models.dynamics.rigidformer.sl_recipe:build_rigidformer_trainer`. Its +tensor-only archive contains `positions[E,T,O,P,3]`, `properties[E,O,A]` (or +`[E,O,P,A]`) and `step_dt`; optional valid lengths retain the source field +names. This regular episode-uniform sampler is separate from the reference's +event-biased Box2D sampling. The source clones under an explicit resident byte +budget, returns owned bounded batches and never crosses episode boundaries. +Reference geometry is supplied explicitly and point identities remain aligned +across time. + +Importing the adapter does not load the optional model or simulation. The +current upstream nearest-neighbour feature uses a dense pairwise distance +tensor, so the adapter checks its estimated byte size before forwarding. This +bound is for that tensor alone, not total process/accelerator memory. + +Finite padding is supported with unequal object counts and unequal point +counts when the padded point axis is at most 128. Each supplied point count +must be at least `num_anchors` (four by default), including placeholder objects +outside `object_lens`; fixed anchors must index valid points. Nonfinite padding +is unsupported. The pinned PointNet hierarchy uses unmasked gathered KNN +neighbors when the padded point axis exceeds 128, so unequal point counts in +that layout are explicitly rejected by both the objective path and rollout. +Full point lengths above 128 remain available with unchanged upstream math. +The 48-point Box2D reference does not enter the restricted layout. Focused +optional-model tests change finite padded geometry, attributes and targets and +check objective loss/gradients and valid rollout points at 48 and 128 points. + +## Source-bound evidence + +The relocation preserves the pinned upstream model, recipe math, parameter keys, +geometry schema, fixed seeds, budgets and acceptance gates above. New runs bind +`candidate_sources_sha256` in `manifest.json` to the executing model, adapters, +shared trainer and private runner. Existing reports are immutable historical +evidence for their recorded source, data and runtime; they are not relabeled as +results for a newer commit. A short paired run verifies relocation parity but +continues to report pending full acceptance under the unchanged 5-seed/1,000-step +threshold. diff --git a/extensions/embodichain_models/README.md b/extensions/embodichain_models/README.md index 7509d4f05d..24054af577 100644 --- a/extensions/embodichain_models/README.md +++ b/extensions/embodichain_models/README.md @@ -1,7 +1,8 @@ # Official models `embodichain_models` contains official model implementations grouped by model -family, including `embodichain_models.dynamics.nerd`. It ships in the same +family, including `embodichain_models.dynamics.nerd` and +`embodichain_models.dynamics.rigidformer`. It ships in the same `embodichain` wheel and uses that distribution's version and dependencies. There is no separate model distribution, version file or model registry. @@ -35,7 +36,7 @@ Use the following canonical packages for the four current model directions: | Ordinary state/visual BC | `embodichain_models.policy.bc` | Target for the BC follow-up; not created yet. | | NeRD | `embodichain_models.dynamics.nerd` | Implemented in this change. | | Robot world model (RWM) | `embodichain_models.dynamics.rwm` | Target for the RWM follow-up; not created yet. | -| RigidFormer | `embodichain_models.dynamics.rigidformer` | Target for the RigidFormer follow-up; not created yet. | +| RigidFormer | `embodichain_models.dynamics.rigidformer` | Implemented; optional pinned upstream dependency. | RWM predicts action-conditioned robot state transitions, so it belongs to `dynamics`. World models also describe future-state/observation prediction; @@ -63,3 +64,10 @@ resources remain in the corresponding example. Shared MLPs and normalizers remain in core. All four directions use the same SL contracts regardless of their domain. Create a family only when its implementation lands, without empty placeholder directories, eager cross-family imports or a new registry. + +RigidFormer owns aligned point-cloud data windows, its explicit reference-geometry +adapter, valid-object weighted objective and `sl_recipe:build_rigidformer_trainer`. +The optional upstream import occurs at model construction. Its pinned Box2D +reference experiment and configuration presets belong to +`examples/learning/rigidformer/`; their frozen protocol and historical results +retain their original source/data identities. diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/THIRD_PARTY_LICENSES b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/THIRD_PARTY_LICENSES new file mode 100644 index 0000000000..67485bf8fe --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/THIRD_PARTY_LICENSES @@ -0,0 +1,32 @@ +RigidFormer +=========== + +Optional upstream model and Box2D reference recipe: +https://github.com/lucidrains/rigidformer +Revision: 4ab4edac7271b49484748d6bf8086a497f13272d + +The upstream source remains in its separately installed dependency and explicit +reference checkout. EmbodiChain does not vendor those implementation files. +Their complete license follows. + +MIT License + +Copyright (c) 2026 Phil Wang + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/__init__.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/__init__.py new file mode 100644 index 0000000000..36635ffd16 --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/__init__.py @@ -0,0 +1,50 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Lazy RigidFormer geometry prediction and common supervised integration.""" + +from __future__ import annotations + +from importlib import import_module +from typing import Any + +__all__ = [ + "RIGIDFORMER_REVISION", + "RigidFormerCfg", + "RigidFormerModel", + "RigidBodySequenceSource", + "RigidFormerObjective", + "rigidformer_batch", + "build_rigidformer_trainer", +] + +_EXPORT_MODULES = { + "RIGIDFORMER_REVISION": "model", + "RigidFormerCfg": "model", + "RigidFormerModel": "model", + "RigidBodySequenceSource": "data", + "RigidFormerObjective": "supervised", + "rigidformer_batch": "supervised", + "build_rigidformer_trainer": "sl_recipe", +} + + +def __getattr__(name: str) -> Any: + if name not in _EXPORT_MODULES: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + value = getattr(import_module(f".{_EXPORT_MODULES[name]}", __name__), name) + globals()[name] = value + return value diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/data.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/data.py new file mode 100644 index 0000000000..ee3daea2e6 --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/data.py @@ -0,0 +1,209 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Owned causal point-cloud windows from a fixed resident geometry corpus.""" + +from __future__ import annotations + +import math + +import torch + +__all__ = ["RigidBodySequenceSource"] + + +class RigidBodySequenceSource: + """Sample previous/current/next geometry without crossing episode bounds. + + Args: + positions: Float CPU points with axes ``[episode, time, object, point, 3]``. + properties: Float CPU attributes ``[episode, object, attributes]`` or + ``[episode, object, point, attributes]``. + step_dt: Seconds between stored frames. + source_id: Immutable corpus identity recorded with downstream batches. + stride: Number of stored frames between model input steps. + episode_lengths: Optional number of real frames per episode. + object_lens: Optional number of real objects per episode. + object_point_lens: Optional number of real points per episode/object. + max_resident_bytes: Bound on cloned resident tensors. + max_batch_bytes: Bound on returned geometry payloads. + + Points retain their material identities across time. Every returned tensor + owns its data independently of the source and later sampling. Sampler RNG + is caller-owned and can be restored with a learner checkpoint. + """ + + def __init__( + self, + positions: torch.Tensor, + properties: torch.Tensor, + *, + step_dt: float, + source_id: str, + stride: int = 1, + episode_lengths: torch.Tensor | None = None, + object_lens: torch.Tensor | None = None, + object_point_lens: torch.Tensor | None = None, + max_resident_bytes: int = 512 * 1024 * 1024, + max_batch_bytes: int = 64 * 1024 * 1024, + ) -> None: + if positions.ndim != 5 or positions.shape[-1] != 3: + raise ValueError("positions must be [episode, time, object, point, 3]") + if min(positions.shape[:-1]) < 1 or not positions.is_floating_point(): + raise ValueError("positions must be nonempty floating-point tensors") + if ( + properties.ndim not in (3, 4) + or properties.shape[:2] != (positions.shape[0], positions.shape[2]) + or (properties.ndim == 4 and properties.shape[2] != positions.shape[3]) + or not properties.is_floating_point() + ): + raise ValueError("properties must match the episode/object/point layout") + if type(stride) is not int or stride < 1: + raise ValueError("stride must be a positive integer") + if not math.isfinite(step_dt) or step_dt <= 0 or not source_id: + raise ValueError("Positive finite step_dt and source_id are required") + tensors = [positions, properties] + tensors += [ + v + for v in (episode_lengths, object_lens, object_point_lens) + if v is not None + ] + if any(v.device.type != "cpu" for v in tensors): + raise ValueError("Resident geometry must be on CPU") + lengths = ( + torch.full((positions.shape[0],), positions.shape[1], dtype=torch.long) + if episode_lengths is None + else episode_lengths + ) + if ( + lengths.shape != positions.shape[:1] + or lengths.dtype != torch.long + or (lengths > positions.shape[1]).any() + or (lengths < 2 * stride + 1).any() + ): + raise ValueError("Every episode needs enough valid frames for one window") + if object_lens is not None and ( + object_lens.shape != positions.shape[:1] + or object_lens.dtype != torch.long + or (object_lens < 1).any() + or (object_lens > positions.shape[2]).any() + ): + raise ValueError("Invalid object_lens") + if object_point_lens is not None and ( + object_point_lens.shape != (positions.shape[0], positions.shape[2]) + or object_point_lens.dtype != torch.long + or (object_point_lens < 1).any() + or (object_point_lens > positions.shape[3]).any() + ): + raise ValueError("Invalid object_point_lens") + resident = [positions, properties, lengths] + resident += [v for v in (object_lens, object_point_lens) if v is not None] + required = sum(v.numel() * v.element_size() for v in resident) + if required > max_resident_bytes or max_batch_bytes < 1: + raise ValueError("Geometry exceeds the configured memory budget") + self.positions, self.properties = positions.clone(), properties.clone() + self.episode_lengths = lengths.clone() + self.object_lens = None if object_lens is None else object_lens.clone() + self.object_point_lens = ( + None if object_point_lens is None else object_point_lens.clone() + ) + self.step_dt, self.stride, self.source_id = float(step_dt), stride, source_id + self.max_batch_bytes = max_batch_bytes + + def get_batch( + self, episode_indices: torch.Tensor, anchors: torch.Tensor + ) -> dict[str, torch.Tensor]: + """Return owned windows at explicit current-frame anchors. + + Args: + episode_indices: CPU int64 episode indices, ``[B]``. + anchors: CPU int64 current-frame indices, ``[B]``. + + Returns: + Upstream-compatible previous/current/reference/target geometry. + """ + if ( + episode_indices.ndim != 1 + or episode_indices.numel() < 1 + or anchors.shape != episode_indices.shape + or any( + v.dtype != torch.long or v.device.type != "cpu" + for v in (episode_indices, anchors) + ) + ): + raise ValueError( + "Episode and anchor indices must be nonempty CPU int64 vectors" + ) + if (episode_indices < 0).any() or ( + episode_indices >= self.positions.shape[0] + ).any(): + raise ValueError("Episode index out of range") + lengths = self.episode_lengths[episode_indices] + if (anchors < self.stride).any() or (anchors + self.stride >= lengths).any(): + raise ValueError("Geometry windows cannot cross an episode boundary") + batch = episode_indices.numel() + point_bytes = self.positions[0, 0].numel() * self.positions.element_size() + property_bytes = self.properties[0].numel() * self.properties.element_size() + row_bytes = ( + 4 * point_bytes + property_bytes + self.positions.element_size() + 16 + ) + row_bytes += 8 if self.object_lens is not None else 0 + row_bytes += ( + 8 * self.positions.shape[2] if self.object_point_lens is not None else 0 + ) + if batch * row_bytes > self.max_batch_bytes: + raise ValueError("Returned geometry exceeds max_batch_bytes") + previous = self.positions[episode_indices, anchors - self.stride] + result = { + "object_pos_prev": previous, + "object_pos": self.positions[episode_indices, anchors], + "object_pos_next": self.positions[episode_indices, anchors + self.stride], + "object_first_frame_pos": previous.clone(), + "vertex_properties": self.properties[episode_indices], + "delta_times": torch.full( + (batch,), self.step_dt * self.stride, dtype=self.positions.dtype + ), + "episode_index": episode_indices.clone(), + "frame_index": anchors.clone(), + } + for key in ("object_lens", "object_point_lens"): + value = getattr(self, key) + if value is not None: + result[key] = value[episode_indices] + return result + + def sample_batch( + self, batch_size: int, *, generator: torch.Generator | None = None + ) -> dict[str, torch.Tensor]: + """Sample episode-uniform windows with a caller-owned CPU RNG. + + Args: + batch_size: Positive number of windows. + generator: Optional CPU sampling generator. + + Returns: + Owned geometry windows in the same layout as :meth:`get_batch`. + """ + if type(batch_size) is not int or batch_size < 1: + raise ValueError("batch_size must be positive") + episodes = torch.randint( + self.positions.shape[0], (batch_size,), generator=generator + ) + count = self.episode_lengths[episodes] - 2 * self.stride + anchors = ( + self.stride + (torch.rand(batch_size, generator=generator) * count).long() + ) + return self.get_batch(episodes, anchors) diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/model.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/model.py new file mode 100644 index 0000000000..8150294369 --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/model.py @@ -0,0 +1,269 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Optional RigidFormer adapter with an explicit reference-geometry contract.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from typing import Any + +import torch +from torch import nn + +from embodichain.utils import configclass + +__all__ = ["RIGIDFORMER_REVISION", "RigidFormerCfg", "RigidFormerModel"] + +RIGIDFORMER_REVISION = "4ab4edac7271b49484748d6bf8086a497f13272d" + + +@configclass +class RigidFormerCfg: + """Configure the optional upstream network and its distance-buffer budget. + + The defaults match the pinned upstream Box2D experiment. The byte budget + bounds the dense nearest-neighbour distance tensor, not total GPU memory. + Each distance counts at least four bytes because autocast may promote + reduced-precision geometry to an FP32 ``cdist`` result. + + Args: + dim: Object/anchor feature width. + dim_head: Per-head attention width. + heads: Attention heads per layer. + num_register_tokens: Additional object-attention register tokens. + object_self_attn_depth: Object interaction layers. + anchor_cross_attn_depth: Anchor-to-object interaction layers. + object_hidden_layers: Object-layer outputs consumed by anchor layers. + num_anchors: Material points used to advance each rigid object. + vertex_properties_dim: Number of object/point attribute channels. + max_pairwise_bytes: Maximum dense nearest-neighbour distance bytes. + """ + + dim: int = 128 + dim_head: int = 32 + heads: int = 4 + num_register_tokens: int = 4 + object_self_attn_depth: int = 3 + anchor_cross_attn_depth: int = 3 + object_hidden_layers: tuple[int, ...] = (0, 1, 3) + num_anchors: int = 4 + vertex_properties_dim: int = 3 + max_pairwise_bytes: int = 256 * 1024 * 1024 + + +class RigidFormerModel(nn.Module): + """Wrap the pinned RigidFormer architecture without importing simulation. + + Args: + cfg: Upstream architecture and per-call pairwise-distance byte budget. + + Reference geometry and point identities must remain aligned with the + current/previous point clouds. The adapter exposes forward prediction and + parameter training; action gradients and a physics-solver interface are + outside its contract. Finite padding with unequal point lengths is supported + only when the padded point axis is at most 128; the pinned upstream + hierarchical encoder does not preserve masking beyond that boundary. + """ + + action_conditioned: bool = False + qualified_state_gradients: bool = False + + def __init__(self, cfg: RigidFormerCfg | None = None) -> None: + super().__init__() + self.cfg = cfg if cfg is not None else RigidFormerCfg() + self.cfg.validate() + if self.cfg.max_pairwise_bytes < 1: + raise ValueError("max_pairwise_bytes must be positive") + try: + from rigidformer import Rigidformer + except ImportError as exc: + raise ImportError( + "RigidFormer requires the optional rigidformer dependency; " + "install embodichain[rigidformer] (rigidformer==0.2.0). " + f"For frozen reference revision {RIGIDFORMER_REVISION}, use " + "examples/learning/rigidformer/requirements-reference.txt" + ) from exc + fields = self.cfg.to_dict() + fields.pop("max_pairwise_bytes") + self.network = Rigidformer(**fields) + + def _check(self, values: Mapping[str, Any]) -> None: + current = values["object_pos"] + if current.ndim != 4 or current.shape[-1] != 3: + raise ValueError("object_pos must have shape [batch, object, point, 3]") + if min(current.shape[:3]) < 1: + raise ValueError("Point batches must contain objects and points") + for key in ("object_pos_prev", "object_first_frame_pos"): + if values[key].shape != current.shape: + raise ValueError(f"{key} must match the current point layout") + target = values.get("object_pos_next") + if target is not None and target.shape != current.shape: + raise ValueError("object_pos_next must match the current point layout") + delta_times = values["delta_times"] + if delta_times.shape != current.shape[:1]: + raise ValueError("delta_times must have shape [batch]") + if not torch.isfinite(delta_times).all() or (delta_times <= 0).any(): + raise ValueError("delta_times must be finite and positive") + batch, objects, points = current.shape[:3] + object_lens = values.get("object_lens") + point_lens = values.get("object_point_lens") + for name, lengths, shape, maximum, minimum in ( + ("object_lens", object_lens, (batch,), objects, 1), + ( + "object_point_lens", + point_lens, + (batch, objects), + points, + self.cfg.num_anchors, + ), + ): + if lengths is not None and ( + lengths.shape != shape + or lengths.dtype != torch.long + or lengths.device != current.device + or (lengths < minimum).any() + or (lengths > maximum).any() + ): + raise ValueError( + f"{name} must be int64 {shape} counts in [{minimum}, {maximum}] " + "on the geometry device, including padded objects" + ) + if point_lens is not None and points > 128 and (point_lens != points).any(): + raise ValueError( + "Unequal object_point_lens requires a padded point axis <= 128; " + "the pinned upstream PointNet FPS/KNN leaks finite padding above 128" + ) + anchors = values.get("anchor_indices") + if anchors is not None: + valid_points = points if point_lens is None else point_lens[..., None] + if ( + anchors.ndim != 3 + or anchors.shape[:2] != (batch, objects) + or anchors.shape[-1] < 1 + or anchors.dtype != torch.long + or anchors.device != current.device + or (anchors < 0).any() + or (anchors >= valid_points).any() + ): + raise ValueError( + "anchor_indices must index valid points in every object" + ) + # cdist may run in FP32 under autocast, including for FP16/BF16 inputs. + distance_bytes = ( + batch * (objects * points) ** 2 * max(current.element_size(), 4) + ) + if distance_bytes > self.cfg.max_pairwise_bytes: + raise ValueError( + f"Pairwise distances require {distance_bytes} bytes, " + f"exceeding {self.cfg.max_pairwise_bytes}" + ) + + def forward( + self, + *, + object_pos: torch.Tensor, + object_pos_prev: torch.Tensor, + object_first_frame_pos: torch.Tensor, + vertex_properties: torch.Tensor, + delta_times: torch.Tensor, + object_pos_next: torch.Tensor | None = None, + anchor_indices: torch.Tensor | None = None, + object_point_lens: torch.Tensor | None = None, + object_lens: torch.Tensor | None = None, + ) -> Any: + """Predict next points, or return upstream loss with an explicit target. + + Args: + object_pos: Current points with axes ``[B, O, P, 3]``. + object_pos_prev: Previous points with identical point identities. + object_first_frame_pos: Reference points for rigid projection. + vertex_properties: Per-object or per-point attributes. + delta_times: Positive integration intervals in seconds, ``[B]``. + object_pos_next: Optional training target, never a model input. + anchor_indices: Optional fixed per-object anchor indices. + object_point_lens: Optional valid point counts, ``[B, O]``, each at + least ``cfg.num_anchors`` including padded objects. Unequal + counts require a padded point axis of at most 128. + object_lens: Optional valid object counts, ``[B]``. + + Returns: + Upstream prediction namedtuple, or scalar loss and loss breakdown. + """ + values = dict( + object_pos=object_pos, + object_pos_prev=object_pos_prev, + object_first_frame_pos=object_first_frame_pos, + vertex_properties=vertex_properties, + delta_times=delta_times, + object_pos_next=object_pos_next, + anchor_indices=anchor_indices, + object_point_lens=object_point_lens, + object_lens=object_lens, + ) + self._check(values) + return self.network(**values) + + def rollout( + self, + object_positions: Sequence[torch.Tensor], + *, + vertex_properties: torch.Tensor, + delta_times: torch.Tensor, + num_steps: int, + anchor_indices: torch.Tensor | None = None, + object_point_lens: torch.Tensor | None = None, + object_lens: torch.Tensor | None = None, + ) -> list[torch.Tensor]: + """Predict a fixed-step trajectory while preserving the first geometry. + + Args: + object_positions: At least two aligned initial frames. + vertex_properties: Static object/point attributes. + delta_times: Fixed positive time interval per batch row. + num_steps: Number of new frames, excluding the supplied history. + anchor_indices: Optional initial anchors reused throughout rollout. + object_point_lens: Optional valid point counts with the same masking + limits as :meth:`forward`. + object_lens: Optional valid object counts. + + Returns: + Initial frames followed by the predicted frames. + """ + if len(object_positions) < 2 or type(num_steps) is not int or num_steps < 1: + raise ValueError("Rollout requires two frames and positive num_steps") + self._check( + dict( + object_pos=object_positions[-1], + object_pos_prev=object_positions[-2], + object_first_frame_pos=object_positions[0], + delta_times=delta_times, + anchor_indices=anchor_indices, + object_point_lens=object_point_lens, + object_lens=object_lens, + ) + ) + from rigidformer import RigidformerRolloutWrapper + + return RigidformerRolloutWrapper(self.network)( + object_positions=list(object_positions), + vertex_properties=vertex_properties, + delta_times=delta_times, + num_steps=num_steps, + anchor_indices=anchor_indices, + object_point_lens=object_point_lens, + object_lens=object_lens, + ) diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/sl_recipe.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/sl_recipe.py new file mode 100644 index 0000000000..2bb34b287d --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/sl_recipe.py @@ -0,0 +1,109 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Standalone RigidFormer geometry recipe for the shared SL trainer.""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any + +import torch + +from embodichain.learning.sl import SupervisedBatch + +from .supervised import RigidFormerObjective, rigidformer_batch + +if TYPE_CHECKING: + from embodichain.learning.sl.trainer import SupervisedTrainer + +__all__ = ["build_rigidformer_trainer"] + + +class _RigidFormerSource: + def __init__(self, source: Any) -> None: + self.source = source + self.source_id = source.source_id + + def sample(self, batch_size: int, *, generator: torch.Generator) -> SupervisedBatch: + raw = self.source.sample_batch(batch_size, generator=generator) + return rigidformer_batch(raw, metadata={"source_id": self.source_id}) + + +def build_rigidformer_trainer(config: Mapping[str, Any]) -> SupervisedTrainer: + """Build an offline geometry recipe for the unified ``train-sl`` command. + + Args: + config: Recipe mapping with ``data.path`` pointing to a tensor-only + PyTorch archive containing ``positions``, ``properties``, and + ``step_dt``. Optional geometry lengths use their source field names. + The ``model``, ``optimizer``, and ``trainer`` mappings configure the + upstream network, AdamW, and shared trainer, respectively. + + Returns: + Shared trainer consuming bounded causal geometry windows. + """ + from pathlib import Path + import hashlib + + from .data import RigidBodySequenceSource + from .model import RigidFormerCfg, RigidFormerModel + from embodichain.learning.sl import SupervisedLearner + from embodichain.learning.sl.trainer import SupervisedTrainer, SupervisedTrainerCfg + + data_cfg = dict(config["data"]) + path = Path(data_cfg.pop("path")) + source_id = data_cfg.pop("source_id", None) + digest = hashlib.sha256() + with path.open("rb") as handle: + for block in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(block) + fingerprint = digest.hexdigest() + payload = torch.load(path, map_location="cpu", weights_only=True, mmap=True) + source = RigidBodySequenceSource( + payload["positions"], + payload["properties"], + step_dt=payload["step_dt"], + source_id=source_id or fingerprint, + episode_lengths=payload.get("episode_lengths"), + object_lens=payload.get("object_lens"), + object_point_lens=payload.get("object_point_lens"), + **data_cfg, + ) + trainer_cfg = SupervisedTrainerCfg(**config.get("trainer", {})) + torch.manual_seed(trainer_cfg.seed) + model = RigidFormerModel(RigidFormerCfg(**config.get("model", {}))).to( + trainer_cfg.device + ) + optimizer_cfg = {"lr": 3e-4, "weight_decay": 1e-2, **config.get("optimizer", {})} + learner = SupervisedLearner( + model, + RigidFormerObjective(), + torch.optim.AdamW(model.parameters(), **optimizer_cfg), + max_grad_norm=1.0, + ) + return SupervisedTrainer( + learner, + _RigidFormerSource(source), + trainer_cfg, + signature={ + "geometry_sha256": fingerprint, + "model": model.cfg.to_dict(), + "optimizer": optimizer_cfg, + "data": data_cfg, + "objective": "rigidformer-reference-v1", + }, + ) diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/supervised.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/supervised.py new file mode 100644 index 0000000000..8634625a2d --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/supervised.py @@ -0,0 +1,98 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""RigidFormer objective and conversion of explicit geometry supervision.""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +import torch +from torch import nn + +from embodichain.learning.sl import LossResult, SupervisedBatch + +__all__ = ["RigidFormerObjective", "rigidformer_batch"] + +_INPUTS = ( + "object_pos", + "object_pos_prev", + "object_first_frame_pos", + "vertex_properties", + "delta_times", +) +_OPTIONAL = ("anchor_indices", "object_point_lens", "object_lens") + + +def rigidformer_batch( + values: Mapping[str, torch.Tensor], *, metadata: Mapping[str, Any] | None = None +) -> SupervisedBatch: + """Separate next-frame targets from causal RigidFormer inputs. + + Args: + values: Owned geometry tensors in the upstream keyword layout. + metadata: Optional source identity and frame provenance. + + Returns: + A supervised batch with target geometry outside the input mapping. + """ + missing = set((*_INPUTS, "object_pos_next")) - values.keys() + if missing: + raise ValueError(f"Missing geometry fields: {sorted(missing)}") + inputs = {key: values[key] for key in (*_INPUTS, *_OPTIONAL) if key in values} + context = dict(metadata or {}) + context.update( + {key: values[key] for key in ("episode_index", "frame_index") if key in values} + ) + return SupervisedBatch( + inputs=inputs, + targets=values["object_pos_next"], + metadata=context, + ) + + +class RigidFormerObjective: + """Preserve upstream anchor position/acceleration losses and their weights.""" + + def compute_loss(self, model: nn.Module, batch: SupervisedBatch) -> LossResult: + """Compute loss without placing the next-frame target in model inputs. + + Args: + model: RigidFormer-compatible model returning loss and breakdown. + batch: Explicit geometry inputs and next-frame targets. + + Returns: + Scalar loss and upstream position/acceleration diagnostic terms, + weighted by the valid object count used by the upstream mean. + """ + if batch.masks: + raise ValueError( + "RigidFormer masks use object_lens/object_point_lens inputs; " + "arbitrary temporal loss masks are unsupported" + ) + loss, terms = model(**batch.inputs, object_pos_next=batch.targets) + object_lens = batch.inputs.get("object_lens") + weight = ( + float(object_lens.sum()) + if object_lens is not None + else float(batch.targets.shape[0] * batch.targets.shape[1]) + ) + return LossResult( + loss=loss, + metrics={"acceleration": terms.acceleration, "position": terms.position}, + weight=weight, + ) diff --git a/pyproject.toml b/pyproject.toml index aad3274ea5..dbfe681233 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -56,6 +56,13 @@ dependencies = [ ] [project.optional-dependencies] +# Model construction imports the separately installed upstream package lazily. +# Box2D/fire support the example-owned frozen reference workflow. +rigidformer = [ + "rigidformer==0.2.0", + "box2d-py", + "fire", +] policy-deploy = [ # ORT 1.27+ defaults to CUDA 13; EmbodiChain's supported stack uses CUDA 12. "onnxruntime-gpu>=1.20,<1.27", diff --git a/setup.py b/setup.py index 6ea584a60b..8863d8bcbe 100644 --- a/setup.py +++ b/setup.py @@ -185,6 +185,7 @@ def main(): package_data={ "embodichain": ["VERSION"], "embodichain_models.dynamics.nerd": ["THIRD_PARTY_LICENSES"], + "embodichain_models.dynamics.rigidformer": ["THIRD_PARTY_LICENSES"], "embodichain.data": ["default_materials.json"], "embodichain_benchmarks.motion_generation": ["suites/*.yaml"], "embodichain.gen_sim.simready_pipeline.configs": ["*.json"], diff --git a/tests/benchmark/motion_generation/test_embodiment_runtime.py b/tests/benchmark/motion_generation/test_embodiment_runtime.py index 5549b17e8b..120c81b502 100644 --- a/tests/benchmark/motion_generation/test_embodiment_runtime.py +++ b/tests/benchmark/motion_generation/test_embodiment_runtime.py @@ -253,7 +253,7 @@ def test_gripper_and_domain_overrides_cannot_change_component_commands( def test_official_ur5_component_uses_robot_catalog(monkeypatch, tmp_path: Path) -> None: from embodichain.lab.sim import cfg as cfg_module - from embodichain_robots import ur_robot + from embodichain_robots.universal_robots import ur_robot def local_asset(path: str) -> str: target = tmp_path / path diff --git a/tests/cli/test_list_robot.py b/tests/cli/test_list_robot.py index 648e73b6aa..ff65ecd0a0 100644 --- a/tests/cli/test_list_robot.py +++ b/tests/cli/test_list_robot.py @@ -248,8 +248,10 @@ def test_list_robot_wraps_long_catalog_metadata( def test_list_robot_prints_control_metadata_without_loading_robot( catalogs: list[Any], capsys: pytest.CaptureFixture[str], + monkeypatch: pytest.MonkeyPatch, ) -> None: """A provider's control summary is printed directly from its catalog.""" + monkeypatch.setenv("COLUMNS", "120") spec = { "class_path": "absent_robots:RobotCfg", "dof": 7, diff --git a/tests/examples/rigidformer_reference/__init__.py b/tests/examples/rigidformer_reference/__init__.py new file mode 100644 index 0000000000..b4c74d22b8 --- /dev/null +++ b/tests/examples/rigidformer_reference/__init__.py @@ -0,0 +1,22 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + + +"""RigidFormer reference workflow checks.""" + +from __future__ import annotations + +__all__: list[str] = [] diff --git a/tests/examples/rigidformer_reference/test_acceptance.py b/tests/examples/rigidformer_reference/test_acceptance.py new file mode 100644 index 0000000000..5bbde0bdc3 --- /dev/null +++ b/tests/examples/rigidformer_reference/test_acceptance.py @@ -0,0 +1,171 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Exercise immutable reference provenance and paired acceptance failure gates.""" + +from __future__ import annotations + +import importlib.metadata +import importlib.util +import json +from pathlib import Path +import subprocess +import sys + +import pytest + + +@pytest.fixture +def acceptance(): + source = ( + Path(__file__).resolve().parents[3] + / "examples/learning/rigidformer/_support/_acceptance.py" + ) + spec = importlib.util.spec_from_file_location("rigidformer_acceptance_test", source) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +@pytest.fixture +def reference_sources(tmp_path, acceptance, monkeypatch): + checkout = tmp_path / "reference" + installed = tmp_path / "installed" + files = { + "examples/box2d_dataset.py": "# pinned dataset\n", + "examples/train_box2d.py": "# pinned training recipe\n", + "rigidformer/__init__.py": "# pinned package exports\n", + "rigidformer/rigidformer.py": "# pinned model\n", + "rigidformer/rotary_3d.py": "# pinned model helper\n", + } + for name, source in files.items(): + path = checkout / name + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(source) + if name.startswith("rigidformer/"): + destination = installed / Path(name).relative_to("rigidformer") + destination.parent.mkdir(parents=True, exist_ok=True) + destination.write_text(source) + subprocess.run(["git", "init", "-q", str(checkout)], check=True) + subprocess.run(["git", "-C", str(checkout), "add", "."], check=True) + subprocess.run( + [ + "git", + "-C", + str(checkout), + "-c", + "user.name=Test", + "-c", + "user.email=test@example.invalid", + "commit", + "-qm", + "Pinned fixture", + ], + check=True, + ) + revision = subprocess.check_output( + ["git", "-C", str(checkout), "rev-parse", "HEAD"], text=True + ).strip() + monkeypatch.setattr(acceptance, "RIGIDFORMER_REVISION", revision) + assert set(acceptance._verify_sources(checkout, installed)) == set(files) + return checkout, installed + + +@pytest.mark.parametrize( + "name", + [ + "examples/box2d_dataset.py", + "examples/train_box2d.py", + "rigidformer/rigidformer.py", + "rigidformer/rotary_3d.py", + "rigidformer/__init__.py", + ], +) +def test_reference_rejects_dirty_sources_even_when_head_and_installed_bytes_match( + acceptance, reference_sources, name +): + checkout, installed = reference_sources + (checkout / name).write_text("# changed after commit\n") + if name.startswith("rigidformer/"): + (installed / Path(name).relative_to("rigidformer")).write_text( + "# changed after commit\n" + ) + with pytest.raises( + ValueError, match="Reference source differs from pinned Git blob" + ): + acceptance._verify_sources(checkout, installed) + + +def test_reference_rejects_installed_model_divergence(acceptance, reference_sources): + checkout, installed = reference_sources + (installed / "rigidformer.py").write_text("# installed model changed\n") + with pytest.raises( + ValueError, match="Installed source differs from pinned Git blob" + ): + acceptance._verify_sources(checkout, installed) + + +def test_loss_curve_gate_checks_every_update_and_complete_finite_history(acceptance): + assert acceptance._loss_curve_gate([1.0, 2.0, 3.0], [1.0, 2.0001, 3.0], steps=3)[ + "passed" + ] + changed = acceptance._loss_curve_gate([1.0, 2.0, 3.0], [1.0, 2.1, 3.0], steps=3) + assert not changed["passed"] + assert changed["first_mismatched_step"] == 2 + assert changed["mismatched_steps"] == 1 + assert not acceptance._loss_curve_gate([1.0, 2.0], [1.0, 2.0], steps=3)["passed"] + assert not acceptance._loss_curve_gate( + [1.0, float("nan")], [1.0, float("nan")], steps=2 + )["passed"] + + +@pytest.mark.parametrize( + "failure", + [ + RuntimeError("device setup failed"), + importlib.metadata.PackageNotFoundError("rigidformer"), + ValueError("Expected reference revision"), + SystemExit("Acceptance incomplete or failed; inspect summary.json"), + ], +) +def test_acceptance_records_setup_and_qualification_failures_after_fresh_admission( + acceptance, monkeypatch, tmp_path, failure +): + output = tmp_path / "fresh" + monkeypatch.setattr( + sys, + "argv", + [ + "acceptance.py", + "--reference", + str(tmp_path / "reference"), + "--output", + str(output), + ], + ) + + def fail(args): + raise failure + + monkeypatch.setattr(acceptance, "_run", fail) + with pytest.raises(type(failure)): + acceptance._main() + recorded = json.loads((output / "failure.json").read_text()) + assert recorded == {"type": type(failure).__name__, "message": str(failure)} + (output / "failure.json").write_text("existing result") + with pytest.raises(FileExistsError): + acceptance._main() + assert (output / "failure.json").read_text() == "existing result" diff --git a/tests/learning/models/test_rigidformer.py b/tests/learning/models/test_rigidformer.py new file mode 100644 index 0000000000..e14db50f1d --- /dev/null +++ b/tests/learning/models/test_rigidformer.py @@ -0,0 +1,402 @@ +# ---------------------------------------------------------------------------- +# 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 the SL adapter against the independent optional upstream model.""" + +from __future__ import annotations + +from copy import deepcopy +import subprocess +import sys + +import pytest +import torch + +from embodichain_models.dynamics.rigidformer.model import ( + RigidFormerCfg, + RigidFormerModel, +) +from embodichain_models.dynamics.rigidformer.sl_recipe import build_rigidformer_trainer +from embodichain.learning.sl import SupervisedLearner +from embodichain_models.dynamics.rigidformer.supervised import ( + RigidFormerObjective, + rigidformer_batch, +) + + +@pytest.fixture +def geometry(): + generator = torch.Generator().manual_seed(19) + previous = torch.randn(2, 2, 32, 3, generator=generator) + return dict( + object_pos=previous + 0.01, + object_pos_prev=previous, + object_first_frame_pos=previous.clone(), + object_pos_next=previous + 0.025, + vertex_properties=torch.ones(2, 2, 3), + delta_times=torch.full((2,), 0.1), + anchor_indices=torch.tensor([0, 8, 16, 24]).expand(2, 2, 4), + ) + + +def test_rigidformer_import_is_lazy_without_optional_model_or_simulation(): + code = """ +import importlib.abc +import sys +class BlockOptional(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if fullname.split('.')[0] in {'rigidformer', 'dexsim', 'newton', 'warp'}: + raise ImportError('blocked ' + fullname) +sys.meta_path.insert(0, BlockOptional()) +from embodichain_models.dynamics.rigidformer import ( + RigidFormerModel, RigidFormerObjective, RigidBodySequenceSource, + build_rigidformer_trainer, +) +assert "rigidformer" not in sys.modules +try: + RigidFormerModel() +except ImportError as exc: + assert 'optional rigidformer dependency' in str(exc) + assert 'embodichain[rigidformer]' in str(exc) + assert 'requirements-reference.txt' in str(exc) +else: + raise AssertionError('Missing optional dependency must be explicit') +""" + subprocess.run( + [sys.executable, "-c", code], check=True, capture_output=True, text=True + ) + + +def test_rigidformer_objective_separates_targets_and_preserves_reference_update( + geometry, +): + upstream = pytest.importorskip("rigidformer") + torch.set_num_threads(2) + cfg = RigidFormerCfg( + dim=16, + dim_head=6, + heads=2, + num_register_tokens=2, + object_self_attn_depth=1, + anchor_cross_attn_depth=2, + object_hidden_layers=(0, 1), + ) + candidate = RigidFormerModel(cfg) + reference_cfg = cfg.to_dict() + reference_cfg.pop("max_pairwise_bytes") + reference = upstream.Rigidformer(**reference_cfg) + reference.load_state_dict(deepcopy(candidate.network.state_dict())) + reference_optimizer = torch.optim.AdamW(reference.parameters(), lr=3e-4) + optimizer = torch.optim.AdamW(candidate.parameters(), lr=3e-4) + batch = rigidformer_batch(geometry, metadata={"source_id": "test"}) + assert "object_pos_next" not in batch.inputs + assert batch.targets is geometry["object_pos_next"] + reference_optimizer.zero_grad(set_to_none=True) + state = torch.get_rng_state() + loss, _ = reference(**geometry) + loss.backward() + torch.nn.utils.clip_grad_norm_(reference.parameters(), 1.0) + reference_optimizer.step() + torch.set_rng_state(state) + metrics = SupervisedLearner( + candidate, RigidFormerObjective(), optimizer, max_grad_norm=1.0 + ).update(batch) + assert metrics["loss"] == float(loss) + assert metrics["weight"] == 4.0 # Two scenes, each with two valid objects. + for name, expected in reference.state_dict().items(): + torch.testing.assert_close( + candidate.network.state_dict()[name], expected, atol=0, rtol=0 + ) + + +@pytest.mark.parametrize("object_lens", [None, (1, 2)]) +def test_rigidformer_split_scene_weighted_loss_and_gradients_match_upstream_batch( + geometry, object_lens +): + pytest.importorskip("rigidformer") + torch.set_num_threads(2) + torch.manual_seed(31) + candidate = RigidFormerModel( + RigidFormerCfg( + dim=16, + dim_head=6, + heads=2, + num_register_tokens=2, + object_self_attn_depth=1, + anchor_cross_attn_depth=2, + object_hidden_layers=(0, 1), + ) + ).eval() + # Fixed valid anchors remove FPS randomness; the scenes have different + # target accelerations so treating their local means equally is observable. + geometry["object_pos_next"][0] += torch.tensor([0.04, -0.03, 0.02]) + geometry["object_pos_next"][1] += torch.tensor([-0.02, 0.04, -0.01]) + if object_lens is not None: + geometry["object_lens"] = torch.tensor(object_lens) + reference = deepcopy(candidate.network) + expected_loss, _ = reference(**geometry) + expected_loss.backward() + expected_gradients = { + name: parameter.grad.clone() + for name, parameter in reference.named_parameters() + if parameter.grad is not None + } + objective = RigidFormerObjective() + results = [ + objective.compute_loss( + candidate, + rigidformer_batch( + {key: value[row : row + 1] for key, value in geometry.items()} + ), + ) + for row in range(2) + ] + assert [result.weight for result in results] == ( + [2.0, 2.0] if object_lens is None else [1.0, 2.0] + ) + actual_loss = sum(result.loss * result.weight for result in results) / sum( + result.weight for result in results + ) + torch.testing.assert_close(actual_loss, expected_loss, atol=1e-5, rtol=1e-4) + actual_loss.backward() + for name, parameter in candidate.network.named_parameters(): + if name in expected_gradients: + torch.testing.assert_close( + parameter.grad, expected_gradients[name], atol=1e-5, rtol=1e-4 + ) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32, torch.float64]) +def test_rigidformer_reference_geometry_and_distance_budget_are_explicit( + geometry, dtype +): + pytest.importorskip("rigidformer") + model = RigidFormerModel(RigidFormerCfg(max_pairwise_bytes=16)) + with pytest.raises(ValueError, match="Pairwise distances"): + model(**geometry) + # CPU autocast promotes BF16 cdist inputs to an FP32 distance matrix. + previous = geometry["object_pos_prev"][:1, :1, :4].to(dtype=dtype) + tiny_geometry = { + "object_pos": previous + 0.01, + "object_pos_prev": previous, + "object_first_frame_pos": previous.clone(), + "vertex_properties": torch.ones(1, 1, 3, dtype=dtype), + "delta_times": torch.tensor([0.1]), + "anchor_indices": torch.arange(4).reshape(1, 1, 4), + } + with torch.no_grad(), torch.autocast("cpu", dtype=torch.bfloat16): + flat = previous.flatten(1, 2) + distances = torch.cdist(flat, flat) + expected_dtype = torch.float64 if dtype == torch.float64 else torch.float32 + assert distances.dtype == expected_dtype + required_bytes = distances.numel() * distances.element_size() + model.cfg.max_pairwise_bytes = required_bytes - 1 + with pytest.raises(ValueError, match=f"require {required_bytes} bytes"): + model(**tiny_geometry) + del geometry["object_first_frame_pos"] + with pytest.raises(ValueError, match="Missing geometry"): + rigidformer_batch(geometry) + + +def test_rigidformer_rollout_matches_reference_and_retains_rigidity(geometry): + upstream = pytest.importorskip("rigidformer") + torch.set_num_threads(2) + model = RigidFormerModel( + RigidFormerCfg( + dim=16, + dim_head=6, + heads=2, + object_self_attn_depth=1, + anchor_cross_attn_depth=2, + object_hidden_layers=(0, 1), + ) + ) + initial = [geometry["object_pos_prev"], geometry["object_pos"]] + kwargs = dict( + object_positions=initial, + vertex_properties=geometry["vertex_properties"], + delta_times=geometry["delta_times"], + anchor_indices=geometry["anchor_indices"], + num_steps=3, + ) + with torch.no_grad(): + expected = upstream.RigidformerRolloutWrapper(model.network)(**kwargs) + actual = model.rollout(**kwargs) + assert len(actual) == 5 + torch.testing.assert_close( + torch.stack(actual), torch.stack(expected), atol=0, rtol=0 + ) + # Rigid projection preserves pairwise distances within each object. + reference_distances = torch.cdist( + initial[0], initial[0], compute_mode="donot_use_mm_for_euclid_dist" + ) + torch.testing.assert_close( + torch.cdist( + actual[-1], actual[-1], compute_mode="donot_use_mm_for_euclid_dist" + ), + reference_distances, + atol=2e-5, + rtol=2e-5, + ) + + +@pytest.mark.parametrize("points", [48, 128]) +def test_rigidformer_unequal_lengths_ignore_finite_padding_in_objective_and_rollout( + points, +): + pytest.importorskip("rigidformer") + torch.set_num_threads(2) + torch.manual_seed(31) + model = RigidFormerModel( + RigidFormerCfg( + dim=16, + dim_head=6, + heads=2, + object_self_attn_depth=1, + anchor_cross_attn_depth=2, + object_hidden_layers=(0, 1), + ) + ).eval() + previous = torch.randn(2, 3, points, 3) * 0.1 + torch.tensor([0.0, 0.0, 0.4]) + values = dict( + object_pos_prev=previous, + object_pos=previous + 0.01, + object_first_frame_pos=previous.clone(), + object_pos_next=previous + 0.025, + vertex_properties=torch.randn(2, 3, points, 3), + delta_times=torch.full((2,), 0.1), + object_lens=torch.tensor([2, 3]), + object_point_lens=torch.tensor([[16, 32, points], [32, points, 16]]), + ) + valid = (torch.arange(3)[None, :, None] < values["object_lens"][:, None, None]) & ( + torch.arange(points)[None, None, :] < values["object_point_lens"][..., None] + ) + changed = {key: value.clone() for key, value in values.items()} + for key in ( + "object_pos_prev", + "object_pos", + "object_first_frame_pos", + "object_pos_next", + "vertex_properties", + ): + changed[key][~valid] = torch.randn_like(changed[key][~valid]) * 3 + rng = torch.get_rng_state() + objective = RigidFormerObjective() + expected = objective.compute_loss(model, rigidformer_batch(values)) + assert expected.weight == 5.0 # Valid object counts are two and three. + expected.loss.backward() + gradients = { + name: parameter.grad.clone() + for name, parameter in model.named_parameters() + if parameter.grad is not None + } + model.zero_grad(set_to_none=True) + torch.set_rng_state(rng) + actual = objective.compute_loss(model, rigidformer_batch(changed)) + assert actual.weight == expected.weight + actual.loss.backward() + torch.testing.assert_close(actual.loss, expected.loss, atol=0, rtol=0) + for name, parameter in model.named_parameters(): + if name in gradients: + torch.testing.assert_close(parameter.grad, gradients[name], atol=0, rtol=0) + trajectories = [] + for inputs in (values, changed): + torch.set_rng_state(rng) + with torch.no_grad(): + trajectories.append( + model.rollout( + [inputs["object_pos_prev"], inputs["object_pos"]], + vertex_properties=inputs["vertex_properties"], + delta_times=inputs["delta_times"], + object_lens=inputs["object_lens"], + object_point_lens=inputs["object_point_lens"], + num_steps=3, + ) + ) + torch.testing.assert_close( + torch.stack(trajectories[0])[:, valid], + torch.stack(trajectories[1])[:, valid], + atol=0, + rtol=0, + ) + + +def test_rigidformer_rejects_unsupported_point_masking_and_invalid_anchors(geometry): + pytest.importorskip("rigidformer") + model = RigidFormerModel() + values = { + key: value.repeat_interleave(5, dim=2) if value.ndim == 4 else value + for key, value in geometry.items() + } + values["object_point_lens"] = torch.tensor([[16, 160], [64, 160]]) + with pytest.raises(ValueError, match="PointNet FPS/KNN leaks finite padding"): + RigidFormerObjective().compute_loss(model, rigidformer_batch(values)) + with pytest.raises(ValueError, match="PointNet FPS/KNN leaks finite padding"): + model.rollout( + [values["object_pos_prev"], values["object_pos"]], + vertex_properties=values["vertex_properties"], + delta_times=values["delta_times"], + object_point_lens=values["object_point_lens"], + num_steps=1, + ) + geometry["object_point_lens"] = torch.tensor([[16, 32], [32, 16]]) + with pytest.raises(ValueError, match="anchor_indices must index valid points"): + model(**geometry) + + +def test_rigidformer_offline_recipe_restores_the_shared_training_boundary(tmp_path): + pytest.importorskip("rigidformer") + torch.set_num_threads(2) + generator = torch.Generator().manual_seed(9) + first = torch.randn(2, 1, 2, 32, 3, generator=generator) + positions = first + torch.arange(6).reshape(1, 6, 1, 1, 1) * 0.01 + path = tmp_path / "geometry.pt" + torch.save( + {"positions": positions, "properties": torch.ones(2, 2, 3), "step_dt": 0.1}, + path, + ) + config = { + "data": {"path": str(path)}, + "model": { + "dim": 16, + "dim_head": 6, + "heads": 2, + "object_self_attn_depth": 1, + "anchor_cross_attn_depth": 2, + "object_hidden_layers": (0, 1), + }, + "trainer": { + "total_updates": 2, + "batch_size": 2, + "output_dir": str(tmp_path / "run"), + }, + } + torch.manual_seed(31) + uninterrupted = build_rigidformer_trainer(config) + uninterrupted.fit() + expected = deepcopy(uninterrupted.learner.model.state_dict()) + torch.manual_seed(31) + config["trainer"]["output_dir"] = str(tmp_path / "interrupted") + interrupted = build_rigidformer_trainer(config) + result = interrupted.fit(stop_after=1) + config["trainer"]["output_dir"] = str(tmp_path / "resumed") + resumed = build_rigidformer_trainer(config) + resumed.resume(result["checkpoint"]) + resumed.fit() + for name, value in expected.items(): + torch.testing.assert_close( + resumed.learner.model.state_dict()[name], value, atol=0, rtol=0 + ) diff --git a/tests/learning/models/test_rigidformer_data.py b/tests/learning/models/test_rigidformer_data.py new file mode 100644 index 0000000000..992a7425fc --- /dev/null +++ b/tests/learning/models/test_rigidformer_data.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. +# ---------------------------------------------------------------------------- + +"""Causal geometry window ownership and episode-boundary coverage.""" + +from __future__ import annotations + +import pytest +import torch + +from embodichain_models.dynamics.rigidformer.data import RigidBodySequenceSource + + +def test_geometry_windows_keep_point_identity_and_owned_storage(): + # Every point coordinate encodes its episode/time index independently. + positions = torch.arange(2 * 7 * 2 * 4 * 3).reshape(2, 7, 2, 4, 3).float() + expected = positions.clone() + source = RigidBodySequenceSource( + positions, + torch.ones(2, 2, 3), + step_dt=0.05, + source_id="fixed-corpus", + stride=2, + episode_lengths=torch.tensor([5, 7]), + ) + positions.zero_() + batch = source.get_batch(torch.tensor([0, 1]), torch.tensor([2, 4])) + torch.testing.assert_close(batch["object_pos_prev"], expected[[0, 1], [0, 2]]) + torch.testing.assert_close(batch["object_pos"], expected[[0, 1], [2, 4]]) + torch.testing.assert_close(batch["object_pos_next"], expected[[0, 1], [4, 6]]) + torch.testing.assert_close( + batch["object_first_frame_pos"], batch["object_pos_prev"] + ) + torch.testing.assert_close(batch["delta_times"], torch.full((2,), 0.1)) + batch["object_pos_next"].zero_() + again = source.get_batch(torch.tensor([0]), torch.tensor([2])) + torch.testing.assert_close(again["object_pos_next"][0], expected[0, 4]) + with pytest.raises(ValueError, match="episode boundary"): + source.get_batch(torch.tensor([0]), torch.tensor([3])) + + +def test_geometry_sampling_replays_rng_and_preserves_valid_lengths(): + source = RigidBodySequenceSource( + torch.randn(2, 8, 3, 5, 3), + torch.randn(2, 3, 5, 2), + step_dt=0.1, + source_id="ragged", + episode_lengths=torch.tensor([3, 8]), + object_lens=torch.tensor([1, 3]), + object_point_lens=torch.tensor([[4, 5, 5], [3, 4, 5]]), + ) + first = source.sample_batch(20, generator=torch.Generator().manual_seed(3)) + second = source.sample_batch(20, generator=torch.Generator().manual_seed(3)) + for key in first: + torch.testing.assert_close(first[key], second[key], rtol=0, atol=0) + assert set(first["object_lens"].tolist()) == {1, 3} + assert first["vertex_properties"].shape == (20, 3, 5, 2) + + +def test_geometry_memory_limits_precede_payload_delivery(): + positions, props = torch.zeros(1, 4, 2, 4, 3), torch.ones(1, 2, 3) + with pytest.raises(ValueError, match="memory budget"): + RigidBodySequenceSource( + positions, props, step_dt=0.1, source_id="bounded", max_resident_bytes=1 + ) + source = RigidBodySequenceSource( + positions, props, step_dt=0.1, source_id="bounded", max_batch_bytes=16 + ) + with pytest.raises(ValueError, match="max_batch_bytes"): + source.sample_batch(1)