all lessons/ world_models/ 15 · training and scaling systemslesson 15 / 16

Training and scaling systems

The unit of data is a time-aligned trajectory, not a shuffled frame. Resolution, frame rate, context, horizon, views, candidates, refinements, and uncertainty samples all multiply into the real training and serving bill.

Where we are
Earlier lessons assumed clean synchronized trajectories and affordable imagination. Now we follow Countertop Assistant from raw sensor packets through alignment, tokenization, sampling, distributed training, rollout search, and a stateful service with a hard control deadline — and do the arithmetic that decides the architecture.
Forced by 14A dense, asynchronous, embodied trajectory. This stepAccount for tokens, compute, rollout throughput, and the serving deadline. Forces 16A trained system must be evaluated for what it actually guarantees.

1 · Specify the trajectory contract before the architecture

A static image example is basically (x,y). A world-model example is an interval of physical history — and that difference drives everything in this lesson. At minimum it needs timestamped observations for every sensor, commanded and executed actions, episode boundaries, reward or progress, safety costs, validity masks, calibration, and embodiment metadata. That record is the evidence from which action-conditioned dynamics are identified, so if the record is wrong, a sophisticated model just learns a sophisticated version of the wrong causal process.

Continue the running product. Countertop Assistant must put the red ceramic mug on a tray within 20 seconds. A raw episode holds two RGB cameras, depth, 200 Hz joints, 1 kHz tactile, 50 Hz commanded actions, controller-reported executed actions, the language instruction, annotations when available, force-limit events, and outcome. The dataset contract should state, for every field:

Store raw immutable packets and build derived aligned views through versioned transformations, so a later discovery of a 22 ms actuator delay creates a new dataset version rather than an irreversible relabeling. And split train/eval by environment, object instance, time period, and collection policy where leakage matters — random clip splitting can drop adjacent moments of one episode into both sets and manufacture an illusion of generalization.

2 · Build the data factory around causality

The ingestion path should be linear and auditable:

raw packets → clock correction → calibration/frame transforms → action execution alignment → episode segmentation → validity masks → derived state/labels → chunk index → immutable manifest → train/eval sampling views

Clock correction maps device clocks to one monotonic timeline and keeps uncertainty; calibration puts points and poses into named frames; action alignment separates what was requested from what the controller applied; segmentation forbids prediction across resets, cuts, teleports, or takeover; validity masks tell the learner a sensor is missing instead of reading zeros as measurements. The final manifest should make every training sequence reproducible from content hashes and transformation versions.

Automated tests earn their keep here: reproject depth into RGB and track pixel error, integrate joint velocity against encoder displacement, verify force changes follow (not precede) contact actions, check conservation-like quantities and workspace bounds, and plot timestamp gaps, saturation, exposure, and missingness by robot revision. Then actually watch a stratified sample of synchronized episodes — humans catch reversed events and coordinate mistakes summary statistics miss. And note a bias specific to this product: demonstrations approach visible, easy mugs and abort dangerous grasps, so passive data never identifies what happens under unsupported actions. The collection plan needs safe interventions — small approach perturbations, varied grip forces, deliberate viewpoint changes — plus randomized simulator data, and it should record the behavior-policy probability so evaluation can separate world frequency from collector preference.

3 · Define the training example without leaking the future

Choose a context interval [τ−C, τ], an executed action sequence starting at τ, and targets through τ+H. An encoder infers belief from the past, dynamics predicts future latent state under actions, and heads predict observations, task variables, rewards, costs, and uncertainty. A common objective:

L = λobs Lobs + λdyn Llatent + λrew Lreward + λcost Lsafety + λstate Ltask-state + λKL Lregularization + λcal Lcalibration

The coefficients encode product priorities and gradient scale, so don’t guess them once and forget: normalize by valid target count, log per-head gradient norms, and check whether a big reconstruction term is quietly suppressing rare collision learning. If pixel prediction can be solved by texture without dynamics, add action-effect, inverse-dynamics, temporal-consistency, permanence, and task-state objectives; if predicting all pixels wastes capacity, decode pixels on a subset while training structured heads every step. Loss frequency matters as much as its coefficient — a collision head that sees one rare label per batch gets noisy gradients even with a huge weight, so build event-rich microbatches, accumulate rare-head gradients, or use a separate replay stream while keeping ordinary examples to control false alarms. Watch the Pareto curves: lowering reconstruction loss can worsen contact calibration, and maximizing reward prediction can discard geometry a new goal needs — a shared representation is justified only when transfer beats this interference. And the sample must mimic deployment causality: bidirectional attention inside the past is fine, but future frames, future normalization statistics, completed-episode labels, and future-smoothed poses must never enter the online state, so train with realistic dropout and delay and gradually test free rollouts, because a model trained only on true states never meets its own errors.

