actaug/code/README.md

# actaug — perturbed-object action re-earning for robocasa episodes

Standalone handoff package. **No imports from any other project directory** — only
installed libraries (robocasa / robosuite / robomimic / numpy / pandas / zmq /
torch / imageio). Reference implementations it was distilled from (read-only):
`/home/nvidia/jonghoon/idm_eval/scripts/gripper_state_replay_sandbox/replay/`
(`gripper_state_replay.py` hybrid mode, `grip_perturb.py`).

Physical location `/lp-dev/jonghoon/actaug` (symlinked at
`/home/nvidia/jonghoon/actaug`).

## What it does

For an arbitrary episode of a robocasa LeRobot export (per-episode
`extras/episode_XXXXXX/{model.xml.gz, states.npz, ep_meta.json}` + 12-d actions in
`data/chunk-XXX/episode_XXXXXX.parquet`):

1. **Reset** the MuJoCo sim to the recorded initial state — the episode's own
   compiled `model.xml` + flattened state via `EnvRobosuite.reset_to`, so the
   scene / layout / cameras are bit-identical to the source episode.
2. **Perturb** the manipulated object (`obj_joint0` free joint): yaw about the
   WORLD vertical + xy translation, then a short physics settle. Random draws:
   yaw ~ U[-rot_deg, rot_deg], translation direction uniform on the circle,
   radius ~ U[trans_min_m, trans_m].
3. **Replay** the recorded actions until `grasp_t - pre_buffer`
   (`grasp_t` = first recorded gripper-close command).
4. **Policy handoff**: GR00T-N1.5 grasps (chunks of `exec_h=16` env steps,
   stop when the object is lifted > 3 cm or budget runs out), then the recorded
   transport actions resume at `grasp_t + resume_offset`.
5. **Sample N times** per perturbation. Each attempt is fully seeded: the obs
   dict carries `_policy_action_seed` (honored by myGR00T `policy.get_action`),
   so attempts are diverse AND reproducible.
6. **Filter on task success** (`env.is_success()["task"]`). Successful attempts
   are dumped under `accepted/` with executed `actions.npy` (parquet order),
   `states.npz` (16-d observation.state layout), and mp4s from the SAME three
   cameras as the dataset (`robot0_agentview_left/right`, `robot0_eye_in_hand`,
   256×256) + a side-by-side `3cam.mp4` + `meta.json`.

## Two grasp-re-earning algorithms (pick via YAML `mode:`)

**`hybrid` (unconstrained)** — GT actions to `grasp_t − pre_buffer`, then the
policy commands the FULL action space (EEF pose + gripper through the OSC
controller) until the object lifts, then recorded transport resumes. Handles
translated objects (the policy can chase them). `configs/unconstrained.yaml`.

**`constrained`** — the sim_action_aug algorithm
(`sim_action_aug/code/rollout/r02_rollout.py`, `ws/ROLLOUT_REPORT.md` §3)
ported to robocasa: every arm joint NOT listed in `constrained.policy_joints`
(+ mobile base + torso) is PINNED to the recorded MuJoCo state path every
physics substep (zero drift, the grip_state/grip_policy machinery). The policy
owns ONLY the listed joints (default `[6, 7]`, i.e. wrist pitch + roll, like
the sim_action_aug pipeline) and the gripper:

- **each joint j in `policy_joints`**: `recorded + servo` — the policy's delta
  rotation projected onto joint j's world axis, capped `step_cap` rad/step and
  `servo_caps[j]` rad total (defaults: 6→0.30, 7→0.60);
- **joint 7** (when listed) additionally: `recorded + dj7_init + servo` — `dj7_init` is a searched
  grid (`constrained.dj7_init_deg`, slewed in over ≤48 frames instead of
  teleported, per r02 phase R), with the (−π/2, π/2] wrap (2-jaw symmetry)
  and `|dj7_init + servo| ≤ j7_total_cap`;
- **the gripper**: `action.gripper_close` in one of three interpretations —
  `bit` (binarized, robocasa-native), `threshold` (latch full close once
  ≥ `grip_tau`; this is what unlocked ep 0009 in sim_action_aug),
  `continuous` (ctrl interpolated open→close), `recorded` (diagnostic: the
  recorded close bit, = grip_state).

