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.pydirectly — see thetrack-and-build-datasetskill.
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/data0dirs — 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/CKPTstill 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/*_boxmarkdata + 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 inexternal/because it's coupled tosam3/cotrackerthere and is referenced by 13 files (incl. the liverl/rewards/track_flow.pyreward, which subprocesses it by absolute path). Stage-1 scripts here import it viasys.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.pybuild_3pointr_fill_gaps.py(the original generic builder, superseded by the task-specificbuild_*_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 archivedbuild_3pointr_dataset.pyby its old path — they're part of the same retired ecosystem (not moved; harmless dead refs).