merge: socket-peg sim and imf attnres into main

# Conflicts:
#	roboimi/demos/vla_scripts/eval_vla.py
#	roboimi/envs/double_base.py

docs/superpowers/plans/2026-04-23-sim-air-insert-ring-bar.md

@@ -0,0 +1,311 @@
+# sim_air_insert_ring_bar Implementation Plan
+
+> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.
+
+**Goal:** Add an independent dual-Diana MuJoCo task `sim_air_insert_ring_bar` with a square ring block, a square bar block, staged rewards, strict finite-geometry in-air insertion success detection, and a task-specific scripted policy.
+
+**Architecture:** Reuse the current dual-Diana EE-control stack and environment factory, but add a task-specific scene XML, robot asset entrypoint, sampling helpers, and a new task-specific environment module. Keep `sim_transfer` untouched while introducing pure-Python geometry helpers and focused tests so reward/success behavior can be regression tested without requiring a full MuJoCo rollout in every test.
+
+**Tech Stack:** Python, unittest, MuJoCo XML assets, existing dual-Diana environment classes, Hydra-compatible task naming/config patterns.
+
+---
+
+## File Structure / Responsibilities
+
+- **Create:** `roboimi/assets/models/manipulators/DianaMed/ring_bar_objects.xml`
+  - Defines the rigid ring body and bar body, each with a free joint and stable box-based geoms.
+- **Create:** `roboimi/assets/models/manipulators/DianaMed/bi_diana_ring_bar_ee.xml`
+  - Scene entrypoint that includes the shared world/table/robot assets plus the new object XML.
+- **Modify:** `roboimi/assets/robots/diana_med.py`
+  - Add a task-specific robot asset class for the new scene XML without changing existing `BiDianaMed` behavior.
+- **Modify:** `roboimi/utils/act_ex_utils.py`
+  - Add deterministic helpers to sample left/right planar placement regions for ring and bar objects.
+- **Modify:** `roboimi/utils/constants.py`
+  - Register the new task name and default metadata.
+- **Create:** `roboimi/envs/double_air_insert_env.py`
+  - New task-specific environment, finite-geometry success helpers, reset logic, reward logic, and task factory branch.
+- **Modify:** `roboimi/envs/double_pos_ctrl_env.py`
+  - Route `make_sim_env()` to the new task-specific environment while keeping current `sim_transfer` logic unchanged.
+- **Create:** `roboimi/demos/diana_air_insert_policy.py`
+  - Task-specific waypoint/open-loop scripted policy for grasp-lift-align-insert.
+- **Modify:** `roboimi/demos/vla_scripts/eval_vla.py`
+  - Reset the new task with the correct sampled task state instead of assuming a single transfer box pose.
+- **Create:** `tests/test_air_insert_env.py`
+  - Focused unit tests for sampling, reset helpers, reward progression, and strict success detection.
+- **Modify:** `tests/test_eval_vla_headless.py`
+  - Add coverage that headless evaluation dispatches the correct reset sampler for the new task.
+- **Modify:** `tests/test_robot_asset_paths.py`
+  - Verify the new robot asset class resolves its XML path correctly independent of cwd.
+
+---
+
+### Task 1: Add failing tests for task registration, samplers, and asset wiring
+
+**Files:**
+- Create: `tests/test_air_insert_env.py`
+- Modify: `tests/test_eval_vla_headless.py`
+- Modify: `tests/test_robot_asset_paths.py`
+- Modify: `roboimi/utils/act_ex_utils.py` (later in implementation)
+- Modify: `roboimi/utils/constants.py` (later in implementation)
+- Modify: `roboimi/assets/robots/diana_med.py` (later in implementation)
+- Modify: `roboimi/envs/double_pos_ctrl_env.py` (later in implementation)
+- Create: `roboimi/envs/double_air_insert_env.py` (minimal stub in this task)
+
+- [ ] **Step 1: Write failing tests for task config and sampling helpers**
+
+Add tests in `tests/test_air_insert_env.py` covering:
+- `SIM_TASK_CONFIGS['sim_air_insert_ring_bar']` exists
+- `sample_air_insert_ring_bar_pose()` (or equivalent helper) returns ring/bar positions with fixed z and correct left/right planar ranges
+- output structure is explicit and easy for reset/eval code to consume
+
+- [ ] **Step 2: Write failing tests for environment factory dispatch and robot asset resolution**
+
+Add tests covering:
+- `make_sim_env('sim_air_insert_ring_bar', headless=True)` dispatches to the new environment with rendering disabled
+- a new robot asset class resolves the new XML path independent of cwd, similar to the existing `BiDianaMed` test pattern
+
+- [ ] **Step 3: Write failing tests for eval reset helper dispatch**
+
+Extend `tests/test_eval_vla_headless.py` so headless eval can reset the new task using the new sampler instead of hard-coding `sample_transfer_pose()`.
+
+- [ ] **Step 4: Run the targeted tests to verify they fail for the expected missing-feature reasons**
+
+Run:
+`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env tests.test_eval_vla_headless tests.test_robot_asset_paths -v`
+
+Expected:
+- FAIL because the new task config/helper/class/dispatch branch does not exist yet
+
+- [ ] **Step 5: Implement the minimal production code to satisfy the new task registration and helper tests**
+
+Implement only enough to make the new tests pass:
+- add new task config entry
+- add the new placement sampler
+- add the new robot asset class
+- create a minimal importable `double_air_insert_env.py` stub and class/function surface needed for factory dispatch tests
+- add the factory dispatch branch / headless wiring
+- update eval reset dispatch for the new task
+
+- [ ] **Step 6: Re-run the targeted tests to verify they pass**
+
+Run:
+`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env tests.test_eval_vla_headless tests.test_robot_asset_paths -v`
+
+Expected:
+- PASS for the new registration/sampler/dispatch/asset tests
+
+- [ ] **Step 7: Commit Task 1**
+
+Run:
+`git add tests/test_air_insert_env.py tests/test_eval_vla_headless.py tests/test_robot_asset_paths.py roboimi/utils/act_ex_utils.py roboimi/utils/constants.py roboimi/assets/robots/diana_med.py roboimi/envs/double_pos_ctrl_env.py roboimi/envs/double_air_insert_env.py roboimi/demos/vla_scripts/eval_vla.py && git commit -m "feat(env): register sim air insert ring bar task"`
+
+---
+
+### Task 2: Add the MuJoCo ring+bar scene assets and reset helpers
+
+**Files:**
+- Create: `roboimi/assets/models/manipulators/DianaMed/ring_bar_objects.xml`
+- Create: `roboimi/assets/models/manipulators/DianaMed/bi_diana_ring_bar_ee.xml`
+- Create or Modify: `roboimi/envs/double_air_insert_env.py`
+- Modify: `tests/test_air_insert_env.py`
+
+- [ ] **Step 1: Write failing tests for object reset helpers and scene-specific joint naming assumptions**
+
+In `tests/test_air_insert_env.py`, add unit tests for helper functions that:
+- write ring pose to `ring_block_joint` from the named task-state mapping
+- write bar pose to `bar_block_joint` from the named task-state mapping
+- read back `env_state` as a stable 14D vector `[ring_pos, ring_quat, bar_pos, bar_quat]`
+
+Use fake `mj_data` objects so tests stay fast and deterministic.
+
+- [ ] **Step 2: Run the focused test slice and verify it fails**
+
+Run:
+`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env -v`
+
+Expected:
+- FAIL because reset/state helper functions and joint conventions are not implemented yet
+
+- [ ] **Step 3: Implement the scene XML files and reset/state helper code**
+
+Implement:
+- the object XML with one rigid ring body and one rigid bar body
+- the task scene XML entrypoint using the shared world/table/robot includes
+- reset helper(s) in `double_air_insert_env.py` that set qpos for both free joints with fixed quaternions
+- task-state accessor(s) returning both object poses in a stable structure
+
+- [ ] **Step 4: Re-run the focused test slice and verify it passes**
+
+Run:
+`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env -v`
+
+Expected:
+- PASS for reset/state helper tests
+
+- [ ] **Step 5: Commit Task 2**
+
+Run:
+`git add roboimi/assets/models/manipulators/DianaMed/ring_bar_objects.xml roboimi/assets/models/manipulators/DianaMed/bi_diana_ring_bar_ee.xml roboimi/envs/double_air_insert_env.py tests/test_air_insert_env.py && git commit -m "feat(scene): add ring and bar insertion scene assets"`
+
+---
+
+### Task 3: Implement strict reward and finite-geometry success detection
+
+**Files:**
+- Modify: `roboimi/envs/double_air_insert_env.py`
+- Modify: `tests/test_air_insert_env.py`
+
+- [ ] **Step 1: Write failing tests for reward stages and strict success detection**
+
+Add tests in `tests/test_air_insert_env.py` for:
+- left contact stage reward
+- right contact stage reward
+- ring lifted off table stage
+- bar lifted off table stage
+- positive success case where a finite bar truly passes through the aperture
+- negative case where the centerline would pass but the finite square body would clip
+- negative case where the bar has not crossed the ring thickness direction enough
+- negative case where one/both objects are still on the table
+
+Structure the tests around pure helper functions and light fake contact/state objects so the geometry logic is directly regression tested.
+
+- [ ] **Step 2: Run the focused tests and verify they fail for missing reward/success logic**
+
+Run:
+`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env -v`
+
+Expected:
+- FAIL because the staged reward and finite-geometry insertion logic are not implemented yet
+
+- [ ] **Step 3: Implement minimal strict success helpers and reward logic**
+
+Implement in `roboimi/envs/double_air_insert_env.py`:
+- pure helper(s) for transforming bar geometry into ring-local coordinates
+- finite-geometry insertion predicate (not centerline-only)
+- table-contact / airborne checks
+- staged reward function returning the highest achieved stage with `max_reward = 5`
+
+- [ ] **Step 4: Re-run the focused tests to verify the logic passes**
+
+Run:
+`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env -v`
+
+Expected:
+- PASS for reward and success-detection regression tests
+
+- [ ] **Step 5: Commit Task 3**
+
+Run:
+`git add roboimi/envs/double_air_insert_env.py tests/test_air_insert_env.py && git commit -m "feat(env): add strict air insertion reward and success logic"`
+
+---
+
+### Task 4: Add the scripted policy and integration smoke coverage
+
+**Files:**
+- Create: `roboimi/demos/diana_air_insert_policy.py`
+- Modify: `roboimi/demos/diana_record_sim_episodes.py`
+- Modify: `tests/test_air_insert_env.py`
+- Optionally Modify: `roboimi/demos/vla_scripts/eval_vla.py` (only if integration gaps remain after Task 1)
+
+- [ ] **Step 1: Write failing tests for scripted-policy action shape and basic generation**
+
+Add tests covering:
+- the new policy produces a 16D action
+- trajectory generation accepts sampled named task state without error
+- the first action is a valid open-gripper safe pose command
+- a deterministic nominal smoke path (with canonical sampled state or fake env shim) reaches the intended terminal interface contract without shape/reward mismatches
+
+Keep the tests unit-level; do not require a full MuJoCo rollout for every assertion.
+
+- [ ] **Step 2: Write failing tests for the scripted rollout entrypoint and a real headless smoke path**
+
+Add coverage for both:
+- the standard scripted rollout entrypoint (`roboimi/demos/diana_record_sim_episodes.py`) can select the new task sampler/policy instead of remaining sim_transfer-only
+- a deterministic integration/smoke test that instantiates `make_sim_env('sim_air_insert_ring_bar', headless=True)`, resets with sampled named task state, and steps a few actions or scripted-policy outputs using the real task XML and task-specific wiring
+
+- [ ] **Step 3: Run the scripted-policy tests and verify they fail**
+
+Run:
+`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env -v`
+
+Expected:
+- FAIL because the new scripted policy does not exist yet
+
+- [ ] **Step 4: Implement the waypoint-based scripted policy**
+
+Implement a conservative open-loop policy with phases:
+- safe wait pose
+- above-target approach
+- descend + grasp
+- dual lift
+- airborne meeting alignment
+- bar push-through insertion
+
+Use fixed orientations for version 1 and follow the existing repository style from `diana_policy.py`.
+
+- [ ] **Step 5: Re-run the scripted-policy tests to verify they pass**
+
+Run:
+`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env -v`
+
+Expected:
+- PASS for scripted-policy tests
+
+- [ ] **Step 6: Run the combined verification suite for this feature**
+
+Run:
+`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env tests.test_eval_vla_headless tests.test_eval_vla_rollout_artifacts tests.test_train_vla_rollout_validation tests.test_robot_asset_paths -v`
+
+Expected:
+- PASS with 0 failures
+
+- [ ] **Step 6b: Run the mandatory real headless smoke check**
+
+Run a focused smoke command that instantiates the real task, resets with sampled state, and steps a few actions using the new scripted policy or a deterministic action sequence.
+
+Example command (adjust module/test helper if needed):
+`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env.AirInsertEnvSmokeTest -v`
+
+Expected:
+- PASS, proving the real XML/assets/env wiring instantiate and step correctly in headless mode
+
+- [ ] **Step 7: Commit Task 4**
+
+Run:
+`git add roboimi/demos/diana_air_insert_policy.py tests/test_air_insert_env.py tests/test_eval_vla_headless.py tests/test_robot_asset_paths.py roboimi/demos/vla_scripts/eval_vla.py && git commit -m "feat(policy): add scripted air insertion policy"`
+
+---
+
+### Task 5: Final verification and implementation review
+
+**Files:**
+- Review all files touched above
+
+- [ ] **Step 1: Run fresh end-to-end verification before claiming completion**
+
+Run:
+`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env tests.test_eval_vla_headless tests.test_robot_asset_paths -v`
+
+Expected:
+- PASS with 0 failures
+
+- [ ] **Step 2: Inspect git status and recent commits**
+
+Run:
+`git status --short && git log --oneline --decorate -n 8`
+
+Expected:
+- only intended feature files modified / committed
+
+- [ ] **Step 3: Request final code review for the completed feature**
+
+Use the requesting-code-review skill against the full diff from the feature branch starting point to current HEAD.
+
+- [ ] **Step 4: Address any review findings and re-run verification if code changes**
+
+If fixes are made, repeat the unittest command from Step 1.
+
+- [ ] **Step 5: Hand off using finishing-a-development-branch**
+
+After verification and review, use the finishing-a-development-branch skill to decide merge / PR / cleanup.

