roboimi/docs/superpowers/plans/2026-04-30-socket-peg-air-insert.md

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.