4 · Count tokens from the physical sampling plan

Here is where the trajectory’s cost becomes concrete. Let view v have resolution Hᵥ×Wᵥ, patch pᵥ, frame rate fᵥ, and context C. With one token per patch:

Nvisual = Σᵥ ceil(Hᵥ/pᵥ) ceil(Wᵥ/pᵥ) ceil(fᵥ C) Ntotal = Nvisual + Ndepth + Nproprio + Ntactile + Naction + Nlanguage + Nspecial

And that is only the input count. Autoregressive future tokens add horizon; dense self-attention materializes interactions proportional to per layer, while MLP and projection scale roughly linearly in tokens but quadratically in width. Space–time factorization, latent bottlenecks, recurrent state, sparse memory, and multirate tokenization all change these terms — they are architectural responses to a measurement plan, not free choices.

Trajectory token accountant
Raise resolution, frame rate, and horizon. Dense attention pairs grow quadratically, motivating compression, factorization, hierarchy, and sparse memory.
Interactive compute accountant.
trajectory tokens
attention pairs
pressure
response

Work the token and storage bill

Take an 8-second context. Two 256×256 RGB views at 12 Hz, patch 16, give 2 × 96 × 16 × 16 = 49,152 raw visual tokens. A depth stream compressed to 64 tokens/frame adds 96 × 64 = 6,144. Rather than tokenize all 200 Hz joints, a multirate encoder summarizes ten samples into 16 tokens at 20 Hz, adding 8 × 20 × 16 = 2,560. Tactile windows become 32 tokens at 20 Hz (5,120), actions add 8 per 20 Hz step (1,280), language and special state use 256 — 64,512 tokens of context.

Now the sticker shock: dense attention over that sequence would expose 64,512² ≈ 4.16 billion token pairs per layer. At 24 layers and just two bytes per stored attention value, the naive attention maps alone are roughly 200 GB for one sample before gradients and optimizer state. That single number forces the architecture: encode each frame spatially, pool to maybe 64 scene tokens, keep object/contact memory, and run temporal attention over compressed state — 96 timesteps × 128 fused tokens is 12,288 temporal tokens, and windowing or recurrent belief shrinks it further. Raw storage tells a different story: two RGB streams at ~1.5 MB/s each are 24 MB per 8 s, depth at 0.8 MB/s is 6.4 MB, and numerics/actions/audio/labels add ~4 MB — about 34.4 MB of raw episode per example, so ten million non-overlapping examples would be 344 TB before replication. In practice contexts overlap, so store each episode once and index windows rather than duplicating bytes.

5 · Count training compute, not just dataset size

For a dense Transformer, a first-order training estimate is 6PT FLOPs, with P active parameters and T training tokens — it hides attention, sparsity, multimodal encoders, and recomputation, but it catches order-of-magnitude blunders. Say compression yields 12,000 model tokens per sample, the usable corpus is 8 million windows per epoch, and training runs 2.5 effective epochs: T = 12,000 × 8,000,000 × 2.5 = 240 billion tokens. A 1.2B-active-parameter model then costs:

Ftrain ≈ 6 × 1.2×10⁹ × 240×10⁹ = 1.728×10²¹ FLOPs

At an optimistic sustained 35% of 800 accelerators rated 300 TFLOP/s, throughput is 0.35 × 800 × 300×10¹² = 8.4×10¹⁶ FLOP/s, giving an idealized ~20,571 s ≈ 5.7 hours. That number is suspiciously low on purpose — the formula omits vision encoders, attention terms, evaluation, data stalls, checkpointing, failed runs, and hyperparameter sweeps. Stating assumptions is the lesson; a realistic plan multiplies the core run by measured model FLOPs and an experimentation factor of often 5–20×. Also count optimizer memory: mixed-precision parameters, gradients, and Adam states run ~12–20 bytes/parameter, so 1.2B needs 14–24 GB before activations — and video activations often dominate. Activation checkpointing trades compute for memory; fully sharded data parallelism divides model state; tensor parallelism splits big matmuls; sequence/context parallelism splits long token axes; pipeline parallelism divides layers — and the best mixture follows the measured bottleneck, not fashion. Finally, feed the accelerators: 800 workers at 40 MB/s of decoded examples is 32 GB/s, so shard episodes for sequential reads, cache deterministic encodings only when the tokenizer version is fixed, prefetch asynchronously, and measure idle time — a cache that saves encoder work also creates a versioned derived dataset that must be invalidated, not silently reused, when augmentation changes.

