Add native SmolVLA model integration

docs/superpowers/plans/2026-05-25-native-smolvla-model-migration.md

@@ -0,0 +1,184 @@
+# Native SmolVLA Model Migration Implementation Plan
+
+> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.
+
+**Goal:** Migrate only the SmolVLA model core into RoboIMI and expose it as a native VLA agent without importing the full LeRobot package.
+
+**Architecture:** Add a focused `roboimi.vla.models.smolvla` model package and a `SmolVLANativeAgent` wrapper that speaks the existing RoboIMI train/eval interface. Keep normalization, queues, camera ordering, and task fallback in the agent; keep model math in the migrated model package.
+
+**Tech Stack:** PyTorch, Transformers, Hydra/OmegaConf, unittest/pytest-style tests, existing RoboIMI normalization and VLA scripts.
+
+---
+
+## File Structure
+
+- Create `roboimi/vla/models/smolvla/__init__.py`: package exports.
+- Create `roboimi/vla/models/smolvla/configuration.py`: lightweight config dataclass.
+- Create `roboimi/vla/models/smolvla/modeling.py`: model utilities and `VLAFlowMatching`.
+- Create `roboimi/vla/models/smolvla/smolvlm_with_expert.py`: adapted VLM/expert core.
+- Create `roboimi/vla/agent_smolvla_native.py`: RoboIMI agent wrapper.
+- Create `roboimi/vla/conf/agent/smolvla_native.yaml`: Hydra config.
+- Create `tests/test_smolvla_native_agent.py`: TDD tests for wrapper behavior.
+- Create `tests/test_smolvla_native_modeling.py`: TDD tests for utilities/config where useful.
+
+## Task 1: Wrapper behavior tests and minimal agent skeleton
+
+**Files:**
+- Create: `tests/test_smolvla_native_agent.py`
+- Create: `roboimi/vla/agent_smolvla_native.py`
+
+- [ ] **Step 1: Write failing tests for task fallback, camera order, queue, and fake model calls**
+
+Implement tests using fake tokenizer and fake model. Test names:
+
+- `test_compute_loss_orders_cameras_tokenizes_task_and_calls_native_model`
+- `test_unknown_task_uses_configured_task_description`
+- `test_missing_camera_raises_clear_error`
+- `test_predict_action_chunk_denormalizes_fake_model_output`
+- `test_select_action_uses_action_queue_before_recomputing`
+
+- [ ] **Step 2: Run tests and verify import failure**
+
+Run: `python -m unittest tests.test_smolvla_native_agent -v`
+
+Expected: FAIL because `roboimi.vla.agent_smolvla_native` does not exist.
+
+- [ ] **Step 3: Implement minimal `SmolVLANativeAgent` skeleton**
+
+Implement constructor injection for `model` and `tokenizer`, plus `_resolve_task`, `_order_images`, `_tokenize_tasks`, `reset`, `_prepare_observation_batch`, `compute_loss`, `predict_action_chunk`, and `select_action`. Use fake model in tests; real model construction can raise a clear ImportError until Task 3.
+
+- [ ] **Step 4: Run tests and verify pass**
+
+Run: `python -m unittest tests.test_smolvla_native_agent -v`
+
+Expected: PASS.
+
+## Task 2: Config and modeling utility tests
+
+**Files:**
+- Create: `tests/test_smolvla_native_modeling.py`
+- Create: `roboimi/vla/models/smolvla/configuration.py`
+- Create: `roboimi/vla/models/smolvla/modeling.py`
+- Create: `roboimi/vla/models/smolvla/__init__.py`
+
+- [ ] **Step 1: Write failing tests**
+
+Cover:
+
+- `NativeSmolVLAConfig` validates `n_action_steps <= chunk_size`.
+- `pad_vector` pads and rejects truncation.
+- `resize_with_pad` preserves batch/channel shape and output size.
+- `make_att_2d_masks` matches prefix-LM mask semantics.
+
+- [ ] **Step 2: Run tests and verify failure**
+
+Run: `python -m unittest tests.test_smolvla_native_modeling -v`
+
+Expected: FAIL because package/modeling code is missing.
+
+- [ ] **Step 3: Implement config and utility functions**
+
+Port minimal functions from external SmolVLA while removing LeRobot imports.
+
+- [ ] **Step 4: Run tests and verify pass**
+
+Run: `python -m unittest tests.test_smolvla_native_modeling -v`
+
+Expected: PASS.
+
+## Task 3: Migrate model core
+
+**Files:**
+- Modify: `roboimi/vla/models/smolvla/modeling.py`
+- Create: `roboimi/vla/models/smolvla/smolvlm_with_expert.py`
+- Modify: `roboimi/vla/agent_smolvla_native.py`
+
+- [ ] **Step 1: Write fake-backed integration tests for lazy real model construction**
+
+Patch `VLAFlowMatching` and tokenizer loader so `SmolVLANativeAgent(model=None, tokenizer=None)` constructs a fake model without real downloads. Verify config fields are passed.
+
+- [ ] **Step 2: Run test and verify failure**
+
+Run: `python -m unittest tests.test_smolvla_native_agent -v`
+
+Expected: FAIL because real construction path is not implemented.
+
+- [ ] **Step 3: Port `SmolVLMWithExpertModel` and `VLAFlowMatching`**
+
+Copy only model-core logic from external files, replacing LeRobot dependencies with local config and local helpers. Preserve `forward`, `sample_actions`, `denoise_step`, image resize/pad, state/action padding, and language token inputs.
+
+- [ ] **Step 4: Wire agent lazy construction**
+
+If `model` is not supplied, create `NativeSmolVLAConfig`, tokenizer, and `VLAFlowMatching`. Keep `model` injection path for tests.
+
+- [ ] **Step 5: Run unit tests**
+
+Run:
+
+```bash
+python -m unittest tests.test_smolvla_native_agent tests.test_smolvla_native_modeling -v
+```
+
+Expected: PASS.
+
+## Task 4: Hydra config
+
+**Files:**
+- Create: `roboimi/vla/conf/agent/smolvla_native.yaml`
+- Modify: `tests/test_smolvla_native_agent.py`
+
+- [ ] **Step 1: Add failing Hydra compose test**
+
+Test composing `agent=smolvla_native` exposes `_target_`, `action_dim=16`, `obs_dim=16`, `camera_names=${data.camera_names}`, and model config fields.
+
+- [ ] **Step 2: Run and verify failure**
+
+Run: `python -m unittest tests.test_smolvla_native_agent -v`
+
+Expected: FAIL because yaml is missing.
+
+- [ ] **Step 3: Add yaml config**
+
+Add `smolvla_native.yaml` with `_target_: roboimi.vla.agent_smolvla_native.SmolVLANativeAgent` and sane defaults.
+
+- [ ] **Step 4: Run test and verify pass**
+
+Run: `python -m unittest tests.test_smolvla_native_agent -v`
+
+Expected: PASS.
+
+## Task 5: Regression checks
+
+**Files:**
+- No production changes expected unless tests reveal integration gaps.
+
+- [ ] **Step 1: Run focused existing tests**
+
+Run:
+
+```bash
+python -m unittest tests.test_smolvla_prefix_encoder tests.test_smolvla_imf_agent tests.test_eval_vla_execution -v
+```
+
+Expected: PASS.
+
+- [ ] **Step 2: Search for forbidden import**
+
+Run:
+
+```bash
+rg -n "import lerobot|from lerobot" roboimi/vla/models/smolvla roboimi/vla/agent_smolvla_native.py
+```
+
+Expected: no matches.
+
+- [ ] **Step 3: Commit**
+
+Run:
+
+```bash
+git add roboimi/vla/models/smolvla roboimi/vla/agent_smolvla_native.py roboimi/vla/conf/agent/smolvla_native.yaml tests/test_smolvla_native_agent.py tests/test_smolvla_native_modeling.py docs/superpowers/specs/2026-05-25-native-smolvla-model-migration-design.md docs/superpowers/plans/2026-05-25-native-smolvla-model-migration.md
+git commit -m "feat(vla): add native SmolVLA model agent"
+```
+
+Expected: commit succeeds.

docs/superpowers/specs/2026-05-25-native-smolvla-model-migration-design.md

@@ -0,0 +1,86 @@
+# Native SmolVLA Model Migration Design
+
+## Goal
+
+Migrate only the SmolVLA model core from `/data/lerobot-imf-attnres-exp/lerobot-imf-attnres` into this `roboimi` project so Diana simulation can train/evaluate through the existing `roboimi.vla` agent interface without importing or installing the full LeRobot tree.
+
+## Non-goals
+
+- Do not vendor the full `lerobot` package.
+- Do not change the current Python environment to LeRobot 0.5.x requirements.
+- Do not alter Diana environment action semantics.
+- Do not change `train_vla.py` or `eval_vla.py` main control flow unless a minimal compatibility hook is unavoidable.
+- Do not implement RTC in the first pass.
+
+## Architecture
+
+Add a native RoboIMI SmolVLA package under `roboimi/vla/models/smolvla/`. It will contain a lightweight config dataclass, a minimally adapted `SmolVLMWithExpertModel`, and a `VLAFlowMatching` implementation. A new `SmolVLANativeAgent` wraps the model with the existing RoboIMI agent contract: `compute_loss`, `predict_action_chunk`, `select_action`, `reset`, and `get_normalization_stats`.
+
+The new code will preserve the original model math where practical, but replace LeRobot framework dependencies with local constants, tokenizer handling, queue management, and `roboimi.vla.models.normalization.NormalizationModule`.
+
+## Data flow
+
+Training batch input remains the current RoboIMI format:
+
+```python
+{
+  "images": {cam: Tensor[B, T, C, H, W]},
+  "qpos": Tensor[B, T, 16],
+  "action": Tensor[B, H, 16],
+  "action_is_pad": optional BoolTensor[B, H],
+  "task": optional list[str],
+}
+```
+
+`SmolVLANativeAgent` normalizes `qpos` and `action` using RoboIMI dataset stats, tokenizes task strings, and passes images/state/language/action into the native SmolVLA model. In inference, the model returns normalized action chunks; the agent denormalizes them to 16-dim Diana EE actions.
+
+## Components
+
+### `roboimi/vla/models/smolvla/configuration.py`
+
+Defines `NativeSmolVLAConfig`, a small dataclass with fields needed by the model: state/action dimensions, padding dimensions, image resize target, tokenizer settings, VLM model name, VLM loading flags, expert layer configuration, sampling steps, dtype/device behavior, and compile flags.
+
+### `roboimi/vla/models/smolvla/smolvlm_with_expert.py`
+
+Migrates the model-core helper from the old repo. It should depend only on PyTorch and Transformers. It will expose `SmolVLMWithExpertModel` and attention helpers.
+
+### `roboimi/vla/models/smolvla/modeling.py`
+
+Defines model-core utilities (`resize_with_pad`, `pad_vector`, `make_att_2d_masks`, sinusoidal time embedding) plus `VLAFlowMatching`, with `forward` and `sample_actions`.
+
+### `roboimi/vla/agent_smolvla_native.py`
+
+RoboIMI-native agent wrapper. It owns tokenizer, normalization, task fallback, camera ordering, observation/action queues, loss masking, and denormalized rollout actions.
+
+### `roboimi/vla/conf/agent/smolvla_native.yaml`
+
+Hydra config for the native model, defaulting to Diana 16-dim state/action and configurable camera names.
+
+## Error handling
+
+- Missing configured camera raises `ValueError` with expected/missing names.
+- Task list length not matching batch size raises `ValueError`.
+- State/action dimensions exceeding configured `max_state_dim` / `max_action_dim` raises `ValueError`.
+- Missing Transformers SmolVLM classes raises `ImportError` explaining the required package.
+- Invalid action chunk shape raises `RuntimeError` explaining expected `(B,H,A)`.
+
+## Testing
+
+Use TDD with fake VLM/tokenizer/model components first, so tests do not download or instantiate the real SmolVLM. Cover:
+
+1. Config and Hydra instantiation with fake injected components.
+2. Camera ordering and missing-camera errors.
+3. Task fallback and newline/tokenization behavior.
+4. Loss path shape/mask behavior using a fake native model.
+5. `predict_action_chunk` normalization/denormalization and shape.
+6. `select_action` queue behavior.
+
+A later smoke test may instantiate the real model in an environment that already has the required Transformers/weights, but the first implementation must pass without external downloads.
+
+## Acceptance criteria
+
+- `agent=smolvla_native` can be composed by Hydra.
+- Unit tests pass without importing external `/data/.../src/lerobot`.
+- The new production code has no `import lerobot`.
+- The agent accepts the same batch/observation structure used by current train/eval scripts.
+- The agent emits 16-dim denormalized Diana EE actions for rollout.

roboimi/demos/vla_scripts/eval_vla.py

@@ -160,6 +160,8 @@ def _normalize_resize_shape(shape) -> Optional[tuple[int, int]]:
 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
+    if 'eval_image_resize_shape' in agent_cfg:
+        return _normalize_resize_shape(agent_cfg.get('eval_image_resize_shape'))
     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:
@@ -242,13 +244,14 @@ def _enqueue_predicted_actions(
     predicted_actions: Any,
     obs_horizon: int,
     num_action_steps: int,
+    action_chunk_start: Optional[int] = None,
 ) -> None:
     if isinstance(predicted_actions, np.ndarray):
         predicted_actions = torch.from_numpy(predicted_actions)
     if predicted_actions.ndim == 2:
         predicted_actions = predicted_actions.unsqueeze(0)
 
