docs(sim): save socket peg air insert notes

MUJOCO_LOG.TXT

@@ -0,0 +1,72 @@
+Sat Apr 25 10:03:15 2026
+WARNING: OpenGL error 0x502 in or before mjr_makeContext
+
+Sat Apr 25 10:12:58 2026
+WARNING: OpenGL error 0x502 in or before mjr_makeContext
+
+Sat Apr 25 10:13:59 2026
+WARNING: OpenGL error 0x502 in or before mjr_makeContext
+
+Sat Apr 25 10:14:52 2026
+WARNING: OpenGL error 0x502 in or before mjr_makeContext
+
+Sat Apr 25 10:15:27 2026
+WARNING: OpenGL error 0x502 in or before mjr_makeContext
+
+Sat Apr 25 10:16:03 2026
+WARNING: OpenGL error 0x502 in or before mjr_makeContext
+
+Sat Apr 25 10:18:11 2026
+WARNING: OpenGL error 0x502 in or before mjr_makeContext
+
+Sat Apr 25 10:19:03 2026
+WARNING: OpenGL error 0x502 in or before mjr_makeContext
+
+Sat Apr 25 10:20:06 2026
+WARNING: OpenGL error 0x502 in or before mjr_makeContext
+
+Sat Apr 25 10:20:38 2026
+WARNING: OpenGL error 0x502 in or before mjr_makeContext
+
+Sat Apr 25 10:21:23 2026
+WARNING: OpenGL error 0x502 in or before mjr_makeContext
+
+Sat Apr 25 10:21:55 2026
+WARNING: OpenGL error 0x502 in or before mjr_makeContext
+
+Sat Apr 25 10:22:29 2026
+WARNING: OpenGL error 0x502 in or before mjr_makeContext
+
+Sat Apr 25 10:23:01 2026
+WARNING: OpenGL error 0x502 in or before mjr_makeContext
+
+Sat Apr 25 10:23:33 2026
+WARNING: OpenGL error 0x502 in or before mjr_makeContext
+
+Sat Apr 25 10:25:01 2026
+WARNING: OpenGL error 0x502 in or before mjr_makeContext
+
+Sat Apr 25 10:25:59 2026
+WARNING: OpenGL error 0x502 in or before mjr_makeContext
+
+Sat Apr 25 10:26:52 2026
+WARNING: OpenGL error 0x502 in or before mjr_makeContext
+
+Sat Apr 25 10:27:20 2026
+WARNING: OpenGL error 0x502 in or before mjr_makeContext
+
+Sat Apr 25 10:27:49 2026
+WARNING: OpenGL error 0x502 in or before mjr_makeContext
+
+Sat Apr 25 10:28:18 2026
+WARNING: OpenGL error 0x502 in or before mjr_makeContext
+
+Sat Apr 25 10:37:53 2026
+WARNING: OpenGL error 0x502 in or before mjr_makeContext
+
+Sat Apr 25 10:38:56 2026
+WARNING: OpenGL error 0x502 in or before mjr_makeContext
+
+Sat Apr 25 10:41:12 2026
+WARNING: OpenGL error 0x502 in or before mjr_makeContext
+

docs/superpowers/plans/2026-04-01-imf-attnres-policy-migration.md

