| |
| |
| @@ -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): |
| |
| |
| @@ -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) |
|
|
