diff --git a/outputs/libero_diana_videos/sim_libero_object_tabletop_front.mp4 b/outputs/libero_diana_videos/sim_libero_object_tabletop_front.mp4 new file mode 100644 index 0000000..9fb1514 Binary files /dev/null and b/outputs/libero_diana_videos/sim_libero_object_tabletop_front.mp4 differ diff --git a/outputs/libero_diana_videos/sim_libero_spatial_tabletop_front.mp4 b/outputs/libero_diana_videos/sim_libero_spatial_tabletop_front.mp4 new file mode 100644 index 0000000..3e10d4f Binary files /dev/null and b/outputs/libero_diana_videos/sim_libero_spatial_tabletop_front.mp4 differ diff --git a/roboimi/demos/vla_scripts/train_vla.py b/roboimi/demos/vla_scripts/train_vla.py index f42850a..ed14861 100644 --- a/roboimi/demos/vla_scripts/train_vla.py +++ b/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 的余弦退火 - progress = float(step - warmup_steps) / float(max(1, max_steps - warmup_steps)) - cosine_decay = 0.5 * (1.0 + math.cos(math.pi * progress)) - return max(min_lr_ratio, cosine_decay) + # 与 LeRobot SmolVLA/PI0 的 CosineDecayWithWarmupSchedulerConfig 对齐： + # 1) 余弦衰减步数是固定的 num_decay_steps（SmolVLA 默认 30k）， + # 2) 超过 decay_steps 后 clamp 到 decay_lr，不能继续 cos() 进入下一周期； + # 否则 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)) - grad_clip = float(cfg.train.get('grad_clip', 1.0)) - - optimizer = build_training_optimizer(agent, lr=cfg.train.lr, weight_decay=weight_decay) - log.info(f"🔧 优化器: AdamW (学习率={cfg.train.lr}, weight_decay={weight_decay})") + training_recipe = resolve_training_recipe(cfg) + weight_decay = training_recipe['weight_decay'] + grad_clip = training_recipe['grad_clip'] + train_lr = training_recipe['lr'] + 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)) - scheduler_type = cfg.train.get('scheduler_type', 'cosine') - min_lr = float(cfg.train.get('min_lr', 1e-6)) + warmup_steps = training_recipe['warmup_steps'] + scheduler_type = training_recipe['scheduler_type'] + 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 ) diff --git a/roboimi/vla/agent_smolvla_native.py b/roboimi/vla/agent_smolvla_native.py index d73b74d..571d38d 100644 --- a/roboimi/vla/agent_smolvla_native.py +++ b/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' diff --git a/roboimi/vla/conf/agent/smolvla_native.yaml b/roboimi/vla/conf/agent/smolvla_native.yaml index 283253d..a285066 100644 --- a/roboimi/vla/conf/agent/smolvla_native.yaml +++ b/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" -chunk_size: 16 +normalization_type: "gaussian" +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} diff --git a/tests/test_smolvla_native_agent.py b/tests/test_smolvla_native_agent.py index 237e857..2424110 100644 --- a/tests/test_smolvla_native_agent.py +++ b/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.n_action_steps, 8) + self.assertEqual(cfg.agent.chunk_size, 32) + 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() diff --git a/tests/test_train_vla_transformer_optimizer.py b/tests/test_train_vla_transformer_optimizer.py index 3c893f4..63ecc7a 100644 --- a/tests/test_train_vla_transformer_optimizer.py +++ b/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)