[yuecideng]

Description

Train prepared MOVi geometry through the existing supervised learner/trainer, preserving author shifted anchor losses, grouped AdamW, epoch scheduling and ramped FP32 EMA. Immutable geometry, splits, scales and committed anchor/FPS caches bind source recovery; a compact index retains the author window mapping without repeated path dictionaries.

Add ordered SupervisedTrainingState participants that commit after optimizer success and retain rank-local checkpoint/validation state. Preserve main's #877 source-recovery preflight. DDP buffer broadcast is opt-in; empty/default recipes retain checkpoint formats 1/2, while required participant/buffer state uses 3/4. Raw/EMA export remains distinct from importing official optimizer state.

Stack

• Layer: 2/3
• Base: main
• Depends on: #888 (merged)
• Followed by: #890

Type of change

• New feature (non-breaking)

Validation

• After the model review fixes, 86 affected CPU tests and the two-rank NCCL recovery/weighted-update test passed. The full-model continuous versus resumed run matched exactly, and the three-update official Flash comparison remained exact.
• 126 affected CPU SL/model/context checks pass. The complete stack's six CUDA tests also pass, including real two-rank NCCL/SyncBatchNorm, unequal 3/2-object weighted supervision and exact rank recovery across three epochs/two warmup epochs.
• The concatenated-batch oracle disables stochastic dropout; recovery retains author dropout. Maximum two-rank state error versus the oracle is 1.79e-7, within atol=1e-5/rtol=1e-4. Continuous/resumed model, optimizer, EMA, schedule, RNG, cursor and buffers match exactly.
• The installed wheel verifies source bytes and inference equality, runs the public train-sl CLI through recovery, and strictly reloads raw/EMA exports without a repository PYTHONPATH.
• API coverage: 3284/3284. Black and context checks pass. Coverage adds concrete legacy-checkpoint, sampling/index, manifest-integrity and real model-recovery cases.
• Index-only CPU measurements at 1/32/960 scenes report latency and Python allocations separately; no GPU/training-throughput claim is made.

The development Blackwell host used process-local NCCL 2.29.7 built for sm_120 with CUDA 12.8. Its inherited NCCL 2.26.2 also failed an independent parameter-layout-only DDP broadcast. This correctness validation uses shared-memory transport with P2P/IB disabled; the dependent acceptance protocol records reproduction details.

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: 4/5

[Medium impact] The PR adds a RigidFormer training workflow and optional recovery features to shared supervised learning.

The PR appears safe to merge, with a non-blocking correction needed to preserve training settings in exported weights.

Findings

1.  Exported settings lose overrides ▶

Fix with agent prompt

### Issue 1
extensions/embodichain_models/embodichain_models/dynamics/rigidformer/checkpoint.py:131
When a recipe overrides training settings or sampling, `export_rigidformer_checkpoint` still writes the unchanged `reference_config`. Changed loss weights or step-size choices disappear from the exported file, leaving readers with settings that do not describe how the weights were trained. Copy the reference configuration and merge the saved `training` and `sampling` settings before exporting it.

---

For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.

Summary

Adds official RigidFormer training through the shared supervised trainer.

• Shared training runs can save and restore extra update state.
• Prepared MOVi scenes feed training as verified, rank-sharded windows.
• RigidFormer trains from shifted geometry and exports raw or EMA weights.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Verified MOVi scenes] --> B[Shifted geometry batch]
  B --> C[Anchor losses]
  C --> D[Optimizer update]
  D --> E[Commit source cursor]
  E --> F[Advance epoch schedule]
  F --> G[Update EMA]
  G --> H[Save rank-local checkpoint]
  H --> I[Export raw or EMA weights]

Loading

_{Reviews (1) · Last reviewed commit: "fix(sl): harden RigidFormer training and..." · Reviewed by Greptile}

[greptile-apps]

Exported settings lose overrides

When a recipe overrides training settings or sampling, export_rigidformer_checkpoint still writes the unchanged reference_config. Changed loss weights or step-size choices disappear from the exported file, leaving readers with settings that do not describe how the weights were trained. Copy the reference configuration and merge the saved training and sampling settings before exporting it.

Knowledge Base Used: RigidFormer dynamics models

Prompt To Fix With AI

This is a comment left during a code review.
Path: extensions/embodichain_models/embodichain_models/dynamics/rigidformer/checkpoint.py
Line: 131

Comment:
**Exported settings lose overrides**

When a recipe overrides training settings or sampling, `export_rigidformer_checkpoint` still writes the unchanged `reference_config`. Changed loss weights or step-size choices disappear from the exported file, leaving readers with settings that do not describe how the weights were trained. Copy the reference configuration and merge the saved `training` and `sampling` settings before exporting it.

**Knowledge Base Used:** [RigidFormer dynamics models](https://app.greptile.com/dexforce/-/custom-context/knowledge-base/dexforce/embodichain/-/docs/rigidformer-dynamics.md)

---

For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.