# Socket/Peg Air Insert Implementation Plan

> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.

**Goal:** Convert the air-insert task from ring/bar objects to the XML-defined socket/peg objects and rename the canonical task to `sim_air_insert_socket_peg`.

**Architecture:** Keep the existing Diana dual-arm air-insert environment and policy flow, but replace all task-state keys, MuJoCo joint/geom names, reward names, and registration names with socket/peg semantics. Use the socket's internal `pin` geom as the success detector: reward reaches max only when `red_peg` contacts `pin` while both objects are airborne. Keep a narrowly scoped backward-compatible robot class alias only if needed to avoid import breakage, but make socket/peg the canonical name.

**Tech Stack:** Python unittest, NumPy, MuJoCo XML assets, existing `PolicyBase` interpolation behavior.

---

### Task 1: Rename task registration and reset state API

**Files:**
- Modify: `roboimi/utils/constants.py`
- Modify: `roboimi/utils/act_ex_utils.py`
- Modify: `roboimi/demos/diana_record_sim_episodes.py`
- Modify: `roboimi/demos/vla_scripts/eval_vla.py`
- Modify: `roboimi/envs/double_pos_ctrl_env.py`
- Test: `tests/test_air_insert_env.py`
- Test: `tests/test_eval_vla_headless.py`

- [ ] Write/adjust tests to expect `sim_air_insert_socket_peg`, `sample_air_insert_socket_peg_state()`, and task-state keys `socket_pos/socket_quat/peg_pos/peg_quat`.
- [ ] Run targeted tests and confirm they fail before implementation.
- [ ] Implement the canonical task-name and sampler rename across registration, rollout, and eval paths.
- [ ] Run targeted tests and confirm they pass.

### Task 2: Switch environment helpers/reward to socket/peg MuJoCo names

**Files:**
- Modify: `roboimi/envs/double_air_insert_env.py`
- Test: `tests/test_air_insert_env.py`

- [ ] Write/adjust tests for `set_socket_peg_task_state()`, `get_socket_peg_env_state()`, joint names `blue_socket_joint/red_peg_joint`, geom names `socket-1..4`, `red_peg`, and success contact `red_peg` ↔ `pin`.
- [ ] Run targeted tests and confirm they fail before implementation.
- [ ] Implement helper and reward rename; remove ring/bar aperture geometry success logic.
- [ ] Run targeted tests and confirm they pass.

### Task 3: Switch robot/XML asset names to socket/peg

**Files:**
- Move/Modify: `roboimi/assets/models/manipulators/DianaMed/ring_bar_objects.xml` -> `roboimi/assets/models/manipulators/DianaMed/socket_peg_objects.xml`
- Move/Modify: `roboimi/assets/models/manipulators/DianaMed/bi_diana_ring_bar_ee.xml` -> `roboimi/assets/models/manipulators/DianaMed/bi_diana_socket_peg_ee.xml`
- Modify: `roboimi/assets/robots/diana_med.py`
- Test: `tests/test_robot_asset_paths.py`
- Test: `tests/test_air_insert_env.py`

- [ ] Write/adjust tests to expect `BiDianaMedSocketPeg` and `bi_diana_socket_peg_ee.xml` including `socket_peg_objects.xml`.
- [ ] Run targeted tests and confirm they fail before implementation.
- [ ] Rename XML assets, update include/model names, set default object positions on the table, and update robot class/path.
- [ ] Run targeted tests and confirm they pass.

### Task 4: Switch scripted policy to socket/peg

**Files:**
- Modify: `roboimi/demos/diana_air_insert_policy.py`
- Test: `tests/test_air_insert_env.py`

- [ ] Write/adjust tests to ensure left hand grasps socket, right hand grasps peg, and the insertion trajectory drives peg toward socket/pin.
- [ ] Run targeted policy tests and confirm they fail before implementation.
- [ ] Rename policy constants/locals from ring/bar to socket/peg, keep the existing timing/interpolation style, and use socket hold + peg insert waypoints.
- [ ] Run targeted policy tests and confirm they pass.

### Task 5: Full verification

**Files:**
- All modified files above.

- [ ] Run `python -m unittest tests.test_air_insert_env tests.test_robot_asset_paths tests.test_eval_vla_headless -v`.
- [ ] Run a grep check that canonical code no longer references `sim_air_insert_ring_bar`, `sample_air_insert_ring_bar_state`, `ring_block`, or `bar_block` except deliberate compatibility aliases/comments if any.
- [ ] Report exact verification output and any remaining caveats.