| # Stack-Can: Data Collection + Training Pipeline |
|
|
| End-to-end recipe for the new `robot_new/Stack-Can` Stack-Can experiments |
| (50 HF success demos + 25 BC-failure demos at 384×384). Covers: |
|
|
| 1. Generating training videos |
| 2. Building the BC DP zarr + training BC + sim eval |
| 3. Building the WM training CSV + LoRA fine-tuning Wan2.2-TI2V-5B |
| 4. Building the 3pointr tracking dataset (SAM3 + CoTracker3 + MoGe2) |
| 5. Training the **history-conditioned** 3pointr flow predictor |
| (input: 9-frame history; output: 48-step future; matches WM horizon) |
|
|
| All paths assume the repo layout under `/home/hongyi/PWM/external/`. Conda |
| envs used: `pwm` (collection / WM / build), `pointr` (3pointr training). |
|
|
| --- |
|
|
| ## 0. Conventions |
|
|
| - Success videos: `data/robot_new/Stack-Can/tableHF_ep000-049.{mp4,*_pred.mp4,*_actions.npz}` |
| built by `replay_hdf5_to_diffsynth.py` from the public HF EgoVLA-Humanoid-Sim |
| hdf5s. Resolution 384×384, room=1, table=1. |
| - Failure videos: `data/robot_new/Stack-Can/table2_failure_ep085-109.{mp4,*_pred.mp4,*_actions.npz}` |
| collected by `EgoVLA_Release/collect_failures_random_384.py` running BC with |
| action-noise injection in the same env. 384×384. |
| - 38-d action format throughout: concat(left_ee[7], right_ee[7], left_qpos[12], right_qpos[12]). |
| 50-d proprio is the full hdf5 `observations/qpos`. |
| - 3pointr per-episode timestep_K layout (compatible with `3PoinTr_repo`): |
| `K_pos.pth` (N,3) cam-frame xyz, `K_cum_flow.pth` (N,3), `K_visibility.pth` (N,) bool, |
| `K_rgb.png`, `K_masked_rgb.png`, plus `metadata.json` with intrinsics. |
|
|
| --- |
|
|
| ## 1. Generating training videos |
|
|
| ### 1.1 Success demos from HF hdf5s |
|
|
| ```bash |
| cd /home/hongyi/PWM/EgoVLA_Release |
| /home/hongyi/miniconda3/envs/pwm/bin/python replay_hdf5_to_diffsynth.py \ |
| --hdf5_dir /data0/hongyi/EgoVLA_SIM/Stack-Can \ |
| --task Humanoid-Stack-Can-v0 \ |
| --out_dir /home/hongyi/PWM/external/DiffSynth-Studio/data/robot_new/Stack-Can \ |
| --start 0 --end 50 \ |
| --rand_offset 0 |
| ``` |
|
|
| Per-episode produces `tableHF_ep{N:03d}.mp4` (RGB, dumped from hdf5 `observations/images/main`), |
| `tableHF_ep{N:03d}_pred.mp4` (sim-rendered hand-only, stretched to 384×384), |
| and `tableHF_ep{N:03d}_actions.npz` (multi-step EE + arm/hand qpos + 38-d robot_qpos). |
| |
| ### 1.2 Failure demos (BC + action noise) |
| |
| ```bash |
| cd /home/hongyi/PWM/EgoVLA_Release |
| /home/hongyi/miniconda3/envs/pwm/bin/python collect_failures_random_384.py \ |
| --task Humanoid-Stack-Can-v0 \ |
| --ckpt <BC_CKPT.ckpt> \ |
| --start_ep 85 --end_ep 110 \ |
| --noise_scale 0.05 \ |
| --rand_offset 4000 |
| ``` |
| |
| Same env (room=1, table=1), bottle replaced with orange-soda Fanta. |
| Outputs into `data/robot_new/Stack-Can/table2_failure_ep*.mp4` etc. |
|
|
| --- |
|
|
| ## 2. BC Diffusion-Policy |
|
|
| ### 2.1 Build the zarr from the first 50 success demos |
|
|
| ```bash |
| cd /home/hongyi/PWM/EgoVLA_Release |
| /home/hongyi/miniconda3/envs/pwm/bin/python dp_results/data/build_robot_new_zarr.py \ |
| --src_dir /home/hongyi/PWM/external/DiffSynth-Studio/data/robot_new/Stack-Can \ |
| --pattern "tableHF_ep0[0-4]?.mp4" \ |
| --out /home/hongyi/PWM/policy_rl/data/stack_can_robot_new_50demos.zarr |
| ``` |
|
|
| ### 2.2 Train BC (DP, h48) |
|
|
| `config/dp_robot_new_stackcan_h48.yaml` sets horizon=48, n_obs_steps=6, |
| n_action_steps=24. Output: `dp_results/runs/stack_can_bc_50demos_robot_new/`. |
|
|
| ### 2.3 Eval in sim (in-distribution + OOD) |
|
|
| ```bash |
| # Case A: room=1, table=1 (in-distribution) |
| /home/hongyi/miniconda3/envs/pwm/bin/python policy_rl/eval_in_sim.py \ |
| --task Humanoid-Stack-Can-v0 \ |
| --ckpt dp_results/runs/.../latest.ckpt \ |
| --num_episodes 20 --n_action_exec 24 \ |
| --room_idx 1 --table_idx 1 \ |
| --rand_offset 4000 \ |
| --output_dir dp_results/eval_out/.../Case_A_r1t1 \ |
| --save_video |
| |
| # Case B: room=2, table=2 (OOD) |
| ... --room_idx 2 --table_idx 2 \ |
| --output_dir dp_results/eval_out/.../Case_B_r2t2 ... |
| ``` |
|
|
| `eval_in_sim.py` accepts `--room_idx`/`--table_idx` (env-randomization handles |
| appearance changes) and uses sequential `randomize_idx` starting at `rand_offset`. |
|
|
| --- |
|
|
| ## 3. WM (Wan2.2-TI2V-5B LoRA) |
|
|
| ### 3.1 Build the metadata CSV |
|
|
| CSV format: `video,pred_video,prompt`. Paths are relative to `--dataset_base_path` |
| (set to `data/robot_new`). |
|
|
| ```bash |
| cd /home/hongyi/PWM/external/DiffSynth-Studio |
| /home/hongyi/miniconda3/envs/pwm/bin/python <<'PY' |
| import csv |
| prompt = "A humanoid robot stacks a can with its dexterous hands." |
| with open("data/robot_new_stackcan_50s_25f.csv", "w") as f: |
| w = csv.writer(f); w.writerow(["video","pred_video","prompt"]) |
| for i in range(50): |
| w.writerow([f"Stack-Can/tableHF_ep{i:03d}.mp4", |
| f"Stack-Can/tableHF_ep{i:03d}_pred.mp4", prompt]) |
| for i in range(85, 110): |
| w.writerow([f"Stack-Can/table2_failure_ep{i:03d}.mp4", |
| f"Stack-Can/table2_failure_ep{i:03d}_pred.mp4", prompt]) |
| PY |
| ``` |
|
|
| ### 3.2 Train the LoRA on GPUs 0,1,2 |
|
|
| ```bash |
| cd /home/hongyi/PWM/external/DiffSynth-Studio |
| nohup env CUDA_VISIBLE_DEVICES=0,1,2 NGPU=3 \ |
| DATASET_BASE=data/robot_new \ |
| DATASET_META=data/robot_new_stackcan_50s_25f.csv \ |
| OUTPUT_PATH=./models/train/Wan2.2-TI2V-5B-Robot_new_stackcan_50s_25f_lora \ |
| EPOCHS=50 \ |
| bash train_robot.sh > logs/wm_robot_new_stackcan_50s_25f.log 2>&1 & |
| ``` |
|
|
| Defaults from `train_robot.sh`: 320×512×49f, lr 1e-4, LoRA rank 32, save every 1000 steps. |
|
|
| --- |
|
|
| ## 4. 3pointr tracking dataset (SAM3 + CoTracker3 + MoGe2) |
|
|
| Each video is processed: |
| SAM3 text prompt → propagate frame-0 mask → sample 512 query points → |
| CoTracker3 (offline) → MoGe2 metric depth → lift to 3D camera frame → |
| write per-timestep .pth files + RGB + masked RGB + metadata.json. |
|
|
| ### 4.1 Initial build (one shot) |
|
|
| ```bash |
| cd /home/hongyi/PWM/external |
| nohup env CUDA_VISIBLE_DEVICES=3 \ |
| bash _build_50s25f.sh > logs/build_50s25f.log 2>&1 & |
| ``` |
|
|
| `_build_50s25f.sh` runs `build_3pointr_dataset.py` twice: |
| - Phase 1 success videos → output IDs `0001..0050` with start_idx=1 |
| - Phase 2 failure videos → output IDs `0051..0075` with start_idx=51 |
|
|
| Default prompt: `"an orange soda bottle on the table"` (works for some |
| videos). Episodes where SAM3 detects 0 objects are simply skipped, leaving |
| gaps in the 0001..0075 numbering. |
|
|
| ### 4.2 Fill gaps with a different prompt |
|
|
| If SAM3 misses some videos, re-run only the missing IDs with another prompt: |
|
|
| ```bash |
| cd /home/hongyi/PWM/external |
| nohup env CUDA_VISIBLE_DEVICES=3 \ |
| /home/hongyi/miniconda3/envs/pwm/bin/python -u build_3pointr_fill_gaps.py \ |
| --prompts "a bottle can on the table" \ |
| > logs/build_50s25f_fill.log 2>&1 & |
| ``` |
|
|
| `build_3pointr_fill_gaps.py` walks IDs 1..75, skips ones already built, |
| re-runs only the gaps. You can pass multiple `--prompts P1 P2 P3` and it will |
| try each in order until one succeeds. Output mapping is preserved |
| (0001 ↔ tableHF_ep000, …, 0075 ↔ table2_failure_ep109). |
| |
| ### 4.3 Visualize first-frame points |
| |
| ```bash |
| /home/hongyi/miniconda3/envs/pwm/bin/python /home/hongyi/PWM/external/visualize_3pointr_first_frames.py \ |
| --cols 10 --panel_h 192 --panel_w 192 \ |
| --out /home/hongyi/PWM/external/3pointr_data/stack_can_50s_25f_first_frames.png |
| ``` |
| |
| Produces a 10×8 grid PNG (192×192 panels) — for each episode 0001..0075, |
| shows the first-frame RGB with the sampled query points overlaid (red = |
| visible at frame 0, gray = invisible). Missing IDs become labeled black panels. |
|
|
| --- |
|
|
| ## 5. History-conditioned 3pointr flow predictor |
|
|
| Input: per-point 9-frame history `(x, y, z, vis_bit)` (= 36 dims/point). |
| Output: 48-step future absolute positions, matching the WM horizon. |
| History-side dropouts are handled by **carry-forward fill** + **explicit |
| visibility flag** so the model can discount unreliable frames. |
|
|
| ### 5.1 Files |
|
|
| - `3PoinTr_repo/util/training_utils/arch_history.py::HistoryFlowModel` — |
| thin model: per-point MLP-embed concat-history → self-attention over N |
| point tokens → flow head outputs (B, N, 3*48) → reshape to (B, 48, N, 3). |
| ~2.3M params at dim=256/depth=2. |
| - `3PoinTr_repo/scripts_training/pretrain_flow_history.py` — sliding-window |
| dataset (`HistoryWindowDataset`) + train loop. Pads short history with the |
| current frame, masks padded future via `step_mask`, applies |
| `compute_start_sphere_params` on the **current frame** for normalization. |
| - `launch_3pointr_train_h9p48.sh` — end-to-end launcher: auto-splits |
| 90/10 via `splitter.py`, then runs `pretrain_flow_history.py`. |
| |
| ### 5.2 Train |
| |
| ```bash |
| cd /home/hongyi/PWM/external |
| nohup env CUDA_VISIBLE_DEVICES=3 DEVICE=cuda:0 EPOCHS=200 BATCH_SIZE=16 STRIDE=2 \ |
| bash launch_3pointr_train_h9p48.sh > logs/train_3pointr_h9p48.log 2>&1 & |
| ``` |
| |
| Knobs (env vars): |
| - `HISTORY_LEN` (default 9) — current + (HISTORY_LEN-1) prior frames |
| - `HORIZON` (default 48) — future steps to predict |
| - `STRIDE` (default 2) — sliding-window stride over current-frame index |
| - `EPOCHS`, `BATCH_SIZE`, `DEVICE` |
| |
| Output: `3pointr_models/stack_can_50s_25f_h9p48_512pts_flow_backbone.pth`. |
| |
| ### 5.3 What's stored in the checkpoint |
| |
| ``` |
| { |
| "epoch": int, |
| "model_state_dict": <HistoryFlowModel state>, |
| "args": <vars(args)>, |
| "history_len": 9, |
| "horizon": 48, |
| "best_val_flow_epe_mm": float, |
| ... |
| } |
| ``` |
| |
| To reload: |
| |
| ```python |
| import torch |
| from util.training_utils.arch_history import HistoryFlowModel |
| ck = torch.load("stack_can_50s_25f_h9p48_512pts_flow_backbone.pth") |
| m = HistoryFlowModel(history_len=ck["history_len"], horizon=ck["horizon"], |
| dim=256, depth=2, num_heads=4) |
| m.load_state_dict(ck["model_state_dict"]) |
| ``` |
| |
| Inference batch dict (after normalizing by `compute_start_sphere_params(x_curr)`): |
| |
| ```python |
| batch = { |
| "x_hist": x_hist_n, # (B, 9, N, 3) normalized |
| "x_hist_vis": x_hist_vis, # (B, 9, N) bool |
| } |
| out = m(batch, mode="eval") |
| pred_pos = out["pred_pos"] # (B, 48, N, 3) normalized abs positions |
| ``` |
| |
| De-normalize via `sim_transform_inverse_points(pred_pos, scales=1/scales, translations=-center)`. |
| |
| --- |
| |
| ## 6. Provenance / quick links |
| |
| - WM CSV: `DiffSynth-Studio/data/robot_new_stackcan_50s_25f.csv` |
| - BC zarr: `policy_rl/data/stack_can_robot_new_50demos.zarr` |
| - 3pointr data: `external/3pointr_data/no_actions/stack_can_50s_25f/{0001..0075}/` |
| - 3pointr first-frame grid: `external/3pointr_data/stack_can_50s_25f_first_frames.png` |
| - 3pointr split: `external/3pointr_data/stack_can_50s_25f_flow.json` |
| - 3pointr models: `external/3pointr_models/stack_can_50s_25f_h9p48_*` |
| - WM LoRA: `external/DiffSynth-Studio/models/train/Wan2.2-TI2V-5B-Robot_new_stackcan_50s_25f_lora/` |
|
|
