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

roboimi/demos/vla_scripts/eval_vla.py

@@ -1497,7 +1497,12 @@ def _build_parallel_worker_payloads(
         num_episodes=int(eval_cfg.num_episodes),
         num_workers=requested_workers,
     )
-    box_poses = _plan_episode_box_poses(int(eval_cfg.num_episodes))
+    task_name = str(eval_cfg.get('task_name', ''))
+    box_poses = (
+        _plan_episode_box_poses(int(eval_cfg.num_episodes))
+        if 'sim_transfer' in task_name
+        else None
+    )
     resolved_cfg = OmegaConf.to_container(cfg, resolve=True)
     payloads = []
     workers_dir = None
@@ -1521,10 +1526,14 @@ def _build_parallel_worker_payloads(
             'worker_index': int(worker_index),
             'artifact_dir': str(worker_artifact_dir) if worker_artifact_dir is not None else None,
             'episode_plans': [
-                {
+                (
-                    'episode_index': int(episode_index),
+                    {
-                    'box_pos': np.asarray(box_poses[episode_index], dtype=np.float32).tolist(),
+                        'episode_index': int(episode_index),
-                }
+                        'box_pos': np.asarray(box_poses[episode_index], dtype=np.float32).tolist(),
+                    }
+                    if box_poses is not None
+                    else {'episode_index': int(episode_index)}
+                )
                 for episode_index in episode_indices
             ],
         })

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

@@ -445,6 +445,35 @@ class EvalVLAExecutionTest(unittest.TestCase):
         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,