# 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
```bash
# 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).