Add official RigidFormer architecture and weight import

agent_context/MAP.yaml

@@ -270,10 +270,10 @@ 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, RigidFormer, 神经动力学, 动力学模型训练, NeRD 示例, RigidFormer 示例]
   keywords: [NeRDStateCodec, NeRDModel, DynamicsGenerationCfg,
     DynamicsShardWriter, DynamicsWindowDataset, DynamicsTrainer, embodichain_models,
-    learning/nerd, nerd reference assets,
+    learning/nerd, nerd reference assets, RigidFormerModel,
     官方训练精度对比]
   paths: [topics/neural-dynamics/neural-dynamics.md]
   source_of_truth:
@@ -282,6 +282,7 @@ topics:
   - examples/learning/nerd/_resources.py
   - examples/learning/nerd/_support/
   - examples/learning/nerd/configs/
+  - extensions/embodichain_models/embodichain_models/dynamics/rigidformer/
   watch_paths: [extensions/embodichain_models/, examples/learning/nerd/,
     docs/source/features/neural_dynamics.md,
     docs/source/api_reference/model_catalogs.rst]

agent_context/topics/neural-dynamics/neural-dynamics.md

@@ -40,6 +40,19 @@ 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.
 
+## RigidFormer author model
+
+The native family lives in
+`extensions/embodichain_models/embodichain_models/dynamics/rigidformer/`.
+`model.py` owns geometry/configuration; `_native/` preserves licensed author
+code. Checkpoint import validates recorded architecture and named tensors.
+The community `rigidformer` package is not a runtime dependency.
+
+Integration dt=1 differs from physical seconds. Geometry prediction requires
+CUDA and at least two valid objects; construction and weight inspection support
+CPU. Attention backend selection is explicit. The model predicts geometry
+without action inputs and does not qualify state/action gradients.
+
 ## Resource and physical identity
 
 `prepare --robot cartpole|ant --resource-root DIR` downloads the pinned dataset,

docs/source/api_reference/index.rst

@@ -94,6 +94,7 @@ documentation. CI runs this same checker after style checks and before tests.
 
    robot_catalogs
    model_catalogs
+   model_rigidformer
 
 
 .. toctree::

docs/source/api_reference/model_rigidformer.rst

@@ -0,0 +1,70 @@
+RigidFormer geometry dynamics
+=============================
+
+The model family implements the author's RigidFormer architecture from revision
+160fe0852c51cffae7b7440452b58d0b49db6df6. Licensed implementation code is private
+to the model package; public adapters do not import the research showcase or a
+simulator. The optional extra installs numerical/data dependencies, without the
+community rigidformer package.
+
+Geometry uses aligned material points, explicit reference geometry, valid
+object/vertex masks and physical mass/friction/restitution. The author KNN path
+requires CUDA and at least two valid objects per scene. Model construction and
+weight inspection can run on CPU. The geometry integration unit dt=1 is distinct
+from physical intervals step_size / 240 seconds.
+
+Attention uses explicit SDPA or an optional compatible Flash-Attention build.
+Backend selection is explicit in model configuration. Inference rollout restores
+the caller's module modes and does not retain a gradient graph.
+
+Public model family
+-------------------
+
+.. currentmodule:: embodichain_models.dynamics.rigidformer
+
+.. autosummary::
+
+   RIGIDFORMER_REVISION
+   RigidFormerCfg
+   RigidFormerModel
+   load_rigidformer_checkpoint
+
+.. autodata:: RIGIDFORMER_REVISION
+
+.. autoclass:: RigidFormerCfg
+   :members:
+
+.. autoclass:: RigidFormerModel
+   :members:
+
+.. autofunction:: load_rigidformer_checkpoint
+
+Architecture and geometry contract
+----------------------------------
+
+.. currentmodule:: embodichain_models.dynamics.rigidformer.model
+
+.. autosummary::
+
+   RIGIDFORMER_REVISION
+   RigidFormerCfg
+   RigidFormerModel
+
+.. autodata:: RIGIDFORMER_REVISION
+
+.. autoclass:: RigidFormerCfg
+   :members:
+
+.. autoclass:: RigidFormerModel
+   :members:
+
+Official inference weights
+--------------------------
+
+.. currentmodule:: embodichain_models.dynamics.rigidformer.checkpoint
+
+.. autosummary::
+
+   load_rigidformer_checkpoint
+
+.. autofunction:: load_rigidformer_checkpoint

