lerobot044_pi05_train_env

lerobot 0.4.4 에서 pi0.5 (pi05) 정책을 파인튜닝하기 위한 환경 설치 가이드. A100 등 새 서버에서 그대로 재현해 여러 task를 학습할 수 있도록 일반화한 문서입니다.

검증 환경: 2× RTX A6000 (48GB), CUDA driver 12.x, conda. torch 2.10 (cu12.8) 휠은 A100(sm_80)에서도 그대로 동작합니다.

---

⚠️ 핵심 함정 2가지 (둘 다 안 하면 학습이 즉시 실패)

1. 패치된 transformers가 필요하다. 그냥 pip install lerobot==0.4.4 만 하거나 lerobot[transformers-dep] 로 설치하면 pi05 생성 시 ValueError: An incorrect transformer version is used 로 죽습니다. pi0/pi05 는 lerobot 이 패치한 transformers 브랜치(fix/lerobot_openpi)를 요구합니다 → pi extra 로 설치.
2. lerobot/pi05_base 리비전을 고정해야 한다. 허브의 최신(main) 전처리기는 relative_actions_processor 스텝을 쓰는데 이건 lerobot 0.4.4 에 존재하지 않습니다 (나중 버전에서 추가됨). 그대로 쓰면 KeyError: Processor step 'relative_actions_processor' not found in registry 로 죽습니다. → 0.4.4 시기의 리비전 9e55186ad3 을 받아서 로컬 경로로 지정.

---

1. Conda 환경 생성

conda create -n lerobot044 python=3.10 -y
conda activate lerobot044

2. lerobot 0.4.4 + pi extra 설치 (핵심)

# lerobot core + 패치 transformers(fix/lerobot_openpi) + scipy 를 한 번에 설치
pip install "lerobot[pi]==0.4.4"

pi extra 정의(lerobot 0.4.4 pyproject.toml):

pi = ["transformers @ git+https://github.com/huggingface/transformers.git@fix/lerobot_openpi",
      "scipy>=1.10.1,<1.15"]

설치 시 git + 네트워크 필요(transformers 를 소스에서 빌드). 몇 분 걸립니다.

이미 다른 transformers 가 깔린 환경이라면 아래로 강제 교체할 수 있습니다:

pip install "transformers @ git+https://github.com/huggingface/transformers.git@fix/lerobot_openpi" "scipy>=1.10.1,<1.15"

3. 설치 검증

python - <<'PY'
import lerobot, torch, transformers
print("lerobot", lerobot.__version__)          # 0.4.4
print("torch", torch.__version__, "cuda", torch.cuda.is_available(), "ndev", torch.cuda.device_count())
print("transformers", transformers.__version__) # 4.53.3 (패치 버전)
from lerobot.policies.pi05.modeling_pi05 import PI05Policy
from transformers.models.siglip import check
print("siglip check ->", check.check_whether_transformers_replace_is_installed_correctly())  # True 여야 함
PY

siglip check -> True 가 나오면 함정 #1 통과.

4. pi05_base 리비전 고정 다운로드 (핵심)

# 0.4.4 호환 리비전(relative_actions_processor 추가 이전)을 로컬로 받는다
python - <<'PY'
from huggingface_hub import snapshot_download
p = snapshot_download("lerobot/pi05_base", revision="9e55186ad3",
                      local_dir="/DATA/pi05_base_9e55186")  # 로컬 nvme 경로로 변경
print("DONE", p)
PY

• 리비전 9e55186ad3 (2026-01-22) = 커밋 7de663972b("Add relative action processor steps", 2026-06-03) 직전.
• 이후 학습에서 --policy.pretrained_path=/DATA/pi05_base_9e55186 로 이 로컬 경로를 지정.
• (참고) 이 리비전 전처리기 스텝: rename_observations_processor, to_batch_processor, normalizer_processor, pi05_prepare_state_tokenizer_processor_step, tokenizer_processor, device_processor — 전부 0.4.4 레지스트리에 존재.

---

5. 학습 명령 (일반 템플릿, 멀티-GPU)

여러 task에 재사용하도록 변수만 바꾸면 됩니다. DDP는 accelerate launch 사용.

DATAROOT=/DATA/datasets/<task>          # LeRobotDataset v3.0 로컬 경로 (meta/info.json 포함)
OUTDIR=/DATA/outputs/pi05_<task>
BASE=/DATA/pi05_base_9e55186            # 4단계에서 받은 로컬 pi05_base

