[yuecideng]

Description

Add robot world-model fitting, graph-connected prediction and bounded transition replay as the installable embodichain_models.dynamics.rwm family. Synchronize with main@c5b1fdfc2, including the merged #855 foundation and use its shared SL trainer and durable recovery contracts.

The model predicts frame-local robot-state increments from state, action and optional physical context, then integrates the next state in FP32. Its state math, model, objective, predictor, replay and installable sl_recipe now live together in the model extension. Controlled corpus generation and fitting move to examples/learning/rwm; core SL and the generic data pipeline do not import the model family. No simulator dependency, compatibility shim, separate registry or parallel trainer is introduced.

RWM belongs to the dynamics domain because its output is future robot state. There is no parallel world_models namespace. The migration preserves action/timestep/joint identity, frozen normalization, checkpoint state keys and physical arithmetic.

Refs #848.

Entry points and ownership

conda activate py311
OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 embodichain run-example learning/rwm \
  --output outputs/rwm-controlled-seed0 --seed 0

The example writes a resolved recipe for embodichain train-sl --config FILE. Its factory is embodichain_models.dynamics.rwm.sl_recipe:build_rwm_trainer, also exported from the family root. See the showcase.

Replay owns detached pre-reset transitions and returns independent episode-safe windows. The predictor preserves state/action graphs while allowing model parameters/statistics to be frozen. Invalid padded actions/context are sanitized before unrolling; invalid targets are selected out before error arithmetic and backward. Collection, task rewards, APG orchestration and physical reset remain caller-owned.

Main synchronization

Main synchronization at e6a8ce525: all 12 model/example execution files remain byte/AST-identical; shared SL, vendor/CI and release code match main. Current focused validation passes 126 production/metadata/context checks plus four staged-factory/CLI checks. The real historical 300-update checkpoint reproduces its metrics and gradients under the canonical namespace. An unrelated four-case simulation-help selection still hits this environment's existing missing open3d.cpu; its failed log is preserved and is not called green.

Validation

Current branch source: e6a8ce5251bd5826629f67048509cbb89ba72d43; prior qualified source: dcb76691103b168d741db1202378d169e796f543; implementation: 07c51a24a; normal base merge: c38cd89b6.

• 135 focused CPU checks pass; the controlled GPU/BF16/replay regression passes on physical GPU 3. The initial accidentally unisolated baseline GPU invocation is retained in evidence and is not counted as the controlled GPU result.
• Seeds 0/1/2 each complete 300 controlled updates. Final/initial held-out error ratios remain 0.0037039333 / 0.0011929718 / 0.0013287687, matching the earlier implementation. The 29-joint/three-context layout and action-gradient check also passes. This is controlled fitting, not SuperMimic or robot-policy reproduction.
• A real step-150 checkpoint from old source 4a5643875 loads under the new namespace. Predictions, state/action gradients and state tensors match exactly; resuming to update 151 reproduces the recorded update metrics.
• The actual wheel works without repository fallbacks or simulator/example imports: two updates and recovery to four match uninterrupted training. Installed run-example learning/rwm also completes the controlled fit.
• In the final combined DDP/model wheel, installed RWM two-to-four recovery matches the full run's model/optimizer/sampler/RNG/metrics. A real two-worker three-step RWM objective with NaN padding matches independent full-batch loss/gradients/updates; maximum parameter difference is 9.32e-10.
• Black 26.10.1, context checks and API coverage 3190/3190 pass. RWM/SL real pages build strictly with zero warnings. The full-site build remains blocked in this environment by the existing open3d.cpu dependency, whose failed log is preserved.
• Independent architecture/API/AI/security review found no unresolved actionable defect. Source equivalence, episode/replay ownership, padding and checkpoint contracts were inspected as well as tested.

Runtime results remain bound to their recorded source hashes. No long policy/APG training, native learned-solver rollout, ODS integration or new convergence claim is included.

CI integration repairs

Current review head: 2f5e450dcd491fc16ddce0427ac7623be48b55dd. 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.

This branch also renames the new replay test to a unique module basename for default pytest collection and adds the actual RWM architecture node, source evidence, API link and SL connections. The complete architecture snapshot is refreshed under the existing maintenance contract; 94 focused Python checks, five frontend data checks and the frontend build pass. Model code and numerical qualification are unchanged.

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 RigidFormer Add RigidFormer dynamics models and SL showcase #858 remain sibling changes. Their combined installation and training were checked without adding them to this PR's diff.
• The family ships in the main embodichain wheel/version. No new mandatory dependency is added.
• Merge requires the requesting user's review; keep auto-merge disabled.

Type of change

• New feature, extension migration and documentation update.

Checklist

• Run black . with Black 26.10.1.
• Review context ownership and add the robot-world-model navigation owner.
• Document canonical exports and pass API coverage.
• Reuse production fitting/gradient/replay/padding/recovery coverage.
• Validate old checkpoints, actual wheel entry points and DDP integration.
• Preserve source-bound evidence and review-required merging.