Project-IRA/TPSoSe2026_Pi05_LeRobot_SO101_Finetuning_V7_Full_V2

Pi0.5 SO-101 Multi-Task — V7 Full V2 (recommended)

This is the recommended Project-IRA model. Pi0.5 fine-tuned on all four tasks with image augmentation enabled. Works really well on the physical arm.

Demo

Pi0.5 sorting lego bricks onto colour-matched plates on the physical SO-101.

Part of Project-IRA — Interactive Robotic Arm. Code: https://github.com/Project-IRA/interactive-robotic-arm

Base model	lerobot/pi05_base
Robot	SO-101 follower (6-DOF)
Training data	Project-IRA/TPSoSe2026_Dataset_Full_Merged_Final_LeRobot_SO101_V1
Recommended checkpoint	008000
Inputs	desk_view + wrist_left camera images, 6-dim joint state, English instruction
Outputs	6-dim continuous action chunks

Quality

Works really well across all four tasks. Best model produced by the project.

Checkpoint 008000 is the recommended one. Consistent with the V7 run, later checkpoints degrade — checkpoints beyond roughly step 20000 were unusable on the real robot. More steps is not better here.

Training

Run name pi05_6gpu_fsdp_V2 (SLURM job 2166960). 6x L40S, FSDP FULL_SHARD — plain DDP OOMs, because the optimizer state is replicated per rank.

Setting	Value
Base	lerobot/pi05_base
Dataset	930-episode merged set
Steps	30000 target, --save_freq=2000; checkpoints 002000-016000 uploaded
Batch size	32 per GPU (effective 192 across 6 GPUs)
Precision	--policy.dtype=float32 + FSDP mixed_precision: bf16
Gradient checkpointing	on
Image augmentation	--dataset.image_transforms.enable=true — the change vs V7
Vision encoder	unfrozen
Parameters	4,143,404,816 — all trainable (num_learnable_params == num_total_params)
Optimizer	AdamW, betas (0.9, 0.95), grad clip 1.0
LR schedule	peak 2.5e-05, 1000 warmup steps, cosine decay to 2.5e-06 over 30000 steps
Normalization	ACTION: MEAN_STD, STATE: MEAN_STD, VISUAL: IDENTITY
Camera keys	native wrist_left / desk_view (no rename_map)
SLURM	--gres=gpu:L40S:6 --cpus-per-task=96 --mem=540G --time=96:00:00

--policy.dtype=float32 is required under FSDP; bf16 comes from the accelerate config's mixed_precision instead. Setting the policy dtype to bfloat16 directly breaks FSDP here.

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

accelerate launch --config_file ~/fsdp_config.yaml $(which lerobot-train) \
    --dataset.repo_id=TPSoSe2026_Dataset_Full_Merged_Final_LeRobot_SO101 \
    --dataset.image_transforms.enable=true \
    --policy.type=pi05 \
    --policy.pretrained_path=lerobot/pi05_base \
    --policy.normalization_mapping='{"ACTION": "MEAN_STD", "STATE": "MEAN_STD", "VISUAL": "IDENTITY"}' \
    --policy.compile_model=false \
    --policy.gradient_checkpointing=true \
    --policy.dtype=float32 \
    --policy.freeze_vision_encoder=false \
    --policy.train_expert_only=false \
    --policy.device=cuda \
    --policy.push_to_hub=false \
    --batch_size=32 --steps=30000 --save_freq=2000 \
    --num_workers=16 --tolerance_s=0.01 \
    --output_dir=outputs/train/pi05_6gpu_fsdp_V2 \
    --job_name=pi05_fine_full_V2 \
    --wandb.enable=false

The FSDP config (~/fsdp_config.yaml):

compute_environment: LOCAL_MACHINE
distributed_type: FSDP
downcast_bf16: 'no'
fsdp_config:
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_backward_prefetch: BACKWARD_PRE
  fsdp_cpu_ram_efficient_loading: true
  fsdp_forward_prefetch: false
  fsdp_offload_params: false
  fsdp_sharding_strategy: FULL_SHARD
  fsdp_state_dict_type: SHARDED_STATE_DICT
  fsdp_sync_module_states: true
  fsdp_use_orig_params: true
mixed_precision: bf16
num_machines: 1
num_processes: 6

Pi0.5 requires lerobot[pi] installed from GitHub main, not the PyPI release.

Usage

Unusual repository layout — from_pretrained("Project-IRA/...") on the repo ID will not work. Model files are nested under outputs_V8/; that folder name comes from the training run and deliberately does not match the repo's V7_Full_V2 name.

outputs_V8/train/pi05_6gpu_fsdp_V2/checkpoints/<step>/pretrained_model/   <- weights
outputs_V8/train/pi05_6gpu_fsdp_V2/checkpoints/<step>/training_state/     <- resume only

Checkpoints present: 002000, 004000, 006000, 008000, 010000, 012000, 014000, 016000. The run itself went to 30000 steps, but only these eight were uploaded — the later ones were not useful (see Quality).

