feat(vla): add ACT policy for socket peg

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

@@ -0,0 +1,49 @@
+# 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-05-08-act-socket-peg-design.md

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

@@ -0,0 +1,200 @@
+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/demos/diana_air_insert_policy.py

@@ -12,12 +12,11 @@ class TestAirInsertPolicy(PolicyBase):
     LEGACY_GRASP_STRATEGY = "legacy"
     SOCKET_HOLD_Z = 0.85
     PEG_INSERT_START_OFFSET = np.array([0.105, 0.0, 0.0], dtype=np.float64)
-    INSERT_START_T = 650
-    INSERT_END_T = 730
+    INSERT_END_T = 580
     LEFT_SOCKET_GRIPPER_CLOSED = -100
     RIGHT_PEG_GRIPPER_CLOSED = -100
     SOCKET_APPROACH_Z = 1.05
-    EPISODE_END_T = 1000
+    EPISODE_END_T = 600
 
     def __init__(self, inject_noise=False, grasp_strategy=SOCKET_OUTER_GRASP_STRATEGY):
         super().__init__(inject_noise=inject_noise)
@@ -120,13 +119,7 @@ class TestAirInsertPolicy(PolicyBase):
                 "gripper": self.LEFT_SOCKET_GRIPPER_CLOSED,
             },
             {
-                "t": 450,
-                "xyz": socket_hold_action,
-                "quat": left_pick_quat,
-                "gripper": self.LEFT_SOCKET_GRIPPER_CLOSED,
-            },
-            {
-                "t": 750,
+                "t": 350,
                 "xyz": socket_hold_action,
                 "quat": left_pick_quat,
                 "gripper": self.LEFT_SOCKET_GRIPPER_CLOSED,
@@ -165,19 +158,13 @@ class TestAirInsertPolicy(PolicyBase):
                 "gripper": self.RIGHT_PEG_GRIPPER_CLOSED,
             },
             {
-                "t": 450,
+                "t": 350,
                 "xyz": peg_init_xyz,
                 "quat": right_pick_quat,
                 "gripper": self.RIGHT_PEG_GRIPPER_CLOSED,
             },
             {
-                "t": 550,
-                "xyz": peg_lift_center,
-                "quat": right_pick_quat,
-                "gripper": self.RIGHT_PEG_GRIPPER_CLOSED,
-            },
-            {
-                "t": self.INSERT_START_T,
+                "t": 450,
                 "xyz": peg_lift_center,
                 "quat": right_pick_quat,
                 "gripper": self.RIGHT_PEG_GRIPPER_CLOSED,
@@ -188,12 +175,6 @@ class TestAirInsertPolicy(PolicyBase):
                 "quat": right_pick_quat,
                 "gripper": self.RIGHT_PEG_GRIPPER_CLOSED,
             },
-            {
-                "t": 750,
-                "xyz": peg_insert_end_center,
-                "quat": right_pick_quat,
-                "gripper": self.RIGHT_PEG_GRIPPER_CLOSED,
-            },
             {
                 "t": self.EPISODE_END_T,
                 "xyz": peg_insert_end_center,

roboimi/demos/vla_scripts/train_vla.py

@@ -702,10 +702,28 @@ 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 = int(cfg.train.get('rollout_num_episodes', 1))
+            rollout_cfg.eval.num_episodes = rollout_num_episodes
+            rollout_cfg.eval.num_workers = rollout_num_workers
             rollout_cfg.eval.headless = True
-            rollout_cfg.eval.device = 'cpu'
+            rollout_cfg.eval.device = rollout_device
+            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
@@ -716,9 +734,11 @@ def _run_training(cfg: DictConfig):
             )
 
             log.info(
-                "🎯 开始 checkpoint rollout 验证: %s (episodes=%s, headless=True)",
+                "🎯 开始 checkpoint rollout 验证: %s (episodes=%s, device=%s, workers=%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_base.py

@@ -91,7 +91,6 @@ 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)
@@ -104,6 +103,7 @@ class DualDianaMed(MujocoEnv):
 
             super().step(action)
         self.base_time = time.time() - ctrl_cur_time
+        self.obs = self._get_obs()
             
 
     def preStep(self, action):

roboimi/vla/agent_act.py

@@ -0,0 +1,214 @@
+"""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

@@ -0,0 +1,36 @@
+# @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,6 +29,11 @@ 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,6 +2,10 @@
 # 评估配置
 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

@@ -0,0 +1,257 @@
+"""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

@@ -0,0 +1,249 @@
+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_eval_vla_execution.py

@@ -1,5 +1,11 @@
 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
 
 
@@ -14,6 +20,48 @@ 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()
@@ -23,6 +71,691 @@ 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