Transport = relative replay: everything pinned, dj7 frozen at hand-off,
gripper held closed until the recorded release frame, then the recorded
opening schedule (never re-closes), plus 30 settle frames.

**rel-6D weld (on by default, `constrained.rel6d`)** — sim_action_aug's simfix,
required here too: a pinned (infinite-impedance) arm carry slips out of the
friction grasp (the old robocasa pin-sweep peaked at ~33% success; reproduced
on ep0: the pinch engages, then the bottle squeezes out during the lift, with
recorded gripper timing and no policy at all). So an INACTIVE
`<weld body1="gripper0_right_eef" body2="<obj root>"/>` is injected into the
episode XML; it engages at the CURRENT relative pose once both jaw pads touch
the object while it is lifted ≥ `weld_lift_m` under a close command, and it
drops the moment the gripper is commanded open — the release is a genuine
physical fall. Engage/release frames are logged per attempt (`weld_engage_t`,
`weld_release_t`). Tuning matters: with a 15 mm engage gate + soft solref 0.02
the bottle sagged ~19 mm in the jaw, caught the cabinet-shelf edge on the way
in, and fell out at release; `weld_lift_m: 0.004` + `weld_solref: "0.005 1"`
(the defaults) fixed it — recorded-gripper baseline then reproduces the GT
success exactly.

`configs/constrained.yaml`. Caveat: the pinned arm cannot chase a TRANSLATED
object — use it for yaw-dominant perturbations / small translations, and
`hybrid` for the 3–7 cm displacement regime.

Note: attempts count = `n_perturbs × len(dj7_init_deg) × n_samples` for
constrained, `n_perturbs × n_samples` for the others. The wrist servo uses the
policy's raw delta-rotation command (controller units) — direction is what
matters; the per-step/total caps bound the magnitude exactly as in r02.

## YAML reference (every key)

A YAML passed as `--config` sets argparse defaults; any explicit CLI flag
overrides it. Top-level keys mirror `rollout.py` flags; the `constrained:`
block is passed to `run_constrained_attempt` verbatim.

### Top-level keys (both algorithms)