Fetch just the recommended checkpoint's weights (~13 GB instead of 298 GB):

hf download Project-IRA/TPSoSe2026_Pi05_LeRobot_SO101_Finetuning_V7_Full_V2 \
  --include 'outputs_V8/train/pi05_6gpu_fsdp_V2/checkpoints/008000/pretrained_model/*' \
  --local-dir ./pi05_v7v2

The repo totals ~298 GB because training_state/ (optimizer moments, scheduler, RNG) is published beside every checkpoint. You do not need it for inference — the --include filter above skips it.

from lerobot.policies.pi0.modeling_pi0 import PI0Policy   # pi05 shares the PI0 module

policy = PI0Policy.from_pretrained("<local path to the checkpoint's pretrained_model/>")
policy = policy.to("cuda").eval()

Inference dtype: checkpoints are saved from a bfloat16 training run. If you hit GPU OOM at inference, confirm the loaded policy is in bfloat16 and not silently upcast to float32.

Pi0.5 is ~4B parameters. On-robot inference from the robot PC is impractical; we served it over the asynchronous gRPC inference server shipped in the code repo (https://github.com/Project-IRA/interactive-robotic-arm) and ran the robot as a thin client.

Robot setup

Robot	SO-101 follower arm (6-DOF), robot_type: so_follower
Teleoperation	SO-101 leader arm
Control frequency	30 fps
State / action space	6-dim: shoulder_pan.pos, shoulder_lift.pos, elbow_flex.pos, wrist_flex.pos, wrist_roll.pos, gripper.pos
Camera observation.images.desk_view	800x600, h264 (recording)
Camera observation.images.wrist_left	640x480, h264 (recording)

Inference note: both cameras are run at 640x480 during inference, not at their recording resolutions, to reduce the payload sent to the inference server.

Environment notes

All training ran on a SLURM cluster with L40S GPUs. Two environment details were required and are easy to miss when reproducing:

• ffmpeg libraries for torchcodec. A minimal conda env supplies the shared libraries that torchcodec discovers at runtime: export LD_LIBRARY_PATH=$CONDA_PREFIX/envs/ffmpeg_libs_v8/lib:<venv>/lib/python3.12/site-packages/nvidia/npp/lib:$LD_LIBRARY_PATH
• --tolerance_s=0.01 on every run, to accommodate timestamp jitter in the recorded episodes.

Multi-GPU runs additionally set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True. Datasets and the virtualenv were copied to node-local /scratch before training rather than read from shared storage.

No Weights & Biases logging was enabled for any run (--wandb.enable=false), so there are no public training curves — the job.*.err SLURM logs are the record.

Tasks and prompts

The model is conditioned on English natural-language instructions. Prompt phrasing was varied roughly every 10 episodes during recording, giving 93 distinct prompts in the merged dataset. Use one of the training prompts verbatim for best results — the full lists are on the dataset card.

Limitations

• Behaviour cloning. The policy imitates teleoperated demonstrations and has no notion of recovery beyond what was demonstrated. It is susceptible to covariate shift and can fail to recover from states outside the demonstration distribution.
• Recovery data is incidental, not systematic. Recovery behaviour appears in the data only where the operator happened to make and correct a mistake during recording; no recovery episodes were scripted deliberately.
• Single environment. All data comes from one lab desk with one lighting setup, one camera geometry, and one set of physical objects. Expect degradation elsewhere.
• Prompt sensitivity. Language conditioning was trained on a fixed set of phrasings (listed in the dataset card). Prompts far from those phrasings may behave unpredictably.
• No formal evaluation. Quality assessments below are qualitative, from operators observing rollouts on the physical arm. There are no success-rate numbers.
• Not safety-rated. Supervise all physical execution and keep the workspace clear.

Upstream licensing & attribution

This model is a derivative work of Apache-2.0 licensed components:

Component	Upstream	License
LeRobot framework	https://github.com/huggingface/lerobot	Apache-2.0
lerobot/pi05_base (Physical Intelligence, openpi)	https://github.com/Physical-Intelligence/openpi	Apache-2.0

Apache-2.0 permits relicensing derivative works. We retain the upstream copyright notices, license text, and NOTICE files for the incorporated material, as Apache-2.0 Section 4 requires. The upstream components remain under Apache-2.0 — only this project's own contributions (the fine-tuned weights and training configuration) are offered under CC BY-SA 4.0.

CC BY-SA 4.0 was chosen because it is share-alike: derivatives must be released under the same licence, so this work cannot be taken closed-source. The project's source code lives in a separate repository under its own licence — see https://github.com/Project-IRA/interactive-robotic-arm.

Citation

@misc{project_ira_2026,
  title        = {Project-IRA: Interactive Robotic Arm},
  author       = {Baten, Cleo and Keppler, Bela and Sapper, Jonas},
  year         = {2026},
  howpublished = {\url{https://huggingface.co/Project-IRA}},
  note         = {Code: \url{https://github.com/Project-IRA/interactive-robotic-arm}}
}