| # 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). |
| |
