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()