diff --git a/roboimi/demos/vla_scripts/eval_vla.py b/roboimi/demos/vla_scripts/eval_vla.py index fc30df2..bc89378 100644 --- a/roboimi/demos/vla_scripts/eval_vla.py +++ b/roboimi/demos/vla_scripts/eval_vla.py @@ -113,6 +113,7 @@ def prepare_observation( obs: Dict, camera_names: list, image_resize_shape: Optional[tuple[int, int]] = (224, 224), + task_description: Optional[str] = None, ) -> Dict: """ 将环境观测转换为 agent 格式。 @@ -139,7 +140,32 @@ def prepare_observation( # 转换 qpos: numpy -> tensor qpos = torch.from_numpy(obs['qpos']).float() - return {'qpos': qpos, 'images': images} + observation = {'qpos': qpos, 'images': images} + if 'task' in obs: + observation['task'] = obs['task'] + elif task_description is not None: + observation['task'] = task_description + return observation + + +def _normalize_resize_shape(shape) -> Optional[tuple[int, int]]: + if shape is None: + return None + normalized = tuple(int(v) for v in shape) + if len(normalized) != 2: + raise ValueError(f'image resize shape must contain exactly two values, got {normalized}') + return normalized + + +def _resolve_eval_image_resize_shape(cfg: DictConfig) -> Optional[tuple[int, int]]: + image_resize_shape = cfg.get('data', {}).get('image_resize_shape', (224, 224)) + agent_cfg = cfg.agent + for backbone_key in ('vision_backbone', 'condition_encoder'): + backbone_cfg = agent_cfg.get(backbone_key, None) + if backbone_cfg is not None and 'eval_image_resize_shape' in backbone_cfg: + image_resize_shape = backbone_cfg.get('eval_image_resize_shape') + break + return _normalize_resize_shape(image_resize_shape) def _resolve_policy_camera_names(cfg: DictConfig) -> list[str]: @@ -159,6 +185,7 @@ def _new_local_policy_queues(obs_horizon: int) -> dict[str, deque]: return { 'qpos': deque(maxlen=int(obs_horizon)), 'images': deque(maxlen=int(obs_horizon)), + 'task': deque(maxlen=int(obs_horizon)), 'action': deque(), } @@ -174,6 +201,8 @@ def _populate_local_policy_queues( camera_name: image.detach().clone() for camera_name, image in observation['images'].items() }) + if 'task' in observation: + queues['task'].append(observation['task']) def _prepare_local_policy_batch( @@ -202,7 +231,10 @@ def _prepare_local_policy_batch( ).unsqueeze(0) for camera_name in ordered_camera_names } - return {'qpos': batch_qpos, 'images': batch_images} + batch = {'qpos': batch_qpos, 'images': batch_images} + if queues.get('task'): + batch['task'] = [list(queues['task'])[-1]] + return batch def _enqueue_predicted_actions( @@ -226,23 +258,29 @@ def _enqueue_predicted_actions( def _serialize_policy_batch(batch: Dict[str, torch.Tensor]) -> dict[str, Any]: - return { + serialized = { 'qpos': batch['qpos'].detach().cpu().numpy().astype(np.float32, copy=True), 'images': { camera_name: image.detach().cpu().numpy().astype(np.float32, copy=True) for camera_name, image in batch['images'].items() }, } + if 'task' in batch: + serialized['task'] = batch['task'] + return serialized def _deserialize_policy_batch(batch: dict[str, Any], device: str) -> Dict[str, torch.Tensor]: - return { + deserialized = { 'qpos': torch.as_tensor(batch['qpos'], dtype=torch.float32, device=device), 'images': { camera_name: torch.as_tensor(image, dtype=torch.float32, device=device) for camera_name, image in batch['images'].items() }, } + if 'task' in batch: + deserialized['task'] = batch['task'] + return deserialized class _LocalPolicyRunner: @@ -1015,6 +1053,8 @@ def _run_eval_episode_plans( eval_cfg = cfg.eval device = str(eval_cfg.device) camera_names = list(eval_cfg.camera_names) + image_resize_shape = _resolve_eval_image_resize_shape(cfg) + task_description = eval_cfg.get('task_description', None) artifact_paths = artifact_paths or _resolve_artifact_paths(eval_cfg) video_recorder = _RolloutVideoRecorder( output_path=artifact_paths['video_mp4'], @@ -1089,7 +1129,12 @@ def _run_eval_episode_plans( video_recorder.write(video_frame) # 准备给 agent 的观测 - observation = prepare_observation(obs, camera_names) + observation = prepare_observation( + obs, + camera_names, + image_resize_shape=image_resize_shape, + task_description=task_description, + ) end_preprocess = time.perf_counter() # 选择动作（本地 agent 或远端 inference server） diff --git a/roboimi/demos/vla_scripts/train_vla.py b/roboimi/demos/vla_scripts/train_vla.py index 2607fa3..e7fac56 100644 --- a/roboimi/demos/vla_scripts/train_vla.py +++ b/roboimi/demos/vla_scripts/train_vla.py @@ -184,6 +184,65 @@ def get_lr_schedule_with_warmup(optimizer, warmup_steps, max_steps, scheduler_ty return LambdaLR(optimizer, lr_lambda) +def _instantiate_dataset(cfg, dataset_image_resize_shape, episode_indices=None): + kwargs = {'image_resize_shape': dataset_image_resize_shape} + if episode_indices is not None: + kwargs['episode_indices'] = episode_indices + return instantiate(cfg.data, **kwargs) + + +def build_train_val_datasets(cfg, dataset_image_resize_shape): + val_episode_indices = cfg.train.get('val_episode_indices', None) + if val_episode_indices: + dataset = _instantiate_dataset(cfg, dataset_image_resize_shape) + available_episode_indices = list(getattr(dataset, 'available_episode_indices', [])) + if not available_episode_indices: + raise ValueError('显式 val_episode_indices 需要数据集暴露 available_episode_indices') + + requested_val_episode_indices = sorted(int(idx) for idx in val_episode_indices) + available_set = set(available_episode_indices) + missing = sorted(set(requested_val_episode_indices) - available_set) + if missing: + raise ValueError( + f'val_episode_indices {missing} 不存在于数据集可用 episodes {available_episode_indices}' + ) + + val_set = set(requested_val_episode_indices) + train_episode_indices = [ + idx for idx in available_episode_indices + if idx not in val_set + ] + if not train_episode_indices: + raise ValueError('显式 val_episode_indices 不能覆盖全部 episodes，训练集将为空') + + train_dataset = _instantiate_dataset( + cfg, + dataset_image_resize_shape, + episode_indices=train_episode_indices, + ) + val_dataset = _instantiate_dataset( + cfg, + dataset_image_resize_shape, + episode_indices=requested_val_episode_indices, + ) + return dataset, train_dataset, val_dataset, requested_val_episode_indices + + dataset = _instantiate_dataset(cfg, dataset_image_resize_shape) + val_split = float(cfg.train.get('val_split', 0.1)) + seed = int(cfg.train.get('seed', 42)) + val_size = int(len(dataset) * val_split) + train_size = len(dataset) - val_size + if val_size > 0: + train_dataset, val_dataset = random_split( + dataset, + [train_size, val_size], + generator=torch.Generator().manual_seed(seed) + ) + else: + train_dataset, val_dataset = dataset, None + return dataset, train_dataset, val_dataset, None + + def build_training_optimizer(agent, lr, weight_decay): """为训练脚本构建优化器，优先复用任意 head 自带的参数分组。""" trainable_params = [param for param in agent.parameters() if param.requires_grad] @@ -369,6 +428,11 @@ def _run_training(cfg: DictConfig): _configure_cuda_runtime(cfg) swanlab_module = _init_swanlab(cfg) try: + action_mse_val_freq_epochs = int(cfg.train.get('action_mse_val_freq_epochs', 0) or 0) + explicit_val_episode_indices = cfg.train.get('val_episode_indices', None) + if action_mse_val_freq_epochs > 0 and not explicit_val_episode_indices: + raise ValueError('action_mse_val_freq_epochs > 0 requires train.val_episode_indices') + # 创建检查点目录 run_output_dir = _resolve_run_output_dir() checkpoint_dir = run_output_dir / "checkpoints" @@ -381,33 +445,35 @@ def _run_training(cfg: DictConfig): log.info("📦 加载数据集...") try: dataset_image_resize_shape = cfg.data.get('image_resize_shape', (224, 224)) - vision_backbone_cfg = cfg.agent.get('vision_backbone', None) - if vision_backbone_cfg is not None and 'dataset_image_resize_shape' in vision_backbone_cfg: - dataset_image_resize_shape = vision_backbone_cfg.get('dataset_image_resize_shape') - dataset = instantiate( - cfg.data, - image_resize_shape=dataset_image_resize_shape, + for backbone_key in ('vision_backbone', 'condition_encoder'): + backbone_cfg = cfg.agent.get(backbone_key, None) + if backbone_cfg is not None and 'dataset_image_resize_shape' in backbone_cfg: + dataset_image_resize_shape = backbone_cfg.get('dataset_image_resize_shape') + break + dataset, train_dataset, val_dataset, explicit_val_episode_indices = ( + build_train_val_datasets(cfg, dataset_image_resize_shape) ) log.info(f"✅ 数据集加载成功。总样本数: {len(dataset)}") except Exception as e: log.error(f"❌ 数据集加载失败: {e}") raise - # 训练/验证集划分 - val_split = float(cfg.train.get('val_split', 0.1)) - seed = int(cfg.train.get('seed', 42)) - val_size = int(len(dataset) * val_split) - train_size = len(dataset) - val_size - if val_size > 0: - train_dataset, val_dataset = random_split( - dataset, - [train_size, val_size], - generator=torch.Generator().manual_seed(seed) + if explicit_val_episode_indices is not None: + log.info( + "✅ 数据集划分: 训练集=%s, 验证集=%s (显式 held-out episodes=%s)", + len(train_dataset), + len(val_dataset), + explicit_val_episode_indices, ) - log.info(f"✅ 数据集划分: 训练集={train_size}, 验证集={val_size} (验证比例={val_split})") else: - train_dataset, val_dataset = dataset, None - log.info("✅ 数据集划分: 全部用于训练, 验证集=0 (验证比例=0)") + val_split = float(cfg.train.get('val_split', 0.1)) + val_size = len(val_dataset) if val_dataset is not None else 0 + if val_size > 0: + log.info( + f"✅ 数据集划分: 训练集={len(train_dataset)}, 验证集={val_size} (验证比例={val_split})" + ) + else: + log.info("✅ 数据集划分: 全部用于训练, 验证集=0 (验证比例=0)") train_batch_size = int(cfg.train.batch_size) train_drop_last = len(train_dataset) >= train_batch_size @@ -652,12 +718,16 @@ def _run_training(cfg: DictConfig): if key in batch_data: images[cam_name] = batch_data[key] - return { + agent_input = { 'images': images, 'qpos': batch_data['observation.state'], # SimpleRobotDataset 使用 observation.state 'action': batch_data['action'], 'action_is_pad': batch_data.get('action_is_pad', None) # 传递padding mask } + if 'task' in batch_data: + agent_input['task'] = batch_data['task'] + + return agent_input def save_checkpoint(checkpoint_path: Path, step: int, loss_value, val_loss=None, rollout_avg_reward=None): agent_stats = agent.get_normalization_stats() @@ -904,6 +974,31 @@ def _run_training(cfg: DictConfig): completed_steps // steps_per_epoch if steps_per_epoch > 0 else 0 ) + should_run_action_mse_val = ( + val_loader is not None + and explicit_val_episode_indices is not None + and action_mse_val_freq_epochs > 0 + and steps_per_epoch > 0 + and completed_steps % steps_per_epoch == 0 + and completed_epoch > 0 + and completed_epoch % action_mse_val_freq_epochs == 0 + ) + if should_run_action_mse_val: + val_loss = run_validation() + if val_loss is not None: + log.info( + f"步骤 {step}/{cfg.train.max_steps} | Epoch {completed_epoch} " + f"held-out action MSE: {val_loss:.6f}" + ) + _log_to_swanlab( + swanlab_module, + { + 'val/action_mse': val_loss, + 'val/epoch': completed_epoch, + }, + step=step, + ) + should_run_epoch_rollout = ( rollout_validation_enabled and steps_per_epoch > 0 diff --git a/roboimi/vla/agent_smolvla_conditioned.py b/roboimi/vla/agent_smolvla_conditioned.py new file mode 100644 index 0000000..bf7d078 --- /dev/null +++ b/roboimi/vla/agent_smolvla_conditioned.py @@ -0,0 +1,277 @@ +from __future__ import annotations + +from collections import deque +from typing import Dict, Optional, Sequence + +import torch +import torch.nn as nn + +from roboimi.vla.agent_imf import IMFVLAAgent +from roboimi.vla.models.normalization import NormalizationModule + + +class SmolVLAIMFAttnResAgent(IMFVLAAgent): + """IMF-AttnRes action expert conditioned by SmolVLA-style VLM prefix tokens. + + Unlike ``VLAAgent`` this agent does not concatenate ResNet features and raw + state at every observation step. A ``condition_encoder`` receives images, + normalized state, and variable task language, and returns a condition token + sequence `(B, S, D)` that is passed directly to the IMF head. + """ + + def __init__( + self, + condition_encoder: nn.Module, + action_encoder, + head, + action_dim: int, + obs_dim: int, + pred_horizon: int = 16, + obs_horizon: int = 2, + diffusion_steps: int = 100, + inference_steps: int = 1, + num_cams: int = 3, + camera_names: Optional[Sequence[str]] = None, + dataset_stats=None, + normalization_type: str = 'min_max', + num_action_steps: int = 8, + head_type: str = 'transformer', + condition_dim: int | None = None, + condition_sequence_length: int | None = None, + task_description: str | None = None, + **kwargs, + ) -> None: + # Intentionally bypass VLAAgent.__init__; its condition dimensions are + # tied to ResNet-style visual+state concatenation. + nn.Module.__init__(self) + if inference_steps != 1: + raise ValueError( + 'SmolVLAIMFAttnResAgent only supports one-step IMF inference; ' + f'inference_steps must be 1, got {inference_steps}.' + ) + if head_type != 'transformer': + raise ValueError(f'SmolVLAIMFAttnResAgent requires head_type="transformer", got {head_type!r}') + del diffusion_steps, kwargs + + 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.inference_steps = 1 + self.head_type = head_type + self.camera_names = tuple(camera_names) if camera_names is not None else None + if self.camera_names is not None and len(self.camera_names) != self.num_cams: + raise ValueError(f'camera_names length({len(self.camera_names)}) does not match num_cams({self.num_cams})') + + self.normalization = NormalizationModule(stats=dataset_stats, normalization_type=normalization_type) + self.condition_encoder = condition_encoder + # Compatibility aliases used by some utilities/tests. + self.vision_encoder = condition_encoder + self.action_encoder = action_encoder + self.state_encoder = None + self.task_description = task_description + + encoder_dim = getattr(condition_encoder, 'output_dim', None) + if encoder_dim is None: + encoder_dim = getattr(condition_encoder, 'joint_output_dim', None) + if condition_dim is None: + if encoder_dim is None: + raise ValueError('condition_dim must be provided when condition_encoder has no output_dim') + condition_dim = int(encoder_dim) + self.per_step_cond_dim = int(condition_dim) + self.raw_per_step_cond_dim = self.per_step_cond_dim + + encoder_seq_len = getattr(condition_encoder, 'condition_sequence_length', None) + if condition_sequence_length is None: + if encoder_seq_len is None: + raise ValueError( + 'condition_sequence_length must be provided when condition_encoder has no condition_sequence_length' + ) + condition_sequence_length = int(encoder_seq_len) + self.condition_sequence_length = int(condition_sequence_length) + self.condition_tokens_per_step = self.condition_sequence_length + self.global_cond_dim = self.per_step_cond_dim * self.condition_sequence_length + + if isinstance(head, nn.Module): + self.noise_pred_net = head + else: + self.noise_pred_net = head( + input_dim=self.action_dim, + output_dim=self.action_dim, + horizon=self.pred_horizon, + n_obs_steps=self.condition_sequence_length, + cond_dim=self.per_step_cond_dim, + ) + 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: + names = tuple(sorted(images.keys())) + if len(names) != self.num_cams: + raise ValueError(f'image camera count({len(names)}) does not match num_cams({self.num_cams})') + return {name: images[name] for name in names} + missing = [name for name in self.camera_names if name not in images] + if missing: + raise ValueError(f'image condition missing required cameras. missing={missing}, expected={list(self.camera_names)}') + return {name: images[name] for name in self.camera_names} + + def _resolve_task(self, task, batch_size: int): + def is_missing_task(item) -> bool: + return item is None or (isinstance(item, str) and item.strip().lower() in {'', 'unknown'}) + + def fallback_or(item): + if self.task_description is not None and is_missing_task(item): + return self.task_description + if item is None: + return '' + return item + + if task is None: + task = self.task_description + if task is None: + return [''] * batch_size + if isinstance(task, str): + return [fallback_or(task)] * batch_size + task = list(task) + if len(task) != batch_size: + raise ValueError(f'task batch size ({len(task)}) must match batch size ({batch_size})') + return [fallback_or(item) for item in task] + + def _build_cond(self, images: Dict[str, torch.Tensor], states: torch.Tensor, task=None) -> torch.Tensor: + ordered_images = self._order_images(images) + batch_size = states.shape[0] + tasks = self._resolve_task(task, batch_size=batch_size) + cond = self.condition_encoder(ordered_images, state=states, task=tasks) + if cond.ndim != 3: + raise RuntimeError(f'condition_encoder must return (B,S,D), got {tuple(cond.shape)}') + if cond.shape[0] != batch_size: + raise RuntimeError(f'condition batch mismatch: got {cond.shape[0]}, expected {batch_size}') + if cond.shape[1] != self.condition_sequence_length: + raise RuntimeError( + f'condition sequence length mismatch: got {cond.shape[1]}, expected {self.condition_sequence_length}' + ) + if cond.shape[2] != self.per_step_cond_dim: + raise RuntimeError(f'condition dim mismatch: got {cond.shape[2]}, expected {self.per_step_cond_dim}') + head_dtype = next(self.noise_pred_net.parameters()).dtype + return cond.to(dtype=head_dtype) + + def compute_loss(self, batch): + actions, states, images = batch['action'], batch['qpos'], batch['images'] + action_is_pad = batch.get('action_is_pad', None) + batch_size = actions.shape[0] + + states = self.normalization.normalize_qpos(states) + actions = self.normalization.normalize_action(actions) + cond = self._build_cond(images, states, task=batch.get('task', None)) + + x = actions + e = torch.randn_like(x) + t = torch.rand(batch_size, device=x.device, dtype=x.dtype) + r = torch.rand(batch_size, device=x.device, dtype=x.dtype) + t, r = torch.maximum(t, r), torch.minimum(t, r) + + t_broadcast = self._broadcast_batch_time(t, x) + z_t = (1 - t_broadcast) * x + t_broadcast * e + + v = self.fn(z_t, t, t, cond=cond) + u, du_dt = self._compute_u_and_du_dt(z_t, r, t, cond=cond, v=v) + V = self._compound_velocity(u, du_dt, r, t) + target = e - x + + loss = nn.functional.mse_loss(V, target, reduction='none') + if action_is_pad is not None: + mask = (~action_is_pad).unsqueeze(-1).to(loss.dtype) + valid_count = mask.sum() * loss.shape[-1] + loss = (loss * mask).sum() / valid_count.clamp_min(1.0) + else: + loss = loss.mean() + return loss + + @torch.no_grad() + def predict_action(self, images, proprioception, task=None): + batch_size = proprioception.shape[0] + proprioception = self.normalization.normalize_qpos(proprioception) + cond = self._build_cond(images, proprioception, task=task) + z_t = torch.randn((batch_size, self.pred_horizon, self.action_dim), device=cond.device, dtype=cond.dtype) + action = self._sample_one_step(z_t, cond=cond) + return self.normalization.denormalize_action(action) + + @torch.no_grad() + def predict_action_chunk(self, batch: Dict[str, torch.Tensor]) -> torch.Tensor: + return self.predict_action(batch['images'], batch['qpos'], task=batch.get('task', None)) + + def reset(self): + self._queues = { + 'qpos': deque(maxlen=self.obs_horizon), + 'images': deque(maxlen=self.obs_horizon), + 'task': deque(maxlen=self.obs_horizon), + 'action': deque(maxlen=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()}) + if 'task' in observation: + self._queues['task'].append(observation['task']) + + def _prepare_observation_batch(self) -> Dict[str, torch.Tensor]: + qpos_list = list(self._queues['qpos']) + if not qpos_list: + raise ValueError('observation queue is empty.') + 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('image queue is empty.') + while len(images_list) < self.obs_horizon: + images_list.append(images_list[-1]) + names = self.camera_names if self.camera_names is not None else tuple(sorted(images_list[0].keys())) + batch_images = { + name: torch.stack([item[name] for item in images_list], dim=0).unsqueeze(0) + for name in names + } + batch = {'qpos': batch_qpos, 'images': batch_images} + if self._queues['task']: + batch['task'] = [list(self._queues['task'])[-1]] + elif self.task_description is not None: + batch['task'] = [self.task_description] + return batch + + @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() + + def get_normalization_stats(self): + return self.normalization.get_stats() diff --git a/roboimi/vla/conf/agent/smolvla_imf_attnres.yaml b/roboimi/vla/conf/agent/smolvla_imf_attnres.yaml new file mode 100644 index 0000000..eac0b05 --- /dev/null +++ b/roboimi/vla/conf/agent/smolvla_imf_attnres.yaml @@ -0,0 +1,46 @@ +# @package agent +defaults: + - /backbone@condition_encoder: smolvla_prefix_encoder + - /modules@action_encoder: identity_action_encoder + - /head: imf_transformer1d + - _self_ + +_target_: roboimi.vla.agent_smolvla_conditioned.SmolVLAIMFAttnResAgent + +action_dim: 16 +obs_dim: 16 +normalization_type: "min_max" +pred_horizon: 16 +obs_horizon: 2 +num_action_steps: 8 +camera_names: ${data.camera_names} +num_cams: ${len:${agent.camera_names}} + +# Optional fallback language instruction used only when a batch/observation +# does not provide `task`. Keep null by default so this architecture remains +# generic and dataset/eval can supply variable language. +task_description: null +condition_dim: 960 +condition_sequence_length: 241 + +condition_encoder: + num_cameras: ${agent.num_cams} + camera_names: ${agent.camera_names} + +diffusion_steps: 100 +inference_steps: 1 +head_type: "transformer" + +head: + input_dim: ${agent.action_dim} + output_dim: ${agent.action_dim} + horizon: ${agent.pred_horizon} + n_obs_steps: ${agent.condition_sequence_length} + cond_dim: ${agent.condition_dim} + causal_attn: false + time_as_cond: true + obs_as_cond: true + n_cond_layers: 0 + backbone_type: attnres_full + n_head: 1 + n_kv_head: 1 diff --git a/roboimi/vla/conf/backbone/smolvla_prefix_encoder.yaml b/roboimi/vla/conf/backbone/smolvla_prefix_encoder.yaml new file mode 100644 index 0000000..c74c27b --- /dev/null +++ b/roboimi/vla/conf/backbone/smolvla_prefix_encoder.yaml @@ -0,0 +1,19 @@ +_target_: roboimi.vla.models.backbones.smolvla_prefix_encoder.SmolVLAPrefixEncoder + +model_name: HuggingFaceTB/SmolVLM2-500M-Video-Instruct +load_vlm_weights: true +local_files_only: false +num_vlm_layers: 16 +freeze_vlm: true +freeze_vision_encoder: true +train_state_proj: true +max_state_dim: 32 +resize_imgs_with_padding: [512, 512] +tokenizer_max_length: 48 +pad_language_to: max_length +run_text_model: true +camera_names: [r_vis, top, front] +num_cameras: 3 + +dataset_image_resize_shape: null +eval_image_resize_shape: null diff --git a/roboimi/vla/conf/eval/eval.yaml b/roboimi/vla/conf/eval/eval.yaml index b2e38a2..e3f34ee 100644 --- a/roboimi/vla/conf/eval/eval.yaml +++ b/roboimi/vla/conf/eval/eval.yaml @@ -9,6 +9,7 @@ server_startup_timeout_s: 300.0 # parent 等待 inference server 就绪的超时 max_timesteps: 700 # 每回合最大时间步 device: ${train.device} # 与训练保持一致 task_name: "sim_transfer" # 环境任务名称 +task_description: null # 可选语言指令；环境 obs 没有 task 时注入给语言条件策略 # ==================== # 策略执行参数 diff --git a/roboimi/vla/data/simpe_robot_dataset.py b/roboimi/vla/data/simpe_robot_dataset.py index c1c5a93..70c39bc 100644 --- a/roboimi/vla/data/simpe_robot_dataset.py +++ b/roboimi/vla/data/simpe_robot_dataset.py @@ -4,6 +4,7 @@ from torch.utils.data import Dataset from typing import List, Dict, Union, Optional, Sequence from pathlib import Path from collections import OrderedDict +import re class SimpleRobotDataset(Dataset): @@ -24,6 +25,7 @@ class SimpleRobotDataset(Dataset): camera_names: List[str] = None, image_resize_shape: Optional[Sequence[int]] = (224, 224), max_open_files: int = 64, + episode_indices: Optional[Sequence[int]] = None, ): """ Args: @@ -33,6 +35,7 @@ class SimpleRobotDataset(Dataset): camera_names: 相机名称列表，如 ["r_vis", "top", "front"] image_resize_shape: 图像缩放尺寸 (W, H)；为 None 时保留原始分辨率 max_open_files: 每个 worker 最多缓存的 HDF5 文件句柄数 + episode_indices: 可选的原始 episode 编号子集 HDF5 文件格式: - action: [T, action_dim] @@ -48,6 +51,9 @@ class SimpleRobotDataset(Dataset): ) self.max_open_files = max(1, int(max_open_files)) self._file_cache: "OrderedDict[str, h5py.File]" = OrderedDict() + self.requested_episode_indices = ( + None if episode_indices is None else tuple(sorted(int(idx) for idx in episode_indices)) + ) self.dataset_dir = Path(dataset_dir) if not self.dataset_dir.exists(): @@ -59,6 +65,18 @@ class SimpleRobotDataset(Dataset): self.hdf5_files = sorted(self.dataset_dir.glob("episode_*.hdf5")) if not self.hdf5_files: raise FileNotFoundError(f"在 {dataset_dir} 中未找到 HDF5 文件") + if self.requested_episode_indices is not None: + requested = set(self.requested_episode_indices) + filtered = [] + for hdf5_path in self.hdf5_files: + match = re.search(r'episode_(\d+)$', hdf5_path.stem) + if match and int(match.group(1)) in requested: + filtered.append(hdf5_path) + self.hdf5_files = filtered + if not self.hdf5_files: + raise FileNotFoundError( + f"在 {dataset_dir} 中未找到 episode_indices={sorted(requested)} 对应的 HDF5 文件" + ) # 构建 episode 索引（只存储元数据，不加载数据） self.episodes = {} @@ -66,14 +84,18 @@ class SimpleRobotDataset(Dataset): for ep_idx, hdf5_path in enumerate(self.hdf5_files): with h5py.File(hdf5_path, 'r') as f: T = f['action'].shape[0] + dataset_episode_idx = ep_idx + match = re.search(r'episode_(\d+)$', hdf5_path.stem) + if match: + dataset_episode_idx = int(match.group(1)) start_idx = len(self.frame_meta) for t in range(T): self.frame_meta.append({ - "ep_idx": ep_idx, + "ep_idx": dataset_episode_idx, "frame_idx": t, "hdf5_path": hdf5_path, }) - self.episodes[ep_idx] = list(range(start_idx, len(self.frame_meta))) + self.episodes[dataset_episode_idx] = list(range(start_idx, len(self.frame_meta))) print(f"懒加载模式: {len(self.hdf5_files)} 个 episodes, 共 {len(self.frame_meta)} 帧") @@ -227,6 +249,10 @@ class SimpleRobotDataset(Dataset): """获取所有相机键名 (LeRobotDataset 格式)""" return [f"observation.{cam_name}" for cam_name in self.camera_names] + @property + def available_episode_indices(self) -> List[int]: + return sorted(self.episodes.keys()) + @property def camera_info(self) -> dict: """获取相机信息""" diff --git a/roboimi/vla/models/backbones/smolvla_prefix_encoder.py b/roboimi/vla/models/backbones/smolvla_prefix_encoder.py new file mode 100644 index 0000000..4f7e3c7 --- /dev/null +++ b/roboimi/vla/models/backbones/smolvla_prefix_encoder.py @@ -0,0 +1,410 @@ +from __future__ import annotations + +import warnings +from typing import Dict, Sequence + +import torch +import torch.nn.functional as F +from torch import nn + +from roboimi.vla.core.interfaces import VLABackbone + +try: # pragma: no cover - exercised by tests via monkeypatch/fakes + from transformers import AutoModelForImageTextToText, AutoTokenizer +except Exception: # pragma: no cover + AutoModelForImageTextToText = None + AutoTokenizer = None + + +def _resize_with_pad(img: torch.Tensor, width: int, height: int, pad_value: float = 0.0) -> torch.Tensor: + if img.ndim != 4: + raise ValueError(f'expected image tensor shaped (B,C,H,W), got {tuple(img.shape)}') + cur_height, cur_width = img.shape[2:] + ratio = max(cur_width / width, cur_height / height) + resized_height = int(cur_height / ratio) + resized_width = int(cur_width / ratio) + resized = F.interpolate(img, size=(resized_height, resized_width), mode='bilinear', align_corners=False) + pad_height = max(0, int(height - resized_height)) + pad_width = max(0, int(width - resized_width)) + return F.pad(resized, (pad_width, 0, pad_height, 0), value=pad_value) + + +def _pad_vector(vector: torch.Tensor, new_dim: int) -> torch.Tensor: + if vector.shape[-1] == new_dim: + return vector + if vector.shape[-1] > new_dim: + raise ValueError(f'cannot pad vector with dim {vector.shape[-1]} to smaller dim {new_dim}') + shape = list(vector.shape) + shape[-1] = int(new_dim) + padded = torch.zeros(*shape, dtype=vector.dtype, device=vector.device) + padded[..., : vector.shape[-1]] = vector + return padded + + +class SmolVLAPrefixEncoder(VLABackbone): + """SmolVLA-compatible VLM prefix encoder for RoboIMI action experts. + + This module intentionally extracts only the conditioning path from SmolVLA: + multiview image tokens, variable language task tokens, and one projected + state token. It freezes the pretrained VLM by default and returns a fixed + condition token sequence `(B, S, D)` that can be consumed by existing IMF / + transformer-style action heads. + """ + + def __init__( + self, + model_name: str = 'HuggingFaceTB/SmolVLM2-500M-Video-Instruct', + *, + model_name_or_path: str | None = None, + vlm: nn.Module | None = None, + tokenizer=None, + load_vlm_weights: bool = True, + local_files_only: bool = False, + num_vlm_layers: int = 16, + freeze_vlm: bool = True, + freeze_vision_encoder: bool = True, + train_state_proj: bool = True, + max_state_dim: int = 32, + resize_imgs_with_padding: Sequence[int] | None = (512, 512), + dataset_image_resize_shape: Sequence[int] | None = None, + eval_image_resize_shape: Sequence[int] | None = None, + tokenizer_max_length: int = 48, + pad_language_to: str = 'max_length', + run_text_model: bool = True, + camera_names: Sequence[str] = ('r_vis', 'top', 'front'), + num_cameras: int | None = None, + ) -> None: + super().__init__() + if model_name_or_path is not None: + model_name = model_name_or_path + self.model_name = str(model_name) + self.camera_names = tuple(camera_names) + self.num_cameras = int(num_cameras) if num_cameras is not None else len(self.camera_names) + if len(self.camera_names) != self.num_cameras: + raise ValueError( + f'camera_names length ({len(self.camera_names)}) must match num_cameras ({self.num_cameras})' + ) + self.load_vlm_weights = bool(load_vlm_weights) + self.local_files_only = bool(local_files_only) + self.num_vlm_layers_requested = int(num_vlm_layers) + self.freeze_vlm = bool(freeze_vlm) + self.freeze_vision_encoder = bool(freeze_vision_encoder) + self.max_state_dim = int(max_state_dim) + self.resize_imgs_with_padding = self._normalize_resize_shape(resize_imgs_with_padding) + self.dataset_image_resize_shape = self._normalize_resize_shape(dataset_image_resize_shape) + self.eval_image_resize_shape = self._normalize_resize_shape(eval_image_resize_shape) + self.tokenizer_max_length = int(tokenizer_max_length) + self.pad_language_to = str(pad_language_to) + self.run_text_model = bool(run_text_model) + self._warned_trainable_frozen_text_model = False + + if vlm is None: + if AutoModelForImageTextToText is None: + raise ImportError('transformers AutoModelForImageTextToText is required for SmolVLAPrefixEncoder') + if not self.load_vlm_weights: + raise ValueError('SmolVLAPrefixEncoder currently requires load_vlm_weights=True') + vlm = AutoModelForImageTextToText.from_pretrained( + self.model_name, + torch_dtype='bfloat16', + low_cpu_mem_usage=True, + local_files_only=self.local_files_only, + ) + self.vlm = vlm + + if tokenizer is None: + if AutoTokenizer is None: + raise ImportError('transformers AutoTokenizer is required for SmolVLAPrefixEncoder') + tokenizer = AutoTokenizer.from_pretrained(self.model_name, local_files_only=self.local_files_only) + self.tokenizer = tokenizer + + if self.num_vlm_layers_requested > 0: + text_layers = self.vlm.model.text_model.layers + if self.num_vlm_layers_requested > len(text_layers): + raise ValueError( + f'num_vlm_layers ({self.num_vlm_layers_requested}) exceeds available text layers ' + f'({len(text_layers)})' + ) + self.vlm.model.text_model.layers = text_layers[: self.num_vlm_layers_requested] + self.num_vlm_layers = len(self.vlm.model.text_model.layers) + text_config = getattr(self.vlm.config, 'text_config', None) + if text_config is not None and hasattr(text_config, 'num_hidden_layers'): + text_config.num_hidden_layers = self.num_vlm_layers + + hidden_size = int(self.vlm.config.text_config.hidden_size) + self._output_dim = hidden_size + self.state_proj = nn.Linear(self.max_state_dim, hidden_size) + for param in self.state_proj.parameters(): + param.requires_grad = bool(train_state_proj) + + self.last_prefix_pad_mask: torch.Tensor | None = None + self.last_prefix_att_mask: torch.Tensor | None = None + self.last_attention_2d_mask: torch.Tensor | None = None + self.last_position_ids: torch.Tensor | None = None + self._configured_condition_sequence_length = self._infer_configured_condition_sequence_length() + + self.set_requires_grad() + + @staticmethod + def _normalize_resize_shape(shape: Sequence[int] | None) -> tuple[int, int] | None: + if shape is None: + return None + normalized = tuple(int(v) for v in shape) + if len(normalized) != 2: + raise ValueError(f'resize_imgs_with_padding must contain exactly two values, got {normalized}') + return normalized + + @property + def output_dim(self) -> int: + return self._output_dim + + @property + def joint_output_dim(self) -> int: + return self._output_dim + + def _infer_configured_condition_sequence_length(self) -> int: + image_tokens_per_camera = 0 + config = getattr(self.vlm, 'config', None) + vision_config = getattr(config, 'vision_config', None) + if vision_config is not None: + image_size = int(getattr(vision_config, 'image_size', 0) or 0) + patch_size = int(getattr(vision_config, 'patch_size', 0) or 0) + scale_factor = int(getattr(config, 'scale_factor', 1) or 1) + if image_size > 0 and patch_size > 0 and scale_factor > 0: + image_tokens_per_camera = int(((image_size // patch_size) ** 2) / (scale_factor**2)) + return self.num_cameras * image_tokens_per_camera + self.tokenizer_max_length + 1 + + @property + def tokens_per_step(self) -> int: + if self.last_prefix_pad_mask is not None: + return int(self.last_prefix_pad_mask.shape[1]) + return int(self._configured_condition_sequence_length) + + @property + def condition_sequence_length(self) -> int: + return self.tokens_per_step + + @staticmethod + def _freeze_module(module) -> None: + if module is None: + return + if hasattr(module, 'eval'): + module.eval() + parameters = getattr(module, 'parameters', None) + if callable(parameters): + for param in parameters(): + param.requires_grad = False + + def set_requires_grad(self) -> None: + if self.freeze_vision_encoder: + self._freeze_module(self.vlm.model.vision_model) + if self.freeze_vlm: + self._freeze_module(self.vlm) + self._freeze_module(getattr(self.vlm.model, 'vision_model', None)) + self._freeze_module(getattr(self.vlm.model, 'connector', None)) + self._freeze_module(getattr(self.vlm.model, 'text_model', None)) + + def train(self, mode: bool = True): + super().train(mode) + if self.freeze_vlm: + self.vlm.eval() + elif self.freeze_vision_encoder: + self.vlm.model.vision_model.eval() + return self + + def _ordered_camera_names(self, images: Dict[str, torch.Tensor]) -> tuple[str, ...]: + missing = [name for name in self.camera_names if name not in images] + if missing: + raise ValueError(f'image input missing required cameras. missing={missing}, expected={list(self.camera_names)}') + return self.camera_names + + def _batch_size_from_images(self, images: Dict[str, torch.Tensor]) -> int: + return int(next(iter(images.values())).shape[0]) + + def _normalize_tasks(self, task, batch_size: int) -> list[str]: + if task is None: + tasks = [''] * batch_size + elif isinstance(task, str): + tasks = [task] * batch_size + elif isinstance(task, tuple): + tasks = list(task) + elif isinstance(task, list): + tasks = task + else: + raise TypeError(f'task must be str/list/tuple/None, got {type(task)!r}') + if len(tasks) != batch_size: + raise ValueError(f'task batch size ({len(tasks)}) must match image batch size ({batch_size})') + return [item if item.endswith('\n') else f'{item}\n' for item in tasks] + + def tokenize_task(self, task, batch_size: int, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]: + tasks = self._normalize_tasks(task, batch_size) + old_padding_side = getattr(self.tokenizer, 'padding_side', None) + if old_padding_side is not None: + self.tokenizer.padding_side = 'right' + try: + tokenized = self.tokenizer( + tasks, + padding=self.pad_language_to, + max_length=self.tokenizer_max_length, + return_tensors='pt', + truncation=True, + ) + finally: + if old_padding_side is not None: + self.tokenizer.padding_side = old_padding_side + tokens = tokenized['input_ids'].to(device=device) + masks = tokenized['attention_mask'].to(device=device, dtype=torch.bool) + return tokens, masks + + def prepare_images(self, images: Dict[str, torch.Tensor]) -> tuple[list[torch.Tensor], list[torch.Tensor]]: + camera_names = self._ordered_camera_names(images) + reference = images[camera_names[0]] + if reference.ndim != 5: + raise ValueError(f'expected image tensor shaped (B,T,C,H,W), got {tuple(reference.shape)}') + batch_size = reference.shape[0] + prepared: list[torch.Tensor] = [] + masks: list[torch.Tensor] = [] + for camera_name in camera_names: + image = images[camera_name] + if image.shape[:2] != reference.shape[:2] or image.shape[2] != reference.shape[2]: + raise ValueError(f'camera {camera_name!r} shape {tuple(image.shape)} does not match reference {tuple(reference.shape)}') + image = image[:, -1].contiguous().float().clamp(0.0, 1.0) + if self.resize_imgs_with_padding is not None: + image = _resize_with_pad(image, *self.resize_imgs_with_padding, pad_value=0.0) + image = image * 2.0 - 1.0 + prepared.append(image) + masks.append(torch.ones(batch_size, dtype=torch.bool, device=image.device)) + return prepared, masks + + def prepare_state(self, state: torch.Tensor) -> torch.Tensor: + if state.ndim > 2: + state = state[:, -1, :] + return _pad_vector(state.float(), self.max_state_dim) + + def embed_image(self, image: torch.Tensor) -> torch.Tensor: + if hasattr(self.vlm.model, 'get_image_features'): + pixel_values = image[:, None, ...] + pixel_attention_mask = torch.ones( + image.shape[0], + 1, + image.shape[2], + image.shape[3], + dtype=torch.bool, + device=image.device, + ) + with torch.set_grad_enabled( + torch.is_grad_enabled() and not self.freeze_vlm and not self.freeze_vision_encoder + ): + hidden = self.vlm.model.get_image_features( + pixel_values=pixel_values, + pixel_attention_mask=pixel_attention_mask, + return_dict=True, + ).pooler_output + else: + vision_model = self.vlm.model.vision_model + with torch.set_grad_enabled( + torch.is_grad_enabled() and not self.freeze_vlm and not self.freeze_vision_encoder + ): + patch_attention_mask = torch.ones( + image.shape[0], + image.shape[2] // self.vlm.config.vision_config.patch_size, + image.shape[3] // self.vlm.config.vision_config.patch_size, + dtype=torch.bool, + device=image.device, + ) if hasattr(self.vlm.config, 'vision_config') and hasattr(self.vlm.config.vision_config, 'patch_size') else None + hidden = vision_model( + pixel_values=image.to(dtype=vision_model.dtype), + patch_attention_mask=patch_attention_mask, + ).last_hidden_state + hidden = self.vlm.model.connector(hidden) + return hidden + + def embed_language_tokens(self, tokens: torch.Tensor) -> torch.Tensor: + return self.vlm.model.text_model.get_input_embeddings()(tokens) + + @staticmethod + def make_att_2d_masks(pad_masks: torch.Tensor, att_masks: torch.Tensor) -> torch.Tensor: + cumsum = torch.cumsum(att_masks, dim=1) + att_2d = cumsum[:, None, :] <= cumsum[:, :, None] + pad_2d = pad_masks[:, None, :] * pad_masks[:, :, None] + return att_2d & pad_2d + + def embed_prefix(self, images: Dict[str, torch.Tensor], state: torch.Tensor, task=None) -> torch.Tensor: + prepared_images, image_masks = self.prepare_images(images) + batch_size = prepared_images[0].shape[0] + device = prepared_images[0].device + + embs: list[torch.Tensor] = [] + pad_masks: list[torch.Tensor] = [] + att_masks: list[int] = [] + + for image, image_mask in zip(prepared_images, image_masks, strict=False): + img_emb = self.embed_image(image) + img_emb = img_emb * torch.tensor( + img_emb.shape[-1] ** 0.5, + dtype=img_emb.dtype, + device=img_emb.device, + ) + num_img_tokens = img_emb.shape[1] + embs.append(img_emb) + pad_masks.append(image_mask[:, None].expand(batch_size, num_img_tokens)) + att_masks += [0] * num_img_tokens + + lang_tokens, lang_masks = self.tokenize_task(task, batch_size=batch_size, device=device) + lang_emb = self.embed_language_tokens(lang_tokens) + lang_emb = lang_emb * (lang_emb.shape[-1] ** 0.5) + embs.append(lang_emb) + pad_masks.append(lang_masks) + att_masks += [0] * lang_emb.shape[1] + + state = self.prepare_state(state).to(device=device) + state_emb = self.state_proj(state).unsqueeze(1) + embs.append(state_emb) + pad_masks.append(torch.ones(batch_size, 1, dtype=torch.bool, device=device)) + att_masks += [1] + + prefix = torch.cat(embs, dim=1) + prefix_pad_mask = torch.cat(pad_masks, dim=1) + prefix_att_mask = torch.tensor(att_masks, dtype=torch.bool, device=device)[None, :].expand(batch_size, -1) + + self.last_prefix_pad_mask = prefix_pad_mask + self.last_prefix_att_mask = prefix_att_mask + self.last_attention_2d_mask = self.make_att_2d_masks(prefix_pad_mask, prefix_att_mask) + self.last_position_ids = torch.cumsum(prefix_pad_mask, dim=1) - 1 + return prefix + + def encode_prefix(self, prefix: torch.Tensor) -> torch.Tensor: + if not self.run_text_model: + return prefix + if self.last_attention_2d_mask is None or self.last_position_ids is None: + raise RuntimeError('encode_prefix requires masks from embed_prefix') + text_model = self.vlm.model.text_model + text_dtype = getattr(text_model, 'dtype', prefix.dtype) + if ( + self.freeze_vlm + and not self._warned_trainable_frozen_text_model + and any(param.requires_grad for param in text_model.parameters()) + ): + warnings.warn( + 'freeze_vlm=True but text_model has trainable parameters; temporarily disabling ' + 'their gradients while keeping gradients for trainable prefix inputs.', + RuntimeWarning, + ) + self._warned_trainable_frozen_text_model = True + with torch.set_grad_enabled(torch.is_grad_enabled()): + outputs = text_model( + inputs_embeds=prefix.to(dtype=text_dtype), + attention_mask=self.last_attention_2d_mask[:, None, :, :], + position_ids=self.last_position_ids, + use_cache=False, + return_dict=True, + ) + return outputs.last_hidden_state + + def forward(self, images: Dict[str, torch.Tensor], state: torch.Tensor | None = None, task=None) -> torch.Tensor: + if state is None: + raise ValueError('SmolVLAPrefixEncoder.forward requires `state` for SmolVLA-compatible state token') + prefix = self.embed_prefix(images=images, state=state, task=task) + return self.encode_prefix(prefix) + + +SmolVLMPrefixEncoder = SmolVLAPrefixEncoder diff --git a/tests/test_eval_vla_execution.py b/tests/test_eval_vla_execution.py index f9d7a44..db0fa7a 100644 --- a/tests/test_eval_vla_execution.py +++ b/tests/test_eval_vla_execution.py @@ -129,6 +129,38 @@ class EvalVLAExecutionTest(unittest.TestCase): ["r_vis", "top", "front"], ) + def test_resolve_eval_image_resize_shape_prefers_condition_encoder_override(self): + cfg = OmegaConf.create( + { + "agent": { + "condition_encoder": { + "eval_image_resize_shape": None, + }, + }, + "data": { + "image_resize_shape": [224, 224], + }, + } + ) + + self.assertIsNone(eval_vla._resolve_eval_image_resize_shape(cfg)) + + def test_resolve_eval_image_resize_shape_prefers_vision_backbone_override(self): + cfg = OmegaConf.create( + { + "agent": { + "vision_backbone": { + "eval_image_resize_shape": [256, 256], + }, + }, + "data": { + "image_resize_shape": [224, 224], + }, + } + ) + + self.assertEqual(eval_vla._resolve_eval_image_resize_shape(cfg), (256, 256)) + def test_build_episode_plans_without_box_poses_keeps_serial_sampling_lazy(self): plans = eval_vla._build_episode_plans(num_episodes=3) @@ -168,6 +200,58 @@ class EvalVLAExecutionTest(unittest.TestCase): np.array([[[[1.0]]], [[[1.0]]], [[[1.0]]]], dtype=np.float32), ) + def test_prepare_local_policy_batch_keeps_latest_variable_task_when_present(self): + queues = eval_vla._new_local_policy_queues(obs_horizon=2) + first_observation = { + "qpos": torch.tensor([1.0, 2.0], dtype=torch.float32), + "images": {"front": torch.tensor([[[1.0]]], dtype=torch.float32)}, + "task": "pick the red cube", + } + second_observation = { + "qpos": torch.tensor([3.0, 4.0], dtype=torch.float32), + "images": {"front": torch.tensor([[[2.0]]], dtype=torch.float32)}, + "task": "insert the peg into the socket", + } + + eval_vla._populate_local_policy_queues(queues, first_observation) + eval_vla._populate_local_policy_queues(queues, second_observation) + batch = eval_vla._prepare_local_policy_batch( + queues, + obs_horizon=2, + camera_names=["front"], + ) + + self.assertEqual(batch["task"], ["insert the peg into the socket"]) + + def test_prepare_local_policy_batch_omits_task_for_legacy_observations(self): + queues = eval_vla._new_local_policy_queues(obs_horizon=2) + 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=2, + camera_names=["front"], + ) + + self.assertNotIn("task", batch) + + def test_serialize_deserialize_policy_batch_preserves_task(self): + batch = { + "qpos": torch.zeros(1, 2, 2, dtype=torch.float32), + "images": {"front": torch.zeros(1, 2, 1, 1, 1, dtype=torch.float32)}, + "task": ["pick the red cube", "insert the peg into the socket"], + } + + serialized = eval_vla._serialize_policy_batch(batch) + deserialized = eval_vla._deserialize_policy_batch(serialized, device="cpu") + + self.assertEqual(serialized["task"], batch["task"]) + self.assertEqual(deserialized["task"], batch["task"]) + def test_enqueue_predicted_actions_uses_executable_slice(self): queues = eval_vla._new_local_policy_queues(obs_horizon=2) predicted_actions = torch.tensor( diff --git a/tests/test_eval_vla_headless.py b/tests/test_eval_vla_headless.py index 768965f..0dd469f 100644 --- a/tests/test_eval_vla_headless.py +++ b/tests/test_eval_vla_headless.py @@ -15,6 +15,7 @@ class _FakeAgent: def __init__(self): self.reset_calls = 0 self.last_observation = None + self.observation_shapes = [] def eval(self): return self @@ -27,6 +28,7 @@ class _FakeAgent: def select_action(self, observation): self.last_observation = observation + self.observation_shapes.append(tuple(observation["images"]["front"].shape)) return torch.zeros(16) @@ -108,6 +110,23 @@ class EvalVLAHeadlessTest(unittest.TestCase): self.assertEqual(tuple(prepared["images"]["front"].shape), (3, 8, 8)) self.assertEqual(tuple(prepared["qpos"].shape), (16,)) + def test_prepare_observation_preserves_task_when_present(self): + obs = { + "images": { + "front": np.zeros((4, 4, 3), dtype=np.uint8), + }, + "qpos": np.zeros(16, dtype=np.float32), + "task": "insert the peg into the socket", + } + + prepared = eval_vla.prepare_observation( + obs, + ["front"], + image_resize_shape=None, + ) + + self.assertEqual(prepared["task"], "insert the peg into the socket") + def test_headless_eval_sets_mujoco_gl_to_egl_when_display_missing(self): cfg = OmegaConf.create({"eval": {"headless": True}}) with mock.patch.dict(eval_vla.os.environ, {}, clear=True): @@ -125,6 +144,8 @@ class EvalVLAHeadlessTest(unittest.TestCase): self.assertIn("headless", eval_cfg) self.assertFalse(eval_cfg.headless) + self.assertIn("task_description", eval_cfg) + self.assertIsNone(eval_cfg.task_description) def test_make_sim_env_accepts_headless_and_disables_render(self): fake_env = object() @@ -291,6 +312,74 @@ class EvalVLAHeadlessTest(unittest.TestCase): self.assertIsNotNone(fake_agent.last_observation) self.assertIn("front", fake_agent.last_observation["images"]) + def test_eval_main_uses_condition_encoder_resize_override(self): + fake_env = _FakeEnv() + fake_agent = _FakeAgent() + cfg = OmegaConf.create( + { + "agent": { + "condition_encoder": { + "eval_image_resize_shape": None, + }, + }, + "data": { + "image_resize_shape": [224, 224], + }, + "eval": { + "ckpt_path": "checkpoints/vla_model_best.pt", + "num_episodes": 1, + "max_timesteps": 1, + "device": "cpu", + "task_name": "sim_transfer", + "camera_names": ["front"], + "use_smoothing": False, + "smooth_alpha": 0.3, + "verbose_action": False, + "headless": True, + }, + } + ) + + with mock.patch.object(eval_vla, "load_checkpoint", return_value=(fake_agent, None)), \ + mock.patch.object(eval_vla, "make_sim_env", return_value=fake_env), \ + mock.patch.object(eval_vla, "sample_transfer_pose", return_value=np.array([0.1, 0.2, 0.3])), \ + mock.patch.object(eval_vla, "execute_policy_action"), \ + mock.patch.object(eval_vla, "tqdm", side_effect=lambda iterable, **kwargs: iterable): + eval_vla.main.__wrapped__(cfg) + + self.assertEqual(fake_agent.observation_shapes, [(3, 8, 8)]) + + def test_eval_main_injects_configured_task_description_when_env_omits_task(self): + fake_env = _FakeEnv() + fake_agent = _FakeAgent() + cfg = OmegaConf.create( + { + "agent": {}, + "eval": { + "ckpt_path": "checkpoints/vla_model_best.pt", + "num_episodes": 1, + "max_timesteps": 1, + "device": "cpu", + "task_name": "sim_transfer", + "task_description": "pick the red cube", + "camera_names": ["front"], + "use_smoothing": False, + "smooth_alpha": 0.3, + "verbose_action": False, + "headless": True, + }, + } + ) + + with mock.patch.object(eval_vla, "load_checkpoint", return_value=(fake_agent, None)), \ + mock.patch.object(eval_vla, "make_sim_env", return_value=fake_env), \ + mock.patch.object(eval_vla, "sample_transfer_pose", return_value=np.array([0.1, 0.2, 0.3])), \ + mock.patch.object(eval_vla, "execute_policy_action"), \ + mock.patch.object(eval_vla, "tqdm", side_effect=lambda iterable, **kwargs: iterable): + eval_vla.main.__wrapped__(cfg) + + self.assertEqual(fake_agent.last_observation["task"], "pick the red cube") + def test_run_eval_returns_average_reward_summary(self): reward_sequences = [ [1.0, 2.0], diff --git a/tests/test_simple_robot_dataset_episode_filter.py b/tests/test_simple_robot_dataset_episode_filter.py new file mode 100644 index 0000000..f94e5cc --- /dev/null +++ b/tests/test_simple_robot_dataset_episode_filter.py @@ -0,0 +1,42 @@ +import tempfile +import unittest +from pathlib import Path + +import h5py +import numpy as np + + +def _write_episode(path: Path, length: int = 2): + with h5py.File(path, 'w') as f: + f.create_dataset('action', data=np.zeros((length, 16), dtype=np.float32)) + obs = f.create_group('observations') + obs.create_dataset('qpos', data=np.zeros((length, 16), dtype=np.float32)) + images = obs.create_group('images') + for cam_name in ('l_vis', 'r_vis', 'front'): + images.create_dataset(cam_name, data=np.zeros((length, 4, 4, 3), dtype=np.uint8)) + + +class SimpleRobotDatasetEpisodeFilterTest(unittest.TestCase): + def test_filters_by_original_episode_indices_and_exposes_available_indices(self): + from roboimi.vla.data.simpe_robot_dataset import SimpleRobotDataset + + with tempfile.TemporaryDirectory() as tmpdir: + root = Path(tmpdir) + _write_episode(root / 'episode_0.hdf5') + _write_episode(root / 'episode_2.hdf5') + _write_episode(root / 'episode_10.hdf5') + + dataset = SimpleRobotDataset( + root, + camera_names=['l_vis', 'r_vis', 'front'], + image_resize_shape=None, + episode_indices=[10, 2], + ) + + self.assertEqual(dataset.available_episode_indices, [2, 10]) + self.assertEqual(set(dataset.episodes.keys()), {2, 10}) + self.assertEqual(len(dataset), 4) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/test_smolvla_imf_agent.py b/tests/test_smolvla_imf_agent.py new file mode 100644 index 0000000..c93bcea --- /dev/null +++ b/tests/test_smolvla_imf_agent.py @@ -0,0 +1,311 @@ +import contextlib +import importlib +import importlib.machinery +import sys +import types +import unittest +from pathlib import Path +from unittest import mock + +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 +from torch import nn + + +_REPO_ROOT = Path(__file__).resolve().parents[1] +_CONFIG_DIR = str((_REPO_ROOT / 'roboimi/vla/conf').resolve()) +_CAMERA_NAMES = ('r_vis', 'top', 'front') +_MISSING = object() + + +class _RecordingHead(nn.Module): + def __init__(self): + super().__init__() + self.scale = nn.Parameter(torch.tensor(0.5)) + self.calls = [] + + @staticmethod + def _broadcast(value, reference): + while value.ndim < reference.ndim: + value = value.unsqueeze(-1) + return value + + def forward(self, sample, r, t, cond=None): + self.calls.append({ + 'sample': sample.detach().clone(), + 'r': r.detach().clone(), + 't': t.detach().clone(), + 'cond': None if cond is None else cond.detach().clone(), + }) + cond_term = 0.0 if cond is None else cond.mean(dim=(1, 2), keepdim=True) + return self.scale * sample + self._broadcast(r, sample) + 2.0 * self._broadcast(t, sample) + cond_term + + +class _TaskAwareConditionEncoder(nn.Module): + output_dim = 4 + condition_sequence_length = 3 + tokens_per_step = 3 + joint_output_dim = 4 + camera_names = _CAMERA_NAMES + num_cameras = 3 + + def __init__(self, **kwargs): + super().__init__() + self.constructor_kwargs = dict(kwargs) + self.bias = nn.Parameter(torch.tensor(0.0)) + self.calls = [] + + def forward(self, images, state, task=None): + self.calls.append({'task': task, 'state': state.detach().clone(), 'image_keys': tuple(images.keys())}) + batch_size = state.shape[0] + state_last = state[:, -1, 0] + if isinstance(task, str): + task_lengths = torch.full((batch_size,), float(len(task)), dtype=state.dtype, device=state.device) + else: + task_lengths = torch.tensor([float(len(item)) for item in task], dtype=state.dtype, device=state.device) + image_marker = images['r_vis'][:, -1].mean(dim=(1, 2, 3)) + token0 = torch.stack([state_last, task_lengths, image_marker, torch.ones_like(state_last)], dim=-1) + token1 = token0 + 1.0 + token2 = token0 + 2.0 + return torch.stack([token0, token1, token2], dim=1) + self.bias + + +class _BF16TaskAwareConditionEncoder(_TaskAwareConditionEncoder): + def forward(self, images, state, task=None): + return super().forward(images, state, task=task).to(dtype=torch.bfloat16) + + +class _StubIMFHead(nn.Module): + def __init__(self, input_dim, output_dim, horizon, n_obs_steps, cond_dim, **kwargs): + super().__init__() + self.constructor_kwargs = { + 'input_dim': input_dim, + 'output_dim': output_dim, + 'horizon': horizon, + 'n_obs_steps': n_obs_steps, + 'cond_dim': cond_dim, + **kwargs, + } + self.proj = nn.Linear(input_dim, output_dim) + self.cond_obs_emb = nn.Linear(cond_dim, max(cond_dim, 1)) + + def forward(self, sample, r, t, cond=None): + return torch.zeros_like(sample) + + def get_optim_groups(self, weight_decay): + return [ + {'params': [self.proj.weight], 'weight_decay': weight_decay}, + {'params': [self.proj.bias, self.cond_obs_emb.weight, self.cond_obs_emb.bias], 'weight_decay': 0.0}, + ] + + +@contextlib.contextmanager +def _stub_optional_modules(include_head=False, include_condition_encoder=False): + previous = {} + + def inject(name, module): + if name not in previous: + previous[name] = sys.modules.get(name, _MISSING) + sys.modules[name] = module + + diffusers_module = types.ModuleType('diffusers') + schedulers_module = types.ModuleType('diffusers.schedulers') + ddpm_module = types.ModuleType('diffusers.schedulers.scheduling_ddpm') + ddim_module = types.ModuleType('diffusers.schedulers.scheduling_ddim') + + class _FakeScheduler: + def __init__(self, num_train_timesteps=100, **kwargs): + self.config = types.SimpleNamespace(num_train_timesteps=num_train_timesteps) + + ddpm_module.DDPMScheduler = _FakeScheduler + ddim_module.DDIMScheduler = _FakeScheduler + diffusers_module.DDPMScheduler = _FakeScheduler + diffusers_module.DDIMScheduler = _FakeScheduler + diffusers_module.schedulers = schedulers_module + + try: + inject('diffusers', diffusers_module) + inject('diffusers.schedulers', schedulers_module) + inject('diffusers.schedulers.scheduling_ddpm', ddpm_module) + inject('diffusers.schedulers.scheduling_ddim', ddim_module) + if include_head: + import roboimi.vla.models.heads as heads_package + head_module = types.ModuleType('roboimi.vla.models.heads.imf_transformer1d') + head_module.IMFTransformer1D = _StubIMFHead + inject('roboimi.vla.models.heads.imf_transformer1d', head_module) + setattr(heads_package, 'imf_transformer1d', head_module) + if include_condition_encoder: + module = types.ModuleType('tests.fake_smolvla_condition_encoder') + module.TaskAwareConditionEncoder = _TaskAwareConditionEncoder + module.BF16TaskAwareConditionEncoder = _BF16TaskAwareConditionEncoder + inject('tests.fake_smolvla_condition_encoder', 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 [])) + + +class SmolVLAIMFAgentTest(unittest.TestCase): + def test_compute_loss_and_predict_action_pass_variable_task_to_condition_encoder(self): + from roboimi.vla.agent_smolvla_conditioned import SmolVLAIMFAttnResAgent + + condition_encoder = _BF16TaskAwareConditionEncoder() + head = _RecordingHead() + agent = SmolVLAIMFAttnResAgent( + condition_encoder=condition_encoder, + action_encoder=nn.Identity(), + head=head, + action_dim=2, + obs_dim=1, + pred_horizon=3, + obs_horizon=2, + diffusion_steps=10, + inference_steps=1, + num_cams=3, + camera_names=_CAMERA_NAMES, + num_action_steps=2, + head_type='transformer', + ) + images = { + 'r_vis': torch.full((2, 2, 1, 2, 2), 1.0), + 'top': torch.full((2, 2, 1, 2, 2), 2.0), + 'front': torch.full((2, 2, 1, 2, 2), 3.0), + } + qpos = torch.tensor([[[1.0], [2.0]], [[3.0], [4.0]]]) + actions = torch.zeros(2, 3, 2) + tasks = ['short', 'longer task'] + loss = agent.compute_loss({'images': images, 'qpos': qpos, 'action': actions, 'task': tasks}) + self.assertTrue(torch.isfinite(loss)) + self.assertEqual(condition_encoder.calls[-1]['task'], tasks) + self.assertEqual(head.calls[-1]['cond'].shape, (2, 3, 4)) + self.assertTrue(torch.allclose(head.calls[-1]['cond'][:, 0, 1], torch.tensor([5.0, 11.0]))) + + with mock.patch('roboimi.vla.agent_imf.torch.randn', return_value=torch.zeros(2, 3, 2)): + pred = agent.predict_action(images, qpos, task=tasks) + self.assertEqual(pred.shape, (2, 3, 2)) + self.assertEqual(condition_encoder.calls[-1]['task'], tasks) + + def test_condition_tokens_are_cast_to_action_head_dtype(self): + from roboimi.vla.agent_smolvla_conditioned import SmolVLAIMFAttnResAgent + + condition_encoder = _TaskAwareConditionEncoder() + head = _RecordingHead() + agent = SmolVLAIMFAttnResAgent( + condition_encoder=condition_encoder, + action_encoder=nn.Identity(), + head=head, + action_dim=2, + obs_dim=1, + pred_horizon=3, + obs_horizon=2, + diffusion_steps=10, + inference_steps=1, + num_cams=3, + camera_names=_CAMERA_NAMES, + num_action_steps=2, + head_type='transformer', + ) + images = { + 'r_vis': torch.full((1, 2, 1, 2, 2), 1.0), + 'top': torch.full((1, 2, 1, 2, 2), 2.0), + 'front': torch.full((1, 2, 1, 2, 2), 3.0), + } + qpos = torch.tensor([[[1.0], [2.0]]], dtype=torch.float32) + + cond = agent._build_cond(images, qpos, task=['pick']) + + self.assertEqual(cond.dtype, head.scale.dtype) + + def test_unknown_dataset_task_uses_configured_task_description(self): + from roboimi.vla.agent_smolvla_conditioned import SmolVLAIMFAttnResAgent + + condition_encoder = _TaskAwareConditionEncoder() + head = _RecordingHead() + agent = SmolVLAIMFAttnResAgent( + condition_encoder=condition_encoder, + action_encoder=nn.Identity(), + head=head, + action_dim=2, + obs_dim=1, + pred_horizon=3, + obs_horizon=2, + diffusion_steps=10, + inference_steps=1, + num_cams=3, + camera_names=_CAMERA_NAMES, + num_action_steps=2, + head_type='transformer', + task_description='insert the peg into the socket', + ) + images = { + 'r_vis': torch.full((2, 2, 1, 2, 2), 1.0), + 'top': torch.full((2, 2, 1, 2, 2), 2.0), + 'front': torch.full((2, 2, 1, 2, 2), 3.0), + } + qpos = torch.tensor([[[1.0], [2.0]], [[3.0], [4.0]]], dtype=torch.float32) + + agent._build_cond(images, qpos, task=['unknown', '']) + + self.assertEqual( + condition_encoder.calls[-1]['task'], + ['insert the peg into the socket', 'insert the peg into the socket'], + ) + + def test_hydra_config_instantiates_smolvla_imf_attnres_with_condition_encoder_contract(self): + cfg = _compose_cfg(overrides=[ + 'agent=smolvla_imf_attnres', + 'agent.condition_encoder._target_=tests.fake_smolvla_condition_encoder.TaskAwareConditionEncoder', + 'agent.condition_dim=4', + 'agent.condition_sequence_length=3', + 'agent.head.n_layer=1', + 'agent.head.n_emb=16', + ]) + self.assertEqual(cfg.agent._target_, 'roboimi.vla.agent_smolvla_conditioned.SmolVLAIMFAttnResAgent') + self.assertEqual(cfg.agent.head.cond_dim, cfg.agent.condition_dim) + self.assertEqual(cfg.agent.head.n_obs_steps, cfg.agent.condition_sequence_length) + + with _stub_optional_modules(include_head=True, include_condition_encoder=True): + agent = instantiate(cfg.agent) + + self.assertEqual(agent.per_step_cond_dim, 4) + self.assertEqual(agent.condition_sequence_length, 3) + self.assertIsInstance(agent.noise_pred_net, _StubIMFHead) + self.assertEqual(agent.noise_pred_net.constructor_kwargs['cond_dim'], 4) + self.assertEqual(agent.noise_pred_net.constructor_kwargs['n_obs_steps'], 3) + + def test_hydra_config_exposes_smolvla_pretrained_vlm_defaults(self): + cfg = _compose_cfg(overrides=[ + 'agent=smolvla_imf_attnres', + ]) + + self.assertEqual(cfg.agent.condition_encoder.model_name, 'HuggingFaceTB/SmolVLM2-500M-Video-Instruct') + self.assertTrue(cfg.agent.condition_encoder.load_vlm_weights) + self.assertEqual(cfg.agent.condition_encoder.num_vlm_layers, 16) + self.assertTrue(cfg.agent.condition_encoder.freeze_vlm) + self.assertTrue(cfg.agent.condition_encoder.freeze_vision_encoder) + self.assertTrue(cfg.agent.condition_encoder.run_text_model) + self.assertEqual(cfg.agent.condition_encoder.max_state_dim, 32) + self.assertIsNone(cfg.agent.condition_encoder.dataset_image_resize_shape) + self.assertIsNone(cfg.agent.condition_encoder.eval_image_resize_shape) + self.assertEqual(cfg.agent.condition_dim, 960) + self.assertEqual(cfg.agent.condition_sequence_length, 241) + self.assertEqual(cfg.agent.head.cond_dim, 960) + self.assertEqual(cfg.agent.head.n_obs_steps, 241) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/test_smolvla_prefix_encoder.py b/tests/test_smolvla_prefix_encoder.py new file mode 100644 index 0000000..5f41cf8 --- /dev/null +++ b/tests/test_smolvla_prefix_encoder.py @@ -0,0 +1,328 @@ +import types +import unittest +from unittest import mock + +import torch +from torch import nn + + +class _FakeVisionOutput: + def __init__(self, last_hidden_state): + self.last_hidden_state = last_hidden_state + + +class _FakeVisionModel(nn.Module): + def __init__(self, hidden_size=4): + super().__init__() + self.dtype = torch.float32 + self.scale = nn.Parameter(torch.tensor(1.0)) + self.calls = [] + self.hidden_size = hidden_size + + def forward(self, pixel_values=None, patch_attention_mask=None): + self.calls.append({ + 'pixel_values': pixel_values.detach().clone(), + 'patch_attention_mask': patch_attention_mask, + }) + pooled = pixel_values.mean(dim=(2, 3)) * self.scale + tokens = torch.stack([pooled, pooled + 1.0], dim=1) + return _FakeVisionOutput(tokens) + + +class _FakeConnector(nn.Module): + def __init__(self, in_dim=3, out_dim=4): + super().__init__() + self.proj = nn.Linear(in_dim, out_dim, bias=False) + with torch.no_grad(): + self.proj.weight.copy_( + torch.tensor( + [ + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, 0.0, 1.0], + [1.0, 1.0, 1.0], + ] + ) + ) + + def forward(self, x): + return self.proj(x) + + +class _FakeTextModel(nn.Module): + def __init__(self, vocab_size=32, hidden_size=4, num_layers=6): + super().__init__() + self.embed = nn.Embedding(vocab_size, hidden_size) + self.layers = nn.ModuleList([nn.Linear(hidden_size, hidden_size) for _ in range(num_layers)]) + self.norm = nn.Identity() + self.forward_calls = [] + with torch.no_grad(): + for idx in range(vocab_size): + self.embed.weight[idx].fill_(float(idx)) + + def get_input_embeddings(self): + return self.embed + + def forward( + self, + input_ids=None, + attention_mask=None, + position_ids=None, + past_key_values=None, + inputs_embeds=None, + use_cache=None, + return_dict=True, + cache_position=None, + **kwargs, + ): + del input_ids, past_key_values, use_cache, cache_position, kwargs + self.forward_calls.append({ + 'attention_mask': None if attention_mask is None else attention_mask.detach().clone(), + 'position_ids': None if position_ids is None else position_ids.detach().clone(), + 'inputs_embeds': inputs_embeds.detach().clone(), + }) + hidden = inputs_embeds + for layer in self.layers: + hidden = layer(hidden) + hidden = self.norm(hidden) + if return_dict: + return types.SimpleNamespace(last_hidden_state=hidden) + return (hidden,) + + +class _FakeVLM(nn.Module): + def __init__(self): + super().__init__() + text_config = types.SimpleNamespace(hidden_size=4, head_dim=2, num_attention_heads=2, num_key_value_heads=1) + self.config = types.SimpleNamespace(text_config=text_config) + self.model = types.SimpleNamespace( + vision_model=_FakeVisionModel(hidden_size=4), + connector=_FakeConnector(in_dim=3, out_dim=4), + text_model=_FakeTextModel(hidden_size=4, num_layers=6), + ) + + +class _FakeTokenizer: + fake_image_token_id = 29 + global_image_token_id = 30 + + def __init__(self): + self.padding_side = 'left' + self.calls = [] + + def __call__(self, text, *, padding, max_length, return_tensors, truncation): + self.calls.append({ + 'text': list(text), + 'padding': padding, + 'max_length': max_length, + 'return_tensors': return_tensors, + 'truncation': truncation, + 'padding_side': self.padding_side, + }) + batch = len(text) + ids = torch.zeros(batch, max_length, dtype=torch.long) + mask = torch.zeros(batch, max_length, dtype=torch.bool) + for row, item in enumerate(text): + del item + ids[row, :3] = torch.tensor([1, 2, 3]) + mask[row, :3] = True + return {'input_ids': ids, 'attention_mask': mask} + + +class SmolVLAPrefixEncoderTest(unittest.TestCase): + def test_loads_pretrained_vlm_with_local_files_only_crops_layers_and_freezes_vlm(self): + from roboimi.vla.models.backbones.smolvla_prefix_encoder import SmolVLAPrefixEncoder + + fake_vlm = _FakeVLM() + fake_tokenizer = _FakeTokenizer() + with mock.patch( + 'roboimi.vla.models.backbones.smolvla_prefix_encoder.AutoModelForImageTextToText.from_pretrained', + return_value=fake_vlm, + ) as model_loader, mock.patch( + 'roboimi.vla.models.backbones.smolvla_prefix_encoder.AutoTokenizer.from_pretrained', + return_value=fake_tokenizer, + ) as tokenizer_loader: + encoder = SmolVLAPrefixEncoder( + model_name='HuggingFaceTB/SmolVLM2-500M-Video-Instruct', + local_files_only=True, + num_vlm_layers=2, + freeze_vlm=True, + max_state_dim=5, + tokenizer_max_length=4, + camera_names=('r_vis', 'top'), + resize_imgs_with_padding=None, + ) + + model_loader.assert_called_once() + self.assertEqual(model_loader.call_args.kwargs['local_files_only'], True) + self.assertEqual(model_loader.call_args.kwargs['torch_dtype'], 'bfloat16') + tokenizer_loader.assert_called_once_with( + 'HuggingFaceTB/SmolVLM2-500M-Video-Instruct', + local_files_only=True, + ) + self.assertEqual(len(encoder.vlm.model.text_model.layers), 2) + self.assertTrue(all(not p.requires_grad for p in encoder.vlm.parameters())) + self.assertFalse(encoder.vlm.training) + encoder.train() + self.assertFalse(encoder.vlm.training) + + def test_accepts_train_eval_resize_compatibility_fields(self): + from roboimi.vla.models.backbones.smolvla_prefix_encoder import SmolVLAPrefixEncoder + + encoder = SmolVLAPrefixEncoder( + vlm=_FakeVLM(), + tokenizer=_FakeTokenizer(), + num_vlm_layers=3, + camera_names=('r_vis', 'top'), + resize_imgs_with_padding=None, + dataset_image_resize_shape=None, + eval_image_resize_shape=(640, 480), + ) + + self.assertIsNone(encoder.dataset_image_resize_shape) + self.assertEqual(encoder.eval_image_resize_shape, (640, 480)) + + def test_embed_prefix_uses_variable_tasks_camera_order_last_state_and_smolvla_masks(self): + from roboimi.vla.models.backbones.smolvla_prefix_encoder import SmolVLAPrefixEncoder + + fake_vlm = _FakeVLM() + fake_tokenizer = _FakeTokenizer() + encoder = SmolVLAPrefixEncoder( + vlm=fake_vlm, + tokenizer=fake_tokenizer, + num_vlm_layers=3, + freeze_vlm=True, + train_state_proj=True, + max_state_dim=5, + tokenizer_max_length=4, + camera_names=('r_vis', 'top'), + resize_imgs_with_padding=None, + ) + with torch.no_grad(): + encoder.state_proj.weight.zero_() + encoder.state_proj.bias.zero_() + encoder.state_proj.weight[:, :4] = torch.eye(4) + + images = { + 'top': torch.full((2, 2, 3, 2, 2), 0.75), + 'r_vis': torch.full((2, 2, 3, 2, 2), 0.25), + } + state = torch.tensor( + [ + [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], + [[7.0, 8.0, 9.0], [10.0, 11.0, 12.0]], + ] + ) + tasks = ['pick red cube', 'open drawer'] + out = encoder.embed_prefix(images=images, state=state, task=tasks) + + # 2 cameras * 2 image tokens + 4 language tokens + 1 state token + self.assertEqual(out.shape, (2, 9, 4)) + self.assertEqual(encoder.output_dim, 4) + self.assertEqual(encoder.tokens_per_step, 9) + self.assertEqual(encoder.condition_sequence_length, 9) + + self.assertEqual(fake_tokenizer.calls[-1]['text'], ['pick red cube\n', 'open drawer\n']) + self.assertEqual(fake_tokenizer.calls[-1]['padding'], 'max_length') + self.assertEqual(fake_tokenizer.calls[-1]['max_length'], 4) + self.assertEqual(fake_tokenizer.calls[-1]['padding_side'], 'right') + + # Camera order is r_vis then top, and pixels are mapped [0,1] -> [-1,1]. + first_camera_pixels = fake_vlm.model.vision_model.calls[0]['pixel_values'] + second_camera_pixels = fake_vlm.model.vision_model.calls[1]['pixel_values'] + self.assertTrue(torch.allclose(first_camera_pixels, torch.full((2, 3, 2, 2), -0.5))) + self.assertTrue(torch.allclose(second_camera_pixels, torch.full((2, 3, 2, 2), 0.5))) + + # Last token is padded last state projected from [4,5,6,0,0] and [10,11,12,0,0]. + self.assertTrue(torch.allclose(out[0, -1], torch.tensor([4.0, 5.0, 6.0, 0.0]))) + self.assertTrue(torch.allclose(out[1, -1], torch.tensor([10.0, 11.0, 12.0, 0.0]))) + + pad_mask, att_mask = encoder.last_prefix_pad_mask, encoder.last_prefix_att_mask + self.assertEqual(pad_mask.shape, (2, 9)) + self.assertEqual(att_mask.shape, (2, 9)) + self.assertTrue(torch.all(att_mask[:, :8] == 0)) + self.assertTrue(torch.all(att_mask[:, -1] == 1)) + + def test_forward_can_run_cropped_frozen_text_model_over_prefix_tokens(self): + from roboimi.vla.models.backbones.smolvla_prefix_encoder import SmolVLAPrefixEncoder + + fake_vlm = _FakeVLM() + fake_tokenizer = _FakeTokenizer() + encoder = SmolVLAPrefixEncoder( + vlm=fake_vlm, + tokenizer=fake_tokenizer, + num_vlm_layers=2, + freeze_vlm=True, + max_state_dim=5, + tokenizer_max_length=4, + camera_names=('r_vis',), + resize_imgs_with_padding=None, + run_text_model=True, + ) + images = { + 'r_vis': torch.full((2, 1, 3, 2, 2), 0.25), + } + state = torch.tensor([[[1.0, 2.0, 3.0]], [[4.0, 5.0, 6.0]]]) + + out = encoder(images, state=state, task=['pick', 'place']) + + # 1 camera * 2 image tokens + 4 language tokens + 1 state token. + self.assertEqual(out.shape, (2, 7, 4)) + self.assertEqual(len(fake_vlm.model.text_model.layers), 2) + self.assertEqual(len(fake_vlm.model.text_model.forward_calls), 1) + text_call = fake_vlm.model.text_model.forward_calls[-1] + self.assertEqual(tuple(text_call['attention_mask'].shape), (2, 1, 7, 7)) + self.assertEqual(tuple(text_call['position_ids'].shape), (2, 7)) + self.assertTrue(torch.all(out != 0)) + + def test_frozen_vlm_still_backpropagates_through_text_model_to_state_projection(self): + from roboimi.vla.models.backbones.smolvla_prefix_encoder import SmolVLAPrefixEncoder + + fake_vlm = _FakeVLM() + encoder = SmolVLAPrefixEncoder( + vlm=fake_vlm, + tokenizer=_FakeTokenizer(), + num_vlm_layers=3, + freeze_vlm=True, + train_state_proj=True, + max_state_dim=5, + tokenizer_max_length=4, + camera_names=('r_vis',), + resize_imgs_with_padding=None, + run_text_model=True, + ) + images = { + 'r_vis': torch.full((2, 1, 3, 2, 2), 0.25), + } + state = torch.tensor([[[1.0, 2.0, 3.0]], [[4.0, 5.0, 6.0]]]) + + out = encoder(images, state=state, task=['pick', 'place']) + out[:, -1].sum().backward() + + self.assertIsNotNone(encoder.state_proj.weight.grad) + self.assertGreater(float(encoder.state_proj.weight.grad.abs().sum()), 0.0) + self.assertTrue(all(param.grad is None for param in fake_vlm.parameters())) + + def test_embed_prefix_rejects_missing_camera_and_task_batch_mismatch(self): + from roboimi.vla.models.backbones.smolvla_prefix_encoder import SmolVLAPrefixEncoder + + encoder = SmolVLAPrefixEncoder( + vlm=_FakeVLM(), + tokenizer=_FakeTokenizer(), + num_vlm_layers=3, + camera_names=('r_vis', 'top'), + resize_imgs_with_padding=None, + max_state_dim=5, + tokenizer_max_length=4, + ) + images = {'r_vis': torch.rand(2, 1, 3, 2, 2)} + state = torch.rand(2, 1, 3) + with self.assertRaisesRegex(ValueError, 'missing.*top'): + encoder.embed_prefix(images=images, state=state, task=['a', 'b']) + images['top'] = torch.rand(2, 1, 3, 2, 2) + with self.assertRaisesRegex(ValueError, 'task batch'): + encoder.embed_prefix(images=images, state=state, task=['only one']) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/test_train_vla_rollout_validation.py b/tests/test_train_vla_rollout_validation.py index 2afd791..295ea80 100644 --- a/tests/test_train_vla_rollout_validation.py +++ b/tests/test_train_vla_rollout_validation.py @@ -14,6 +14,8 @@ from roboimi.demos.vla_scripts import eval_vla, train_vla class _FakeDataset: + available_episode_indices = [0, 1] + def __len__(self): return 4 @@ -29,6 +31,13 @@ class _FakeLoader: return iter(self._batches) +class _FakeValDataset(_FakeDataset): + available_episode_indices = [1] + + def __len__(self): + return 2 + + class _FakeOptimizer: def __init__(self, lr=1e-3): self.param_groups = [{'lr': lr}] @@ -91,6 +100,16 @@ class _FakeAgent(nn.Module): return {} +class _CapturingAgent(_FakeAgent): + def __init__(self): + super().__init__() + self.compute_loss_inputs = [] + + def compute_loss(self, agent_input): + self.compute_loss_inputs.append(agent_input) + return (self.weight - torch.tensor(0.5)).pow(2) + + class _SequentialLossAgent(nn.Module): def __init__(self, losses): super().__init__() @@ -150,6 +169,94 @@ class _FakeEvalEnv: class TrainVLARolloutValidationTest(unittest.TestCase): + def test_run_training_passes_variable_batch_task_to_agent_input(self): + cfg = OmegaConf.create( + { + 'train': { + 'device': 'cpu', + 'batch_size': 2, + 'num_workers': 0, + 'val_split': 0.0, + 'seed': 0, + 'lr': 1e-3, + 'max_steps': 1, + 'log_freq': 100, + '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': 0, + 'rollout_validate_on_checkpoint': False, + 'rollout_num_episodes': 1, + }, + 'data': { + 'camera_names': ['front'], + 'dataset_dir': 'unused', + }, + 'agent': { + '_target_': 'fake.agent', + 'normalization_type': 'min_max', + }, + 'eval': { + 'ckpt_path': 'unused.pt', + 'num_episodes': 1, + 'max_timesteps': 1, + 'device': 'cpu', + 'task_name': 'sim_transfer', + 'camera_names': ['front'], + 'use_smoothing': False, + 'smooth_alpha': 0.3, + 'verbose_action': False, + 'headless': True, + }, + 'experiment': {}, + } + ) + agent = _CapturingAgent() + batch_task = ['pick the red cube', 'insert the peg into the socket'] + + def fake_instantiate(config_node, **_kwargs): + if config_node is cfg.data: + return _FakeDataset() + if config_node is cfg.agent: + return agent + raise AssertionError(f'unexpected instantiate config: {config_node!r}') + + def fake_dataloader(_dataset, *, shuffle, **_kwargs): + del shuffle, _kwargs + return _FakeLoader( + { + 'observation.front': torch.zeros(2, 2, 3, 4, 4), + 'observation.state': torch.zeros(2, 2, 4), + 'action': torch.zeros(2, 8, 2), + 'action_is_pad': torch.zeros(2, 8, dtype=torch.bool), + 'task': list(batch_task), + }, + 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): + train_vla._run_training(cfg) + finally: + os.chdir(previous_cwd) + + self.assertEqual(len(agent.compute_loss_inputs), 1) + self.assertEqual(agent.compute_loss_inputs[0]['task'], batch_task) + def test_default_train_config_uses_full_dataset_and_epoch_rollout_validation(self): cfg = OmegaConf.load(Path('roboimi/vla/conf/config.yaml')) @@ -162,6 +269,39 @@ class TrainVLARolloutValidationTest(unittest.TestCase): self.assertIsNone(cfg.train.rollout_num_workers) self.assertIsNone(cfg.train.rollout_cuda_devices) + def test_explicit_val_episode_indices_builds_held_out_dataset(self): + cfg = OmegaConf.create( + { + 'train': { + 'val_episode_indices': [1], + 'val_split': 0.0, + 'seed': 42, + }, + 'data': {}, + } + ) + instantiate_calls = [] + + def fake_instantiate(config_node, **kwargs): + del config_node + instantiate_calls.append(dict(kwargs)) + if kwargs.get('episode_indices') == [1]: + return _FakeValDataset() + return _FakeDataset() + + with mock.patch.object(train_vla, 'instantiate', side_effect=fake_instantiate): + dataset, train_dataset, val_dataset, explicit = train_vla.build_train_val_datasets( + cfg, + dataset_image_resize_shape=None, + ) + + self.assertIsInstance(dataset, _FakeDataset) + self.assertIsInstance(train_dataset, _FakeDataset) + self.assertIsInstance(val_dataset, _FakeValDataset) + self.assertEqual(explicit, [1]) + self.assertEqual(instantiate_calls[1]['episode_indices'], [0]) + self.assertEqual(instantiate_calls[2]['episode_indices'], [1]) + def test_run_training_rollout_validation_propagates_gpu_parallel_settings(self): cfg = OmegaConf.create( @@ -340,6 +480,93 @@ class TrainVLARolloutValidationTest(unittest.TestCase): self.assertIn('image_resize_shape', captured_dataset_kwargs) self.assertIsNone(captured_dataset_kwargs['image_resize_shape']) + def test_training_passes_condition_encoder_image_resize_override_to_dataset_instantiation(self): + cfg = OmegaConf.create( + { + 'agent': { + 'condition_encoder': { + 'dataset_image_resize_shape': None, + }, + 'normalization_type': 'min_max', + }, + 'data': { + 'dataset_dir': 'unused', + 'camera_names': ['front'], + 'image_resize_shape': [224, 224], + }, + 'train': { + 'batch_size': 2, + 'lr': 1e-4, + 'max_steps': 0, + 'device': 'cpu', + 'disable_cudnn': False, + 'num_workers': 0, + 'val_split': 0.0, + 'seed': 42, + 'log_freq': 1, + 'save_freq': 10, + 'use_swanlab': False, + 'rollout_val_freq_epochs': 0, + 'rollout_validate_on_checkpoint': False, + 'rollout_num_episodes': 1, + 'warmup_steps': 1, + 'scheduler_type': 'constant', + 'min_lr': 1e-6, + 'weight_decay': 1e-5, + 'grad_clip': 1.0, + 'pretrained_ckpt': None, + }, + 'eval': { + 'ckpt_path': 'unused.pt', + 'num_episodes': 1, + 'headless': True, + 'device': 'cpu', + 'verbose_action': False, + }, + 'experiment': {}, + } + ) + captured_dataset_kwargs = {} + + def fake_instantiate(config_node, **kwargs): + if config_node is cfg.data: + captured_dataset_kwargs.update(kwargs) + 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, '_init_swanlab', return_value=None), \ + mock.patch.object(train_vla, '_finish_swanlab', return_value=None), \ + mock.patch.object(train_vla.torch, 'save', return_value=None): + train_vla._run_training(cfg) + finally: + os.chdir(previous_cwd) + + self.assertIn('image_resize_shape', captured_dataset_kwargs) + self.assertIsNone(captured_dataset_kwargs['image_resize_shape']) + def test_eval_main_delegates_to_plain_run_eval_helper(self): cfg = OmegaConf.create( { diff --git a/tests/test_train_vla_transformer_optimizer.py b/tests/test_train_vla_transformer_optimizer.py index 56ee107..3c893f4 100644 --- a/tests/test_train_vla_transformer_optimizer.py +++ b/tests/test_train_vla_transformer_optimizer.py @@ -422,3 +422,45 @@ class TrainVLATransformerOptimizerTest(unittest.TestCase): if __name__ == '__main__': unittest.main() + +class TrainVLASmolVLAOptimizerTest(unittest.TestCase): + def test_build_training_optimizer_excludes_frozen_vlm_parameters_and_keeps_state_proj_and_head(self): + module = TrainVLATransformerOptimizerTest()._load_train_vla_module() + + class _Head(nn.Module): + def __init__(self): + super().__init__() + self.proj = nn.Linear(2, 2) + + def get_optim_groups(self, weight_decay): + return [{'params': list(self.parameters()), 'weight_decay': weight_decay}] + + class _ConditionEncoder(nn.Module): + def __init__(self): + super().__init__() + self.vlm = nn.Linear(2, 2) + for param in self.vlm.parameters(): + param.requires_grad = False + self.state_proj = nn.Linear(2, 2) + + class _Agent(nn.Module): + def __init__(self): + super().__init__() + self.noise_pred_net = _Head() + self.condition_encoder = _ConditionEncoder() + + agent = _Agent() + with mock.patch.object(module, 'AdamW', RecordingAdamW): + optimizer = module.build_training_optimizer(agent, lr=1e-4, weight_decay=0.01) + + names_by_param_id = {id(param): name for name, param in agent.named_parameters()} + optimizer_names = { + names_by_param_id[id(param)] + for group in optimizer.param_groups + for param in group['params'] + } + self.assertIn('condition_encoder.state_proj.weight', optimizer_names) + self.assertIn('condition_encoder.state_proj.bias', optimizer_names) + self.assertIn('noise_pred_net.proj.weight', optimizer_names) + self.assertNotIn('condition_encoder.vlm.weight', optimizer_names) + self.assertNotIn('condition_encoder.vlm.bias', optimizer_names)