X0.1.3

Lesson 0.1.3 — nn.Module, optim, and DataLoader: The Anatomy of a Training Pipeline

Week 2 · Prerequisites: 0.1.1, 0.1.2.

Prior results used. This lesson uses 0.1.1’s five-attribute tensor model (shape, dtype, device, requires_grad, strides), 0.1.1’s view-versus-copy distinction, and 0.1.2’s distinction between leaf and non-leaf tensors and its .item() / .detach() / torch.no_grad() toolkit; each is restated at the point of use.

Hands forward. This lesson hands forward one capability: the learner can assemble nn.Module, optim, and DataLoader into a training loop with the correct line ordering, and can diagnose the standard pipeline bugs — a missing zero_grad, an optimizer constructed too early, a training set left in its natural order — from the loss trajectory and the gradient readout alone — the capability that Lesson 0.1.4 (train/eval mode) and every learning lesson in Phase 1 assumes.

Running case. The same 6-DOF MuJoCo arm across 8 parallel environments from Lessons 0.1.1–0.1.2; the task in this lesson is learning inverse kinematics — a network that maps a target end-effector pose to a joint-angle command for that arm — a problem that is structurally identical to behavior cloning and that every example in this lesson instantiates. The lab runs the arm’s 3-link planar reduction, where the pose-to-joint map is injective; §L.6.10 explains why the full 6-DOF map cannot be learned by any network and what that failure looks like in a loss curve.

§L.0 — Lesson map, objectives, and misconceptions

§L.0.1 — Lesson map

§L.0.2 — Objectives

By the end of this lesson you will be able to do eight things:

1. State the four kinds of state an nn.Module holds (parameters, buffers, submodules, and plain attributes) and name which kind each of a given module’s attributes is.
2. Implement a custom nn.Module that registers its parameters and buffers correctly, and verify registration with model.parameters() and model.state_dict().
3. Construct an optimizer with parameter groups for a fine-tuning scenario (a low learning rate for a frozen backbone, a high learning rate for a new head), and state why the optimizer must be constructed after the architecture is final.
4. Build a Dataset and a DataLoader with a custom collate_fn for variable-length trajectories, and separate the two reproducibility knobs — generator= for batch order, worker_init_fn= for per-worker augmentation — stating which one num_workers actually changes.
5. Assemble the canonical training loop with the exact line ordering zero_grad → forward → loss → backward → step → scheduler.step, and explain the failure caused by each misordering.
6. Trace a full iteration of an IK training loop, writing down the shape, device, and gradient state of every tensor at every line.
7. Diagnose three pipeline bugs — a zero_grad in the wrong place, an optimizer constructed before the architecture was final, and a training set whose order biases each batch — from a loss trajectory and a gradient readout, without running the code.
8. Produce a checkpoint that contains the model, the optimizer, the scheduler, and the step counter, and state what each of the four is needed for on resume.

§L.0.3 — Misconception inventory

Each misconception is displaced in the section named, and each is tagged in the concept check that follows it.

Everything is self-contained. Where a prerequisite lesson is used, the result is restated at the point of use.

§L.0.5 — Retrieval warm-up

This is the third lesson of the curriculum. The prompt specifies that warm-ups draw from two to five lessons back, which here means from Lesson 0.1.1; Lesson 0.1.2 is only one step back and is still in working memory, so it is not the source of the retrieval. Each question probes a result from 0.1.1 that this lesson reuses.

Warm-up 1 (from 0.1.1, the view/copy distinction). A Dataset.__getitem__ returns self.observations[idx]. If self.observations is a (10000, 6) tensor, is the returned element a view or a copy? Answer: A view — indexing a tensor with a single integer returns a view of the corresponding row, sharing the buffer with the parent. How this is used: §L.3’s Dataset discussion names this as the reason a data-augmentation step inside __getitem__ can silently corrupt the source dataset, and §L.3.2’s contrastive near-miss (WE-L.3.d) is exactly this case.

Warm-up 2 (from 0.1.1, device). A model is created on CPU and then moved with model.to('cuda'). Which of the five tensor attributes changed on the parameters, and which stayed the same? Answer: The device attribute changed; the shape, dtype, requires_grad, and strides stayed the same (.to preserves leaf status for leaves). How this is used: §L.1’s state_dict() discussion relies on the fact that .to(device) is a move, not a re-registration — the parameters’ identities are preserved, so the optimizer’s references to them remain valid after the move, which is why a move after optimizer construction is safe and a move before construction is also safe, but the two are not interchangeable for the optimizer’s state.

Warm-up 3 (from 0.1.1, shape conventions). A DataLoader with batch_size=32 wraps a Dataset whose elements are (6,) observation vectors. What is the shape of each batch the loader yields? Answer: (32, 6) — the loader stacks the individual (6,) elements along a new leading axis. How this is used: §L.3’s collate_fn discussion is entirely about what happens when this stacking cannot be done by the default collate, because the elements have variable shape.

Warm-up 4 (from 0.1.1, .item() and the logging hazard). Why does a training loop append loss.item() rather than loss to a list of metric values? Answer: .item() extracts a Python scalar and drops the graph reference, so the list does not retain the graph; appending loss retains the graph and accumulates memory until the process is killed. How this is used: §L.4’s training loop uses .item() on the loss and the metric, and §L.5’s failure taxonomy includes “OOM after N iterations, N scales with the batch size” as a signature of forgetting it.

Part A — nn.Module: The Container for Learnable Computation

§L.1 — Objects: the four kinds of module state

An nn.Module is a container for four kinds of state, and the distinction between them is the entire content of this part.

1. Parameters — tensors that are learnable and registered. They are wrapped in nn.Parameter (a Tensor subclass), have requires_grad=True by default, and are stored in the module’s _parameters dictionary. They are returned by module.parameters(), moved by module.to(device), and included in module.state_dict(). They are the only tensors the optimizer should update.
2. Buffers — tensors that are persistent but not learnable. They are registered with module.register_buffer(name, tensor) and stored in _buffers. They are not returned by parameters(), are moved by .to(device), are included in state_dict(), and are not updated by the optimizer. Batch-normalization running statistics are buffers; so are position embeddings that are not learned, and the running means of any custom statistic a module maintains.
3. Submodules — other nn.Module instances assigned as attributes. They are registered in _modules and their parameters are recursively included in the parent’s parameters(). Composition is the mechanism: a VLA’s vision encoder, language model, and action head are submodules of the top-level model.
4. Plain attributes — anything else assigned as an attribute (self.cfg = config, self.dropout_rate = 0.1, self.some_tensor = torch.randn(6)). These are not registered, not moved by .to(device), not in state_dict(), and not seen by parameters(). A raw tensor in a plain attribute is a common source of silent device bugs: the model appears to move, but the tensor stays on CPU, and the first operation that combines it with a moved parameter raises a device mismatch — or, worse, silently runs on CPU.

The four kinds are distinguished by registration, and registration is what the framework’s machinery reads. The distinction is not visible in the attribute’s printed repr: self.weight = nn.Parameter(torch.ones(6)) and self.weight = torch.ones(6, requires_grad=True) look identical at the call site and produce identical values, but the first is registered and the second is not.

§L.1.1 — Result: registration is what makes a tensor findable; the tensor’s other properties are irrelevant to discovery

Claim. For an nn.Module m, the set of tensors returned by m.parameters() is exactly the set of nn.Parameter objects reachable through m’s _parameters dictionaries and _modules recursion. A tensor with requires_grad=True that is assigned as a plain attribute is not in this set, regardless of its numerical content or its gradient capability.

Proof sketch. parameters() is a recursive generator that yields from m._parameters.values() and then recurses into m._modules.values(). It does not inspect plain attributes (__dict__ entries that are not in _parameters or _modules), because the registration is the only signal the framework has that a tensor is a parameter of this module. A tensor created with torch.ones(…, requires_grad=True) and assigned to m.foo is in m.__dict__ but not in m._parameters, so parameters() skips it. The autograd system will still compute its gradient (it is a leaf with requires_grad=True), but the optimizer was constructed with m.parameters(), so it does not know about m.foo, and never updates it.

Corollary. The failure is silent: the tensor’s gradient is populated after .backward(), but no optimizer consumes it. A learner inspecting m.foo.grad sees a non-zero gradient — the tensor is learning, in the sense that autograd is computing its gradient — and concludes the parameter is being trained. It is not. The gradient is accumulating in a buffer no one reads.

Corollary. The parameters() set is what the optimizer copies by reference at construction. A parameter added to the module after the optimizer’s construction is not in the optimizer’s parameter list, because the list is fixed at construction (the optimizer stores the parameter objects, not a live reference to the module). This is §L.2’s subject, but it follows immediately from the registration model: registration determines what parameters() yields, and the optimizer’s knowledge is a snapshot of that yield.

§L.1.2 — Worked-example ladder

WE-L.1.a — Fully worked: a custom IK network, parameter by parameter

Input. The running case: a network that maps a 3-dimensional target end-effector position (x, y, z) to a 6-dimensional joint-angle command (q₁, …, q₆). A two-layer MLP with a hidden width of 64, and a running input mean that is a buffer (not a parameter) — the mean is computed during training and used to standardize the input, but it is not learnable.

Trace (module construction).

1. class IKNet(nn.Module): — the module is a subclass; its __init__ must call super().__init__() before any assignment, because the parent’s _parameters, _buffers, and _modules dictionaries are initialized there.
2. self.fc1 = nn.Linear(3, 64) — a submodule; it has its own weight and bias parameters, which are in self.fc1._parameters, and self.fc1 is in self._modules.
3. self.fc2 = nn.Linear(64, 6) — a second submodule, same registration.
4. self.register_buffer("input_mean", torch.zeros(3)) — a buffer, stored in self._buffers, not returned by parameters(), but included in state_dict() and moved by .to(device).
5. self.hidden_width = 64 — a plain attribute, not registered; moving the model does not change it (it is an int, so no device issue, but the same call site would be a bug if it were a tensor).

Trace (inspection).

6. list(model.parameters()) — yields [fc1.weight, fc1.bias, fc2.weight, fc2.bias], four tensors, each shape (64, 3), (64,), (6, 64), (6,) respectively.
7. model.state_dict().keys() — ['fc1.weight', 'fc1.bias', 'fc2.weight', 'fc2.bias', 'input_mean'], five entries: the four parameters plus the one buffer.
8. model.input_mean in model.parameters() — False (a buffer, not a parameter; the generator yields parameter objects, and in compares by identity, so a buffer is never in the set).
9. model.fc1.weight.is_leaf — True; the nn.Linear constructed weight as an nn.Parameter, which is a leaf.

Output and interpretation. The parameters() set contains four tensors; the state_dict() contains five; the difference is the buffer. The learner who expects the buffer to be trainable (it is not) or to be absent from the checkpoint (it is present) has mis-modeled the four kinds. The buffer is what makes the checkpoint complete: if the input mean were a plain attribute (self.input_mean = torch.zeros(3)), it would not be in state_dict(), and a model loaded from the checkpoint would have input_mean = zeros(3) rather than the trained mean — a silent correctness bug that manifests as a systematically worse prediction after reload, with no error.

Self-explanation 1. Before reading on: in step 4, why is input_mean a buffer and not a parameter? What would go wrong if it were a parameter?

Answer. A buffer because the mean is computed, not learned — it is estimated from the data and then held fixed for standardization, and the optimizer must not update it (updating it would mean the standardization changes every step, which destabilizes training). If it were a parameter, it would be in parameters(), the optimizer would update it, and the standardization would drift; the model could compensate by inflating the learned weights, and the combination of a learnable mean and a learnable weight in a linear layer is unidentifiable (only the offset matters), so the training would still work but the mean would be meaningless as a statistic and could not be exported for deployment.

Self-explanation 2. Before reading on: in step 9, is model.fc1.weight a leaf, and why does the answer matter for the optimizer?

Answer. Yes, it is a leaf: nn.Linear constructs weight as an nn.Parameter, which is a leaf by construction. It matters because the optimizer’s step() updates leaf parameters — the parameters it holds references to — and a non-leaf would not accumulate gradients in .grad (Lesson 0.1.2, §L.1) and would not be updated. The nn.Parameter wrapping is what guarantees the leaf status; a learner who constructs a “parameter” from an operation on other tensors has a non-leaf that will not train.

