EVALUATION_PIPELINE.md

# RoboPRO evaluation pipeline on GB10 — full reference (to recreate it)

This is the concrete, runnable description of the eval pipeline. On **this GB10 cluster** the whole
thing is already built and **group-readable** at `/shared_work/jack/eval_root` — same-cluster teams
can run it in place or copy it. To rebuild elsewhere, see the **Recreate checklist** at the end.

Companion docs: `GB10_EVAL_PIPELINE.md` (issues/gotchas + validated config), `ROBOPRO_EVAL_CROSSNODE_GB10.md`
(model+sim on separate nodes), `ROBOPRO_EVAL_SIM_AS_SERVER.md` (sim-as-server direction).

---

## 0. Layout
```
/shared_work/jack/eval_root
├── env.sh                       # sets EVAL_ROOT, BUNDLE, HF_HOME, XVLA_PY, ROBOTWIN_PY, sources /shared_work/jack/.env (HF token)
├── code/RoboPRO/                # the benchmark: benchmark/ (bench_task_config/*.yml + assets) + customized_robotwin/ (sim + scripts)
│   └── customized_robotwin/script/{eval_policy_client.py, policy_model_server.py}  + policy/dxvla/deploy_policy.py
├── code/DA3-XVLA/               # model code (DA3_XVLA_ROOT) — models/ etc.
├── envs/{robotwin, xvla, robotwin_b1}   # the two run envs + the SAPIEN 3.0.0b1 swap pkg
├── checkpoints/ , ablation_ckpts/       # model weights (k320, giants, geostack, ...)
├── slurm_eval_node.sh           # per-NODE worker entrypoint (self-stage -> b1 swap -> launch lease workers)
├── stage_k320_node.sh , stage_giants_node.sh   # node-local stagers
├── organize_results.py          # raw metrics -> runs/<model>/<scene>/<task>/seed<N>.mp4 + results.txt
├── orchestrate_*.sh             # multi-node drivers (clutter sweep, parallel-giants, fill-in)
└── runs/<RUN>/                  # per-run output: queue.txt, queue.txt.leases/, eval_result/, logs
/shared_work/jack/robopro-eval-pipeline   # = $BUNDLE
├── harness/eval_worker_steal_dualgpu.sh  # the per-WORKER lease loop (model server + sim client)
└── seeds/                       # pre-validated seed banks + queue files (clean + clutter d6..d15)
```

## 1. Architecture (per node)
Two processes over a localhost TCP socket (**4-byte big-endian length + JSON-with-numpy**):
- **Model server** (`policy_model_server.py`, `xvla` env) — loads the policy via `get_model()` once, listens, returns a 30-step action chunk per observation.
- **Sim client** (`eval_policy_client.py`, `robotwin` env) — runs SAPIEN + the RoboPRO task **in-process**, sends `(3-cam RGB + proprio + instruction)` each step, executes the returned actions.

The model is the server, the sim connects (default). Cross-node and sim-as-server variants are env-flag selectable — see the companion docs.

## 2. Node-local staging (the key perf trick)
NFS mmap of envs/ckpts across many nodes stalls. So every node **self-stages to `/tmp/dxvla_local`**
on first job and runs from there; only one-time bulk rsyncs touch NFS.
- `stage_k320_node.sh` — base: `envs/{xvla,robotwin}`, `code/`, the DA3-BASE HF backbone, curobo build, asset symlink fixes.
- `stage_giants_node.sh` — per-model giant checkpoints + their HF backbones; keeps only `GIANT_KEEP`, prunes the rest (cache self-bounds ~40 GB/node; node disks are 3.6 TB).
- The model code can also be read from NFS via `GIANT_CODE_DIR` (small .py imports fine over NFS; only weights need node-local).

## 3. Per-node worker: `slurm_eval_node.sh`
Driven entirely by env vars. On each node it: (1) self-stages if cold, (2) **swaps SAPIEN to 3.0.0b1**
if `DXVLA_SAPIEN_B1=1`, (3) fixes asset symlinks, (4) refresh-rsyncs `script/`+`policy/`+`envs/`+bench
code from NFS (so edits propagate without a re-stage), (5) exports the eval knobs, (6) launches
`WORKERS_PER_GPU` copies of the lease worker and `wait`s. Run it as the SLURM payload:
```bash
srun --partition=gb10 --nodelist=<nodes> --nodes=N --ntasks-per-node=1 --gres=gpu:1 \
     --cpus-per-task=18 bash /shared_work/jack/eval_root/slurm_eval_node.sh
```