@@ -216,6 +216,31 @@ class EvalVLAHeadlessTest(unittest.TestCase):
         self.assertIsNotNone(env.top)
         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()
@@ -323,6 +348,193 @@ class EvalVLAHeadlessTest(unittest.TestCase):
         self.assertAlmostEqual(summary["avg_reward"], 3.75)
         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"),

tests/test_eval_vla_rollout_artifacts.py

@@ -102,10 +102,8 @@ 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)
 
@@ -135,8 +133,6 @@ 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,
@@ -180,14 +176,12 @@ 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]))
@@ -224,120 +218,267 @@ 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_exports_front_trajectory_images_without_video_dependency(self):
-        actions = [
-            np.arange(16, dtype=np.float32),
-            np.arange(16, dtype=np.float32) + 10.0,
-            np.arange(16, dtype=np.float32) + 100.0,
-            np.arange(16, dtype=np.float32) + 110.0,
+    def test_run_eval_parallel_rejects_trajectory_and_video_exports(self):
+        unsupported_flags = [
+            "record_video",
+            "save_trajectory",
+            "save_trajectory_npz",
         ]
-        fake_agent = _FakeAgent(actions)
-        fake_env = _FakeEnv()
 
-        with tempfile.TemporaryDirectory() as tmpdir:
+        for flag_name in unsupported_flags:
+            with self.subTest(flag_name=flag_name):
                 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,
+                        "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,
+                            "save_artifacts": True,
+                            flag_name: True,
                         },
                     }
                 )
 
-            trajectory_image_calls = []
+                with self.assertRaisesRegex(ValueError, flag_name):
+                    eval_vla._run_eval_parallel(cfg)
 
-            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(
+    def test_run_eval_parallel_writes_merged_summary_timing_and_worker_dirs(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            cfg = OmegaConf.create(
                 {
-                        'output_path': output_path,
-                        'camera_name': camera_name,
-                        'raw_actions': [np.array(action, copy=True) for action in raw_actions],
+                    "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,
+                    },
                 }
             )
-                if output_path is None:
-                    return None
-                output_path = Path(output_path)
-                output_path.parent.mkdir(parents=True, exist_ok=True)
-                output_path.write_bytes(b'fake-png')
-                return str(output_path)
+
+            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,
-                'load_checkpoint',
-                return_value=(fake_agent, None),
+                "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),
+                ],
             ), mock.patch.object(
                 eval_vla,
-                'make_sim_env',
-                return_value=fake_env,
-            ), mock.patch.object(
-                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)
+                "_run_spawn_jobs",
+                side_effect=fake_run_spawn_jobs,
+            ):
+                summary = eval_vla._run_eval_parallel(cfg)
 
-            self.assertEqual(save_trajectory_image_mock.call_count, 2)
-            open_video_writer_mock.assert_not_called()
-            self.assertIsNone(summary['artifacts']['video_mp4'])
-            self.assertEqual(summary['artifacts']['trajectory_image_camera_name'], 'front')
-            self.assertEqual(
-                [call['camera_name'] for call in trajectory_image_calls],
-                ['front', 'front'],
+            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)
+
+    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,
+                    },
+                }
             )
 
-            first_episode_path = Path(summary['episodes'][0]['artifact_paths']['trajectory_image'])
-            second_episode_path = Path(summary['episodes'][1]['artifact_paths']['trajectory_image'])
-            self.assertTrue(first_episode_path.exists())
-            self.assertTrue(second_episode_path.exists())
-            self.assertNotEqual(first_episode_path, second_episode_path)
-            self.assertEqual(first_episode_path.parent, Path(tmpdir))
-            self.assertEqual(second_episode_path.parent, Path(tmpdir))
+            def fake_run_cuda_parallel_processes(server_payloads, worker_payloads):
+                self.assertEqual(len(server_payloads), 1)
+                self.assertEqual(server_payloads[0]["device_index"], 0)
+                self.assertEqual([payload["server_index"] for payload in worker_payloads], [0, 0])
+                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],
+                        },
+                    },
+                ]
 
-            np.testing.assert_array_equal(trajectory_image_calls[0]['raw_actions'][0], actions[0])
-            np.testing.assert_array_equal(trajectory_image_calls[0]['raw_actions'][1], actions[1])
-            np.testing.assert_array_equal(trajectory_image_calls[1]['raw_actions'][0], actions[2])
-            np.testing.assert_array_equal(trajectory_image_calls[1]['raw_actions'][1], actions[3])
+            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),
+                ],
+            ), 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_train_vla_rollout_validation.py

@@ -158,6 +158,101 @@ 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(