@@ -1,268 +0,0 @@
-# IMF-AttnRes Policy Migration 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:** 将 external `diffusion_policy@185ed659` 的 IMF-AttnRes 模型、训练目标和一步推理机制迁移到 RoboIMI，并在保持三相机视觉条件输入与现有训练/rollout 工作流的前提下启动同参数训练。
-
-**Architecture:** 保留 RoboIMI 现有 ResNet 三相机观测编码、normalization、queue-based online rollout 和训练脚本；新增 AttnRes 组件与 IMF transformer head，并新增 IMF 专用 agent 以覆盖 DDPM loss / DDIM inference 语义。训练脚本只做最小接线修改，让新 head/agent 能用现有 optimizer、checkpoint、SwanLab 和 headless rollout。
-
-**Tech Stack:** PyTorch, Hydra, diffusers schedulers (仅保留兼容初始化), MuJoCo rollout, unittest, SwanLab
-
----
-
-## File Map
-
-### New files
-- `roboimi/vla/models/heads/attnres_transformer_components.py` — 本地 IMF AttnRes 基础组件
-- `roboimi/vla/models/heads/imf_transformer1d.py` — IMF transformer head，暴露 `forward(sample, r, t, cond=None)`
-- `roboimi/vla/agent_imf.py` — IMF 专用 VLA agent，复用现有观测/队列/normalization 逻辑并覆盖 loss / inference
-- `roboimi/vla/conf/head/imf_transformer1d.yaml` — IMF head 配置
-- `roboimi/vla/conf/agent/resnet_imf_attnres.yaml` — IMF agent + backbone/head 组合配置
-- `tests/test_imf_transformer1d_external_alignment.py` — external `185ed659` 对齐测试
-- `tests/test_imf_vla_agent.py` — IMF agent 的 loss / inference / queue 语义测试
-
-### Modified files
-- `roboimi/demos/vla_scripts/train_vla.py` — 优化器参数分组接线；确保新 agent 能无缝训练
-- `roboimi/vla/conf/config.yaml` — 保持默认配置不变，仅支持通过 override 启用 IMF agent
-- `tests/test_train_vla_transformer_optimizer.py` — 覆盖 IMF head 的 optimizer-group 行为
-- （如需要）`roboimi/vla/models/heads/__init__.py` 或相近导出文件 — 暴露新 head
-
----
-
-### Task 1: 写 IMF transformer 对齐测试
-
-**Files:**
-- Create: `tests/test_imf_transformer1d_external_alignment.py`
-- Reference: `/home/droid/project/diffusion_policy/diffusion_policy/model/diffusion/attnres_transformer_components.py`
-- Reference: `/home/droid/project/diffusion_policy/diffusion_policy/model/diffusion/imf_transformer_for_diffusion.py`
-
-- [ ] **Step 1: 写失败测试，验证 local IMF head 与 external `185ed659` 的 state-dict key、前向 shape、forward 数值、optim groups 对齐**
-
-```python
-with torch.no_grad():
-    external_out = external_model(sample=sample, r=r, t=t, cond=cond)
-    local_out = local_model(sample=sample, r=r, t=t, cond=cond)
-assert torch.allclose(local_out, external_out, atol=1e-6, rtol=1e-5)
-```
-
-- [ ] **Step 2: 运行单测，确认当前失败**
-
-Run: `python -m unittest tests.test_imf_transformer1d_external_alignment -v`
-Expected: FAIL，提示 `imf_transformer1d` / `attnres` 模块不存在
-
-- [ ] **Step 3: 若测试需要复用现有 external-loader 逻辑，则从 `tests/test_transformer1d_external_alignment.py` 复制最小必要 helper，避免重复依赖 session context**
-
-- [ ] **Step 4: 提交测试骨架**
-
-```bash
-git add tests/test_imf_transformer1d_external_alignment.py
-git commit -m "test: add IMF transformer external alignment coverage"
-```
-
-### Task 2: 实现 AttnRes 组件与 IMF transformer head
-
-**Files:**
-- Create: `roboimi/vla/models/heads/attnres_transformer_components.py`
-- Create: `roboimi/vla/models/heads/imf_transformer1d.py`
-- Modify: `tests/test_imf_transformer1d_external_alignment.py`
-
-- [ ] **Step 1: 按 external `185ed659` 迁移 AttnRes 基础组件，保持命名和参数语义一致**
-
-必须包含：
-- `RMSNorm`
-- `RMSNormNoWeight`
-- `precompute_rope_freqs`
-- `apply_rope`
-- `GroupedQuerySelfAttention`
-- `SwiGLUFFN`
-- `AttnResOperator`
-- `AttnResSubLayer`
-- `AttnResTransformerBackbone`
-
-- [ ] **Step 2: 在 `imf_transformer1d.py` 中实现本地 IMF head**
-
-必须满足：
-- `forward(sample, r, t, cond=None)`
-- 默认支持 `backbone_type='attnres_full'`
-- token 序列为 `[r_token, t_token, cond_tokens..., sample_tokens...]`
-- 输出只切回 sample token 段
-- 保留 `get_optim_groups()` 供 AdamW 分组
-
-- [ ] **Step 3: 运行对齐测试，修正 state-dict key / init / no-decay 参数分组不一致问题**
-
-Run: `python -m unittest tests.test_imf_transformer1d_external_alignment -v`
-Expected: PASS
-
-- [ ] **Step 4: 提交模型组件实现**
-
-```bash
-git add roboimi/vla/models/heads/attnres_transformer_components.py \
-        roboimi/vla/models/heads/imf_transformer1d.py \
-        tests/test_imf_transformer1d_external_alignment.py
-git commit -m "feat: add IMF AttnRes transformer head"
-```
-
-### Task 3: 写 IMF agent 行为测试
-
-**Files:**
-- Create: `tests/test_imf_vla_agent.py`
-- Reference: `roboimi/vla/agent.py`
-- Reference: `tests/test_resnet_transformer_agent_wiring.py`
-
-- [ ] **Step 1: 写失败测试，覆盖 IMF agent 的核心契约**
-
-需要覆盖：
-1. `compute_loss()` 接受当前 batch 结构并返回标量 loss
-2. `predict_action()` 输出 `(B, pred_horizon, action_dim)`
-3. `select_action()` 仍按 queue/chunk 语义工作
-4. `predict_action()` 不走 DDIM 多步循环，而是只触发一步 IMF sample
-5. `action_is_pad` 存在时仅在有效 action 上计 loss
-
-- [ ] **Step 2: 用 stub backbone / stub head 记录调用参数，验证 `r,t,cond` 的传递与 observation conditioning 维度正确**
-
-```python
-self.assertEqual(recorded['cond'].shape, (B, obs_horizon, expected_cond_dim))
-self.assertTrue(torch.allclose(recorded['r'], torch.zeros(B)))
-self.assertTrue(torch.allclose(recorded['t'], torch.ones(B)))
-```
-
-- [ ] **Step 3: 运行测试，确认当前失败**
-
-Run: `python -m unittest tests.test_imf_vla_agent -v`
-Expected: FAIL，提示 `roboimi.vla.agent_imf` 不存在
-
-- [ ] **Step 4: 提交测试骨架**
-
-```bash
-git add tests/test_imf_vla_agent.py
-git commit -m "test: add IMF VLA agent behavior coverage"
-```
-
-### Task 4: 实现 IMF agent 与 Hydra 接线
-
-**Files:**
-- Create: `roboimi/vla/agent_imf.py`
-- Create: `roboimi/vla/conf/head/imf_transformer1d.yaml`
-- Create: `roboimi/vla/conf/agent/resnet_imf_attnres.yaml`
-- Modify: `roboimi/demos/vla_scripts/train_vla.py`
-- Modify: `tests/test_train_vla_transformer_optimizer.py`
-- Modify: `tests/test_imf_vla_agent.py`
-
-- [ ] **Step 1: 以 `VLAAgent` 为基础实现 `IMFVLAAgent`**
-
-实现策略：
-- 复用 `VLAAgent.__init__`、`_build_cond()`、`reset()`、`_populate_queues()`、`_prepare_observation_batch()`、`select_action()`、`get_normalization_stats()`
-- 覆盖：
-  - `compute_loss()` -> IMF objective
-  - `predict_action()` -> one-step sample
-- 提供内部 helper：
-  - `_broadcast_batch_time`
-  - `_apply_conditioning`（如需）
-  - `_compute_u_and_du_dt`
-  - `_compound_velocity`
-  - `_sample_one_step`
-
-- [ ] **Step 2: 在 JVP 路径中加入 CUDA math SDPA fallback，保持 external repo 的稳定性策略**
-
-- [ ] **Step 3: 新增 Hydra 配置，让 `agent=resnet_imf_attnres` 可实例化**
-
-关键默认值：
-- `_target_: roboimi.vla.agent_imf.IMFVLAAgent`
-- `head._target_: roboimi.vla.models.heads.imf_transformer1d.IMFTransformer1D`
-- `head.backbone_type: attnres_full`
-- `head.causal_attn: false`
-- `head.time_as_cond: true`
-- `head.n_cond_layers: 0`
-- `inference_steps: 1`
-- `camera_names: ${data.camera_names}`
-- `vision_backbone.camera_names: ${agent.camera_names}`
-
-- [ ] **Step 4: 让训练脚本对任何带 `get_optim_groups()` 的 head 复用参数分组，而不是硬编码旧 transformer head_type**
-
-推荐最小改法：
-```python
-use_head_groups = callable(getattr(noise_pred_net, 'get_optim_groups', None))
-```
-
-- [ ] **Step 5: 运行测试并修复 wiring 问题**
-
-Run:
-- `python -m unittest tests.test_imf_vla_agent -v`
-- `python -m unittest tests.test_train_vla_transformer_optimizer -v`
-
-Expected: PASS
-
-- [ ] **Step 6: 提交 agent / config / train-script 接线**
-
-```bash
-git add roboimi/vla/agent_imf.py \
-        roboimi/vla/conf/head/imf_transformer1d.yaml \
-        roboimi/vla/conf/agent/resnet_imf_attnres.yaml \
-        roboimi/demos/vla_scripts/train_vla.py \
-        tests/test_imf_vla_agent.py \
-        tests/test_train_vla_transformer_optimizer.py
-git commit -m "feat: add IMF VLA agent and training wiring"
-```
-
-### Task 5: 集成验证与训练启动
-
-**Files:**
-- Modify: none required unless验证暴露真实问题
-- Use run artifacts under: `runs/`
-
-- [ ] **Step 1: 运行聚焦测试集**
-
-Run:
-```bash
-python -m unittest \
-  tests.test_imf_transformer1d_external_alignment \
-  tests.test_imf_vla_agent \
-  tests.test_resnet_transformer_agent_wiring \
-  tests.test_train_vla_transformer_optimizer -v
-```
-Expected: PASS
-
-- [ ] **Step 2: 运行一个最小 GPU 训练冒烟任务（不必长跑）**
-
-Run:
-```bash
-/home/droid/.conda/envs/roboimi/bin/python roboimi/demos/vla_scripts/train_vla.py \
-  agent=resnet_imf_attnres \
-  data.dataset_dir=/home/droid/project/diana_sim/sim_transfer \
-  data.camera_names=[r_vis,top,front] \
-  train.device=cuda train.max_steps=2 train.batch_size=4 train.num_workers=2 \
-  train.use_swanlab=false train.rollout_val_freq_epochs=0
-```
-Expected: 成功完成 2 steps，生成 checkpoint / log，无 shape 或 JVP 错误
-
-- [ ] **Step 3: 用正式参数启动 IMF 训练**
-
-Run:
-```bash
-/home/droid/.conda/envs/roboimi/bin/python roboimi/demos/vla_scripts/train_vla.py \
-  agent=resnet_imf_attnres \
-  data.dataset_dir=/home/droid/project/diana_sim/sim_transfer \
-  data.camera_names=[r_vis,top,front] \
-  train.device=cuda train.val_split=0.0 train.seed=42 \
-  train.batch_size=80 train.lr=5e-4 train.num_workers=12 train.max_steps=150000 \
-  train.log_freq=100 train.save_freq=10000 train.use_swanlab=true \
-  train.swanlab_project=roboimi-vla \
-  train.rollout_val_freq_epochs=5 train.rollout_validate_on_checkpoint=false \
-  train.rollout_num_episodes=5 train.warmup_steps=2000 \
-  train.scheduler_type=cosine train.min_lr=1e-6 train.weight_decay=1e-5 train.grad_clip=1.0 \
-  agent.pred_horizon=16 agent.inference_steps=1 \
-  agent.head.n_emb=384 agent.head.n_layer=18 agent.head.n_head=1 agent.head.n_kv_head=1 \
-  agent.vision_backbone.pretrained_backbone_weights=null \
-  agent.vision_backbone.freeze_backbone=false \
-  agent.vision_backbone.use_separate_rgb_encoder_per_camera=true
-```
-Expected: 训练启动成功，SwanLab 记录完整 config，5 epoch 一次 headless rollout
-
-- [ ] **Step 4: 记录 run 路径、训练 PID、SwanLab 运行名并向用户汇报**
-
-- [ ] **Step 5: 提交最终收尾改动（如果 smoke fix 需要额外 patch）**
-
-```bash
-git add <changed files>
-git commit -m "chore: verify IMF AttnRes training launch"
-```

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/plans/2026-04-30-socket-peg-air-insert.md