docs/superpowers/specs/2026-04-23-sim-air-insert-ring-bar-design.md

@@ -0,0 +1,316 @@
+# sim_air_insert_ring_bar Design
+
+## Summary
+
+Add a new independent MuJoCo simulation task named `sim_air_insert_ring_bar` that keeps the existing dual-Diana tabletop setup but replaces the single transfer box with two randomized objects:
+
+- a square ring block grasped by the left arm
+- a square bar block grasped by the right arm
+
+The task is to pick both objects off the table and complete an in-air insertion where the bar truly passes through the ring aperture. The existing `sim_transfer` task must remain unchanged.
+
+## Goals
+
+- Reuse the current dual-Diana EE-control simulation stack
+- Keep the same table/base robot arrangement as the existing transfer task
+- Add an independent task entrypoint and scene definition
+- Randomize planar placement of both objects within left/right task-specific regions
+- Implement reward staging for contact, lift, and successful in-air insertion
+- Add a scripted policy that performs pick, lift, align, and in-air insertion
+- Preserve compatibility with existing environment creation, evaluation, and rollout patterns
+
+## Non-Goals
+
+- No random yaw in the first version
+- No visual servoing or closed-loop insertion controller
+- No general multi-task environment framework refactor
+- No guarantee that the VLA training stack is immediately tuned for this new task
+- No replacement or behavior change for `sim_transfer`
+
+## Task Name
+
+Use a new task name:
+
+- `sim_air_insert_ring_bar`
+
+This task should be exposed alongside `sim_transfer`, not as a replacement.
+
+## Scene Geometry
+
+### Shared Base Scene
+
+Keep the dual Diana robot, the table, and the existing camera layout conceptually unchanged.
+
+### Ring Block
+
+Represent the square ring as a rigid free body composed from simple MuJoCo box geoms rather than an external mesh.
+
+Dimensions:
+
+- outer side length: 68 mm
+- inner aperture side length: 32 mm
+- thickness: 18 mm
+- ring wall width: 18 mm
+
+The ring should behave as a single object body with a single free joint.
+
+### Bar Block
+
+Represent the bar as a rigid free body with a single box geom.
+
+Dimensions:
+
+- length: 90 mm
+- cross-section: 18 mm x 18 mm
+
+The bar should also be a single free-joint body.
+
+## Initial Placement / Reset
+
+The first version uses position-only randomization with fixed orientation. Reset sampling stays **caller-driven**, matching the existing `sim_transfer` usage pattern in rollout/eval code: a helper samples task state, then callers pass that state into `env.reset(...)`.
+
+Use an explicit sampled task-state structure with named fields:
+
+- `ring_pos`: 3D position
+- `ring_quat`: fixed 4D quaternion for version 1
+- `bar_pos`: 3D position
+- `bar_quat`: fixed 4D quaternion for version 1
+
+Behavior:
+
+- ring block: randomized only in a left-side planar sampling region
+- bar block: randomized only in a right-side planar sampling region
+- both objects start flat on the table
+- both objects use fixed orientation at reset
+- no random yaw, tilt, or flip in this version
+
+The sampling regions should be chosen conservatively so that:
+
+- the left arm can comfortably reach and grasp the ring
+- the right arm can comfortably reach and grasp the bar
+- scripted open-loop pick trajectories remain feasible
+
+## Control / Action Interface
+
+Reuse the current 16D EE-space action convention already used by the dual-Diana position-control environment:
+
+- left arm EE pose: 7D (`xyz + quat`)
+- right arm EE pose: 7D (`xyz + quat`)
+- left gripper command: 1D
+- right gripper command: 1D
+
+The new task should continue using EE targets transformed through the existing IK-based control path.
+
+## Environment Structure
+
+Implement this as a new task-specific environment path while reusing the existing dual-Diana simulation base where possible.
+
+Expected responsibilities:
+
+- scene instantiation for the ring+bar setup
+- task reset for randomized object placement
+- environment-state accessors for both objects
+- reward computation
+- in-air insertion success detection
+
+The environment factory must dispatch by task name and leave the `sim_transfer` branch unchanged.
+
+## Observation / Environment State
+
+The task should retain the current observation structure style used by the dual-Diana environment:
+
+- `qpos`
+- multi-camera images
+
+For task state access, the environment should expose a stable `env_state` vector with this exact order:
+
+- `ring_pos[0:3]`
+- `ring_quat[3:7]`
+- `bar_pos[7:10]`
+- `bar_quat[10:14]`
+
+This 14D state should be sufficient for scripted-policy debugging and future rollout analysis, while reset itself remains caller-driven via the named task-state helper structure above.
+
+## Reward Design
+
+Use staged rewards in the same spirit as the current task, returning the highest achieved stage rather than accumulating one-time sparse bonuses per event.
+
+Maximum reward:
+
+- `max_reward = 5`
+
+Reward stages:
+
+1. left gripper touches the ring block
+2. right gripper touches the bar block
+3. ring block is lifted off the table
+4. bar block is lifted off the table
+5. while both objects are off the table, the bar truly passes through the ring aperture
+
+Notes:
+
+- contact rewards are intended as grasp-progress stages
+- lift rewards require the object to be off the table, not merely touched
+- final success reward only applies when both objects are airborne
+
+## Success Detection
+
+Success must **not** be based on a centerline-only check.
+
+A centerline-only test is insufficient because:
+
+- the bar has thickness, so a centerline can pass through while the body cannot
+- a square bar with imperfect orientation can have its centerline inside the aperture while its corners still collide with the ring
+
+### Required Success Semantics
+
+A successful insertion requires all of the following:
+
+1. the ring is off the table
+2. the bar is off the table
+3. the bar has actually crossed through the ring thickness direction
+4. the bar’s finite square cross-section fits through the square aperture during that crossing
+
+### Recommended Detection Approach
+
+Use a task-level geometric check in Python rather than relying on contact alone.
+
+Implementation intent:
+
+- transform the bar geometry into the ring’s local frame
+- reason about the bar as a finite oriented box (not a line)
+- verify that the bar has crossed the ring thickness direction
+- verify that the portion of the bar passing the aperture fits within the inner square opening, accounting for the bar’s cross-section and orientation
+
+This geometric check is the primary success test.
+
+### Role of Contacts
+
+Contacts may still be used for:
+
+- grasp-stage rewards
+- debugging / diagnostics
+
+But contact alone should **not** be the sole criterion for insertion success, since:
+
+- a true clean insertion may have limited aperture-wall contact
+- persistent contact can also happen while the bar is jammed and not actually inserted
+
+## Scripted Policy
+
+Add a new task-specific scripted policy for `sim_air_insert_ring_bar`.
+
+### Policy Intent
+
+The first version prioritizes a conservative, reliable open-loop demonstration rather than an optimized trajectory.
+
+### Action Phases
+
+The scripted policy should follow these phases:
+
+1. move both arms to safe initial / waiting poses with grippers open
+2. move left arm above the ring and right arm above the bar
+3. descend and grasp the assigned objects
+4. lift both objects clear of the table
+5. move both objects to an airborne meeting region above the table
+6. hold the ring stably while aligning the bar with the aperture
+7. push the bar along the intended insertion direction until the geometric success condition is met
+
+### Grasp Assignment
+
+- left arm: ring only
+- right arm: bar only
+
+### Motion Style
+
+Keep the current repository style:
+
+- waypoint-based trajectory definition
+- open-loop interpolation between waypoints
+- fixed grasp orientation in the first version
+
+No adaptive replanning is required for the first version.
+
+## Files / Integration Scope
+
+The implementation is expected to add task-specific files rather than broadly refactoring the codebase.
+
+Likely additions / changes:
+
+- a new MuJoCo scene XML for the ring+bar task
+- one or more XML fragments defining the two new objects
+- a new task-specific dual-Diana environment file
+- robot asset wiring for the new scene XML
+- reset sampling helpers for the new task
+- task registration in constants / environment factory paths
+- a new scripted policy file
+- focused tests for task creation, reset, rewards, success detection, and scripted policy shape/smoke behavior
+
+## Testing Requirements
+
+At minimum, add regression coverage for:
+
+### Environment Creation
+
+- the new task can be created via the task factory
+- the existing `sim_transfer` task remains unchanged
+
+### Reset / Sampling
+
+- ring reset positions are inside the left sampling region
+- bar reset positions are inside the right sampling region
+- reset orientation is fixed as intended
+
+### Environment State
+
+- environment-state access returns both object poses in the expected structure
+
+### Success Detection
+
+Must include both positive and negative cases.
+
+Positive case:
+
+- a configuration where the finite bar truly passes through the ring aperture is detected as success
+
+Negative cases:
+
+- centerline-inside but finite body would clip the aperture
+- not enough depth / not actually crossing the ring thickness direction
+- one or both objects still on the table
+
+### Reward Logic
+
+- left contact stage
+- right contact stage
+- ring lift stage
+- bar lift stage
+- final success stage with `max_reward = 5`
+
+### Scripted Policy
+
+At minimum:
+
+- policy emits valid 16D actions
+- trajectory generation does not error
+- rollout smoke path can step through the new environment
+
+## Risks / Constraints
+
+- MuJoCo contact naming must remain stable enough for stage rewards
+- geometric insertion checks must be strict enough to avoid false positives but not so brittle that numerically valid insertions are missed
+- scripted open-loop insertion may require conservative alignment and lift heights to keep the first version reliable
+
+## Acceptance Criteria
+
+The feature is complete when all of the following are true:
+
+- `sim_air_insert_ring_bar` is creatable as an independent task
+- the scene contains the dual Diana, table, ring block, and bar block
+- reset randomizes ring and bar positions in left/right planar regions with fixed orientation
+- the environment exposes task state for both objects
+- staged rewards progress to `max_reward = 5`
+- final success is based on finite-geometry insertion semantics, not a centerline-only shortcut
+- a new scripted policy can execute the intended pick-lift-align-insert behavior in the new environment
+- a canonical nominal smoke path (unit-level or deterministic integration-level) exists for the new scripted-policy interface so success is not judged purely by interpretation
+- existing `sim_transfer` behavior is preserved

