Implement v1 training pipeline: CenterNet-style detector

Add the anchor-free, center-heatmap detector CLAUDE.md §3.3 specifies:
- targets.py: encodes ground-truth boxes into a per-class Gaussian
  heatmap + wh regression target + center mask, using the standard
  CornerNet/CenterNet gaussian-radius formulation.
- model.py: a small conv backbone (stride 4) with heatmap (sigmoid)
  and wh regression heads - exactly the two outputs CLAUDE.md
  specifies, sized with the 66ms/tick budget in mind.
- losses.py: modified focal loss (heatmap) + masked L1 (wh), combined
  with the standard CenterNet wh_weight=0.1.
- client.py / dataset.py: pull a promoted dataset version + its set's
  class list from the labeling backend and turn it into a
  torch.utils.data.Dataset, reading images from local disk (same
  machine as FRAMES_ROOT).
- train.py: wires it into a basic DataLoader -> train loop ->
  per-epoch checkpoint.

The exact loss weighting/architecture sizing is a reasonable, standard
v1 default, not a tuned final answer - CLAUDE.md's own roadmap flags
the exact formulation as still open; this is the starting point to
iterate from.

16/16 unit tests pass on CPU with synthetic data (target/model/loss
correctness, dataset-version parsing). Beyond that, ran a real
end-to-end smoke test: live labeling backend -> promoted dataset
version -> `spelunkai-train` actually training one real epoch against
it and writing a checkpoint. Not verified: multi-epoch convergence on
real data, which needs jai's GPU and real labeled frames.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

training/README.md

@@ -2,18 +2,82 @@
 
 Per-set training pipeline for the anchor-free, center-heatmap-based (CenterNet-style)
 CNN detectors — one model per label set (Enemy, Items, Traps, ...). Custom, small/
-efficient architectures sized around the overall 66ms/tick inference budget. Includes
+efficient architectures sized around the overall 66ms/tick inference budget.
-KPI evaluation (precision, recall, mAP-equivalent for the heatmap formulation) against
-held-out labeled data.
 
 See CLAUDE.md §3.3–3.4 for full requirements.
 
-**Status:** not yet implemented.
+**Status:** core v1 implemented — target encoding, model, loss, dataset loader
+(pulls a promoted dataset version from `labeling/backend`), and a training script.
+**Not yet implemented/verified:** actual multi-epoch training on a real dataset
+(needs `jai`'s GPU and real labeled data — nothing to train on yet), KPI evaluation
+(precision/recall/mAP-equivalent), and inference/export (ONNX etc. for the runtime
+loop). The exact loss weighting and architecture sizing are a reasonable v1 default,
+not a tuned final answer — CLAUDE.md's own roadmap flags the exact formulation as an
+open item; treat this as the starting point to iterate from once real training runs
+are possible.
+
+## How it works
+
+- `targets.py` — encodes ground-truth boxes into training targets: a per-class
+  Gaussian center-point heatmap (radius scaled to box size via the standard
+  CornerNet/CenterNet formula, so overlap penalties roughly track IoU) plus a
+  width/height regression target and a mask marking which pixels are actual object
+  centers.
+- `model.py` — `CenterNetDetector`: a small conv backbone (stride 4 output, matching
+  `OUTPUT_STRIDE`) with two 1x1-conv heads — a per-class heatmap (sigmoid) and a
+  width/height regression head — exactly the two outputs CLAUDE.md §3.3 specifies,
+  nothing extra (no separate offset head).
+- `losses.py` — modified focal loss for the heatmap + masked L1 loss for width/height
+  (masked to object-center pixels only), combined with the standard CenterNet
+  `wh_weight=0.1`.
+- `client.py` — pulls a promoted dataset version (`GET /dataset-versions/{id}`) and
+  its set's class list (`GET /sets/{id}`) from the labeling backend over plain HTTP
+  (stdlib `urllib`, no extra dependency for two GET requests).
+- `dataset.py` — `CenterNetDataset`: a `torch.utils.data.Dataset` that reads each
+  frame's image straight from local disk (training runs on the same machine/
+  `FRAMES_ROOT` as the labeling backend, so no need to re-download images) and
+  encodes its labels via `targets.py`.
+- `train.py` — wires it all together: fetch → `DataLoader` → train loop → checkpoint
+  (`checkpoints/last.pt`) after every epoch.
 
 ## Setup
 
 ```
 python -m venv .venv
 source .venv/bin/activate
-pip install -e .
+pip install -e ".[dev]"
 ```
