[yuecideng]

Description

Add optional RigidFormer geometry prediction and fitting as embodichain_models.dynamics.rigidformer, using the shared SL trainer from main@c5b1fdfc2, including the merged #855 foundation.

The extension owns the lazy upstream adapter, aligned geometry windows, objective and installable recipe. Reference qualification moves to examples/learning/rigidformer, with private support outside the installable model package. Reference geometry, point identities, object/point counts and positive integration intervals remain explicit. Core SL and the generic data pipeline do not import this model; no separate registry or trainer is introduced.

The adapter predicts geometry without action inputs. It does not qualify action gradients, APG or native physics-solver deployment. The loss aggregation weight counts valid objects, matching the upstream denominator under DDP.

Optional dependency and entry points

python -m pip install 'embodichain[rigidformer]'
embodichain train-sl --config examples/learning/rigidformer/configs/sl.train.yaml
embodichain run-example learning/rigidformer --help

Set data.path to the tensor geometry archive and use a fresh output directory. The installable factory is embodichain_models.dynamics.rigidformer.sl_recipe:build_rigidformer_trainer. See the showcase.

The published extra pins rigidformer==0.2.0. The official PyPI wheel's four runtime Python files, license, Python requirement and dependency specifiers were verified against Git revision 4ab4edac7271b49484748d6bf8086a497f13272d. The exact Git installation remains in the showcase's requirements-reference.txt; published project/wheel metadata contains no direct URL dependency. The unchanged release metadata gate and actual wheel validator pass. Frozen reference qualification still checks executed checkout/installed bytes against the pinned Git blobs.

Complete upstream MIT attribution ships in the model family's THIRD_PARTY_LICENSES.

Main synchronization

Main synchronization at b5f4bf8e2: all 14 model/example execution files remain byte/AST-identical, including the conservative BF16 distance budget and the verified rigidformer==0.2.0 optional dependency. Main's SL/vendor/CI/release rules are preserved. Current focused validation passes 158 production/metadata/context checks plus four staged-factory/CLI checks; historical one-to-two-update recovery and the original dimension-128 checkpoint match upstream exactly. The same unrelated open3d.cpu simulation-help limitation is retained in evidence.

Validation and review fixes

Current branch source: b5f4bf8e2c267e1956d4601097c57126b3495917; prior qualified source: c29b2c6f09c48f1c024822511b2da94b1d992f80; extension migration: 5e0b3b144; memory guard correction: 9b53fa759; normal base merge: ddb7c5931.

• Initial migration validation: 47 production/reference/SL checks, 88 context/docs/example checks and 3 package checks pass. Model/data implementations retain baseline bytes, and extracted objectives/recipe/reference math match baseline ASTs after import mapping.
• On GPU 5, friction and collision each run a paired seed-42 eight-update representative check: every paired loss differs by zero; gradients, updates and held-out rollout endpoints match. Generated corpora match the frozen corpus hashes. The original five-seed/1,000-update gate remains pending for these short runs; thresholds were not reduced.
• Historical source 8fe234160 retains the full friction/collision × seeds42–46 × reference/SL ×1,000-update qualification. Earlier review hardened executed-source checks, full paired curves, failure records and unequal-padding model tests. These immutable records are not relabeled as reruns of the relocation.
• Independent AI review found a real BF16-autocast memory-bound defect: two-byte geometry can produce a four-byte cdist buffer. A real red-to-green regression now budgets at least FP32, retaining eight-byte FP64 accounting and adding no per-call tensor allocation. Twenty-six focused checks pass; default FP32 predictions are bitwise unchanged.
• Combined release checks found the prior VCS extra was unpublishable. Moving it to an exact public-index pin with proven source equivalence fixes the unchanged gate. Twenty focused release/model checks and the actual independent wheel metadata validator pass; the PyPI wheel itself performs a real model update/prediction in isolation.
• In the combined installed wheel, SDK-blocked canonical imports, geometry factory, two-to-four update recovery and example CLI work after deleting build source. Two-worker variable-object training matches the independent full batch; maximum parameter difference is 1.50e-8. Actual wheel metadata has no direct URLs and installs the reference requirements alongside the example.
• Black 26.10.1, context checks and API coverage 3182/3182 pass. The full-site Sphinx dummy build succeeds with 953 existing site warnings; no whole-site strict claim is made. The review's memory-bound finding is fixed and independently verified.

The intermediate budget bounds the dense nearest-neighbor distance buffer, not total process/GPU memory. Unequal point counts require padded P<=128 with the pinned hierarchical encoder. Preserve original source/data/normalization identities; offline parity does not establish paper-level accuracy or physical robot performance.

CI integration repairs

Current review head: ac69eb828bb66b419c461de0efb593605b212210. The unchanged main CI exposed two tests left behind by the vendor migration: narrow-terminal label wrapping and a retired UR mock import. The exact two-file maintenance commit b1d8ab415 is reviewed separately in #878 and included here; semantic assertions, CI selection/gates and production behavior are unchanged. Its original two failures reproduce and the corrected related modules pass 42 tests.

Dependencies and review

• Base: main, synchronized at c5b1fdfc28ea47376933d95207b7894b7e6d54f1. Add NeRD model extension and an SL training showcase #855, Add shared SL updates, offline training and recovery #857 and Route NeRD training through the shared SL learner #860 are merged.
• DDP Add single-node DDP training to shared SL #864 and RWM Add robot world models through the model extension and SL #859 remain siblings; the combined final wheel and model oracles were validated separately from this focused PR.
• Ordinary SL imports keep the optional upstream model and simulation unloaded.
• Merge requires the requesting user's review; keep auto-merge disabled.

Type of change

• New feature, extension migration, memory-budget fix and packaging compatibility.

Checklist

• Run black . with Black 26.10.1.
• Review context and document model, example and optional dependency ownership.
• Document exports and pass API coverage.
• Reuse independent reference/gradient/padding/recovery oracles and close the actual BF16 budget gap.
• Pass unchanged release requirements and actual wheel metadata checks.
• Preserve MIT attribution, frozen evidence and review-required merging.

[greptile-apps]

Confidence Score: 5/5

[Medium impact] The PR appears safe to merge; no new blocking issues were found.

Summary

Adds optional RigidFormer geometry prediction and training through the shared SL trainer.

• RigidFormer predicts future point geometry and trains through the shared supervised-learning recipe.
• A separate Box2D example compares the shared trainer with a pinned RigidFormer reference.
• Two existing test fixtures use the current robot import path and a fixed terminal width.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Geometry archive] --> B[RigidBodySequenceSource]
  B --> C[rigidformer_batch]
  C --> D[Shared SL trainer]
  D --> E[RigidFormerObjective]
  E --> F[RigidFormerModel]
  F --> G[Optional upstream model]

Loading

_{Reviews (7) · Last reviewed commit: "Include vendor catalog test compatibilit..." · Reviewed by Greptile}