@@ -0,0 +1,72 @@
+# 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.

docs/superpowers/specs/2026-04-01-imf-attnres-policy-design.md

@@ -1,272 +0,0 @@
-# IMF-AttnRes Policy Migration Design
-
-**Date:** 2026-04-01
-**Status:** Approved in chat, written spec pending review
-
-## Goal
-
-将 `/home/droid/project/diffusion_policy` 中提交 `185ed659` 的 IMF-AttnRes diffusion policy 迁移到当前 `roboimi` 仓库，作为当前 DiT / Transformer diffusion policy 的替代训练选项；同时迁移其训练目标与一步推理机制，并保持 RoboIMI 现有的仿真环境、三相机视觉输入、数据集格式、训练脚本和 rollout 验证工作流可继续使用。
-
-## Non-Goals
-
-- 不迁移 external repo 中与当前任务无关的 obs encoder、dataset、env wrapper、PushT 专用逻辑。
-- 不强行复刻 external repo 中全部目录结构；仅迁移当前 RoboIMI 训练所必需的模型、loss、inference 语义。
-- 不在本次工作中同时保留旧 DiT 为默认训练目标；旧配置继续可用，但新模型单独提供 config 入口。
-
-## User-Confirmed Requirements
-
-1. 迁移对象是 `185ed659` 中的 **IMF-AttnRes 模型相关代码**。
-2. 不只是迁移骨架，还要迁移：
-   - **训练目标**
-   - **一步推理机制**
-3. 视觉输入与当前 RoboIMI diffusion policy 一致：
-   - 使用三个相机图像作为条件输入
-   - 图像观测必须作为条件，而不是拼进输出预测目标
-4. 当前任务里，IMF policy 用来替代现有 DiT/Transformer diffusion policy 训练。
-5. 训练参数沿用最近一次训练的大体设置（后续由训练命令显式覆盖），但推理方式改为 IMF 的 one-step 机制。
-6. 用户接受 IMF 中“全注意力 / 非因果注意力”的实现约束。
-
-## External Source of Truth
-
-迁移语义以 external repo 的以下文件为准：
-
-- `diffusion_policy/model/diffusion/attnres_transformer_components.py`
-- `diffusion_policy/model/diffusion/imf_transformer_for_diffusion.py`
-- `diffusion_policy/policy/imf_transformer_hybrid_image_policy.py`
-- 参考配置：`image_pusht_diffusion_policy_dit_imf_attnres_full.yaml`
-
-其中最关键的差异是：该策略并非 DDPM/DDIM 多步去噪，而是 IMF 训练目标 + one-step 推理。
-
-## Current RoboIMI Baseline
-
-当前 RoboIMI 中与该任务直接相关的基线如下：
-
-- 视觉编码：`ResNetDiffusionBackbone`
-  - 三相机：`r_vis`, `top`, `front`
-  - 每个时间步将相机特征与 `qpos` 拼接为 per-step condition
-- 策略主体：`VLAAgent`
-  - `compute_loss()` 使用 DDPM 噪声预测损失
-  - `predict_action()` 使用 DDIM 多步采样
-  - 在线控制通过动作队列机制在 `select_action()` 中按 chunk 触发预测
-- 训练脚本：`roboimi/demos/vla_scripts/train_vla.py`
-  - 支持 GPU 训练、SwanLab 日志、headless rollout 验证
-
-因此，本次迁移的核心不是换视觉 backbone，而是替换 **head + loss + inference semantics**。
-
-## Recommended Integration Approach
-
-采用 **最小侵入式集成**：
-
-1. **保留当前 RoboIMI 的视觉编码、数据读取、rollout/eval、训练脚本主框架**。
-2. **新增 IMF 专用 head 模块**，在 RoboIMI 内本地实现：
-   - AttnRes 组件
-   - IMF transformer 主体
-3. **新增 IMF 专用 agent**，复用当前 `VLAAgent` 的：
-   - 归一化逻辑
-   - 相机顺序管理
-   - 观测缓存 / 动作 chunk 缓存
-   - rollout 接口
-   但覆盖：
-   - `compute_loss()`
-   - `predict_action()`
-4. **新增独立 Hydra config**，让 IMF policy 作为新的 agent 选项，不破坏已有 resnet_transformer / gr00t_dit 配置。
-
-这样做的原因：
-
-- 迁移 IMF 语义时不必把当前 DDPM agent 搅乱；
-- rollout / eval / checkpoint 逻辑仍然可复用；
-- 便于和现有 Transformer / DiT 直接做 A/B 对比训练。
-
-## Architecture
-
-### 1. Observation / Conditioning Path
-
-沿用当前 RoboIMI 的视觉路径：
-
-- 输入观测：`images={r_vis, top, front}` + `qpos`
-- `ResNetDiffusionBackbone` 对每个相机编码，得到 per-camera feature
-- `state_encoder` 编码 `qpos`
-- 将三相机特征与 state feature 按时间步拼接，形成 `per_step_cond`
-
-这里不迁移 external repo 的 obs_encoder 实现；我们只对齐 **“图像作为条件 token 输入 transformer”** 这一语义。
-
-### 2. Condition Tokenization
-
-对齐 external IMF transformer 的 token 使用方式：
-
-- action trajectory token：由 `(B, pred_horizon, action_dim)` 通过线性层映射到 `n_emb`
-- time token：两个标量 `r` 与 `t`，分别通过 sinusoidal embedding + linear projection 得到 token
-- observation token：`per_step_cond` 通过线性层映射到 `n_emb`
-- 最终 token 序列为：
-  - `[r_token, t_token, obs_cond_tokens..., action_tokens...]`
-
-在当前任务中，obs token 数量等于 `obs_horizon`，且图像观测始终作为条件输入。
-
-### 3. IMF-AttnRes Backbone
-
-在 RoboIMI 内新增 AttnRes backbone 实现，保持 external commit 的关键语义：
-
-- `RMSNorm` / `RMSNormNoWeight`
-- RoPE
-- Grouped Query Self-Attention
-- SwiGLU FFN
-- AttnRes operator / residual source aggregation
-- `AttnResTransformerBackbone`
-
-并保持：
-
-- **full attention**（不使用因果注意力）
-- `backbone_type='attnres_full'`
-- 输出仅切回 action token 部分，再经过最终 norm + head 得到 velocity-like 输出
-
-### 4. Training Objective
-
-训练目标从当前 DDPM epsilon prediction 改为 external IMF 目标：
-
-给定真实轨迹 `x` 与随机噪声 `e`：
-
-1. 采样 `t ~ U(0,1)`、`r ~ U(0,1)`，并排序为 `t >= r`
-2. 构造插值状态：
-   - `z_t = (1 - t) x + t e`
-3. 用模型计算：
-   - `v = f(z_t, t, t, cond)`
-4. 对 `g(z, r, t) = f(z, r, t, cond)` 做 JVP，得到：
-   - `u, du_dt`
-5. 构造 compound velocity：
-   - `V = u + (t - r) * du_dt`
-6. 目标为：
-   - `target = e - x`
-7. 用 action 维度上的 MSE 作为最终损失
-
-RoboIMI 现有 batch 中的 `action_is_pad` 仍要保留支持；如果存在 padding，只在有效 action 上计算损失。
-
-### 5. One-Step Inference
-
-推理改为 external IMF 的一步采样语义：
-
-1. 从标准高斯初始化 action trajectory `z_t`
-2. 计算 `u = f(z_t, r=0, t=1, cond)`
-3. 一步更新：
-   - `x_hat = z_t - (t-r) * u = z_t - u`
-4. 反归一化得到动作序列
-
-这意味着：
-
-- `num_inference_steps` 对 IMF policy 固定为 `1`
-- 不再调用 DDIM scheduler 的多步 `step()`
-- 在线控制中仍沿用当前 chunk 机制：
-  - 动作队列为空时触发一次 `predict_action_chunk()`
-  - 取预测序列中 `[obs_horizon-1 : obs_horizon-1+num_action_steps]` 这一段入队
-
-也就是说，**触发模型前向的规则不变，改变的是每次触发后的动作序列生成方式**。
-
-## API / Code Structure
-
-计划中的主要代码边界如下：
-
-- `roboimi/vla/models/heads/attnres_transformer_components.py`
-  - IMF AttnRes 基础组件
-- `roboimi/vla/models/heads/imf_transformer1d.py`
-  - RoboIMI 版本 IMF transformer head
-  - 对外暴露 `forward(sample, r, t, cond=None)`
-  - 暴露 `get_optim_groups()` 供 AdamW 分组使用
-- `roboimi/vla/agent_imf.py`
-  - 复用 `VLAAgent` 的观测处理 / normalization / queue 基础设施
-  - 覆盖 IMF 的训练损失与 one-step 预测逻辑
-- Hydra config
-  - `roboimi/vla/conf/head/imf_transformer1d.yaml`
-  - `roboimi/vla/conf/agent/resnet_imf_attnres.yaml`
-
-训练脚本主流程尽量不改；只要求它能 instantiate 新 agent 并继续使用当前 rollout / checkpoint / swanlab 逻辑。
-
-## Compatibility Decisions
-
-## Initial Config Defaults To Preserve
-
-为避免迁移时语义漂移，首版 IMF 配置默认值明确固定为：
-
-- `backbone_type: attnres_full`
-- `n_head: 1`
-- `n_kv_head: 1`
-- `n_cond_layers: 0`
-- `time_as_cond: true`
-- `causal_attn: false`
-- `num_inference_steps: 1`
-
-这些默认值与 external `185ed659` 的 IMF-AttnRes 使用方式保持一致；后续调参可以覆盖，但首版迁移必须先以该语义跑通。
-
-### Reuse From RoboIMI
-
-保留：
-
-- 三相机数据读取方式
-- ResNet visual backbone
-- qpos / action normalization
-- 训练循环、优化器、scheduler、SwanLab、headless rollout
-- `select_action()` 的在线 chunk 执行方式
-
-### Replace With External IMF Semantics
-
-替换：
-
-- transformer head 实现
-- diffusion training objective
-- inference sampling semantics
-
-### Intentionally Not Mirrored 1:1
-
-不强行与 external repo 一致的部分：
-
-- external repo 的整体 policy 基类继承体系
-- external repo 的 obs encoder 模块树
-- external repo 的 normalizer / mask generator 框架
-
-原因是当前 RoboIMI 已有稳定的数据接口和 rollout 流程，直接嫁接进去更稳。
-
-## Testing / Verification Strategy
-
-迁移完成后至少验证以下内容：
-
-1. **单元 / 冒烟验证**
-   - IMF head 前向 shape 正确
-   - IMF agent `compute_loss()` 在真实 batch 上可前向、反向
-   - IMF agent `predict_action()` 能输出 `(B, pred_horizon, action_dim)`
-2. **训练链路验证**
-   - 使用 GPU 跑一个短训练任务，确认：
-     - dataloader 正常
-     - optimizer / lr scheduler 正常
-     - SwanLab 正常记录配置和训练指标
-3. **rollout 验证**
-   - 训练中周期性 headless rollout 能跑通
-   - 环境仍按 EE-style `step()` 接收动作
-4. **最终交付**
-   - 用用户指定的同类超参数启动正式训练
-
-## Risks and Mitigations
-
-### Risk 1: JVP 在 CUDA 注意力内核上不稳定
-
-缓解：沿用 external repo 的策略，在 JVP 路径上切换到 math SDP kernel，必要时 fallback 到 `torch.autograd.functional.jvp`。同时，JVP 的切线构造与 `u, du_dt` 计算流程必须严格对齐 external source，不在本次迁移中自行改写其数学语义。
-
-### Risk 2: Optimizer 参数分组遗漏新模块
-
-缓解：IMF head 提供 `get_optim_groups()`，并在训练脚本中按“只要 head 提供该接口就使用”的策略统一处理，而不是绑定旧 `head_type`。
-
-### Risk 3: 现有 rollout 逻辑假定 DDIM 多步采样
-
-缓解：保持 `select_action()` / `predict_action_chunk()` 接口不变，只替换 `predict_action()` 内部实现，确保 eval 代码无需理解 IMF 细节。
-
-### Risk 4: 训练命令参数与新 config 不一致
-
-缓解：新增独立 agent config，并保留此前训练参数作为显式 CLI override 模板。
-
-## Success Criteria
-
-以下条件全部满足，视为本次迁移成功：
-
-1. RoboIMI 中新增 IMF-AttnRes policy，可通过 Hydra config 单独启用。
-2. 训练时使用 external IMF 的 loss，而不是当前 DDPM epsilon loss。
-3. 推理时使用 one-step IMF 采样，而不是 DDIM 多步采样。
-4. 三相机图像始终作为条件输入参与模型前向。
-5. 在线 rollout 能在 headless 仿真环境中跑通。
-6. 能按最近一次实验参数模板成功启动训练。

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,7 @@
           <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="left_side" pos="-0.55 0.85 0.85" fovy="65" 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