roboimi/assets/models/manipulators/DianaMed/BiDianaMed_rethink.xml

@@ -76,7 +76,7 @@
                                                 <body name="ee_cam_left" pos="0.00 0.046 -0.075" euler="0.0 0.0 -0.0">
                                                         <inertial pos="0 0 0" quat="1 0 0 0" mass="0" diaginertia="0 0 0" />
                                                         <geom type="mesh" contype="1" conaffinity="1" group="1" rgba="0.69804 0.69804 0.69804 1" mesh="realsense_cam" />
-                                                        <camera name="rs_cam_left" mode="fixed" pos="0.0 0.0 -0.25" euler="0.0 9.4 0.0 " fovy="15" resolution="1920 1200"/>
+                                                        <camera name="rs_cam_left" mode="fixed" pos="0.0 0.0 0.01" euler="0.0 9.4 0.0 " fovy="50" resolution="1920 1200"/>
                                                 </body>
                                             </body>
                                             <body name="l_finger_left" pos="0 0.01 0.0444">

roboimi/assets/models/manipulators/DianaMed/bi_diana_socket_peg_ee.xml

@@ -0,0 +1,6 @@
+<mujoco model="bi_diana_socket_peg">
+    <include file="./empty_world.xml" />
+    <include file="./table_square.xml" />
+    <include file="./socket_peg_objects.xml" />
+    <include file="./BiDianaMed_rethink.xml" />
+</mujoco>

roboimi/assets/models/manipulators/DianaMed/socket_peg_objects.xml

@@ -0,0 +1,19 @@
+<mujoco model="socket_peg_objects">
+  <worldbody>
+     <body name="peg" pos="0.12 0.90 0.46">
+          <joint name="red_peg_joint" type="free" frictionloss="0.01" />
+          <inertial pos="0 0 0" mass="0.05" diaginertia="0.002 0.002 0.002" />
+          <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.005 0.0001" pos="0 0 0" size="0.06 0.01 0.01" type="box" name="red_peg" rgba="1 0 0 1" />
+     </body>
+
+     <body name="socket" pos="-0.12 0.90 0.472">
+          <joint name="blue_socket_joint" type="free" frictionloss="0.01" />
+          <inertial pos="0 0 0" mass="0.05" diaginertia="0.002 0.002 0.002" />
+          <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.05 0.001" pos="0 0 -0.02" size="0.06 0.018 0.002" type="box" name="socket-1" rgba="0 0 1 1" />
+          <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.05 0.001" pos="0 0 0.02" size="0.06 0.018 0.002" type="box" name="socket-2" rgba="0 0 1 1" />
+          <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.05 0.001" pos="0 0.02 0" size="0.06 0.002 0.018" type="box" name="socket-3" rgba="0 0 1 1" />
+          <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.05 0.001" pos="0 -0.02 0" size="0.06 0.002 0.018" type="box" name="socket-4" rgba="0 0 1 1" />
+          <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.005 0.0001" pos="0 0 0" size="0.04 0.01 0.01" type="box" name="pin" rgba="1 0 0 1" />
+     </body>
+  </worldbody>
+</mujoco>

roboimi/assets/models/manipulators/DianaMed/table_square.xml

@@ -7,7 +7,6 @@
           <geom name="table" condim="4" contype="1" conaffinity="1"  type="box" rgba="0.4 0.4 0.4 1" size="0.62 0.62 0.01" density="1500" friction="0.9 0.9 0.9"/>
       </body>
       <camera name="top" pos="0.0 1.0 2.0" fovy="44" mode="targetbody" target="table"/>
-      <camera name="angle" pos="0.0 0.0 2.0" fovy="37" mode="targetbody" target="table"/>
       <camera name="front" pos="0 0 0.8" fovy="65" mode="fixed" quat="0.7071 0.7071 0 0"/>
   </worldbody>
 </mujoco>

roboimi/assets/robots/diana_med.py

@@ -91,3 +91,39 @@ class BiDianaMed(ArmBase):
         """ Robot's init joint position. """
         return np.array([0.0, 0.0, 0.0, 1.57, 0.0, 0.0, 0.0])
 
+
+class BiDianaMedSocketPeg(ArmBase):
+    def __init__(self):
+        super().__init__(
+            name="Bidiana_socket_peg",
+            urdf_path="roboimi/assets/models/manipulators/DianaMed/DualDianaMed.urdf",
+            xml_path="roboimi/assets/models/manipulators/DianaMed/bi_diana_socket_peg_ee.xml",
+            gripper=None
+        )
+        self.left_arm = self.Arm(self, 'single', self.urdf_path)
+        self.left_arm.set_Arm_base_link('left_base_link')
+        self.left_arm.set_Arm_ee_link('left_link7')
+        self.left_arm.InitKDL
+        self.left_arm.joint_index = ['l_j1','l_j2','l_j3','l_j4','l_j5','l_j6','l_j7']
+        self.left_arm.gripper_index = ['l_finger_joint_left','r_finger_joint_left']
+        self.left_arm.actuator_index = ['a1_l','a2_l','a3_l','a4_l','a5_l','a6_l','a7_l','gripper_left']
+        self.left_arm.setArmInitPose(self.init_qpos)
+        self.arms.append(self.left_arm)
+        self.right_arm = self.Arm(self,'single', self.urdf_path)
+        self.right_arm.set_Arm_base_link('right_base_link')
+        self.right_arm.set_Arm_ee_link('right_link7')
+        self.right_arm.InitKDL
+        self.right_arm.joint_index = ['r_j1','r_j2','r_j3','r_j4','r_j5','r_j6','r_j7']
+        self.right_arm.gripper_index = ['l_finger_joint_right','r_finger_joint_right']
+        self.right_arm.actuator_index = ['a1_r','a2_r','a3_r','a4_r','a5_r','a6_r','a7_r','gripper_right']
+        self.right_arm.setArmInitPose(self.init_qpos)
+        self.arms.append(self.right_arm)
+        self.jnt_num = self.left_arm.jnt_num + self.right_arm.jnt_num
+        self.kp = 500 * np.ones(self.jnt_num)
+        self.kd = 44.57 * np.ones(self.jnt_num)
+
+    @property
+    def init_qpos(self):
+        """ Robot's init joint position. """
+        return np.array([0.0, 0.0, 0.0, 1.57, 0.0, 0.0, 0.0])
+    

roboimi/demos/diana_air_insert_policy.py

