# training/ — General BC Training

## Purpose
Core training pipeline for PARA. Ported and cleaned from the prototype in `3dkeygrip/volume_dino_tracks/`.

## Entry Points
- `train.py` — main training script
- `model.py` — `TrajectoryHeatmapPredictor` model definition
- `data.py` — `RealTrajectoryDataset` dataset
- `utils.py` — geometry helpers (3D recovery from 2D + height)

## Training Loop Notes
- Loss = CE on 2D pixel (flatten H×W, use GT pixel as target) + CE on height bin + CE on gripper bin
- Gripper/rotation head uses **teacher forcing**: supervised at GT pixel location during training, decoded at argmax pixel during inference
- Height range (MIN_HEIGHT / MAX_HEIGHT) is computed from dataset statistics — update in `model.py` before training
- Backbone can be frozen (`freeze_backbone=True`) for fast iteration

## Running

```bash
# Mac (MPS, small batch for debugging)
python training/train.py --dataset_root /path/to/data --batch_size 2

# Lab server (CUDA, full run)
python training/train.py --dataset_root /path/to/data --batch_size 16
```

## Known Issues / Notes from Prototype
- MIN_HEIGHT / MAX_HEIGHT currently hardcoded to same value (0.043347) in prototype — needs proper range from data stats
- `TkAgg` matplotlib backend used in prototype — may need to change to `Agg` on headless server
- DINO weights path is hardcoded in prototype model.py — parameterize via config