Add native SmolVLA model integration
This commit is contained in:
Logic
@@ -129,6 +129,23 @@ class EvalVLAExecutionTest(unittest.TestCase):
|
||||
["r_vis", "top", "front"],
|
||||
)
|
||||
|
||||
def test_resolve_eval_image_resize_shape_prefers_agent_top_level_override(self):
|
||||
cfg = OmegaConf.create(
|
||||
{
|
||||
"agent": {
|
||||
"eval_image_resize_shape": None,
|
||||
"condition_encoder": {
|
||||
"eval_image_resize_shape": [256, 256],
|
||||
},
|
||||
},
|
||||
"data": {
|
||||
"image_resize_shape": [224, 224],
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
self.assertIsNone(eval_vla._resolve_eval_image_resize_shape(cfg))
|
||||
|
||||
def test_resolve_eval_image_resize_shape_prefers_condition_encoder_override(self):
|
||||
cfg = OmegaConf.create(
|
||||
{
|
||||
@@ -270,6 +287,25 @@ class EvalVLAExecutionTest(unittest.TestCase):
|
||||
np.testing.assert_array_equal(queues["action"].popleft().numpy(), np.array([20.0], dtype=np.float32))
|
||||
np.testing.assert_array_equal(queues["action"].popleft().numpy(), np.array([30.0], dtype=np.float32))
|
||||
|
||||
def test_enqueue_predicted_actions_honors_explicit_chunk_start(self):
|
||||
queues = eval_vla._new_local_policy_queues(obs_horizon=2)
|
||||
predicted_actions = torch.tensor(
|
||||
[[[10.0], [20.0], [30.0], [40.0]]],
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
eval_vla._enqueue_predicted_actions(
|
||||
queues,
|
||||
predicted_actions=predicted_actions,
|
||||
obs_horizon=2,
|
||||
num_action_steps=2,
|
||||
action_chunk_start=0,
|
||||
)
|
||||
|
||||
self.assertEqual(len(queues["action"]), 2)
|
||||
np.testing.assert_array_equal(queues["action"].popleft().numpy(), np.array([10.0], dtype=np.float32))
|
||||
np.testing.assert_array_equal(queues["action"].popleft().numpy(), np.array([20.0], dtype=np.float32))
|
||||
|
||||
def test_remote_policy_runner_only_requests_server_inference_when_local_action_queue_is_empty(self):
|
||||
request_queue = _FakeQueue()
|
||||
response_queue = _FakeQueue(
|
||||
@@ -288,6 +324,7 @@ class EvalVLAExecutionTest(unittest.TestCase):
|
||||
camera_names=["front"],
|
||||
obs_horizon=2,
|
||||
num_action_steps=2,
|
||||
action_chunk_start=0,
|
||||
)
|
||||
first_observation = {
|
||||
"qpos": torch.tensor([1.0, 2.0], dtype=torch.float32),
|
||||
@@ -315,8 +352,46 @@ class EvalVLAExecutionTest(unittest.TestCase):
|
||||
self.assertEqual(request_queue.put_calls[0]["type"], "predict_chunk")
|
||||
self.assertEqual(request_queue.put_calls[0]["worker_index"], 3)
|
||||
self.assertEqual(request_queue.put_calls[0]["server_index"], 1)
|
||||
np.testing.assert_array_equal(first_action.numpy(), np.array([20.0], dtype=np.float32))
|
||||
np.testing.assert_array_equal(second_action.numpy(), np.array([30.0], dtype=np.float32))
|
||||
np.testing.assert_array_equal(first_action.numpy(), np.array([10.0], dtype=np.float32))
|
||||
np.testing.assert_array_equal(second_action.numpy(), np.array([20.0], dtype=np.float32))
|
||||
|
||||
def test_remote_eval_worker_passes_agent_action_chunk_start_to_remote_runner(self):
|
||||
cfg = OmegaConf.create(
|
||||
{
|
||||
"agent": {
|
||||
"obs_horizon": 2,
|
||||
"num_action_steps": 2,
|
||||
"action_chunk_start": 0,
|
||||
"camera_names": ["front"],
|
||||
},
|
||||
"eval": {
|
||||
"obs_horizon": 2,
|
||||
"num_queries": 2,
|
||||
"response_timeout_s": 3.0,
|
||||
"camera_names": ["front"],
|
||||
},
|
||||
}
|
||||
)
|
||||
captured = {}
|
||||
|
||||
class CapturingRunner:
|
||||
def __init__(self, **kwargs):
|
||||
captured.update(kwargs)
|
||||
|
||||
with mock.patch.object(eval_vla, "_RemotePolicyRunner", CapturingRunner), \
|
||||
mock.patch.object(eval_vla, "_run_eval_episode_plans", return_value={"ok": True}) as run_plans:
|
||||
result = eval_vla._run_remote_eval_worker(
|
||||
cfg,
|
||||
episode_plans=[{"episode_index": 0}],
|
||||
worker_index=1,
|
||||
server_index=2,
|
||||
request_queue=_FakeQueue(),
|
||||
response_queue=_FakeQueue(),
|
||||
)
|
||||
|
||||
self.assertEqual(result, {"ok": True})
|
||||
self.assertEqual(captured["action_chunk_start"], 0)
|
||||
run_plans.assert_called_once()
|
||||
|
||||
def test_merge_worker_summaries_sorts_episodes_and_recomputes_aggregates(self):
|
||||
worker_summaries = [
|
||||
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -394,6 +394,23 @@ class TrainVLARolloutValidationTest(unittest.TestCase):
|
||||
self.assertTrue(rollout_cfg.eval.save_summary_json)
|
||||
self.assertTrue(rollout_cfg.eval.save_trajectory_image)
|
||||
|
||||
def test_resolve_dataset_image_resize_shape_prefers_agent_top_level_override(self):
|
||||
cfg = OmegaConf.create(
|
||||
{
|
||||
'agent': {
|
||||
'dataset_image_resize_shape': None,
|
||||
'vision_backbone': {
|
||||
'dataset_image_resize_shape': [256, 256],
|
||||
},
|
||||
},
|
||||
'data': {
|
||||
'image_resize_shape': [224, 224],
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
self.assertIsNone(train_vla._resolve_dataset_image_resize_shape(cfg))
|
||||
|
||||
def test_training_passes_backbone_image_resize_override_to_dataset_instantiation(self):
|
||||
cfg = OmegaConf.create(
|
||||
{
|
||||
|
||||
Reference in New Issue
Block a user