From 2ac926d427bf232cd478f3bf0a93c380787b863d Mon Sep 17 00:00:00 2001 From: Logic Date: Mon, 25 May 2026 23:24:36 +0800 Subject: [PATCH] Add native SmolVLA model integration --- ...26-05-25-native-smolvla-model-migration.md | 184 ++++++ ...5-native-smolvla-model-migration-design.md | 86 +++ roboimi/demos/vla_scripts/eval_vla.py | 10 +- roboimi/demos/vla_scripts/train_vla.py | 20 +- roboimi/vla/agent_smolvla_native.py | 395 ++++++++++++ roboimi/vla/conf/agent/smolvla_native.yaml | 35 ++ roboimi/vla/models/smolvla/__init__.py | 13 + roboimi/vla/models/smolvla/configuration.py | 75 +++ roboimi/vla/models/smolvla/modeling.py | 421 +++++++++++++ .../vla/models/smolvla/smolvlm_with_expert.py | 571 ++++++++++++++++++ tests/test_eval_vla_execution.py | 79 ++- tests/test_smolvla_native_agent.py | 435 +++++++++++++ tests/test_smolvla_native_modeling.py | 121 ++++ tests/test_train_vla_rollout_validation.py | 17 + 14 files changed, 2453 insertions(+), 9 deletions(-) create mode 100644 docs/superpowers/plans/2026-05-25-native-smolvla-model-migration.md create mode 100644 docs/superpowers/specs/2026-05-25-native-smolvla-model-migration-design.md create mode 100644 roboimi/vla/agent_smolvla_native.py create mode 100644 roboimi/vla/conf/agent/smolvla_native.yaml create mode 100644 roboimi/vla/models/smolvla/__init__.py create mode 100644 roboimi/vla/models/smolvla/configuration.py create mode 100644 roboimi/vla/models/smolvla/modeling.py create mode 100644 roboimi/vla/models/smolvla/smolvlm_with_expert.py create mode 100644 tests/test_smolvla_native_agent.py create mode 100644 tests/test_smolvla_native_modeling.py diff --git a/docs/superpowers/plans/2026-05-25-native-smolvla-model-migration.md b/docs/superpowers/plans/2026-05-25-native-smolvla-model-migration.md new file mode 100644 index 0000000..3c222cf --- /dev/null +++ b/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. diff --git a/docs/superpowers/specs/2026-05-25-native-smolvla-model-migration-design.md b/docs/superpowers/specs/2026-05-25-native-smolvla-model-migration-design.md new file mode 100644 index 0000000..29316a8 --- /dev/null +++ b/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. diff --git a/roboimi/demos/vla_scripts/eval_vla.py b/roboimi/demos/vla_scripts/eval_vla.py index bc89378..3d35d4a 100644 --- a/roboimi/demos/vla_scripts/eval_vla.py +++ b/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( diff --git a/roboimi/demos/vla_scripts/train_vla.py b/roboimi/demos/vla_scripts/train_vla.py index e7fac56..f42850a 100644 --- a/roboimi/demos/vla_scripts/train_vla.py +++ b/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) ) diff --git a/roboimi/vla/agent_smolvla_native.py b/roboimi/vla/agent_smolvla_native.py new file mode 100644 index 0000000..d73b74d --- /dev/null +++ b/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() diff --git a/roboimi/vla/conf/agent/smolvla_native.yaml b/roboimi/vla/conf/agent/smolvla_native.yaml new file mode 100644 index 0000000..283253d --- /dev/null +++ b/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] diff --git a/roboimi/vla/models/smolvla/__init__.py b/roboimi/vla/models/smolvla/__init__.py new file mode 100644 index 0000000..06de9c9 --- /dev/null +++ b/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", +] diff --git a/roboimi/vla/models/smolvla/configuration.py b/roboimi/vla/models/smolvla/configuration.py new file mode 100644 index 0000000..58d99af --- /dev/null +++ b/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 diff --git a/roboimi/vla/models/smolvla/modeling.py b/roboimi/vla/models/smolvla/modeling.py new file mode 100644 index 0000000..6349080 --- /dev/null +++ b/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) diff --git a/roboimi/vla/models/smolvla/smolvlm_with_expert.py b/roboimi/vla/models/smolvla/smolvlm_with_expert.py new file mode 100644 index 0000000..c464819 --- /dev/null +++ b/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 diff --git a/tests/test_eval_vla_execution.py b/tests/test_eval_vla_execution.py index db0fa7a..daea8dc 100644 --- a/tests/test_eval_vla_execution.py +++ b/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 = [ diff --git a/tests/test_smolvla_native_agent.py b/tests/test_smolvla_native_agent.py new file mode 100644 index 0000000..237e857 --- /dev/null +++ b/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() diff --git a/tests/test_smolvla_native_modeling.py b/tests/test_smolvla_native_modeling.py new file mode 100644 index 0000000..cf05f8e --- /dev/null +++ b/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() diff --git a/tests/test_train_vla_rollout_validation.py b/tests/test_train_vla_rollout_validation.py index 295ea80..60b31b9 100644 --- a/tests/test_train_vla_rollout_validation.py +++ b/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( {