docs: refine IMF migration spec defaults

docs: add IMF-AttnRes migration implementation plan
docs: add IMF-AttnRes migration design spec
2026-04-01 22:55:50 +08:00 · 2026-04-01 22:52:39 +08:00 · 2026-04-01 22:49:10 +08:00
23 changed files with 600 additions and 1924 deletions
@@ -1,72 +0,0 @@
 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
@@ -0,0 +1,268 @@
 # 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"
 ```
@@ -1,311 +0,0 @@
 # 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.
@@ -1,72 +0,0 @@
 # 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.
@@ -0,0 +1,272 @@
 # 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. 能按最近一次实验参数模板成功启动训练。
@@ -1,316 +0,0 @@
 # 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
@@ -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.01" euler="0.0 9.4 0.0 " fovy="50" resolution="1920 1200"/>
+                                                        <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"/>
                                                </body>
                                            </body>
                                            <body name="l_finger_left" pos="0 0.01 0.0444">
@@ -1,6 +0,0 @@
 <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>
@@ -1,19 +0,0 @@
 <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>
@@ -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="left_side" pos="-0.55 0.85 0.85" fovy="65" mode="targetbody" target="table"/>
+      <camera name="angle" pos="0.0 0.0 2.0" fovy="37" mode="targetbody" target="table"/>
      <camera name="front" pos="0 0 0.8" fovy="65" mode="fixed" quat="0.7071 0.7071 0 0"/>
  </worldbody>
 </mujoco>
@@ -91,39 +91,3 @@ class BiDianaMed(ArmBase):
        """ Robot's init joint position. """
        return np.array([0.0, 0.0, 0.0, 1.57, 0.0, 0.0, 0.0])
 class BiDianaMedSocketPeg(ArmBase):
    def __init__(self):
        super().__init__(
            name="Bidiana_socket_peg",
            urdf_path="roboimi/assets/models/manipulators/DianaMed/DualDianaMed.urdf",
            xml_path="roboimi/assets/models/manipulators/DianaMed/bi_diana_socket_peg_ee.xml",
            gripper=None
        )
        self.left_arm = self.Arm(self, 'single', self.urdf_path)
        self.left_arm.set_Arm_base_link('left_base_link')
        self.left_arm.set_Arm_ee_link('left_link7')
        self.left_arm.InitKDL
        self.left_arm.joint_index = ['l_j1','l_j2','l_j3','l_j4','l_j5','l_j6','l_j7']
        self.left_arm.gripper_index = ['l_finger_joint_left','r_finger_joint_left']
        self.left_arm.actuator_index = ['a1_l','a2_l','a3_l','a4_l','a5_l','a6_l','a7_l','gripper_left']
        self.left_arm.setArmInitPose(self.init_qpos)
        self.arms.append(self.left_arm)
        self.right_arm = self.Arm(self,'single', self.urdf_path)
        self.right_arm.set_Arm_base_link('right_base_link')
        self.right_arm.set_Arm_ee_link('right_link7')
        self.right_arm.InitKDL
        self.right_arm.joint_index = ['r_j1','r_j2','r_j3','r_j4','r_j5','r_j6','r_j7']
        self.right_arm.gripper_index = ['l_finger_joint_right','r_finger_joint_right']
        self.right_arm.actuator_index = ['a1_r','a2_r','a3_r','a4_r','a5_r','a6_r','a7_r','gripper_right']
        self.right_arm.setArmInitPose(self.init_qpos)
        self.arms.append(self.right_arm)
        self.jnt_num = self.left_arm.jnt_num + self.right_arm.jnt_num
        self.kp = 500 * np.ones(self.jnt_num)
        self.kd = 44.57 * np.ones(self.jnt_num)
    @property
    def init_qpos(self):
        """ Robot's init joint position. """
        return np.array([0.0, 0.0, 0.0, 1.57, 0.0, 0.0, 0.0])
@@ -1,203 +0,0 @@
 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,
            },
        ]
@@ -2,11 +2,9 @@ import time
 import os
 import numpy as np
 from roboimi.envs.double_pos_ctrl_env import make_sim_env
-from roboimi.demos.diana_air_insert_policy import TestAirInsertPolicy
+from diana_policy import TestPickAndTransferPolicy
 from roboimi.demos.diana_policy import TestPickAndTransferPolicy
 import cv2
-from roboimi.utils.act_ex_utils import sample_air_insert_socket_peg_state, sample_transfer_pose
+from roboimi.utils.act_ex_utils import sample_transfer_pose
 from roboimi.utils.constants import SIM_TASK_CONFIGS
 from roboimi.utils.streaming_episode_writer import StreamingEpisodeWriter
 import pathlib
@@ -14,34 +12,16 @@ HOME_PATH = str(pathlib.Path(__file__).parent.resolve())
 DATASET_DIR = HOME_PATH + '/dataset'
-def sample_task_state(task_name):
+def main():
-    if task_name == 'sim_transfer':
+    task_name = 'sim_transfer'
-        return sample_transfer_pose()
+    dataset_dir = DATASET_DIR + '/sim_transfer'   #SIM_TASK_CONFIGS[task_name]['dataset_dir']
-    if task_name == 'sim_air_insert_socket_peg':
+    num_episodes = 100       #SIM_TASK_CONFIGS[task_name]['num_episodes']
        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 = task_cfg['episode_len']
+    episode_len  = 700      #SIM_TASK_CONFIGS[task_name]['episode_len']
-    camera_names = task_cfg['camera_names']
+    camera_names = ['angle','r_vis', 'top', 'front']   #SIM_TASK_CONFIGS[task_name]['camera_names']
    image_size = (256, 256)
-    if task_name in {'sim_transfer', 'sim_air_insert_socket_peg'}:
+    if task_name == 'sim_transfer':
        print(task_name)
    else:
        raise NotImplementedError
@@ -49,7 +29,7 @@ def main(task_name='sim_transfer'):
    success = []
    env = make_sim_env(task_name)
-    policy = make_policy(task_name, inject_noise=inject_noise)
+    policy = TestPickAndTransferPolicy(inject_noise)
    # 等待osmesa完全启动后再开始收集数据
    print("等待osmesa线程启动...")
@@ -61,8 +41,8 @@ def main(task_name='sim_transfer'):
        max_reward = float('-inf')
        print(f'\n{episode_idx=}')
        print('Rollout out EE space scripted policy')
-        task_state = sample_task_state(task_name)
+        box_pos = sample_transfer_pose()
-        env.reset(task_state)
+        env.reset(box_pos)
        episode_writer = StreamingEpisodeWriter(
            dataset_path=os.path.join(dataset_dir, f'episode_{episode_idx}.hdf5'),
            max_timesteps=episode_len,
@@ -70,7 +50,7 @@ def main(task_name='sim_transfer'):
            image_size=image_size,
        )
        for step in range(episode_len):
-            raw_action = policy.predict(task_state, step)
+            raw_action = policy.predict(box_pos,step)
            env.step(raw_action)
            env.render()
            sum_reward += env.rew
@@ -26,10 +26,7 @@ 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 (
+from roboimi.utils.act_ex_utils import sample_transfer_pose
    sample_air_insert_socket_peg_state,
    sample_transfer_pose,
 )
 from roboimi.vla.eval_utils import execute_policy_action
 sys.path.append(os.getcwd())
@@ -488,14 +485,6 @@ 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 评估
@@ -560,8 +549,8 @@ def _run_eval(cfg: DictConfig):
            print(f"回合 {episode_idx + 1}/{eval_cfg.num_episodes}")
            print(f"{'='*60}\n")
-            task_state = _sample_task_reset_state(str(eval_cfg.task_name))
+            box_pos = sample_transfer_pose()
-            env.reset(task_state)
+            env.reset(box_pos)
            # 为新回合重置 agent 队列
            agent.reset()
@@ -1,154 +0,0 @@
 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())
@@ -52,7 +52,7 @@ class DualDianaMed(MujocoEnv):
        self.r_vis = None
        self.l_vis = None
        self.top = None
-        self.left_side = None
+        self.angle = 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']['left_side'] = self.left_side
+        obs['images']['angle'] = self.angle
        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']['left_side'] = self.left_side
+        obs['images']['angle'] = self.angle
        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 == 'left_side':
+        elif self.cam == 'angle':
-            return self.left_side
+            return self.angle
        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="left_side")
+            img_renderer.update_scene(self.mj_data,camera="angle")
-            self.left_side = img_renderer.render()
+            self.angle = img_renderer.render()
-            self.left_side = self.left_side[:, :, ::-1]
+            self.angle = self.angle[:, :, ::-1]
            img_renderer.update_scene(self.mj_data,camera="front")
            self.front = img_renderer.render()
            self.front = self.front[:, :, ::-1]
@@ -73,14 +73,14 @@ class DualDianaMed_Pos_Ctrl(DualDianaMed):
        self.mj_data.joint('red_box_joint').qpos[6] = 0.0
        super().reset()
        self.top = None
-        self.left_side = None
+        self.angle = 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.left_side)==type(None)
+               or type(self.angle)==type(None) 
               or type(self.r_vis)==type(None)
               or type(self.front)==type(None)):
                time.sleep(0.001)
@@ -134,18 +134,6 @@ class DualDianaMed_Pos_Ctrl(DualDianaMed):
 def make_sim_env(task_name, headless=False):
    if task_name == 'sim_air_insert_socket_peg':
        from roboimi.assets.robots.diana_med import BiDianaMedSocketPeg
        from roboimi.envs.double_air_insert_env import DualDianaMed_Air_Insert
        env = DualDianaMed_Air_Insert(
            robot=BiDianaMedSocketPeg(),
            is_render=not headless,
            control_freq=30,
            is_interpolate=True,
            cam_view='left_side'
        )
        return env
    if 'sim_transfer' in task_name:
        from roboimi.assets.robots.diana_med import BiDianaMed
        env = DualDianaMed_Pos_Ctrl(
@@ -153,7 +141,7 @@ def make_sim_env(task_name, headless=False):
            is_render=not headless,
            control_freq=30,
            is_interpolate=True,
-            cam_view='left_side'
+            cam_view='angle'
        )
        return env
    else:
@@ -1,6 +1,5 @@
 import numpy as np
 def sample_insertion_pose():
        # Peg
        x_range = [0.1, 0.2]
@@ -37,22 +36,3 @@ def sample_transfer_pose():
        return box_position
 def sample_air_insert_socket_peg_state():
        socket_position = np.random.uniform(
            low=np.array([-0.20, 0.80, 0.472], dtype=np.float32),
            high=np.array([-0.10, 1.00, 0.472], dtype=np.float32),
        )
        peg_position = np.random.uniform(
            low=np.array([0.10, 0.80, 0.46], dtype=np.float32),
            high=np.array([0.20, 1.00, 0.46], dtype=np.float32),
        )
        socket_quat = np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32)
        peg_quat = np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32)
        return {
            "socket_pos": socket_position.astype(np.float32, copy=False),
            "socket_quat": socket_quat,
            "peg_pos": peg_position.astype(np.float32, copy=False),
            "peg_quat": peg_quat,
        }
@@ -23,13 +23,6 @@ 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'
    },
 }
@@ -59,3 +52,13 @@ 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
@@ -1,510 +0,0 @@
 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()
@@ -36,8 +36,8 @@ class _FakeEnv:
        self.render_calls = 0
        self.reset_calls = []
-    def reset(self, task_state):
+    def reset(self, box_pos):
-        self.reset_calls.append(task_state)
+        self.reset_calls.append(np.array(box_pos))
    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="left_side",
+            cam_view="angle",
        )
    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 = "left_side"
+        env.cam = "angle"
        env.r_vis = None
        env.l_vis = None
        env.top = None
-        env.left_side = None
+        env.angle = 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.left_side)
+        self.assertIsNotNone(env.angle)
        self.assertIsNotNone(env.front)
    def test_eval_main_headless_skips_render_and_still_executes_policy(self):
@@ -254,69 +254,6 @@ 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()
@@ -4,7 +4,7 @@ import unittest
 from pathlib import Path
 from unittest import mock
-from roboimi.assets.robots import diana_med
+from roboimi.assets.robots.diana_med import BiDianaMed
 class _FakeKDL:
@@ -24,7 +24,6 @@ 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'
@@ -59,47 +58,6 @@ 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()
Author	SHA1	Message	Date
Logic	8d6060224a	docs: refine IMF migration spec defaults	2026-04-01 22:55:50 +08:00
Logic	8a8193fe7e	docs: add IMF-AttnRes migration implementation plan	2026-04-01 22:52:39 +08:00
Logic	1a92c5e8a6	docs: add IMF-AttnRes migration design spec	2026-04-01 22:49:10 +08:00