feat(vla): add SmolVLA conditioned agent with prefix encoder and train/eval validation

Introduce SmolVLA-style VLM prefix encoder backbone and IMF-AttnRes
conditioned agent. Add episode-level train/val split, action MSE
validation in training, and headless eval support.

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）

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)
+            for backbone_key in ('vision_backbone', 'condition_encoder'):
-            if vision_backbone_cfg is not None and 'dataset_image_resize_shape' in vision_backbone_cfg:
+                backbone_cfg = cfg.agent.get(backbone_key, None)
-                dataset_image_resize_shape = vision_backbone_cfg.get('dataset_image_resize_shape')
+                if backbone_cfg is not None and 'dataset_image_resize_shape' in backbone_cfg:
-            dataset = instantiate(
+                    dataset_image_resize_shape = backbone_cfg.get('dataset_image_resize_shape')
-                cfg.data,
+                    break
-                image_resize_shape=dataset_image_resize_shape,
+            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
 
-        # 训练/验证集划分
+        if explicit_val_episode_indices is not None:
-        val_split = float(cfg.train.get('val_split', 0.1))
+            log.info(
-        seed = int(cfg.train.get('seed', 42))
+                "✅ 数据集划分: 训练集=%s, 验证集=%s (显式 held-out episodes=%s)",
-        val_size = int(len(dataset) * val_split)
+                len(train_dataset),
-        train_size = len(dataset) - val_size
+                len(val_dataset),
-        if val_size > 0:
+                explicit_val_episode_indices,
-            train_dataset, val_dataset = random_split(
-                dataset,
-                [train_size, val_size],
-                generator=torch.Generator().manual_seed(seed)
             )
-            log.info(f"✅ 数据集划分: 训练集={train_size}, 验证集={val_size} (验证比例={val_split})")
         else:
-            train_dataset, val_dataset = dataset, None
+            val_split = float(cfg.train.get('val_split', 0.1))
-            log.info("✅ 数据集划分: 全部用于训练, 验证集=0 (验证比例=0)")
+            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

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()

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

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

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 时注入给语言条件策略
 
 # ====================
 # 策略执行参数

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:
         """获取相机信息"""

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

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(

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],

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()

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()

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()

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(
             {

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)