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

feat(vla): add SmolVLA conditioning and experiment artifacts

Browse Source
This commit is contained in:
Logic
2026-07-31 10:11:04 +08:00
parent acbd7c605a
commit 5ae9f5fa48
175 changed files with 317471 additions and 70 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_train_vla_rollout_validation.py
+183
View File
@@ -18,6 +18,17 @@ class _FakeDataset:
return 4
class _FakeEpisodeDataset(_FakeDataset):
def __init__(self, episode_indices=None):
self.available_episode_indices = [100, 101]
self.episode_indices = None if episode_indices is None else list(episode_indices)
def __len__(self):
if self.episode_indices is None:
return 8
return 4 * len(self.episode_indices)
class _FakeLoader:
def __init__(self, batch, length=1):
self._batches = [batch] * length
@@ -90,6 +101,9 @@ class _FakeAgent(nn.Module):
def get_normalization_stats(self):
return {}
def predict_action_chunk(self, agent_input):
return agent_input['action']
class _SequentialLossAgent(nn.Module):
def __init__(self, losses):
@@ -340,6 +354,175 @@ class TrainVLARolloutValidationTest(unittest.TestCase):
self.assertIn('image_resize_shape', captured_dataset_kwargs)
self.assertIsNone(captured_dataset_kwargs['image_resize_shape'])
def test_explicit_val_episode_indices_split_train_and_val_datasets(self):
cfg = OmegaConf.create(
{
'agent': {
'normalization_type': 'min_max',
},
'data': {
'dataset_dir': 'unused',
'camera_names': ['front'],
},
'train': {
'batch_size': 2,
'lr': 1e-4,
'max_steps': 0,
'device': 'cpu',
'disable_cudnn': False,
'num_workers': 0,
'val_split': 0.0,
'val_episode_indices': [100],
'action_mse_val_freq_epochs': 1,
'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,
'resume_ckpt': None,
},
'eval': {
'ckpt_path': 'unused.pt',
'num_episodes': 1,
'headless': True,
'device': 'cpu',
'verbose_action': False,
},
'experiment': {},
}
)
instantiate_kwargs = []
def fake_instantiate(config_node, **kwargs):
if config_node is cfg.data:
instantiate_kwargs.append(dict(kwargs))
return _FakeEpisodeDataset(kwargs.get('episode_indices'))
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)
episode_indices_calls = [kwargs.get('episode_indices') for kwargs in instantiate_kwargs]
self.assertEqual(episode_indices_calls, [None, [101], [100]])
def test_action_mse_requires_explicit_val_episode_indices(self):
cfg = OmegaConf.create(
{
'agent': {
'normalization_type': 'min_max',
},
'data': {
'dataset_dir': 'unused',
'camera_names': ['front'],
},
'train': {
'batch_size': 2,
'lr': 1e-4,
'max_steps': 0,
'device': 'cpu',
'disable_cudnn': False,
'num_workers': 0,
'val_split': 0.25,
'val_episode_indices': None,
'action_mse_val_freq_epochs': 1,
'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,
'resume_ckpt': None,
},
'eval': {
'ckpt_path': 'unused.pt',
'num_episodes': 1,
'headless': True,
'device': 'cpu',
'verbose_action': False,
},
'experiment': {},
}
)
def fake_instantiate(config_node, **_kwargs):
if config_node is cfg.data:
return _FakeDataset()
if config_node is cfg.agent:
return _FakeAgent()
raise AssertionError(f'unexpected instantiate config: {config_node!r}')
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, '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, '_init_swanlab', return_value=None), \
mock.patch.object(train_vla, '_finish_swanlab', return_value=None):
with self.assertRaisesRegex(ValueError, 'val_episode_indices'):
train_vla._run_training(cfg)
finally:
os.chdir(previous_cwd)
def test_compute_action_mse_validation_masks_padding(self):
agent = _FakeAgent()
val_loader = _FakeLoader(
{
'observation.front': torch.zeros(1, 2, 3, 2, 2),
'observation.state': torch.zeros(1, 2, 4),
'action': torch.tensor([[[1.0, 1.0], [5.0, 5.0]]]),
'action_is_pad': torch.tensor([[False, True]]),
},
length=1,
)
mse = train_vla.compute_action_mse_validation(agent, val_loader, device='cpu')
self.assertEqual(mse, 0.0)
def test_eval_main_delegates_to_plain_run_eval_helper(self):
cfg = OmegaConf.create(
{