export CUDA_VISIBLE_DEVICES=0,1        # A100 여러 장이면 0,1,2,3 ...
NPROC=$(echo "$CUDA_VISIBLE_DEVICES" | tr ',' '\n' | wc -l)
EFFECTIVE_BATCH=32                      # 총(effective) 배치. per-GPU = EFFECTIVE_BATCH/NPROC
BATCH=$(( EFFECTIVE_BATCH / NPROC ))

accelerate launch --num_processes="$NPROC" --mixed_precision=bf16 \
  "$(which lerobot-train)" \
  --dataset.repo_id=<hub/label-only-ok> \
  --dataset.root="$DATAROOT" \
  --dataset.video_backend=pyav \
  --policy.type=pi05 \
  --policy.pretrained_path="$BASE" \
  --policy.device=cuda \
  --policy.dtype=bfloat16 \
  --policy.compile_model=true \
  --policy.gradient_checkpointing=true \
  --policy.freeze_vision_encoder=false \
  --policy.train_expert_only=false \
  --policy.tokenizer_max_length=200 \
  --policy.push_to_hub=false \
  --steps=30000 --save_freq=5000 \
  --batch_size="$BATCH" --num_workers=8 --seed=1000 \
  --output_dir="$OUTDIR" --job_name=pi05_<task> --wandb.enable=false

배치 규칙 (중요)

lerobot DDP 에서 --batch_size 는 프로세스(GPU)당 값입니다. effective 배치 = batch_size × GPU 수. 예) 총 배치 32 를 2-GPU로 → --batch_size=16. A100 80GB 는 여유가 많으니 per-GPU 배치를 올려도 됩니다.

하이퍼파라미터 (pi05 기본값 = 아래와 동일, 별도 지정 불필요)

• optimizer: AdamW, lr 2.5e-5, betas (0.9, 0.95), wd 0.01, grad_clip 1.0
• scheduler: cosine, warmup 1000, decay 30000 → 2.5e-6
• chunk_size / n_action_steps = 50, n_obs_steps = 1
• normalization: STATE/ACTION = QUANTILES, VISUAL = IDENTITY

위 값은 pi05_base config 에서 상속되므로 명령에 넣지 않아도 재현됩니다. --steps 를 바꾸면 scheduler decay 도 같이 맞추고 싶을 때만 --policy.scheduler_decay_steps=<steps> 를 추가하세요.

---

6. 데이터셋 요구사항

• LeRobotDataset v3.0 (meta/info.json 의 codebase_version: v3.0).
• 카메라/상태/액션 feature 는 policy 입력과 맞아야 함. (예시로 검증한 UR7e abs_ee 셋: observation.images.third_person, observation.images.eye_in_hand (3,480,640), observation.state (24,), action (22,), 50 fps.)
• 액션/상태 차원은 pi05 의 max_state_dim=32, max_action_dim=32 이내면 됨(패딩 처리).
• 로컬 nvme 에 두고 --dataset.root 로 직접 가리키는 걸 권장(느린 네트워크 스토리지 회피).

---

7. 트러블슈팅

증상	원인	해결
ValueError: An incorrect transformer version is used	패치 안 된 transformers	lerobot[pi] 로 설치(함정 #1)
KeyError: 'relative_actions_processor' not found in registry	pi05_base main 리비전이 0.4.4 와 불일치	리비전 9e55186ad3 사용(함정 #2)
첫 스텝 전 AUTOTUNE/triton_mm... 로그가 한참 나옴	compile_model=true 의 torch.compile 워밍업(정상)	기다리기. 로그가 Training: N/30000 loss=... 로 바뀌면 실제 스텝 시작. 싫으면 --policy.compile_model=false
CUDA out of memory	per-GPU 배치/모델이 큼	EFFECTIVE_BATCH 낮추기, GPU 늘리기, gradient_checkpointing=true 유지, 필요시 compile_model=false

부록: 검증된 패키지 버전

lerobot        0.4.4
transformers   4.53.3   (git: fix/lerobot_openpi, 패치)
tokenizers     0.21.4
scipy          1.14.1
torch          2.10.0+cu128
accelerate     1.15.0
av (pyav)      15.1.0
torchcodec     0.10.0
python         3.10