feat(train): align native SmolVLA recipe

roboimi/demos/vla_scripts/train_vla.py

@@ -71,6 +71,19 @@ from hydra.utils import instantiate
 
 log = logging.getLogger(__name__)
 
+
+SMOLVLA_NATIVE_TRAINING_PRESET = {
+    'lr': 1e-4,
+    'betas': (0.9, 0.95),
+    'eps': 1e-8,
+    'weight_decay': 1e-10,
+    'grad_clip': 10.0,
+    'warmup_steps': 1000,
+    'scheduler_type': 'cosine',
+    'scheduler_decay_steps': 30000,
+    'scheduler_decay_lr': 2.5e-6,
+}
+
 # 注册列表长度解析器（用于配置中如 ${len:${data.camera_names}}）
 if not OmegaConf.has_resolver("len"):
     OmegaConf.register_new_resolver("len", lambda x: len(x))
@@ -154,7 +167,7 @@ def get_lr_schedule_with_warmup(optimizer, warmup_steps, max_steps, scheduler_ty
     Args:
         optimizer: PyTorch 优化器
         warmup_steps: 预热步数
-        max_steps: 总训练步数
+        max_steps: 余弦衰减步数
         scheduler_type: 预热后的调度器类型 ('cosine' 或 'constant')
         min_lr: 最小学习率（用于余弦衰减）
 
@@ -167,16 +180,24 @@ def get_lr_schedule_with_warmup(optimizer, warmup_steps, max_steps, scheduler_ty
     min_lr_ratio = min_lr / base_lr if base_lr > 0 else 0.0
 
     def lr_lambda(step):
-        # 预热阶段：从 0 线性增加到 1
+        # LeRobot CosineDecayWithWarmupSchedulerConfig 的线性预热：
+        # 从一个很小的非零 LR 开始，避免首步完全为 0。
         if step < warmup_steps:
-            return float(step) / float(max(1, warmup_steps))
+            if step <= 0:
+                return 1.0 / float(max(1, warmup_steps + 1))
+            frac = 1.0 - float(step) / float(max(1, warmup_steps))
+            return (1.0 / float(max(1, warmup_steps + 1)) - 1.0) * frac + 1.0
 
         # 预热后阶段
         if scheduler_type == 'cosine':
-            # 从 1 到 min_lr_ratio 的余弦退火
+            # 与 LeRobot SmolVLA/PI0 的 CosineDecayWithWarmupSchedulerConfig 对齐：
-            progress = float(step - warmup_steps) / float(max(1, max_steps - warmup_steps))
+            # 1) 余弦衰减步数是固定的 num_decay_steps（SmolVLA 默认 30k），
-            cosine_decay = 0.5 * (1.0 + math.cos(math.pi * progress))
+            # 2) 超过 decay_steps 后 clamp 到 decay_lr，不能继续 cos() 进入下一周期；
-            return max(min_lr_ratio, cosine_decay)
+            #    否则 150k 训练会显示成“正弦波”式反复升降。
+            decay_steps = max(1, int(max_steps))
+            clamped_step = min(max(int(step), 0), decay_steps)
+            cosine_decay = 0.5 * (1.0 + math.cos(math.pi * clamped_step / decay_steps))
+            return (1.0 - min_lr_ratio) * cosine_decay + min_lr_ratio
         else:
             # 恒定学习率
             return 1.0
@@ -184,6 +205,46 @@ def get_lr_schedule_with_warmup(optimizer, warmup_steps, max_steps, scheduler_ty
     return LambdaLR(optimizer, lr_lambda)
 
 
+def _is_smolvla_native_agent_config(agent_cfg) -> bool:
+    target = str(agent_cfg.get('_target_', '')) if hasattr(agent_cfg, 'get') else ''
+    return target == 'roboimi.vla.agent_smolvla_native.SmolVLANativeAgent'
+
+
+def resolve_training_recipe(cfg):
+    """Return optimizer/scheduler knobs, using LeRobot SmolVLA defaults for smolvla_native."""
+    if _is_smolvla_native_agent_config(cfg.agent):
+        preset = SMOLVLA_NATIVE_TRAINING_PRESET
+        scheduler_steps = int(cfg.train.get(
+            'scheduler_decay_steps',
+            cfg.train.get('lr_scheduler_steps', cfg.train.max_steps),
+        ))
+        return {
+            'lr': float(preset['lr']),
+            'betas': tuple(preset['betas']),
+            'eps': float(preset['eps']),
+            'weight_decay': float(preset['weight_decay']),
+            'grad_clip': float(preset['grad_clip']),
+            'warmup_steps': int(preset['warmup_steps']),
+            'scheduler_type': str(preset['scheduler_type']),
+            'scheduler_steps': scheduler_steps,
+            'min_lr': float(preset['scheduler_decay_lr']),
+            'preset_name': 'lerobot_smolvla',
+        }
+
+    return {
+        'lr': float(cfg.train.lr),
+        'betas': (0.9, 0.999),
+        'eps': 1e-8,
+        'weight_decay': float(cfg.train.get('weight_decay', 1e-5)),
+        'grad_clip': float(cfg.train.get('grad_clip', 1.0)),
+        'warmup_steps': int(cfg.train.get('warmup_steps', 500)),
+        'scheduler_type': str(cfg.train.get('scheduler_type', 'cosine')),
+        'scheduler_steps': int(cfg.train.max_steps),
+        'min_lr': float(cfg.train.get('min_lr', 1e-6)),
+        'preset_name': 'config',
+    }
+
+
 def _instantiate_dataset(cfg, dataset_image_resize_shape, episode_indices=None):
     kwargs = {'image_resize_shape': dataset_image_resize_shape}
     if episode_indices is not None:
@@ -256,7 +317,7 @@ def build_train_val_datasets(cfg, dataset_image_resize_shape):
     return dataset, train_dataset, val_dataset, None
 
 
-def build_training_optimizer(agent, lr, weight_decay):
+def build_training_optimizer(agent, lr, weight_decay, betas=(0.9, 0.999), eps=1e-8):
     """为训练脚本构建优化器，优先复用任意 head 自带的参数分组。"""
     trainable_params = [param for param in agent.parameters() if param.requires_grad]
     noise_pred_net = getattr(agent, 'noise_pred_net', None)
@@ -264,7 +325,7 @@ def build_training_optimizer(agent, lr, weight_decay):
     use_head_groups = callable(get_optim_groups)
 
     if not use_head_groups:
-        return AdamW(trainable_params, lr=lr, weight_decay=weight_decay)
+        return AdamW(trainable_params, lr=lr, weight_decay=weight_decay, betas=betas, eps=eps)
 
     head_groups = []
     grouped_param_ids = set()
@@ -306,7 +367,7 @@ def build_training_optimizer(agent, lr, weight_decay):
     if grouped_param_ids != all_trainable_param_ids:
         raise ValueError('Optimizer parameter groups must include each trainable parameter exactly once')
 
-    return AdamW(optim_groups, lr=lr, weight_decay=weight_decay)
+    return AdamW(optim_groups, lr=lr, weight_decay=weight_decay, betas=betas, eps=eps)
 
 
 def _init_swanlab(cfg):
@@ -353,7 +414,12 @@ def _init_swanlab(cfg):
         init_kwargs['experiment_name'] = run_name
 
     try:
-        swanlab.init(**init_kwargs)
+        run = swanlab.init(**init_kwargs)
+        if run is not None:
+            try:
+                setattr(swanlab, "_roboimi_swanlab_run", run)
+            except Exception:
+                pass
     except Exception as exc:
         raise RuntimeError(
             f"SwanLab logging is enabled, but SwanLab init/login failed: {exc}"
@@ -362,6 +428,24 @@ def _init_swanlab(cfg):
     return swanlab
 
 
+def _log_swanlab_init_details(swanlab_module):
+    if swanlab_module is None:
+        return
+    run = getattr(swanlab_module, "_roboimi_swanlab_run", None)
+    if run is None:
+        run = getattr(swanlab_module, "run", None)
+    if run is None:
+        return
+
+    url = getattr(run, "url", None)
+    swanlog_dir = getattr(run, "swanlog_dir", None)
+    log.info(
+        "🦢 SwanLab initialized%s%s",
+        f" | url={url}" if url else "",
+        f" | swanlog_dir={swanlog_dir}" if swanlog_dir else "",
+    )
+
+
 def _log_to_swanlab(swanlab_module, payload, step=None):
     if swanlab_module is None:
         return
@@ -440,6 +524,7 @@ def _run_training(cfg: DictConfig):
     log.info(f"🚀 开始 VLA 训练 (设备: {cfg.train.device})")
     _configure_cuda_runtime(cfg)
     swanlab_module = _init_swanlab(cfg)
+    _log_swanlab_init_details(swanlab_module)
     try:
         action_mse_val_freq_epochs = int(cfg.train.get('action_mse_val_freq_epochs', 0) or 0)
         explicit_val_episode_indices = cfg.train.get('val_episode_indices', None)
@@ -609,21 +694,35 @@ def _run_training(cfg: DictConfig):
         # =========================================================================
         # 4. 设置优化器与学习率调度器
         # =========================================================================
-        weight_decay = float(cfg.train.get('weight_decay', 1e-5))
+        training_recipe = resolve_training_recipe(cfg)
-        grad_clip = float(cfg.train.get('grad_clip', 1.0))
+        weight_decay = training_recipe['weight_decay']
-
+        grad_clip = training_recipe['grad_clip']
-        optimizer = build_training_optimizer(agent, lr=cfg.train.lr, weight_decay=weight_decay)
+        train_lr = training_recipe['lr']
-        log.info(f"🔧 优化器: AdamW (学习率={cfg.train.lr}, weight_decay={weight_decay})")
+        optimizer = build_training_optimizer(
+            agent,
+            lr=train_lr,
+            weight_decay=weight_decay,
+            betas=training_recipe['betas'],
+            eps=training_recipe['eps'],
+        )
+        log.info(
+            "🔧 优化器: AdamW (preset=%s, 学习率=%s, betas=%s, eps=%s, weight_decay=%s)",
+            training_recipe['preset_name'],
+            train_lr,
+            training_recipe['betas'],
+            training_recipe['eps'],
+            weight_decay,
+        )
 
         # 设置带预热的学習率调度器
-        warmup_steps = int(cfg.train.get('warmup_steps', 500))
+        warmup_steps = training_recipe['warmup_steps']
-        scheduler_type = cfg.train.get('scheduler_type', 'cosine')
+        scheduler_type = training_recipe['scheduler_type']
-        min_lr = float(cfg.train.get('min_lr', 1e-6))
+        min_lr = training_recipe['min_lr']
 
         scheduler = get_lr_schedule_with_warmup(
             optimizer,
             warmup_steps=warmup_steps,
-            max_steps=cfg.train.max_steps,
+            max_steps=training_recipe['scheduler_steps'],
             scheduler_type=scheduler_type,
             min_lr=min_lr
         )

roboimi/vla/agent_smolvla_native.py

@@ -192,7 +192,10 @@ class SmolVLANativeAgent(nn.Module):
         return [fallback_or(item) for item in task]
 
     def _tokenize_tasks(self, task, batch_size: int, device: torch.device) -> dict:
-        tasks = [str(text) + '\n' for text in self._resolve_task(task, batch_size)]
+        tasks = []
+        for text in self._resolve_task(task, batch_size):
+            text = str(text)
+            tasks.append(text if text.endswith('\n') else f'{text}\n')
         old_padding_side = getattr(self.tokenizer, 'padding_side', None)
         if old_padding_side is not None:
             self.tokenizer.padding_side = 'right'

roboimi/vla/conf/agent/smolvla_native.yaml

@@ -7,11 +7,11 @@ tokenizer_name: ${agent.model_config.vlm_model_name}
 
 action_dim: 16
 obs_dim: 16
-normalization_type: "min_max"
+normalization_type: "gaussian"
-chunk_size: 16
+chunk_size: 32
 pred_horizon: ${agent.chunk_size}
 obs_horizon: 2
-n_action_steps: 8
+n_action_steps: 16
 num_action_steps: ${agent.n_action_steps}
 action_chunk_start: 0
 camera_names: ${data.camera_names}

tests/test_smolvla_native_agent.py

@@ -416,12 +416,15 @@ class SmolVLANativeAgentTest(unittest.TestCase):
         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(cfg.agent.normalization_type, 'gaussian')
         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.chunk_size, 32)
-        self.assertEqual(cfg.agent.n_action_steps, 8)
+        self.assertEqual(cfg.agent.n_action_steps, 16)
         self.assertEqual(cfg.agent.model_config.max_state_dim, 32)
         self.assertEqual(cfg.agent.model_config.max_action_dim, 32)
+        self.assertEqual(cfg.agent.model_config.chunk_size, 32)
+        self.assertEqual(cfg.agent.model_config.n_action_steps, 16)
         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])
@@ -430,6 +433,13 @@ class SmolVLANativeAgentTest(unittest.TestCase):
         self.assertNotIn('tokenizer_name', cfg.agent.model_config)
         self.assertEqual(cfg.agent.model_config.vlm_model_name, 'HuggingFaceTB/SmolVLM2-500M-Video-Instruct')
 
+    def test_tokenize_tasks_preserves_existing_single_newline(self):
+        agent = _make_agent(model_config={'resize_imgs_with_padding': None, 'pad_language_to': 'max_length', 'tokenizer_max_length': 12})
+
+        agent._tokenize_tasks(['already newline\n', 'needs newline'], batch_size=2, device=torch.device('cpu'))
+
+        self.assertEqual(agent.tokenizer.calls[-1]['texts'], ['already newline\n', 'needs newline\n'])
+
 
 if __name__ == '__main__':
     unittest.main()

tests/test_train_vla_transformer_optimizer.py

@@ -39,6 +39,17 @@ class FakeLoader:
         return iter(())
 
 
+class FakeTqdm:
+    def __init__(self, iterable, **_kwargs):
+        self.iterable = iterable
+
+    def __iter__(self):
+        return iter(self.iterable)
+
+    def set_postfix(self, *_args, **_kwargs):
+        return None
+
+
 class FakeScheduler:
     def state_dict(self):
         return {}
@@ -46,13 +57,18 @@ class FakeScheduler:
     def load_state_dict(self, state_dict):
         return None
 
+    def step(self):
+        return None
+
 
 class RecordingAdamW:
     created = []
 
-    def __init__(self, params, lr, weight_decay):
+    def __init__(self, params, lr, weight_decay, betas=(0.9, 0.999), eps=1e-8):
         self.lr = lr
         self.weight_decay = weight_decay
+        self.betas = betas
+        self.eps = eps
         self.param_groups = self._normalize_param_groups(params, lr, weight_decay)
         RecordingAdamW.created.append(self)
 
@@ -79,6 +95,12 @@ class RecordingAdamW:
     def load_state_dict(self, state_dict):
         return None
 
+    def zero_grad(self):
+        return None
+
+    def step(self):
+        return None
+
 
 class RecordingTransformerHead(nn.Module):
     def __init__(self):
@@ -324,7 +346,7 @@ class TrainVLATransformerOptimizerTest(unittest.TestCase):
                      mock.patch.object(module, 'get_lr_schedule_with_warmup', return_value=FakeScheduler()), \
                      mock.patch.object(module, 'AdamW', RecordingAdamW), \
                      mock.patch.object(module.torch, 'save', return_value=None), \
-                     mock.patch.object(module, 'tqdm', side_effect=lambda iterable, **kwargs: iterable):
+                     mock.patch.object(module, 'tqdm', side_effect=lambda iterable, **kwargs: FakeTqdm(iterable, **kwargs)):
                     module.main(cfg)
             finally:
                 os.chdir(previous_cwd)
@@ -404,7 +426,7 @@ class TrainVLATransformerOptimizerTest(unittest.TestCase):
                      mock.patch.object(module, 'get_lr_schedule_with_warmup', return_value=FakeScheduler()), \
                      mock.patch.object(module, 'AdamW', RecordingAdamW), \
                      mock.patch.object(module.torch, 'save', return_value=None), \
-                     mock.patch.object(module, 'tqdm', side_effect=lambda iterable, **kwargs: iterable):
+                     mock.patch.object(module, 'tqdm', side_effect=lambda iterable, **kwargs: FakeTqdm(iterable, **kwargs)):
                     module.main(cfg)
             finally:
                 os.chdir(previous_cwd)
@@ -464,3 +486,117 @@ class TrainVLASmolVLAOptimizerTest(unittest.TestCase):
         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)
