augment/vace/README.md

VACE augmentation — 8 objects × 32 episodes, 3 cameras

Swaps a target mesh into RoboCasa episodes with the VACE video-diffusion model, producing the same folder layout as the MimicGen augmentation run so the two baselines can be compared file-for-file.

Same-category first, random for the remainder. For donut_5, episodes that already contain a donut (donut_8, donut_9, …) are used first — those are the swaps VACE was designed for. Only after that pool runs out does the rest of the 32 come from randomly drawn episodes of any category (apple → donut_5, milk → donut_5, …). The split is recorded per assignment as fill: same_category | random.

How far the same-category pool goes, after make_gt_masks.py filled in the 58 missing same-category masks (episodes with a complete set of GT masks, excluding the target mesh):

target	same-category pool	same-category used	random fill	reaches 32 same-category?
MeasuringCup009	57	32	0	yes
steak_8	18	18	14	no
SoapDispenser010	5	5	27	no
teapot_6	5	5	27	no
teapot_7	5	5	27	no
donut_5	3	3	29	no
Jar023	2	2	30	no
SyrupBottle005	2	2	30	no
total		72	184

Only MeasuringCup009 can be filled from its own category. This is a property of the dataset, not of the masks: across all 2,517 episodes there are only 5 teapot episodes, 3 donut, 2 jar, 2 syrup_bottle. Generating the other 1,830 missing masks would not move these numbers — only 58 of the 1,888 were same-category for these eight targets (55 measuring_cup, 3 steak), and those are already done.

Run

./run_aug32.sh <work_dir> <out_dir> <gpu> [num_wan_workers]

# one GPU
./run_aug32.sh /scratch/vace_aug32 \
    /home/nvidia/minha/pose6daug/baseline/mimicgen/output_aug_256 2

# four GPUs: split-wan writes four manifests, one worker each
./run_aug32.sh /scratch/vace_aug32 /path/out 2 4     # then move the workers onto 4 GPUs by hand

Every stage is idempotent — re-running after an interruption resumes rather than redoing.

stage	what it does	cost
1 build	pick sources, same-category first	seconds
2 prep	render the target mesh's reference views, cut frames, per-view manifests	~4 s / assignment
3 split	chop into 81-frame WAN jobs	seconds
4 wan	the diffusion pass	everything
5 stitch	glue each camera's chunks back together	seconds
6 collect	copy into the output layout	seconds

Runtime. 256 assignments × 3 cameras ≈ 3,084 WAN jobs. At the measured 51.4 s/job that is ~44 h on one GPU, ~11 h on four. Levers, in order of how much they buy and how little they cost:

• --offload_model False is already passed by the script. The upstream default is True whenever world_size == 1, which moves the DiT to CPU and back around every job: 81.8 s vs 51.4 s measured. Never drop this flag.
• --sample_steps 25 (upstream default 50) → 32.1 s/job, ~27 h on one GPU. Quality trade-off; look at the plates before committing.
• Dropping the two exo cameras cuts the job count to a third. At 256×256 the target-object mask is 0.2 % of an exo frame and VACE tends to erase the object instead of replacing it; only robot0_eye_in_hand (9 %) reliably gets a real swap. If the exo views are not needed, edit CAMS in ../vace_object_insert_batch.py.

Output

<out_dir>/<task>-ep<N>_<object the episode had>_<object it becomes>/
    robot0_agentview_left/robot0_agentview_left.mp4
    robot0_agentview_right/robot0_agentview_right.mp4
    robot0_eye_in_hand/robot0_eye_in_hand.mp4
    .done

e.g. c2cab-ep000151_donut_donut_5 (same-category) and snk2c-ep000004_egg_teapot_7 (random fill). Task tags are c2cab, c2sto, drw2c, snk2c, tst2c; the hyphen keeps the name to three underscore-separated parts, matching MimicGen's ep098_apple_teapot_6. .done carries the real task, episode, source mesh, target, language instruction, and which fill it came from.

Files

file
build_aug32_assignments.py	source selection — the only new logic
collect_outputs.py	VACE plates → the output layout
run_aug32.sh	all six stages
SETUP.md	start here on a new machine — the setup path, in order
preflight.sh	checks environments, weights, datasets, GT-mask coverage, GPU
make_gt_masks.py	re-renders the missing GT masks from the simulator
vace_wan_inference_batch.patch	the upstream diff, for porting onto a fresh clone
vendor/render_mesh_views_side.py	mesh reference-view renderer, vendored from the original machine

Stages 2–5 are ../vace_object_insert_batch.py and ../VACE/vace/vace_wan_inference.py unmodified; this directory adds selection and collection around them.

Prerequisites

Run ./preflight.sh — it checks every item below and prints GT-mask coverage and GPU state. Nothing here needs to be installed or cloned; it all already exists on this box and is referenced by absolute path, the same convention the other baselines use.

Environments (do NOT create a venv)

what	python
batch driver (vace_object_insert_batch.py, the builder, the collector)	/lp-dev/jonghoon/robocasa_calib/envs/robocasa/bin/python (3.10, mujoco + robocasa)
WAN inference	/lp-dev/jonghoon/VACE/venv/bin/python (3.12, torch 2.5.1+cu124, flash_attn 2.8.3)
mesh reference-view rendering (called by prep-worker)	/lp-dev/jonghoon/isaac-gr00t/venv/bin/python — override with VACE_RENDER_PYTHON; it only needs trimesh, pyrender, numpy, pillow (measured: 4.12.1 / 0.1.45 / 1.26.4 / 12.1.1), so any env with those will do