-    start = int(obs_horizon) - 1
+    start = int(obs_horizon) - 1 if action_chunk_start is None else int(action_chunk_start)
     end = start + int(num_action_steps)
     executable_actions = predicted_actions[:, start:end]
     for action_index in range(executable_actions.shape[1]):
@@ -316,6 +319,7 @@ class _RemotePolicyRunner:
         camera_names: list[str],
         obs_horizon: int,
         num_action_steps: int,
+        action_chunk_start: Optional[int] = None,
         response_timeout_s: float = 30.0,
     ):
         self.worker_index = int(worker_index)
@@ -325,6 +329,7 @@ class _RemotePolicyRunner:
         self.camera_names = list(camera_names)
         self.obs_horizon = int(obs_horizon)
         self.num_action_steps = int(num_action_steps)
+        self.action_chunk_start = None if action_chunk_start is None else int(action_chunk_start)
         self.response_timeout_s = float(response_timeout_s)
         self.local_queues = _new_local_policy_queues(self.obs_horizon)
         self.uses_local_model = False
@@ -371,6 +376,7 @@ class _RemotePolicyRunner:
                 predicted_actions=response['actions'],
                 obs_horizon=self.obs_horizon,
                 num_action_steps=self.num_action_steps,
+                action_chunk_start=self.action_chunk_start,
             )
 
         if not self.local_queues['action']:
@@ -1407,6 +1413,7 @@ def _run_remote_eval_worker(
     eval_cfg = cfg.eval
     agent_cfg = cfg.agent
     num_action_steps = int(agent_cfg.get('num_action_steps', eval_cfg.get('num_queries', 1)))
+    action_chunk_start = agent_cfg.get('action_chunk_start', None)
     policy_runner = _RemotePolicyRunner(
         worker_index=worker_index,
         server_index=server_index,
@@ -1415,6 +1422,7 @@ def _run_remote_eval_worker(
         camera_names=_resolve_policy_camera_names(cfg),
         obs_horizon=int(agent_cfg.get('obs_horizon', eval_cfg.obs_horizon)),
         num_action_steps=num_action_steps,
+        action_chunk_start=action_chunk_start,
         response_timeout_s=float(eval_cfg.get('response_timeout_s', 300.0)),
     )
     return _run_eval_episode_plans(

roboimi/demos/vla_scripts/train_vla.py

@@ -191,6 +191,19 @@ def _instantiate_dataset(cfg, dataset_image_resize_shape, episode_indices=None):
     return instantiate(cfg.data, **kwargs)
 
 
+def _resolve_dataset_image_resize_shape(cfg):
+    dataset_image_resize_shape = cfg.data.get('image_resize_shape', (224, 224))
+    agent_cfg = cfg.agent
+    if 'dataset_image_resize_shape' in agent_cfg:
+        return agent_cfg.get('dataset_image_resize_shape')
+    for backbone_key in ('vision_backbone', 'condition_encoder'):
+        backbone_cfg = agent_cfg.get(backbone_key, None)
+        if backbone_cfg is not None and 'dataset_image_resize_shape' in backbone_cfg:
+            dataset_image_resize_shape = backbone_cfg.get('dataset_image_resize_shape')
+            break
+    return dataset_image_resize_shape
+
+
 def build_train_val_datasets(cfg, dataset_image_resize_shape):
     val_episode_indices = cfg.train.get('val_episode_indices', None)
     if val_episode_indices:
@@ -444,12 +457,7 @@ def _run_training(cfg: DictConfig):
         # =========================================================================
         log.info("📦 加载数据集...")
         try:
-            dataset_image_resize_shape = cfg.data.get('image_resize_shape', (224, 224))
-            for backbone_key in ('vision_backbone', 'condition_encoder'):
-                backbone_cfg = cfg.agent.get(backbone_key, None)
-                if backbone_cfg is not None and 'dataset_image_resize_shape' in backbone_cfg:
-                    dataset_image_resize_shape = backbone_cfg.get('dataset_image_resize_shape')
-                    break
+            dataset_image_resize_shape = _resolve_dataset_image_resize_shape(cfg)
             dataset, train_dataset, val_dataset, explicit_val_episode_indices = (
                 build_train_val_datasets(cfg, dataset_image_resize_shape)
             )

roboimi/vla/agent_smolvla_native.py

@@ -0,0 +1,395 @@
+from __future__ import annotations
+
+from collections import deque
+from dataclasses import fields
+from typing import Dict, Optional, Sequence
+
+import torch
+import torch.nn as nn
+
+from roboimi.vla.models.normalization import NormalizationModule
+
+
+class SmolVLANativeAgent(nn.Module):
+    """RoboIMI wrapper for the native SmolVLA flow-matching model.
+
+    The wrapper owns RoboIMI-facing concerns only: dataset normalization,
+    camera ordering, language fallback/tokenization, rollout queues, and action
+    denormalization.  The native SmolVLA model itself is imported lazily so unit
+    tests can inject fakes without downloading or loading a real VLM.
+    """
+
+    def __init__(
+        self,
+        model: Optional[nn.Module] = None,
+        tokenizer=None,
+        action_dim: int = 16,
+        obs_dim: int = 16,
+        chunk_size: int = 16,
+        n_action_steps: int = 8,
+        obs_horizon: int = 1,
+        action_chunk_start: int = 0,
+        num_cams: int = 3,
+        camera_names: Optional[Sequence[str]] = None,
+        dataset_stats=None,
+        normalization_type: str = 'min_max',
+        task_description: Optional[str] = None,
+        model_config: Optional[dict] = None,
+        tokenizer_name: Optional[str] = None,
+        **kwargs,
+    ) -> None:
+        super().__init__()
+        del kwargs
+        self.action_dim = int(action_dim)
+        self.obs_dim = int(obs_dim)
+        self.chunk_size = int(chunk_size)
+        self.pred_horizon = self.chunk_size
+        self.n_action_steps = int(n_action_steps)
+        self.num_action_steps = self.n_action_steps
+        self.obs_horizon = int(obs_horizon)
+        self.action_chunk_start = int(action_chunk_start)
+        self.num_cams = int(num_cams)
+        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})')
+        if self.n_action_steps < 1:
+            raise ValueError('n_action_steps must be >= 1')
+        if self.action_chunk_start < 0:
+            raise ValueError('action_chunk_start must be >= 0')
+        if self.action_chunk_start + self.n_action_steps > self.chunk_size:
+            raise ValueError('action_chunk_start + n_action_steps must be <= chunk_size')
+
+        self.normalization = NormalizationModule(stats=dataset_stats, normalization_type=normalization_type)
+        self.task_description = task_description
+        self.model_config = dict(model_config or {})
+        self.max_state_dim = int(self.model_config.get('max_state_dim', self.obs_dim))
+        self.max_action_dim = int(self.model_config.get('max_action_dim', self.action_dim))
+        self.resize_imgs_with_padding = self._normalize_resize_shape(
+            self.model_config.get('resize_imgs_with_padding', self.model_config.get('image_resize_shape', (512, 512)))
+        )
+        self.tokenizer_max_length = int(self.model_config.get('tokenizer_max_length', 48))
+        self.pad_language_to = str(self.model_config.get('pad_language_to', 'longest'))
+        self.tokenizer_name = tokenizer_name
+        self.tokenizer = tokenizer if tokenizer is not None else self._build_tokenizer(tokenizer_name)
+        self.model = model if model is not None else self._build_model()
+        self.reset()
+
+    @staticmethod
+    def _normalize_resize_shape(shape):
+        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
+
+    def _build_tokenizer(self, tokenizer_name):
+        if tokenizer_name is None:
+            tokenizer_name = self.model_config.get('tokenizer_name') or self.model_config.get('vlm_model_name')
+        if tokenizer_name is None:
+            raise ImportError('A tokenizer or tokenizer_name/model_config.vlm_model_name is required for SmolVLANativeAgent')
+        try:
+            from transformers import AutoTokenizer
+        except ImportError as exc:
+            raise ImportError('SmolVLANativeAgent requires transformers to construct the real tokenizer') from exc
+        return AutoTokenizer.from_pretrained(tokenizer_name)
+
+    def _native_config_kwargs(self) -> dict:
+        from roboimi.vla.models.smolvla.configuration import NativeSmolVLAConfig
+
+        cfg_kwargs = dict(self.model_config)
+
+        # Backwards-compatible aliases from earlier RoboIMI-facing drafts.  The
+        # native model core intentionally keeps only SmolVLA model fields.
+        if 'image_resize_shape' in cfg_kwargs and 'resize_imgs_with_padding' not in cfg_kwargs:
+            cfg_kwargs['resize_imgs_with_padding'] = cfg_kwargs['image_resize_shape']
+        if 'freeze_vlm' in cfg_kwargs and 'train_expert_only' not in cfg_kwargs:
+            cfg_kwargs['train_expert_only'] = bool(cfg_kwargs['freeze_vlm'])
+
+        for wrapper_only_key in (
+            'state_dim',
+            'action_dim',
+            'tokenizer_name',
+            'freeze_vlm',
+            'image_resize_shape',
+            'num_cameras',
+        ):
+            cfg_kwargs.pop(wrapper_only_key, None)
+
+        cfg_kwargs.setdefault('max_state_dim', self.max_state_dim)
+        cfg_kwargs.setdefault('max_action_dim', self.max_action_dim)
+        cfg_kwargs.setdefault('chunk_size', self.chunk_size)
+        cfg_kwargs.setdefault('n_action_steps', self.n_action_steps)
+        cfg_kwargs['resize_imgs_with_padding'] = self._normalize_resize_shape(
+            cfg_kwargs.get('resize_imgs_with_padding', self.resize_imgs_with_padding)
+        )
+
+        allowed_keys = {field.name for field in fields(NativeSmolVLAConfig)}
+        return {key: value for key, value in cfg_kwargs.items() if key in allowed_keys}
+
+    def _build_model(self):
+        try:
+            from roboimi.vla.models.smolvla.configuration import NativeSmolVLAConfig
+            from roboimi.vla.models.smolvla.modeling import VLAFlowMatching
+        except ImportError as exc:
+            raise ImportError(
+                'Native SmolVLA model modules are unavailable. Pass injected model/tokenizer fakes in tests, '
+                'or add roboimi.vla.models.smolvla.configuration/modeling for real construction.'
+            ) from exc
+        config = NativeSmolVLAConfig(**self._native_config_kwargs())
+        return VLAFlowMatching(config=config)
+
+    def _get_model_device(self) -> torch.device:
+        try:
+            return next(self.model.parameters()).device
+        except (StopIteration, AttributeError):
+            return torch.device('cpu')
+
+    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 batch missing required cameras. missing={missing}, expected={list(self.camera_names)}')
+        return {name: images[name] for name in self.camera_names}
+
+    @staticmethod
+    def _is_missing_task(item) -> bool:
+        return item is None or (isinstance(item, str) and item.strip().lower() in {'', 'unknown'})
+
+    def _resolve_task(self, task, batch_size: int):
+        def fallback_or(item):
+            if self._is_missing_task(item):
+                if self.task_description is not None:
+                    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 _tokenize_tasks(self, task, batch_size: int, device: torch.device) -> dict:
+        tasks = [str(text) + '\n' for text in self._resolve_task(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,
+                truncation=True,
+                return_tensors='pt',
+            )
+        finally:
+            if old_padding_side is not None:
+                self.tokenizer.padding_side = old_padding_side
+        return {
+            'lang_tokens': tokenized['input_ids'].to(device=device),
+            'lang_masks': tokenized['attention_mask'].to(device=device, dtype=torch.bool),
+        }
+
+    @staticmethod
+    def _pad_vector(vector: torch.Tensor, new_dim: int) -> torch.Tensor:
+        current_dim = vector.shape[-1]
+        if current_dim == new_dim:
+            return vector
+        if current_dim > new_dim:
+            raise ValueError(f'cannot pad vector with dim {current_dim} to smaller dim {new_dim}')
+        padded_shape = list(vector.shape)
+        padded_shape[-1] = int(new_dim)
+        padded = torch.zeros(*padded_shape, dtype=vector.dtype, device=vector.device)
+        padded[..., :current_dim] = vector
+        return padded
+
+    @staticmethod
+    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)}')
+        import torch.nn.functional as F
+
+        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 _prepare_native_images(self, images: Dict[str, torch.Tensor]) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
+        ordered_images = self._order_images(images)
+        prepared_images: list[torch.Tensor] = []
+        img_masks: list[torch.Tensor] = []
+        reference_batch_size: int | None = None
+        for camera_name, image in ordered_images.items():
+            if image.ndim == 5:
+                image = image[:, -1]
+            elif image.ndim != 4:
+                raise ValueError(
+                    f'image for camera {camera_name!r} must be shaped (B,T,C,H,W) or (B,C,H,W), got {tuple(image.shape)}'
+                )
+            if reference_batch_size is None:
+                reference_batch_size = int(image.shape[0])
+            elif int(image.shape[0]) != reference_batch_size:
+                raise ValueError(f'image batch size mismatch for camera {camera_name!r}')
+
+            image = image.contiguous().float().clamp(0.0, 1.0)
+            if self.resize_imgs_with_padding is not None:
+                image = self._resize_with_pad(image, *self.resize_imgs_with_padding, pad_value=0.0)
+            image = image * 2.0 - 1.0
+            prepared_images.append(image)
+            img_masks.append(torch.ones(image.shape[0], dtype=torch.bool, device=image.device))
+        return prepared_images, img_masks
+
+    def _prepare_native_state(self, states: torch.Tensor) -> torch.Tensor:
+        states = self.normalization.normalize_qpos(states)
+        if states.ndim > 2:
+            states = states[:, -1, :]
+        return self._pad_vector(states.float(), self.max_state_dim)
+
+    def _prepare_native_actions(self, actions: torch.Tensor) -> torch.Tensor:
+        actions = self.normalization.normalize_action(actions)
+        return self._pad_vector(actions.float(), self.max_action_dim)
+
+    def _prepare_model_inputs(self, batch: Dict[str, torch.Tensor], include_actions: bool = False) -> dict:
+        states = batch['qpos']
+        batch_size = states.shape[0]
+        device = states.device
+        images, img_masks = self._prepare_native_images(batch['images'])
+        inputs = {
+            'images': images,
+            'img_masks': img_masks,
+            'state': self._prepare_native_state(states),
+        }
+        inputs.update(self._tokenize_tasks(batch.get('task', None), batch_size, device))
+        if include_actions:
+            inputs['actions'] = self._prepare_native_actions(batch['action'])
+        return inputs
+
+    def _reduce_native_losses(self, losses: torch.Tensor, action_is_pad: torch.Tensor | None = None) -> torch.Tensor:
+        if losses.ndim == 0:
+            return losses
+        if losses.shape[-1] < self.action_dim:
+            raise RuntimeError(f'loss action dim mismatch: got {losses.shape[-1]}, expected at least {self.action_dim}')
+        losses = losses[..., : self.action_dim]
+        if action_is_pad is None:
+            return losses.mean()
+        if losses.ndim != 3:
+            raise RuntimeError(f'action padding mask requires per-element losses shaped (B,H,A), got {tuple(losses.shape)}')
+        if tuple(action_is_pad.shape) != tuple(losses.shape[:2]):
+            raise RuntimeError(
+                f'action_is_pad shape {tuple(action_is_pad.shape)} does not match loss prefix {tuple(losses.shape[:2])}'
+            )
+        mask = (~action_is_pad).to(device=losses.device, dtype=losses.dtype).unsqueeze(-1)
+        denom = (mask.sum() * losses.shape[-1]).clamp_min(1.0)
+        return (losses * mask).sum() / denom
+
+    def compute_loss(self, batch):
+        inputs = self._prepare_model_inputs(batch, include_actions=True)
+        action_is_pad = batch.get('action_is_pad', None)
+        output = self.model(**inputs)
+        if isinstance(output, dict):
+            if 'loss' not in output:
+                raise RuntimeError('SmolVLA model forward returned a dict without loss')
+            return output['loss']
+        if torch.is_tensor(output):
+            return self._reduce_native_losses(output, action_is_pad=action_is_pad)
+        if hasattr(output, 'loss'):
+            return output.loss
+        raise RuntimeError(f'Unsupported SmolVLA forward output type: {type(output)!r}')
+
+    @torch.no_grad()
+    def predict_action_chunk(self, batch: Dict[str, torch.Tensor]) -> torch.Tensor:
+        inputs = self._prepare_model_inputs(batch, include_actions=False)
+        if not hasattr(self.model, 'sample_actions'):
+            raise RuntimeError('SmolVLA model must implement sample_actions for inference')
+        actions = self.model.sample_actions(**inputs)
+        if not torch.is_tensor(actions) or actions.ndim != 3:
+            raise RuntimeError(f'sample_actions must return (B,H,A), got {type(actions)!r} {getattr(actions, "shape", None)}')
+        if actions.shape[-1] < self.action_dim:
+            raise RuntimeError(f'action dim mismatch: got {actions.shape[-1]}, expected at least {self.action_dim}')
+        actions = actions[..., : self.action_dim]
+        return self.normalization.denormalize_action(actions)
+
+    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.n_action_steps),
+        }
+
+    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.action_chunk_start
+            end = start + self.n_action_steps
+            if actions.shape[1] < end:
+                raise RuntimeError(f'action chunk too short: got {actions.shape[1]}, need at least {end}')
+            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_native.yaml

