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

# 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.