prep-worker renders the target mesh's reference views by importing render_views from render_mesh_views_side.py, which lived outside this tree on the original machine. A copy ships in vendor/render_mesh_views_side.py; point at it with

export VACE_RENDER_SCRIPT_DIR=$PWD/aug32/vendor
export VACE_RENDER_PYTHON=/path/to/a/python/with/trimesh+pyrender

VACE code and weights (already cloned — do NOT clone again)

../VACE/ is a fork of https://github.com/ali-vilab/VACE.git carrying the batch flags (--batch_manifest / --skip_existing / --batch_continue_on_error) this pipeline depends on; a fresh upstream clone will not have them. Weights come in through ../VACE/models -> /lp-dev/jonghoon/VACE-home-models (31 GB: Wan2.1-VACE-1.3B, VACE-Annotators). ~27 GB of GPU memory is needed with --offload_model False.

Datasets

path	what
/home/nvidia/jonghoon/robocasa_full/pickplace_target_human/	source episodes — videos/chunk-000/observation.images.<camera>/episode_NNNNNN.mp4 and extras/episode_NNNNNN/ep_meta.json (symlink to /lp-dev/...)
/home/nvidia/jonghoon/robocasa_full/aug_materials/robocasa/<task>/ep<N>/	mask_gt_target_object_<camera>.mp4 — the GT masks
/lp-dev/jonghoon/robocasa_calib/repos/robocasa/robocasa/models/assets/objects/{objaverse,lightwheel}/	target meshes

GT-mask coverage, measured: PickPlaceCounterToCabinet 502, CounterToStove 50, DrawerToCounter 43, SinkToCounter 34, ToasterToCounter 0 — 629 usable episodes, and the builder draws only from those.

Porting to another machine

Do not clone ali-vilab/VACE fresh — upstream has no --batch_manifest, which every stage past split-wan depends on.

Copy the whole baseline/vace/ directory, not just aug32/. Nothing outside it is needed except the 31 GB of WAN weights. To rebuild from a clone instead: git clone https://github.com/ali-vilab/VACE.git && git checkout 48eb44f, then git apply vace_wan_inference_batch.patch, and copy build_vace_wan_batch_manifest.py, vace_inference_full.py and this aug32/ directory alongside it.

Unpack the data

cd baseline/vace && mkdir -p data && cd data
tar -xf  ../datasets_pack/source_episodes.tar       # 2.7 GB -> pickplace_target_human/
tar -xzf ../datasets_pack/aug_materials_ext.tar.gz  # 121 MB -> aug_materials_ext/
tar -xzf ../datasets_pack/target_meshes.tar.gz      #  14 MB -> objaverse/, lightwheel/
cd ..

Point everything at it

Every path is an environment variable with the original machine's value as the default, so this is the whole configuration:

variable	default	set it to
VACE_ROBOCASA_ROOT	original machine	$PWD/data/pickplace_target_human
VACE_AUG_MATERIALS	original machine	$PWD/data/aug_materials_ext
VACE_ASSET_ROOTS	original machine	$PWD/data/lightwheel:$PWD/data/objaverse
VACE_PY_ROBOCASA	original machine	a python with robocasa, mujoco, trimesh, pyrender
VACE_PY_WAN	original machine	a python with torch 2.5.1+cu124 and flash_attn
VACE_RENDER_SCRIPT_DIR	aug32/vendor/	— already inside the tree
VACE_RENDER_PYTHON	VACE_PY_ROBOCASA	— only if that env lacks trimesh/pyrender

export VACE_ROBOCASA_ROOT=$PWD/data/pickplace_target_human
export VACE_AUG_MATERIALS=$PWD/data/aug_materials_ext
export VACE_ASSET_ROOTS=$PWD/data/lightwheel:$PWD/data/objaverse
export VACE_PY_ROBOCASA=/path/to/robocasa/python
export VACE_PY_WAN=/path/to/wan/python
./aug32/preflight.sh          # says MISS for anything still wrong

The mesh reference-view renderer (render_mesh_views_side.py, which prep-worker imports) is vendored at aug32/vendor/ and defaults to the driver's own interpreter, so it needs no separate environment — trimesh and pyrender are its only requirements beyond numpy/PIL.

Still needed from elsewhere

The 31 GB of WAN weights (Wan2.1-VACE-1.3B, VACE-Annotators) — link them in at VACE/models. And, only if you want to re-run make_gt_masks.py there, the meshes of objects other than the eight targets, since that rebuilds each episode's own scene.

Gotchas

• An episode without GT masks makes prep-worker raise FileNotFoundError. The builder filters on this, so it cannot happen through run_aug32.sh — but it will if you hand-edit assignments.jsonl.
• The WAN stage must run with vace/VACE as the working directory (imports and --ckpt_dir are relative to it). run_aug32.sh does the cd.
• --save_file in vace_wan_inference.py is written relative to the current directory, not --save_dir. The batch path does not use it; a one-off call should pass an absolute path.