--- a/src/openpi/training/config.py +++ b/src/openpi/training/config.py @@ -968,6 +968,193 @@ # RoboArena & PolaRiS configs. *roboarena_config.get_roboarena_configs(), *polaris_config.get_polaris_configs(), + # + # Finetunes used in "Which VLAs Learn a Grasp That Transfers to New Object Positions?". + # + # pi0.5 LoRA on LIBERO's own 50 LIBERO-Object demos x 10 tasks (the whole "pick up the X and place + # it in the basket" task, native 128x128 frames, converted with + # training/convert_libero_original_to_lerobot.py): pi05_base init, LoRA on both experts, batch 24, + # peak lr 5e-5 cosine to 1e-6 over 12,000 steps, every 1,000th checkpoint kept. + # + TrainConfig( + name="pi05_libero_object_orig", + model=pi0_config.Pi0Config( + pi05=True, + action_horizon=10, + discrete_state_input=False, + paligemma_variant="gemma_2b_lora", + action_expert_variant="gemma_300m_lora", + ), + data=LeRobotLiberoDataConfig( + repo_id="libero/object_original", + base_config=DataConfig(prompt_from_task=True), + extra_delta_transform=False, + ), + weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"), + freeze_filter=pi0_config.Pi0Config( + pi05=True, + action_horizon=10, + discrete_state_input=False, + paligemma_variant="gemma_2b_lora", + action_expert_variant="gemma_300m_lora", + ).get_freeze_filter(), + ema_decay=None, + batch_size=24, + lr_schedule=_optimizer.CosineDecaySchedule( + warmup_steps=200, peak_lr=5e-5, decay_steps=12_000, decay_lr=1e-6 + ), + optimizer=_optimizer.AdamW(clip_gradient_norm=1.0), + num_train_steps=12_000, + save_interval=1000, + keep_period=1000, + num_workers=16, + wandb_enabled=False, + ), + # + # Data ladder on the same demonstrations. Each rung trains on a stratified subset (every task + # keeps its share; at 2 % that is one demo per task) for the SAME 2,000 steps at batch 24, so only + # the pool the 48,000 training frames are drawn from shrinks. + # + # 50 % of the demos: 250 episodes. + TrainConfig( + name="pi05_libero_object_p50", + model=pi0_config.Pi0Config( + pi05=True, + action_horizon=10, + discrete_state_input=False, + paligemma_variant="gemma_2b_lora", + action_expert_variant="gemma_300m_lora", + ), + data=LeRobotLiberoDataConfig( + repo_id="libero/object_orig_p50", + base_config=DataConfig(prompt_from_task=True), + extra_delta_transform=False, + ), + weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"), + freeze_filter=pi0_config.Pi0Config( + pi05=True, + action_horizon=10, + discrete_state_input=False, + paligemma_variant="gemma_2b_lora", + action_expert_variant="gemma_300m_lora", + ).get_freeze_filter(), + ema_decay=None, + batch_size=24, + lr_schedule=_optimizer.CosineDecaySchedule( + warmup_steps=200, peak_lr=5e-5, decay_steps=2_000, decay_lr=1e-6 + ), + optimizer=_optimizer.AdamW(clip_gradient_norm=1.0), + num_train_steps=2_000, + save_interval=500, + keep_period=500, + num_workers=16, + wandb_enabled=False, + ), + # 25 % of the demos: 125 episodes. + TrainConfig( + name="pi05_libero_object_p25", + model=pi0_config.Pi0Config( + pi05=True, + action_horizon=10, + discrete_state_input=False, + paligemma_variant="gemma_2b_lora", + action_expert_variant="gemma_300m_lora", + ), + data=LeRobotLiberoDataConfig( + repo_id="libero/object_orig_p25", + base_config=DataConfig(prompt_from_task=True), + extra_delta_transform=False, + ), + weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"), + freeze_filter=pi0_config.Pi0Config( + pi05=True, + action_horizon=10, + discrete_state_input=False, + paligemma_variant="gemma_2b_lora", + action_expert_variant="gemma_300m_lora", + ).get_freeze_filter(), + ema_decay=None, + batch_size=24, + lr_schedule=_optimizer.CosineDecaySchedule( + warmup_steps=200, peak_lr=5e-5, decay_steps=2_000, decay_lr=1e-6 + ), + optimizer=_optimizer.AdamW(clip_gradient_norm=1.0), + num_train_steps=2_000, + save_interval=500, + keep_period=500, + num_workers=16, + wandb_enabled=False, + ), + # 10 % of the demos: 50 episodes. + TrainConfig( + name="pi05_libero_object_p10", + model=pi0_config.Pi0Config( + pi05=True, + action_horizon=10, + discrete_state_input=False, + paligemma_variant="gemma_2b_lora", + action_expert_variant="gemma_300m_lora", + ), + data=LeRobotLiberoDataConfig( + repo_id="libero/object_orig_p10", + base_config=DataConfig(prompt_from_task=True), + extra_delta_transform=False, + ), + weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"), + freeze_filter=pi0_config.Pi0Config( + pi05=True, + action_horizon=10, + discrete_state_input=False, + paligemma_variant="gemma_2b_lora", + action_expert_variant="gemma_300m_lora", + ).get_freeze_filter(), + ema_decay=None, + batch_size=24, + lr_schedule=_optimizer.CosineDecaySchedule( + warmup_steps=200, peak_lr=5e-5, decay_steps=2_000, decay_lr=1e-6 + ), + optimizer=_optimizer.AdamW(clip_gradient_norm=1.0), + num_train_steps=2_000, + save_interval=500, + keep_period=500, + num_workers=16, + wandb_enabled=False, + ), + # 2 % of the demos: 10 episodes (one per task). + TrainConfig( + name="pi05_libero_object_p02", + model=pi0_config.Pi0Config( + pi05=True, + action_horizon=10, + discrete_state_input=False, + paligemma_variant="gemma_2b_lora", + action_expert_variant="gemma_300m_lora", + ), + data=LeRobotLiberoDataConfig( + repo_id="libero/object_orig_p02", + base_config=DataConfig(prompt_from_task=True), + extra_delta_transform=False, + ), + weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"), + freeze_filter=pi0_config.Pi0Config( + pi05=True, + action_horizon=10, + discrete_state_input=False, + paligemma_variant="gemma_2b_lora", + action_expert_variant="gemma_300m_lora", + ).get_freeze_filter(), + ema_decay=None, + batch_size=24, + lr_schedule=_optimizer.CosineDecaySchedule( + warmup_steps=200, peak_lr=5e-5, decay_steps=2_000, decay_lr=1e-6 + ), + optimizer=_optimizer.AdamW(clip_gradient_norm=1.0), + num_train_steps=2_000, + save_interval=500, + keep_period=500, + num_workers=16, + wandb_enabled=False, + ), ] if len({config.name for config in _CONFIGS}) != len(_CONFIGS): --- a/scripts/train.py +++ b/scripts/train.py @@ -159,6 +159,17 @@ params = state.params.filter(config.trainable_filter) updates, new_opt_state = state.tx.update(grads, state.opt_state, params) + + # Drop the update if anything in it is non-finite. A freshly compiled train_step can + # occasionally overflow -- one of 320 loss entries goes inf, and re-running the identical + # computation in another process is finite -- and a single nan update poisons the weights and + # then gets written over the good checkpoints by the next save. Skipping leaves params and the + # optimizer moments exactly as they were; `skipped` in the log line says how often it happens. + grad_norm = optax.global_norm(grads) + is_finite = jnp.isfinite(grad_norm) & jnp.isfinite(loss) + updates = jax.tree.map(lambda u: jnp.where(is_finite, u, jnp.zeros_like(u)), updates) + new_opt_state = jax.tree.map(lambda new, old: jnp.where(is_finite, new, old), new_opt_state, state.opt_state) + new_params = optax.apply_updates(params, updates) # Update the model in place and return the new full state. @@ -184,9 +195,11 @@ ), ) info = { - "loss": loss, - "grad_norm": optax.global_norm(grads), + # Masked so that one skipped step cannot turn a whole log window's mean into inf/nan. + "loss": jnp.where(is_finite, loss, 0.0), + "grad_norm": jnp.where(is_finite, grad_norm, 0.0), "param_norm": optax.global_norm(kernel_params), + "skipped": (~is_finite).astype(jnp.float32), } return new_state, info @@ -265,6 +278,9 @@ reduced_info = jax.device_get(jax.tree.map(jnp.mean, stacked_infos)) info_str = ", ".join(f"{k}={v:.4f}" for k, v in reduced_info.items()) pbar.write(f"Step {step}: {info_str}") + # tqdm_loggable swallows pbar.write, so with wandb off the run would produce no loss + # curve at all; log it through the normal logger too. + logging.info(f"Step {step}: {info_str}") wandb.log(reduced_info, step=step) infos = [] batch = next(data_iter)