6 · Sampling decides which physics gets gradient

Uniform time sampling overweights stationary views and routine success. Contact onset may fill 100 ms of a 20-second episode — 0.5% of the time — yet decide success. So mix strata: common traffic for calibration, event-centered windows for transitions, failure/recovery windows, long contexts for memory, high-action windows for system identification, and uncertain/disagreement cases from deployed models. And introduce difficulty in causal order. The tempting alternative — max everything (longest horizon, highest resolution, widest action set, all sensors, perfectly balanced environments) from the first update — produces a hard loss that never reveals which prerequisite failed. A linear curriculum changes one difficulty source only after the previous one carries signal:

  1. Make local change work before demanding memory. Start the horizon short and observation-corrected; scale to long free rollout only once executed actions reliably produce the right immediate change — otherwise a bad long rollout can’t tell you whether the defect began at step one or step twenty.
  2. Learn structure before spending capacity on texture. Start coarse, then add the fine contact and identity cues that change decisions; starting at max resolution lets texture eat capacity while geometry stays weak. Diagnostic: cut appearance detail — if task-state prediction collapses, the model was leaning on texture.
  3. Identify supported actions before extrapolating. Start inside behavior support where the data has evidence, then add safe interventions and edge cases; pushing into unsupported actions early is not just unstable, on hardware it can be unsafe to collect.
  4. Establish a synchronized sensory core before teaching degradation. Start with trusted-clock modalities, then introduce noise, correlated dropout, and missingness; add noise before the core is identifiable and the error source becomes impossible to localize.
  5. Learn ordinary frequency before stressing the tail. Start with common environments, then extend to the long tail and controlled shifts; balanced training builds capability, but treating that balance as real-world frequency destroys probability calibration.

The implementation isn’t five separate curricula but one measurable promotion rule: advance an axis when its prerequisite passes a held-out diagnostic — local transition error before horizon, structured state before resolution, intervention curves before unsupported planning, sensor-subset calibration before dropout deployment, common-regime calibration before tail reweighting — and when a metric later regresses, roll back along that dependency chain rather than adding more mixed data. Concretely: if low-friction lifts are oversampled 20×, train on them but calibrate probabilities on the true deployment mix (or apply importance weights), and keep two dashboards — conditional capability by slice, and expected performance under production frequency — while making sure hard-example mining doesn’t spiral into letting inherently stochastic cases consume all capacity.

7 · Train state, dynamics, and uncertainty in stages — but test jointly

A practical curriculum pretrains modality encoders, learns short action-conditioned transitions, extends free-rollout horizon, adds task and safety heads, then fine-tunes with planner-generated hard cases. Staging eases optimization, but frozen interfaces can lock in the wrong sufficient statistic, so joint fine-tuning must verify that a state good for reconstruction is also good for action effects and risk. Scheduled sampling replaces some true previous states with predictions to expose compounding error; overshooting predicts several future latents from one anchor to encourage longer consistency; contrastive and permanence losses preserve identity; distributional heads and ensembles carry uncertainty — and calibrate after the full pipeline, because planner selection shifts the action distribution and usually reveals overconfidence. Track error versus horizon rather than one average, separate one-step posterior reconstruction from prior free rollout, and report losses per modality, action magnitude, contact mode, scene, robot, and policy — a falling total loss can hide rising collision false negatives. Evaluate old checkpoints on a frozen golden set and new data on a freshness set to tell regression apart from distribution drift.

8 · Planning creates a second, multiplicative compute bill

Training throughput and imagination throughput are different animals. Planning cost depends on transition evaluations, decoding, candidate generation, and synchronization. For sampling-based MPC:

transition calls per decision = K candidates × H latent steps × I refinements × E ensemble members × S stochastic samples decoded outputs = selected candidates/steps only, if latent costs are available end-to-end latency = encode/correct + propose + rollout + score + select + safety margin

Work the rollout bill

At 10 Hz the robot has a 100 ms decision period. Reserve 20 ms for sensor alignment/state correction, 15 ms for motion controller and jitter, 10 ms safety margin — leaving 55 ms to plan. CEM-MPC evaluates K=256 candidate chunks for H=20 latent steps, refines I=4 times, uses E=3 ensemble members and S=2 stochastic futures:

256 × 20 × 4 × 3 × 2 = 122,880 latent transition evaluations per decision at 10 decisions/s = 1,228,800 evaluations/s per robot

