# 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/`