## 4. Work-stealing lease queue (how N nodes cooperate + resume)
All workers (across all nodes) share **one run dir** and steal from **`runs/<RUN>/queue.txt`** (lines = `scene task`):
- `queue.txt.leases/` is a **directory of per-cell lock dirs** (hashed). A worker claims a cell by an
  **atomic `mkdir`**; it heartbeats a timestamp; a **reaper** reclaims leases with no heartbeat for
  `LEASE_TTL` (or whose owner died). → `rm -rf queue.txt.leases` to force a clean re-lease; `rm -f` does NOT clear it.
- **Resume is free**: with `EVAL_RESUME_SEEDS=1`, before running a cell the worker counts already-recorded
  seeds in `eval_result/.../_metrics.jsonl` and runs only the missing ones. Re-launching the same run dir
  skips finished work — that's how the sweeps pause/resume and how you "fill in" extra seeds (just bump
  `EVAL_TEST_NUM` and rerun).
- One worker = one persistent model server + a loop of: claim cell → run `eval_policy_client` for the
  missing seeds → mark done → next cell. `WORKERS_PER_GPU=1` (MPS is broken on Blackwell; sharing corrupts the sim).

## 5. Env-var reference
**Per-run (you set these / the orchestrator sets them):**
| var | meaning |
|---|---|
| `STUDY_RUN` | run dir name under `runs/` (shared lease queue + results) |
| `CKPT_NAME` | checkpoint dir name under `checkpoints/` or `ablation_ckpts/` |
| `GIANT_CODE_DIR` | model code dir → `DA3_XVLA_ROOT` (for ablation/giant/geostack models; omit for k320) |
| `BENCH_CONFIG` | `clean` or clutter `d6..d15` (→ task_config `bench_demo_<scene>_<CONFIG>`) |
| `SEEDS_ROOT` | seed bank dir (`USE_EVAL_SEEDS=1` reads `<root>/<task>/<task_config>.txt`) |
| `EVAL_TEST_NUM` | seeds per cell (clean usually 5/10; clutter 2) |
| `WORKERS_PER_GPU` | **1** (do not raise; no MPS) |
| `DXVLA_SAPIEN_B1` | **1** → swap to SAPIEN 3.0.0b1 (required for correct office/kitchen) |

**Validated defaults (exported by `slurm_eval_node.sh`, override via `${VAR:-default}`):**
`DXVLA_DENOISER=optix` (OIDN no-op on Blackwell) · `XVLA_POSED_DA3=1` · `XVLA_RGB_INPUT=1` ·
`XVLA_SKIP_GEOMETRY_REINIT=1` · `DXVLA_ACTION_SMOOTH=0.7` `DXVLA_SMOOTH_ROT=1` (cross-chunk smoothing) ·
`DXVLA_PRIMARY_CAM=countertop` `DXVLA_VIDEO_CAM=countertop_camera` · `DXVLA_FAKE_PROPRIO=1` ·
`DXVLA_RT_SPP=32` · `DXVLA_EXEC_ACTIONS` (0=execute full 30; 20=receding-horizon) ·
`EVAL_RESUME_SEEDS=1` `QUEUE_MODE=lease` `LEASE_TTL` · `STALL_TIMEOUT` `PER_TASK_TIMEOUT` ·
`CUROBO_IK_FAIL_RETURN` · `EVAL_FAST_CTRL=1` · `VK_ICD_FILENAMES=.../nvidia_icd.json` ·
node-local `WARP_CACHE_PATH`/`CUDA_CACHE_PATH`. Cross-node: `MODEL_SERVER_BIND`/`MODEL_SERVER_HOST`.
Sim-as-server: `SIM_AS_SERVER`/`MODEL_AS_CLIENT` + `SIM_SERVER_HOST`/`SIM_SERVER_PORT`.

