Logic
215 lines
9.7 KiB
Python
215 lines
9.7 KiB
Python
"""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()
|