@@ -90,4 +90,40 @@ class BiDianaMed(ArmBase):
     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])
+
+
+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():
-    task_name = 'sim_transfer'
-    dataset_dir = DATASET_DIR + '/sim_transfer'   #SIM_TASK_CONFIGS[task_name]['dataset_dir']
-    num_episodes = 100       #SIM_TASK_CONFIGS[task_name]['num_episodes']
+def sample_task_state(task_name):
+    if task_name == 'sim_transfer':
+        return sample_transfer_pose()
+    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']
-    camera_names = ['angle','r_vis', 'top', 'front']   #SIM_TASK_CONFIGS[task_name]['camera_names']
+    episode_len = task_cfg['episode_len']
+    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()
-        env.reset(box_pos)
+        task_state = sample_task_state(task_name)
+        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

@@ -26,7 +26,10 @@ from hydra.utils import instantiate
 from einops import rearrange
 
 from roboimi.envs.double_pos_ctrl_env import make_sim_env
-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.vla.eval_utils import execute_policy_action
 
 sys.path.append(os.getcwd())
@@ -485,6 +488,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 评估
@@ -549,8 +560,8 @@ def _run_eval(cfg: DictConfig):
             print(f"回合 {episode_idx + 1}/{eval_cfg.num_episodes}")
             print(f"{'='*60}\n")
 
