# reward_model/scripts

Helper scripts for the **TrackFlow / 3-PoinTr point-flow reward model**. The pipeline goes:

```
demos ──► [1 track] ──► tracks_labeled.npz ──► [2 build] ──► 3-PoinTr dataset
                                                                   │
                                              [3 train] ◄──────────┘
                                                   │
                                        <task>_flow_backbone.pth
                                                   │
                            ┌──────────────────────┼───────────────────────┐
                       [4 overlay]           used by TrackFlow         [5 compare]
                   GT-vs-pred videos        reward in RL (rl/rewards/   Robometer vs
                                            track_flow.py)              TrackFlow AUC
```

**Envs:** `pwm` = tracking / overlays / scoring · `pointr` = 3-PoinTr training.
**Data** lives in `reward_model/data/` → symlink to `/data0/hongyi/reward_model_data/`
(holds all `3pointr_data_*` + `track_cache/`); **models** in `../3pointr_models/`.
Reusable how-to: `/home/hongyi/PWM/skills/reward_model/`.

---

## 1. Track demos (SAM3 + CoTracker3) → `tracks_labeled.npz`
Segment object(s) by text prompt on frame 0, propagate the mask, track a point grid. Loads
SAM3+CoTracker **once** and loops all demos. Output per demo: `tracks(T,N,2)@384`, `vis(T,N)`,
`labels(N,)` (0=obj0, 1=obj1), `n_<obj>`, plus `input.mp4` + `combined_overlay.mp4`.

| file | what it does |
|---|---|
| `build_stackcan_tracks.py` | Stack-Can: tracks **red can** + **plate** on the 100 `tableHF` demos. The most evolved tracker — squashes 1280×720→384, **grounds SAM3 on a short 16-frame clip** (full-video grounding is flaky), sweeps **prompt-variants × anchor frames [0,16,8,32]**, has the plate prompt fallback list (`white round plate`/`white circle`/`white disc`/`white object on the table`). `--demos ep002,ep007` re-tracks only those; `--stride`, `--n_overlay` flags. → `/data0/hongyi/reward_model_data/3pointr_data_stack_can_2d_canplate/` |
| `build_boxmark_tracks.py` | Push-Box: tracks **cheez-it box** + **red cross mark** on the WM-domain 384 demos (`rise_data_384_pushbox_dualmask/.../videos_small`). Simpler/earlier version of the above (no short-clip grounding / anchor sweep). → `/data0/hongyi/reward_model_data/3pointr_data_push_box_2d_boxmark/` |

**Prompt-tuning + QA helpers:**
| file | what it does |
|---|---|
| `_probe_stackcan_plate_v2.py` | Probe which SAM3 prompts recover a hard-to-detect object on the *failed* demos (tries a prompt list × anchors, reports per-prompt hit count). Use **before** re-tracking to pick prompts. |
| `_stackcan_montage.py` | Build a 10×10 grid PNG of the t=0 tracked-overlay of every demo — quick visual QA of a whole tracked set (spot wrong-object / missed detections). |

> Single-video tracking (no dataset) uses `external/track_objects_in_video.py` directly — see the
> `track-and-build-dataset` skill.

---

## 2. Build the 3-PoinTr dataset ← `tracks_labeled.npz`
Pure tensor reshaping (fast, CPU). For each demo writes the layout `pretrain_flow.py` reads:
`timestep_k/{k_pos = tracks[0], k_cum_flow = tracks[k]−tracks[0], k_visibility}.pth`, plus
`labels.pth` and `metadata.json`. Prints **`min_N`** → the `--subsample` ceiling. Drops bad demos.

| file | what it does |
|---|---|
| `build_stackcan_3pointr.py` | can+plate tracks → `/data0/hongyi/reward_model_data/3pointr_data_stack_can_canplate/no_actions/stack_can_canplate/`. Excludes `ep042` (false-positive plate). |
| `build_boxmark_3pointr.py` | box+mark tracks → `/data0/hongyi/reward_model_data/3pointr_data_push_box_boxmark/no_actions/push_box_boxmark/`. |

> Note the folder naming: **tracks** in `..._2d_boxmark` / `..._2d_canplate`; **built dataset** in
> `..._boxmark` / `..._canplate` (no `_2d`). Two different `/data0` dirs — easy to mix up.

---

## 3. Train the flow reward model
| file | what it does |
|---|---|
| `launch_3pointr_stackcan_canplate.sh` | One-shot wrapper: runs `3PoinTr_repo/scripts_training/splitter.py` (90/10 split) then `pretrain_flow.py` with the validated box+mark recipe (`subsample 150`, `--point_dim 2 --use_point_labels --no_img_cond`, 5000 ep, `flow_loss_w 10`). Env `pointr`, GPU via `GPU=`. → `../3pointr_models/stack_can_canplate_150pts_384_flow_backbone{,_final,_snap}.pth`. Copy + edit `DATA_ROOT`/`TASK`/`SUBSAMPLE` for a new task. |

---