@@ -0,0 +1,35 @@
+# @package agent
+_target_: roboimi.vla.agent_smolvla_native.SmolVLANativeAgent
+
+model: null
+tokenizer: null
+tokenizer_name: ${agent.model_config.vlm_model_name}
+
+action_dim: 16
+obs_dim: 16
+normalization_type: "min_max"
+chunk_size: 16
+pred_horizon: ${agent.chunk_size}
+obs_horizon: 2
+n_action_steps: 8
+num_action_steps: ${agent.n_action_steps}
+action_chunk_start: 0
+camera_names: ${data.camera_names}
+num_cams: ${len:${agent.camera_names}}
+task_description: null
+# SmolVLA performs its own SigLIP-style resize+pad inside the wrapper/core.
+# Keeping dataset/eval resize disabled avoids lossy double resizing.
+dataset_image_resize_shape: null
+eval_image_resize_shape: null
+
+model_config:
+  max_state_dim: 32
+  max_action_dim: 32
+  chunk_size: ${agent.chunk_size}
+  n_action_steps: ${agent.n_action_steps}
+  vlm_model_name: HuggingFaceTB/SmolVLM2-500M-Video-Instruct
+  load_vlm_weights: true
+  train_expert_only: true
+  freeze_vision_encoder: true
+  num_vlm_layers: 16
+  resize_imgs_with_padding: [512, 512]

roboimi/vla/models/smolvla/__init__.py

@@ -0,0 +1,13 @@
+"""Native SmolVLA model core."""
+
+from .configuration import NativeSmolVLAConfig, SmolVLAConfig
+from .modeling import VLAFlowMatching, make_att_2d_masks, pad_vector, resize_with_pad
+
+__all__ = [
+    "NativeSmolVLAConfig",
+    "SmolVLAConfig",
+    "VLAFlowMatching",
+    "make_att_2d_masks",
+    "pad_vector",
+    "resize_with_pad",
+]

roboimi/vla/models/smolvla/configuration.py

