Files
roboimi/docs/superpowers/plans/2026-05-25-native-smolvla-model-migration.md
2026-05-25 23:24:36 +08:00

6.7 KiB

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:

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:

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:

rg -n "import lerobot|from lerobot" roboimi/vla/models/smolvla roboimi/vla/agent_smolvla_native.py

Expected: no matches.

  • Step 3: Commit

Run:

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.