| 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 | |
| _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( | |