Supervised Learning

Supervised Learning#

Use embodichain.learning.sl to train a model from explicit inputs and supervision, with a shared optimizer-update loop, validation, checkpoints and single-node distributed training. The same infrastructure supports imitation policies and learned dynamics. It does not require a simulator or an RL environment.

A model recipe defines the network, data encoding and loss. The SL layer runs that recipe without assuming an observation layout, action representation or sequence length. Shared networks live in learning.models; concrete models and their recipes live in the bundled embodichain_models extension. See Integrated models to choose a model and find its data preparation and training workflow.

Components and data flow#

Component

Responsibility

SupervisedBatch

Structured inputs, targets, named masks and host-side metadata

SupervisedSource

Sample an owned batch with an explicit CPU sampling generator and a stable source identity

SupervisedObjective

Compute a differentiable scalar loss and return a LossResult with metrics and effective supervision weight

SupervisedLearner

Run forward/backward, check and clip gradients, and commit one optimizer update

SupervisedTrainer

Schedule sampling, updates, validation, metrics and checkpoint recovery

SupervisedTrainingState

Integrate additional committed state, such as a source cursor or model average, into updates and recovery

 
        flowchart LR
    Config["Configuration"] --> Factory["Recipe factory"]
    Factory --> Source["Source"]
    Factory --> Trainer["Trainer"]
    Source --> Batch["Batch: inputs, targets, masks"]
    Batch --> Learner["Learner"]
    Objective["Model + Objective + Optimizer"] --> Learner
    Trainer --> Source
    Trainer --> Learner
    Trainer --> Validation["Validation callback"]
    Trainer --> Artifacts["Metrics and checkpoints"]

For each update, the trainer samples a batch and moves its tensor trees to the configured device. The learner computes the loss, backpropagates, checks/clips gradients and steps the optimizer. After success, the update count and sampling RNG advance; registered training states update in order, followed by the learning rate scheduler. Validation, metric logging and periodic saving then run when due.

The objective owns masks, normalization and multi-step integration. SupervisedBatch.to(device) preserves gradient paths; the trainer does not detach rollout histories. Nonfinite losses or gradients prevent an optimizer commit. The whole iteration is not a rollback transaction: an error after the optimizer step does not undo that step or refund data already delivered.

See the SL API for complete signatures and configuration fields.

Integrated models#

Three model families provide installable factories for the shared trainer. They ship in the main embodichain wheel; datasets, pretrained weights and optional runtime dependencies are prepared through each model’s workflow.

Model

Use case

Prediction inputs

Prediction output

NeRD

Robot dynamics with history

States, joint efforts and optional contact history

State changes decoded into the next robot state

RWM

Action-conditioned robot dynamics

Current state, action and optional physical context

Next robot state at a fixed control period

RigidFormer

Interacting rigid-object dynamics

Geometry history, reference geometry and physical properties

Future rigid-object geometry

ACT: visual action-chunk policy#

Action Chunking with Transformers (ACT) learns actions from demonstrations. The integrated policy combines a ResNet-18 image encoder, a Transformer and a conditional VAE. Training uses demonstrated action chunks with a masked L1 loss and a KL term; inference predicts an action chunk from current RGB and robot state without requiring future observations.

The ACT example uses LeRobot 0.6.1’s official dataset loader, processors and training workflow on Python 3.12. EmbodiChain collects demonstrations and performs checkpoint selection and GPU simulator evaluation. Training uses a training-only dataset; a separately recorded validation corpus and independent development/test scenes remain outside training statistics. Native policy and processor artifacts replace the historical custom ACT checkpoint format. See the ACT workflow and visual ACT API for adapters and requirements.

NeRD: history-based robot dynamics#

Neural Robot Dynamics (NeRD) uses a causal Transformer to predict state changes from histories of robot states, generalized efforts and configured contact observations. A physical codec defines the feature representation and decodes predictions into native robot states. Training minimizes normalized state-change error using statistics computed from training windows.

The shipped reference systems are Cartpole and Ant. Their committed trajectory shards, codec and physical metadata must match the model; the network runs on CPU or CUDA without starting a simulator. Start from the Cartpole recipe or Ant recipe after preparing data and updating paths as described in the NeRD guide.

RWM: action-conditioned robot state prediction#

Robot World Model (RWM) uses an MLP to predict frame-local state increments and integrates them into the next robot state. Its supervised objective rolls out multiple steps autoregressively and compares them with measured targets while retaining state/action gradients. Use it when measured state/action transitions and a fixed control period define the prediction problem.

The recipe reads a tensor corpus with initial states, actions, future state targets and validity masks. Action semantics, control period and ordered joints are explicit metadata. The CPU/CUDA example generates controlled fitting data and an SL recipe.json; these checks qualify fitting and interfaces rather than robot policy performance. See the RWM guide and RWM example for corpus preparation and training.

RigidFormer: rigid geometry prediction#

RigidFormer combines a geometry encoder and attention to predict anchor motion, then applies rigid projection to produce the next object geometry. Its current adapter predicts motion from observed geometry and physical properties without an action input. Training uses position and acceleration losses on anchors before and after rigid projection, with model-specific augmentation and an exponential moving average (EMA) of model parameters.

The MOVi-Sphere, MOVi-A and MOVi-B recipes consume prepared, aligned material-point sequences, reference geometry, anchors and physical properties. Geometry prediction requires CUDA and the embodichain[rigidformer] dependencies. The public rollout() runs without autograd; differentiable action-based control is not part of this adapter. Start from the MOVi-A recipe after following the RigidFormer guide for preparation, runtime requirements and evaluation.