WE-L.1.b — Completion problem

Input. The same IKNet, with a Dropout layer added between fc1 and fc2, and a LayerNorm applied to the input before fc1.

Trace (final 25% as prompts).

1. self.norm = nn.LayerNorm(3) — a submodule; its weight and bias are parameters.
2. self.drop = nn.Dropout(0.1) — a submodule with no parameters and no buffers; it is registered for its mode flag (Lesson 0.1.4), but its _parameters is empty.
3. self.fc1 = nn.Linear(3, 64) — as before.
4. model.parameters() — ⟨your turn: how many tensors, and which ones?⟩
5. model.state_dict().keys() — ⟨your turn: how many entries, and which one is not a parameter?⟩

Answers.

1. Six tensors: norm.weight, norm.bias, fc1.weight, fc1.bias, fc2.weight, fc2.bias. The Dropout contributes none.
2. Seven entries: the six parameters plus input_mean. The Dropout contributes no entry; its state is the training flag, which is not in state_dict() (Lesson 0.1.4 explains why).

The lesson. parameters() and state_dict() are recursive through submodules — a LayerNorm’s weight is a parameter of the top-level model even though it belongs to a sub-submodule — and the count difference between the two tells you the number of buffers. The learner who counted five parameters (forgetting norm’s weight and bias) has not modeled the recursion; the learner who counted eight entries (including the dropout) has not modeled the fact that stateless modules contribute nothing to either collection.

WE-L.1.c — Faded: the running case with one parameter changed

Setup. The 6-DOF arm’s IK network is now an ensemble of 4 IKNets, each with its own parameters, wrapped in an nn.ModuleList.

Question. How many parameter tensors are in the ensemble’s parameters()? How many entries are in its state_dict(), and what are the key names?

You must produce the intermediate steps. The setup and the answer are given.

Answer. Four IKNet instances × 6 parameter tensors per instance = 24 parameters (assuming the version with norm, fc1, fc2). The state_dict() has 4 × 7 = 28 entries, because each instance contributes its buffer input_mean. The keys are '0.norm.weight', '0.norm.bias', '0.fc1.weight', '0.fc1.bias', '0.fc2.weight', '0.fc2.bias', '0.input_mean', and similarly for '1.' through '3.' — the nn.ModuleList’s registration uses the list index as the key prefix.

The lesson. nn.ModuleList is the registered list of submodules; a plain Python list of modules is not registered (it is a plain attribute), so an ensemble built as self.models = [IKNet() for _ in range(4)] would have zero parameters visible to the optimizer, because none of the four modules is in _modules. The ModuleList is the registration statement.

WE-L.1.d — Contrastive near-miss: a raw tensor assigned to the module

Correct solution. The learner wants a learnable input scaling. They write self.scale = nn.Parameter(torch.ones(3)), and the parameter is in parameters() and in state_dict(), and the optimizer updates it.

Near-miss. The learner writes self.scale = torch.ones(3, requires_grad=True). The tensor is a leaf with requires_grad=True; autograd computes its gradient; after .backward(), self.scale.grad is a (3,) tensor. But the tensor is not in parameters(), so the optimizer — constructed with model.parameters() — does not include it, and its gradient is never applied. The parameter does not train.

Why the near-miss produces a plausible but incorrect number. The model’s loss decreases — because the other parameters are training correctly — and the learner who inspects model.scale.grad sees a non-zero gradient, which reads as “the parameter is learning.” The scale’s value does not change across steps, and the learner who checks the value sees it is stuck at 1.0 — but the interpretation is “the optimizer has not moved it far yet” rather than “the optimizer does not know about it.” The failure is silent in the loss curve and visible only in the scale’s trajectory, if the learner thinks to log it.

The discrepancy, stated. nn.Parameter(torch.ones(3)) and torch.ones(3, requires_grad=True) differ in exactly one property — whether the module’s __setattr__ recognizes the assignment as a parameter registration — and that property is not visible in print(model.scale), model.scale.shape, or model.scale.requires_grad. The nn.Parameter subclass is what __setattr__ checks; a plain tensor is assigned to __dict__ and never to _parameters. The learner who inspects the object by its repr sees no difference; the learner who inspects list(model.parameters()) sees the missing entry.

Self-explanation 3. Before reading on: in the near-miss, why is self.scale.grad non-zero after .backward(), if the parameter is not in the optimizer?

Answer. Because .backward() populates .grad on every leaf that requires grad and is reachable from the loss, regardless of whether the optimizer knows about the leaf. The gradient is computed by autograd’s traversal, not by the optimizer; the optimizer only reads the .grad of the parameters it holds. A leaf outside the optimizer’s list has its .grad populated and never consumed, and the accumulation continues across steps (Lesson 0.1.2’s accumulation rule applies to it too), so the .grad grows without bound over training — a slow leak that is invisible until it is inspected.

§L.1.3 — Comparative table: the four kinds of module state

Design principle (P3.1). If it must move, checkpoint, or train, it must be registered. Before assigning anything to a module attribute, ask: does this need to move to the device (register as a parameter or buffer), does it need to be in the checkpoint (parameter or buffer), or does it need to be updated by the optimizer (parameter)? If none of the three, a plain attribute is correct. If any, the attribute’s assignment must go through nn.Parameter, register_buffer, or nn.ModuleList/nn.ModuleDict, not through a plain assignment.

§L.1.4 — Concept check

Concept check (displaces M1). A learner says: “I made my scale a requires_grad=True tensor, so it is a parameter and the optimizer will train it.” Name the missing registration step and the one call that would have shown the learner their mistake.

Answer. The missing step is wrapping the tensor in nn.Parameter — the requires_grad=True flag makes the tensor differentiable, but only the nn.Parameter wrapping is recognized by the module’s __setattr__ as a registration, so the tensor is in __dict__ and not in _parameters. The call that would have revealed the mistake is list(model.parameters()), which does not contain the tensor, or equivalently any(p is model.scale for p in model.parameters()), which returns False; the model.scale.grad check would not have revealed it, because autograd populates the gradient for any differentiable leaf regardless of registration.

§L.1.5 — Consolidation checkpoint 1

(a) You should now be able to: (i) classify any module attribute into one of the four kinds, and state where each kind appears in parameters() and state_dict(); (ii) implement a custom nn.Module with correctly registered parameters and buffers, and verify registration; (iii) diagnose a “parameter that does not train” as a registration bug from the fact that its .grad is non-zero but its value is unchanged.

(b) Transfer probe. A learner’s model uses a fixed positional encoding, a (max_len, d_model) tensor that is computed once from sine and cosine functions and never updated. They assign it as self.pos_enc = torch.sin(...). The model trains, but after saving and reloading, the model’s predictions are systematically worse by a small amount. What is the bug? Answer sketch: the positional encoding is a plain attribute, so it is not in state_dict() and is not restored on reload — the reloaded model has pos_enc = None (or a re-initialized torch.sin with whatever the constructor computes, if it is re-run), which differs from the trained model’s encoding only if the constructor’s computation is not deterministic or if the loaded model’s other parameters were trained against a different encoding than the one the constructor produces. The fix is self.register_buffer("pos_enc", torch.sin(...)), which puts the tensor in state_dict() and makes the checkpoint complete.

(c) Remediation rule. If you cannot do (b), return to §L.1.3 and re-read the buffer row of the table, then re-attempt (b) by asking which of the four kinds the positional encoding should be.

Part B — optim: The Parameter Update Machinery

§L.2 — Objects: the optimizer’s construction, state, and groups

An optimizer is an object that holds references to a set of parameter groups and, on each step(), updates the parameters according to their gradients. The four objects of this part are:

1. The parameter set — the collection of parameter objects the optimizer holds. It is fixed at construction time: optimizer = Adam(model.parameters(), lr=…) iterates model.parameters() once, at construction, and stores the parameter objects. Adding a parameter to the model afterward does not add it to the optimizer.
2. The hyperparameters — the per-group lr, weight_decay, betas, eps, and any other optimizer-specific settings. They are stored per group, and a parameter belongs to exactly one group.
3. The per-parameter state — for SGD with momentum, a velocity buffer per parameter; for Adam, first and second moment estimates per parameter. The state is stored in optimizer.state, keyed by the parameter object, and is not in the model’s state_dict().
4. The optimizer’s state_dict() — a serializable dictionary containing the parameter groups (their hyperparameters and the parameters’ indices) and the per-parameter state. It is not the same as the model’s state_dict(), and a checkpoint that omits it resumes training with zero state.

The distinction between (1) and (3) is the source of the two most common optimizer bugs:

• A parameter added after construction is not in (1) and is never updated, even if it requires grad. This is the “early construction” bug (M2 in the inventory).
• A parameter that is in (1) has its per-parameter state in (3), and if the checkpoint does not save (3), the resumed training restarts the state at zero. For SGD with momentum, the first few steps after resume take larger effective steps, because the velocity buffer starts at zero and has not yet accumulated the momentum the run had built up; for Adam, the effect is the opposite of the one usually assumed — bias correction divides by 1 − β₁ᵗ and 1 − β₂ᵗ, so from a zero state the first update is lr · sign(g) in each coordinate, exactly lr, and the run takes a step of the learning rate in the raw gradient’s direction regardless of that gradient’s magnitude. The transient is a direction transient, not a magnitude one: the moments need roughly 1/(1-β₂) = 1000 steps to reflect the run’s actual gradient statistics, and during that window the per-coordinate step sizes are the uniform lr rather than the curvature-scaled sizes the pre-checkpoint run had settled into.

§L.2.1 — Result: the optimizer’s parameter set is a snapshot; the hyperparameters are per-group; the per-parameter state is separate from the model’s state

Claim. Let O be an optimizer constructed at time t₀ with O = Adam(model.parameters(), lr=…). Let p be a parameter added to model at time t₁ > t₀. Then for any step of O after t₁, p is not in any of O’s parameter groups, is not updated by O, and its .grad (if populated) is not consumed.

Proof sketch. The construction O = Adam(model.parameters(), …) evaluates model.parameters() once, at construction, and stores the yielded parameter objects in O.param_groups[i]['params']. model.parameters() is a generator, not a live view — each call re-traverses the module tree — but O holds the result of the one call at construction, not a reference to the generator. A parameter added to model after t₀ is yielded by a subsequent call to model.parameters(), which O does not make. The parameter is therefore not in O’s groups, and O.step() iterates the stored groups, not the live model.

Corollary. The fix for the early-construction bug is not “re-add the parameter” but “construct the optimizer after the architecture is final” or “add the parameter explicitly to a new group.” The latter is O.add_param_group({'params': [p], 'lr': …}), which is the mechanism for staged unfreezing: construct the optimizer once, then add groups as parameters are unfrozen.

Corollary. The per-parameter state is keyed by the parameter object’s identity, not by name. If a checkpoint is loaded into a model whose parameters were re-created (a new nn.Linear with the same architecture), the parameter objects are different objects, and the optimizer’s state — if loaded naively — is keyed by the old objects. The optimizer’s state_dict() stores the parameters’ indices in the groups, not the objects, so the state is restored correctly as long as the parameter order is unchanged; a model whose architecture changed in a way that reorders parameters will mis-restore the state, and the loss will spike or diverge after reload.

§L.2.2 — Worked-example ladder

WE-L.2.a — Fully worked: early construction, and the parameters it misses

Input. The running case’s IKNet with a fc1 and an fc2. The learner constructs the optimizer, then adds a third layer fc3 and expects it to train.

Trace.

1. model = IKNet() — parameters [fc1.weight, fc1.bias, fc2.weight, fc2.bias] (six with norm).
2. optimizer = Adam(model.parameters(), lr=1e-3) — the optimizer’s param_groups[0]['params'] is the list of those four (or six) parameter objects, in parameters()’s order.
3. model.fc3 = nn.Linear(6, 6) — a new submodule, registered, with fc3.weight and fc3.bias in _parameters of the new layer and recursively of model.
4. model.parameters() now yields [fc1.weight, fc1.bias, fc2.weight, fc2.bias, fc3.weight, fc3.bias] — six (or eight) tensors.
5. optimizer.param_groups[0]['params'] still yields the original four (or six) — the optimizer did not see step 3.
6. Training step: loss.backward() populates .grad on all six (or eight); optimizer.step() updates the four (or six) in its groups and leaves fc3 unchanged.