+
+On a machine without an NVIDIA GPU (e.g. for running the tests), install the CPU
+build explicitly to avoid pulling a multi-GB CUDA wheel:
+
+```
+pip install torch --index-url https://download.pytorch.org/whl/cpu
+pip install -e ".[dev]"
+```
+
+On `jai`, install the CUDA build matching its driver/CUDA version instead (see
+[pytorch.org](https://pytorch.org/get-started/locally/)).
+
+## Training
+
+```
+spelunkai-train \
+  --backend-url http://127.0.0.1:8000 \
+  --set-id 1 --dataset-version-id 1 \
+  --images-root /path/to/FRAMES_ROOT \
+  --output-dir checkpoints
+```
+
+## Testing
+
+```
+pytest
+```
+
+Everything here is testable on CPU with synthetic data and doesn't need the labeling
+backend running: target encoding (peak placement, wh/mask correctness), the model's
+forward-pass output shapes, loss behavior (zero for a perfect prediction, masked
+correctly), the dataset-version JSON parsing, and `CenterNetDataset.__getitem__`
+against a tiny generated image. **Not verified here:** an actual multi-epoch training
+run converging on real data — that needs `jai`'s GPU and a real promoted dataset.

training/pyproject.toml

@@ -3,7 +3,17 @@ name = "spelunkai-training"
 version = "0.0.0"
 description = "SpelunkAI detector training pipeline (anchor-free, per-set CNNs)"
 requires-python = ">=3.10"
-dependencies = []
+dependencies = [
+    "torch>=2.0",
+    "numpy>=1.24",
+    "pillow>=10.0",
+]
+
+[project.optional-dependencies]
+dev = ["pytest>=8"]
+
+[project.scripts]
+spelunkai-train = "spelunkai_training.train:main"
 
 [build-system]
 requires = ["setuptools>=68"]

training/src/spelunkai_training/client.py

@@ -0,0 +1,85 @@
+"""Fetches a promoted dataset version (and its set's class list) from the
+labeling backend. Uses the stdlib HTTP client since this is a couple of
+simple GET requests - not worth adding a dependency for.
+"""
+from __future__ import annotations
+
+import json
+import urllib.request
+from dataclasses import dataclass
+from typing import Dict, List
+
+
+@dataclass
+class LabelData:
+    sub_class_id: int
+    x: float
+    y: float
+    width: float
+    height: float
+
+
+@dataclass
+class FrameData:
+    id: int
+    session_name: str
+    frame_index: int
+    image_path: str
+    width: int
+    height: int
+    labels: List[LabelData]
+
+
+@dataclass
+class DatasetVersionData:
+    id: int
+    set_id: int
+    name: str
+    frames: List[FrameData]
+
+
+def parse_dataset_version(payload: Dict) -> DatasetVersionData:
+    frames = []
+    for entry in payload["frames"]:
+        frame = entry["frame"]
+        labels = [
+            LabelData(sub_class_id=l["sub_class_id"], x=l["x"], y=l["y"], width=l["width"], height=l["height"])
+            for l in entry["labels"]
+        ]
+        frames.append(FrameData(
+            id=frame["id"],
+            session_name=frame["session_name"],
+            frame_index=frame["frame_index"],
+            image_path=frame["image_path"],
+            width=frame["width"],
+            height=frame["height"],
+            labels=labels,
+        ))
+    return DatasetVersionData(id=payload["id"], set_id=payload["set_id"], name=payload["name"], frames=frames)
+
+
+def build_class_mapping(sets_payload: Dict) -> Dict[int, int]:
+    """sub_class_id -> stable class index (0..N-1), ordered by sub_class_id so
+    the mapping doesn't shift between training runs on different dataset
+    versions of the same set."""
+    sub_class_ids = sorted(
+        sub_class["id"]
+        for main_class in sets_payload["main_classes"]
+        for sub_class in main_class["sub_classes"]
+    )
+    return {sub_class_id: index for index, sub_class_id in enumerate(sub_class_ids)}
+
+
+def _get_json(url: str) -> Dict:
+    with urllib.request.urlopen(url, timeout=30) as resp:
+        return json.loads(resp.read().decode("utf-8"))
+
+
+def fetch_dataset_version(base_url: str, version_id: int) -> DatasetVersionData:
+    payload = _get_json(f"{base_url.rstrip('/')}/dataset-versions/{version_id}")
+    return parse_dataset_version(payload)
+
+
+def fetch_class_mapping(base_url: str, set_id: int) -> Dict[int, int]:
+    payload = _get_json(f"{base_url.rstrip('/')}/sets/{set_id}")
+    return build_class_mapping(payload)

training/src/spelunkai_training/dataset.py

@@ -0,0 +1,64 @@
+"""PyTorch Dataset wrapping a fetched dataset version: loads each frame's
+image from local disk - training runs on the same machine/FRAMES_ROOT as the
+labeling backend, so images are read directly rather than re-downloaded over
+HTTP - and encodes its labels into CenterNet targets via `targets.py`.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Dict
+
+import numpy as np
+import torch
+from PIL import Image
+from torch.utils.data import Dataset
+
+from .client import DatasetVersionData
+from .model import OUTPUT_STRIDE
+from .targets import Box, encode_targets
+
+
+class CenterNetDataset(Dataset):
+    def __init__(
+        self,
+        dataset_version: DatasetVersionData,
+        images_root: Path,
+        class_index_by_sub_class_id: Dict[int, int],
+        output_stride: int = OUTPUT_STRIDE,
+    ):
+        self.frames = dataset_version.frames
+        self.images_root = Path(images_root)
+        self.class_index_by_sub_class_id = class_index_by_sub_class_id
+        self.output_stride = output_stride
+        self.num_classes = len(set(class_index_by_sub_class_id.values()))
+
+    def __len__(self) -> int:
+        return len(self.frames)
+
+    def __getitem__(self, index: int) -> Dict[str, torch.Tensor]:
+        frame = self.frames[index]
+        image = Image.open(self.images_root / frame.image_path).convert("RGB")
+        image_tensor = torch.from_numpy(np.asarray(image, dtype=np.float32) / 255.0).permute(2, 0, 1)
+
+        boxes = [
+            Box(
+                class_index=self.class_index_by_sub_class_id[label.sub_class_id],
+                x=label.x, y=label.y, width=label.width, height=label.height,
+            )
+            for label in frame.labels
+            if label.sub_class_id in self.class_index_by_sub_class_id
+        ]
+        targets = encode_targets(
+            boxes=boxes,
+            num_classes=self.num_classes,
+            image_width=frame.width,
+            image_height=frame.height,
+            output_stride=self.output_stride,
+        )
+
+        return {
+            "image": image_tensor,
+            "heatmap": torch.from_numpy(targets.heatmap),
+            "wh": torch.from_numpy(targets.wh),
+            "mask": torch.from_numpy(targets.mask),
+        }

training/src/spelunkai_training/losses.py

@@ -0,0 +1,51 @@
+"""Loss functions for the CenterNet-style detector: a modified focal loss
+for the heatmap (Law & Deng 2018 / Zhou et al. 2019) and a masked L1 loss
+for width/height regression, applied only at ground-truth object centers.
+"""
+from __future__ import annotations
+
+from typing import Dict, Tuple
+
+import torch
+
+
+def focal_loss(pred: torch.Tensor, target: torch.Tensor, alpha: float = 2.0, beta: float = 4.0) -> torch.Tensor:
+    """Pixel-wise modified focal loss, normalized by the number of objects (peaks == 1)."""
+    eps = 1e-6
+    pred = pred.clamp(eps, 1 - eps)
+
+    pos_mask = target.eq(1).float()
+    neg_mask = target.lt(1).float()
+    neg_weights = torch.pow(1 - target, beta)
+
+    pos_loss = torch.log(pred) * torch.pow(1 - pred, alpha) * pos_mask
+    neg_loss = torch.log(1 - pred) * torch.pow(pred, alpha) * neg_weights * neg_mask
+
+    num_pos = pos_mask.sum()
+    loss = -(pos_loss.sum() + neg_loss.sum())
+    if num_pos > 0:
+        loss = loss / num_pos
+    return loss
+
+
+def masked_l1_loss(pred: torch.Tensor, target: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+    """L1 loss over `pred`/`target` (both `(B, 2, H, W)`), counted only where `mask` (`(B, H, W)`) is 1."""
+    mask = mask.unsqueeze(1)  # (B, 1, H, W), broadcasts over the 2 wh channels
+    num_valid = mask.sum().clamp(min=1.0)
+    loss = torch.abs(pred - target) * mask
+    return loss.sum() / num_valid
+
+
+def detection_loss(
+    pred_heatmap: torch.Tensor,
+    pred_wh: torch.Tensor,
+    target_heatmap: torch.Tensor,
+    target_wh: torch.Tensor,
+    mask: torch.Tensor,
+    wh_weight: float = 0.1,
+) -> Tuple[torch.Tensor, Dict[str, float]]:
+    """Combined loss; `wh_weight=0.1` matches the standard CenterNet size-loss weighting."""
+    hm_loss = focal_loss(pred_heatmap, target_heatmap)
+    wh_loss = masked_l1_loss(pred_wh, target_wh, mask)
+    total = hm_loss + wh_weight * wh_loss
+    return total, {"heatmap_loss": hm_loss.item(), "wh_loss": wh_loss.item(), "total_loss": total.item()}

training/src/spelunkai_training/model.py

@@ -0,0 +1,69 @@
+"""A small anchor-free, center-heatmap-based detector (CenterNet-style), per
+CLAUDE.md §3.3: predicts a per-class center-point heatmap plus direct
+width/height regression from each detected center.
+
+Deliberately small/efficient - this is a v1 architecture sized with the
+66ms/tick budget in mind, given multiple such models may run per tick.
+Tune channel widths/depth from here once real timing numbers are measured
+on the target GPU (RTX 3070 Ti); the architecture itself (stride, head
+shapes) is what the rest of the pipeline depends on, not these specific
+widths.
+"""
+from __future__ import annotations
+
+from typing import Tuple
+
+import torch
+from torch import nn
+
+OUTPUT_STRIDE = 4
+
+
+class ConvBlock(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int, stride: int = 1):
+        super().__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False)
+        self.bn = nn.BatchNorm2d(out_channels)
+        self.act = nn.ReLU(inplace=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.act(self.bn(self.conv(x)))
+
+
+class CenterNetDetector(nn.Module):
+    def __init__(self, num_classes: int, base_channels: int = 32):
+        super().__init__()
+        c = base_channels
+        self.backbone = nn.Sequential(
+            ConvBlock(3, c // 2, stride=1),
+            ConvBlock(c // 2, c, stride=2),  # stride 2
+            ConvBlock(c, c * 2, stride=2),    # stride 4 (= OUTPUT_STRIDE)
+            ConvBlock(c * 2, c * 2, stride=1),
+            ConvBlock(c * 2, c * 2, stride=1),
+        )
+
+        head_channels = c * 2
+        self.heatmap_head = nn.Sequential(
+            ConvBlock(head_channels, head_channels, stride=1),
+            nn.Conv2d(head_channels, num_classes, kernel_size=1),
+        )
+        self.wh_head = nn.Sequential(
+            ConvBlock(head_channels, head_channels, stride=1),
+            nn.Conv2d(head_channels, 2, kernel_size=1),
+        )
+
+        self._init_heatmap_bias()
+
+    def _init_heatmap_bias(self) -> None:
+        # Standard CenterNet trick: bias the heatmap head's last conv so
+        # predictions start near-zero everywhere (most pixels are
+        # background) instead of ~0.5, which noticeably stabilizes early
+        # focal-loss training.
+        final_conv = self.heatmap_head[-1]
+        nn.init.constant_(final_conv.bias, -2.19)  # sigmoid(-2.19) ~= 0.1
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        features = self.backbone(x)
+        heatmap = torch.sigmoid(self.heatmap_head(features))
+        wh = self.wh_head(features)
+        return heatmap, wh

training/src/spelunkai_training/targets.py

@@ -0,0 +1,108 @@
+"""CenterNet-style target encoding: converts pixel-space bounding boxes into
+a per-class center-point heatmap plus width/height regression targets, per
+CLAUDE.md §3.3 ("a center-point heatmap per class + direct width/height
+regression from each detected center").
+
+Uses the standard CornerNet/CenterNet gaussian-radius formulation (Law &
+Deng 2018; Zhou et al. 2019), so overlapping-box penalties in the heatmap
+loss roughly track IoU instead of using an arbitrary fixed-radius blob.
+"""
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import List
+
+import numpy as np
+
+
+@dataclass
+class Box:
+    class_index: int
+    x: float
+    y: float
+    width: float
+    height: float
+
+
+@dataclass
+class Targets:
+    heatmap: np.ndarray  # (num_classes, out_h, out_w), float32, in [0, 1]
+    wh: np.ndarray        # (2, out_h, out_w), float32, pixel units (image space)
+    mask: np.ndarray      # (out_h, out_w), float32, 1 at object centers else 0
+
+
+def gaussian_radius(height: float, width: float, min_overlap: float = 0.7) -> float:
+    a1 = 1.0
+    b1 = height + width
+    c1 = width * height * (1 - min_overlap) / (1 + min_overlap)
+    r1 = (b1 + math.sqrt(b1 ** 2 - 4 * a1 * c1)) / 2
+
+    a2 = 4.0
+    b2 = 2 * (height + width)
+    c2 = (1 - min_overlap) * width * height
+    r2 = (b2 + math.sqrt(b2 ** 2 - 4 * a2 * c2)) / 2
+
+    a3 = 4.0 * min_overlap
+    b3 = -2 * min_overlap * (height + width)
+    c3 = (min_overlap - 1) * width * height
+    r3 = (b3 + math.sqrt(b3 ** 2 - 4 * a3 * c3)) / 2
+
+    return min(r1, r2, r3)
+
+
+def _gaussian_2d(diameter: int, sigma: float) -> np.ndarray:
+    radius = (diameter - 1) / 2
+    y, x = np.ogrid[-radius:radius + 1, -radius:radius + 1]
+    gaussian = np.exp(-(x * x + y * y) / (2 * sigma * sigma))
+    gaussian[gaussian < np.finfo(gaussian.dtype).eps * gaussian.max()] = 0
+    return gaussian
+
+
+def _draw_gaussian(heatmap: np.ndarray, center_x: int, center_y: int, radius: int) -> None:
+    diameter = 2 * radius + 1
+    gaussian = _gaussian_2d(diameter, sigma=diameter / 6)
+
+    out_h, out_w = heatmap.shape
+    left, right = min(center_x, radius), min(out_w - center_x, radius + 1)
+    top, bottom = min(center_y, radius), min(out_h - center_y, radius + 1)
+    if left + right <= 0 or top + bottom <= 0:
+        return
+
+    masked_heatmap = heatmap[center_y - top:center_y + bottom, center_x - left:center_x + right]
+    masked_gaussian = gaussian[radius - top:radius + bottom, radius - left:radius + right]
+    np.maximum(masked_heatmap, masked_gaussian, out=masked_heatmap)
+
+
+def encode_targets(
+    boxes: List[Box],
+    num_classes: int,
+    image_width: int,
+    image_height: int,
+    output_stride: int,
+) -> Targets:
+    out_w = image_width // output_stride
+    out_h = image_height // output_stride
+
+    heatmap = np.zeros((num_classes, out_h, out_w), dtype=np.float32)
+    wh = np.zeros((2, out_h, out_w), dtype=np.float32)
+    mask = np.zeros((out_h, out_w), dtype=np.float32)
+
+    for box in boxes:
+        if not (0 <= box.class_index < num_classes):
+            raise ValueError(f"class_index {box.class_index} out of range for num_classes={num_classes}")
+
+        center_x_out = (box.x + box.width / 2) / output_stride
+        center_y_out = (box.y + box.height / 2) / output_stride
+        cx, cy = int(center_x_out), int(center_y_out)
+        if not (0 <= cx < out_w and 0 <= cy < out_h):
+            continue  # center falls outside the feature map - shouldn't normally happen
+
+        radius = max(0, int(gaussian_radius(box.height / output_stride, box.width / output_stride)))
+        _draw_gaussian(heatmap[box.class_index], cx, cy, radius)
+
+        wh[0, cy, cx] = box.width
+        wh[1, cy, cx] = box.height
+        mask[cy, cx] = 1.0
+
+    return Targets(heatmap=heatmap, wh=wh, mask=mask)

training/src/spelunkai_training/train.py

@@ -0,0 +1,68 @@
+"""Training entry point: fetch a promoted dataset version from the labeling
+backend, train the CenterNet-style detector on it, checkpoint after every
+epoch. Meant to run on `jai` (GPU) - CPU works too, just much slower.
+"""
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import torch
+from torch.utils.data import DataLoader
+
+from .client import fetch_class_mapping, fetch_dataset_version
+from .dataset import CenterNetDataset
+from .losses import detection_loss
+from .model import CenterNetDetector
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(prog="spelunkai-train")
+    parser.add_argument("--backend-url", default="http://127.0.0.1:8000")
+    parser.add_argument("--dataset-version-id", type=int, required=True)
+    parser.add_argument("--set-id", type=int, required=True)
+    parser.add_argument(
+        "--images-root", type=Path, required=True,
+        help="Local FRAMES_ROOT the labeling backend serves images from",
+    )
+    parser.add_argument("--output-dir", type=Path, default=Path("checkpoints"))
+    parser.add_argument("--epochs", type=int, default=50)
+    parser.add_argument("--batch-size", type=int, default=8)
+    parser.add_argument("--lr", type=float, default=1e-3)
+    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    args = parser.parse_args()
+
+    class_index_by_sub_class_id = fetch_class_mapping(args.backend_url, args.set_id)
+    dataset_version = fetch_dataset_version(args.backend_url, args.dataset_version_id)
+    dataset = CenterNetDataset(dataset_version, args.images_root, class_index_by_sub_class_id)
+    loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, num_workers=2)
+
+    model = CenterNetDetector(num_classes=dataset.num_classes).to(args.device)
+    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
+
+    args.output_dir.mkdir(parents=True, exist_ok=True)
+
+    for epoch in range(args.epochs):
+        model.train()
+        epoch_loss = 0.0
+        for batch in loader:
+            images = batch["image"].to(args.device)
+            target_heatmap = batch["heatmap"].to(args.device)
+            target_wh = batch["wh"].to(args.device)
+            mask = batch["mask"].to(args.device)
+
+            pred_heatmap, pred_wh = model(images)
+            loss, _ = detection_loss(pred_heatmap, pred_wh, target_heatmap, target_wh, mask)
+
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            epoch_loss += loss.item()
+
+        avg_loss = epoch_loss / max(1, len(loader))
+        print(f"epoch {epoch + 1}/{args.epochs} - loss {avg_loss:.4f}")
+        torch.save(model.state_dict(), args.output_dir / "last.pt")
+
+
+if __name__ == "__main__":
+    main()

training/tests/test_client.py

@@ -0,0 +1,44 @@
+from spelunkai_training.client import build_class_mapping, parse_dataset_version
+
+
+def test_parse_dataset_version():
+    payload = {
+        "id": 1,
+        "set_id": 2,
+        "name": "v1",
+        "frames": [
+            {
+                "frame": {
+                    "id": 10, "session_name": "run01", "frame_index": 0,
+                    "image_path": "run01_frames/frame_000000.png", "width": 1280, "height": 720,
+                },
+                "labels": [
+                    {"sub_class_id": 5, "x": 1.0, "y": 2.0, "width": 3.0, "height": 4.0},
+                ],
+            },
+        ],
+    }
+
+    result = parse_dataset_version(payload)
+
+    assert result.id == 1
+    assert result.set_id == 2
+    assert result.name == "v1"
+    assert len(result.frames) == 1
+    frame = result.frames[0]
+    assert frame.image_path == "run01_frames/frame_000000.png"
+    assert len(frame.labels) == 1
+    assert frame.labels[0].sub_class_id == 5
+
+
+def test_build_class_mapping_is_stable_and_ordered_by_id():
+    payload = {
+        "main_classes": [
+            {"id": 1, "name": "Enemy", "sub_classes": [{"id": 5, "name": "Snake"}, {"id": 2, "name": "Bat"}]},
+            {"id": 2, "name": "Hazard", "sub_classes": [{"id": 9, "name": "Spikes"}]},
+        ],
+    }
+
+    mapping = build_class_mapping(payload)
+
+    assert mapping == {2: 0, 5: 1, 9: 2}

training/tests/test_dataset.py

@@ -0,0 +1,69 @@
+from pathlib import Path
+
+from PIL import Image
+
+from spelunkai_training.client import DatasetVersionData, FrameData, LabelData
+from spelunkai_training.dataset import CenterNetDataset
+from spelunkai_training.model import OUTPUT_STRIDE
+
+
+def _write_fake_image(path: Path, width: int, height: int) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    Image.new("RGB", (width, height), color=(10, 20, 30)).save(path)
+
+
+def test_dataset_getitem_shapes(tmp_path):
+    images_root = tmp_path / "frames"
+    _write_fake_image(images_root / "run01_frames" / "frame_000000.png", width=64, height=64)
+
+    dataset_version = DatasetVersionData(
+        id=1, set_id=1, name="v1",
+        frames=[
+            FrameData(
+                id=1, session_name="run01", frame_index=0,
+                image_path="run01_frames/frame_000000.png", width=64, height=64,
+                labels=[LabelData(sub_class_id=7, x=8, y=8, width=16, height=16)],
+            ),
+        ],
+    )
+
+    dataset = CenterNetDataset(
+        dataset_version=dataset_version,
+        images_root=images_root,
+        class_index_by_sub_class_id={7: 0},
+    )
+
+    assert len(dataset) == 1
+    sample = dataset[0]
+
+    out_size = 64 // OUTPUT_STRIDE
+    assert sample["image"].shape == (3, 64, 64)
+    assert sample["heatmap"].shape == (1, out_size, out_size)
+    assert sample["wh"].shape == (2, out_size, out_size)
+    assert sample["mask"].shape == (out_size, out_size)
+    assert sample["mask"].sum().item() == 1.0
+
+
+def test_dataset_ignores_labels_for_unknown_sub_classes(tmp_path):
+    images_root = tmp_path / "frames"
+    _write_fake_image(images_root / "run01_frames" / "frame_000000.png", width=32, height=32)
+
+    dataset_version = DatasetVersionData(
+        id=1, set_id=1, name="v1",
+        frames=[
+            FrameData(
+                id=1, session_name="run01", frame_index=0,
+                image_path="run01_frames/frame_000000.png", width=32, height=32,
+                labels=[LabelData(sub_class_id=999, x=0, y=0, width=4, height=4)],
+            ),
+        ],
+    )
+
+    dataset = CenterNetDataset(
+        dataset_version=dataset_version,
+        images_root=images_root,
+        class_index_by_sub_class_id={7: 0},
+    )
+
+    sample = dataset[0]
+    assert sample["mask"].sum().item() == 0.0

training/tests/test_losses.py

@@ -0,0 +1,58 @@
+import pytest
+import torch
+
+from spelunkai_training.losses import detection_loss, focal_loss, masked_l1_loss
+
+
+def test_focal_loss_is_zero_for_a_perfect_prediction():
+    target = torch.zeros(1, 1, 8, 8)
+    target[0, 0, 3, 3] = 1.0
+    pred = target.clone()
+
+    loss = focal_loss(pred, target)
+    assert torch.isclose(loss, torch.tensor(0.0), atol=1e-6)
+
+
+def test_focal_loss_is_positive_for_a_wrong_prediction():
+    target = torch.zeros(1, 1, 8, 8)
+    target[0, 0, 3, 3] = 1.0
+    pred = torch.full_like(target, 0.1)
+
+    loss = focal_loss(pred, target)
+    assert loss.item() > 0
+
+
+def test_masked_l1_loss_only_counts_masked_pixels():
+    pred = torch.zeros(1, 2, 4, 4)
+    target = torch.zeros(1, 2, 4, 4)
+    mask = torch.zeros(1, 4, 4)
+
+    # mismatch outside the mask - should not affect the loss
+    pred[0, :, 0, 0] = 100.0
+    loss = masked_l1_loss(pred, target, mask)
+    assert torch.isclose(loss, torch.tensor(0.0))
+
+    # mismatch inside the mask - should affect the loss. Both wh channels
+    # contribute (|5-3|=2 each), summed per masked pixel then divided by the
+    # number of masked pixels (1) -> 4, not 2.
+    mask[0, 1, 1] = 1.0
+    pred[0, :, 1, 1] = 5.0
+    target[0, :, 1, 1] = 3.0
+    loss = masked_l1_loss(pred, target, mask)
+    assert torch.isclose(loss, torch.tensor(4.0))
+
+
+def test_detection_loss_combines_both_terms():
+    heatmap_target = torch.zeros(1, 1, 8, 8)
+    heatmap_target[0, 0, 2, 2] = 1.0
+    wh_target = torch.zeros(1, 2, 8, 8)
+    wh_target[0, :, 2, 2] = 10.0
+    mask = torch.zeros(1, 8, 8)
+    mask[0, 2, 2] = 1.0
+
+    pred_heatmap = torch.full_like(heatmap_target, 0.3)
+    pred_wh = torch.zeros_like(wh_target)
+
+    total, parts = detection_loss(pred_heatmap, pred_wh, heatmap_target, wh_target, mask)
+    assert total.item() > 0
+    assert parts["total_loss"] == pytest.approx(parts["heatmap_loss"] + 0.1 * parts["wh_loss"], rel=1e-4)

training/tests/test_model.py

@@ -0,0 +1,34 @@
+import torch
+
+from spelunkai_training.model import OUTPUT_STRIDE, CenterNetDetector
+
+
+def test_forward_pass_output_shapes():
+    model = CenterNetDetector(num_classes=3, base_channels=8)
+    model.eval()
+
+    x = torch.rand(2, 3, 64, 64)
+    with torch.no_grad():
+        heatmap, wh = model(x)
+
+    expected_size = 64 // OUTPUT_STRIDE
+    assert heatmap.shape == (2, 3, expected_size, expected_size)
+    assert wh.shape == (2, 2, expected_size, expected_size)
+
+
+def test_heatmap_output_is_in_unit_range():
+    model = CenterNetDetector(num_classes=2, base_channels=8)
+    model.eval()
+
+    x = torch.rand(1, 3, 32, 32)
+    with torch.no_grad():
+        heatmap, _ = model(x)
+
+    assert heatmap.min() >= 0.0
+    assert heatmap.max() <= 1.0
+
+
+def test_real_capture_resolution_is_divisible_by_output_stride():
+    # Spelunky Classic HD's fixed capture resolution (CLAUDE.md §3.1).
+    assert 1280 % OUTPUT_STRIDE == 0
+    assert 720 % OUTPUT_STRIDE == 0

training/tests/test_targets.py

@@ -0,0 +1,43 @@
+import numpy as np
+
+from spelunkai_training.targets import Box, encode_targets, gaussian_radius
+
+
+def test_gaussian_radius_is_positive_for_a_reasonable_box():
+    assert gaussian_radius(height=20, width=15) > 0
+
+
+def test_encode_targets_places_peak_at_box_center():
+    box = Box(class_index=0, x=8, y=16, width=8, height=8)  # center = (12, 20)
+    targets = encode_targets([box], num_classes=1, image_width=64, image_height=64, output_stride=4)
+
+    # center in output space: (12/4, 20/4) = (3, 5) -> (x=3, y=5)
+    assert targets.heatmap.shape == (1, 16, 16)
+    assert targets.heatmap[0, 5, 3] == 1.0
+    assert targets.heatmap.max() == 1.0
+
+
+def test_encode_targets_sets_wh_and_mask_only_at_center():
+    box = Box(class_index=0, x=8, y=16, width=8, height=12)
+    targets = encode_targets([box], num_classes=1, image_width=64, image_height=64, output_stride=4)
+
+    assert targets.wh[0, 5, 3] == 8
+    assert targets.wh[1, 5, 3] == 12
+    assert targets.mask[5, 3] == 1.0
+    assert targets.mask.sum() == 1.0
+
+
+def test_encode_targets_rejects_out_of_range_class_index():
+    box = Box(class_index=5, x=0, y=0, width=4, height=4)
+    try:
+        encode_targets([box], num_classes=2, image_width=32, image_height=32, output_stride=4)
+        assert False, "expected ValueError"
+    except ValueError:
+        pass
+
+
+def test_encode_targets_with_no_boxes_is_all_zero():
+    targets = encode_targets([], num_classes=3, image_width=32, image_height=32, output_stride=4)
+    assert targets.heatmap.shape == (3, 8, 8)
+    assert np.all(targets.heatmap == 0)
+    assert np.all(targets.mask == 0)