yrpark/pi05-ur7e-pour_cup-train-env

lerobot π₀.₅ (pi05) fine-tuning 환경 — DexSteer/isaaclab-ur7e-pour_cup

UR7e + RH5D 핸드(Isaac Lab) 데이터셋으로 π₀.₅ 를 fine-tuning 하기 위한 재현 가능한 설치/학습 스크립트입니다. GPU 1개(A100 80GB)와 GPU 2~4개(DDP) 환경 모두 같은 스크립트로 동작합니다.

검증된 조합 (2026-09-07, Ubuntu 24.04, NVIDIA driver 575 / CUDA 12.9, A100-80GB):

항목	값
Python	3.12 (uv venv)
lerobot	0.6.1 (PyPI)
torch / torchvision / torchcodec	2.11.0+cu128 / 0.26.0 / 0.11.1
transformers / accelerate	5.5.4 / 1.14.0
설치 시간	uv 기준 약 40초 (torch 휠 다운로드 포함)

정확한 전체 버전은 requirements.lock.txt 에 있습니다.

---

0. 가장 빠른 경로 (요약)

# 1) 이 폴더를 영구 저장소로 복사 (예: /data/pi05-ur7e). /home/work 는 세션 종료 시 삭제됨.
git clone <이 저장소> /data/pi05-ur7e   # 또는 hf download yrpark/pi05-ur7e-pour_cup-train-env --local-dir /data/pi05-ur7e
cd /data/pi05-ur7e

# 2) 설치 (uv 자동 설치 → Python 3.12 venv → lerobot[pi,training]==0.6.1 + torch cu128 → 검증)
bash setup.sh
bash setup_extra.sh    # torchcodec 활성화(시스템 ffmpeg + nvidia-npp-cu12 + env.sh LD_LIBRARY_PATH). 생략 시 pyav 로 자동 대체

# 3) HF 로그인 (PaliGemma 토크나이저가 gated 라서 반드시 필요)
#    https://huggingface.co/google/paligemma-3b-pt-224 에서 라이선스 동의 후
hf auth login            # 또는 export HF_TOKEN=hf_xxx  (env.sh 가 로그인 토큰을 자동으로 찾습니다)

# 4) 데이터셋(0.33 GB) + pi05_base(14 GB) + 토크나이저 미리 받기
bash download.sh

# 5) 학습
bash train.sh                 # GPU 1개
NUM_GPUS=4 bash train.sh      # GPU 4개 (DDP)

setup.sh 만 GPU 없이도 동작하고, 나머지는 source env.sh 로 venv 가 활성화된 셸에서 실행됩니다.

---

1. 파일 구성

파일	역할
env.sh	공통 환경 변수. WORKDIR, HF 캐시 경로, PYTHONPATH 제거, venv 활성화
setup.sh	uv 설치 → venv 생성 → lerobot 설치 → scripts/check_env.py 로 검증
setup_extra.sh	torchcodec 활성화용 추가 설치: 시스템 ffmpeg + nvidia-npp-cu12(libnppicc.so.12) + env.sh 에 NPP LD_LIBRARY_PATH 주입. torch 2.11+cu128 휠에 NPP 가 없어 torchcodec cu 빌드가 로드 실패하는 문제를 해결. setup.sh 다음에 실행. 없으면 VIDEO_BACKEND=pyav 로 대체
download.sh	데이터셋 / lerobot/pi05_base / PaliGemma 토크나이저 사전 다운로드 및 디코딩 테스트
train.sh	학습 실행. GPU 수를 감지해 단일 GPU / accelerate launch DDP 자동 분기. RESUME=1 지원. 로그는 logs/<job>.log
scripts/check_env.py	버전·CUDA·비디오 백엔드·HF 토큰 점검
scripts/smoke_test.py	토크나이저 없이 pi05_base 를 로드해 forward/backward 메모리·속도 측정
upload_to_hf.sh	이 폴더를 HF Hub 에 업로드 (.venv, cache, outputs 제외)
requirements.lock.txt	uv pip freeze 결과. LOCK=1 bash setup.sh 로 완전 동일 재현

---

2. 왜 이 구성이 가장 빠른가

1. uv + PyPI 릴리스 고정 — conda 보다 훨씬 빠르고(수십 초), lerobot==0.6.1 로 버전이 고정되어 다른 머신에서도 같은 결과가 나옵니다. torch 는 --torch-backend=cu128 로 CUDA 휠을 직접 받습니다.
2. NGC/Backend.AI 컨테이너의 시스템 torch 를 쓰지 않음 — 컨테이너의 torch 2.7.0a0 / numpy 1.26 / transformers 4.55 는 lerobot 0.6.1 요구사항(numpy ≥ 2, transformers ≥ 5.4)과 충돌합니다. env.sh 가 PYTHONPATH 를 비워 시스템 site-packages 가 venv 로 새어 들어오지 않게 합니다.
3. 데이터셋이 이미 pi05 호환 — LeRobot v3.0 포맷, meta/stats.json 에 q01/q99 가 있어 π₀.₅ 기본 QUANTILES 정규화가 바로 됩니다. 재변환·재통계 불필요.
4. 캐시를 WORKDIR/cache 로 고정 — 영구 마운트에 두면 세션이 바뀌어도 14 GB 모델을 다시 받지 않습니다.
5. 비디오 디코딩 — AV1(dav1d) 영상이며 torchcodec(시스템 ffmpeg 6.1) 과 pyav 둘 다 검증됨. 8 worker 기준 torchcodec ≈ 2,000 samples/s, pyav ≈ 1,550 samples/s 라 학습 병목이 아닙니다.