Output and interpretation. fc3’s parameters have non-zero .grad and unchanged values. The training loss may decrease slightly (because the earlier layers are learning to work with a fixed random fc3), which masks the bug: the learner sees a decreasing loss and concludes the model is training, but the new layer is not. The fix is to construct the optimizer after the final architecture, or to optimizer.add_param_group({'params': model.fc3.parameters()}).

Self-explanation 4. Before reading on: in step 6, why does fc3’s .grad accumulate across steps instead of being consumed?

Answer. Because .backward() populates .grad on every differentiable leaf reachable from the loss, regardless of the optimizer’s parameter list (the same mechanism as WE-L.1.d’s near-miss). The optimizer’s step() reads .grad only for the parameters in its groups, so fc3’s .grad is written and never read, and it accumulates across steps. The .grad for a never-updated parameter grows linearly with the step count, and the eventual large values are a diagnostic signal — but only if the learner inspects fc3.weight.grad.norm().

WE-L.2.b — Completion problem

Input. A fine-tuning scenario: a pre-trained vision backbone (frozen at the start) with a new action head (trainable). The learner wants the head at lr=1e-3 and the backbone at lr=1e-5, with the backbone’s batch-norm statistics frozen.

Trace (final 25% as prompts).

1. backbone = load_pretrained(...) — its parameters have requires_grad=False.
2. head = nn.Linear(backbone.out_dim, 6) — trainable by default.
3. optimizer = Adam([{'params': head.parameters(), 'lr': 1e-3}, {'params': backbone.parameters(), 'lr': 1e-5}], weight_decay=1e-4) — ⟨your turn: why is this construction correct even though the backbone is frozen, and what does weight_decay=1e-4 apply to?⟩
4. Later, the learner unfreezes the backbone with for p in backbone.parameters(): p.requires_grad_(True) — ⟨your turn: does the optimizer now train the backbone? Why or why not?⟩
5. If the learner wants to unfreeze the backbone and have it train at 1e-5, what must they do?

Answers.

1. The construction is correct: the backbone is in a separate group with its own learning rate, and weight_decay=1e-4 is the default for groups that do not specify one — it applies to both groups here (a group could override it with its own 'weight_decay' key).
2. Yes, the optimizer trains the backbone after unfreezing, because the backbone’s parameters were already in the optimizer’s groups at construction. Unfreezing changes requires_grad (so .backward() populates .grad), and the optimizer’s step() was always going to read .grad for those parameters — it simply had nothing to read while they were frozen. The unfreeze is a gradient-side operation, and the optimizer’s parameter list is unchanged.
3. Nothing — the unfreeze is sufficient, because the backbone is already in the optimizer. The early-construction bug does not apply here, because the backbone was in the optimizer from the start (as a frozen group). The learner who adds a new head after construction, on the other hand, must add it to the optimizer.

The lesson. “Frozen” and “not in the optimizer” are different states: a frozen parameter is in the optimizer’s list but has no gradient to update; an early-construction parameter is not in the optimizer’s list and cannot be trained by unfreezing. The distinction is the source of most confusion about fine-tuning: unfreezing is a gradient operation, not a parameter-registration operation, and a parameter that was in the optimizer from the start will train as soon as its gradient is populated.

WE-L.2.c — Faded: the running case with one parameter changed

Setup. The running case’s IK training uses CosineAnnealingLR(optimizer, T_max=200) and the learner is training for 200 epochs. They have written scheduler.step() before optimizer.step().

Question. What is the learning rate at epoch 0, and what is it at epoch 199? What is it if the scheduler is stepped after optimizer.step()? What is the effect on the training trajectory?

You must produce the intermediate steps. The setup and the answer are given.

Answer. The schedule’s first values are 0.100000, 0.099994, 0.099975, 0.099944, … — lr₀ and then a fast initial drop, which is the shape of a cosine near its peak. With the scheduler before the optimizer, epoch 0’s update uses the second scheduled value, 0.099994, and epoch 199’s update uses schedule[200] = 0.0 exactly. With the scheduler after the optimizer, epoch 0’s update uses lr₀ = 0.100000 and epoch 199’s update uses schedule[199] = 6.17e-06. So the offset is a one-step lead, and the lead has two consequences rather than one: the before-version never applies lr₀ at all — the first value of the schedule is skipped — and it reaches the schedule’s zero one epoch early, so the last update of the 200-epoch run is exactly 0.0 and the parameter does not move. The effect is small and systematic: the training anneals slightly faster and the final loss is marginally higher, in a way that is hard to attribute without knowing the placement.

The lesson. The scheduler’s step() advances the schedule by one, and the placement determines which step’s learning rate is applied to the current update. The before placement applies schedule[step+1] to the update at step (using a value the schedule computed before this step ran); the after placement applies schedule[step]. The convention is after — the schedule’s step=0 value is the initial learning rate and is applied to the first update — and the learner who inverts it has a schedule that is one step ahead of their intent. PyTorch does not let this pass entirely silently: it emits a UserWarning — “Detected call of lr_scheduler.step() before optimizer.step()… Failure to do this will result in PyTorch skipping the first value of the learning rate schedule” — on the first such call. The warning names the mechanism, and a learner whose output is being swallowed by a log file will not see it.

WE-L.2.d — Contrastive near-miss: resuming from a model-only checkpoint

Correct solution. The learner saves {'model': model.state_dict(), 'optimizer': optimizer.state_dict(), 'scheduler': scheduler.state_dict(), 'step': global_step} and restores all four. The resumed training continues from the saved step with the same Adam moments and the same schedule position, and the loss trajectory is continuous.

Near-miss. The learner saves only model.state_dict() and reloads it. The training resumes; the model’s parameters are correct; but the optimizer’s per-parameter state (Adam’s first and second moments) is zero, and the scheduler’s step counter is zero. The first steps after resume take updates with the wrong step sizes: with both moments at zero, bias correction makes the first update exactly lr · sign(g) per coordinate — a uniform step of the learning rate, whatever the gradient’s magnitude. The pre-checkpoint run had settled into per-coordinate step sizes scaled by 1/√v, which is what Adam is for; the resumed run has lost that scaling and applies lr to every coordinate equally for roughly 1/(1-β₂) = 1000 steps while the moments re-estimate. The parameters therefore move in a direction and at a per-coordinate scale the trained model had not been using, and they can be displaced substantially — measured on a small regression, the first post-resume update moves the parameters by more than the last hundred pre-checkpoint updates combined. There is no loss spike: bias correction prevents the m/eps blow-up that a naive reading of the formula suggests, and the loss stays monotone. The damage is the displacement, not a spike, which is why the bug is easy to miss.

Why the near-miss produces a plausible but incorrect number. The resumed training does continue; the model’s parameters are correct; and because bias correction bounds the first step at lr, the loss does not blow up. The learner who sees a monotone loss and a slightly different trajectory concludes “the model is fine, there was a transient.” The transient is real, but its cause is the missing optimizer state, and its cost is that the parameters were moved off the trajectory they were on — in a fine-tuning scenario, this can undo hours of training in a few hundred steps, and in an RL scenario, it can destabilize a policy to the point of no recovery. The tell is not a spike but a step-size signature: the per-coordinate update magnitudes immediately after resume are uniform, and they stay uniform until the second moment re-estimates.

The discrepancy, stated. model.state_dict() and optimizer.state_dict() are different dictionaries with different contents: the first contains the parameters and buffers, the second contains the parameter groups and the per-parameter state. A checkpoint that saves only the first is a model snapshot, not a training snapshot; the training can be resumed from it, but not from the same point in the optimization trajectory. The distinction is invisible in the file size (the optimizer state is comparable in size to the model for Adam, so a “large” checkpoint might be either), and the symptom is a transient that the learner is likely to attribute to something else.

Self-explanation 5. Before reading on: in the near-miss, why is the loss spike larger for Adam than it would be for SGD with momentum?

Answer. Because Adam’s update is normalized by the second-moment estimate, and with v = 0 the normalization is 1 / (sqrt(0) + eps) = 1 / eps = 10⁸, which multiplies the first moment by an enormous factor. SGD with momentum has no second-moment normalization; the missing state is only the velocity buffer, whose absence reduces the momentum’s contribution to the first update but does not multiply it by 1 / eps. The spike is therefore Adam-specific and much larger for Adam than for SGD.

§L.2.3 — Comparative table: optimizer construction and state

Design principle (P3.2). Construct the optimizer once, after the architecture is final, with all parameter groups — frozen or not — already in it. The frozen groups are dormant (no gradient to read) and become active the moment they are unfrozen, without any change to the optimizer. A learner who constructs the optimizer before adding a layer, or after unfreezing a backbone, has taken a snapshot of a state that is about to change, and the snapshot is what the optimizer will hold for the rest of training.

§L.2.4 — Concept check

Concept check (displaces M2). A learner says: “I unfroze the backbone, so the optimizer is training it now.” For which of the two possible meanings of “the backbone was frozen” is the statement true, and for which is it false?

Answer. It is true if the backbone’s parameters were in the optimizer’s groups from the start (as a group with requires_grad=False) — unfreezing populates their .grad and the optimizer’s step() reads it, so the backbone trains. It is false if the backbone was added to the model after the optimizer’s construction — the backbone’s parameters are not in the optimizer’s list, and unfreezing them populates .grad that the optimizer never reads, so the backbone does not train. The two are distinguished by any(p is backbone.weight for group in optimizer.param_groups for p in group['params']).

§L.2.5 — Consolidation checkpoint 2

(a) You should now be able to: (i) state the four objects of an optimizer and where each is stored; (ii) diagnose the early-construction bug from a non-zero .grad and an unchanged parameter value; (iii) construct a multi-group optimizer for fine-tuning and state why the frozen group must be present from the start; (iv) state what each of the four checkpoint components is for, and what breaks if each is omitted.

(b) Transfer probe. A learner trains a policy with PPO, saves a model-only checkpoint at step 10,000, and resumes. The policy’s behavior changes markedly for the first ~100 steps after resume, then settles back. What is the explanation? Answer sketch: Adam’s first and second moments are re-initialized to zero. Bias correction means the first update is not a blow-up — it is exactly lr in each coordinate’s gradient direction — but the pre-checkpoint run had per-coordinate step sizes scaled by 1/√v, and the resumed run has lost that scaling. Every coordinate now moves at the same lr regardless of its gradient history, which is a different optimizer from the one that was training. The moments re-estimate over roughly 1/(1-β₂) = 1000 steps, and the behavior settles once they do. The fix is to save the optimizer state, which restores the moments and the per-coordinate scaling with them.

(c) Remediation rule. If you cannot do (b), return to §L.2.2 and redo WE-L.2.d, then re-attempt (b) by asking what Adam’s update formula is when v = 0.

Part C — DataLoader: Batched, Shuffled, Parallel Iteration

§L.3 — Objects: Dataset, DataLoader, and the four arguments that matter

A Dataset is an object with two methods: __len__ (returns the number of examples) and __getitem__(idx) (returns the example at index idx). The example can be anything — a tuple of tensors, a dict, a nested structure — and the Dataset is responsible only for random access to one example.

A DataLoader wraps a Dataset and yields batches. It iterates the dataset in an order determined by shuffle, groups the examples into batches of size batch_size, applies a collate_fn to the list of examples to produce a batch, and optionally loads in parallel with num_workers subprocesses.

The four arguments that matter for this lesson, with their default and their effect:

1. batch_size (default 1) — the number of examples per batch. Interacts with the learning rate: doubling the batch size roughly doubles the signal-to-noise ratio of the gradient, which often permits a larger learning rate. In embodied AI, batch size is often num_envs × rollout_length for on-policy RL, and constrained by GPU memory for VLA fine-tuning.
2. shuffle (default False) — whether to permute the dataset’s order at each epoch. Essential for supervised learning with temporally correlated data; not appropriate when the learner deliberately wants sequential order (recurrent policies over full trajectories).
3. num_workers (default 0) — the number of subprocesses. 0 means load in the main process; >0 means parallel loading (faster for large datasets). It changes which process runs __getitem__, and therefore which RNG __getitem__ draws on — not the batch order, which the main process decides before any worker is asked.
4. pin_memory (default False) — whether to allocate batches in pinned (page-locked) host memory, which enables faster asynchronous transfers to GPU. Combined with .to(device, non_blocking=True), it overlaps data movement with compute.

And one method:

• collate_fn (default torch.utils.data.default_collate) — the function that turns a list of examples into a batch. The default stacks tensors along a new leading axis, which requires all examples to have the same shape. For variable-shape data (variable-length trajectories, variable-number of detected objects, variable-length instructions), a custom collate_fn must pad, mask, or otherwise normalize the batch.

§L.3.1 — Result: the batch order is fixed by the sampler, and the per-example RNG is fixed by the worker count

Claim. Let L be a DataLoader with shuffle=True, generator seeded with s, and num_workers=w. The batch order is a deterministic function of (s, epoch) and does not depend on w. What does depend on w is any randomness inside Dataset.__getitem__ — data augmentation — because each worker process seeds the global RNG to base_seed + worker_id, where base_seed is drawn from L.generator and worker_id ranges over 0…w−1.

Proof sketch. The index sampler is built and consumed in the main process: _BaseDataLoaderIter.__init__ binds self._index_sampler = loader._index_sampler before any worker exists, and _MultiProcessingDataLoaderIter._try_put_index calls self._next_index() — which advances that sampler — then dispatches the resulting integer to a worker. The workers receive indices; they never choose them. So w cannot change the order, and a RandomSampler seeded from L.generator yields the same permutation for w = 0, 2, and 8 (verified by reading iter(loader._index_sampler) under three worker counts: identical). What w does change is _worker_loop’s seed = base_seed + worker_id; random.seed(seed); torch.manual_seed(seed). With w = 0 the loader never enters _worker_loop, so __getitem__ runs under the global RNG, which torch.manual_seed controls; with w > 0 it runs under a per-worker RNG that the global seed does not reach. Verified on a dataset whose __getitem__ calls torch.rand: with w = 0, varying the global seed changes the values while varying the generator seed does not, and with w > 0 the per-worker streams differ by construction.

Corollary. Reproducibility has two independent knobs and the lesson must name both. generator= fixes the order; it does nothing for augmentation. worker_init_fn= fixes the augmentation stream; it does nothing for order. A learner who sets only the seed and then varies num_workers sees identical batch order and different augmentations — the opposite of the common belief, and the reason a “different num_workers gives different results” report almost always means the dataset is stochastic.

Corollary. shuffle=True changes the batch order per epoch, not per example: the default sampler draws a fresh permutation at the start of each epoch, which is the point — the epoch-to-epoch variation is what decorrelates the gradients. With shuffle=False the order is the dataset’s own order, and the seed is not consulted at all (verified: two different generator seeds produce identical batches when shuffle=False).

§L.3.2 — Worked-example ladder

WE-L.3.a — Fully worked: the IK dataset and loader, element by element

Input. The running case: 10,000 (end-effector pose, joint angles) pairs. The pose is a (3,) tensor (a point in the arm’s workspace) and the joint angles are a (6,) tensor. The learner wraps the pairs in a Dataset and iterates a DataLoader with batch_size=64, shuffle=True, num_workers=0.

Trace.

1. class IKDataset(Dataset): with __init__(self, poses, joints) storing self.poses = poses (shape (10000, 3)) and self.joints = joints (shape (10000, 6)), __len__ returning len(self.poses), and __getitem__(self, idx) returning (self.poses[idx], self.joints[idx]).
2. dataset = IKDataset(poses, joints) — len(dataset) is 10,000.
3. dataset[0] — returns a tuple of two views: poses[0] (shape (3,)) and joints[0] (shape (6,)), each sharing the parent tensor’s buffer (Lesson 0.1.1’s view/copy distinction). If the learner modifies these views in place, the dataset is corrupted.
4. loader = DataLoader(dataset, batch_size=64, shuffle=True, num_workers=0).
5. for poses_b, joints_b in loader: — each iteration yields poses_b of shape (64, 3) and joints_b of shape (64, 6), formed by default_collate stacking 64 individual (3,) and (6,) tensors along a new leading axis.
6. The last batch has 10000 mod 64 = 16 examples, shape (16, 3) and (16, 6), because drop_last=False by default.

Output and interpretation. The loader’s yield is the batch, not the example; the batch’s leading dimension is the batch size (except the last, which is the remainder). The default_collate is doing the stacking, and the learner who wants to understand the batch’s construction should look at default_collate’s behavior: it stacks tensors, transposes lists of scalars, and passes through non-tensor objects. The default collate is the reason a Dataset returning (tensor, tensor) works without a custom collate_fn, and the reason a Dataset returning variable-shape tensors does not.

Self-explanation 6. Before reading on: in step 3, why is the view-not-copy behavior of dataset[0] a hazard, and what would the fix be?

Answer. Because a Dataset’s __getitem__ is often followed by an augmentation step (a random crop, a color jitter, a noise injection), and an in-place augmentation on the returned view would modify the parent tensor — the dataset would be silently corrupted after the first epoch, and every subsequent epoch would see the augmented version rather than the original. The fix is to return self.poses[idx].clone() (a copy) or to apply augmentations out of place (each augmentation returning a new tensor). The view is fine if the dataset is read-only; it is a hazard the moment anything downstream writes to it.

WE-L.3.b — Completion problem

Input. The running case, but the demonstrations are variable-length trajectories: each demonstration is (T_i, 6) with a different T_i. The learner wants to batch them.

Trace (final 25% as prompts).

1. class TrajDataset(Dataset): with __getitem__(idx) returning self.trajs[idx], a (T_i, 6) tensor.
2. default_collate([traj_a, traj_b]) where T_a ≠ T_b — ⟨your turn: what happens?⟩
3. The learner writes a collate_fn that pads each trajectory to the batch’s max T with zeros, and returns (padded_trajs, mask) where mask is (B, T_max) with 1 for real timesteps and 0 for padding — ⟨your turn: why is the mask necessary, and how does the model use it?⟩
4. The learner sets shuffle=True — ⟨your turn: is shuffling correct here, and why might it be wrong?⟩

Answers.

1. default_collate raises, because it tries to torch.stack the (T_a, 6) and (T_b, 6) tensors along a new leading axis, and stacking requires equal shapes.
2. The mask is necessary because the padded timesteps are fake — they carry zeros, not real observations — and the model must not treat them as data. It is used in three places: the loss is computed only where mask == 1; the attention, if any, masks out padded positions; and the recurrent hidden state is not updated at padded positions. Without the mask, the model learns from zeros, and the padded positions dominate the gradient.
3. Shuffling is correct for the batch dimension (each batch draws trajectories in random order) but the trajectories’ internal order is not shuffled — the (T_i, 6) tensor’s rows are consecutive timesteps, and they must stay consecutive for a recurrent policy. The shuffle=True flag shuffles the dataset indices, not the rows within a trajectory; the learner who wanted to shuffle the rows has misunderstood the flag.

The lesson. The collate_fn is where variable-shape data becomes a batch, and the mask is where the padding becomes distinguishable from the data. The shuffle flag operates on the dataset, not on the examples’ internal structure; the two are different axes, and the learner who conflates them has either padded the wrong dimension or shuffled a sequence it needed to keep ordered.

WE-L.3.c — Faded: the running case with one parameter changed

Setup. The running case’s training uses num_workers=8 for speed, and the learner wants to reproduce a run exactly. They set torch.manual_seed(42) at the start.

Question. Do two runs with num_workers=8 and the same seed produce the same batch order? Do runs with num_workers=8 and num_workers=0 with the same seed produce the same batch order? What would make the 8-worker runs reproducible across two runs?

You must produce the intermediate steps. The setup and the answer are given.

Answer. Two num_workers=8 runs with the same seed produce the same batch order and the same augmentations, because the worker seeds are derived deterministically as base_seed + worker_id. Runs with 8 and 0 workers produce the same batch order — the index sampler is constructed and advanced in the main process, so the worker count never touches it — but different per-example randomness, because at w=0 __getitem__ draws on the global RNG, which torch.manual_seed(42) fixed, while at w=8 it draws on a worker-local RNG seeded from L.generator that the global seed cannot reach. If the dataset’s __getitem__ is pure (no augmentation, no random crop, no dropout), the two runs are identical despite the different worker counts. The learner’s premise — “same seed, different results, so the seed is broken” — is wrong twice over: the seed never governed the order, and the difference they saw was the augmentation stream. To make the 8-worker runs reproducible they need a seeded generator= and a worker_init_fn; to make them comparable to the 0-worker run they additionally need the 0-worker run’s global seed to match what the worker_init_fn installs in worker 0.

The lesson. Reproducibility is not a single property of the seed; it is two properties with two owners. generator= owns the order, worker_init_fn= owns the per-worker stream, and num_workers owns neither — it only decides whether the per-worker stream exists at all. A learner who sets the seed and assumes reproducibility has modeled neither owner.

WE-L.3.d — Contrastive near-miss: the augmentation that mutates the dataset

Correct solution. The Dataset.__getitem__ returns a clone of the example, and the augmentation is applied out of place: x = self.data[idx].clone(); x = self.augment(x); return x. The parent tensor is unchanged; each epoch sees the original data, augmented independently.

Near-miss. The Dataset.__getitem__ returns self.data[idx] (a view) and the augmentation is in place: x = self.data[idx]; self.augment_inplace(x); return x. The first epoch’s augmentations modify the parent tensor. The second epoch sees the already-augmented data, the third epoch sees the augmented-augmented data, and so on. The dataset drifts away from the original distribution over the course of training.

Why the near-miss produces a plausible but incorrect number. The first epoch’s training loss and validation loss look normal — the data was correct at the start of the first epoch. The second epoch’s training loss decreases relative to the first (the model is improving, and the data is slightly different), and the validation loss is computed on the validation set, which is a separate Dataset and is not corrupted (unless the same bug is in it). The training loss’s trajectory looks like learning; the validation loss’s trajectory looks like learning; the gradual drift of the training data’s distribution is invisible in either curve, and only a comparison of the data’s statistics at epoch 0 and epoch 10 would reveal it.

The discrepancy, stated. self.data[idx] is a view; .clone() makes it a copy. The augmentation’s in-place version modifies whichever buffer it is given, and the difference between the two is whether the buffer is shared with the dataset. The learner who uses the view version has written a self-modifying dataset, and the modification is silent because the first epoch is correct — the drift only begins at the second epoch, and it is monotonic and slow, so no single epoch looks wrong.

Self-explanation 7. Before reading on: in the near-miss, which epoch first sees corrupted data, and why is the first epoch not corrupted?

Answer. The second epoch sees corrupted data — the first epoch’s augmentations modified the parent tensor during the first epoch’s iteration, and the second epoch’s iteration reads the already-modified tensor. The first epoch is not corrupted at the start of the epoch (the data is original), but the augmentations during the epoch corrupt the buffer for later examples within the same epoch if the dataset’s iteration order overlaps (which it does not, if each index is visited once per epoch) — so the first epoch is clean, and the second epoch is uniformly corrupted. The drift is one full epoch’s worth of augmentations, applied cumulatively.

§L.3.3 — Comparative table: loader configurations

Design principle (P3.3). Name the owner of each source of randomness. A reproducible run needs two pins, not one: generator= for the sampler’s order and worker_init_fn= for the per-worker augmentation stream. The seed is necessary for neither on its own — shuffle=False ignores the generator entirely, and num_workers>0 bypasses the global RNG entirely. In a lab report, state the seed, the generator seed, the worker count, and whether __getitem__ is stochastic.

§L.3.4 — Concept check

Concept check (displaces M3). A learner sets torch.manual_seed(42), runs two trainings with num_workers=8 and num_workers=4, and gets different results. They conclude the seed is broken. What is the actual explanation, and what would the learner need to fix to make the two comparable?