extensions/embodichain_models/README.md

@@ -35,7 +35,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` | Implemented. |
-| RigidFormer | `embodichain_models.dynamics.rigidformer` | Target for the RigidFormer follow-up; not created yet. |
+| RigidFormer | `embodichain_models.dynamics.rigidformer` | Native author architecture and strict official weights. |
 
 RWM predicts action-conditioned robot state transitions, so it belongs to
 `dynamics`. World models also describe future-state/observation prediction;

extensions/embodichain_models/embodichain_models/dynamics/rigidformer/REFERENCE_SOURCES.json

@@ -0,0 +1,28 @@
+{
+  "files": {
+    "src/data/gpu_augmentation.py": "944ddafab989467755a5efd8306a939e1370e191ccd9457f3d5885dd7508e8e4",
+    "src/data/multi_dataloader.py": "0ba2cc56a994ed4fd03f196c0921d5f4ab536b9fc7b68b481531f1ebfab08f8c",
+    "src/data/object_scale.py": "e68ba74455e92b23e14a067d3cdacc5daaf713596e137883c11a0e812e215d6d",
+    "src/data/partial_view.py": "e0c9c36d385bae969f3a7c32ebba16469360a67735bad10abd744981337b5ac6",
+    "src/data/scene_frames.py": "e11b454bf29a631a182618892d02ecbe3117c24e91a365b4ed64bd5407bffdee",
+    "src/data/shape_completion.py": "7bf093e54b687e8ccaa10ef384aaab5962350e0d3c5287b302c6cff0cf9691f7",
+    "src/models/anchor_vertex_pooling.py": "4162975292918bf7a91d3f3cd4542d07f16062cba245511c306a4131ccf0fe8b",
+    "src/models/decoder.py": "4d8d2152f09b2ebbc04747ad426ff6e81db689e7a39290477e09234bee569ec5",
+    "src/models/encoder.py": "4764f3ff377708b128e42a403eae573ff2e5dc6a139eab7aa7948a13b43d9a49",
+    "src/models/geometry_encoders.py": "78b2007a54d8e3b8ae55bb234dc5dffac601e065b717d17ec66002c7802ee9ef",
+    "src/models/multi_scale_predictor.py": "8e634f4b745e1458c5f5f48a5da37a985446f40fa1730f07cbac6506dfd937b5",
+    "src/models/physformer_rope.py": "66a0743e4f4eef324694066effa5b9ee7076c02870d240888f8bad611ab492a4",
+    "src/models/predictor.py": "05490a8d6c7bc8915fb758ad9226509713b10bf5c3478431583d0cd215253cbf",
+    "src/models/se3_attention.py": "18117ace2b2e9077bacadcf187a1932cd62c38eb15165f5ab9ef6306ba638a07",
+    "src/models/state.py": "85952c21b7b4cf86a6e01dfac80ef7fd40f78f2093b581f6fc9a376e49af0bfa",
+    "src/models/transform_layer.py": "39acea7d4b35454ff106fed7e047111d7dcd01606af31586efc4ce81239700c5",
+    "src/train.py": "4e3dab822259bf0d4e06a512ed537d9a7a9a2855eb61434571e42401481166e1",
+    "src/training/evaluator.py": "abb585f50c3d78adb04cfa76a67d65d51d11d2d4d08e1999304e6a0516ce1a86",
+    "src/utils/anchor_rotation.py": "9d732ce7f236330e1da91f83e0046784ede3999b4ab340b1a1c595f160b7bd89",
+    "src/utils/cuda_knn.py": "0db3acf65d76a05feb2fd5c920517f8c3d5804b04bc526c345912c0172b848f4",
+    "src/utils/intersection_handler.py": "9bb9ca50513faf16fab55dac596381e01a447e9dc1ea19a1632ad43a7a0936f8",
+    "src/utils/noise_utils.py": "f85be4e298a26be30c67fc5b91ee5137a3f475851f71aed820537570b03fb34d",
+    "src/utils/weight_ema.py": "f0599c2378f56ff9e6526b9dbc725a280186a5aa4c8c02c6fd38f7212fafbd81"
+  },
+  "revision": "160fe0852c51cffae7b7440452b58d0b49db6df6"
+}

extensions/embodichain_models/embodichain_models/dynamics/rigidformer/THIRD_PARTY_LICENSES

@@ -0,0 +1,25 @@
+RigidFormer author implementation
+https://github.com/Frank-ZY-Dou/Dynamics-Modeling/tree/160fe0852c51cffae7b7440452b58d0b49db6df6/RigidFormer
+
+MIT License
+
+Copyright (c) 2026 Zhiyang Dou
+
+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.
+

extensions/embodichain_models/embodichain_models/dynamics/rigidformer/__init__.py

@@ -0,0 +1,42 @@
+# ----------------------------------------------------------------------------
+# Copyright (c) 2021-2026 DexForce Technology Co., Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ----------------------------------------------------------------------------
+
+"""Author-based rigid geometry prediction with lazy optional dependencies."""
+
+from __future__ import annotations
+from importlib import import_module
+from typing import Any
+
+__all__ = [
+    "RIGIDFORMER_REVISION",
+    "RigidFormerCfg",
+    "RigidFormerModel",
+    "load_rigidformer_checkpoint",
+]
+_EXPORTS = {
+    "RIGIDFORMER_REVISION": "model",
+    "RigidFormerCfg": "model",
+    "RigidFormerModel": "model",
+    "load_rigidformer_checkpoint": "checkpoint",
+}
+
+
+def __getattr__(name: str) -> Any:
+    if name not in _EXPORTS:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    value = getattr(import_module(f".{_EXPORTS[name]}", __name__), name)
+    globals()[name] = value
+    return value

extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/__init__.py

@@ -0,0 +1,10 @@
+# Derived from Zhiyang Dou's RigidFormer, revision 160fe0852c51cffae7b7440452b58d0b49db6df6.
+# Copyright (c) 2026 Zhiyang Dou. SPDX-License-Identifier: MIT
+# Adaptations: package-relative imports; framework-owned public interfaces.
+# The complete MIT license is in ../THIRD_PARTY_LICENSES.
+
+"""Private MIT-licensed author implementation; use the public model API."""
+
+from __future__ import annotations
+
+__all__ = []

extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/attention.py

@@ -0,0 +1,51 @@
+# ----------------------------------------------------------------------------
+# 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.
+# ----------------------------------------------------------------------------
+
+"""Explicit SDPA and optional Flash-Attention backends for author attention."""
+
+from __future__ import annotations
+import torch
+from torch.nn import functional as F
+
+
+def attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    *,
+    dropout_p: float = 0.0,
+    softmax_scale: float | None = None,
+    backend: str = "sdpa",
+) -> torch.Tensor:
+    if backend == "flash":
+        try:
+            from flash_attn import flash_attn_func
+        except ImportError as exc:
+            raise ImportError(
+                "The flash backend requires a Flash-Attention build matching the installed Torch/CUDA stack."
+            ) from exc
+        return flash_attn_func(
+            q, k, v, dropout_p=dropout_p, softmax_scale=softmax_scale
+        )
+    if backend != "sdpa":
+        raise ValueError(f"Unsupported attention backend: {backend!r}")
+    return F.scaled_dot_product_attention(
+        q.transpose(1, 2),
+        k.transpose(1, 2),
+        v.transpose(1, 2),
+        dropout_p=dropout_p,
+        scale=softmax_scale,
+    ).transpose(1, 2)

extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/__init__.py

@@ -0,0 +1,10 @@
+# Derived from Zhiyang Dou's RigidFormer, revision 160fe0852c51cffae7b7440452b58d0b49db6df6.
+# Copyright (c) 2026 Zhiyang Dou. SPDX-License-Identifier: MIT
+# Adaptations: package-relative imports; framework-owned public interfaces.
+# The complete MIT license is in ../THIRD_PARTY_LICENSES.
+
+"""Private MIT-licensed author implementation; use the public model API."""
+
+from __future__ import annotations
+
+__all__ = []