actaug/code/GUIDE.md

Pose6DAug — simulator action augmentation (RoboCasa), reproduction guide

This directory is the data-generation code that produces the augmented training episodes: a source RoboCasa PickPlaceCounterToCabinet episode has its object swapped for one of 8 textured SAM3D reconstructions (visual + CoACD collision, true GT size per axis, GT mass), and the grasp is re-earned in simulation by the policy (constrained inference: the policy owns only wrist roll joint 7 + gripper; every other joint, the base and torso are pinned to the recording). Only task-successful rollouts are kept. No pose estimation, no inpainting: everything is sim.

Checkpoints

use	checkpoint	where
GR00T grasp policy for augmentation, and the Table-2 baseline	pnpcountertocab_all502_gbs64_wandb_60k_save5k_20260427_193901/checkpoint-60000	base/ in this repo (https://huggingface.co/Ronaldo-GOAT/pose6daug/tree/main/base)
π0.5 grasp policy (optional variant)	pi05_ctc502_60k_rawparams (openpi, config pi05_robocasa_target_PickPlaceCounterToCabinet)	https://huggingface.co/Ronaldo-GOAT/transfer (checkpoint + patched openpi code, see its EVAL_SEEDS.md)
finetune start point for training on the augmented data	same base/ checkpoint-60000	this repo

Training after augmentation: code/train/train_groot15_single_dataset.sh with BASE_CKPT=<base/> (GR00T 1.5, new_embodiment, single_panda_gripper, global batch 64).

Inputs

• Source dataset: RoboCasa LeRobot export pickplace_target_human/PickPlaceCounterToCabinet (502 episodes, 20 fps, extras/episode_XXXXXX/{model.xml.gz,states.npz,ep_meta.json}; the per-episode compiled MJCF + flattened sim states are what make exact resets possible).
• RoboCasa / robosuite checkouts (robosuite_robocasa_dev fork) and the mimicgen conda env for the sim client (actaug_core.make_env uses robomimic EnvRobosuite, reset_to).
• The 8 targets (robocasa_coacd/canonical/<target>/): SAM3D original.obj + material_0.png, already rotated into the GT canonical frame (aligned_manifest.json), CoACD hulls hull_NN.obj (coacd_manifest.json, max 16 hulls, threshold 0.05). Per-axis size reference true_gt_extents_m.json (= the GT mesh bounds × the <mesh scale> in the GT model.xml) and true_gt_mass.json (GT compiled mass → hull density).

Pipeline (what robocasa_coacd/actaug_runner_coacd_prio.py does per (source episode, target))

0. swap_object_in_xml (GT replacement) → swap_full_to_sam3d: visual → SAM3D mesh, collision → CoACD hulls, per-axis scale TRUE_GT/sam_extents × s_jit, s_jit ~ U[0.96,1.04] drawn once per pair (model compiled once; attempts reset state only), hull density = GT mass / hull volume, hull contact params priority="2" + explicit GT friction/solref/solimp (the "prio" fix: without it MuJoCo averages the hull's stiff solref with the soft finger solref and the fingers sink up to ~10 mm into the object; with it held-grasp penetration is < 0.2 mm).
1. Place upright (pure-yaw body quat, base 2 mm above the counter, settle 60 steps); reject if the object tips (up_cos < 0.80) or the bottom is off the support by > 3 cm.
2. Replay the recorded actions verbatim; if the task succeeds keep them (stage=replay).
3. Otherwise R/T/S + policy grasp: yaw ∈ {0, 90, 180}° ± 10°, searched wrist offset dj7 ∈ {0, ±30, ±60}°, policy handoff start_off=8 steps before the recorded grasp, threshold gripper latch, rel-6D weld for the carry, up to --cap attempts; first is_success()["task"] wins.
4. Dump: actions.npy (T,12, parquet order base-first), states.npz (T,16), left/right/wrist/3cam.mp4 (256², 20 fps), obj_track.npz, meta.json (success, stage, seeds, s_jit, instruction).

Feasibility pool (heightfilter/): a source is eligible for a target if its object height is within ±3 cm of the target's placed height (heights.csv, computed from the recorded grasp); same-category sources are ranked first. Base movers (23 source episodes whose mobile base drives) are excluded (source_movers.json), and curation additionally drops any dump with a nonzero base action, > 5 mm base drift or > 0.5° yaw drift.

Running

# 1. policy servers (GR00T, one per GPU; ports 8905/8906 used by the launch scripts)
bash launch_server.sh 2 8905 <base/checkpoint dir>; bash launch_server.sh 3 8906 <base/checkpoint dir>
# 2. generation: 1817 static pairs, cap 24 attempts, N workers on GPU 2/3 (EGL id 0/1 = physical 2/3 on the source node)
N=40 bash robocasa_coacd/launch_prio.sh          # -> outputs/robocasa_aug1024_prio/
# 3. curation (the launch script starts curate_prio.sh, which also stops the workers at 128/object):
#    outputs/robocasa_aug256_prio/ (32/object) and outputs/robocasa_aug1024_prio_curated/ (128/object), symlinks
# 4. packaging to LeRobot (language rewritten to the target object, intra-only h264):
python robocasa_coacd/pack_lerobot.py outputs/robocasa_aug1024_prio_curated <gr00t_views>/<name>
# 5. merge with the base 502 (videos symlinked) and finetune from base/checkpoint-60000
python merge_gr00t_view_datasets.py --source <base502> --source <aug> --output <merged>

π0.5 variant: launch_pi05_servers.sh (openpi servers on GPU 2/3, ports 9800/9801) then N=40 bash launch_pi05.sh → outputs/robocasa_aug1024_prio_pi05*. pi05_client.py maps the GR00T obs dict to the openpi request (left + wrist image with the eval's mp4 round-trip, 16-d state in openpi order, client-side flow-matching noise seeded per query) and the (10,12) robosuite-order action chunk back to the GR00T chunk dict, so constrained.py is unchanged.

Verification tools included

probe_penetration*.py / prio_batch_probe.py / inspect_contact_params.py (true finger↔object depth, contact parameter dump), _scale_sheet.png / _inscene_sheet.png (GT vs SAM3D silhouettes and in-scene renders), priotest/FINAL_TABLE.txt (baseline vs prio depth/success on 12 pairs).

Evaluation (Table 2)

code/eval/table2_exact160/eval_exact160.sh <tag> <checkpoint>: 160-episode exact-replay bank (8 SAM3D targets × 20 fixed initial states, bank included), 8 workers × 20, server seed 42, policy_seed 42. Reference: base/checkpoint-60000 = 15/160 (9.4%), ours = 34/160 (21.2%). Paths at the top of the script are overridable by env (MYGROOT_ROOT, MYGROOT_PY, ROBOCASA_PY, GPUS, EVAL_OUT).

Notes

• Scripts carry absolute paths of the source node (/lp-dev/jonghoon/...); they are the exact artifacts that generated the data, not a portable package.
• Seeds: per-pair seed0 = ep*100 + rep, attempt seed seed0*100 + n, policy noise seed attempt_seed*100003 + n_infer (GR00T honors _policy_action_seed; the π0.5 client draws the noise itself from the same value).