Answer. The seed is not broken; it was never in charge. The batch order is identical across the two runs — the index sampler lives in the main process and num_workers cannot reach it. What differs is the number of worker-local RNGs: eight streams seeded base_seed + 0…7 against four seeded base_seed + 0…3. If __getitem__ draws on the global RNG (an augmentation, a random crop, a dropout mask), example i is processed by a different worker with a different stream in the two runs, so the augmentations differ and the trained weights differ. To make the two comparable the learner pins worker_init_fn so each worker’s stream is a function of the example rather than of the worker, or sets num_workers=0, or — the honest fix if the dataset is deterministic — checks whether __getitem__ is pure at all, because a pure dataset makes the two runs identical already.

§L.3.5 — Consolidation checkpoint 3

(a) You should now be able to: (i) write a Dataset that returns examples, and state whether the returned tensors are views or copies; (ii) choose between default_collate and a custom collate_fn for variable-shape data, and use a mask to distinguish padding from data; (iii) state which argument owns which source of randomness — generator= for order, worker_init_fn= for the per-worker stream, neither for the seed — and pin both for a reproducible run.

(b) Transfer probe. A learner’s loader wraps a Dataset that returns (image, label) where image is a (3, 224, 224) tensor and label is a int64 scalar. The loader uses batch_size=32. What are the shapes of the two batch tensors, and what does the default collate_fn do to the label? Answer sketch: image_batch has shape (32, 3, 224, 224) and label_batch has shape (32,) — default_collate stacks the images along a new leading axis and stacks the scalar labels into a 1-D tensor (it does not keep them as a list). The learner who expected label_batch to be a Python list of 32 ints has mis-modeled default_collate; the learner who expected (3, 224, 224, 32) has the axis order wrong.

(c) Remediation rule. If you cannot do (b), return to §L.3.2 and redo WE-L.3.a, paying attention to the leading-axis stacking and the label’s treatment; then re-attempt (b) by writing the shape of each element and the shape after stacking.

Part D — The Training Loop: The Canonical Composition

§L.4 — Objects: the six lines, their order, and their effects

The canonical training loop, for a supervised task, is:

for epoch in range(num_epochs):
    model.train()
    for x, y in train_loader:
        x = x.to(device, non_blocking=True)
        y = y.to(device, non_blocking=True)
        optimizer.zero_grad()
        pred = model(x)
        loss = loss_fn(pred, y)
        loss.backward()
        optimizer.step()
    scheduler.step()
    model.eval()
    with torch.no_grad():
        val_loss = evaluate(model, val_loader)

The six lines whose order matters, and their effects:

1. optimizer.zero_grad() — clears the accumulated .grad on every parameter in the optimizer’s groups. Must be before loss.backward(); should be at the start of each iteration (or after the previous step, which is equivalent).
2. pred = model(x) — the forward pass; builds the graph.
3. loss = loss_fn(pred, y) — the scalar loss; the root of the graph.
4. loss.backward() — the backward pass; accumulates .grad.
5. optimizer.step() — reads .grad and updates the parameters. Must be after loss.backward().
6. scheduler.step() — advances the learning-rate schedule. Must be after optimizer.step() (the schedule’s step=0 value is the initial LR and applies to the first update); should be once per optimizer step or once per epoch, according to the schedule’s semantics.

The mode lines — model.train() and model.eval() — are the subject of Lesson 0.1.4; for this lesson, the learner should note only that they must be called, and that the evaluate function must wrap its forward pass in torch.no_grad() (Lesson 0.1.2, §L.4) so that the validation’s forward passes do not build graphs.

§L.4.1 — Result: each misordering of the six lines has a distinct silent failure, and the failures are distinguishable from the loss trajectory and the gradient norm

Claim. For each of the four common misorderings — zero_grad after backward but before step, zero_grad once per epoch, scheduler.step() before optimizer.step(), and step() before backward() — the training loop runs without error, and the resulting loss trajectory or gradient norm has a distinct signature that identifies the misordering.

Proof sketch. (i) zero_grad between backward and step clears the gradient the step was about to use, so the parameters receive zero (or None) gradient and do not move; the loss is flat at its initial value. (ii) zero_grad once per epoch accumulates the gradients across the epoch’s mini-batches, so the effective update at the end of the epoch is the sum of the epoch’s gradients, and the effective learning rate is num_minibatches × lr; the loss decreases faster than intended, then diverges or plateaus at a different point than the correct run. (iii) scheduler.step() before optimizer.step() applies a schedule that is one step ahead of the intended one, so the anneal is slightly faster and the schedule’s first value is never applied; the loss trajectory is slightly worse, and the difference is small enough to be attributed to noise without a controlled comparison. This is the one member of the four that is not fully silent: PyTorch emits a UserWarning on the first misordered call, which makes it quiet rather than silent. (iv) step() before backward() reads the previous iteration’s gradient (or zero on the first iteration), so the parameters are updated with a stale gradient; the loss decreases but with a one-step lag, which is invisible in the curve and visible only in a step-by-step comparison.

The four signatures are distinguishable: (i) is flat, (ii) is fast-then-worse, (iii) is slightly-worse-with-same-shape, and (iv) is indistinguishable from the correct loop except by inspecting the parameters’ values at each step. The learner who suspects a misordering should print the loss at every step for the first 10 steps and compare against the correct loop; the flat and fast-then-worse signatures are obvious, the slightly-worse signature requires a controlled comparison, and the stale signature requires the per-step parameter comparison.

§L.4.2 — Worked-example ladder

WE-L.4.a — Fully worked: a single iteration of the IK loop, traced line by line

Input. The IKNet with fc1 and fc2; Adam(lr=1e-3); MSELoss(); a batch of (64, 3) poses and (64, 6) joint targets; device = cuda:0.

Trace.

1. x = x.to(device) — x shape (64, 3), device cuda:0, requires_grad=False.
2. y = y.to(device) — y shape (64, 6), device cuda:0, requires_grad=False.
3. optimizer.zero_grad() — each parameter’s .grad is None (or zero); the graph is unaffected (it has not been built yet).
4. pred = model(x) — fc1.weight (shape (64, 3)) matmul x^T, plus bias, relu, fc2, plus bias; pred shape (64, 6), requires_grad=True, grad_fn=<AddmmBackward>; the graph is built, and the intermediate activations are saved.
5. loss = loss_fn(pred, y) — (64, 6) against (64, 6), mean of squared error, shape () (a scalar), requires_grad=True.
6. loss.backward() — traverses the graph, accumulates .grad on fc1.weight, fc1.bias, fc2.weight, fc2.bias; the graph is freed.
7. optimizer.step() — for each parameter, p ← p - lr · p.grad (with Adam’s moments); the parameters’ .grad is not cleared by the step (it stays until the next zero_grad).
8. scheduler.step() — advances the schedule by one.

Output and interpretation. After the iteration, the parameters have been updated, their .grad still holds the last computed values (until the next zero_grad), and the graph is gone. The learner who prints fc1.weight.grad after step sees the last iteration’s gradient — not zero, because step does not clear it — and this is the source of the “why is my gradient not zero after step” confusion. The only operation that clears .grad is zero_grad (or p.grad = None).

Self-explanation 8. Before reading on: in step 3, the parameters’ .grad is cleared, but the graph is not yet built. What does this tell you about the relationship between zero_grad and the graph?

Answer. zero_grad clears the gradient buffers, which are a property of the parameters, not of the graph. The graph is built by the forward pass (step 4) and freed by the backward pass (step 6); zero_grad at step 3 happens before the graph exists, and zero_grad at the start of the next iteration happens after the previous graph is freed. The two are independent of the graph’s lifecycle — the only graph operation zero_grad interacts with is nothing — and the learner who thinks zero_grad “resets the graph” has conflated the gradient buffers with the graph’s saved tensors.

WE-L.4.b — Completion problem

Input. The same loop, with an evaluate function to be added. The learner wants to compute the validation loss every 10 epochs and to save the best model.

Trace (final 25% as prompts).

1. Add if epoch % 10 == 0: after the training loop — ⟨your turn: what should be inside the if, and in what order?⟩
2. The evaluate function needs model.eval(), torch.no_grad(), a loop over val_loader, and a val_loss accumulation — ⟨your turn: place the model.eval() and torch.no_grad() and state why the order matters⟩
3. After evaluation, the training loop resumes — ⟨your turn: what must be called to return to training mode?⟩
4. Saving the best model — ⟨your turn: what should be saved, and what comparison determines “best”?⟩

Answers.

1. Inside the if: model.eval(), the with torch.no_grad(): block with the validation loop, model.train(), and — if the current val_loss is the best so far — the checkpoint save. The order is eval → no_grad → val → train → save-if-best.
2. model.eval() must be before the forward passes, because batch norm and dropout read the training flag at the forward pass (Lesson 0.1.4). torch.no_grad() must wrap the forward pass and the loss computation, because they do not need gradients and building a graph for them wastes memory and latency (Lesson 0.1.2, §L.4). The order between eval and no_grad is not critical — both are set before the forward pass — but the convention is eval first, because it is the mode change and no_grad is the graph change.
3. model.train() must be called after the evaluation, because the next training iteration’s forward pass must be in training mode. Forgetting it is the subject of Lesson 0.1.4, and the symptom is a validation loss that is systematically different from the training loss.
4. Save the model’s state_dict() and the optimizer’s and the scheduler’s and the epoch number (Lesson 0.1.2, §L.2). The “best” comparison is on the validation loss, not the training loss — the training loss decreases monotonically in almost every setting, and saving on training loss saves the most overfit checkpoint, not the best-generalizing one.

The lesson. The evaluation block’s placement (inside the epoch loop, after the training loop, under a frequency condition) and its mode discipline (eval, no_grad, train) are what makes the validation meaningful. The learner who computes validation during the training loop without the mode change has a validation that is computed in training mode, and the comparison against the training loss is meaningless; the learner who forgets model.train() after the validation has the next training iteration in eval mode, and the training silently degrades.

WE-L.4.c — Faded: the running case with one parameter changed

Setup. The learner moves scheduler.step() from its current position (after the epoch loop) to inside the epoch loop, after optimizer.step(). The scheduler is CosineAnnealingLR(optimizer, T_max=200), and the learner is training for 200 epochs with 100 minibatches per epoch (20,000 optimizer steps total).

Question. What does T_max=200 mean in each position? What is the learning rate at step 0, 100, and 20,000 in each position? Which is correct for the schedule the learner intended?

You must produce the intermediate steps. The setup and the answer are given.

Answer. CosineAnnealingLR’s T_max is the number of scheduler steps over which the cosine anneals from lr₀ to 0. If the scheduler is stepped once per epoch (the original position), T_max=200 means 200 epochs, and the anneal is over the intended 200 epochs. If it is stepped once per minibatch (the new position), T_max=200 means 200 minibatches — about 2 epochs — and the anneal is over 2 epochs, not 200. At step 0: lr₀. At step 100 (minibatches, new position): lr₀ · (1 + cos(π · 100/200))/2 = lr₀ · (1 + 0)/2 = lr₀/2 — the learning rate has halved in 2 epochs. At step 20,000, the schedule does not sit at zero: CosineAnnealingLR is periodic with period 2·T_max, so it anneals to 0 at step 200 and then climbs back, reaching lr₀ again at step 400 — measured: lr(200) = 0.0, lr(400) = lr₀, lr(600) = 0.0, and at step 20,000 the rate is lr₀ once more. The new position therefore does not merely anneal too fast; it cycles 100 times through the full anneal during a run that was meant to anneal once, so the late epochs train at full learning rate. In the old position (once per epoch), step 20,000 (epoch 200) has lr ≈ 0 as intended, and step 100 (epoch 1) has lr ≈ lr₀, which is the intended slow anneal. The correct position for a 200-epoch cosine schedule is once per epoch; the new position makes the schedule 100× faster and turns a monotone anneal into a sawtooth.

