---
phase: 03-training-energy-tests
plan: 01
type: execute
wave: 1
depends_on: []
files_modified:
  - src/pcn_torch/energy.py
  - src/pcn_torch/trainer.py
  - pyproject.toml
autonomous: true

must_haves:
  truths:
    - "train_pcn(network, dataloader, config) trains the network and returns a TrainHistory with energy, loss, and accuracy data"
    - "All training operations run under torch.no_grad() with no autograd graph construction"
    - "Inference loop updates latents for T_infer steps using snapshot errors and local gradients"
    - "Learning loop updates weights for T_learn steps with batch-averaged (1/B) gradients via .data in-place mutation"
    - "Classification targets are one-hot encoded internally; regression targets pass through unchanged"
    - "Energy is computed per-step (rolling window), per-batch (full history), and per-epoch (full history)"
    - "test_pcn evaluates a trained network on a test DataLoader and returns accuracy/energy metrics"
    - "RichCallback provides default console progress display; TrainCallback base class enables custom hooks"
  artifacts:
    - path: "src/pcn_torch/energy.py"
      provides: "Energy computation from PCNErrors"
      exports: ["compute_energy", "compute_energy_per_layer"]
    - path: "src/pcn_torch/trainer.py"
      provides: "Training loops, config, history, callbacks"
      exports: ["TrainConfig", "TrainHistory", "EnergyHistory", "TrainCallback", "RichCallback", "train_pcn", "test_pcn"]
    - path: "pyproject.toml"
      provides: "rich>=14.0 dependency"
      contains: "rich>=14.0"
  key_links:
    - from: "src/pcn_torch/trainer.py"
      to: "src/pcn_torch/energy.py"
      via: "compute_energy(errors) called in inference and learning loops"
      pattern: "from pcn_torch\\.energy import compute_energy"
    - from: "src/pcn_torch/trainer.py"
      to: "src/pcn_torch/network.py"
      via: "model.compute_errors(), model.init_latents(), model.predict(), model.latents"
      pattern: "from pcn_torch\\.network import"
    - from: "train_pcn inference loop"
      to: "model.latents[l]"
      via: "in-place latent update: model.latents[l] -= lr_infer * grad_x"
      pattern: "model\\.latents\\[.*\\] -="
    - from: "train_pcn learning loop"
      to: "layer.weight.data"
      via: "in-place weight update: weights[l].data -= lr_learn * grad_W"
      pattern: "weights\\[.*\\]\\.data -="
---

<objective>
Implement the complete PCN training system: energy computation module, training configuration and history dataclasses, callback system with rich console default, and the two-phase training loop (inference + learning) following Algorithm 3 from arXiv:2506.06332v1.

Purpose: This is the core algorithmic implementation -- without it, the network cannot learn. The inference loop drives latent variables toward consistency, the learning loop updates weights via local Hebbian-like rules, and energy tracking provides observability.

Output: Two new modules (energy.py, trainer.py) and updated pyproject.toml with rich dependency.
</objective>

<execution_context>
@C:\Users\RDCHLEM1\.claude/get-shit-done/workflows/execute-plan.md
@C:\Users\RDCHLEM1\.claude/get-shit-done/templates/summary.md
</execution_context>

<context>
@.planning/PROJECT.md
@.planning/ROADMAP.md
@.planning/STATE.md
@.planning/phases/03-training-energy-tests/03-RESEARCH.md
@.planning/phases/02-core-model/02-01-SUMMARY.md
@src/pcn_torch/network.py
@src/pcn_torch/layers.py
@src/pcn_torch/__init__.py
@pyproject.toml
</context>

<tasks>

<task type="auto">
  <name>Task 1: Create energy computation module</name>
  <files>src/pcn_torch/energy.py</files>
  <action>
Create `src/pcn_torch/energy.py` with two functions:

1. **`compute_energy(errors: PCNErrors) -> float`**
   - Paper formula: E_batch = (1/B) * sum_b [ (1/2) * sum_l ||eps_b^(l)||^2 + (1/2) * ||eps_sup_b||^2 ]
   - Get batch size B from `errors.errors[0].shape[0]`
   - Sum 0.5 * (e.float() * e.float()).sum().item() / B for each e in errors.errors
   - If errors.supervised_error is not None, add 0.5 * (sf.float() * sf.float()).sum().item() / B
   - Cast to .float() BEFORE squaring to prevent float16 overflow under autocast
   - Return Python float

