[yuecideng]

Description

Add the author's RigidFormer geometry architecture at 160fe0852c51cffae7b7440452b58d0b49db6df6 as embodichain_models.dynamics.rigidformer, with its MIT attribution and three MOVi configurations. Strict official tensor/config imports preserve named architecture, shapes and dtypes; the model uses explicit Torch SDPA or a compatible Flash-Attention backend.

This layer owns model construction, geometry validation and inference. CUDA geometry requires at least two valid objects; CPU construction and weight inspection remain available. Model integration uses dt=1, independently of physical seconds.

Stack

• Layer: 1/3
• Base: main
• Depends on: none
• Followed by: #889

Type of change

• New feature (non-breaking)

Validation

• The rigidformer extra declares only einops and roma; Torch selection remains governed by the project environment. Release-metadata checks (8 passed) and the built wheel confirm there is no model-specific Torch requirement.

• 53 focused model, context and release-metadata checks pass, including native CUDA checkpoint/rollout and finite-padding coverage above 128 points.

• Its wheel installs outside the checkout with no repository PYTHONPATH. Installed source bytes, configs/licenses and all three presets are checked; official-weight predictions match the frozen source baseline exactly.

• API coverage: 3271/3271. Context routing/checks and Black 26.10.1 pass.

• Focused model coverage protects independent loading/padding/import contracts. Shared training and reference qualification are supplied by the dependent layers.

Review follow-up

• CPU-only model construction and strict checkpoint import now work without CUDA; geometry execution still requires CUDA. The rebuilt model-only wheel was installed outside the checkout to verify this behavior, the license and the optional dependencies.
• Inference-mode rollout matches ordinary rollout without storing inference tensors in the rotary cache. Differently configured models retain independent rotary pairing.
• Post-fix validation: 86 affected CPU tests, two model CUDA tests and the two-rank NCCL test passed. An independent official Sphere/Flash-Attention comparison passed three updates with zero recorded prediction, rollout, loss, gradient, parameter, optimizer, EMA and scheduler differences. Full-model continuous versus resumed training also matched exactly.
• The pinned author src/train.py does not forward decoder.num_register_tokens to the decoder. The adapter deliberately keeps the author's 16-token default even when a source config names another value, preserving the layout of official checkpoints.
• Affected neural-dynamics context was reviewed; its existing model ownership and CUDA geometry guidance remain accurate.

Checklist

• Ran black . with Black 26.10.1.
• Reviewed and updated affected documentation and agent context.
• Public API changes have useful API documentation.
• Verified proportional behavioral coverage and reviewed test necessity.
• Updated optional dependencies and packaged resources where applicable.

[greptile-apps]

Confidence Score: 5/5

[Medium impact] The reviewed changes appear safe to merge, with no outstanding findings.

Summary

Adds the pinned author's RigidFormer architecture, three MOVi presets, strict weight imports, and a public geometry API.

• RigidFormer predicts rigid geometry from CUDA batches and can roll out future frames.
• Official RigidFormer weights load into a matching native model.

No new actionable issues were found in the changes since the previous review.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  Config[Saved configuration or MOVi preset] --> Model[RigidFormerModel]
  Weights[Official named tensors] --> Import[Strict checkpoint import]
  Import --> Model
  Geometry[Geometry, masks, anchors and physics] --> Check[Check geometry contract]
  Check --> CUDA[CUDA geometry prediction]
  Model --> CUDA
  CUDA --> Output[Predicted frames or free rollout]

Loading

_{Reviews (2) · Last reviewed commit: "fix(model): isolate RigidFormer inferenc..." · Reviewed by Greptile}