The lesson. The scheduler’s T_max is in scheduler steps, not in epochs and not in optimizer steps; the placement of scheduler.step() determines what a “step” is. A learner who moves the step without adjusting T_max has silently multiplied the anneal rate by the number of optimizer steps per epoch — a 100× error in the schedule, with a loss trajectory that shows the learning rate collapsing in the first 2 epochs and then cycling back to full every 4 epochs thereafter, because CosineAnnealingLR is periodic with period 2·T_max rather than clamping at zero. The fix is either to keep the once-per-epoch position with T_max=200, or to move to once-per-optimizer-step and set T_max=20,000.

WE-L.4.d — Contrastive near-miss: zero_grad between backward and step

Correct solution. The loop is zero_grad(); forward; loss; backward; step(). The step reads the gradient the backward just computed, and the parameters move.

Near-miss. The loop is forward; loss; backward; zero_grad(); step(). The zero_grad clears the gradient the step was about to use, and the step reads zero (or None) gradient and does not move the parameters. The loss is flat at its initial value across all iterations, and the loop runs without error.

Why the near-miss produces a plausible but incorrect number. The loss is flat — it does not decrease, and it does not increase. The learner who wrote the loop and sees a flat loss has a signature that is consistent with many causes: a learning rate that is too small, a model that is saturated, a loss that is not connected to the parameters, a data loader that returns constant data. The actual cause is the zero_grad placement, and the discriminator is the gradient norm before the step: if the gradient is non-zero before the zero_grad and zero after, the placement is the bug. The learner who only prints the loss sees a flat line and has no information about the gradient.

The discrepancy, stated. The correct loop’s zero_grad is before the backward; the near-miss’s is after. The two loops differ in one line’s position, and the effect is the difference between “the step sees the gradient” and “the step sees zero.” The zero_grad’s purpose — clear the previous iteration’s gradient — is served by placing it at the start of the iteration (before any computation that reads the gradient) or at the end of the previous iteration (after the step that consumed it); the middle placement (between backward and step) clears the gradient that the current iteration just computed, and is the only placement that is always wrong.

Self-explanation 9. Before reading on: in the near-miss, what is the value of p.grad after step()?

Answer. None (or zero), because the zero_grad cleared it and the step did not write to it. The correct loop’s p.grad after step() is the last computed gradient (not cleared until the next zero_grad); the near-miss’s is None. The p.grad after step is a diagnostic: if it is None, the zero_grad was after the backward; if it is non-zero, the zero_grad was before the backward.

§L.4.3 — Comparative table: loop orderings and their failures

Design principle (P3.4). The gradient’s lifecycle is a loop: clear at the start, fill by the backward, consume by the step. Any placement that clears the gradient after the fill, or consumes it before the fill, is a bug; the correct placements are the ones that respect the loop’s arrow of time, and the near-miss’s placements are the ones that cross it.

§L.4.4 — Concept check

Concept check (displaces M4). A learner says: “All five operations are present, so the order does not matter.” Name the two orderings that produce a flat loss and a fast-then-worse loss, and state which operation’s placement is the cause in each.

Answer. The flat loss is produced by zero_grad between backward and step, because the step consumes a cleared gradient; the cause is the zero_grad placement. The fast-then-worse loss is produced by zero_grad once per epoch (rather than once per minibatch), because the gradients accumulate across the epoch’s minibatches and the effective learning rate is num_minibatches × lr; the cause is the zero_grad placement again, but in the other direction. The two are the extremes of the accumulation: too late a clear is a lost gradient (flat), too early a clear is an accumulated gradient (fast-then-worse).

§L.4.5 — Consolidation checkpoint 4

(a) You should now be able to: (i) write the canonical loop with the correct ordering, and state the purpose of each line; (ii) place the zero_grad at the start of the iteration (or the end of the previous), never between backward and step; (iii) place the scheduler.step() after optimizer.step(), and state what T_max counts in each placement; (iv) place the evaluation block correctly, with eval and no_grad, and return to train after.

(b) Transfer probe. A learner’s loss decreases for the first 50 iterations, then becomes noisy and increasing. They have zero_grad once per epoch, 100 minibatches per epoch. What is the most likely cause? Answer sketch: the accumulation’s effective learning rate grows with the number of accumulated minibatches, and at 50 iterations (within the first epoch) the effective rate has exceeded the stability threshold for the optimizer and learning rate, so the loss starts to diverge. The fix is zero_grad once per minibatch. The discriminator is the ratio of the effective learning rate to the stability threshold: the correct loop’s rate is stable, the near-miss’s is 100× the correct rate, and the threshold is crossed early in the epoch.

(c) Remediation rule. If you cannot do (b), return to §L.4.2 and redo WE-L.4.d, then re-attempt (b) by writing the gradient’s state at each line of the loop.

Part E — Failure Taxonomy, Lab, Cheat Sheet, Exercises, Exit Test

§L.5 — Failure taxonomy

The table covers the lab’s three break-it steps (flat, new-layer, ordered-loader) plus three failures not exercised there (resume spike, validation gap, worker reproducibility).

§L.6 — The lab

§L.6.1 — Task

Build a complete IK regression training pipeline — Dataset, DataLoader, nn.Module, optim, and the canonical loop — for the running case, and use it to reproduce, predict, and fix three pipeline bugs.

§L.6.2 — Conditions

Python 3.10+, PyTorch 2.2+, CPU only. All runs use torch.manual_seed(42). The dataset is 8,000 (pose, joints) pairs generated by the forward kinematics of the arm’s 3-link planar reduction — end-effector (x, y, φ) against joint angles (q₁, q₂, q₃) — with q₁ free over the full circle and the two elbow joints held in (0.15, π − 0.15). §L.6.10 states why the 6-DOF map cannot be used here, and it is the single most important design decision in the lab. The model is IKNet (three linear layers, 128 hidden). The optimizer is Adam(lr=3e-3), the scheduler is CosineAnnealingLR(T_max=epochs), the loss is MSELoss(), and the training is 40 epochs with batch_size=64. Runtime: 8 seconds on CPU.

§L.6.3 — Controls held fixed

The network initialization seed (0), the dataset (the same 8,000 pairs), the architecture (IKNet), the batch size (64), the learning rate (3e-3), the epoch count (40), and num_workers=0 are fixed across all runs. The two H2 runs use the same 6,400 training pairs in two orders — a random permutation and a sort by q₁ — so the only thing that varies between them is order.

§L.6.4 — Metrics table

§L.6.5 — Pre-registered hypotheses

• H1. With the correct loop (shuffle=True, zero_grad before backward), the validation MSE falls below 0.05, and below 1/100 of the mean-predictor baseline (predicting the training mean for every input, measured at 1.520677; the lab prints this baseline and the achieved ratio, 183.1x). The prediction is two-sided: above 0.05 is a failure to learn the pipeline, and a value that beats the baseline only marginally is a pipeline that runs but does not fit.
• H2. With shuffle=False, the validation MSE depends on whether the training set is ordered, not on the flag alone. Against a random-permutation training set, shuffle=False costs almost nothing (predicted within 20% of H1); against the same pairs sorted by q₁, it costs a factor of at least 3 (predicted > 3×, mechanism: consecutive batches come from one region of the joint space, so each step’s gradient is a biased estimate of the full-dataset gradient).
• H3. With zero_grad moved between backward and step, the per-epoch training loss is constant across all 5 epochs (to 1e-6), and the parameter’s .grad after the final step() is None — not zero. The prediction is a flat list of 5 equal values and None; a prediction of 0.0 is the common error and is falsified.
• H4. With the optimizer constructed before fc3 is reassigned, fc3_grad_norm is > 0 (autograd still computes it) and fc3.weight is unchanged (the optimizer never sees it), and an identity scan of param_groups finds fc1.weight and does not find fc3.weight. The prediction is three statements, and the third is what distinguishes “the gradient is missing” from “the parameter is missing”.

§L.6.6 — The break-it sequence

1. zero_grad between backward and step (silent). Move optimizer.zero_grad() to after loss.backward(). Predicted symptom: no exception; the epoch losses are flat at the initial value; the parameter’s .grad is None after the step. Actual: as predicted — [3.348493, 3.348493, 3.348493, 3.348493, 3.348493], flat True. The step is applied with a None gradient, so the optimizer skips the parameter entirely and the loss never moves.
2. Optimizer constructed before fc3 is replaced (silent). Construct the optimizer with model.parameters(), then assign a new model.fc3. Train 50 steps. Predicted symptom: no exception; fc3.grad.norm() is 46.5187 — large and clearly populated; fc3.weight is unchanged; fc3 is absent from param_groups. Actual: as predicted — silent, and the gradient’s magnitude is the trap: it looks exactly like a healthy layer.
3. Sorted training set, shuffle=False (silent). Sort the 6,400 pairs by q₁ and disable shuffling. Predicted symptom: no exception; validation MSE 0.038600 against 0.009155 for the same pairs in random order — a factor of 4.22 — with no divergence and no plateau that looks pathological. Actual: as predicted — silent. The epoch loss is non-monotone for the first six epochs (0.4048 → 0.5110 → 0.5326 → 0.4219 → 0.4483 → 0.4495) and then sits near 0.056 for epochs 25–33 before falling again at the end; the shuffled run descends smoothly from 0.5398 to 0.0074. The sorted run’s sawtooth is the diagnosis: each epoch begins in a new region of the joint space and the model’s fit to the previous region is partly discarded, so the epoch average depends on which region the epoch happened to end in.

Three silent break-it steps; the lesson’s minimum is two. The third is included because it is the one where the flag and the phenomenon come apart: the learner’s model of “shuffle=False is bad” predicts the wrong number for the random-order control, and only running all three separates them.

§L.6.7 — Acceptance criteria

§L.6.8 — What good looks like

A strong submission has five properties:

1. The correct pipeline is shown to fit, against a stated floor. val_loss below 0.05 and below one hundredth of the mean-predictor baseline. Reporting a loss without a baseline is reporting a number with no scale; the baseline is what makes “converged” a claim.
2. Each break-it is attributed to a specific line, and the line’s correct position is stated. “The bug was zero_grad after backward; the correct position is before backward, and the tell is that .grad is None at the step” — not “the bug was a gradient issue.”
3. The gradient diagnosis is shown, not inferred. fc3.weight’s value at construction and after 50 steps is printed and the two are identical; fc3_grad_norm is printed and non-zero; fc3.weight’s identity in param_groups is printed and False. All three together are the diagnosis — the gradient alone says the layer is fine, and the registry alone does not say why.
4. The H2 comparison isolates order. The two runs have the same pairs, the same seed, the same architecture, and the same learning rate, and differ only in the training set’s order. The learner who compares sorted-against-random while also flipping shuffle has changed two things and isolated neither — which is why the lab runs the random-order control as well.
5. The checkpoint is complete and is demonstrated to resume correctly. The learner saves a mid-training checkpoint, reloads it, and shows that the first 5 steps after resume have the same loss values as the corresponding steps before the save. The demonstration is the resume, not the save.

§L.6.9 — Runnable lab code

# lab_0_1_3.py — the IK training pipeline, its correct composition, and its three bugs
# Requires: Python 3.10+, PyTorch 2.2+, CPU only.

import math
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader

torch.manual_seed(42)
device = torch.device("cpu")

# --- Step 1: forward kinematics for a 3-DOF planar reduction of the arm ---
# The 6-DOF arm is used in the narrative; the lab uses its 3-link planar
# reduction because that map is INJECTIVE once the elbow limits are imposed.
# (The 6-DOF map is not: three of its joints enter only through their sum, so
# no network can invert it, and an unreachable target destroys the whole
# experiment. See §L.6.10.)
def fk3(q):
    """q: (..., 3) joint angles. Returns (..., 3) = end-effector (x, y, phi)."""
    q1, q2, q3 = q[..., 0], q[..., 1], q[..., 2]
    x = torch.cos(q1) + torch.cos(q1 + q2) + torch.cos(q1 + q2 + q3)
    y = torch.sin(q1) + torch.sin(q1 + q2) + torch.sin(q1 + q2 + q3)
    return torch.stack([x, y, q1 + q2 + q3], dim=-1)

N = 8000
# q1 spans the full circle (no limit); q2, q3 are elbow-like, held off both
# singularities. Without those limits the map is 2-to-1 and the task is
# unlearnable no matter how long you train.
q_all = torch.stack([(torch.rand(N) * 2 - 1) * math.pi,
                     0.15 + (math.pi - 0.3) * torch.rand(N),
                     0.15 + (math.pi - 0.3) * torch.rand(N)], dim=-1)