@@ -0,0 +1,203 @@
+import numpy as np
+from pyquaternion import Quaternion
+
+from roboimi.demos.diana_policy import PolicyBase
+
+
+class TestAirInsertPolicy(PolicyBase):
+    ACTION_OBJECT_Z_OFFSET = 0.078
+    SOCKET_GRASP_OFFSET = np.array([0.0, 0.0, 0.0], dtype=np.float64)
+    PEG_GRASP_OFFSET = np.array([0.0, 0.0, 0.0], dtype=np.float64)
+    SOCKET_OUTER_GRASP_STRATEGY = "socket_outer"
+    LEGACY_GRASP_STRATEGY = "legacy"
+    SOCKET_HOLD_Z = 0.85
+    PEG_INSERT_START_OFFSET = np.array([0.105, 0.0, 0.0], dtype=np.float64)
+    INSERT_START_T = 650
+    INSERT_END_T = 730
+    LEFT_SOCKET_GRIPPER_CLOSED = -100
+    RIGHT_PEG_GRIPPER_CLOSED = -100
+    SOCKET_APPROACH_Z = 1.05
+    EPISODE_END_T = 1000
+
+    def __init__(self, inject_noise=False, grasp_strategy=SOCKET_OUTER_GRASP_STRATEGY):
+        super().__init__(inject_noise=inject_noise)
+        valid_strategies = {
+            self.SOCKET_OUTER_GRASP_STRATEGY,
+            self.LEGACY_GRASP_STRATEGY,
+        }
+        if grasp_strategy not in valid_strategies:
+            raise ValueError(
+                f"Unsupported air insert grasp_strategy={grasp_strategy!r}; "
+                f"expected one of {sorted(valid_strategies)}"
+            )
+        self.grasp_strategy = grasp_strategy
+
+    def generate_trajectory(self, task_state):
+        return self._generate_socket_peg_trajectory(task_state)
+
+    def _generate_socket_peg_trajectory(self, task_state):
+        socket_xyz = np.asarray(task_state["socket_pos"], dtype=np.float64)
+        peg_xyz = np.asarray(task_state["peg_pos"], dtype=np.float64)
+
+        init_mocap_pose_left = np.array(
+            [
+                -0.17297014,
+                1.00485877,
+                1.32773627,
+                7.06825181e-01,
+                8.20281078e-06,
+                -7.07388269e-01,
+                -5.20399313e-06,
+            ],
+            dtype=np.float64,
+        )
+        init_mocap_pose_right = np.array(
+            [
+                0.17297014,
+                0.9951369,
+                1.32773623,
+                2.59463975e-06,
+                7.07388269e-01,
+                5.59551158e-06,
+                7.06825181e-01,
+            ],
+            dtype=np.float64,
+        )
+
+        left_init_quat = Quaternion(init_mocap_pose_left[3:])
+        right_init_quat = Quaternion(init_mocap_pose_right[3:])
+
+        left_pick_quat = (
+            left_init_quat * Quaternion(axis=[0.0, 1.0, 0.0], degrees=45)
+        ).elements
+        right_pick_quat = (
+            right_init_quat * Quaternion(axis=[0.0, 1.0, 0.0], degrees=45)
+        ).elements
+
+        socket_hold_action = np.array(
+            [socket_xyz[0] - 0.078, socket_xyz[1], self.SOCKET_HOLD_Z], dtype=np.float64
+        )
+
+        peg_init_xyz = peg_xyz + np.array(
+            [0.078, 0.0, self.ACTION_OBJECT_Z_OFFSET + 0.01]
+        )
+        peg_lift_center = np.array(
+            [peg_xyz[0] + 0.078, socket_hold_action[1], self.SOCKET_HOLD_Z - 0.01],
+            dtype=np.float64,
+        )
+        # The front camera looks along +Y, so visual right-to-left insertion is
+        # world +X -> -X.  With the socket XML in identity orientation, its
+        # tunnel axis is local/world X, so the peg approaches from +X and stops
+        # when its leading face reaches the socket's internal pin.
+        peg_insert_end_center = np.array(
+            [
+                socket_hold_action[0] + 0.078 * 2 + 0.04 + 0.06 - 0.01,
+                socket_hold_action[1],
+                self.SOCKET_HOLD_Z - 0.01,
+            ],
+            dtype=np.float64,
+        )
+
+        self.left_trajectory = [
+            {
+                "t": 1,
+                "xyz": init_mocap_pose_left[:3],
+                "quat": init_mocap_pose_left[3:],
+                "gripper": 100,
+            },
+            {
+                "t": 130,
+                "xyz": socket_xyz
+                + np.array([-0.078, 0.0, self.ACTION_OBJECT_Z_OFFSET]),
+                "quat": left_pick_quat,
+                "gripper": 100,
+            },
+            {
+                "t": 180,
+                "xyz": socket_xyz
+                + np.array([-0.078, 0.0, self.ACTION_OBJECT_Z_OFFSET]),
+                "quat": left_pick_quat,
+                "gripper": self.LEFT_SOCKET_GRIPPER_CLOSED,
+            },
+            {
+                "t": 450,
+                "xyz": socket_hold_action,
+                "quat": left_pick_quat,
+                "gripper": self.LEFT_SOCKET_GRIPPER_CLOSED,
+            },
+            {
+                "t": 750,
+                "xyz": socket_hold_action,
+                "quat": left_pick_quat,
+                "gripper": self.LEFT_SOCKET_GRIPPER_CLOSED,
+            },
+            {
+                "t": self.EPISODE_END_T,
+                "xyz": socket_hold_action,
+                "quat": left_pick_quat,
+                "gripper": self.LEFT_SOCKET_GRIPPER_CLOSED,
+            },
+        ]
+
+        self.right_trajectory = [
+            {
+                "t": 1,
+                "xyz": init_mocap_pose_right[:3],
+                "quat": init_mocap_pose_right[3:],
+                "gripper": 100,
+            },
+            {
+                "t": 80,
+                "xyz": peg_init_xyz,
+                "quat": right_pick_quat,
+                "gripper": 100,
+            },
+            {
+                "t": 150,
+                "xyz": peg_init_xyz,
+                "quat": right_pick_quat,
+                "gripper": 100,
+            },
+            {
+                "t": 180,
+                "xyz": peg_init_xyz,
+                "quat": right_pick_quat,
+                "gripper": self.RIGHT_PEG_GRIPPER_CLOSED,
+            },
+            {
+                "t": 450,
+                "xyz": peg_init_xyz,
+                "quat": right_pick_quat,
+                "gripper": self.RIGHT_PEG_GRIPPER_CLOSED,
+            },
+            {
+                "t": 550,
+                "xyz": peg_lift_center,
+                "quat": right_pick_quat,
+                "gripper": self.RIGHT_PEG_GRIPPER_CLOSED,
+            },
+            {
+                "t": self.INSERT_START_T,
+                "xyz": peg_lift_center,
+                "quat": right_pick_quat,
+                "gripper": self.RIGHT_PEG_GRIPPER_CLOSED,
+            },
+            {
+                "t": self.INSERT_END_T,
+                "xyz": peg_insert_end_center,
+                "quat": right_pick_quat,
+                "gripper": self.RIGHT_PEG_GRIPPER_CLOSED,
+            },
+            {
+                "t": 750,
+                "xyz": peg_insert_end_center,
+                "quat": right_pick_quat,
+                "gripper": self.RIGHT_PEG_GRIPPER_CLOSED,
+            },
+            {
+                "t": self.EPISODE_END_T,
+                "xyz": peg_insert_end_center,
+                "quat": right_pick_quat,
+                "gripper": self.RIGHT_PEG_GRIPPER_CLOSED,
+            },
+        ]

roboimi/demos/diana_record_sim_episodes.py

@@ -2,9 +2,11 @@ import time
 import os
 import numpy as np
 from roboimi.envs.double_pos_ctrl_env import make_sim_env
-from diana_policy import TestPickAndTransferPolicy
+from roboimi.demos.diana_air_insert_policy import TestAirInsertPolicy
+from roboimi.demos.diana_policy import TestPickAndTransferPolicy
 import cv2
-from roboimi.utils.act_ex_utils import sample_transfer_pose
+from roboimi.utils.act_ex_utils import sample_air_insert_socket_peg_state, sample_transfer_pose
+from roboimi.utils.constants import SIM_TASK_CONFIGS
 from roboimi.utils.streaming_episode_writer import StreamingEpisodeWriter
 
 import pathlib
@@ -12,16 +14,34 @@ HOME_PATH = str(pathlib.Path(__file__).parent.resolve())
 DATASET_DIR = HOME_PATH + '/dataset'
 
 
-def main():
+def sample_task_state(task_name):
-    task_name = 'sim_transfer'
+    if task_name == 'sim_transfer':
-    dataset_dir = DATASET_DIR + '/sim_transfer'   #SIM_TASK_CONFIGS[task_name]['dataset_dir']
+        return sample_transfer_pose()
-    num_episodes = 100       #SIM_TASK_CONFIGS[task_name]['num_episodes']
+    if task_name == 'sim_air_insert_socket_peg':
+        return sample_air_insert_socket_peg_state()
+    raise NotImplementedError(f'Unsupported scripted rollout task: {task_name}')
+
+
+def make_policy(task_name, inject_noise=False, grasp_strategy=None):
+    if task_name == 'sim_transfer':
+        return TestPickAndTransferPolicy(inject_noise)
+    if task_name == 'sim_air_insert_socket_peg':
+        if grasp_strategy is None:
+            return TestAirInsertPolicy(inject_noise)
+        return TestAirInsertPolicy(inject_noise, grasp_strategy=grasp_strategy)
+    raise NotImplementedError(f'Unsupported scripted rollout task: {task_name}')
+
+
+def main(task_name='sim_transfer'):
+    task_cfg = SIM_TASK_CONFIGS[task_name]
+    dataset_dir = task_cfg['dataset_dir']
+    num_episodes = 100
     inject_noise = False
 
-    episode_len  = 700      #SIM_TASK_CONFIGS[task_name]['episode_len']
+    episode_len = task_cfg['episode_len']
-    camera_names = ['angle','r_vis', 'top', 'front']   #SIM_TASK_CONFIGS[task_name]['camera_names']
+    camera_names = task_cfg['camera_names']
     image_size = (256, 256)
-    if task_name == 'sim_transfer':
+    if task_name in {'sim_transfer', 'sim_air_insert_socket_peg'}:
         print(task_name)
     else:
         raise NotImplementedError
@@ -29,7 +49,7 @@ def main():
     success = []
     
     env = make_sim_env(task_name)
-    policy = TestPickAndTransferPolicy(inject_noise)
+    policy = make_policy(task_name, inject_noise=inject_noise)
 
     # 等待osmesa完全启动后再开始收集数据
     print("等待osmesa线程启动...")
@@ -41,8 +61,8 @@ def main():
         max_reward = float('-inf')
         print(f'\n{episode_idx=}')
         print('Rollout out EE space scripted policy')
-        box_pos = sample_transfer_pose()
+        task_state = sample_task_state(task_name)
-        env.reset(box_pos)
+        env.reset(task_state)
         episode_writer = StreamingEpisodeWriter(
             dataset_path=os.path.join(dataset_dir, f'episode_{episode_idx}.hdf5'),
             max_timesteps=episode_len,
@@ -50,7 +70,7 @@ def main():
             image_size=image_size,
         )
         for step in range(episode_len):
-            raw_action = policy.predict(box_pos,step)
+            raw_action = policy.predict(task_state, step)
             env.step(raw_action)
             env.render()
             sum_reward += env.rew

roboimi/demos/vla_scripts/eval_vla.py

@@ -25,7 +25,11 @@ from omegaconf import DictConfig, OmegaConf
 from hydra.utils import instantiate
 from einops import rearrange
 
-from roboimi.utils.act_ex_utils import sample_transfer_pose
+from roboimi.envs.double_pos_ctrl_env import make_sim_env
+from roboimi.utils.act_ex_utils import (
+    sample_air_insert_socket_peg_state,
+    sample_transfer_pose,
+)
 from roboimi.vla.eval_utils import execute_policy_action
 
 sys.path.append(os.getcwd())
@@ -630,6 +634,14 @@ def _close_env(env):
         viewer.close()
 
 