2. **`compute_energy_per_layer(errors: PCNErrors) -> list[float]`**
   - Same formula but returns a list with one value per layer
   - Each entry: 0.5 * (e.float() * e.float()).sum().item() / B
   - Append supervised error energy as the final element if present

Use `from __future__ import annotations` and TYPE_CHECKING guard for PCNErrors import (ruff TCH001 compliance). Import Tensor from torch if needed.

Do NOT use torch.norm or torch.sum -- use element-wise multiply + sum for clarity and to match the paper exactly.
  </action>
  <verify>
Run: `python -c "from pcn_torch.energy import compute_energy, compute_energy_per_layer; print('OK')"`
Run: `ruff check src/pcn_torch/energy.py`
Run: `mypy src/pcn_torch/energy.py`
  </verify>
  <done>
compute_energy returns a scalar float from PCNErrors. compute_energy_per_layer returns a list of floats. Both cast to float32 before squaring. Both pass ruff and mypy.
  </done>
</task>

<task type="auto">
  <name>Task 2: Create trainer module with config, history, callbacks, and training loops</name>
  <files>src/pcn_torch/trainer.py, pyproject.toml</files>
  <action>
**Step A: Add rich dependency to pyproject.toml**

Add `"rich>=14.0"` to the `[project] dependencies` list in pyproject.toml.

**Step B: Create `src/pcn_torch/trainer.py`** with the following components:

### TrainCallback base class

```python
class TrainCallback:
    """Base class for training callbacks. Override methods you need."""
    def on_train_start(self, config: TrainConfig, history: TrainHistory) -> None: ...
    def on_epoch_start(self, epoch: int, num_epochs: int) -> None: ...
    def on_batch_start(self, batch: int, num_batches: int) -> None: ...
    def on_inference_step(self, step: int, energy: float) -> None: ...
    def on_learning_step(self, step: int, energy: float) -> None: ...
    def on_batch_end(self, batch: int, batch_energy: float) -> None: ...
    def on_epoch_end(self, epoch: int, history: TrainHistory) -> None: ...
    def on_train_end(self, history: TrainHistory) -> None: ...
```

All methods have `pass` body (no-op defaults).

### EnergyHistory dataclass

```python
@dataclass
class EnergyHistory:
    per_batch: list[float] = field(default_factory=list)
    per_epoch: list[float] = field(default_factory=list)
```

### TrainHistory dataclass

```python
@dataclass
class TrainHistory:
    energy: EnergyHistory = field(default_factory=EnergyHistory)
    train_loss: list[float] = field(default_factory=list)
    train_accuracy: list[float] = field(default_factory=list)
    test_loss: list[float] = field(default_factory=list)
    test_accuracy: list[float] = field(default_factory=list)
    config: TrainConfig | None = None
```

### TrainConfig dataclass

```python
@dataclass
class TrainConfig:
    task: str = "classification"
    T_infer: int = 50
    T_learn: int | None = None  # None means T_learn = batch_size
    lr_infer: float = 0.05
    lr_learn: float = 0.005
    num_epochs: int = 4
    early_stop_threshold: float | None = None
    energy_window_size: int = 10
    use_mixed_precision: bool = True
    callback: TrainCallback | None = None  # None -> RichCallback
```

With `__post_init__` validation:
- task must be "classification" or "regression"
- T_infer >= 1
- T_learn is None or >= 1

### RichCallback class (subclasses TrainCallback)

Use `rich.progress.Progress` and `rich.console.Console` for:
- Epoch progress bar
- Batch progress bar (nested)
- Print epoch summary at epoch end: epoch, energy, accuracy (if classification)
- on_train_end prints final summary

Wrap `import rich` in a try/except; if rich is not available, fall back to a `_PlainCallback` that uses plain print statements. This is defensive -- rich WILL be installed (it is in dependencies), but belt-and-suspenders.

### _check_no_autograd(model) internal function