-            box_pos = sample_transfer_pose()
-            env.reset(box_pos)
+            task_state = _sample_task_reset_state(str(eval_cfg.task_name))
+            env.reset(task_state)
 
             # 为新回合重置 agent 队列
             agent.reset()

roboimi/envs/double_air_insert_env.py

@@ -0,0 +1,154 @@
+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.left_side = None
+        self.r_vis = None
+        self.front = None
+        self.cam_flage = True
+        while self.cam_flage:
+            if (
+                type(self.top) == type(None)
+                or type(self.left_side) == type(None)
+                or type(self.r_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,7 @@ class DualDianaMed(MujocoEnv):
         self.r_vis = None
         self.l_vis = None
         self.top = None
-        self.angle = None
+        self.left_side = None
         self.front = None
         self.obs = None
         
@@ -166,7 +166,7 @@ class DualDianaMed(MujocoEnv):
         obs['action'] = self.compute_qpos
         obs['images'] = dict()
         obs['images']['top']   = self.top
-        obs['images']['angle'] = self.angle
+        obs['images']['left_side'] = self.left_side
         obs['images']['r_vis'] = self.r_vis
         obs['images']['l_vis'] = self.l_vis
         obs['images']['front'] = self.front
@@ -176,7 +176,7 @@ class DualDianaMed(MujocoEnv):
         obs = collections.OrderedDict()
         obs['images'] = dict()
         obs['images']['top']   = self.top
-        obs['images']['angle'] = self.angle
+        obs['images']['left_side'] = self.left_side
         obs['images']['r_vis'] = self.r_vis
         obs['images']['l_vis'] = self.l_vis
         obs['images']['front'] = self.front
@@ -199,8 +199,8 @@ class DualDianaMed(MujocoEnv):
     def cam_view(self):
         if self.cam == 'top':
             return self.top
-        elif self.cam == 'angle':
-            return self.angle
+        elif self.cam == 'left_side':
+            return self.left_side
         elif self.cam == 'r_vis':
             return self.r_vis
         elif self.cam == 'l_vis':
@@ -226,9 +226,9 @@ class DualDianaMed(MujocoEnv):
             img_renderer.update_scene(self.mj_data,camera="top")
             self.top = img_renderer.render()
             self.top = self.top[:, :, ::-1]
-            img_renderer.update_scene(self.mj_data,camera="angle")
-            self.angle = img_renderer.render()
-            self.angle = self.angle[:, :, ::-1]
+            img_renderer.update_scene(self.mj_data,camera="left_side")
+            self.left_side = img_renderer.render()
+            self.left_side = self.left_side[:, :, ::-1]
             img_renderer.update_scene(self.mj_data,camera="front")
             self.front = img_renderer.render()
             self.front = self.front[:, :, ::-1]

roboimi/envs/double_pos_ctrl_env.py

@@ -34,19 +34,19 @@ class DualDianaMed_Pos_Ctrl(DualDianaMed):
             is_interpolate=is_interpolate,
             cam_view=cam_view
         )
-        
+
         self.max_reward = 4
-        
+
         self.cam_start()
 
-        
+
     def step(self,action=np.zeros(16)):
         action_left  = self.ik_solve(action[:3],action[3:7],self.arm_left)
         action_right = self.ik_solve(action[7:10],action[10:14],self.arm_right)
         action = np.hstack((action_left,action_right,action[14:]))
         super().step(action)
         self.rew = self._get_reward()
-        
+
     def step_jnt(self,action):
         super().step(action)
 
@@ -63,8 +63,8 @@ class DualDianaMed_Pos_Ctrl(DualDianaMed):
         return Arm.kdl_solver.ikSolver(p_goal, mat_goal, Arm.arm_qpos)
 
     def reset(self,box_pos):
-        
-        self.mj_data.joint('red_box_joint').qpos[0] = box_pos[0] 
+
+        self.mj_data.joint('red_box_joint').qpos[0] = box_pos[0]
         self.mj_data.joint('red_box_joint').qpos[1] = box_pos[1]
         self.mj_data.joint('red_box_joint').qpos[2] = box_pos[2]
         self.mj_data.joint('red_box_joint').qpos[3] = 1.0
@@ -73,22 +73,22 @@ 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.left_side = None
         self.r_vis = None
         self.front = None
         self.cam_flage = True
         t=0
         while self.cam_flage:
-            if(type(self.top)==type(None) 
-               or type(self.angle)==type(None) 
+            if(type(self.top)==type(None)
+               or type(self.left_side)==type(None)
                or type(self.r_vis)==type(None)
                or type(self.front)==type(None)):
                 time.sleep(0.001)
                 t+=1
             else:
                 self.cam_flage=False
-        
-        
+
+
 
     def preStep(self, action):
         if isinstance(action,np.ndarray) and len(action)==16:
@@ -101,7 +101,7 @@ class DualDianaMed_Pos_Ctrl(DualDianaMed):
         for i in range(3):
             box_pose[i] = cp.deepcopy(self.mj_data.joint('red_box_joint').qpos[i])
         return box_pose
-    
+
 
     def _get_reward(self):
         all_contact_pairs = []
@@ -124,16 +124,28 @@ class DualDianaMed_Pos_Ctrl(DualDianaMed):
         reward = 0
         if touch_right_gripper and not touch_table:
             reward = 1
-        if touch_right_gripper and not box_touch_table: 
+        if touch_right_gripper and not box_touch_table:
             reward = 2
         if touch_left_gripper: # attempted transfer
             reward = 3
         if touch_left_gripper and not box_touch_table: # successful transfer
             reward = 4
         return reward
-        
+
 
 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='left_side'
+        )
+        return env
     if 'sim_transfer' in task_name:
         from roboimi.assets.robots.diana_med import BiDianaMed
         env = DualDianaMed_Pos_Ctrl(
@@ -141,7 +153,7 @@ def make_sim_env(task_name, headless=False):
             is_render=not headless,
             control_freq=30,
             is_interpolate=True,
-            cam_view='angle'
+            cam_view='left_side'
         )
         return env
     else:
@@ -167,4 +179,4 @@ 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]
@@ -35,4 +36,23 @@ def sample_transfer_pose():
         box_position = np.random.uniform(ranges[:, 0], ranges[:, 1])
 
 
