diff --git a/src/openpi/training/config.py b/src/openpi/training/config.py index 4ca47e1..cb06175 100644 --- a/src/openpi/training/config.py +++ b/src/openpi/training/config.py @@ -826,6 +826,36 @@ _CONFIGS = [ batch_size=64, ), # + # Fine-tuning Kinova captures (plush-to-box). + # + TrainConfig( + name="pi05_kinova_plush_box", + model=pi0_config.Pi0Config( + pi05=True, + action_dim=32, + action_horizon=10, + ), + data=LeRobotLiberoDataConfig( + repo_id="local/kinova_plush_box", + base_config=DataConfig(prompt_from_task=True), + extra_delta_transform=True, + ), + weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"), + lr_schedule=_optimizer.CosineDecaySchedule( + warmup_steps=1_000, + peak_lr=2.5e-5, + decay_steps=8_000, + decay_lr=2.5e-6, + ), + num_train_steps=8_000, + batch_size=16, + num_workers=0, + log_interval=50, + save_interval=1_000, + keep_period=5_000, + wandb_enabled=False, + ), + # # Fine-tuning DROID configs. # TrainConfig(