README.md

# pose6daug — augmentation, dataset build, training and evaluation scripts

The scripts behind the MimicGen and VACE augmentation baselines for RoboCasa
PickPlaceCounterToCabinet, and the GR00T 1.5 fine-tuning and exact-replay evaluation run on
top of them. Archived as-run: paths are absolute and point at one particular machine, so treat
this as a record of the procedure rather than a turnkey package. See **Paths to change** below.

```
augment/mimicgen/   generate episodes with MimicGen, render preview videos
augment/vace/       VACE object-swap augmentation (assignment builder, runner, GT masks)
augment/actaug/     actaug episode folders -> gr00t_views (no simulator needed)
dataset/            convert generated episodes into a gr00t_views (LeRobot v2.1) dataset
train/              GR00T 1.5 fine-tuning launcher
eval/               exact-state replay evaluation -- read eval/EVALUATION.md for the protocol
ops/                checkpoint retention, optimizer pruning, eval-on-checkpoint triggers
tools/              dataset sanity checks
```

## Pipeline

```
MimicGen generation ──► demo.hdf5 (MuJoCo states + actions, no pixels)
                          │
                ┌─────────┴──────────┐
                ▼                    ▼
       preview mp4 (3 views)   gr00t_views dataset  ──► fine-tune ──► exact-replay eval
```

VACE skips the simulator: it repaints pixels and keeps the source trajectory, so its builder
copies the source parquet instead of replaying.

### 1. Augment

`augment/mimicgen/run_generation.sh` launches one MimicGen process per worker from a JSON
config per (object, worker); `config_template.json` is the template. Two knobs matter:

- `guarantee=false` makes `num_trials` mean **attempts**, not successes. Giving every object
  the same attempt budget is what produces the generator's natural yield instead of a
  per-object quota.
- `obj_registries` must include `aigen` for objects under `aigen_objs/` (e.g. `wine_5`), or
  `sample_kitchen_object_helper` raises a bare `ValueError`.

`snapshot_episode_times.sh` records the per-episode temporary filenames while they exist:
MimicGen writes each success to `tmp/date_..._time_HH_MM_SS.hdf5`, then `merge_all_hdf5`
sorts by timestamp and deletes the folder. That snapshot is the only record of when each
episode was produced, and `select_by_generation_time.py` joins it back to order episodes
globally across workers. The poller can only ever miss a worker's **last** file, which the
selector pads with the merged file's mtime.

### 2. Build the dataset

`dataset/render_to_gr00t.py` replays each episode in MuJoCo, renders three cameras and writes
parquet + videos + meta. `dataset/run_convert.sh` shards it (each shard needs its own
`GEN_DIR` and `DATASET_OUT`, or the glob picks up the others), then
`merge_gr00t_view_datasets.py` merges and `repair_mimicgen_task_ids.py` fixes task ids.

`repair_mimicgen_task_ids.py` is not optional: the writer stores every parquet task column as
0 while `episodes.jsonl` holds the intended language, so without it every episode trains as
task 0 and the language conditioning silently collapses.

### 3. Train

```bash
DATASET_PATH=<gr00t_views dataset> DATASETNAME=<name> \
GPUS=2,3 PER_GPU_BATCH=32 MAX_STEPS=30000 SAVE_STEPS=5000 \
bash train/train_groot15_single_dataset.sh
```

Fine-tunes the action-head projector and diffusion head from a base checkpoint; the backbone
stays frozen. `RESUME=1` picks up the newest checkpoint in the output directory.

### 4. Evaluate

The full protocol -- episode set, seeds, which checkpoints to compare, how to read the
stage flags -- is in [`eval/EVALUATION.md`](eval/EVALUATION.md). The short version:

```bash
MODEL_PATH=<checkpoint> MYGROOT_ROOT=<myGR00T tree> \
EVAL_CLIENT=<eval/eval_robocasa_replay_state_grasp.py> \
REPLAY_STATE_ROOT=<replay set> N_EPISODES=160 GPUS=4,5,6,7 POLICY_SEED=12345 SEED_BASE=42 \
bash eval/eval_groot15_exact_replay.sh
```

Each episode restores a saved scene XML and flattened MuJoCo state, so runs are comparable
across checkpoints. Two independent seeds:

- `POLICY_SEED` — the policy server's global RNG **and** the per-step action seed
  (`policy_seed + episode_index × stride + step`). This is the one that changes the rollout.
- `SEED_BASE` — the env seed, `SEED_BASE + worker_id` per worker. Exact replay overwrites the
  scene immediately after, so it should not affect the initial state.

Worker count follows the GPU list, and episodes are split evenly across workers — so the same
`SEED_BASE` with a different GPU count gives each episode a different worker seed.

Client variants:

| file | adds |
|---|---|
| `eval_robocasa_replay_state_3view.py` | saves the wrist (ego) view; `RecordVideo` only captures `robot0_agentview_center` |
| `eval_robocasa_replay_state_grasp.py` | the above, plus `--episode_indices` for an arbitrary subset, and per-episode `grasped` / `lifted` / `in_cab` stage flags |

Stock RoboCasa success is `obj_inside_of(cab) and gripper_obj_far` — a single boolean, which
says nothing about where a failed episode broke down. The stage flags come from
`_check_grasp`, a 3 cm rise in the object's body z, and `OU.obj_inside_of`.

`organize_videos.py` regroups the output into one folder per global episode
(`rollouts/episode_NNNNNN/{center.mp4, wrist.mp4, info.txt}`); `RecordVideo` names files by
worker-local index, which does not match the global index.

## The action-order bug — check any dataset before training on it

RoboCasa's simulator exports actions **arm-first**:

```
[eef_pos(3), eef_rot(3), gripper, base(3), torso, base_mode]
```

A gr00t_views dataset declares them **base-first**:

```
base_motion[0:4]  control_mode[4:5]  eef_pos[5:8]  eef_rot[8:11]  gripper_close[11:12]
```

A builder that re-derives actions from a simulator rollout has to reorder; one that copies
rows from an existing LeRobot dataset does not. Copying them through unchanged puts arm motion
on the base-motion channel and `base_mode` on the gripper. Nothing errors, the loss converges
to a small value, and the policy drives the base away from the counter and never closes the
gripper — **0/160** on an exact-replay eval whose base checkpoint scored 11/160.

```python
actions = raw_actions[:, [7, 8, 9, 10, 11, 0, 1, 2, 3, 4, 5, 6]]
```

`tools/check_action_layout.py` tells the two layouts apart from the data alone (the mobile
base never moves in this task, so the constant dimensions give it away), and
`dataset/repair_action_order.py` fixes an already-built dataset in place — parquet plus the
action entry of `meta/stats.json`, no re-render. It refuses to run on a dataset that does not
look like simulator order, so it cannot be applied twice. Read `tools/ACTION_LAYOUT.md` first.

## Paths to change

Every script hard-codes absolute paths from the machine this was run on. At minimum:

| what | appears as |
|---|---|
| RoboCasa / robosuite checkouts | `/lp-dev/jonghoon/robocasa_calib/repos/...` |
| MimicGen env + augmentation code | `/lp-dev/jonghoon/mimicgen_augment/...` |
| myGR00T tree and conda envs | `/data/minha/pose6daug/train_robocasa/myGR00T`, `/data/nvidia/gripper_augmentator/conda-envs/...` |
| base checkpoint | `/lp-dev/jonghoon/myGR00T_outputs/pnpcountertocab_all502_.../checkpoint-60000` |
| replay sets and eval output | `/lp-dev/jonghoon/isaac-gr00t/eval_results/...` |
| generation scratch | `/tmp/claude-.../scratchpad/...` |

`eval/eval_groot15_exact_replay.sh` defaults `MYGROOT_ROOT` to a path that no longer exists;
pass it explicitly. No credentials are embedded — the training launcher reads
`WANDB_API_KEY` from an env var or a file path you supply.

## Attribution

`eval/eval_groot15_exact_replay.sh` and the upstream of `dataset/render_to_gr00t.py` and the
eval clients come from the pose6daug project's shared tree; the copies here carry the fixes
described above (action reorder, mesh-swap bbox fix, target sharding, wrist-view capture,
stage logging).

## actaug

`augment/actaug/build_actaug_gr00t.py` converts actaug's per-episode folders
(`actions.npy`, `states.npz`, `left/right/wrist.mp4`, `meta.json`) into a gr00t_views dataset.
Like the VACE builder it needs no simulator -- the videos and states already exist.

Two things it handles:

- **Instruction rewrite.** actaug keeps the *source* episode's language (65 distinct strings,
  e.g. "Pick the wine ..." on a SoapDispenser010 episode). The augmented object is the one in
  the folder name, so the instruction is regenerated from that, giving the same 7 task strings
  the other datasets use.
- **Action order.** actaug already writes base-first, matching `modality.json`, so no reorder
  is needed -- unlike the MimicGen path. Its actions do exercise the mobile base in some
  episodes, which the other datasets never do, so `tools/check_action_layout.py` cannot
  classify it (that check assumes a static base). Verify by hand there instead.

## Related artifacts

- `mlnha/mimicgen-batch64-30k-ckpts` — checkpoints from the fixed-action-order run
- `mlnha/vace-batch64-30k-ckpts` — VACE run checkpoints
- `mlnha/mimicgen-pi05-aug256` — MimicGen augmentation for the pi0.5 hard-object set