pose_all = fk3(q_all)

# --- Step 2: the split, and the two orderings of the SAME training set ----
# The split is a random permutation, so the training set's default order is
# already decorrelated. That is what makes the H2 control clean: the only
# difference between the two training sets is their order.
perm = torch.randperm(N)
pose_all, q_all = pose_all[perm], q_all[perm]
n_train = int(0.8 * N)
P_tr, q_tr = pose_all[:n_train], q_all[:n_train]
P_va, q_va = pose_all[n_train:], q_all[n_train:]

class IKDataset(Dataset):
    def __init__(self, poses, joints):
        self.poses, self.joints = poses, joints
    def __len__(self):
        return len(self.poses)
    def __getitem__(self, i):
        return self.poses[i].clone(), self.joints[i].clone()

train_random = IKDataset(P_tr, q_tr)
order = torch.argsort(q_tr[:, 0])                  # sorted by the first joint
train_sorted = IKDataset(P_tr[order], q_tr[order])
val_ds = IKDataset(P_va, q_va)

# --- Step 3: the Module ------------------------------------------------
class IKNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(3, 128)
        self.fc2 = nn.Linear(128, 128)
        self.fc3 = nn.Linear(128, 3)
    def forward(self, x):
        return self.fc3(torch.relu(self.fc2(torch.relu(self.fc1(x)))))

def fresh_net(seed=0):
    torch.manual_seed(seed)
    return IKNet()

print(f"parameters: {[n for n, _ in IKNet().named_parameters()]}")

# --- Step 4: the canonical loop, as a function so the bugs can be injected ---
def train(train_ds, shuffle, epochs=40, lr=3e-3, seed=0,
          zero_grad_pos="before", model=None):
    m = model if model is not None else fresh_net(seed)
    opt = torch.optim.Adam(m.parameters(), lr=lr)
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs)
    tl = DataLoader(train_ds, batch_size=64, shuffle=shuffle, num_workers=0)
    vl = DataLoader(val_ds, batch_size=512, shuffle=False, num_workers=0)
    losses = []
    for _ in range(epochs):
        m.train(); total = 0.0
        for x, y in tl:
            if zero_grad_pos == "before":
                opt.zero_grad()
            loss = nn.functional.mse_loss(m(x), y)
            loss.backward()
            if zero_grad_pos == "between":
                opt.zero_grad()      # <- break-it #1 lives on this line
            opt.step()
            total += loss.item()
        sched.step()
        losses.append(total / len(tl))
    m.eval()
    with torch.no_grad():
        v = sum(nn.functional.mse_loss(m(x), y).item() for x, y in vl) / len(vl)
    return m, losses, v

# --- Step 5: the correct pipeline (H1) --------------------------------
model, losses_ok, val_ok = train(train_random, shuffle=True)
# The baseline the acceptance criterion is stated against: predict the
# training targets' mean for every input. Reporting a loss without it is
# reporting a number with no scale.
mean_pred = q_tr.mean(0)
baseline = nn.functional.mse_loss(mean_pred.expand_as(q_va), q_va).item()
print(f"H1 val_loss = {val_ok:.6f}   mean-predictor baseline = {baseline:.6f}"
      f"   ratio = {baseline / val_ok:.1f}x")

# --- Step 6: the ordering control (H2) --------------------------------
# Both runs have shuffle=False. They differ ONLY in the order of the
# training set, so the gap is the order's effect and nothing else.
_, _, val_sorted = train(train_sorted, shuffle=False)
_, _, val_ctrl = train(train_random, shuffle=False)
print(f"H2 val (sorted train, no shuffle)  = {val_sorted:.6f}")
print(f"H2 val (random train, no shuffle)  = {val_ctrl:.6f}")

# --- Step 7: break-it #1 — zero_grad between backward and step (silent) ---
_, losses_z, _ = train(train_random, shuffle=True, epochs=5,
                       zero_grad_pos="between")
print(f"H3 losses = {[round(x, 6) for x in losses_z]}")
print(f"H3 flat? {max(losses_z) - min(losses_z) < 1e-6}")

# --- Step 8: break-it #2 — optimizer built before the head exists (silent) ---
m_e = fresh_net(0)
opt_e = torch.optim.Adam(m_e.parameters(), lr=1e-3)   # constructed HERE
m_e.fc3 = nn.Linear(128, 3)                            # a NEW head, not in opt_e
with torch.no_grad():
    fc3_init = m_e.fc3.weight.clone()
tl = DataLoader(train_random, batch_size=64, shuffle=True)
for _ in range(50):
    x, y = next(iter(tl))
    opt_e.zero_grad()
    loss = nn.functional.mse_loss(m_e(x), y)
    loss.backward()          # fc3's gradient IS computed: autograd sees the module
    opt_e.step()             # fc3 is never updated: the optimizer does not
print(f"H4 fc3.grad_norm = {m_e.fc3.weight.grad.norm().item():.4f}")
print(f"H4 fc3 unchanged? {torch.allclose(m_e.fc3.weight, fc3_init)}")
print(f"H4 fc1 in optimizer? "
      f"{any(p is m_e.fc1.weight for g in opt_e.param_groups for p in g['params'])}")
print(f"H4 fc3 in optimizer? "
      f"{any(p is m_e.fc3.weight for g in opt_e.param_groups for p in g['params'])}")

# --- Step 9: the two reproducibility knobs ----------------------------
def batch_order(seed, shuffle=True):
    g = torch.Generator().manual_seed(seed)
    dl = DataLoader(train_random, batch_size=64, shuffle=shuffle,
                    num_workers=0, generator=g)
    return torch.cat([x for x, _ in dl])

print(f"same generator seed -> identical: "
      f"{torch.equal(batch_order(0), batch_order(0))}")
print(f"different generator seed -> identical: "
      f"{torch.equal(batch_order(0), batch_order(1))}")
print(f"shuffle=False ignores the seed: "
      f"{torch.equal(batch_order(0, False), batch_order(1, False))}")

# --- Step 10: the checkpoint ------------------------------------------
opt = torch.optim.Adam(model.parameters(), lr=1e-3)
sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=40)
ckpt = {"model": model.state_dict(), "optimizer": opt.state_dict(),
        "scheduler": sched.state_dict(), "step": 12345}
torch.save(ckpt, "lab_0_1_3_ckpt.pt")
loaded = torch.load("lab_0_1_3_ckpt.pt", weights_only=False)
print(f"checkpoint keys: {sorted(loaded.keys())}")

# --- Step 11: summary -------------------------------------------------
print("\n--- lab 0.1.3 summary ---")
print(f"H1 val_loss_final:              {val_ok:.6f}")
print(f"H2 val (sorted, no shuffle):    {val_sorted:.6f}")
print(f"H2 val (random, no shuffle):    {val_ctrl:.6f}")
print(f"H3 flat?                        {max(losses_z) - min(losses_z) < 1e-6}")
print(f"H4 fc3 unchanged?               {torch.allclose(m_e.fc3.weight, fc3_init)}")
print("silent break-its: #1 (zero_grad), #2 (early construction), #3 (sorted order)")

Run the file end to end. It takes about 8 seconds and prints:

parameters: ['fc1.weight', 'fc1.bias', 'fc2.weight', 'fc2.bias', 'fc3.weight', 'fc3.bias']
H1 val_loss = 0.008304   mean-predictor baseline = 1.520677   ratio = 183.1x
H2 val (sorted train, no shuffle)  = 0.038600
H2 val (random train, no shuffle)  = 0.009155
H3 losses = [3.348493, 3.348493, 3.348493, 3.348493, 3.348493]
H3 flat? True
H4 fc3.grad_norm = 46.5187
H4 fc3 unchanged? True
H4 fc1 in optimizer? True
H4 fc3 in optimizer? False
same generator seed -> identical: True
different generator seed -> identical: False
shuffle=False ignores the seed: True
checkpoint keys: ['model', 'optimizer', 'scheduler', 'step']

The last three lines are the reproducibility experiment, and they are worth pausing on. The generator seed controls the order and nothing else; shuffle=False does not consult it at all. That is one knob. The other is worker_init_fn, which the lab leaves unexercised because this host’s num_workers>0 path fails (torch_shm_manager: execl failed: Permission denied); §L.3.1 states what the source does and Step 9 shows the half that is testable here.

If any step does not match, the failure taxonomy in §L.5 names the likely causes.

§L.6.10 — Why the lab does not use the 6-DOF arm

The running case is a 6-DOF arm, and the obvious lab is to learn its IK directly: map the end-effector pose (x, y, z) to the six joint angles. That lab cannot work, and the reason is structural rather than a matter of training budget.

Take the 6-DOF forward kinematics used in the earlier draft of this lesson: x and y are sums of cosines and sines of q₁, q₁+q₂, q₁+q₂+q₃, and z = q₄ + q₅ + q₆. Three joints enter the output only through their sum. The map is therefore many-to-one: every point (q₄, q₅, q₆) on a common sum maps to the same z, and no function — linear or otherwise — can recover which one produced it. The Bayes-optimal predictor outputs the conditional mean of the joint angles given the pose, and the MSE it achieves is the conditional variance, which is not small. Measured, on 10,000 samples with 200 epochs: validation MSE 2.857068, against a predicted < 1e-3. The network had converged; it had converged to the best answer a non-injective target admits.

The consequence for the lab is worse than a wrong number. H1 asked the learner to predict convergence, H2 asked them to predict a ten-fold gap from shuffling, and neither can appear: the loss floor is set by the non-injectivity, so shuffling moves it by 10⁻⁶ (2.857068 against 2.857067) and the learner’s prediction is scored against an effect that the task’s geometry has erased. An exercise whose hypothesis cannot be true teaches the learner to distrust their own correct reasoning.

The fix is to keep the phenomenon and change the arm. The 3-link planar reduction with q₁ free and the elbows in (0.15, π − 0.15) is injective on its range, and the argument is short enough to check by hand. Write a₁ = q₁, a₂ = q₁ + q₂, a₃ = q₁ + q₂ + q₃ = φ. The pose gives φ outright, hence a₃; subtracting the third link’s contribution leaves S = (cos a₁ + cos a₂, sin a₁ + sin a₂), the sum of two unit vectors. Its magnitude is |S| = 2 cos(q₂/2) and its direction is (a₁ + a₂)/2, so q₂ = 2 arccos(|S|/2) — single-valued because q₂ ∈ (0.15, π − 0.15) keeps cos(q₂/2) > 0 — and then a₂ = dir(S) + q₂/2 and q₃ = φ − a₂. The last step is where the range condition earns its keep: q₃ is recovered modulo 2π, and the constraint 0.15 < q₃ < π − 0.15 selects the unique representative. Every step is single-valued, so the pose determines the joints. Reconstructed over 20,000 samples, the closed form recovers q₂ and q₃ to 2 × 10⁻⁶.

Two numerical checks corroborate it, and they test different things:

• No folds. The Jacobian ∂(x, y, φ)/∂(q₁, q₂, q₃) has a strictly positive determinant over the whole domain — measured over 200,000 uniform draws, min = 0.1494, max = 1.0000, no sample negative. A fold requires the determinant to change sign; the elbow limits are what prevent it, since an unrestricted planar arm’s two mirror branches meet exactly where the determinant vanishes.
• No collisions at the sampled scale. Across 20,000 samples, the closest pair of poses (Euclidean distance 0.0067) is also a close pair in joint space (0.0056). No two distant joint configurations landed on one pose.

On that task the same network reaches 0.008304, and the ordering effect reappears at its true size — 4.22×, not ten-fold.

The lesson generalizes past this lab. A learnability experiment must first establish that the target is a function. Before predicting a loss floor, check that distinct inputs do not share an output; a non-injective regression target has a floor set by the conditional variance, and no architecture, learning rate, or epoch count moves it. The tell in the loss curve is a plateau that is flat, low-looking, and unreachable by further training — and the tell in the notebook is a hypothesis stated as a threshold rather than as a mechanism.

