JackieMM/robotwin-pi05-arx-training

Train pi0.5 on another server

This directory is the reproducible training entry point for the ARX X5 real-robot data. It is separate from the ROS/CAN deployment tools. The default command downloads the complete 71-episode two-camera dataset and starts LoRA fine-tuning; it does not connect to a robot.

1. Prepare OpenPI

Use the pinned RoboTwin/XPolicyLab checkout used by this project. The helper below clones the public RoboTwin repository and its pinned XPolicyLab submodule, then installs OpenPI with uv:

./setup_openpi.sh /data/robotwin_arx

The resulting OpenPI path is /data/robotwin_arx/RoboTwin/XPolicyLab/policy/Pi_05/openpi. The machine needs Python 3.11+, uv, NVIDIA drivers, and a CUDA 12.x runtime. For a 24 GB RTX 4090, start with one GPU, batch size 1, and XLA_PYTHON_CLIENT_ALLOCATOR=platform.

2. Train from Hugging Face

Authenticate with a read token (do not put it in a repository file), then run:

hf auth login
./train_from_hf.sh \
  --openpi-root /data/robotwin_arx/RoboTwin/XPolicyLab/policy/Pi_05/openpi \
  --work-dir /data/arx_pi05 \
  --dataset-id JackieMM/LQ-20260915-1-all71-2cam \
  --gpu 0 \
  --steps 30000 \
  --save-interval 5000 \
  --keep-period 5000

The script downloads the LeRobot dataset, copies its matching norm_stats.json, loads the pi0.5 DROID parameters, and invokes train_pi05_lora.py with only cam_high and cam_right_wrist. It refuses to overwrite an existing checkpoint directory. To resume an interrupted run, pass --resume and use a new invocation with the same work directory; the checkpoint directory is then reused.

For two GPUs, pass --gpu 0,1 --fsdp-devices 2. FSDP can reduce per-GPU memory but may be slower than one GPU because of inter-GPU communication; use one GPU unless memory requires sharding.

Dataset choices

• Full data: JackieMM/LQ-20260915-1-all71-2cam (71 episodes, 25,920 frames).
• First 50: JackieMM/LQ-20260915-1-first50-2cam (50 episodes, 18,191 frames).

The action/state ABI is 14-dimensional. The left-wrist image is intentionally absent; OpenPI masks it instead of duplicating another camera. Checkpoint files are local artifacts and are not included in the source repository.