Add native SmolVLA model integration
This commit is contained in:
Logic
@@ -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.
|
||||
@@ -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.
|
||||
Reference in New Issue
Block a user