```python
def _check_no_autograd(model: PredictiveCodingNetwork) -> None:
```

- Check `torch.is_grad_enabled()` -> warn with RuntimeWarning
- Check any `lat.grad_fn is not None` for lat in model.latents -> warn with RuntimeWarning
- Called once on first batch, first inference step

### train_pcn function

```python
def train_pcn(
    model: PredictiveCodingNetwork,
    dataloader: DataLoader,
    config: TrainConfig,
) -> TrainHistory:
```

Follow the Algorithm 3 translation from RESEARCH.md exactly:

**Outer loop** (epochs):
- For each epoch 0..num_epochs-1:

**Batch loop:**
- For each (x_batch, y_batch) from dataloader:
  - B = x_batch.shape[0]
  - Flatten: x_batch = x_batch.view(B, -1).to(device)
  - Classification one-hot encoding: if task=="classification" and y_batch is 1D integer, apply F.one_hot(y_batch.long(), num_classes).float().to(device). If y_batch is already 2D float, use as-is.
  - Regression: y_batch = y_batch.to(device)
  - model.init_latents(B)
  - Build weight references: `weights = [layer.weight for layer in model.layers] + [model._readout.weight]` (with appropriate type: ignore)
  - Resolve T_learn: `T_learn = config.T_learn if config.T_learn is not None else B`

**Inference phase** (under torch.no_grad() + optional autocast):
- autocast context: `torch.amp.autocast(device_type="cuda")` if use_mixed_precision and device.type == "cuda", else `nullcontext()`
- Rolling window: `deque(maxlen=config.energy_window_size)`
- For t in range(T_infer):
  - First batch, first step: call `_check_no_autograd(model)`
  - errors = model.compute_errors(x_batch, y_batch)
  - step_energy = compute_energy(errors)
  - energy_window.append(step_energy)
  - callback.on_inference_step(t, step_energy)
  - Build extended_errors: list(errors.errors) + [errors.top_error or zeros_like(model.latents[-1])]
  - Latent update for idx in range(len(model.latents)):
    - `grad_x = extended_errors[idx + 1] - (errors.gm_errors[idx] @ weights[idx])`
    - `model.latents[idx] -= config.lr_infer * grad_x`
  - Early stopping: if threshold set and |window[-1] - window[-2]| < threshold, break

**Learning phase** (under torch.no_grad(), NO autocast):
- states = [x_batch, *model.latents]
- For t in range(T_learn):
  - errors = model.compute_errors(x_batch, y_batch)
  - step_energy = compute_energy(errors)
  - callback.on_learning_step(t, step_energy)
  - Weight update for idx in range(len(model.layers)):
    - `grad_W = -(errors.gm_errors[idx].T @ states[idx + 1]) / B`
    - `weights[idx].data -= config.lr_learn * grad_W`
  - Readout update if errors.supervised_error is not None:
    - `grad_W_out = (errors.supervised_error.T @ model.latents[-1]) / B`
    - `weights[-1].data -= config.lr_learn * grad_W_out`

**After learning phase:**
- batch_energy = compute_energy(model.compute_errors(x_batch, y_batch))
- history.energy.per_batch.append(batch_energy)
- epoch_energy += batch_energy * B
- If classification: compute accuracy from model.predict()
- epoch_total += B
- callback.on_batch_end(batch_idx, batch_energy)

**Epoch end:**
- epoch_mean_energy = epoch_energy / epoch_total
- history.energy.per_epoch.append(epoch_mean_energy)
- history.train_loss.append(epoch_mean_energy)
- If classification: history.train_accuracy.append(epoch_correct / epoch_total)
- callback.on_epoch_end(epoch, history)

**Train end:**
- callback.on_train_end(history)
- return history

### test_pcn function

```python
def test_pcn(
    model: PredictiveCodingNetwork,
    dataloader: DataLoader,
    config: TrainConfig,
) -> dict[str, float]:
```

- Set the model to evaluation mode via its `.eval()` method
- Under torch.no_grad():
  - For each (x_batch, y_batch):
    - Flatten + device transfer (same as train_pcn)
    - One-hot encode if classification (same as train_pcn)
    - model.init_latents(B)
    - Run ONLY the inference loop (same as train_pcn inference phase, with autocast)
    - After inference: y_hat = model.predict()
    - If classification: count correct predictions (argmax comparison)
    - Accumulate total energy and count
