# 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= DATASETNAME= \ 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= MYGROOT_ROOT= \ EVAL_CLIENT= \ REPLAY_STATE_ROOT= 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