@@ -0,0 +1,75 @@
+"""Native SmolVLA configuration, independent of LeRobot."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass
+class NativeSmolVLAConfig:
+    """Lightweight configuration for the native SmolVLA model core.
+
+    This intentionally keeps only model-core fields needed by
+    :class:`VLAFlowMatching`; policy/dataset/optimizer concerns stay outside of
+    the native model package.
+    """
+
+    # Input / output structure.
+    n_obs_steps: int = 1
+    chunk_size: int = 50
+    n_action_steps: int = 50
+
+    # Shorter state and action vectors are padded before entering the model.
+    max_state_dim: int = 32
+    max_action_dim: int = 32
+
+    # Image preprocessing.
+    resize_imgs_with_padding: tuple[int, int] = (512, 512)
+
+    # Tokenizer / decoding.
+    tokenizer_max_length: int = 48
+    num_steps: int = 10
+
+    # Attention utils.
+    use_cache: bool = True
+
+    # Finetuning settings.
+    freeze_vision_encoder: bool = True
+    train_expert_only: bool = True
+    train_state_proj: bool = True
+
+    # VLM / expert construction settings.
+    vlm_model_name: str = "HuggingFaceTB/SmolVLM2-500M-Video-Instruct"
+    load_vlm_weights: bool = False
+    add_image_special_tokens: bool = False
+    attention_mode: str = "cross_attn"
+    prefix_length: int = -1
+    pad_language_to: str = "longest"
+    num_expert_layers: int = -1
+    num_vlm_layers: int = 16
+    self_attn_every_n_layers: int = 2
+    expert_width_multiplier: float = 0.75
+
+    # Flow-matching timestep embedding.
+    min_period: float = 4e-3
+    max_period: float = 4.0
+
+    # Runtime settings.
+    device: str | None = None
+    compile_model: bool = False
+    compile_mode: str = "max-autotune"
+
+    # RTC hook placeholder; the native core does not implement RTC, but keeping
+    # the field lets migrated call sites pass configs through unchanged.
+    rtc_config: object | None = None
+
+    def __post_init__(self) -> None:
+        if self.n_action_steps > self.chunk_size:
+            raise ValueError(
+                "The chunk size is the upper bound for the number of action steps per model invocation. "
+                f"Got {self.n_action_steps} for `n_action_steps` and {self.chunk_size} for `chunk_size`."
+            )
+
+
+# Compatibility alias for migrated code that still imports SmolVLAConfig.
+SmolVLAConfig = NativeSmolVLAConfig

roboimi/vla/models/smolvla/modeling.py

@@ -0,0 +1,421 @@
+"""Native SmolVLA flow-matching core.
+
+This module ports the lightweight helpers and ``VLAFlowMatching`` from the
+LeRobot SmolVLA implementation while avoiding any LeRobot imports.  The heavy
+Transformers-backed VLM/expert is optional and can be injected for tests.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import TypedDict
+
+try:
+    from typing import Unpack
+except ImportError:  # Python 3.10
+    from typing_extensions import Unpack
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor, nn
+
+from .configuration import NativeSmolVLAConfig
+
+
+class ActionSelectKwargs(TypedDict, total=False):
+    inference_delay: int | None
+    prev_chunk_left_over: Tensor | None
+    execution_horizon: int | None
+
+
+def create_sinusoidal_pos_embedding(
+    time: torch.Tensor,
+    dimension: int,
+    min_period: float,
+    max_period: float,
+    device: torch.device | str = "cpu",
+) -> Tensor:
+    """Compute sine/cosine positional embeddings for scalar positions."""
+    if dimension % 2 != 0:
+        raise ValueError(f"dimension ({dimension}) must be divisible by 2")
+    if time.ndim != 1:
+        raise ValueError("The time tensor is expected to be of shape `(batch_size, )`.")
+
+    fraction = torch.linspace(0.0, 1.0, dimension // 2, dtype=torch.float64, device=device)
+    period = min_period * (max_period / min_period) ** fraction
+    scaling_factor = 1.0 / period * 2 * math.pi
+    sin_input = scaling_factor[None, :] * time[:, None].to(dtype=torch.float64)
+    pos_emb = torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1)
+    return pos_emb
+
+
+def make_att_2d_masks(pad_masks: Tensor, att_masks: Tensor) -> Tensor:
+    """Build Big Vision-style 2D prefix-LM attention masks.
+
+    ``pad_masks`` is ``bool[B, N]`` and marks valid tokens. ``att_masks`` is
+    ``bool/int[B, N]`` where cumulative increments begin new causal groups.  A
+    query token can attend to valid key tokens whose cumulative attention group
+    is less than or equal to the query's group.
+    """
+    if att_masks.ndim != 2:
+        raise ValueError(att_masks.ndim)
+    if pad_masks.ndim != 2:
+        raise ValueError(pad_masks.ndim)
+
+    cumsum = torch.cumsum(att_masks.to(dtype=torch.long), dim=1)
+    att_2d_masks = cumsum[:, None, :] <= cumsum[:, :, None]
+    pad_2d_masks = pad_masks[:, None, :].bool() & pad_masks[:, :, None].bool()
+    return att_2d_masks & pad_2d_masks
+
+
+def resize_with_pad(img: Tensor, width: int, height: int, pad_value: float = -1) -> Tensor:
+    """Resize a BCHW image batch preserving aspect ratio, then top/left pad."""
+    if img.ndim != 4:
+        raise ValueError(f"(b,c,h,w) expected, but {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_img = 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_img, (pad_width, 0, pad_height, 0), value=pad_value)
+
+
+def pad_vector(vector: Tensor, new_dim: int) -> Tensor:
+    """Pad a vector-like tensor's last dimension with zeros to ``new_dim``."""
+    current_dim = vector.shape[-1]
+    if current_dim == new_dim:
+        return vector
+    if current_dim > new_dim:
+        raise ValueError(f"Cannot pad vector with current dimension {current_dim} to smaller target dimension {new_dim}.")
+
+    shape = list(vector.shape)
+    shape[-1] = new_dim
+    new_vector = torch.zeros(*shape, dtype=vector.dtype, device=vector.device)
+    new_vector[..., :current_dim] = vector
+    return new_vector
+
+
+def pad_tensor(tensor: Tensor, max_len: int, pad_value: int | float) -> Tensor:
+    """Pad a tensor along sequence dimension to ``max_len``."""
+    bsize, seq_len = tensor.shape[:2]
+    if seq_len >= max_len:
+        return tensor
+    padded_tensor = torch.full(
+        (bsize, max_len, *tensor.shape[2:]),
+        pad_value,
+        dtype=tensor.dtype,
+        device=tensor.device,
+    )
+    padded_tensor[:, :seq_len] = tensor
+    return padded_tensor
+
+
+class VLAFlowMatching(nn.Module):
+    """SmolVLA flow-matching action head around a VLM plus action expert."""
+
+    def __init__(
+        self,
+        config: NativeSmolVLAConfig,
+        rtc_processor: object | None = None,
+        vlm_with_expert: nn.Module | None = None,
+    ):
+        super().__init__()
+        self.config = config
+
+        if vlm_with_expert is None:
+            from .smolvlm_with_expert import SmolVLMWithExpertModel
+
+            vlm_with_expert = SmolVLMWithExpertModel(
+                model_id=self.config.vlm_model_name,
+                freeze_vision_encoder=self.config.freeze_vision_encoder,
+                train_expert_only=self.config.train_expert_only,
+                load_vlm_weights=self.config.load_vlm_weights,
+                attention_mode=self.config.attention_mode,
+                num_expert_layers=self.config.num_expert_layers,
+                num_vlm_layers=self.config.num_vlm_layers,
+                self_attn_every_n_layers=self.config.self_attn_every_n_layers,
+                expert_width_multiplier=self.config.expert_width_multiplier,
+                device=self.config.device if self.config.device is not None else "auto",
+            )
+        self.vlm_with_expert = vlm_with_expert
+
+        vlm_hidden_size = self.vlm_with_expert.config.text_config.hidden_size
+        expert_hidden_size = self.vlm_with_expert.expert_hidden_size
+        self.state_proj = nn.Linear(self.config.max_state_dim, vlm_hidden_size)
+        self.action_in_proj = nn.Linear(self.config.max_action_dim, expert_hidden_size)
+        self.action_out_proj = nn.Linear(expert_hidden_size, self.config.max_action_dim)
+        self.action_time_mlp_in = nn.Linear(expert_hidden_size * 2, expert_hidden_size)
+        self.action_time_mlp_out = nn.Linear(expert_hidden_size, expert_hidden_size)
+
+        self.set_requires_grad()
+        tokenizer = self.vlm_with_expert.processor.tokenizer
+        self.fake_image_token = tokenizer.fake_image_token_id
+        self.global_image_token = tokenizer.global_image_token_id
+        self.global_image_start_token = torch.tensor(
+            [self.fake_image_token, self.global_image_token], dtype=torch.long
+        )
+        self.add_image_special_tokens = self.config.add_image_special_tokens
+        self.image_end_token = torch.tensor([self.fake_image_token], dtype=torch.long)
+        self.prefix_length = self.config.prefix_length
+        self.rtc_processor = rtc_processor
+
+        if config.compile_model:
+            torch.set_float32_matmul_precision("high")
+            self.sample_actions = torch.compile(self.sample_actions, mode=config.compile_mode)
+            self.forward = torch.compile(self.forward, mode=config.compile_mode)
+
+    def _rtc_enabled(self) -> bool:
+        return bool(self.config.rtc_config is not None and getattr(self.config.rtc_config, "enabled", False))
+
+    def set_requires_grad(self) -> None:
+        for params in self.state_proj.parameters():
+            params.requires_grad = self.config.train_state_proj
+
+    def sample_noise(self, shape: tuple[int, ...] | torch.Size, device: torch.device | str) -> Tensor:
+        return torch.normal(mean=0.0, std=1.0, size=shape, dtype=torch.float32, device=device)
+
+    def sample_time(self, bsize: int, device: torch.device | str) -> Tensor:
+        beta_dist = torch.distributions.Beta(concentration1=1.5, concentration0=1.0)
+        time_beta = beta_dist.sample((bsize,)).to(device=device, dtype=torch.float32)
+        return time_beta * 0.999 + 0.001
+
+    def _vlm_device(self) -> torch.device:
+        vlm = getattr(self.vlm_with_expert, "vlm", None)
+        return getattr(vlm, "device", next(self.parameters()).device)
+
+    def embed_prefix(
+        self,
+        images: list[Tensor],
+        img_masks: list[Tensor],
+        lang_tokens: Tensor,
+        lang_masks: Tensor,
+        state: Tensor,
+    ) -> tuple[Tensor, Tensor, Tensor]:
+        """Embed images, language tokens, and robot state as prefix tokens."""
+        embs: list[Tensor] = []
+        pad_masks: list[Tensor] = []
+        att_masks: list[int] = []
+
+        for img, img_mask in zip(images, img_masks, strict=False):
+            if self.add_image_special_tokens:
+                image_start_token = (
+                    self.vlm_with_expert.embed_language_tokens(self.global_image_start_token.to(device=self._vlm_device()))
+                    .unsqueeze(0)
+                    .expand(img.shape[0], -1, -1)
+                )
+                image_start_mask = torch.ones_like(image_start_token[:, :, 0], dtype=torch.bool)
+                embs.append(image_start_token)
+                pad_masks.append(image_start_mask)
+                att_masks += [0] * image_start_mask.shape[-1]
+
+            img_emb = self.vlm_with_expert.embed_image(img)
+            img_emb = img_emb * torch.tensor(img_emb.shape[-1] ** 0.5, dtype=img_emb.dtype, device=img_emb.device)
+            bsize, num_img_embs = img_emb.shape[:2]
+            img_mask = img_mask.to(device=img_emb.device, dtype=torch.bool)[:, None].expand(bsize, num_img_embs)
+            embs.append(img_emb)
+            pad_masks.append(img_mask)
+            att_masks += [0] * num_img_embs
+
+            if self.add_image_special_tokens:
+                image_end_token = (
+                    self.vlm_with_expert.embed_language_tokens(self.image_end_token.to(device=self._vlm_device()))
+                    .unsqueeze(0)
+                    .expand(img.shape[0], -1, -1)
+                )
+                image_end_mask = torch.ones_like(image_end_token[:, :, 0], dtype=torch.bool)
+                embs.append(image_end_token)
+                pad_masks.append(image_end_mask)
+                att_masks += [0] * image_end_mask.shape[1]
+
+        lang_emb = self.vlm_with_expert.embed_language_tokens(lang_tokens)
+        lang_emb = lang_emb * math.sqrt(lang_emb.shape[-1])
+        embs.append(lang_emb)
+        pad_masks.append(lang_masks.to(device=lang_emb.device, dtype=torch.bool))
+        att_masks += [0] * lang_emb.shape[1]
+
+        state_emb = self.state_proj(state)
+        state_emb = state_emb[:, None, :] if state_emb.ndim == 2 else state_emb
+        embs.append(state_emb)
+        state_mask = torch.ones(state_emb.shape[:2], dtype=torch.bool, device=state_emb.device)
+        pad_masks.append(state_mask)
+        att_masks += [1] * state_emb.shape[1]
+
+        all_embs = torch.cat(embs, dim=1)
+        all_pad_masks = torch.cat(pad_masks, dim=1)
+        all_att_masks = torch.tensor(att_masks, dtype=torch.bool, device=all_pad_masks.device)[None, :]
+
+        if self.prefix_length > 0 and all_pad_masks.shape[1] < self.prefix_length:
+            all_embs = pad_tensor(all_embs, self.prefix_length, pad_value=0)
+            all_pad_masks = pad_tensor(all_pad_masks, self.prefix_length, pad_value=0)
+            all_att_masks = pad_tensor(all_att_masks, self.prefix_length, pad_value=0)
+
+        all_att_masks = all_att_masks.expand(all_pad_masks.shape[0], -1)
+        return all_embs, all_pad_masks, all_att_masks
+
+    def embed_suffix(self, noisy_actions: Tensor, timestep: Tensor) -> tuple[Tensor, Tensor, Tensor]:
+        """Embed noisy action tokens and timestep for the expert suffix."""
+        action_emb = self.action_in_proj(noisy_actions)
+        device = action_emb.device
+        bsize = action_emb.shape[0]
+        dtype = action_emb.dtype
+
+        time_emb = create_sinusoidal_pos_embedding(
+            timestep,
+            self.vlm_with_expert.expert_hidden_size,
+            self.config.min_period,
+            self.config.max_period,
+            device=device,
+        ).to(dtype=dtype)
+        time_emb = time_emb[:, None, :].expand_as(action_emb)
+        action_time_emb = torch.cat([action_emb, time_emb], dim=2)
+        action_time_emb = self.action_time_mlp_in(action_time_emb)
+        action_time_emb = F.silu(action_time_emb)
+        action_time_emb = self.action_time_mlp_out(action_time_emb)
+
+        action_time_mask = torch.ones(action_time_emb.shape[:2], dtype=torch.bool, device=device)
+        att_masks = torch.ones(bsize, self.config.chunk_size, dtype=torch.bool, device=device)
+        return action_time_emb, action_time_mask, att_masks
+
+    def forward(
+        self,
+        images: list[Tensor],
+        img_masks: list[Tensor],
+        lang_tokens: Tensor,
+        lang_masks: Tensor,
+        state: Tensor,
+        actions: Tensor,
+        noise: Tensor | None = None,
+        time: Tensor | None = None,
+    ) -> Tensor:
+        """Run a training forward pass and return per-element flow loss."""
+        if noise is None:
+            noise = self.sample_noise(actions.shape, actions.device)
+        if time is None:
+            time = self.sample_time(actions.shape[0], actions.device)
+
+        time_expanded = time[:, None, None]
+        x_t = time_expanded * noise + (1 - time_expanded) * actions
+        u_t = noise - actions
+
+        prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix(
+            images, img_masks, lang_tokens, lang_masks, state=state
+        )
+        suffix_embs, suffix_pad_masks, suffix_att_masks = self.embed_suffix(x_t, time)
+
+        pad_masks = torch.cat([prefix_pad_masks, suffix_pad_masks], dim=1)
+        att_masks = torch.cat([prefix_att_masks, suffix_att_masks], dim=1)
+        att_2d_masks = make_att_2d_masks(pad_masks, att_masks)
+        position_ids = torch.cumsum(pad_masks, dim=1) - 1
+
+        (_, suffix_out), _ = self.vlm_with_expert.forward(
+            attention_mask=att_2d_masks,
+            position_ids=position_ids,
+            past_key_values=None,
+            inputs_embeds=[prefix_embs, suffix_embs],
+            use_cache=False,
+            fill_kv_cache=False,
+        )
+        suffix_out = suffix_out[:, -self.config.chunk_size :].to(dtype=torch.float32)
+        v_t = self.action_out_proj(suffix_out)
+        return F.mse_loss(u_t, v_t, reduction="none")
+
+    def sample_actions(
+        self,
+        images: list[Tensor],
+        img_masks: list[Tensor],
+        lang_tokens: Tensor,
+        lang_masks: Tensor,
+        state: Tensor,
+        noise: Tensor | None = None,
+        **kwargs: Unpack[ActionSelectKwargs],
+    ) -> Tensor:
+        """Sample an action chunk with Euler integration over the flow field."""
+        bsize = state.shape[0]
+        device = state.device
+        if noise is None:
+            actions_shape = (bsize, self.config.chunk_size, self.config.max_action_dim)
+            noise = self.sample_noise(actions_shape, device)
+
+        prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix(
+            images, img_masks, lang_tokens, lang_masks, state=state
+        )
+        prefix_att_2d_masks = make_att_2d_masks(prefix_pad_masks, prefix_att_masks)
+        prefix_position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1
+        _, past_key_values = self.vlm_with_expert.forward(
+            attention_mask=prefix_att_2d_masks,
+            position_ids=prefix_position_ids,
+            past_key_values=None,
+            inputs_embeds=[prefix_embs, None],
+            use_cache=self.config.use_cache,
+            fill_kv_cache=True,
+        )
+
+        dt = -1.0 / self.config.num_steps
+        x_t = noise
+        for step in range(self.config.num_steps):
+            time = 1.0 + step * dt
+            time_tensor = torch.tensor(time, dtype=torch.float32, device=device).expand(bsize)
+
+            def denoise_step_partial_call(input_x_t: Tensor, current_timestep: Tensor = time_tensor) -> Tensor:
+                return self.denoise_step(
+                    x_t=input_x_t,
+                    prefix_pad_masks=prefix_pad_masks,
+                    past_key_values=past_key_values,
+                    timestep=current_timestep,
+                )
+
+            if self._rtc_enabled() and self.rtc_processor is not None:
+                v_t = self.rtc_processor.denoise_step(
+                    x_t=x_t,
+                    prev_chunk_left_over=kwargs.get("prev_chunk_left_over"),
+                    inference_delay=kwargs.get("inference_delay"),
+                    time=time,
+                    original_denoise_step_partial=denoise_step_partial_call,
+                    execution_horizon=kwargs.get("execution_horizon"),
+                )
+            else:
+                v_t = denoise_step_partial_call(x_t)
+            x_t = x_t + dt * v_t
+
+            if self.rtc_processor is not None and getattr(self.rtc_processor, "is_debug_enabled", lambda: False)():
+                self.rtc_processor.track(time=time, x_t=x_t, v_t=v_t)
+
+        return x_t
+
+    def denoise_step(
+        self,
+        prefix_pad_masks: Tensor,
+        past_key_values: object,
+        x_t: Tensor,
+        timestep: Tensor,
+    ) -> Tensor:
+        """Apply one denoising step at a given timestep."""
+        suffix_embs, suffix_pad_masks, suffix_att_masks = self.embed_suffix(x_t, timestep)
+        suffix_len = suffix_pad_masks.shape[1]
+        batch_size = prefix_pad_masks.shape[0]
+        prefix_len = prefix_pad_masks.shape[1]
+        prefix_pad_2d_masks = prefix_pad_masks[:, None, :].expand(batch_size, suffix_len, prefix_len)
+        suffix_att_2d_masks = make_att_2d_masks(suffix_pad_masks, suffix_att_masks)
+        full_att_2d_masks = torch.cat([prefix_pad_2d_masks, suffix_att_2d_masks], dim=2)
+        prefix_offsets = torch.sum(prefix_pad_masks, dim=-1)[:, None]
+        position_ids = prefix_offsets + torch.cumsum(suffix_pad_masks, dim=1) - 1
+
+        outputs_embeds, _ = self.vlm_with_expert.forward(
+            attention_mask=full_att_2d_masks,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=[None, suffix_embs],
+            use_cache=self.config.use_cache,
+            fill_kv_cache=False,
+        )
+        suffix_out = outputs_embeds[1][:, -self.config.chunk_size :].to(dtype=torch.float32)
+        return self.action_out_proj(suffix_out)

roboimi/vla/models/smolvla/smolvlm_with_expert.py

@@ -0,0 +1,571 @@
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+from typing import TYPE_CHECKING
+
+import torch
+from torch import nn
+
+if TYPE_CHECKING:
+    from transformers import (
+        AutoConfig,
+        AutoModel,
+        AutoModelForImageTextToText,
+        AutoProcessor,
+        SmolVLMForConditionalGeneration,
+    )
+
+
+def _require_transformers():
+    try:
+        from transformers import (
+            AutoConfig,
+            AutoModel,
+            AutoModelForImageTextToText,
+            AutoProcessor,
+            SmolVLMForConditionalGeneration,
+        )
+    except ImportError as exc:
+        raise ImportError(
+            "Native SmolVLA requires the optional `transformers` package to construct "
+            "SmolVLMWithExpertModel. Install transformers or inject `vlm_with_expert` "
+            "when constructing VLAFlowMatching."
+        ) from exc
+    return AutoConfig, AutoModel, AutoModelForImageTextToText, AutoProcessor, SmolVLMForConditionalGeneration
+
+
+def apply_rope(x, positions, max_wavelength=10_000):
+    """
+    Applies RoPE positions [B, L] to x [B, L, H, D].
+    """
+    d_half = x.shape[-1] // 2
+    device = x.device
+    dtype = x.dtype
+    x = x.to(torch.float32)
+
+    freq_exponents = (2.0 / x.shape[-1]) * torch.arange(d_half, dtype=torch.float32, device=device)
+    timescale = max_wavelength**freq_exponents
+    radians = positions[..., None].to(torch.float32) / timescale[None, None, :].to(torch.float32)
+
+    radians = radians[..., None, :]
+
+    sin = torch.sin(radians)  # .to(dtype=dtype)
+    cos = torch.cos(radians)  # .to(dtype=dtype)
+
+    x1, x2 = x.split(d_half, dim=-1)
+    res = torch.empty_like(x)
+    res[..., :d_half] = x1 * cos - x2 * sin
+    res[..., d_half:] = x2 * cos + x1 * sin
+
+    return res.to(dtype)
+
+
+def get_intermediate_size(hidden_dim, ffn_dim_multiplier=4, multiple_of=256):
+    hidden_dim = int(2 * hidden_dim / 3)
+    hidden_dim = int(ffn_dim_multiplier * hidden_dim)
+    hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
+    return hidden_dim
+
+
+class SmolVLMWithExpertModel(nn.Module):
+    def __init__(
+        self,
+        model_id: str = "HuggingFaceTB/SmolVLM2-500M-Video-Instruct",
+        load_vlm_weights: bool = True,
+        train_expert_only: bool = True,
+        freeze_vision_encoder: bool = False,
+        attention_mode: str = "self_attn",
+        num_expert_layers: int = -1,
+        num_vlm_layers: int = -1,
+        self_attn_every_n_layers: int = -1,
+        expert_width_multiplier: float = 0.5,
+        device: str = "auto",
+    ):
+        super().__init__()
+        AutoConfig, AutoModel, AutoModelForImageTextToText, AutoProcessor, SmolVLMForConditionalGeneration = _require_transformers()
+        if load_vlm_weights:
+            print(f"Loading  {model_id} weights ...")
+            self.vlm = AutoModelForImageTextToText.from_pretrained(
+                model_id,
+                torch_dtype="bfloat16",
+                low_cpu_mem_usage=True,
+            )
+            config = self.vlm.config
+        else:
+            config = AutoConfig.from_pretrained(model_id)
+            self.vlm = SmolVLMForConditionalGeneration(config=config)
+        self.processor = AutoProcessor.from_pretrained(model_id)
+        if num_vlm_layers > 0:
+            print(f"Reducing the number of VLM layers to {num_vlm_layers} ...")
+            self.get_vlm_model().text_model.layers = self.get_vlm_model().text_model.layers[:num_vlm_layers]
+        self.num_vlm_layers = len(self.get_vlm_model().text_model.layers)
+        self.config = config
+        # Smaller lm expert
+        lm_expert_config = copy.deepcopy(config.text_config)
+        hidden_size = lm_expert_config.hidden_size
+        lm_expert_config.hidden_size = int(hidden_size * expert_width_multiplier)  # hidden_size // 2
+        lm_expert_config.intermediate_size = get_intermediate_size(int(hidden_size * expert_width_multiplier))
+        lm_expert_config.num_hidden_layers = self.num_vlm_layers
+        if num_expert_layers > 0:
+            assert len(self.get_vlm_model().text_model.layers) % num_expert_layers == 0, (
+                f"Number of layers in the VLM {len(self.get_vlm_model().text_model.layers)} are not multiple of num_expert_layers {num_expert_layers}"
+            )
+            lm_expert_config.num_hidden_layers = num_expert_layers
+        self.lm_expert = AutoModel.from_config(lm_expert_config)
+
+        self.num_expert_layers = len(self.lm_expert.layers)
+        self.self_attn_every_n_layers = self_attn_every_n_layers
+        if "cross" in attention_mode:
+            # Reshape qkv projections to have the same input dimension as the vlm
+            for layer_idx in range(len(self.lm_expert.layers)):
+                if self.self_attn_every_n_layers > 0 and layer_idx % self.self_attn_every_n_layers == 0:
+                    continue
+                self.lm_expert.layers[layer_idx].self_attn.k_proj = nn.Linear(
+                    config.text_config.num_key_value_heads * config.text_config.head_dim,
+                    lm_expert_config.num_key_value_heads * lm_expert_config.head_dim,
+                    bias=lm_expert_config.attention_bias,
+                )
+                self.lm_expert.layers[layer_idx].self_attn.v_proj = nn.Linear(
+                    config.text_config.num_key_value_heads * config.text_config.head_dim,
+                    lm_expert_config.num_key_value_heads * lm_expert_config.head_dim,
+                    bias=lm_expert_config.attention_bias,
+                )
+        # Remove unused embed_tokens
+        self.lm_expert.embed_tokens = None
+
+        self.num_attention_heads = self.config.text_config.num_attention_heads
+        self.num_key_value_heads = self.config.text_config.num_key_value_heads
+
+        self.freeze_vision_encoder = freeze_vision_encoder
+        self.train_expert_only = train_expert_only
+        self.attention_mode = attention_mode
+        self.expert_hidden_size = lm_expert_config.hidden_size
+        self.set_requires_grad()
+
+    def get_vlm_model(self):
+        return self.vlm.model
+
+    def set_requires_grad(self):
+        if self.freeze_vision_encoder:
+            self.get_vlm_model().vision_model.eval()
+            for params in self.get_vlm_model().vision_model.parameters():
+                params.requires_grad = False
+        if self.train_expert_only:
+            self.vlm.eval()
+            for params in self.vlm.parameters():
+                params.requires_grad = False
+        else:
+            # To avoid unused params issue with distributed training
+            last_layers = [self.num_vlm_layers - 1]
+            if (
+                self.num_vlm_layers != self.num_expert_layers
+                and self.num_vlm_layers % self.num_expert_layers == 0
+            ):
+                last_layers.append(self.num_vlm_layers - 2)
+            frozen_layers = [
+                "lm_head",
+                "text_model.model.norm.weight",
+            ]
+            for layer in last_layers:
+                frozen_layers.append(f"text_model.model.layers.{layer}.")
+
+            for name, params in self.vlm.named_parameters():
+                if any(k in name for k in frozen_layers):
+                    params.requires_grad = False
+        # To avoid unused params issue with distributed training
+        for name, params in self.lm_expert.named_parameters():
+            if "lm_head" in name:
+                params.requires_grad = False
+
+    def train(self, mode: bool = True):
+        super().train(mode)
+
+        if self.freeze_vision_encoder:
+            self.get_vlm_model().vision_model.eval()
+
+        if self.train_expert_only:
+            self.vlm.eval()
+
+    def embed_image(self, image: torch.Tensor):
+        patch_attention_mask = None
+        # Get sequence from the vision encoder
+        image_hidden_states = (
+            self.get_vlm_model()
+            .vision_model(
+                pixel_values=image.to(dtype=self.get_vlm_model().vision_model.dtype),
+                patch_attention_mask=patch_attention_mask,
+            )
+            .last_hidden_state
+        )
+        # Modality projection & resampling
+        image_hidden_states = self.get_vlm_model().connector(image_hidden_states)
+        return image_hidden_states
+
+    def embed_language_tokens(self, tokens: torch.Tensor):
+        return self.get_vlm_model().text_model.get_input_embeddings()(tokens)
+
+    def forward_attn_layer(
+        self,
+        model_layers,
+        inputs_embeds,
+        layer_idx,
+        position_ids,
+        attention_mask,
+        batch_size,
+        head_dim,
+        use_cache: bool = True,
+        fill_kv_cache: bool = True,
+        past_key_values=None,
+    ) -> list[torch.Tensor]:
+        query_states = []
+        key_states = []
+        value_states = []
+        for i, hidden_states in enumerate(inputs_embeds):
+            layer = model_layers[i][layer_idx]
+            if hidden_states is None or layer is None:
+                continue
+            hidden_states = layer.input_layernorm(hidden_states)
+
+            input_shape = hidden_states.shape[:-1]
+            hidden_shape = (*input_shape, -1, layer.self_attn.head_dim)
+
+            hidden_states = hidden_states.to(dtype=layer.self_attn.q_proj.weight.dtype)
+            query_state = layer.self_attn.q_proj(hidden_states).view(hidden_shape)
+            key_state = layer.self_attn.k_proj(hidden_states).view(hidden_shape)
+            value_state = layer.self_attn.v_proj(hidden_states).view(hidden_shape)
+
+            query_states.append(query_state)
+            key_states.append(key_state)
+            value_states.append(value_state)
+
+        # B,L,H,D with L sequence length, H number of heads, D head dim
+        # concatenate on the number of embeddings/tokens
+        query_states = torch.cat(query_states, dim=1)
+        key_states = torch.cat(key_states, dim=1)
+        value_states = torch.cat(value_states, dim=1)
+        seq_len = query_states.shape[1]
+        if seq_len < position_ids.shape[1]:
+            _position_ids = position_ids[:, :seq_len]
+            _attention_mask = attention_mask[:, :seq_len, :seq_len]
+        else:
+            _position_ids = position_ids
+            _attention_mask = attention_mask
+
+        attention_mask_ = _attention_mask
+        position_ids_ = _position_ids
+
+        query_states = apply_rope(query_states, position_ids_)
+        key_states = apply_rope(key_states, position_ids_)
+
+        if use_cache and past_key_values is None:
+            past_key_values = {}
+
+        if use_cache:
+            if fill_kv_cache:
+                past_key_values[layer_idx] = {
+                    "key_states": key_states,
+                    "value_states": value_states,
+                }
+            else:
+                # TODO here, some optimization can be done - similar to a `StaticCache` we can declare the `max_len` before.
+                # so we create an empty cache, with just one cuda malloc, and if (in autoregressive case) we reach
+                # the max len, then we (for instance) double the cache size. This implementation already exists
+                # in `transformers`. (molbap)
+                key_states = torch.cat([past_key_values[layer_idx]["key_states"], key_states], dim=1)
+                value_states = torch.cat([past_key_values[layer_idx]["value_states"], value_states], dim=1)
+
+        attention_interface = self.get_attention_interface()
+
+        att_output = attention_interface(
+            attention_mask_, batch_size, head_dim, query_states, key_states, value_states
+        )
+        return [att_output], past_key_values
+
+    def forward_cross_attn_layer(
+        self,
+        model_layers,
+        inputs_embeds,
+        layer_idx,
+        position_ids,
+        attention_mask,
+        batch_size,
+        head_dim,
+        use_cache: bool = True,
+        fill_kv_cache: bool = True,
+        past_key_values=None,
+    ) -> list[torch.Tensor]:
+        attention_interface = self.get_attention_interface()
+
+        att_outputs = []
+        assert len(inputs_embeds) == 2 or (use_cache and past_key_values is not None and not fill_kv_cache), (
+            f"Both len(inputs_embeds) == {len(inputs_embeds)} and past_key_values is {past_key_values}"
+        )
+
+        if len(inputs_embeds) == 2 and not past_key_values:
+            # Prefix attention
+            seq_len = inputs_embeds[0].shape[1]
+            position_id, expert_position_id = position_ids[:, :seq_len], position_ids[:, seq_len:]
+            prefix_attention_mask = attention_mask[:, :seq_len, :seq_len]
+
+            layer = model_layers[0][layer_idx]
+
+            hidden_states = layer.input_layernorm(inputs_embeds[0])
+
+            input_shape = hidden_states.shape[:-1]
+            hidden_shape = (*input_shape, -1, layer.self_attn.head_dim)
+
+            hidden_states = hidden_states.to(dtype=layer.self_attn.q_proj.weight.dtype)
+            query_state = layer.self_attn.q_proj(hidden_states).view(hidden_shape)
+            key_state = layer.self_attn.k_proj(hidden_states).view(hidden_shape)
+            value_states = layer.self_attn.v_proj(hidden_states).view(hidden_shape)
+
+            # B,L,H,D with L sequence length, H number of heads, D head dim
+            query_states = apply_rope(query_state, position_id)
+            key_states = apply_rope(key_state, position_id)
+
+            att_output = attention_interface(
+                prefix_attention_mask, batch_size, head_dim, query_states, key_states, value_states
+            )
+            att_outputs.append(att_output)
+        else:
+            expert_position_id = position_ids
+
+        if use_cache and past_key_values is None:
+            past_key_values = {}
+
+        if use_cache:
+            if fill_kv_cache:
+                past_key_values[layer_idx] = {
+                    "key_states": key_states,
+                    "value_states": value_states,
+                }
+            else:
+                # TODO here, some optimization can be done - similar to a `StaticCache` we can declare the `max_len` before.
+                # so we create an empty cache, with just one cuda malloc, and if (in autoregressive case) we reach
+                # the max len, then we (for instance) double the cache size. This implementation already exists
+                # in `transformers`. (molbap)
+                key_states = past_key_values[layer_idx]["key_states"]
+                value_states = past_key_values[layer_idx]["value_states"]
+
+        # Expert
+        expert_layer = model_layers[1][layer_idx]
+        if expert_layer is not None:
+            expert_hidden_states = expert_layer.input_layernorm(inputs_embeds[1])
+
+            expert_input_shape = expert_hidden_states.shape[:-1]
+            expert_hidden_shape = (*expert_input_shape, -1, expert_layer.self_attn.head_dim)
+
+            expert_hidden_states = expert_hidden_states.to(dtype=expert_layer.self_attn.q_proj.weight.dtype)
+            expert_query_state = expert_layer.self_attn.q_proj(expert_hidden_states).view(expert_hidden_shape)
+
+            _key_states = key_states.to(dtype=expert_layer.self_attn.k_proj.weight.dtype).view(
+                *key_states.shape[:2], -1
+            )
+            expert_key_states = expert_layer.self_attn.k_proj(_key_states).view(
+                *_key_states.shape[:-1], -1, expert_layer.self_attn.head_dim
+            )  # k_proj should have same dim as kv
+
+            _value_states = value_states.to(dtype=expert_layer.self_attn.v_proj.weight.dtype).view(
+                *value_states.shape[:2], -1
+            )
+            expert_value_states = expert_layer.self_attn.v_proj(_value_states).view(
+                *_value_states.shape[:-1], -1, expert_layer.self_attn.head_dim
+            )
+
+            expert_position_id = (
+                expert_position_id - torch.min(expert_position_id, dim=1, keepdim=True).values
+            )  # start from 0
+            expert_attention_mask = attention_mask[
+                :, -inputs_embeds[1].shape[1] :, : expert_key_states.shape[1] :
+            ]  # take into account kv
+
+            expert_query_states = apply_rope(expert_query_state, expert_position_id)
+
+            att_output = attention_interface(
+                expert_attention_mask,
+                batch_size,
+                head_dim,
+                expert_query_states,
+                expert_key_states,
+                expert_value_states,
+            )
+            att_outputs.append(att_output)
+        else:
+            att_outputs.append(None)
+
+        # att_output = att_output.to(dtype=models[i].dtype)
+        return att_outputs, past_key_values
+
+    def get_model_layers(self, models: list) -> list:
+        vlm_layers = []
+        expert_layers = []
+        multiple_of = self.num_vlm_layers // self.num_expert_layers
+        for i in range(self.num_vlm_layers):
+            if multiple_of > 0 and i > 0 and i % multiple_of != 0:
+                expert_layer = None
+            else:
+                expert_layer_index = i // multiple_of if multiple_of > 0 else i
+                expert_layer = models[1].layers[expert_layer_index]
+            vlm_layers.append(models[0].layers[i])
+            expert_layers.append(expert_layer)
+        return [vlm_layers, expert_layers]
+
+    def forward(
+        self,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: list[torch.FloatTensor] | None = None,
+        inputs_embeds: list[torch.FloatTensor] = None,
+        use_cache: bool | None = None,
+        fill_kv_cache: bool | None = None,
+    ):
+        models = [self.get_vlm_model().text_model, self.lm_expert]
+        model_layers = self.get_model_layers(models)
+        for hidden_states in inputs_embeds:
+            # TODO this is very inefficient
+            # dtype is always the same, batch size too (if > 1 len)
+            # device could be trickier in multi gpu edge cases but that's it
+            if hidden_states is None:
+                continue
+            batch_size = hidden_states.shape[0]
+
+        # RMSNorm
+        num_layers = self.num_vlm_layers
+        head_dim = self.vlm.config.text_config.head_dim
+        for layer_idx in range(num_layers):
+            if (
+                fill_kv_cache
+                or "cross" not in self.attention_mode
+                or (self.self_attn_every_n_layers > 0 and layer_idx % self.self_attn_every_n_layers == 0)
+            ):
+                att_outputs, past_key_values = self.forward_attn_layer(
+                    model_layers,
+                    inputs_embeds,
+                    layer_idx,
+                    position_ids,
+                    attention_mask,
+                    batch_size,
+                    head_dim,
+                    use_cache=use_cache,
+                    fill_kv_cache=fill_kv_cache,
+                    past_key_values=past_key_values,
+                )
+            else:
+                att_outputs, past_key_values = self.forward_cross_attn_layer(
+                    model_layers,
+                    inputs_embeds,
+                    layer_idx,
+                    position_ids,
+                    attention_mask,
+                    batch_size,
+                    head_dim,
+                    use_cache=use_cache,
+                    fill_kv_cache=fill_kv_cache,
+                    past_key_values=past_key_values,
+                )
+            outputs_embeds = []
+            start = 0
+            for i, hidden_states in enumerate(inputs_embeds):
+                layer = model_layers[i][layer_idx]
+                att_output = (
+                    att_outputs[i] if i < len(att_outputs) else att_outputs[0]
+                )  # in case of self_attn
+                if hidden_states is not None:
+                    if layer is None:
+                        outputs_embeds.append(hidden_states)
+                        continue
+                    end = start + hidden_states.shape[1]
+
+                    if att_output.dtype != layer.self_attn.o_proj.weight.dtype:
+                        att_output = att_output.to(layer.self_attn.o_proj.weight.dtype)
+                    att_out = att_output[:, start:end]
+                    out_emb = layer.self_attn.o_proj(att_out)
+
+                    out_emb += hidden_states
+                    after_first_residual = out_emb.clone()
+
+                    out_emb = layer.post_attention_layernorm(out_emb)
+                    out_emb = layer.mlp(out_emb)
+
+                    out_emb += after_first_residual
+
+                    outputs_embeds.append(out_emb)
+
+                    start = end if len(att_outputs) == 1 else 0
+                else:
+                    outputs_embeds.append(None)
+
+            inputs_embeds = outputs_embeds
+
+        # final norm
+        outputs_embeds = []
+        for i, hidden_states in enumerate(inputs_embeds):
+            if hidden_states is not None:
+                out_emb = models[i].norm(hidden_states)
+                outputs_embeds.append(out_emb)
+            else:
+                outputs_embeds.append(None)
+        return outputs_embeds, past_key_values
+
+    def get_attention_interface(self):
+        attention_interface = self.eager_attention_forward
+        return attention_interface
+
+    def eager_attention_forward(
+        self, attention_mask, batch_size, head_dim, query_states, key_states, value_states
+    ):
+        num_att_heads = self.num_attention_heads
+        num_key_value_heads = self.num_key_value_heads
+        num_key_value_groups = num_att_heads // num_key_value_heads
+
+        sequence_length = key_states.shape[1]
+
+        key_states = key_states[:, :, :, None, :].expand(
+            batch_size, sequence_length, num_key_value_heads, num_key_value_groups, head_dim
+        )
+        key_states = key_states.reshape(
+            batch_size, sequence_length, num_key_value_heads * num_key_value_groups, head_dim
+        )
+
+        value_states = value_states[:, :, :, None, :].expand(
+            batch_size, sequence_length, num_key_value_heads, num_key_value_groups, head_dim
+        )
+        value_states = value_states.reshape(
+            batch_size, sequence_length, num_key_value_heads * num_key_value_groups, head_dim
+        )
+
+        # Attention here is upcasted to float32 to match the original eager implementation.
+        query_states = query_states.to(dtype=torch.float32)
+        key_states = key_states.to(dtype=torch.float32)
+
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.transpose(1, 2)
+
+        att_weights = torch.matmul(query_states, key_states.transpose(2, 3))
+        att_weights *= head_dim**-0.5
+
+        att_weights = att_weights.to(dtype=torch.float32)
+        big_neg = torch.finfo(att_weights.dtype).min  # -2.3819763e38  # See gemma/modules.py
+        masked_att_weights = torch.where(attention_mask[:, None, :, :], att_weights, big_neg)
+        probs = nn.functional.softmax(masked_att_weights, dim=-1)
+        probs = probs.to(dtype=value_states.dtype)
+
+        att_output = torch.matmul(probs, value_states.permute(0, 2, 1, 3))
+
+        att_output = att_output.permute(0, 2, 1, 3)
+        # we use -1 because sequence length can change
+        att_output = att_output.reshape(batch_size, -1, num_key_value_heads * num_key_value_groups * head_dim)
+
+        return att_output

tests/test_eval_vla_execution.py

@@ -129,6 +129,23 @@ class EvalVLAExecutionTest(unittest.TestCase):
             ["r_vis", "top", "front"],
         )
 