So “the model takes 4 ms” is meaningless without batch shape and this multiplier. If an accelerator runs a batch of 256 one-step transitions in 0.45 ms, sequential horizon and four refinements already cost 0.45 × 20 × 4 × 3 × 2 = 216 ms — missing the deadline before overhead. The fixes are all reductions of that product: fewer candidates or ensemble samples, multi-step chunks, a distilled value model, parallel ensemble members, warm-starting from the previous plan, shared deterministic prefixes, or calling the full planner only at uncertainty/event boundaries. A revised config — 128 candidates, 10 two-step chunks, 2 refinements, 2 ensemble members — is 5,120 calls, and at 0.45 ms/step its serial depth of 10 × 2 × 2 = 40 batches is 18 ms, leaving time to score. Decode structured risk every step but render RGB only for the last few candidates, and cache the belief and static scene — never re-encode eight seconds of video per branch.

Hidden multiplier
Quote context encoding + horizon × candidates × refinement × ensembles × stochastic samples × decisions per second. Add denoising steps for diffusion and tree nodes/expansions for search. P50 single-forward latency is not a capacity plan.

9 · Serving is a stateful streaming control system

Each robot owns a belief instance keyed by model, calibration, and embodiment version. The service ingests out-of-order packets, advances state causally, branches candidate futures, returns action scores before the deadline, and records uncertainty/fallback events — while handling duplicates, packet loss, reconnect, reset, stale-state eviction, and model migration. A stateless request carrying only “the latest frame” throws away exactly the memory a world model exists to provide. Specify latency as a distribution and a deadline-miss policy: measure P50, P95, P99, and worst bursts under realistic concurrency, and if the plan is late, the robot continues a verified safe controller or stops — it does not execute a stale action. Apply backpressure by dropping redundant visualization frames before force or action events, and let admission control preserve per-agent cadence rather than letting one long rollout starve others. Deployment optimizations — caching encodings, batching compatible rollout stages, quantizing, distilling, compiling fixed shapes, separating slow semantic from fast belief updates, placing compute near the robot — each help, but every approximation must be re-checked for calibration and rare risk: quantization that barely changes average error can still erase a small collision logit.

10 · Close the data flywheel without training on your own fiction

Log the physical observation after every planned intervention, not just the imagined trajectory, and join plan, uncertainty, executed action, controller overrides, outcome, and model version. Prioritize cases with external model error, ensemble disagreement, fallback, near-miss, novel state, or human correction. A planner can propose useful candidate actions, but its predicted futures are not ground truth — replay them in a trusted simulator or a controlled real setting before labeling consequences, or you train on your own fiction. Version data, code, architecture, loss, tokenizer, calibration, and evaluation together, and keep a model registry recording the deployment envelope: supported robots, sensor set, context, trusted horizon, latency hardware, calibration thresholds, and known unsafe slices. Promotion should require offline gates, shadow-mode comparison, a bounded canary, and a rollback path — systems discipline is part of model quality, because irreproducible data and stale beliefs produce behavioral failures.

11 · A first-principles scaling review

  1. Trajectory: Are observation and executed action aligned on physical time, with cuts and invalidity explicit?
  2. Information: Which hidden variables need raw resolution, and which can use objects, latents, or multirate summaries?
  3. Tokens: Write tokens per modality × rate × duration × views, then future and batch counts.
  4. Training: Compute corpus tokens, model FLOPs, optimizer/activation memory, experimental multiplier, and data bandwidth.
  5. Sampling: Preserve production frequencies for calibration while ensuring rare causal transitions receive gradient.
  6. Rollouts: Multiply horizon, candidates, refinements, ensembles, stochastic samples, and decision frequency.
  7. Serving: Budget state correction, planning, control, jitter, safety margin, concurrency, and fallback.
  8. Flywheel: Collect external outcomes at the model’s error frontier and keep a frozen, leakage-resistant evaluation suite.

Where this points next

A data and compute system can efficiently train the wrong abstraction, produce calibrated pixels but incorrect action effects, or let a planner exploit tiny errors. Lesson 16 evaluates Countertop Assistant from observation alignment through hidden state, causal interventions, uncertainty, closed-loop utility, OOD behavior, system deadlines, and product safety — then turns the complete series into a reusable design-interview framework.

Takeaway
Preserve the trajectory before scaling the model. Compute every modality token, every training pass, and every imagined branch. Sample rare dynamics without forgetting deployment frequencies. Architect training parallelism around actual memory and compute, and architect serving around persistent belief and a deadline-miss policy. Scaling is the discipline of making the whole causal loop affordable and reproducible.

Interview prompts