mlnha/pose6daug-scripts

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

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. The short version:

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.

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