| key | default | meaning |
|---|---|---|
| `mode` | hybrid | `gt` (pure replay sanity) / `hybrid` / `policy` (end-to-end) / `constrained` |
| `dataset` | PickPlaceCounterToCabinet export | robocasa LeRobot export root |
| `episode` | — (required) | episode index |
| `perturb` | "" | explicit `"dyaw_deg,dx_m,dy_m"`; overrides random draws |
| `n_perturbs` | 0 | number of RANDOM perturbations (0 + no `perturb` = nominal) |
| `rot_deg` / `rot_min_deg` | 60 / 0 | yaw magnitude ~ U[rot_min_deg, rot_deg], sign random |
| `trans_m` / `trans_min_m` | 0.05 / 0 | translation radius ~ U[trans_min_m, trans_m], direction uniform |
| `seed` | 0 | seeds BOTH the perturbation draws and the per-attempt policy seeds |
| `n_samples` | 1 | **N**: policy samples per initial condition (rejection sampling) |
| `pre_buffer` | 10 | **the buffer**: policy inference starts at `grasp_t − pre_buffer` frames (`grasp_t` = first recorded gripper-close). hybrid-mode flag; constrained mode reads it from its own block (below) |
| `policy_budget` | 200 | hybrid only: max policy env-steps to achieve the lift |
| `resume_offset` | 0 | hybrid only: recorded transport resumes at `grasp_t + resume_offset` |
| `exec_h` | 16 | hybrid only: env steps executed per policy chunk (inference stride) |
| `host` / `port` | 127.0.0.1 / 8801 | GR00T server address |
| `save` | success | `success` (dump accepted only) / `all` (+ rejected/) / `none` |
| `fps` | 20 | output video fps (= dataset rate; don't change unless re-timing) |
| `out` | outputs/ep\<N\>_\<mode\> | output dir |

### `constrained:` block

| key | default | meaning |
|---|---|---|
| `pre_buffer` (alias `start_off`) | 4 | **the buffer**: inference starts at `grasp_t − pre_buffer`; the dj7_init ramp is scheduled BEFORE this point so it never eats into the policy window |
| `k_post` | 12 | policy keeps control until `grasp_t + k_post`, then hand-off |
| `replan` | 8 | **inference stride**: control steps between policy queries; each query returns a 16-step chunk, of which the first `replan` steps are executed |
| `policy_joints` | [6, 7] | **the joint constraint**: 1-indexed arm joints the policy may servo; everything else (joints, base, torso) stays pinned to the recording. Gripper is ALWAYS the policy's unless `gripper_mode: recorded` |
| `step_cap` | 0.05 | **inference speed**: max servo change per 20 Hz control step (rad) for every policy joint; also paces the dj7_init ramp. 0.05 ≈ 1 rad/s (rate-matched to sim_action_aug) |
| `servo_caps` | {6: 0.30, 7: 0.60} | total authority per joint (rad) — how FAR, not how fast |
| `servo_cap_default` | 0.30 | cap for listed joints missing from `servo_caps` |
| `j7_total_cap` | π | joint 7 only: bound on `dj7_init + servo` |
| `dj7_init_deg` | [0] | searched initial wrist-roll offsets (deg); grid × `n_samples` = attempts per perturbation. Use `[0, 90, -90]` for 90°+ object rotations |
| `gripper_mode` | threshold | `bit` / `threshold` / `continuous` / `recorded` (see algorithm section) |
| `grip_tau` | 0.35 | threshold mode: latch full close once policy's close fraction ≥ τ (use ~0.25 for the base model) |
| `embodiment` | robocasa | `robocasa` (finetuned head) / `oxe_droid` (BASE model via DROID head) |
| `oxe_grip_flip` | false | oxe_droid: invert gripper polarity if the base model's convention is reversed |
| `rel6d` | true | the grasp weld (mandatory for reliable carry — see weld section) |
| `weld_lift_m` | 0.004 | weld engages once both pads touch + object lifted this much + close commanded |
| `weld_solref` | "0.005 1" | weld stiffness injected into the XML |
| `weld_open_frac` / `weld_open_floor` | 0.5 / 0.10 | weld releases when the close command drops below `max(frac·held, floor)` |
| `settle_frames` | 30 | pinned hold after the last frame so the final pose is a rest pose |
| `qvel_mode` | finitediff | qvel written for pinned joints each substep (`zero` = no injected velocity) |

## Code documentation (what each file does)

### `actaug_core.py` — the shared library (no CLI)
Everything both algorithms need, importable, side-effect free:
- **Episode I/O**: `load_episode(root, idx)` → dict with the compiled
  `model.xml`, the full flattened MuJoCo `states` (T×(1+nq+nv)), `ep_meta`,
  the 12-d parquet `actions`, the language `instruction` (tasks.jsonl), and
  `env_args`. `grasp_frame(actions)` = first recorded gripper-close.
- **Env**: `make_env(env_args)` builds the robomimic `EnvRobosuite` (cached per
  env name — construction is minutes); `reset_episode(env, ep)` = `reset_to`
  with the episode's own XML + state 0.
- **Perturbation**: `apply_object_perturbation(env, dyaw, dx, dy)` rotates the
  object's free joint about the world vertical and translates it in-sim;
  `sample_perturbation(rng, ...)` draws random ones; `settle(env)` lets physics
  rest after.
- **Policy bridges**: `groot_obs` / `groot_obs_oxe` build the obs dict
  (3 camera renders + proprio + instruction + `_policy_action_seed`) for the
  finetuned robocasa head / the base model's DROID head;
  `policy_chunk_to_env` / `oxe_chunk_to_env` decode a returned action chunk
  into robosuite env order; `to_env_action` / `env_to_parquet_action` convert
  between parquet layout and env layout.
- **PolicyClient**: minimal ZMQ REQ client for the myGR00T inference server
  (torch.save wire format).
- **Predicates & output**: `obj_z`, `is_success`, `render3`, `groot_state_vec`
  (16-d observation.state), `write_mp4`.

### `rollout.py` — the CLI driver
Parses CLI + `--config` YAML, loads the episode, builds the env, draws the
perturbations, then loops `perturbation × dj7_grid × n_samples`:
- `run_attempt(...)` implements `gt` / `policy` / `hybrid` (GT approach via
  `env.step` → `run_policy` chunks until lift → GT transport), recording every
  executed step;
- for `mode: constrained` it injects the weld into the episode XML
  (`inject_weld`), resets + perturbs, and delegates to
  `constrained.run_constrained_attempt`;
- `save_dump(...)` writes each kept attempt: `actions.npy` (parquet order),
  `states.npz`, `left/right/wrist.mp4` + `3cam.mp4`, `obj_track.npz`
  (executed-vs-recorded object positions, constrained only), `meta.json`;
  `results.csv` is re-flushed after every attempt.

### `constrained.py` — the constrained algorithm
- `inject_weld(xml, solref)`: adds the inactive
  `<weld body1="gripper0_right_eef" body2="<object root>"/>` equality to the
  episode XML (pure text transform, idempotent).
- `PinRig(env)`: all model addressing for one compiled scene — arm/base/torso
  joint qpos/dof addresses, finger actuators (+ open/close ctrl targets),
  object vs gripper geom sets, left/right pad geom sets, joint world axes,
  weld eq id. Methods: `pin_step` (one control frame: lerp the pinned joints
  across physics substeps while gripper + object integrate), `grip_ctrl`,
  `grip_contact`, `pad_contacts`, `weld_engage/release/relpose/active`.
- `run_constrained_attempt(env, ep, client, dj7_init_deg, seed, cfg, record)`:
  the four phases — pinned approach (with the dj7_init ramp), policy window
  (per-joint servo + gripper, chunk every `replan` steps, per-call action
  seed), relative-replay transport (frozen offsets, recorded opening), settle —
  plus the weld engage/release bookkeeping and the result row.

### `compose_2x3.py` — comparison videos
`compose_2x3.py <episode> <dump_dir> <out.mp4>`: row 1 = the ORIGINAL
episode's three dataset videos, row 2 = the augmented rollout's same three
views; the shorter row freezes on its last frame.

### Shell scripts
- `launch_server.sh [gpu] [port] [model] [embodiment] [data_config]` —
  detached GR00T inference server (defaults = the finetuned 60k checkpoint;
  pass the HF snapshot + `oxe_droid oxe_droid` for the base model).
- `run_sanity.sh [gpu] [port] [ep]` — GT replay sanity → hybrid nominal →
  3 random perturbations × 3 samples.

### `configs/`
`unconstrained.yaml` (hybrid preset), `constrained.yaml` (fully commented
constrained preset), `asym_base_j6.yaml` / `asym_bigrot.yaml` (the base-model
asymmetric-object example runs, standard and 90–150° rotation).

## Interpreters / environment

- sim client: `/lp-dev/jonghoon/mimicgen_augment/envs/mimicgen/bin/python`
  with `MUJOCO_GL=egl MUJOCO_EGL_DEVICE_ID=<gpu>` (osmesa is broken on this box).
- policy server: `/data/nvidia/gripper_augmentator/conda-envs/mygr00t/bin/python`
  running `myGR00T/scripts/inference_service.py` (see `launch_server.sh`;
  `HF_HUB_OFFLINE=1` is mandatory). Default checkpoint:
  `/lp-dev/jonghoon/myGR00T_outputs/pnpcountertocab_all502_gbs64_wandb_60k_save5k_20260427_193901/checkpoint-60000`
  (GR00T-N1.5 finetuned on all 502 PnPCounterToCabinet episodes — the same
  dataset these rollouts perturb).
- **BASE (non-finetuned) GR00T-N1.5**: the base checkpoint has NO robocasa
  embodiment head (only gr1 / oxe_droid / agibot_genie1), so it is served via
  its **oxe_droid (DROID) head** with a best-effort obs/action mapping
  (`groot_obs_oxe` / `oxe_chunk_to_env` in actaug_core, same approach as the
  old gripper_state_replay code):
  `bash launch_server.sh <gpu> 8802 <hf-snapshot-dir> oxe_droid oxe_droid`
  and set `constrained.embodiment: oxe_droid` (+ optionally lower
  `grip_tau` to ~0.25 — the base policy's close commands are weaker).
  Expect lower sample efficiency; that's what `n_samples` is for.
- default dataset:
  `/lp-dev/jonghoon/robocasa_full/pickplace_target_human/PickPlaceCounterToCabinet`
  (502 eps, PandaOmron, HYBRID_MOBILE_BASE, fps 20).

## Usage

```bash
# one-time server (GPU 5, port 8801)
bash launch_server.sh 5 8801

SIMPY=/lp-dev/jonghoon/mimicgen_augment/envs/mimicgen/bin/python
export MUJOCO_GL=egl MUJOCO_EGL_DEVICE_ID=2

# sanity: unperturbed GT replay must succeed
$SIMPY rollout.py --episode 0 --mode gt

# the pipeline: 5 random perturbations x 8 policy samples, keep successes
$SIMPY rollout.py --episode 0 --mode hybrid \
    --n_perturbs 5 --n_samples 8 --rot_deg 60 --trans_m 0.05 \
    --pre_buffer 10 --policy_budget 200 --port 8801 --seed 7

# explicit perturbation (yaw +45deg, +3cm x, -2cm y)
$SIMPY rollout.py --episode 0 --mode hybrid --perturb "45,0.03,-0.02" \
    --n_samples 8 --port 8801

# constrained algorithm, all knobs from YAML (CLI flags override YAML values)
$SIMPY rollout.py --config configs/constrained.yaml --episode 0

# or everything at once
bash run_sanity.sh 2 8801 0
```

Outputs: `outputs/<run>/results.csv` (one row per attempt: perturbation, seed,
peak lift, grasped, success, wall time) + `accepted/pertXX_sYY/` dumps.
`--save all` also dumps failures under `rejected/`.

## Controlling how FAST the policy moves (inference speed)

The policy-driven motion speed is set entirely by per-control-step caps in the
`constrained:` block — the sim always runs at the dataset's 20 Hz, so these
caps ARE the deg/s of the wrist during the inference window:

| key | meaning | default | rad/s @20 Hz |
|---|---|---|---|
| `step_cap` | max servo change per control step, every policy joint; also paces the `dj7_init` ramp | 0.05 | 1.0 (≈57°/s) |
| `servo_caps: {j: cap}` | TOTAL authority per joint (how far, not how fast) | 6→0.30, 7→0.60 | — |
| `replan` | control steps between policy queries (reaction latency, not speed) | 8 | — |

History: the port initially reused r02's 0.10 rad/step, but r02 ran at 9.2 Hz —
at robocasa's 20 Hz that is ~2 rad/s and the wrist visibly whips. 0.05 is
rate-matched to the original. Raise/lower `step_cap` in the YAML to make the
policy's wrist motion faster/slower; the gripper close speed is the actuator's
own dynamics (same as GT) and is not affected.

For `hybrid` mode the analogous knob is `exec_h` (env steps executed per chunk)
— the policy's actions are executed at the recorded control rate either way.

## Knobs that matter

- `--pre_buffer` — how early the policy takes over. 10 works for ≤5 cm
  translations; increase (or use `--mode policy`) for larger displacements,
  since the recorded approach aims at the ORIGINAL object location.
- `--policy_budget` — max policy env-steps to achieve the lift (default 200).
- `--resume_offset` — where recorded transport resumes. 0 = at the recorded
  grasp frame. The transport is delta-eef, so it carries the object from
  wherever the policy lifted it.
- Perturbation ranges: prior sweeps (`grip_perturb.py`) used yaw up to ±180°,
  translations ≤ 7 cm, dz ≤ 2 cm. Large yaws are fine for rotationally
  symmetric objects; translations > ~7 cm start leaving the reachable
  counter region.

## Gotchas

- Env construction takes ~2–4 min (kitchen scene compile) — it is cached per
  process; batch many attempts/episodes per process launch.
- `EnvRobosuite.reset_to` internally remaps the export's absolute asset paths
  (`/root/robocasa/...`, `/opt/conda/...`) to the local robocasa/robosuite
  installs — only the mimicgen env's robocasa checkout has all assets; the
  `robocasa_calib` env historically had a geom mismatch. Use the mimicgen python.
- GPU 4 is often another user's — check `nvidia-smi` before picking GPUs.
- `obj_pos` observable + `obj_joint0` joint exist in all robocasa kitchen envs;
  success predicate is the env's own `is_success()["task"]`.