Wr3ck1Am/pi05-real-single-task-lora-ckpts

Pi05 Single-Task LoRA — Real-World Tasks

Per-task LoRA fine-tunes of the pi0.5 (PaliGemma + action expert) policy on real-robot data from harrywang01/real_world_task.

Tasks

Task	Instruction
wipe_new	"Brush the two plates with the brushes, and each brush can only be used once."
can	"Swap the positions of the pink and white soda cans by using another empty location as a buffer."
peg_new	"Pick up the peg from its starting position and insert it into the hole below that has exactly the same length. Each hole can only be attempted once."
uncover_new	"Lift the cups to find the small cube hidden underneath, and each cup may only be lifted once."
pushT_new	"Push the T-shaped block to its target pose marked on the table."

Recipe

• Action dim 8 (7-D Δpose + binarized gripper appended via converter scripts/convert_realworld_hf_to_robomimic.py --gripper_in_action).
• LoRA r=32 (vision/expert) + r=24 (PaliGemma language); see config.yaml.
• Single-GPU H100, bs=32 × grad_accum=2 (effective 64), lr=1.5e-4, num_epochs=20, checkpoint_every=1.
• DataLoader: num_workers=4, persistent_workers=True.
• train_pi05_image_workspace.py from KuanchengWang/diffusion_policy@jinglin.

Layout

<task>/
  checkpoints/
    epoch=XXXX-val=Y.YYYY.ckpt        # top-k by val_loss
    milestone_epochXXXX.ckpt          # every 2 epochs
    latest.ckpt                       # last-completed epoch
  config.yaml                         # hydra config for this run

Loading

import torch
ckpt = torch.load("can/checkpoints/epoch=0019-val=0.0296.ckpt",
                  map_location="cpu", weights_only=False)
# ckpt = {"cfg": OmegaConf, "state_dicts": {"model": ..., "ema_model": ...},
#         "pickles": {"global_step": ..., "epoch": ...}}

Each .ckpt contains the LoRA-only state dict (~29 MB). The base PaliGemma weights stay loaded from HuggingFace at eval time.