# YAM — training

How the single-arm + wrist model trains. Architecture, data pipeline,
wandb, background-launch pattern.

## Architecture (current: `yam_wrist_xattn_v1`)

`DinoVolumeQueryBimanual` (in `~/yam_para/code/model_yam_single.py`):

- **Backbone**: DINOv3 ViT-S/16+, `embed_dim=384`. Loaded from
  `~/yam_para/dinov3/weights/dinov3_vits16plus_pretrain_lvd1689m-4057cbaa.pth`.
- **Input**: 504×504 RGB (`IMG_SIZE = 504`). Two views (scene + wrist).
- **Patch grid**: 56×56 (`PRED_SIZE = 56`).
- **CLS concat**: scene-CLS ⊕ wrist-CLS into the input projection
  (instead of a single-view CLS).
- **Cross-attention** (when `n_xattn_layers > 0`): bidirectional
  `CrossAttentionBlock` — scene attends to wrist, then wrist attends to
  updated scene. Repeated `n_xattn_layers` times.
- **Volume head**: factorized YX × Z × T (per timestep + arm). Voxel
  bins via `bin_centers_z` (Z = height bins).
- **Rotation head**: K-means quaternion centroids (`K = 64`).
- **Wrist abstain**: `F_oof_token = nn.Parameter(torch.randn(d_feat) *
  0.02)` appended to wrist logits as an out-of-frustum class. CE target
  = `Pw * Pw` (abstain_idx) when ground-truth lifts off-image.

Cross-attention block sketch (full code in `model_yam_single.py`):

```python
class CrossAttentionBlock(nn.Module):
    def forward(self, x_s, x_w):
        # scene attends to wrist
        a_s, _ = self.attn_s(self.norm_q_s(x_s),
                             self.norm_kv_s(x_w),
                             self.norm_kv_s(x_w))
        x_s = x_s + a_s
        x_s = x_s + self.mlp_s(self.norm2_s(x_s))
        # wrist attends to UPDATED scene
        a_w, _ = self.attn_w(self.norm_q_w(x_w),
                             self.norm_kv_w(x_s),
                             self.norm_kv_w(x_s))
        x_w = x_w + a_w
        x_w = x_w + self.mlp_w(self.norm2_w(x_w))
        return x_s, x_w
```

## Data pipeline

`~/yam_para/code/data_yam_single.py`:

- Loads scene + wrist PNG per frame from `place_mug_wrist_cam/<ep>/<cam>/png/`.
- Loads per-frame T_wrist_in_world from `lowdim/<frame>.pkl`.
- Computes per-frame wrist projection lazily in `__getitem__`
  (`_project_into_wrist_at_cur`). Cache the K_wrist + T_world2wrist
  factorization across frames in the same episode.
- Abstain target = `Pw * Pw` when the GT 3D point projects outside
  `(0..Pw) × (0..Pw)`.

PKL layout: see `conversion.md`. Image norm = ImageNet mean/std on
[0,1] inside the dataset loader (never re-normalize in the model — see
`vault/para/memory.md`).

## Launch — background, not in the agent tab

**Default: run training in background via `nohup`**, not live in a tmux
pane. The agent's tmux panes are for short-lived monitoring; long
training pollutes scrollback and risks loss on disconnect.

```bash
ssh dev   # puget
cd ~/yam_para/code

nohup \
  CUDA_VISIBLE_DEVICES=1 \
  DINO_REPO_DIR=$HOME/yam_para/dinov3 \
  DINO_WEIGHTS_PATH=$HOME/yam_para/dinov3/weights/dinov3_vits16plus_pretrain_lvd1689m-4057cbaa.pth \
  python train_yam_single.py \
    --run_name yam_wrist_xattn_v2 \
    --data_root ~/yam_para/data/place_mug_wrist_cam \
    --epochs 200 \
    --batch_size 24 \
    --use_wrist 1 \
    --n_xattn_layers 2 \
    --xattn_heads 4 \
    --wrist_loss_weight 1.0 \
    --vis_every_steps 1000 \
  > /tmp/train_xattn_v2.log 2>&1 &
echo $! > /tmp/train_xattn_v2.pid
```

Check progress without consuming context:

```bash
tail -n 40 /tmp/train_xattn_v2.log
ps -p $(cat /tmp/train_xattn_v2.pid)
```

