File size: 2,473 Bytes
787bc28 | 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 43 44 45 46 47 48 49 50 51 52 53 54 | # 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,
),
|