---

3. 데이터셋 요약

항목	값
포맷	LeRobot v3.0, 150 episodes, 69,642 frames, 50 fps, 0.33 GB
카메라	observation.images.third_person, observation.images.eye_in_hand (480×640, AV1) → 모델 내부에서 224×224 로 리사이즈
state	24-dim (UR7e 6 관절 + RH5D 핸드 18 관절)
action	19-dim
task	"Grasp the cup, lift it, and pour it out." (단일 task)

주의: action 차원 0~5 는 전 프레임에서 상수 0, 차원 6 은 상수 1.5 입니다. 실제로 변하는 값은 차원 7~18(핸드 12개)뿐입니다. 즉 이 데이터의 action 에는 팔(UR7e) 명령이 들어 있지 않습니다. 학습 자체는 문제 없이 되지만(상수 차원은 정규화에서 eps 로 처리되고 loss 가 0 에 수렴), 배포 시 팔 제어가 필요하다면 데이터셋 제작 측에 action 정의를 확인하세요. download.sh 실행 후 python -c 로 cache/lerobot/.../meta/stats.json 을 열어보면 확인할 수 있습니다.

---

4. 학습 설정 설명 (train.sh)

기본 명령(단일 GPU)은 lerobot 공식 π₀.₅ 가이드의 80 GB 설정을 이 데이터셋에 맞춘 것입니다:

lerobot-train \
  --dataset.repo_id=DexSteer/isaaclab-ur7e-pour_cup \
  --dataset.video_backend=torchcodec \
  --policy.type=pi05 \
  --policy.pretrained_path=lerobot/pi05_base \
  --policy.normalization_mapping='{"ACTION": "QUANTILES", "STATE": "QUANTILES", "VISUAL": "IDENTITY"}' \
  --policy.n_action_steps=25 --policy.empty_cameras=1 \
  --policy.freeze_vision_encoder=false --policy.train_expert_only=false \
  --policy.gradient_checkpointing=true --policy.dtype=bfloat16 --policy.device=cuda \
  --policy.push_to_hub=false \
  --optimizer.lr=2.5e-5 --batch_size=32 --num_workers=8 --steps=30000 \
  --save_freq=2500 --log_freq=50 --seed=1000 \
  --output_dir=outputs/pi05_ur7e_pour_cup --job_name=pi05_ur7e_pour_cup

• --policy.pretrained_path 는 가중치만 로드하고 feature 이름/차원은 데이터셋에서 가져옵니다. 그래서 rename_map 이 필요 없고, state 24 / action 19 는 내부에서 32 로 패딩됩니다. (--policy.path 와 다름 — 그쪽은 config 까지 상속)
• --policy.empty_cameras=1 — π₀.₅ 는 3 카메라 슬롯으로 사전학습되었으므로 2 카메라 데이터에 빈 슬롯 1개를 마스킹으로 채웁니다.
• --policy.n_action_steps=25 — 추론 시 실행하는 step 수(50 fps → 0.5 s). 학습 결과에는 영향 없음. pretrained_path 를 쓰면 기본값(50)으로 리셋되므로 명시.
• 정규화: 기본 QUANTILES. 이 데이터셋의 q01/q99 는 사실상 min/max 와 같아(기록 시 per-episode 집계) min-max 스케일링과 동일하게 동작합니다. 이상치가 걱정되면 NORM_MODE=mean_std bash train.sh.
• 이 데이터셋에서 --rename_map 은 쓰지 마세요. lerobot 0.6.1 에서 pretrained_path + rename_map 조합은 첫 배치에서 "All image features are missing" 으로 실패합니다.

GPU 메모리 / 속도 (A100 80GB, bf16, scripts/smoke_test.py 측정)

batch / GPU	grad ckpt	peak mem	step 시간
16	on	37.6 GiB	2.7 s
32	on	45.4 GiB	5.1 s
64	on	60.7 GiB	9.9 s
16	off	OOM (80 GB)	-

→ 80 GB 카드는 BATCH_SIZE=32(기본) 또는 64, 40/48 GB 카드는 BATCH_SIZE=16 (+ 필요 시 TRAIN_EXPERT_ONLY=true). gradient checkpointing 은 끄지 마세요. 끄면 80 GB 에서도 batch 16 이 OOM 납니다(4.1B 파라미터 + AdamW 상태만 ~33 GiB). 실제 lerobot-train 30 step 검증(batch 16, 토크나이저 포함): 2.47 s/step, 37 GiB, loss 0.70 → 0.14, 체크포인트 저장 정상. 단일 GPU 로 30k step(batch 32) 은 약 42 h 이므로 다중 GPU 권장.

다중 GPU (2~4개)