-        return box_position
+        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

@@ -23,6 +23,13 @@ SIM_TASK_CONFIGS = {
         'camera_names': ['top','r_vis','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="left_side",
+        )
+
+    def test_diana_table_scene_uses_left_side_camera_instead_of_angle(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.assertIn("left_side", cameras, "DianaMed scene should expose the left-side task camera")
+        left_side_pos = np.fromstring(cameras["left_side"]["pos"], sep=" ")
+        self.assertLess(float(left_side_pos[0]), 0.0)
+        self.assertEqual(cameras["left_side"].get("mode"), "targetbody")
+        self.assertEqual(cameras["left_side"].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):
-        self.reset_calls.append(np.array(box_pos))
+    def reset(self, task_state):
+        self.reset_calls.append(task_state)
 
     def _get_image_obs(self):
         self.image_obs_calls += 1
@@ -114,7 +114,7 @@ class EvalVLAHeadlessTest(unittest.TestCase):
             is_render=False,
             control_freq=30,
             is_interpolate=True,
-            cam_view="angle",
+            cam_view="left_side",
         )
 
     def test_camera_viewer_headless_updates_images_without_gui_calls(self):
@@ -123,11 +123,11 @@ class EvalVLAHeadlessTest(unittest.TestCase):
         env.mj_data = object()
         env.exit_flag = False
         env.is_render = False
-        env.cam = "angle"
+        env.cam = "left_side"
         env.r_vis = None
         env.l_vis = None
         env.top = None
-        env.angle = None
+        env.left_side = None
         env.front = None
 
         with mock.patch(
@@ -144,7 +144,7 @@ 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.left_side)
         self.assertIsNotNone(env.front)
 
     def test_eval_main_headless_skips_render_and_still_executes_policy(self):
@@ -254,6 +254,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()