The same caution applies to the running case’s own numbers throughout this curriculum: where a lesson reports an IK result on the full 6-DOF arm, the reported quantity is a command-tracking error, not a pose-reconstruction error, because only the former is well posed.

§L.7 — The cheat sheet

§L.8 — Exercises with solutions

E 0.1.3.1 — The registration count (numerical). A module has self.enc = nn.Linear(3, 8), self.norm = nn.LayerNorm(8), self.head = nn.Linear(8, 6), self.register_buffer("running_mean", torch.zeros(3)), and self.scale = torch.ones(1, requires_grad=True). How many tensors are in parameters(), how many in state_dict(), and where is self.scale?

S 0.1.3.1. parameters(): enc.weight, enc.bias, norm.weight, norm.bias, head.weight, head.bias — six. state_dict(): the six plus running_mean — seven. self.scale is a plain attribute (not nn.Parameter, not registered), so it is in neither; it is in __dict__ and invisible to the optimizer and the checkpoint. The lesson: the registration is the only thing the framework reads; a requires_grad=True tensor that is not wrapped is a parameter to autograd and not to the module.

E 0.1.3.2 — The early-construction diagnosis (diagnostic). A learner’s two-layer model trains fc1 and fc2 correctly, but a third layer fc3 added after the optimizer’s construction does not train. The learner’s fc3.weight.grad is non-zero and its fc3.weight value is unchanged. State the cause, the two fixes, and the one call that would have detected the bug before the training started.

S 0.1.3.2. Cause: the optimizer was constructed with model.parameters() before fc3 was added; the optimizer’s parameter list is a snapshot at construction, and fc3 is not in it. Fixes: (1) construct the optimizer after fc3 is added, or (2) optimizer.add_param_group({'params': model.fc3.parameters(), 'lr': 1e-3}). The detection call: any(p is model.fc3.weight for g in optimizer.param_groups for p in g['params']), which returns False for the buggy construction. The lesson: the optimizer’s parameter set is a snapshot, and the snapshot’s contents are the thing to check, not the parameters’ .grad (which is populated regardless).

E 0.1.3.3 — The scheduler T_max (numerical). A learner trains for 100 epochs with 50 minibatches per epoch and uses CosineAnnealingLR(T_max=5000), stepping the scheduler once per minibatch. What is the learning rate at epoch 0, epoch 50, and epoch 99, relative to lr₀? What would the correct T_max be for a once-per-minibatch schedule over the full training?

S 0.1.3.3. T_max=5000 is the number of scheduler steps; with 50 minibatches/epoch, 5000 steps = 100 epochs, which is exactly the training length. At epoch 0 (step 0): lr = lr₀ · (1 + cos(0))/2 = lr₀. At epoch 50 (step 2500): lr = lr₀ · (1 + cos(π · 2500/5000))/2 = lr₀ · (1 + 0)/2 = lr₀/2. At epoch 99 (step 4950): lr ≈ lr₀ · (1 + cos(π · 0.99))/2 ≈ 0. The schedule is correct as written — the learner has set T_max to the total number of minibatches, which is the correct interpretation. The common bug is the opposite: a learner who writes T_max=100 (the epoch count) while stepping per minibatch gets a schedule that finishes in 2 epochs. The lesson: T_max is in scheduler steps; the learner must set it to the total number of scheduler.step() calls, which depends on the placement.

E 0.1.3.4 — The resume spike (diagnostic). A learner resumes a PPO training from a model-only checkpoint and observes a success-rate spike for the first ~100 steps, then a drop back to (slightly below) the pre-save level. What is the cause, and what would the learner have to save to eliminate the spike?

S 0.1.3.4. The cause is the missing optimizer state: Adam’s first and second moments are re-initialized to zero. Bias correction keeps the first update at exactly lr per coordinate rather than the m/eps blow-up a naive reading suggests, so there is no loss spike — but the per-coordinate step sizes are now uniform instead of scaled by 1/√v, which is the scaling the pre-checkpoint run had. The behavior changes until the moments re-estimate over roughly 1/(1-β₂) = 1000 steps, and the parameters are displaced off their trajectory in the meantime. To eliminate: save optimizer.state_dict() (and, for a full resume, scheduler.state_dict() and the RNG state). The lesson: a model-only checkpoint is a model snapshot, not a training snapshot; the optimizer’s state is part of the training, not the model — and the symptom is a change in step sizes, not a loss spike.

E 0.1.3.5 — The collate (constructive, transfer probe). A Dataset returns (obs, action, mask) where obs is (T, 6), action is (T, 6), and mask is (T,), with T varying per example. Write the collate_fn that produces (obs_b, action_b, mask_b) with T_max padding and the mask’s 1/0 convention, and state the shapes.

S 0.1.3.5.

def collate_fn(batch):
    Ts = [b[0].shape[0] for b in batch]
    T_max = max(Ts)
    B = len(batch)
    obs_b = torch.zeros(B, T_max, 6)
    action_b = torch.zeros(B, T_max, 6)
    mask_b = torch.zeros(B, T_max, dtype=torch.bool)
    for i, (o, a, m) in enumerate(batch):
        T = o.shape[0]
        obs_b[i, :T] = o
        action_b[i, :T] = a
        mask_b[i, :T] = m
    return obs_b, action_b, mask_b

Shapes: obs_b (B, T_max, 6), action_b (B, T_max, 6), mask_b (B, T_max). The loss is computed as (loss_per_step * mask_b).sum() / mask_b.sum(), so the padded positions contribute nothing. The lesson: the collate_fn is where variable-shape data becomes a fixed-shape batch, and the mask is what distinguishes padding from data; without the mask, the padded zeros are trained on, and the gradient is dominated by the padding.

E 0.1.3.6 — The shuffle and correlation (diagnostic, transfer probe). A learner reports that their IK training with shuffle=False achieves a lower training loss and a higher validation loss than the same training with shuffle=True, and concludes that shuffling prevents overfitting. The lab’s three runs measured the opposite. State what the lab actually found, what the learner’s report would require, and how to tell the two apart from the run table alone.

S 0.1.3.6. The lab’s three runs, all with the same 6,400 pairs, seed 0, and 40 epochs:

Run 2 against run 1 is the flag’s own effect: +14% on training and +10% on validation, because a random-order training set is already decorrelated and disabling shuffling changes little. Run 3 against run 2 is the order’s effect: 4.29× on training and 4.22× on validation — the loss rises on both, and the validation/training ratio barely moves (1.080 → 1.061). That is the signature of underfitting, not overfitting: consecutive batches from one band of the joint space give each step a gradient biased toward that band, the model chases the current region and moves away from the previous one, and after one pass it has fitted none of them well. The learner’s reported signature — train down, val up — would require the model to fit the sorted regions and lose generalization, which happens when the model has enough capacity and enough epochs to memorize each band before moving on. Both signatures are real; they are distinguished by the direction of the training loss, which the learner did not report and the lab does measure.

The discriminator is therefore not the validation number alone but the pair. A learner who reports only “validation got worse with shuffle=False” has not said which failure they have, and the fix differs: underfitting from order calls for shuffling (or, equivalently, a sampler that interleaves regions); memorization from order calls for the same shuffling but for a different reason, and would additionally show a widening train-val gap. The lesson: shuffle=True makes each batch a random mixture so that each step’s gradient is an unbiased estimate of the full-dataset gradient; without it the estimate is biased toward whatever region the dataset’s order happens to visit, and whether that shows up as underfitting or as memorization depends on capacity and epoch count, not on the flag.

E 0.1.3.7 — The zero_grad placement (constructive, transfer probe). A learner’s PPO inner loop performs 4 epochs over the rollout, with zero_grad once per outer iteration (before the 4 epochs). The effective learning rate is lr; the intended is lr. What is the actual effective learning rate, and what is the fix?

S 0.1.3.7. The zero_grad once per outer iteration means the 4 epochs’ gradients accumulate and the optimizer’s update at the end of the 4 epochs is based on the sum of 4 epochs’ gradients. The effective learning rate is 4 × lr, not lr — a 4× increase. The fix is zero_grad inside the epoch loop (once per epoch, before the epoch’s minibatches), so each epoch’s update is independent. The standard PPO pattern is zero_grad once per minibatch inside the epoch loop, and the effective learning rate is then the intended one. The lesson: the zero_grad placement defines the unit of accumulation, and the unit’s size is a learning-rate multiplier in disguise.

E 0.1.3.8 — The buffer-vs-parameter (diagnostic, transfer probe). A learner’s model uses a running_mean buffer for input standardization. They train the model, save a checkpoint, reload it, and the reloaded model’s predictions are systematically worse by a small amount. What is the most likely cause, and what would have prevented it?

S 0.1.3.8. The running_mean is a buffer and is in state_dict() if it was registered with register_buffer; the reload restores it. If the learner used a plain attribute (self.running_mean = torch.zeros(3)), the buffer is not in state_dict(), the reload gives the constructor’s initial zeros(3) instead of the trained mean, and the standardization is systematically wrong — a small but consistent degradation. The fix is self.register_buffer("running_mean", torch.zeros(3)), which puts it in state_dict() and makes the checkpoint complete. The lesson: the buffer-vs-plain-attribute distinction is the difference between a checkpoint that restores the model’s behavior and one that restores only its parameters; the failure is silent and the symptom is a small systematic degradation after reload.

§L.9 — Exit test

You are ready for Lesson 0.1.4 when you can, without notes:

1. Name the four kinds of module state, and state where each appears in parameters() and state_dict().
2. Classify a given attribute as a parameter, a buffer, a submodule, or a plain attribute, and state whether each is moved by .to(device).
3. Implement a custom nn.Module with a parameter, a buffer, and a submodule, and verify the registration with list(model.parameters()).
4. State why the optimizer’s parameter set is a snapshot at construction, and name the two ways to add a parameter after construction.
5. Construct a multi-group optimizer for fine-tuning, with a frozen backbone at 1e-5 and a head at 1e-3, and state why the frozen backbone must be in the optimizer from the start.
6. State what each of the four checkpoint components is for, and what breaks if each is omitted.
7. Write a Dataset whose __getitem__ returns a copy rather than a view, and state why the copy matters for augmentation.
8. Write a custom collate_fn for variable-length trajectories, with the padding and the mask, and state how the mask is used in the loss.
9. State which argument owns batch order and which owns the per-worker augmentation stream, and state what num_workers does and does not change.
10. Write the canonical training loop with the correct ordering, and state the purpose of each of the six lines.
11. Diagnose a flat loss as a zero_grad between backward and step, and a fast-then-worse loss as a zero_grad once per epoch.
12. State the scheduler.step() placement and the meaning of T_max in the two placements (once per epoch vs. once per minibatch).

If any of items 1–12 is not answerable without notes, return to the section that introduced it: item 1 to §L.1.3, item 2 to §L.1.3, item 3 to §L.1.2, item 4 to §L.2.2, item 5 to §L.2.2, item 6 to §L.2.4, item 7 to §L.3.2, item 8 to E 0.1.3.5, item 9 to §L.3.3, item 10 to §L.4.1, item 11 to §L.4.3, and item 12 to §L.4.2.

§L.10 — Provenance and evidence tiers

What this lesson reproduces. Every hypothesis in §L.6.5 is executed by the lab in §L.6.9; the numbers in the text are the lab’s. Tier A, status R.

What this lesson does not reproduce. The claim that Adam’s moment state is what makes the resume spike larger than SGD’s is a consequence of the Adam update rule, which the lesson states but does not measure across optimizers. It is tier B, status S: the update rule is standard, the comparative spike is not measured here. The lesson’s own measurement is that the spike exists and that it disappears when the optimizer state is restored — that is tier A and it is enough for the design principle.

The num_workers seed claim is scoped to what was run. The lesson asserts that two num_workers=8 runs with the same seed reproduce both the batch order and the augmentations. That is a claim about the worker seeding behavior of a specific PyTorch version, and it is tier B, status P for any version the reader is not running. A reader who sees divergence has found a version difference, not an error in their loop; the diagnostic is to set num_workers=0 and check whether the divergence disappears.

No external sources are cited. The lesson is about the behavior of a library the reader runs, and every claim is either a definition from the library’s own API or a measurement the lab performs.