feat(vla): add SmolVLA conditioning and experiment artifacts

tests/test_train_vla_swanlab_logging.py

@@ -113,6 +113,20 @@ class FakeAgent(nn.Module):
     def get_normalization_stats(self):
         return {}
 
+    def predict_action_chunk(self, agent_input):
+        return agent_input['action']
+
+
+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 FakeSwanLab:
     def __init__(self, init_error=None, log_errors=None, finish_error=None, image_errors=None):
@@ -339,6 +353,8 @@ class TrainVLASwanLabLoggingTest(unittest.TestCase):
                 batch_size=2,
                 num_workers=0,
                 val_split=0.25,
+                val_episode_indices=None,
+                action_mse_val_freq_epochs=0,
                 seed=0,
                 lr=1e-3,
                 max_steps=2,
@@ -442,6 +458,8 @@ class TrainVLASwanLabLoggingTest(unittest.TestCase):
                         'batch_size': 2,
                         'num_workers': 0,
                         'val_split': 0.25,
+                        'val_episode_indices': None,
+                        'action_mse_val_freq_epochs': 0,
                         'seed': 0,
                         'lr': 1e-3,
                         'max_steps': 2,
@@ -539,6 +557,58 @@ class TrainVLASwanLabLoggingTest(unittest.TestCase):
 
         self.assertEqual(fake_swanlab.finish_calls, 1)
 
+    def test_run_training_logs_action_mse_only_for_explicit_val_episode_validation(self):
+        module = self._load_train_vla_module()
+        run_training = self._get_run_training(module)
+        cfg = self._make_cfg()
+        cfg.train.val_split = 0.0
+        cfg.train.val_episode_indices = [100]
+        cfg.train.action_mse_val_freq_epochs = 1
+        cfg.train.max_steps = 4
+        cfg.train.save_freq = 100
+        cfg.train.rollout_val_freq_epochs = 0
+        agent = FakeAgent()
+        fake_swanlab = FakeSwanLab()
+        real_import_module = importlib.import_module
+
+        def fake_instantiate(config_node, **kwargs):
+            if config_node is cfg.data:
+                return FakeEpisodeDataset(kwargs.get('episode_indices'))
+            if config_node is cfg.agent:
+                return agent
+            raise AssertionError(f'unexpected instantiate config: {config_node!r}')
+
+        def fake_import_module(name, package=None):
+            if name == 'swanlab':
+                return fake_swanlab
+            return real_import_module(name, package)
+
+        def loader_factory(_dataset, *, shuffle, **_kwargs):
+            batch = {
+                'observation.front': torch.zeros(1, 2, 3, 2, 2),
+                'observation.state': torch.zeros(1, 2, 4),
+                'action': torch.zeros(1, 2, 2),
+                'action_is_pad': torch.zeros(1, 2, dtype=torch.bool),
+            }
+            return FakeLoader(batch)
+
+        with tempfile.TemporaryDirectory() as tempdir:
+            previous_cwd = os.getcwd()
+            try:
+                os.chdir(tempdir)
+                with mock.patch.object(module, 'instantiate', side_effect=fake_instantiate), \
+                     mock.patch.object(module, 'DataLoader', side_effect=loader_factory), \
+                     mock.patch.object(module, 'get_lr_schedule_with_warmup', return_value=FakeScheduler()), \
+                     mock.patch.object(module, 'tqdm', side_effect=lambda iterable, **kwargs: FakeProgressBar(iterable)), \
+                     mock.patch.object(module.torch, 'save', return_value=None), \
+                     mock.patch.object(module.importlib, 'import_module', side_effect=fake_import_module):
+                    run_training(cfg)
+            finally:
+                os.chdir(previous_cwd)
+
+        logged_keys = set().union(*(payload.keys() for payload, _step in fake_swanlab.log_calls))
+        self.assertIn('val/action_mse', logged_keys)
+
     def test_run_training_warns_and_continues_when_swanlab_log_and_finish_fail(self):
         module = self._load_train_vla_module()
         run_training = self._get_run_training(module)