Logic
122 lines
4.8 KiB
Python
122 lines
4.8 KiB
Python
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()
|