## 6. Running it
**One config (e.g. k320, clean, 5 seeds) on a node pool** — set up the run dir + queue, then srun the pool:
```bash
source /shared_work/jack/eval_root/env.sh
RUN=k320_clean; mkdir -p runs/$RUN/eval_result
cp $BUNDLE/seeds/k320_first10_queue.txt runs/$RUN/queue.txt; rm -rf runs/$RUN/queue.txt.leases
DXVLA_SAPIEN_B1=1 STUDY_RUN=$RUN CKPT_NAME=da3-xvla-k320-ckpt90k BENCH_CONFIG=clean \
  SEEDS_ROOT=$BUNDLE/seeds/k320_first10_seeds EVAL_TEST_NUM=5 WORKERS_PER_GPU=1 \
  srun --partition=gb10 --nodelist=trt-gb10-[3-15] --nodes=13 --ntasks-per-node=1 \
       --gres=gpu:1 --cpus-per-task=18 -t 600 --kill-on-bad-exit=0 \
       bash /shared_work/jack/eval_root/slurm_eval_node.sh
python organize_results.py $RUN k320     # -> runs/k320/<scene>/<task>/seed<N>.mp4 + results.txt
```
All 13 workers share `runs/$RUN/queue.txt` and steal cells; resume by rerunning the same block.
Giants/geostack: add `GIANT_CODE_DIR=<code dir>` (+ `XVLA_DA3_NATIVE_INPUT=1` for the K=160 giants;
geostack uses `DXVLA_GEOSTACK_STEP=20000` and reloads DA3-Large — see its `EVAL_INTERFACE.md`).
Clutter sweep / parallel multi-model / fill-in: see `orchestrate_clutter_k320.sh`,
`orchestrate_parallel_giants.sh`, `orchestrate_fillin.sh` (patterns below).

## 6b. Reproduce OUR exact eval config
Copy these verbatim to match our numbers (k320 clean reference = **58% SR / 49% HSR**):
```bash
DXVLA_SAPIEN_B1=1             # SAPIEN 3.0.0b1 (NOT 3.0.3) — required for correct office/kitchen physics
WORKERS_PER_GPU=1            # no MPS on Blackwell; >1 corrupts the sim
DXVLA_DENOISER=optix         # OptiX denoiser; OIDN is a SILENT no-op on Blackwell (don't use it)
DXVLA_RT_SPP=32              # ray-trace samples/pixel (render quality)
XVLA_POSED_DA3=1             # feed camera extrinsics+intrinsics to DA3 (geometry models)
XVLA_RGB_INPUT=1
XVLA_SKIP_GEOMETRY_REINIT=1  # keep trained geometry weights (do NOT re-init them)
DXVLA_FAKE_PROPRIO=1         # proprio = last commanded action (we eval without true proprio)
DXVLA_ACTION_SMOOTH=0.7      # cross-chunk EMA smoothing (see 6d)
DXVLA_SMOOTH_ROT=1           # smooth xyz + rotation, not just translation
DXVLA_EXEC_ACTIONS=0         # 0 = execute the full 30-action chunk; 20 = receding-horizon
DXVLA_PRIMARY_CAM=countertop  DXVLA_VIDEO_CAM=countertop_camera   # overhead main camera
USE_EVAL_SEEDS=1  EVAL_RESUME_SEEDS=1  EVAL_TEST_NUM=<N>          # seeded + resume
```
k320 used `DXVLA_EXEC_ACTIONS=0` (full 30-chunk); the giant ablations used `=20` (generate 30 / execute 20 / re-query). Keep it consistent within any comparison. These are also the `slurm_eval_node.sh` defaults (override via env).

## 6c. Seed banks — use the SAME seeds as us
Everything under **`$BUNDLE/seeds/`** (`$BUNDLE=/shared_work/jack/robopro-eval-pipeline`, group-readable). Set `SEEDS_ROOT` + `USE_EVAL_SEEDS=1`; the eval reads `<SEEDS_ROOT>/<task>/<task_config>.txt` — space-separated, pre-validated **solvable** episode seeds.
| eval | `SEEDS_ROOT` | queue file | seeds/task |
|---|---|---|---|
| **clean** | `k320_first10_seeds` | `k320_first10_queue.txt` (79 tasks) | 10 (40000–40009) |
| **clutter d6–d15** | `clutter_seeds_hf` | `clutter_queue_hf_d<N>.txt` | 2 per (task, level) |
| clean seeds 5–9 (10-seed fill) | `k320_seeds5to9` | `k320_first10_queue.txt` | 5 (40005–40009) |

Example: `k320_first10_seeds/put_cup_in_box/bench_demo_study_clean.txt` → `40000 40001 … 40009`. To reproduce our **exact episodes**, point `SEEDS_ROOT` at these and use the same `EVAL_TEST_NUM` (clean = 5 or 10; clutter = 2). The clutter banks were rebuilt to match the public vanilla reference (`JackLiu0406/xvla-vanilla-backup/.../2026-05-01_xvla_eval_clutter`), so clutter SR is directly comparable to it.

