From ff7f4a1b03c600719d7421ced6c3ac84fc058776 Mon Sep 17 00:00:00 2001 From: Logic Date: Fri, 31 Jul 2026 10:11:04 +0800 Subject: [PATCH] feat(vla): add ACT policy for socket peg --- .../plans/2026-05-08-act-socket-peg.md | 49 ++++ .../specs/2026-05-08-act-socket-peg-design.md | 78 ++++++ roboimi/demos/vla_scripts/eval_vla.py | 19 +- roboimi/vla/agent_act.py | 214 +++++++++++++++ roboimi/vla/conf/agent/act_resnet.yaml | 36 +++ roboimi/vla/models/heads/act.py | 257 ++++++++++++++++++ tests/test_act_agent.py | 249 +++++++++++++++++ tests/test_eval_vla_execution.py | 29 ++ 8 files changed, 926 insertions(+), 5 deletions(-) create mode 100644 docs/superpowers/plans/2026-05-08-act-socket-peg.md create mode 100644 docs/superpowers/specs/2026-05-08-act-socket-peg-design.md create mode 100644 roboimi/vla/agent_act.py create mode 100644 roboimi/vla/conf/agent/act_resnet.yaml create mode 100644 roboimi/vla/models/heads/act.py create mode 100644 tests/test_act_agent.py diff --git a/docs/superpowers/plans/2026-05-08-act-socket-peg.md b/docs/superpowers/plans/2026-05-08-act-socket-peg.md new file mode 100644 index 0000000..946161b --- /dev/null +++ b/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. diff --git a/docs/superpowers/specs/2026-05-08-act-socket-peg-design.md b/docs/superpowers/specs/2026-05-08-act-socket-peg-design.md new file mode 100644 index 0000000..b720a88 --- /dev/null +++ b/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. + diff --git a/roboimi/demos/vla_scripts/eval_vla.py b/roboimi/demos/vla_scripts/eval_vla.py index fc30df2..f6307fa 100644 --- a/roboimi/demos/vla_scripts/eval_vla.py +++ b/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 ], }) diff --git a/roboimi/vla/agent_act.py b/roboimi/vla/agent_act.py new file mode 100644 index 0000000..0581787 --- /dev/null +++ b/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() diff --git a/roboimi/vla/conf/agent/act_resnet.yaml b/roboimi/vla/conf/agent/act_resnet.yaml new file mode 100644 index 0000000..1775cbe --- /dev/null +++ b/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" diff --git a/roboimi/vla/models/heads/act.py b/roboimi/vla/models/heads/act.py new file mode 100644 index 0000000..86d47bf --- /dev/null +++ b/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 diff --git a/tests/test_act_agent.py b/tests/test_act_agent.py new file mode 100644 index 0000000..e94fe97 --- /dev/null +++ b/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() diff --git a/tests/test_eval_vla_execution.py b/tests/test_eval_vla_execution.py index f9d7a44..4526a7e 100644 --- a/tests/test_eval_vla_execution.py +++ b/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,