## 4. Pred-flow overlay videos
Load a trained flow backbone, predict the point flow from frame-0 GT points (+ per-point labels for
2-object models), overlay **GREEN = GT CoTracker tracks, RED = 3-PoinTr prediction**, annotate
per-frame mean px error. Env `pwm`. `--ep epNNN` (repeatable). Output `pred_flow_<task>/pred_flow_<ep>.mp4`.

| file | model / domain |
|---|---|
| `gen_pred_flow_stackcan.py` | `stack_can_canplate_150pts_384` (can+plate, labels); predicts @384, draws on native 1280×720 |
| `gen_pred_flow_pushbox_boxmark.py` | `push_box_boxmark_150pts_384` (box+mark, labels); draws on the 384 `input.mp4`, 2× upscaled |
| `gen_pred_flow_pushbox.py` | `push_box_2d_150pts_384` (single box, no labels); native 1280×720 |
| `gen_pred_flow_opendrawer.py` | `open_drawer_2d_wb_150pts` (wooden item, single); 384 `input_short.mp4`, 2× |

Common logic: infer horizon from `point_flow_head.shape // point_dim` (don't trust `args["horizon"]`);
2-object models **require** `batch["point_labels"]`; subsample to ~150 preserving label order.

| `overlay_object_tracks_on_pred.py` | A *different* overlay: draws 3-PoinTr **object tracks onto the World-Model's `_pred.mp4` hand-prediction video** (→ `_pred_handobject.mp4`), preserving the black hand silhouettes. **3D-era** (projects 3D pos via stored intrinsics from the dataset `metadata.json`) — pairs with the older history-conditioned 3pointr, not the current 2-D models; adapt the projection to reuse with 2-D tracks. Documented in the `gen-pred-flow-overlay` skill. |

### Eval / debug tools (generic — operate on a *dataset* split, not a demo video)
| file | what it does |
|---|---|
| `3pointr_eval.py` | Evaluate a trained backbone's **flow-prediction error** on the val split (+ 1 train episode). The only flow-error metric tool. |
| `3pointr_predict_and_viz.py` | Predict + visualize trajectories on a val episode from the dataset tree. |
| `3pointr_smoke_test.py` | Load a tiny data subset + run a forward pass — quick sanity check after code changes. |
> ⚠️ Their **default** `DATA_ROOT`/`CKPT` still point at the old 512-pt Stack-Can set
> (`external/3pointr_data`, `stack_can_flowonly_512pts`). Repoint them (edit or args) for the
> current `*_canplate` / `*_boxmark` data + models before use.

---

## 5. Compare reward models
| file | what it does |
|---|---|
| `compare_robo_trackflow.py` | Score 10 success + 10 failure videos per task with **Robometer** (max success) and **TrackFlow** (`mode="progress_gated"`, distance-weighted), report success-vs-failure **ROC-AUC**. One task per process (`--task Open-Drawer\|Push-Box\|Stack-Can`, `TF_GPU=`) → parallelize across GPUs. Needs the Robometer server up (`dp_results/rl/start_robometer_server.sh`, port 8801). Writes `robo_vs_trackflow_<task>.json` (per-video rows + AUC). |

---

## Lives in `external/` (not here)
- `external/track_objects_in_video.py` — the **SAM3 + CoTracker3 tracking primitive** (segment by
  text prompt → propagate mask → track point grid). Kept in `external/` because it's coupled to
  `sam3`/`cotracker` there and is referenced by **13 files** (incl. the live `rl/rewards/track_flow.py`
  reward, which subprocesses it by absolute path). Stage-1 scripts here import it via
  `sys.path.insert(0, "/home/hongyi/PWM/external")`.

## archive/
Superseded / one-off scripts, kept for reference, not referenced by any active code:
- **old launchers/overlay/probes:** `launch_3pointr_train.sh`, `launch_3pointr_train_h9p48.sh`
  (old 512-pt launchers), `gen_box_mark_overlay.py` (GT-only overlay), `_probe_stackcan_prompts.py`,
  `_probe_stackcan_robust.py` (early prompt probes).
- **old 3D-era generic pipeline (Apr):** `build_3pointr_dataset.py` + `build_3pointr_dataset_fallback.py`
  + `build_3pointr_fill_gaps.py` (the original generic builder, superseded by the task-specific
  `build_*_3pointr.py`), `tracks_to_3pointr.py` (3D tracks→dataset converter).
- **old SAM3 tests / preprocessing:** `sam3_cotracker_sanity.py`, `sam3_test_prompts.py`,
  `stretch_robot_new_to_1280x720.py`.
- **old viz:** `visualize_3pointr_episode.py`, `visualize_3pointr_first_frames.py`, `visualize_tracks.py`.

> Note: a couple of *other* old build scripts still in `external/` (`_build_50s25f.sh`,
> `fix_metadata_intrinsics.py`) reference the archived `build_3pointr_dataset.py` by its old path —
> they're part of the same retired ecosystem (not moved; harmless dead refs).