## 6d. Action smoothing (the 0.7)
`policy/dxvla/deploy_policy.py` (~L279) applies **cross-chunk EMA smoothing** so consecutive 30-action chunks don't jerk at the boundary (we saw a visible gap without it):
- `DXVLA_ACTION_SMOOTH=0.7` — EMA weight α. Each new chunk's first action = `α·(prev chunk's last EXECUTED action) + (1-α)·(raw first action)`, carried across chunks via `model._last_smoothed_action`.
- `DXVLA_SMOOTH_ROT=1` — smooth **xyz + rot6d** (`cont_idx = [0:9]+[10:19]`, skipping the gripper at idx 9/19); `=0` → translation only, orientation raw.
- `DXVLA_ACTION_SMOOTH=0` disables it. **All our final numbers used `0.7` + `SMOOTH_ROT=1`.**

## 6e. Swap in YOUR model (the wrapper)
The policy is loaded by **`policy/dxvla/deploy_policy.py` → `get_model(usr_args)`** (L230), which builds the **`DXVLA` wrapper in `policy/dxvla/dxvla_model.py`** (`XVLA.from_pretrained`, the DA3 weight reload, and the per-step `generate_actions`). `deploy_policy.yml` is its config. The model **server** (`policy_model_server.py`) calls this wrapper; the sim/socket/lease-queue/scoring are model-agnostic.

To run a different model, replace the wrapper:
1. Provide a `get_model(usr_args)` returning an object with
   **`generate_actions(input_ids, image_input, image_mask, domain_id, proprio, ...) -> [B, num_actions, action_dim]`**
   — inputs arrive already tokenized/preprocessed (`image_input` = `[B, V, 3, H, W]`, language in `input_ids`, `proprio`); output is a 30-step action chunk (`ee6d`: 10-D = xyz + rot6d + gripper).
2. Either drop your code under `policy/<name>/` and pass `--config policy/<name>/deploy_policy.yml --policy_name <name>`, or reuse the `dxvla` wrapper and just set `CKPT_NAME` (+ `GIANT_CODE_DIR` for a bundled-code model).
3. Post-load steps (e.g. reload a frozen backbone, set a schedule) go inside `get_model` — see the GeoStack branch in `dxvla_model.py` and each checkpoint's `EVAL_INTERFACE.md`.

## 7. Output & scoring
Raw per-episode metrics stream to `runs/<RUN>/eval_result/<task>/dxvla/bench_demo_<scene>_<cfg>/.../_metrics.jsonl`
(`{"task","seed","success","hard_success",...}`). `organize_results.py <RUN> <model_name>` dedups latest-per-seed
and writes:
- `runs/<model>/<scene>/<task>/seed<N>.mp4` (videos)
- `runs/<model>/<scene>/results.txt` (per-task SR/HSR)
- `runs/<model>/results.txt` (per-scene SR/HSR + TOTAL)
SR = success rate; HSR = "hard" success (success **and** no collision).

## 8. Orchestrator patterns (the 3 drivers)
- **Single-config sweep** (`orchestrate_clutter_k320.sh`): for each `BENCH_CONFIG` in a list, one blocking
  `srun` over the pool on its queue; resume-skips completed configs; aggregates SR per config.
- **Parallel multi-model** (`orchestrate_parallel_giants.sh`): N dedicated 1-node workers (one model each,
  continuous) + an M-node pool that sweeps the models in sequence, all cooperating per model via the shared
  lease queue — every model makes progress at once.
- **Fill-in / resume** (`orchestrate_fillin.sh`): rerun a run dir at a higher `EVAL_TEST_NUM` (resume adds only
  the missing seeds), or run a disjoint seed bank and merge — no redundant episodes.

## 9. Recreate checklist (to run on a DIFFERENT machine)
1. **Benchmark code+assets**: RoboPRO is a customized RoboTwin 2.0 — get RoboTwin 2.0 + its assets (~15 GB),
   apply the `customized_robotwin` deltas (this repo's `code/RoboPRO`).
2. **Two envs**: `robotwin` (SAPIEN, curobo, gymnasium...) and `xvla` (torch, transformers, the model code).
   On aarch64/CUDA-13/Blackwell, **build SAPIEN 3.0.0b1 from source** (see `SAPIEN_B1_ATTEMPT.md`); on x86 you
   may not need the b1 swap but pin the simulator version to match reference numbers.
3. **Checkpoints**: the model weights + their code dirs (+ `EVAL_INTERFACE.md` per ckpt).
4. **Seeds**: pre-validated seed banks (`$BUNDLE/seeds/...`) or regenerate with `precollect_eval_seeds.py`.
5. **Scripts**: `slurm_eval_node.sh`, `stage_*_node.sh`, `harness/eval_worker_steal_dualgpu.sh`,
   `organize_results.py`, the orchestrators — adjust node names/partition/paths.
6. Apply the **GB10 gotchas** in `GB10_EVAL_PIPELINE.md` (WPG=1, OptiX, node-local staging, ffmpeg on PATH, etc.).
```