File size: 1,267 Bytes
d27c910 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 | 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(
|