Train official RigidFormer through shared SL

agent_context/MAP.yaml

@@ -273,7 +273,7 @@ topics:
   aliases: [neural dynamics, neural robot dynamics, NeRD, RigidFormer, 神经动力学, 动力学模型训练, NeRD 示例, RigidFormer 示例]
   keywords: [NeRDStateCodec, NeRDModel, DynamicsGenerationCfg,
     DynamicsShardWriter, DynamicsWindowDataset, DynamicsTrainer, embodichain_models,
-    learning/nerd, nerd reference assets, RigidFormerModel,
+    learning/nerd, nerd reference assets, RigidFormerModel, MOViSequenceSource, build_rigidformer_trainer,
     官方训练精度对比]
   paths: [topics/neural-dynamics/neural-dynamics.md]
   source_of_truth:

agent_context/topics/neural-dynamics/neural-dynamics.md

@@ -47,11 +47,18 @@ The native family lives in
 `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.
+`data.py` reads verified packed scenes; `supervised.py` separates shifted
+conditioning and anchor targets; `sl_recipe.py` constructs the shared trainer
+with committed sampling, epoch scheduling and EMA participants.
+
+Integration dt=1 differs from physical seconds. Stored vertices are unscaled;
+metadata sizes apply about each origin once. Geometry prediction requires CUDA
+and at least two valid objects. Attention backend and source semantics belong
+to recipe/recovery identity. Source-owned epoch sampling does not claim the
+legacy research DataLoader's worker-RNG trajectory.
+
+Numerical agreement, recovery and full benchmark/retraining qualification
+remain separate evidence requirements.
 
 ## Resource and physical identity
 

agent_context/topics/supervised-learning/supervised-learning.md

@@ -42,8 +42,9 @@ Objectives own loss masks, normalization, integration and multi-step graph
 construction. The optimizer does not infer action semantics or detach histories.
 Under DDP, `LossResult.loss` is a local mean and `weight` is its effective valid
 supervision count; the learner forms the global weighted mean. Batch size is per
-rank, the learning rate is unchanged, and rank-local sampling streams draw with
-replacement, so samples may overlap. All ranks call fitting, saving and recovery;
+rank and the learning rate is unchanged. Sampling belongs to the source: tensor
+sources draw with replacement, while epoch-traversal sources may shard windows.
+All ranks call fitting, saving and recovery;
 only rank 0 writes artifacts or runs validation through the unwrapped model.
 Schema, action/frame/timestep and numerical-policy identifiers belong in each
 recipe's checkpoint signature. Tensor sources additionally bind content identity.
@@ -57,8 +58,10 @@ changes. Live buffers and delivery credits are not restored: such sources may
 train/save snapshots, but exact continuation is rejected. Acquisition, waits,
 refill and delivery debits stay source/caller-owned; encoding failures do not
 refund a delivered context.
-Distributed checkpoints retain each rank's RNG, sampler and buffers and require
-the same world size for exact recovery. Inside rank 0's validation callback only,
+Distributed checkpoints retain each rank's RNG, sampler, buffers and participant
+state and require the same world size for exact recovery. Rank-local snapshots
+are transferred to rank 0 in bounded chunks to limit device memory. Inside rank
+0's validation callback only,
 a save uses the previously gathered committed-update snapshot without a new
 collective; ordinary save calls require every rank to participate.
 
@@ -80,3 +83,9 @@ qualification belongs to each model's acceptance recipe. The distributed tests
 exercise weighted global updates and fixed-world-size recovery with CPU/Gloo;
 sequence tests reuse real CPU source fixtures for masks, provenance, owned output,
 quota and recovery admission. Examples use launch smoke checks, not tutorial tests.
+
+Factories may register ordered `SupervisedTrainingState` participants for source
+cursors, schedules and model averages. They commit after optimizer success,
+produce owned rank-local snapshots, validate before live restoration and restore
+after validation callbacks. Empty participants retain existing checkpoint
+compatibility. Model-specific participant math stays in its model family.

docs/source/api_reference/embodichain/embodichain.learning.sl.rst

@@ -30,8 +30,8 @@ The example factory uses a small masked MLP regression corpus generated locally.
 For ``WORLD_SIZE > 1``, the workflow initializes the process group before invoking
 the recipe factory. It resolves ``trainer.device`` to ``cpu`` or
 ``cuda:LOCAL_RANK``. ``trainer.batch_size`` is per worker; the global number of
-drawn examples is ``batch_size * world_size``. Each worker samples independently
-with replacement, so examples can overlap. The learning rate is not scaled
+drawn examples is ``batch_size * world_size``. Sampling semantics belong to each source: the tensor example samples independently
+with replacement, while explicit offline adapters may shard an epoch traversal. The learning rate is not scaled
 automatically.
 
 Objectives return a local mean loss and its valid supervision count in
@@ -99,6 +99,7 @@ embodichain.learning.sl
    LossResult
    SupervisedBatch
    SupervisedObjective
+   SupervisedTrainingState
    SupervisedLearner
    SupervisedSource
    TensorBatchSource
@@ -116,6 +117,9 @@ embodichain.learning.sl
 .. autoclass:: SupervisedObjective
    :members:
 
+.. autoclass:: SupervisedTrainingState
+   :members:
+
 .. autoclass:: SupervisedLearner
    :members:
 
@@ -234,6 +238,7 @@ embodichain.learning.sl.types
    SupervisedBatch
    LossResult
    SupervisedObjective
+   SupervisedTrainingState
 
 .. autoclass:: SupervisedBatch
    :members:
@@ -244,6 +249,9 @@ embodichain.learning.sl.types
 .. autoclass:: SupervisedObjective
    :members:
 
+.. autoclass:: SupervisedTrainingState
+   :members:
+
 embodichain.learning.sl.workflow
 --------------------------------
 
@@ -265,3 +273,17 @@ Concrete models own their physical codecs, objectives and source adapters in
 the :doc:`official model library </api_reference/model_catalogs>`. The NeRD
 recipe connects committed trajectory windows to this common trainer without
 loading simulation or requiring the example runner.
+
+Recoverable update participants
+-------------------------------
+
+Factories may pass ordered ``training_states`` implementing
+``SupervisedTrainingState``. Participants commit after optimizer success; their
+owned state is saved per rank, validated before live recovery mutation, and
+restored after validation callbacks. This supports data cursors and model
+averages without model-specific trainer subclasses. Default empty participants
+retain compatibility with existing checkpoints.
+
+Checkpoint formats 3/4 mark nonempty training participants or nondefault buffer
+broadcasting, so older readers reject required state instead of discarding it.
+Default recipes continue to write/read formats 1/2.

docs/source/api_reference/model_rigidformer.rst

@@ -14,7 +14,7 @@ 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
+Backend identity is part of recipe/recovery semantics. Inference rollout restores
 the caller's module modes and does not retain a gradient graph.
 
 Public model family
@@ -24,10 +24,14 @@ Public model family
 
 .. autosummary::
 
-   RIGIDFORMER_REVISION
    RigidFormerCfg
    RigidFormerModel
    load_rigidformer_checkpoint
+   export_rigidformer_checkpoint
+   MOViSequenceSource
+   RigidFormerObjective
+   rigidformer_batch
+   build_rigidformer_trainer
 
 .. autodata:: RIGIDFORMER_REVISION
 
@@ -39,14 +43,25 @@ Public model family
 
 .. autofunction:: load_rigidformer_checkpoint
 
+.. autofunction:: export_rigidformer_checkpoint
+
+.. autoclass:: MOViSequenceSource
+   :members:
+
+.. autoclass:: RigidFormerObjective
+   :members:
+
+.. autofunction:: rigidformer_batch
+
+.. autofunction:: build_rigidformer_trainer
+
 Architecture and geometry contract
 ----------------------------------
 
 .. currentmodule:: embodichain_models.dynamics.rigidformer.model
 
 .. autosummary::
 
-   RIGIDFORMER_REVISION
    RigidFormerCfg
    RigidFormerModel
 
@@ -58,13 +73,78 @@ Architecture and geometry contract
 .. autoclass:: RigidFormerModel
    :members:
 
-Official inference weights
---------------------------
+Official weights and native SL export
+-------------------------------------
 
 .. currentmodule:: embodichain_models.dynamics.rigidformer.checkpoint
 
 .. autosummary::
 
    load_rigidformer_checkpoint
+   export_rigidformer_checkpoint
 
 .. autofunction:: load_rigidformer_checkpoint
+
+.. autofunction:: export_rigidformer_checkpoint
+
+Prepared packed scenes
+----------------------
+
+.. currentmodule:: embodichain_models.dynamics.rigidformer.data
+
+.. autosummary::
+
+   MOViSequenceSource
+
+.. autoclass:: MOViSequenceSource
+   :members:
+
+Shifted supervision and anchor losses
+-------------------------------------
+
+.. currentmodule:: embodichain_models.dynamics.rigidformer.supervised
+
+.. autosummary::
+
+   RigidFormerObjective
+   rigidformer_batch
+
+.. autoclass:: RigidFormerObjective
+   :members:
+
+.. autofunction:: rigidformer_batch
+
+Recoverable EMA state
+---------------------
+
+.. currentmodule:: embodichain_models.dynamics.rigidformer.ema
+
+.. autosummary::
+
+   RigidFormerEMA
+
+.. autoclass:: RigidFormerEMA
+   :members:
+
+Shared SL training recipe
+-------------------------
+
+.. currentmodule:: embodichain_models.dynamics.rigidformer.sl_recipe
+
+.. autosummary::
+
+   build_rigidformer_trainer
+
+.. autofunction:: build_rigidformer_trainer
+
+Recovery and experiment boundaries
+----------------------------------
+
+The source is explicitly registered as a training-state participant so its
+epoch/cursor commits after optimizer success. Prepared files are content-checked;
+source identity also binds splits, scaling, sampling and seed. EMA shadows remain
+FP32 and retain the author's warmup policy. Official weight import is distinct
+from native optimizer recovery; export selects EMA or raw weights explicitly.
+
+Applications own acquisition and independent reference experiments. Short
+numerical comparisons do not establish published accuracy.

embodichain/learning/sl/__init__.py

@@ -21,7 +21,12 @@
 from importlib import import_module
 from typing import TYPE_CHECKING, Any
 
-from .types import LossResult, SupervisedBatch, SupervisedObjective
+from .types import (
+    LossResult,
+    SupervisedBatch,
+    SupervisedObjective,
+    SupervisedTrainingState,
+)
 from .distributed import distributed_run
 from .learner import SupervisedLearner
 from .sources import SupervisedSource, TensorBatchSource
@@ -34,6 +39,7 @@
     "LossResult",
     "SupervisedBatch",
     "SupervisedObjective",
+    "SupervisedTrainingState",
     "SupervisedLearner",
     "SupervisedSource",
     "TensorBatchSource",

embodichain/learning/sl/distributed.py

@@ -21,8 +21,10 @@
 from collections.abc import Callable, Iterator
 from contextlib import contextmanager
 from datetime import timedelta
+import io
 import math
 import os
+import pickle
 from typing import Any, TypeVar
 
 import torch
@@ -159,8 +161,53 @@ def agree(self, value: Any, stage: str) -> None:
     def gather(self, value: Any) -> list[Any] | None:
         if not self.enabled:
             return [value]
-        # Keep the small control-state path on all-gather collectives, which
-        # are supported by both NCCL and Gloo. Non-root copies are temporary.
-        values: list[Any] = [None] * self.world_size
-        dist.all_gather_object(values, value)
-        return values if self.rank == 0 else None
+
+        # Participant snapshots can contain a full model EMA. Broadcasting one
+        # bounded chunk at a time keeps NCCL's device allocation independent of
+        # both the snapshot size and world size; only rank zero retains copies.
+        def serialize() -> io.BytesIO | None:
+            if self.rank == 0:
+                return None
+            stream = io.BytesIO()
+            pickle.dump(value, stream, protocol=pickle.HIGHEST_PROTOCOL)
+            return stream
+
+        payload = self.run(serialize, "checkpoint state serialization")
+        sizes: list[int] = [0] * self.world_size
+        dist.all_gather_object(sizes, payload.tell() if payload is not None else 0)
+        chunk_size = min(max(sizes), 16 * 1024 * 1024)
+        buffer = self.run(
+            lambda: torch.empty(
+                chunk_size, dtype=torch.uint8, device=self.control_device
+            ),
+            "checkpoint transfer buffer",
+        )
+        values: list[Any] | None = [value] if self.rank == 0 else None
+        sender = (
+            torch.frombuffer(payload.getbuffer(), dtype=torch.uint8)
+            if payload is not None
+            else None
+        )
+        for source in range(1, self.world_size):
+            received = io.BytesIO() if self.rank == 0 else None
+            failure = None
+            for offset in range(0, sizes[source], chunk_size):
+                count = min(chunk_size, sizes[source] - offset)
+                if self.rank == source:
+                    assert sender is not None
+                    buffer[:count].copy_(sender[offset : offset + count])
+                dist.broadcast(buffer[:count], src=source)
+                if received is not None and failure is None:
+                    try:
+                        received.write(buffer[:count].cpu().numpy().tobytes())
+                    except Exception as error:
+                        failure = error
+            if received is not None and failure is None:
+                try:
+                    received.seek(0)
+                    assert values is not None
+                    values.append(pickle.load(received))
+                except Exception as error:
+                    failure = error
+            self.check(failure, "checkpoint state collection")
+        return values

embodichain/learning/sl/learner.py

@@ -63,6 +63,9 @@ class SupervisedLearner:
         numerical_context: Context factory spanning forward, backward and clipping.
         ddp_find_unused_parameters: Enable graph traversal for objectives that
             omit trainable parameters. Applies only with an initialized group.
+        ddp_broadcast_buffers: Broadcast rank-zero buffers before forwards.
+            Defaults to the existing rank-local policy; explicit reference
+            recipes can request the PyTorch DDP broadcast policy.
 
     Sampling, normalizer updates, learning-rate schedules and model publication
     belong to callers. An update is committed only after optimizer.step succeeds.
@@ -81,6 +84,7 @@ def __init__(
         ) = None,
         numerical_context: Callable[[], AbstractContextManager[Any]] | None = None,
         ddp_find_unused_parameters: bool = False,
+        ddp_broadcast_buffers: bool = False,
     ) -> None:
         if not math.isfinite(max_grad_norm) or max_grad_norm < 0:
             raise ValueError("max_grad_norm must be finite and nonnegative")
@@ -92,6 +96,7 @@ def __init__(
         self.numerical_context = numerical_context or nullcontext
         self.num_updates = 0
         self.ddp_find_unused_parameters = ddp_find_unused_parameters
+        self.ddp_broadcast_buffers = ddp_broadcast_buffers
         self._ddp: DistributedDataParallel | None = None
         self._distributed_metrics: tuple[str, ...] | None = None
 
@@ -103,7 +108,7 @@ def _configure_distributed(self) -> _DistributedContext:
                 device_ids=(
                     [context.device.index] if context.device.type == "cuda" else None
                 ),
-                broadcast_buffers=False,
+                broadcast_buffers=self.ddp_broadcast_buffers,
                 find_unused_parameters=self.ddp_find_unused_parameters,
             )
         return context