GPU assignment: puget has 2 GPUs. CUDA_VISIBLE_DEVICES=1 (the second
GPU) is the convention for background training to leave GPU 0 free for
interactive work / deploy.

## Wandb

Project: `yam_single`. Each run uses `--run_name <name>`.

### Logged scalars (per-step / per-epoch)

| Name | Meaning |
|---|---|
| `loss/total` | Sum of all term losses |
| `loss/yx_scene` | CE on scene pixel-grid |
| `loss/z` | CE on height bins |
| `loss/t` | CE on per-timestep index |
| `loss/rot` | CE on K-rot bins |
| `loss/gripper` | CE on gripper |
| `loss/yx_wrist` | CE on wrist pixel-grid (only when `--use_wrist 1`) |
| `wrist/gt_abstain_frac` | Fraction of GTs that abstain (project off-image) |
| `wrist/pred_abstain_frac` | Fraction of preds choosing abstain |
| `wrist/in_frust_acc` | Top-1 acc on the in-frustum subset |
| `acc/yx_top1`, `acc/z_top1`, `acc/t_top1`, ... | Per-head top-1 |

### Logged viz panels (every `--vis_every_steps`)

Defaults `--vis_every_steps 1000` (every ~1k steps). Set to `100000` to
effectively disable.

| Panel | What it shows |
|---|---|
| `viz/scene_keypoints` | Scene RGB with 30 rainbow KP + line strip overlaid |
| `viz/wrist_keypoints` | Wrist RGB with 30 KP + line; abstain frames marked |
| `viz/scene_heatmap` | Per-T marginal YX heatmap (scene) |
| `viz/wrist_heatmap` | Per-T marginal YX heatmap (wrist) with abstain prob bar |
| `viz/dino_pca_scene` | First-3-PC PCA of scene patch features as RGB |
| `viz/dino_pca_wrist` | First-3-PC PCA of wrist patch features |

The keypoint + heatmap code is shared with deploy's rerun panels — see
`rerun-panels.md`.

## Checkpoints

`~/yam_para/checkpoints/<run_name>/`:

- `latest.pth` — always the most recent (atomic rename on save).
- `epoch_N.pth` — every `--checkpoint_every_epochs` (default 10).
- `best.pth` — **not used for inference**. Per `vault/para/memory.md`,
  val loss is multimodal-noisy on this family; latest > best.

## Current runs of note

| Run | Notes | Status |
|---|---|---|
| `yam_wrist_xattn_v1` | 12 h budget, 2 xattn layers, used in current deploys | Finished (epoch 138-200ish) |
| earlier: `yam_single_v*` | Pre-wrist baselines | Archived |

## 2026-06-10 canonical experiments — volume_based × baseline

The pickplace_redo_slow comparison report (see `evaluation.md`). All
commands run on puget; checkpoints land in
`~/yam_para/checkpoints/<run_name>/`.

### Volume-based (canonical) — train + finetune

In-distribution train (~80 min):

```bash
cd ~/yam_para && source .venv/bin/activate
DINO_REPO_DIR=$HOME/yam_para/dinov3 \
DINO_WEIGHTS_PATH=$HOME/yam_para/dinov3/weights/dinov3_vits16plus_pretrain_lvd1689m-4057cbaa.pth \
CUDA_VISIBLE_DEVICES=0 \
python ~/lab/para/robot/yam/train.py \
  --name pickplace_redo_slow_pastn8 \
  --data_root ~/yam_para/data/pickplace_redo_slow \
  --multiview scene,wrist \
  --temporal_head flat_mlp --temporal_mlp_hidden 512 \
  --mv_past_n 8 --mv_past_enc_dim 64 \
  --img_size 504 --pred_size 128 \
  --n_height_bins 64 --feat_dim 32 --height_enc_dim 16 --voxel_mlp_hidden 64 \
  --n_window 32 --frame_stride 4 \
  --batch_size 2 --num_workers 4 \
  --lr 3e-4 --lr_min 1e-5 --backbone_lr_mult 0.1 --lr_warmup 500 --lr_schedule cosine \
  --max_iters 10000 --vis_every_iters 500 --ckpt_every_iters 200
```

5-ep OOD finetune (~16 min, 2k iters at lr=1e-4):

