From 65f42a2dbd12dbaa19dd08fe59f38ad38a7e80f7 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Tue, 27 Oct 2026 12:07:53 +0000 Subject: [PATCH 1/3] feat(models): add the official RigidFormer architecture --- agent_context/MAP.yaml | 5 +- .../topics/neural-dynamics/neural-dynamics.md | 13 + docs/source/api_reference/index.rst | 1 + .../api_reference/model_rigidformer.rst | 70 + extensions/embodichain_models/README.md | 2 +- .../rigidformer/REFERENCE_SOURCES.json | 28 + .../dynamics/rigidformer/THIRD_PARTY_LICENSES | 25 + .../dynamics/rigidformer/__init__.py | 42 + .../dynamics/rigidformer/_native/__init__.py | 10 + .../dynamics/rigidformer/_native/attention.py | 51 + .../rigidformer/_native/models/__init__.py | 10 + .../_native/models/anchor_vertex_pooling.py | 136 + .../rigidformer/_native/models/decoder.py | 1370 +++++++ .../rigidformer/_native/models/encoder.py | 937 +++++ .../_native/models/geometry_encoders.py | 254 ++ .../_native/models/multi_scale_predictor.py | 885 +++++ .../_native/models/physformer_rope.py | 826 +++++ .../rigidformer/_native/models/predictor.py | 282 ++ .../_native/models/se3_attention.py | 336 ++ .../rigidformer/_native/models/state.py | 423 +++ .../_native/models/transform_layer.py | 418 +++ .../dynamics/rigidformer/_native/network.py | 3148 +++++++++++++++++ .../rigidformer/_native/utils/__init__.py | 10 + .../_native/utils/anchor_rotation.py | 31 + .../rigidformer/_native/utils/cuda_knn.py | 671 ++++ .../_native/utils/intersection_handler.py | 232 ++ .../rigidformer/_native/utils/weight_ema.py | 92 + .../dynamics/rigidformer/checkpoint.py | 78 + .../dynamics/rigidformer/configs/movi_a.yaml | 284 ++ .../dynamics/rigidformer/configs/movi_b.yaml | 284 ++ .../rigidformer/configs/movi_sphere.yaml | 284 ++ .../dynamics/rigidformer/model.py | 254 ++ pyproject.toml | 5 + setup.py | 5 + tests/models/dynamics/__init__.py | 19 + .../models/dynamics/_rigidformer_fixtures.py | 60 + .../dynamics/test_rigidformer_official.py | 98 + 37 files changed, 11676 insertions(+), 3 deletions(-) create mode 100644 docs/source/api_reference/model_rigidformer.rst create mode 100644 extensions/embodichain_models/embodichain_models/dynamics/rigidformer/REFERENCE_SOURCES.json create mode 100644 extensions/embodichain_models/embodichain_models/dynamics/rigidformer/THIRD_PARTY_LICENSES create mode 100644 extensions/embodichain_models/embodichain_models/dynamics/rigidformer/__init__.py create mode 100644 extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/__init__.py create mode 100644 extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/attention.py create mode 100644 extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/__init__.py create mode 100644 extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/anchor_vertex_pooling.py create mode 100644 extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/decoder.py create mode 100644 extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/encoder.py create mode 100644 extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/geometry_encoders.py create mode 100644 extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/multi_scale_predictor.py create mode 100644 extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/physformer_rope.py create mode 100644 extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/predictor.py create mode 100644 extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/se3_attention.py create mode 100644 extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/state.py create mode 100644 extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/transform_layer.py create mode 100644 extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/network.py create mode 100644 extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/utils/__init__.py create mode 100644 extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/utils/anchor_rotation.py create mode 100644 extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/utils/cuda_knn.py create mode 100644 extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/utils/intersection_handler.py create mode 100644 extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/utils/weight_ema.py create mode 100644 extensions/embodichain_models/embodichain_models/dynamics/rigidformer/checkpoint.py create mode 100644 extensions/embodichain_models/embodichain_models/dynamics/rigidformer/configs/movi_a.yaml create mode 100644 extensions/embodichain_models/embodichain_models/dynamics/rigidformer/configs/movi_b.yaml create mode 100644 extensions/embodichain_models/embodichain_models/dynamics/rigidformer/configs/movi_sphere.yaml create mode 100644 extensions/embodichain_models/embodichain_models/dynamics/rigidformer/model.py create mode 100644 tests/models/dynamics/__init__.py create mode 100644 tests/models/dynamics/_rigidformer_fixtures.py create mode 100644 tests/models/dynamics/test_rigidformer_official.py diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index 9f59109509..548b67337c 100644 --- a/agent_context/MAP.yaml +++ b/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] diff --git a/agent_context/topics/neural-dynamics/neural-dynamics.md b/agent_context/topics/neural-dynamics/neural-dynamics.md index 7d449ccd98..1a05799ba9 100644 --- a/agent_context/topics/neural-dynamics/neural-dynamics.md +++ b/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, diff --git a/docs/source/api_reference/index.rst b/docs/source/api_reference/index.rst index 2403b38c45..b209d3e698 100644 --- a/docs/source/api_reference/index.rst +++ b/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:: diff --git a/docs/source/api_reference/model_rigidformer.rst b/docs/source/api_reference/model_rigidformer.rst new file mode 100644 index 0000000000..61db20c7c6 --- /dev/null +++ b/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 diff --git a/extensions/embodichain_models/README.md b/extensions/embodichain_models/README.md index bf4f28d533..bc3980232c 100644 --- a/extensions/embodichain_models/README.md +++ b/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; diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/REFERENCE_SOURCES.json b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/REFERENCE_SOURCES.json new file mode 100644 index 0000000000..d324431f65 --- /dev/null +++ b/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" +} 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..2056512783 --- /dev/null +++ b/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. + 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..d74413c98d --- /dev/null +++ b/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 diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/__init__.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/__init__.py new file mode 100644 index 0000000000..9071a4c0b7 --- /dev/null +++ b/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__ = [] diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/attention.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/attention.py new file mode 100644 index 0000000000..6f101be2e6 --- /dev/null +++ b/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) diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/__init__.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/__init__.py new file mode 100644 index 0000000000..9071a4c0b7 --- /dev/null +++ b/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__ = [] diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/anchor_vertex_pooling.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/anchor_vertex_pooling.py new file mode 100644 index 0000000000..6f176691c4 --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/anchor_vertex_pooling.py @@ -0,0 +1,136 @@ +# 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. + +""" +Anchor-Vertex Pooling (AVP) — collision-aware anchor feature enrichment. + +Each anchor aggregates per-vertex features from its spatial neighborhood using +exponential distance weighting exp(-||a - v|| / sigma). This gives anchors +collision-aware representations without requiring vertex-level attention +(O(Na*Nv) instead of O(Nv^2)). + +Properties: +- Translation consistent: weights depend only on the relative distance ||anchor - vertex|| +- Permutation equivariant: operates per-object independently +- Learnable sigma: the model learns the neighborhood radius +""" + +from __future__ import annotations + +import torch +import torch.nn as nn +import logging + +logger = logging.getLogger(__name__) + + +class AnchorVertexPooling(nn.Module): + """ + Distance-based attention pooling from vertices to anchors. + + For each anchor, computes exponential distance weights exp(-||a - v|| / sigma) + over all vertices, then pools per-vertex features with the normalized weights. + + Args: + vertex_dim: Dimension of per-vertex features from the encoder backbone (e.g. 1024) + output_dim: Output dimension per anchor (e.g. 256) + init_sigma: Initial value of the learnable kernel radius sigma + """ + + def __init__( + self, vertex_dim: int = 1024, output_dim: int = 256, init_sigma: float = 0.1 + ): + super().__init__() + self.vertex_dim = vertex_dim + self.output_dim = output_dim + + # Learnable neighborhood radius (log-space for numerical stability) + # sigma = exp(log_sigma), so sigma is always positive + self.log_sigma = nn.Parameter(torch.tensor(float(init_sigma)).log()) + + # Projection: vertex_dim -> output_dim + self.proj = nn.Sequential( + nn.Linear(vertex_dim, output_dim), + nn.SiLU(), + nn.Linear(output_dim, output_dim), + ) + + # Zero-init last layer: the module starts out contributing near-zero features + nn.init.zeros_(self.proj[-1].weight) + nn.init.zeros_(self.proj[-1].bias) + + self._logged = False + + def forward( + self, + vertex_features: torch.Tensor, # (B_flat, Nv, vertex_dim) + vertex_positions: torch.Tensor, # (B_flat, Nv, 3) + anchor_positions: torch.Tensor, # (B_flat, Na, 3) + vtx_mask: torch.Tensor, # (B_flat, Nv) bool + ) -> torch.Tensor: + """ + Args: + vertex_features: (B_flat, Nv, vertex_dim) — per-vertex backbone features; + the caller must transpose from the encoder's (B_flat, vertex_dim, Nv) + vertex_positions: (B_flat, Nv, 3) — current frame vertex positions + anchor_positions: (B_flat, Na, 3) — current frame anchor positions + vtx_mask: (B_flat, Nv) — True for valid vertices, False for padding + + Returns: + anchor_features: (B_flat, Na, output_dim) + """ + B_flat, Nv, D = vertex_features.shape + Na = anchor_positions.shape[1] + + # Compute pairwise spatial distance: (B_flat, Na, Nv) + # anchor_positions: (B_flat, Na, 1, 3) - vertex_positions: (B_flat, 1, Nv, 3) + diff = anchor_positions.unsqueeze(2) - vertex_positions.unsqueeze( + 1 + ) # (B_flat, Na, Nv, 3) + dist = diff.norm(dim=-1) # (B_flat, Na, Nv) + + # Exponential distance weights + sigma = self.log_sigma.exp() # always positive + attn = torch.exp(-dist / (sigma + 1e-8)) # (B_flat, Na, Nv) + + # Mask padding vertices (set weight to 0) + attn = attn * vtx_mask.unsqueeze(1).float() # (B_flat, Na, Nv) + + # Normalize (sum over vertices) + attn_sum = attn.sum(dim=-1, keepdim=True) # (B_flat, Na, 1) + attn = attn / (attn_sum + 1e-8) # (B_flat, Na, Nv) + + # Weighted pooling of vertex features + pooled = torch.bmm(attn, vertex_features) # (B_flat, Na, vertex_dim) + + # Project to output dimension + output = self.proj(pooled) # (B_flat, Na, output_dim) + + # One-time diagnostic log + if not self._logged: + n_valid_vtx = vtx_mask.sum().item() + n_total_vtx = vtx_mask.numel() + n_fully_masked = ( + (vtx_mask.sum(dim=-1) == 0).sum().item() + ) # objects with all vertices masked + attn_max_on_masked = ( + (attn * (~vtx_mask).unsqueeze(1).float()).max().item() + ) # expected 0 + output_has_nan = torch.isnan(output).any().item() + logger.info( + f"[AVP] Anchor-Vertex Pooling: " + f"vertex_feat=({B_flat},{Nv},{D}), anchors={Na}, " + f"sigma={sigma.item():.4f}, output=({B_flat},{Na},{self.output_dim})" + ) + logger.info( + f"[AVP] Masking check: valid_vtx={n_valid_vtx}/{n_total_vtx} " + f"({100*n_valid_vtx/max(n_total_vtx,1):.1f}%), " + f"fully_masked_objs={n_fully_masked}/{B_flat}, " + f"attn_on_padding={attn_max_on_masked:.2e} (expected 0), " + f"has_nan={output_has_nan}" + ) + self._logged = True + + return output diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/decoder.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/decoder.py new file mode 100644 index 0000000000..ccc02b3fab --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/decoder.py @@ -0,0 +1,1370 @@ +# 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. + +""" +Bidirectional Transformer decoder inspired by CUT3R. +""" + +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.nn.functional as F +from ..attention import attention as flash_attn_func +from typing import Optional, Tuple, List +import math +import os +import sys +import logging +from pathlib import Path + +logger = logging.getLogger(__name__) + + +from embodichain_models.dynamics.rigidformer._native.models.physformer_rope import ( + PhysFormerRoPE, + apply_rotary_emb, +) + + +class RMSNorm(nn.Module): + """ + Root Mean Square Layer Normalization (RMSNorm). + More efficient than LayerNorm, used in modern transformers like LLaMA. + + Reference: Zhang and Sennrich 2019 - "Root Mean Square Layer Normalization" + """ + + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # RMS normalization: x / RMS(x) * weight + # RMS(x) = sqrt(mean(x^2) + eps) + rms = torch.sqrt(torch.mean(x**2, dim=-1, keepdim=True) + self.eps) + x_normed = x / rms + return self.weight * x_normed + + +def get_norm_layer(norm_type: str, dim: int) -> nn.Module: + """ + Factory function to create normalization layer. + + Args: + norm_type: 'layer_norm' or 'rms_norm' + dim: Dimension to normalize + + Returns: + Normalization module + """ + if norm_type == "rms_norm": + return RMSNorm(dim) + elif norm_type == "layer_norm": + return nn.LayerNorm(dim) + else: + raise ValueError( + f"Unknown norm_type: {norm_type}. Choose 'layer_norm' or 'rms_norm'." + ) + + +class BidirectionalDecoder(nn.Module): + """ + Bidirectional decoder following CUT3R's architecture. + + - Returns all layer outputs for multi-scale fusion + - Bidirectional cross-attention between state and object tokens + - RMSNorm and QK-normalization for training stability (configurable) + - Optional elementwise gated attention (Qwen3 style) + """ + + def __init__( + self, + state_dim: int = 768, + num_layers: int = 12, + num_heads: int = 12, + feedforward_dim: int = None, # None: 2.5x state_dim (SwiGLU) + dropout: float = 0.1, + activation: str = "swiglu", + return_all_layers: bool = True, # For multi-scale fusion + use_learnable_pe: bool = False, + use_rope: bool = True, + max_state_tokens: int = 1024, + max_obs_tokens: int = 256, + use_self_attention: bool = False, # Optional state self-attention + norm_type: str = "rms_norm", # 'layer_norm' or 'rms_norm' + qk_norm: bool = True, # QK-normalization for training stability + num_register_tokens: int = 16, # Learnable register tokens (RenderFormer) + center_rope_positions: bool = False, # Subtract the valid-object anchor centroid before RoPE (registers sit at the centroid) + rope_set_aggregation: bool = False, # Set aggregation for multi-point RoPE + rope_dim: int = None, # Custom RoPE dimension (None = auto-compute as head_dim//6*6) + dim_per_coord_3d: int = None, # Dim per coordinate for single-point (3D) RoPE + dim_per_coord_multi: int = None, # Dim per coordinate for multi-point RoPE + layer_indices_to_save: set = None, # Save only these layer outputs to reduce memory + elementwise_gate: bool = False, # Elementwise gated attention (Qwen3 style) + state_update_enabled: bool = True, # False: state is fixed learnable embeddings + objects_self_attention: bool = False, # Objects self-attention for object-object interaction + use_se3: bool = False, # SE(3) pairwise bias for object self-attention + # Block Attention Residuals (Kimi 2026, arXiv:2603.15031) + block_attn_res: bool = False, # Enable Block AttnRes for inter-layer residuals + block_attn_res_block_size: int = 4, # Layers per block (S in the paper) + block_attn_res_final_aggregate: bool = False, # Return the final depth aggregation as the last layer's features + film_conditioning: bool = False, # Per-layer FiLM conditioning on step size + ): + super().__init__() + + self.state_dim = state_dim + self.num_layers = num_layers + self.return_all_layers = return_all_layers + + # Optionally save only selected layer outputs to reduce memory + if layer_indices_to_save is not None: + for idx in layer_indices_to_save: + assert ( + 0 <= idx <= num_layers + ), f"Invalid layer index {idx} to save. Valid range: [0, {num_layers}]" + self.layer_indices_to_save = layer_indices_to_save + else: + # Default: save all layers if return_all_layers is True + self.layer_indices_to_save = ( + set(range(num_layers + 1)) if return_all_layers else set() + ) + self.use_rope = bool(use_rope) + self.rope_set_aggregation = rope_set_aggregation + # If RoPE is enabled, disable learnable positional embeddings in decoder inputs + self.use_learnable_pe = bool(use_learnable_pe and (not self.use_rope)) + assert self.use_rope, "RoPE must be enabled for decoder" + if feedforward_dim is None: + feedforward_dim = int(2.5 * state_dim) + + # Register tokens (RenderFormer): learnable tokens that provide scratch + # space for global information and reduce embedding noise. + self.num_register_tokens = num_register_tokens + self.center_rope_positions = bool(center_rope_positions) + if num_register_tokens > 0: + self.register_tokens = nn.Parameter( + torch.randn(1, num_register_tokens, state_dim) * 0.02 + ) + # Register positions are created in forward() to match the dimensionality + # of obj_anchor_pos (3D single-point or Nq*3 D multi-point RoPE). + else: + self.register_tokens = None + + # Transformer decoder layers + self.state_update_enabled = state_update_enabled + self.objects_self_attention = objects_self_attention + if not state_update_enabled: + print( + f"[Decoder] State update disabled: state acts as learnable embeddings (objects attend to state only)" + ) + if objects_self_attention: + print( + f"[Decoder] Objects self-attention enabled: direct object-to-object interaction" + ) + self.layers = nn.ModuleList( + [ + TransformerDecoderLayer( + state_dim, + num_heads, + feedforward_dim, + dropout, + activation, + use_self_attention=use_self_attention, + use_rope=self.use_rope, + norm_type=norm_type, + qk_norm=qk_norm, + rope_set_aggregation=self.rope_set_aggregation, + rope_dim=rope_dim, + dim_per_coord_3d=dim_per_coord_3d, + dim_per_coord_multi=dim_per_coord_multi, + elementwise_gate=elementwise_gate, + state_update_enabled=state_update_enabled, + objects_self_attention=objects_self_attention, + use_se3=use_se3, + film_conditioning=film_conditioning, + ) + for _ in range(num_layers) + ] + ) + + # Final layer normalization + self.state_norm = get_norm_layer(norm_type, state_dim) + self.obs_norm = get_norm_layer(norm_type, state_dim) + + # Block Attention Residuals (Kimi 2026, arXiv:2603.15031) + # Replaces standard inter-layer residuals with softmax attention over block representations + # h_l = sum_i alpha_{i->l} * v_i (attention over depth) instead of h_l = h_{l-1} + f_l(h_{l-1}) + self.block_attn_res = block_attn_res + self.block_attn_res_block_size = block_attn_res_block_size + # With return_all_layers the last returned features are the last layer's own output; + # block_attn_res_final_aggregate returns the final depth aggregation instead (off by default). + self.block_attn_res_final_aggregate = bool(block_attn_res_final_aggregate) + if self.block_attn_res_final_aggregate and not block_attn_res: + raise ValueError("block_attn_res_final_aggregate requires block_attn_res") + if self.block_attn_res_final_aggregate: + print( + "[Decoder] Block AttnRes final aggregation returned as the last layer's features" + ) + if block_attn_res: + N_blocks = max(1, num_layers // block_attn_res_block_size) + print( + f"[Decoder] Block AttnRes enabled: {num_layers} layers, block_size={block_attn_res_block_size}, N_blocks={N_blocks}" + ) + # Per-layer pseudo-query w_l (d-dimensional learnable vector) for depth attention + self.attn_res_queries = nn.ParameterList( + [nn.Parameter(torch.zeros(state_dim)) for _ in range(num_layers)] + ) + # Shared RMSNorm for keys in depth attention (prevents magnitude bias) + self.attn_res_key_norm = RMSNorm(state_dim) + + def forward( + self, + state: Optional[torch.Tensor], + obj_features: torch.Tensor, + state_pos: Optional[torch.Tensor] = None, + obj_anchor_pos: Optional[torch.Tensor] = None, + mask: Optional[torch.Tensor] = None, + se3_bias: Optional[ + torch.Tensor + ] = None, # (B, H, Nobj, Nobj) pairwise SE(3) bias + time_features: Optional[ + torch.Tensor + ] = None, # (B, 2) [s, s^2] for per-layer FiLM + ) -> Tuple[List[torch.Tensor], List[torch.Tensor]]: + """ + Process state and per-object features through bidirectional interaction. + Register tokens (RenderFormer) are prepended to object features. + + When state=None (no_global_state mode): + - Skips all state cross-attention + - Runs objects self-attention + objects FFN only + - Returns dummy state entries + + Args: + state: State tokens (B, N_state, D) or None for no-state mode + obj_features: Per-object encoder features (B, N_obj, D) + state_pos: State positions for RoPE (B, N_state, 3) or None + obj_anchor_pos: Object anchor positions for multi-point RoPE (B, N_obj, Nq*3) + mask: Optional attention mask + time_features: (B, 2) [s, s^2] for per-layer FiLM conditioning + + Returns: + If return_all_layers: (all_state_layers, all_obj_layers) + Else: (final_state, final_obj_features) + Register tokens are stripped from obj_features before returning. + """ + no_state = state is None + + # Store selected layer outputs for memory efficiency + all_states = [] if self.return_all_layers else None + all_obj_features = [] if self.return_all_layers else None + + # Rope-only: expect explicit obj_anchor_pos; state_pos only needed when state is used + if obj_anchor_pos is None: + raise ValueError( + "decoder.forward expects obj_anchor_pos for rope-only mode" + ) + if not no_state and state_pos is None: + raise ValueError("decoder.forward expects state_pos when state is provided") + + if self.center_rope_positions: + # Rotary angles of register tokens are zero, so register<->object attention would otherwise + # depend on absolute object positions. Expressing every position relative to the centroid of + # the valid objects' anchors makes the whole decoder exactly translation invariant. + Bc, Nc, Pc = obj_anchor_pos.shape + pos = obj_anchor_pos.view(Bc, Nc, Pc // 3, 3) + if mask is not None: + valid = (mask.reshape(Bc, -1)[:, -Nc:] == 0).to( + pos.dtype + ) # (B, Nobj): additive mask 0 = valid + else: + valid = torch.ones(Bc, Nc, dtype=pos.dtype, device=pos.device) + w = valid.view(Bc, Nc, 1, 1) + centroid = (pos * w).sum(dim=(1, 2)) / ( + w.sum(dim=(1, 2)) * pos.shape[2] + ).clamp_min( + 1.0 + ) # (B, 3) + obj_anchor_pos = (pos - centroid.view(Bc, 1, 1, 3)).reshape(Bc, Nc, Pc) + if not getattr(self, "_center_logged", False): + self._center_logged = True + print( + f"[Decoder] RoPE positions centred at the valid-object anchor centroid ACTIVE " + f"(|centroid| max={centroid.abs().max().item():.2f} m; registers at the centroid)", + flush=True, + ) + # Prepend register tokens to object features (RenderFormer) + if self.num_register_tokens > 0: + B = obj_features.shape[0] + register_tokens_expanded = self.register_tokens.expand( + B, -1, -1 + ) # (B, N_reg, D) + obj_features_with_reg = torch.cat( + [register_tokens_expanded, obj_features], dim=1 + ) + + # Register positions match the dimensionality of obj_anchor_pos (3 or Nq*3) + pos_dim = obj_anchor_pos.shape[-1] # Either 3 or Nq*3 + register_positions = torch.zeros( + B, + self.num_register_tokens, + pos_dim, + device=obj_anchor_pos.device, + dtype=obj_anchor_pos.dtype, + ) + obj_anchor_pos_with_reg = torch.cat( + [register_positions, obj_anchor_pos], dim=1 + ) + + # Extend mask to cover register tokens (always valid) + if mask is not None: + # mask: (B, 1, 1, Nobj) SDPA additive -> (B, 1, 1, Nreg+Nobj); + # register entries are 0.0 (attend normally) + register_mask = torch.zeros( + B, + 1, + 1, + self.num_register_tokens, + device=mask.device, + dtype=mask.dtype, + ) + mask_with_reg = torch.cat([register_mask, mask], dim=-1) + else: + mask_with_reg = None + + # Extend SE(3) bias to include register tokens (zero bias = attend freely) + if se3_bias is not None and self.num_register_tokens > 0: + N_reg = self.num_register_tokens + N_obj = se3_bias.shape[-1] # original Nobj + H = se3_bias.shape[1] + # Pad: (B, H, Nobj, Nobj) -> (B, H, Nreg+Nobj, Nreg+Nobj) + # Register tokens have zero bias (vanilla attention) + se3_bias = F.pad(se3_bias, (N_reg, 0, N_reg, 0), value=0.0) + else: + obj_features_with_reg = obj_features + obj_anchor_pos_with_reg = obj_anchor_pos + mask_with_reg = mask + + # Process through each layer with register tokens + state_current = state + obj_current = obj_features_with_reg + + # Create a dummy state placeholder for no-state mode (used in returned lists) + if no_state: + B = obj_features.shape[0] + _dummy_state = torch.zeros( + B, + 1, + self.state_dim, + device=obj_features.device, + dtype=obj_features.dtype, + ) + + # Add initial layer (layer 0) if needed + if self.return_all_layers and 0 in self.layer_indices_to_save: + all_states.append(_dummy_state if no_state else state) + all_obj_features.append( + obj_features + ) # Store original without register tokens + + # Block Attention Residuals path + if self.block_attn_res: + # Block AttnRes (Kimi 2026): replace inter-layer residuals with + # softmax attention over completed block sums + partial block sum + S = self.block_attn_res_block_size + completed_blocks = [obj_current] # b_0 = initial obj features (embedding) + partial_block = None # Running intra-block sum + layer_in_block = 0 + + for layer_idx, layer in enumerate(self.layers): + # Block boundary: commit partial sum + if layer_in_block == S: + completed_blocks.append(partial_block) + partial_block = None + layer_in_block = 0 + + # Compute input via depth attention over block representations + obj_current = self._block_attn_res_aggregate( + completed_blocks, partial_block, self.attn_res_queries[layer_idx] + ) + + # Run layer (layer uses internal residuals for sublayers) + state_current, obj_output = layer( + state_current, + obj_current, + state_pos, + obj_anchor_pos_with_reg, + mask_with_reg, + se3_bias=se3_bias, + time_features=time_features, + ) + + # Compute layer delta: f_l(h_l) = output - input + # The layer output includes internal residuals, so + # delta = obj_output - obj_current is the "function output" for AttnRes + layer_delta = obj_output - obj_current + + # Accumulate delta into partial block sum + partial_block = ( + layer_delta + if partial_block is None + else partial_block + layer_delta + ) + layer_in_block += 1 + + # Save layer outputs if needed + if ( + self.return_all_layers + and (layer_idx + 1) in self.layer_indices_to_save + ): + # For output, use the full obj_output (input + delta) + if self.num_register_tokens > 0: + obj_out_no_reg = obj_output[:, self.num_register_tokens :, :] + else: + obj_out_no_reg = obj_output + all_states.append(_dummy_state if no_state else state_current) + all_obj_features.append(obj_out_no_reg) + + # Final aggregation: attend over all blocks (including last partial) + if partial_block is not None: + completed_blocks.append(partial_block) + # Reuse the last layer's query for the final aggregation + obj_current = self._block_attn_res_aggregate( + completed_blocks, None, self.attn_res_queries[-1] + ) + if self.return_all_layers and self.block_attn_res_final_aggregate: + if self.num_layers not in self.layer_indices_to_save: + raise ValueError( + "block_attn_res_final_aggregate needs the last layer among the saved layers" + ) + all_obj_features[-1] = ( + obj_current[:, self.num_register_tokens :, :] + if self.num_register_tokens > 0 + else obj_current + ) + + else: + # Standard residual path + for layer_idx, layer in enumerate(self.layers): + state_current, obj_current = layer( + state_current, + obj_current, + state_pos, + obj_anchor_pos_with_reg, + mask_with_reg, + se3_bias=se3_bias, + time_features=time_features, + ) + + if ( + self.return_all_layers + and (layer_idx + 1) in self.layer_indices_to_save + ): + if self.num_register_tokens > 0: + obj_current_no_reg = obj_current[ + :, self.num_register_tokens :, : + ] + else: + obj_current_no_reg = obj_current + all_states.append(_dummy_state if no_state else state_current) + all_obj_features.append(obj_current_no_reg) + + # Apply final normalization + if self.return_all_layers: + if not no_state: + all_states[-1] = self.state_norm(all_states[-1]) + all_obj_features[-1] = self.obs_norm(all_obj_features[-1]) + return all_states, all_obj_features + else: + if self.num_register_tokens > 0: + obj_current_no_reg = obj_current[:, self.num_register_tokens :, :] + else: + obj_current_no_reg = obj_current + state_final = ( + self.state_norm(state_current) if not no_state else _dummy_state + ) + obj_final = self.obs_norm(obj_current_no_reg) + return state_final, obj_final + + def _block_attn_res_aggregate( + self, + completed_blocks: list, + partial_block: Optional[torch.Tensor], + w_l: nn.Parameter, + ) -> torch.Tensor: + """ + Block AttnRes aggregation (Eq. 6 from Kimi 2026). + + Computes softmax attention over completed block representations + optional + partial block sum, using per-layer pseudo-query w_l. + + phi(q, k) = exp(q^T RMSNorm(k)) + alpha_{i->l} = phi(w_l, k_i) / sum_j phi(w_l, k_j) + h_l = sum_i alpha_{i->l} * v_i + + Args: + completed_blocks: List of block representations [b_0, b_1, ...] + partial_block: Current intra-block partial sum (None if at block boundary) + w_l: Learnable pseudo-query vector (d,) + + Returns: + Aggregated input for next layer (B, N_tokens, D) + """ + if partial_block is not None: + sources = completed_blocks + [partial_block] + else: + sources = completed_blocks + + # Stack: (num_sources, B, N_tokens, D) + V = torch.stack(sources, dim=0) + # RMSNorm on keys prevents magnitude bias between complete blocks and partial sums + K = self.attn_res_key_norm(V) + # Compute attention logits: w_l^T RMSNorm(k_i) for each source + # w_l: (D,), K: (num_sources, B, N_tokens, D) -> logits: (num_sources, B, N_tokens) + logits = torch.einsum("d, n b t d -> n b t", w_l, K) + alpha = logits.softmax(dim=0) + # Weighted sum: (num_sources, B, N_tokens) x (num_sources, B, N_tokens, D) -> (B, N_tokens, D) + return torch.einsum("n b t, n b t d -> b t d", alpha, V) + + +class TransformerDecoderLayer(nn.Module): + """ + Single decoder layer of bidirectional cross-attention, following CUT3R and + RenderFormer (RMSNorm, QK-normalization, optional gated attention). + + When state_update_enabled=False: + - State is treated as fixed learnable embeddings (never updated) + - Only the Objects <- State cross-attention runs + - State self-attention and State <- Objects are skipped, saving roughly + half the computation + + When objects_self_attention=True: + - Objects self-attention runs before Objects <- State, enabling direct + object-to-object interaction + """ + + def __init__( + self, + d_model: int, + num_heads: int, + feedforward_dim: int, + dropout: float = 0.1, + activation: str = "swiglu", + use_self_attention: bool = False, + drop_path_prob: float = 0.1, + use_rope: bool = True, + norm_type: str = "rms_norm", # 'layer_norm' or 'rms_norm' + qk_norm: bool = True, + rope_set_aggregation: bool = False, # Set aggregation for multi-point RoPE + rope_dim: int = None, # Custom RoPE dimension + dim_per_coord_3d: int = None, # Dim per coordinate for single-point (3D) RoPE + dim_per_coord_multi: int = None, # Dim per coordinate for multi-point RoPE + elementwise_gate: bool = False, # Elementwise gated attention (Qwen3 style) + state_update_enabled: bool = True, # False: state is fixed learnable embeddings + objects_self_attention: bool = False, # Objects self-attention for object-object interaction + use_se3: bool = False, # SE(3) pairwise bias replaces RoPE in object self-attention + film_conditioning: bool = False, # Per-layer FiLM conditioning on step size + ): + super().__init__() + + self.use_self_attention = use_self_attention + self.use_rope = bool(use_rope) + self.state_update_enabled = state_update_enabled + self.objects_self_attention = objects_self_attention + self._objects_self_attention_logged = False # One-time log flag + + # Cross-attention with optional RoPE and QK-normalization + self.state_cross_attn = CrossAttentionWithRoPE( + d_model=d_model, + num_heads=num_heads, + dropout=dropout, + use_rope=self.use_rope, + qk_norm=qk_norm, + rope_set_aggregation=rope_set_aggregation, + rope_dim=rope_dim, + dim_per_coord_3d=dim_per_coord_3d, + dim_per_coord_multi=dim_per_coord_multi, + elementwise_gate=elementwise_gate, + ) + self.obs_cross_attn = CrossAttentionWithRoPE( + d_model=d_model, + num_heads=num_heads, + dropout=dropout, + use_rope=self.use_rope, + qk_norm=qk_norm, + rope_set_aggregation=rope_set_aggregation, + rope_dim=rope_dim, + dim_per_coord_3d=dim_per_coord_3d, + dim_per_coord_multi=dim_per_coord_multi, + elementwise_gate=elementwise_gate, + ) + + # Optional state self-attention (disabled by default) + if use_self_attention: + self.state_self_attention = SelfAttentionWithRoPE( + d_model=d_model, + num_heads=num_heads, + dropout=dropout, + use_rope=self.use_rope, + qk_norm=qk_norm, + rope_set_aggregation=rope_set_aggregation, + rope_dim=rope_dim, + dim_per_coord_3d=dim_per_coord_3d, + dim_per_coord_multi=dim_per_coord_multi, + elementwise_gate=elementwise_gate, + ) + self.state_norm0 = get_norm_layer( + norm_type, d_model + ) # Norm for state self-attention + + # Objects self-attention for object-to-object interaction. Spatial encoding + # uses multi-point RoPE over anchor positions, or an SE(3) pairwise bias + # when use_se3=True. + if objects_self_attention: + if use_se3: + from embodichain_models.dynamics.rigidformer._native.models.se3_attention import ( + SelfAttentionWithSE3Bias, + ) + + self.obj_self_attn = SelfAttentionWithSE3Bias( + d_model=d_model, + num_heads=num_heads, + dropout=dropout, + qk_norm=qk_norm, + elementwise_gate=elementwise_gate, + ) + logger.info( + f"[Decoder Init] obj_self_attn = SelfAttentionWithSE3Bias (no RoPE, pairwise bias)" + ) + else: + self.obj_self_attn = SelfAttentionWithRoPE( + d_model=d_model, + num_heads=num_heads, + dropout=dropout, + use_rope=self.use_rope, + qk_norm=qk_norm, + rope_set_aggregation=rope_set_aggregation, + rope_dim=rope_dim, + dim_per_coord_3d=dim_per_coord_3d, + dim_per_coord_multi=dim_per_coord_multi, + elementwise_gate=elementwise_gate, + ) + self.obj_norm0 = get_norm_layer( + norm_type, d_model + ) # Norm for objects self-attention + + # Feedforward networks + self.state_ffn = FeedForward(d_model, feedforward_dim, dropout, activation) + self.obs_ffn = FeedForward(d_model, feedforward_dim, dropout, activation) + + # Layer normalization (pre-norm) + self.state_norm1 = get_norm_layer(norm_type, d_model) + self.state_norm2 = get_norm_layer(norm_type, d_model) + self.obs_norm1 = get_norm_layer(norm_type, d_model) + self.obs_norm2 = get_norm_layer(norm_type, d_model) + + # Per-layer FiLM conditioning on step size + self.layer_film = None + if film_conditioning: + from embodichain_models.dynamics.rigidformer._native.models.encoder import ( + StepSizeFiLM, + ) + + self.layer_film = StepSizeFiLM(time_dim=2, feature_dim=d_model) + + # Dropout + self.dropout = nn.Dropout(dropout) + self.drop_path = ( + DropPath(drop_path_prob) + if drop_path_prob and drop_path_prob > 0.0 + else nn.Identity() + ) + + def forward( + self, + state: torch.Tensor, + obj_features: torch.Tensor, + state_pos: Optional[torch.Tensor] = None, + obj_anchor_pos: Optional[torch.Tensor] = None, + mask: Optional[torch.Tensor] = None, + se3_bias: Optional[ + torch.Tensor + ] = None, # (B, H, Nobj, Nobj) pairwise SE(3) bias + time_features: Optional[ + torch.Tensor + ] = None, # (B, 2) [s, s^2] for per-layer FiLM + ) -> tuple: + """ + Bidirectional or unidirectional cross-attention. + + When state_update_enabled=True (default): + - Full bidirectional interaction: state and objects attend to each other + - Optional state self-attention + - Both state and objects are updated + + When state_update_enabled=False: + - Unidirectional: Objects <- State only + - State remains fixed learnable embeddings + + Args: + state: State tokens (B, N_state, D) + obj_features: Per-object encoder features (B, N_obj, D) + state_pos: State 3D positions for RoPE (B, N_state, 3) + obj_anchor_pos: Object anchor positions for RoPE (B, N_obj, Nq*3) + mask: Optional attention mask + + Returns: + Updated state and obj_features + """ + no_state = state is None + + if self.state_update_enabled and not no_state: + # Full bidirectional mode + # Optional state self-attention with RoPE. The object mask does not + # apply to state tokens; pass None. + if self.use_self_attention: + state_sa = self.state_norm0(state) + state_sa = self.state_self_attention( + query=state_sa, + key_value=state_sa, + pos_q=state_pos, # Use 3D grid positions for RoPE + pos_k=state_pos, + attn_mask=None, # Object mask does not apply to state tokens + ) + state = state + self.drop_path(self.dropout(state_sa)) + + # Bidirectional cross-attention + # State <- Objects: state queries per-object features + state_normed = self.state_norm1(state) + obj_normed = self.obs_norm1(obj_features) + state2 = self.state_cross_attn( + query=state_normed, + key_value=obj_normed, + pos_q=state_pos, + pos_k=obj_anchor_pos, + attn_mask=mask, + ) + state = state + self.drop_path(self.dropout(state2)) + + # Objects <- State: per-object features query state + obj2 = self.obs_cross_attn( + query=obj_normed, + key_value=state_normed, + pos_q=obj_anchor_pos, + pos_k=state_pos, + attn_mask=None, # State keys are all valid, no need to mask + ) + obj_features = obj_features + self.drop_path(self.dropout(obj2)) + + # Per-layer FiLM conditioning on step size + if self.layer_film is not None and time_features is not None: + obj_features = self.layer_film(obj_features, time_features) + + # Feedforward networks with residual connections + state = state + self.drop_path(self.state_ffn(self.state_norm2(state))) + obj_features = obj_features + self.drop_path( + self.obs_ffn(self.obs_norm2(obj_features)) + ) + + else: + # Unidirectional / no-state mode + # no_state=True: objects self-attention + objects FFN only + # no_state=False: objects self-attention + Objects <- State + objects FFN + + # Objects self-attention runs before Objects <- State so objects + # exchange information first + if self.objects_self_attention: + # One-time runtime confirmation log + if not self._objects_self_attention_logged: + mode_str = "NO-STATE" if no_state else "UNIDIRECTIONAL" + print( + f"[Decoder Runtime] Objects self-attention active ({mode_str}): shape={obj_features.shape}, se3_bias={'yes' if se3_bias is not None else 'no'}" + ) + self._objects_self_attention_logged = True + obj_sa = self.obj_norm0(obj_features) + if se3_bias is not None: + # SE(3) path: use pairwise bias instead of RoPE + obj_sa = self.obj_self_attn( + query=obj_sa, + key_value=obj_sa, + se3_bias=se3_bias, + attn_mask=mask, + ) + else: + # ARoPE/OBB/PCA path: use RoPE with anchor positions + obj_sa = self.obj_self_attn( + query=obj_sa, + key_value=obj_sa, + pos_q=obj_anchor_pos, # Multi-point anchor positions for RoPE + pos_k=obj_anchor_pos, + attn_mask=mask, # Mask padding objects + ) + obj_features = obj_features + self.drop_path(self.dropout(obj_sa)) + + # Objects <- State cross-attention (skipped when no_state=True) + if not no_state: + state_normed = self.state_norm1(state) + obj_normed = self.obs_norm1(obj_features) + + # Objects <- State: per-object features query state (learnable context) + obj2 = self.obs_cross_attn( + query=obj_normed, + key_value=state_normed, + pos_q=obj_anchor_pos, + pos_k=state_pos, + attn_mask=None, # State keys are all valid, no need to mask + ) + obj_features = obj_features + self.drop_path(self.dropout(obj2)) + + # Per-layer FiLM conditioning on step size + # Applied after attention, before the FFN. Padding objects may receive + # non-zero features from the FiLM beta bias; the SDPA -inf mask keeps + # valid objects from attending to them. + if self.layer_film is not None and time_features is not None: + obj_features = self.layer_film(obj_features, time_features) + + # Objects FFN + obj_features = obj_features + self.drop_path( + self.obs_ffn(self.obs_norm2(obj_features)) + ) + + return state, obj_features + + +class SelfAttentionWithRoPE(nn.Module): + """ + Self-attention that applies RoPE using 3D positions. + + QK-normalization (RenderFormer) improves training stability; elementwise + gated attention (Qwen3) adds non-linearity and sparsity. + """ + + def __init__( + self, + d_model: int, + num_heads: int, + dropout: float = 0.1, + use_rope: bool = True, + qk_norm: bool = True, # QK-normalization (RenderFormer) + rope_set_aggregation: bool = False, + rope_dim: int = None, # Custom RoPE dimension (None = auto-compute) + dim_per_coord_3d: int = None, # Dim per coordinate for single-point (3D) RoPE + dim_per_coord_multi: int = None, # Dim per coordinate for multi-point RoPE + elementwise_gate: bool = False, # Elementwise gated attention (Qwen3 style) + ): + super().__init__() + assert d_model % num_heads == 0, "d_model must be divisible by num_heads" + self.d_model = d_model + self.num_heads = num_heads + self.head_dim = d_model // num_heads + self.scale = self.head_dim**-0.5 + self.use_rope = bool(use_rope) + self.qk_norm = qk_norm + self.elementwise_gate = elementwise_gate + + # Projections; with elementwise_gate, q_proj outputs 2x dims (Q and gate scores) + if self.elementwise_gate: + self.q_proj = nn.Linear(d_model, d_model * 2, bias=True) + else: + self.q_proj = nn.Linear(d_model, d_model, bias=True) + self.k_proj = nn.Linear(d_model, d_model, bias=True) + self.v_proj = nn.Linear(d_model, d_model, bias=True) + self.out_proj = nn.Linear(d_model, d_model, bias=True) + self.dropout = nn.Dropout(dropout) + self._attn_dispatch_logged = False # One-time attention dispatch log + + # QK-normalization: per-head RMSNorm (RenderFormer) + if self.qk_norm: + self.q_norm = RMSNorm(self.head_dim) + self.k_norm = RMSNorm(self.head_dim) + + # 3D rotary embedding; only the first rope_dim dims are rotated + if rope_dim is None: + # Auto-compute: multiple of 6, capped at head_dim + self.rope_dim = max(6, (self.head_dim // 6) * 6) + self.rope_dim = min(self.rope_dim, self.head_dim) + else: + # Custom rope_dim (must be a multiple of 6 for 3D RoPE) + assert ( + rope_dim % 6 == 0 + ), f"rope_dim must be multiple of 6 for 3D RoPE, got {rope_dim}" + assert ( + rope_dim <= self.head_dim + ), f"rope_dim {rope_dim} exceeds head_dim {self.head_dim}" + self.rope_dim = rope_dim + + self.rope3d = PhysFormerRoPE( + dim=self.rope_dim, + max_freq=10.0, + use_3d=True, + learnable_freqs=False, + use_set_aggregation=rope_set_aggregation, + dim_per_coord_3d=dim_per_coord_3d, # Pass to PhysFormerRoPE + dim_per_coord_multi=dim_per_coord_multi, # Pass to PhysFormerRoPE + ) + + def _reshape_to_heads(self, x: torch.Tensor) -> torch.Tensor: + B, N, D = x.shape + x = x.view(B, N, self.num_heads, self.head_dim).permute( + 0, 2, 1, 3 + ) # (B,H,N,Hd) + return x + + def _merge_heads(self, x: torch.Tensor) -> torch.Tensor: + B, H, N, Hd = x.shape + return x.permute(0, 2, 1, 3).contiguous().view(B, N, H * Hd) + + def forward( + self, + query: torch.Tensor, + key_value: torch.Tensor, + pos_q: torch.Tensor, + pos_k: torch.Tensor, + attn_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """ + Self-attention using Flash Attention, or SDPA when a mask is provided. + + Args: + query: Input tokens (B, N, D) + key_value: Same as query for self-attention (B, N, D) + pos_q: 3D positions (B, N, 3) or multi-point positions (B, N, Nq*3) + pos_k: Same as pos_q for self-attention + attn_mask: Optional additive mask; when present, attention runs + through SDPA since Flash does not support masks + + Returns: + Output tokens (B, N, D) + """ + B, N, _ = query.shape + + # Q projection with optional elementwise gate + if self.elementwise_gate: + q_gate = self.q_proj(query) # (B, N, d_model * 2) + q_states, gate_score = torch.split(q_gate, self.d_model, dim=-1) + q = q_states.view( + B, N, self.num_heads, self.head_dim + ) # (B, N, H, D) - Flash expects (B, S, H, D) + gate_score = gate_score.view( + B, N, self.num_heads, self.head_dim + ) # (B, N, H, D) + else: + q = self.q_proj(query).view( + B, N, self.num_heads, self.head_dim + ) # (B, N, H, D) + gate_score = None + + k = self.k_proj(key_value).view( + B, N, self.num_heads, self.head_dim + ) # (B, N, H, D) + v = self.v_proj(key_value).view( + B, N, self.num_heads, self.head_dim + ) # (B, N, H, D) + + # QK-normalization (RenderFormer) + if self.qk_norm: + # RMSNorm normalizes the last (per-head) dimension + q = self.q_norm(q) + k = self.k_norm(k) + + # RoPE (single-point 3D or multi-point) + if pos_q is None or pos_k is None: + raise ValueError("SelfAttentionWithRoPE expects pos_q and pos_k") + + # Single-point (3D) vs multi-point positions + if pos_q.shape[-1] == 3: + rope_q = self.rope3d.forward_3d(pos_q) # (B, N, rope_dim) + else: + rope_q = self.rope3d.forward_sa_rope( + pos_q + ) # SA-RoPE: symmetric anchor pooling + + if pos_k is pos_q: + rope_k = rope_q # identical positions -> identical angles (computed once) + elif pos_k.shape[-1] == 3: + rope_k = self.rope3d.forward_3d(pos_k) # (B, N, rope_dim) + else: + rope_k = self.rope3d.forward_sa_rope( + pos_k + ) # SA-RoPE: symmetric anchor pooling + + # The angles stay at (B, 1, N, rope_dim) and broadcast over heads inside apply_rotary_emb. cos/sin + # are elementwise, so this equals expanding to (B, H, N, rope_dim) first, at 1/num_heads of the work. + rope_q_T = rope_q.unsqueeze(1) + rope_k_T = rope_k.unsqueeze(1) if rope_k is not rope_q else rope_q_T + + # Apply rotation + q_rot = q[..., : self.rope_dim] + k_rot = k[..., : self.rope_dim] + q_pass = q[..., self.rope_dim :] + k_pass = k[..., self.rope_dim :] + + # apply_rotary_emb operates on (B, H, N, D): transpose for the rotation, + # then transpose back to the (B, N, H, D) layout Flash expects. + + q_rot = q_rot.transpose(1, 2) # (B, H, N, D) + k_rot = k_rot.transpose(1, 2) + + q_rot = apply_rotary_emb(rope_q_T, q_rot) + k_rot = apply_rotary_emb(rope_k_T, k_rot) + + # Permute back to (B, N, H, D) + q_rot = q_rot.transpose(1, 2) + k_rot = k_rot.transpose(1, 2) + + q = torch.cat([q_rot, q_pass], dim=-1) + k = torch.cat([k_rot, k_pass], dim=-1) + + # flash_attn_func does not support additive masks: dispatch to SDPA when a + # mask is provided, Flash Attention otherwise. + if attn_mask is not None: + if not self._attn_dispatch_logged: + print( + f"[SelfAttentionWithRoPE] SDPA path (mask active): Q={q.shape}, K={k.shape}, mask={attn_mask.shape}, " + f"mask_min={attn_mask.min().item():.0f}, n_masked={(attn_mask == float('-inf')).sum().item()}" + ) + self._attn_dispatch_logged = True + # SDPA expects (B, H, N, D) layout + q_sdpa = q.transpose(1, 2) # (B, H, N, D) + k_sdpa = k.transpose(1, 2) + v_sdpa = v.transpose(1, 2) + out = F.scaled_dot_product_attention( + q_sdpa, + k_sdpa, + v_sdpa, + attn_mask=attn_mask, + dropout_p=self.dropout.p if self.training else 0.0, + scale=self.scale, + ) # (B, H, N, D) + out = out.transpose(1, 2) # Back to (B, N, H, D) + else: + if not self._attn_dispatch_logged: + print(f"[SelfAttentionWithRoPE] Flash path (no mask): Q={q.shape}") + self._attn_dispatch_logged = True + # Flash Attention runs in bf16/fp16; the output is cast back to the + # original dtype immediately. + original_dtype = query.dtype + attn_dtype = ( + torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 + ) + q_half = q.to(attn_dtype) + k_half = k.to(attn_dtype) + v_half = v.to(attn_dtype) + out = flash_attn_func( + q_half, + k_half, + v_half, + dropout_p=self.dropout.p if self.training else 0.0, + softmax_scale=self.scale, + backend=getattr(self, "attention_backend", "sdpa"), + ) + out = out.to(original_dtype) + + if self.elementwise_gate: + out = out * torch.sigmoid(gate_score) + + # Merge heads: (B, N, H, D) -> (B, N, H*D) + out = out.reshape(B, N, self.num_heads * self.head_dim) + out = self.out_proj(out) + return out + + +class CrossAttentionWithRoPE(nn.Module): + """ + Cross-attention that applies RoPE to queries and keys. + + QK-normalization (RenderFormer) improves training stability; elementwise + gated attention (Qwen3) adds non-linearity and sparsity. + """ + + def __init__( + self, + d_model: int, + num_heads: int, + dropout: float = 0.1, + use_rope: bool = True, + qk_norm: bool = True, # QK-normalization (RenderFormer) + rope_set_aggregation: bool = False, + rope_dim: int = None, # Custom RoPE dimension (None = auto-compute) + dim_per_coord_3d: int = None, # Dim per coordinate for single-point (3D) RoPE + dim_per_coord_multi: int = None, # Dim per coordinate for multi-point RoPE + elementwise_gate: bool = False, # Elementwise gated attention (Qwen3 style) + ): + super().__init__() + assert d_model % num_heads == 0, "d_model must be divisible by num_heads" + self.d_model = d_model + self.num_heads = num_heads + self.head_dim = d_model // num_heads + self.scale = self.head_dim**-0.5 + self.use_rope = bool(use_rope) + self.qk_norm = qk_norm + self.elementwise_gate = elementwise_gate + + # Projections; with elementwise_gate, q_proj outputs 2x dims (Q and gate scores) + if self.elementwise_gate: + self.q_proj = nn.Linear(d_model, d_model * 2, bias=True) + else: + self.q_proj = nn.Linear(d_model, d_model, bias=True) + self.k_proj = nn.Linear(d_model, d_model, bias=True) + self.v_proj = nn.Linear(d_model, d_model, bias=True) + self.out_proj = nn.Linear(d_model, d_model, bias=True) + self.dropout = nn.Dropout(dropout) + self._attn_dispatch_logged = False # One-time attention dispatch log + + # QK-normalization: per-head RMSNorm (RenderFormer) + if self.qk_norm: + self.q_norm = RMSNorm(self.head_dim) + self.k_norm = RMSNorm(self.head_dim) + + # 3D rotary embedding; only the first rope_dim dims are rotated + if rope_dim is None: + # Auto-compute: multiple of 6, capped at head_dim + self.rope_dim = max(6, (self.head_dim // 6) * 6) + self.rope_dim = min(self.rope_dim, self.head_dim) + else: + # Custom rope_dim (must be a multiple of 6 for 3D RoPE) + assert ( + rope_dim % 6 == 0 + ), f"rope_dim must be multiple of 6 for 3D RoPE, got {rope_dim}" + assert ( + rope_dim <= self.head_dim + ), f"rope_dim {rope_dim} exceeds head_dim {self.head_dim}" + self.rope_dim = rope_dim + + self.rope3d = PhysFormerRoPE( + dim=self.rope_dim, + max_freq=10.0, + use_3d=True, + learnable_freqs=False, + use_set_aggregation=rope_set_aggregation, + dim_per_coord_3d=dim_per_coord_3d, # Pass to PhysFormerRoPE + dim_per_coord_multi=dim_per_coord_multi, # Pass to PhysFormerRoPE + ) + + def _reshape_to_heads(self, x: torch.Tensor) -> torch.Tensor: + B, N, D = x.shape + x = x.view(B, N, self.num_heads, self.head_dim).permute( + 0, 2, 1, 3 + ) # (B,H,N,Hd) + return x + + def _merge_heads(self, x: torch.Tensor) -> torch.Tensor: + B, H, N, Hd = x.shape + return x.permute(0, 2, 1, 3).contiguous().view(B, N, H * Hd) + + def forward( + self, + query: torch.Tensor, + key_value: torch.Tensor, + pos_q: torch.Tensor, + pos_k: torch.Tensor, + attn_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """ + Cross-attention using Flash Attention v2. + """ + B, Nq, _ = query.shape + _, Nk, _ = key_value.shape + + # Projections => (B, N, H, D) for Flash + if self.elementwise_gate: + q_gate = self.q_proj(query) + q_states, gate_score = torch.split(q_gate, self.d_model, dim=-1) + q = q_states.view(B, Nq, self.num_heads, self.head_dim) + gate_score = gate_score.view(B, Nq, self.num_heads, self.head_dim) + else: + q = self.q_proj(query).view(B, Nq, self.num_heads, self.head_dim) + gate_score = None + + k = self.k_proj(key_value).view(B, Nk, self.num_heads, self.head_dim) + v = self.v_proj(key_value).view(B, Nk, self.num_heads, self.head_dim) + + # QK Norm + if self.qk_norm: + q = self.q_norm(q) + k = self.k_norm(k) + + # RoPE + if pos_q is None or pos_k is None: + raise ValueError("CrossAttentionWithRoPE expects pos_q and pos_k") + + # Single-point (3D) vs multi-point positions + if pos_q.shape[-1] == 3: + rope_q = self.rope3d.forward_3d(pos_q) + else: + rope_q = self.rope3d.forward_sa_rope( + pos_q + ) # SA-RoPE: symmetric anchor pooling + + if pos_k is pos_q: + rope_k = rope_q # identical positions -> identical angles (computed once) + elif pos_k.shape[-1] == 3: + rope_k = self.rope3d.forward_3d(pos_k) + else: + rope_k = self.rope3d.forward_sa_rope( + pos_k + ) # SA-RoPE: symmetric anchor pooling + + # (B, 1, N, rope_dim) angles broadcast over heads (see SelfAttentionWithRoPE) + rope_q_T = rope_q.unsqueeze(1) + rope_k_T = rope_k.unsqueeze(1) if rope_k is not rope_q else rope_q_T + + # apply_rotary_emb operates on (B, H, N, D): transpose for the rotation and back + q_rot = q[..., : self.rope_dim].transpose(1, 2) + k_rot = k[..., : self.rope_dim].transpose(1, 2) + + q_rot = apply_rotary_emb(rope_q_T, q_rot) + k_rot = apply_rotary_emb(rope_k_T, k_rot) + + q_rot = q_rot.transpose(1, 2) + k_rot = k_rot.transpose(1, 2) + + q_pass = q[..., self.rope_dim :] + k_pass = k[..., self.rope_dim :] + q = torch.cat([q_rot, q_pass], dim=-1) + k = torch.cat([k_rot, k_pass], dim=-1) + + # flash_attn_func does not support additive masks: dispatch to SDPA when a + # mask is provided, Flash Attention otherwise. + if attn_mask is not None: + if not self._attn_dispatch_logged: + n_inf = (attn_mask == float("-inf")).sum().item() + n_total = attn_mask.numel() + print( + f"[CrossAttentionWithRoPE] SDPA path (mask active): Q={q.shape}, K={k.shape}, " + f"mask={attn_mask.shape}, masked={n_inf}/{n_total} ({100*n_inf/max(n_total,1):.1f}%)" + ) + self._attn_dispatch_logged = True + # SDPA expects (B, H, N, D) layout + q_sdpa = q.transpose(1, 2) # (B, H, Nq, D) + k_sdpa = k.transpose(1, 2) # (B, H, Nk, D) + v_sdpa = v.transpose(1, 2) # (B, H, Nk, D) + out = F.scaled_dot_product_attention( + q_sdpa, + k_sdpa, + v_sdpa, + attn_mask=attn_mask, + dropout_p=self.dropout.p if self.training else 0.0, + scale=self.scale, + ) # (B, H, Nq, D) + out = out.transpose(1, 2) # Back to (B, Nq, H, D) + else: + if not self._attn_dispatch_logged: + print( + f"[CrossAttentionWithRoPE] Flash path (no mask): Q={q.shape}, K={k.shape}" + ) + self._attn_dispatch_logged = True + # Flash Attention runs in bf16/fp16; the output is cast back to the + # original dtype immediately. + original_dtype = query.dtype + attn_dtype = ( + torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 + ) + q_half = q.to(attn_dtype) + k_half = k.to(attn_dtype) + v_half = v.to(attn_dtype) + out = flash_attn_func( + q_half, + k_half, + v_half, + dropout_p=self.dropout.p if self.training else 0.0, + softmax_scale=self.scale, + backend=getattr(self, "attention_backend", "sdpa"), + ) + out = out.to(original_dtype) + + if self.elementwise_gate: + out = out * torch.sigmoid(gate_score) + + out = out.reshape(B, Nq, self.num_heads * self.head_dim) + out = self.out_proj(out) + return out + + +class FeedForward(nn.Module): + """Feedforward network with activation and dropout. + Supports GELU/RELU and SwiGLU (default). For SwiGLU, hidden_dim is the inner dimension. + """ + + def __init__( + self, + d_model: int, + hidden_dim: int, + dropout: float = 0.1, + activation: str = "swiglu", + ): + super().__init__() + self.activation_name = activation + self.dropout = nn.Dropout(dropout) + if activation == "swiglu": + self.fc1 = nn.Linear(d_model, hidden_dim) + self.fc1_gate = nn.Linear(d_model, hidden_dim) + self.fc2 = nn.Linear(hidden_dim, d_model) + else: + self.fc1 = nn.Linear(d_model, hidden_dim) + self.fc2 = nn.Linear(hidden_dim, d_model) + if activation == "gelu": + self.act = nn.GELU() + elif activation == "relu": + self.act = nn.ReLU() + elif activation == "silu": + self.act = nn.SiLU() + else: + raise ValueError(f"Unknown activation: {activation}") + + def forward(self, x): + if self.activation_name == "swiglu": + a = self.fc1(x) + b = torch.nn.functional.silu(self.fc1_gate(x)) + y = a * b + y = self.dropout(y) + y = self.fc2(y) + y = self.dropout(y) + return y + else: + y = self.fc1(x) + y = self.act(y) + y = self.dropout(y) + y = self.fc2(y) + y = self.dropout(y) + return y + + +class DropPath(nn.Module): + """Stochastic Depth per sample (DropPath).""" + + def __init__(self, drop_prob: float = 0.0): + super().__init__() + self.drop_prob = drop_prob + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.drop_prob == 0.0 or not self.training: + return x + keep_prob = 1 - self.drop_prob + shape = (x.shape[0],) + (1,) * (x.ndim - 1) + random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device) + random_tensor.floor_() + return x.div(keep_prob) * random_tensor + + +if __name__ == "__main__": + """Test the decoder module.""" + + print("Testing BidirectionalDecoder...") + decoder = BidirectionalDecoder( + state_dim=512, num_layers=4, num_heads=8, feedforward_dim=2048 + ) + + # Create dummy inputs + batch_size = 2 + state = torch.randn(batch_size, 256, 512) # 256 state tokens + observation = torch.randn(batch_size, 10, 512) # 10 observation tokens + + # Forward pass + output = decoder(state, observation) + print(f"Output shape: {output.shape}") # Should be (2, 10, 512) + + print("\nDecoder test successful!") diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/encoder.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/encoder.py new file mode 100644 index 0000000000..d25e3403e6 --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/encoder.py @@ -0,0 +1,937 @@ +# 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. + +""" +PointNet-based encoder for rigid body representation. + +Contains the per-object point cloud encoder (PointNet backbone with spatial +and feature transforms), multi-strategy global pooling, and step-size FiLM +conditioning. +""" + +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.nn.functional as F +from typing import Tuple, Optional + + +class STN3d(nn.Module): + """Spatial Transformer Network for 3D point alignment.""" + + def __init__(self, input_dim: int = 3): + super().__init__() + self.input_dim = input_dim + + # Convolutional layers + self.conv1 = nn.Conv1d(input_dim, 64, 1, bias=False) + self.conv2 = nn.Conv1d(64, 128, 1, bias=False) + self.conv3 = nn.Conv1d(128, 1024, 1, bias=False) + + # Batch normalization + self.bn1 = nn.BatchNorm1d(64) + self.bn2 = nn.BatchNorm1d(128) + self.bn3 = nn.BatchNorm1d(1024) + + # Fully connected layers + self.fc1 = nn.Linear(1024, 512, bias=False) + self.fc2 = nn.Linear(512, 256, bias=False) + self.fc3 = nn.Linear(256, input_dim * input_dim, bias=False) + + self.bn4 = nn.BatchNorm1d(512) + self.bn5 = nn.BatchNorm1d(256) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Args: + x: Input point cloud (B, input_dim, N) + Returns: + Transformation matrix (B, input_dim, input_dim) + """ + batch_size = x.size(0) + + # Extract features (use SiLU for smoother nonlinearity) + x = F.silu(self.bn1(self.conv1(x))) + x = F.silu(self.bn2(self.conv2(x))) + x = F.silu(self.bn3(self.conv3(x))) + + # Global feature via max pooling + x = torch.max(x, 2, keepdim=False)[0] + + # Transform to matrix + x = F.silu(self.bn4(self.fc1(x))) + x = F.silu(self.bn5(self.fc2(x))) + x = self.fc3(x) + + # Add identity matrix for stability + identity = torch.eye(self.input_dim, dtype=x.dtype, device=x.device) + identity = identity.view(1, -1).repeat(batch_size, 1) + x = x + identity + x = x.view(batch_size, self.input_dim, self.input_dim) + + return x + + +def grouped_masked_batch_norm( + bn: nn.Module, x: torch.Tensor, mask: torch.Tensor, groups: int +) -> torch.Tensor: + """Training-mode BatchNorm over x (G*M, C, N) with separate statistics for each of G groups. + + Group g is the contiguous block of rows [g*M, (g+1)*M); its statistics use only the masked-valid + positions (mask (G*M, N), True = valid). This equals G separate calls of the packed masked BatchNorm + (RigidBodyEncoder._bn) on the G blocks, including the running-statistics updates (applied in the + closed form of G sequential momentum updates), without packing: fixed shapes and no host syncs. Under + torch.distributed the per-group sums are all-reduced with an autograd-aware collective, so both the + statistics and their gradients are those of SyncBatchNorm applied per group. Padded positions of the + output are zero, as in the packed path. + Sums run in fp32 with fp64 accumulation (no fp64 copy of x); results match the packed path up to + floating-point rounding (different reduction order).""" + GM, C, N = x.shape + if GM % groups != 0: + raise ValueError( + f"grouped_masked_batch_norm: {GM} rows are not divisible into {groups} groups" + ) + M = GM // groups + if mask.shape != (GM, N): + raise ValueError( + f"grouped_masked_batch_norm: mask must be {(GM, N)}; got {tuple(mask.shape)}" + ) + m = mask.view(groups, M, 1, N) + xf = x.float().view(groups, M, C, N) + xm = torch.where(m, xf, torch.zeros((), dtype=xf.dtype, device=x.device)) + s1 = xm.sum(dim=(1, 3), dtype=torch.float64) # (G, C) + s2 = (xm * xm).sum(dim=(1, 3), dtype=torch.float64) # (G, C) + cnt = m.sum(dim=(1, 3), dtype=torch.float64) # (G, 1) + if torch.distributed.is_available() and torch.distributed.is_initialized(): + from torch.distributed.nn.functional import all_reduce as _ar_autograd + + packed = _ar_autograd( + torch.cat([s1, s2, cnt], dim=1) + ) # (G, 2C+1), gradients all-reduced too + s1, s2, cnt = packed[:, :C], packed[:, C : 2 * C], packed[:, 2 * C :] + cnt_c = cnt.clamp_min(1.0) + mean = s1 / cnt_c # (G, C) + var = (s2 / cnt_c - mean * mean).clamp_min(0.0) # biased (normalisation), (G, C) + if bn.track_running_stats and bn.running_mean is not None: + with torch.no_grad(): + mom = bn.momentum if bn.momentum is not None else 0.1 + unbiased = var * (cnt / (cnt - 1.0).clamp_min(1.0)) + # G sequential updates r <- (1-mom) r + mom * stat_g, in closed form + w = mom * (1.0 - mom) ** torch.arange( + groups - 1, -1, -1, dtype=torch.float64, device=x.device + ) + decay = (1.0 - mom) ** groups + rm = bn.running_mean.double() * decay + ( + w.view(groups, 1) * mean.detach() + ).sum(0) + rv = bn.running_var.double() * decay + ( + w.view(groups, 1) * unbiased.detach() + ).sum(0) + bn.running_mean.copy_(rm.to(bn.running_mean.dtype)) + bn.running_var.copy_(rv.to(bn.running_var.dtype)) + bn.num_batches_tracked.add_(groups) + inv = torch.rsqrt(var.float() + bn.eps).view(groups, 1, C, 1) + y = (xf - mean.float().view(groups, 1, C, 1)) * inv + if bn.affine: + y = y * bn.weight.float().view(1, 1, C, 1) + bn.bias.float().view(1, 1, C, 1) + y = torch.where(m, y, torch.zeros((), dtype=y.dtype, device=y.device)) + return y.view(GM, C, N).to(x.dtype) + + +def grouped_rows_batch_norm( + bn: nn.Module, x: torch.Tensor, group_sizes: list +) -> torch.Tensor: + """Training-mode BatchNorm over packed rows x (P, C) with separate statistics for G contiguous row blocks of the + given sizes: the packed counterpart of grouped_masked_batch_norm (same fp64 accumulation, autograd-aware + all-reduce of the per-group sums under torch.distributed, closed-form running-statistics updates of G + sequential momentum steps).""" + G = len(group_sizes) + P, C = x.shape + xf = x.float() + blocks = xf.split(group_sizes, dim=0) + s1 = torch.stack([b.sum(dim=0, dtype=torch.float64) for b in blocks]) # (G, C) + s2 = torch.stack( + [(b * b).sum(dim=0, dtype=torch.float64) for b in blocks] + ) # (G, C) + cnt = torch.tensor(group_sizes, dtype=torch.float64, device=x.device).view(G, 1) + if torch.distributed.is_available() and torch.distributed.is_initialized(): + from torch.distributed.nn.functional import all_reduce as _ar_autograd + + packed = _ar_autograd(torch.cat([s1, s2, cnt], dim=1)) + s1, s2, cnt = packed[:, :C], packed[:, C : 2 * C], packed[:, 2 * C :] + cnt_c = cnt.clamp_min(1.0) + mean = s1 / cnt_c + var = (s2 / cnt_c - mean * mean).clamp_min(0.0) + if bn.track_running_stats and bn.running_mean is not None: + with torch.no_grad(): + mom = bn.momentum if bn.momentum is not None else 0.1 + unbiased = var * (cnt / (cnt - 1.0).clamp_min(1.0)) + w = mom * (1.0 - mom) ** torch.arange( + G - 1, -1, -1, dtype=torch.float64, device=x.device + ) + decay = (1.0 - mom) ** G + rm = bn.running_mean.double() * decay + (w.view(G, 1) * mean.detach()).sum( + 0 + ) + rv = bn.running_var.double() * decay + ( + w.view(G, 1) * unbiased.detach() + ).sum(0) + bn.running_mean.copy_(rm.to(bn.running_mean.dtype)) + bn.running_var.copy_(rv.to(bn.running_var.dtype)) + bn.num_batches_tracked.add_(G) + # Per-block broadcast of the group statistics (the same elementwise ops as indexing them per row, whose + # backward is a heavily duplicated index_put) + inv, mean_f = torch.rsqrt(var.float() + bn.eps), mean.float() # (G, C) + y = torch.cat( + [(b - mean_f[g]) * inv[g] for g, b in enumerate(blocks)], dim=0 + ) # (P, C) + if bn.affine: + y = y * bn.weight.float() + bn.bias.float() + return y.to(x.dtype) + + +def hierarchical_subset_size( + n_valid: torch.Tensor, level: int, small: str +) -> torch.Tensor: + """Vertices pooled by subsampled level `level` (0-based) of an object with n_valid valid vertices: the first + n_valid // 2^(level+1) of its FPS ordering (50/25/12.5% for levels 0/1/2). small = 'all_below_8': a level + with fewer than 8 vertices uses all of them; 'fraction': the stated fraction for every object, at least 1 vertex. + """ + k = n_valid // (2 ** (level + 1)) + if small == "all_below_8": + return torch.where(k < 8, n_valid, k) + if small == "fraction": + return k.clamp_min(1) + raise ValueError( + f"subset_small_objects must be 'all_below_8' or 'fraction', got {small!r}" + ) + + +class EncoderPack: + """Packed layout of the valid vertices of an encoder call (rows = objects, N = vertex slots): flat indices of the + valid vertices, of every hierarchical level's subset, and the per-group row counts for grouped BatchNorm. + The masks do not change over a batch or a rollout, so the pack is built once (a few host syncs) and reused. + """ + + def __init__( + self, + valid_mask: torch.Tensor, + vtx_rank: Optional[torch.Tensor], + num_levels: int, + subsets: bool, + groups: int = 1, + small: str = "all_below_8", + ): + R, N = valid_mask.shape + if R % groups != 0: + raise ValueError( + f"EncoderPack: {R} rows are not divisible into {groups} groups" + ) + self.R, self.N, self.groups = R, N, groups + mask = valid_mask.bool() + self.idx = mask.reshape(-1).nonzero(as_tuple=True)[0] # (P,) sorted + rows_per_group = R // groups + g_of = torch.div( + torch.div(self.idx, N, rounding_mode="floor"), + rows_per_group, + rounding_mode="floor", + ) + self.group_sizes = ( + torch.bincount(g_of, minlength=groups).tolist() + if groups > 1 + else [int(self.idx.numel())] + ) + self.level_idx, self.level_rows, self.level_masks, self.level_group_sizes = ( + [], + [], + [], + [], + ) + n_valid = mask.sum(dim=1) + # Segment layout for pooling over the packed rows (packed_pooling='segment'): row r's valid vertices are + # the contiguous packed rows of segment r (idx is sorted), lengths = their counts, seg = segment per row. + self.lengths, self.seg = n_valid, torch.div(self.idx, N, rounding_mode="floor") + self.level_lengths, self.level_seg = [], [] + for level in range(num_levels): + if subsets: + if vtx_rank is None: + raise ValueError("EncoderPack: hierarchical subsets need vtx_rank") + k = hierarchical_subset_size(n_valid, level, small) + lmask = mask & (vtx_rank < k.unsqueeze(1)) + else: + lmask = mask + li = lmask.reshape(-1).nonzero(as_tuple=True)[0] # sorted, subset of idx + # positions of the level's vertices inside the packed rows (idx is sorted, li a subset of it) + self.level_rows.append(torch.searchsorted(self.idx, li)) + self.level_idx.append(li) + self.level_masks.append(lmask) + lg = torch.div( + torch.div(li, N, rounding_mode="floor"), + rows_per_group, + rounding_mode="floor", + ) + self.level_group_sizes.append( + torch.bincount(lg, minlength=groups).tolist() + if groups > 1 + else [int(li.numel())] + ) + self.level_lengths.append(lmask.sum(dim=1)) + self.level_seg.append(torch.div(li, N, rounding_mode="floor")) + self.mask = mask + + +class RigidBodyEncoder(nn.Module): + """ + PointNet encoder for rigid body representation. + + Follows the original PointNet architecture (STN, feature transform, + shared MLPs), with two pooling variants beyond plain max pooling: + - Enhanced pooling: max + mean + std for richer global features + - Hierarchical pooling: multi-scale aggregation (PointNet++ style) + """ + + def __init__( + self, + input_dim: int = 9, # 3 for dist_vec + 3 for velocity + 3 for relative_pos + hidden_dims: list = None, # Will use original PointNet dims by default + output_dim: int = 1024, # Final fused output dim used by downstream modules (MultiBodyEncoder maps to this) + backbone_dim: int = None, # Internal PointNet feature width before pooling/ fusion; defaults to output_dim + use_stn: bool = True, # Enable STN like original + use_batch_norm: bool = True, + feature_transform: bool = True, # Enable feature transform like original + feature_transform_dim: int = 64, # K=64 in original paper + pooling_type: str = "enhanced", # 'max' | 'enhanced' (max+mean+std) | 'hierarchical' (PointNet++) + num_hierarchical_levels: int = 2, + masked_pooling: bool = False, # Exact set-function pooling / BatchNorm over valid vertices only + freeze_bn_stats: bool = False, # Training uses the stored running statistics (fixed per-channel affine); no re-estimation + height_input: bool = False, # Extra input channel: vertex height above the ground plane (absolute z) + hierarchical_subsets: bool = False, # Levels pool the first 50/25/12.5/... percent of an order-independent FPS ordering + packed_pooling: str = "scatter", # Packed path only: 'scatter' (pool the padded layout) | 'segment' (pool the packed rows) + subset_small_objects: str = "all_below_8", # Subset levels below 8 vertices: 'all_below_8' (all vertices) | 'fraction' + ): + super().__init__() + + self.input_dim = input_dim + self.output_dim = output_dim + self.backbone_dim = backbone_dim if backbone_dim is not None else output_dim + self.use_stn = use_stn + self.feature_transform = feature_transform + self.feature_transform_dim = feature_transform_dim + self.pooling_type = pooling_type + self.num_hierarchical_levels = num_hierarchical_levels + # Set-function mode: pooling and BatchNorm statistics ignore padded vertices, + # and every hierarchical level pools over the full valid set. + self.masked_pooling = masked_pooling + self.freeze_bn_stats = bool(freeze_bn_stats) + self.height_input = bool(height_input) + self.hierarchical_subsets = bool(hierarchical_subsets) + if packed_pooling not in ("scatter", "segment"): + raise ValueError( + f"packed_pooling must be 'scatter' or 'segment', got {packed_pooling!r}" + ) + self.packed_pooling = packed_pooling + if subset_small_objects not in ("all_below_8", "fraction"): + raise ValueError( + f"subset_small_objects must be 'all_below_8' or 'fraction', got {subset_small_objects!r}" + ) + self.subset_small_objects = subset_small_objects + + # Use original PointNet architecture dimensions if not specified: + # conv(64) -> conv(64) -> [feature transform] -> conv(64) -> conv(128) -> conv(1024) + if hidden_dims is None: + hidden_dims = [64, 128, 1024] + + # Spatial Transformer Network (for input alignment) + if use_stn: + self.stn = STN3d(input_dim) + + # First set of convolutions (before feature transform) + self.conv1 = nn.Conv1d(input_dim, 64, 1, bias=not use_batch_norm) + self.bn1 = nn.BatchNorm1d(64) if use_batch_norm else nn.Identity() + + self.conv2 = nn.Conv1d(64, 64, 1, bias=not use_batch_norm) + self.bn2 = nn.BatchNorm1d(64) if use_batch_norm else nn.Identity() + + # Feature transform network at 64-dimensional features + if feature_transform: + self.feature_stn = STN3d(feature_transform_dim) + + # Second set of convolutions (after feature transform) + self.conv3 = nn.Conv1d(64, 64, 1, bias=not use_batch_norm) + self.bn3 = nn.BatchNorm1d(64) if use_batch_norm else nn.Identity() + + self.conv4 = nn.Conv1d(64, 128, 1, bias=not use_batch_norm) + self.bn4 = nn.BatchNorm1d(128) if use_batch_norm else nn.Identity() + + self.conv5 = nn.Conv1d(128, self.backbone_dim, 1, bias=not use_batch_norm) + self.bn5 = ( + nn.BatchNorm1d(self.backbone_dim) if use_batch_norm else nn.Identity() + ) + + # Enhanced pooling: project (max+mean+std) back to backbone_dim + if pooling_type == "enhanced": + self.pooling_proj = nn.Sequential( + nn.Linear(self.backbone_dim * 3, self.backbone_dim * 2), + nn.ReLU(), + nn.Linear(self.backbone_dim * 2, self.backbone_dim), + ) + + # Hierarchical pooling: multi-scale feature aggregation (PointNet++ inspired) + # Designed to handle variable vertex counts robustly + if pooling_type == "hierarchical": + self.hierarchical_layers = nn.ModuleList() + for i in range(num_hierarchical_levels): + # Each level: Conv1d for local feature extraction + self.hierarchical_layers.append( + nn.Sequential( + nn.Conv1d(self.backbone_dim, self.backbone_dim, 1), + ( + nn.BatchNorm1d(self.backbone_dim) + if use_batch_norm + else nn.Identity() + ), + nn.ReLU(), + ) + ) + # Project concatenated multi-scale features + # Input: global (backbone_dim*2) + level1 (backbone_dim*2) + ... + levelN (backbone_dim*2) + # Total: backbone_dim * 2 * (1 + num_hierarchical_levels) + proj_input_dim = self.backbone_dim * 2 * (1 + num_hierarchical_levels) + self.hierarchical_proj = nn.Sequential( + nn.Linear(proj_input_dim, self.backbone_dim * 2), + nn.ReLU(), + nn.Dropout(0.1), # Dropout for regularization + nn.Linear(self.backbone_dim * 2, self.backbone_dim), + ) + + def _bn( + self, + bn: nn.Module, + x: torch.Tensor, + mask: Optional[torch.Tensor], + mask_idx: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + groups: int = 1, + ) -> torch.Tensor: + """BatchNorm over (B, C, N) whose training statistics come from valid vertices only. + + The valid vertices are packed into one (1, C, P) tensor and normalised by the stock + (Sync)BatchNorm module, so batch statistics, running statistics and the gradient + through the statistics are exactly those of a BatchNorm applied to the un-padded + set. Padded rows of the output are zero; the masked pooling ignores them. + Evaluation uses the running statistics, a per-channel affine map that is already + independent of padding. + + mask_idx: optional mask.nonzero(as_tuple=True), shared by calls that use the same mask. Indexing + with it is what boolean-mask indexing does internally, so results are identical; it only avoids + recomputing the nonzero (a host sync) for the pack and again for the scatter. + + groups > 1: the rows are G contiguous blocks (e.g. G time steps batched together) whose statistics + must stay separate; grouped_masked_batch_norm computes them per block without packing. + """ + is_bn = isinstance(bn, torch.nn.modules.batchnorm._BatchNorm) + if self.freeze_bn_stats and self.training and is_bn: + # Fine-tuning a converged model: keep the stored statistics (the downstream weights + # were trained against them), so training and evaluation normalise identically and + # padded rows cannot influence anything. The affine parameters remain trainable. + if not getattr(self, "_frozen_bn_logged", False): + self._frozen_bn_logged = True + print( + f"[Encoder] BatchNorm statistics FROZEN in training (module={type(bn).__name__}; stored running stats, affine trainable)", + flush=True, + ) + return torch.nn.functional.batch_norm( + x, + bn.running_mean, + bn.running_var, + bn.weight, + bn.bias, + False, + 0.0, + bn.eps, + ) + if not (self.masked_pooling and mask is not None and self.training and is_bn): + return bn(x) + if not getattr(self, "_packed_bn_logged", False): + self._packed_bn_logged = True + print( + f"[Encoder] packed masked BatchNorm ACTIVE: module={type(bn).__name__}, valid={int(mask.sum())}/{mask.numel()} vertices", + flush=True, + ) + if groups > 1: + return grouped_masked_batch_norm(bn, x, mask, groups) + if mask_idx is None: + mask_idx = mask.nonzero(as_tuple=True) + xt = x.transpose(1, 2) # (B, N, C) + packed = xt[mask_idx] # (P, C), same rows and order as xt[mask] + y = bn(packed.t().unsqueeze(0)).squeeze(0).t() # (P, C) + out = xt.new_zeros(xt.shape) + out[mask_idx] = y.to(out.dtype) + return out.transpose(1, 2) + + @staticmethod + def _masked_max(f: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: + neg = torch.finfo(f.dtype).min + out = f.masked_fill(~mask.unsqueeze(1), neg).max(dim=2)[0] + return torch.where(mask.any(dim=1, keepdim=True), out, torch.zeros_like(out)) + + @staticmethod + def _masked_mean(f: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: + # Accumulate in fp32 so the mean does not depend on the padded length under autocast. + m = mask.unsqueeze(1).float() + return ((f.float() * m).sum(dim=2) / m.sum(dim=2).clamp_min(1.0)).to(f.dtype) + + @staticmethod + def _segment_max_mean(h: torch.Tensor, lengths: torch.Tensor, seg: torch.Tensor): + """_masked_max and _masked_mean of the packed rows h (P, C) without the padded layout: segment r is the + contiguous block of lengths[r] rows (seg = segment of every row). The max has the same value and passes its + gradient to the first maximal row of the segment, which is the index torch.max returns on the padded layout + (the packed order is the slot order); the mean accumulates in fp32. Empty segments give zeros. + """ + P = h.shape[0] + if P >= 2**24: + raise ValueError( + f"_segment_max_mean: {P} rows exceed the exact fp32 row positions" + ) + with torch.no_grad(): + mx = torch.segment_reduce( + h, "max", lengths=lengths, axis=0, unsafe=True + ) # (R, C) + pos = torch.arange(P, device=h.device, dtype=torch.float32).unsqueeze(1) + first = torch.segment_reduce( + torch.where(h == mx.index_select(0, seg), pos, float(P)), + "min", + lengths=lengths, + axis=0, + unsafe=True, + ) + first = first.clamp_max(P - 1).long() # (R, C) + mx = torch.where( + (lengths > 0).unsqueeze(1), + h.gather(0, first), + torch.zeros((), dtype=h.dtype, device=h.device), + ) + mean = ( + torch.segment_reduce(h.float(), "sum", lengths=lengths, axis=0, unsafe=True) + / lengths.clamp_min(1).unsqueeze(1).float() + ).to(h.dtype) + return mx, mean + + def _bn_rows( + self, bn: nn.Module, x: torch.Tensor, group_sizes: list, groups: int + ) -> torch.Tensor: + """BatchNorm over packed valid rows x (P, C): the statistics of the masked path without packing/scattering.""" + if not isinstance(bn, torch.nn.modules.batchnorm._BatchNorm): + return bn(x) + if self.freeze_bn_stats and self.training: + return torch.nn.functional.batch_norm( + x, + bn.running_mean, + bn.running_var, + bn.weight, + bn.bias, + False, + 0.0, + bn.eps, + ) + if self.training and groups > 1: + return grouped_rows_batch_norm(bn, x, group_sizes) + return bn(x) + + @staticmethod + def _rows_linear(conv: nn.Conv1d, x: torch.Tensor) -> torch.Tensor: + return F.linear(x, conv.weight.squeeze(-1), conv.bias) + + def _forward_packed(self, x: torch.Tensor, pack: "EncoderPack", bn_groups: int): + """Encoder over the valid vertices only (model.encoder.packed): every pointwise layer runs on the packed rows, + each hierarchical level only on its own vertex subset, and the features are scattered back (zeros at + padding) for the unchanged masked pooling and for anchor-vertex pooling. Equal to the padded path at every + valid vertex up to floating-point rounding (different GEMM shapes / reduction order). + """ + if ( + not (self.masked_pooling and self.pooling_type == "hierarchical") + or self.use_stn + or self.feature_transform + ): + raise ValueError( + "the packed encoder path needs masked hierarchical pooling without STN / feature transform" + ) + R, N, Cin = x.shape + if (R, N) != (pack.R, pack.N) or bn_groups != pack.groups: + raise ValueError( + f"EncoderPack built for {(pack.R, pack.N, pack.groups)}, called with {(R, N, bn_groups)}" + ) + if not getattr(self, "_packed_path_logged", False): + self._packed_path_logged = True + print( + f"[Encoder] packed vertex path ACTIVE: {int(pack.idx.numel())}/{R * N} vertex slots valid, " + f"levels {[int(i.numel()) for i in pack.level_idx]}", + flush=True, + ) + gs = pack.group_sizes + h = x.reshape(R * N, Cin).index_select(0, pack.idx) + h = F.silu( + self._bn_rows(self.bn1, self._rows_linear(self.conv1, h), gs, bn_groups) + ) + h = F.silu( + self._bn_rows(self.bn2, self._rows_linear(self.conv2, h), gs, bn_groups) + ) + h = F.silu( + self._bn_rows(self.bn3, self._rows_linear(self.conv3, h), gs, bn_groups) + ) + h = F.silu( + self._bn_rows(self.bn4, self._rows_linear(self.conv4, h), gs, bn_groups) + ) + h = self._bn_rows( + self.bn5, self._rows_linear(self.conv5, h), gs, bn_groups + ) # (P, C) + C = h.shape[1] + # Per-vertex features in the padded layout (zeros at padding) for anchor-vertex pooling + per_point_features = ( + h.new_zeros(R * N, C) + .index_copy(0, pack.idx, h) + .view(R, N, C) + .transpose(1, 2) + ) + segment = self.packed_pooling == "segment" + if segment and not getattr(self, "_segment_pool_logged", False): + self._segment_pool_logged = True + print( + "[Encoder] packed pooling: segment max/mean over the packed rows (no padded level tensors)", + flush=True, + ) + if segment: + feats = [ + torch.cat(self._segment_max_mean(h, pack.lengths, pack.seg), dim=1) + ] + else: + mask = pack.mask + feats = [ + torch.cat( + [ + self._masked_max(per_point_features, mask), + self._masked_mean(per_point_features, mask), + ], + dim=1, + ) + ] + for level, layer in enumerate(self.hierarchical_layers): + hl = self._rows_linear(layer[0], h.index_select(0, pack.level_rows[level])) + hl = layer[2]( + self._bn_rows(layer[1], hl, pack.level_group_sizes[level], bn_groups) + ) + if segment: + feats.append( + torch.cat( + self._segment_max_mean( + hl, pack.level_lengths[level], pack.level_seg[level] + ), + dim=1, + ) + ) + continue + lf = ( + hl.new_zeros(R * N, C) + .index_copy(0, pack.level_idx[level], hl) + .view(R, N, C) + .transpose(1, 2) + ) + lmask = pack.level_masks[level] + feats.append( + torch.cat( + [self._masked_max(lf, lmask), self._masked_mean(lf, lmask)], dim=1 + ) + ) + latent = self.hierarchical_proj(torch.cat(feats, dim=1)) + return latent, per_point_features, None + + def forward( + self, + dist_vec_norm: torch.Tensor, + velocity_norm: torch.Tensor, + relative_pos: torch.Tensor, + phys_params: Optional[torch.Tensor] = None, + time_step: Optional[torch.Tensor] = None, + valid_mask: Optional[torch.Tensor] = None, + height: Optional[torch.Tensor] = None, + vtx_rank: Optional[torch.Tensor] = None, + bn_groups: int = 1, + pack: Optional["EncoderPack"] = None, + ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + """ + Encode rigid body point cloud into latent representation. + Follows the original PointNet architecture. + + Args: + dist_vec_norm: Normalized distance vectors (B, N, 3) where N is number of vertices + velocity_norm: Normalized velocities (B, N, 3) + relative_pos: Position relative to reference pose (B, N, 3) + phys_params: Optional physics parameters (B, N, 3) - [mass, friction, restitution] broadcasted to all vertices + time_step: Optional time step size (B, N, 1) - skip_step_size broadcasted to all vertices + valid_mask: Optional vertex validity mask (B, N) - True for valid vertices, False for padding + + Returns: + latent: Global latent representation (B, output_dim) + per_point_features: Per-point features before pooling (B, output_dim, N) + trans_feat: Feature transformation matrix if feature_transform=True + """ + B, N, _ = dist_vec_norm.shape + + # Concatenate all input features: (B, N, input_dim) + # Base: 9D (dist_vec + velocity + rel_pos) + # +phys_params: +3D = 12D + # +time_step: +2D (s and s^2) = 14D + features = [dist_vec_norm, velocity_norm, relative_pos] + + if phys_params is not None: + features.append(phys_params) + + if time_step is not None: + features.append(time_step) + if self.height_input: + if height is None: + raise ValueError( + "RigidBodyEncoder(height_input=True) requires the `height` tensor (B, N, 1)" + ) + features.append(height) # last channel: height above the ground plane + if not getattr(self, "_height_logged", False): + self._height_logged = True + print( + f"[Encoder] height-above-ground input channel ACTIVE (input_dim={self.input_dim}, " + f"height range [{height.min().item():.2f}, {height.max().item():.2f}] m)", + flush=True, + ) + x = torch.cat(features, dim=-1) + + if pack is not None: + return self._forward_packed(x, pack, bn_groups) + + # Transpose for Conv1d: (B, N, C) -> (B, C, N) + x = x.transpose(1, 2) + + # Apply input spatial transformer if enabled + if self.use_stn: + trans = self.stn(x) + x = x.transpose(2, 1) # (B, C, N) -> (B, N, C) + x = torch.bmm(x, trans) # Apply transformation + x = x.transpose(2, 1) # (B, N, C) -> (B, C, N) + + # The five backbone BatchNorms share one vertex mask: compute its index set once (training with + # masked statistics only; evaluation and frozen statistics never pack). + _bn_idx = ( + valid_mask.nonzero(as_tuple=True) + if ( + self.masked_pooling + and valid_mask is not None + and self.training + and not self.freeze_bn_stats + and bn_groups == 1 + ) + else None + ) + + # First block of convolutions + x = F.silu(self._bn(self.bn1, self.conv1(x), valid_mask, _bn_idx, bn_groups)) + x = F.silu(self._bn(self.bn2, self.conv2(x), valid_mask, _bn_idx, bn_groups)) + + # Apply feature transform at 64d features + trans_feat = None + if self.feature_transform: + trans_feat = self.feature_stn(x) + x_transformed = x.transpose(2, 1) # (B, 64, N) -> (B, N, 64) + x_transformed = torch.bmm(x_transformed, trans_feat) # Apply transformation + x = x_transformed.transpose(2, 1) # (B, N, 64) -> (B, 64, N) + + # Second block of convolutions (after feature transform) + x = F.silu(self._bn(self.bn3, self.conv3(x), valid_mask, _bn_idx, bn_groups)) + x = F.silu(self._bn(self.bn4, self.conv4(x), valid_mask, _bn_idx, bn_groups)) + x = self._bn( + self.bn5, self.conv5(x), valid_mask, _bn_idx, bn_groups + ) # No ReLU on last layer before pooling + + # Store per-point features + per_point_features = x # (B, backbone_dim, N) + + # With zero-valued padding, standard pooling is safe: zeros rarely + # dominate the max and only mildly dilute the mean, so valid_mask is + # accepted for interface compatibility but not applied here. + if self.masked_pooling and valid_mask is not None: + _key = "_masked_pool_logged_" + ("train" if self.training else "eval") + if not getattr(self, _key, False): + setattr(self, _key, True) + print( + f"[Encoder] masked set-function pooling ACTIVE ({'train' if self.training else 'eval'} mode): " + f"pooling={self.pooling_type}, levels={getattr(self, 'num_hierarchical_levels', None)}, " + f"valid={int(valid_mask.sum())}/{valid_mask.numel()} vertices, mask dtype={valid_mask.dtype}", + flush=True, + ) + mask = valid_mask.bool() + if self.pooling_type == "max": + latent = self._masked_max(per_point_features, mask) + elif self.pooling_type == "enhanced": + mean_feat = self._masked_mean(per_point_features, mask) + m = mask.unsqueeze(1).to(per_point_features.dtype) + var = (((per_point_features - mean_feat.unsqueeze(-1)) ** 2) * m).sum( + dim=2 + ) / (m.sum(dim=2) - 1).clamp_min(1.0) + combined = torch.cat( + [ + self._masked_max(per_point_features, mask), + mean_feat, + var.clamp_min(1e-12).sqrt(), + ], + dim=1, + ) + latent = self.pooling_proj(combined) + elif self.pooling_type == "hierarchical": + # Every level pools over the full valid set: a set function of the + # vertices, independent of their order, count, and padding. + feats = [ + torch.cat( + [ + self._masked_max(per_point_features, mask), + self._masked_mean(per_point_features, mask), + ], + dim=1, + ) + ] + if self.hierarchical_subsets: + if vtx_rank is None: + raise ValueError( + "RigidBodyEncoder(hierarchical_subsets=True) requires `vtx_rank` (B, N)" + ) + n_valid = mask.sum(dim=1) # (B,) + if not getattr(self, "_subsets_logged", False): + self._subsets_logged = True + print( + f"[Encoder] hierarchical FPS subsets ACTIVE: levels={len(self.hierarchical_layers)}, " + f"fractions={[1.0 / 2 ** (l + 1) for l in range(len(self.hierarchical_layers))]}, " + f"small objects={self.subset_small_objects}", + flush=True, + ) + for level, layer in enumerate(self.hierarchical_layers): + if self.hierarchical_subsets: + # n // 2^(level+1) vertices; small objects per subset_small_objects + k = hierarchical_subset_size( + n_valid, level, self.subset_small_objects + ) + level_mask = mask & (vtx_rank < k.unsqueeze(1)) + else: + level_mask = mask + h = layer[0](per_point_features) + h = layer[2](self._bn(layer[1], h, level_mask, None, bn_groups)) + feats.append( + torch.cat( + [ + self._masked_max(h, level_mask), + self._masked_mean(h, level_mask), + ], + dim=1, + ) + ) + latent = self.hierarchical_proj(torch.cat(feats, dim=1)) + else: + raise ValueError(f"Unknown pooling_type: {self.pooling_type}") + return latent, per_point_features, trans_feat + + # Unmasked pooling over all slots (PADDING_VALUE=0). + if self.pooling_type == "max": + latent = torch.max(per_point_features, dim=2)[0] # (B, backbone_dim) + + elif self.pooling_type == "enhanced": + max_feat = torch.max(per_point_features, dim=2)[0] # (B, backbone_dim) + mean_feat = torch.mean(per_point_features, dim=2) # (B, backbone_dim) + std_feat = torch.std(per_point_features, dim=2) # (B, backbone_dim) + + combined = torch.cat( + [max_feat, mean_feat, std_feat], dim=1 + ) # (B, backbone_dim*3) + latent = self.pooling_proj(combined) # (B, backbone_dim) + + elif self.pooling_type == "hierarchical": + # Hierarchical: multi-scale feature aggregation (PointNet++ inspired) + global_max = torch.max(per_point_features, dim=2)[0] # (B, backbone_dim) + global_mean = torch.mean(per_point_features, dim=2) # (B, backbone_dim) + global_feat = torch.cat( + [global_max, global_mean], dim=1 + ) # (B, backbone_dim*2) + + local_features = [global_feat] + + for level, layer in enumerate(self.hierarchical_layers): + stride = 2 ** (level + 1) + sampled_feat = per_point_features[ + :, :, ::stride + ] # (B, backbone_dim, N//stride) + + if sampled_feat.shape[2] < 8: + sampled_feat = per_point_features + + local_feat = layer(sampled_feat) # (B, backbone_dim, num_sampled) + local_max = torch.max(local_feat, dim=2)[0] # (B, backbone_dim) + local_mean = torch.mean(local_feat, dim=2) # (B, backbone_dim) + + local_combined = torch.cat( + [local_max, local_mean], dim=1 + ) # (B, backbone_dim*2) + local_features.append(local_combined) + + all_features = torch.cat(local_features, dim=1) + latent = self.hierarchical_proj(all_features) # (B, backbone_dim) + + else: + raise ValueError(f"Unknown pooling_type: {self.pooling_type}") + + return latent, per_point_features, trans_feat + + +class StepSizeFiLM(nn.Module): + """ + Feature-wise Linear Modulation (FiLM) conditioned on step-size. + + Generates per-feature scale (γ) and shift (β) from [s, s²]: + output = features * (1 + γ) + β + + This gives the step size a stronger influence than concatenating + [s, s²] as two extra dimensions of the encoder input. + + Reference: Perez et al., "FiLM: Visual Reasoning with a General + Conditioning Layer", AAAI 2018. + """ + + def __init__(self, time_dim: int = 2, feature_dim: int = 768): + super().__init__() + self.film_net = nn.Sequential( + nn.Linear(time_dim, 256), + nn.SiLU(), + nn.Linear(256, feature_dim * 2), # γ and β + ) + # Zero-init the head so γ = β = 0 and FiLM starts as the identity + nn.init.zeros_(self.film_net[-1].weight) + nn.init.zeros_(self.film_net[-1].bias) + + def forward( + self, features: torch.Tensor, time_features: torch.Tensor + ) -> torch.Tensor: + """ + Args: + features: (B, Nobj, D) or (B*Nobj, Nq, D) — encoder/predictor features + time_features: (B, 2) — [s, s²] per sample + + Returns: + Modulated features, same shape as input + """ + gamma_beta = self.film_net(time_features) # (B, 2*D) + D = features.shape[-1] + gamma = gamma_beta[..., :D] # (B, D) + beta = gamma_beta[..., D:] # (B, D) + + # Broadcast over spatial dimensions (Nobj or Nq) + ndim_diff = features.dim() - gamma.dim() + for _ in range(ndim_diff): + gamma = gamma.unsqueeze(-2) # (B, 1, D) or (B, 1, 1, D) + beta = beta.unsqueeze(-2) + + return features * (1.0 + gamma) + beta diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/geometry_encoders.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/geometry_encoders.py new file mode 100644 index 0000000000..edc4dd75b3 --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/geometry_encoders.py @@ -0,0 +1,254 @@ +# 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. + +""" +Geometry-based positional encoders for the ARoPE ablation study. + +These encoders produce 12D geometry descriptors that replace the ARoPE +anchor-based 12D vector (4 anchors x 3D) in the RoPE encoding path. The +downstream RoPE encoding (12D -> 96D) is reused unchanged. + +Baselines: + - OBBEncoder: OBB center + 3 half-axis vectors (12D) + - PCAFrameEncoder: centroid + sqrt(eigenvalue)-scaled eigenvectors (12D) + +Only the 12D positional descriptor changes; the encoder, decoder, loss, +Kabsch alignment, and training schedule stay identical. +""" + +from __future__ import annotations + +import torch +import torch.nn as nn +from typing import Optional + + +class OBBEncoder(nn.Module): + """ + Oriented Bounding Box (OBB) geometry encoder. + + For each object's point cloud, computes: + g_i^obb = [center(3D), half_axis_1(3D), half_axis_2(3D), half_axis_3(3D)] ∈ R^12 + + The half-axis vectors are the PCA eigenvectors scaled by the half-extent of the + projection along each axis (i.e., half the range of points projected onto each + eigenvector). This captures the bounding box shape and orientation. + + Handles sign ambiguity by: + 1. Enforcing a right-handed frame (det > 0) + 2. Temporal sign continuity with a reference frame (if provided) + """ + + def __init__(self): + super().__init__() + self._prev_axes = None # Cache for temporal sign continuity + + @torch.no_grad() + @torch.cuda.amp.custom_fwd(cast_inputs=torch.float32) + def forward( + self, + points: torch.Tensor, + vtx_mask: torch.Tensor, + obj_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """ + Args: + points: (B, Nobj, Nv, 3) - current frame positions + vtx_mask: (B, Nobj, Nv) - valid vertex mask (bool) + obj_mask: (B, Nobj) - valid object mask (bool), optional + + Returns: + feat: (B, Nobj, 12) - OBB geometry descriptor + """ + B, Nobj, Nv, _ = points.shape + device = points.device + dtype = points.dtype + + feat = torch.zeros(B, Nobj, 12, device=device, dtype=torch.float32) + + # torch.linalg.eigh does not support bfloat16/float16, and AMP autocast + # re-casts matmuls even under custom_fwd, so autocast is fully disabled + # for the geometry computation. + with torch.amp.autocast("cuda", enabled=False): + points_f32 = points.float() + + # Process each (batch, object) pair + for b in range(B): + for i in range(Nobj): + if obj_mask is not None and not obj_mask[b, i]: + continue # padding object: zero feature + + mask_i = vtx_mask[b, i] # (Nv,) + pts = points_f32[b, i][mask_i] # (N_valid, 3) in FP32 + + if pts.shape[0] < 3: + continue # degenerate: zero feature + + # Centroid + center = pts.mean(dim=0) # (3,) + + # PCA of the covariance matrix + pts_c = pts - center + cov = (pts_c.T @ pts_c) / max(pts.shape[0] - 1, 1) + cov = cov.float() # eigh requires FP32 + evals, evecs = torch.linalg.eigh(cov) # ascending order + + # Sort descending + idx = torch.argsort(evals, descending=True) + evals = evals[idx] + evecs = evecs[:, idx] + + # Right-handed frame + if torch.det(evecs) < 0: + evecs[:, -1] = -evecs[:, -1] + + # Temporal sign continuity + if ( + self._prev_axes is not None + and b < self._prev_axes.shape[0] + and i < self._prev_axes.shape[1] + ): + ref = self._prev_axes[b, i] # (3, 3) + for k in range(3): + if torch.dot(evecs[:, k], ref[:, k]) < 0: + evecs[:, k] = -evecs[:, k] + + # Half-extent along each axis (half of projection range) + proj = pts_c @ evecs # (N_valid, 3) + half_extents = ( + proj.max(dim=0).values - proj.min(dim=0).values + ) / 2 # (3,) + + # Scale axes by half-extents + axis1 = half_extents[0] * evecs[:, 0] + axis2 = half_extents[1] * evecs[:, 1] + axis3 = half_extents[2] * evecs[:, 2] + + feat[b, i] = torch.cat([center, axis1, axis2, axis3], dim=0) + + # Cache axes for temporal continuity (detached from the graph) + # by extracting axis columns from feat (B, Nobj, 12) + with torch.no_grad(): + axes_cache = torch.zeros(B, Nobj, 3, 3, device=device, dtype=torch.float32) + for b in range(B): + for i in range(Nobj): + f = feat[b, i] + # Normalize cached axes (direction only) + for k in range(3): + v = f[3 + k * 3 : 6 + k * 3] + norm = v.norm() + if norm > 1e-8: + axes_cache[b, i, :, k] = v / norm + self._prev_axes = axes_cache + + return feat + + +class PCAFrameEncoder(nn.Module): + """ + PCA (covariance) Frame geometry encoder. + + For each object's point cloud, computes: + g_i^pca = [center(3D), sqrt(l_1)*u_1(3D), sqrt(l_2)*u_2(3D), sqrt(l_3)*u_3(3D)] ∈ R^12 + + This encodes second-order shape moments (covariance eigenframe), distinct from + OBB which uses half-extents of the bounding box projection. + + Same sign ambiguity handling as OBBEncoder. + """ + + def __init__(self): + super().__init__() + self._prev_axes = None + + @torch.no_grad() + @torch.cuda.amp.custom_fwd(cast_inputs=torch.float32) + def forward( + self, + points: torch.Tensor, + vtx_mask: torch.Tensor, + obj_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """ + Args: + points: (B, Nobj, Nv, 3) - current frame positions + vtx_mask: (B, Nobj, Nv) - valid vertex mask (bool) + obj_mask: (B, Nobj) - valid object mask (bool), optional + + Returns: + feat: (B, Nobj, 12) - PCA frame geometry descriptor + """ + B, Nobj, Nv, _ = points.shape + device = points.device + dtype = points.dtype + eps = 1e-6 + + feat = torch.zeros(B, Nobj, 12, device=device, dtype=torch.float32) + + # torch.linalg.eigh does not support bfloat16/float16, and AMP autocast + # re-casts matmuls even under custom_fwd, so autocast is fully disabled + # for the geometry computation. + with torch.amp.autocast("cuda", enabled=False): + points_f32 = points.float() + + for b in range(B): + for i in range(Nobj): + if obj_mask is not None and not obj_mask[b, i]: + continue + + mask_i = vtx_mask[b, i] + pts = points_f32[b, i][mask_i] # (N_valid, 3) in FP32 + + if pts.shape[0] < 3: + continue + + center = pts.mean(dim=0) + pts_c = pts - center + cov = (pts_c.T @ pts_c) / max(pts.shape[0] - 1, 1) + + cov = cov.float() # eigh requires FP32 + evals, evecs = torch.linalg.eigh(cov) + idx = torch.argsort(evals, descending=True) + evals = evals[idx] + evecs = evecs[:, idx] + + # Right-handed frame + if torch.det(evecs) < 0: + evecs[:, -1] = -evecs[:, -1] + + # Temporal sign continuity + if ( + self._prev_axes is not None + and b < self._prev_axes.shape[0] + and i < self._prev_axes.shape[1] + ): + ref = self._prev_axes[b, i] + for k in range(3): + if torch.dot(evecs[:, k], ref[:, k]) < 0: + evecs[:, k] = -evecs[:, k] + + # Scale by sqrt(eigenvalue) — standard deviation along each axis + scales = torch.sqrt(torch.clamp(evals, min=eps)) + + axis1 = scales[0] * evecs[:, 0] + axis2 = scales[1] * evecs[:, 1] + axis3 = scales[2] * evecs[:, 2] + + feat[b, i] = torch.cat([center, axis1, axis2, axis3], dim=0) + + # Cache axes for temporal continuity + with torch.no_grad(): + axes_cache = torch.zeros(B, Nobj, 3, 3, device=device, dtype=torch.float32) + for b in range(B): + for i in range(Nobj): + f = feat[b, i] + for k in range(3): + v = f[3 + k * 3 : 6 + k * 3] + norm = v.norm() + if norm > 1e-8: + axes_cache[b, i, :, k] = v / norm + self._prev_axes = axes_cache + + return feat diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/multi_scale_predictor.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/multi_scale_predictor.py new file mode 100644 index 0000000000..d01fcbdc34 --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/multi_scale_predictor.py @@ -0,0 +1,885 @@ +# 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. + +""" +Multi-scale prediction head with cross-attention to decoder features. +Query-based predictor using RoPE-aware attention for position encoding. +""" + +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.nn.functional as F +from typing import List, Optional, Tuple + +# Import RoPE modules for position-aware attention +from embodichain_models.dynamics.rigidformer._native.models.decoder import ( + CrossAttentionWithRoPE, + SelfAttentionWithRoPE, +) + + +class QueryBasedMultiScalePredictor(nn.Module): + """ + Query-based predictor with cross-attention to multi-scale features. + Anchors attend to multi-scale decoder features instead of pooling them. + """ + + def __init__( + self, + state_dim: int = 768, + query_dim: int = 3, + hidden_dim: int = 256, + output_dim: int = 3, + layer_indices: List[int] = [0, 6, 9, 12], + num_heads: int = 8, + dropout: float = 0.1, + use_ref_pose: bool = False, # Whether to use reference pose + norm_type: str = "layer_norm", # 'layer_norm' | 'rms_norm' + qk_norm: bool = False, # normalize q/k before attention + fusion_activation: str = "swiglu", # 'swiglu' | 'gelu' | 'silu' + use_rope: bool = True, # Position-aware attention with RoPE + rope_dim: int = None, # Custom RoPE dimension + dim_per_coord_3d: int = None, # Dim per coordinate for single-point (3D) RoPE + dim_per_coord_multi: int = None, # Dim per coordinate for multi-point RoPE + layer_index_mapping: dict = None, # Maps original layer index to saved position + elementwise_gate: bool = False, # Elementwise gated attention + cross_object_attention: bool = False, # Cross-object anchor attention + cross_object_kv: bool = False, # Cross-object key-value: anchors query all objects + no_global_state: bool = False, # Skip state modules entirely when True + ): + super().__init__() + self.cross_object_attention = cross_object_attention + self.cross_object_kv = cross_object_kv + self.no_global_state = no_global_state + self._cross_object_attention_logged = False # One-time log flag + self._cross_object_kv_logged = False # One-time log flag + self._p0_mask_logged = False # One-time padding mask log + self._p1_no_state_logged = False # One-time no-state log + if cross_object_attention: + print( + f"[Predictor] Cross-object anchor attention enabled: anchors from different objects can interact" + ) + if cross_object_kv: + print( + f"[Predictor] Cross-object key-value enabled: anchors query all objects' features" + ) + if no_global_state: + print( + f"[Predictor] no_global_state=True: state modules omitted, fusion uses object features only" + ) + + self.layer_indices = layer_indices + self.num_scales = len(layer_indices) + + # Maps original layer index to its position in the saved feature list. + # If None, all layers are assumed saved. + if layer_index_mapping is not None: + self.layer_index_mapping = layer_index_mapping + else: + # Default: assume sequential saving (all layers saved) + self.layer_index_mapping = {idx: idx for idx in layer_indices} + self.use_ref_pose = use_ref_pose + self.use_rope = use_rope + self.hidden_dim = hidden_dim + + # Scale-specific projections for object features + self.scale_projections = nn.ModuleList( + [nn.Linear(state_dim, hidden_dim) for _ in range(self.num_scales)] + ) + + # Scale-specific projections for state features (skipped when no_global_state) + if not no_global_state: + self.state_scale_projections = nn.ModuleList( + [nn.Linear(state_dim, hidden_dim) for _ in range(self.num_scales)] + ) + + # Query projection; query_dim already includes the ref-pose concatenation + self.query_proj = nn.Linear(query_dim, hidden_dim) + + # Cross-attention to object features for each scale with RoPE + if use_rope: + self.cross_attentions = nn.ModuleList( + [ + CrossAttentionWithRoPE( + d_model=hidden_dim, + num_heads=num_heads, + dropout=dropout, + use_rope=True, + rope_dim=rope_dim, + dim_per_coord_3d=dim_per_coord_3d, + dim_per_coord_multi=dim_per_coord_multi, + elementwise_gate=elementwise_gate, + ) + for _ in range(self.num_scales) + ] + ) + # Cross-attention to state features (skipped when no_global_state) + if not no_global_state: + self.state_cross_attentions = nn.ModuleList( + [ + CrossAttentionWithRoPE( + d_model=hidden_dim, + num_heads=num_heads, + dropout=dropout, + use_rope=True, + rope_dim=rope_dim, + dim_per_coord_3d=dim_per_coord_3d, + dim_per_coord_multi=dim_per_coord_multi, + elementwise_gate=elementwise_gate, + ) + for _ in range(self.num_scales) + ] + ) + # Self-attention for anchors within each object (with RoPE) + self.self_attentions = nn.ModuleList( + [ + SelfAttentionWithRoPE( + d_model=hidden_dim, + num_heads=num_heads, + dropout=dropout, + use_rope=True, + rope_dim=rope_dim, + dim_per_coord_3d=dim_per_coord_3d, + dim_per_coord_multi=dim_per_coord_multi, + elementwise_gate=elementwise_gate, + ) + for _ in range(self.num_scales) + ] + ) + else: + # Standard attention + self.cross_attentions = nn.ModuleList( + [ + nn.MultiheadAttention( + hidden_dim, num_heads, dropout=dropout, batch_first=True + ) + for _ in range(self.num_scales) + ] + ) + if not no_global_state: + self.state_cross_attentions = nn.ModuleList( + [ + nn.MultiheadAttention( + hidden_dim, num_heads, dropout=dropout, batch_first=True + ) + for _ in range(self.num_scales) + ] + ) + # Self-attention for anchors (standard attention) + self.self_attentions = nn.ModuleList( + [ + nn.MultiheadAttention( + hidden_dim, num_heads, dropout=dropout, batch_first=True + ) + for _ in range(self.num_scales) + ] + ) + + # Norm modules (configurable type) for norm-first and qk-norm + Norm = ( + nn.LayerNorm + if norm_type == "layer_norm" + else getattr(nn, "RMSNorm", nn.LayerNorm) + ) + # For object attention + self.query_in_norms = nn.ModuleList( + [Norm(hidden_dim) for _ in range(self.num_scales)] + ) + self.kv_in_norms = nn.ModuleList( + [Norm(hidden_dim) for _ in range(self.num_scales)] + ) + self.post_attn_norms = nn.ModuleList( + [Norm(hidden_dim) for _ in range(self.num_scales)] + ) + # For state attention (skipped when no_global_state) + if not no_global_state: + self.state_kv_in_norms = nn.ModuleList( + [Norm(hidden_dim) for _ in range(self.num_scales)] + ) + self.state_post_attn_norms = nn.ModuleList( + [Norm(hidden_dim) for _ in range(self.num_scales)] + ) + # For self-attention (within same object) + self.self_attn_in_norms = nn.ModuleList( + [Norm(hidden_dim) for _ in range(self.num_scales)] + ) + self.self_attn_post_norms = nn.ModuleList( + [Norm(hidden_dim) for _ in range(self.num_scales)] + ) + self._use_qk_norm = qk_norm + + # Cross-object anchor self-attention + # Allows anchors from different objects to interact directly + if cross_object_attention: + if use_rope: + self.cross_obj_self_attentions = nn.ModuleList( + [ + SelfAttentionWithRoPE( + d_model=hidden_dim, + num_heads=num_heads, + dropout=dropout, + use_rope=True, + rope_dim=rope_dim, + dim_per_coord_3d=dim_per_coord_3d, + dim_per_coord_multi=dim_per_coord_multi, + elementwise_gate=elementwise_gate, + ) + for _ in range(self.num_scales) + ] + ) + else: + self.cross_obj_self_attentions = nn.ModuleList( + [ + nn.MultiheadAttention( + hidden_dim, num_heads, dropout=dropout, batch_first=True + ) + for _ in range(self.num_scales) + ] + ) + self.cross_obj_in_norms = nn.ModuleList( + [Norm(hidden_dim) for _ in range(self.num_scales)] + ) + self.cross_obj_post_norms = nn.ModuleList( + [Norm(hidden_dim) for _ in range(self.num_scales)] + ) + + # Fusion and prediction + # With no_global_state, fusion uses object features only (no dead zero channels) + if no_global_state: + fusion_input_dim = hidden_dim * self.num_scales # Only object features + print( + f"[Predictor] fusion_input_dim = {fusion_input_dim} (object features only)" + ) + else: + fusion_input_dim = ( + hidden_dim * self.num_scales * 2 + ) # × 2 for object + state + if use_ref_pose: + # Add extra projection for reference pose in fusion + # Reference pose is always 3D (positions) + self.ref_pose_proj = nn.Linear(3, hidden_dim) + fusion_input_dim += hidden_dim + print(f"[Predictor] fusion_input_dim with ref_pose = {fusion_input_dim}") + + # Fusion head: norm-first + SwiGLU (default) + def _fusion_block(in_dim: int, out_dim: int) -> nn.Sequential: + if fusion_activation == "swiglu": + return nn.Sequential( + ( + nn.LayerNorm(in_dim) + if norm_type == "layer_norm" + else getattr(nn, "RMSNorm", nn.LayerNorm)(in_dim) + ), + nn.Linear(in_dim, out_dim), + nn.Linear(in_dim, out_dim), # gate + ) + else: + act = nn.GELU() if fusion_activation == "gelu" else nn.SiLU() + return nn.Sequential( + ( + nn.LayerNorm(in_dim) + if norm_type == "layer_norm" + else getattr(nn, "RMSNorm", nn.LayerNorm)(in_dim) + ), + nn.Linear(in_dim, out_dim), + act, + ) + + # Build two-stage fusion: [fusion_dim -> hidden_dim*2 -> hidden_dim] then linear to output + self.fusion_dim1 = hidden_dim * 2 + self.fusion_norm1 = ( + nn.LayerNorm(fusion_input_dim) + if norm_type == "layer_norm" + else getattr(nn, "RMSNorm", nn.LayerNorm)(fusion_input_dim) + ) + self.fusion_w1 = nn.Linear(fusion_input_dim, self.fusion_dim1) + self.fusion_g1 = nn.Linear(fusion_input_dim, self.fusion_dim1) + self.fusion_drop = nn.Dropout(dropout) + self.fusion_norm2 = ( + nn.LayerNorm(self.fusion_dim1) + if norm_type == "layer_norm" + else getattr(nn, "RMSNorm", nn.LayerNorm)(self.fusion_dim1) + ) + self.fusion_w2 = nn.Linear(self.fusion_dim1, hidden_dim) + self.fusion_g2 = nn.Linear(self.fusion_dim1, hidden_dim) + self.fusion_out = nn.Linear(hidden_dim, output_dim) + + # Lightweight penetration head for vertex-level sign prediction: +1 (outside) + # vs -1 (inside) for all vertices, computed from the last-layer object + # feature, a global state feature, and the vertex position. + # Two 32-d projections + a 16-d position projection -> 1-d tanh score. + self._pen_obj_proj = nn.Linear(state_dim, 32) + self._pen_state_proj = nn.Linear(state_dim, 32) + self._pen_pos_proj = nn.Linear(3, 16) + self._pen_head_norm = ( + nn.LayerNorm(32 + 32 + 16) + if norm_type == "layer_norm" + else getattr(nn, "RMSNorm", nn.LayerNorm)(32 + 32 + 16) + ) + self._pen_head = nn.Linear(32 + 32 + 16, 1) + + # torch.nn.RMSNorm built without an explicit eps uses torch.finfo(input.dtype).eps: 0.0078125 for the bf16 inputs + # four module groups receive under autocast in training, 1.19e-7 for fp32 inputs, and the evaluation rollouts run + # without autocast. pin_rmsnorm_eps() gives every module the eps of the dtype it receives in bf16-autocast + # training (input dtypes measured in bf16-autocast training), so evaluation computes the trained function; + # bf16-autocast training itself is unchanged. + _RMSNORM_BF16_TRAINING_INPUT = ( + "query_in_norms", + "kv_in_norms", + "post_attn_norms", + "fusion_norm2", + ) + _RMSNORM_FP32_TRAINING_INPUT = ( + "self_attn_in_norms", + "self_attn_post_norms", + "cross_obj_in_norms", + "cross_obj_post_norms", + "fusion_norm1", + "_pen_head_norm", + ) + + def pin_rmsnorm_eps(self) -> int: + bf16, fp32 = torch.finfo(torch.bfloat16).eps, torch.finfo(torch.float32).eps + n = 0 + for name, m in self.named_modules(): + if not isinstance(m, nn.RMSNorm): + continue + top = name.split(".")[0] + if top in self._RMSNORM_BF16_TRAINING_INPUT: + m.eps = bf16 + elif top in self._RMSNORM_FP32_TRAINING_INPUT: + m.eps = fp32 + else: + raise ValueError( + f"pin_rmsnorm_eps: training input dtype of RMSNorm '{name}' was not measured" + ) + n += 1 + return n + + def forward( + self, + multi_scale_features: List[torch.Tensor], + query: torch.Tensor, + ref_pose: Optional[torch.Tensor] = None, + query_pos: Optional[torch.Tensor] = None, # Anchor 3D positions for RoPE + key_pos: Optional[torch.Tensor] = None, # Object anchor positions for RoPE + multi_scale_state_features: Optional[ + List[torch.Tensor] + ] = None, # State features from decoder + state_pos: Optional[torch.Tensor] = None, # State 3D grid positions for RoPE + obj_mask: Optional[ + torch.Tensor + ] = None, # (B, Nobj) bool mask for valid objects + ) -> torch.Tensor: + """ + Query-based prediction using cross-attention to multi-scale object and state features. + + Supports two input formats: + 1. Grouped: query (B, Nobj, Nq, feature_dim) - preferred for per-object attention + 2. Flat: query (B, M, feature_dim) where M=Nobj*Nq + + Args: + multi_scale_features: List of object features from decoder layers (B, Nobj, D) + query: Query features + - Grouped: (B, Nobj, Nq, 9/12) - dist_vec + velocity + rel_pos + - Flat: (B, Nobj*Nq, 9/12) + ref_pose: Reference pose + - Grouped: (B, Nobj, Nq, 3) + - Flat: (B, Nobj*Nq, 3) + query_pos: Anchor 3D positions for RoPE + - Grouped: (B, Nobj, Nq, 3) + - Flat: (B, Nobj*Nq, 3) + key_pos: Object anchor positions (B, Nobj, Nq*3) for multi-point RoPE + (RenderFormer-style: flattened multi-point representation) + multi_scale_state_features: List of state features from decoder layers (B, N_state, D) + state_pos: State 3D grid positions (B, N_state, 3) for RoPE + + Returns: + Predictions matching input format: + - Grouped input -> (B, Nobj, Nq, output_dim) + - Flat input -> (B, Nobj*Nq, output_dim) + """ + # Detect input format + is_grouped = query.dim() == 4 # (B, Nobj, Nq, D) vs (B, M, D) + + if is_grouped: + B, Nobj, Nq, query_dim = query.shape + M = Nobj * Nq + else: + B, M, query_dim = query.shape + # Infer Nobj from key_pos or from the first available feature layer + if key_pos is not None: + Nobj = key_pos.shape[1] + else: + # Safely access the first available layer in the saved features + # (layer 0 might not be saved with memory optimization) + Nobj = ( + multi_scale_features[0].shape[1] + if len(multi_scale_features) > 0 + else 1 + ) + Nq = M // Nobj + assert M == Nobj * Nq, f"M={M} must be divisible by Nobj={Nobj}" + + # Select object features from specified layers using the mapping + # The mapping converts from original layer index to position in saved list + selected_features = [ + multi_scale_features[self.layer_index_mapping[i]] + for i in self.layer_indices + ] + + # Select state features from specified layers (if provided) + selected_state_features = None + if multi_scale_state_features is not None: + selected_state_features = [ + multi_scale_state_features[self.layer_index_mapping[i]] + for i in self.layer_indices + ] + + # Convert to grouped format for processing + if is_grouped: + # Already grouped + query_grouped = query # (B, Nobj, Nq, query_dim) + ref_pose_grouped = ref_pose # (B, Nobj, Nq, 3) + query_pos_grouped = ( + query_pos if query_pos is not None else None + ) # (B, Nobj, Nq, 3) + else: + # Reshape flat to grouped + query_grouped = query.view( + B, Nobj, Nq, query_dim + ) # (B, Nobj, Nq, query_dim) + if ref_pose is not None: + ref_pose_grouped = ref_pose.view(B, Nobj, Nq, 3) + if query_pos is not None: + query_pos_grouped = query_pos.view(B, Nobj, Nq, 3) + + # Prepare query with reference pose + assert ref_pose_grouped is not None + query_input_grouped = torch.cat( + [query_grouped, ref_pose_grouped], dim=-1 + ) # (B, Nobj, Nq, query_dim+3) + + # Reshape for batched per-object attention: (B*Nobj, Nq, query_dim+3) + query_input = query_input_grouped.view(B * Nobj, Nq, -1) + + # Validate query dimensionality against the projection layer + import logging + + logger = logging.getLogger(__name__) + expected_dim = self.query_proj.in_features + actual_dim = query_input.shape[-1] + if expected_dim != actual_dim: + logger.error( + f"Dimension mismatch: query_proj expects {expected_dim}D input, got {actual_dim}D" + ) + logger.error( + f" query shape: {query.shape}, ref_pose shape: {ref_pose.shape if ref_pose is not None else None}" + ) + raise RuntimeError( + f"Predictor input dimension mismatch: expected {expected_dim}, got {actual_dim}" + ) + + # Project query: (B*Nobj, Nq, query_dim+3) -> (B*Nobj, Nq, hidden_dim) + query_proj = self.query_proj(query_input) # (B*Nobj, Nq, hidden_dim) + + # Process each scale with cross-attention (with RoPE if enabled) + # For batched per-object attention, we reshape features to (B*Nobj, 1, hidden_dim) + # so each object's features can be attended by its Nq anchors + # When cross_object_kv=True, each anchor can attend to ALL objects' features (B*Nobj, Nobj, hidden_dim) + attended_features = [] + attended_state_features = [] + + # Build attention masks for cross-object attention from obj_mask + kv_mask_for_cross_obj = None # For cross-object KV attention + anchor_mask_for_global = None # For global anchor self-attention + if obj_mask is not None and ( + self.cross_object_kv or self.cross_object_attention + ): + # kv mask for cross-object KV: each anchor group may only attend to + # valid objects' features; obj_mask (B, Nobj) expands to all anchor groups + kv_mask_expanded = obj_mask.unsqueeze(1).expand( + B, Nobj, Nobj + ) # (B, Nobj, Nobj) + kv_mask_expanded = kv_mask_expanded.reshape( + B * Nobj, Nobj + ) # (B*Nobj, Nobj) + # SDPA additive format: (B*Nobj, 1, Nq, Nobj); True=attend becomes 0.0 + kv_mask_for_cross_obj = ( + kv_mask_expanded.unsqueeze(1).unsqueeze(1).expand(-1, 1, Nq, -1) + ) # (B*Nobj, 1, Nq, Nobj) + kv_mask_for_cross_obj = torch.where( + kv_mask_for_cross_obj, 0.0, float("-inf") + ) # SDPA additive mask + + # Mask for global anchor self-attention: expand obj_mask to per-anchor + # (B, Nobj) -> (B, Nobj*Nq) + anchor_mask_flat = ( + obj_mask.unsqueeze(-1).expand(B, Nobj, Nq).reshape(B, Nobj * Nq) + ) # (B, Nobj*Nq) + # SDPA mask: (B, 1, Nobj*Nq, Nobj*Nq); only valid anchors attend to valid anchors + anchor_mask_2d = anchor_mask_flat.unsqueeze(1) & anchor_mask_flat.unsqueeze( + 2 + ) # (B, Nobj*Nq, Nobj*Nq) + anchor_mask_for_global = torch.where( + anchor_mask_2d.unsqueeze(1), 0.0, float("-inf") + ) # (B, 1, Nobj*Nq, Nobj*Nq) + + if not self._p0_mask_logged: + n_valid = obj_mask.sum(dim=1).float().mean().item() + n_pad = Nobj - n_valid + print( + f"[Predictor] Padding mask active: {n_valid:.0f} valid / {n_pad:.0f} padding objects, " + f"{n_pad*Nq:.0f} padded anchors masked in attention" + ) + self._p0_mask_logged = True + + for i, feat in enumerate(selected_features): + # Project features: (B, Nobj, hidden_dim) + feat_proj = self.scale_projections[i](feat) # (B, Nobj, hidden_dim) + + if self.cross_object_kv: + # === Cross-Object Key-Value Mode === + # Each anchor can attend to ALL objects' features + # Expand feat_proj for all objects: (B, Nobj, D) -> (B*Nobj, Nobj, D) + # Each of the B*Nobj anchor-groups sees all Nobj object features + kv_batched = ( + feat_proj.unsqueeze(1) + .expand(B, Nobj, Nobj, -1) + .reshape(B * Nobj, Nobj, -1) + ) + # kv_batched: (B*Nobj, Nobj, hidden_dim) + + # One-time runtime confirmation log + if not self._cross_object_kv_logged: + print( + f"[Predictor Runtime] Cross-object KV active: each anchor queries {Nobj} objects' features" + ) + self._cross_object_kv_logged = True + else: + # === Standard Per-Object Mode === + # Each object has 1 feature vector that Nq anchors will attend to + kv_batched = feat_proj.view(B * Nobj, 1, -1) # (B*Nobj, 1, hidden_dim) + + # Norm-first with optional qk-norm + q_in = self.query_in_norms[i](query_proj) # (B*Nobj, Nq, hidden_dim) + kv_in = self.kv_in_norms[i](kv_batched) if self._use_qk_norm else kv_batched + + # Per-object attention: Nq anchors attend to object feature(s) + if self.use_rope: + # Use RoPE-aware cross-attention + if query_pos_grouped is None or key_pos is None: + raise ValueError( + "QueryBasedMultiScalePredictor with use_rope=True requires " + "query_pos (anchor positions) and key_pos (object positions)" + ) + # Reshape positions for batched per-object + # query_pos: individual anchor positions (B, Nobj, Nq, 3) -> (B*Nobj, Nq, 3) + query_pos_batched = query_pos_grouped.view( + B * Nobj, Nq, 3 + ) # (B*Nobj, Nq, 3) + + if self.cross_object_kv: + # Key positions for ALL objects: (B, Nobj, Nq*3) -> (B*Nobj, Nobj, Nq*3) + # Each anchor-group sees all objects' anchor positions + key_pos_batched = ( + key_pos.unsqueeze(1) + .expand(B, Nobj, Nobj, -1) + .reshape(B * Nobj, Nobj, -1) + ) + # key_pos_batched: (B*Nobj, Nobj, Nq*3) + else: + # key_pos: flattened multi-point anchors (B, Nobj, Nq*3) -> (B*Nobj, 1, Nq*3) + key_pos_batched = key_pos.view( + B * Nobj, 1, -1 + ) # (B*Nobj, 1, Nq*3) for multi-point RoPE + + attn_out = self.cross_attentions[i]( + query=q_in, # (B*Nobj, Nq, hidden_dim) + key_value=kv_in, # (B*Nobj, 1 or Nobj, hidden_dim) + pos_q=query_pos_batched, # (B*Nobj, Nq, 3) - individual anchor positions + pos_k=key_pos_batched, # (B*Nobj, 1 or Nobj, Nq*3) - multi-point anchors + attn_mask=kv_mask_for_cross_obj if self.cross_object_kv else None, + ) # (B*Nobj, Nq, hidden_dim) + else: + # Standard attention + attn_out, _ = self.cross_attentions[i]( + query=q_in, key=kv_in, value=kv_in + ) # (B*Nobj, Nq, hidden_dim) + + # Residual + post norm + attn_out = self.post_attn_norms[i]( + attn_out + query_proj + ) # (B*Nobj, Nq, hidden_dim) + attended_features.append(attn_out) + + # === Cross-attention to state features (skipped when no_global_state) === + if not self.no_global_state and selected_state_features is not None: + state_feat = selected_state_features[i] # (B, N_state, D) + state_feat_proj = self.state_scale_projections[i]( + state_feat + ) # (B, N_state, hidden_dim) + + # Expand state features for all anchors across all objects + # state_feat_proj: (B, N_state, hidden_dim) -> (B*Nobj, N_state, hidden_dim) + state_feat_expanded = ( + state_feat_proj.unsqueeze(1) + .expand(B, Nobj, -1, -1) + .reshape(B * Nobj, -1, self.hidden_dim) + ) + + # Norm + state_kv_in = ( + self.state_kv_in_norms[i](state_feat_expanded) + if self._use_qk_norm + else state_feat_expanded + ) + + # Cross-attention: anchors attend to state + if self.use_rope: + if state_pos is None: + raise ValueError( + "state_pos is required when using RoPE for state attention" + ) + + # Expand state_pos for all objects: (B, N_state, 3) -> (B*Nobj, N_state, 3) + state_pos_expanded = ( + state_pos.unsqueeze(1) + .expand(B, Nobj, -1, -1) + .reshape(B * Nobj, -1, 3) + ) + + state_attn_out = self.state_cross_attentions[i]( + query=q_in, # (B*Nobj, Nq, hidden_dim) + key_value=state_kv_in, # (B*Nobj, N_state, hidden_dim) + pos_q=query_pos_batched, # (B*Nobj, Nq, 3) - anchor positions + pos_k=state_pos_expanded, # (B*Nobj, N_state, 3) - state grid positions + attn_mask=None, + ) # (B*Nobj, Nq, hidden_dim) + else: + state_attn_out, _ = self.state_cross_attentions[i]( + query=q_in, key=state_kv_in, value=state_kv_in + ) + + # Residual + post norm + state_attn_out = self.state_post_attn_norms[i]( + state_attn_out + query_proj + ) + attended_state_features.append(state_attn_out) + else: + if not self.no_global_state: + # Without a global state input: append zeros to keep fusion dimensions consistent + attended_state_features.append( + torch.zeros_like(attended_features[-1]) + ) + # With no_global_state, attended_state_features stays empty + + # === Self-attention among anchors within each object === + # With no_global_state, only object features feed the self-attention + if self.no_global_state or len(attended_state_features) <= i: + sa_input = attended_features[ + -1 + ] # (B*Nobj, Nq, hidden_dim), object features only + else: + sa_input = ( + attended_features[-1] + attended_state_features[-1] + ) # (B*Nobj, Nq, hidden_dim) + + # Norm-first + sa_input_normed = self.self_attn_in_norms[i]( + sa_input + ) # (B*Nobj, Nq, hidden_dim) + + # Self-attention: anchors attend to each other within the same object + if self.use_rope: + sa_out = self.self_attentions[i]( + query=sa_input_normed, # (B*Nobj, Nq, hidden_dim) + key_value=sa_input_normed, # (B*Nobj, Nq, hidden_dim) + pos_q=query_pos_batched, # (B*Nobj, Nq, 3) - anchor positions + pos_k=query_pos_batched, # (B*Nobj, Nq, 3) - same positions for self-attn + attn_mask=None, + ) # (B*Nobj, Nq, hidden_dim) + else: + sa_out, _ = self.self_attentions[i]( + query=sa_input_normed, key=sa_input_normed, value=sa_input_normed + ) # (B*Nobj, Nq, hidden_dim) + + # Residual + post norm + sa_out = self.self_attn_post_norms[i]( + sa_out + sa_input + ) # (B*Nobj, Nq, hidden_dim) + + # === Cross-object anchor self-attention === + # Allows anchors from different objects to interact directly + if self.cross_object_attention: + # One-time runtime confirmation log + if not self._cross_object_attention_logged: + print( + f"[Predictor Runtime] Cross-object anchor attention active: B={B}, Nobj={Nobj}, Nq={Nq}, total_anchors={Nobj*Nq}" + ) + self._cross_object_attention_logged = True + # Reshape from per-object to global: (B*Nobj, Nq, D) -> (B, Nobj*Nq, D) + sa_out_global = sa_out.view( + B, Nobj * Nq, -1 + ) # (B, Nobj*Nq, hidden_dim) + + # Reshape anchor positions: (B, Nobj, Nq, 3) -> (B, Nobj*Nq, 3) + query_pos_global = query_pos_grouped.view( + B, Nobj * Nq, 3 + ) # (B, Nobj*Nq, 3) + + # Norm-first + cross_obj_input = self.cross_obj_in_norms[i]( + sa_out_global + ) # (B, Nobj*Nq, hidden_dim) + + # Cross-object self-attention: all anchors attend to all anchors + if self.use_rope: + cross_obj_out = self.cross_obj_self_attentions[i]( + query=cross_obj_input, # (B, Nobj*Nq, hidden_dim) + key_value=cross_obj_input, # (B, Nobj*Nq, hidden_dim) + pos_q=query_pos_global, # (B, Nobj*Nq, 3) - all anchor positions + pos_k=query_pos_global, # (B, Nobj*Nq, 3) - same for self-attn + attn_mask=anchor_mask_for_global, # Mask padding anchors + ) # (B, Nobj*Nq, hidden_dim) + else: + cross_obj_out, _ = self.cross_obj_self_attentions[i]( + query=cross_obj_input, + key=cross_obj_input, + value=cross_obj_input, + ) # (B, Nobj*Nq, hidden_dim) + + # Residual + post norm + cross_obj_out = self.cross_obj_post_norms[i]( + cross_obj_out + sa_out_global + ) # (B, Nobj*Nq, hidden_dim) + + # Reshape back to per-object: (B, Nobj*Nq, D) -> (B*Nobj, Nq, D) + sa_out = cross_obj_out.view( + B * Nobj, Nq, -1 + ) # (B*Nobj, Nq, hidden_dim) + + # Replace this scale's object features with the enhanced version + attended_features[-1] = sa_out + + # Concatenate multi-scale features for fusion + # With no_global_state, only object features are concatenated + if self.no_global_state: + multi_scale_output = torch.cat(attended_features, dim=-1) + # (B*Nobj, Nq, hidden_dim * num_scales) + if not self._p1_no_state_logged: + print( + f"[Predictor] Fusion uses object features only: {multi_scale_output.shape[-1]}D" + ) + self._p1_no_state_logged = True + else: + multi_scale_output = torch.cat( + attended_features + attended_state_features, dim=-1 + ) + # (B*Nobj, Nq, hidden_dim * num_scales * 2) + + # Add reference pose to fusion if enabled + if self.use_ref_pose: + # ref_pose_grouped: (B, Nobj, Nq, 3) -> (B*Nobj, Nq, 3) + ref_pose_batched = ref_pose_grouped.view(B * Nobj, Nq, 3) + ref_pose_features = self.ref_pose_proj( + ref_pose_batched + ) # (B*Nobj, Nq, hidden_dim) + multi_scale_output = torch.cat( + [multi_scale_output, ref_pose_features], dim=-1 + ) + # (B*Nobj, Nq, hidden_dim * num_scales + hidden_dim) + + # Final fusion and prediction (SwiGLU blocks) + # All operations are per-anchor: (B*Nobj, Nq, dim) + # Block 1 + x = self.fusion_norm1(multi_scale_output) + if hasattr(torch.nn.functional, "silu"): + x = self.fusion_drop( + torch.nn.functional.silu(self.fusion_g1(x)) * self.fusion_w1(x) + ) + else: + x = self.fusion_drop(torch.sigmoid(self.fusion_g1(x)) * self.fusion_w1(x)) + # Block 2 + x = self.fusion_norm2(x) + if hasattr(torch.nn.functional, "silu"): + x = self.fusion_drop( + torch.nn.functional.silu(self.fusion_g2(x)) * self.fusion_w2(x) + ) + else: + x = self.fusion_drop(torch.sigmoid(self.fusion_g2(x)) * self.fusion_w2(x)) + if getattr(self, "fp32_output", False): + # model.predictor.fp32_output: the acceleration read-out in fp32 also under bf16 autocast (training), as + # in the fp32 evaluation rollout + with torch.autocast(device_type=x.device.type, enabled=False): + predictions = self.fusion_out(x.float()) # (B*Nobj, Nq, output_dim) + else: + predictions = self.fusion_out(x) # (B*Nobj, Nq, output_dim) + + # Reshape output to match input format + if is_grouped: + # Return grouped format: (B, Nobj, Nq, output_dim) + predictions = predictions.view(B, Nobj, Nq, -1) + else: + # Return flat format: (B, Nobj*Nq, output_dim) + predictions = predictions.view(B, M, -1) + + return predictions + + def predict_penetration_vertices( + self, + multi_scale_features: List[torch.Tensor], + multi_scale_state_features: Optional[List[torch.Tensor]], + pos_full: torch.Tensor, # (B, Nobj, Nv, 3) + ) -> torch.Tensor: + """ + Predict penetration sign for ALL vertices (vertex-level supervision). + + Returns penetration score in [-1, 1] where: + - +1 indicates vertex is outside other objects (no penetration) + - -1 indicates vertex is inside other objects (penetration) + + Args: + multi_scale_features: List of object features from decoder layers (B, Nobj, D) + multi_scale_state_features: List of state features from decoder layers (B, N_state, D) or None + pos_full: Vertex positions (B, Nobj, Nv, 3) - AFTER Kabsch transform + + Returns: + Penetration scores (B, Nobj, Nv, 1) with tanh activation, requires_grad=True for training + """ + # Use layer_index_mapping to get the correct position in the saved features list + last_layer_idx = self.layer_indices[-1] + mapped_idx = self.layer_index_mapping[last_layer_idx] + + obj_last = multi_scale_features[mapped_idx] # (B, Nobj, D) + B, Nobj, _ = obj_last.shape + obj_proj = self._pen_obj_proj(obj_last) # (B, Nobj, 32) + + if ( + multi_scale_state_features is not None + and len(multi_scale_state_features) > 0 + ): + state_last = multi_scale_state_features[mapped_idx] # (B, N_state, D) + state_global = state_last.mean(dim=1) # (B, D) + else: + state_global = obj_last.mean(dim=1) * 0.0 # (B, D) + state_proj = self._pen_state_proj(state_global) # (B, 32) + state_proj = state_proj.unsqueeze(1).expand(-1, Nobj, -1) # (B, Nobj, 32) + + Nv = pos_full.shape[2] + pos_proj = self._pen_pos_proj(pos_full) # (B, Nobj, Nv, 16) + obj_proj_v = obj_proj.unsqueeze(2).expand(-1, -1, Nv, -1) # (B, Nobj, Nv, 32) + state_proj_v = state_proj.unsqueeze(2).expand( + -1, -1, Nv, -1 + ) # (B, Nobj, Nv, 32) + + pen_feat = torch.cat( + [obj_proj_v, state_proj_v, pos_proj], dim=-1 + ) # (B, Nobj, Nv, 80) + pen_feat = self._pen_head_norm(pen_feat) + pen = torch.tanh(self._pen_head(pen_feat)) # (B, Nobj, Nv, 1) + return pen diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/physformer_rope.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/physformer_rope.py new file mode 100644 index 0000000000..7746b6461e --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/physformer_rope.py @@ -0,0 +1,826 @@ +# 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. + +""" +RoPE (Rotary Position Embedding) for PhysFormer. +Adapted from RenderFormer's implementation for 3D point clouds and physics simulation. +""" + +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.amp import autocast +from einops import rearrange, repeat +import math +from typing import Optional, Tuple, List + + +def rotate_half(x): + """Rotate interleaved channel pairs: (x0, x1, x2, x3, ...) -> (-x1, x0, -x3, x2, ...). + + Split into pairs, unbind, stack, merge (the einops rearrangement, without its per-call pattern parsing). + """ + x1, x2 = x.unflatten(-1, (-1, 2)).unbind(dim=-1) + return torch.stack((-x2, x1), dim=-1).flatten(-2) + + +# Per-step memo of rotary angles and of their cos/sin. Within one model step every attention call that sees +# the same position tensor recomputes the same angles (16 calls per step in the decoder and predictor); the +# memo computes each once. Entries are keyed by the input's storage pointer, shape, strides, dtype and version +# counter, and keep a reference to it, so a hit can only return values computed from identical memory (an +# in-place change bumps the version; the pinned storage cannot be reused). Only inputs that do not require +# gradients are memoized, so autograd is unaffected and results are bit-identical. The model clears the memo +# at the start of every step (rope_memo_reset), which also keeps CUDA-graph capture self-contained. +_ROPE_MEMO = {} +ROPE_MEMO_ENABLED = True + + +def rope_memo_reset(): + _ROPE_MEMO.clear() + + +_ROPE_MEMO_MAX = 512 + + +def _memo_put(key, value): + if ( + len(_ROPE_MEMO) >= _ROPE_MEMO_MAX + ): # paths that never reset the memo cannot grow it without bound + _ROPE_MEMO.clear() + _ROPE_MEMO[key] = value + + +def _memo_key(tag, x, extra=()): + return ( + tag, + x.device, + x.data_ptr(), + tuple(x.shape), + x.stride(), + x.dtype, + x._version, + ) + tuple(extra) + + +def _memo_usable(*tensors): + # Under torch.compile the graph de-duplicates identical subexpressions itself; the memo (which reads storage + # pointers) is skipped there. + if torch.compiler.is_compiling(): + return False + return ROPE_MEMO_ENABLED and not any(t.requires_grad for t in tensors) + + +def _cos_sin(freqs): + if _memo_usable(freqs): + key = _memo_key("cs", freqs) + hit = _ROPE_MEMO.get(key) + if hit is not None: + return hit[1], hit[2] + c, s_ = freqs.cos(), freqs.sin() + _memo_put(key, (freqs, c, s_)) + return c, s_ + return freqs.cos(), freqs.sin() + + +@autocast("cuda", enabled=False) +def apply_rotary_emb(freqs, t, scale=1.0): + """ + Apply rotary embeddings to tensor t using frequency tensor freqs. + + Args: + freqs: Frequency tensor + t: Input tensor to apply RoPE to + scale: Scaling factor + """ + dtype = t.dtype + rot_dim = freqs.shape[-1] + + assert ( + rot_dim <= t.shape[-1] + ), f"Feature dimension {t.shape[-1]} is not sufficient for rotation dimension {rot_dim}" + + # Apply to the first rot_dim dimensions + t_rot = t[..., :rot_dim] + t_pass = t[..., rot_dim:] + + # Apply rotary embeddings (x * 1.0 == x exactly, so the unit scale is skipped) + cos, sin = _cos_sin(freqs) + if scale == 1.0: + t_rot = (t_rot * cos) + (rotate_half(t_rot) * sin) + else: + t_rot = (t_rot * cos * scale) + (rotate_half(t_rot) * sin * scale) + + return torch.cat([t_rot, t_pass], dim=-1).type(dtype) + + +class PhysFormerRoPE(nn.Module): + """ + Rotary Position Embedding for PhysFormer's 3D point clouds. + Handles both state tokens and vertex positions. + """ + + # Pair sin/cos channels the way rotate_half pairs them (interleaved), which + # keeps the relative-position property; the half-split layout does not. A + # model is evaluated with the layout it was trained with. + INTERLEAVED_PAIRING = False + _PAIRING_LOGGED = False + + def _paired_angles(self, angles: torch.Tensor) -> torch.Tensor: + cls = type(self) + if cls.INTERLEAVED_PAIRING: + if not cls._PAIRING_LOGGED: + cls._PAIRING_LOGGED = True + print( + f"[RoPE] interleaved rotary pairing ACTIVE ({angles.shape[-1]} angles -> {2 * angles.shape[-1]} dims)", + flush=True, + ) + return torch.repeat_interleave(angles, 2, dim=-1) + return torch.cat([angles, angles], dim=-1) + + def __init__( + self, + dim: int, + max_freq: float = 10.0, + use_3d: bool = True, + learnable_freqs: bool = False, + include_time: bool = False, + use_set_aggregation: bool = False, + dim_per_coord_3d: int = None, # Explicit dim per coordinate for single-point (3D) mode + dim_per_coord_multi: int = None, # Explicit dim per coordinate for multi-point mode + ): + """ + Args: + dim: Dimension of the embeddings (rope_dim, e.g., 96) + max_freq: Maximum frequency for sinusoidal embeddings + use_3d: Whether to use 3D positional encoding (x, y, z) + learnable_freqs: Whether to make frequencies learnable + include_time: Whether to include time dimension (for dynamics) + use_set_aggregation: Whether to use set aggregation for multi-point encoding + If True: encode each point independently then aggregate + If False: use dim_per_coord_multi to encode all coords without truncation + dim_per_coord_3d: Explicit dimension per coordinate for single-point 3D mode. + If None, auto-compute: dim // (3 * 2) + Example: 32 for 3 coords × 32 = 96 + dim_per_coord_multi: Explicit dimension per coordinate for multi-point mode. + If None, will raise error when forward_multi_point is called + Example: 4 for 24 coords × 4 = 96 + """ + super().__init__() + + self.dim = dim + self.use_3d = use_3d + self.include_time = include_time + self.use_set_aggregation = use_set_aggregation + self.max_freq = max_freq # Store for dynamic freq generation + + # Determine how many dimensions per coordinate (for forward_3d) + if use_3d: + if include_time: + # 4D: (x, y, z, t) + self.n_coords = 4 + if dim_per_coord_3d is not None: + # Use explicit configuration + self.dim_per_coord_3d = dim_per_coord_3d + # Verify correctness + expected = self.n_coords * self.dim_per_coord_3d + if expected != dim: + raise ValueError( + f"dim_per_coord_3d mismatch: {self.n_coords} coords × {dim_per_coord_3d} = {expected}, " + f"but rope_dim={dim}" + ) + else: + # Auto-compute + assert ( + dim % (self.n_coords * 2) == 0 + ), f"dim {dim} must be divisible by {self.n_coords * 2}" + self.dim_per_coord_3d = dim // (self.n_coords * 2) + else: + # 3D: (x, y, z) + self.n_coords = 3 + if dim_per_coord_3d is not None: + # Use explicit configuration + self.dim_per_coord_3d = dim_per_coord_3d + # Verify correctness + expected = self.n_coords * self.dim_per_coord_3d + if expected != dim: + raise ValueError( + f"dim_per_coord_3d mismatch: {self.n_coords} coords × {dim_per_coord_3d} = {expected}, " + f"but rope_dim={dim}" + ) + else: + # Auto-compute + assert ( + dim % (self.n_coords * 2) == 0 + ), f"dim {dim} must be divisible by {self.n_coords * 2}" + self.dim_per_coord_3d = dim // (self.n_coords * 2) + else: + raise ValueError( + "PhysFormerRoPE now supports only 3D (or 4D with time) positions" + ) + + # For multi-point mode: explicit dim_per_coord (optional) + self.dim_per_coord_multi = dim_per_coord_multi # Can be None + + # Create frequency bands for 3D mode + # dim_per_coord_3d is the total dimension per coordinate (including sin+cos) + # so freq_len = dim_per_coord_3d // 2 + freq_len_3d = self.dim_per_coord_3d // 2 + if self.dim_per_coord_3d % 2 != 0: + raise ValueError( + f"dim_per_coord_3d ({self.dim_per_coord_3d}) must be even for sin/cos pairs" + ) + + freqs = torch.logspace(0, math.log10(max_freq), freq_len_3d) + + if learnable_freqs: + self.freqs = nn.Parameter(freqs, requires_grad=True) + else: + self.register_buffer("freqs", freqs) + + # Optional learnable position embedding for state tokens (unused by default) + self.state_pos_embed = None + + def _freqs_signature(self): + """Value signature of self.freqs, so layers with identical frequencies share memo entries. + Recomputed (one small device-to-host copy) only when the buffer's storage or version changes. + """ + f = self.freqs + state = (f.data_ptr(), f._version, f.dtype, f.device) + cached = getattr(self, "_freqs_sig_cache", None) + if cached is None or cached[0] != state: + sig = ( + f.detach().cpu().numpy().tobytes(), + str(f.dtype), + tuple(f.shape), + ) # exact values + self._freqs_sig_cache = (state, sig) + return self._freqs_sig_cache[1] + + def _memo_extra(self): + return ( + self._freqs_signature(), + type(self).INTERLEAVED_PAIRING, + self.n_coords, + self.include_time, + self.dim_per_coord_3d, + ) + + def forward_3d(self, positions: torch.Tensor) -> torch.Tensor: + if _memo_usable(positions, self.freqs): + key = _memo_key("3d", positions, self._memo_extra()) + hit = _ROPE_MEMO.get(key) + if hit is None: + hit = (positions, self._forward_3d(positions)) + _memo_put(key, hit) + return hit[1] + return self._forward_3d(positions) + + def _forward_3d(self, positions: torch.Tensor) -> torch.Tensor: + """ + Compute RoPE for 3D positions (single point). + + Args: + positions: 3D positions of shape (..., 3) or (..., 4) if include_time + + Returns: + RoPE embeddings of shape (..., dim) + """ + device = positions.device + dtype = positions.dtype + + # Split positions into coordinates + if self.include_time: + assert positions.shape[-1] == 4, "Expected 4D positions (x, y, z, t)" + coords = positions # (..., 4) + else: + assert positions.shape[-1] == 3, "Expected 3D positions (x, y, z)" + coords = positions # (..., 3) + + # Compute frequencies for each coordinate + embeddings = [] + for i in range(self.n_coords): + coord = coords[..., i : i + 1] # (..., 1) + # Outer product with frequencies (use 3D freqs) + angles = coord * self.freqs # (..., dim_per_coord_3d) + # Create sin/cos pairs + emb = self._paired_angles(angles) # (..., 2*dim_per_coord_3d) + embeddings.append(emb) + + # Concatenate all coordinate embeddings + rope_emb = torch.cat(embeddings, dim=-1) # (..., dim) + + return rope_emb + + def forward_multi_point(self, positions: torch.Tensor) -> torch.Tensor: + """ + Compute RoPE for multiple 3D points per token (RenderFormer-style). + + RenderFormer encodes 3 triangle vertices (9D); here Nq anchor points + (Nq*3 D) represent each object. + + Two modes: + 1. Set aggregation (use_set_aggregation=True): encode each point + independently with the full 3D encoding, then mean-pool over points. + No truncation; the result is invariant to point order. + NOTE: forward_3d returns rotary *angles*, which are linear in the + position, so the mean over anchors equals the angles of the anchor + centroid: ARoPE({x_k}) == RoPE(mean_k x_k). The descriptor carries + the anchor mean position only; shape/extent information reaches + attention through the anchor features (AVP / encoder), not through + these phases. + 2. Concatenation (use_set_aggregation=False): encode every coordinate with + dim_per_coord_multi dims and concatenate, spreading rope_dim across + all Nq*3 coordinates. + + Args: + positions: Multi-point positions of shape (..., Nq*3), formatted as + concat([p0_x, p0_y, p0_z, p1_x, p1_y, p1_z, ...]) + + Returns: + RoPE embeddings of shape (..., dim) + + Example: + 8 anchors: input (B, Nobj, 24) -> output (B, Nobj, dim) + """ + device = positions.device + dtype = positions.dtype + + if self.use_set_aggregation: + # ===== Set aggregation mode ===== + # Encode each anchor independently, then aggregate + + # Reshape: (..., Nq*3) -> (..., Nq, 3) + orig_shape = positions.shape[:-1] # Save original batch shape + n_coords_total = positions.shape[-1] + + assert ( + n_coords_total % 3 == 0 + ), f"Multi-point positions must be divisible by 3, got {n_coords_total}" + Nq = n_coords_total // 3 + + # Reshape to separate anchors: (..., Nq, 3) + positions_reshaped = positions.view(*orig_shape, Nq, 3) + + # Flatten batch dimensions: (..., Nq, 3) -> (total_points, 3) + positions_flat = positions_reshaped.reshape(-1, 3) + + # Encode each anchor independently + rope_per_anchor = self.forward_3d(positions_flat) # (total_points, dim) + + # Reshape back: (total_points, dim) -> (..., Nq, dim) + rope_per_anchor = rope_per_anchor.view(*orig_shape, Nq, -1) + + # Aggregate via mean pooling: (..., Nq, dim) -> (..., dim) + rope_emb = rope_per_anchor.mean(dim=-2) + + else: + # ===== Concatenation mode: explicit dim_per_coord, no truncation ===== + # Use explicit dim_per_coord to encode all coords without truncation + + n_coords_total = positions.shape[-1] + + # Use explicit dim_per_coord from config (if provided) + if self.dim_per_coord_multi is not None: + dim_per_coord = self.dim_per_coord_multi + + # Verify configuration correctness + expected_dim = n_coords_total * dim_per_coord + if expected_dim != self.dim: + raise ValueError( + f"Dimension mismatch! rope_dim={self.dim} but " + f"n_coords_total({n_coords_total}) × dim_per_coord({dim_per_coord}) = {expected_dim}. " + f"Please adjust dim_per_coord in config to: {self.dim // n_coords_total}" + ) + else: + # Fallback: auto-compute to avoid truncation + if self.dim % n_coords_total != 0: + raise ValueError( + f"rope_dim ({self.dim}) must be divisible by n_coords_total ({n_coords_total}) " + f"when dim_per_coord is not explicitly provided. " + f"Please set dim_per_coord={self.dim // n_coords_total} in config." + ) + dim_per_coord = self.dim // n_coords_total + + # Frequency vector length (each coord -> sin+cos pairs) + freq_len = dim_per_coord // 2 + if dim_per_coord % 2 != 0: + raise ValueError( + f"dim_per_coord ({dim_per_coord}) must be even (for sin/cos pairs). " + f"Got odd value." + ) + + # Generate frequencies for this specific dim_per_coord + freqs_multi = torch.logspace( + 0, math.log10(self.max_freq), freq_len, device=device, dtype=dtype + ) + + # Encode each coordinate + embeddings = [] + for i in range(n_coords_total): + coord = positions[..., i : i + 1] # (..., 1) + # Outer product with frequencies + angles = coord * freqs_multi # (..., freq_len) + # Create sin/cos pairs + emb = self._paired_angles(angles) # (..., dim_per_coord) + embeddings.append(emb) + + # Concatenate all coordinate embeddings + rope_emb = torch.cat( + embeddings, dim=-1 + ) # (..., n_coords_total * dim_per_coord) + + # Verify: should exactly match self.dim + assert rope_emb.shape[-1] == self.dim, ( + f"RoPE encoding dimension mismatch: expected {self.dim}, got {rope_emb.shape[-1]}. " + f"n_coords={n_coords_total}, dim_per_coord={dim_per_coord}" + ) + + return rope_emb + + def forward_sa_rope(self, positions: torch.Tensor) -> torch.Tensor: + if _memo_usable(positions, self.freqs): + key = _memo_key("sa", positions, self._memo_extra()) + hit = _ROPE_MEMO.get(key) + if hit is None: + hit = (positions, self._forward_sa_rope(positions)) + _memo_put(key, hit) + return hit[1] + return self._forward_sa_rope(positions) + + def _forward_sa_rope(self, positions: torch.Tensor) -> torch.Tensor: + """ + SA-RoPE: Symmetric Anchor RoPE (Deep Sets style). + + Encode each anchor independently via 3D RoPE (high-resolution: + dim_per_coord_3d dims per coordinate), then mean-pool over anchors. + Produces a per-object positional descriptor that is **strictly + permutation-invariant** to anchor ordering within each object. + + Mathematical form: the descriptor is the mean of the per-anchor rotary + angles a(q_ik) = ω ⊗ q_ik (a linear map of the anchor position), + ā_i = (1/Nq) Σ_k a(q_ik) = a( (1/Nq) Σ_k q_ik ), + and attention applies the rotation R(ā_i) to the query/key. Because + the phases are linear, this equals single-point 3D RoPE evaluated at the + anchor centroid; the mean is the Deep Sets aggregation that guarantees + invariance under any anchor permutation π ∈ S_{Nq}. It does not encode + the anchors' spread (a set scaled about its mean gives the same ā_i). + + Args: + positions: Multi-point positions (..., Nq*3) where Nq is + the number of anchor points per object. + Format: [q1_x, q1_y, q1_z, q2_x, q2_y, q2_z, ...] + + Returns: + RoPE embeddings (..., dim) — permutation-invariant over anchors + """ + orig_shape = positions.shape[:-1] # (...,) + n_coords = positions.shape[-1] + assert ( + n_coords % 3 == 0 + ), f"SA-RoPE: positions last dim must be multiple of 3, got {n_coords}" + Nq = n_coords // 3 + + # Reshape: (..., Nq*3) -> (..., Nq, 3) + pos_per_anchor = positions.view(*orig_shape, Nq, 3) + + # Encode each anchor independently via 3D RoPE + # Flatten for batch processing: (..., Nq, 3) -> (total_anchors, 3) + pos_flat = pos_per_anchor.reshape(-1, 3) + rope_per_anchor = self.forward_3d(pos_flat) # (total_anchors, rope_dim) + + # Reshape back: (total_anchors, rope_dim) -> (..., Nq, rope_dim) + rope_per_anchor = rope_per_anchor.view(*orig_shape, Nq, -1) + + # Mean pool over anchors: (..., Nq, rope_dim) -> (..., rope_dim) + # This is the Deep Sets symmetric aggregation — strictly permutation-invariant + rope_emb = rope_per_anchor.mean(dim=-2) + + return rope_emb + + def forward( + self, + positions: Optional[torch.Tensor] = None, + batch_size: Optional[int] = None, + seq_len: Optional[int] = None, + ) -> torch.Tensor: + """ + Flexible forward function that handles different input types. + + Args: + positions: 3D/4D positions for vertices + batch_size: Batch size for generating state token positions + seq_len: Sequence length for state tokens + + Returns: + RoPE embeddings + """ + if positions is not None: + return self.forward_3d(positions) + else: + raise ValueError("Must provide positions (3D or 4D) for RoPE") + + +class PhysFormerAttentionWithRoPE(nn.Module): + """ + Multi-head attention with RoPE for PhysFormer. + Handles both state-body and body-body interactions. + """ + + def __init__( + self, + hidden_dim: int, + num_heads: int, + rope_dim: Optional[int] = None, + use_3d_rope: bool = True, + include_time: bool = False, + dropout: float = 0.1, + ): + super().__init__() + + self.hidden_dim = hidden_dim + self.num_heads = num_heads + self.head_dim = hidden_dim // num_heads + + # Multi-head attention + self.attention = nn.MultiheadAttention( + hidden_dim, num_heads, dropout=dropout, batch_first=True + ) + + # RoPE + self.rope_dim = rope_dim or self.head_dim + self.rope = PhysFormerRoPE( + dim=self.rope_dim, use_3d=use_3d_rope, include_time=include_time + ) + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + query_pos: Optional[torch.Tensor] = None, + key_pos: Optional[torch.Tensor] = None, + attn_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """ + Forward pass with RoPE. + + Args: + query: Query tensor (B, N_q, D) + key: Key tensor (B, N_k, D) + value: Value tensor (B, N_k, D) + query_pos: Query positions for RoPE + key_pos: Key positions for RoPE + attn_mask: Attention mask + + Returns: + Attended features (B, N_q, D) + """ + B, N_q, D = query.shape + _, N_k, _ = key.shape + + # Reshape for multi-head attention + query = query.view(B, N_q, self.num_heads, self.head_dim).transpose(1, 2) + key = key.view(B, N_k, self.num_heads, self.head_dim).transpose(1, 2) + value = value.view(B, N_k, self.num_heads, self.head_dim).transpose(1, 2) + + # Apply RoPE if positions are provided + if query_pos is not None: + query_rope = self.rope(positions=query_pos) + query_rope = query_rope.view(B, N_q, self.num_heads, -1).transpose(1, 2) + query = apply_rotary_emb(query_rope, query) + + if key_pos is not None: + key_rope = self.rope(positions=key_pos) + key_rope = key_rope.view(B, N_k, self.num_heads, -1).transpose(1, 2) + key = apply_rotary_emb(key_rope, key) + + # Reshape back for standard attention + query = query.transpose(1, 2).reshape(B, N_q, D) + key = key.transpose(1, 2).reshape(B, N_k, D) + value = value.transpose(1, 2).reshape(B, N_k, D) + + # Apply attention + out, _ = self.attention(query, key, value, attn_mask=attn_mask) + + return out + + +class PhysFormerCrossAttentionBlock(nn.Module): + """ + Cross-attention block with RoPE for state-body interactions. + """ + + def __init__( + self, + state_dim: int = 512, + body_dim: int = 256, + num_heads: int = 8, + rope_dim: Optional[int] = None, + feedforward_dim: int = 2048, + dropout: float = 0.1, + ): + super().__init__() + + # Project to common dimension + self.common_dim = max(state_dim, body_dim) + self.state_proj = ( + nn.Linear(state_dim, self.common_dim) + if state_dim != self.common_dim + else nn.Identity() + ) + self.body_proj = ( + nn.Linear(body_dim, self.common_dim) + if body_dim != self.common_dim + else nn.Identity() + ) + + # State -> Body attention (state queries body features) + self.state_to_body_attn = PhysFormerAttentionWithRoPE( + self.common_dim, + num_heads, + rope_dim=rope_dim, + use_3d_rope=False, # State uses 1D encoding + dropout=dropout, + ) + + # Body -> State attention (body queries state features) + self.body_to_state_attn = PhysFormerAttentionWithRoPE( + self.common_dim, + num_heads, + rope_dim=rope_dim, + use_3d_rope=True, # Body uses 3D encoding + dropout=dropout, + ) + + # Feedforward networks + self.state_ffn = nn.Sequential( + nn.Linear(self.common_dim, feedforward_dim), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(feedforward_dim, self.common_dim), + nn.Dropout(dropout), + ) + + self.body_ffn = nn.Sequential( + nn.Linear(self.common_dim, feedforward_dim), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(feedforward_dim, self.common_dim), + nn.Dropout(dropout), + ) + + # Layer norms + self.state_norm1 = nn.LayerNorm(self.common_dim) + self.state_norm2 = nn.LayerNorm(self.common_dim) + self.body_norm1 = nn.LayerNorm(self.common_dim) + self.body_norm2 = nn.LayerNorm(self.common_dim) + + def forward( + self, + state_features: torch.Tensor, # (B, N_state, D_state) + body_features: torch.Tensor, # (B, N_body, D_body) + body_positions: torch.Tensor, # (B, N_body, 3) + state_indices: Optional[torch.Tensor] = None, # (B, N_state) + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Bidirectional cross-attention with RoPE. + """ + # Project to common dimension + state_feat = self.state_proj(state_features) + body_feat = self.body_proj(body_features) + + # State -> Body: State queries body features + state_feat2 = self.state_to_body_attn( + query=state_feat, + key=body_feat, + value=body_feat, + query_pos=state_indices, # 1D indices for state + key_pos=body_positions, # 3D positions for bodies + ) + state_feat = state_feat + state_feat2 + state_feat = self.state_norm1(state_feat) + + # Body -> State: Body queries state features + body_feat2 = self.body_to_state_attn( + query=body_feat, + key=state_feat, + value=state_feat, + query_pos=body_positions, # 3D positions for bodies + key_pos=state_indices, # 1D indices for state + ) + body_feat = body_feat + body_feat2 + body_feat = self.body_norm1(body_feat) + + # Feedforward + state_feat = state_feat + self.state_ffn(state_feat) + state_feat = self.state_norm2(state_feat) + + body_feat = body_feat + self.body_ffn(body_feat) + body_feat = self.body_norm2(body_feat) + + return state_feat, body_feat + + +class PhysFormerDecoderWithRoPE(nn.Module): + """ + Complete decoder with RoPE for PhysFormer. + Combines ideas from CUT3R and RenderFormer. + """ + + def __init__( + self, + state_dim: int = 512, + body_dim: int = 256, + num_layers: int = 12, + num_heads: int = 8, + rope_dim: Optional[int] = None, + feedforward_dim: int = 2048, + dropout: float = 0.1, + include_time: bool = True, # For dynamics + return_all_layers: bool = True, + ): + super().__init__() + + self.num_layers = num_layers + self.return_all_layers = return_all_layers + self.include_time = include_time + + # Create layers + self.layers = nn.ModuleList( + [ + PhysFormerCrossAttentionBlock( + state_dim=state_dim if i == 0 else max(state_dim, body_dim), + body_dim=body_dim if i == 0 else max(state_dim, body_dim), + num_heads=num_heads, + rope_dim=rope_dim, + feedforward_dim=feedforward_dim, + dropout=dropout, + ) + for i in range(num_layers) + ] + ) + + def forward( + self, + state_features: torch.Tensor, + body_features: torch.Tensor, + body_positions: torch.Tensor, + time_step: Optional[float] = None, + ) -> Tuple[List[torch.Tensor], List[torch.Tensor]]: + """ + Forward pass through all layers with RoPE. + + Args: + state_features: State token features (B, N_state, D_state) + body_features: Body vertex features (B, N_body*8, D_body) + body_positions: Body vertex 3D positions (B, N_body*8, 3) + time_step: Optional time for dynamics (scalar or tensor) + + Returns: + All layer outputs if return_all_layers, else final outputs + """ + B, N_state, _ = state_features.shape + + # Generate state indices + state_indices = torch.arange(N_state, device=state_features.device) + state_indices = state_indices.unsqueeze(0).expand(B, -1) + + # Add time dimension to positions if needed + if self.include_time and time_step is not None: + if isinstance(time_step, (int, float)): + time_tensor = torch.full( + (B, body_positions.shape[1], 1), + time_step, + device=body_positions.device, + dtype=body_positions.dtype, + ) + else: + time_tensor = time_step.view(B, 1, 1).expand( + B, body_positions.shape[1], 1 + ) + + body_positions = torch.cat( + [body_positions, time_tensor], dim=-1 + ) # (B, N, 4) + + # Store outputs + all_states = [] + all_bodies = [] + + # Process through layers + for layer in self.layers: + state_features, body_features = layer( + state_features, body_features, body_positions, state_indices + ) + + all_states.append(state_features) + all_bodies.append(body_features) + + if self.return_all_layers: + return all_states, all_bodies + else: + return state_features, body_features diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/predictor.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/predictor.py new file mode 100644 index 0000000000..af1d41782b --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/predictor.py @@ -0,0 +1,282 @@ +# 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. + +""" +Prediction heads for the dynamics model: plain MLP, multi-scale, and +residual variants. +""" + +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.nn.functional as F +from typing import List, Optional + + +class DynamicsPredictor(nn.Module): + """ + MLP predictor that takes state features and query points to predict dynamics. + """ + + def __init__( + self, + state_dim: int = 512, + query_dim: int = 3, # 3D position + hidden_dims: List[int] = [256, 128, 64], + output_dim: int = 3, # 3D acceleration/displacement + activation: str = "relu", + use_batch_norm: bool = False, + dropout: float = 0.1, + ): + super().__init__() + + self.state_dim = state_dim + self.query_dim = query_dim + self.output_dim = output_dim + + # Build MLP layers + layers = [] + input_dim = state_dim + query_dim + + for hidden_dim in hidden_dims: + layers.append(nn.Linear(input_dim, hidden_dim)) + + if use_batch_norm: + layers.append(nn.BatchNorm1d(hidden_dim)) + + if activation == "relu": + layers.append(nn.SiLU()) + elif activation == "gelu": + layers.append(nn.SiLU()) + elif activation == "tanh": + layers.append(nn.Tanh()) + else: + raise ValueError(f"Unknown activation: {activation}") + + if dropout > 0: + layers.append(nn.Dropout(dropout)) + + input_dim = hidden_dim + + # Output layer + layers.append(nn.Linear(hidden_dims[-1], output_dim)) + layers.append(nn.Tanh()) # bounded output + self.output_scale = nn.Parameter(torch.ones(1) * 1.0) + + self.mlp = nn.Sequential(*layers) + + def forward( + self, state_features: torch.Tensor, query_points: torch.Tensor + ) -> torch.Tensor: + """ + Predict dynamics for query points given state features. + + Args: + state_features: Enhanced state features (B, N_state, D_state) or (B, D_state) + query_points: Query vertex positions (B, N_query, 3) + + Returns: + Predicted accelerations/displacements (B, N_query, 3) + """ + B = query_points.shape[0] + N_query = query_points.shape[1] + + # If state_features has sequence dimension, pool it + if state_features.dim() == 3: + # Global max pooling over state tokens + state_features = torch.max(state_features, dim=1)[0] # (B, D_state) + + # Expand state features to match number of query points + state_features_expanded = state_features.unsqueeze(1).expand(B, N_query, -1) + + # Concatenate state features with query points + combined = torch.cat([state_features_expanded, query_points], dim=-1) + + # Reshape for batch processing through MLP + combined_flat = combined.view(B * N_query, -1) + + # Apply MLP + output_flat = self.mlp(combined_flat) * self.output_scale + + # Reshape back + output = output_flat.view(B, N_query, self.output_dim) + + return output + + +class MultiScalePredictor(nn.Module): + """ + Multi-scale predictor that processes features at different resolutions. + """ + + def __init__( + self, + state_dim: int = 512, + query_dim: int = 3, + scales: List[int] = [64, 128, 256], + output_dim: int = 3, + activation: str = "relu", + dropout: float = 0.1, + ): + super().__init__() + + self.scales = scales + + # Create predictors for each scale + self.scale_predictors = nn.ModuleList( + [ + DynamicsPredictor( + state_dim=state_dim, + query_dim=query_dim, + hidden_dims=[scale, scale // 2], + output_dim=output_dim, + activation=activation, + dropout=dropout, + ) + for scale in scales + ] + ) + + # Fusion layer + self.fusion = nn.Linear(output_dim * len(scales), output_dim) + + def forward( + self, state_features: torch.Tensor, query_points: torch.Tensor + ) -> torch.Tensor: + """ + Multi-scale prediction with fusion. + """ + predictions = [] + + for predictor in self.scale_predictors: + pred = predictor(state_features, query_points) + predictions.append(pred) + + # Concatenate all scale predictions + combined = torch.cat(predictions, dim=-1) + + # Fuse predictions + output = self.fusion(combined) + + return output + + +class ResidualPredictor(nn.Module): + """ + Predictor with residual connections for better gradient flow. + """ + + def __init__( + self, + state_dim: int = 512, + query_dim: int = 3, + hidden_dim: int = 256, + num_layers: int = 3, + output_dim: int = 3, + activation: str = "relu", + dropout: float = 0.1, + ): + super().__init__() + + self.input_projection = nn.Linear(state_dim + query_dim, hidden_dim) + + # Residual blocks + self.res_blocks = nn.ModuleList( + [ResidualBlock(hidden_dim, activation, dropout) for _ in range(num_layers)] + ) + + self.output_projection = nn.Linear(hidden_dim, output_dim) + + def forward( + self, state_features: torch.Tensor, query_points: torch.Tensor + ) -> torch.Tensor: + """Forward pass with residual connections.""" + B = query_points.shape[0] + N_query = query_points.shape[1] + + # Pool state features if needed + if state_features.dim() == 3: + state_features = torch.max(state_features, dim=1)[0] + + # Expand and concatenate + state_expanded = state_features.unsqueeze(1).expand(B, N_query, -1) + combined = torch.cat([state_expanded, query_points], dim=-1) + + # Initial projection + x = self.input_projection(combined) + + # Apply residual blocks + for block in self.res_blocks: + x = block(x) + + # Output projection + output = self.output_projection(x) + + return output + + +class ResidualBlock(nn.Module): + """Single residual block.""" + + def __init__(self, hidden_dim: int, activation: str = "relu", dropout: float = 0.1): + super().__init__() + + self.fc1 = nn.Linear(hidden_dim, hidden_dim) + self.fc2 = nn.Linear(hidden_dim, hidden_dim) + + if activation == "relu": + self.activation = nn.ReLU() + elif activation == "gelu": + self.activation = nn.GELU() + else: + self.activation = nn.ReLU() + + self.dropout = nn.Dropout(dropout) + self.norm1 = nn.LayerNorm(hidden_dim) + self.norm2 = nn.LayerNorm(hidden_dim) + + def forward(self, x): + residual = x + + x = self.norm1(x) + x = self.fc1(x) + x = self.activation(x) + x = self.dropout(x) + + x = self.norm2(x) + x = self.fc2(x) + x = self.dropout(x) + + return x + residual + + +if __name__ == "__main__": + """Test predictor modules.""" + + print("Testing DynamicsPredictor...") + predictor = DynamicsPredictor( + state_dim=512, query_dim=3, hidden_dims=[256, 128, 64], output_dim=3 + ) + + # Test inputs + batch_size = 4 + state_features = torch.randn(batch_size, 10, 512) # With sequence dimension + query_points = torch.randn(batch_size, 8, 3) # 8 vertices + + output = predictor(state_features, query_points) + print(f"Output shape: {output.shape}") # Should be (4, 8, 3) + + print("\nTesting MultiScalePredictor...") + multi_predictor = MultiScalePredictor(state_dim=512, query_dim=3) + output_multi = multi_predictor(state_features, query_points) + print(f"Multi-scale output shape: {output_multi.shape}") # (4, 8, 3) + + print("\nTesting ResidualPredictor...") + res_predictor = ResidualPredictor(state_dim=512, query_dim=3) + output_res = res_predictor(state_features, query_points) + print(f"Residual output shape: {output_res.shape}") # (4, 8, 3) + + print("\nPredictor tests successful!") diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/se3_attention.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/se3_attention.py new file mode 100644 index 0000000000..3eaf581004 --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/se3_attention.py @@ -0,0 +1,336 @@ +# 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. + +""" +Pairwise relative SE(3) attention bias. + +Baseline for the positional-encoding ablation: replaces ARoPE's anchor-based +rotary position encoding with pairwise relative SE(3) features between +object pairs. + +Design notes: +- Translation uses the current object center c_i (not the Kabsch t_i) +- Relative translations are normalized by the RMS scale s_i +- 7D pairwise feature: [t_ij, omega_ij, log(s_j / s_i)] +- Last MLP layer is zero-initialized, so attention starts unbiased +- Padding pairs receive a finite -1e4 bias (avoids NaN from -inf) +- Outputs are re-masked by obj_mask after attention + +Used only when pe_type='se3'. +""" + +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.nn.functional as F +import roma +import logging +from typing import Optional + +logger = logging.getLogger(__name__) + + +def extract_object_pose_features( + curr_pos: torch.Tensor, # (B, Nobj, Nv, 3) - current frame positions + ref_pose: torch.Tensor, # (B, Nobj, Nv, 3) - reference pose (frame 0) + vtx_mask: torch.Tensor, # (B, Nobj, Nv) - valid vertex mask + obj_mask: torch.Tensor, # (B, Nobj) - valid object mask + eps: float = 1e-6, +) -> tuple: + """ + Extract per-object rigid pose features using center-based Kabsch. + + For each valid object i: + 1. Compute centers: c_ref = mean(X_ref[valid]), c_i = mean(X_t[valid]) + 2. Kabsch on centered points: R_i = Kabsch(X_ref - c_ref, X_t - c_i) + 3. Scale: s_i = RMS radius of reference point cloud + + Returns: + rotmats: (B, Nobj, 3, 3) - rotation matrices (identity for padding) + centers: (B, Nobj, 3) - current frame centers (zero for padding) + scales: (B, Nobj, 1) - RMS radius of ref point cloud (1.0 for padding) + """ + B, Nobj, Nv, _ = curr_pos.shape + device = curr_pos.device + dtype = curr_pos.dtype + + rotmats = torch.eye(3, device=device, dtype=dtype)[None, None].repeat(B, Nobj, 1, 1) + centers = torch.zeros(B, Nobj, 3, device=device, dtype=dtype) + scales = torch.ones(B, Nobj, 1, device=device, dtype=dtype) + + for b in range(B): + for i in range(Nobj): + if not obj_mask[b, i]: + continue + + valid = vtx_mask[b, i].bool() + n_valid = valid.sum() + if n_valid == 0: + continue + + Xr = ref_pose[b, i][valid] # (Ni, 3) - reference + Xt = curr_pos[b, i][valid] # (Ni, 3) - current + + # Compute centers + c_ref = Xr.mean(dim=0) # (3,) + c_cur = Xt.mean(dim=0) # (3,) + + # The translation feature is the current center, not the Kabsch t_i + centers[b, i] = c_cur + + if n_valid < 3: + # Not enough points for Kabsch, keep identity R + continue + + # Center the point clouds + Xr0 = Xr - c_ref # (Ni, 3) + Xt0 = Xt - c_cur # (Ni, 3) + + # Kabsch alignment on centered points yields the rotation only + try: + Ri, _ = roma.rigid_points_registration(Xr0, Xt0) + rotmats[b, i] = Ri + except Exception: + # Degenerate case (e.g., highly symmetric object), keep identity + pass + + # RMS radius of reference point cloud (scale normalization) + si = torch.sqrt(torch.clamp(Xr0.pow(2).sum(dim=-1).mean(), min=eps)) + scales[b, i, 0] = si + + return rotmats, centers, scales + + +class SE3PairwiseBias(nn.Module): + """ + Compute pairwise relative SE(3) attention bias between object pairs. + + 7D pairwise feature: + t_ij = R_i^T (c_j - c_i) / s_i — scale-normalized relative translation (3D) + omega_ij = log_SO3(R_i^T R_j) — relative rotation as axis-angle (3D) + log(s_j / s_i) — log scale ratio (1D) + + An MLP maps the 7D feature to a per-head bias; its last layer is zero-initialized. + """ + + def __init__( + self, num_heads: int, hidden_dim: int = 64, use_scale_ratio: bool = True + ): + super().__init__() + self.num_heads = num_heads + self.use_scale_ratio = use_scale_ratio + d_in = 7 if use_scale_ratio else 6 + + self.mlp = nn.Sequential( + nn.Linear(d_in, hidden_dim), + nn.GELU(), + nn.Linear(hidden_dim, hidden_dim), + nn.GELU(), + nn.Linear(hidden_dim, num_heads), + ) + + # Zero-init last layer: start as vanilla attention (no geometric bias) + nn.init.zeros_(self.mlp[-1].weight) + nn.init.zeros_(self.mlp[-1].bias) + + self._logged = False + + def forward( + self, + rotmats: torch.Tensor, # (B, Nobj, 3, 3) + centers: torch.Tensor, # (B, Nobj, 3) + scales: torch.Tensor, # (B, Nobj, 1) + obj_mask: torch.Tensor, # (B, Nobj) + ) -> torch.Tensor: + """ + Compute pairwise SE(3) attention bias. + + Returns: + bias: (B, num_heads, Nobj, Nobj) - additive attention bias + """ + B, Nobj = obj_mask.shape + + # R_i^T: (B, Nobj, 3, 3) + Ri_T = rotmats.transpose(-1, -2) + + # Relative translation: t_ij = R_i^T (c_j - c_i) / s_i + ci = centers.unsqueeze(2) # (B, Nobj, 1, 3) + cj = centers.unsqueeze(1) # (B, 1, Nobj, 3) + delta_c = cj - ci # (B, Nobj, Nobj, 3) + + Ri_T_expanded = Ri_T.unsqueeze(2) # (B, Nobj, 1, 3, 3) + t_ij = torch.matmul(Ri_T_expanded, delta_c.unsqueeze(-1)).squeeze( + -1 + ) # (B, Nobj, Nobj, 3) + + # Scale normalization: divide by s_i + t_ij = t_ij / scales.clamp_min(1e-6).unsqueeze(2) # (B, Nobj, 1, 1) broadcast + + # Relative rotation: R_ij = R_i^T @ R_j, then log map to axis-angle + Rj = rotmats.unsqueeze(1) # (B, 1, Nobj, 3, 3) + R_ij = torch.matmul(Ri_T_expanded, Rj) # (B, Nobj, Nobj, 3, 3) + omega_flat = roma.rotmat_to_rotvec(R_ij.reshape(-1, 3, 3)) + omega_ij = omega_flat.reshape(B, Nobj, Nobj, 3) + + # Build feature vector + feats = [t_ij, omega_ij] + + if self.use_scale_ratio: + # log(s_j / s_i) = log(s_j) - log(s_i) + log_s = scales.clamp_min(1e-6).log().squeeze(-1) # (B, Nobj) + log_ratio = log_s.unsqueeze(2) - log_s.unsqueeze(1) # (B, Nobj, Nobj) + feats.append(log_ratio.unsqueeze(-1)) + + phi_ij = torch.cat(feats, dim=-1) # (B, Nobj, Nobj, 7) + + # MLP to per-head bias + bias = self.mlp(phi_ij) # (B, Nobj, Nobj, H) + bias = bias.permute(0, 3, 1, 2) # (B, H, Nobj, Nobj) + + # Mask padding objects: use -1e4 instead of -inf (avoids NaN issues) + pair_mask = obj_mask[:, :, None] & obj_mask[:, None, :] # (B, Nobj, Nobj) + neg = torch.full_like(bias, -1e4) + bias = torch.where(pair_mask.unsqueeze(1), bias, neg) + + if not self._logged: + n_valid = obj_mask.sum(dim=1).float().mean().item() + valid_bias = bias[pair_mask.unsqueeze(1).expand_as(bias)] + logger.info( + f"[SE3 Bias] Computed pairwise SE(3) bias: " + f"shape={tuple(bias.shape)}, " + f"avg_valid_objects={n_valid:.1f}, " + f"phi_ij dim={phi_ij.shape[-1]}, " + f"bias range=[{valid_bias.min().item():.6f}, {valid_bias.max().item():.6f}] " + f"(zero-init: expected near 0)" + ) + self._logged = True + + return bias + + +class SelfAttentionWithSE3Bias(nn.Module): + """ + Self-attention with pairwise SE(3) attention bias (no RoPE). + + Drop-in replacement for SelfAttentionWithRoPE when pe_type='se3'. + Uses PyTorch SDPA instead of Flash Attention (FA2 doesn't support custom attn_bias). + """ + + def __init__( + self, + d_model: int, + num_heads: int, + dropout: float = 0.1, + qk_norm: bool = True, + elementwise_gate: bool = False, + ): + super().__init__() + assert d_model % num_heads == 0, "d_model must be divisible by num_heads" + self.d_model = d_model + self.num_heads = num_heads + self.head_dim = d_model // num_heads + self.scale = self.head_dim**-0.5 + self.qk_norm = qk_norm + self.elementwise_gate = elementwise_gate + + # Projections + if self.elementwise_gate: + self.q_proj = nn.Linear(d_model, d_model * 2, bias=True) + else: + self.q_proj = nn.Linear(d_model, d_model, bias=True) + self.k_proj = nn.Linear(d_model, d_model, bias=True) + self.v_proj = nn.Linear(d_model, d_model, bias=True) + self.out_proj = nn.Linear(d_model, d_model, bias=True) + self.dropout_p = dropout + + # QK-normalization + if self.qk_norm: + from embodichain_models.dynamics.rigidformer._native.models.decoder import ( + RMSNorm, + ) + + self.q_norm = RMSNorm(self.head_dim) + self.k_norm = RMSNorm(self.head_dim) + + self._logged = False + + def forward( + self, + query: torch.Tensor, # (B, N, D) + key_value: torch.Tensor, # (B, N, D) - same as query for self-attn + se3_bias: torch.Tensor, # (B, H, N, N) - pairwise SE(3) bias + attn_mask: Optional[torch.Tensor] = None, # (B, 1, 1, N) object mask + ) -> torch.Tensor: + """ + Self-attention with SE(3) pairwise bias (no RoPE). + """ + B, N, _ = query.shape + + # Q projection + if self.elementwise_gate: + q_gate = self.q_proj(query) + q_states, gate_score = torch.split(q_gate, self.d_model, dim=-1) + q = q_states.view(B, N, self.num_heads, self.head_dim) + gate_score = gate_score.view(B, N, self.num_heads, self.head_dim) + else: + q = self.q_proj(query).view(B, N, self.num_heads, self.head_dim) + gate_score = None + + k = self.k_proj(key_value).view(B, N, self.num_heads, self.head_dim) + v = self.v_proj(key_value).view(B, N, self.num_heads, self.head_dim) + + # QK-normalization + if self.qk_norm: + q = self.q_norm(q) + k = self.k_norm(k) + + # Transpose to (B, H, N, D) + q = q.transpose(1, 2) + k = k.transpose(1, 2) + v = v.transpose(1, 2) + + # The SE(3) bias serves directly as the additive attention bias; padding + # pairs already carry -1e4, so no separate padding mask is needed. + attn_bias = se3_bias.to(dtype=q.dtype) + + # PyTorch SDPA + dropout_p = self.dropout_p if self.training else 0.0 + out = F.scaled_dot_product_attention( + q, k, v, attn_mask=attn_bias, dropout_p=dropout_p, scale=self.scale + ) # (B, H, N, D) + + # Safety: replace any NaN from degenerate softmax + out = torch.nan_to_num(out, nan=0.0) + + # Transpose back to (B, N, H, D) + out = out.transpose(1, 2) + + if self.elementwise_gate: + out = out * torch.sigmoid(gate_score) + + # Merge heads + out = out.reshape(B, N, self.d_model) + out = self.out_proj(out) + + # Zero out padding query outputs after out_proj so its bias is zeroed as well + if attn_mask is not None: + if attn_mask.dim() == 4: + query_mask = attn_mask[:, 0, 0, :] # (B, N) + else: + query_mask = attn_mask # (B, N) + out = out * query_mask.unsqueeze(-1).to(out.dtype) # (B, N, 1) + + if not self._logged: + valid_bias = se3_bias[se3_bias > -1e3] + if valid_bias.numel() > 0: + logger.info( + f"[SE3 SelfAttn] Forward: " + f"q={tuple(q.shape)}, se3_bias={tuple(se3_bias.shape)}, " + f"valid bias range=[{valid_bias.min().item():.6f}, {valid_bias.max().item():.6f}]" + ) + self._logged = True + + return out diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/state.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/state.py new file mode 100644 index 0000000000..18fd45f5f4 --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/state.py @@ -0,0 +1,423 @@ +# 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. + +""" +State representation module inspired by CUT3R. +Maintains persistent state across time steps for multi-body dynamics. +""" + +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.nn.functional as F +from typing import Optional, Tuple +import math + + +class StateRepresentation(nn.Module): + """ + Persistent state representation for multi-body dynamics, following + CUT3R's state management: + - Register tokens stored in an nn.Embedding, as in CUT3R + - Tokens laid out on a 3D grid whose coordinates drive RoPE + - Rotary-only position encoding (no additive position embeddings) + """ + + def __init__( + self, + num_tokens: int = 216, # 6x6x6 = 216 tokens on a 3D grid + token_dim: int = 768, + position_encoding: str = "learnable_3d", # rope-only: use 3D grid positions + rope_dim: Optional[int] = None, + ): + super().__init__() + + self.num_tokens = num_tokens + self.token_dim = token_dim + self.position_encoding = position_encoding + + # Register tokens (no additive position embeddings; rope-only) + self.register_tokens = nn.Embedding(num_tokens, token_dim) + + # Small random initialization, following CUT3R + nn.init.normal_(self.register_tokens.weight, mean=0.0, std=0.02) + + # 3D grid positions for RoPE (rope-only) + if position_encoding != "learnable_3d": + raise ValueError( + "StateRepresentation now supports only 'learnable_3d' (rope-only)" + ) + if num_tokens == 216: + self.grid_3d = (6, 6, 6) + else: + self.grid_3d = self._compute_3d_grid(num_tokens) + + # Layer normalization + self.norm = nn.LayerNorm(token_dim) + + def _compute_3d_grid(self, n): + """Compute 3D grid dimensions for n tokens.""" + # Try to make roughly cubic + cube_root = int(round(n ** (1 / 3))) + for x in range(cube_root, 0, -1): + if n % x == 0: + remaining = n // x + sqrt_rem = int(round(remaining**0.5)) + for y in range(sqrt_rem, 0, -1): + if remaining % y == 0: + z = remaining // y + return (x, y, z) + return (n, 1, 1) + + def get_state_and_positions( + self, batch_size: int, device: torch.device + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Get initial state tokens and 3D grid positions for RoPE (rope-only). + + Returns: + state_tokens: (B, num_tokens, token_dim) + positions: (B, num_tokens, 3) in [-1, 1] + """ + # Get state tokens using embedding + indices = torch.arange(self.num_tokens, dtype=torch.long, device=device) + state_tokens = self.register_tokens(indices) # (num_tokens, token_dim) + state_tokens = state_tokens.unsqueeze(0).expand(batch_size, -1, -1) + + # Normalize tokens + state_tokens = self.norm(state_tokens) + + # Build 3D grid positions normalized to [-1, 1] + gx, gy, gz = self.grid_3d + xs = torch.linspace(-1, 1, steps=gx, device=device, dtype=state_tokens.dtype) + ys = torch.linspace(-1, 1, steps=gy, device=device, dtype=state_tokens.dtype) + zs = torch.linspace(-1, 1, steps=gz, device=device, dtype=state_tokens.dtype) + grid = torch.stack(torch.meshgrid(xs, ys, zs, indexing="ij"), dim=-1).view( + -1, 3 + ) # (num_tokens,3) + if grid.shape[0] != self.num_tokens: + # Tile or truncate to match the token count + if grid.shape[0] < self.num_tokens: + reps = (self.num_tokens + grid.shape[0] - 1) // grid.shape[0] + grid = grid.repeat(reps, 1)[: self.num_tokens] + else: + grid = grid[: self.num_tokens] + positions = grid.unsqueeze(0).expand(batch_size, -1, -1) # (B,N,3) + + return state_tokens, positions + + def get_initial_state(self, batch_size: int) -> torch.Tensor: + """ + Return the initial state tokens only (grid positions discarded). + """ + device = self.register_tokens.weight.device + state, _ = self.get_state_and_positions(batch_size, device) + return state + + def forward( + self, state: Optional[torch.Tensor] = None, batch_size: Optional[int] = None + ) -> torch.Tensor: + """ + Get state representation (rope-only). + + Args: + state: Current state if available (B, num_tokens, token_dim) + batch_size: Batch size if state is None + + Returns: + State tokens (normalized) + """ + if state is None: + if batch_size is None: + raise ValueError("Either state or batch_size must be provided") + return self.get_initial_state(batch_size) + + # Normalize existing state + state = self.norm(state) + + return state + + +class StateUpdateModule(nn.Module): + """ + Module for updating state based on new observations. + Implements the State-Update mechanism from CUT3R. + """ + + def __init__( + self, + state_dim: int = 512, + observation_dim: int = 512, + num_heads: int = 8, + dropout: float = 0.1, + ): + super().__init__() + + self.state_dim = state_dim + self.observation_dim = observation_dim + + # Project observations to state dimension if needed + if observation_dim != state_dim: + self.obs_projection = nn.Linear(observation_dim, state_dim) + else: + self.obs_projection = nn.Identity() + + # Multi-head attention for state update + self.cross_attention = nn.MultiheadAttention( + embed_dim=state_dim, num_heads=num_heads, dropout=dropout, batch_first=True + ) + + # Gating mechanism for state update + self.gate = nn.Sequential(nn.Linear(state_dim * 2, state_dim), nn.Sigmoid()) + + # Feedforward network + self.ffn = nn.Sequential( + nn.Linear(state_dim, state_dim * 4), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(state_dim * 4, state_dim), + ) + + # Layer normalization + self.norm1 = nn.LayerNorm(state_dim) + self.norm2 = nn.LayerNorm(state_dim) + + def forward(self, state: torch.Tensor, observation: torch.Tensor) -> torch.Tensor: + """ + Update state based on new observation. + + Args: + state: Current state (B, num_tokens, state_dim) + observation: New observation features (B, obs_len, obs_dim) + + Returns: + Updated state (B, num_tokens, state_dim) + """ + # Project observation to state dimension + obs_proj = self.obs_projection(observation) + + # Cross-attention: state attends to observation + attn_output, _ = self.cross_attention(query=state, key=obs_proj, value=obs_proj) + + # Gated update mechanism + gate_input = torch.cat([state, attn_output], dim=-1) + gate = self.gate(gate_input) + + # Apply gated update + state = state + gate * self.norm1(attn_output) + + # Feedforward network with residual connection + state = state + self.ffn(self.norm2(state)) + + return state + + +class StateReadoutModule(nn.Module): + """ + Module for reading information from state. + Implements the State-Readout mechanism from CUT3R. + """ + + def __init__( + self, + state_dim: int = 512, + query_dim: int = 512, + output_dim: int = 512, + num_heads: int = 8, + dropout: float = 0.1, + ): + super().__init__() + + self.state_dim = state_dim + self.query_dim = query_dim + self.output_dim = output_dim + + # Project query to state dimension if needed + if query_dim != state_dim: + self.query_projection = nn.Linear(query_dim, state_dim) + else: + self.query_projection = nn.Identity() + + # Multi-head attention for state readout + self.cross_attention = nn.MultiheadAttention( + embed_dim=state_dim, num_heads=num_heads, dropout=dropout, batch_first=True + ) + + # Output projection + self.output_projection = nn.Sequential( + nn.Linear(state_dim, output_dim), nn.LayerNorm(output_dim) + ) + + # Feedforward network + self.ffn = nn.Sequential( + nn.Linear(state_dim, state_dim * 4), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(state_dim * 4, state_dim), + ) + + # Layer normalization + self.norm1 = nn.LayerNorm(state_dim) + self.norm2 = nn.LayerNorm(state_dim) + + def forward(self, state: torch.Tensor, query: torch.Tensor) -> torch.Tensor: + """ + Read information from state based on query. + + Args: + state: Current state (B, num_tokens, state_dim) + query: Query features (B, query_len, query_dim) + + Returns: + Readout features (B, query_len, output_dim) + """ + # Project query to state dimension + query_proj = self.query_projection(query) + + # Cross-attention: query attends to state + attn_output, attention_weights = self.cross_attention( + query=query_proj, key=state, value=state + ) + + # Residual connection and normalization + readout = query_proj + self.norm1(attn_output) + + # Feedforward network with residual connection + readout = readout + self.ffn(self.norm2(readout)) + + # Project to output dimension + output = self.output_projection(readout) + + return output + + +class PersistentStateManager(nn.Module): + """ + Complete state management system combining initialization, + update, and readout mechanisms. + """ + + def __init__( + self, + num_tokens: int = 256, + token_dim: int = 512, + observation_dim: int = 512, + output_dim: int = 512, + num_heads: int = 8, + dropout: float = 0.1, + learnable_init: bool = True, + ): + super().__init__() + + # State representation + self.state_repr = StateRepresentation( + num_tokens=num_tokens, token_dim=token_dim + ) + + # State update module + self.state_update = StateUpdateModule( + state_dim=token_dim, + observation_dim=observation_dim, + num_heads=num_heads, + dropout=dropout, + ) + + # State readout module + self.state_readout = StateReadoutModule( + state_dim=token_dim, + query_dim=observation_dim, + output_dim=output_dim, + num_heads=num_heads, + dropout=dropout, + ) + + # Current state storage (for sequential processing) + self.register_buffer("current_state", None) + + def reset_state(self, batch_size: int): + """Reset state to initial values.""" + self.current_state = self.state_repr.get_initial_state(batch_size) + + def forward( + self, + observation: torch.Tensor, + query: Optional[torch.Tensor] = None, + state: Optional[torch.Tensor] = None, + return_updated_state: bool = True, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Process observation and update/read from state. + + Args: + observation: New observation features (B, obs_len, obs_dim) + query: Query for readout (B, query_len, obs_dim). + If None, use observation as query. + state: External state if provided. If None, use internal state. + return_updated_state: Whether to return the updated state + + Returns: + readout: Features read from state (B, query_len, output_dim) + updated_state: Updated state if return_updated_state=True + """ + batch_size = observation.shape[0] + + # Get current state + if state is not None: + current_state = state + elif self.current_state is not None: + current_state = self.current_state + else: + current_state = self.state_repr.get_initial_state(batch_size) + + # Update state with new observation + updated_state = self.state_update(current_state, observation) + + # Store updated state internally + self.current_state = updated_state + + # Read from updated state + if query is None: + query = observation + readout = self.state_readout(updated_state, query) + + if return_updated_state: + return readout, updated_state + else: + return readout + + +if __name__ == "__main__": + """Test the state management modules.""" + + print("Testing StateRepresentation...") + state_repr = StateRepresentation(num_tokens=256, token_dim=512) + initial_state = state_repr.get_initial_state(batch_size=4) + print(f"Initial state shape: {initial_state.shape}") # (4, 256, 512) + + print("\nTesting StateUpdateModule...") + state_update = StateUpdateModule(state_dim=512, observation_dim=256) + observation = torch.randn(4, 10, 256) # 10 observation tokens + updated_state = state_update(initial_state, observation) + print(f"Updated state shape: {updated_state.shape}") # (4, 256, 512) + + print("\nTesting StateReadoutModule...") + state_readout = StateReadoutModule(state_dim=512, query_dim=256, output_dim=512) + query = torch.randn(4, 5, 256) # 5 query tokens + readout = state_readout(updated_state, query) + print(f"Readout shape: {readout.shape}") # (4, 5, 512) + + print("\nTesting PersistentStateManager...") + manager = PersistentStateManager( + num_tokens=256, token_dim=512, observation_dim=256, output_dim=512 + ) + + # Process sequence + for t in range(3): + obs = torch.randn(4, 10, 256) + readout, state = manager(obs, return_updated_state=True) + print(f"Time {t}: Readout shape {readout.shape}, State shape {state.shape}") + + print("\nState management test successful!") diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/transform_layer.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/transform_layer.py new file mode 100644 index 0000000000..6ddcf558e3 --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/transform_layer.py @@ -0,0 +1,418 @@ +# 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. + +""" +Transform layer for registering previous frame to current frame using Kabsch algorithm. +Handles batched rigid body registration with known vertex correspondences. + +Supports two modes: +- Standard Kabsch: Non-differentiable SVD (requires STE for gradient flow) +- RoMa Kabsch: Fully differentiable (gradients flow through transformation) +""" + +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.nn.functional as F + +try: + import roma + + ROMA_AVAILABLE = True +except ImportError: + ROMA_AVAILABLE = False + roma = None + + +class TransformLayer(nn.Module): + """ + Registers previous frame positions to current frame positions using optimal rigid transformation. + Uses Kabsch algorithm for finding the best rotation and translation. + + Input shapes: + batch_last_position: (B, N, 3) - positions from previous frame + batch_current_position: (B, N, 3) - positions from current frame (target) + + Output shape: + (B, N, 3) - transformed previous positions aligned to current frame + """ + + def __init__(self, use_scale=False, regularize_rotation=False, if_roma=False): + """ + Args: + use_scale: If True, use Umeyama algorithm (similarity transform with scale). + If False, use Kabsch algorithm (rigid transform only). + regularize_rotation: If True, add small regularization to avoid numerical issues. + if_roma: If True, use RoMa's differentiable rigid_points_registration. + If False, use standard SVD-based Kabsch (non-differentiable). + """ + super().__init__() + self.use_scale = use_scale + self.regularize_rotation = regularize_rotation + self.if_roma = if_roma + + if self.if_roma: + if not ROMA_AVAILABLE: + raise ImportError( + "RoMa library is not installed. Install with: pip install roma" + ) + if self.use_scale: + raise ValueError( + "RoMa mode does not support scale (Umeyama). Set use_scale=False." + ) + + def forward(self, batch_last_position, batch_current_position): + """ + Register batch_last_position to batch_current_position. + + Args: + batch_last_position: (B, N, 3) - source positions to be transformed + batch_current_position: (B, N, 3) - target positions to align to + + Returns: + transformed_positions: (B, N, 3) - batch_last_position after rigid transformation + """ + B, N, D = batch_last_position.shape + assert D == 3, f"Expected 3D points, got {D}D" + assert batch_last_position.shape == batch_current_position.shape + + # RoMa mode: fully differentiable Kabsch + if self.if_roma: + # roma.rigid_points_registration: X -> Y + # Returns R (B,3,3) and t_vec (B,3) such that: Y ≈ (R @ X.T).T + t + R, t_vec = roma.rigid_points_registration( + batch_last_position, batch_current_position # Source (X) # Target (Y) + ) + # Apply transformation: Y = X @ R.T + t + # R: (B,3,3), batch_last_position: (B,N,3) + # X @ R.T: (B,N,3) @ (B,3,3).T = (B,N,3) + transformed = torch.bmm(batch_last_position, R.transpose(1, 2)) # (B,N,3) + transformed = transformed + t_vec.unsqueeze(1) # (B,N,3) + (B,1,3) + return transformed + + # Compute centroids + centroid_last = batch_last_position.mean(dim=1, keepdim=True) # (B, 1, 3) + centroid_current = batch_current_position.mean(dim=1, keepdim=True) # (B, 1, 3) + + # Center the point clouds + last_centered = batch_last_position - centroid_last # (B, N, 3) + current_centered = batch_current_position - centroid_current # (B, N, 3) + + # Compute covariance matrix H = X^T @ Y + # where X is source (last) and Y is target (current) + H = torch.bmm(last_centered.transpose(1, 2), current_centered) # (B, 3, 3) + + # SVD decomposition with autocast disabled (SVD requires float32) + with torch.amp.autocast("cuda", enabled=False): + H_float32 = H.float() # Cast to float32 for SVD + U, S, Vt = torch.linalg.svd( + H_float32 + ) # U: (B, 3, 3), S: (B, 3), Vt: (B, 3, 3) + # Cast back to original dtype + U, S, Vt = U.to(H.dtype), S.to(H.dtype), Vt.to(H.dtype) + + # Compute rotation matrix R = V @ U^T + R = torch.bmm(Vt.transpose(1, 2), U.transpose(1, 2)) # (B, 3, 3) + + # Handle reflection case (ensure proper rotation with det(R) = 1) + # Compute determinant in float32 to avoid half precision issues + with torch.amp.autocast("cuda", enabled=False): + R_float32 = R.float() + det_R = torch.det(R_float32) # (B,) + det_R = det_R.to(R.dtype) + + # Create a diagonal matrix for flipping if needed + flip_matrix = ( + torch.eye(3, device=R.device).unsqueeze(0).repeat(B, 1, 1) + ) # (B, 3, 3) + need_flip = det_R < 0 # (B,) + + if need_flip.any(): + # Flip the sign of the last singular vector + flip_matrix[need_flip, 2, 2] = -1 + Vt_corrected = torch.bmm(flip_matrix, Vt) + R = torch.bmm(Vt_corrected.transpose(1, 2), U.transpose(1, 2)) + + # Apply regularization if requested (helps with numerical stability) + if self.regularize_rotation: + # Re-orthogonalize using another SVD + with torch.amp.autocast("cuda", enabled=False): + R_float32 = R.float() + U_reg, S_reg, Vt_reg = torch.linalg.svd(R_float32) + R = torch.bmm(U_reg, Vt_reg).to(R.dtype) + + # Compute scale if using Umeyama algorithm + if self.use_scale: + # Compute optimal scale + # s = trace(S) / trace(X^T @ X) + numerator = S.sum(dim=1) # (B,) + denominator = (last_centered**2).sum(dim=[1, 2]) # (B,) + scale = ( + (numerator / (denominator + 1e-8)).unsqueeze(1).unsqueeze(2) + ) # (B, 1, 1) + else: + scale = torch.ones(B, 1, 1, device=R.device) + + # Compute translation + # t = centroid_current - scale * R @ centroid_last + rotated_centroid = torch.bmm(R, centroid_last.transpose(1, 2)).transpose( + 1, 2 + ) # (B, 1, 3) + translation = centroid_current - scale * rotated_centroid # (B, 1, 3) + + # Apply transformation: y = scale * R @ x + t + transformed = ( + scale * torch.bmm(batch_last_position, R.transpose(1, 2)) + translation + ) + + return transformed + + +class DifferentiableTransformLayer(TransformLayer): + """ + Learnable transform layer that predicts rotation (6D) and translation. + + Modes: + - absolute: predict the full R/t mapping last -> current + - residual: reserved for composing a residual with rigid alignment (not implemented) + """ + + def __init__( + self, + use_scale=False, + hidden_dim=256, + use_6d_rotation=True, + learnable_mode="absolute", + return_params=True, + ): + super().__init__(use_scale=use_scale) + self.use_6d_rotation = use_6d_rotation + self.learnable_mode = learnable_mode # 'absolute' | 'residual' + self.return_params = return_params + + # Shallow PointNet over all vertices for each body (last/current) + # Encode centered coordinates with 1x1 convs + BN + ReLU, global max-pool + # last_feat, curr_feat each -> (B, 256) + self.pn_conv1 = nn.Conv1d(3, 64, 1) + self.pn_bn1 = nn.BatchNorm1d(64) + self.pn_conv2 = nn.Conv1d(64, 128, 1) + self.pn_bn2 = nn.BatchNorm1d(128) + self.pn_conv3 = nn.Conv1d(128, 256, 1) + self.pn_bn3 = nn.BatchNorm1d(256) + + # Head MLP maps fused global features (+ centroids for translation prior) to rot6d + t + # Input dim: 256 (last) + 256 (current) + 6 (centroid_last/current) + head_in_dim = 256 + 256 + 6 + out_dim = 9 # 6D rotation + 3D translation + self.head_mlp = nn.Sequential( + nn.Linear(head_in_dim, hidden_dim), + nn.ReLU(), + nn.Linear(hidden_dim, hidden_dim), + nn.ReLU(), + nn.Linear(hidden_dim, out_dim), + ) + + @staticmethod + def rotation_6d_to_matrix(rot6d: torch.Tensor) -> torch.Tensor: + a1 = rot6d[:, 0:3] + a2 = rot6d[:, 3:6] + b1 = F.normalize(a1, dim=1) + a2_proj = a2 - (b1 * (a2 * b1).sum(dim=1, keepdim=True)) + b2 = F.normalize(a2_proj, dim=1) + b3 = torch.cross(b1, b2, dim=1) + return torch.stack([b1, b2, b3], dim=2) # (B,3,3) + + def forward(self, batch_last_position, batch_current_position): + B, N, D = batch_last_position.shape + assert D == 3, f"Expected 3D points, got {D}D" + + # Compute centroids for stability and translation prior + centroid_last = batch_last_position.mean(dim=1) # (B,3) + centroid_current = batch_current_position.mean(dim=1) # (B,3) + + # Center the point clouds (invariant to global translation) + last_centered = batch_last_position - centroid_last.unsqueeze(1) # (B,N,3) + current_centered = batch_current_position - centroid_current.unsqueeze( + 1 + ) # (B,N,3) + + # PointNet feature extraction for both clouds + def encode_points(x_centered: torch.Tensor) -> torch.Tensor: + # x_centered: (B,N,3) -> (B,3,N) + x = x_centered.transpose(1, 2) + x = F.relu(self.pn_bn1(self.pn_conv1(x))) + x = F.relu(self.pn_bn2(self.pn_conv2(x))) + x = F.relu(self.pn_bn3(self.pn_conv3(x))) + # Global max pooling over points -> (B,256) + x = torch.amax(x, dim=2) + return x + + last_feat = encode_points(last_centered) + curr_feat = encode_points(current_centered) + + # Fuse global features and centroids + fused = torch.cat( + [last_feat, curr_feat, centroid_last, centroid_current], dim=1 + ) # (B, 518) + + # Predict rotation (6D) and translation (3D) + pred = self.head_mlp(fused) # (B,9) + rot6d = pred[:, :6] + t = pred[:, 6:9].unsqueeze(1) # (B,1,3) + R_pred = self.rotation_6d_to_matrix(rot6d) # (B,3,3) + + if self.learnable_mode == "absolute": + corrected = torch.bmm(batch_last_position, R_pred.transpose(1, 2)) + t + + return corrected, {"rot6d": rot6d, "trans": t.squeeze(1)} + + +class RobustTransformLayer(TransformLayer): + """ + Robust variant that downweights outliers via iteratively reweighted + Kabsch alignment (weights decay exponentially with residual distance). + """ + + def __init__(self, use_scale=False, num_iterations=3, outlier_threshold=0.1): + super().__init__(use_scale=use_scale) + self.num_iterations = num_iterations + self.outlier_threshold = outlier_threshold + + def forward(self, batch_last_position, batch_current_position): + B, N, _ = batch_last_position.shape + + # Initialize weights (all points equally weighted) + weights = torch.ones(B, N, 1, device=batch_last_position.device) + + transformed = batch_last_position.clone() + + for iteration in range(self.num_iterations): + # Weighted centroids + weights_norm = weights / (weights.sum(dim=1, keepdim=True) + 1e-8) + centroid_last = (transformed * weights_norm).sum(dim=1, keepdim=True) + centroid_current = (batch_current_position * weights_norm).sum( + dim=1, keepdim=True + ) + + # Center with weighted centroids + last_centered = transformed - centroid_last + current_centered = batch_current_position - centroid_current + + # Weighted covariance + last_weighted = last_centered * weights_norm.sqrt() + current_weighted = current_centered * weights_norm.sqrt() + H = torch.bmm(last_weighted.transpose(1, 2), current_weighted) + + # SVD and compute transformation with autocast disabled + with torch.amp.autocast("cuda", enabled=False): + H_float32 = H.float() + U, S, Vt = torch.linalg.svd(H_float32) + U, S, Vt = U.to(H.dtype), S.to(H.dtype), Vt.to(H.dtype) + + R = torch.bmm(Vt.transpose(1, 2), U.transpose(1, 2)) + + # Fix reflection + # Compute determinant in float32 to avoid half precision issues + with torch.amp.autocast("cuda", enabled=False): + R_float32 = R.float() + det_R = torch.det(R_float32) + det_R = det_R.to(R.dtype) + flip_matrix = torch.eye(3, device=R.device).unsqueeze(0).repeat(B, 1, 1) + need_flip = det_R < 0 + if need_flip.any(): + flip_matrix[need_flip, 2, 2] = -1 + Vt_corrected = torch.bmm(flip_matrix, Vt) + R = torch.bmm(Vt_corrected.transpose(1, 2), U.transpose(1, 2)) + + # Translation + translation = centroid_current - torch.bmm( + R, centroid_last.transpose(1, 2) + ).transpose(1, 2) + + # Apply transformation + transformed = ( + torch.bmm(batch_last_position, R.transpose(1, 2)) + translation + ) + + # Update weights based on distances (downweight outliers) + if iteration < self.num_iterations - 1: + distances = torch.norm( + transformed - batch_current_position, dim=2, keepdim=True + ) + weights = torch.exp(-distances / self.outlier_threshold) + + return transformed + + +class RoMaTransformLayer(nn.Module): + """ + Differentiable Kabsch registration using RoMa library. + + Key difference from TransformLayer: + - Fully differentiable (gradients flow through R and t) + - No need for STE (Straight-Through Estimator) + - Uses roma.rigid_points_registration() which handles SVD gradient + + Note: Requires 'roma' library to be installed. + """ + + def __init__(self): + super().__init__() + if not ROMA_AVAILABLE: + raise RuntimeError( + "RoMa library not available. Install with: pip install roma\n" + "RoMa is required for differentiable Kabsch registration." + ) + + def forward(self, batch_last_position, batch_current_position): + """ + Register batch_last_position to batch_current_position using RoMa. + + Args: + batch_last_position: (B, N, 3) - source positions to be transformed + batch_current_position: (B, N, 3) - target positions to align to + + Returns: + transformed_positions: (B, N, 3) - transformed positions (differentiable) + """ + B, N, D = batch_last_position.shape + assert D == 3, f"Expected 3D points, got {D}D" + assert batch_last_position.shape == batch_current_position.shape + + # RoMa's rigid_points_registration returns (R, t_vec) + # R: (B, 3, 3) rotation matrix + # t_vec: (B, 3) translation vector + # Formula: y = R @ x + t + R, t_vec = roma.rigid_points_registration( + batch_last_position, batch_current_position # source # target + ) + + # Apply transformation: y = x @ R^T + t + # (B, N, 3) @ (B, 3, 3) -> (B, N, 3) + transformed = torch.bmm( + batch_last_position, R.transpose(1, 2) + ) + t_vec.unsqueeze(1) + + return transformed + + +# Convenience function for creating the layer +def create_transform_layer(method="kabsch", if_roma=False, **kwargs): + """ + Factory function to create transform layer. + + Args: + method: Only 'kabsch' is supported + if_roma: If True, use RoMa's differentiable Kabsch; otherwise use standard SVD Kabsch + **kwargs: Additional arguments (e.g., regularize_rotation) + + Returns: + TransformLayer instance with Kabsch algorithm + """ + if method != "kabsch": + raise ValueError(f"Only method='kabsch' is supported, got '{method}'") + + return TransformLayer(use_scale=False, if_roma=if_roma, **kwargs) diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/network.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/network.py new file mode 100644 index 0000000000..b36b2d9d45 --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/network.py @@ -0,0 +1,3148 @@ +# 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. + +"""Author model composition, extracted from src/train.py; no training CLI imports.""" + +from __future__ import annotations + +import os +import math +import logging +import numpy as np +import torch +from torch import nn +from torch.nn import functional as F +from .models.encoder import RigidBodyEncoder, EncoderPack +from .models.state import PersistentStateManager +from .models.decoder import BidirectionalDecoder +from .models.predictor import DynamicsPredictor +from .models.transform_layer import create_transform_layer +from .models.physformer_rope import rope_memo_reset +from .utils.intersection_handler import handle_intersections +from .utils.cuda_knn import ( + compute_nearest_distance_to_list_cuda, + compute_knn_all_objects_vectorized, + knn_pack, + compute_knn_packed, +) + +try: + import roma + + ROMA_AVAILABLE = True +except ImportError: + roma = None + ROMA_AVAILABLE = False +logger = logging.getLogger(__name__) +DEBUG_NAN = os.environ.get("DEBUG_NAN", "0") == "1" + + +def distance_2_plane(points, plane_point, plane_axis, body_dist_vec): + """ + Check if points are closer to a plane than to another body, and update distance vectors. + + Args: + points: Tensor of shape (B, M, 3) containing point positions + plane_point: List [x, y, z] defining a point on the plane (e.g., [0,0,0]) + plane_axis: String "x", "y", or "z" defining the plane normal direction + body_dist_vec: Tensor of shape (B, M, 3) containing vectors to nearest points on other body + + Returns: + Updated distance vectors (B, M, 3) where vectors to plane replace body vectors when closer + """ + B, M, _ = points.shape + device = points.device + + # Get the axis index + axis_map = {"x": 0, "y": 1, "z": 2} + if plane_axis not in axis_map: + raise ValueError(f"Invalid plane_axis: {plane_axis}. Must be 'x', 'y', or 'z'") + axis_idx = axis_map[plane_axis] + + # Get plane position on the specified axis + plane_pos = plane_point[axis_idx] + + # Calculate distance to plane for each point + # Distance is just the absolute difference in the specified axis + point_axis_coords = points[:, :, axis_idx] # (B, M) + dist_to_plane = torch.abs(point_axis_coords - plane_pos) # (B, M) + + # Calculate distance from body vectors (norm of the vectors) + dist_to_body = torch.norm(body_dist_vec, dim=2) # (B, M) + + # Create vectors from points to plane + # The vector points from the point to its projection on the plane + plane_vectors = torch.zeros_like(points) # (B, M, 3) + plane_vectors[:, :, axis_idx] = ( + plane_pos - point_axis_coords + ) # Only non-zero in the plane normal direction + + # Create mask for points closer to plane than to body + closer_to_plane = dist_to_plane < dist_to_body # (B, M) + + # Expand mask to match vector dimensions + closer_to_plane_expanded = closer_to_plane.unsqueeze(-1).expand_as( + body_dist_vec + ) # (B, M, 3) + + # Update vectors: use plane vectors where closer to plane, otherwise keep body vectors + updated_vectors = torch.where( + closer_to_plane_expanded, plane_vectors, body_dist_vec + ) + + return updated_vectors + + +def _object_centered(ref_pose: torch.Tensor, vmask: torch.Tensor) -> torch.Tensor: + """ref_pose (B, Nobj, Nv, 3) minus each object's x/y centroid over valid vertices vmask (B, Nobj, Nv). + + Only x and y are centred: horizontal translation is a symmetry of the dynamics, height above + the ground plane (z = 0) is not, so z is kept absolute.""" + m = vmask.to(ref_pose.dtype).unsqueeze(-1) + c = (ref_pose * m).sum(dim=2, keepdim=True) / m.sum(dim=2, keepdim=True).clamp_min( + 1.0 + ) + c = torch.cat([c[..., :2], torch.zeros_like(c[..., 2:3])], dim=-1) + return ref_pose - c + + +class CUT3RDynamicsModel(nn.Module): + """ + Complete model combining encoder, state management, decoder, and predictor. + """ + + def __init__(self, config): + super().__init__() + + # Store config for later use + self.config = config + + # Set device + self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # Physics-parameter and time-step flags (read before submodule construction) + # Generalization arm: anchor reference poses expressed relative to each object's + # frame-0 centroid (removes absolute translation from the predictor query), and + # rotary pairing that matches rotate_half so attention depends on displacement only. + self._ref_pose_object_centered = bool( + config["model"].get("ref_pose_object_centered", False) + ) + from embodichain_models.dynamics.rigidformer._native.models.physformer_rope import ( + PhysFormerRoPE as _RopeCls, + ) + + _RopeCls.INTERLEAVED_PAIRING = bool( + config["model"].get("rope_interleaved_pairing", False) + ) + + self._if_phys_param = config.get("model", {}).get("if_phys_param", False) + self._if_time_step = config.get("model", {}).get("if_time_step", False) + self._film_conditioning = config.get("model", {}).get( + "film_conditioning", False + ) + if self._if_phys_param: + print( + f"[Model] Physics parameters enabled: mass, friction, restitution (+3D to encoder/predictor)" + ) + if self._if_time_step: + if self._film_conditioning: + print( + f"[Model] Time step enabled via FiLM only (NOT concatenated to encoder input)" + ) + else: + print( + f"[Model] Time step enabled: skip_step_size (s and s², +2D to encoder/predictor)" + ) + + # Rollout state reset: reset state tokens at each step to avoid train-test mismatch + # When training with short sequences (e.g., 4 frames), state tokens only see 2 updates. + # During long rollouts (e.g., 400 frames), state drift causes OOD behavior. + # Setting this to True resets state tokens at each rollout step to match training distribution. + self._rollout_state_reset = config.get("model", {}).get( + "rollout_state_reset", False + ) + if self._rollout_state_reset: + print( + f"[Model] Rollout state reset ENABLED: state tokens will be reset at each rollout step" + ) + + # KNN acceleration: GPU-parallel grouped processing + self._knn_num_groups = config.get("data", {}).get("knn_num_groups", 1) + # Valid-vertex packing (model.knn_packed, model.encoder.packed): the distance features and the encoder's + # pointwise layers run over the valid vertices only instead of every padded (object slot x vertex) entry. + # Equal to the padded paths up to floating-point rounding; the packs are built once per batch / rollout. + self._knn_packed = bool(config.get("model", {}).get("knn_packed", False)) + self._encoder_packed = bool( + config.get("model", {}).get("encoder", {}).get("packed", False) + ) + if self._knn_packed or self._encoder_packed: + print( + f"[Model] valid-vertex packing: knn_packed={self._knn_packed}, encoder.packed={self._encoder_packed}" + ) + if self._knn_packed and self._knn_num_groups > 1: + raise ValueError( + "model.knn_packed replaces data.knn_num_groups; set knn_num_groups to 1" + ) + # model.encoder.packed_pooling: 'segment' pools the packed rows directly (segment max / mean) instead of + # scattering every level back to the padded layout; same values up to fp32 summation order in the mean. + if ( + config.get("model", {}).get("encoder", {}).get("packed_pooling", "scatter") + != "scatter" + and not self._encoder_packed + ): + raise ValueError( + "model.encoder.packed_pooling applies to the packed encoder path; set model.encoder.packed" + ) + if self._knn_num_groups > 1: + print( + f"[Model] KNN acceleration enabled: num_groups={self._knn_num_groups}" + ) + + # model.kabsch_reference: the mesh the per-step Kabsch fit transforms. 'advancing' (default) = the + # previous step's result, so every step re-transforms an already transformed mesh and rounding accumulates; + # 'fixed' = one unchanged reference shape for the whole sequence / rollout (the first current frame), equal + # in exact arithmetic (the fit is equivariant) without the accumulation. + _kref = config.get("model", {}).get("kabsch_reference", "advancing") + if _kref not in ("fixed", "advancing"): + raise ValueError( + f"model.kabsch_reference must be 'fixed' or 'advancing', got {_kref!r}" + ) + self._kabsch_reference_fixed = _kref == "fixed" + if self._kabsch_reference_fixed: + print( + "[Model] Kabsch reference FIXED: every step fits the same reference mesh" + ) + + # model.intersection_handling.after_kabsch_loss: which state the position-after-Kabsch loss supervises. + # 'rigid' = the Kabsch output (the state the acceleration-after-Kabsch loss already uses), 'handled' + # (default) = the state after the ground/object correction. The rollout state is the corrected one either way. + _akl = ( + config.get("model", {}) + .get("intersection_handling", {}) + .get("after_kabsch_loss", "handled") + ) + if _akl not in ("rigid", "handled"): + raise ValueError( + f"model.intersection_handling.after_kabsch_loss must be 'rigid' or 'handled', got {_akl!r}" + ) + self._after_kabsch_loss_rigid = _akl == "rigid" + if self._after_kabsch_loss_rigid: + print( + "[Model] position-after-Kabsch loss supervises the Kabsch output (before intersection handling)" + ) + + # Time-batched training forward (training.time_batched_forward): the network part of all teacher-forced + # steps runs as one batch (see _forward_multi_object). Only valid when no step's network input depends on + # an earlier step's prediction or state; the settings that would create such a dependency are refused. + self._time_batched_forward = bool( + config.get("training", {}).get("time_batched_forward", False) + ) + if self._time_batched_forward: + _enc = config.get("model", {}).get("encoder", {}) + _bad = [] + if config.get("model", {}).get( + "if_penetration", + config.get("model", {}).get("if_directional_distance", False), + ): + _bad.append( + "model.if_penetration/if_directional_distance (uses the previous step's decoder output)" + ) + if not _enc.get("masked_pooling", False): + _bad.append( + "model.encoder.masked_pooling=false (BatchNorm statistics would pool across steps)" + ) + if _enc.get("use_stn", False) or _enc.get("feature_transform", False): + _bad.append( + "encoder STN / feature transform (their BatchNorms would pool across steps)" + ) + if _bad: + raise ValueError( + "training.time_batched_forward is incompatible with: " + + "; ".join(_bad) + ) + print( + "[Model] time-batched training forward ACTIVE (network part of all teacher-forced steps in one " + "batch; per-step BatchNorm statistics)", + flush=True, + ) + + # Debug mode: from data.random_step_size.if_debug or training.loss_weights.if_debug + self._debug_mode = config.get("data", {}).get("random_step_size", {}).get( + "if_debug", False + ) or config.get("training", {}).get("loss_weights", {}).get("if_debug", False) + + # Config Consistency Check: Multi-timescale training requires time_step feature + random_step_size_config = config.get("data", {}).get("random_step_size", {}) + if random_step_size_config.get("enable", False): + if not self._if_time_step: + raise ValueError( + "Config inconsistency detected!\n" + " data.random_step_size.enable = True (multi-timescale training enabled)\n" + " model.if_time_step = False (time step feature disabled)\n" + "\n" + "Multi-timescale training requires the network to receive step_size as input.\n" + "Please set model.if_time_step = true in your config file." + ) + print( + f"[Model] Multi-timescale training enabled: step_size will be used as condition feature" + ) + + # Initialize components (multi-object only) + + # ============================================================ + # ARCHITECTURAL INVARIANTS — Translation Consistency & Permutation Equivariance + # These properties are fundamental to the model. Enforce at init time. + # ============================================================ + + # 1. State tokens permanently removed — fixed 3D grid breaks translation consistency + self._no_global_state = True + self.state_repr = None + assert ( + config["model"].get("no_global_state", True) is True + ), "no_global_state must be True: state tokens on a fixed 3D grid break translation consistency" + + # 2. rope_set_aggregation must be True — concat mode is anchor-order sensitive + assert config["model"]["decoder"].get("rope_set_aggregation", False) is True, ( + "rope_set_aggregation must be True: the concatenation mode is anchor-order sensitive, " + "breaking permutation invariance of ARoPE within each object" + ) + + # 3. No sequence-index positional embeddings — would break object permutation equivariance + assert not config["model"]["decoder"].get( + "use_learnable_pe", False + ), "use_learnable_pe must be False: sequence-index PE breaks object permutation equivariance" + + # ============================================================ + + # Multi-scale predictor reads features from selected decoder layers (CUT3R-style) + from embodichain_models.dynamics.rigidformer._native.models.multi_scale_predictor import ( + QueryBasedMultiScalePredictor, + ) + + # Adjust layer indices based on actual number of decoder layers (needed before decoder init) + num_decoder_layers = config["model"]["decoder"].get("num_layers", 12) + if num_decoder_layers >= 12: + layer_indices = [0, 6, 9, 12] # CUT3R's layers for 12-layer model + elif num_decoder_layers >= 6: + layer_indices = [0, 2, 4, num_decoder_layers] # For 6+ layer model + elif num_decoder_layers >= 4: + # For 4-layer model + layer_indices = [0, 1, 2, num_decoder_layers] # Use all available layers + else: + # For 2-layer model (tiny), use all available layers + layer_indices = list( + range(num_decoder_layers + 1) + ) # [0, 1, 2] for 2 layers + # model.predictor.layer_indices overrides the rule above, e.g. [0, 1, 2, 8] for a decoder grown from 4 to 8 + # layers (scripts/grow_decoder.py) that reads the same layers as the 4-layer model + if config["model"]["predictor"].get("layer_indices") is not None: + layer_indices = [ + int(i) for i in config["model"]["predictor"]["layer_indices"] + ] + if ( + sorted(set(layer_indices)) != layer_indices + or layer_indices[-1] != num_decoder_layers + ): + raise ValueError( + f"model.predictor.layer_indices must be increasing and end with the last decoder layer " + f"({num_decoder_layers}); got {layer_indices}" + ) + + # Verify that all layer indices are valid for the decoder + max_valid_index = ( + num_decoder_layers # 0 to num_decoder_layers inclusive (initial + N layers) + ) + for idx in layer_indices: + assert 0 <= idx <= max_valid_index, ( + f"Invalid layer index {idx} for decoder with {num_decoder_layers} layers. " + f"Valid range: [0, {max_valid_index}]" + ) + + # Store layer indices for decoder to only save needed layers + self.predictor_layer_indices = set(layer_indices) # Use set for O(1) lookup + + # Create mapping from original layer index to position in saved list + # For example: if layer_indices = [0, 1, 2, 4], then: + # layer_index_mapping = {0: 0, 1: 1, 2: 2, 4: 3} + self.layer_index_mapping = { + idx: i for i, idx in enumerate(sorted(layer_indices)) + } + + # Gated Attention configuration (Qwen3-style) + gated_attention_config = config["model"]["decoder"].get("gated_attention", {}) + gated_attention_enable = gated_attention_config.get("enable", False) + gated_attention_type = gated_attention_config.get("type", "elementwise") + gated_attention_debug = gated_attention_config.get("if_debug", False) + # Currently only support elementwise gate + if gated_attention_enable and gated_attention_type != "elementwise": + raise ValueError( + f"Unsupported gated_attention type: {gated_attention_type}. Only 'elementwise' is supported." + ) + self._elementwise_gate = ( + gated_attention_enable and gated_attention_type == "elementwise" + ) + + # Debug print for gated attention (print to console since FileHandler not yet set) + if gated_attention_debug: + msg = ( + f"[Gated Attention] enable={gated_attention_enable}, type={gated_attention_type}, " + f"elementwise_gate={self._elementwise_gate}" + ) + print(msg) # Console output (FileHandler not ready at model init time) + logger.info( + msg + ) # Also log (will appear in file after Trainer sets FileHandler) + + # State update mode: when False, state is just learnable embeddings (no bidirectional update) + # Use with rollout_state_reset=True to avoid train-test mismatch + self._state_update_enabled = config["model"]["decoder"].get( + "state_update_enabled", True + ) + + # Objects Self-Attention: Enable direct object-to-object interaction + # Useful for physics simulation where objects need to exchange information + self._objects_self_attention = config["model"]["decoder"].get( + "objects_self_attention", False + ) + + # PE type must be read before decoder init (use_se3 flag needs it) + self._pe_type = config["model"]["decoder"].get("pe_type", "arope") + + # Initialize decoder with layer_indices_to_save for memory efficiency + # Block AttnRes config (Kimi 2026) + _block_attn_res_cfg = config["model"]["decoder"].get("block_attn_res", {}) + _block_attn_res_enable = ( + _block_attn_res_cfg.get("enable", False) + if isinstance(_block_attn_res_cfg, dict) + else bool(_block_attn_res_cfg) + ) + _block_attn_res_block_size = ( + _block_attn_res_cfg.get("block_size", 4) + if isinstance(_block_attn_res_cfg, dict) + else 4 + ) + _block_attn_res_final = ( + bool(_block_attn_res_cfg.get("final_aggregate", False)) + if isinstance(_block_attn_res_cfg, dict) + else False + ) + + self.decoder = BidirectionalDecoder( + state_dim=config["model"]["state"].get("token_dim", 768), + num_layers=config["model"]["decoder"].get( + "num_layers", 12 + ), # CUT3R uses 12 + num_heads=config["model"]["decoder"].get("num_heads", 12), + feedforward_dim=config["model"]["decoder"].get("feedforward_dim", 3072), + dropout=config["model"]["decoder"]["dropout"], + return_all_layers=True, # For multi-scale fusion + use_self_attention=config["model"]["decoder"].get( + "use_self_attention", False + ), # Optional state SA + center_rope_positions=config["model"].get("rope_center_positions", False), + rope_set_aggregation=config["model"]["decoder"].get( + "rope_set_aggregation", False + ), # Set aggregation for RoPE + rope_dim=config["model"]["decoder"].get( + "rope_dim", None + ), # Custom RoPE dimension (None = auto-compute) + dim_per_coord_3d=config["model"]["decoder"].get( + "dim_per_coord_3d", None + ), # Dim per coordinate for single-point (3D) RoPE + dim_per_coord_multi=config["model"]["decoder"].get( + "dim_per_coord_multi", None + ), # Dim per coordinate for multi-point RoPE + layer_indices_to_save=self.predictor_layer_indices, # Only save layers needed by predictor + elementwise_gate=self._elementwise_gate, # Qwen3: elementwise gated attention + state_update_enabled=self._state_update_enabled, # When False, state is learnable embeddings only + objects_self_attention=self._objects_self_attention, # Enable Objects Self-Attention for object-object interaction + use_se3=( + self._pe_type == "se3" + ), # SE(3) pairwise bias for obj self-attention + block_attn_res=_block_attn_res_enable, # Block Attention Residuals (Kimi 2026) + block_attn_res_block_size=_block_attn_res_block_size, # Layers per block + block_attn_res_final_aggregate=_block_attn_res_final, # final depth aggregation as the last features + film_conditioning=( + config["model"]["decoder"].get("film_conditioning", False) + and self._film_conditioning + and self._if_time_step + ), # Per-layer FiLM: requires decoder.film_conditioning + model.film_conditioning + model.if_time_step + ) + + # ============ Anchor-Vertex Pooling (AVP) ============ + avp_cfg = config.get("model", {}).get("avp", {}) + self._avp_enabled = avp_cfg.get("enable", False) + self._avp_output_dim = avp_cfg.get("output_dim", 256) + if self._avp_enabled: + from embodichain_models.dynamics.rigidformer._native.models.anchor_vertex_pooling import ( + AnchorVertexPooling, + ) + + _avp_vertex_dim = ( + config.get("model", {}) + .get("encoder", {}) + .get( + "backbone_dim", + config.get("model", {}).get("encoder", {}).get("output_dim", 1024), + ) + ) + self.avp = AnchorVertexPooling( + vertex_dim=_avp_vertex_dim, + output_dim=self._avp_output_dim, + init_sigma=avp_cfg.get("init_sigma", 0.1), + ) + print( + f"[Model] AVP ENABLED: vertex_dim={_avp_vertex_dim}, output_dim={self._avp_output_dim}, " + f"init_sigma={avp_cfg.get('init_sigma', 0.1)}" + ) + else: + self.avp = None + + # Dynamic predictor input_dim calculation: + # Base: 12D (9D query: dist_vec+vel+rel_pos + 3D ref_pose) + # + phys_params (if if_phys_param=true): +3D = 15D + # + time_step (if if_time_step=true AND NOT film): +2D = 17D + # + AVP (if avp.enable=true): +avp_output_dim + predictor_input_dim = 12 # Base: 9D query + 3D ref_pose + if self._if_phys_param: + predictor_input_dim += 3 # Physics params + if self._if_time_step and not self._film_conditioning: + predictor_input_dim += ( + 2 # Time step (s and s²) — skipped when FiLM handles it + ) + if self._avp_enabled: + predictor_input_dim += self._avp_output_dim # AVP features per anchor + print( + f"[Model] Predictor input_dim: {predictor_input_dim} " + f"(base 12 + phys {3 if self._if_phys_param else 0} + AVP {self._avp_output_dim})" + ) + + self.predictor = QueryBasedMultiScalePredictor( + state_dim=config["model"]["state"].get("token_dim", 768), + query_dim=predictor_input_dim, # Dynamically computed based on flags + hidden_dim=config["model"]["predictor"].get("hidden_dims", [512, 256])[0], + output_dim=config["model"]["predictor"]["output_dim"], + layer_indices=layer_indices, + layer_index_mapping=self.layer_index_mapping, # Pass the mapping for memory-optimized decoder + num_heads=config["model"]["predictor"].get( + "num_heads", config["model"]["decoder"].get("num_heads", 8) + ), # Use predictor's num_heads if specified + dropout=config["model"]["predictor"]["dropout"], + use_ref_pose=config["model"]["predictor"].get( + "use_ref_pose", True + ), # Enable ref pose + norm_type=config["model"]["predictor"].get("norm_type", "layer_norm"), + qk_norm=config["model"]["predictor"].get("qk_norm", False), + fusion_activation="swiglu", + use_rope=config["model"]["predictor"].get( + "use_rope", True + ), # Enable RoPE for position-aware attention + rope_dim=config["model"]["decoder"].get( + "rope_dim", None + ), # Use same RoPE dim as decoder + dim_per_coord_3d=config["model"]["decoder"].get( + "dim_per_coord_3d", None + ), # Single-Point RoPE config + dim_per_coord_multi=config["model"]["decoder"].get( + "dim_per_coord_multi", None + ), # Multi-Point RoPE config + elementwise_gate=self._elementwise_gate, # Pass elementwise_gate config + cross_object_attention=config["model"]["predictor"].get( + "cross_object_attention", False + ), # Cross-object anchor attention + cross_object_kv=config["model"]["predictor"].get( + "cross_object_kv", False + ), # Cross-object key-value attention + no_global_state=config["model"].get( + "no_global_state", False + ), # Removes unused state input channels + ) + # model.predictor.rmsnorm_eps: training_dtype -> the predictor's RMSNorm modules keep the eps they have in + # bf16-autocast training also when evaluated in fp32 (their default eps follows the input dtype; see + # pin_rmsnorm_eps). Applied here, so every construction of the model computes the configured function. + _rms_eps_mode = config["model"]["predictor"].get("rmsnorm_eps") + if _rms_eps_mode == "training_dtype": + _n_pinned = self.predictor.pin_rmsnorm_eps() + logger.info( + f"[Predictor] RMSNorm eps pinned to the bf16-training values on {_n_pinned} modules" + ) + elif _rms_eps_mode not in (None, "default"): + raise ValueError( + "model.predictor.rmsnorm_eps must be 'training_dtype' or 'default'" + ) + # model.predictor.fp32_output: the acceleration read-out layer in fp32 under bf16 training autocast (evaluation + # runs it in fp32 already, so evaluation is unchanged) + self.predictor.fp32_output = bool( + config["model"]["predictor"].get("fp32_output", False) + ) + if self.predictor.fp32_output: + logger.info( + "[Predictor] acceleration read-out (fusion_out) in fp32 during training" + ) + + # Store attention flags for reference + self._cross_object_attention = config["model"]["predictor"].get( + "cross_object_attention", False + ) + self._cross_object_kv = config["model"]["predictor"].get( + "cross_object_kv", False + ) + + # Log attention configuration summary + logger.info("=" * 60) + logger.info("[Attention Config Summary]") + logger.info( + f" Decoder Objects Self-Attention: {'ENABLED' if self._objects_self_attention else 'DISABLED'}" + ) + logger.info( + f" Predictor Cross-Object Anchor Attention: {'ENABLED' if self._cross_object_attention else 'DISABLED'}" + ) + logger.info( + f" Predictor Cross-Object Key-Value: {'ENABLED' if self._cross_object_kv else 'DISABLED'}" + ) + logger.info("=" * 60) + + # Store actual dt for video/visualization + self.dt = config["data"].get("dt", 0.01) # Actual timestep (0.01 = 100 FPS) + # Verlet integration always uses dt=1; the network learns accelerations in this unit system. + self.dt_verlet = 1.0 + + # Create transform layer (configurable) + transform_cfg = config.get("model", {}).get("transform", {}) + transform_method = transform_cfg.get("method", "kabsch") + self.learnable_transform = transform_method == "learnable" + + # RoMa differentiable Kabsch (if enabled, no STE needed) + self.if_roma = bool(transform_cfg.get("if_roma", False)) + if self.if_roma and not ROMA_AVAILABLE: + raise ImportError( + "RoMa is required when if_roma=true. Install with: pip install roma" + ) + + # Build kwargs per method to avoid passing unsupported args + if self.learnable_transform: + tl_kwargs = { + "use_6d_rotation": transform_cfg.get("use_6d_rotation", True), + "learnable_mode": transform_cfg.get("learnable_mode", "absolute"), + } + else: + tl_kwargs = { + "regularize_rotation": transform_cfg.get("regularize_rotation", True) + } + self.transform_layer = create_transform_layer( + method=transform_method, if_roma=self.if_roma, **tl_kwargs + ) + + # ============== Multi-object (MOVi) specific modules ============== + self.enable_multi_object = bool( + self.config.get("data", {}).get("multi_object", False) + ) + if self.enable_multi_object: + # Cache for mesh faces (dataset_name:obj_idx -> faces array) + self._faces_cache = {} + + # Shared PointNet encoder for arbitrary number of objects (weights shared) + enc_cfg = config.get("model", {}).get("encoder", {}) + # Dynamic encoder input_dim calculation: + # Base: 9D (dist_vec + velocity + rel_pos) + # + phys_params (if if_phys_param=true): +3D + # + time_step (if if_time_step=true and no FiLM): +2D (s and s²) + encoder_input_dim = 9 # Base + if self._if_phys_param: + encoder_input_dim += 3 # Physics params + if self._if_time_step and not self._film_conditioning: + encoder_input_dim += ( + 2 # Time step (s and s²) — skipped when FiLM handles it + ) + if enc_cfg.get("height_input", False): + encoder_input_dim += 1 # Height above the ground plane (absolute z) + self._encoder_input_dim = encoder_input_dim # Save for logging + + self.shared_obj_encoder = RigidBodyEncoder( + input_dim=encoder_input_dim, + hidden_dims=enc_cfg.get("hidden_dims", [512, 1024]), + output_dim=enc_cfg.get("backbone_dim", enc_cfg.get("output_dim", 896)), + backbone_dim=enc_cfg.get( + "backbone_dim", enc_cfg.get("output_dim", 896) + ), + use_stn=enc_cfg.get("use_stn", False), + use_batch_norm=enc_cfg.get("use_batch_norm", True), + masked_pooling=enc_cfg.get("masked_pooling", False), + freeze_bn_stats=enc_cfg.get("freeze_bn_stats", False), + height_input=enc_cfg.get("height_input", False), + hierarchical_subsets=enc_cfg.get("hierarchical_subsets", False), + packed_pooling=enc_cfg.get("packed_pooling", "scatter"), + subset_small_objects=enc_cfg.get("subset_small_objects", "all_below_8"), + feature_transform=enc_cfg.get("feature_transform", False), + feature_transform_dim=enc_cfg.get("feature_transform_dim", 64), + pooling_type=enc_cfg.get( + "pooling_type", "max" + ), # New: enhanced pooling + num_hierarchical_levels=enc_cfg.get( + "num_hierarchical_levels", 2 + ), # New: hierarchical levels + ) + # Project per-object latent to model token dim + token_dim = config["model"]["state"].get("token_dim", 768) + enc_out_dim = enc_cfg.get("backbone_dim", enc_cfg.get("output_dim", 896)) + # Project PointNet encoder output (per-object) to token dimension for Decoder + self.obj_feature_proj = nn.Linear(enc_out_dim, token_dim) + # Anchor query dimension is always 9D (dist_vec + vel + rel_pos) + anchor_query_dim = 9 + # Projection for anchor query features (for Predictor) + self.anchor_query_proj = nn.Linear(anchor_query_dim, token_dim) + + # FiLM conditioning: modulate encoder output with [s, s²] instead of + # concatenating step size to the input features + if self._film_conditioning and self._if_time_step: + from embodichain_models.dynamics.rigidformer._native.models.encoder import ( + StepSizeFiLM, + ) + + self.encoder_film = StepSizeFiLM(time_dim=2, feature_dim=token_dim) + print( + f"[Model] FiLM conditioning ENABLED: step-size modulates encoder output (dim={token_dim})" + ) + print( + f"[Model] Encoder input: {encoder_input_dim}D (time NOT concatenated, handled by FiLM)" + ) + print( + f"[Model] Predictor query: {predictor_input_dim}D (time NOT in query)" + ) + else: + self.encoder_film = None + if self._film_conditioning and not self._if_time_step: + print( + f"[Model] WARNING: film_conditioning=True but if_time_step=False; FiLM disabled" + ) + + # ============== PE Type for ARoPE Ablation ============== + # pe_type controls what 12D positional descriptor is fed to the decoder's RoPE path. + # "arope" (default): original 4 anchors × 3D = 12D + # "obb": OBB center + 3 half-axis vectors = 12D + # "pca": PCA centroid + √λ·eigenvectors = 12D + self._pe_type = config["model"]["decoder"].get("pe_type", "arope") + _VALID_PE_TYPES = ("arope", "obb", "pca", "se3") + if self._pe_type not in _VALID_PE_TYPES: + raise ValueError( + f"Invalid pe_type='{self._pe_type}'. Must be one of {_VALID_PE_TYPES}." + ) + _pe_msg = f"[PE Type] Using pe_type='{self._pe_type}' for decoder/predictor positional encoding." + print(_pe_msg) + logger.info(_pe_msg) + + self._geometry_encoder = None + if self._pe_type in ("obb", "pca"): + from embodichain_models.dynamics.rigidformer._native.models.geometry_encoders import ( + OBBEncoder, + PCAFrameEncoder, + ) + + if self._pe_type == "obb": + self._geometry_encoder = OBBEncoder() + logger.info( + "[PE Type] Initialized OBBEncoder (center + 3 half-axis vectors -> 12D)" + ) + elif self._pe_type == "pca": + self._geometry_encoder = PCAFrameEncoder() + logger.info( + "[PE Type] Initialized PCAFrameEncoder (centroid + scaled eigenvectors -> 12D)" + ) + + # SE(3) pairwise bias module + self._se3_bias_module = None + if self._pe_type == "se3": + from embodichain_models.dynamics.rigidformer._native.models.se3_attention import ( + SE3PairwiseBias, + ) + + _num_heads = config["model"]["decoder"].get("num_heads", 6) + _se3_hidden = config["model"]["decoder"].get("se3_hidden_dim", 64) + self._se3_bias_module = SE3PairwiseBias( + num_heads=_num_heads, hidden_dim=_se3_hidden, use_scale_ratio=True + ) + logger.info( + f"[PE Type] Initialized SE3PairwiseBias (7D -> {_se3_hidden} -> {_num_heads} heads, zero-init, scale-norm)" + ) + + # ============== RoPE frequency configuration summary ============== + _rope_dim = config["model"]["decoder"].get("rope_dim", None) + _dim_per_coord_3d = config["model"]["decoder"].get("dim_per_coord_3d", None) + _dim_per_coord_multi = config["model"]["decoder"].get( + "dim_per_coord_multi", None + ) + _hidden_dim = config["model"]["decoder"].get("hidden_dim", 768) + _num_heads = config["model"]["decoder"].get("num_heads", 6) + _head_dim = _hidden_dim // _num_heads if _num_heads > 0 else 0 + _anchor_k = config.get("data", {}).get("anchor_k", 4) + _n_coords = _anchor_k * 3 + _rope_msg = ( + f"[RoPE Freq] rope_dim={_rope_dim} | " + f"dim_per_coord_multi={_dim_per_coord_multi} (freq/coord) | " + f"dim_per_coord_3d={_dim_per_coord_3d} | " + f"hidden_dim={_hidden_dim} | num_heads={_num_heads} | head_dim={_head_dim} | " + f"anchor_k={_anchor_k} | n_coords={_n_coords} | " + f"check: {_n_coords}x{_dim_per_coord_multi}={_n_coords * _dim_per_coord_multi if _dim_per_coord_multi else '?'} " + f"{'== ' + str(_rope_dim) if _dim_per_coord_multi and _n_coords * _dim_per_coord_multi == _rope_dim else '!= ' + str(_rope_dim) + ' (MISMATCH)'}" + ) + print(_rope_msg) + logger.info(_rope_msg) + + def forward(self, batch, return_all=False, teacher_forcing_ratio=1.0): + """ + Forward pass through the model. + + Must remain consistent with rollout(): + 1. Reference pose is frame 0 of the full sequence (not batch[:, 0]) + 2. Input: current position only (no velocity to encoder) + 3. Output: acceleration, integrated via Verlet to get the next position + 4. Predictor maps (current_pos, ref_pose_frame0) -> acceleration + 5. Verlet: x_{t+1} = a*dt² + 2*x_t - x_{t-1} + + Args: + batch: Dictionary with body positions and velocities + return_all: Whether to return intermediate results + teacher_forcing_ratio: Probability of using ground truth (1.0 = always teacher, 0.0 = always student) + + Returns: + predictions: Dictionary with predicted positions, velocities, etc. + """ + # Multi-object path only + if "obj_positions" in batch and getattr(self, "enable_multi_object", False): + return self._forward_multi_object( + batch, + return_all=return_all, + teacher_forcing_ratio=teacher_forcing_ratio, + ) + else: + raise ValueError("Expected multi-object batch with 'obj_positions'.") + + # ================= Multi-object forward (MOVi) ================= + def _forward_multi_object(self, batch, return_all=False, teacher_forcing_ratio=1.0): + """ + Full multi-object forward path using shared per-object PointNet encoder and + mask-aware per-vertex queries across all objects. + """ + time_step_id = batch["time_step_id"] # (B,) + # network inputs: the rigid-noise augmentation's perturbed frames when present (the losses keep comparing + # against batch['obj_positions']) + obj_pos = batch.get( + "obj_positions_input", batch["obj_positions"] + ) # (B,T,Nobj,Nv,3) + obj_mask = batch["valid_obj_mask"] # (B,Nobj) + vtx_mask = batch["valid_vtx_mask"] # (B,T,Nobj,Nv) + ref_obj_positions = batch["ref_obj_positions"] # (B?,Nobj,Nv,3) or (Nobj,Nv,3) + + # Read physics parameters and time step if enabled + phys_params = None # (B, Nobj, 3): [mass, friction, restitution] + skip_step_size_batch = None # (B,) + if self._if_phys_param: + if "phys_params" not in batch: + raise KeyError( + "Batch is missing 'phys_params' required when if_phys_param=true" + ) + phys_params = batch["phys_params"].to(obj_pos.device) # (B, Nobj, 3) + if self._if_time_step: + if "skip_step_size" not in batch: + raise KeyError( + "Batch is missing 'skip_step_size' required when if_time_step=true" + ) + skip_step_size_batch = batch["skip_step_size"].to(obj_pos.device) # (B,) + + B, T, Nobj, Nv, _ = obj_pos.shape + device = obj_pos.device + + # Normalize reference pose to (B,Nobj,Nv,3) + if ref_obj_positions.dim() == 3: + ref_pose = ref_obj_positions.unsqueeze(0).expand(B, -1, -1, -1).to(device) + else: + ref_pose = ref_obj_positions.to(device) # (B,Nobj,Nv,3) + + # ========== Center scene by first frame ========== + # If enabled, subtract the global center of the first frame from all coordinates; + # relative distances and object shapes are unaffected. + if self.config.get("data", {}).get("center_scene_by_t0", False): + for b in range(B): + vmask_b0 = vtx_mask[b, 0] # (Nobj, Nv) + v_t0 = obj_pos[b, 0][vmask_b0] # Valid vertices at t=0 + if v_t0.shape[0] > 0: + center_t0 = v_t0.mean(dim=0) # (3,) + + # Log centering (only once) + if not hasattr(self, "_center_t0_logged_train"): + logger.info( + f"[Center-T0 TRAIN] Enabled. Sample b={b}: center_t0={center_t0.cpu().numpy().tolist()}, " + f"obj_pos range before: [{obj_pos[b].min().item():.4f}, {obj_pos[b].max().item():.4f}]" + ) + + # Shift positions + obj_pos[b] = obj_pos[b] - center_t0 + ref_pose[b] = ref_pose[b] - center_t0 + + if not hasattr(self, "_center_t0_logged_train"): + logger.info( + f"[Center-T0 TRAIN] After centering b={b}: " + f"obj_pos range: [{obj_pos[b].min().item():.4f}, {obj_pos[b].max().item():.4f}], " + f"ref_pose range: [{ref_pose[b].min().item():.4f}, {ref_pose[b].max().item():.4f}]" + ) + if not hasattr(self, "_center_t0_logged_train"): + self._center_t0_logged_train = True + logger.info( + f"[Center-T0 TRAIN] Scene centering applied to obj_pos and ref_pose for batch of {B} samples" + ) + + # Require anchors for query and supervision (no fallback) + if "anchor_idx" not in batch: + raise KeyError( + "Batch is missing 'anchor_idx' required for anchor-based querying" + ) + anchor_idx = batch["anchor_idx"].to(device).long() # (B,Nobj,Nq) + if ( + anchor_idx.dim() != 3 + or anchor_idx.shape[0] != B + or anchor_idx.shape[1] != Nobj + ): + raise ValueError( + f"anchor_idx must be (B,Nobj,Nq); got {tuple(anchor_idx.shape)} vs (B={B},Nobj={Nobj})" + ) + Nq = int(anchor_idx.shape[2]) + + # Initialize persistent state tokens and 3D positions for RoPE + # When no_global_state=True, pass None to decoder (skips all state cross-attention) + if self._no_global_state: + state_tokens = None + state_pos = None + else: + state_tokens, state_pos = self.state_repr.get_state_and_positions(B, device) + + # Storage across time steps + all_pred_pos_for_rollout = ( + [] + ) # STE positions for autoregressive input (B, Nobj, Nv, 3) + all_pred_pos_raw = [] # Raw positions for loss supervision (B, Nobj, Nv, 3) + all_pred_pos_raw_norm = [] # Raw normalized positions for loss (before Kabsch) + all_pred_pos_icp_norm = ( + [] + ) # ICP/Kabsch-transformed normalized positions for loss (after Kabsch) + all_accel_no_icp_norm_anchor_seq = [] # (B, Nobj, Nq, 3) + all_accel_icp_norm_seq = [] + all_pen_pred_vertices_seq = ( + [] + ) # (B, Nobj, Nv, 1) per step if enabled (ALL vertices) + + # Registration refs per object (initialized with GT at the first step) + registration_refs = None + + # Previous step's features for penetration prediction (t>1) + prev_obj_layers = None + prev_state_layers = None + + # ========== ANCHOR VALIDATION (whole window, vectorized, one host sync) ========== + # Every anchor of a valid object must sit on a valid vertex at every frame the loop reads + # (vtx_mask[:, t] for t in 1..T-2). Checked once per batch instead of per (t, object). + if T > 2: + if ( + anchor_idx.dim() != 3 + or anchor_idx.shape[0] != B + or anchor_idx.shape[2] != Nq + ): + raise RuntimeError( + f"anchor_idx shape mismatch: got {tuple(anchor_idx.shape)}, expected (B={B},Nobj,Nq={Nq})" + ) + _vmask_win = vtx_mask[:, 1 : T - 1] # (B,Tw,Nobj,Nv) + _numv_win = _vmask_win.sum(dim=3) # (B,Tw,Nobj) + _idx_win = anchor_idx.unsqueeze(1).expand( + -1, T - 2, -1, -1 + ) # (B,Tw,Nobj,Nq) + _at_anchor = torch.gather(_vmask_win, 3, _idx_win) # (B,Tw,Nobj,Nq) + _ovalid = obj_mask.unsqueeze(1) # (B,1,Nobj) + _zero_bad = _ovalid & (_numv_win == 0) + _short_bad = _ovalid & (_numv_win < Nq) + _anch_bad = (~_at_anchor) & _ovalid.unsqueeze(-1) + _flags = torch.stack([_zero_bad.any(), _short_bad.any(), _anch_bad.any()]) + if bool( + _flags.any().item() + ): # the single sync; everything below is the error path + zb, sb, ab = [bool(x.item()) for x in _flags] + _where = ( + _zero_bad if zb else (_short_bad if sb else _anch_bad.any(dim=-1)) + ).nonzero(as_tuple=False) + _b0, _tw0, _i0 = [int(v) for v in _where[0][:3]] + _nv0 = int(_numv_win[_b0, _tw0, _i0].item()) + ds_name = batch.get("dataset_name", None) + ds_name_str = ( + ds_name + if isinstance(ds_name, str) + else (ds_name[0] if hasattr(ds_name, "__getitem__") else "unknown") + ) + logger.error( + f"[Anchor Debug] dataset={ds_name_str} t={_tw0 + 1} obj={_i0} sample={_b0} Nv={Nv} anchor_k={Nq} " + f"valid_vertices={_nv0} zero_any={zb} short_any={sb} bad_anchor_any={ab}" + ) + if zb: + raise RuntimeError( + f"Zero valid vertices encountered at t={_tw0 + 1}, obj={_i0} for some valid samples; dataset likely corrupted." + ) + if sb: + raise RuntimeError( + f"Not enough valid vertices for anchors at t={_tw0 + 1}, obj={_i0}: valid={_nv0} < anchor_k={Nq}" + ) + raise RuntimeError( + f"Invalid anchors: anchor indices refer to invalid/padded vertices at t={_tw0 + 1}, obj={_i0}. " + f"Consider reducing anchor_k or ensuring consistent Nv per frame." + ) + + # Object slots that hold a valid object in at least one sample (fixed for the batch). The per-object + # KNN skips the others: they are padding everywhere and their distance vectors are zeroed anyway. + _knn_active_slots = ( + torch.nonzero(obj_mask.any(dim=0)).flatten().tolist() + if self._knn_num_groups <= 1 + else None + ) + # Constant scalars of the decoder attention mask and the valid-object index set used by the Kabsch + # fit: both are fixed for the batch, so they are built once instead of at every step (the per-step + # versions cost two host-to-device copies and one nonzero sync each step). + _attn_zero_fwd = torch.tensor(0.0, dtype=torch.float32, device=obj_mask.device) + _attn_ninf_fwd = torch.tensor( + float("-inf"), dtype=torch.float32, device=obj_mask.device + ) + _kabsch_valid_idx_fwd = obj_mask.reshape(-1).nonzero(as_tuple=True)[0] + # Valid-vertex packs (model.knn_packed / model.encoder.packed), fixed for the batch: every frame of a + # sequence has the same vertex mask, which is checked once here. + _knn_pack_fwd = _enc_pack_fwd = None + if self._knn_packed or self._encoder_packed: + _vm1 = vtx_mask[:, 1] + if not bool((vtx_mask[:, 1 : T - 1] == _vm1.unsqueeze(1)).all()): + raise ValueError( + "valid-vertex packing needs the same vertex mask in every frame of a sequence" + ) + if self._knn_packed: + _knn_pack_fwd = knn_pack(_vm1, obj_mask) + if self._encoder_packed: + _enc_m = self.shared_obj_encoder + _vr0 = batch.get("vtx_rank") + _enc_pack_fwd = EncoderPack( + _vm1.reshape(B * Nobj, Nv), + _vr0.reshape(B * Nobj, Nv) if _vr0 is not None else None, + len(_enc_m.hierarchical_layers), + _enc_m.hierarchical_subsets, + 1, + small=_enc_m.subset_small_objects, + ) + + def _net_step( + t, + curr_pos_bt, + prev_pos_bt, + B, + Nobj, + Nv, + Nq, + obj_mask, + vtx_mask_t, + anchor_idx, + ref_pose, + phys_params, + skip_step_size_batch, + vtx_rank_all, + state_tokens, + bn_groups, + knn_chunks, + knn_slots, + enc_pack=None, + knn_pk=None, + ): + """Network part of one step (distance features -> encoder -> AVP -> decoder -> predictor): + per-anchor accelerations for the batch given current/previous positions. Called once per step, + or once for all teacher-forced steps stacked along the batch (bn_groups = knn_chunks = steps). + """ + # Velocities + vel_bt = curr_pos_bt - prev_pos_bt # (B,Nobj,Nv,3) + + # Build per-object features and per-object latents + per_obj_latents = [] # each (B, D) + # Build per-object normalized query features and ref poses (padded) + per_obj_vel_norm = [] + per_obj_relpos_norm = [] + per_obj_distvec_norm = [] + + # Prepare far-away filler for invalid vertices in reference lists + FAR = 1e6 + + # (Anchor validation runs once per batch before the time loop: it depends only on the + # ground-truth masks and anchor indices, never on predictions.) + + # ========== VECTORIZED KNN FOR ALL OBJECTS AT ONCE ========== + # curr_pos_bt: (B, Nobj, Nv, 3), vtx_mask_t: (B, Nobj, Nv) + vmask_t = vtx_mask_t # (B, Nobj, Nv) + + # Compute KNN distance vectors for ALL objects simultaneously + # Only computes for valid objects and valid vertices! + # Returns: (B, Nobj, Nv, 3) - vector from each vertex to nearest point on OTHER objects + if knn_pk is not None: + # packed valid vertices; with batched time steps the step-sized chunks share one pack + _cs = B // knn_chunks + dist_vec_all = ( + compute_knn_packed(curr_pos_bt, knn_pk) + if knn_chunks == 1 + else torch.cat( + [ + compute_knn_packed( + curr_pos_bt[_c * _cs : (_c + 1) * _cs], knn_pk + ) + for _c in range(knn_chunks) + ], + dim=0, + ) + ) + elif knn_chunks == 1: + dist_vec_all = compute_knn_all_objects_vectorized( + curr_pos_bt, + vmask_t, + FAR, + valid_obj_mask=obj_mask, + num_groups=self._knn_num_groups, + active_slots=knn_slots, + ) + else: + # batched time steps: one call per step-sized chunk (memory of the per-object cdist stays + # that of a single step) + _cs = B // knn_chunks + dist_vec_all = torch.cat( + [ + compute_knn_all_objects_vectorized( + curr_pos_bt[_c * _cs : (_c + 1) * _cs], + vmask_t[_c * _cs : (_c + 1) * _cs], + FAR, + valid_obj_mask=obj_mask[_c * _cs : (_c + 1) * _cs], + num_groups=self._knn_num_groups, + active_slots=knn_slots, + ) + for _c in range(knn_chunks) + ], + dim=0, + ) + + # Apply plane constraint adjustment (vectorized) + # Reshape for distance_2_plane: (B, Nobj*Nv, 3) + # Use reshape() instead of view() to handle non-contiguous tensors + curr_pos_flat = curr_pos_bt.reshape(B, Nobj * Nv, 3) + dist_vec_flat = dist_vec_all.reshape(B, Nobj * Nv, 3) + dist_vec2_flat = distance_2_plane( + curr_pos_flat, [0, 0, 0], "z", dist_vec_flat + ) + dist_vec2_stacked = dist_vec2_flat.reshape(B, Nobj, Nv, 3) + + # Relative position (vectorized): curr_pos - ref_pose + relpos_stacked = curr_pos_bt - ref_pose # (B, Nobj, Nv, 3) + + # Velocity (already in correct shape) + vel_stacked = vel_bt # (B, Nobj, Nv, 3) + + # Noise augmentation (vectorized) + if self.training: + aug_config = self.config["data"].get("augmentation", {}) + if aug_config.get("dist_vec_noise", 0) > 0: + std = aug_config["dist_vec_noise"] + dist_vec2_stacked = ( + dist_vec2_stacked + torch.randn_like(dist_vec2_stacked) * std + ) + if aug_config.get("velocity_noise", 0) > 0: + std = aug_config["velocity_noise"] + vel_stacked = vel_stacked + torch.randn_like(vel_stacked) * std + if aug_config.get("rel_pos_noise", 0) > 0: + std = aug_config["rel_pos_noise"] + relpos_stacked = ( + relpos_stacked + torch.randn_like(relpos_stacked) * std + ) + + # Apply vertex masks BEFORE normalization (vectorized across all objects) + vmask_all = vtx_mask_t.unsqueeze(-1) # (B,Nobj,Nv,1) + dist_vec2_stacked = torch.where( + vmask_all, dist_vec2_stacked, torch.zeros_like(dist_vec2_stacked) + ) + vel_stacked = torch.where( + vmask_all, vel_stacked, torch.zeros_like(vel_stacked) + ) + relpos_stacked = torch.where( + vmask_all, relpos_stacked, torch.zeros_like(relpos_stacked) + ) + + # Apply network-predicted directional sign to ALL vertices' dist_vec (before encoder) + # Only for t>1 (may have penetration in student forcing with predicted positions) + # For t=1, GT positions (no penetration, sign=+1 implicit) + # Use if_penetration as master switch (fallback to if_directional_distance for backwards compatibility) + penetration_enabled = self.config["model"].get( + "if_penetration", + self.config["model"].get("if_directional_distance", False), + ) + if penetration_enabled and t > 1: + # Predict penetration sign using PREVIOUS step's features and CURRENT positions + # curr_pos_bt was already selected by teacher/student forcing above + # In student forcing mode: curr_pos_bt = all_pred_pos_for_rollout[-1] = blended_next (after Kabsch transform) + # GT supervision will use predictions['obj_positions_ste'] to match + pen_pred_all_vertices = self.predictor.predict_penetration_vertices( + prev_obj_layers, # Use previous step's decoder outputs + prev_state_layers, # Use previous step's state features + pos_full=curr_pos_bt, # AFTER Kabsch transform (student forcing) or GT (teacher forcing/t=1) + ) # (B, Nobj, Nv, 1) in [-1,1] + + # Apply predicted sign to dist_vec for ALL vertices + pen_sign_all = torch.sign( + pen_pred_all_vertices + 1e-6 + ) # (B, Nobj, Nv, 1) in {-1,+1} + dist_vec2_stacked = dist_vec2_stacked * pen_sign_all # (B, Nobj, Nv, 3) + + # Store predicted signs for loss supervision (will compute GT later in validation) + all_pen_pred_vertices_seq.append(pen_pred_all_vertices) + elif penetration_enabled and t == 1: + # First step: GT positions, no penetration prediction + # Append placeholder (all +1) to maintain time dimension alignment + # This will NOT be used in loss (masked out), but keeps dimensions consistent + B, Nobj, Nv = ( + dist_vec2_stacked.shape[0], + dist_vec2_stacked.shape[1], + dist_vec2_stacked.shape[2], + ) + placeholder = torch.ones( + (B, Nobj, Nv, 1), dtype=torch.float32, device=self.device + ) + all_pen_pred_vertices_seq.append(placeholder) + + # No normalization — use raw values directly + dist_vec2_norm_stacked = dist_vec2_stacked + vel_norm_stacked = vel_stacked + relpos_norm_stacked = relpos_stacked + + # Broadcast physics parameters and time_step to vertex dimension if enabled + phys_broadcast = None # (B, Nobj, Nv, 3) + time_broadcast = None # (B, Nobj, Nv, 2) - contains [s, s²] + + if self._if_phys_param: + # Broadcast per-object physics params (B,Nobj,3) to all vertices (B,Nobj,Nv,3) + phys_broadcast = phys_params.unsqueeze(2).expand( + B, Nobj, Nv, 3 + ) # (B,Nobj,Nv,3) + # Mask with vertex mask to ensure padding vertices get zero physics params + phys_broadcast = torch.where( + vmask_all, phys_broadcast, torch.zeros_like(phys_broadcast) + ) + + if self._if_time_step: + # Broadcast both s and s² to all objects and vertices (B,Nobj,Nv,2) + # Input both s and s² because: + # - s: original step_size for temporal context + # - s²: the scaling factor that will be applied to network output + # - Network can leverage both for better learning + step_size_s = skip_step_size_batch.float() # (B,) + step_size_sq = (skip_step_size_batch**2).float() # (B,) + # Stack s and s² -> (B, 2) + time_features = torch.stack( + [step_size_s, step_size_sq], dim=-1 + ) # (B, 2) + time_broadcast = time_features.view(B, 1, 1, 2).expand( + B, Nobj, Nv, 2 + ) # (B,Nobj,Nv,2) + # Time step doesn't need masking (it's a scalar property of the batch, not vertex-specific) + # But for consistency, we can zero it for padding vertices + # vmask_all is (B,Nobj,Nv,1), expand to match (B,Nobj,Nv,2) + vmask_time = vmask_all.expand(-1, -1, -1, 2) # (B,Nobj,Nv,2) + time_broadcast = torch.where( + vmask_time, time_broadcast, torch.zeros_like(time_broadcast) + ) + + # PointNet Encoder: process each object's Nv vertices -> per-object global feature + # Input dimension depends on flags: + # Base: 9D (dist_vec + velocity + rel_pos) + # +if_phys_param=true: +3D + # +if_time_step=true (no FiLM): +2D (s and s²) + # Split into components and reshape to (B*Nobj, Nv, D) for encoder + dist_reshaped = dist_vec2_norm_stacked.view(B * Nobj, Nv, 3) + vel_reshaped = vel_norm_stacked.view(B * Nobj, Nv, 3) + relpos_reshaped = relpos_norm_stacked.view(B * Nobj, Nv, 3) + + phys_reshaped = None + if self._if_phys_param: + phys_reshaped = phys_broadcast.view(B * Nobj, Nv, 3) + + time_reshaped = None + if self._if_time_step and not self._film_conditioning: + time_reshaped = time_broadcast.view(B * Nobj, Nv, 2) # (s, s²) + + # Prepare vertex mask for encoder pooling + # vmask_all is (B,Nobj,Nv,1), need (B*Nobj, Nv) + vtx_mask_for_encoder = vmask_all.squeeze(-1).reshape( + B * Nobj, Nv + ) # (B*Nobj, Nv) + height_reshaped = None + if self.shared_obj_encoder.height_input: + _h = torch.where( + vmask_all, + curr_pos_bt[..., 2:3], + torch.zeros_like(curr_pos_bt[..., 2:3]), + ) + height_reshaped = _h.reshape(B * Nobj, Nv, 1) + rank_reshaped = None + if self.shared_obj_encoder.hierarchical_subsets: + if vtx_rank_all is None: + raise KeyError( + "model.encoder.hierarchical_subsets requires data.fps_order: true (batch lacks 'vtx_rank')" + ) + rank_reshaped = vtx_rank_all.reshape(B * Nobj, Nv) + + # Call encoder with all features (encoder will concatenate them internally) + # AVP: capture per-point features if AVP is enabled + obj_features_encoded, _per_vtx_feat, _ = self.shared_obj_encoder( + dist_reshaped, + vel_reshaped, + relpos_reshaped, + phys_params=phys_reshaped, + time_step=time_reshaped, + valid_mask=vtx_mask_for_encoder, + height=height_reshaped, + vtx_rank=rank_reshaped, + bn_groups=bn_groups, + pack=enc_pack, + ) # obj_features_encoded: (B*Nobj, enc_out_dim), _per_vtx_feat: (B*Nobj, backbone_dim, Nv) + + # AVP: Anchor-Vertex Pooling — enrich anchor queries with collision-aware features + _avp_anchor_features = ( + None # Will be (B, Nobj, Nq, avp_output_dim) if AVP enabled + ) + if self.avp is not None and _per_vtx_feat is not None: + # _per_vtx_feat is (B*Nobj, backbone_dim, Nv) — Conv1d channels-first format + # Transpose to (B*Nobj, Nv, backbone_dim) for AVP + _pvf = _per_vtx_feat.transpose(1, 2) # (B*Nobj, Nv, backbone_dim) + # Anchor positions in flat format + _curr_pos_flat = curr_pos_bt.reshape(B * Nobj, Nv, 3) # (B*Nobj, Nv, 3) + _anchor_pos_flat = torch.gather( + _curr_pos_flat, + dim=1, + index=anchor_idx.reshape(B * Nobj, -1) + .unsqueeze(-1) + .expand(-1, -1, 3), + ) # (B*Nobj, Nq, 3) + _avp_anchor_features = self.avp( + _pvf, _curr_pos_flat, _anchor_pos_flat, vtx_mask_for_encoder + ) # (B*Nobj, Nq, avp_output_dim) + _avp_anchor_features = _avp_anchor_features.view( + B, Nobj, anchor_idx.shape[2], -1 + ) # (B, Nobj, Nq, avp_output_dim) + # Zero out padding objects (proj bias makes padding non-zero; consistent with other features) + _avp_anchor_features = ( + _avp_anchor_features * obj_mask.unsqueeze(-1).unsqueeze(-1).float() + ) + + obj_features_encoded = obj_features_encoded.view( + B, Nobj, -1 + ) # (B, Nobj, enc_out_dim) + + # Project to token dimension for Decoder + obj_features = self.obj_feature_proj(obj_features_encoded) # (B, Nobj, D) + + # FiLM: modulate encoder features with step-size [s, s²] + if self.encoder_film is not None and skip_step_size_batch is not None: + time_feat_film = torch.stack( + [skip_step_size_batch.float(), (skip_step_size_batch**2).float()], + dim=-1, + ) # (B, 2) + obj_features = self.encoder_film( + obj_features, time_feat_film + ) # (B, Nobj, D) + + # Object anchor positions for RoPE (like RenderFormer's triangle approach) + # RenderFormer: Triangle = 3 vertices (9D) + # Our approach: Object = Nq anchors (Nq*3 D) + # Gather anchor positions: (B, Nobj, Nq, 3) + # Replace padding objects' positions with zeros before the gather: + # padding objects have anchor_idx=0, which would gather arbitrary values into RoPE inputs. + obj_mask_pos = obj_mask.unsqueeze(-1).unsqueeze(-1) # (B, Nobj, 1, 1) + safe_curr_pos_bt = torch.where( + obj_mask_pos, curr_pos_bt, torch.zeros_like(curr_pos_bt) + ) + + idx_exp_3_for_decoder = anchor_idx.unsqueeze(-1).expand( + -1, -1, -1, 3 + ) # (B, Nobj, Nq, 3) + obj_anchor_positions = torch.gather( + safe_curr_pos_bt, dim=2, index=idx_exp_3_for_decoder + ) # (B, Nobj, Nq, 3) + # Flatten to (B, Nobj, Nq*3) for multi-point RoPE + Nq = anchor_idx.shape[2] # Number of anchors per object + + # ============== PE Type Routing (forward) ============== + if self._pe_type in ("arope", "se3"): + # SE(3) uses arope anchors for cross-attention RoPE; SE(3) bias computed separately + obj_anchor_flat = obj_anchor_positions.view( + B, Nobj, Nq * 3 + ) # (B, Nobj, 12) + elif self._pe_type in ("obb", "pca"): + obj_anchor_flat = self._geometry_encoder( + safe_curr_pos_bt, vtx_mask_t, obj_mask=obj_mask + ) # (B, Nobj, 12) + else: + raise ValueError(f"Unknown pe_type='{self._pe_type}' in forward().") + + # ============== SE(3) Pairwise Bias Computation ============== + _se3_bias = None + if self._pe_type == "se3" and self._se3_bias_module is not None: + from embodichain_models.dynamics.rigidformer._native.models.se3_attention import ( + extract_object_pose_features, + ) + + _rotmats, _centers, _scales = extract_object_pose_features( + safe_curr_pos_bt, ref_pose, vtx_mask_t, obj_mask + ) + _se3_bias = self._se3_bias_module(_rotmats, _centers, _scales, obj_mask) + + # Decoder attention mask: SDPA additive format (0=attend, -inf=mask) + # Prevents padding objects from participating in Decoder self-attention. + # Combined with feature zeroing to ensure clean inputs. + obj_attn_mask = torch.where( + obj_mask.unsqueeze(1).unsqueeze(1), # (B, 1, 1, Nobj) + _attn_zero_fwd, + _attn_ninf_fwd, + ) # (B, 1, 1, Nobj) float32 + + # Zero out padding object features and anchor positions BEFORE Decoder + # to prevent garbage from padding objects from destabilizing attention + pad_mask_float = obj_mask.unsqueeze(-1).to(dtype=obj_features.dtype) + obj_features = obj_features * pad_mask_float # (B, Nobj, D) + obj_anchor_flat = obj_anchor_flat * obj_mask.unsqueeze(-1).to( + dtype=obj_anchor_flat.dtype + ) + + # Decoder interaction with per-object features + # Per-layer FiLM: pass time_features for step-size conditioning in each layer + _decoder_time_feat = ( + time_feat_film + if (self._film_conditioning and self._if_time_step) + else None + ) + all_state_layers, all_obj_layers = self.decoder( + state_tokens, + obj_features, + state_pos, + obj_anchor_flat, + mask=obj_attn_mask, # SDPA additive mask for padding objects + se3_bias=_se3_bias, + time_features=_decoder_time_feat, # Per-layer FiLM + ) + + # Update state for next step (skip when no_global_state) + if not self._no_global_state: + state_tokens = all_state_layers[-1] + + # Build anchor query for Predictor (sample at anchor points Nq) + # Query dimension depends on flags: + # Base: 9D (dist_vec + vel + rel_pos) + # +if_phys_param=true: +3D = 12D + # +if_time_step=true (no FiLM): +2D (s and s²) + idx_exp_3 = anchor_idx.unsqueeze(-1).expand(-1, -1, -1, 3) # (B,Nobj,Nq,3) + distvec_at_anchors = torch.gather( + dist_vec2_norm_stacked, dim=2, index=idx_exp_3 + ) # (B,Nobj,Nq,3) + vel_at_anchors = torch.gather( + vel_norm_stacked, dim=2, index=idx_exp_3 + ) # (B,Nobj,Nq,3) + relpos_at_anchors = torch.gather( + relpos_norm_stacked, dim=2, index=idx_exp_3 + ) # (B,Nobj,Nq,3) + + # Build query features list + query_features = [distvec_at_anchors, vel_at_anchors, relpos_at_anchors] + + # Add physics parameters at anchors if enabled + if self._if_phys_param: + phys_at_anchors = torch.gather( + phys_broadcast, dim=2, index=idx_exp_3 + ) # (B,Nobj,Nq,3) + query_features.append(phys_at_anchors) + + # Add time step at anchors if enabled AND not using FiLM + if self._if_time_step and not self._film_conditioning: + idx_exp_2 = anchor_idx.unsqueeze(-1).expand( + -1, -1, -1, 2 + ) # (B,Nobj,Nq,2) + time_at_anchors = torch.gather( + time_broadcast, dim=2, index=idx_exp_2 + ) # (B,Nobj,Nq,2) + query_features.append(time_at_anchors) + + # AVP: Append collision-aware anchor features to query + if _avp_anchor_features is not None: + query_features.append( + _avp_anchor_features + ) # (B, Nobj, Nq, avp_output_dim) + + # Concatenate all query features + anchor_query = torch.cat( + query_features, dim=-1 + ) # (B,Nobj,Nq, 9/12 + avp_dim) + + # Extract anchor positions for RoPE (current frame positions at anchors) + anchor_pos_at_curr = torch.gather( + curr_pos_bt, dim=2, index=idx_exp_3 + ) # (B,Nobj,Nq,3) + + # Penetration sign was already applied to all vertices' dist_vec before the encoder; + # the anchor query therefore carries the sign information as well. + + _ref_src = ( + _object_centered(ref_pose, vmask_all.squeeze(-1).bool()) + if self._ref_pose_object_centered + else ref_pose + ) + if self._ref_pose_object_centered and not getattr( + self, "_refpose_logged_fwd", False + ): + self._refpose_logged_fwd = True + print( + f"[RefPose] object-centred (x/y only, z absolute) reference pose ACTIVE (forward): |shift| max={((ref_pose - _ref_src).abs().max()).item():.3f} m", + flush=True, + ) + ref_pose_at_anchors = torch.gather( + _ref_src, dim=2, index=idx_exp_3 + ) # (B,Nobj,Nq,3) + + # Check for NaN/Inf before predictor + if DEBUG_NAN and ( + torch.isnan(anchor_query).any() or torch.isinf(anchor_query).any() + ): + logger.error(f"NaN/Inf detected in anchor_query!") + logger.error( + f" distvec_at_anchors: nan={torch.isnan(distvec_at_anchors).sum().item()}, inf={torch.isinf(distvec_at_anchors).sum().item()}" + ) + logger.error( + f" vel_at_anchors: nan={torch.isnan(vel_at_anchors).sum().item()}, inf={torch.isinf(vel_at_anchors).sum().item()}" + ) + logger.error( + f" relpos_at_anchors: nan={torch.isnan(relpos_at_anchors).sum().item()}, inf={torch.isinf(relpos_at_anchors).sum().item()}" + ) + + if DEBUG_NAN and ( + torch.isnan(ref_pose_at_anchors).any() + or torch.isinf(ref_pose_at_anchors).any() + ): + logger.error(f"NaN/Inf detected in ref_pose_at_anchors!") + logger.error( + f" ref_pose_at_anchors range: [{ref_pose_at_anchors.min().item():.4f}, {ref_pose_at_anchors.max().item():.4f}]" + ) + + for layer_idx, obj_layer in enumerate(all_obj_layers if DEBUG_NAN else ()): + if torch.isnan(obj_layer).any() or torch.isinf(obj_layer).any(): + logger.error(f"NaN/Inf detected in all_obj_layers[{layer_idx}]!") + logger.error( + f" Shape: {obj_layer.shape}, nan={torch.isnan(obj_layer).sum().item()}, inf={torch.isinf(obj_layer).sum().item()}" + ) + + # Predict accelerations per anchor using grouped format (B, Nobj, Nq, *) + accel_grouped = self.predictor( + all_obj_layers, + anchor_query, # (B, Nobj, Nq, 9/12) + ref_pose=ref_pose_at_anchors, # (B, Nobj, Nq, 3) + query_pos=anchor_pos_at_curr, # (B, Nobj, Nq, 3) + key_pos=obj_anchor_flat, # (B, Nobj, Nq*3) - multi-point RoPE + multi_scale_state_features=( + None if self._no_global_state else all_state_layers + ), + state_pos=None if self._no_global_state else state_pos, + obj_mask=obj_mask, # Object validity mask for padding + ) # Returns (B, Nobj, Nq, 3) + + # Check for NaN/Inf after predictor + if DEBUG_NAN and ( + torch.isnan(accel_grouped).any() or torch.isinf(accel_grouped).any() + ): + logger.error(f"NaN/Inf detected in predictor output!") + logger.error( + f" accel_grouped: nan={torch.isnan(accel_grouped).sum().item()}, inf={torch.isinf(accel_grouped).sum().item()}" + ) + logger.error( + f" accel_grouped range: min={accel_grouped[~torch.isnan(accel_grouped)].min().item() if (~torch.isnan(accel_grouped)).any() else 'all nan'}, max={accel_grouped[~torch.isnan(accel_grouped)].max().item() if (~torch.isnan(accel_grouped)).any() else 'all nan'}" + ) + + return ( + vmask_t, + accel_grouped, + all_obj_layers, + all_state_layers, + state_tokens, + ) + + # ========== TIME-BATCHED NETWORK PASS (training.time_batched_forward) ========== + # With pure teacher forcing every step's network input is built from ground-truth frames, so the network + # part of all T-2 steps runs once on the steps stacked along the batch (step-major: rows [g*B, (g+1)*B) + # belong to step g+1). The encoder keeps per-step BatchNorm statistics (grouped masked BN: same + # statistics, gradients and running-statistics updates as step-by-step calls), and the per-step physics + # (Verlet, Kabsch, ground handling) below is unchanged. Equal to the step-by-step path up to + # floating-point rounding and the dropout random stream. + _tb_accel = None + if self.training and self._time_batched_forward and T > 2: + if teacher_forcing_ratio < 1.0: + raise ValueError( + "training.time_batched_forward requires teacher_forcing_ratio == 1 (inputs must be ground truth)" + ) + Tw = T - 2 + rope_memo_reset() + + def _rep(x): + return ( + x.unsqueeze(0) + .expand(Tw, *x.shape) + .reshape(Tw * x.shape[0], *x.shape[1:]) + ) + + _ref_b = ( + _rep(ref_pose) + if ( + torch.is_tensor(ref_pose) + and ref_pose.dim() == 4 + and ref_pose.shape[0] == B + ) + else ref_pose + ) + _vr = batch.get("vtx_rank") + _, _acc_s, _, _, _ = _net_step( + None, + obj_pos[:, 1 : T - 1].transpose(0, 1).reshape(Tw * B, Nobj, Nv, 3), + obj_pos[:, 0 : T - 2].transpose(0, 1).reshape(Tw * B, Nobj, Nv, 3), + Tw * B, + Nobj, + Nv, + Nq, + _rep(obj_mask), + vtx_mask[:, 1 : T - 1].transpose(0, 1).reshape(Tw * B, Nobj, Nv), + _rep(anchor_idx), + _ref_b, + _rep(phys_params) if phys_params is not None else None, + ( + _rep(skip_step_size_batch) + if skip_step_size_batch is not None + else None + ), + _rep(_vr) if _vr is not None else None, + None, + Tw, + Tw, + _knn_active_slots, + enc_pack=( + EncoderPack( + vtx_mask[:, 1 : T - 1] + .transpose(0, 1) + .reshape(Tw * B * Nobj, Nv), + ( + _rep(_vr).reshape(Tw * B * Nobj, Nv) + if _vr is not None + else None + ), + len(self.shared_obj_encoder.hierarchical_layers), + self.shared_obj_encoder.hierarchical_subsets, + Tw, + small=self.shared_obj_encoder.subset_small_objects, + ) + if self._encoder_packed + else None + ), + knn_pk=_knn_pack_fwd, + ) + _tb_accel = _acc_s.view(Tw, B, *_acc_s.shape[1:]) + + # Time loop predicts for steps t in [1, T-2] -> next positions x_{t+1} + for t in range(1, T - 1): + rope_memo_reset() # rotary angles are memoized within one step only + # Teacher vs student forcing (no special-case first step) + use_teacher_forcing = torch.rand(1).item() < teacher_forcing_ratio + + if use_teacher_forcing or len(all_pred_pos_for_rollout) == 0: + prev_pos_bt = obj_pos[:, t - 1] # (B,Nobj,Nv,3) + curr_pos_bt = obj_pos[:, t] + else: + # Use last predicted (STE blended) as current for autoregressive stability + curr_pos_bt = all_pred_pos_for_rollout[-1] # (B,Nobj,Nv,3) + if len(all_pred_pos_for_rollout) >= 2: + prev_pos_bt = all_pred_pos_for_rollout[-2] + else: + prev_pos_bt = obj_pos[:, t - 1] + + if _tb_accel is not None: + # time-batched: this step's accelerations were computed with all teacher-forced steps + vmask_t = vtx_mask[:, t] + accel_grouped = _tb_accel[t - 1] + all_obj_layers = all_state_layers = None + else: + ( + vmask_t, + accel_grouped, + all_obj_layers, + all_state_layers, + state_tokens, + ) = _net_step( + t, + curr_pos_bt, + prev_pos_bt, + B, + Nobj, + Nv, + Nq, + obj_mask, + vtx_mask[:, t], + anchor_idx, + ref_pose, + phys_params, + skip_step_size_batch, + batch.get("vtx_rank"), + state_tokens, + 1, + 1, + _knn_active_slots, + enc_pack=_enc_pack_fwd, + knn_pk=_knn_pack_fwd, + ) + # Already in grouped format, no need to reshape + accel_flat = accel_grouped # (B, Nobj, Nq, 3) + accel_base_norm = ( + accel_flat # Network output: base acceleration (step_size-independent) + ) + + # Verlet integration (dt_verlet=1) + dt_squared = self.dt_verlet * self.dt_verlet + with torch.amp.autocast("cuda", enabled=False): + # No normalizer — network output is raw acceleration + accel_base_denorm = accel_base_norm + + # ============================================================ + # Explicit step_size² scaling for multi-timescale training + # ============================================================ + # Physical reasoning: + # - Network outputs "base acceleration" a_base (step_size-independent) + # - Actual acceleration = a_base × step_size² + # - This explicit scaling removes the burden from network to learn s² relationship + # - Benefit: Better generalization to unseen step_sizes, easier optimization + # ============================================================ + if self._if_time_step and skip_step_size_batch is not None: + step_size_sq = (skip_step_size_batch**2).view( + B, 1, 1, 1 + ) # (B,1,1,1) + accel_scaled_denorm = ( + accel_base_denorm * step_size_sq + ) # (B,Nobj,Nq,3) + # Also scale normalized version for loss computation + accel_no_icp_norm_anchor = accel_base_norm * step_size_sq + else: + accel_scaled_denorm = accel_base_denorm + accel_no_icp_norm_anchor = accel_base_norm + + # Build base next positions for all vertices as zero-accel extrapolation + # No detach: allow gradients to flow through position history for long-horizon training + next_pos_bt = 2 * curr_pos_bt - prev_pos_bt # (B,Nobj,Nv,3) + + # Vectorized scatter using scatter_reduce (faster than per-object loop) + # Step 1: Gather all anchor positions at once + anchor_idx_exp = anchor_idx.unsqueeze(-1).expand( + -1, -1, -1, 3 + ) # (B,Nobj,Nq,3) + curr_anchors = torch.gather( + curr_pos_bt, dim=2, index=anchor_idx_exp + ) # (B,Nobj,Nq,3) + prev_anchors = torch.gather( + prev_pos_bt, dim=2, index=anchor_idx_exp + ) # (B,Nobj,Nq,3) + + # Step 2: Compute next positions at anchors using Verlet (with scaled acceleration) + next_anchors = ( + accel_scaled_denorm * dt_squared + 2 * curr_anchors - prev_anchors + ) # (B,Nobj,Nq,3) + + # Step 3: Compute updates (delta from base prediction) + base_anchors = torch.gather( + next_pos_bt, dim=2, index=anchor_idx_exp + ) # (B,Nobj,Nq,3) + anchor_updates = next_anchors - base_anchors # (B,Nobj,Nq,3) + + # Mask padding objects' updates to zero + # Padding objects have anchor_idx=0 which would scatter to wrong positions + obj_mask_scatter = ( + obj_mask.unsqueeze(-1).unsqueeze(-1).float() + ) # (B,Nobj,1,1) + anchor_updates = ( + anchor_updates * obj_mask_scatter + ) # padding objects -> zero updates + + # Step 4: Vectorized scatter using scatter_reduce + # Flatten batch and object dimensions for efficient scatter + # Use reshape() to handle non-contiguous tensors + B_Nobj_Nv = B * Nobj * Nv + next_pos_flat = next_pos_bt.reshape(B_Nobj_Nv, 3) # (B*Nobj*Nv, 3) + + # Create flat indices for scatter_reduce + # anchor_idx is (B, Nobj, Nq), we need indices in flat (B*Nobj*Nv) space + obj_offsets = ( + torch.arange(Nobj, device=anchor_idx.device).view(1, Nobj, 1) * Nv + ) # (1,Nobj,1) + batch_offsets = torch.arange(B, device=anchor_idx.device).view( + B, 1, 1 + ) * ( + Nobj * Nv + ) # (B,1,1) + flat_anchor_idx = ( + anchor_idx + obj_offsets + batch_offsets + ) # (B,Nobj,Nq) + flat_anchor_idx = flat_anchor_idx.reshape(B * Nobj * Nq) # (B*Nobj*Nq,) + + # Flatten updates + anchor_updates_flat = anchor_updates.reshape( + B * Nobj * Nq, 3 + ) # (B*Nobj*Nq, 3) + + # Apply updates using scatter_reduce (more efficient than loop) + next_pos_flat = torch.scatter_reduce( + next_pos_flat, + dim=0, + index=flat_anchor_idx.unsqueeze(-1).expand(-1, 3), + src=anchor_updates_flat, + reduce="sum", + include_self=True, + ) + + # Reshape back + next_pos_bt = next_pos_flat.reshape(B, Nobj, Nv, 3) + + # Per-object rigid projection using anchors (Kabsch on anchors), then apply to full Nv + if registration_refs is None: + registration_refs = curr_pos_bt.detach().clone() # (B,Nobj,Nv,3) + + # Vectorized Kabsch: Process ONLY valid objects (filter padding BEFORE SVD) + # Step 1: Gather anchor positions for all objects at once + anchor_idx_exp = anchor_idx.unsqueeze(-1).expand( + -1, -1, -1, 3 + ) # (B,Nobj,Nq,3) + last_anchor_all = torch.gather( + registration_refs, dim=2, index=anchor_idx_exp + ) # (B,Nobj,Nq,3) + curr_anchor_all = torch.gather( + next_pos_bt, dim=2, index=anchor_idx_exp + ) # (B,Nobj,Nq,3) + + # Reshape to (B*Nobj, Nq, 3) for batch processing + B_Nobj = B * Nobj + last_anchor_flat = last_anchor_all.view(B_Nobj, Nq, 3) + curr_anchor_flat = curr_anchor_all.view(B_Nobj, Nq, 3) + + # Filter padding objects before SVD; they would produce garbage R/t + obj_mask_flat = obj_mask.view(B_Nobj) # (B*Nobj,) + valid_idx = _kabsch_valid_idx_fwd # indices of valid objects (fixed within the forward) + n_valid = valid_idx.shape[0] + + # Initialize R/t for ALL objects (identity R, zero t) + R_flat = ( + torch.eye(3, device=device, dtype=last_anchor_flat.dtype) + .unsqueeze(0) + .expand(B_Nobj, -1, -1) + .clone() + ) + t_vec_flat = torch.zeros( + B_Nobj, 1, 3, device=device, dtype=last_anchor_flat.dtype + ) + + if n_valid > 0: + # Gather only valid objects' anchors + last_anchor_valid = last_anchor_flat[valid_idx] # (Nvalid, Nq, 3) + curr_anchor_valid = curr_anchor_flat[valid_idx] # (Nvalid, Nq, 3) + + if self.if_roma: + # RoMa mode: batch process only valid objects + # RoMa requires FP32 for numerical stability + with torch.amp.autocast("cuda", enabled=False): + R_valid, t_vec_valid = roma.rigid_points_registration( + last_anchor_valid.float(), curr_anchor_valid.float() + ) + R_valid = R_valid.float() # Ensure FP32 for downstream + t_vec_valid = t_vec_valid.float().unsqueeze(1) # (Nvalid, 1, 3) + + else: + # Standard SVD Kabsch: batch process only valid objects + centroid_last = last_anchor_valid.mean( + dim=1, keepdim=True + ) # (Nvalid,1,3) + centroid_curr = curr_anchor_valid.mean( + dim=1, keepdim=True + ) # (Nvalid,1,3) + last_c = last_anchor_valid - centroid_last + curr_c = curr_anchor_valid - centroid_curr + + # Batch SVD for valid objects only + H = torch.bmm(last_c.transpose(1, 2), curr_c) # (Nvalid,3,3) + with torch.amp.autocast("cuda", enabled=False): + H_float = H.float() + U, S, Vt = torch.linalg.svd(H_float) # Batch SVD! + U = U.to(H.dtype) + Vt = Vt.to(H.dtype) + R_valid = torch.bmm( + Vt.transpose(1, 2), U.transpose(1, 2) + ) # (Nvalid,3,3) + + # Fix reflections + with torch.amp.autocast("cuda", enabled=False): + det_R = torch.det(R_valid.float()).to(R_valid.dtype) + need_flip = det_R < 0 + if need_flip.any(): + Fm = ( + torch.eye(3, device=R_valid.device, dtype=R_valid.dtype) + .unsqueeze(0) + .expand(n_valid, -1, -1) + .clone() + ) + Fm[need_flip, 2, 2] = -1 + R_valid = torch.bmm( + torch.bmm(Vt.transpose(1, 2), Fm), U.transpose(1, 2) + ) + + # Translation + rotated_centroid = torch.bmm( + R_valid, centroid_last.transpose(1, 2) + ).transpose(1, 2) + t_vec_valid = centroid_curr - rotated_centroid # (Nvalid,1,3) + + # Scatter valid R/t back to full tensors + R_flat[valid_idx] = R_valid + t_vec_flat[valid_idx] = t_vec_valid + + # Reshape back to (B,Nobj,3,3) and (B,Nobj,1,3) + R_all = R_flat.view(B, Nobj, 3, 3) + t_vec_all = t_vec_flat.view(B, Nobj, 1, 3) + + # ========== VECTORIZED TRANSFORMATION APPLICATION ========== + # Apply R, t to all objects in a single bmm call (no per-object loop) + # registration_refs: (B, Nobj, Nv, 3) + # R_all: (B, Nobj, 3, 3) + # t_vec_all: (B, Nobj, 1, 3) + + # Reshape for batch bmm: (B, Nobj, Nv, 3) -> (B*Nobj, Nv, 3) + last_flat = registration_refs.view(B * Nobj, Nv, 3) + R_flat = R_all.view(B * Nobj, 3, 3) + t_flat = t_vec_all.view(B * Nobj, 1, 3) + + # Single bmm for ALL objects: y = R @ x + t (RoMa convention) + # last_flat.transpose(1, 2): (B*Nobj, 3, Nv) + # bmm result: (B*Nobj, 3, Nv) + # transpose back + t: (B*Nobj, Nv, 3) + y_flat = ( + torch.bmm(R_flat, last_flat.transpose(1, 2)).transpose(1, 2) + + t_flat + ) + # Use reshape() because transpose() makes tensor non-contiguous + yproj_bt = y_flat.reshape( + B, Nobj, Nv, 3 + ).float() # FP32 for numerical stability + + # Raw position BEFORE Kabsch/STE (for position loss supervision) + next_pos_raw_norm = next_pos_bt + + # Apply gradient strategy based on transform type + if self.if_roma: + # RoMa mode: fully differentiable, gradients flow through Kabsch + blended_next = yproj_bt + else: + # Standard mode: straight-through estimator (STE) + # Use Kabsch value but only pass gradient to original prediction + blended_next = next_pos_bt + (yproj_bt - next_pos_bt).detach() + + # Recompute accelerations from blended positions (for acceleration_after_icp loss) + accel_icp_denorm = ( + blended_next - 2 * curr_pos_bt + prev_pos_bt + ) / dt_squared + accel_icp_norm = accel_icp_denorm + + # ========== INTERSECTION HANDLING ========== + # Resolve object-ground and object-object penetrations (STE-differentiable) + intersection_cfg = self.config.get("model", {}).get( + "intersection_handling", {} + ) + rigid_next = ( + blended_next # Kabsch output before the ground/object correction + ) + blended_next = handle_intersections( + blended_next, obj_mask, vmask_t, intersection_cfg + ) + + # Update registration refs to the post-correction position (consistent with rollout); with + # model.kabsch_reference 'fixed' every step fits the same reference mesh (the first current frame) + if not self._kabsch_reference_fixed: + registration_refs = blended_next.detach() + + # Store positions: + # - STE blended positions for autoregressive rollout (more stable) + # - Raw positions for loss supervision (direct gradient) + # - ICP-transformed positions for position_after_icp loss + all_pred_pos_for_rollout.append(blended_next) # For next step input + all_pred_pos_raw.append(next_pos_bt) # For position_before_icp loss + all_pred_pos_raw_norm.append( + next_pos_raw_norm + ) # For position_before_icp loss (normalized) + # ICP/Kabsch-transformed positions for position_after_icp loss. after_kabsch_loss 'rigid': the Kabsch + # output itself, the state the acceleration_after_icp loss uses (the stated after-Kabsch losses); + # 'handled': the state after the intersection correction. + blended_next_norm = ( + rigid_next if self._after_kabsch_loss_rigid else blended_next + ) + all_pred_pos_icp_norm.append( + blended_next_norm + ) # For position_after_icp loss + all_accel_no_icp_norm_anchor_seq.append(accel_no_icp_norm_anchor) + + # Save current features for next step's penetration prediction + prev_obj_layers = all_obj_layers + prev_state_layers = all_state_layers + all_accel_icp_norm_seq.append(accel_icp_norm) + + # Stack time dimension + # Design principle: + # - Position outputs: RAW predictions (before Kabsch) AND ICP-transformed (after Kabsch) + # - Acceleration outputs: before and after Kabsch transform + # - Autoregressive rollout: uses ICP-transformed positions internally for stability + output = { + "obj_positions_denorm": torch.stack( + all_pred_pos_raw, dim=1 + ), # (B,T-2,Nobj,Nv,3) - raw (before Kabsch) + "obj_positions_norm": torch.stack( + all_pred_pos_raw_norm, dim=1 + ), # (B,T-2,Nobj,Nv,3) - raw normalized (before Kabsch) + "obj_positions_icp_norm": torch.stack( + all_pred_pos_icp_norm, dim=1 + ), # (B,T-2,Nobj,Nv,3) - ICP-transformed normalized (after Kabsch) + "obj_positions_ste": torch.stack( + all_pred_pos_for_rollout, dim=1 + ), # (B,T-2,Nobj,Nv,3) - ICP blended (matches rollout) + "obj_accelerations_no_icp_norm_anchor": torch.stack( + all_accel_no_icp_norm_anchor_seq, dim=1 + ), # before Kabsch + "obj_accelerations_norm": torch.stack( + all_accel_icp_norm_seq, dim=1 + ), # after Kabsch + } + + # Include penetration predictions for ALL vertices if collected + penetration_enabled = self.config["model"].get( + "if_penetration", self.config["model"].get("if_directional_distance", False) + ) + if penetration_enabled and len(all_pen_pred_vertices_seq) > 0: + try: + output["obj_pen_pred_vertices"] = torch.stack( + all_pen_pred_vertices_seq, dim=1 + ) # (B,T-2,Nobj,Nv,1) + except Exception as _e: + logger.warning(f"Failed to stack pen predictions: {_e}") + + if return_all: + # Not storing states for multi-object yet + pass + + return output + + def rollout( + self, initial_batch, num_steps, dt=0.01, return_extras=True, warmup_frames=2 + ): + """ + Rollout predictions from initial conditions (multi-object). + + Args: + initial_batch: dict from dataloader containing at least + - obj_positions (B, T0, Nobj, Nv, 3) + - valid_obj_mask (B, Nobj) + - valid_vtx_mask (B, T0, Nobj, Nv) + - ref_obj_positions (B or 1, Nobj, Nv, 3) [optional] + num_steps: total frames to output (including provided GT warmup frames) + dt: timestep (unused; integration handled in model forward) + warmup_frames: number of GT warmup frames (default=2 for our model, HopNet uses 4) + Returns: + dict with 'obj_positions' of shape (B, num_steps, Nobj, Nv, 3) + and optionally 'obj_body_dist_vec2' and 'obj_accelerations_anchor' if return_extras=True + """ + self.eval() + with torch.no_grad(): + if "obj_positions" not in initial_batch: + raise ValueError( + "rollout expects multi-object 'obj_positions' in initial_batch" + ) + + obj_pos_all = initial_batch["obj_positions"] # (B,T0,Nobj,Nv,3) + valid_obj_mask = initial_batch["valid_obj_mask"] # (B,Nobj) + valid_vtx_mask_all = initial_batch["valid_vtx_mask"] # (B,T0,Nobj,Nv) + + # Read physics parameters and time step if enabled + phys_params = None # (B, Nobj, 3): [mass, friction, restitution] + skip_step_size_batch = None # (B,) + if self._if_phys_param: + if "phys_params" not in initial_batch: + raise KeyError( + "rollout requires 'phys_params' in initial_batch when if_phys_param=true" + ) + phys_params = initial_batch["phys_params"].to( + obj_pos_all.device + ) # (B, Nobj, 3) + if self._if_time_step: + if "skip_step_size" not in initial_batch: + raise KeyError( + "rollout requires 'skip_step_size' in initial_batch when if_time_step=true" + ) + skip_step_size_batch = initial_batch["skip_step_size"].to( + obj_pos_all.device + ) # (B,) + + # Optional debug logging of the rollout step size + if_debug = ( + self.config.get("data", {}) + .get("random_step_size", {}) + .get("if_debug", False) + ) + if if_debug: + logger.info( + f"[DEBUG Rollout] step_size for rollout: {skip_step_size_batch.cpu().numpy().tolist()}, num_steps={num_steps}" + ) + + # Validate step_size is within training range + # Check if multi-timescale training is enabled and validate step_size + random_step_config = self.config.get("data", {}).get( + "random_step_size", {} + ) + if random_step_config.get("enable", False): + trained_step_sizes = random_step_config.get( + "step_sizes", [1] + ) # e.g., [1, 2, 5, 10] + for b_idx, step_size_val in enumerate( + skip_step_size_batch.cpu().numpy() + ): + if step_size_val not in trained_step_sizes: + logger.warning( + f"[Rollout Warning] Sample {b_idx}: step_size={step_size_val} is NOT in trained range {trained_step_sizes}!\n" + f"Model was only trained on step_sizes {trained_step_sizes}. Prediction may be inaccurate due to extrapolation." + ) + elif if_debug: + logger.info( + f"[Rollout Validation] Sample {b_idx}: step_size={step_size_val} is within trained range {trained_step_sizes}" + ) + + B, T0, Nobj, Nv, _ = obj_pos_all.shape + device = obj_pos_all.device + + # Require at least warmup_frames GT frames (minimum 2 for Verlet initialization) + min_required = max(2, warmup_frames) + if T0 < min_required: + raise ValueError( + f"rollout requires at least {min_required} GT frames (warmup_frames={warmup_frames}), got {T0}" + ) + + # Reference pose: MUST be global frame 0 (same as training) + if "ref_obj_positions" not in initial_batch: + raise ValueError( + "rollout requires 'ref_obj_positions' (global frame 0) in initial_batch" + ) + ref_pose = initial_batch["ref_obj_positions"] # (B?,Nobj,Nv,3) + if ref_pose.dim() == 3: + ref_pose = ref_pose.unsqueeze(0).expand(B, -1, -1, -1).to(device) + else: + ref_pose = ref_pose.to(device) + + # ========== Center scene by first frame (eval/rollout alignment) ========== + # If enabled, subtract the global center of the first frame from all coordinates + # so the rollout coordinate system matches the training setup. + # center_t0 is stored per sample so visualization can apply the same shift to disk-loaded GT. + _center_t0_per_sample = ( + [] + ) # Store per-sample center for visualization GT alignment + if self.config.get("data", {}).get("center_scene_by_t0", False): + for b in range(B): + vmask_b0 = valid_vtx_mask_all[b, 0] # (Nobj, Nv) + v_t0 = obj_pos_all[b, 0][vmask_b0] # Valid vertices at t=0 + if v_t0.shape[0] > 0: + center_t0 = v_t0.mean(dim=0) # (3,) + _center_t0_per_sample.append(center_t0.cpu().numpy()) + + logger.info( + f"[Center-T0 ROLLOUT] b={b}: center_t0={center_t0.cpu().numpy().tolist()}, " + f"obj_pos range before: [{obj_pos_all[b].min().item():.4f}, {obj_pos_all[b].max().item():.4f}]" + ) + + # Shift positions + obj_pos_all[b] = obj_pos_all[b] - center_t0 + ref_pose[b] = ref_pose[b] - center_t0 + + logger.info( + f"[Center-T0 ROLLOUT] b={b} after: " + f"obj_pos range: [{obj_pos_all[b].min().item():.4f}, {obj_pos_all[b].max().item():.4f}]" + ) + else: + _center_t0_per_sample.append(None) + logger.info( + f"[Center-T0 ROLLOUT] Applied centering to {len(_center_t0_per_sample)} samples" + ) + + # Require anchors for query and supervision (no fallback) + if "anchor_idx" not in initial_batch: + raise KeyError( + "rollout requires 'anchor_idx' in initial_batch for anchor-based querying" + ) + anchor_idx = initial_batch["anchor_idx"].to(device).long() # (B,Nobj,Nq) + if ( + anchor_idx.dim() != 3 + or anchor_idx.shape[0] != B + or anchor_idx.shape[1] != Nobj + ): + raise ValueError( + f"anchor_idx must be (B,Nobj,Nq); got {tuple(anchor_idx.shape)} vs (B={B},Nobj={Nobj})" + ) + Nq = int(anchor_idx.shape[2]) + + # Initialize previous/current positions (use last 2 frames of warmup for Verlet) + # For warmup_frames=2: use frame 0, 1 + # For warmup_frames=4: use frame 2, 3 (GT frames 0,1,2,3 then predict from frame 4) + prev_idx = warmup_frames - 2 + curr_idx = warmup_frames - 1 + prev_pos_bt = obj_pos_all[:, prev_idx] # (B,Nobj,Nv,3) + curr_pos_bt = obj_pos_all[:, curr_idx] + registration_refs = curr_pos_bt.detach().clone() + + # Initialize persistent state tokens (carry across steps like training) + # When no_global_state=True, pass None to decoder + if self._no_global_state: + state_tokens = None + state_pos = None + else: + state_tokens, state_pos = self.state_repr.get_state_and_positions( + B, device + ) + + # Collect frames: include ALL GT warmup frames first + out_seq = [obj_pos_all[:, i] for i in range(warmup_frames)] + + # Steps to predict beyond the provided GT frames in the output horizon + remaining = max(0, num_steps - warmup_frames) + + # Collect per-step input body_dist_vec2 for visualization (curr frame at each step) + # Only collect if return_extras=True (video gen) to save memory during RMSE eval + all_dist_vec2_seq = ( + [] if return_extras else None + ) # list of (B, Nobj, Nv, 3) + # Collect predicted anchor accelerations (denormalized) for visualization + all_accel_anchor_denorm_seq = ( + [] if return_extras else None + ) # list of (B, Nobj, Nq, 3) + + # Previous step's features for penetration prediction (step>0) + prev_obj_layers_rollout = None + prev_state_layers_rollout = None + + # Optional CUDA-graph fast path (RIGIDFORMER_CUDA_GRAPH=1): after the + # early steps whose vertex-mask index is still advancing, every step + # executes an identical op stream on fixed shapes, so one step is + # captured as a CUDA graph and replayed for the remainder. Cross-step + # tensors are kept in fixed-address buffers that the graph itself + # rotates, which makes each replay self-contained. + use_cuda_graph = ( + os.environ.get("RIGIDFORMER_CUDA_GRAPH", "0") == "1" + and not return_extras + and torch.cuda.is_available() + ) + # Constant scalars used inside the step; created once so the step body + # performs no host-to-device transfers (required for graph capture). + _attn_zero = torch.tensor(0.0, dtype=torch.float32, device=device) + _attn_ninf = torch.tensor(float("-inf"), dtype=torch.float32, device=device) + # Kabsch padding filter: valid_obj_mask is fixed for the whole rollout, + # so the valid-object index set is computed once (nonzero would force a + # host sync every step and cannot run under graph capture). + _kabsch_mask_flat = valid_obj_mask.reshape(B * Nobj) + _kabsch_valid_idx = _kabsch_mask_flat.nonzero(as_tuple=True)[0] + _kabsch_n_valid = int(_kabsch_valid_idx.shape[0]) + use_cuda_graph = use_cuda_graph and bool(getattr(self, "if_roma", False)) + graph_obj = None + graph_obj2 = None + graph_capture_step = T0 - 1 # first step index with the saturated mask + if use_cuda_graph and remaining <= graph_capture_step + 1: + use_cuda_graph = False + + # Progress bar for rollout (only show on rank 0) + import torch.distributed as dist + from tqdm import tqdm + + show_pbar = True + if dist.is_available() and dist.is_initialized(): + show_pbar = dist.get_rank() == 0 + + # Object slots holding a valid object in at least one sample; fixed for the rollout, so + # computed once here (a host sync outside the step loop and any graph capture). + _knn_active_slots_r = ( + torch.nonzero(valid_obj_mask.any(dim=0)).flatten().tolist() + if self._knn_num_groups <= 1 + else None + ) + # Valid-vertex packs, one per mask index the loop uses (min(warmup_frames - 1 + step, T0 - 1)); built here, outside + # the step loop and any graph capture. + _knn_packs_r, _enc_packs_r = {}, {} + if self._knn_packed or self._encoder_packed: + _vr_r = initial_batch.get("vtx_rank") + for _mi in range(1, valid_vtx_mask_all.shape[1]): + _vm_r = valid_vtx_mask_all[:, _mi] + if self._knn_packed: + _knn_packs_r[_mi] = knn_pack(_vm_r, valid_obj_mask) + if self._encoder_packed: + _enc_packs_r[_mi] = EncoderPack( + _vm_r.reshape(B * Nobj, Nv), + _vr_r.reshape(B * Nobj, Nv) if _vr_r is not None else None, + len(self.shared_obj_encoder.hierarchical_layers), + self.shared_obj_encoder.hierarchical_subsets, + 1, + small=self.shared_obj_encoder.subset_small_objects, + ) + + # Anchor validation. Anchors and vertex masks are fixed for the whole + # rollout, so every per-step check is equivalent to checking each mask + # index once up front; this keeps the step loop free of host syncs. + for _mask_idx in range(1, T0): + for i in range(Nobj): + vmask_curr = valid_vtx_mask_all[:, _mask_idx, i] # (B,Nv) + idx_i_chk = anchor_idx[:, i] # (B,Nq) + mask_at_anchors = torch.gather(vmask_curr, dim=1, index=idx_i_chk) + obj_valid_b = valid_obj_mask[:, i] # (B,) + num_valid = vmask_curr.sum(dim=1) # (B,) + if int(((num_valid == 0) & obj_valid_b).sum().item()) > 0: + raise RuntimeError( + f"Zero valid vertices for rollout at mask_idx={_mask_idx}, obj={i} for valid samples" + ) + need = (num_valid < Nq) & obj_valid_b + if bool(need.any()): + min_valid = ( + int(num_valid[obj_valid_b].min().item()) + if bool(obj_valid_b.any()) + else 0 + ) + raise RuntimeError( + f"Not enough valid vertices for rollout at mask_idx={_mask_idx}, obj={i}: min_valid={min_valid} < anchor_k={Nq}" + ) + bad_mask = (~mask_at_anchors) & obj_valid_b.unsqueeze(1) + if bool(bad_mask.any()): + bad = bad_mask.sum().item() + raise RuntimeError( + f"Invalid anchors for rollout: {bad} indices invalid at mask_idx={_mask_idx}, obj={i}." + ) + + pbar = tqdm( + range(remaining), desc="Rollout", disable=not show_pbar, leave=False + ) + for step in pbar: + rope_memo_reset() # rotary angles are memoized within one step only (keeps graph capture self-contained) + _capturing = use_cuda_graph and step == graph_capture_step + if _capturing: + # Move cross-step carries into fixed-address buffers, then + # record this step. Kernels are only recorded during capture, + # so the step is materialized by an immediate replay below. + prev_pos_bt = prev_pos_bt.clone() + curr_pos_bt = curr_pos_bt.clone() + registration_refs = registration_refs.clone() + if prev_obj_layers_rollout is not None: + prev_obj_layers_rollout = [ + x.clone() for x in prev_obj_layers_rollout + ] + if prev_state_layers_rollout is not None: + prev_state_layers_rollout = [ + x.clone() for x in prev_state_layers_rollout + ] + torch.cuda.synchronize() + graph_obj = torch.cuda.CUDAGraph() + _graph_stream = torch.cuda.Stream() + _prev_stream = torch.cuda.current_stream() + torch.cuda.set_stream(_graph_stream) + # thread_local: NCCL watchdog / DataLoader threads must not + # poison the capture under DDP evaluation. + graph_obj.capture_begin(capture_error_mode="thread_local") + # ========== STATE RESET (optional, to avoid train-test mismatch) ========== + # When rollout_state_reset=True, reset state tokens at each step to match + # training distribution (training only sees ~2 state updates per sequence). + # This prevents state drift during long rollouts (e.g., 400 frames). + if self._rollout_state_reset: + if self._no_global_state: + state_tokens = None + state_pos = None + else: + state_tokens, state_pos = ( + self.state_repr.get_state_and_positions(B, device) + ) + + # Compute velocities from last two frames + vel_bt = curr_pos_bt - prev_pos_bt # (B,Nobj,Nv,3) + + # Prepare far-away filler for invalid vertices in reference lists + FAR = 1e6 + + mask_idx_curr = min( + warmup_frames - 1 + step, T0 - 1 + ) # current frame's mask + + # ========== VECTORIZED KNN FOR ALL OBJECTS AT ONCE ========== + # curr_pos_bt: (B, Nobj, Nv, 3) + # valid_vtx_mask_all[:, mask_idx_curr]: (B, Nobj, Nv) + vmask_step = valid_vtx_mask_all[:, mask_idx_curr] # (B, Nobj, Nv) + + # Compute KNN distance vectors for ALL objects simultaneously + # Only computes for valid objects and valid vertices! + if self._knn_packed: + dist_vec_all = compute_knn_packed( + curr_pos_bt, _knn_packs_r[mask_idx_curr] + ) + else: + dist_vec_all = compute_knn_all_objects_vectorized( + curr_pos_bt, + vmask_step, + FAR, + valid_obj_mask=valid_obj_mask, + num_groups=self._knn_num_groups, + active_slots=_knn_active_slots_r, + ) + + # Apply plane constraint adjustment (vectorized) + # Use reshape() instead of view() to handle non-contiguous tensors + curr_pos_flat = curr_pos_bt.reshape(B, Nobj * Nv, 3) + dist_vec_flat = dist_vec_all.reshape(B, Nobj * Nv, 3) + dist_vec2_flat = distance_2_plane( + curr_pos_flat, [0, 0, 0], "z", dist_vec_flat + ) + dist_vec2_bt = dist_vec2_flat.reshape(B, Nobj, Nv, 3) + + # Store raw dist_vec2 for visualization (only if return_extras=True) + if return_extras: + all_dist_vec2_seq.append(dist_vec2_bt) + + # Relative position (vectorized): curr_pos - ref_pose + relpos_cat = curr_pos_bt - ref_pose # (B, Nobj, Nv, 3) + + # Velocity (already in correct shape) + vel_cat = vel_bt # (B, Nobj, Nv, 3) + + # Dist vectors + distvec_cat = dist_vec2_bt # (B, Nobj, Nv, 3) + + # Get current step's vertex mask for all objects (vectorized) + mask_idx_curr_all = mask_idx_curr + vmask_curr_all = valid_vtx_mask_all[:, mask_idx_curr_all] # (B,Nobj,Nv) + vmask_curr_all_3d = vmask_curr_all.unsqueeze(-1) # (B,Nobj,Nv,1) + + # Apply masks BEFORE normalization (vectorized across all objects) + distvec_cat = torch.where( + vmask_curr_all_3d, distvec_cat, torch.zeros_like(distvec_cat) + ) + vel_cat = torch.where( + vmask_curr_all_3d, vel_cat, torch.zeros_like(vel_cat) + ) + relpos_cat = torch.where( + vmask_curr_all_3d, relpos_cat, torch.zeros_like(relpos_cat) + ) + + # Apply network-predicted directional sign to ALL vertices' dist_vec (before normalization) + # Only for step>0 (student forcing steps with predicted positions, may have penetration) + # For step=0, use GT positions from warmup (no penetration, sign=+1) + penetration_enabled = self.config["model"].get( + "if_penetration", + self.config["model"].get("if_directional_distance", False), + ) + if penetration_enabled and step > 0: + # Predict penetration sign using PREVIOUS step's features and CURRENT predicted positions + # curr_pos_bt = blended_next (AFTER Kabsch + intersection handling) + pen_pred_all_vertices = self.predictor.predict_penetration_vertices( + prev_obj_layers_rollout, # Use previous step's decoder outputs + prev_state_layers_rollout, # Use previous step's state features + pos_full=curr_pos_bt, # AFTER Kabsch transform (matches training forward) + ) # (B, Nobj, Nv, 1) in [-1,1] + + # Apply predicted sign to dist_vec for ALL vertices + pen_sign_all = torch.sign( + pen_pred_all_vertices + 1e-6 + ) # (B, Nobj, Nv, 1) in {-1,+1} + distvec_cat = distvec_cat * pen_sign_all # (B, Nobj, Nv, 3) + + # No normalization — raw values, matching the forward path + # (features are masked once, then the penetration sign is applied) + + # Broadcast physics parameters and time_step to vertex dimension if enabled (rollout) + phys_broadcast = None # (B, Nobj, Nv, 3) + time_broadcast = None # (B, Nobj, Nv, 2) + + if self._if_phys_param: + # Broadcast per-object physics params (B,Nobj,3) to all vertices (B,Nobj,Nv,3) + phys_broadcast = phys_params.unsqueeze(2).expand( + B, Nobj, Nv, 3 + ) # (B,Nobj,Nv,3) + # Mask with vertex mask to ensure padding vertices get zero physics params + phys_broadcast = torch.where( + vmask_curr_all_3d, + phys_broadcast, + torch.zeros_like(phys_broadcast), + ) + + if self._if_time_step: + # Broadcast both s and s² to all objects and vertices (B,Nobj,Nv,2) + # Must match the forward path + step_size_s = skip_step_size_batch.float() # (B,) + step_size_sq = (skip_step_size_batch**2).float() # (B,) + time_features = torch.stack( + [step_size_s, step_size_sq], dim=-1 + ) # (B, 2) + time_broadcast = time_features.view(B, 1, 1, 2).expand( + B, Nobj, Nv, 2 + ) # (B,Nobj,Nv,2) + # Mask for consistency + # vmask_curr_all_3d is (B,Nobj,Nv,1), expand to match (B,Nobj,Nv,2) + vmask_time = vmask_curr_all_3d.expand( + -1, -1, -1, 2 + ) # (B,Nobj,Nv,2) + time_broadcast = torch.where( + vmask_time, time_broadcast, torch.zeros_like(time_broadcast) + ) + + # PointNet Encoder: process each object's Nv vertices -> per-object global feature + # Input dimension depends on flags: + # Base: 9D (dist_vec + velocity + rel_pos) + # +if_phys_param=true: +3D + # +if_time_step=true (no FiLM): +2D (s and s²) + # Split into components and reshape to (B*Nobj, Nv, D) for encoder + dist_reshaped = distvec_cat.view(B * Nobj, Nv, 3) + vel_reshaped = vel_cat.view(B * Nobj, Nv, 3) + relpos_reshaped = relpos_cat.view(B * Nobj, Nv, 3) + + # Optional features + + phys_reshaped = None + if self._if_phys_param: + phys_reshaped = phys_broadcast.view(B * Nobj, Nv, 3) + + time_reshaped = None + if self._if_time_step and not self._film_conditioning: + time_reshaped = time_broadcast.view(B * Nobj, Nv, 2) # (s, s²) + + # Prepare vertex mask for encoder pooling + # vmask_curr_all_3d is (B,Nobj,Nv,1), need (B*Nobj, Nv) + vtx_mask_for_encoder = vmask_curr_all_3d.squeeze(-1).reshape( + B * Nobj, Nv + ) # (B*Nobj, Nv) + height_reshaped = None + if self.shared_obj_encoder.height_input: + _h = torch.where( + vmask_curr_all_3d, + curr_pos_bt[..., 2:3], + torch.zeros_like(curr_pos_bt[..., 2:3]), + ) + height_reshaped = _h.reshape(B * Nobj, Nv, 1) + rank_reshaped = None + if self.shared_obj_encoder.hierarchical_subsets: + if "vtx_rank" not in initial_batch: + raise KeyError( + "model.encoder.hierarchical_subsets requires data.fps_order: true (batch lacks 'vtx_rank')" + ) + rank_reshaped = initial_batch["vtx_rank"].reshape(B * Nobj, Nv) + + # Call encoder with all features (encoder will concatenate them internally) + # AVP: capture per-point features if AVP is enabled + obj_features_encoded, _per_vtx_feat, _ = self.shared_obj_encoder( + dist_reshaped, + vel_reshaped, + relpos_reshaped, + phys_params=phys_reshaped, + time_step=time_reshaped, + valid_mask=vtx_mask_for_encoder, + height=height_reshaped, + vtx_rank=rank_reshaped, + pack=( + _enc_packs_r.get(mask_idx_curr_all) + if self._encoder_packed + else None + ), + ) # obj_features_encoded: (B*Nobj, enc_out_dim), _per_vtx_feat: (B*Nobj, backbone_dim, Nv) + + # AVP: Anchor-Vertex Pooling — enrich anchor queries with collision-aware features + _avp_anchor_features = ( + None # Will be (B, Nobj, Nq, avp_output_dim) if AVP enabled + ) + if self.avp is not None and _per_vtx_feat is not None: + # _per_vtx_feat is (B*Nobj, backbone_dim, Nv) — Conv1d channels-first format + # Transpose to (B*Nobj, Nv, backbone_dim) for AVP + _pvf = _per_vtx_feat.transpose(1, 2) # (B*Nobj, Nv, backbone_dim) + # Anchor positions in flat format + _curr_pos_flat = curr_pos_bt.reshape( + B * Nobj, Nv, 3 + ) # (B*Nobj, Nv, 3) + _anchor_pos_flat = torch.gather( + _curr_pos_flat, + dim=1, + index=anchor_idx.reshape(B * Nobj, -1) + .unsqueeze(-1) + .expand(-1, -1, 3), + ) # (B*Nobj, Nq, 3) + _avp_anchor_features = self.avp( + _pvf, _curr_pos_flat, _anchor_pos_flat, vtx_mask_for_encoder + ) # (B*Nobj, Nq, avp_output_dim) + _avp_anchor_features = _avp_anchor_features.view( + B, Nobj, anchor_idx.shape[2], -1 + ) # (B, Nobj, Nq, avp_output_dim) + # Zero out padding objects (proj bias makes padding non-zero; consistent with other features) + _avp_anchor_features = ( + _avp_anchor_features + * valid_obj_mask.unsqueeze(-1).unsqueeze(-1).float() + ) + + obj_features_encoded = obj_features_encoded.view( + B, Nobj, -1 + ) # (B, Nobj, enc_out_dim) + + # Project to token dimension for Decoder + obj_features = self.obj_feature_proj( + obj_features_encoded + ) # (B, Nobj, D) + + # FiLM: modulate encoder features with step-size [s, s²] + if self.encoder_film is not None and skip_step_size_batch is not None: + time_feat_film = torch.stack( + [ + skip_step_size_batch.float(), + (skip_step_size_batch**2).float(), + ], + dim=-1, + ) # (B, 2) + obj_features = self.encoder_film( + obj_features, time_feat_film + ) # (B, Nobj, D) + + # Object anchor positions for RoPE + # Same approach as forward(): gather anchors, flatten to multi-point for RoPE. + # Replace padding objects' positions with zeros before the gather: + # padding objects have anchor_idx=0, which would gather arbitrary values into RoPE inputs. + obj_mask_pos = valid_obj_mask.unsqueeze(-1).unsqueeze( + -1 + ) # (B, Nobj, 1, 1) + safe_curr_pos_bt = torch.where( + obj_mask_pos, curr_pos_bt, torch.zeros_like(curr_pos_bt) + ) + + idx_exp_3_for_decoder = anchor_idx.unsqueeze(-1).expand( + -1, -1, -1, 3 + ) # (B, Nobj, Nq, 3) + obj_anchor_positions = torch.gather( + safe_curr_pos_bt, dim=2, index=idx_exp_3_for_decoder + ) # (B, Nobj, Nq, 3) + # Flatten to (B, Nobj, Nq*3) for multi-point RoPE + Nq = anchor_idx.shape[2] # Number of anchors per object + + # ============== PE Type Routing (rollout) ============== + # valid_vtx_mask_t is needed for geometry encoder + valid_vtx_mask_t = valid_vtx_mask_all[:, mask_idx_curr] # (B, Nobj, Nv) + if self._pe_type in ("arope", "se3"): + # SE(3) uses arope anchors for cross-attention RoPE; SE(3) bias computed separately + obj_anchor_flat = obj_anchor_positions.view( + B, Nobj, Nq * 3 + ) # (B, Nobj, 12) + elif self._pe_type in ("obb", "pca"): + obj_anchor_flat = self._geometry_encoder( + safe_curr_pos_bt, valid_vtx_mask_t, obj_mask=valid_obj_mask + ) # (B, Nobj, 12) + else: + raise ValueError(f"Unknown pe_type='{self._pe_type}' in rollout().") + + # ============== SE(3) Pairwise Bias Computation (rollout) ============== + _se3_bias = None + if self._pe_type == "se3" and self._se3_bias_module is not None: + from embodichain_models.dynamics.rigidformer._native.models.se3_attention import ( + extract_object_pose_features, + ) + + _rotmats, _centers, _scales = extract_object_pose_features( + safe_curr_pos_bt, ref_pose, valid_vtx_mask_t, valid_obj_mask + ) + _se3_bias = self._se3_bias_module( + _rotmats, _centers, _scales, valid_obj_mask + ) + + # Decoder: SDPA additive mask + feature zeroing (same as forward) + obj_attn_mask = torch.where( + valid_obj_mask.unsqueeze(1).unsqueeze(1), _attn_zero, _attn_ninf + ) + pad_mask_float = valid_obj_mask.unsqueeze(-1).to( + dtype=obj_features.dtype + ) + obj_features = obj_features * pad_mask_float + obj_anchor_flat = obj_anchor_flat * valid_obj_mask.unsqueeze(-1).to( + dtype=obj_anchor_flat.dtype + ) + + # Per-layer FiLM: pass time_features for step-size conditioning in each layer + _decoder_time_feat = ( + time_feat_film + if (self._film_conditioning and self._if_time_step) + else None + ) + all_state_layers, all_obj_layers = self.decoder( + state_tokens, + obj_features, + state_pos, + obj_anchor_flat, + mask=obj_attn_mask, + se3_bias=_se3_bias, + time_features=_decoder_time_feat, # Per-layer FiLM + ) + # No need to remove any tokens - all_obj_layers directly contains object features + + # Update state (skip when no_global_state) + if not self._no_global_state: + state_tokens = all_state_layers[-1] + + # Build anchor query for Predictor (sample at anchor points Nq) + # Query dimension depends on flags: + # Base: 9D (dist_vec + vel + rel_pos) + # +if_phys_param=true: +3D = 12D + # +if_time_step=true (no FiLM): +2D (s and s²) + idx_exp_3 = anchor_idx.unsqueeze(-1).expand( + -1, -1, -1, 3 + ) # (B,Nobj,Nq,3) + distvec_at_anchors = torch.gather(distvec_cat, dim=2, index=idx_exp_3) + vel_at_anchors = torch.gather(vel_cat, dim=2, index=idx_exp_3) + relpos_at_anchors = torch.gather(relpos_cat, dim=2, index=idx_exp_3) + + # Build query features list + query_features = [distvec_at_anchors, vel_at_anchors, relpos_at_anchors] + + # Add physics parameters at anchors if enabled + if self._if_phys_param: + phys_at_anchors = torch.gather( + phys_broadcast, dim=2, index=idx_exp_3 + ) # (B,Nobj,Nq,3) + query_features.append(phys_at_anchors) + + # Add time step at anchors if enabled AND not using FiLM + if self._if_time_step and not self._film_conditioning: + idx_exp_2 = anchor_idx.unsqueeze(-1).expand( + -1, -1, -1, 2 + ) # (B,Nobj,Nq,2) + time_at_anchors = torch.gather( + time_broadcast, dim=2, index=idx_exp_2 + ) # (B,Nobj,Nq,2) + query_features.append(time_at_anchors) + + # AVP: Append collision-aware anchor features to query + if _avp_anchor_features is not None: + query_features.append( + _avp_anchor_features + ) # (B, Nobj, Nq, avp_output_dim) + + # Concatenate all query features + anchor_query = torch.cat( + query_features, dim=-1 + ) # (B,Nobj,Nq, 9/12 + avp_dim) + + # Extract anchor positions for RoPE (current frame positions at anchors) + anchor_pos_at_curr = torch.gather( + curr_pos_bt, dim=2, index=idx_exp_3 + ) # (B,Nobj,Nq,3) + + # Penetration sign was already applied to all vertices' dist_vec; + # the anchor query therefore carries the sign information as well. + + _ref_src = ( + _object_centered(ref_pose, vmask_step.bool()) + if self._ref_pose_object_centered + else ref_pose + ) + if self._ref_pose_object_centered and not getattr( + self, "_refpose_logged_roll", False + ): + self._refpose_logged_roll = True + print( + f"[RefPose] object-centred (x/y only, z absolute) reference pose ACTIVE (rollout): |shift| max={((ref_pose - _ref_src).abs().max()).item():.3f} m", + flush=True, + ) + ref_pose_at_anchors = torch.gather( + _ref_src, dim=2, index=idx_exp_3 + ) # (B,Nobj,Nq,3) + + # Predict accelerations at anchors using grouped format (B, Nobj, Nq, *) + accel_norm_anchor = self.predictor( + all_obj_layers, + anchor_query, # (B, Nobj, Nq, 9/12) + ref_pose=ref_pose_at_anchors, # (B, Nobj, Nq, 3) + query_pos=anchor_pos_at_curr, # (B, Nobj, Nq, 3) + key_pos=obj_anchor_flat, # (B, Nobj, Nq*3) - multi-point RoPE + multi_scale_state_features=( + None if self._no_global_state else all_state_layers + ), + state_pos=None if self._no_global_state else state_pos, + obj_mask=valid_obj_mask, # Object validity mask for padding + ) # Returns (B, Nobj, Nq, 3) + + # Force FP32 for Verlet integration and position computation + # (must match the training forward path for numerical stability) + with torch.amp.autocast("cuda", enabled=False): + accel_base_denorm = accel_norm_anchor.float() + + # Explicit step_size² scaling for multi-timescale rollout + # (must match the forward path) + if self._if_time_step and skip_step_size_batch is not None: + step_size_sq = (skip_step_size_batch.float() ** 2).view( + B, 1, 1, 1 + ) # (B,1,1,1) FP32 + accel_scaled_denorm = ( + accel_base_denorm * step_size_sq + ) # FP32 * FP32 = FP32 + else: + accel_scaled_denorm = accel_base_denorm + + # Store scaled acceleration for visualization (matches forward behavior) + if return_extras: + all_accel_anchor_denorm_seq.append(accel_scaled_denorm) + + # Verlet integration on anchors; initialize full next positions as zero-accel extrapolation + dt_squared = self.dt_verlet * self.dt_verlet + # Positions must be FP32 for numerical stability + curr_pos_bt_f = curr_pos_bt.float() + prev_pos_bt_f = prev_pos_bt.float() + next_pos_bt = ( + 2 * curr_pos_bt_f - prev_pos_bt_f + ) # (B,Nobj,Nv,3) - FP32 + + # Vectorized scatter using scatter_reduce (same optimization as forward) + # Step 1: Gather all anchor positions at once (use FP32 positions) + anchor_idx_exp = anchor_idx.unsqueeze(-1).expand( + -1, -1, -1, 3 + ) # (B,Nobj,Nq,3) + curr_anchors = torch.gather( + curr_pos_bt_f, dim=2, index=anchor_idx_exp + ) # (B,Nobj,Nq,3) + prev_anchors = torch.gather( + prev_pos_bt_f, dim=2, index=anchor_idx_exp + ) # (B,Nobj,Nq,3) + + # Step 2: Compute next positions at anchors using Verlet (with scaled acceleration) + next_anchors = ( + accel_scaled_denorm * dt_squared + + 2 * curr_anchors + - prev_anchors + ) # (B,Nobj,Nq,3) + + # Step 3: Compute updates (delta from base prediction) + base_anchors = torch.gather( + next_pos_bt, dim=2, index=anchor_idx_exp + ) # (B,Nobj,Nq,3) + anchor_updates = next_anchors - base_anchors # (B,Nobj,Nq,3) + + # Mask padding objects' updates to zero + # Padding objects have anchor_idx=0 which would scatter to wrong positions + obj_mask_scatter = ( + valid_obj_mask.unsqueeze(-1).unsqueeze(-1).float() + ) # (B,Nobj,1,1) + anchor_updates = ( + anchor_updates * obj_mask_scatter + ) # padding objects -> zero updates + + # Step 4: Vectorized scatter using scatter_reduce + # Flatten batch and object dimensions for efficient scatter + # Use reshape() to handle non-contiguous tensors + B_Nobj_Nv = B * Nobj * Nv + next_pos_flat = next_pos_bt.reshape(B_Nobj_Nv, 3) # (B*Nobj*Nv, 3) + + # Create flat indices for scatter_reduce + obj_offsets = ( + torch.arange(Nobj, device=anchor_idx.device).view(1, Nobj, 1) + * Nv + ) # (1,Nobj,1) + batch_offsets = torch.arange(B, device=anchor_idx.device).view( + B, 1, 1 + ) * ( + Nobj * Nv + ) # (B,1,1) + flat_anchor_idx = ( + anchor_idx + obj_offsets + batch_offsets + ) # (B,Nobj,Nq) + flat_anchor_idx = flat_anchor_idx.reshape( + B * Nobj * Nq + ) # (B*Nobj*Nq,) + + # Flatten updates + anchor_updates_flat = anchor_updates.reshape( + B * Nobj * Nq, 3 + ) # (B*Nobj*Nq, 3) + + # Apply updates using scatter_reduce (more efficient than loop) + next_pos_flat = torch.scatter_reduce( + next_pos_flat, + dim=0, + index=flat_anchor_idx.unsqueeze(-1).expand(-1, 3), + src=anchor_updates_flat, + reduce="sum", + include_self=True, + ) + + # Reshape back + next_pos_bt = next_pos_flat.reshape(B, Nobj, Nv, 3) + + # Per-object rigid projection using anchors; apply to full Nv + # Vectorized Kabsch: Process ONLY valid objects (filter padding BEFORE SVD) + # Step 1: Gather anchor positions for all objects at once + anchor_idx_exp = anchor_idx.unsqueeze(-1).expand( + -1, -1, -1, 3 + ) # (B,Nobj,Nq,3) + last_anchor_all = torch.gather( + registration_refs, dim=2, index=anchor_idx_exp + ) # (B,Nobj,Nq,3) + curr_anchor_all = torch.gather( + next_pos_bt, dim=2, index=anchor_idx_exp + ) # (B,Nobj,Nq,3) + + # Reshape to (B*Nobj, Nq, 3) for batch processing + B_Nobj = B * Nobj + last_anchor_flat = last_anchor_all.view(B_Nobj, Nq, 3) + curr_anchor_flat = curr_anchor_all.view(B_Nobj, Nq, 3) + + # Filter padding objects before SVD; they would produce garbage R/t + # (indices precomputed before the loop; the mask never changes) + valid_idx = _kabsch_valid_idx + n_valid = _kabsch_n_valid + + # Initialize R/t for ALL objects (identity R, zero t) + # FP32 so the RoMa FP32 output is not downcast during scatter + R_flat = ( + torch.eye(3, device=device, dtype=torch.float32) + .unsqueeze(0) + .expand(B_Nobj, -1, -1) + .clone() + ) + t_vec_flat = torch.zeros( + B_Nobj, 1, 3, device=device, dtype=torch.float32 + ) + + if n_valid > 0: + # Gather only valid objects' anchors + last_anchor_valid = last_anchor_flat[valid_idx] # (Nvalid, Nq, 3) + curr_anchor_valid = curr_anchor_flat[valid_idx] # (Nvalid, Nq, 3) + + if _capturing: + # The SVD inside the rigid solve is not graph-capturable, so + # the step is recorded as two graphs around an eager solve. + # Close graph 1 here and execute it to materialize its outputs. + graph_obj.capture_end() + graph_obj.replay() + torch.cuda.current_stream().synchronize() + + if self.if_roma: + # RoMa mode: batch process only valid objects + # RoMa requires FP32 for numerical stability + with torch.amp.autocast("cuda", enabled=False): + R_valid, t_vec_valid = roma.rigid_points_registration( + last_anchor_valid.float(), curr_anchor_valid.float() + ) + R_valid = R_valid.float() # Ensure FP32 for downstream + t_vec_valid = t_vec_valid.float().unsqueeze(1) # (Nvalid, 1, 3) + + else: + # Standard SVD Kabsch: batch process only valid objects + centroid_last = last_anchor_valid.mean( + dim=1, keepdim=True + ) # (Nvalid,1,3) + centroid_curr = curr_anchor_valid.mean( + dim=1, keepdim=True + ) # (Nvalid,1,3) + last_c = last_anchor_valid - centroid_last + curr_c = curr_anchor_valid - centroid_curr + + # Batch SVD for valid objects only + H = torch.bmm(last_c.transpose(1, 2), curr_c) # (Nvalid,3,3) + with torch.amp.autocast("cuda", enabled=False): + H_float = H.float() + U, S, Vt = torch.linalg.svd(H_float) # Batch SVD! + U = U.to(H.dtype) + Vt = Vt.to(H.dtype) + R_valid = torch.bmm( + Vt.transpose(1, 2), U.transpose(1, 2) + ) # (Nvalid,3,3) + + # Fix reflections + with torch.amp.autocast("cuda", enabled=False): + det_R = torch.det(R_valid.float()).to(R_valid.dtype) + need_flip = det_R < 0 + if need_flip.any(): + Fm = ( + torch.eye(3, device=R_valid.device, dtype=R_valid.dtype) + .unsqueeze(0) + .expand(n_valid, -1, -1) + .clone() + ) + Fm[need_flip, 2, 2] = -1 + R_valid = torch.bmm( + torch.bmm(Vt.transpose(1, 2), Fm), U.transpose(1, 2) + ) + + # Translation + rotated_centroid = torch.bmm( + R_valid, centroid_last.transpose(1, 2) + ).transpose(1, 2) + t_vec_valid = centroid_curr - rotated_centroid # (Nvalid,1,3) + + if _capturing: + # Fixed-address buffers carry the solve result into graph 2. + _rt_R_buf = R_valid.clone() + _rt_t_buf = t_vec_valid.clone() + R_valid = _rt_R_buf + t_vec_valid = _rt_t_buf + graph_obj2 = torch.cuda.CUDAGraph() + graph_obj2.capture_begin( + pool=graph_obj.pool(), capture_error_mode="thread_local" + ) + elif graph_obj2 is not None: + _rt_R_buf.copy_(R_valid) + _rt_t_buf.copy_(t_vec_valid) + R_valid = _rt_R_buf + t_vec_valid = _rt_t_buf + + # Scatter valid R/t back to full tensors + R_flat[valid_idx] = R_valid + t_vec_flat[valid_idx] = t_vec_valid + + # Reshape back to (B,Nobj,3,3) and (B,Nobj,1,3) + R_all = R_flat.view(B, Nobj, 3, 3) + t_vec_all = t_vec_flat.view(B, Nobj, 1, 3) + + # ========== VECTORIZED TRANSFORMATION APPLICATION ========== + # Apply R, t to all objects in a single bmm call (no per-object loop) + # registration_refs: (B, Nobj, Nv, 3) + # R_all: (B, Nobj, 3, 3) + # t_vec_all: (B, Nobj, 1, 3) + + # Reshape for batch bmm: (B, Nobj, Nv, 3) -> (B*Nobj, Nv, 3) + last_flat = registration_refs.view(B * Nobj, Nv, 3) + R_flat = R_all.view(B * Nobj, 3, 3) + t_flat = t_vec_all.view(B * Nobj, 1, 3) + + # Single bmm for ALL objects: y = R @ x + t (RoMa convention) + # Convert to FP32 to prevent BF16 precision loss during autoregressive rollout + y_flat = ( + torch.bmm(R_flat, last_flat.transpose(1, 2)).transpose(1, 2) + + t_flat + ) + # Use reshape() because transpose() makes tensor non-contiguous + yproj_bt = y_flat.reshape( + B, Nobj, Nv, 3 + ).float() # FP32 for numerical stability + + # Apply gradient strategy based on transform type + if self.if_roma: + # RoMa mode: fully differentiable, gradients flow through Kabsch + blended_next = yproj_bt + else: + # Standard mode: straight-through estimator (STE) + blended_next = next_pos_bt + (yproj_bt - next_pos_bt).detach() + + # ========== INTERSECTION HANDLING (EVAL) ========== + intersection_cfg = self.config.get("model", {}).get( + "intersection_handling", {} + ) + blended_next = handle_intersections( + blended_next, valid_obj_mask, vmask_step, intersection_cfg + ) + + if _capturing: + # Rotate the carry buffers in place inside the captured region + # so every replay advances the state by one step. + prev_pos_bt.copy_(curr_pos_bt) + curr_pos_bt.copy_(blended_next) + if not self._kabsch_reference_fixed: + registration_refs.copy_(blended_next) + if prev_obj_layers_rollout is not None: + for _buf, _new in zip(prev_obj_layers_rollout, all_obj_layers): + _buf.copy_(_new) + if prev_state_layers_rollout is not None: + for _buf, _new in zip( + prev_state_layers_rollout, all_state_layers + ): + _buf.copy_(_new) + graph_obj2.capture_end() + torch.cuda.set_stream(_prev_stream) + torch.cuda.synchronize() + graph_obj2.replay() # capture only records; execute part 2 now + torch.cuda.synchronize() + out_seq.append(curr_pos_bt.clone()) + pbar.close() + break + # Registration refs use the post-collision position, matching the forward path ('advancing'); + # model.kabsch_reference 'fixed' keeps the last warm-up frame as the reference for every step + if not self._kabsch_reference_fixed: + registration_refs = blended_next.detach() + + out_seq.append(blended_next) + + # Save current features for next step's penetration prediction + prev_obj_layers_rollout = all_obj_layers + prev_state_layers_rollout = all_state_layers + + prev_pos_bt = curr_pos_bt + curr_pos_bt = blended_next + + if graph_obj is not None: + target_len = warmup_frames + remaining + while len(out_seq) < target_len: + graph_obj.replay() + with torch.amp.autocast("cuda", enabled=False): + _R_new, _t_new = roma.rigid_points_registration( + last_anchor_valid.float(), curr_anchor_valid.float() + ) + _rt_R_buf.copy_(_R_new.float()) + _rt_t_buf.copy_(_t_new.float().unsqueeze(1)) + graph_obj2.replay() + out_seq.append(curr_pos_bt.clone()) + + if graph_obj is not None: + # Tear the graphs down and return their private pools to the + # allocator: the evaluator calls rollout once per batch and per + # step size, and retired pools would otherwise accumulate. + graph_obj.reset() + if graph_obj2 is not None: + graph_obj2.reset() + graph_obj = None + graph_obj2 = None + if os.environ.get("RIGIDFORMER_GRAPH_KEEPCACHE", "0") != "1": + torch.cuda.empty_cache() + + full = torch.stack(out_seq, dim=1) # (B, 2+remaining, Nobj, Nv, 3) + + # Single numerical-health check per rollout (one host sync total, not + # one per step): any non-finite prediction aborts loudly here. + if not torch.isfinite(full).all(): + raise RuntimeError( + "Non-finite values in rollout output; " + "inspect the checkpoint and input data." + ) + + # ========== Map back to original absolute space ========== + # If center_scene_by_t0 was applied, add center_t0 back to predictions + # so that the output is in the original world coordinate frame; + # centering remains a purely internal operation of the model. + if ( + self.config.get("data", {}).get("center_scene_by_t0", False) + and len(_center_t0_per_sample) > 0 + ): + for b in range(B): + if _center_t0_per_sample[b] is not None: + center_t0_tensor = torch.from_numpy( + _center_t0_per_sample[b] + ).to(full.device) + full[b] = full[b] + center_t0_tensor + logger.info( + f"[Center-T0 ROLLOUT] Mapped predictions back to original absolute space" + ) + + output = {"obj_positions": full} + + # Return body_dist_vec2 and accelerations only if requested (for visualization) + # Skip during RMSE eval to save ~4GB memory on long rollouts + if return_extras and remaining > 0: + output["obj_body_dist_vec2"] = torch.stack(all_dist_vec2_seq, dim=1) + output["obj_accelerations_anchor"] = torch.stack( + all_accel_anchor_denorm_seq, dim=1 + ) + + return output diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/utils/__init__.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/utils/__init__.py new file mode 100644 index 0000000000..9071a4c0b7 --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/utils/__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__ = [] diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/utils/anchor_rotation.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/utils/anchor_rotation.py new file mode 100644 index 0000000000..3cca34ec4d --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/utils/anchor_rotation.py @@ -0,0 +1,31 @@ +# 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. + +"""Rotation part of an anchor-position error, in units of the object's size.""" + +from __future__ import annotations +import torch + + +def centered_anchor_error( + pred: torch.Tensor, gt: torch.Tensor, min_radius: float = 1e-3 +) -> torch.Tensor: + """pred, gt: (..., Nq, 3) anchor positions of rigid objects. Returns (..., Nq, 3): the error without its mean over the + anchors (the translation part), divided by the RMS distance of the ground-truth anchors from their centroid. + + For a rotation error of angle theta and no other error, the RMS norm of the result is 2 sin(theta / 2) times the + RMS component of the centred anchors perpendicular to the rotation axis over their RMS radius: about theta, and the + same for a small and a large object.""" + gt_c = gt - gt.mean(dim=-2, keepdim=True) + pred_c = pred - pred.mean(dim=-2, keepdim=True) + radius = ( + gt_c.pow(2) + .sum(dim=-1) + .mean(dim=-1, keepdim=True) + .sqrt() + .unsqueeze(-1) + .clamp_min(min_radius) + ) + return (pred_c - gt_c) / radius diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/utils/cuda_knn.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/utils/cuda_knn.py new file mode 100644 index 0000000000..e9ff1cb077 --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/utils/cuda_knn.py @@ -0,0 +1,671 @@ +# 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. + +from __future__ import annotations + +import torch +import torch.nn.functional as F +from typing import List, Tuple, Optional + +# CUDA is mandatory for this module; there is intentionally no CPU fallback. +assert torch.cuda.is_available(), "cuda_knn requires a CUDA-capable device." + + +def compute_nearest_distance_to_list_cuda( + tensor_a: torch.Tensor, + tensor_list: List[torch.Tensor], + batch_mode: bool = False, + return_vectors: bool = False, + return_tensor_indices: bool = True, +) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + """ + CUDA-accelerated parallel computation of nearest distances from query + points to a list of reference point sets. All tensors must reside on the + same device. + + Args: + tensor_a: Query tensor + - If batch_mode=False: shape [N, D] + - If batch_mode=True: shape [B, N, D] + tensor_list: List of reference tensors + - If batch_mode=False: [M_i, D] each + - If batch_mode=True: [B, M_i, D] each + batch_mode: Whether to use batch processing + return_vectors: If True, also return vectors from query to nearest points + return_tensor_indices: If False, skip the per-list-entry loop that maps the nearest point to + the index of its list entry (returned as None); distances and vectors are unaffected + + Returns: + min_distances: Minimum distances [N] or [B, N] + min_indices: Tensor list indices [N] or [B, N] + vectors: (optional) Vectors from query to nearest points [N, D] or [B, N, D] + vectors = nearest_point - query_point + """ + # All tensors must reside on the same device + device = tensor_a.device + for i, t in enumerate(tensor_list): + if t.device != device: + raise RuntimeError( + f"All tensors must be on the same device: tensor_a is on {device}, tensor_list[{i}] is on {t.device}" + ) + + if len(tensor_list) == 0: + raise ValueError("tensor_list cannot be empty") + + device = tensor_a.device + original_batch_mode = batch_mode + + if not batch_mode: + # Add batch dimension for unified processing + tensor_a = tensor_a.unsqueeze(0) # [1, N, D] + tensor_list = [t.unsqueeze(0) for t in tensor_list] # Each becomes [1, M_i, D] + + B, N, D = tensor_a.shape + + # Strategy 1: Concatenate all references for single parallel computation + ref_sizes = [ref.shape[1] for ref in tensor_list] + all_refs = torch.cat(tensor_list, dim=1) # [B, total_M, D] + + # Single CUDA-parallel distance computation + dist_matrix = torch.cdist(tensor_a, all_refs, p=2.0) # [B, N, total_M] + + # Find minimum distances and indices in parallel + min_dists, min_idx = dist_matrix.min(dim=2) # Both [B, N] + + # Convert concatenated indices to tensor list indices + min_tensor_indices = None + if return_tensor_indices: + min_tensor_indices = torch.zeros_like(min_idx, dtype=torch.long) + cumsum = 0 + for i, size in enumerate(ref_sizes): + mask = (min_idx >= cumsum) & (min_idx < cumsum + size) + min_tensor_indices[mask] = i + cumsum += size + + # Compute vectors if requested + vectors = None + if return_vectors: + # Get the nearest points using advanced indexing + batch_indices = torch.arange(B, device=device).view(B, 1).expand(B, N) + nearest_points = all_refs[batch_indices, min_idx] # [B, N, D] + # Vector from query to nearest point + vectors = nearest_points - tensor_a # [B, N, D] + + if not original_batch_mode: + min_dists = min_dists.squeeze(0) + if min_tensor_indices is not None: + min_tensor_indices = min_tensor_indices.squeeze(0) + if vectors is not None: + vectors = vectors.squeeze(0) + + if return_vectors: + return min_dists, min_tensor_indices, vectors + else: + return min_dists, min_tensor_indices + + +def compute_knn_all_objects_vectorized( + positions: torch.Tensor, + valid_vtx_mask: torch.Tensor, + FAR: float = 1e6, + valid_obj_mask: torch.Tensor = None, + num_groups: int = 1, + active_slots: Optional[List[int]] = None, +) -> torch.Tensor: + """ + Efficient batch-mode KNN for multi-object collision detection. + + Args: + positions: (B, Nobj, Nv, 3) - vertex positions for all objects + valid_vtx_mask: (B, Nobj, Nv) - True for valid vertices + FAR: Distance offset for invalid vertices (default 1e6) + valid_obj_mask: (B, Nobj) - True for valid objects + num_groups: Number of groups to split objects into for GPU-parallel processing. + num_groups=1: per-object loop (Nobj iterations) + num_groups>1: GPU-parallel within each group (num_groups iterations) + Example: Nobj=564, num_groups=6 -> 94 objects processed in parallel per group + active_slots: Optional list of object slots that hold a valid object in at least one batch element + (per-object path only). Slots outside it are padding in every sample and their output + is zeroed either way, so they are not searched; valid objects still search every other slot. + + Returns: + dist_vectors: (B, Nobj, Nv, 3) - vector from each vertex to nearest point on OTHER objects + """ + B, Nobj, Nv, D = positions.shape + device = positions.device + + if Nobj < 2: + raise RuntimeError( + "compute_knn_all_objects_vectorized requires at least 2 objects" + ) + + if valid_obj_mask is None: + valid_obj_mask = valid_vtx_mask.any(dim=2) # (B, Nobj) + + dist_vectors = torch.zeros(B, Nobj, Nv, D, device=device) + + # Pre-compute masked positions (FAR offset for invalid vertices) + masked_positions = positions.clone() + vtx_mask_3d = valid_vtx_mask.unsqueeze(-1) # (B, Nobj, Nv, 1) + masked_positions = torch.where( + vtx_mask_3d, masked_positions, masked_positions + FAR + ) + + # ========== GPU-PARALLEL GROUPED PROCESSING ========== + if num_groups > 1 and Nobj > num_groups: + group_size = (Nobj + num_groups - 1) // num_groups + + # All reference points (all objects, masked): (B, Nobj*Nv, 3) + all_refs_flat = masked_positions.reshape(B, Nobj * Nv, D) + + for group_idx in range(num_groups): + g_start = group_idx * group_size + g_end = min(g_start + group_size, Nobj) + G = g_end - g_start # Objects in this group + + # Query points for this group: (B, G, Nv, 3) -> (B, G*Nv, 3) + group_queries = positions[:, g_start:g_end].reshape(B, G * Nv, D) + + # Compute full distance matrix: (B, G*Nv, Nobj*Nv) + dist_matrix = torch.cdist(group_queries, all_refs_flat, p=2.0) + + # Create self-exclusion mask: for each query vertex, mask out vertices from same object + # Query vertex q belongs to object (g_start + q // Nv) + # Ref vertex r belongs to object (r // Nv) + query_obj_idx = ( + torch.arange(G, device=device).unsqueeze(1).expand(G, Nv).reshape(-1) + + g_start + ) # (G*Nv,) + ref_obj_idx = ( + torch.arange(Nobj, device=device) + .unsqueeze(1) + .expand(Nobj, Nv) + .reshape(-1) + ) # (Nobj*Nv,) + + # Self mask: (G*Nv, Nobj*Nv), True where query and ref belong to same object + self_mask = query_obj_idx.unsqueeze(1) == ref_obj_idx.unsqueeze( + 0 + ) # (G*Nv, Nobj*Nv) + + # Apply FAR to self-distances so they're not selected as minimum + dist_matrix = dist_matrix + self_mask.float() * FAR # (B, G*Nv, Nobj*Nv) + + # Find minimum distance and corresponding ref index + min_dists, min_indices = dist_matrix.min(dim=2) # (B, G*Nv) + + # Get nearest points using advanced indexing + batch_idx = torch.arange(B, device=device).view(B, 1).expand(B, G * Nv) + nearest_points = all_refs_flat[batch_idx, min_indices] # (B, G*Nv, 3) + + # Compute vectors from query to nearest + vectors = nearest_points - group_queries # (B, G*Nv, 3) + + # Reshape back and store: (B, G*Nv, 3) -> (B, G, Nv, 3) + dist_vectors[:, g_start:g_end] = vectors.reshape(B, G, Nv, D) + + else: + # ========== PER-OBJECT PROCESSING ========== + # The references of object i are all other slots in slot order: the flattened (B, Nobj*Nv, D) + # tensor without object i's block, the same tensor as concatenating the other slots one by one. + refs_flat = masked_positions.reshape(B, Nobj * Nv, D) + for i in (range(Nobj) if active_slots is None else active_slots): + body_pos = positions[:, i] # (B, Nv, 3) + ref_list = [refs_flat[:, : i * Nv], refs_flat[:, (i + 1) * Nv :]] + + _, _, body_dist_vec = compute_nearest_distance_to_list_cuda( + body_pos, + ref_list, + batch_mode=True, + return_vectors=True, + return_tensor_indices=False, + ) + dist_vectors[:, i] = body_dist_vec + + # Zero out invalid query vertices + vtx_mask_expanded = valid_vtx_mask.unsqueeze(-1) + dist_vectors = torch.where( + vtx_mask_expanded, dist_vectors, torch.zeros_like(dist_vectors) + ) + + # Zero out invalid objects + obj_mask_expanded = valid_obj_mask.unsqueeze(-1).unsqueeze(-1) + dist_vectors = torch.where( + obj_mask_expanded, dist_vectors, torch.zeros_like(dist_vectors) + ) + + return dist_vectors + + +def knn_pack( + valid_vtx_mask: torch.Tensor, valid_obj_mask: Optional[torch.Tensor] = None +) -> dict: + """Packed layout of every sample's valid vertices for compute_knn_packed. + + The masks do not change over a batch or a rollout, so this is computed once (one host sync for the packed + width) and reused at every step. valid_vtx_mask (B, Nobj, Nv); valid_obj_mask (B, Nobj) or None. + Packed order = slot order, then vertex order: the order in which the per-object path lists candidates. + """ + B, Nobj, Nv = valid_vtx_mask.shape + vm = valid_vtx_mask.bool() + if valid_obj_mask is not None: + vm = vm & valid_obj_mask.bool().unsqueeze(-1) + n_obj = vm.any(dim=2).sum(dim=1) + if bool((n_obj < 2).any()): + raise ValueError( + "compute_knn_packed: every sample needs at least two objects with valid vertices" + ) + flat = vm.reshape(B, Nobj * Nv) + counts = flat.sum(dim=1) + n = int(counts.max()) + order = torch.argsort((~flat).to(torch.int8), dim=1, stable=True)[ + :, :n + ] # (B, n) valid first + valid = torch.arange(n, device=flat.device).unsqueeze(0) < counts.unsqueeze( + 1 + ) # (B, n) + obj = torch.where( + valid, torch.div(order, Nv, rounding_mode="floor"), torch.full_like(order, -1) + ) + return {"order": order, "valid": valid, "obj": obj, "shape": (B, Nobj, Nv)} + + +def compute_knn_packed( + positions: torch.Tensor, pack: dict, block_elems: int = 1 << 27 +) -> torch.Tensor: + """Vector from every valid vertex to the nearest valid vertex of any OTHER object, over the packed valid + vertices only (padded vertices and padded object slots never enter the distance matrix). + + Same candidates in the same order and the same first-minimum tie-breaking as compute_knn_all_objects_vectorized, + and the squared distance is cdist's own construction ([-2x, |x|^2, 1] . [y, 1, |y|^2]) without the clamp and + square root over the whole matrix (monotone, so the argmin is unchanged). Results agree with the per-object + path up to the floating-point rounding of the distance values (a near-tie may pick a different, equally near + vertex). positions (B, Nobj, Nv, D); pack from knn_pack; returns (B, Nobj, Nv, D), zero at invalid entries. + """ + B, Nobj, Nv = pack["shape"] + D = positions.shape[-1] + order, valid, obj = pack["order"], pack["valid"], pack["obj"] + n = order.shape[1] + with torch.autocast(device_type="cuda", enabled=False): + flat = positions.float().reshape(B, Nobj * Nv, D) + P = torch.gather(flat, 1, order.unsqueeze(-1).expand(B, n, D)) # (B, n, D) + sq = P.pow(2).sum(dim=-1, keepdim=True) + ones = torch.ones_like(sq) + Xq = torch.cat([P.mul(-2), sq, ones], dim=-1) # (B, n, D+2) + Yr = torch.cat([P, ones, sq], dim=-1).transpose(1, 2) # (B, D+2, n) + rows = max(1, min(n, block_elems // max(1, B * n))) + idx = torch.empty(B, n, dtype=torch.long, device=P.device) + for s0 in range(0, n, rows): + s1 = min(n, s0 + rows) + d2 = torch.bmm(Xq[:, s0:s1], Yr) # (B, rows, n) + d2.masked_fill_( + (obj[:, s0:s1, None] == obj[:, None, :]) | ~valid[:, None, :], + float("inf"), + ) + idx[:, s0:s1] = d2.argmin(dim=2) + nearest = torch.gather(P, 1, idx.unsqueeze(-1).expand(B, n, D)) + vec = torch.where( + valid.unsqueeze(-1), nearest - P, torch.zeros((), device=P.device) + ) + out = torch.zeros(B, Nobj * Nv, D, device=P.device, dtype=torch.float32) + out.scatter_(1, order.unsqueeze(-1).expand(B, n, D), vec) + return out.reshape(B, Nobj, Nv, D).to(positions.dtype) + + +def compute_nearest_distance_to_list_cuda_chunked( + tensor_a: torch.Tensor, tensor_list: List[torch.Tensor], chunk_size: int = 10000 +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Memory-efficient CUDA version using chunked processing for large tensors. + Processes distance computation in chunks to avoid OOM on large datasets. + + Args: + tensor_a: Query tensor on CUDA, shape [N, D] + tensor_list: List of reference tensors on CUDA, each [M_i, D] + chunk_size: Maximum number of reference points to process at once + + Returns: + min_distances: shape [N] + min_tensor_indices: shape [N] + """ + if not tensor_a.is_cuda: + raise RuntimeError("tensor_a must be on CUDA; CPU execution is not supported.") + + for i, t in enumerate(tensor_list): + if not t.is_cuda: + raise RuntimeError( + f"tensor_list[{i}] must be on CUDA; CPU execution is not supported." + ) + + N, D = tensor_a.shape + device = tensor_a.device + + # Initialize results + global_min_dists = torch.full((N,), float("inf"), device=device) + global_tensor_idx = torch.zeros(N, dtype=torch.long, device=device) + + for tensor_idx, ref_tensor in enumerate(tensor_list): + M = ref_tensor.shape[0] + + # Process this reference tensor in chunks + chunk_min_dists = torch.full((N,), float("inf"), device=device) + + for start_idx in range(0, M, chunk_size): + end_idx = min(start_idx + chunk_size, M) + ref_chunk = ref_tensor[start_idx:end_idx] + + # Compute distances for this chunk + dist_matrix = torch.cdist(tensor_a, ref_chunk, p=2.0) # [N, chunk_size] + chunk_dists, _ = dist_matrix.min(dim=1) # [N] + + # Update chunk minimums + chunk_min_dists = torch.minimum(chunk_min_dists, chunk_dists) + + # Update global minimums + update_mask = chunk_min_dists < global_min_dists + global_min_dists[update_mask] = chunk_min_dists[update_mask] + global_tensor_idx[update_mask] = tensor_idx + + return global_min_dists, global_tensor_idx + + +def compute_nearest_distance_to_list_cuda_optimized( + tensor_a: torch.Tensor, tensor_list: List[torch.Tensor], max_memory_mb: float = 1000 +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Adaptive CUDA implementation that chooses strategy based on memory constraints. + + Args: + tensor_a: Query tensor on CUDA, shape [N, D] + tensor_list: List of reference tensors on CUDA + max_memory_mb: Maximum memory to use in MB + + Returns: + min_distances: shape [N] + min_tensor_indices: shape [N] + """ + if not tensor_a.is_cuda: + raise RuntimeError( + "All tensors must be on CUDA; CPU execution is not supported." + ) + + N, D = tensor_a.shape + + # Estimate memory requirement + total_ref_points = sum(t.shape[0] for t in tensor_list) + memory_needed_mb = (N * total_ref_points * 4) / (1024 * 1024) # float32 = 4 bytes + + if memory_needed_mb <= max_memory_mb: + # Use fast concatenated version + return compute_nearest_distance_to_list_cuda( + tensor_a, tensor_list, batch_mode=False + ) + else: + # Use memory-efficient chunked version + chunk_size = int(max_memory_mb * 1024 * 1024 / (N * 4)) + return compute_nearest_distance_to_list_cuda_chunked( + tensor_a, tensor_list, chunk_size + ) + + +# Self-tests +if __name__ == "__main__": + print("=" * 60) + print("Testing CUDA-accelerated nearest distance computation") + print("=" * 60) + + # Create test data ON CUDA + tensor_a = torch.rand(1000, 3).cuda() # 1000 query points + tensor_list = [ + torch.rand(500, 3).cuda(), # 500 reference points + torch.rand(800, 3).cuda(), # 800 reference points + torch.rand(1200, 3).cuda(), # 1200 reference points + ] + + print("\nTest 1: Basic CUDA version") + try: + min_dists, min_indices = compute_nearest_distance_to_list_cuda( + tensor_a, tensor_list, batch_mode=False + ) + print(f"Success. Min distances shape: {min_dists.shape}") + print(f" First 5 distances: {min_dists[:5]}") + print(f" Tensor indices: {min_indices[:5]}") + except Exception as e: + print(f"Error: {e}") + + print("\nTest 1b: With vector output") + try: + min_dists, min_indices, vectors = compute_nearest_distance_to_list_cuda( + tensor_a, tensor_list, batch_mode=False, return_vectors=True + ) + print(f"Success. Vectors shape: {vectors.shape}") + print(f" First vector: {vectors[0]}") + # Verify vector length matches distance + vec_norm = vectors[0].norm() + print(f" Vector norm: {vec_norm:.6f}, Distance: {min_dists[0]:.6f}") + print(f" Match: {torch.allclose(vec_norm, min_dists[0], rtol=1e-5)}") + except Exception as e: + print(f"Error: {e}") + + print("\nTest 2: Batch mode") + tensor_a_batch = torch.rand(4, 100, 3).cuda() # Batch of 4 + tensor_list_batch = [ + torch.rand(4, 50, 3).cuda(), + torch.rand(4, 80, 3).cuda(), + ] + + try: + min_dists_batch, min_indices_batch = compute_nearest_distance_to_list_cuda( + tensor_a_batch, tensor_list_batch, batch_mode=True + ) + print(f"Success. Batch distances shape: {min_dists_batch.shape}") + except Exception as e: + print(f"Error: {e}") + + print("\nTest 3: Chunked version for large tensors") + try: + min_dists_chunked, min_indices_chunked = ( + compute_nearest_distance_to_list_cuda_chunked( + tensor_a, tensor_list, chunk_size=100 + ) + ) + print(f"Success. Chunked distances shape: {min_dists_chunked.shape}") + + # Verify chunked gives same result as full + diff = (min_dists - min_dists_chunked).abs().max() + print(f" Difference from full version: {diff:.6f}") + except Exception as e: + print(f"Error: {e}") + + print("\nTest 4: Testing with CPU tensor (should fail)") + tensor_cpu = torch.rand(10, 3) # CPU tensor + try: + min_dists, min_indices = compute_nearest_distance_to_list_cuda( + tensor_cpu, tensor_list, batch_mode=False + ) + print("FAILED: expected an error for the CPU tensor") + except RuntimeError as e: + print(f"Correctly rejected CPU tensor: {e}") + + print("\nTest 5: Performance comparison") + import time + + # Large test case + large_query = torch.rand(5000, 3).cuda() + large_refs = [torch.rand(2000, 3).cuda() for _ in range(5)] + + # Time the optimized version + start = time.time() + result = compute_nearest_distance_to_list_cuda_optimized(large_query, large_refs) + cuda_time = time.time() - start + + print(f"Processed 5000 queries against 10000 references in {cuda_time:.3f}s") + print( + f" Throughput: {5000 * 10000 / cuda_time / 1e6:.1f} million distance computations/second" + ) + + print("\n" + "=" * 60) + print("Test 6: Vectorized multi-object KNN (correctness verification)") + print("=" * 60) + + # Create test data: B batches, Nobj objects, Nv vertices each + B, Nobj, Nv = 4, 5, 100 + FAR = 1e6 + + positions = torch.rand(B, Nobj, Nv, 3).cuda() + # Create a valid mask (80% of vertices are valid) + valid_mask = torch.rand(B, Nobj, Nv).cuda() > 0.2 + + print(f" Test setup: B={B}, Nobj={Nobj}, Nv={Nv}") + print(f" Valid vertices per object: {valid_mask.float().mean(dim=(0,2)) * Nv}") + + # Method 1: loop-based reference implementation + print("\n Method 1: Loop-based KNN...") + dist_vectors_loop = torch.zeros(B, Nobj, Nv, 3, device=positions.device) + + start = time.time() + for i in range(Nobj): + body_pos = positions[:, i] # (B, Nv, 3) + ref_list = [] + for j in range(Nobj): + if j == i: + continue + other_pos = positions[:, j].clone() + other_vmask = valid_mask[:, j].unsqueeze(-1) # (B, Nv, 1) + other_pos = torch.where(other_vmask, other_pos, other_pos + FAR) + ref_list.append(other_pos) + + _, _, body_dist_vec = compute_nearest_distance_to_list_cuda( + body_pos, ref_list, batch_mode=True, return_vectors=True + ) + dist_vectors_loop[:, i] = body_dist_vec + loop_time = time.time() - start + print(f" Loop-based completed in {loop_time*1000:.2f}ms") + + # Method 2: Vectorized approach + print("\n Method 2: Vectorized KNN...") + start = time.time() + dist_vectors_vec = compute_knn_all_objects_vectorized(positions, valid_mask, FAR) + vec_time = time.time() - start + print(f" Vectorized completed in {vec_time*1000:.2f}ms") + + # Compare results + print("\n Comparing results...") + + # Only compare valid query vertices (invalid queries may have different results) + valid_query_mask = valid_mask.unsqueeze(-1).expand_as(dist_vectors_loop) + + diff = (dist_vectors_loop - dist_vectors_vec).abs() + diff_valid = diff[valid_query_mask] + + max_diff = diff_valid.max().item() if diff_valid.numel() > 0 else 0 + mean_diff = diff_valid.mean().item() if diff_valid.numel() > 0 else 0 + + print(f" Max difference (valid vertices): {max_diff:.6e}") + print(f" Mean difference (valid vertices): {mean_diff:.6e}") + + if max_diff < 1e-5: + print(f" PASSED: Results match within tolerance (1e-5)") + else: + print(f" FAILED: Results differ by more than tolerance") + # Debug: find where they differ + diff_mask = diff > 1e-5 + num_diff = diff_mask.sum().item() + print(f" Number of differing elements: {num_diff}") + + # Speedup + speedup = loop_time / vec_time if vec_time > 0 else float("inf") + print(f"\n Speedup: {speedup:.2f}x") + + # Test with larger data + print("\n Testing with larger data (B=8, Nobj=10, Nv=500)...") + B2, Nobj2, Nv2 = 8, 10, 500 + positions2 = torch.rand(B2, Nobj2, Nv2, 3).cuda() + valid_mask2 = torch.rand(B2, Nobj2, Nv2).cuda() > 0.1 + + # Warmup + _ = compute_knn_all_objects_vectorized(positions2, valid_mask2, FAR) + torch.cuda.synchronize() + + start = time.time() + _ = compute_knn_all_objects_vectorized(positions2, valid_mask2, FAR) + torch.cuda.synchronize() + vec_time2 = time.time() - start + print(f" Vectorized (large): {vec_time2*1000:.2f}ms") + print(f" Throughput: {B2 * Nobj2 * Nv2 / vec_time2 / 1e6:.2f}M vertices/sec") + + print("\n" + "=" * 60) + print("Test 7: Invalid vertex/object masking verification") + print("=" * 60) + + B, Nobj, Nv = 2, 4, 50 + positions = torch.rand(B, Nobj, Nv, 3).cuda() + + # Create mask with some fully invalid objects and some invalid vertices + valid_vtx_mask = torch.ones(B, Nobj, Nv, dtype=torch.bool, device="cuda") + valid_obj_mask = torch.ones(B, Nobj, dtype=torch.bool, device="cuda") + + # Make object 2 fully invalid (no valid vertices) + valid_vtx_mask[:, 2, :] = False + valid_obj_mask[:, 2] = False + + # Make some vertices invalid in valid objects + valid_vtx_mask[0, 0, 10:20] = False # batch 0, object 0, vertices 10-19 invalid + valid_vtx_mask[1, 1, 30:40] = False # batch 1, object 1, vertices 30-39 invalid + + print(f" Setup: B={B}, Nobj={Nobj}, Nv={Nv}") + print(f" Object 2 is fully invalid (all vertices masked)") + print(f" Some vertices in valid objects are also masked") + + dist_vectors = compute_knn_all_objects_vectorized( + positions, valid_vtx_mask, FAR, valid_obj_mask + ) + + # Check 1: Invalid object (object 2) should have ALL zero distance vectors + obj2_dist = dist_vectors[:, 2, :, :] # (B, Nv, 3) + obj2_is_zero = (obj2_dist.abs() < 1e-10).all() + print( + f"\n Check 1: Invalid object (obj 2) has zero dist_vec: {obj2_is_zero.item()}" + ) + if obj2_is_zero: + print(" PASSED: Invalid object has all zeros") + else: + print( + f" FAILED: Non-zero values found: max={obj2_dist.abs().max().item():.6e}" + ) + + # Check 2: Invalid vertices in valid objects should have zero distance vectors + invalid_vtx_dist_0 = dist_vectors[0, 0, 10:20, :] # batch 0, obj 0, vtx 10-19 + invalid_vtx_dist_1 = dist_vectors[1, 1, 30:40, :] # batch 1, obj 1, vtx 30-39 + + vtx0_is_zero = (invalid_vtx_dist_0.abs() < 1e-10).all() + vtx1_is_zero = (invalid_vtx_dist_1.abs() < 1e-10).all() + + print(f"\n Check 2: Invalid vertices have zero dist_vec:") + print(f" batch 0, obj 0, vtx 10-19: {vtx0_is_zero.item()}") + print(f" batch 1, obj 1, vtx 30-39: {vtx1_is_zero.item()}") + if vtx0_is_zero and vtx1_is_zero: + print(" PASSED: Invalid vertices have all zeros") + else: + print(f" FAILED: Non-zero values found") + + # Check 3: Valid vertices in valid objects should have NON-zero distance vectors + # (at least most of them, since they should find nearest neighbors on other objects) + valid_vtx_dist_0 = dist_vectors[0, 0, :10, :] # batch 0, obj 0, vtx 0-9 (valid) + valid_vtx_dist_1 = dist_vectors[0, 1, :30, :] # batch 0, obj 1, vtx 0-29 (valid) + + valid_has_nonzero_0 = (valid_vtx_dist_0.abs() > 1e-10).any() + valid_has_nonzero_1 = (valid_vtx_dist_1.abs() > 1e-10).any() + + print(f"\n Check 3: Valid vertices have NON-zero dist_vec (found neighbors):") + print(f" batch 0, obj 0, vtx 0-9: {valid_has_nonzero_0.item()}") + print(f" batch 0, obj 1, vtx 0-29: {valid_has_nonzero_1.item()}") + if valid_has_nonzero_0 and valid_has_nonzero_1: + print(" PASSED: Valid vertices have non-zero distances") + else: + print(" FAILED: Valid vertices should find neighbors on other objects") + + print("\n" + "=" * 60) + print("All tests completed! CUDA acceleration is working.") diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/utils/intersection_handler.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/utils/intersection_handler.py new file mode 100644 index 0000000000..a50ce8ca8e --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/utils/intersection_handler.py @@ -0,0 +1,232 @@ +# 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. + +""" +Differentiable intersection handling. + +Resolves object-ground and object-object penetrations by translating whole +objects, with a straight-through estimator (STE) keeping the correction +differentiable. The correction formulation follows UNIC (arXiv:2603.25580, +Sec. 3.3), Eq. 8: + + v_corrected = v + detach(target_position - v) + Forward: moves v to target_position + Backward: gradients pass through to v unchanged + +Whereas UNIC drags individual garment vertices, rigid objects are translated +as a whole: + - Ground plane: shift the object upward when its lowest vertex falls + below ground_z + buffer + - Object-object: push objects apart when their bounding spheres overlap +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F +import logging + +logger = logging.getLogger(__name__) + + +def resolve_ground_penetration( + positions: torch.Tensor, + obj_mask: torch.Tensor, + vtx_mask: torch.Tensor, + ground_z: float = 0.0, + buffer: float = 0.005, +) -> torch.Tensor: + """ + Resolve object-ground penetrations by translating objects upward. + + For each valid object, find the LOWEST valid vertex. If min_z < ground_z + buffer, + shift ALL vertices of that object upward by the deficit. + This guarantees no valid vertex penetrates the ground, while preserving rigidity. + + Uses STE trick: correction is detached so gradients flow through original positions. + + Args: + positions: (B, Nobj, Nv, 3) predicted vertex positions + obj_mask: (B, Nobj) boolean mask for valid objects + vtx_mask: (B, Nobj, Nv) boolean mask for valid vertices + ground_z: ground plane height (z-up coordinate system) + buffer: relaxation buffer above ground plane + + Returns: + corrected_positions: (B, Nobj, Nv, 3) with ground penetrations resolved + """ + B, Nobj, Nv, _ = positions.shape + target_z = ground_z + buffer + + # Find per-object minimum z among valid vertices + # Set invalid vertices to +inf so they don't affect min + z_vals = positions[..., 2] # (B, Nobj, Nv) + z_vals_masked = z_vals.masked_fill( + ~vtx_mask, float("inf") + ) # invalid vertices -> +inf + min_z, _ = z_vals_masked.min(dim=2) # (B, Nobj) lowest valid vertex z per object + + # Find objects that penetrate the ground + penetration_depth = target_z - min_z # positive means below ground + # Padding objects have min_z=+inf, giving depth=-inf; clamp to >= 0 before + # the expand below so 0*(-inf) can never produce NaN + penetration_depth = penetration_depth.clamp(min=0.0) + is_penetrating = (penetration_depth > 0) & obj_mask # (B, Nobj) + + # No early exit on "nothing penetrates": the correction below is identically + # zero in that case, and the shortcut would cost a host sync per call. + + # Compute correction: shift entire object upward by penetration_depth + correction = torch.zeros_like(positions) # (B, Nobj, Nv, 3) + # Only z-component needs correction + correction[..., 2] = penetration_depth.unsqueeze(-1).expand( + -1, -1, Nv + ) # (B, Nobj, Nv) + + # Zero out correction for non-penetrating objects + correction = correction * is_penetrating.float().unsqueeze(-1).unsqueeze( + -1 + ) # (B, Nobj, Nv, 3) + + # STE trick (UNIC Eq. 8): forward moves position, backward passes gradient through + corrected = positions + correction.detach() + + return corrected + + +def resolve_object_penetration( + positions: torch.Tensor, + obj_mask: torch.Tensor, + vtx_mask: torch.Tensor, + buffer: float = 0.005, +) -> torch.Tensor: + """ + Resolve object-object penetrations using bounding sphere overlap detection. + + For each pair of valid objects, compute COM and bounding sphere radius. + If spheres overlap, push both objects apart along the COM-COM direction. + + Uses STE trick for differentiability. + + Args: + positions: (B, Nobj, Nv, 3) predicted vertex positions + obj_mask: (B, Nobj) boolean mask for valid objects + vtx_mask: (B, Nobj, Nv) boolean mask for valid vertices + buffer: relaxation buffer between objects + + Returns: + corrected_positions: (B, Nobj, Nv, 3) with object penetrations resolved + """ + B, Nobj, Nv, _ = positions.shape + + if Nobj < 2: + return positions + + # Compute per-object COM and bounding sphere radius + vtx_weight = vtx_mask.float().unsqueeze(-1) # (B, Nobj, Nv, 1) + num_valid = vtx_weight.sum(dim=2).clamp(min=1) # (B, Nobj, 1) + com = (positions * vtx_weight).sum(dim=2) / num_valid # (B, Nobj, 3) + + # Bounding sphere radius: max distance from COM to any valid vertex + diff_to_com = positions - com.unsqueeze(2) # (B, Nobj, Nv, 3) + dist_to_com = diff_to_com.norm(dim=-1) # (B, Nobj, Nv) + # Mask invalid vertices with 0 distance + dist_to_com = dist_to_com * vtx_mask.float() # (B, Nobj, Nv) + radius = dist_to_com.max(dim=2).values # (B, Nobj) + + # Accumulate corrections per object + correction = torch.zeros(B, Nobj, 3, device=positions.device, dtype=positions.dtype) + + # Check all pairs (i, j) where i < j + for i in range(Nobj): + for j in range(i + 1, Nobj): + # Both objects must be valid + pair_valid = obj_mask[:, i] & obj_mask[:, j] # (B,) + if not pair_valid.any(): + continue + + # COM distance + com_diff = com[:, j] - com[:, i] # (B, 3) + com_dist = com_diff.norm(dim=-1).clamp(min=1e-8) # (B,) + + # Minimum allowed distance + min_dist = radius[:, i] + radius[:, j] + buffer # (B,) + + # Overlap amount + overlap = min_dist - com_dist # (B,) positive means penetrating + is_overlapping = (overlap > 0) & pair_valid # (B,) + + if not is_overlapping.any(): + continue + + # Push direction: from i to j (normalized) + push_dir = com_diff / com_dist.unsqueeze(-1) # (B, 3) + + # Each object gets pushed by half the overlap + push_amount = (overlap / 2).clamp(min=0) * is_overlapping.float() # (B,) + + correction[:, i] -= ( + push_amount.unsqueeze(-1) * push_dir + ) # push i away from j + correction[:, j] += ( + push_amount.unsqueeze(-1) * push_dir + ) # push j away from i + + if (correction.abs() < 1e-10).all(): + return positions + + # Expand correction to per-vertex: (B, Nobj, 3) -> (B, Nobj, Nv, 3) + correction_expanded = correction.unsqueeze(2).expand(-1, -1, Nv, -1) + + # STE trick + corrected = positions + correction_expanded.detach() + + return corrected + + +def handle_intersections( + positions: torch.Tensor, + obj_mask: torch.Tensor, + vtx_mask: torch.Tensor, + config: dict, +) -> torch.Tensor: + """ + Main entry point: apply all intersection resolution steps. + + Args: + positions: (B, Nobj, Nv, 3) predicted vertex positions + obj_mask: (B, Nobj) boolean mask + vtx_mask: (B, Nobj, Nv) boolean mask for valid vertices + config: intersection_handling config dict + + Returns: + corrected positions: (B, Nobj, Nv, 3) + """ + if not config.get("enable", False): + return positions + + ground_cfg = config.get("ground", {}) + object_cfg = config.get("object", {}) + + # Step 1: Resolve ground penetrations + if ground_cfg.get("enable", True): + positions = resolve_ground_penetration( + positions, + obj_mask, + vtx_mask, + ground_z=ground_cfg.get("z", 0.0), + buffer=ground_cfg.get("buffer", 0.005), + ) + + # Step 2: Resolve object-object penetrations + if object_cfg.get("enable", False): + positions = resolve_object_penetration( + positions, + obj_mask, + vtx_mask, + buffer=object_cfg.get("buffer", 0.005), + ) + + return positions diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/utils/weight_ema.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/utils/weight_ema.py new file mode 100644 index 0000000000..3e7951c387 --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/utils/weight_ema.py @@ -0,0 +1,92 @@ +# 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. + +"""Exponential moving average (EMA) of a model's weights. + +The average covers every floating-point entry of the module's state dict (parameters and +floating buffers such as BatchNorm running statistics); integer buffers (e.g. BatchNorm's +num_batches_tracked) are copied. The shadow copy is kept in fp32. The effective decay ramps as +min(decay, (1 + n) / (10 + n)) over the first updates so the average is not biased toward the +initial weights. +""" + +from __future__ import annotations +from typing import Dict, Optional + +import torch + + +class WeightEMA: + def __init__(self, module: torch.nn.Module, decay: float): + if not 0.0 < float(decay) < 1.0: + raise ValueError(f"EMA decay must be in (0, 1); got {decay}") + self.module = module + self.decay = float(decay) + self.num_updates = 0 + self._keys = list(module.state_dict(keep_vars=True).keys()) + self.shadow: Dict[str, torch.Tensor] = {} + self.reset() + + def _state(self) -> Dict[str, torch.Tensor]: + sd = self.module.state_dict(keep_vars=True) + if list(sd.keys()) != self._keys: + raise RuntimeError( + "WeightEMA: the module's state-dict keys changed after construction" + ) + return sd + + @torch.no_grad() + def reset(self) -> None: + """Restart the average from the module's current weights (e.g. after loading a checkpoint).""" + self.shadow = { + k: ( + v.detach().float().clone() + if v.is_floating_point() + else v.detach().clone() + ) + for k, v in self._state().items() + } + self.num_updates = 0 + + @torch.no_grad() + def update(self) -> None: + self.num_updates += 1 + d = min(self.decay, (1.0 + self.num_updates) / (10.0 + self.num_updates)) + sh, src = [], [] + for k, v in self._state().items(): + if v.is_floating_point(): + sh.append(self.shadow[k]) + src.append(v.detach().float()) + else: + self.shadow[k].copy_(v.detach()) + torch._foreach_mul_(sh, d) + torch._foreach_add_(sh, src, alpha=1.0 - d) + + def state_dict(self) -> Dict[str, torch.Tensor]: + """Averaged weights in each entry's own dtype (loadable with module.load_state_dict).""" + return { + k: ( + self.shadow[k].to(v.dtype) + if v.is_floating_point() + else self.shadow[k].clone() + ) + for k, v in self._state().items() + } + + @torch.no_grad() + def load_state_dict( + self, sd: Dict[str, torch.Tensor], num_updates: Optional[int] = None + ) -> None: + missing = [k for k in self._keys if k not in sd] + if missing: + raise KeyError( + f"WeightEMA.load_state_dict: {len(missing)} keys missing, e.g. {missing[:3]}" + ) + for k in self._keys: + self.shadow[k].copy_( + sd[k].to(device=self.shadow[k].device, dtype=self.shadow[k].dtype) + ) + if num_updates is not None: + self.num_updates = int(num_updates) diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/checkpoint.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/checkpoint.py new file mode 100644 index 0000000000..8f5fe6ea22 --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/checkpoint.py @@ -0,0 +1,78 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Strict import of official RigidFormer weights without optimizer conversion.""" + +from __future__ import annotations +from collections.abc import Mapping +from pathlib import Path +import torch +from .model import RigidFormerCfg, RigidFormerModel + +__all__ = ["load_rigidformer_checkpoint"] + + +def load_rigidformer_checkpoint( + path: str | Path, + *, + device: str | torch.device = "cpu", + attention_backend: str = "sdpa", +) -> RigidFormerModel: + """Load an official model/EMA weights file and its architecture. + + Args: + path: Official tensor/config checkpoint; decoded with weights_only=True. + device: Destination device, independent of the file's saved device. + attention_backend: Explicit native attention implementation. + + Returns: + Evaluation model with every author weight loaded strictly. + + This is a weights import, not recovery of an official training process. + """ + payload = torch.load(path, map_location="cpu", weights_only=True, mmap=True) + if ( + not isinstance(payload, Mapping) + or "model_state_dict" not in payload + or "config" not in payload + ): + raise ValueError( + "Expected official model_state_dict and config; community/SL checkpoints are distinct formats" + ) + cfg = RigidFormerCfg( + architecture=dict(payload["config"]), attention_backend=attention_backend + ) + model = RigidFormerModel(cfg) + incoming = payload["model_state_dict"] + if not isinstance(incoming, Mapping) or not all( + isinstance(k, str) and isinstance(v, torch.Tensor) for k, v in incoming.items() + ): + raise ValueError("Official weights must be a named tensor mapping") + prefix = all(key.startswith("module.") for key in incoming) + state = { + key.removeprefix("module.") if prefix else key: value + for key, value in incoming.items() + } + expected = model.network.state_dict() + if state.keys() != expected.keys(): + raise ValueError( + "Official checkpoint keys differ from the configured native architecture" + ) + for name, value in expected.items(): + if state[name].shape != value.shape or state[name].dtype != value.dtype: + raise ValueError(f"Official checkpoint tensor differs: {name}") + model.network.load_state_dict(state, strict=True) + return model.to(device).eval() diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/configs/movi_a.yaml b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/configs/movi_a.yaml new file mode 100644 index 0000000000..b2af2734b2 --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/configs/movi_a.yaml @@ -0,0 +1,284 @@ +# RigidFormer on MOVi-A. Data: datasets/MOVi-A (data_generation/README.md), scene lists +# dataset_split_MoviA_1200.json (960 / 120 / 120), object sizes data_generation/object_scales_MoviA.json. +# Training: 300 epochs on 4 GPUs x 32 scenes. +logging: + use_tensorboard: true + file_logging: true +experiment_tag: rigidformer_movia +paths: + log_dir: ./exps/logs_multi + checkpoint_dir: ./exps/checkpoints_multi +data: + dataset_path: ./datasets/MOVi-A + split_file: dataset_split_MoviA_1200.json + scene_set: complete + multi_object: true + anchor_k: 4 + num_workers: 4 + sequence_length: 8 + train_split: 0.99 + val_split: 0.01 + test_split: 0.0 + skip_first_frames: 2 + load_gt_motion: false + random_step_size: + enable: true + step_sizes: + - 1 + - 5 + - 10 + random_start_frame: true + sampling_weights: + - 0.34 + - 0.33 + - 0.33 + eval_step_sizes: + - 10 + - 5 + - 1 + max_vertices_per_object: null + max_objects_per_frame: 16 + augmentation: + random_walk: + enable: false + start_epoch: 5 + noise_std: 0.0 + independent: + enable: false + velocity_noise: 0.0 + dist_vec_noise: 0.0 + rel_pos_noise: 0.0 + rotation_augmentation: + enable: true + angle_step_size: 5 + prob_use_rotation: 0.5 + object_permutation: + enable: true + prob_apply: 0.5 + if_debug: false + anchor_resampling: + enable: false + prob: 0.5 + anchor_fps_start: centroid_farthest + fps_order: true + object_scale_file: data_generation/object_scales_MoviA.json +model: + if KV-caching: false + no_global_state: true + if_penetration: false + if_directional_distance: false + if_phys_param: true + if_time_step: true + film_conditioning: true + avp: + enable: true + output_dim: 256 + init_sigma: 0.1 + encoder: + input_dim: 12 + hidden_dims: + - 1024 + - 1024 + output_dim: 768 + backbone_dim: 1024 + fusion_hidden_dim: 1024 + use_batch_norm: true + use_stn: false + feature_transform: false + feature_transform_dim: 64 + pooling_type: hierarchical + num_hierarchical_levels: 3 + masked_pooling: true + hierarchical_subsets: true + packed: true + packed_pooling: segment + subset_small_objects: fraction + freeze_bn_stats: false + state: + token_dim: 768 + decoder: + num_layers: 4 + num_heads: 6 + hidden_dim: 768 + state_update_enabled: false + rope_set_aggregation: true + rope_dim: 96 + dim_per_coord_3d: 32 + dim_per_coord_multi: 8 + feedforward_dim: null + dropout: 0.1 + return_all_layers: true + activation: swiglu + use_self_attention: false + norm_type: rms_norm + qk_norm: true + num_register_tokens: 16 + cross_attention: + enable: true + num_heads: 6 + gated_attention: + enable: true + type: elementwise + if_debug: false + objects_self_attention: true + film_conditioning: true + block_attn_res: + enable: true + block_size: 4 + final_aggregate: true + predictor: + rmsnorm_eps: training_dtype + fp32_output: true + input_dim: 271 + hidden_dims: + - 768 + - 256 + output_dim: 3 + dropout: 0.1 + use_ref_pose: true + norm_type: rms_norm + qk_norm: true + cross_object_attention: true + cross_object_kv: true + transform: + method: kabsch + if_roma: true + integrator: + type: improved_verlet + use_damping: false + damping_factor: -100 + intersection_handling: + enable: true + after_kabsch_loss: rigid + ground: + enable: true + z: 0.0 + buffer: -0.01 + object: + enable: false + buffer: 0.005 + exp1: + output_type: raw_acceleration + ref_pose_object_centered: true + rope_interleaved_pairing: true + rope_center_positions: true + kabsch_reference: fixed + knn_packed: true +training: + enable_validation: false + max_epochs: 300 + scheduled_sampling: + enabled: false + mode: three_stage + teacher_epochs: 9999 + ramping_epochs: 0 + optimizer: + type: AdamW + lr: 1e-4 + fusion_lr_mult: 10.0 + betas: + - 0.9 + - 0.999 + weight_decay: 0.01 + decay_matrices_only: true + gradient_clip: 1.0 + scheduler: + type: cosine + warmup_epochs: 10 + start_factor: 0.1 + eta_min: 1e-6 + loss_weights: + if_time_step_scale: true + global_anchor_mean: true + if_debug: false + position_before_icp: 10 + position_after_icp: 10 + acceleration_after_icp: 1 + acceleration_before_icp: 1 + penetration_sign: 0.0 + rt_rot: 0.0 + rt_trans: 0.0 + velocity: 0.0 + collision: 0.0 + physics: 0.0 + edge_length: 0.0 + checkpoint: + best_metric: + - p_rmse_@200/s=1 + - p_rmse_@300/s=1 + - p_rmse_@400/s=1 + save_every_n_epochs: 5 + save_video: false + video_every_n_epochs: 1 + original_video_num_frames: 480 + original_video_fps: 240 + num_videos_per_split: + train: 2 + val: 0 + test: 2 + if_overlay: false + gt_transparency: 0.1 + arrow_enable: true + arrow_scale: 0.01 + arrow_subsample: 1 + arrow_thickness: 0.0007 + point_size: 0.0001 + camera: + mode: fixed + center: + - 0.0 + - 0.0 + - 1.0 + eye: + - 0.0 + - -15.0 + - 10.0 + auto_direction: + - 2.0 + - -2.0 + - 2.0 + auto_scale: 3.0 + ground: + size: 10.0 + color: + - 0.8 + - 0.8 + - 0.8 + metrics_test: + enable: true + initial_eval: false + print_every_n_epochs: 1 + eval_num_frames: 480 + eval_warmup_frames: 4 + hopnet_checkpoints: + - 25 + - 50 + - 75 + - 100 + - 200 + - 300 + - 400 + position_rmse: true + position_final_frame: true + position_per_timestep: true + rotation_error_deg: true + rotation_final_frame: true + rotation_per_timestep: true + position_estimator: hopnet + train_batch_size: 32 + val_batch_size: 18 + eval_batch_size: 18 + ema: + enable: true + decay: 0.9999 + save_every_n_epochs: 5 + rank_seeded_rng: true + time_batched_forward: true + compile: + enable: true + modules: + - decoder + - predictor + mode: default + dynamic: false + diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/configs/movi_b.yaml b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/configs/movi_b.yaml new file mode 100644 index 0000000000..10a281e3c2 --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/configs/movi_b.yaml @@ -0,0 +1,284 @@ +# RigidFormer on MOVi-B. Data: datasets/MOVi-B (data_generation/README.md), scene lists +# dataset_split_MoviB_1200.json (960 / 120 / 120), object sizes data_generation/object_scales_MoviB.json. +# Training: 300 epochs on 6 GPUs x 21 scenes. +logging: + use_tensorboard: true + file_logging: true +experiment_tag: rigidformer_movib +paths: + log_dir: ./exps/logs_multi + checkpoint_dir: ./exps/checkpoints_multi +data: + dataset_path: ./datasets/MOVi-B + split_file: dataset_split_MoviB_1200.json + scene_set: complete + multi_object: true + anchor_k: 4 + num_workers: 4 + sequence_length: 8 + train_split: 0.99 + val_split: 0.01 + test_split: 0.0 + skip_first_frames: 2 + load_gt_motion: false + random_step_size: + enable: true + step_sizes: + - 1 + - 5 + - 10 + random_start_frame: true + sampling_weights: + - 0.34 + - 0.33 + - 0.33 + eval_step_sizes: + - 10 + - 5 + - 1 + max_vertices_per_object: null + max_objects_per_frame: 16 + augmentation: + random_walk: + enable: false + start_epoch: 5 + noise_std: 0.0 + independent: + enable: false + velocity_noise: 0.0 + dist_vec_noise: 0.0 + rel_pos_noise: 0.0 + rotation_augmentation: + enable: true + angle_step_size: 5 + prob_use_rotation: 0.5 + object_permutation: + enable: true + prob_apply: 0.5 + if_debug: false + anchor_resampling: + enable: false + prob: 0.5 + anchor_fps_start: centroid_farthest + fps_order: true + object_scale_file: data_generation/object_scales_MoviB.json +model: + if KV-caching: false + no_global_state: true + if_penetration: false + if_directional_distance: false + if_phys_param: true + if_time_step: true + film_conditioning: true + avp: + enable: true + output_dim: 256 + init_sigma: 0.1 + encoder: + input_dim: 12 + hidden_dims: + - 1024 + - 1024 + output_dim: 768 + backbone_dim: 1024 + fusion_hidden_dim: 1024 + use_batch_norm: true + use_stn: false + feature_transform: false + feature_transform_dim: 64 + pooling_type: hierarchical + num_hierarchical_levels: 3 + masked_pooling: true + hierarchical_subsets: true + packed: true + packed_pooling: segment + subset_small_objects: fraction + freeze_bn_stats: false + state: + token_dim: 768 + decoder: + num_layers: 4 + num_heads: 6 + hidden_dim: 768 + state_update_enabled: false + rope_set_aggregation: true + rope_dim: 96 + dim_per_coord_3d: 32 + dim_per_coord_multi: 8 + feedforward_dim: null + dropout: 0.1 + return_all_layers: true + activation: swiglu + use_self_attention: false + norm_type: rms_norm + qk_norm: true + num_register_tokens: 16 + cross_attention: + enable: true + num_heads: 6 + gated_attention: + enable: true + type: elementwise + if_debug: false + objects_self_attention: true + film_conditioning: true + block_attn_res: + enable: true + block_size: 4 + final_aggregate: true + predictor: + rmsnorm_eps: training_dtype + fp32_output: true + input_dim: 271 + hidden_dims: + - 768 + - 256 + output_dim: 3 + dropout: 0.1 + use_ref_pose: true + norm_type: rms_norm + qk_norm: true + cross_object_attention: true + cross_object_kv: true + transform: + method: kabsch + if_roma: true + integrator: + type: improved_verlet + use_damping: false + damping_factor: -100 + intersection_handling: + enable: true + after_kabsch_loss: rigid + ground: + enable: true + z: 0.0 + buffer: -0.01 + object: + enable: false + buffer: 0.005 + exp1: + output_type: raw_acceleration + ref_pose_object_centered: true + knn_packed: true + rope_interleaved_pairing: true + rope_center_positions: true + kabsch_reference: fixed +training: + enable_validation: false + max_epochs: 300 + scheduled_sampling: + enabled: false + mode: three_stage + teacher_epochs: 9999 + ramping_epochs: 0 + optimizer: + type: AdamW + lr: 1e-4 + fusion_lr_mult: 10.0 + betas: + - 0.9 + - 0.999 + weight_decay: 0.01 + decay_matrices_only: true + gradient_clip: 1.0 + scheduler: + type: cosine + warmup_epochs: 10 + start_factor: 0.1 + eta_min: 1e-6 + loss_weights: + if_time_step_scale: true + global_anchor_mean: true + if_debug: false + position_before_icp: 10 + position_after_icp: 10 + acceleration_after_icp: 1 + acceleration_before_icp: 1 + penetration_sign: 0.0 + rt_rot: 0.0 + rt_trans: 0.0 + velocity: 0.0 + collision: 0.0 + physics: 0.0 + edge_length: 0.0 + checkpoint: + best_metric: + - p_rmse_@200/s=1 + - p_rmse_@300/s=1 + - p_rmse_@400/s=1 + save_every_n_epochs: 5 + save_video: false + video_every_n_epochs: 1 + original_video_num_frames: 480 + original_video_fps: 240 + num_videos_per_split: + train: 2 + val: 0 + test: 2 + if_overlay: false + gt_transparency: 0.1 + arrow_enable: true + arrow_scale: 0.01 + arrow_subsample: 1 + arrow_thickness: 0.0007 + point_size: 0.0001 + camera: + mode: fixed + center: + - 0.0 + - 0.0 + - 1.0 + eye: + - 0.0 + - -15.0 + - 10.0 + auto_direction: + - 2.0 + - -2.0 + - 2.0 + auto_scale: 3.0 + ground: + size: 10.0 + color: + - 0.8 + - 0.8 + - 0.8 + metrics_test: + enable: true + initial_eval: false + print_every_n_epochs: 1 + eval_num_frames: 480 + eval_warmup_frames: 4 + hopnet_checkpoints: + - 25 + - 50 + - 75 + - 100 + - 200 + - 300 + - 400 + position_rmse: true + position_final_frame: true + position_per_timestep: true + rotation_error_deg: true + rotation_final_frame: true + rotation_per_timestep: true + position_estimator: hopnet + train_batch_size: 21 + val_batch_size: 18 + eval_batch_size: 18 + ema: + enable: true + decay: 0.9999 + save_every_n_epochs: 5 + rank_seeded_rng: true + time_batched_forward: false + compile: + enable: true + modules: + - decoder + - predictor + mode: default + dynamic: false + diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/configs/movi_sphere.yaml b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/configs/movi_sphere.yaml new file mode 100644 index 0000000000..f028b136c3 --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/configs/movi_sphere.yaml @@ -0,0 +1,284 @@ +# RigidFormer on MOVi-Sphere. Data: datasets/MOVi-Sphere (data_generation/README.md), scene lists +# dataset_split_MoviSphere_1200.json (960 / 120 / 120), object sizes data_generation/object_scales_MoviSphere.json. +# Training: 300 epochs on 4 GPUs x 32 scenes. +logging: + use_tensorboard: true + file_logging: true +experiment_tag: rigidformer_movisphere +paths: + log_dir: ./exps/logs_multi + checkpoint_dir: ./exps/checkpoints_multi +data: + dataset_path: ./datasets/MOVi-Sphere + split_file: dataset_split_MoviSphere_1200.json + scene_set: complete + multi_object: true + anchor_k: 4 + num_workers: 4 + sequence_length: 8 + train_split: 0.99 + val_split: 0.01 + test_split: 0.0 + skip_first_frames: 2 + load_gt_motion: false + random_step_size: + enable: true + step_sizes: + - 1 + - 5 + - 10 + random_start_frame: true + sampling_weights: + - 0.34 + - 0.33 + - 0.33 + eval_step_sizes: + - 10 + - 5 + - 1 + max_vertices_per_object: null + max_objects_per_frame: 16 + augmentation: + random_walk: + enable: false + start_epoch: 5 + noise_std: 0.0 + independent: + enable: false + velocity_noise: 0.0 + dist_vec_noise: 0.0 + rel_pos_noise: 0.0 + rotation_augmentation: + enable: true + angle_step_size: 5 + prob_use_rotation: 0.5 + object_permutation: + enable: true + prob_apply: 0.5 + if_debug: false + anchor_resampling: + enable: false + prob: 0.5 + anchor_fps_start: centroid_farthest + fps_order: true + object_scale_file: data_generation/object_scales_MoviSphere.json +model: + if KV-caching: false + no_global_state: true + if_penetration: false + if_directional_distance: false + if_phys_param: true + if_time_step: true + film_conditioning: true + avp: + enable: true + output_dim: 256 + init_sigma: 0.1 + encoder: + input_dim: 12 + hidden_dims: + - 1024 + - 1024 + output_dim: 768 + backbone_dim: 1024 + fusion_hidden_dim: 1024 + use_batch_norm: true + use_stn: false + feature_transform: false + feature_transform_dim: 64 + pooling_type: hierarchical + num_hierarchical_levels: 3 + masked_pooling: true + hierarchical_subsets: true + packed: true + packed_pooling: segment + subset_small_objects: fraction + freeze_bn_stats: false + state: + token_dim: 768 + decoder: + num_layers: 4 + num_heads: 6 + hidden_dim: 768 + state_update_enabled: false + rope_set_aggregation: true + rope_dim: 96 + dim_per_coord_3d: 32 + dim_per_coord_multi: 8 + feedforward_dim: null + dropout: 0.1 + return_all_layers: true + activation: swiglu + use_self_attention: false + norm_type: rms_norm + qk_norm: true + num_register_tokens: 16 + cross_attention: + enable: true + num_heads: 6 + gated_attention: + enable: true + type: elementwise + if_debug: false + objects_self_attention: true + film_conditioning: true + block_attn_res: + enable: true + block_size: 4 + final_aggregate: true + predictor: + rmsnorm_eps: training_dtype + fp32_output: true + input_dim: 271 + hidden_dims: + - 768 + - 256 + output_dim: 3 + dropout: 0.1 + use_ref_pose: true + norm_type: rms_norm + qk_norm: true + cross_object_attention: true + cross_object_kv: true + transform: + method: kabsch + if_roma: true + integrator: + type: improved_verlet + use_damping: false + damping_factor: -100 + intersection_handling: + enable: true + after_kabsch_loss: rigid + ground: + enable: true + z: 0.0 + buffer: -0.01 + object: + enable: false + buffer: 0.005 + exp1: + output_type: raw_acceleration + ref_pose_object_centered: true + rope_interleaved_pairing: true + rope_center_positions: true + kabsch_reference: fixed + knn_packed: true +training: + enable_validation: false + max_epochs: 300 + scheduled_sampling: + enabled: false + mode: three_stage + teacher_epochs: 9999 + ramping_epochs: 0 + optimizer: + type: AdamW + lr: 1e-4 + fusion_lr_mult: 10.0 + betas: + - 0.9 + - 0.999 + weight_decay: 0.01 + decay_matrices_only: true + gradient_clip: 1.0 + scheduler: + type: cosine + warmup_epochs: 10 + start_factor: 0.1 + eta_min: 1e-6 + loss_weights: + if_time_step_scale: true + global_anchor_mean: true + if_debug: false + position_before_icp: 10 + position_after_icp: 10 + acceleration_after_icp: 1 + acceleration_before_icp: 1 + penetration_sign: 0.0 + rt_rot: 0.0 + rt_trans: 0.0 + velocity: 0.0 + collision: 0.0 + physics: 0.0 + edge_length: 0.0 + checkpoint: + best_metric: + - p_rmse_@200/s=1 + - p_rmse_@300/s=1 + - p_rmse_@400/s=1 + save_every_n_epochs: 5 + save_video: false + video_every_n_epochs: 1 + original_video_num_frames: 480 + original_video_fps: 240 + num_videos_per_split: + train: 2 + val: 0 + test: 2 + if_overlay: false + gt_transparency: 0.1 + arrow_enable: true + arrow_scale: 0.01 + arrow_subsample: 1 + arrow_thickness: 0.0007 + point_size: 0.0001 + camera: + mode: fixed + center: + - 0.0 + - 0.0 + - 1.0 + eye: + - 0.0 + - -15.0 + - 10.0 + auto_direction: + - 2.0 + - -2.0 + - 2.0 + auto_scale: 3.0 + ground: + size: 10.0 + color: + - 0.8 + - 0.8 + - 0.8 + metrics_test: + enable: true + initial_eval: false + print_every_n_epochs: 1 + eval_num_frames: 480 + eval_warmup_frames: 4 + hopnet_checkpoints: + - 25 + - 50 + - 75 + - 100 + - 200 + - 300 + - 400 + position_rmse: true + position_final_frame: true + position_per_timestep: true + rotation_error_deg: true + rotation_final_frame: true + rotation_per_timestep: true + position_estimator: hopnet + train_batch_size: 32 + val_batch_size: 18 + eval_batch_size: 18 + ema: + enable: true + decay: 0.9999 + save_every_n_epochs: 5 + rank_seeded_rng: true + time_batched_forward: true + compile: + enable: true + modules: + - decoder + - predictor + mode: default + dynamic: false + 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..59c4baaa1e --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/model.py @@ -0,0 +1,254 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Native author-based RigidFormer and explicit geometry/time contracts.""" + +from __future__ import annotations +from collections.abc import Mapping +from copy import deepcopy +from importlib.resources import files +from typing import Any +import torch +from torch import nn +from embodichain.utils import configclass + +__all__ = ["RIGIDFORMER_REVISION", "RigidFormerCfg", "RigidFormerModel"] +RIGIDFORMER_REVISION = "160fe0852c51cffae7b7440452b58d0b49db6df6" + + +@configclass +class RigidFormerCfg: + """Configure a native model using an author architecture. + + Args: + preset: Author MOVi architecture/configuration to load. + architecture: Optional full author configuration. Checkpoint imports + supply their saved configuration rather than guessing a preset. + attention_backend: Explicit sdpa or optional flash backend. + max_batch_bytes: Bound on geometry and context tensors before forwarding. + + Physical intervals are step_size / frame_rate seconds. The author model + integrates in its learned unit system with dt=1; these are distinct units. + """ + + preset: str = "movi_b" + architecture: dict[str, Any] | None = None + attention_backend: str = "sdpa" + max_batch_bytes: int = 512 * 1024 * 1024 + + def reference_config(self) -> dict[str, Any]: + """Return an owned author configuration without opening data assets.""" + self.validate() + if self.preset not in ("movi_b", "movi_a", "movi_sphere"): + raise ValueError("preset must be movi_b, movi_a, or movi_sphere") + if self.attention_backend not in ("sdpa", "flash"): + raise ValueError("attention_backend must be sdpa or flash") + if type(self.max_batch_bytes) is not int or self.max_batch_bytes < 1: + raise ValueError("max_batch_bytes must be a positive integer") + if self.architecture is not None: + result = deepcopy(self.architecture) + else: + import yaml + + result = yaml.safe_load( + files(__package__) + .joinpath("configs", self.preset + ".yaml") + .read_text() + ) + if not result.get("model", {}).get("no_global_state", False): + raise ValueError( + "The supported author MOVi family requires no_global_state=true" + ) + if result["model"].get("kabsch_reference") != "fixed": + raise ValueError( + "The supported author MOVi family requires fixed Kabsch reference geometry" + ) + return result + + +class RigidFormerModel(nn.Module): + """Predict rigid geometry with the author architecture and weight layout. + + Args: + cfg: Native architecture, optional attention backend and geometry budget. + + Construction loads optional numerical dependencies, but no data, weights, + examples or simulation. Inputs follow the author's batch layout. Training + uses shifted conditioning frames; free rollout takes only warmup frames. + """ + + action_conditioned: bool = False + qualified_state_gradients: bool = False + + def __init__(self, cfg: RigidFormerCfg | None = None) -> None: + super().__init__() + self.cfg = cfg or RigidFormerCfg() + self.reference_config = self.cfg.reference_config() + try: + import roma + from ._native.network import CUT3RDynamicsModel + except ImportError as exc: + raise ImportError( + "Native RigidFormer requires embodichain[rigidformer] (einops and roma); the external rigidformer package is not used." + ) from exc + self.network = CUT3RDynamicsModel(deepcopy(self.reference_config)) + for module in self.network.modules(): + module.attention_backend = self.cfg.attention_backend + + def _check(self, batch: Mapping[str, Any], *, min_frames: int) -> None: + positions = batch.get("obj_positions") + if ( + not isinstance(positions, torch.Tensor) + or positions.ndim != 5 + or positions.shape[-1] != 3 + or min(positions.shape[:4]) < 1 + ): + raise ValueError("obj_positions must be floating geometry [B,T,O,P,3]") + if not positions.is_floating_point() or positions.shape[1] < min_frames: + raise ValueError(f"Geometry requires at least {min_frames} frames") + b, t, o, p, _ = positions.shape + for key, shape in ( + ("valid_obj_mask", (b, o)), + ("valid_vtx_mask", (b, t, o, p)), + ): + value = batch.get(key) + if ( + not isinstance(value, torch.Tensor) + or value.shape != shape + or value.dtype != torch.bool + or value.device != positions.device + ): + raise ValueError( + f"{key} must be a bool tensor {shape} on the geometry device" + ) + valid = batch["valid_obj_mask"] + if not (valid.sum(dim=-1) >= 2).all(): + raise ValueError( + "The author nearest-body geometry requires at least two valid objects per scene" + ) + anchor = batch.get("anchor_idx") + count = self.reference_config["data"]["anchor_k"] + if ( + not isinstance(anchor, torch.Tensor) + or anchor.shape != (b, o, count) + or anchor.dtype != torch.long + or anchor.device != positions.device + ): + raise ValueError( + "anchor_idx must be int64 [B,O,anchor_k] on the geometry device" + ) + if (anchor < 0).any() or (anchor >= p).any(): + raise ValueError("Anchors must index stored points") + anchor_mask = torch.gather( + batch["valid_vtx_mask"], 3, anchor[:, None].expand(b, t, o, count) + ) + if not (anchor_mask | ~valid[:, None, :, None]).all(): + raise ValueError( + "Every valid object anchor must index valid points at every frame" + ) + if not torch.isfinite(positions).all(): + raise ValueError("Geometry, including padding, must be finite") + ref = batch.get("ref_obj_positions") + if ( + not isinstance(ref, torch.Tensor) + or ref.shape != (b, o, p, 3) + or ref.device != positions.device + or not torch.isfinite(ref).all() + ): + raise ValueError( + "ref_obj_positions must be finite aligned geometry [B,O,P,3]" + ) + props = batch.get("phys_params") + if ( + not isinstance(props, torch.Tensor) + or props.shape != (b, o, 3) + or props.device != positions.device + or not torch.isfinite(props).all() + ): + raise ValueError( + "phys_params must be finite [B,O,3] mass/friction/restitution" + ) + step = batch.get("skip_step_size") + if ( + not isinstance(step, torch.Tensor) + or step.shape != (b,) + or step.device != positions.device + or not torch.isfinite(step).all() + or (step <= 0).any() + ): + raise ValueError("skip_step_size must be positive finite [B]") + required = sum( + v.numel() * v.element_size() + for v in batch.values() + if isinstance(v, torch.Tensor) + ) + if required > self.cfg.max_batch_bytes: + raise ValueError("Input geometry exceeds max_batch_bytes") + if positions.device.type != "cuda": + raise ValueError( + "The author geometry/KNN path requires CUDA; CPU supports import and checkpoint inspection only" + ) + + def forward( + self, batch: Mapping[str, Any], *, teacher_forcing_ratio: float = 1.0 + ) -> dict[str, torch.Tensor]: + """Predict shifted training frames using explicit causal conditioning. + + Args: + batch: Author-layout geometry, masks, anchors and physical context. + teacher_forcing_ratio: Ratio in [0,1]; one matches released recipes. + + Returns: + Author prediction fields with frames aligned to positions[:,2:]. + """ + self._check(batch, min_frames=3) + if not 0.0 <= teacher_forcing_ratio <= 1.0: + raise ValueError("teacher_forcing_ratio must be in [0,1]") + return self.network(dict(batch), teacher_forcing_ratio=teacher_forcing_ratio) + + def rollout( + self, batch: Mapping[str, Any], *, num_frames: int, warmup_frames: int = 4 + ) -> dict[str, torch.Tensor]: + """Predict a free trajectory from warmup geometry. + + Args: + batch: Geometry containing exactly the declared warmup frames. + num_frames: Total returned frames including warmup. + warmup_frames: Initial ground-truth frames; four matches evaluation. + + Returns: + Author rollout fields. This inference method does not retain autograd. + """ + self._check(batch, min_frames=warmup_frames) + if batch["obj_positions"].shape[1] != warmup_frames: + raise ValueError("Rollout inputs must contain exactly warmup_frames") + if ( + type(num_frames) is not int + or num_frames <= warmup_frames + or warmup_frames < 2 + ): + raise ValueError("num_frames must exceed at least two warmup frames") + modes = [(m, m.training) for m in self.network.modules()] + try: + return self.network.rollout( + dict(batch), + num_steps=num_frames, + warmup_frames=warmup_frames, + return_extras=False, + ) + finally: + for module, mode in modes: + module.training = mode diff --git a/pyproject.toml b/pyproject.toml index bf33f43e5f..45e3e66893 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -57,6 +57,11 @@ dependencies = [ ] [project.optional-dependencies] +rigidformer = [ + "torch>=2.6", + "einops>=0.8.1", + "roma==1.5.4", +] 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 ca48a141b2..64144ae97c 100644 --- a/setup.py +++ b/setup.py @@ -185,6 +185,11 @@ def main(): package_data={ "embodichain": ["VERSION"], "embodichain_models.dynamics.nerd": ["THIRD_PARTY_LICENSES"], + "embodichain_models.dynamics.rigidformer": [ + "THIRD_PARTY_LICENSES", + "REFERENCE_SOURCES.json", + "configs/*.yaml", + ], "embodichain.data": ["default_materials.json"], "embodichain.compute.kinematics._warp": [ "ssik_LICENSE.txt", diff --git a/tests/models/dynamics/__init__.py b/tests/models/dynamics/__init__.py new file mode 100644 index 0000000000..45d11c7ba5 --- /dev/null +++ b/tests/models/dynamics/__init__.py @@ -0,0 +1,19 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +"""Dynamics model contract tests.""" diff --git a/tests/models/dynamics/_rigidformer_fixtures.py b/tests/models/dynamics/_rigidformer_fixtures.py new file mode 100644 index 0000000000..a02f93a9be --- /dev/null +++ b/tests/models/dynamics/_rigidformer_fixtures.py @@ -0,0 +1,60 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import torch +from embodichain_models.dynamics.rigidformer import RigidFormerCfg + + +def tiny_cfg() -> RigidFormerCfg: + cfg = RigidFormerCfg().reference_config() + m = cfg["model"] + m["encoder"].update( + output_dim=96, backbone_dim=32, hidden_dims=[32, 32], fusion_hidden_dim=32 + ) + m["state"]["token_dim"] = 96 + m["decoder"].update( + hidden_dim=96, + num_heads=1, + num_layers=1, + dropout=0.0, + num_register_tokens=2, + block_attn_res={"enable": False}, + ) + m["predictor"].update(hidden_dims=[96, 32], dropout=0.0) + m["avp"]["output_dim"] = 16 + return RigidFormerCfg(architecture=cfg) + + +def geometry_batch(device="cpu", points=8): + gen = torch.Generator().manual_seed(432) + ref = torch.randn(2, 2, points, 3, generator=gen) * 0.1 + ref[..., 2] += 1.0 + pos = torch.stack([ref + i * 0.002 for i in range(4)], 1) + v = { + "obj_positions": pos, + "valid_obj_mask": torch.tensor([[True, True], [True, False]]), + "valid_vtx_mask": torch.ones(2, 4, 2, points, dtype=torch.bool), + "ref_obj_positions": ref, + "anchor_idx": torch.arange(4).expand(2, 2, 4).clone(), + "phys_params": torch.tensor([1.0, 0.8, 0.3]).expand(2, 2, 3).clone(), + "skip_step_size": torch.tensor([1.0, 5.0]), + "dt": torch.ones(2), + "time_step_id": torch.zeros(2, dtype=torch.long), + "vtx_rank": torch.arange(points).expand(2, 2, points).clone(), + } + return {k: x.to(device) for k, x in v.items()} diff --git a/tests/models/dynamics/test_rigidformer_official.py b/tests/models/dynamics/test_rigidformer_official.py new file mode 100644 index 0000000000..e279cbe47a --- /dev/null +++ b/tests/models/dynamics/test_rigidformer_official.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. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import subprocess +import sys +import pytest +import torch +from ._rigidformer_fixtures import tiny_cfg, geometry_batch + + +def test_import_does_not_require_community_model_or_simulator(): + code = """import sys +from embodichain_models.dynamics.rigidformer import RigidFormerCfg +assert 'rigidformer' not in sys.modules +assert 'flash_attn' not in sys.modules +assert not any(n.startswith('embodichain.lab.sim') for n in sys.modules) +assert RigidFormerCfg().reference_config()['model']['kabsch_reference'] == 'fixed' +""" + subprocess.run([sys.executable, "-c", code], check=True) + + +@pytest.mark.gpu +def test_native_forward_checkpoint_and_rollout_with_variable_objects(tmp_path): + pytest.importorskip("roma") + from embodichain_models.dynamics.rigidformer import ( + RigidFormerModel, + load_rigidformer_checkpoint, + ) + + torch.manual_seed(42) + model = RigidFormerModel(tiny_cfg()).cuda().eval() + values = geometry_batch("cuda") + values["valid_obj_mask"][:] = True + for key in ("obj_positions", "valid_vtx_mask"): + values[key] = torch.cat((values[key], values[key][:, :, :1].clone()), dim=2) + for key in ("ref_obj_positions", "anchor_idx", "phys_params", "vtx_rank"): + values[key] = torch.cat((values[key], values[key][:, :1].clone()), dim=1) + values["obj_positions"][:, :, 2, :, 0] += 1 + values["ref_obj_positions"][:, 2, :, 0] += 1 + values["valid_obj_mask"] = torch.tensor( + [[True, True, True], [True, True, False]], device="cuda" + ) + with torch.no_grad(): + out = model(values) + assert torch.isfinite(out["obj_positions_denorm"]).all() + weights = tmp_path / "official-layout.pt" + torch.save( + { + "config": model.reference_config, + "model_state_dict": { + k: v.cpu() for k, v in model.network.state_dict().items() + }, + }, + weights, + ) + restored = load_rigidformer_checkpoint(weights, device="cuda") + with torch.no_grad(): + same = restored(values) + torch.testing.assert_close( + out["obj_positions_denorm"], same["obj_positions_denorm"], rtol=0, atol=0 + ) + rollout = restored.rollout(values, num_frames=8, warmup_frames=4) + assert rollout["obj_positions"].shape == (2, 8, 3, 8, 3) + assert torch.isfinite(rollout["obj_positions"]).all() + + +@pytest.mark.gpu +def test_native_packed_points_ignore_finite_padding_above_128(): + pytest.importorskip("roma") + from embodichain_models.dynamics.rigidformer import RigidFormerModel + + model = RigidFormerModel(tiny_cfg()).cuda().eval() + batch = geometry_batch("cuda", points=160) + batch["valid_obj_mask"][:] = True + batch["valid_vtx_mask"][:, :, 0, 96:] = False + with torch.no_grad(): + expected = model(batch)["obj_positions_denorm"] + changed = {k: v.clone() for k, v in batch.items()} + changed["obj_positions"][:, :, 0, 96:] += 10000 + changed["ref_obj_positions"][:, 0, 96:] -= 10000 + actual = model(changed)["obj_positions_denorm"] + valid = batch["valid_vtx_mask"][:, 2:] + torch.testing.assert_close(actual[valid], expected[valid], rtol=0, atol=0) From 5337f57453700c0e1662a39c15f7d021983e96ee Mon Sep 17 00:00:00 2001 From: yuecideng Date: Tue, 27 Oct 2026 15:54:47 +0000 Subject: [PATCH 2/3] fix(packaging): remove RigidFormer-specific Torch constraint --- pyproject.toml | 1 - 1 file changed, 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 45e3e66893..e0d6695510 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -58,7 +58,6 @@ dependencies = [ [project.optional-dependencies] rigidformer = [ - "torch>=2.6", "einops>=0.8.1", "roma==1.5.4", ] From acb01d18e28832557dbd4bff39ede7b6d12bb02f Mon Sep 17 00:00:00 2001 From: yuecideng Date: Wed, 28 Oct 2026 01:29:22 +0000 Subject: [PATCH 3/3] fix(model): isolate RigidFormer inference and rotary state --- .../_native/models/physformer_rope.py | 15 ++-- .../dynamics/rigidformer/_native/network.py | 10 ++- .../rigidformer/_native/utils/cuda_knn.py | 4 +- .../models/dynamics/_rigidformer_fixtures.py | 1 - .../dynamics/test_rigidformer_official.py | 70 +++++++++++++++++++ 5 files changed, 90 insertions(+), 10 deletions(-) diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/physformer_rope.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/physformer_rope.py index 7746b6461e..4317984b3f 100644 --- a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/physformer_rope.py +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/models/physformer_rope.py @@ -68,10 +68,13 @@ def _memo_key(tag, x, extra=()): def _memo_usable(*tensors): # Under torch.compile the graph de-duplicates identical subexpressions itself; the memo (which reads storage - # pointers) is skipped there. - if torch.compiler.is_compiling(): + # pointers) is skipped there. Inference tensors have no version counter, + # and caching an inference-mode result could contaminate a later autograd run. + if torch.compiler.is_compiling() or torch.is_inference_mode_enabled(): return False - return ROPE_MEMO_ENABLED and not any(t.requires_grad for t in tensors) + return ROPE_MEMO_ENABLED and not any( + t.requires_grad or t.is_inference() for t in tensors + ) def _cos_sin(freqs): @@ -126,12 +129,11 @@ class PhysFormerRoPE(nn.Module): # Pair sin/cos channels the way rotate_half pairs them (interleaved), which # keeps the relative-position property; the half-split layout does not. A # model is evaluated with the layout it was trained with. - INTERLEAVED_PAIRING = False _PAIRING_LOGGED = False def _paired_angles(self, angles: torch.Tensor) -> torch.Tensor: cls = type(self) - if cls.INTERLEAVED_PAIRING: + if self.interleaved_pairing: if not cls._PAIRING_LOGGED: cls._PAIRING_LOGGED = True print( @@ -171,6 +173,7 @@ def __init__( """ super().__init__() + self.interleaved_pairing = False self.dim = dim self.use_3d = use_3d self.include_time = include_time @@ -263,7 +266,7 @@ def _freqs_signature(self): def _memo_extra(self): return ( self._freqs_signature(), - type(self).INTERLEAVED_PAIRING, + self.interleaved_pairing, self.n_coords, self.include_time, self.dim_per_coord_3d, diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/network.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/network.py index b36b2d9d45..0b20da64d2 100644 --- a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/network.py +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/network.py @@ -133,7 +133,7 @@ def __init__(self, config): PhysFormerRoPE as _RopeCls, ) - _RopeCls.INTERLEAVED_PAIRING = bool( + rope_interleaved_pairing = bool( config["model"].get("rope_interleaved_pairing", False) ) @@ -429,6 +429,8 @@ def __init__(self, config): else False ) + # The pinned author implementation uses the decoder's 16-token default, + # even if the source config contains another num_register_tokens value. self.decoder = BidirectionalDecoder( state_dim=config["model"]["state"].get("token_dim", 768), num_layers=config["model"]["decoder"].get( @@ -790,6 +792,12 @@ def __init__(self, config): print(_rope_msg) logger.info(_rope_msg) + # Pairing belongs to this model, not to the RoPE class shared by all + # model instances. This preserves each checkpoint's saved policy. + for module in self.modules(): + if isinstance(module, _RopeCls): + module.interleaved_pairing = rope_interleaved_pairing + def forward(self, batch, return_all=False, teacher_forcing_ratio=1.0): """ Forward pass through the model. diff --git a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/utils/cuda_knn.py b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/utils/cuda_knn.py index e9ff1cb077..ebbe24e3c5 100644 --- a/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/utils/cuda_knn.py +++ b/extensions/embodichain_models/embodichain_models/dynamics/rigidformer/_native/utils/cuda_knn.py @@ -9,8 +9,8 @@ import torch.nn.functional as F from typing import List, Tuple, Optional -# CUDA is mandatory for this module; there is intentionally no CPU fallback. -assert torch.cuda.is_available(), "cuda_knn requires a CUDA-capable device." +# The public model checks for CUDA when geometry is executed. Keep this module +# importable without CUDA so model construction and checkpoint inspection work. def compute_nearest_distance_to_list_cuda( diff --git a/tests/models/dynamics/_rigidformer_fixtures.py b/tests/models/dynamics/_rigidformer_fixtures.py index a02f93a9be..d15db1a84c 100644 --- a/tests/models/dynamics/_rigidformer_fixtures.py +++ b/tests/models/dynamics/_rigidformer_fixtures.py @@ -32,7 +32,6 @@ def tiny_cfg() -> RigidFormerCfg: num_heads=1, num_layers=1, dropout=0.0, - num_register_tokens=2, block_attn_res={"enable": False}, ) m["predictor"].update(hidden_dims=[96, 32], dropout=0.0) diff --git a/tests/models/dynamics/test_rigidformer_official.py b/tests/models/dynamics/test_rigidformer_official.py index e279cbe47a..8f250de780 100644 --- a/tests/models/dynamics/test_rigidformer_official.py +++ b/tests/models/dynamics/test_rigidformer_official.py @@ -16,6 +16,8 @@ from __future__ import annotations +import os +from pathlib import Path import subprocess import sys import pytest @@ -34,6 +36,67 @@ def test_import_does_not_require_community_model_or_simulator(): subprocess.run([sys.executable, "-c", code], check=True) +def test_cpu_only_model_construction_and_checkpoint_import(tmp_path): + pytest.importorskip("roma") + code = """import importlib.util +import sys +import torch +from embodichain_models.dynamics.rigidformer import RigidFormerModel, load_rigidformer_checkpoint + +assert not torch.cuda.is_available() +spec = importlib.util.spec_from_file_location('fixtures', sys.argv[2]) +fixtures = importlib.util.module_from_spec(spec) +spec.loader.exec_module(fixtures) +model = RigidFormerModel(fixtures.tiny_cfg()) +torch.save({'config': model.reference_config, 'model_state_dict': model.network.state_dict()}, sys.argv[1]) +restored = load_rigidformer_checkpoint(sys.argv[1]) +assert all(v.device.type == 'cpu' for v in restored.state_dict().values()) +batch = fixtures.geometry_batch() +batch['valid_obj_mask'][:] = True +try: + restored(batch) +except ValueError as exc: + assert 'requires CUDA' in str(exc) +else: + raise AssertionError('CPU geometry execution must be rejected') +""" + env = os.environ.copy() + env["CUDA_VISIBLE_DEVICES"] = "" + subprocess.run( + [ + sys.executable, + "-c", + code, + str(tmp_path / "weights.pt"), + str(Path(__file__).with_name("_rigidformer_fixtures.py")), + ], + check=True, + env=env, + ) + + +def test_rope_pairing_is_owned_by_each_model(): + pytest.importorskip("roma") + from embodichain_models.dynamics.rigidformer import RigidFormerCfg, RigidFormerModel + from embodichain_models.dynamics.rigidformer._native.models.physformer_rope import ( + PhysFormerRoPE, + ) + + first = RigidFormerModel(tiny_cfg()) + first_rope = next(m for m in first.modules() if isinstance(m, PhysFormerRoPE)) + positions = torch.tensor([[0.1, 0.2, 0.4]]) + expected = first_rope.forward_3d(positions).clone() + config = tiny_cfg().reference_config() + config["model"]["rope_interleaved_pairing"] = False + second = RigidFormerModel(RigidFormerCfg(architecture=config)) + second_rope = next(m for m in second.modules() if isinstance(m, PhysFormerRoPE)) + + torch.testing.assert_close( + first_rope.forward_3d(positions), expected, rtol=0, atol=0 + ) + assert not torch.equal(second_rope.forward_3d(positions), expected) + + @pytest.mark.gpu def test_native_forward_checkpoint_and_rollout_with_variable_objects(tmp_path): pytest.importorskip("roma") @@ -77,6 +140,13 @@ def test_native_forward_checkpoint_and_rollout_with_variable_objects(tmp_path): rollout = restored.rollout(values, num_frames=8, warmup_frames=4) assert rollout["obj_positions"].shape == (2, 8, 3, 8, 3) assert torch.isfinite(rollout["obj_positions"]).all() + with torch.inference_mode(): + inference = restored.rollout(values, num_frames=8, warmup_frames=4) + torch.testing.assert_close( + inference["obj_positions"], rollout["obj_positions"], rtol=0, atol=0 + ) + # Inference-mode results must not leave inference tensors in the RoPE memo. + restored(values)["obj_positions_denorm"].sum().backward() @pytest.mark.gpu