- Return {"accuracy": ..., "energy": ...}

**CRITICAL implementation notes (from RESEARCH.md pitfalls):**
- Use `weights[idx].data -= ...` for in-place weight mutation (Pitfall 4)
- Divide weight gradients by B (batch averaging) but NOT latent gradients (Pitfall 2)
- Compute ALL errors first via compute_errors() before updating ANY latent (Pitfall 1 -- synchronous updates)
- NO autocast during learning phase (only inference) per paper reference code
- Cast to .float() in energy computation to avoid float16 overflow (Pitfall 5)
- Use variable name `idx` not `l` for loop variables (ruff E741)

**Coding standards:**
- `from __future__ import annotations`
- TYPE_CHECKING guard for DataLoader import (ruff TCH001)
- Match established patterns from network.py: type: ignore comments where needed for nn.ModuleList iteration
  </action>
  <verify>
Run: `python -c "from pcn_torch.trainer import train_pcn, test_pcn, TrainConfig, TrainHistory, EnergyHistory, TrainCallback, RichCallback; print('OK')"`
Run: `ruff check src/pcn_torch/trainer.py`
Run: `mypy src/pcn_torch/trainer.py`

Functional smoke test:
```python
python -c "
import torch
from torch.utils.data import DataLoader, TensorDataset
from pcn_torch import PredictiveCodingNetwork
from pcn_torch.trainer import train_pcn, test_pcn, TrainConfig

torch.manual_seed(42)
net = PredictiveCodingNetwork(dims=[20, 10, 5])
X = torch.randn(16, 20)
y = torch.randint(0, 5, (16,))
dl = DataLoader(TensorDataset(X, y), batch_size=8)
config = TrainConfig(T_infer=3, T_learn=2, num_epochs=1, use_mixed_precision=False)
history = train_pcn(net, dl, config)
print(f'Batches: {len(history.energy.per_batch)}')
print(f'Epochs: {len(history.energy.per_epoch)}')
print(f'Accuracy: {history.train_accuracy}')
result = test_pcn(net, dl, config)
print(f'Test: {result}')
print('SMOKE TEST PASSED')
"
```
  </verify>
  <done>
trainer.py contains TrainConfig, TrainHistory, EnergyHistory, TrainCallback, RichCallback, train_pcn, and test_pcn. train_pcn implements the full Algorithm 3 two-phase loop (inference + learning) under torch.no_grad(). test_pcn runs inference-only evaluation. pyproject.toml includes rich>=14.0. The smoke test completes without error.
  </done>
</task>

</tasks>

<verification>
1. `ruff check src/pcn_torch/energy.py src/pcn_torch/trainer.py` passes with no errors
2. `mypy src/pcn_torch/energy.py src/pcn_torch/trainer.py` passes with no errors
3. Smoke test: train_pcn trains a small network for 1 epoch, returns TrainHistory with populated energy.per_batch, energy.per_epoch, and train_accuracy
4. Smoke test: test_pcn evaluates and returns accuracy and energy metrics
5. `python -c "from pcn_torch.energy import compute_energy"` succeeds
6. `python -c "from pcn_torch.trainer import train_pcn, test_pcn, TrainConfig"` succeeds
</verification>

<success_criteria>
- energy.py exports compute_energy and compute_energy_per_layer with float32-safe computation
- trainer.py exports all 7 public names (TrainConfig, TrainHistory, EnergyHistory, TrainCallback, RichCallback, train_pcn, test_pcn)
- train_pcn implements Algorithm 3 with correct inference (latent updates) and learning (weight updates) phases
- All operations under torch.no_grad(), no autograd graph built
- Classification (one-hot) and regression both supported via config.task
- Energy tracked at per-step (rolling window), per-batch, and per-epoch granularity
- RichCallback provides default console output
- test_pcn runs inference-only evaluation
- ruff and mypy pass on both new files
</success_criteria>

<output>
After completion, create `.planning/phases/03-training-energy-tests/03-01-SUMMARY.md`
</output>