NUM_GPUS=4 bash train.sh            # accelerate launch --multi_gpu --num_processes=4 $(which lerobot-train) ...
NUM_GPUS=2 BATCH_SIZE=64 bash train.sh

• lerobot 0.6.1 은 accelerate launch 를 launcher 로만 사용하고 mixed precision 은 --policy.dtype=bfloat16 에서 자동으로 bf16 이 됩니다. accelerate config 를 돌릴 필요 없습니다.
• 유효 배치 = BATCH_SIZE × NUM_GPUS. lerobot 은 lr/step 을 자동 스케일하지 않으므로 train.sh 는 STEPS 를 30000 / NUM_GPUS 로 기본 축소합니다(4 GPU → 7,500 step, 유효 배치 128, ≈ 14 epoch). 총 샘플 수는 단일 GPU 와 동일하고 wall-clock 은 약 1/NUM_GPUS. 필요하면 STEPS=, LR= 로 직접 지정하세요.
• π₀.₅ 스케줄러(warmup 1k, cosine decay 30k)는 --steps 가 30k 보다 작으면 자동으로 비례 축소됩니다.
• 체크포인트는 main process 만 저장하며 outputs/<job>/checkpoints/<step>/pretrained_model/ 에 lerobot-eval, PI05Policy.from_pretrained 로 바로 쓸 수 있는 형태로 남습니다.
• 같은 노드에서 두 job 을 동시에 돌리면 MASTER_PORT=29501 처럼 포트를 바꾸세요.
• 4 GPU 에서 NCCL 이 hang 하면 NCCL_P2P_DISABLE=1 또는 NCCL_IB_DISABLE=1 을 시도하세요 (컨테이너 환경에서 흔한 문제).

재개 / 기타

RESUME=1 bash train.sh                       # outputs/<job>/checkpoints/last 에서 재개 (같은 JOB_NAME 으로 새로 시작하면 덮어쓰기 방지로 거부됨)
JOB_NAME=exp2 TRAIN_EXPERT_ONLY=true bash train.sh   # 별도 job, VLM 동결 + action expert 만 학습
WANDB_ENABLE=true WANDB_PROJECT=pi05 bash train.sh
EXTRA_ARGS="--policy.rtc_training_max_delay=10" bash train.sh   # 학습-시 RTC (real-time chunking)
COMPILE=true bash train.sh                   # torch.compile (첫 step 컴파일 수 분 소요)

---

5. 다른 머신에 설치할 때 체크리스트

1. 영구 경로 — WORKDIR 를 마운트된 데이터 폴더로: export WORKDIR=/data/pi05-ur7e (기본은 스크립트 폴더). 캐시·출력이 모두 그 아래로 갑니다.
2. CUDA 휠 — 드라이버가 CUDA 12.8 미만이면 TORCH_BACKEND=cu126 bash setup.sh, 최신이면 cu130 또는 auto.
3. ffmpeg — apt install ffmpeg (Ubuntu 22.04/24.04 의 ffmpeg 는 libdav1d 포함). 설치 불가하면 VIDEO_BACKEND=pyav bash train.sh (pyav 는 자체 ffmpeg 를 내장, 약 25% 느리지만 충분).
4. Python 3.12 — 시스템에 없어도 uv venv --python 3.12 가 자동 다운로드합니다.
5. HF 토큰 — PaliGemma 라이선스 동의 + hf auth login(또는 HF_TOKEN). env.sh 가 ~/.cache/huggingface/token 을 HF_HOME 으로 복사해 줍니다. 없으면 download.sh/train.sh 가 토크나이저 단계에서 실패합니다. 토큰이 read 권한만 있어도 학습은 가능하고, upload_to_hf.sh 에는 write 가 필요합니다.
6. 오프라인 노드 — 온라인 노드에서 bash download.sh 를 먼저 돌리고 cache/ 를 통째로 복사한 뒤 HF_HUB_OFFLINE=1 bash train.sh.
7. lerobot main 브랜치(--parallelism.dp_shard 등 FSDP 신기능)가 필요하면 LEROBOT_SOURCE=git bash setup.sh. 단, 그 경우 다중 GPU 실행 방식이 torchrun/--parallelism.* 로 바뀌므로 train.sh 의 accelerate 분기를 수정해야 합니다.

---

6. 검증 방법

source env.sh
python scripts/check_env.py                  # 버전/CUDA/디코더/토큰
python scripts/smoke_test.py --batch-sizes 16,32   # 토크나이저 없이 GPU 메모리 확인 (모델 로드 ~2.5분)
DRY_RUN=1 NUM_GPUS=4 bash train.sh           # 실제 명령만 출력

7. 학습 후

# 체크포인트 로드 (config 포함)
lerobot-eval --policy.path=outputs/pi05_ur7e_pour_cup/checkpoints/last/pretrained_model ...
# Hub 업로드
hf upload yrpark/pi05-ur7e-pour_cup outputs/pi05_ur7e_pour_cup/checkpoints/last/pretrained_model .

참고 문서: lerobot docs/source/pi05.mdx, multi_gpu_training.mdx, rename_map.mdx (v0.6.1 태그 기준).