pose6daug-scripts / augment /vace /DATASET_CONVERSION.md
mlnha's picture
pose6daug augmentation / dataset / training / eval scripts
5d08972 verified
|
Raw History Blame Contribute Delete
8 kB

Turning VACE output into a trainable dataset β€” everything needed

After run_aug32.sh finishes you have per-episode videos. Training needs a LeRobot v2.1 "gr00t_views" dataset. This is the full inventory for doing that conversion on a machine that has nothing but this repo.

Written 2026-09-19. The MimicGen side of the same job is documented in baseline/mimicgen/TRAINING_DATASET_PREP.md; the two are not the same procedure β€” see "Why VACE is the easy case" below.


Why VACE is the easy case

VACE repaints pixels; it never changes the trajectory. So the augmented episode's state/action rows are the source episode's own rows, unchanged. The builder copies that parquet, rewrites indices and task ids, and links the new videos in. No simulator, no re-render, no MuJoCo.

MimicGen is the opposite: every episode is a new rollout, so its parquet has to be rebuilt from the generation HDF5 by replaying it in the simulator (~2 h for 256 episodes).

VACE MimicGen
builder build_vace_objwise_gr00t_dataset.py render_to_gr00t.py
parquet copied from the source episode rebuilt from HDF5
video symlinked from the VACE output re-rendered in MuJoCo
needs robosuite/robocasa/mujoco no yes
runtime for 256 eps minutes ~2 h

What to bring

1. The builder and its checker

train_robocasa/scripts/dataset_build/build_vace_objwise_gr00t_dataset.py
train_robocasa/scripts/dataset_build/repair_mimicgen_task_ids.py      # also needed here
train_robocasa/scripts/dataset_build/verify_geom_objwise_gr00t_dataset.py
train_robocasa/scripts/dataset_build/merge_gr00t_view_datasets.py     # only if merging with the 502 base
train_robocasa/scripts/dataset_build/verify_merged_gr00t_dataset.py

Only numpy and pandas beyond the standard library, plus ffmpeg on PATH for the optional three-view preview (--skip-three-view avoids even that).

2. Inputs

--flag what where it comes from
--vace-root the VACE output tree, one directory per assignment holding work_<view>/<inference_dir>/out_video.mp4 run_aug32.sh's work dir
--manifest assignments.jsonl (repeatable) aug32/build_aug32_assignments.py
--source-root the original RoboCasa dataset β€” the parquet is taken from here datasets_pack/source_episodes.tar
--canonical-meta-source a dataset whose info.json / modality.json / view_manifest.json contract is copied dataset_build/canonical_meta/, bundled here
--output where the dataset is written your choice

The defaults all point at the original machine, so pass all five explicitly.

The 502-episode gr00t_views set is NOT needed

Neither for training (we do not merge) nor for conversion:

  • its meta/ was the only thing --canonical-meta-source reads, and those six files are bundled at dataset_build/canonical_meta/ (96 KB);
  • the state/action rows come from --source-root, i.e. pickplace_target_human, whose parquet already carries state 16-d and action 12-d β€” the same schema as gr00t_views.

--canonical-meta-source expects a dataset directory and reads <dir>/meta/, so dataset_build/canonical_dataset/ is set up for exactly that β€” its meta is a symlink to canonical_meta. Pass --canonical-meta-source dataset_build/canonical_dataset.

3. Manifest fields the builder reads

case_dir, episode_index, provenance, and the whole source block (episode, task, lang, source_mesh, source_category, per-camera videos). build_aug32_assignments.py writes all of them.

provenance decides what gets in:

row.setdefault("provenance", "existing")
if row["provenance"] == "random_fill" and not include_random_fill:
    continue          # dropped

Our builder marks same-category assignments existing and the random remainder random_fill. With 32 per target that is 59 same-category and 197 random, so the default keeps only 59. Pass --include-random-fill to get all 256.


Run it

python build_vace_objwise_gr00t_dataset.py \
    --vace-root        /path/to/vace_aug32_work \
    --manifest         /path/to/vace_aug32_work/assignments.jsonl \
    --source-root      $PWD/data/pickplace_target_human \
    --canonical-meta-source $PWD/dataset_build/canonical_dataset \
    --output           /path/to/gr00t_views/vace_per_target_32_256eps \
    --include-random-fill \
    --skip-three-view

python repair_mimicgen_task_ids.py --dataset /path/to/gr00t_views/vace_per_target_32_256eps
python verify_geom_objwise_gr00t_dataset.py --dataset /path/to/gr00t_views/vace_per_target_32_256eps

--copy-videos materialises the videos instead of symlinking β€” needed if the dataset will be moved away from the VACE output tree.

repair_mimicgen_task_ids.py is not optional

Its docstring: the writer "stores all parquet task columns as 0 while episodes.jsonl stores the intended language per episode", and training reads the task id from the parquet β€” so without this pass every augmented episode trains as task 0 and the language conditioning silently collapses. It also renames the modality keys to the canonical left_view / right_view / wrist_view that single_panda_gripper expects.


What the result must look like

<output>/
  meta/{info,modality,episodes,tasks,stats}.json      all five, asserted by the loader
  data/chunk-000/episode_NNNNNN.parquet
  videos/chunk-000/observation.images.robot0_agentview_left/episode_NNNNNN.mp4
                  .../robot0_agentview_right/...   .../robot0_eye_in_hand/...
  extras/episode_NNNNNN/...                         eval only, training never reads it

parquet columns (11): observation.state (list, 16), action (list, 12), next.reward, next.done, timestamp, frame_index, episode_index, index, task_index, annotation.human.task_description (int index into tasks.jsonl, not the string), annotation.human.task_name.

modality.json slicing:

state   base_position 0:3 | base_rotation 3:7 (quat) | EE_pos_rel 7:10
        | EE_rot_rel 10:14 (quat) | gripper_qpos 14:16
action  base_motion 0:4 | control_mode 4:5 | EE_pos 5:8 | EE_rot 8:11 | gripper_close 11:12
video   left_view / right_view / wrist_view -> observation.images.robot0_agentview_{left,right}
                                               and observation.images.robot0_eye_in_hand

Then training

train_robocasa/ is a separate self-contained bundle (49 MB + the mygr00t conda env and a 7.1 GB base checkpoint, both referenced by absolute path). Point it at the dataset by absolute path rather than by name:

cd train_robocasa
DATASET_PATH=/path/to/gr00t_views/vace_per_target_32_256eps \
GPUS=0,1,2,3 MAX_STEPS=30000 SAVE_STEPS=5000 PER_GPU_BATCH=16 \
bash scripts/train_groot15_single_dataset.sh

DATASETNAME (the alternative) resolves against /lp-dev/jonghoon/robocasa_full/gr00t_views/, which will not exist elsewhere.

Check first with bash scripts/smoke_test.sh <gpu> β€” imports, dataset load, and a real 20-step fine-tune.


Checklist for the other machine

  • run_aug32.sh finished; every assignment has three out_video.mp4
  • assignments.jsonl carries provenance (re-run build_aug32_assignments.py if it predates 2026-09-19)
  • source_episodes.tar unpacked β€” the parquet (data/) is read from it; the archive was repacked 2026-09-19 to include it, an earlier copy has only videos/ and extras/
  • dataset_build/canonical_dataset/ present (its meta symlink must resolve; use cp -rL when moving)
  • build_vace_objwise_gr00t_dataset.py + repair_mimicgen_task_ids.py + a verifier copied over
  • python with numpy, pandas; ffmpeg unless --skip-three-view
  • decided: --include-random-fill (256) or not (59)
  • decided: merge with the 502 base or train on the augmented set alone