File size: 9,045 Bytes
dd253d4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 | # 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).
|