chore(eval): preserve parallel evaluation prototype

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

@@ -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.

docs/superpowers/plans/2026-05-08-act-socket-peg.md

@@ -1,49 +0,0 @@
-# ACT Socket Peg 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 a local ACT policy/model to RoboIMI and launch training on the socket peg dataset with three 224×224 camera views.
-
-**Architecture:** Implement a self-contained ACT agent and head that reuse the existing ResNet multiview backbone, dataset, training loop, normalization, and checkpointing. The ACT model uses a posterior transformer encoder for latent z and a transformer decoder with learned action queries for action chunks.
-
-**Tech Stack:** Python, PyTorch, Hydra/OmegaConf, unittest, HDF5 dataset via existing `SimpleRobotDataset`.
-
----
-
-## File Structure
-- Create `roboimi/vla/models/heads/act.py`: local ACT model/head implementation, no imports from external ACT repository.
-- Create `roboimi/vla/agent_act.py`: VLA-compatible ACT agent wrapper with normalization, condition building, loss, and inference queues.
-- Create `roboimi/vla/conf/agent/act_resnet.yaml`: Hydra agent config for three-camera ACT with 224×224 images.
-- Create `tests/test_act_agent.py`: model/agent unit tests.
-- Modify no external ACT code and do not add vendored ACT files.
-
-## Tasks
-
-### Task 1: Add ACT model/head tests
-- [ ] Write tests in `tests/test_act_agent.py` that define a lightweight fake vision backbone emitting deterministic camera tokens.
-- [ ] Test `ACTAgent.compute_loss()` returns a scalar tensor and backpropagates through the head.
-- [ ] Test masked L1 ignores padded timesteps by comparing all-padded vs partially valid batches for finite loss behavior.
-- [ ] Test `ACTAgent.predict_action()` returns `(B,pred_horizon,action_dim)`.
-- [ ] Run `python -m unittest tests.test_act_agent -v` and confirm tests fail because `roboimi.vla.agent_act` does not exist.
-
-### Task 2: Implement local ACT head and ACT agent
-- [ ] Create `roboimi/vla/models/heads/act.py` with `ACTPolicyHead`, sinusoidal table helper, KL helper, and transformer layers using `batch_first=True` PyTorch modules.
-- [ ] Create `roboimi/vla/agent_act.py` with `ACTAgent` implementing existing training/inference API.
-- [ ] Reuse `NormalizationModule` and camera ordering checks from `VLAAgent` behavior.
-- [ ] Run `python -m unittest tests.test_act_agent -v` and fix until green.
-
-### Task 3: Add Hydra config and wiring tests
-- [ ] Add `roboimi/vla/conf/agent/act_resnet.yaml` using existing `resnet_diffusion` backbone with `output_tokens_per_camera=true` and `camera_names=${data.camera_names}`.
-- [ ] Extend `tests/test_act_agent.py` with a Hydra compose/instantiate test using reduced backbone/head sizes and `data.camera_names='[l_vis,r_vis,front]'`.
-- [ ] Run `python -m unittest tests.test_act_agent -v` and `python -m unittest tests.test_resnet_transformer_agent_wiring -v`.
-
-### Task 4: Verify socket peg data path and training smoke test
-- [ ] Run a dataset sample check against `/data/roboimi_datasets/sim_air_insert_socket_peg` with `camera_names=[l_vis,r_vis,front]` and `image_resize_shape=[224,224]`.
-- [ ] Run a short CPU or GPU training smoke test with `agent=act_resnet`, `train.max_steps=2`, `train.num_workers=0`, pretrained backbone disabled, and reduced head sizes if needed.
-- [ ] Record exact command and output snippet.
-
-### Task 5: Launch real ACT socket peg training
-- [ ] Create a run directory under `runs/` with timestamped name.
-- [ ] Start training using `/home/droid/.conda/envs/roboimi/bin/python roboimi/demos/vla_scripts/train_vla.py agent=act_resnet data.dataset_dir=/data/roboimi_datasets/sim_air_insert_socket_peg data.camera_names='[l_vis,r_vis,front]' data.image_resize_shape='[224,224]'` plus selected training hyperparameters.
-- [ ] Redirect output to `train.log` and store PID in `train.pid`.
-- [ ] Tail log to verify dataset loads, agent initializes, and first loss is produced.

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

@@ -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

docs/superpowers/specs/2026-05-08-act-socket-peg-design.md

@@ -1,78 +0,0 @@
-# ACT Socket Peg Policy Design
-
-## Goal
-Add a local ACT-style policy/model to the existing VLA training stack and start training it on `/data/roboimi_datasets/sim_air_insert_socket_peg` using three 224×224 camera views.
-
-## Constraints
-- Base work is on branch `feat-act-socket-peg`, created from current `main` (`acbd7c605a8d203a774f2f47cff8094c05d9325e`).
-- Do not vendor or import ACT repository environment, dataset, training, or utility code.
-- Reimplement only the model/policy logic needed for this repo: CVAE action encoder, learned action queries, transformer conditioning, KL + L1 loss, and inference from prior.
-- Keep existing dataset/training loop style and checkpoint format.
-
-## Data
-The socket peg dataset is HDF5 under `/data/roboimi_datasets/sim_air_insert_socket_peg`. Episodes contain:
-- `action`: `(600, 16)`, `float32`
-- `observations/qpos`: `(600, 16)`, `float32`
-- `observations/images/l_vis`, `r_vis`, `front`: `(600, 256, 256, 3)`, `uint8`
-- attrs include `camera_names=[l_vis,r_vis,front]`, `image_height=256`, `image_width=256`, `sim=True`.
-
-Training config must use `data.camera_names='[l_vis,r_vis,front]'` and `data.image_resize_shape=[224,224]`. Existing dataset code already resizes HWC uint8 frames to `(C,224,224)` float tensors in `[0,1]`.
-
-## Architecture
-Create `roboimi/vla/agent_act.py` with `ACTAgent`, an `nn.Module` that follows the existing agent API:
-- `compute_loss(batch)` accepts `images`, `qpos`, `action`, `action_is_pad`.
-- `predict_action(images, proprioception)` returns denormalized `(B,pred_horizon,action_dim)` chunks.
-- `predict_action_chunk`, `select_action`, and queue handling mirror the existing inference contract.
-- `get_normalization_stats()` returns the current normalization module stats.
-
-Create `roboimi/vla/models/heads/act.py` containing focused, local ACT model classes:
-- Sinusoidal positional table helper.
-- Transformer encoder for posterior `z` from `[CLS, qpos, action sequence]` with padding mask.
-- Transformer decoder/action-query module conditioned on visual tokens, current qpos, and latent token.
-- `ACTPolicyHead` returning action predictions and latent `(mu, logvar)`.
-
-To avoid copying ACT's DETR image backbone code, reuse this repo's `ResNetDiffusionBackbone`. Configure it with `output_tokens_per_camera=true`, so each observation step emits one token per camera. The ACT agent builds memory tokens by concatenating camera visual tokens, a qpos token, and a latent token.
-
-## Loss
-During training:
-1. Normalize qpos/action with existing `NormalizationModule`.
-2. Keep only `num_queries == pred_horizon` actions.
-3. Encode posterior `z` from normalized current qpos and normalized target action sequence.
-4. Predict action chunk from image/qpos/latent tokens.
-5. Compute masked L1 over non-padded action timesteps.
-6. Add `kl_weight * KL(N(mu,sigma), N(0,I))`.
-
-Inference uses zero latent prior and denormalizes predicted actions.
-
-## Config
-Add `roboimi/vla/conf/agent/act_resnet.yaml`:
-- `_target_: roboimi.vla.agent_act.ACTAgent`
-- `action_dim=16`, `obs_dim=16`
-- `pred_horizon=16`, `obs_horizon=1` by default, `num_action_steps=8`
-- `camera_names: ${data.camera_names}`, `num_cams: 3`
-- ResNet backbone with `input_shape=[3,224,224]`, `output_tokens_per_camera=true`, three cameras.
-- ACT head hyperparameters small enough for a training smoke test and usable as defaults: `hidden_dim=256`, `nheads=8`, `enc_layers=4`, `dec_layers=6`, `dim_feedforward=2048`, `latent_dim=32`, `dropout=0.1`, `kl_weight=10.0`.
-
-Training command should override dataset path and camera names:
-```bash
-/home/droid/.conda/envs/roboimi/bin/python roboimi/demos/vla_scripts/train_vla.py \
-  agent=act_resnet \
-  data.dataset_dir=/data/roboimi_datasets/sim_air_insert_socket_peg \
-  data.camera_names='[l_vis,r_vis,front]' \
-  data.image_resize_shape='[224,224]' \
-  train.device=cuda \
-  train.num_workers=8 \
-  train.batch_size=32 \
-  train.max_steps=100000 \
-  train.use_swanlab=true \
-  train.swanlab_project=roboimi-vla \
-  train.swanlab_run_name=act-socket-peg-224-$(date +%Y%m%d-%H%M%S)
-```
-
-## Tests
-Add unit coverage without requiring ACT repo code:
-- Hydra config can instantiate `agent=act_resnet` with stubbed torchvision/diffusers-like dependencies where needed.
-- ACT loss returns a scalar, masks padded actions, and produces gradients.
-- ACT prediction returns `(B,pred_horizon,action_dim)` and honors configured camera ordering.
-- Dataset config/sample for socket peg cameras returns three resized `(obs_horizon,3,224,224)` tensors.
-

render.py

@@ -1,200 +0,0 @@
-import mujoco
-from mujoco import viewer
-import sys
-import numpy as np
-import time
-import threading
-
-
-class MjBasicRenderer:
-    def __new__(cls, *args, **kwargs):
-        return super().__new__(cls)
-
-    def __init__(self, mj_model=None, mj_data=None):
-        # keyboard flag
-        self.render_paused = True
-        self.exit_flag = False
-        # init param
-        self.mj_model = mj_model
-        self.mj_data = mj_data
-        self.renderer = "viewer"  # default
-        self.viewer = None
-        self._image = None
-
-        # Set up mujoco viewer
-        self.image_renderer = mujoco.Renderer(self.mj_model)
-
-    def __del__(self):
-        pass
-
-    def _init_renderer(self):
-        """Initialize renderer, choose official renderer with "viewer"(joined from version 2.3.3),
-        another renderer with "mujoco_viewer"
-        """
-
-        def key_callback(keycode):
-            if keycode == 32:  # space
-                self.render_paused = not self.render_paused
-            elif keycode == 256:  # escape
-                self.exit_flag = not self.exit_flag
-
-        if self.renderer == "viewer":
-            # This function does not block, allowing user code to continue execution.
-            self.viewer = viewer.launch_passive(
-                self.mj_model,
-                self.mj_data,
-                key_callback=key_callback,
-                show_left_ui=False,
-                show_right_ui=False,
-            )
-            self.set_renderer_config()
-        else:
-            raise ValueError("Invalid renderer for some reason.")
-
-    def render(self):
-        """mujoco render"""
-        if self.viewer is not None and self.render_paused is True:
-            if self.viewer.is_running() and self.exit_flag is False:
-                self.viewer: viewer.Handle
-                self.viewer.sync()
-            else:
-                self.viewer.close()
-
-    def set_renderer_config(self):
-        """Setup mujoco global config while using viewer as renderer.
-        It should be noted that the render thread need locked.
-        """
-        self.viewer.cam.lookat = np.array([0.4, 0, 0.5])
-        self.viewer.cam.azimuth -= 0.005
-        with self.viewer.lock():
-            self.viewer.opt.flags[mujoco.mjtVisFlag.mjVIS_CONTACTPOINT] = int(
-                self.mj_data.time % 2
-            )
-
-
-try:
-    import cv2
-except ImportError:
-    print("Could not import cv2, please install it to enable camera viewer.")
-
-
-class MjMultiRenderer(MjBasicRenderer):
-
-    # __slots__=('mj_model','mj_data','renderer','enable_camera_viewer')
-    def __new__(cls, *args, **kwargs):
-        return super().__new__(cls)
-
-    def __init__(
-        self,
-        mj_model=None,
-        mj_data=None,
-        renderer=None,
-        enable_camera_viewer=False,
-        enable_depth=False,
-    ):
-        super().__init__(mj_model, mj_data)
-        self._depth = None
-        self.renderer = renderer
-        self._init_renderer()
-        self.enable_camera_viewer = enable_camera_viewer
-        if self.enable_camera_viewer:
-            self.enable_depth = enable_depth
-            self._init_window()
-        else:
-            self.enable_depth = False
-            print("No Camera View")
-
-    def __del__(self):
-        self.close()
-
-    def _init_renderer(self):
-        """
-        Initialize renderer, choose official renderer with "viewer"(joined from version 2.3.3)
-        """
-        if self.renderer == "unity":
-            # TODO: Support unity renderer.
-            raise ValueError("Unity renderer init failed for no supporting reason")
-        elif self.renderer == "viewer":
-            super()._init_renderer()
-            print("mujoco viewer init !")
-        else:
-            raise ValueError("renderer init failed for some reason.")
-
-    def _init_window(self, name="Camera view"):
-        if not self.enable_depth:
-            cv2.namedWindow(name, cv2.WINDOW_NORMAL)
-        else:
-            cv2.namedWindow(name, cv2.WINDOW_NORMAL)
-            cv2.namedWindow("Camera depth view", cv2.WINDOW_NORMAL)
-
-    def render(self):
-        """render mujoco"""
-        if self.renderer == "viewer":
-            super().render()
-        elif self.renderer == "unity":
-            # TODO: Support unity renderer.
-            raise ValueError("Unity renderer not supported now.")
-        else:
-            raise ValueError("Invalid renderer for some reason.")
-
-    def render(self):
-        """mujoco render"""
-        if self.viewer is not None and self.render_paused is True:
-            if self.viewer.is_running() and self.exit_flag is False:
-                self.viewer: viewer.Handle
-                self.viewer.sync()
-            else:
-                self.viewer.close()
-
-    def camera_render(self, cam=None):
-        if self.enable_camera_viewer:
-            if not self.enable_depth:
-                rgb, depth = self.render_from_camera(cam)
-                rgb = cv2.resize(rgb, (1920, 1600))
-                cv2.imshow("Camera view", rgb)
-                cv2.waitKey(1)
-
-            else:
-                rgb, depth = self.render_from_camera(cam)
-                cv2.imshow("Camera view", rgb)
-                cv2.imshow("Camera depth view", depth)
-                cv2.waitKey(1)
-
-        else:
-            print("camera info disable")
-            return
-
-    def render_from_camera(self, cam=None):
-        self.image_renderer.update_scene(self.mj_data, camera=cam)
-        if self.enable_depth is True:
-            self.image_renderer.enable_depth_rendering()
-            org = self.image_renderer.render()
-            depth = org[:, :]
-            self.image_renderer.disable_depth_rendering()
-            org = self.image_renderer.render()
-            image = org[:, :, ::-1]
-        else:
-            org = self.image_renderer.render()
-            image = org[:, :, ::-1]
-            depth = np.zeros([240, 320])
-        return image, depth
-
-    def close(self):
-        """close the environment."""
-        if self.enable_camera_viewer and self.viewer.is_running() == False:
-            cv2.destroyAllWindows()
-        self.viewer.close()
-        # sys.exit(0)
-
-    # def get_cam_intrinsic(self, fovy=45.0, width=320, height=240):
-    #     aspect = width * 1.0 / height
-    #     fovx = np.degrees(2 * np.arctan(aspect * np.tan(np.radians(fovy / 2))))
-
-    #     cx = 0.5 * width
-    #     cy = 0.5 * height
-    #     fx = cx / np.tan(fovx * np.pi / 180 * 0.5)
-    #     fy = cy / np.tan(fovy * np.pi / 180 * 0.5)
-
-    #     K = np.array([[fx, 0, cx],
-    #                   [0, fy, cy],
-    #                   [0, 0, 1]], dtype=np.float32)

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