+    def test_resolve_eval_image_resize_shape_prefers_agent_top_level_override(self):
+        cfg = OmegaConf.create(
+            {
+                "agent": {
+                    "eval_image_resize_shape": None,
+                    "condition_encoder": {
+                        "eval_image_resize_shape": [256, 256],
+                    },
+                },
+                "data": {
+                    "image_resize_shape": [224, 224],
+                },
+            }
+        )
+
+        self.assertIsNone(eval_vla._resolve_eval_image_resize_shape(cfg))
+
     def test_resolve_eval_image_resize_shape_prefers_condition_encoder_override(self):
         cfg = OmegaConf.create(
             {
@@ -270,6 +287,25 @@ class EvalVLAExecutionTest(unittest.TestCase):
         np.testing.assert_array_equal(queues["action"].popleft().numpy(), np.array([20.0], dtype=np.float32))
         np.testing.assert_array_equal(queues["action"].popleft().numpy(), np.array([30.0], dtype=np.float32))
 
+    def test_enqueue_predicted_actions_honors_explicit_chunk_start(self):
+        queues = eval_vla._new_local_policy_queues(obs_horizon=2)
+        predicted_actions = torch.tensor(
+            [[[10.0], [20.0], [30.0], [40.0]]],
+            dtype=torch.float32,
+        )
+
+        eval_vla._enqueue_predicted_actions(
+            queues,
+            predicted_actions=predicted_actions,
+            obs_horizon=2,
+            num_action_steps=2,
+            action_chunk_start=0,
+        )
+
+        self.assertEqual(len(queues["action"]), 2)
+        np.testing.assert_array_equal(queues["action"].popleft().numpy(), np.array([10.0], dtype=np.float32))
+        np.testing.assert_array_equal(queues["action"].popleft().numpy(), np.array([20.0], dtype=np.float32))
+
     def test_remote_policy_runner_only_requests_server_inference_when_local_action_queue_is_empty(self):
         request_queue = _FakeQueue()
         response_queue = _FakeQueue(
@@ -288,6 +324,7 @@ class EvalVLAExecutionTest(unittest.TestCase):
             camera_names=["front"],
             obs_horizon=2,
             num_action_steps=2,
+            action_chunk_start=0,
         )
         first_observation = {
             "qpos": torch.tensor([1.0, 2.0], dtype=torch.float32),
@@ -315,8 +352,46 @@ class EvalVLAExecutionTest(unittest.TestCase):
         self.assertEqual(request_queue.put_calls[0]["type"], "predict_chunk")
         self.assertEqual(request_queue.put_calls[0]["worker_index"], 3)
         self.assertEqual(request_queue.put_calls[0]["server_index"], 1)
-        np.testing.assert_array_equal(first_action.numpy(), np.array([20.0], dtype=np.float32))
-        np.testing.assert_array_equal(second_action.numpy(), np.array([30.0], dtype=np.float32))
+        np.testing.assert_array_equal(first_action.numpy(), np.array([10.0], dtype=np.float32))
+        np.testing.assert_array_equal(second_action.numpy(), np.array([20.0], dtype=np.float32))
+
+    def test_remote_eval_worker_passes_agent_action_chunk_start_to_remote_runner(self):
+        cfg = OmegaConf.create(
+            {
+                "agent": {
+                    "obs_horizon": 2,
+                    "num_action_steps": 2,
+                    "action_chunk_start": 0,
+                    "camera_names": ["front"],
+                },
+                "eval": {
+                    "obs_horizon": 2,
+                    "num_queries": 2,
+                    "response_timeout_s": 3.0,
+                    "camera_names": ["front"],
+                },
+            }
+        )
+        captured = {}
+
+        class CapturingRunner:
+            def __init__(self, **kwargs):
+                captured.update(kwargs)
+
+        with mock.patch.object(eval_vla, "_RemotePolicyRunner", CapturingRunner), \
+             mock.patch.object(eval_vla, "_run_eval_episode_plans", return_value={"ok": True}) as run_plans:
+            result = eval_vla._run_remote_eval_worker(
+                cfg,
+                episode_plans=[{"episode_index": 0}],
+                worker_index=1,
+                server_index=2,
+                request_queue=_FakeQueue(),
+                response_queue=_FakeQueue(),
+            )
+
+        self.assertEqual(result, {"ok": True})
+        self.assertEqual(captured["action_chunk_start"], 0)
+        run_plans.assert_called_once()
 
     def test_merge_worker_summaries_sorts_episodes_and_recomputes_aggregates(self):
         worker_summaries = [

tests/test_smolvla_native_agent.py

@@ -0,0 +1,435 @@
+import contextlib
+import sys
+import types
+import unittest
+from pathlib import Path
+
+import torch
+from hydra import compose, initialize_config_dir
+from hydra.core.global_hydra import GlobalHydra
+from 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()
+
+
+def _stats():
+    return {
+        'qpos_min': [0.0, 10.0],
+        'qpos_max': [10.0, 20.0],
+        'action_min': [-10.0, 10.0],
+        'action_max': [10.0, 30.0],
+    }
+
+
+class FakeTokenizer:
+    def __init__(self):
+        self.calls = []
+
+    def __call__(self, texts, **kwargs):
+        self.calls.append({'texts': list(texts), 'kwargs': dict(kwargs)})
+        batch = len(texts)
+        if kwargs.get('padding') == 'max_length' and kwargs.get('max_length') is not None:
+            max_len = int(kwargs['max_length'])
+        else:
+            max_len = max(len(text) for text in texts) if texts else 0
+        input_ids = torch.zeros(batch, max_len, dtype=torch.long)
+        attention_mask = torch.zeros(batch, max_len, dtype=torch.long)
+        for i, text in enumerate(texts):
+            encoded = [ord(ch) % 127 for ch in text]
+            if kwargs.get('truncation') and len(encoded) > max_len:
+                encoded = encoded[:max_len]
+            if encoded:
+                input_ids[i, : len(encoded)] = torch.tensor(encoded, dtype=torch.long)
+                attention_mask[i, : len(encoded)] = 1
+        return {'input_ids': input_ids, 'attention_mask': attention_mask}
+
+
+class FakeNativeModel(nn.Module):
+    def __init__(self, action_dim=2, chunk_size=3):
+        super().__init__()
+        self.anchor = nn.Parameter(torch.tensor(0.0))
+        self.action_dim = action_dim
+        self.chunk_size = chunk_size
+        self.forward_calls = []
+        self.sample_calls = []
+        self.sample_return = torch.tensor([[[-1.0, 1.0], [0.0, 0.0], [1.0, -1.0]]])
+
+    def forward(self, **kwargs):
+        self.forward_calls.append({k: _clone(v) for k, v in kwargs.items()})
+        actions = kwargs['actions']
+        loss = (actions**2).sum(dim=-1)
+        if 'action_is_pad' in kwargs and kwargs['action_is_pad'] is not None:
+            mask = (~kwargs['action_is_pad']).to(loss.dtype)
+            loss = (loss * mask).sum() / mask.sum().clamp_min(1.0)
+        else:
+            loss = loss.mean()
+        return {'loss': loss + self.anchor}
+
+    def sample_actions(self, **kwargs):
+        self.sample_calls.append({k: _clone(v) for k, v in kwargs.items()})
+        return self.sample_return.to(device=kwargs['state'].device, dtype=kwargs['state'].dtype)
+
+
+def _clone(value):
+    if torch.is_tensor(value):
+        return value.detach().clone()
+    if isinstance(value, dict):
+        return {k: _clone(v) for k, v in value.items()}
+    if isinstance(value, list):
+        return list(value)
+    return value
+
+
+def _make_agent(**kwargs):
+    from roboimi.vla.agent_smolvla_native import SmolVLANativeAgent
+
+    params = dict(
+        model=FakeNativeModel(),
+        tokenizer=FakeTokenizer(),
+        action_dim=2,
+        obs_dim=2,
+        chunk_size=3,
+        n_action_steps=2,
+        obs_horizon=2,
+        camera_names=_CAMERA_NAMES,
+        num_cams=3,
+        dataset_stats=_stats(),
+        normalization_type='min_max',
+        task_description='fallback task',
+        model_config={'resize_imgs_with_padding': None},
+    )
+    params.update(kwargs)
+    return SmolVLANativeAgent(**params)
+
+
+def _batch(task=None):
+    images = {
+        'front': torch.full((2, 2, 1, 2, 2), 3.0),
+        'r_vis': torch.full((2, 2, 1, 2, 2), 1.0),
+        'top': torch.full((2, 2, 1, 2, 2), 2.0),
+    }
+    batch = {
+        'images': images,
+        'qpos': torch.tensor([[[0.0, 10.0], [10.0, 20.0]], [[5.0, 15.0], [10.0, 10.0]]]),
+        'action': torch.tensor([[[-10.0, 10.0], [0.0, 20.0], [10.0, 30.0]], [[0.0, 20.0], [10.0, 10.0], [-10.0, 30.0]]]),
+        'action_is_pad': torch.tensor([[False, False, True], [False, True, True]]),
+    }
+    if task is not None:
+        batch['task'] = task
+    return batch
+
+
+@contextlib.contextmanager
+def _stub_native_modules():
+    old_pkg = sys.modules.get('roboimi.vla.models.smolvla', _MISSING)
+    old_conf = sys.modules.get('roboimi.vla.models.smolvla.configuration', _MISSING)
+    try:
+        pkg = types.ModuleType('roboimi.vla.models.smolvla')
+        conf = types.ModuleType('roboimi.vla.models.smolvla.configuration')
+
+        class NativeSmolVLAConfig:
+            def __init__(self, **kwargs):
+                self.kwargs = kwargs
+
+        conf.NativeSmolVLAConfig = NativeSmolVLAConfig
+        sys.modules['roboimi.vla.models.smolvla'] = pkg
+        sys.modules['roboimi.vla.models.smolvla.configuration'] = conf
+        yield
+    finally:
+        for name, old in [
+            ('roboimi.vla.models.smolvla.configuration', old_conf),
+            ('roboimi.vla.models.smolvla', old_pkg),
+        ]:
+            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 StrictSignatureNativeModel(nn.Module):
+    def __init__(self, action_dim=2, max_action_dim=4, chunk_size=3):
+        super().__init__()
+        self.anchor = nn.Parameter(torch.tensor(0.0))
+        self.action_dim = action_dim
+        self.max_action_dim = max_action_dim
+        self.chunk_size = chunk_size
+        self.forward_calls = []
+        self.sample_calls = []
+
+    def forward(self, images, img_masks, lang_tokens, lang_masks, state, actions):
+        self.forward_calls.append({
+            'images': _clone(images),
+            'img_masks': _clone(img_masks),
+            'lang_tokens': _clone(lang_tokens),
+            'lang_masks': _clone(lang_masks),
+            'state': _clone(state),
+            'actions': _clone(actions),
+        })
+        # Per-element losses in the native core's max_action_dim space.
+        return actions.pow(2) + self.anchor
+
+    def sample_actions(self, images, img_masks, lang_tokens, lang_masks, state):
+        self.sample_calls.append({
+            'images': _clone(images),
+            'img_masks': _clone(img_masks),
+            'lang_tokens': _clone(lang_tokens),
+            'lang_masks': _clone(lang_masks),
+            'state': _clone(state),
+        })
+        batch_size = state.shape[0]
+        out = torch.zeros(batch_size, self.chunk_size, self.max_action_dim, device=state.device, dtype=state.dtype)
+        out[..., : self.action_dim] = torch.tensor(
+            [[[-1.0, 1.0], [0.0, 0.0], [1.0, -1.0]]],
+            device=state.device,
+            dtype=state.dtype,
+        ).expand(batch_size, -1, -1)
+        out[..., self.action_dim :] = 123.0
+        return out
+
+
+class SmolVLANativeAgentTest(unittest.TestCase):
+    def test_compute_loss_orders_normalizes_tokenizes_and_passes_action_pad(self):
+        agent = _make_agent()
+        batch = _batch(task=['pick', 'place'])
+
+        loss = agent.compute_loss(batch)
+
+        self.assertTrue(torch.isfinite(loss))
+        call = agent.model.forward_calls[-1]
+        self.assertIsInstance(call['images'], list)
+        self.assertEqual(len(call['images']), len(_CAMERA_NAMES))
+        self.assertTrue(torch.allclose(call['images'][0], torch.full((2, 1, 2, 2), 1.0)))
+        self.assertTrue(torch.allclose(call['state'], torch.tensor([[1.0, 1.0], [1.0, -1.0]])))
+        self.assertTrue(torch.allclose(call['actions'][0], torch.tensor([[-1.0, -1.0], [0.0, 0.0], [1.0, 1.0]])))
+        self.assertNotIn('action_is_pad', call)
+        self.assertEqual(agent.tokenizer.calls[-1]['texts'], ['pick\n', 'place\n'])
+        self.assertIn('lang_tokens', call)
+        self.assertIn('lang_masks', call)
+        self.assertEqual(call['lang_tokens'].dtype, torch.long)
+        self.assertEqual(call['lang_masks'].dtype, torch.bool)
+
+    def test_task_fallback_for_missing_empty_and_unknown_but_valid_task_wins(self):
+        agent = _make_agent(task_description='default instruction')
+
+        agent.compute_loss(_batch())
+        self.assertEqual(agent.tokenizer.calls[-1]['texts'], ['default instruction\n', 'default instruction\n'])
+
+        agent.compute_loss(_batch(task=[None, '']))
+        self.assertEqual(agent.tokenizer.calls[-1]['texts'], ['default instruction\n', 'default instruction\n'])
+
+        agent.compute_loss(_batch(task=['unknown', 'valid task']))
+        self.assertEqual(agent.tokenizer.calls[-1]['texts'], ['default instruction\n', 'valid task\n'])
+
+    def test_missing_camera_raises_value_error(self):
+        agent = _make_agent()
+        batch = _batch(task=['pick', 'place'])
+        del batch['images']['top']
+
+        with self.assertRaisesRegex(ValueError, 'missing.*top'):
+            agent.compute_loss(batch)
+
+    def test_predict_action_chunk_samples_and_denormalizes_actions(self):
+        agent = _make_agent()
+        batch = _batch(task=['pick', 'place'])
+        batch.pop('action')
+        batch.pop('action_is_pad')
+        agent.model.sample_return = torch.tensor([[[-1.0, 1.0], [0.0, 0.0], [1.0, -1.0]], [[0.5, -0.5], [-0.5, 0.5], [0.0, 1.0]]])
+
+        actions = agent.predict_action_chunk(batch)
+
+        self.assertTrue(torch.allclose(actions[0], torch.tensor([[-10.0, 30.0], [0.0, 20.0], [10.0, 10.0]])))
+        self.assertEqual(actions.shape, (2, 3, 2))
+        self.assertEqual(len(agent.model.sample_calls), 1)
+        self.assertTrue(torch.allclose(agent.model.sample_calls[-1]['state'][0], torch.tensor([1.0, 1.0])))
+
+    def test_select_action_queues_actions_and_defaults_chunk_start_zero(self):
+        agent = _make_agent(n_action_steps=2, action_chunk_start=0)
+        agent.model.sample_return = torch.tensor([[[-1.0, -1.0], [0.0, 0.0], [1.0, 1.0]]])
+        obs = {
+            'images': {name: torch.full((1, 2, 2), float(i + 1)) for i, name in enumerate(_CAMERA_NAMES)},
+            'qpos': torch.tensor([0.0, 10.0]),
+            'task': 'rollout task',
+        }
+
+        action0 = agent.select_action(obs)
+        action1 = agent.select_action(obs)
+
+        self.assertEqual(len(agent.model.sample_calls), 1)
+        self.assertTrue(torch.allclose(action0, torch.tensor([-10.0, 10.0])))
+        self.assertTrue(torch.allclose(action1, torch.tensor([0.0, 20.0])))
+        self.assertEqual(agent.action_chunk_start, 0)
+
+    def test_get_normalization_stats_returns_stats(self):
+        stats = _make_agent().get_normalization_stats()
+        self.assertEqual(stats['normalization_type'], 'min_max')
+        self.assertEqual(stats['qpos_min'], [0.0, 10.0])
+
+
+    def test_agent_adapts_roboimi_batch_to_native_vla_signature_and_reduces_loss(self):
+        from roboimi.vla.agent_smolvla_native import SmolVLANativeAgent
+
+        model = StrictSignatureNativeModel(action_dim=2, max_action_dim=4, chunk_size=3)
+        agent = SmolVLANativeAgent(
+            model=model,
+            tokenizer=FakeTokenizer(),
+            action_dim=2,
+            obs_dim=2,
+            chunk_size=3,
+            n_action_steps=2,
+            obs_horizon=2,
+            camera_names=_CAMERA_NAMES,
+            num_cams=3,
+            dataset_stats=_stats(),
+            normalization_type='min_max',
+            model_config={'max_state_dim': 4, 'max_action_dim': 4, 'resize_imgs_with_padding': None},
+        )
+
+        batch = _batch(task=['pick', 'place'])
+        batch['images'] = {
+            'front': torch.full((2, 2, 1, 2, 2), 0.6),
+            'r_vis': torch.full((2, 2, 1, 2, 2), 0.2),
+            'top': torch.full((2, 2, 1, 2, 2), 0.4),
+        }
+
+        loss = agent.compute_loss(batch)
+
+        self.assertEqual(loss.ndim, 0)
+        self.assertTrue(torch.isfinite(loss))
+        call = model.forward_calls[-1]
+        self.assertIsInstance(call['images'], list)
+        self.assertEqual(len(call['images']), len(_CAMERA_NAMES))
+        self.assertEqual(tuple(call['images'][0].shape), (2, 1, 2, 2))
+        self.assertTrue(torch.allclose(call['images'][0], torch.full((2, 1, 2, 2), -0.6)))
+        self.assertTrue(torch.allclose(call['images'][1], torch.full((2, 1, 2, 2), -0.2)))
+        self.assertTrue(torch.allclose(call['images'][2], torch.full((2, 1, 2, 2), 0.2)))
+        self.assertEqual(len(call['img_masks']), len(_CAMERA_NAMES))
+        self.assertTrue(torch.equal(call['img_masks'][0], torch.ones(2, dtype=torch.bool)))
+        self.assertTrue(torch.equal(call['lang_tokens'], model.forward_calls[-1]['lang_tokens']))
+        self.assertEqual(call['lang_tokens'].dtype, torch.long)
+        self.assertEqual(call['lang_masks'].dtype, torch.bool)
+        self.assertEqual(tuple(call['state'].shape), (2, 4))
+        self.assertTrue(torch.allclose(call['state'][0], torch.tensor([1.0, 1.0, 0.0, 0.0])))
+        self.assertEqual(tuple(call['actions'].shape), (2, 3, 4))
+        self.assertTrue(torch.allclose(call['actions'][0, :, :2], torch.tensor([[-1.0, -1.0], [0.0, 0.0], [1.0, 1.0]])))
+        self.assertTrue(torch.allclose(call['actions'][..., 2:], torch.zeros(2, 3, 2)))
+        # Valid entries: sample0 steps 0,1 and sample1 step 0, only first action_dim losses are reduced.
+        expected = (2.0 + 0.0 + 0.0) / (3 * 2)
+        self.assertTrue(torch.allclose(loss, torch.tensor(expected)))
+
+    def test_predict_action_chunk_accepts_native_max_action_dim_and_crops_before_denorm(self):
+        from roboimi.vla.agent_smolvla_native import SmolVLANativeAgent
+
+        model = StrictSignatureNativeModel(action_dim=2, max_action_dim=4, chunk_size=3)
+        agent = SmolVLANativeAgent(
+            model=model,
+            tokenizer=FakeTokenizer(),
+            action_dim=2,
+            obs_dim=2,
+            chunk_size=3,
+            n_action_steps=2,
+            obs_horizon=2,
+            camera_names=_CAMERA_NAMES,
+            num_cams=3,
+            dataset_stats=_stats(),
+            normalization_type='min_max',
+            model_config={'max_state_dim': 4, 'max_action_dim': 4, 'resize_imgs_with_padding': None},
+        )
+        batch = _batch(task=['pick', 'place'])
+        batch.pop('action')
+        batch.pop('action_is_pad')
+
+        actions = agent.predict_action_chunk(batch)
+
+        self.assertEqual(actions.shape, (2, 3, 2))
+        self.assertTrue(torch.allclose(actions[0], torch.tensor([[-10.0, 30.0], [0.0, 20.0], [10.0, 10.0]])))
+        call = model.sample_calls[-1]
+        self.assertEqual(tuple(call['state'].shape), (2, 4))
+        self.assertEqual(len(call['images']), len(_CAMERA_NAMES))
+
+    def test_build_model_filters_and_maps_non_core_config_fields(self):
+        from roboimi.vla.agent_smolvla_native import SmolVLANativeAgent
+        from roboimi.vla.models.smolvla.configuration import NativeSmolVLAConfig
+
+        captured = {}
+
+        class BuildOnlyAgent(SmolVLANativeAgent):
+            def _build_tokenizer(self, tokenizer_name):
+                return FakeTokenizer()
+
+            def _build_model(self):
+                cfg_kwargs = self._native_config_kwargs()
+                captured.update(cfg_kwargs)
+                return FakeNativeModel(action_dim=2, chunk_size=3)
+
+        agent = BuildOnlyAgent(
+            action_dim=2,
+            obs_dim=2,
+            chunk_size=3,
+            n_action_steps=2,
+            obs_horizon=2,
+            camera_names=_CAMERA_NAMES,
+            num_cams=3,
+            dataset_stats=_stats(),
+            normalization_type='min_max',
+            model_config={
+                'state_dim': 2,
+                'action_dim': 2,
+                'max_state_dim': 4,
+                'max_action_dim': 4,
+                'tokenizer_name': 'dummy-tokenizer',
+                'freeze_vlm': True,
+                'image_resize_shape': [512, 320],
+                'num_cameras': 3,
+                'load_vlm_weights': False,
+            },
+        )
+        self.assertIsNotNone(agent.model)
+        config = NativeSmolVLAConfig(**captured)
+        self.assertEqual(config.resize_imgs_with_padding, (512, 320))
+        self.assertEqual(config.max_state_dim, 4)
+        self.assertEqual(config.max_action_dim, 4)
+        self.assertNotIn('state_dim', captured)
+        self.assertNotIn('action_dim', captured)
+        self.assertNotIn('tokenizer_name', captured)
+        self.assertNotIn('freeze_vlm', captured)
+        self.assertNotIn('image_resize_shape', captured)
+        self.assertNotIn('num_cameras', captured)
+
+    def test_hydra_config_target_and_key_fields(self):
+        with _stub_native_modules():
+            cfg = _compose_cfg(overrides=['agent=smolvla_native'])
+        self.assertEqual(cfg.agent._target_, 'roboimi.vla.agent_smolvla_native.SmolVLANativeAgent')
+        self.assertEqual(cfg.agent.action_dim, 16)
+        self.assertEqual(cfg.agent.obs_dim, 16)
+        self.assertEqual(list(cfg.agent.camera_names), list(cfg.data.camera_names))
+        self.assertEqual(cfg.agent.num_cams, len(cfg.data.camera_names))
+        self.assertEqual(cfg.agent.chunk_size, 16)
+        self.assertEqual(cfg.agent.n_action_steps, 8)
+        self.assertEqual(cfg.agent.model_config.max_state_dim, 32)
+        self.assertEqual(cfg.agent.model_config.max_action_dim, 32)
+        self.assertIsNone(cfg.agent.dataset_image_resize_shape)
+        self.assertIsNone(cfg.agent.eval_image_resize_shape)
+        self.assertEqual(list(cfg.agent.model_config.resize_imgs_with_padding), [512, 512])
+        self.assertNotIn('state_dim', cfg.agent.model_config)
+        self.assertNotIn('action_dim', cfg.agent.model_config)
+        self.assertNotIn('tokenizer_name', cfg.agent.model_config)
+        self.assertEqual(cfg.agent.model_config.vlm_model_name, 'HuggingFaceTB/SmolVLM2-500M-Video-Instruct')
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/test_smolvla_native_modeling.py

@@ -0,0 +1,121 @@
+import unittest
+from types import SimpleNamespace
+
+import torch
+from torch import nn
+
+from roboimi.vla.models.smolvla import NativeSmolVLAConfig
+from roboimi.vla.models.smolvla.modeling import (
+    VLAFlowMatching,
+    make_att_2d_masks,
+    pad_vector,
+    resize_with_pad,
+)
+
+
+class FakeTokenizer:
+    fake_image_token_id = 32000
+    global_image_token_id = 32001
+
+
+class FakeProcessor:
+    tokenizer = FakeTokenizer()
+
+
+class FakeVLMWithExpert(nn.Module):
+    def __init__(self, vlm_hidden_size=8, expert_hidden_size=6, image_tokens=3, vocab_size=64):
+        super().__init__()
+        self.config = SimpleNamespace(text_config=SimpleNamespace(hidden_size=vlm_hidden_size))
+        self.expert_hidden_size = expert_hidden_size
+        self.image_tokens = image_tokens
+        self.processor = FakeProcessor()
+        self.vlm = SimpleNamespace(device=torch.device("cpu"))
+        self.image_proj = nn.Linear(3, vlm_hidden_size)
+        self.token_emb = nn.Embedding(vocab_size, vlm_hidden_size)
+        self.suffix_proj = nn.Linear(expert_hidden_size, expert_hidden_size)
+
+    def embed_image(self, image):
+        # Deterministic lightweight image embedding: pool pixels, then repeat.
+        pooled = image.mean(dim=(-1, -2)).to(dtype=torch.float32)
+        return self.image_proj(pooled).unsqueeze(1).expand(-1, self.image_tokens, -1)
+
+    def embed_language_tokens(self, tokens):
+        return self.token_emb(tokens)
+
+    def forward(self, attention_mask, position_ids, past_key_values, inputs_embeds, use_cache, fill_kv_cache):
+        prefix_embs, suffix_embs = inputs_embeds
+        prefix_out = prefix_embs if prefix_embs is not None else None
+        suffix_out = self.suffix_proj(suffix_embs) if suffix_embs is not None else None
+        cache = ("fake-cache",) if fill_kv_cache else past_key_values
+        return (prefix_out, suffix_out), cache
+
+
+class NativeSmolVLAModelingTest(unittest.TestCase):
+    def test_config_rejects_action_steps_greater_than_chunk_size(self):
+        with self.assertRaisesRegex(ValueError, "n_action_steps"):
+            NativeSmolVLAConfig(chunk_size=2, n_action_steps=3)
+
+    def test_pad_vector_pads_returns_original_for_equal_and_rejects_truncation(self):
+        vector = torch.tensor([[1.0, 2.0, 3.0]])
+        padded = pad_vector(vector, 5)
+        self.assertEqual(tuple(padded.shape), (1, 5))
+        torch.testing.assert_close(padded, torch.tensor([[1.0, 2.0, 3.0, 0.0, 0.0]]))
+
+        same = pad_vector(vector, 3)
+        self.assertIs(same, vector)
+
+        with self.assertRaisesRegex(ValueError, "target dimension"):
+            pad_vector(vector, 2)
+
+    def test_resize_with_pad_returns_requested_spatial_size(self):
+        img = torch.arange(2 * 3 * 4 * 8, dtype=torch.float32).reshape(2, 3, 4, 8)
+        resized = resize_with_pad(img, width=10, height=10, pad_value=-1)
+        self.assertEqual(tuple(resized.shape), (2, 3, 10, 10))
+
+    def test_make_att_2d_masks_implements_prefix_lm_semantics(self):
+        pad_masks = torch.tensor([[True, True, True, True, False]])
+        # First two tokens are bidirectional prefix, later valid tokens are causal.
+        att_masks = torch.tensor([[False, False, True, True, True]])
+        mask = make_att_2d_masks(pad_masks, att_masks)
+        expected = torch.tensor(
+            [[
+                [True, True, False, False, False],
+                [True, True, False, False, False],
+                [True, True, True, False, False],
+                [True, True, True, True, False],
+                [False, False, False, False, False],
+            ]]
+        )
+        torch.testing.assert_close(mask, expected)
+
+    def test_vla_flow_matching_forward_and_sample_actions_with_fake_vlm(self):
+        torch.manual_seed(0)
+        config = NativeSmolVLAConfig(
+            chunk_size=4,
+            n_action_steps=4,
+            max_state_dim=5,
+            max_action_dim=3,
+            num_steps=2,
+            prefix_length=8,
+            add_image_special_tokens=False,
+        )
+        fake_vlm = FakeVLMWithExpert(vlm_hidden_size=8, expert_hidden_size=6)
+        model = VLAFlowMatching(config, vlm_with_expert=fake_vlm)
+
+        bsize = 2
+        images = [torch.randn(bsize, 3, 6, 6)]
+        img_masks = [torch.tensor([True, False])]
+        lang_tokens = torch.tensor([[1, 2, 3], [4, 5, 0]])
+        lang_masks = torch.tensor([[True, True, True], [True, True, False]])
+        state = torch.randn(bsize, config.max_state_dim)
+        actions = torch.randn(bsize, config.chunk_size, config.max_action_dim)
+
+        losses = model(images, img_masks, lang_tokens, lang_masks, state, actions)
+        self.assertEqual(tuple(losses.shape), (bsize, config.chunk_size, config.max_action_dim))
+
+        sampled = model.sample_actions(images, img_masks, lang_tokens, lang_masks, state)
+        self.assertEqual(tuple(sampled.shape), (bsize, config.chunk_size, config.max_action_dim))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_train_vla_rollout_validation.py

@@ -394,6 +394,23 @@ class TrainVLARolloutValidationTest(unittest.TestCase):
         self.assertTrue(rollout_cfg.eval.save_summary_json)
         self.assertTrue(rollout_cfg.eval.save_trajectory_image)
 
+    def test_resolve_dataset_image_resize_shape_prefers_agent_top_level_override(self):
+        cfg = OmegaConf.create(
+            {
+                'agent': {
+                    'dataset_image_resize_shape': None,
+                    'vision_backbone': {
+                        'dataset_image_resize_shape': [256, 256],
+                    },
+                },
+                'data': {
+                    'image_resize_shape': [224, 224],
+                },
+            }
+        )
+
+        self.assertIsNone(train_vla._resolve_dataset_image_resize_shape(cfg))
+
     def test_training_passes_backbone_image_resize_override_to_dataset_instantiation(self):
         cfg = OmegaConf.create(
             {