JiajunLI/roboimi
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(vla): add SmolVLA conditioned agent with prefix encoder and train/eval validation

Browse Source 
Introduce SmolVLA-style VLM prefix encoder backbone and IMF-AttnRes
conditioned agent. Add episode-level train/val split, action MSE
validation in training, and headless eval support.
This commit is contained in:
Logic
2026-05-23 22:35:47 +08:00
parent acbd7c605a
commit d94eb8f70b
15 changed files with 2069 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_eval_vla_execution.py
+84
View File
@@ -129,6 +129,38 @@ class EvalVLAExecutionTest(unittest.TestCase):
["r_vis", "top", "front"],
)
def test_resolve_eval_image_resize_shape_prefers_condition_encoder_override(self):
cfg = OmegaConf.create(
{
"agent": {
"condition_encoder": {
"eval_image_resize_shape": None,
},
},
"data": {
"image_resize_shape": [224, 224],
},
}
)
self.assertIsNone(eval_vla._resolve_eval_image_resize_shape(cfg))
def test_resolve_eval_image_resize_shape_prefers_vision_backbone_override(self):
cfg = OmegaConf.create(
{
"agent": {
"vision_backbone": {
"eval_image_resize_shape": [256, 256],
},
},
"data": {
"image_resize_shape": [224, 224],
},
}
)
self.assertEqual(eval_vla._resolve_eval_image_resize_shape(cfg), (256, 256))
def test_build_episode_plans_without_box_poses_keeps_serial_sampling_lazy(self):
plans = eval_vla._build_episode_plans(num_episodes=3)
@@ -168,6 +200,58 @@ class EvalVLAExecutionTest(unittest.TestCase):
np.array([[[[1.0]]], [[[1.0]]], [[[1.0]]]], dtype=np.float32),
)
def test_prepare_local_policy_batch_keeps_latest_variable_task_when_present(self):
queues = eval_vla._new_local_policy_queues(obs_horizon=2)
first_observation = {
"qpos": torch.tensor([1.0, 2.0], dtype=torch.float32),
"images": {"front": torch.tensor([[[1.0]]], dtype=torch.float32)},
"task": "pick the red cube",
}
second_observation = {
"qpos": torch.tensor([3.0, 4.0], dtype=torch.float32),
"images": {"front": torch.tensor([[[2.0]]], dtype=torch.float32)},
"task": "insert the peg into the socket",
}
eval_vla._populate_local_policy_queues(queues, first_observation)
eval_vla._populate_local_policy_queues(queues, second_observation)
batch = eval_vla._prepare_local_policy_batch(
queues,
obs_horizon=2,
camera_names=["front"],
)
self.assertEqual(batch["task"], ["insert the peg into the socket"])
def test_prepare_local_policy_batch_omits_task_for_legacy_observations(self):
queues = eval_vla._new_local_policy_queues(obs_horizon=2)
observation = {
"qpos": torch.tensor([1.0, 2.0], dtype=torch.float32),
"images": {"front": torch.tensor([[[1.0]]], dtype=torch.float32)},
}
eval_vla._populate_local_policy_queues(queues, observation)
batch = eval_vla._prepare_local_policy_batch(
queues,
obs_horizon=2,
camera_names=["front"],
)
self.assertNotIn("task", batch)
def test_serialize_deserialize_policy_batch_preserves_task(self):
batch = {
"qpos": torch.zeros(1, 2, 2, dtype=torch.float32),
"images": {"front": torch.zeros(1, 2, 1, 1, 1, dtype=torch.float32)},
"task": ["pick the red cube", "insert the peg into the socket"],
}
serialized = eval_vla._serialize_policy_batch(batch)
deserialized = eval_vla._deserialize_policy_batch(serialized, device="cpu")
self.assertEqual(serialized["task"], batch["task"])
self.assertEqual(deserialized["task"], batch["task"])
def test_enqueue_predicted_actions_uses_executable_slice(self):
queues = eval_vla._new_local_policy_queues(obs_horizon=2)
predicted_actions = torch.tensor(
tests/test_eval_vla_headless.py
+89
View File
@@ -15,6 +15,7 @@ class _FakeAgent:
def __init__(self):
self.reset_calls = 0
self.last_observation = None
self.observation_shapes = []
def eval(self):
return self
@@ -27,6 +28,7 @@ class _FakeAgent:
def select_action(self, observation):
self.last_observation = observation
self.observation_shapes.append(tuple(observation["images"]["front"].shape))
return torch.zeros(16)
@@ -108,6 +110,23 @@ class EvalVLAHeadlessTest(unittest.TestCase):
self.assertEqual(tuple(prepared["images"]["front"].shape), (3, 8, 8))
self.assertEqual(tuple(prepared["qpos"].shape), (16,))
def test_prepare_observation_preserves_task_when_present(self):
obs = {
"images": {
"front": np.zeros((4, 4, 3), dtype=np.uint8),
},
"qpos": np.zeros(16, dtype=np.float32),
"task": "insert the peg into the socket",
}
prepared = eval_vla.prepare_observation(
obs,
["front"],
image_resize_shape=None,
)
self.assertEqual(prepared["task"], "insert the peg into the socket")
def test_headless_eval_sets_mujoco_gl_to_egl_when_display_missing(self):
cfg = OmegaConf.create({"eval": {"headless": True}})
with mock.patch.dict(eval_vla.os.environ, {}, clear=True):
@@ -125,6 +144,8 @@ class EvalVLAHeadlessTest(unittest.TestCase):
self.assertIn("headless", eval_cfg)
self.assertFalse(eval_cfg.headless)
self.assertIn("task_description", eval_cfg)
self.assertIsNone(eval_cfg.task_description)
def test_make_sim_env_accepts_headless_and_disables_render(self):
fake_env = object()
@@ -291,6 +312,74 @@ class EvalVLAHeadlessTest(unittest.TestCase):
self.assertIsNotNone(fake_agent.last_observation)
self.assertIn("front", fake_agent.last_observation["images"])
def test_eval_main_uses_condition_encoder_resize_override(self):
fake_env = _FakeEnv()
fake_agent = _FakeAgent()
cfg = OmegaConf.create(
{
"agent": {
"condition_encoder": {
"eval_image_resize_shape": None,
},
},
"data": {
"image_resize_shape": [224, 224],
},
"eval": {
"ckpt_path": "checkpoints/vla_model_best.pt",
"num_episodes": 1,
"max_timesteps": 1,
"device": "cpu",
"task_name": "sim_transfer",
"camera_names": ["front"],
"use_smoothing": False,
"smooth_alpha": 0.3,
"verbose_action": False,
"headless": True,
},
}
)
with mock.patch.object(eval_vla, "load_checkpoint", return_value=(fake_agent, None)), \
mock.patch.object(eval_vla, "make_sim_env", return_value=fake_env), \
mock.patch.object(eval_vla, "sample_transfer_pose", return_value=np.array([0.1, 0.2, 0.3])), \
mock.patch.object(eval_vla, "execute_policy_action"), \
mock.patch.object(eval_vla, "tqdm", side_effect=lambda iterable, **kwargs: iterable):
eval_vla.main.__wrapped__(cfg)
self.assertEqual(fake_agent.observation_shapes, [(3, 8, 8)])
def test_eval_main_injects_configured_task_description_when_env_omits_task(self):
fake_env = _FakeEnv()
fake_agent = _FakeAgent()
cfg = OmegaConf.create(
{
"agent": {},
"eval": {
"ckpt_path": "checkpoints/vla_model_best.pt",
"num_episodes": 1,
"max_timesteps": 1,
"device": "cpu",
"task_name": "sim_transfer",
"task_description": "pick the red cube",
"camera_names": ["front"],
"use_smoothing": False,
"smooth_alpha": 0.3,
"verbose_action": False,
"headless": True,
},
}
)
with mock.patch.object(eval_vla, "load_checkpoint", return_value=(fake_agent, None)), \
mock.patch.object(eval_vla, "make_sim_env", return_value=fake_env), \
mock.patch.object(eval_vla, "sample_transfer_pose", return_value=np.array([0.1, 0.2, 0.3])), \
mock.patch.object(eval_vla, "execute_policy_action"), \
mock.patch.object(eval_vla, "tqdm", side_effect=lambda iterable, **kwargs: iterable):
eval_vla.main.__wrapped__(cfg)
self.assertEqual(fake_agent.last_observation["task"], "pick the red cube")
def test_run_eval_returns_average_reward_summary(self):
reward_sequences = [
[1.0, 2.0],
tests/test_simple_robot_dataset_episode_filter.py
+42
View File
@@ -0,0 +1,42 @@
import tempfile
import unittest
from pathlib import Path
import h5py
import numpy as np
def _write_episode(path: Path, length: int = 2):
with h5py.File(path, 'w') as f:
f.create_dataset('action', data=np.zeros((length, 16), dtype=np.float32))
obs = f.create_group('observations')
obs.create_dataset('qpos', data=np.zeros((length, 16), dtype=np.float32))
images = obs.create_group('images')
for cam_name in ('l_vis', 'r_vis', 'front'):
images.create_dataset(cam_name, data=np.zeros((length, 4, 4, 3), dtype=np.uint8))
class SimpleRobotDatasetEpisodeFilterTest(unittest.TestCase):
def test_filters_by_original_episode_indices_and_exposes_available_indices(self):
from roboimi.vla.data.simpe_robot_dataset import SimpleRobotDataset
with tempfile.TemporaryDirectory() as tmpdir:
root = Path(tmpdir)
_write_episode(root / 'episode_0.hdf5')
_write_episode(root / 'episode_2.hdf5')
_write_episode(root / 'episode_10.hdf5')
dataset = SimpleRobotDataset(
root,
camera_names=['l_vis', 'r_vis', 'front'],
image_resize_shape=None,
episode_indices=[10, 2],
)
self.assertEqual(dataset.available_episode_indices, [2, 10])
self.assertEqual(set(dataset.episodes.keys()), {2, 10})
self.assertEqual(len(dataset), 4)
if __name__ == '__main__':
unittest.main()
tests/test_smolvla_imf_agent.py
+311
View File
@@ -0,0 +1,311 @@
import contextlib
import importlib
import importlib.machinery
import sys
import types
import unittest
from pathlib import Path
from unittest import mock
import torch
from hydra import compose, initialize_config_dir
from hydra.core.global_hydra import GlobalHydra
from hydra.utils import instantiate
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()
class _RecordingHead(nn.Module):
def __init__(self):
super().__init__()
self.scale = nn.Parameter(torch.tensor(0.5))
self.calls = []
@staticmethod
def _broadcast(value, reference):
while value.ndim < reference.ndim:
value = value.unsqueeze(-1)
return value
def forward(self, sample, r, t, cond=None):
self.calls.append({
'sample': sample.detach().clone(),
'r': r.detach().clone(),
't': t.detach().clone(),
'cond': None if cond is None else cond.detach().clone(),
})
cond_term = 0.0 if cond is None else cond.mean(dim=(1, 2), keepdim=True)
return self.scale * sample + self._broadcast(r, sample) + 2.0 * self._broadcast(t, sample) + cond_term
class _TaskAwareConditionEncoder(nn.Module):
output_dim = 4
condition_sequence_length = 3
tokens_per_step = 3
joint_output_dim = 4
camera_names = _CAMERA_NAMES
num_cameras = 3
def __init__(self, **kwargs):
super().__init__()
self.constructor_kwargs = dict(kwargs)
self.bias = nn.Parameter(torch.tensor(0.0))
self.calls = []
def forward(self, images, state, task=None):
self.calls.append({'task': task, 'state': state.detach().clone(), 'image_keys': tuple(images.keys())})
batch_size = state.shape[0]
state_last = state[:, -1, 0]
if isinstance(task, str):
task_lengths = torch.full((batch_size,), float(len(task)), dtype=state.dtype, device=state.device)
else:
task_lengths = torch.tensor([float(len(item)) for item in task], dtype=state.dtype, device=state.device)
image_marker = images['r_vis'][:, -1].mean(dim=(1, 2, 3))
token0 = torch.stack([state_last, task_lengths, image_marker, torch.ones_like(state_last)], dim=-1)
token1 = token0 + 1.0
token2 = token0 + 2.0
return torch.stack([token0, token1, token2], dim=1) + self.bias
class _BF16TaskAwareConditionEncoder(_TaskAwareConditionEncoder):
def forward(self, images, state, task=None):
return super().forward(images, state, task=task).to(dtype=torch.bfloat16)
class _StubIMFHead(nn.Module):
def __init__(self, input_dim, output_dim, horizon, n_obs_steps, cond_dim, **kwargs):
super().__init__()
self.constructor_kwargs = {
'input_dim': input_dim,
'output_dim': output_dim,
'horizon': horizon,
'n_obs_steps': n_obs_steps,
'cond_dim': cond_dim,
**kwargs,
}
self.proj = nn.Linear(input_dim, output_dim)
self.cond_obs_emb = nn.Linear(cond_dim, max(cond_dim, 1))
def forward(self, sample, r, t, cond=None):
return torch.zeros_like(sample)
def get_optim_groups(self, weight_decay):
return [
{'params': [self.proj.weight], 'weight_decay': weight_decay},
{'params': [self.proj.bias, self.cond_obs_emb.weight, self.cond_obs_emb.bias], 'weight_decay': 0.0},
]
@contextlib.contextmanager
def _stub_optional_modules(include_head=False, include_condition_encoder=False):
previous = {}
def inject(name, module):
if name not in previous:
previous[name] = sys.modules.get(name, _MISSING)
sys.modules[name] = module
diffusers_module = types.ModuleType('diffusers')
schedulers_module = types.ModuleType('diffusers.schedulers')
ddpm_module = types.ModuleType('diffusers.schedulers.scheduling_ddpm')
ddim_module = types.ModuleType('diffusers.schedulers.scheduling_ddim')
class _FakeScheduler:
def __init__(self, num_train_timesteps=100, **kwargs):
self.config = types.SimpleNamespace(num_train_timesteps=num_train_timesteps)
ddpm_module.DDPMScheduler = _FakeScheduler
ddim_module.DDIMScheduler = _FakeScheduler
diffusers_module.DDPMScheduler = _FakeScheduler
diffusers_module.DDIMScheduler = _FakeScheduler
diffusers_module.schedulers = schedulers_module
try:
inject('diffusers', diffusers_module)
inject('diffusers.schedulers', schedulers_module)
inject('diffusers.schedulers.scheduling_ddpm', ddpm_module)
inject('diffusers.schedulers.scheduling_ddim', ddim_module)
if include_head:
import roboimi.vla.models.heads as heads_package
head_module = types.ModuleType('roboimi.vla.models.heads.imf_transformer1d')
head_module.IMFTransformer1D = _StubIMFHead
inject('roboimi.vla.models.heads.imf_transformer1d', head_module)
setattr(heads_package, 'imf_transformer1d', head_module)
if include_condition_encoder:
module = types.ModuleType('tests.fake_smolvla_condition_encoder')
module.TaskAwareConditionEncoder = _TaskAwareConditionEncoder
module.BF16TaskAwareConditionEncoder = _BF16TaskAwareConditionEncoder
inject('tests.fake_smolvla_condition_encoder', module)
yield
finally:
for name, old in reversed(list(previous.items())):
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 SmolVLAIMFAgentTest(unittest.TestCase):
def test_compute_loss_and_predict_action_pass_variable_task_to_condition_encoder(self):
from roboimi.vla.agent_smolvla_conditioned import SmolVLAIMFAttnResAgent
condition_encoder = _BF16TaskAwareConditionEncoder()
head = _RecordingHead()
agent = SmolVLAIMFAttnResAgent(
condition_encoder=condition_encoder,
action_encoder=nn.Identity(),
head=head,
action_dim=2,
obs_dim=1,
pred_horizon=3,
obs_horizon=2,
diffusion_steps=10,
inference_steps=1,
num_cams=3,
camera_names=_CAMERA_NAMES,
num_action_steps=2,
head_type='transformer',
)
images = {
'r_vis': torch.full((2, 2, 1, 2, 2), 1.0),
'top': torch.full((2, 2, 1, 2, 2), 2.0),
'front': torch.full((2, 2, 1, 2, 2), 3.0),
}
qpos = torch.tensor([[[1.0], [2.0]], [[3.0], [4.0]]])
actions = torch.zeros(2, 3, 2)
tasks = ['short', 'longer task']
loss = agent.compute_loss({'images': images, 'qpos': qpos, 'action': actions, 'task': tasks})
self.assertTrue(torch.isfinite(loss))
self.assertEqual(condition_encoder.calls[-1]['task'], tasks)
self.assertEqual(head.calls[-1]['cond'].shape, (2, 3, 4))
self.assertTrue(torch.allclose(head.calls[-1]['cond'][:, 0, 1], torch.tensor([5.0, 11.0])))
with mock.patch('roboimi.vla.agent_imf.torch.randn', return_value=torch.zeros(2, 3, 2)):
pred = agent.predict_action(images, qpos, task=tasks)
self.assertEqual(pred.shape, (2, 3, 2))
self.assertEqual(condition_encoder.calls[-1]['task'], tasks)
def test_condition_tokens_are_cast_to_action_head_dtype(self):
from roboimi.vla.agent_smolvla_conditioned import SmolVLAIMFAttnResAgent
condition_encoder = _TaskAwareConditionEncoder()
head = _RecordingHead()
agent = SmolVLAIMFAttnResAgent(
condition_encoder=condition_encoder,
action_encoder=nn.Identity(),
head=head,
action_dim=2,
obs_dim=1,
pred_horizon=3,
obs_horizon=2,
diffusion_steps=10,
inference_steps=1,
num_cams=3,
camera_names=_CAMERA_NAMES,
num_action_steps=2,
head_type='transformer',
)
images = {
'r_vis': torch.full((1, 2, 1, 2, 2), 1.0),
'top': torch.full((1, 2, 1, 2, 2), 2.0),
'front': torch.full((1, 2, 1, 2, 2), 3.0),
}
qpos = torch.tensor([[[1.0], [2.0]]], dtype=torch.float32)
cond = agent._build_cond(images, qpos, task=['pick'])
self.assertEqual(cond.dtype, head.scale.dtype)
def test_unknown_dataset_task_uses_configured_task_description(self):
from roboimi.vla.agent_smolvla_conditioned import SmolVLAIMFAttnResAgent
condition_encoder = _TaskAwareConditionEncoder()
head = _RecordingHead()
agent = SmolVLAIMFAttnResAgent(
condition_encoder=condition_encoder,
action_encoder=nn.Identity(),
head=head,
action_dim=2,
obs_dim=1,
pred_horizon=3,
obs_horizon=2,
diffusion_steps=10,
inference_steps=1,
num_cams=3,
camera_names=_CAMERA_NAMES,
num_action_steps=2,
head_type='transformer',
task_description='insert the peg into the socket',
)
images = {
'r_vis': torch.full((2, 2, 1, 2, 2), 1.0),
'top': torch.full((2, 2, 1, 2, 2), 2.0),
'front': torch.full((2, 2, 1, 2, 2), 3.0),
}
qpos = torch.tensor([[[1.0], [2.0]], [[3.0], [4.0]]], dtype=torch.float32)
agent._build_cond(images, qpos, task=['unknown', ''])
self.assertEqual(
condition_encoder.calls[-1]['task'],
['insert the peg into the socket', 'insert the peg into the socket'],
)
def test_hydra_config_instantiates_smolvla_imf_attnres_with_condition_encoder_contract(self):
cfg = _compose_cfg(overrides=[
'agent=smolvla_imf_attnres',
'agent.condition_encoder._target_=tests.fake_smolvla_condition_encoder.TaskAwareConditionEncoder',
'agent.condition_dim=4',
'agent.condition_sequence_length=3',
'agent.head.n_layer=1',
'agent.head.n_emb=16',
])
self.assertEqual(cfg.agent._target_, 'roboimi.vla.agent_smolvla_conditioned.SmolVLAIMFAttnResAgent')
self.assertEqual(cfg.agent.head.cond_dim, cfg.agent.condition_dim)
self.assertEqual(cfg.agent.head.n_obs_steps, cfg.agent.condition_sequence_length)
with _stub_optional_modules(include_head=True, include_condition_encoder=True):
agent = instantiate(cfg.agent)
self.assertEqual(agent.per_step_cond_dim, 4)
self.assertEqual(agent.condition_sequence_length, 3)
self.assertIsInstance(agent.noise_pred_net, _StubIMFHead)
self.assertEqual(agent.noise_pred_net.constructor_kwargs['cond_dim'], 4)
self.assertEqual(agent.noise_pred_net.constructor_kwargs['n_obs_steps'], 3)
def test_hydra_config_exposes_smolvla_pretrained_vlm_defaults(self):
cfg = _compose_cfg(overrides=[
'agent=smolvla_imf_attnres',
])
self.assertEqual(cfg.agent.condition_encoder.model_name, 'HuggingFaceTB/SmolVLM2-500M-Video-Instruct')
self.assertTrue(cfg.agent.condition_encoder.load_vlm_weights)
self.assertEqual(cfg.agent.condition_encoder.num_vlm_layers, 16)
self.assertTrue(cfg.agent.condition_encoder.freeze_vlm)
self.assertTrue(cfg.agent.condition_encoder.freeze_vision_encoder)
self.assertTrue(cfg.agent.condition_encoder.run_text_model)
self.assertEqual(cfg.agent.condition_encoder.max_state_dim, 32)
self.assertIsNone(cfg.agent.condition_encoder.dataset_image_resize_shape)
self.assertIsNone(cfg.agent.condition_encoder.eval_image_resize_shape)
self.assertEqual(cfg.agent.condition_dim, 960)
self.assertEqual(cfg.agent.condition_sequence_length, 241)
self.assertEqual(cfg.agent.head.cond_dim, 960)
self.assertEqual(cfg.agent.head.n_obs_steps, 241)
if __name__ == '__main__':
unittest.main()
tests/test_smolvla_prefix_encoder.py
+328
View File
@@ -0,0 +1,328 @@
import types
import unittest
from unittest import mock
import torch
from torch import nn
class _FakeVisionOutput:
def __init__(self, last_hidden_state):
self.last_hidden_state = last_hidden_state
class _FakeVisionModel(nn.Module):
def __init__(self, hidden_size=4):
super().__init__()
self.dtype = torch.float32
self.scale = nn.Parameter(torch.tensor(1.0))
self.calls = []
self.hidden_size = hidden_size
def forward(self, pixel_values=None, patch_attention_mask=None):
self.calls.append({
'pixel_values': pixel_values.detach().clone(),
'patch_attention_mask': patch_attention_mask,
})
pooled = pixel_values.mean(dim=(2, 3)) * self.scale
tokens = torch.stack([pooled, pooled + 1.0], dim=1)
return _FakeVisionOutput(tokens)
class _FakeConnector(nn.Module):
def __init__(self, in_dim=3, out_dim=4):
super().__init__()
self.proj = nn.Linear(in_dim, out_dim, bias=False)
with torch.no_grad():
self.proj.weight.copy_(
torch.tensor(
[
[1.0, 0.0, 0.0],
[0.0, 1.0, 0.0],
[0.0, 0.0, 1.0],
[1.0, 1.0, 1.0],
]
)
)
def forward(self, x):
return self.proj(x)
class _FakeTextModel(nn.Module):
def __init__(self, vocab_size=32, hidden_size=4, num_layers=6):
super().__init__()
self.embed = nn.Embedding(vocab_size, hidden_size)
self.layers = nn.ModuleList([nn.Linear(hidden_size, hidden_size) for _ in range(num_layers)])
self.norm = nn.Identity()
self.forward_calls = []
with torch.no_grad():
for idx in range(vocab_size):
self.embed.weight[idx].fill_(float(idx))
def get_input_embeddings(self):
return self.embed
def forward(
self,
input_ids=None,
attention_mask=None,
position_ids=None,
past_key_values=None,
inputs_embeds=None,
use_cache=None,
return_dict=True,
cache_position=None,
**kwargs,
):
del input_ids, past_key_values, use_cache, cache_position, kwargs
self.forward_calls.append({
'attention_mask': None if attention_mask is None else attention_mask.detach().clone(),
'position_ids': None if position_ids is None else position_ids.detach().clone(),
'inputs_embeds': inputs_embeds.detach().clone(),
})
hidden = inputs_embeds
for layer in self.layers:
hidden = layer(hidden)
hidden = self.norm(hidden)
if return_dict:
return types.SimpleNamespace(last_hidden_state=hidden)
return (hidden,)
class _FakeVLM(nn.Module):
def __init__(self):
super().__init__()
text_config = types.SimpleNamespace(hidden_size=4, head_dim=2, num_attention_heads=2, num_key_value_heads=1)
self.config = types.SimpleNamespace(text_config=text_config)
self.model = types.SimpleNamespace(
vision_model=_FakeVisionModel(hidden_size=4),
connector=_FakeConnector(in_dim=3, out_dim=4),
text_model=_FakeTextModel(hidden_size=4, num_layers=6),
)
class _FakeTokenizer:
fake_image_token_id = 29
global_image_token_id = 30
def __init__(self):
self.padding_side = 'left'
self.calls = []
def __call__(self, text, *, padding, max_length, return_tensors, truncation):
self.calls.append({
'text': list(text),
'padding': padding,
'max_length': max_length,
'return_tensors': return_tensors,
'truncation': truncation,
'padding_side': self.padding_side,
})
batch = len(text)
ids = torch.zeros(batch, max_length, dtype=torch.long)
mask = torch.zeros(batch, max_length, dtype=torch.bool)
for row, item in enumerate(text):
del item
ids[row, :3] = torch.tensor([1, 2, 3])
mask[row, :3] = True
return {'input_ids': ids, 'attention_mask': mask}
class SmolVLAPrefixEncoderTest(unittest.TestCase):
def test_loads_pretrained_vlm_with_local_files_only_crops_layers_and_freezes_vlm(self):
from roboimi.vla.models.backbones.smolvla_prefix_encoder import SmolVLAPrefixEncoder
fake_vlm = _FakeVLM()
fake_tokenizer = _FakeTokenizer()
with mock.patch(
'roboimi.vla.models.backbones.smolvla_prefix_encoder.AutoModelForImageTextToText.from_pretrained',
return_value=fake_vlm,
) as model_loader, mock.patch(
'roboimi.vla.models.backbones.smolvla_prefix_encoder.AutoTokenizer.from_pretrained',
return_value=fake_tokenizer,
) as tokenizer_loader:
encoder = SmolVLAPrefixEncoder(
model_name='HuggingFaceTB/SmolVLM2-500M-Video-Instruct',
local_files_only=True,
num_vlm_layers=2,
freeze_vlm=True,
max_state_dim=5,
tokenizer_max_length=4,
camera_names=('r_vis', 'top'),
resize_imgs_with_padding=None,
)
model_loader.assert_called_once()
self.assertEqual(model_loader.call_args.kwargs['local_files_only'], True)
self.assertEqual(model_loader.call_args.kwargs['torch_dtype'], 'bfloat16')
tokenizer_loader.assert_called_once_with(
'HuggingFaceTB/SmolVLM2-500M-Video-Instruct',
local_files_only=True,
)
self.assertEqual(len(encoder.vlm.model.text_model.layers), 2)
self.assertTrue(all(not p.requires_grad for p in encoder.vlm.parameters()))
self.assertFalse(encoder.vlm.training)
encoder.train()
self.assertFalse(encoder.vlm.training)
def test_accepts_train_eval_resize_compatibility_fields(self):
from roboimi.vla.models.backbones.smolvla_prefix_encoder import SmolVLAPrefixEncoder
encoder = SmolVLAPrefixEncoder(
vlm=_FakeVLM(),
tokenizer=_FakeTokenizer(),
num_vlm_layers=3,
camera_names=('r_vis', 'top'),
resize_imgs_with_padding=None,
dataset_image_resize_shape=None,
eval_image_resize_shape=(640, 480),
)
self.assertIsNone(encoder.dataset_image_resize_shape)
self.assertEqual(encoder.eval_image_resize_shape, (640, 480))
def test_embed_prefix_uses_variable_tasks_camera_order_last_state_and_smolvla_masks(self):
from roboimi.vla.models.backbones.smolvla_prefix_encoder import SmolVLAPrefixEncoder
fake_vlm = _FakeVLM()
fake_tokenizer = _FakeTokenizer()
encoder = SmolVLAPrefixEncoder(
vlm=fake_vlm,
tokenizer=fake_tokenizer,
num_vlm_layers=3,
freeze_vlm=True,
train_state_proj=True,
max_state_dim=5,
tokenizer_max_length=4,
camera_names=('r_vis', 'top'),
resize_imgs_with_padding=None,
)
with torch.no_grad():
encoder.state_proj.weight.zero_()
encoder.state_proj.bias.zero_()
encoder.state_proj.weight[:, :4] = torch.eye(4)
images = {
'top': torch.full((2, 2, 3, 2, 2), 0.75),
'r_vis': torch.full((2, 2, 3, 2, 2), 0.25),
}
state = torch.tensor(
[
[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]],
[[7.0, 8.0, 9.0], [10.0, 11.0, 12.0]],
]
)
tasks = ['pick red cube', 'open drawer']
out = encoder.embed_prefix(images=images, state=state, task=tasks)
# 2 cameras * 2 image tokens + 4 language tokens + 1 state token
self.assertEqual(out.shape, (2, 9, 4))
self.assertEqual(encoder.output_dim, 4)
self.assertEqual(encoder.tokens_per_step, 9)
self.assertEqual(encoder.condition_sequence_length, 9)
self.assertEqual(fake_tokenizer.calls[-1]['text'], ['pick red cube\n', 'open drawer\n'])
self.assertEqual(fake_tokenizer.calls[-1]['padding'], 'max_length')
self.assertEqual(fake_tokenizer.calls[-1]['max_length'], 4)
self.assertEqual(fake_tokenizer.calls[-1]['padding_side'], 'right')
# Camera order is r_vis then top, and pixels are mapped [0,1] -> [-1,1].
first_camera_pixels = fake_vlm.model.vision_model.calls[0]['pixel_values']
second_camera_pixels = fake_vlm.model.vision_model.calls[1]['pixel_values']
self.assertTrue(torch.allclose(first_camera_pixels, torch.full((2, 3, 2, 2), -0.5)))
self.assertTrue(torch.allclose(second_camera_pixels, torch.full((2, 3, 2, 2), 0.5)))
# Last token is padded last state projected from [4,5,6,0,0] and [10,11,12,0,0].
self.assertTrue(torch.allclose(out[0, -1], torch.tensor([4.0, 5.0, 6.0, 0.0])))
self.assertTrue(torch.allclose(out[1, -1], torch.tensor([10.0, 11.0, 12.0, 0.0])))
pad_mask, att_mask = encoder.last_prefix_pad_mask, encoder.last_prefix_att_mask
self.assertEqual(pad_mask.shape, (2, 9))
self.assertEqual(att_mask.shape, (2, 9))
self.assertTrue(torch.all(att_mask[:, :8] == 0))
self.assertTrue(torch.all(att_mask[:, -1] == 1))
def test_forward_can_run_cropped_frozen_text_model_over_prefix_tokens(self):
from roboimi.vla.models.backbones.smolvla_prefix_encoder import SmolVLAPrefixEncoder
fake_vlm = _FakeVLM()
fake_tokenizer = _FakeTokenizer()
encoder = SmolVLAPrefixEncoder(
vlm=fake_vlm,
tokenizer=fake_tokenizer,
num_vlm_layers=2,
freeze_vlm=True,
max_state_dim=5,
tokenizer_max_length=4,
camera_names=('r_vis',),
resize_imgs_with_padding=None,
run_text_model=True,
)
images = {
'r_vis': torch.full((2, 1, 3, 2, 2), 0.25),
}
state = torch.tensor([[[1.0, 2.0, 3.0]], [[4.0, 5.0, 6.0]]])
out = encoder(images, state=state, task=['pick', 'place'])
# 1 camera * 2 image tokens + 4 language tokens + 1 state token.
self.assertEqual(out.shape, (2, 7, 4))
self.assertEqual(len(fake_vlm.model.text_model.layers), 2)
self.assertEqual(len(fake_vlm.model.text_model.forward_calls), 1)
text_call = fake_vlm.model.text_model.forward_calls[-1]
self.assertEqual(tuple(text_call['attention_mask'].shape), (2, 1, 7, 7))
self.assertEqual(tuple(text_call['position_ids'].shape), (2, 7))
self.assertTrue(torch.all(out != 0))
def test_frozen_vlm_still_backpropagates_through_text_model_to_state_projection(self):
from roboimi.vla.models.backbones.smolvla_prefix_encoder import SmolVLAPrefixEncoder
fake_vlm = _FakeVLM()
encoder = SmolVLAPrefixEncoder(
vlm=fake_vlm,
tokenizer=_FakeTokenizer(),
num_vlm_layers=3,
freeze_vlm=True,
train_state_proj=True,
max_state_dim=5,
tokenizer_max_length=4,
camera_names=('r_vis',),
resize_imgs_with_padding=None,
run_text_model=True,
)
images = {
'r_vis': torch.full((2, 1, 3, 2, 2), 0.25),
}
state = torch.tensor([[[1.0, 2.0, 3.0]], [[4.0, 5.0, 6.0]]])
out = encoder(images, state=state, task=['pick', 'place'])
out[:, -1].sum().backward()
self.assertIsNotNone(encoder.state_proj.weight.grad)
self.assertGreater(float(encoder.state_proj.weight.grad.abs().sum()), 0.0)
self.assertTrue(all(param.grad is None for param in fake_vlm.parameters()))
def test_embed_prefix_rejects_missing_camera_and_task_batch_mismatch(self):
from roboimi.vla.models.backbones.smolvla_prefix_encoder import SmolVLAPrefixEncoder
encoder = SmolVLAPrefixEncoder(
vlm=_FakeVLM(),
tokenizer=_FakeTokenizer(),
num_vlm_layers=3,
camera_names=('r_vis', 'top'),
resize_imgs_with_padding=None,
max_state_dim=5,
tokenizer_max_length=4,
)
images = {'r_vis': torch.rand(2, 1, 3, 2, 2)}
state = torch.rand(2, 1, 3)
with self.assertRaisesRegex(ValueError, 'missing.*top'):
encoder.embed_prefix(images=images, state=state, task=['a', 'b'])
images['top'] = torch.rand(2, 1, 3, 2, 2)
with self.assertRaisesRegex(ValueError, 'task batch'):
encoder.embed_prefix(images=images, state=state, task=['only one'])
if __name__ == '__main__':
unittest.main()
tests/test_train_vla_rollout_validation.py
+227
View File
@@ -14,6 +14,8 @@ from roboimi.demos.vla_scripts import eval_vla, train_vla
class _FakeDataset:
available_episode_indices = [0, 1]
def __len__(self):
return 4
@@ -29,6 +31,13 @@ class _FakeLoader:
return iter(self._batches)
class _FakeValDataset(_FakeDataset):
available_episode_indices = [1]
def __len__(self):
return 2
class _FakeOptimizer:
def __init__(self, lr=1e-3):
self.param_groups = [{'lr': lr}]
@@ -91,6 +100,16 @@ class _FakeAgent(nn.Module):
return {}
class _CapturingAgent(_FakeAgent):
def __init__(self):
super().__init__()
self.compute_loss_inputs = []
def compute_loss(self, agent_input):
self.compute_loss_inputs.append(agent_input)
return (self.weight - torch.tensor(0.5)).pow(2)
class _SequentialLossAgent(nn.Module):
def __init__(self, losses):
super().__init__()
@@ -150,6 +169,94 @@ class _FakeEvalEnv:
class TrainVLARolloutValidationTest(unittest.TestCase):
def test_run_training_passes_variable_batch_task_to_agent_input(self):
cfg = OmegaConf.create(
{
'train': {
'device': 'cpu',
'batch_size': 2,
'num_workers': 0,
'val_split': 0.0,
'seed': 0,
'lr': 1e-3,
'max_steps': 1,
'log_freq': 100,
'save_freq': 1000,
'warmup_steps': 1,
'scheduler_type': 'constant',
'min_lr': 0.0,
'grad_clip': 1.0,
'weight_decay': 0.0,
'pretrained_ckpt': None,
'resume_ckpt': None,
'use_swanlab': False,
'rollout_val_freq_epochs': 0,
'rollout_validate_on_checkpoint': False,
'rollout_num_episodes': 1,
},
'data': {
'camera_names': ['front'],
'dataset_dir': 'unused',
},
'agent': {
'_target_': 'fake.agent',
'normalization_type': 'min_max',
},
'eval': {
'ckpt_path': 'unused.pt',
'num_episodes': 1,
'max_timesteps': 1,
'device': 'cpu',
'task_name': 'sim_transfer',
'camera_names': ['front'],
'use_smoothing': False,
'smooth_alpha': 0.3,
'verbose_action': False,
'headless': True,
},
'experiment': {},
}
)
agent = _CapturingAgent()
batch_task = ['pick the red cube', 'insert the peg into the socket']
def fake_instantiate(config_node, **_kwargs):
if config_node is cfg.data:
return _FakeDataset()
if config_node is cfg.agent:
return agent
raise AssertionError(f'unexpected instantiate config: {config_node!r}')
def fake_dataloader(_dataset, *, shuffle, **_kwargs):
del shuffle, _kwargs
return _FakeLoader(
{
'observation.front': torch.zeros(2, 2, 3, 4, 4),
'observation.state': torch.zeros(2, 2, 4),
'action': torch.zeros(2, 8, 2),
'action_is_pad': torch.zeros(2, 8, dtype=torch.bool),
'task': list(batch_task),
},
length=1,
)
with tempfile.TemporaryDirectory() as tempdir:
previous_cwd = os.getcwd()
try:
os.chdir(tempdir)
with mock.patch.object(train_vla, 'instantiate', side_effect=fake_instantiate), \
mock.patch.object(train_vla, 'DataLoader', side_effect=fake_dataloader), \
mock.patch.object(train_vla, 'build_training_optimizer', return_value=_FakeOptimizer(cfg.train.lr)), \
mock.patch.object(train_vla, 'get_lr_schedule_with_warmup', return_value=_FakeScheduler()), \
mock.patch.object(train_vla, 'tqdm', side_effect=lambda iterable, **kwargs: _FakeProgressBar(iterable)), \
mock.patch.object(train_vla.torch, 'save', return_value=None):
train_vla._run_training(cfg)
finally:
os.chdir(previous_cwd)
self.assertEqual(len(agent.compute_loss_inputs), 1)
self.assertEqual(agent.compute_loss_inputs[0]['task'], batch_task)
def test_default_train_config_uses_full_dataset_and_epoch_rollout_validation(self):
cfg = OmegaConf.load(Path('roboimi/vla/conf/config.yaml'))
@@ -162,6 +269,39 @@ class TrainVLARolloutValidationTest(unittest.TestCase):
self.assertIsNone(cfg.train.rollout_num_workers)
self.assertIsNone(cfg.train.rollout_cuda_devices)
def test_explicit_val_episode_indices_builds_held_out_dataset(self):
cfg = OmegaConf.create(
{
'train': {
'val_episode_indices': [1],
'val_split': 0.0,
'seed': 42,
},
'data': {},
}
)
instantiate_calls = []
def fake_instantiate(config_node, **kwargs):
del config_node
instantiate_calls.append(dict(kwargs))
if kwargs.get('episode_indices') == [1]:
return _FakeValDataset()
return _FakeDataset()
with mock.patch.object(train_vla, 'instantiate', side_effect=fake_instantiate):
dataset, train_dataset, val_dataset, explicit = train_vla.build_train_val_datasets(
cfg,
dataset_image_resize_shape=None,
)
self.assertIsInstance(dataset, _FakeDataset)
self.assertIsInstance(train_dataset, _FakeDataset)
self.assertIsInstance(val_dataset, _FakeValDataset)
self.assertEqual(explicit, [1])
self.assertEqual(instantiate_calls[1]['episode_indices'], [0])
self.assertEqual(instantiate_calls[2]['episode_indices'], [1])
def test_run_training_rollout_validation_propagates_gpu_parallel_settings(self):
cfg = OmegaConf.create(
@@ -340,6 +480,93 @@ class TrainVLARolloutValidationTest(unittest.TestCase):
self.assertIn('image_resize_shape', captured_dataset_kwargs)
self.assertIsNone(captured_dataset_kwargs['image_resize_shape'])
def test_training_passes_condition_encoder_image_resize_override_to_dataset_instantiation(self):
cfg = OmegaConf.create(
{
'agent': {
'condition_encoder': {
'dataset_image_resize_shape': None,
},
'normalization_type': 'min_max',
},
'data': {
'dataset_dir': 'unused',
'camera_names': ['front'],
'image_resize_shape': [224, 224],
},
'train': {
'batch_size': 2,
'lr': 1e-4,
'max_steps': 0,
'device': 'cpu',
'disable_cudnn': False,
'num_workers': 0,
'val_split': 0.0,
'seed': 42,
'log_freq': 1,
'save_freq': 10,
'use_swanlab': False,
'rollout_val_freq_epochs': 0,
'rollout_validate_on_checkpoint': False,
'rollout_num_episodes': 1,
'warmup_steps': 1,
'scheduler_type': 'constant',
'min_lr': 1e-6,
'weight_decay': 1e-5,
'grad_clip': 1.0,
'pretrained_ckpt': None,
},
'eval': {
'ckpt_path': 'unused.pt',
'num_episodes': 1,
'headless': True,
'device': 'cpu',
'verbose_action': False,
},
'experiment': {},
}
)
captured_dataset_kwargs = {}
def fake_instantiate(config_node, **kwargs):
if config_node is cfg.data:
captured_dataset_kwargs.update(kwargs)
return _FakeDataset()
if config_node is cfg.agent:
return _FakeAgent()
raise AssertionError(f'unexpected instantiate config: {config_node!r}')
def fake_dataloader(_dataset, *, shuffle, **_kwargs):
del shuffle, _kwargs
return _FakeLoader(
{
'observation.front': torch.zeros(1, 3, 2, 2),
'observation.state': torch.zeros(1, 4),
'action': torch.zeros(1, 2),
'action_is_pad': torch.zeros(1, 1, dtype=torch.bool),
},
length=1,
)
with tempfile.TemporaryDirectory() as tempdir:
previous_cwd = os.getcwd()
try:
os.chdir(tempdir)
with mock.patch.object(train_vla, 'instantiate', side_effect=fake_instantiate), \
mock.patch.object(train_vla, 'DataLoader', side_effect=fake_dataloader), \
mock.patch.object(train_vla, 'build_training_optimizer', return_value=_FakeOptimizer(cfg.train.lr)), \
mock.patch.object(train_vla, 'get_lr_schedule_with_warmup', return_value=_FakeScheduler()), \
mock.patch.object(train_vla, 'tqdm', side_effect=lambda iterable, **kwargs: _FakeProgressBar(iterable)), \
mock.patch.object(train_vla, '_init_swanlab', return_value=None), \
mock.patch.object(train_vla, '_finish_swanlab', return_value=None), \
mock.patch.object(train_vla.torch, 'save', return_value=None):
train_vla._run_training(cfg)
finally:
os.chdir(previous_cwd)
self.assertIn('image_resize_shape', captured_dataset_kwargs)
self.assertIsNone(captured_dataset_kwargs['image_resize_shape'])
def test_eval_main_delegates_to_plain_run_eval_helper(self):
cfg = OmegaConf.create(
{
tests/test_train_vla_transformer_optimizer.py
+42
View File
@@ -422,3 +422,45 @@ class TrainVLATransformerOptimizerTest(unittest.TestCase):
if __name__ == '__main__':
unittest.main()
class TrainVLASmolVLAOptimizerTest(unittest.TestCase):
def test_build_training_optimizer_excludes_frozen_vlm_parameters_and_keeps_state_proj_and_head(self):
module = TrainVLATransformerOptimizerTest()._load_train_vla_module()
class _Head(nn.Module):
def __init__(self):
super().__init__()
self.proj = nn.Linear(2, 2)
def get_optim_groups(self, weight_decay):
return [{'params': list(self.parameters()), 'weight_decay': weight_decay}]
class _ConditionEncoder(nn.Module):
def __init__(self):
super().__init__()
self.vlm = nn.Linear(2, 2)
for param in self.vlm.parameters():
param.requires_grad = False
self.state_proj = nn.Linear(2, 2)
class _Agent(nn.Module):
def __init__(self):
super().__init__()
self.noise_pred_net = _Head()
self.condition_encoder = _ConditionEncoder()
agent = _Agent()
with mock.patch.object(module, 'AdamW', RecordingAdamW):
optimizer = module.build_training_optimizer(agent, lr=1e-4, weight_decay=0.01)
names_by_param_id = {id(param): name for name, param in agent.named_parameters()}
optimizer_names = {
names_by_param_id[id(param)]
for group in optimizer.param_groups
for param in group['params']
}
self.assertIn('condition_encoder.state_proj.weight', optimizer_names)
self.assertIn('condition_encoder.state_proj.bias', optimizer_names)
self.assertIn('noise_pred_net.proj.weight', optimizer_names)
self.assertNotIn('condition_encoder.vlm.weight', optimizer_names)
self.assertNotIn('condition_encoder.vlm.bias', optimizer_names)