Run a first recipe#

From a repository checkout with EmbodiChain installed, run the existing masked regression example on CPU. It generates a fixed synthetic corpus locally and trains a small MLP; no simulator or dataset download is needed.

python -m embodichain train-sl --config examples/learning/sl_ddp/cpu.yaml

Each fresh run needs an artifact-free trainer.output_dir. The supplied recipe uses outputs/sl_ddp_cpu, 200 updates and a batch size of 32. It validates every 50 updates and saves every 100 updates. Copy the configuration and choose a new output directory before starting another independent run.

The configuration selects its factory with:

factory: examples.learning.sl_ddp.recipe:build_trainer

The factory receives the full configuration and returns a SupervisedTrainer. It interprets model, data, objective, optimizer and evaluation settings, then assembles the source, model, objective, optimizer and learner. These sections are recipe-specific; for example, optimizer field names need not match across model families. The trainer section uses SupervisedTrainerCfg.

The example writes run.json, metrics.jsonl, step_100.pt, step_200.pt and last.pt. Its validation MSE is logged as validation/mse. The generic trainer does not automatically select a best checkpoint; a recipe’s validation callback can implement selection using development data, never the test split.

For a local copy of the example, see the CPU configuration, factory source and launch and recovery notes. The example factory is a checkout example, not an installed model-library entry point.

Choose a data source#

Source

Use it for

Recovery considerations

TensorBatchSource

A fixed CPU tensor-tree corpus; samples with replacement

Takes an owned snapshot and binds content and metadata to source identity

SequenceBatchSource

A data-pipeline BatchSource plus an explicit, versioned encoder

Exact resume is opt-in for immutable, unbudgeted resident tensor sources

Model-specific source

Model-owned image decoding, state codecs or window selection

The recipe must provide the required identity and recoverable state

An encoder selects causal inputs, supervision and objective masks explicitly. The sequence bridge preserves canonical masks and provenance under metadata['sequence_source']; it does not infer action meaning or automatically turn future observations into inputs. Units, coordinate frames, action semantics and control periods belong in the data declarations and recipe signature.

Live buffers and delivery quotas are not restored by the common trainer. SequenceBatchSource defaults to rejecting exact resume; immutable=True requires an unchanged resident TensorDataSource (including resident LeRobotDataSource) without delivery quotas. Such sources verify content again before recovery. A live source may train and save a snapshot without supporting exact continuation.

Data production, refill, waits and error recovery remain application-owned. The CLI closes the returned source when it provides close(); programmatic callers manage that consumption lifetime themselves. See Online Data Streaming and the sequence bridge example notes for the data-supply side.

Single-node distributed training#

Run the same CPU recipe with two workers using Gloo, after choosing a fresh output directory in the configuration:

OMP_NUM_THREADS=1 torchrun --standalone --nnodes=1 --nproc-per-node=2 \
  -m embodichain train-sl --config examples/learning/sl_ddp/cpu.yaml

For two GPUs, use the CUDA recipe; each worker uses one GPU through NCCL:

CUDA_VISIBLE_DEVICES=0,1 OMP_NUM_THREADS=1 \
  torchrun --standalone --nnodes=1 --nproc-per-node=2 \
  -m embodichain train-sl --config examples/learning/sl_ddp/cuda.yaml

The CLI initializes the process group before the factory and resolves trainer.device to cpu or cuda:LOCAL_RANK. Programmatic callers can use distributed_run(device) for the same scoped setup. It cleans up only the group it creates.

batch_size is per rank; two workers with 32 examples each draw 64 examples per update. The learning rate is unchanged. Sampling remains source-owned: tensor sources draw independently with replacement, while traversal-based sources may shard windows. Account for each worker’s source/cache memory.

Under DDP, each objective returns a local mean loss and a finite, positive LossResult.weight, representing its effective valid supervision count. The learner forms sum(weight_rank * loss_rank) / sum(weight_rank) and scales the backward pass accordingly. Equally sized batches can have different weights when their masks differ. Objective metrics use the same weighting and must have consistent names across workers and updates.

All ranks participate in ordinary fit(), save() and resume() calls. Only rank 0 writes files and validates with the unwrapped model. A save inside the trainer’s rank-zero validation callback uses the already captured training boundary, without another collective call.

Checkpoints and exact continuation#

A checkpoint stores model and optimizer state, update count, scheduler state, sampling and global RNG states, recipe/source identity and registered training states. Distributed checkpoints also retain each rank’s RNG, sampler and model buffers. Validation preserves training mode, buffers, RNG, sampling state and registered training states; callbacks must not mutate model parameters or other unprotected training resources.

Exact continuation requires the same source, recipe, training configuration and world size. total_updates is an absolute budget, including restored work; changing it is not supported by exact resume. The output directory may change.

For example, after the single-CPU run above, copy cpu.yaml to /tmp/sl-resume.yaml and change only trainer.output_dir to a fresh path such as outputs/sl_cpu_resumed. Keep the 200-update budget and all other settings:

python -m embodichain train-sl --config /tmp/sl-resume.yaml \
  --resume outputs/sl_ddp_cpu/step_100.pt

This continues from update 100 while preserving the original run’s later history. For a checkpoint created by torchrun, use the same worker count to resume. An interrupted run may reuse its original directory only if run.json matches and the complete metric history is not ahead of the checkpoint. Older checkpoints without a run identity require a fresh destination. Resume checks compatibility before restoring live training state; a saved snapshot alone does not establish that its data source supports exact recovery.

Contents