+def _sample_task_reset_state(task_name: str):
+    if task_name == 'sim_air_insert_socket_peg':
+        return sample_air_insert_socket_peg_state()
+    if 'sim_transfer' in task_name:
+        return sample_transfer_pose()
+    raise NotImplementedError(f'Unsupported eval task reset sampling: {task_name}')
+
+
 def _run_eval(cfg: DictConfig):
     """
     使用 agent 内置队列管理的简化版 VLA 评估
@@ -697,8 +709,8 @@ def _run_eval(cfg: DictConfig):
             print(f"回合 {episode_idx + 1}/{eval_cfg.num_episodes}")
             print(f"{'='*60}\n")
 
-            box_pos = sample_transfer_pose()
+            task_state = _sample_task_reset_state(str(eval_cfg.task_name))
-            env.reset(box_pos)
+            env.reset(task_state)
 
             # 为新回合重置 agent 队列
             agent.reset()

roboimi/envs/double_air_insert_env.py

@@ -0,0 +1,157 @@
+import copy as cp
+import time
+
+import numpy as np
+
+from roboimi.envs.double_base import DualDianaMed
+from roboimi.envs.double_pos_ctrl_env import DualDianaMed_Pos_Ctrl
+
+
+SOCKET_JOINT_NAME = "blue_socket_joint"
+PEG_JOINT_NAME = "red_peg_joint"
+REQUIRED_TASK_STATE_KEYS = ("socket_pos", "socket_quat", "peg_pos", "peg_quat")
+SOCKET_GEOM_NAMES = ("socket-1", "socket-2", "socket-3", "socket-4")
+SOCKET_SUCCESS_GEOM_NAMES = ("pin",)
+SOCKET_BODY_GEOM_NAMES = SOCKET_GEOM_NAMES + SOCKET_SUCCESS_GEOM_NAMES
+PEG_GEOM_NAMES = ("red_peg",)
+LEFT_GRIPPER_GEOM_NAMES = (
+    "l_finger_left",
+    "r_finger_left",
+    "l_fingertip_g0_left",
+    "r_fingertip_g0_left",
+    "l_fingerpad_g0_left",
+    "r_fingerpad_g0_left",
+    "l_fingertip_g0_vis_left",
+    "r_fingertip_g0_vis_left",
+)
+RIGHT_GRIPPER_GEOM_NAMES = (
+    "l_finger_right",
+    "r_finger_right",
+    "l_fingertip_g0_right",
+    "r_fingertip_g0_right",
+    "l_fingerpad_g0_right",
+    "r_fingerpad_g0_right",
+    "l_fingertip_g0_vis_right",
+    "r_fingertip_g0_vis_right",
+)
+TABLE_GEOM_NAME = "table"
+
+
+def _set_free_joint_pose(joint, position, quat):
+    joint.qpos[:3] = np.asarray(position, dtype=np.float64)
+    joint.qpos[3:7] = np.asarray(quat, dtype=np.float64)
+
+
+def set_socket_peg_task_state(mj_data, task_state):
+    if not isinstance(task_state, dict) or tuple(task_state.keys()) != REQUIRED_TASK_STATE_KEYS:
+        raise ValueError(
+            "task_state must be an ordered dict-like mapping with keys "
+            "socket_pos, socket_quat, peg_pos, peg_quat"
+        )
+
+    _set_free_joint_pose(
+        mj_data.joint(SOCKET_JOINT_NAME),
+        task_state["socket_pos"],
+        task_state["socket_quat"],
+    )
+    _set_free_joint_pose(
+        mj_data.joint(PEG_JOINT_NAME),
+        task_state["peg_pos"],
+        task_state["peg_quat"],
+    )
+
+
+def get_socket_peg_env_state(mj_data):
+    socket_qpos = cp.deepcopy(np.asarray(mj_data.joint(SOCKET_JOINT_NAME).qpos[:7], dtype=np.float64))
+    peg_qpos = cp.deepcopy(np.asarray(mj_data.joint(PEG_JOINT_NAME).qpos[:7], dtype=np.float64))
+    return np.concatenate([socket_qpos, peg_qpos], dtype=np.float64)
+
+
+def _normalize_contact_pairs(contact_pairs):
+    return {frozenset(pair) for pair in contact_pairs}
+
+
+def _has_any_object_contact(contact_set, object_geom_names, other_geom_names):
+    return any(
+        frozenset((object_geom_name, other_geom_name)) in contact_set
+        for object_geom_name in object_geom_names
+        for other_geom_name in other_geom_names
+    )
+
+
+def _object_is_airborne(contact_set, object_geom_names):
+    return not _has_any_object_contact(contact_set, object_geom_names, (TABLE_GEOM_NAME,))
+
+
+def peg_inserted_into_socket(contact_pairs):
+    contact_set = _normalize_contact_pairs(contact_pairs)
+    return frozenset((PEG_GEOM_NAMES[0], SOCKET_SUCCESS_GEOM_NAMES[0])) in contact_set
+
+
+def compute_air_insert_reward(contact_pairs, env_state=None):
+    del env_state  # kept for API compatibility with rollout/eval code paths
+    contact_set = _normalize_contact_pairs(contact_pairs)
+    reward = 0
+
+    if _has_any_object_contact(contact_set, SOCKET_GEOM_NAMES, LEFT_GRIPPER_GEOM_NAMES):
+        reward += 1
+    if _has_any_object_contact(contact_set, PEG_GEOM_NAMES, RIGHT_GRIPPER_GEOM_NAMES):
+        reward += 1
+
+    socket_airborne = _object_is_airborne(contact_set, SOCKET_BODY_GEOM_NAMES)
+    peg_airborne = _object_is_airborne(contact_set, PEG_GEOM_NAMES)
+    if socket_airborne:
+        reward += 1
+    if peg_airborne:
+        reward += 1
+
+    if socket_airborne and peg_airborne and peg_inserted_into_socket(contact_pairs):
+        reward += 1
+
+    return reward
+
+
+class DualDianaMed_Air_Insert(DualDianaMed_Pos_Ctrl):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.max_reward = 5
+
+    def reset(self, task_state):
+        set_socket_peg_task_state(self.mj_data, task_state)
+        DualDianaMed.reset(self)
+        self.top = None
+        self.r_vis = None
+        self.l_vis = None
+        self.front = None
+        if not self.is_render:
+            self._update_camera_images_sync()
+            return
+        self.cam_flage = True
+        while self.cam_flage:
+            if (
+                type(self.top) == type(None)
+                or type(self.r_vis) == type(None)
+                or type(self.l_vis) == type(None)
+                or type(self.front) == type(None)
+            ):
+                time.sleep(0.001)
+            else:
+                self.cam_flage = False
+
+    def step(self, action=np.zeros(16)):
+        super().step(action)
+        self.rew = self._get_reward()
+        self.obs = self._get_obs()
+
+    def get_env_state(self):
+        return get_socket_peg_env_state(self.mj_data)
+
+    def _get_reward(self):
+        contact_pairs = []
+        for collision_num in range(self.mj_data.ncon):
+            geom1 = self.mj_data.contact[collision_num].geom1
+            geom2 = self.mj_data.contact[collision_num].geom2
+            contact_pairs.append(
+                (self.getID2Name("geom", geom1), self.getID2Name("geom", geom2))
+            )
+        return compute_air_insert_reward(contact_pairs, self.get_env_state())

roboimi/envs/double_base.py

@@ -52,7 +52,6 @@ class DualDianaMed(MujocoEnv):
         self.r_vis = None
         self.l_vis = None
         self.top = None
-        self.angle = None
         self.front = None
         self.obs = None
         
@@ -168,10 +167,9 @@ class DualDianaMed(MujocoEnv):
         obs['qpos'] = self.get_obs_qpos
         obs['action'] = self.compute_qpos
         obs['images'] = dict()
-        obs['images']['top']   = self.top
-        obs['images']['angle'] = self.angle
         obs['images']['r_vis'] = self.r_vis
         obs['images']['l_vis'] = self.l_vis
+        obs['images']['top'] = self.top
         obs['images']['front'] = self.front
         return obs
 
@@ -180,10 +178,9 @@ class DualDianaMed(MujocoEnv):
             self._update_camera_images_sync()
         obs = collections.OrderedDict()
         obs['images'] = dict()
-        obs['images']['top']   = self.top
-        obs['images']['angle'] = self.angle
         obs['images']['r_vis'] = self.r_vis
         obs['images']['l_vis'] = self.l_vis
+        obs['images']['top'] = self.top
         obs['images']['front'] = self.front
         return obs
     
@@ -202,14 +199,12 @@ class DualDianaMed(MujocoEnv):
 
     @property
     def cam_view(self):
-        if self.cam == 'top':
+        if self.cam == 'r_vis':
-            return self.top
-        elif self.cam == 'angle':
-            return self.angle
-        elif self.cam == 'r_vis':
             return self.r_vis
         elif self.cam == 'l_vis':
             return self.l_vis
+        elif self.cam == 'top':
+            return self.top
         elif self.cam == 'front':
             return self.front
         else:
@@ -230,8 +225,6 @@ class DualDianaMed(MujocoEnv):
         self.l_vis = img_renderer.render()[:, :, ::-1]
         img_renderer.update_scene(self.mj_data, camera="top")
         self.top = img_renderer.render()[:, :, ::-1]
-        img_renderer.update_scene(self.mj_data, camera="angle")
-        self.angle = img_renderer.render()[:, :, ::-1]
         img_renderer.update_scene(self.mj_data, camera="front")
         self.front = img_renderer.render()[:, :, ::-1]
 

roboimi/envs/double_pos_ctrl_env.py

@@ -73,8 +73,8 @@ class DualDianaMed_Pos_Ctrl(DualDianaMed):
         self.mj_data.joint('red_box_joint').qpos[6] = 0.0
         super().reset()
         self.top = None
-        self.angle = None
         self.r_vis = None
+        self.l_vis = None
         self.front = None
         if not self.is_render:
             self._update_camera_images_sync()
@@ -83,8 +83,8 @@ class DualDianaMed_Pos_Ctrl(DualDianaMed):
         t=0
         while self.cam_flage:
             if(type(self.top)==type(None)
-               or type(self.angle)==type(None) 
                or type(self.r_vis)==type(None)
+               or type(self.l_vis)==type(None)
                or type(self.front)==type(None)):
                 time.sleep(0.001)
                 t+=1
@@ -137,6 +137,18 @@ class DualDianaMed_Pos_Ctrl(DualDianaMed):
 
 
 def make_sim_env(task_name, headless=False):
+    if task_name == 'sim_air_insert_socket_peg':
+        from roboimi.assets.robots.diana_med import BiDianaMedSocketPeg
+        from roboimi.envs.double_air_insert_env import DualDianaMed_Air_Insert
+
+        env = DualDianaMed_Air_Insert(
+            robot=BiDianaMedSocketPeg(),
+            is_render=not headless,
+            control_freq=30,
+            is_interpolate=True,
+            cam_view='front'
+        )
+        return env
     if 'sim_transfer' in task_name:
         from roboimi.assets.robots.diana_med import BiDianaMed
         env = DualDianaMed_Pos_Ctrl(
@@ -144,7 +156,7 @@ def make_sim_env(task_name, headless=False):
             is_render=not headless,
             control_freq=30,
             is_interpolate=True,
-            cam_view='angle'
+            cam_view='top'
         )
         return env
     else:
@@ -170,4 +182,3 @@ if __name__ == "__main__":
         env.step(action)
         if env.is_render:
             env.render()
-        

roboimi/utils/act_ex_utils.py

@@ -1,5 +1,6 @@
 import numpy as np
 
+
 def sample_insertion_pose():
         # Peg
         x_range = [0.1, 0.2]
@@ -36,3 +37,22 @@ def sample_transfer_pose():
 
 
         return box_position
+
+
+def sample_air_insert_socket_peg_state():
+        socket_position = np.random.uniform(
+            low=np.array([-0.20, 0.80, 0.472], dtype=np.float32),
+            high=np.array([-0.10, 1.00, 0.472], dtype=np.float32),
+        )
+        peg_position = np.random.uniform(
+            low=np.array([0.10, 0.80, 0.46], dtype=np.float32),
+            high=np.array([0.20, 1.00, 0.46], dtype=np.float32),
+        )
+        socket_quat = np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32)
+        peg_quat = np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32)
+        return {
+            "socket_pos": socket_position.astype(np.float32, copy=False),
+            "socket_quat": socket_quat,
+            "peg_pos": peg_position.astype(np.float32, copy=False),
+            "peg_quat": peg_quat,
+        }

roboimi/utils/constants.py

@@ -20,7 +20,14 @@ SIM_TASK_CONFIGS = {
         'dataset_dir': DATASET_DIR + '/sim_transfer',
         'num_episodes': 20,
         'episode_len': 700,
-        'camera_names': ['top','r_vis','front'],
+        'camera_names': ['r_vis', 'top', 'front'],
+        'xml_dir': HOME_PATH + '/assets'
+    },
+    'sim_air_insert_socket_peg': {
+        'dataset_dir': DATASET_DIR + '/sim_air_insert_socket_peg',
+        'num_episodes': 20,
+        'episode_len': 750,
+        'camera_names': ['l_vis', 'r_vis', 'front'],
         'xml_dir': HOME_PATH + '/assets'
     },
 
@@ -52,13 +59,3 @@ PUPPET_GRIPPER_JOINT_NORMALIZE_FN = lambda x: (x - PUPPET_GRIPPER_JOINT_CLOSE) /
 MASTER_GRIPPER_JOINT_UNNORMALIZE_FN = lambda x: x * (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE) + MASTER_GRIPPER_JOINT_CLOSE
 PUPPET_GRIPPER_JOINT_UNNORMALIZE_FN = lambda x: x * (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE) + PUPPET_GRIPPER_JOINT_CLOSE
 MASTER2PUPPET_JOINT_FN = lambda x: PUPPET_GRIPPER_JOINT_UNNORMALIZE_FN(MASTER_GRIPPER_JOINT_NORMALIZE_FN(x))
-
-MASTER_GRIPPER_VELOCITY_NORMALIZE_FN = lambda x: x / (MASTER_GRIPPER_POSITION_OPEN - MASTER_GRIPPER_POSITION_CLOSE)
-PUPPET_GRIPPER_VELOCITY_NORMALIZE_FN = lambda x: x / (PUPPET_GRIPPER_POSITION_OPEN - PUPPET_GRIPPER_POSITION_CLOSE)
-
-MASTER_POS2JOINT = lambda x: MASTER_GRIPPER_POSITION_NORMALIZE_FN(x) * (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE) + MASTER_GRIPPER_JOINT_CLOSE
-MASTER_JOINT2POS = lambda x: MASTER_GRIPPER_POSITION_UNNORMALIZE_FN((x - MASTER_GRIPPER_JOINT_CLOSE) / (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE))
-PUPPET_POS2JOINT = lambda x: PUPPET_GRIPPER_POSITION_NORMALIZE_FN(x) * (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE) + PUPPET_GRIPPER_JOINT_CLOSE
-PUPPET_JOINT2POS = lambda x: PUPPET_GRIPPER_POSITION_UNNORMALIZE_FN((x - PUPPET_GRIPPER_JOINT_CLOSE) / (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE))
-
-MASTER_GRIPPER_JOINT_MID = (MASTER_GRIPPER_JOINT_OPEN + MASTER_GRIPPER_JOINT_CLOSE)/2

tests/test_air_insert_env.py

@@ -0,0 +1,510 @@
+import importlib
+import inspect
+import pathlib
+import unittest
+from unittest import mock
+import xml.etree.ElementTree as ET
+
+import numpy as np
+
+from roboimi.envs.double_pos_ctrl_env import make_sim_env
+from roboimi.utils import act_ex_utils
+from roboimi.utils.constants import SIM_TASK_CONFIGS
+
+
+TASK_NAME = "sim_air_insert_socket_peg"
+
+
+class AirInsertTaskRegistrationTest(unittest.TestCase):
+    def test_sim_task_configs_registers_air_insert_socket_peg(self):
+        self.assertIn(TASK_NAME, SIM_TASK_CONFIGS)
+        self.assertNotIn("sim_air_insert_ring_bar", SIM_TASK_CONFIGS)
+        self.assertEqual(SIM_TASK_CONFIGS[TASK_NAME]["episode_len"], 750)
+        self.assertEqual(SIM_TASK_CONFIGS[TASK_NAME]["camera_names"], ["l_vis", "r_vis", "front"])
+        self.assertTrue(SIM_TASK_CONFIGS[TASK_NAME]["dataset_dir"].endswith("/sim_air_insert_socket_peg"))
+
+    def test_sample_air_insert_socket_peg_state_returns_explicit_named_mapping(self):
+        sampler = getattr(act_ex_utils, "sample_air_insert_socket_peg_state", None)
+        self.assertIsNotNone(
+            sampler,
+            "Expected roboimi.utils.act_ex_utils.sample_air_insert_socket_peg_state()",
+        )
+        self.assertFalse(
+            hasattr(act_ex_utils, "sample_air_insert_ring_bar_state"),
+            "air insert sampler should use socket/peg naming after the task rename",
+        )
+
+        task_state = sampler()
+
+        self.assertEqual(
+            list(task_state.keys()),
+            ["socket_pos", "socket_quat", "peg_pos", "peg_quat"],
+        )
+        self.assertEqual(task_state["socket_pos"].shape, (3,))
+        self.assertEqual(task_state["socket_quat"].shape, (4,))
+        self.assertEqual(task_state["peg_pos"].shape, (3,))
+        self.assertEqual(task_state["peg_quat"].shape, (4,))
+
+    def test_sample_air_insert_socket_peg_state_uses_fixed_quats_and_left_right_planar_ranges(self):
+        sampler = getattr(act_ex_utils, "sample_air_insert_socket_peg_state", None)
+        self.assertIsNotNone(sampler)
+
+        task_state = sampler()
+
+        np.testing.assert_array_equal(task_state["socket_quat"], np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32))
+        np.testing.assert_array_equal(task_state["peg_quat"], np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32))
+        self.assertGreaterEqual(task_state["socket_pos"][0], -0.20)
+        self.assertLessEqual(task_state["socket_pos"][0], -0.05)
+        self.assertGreaterEqual(task_state["socket_pos"][1], 0.70)
+        self.assertLessEqual(task_state["socket_pos"][1], 1.00)
+        self.assertAlmostEqual(float(task_state["socket_pos"][2]), 0.472)
+        self.assertGreaterEqual(task_state["peg_pos"][0], 0.05)
+        self.assertLessEqual(task_state["peg_pos"][0], 0.20)
+        self.assertGreaterEqual(task_state["peg_pos"][1], 0.70)
+        self.assertLessEqual(task_state["peg_pos"][1], 1.00)
+        self.assertAlmostEqual(float(task_state["peg_pos"][2]), 0.46)
+
+    def test_make_sim_env_dispatches_air_insert_socket_peg_headless(self):
+        air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env")
+        air_insert_cls = getattr(air_insert_env, "DualDianaMed_Air_Insert", None)
+        self.assertIsNotNone(air_insert_cls)
+
+        diana_med = importlib.import_module("roboimi.assets.robots.diana_med")
+        socket_peg_robot_cls = getattr(diana_med, "BiDianaMedSocketPeg", None)
+        self.assertIsNotNone(
+            socket_peg_robot_cls,
+            "Expected roboimi.assets.robots.diana_med.BiDianaMedSocketPeg",
+        )
+
+        fake_env = object()
+        with mock.patch.object(
+            diana_med,
+            "BiDianaMedSocketPeg",
+            return_value="robot",
+        ), mock.patch.object(
+            air_insert_env,
+            "DualDianaMed_Air_Insert",
+            return_value=fake_env,
+        ) as env_cls:
+            env = make_sim_env(TASK_NAME, headless=True)
+
+        self.assertIs(env, fake_env)
+        env_cls.assert_called_once_with(
+            robot="robot",
+            is_render=False,
+            control_freq=30,
+            is_interpolate=True,
+            cam_view="front",
+        )
+
+    def test_diana_table_scene_exposes_only_top_and_front_scene_cameras(self):
+        xml_path = (
+            pathlib.Path(__file__).resolve().parents[1]
+            / "roboimi/assets/models/manipulators/DianaMed/table_square.xml"
+        )
+        root = ET.parse(xml_path).getroot()
+        cameras = {camera.attrib["name"]: camera.attrib for camera in root.findall(".//camera")}
+
+        self.assertNotIn("angle", cameras, "DianaMed scene should stop exposing the old angle camera")
+        self.assertNotIn("left_side", cameras, "DianaMed scene should no longer expose left_side")
+        self.assertIn("top", cameras)
+        self.assertIn("front", cameras)
+        self.assertEqual(cameras["top"].get("mode"), "targetbody")
+        self.assertEqual(cameras["top"].get("target"), "table")
+
+
+class AirInsertResetAndStateHelpersTest(unittest.TestCase):
+    def test_set_socket_peg_task_state_writes_free_joint_qpos(self):
+        air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env")
+        setter = getattr(air_insert_env, "set_socket_peg_task_state", None)
+        self.assertIsNotNone(
+            setter,
+            "Expected roboimi.envs.double_air_insert_env.set_socket_peg_task_state",
+        )
+
+        socket_qpos = np.zeros(7, dtype=np.float64)
+        peg_qpos = np.zeros(7, dtype=np.float64)
+
+        class _FakeJoint:
+            def __init__(self, qpos):
+                self.qpos = qpos
+
+        class _FakeData:
+            def joint(self, name):
+                if name == "blue_socket_joint":
+                    return _FakeJoint(socket_qpos)
+                if name == "red_peg_joint":
+                    return _FakeJoint(peg_qpos)
+                raise AssertionError(f"Unexpected joint name: {name}")
+
+        task_state = {
+            "socket_pos": np.array([-0.12, 0.90, 0.472], dtype=np.float64),
+            "socket_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float64),
+            "peg_pos": np.array([0.12, 0.91, 0.46], dtype=np.float64),
+            "peg_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float64),
+        }
+
+        setter(_FakeData(), task_state)
+
+        np.testing.assert_array_equal(
+            socket_qpos,
+            np.array([-0.12, 0.90, 0.472, 1.0, 0.0, 0.0, 0.0], dtype=np.float64),
+        )
+        np.testing.assert_array_equal(
+            peg_qpos,
+            np.array([0.12, 0.91, 0.46, 1.0, 0.0, 0.0, 0.0], dtype=np.float64),
+        )
+
+    def test_get_socket_peg_env_state_returns_stable_14d_vector(self):
+        air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env")
+        getter = getattr(air_insert_env, "get_socket_peg_env_state", None)
+        self.assertIsNotNone(
+            getter,
+            "Expected roboimi.envs.double_air_insert_env.get_socket_peg_env_state",
+        )
+
+        socket_qpos = np.array([-0.12, 0.90, 0.472, 1.0, 0.0, 0.0, 0.0], dtype=np.float64)
+        peg_qpos = np.array([0.12, 0.91, 0.46, 1.0, 0.0, 0.0, 0.0], dtype=np.float64)
+
+        class _FakeJoint:
+            def __init__(self, qpos):
+                self.qpos = qpos
+
+        class _FakeData:
+            def joint(self, name):
+                if name == "blue_socket_joint":
+                    return _FakeJoint(socket_qpos)
+                if name == "red_peg_joint":
+                    return _FakeJoint(peg_qpos)
+                raise AssertionError(f"Unexpected joint name: {name}")
+
+        env_state = getter(_FakeData())
+
+        self.assertEqual(env_state.shape, (14,))
+        np.testing.assert_array_equal(
+            env_state,
+            np.array(
+                [-0.12, 0.90, 0.472, 1.0, 0.0, 0.0, 0.0, 0.12, 0.91, 0.46, 1.0, 0.0, 0.0, 0.0],
+                dtype=np.float64,
+            ),
+        )
+
+    def test_air_insert_env_does_not_script_attach_or_assist_objects_after_reset(self):
+        air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env")
+        env_cls = getattr(air_insert_env, "DualDianaMed_Air_Insert", None)
+        self.assertIsNotNone(env_cls)
+
+        source = inspect.getsource(env_cls)
+
+        self.assertNotIn("_update_scripted_grasped_objects", source)
+        self.assertNotIn("_scripted_", source)
+        self.assertNotIn("_stabilize_ring_grasp", source)
+        self.assertNotIn("_ring_grasp_locked", source)
+        get_reward_source = inspect.getsource(env_cls._get_reward)
+        self.assertNotIn("ring_block", get_reward_source)
+        self.assertNotIn("bar_block", get_reward_source)
+
+    def test_socket_peg_xml_defines_active_socket_and_peg_objects(self):
+        asset_dir = pathlib.Path(__file__).resolve().parents[1] / "roboimi/assets/models/manipulators/DianaMed"
+        xml_path = asset_dir / "socket_peg_objects.xml"
+        self.assertTrue(xml_path.exists(), "socket/peg objects should live in socket_peg_objects.xml")
+        self.assertFalse((asset_dir / "ring_bar_objects.xml").exists(), "old ring_bar_objects.xml should be renamed")
+
+        root = ET.parse(xml_path).getroot()
+        body_names = {body.attrib.get("name") for body in root.findall(".//body")}
+        geom_names = {geom.attrib.get("name") for geom in root.findall(".//geom")}
+        joint_names = {joint.attrib.get("name") for joint in root.findall(".//joint")}
+
+        self.assertIn("socket", body_names)
+        self.assertIn("peg", body_names)
+        self.assertNotIn("ring_block", body_names)
+        self.assertNotIn("bar_block", body_names)
+        self.assertIn("blue_socket_joint", joint_names)
+        self.assertIn("red_peg_joint", joint_names)
+        for geom_name in ("socket-1", "socket-2", "socket-3", "socket-4", "pin", "red_peg"):
+            self.assertIn(geom_name, geom_names)
+
+    def test_socket_peg_wrapper_includes_socket_peg_objects(self):
+        xml_path = (
+            pathlib.Path(__file__).resolve().parents[1]
+            / "roboimi/assets/models/manipulators/DianaMed/bi_diana_socket_peg_ee.xml"
+        )
+        self.assertTrue(xml_path.exists(), "socket/peg wrapper XML should use the new task name")
+        root = ET.parse(xml_path).getroot()
+        includes = [include.attrib.get("file") for include in root.findall(".//include")]
+        self.assertIn("./socket_peg_objects.xml", includes)
+        self.assertNotIn("./ring_bar_objects.xml", includes)
+
+
+class AirInsertRewardAndSuccessTest(unittest.TestCase):
+    @staticmethod
+    def _make_env_state(
+        socket_pos=(0.0, 0.0, 0.472),
+        socket_quat=(1.0, 0.0, 0.0, 0.0),
+        peg_pos=(0.0, 0.0, 0.46),
+        peg_quat=(1.0, 0.0, 0.0, 0.0),
+    ):
+        return np.array([*socket_pos, *socket_quat, *peg_pos, *peg_quat], dtype=np.float64)
+
+    def test_compute_air_insert_reward_counts_left_contact_stage(self):
+        air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env")
+        reward_fn = getattr(air_insert_env, "compute_air_insert_reward", None)
+        self.assertIsNotNone(reward_fn)
+
+        reward = reward_fn(
+            contact_pairs=[
+                ("socket-1", "l_finger_left"),
+                ("socket-1", "table"),
+                ("red_peg", "table"),
+            ],
+            env_state=self._make_env_state(),
+        )
+
+        self.assertEqual(reward, 1)
+
+    def test_compute_air_insert_reward_counts_right_contact_stage(self):
+        air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env")
+        reward_fn = getattr(air_insert_env, "compute_air_insert_reward", None)
+
+        reward = reward_fn(
+            contact_pairs=[
+                ("socket-1", "l_finger_left"),
+                ("red_peg", "l_finger_right"),
+                ("socket-1", "table"),
+                ("red_peg", "table"),
+            ],
+            env_state=self._make_env_state(),
+        )
+
+        self.assertEqual(reward, 2)
+
+    def test_compute_air_insert_reward_counts_lift_stages(self):
+        air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env")
+        reward_fn = getattr(air_insert_env, "compute_air_insert_reward", None)
+
+        reward = reward_fn(
+            contact_pairs=[
+                ("socket-1", "l_finger_left"),
+                ("red_peg", "l_finger_right"),
+            ],
+            env_state=self._make_env_state(),
+        )
+
+        self.assertEqual(reward, 4)
+
+    def test_compute_air_insert_reward_counts_visual_fingertip_contacts_as_gripper_contacts(self):
+        air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env")
+        reward_fn = getattr(air_insert_env, "compute_air_insert_reward", None)
+
+        reward = reward_fn(
+            contact_pairs=[
+                ("socket-3", "r_fingertip_g0_vis_left"),
+                ("red_peg", "l_fingertip_g0_vis_right"),
+            ],
+            env_state=self._make_env_state(),
+        )
+
+        self.assertEqual(
+            reward,
+            4,
+            "visual fingertip geoms are collidable in the Diana XML and should count as gripper-object contacts",
+        )
+
+    def test_peg_inserted_into_socket_uses_pin_contact(self):
+        air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env")
+        success_fn = getattr(air_insert_env, "peg_inserted_into_socket", None)
+        self.assertIsNotNone(
+            success_fn,
+            "Expected roboimi.envs.double_air_insert_env.peg_inserted_into_socket",
+        )
+
+        self.assertTrue(success_fn([("red_peg", "pin")]))
+        self.assertTrue(success_fn([("pin", "red_peg")]))
+        self.assertFalse(success_fn([("red_peg", "socket-1")]))
+
+    def test_compute_air_insert_reward_requires_airborne_success_for_final_point(self):
+        air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env")
+        reward_fn = getattr(air_insert_env, "compute_air_insert_reward", None)
+
+        reward = reward_fn(
+            contact_pairs=[
+                ("socket-1", "l_finger_left"),
+                ("red_peg", "l_finger_right"),
+                ("socket-1", "table"),
+                ("red_peg", "pin"),
+            ],
+            env_state=self._make_env_state(),
+        )
+
+        self.assertEqual(reward, 3)
+
+    def test_compute_air_insert_reward_returns_full_score_on_true_airborne_insert(self):
+        air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env")
+        reward_fn = getattr(air_insert_env, "compute_air_insert_reward", None)
+
+        reward = reward_fn(
+            contact_pairs=[
+                ("socket-1", "l_finger_left"),
+                ("red_peg", "l_finger_right"),
+                ("red_peg", "pin"),
+            ],
+            env_state=self._make_env_state(),
+        )
+
+        self.assertEqual(reward, 5)
+
+
+class AirInsertPolicyAndSmokeTest(unittest.TestCase):
+    @staticmethod
+    def _canonical_task_state():
+        return {
+            "socket_pos": np.array([-0.12, 0.90, 0.472], dtype=np.float32),
+            "socket_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32),
+            "peg_pos": np.array([0.12, 0.90, 0.46], dtype=np.float32),
+            "peg_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32),
+        }
+
+    def test_air_insert_policy_emits_valid_16d_action(self):
+        policy_module = importlib.import_module("roboimi.demos.diana_air_insert_policy")
+        policy_cls = getattr(policy_module, "TestAirInsertPolicy", None)
+        self.assertIsNotNone(policy_cls)
+
+        task_state = act_ex_utils.sample_air_insert_socket_peg_state()
+        policy = policy_cls(inject_noise=False)
+        action = policy.predict(task_state, 0)
+
+        self.assertEqual(action.shape, (16,))
+        np.testing.assert_array_equal(action[-2:], np.array([100, 100]))
+
+    def test_air_insert_policy_inserts_peg_front_view_right_to_left_along_world_x(self):
+        policy_module = importlib.import_module("roboimi.demos.diana_air_insert_policy")
+        policy_cls = getattr(policy_module, "TestAirInsertPolicy", None)
+        self.assertIsNotNone(policy_cls)
+
+        task_state = self._canonical_task_state()
+        policy = policy_cls(inject_noise=False)
+        policy.generate_trajectory(task_state)
+
+        start_waypoint = next(wp for wp in policy.right_trajectory if wp["t"] == policy.INSERT_START_T)
+        end_waypoint = next(wp for wp in policy.right_trajectory if wp["t"] == policy.INSERT_END_T)
+
+        self.assertLess(
+            end_waypoint["xyz"][0],
+            start_waypoint["xyz"][0] - 0.10,
+            "front-view right-to-left peg insertion should decrease world x substantially",
+        )
+        self.assertAlmostEqual(float(end_waypoint["xyz"][1]), float(start_waypoint["xyz"][1]), delta=0.02)
+        expected_insert_end_x = float(task_state["socket_pos"][0] + 0.168)
+        self.assertAlmostEqual(float(end_waypoint["xyz"][0]), expected_insert_end_x, delta=0.02)
+        self.assertGreater(float(start_waypoint["xyz"][2]), 0.70)
+
+    def test_air_insert_policy_default_left_grasps_socket_and_right_grasps_peg(self):
+        policy_module = importlib.import_module("roboimi.demos.diana_air_insert_policy")
+        policy_cls = getattr(policy_module, "TestAirInsertPolicy", None)
+        self.assertIsNotNone(policy_cls)
+
+        task_state = {
+            "socket_pos": np.array([-0.18, 0.78, 0.472], dtype=np.float32),
+            "socket_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32),
+            "peg_pos": np.array([0.16, 0.98, 0.46], dtype=np.float32),
+            "peg_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32),
+        }
+
+        policy = policy_cls(inject_noise=False)
+        policy.generate_trajectory(task_state)
+        left_close = next(wp for wp in policy.left_trajectory if wp["t"] == 180)
+        right_close = next(wp for wp in policy.right_trajectory if wp["t"] == 180)
+        action_z_offset = getattr(policy_cls, "ACTION_OBJECT_Z_OFFSET", 0.11)
+        expected_socket_pick = task_state["socket_pos"] + np.array([-0.078, 0.0, action_z_offset])
+        expected_peg_pick = task_state["peg_pos"] + np.array([0.078, 0.0, action_z_offset + 0.01])
+
+        np.testing.assert_allclose(left_close["xyz"], expected_socket_pick, atol=1e-6)
+        np.testing.assert_allclose(right_close["xyz"], expected_peg_pick, atol=1e-6)
+        self.assertLess(left_close["gripper"], 0, "default policy should close the left gripper on the socket")
+        self.assertLess(right_close["gripper"], 0, "default policy should close the right gripper on the peg")
+
+    def test_air_insert_policy_socket_hold_tracks_socket_xy_without_sweeping_laterally(self):
+        policy_module = importlib.import_module("roboimi.demos.diana_air_insert_policy")
+        policy_cls = getattr(policy_module, "TestAirInsertPolicy", None)
+        self.assertIsNotNone(policy_cls)
+
+        base_state = {
+            "socket_pos": np.array([-0.20, 0.72, 0.472], dtype=np.float32),
+            "socket_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32),
+            "peg_pos": np.array([0.14, 0.76, 0.46], dtype=np.float32),
+            "peg_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32),
+        }
+        shifted_state = dict(base_state)
+        shifted_state["socket_pos"] = np.array([-0.06, 0.99, 0.472], dtype=np.float32)
+
+        base_policy = policy_cls(inject_noise=False)
+        base_policy.generate_trajectory(base_state)
+        shifted_policy = policy_cls(inject_noise=False)
+        shifted_policy.generate_trajectory(shifted_state)
+
+        base_hold = next(wp for wp in base_policy.left_trajectory if wp["t"] == 450)
+        shifted_hold = next(wp for wp in shifted_policy.left_trajectory if wp["t"] == 450)
+        np.testing.assert_allclose(
+            base_hold["xyz"][:2],
+            base_state["socket_pos"][:2] + np.array([-0.078, 0.0]),
+            atol=1e-6,
+        )
+        np.testing.assert_allclose(
+            shifted_hold["xyz"][:2],
+            shifted_state["socket_pos"][:2] + np.array([-0.078, 0.0]),
+            atol=1e-6,
+        )
+
+    def test_air_insert_policy_predicts_through_full_episode_without_exhausting_waypoints(self):
+        policy_module = importlib.import_module("roboimi.demos.diana_air_insert_policy")
+        policy_cls = getattr(policy_module, "TestAirInsertPolicy", None)
+        self.assertIsNotNone(policy_cls)
+
+        task_state = self._canonical_task_state()
+        policy = policy_cls(inject_noise=False)
+
+        for step in range(SIM_TASK_CONFIGS[TASK_NAME]["episode_len"]):
+            action = policy.predict(task_state, step)
+            self.assertEqual(action.shape, (16,))
+
+    def test_scripted_rollout_entrypoint_selects_socket_peg_sampler_and_policy(self):
+        rollout_module = importlib.import_module("roboimi.demos.diana_record_sim_episodes")
+        sampler_fn = getattr(rollout_module, "sample_task_state", None)
+        policy_factory = getattr(rollout_module, "make_policy", None)
+        self.assertIsNotNone(sampler_fn)
+        self.assertIsNotNone(policy_factory)
+
+        task_state = sampler_fn(TASK_NAME)
+        self.assertEqual(list(task_state.keys()), ["socket_pos", "socket_quat", "peg_pos", "peg_quat"])
+
+        policy = policy_factory(TASK_NAME, inject_noise=False)
+        self.assertEqual(policy.__class__.__name__, "TestAirInsertPolicy")
+
+    def test_real_headless_smoke_instantiates_resets_and_steps_new_task_once(self):
+        policy_module = importlib.import_module("roboimi.demos.diana_air_insert_policy")
+        policy_cls = getattr(policy_module, "TestAirInsertPolicy", None)
+        self.assertIsNotNone(policy_cls)
+
+        task_state = act_ex_utils.sample_air_insert_socket_peg_state()
+        env = make_sim_env(TASK_NAME, headless=True)
+        policy = policy_cls(inject_noise=False)
+
+        try:
+            env.reset(task_state)
+            action = policy.predict(task_state, 0)
+            env.step(action)
+            self.assertIsNotNone(env.obs)
+            self.assertIn("qpos", env.obs)
+            self.assertIn("images", env.obs)
+        finally:
+            env.exit_flag = True
+            cam_thread = getattr(env, "cam_thread", None)
+            if cam_thread is not None:
+                cam_thread.join(timeout=1.0)
+            viewer = getattr(env, "viewer", None)
+            if viewer is not None:
+                viewer.close()
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_eval_vla_headless.py

@@ -36,8 +36,8 @@ class _FakeEnv:
         self.render_calls = 0
         self.reset_calls = []
 
-    def reset(self, box_pos):
+    def reset(self, task_state):
-        self.reset_calls.append(np.array(box_pos))
+        self.reset_calls.append(task_state)
 
     def _get_image_obs(self):
         self.image_obs_calls += 1
@@ -74,7 +74,7 @@ class _FakeRenderer:
         self._env = env
         self._frames = [
             np.full((4, 4, 3), fill_value=index, dtype=np.uint8)
-            for index in range(5)
+            for index in range(8)
         ]
         self._index = 0
 
@@ -144,7 +144,7 @@ class EvalVLAHeadlessTest(unittest.TestCase):
             is_render=False,
             control_freq=30,
             is_interpolate=True,
-            cam_view="angle",
+            cam_view="top",
         )
 
     def test_headless_sync_camera_capture_populates_images_without_gui_calls(self):
@@ -153,11 +153,10 @@ class EvalVLAHeadlessTest(unittest.TestCase):
         env.mj_data = object()
         env.exit_flag = False
         env.is_render = False
-        env.cam = 'angle'
+        env.cam = 'top'
         env.r_vis = None
         env.l_vis = None
         env.top = None
-        env.angle = None
         env.front = None
         env._offscreen_renderer = None
 
@@ -176,7 +175,6 @@ class EvalVLAHeadlessTest(unittest.TestCase):
         self.assertIsNotNone(env.r_vis)
         self.assertIsNotNone(env.l_vis)
         self.assertIsNotNone(env.top)
-        self.assertIsNotNone(env.angle)
         self.assertIsNotNone(env.front)
 
     def test_cam_start_skips_background_thread_when_headless(self):
@@ -196,11 +194,10 @@ class EvalVLAHeadlessTest(unittest.TestCase):
         env.mj_data = object()
         env.exit_flag = False
         env.is_render = False
-        env.cam = "angle"
+        env.cam = "top"
         env.r_vis = None
         env.l_vis = None
         env.top = None
-        env.angle = None
         env.front = None
 
         with mock.patch(
@@ -217,7 +214,6 @@ class EvalVLAHeadlessTest(unittest.TestCase):
         self.assertIsNotNone(env.r_vis)
         self.assertIsNotNone(env.l_vis)
         self.assertIsNotNone(env.top)
-        self.assertIsNotNone(env.angle)
         self.assertIsNotNone(env.front)
 
     def test_eval_main_headless_skips_render_and_still_executes_policy(self):
@@ -327,6 +323,69 @@ class EvalVLAHeadlessTest(unittest.TestCase):
         self.assertAlmostEqual(summary["avg_reward"], 3.75)
         self.assertEqual(summary["num_episodes"], 2)
 
+    def test_run_eval_uses_air_insert_sampler_for_socket_peg_task(self):
+        self.assertTrue(
+            hasattr(eval_vla, "sample_air_insert_socket_peg_state"),
+            "Expected eval_vla to expose the new socket/peg reset sampler",
+        )
+
+        fake_env = _FakeEnv()
+        fake_agent = _FakeAgent()
+        sampled_task_state = {
+            "socket_pos": np.array([-0.10, 0.80, 0.47], dtype=np.float32),
+            "socket_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32),
+            "peg_pos": np.array([0.10, 0.82, 0.47], dtype=np.float32),
+            "peg_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32),
+        }
+        cfg = OmegaConf.create(
+            {
+                "agent": {},
+                "eval": {
+                    "ckpt_path": "checkpoints/vla_model_best.pt",
+                    "num_episodes": 1,
+                    "max_timesteps": 1,
+                    "device": "cpu",
+                    "task_name": "sim_air_insert_socket_peg",
+                    "camera_names": ["front"],
+                    "use_smoothing": False,
+                    "smooth_alpha": 0.3,
+                    "verbose_action": False,
+                    "headless": True,
+                },
+            }
+        )
+
+        with mock.patch.object(
+            eval_vla,
+            "load_checkpoint",
+            return_value=(fake_agent, None),
+        ), mock.patch.object(
+            eval_vla,
+            "make_sim_env",
+            return_value=fake_env,
+        ) as make_env, mock.patch.object(
+            eval_vla,
+            "sample_air_insert_socket_peg_state",
+            return_value=sampled_task_state,
+        ) as socket_peg_sampler, mock.patch.object(
+            eval_vla,
+            "sample_transfer_pose",
+            side_effect=AssertionError("sample_transfer_pose should not be used for sim_air_insert_socket_peg"),
+        ), mock.patch.object(
+            eval_vla,
+            "execute_policy_action",
+        ) as execute_policy_action, mock.patch.object(
+            eval_vla,
+            "tqdm",
+            side_effect=lambda iterable, **kwargs: iterable,
+        ):
+            eval_vla._run_eval(cfg)
+
+        make_env.assert_called_once_with("sim_air_insert_socket_peg", headless=True)
+        socket_peg_sampler.assert_called_once_with()
+        execute_policy_action.assert_called_once()
+        self.assertEqual(fake_env.reset_calls, [sampled_task_state])
+
 
 if __name__ == "__main__":
     unittest.main()

tests/test_robot_asset_paths.py

@@ -4,7 +4,7 @@ import unittest
 from pathlib import Path
 from unittest import mock
 
-from roboimi.assets.robots.diana_med import BiDianaMed
+from roboimi.assets.robots import diana_med
 
 
 class _FakeKDL:
@@ -24,6 +24,7 @@ class RobotAssetPathResolutionTest(unittest.TestCase):
         _FakeKDL.reset_calls = []
 
     def test_bidianamed_resolves_robot_asset_paths_independent_of_cwd(self):
+        BiDianaMed = diana_med.BiDianaMed
         repo_root = Path(__file__).resolve().parents[1]
         expected_xml = repo_root / 'roboimi/assets/models/manipulators/DianaMed/bi_diana_transfer_ee.xml'
         expected_urdf = repo_root / 'roboimi/assets/models/manipulators/DianaMed/DualDianaMed.urdf'
@@ -58,6 +59,47 @@ class RobotAssetPathResolutionTest(unittest.TestCase):
         self.assertEqual({Path(path) for path in _FakeKDL.init_calls}, {expected_urdf})
         self.assertTrue(all(Path(path).is_absolute() for path in _FakeKDL.init_calls))
 
+    def test_bidianamed_socket_peg_resolves_robot_asset_paths_independent_of_cwd(self):
+        BiDianaMedSocketPeg = getattr(diana_med, 'BiDianaMedSocketPeg', None)
+        self.assertIsNotNone(
+            BiDianaMedSocketPeg,
+            'Expected roboimi.assets.robots.diana_med.BiDianaMedSocketPeg',
+        )
+
+        repo_root = Path(__file__).resolve().parents[1]
+        expected_xml = repo_root / 'roboimi/assets/models/manipulators/DianaMed/bi_diana_socket_peg_ee.xml'
+        expected_urdf = repo_root / 'roboimi/assets/models/manipulators/DianaMed/DualDianaMed.urdf'
+        xml_calls = []
+
+        def fake_from_xml_path(*, filename, assets=None):
+            xml_calls.append((filename, assets))
+            return object()
+
+        with tempfile.TemporaryDirectory() as tempdir:
+            previous_cwd = os.getcwd()
+            try:
+                os.chdir(tempdir)
+                with mock.patch(
+                    'roboimi.assets.robots.arm_base.mujoco.MjModel.from_xml_path',
+                    side_effect=fake_from_xml_path,
+                ), mock.patch(
+                    'roboimi.assets.robots.arm_base.mujoco.MjData',
+                    return_value=object(),
+                ), mock.patch(
+                    'roboimi.assets.robots.arm_base.KDL_utils',
+                    _FakeKDL,
+                ):
+                    BiDianaMedSocketPeg()
+            finally:
+                os.chdir(previous_cwd)
+
+        self.assertEqual(len(xml_calls), 1)
+        self.assertEqual(Path(xml_calls[0][0]), expected_xml)
+        self.assertTrue(Path(xml_calls[0][0]).is_absolute())
+        self.assertGreaterEqual(len(_FakeKDL.init_calls), 2)
+        self.assertEqual({Path(path) for path in _FakeKDL.init_calls}, {expected_urdf})
+        self.assertTrue(all(Path(path).is_absolute() for path in _FakeKDL.init_calls))
+
 
 if __name__ == '__main__':
     unittest.main()