Appendix B — PyTorch Implementation Quick Reference

The big picture

A tiny story

Think of PyTorch as a very fast luggage system. It will batch, reshape, broadcast, and multiply whatever we hand it. It cannot know that two neighboring frames belong to different episodes, that one H means image height and another means planning horizon, or that a target was supposed to stay connected to a shared encoder.

So this appendix is a cockpit card, not a line-by-line replacement for the official implementation. Its question is simple: which invariant should still be true after every fast tensor move?

The technical backpack

Data and shape contracts

Dataset answers “what is sample 417?”; DataLoader schedules and batches those samples. The dataset or sampler must preserve episode membership, time order, action alignment, and window length. A loader may shuffle whole windows. It must not manufacture a transition by mixing frames inside one.

Storage is not scientific identity. HDF5, Lance, folders, or video can carry the same trajectory contract, while two HDF5 files can differ in preprocessing, splits, frame skip, or action alignment. Record both container and contract.

Before a Transformer call, use a verbal boarding pass:

batch = independent examples
time = ordered latent positions
feature = learned coordinates
visibility = this position may read only itself and earlier positions

Equal sizes do not prove equal roles. If candidate count and time length are both four, swapping them can leave the printed shape unchanged. Carry recognizable episode, time, and candidate markers through a disposable reshape test.

Masking and action conditioning

A causal mask blocks a direct look at later tokens. Change a later context token and verify that earlier outputs stay unchanged. Passing this proves a computational visibility rule, not causal physics.

AdaLN gives actions a route into a Transformer block through shift, scale, and gate modulation. In the frozen implementation, its final modulation generator starts neutral through zero initialization. That does not guarantee the trained model uses actions. Hold latent history fixed, swap consequential actions, and inspect the predictions.

Original LeWM v3 uses one shared trainable visual encoder. The shifted next-observation target is connected: there is no stop-gradient, EMA teacher, or pretrained frozen encoder. A reimplementation that detaches the target has the same forward shape and a different learning system.

SIGReg reference route

Read the frozen SIGReg implementation as five rooms:

collect projected observation embeddings
keep time positions separate; inspect each one across the batch
sample random one-dimensional directions
compare empirical fingerprints with the analytic standard-Gaussian fingerprint
average discrepancies into a differentiable training penalty

Flattening batch and time changes the population question. Gathering across devices changes the empirical population; the frozen LeWM class is explicitly single-GPU. Detaching embeddings changes gradients. Calling the result a p-value changes its role. Finite directions, minibatches, and integration points make this an approximation, not a Gaussianity certificate. Log the penalty beside coordinate spread, covariance spectrum, effective rank, and held-out neighbors.

Rollout and CEM route

Candidate plans can be stacked and evaluated side by side on a GPU. Time within each candidate remains autoregressive: step two consumes the state predicted after step one. Keep environment batch, candidate, model horizon, and latent width distinct.

sample sequences from the current action proposal
roll out candidates with the frozen world model
score terminal latents against the encoded goal
select the configured elite action sequences
refit only the action proposal, then repeat

Elite indices must still point to the action sequences that earned their scores. CEM changes the proposal, not encoder or predictor weights. Continuous bounded actions need declared clipping/update semantics; discrete actions need a categorical proposal. Vectorization saves execution time; it does not remove temporal dependence or model drift.

Precision, checkpoints, and reproducibility

Mixed precision changes how arithmetic executes, not the scientific definition of prediction loss or SIGReg. Watch for non-finite values and compare a short full-precision smoke run. This control need not match bit for bit; it should expose implausible scale or a fragile accumulation path.

A weight file supports inference. A training-resumption checkpoint should also identify optimizer, training position, composed configuration, and any evolving scheduler or gradient-scaler state. Preserve RNG state when continuity requires it, or say it is absent. Also record data identity, code revision, hardware/software, precision, windows, budgets, checkpoint rule, planner settings, goals, seeds, and evaluation episodes.

same data and preprocessing
same model, optimizer, and training budget
same planner and evaluation budget
same episodes and goal rule
declare every intentional difference before reading the score

The trick that fools us

Two labs use the same commit, configuration, seed, checkpoint, and GPU. Their losses match exactly. Then they discover that both samplers pair the last observation of one episode with the first action of the next. The bug is perfectly reproducible.

Repetition confirms a procedure; it does not certify meaning. Challenge episode boundaries, one-step shifts, causal visibility, action sensitivity, gradient recipients, candidate identity, and the separation between world-model training and planning. Conversely, two correct stochastic runs may differ numerically while supporting the same bounded conclusion.

Evidence receipt

PyTorch documentation and Hydra documentation define framework abstractions. LeWorldModel v3 and frozen official code 8edfeb336732b5f3ce7b8b210d0ba370a09e2cac define the inspected graph and interfaces. stable-worldmodel is evidence only for its later active platform, not retroactive evidence about original v3.

Allowed claim: these patterns preserve named data, shape, conditioning, rollout, optimization, precision, and recovery invariants under a recorded configuration. They do not show that a loader understands episodes, a causal mask discovers physics, AdaLN guarantees action use, mixed precision is harmless, or one seed proves reproducibility. Independent reproduction is not established here.

Quick check

1. Which layer must reject a window that crosses an episode boundary?
2. What does the later-token mutation test prove—and what does it leave open?
3. During CEM, which object changes and which model parts stay frozen?