@@ -76,7 +76,7 @@
                                                 <body name="ee_cam_left" pos="0.00 0.046 -0.075" euler="0.0 0.0 -0.0">
                                                         <inertial pos="0 0 0" quat="1 0 0 0" mass="0" diaginertia="0 0 0" />
                                                         <geom type="mesh" contype="1" conaffinity="1" group="1" rgba="0.69804 0.69804 0.69804 1" mesh="realsense_cam" />
-                                                        <camera name="rs_cam_left" mode="fixed" pos="0.0 0.0 0.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">

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

@@ -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>

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

@@ -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>

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

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

roboimi/assets/robots/diana_med.py

@@ -91,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])
-    

roboimi/demos/diana_air_insert_policy.py

@@ -1,184 +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_END_T = 580
-    LEFT_SOCKET_GRIPPER_CLOSED = -100
-    RIGHT_PEG_GRIPPER_CLOSED = -100
-    SOCKET_APPROACH_Z = 1.05
-    EPISODE_END_T = 600
-
-    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": 350,
-                "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": 350,
-                "xyz": peg_init_xyz,
-                "quat": right_pick_quat,
-                "gripper": self.RIGHT_PEG_GRIPPER_CLOSED,
-            },
-            {
-                "t": 450,
-                "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": self.EPISODE_END_T,
-                "xyz": peg_insert_end_center,
-                "quat": right_pick_quat,
-                "gripper": self.RIGHT_PEG_GRIPPER_CLOSED,
-            },
-        ]

roboimi/demos/diana_record_sim_episodes.py

@@ -2,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

roboimi/demos/vla_scripts/train_vla.py

@@ -702,28 +702,10 @@ def _run_training(cfg: DictConfig):
             from roboimi.demos.vla_scripts import eval_vla
 
             rollout_cfg = OmegaConf.create(OmegaConf.to_container(cfg, resolve=False))
-            rollout_num_episodes = int(cfg.train.get('rollout_num_episodes', 1))
-            rollout_device = str(cfg.train.get('rollout_device', cfg.train.device))
-            configured_rollout_workers = cfg.train.get('rollout_num_workers', None)
-            if configured_rollout_workers is None:
-                if rollout_device.startswith('cuda'):
-                    rollout_num_workers = min(max(rollout_num_episodes, 1), 8)
-                else:
-                    rollout_num_workers = 1
-            else:
-                rollout_num_workers = int(configured_rollout_workers)
             rollout_cfg.eval.ckpt_path = str(checkpoint_path)
-            rollout_cfg.eval.num_episodes = rollout_num_episodes
+            rollout_cfg.eval.num_episodes = int(cfg.train.get('rollout_num_episodes', 1))
-            rollout_cfg.eval.num_workers = rollout_num_workers
             rollout_cfg.eval.headless = True
-            rollout_cfg.eval.device = rollout_device
+            rollout_cfg.eval.device = 'cpu'
-            rollout_cfg.eval.cuda_devices = cfg.train.get('rollout_cuda_devices', None)
-            rollout_cfg.eval.response_timeout_s = float(
-                cfg.train.get('rollout_response_timeout_s', 300.0)
-            )
-            rollout_cfg.eval.server_startup_timeout_s = float(
-                cfg.train.get('rollout_server_startup_timeout_s', 300.0)
-            )
             rollout_cfg.eval.verbose_action = False
             rollout_cfg.eval.record_video = False
             rollout_cfg.eval.save_trajectory_image = True
@@ -734,11 +716,9 @@ def _run_training(cfg: DictConfig):
             )
 
             log.info(
-                "🎯 开始 checkpoint rollout 验证: %s (episodes=%s, device=%s, workers=%s, headless=True)",
+                "🎯 开始 checkpoint rollout 验证: %s (episodes=%s, headless=True)",
                 checkpoint_path,
                 rollout_cfg.eval.num_episodes,
-                rollout_cfg.eval.device,
-                rollout_cfg.eval.num_workers,
             )
             return eval_vla._run_eval(rollout_cfg)
 

roboimi/envs/double_air_insert_env.py

@@ -1,157 +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.r_vis = None
-        self.l_vis = None
-        self.front = None
-        if not self.is_render:
-            self._update_camera_images_sync()
-            return
-        self.cam_flage = True
-        while self.cam_flage:
-            if (
-                type(self.top) == type(None)
-                or type(self.r_vis) == type(None)
-                or type(self.l_vis) == type(None)
-                or type(self.front) == type(None)
-            ):
-                time.sleep(0.001)
-            else:
-                self.cam_flage = False
-
-    def step(self, action=np.zeros(16)):
-        super().step(action)
-        self.rew = self._get_reward()
-        self.obs = self._get_obs()
-
-    def get_env_state(self):
-        return get_socket_peg_env_state(self.mj_data)
-
-    def _get_reward(self):
-        contact_pairs = []
-        for collision_num in range(self.mj_data.ncon):
-            geom1 = self.mj_data.contact[collision_num].geom1
-            geom2 = self.mj_data.contact[collision_num].geom2
-            contact_pairs.append(
-                (self.getID2Name("geom", geom1), self.getID2Name("geom", geom2))
-            )
-        return compute_air_insert_reward(contact_pairs, self.get_env_state())

roboimi/envs/double_base.py

@@ -52,6 +52,7 @@ class DualDianaMed(MujocoEnv):
         self.r_vis = None
         self.l_vis = None
         self.top = None
+        self.angle = None
         self.front = None
         self.obs = None
         
@@ -91,6 +92,7 @@ class DualDianaMed(MujocoEnv):
 
     def step(self,action):
         self.compute_qpos = action  #for observation !
+        self.obs = self._get_obs()
         if self.interpolator_left is not None and self.interpolator_right is not None:
             self.interpolator_left.updateInput(action[:7], control_cycle=self.base_time)
             self.interpolator_right.updateInput(action[7:-2], control_cycle=self.base_time)
@@ -103,7 +105,6 @@ class DualDianaMed(MujocoEnv):
 
             super().step(action)
         self.base_time = time.time() - ctrl_cur_time
-        self.obs = self._get_obs()
             
 
     def preStep(self, action):
@@ -167,9 +168,10 @@ class DualDianaMed(MujocoEnv):
         obs['qpos'] = self.get_obs_qpos
         obs['action'] = self.compute_qpos
         obs['images'] = dict()
+        obs['images']['top']   = self.top
+        obs['images']['angle'] = self.angle
         obs['images']['r_vis'] = self.r_vis
         obs['images']['l_vis'] = self.l_vis
-        obs['images']['top'] = self.top
         obs['images']['front'] = self.front
         return obs
 
@@ -178,9 +180,10 @@ class DualDianaMed(MujocoEnv):
             self._update_camera_images_sync()
         obs = collections.OrderedDict()
         obs['images'] = dict()
+        obs['images']['top']   = self.top
+        obs['images']['angle'] = self.angle
         obs['images']['r_vis'] = self.r_vis
         obs['images']['l_vis'] = self.l_vis
-        obs['images']['top'] = self.top
         obs['images']['front'] = self.front
         return obs
     
@@ -199,12 +202,14 @@ class DualDianaMed(MujocoEnv):
 
     @property
     def cam_view(self):
-        if self.cam == 'r_vis':
+        if self.cam == 'top':
+            return self.top
+        elif self.cam == 'angle':
+            return self.angle
+        elif self.cam == 'r_vis':
             return self.r_vis
         elif self.cam == 'l_vis':
             return self.l_vis
-        elif self.cam == 'top':
-            return self.top
         elif self.cam == 'front':
             return self.front
         else:
@@ -225,6 +230,8 @@ class DualDianaMed(MujocoEnv):
         self.l_vis = img_renderer.render()[:, :, ::-1]
         img_renderer.update_scene(self.mj_data, camera="top")
         self.top = img_renderer.render()[:, :, ::-1]
+        img_renderer.update_scene(self.mj_data, camera="angle")
+        self.angle = img_renderer.render()[:, :, ::-1]
         img_renderer.update_scene(self.mj_data, camera="front")
         self.front = img_renderer.render()[:, :, ::-1]
 

roboimi/envs/double_pos_ctrl_env.py

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

roboimi/utils/act_ex_utils.py

@@ -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,
-        }

roboimi/utils/constants.py