```bash
DINO_REPO_DIR=... DINO_WEIGHTS_PATH=... CUDA_VISIBLE_DEVICES=0 \
python ~/lab/para/robot/yam/train.py \
  --name pickplace_redo_slow_pastn8_finetune_ood \
  --data_root ~/yam_para/data/5_ep_ood_finetune_collection \
  --init_from ~/yam_para/checkpoints/pickplace_redo_slow_pastn8/latest.pth \
  --multiview scene,wrist \
  --temporal_head flat_mlp --temporal_mlp_hidden 512 \
  --mv_past_n 8 --mv_past_enc_dim 64 \
  --img_size 504 --pred_size 128 \
  --n_height_bins 64 --feat_dim 32 --height_enc_dim 16 --voxel_mlp_hidden 64 \
  --n_window 32 --frame_stride 4 \
  --batch_size 2 --num_workers 4 \
  --lr 1e-4 --lr_min 1e-5 --backbone_lr_mult 0.1 --lr_warmup 100 --lr_schedule cosine \
  --max_iters 2000 --vis_every_iters 200 --ckpt_every_iters 100
```

### Multi-view baseline — train + finetune

Same encoder (DINOv3 + DA3-style cross-view attn + concat-CLS) as
volume_based, but the spatial head is replaced by a flat (T, 3) XYZ
regression. Run via `train_baseline.py`.

```bash
CUDA_VISIBLE_DEVICES=0 \
python ~/lab/para/robot/yam/train_baseline.py \
  --name pickplace_redo_slow_baseline_xattn4 \
  --data_root ~/yam_para/data/pickplace_redo_slow \
  --multiview scene,wrist --cross_view_layers 4 \
  --img_size 504 --n_window 32 --frame_stride 4 \
  --batch_size 8 --num_workers 4 \
  --lr 3e-4 --lr_min 1e-5 --backbone_lr_mult 0.1 --lr_warmup 500 --lr_schedule cosine \
  --max_iters 10000 --vis_every_iters 500 --ckpt_every_iters 200 \
  --xyz_weight 10.0
```

Followed by the 5-ep OOD finetune (auto-chained today via a poll loop;
manual: same args + `--init_from ...baseline_xattn4/latest.pth
--data_root ...5_ep_ood_finetune_collection --max_iters 2000 --lr 1e-4`).

### Deploy + annotate

Every model is rolled out via
`deploy_with_action_chunking.py` with the eval-name labels that the
report generator picks up:

| model + condition | `--ckpt` | `--eval_name` |
|---|---|---|
| volume_based × in_dist        | `pickplace_redo_slow_pastn8/latest.pth` | `pickplace_redo_slow_pastn8_indist` |
| volume_based × ood_objpos     | `pickplace_redo_slow_pastn8/latest.pth` | `pickplace_redo_slow_pastn8_ood_objpos` |
| volume_based × ood_finetuned  | `pickplace_redo_slow_pastn8_finetune_ood/latest.pth` | `pickplace_redo_slow_pastn8_finetuned_ood` |
| baseline × in_dist            | `pickplace_redo_slow_baseline_xattn4/latest.pth` | `baseline_in_dist` |
| baseline × ood_objpos         | `pickplace_redo_slow_baseline_xattn4/latest.pth` | `baseline_ood` |
| baseline × ood_finetuned      | `pickplace_redo_slow_baseline_xattn4_finetune_ood/latest.pth` | `baseline_finetuned_ood` |

Note `volume_based × in_dist` and `ood_objpos` both use the SAME ckpt —
only the cup position on the table differs. Same for `baseline × in_dist`
vs `ood_objpos`.

Regenerate the report after each annotation pass:

```bash
cd /data/cameron/para/robot/yam
python code/eval_rollout_pages_volume_based.py
```

## Common gotchas

- `--vis_every_steps=0` triggers a zero-div bug; use a large number
  (e.g. `100000`) to effectively disable. See `vault/para/memory.md`.
- If `loss/yx_wrist` is collapsing to ~`log(Pw*Pw + 1)`, the model is
  predicting abstain for everything — usually means the wrist GT
  projection logic disagrees with the live recal, or the OOF token
  weight is too small. Verify with `wrist/gt_abstain_frac` (should be
  ~10-30% on healthy data).
- DINOv3 weights path: must point at the **plus** variant
  (`vits16plus`), not stock `vits16`.