Add native SmolVLA model integration
This commit is contained in:
Logic
@@ -0,0 +1,435 @@
|
||||
import contextlib
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
from hydra import compose, initialize_config_dir
|
||||
from hydra.core.global_hydra import GlobalHydra
|
||||
from omegaconf import OmegaConf
|
||||
from torch import nn
|
||||
|
||||
|
||||
_REPO_ROOT = Path(__file__).resolve().parents[1]
|
||||
_CONFIG_DIR = str((_REPO_ROOT / 'roboimi/vla/conf').resolve())
|
||||
_CAMERA_NAMES = ('r_vis', 'top', 'front')
|
||||
_MISSING = object()
|
||||
|
||||
|
||||
def _stats():
|
||||
return {
|
||||
'qpos_min': [0.0, 10.0],
|
||||
'qpos_max': [10.0, 20.0],
|
||||
'action_min': [-10.0, 10.0],
|
||||
'action_max': [10.0, 30.0],
|
||||
}
|
||||
|
||||
|
||||
class FakeTokenizer:
|
||||
def __init__(self):
|
||||
self.calls = []
|
||||
|
||||
def __call__(self, texts, **kwargs):
|
||||
self.calls.append({'texts': list(texts), 'kwargs': dict(kwargs)})
|
||||
batch = len(texts)
|
||||
if kwargs.get('padding') == 'max_length' and kwargs.get('max_length') is not None:
|
||||
max_len = int(kwargs['max_length'])
|
||||
else:
|
||||
max_len = max(len(text) for text in texts) if texts else 0
|
||||
input_ids = torch.zeros(batch, max_len, dtype=torch.long)
|
||||
attention_mask = torch.zeros(batch, max_len, dtype=torch.long)
|
||||
for i, text in enumerate(texts):
|
||||
encoded = [ord(ch) % 127 for ch in text]
|
||||
if kwargs.get('truncation') and len(encoded) > max_len:
|
||||
encoded = encoded[:max_len]
|
||||
if encoded:
|
||||
input_ids[i, : len(encoded)] = torch.tensor(encoded, dtype=torch.long)
|
||||
attention_mask[i, : len(encoded)] = 1
|
||||
return {'input_ids': input_ids, 'attention_mask': attention_mask}
|
||||
|
||||
|
||||
class FakeNativeModel(nn.Module):
|
||||
def __init__(self, action_dim=2, chunk_size=3):
|
||||
super().__init__()
|
||||
self.anchor = nn.Parameter(torch.tensor(0.0))
|
||||
self.action_dim = action_dim
|
||||
self.chunk_size = chunk_size
|
||||
self.forward_calls = []
|
||||
self.sample_calls = []
|
||||
self.sample_return = torch.tensor([[[-1.0, 1.0], [0.0, 0.0], [1.0, -1.0]]])
|
||||
|
||||
def forward(self, **kwargs):
|
||||
self.forward_calls.append({k: _clone(v) for k, v in kwargs.items()})
|
||||
actions = kwargs['actions']
|
||||
loss = (actions**2).sum(dim=-1)
|
||||
if 'action_is_pad' in kwargs and kwargs['action_is_pad'] is not None:
|
||||
mask = (~kwargs['action_is_pad']).to(loss.dtype)
|
||||
loss = (loss * mask).sum() / mask.sum().clamp_min(1.0)
|
||||
else:
|
||||
loss = loss.mean()
|
||||
return {'loss': loss + self.anchor}
|
||||
|
||||
def sample_actions(self, **kwargs):
|
||||
self.sample_calls.append({k: _clone(v) for k, v in kwargs.items()})
|
||||
return self.sample_return.to(device=kwargs['state'].device, dtype=kwargs['state'].dtype)
|
||||
|
||||
|
||||
def _clone(value):
|
||||
if torch.is_tensor(value):
|
||||
return value.detach().clone()
|
||||
if isinstance(value, dict):
|
||||
return {k: _clone(v) for k, v in value.items()}
|
||||
if isinstance(value, list):
|
||||
return list(value)
|
||||
return value
|
||||
|
||||
|
||||
def _make_agent(**kwargs):
|
||||
from roboimi.vla.agent_smolvla_native import SmolVLANativeAgent
|
||||
|
||||
params = dict(
|
||||
model=FakeNativeModel(),
|
||||
tokenizer=FakeTokenizer(),
|
||||
action_dim=2,
|
||||
obs_dim=2,
|
||||
chunk_size=3,
|
||||
n_action_steps=2,
|
||||
obs_horizon=2,
|
||||
camera_names=_CAMERA_NAMES,
|
||||
num_cams=3,
|
||||
dataset_stats=_stats(),
|
||||
normalization_type='min_max',
|
||||
task_description='fallback task',
|
||||
model_config={'resize_imgs_with_padding': None},
|
||||
)
|
||||
params.update(kwargs)
|
||||
return SmolVLANativeAgent(**params)
|
||||
|
||||
|
||||
def _batch(task=None):
|
||||
images = {
|
||||
'front': torch.full((2, 2, 1, 2, 2), 3.0),
|
||||
'r_vis': torch.full((2, 2, 1, 2, 2), 1.0),
|
||||
'top': torch.full((2, 2, 1, 2, 2), 2.0),
|
||||
}
|
||||
batch = {
|
||||
'images': images,
|
||||
'qpos': torch.tensor([[[0.0, 10.0], [10.0, 20.0]], [[5.0, 15.0], [10.0, 10.0]]]),
|
||||
'action': torch.tensor([[[-10.0, 10.0], [0.0, 20.0], [10.0, 30.0]], [[0.0, 20.0], [10.0, 10.0], [-10.0, 30.0]]]),
|
||||
'action_is_pad': torch.tensor([[False, False, True], [False, True, True]]),
|
||||
}
|
||||
if task is not None:
|
||||
batch['task'] = task
|
||||
return batch
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _stub_native_modules():
|
||||
old_pkg = sys.modules.get('roboimi.vla.models.smolvla', _MISSING)
|
||||
old_conf = sys.modules.get('roboimi.vla.models.smolvla.configuration', _MISSING)
|
||||
try:
|
||||
pkg = types.ModuleType('roboimi.vla.models.smolvla')
|
||||
conf = types.ModuleType('roboimi.vla.models.smolvla.configuration')
|
||||
|
||||
class NativeSmolVLAConfig:
|
||||
def __init__(self, **kwargs):
|
||||
self.kwargs = kwargs
|
||||
|
||||
conf.NativeSmolVLAConfig = NativeSmolVLAConfig
|
||||
sys.modules['roboimi.vla.models.smolvla'] = pkg
|
||||
sys.modules['roboimi.vla.models.smolvla.configuration'] = conf
|
||||
yield
|
||||
finally:
|
||||
for name, old in [
|
||||
('roboimi.vla.models.smolvla.configuration', old_conf),
|
||||
('roboimi.vla.models.smolvla', old_pkg),
|
||||
]:
|
||||
if old is _MISSING:
|
||||
sys.modules.pop(name, None)
|
||||
else:
|
||||
sys.modules[name] = old
|
||||
|
||||
|
||||
def _compose_cfg(overrides=None):
|
||||
if not OmegaConf.has_resolver('len'):
|
||||
OmegaConf.register_new_resolver('len', lambda x: len(x))
|
||||
GlobalHydra.instance().clear()
|
||||
with initialize_config_dir(version_base=None, config_dir=_CONFIG_DIR):
|
||||
return compose(config_name='config', overrides=list(overrides or []))
|
||||
|
||||
|
||||
|
||||
|
||||
class StrictSignatureNativeModel(nn.Module):
|
||||
def __init__(self, action_dim=2, max_action_dim=4, chunk_size=3):
|
||||
super().__init__()
|
||||
self.anchor = nn.Parameter(torch.tensor(0.0))
|
||||
self.action_dim = action_dim
|
||||
self.max_action_dim = max_action_dim
|
||||
self.chunk_size = chunk_size
|
||||
self.forward_calls = []
|
||||
self.sample_calls = []
|
||||
|
||||
def forward(self, images, img_masks, lang_tokens, lang_masks, state, actions):
|
||||
self.forward_calls.append({
|
||||
'images': _clone(images),
|
||||
'img_masks': _clone(img_masks),
|
||||
'lang_tokens': _clone(lang_tokens),
|
||||
'lang_masks': _clone(lang_masks),
|
||||
'state': _clone(state),
|
||||
'actions': _clone(actions),
|
||||
})
|
||||
# Per-element losses in the native core's max_action_dim space.
|
||||
return actions.pow(2) + self.anchor
|
||||
|
||||
def sample_actions(self, images, img_masks, lang_tokens, lang_masks, state):
|
||||
self.sample_calls.append({
|
||||
'images': _clone(images),
|
||||
'img_masks': _clone(img_masks),
|
||||
'lang_tokens': _clone(lang_tokens),
|
||||
'lang_masks': _clone(lang_masks),
|
||||
'state': _clone(state),
|
||||
})
|
||||
batch_size = state.shape[0]
|
||||
out = torch.zeros(batch_size, self.chunk_size, self.max_action_dim, device=state.device, dtype=state.dtype)
|
||||
out[..., : self.action_dim] = torch.tensor(
|
||||
[[[-1.0, 1.0], [0.0, 0.0], [1.0, -1.0]]],
|
||||
device=state.device,
|
||||
dtype=state.dtype,
|
||||
).expand(batch_size, -1, -1)
|
||||
out[..., self.action_dim :] = 123.0
|
||||
return out
|
||||
|
||||
|
||||
class SmolVLANativeAgentTest(unittest.TestCase):
|
||||
def test_compute_loss_orders_normalizes_tokenizes_and_passes_action_pad(self):
|
||||
agent = _make_agent()
|
||||
batch = _batch(task=['pick', 'place'])
|
||||
|
||||
loss = agent.compute_loss(batch)
|
||||
|
||||
self.assertTrue(torch.isfinite(loss))
|
||||
call = agent.model.forward_calls[-1]
|
||||
self.assertIsInstance(call['images'], list)
|
||||
self.assertEqual(len(call['images']), len(_CAMERA_NAMES))
|
||||
self.assertTrue(torch.allclose(call['images'][0], torch.full((2, 1, 2, 2), 1.0)))
|
||||
self.assertTrue(torch.allclose(call['state'], torch.tensor([[1.0, 1.0], [1.0, -1.0]])))
|
||||
self.assertTrue(torch.allclose(call['actions'][0], torch.tensor([[-1.0, -1.0], [0.0, 0.0], [1.0, 1.0]])))
|
||||
self.assertNotIn('action_is_pad', call)
|
||||
self.assertEqual(agent.tokenizer.calls[-1]['texts'], ['pick\n', 'place\n'])
|
||||
self.assertIn('lang_tokens', call)
|
||||
self.assertIn('lang_masks', call)
|
||||
self.assertEqual(call['lang_tokens'].dtype, torch.long)
|
||||
self.assertEqual(call['lang_masks'].dtype, torch.bool)
|
||||
|
||||
def test_task_fallback_for_missing_empty_and_unknown_but_valid_task_wins(self):
|
||||
agent = _make_agent(task_description='default instruction')
|
||||
|
||||
agent.compute_loss(_batch())
|
||||
self.assertEqual(agent.tokenizer.calls[-1]['texts'], ['default instruction\n', 'default instruction\n'])
|
||||
|
||||
agent.compute_loss(_batch(task=[None, '']))
|
||||
self.assertEqual(agent.tokenizer.calls[-1]['texts'], ['default instruction\n', 'default instruction\n'])
|
||||
|
||||
agent.compute_loss(_batch(task=['unknown', 'valid task']))
|
||||
self.assertEqual(agent.tokenizer.calls[-1]['texts'], ['default instruction\n', 'valid task\n'])
|
||||
|
||||
def test_missing_camera_raises_value_error(self):
|
||||
agent = _make_agent()
|
||||
batch = _batch(task=['pick', 'place'])
|
||||
del batch['images']['top']
|
||||
|
||||
with self.assertRaisesRegex(ValueError, 'missing.*top'):
|
||||
agent.compute_loss(batch)
|
||||
|
||||
def test_predict_action_chunk_samples_and_denormalizes_actions(self):
|
||||
agent = _make_agent()
|
||||
batch = _batch(task=['pick', 'place'])
|
||||
batch.pop('action')
|
||||
batch.pop('action_is_pad')
|
||||
agent.model.sample_return = torch.tensor([[[-1.0, 1.0], [0.0, 0.0], [1.0, -1.0]], [[0.5, -0.5], [-0.5, 0.5], [0.0, 1.0]]])
|
||||
|
||||
actions = agent.predict_action_chunk(batch)
|
||||
|
||||
self.assertTrue(torch.allclose(actions[0], torch.tensor([[-10.0, 30.0], [0.0, 20.0], [10.0, 10.0]])))
|
||||
self.assertEqual(actions.shape, (2, 3, 2))
|
||||
self.assertEqual(len(agent.model.sample_calls), 1)
|
||||
self.assertTrue(torch.allclose(agent.model.sample_calls[-1]['state'][0], torch.tensor([1.0, 1.0])))
|
||||
|
||||
def test_select_action_queues_actions_and_defaults_chunk_start_zero(self):
|
||||
agent = _make_agent(n_action_steps=2, action_chunk_start=0)
|
||||
agent.model.sample_return = torch.tensor([[[-1.0, -1.0], [0.0, 0.0], [1.0, 1.0]]])
|
||||
obs = {
|
||||
'images': {name: torch.full((1, 2, 2), float(i + 1)) for i, name in enumerate(_CAMERA_NAMES)},
|
||||
'qpos': torch.tensor([0.0, 10.0]),
|
||||
'task': 'rollout task',
|
||||
}
|
||||
|
||||
action0 = agent.select_action(obs)
|
||||
action1 = agent.select_action(obs)
|
||||
|
||||
self.assertEqual(len(agent.model.sample_calls), 1)
|
||||
self.assertTrue(torch.allclose(action0, torch.tensor([-10.0, 10.0])))
|
||||
self.assertTrue(torch.allclose(action1, torch.tensor([0.0, 20.0])))
|
||||
self.assertEqual(agent.action_chunk_start, 0)
|
||||
|
||||
def test_get_normalization_stats_returns_stats(self):
|
||||
stats = _make_agent().get_normalization_stats()
|
||||
self.assertEqual(stats['normalization_type'], 'min_max')
|
||||
self.assertEqual(stats['qpos_min'], [0.0, 10.0])
|
||||
|
||||
|
||||
def test_agent_adapts_roboimi_batch_to_native_vla_signature_and_reduces_loss(self):
|
||||
from roboimi.vla.agent_smolvla_native import SmolVLANativeAgent
|
||||
|
||||
model = StrictSignatureNativeModel(action_dim=2, max_action_dim=4, chunk_size=3)
|
||||
agent = SmolVLANativeAgent(
|
||||
model=model,
|
||||
tokenizer=FakeTokenizer(),
|
||||
action_dim=2,
|
||||
obs_dim=2,
|
||||
chunk_size=3,
|
||||
n_action_steps=2,
|
||||
obs_horizon=2,
|
||||
camera_names=_CAMERA_NAMES,
|
||||
num_cams=3,
|
||||
dataset_stats=_stats(),
|
||||
normalization_type='min_max',
|
||||
model_config={'max_state_dim': 4, 'max_action_dim': 4, 'resize_imgs_with_padding': None},
|
||||
)
|
||||
|
||||
batch = _batch(task=['pick', 'place'])
|
||||
batch['images'] = {
|
||||
'front': torch.full((2, 2, 1, 2, 2), 0.6),
|
||||
'r_vis': torch.full((2, 2, 1, 2, 2), 0.2),
|
||||
'top': torch.full((2, 2, 1, 2, 2), 0.4),
|
||||
}
|
||||
|
||||
loss = agent.compute_loss(batch)
|
||||
|
||||
self.assertEqual(loss.ndim, 0)
|
||||
self.assertTrue(torch.isfinite(loss))
|
||||
call = model.forward_calls[-1]
|
||||
self.assertIsInstance(call['images'], list)
|
||||
self.assertEqual(len(call['images']), len(_CAMERA_NAMES))
|
||||
self.assertEqual(tuple(call['images'][0].shape), (2, 1, 2, 2))
|
||||
self.assertTrue(torch.allclose(call['images'][0], torch.full((2, 1, 2, 2), -0.6)))
|
||||
self.assertTrue(torch.allclose(call['images'][1], torch.full((2, 1, 2, 2), -0.2)))
|
||||
self.assertTrue(torch.allclose(call['images'][2], torch.full((2, 1, 2, 2), 0.2)))
|
||||
self.assertEqual(len(call['img_masks']), len(_CAMERA_NAMES))
|
||||
self.assertTrue(torch.equal(call['img_masks'][0], torch.ones(2, dtype=torch.bool)))
|
||||
self.assertTrue(torch.equal(call['lang_tokens'], model.forward_calls[-1]['lang_tokens']))
|
||||
self.assertEqual(call['lang_tokens'].dtype, torch.long)
|
||||
self.assertEqual(call['lang_masks'].dtype, torch.bool)
|
||||
self.assertEqual(tuple(call['state'].shape), (2, 4))
|
||||
self.assertTrue(torch.allclose(call['state'][0], torch.tensor([1.0, 1.0, 0.0, 0.0])))
|
||||
self.assertEqual(tuple(call['actions'].shape), (2, 3, 4))
|
||||
self.assertTrue(torch.allclose(call['actions'][0, :, :2], torch.tensor([[-1.0, -1.0], [0.0, 0.0], [1.0, 1.0]])))
|
||||
self.assertTrue(torch.allclose(call['actions'][..., 2:], torch.zeros(2, 3, 2)))
|
||||
# Valid entries: sample0 steps 0,1 and sample1 step 0, only first action_dim losses are reduced.
|
||||
expected = (2.0 + 0.0 + 0.0) / (3 * 2)
|
||||
self.assertTrue(torch.allclose(loss, torch.tensor(expected)))
|
||||
|
||||
def test_predict_action_chunk_accepts_native_max_action_dim_and_crops_before_denorm(self):
|
||||
from roboimi.vla.agent_smolvla_native import SmolVLANativeAgent
|
||||
|
||||
model = StrictSignatureNativeModel(action_dim=2, max_action_dim=4, chunk_size=3)
|
||||
agent = SmolVLANativeAgent(
|
||||
model=model,
|
||||
tokenizer=FakeTokenizer(),
|
||||
action_dim=2,
|
||||
obs_dim=2,
|
||||
chunk_size=3,
|
||||
n_action_steps=2,
|
||||
obs_horizon=2,
|
||||
camera_names=_CAMERA_NAMES,
|
||||
num_cams=3,
|
||||
dataset_stats=_stats(),
|
||||
normalization_type='min_max',
|
||||
model_config={'max_state_dim': 4, 'max_action_dim': 4, 'resize_imgs_with_padding': None},
|
||||
)
|
||||
batch = _batch(task=['pick', 'place'])
|
||||
batch.pop('action')
|
||||
batch.pop('action_is_pad')
|
||||
|
||||
actions = agent.predict_action_chunk(batch)
|
||||
|
||||
self.assertEqual(actions.shape, (2, 3, 2))
|
||||
self.assertTrue(torch.allclose(actions[0], torch.tensor([[-10.0, 30.0], [0.0, 20.0], [10.0, 10.0]])))
|
||||
call = model.sample_calls[-1]
|
||||
self.assertEqual(tuple(call['state'].shape), (2, 4))
|
||||
self.assertEqual(len(call['images']), len(_CAMERA_NAMES))
|
||||
|
||||
def test_build_model_filters_and_maps_non_core_config_fields(self):
|
||||
from roboimi.vla.agent_smolvla_native import SmolVLANativeAgent
|
||||
from roboimi.vla.models.smolvla.configuration import NativeSmolVLAConfig
|
||||
|
||||
captured = {}
|
||||
|
||||
class BuildOnlyAgent(SmolVLANativeAgent):
|
||||
def _build_tokenizer(self, tokenizer_name):
|
||||
return FakeTokenizer()
|
||||
|
||||
def _build_model(self):
|
||||
cfg_kwargs = self._native_config_kwargs()
|
||||
captured.update(cfg_kwargs)
|
||||
return FakeNativeModel(action_dim=2, chunk_size=3)
|
||||
|
||||
agent = BuildOnlyAgent(
|
||||
action_dim=2,
|
||||
obs_dim=2,
|
||||
chunk_size=3,
|
||||
n_action_steps=2,
|
||||
obs_horizon=2,
|
||||
camera_names=_CAMERA_NAMES,
|
||||
num_cams=3,
|
||||
dataset_stats=_stats(),
|
||||
normalization_type='min_max',
|
||||
model_config={
|
||||
'state_dim': 2,
|
||||
'action_dim': 2,
|
||||
'max_state_dim': 4,
|
||||
'max_action_dim': 4,
|
||||
'tokenizer_name': 'dummy-tokenizer',
|
||||
'freeze_vlm': True,
|
||||
'image_resize_shape': [512, 320],
|
||||
'num_cameras': 3,
|
||||
'load_vlm_weights': False,
|
||||
},
|
||||
)
|
||||
self.assertIsNotNone(agent.model)
|
||||
config = NativeSmolVLAConfig(**captured)
|
||||
self.assertEqual(config.resize_imgs_with_padding, (512, 320))
|
||||
self.assertEqual(config.max_state_dim, 4)
|
||||
self.assertEqual(config.max_action_dim, 4)
|
||||
self.assertNotIn('state_dim', captured)
|
||||
self.assertNotIn('action_dim', captured)
|
||||
self.assertNotIn('tokenizer_name', captured)
|
||||
self.assertNotIn('freeze_vlm', captured)
|
||||
self.assertNotIn('image_resize_shape', captured)
|
||||
self.assertNotIn('num_cameras', captured)
|
||||
|
||||
def test_hydra_config_target_and_key_fields(self):
|
||||
with _stub_native_modules():
|
||||
cfg = _compose_cfg(overrides=['agent=smolvla_native'])
|
||||
self.assertEqual(cfg.agent._target_, 'roboimi.vla.agent_smolvla_native.SmolVLANativeAgent')
|
||||
self.assertEqual(cfg.agent.action_dim, 16)
|
||||
self.assertEqual(cfg.agent.obs_dim, 16)
|
||||
self.assertEqual(list(cfg.agent.camera_names), list(cfg.data.camera_names))
|
||||
self.assertEqual(cfg.agent.num_cams, len(cfg.data.camera_names))
|
||||
self.assertEqual(cfg.agent.chunk_size, 16)
|
||||
self.assertEqual(cfg.agent.n_action_steps, 8)
|
||||
self.assertEqual(cfg.agent.model_config.max_state_dim, 32)
|
||||
self.assertEqual(cfg.agent.model_config.max_action_dim, 32)
|
||||
self.assertIsNone(cfg.agent.dataset_image_resize_shape)
|
||||
self.assertIsNone(cfg.agent.eval_image_resize_shape)
|
||||
self.assertEqual(list(cfg.agent.model_config.resize_imgs_with_padding), [512, 512])
|
||||
self.assertNotIn('state_dim', cfg.agent.model_config)
|
||||
self.assertNotIn('action_dim', cfg.agent.model_config)
|
||||
self.assertNotIn('tokenizer_name', cfg.agent.model_config)
|
||||
self.assertEqual(cfg.agent.model_config.vlm_model_name, 'HuggingFaceTB/SmolVLM2-500M-Video-Instruct')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user