+
+    def test_smolvla_native_training_preset_overrides_optimizer_scheduler_and_grad_clip(self):
+        module = TrainVLATransformerOptimizerTest()._load_train_vla_module()
+
+        class _NativeAgent(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.model = nn.Linear(2, 2)
+
+            def to(self, device):
+                return self
+
+            def get_normalization_stats(self):
+                return {}
+
+        agent = _NativeAgent()
+        cfg = TrainVLATransformerOptimizerTest()._make_cfg()
+        cfg.agent = AttrDict(_target_='roboimi.vla.agent_smolvla_native.SmolVLANativeAgent')
+        cfg.train.lr = 9e-4
+        cfg.train.max_steps = 1
+        cfg.train.weight_decay = 0.123
+        cfg.train.grad_clip = 1.0
+        cfg.train.warmup_steps = 7
+        cfg.train.scheduler_type = 'constant'
+        cfg.train.min_lr = 1e-7
+
+        scheduler_calls = []
+        clip_calls = []
+
+        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_scheduler(*args, **kwargs):
+            scheduler_calls.append(kwargs)
+            return FakeScheduler()
+
+        def fake_clip(parameters, max_norm):
+            clip_calls.append(float(max_norm))
+            return torch.tensor(0.0)
+
+        class OneBatchLoader:
+            def __len__(self):
+                return 1
+
+            def __iter__(self):
+                batch = {
+                    'observation.front': torch.zeros(1, 1, 1, 2, 2),
+                    'observation.state': torch.zeros(1, 1, 2),
+                    'action': torch.zeros(1, 1, 2),
+                }
+                return iter([batch])
+
+        def fake_compute_loss(_batch):
+            return agent.model.weight.sum() * 0.0 + torch.tensor(1.0, requires_grad=True)
+
+        agent.compute_loss = fake_compute_loss
+
+        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=lambda *args, **kwargs: OneBatchLoader()), \
+                     mock.patch.object(module, 'get_lr_schedule_with_warmup', side_effect=fake_scheduler), \
+                     mock.patch.object(module, 'AdamW', RecordingAdamW), \
+                     mock.patch.object(module.torch.nn.utils, 'clip_grad_norm_', side_effect=fake_clip), \
+                     mock.patch.object(module.torch, 'save', return_value=None), \
+                     mock.patch.object(module, 'tqdm', side_effect=lambda iterable, **kwargs: FakeTqdm(iterable, **kwargs)):
+                    module.main(cfg)
+            finally:
+                os.chdir(previous_cwd)
+
+        optimizer = RecordingAdamW.created[-1]
+        self.assertEqual(optimizer.lr, 1e-4)
+        self.assertEqual(optimizer.weight_decay, 1e-10)
+        self.assertEqual(optimizer.betas, (0.9, 0.95))
+        self.assertEqual(optimizer.eps, 1e-8)
+        self.assertEqual(scheduler_calls[-1], {
+            'warmup_steps': 1000,
+            'max_steps': 1,
+            'scheduler_type': 'cosine',
+            'min_lr': 2.5e-6,
+        })
+        self.assertEqual(clip_calls, [10.0])
+
+    def test_cosine_scheduler_spans_requested_training_steps_then_clamps(self):
+        module = TrainVLATransformerOptimizerTest()._load_train_vla_module()
+        param = nn.Parameter(torch.tensor(1.0))
+        optimizer = torch.optim.SGD([param], lr=1e-4)
+        base_lr = optimizer.param_groups[0]['lr']
+
+        scheduler = module.get_lr_schedule_with_warmup(
+            optimizer,
+            warmup_steps=1000,
+            max_steps=150000,
+            scheduler_type='cosine',
+            min_lr=2.5e-6,
+        )
+
+        lr_lambda = scheduler.lr_lambdas[0]
+        observed = {}
+        for step in (0, 1, 1000, 30000, 40000, 150000, 160000):
+            observed[step] = base_lr * lr_lambda(step)
+
+        self.assertGreater(observed[1], observed[0])
+        self.assertLess(observed[1000], 1e-4)
+        self.assertGreater(observed[30000], observed[40000])
+        self.assertGreater(observed[40000], observed[150000])
+        self.assertAlmostEqual(observed[150000], 2.5e-6, places=12)
+        self.assertAlmostEqual(observed[160000], 2.5e-6, places=12)