Hugging Face's logo

hchen657Naza
/
PWM

Diffusers
Safetensors
Model card Files
xet
 Community
PWM / PWM_code /reward_model /README_stackcan_pipeline.md
hchen657Naza's picture
hchen657Naza
Add files using upload-large-folder tool
dd253d4 verified
preview code
|
Raw History Blame Contribute Delete
10.5 kB

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 

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) 

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 

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) 

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

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 

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) 

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:

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 

/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 

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:

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)):

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/