@@ -20,14 +20,7 @@ SIM_TASK_CONFIGS = {
         'dataset_dir': DATASET_DIR + '/sim_transfer',
         'num_episodes': 20,
         'episode_len': 700,
-        'camera_names': ['r_vis', 'top', 'front'],
+        '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

roboimi/vla/agent_act.py

@@ -1,214 +0,0 @@
-"""ACT agent wrapper compatible with RoboIMI VLA training/eval scripts."""
-
-from __future__ import annotations
-
-from collections import deque
-from typing import Dict, Optional, Tuple
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from roboimi.vla.models.normalization import NormalizationModule
-
-
-class ACTAgent(nn.Module):
-    def __init__(
-        self,
-        vision_backbone,
-        head,
-        action_dim: int,
-        obs_dim: int,
-        pred_horizon: int = 16,
-        obs_horizon: int = 1,
-        num_cams: int = 3,
-        camera_names: Optional[Tuple[str, ...]] = None,
-        dataset_stats=None,
-        normalization_type: str = "min_max",
-        num_action_steps: int = 8,
-        **_: object,
-    ) -> None:
-        super().__init__()
-        self.action_dim = int(action_dim)
-        self.obs_dim = int(obs_dim)
-        self.pred_horizon = int(pred_horizon)
-        self.obs_horizon = int(obs_horizon)
-        self.num_cams = int(num_cams)
-        self.num_action_steps = int(num_action_steps)
-        self.vision_encoder = vision_backbone
-        self.normalization = NormalizationModule(
-            stats=dataset_stats,
-            normalization_type=normalization_type,
-        )
-
-        agent_camera_names = tuple(camera_names) if camera_names is not None else None
-        backbone_camera_names = getattr(self.vision_encoder, "camera_names", None)
-        backbone_camera_names = tuple(backbone_camera_names) if backbone_camera_names is not None else None
-        backbone_num_cameras = getattr(self.vision_encoder, "num_cameras", None)
-        if backbone_num_cameras is not None and int(backbone_num_cameras) != self.num_cams:
-            raise ValueError(
-                f"agent.num_cams({self.num_cams}) 与 vision_backbone.num_cameras({backbone_num_cameras}) 不一致"
-            )
-        if agent_camera_names is not None and backbone_camera_names is not None and agent_camera_names != backbone_camera_names:
-            raise ValueError(
-                f"agent.camera_names({list(agent_camera_names)}) 与 vision_backbone.camera_names({list(backbone_camera_names)}) 不一致"
-            )
-        self.camera_names = agent_camera_names if agent_camera_names is not None else backbone_camera_names
-        if self.camera_names is not None and len(self.camera_names) != self.num_cams:
-            raise ValueError(f"camera_names 长度({len(self.camera_names)})与 num_cams({self.num_cams})不一致")
-        if self.camera_names is not None:
-            self.vision_encoder.camera_names = self.camera_names
-
-        self.tokens_per_step = int(getattr(self.vision_encoder, "tokens_per_step", 1))
-        base_vision_dim = int(getattr(self.vision_encoder, "output_dim"))
-        self.vision_dim = base_vision_dim if self.tokens_per_step > 1 else base_vision_dim * self.num_cams
-        if isinstance(head, nn.Module):
-            self.policy_head = head
-        else:
-            self.policy_head = head(
-                action_dim=self.action_dim,
-                obs_dim=self.obs_dim,
-                vision_dim=self.vision_dim,
-                num_cams=self.num_cams,
-                pred_horizon=self.pred_horizon,
-                obs_horizon=self.obs_horizon,
-            )
-        # Alias so train_vla.py optimizer grouping can find head groups if added later.
-        self.noise_pred_net = self.policy_head
-        self.reset()
-
-    def _get_model_device(self) -> torch.device:
-        return next(self.parameters()).device
-
-    def _move_to_device(self, data, device: torch.device):
-        if torch.is_tensor(data):
-            return data.to(device)
-        if isinstance(data, dict):
-            return {k: self._move_to_device(v, device) for k, v in data.items()}
-        if isinstance(data, list):
-            return [self._move_to_device(v, device) for v in data]
-        if isinstance(data, tuple):
-            return tuple(self._move_to_device(v, device) for v in data)
-        return data
-
-    def _order_images(self, images: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
-        if self.camera_names is None:
-            camera_names = tuple(sorted(images.keys()))
-            if len(camera_names) != self.num_cams:
-                raise ValueError(f"图像条件相机数量({len(camera_names)})与 num_cams({self.num_cams})不一致")
-            return {cam_name: images[cam_name] for cam_name in camera_names}
-        missing = [cam_name for cam_name in self.camera_names if cam_name not in images]
-        if missing:
-            raise ValueError(f"图像条件缺少必需相机。missing={missing}, expected={list(self.camera_names)}")
-        return {cam_name: images[cam_name] for cam_name in self.camera_names}
-
-    def _build_visual_tokens(self, images: Dict[str, torch.Tensor]) -> torch.Tensor:
-        ordered_images = self._order_images(images)
-        visual = self.vision_encoder(ordered_images)
-        if visual.ndim == 3:
-            if visual.shape[1] < 1:
-                raise RuntimeError("视觉特征时间维为空")
-            return visual[:, -1:, :]
-        if visual.ndim == 4:
-            if visual.shape[1] < 1:
-                raise RuntimeError("视觉特征时间维为空")
-            return visual[:, -1, :, :]
-        raise RuntimeError(f"不支持的视觉特征形状: {tuple(visual.shape)}")
-
-    @staticmethod
-    def _current_qpos(states: torch.Tensor) -> torch.Tensor:
-        if states.ndim == 2:
-            return states
-        if states.ndim == 3:
-            return states[:, -1]
-        raise ValueError(f"qpos must have shape (B,D) or (B,T,D), got {tuple(states.shape)}")
-
-    def compute_loss(self, batch) -> torch.Tensor:
-        actions = batch["action"]
-        states = batch["qpos"]
-        images = batch["images"]
-        action_is_pad = batch.get("action_is_pad", None)
-        states = self.normalization.normalize_qpos(states)
-        actions = self.normalization.normalize_action(actions)
-        qpos = self._current_qpos(states)
-        actions = actions[:, : self.pred_horizon]
-        if action_is_pad is not None:
-            action_is_pad = action_is_pad[:, : self.pred_horizon].to(torch.bool)
-        visual_tokens = self._build_visual_tokens(images)
-        pred_actions, latent_info = self.policy_head(qpos, visual_tokens, actions, action_is_pad)
-        l1 = F.l1_loss(pred_actions, actions, reduction="none")
-        if action_is_pad is not None:
-            mask = (~action_is_pad).unsqueeze(-1).to(l1.dtype)
-            valid_count = (mask.sum() * l1.shape[-1]).clamp_min(1.0)
-            l1_loss = (l1 * mask).sum() / valid_count
-        else:
-            l1_loss = l1.mean()
-        kl = latent_info.get("kl")
-        if kl is None:
-            kl = torch.zeros(1, device=l1_loss.device, dtype=l1_loss.dtype)
-        kl_weight = float(latent_info.get("kl_weight", getattr(self.policy_head, "kl_weight", 0.0)))
-        return l1_loss + kl_weight * kl.squeeze()
-
-    @torch.no_grad()
-    def predict_action(self, images, proprioception):
-        proprioception = self.normalization.normalize_qpos(proprioception)
-        qpos = self._current_qpos(proprioception)
-        visual_tokens = self._build_visual_tokens(images)
-        actions, _ = self.policy_head(qpos, visual_tokens)
-        return self.normalization.denormalize_action(actions)
-
-    def reset(self):
-        self._queues = {
-            "qpos": deque(maxlen=self.obs_horizon),
-            "images": deque(maxlen=self.obs_horizon),
-            "action": deque(maxlen=max(1, self.pred_horizon - self.obs_horizon + 1)),
-        }
-
-    def _populate_queues(self, observation: Dict[str, torch.Tensor]) -> None:
-        if "qpos" in observation:
-            self._queues["qpos"].append(observation["qpos"].clone())
-        if "images" in observation:
-            ordered_images = self._order_images(observation["images"])
-            self._queues["images"].append({k: v.clone() for k, v in ordered_images.items()})
-
-    def _prepare_observation_batch(self) -> Dict[str, torch.Tensor]:
-        qpos_list = list(self._queues["qpos"])
-        if not qpos_list:
-            raise ValueError("观测队列为空，请先调用 _populate_queues 添加观测")
-        while len(qpos_list) < self.obs_horizon:
-            qpos_list.append(qpos_list[-1])
-        batch_qpos = torch.stack(qpos_list, dim=0).unsqueeze(0)
-
-        images_list = list(self._queues["images"])
-        if not images_list:
-            raise ValueError("图像队列为空，请先调用 _populate_queues 添加观测")
-        while len(images_list) < self.obs_horizon:
-            images_list.append(images_list[-1])
-        camera_names = self.camera_names if self.camera_names is not None else tuple(sorted(images_list[0].keys()))
-        batch_images = {
-            cam_name: torch.stack([img[cam_name] for img in images_list], dim=0).unsqueeze(0)
-            for cam_name in camera_names
-        }
-        return {"qpos": batch_qpos, "images": batch_images}
-
-    @torch.no_grad()
-    def select_action(self, observation: Dict[str, torch.Tensor]) -> torch.Tensor:
-        device = self._get_model_device()
-        observation = self._move_to_device(observation, device)
-        self._populate_queues(observation)
-        if len(self._queues["action"]) == 0:
-            batch = self._prepare_observation_batch()
-            actions = self.predict_action_chunk(batch)
-            start = self.obs_horizon - 1
-            end = start + self.num_action_steps
-            executable_actions = actions[:, start:end]
-            for i in range(executable_actions.shape[1]):
-                self._queues["action"].append(executable_actions[:, i].squeeze(0))
-        return self._queues["action"].popleft()
-
-    @torch.no_grad()
-    def predict_action_chunk(self, batch: Dict[str, torch.Tensor]) -> torch.Tensor:
-        return self.predict_action(batch["images"], batch["qpos"])
-
-    def get_normalization_stats(self):
-        return self.normalization.get_stats()

roboimi/vla/conf/agent/act_resnet.yaml

@@ -1,36 +0,0 @@
-# @package agent
-defaults:
-  - /backbone@vision_backbone: resnet_diffusion
-  - _self_
-
-_target_: roboimi.vla.agent_act.ACTAgent
-
-action_dim: 16
-obs_dim: 16
-normalization_type: "min_max"
-
-pred_horizon: 16
-obs_horizon: 1
-num_action_steps: 8
-
-camera_names: ${data.camera_names}
-num_cams: 3
-
-vision_backbone:
-  num_cameras: ${agent.num_cams}
-  camera_names: ${agent.camera_names}
-  input_shape: [3, 224, 224]
-  output_tokens_per_camera: true
-
-head:
-  _target_: roboimi.vla.models.heads.act.ACTPolicyHead
-  _partial_: true
-  hidden_dim: 256
-  nheads: 8
-  enc_layers: 4
-  dec_layers: 6
-  dim_feedforward: 2048
-  latent_dim: 32
-  dropout: 0.1
-  kl_weight: 10.0
-  activation: "gelu"

roboimi/vla/conf/config.yaml

@@ -29,11 +29,6 @@ train:
   rollout_val_freq_epochs: 50  # 每隔多少个 epoch 执行一次 rollout 验证
   rollout_validate_on_checkpoint: false  # 是否在保存 checkpoint 后立即运行 rollout 验证
   rollout_num_episodes: 3  # rollout 验证的回合数
-  rollout_device: ${train.device}  # rollout 使用的设备；默认跟随训练设备
-  rollout_num_workers: null  # rollout 并行 worker 数；null 时 CUDA 自动推断，CPU 保持 1
-  rollout_cuda_devices: null  # rollout CUDA 并行使用的逻辑 device 列表；null 时默认 [0]
-  rollout_response_timeout_s: 300.0  # rollout worker 等待 inference server 响应的超时时间
-  rollout_server_startup_timeout_s: 300.0  # rollout 等待 inference server 就绪的超时时间
 
   # 学习率调度器（带预热）
   warmup_steps: 2000  # 预热步数（Transformer建议更长）

roboimi/vla/conf/eval/eval.yaml

@@ -2,10 +2,6 @@
 # 评估配置
 ckpt_path: "checkpoints/vla_model_best.pt"  # 模型检查点路径
 num_episodes: 3           # 评估回合数
-num_workers: 1            # 并行 worker 数；1 表示保持单进程评估
-cuda_devices: null        # CUDA 并行评估时使用的逻辑设备列表；null 表示默认 [0]
-response_timeout_s: 300.0 # worker 等待 inference server 响应的超时时间（秒）
-server_startup_timeout_s: 300.0 # parent 等待 inference server 就绪的超时时间（秒）
 max_timesteps: 700        # 每回合最大时间步
 device: ${train.device}   # 与训练保持一致
 task_name: "sim_transfer" # 环境任务名称

roboimi/vla/models/heads/act.py

@@ -1,257 +0,0 @@
-"""Local ACT-style CVAE policy head.
-
-This module intentionally reimplements only the ACT model logic needed by the
-RoboIMI VLA training stack. It does not import or vendor the external ACT repo.
-"""
-
-from __future__ import annotations
-
-import math
-from typing import Optional, Tuple
-
-import torch
-import torch.nn as nn
-
-
-def kl_divergence(mu: torch.Tensor, logvar: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    """KL divergence between diagonal posterior N(mu, exp(logvar)) and N(0, I)."""
-    if mu.ndim > 2:
-        mu = mu.view(mu.size(0), -1)
-    if logvar.ndim > 2:
-        logvar = logvar.view(logvar.size(0), -1)
-    klds = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp())
-    total_kld = klds.sum(1).mean(0, keepdim=True)
-    dimension_wise_kld = klds.mean(0)
-    mean_kld = klds.mean(1).mean(0, keepdim=True)
-    return total_kld, dimension_wise_kld, mean_kld
-
-
-def _build_sinusoidal_table(length: int, dim: int) -> torch.Tensor:
-    if length <= 0:
-        raise ValueError(f"length must be positive, got {length}")
-    if dim <= 0:
-        raise ValueError(f"dim must be positive, got {dim}")
-    position = torch.arange(length, dtype=torch.float32).unsqueeze(1)
-    div_term = torch.exp(
-        torch.arange(0, dim, 2, dtype=torch.float32) * (-math.log(10000.0) / max(dim, 1))
-    )
-    table = torch.zeros(length, dim, dtype=torch.float32)
-    table[:, 0::2] = torch.sin(position * div_term)
-    if dim > 1:
-        table[:, 1::2] = torch.cos(position * div_term[: table[:, 1::2].shape[1]])
-    return table.unsqueeze(0)
-
-
-class ACTPolicyHead(nn.Module):
-    """ACT CVAE head using native PyTorch transformer blocks.
-
-    Args follow the existing Hydra style and are intentionally configurable so
-    this implementation is not tied to ACT's original 14-DoF ALOHA setup.
-    """
-
-    def __init__(
-        self,
-        action_dim: int,
-        obs_dim: int,
-        vision_dim: int,
-        num_cams: int,
-        pred_horizon: int,
-        obs_horizon: int = 1,
-        hidden_dim: int = 256,
-        nheads: int = 8,
-        enc_layers: int = 4,
-        dec_layers: int = 6,
-        dim_feedforward: int = 2048,
-        latent_dim: int = 32,
-        dropout: float = 0.1,
-        kl_weight: float = 10.0,
-        activation: str = "gelu",
-        **_: object,
-    ) -> None:
-        super().__init__()
-        self.action_dim = int(action_dim)
-        self.obs_dim = int(obs_dim)
-        self.vision_dim = int(vision_dim)
-        self.num_cams = int(num_cams)
-        self.pred_horizon = int(pred_horizon)
-        self.obs_horizon = int(obs_horizon)
-        self.hidden_dim = int(hidden_dim)
-        self.nheads = int(nheads)
-        self.latent_dim = int(latent_dim)
-        self.kl_weight = float(kl_weight)
-
-        if self.pred_horizon <= 0:
-            raise ValueError(f"pred_horizon must be positive, got {self.pred_horizon}")
-        if self.hidden_dim % self.nheads != 0:
-            raise ValueError(
-                f"hidden_dim({self.hidden_dim}) must be divisible by nheads({self.nheads})"
-            )
-
-        self.cls_embed = nn.Parameter(torch.zeros(1, 1, self.hidden_dim))
-        self.encoder_qpos_proj = nn.Linear(self.obs_dim, self.hidden_dim)
-        self.encoder_action_proj = nn.Linear(self.action_dim, self.hidden_dim)
-        encoder_layer = nn.TransformerEncoderLayer(
-            d_model=self.hidden_dim,
-            nhead=self.nheads,
-            dim_feedforward=int(dim_feedforward),
-            dropout=float(dropout),
-            activation=activation,
-            batch_first=True,
-            norm_first=False,
-        )
-        self.posterior_encoder = nn.TransformerEncoder(encoder_layer, num_layers=int(enc_layers))
-        self.latent_proj = nn.Linear(self.hidden_dim, self.latent_dim * 2)
-
-        self.latent_out_proj = nn.Linear(self.latent_dim, self.hidden_dim)
-        self.decoder_qpos_proj = nn.Linear(self.obs_dim, self.hidden_dim)
-        self.visual_proj = nn.Linear(self.vision_dim, self.hidden_dim)
-        self.memory_type_embed = nn.Embedding(3, self.hidden_dim)  # latent, qpos, visual
-        self.query_embed = nn.Embedding(self.pred_horizon, self.hidden_dim)
-        decoder_layer = nn.TransformerDecoderLayer(
-            d_model=self.hidden_dim,
-            nhead=self.nheads,
-            dim_feedforward=int(dim_feedforward),
-            dropout=float(dropout),
-            activation=activation,
-            batch_first=True,
-            norm_first=False,
-        )
-        self.decoder = nn.TransformerDecoder(decoder_layer, num_layers=int(dec_layers))
-        self.action_head = nn.Linear(self.hidden_dim, self.action_dim)
-
-        self.register_buffer(
-            "posterior_pos_table",
-            _build_sinusoidal_table(self.pred_horizon + 2, self.hidden_dim),
-            persistent=False,
-        )
-        self.register_buffer(
-            "memory_pos_table",
-            _build_sinusoidal_table(1024, self.hidden_dim),
-            persistent=False,
-        )
-        self._reset_parameters()
-
-    def _reset_parameters(self) -> None:
-        nn.init.normal_(self.cls_embed, mean=0.0, std=0.02)
-        for module in self.modules():
-            if isinstance(module, nn.Linear):
-                nn.init.xavier_uniform_(module.weight)
-                if module.bias is not None:
-                    nn.init.zeros_(module.bias)
-            elif isinstance(module, nn.Embedding):
-                nn.init.normal_(module.weight, mean=0.0, std=0.02)
-
-    def _validate_inputs(
-        self,
-        qpos: torch.Tensor,
-        visual_tokens: torch.Tensor,
-        actions: Optional[torch.Tensor],
-        action_is_pad: Optional[torch.Tensor],
-    ) -> None:
-        if qpos.ndim != 2 or qpos.shape[-1] != self.obs_dim:
-            raise ValueError(f"qpos must have shape (B,{self.obs_dim}), got {tuple(qpos.shape)}")
-        if visual_tokens.ndim != 3 or visual_tokens.shape[-1] != self.vision_dim:
-            raise ValueError(
-                f"visual_tokens must have shape (B,N,{self.vision_dim}), got {tuple(visual_tokens.shape)}"
-            )
-        if visual_tokens.shape[0] != qpos.shape[0]:
-            raise ValueError("qpos and visual_tokens batch dimensions must match")
-        if actions is not None:
-            expected = (qpos.shape[0], self.pred_horizon, self.action_dim)
-            if tuple(actions.shape) != expected:
-                raise ValueError(f"actions must have shape {expected}, got {tuple(actions.shape)}")
-            if action_is_pad is not None and tuple(action_is_pad.shape) != expected[:2]:
-                raise ValueError(
-                    f"action_is_pad must have shape {expected[:2]}, got {tuple(action_is_pad.shape)}"
-                )
-
-    def _posterior(
-        self,
-        qpos: torch.Tensor,
-        actions: Optional[torch.Tensor],
-        action_is_pad: Optional[torch.Tensor],
-    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]:
-        batch_size = qpos.shape[0]
-        if actions is None:
-            latent = torch.zeros(batch_size, self.latent_dim, device=qpos.device, dtype=qpos.dtype)
-            return latent, None, None, None
-
-        cls = self.cls_embed.to(dtype=qpos.dtype).expand(batch_size, -1, -1)
-        qpos_token = self.encoder_qpos_proj(qpos).unsqueeze(1)
-        action_tokens = self.encoder_action_proj(actions)
-        tokens = torch.cat([cls, qpos_token, action_tokens], dim=1)
-        pos = self.posterior_pos_table[:, : tokens.shape[1]].to(device=tokens.device, dtype=tokens.dtype)
-        tokens = tokens + pos
-
-        padding_mask = None
-        if action_is_pad is not None:
-            prefix = torch.zeros(
-                batch_size,
-                2,
-                dtype=torch.bool,
-                device=action_is_pad.device,
-            )
-            padding_mask = torch.cat([prefix, action_is_pad.to(torch.bool)], dim=1)
-
-        encoded = self.posterior_encoder(tokens, src_key_padding_mask=padding_mask)
-        latent_info = self.latent_proj(encoded[:, 0])
-        mu, logvar = torch.chunk(latent_info, 2, dim=-1)
-        std = torch.exp(0.5 * logvar)
-        eps = torch.randn_like(std)
-        latent = mu + eps * std
-        kl, _, _ = kl_divergence(mu, logvar)
-        return latent, mu, logvar, kl
-
-    def _memory(self, qpos: torch.Tensor, visual_tokens: torch.Tensor, latent: torch.Tensor) -> torch.Tensor:
-        batch_size = qpos.shape[0]
-        latent_token = self.latent_out_proj(latent).unsqueeze(1)
-        qpos_token = self.decoder_qpos_proj(qpos).unsqueeze(1)
-        visual = self.visual_proj(visual_tokens)
-        memory = torch.cat([latent_token, qpos_token, visual], dim=1)
-
-        if memory.shape[1] > self.memory_pos_table.shape[1]:
-            pos = _build_sinusoidal_table(memory.shape[1], self.hidden_dim).to(
-                device=memory.device,
-                dtype=memory.dtype,
-            )
-        else:
-            pos = self.memory_pos_table[:, : memory.shape[1]].to(
-                device=memory.device,
-                dtype=memory.dtype,
-            )
-        memory = memory + pos
-        type_ids = torch.cat(
-            [
-                torch.zeros(1, dtype=torch.long, device=memory.device),
-                torch.ones(1, dtype=torch.long, device=memory.device),
-                torch.full((memory.shape[1] - 2,), 2, dtype=torch.long, device=memory.device),
-            ]
-        )
-        memory = memory + self.memory_type_embed(type_ids).to(dtype=memory.dtype).unsqueeze(0)
-        if memory.shape[0] != batch_size:
-            raise RuntimeError("internal memory batch shape mismatch")
-        return memory
-
-    def forward(
-        self,
-        qpos: torch.Tensor,
-        visual_tokens: torch.Tensor,
-        actions: Optional[torch.Tensor] = None,
-        action_is_pad: Optional[torch.Tensor] = None,
-    ):
-        self._validate_inputs(qpos, visual_tokens, actions, action_is_pad)
-        latent, mu, logvar, kl = self._posterior(qpos, actions, action_is_pad)
-        memory = self._memory(qpos, visual_tokens, latent)
-        query = self.query_embed.weight.to(dtype=memory.dtype).unsqueeze(0).expand(qpos.shape[0], -1, -1)
-        target = torch.zeros_like(query)
-        decoded = self.decoder(target + query, memory)
-        pred_actions = self.action_head(decoded)
-        if kl is None:
-            kl = torch.zeros(1, device=qpos.device, dtype=qpos.dtype)
-        latent_info = {
-            "mu": mu,
-            "logvar": logvar,
-            "kl": kl,
-            "kl_weight": self.kl_weight,
-        }
-        return pred_actions, latent_info

tests/test_act_agent.py

@@ -1,249 +0,0 @@
-import contextlib
-import sys
-import types
-import unittest
-from pathlib import Path
-
-import torch
-from hydra import compose, initialize_config_dir
-from hydra.core.global_hydra import GlobalHydra
-from hydra.utils import instantiate
-from omegaconf import OmegaConf
-
-
-_REPO_ROOT = Path(__file__).resolve().parents[1]
-_CONFIG_DIR = str((_REPO_ROOT / 'roboimi/vla/conf').resolve())
-_MISSING = object()
-
-
-class FakeVisionBackbone(torch.nn.Module):
-    def __init__(self, output_dim=4, camera_names=('l_vis', 'r_vis', 'front')):
-        super().__init__()
-        self.output_dim = output_dim
-        self.num_cameras = len(camera_names)
-        self.tokens_per_step = self.num_cameras
-        self.camera_names = tuple(camera_names)
-        self.scale = torch.nn.Parameter(torch.tensor(1.0))
-
-    def forward(self, images):
-        features = []
-        for cam_name in self.camera_names:
-            image = images[cam_name]
-            marker = image.mean(dim=(2, 3, 4), keepdim=False).unsqueeze(-1)
-            features.append(marker.repeat(1, 1, self.output_dim) * self.scale)
-        return torch.stack(features, dim=2)
-
-
-class _IdentityCrop:
-    def __init__(self, size):
-        self.size = size
-
-    def __call__(self, x):
-        return x
-
-
-class _FakeResNet(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.conv1 = torch.nn.Conv2d(3, 8, kernel_size=3, padding=1)
-        self.relu1 = torch.nn.ReLU()
-        self.conv2 = torch.nn.Conv2d(8, 16, kernel_size=3, padding=1, stride=2)
-        self.relu2 = torch.nn.ReLU()
-        self.avgpool = torch.nn.AdaptiveAvgPool2d((1, 1))
-        self.fc = torch.nn.Linear(16, 16)
-
-    def forward(self, x):
-        x = self.relu1(self.conv1(x))
-        x = self.relu2(self.conv2(x))
-        x = self.avgpool(x)
-        return self.fc(torch.flatten(x, start_dim=1))
-
-
-@contextlib.contextmanager
-def _stub_torchvision():
-    previous = {}
-
-    def inject(name, module):
-        if name not in previous:
-            previous[name] = sys.modules.get(name, _MISSING)
-        sys.modules[name] = module
-
-    torchvision_module = types.ModuleType('torchvision')
-    models_module = types.ModuleType('torchvision.models')
-    transforms_module = types.ModuleType('torchvision.transforms')
-    models_module.resnet18 = lambda weights=None: _FakeResNet()
-    transforms_module.CenterCrop = _IdentityCrop
-    transforms_module.RandomCrop = _IdentityCrop
-    torchvision_module.models = models_module
-    torchvision_module.transforms = transforms_module
-
-    try:
-        inject('torchvision', torchvision_module)
-        inject('torchvision.models', models_module)
-        inject('torchvision.transforms', transforms_module)
-        yield
-    finally:
-        for name, old in reversed(list(previous.items())):
-            if old is _MISSING:
-                sys.modules.pop(name, None)
-            else:
-                sys.modules[name] = old
-
-
-def _compose_cfg(overrides=None):
-    if not OmegaConf.has_resolver('len'):
-        OmegaConf.register_new_resolver('len', lambda x: len(x))
-    GlobalHydra.instance().clear()
-    with initialize_config_dir(version_base=None, config_dir=_CONFIG_DIR):
-        return compose(config_name='config', overrides=list(overrides or []))
-
-
-def _make_batch(batch_size=2, obs_horizon=2, pred_horizon=4, action_dim=3, obs_dim=5):
-    camera_names = ('l_vis', 'r_vis', 'front')
-    images = {
-        cam_name: torch.full(
-            (batch_size, obs_horizon, 3, 8, 8),
-            float(cam_idx + 1),
-        )
-        for cam_idx, cam_name in enumerate(camera_names)
-    }
-    return {
-        'images': images,
-        'qpos': torch.randn(batch_size, obs_horizon, obs_dim),
-        'action': torch.randn(batch_size, pred_horizon, action_dim),
-        'action_is_pad': torch.zeros(batch_size, pred_horizon, dtype=torch.bool),
-    }
-
-
-class ACTAgentTest(unittest.TestCase):
-    def test_compute_loss_returns_scalar_and_backpropagates(self):
-        from roboimi.vla.agent_act import ACTAgent
-        from roboimi.vla.models.heads.act import ACTPolicyHead
-
-        agent = ACTAgent(
-            vision_backbone=FakeVisionBackbone(output_dim=4),
-            head=ACTPolicyHead(
-                action_dim=3,
-                obs_dim=5,
-                vision_dim=4,
-                num_cams=3,
-                pred_horizon=4,
-                obs_horizon=2,
-                hidden_dim=32,
-                nheads=4,
-                enc_layers=1,
-                dec_layers=1,
-                dim_feedforward=64,
-                latent_dim=8,
-                kl_weight=0.1,
-            ),
-            action_dim=3,
-            obs_dim=5,
-            pred_horizon=4,
-            obs_horizon=2,
-            num_cams=3,
-            camera_names=('l_vis', 'r_vis', 'front'),
-        )
-        loss = agent.compute_loss(_make_batch())
-        self.assertEqual(loss.ndim, 0)
-        self.assertTrue(torch.isfinite(loss))
-        loss.backward()
-        grads = [p.grad for p in agent.parameters() if p.requires_grad]
-        self.assertTrue(any(grad is not None and torch.isfinite(grad).all() for grad in grads))
-
-
-    def test_compute_loss_handles_all_padded_actions_without_nan(self):
-        from roboimi.vla.agent_act import ACTAgent
-        from roboimi.vla.models.heads.act import ACTPolicyHead
-
-        agent = ACTAgent(
-            vision_backbone=FakeVisionBackbone(output_dim=4),
-            head=ACTPolicyHead(
-                action_dim=3,
-                obs_dim=5,
-                vision_dim=4,
-                num_cams=3,
-                pred_horizon=4,
-                obs_horizon=2,
-                hidden_dim=32,
-                nheads=4,
-                enc_layers=1,
-                dec_layers=1,
-                dim_feedforward=64,
-                latent_dim=8,
-                kl_weight=0.1,
-            ),
-            action_dim=3,
-            obs_dim=5,
-            pred_horizon=4,
-            obs_horizon=2,
-            num_cams=3,
-            camera_names=('l_vis', 'r_vis', 'front'),
-        )
-        batch = _make_batch()
-        batch['action_is_pad'][:] = True
-        loss = agent.compute_loss(batch)
-        self.assertEqual(loss.ndim, 0)
-        self.assertTrue(torch.isfinite(loss))
-
-    def test_predict_action_returns_denormalized_chunk_shape(self):
-        from roboimi.vla.agent_act import ACTAgent
-        from roboimi.vla.models.heads.act import ACTPolicyHead
-
-        agent = ACTAgent(
-            vision_backbone=FakeVisionBackbone(output_dim=4),
-            head=ACTPolicyHead(
-                action_dim=3,
-                obs_dim=5,
-                vision_dim=4,
-                num_cams=3,
-                pred_horizon=4,
-                obs_horizon=2,
-                hidden_dim=32,
-                nheads=4,
-                enc_layers=1,
-                dec_layers=1,
-                dim_feedforward=64,
-                latent_dim=8,
-            ),
-            action_dim=3,
-            obs_dim=5,
-            pred_horizon=4,
-            obs_horizon=2,
-            num_cams=3,
-            camera_names=('l_vis', 'r_vis', 'front'),
-        )
-        batch = _make_batch()
-        actions = agent.predict_action(batch['images'], batch['qpos'])
-        self.assertEqual(tuple(actions.shape), (2, 4, 3))
-        self.assertTrue(torch.isfinite(actions).all())
-
-    def test_hydra_instantiates_act_resnet_for_socket_peg_camera_order(self):
-        cfg = _compose_cfg(
-            overrides=[
-                'agent=act_resnet',
-                'data.camera_names=[l_vis,r_vis,front]',
-                'agent.vision_backbone.pretrained_backbone_weights=null',
-                'agent.vision_backbone.input_shape=[3,16,16]',
-                'agent.pred_horizon=4',
-                'agent.obs_horizon=1',
-                'agent.num_action_steps=2',
-                'agent.head.hidden_dim=32',
-                'agent.head.nheads=4',
-                'agent.head.enc_layers=1',
-                'agent.head.dec_layers=1',
-                'agent.head.dim_feedforward=64',
-                'agent.head.latent_dim=8',
-            ]
-        )
-        self.assertEqual(list(cfg.data.camera_names), ['l_vis', 'r_vis', 'front'])
-        self.assertEqual(cfg.agent._target_, 'roboimi.vla.agent_act.ACTAgent')
-        with _stub_torchvision():
-            agent = instantiate(cfg.agent)
-        self.assertEqual(agent.camera_names, ('l_vis', 'r_vis', 'front'))
-        self.assertEqual(agent.pred_horizon, 4)
-        self.assertEqual(agent.vision_encoder.tokens_per_step, 3)
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/test_air_insert_env.py

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

tests/test_eval_vla_execution.py

@@ -1,11 +1,5 @@
 import unittest
-from unittest import mock
 
-import numpy as np
-import torch
-from omegaconf import OmegaConf
-
-from roboimi.demos.vla_scripts import eval_vla
 from roboimi.vla.eval_utils import execute_policy_action
 
 
@@ -20,48 +14,6 @@ class _FakeEnv:
         self.calls.append(("step_jnt", action))
 
 
-class _FakeQueue:
-    def __init__(self, initial_items=None):
-        self.items = list(initial_items or [])
-        self.put_calls = []
-
-    def put(self, item):
-        self.put_calls.append(item)
-        self.items.append(item)
-
-    def get(self, timeout=None):
-        del timeout
-        if not self.items:
-            raise AssertionError("queue unexpectedly empty")
-        return self.items.pop(0)
-
-
-def _make_parallel_cfg(**eval_overrides):
-    eval_cfg = {
-        "ckpt_path": "checkpoints/vla_model_best.pt",
-        "num_episodes": 5,
-        "num_workers": 2,
-        "max_timesteps": 1,
-        "device": "cpu",
-        "task_name": "sim_transfer",
-        "camera_names": ["front"],
-        "use_smoothing": False,
-        "smooth_alpha": 0.3,
-        "verbose_action": False,
-        "headless": True,
-        "artifact_dir": None,
-        "save_artifacts": False,
-        "save_summary_json": False,
-        "save_timing": False,
-        "save_trajectory": False,
-        "save_trajectory_npz": False,
-        "record_video": False,
-        "save_trajectory_image": False,
-    }
-    eval_cfg.update(eval_overrides)
-    return OmegaConf.create({"agent": {}, "eval": eval_cfg})
-
-
 class EvalVLAExecutionTest(unittest.TestCase):
     def test_execute_policy_action_uses_ee_step(self):
         env = _FakeEnv()
@@ -71,691 +23,6 @@ class EvalVLAExecutionTest(unittest.TestCase):
 
         self.assertEqual(env.calls, [("step", action)])
 
-    def test_split_episode_indices_balances_workers(self):
-        self.assertEqual(
-            eval_vla._split_episode_indices(num_episodes=10, num_workers=3),
-            [[0, 1, 2, 3], [4, 5, 6], [7, 8, 9]],
-        )
-
-    def test_normalize_num_workers_caps_worker_count_to_episode_count(self):
-        self.assertEqual(eval_vla._normalize_num_workers(num_workers=5, num_episodes=2), 2)
-
-    def test_plan_episode_box_poses_uses_global_episode_order(self):
-        planned_poses = [
-            np.array([0.1, 0.2, 0.3], dtype=np.float32),
-            np.array([1.1, 1.2, 1.3], dtype=np.float32),
-            np.array([2.1, 2.2, 2.3], dtype=np.float32),
-        ]
-        sampler = mock.Mock(side_effect=planned_poses)
-
-        result = eval_vla._plan_episode_box_poses(num_episodes=3, sampler=sampler)
-
-        self.assertEqual(sampler.call_count, 3)
-        self.assertEqual(len(result), 3)
-        for expected, actual in zip(planned_poses, result):
-            np.testing.assert_array_equal(actual, expected)
-
-    def test_resolve_policy_camera_names_matches_vlaagent_fallback_sorting(self):
-        cfg = OmegaConf.create(
-            {
-                "agent": {
-                    "_target_": "roboimi.vla.agent.VLAAgent",
-                },
-                "eval": {
-                    "camera_names": ["r_vis", "top", "front"],
-                },
-            }
-        )
-
-        self.assertEqual(
-            eval_vla._resolve_policy_camera_names(cfg),
-            ["front", "r_vis", "top"],
-        )
-
-    def test_resolve_policy_camera_names_matches_gr00t_fallback_input_order(self):
-        cfg = OmegaConf.create(
-            {
-                "agent": {
-                    "_target_": "roboimi.vla.agent_gr00t_dit.VLAAgentGr00tDiT",
-                },
-                "eval": {
-                    "camera_names": ["r_vis", "top", "front"],
-                },
-            }
-        )
-
-        self.assertEqual(
-            eval_vla._resolve_policy_camera_names(cfg),
-            ["r_vis", "top", "front"],
-        )
-
-    def test_build_episode_plans_without_box_poses_keeps_serial_sampling_lazy(self):
-        plans = eval_vla._build_episode_plans(num_episodes=3)
-
-        self.assertEqual(
-            plans,
-            [
-                {"episode_index": 0},
-                {"episode_index": 1},
-                {"episode_index": 2},
-            ],
-        )
-
-    def test_prepare_local_policy_batch_pads_latest_observation_to_obs_horizon(self):
-        queues = eval_vla._new_local_policy_queues(obs_horizon=3)
-        observation = {
-            "qpos": torch.tensor([1.0, 2.0], dtype=torch.float32),
-            "images": {
-                "front": torch.tensor([[[1.0]]], dtype=torch.float32),
-            },
-        }
-
-        eval_vla._populate_local_policy_queues(queues, observation)
-        batch = eval_vla._prepare_local_policy_batch(
-            queues,
-            obs_horizon=3,
-            camera_names=["front"],
-        )
-
-        self.assertEqual(tuple(batch["qpos"].shape), (1, 3, 2))
-        self.assertEqual(tuple(batch["images"]["front"].shape), (1, 3, 1, 1, 1))
-        np.testing.assert_array_equal(
-            batch["qpos"][0].cpu().numpy(),
-            np.array([[1.0, 2.0], [1.0, 2.0], [1.0, 2.0]], dtype=np.float32),
-        )
-        np.testing.assert_array_equal(
-            batch["images"]["front"][0].cpu().numpy(),
-            np.array([[[[1.0]]], [[[1.0]]], [[[1.0]]]], dtype=np.float32),
-        )
-
-    def test_enqueue_predicted_actions_uses_executable_slice(self):
-        queues = eval_vla._new_local_policy_queues(obs_horizon=2)
-        predicted_actions = torch.tensor(
-            [[[10.0], [20.0], [30.0], [40.0]]],
-            dtype=torch.float32,
-        )
-
-        eval_vla._enqueue_predicted_actions(
-            queues,
-            predicted_actions=predicted_actions,
-            obs_horizon=2,
-            num_action_steps=2,
-        )
-
-        self.assertEqual(len(queues["action"]), 2)
-        np.testing.assert_array_equal(queues["action"].popleft().numpy(), np.array([20.0], dtype=np.float32))
-        np.testing.assert_array_equal(queues["action"].popleft().numpy(), np.array([30.0], dtype=np.float32))
-
-    def test_remote_policy_runner_only_requests_server_inference_when_local_action_queue_is_empty(self):
-        request_queue = _FakeQueue()
-        response_queue = _FakeQueue(
-            [
-                {
-                    "type": "predict_chunk_result",
-                    "actions": np.asarray([[[10.0], [20.0], [30.0]]], dtype=np.float32),
-                }
-            ]
-        )
-        runner = eval_vla._RemotePolicyRunner(
-            worker_index=3,
-            server_index=1,
-            request_queue=request_queue,
-            response_queue=response_queue,
-            camera_names=["front"],
-            obs_horizon=2,
-            num_action_steps=2,
-        )
-        first_observation = {
-            "qpos": torch.tensor([1.0, 2.0], dtype=torch.float32),
-            "images": {"front": torch.tensor([[[1.0]]], dtype=torch.float32)},
-        }
-        second_observation = {
-            "qpos": torch.tensor([3.0, 4.0], dtype=torch.float32),
-            "images": {"front": torch.tensor([[[2.0]]], dtype=torch.float32)},
-        }
-
-        first_action, first_forward = runner.select_action(
-            first_observation,
-            episode_index=7,
-            timestep=0,
-        )
-        second_action, second_forward = runner.select_action(
-            second_observation,
-            episode_index=7,
-            timestep=1,
-        )
-
-        self.assertTrue(first_forward)
-        self.assertFalse(second_forward)
-        self.assertEqual(len(request_queue.put_calls), 1)
-        self.assertEqual(request_queue.put_calls[0]["type"], "predict_chunk")
-        self.assertEqual(request_queue.put_calls[0]["worker_index"], 3)
-        self.assertEqual(request_queue.put_calls[0]["server_index"], 1)
-        np.testing.assert_array_equal(first_action.numpy(), np.array([20.0], dtype=np.float32))
-        np.testing.assert_array_equal(second_action.numpy(), np.array([30.0], dtype=np.float32))
-
-    def test_merge_worker_summaries_sorts_episodes_and_recomputes_aggregates(self):
-        worker_summaries = [
-            {
-                "avg_inference_fps": 999.0,
-                "avg_control_fps": 999.0,
-                "avg_obs_read_time_ms": 999.0,
-                "avg_total_time_ms": 999.0,
-                "timing_summary": {"count": 999, "model_forward_count": 999},
-                "episodes": [
-                    {
-                        "episode_index": 2,
-                        "episode_reward": 9.0,
-                        "episode_max_reward": 4.0,
-                        "inference_fps": 30.0,
-                        "control_fps": 15.0,
-                    }
-                ],
-                "_merge_state": {
-                    "obs_read_time_ms": [9.0],
-                    "preprocess_time_ms": [1.0],
-                    "inference_time_ms": [3.0],
-                    "env_step_time_ms": [4.0],
-                    "total_time_ms": [10.0],
-                    "model_forward_flags": [False],
-                },
-            },
-            {
-                "avg_inference_fps": 888.0,
-                "avg_control_fps": 888.0,
-                "avg_obs_read_time_ms": 888.0,
-                "avg_total_time_ms": 888.0,
-                "timing_summary": {"count": 888, "model_forward_count": 888},
-                "episodes": [
-                    {
-                        "episode_index": 1,
-                        "episode_reward": 6.0,
-                        "episode_max_reward": 3.0,
-                        "inference_fps": 20.0,
-                        "control_fps": 10.0,
-                    },
-                    {
-                        "episode_index": 0,
-                        "episode_reward": 5.0,
-                        "episode_max_reward": 2.0,
-                        "inference_fps": 10.0,
-                        "control_fps": 5.0,
-                    },
-                ],
-                "_merge_state": {
-                    "obs_read_time_ms": [1.0, 2.0, 12.0],
-                    "preprocess_time_ms": [2.0, 3.0, 4.0],
-                    "inference_time_ms": [4.0, 5.0, 6.0],
-                    "env_step_time_ms": [6.0, 7.0, 8.0],
-                    "total_time_ms": [8.0, 9.0, 20.0],
-                    "model_forward_flags": [True, False, True],
-                },
-            },
-        ]
-        artifact_paths = {
-            "output_dir": "/tmp/merged",
-            "summary_json": "/tmp/merged/rollout_summary.json",
-            "timing_json": "/tmp/merged/timing.json",
-            "trajectory_npz": None,
-            "video_mp4": None,
-            "video_camera_name": None,
-        }
-
-        merged = eval_vla._merge_worker_summaries(worker_summaries, artifact_paths)
-
-        self.assertEqual([episode["episode_index"] for episode in merged["episodes"]], [0, 1, 2])
-        self.assertEqual(merged["episode_rewards"], [5.0, 6.0, 9.0])
-        self.assertEqual(merged["episode_max_rewards"], [2.0, 3.0, 4.0])
-        self.assertAlmostEqual(merged["avg_reward"], 20.0 / 3.0)
-        self.assertAlmostEqual(merged["avg_max_reward"], 3.0)
-        self.assertAlmostEqual(merged["avg_inference_fps"], 20.0)
-        self.assertAlmostEqual(merged["avg_control_fps"], 10.0)
-        self.assertAlmostEqual(merged["avg_obs_read_time_ms"], 6.0)
-        self.assertAlmostEqual(merged["avg_total_time_ms"], 47.0 / 4.0)
-        self.assertEqual(merged["timing_summary"]["count"], 4)
-        self.assertEqual(merged["timing_summary"]["model_forward_count"], 2)
-        self.assertEqual(merged["artifact_dir"], "/tmp/merged")
-        self.assertEqual(merged["artifacts"], artifact_paths)
-
-    def test_build_cuda_server_payloads_uses_round_robin_worker_assignment(self):
-        cfg = _make_parallel_cfg(num_episodes=4, num_workers=4, device="cuda", cuda_devices=[0, 1])
-        artifact_paths = {"output_dir": None}
-
-        with mock.patch.object(
-            eval_vla,
-            "sample_transfer_pose",
-            side_effect=[
-                np.array([0.1, 0.2, 0.3], dtype=np.float32),
-                np.array([0.4, 0.5, 0.6], dtype=np.float32),
-                np.array([0.7, 0.8, 0.9], dtype=np.float32),
-                np.array([1.0, 1.1, 1.2], dtype=np.float32),
-            ],
-        ):
-            worker_payloads, _ = eval_vla._build_parallel_worker_payloads(cfg, artifact_paths)
-
-        server_payloads, assigned_workers = eval_vla._build_cuda_server_payloads(
-            cfg,
-            worker_payloads=worker_payloads,
-            cuda_devices=[0, 1],
-        )
-
-        self.assertEqual([payload["device_index"] for payload in server_payloads], [0, 1])
-        self.assertEqual([payload["worker_index"] for payload in assigned_workers], [0, 1, 2, 3])
-        self.assertEqual([payload["server_index"] for payload in assigned_workers], [0, 1, 0, 1])
-        self.assertEqual(server_payloads[0]["worker_indices"], [0, 2])
-        self.assertEqual(server_payloads[1]["worker_indices"], [1, 3])
-
-    def test_run_eval_parallel_dispatches_episode_splits_and_box_poses(self):
-        cfg = _make_parallel_cfg(num_episodes=5, num_workers=2, artifact_dir="/tmp/parallel-root")
-        planned_poses = [
-            np.array([float(index), float(index) + 0.1, float(index) + 0.2], dtype=np.float32)
-            for index in range(5)
-        ]
-        observed_payloads = []
-
-        def fake_run_spawn_jobs(payloads, max_workers, worker_fn):
-            del worker_fn
-            self.assertEqual(max_workers, 2)
-            observed_payloads.extend(payloads)
-            return [
-                {
-                    "episodes": [
-                        {
-                            "episode_index": 4,
-                            "episode_reward": 5.0,
-                            "episode_max_reward": 5.0,
-                            "inference_fps": 50.0,
-                            "control_fps": 25.0,
-                        },
-                        {
-                            "episode_index": 3,
-                            "episode_reward": 4.0,
-                            "episode_max_reward": 4.0,
-                            "inference_fps": 40.0,
-                            "control_fps": 20.0,
-                        },
-                    ],
-                    "_merge_state": {
-                        "obs_read_time_ms": [4.0, 5.0],
-                        "preprocess_time_ms": [1.0, 1.0],
-                        "inference_time_ms": [2.0, 2.0],
-                        "env_step_time_ms": [3.0, 3.0],
-                        "total_time_ms": [4.0, 5.0],
-                        "model_forward_flags": [True, True],
-                    },
-                },
-                {
-                    "episodes": [
-                        {
-                            "episode_index": 2,
-                            "episode_reward": 3.0,
-                            "episode_max_reward": 3.0,
-                            "inference_fps": 30.0,
-                            "control_fps": 15.0,
-                        },
-                        {
-                            "episode_index": 1,
-                            "episode_reward": 2.0,
-                            "episode_max_reward": 2.0,
-                            "inference_fps": 20.0,
-                            "control_fps": 10.0,
-                        },
-                        {
-                            "episode_index": 0,
-                            "episode_reward": 1.0,
-                            "episode_max_reward": 1.0,
-                            "inference_fps": 10.0,
-                            "control_fps": 5.0,
-                        },
-                    ],
-                    "_merge_state": {
-                        "obs_read_time_ms": [1.0, 2.0, 3.0],
-                        "preprocess_time_ms": [1.0, 1.0, 1.0],
-                        "inference_time_ms": [2.0, 2.0, 2.0],
-                        "env_step_time_ms": [3.0, 3.0, 3.0],
-                        "total_time_ms": [1.0, 2.0, 3.0],
-                        "model_forward_flags": [False, True, False],
-                    },
-                },
-            ]
-
-        with mock.patch.object(
-            eval_vla,
-            "sample_transfer_pose",
-            side_effect=planned_poses,
-        ), mock.patch.object(
-            eval_vla,
-            "_run_spawn_jobs",
-            side_effect=fake_run_spawn_jobs,
-        ):
-            summary = eval_vla._run_eval_parallel(cfg)
-
-        self.assertEqual(len(observed_payloads), 2)
-        self.assertEqual(
-            [[plan["episode_index"] for plan in payload["episode_plans"]] for payload in observed_payloads],
-            [[0, 1, 2], [3, 4]],
-        )
-        for payload in observed_payloads:
-            for plan in payload["episode_plans"]:
-                np.testing.assert_array_equal(
-                    np.asarray(plan["box_pos"], dtype=np.float32),
-                    planned_poses[plan["episode_index"]],
-                )
-        self.assertEqual([episode["episode_index"] for episode in summary["episodes"]], [0, 1, 2, 3, 4])
-        self.assertEqual(summary["episode_rewards"], [1.0, 2.0, 3.0, 4.0, 5.0])
-        self.assertEqual(summary["num_episodes"], 5)
-
-    def test_build_parallel_worker_payloads_keeps_socket_peg_sampling_lazy(self):
-        cfg = _make_parallel_cfg(
-            num_episodes=3,
-            num_workers=2,
-            task_name="sim_air_insert_socket_peg",
-        )
-        artifact_paths = {"output_dir": None}
-
-        with mock.patch.object(
-            eval_vla,
-            "sample_transfer_pose",
-            side_effect=AssertionError("socket-peg parallel eval should not pre-sample transfer poses"),
-        ):
-            worker_payloads, _ = eval_vla._build_parallel_worker_payloads(cfg, artifact_paths)
-
-        episode_plans = [
-            plan
-            for payload in worker_payloads
-            for plan in payload["episode_plans"]
-        ]
-        self.assertEqual(
-            episode_plans,
-            [
-                {"episode_index": 0},
-                {"episode_index": 1},
-                {"episode_index": 2},
-            ],
-        )
-
-    def test_run_eval_parallel_allows_trajectory_images_and_keeps_worker_artifact_paths(self):
-        cfg = _make_parallel_cfg(
-            num_episodes=2,
-            num_workers=2,
-            artifact_dir="/tmp/parallel-images",
-            save_summary_json=True,
-            save_trajectory_image=True,
-        )
-        observed_payloads = []
-
-        def fake_run_spawn_jobs(payloads, max_workers, worker_fn):
-            del worker_fn
-            self.assertEqual(max_workers, 2)
-            observed_payloads.extend(payloads)
-            return [
-                {
-                    "episodes": [
-                        {
-                            "episode_index": 0,
-                            "episode_reward": 1.0,
-                            "episode_max_reward": 1.0,
-                            "inference_fps": 10.0,
-                            "control_fps": 5.0,
-                            "artifact_paths": {
-                                "trajectory_image": f"{payloads[0]['artifact_dir']}/rollout_front_ep01_trajectory.png",
-                            },
-                        },
-                    ],
-                    "_merge_state": {
-                        "obs_read_time_ms": [1.0],
-                        "preprocess_time_ms": [1.0],
-                        "inference_time_ms": [1.0],
-                        "env_step_time_ms": [1.0],
-                        "total_time_ms": [1.0],
-                        "model_forward_flags": [True],
-                    },
-                },
-                {
-                    "episodes": [
-                        {
-                            "episode_index": 1,
-                            "episode_reward": 2.0,
-                            "episode_max_reward": 2.0,
-                            "inference_fps": 20.0,
-                            "control_fps": 10.0,
-                            "artifact_paths": {
-                                "trajectory_image": f"{payloads[1]['artifact_dir']}/rollout_front_ep02_trajectory.png",
-                            },
-                        },
-                    ],
-                    "_merge_state": {
-                        "obs_read_time_ms": [2.0],
-                        "preprocess_time_ms": [2.0],
-                        "inference_time_ms": [2.0],
-                        "env_step_time_ms": [2.0],
-                        "total_time_ms": [2.0],
-                        "model_forward_flags": [False],
-                    },
-                },
-            ]
-
-        with mock.patch.object(
-            eval_vla,
-            "sample_transfer_pose",
-            side_effect=[
-                np.array([0.1, 0.2, 0.3], dtype=np.float32),
-                np.array([0.4, 0.5, 0.6], dtype=np.float32),
-            ],
-        ), mock.patch.object(
-            eval_vla,
-            "_run_spawn_jobs",
-            side_effect=fake_run_spawn_jobs,
-        ):
-            summary = eval_vla._run_eval_parallel(cfg)
-
-        self.assertEqual(len(observed_payloads), 2)
-        self.assertTrue(observed_payloads[0]["artifact_dir"].endswith("workers/worker_00"))
-        self.assertTrue(observed_payloads[1]["artifact_dir"].endswith("workers/worker_01"))
-        self.assertTrue(
-            summary["episodes"][0]["artifact_paths"]["trajectory_image"].endswith(
-                "workers/worker_00/rollout_front_ep01_trajectory.png"
-            )
-        )
-        self.assertTrue(
-            summary["episodes"][1]["artifact_paths"]["trajectory_image"].endswith(
-                "workers/worker_01/rollout_front_ep02_trajectory.png"
-            )
-        )
-
-    def test_run_eval_parallel_surfaces_worker_failures(self):
-        cfg = _make_parallel_cfg(num_episodes=2, num_workers=2)
-
-        with mock.patch.object(
-            eval_vla,
-            "sample_transfer_pose",
-            side_effect=[
-                np.array([0.1, 0.2, 0.3], dtype=np.float32),
-                np.array([0.4, 0.5, 0.6], dtype=np.float32),
-            ],
-        ), mock.patch.object(
-            eval_vla,
-            "_run_spawn_jobs",
-            side_effect=RuntimeError("boom"),
-        ):
-            with self.assertRaisesRegex(RuntimeError, "Parallel rollout worker failed"):
-                eval_vla._run_eval_parallel(cfg)
-
-    def test_run_eval_parallel_cuda_builds_server_payloads_and_merges_worker_results(self):
-        cfg = _make_parallel_cfg(
-            num_episodes=4,
-            num_workers=4,
-            device="cuda",
-            cuda_devices=[0],
-            artifact_dir="/tmp/cuda-root",
-        )
-        observed_server_payloads = []
-        observed_worker_payloads = []
-
-        def fake_run_cuda_parallel_processes(server_payloads, worker_payloads):
-            observed_server_payloads.extend(server_payloads)
-            observed_worker_payloads.extend(worker_payloads)
-            return [
-                {
-                    "episodes": [
-                        {
-                            "episode_index": 2,
-                            "episode_reward": 3.0,
-                            "episode_max_reward": 3.0,
-                            "inference_fps": 30.0,
-                            "control_fps": 15.0,
-                        },
-                        {
-                            "episode_index": 0,
-                            "episode_reward": 1.0,
-                            "episode_max_reward": 1.0,
-                            "inference_fps": 10.0,
-                            "control_fps": 5.0,
-                        },
-                    ],
-                    "_merge_state": {
-                        "obs_read_time_ms": [1.0, 2.0],
-                        "preprocess_time_ms": [1.0, 1.0],
-                        "inference_time_ms": [2.0, 2.0],
-                        "env_step_time_ms": [3.0, 3.0],
-                        "total_time_ms": [4.0, 4.0],
-                        "model_forward_flags": [True, False],
-                    },
-                },
-                {
-                    "episodes": [
-                        {
-                            "episode_index": 3,
-                            "episode_reward": 4.0,
-                            "episode_max_reward": 4.0,
-                            "inference_fps": 40.0,
-                            "control_fps": 20.0,
-                        },
-                        {
-                            "episode_index": 1,
-                            "episode_reward": 2.0,
-                            "episode_max_reward": 2.0,
-                            "inference_fps": 20.0,
-                            "control_fps": 10.0,
-                        },
-                    ],
-                    "_merge_state": {
-                        "obs_read_time_ms": [3.0, 4.0],
-                        "preprocess_time_ms": [1.0, 1.0],
-                        "inference_time_ms": [2.0, 2.0],
-                        "env_step_time_ms": [3.0, 3.0],
-                        "total_time_ms": [4.0, 4.0],
-                        "model_forward_flags": [True, True],
-                    },
-                },
-            ]
-
-        with mock.patch.object(
-            eval_vla,
-            "sample_transfer_pose",
-            side_effect=[
-                np.array([0.1, 0.2, 0.3], dtype=np.float32),
-                np.array([0.4, 0.5, 0.6], dtype=np.float32),
-                np.array([0.7, 0.8, 0.9], dtype=np.float32),
-                np.array([1.0, 1.1, 1.2], dtype=np.float32),
-            ],
-        ), mock.patch.object(
-            eval_vla,
-            "_run_cuda_parallel_processes",
-            side_effect=fake_run_cuda_parallel_processes,
-            create=True,
-        ):
-            summary = eval_vla._run_eval_parallel_cuda(cfg)
-
-        self.assertEqual(len(observed_server_payloads), 1)
-        self.assertEqual(observed_server_payloads[0]["device_index"], 0)
-        self.assertEqual(len(observed_worker_payloads), 4)
-        self.assertTrue(all(payload["server_index"] == 0 for payload in observed_worker_payloads))
-        self.assertEqual([episode["episode_index"] for episode in summary["episodes"]], [0, 1, 2, 3])
-        self.assertEqual(summary["episode_rewards"], [1.0, 2.0, 3.0, 4.0])
-        self.assertEqual(summary["num_episodes"], 4)
-
-    def test_run_eval_parallel_cuda_surfaces_server_failures(self):
-        cfg = _make_parallel_cfg(num_episodes=2, num_workers=2, device="cuda", cuda_devices=[0])
-
-        with mock.patch.object(
-            eval_vla,
-            "sample_transfer_pose",
-            side_effect=[
-                np.array([0.1, 0.2, 0.3], dtype=np.float32),
-                np.array([0.4, 0.5, 0.6], dtype=np.float32),
-            ],
-        ), mock.patch.object(
-            eval_vla,
-            "_run_cuda_parallel_processes",
-            side_effect=RuntimeError("server boom"),
-            create=True,
-        ):
-            with self.assertRaisesRegex(RuntimeError, "Parallel CUDA rollout failed"):
-                eval_vla._run_eval_parallel_cuda(cfg)
-
-    def test_run_spawn_jobs_supports_real_spawn_with_actual_eval_worker_entry(self):
-        payloads = [
-            {"_spawn_probe": True, "probe_value": 1, "worker_index": 0},
-            {"_spawn_probe": True, "probe_value": 2, "worker_index": 1},
-        ]
-
-        results = eval_vla._run_spawn_jobs(
-            payloads=payloads,
-            max_workers=2,
-            worker_fn=eval_vla._run_eval_worker_entry,
-        )
-
-        self.assertEqual(sorted(result["probe_value"] for result in results), [1, 2])
-        self.assertEqual(sorted(result["worker_index"] for result in results), [0, 1])
-
-    def test_cuda_server_and_env_worker_entrypoints_support_real_spawn_probe(self):
-        ctx = eval_vla.multiprocessing.get_context("spawn")
-        request_queue = ctx.Queue()
-        response_queue = ctx.Queue()
-        result_queue = ctx.Queue()
-
-        server = ctx.Process(
-            target=eval_vla._inference_server_main,
-            args=(
-                {
-                    "_spawn_probe": True,
-                    "server_index": 0,
-                    "request_queue": request_queue,
-                    "response_queues": [response_queue],
-                },
-            ),
-        )
-        worker = ctx.Process(
-            target=eval_vla._env_worker_main,
-            args=(
-                {
-                    "_spawn_probe": True,
-                    "worker_index": 0,
-                    "server_index": 0,
-                    "request_queue": request_queue,
-                    "response_queue": response_queue,
-                    "result_queue": result_queue,
-                },
-            ),
-        )
-
-        server.start()
-        worker.start()
-
-        result = result_queue.get(timeout=10.0)
-
-        worker.join(timeout=10.0)
-        request_queue.put({"type": "shutdown_server"})
-        server.join(timeout=10.0)
-
-        self.assertEqual(result["kind"], "worker_result")
-        self.assertEqual(result["summary"]["probe_worker_index"], 0)
-        self.assertEqual(result["summary"]["probe_server_index"], 0)
-        self.assertEqual(result["summary"]["probe_actions"], [[[11.0], [22.0], [33.0]]])
-        self.assertEqual(worker.exitcode, 0)
-        self.assertEqual(server.exitcode, 0)
-
 
 if __name__ == "__main__":
     unittest.main()

tests/test_eval_vla_headless.py

@@ -36,8 +36,8 @@ class _FakeEnv:
         self.render_calls = 0
         self.reset_calls = []
 
-    def reset(self, 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
@@ -74,7 +74,7 @@ class _FakeRenderer:
         self._env = env
         self._frames = [
             np.full((4, 4, 3), fill_value=index, dtype=np.uint8)
-            for index in range(8)
+            for index in range(5)
         ]
         self._index = 0
 
@@ -144,7 +144,7 @@ class EvalVLAHeadlessTest(unittest.TestCase):
             is_render=False,
             control_freq=30,
             is_interpolate=True,
-            cam_view="top",
+            cam_view="angle",
         )
 
     def test_headless_sync_camera_capture_populates_images_without_gui_calls(self):
@@ -153,10 +153,11 @@ class EvalVLAHeadlessTest(unittest.TestCase):
         env.mj_data = object()
         env.exit_flag = False
         env.is_render = False
-        env.cam = 'top'
+        env.cam = 'angle'
         env.r_vis = None
         env.l_vis = None
         env.top = None
+        env.angle = None
         env.front = None
         env._offscreen_renderer = None
 
@@ -175,6 +176,7 @@ class EvalVLAHeadlessTest(unittest.TestCase):
         self.assertIsNotNone(env.r_vis)
         self.assertIsNotNone(env.l_vis)
         self.assertIsNotNone(env.top)
+        self.assertIsNotNone(env.angle)
         self.assertIsNotNone(env.front)
 
     def test_cam_start_skips_background_thread_when_headless(self):
@@ -194,10 +196,11 @@ class EvalVLAHeadlessTest(unittest.TestCase):
         env.mj_data = object()
         env.exit_flag = False
         env.is_render = False
-        env.cam = "top"
+        env.cam = "angle"
         env.r_vis = None
         env.l_vis = None
         env.top = None
+        env.angle = None
         env.front = None
 
         with mock.patch(
@@ -214,33 +217,9 @@ class EvalVLAHeadlessTest(unittest.TestCase):
         self.assertIsNotNone(env.r_vis)
         self.assertIsNotNone(env.l_vis)
         self.assertIsNotNone(env.top)
+        self.assertIsNotNone(env.angle)
         self.assertIsNotNone(env.front)
 
-    def test_dual_diana_step_refreshes_obs_after_physics_step(self):
-        env = DualDianaMed.__new__(DualDianaMed)
-        env.compute_qpos = np.zeros(16)
-        env.interpolator_left = None
-        env.interpolator_right = None
-        env.control_timestep = 0.001
-        env.model_timestep = 0.001
-        env.base_time = 0.0
-        events = []
-
-        def fake_get_obs():
-            events.append("obs")
-            return {"images": {}, "qpos": np.zeros(16, dtype=np.float32)}
-
-        env._get_obs = fake_get_obs
-
-        with mock.patch(
-            "roboimi.envs.double_base.MujocoEnv.step",
-            autospec=True,
-            side_effect=lambda _self, _action: events.append("physics"),
-        ):
-            env.step(np.zeros(16))
-
-        self.assertEqual(events, ["physics", "obs"])
-
     def test_eval_main_headless_skips_render_and_still_executes_policy(self):
         fake_env = _FakeEnv()
         fake_agent = _FakeAgent()
@@ -349,255 +328,5 @@ class EvalVLAHeadlessTest(unittest.TestCase):
         self.assertEqual(summary["num_episodes"], 2)
 
 
-    def test_eval_config_exposes_num_workers_default(self):
-        eval_cfg = OmegaConf.load(Path("roboimi/vla/conf/eval/eval.yaml"))
-
-        self.assertIn("num_workers", eval_cfg)
-        self.assertEqual(eval_cfg.num_workers, 1)
-
-    def test_eval_config_exposes_cuda_devices_default(self):
-        eval_cfg = OmegaConf.load(Path("roboimi/vla/conf/eval/eval.yaml"))
-
-        self.assertIn("cuda_devices", eval_cfg)
-        self.assertIsNone(eval_cfg.cuda_devices)
-
-    def test_eval_config_exposes_parallel_timeout_defaults(self):
-        eval_cfg = OmegaConf.load(Path("roboimi/vla/conf/eval/eval.yaml"))
-
-        self.assertIn("response_timeout_s", eval_cfg)
-        self.assertIn("server_startup_timeout_s", eval_cfg)
-        self.assertEqual(eval_cfg.response_timeout_s, 300.0)
-        self.assertEqual(eval_cfg.server_startup_timeout_s, 300.0)
-
-    def test_run_eval_uses_serial_path_when_num_workers_is_one(self):
-        cfg = OmegaConf.create(
-            {
-                "eval": {
-                    "num_workers": 1,
-                    "num_episodes": 3,
-                }
-            }
-        )
-
-        with mock.patch.object(
-            eval_vla,
-            "_run_eval_serial",
-            return_value={"mode": "serial"},
-        ) as run_eval_serial, mock.patch.object(
-            eval_vla,
-            "_run_eval_parallel",
-        ) as run_eval_parallel:
-            result = eval_vla._run_eval(cfg)
-
-        self.assertEqual(result, {"mode": "serial"})
-        run_eval_serial.assert_called_once_with(cfg)
-        run_eval_parallel.assert_not_called()
-
-    def test_run_eval_uses_serial_path_when_requested_workers_collapse_to_one(self):
-        cfg = OmegaConf.create(
-            {
-                "eval": {
-                    "num_workers": 8,
-                    "num_episodes": 1,
-                }
-            }
-        )
-
-        with mock.patch.object(
-            eval_vla,
-            "_run_eval_serial",
-            return_value={"mode": "serial"},
-        ) as run_eval_serial, mock.patch.object(
-            eval_vla,
-            "_run_eval_parallel",
-        ) as run_eval_parallel:
-            result = eval_vla._run_eval(cfg)
-
-        self.assertEqual(result, {"mode": "serial"})
-        run_eval_serial.assert_called_once_with(cfg)
-        run_eval_parallel.assert_not_called()
-
-    def test_run_eval_parallel_requires_headless_true(self):
-        cfg = OmegaConf.create(
-            {
-                "agent": {},
-                "eval": {
-                    "ckpt_path": "checkpoints/vla_model_best.pt",
-                    "num_episodes": 2,
-                    "num_workers": 2,
-                    "max_timesteps": 1,
-                    "device": "cpu",
-                    "task_name": "sim_transfer",
-                    "camera_names": ["front"],
-                    "use_smoothing": False,
-                    "smooth_alpha": 0.3,
-                    "verbose_action": False,
-                    "headless": False,
-                },
-            }
-        )
-
-        with self.assertRaisesRegex(ValueError, "headless=true"):
-            eval_vla._run_eval_parallel(cfg)
-
-    def test_run_eval_parallel_dispatches_to_cpu_workers_when_device_is_cpu(self):
-        cfg = OmegaConf.create(
-            {
-                "agent": {},
-                "eval": {
-                    "ckpt_path": "checkpoints/vla_model_best.pt",
-                    "num_episodes": 2,
-                    "num_workers": 2,
-                    "max_timesteps": 1,
-                    "device": "cpu",
-                    "task_name": "sim_transfer",
-                    "camera_names": ["front"],
-                    "use_smoothing": False,
-                    "smooth_alpha": 0.3,
-                    "verbose_action": False,
-                    "headless": True,
-                    "cuda_devices": None,
-                },
-            }
-        )
-
-        with mock.patch.object(
-            eval_vla,
-            "_run_eval_parallel_cpu",
-            return_value={"mode": "cpu"},
-            create=True,
-        ) as run_cpu_parallel, mock.patch.object(
-            eval_vla,
-            "_run_eval_parallel_cuda",
-            create=True,
-        ) as run_cuda_parallel:
-            result = eval_vla._run_eval_parallel(cfg)
-
-        self.assertEqual(result, {"mode": "cpu"})
-        run_cpu_parallel.assert_called_once_with(cfg)
-        run_cuda_parallel.assert_not_called()
-
-    def test_run_eval_parallel_dispatches_to_cuda_servers_when_device_is_cuda(self):
-        cfg = OmegaConf.create(
-            {
-                "agent": {},
-                "eval": {
-                    "ckpt_path": "checkpoints/vla_model_best.pt",
-                    "num_episodes": 2,
-                    "num_workers": 2,
-                    "max_timesteps": 1,
-                    "device": "cuda",
-                    "task_name": "sim_transfer",
-                    "camera_names": ["front"],
-                    "use_smoothing": False,
-                    "smooth_alpha": 0.3,
-                    "verbose_action": False,
-                    "headless": True,
-                    "cuda_devices": [0],
-                },
-            }
-        )
-
-        with mock.patch.object(
-            eval_vla,
-            "_run_eval_parallel_cpu",
-            create=True,
-        ) as run_cpu_parallel, mock.patch.object(
-            eval_vla,
-            "_run_eval_parallel_cuda",
-            return_value={"mode": "cuda"},
-            create=True,
-        ) as run_cuda_parallel:
-            result = eval_vla._run_eval_parallel(cfg)
-
-        self.assertEqual(result, {"mode": "cuda"})
-        run_cpu_parallel.assert_not_called()
-        run_cuda_parallel.assert_called_once_with(cfg)
-
-    def test_resolve_cuda_devices_defaults_to_single_logical_gpu(self):
-        cfg = OmegaConf.create(
-            {
-                "device": "cuda",
-                "cuda_devices": None,
-            }
-        )
-
-        self.assertEqual(eval_vla._resolve_cuda_devices(cfg), [0])
-
-    def test_resolve_cuda_devices_rejects_empty_selection(self):
-        cfg = OmegaConf.create(
-            {
-                "device": "cuda",
-                "cuda_devices": [],
-            }
-        )
-
-        with self.assertRaisesRegex(ValueError, "cuda_devices"):
-            eval_vla._resolve_cuda_devices(cfg)
-
-    def test_run_eval_uses_air_insert_sampler_for_socket_peg_task(self):
-        self.assertTrue(
-            hasattr(eval_vla, "sample_air_insert_socket_peg_state"),
-            "Expected eval_vla to expose the new socket/peg reset sampler",
-        )
-
-        fake_env = _FakeEnv()
-        fake_agent = _FakeAgent()
-        sampled_task_state = {
-            "socket_pos": np.array([-0.10, 0.80, 0.47], dtype=np.float32),
-            "socket_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32),
-            "peg_pos": np.array([0.10, 0.82, 0.47], dtype=np.float32),
-            "peg_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32),
-        }
-        cfg = OmegaConf.create(
-            {
-                "agent": {},
-                "eval": {
-                    "ckpt_path": "checkpoints/vla_model_best.pt",
-                    "num_episodes": 1,
-                    "max_timesteps": 1,
-                    "device": "cpu",
-                    "task_name": "sim_air_insert_socket_peg",
-                    "camera_names": ["front"],
-                    "use_smoothing": False,
-                    "smooth_alpha": 0.3,
-                    "verbose_action": False,
-                    "headless": True,
-                },
-            }
-        )
-
-        with mock.patch.object(
-            eval_vla,
-            "load_checkpoint",
-            return_value=(fake_agent, None),
-        ), mock.patch.object(
-            eval_vla,
-            "make_sim_env",
-            return_value=fake_env,
-        ) as make_env, mock.patch.object(
-            eval_vla,
-            "sample_air_insert_socket_peg_state",
-            return_value=sampled_task_state,
-        ) as socket_peg_sampler, mock.patch.object(
-            eval_vla,
-            "sample_transfer_pose",
-            side_effect=AssertionError("sample_transfer_pose should not be used for sim_air_insert_socket_peg"),
-        ), mock.patch.object(
-            eval_vla,
-            "execute_policy_action",
-        ) as execute_policy_action, mock.patch.object(
-            eval_vla,
-            "tqdm",
-            side_effect=lambda iterable, **kwargs: iterable,
-        ):
-            eval_vla._run_eval(cfg)
-
-        make_env.assert_called_once_with("sim_air_insert_socket_peg", headless=True)
-        socket_peg_sampler.assert_called_once_with()
-        execute_policy_action.assert_called_once()
-        self.assertEqual(fake_env.reset_calls, [sampled_task_state])
-
-
 if __name__ == "__main__":
     unittest.main()

tests/test_eval_vla_rollout_artifacts.py

@@ -102,8 +102,10 @@ class EvalVLARolloutArtifactsTest(unittest.TestCase):
         self.assertIn('artifact_dir', eval_cfg)
         self.assertFalse(eval_cfg.save_summary_json)
         self.assertFalse(eval_cfg.save_trajectory_npz)
+        self.assertFalse(eval_cfg.save_trajectory_image)
         self.assertFalse(eval_cfg.record_video)
         self.assertIsNone(eval_cfg.artifact_dir)
+        self.assertIsNone(eval_cfg.trajectory_image_camera_name)
         self.assertIsNone(eval_cfg.video_camera_name)
         self.assertEqual(eval_cfg.video_fps, 30)
 
@@ -133,6 +135,8 @@ class EvalVLARolloutArtifactsTest(unittest.TestCase):
                         'artifact_dir': tmpdir,
                         'save_summary_json': True,
                         'save_trajectory_npz': True,
+                        'save_trajectory_image': True,
+                        'trajectory_image_camera_name': 'front',
                         'record_video': True,
                         'video_camera_name': 'front',
                         'video_fps': 12,
@@ -176,12 +180,14 @@ class EvalVLARolloutArtifactsTest(unittest.TestCase):
             trajectory_path = Path(artifacts['trajectory_npz'])
             summary_path = Path(artifacts['summary_json'])
             video_path = Path(artifacts['video_mp4'])
+            trajectory_image_path = Path(summary['episodes'][0]['artifact_paths']['trajectory_image'])
 
             self.assertEqual(Path(artifacts['output_dir']), Path(tmpdir))
             self.assertEqual(artifacts['video_camera_name'], 'front')
             self.assertTrue(trajectory_path.exists())
             self.assertTrue(summary_path.exists())
             self.assertTrue(video_path.exists())
+            self.assertTrue(trajectory_image_path.exists())
 
             rollout_npz = np.load(trajectory_path)
             np.testing.assert_array_equal(rollout_npz['episode_index'], np.array([0, 0]))
@@ -218,267 +224,120 @@ class EvalVLARolloutArtifactsTest(unittest.TestCase):
                 saved_summary = json.load(fh)
             self.assertEqual(saved_summary['artifacts']['trajectory_npz'], str(trajectory_path))
             self.assertEqual(saved_summary['artifacts']['video_mp4'], str(video_path))
+            self.assertEqual(
+                saved_summary['episodes'][0]['artifact_paths']['trajectory_image'],
+                str(trajectory_image_path),
+            )
             self.assertEqual(saved_summary['episode_rewards'], [3.0])
             self.assertAlmostEqual(summary['avg_reward'], 3.0)
             self.assertIn('avg_obs_read_time_ms', summary)
             self.assertIn('avg_env_step_time_ms', summary)
 
-    def test_run_eval_parallel_rejects_trajectory_and_video_exports(self):
+    def test_run_eval_exports_front_trajectory_images_without_video_dependency(self):
-        unsupported_flags = [
+        actions = [
-            "record_video",
+            np.arange(16, dtype=np.float32),
-            "save_trajectory",
+            np.arange(16, dtype=np.float32) + 10.0,
-            "save_trajectory_npz",
+            np.arange(16, dtype=np.float32) + 100.0,
+            np.arange(16, dtype=np.float32) + 110.0,
         ]
+        fake_agent = _FakeAgent(actions)
+        fake_env = _FakeEnv()
 
-        for flag_name in unsupported_flags:
+        with tempfile.TemporaryDirectory() as tmpdir:
-            with self.subTest(flag_name=flag_name):
+            cfg = OmegaConf.create(
-                cfg = OmegaConf.create(
+                {
+                    'agent': {},
+                    'eval': {
+                        'ckpt_path': 'checkpoints/vla_model_best.pt',
+                        'num_episodes': 2,
+                        'max_timesteps': 2,
+                        'device': 'cpu',
+                        'task_name': 'sim_transfer',
+                        'camera_names': ['top', 'front'],
+                        'use_smoothing': True,
+                        'smooth_alpha': 0.5,
+                        'verbose_action': False,
+                        'headless': True,
+                        'artifact_dir': tmpdir,
+                        'save_trajectory_image': True,
+                        'record_video': False,
+                    },
+                }
+            )
+
+            trajectory_image_calls = []
+
+            def fake_save_rollout_trajectory_image(
+                env,
+                output_path,
+                raw_actions,
+                camera_name,
+                *,
+                line_radius=0.004,
+                max_markers=1500,
+            ):
+                del env, line_radius, max_markers
+                trajectory_image_calls.append(
                     {
-                        "agent": {},
+                        'output_path': output_path,
-                        "eval": {
+                        'camera_name': camera_name,
-                            "ckpt_path": "checkpoints/vla_model_best.pt",
+                        'raw_actions': [np.array(action, copy=True) for action in raw_actions],
-                            "num_episodes": 2,
-                            "num_workers": 2,
-                            "max_timesteps": 1,
-                            "device": "cpu",
-                            "task_name": "sim_transfer",
-                            "camera_names": ["front"],
-                            "use_smoothing": False,
-                            "smooth_alpha": 0.3,
-                            "verbose_action": False,
-                            "headless": True,
-                            "save_artifacts": True,
-                            flag_name: True,
-                        },
                     }
                 )
-
+                if output_path is None:
-                with self.assertRaisesRegex(ValueError, flag_name):
+                    return None
-                    eval_vla._run_eval_parallel(cfg)
+                output_path = Path(output_path)
-
+                output_path.parent.mkdir(parents=True, exist_ok=True)
-    def test_run_eval_parallel_writes_merged_summary_timing_and_worker_dirs(self):
+                output_path.write_bytes(b'fake-png')
-        with tempfile.TemporaryDirectory() as tmpdir:
+                return str(output_path)
-            cfg = OmegaConf.create(
-                {
-                    "agent": {},
-                    "eval": {
-                        "ckpt_path": "checkpoints/vla_model_best.pt",
-                        "num_episodes": 3,
-                        "num_workers": 2,
-                        "max_timesteps": 1,
-                        "device": "cpu",
-                        "task_name": "sim_transfer",
-                        "camera_names": ["front"],
-                        "use_smoothing": False,
-                        "smooth_alpha": 0.3,
-                        "verbose_action": False,
-                        "headless": True,
-                        "artifact_dir": tmpdir,
-                        "save_summary_json": True,
-                        "save_timing": True,
-                    },
-                }
-            )
-
-            def fake_run_spawn_jobs(payloads, max_workers, worker_fn):
-                del max_workers, worker_fn
-                return [
-                    {
-                        "episodes": [
-                            {
-                                "episode_index": 2,
-                                "episode_reward": 3.0,
-                                "episode_max_reward": 3.0,
-                                "inference_fps": 30.0,
-                                "control_fps": 15.0,
-                            }
-                        ],
-                        "_merge_state": {
-                            "obs_read_time_ms": [3.0],
-                            "preprocess_time_ms": [1.0],
-                            "inference_time_ms": [2.0],
-                            "env_step_time_ms": [4.0],
-                            "total_time_ms": [5.0],
-                            "model_forward_flags": [True],
-                        },
-                    },
-                    {
-                        "episodes": [
-                            {
-                                "episode_index": 1,
-                                "episode_reward": 2.0,
-                                "episode_max_reward": 2.0,
-                                "inference_fps": 20.0,
-                                "control_fps": 10.0,
-                            },
-                            {
-                                "episode_index": 0,
-                                "episode_reward": 1.0,
-                                "episode_max_reward": 1.0,
-                                "inference_fps": 10.0,
-                                "control_fps": 5.0,
-                            },
-                        ],
-                        "_merge_state": {
-                            "obs_read_time_ms": [1.0, 2.0],
-                            "preprocess_time_ms": [1.0, 1.0],
-                            "inference_time_ms": [2.0, 2.0],
-                            "env_step_time_ms": [4.0, 4.0],
-                            "total_time_ms": [5.0, 5.0],
-                            "model_forward_flags": [False, True],
-                        },
-                    },
-                ]
 
             with mock.patch.object(
                 eval_vla,
-                "sample_transfer_pose",
+                'load_checkpoint',
-                side_effect=[
+                return_value=(fake_agent, None),
-                    np.array([0.1, 0.2, 0.3], dtype=np.float32),
-                    np.array([0.4, 0.5, 0.6], dtype=np.float32),
-                    np.array([0.7, 0.8, 0.9], dtype=np.float32),
-                ],
             ), mock.patch.object(
                 eval_vla,
-                "_run_spawn_jobs",
+                'make_sim_env',
-                side_effect=fake_run_spawn_jobs,
+                return_value=fake_env,
-            ):
+            ), mock.patch.object(
-                summary = eval_vla._run_eval_parallel(cfg)
+                eval_vla,
+                'sample_transfer_pose',
+                return_value=np.array([0.1, 0.2, 0.3], dtype=np.float32),
+            ), mock.patch.object(
+                eval_vla,
+                'tqdm',
+                side_effect=lambda iterable, **kwargs: iterable,
+            ), mock.patch.object(
+                eval_vla,
+                '_save_rollout_trajectory_image',
+                side_effect=fake_save_rollout_trajectory_image,
+            ) as save_trajectory_image_mock, mock.patch.object(
+                eval_vla,
+                '_open_video_writer',
+            ) as open_video_writer_mock:
+                summary = eval_vla._run_eval(cfg)
 
-            summary_path = Path(tmpdir) / "rollout_summary.json"
+            self.assertEqual(save_trajectory_image_mock.call_count, 2)
-            timing_path = Path(tmpdir) / "timing.json"
+            open_video_writer_mock.assert_not_called()
-            worker_00_dir = Path(tmpdir) / "workers" / "worker_00"
+            self.assertIsNone(summary['artifacts']['video_mp4'])
-            worker_01_dir = Path(tmpdir) / "workers" / "worker_01"
+            self.assertEqual(summary['artifacts']['trajectory_image_camera_name'], 'front')
-
+            self.assertEqual(
-            self.assertTrue(summary_path.exists())
+                [call['camera_name'] for call in trajectory_image_calls],
-            self.assertTrue(timing_path.exists())
+                ['front', 'front'],
-            self.assertTrue(worker_00_dir.is_dir())
-            self.assertTrue(worker_01_dir.is_dir())
-            self.assertEqual(summary["episode_rewards"], [1.0, 2.0, 3.0])
-
-            with summary_path.open("r", encoding="utf-8") as fh:
-                saved_summary = json.load(fh)
-            with timing_path.open("r", encoding="utf-8") as fh:
-                saved_timing = json.load(fh)
-
-            self.assertEqual(saved_summary["episode_rewards"], [1.0, 2.0, 3.0])
-            self.assertEqual(saved_summary["artifact_dir"], tmpdir)
-            self.assertEqual(saved_timing["count"], 3)
-            self.assertEqual(saved_timing["model_forward_count"], 2)
-
-    def test_run_eval_parallel_cuda_writes_merged_summary_timing_and_worker_dirs(self):
-        with tempfile.TemporaryDirectory() as tmpdir:
-            cfg = OmegaConf.create(
-                {
-                    "agent": {},
-                    "eval": {
-                        "ckpt_path": "checkpoints/vla_model_best.pt",
-                        "num_episodes": 3,
-                        "num_workers": 2,
-                        "cuda_devices": [0],
-                        "max_timesteps": 1,
-                        "device": "cuda",
-                        "task_name": "sim_transfer",
-                        "camera_names": ["front"],
-                        "use_smoothing": False,
-                        "smooth_alpha": 0.3,
-                        "verbose_action": False,
-                        "headless": True,
-                        "artifact_dir": tmpdir,
-                        "save_summary_json": True,
-                        "save_timing": True,
-                    },
-                }
             )
 
-            def fake_run_cuda_parallel_processes(server_payloads, worker_payloads):
+            first_episode_path = Path(summary['episodes'][0]['artifact_paths']['trajectory_image'])
-                self.assertEqual(len(server_payloads), 1)
+            second_episode_path = Path(summary['episodes'][1]['artifact_paths']['trajectory_image'])
-                self.assertEqual(server_payloads[0]["device_index"], 0)
+            self.assertTrue(first_episode_path.exists())
-                self.assertEqual([payload["server_index"] for payload in worker_payloads], [0, 0])
+            self.assertTrue(second_episode_path.exists())
-                return [
+            self.assertNotEqual(first_episode_path, second_episode_path)
-                    {
+            self.assertEqual(first_episode_path.parent, Path(tmpdir))
-                        "episodes": [
+            self.assertEqual(second_episode_path.parent, Path(tmpdir))
-                            {
-                                "episode_index": 2,
-                                "episode_reward": 3.0,
-                                "episode_max_reward": 3.0,
-                                "inference_fps": 30.0,
-                                "control_fps": 15.0,
-                            }
-                        ],
-                        "_merge_state": {
-                            "obs_read_time_ms": [3.0],
-                            "preprocess_time_ms": [1.0],
-                            "inference_time_ms": [2.0],
-                            "env_step_time_ms": [4.0],
-                            "total_time_ms": [5.0],
-                            "model_forward_flags": [True],
-                        },
-                    },
-                    {
-                        "episodes": [
-                            {
-                                "episode_index": 1,
-                                "episode_reward": 2.0,
-                                "episode_max_reward": 2.0,
-                                "inference_fps": 20.0,
-                                "control_fps": 10.0,
-                            },
-                            {
-                                "episode_index": 0,
-                                "episode_reward": 1.0,
-                                "episode_max_reward": 1.0,
-                                "inference_fps": 10.0,
-                                "control_fps": 5.0,
-                            },
-                        ],
-                        "_merge_state": {
-                            "obs_read_time_ms": [1.0, 2.0],
-                            "preprocess_time_ms": [1.0, 1.0],
-                            "inference_time_ms": [2.0, 2.0],
-                            "env_step_time_ms": [4.0, 4.0],
-                            "total_time_ms": [5.0, 5.0],
-                            "model_forward_flags": [False, True],
-                        },
-                    },
-                ]
 
-            with mock.patch.object(
+            np.testing.assert_array_equal(trajectory_image_calls[0]['raw_actions'][0], actions[0])
-                eval_vla,
+            np.testing.assert_array_equal(trajectory_image_calls[0]['raw_actions'][1], actions[1])
-                "sample_transfer_pose",
+            np.testing.assert_array_equal(trajectory_image_calls[1]['raw_actions'][0], actions[2])
-                side_effect=[
+            np.testing.assert_array_equal(trajectory_image_calls[1]['raw_actions'][1], actions[3])
-                    np.array([0.1, 0.2, 0.3], dtype=np.float32),
-                    np.array([0.4, 0.5, 0.6], dtype=np.float32),
-                    np.array([0.7, 0.8, 0.9], dtype=np.float32),
-                ],
-            ), mock.patch.object(
-                eval_vla,
-                "_run_cuda_parallel_processes",
-                side_effect=fake_run_cuda_parallel_processes,
-                create=True,
-            ):
-                summary = eval_vla._run_eval_parallel_cuda(cfg)
-
-            summary_path = Path(tmpdir) / "rollout_summary.json"
-            timing_path = Path(tmpdir) / "timing.json"
-            worker_00_dir = Path(tmpdir) / "workers" / "worker_00"
-            worker_01_dir = Path(tmpdir) / "workers" / "worker_01"
-
-            self.assertTrue(summary_path.exists())
-            self.assertTrue(timing_path.exists())
-            self.assertTrue(worker_00_dir.is_dir())
-            self.assertTrue(worker_01_dir.is_dir())
-            self.assertEqual(summary["episode_rewards"], [1.0, 2.0, 3.0])
-
-            with summary_path.open("r", encoding="utf-8") as fh:
-                saved_summary = json.load(fh)
-            with timing_path.open("r", encoding="utf-8") as fh:
-                saved_timing = json.load(fh)
-
-            self.assertEqual(saved_summary["episode_rewards"], [1.0, 2.0, 3.0])
-            self.assertEqual(saved_summary["artifact_dir"], tmpdir)
-            self.assertEqual(saved_timing["count"], 3)
-            self.assertEqual(saved_timing["model_forward_count"], 2)
 
 
 if __name__ == '__main__':

tests/test_robot_asset_paths.py

@@ -4,7 +4,7 @@ import unittest
 from pathlib import Path
 from unittest import mock
 
-from roboimi.assets.robots 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()

tests/test_train_vla_rollout_validation.py

@@ -158,101 +158,6 @@ class TrainVLARolloutValidationTest(unittest.TestCase):
         self.assertGreater(float(cfg.train.lr), 5e-5)
         self.assertGreater(cfg.train.num_workers, 8)
         self.assertEqual(cfg.train.rollout_val_freq_epochs, 50)
-        self.assertEqual(cfg.train.rollout_device, cfg.train.device)
-        self.assertIsNone(cfg.train.rollout_num_workers)
-        self.assertIsNone(cfg.train.rollout_cuda_devices)
-
-
-    def test_run_training_rollout_validation_propagates_gpu_parallel_settings(self):
-        cfg = OmegaConf.create(
-            {
-                'train': {
-                    'device': 'cpu',
-                    'batch_size': 1,
-                    'num_workers': 0,
-                    'val_split': 0.0,
-                    'seed': 0,
-                    'lr': 1e-3,
-                    'max_steps': 2,
-                    'log_freq': 1,
-                    'save_freq': 1000,
-                    'warmup_steps': 1,
-                    'scheduler_type': 'constant',
-                    'min_lr': 0.0,
-                    'grad_clip': 1.0,
-                    'weight_decay': 0.0,
-                    'pretrained_ckpt': None,
-                    'resume_ckpt': None,
-                    'use_swanlab': False,
-                    'rollout_val_freq_epochs': 2,
-                    'rollout_num_episodes': 5,
-                    'rollout_device': 'cuda',
-                    'rollout_num_workers': 4,
-                    'rollout_cuda_devices': [0, 1],
-                    'rollout_response_timeout_s': 123.0,
-                    'rollout_server_startup_timeout_s': 456.0,
-                },
-                'data': {
-                    'camera_names': ['front'],
-                },
-                'agent': {
-                    '_target_': 'fake.agent',
-                },
-                'eval': {
-                    'ckpt_path': 'unused.pt',
-                    'num_episodes': 99,
-                    'max_timesteps': 1,
-                    'device': 'cpu',
-                    'task_name': 'sim_transfer',
-                    'camera_names': ['front'],
-                    'use_smoothing': False,
-                    'smooth_alpha': 0.3,
-                    'verbose_action': False,
-                    'headless': False,
-                },
-            }
-        )
-        rollout_mock = mock.Mock(return_value={'avg_reward': 1.0})
-
-        def fake_instantiate(config_node, **_kwargs):
-            if config_node is cfg.data:
-                return _FakeDataset()
-            if config_node is cfg.agent:
-                return _FakeAgent()
-            raise AssertionError(f'unexpected instantiate config: {config_node!r}')
-
-        def fake_dataloader(_dataset, *, shuffle, **_kwargs):
-            del shuffle, _kwargs
-            return _FakeLoader(
-                {
-                    'observation.front': torch.zeros(1, 3, 2, 2),
-                    'observation.state': torch.zeros(1, 4),
-                    'action': torch.zeros(1, 2),
-                    'action_is_pad': torch.zeros(1, 1, dtype=torch.bool),
-                },
-                length=1,
-            )
-
-        with tempfile.TemporaryDirectory() as tempdir:
-            previous_cwd = os.getcwd()
-            try:
-                os.chdir(tempdir)
-                with mock.patch.object(train_vla, 'instantiate', side_effect=fake_instantiate),                      mock.patch.object(train_vla, 'DataLoader', side_effect=fake_dataloader),                      mock.patch.object(train_vla, 'build_training_optimizer', return_value=_FakeOptimizer(cfg.train.lr)),                      mock.patch.object(train_vla, 'get_lr_schedule_with_warmup', return_value=_FakeScheduler()),                      mock.patch.object(train_vla, 'tqdm', side_effect=lambda iterable, **kwargs: _FakeProgressBar(iterable)),                      mock.patch.object(train_vla.torch, 'save', return_value=None),                      mock.patch.object(eval_vla, '_run_eval', rollout_mock, create=True):
-                    train_vla._run_training(cfg)
-            finally:
-                os.chdir(previous_cwd)
-
-        rollout_cfg = rollout_mock.call_args.args[0]
-        self.assertEqual(rollout_cfg.eval.device, 'cuda')
-        self.assertEqual(rollout_cfg.eval.num_workers, 4)
-        self.assertEqual(list(rollout_cfg.eval.cuda_devices), [0, 1])
-        self.assertEqual(float(rollout_cfg.eval.response_timeout_s), 123.0)
-        self.assertEqual(float(rollout_cfg.eval.server_startup_timeout_s), 456.0)
-        self.assertTrue(rollout_cfg.eval.headless)
-        self.assertEqual(rollout_cfg.eval.num_episodes, 5)
-        self.assertFalse(rollout_cfg.eval.record_video)
-        self.assertTrue(rollout_cfg.eval.save_summary_json)
-        self.assertTrue(rollout_cfg.eval.save_trajectory_image)
 
     def test_training_passes_backbone_image_resize_override_to_dataset_instantiation(self):
         cfg = OmegaConf.create(