Add native SmolVLA model integration

2026-05-25 23:24:36 +08:00
parent d94eb8f70b
commit 2ac926d427
14 changed files with 2453 additions and 9 deletions
@@ -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()