| # Train / eval config for this checkpoint: `pi05_robopro_top_cam_jax` | |
| # | |
| # This is the exact openpi TrainConfig used to fine-tune and to load this | |
| # checkpoint. To use it, add the TrainConfig(...) entry below to the `_CONFIGS` | |
| # list in `src/openpi/training/config.py` of your openpi checkout, then load it | |
| # with `openpi.training.config.get_config("pi05_robopro_top_cam_jax")`. | |
| # | |
| # Symbols referenced (already imported at the top of openpi's config.py): | |
| # TrainConfig, DataConfig, LeRobotAlohaDataConfig | |
| # pi0_config = openpi.models.pi0_config | |
| # _transforms = openpi.transforms | |
| # weight_loaders = openpi.training.weight_loaders | |
| # _optimizer = openpi.training.optimizer | |
| # | |
| # Dataset: robopro top-cam LeRobot v2.1 (`roboreal_lerobot`), robot_type roboreal, | |
| # 25 fps, 14-DoF dual-arm, cams countertop/left/right. Set | |
| # HF_LEROBOT_HOME=<parent> so repo_id `roboreal_lerobot` resolves locally. | |
| # Base weights: JAX pi05_base from gs://openpi-assets (auto-download). | |
| TrainConfig( | |
| name="pi05_robopro_top_cam_jax", | |
| model=pi0_config.Pi0Config(pi05=True), | |
| data=LeRobotAlohaDataConfig( | |
| repo_id="roboreal_lerobot", | |
| # Map the dataset's raw feature keys -> the model's expected keys. | |
| # NOTE: cam_high is fed from the COUNTERTOP (overhead) camera. | |
| repack_transforms=_transforms.Group(inputs=[ | |
| _transforms.RepackTransform({ | |
| "images": { | |
| "cam_high": "observation.images.countertop", | |
| "cam_left_wrist": "observation.images.left", | |
| "cam_right_wrist": "observation.images.right", | |
| }, | |
| "state": "observation.state", | |
| "actions": "action", | |
| "prompt": "prompt", | |
| }) | |
| ]), | |
| base_config=DataConfig( | |
| prompt_from_task=True, | |
| ), | |
| # (defaults inherited from LeRobotAlohaDataConfig:) | |
| # adapt_to_pi=True, use_delta_joint_actions=True | |
| # -> arm joints trained as delta, grippers absolute; | |
| # AbsoluteActions on output => returned actions are ABSOLUTE. | |
| ), | |
| weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"), | |
| lr_schedule=_optimizer.CosineDecaySchedule(decay_steps=30_000), | |
| num_train_steps=30_000, | |
| batch_size=192, # 3 GPUs = 64/GPU (must be divisible by device count) | |
| num_workers=16, | |
| fsdp_devices=1, | |
| ), | |