OpenPi训练标准化流程Docker封装
官方代码拉取
拉取官方代码
git clone --recurse-submodules git@github.com:Physical-Intelligence/openpi.git
# Or if you already cloned the repo:
git submodule update --init --recursive
新建docker文件夹，放置以下内容
Dockerfile
Dockerfile主要训练环境按照官方要求封装，使得不同机器、集群分发的训练环境一致。
# syntax=docker/dockerfile:1.7
FROM nvidia/cuda:12.2.2-cudnn8-runtime-ubuntu22.04
ARG DEBIAN_FRONTEND=noninteractive
ARG INSTALL_RLDS=0
# -----------------------------
# apt 源
# -----------------------------
RUN sed -i \
-e 's@archive.ubuntu.com@mirrors.tuna.tsinghua.edu.cn@g' \
-e 's@security.ubuntu.com@mirrors.tuna.tsinghua.edu.cn@g' \
/etc/apt/sources.list
RUN apt-get update && \
apt-get install -y --no-install-recommends \
git \
git-lfs \
ca-certificates \
curl \
build-essential \
clang \
linux-headers-generic \
ffmpeg \
libgl1 \
libglib2.0-0 \
libsm6 \
libxext6 \
libxrender1 && \
git lfs install --system && \
rm -rf /var/lib/apt/lists/*
# -----------------------------
# uv
# -----------------------------
ENV UV_INSTALL_DIR=/usr/local/bin
ENV UV_PYTHON_INSTALL_MIRROR=https://mirrors.ustc.edu.cn/github-release/astral-sh/python-build-standalone/
ENV UV_DOWNLOAD_URL=https://mirrors.ustc.edu.cn/github-release/astral-sh/uv/LatestRelease/
ENV UV_INDEX_URL=https://pypi.tuna.tsinghua.edu.cn/simple
RUN curl -sL \
https://mirrors.ustc.edu.cn/github-release/astral-sh/uv/LatestRelease/uv-installer.sh \
| sh
# -----------------------------
# Python 环境
# -----------------------------
ENV UV_LINK_MODE=copy
ENV VIRTUAL_ENV=/workspace/openpi/.venv
ENV PATH="${VIRTUAL_ENV}/bin:${PATH}"
ENV PYTHONUNBUFFERED=1
ENV PYTHONDONTWRITEBYTECODE=1
WORKDIR /workspace/openpi
COPY pyproject.toml uv.lock ./
RUN uv venv --python 3.11
RUN uv pip install setuptools wheel
# -----------------------------
# 安装依赖
# -----------------------------
COPY packages/openpi-client/pyproject.toml \
packages/openpi-client/pyproject.toml
COPY packages/openpi-client/src \
packages/openpi-client/src
RUN --mount=type=cache,target=/root/.cache/uv \
if [ "${INSTALL_RLDS}" = "1" ]; then \
GIT_LFS_SKIP_SMUDGE=1 \
uv sync \
--no-install-project \
--no-dev \
--group rlds; \
else \
GIT_LFS_SKIP_SMUDGE=1 \
uv sync \
--no-install-project \
--no-dev; \
fi
# -----------------------------
# openpi源码
# -----------------------------
COPY . .
RUN chmod +x docker/train_pipeline.sh
RUN --mount=type=cache,target=/root/.cache/uv \
GIT_LFS_SKIP_SMUDGE=1 \
uv pip install \
-e .
# -----------------------------
# transformers patch
# -----------------------------
RUN python - <<'PY'
from pathlib import Path
import shutil
import transformers
source = Path(
"/workspace/openpi/src/openpi/models_pytorch/transformers_replace"
)
target = Path(transformers.__file__).parent
if source.exists():
for item in source.iterdir():
destination = target / item.name
if item.is_dir():
shutil.copytree(
item,
destination,
dirs_exist_ok=True,
)
else:
shutil.copy2(item, destination)
print(f"Patched transformers at: {target}")
PY
# -----------------------------
# 数据目录
# -----------------------------
ENV OPENPI_DATA_HOME=/openpi_assets
ENV HF_HOME=/root/.cache/huggingface
RUN mkdir -p \
/data \
/openpi_assets \
/workspace/openpi/checkpoints \
/workspace/openpi/assets \
/workspace/openpi/wandb
ENTRYPOINT ["/workspace/openpi/docker/train_pipeline.sh"]
Docker Compose
Docker Compose主要指定Docker容器的网络通讯，映射目录、镜像、环境变量等标准化部署参数
name: openpi-pi05
services:
train:
image: openpi-train:cuda12.2
build:
context: ..
dockerfile: docker/train.Dockerfile
args:
INSTALL_RLDS: ${INSTALL_RLDS:-0}
init: true
stdin_open: true
tty: true
# 训练程序、数据加载器会使用共享内存
ipc: host
# W&B、Hugging Face、策略服务等访问方便
network_mode: host
working_dir: /workspace/openpi
env_file:
- path: ../.env
required: false
volumes:
# 数据集
- ${DATASET_DIR}:/data:ro
# checkpoint
- ${CHECKPOINT_DIR}:/workspace/openpi/checkpoints
# assets / norm stats
- ${ASSETS_DIR}:/workspace/openpi/assets
# pi0/pi05 权重缓存
- ${OPENPI_CACHE_DIR}:/openpi_assets
# HuggingFace缓存
- ${HF_CACHE_DIR}:/root/.cache/huggingface
# W&B日志
- ${WANDB_DIR}:/workspace/openpi/wandb
environment:
OPENPI_DATA_HOME: /openpi_assets
# hf缓存地址
HF_HOME: /root/.cache/huggingface
# lerobot数据父路径
HF_LEROBOT_HOME: /data
CUDA_VISIBLE_DEVICES: ${CUDA_VISIBLE_DEVICES:-0}
# JAX显存比例
XLA_PYTHON_CLIENT_MEM_FRACTION: ${XLA_MEMORY_FRACTION:-0.90}
WANDB_API_KEY: ${WANDB_API_KEY:-}
WANDB_MODE: ${WANDB_MODE:-online}
NVIDIA_VISIBLE_DEVICES: all
NVIDIA_DRIVER_CAPABILITIES: compute,utility
PYTHONUNBUFFERED: "1"
deploy:
resources:
reservations:
devices:
- driver: nvidia
count: all
capabilities: [gpu]
.env
训练环境变量导入Docker，模型输入输出产物存放地址在该文件中进行改动
USER_ID和GROUP_ID主要考虑模型产物权限问题
可根据宿主机查看：
id -u
id -g
# -----------------------------
# 自动归一化 + 训练 Pipeline
# -----------------------------
TRAIN_CONFIG_NAME=pi05_r2v2_custom
DATASET_REPO_ID=dual_leg_pick_bottle_new_01
R2V2_VIEW=dual
EXP_NAME=dual_leg_pick_bottle_new_01
MODEL_PI05=1
ACTION_DIM=32
OUTPUT_ACTION_DIM=16
ACTION_HORIZON=50
DISCRETE_STATE_INPUT=0
PROMPT_FROM_TASK=1
BASE_CHECKPOINT=gs://openpi-assets/checkpoints/pi05_base/params
TRAIN_BACKEND=jax
FSDP_DEVICES=1
BATCH_SIZE=32
NUM_WORKERS=4
NUM_TRAIN_STEPS=20000
SEED=42
PEAK_LR=2.5e-5
DECAY_LR=2.5e-6
WARMUP_STEPS=1000
LR_DECAY_STEPS=20000
WEIGHT_DECAY=1e-10
CLIP_GRADIENT_NORM=1.0
LOG_INTERVAL=100
SAVE_INTERVAL=1000
KEEP_PERIOD=5000
OVERWRITE=0
RESUME=0
SKIP_NORM_STATS=0
NORM_MAX_FRAMES=
WANDB_ENABLED=1
WANDB_PROJECT=openpi
PYTORCH_GPU_COUNT=1
PYTORCH_TRAINING_PRECISION=bfloat16
# 你的 LeRobot / RLDS 数据目录
DATASET_DIR=/media/eft/F6918E841F1A5101
# 训练结果
CHECKPOINT_DIR=/mnt/data/pi05_outputs/checkpoints
# norm stats 和训练 assets
ASSETS_DIR=/mnt/data/pi05_outputs/assets
# openpi 下载的 pi05_base 等模型缓存
OPENPI_CACHE_DIR=/mnt/data/model_cache/openpi
# Hugging Face 缓存
HF_CACHE_DIR=/mnt/data/model_cache/huggingface
# W&B 日志
WANDB_DIR=/mnt/data/pi05_outputs/wandb
# 单卡一般设置为0
CUDA_VISIBLE_DEVICES=0
XLA_MEMORY_FRACTION=0.90
# 普通 LeRobot 数据不需要RLDS
INSTALL_RLDS=0
WANDB_MODE=online
WANDB_API_KEY=
# proxy
HTTP_PROXY=http://127.0.0.1:7890
HTTPS_PROXY=http://127.0.0.1:7890
NO_PROXY=localhost,127.0.0.1
# 如果是 Clash Meta / v2ray 等
ALL_PROXY=http://127.0.0.1:7890
config.py注入
根据构建好的docker导入的环境变量自动注入对应训练配置文件
def _env_bool(name: str, default: bool) -> bool:
value = os.getenv(name)
if value is None:
return default
normalized = value.strip().lower()
if normalized in {"1", "true", "yes", "on"}:
return True
if normalized in {"0", "false", "no", "off"}:
return False
raise ValueError(f"{name} must be a boolean, got {value!r}")
def _env_config() -> TrainConfig | None:
"""Build an opt-in TrainConfig for the Docker training pipeline."""
name = os.getenv("TRAIN_CONFIG_NAME")
if not name:
return None
repo_id = os.getenv("DATASET_REPO_ID")
if not repo_id:
raise ValueError("DATASET_REPO_ID is required when TRAIN_CONFIG_NAME is set")
view = os.getenv("R2V2_VIEW", "dual").strip().lower()
data_factory = {
"dual": LeRobotR2v2DataConfig,
"right": LeRobotRightR2v2DataConfig,
"front": LeRobotR2v2FrontDataConfig,
}.get(view)
if data_factory is None:
raise ValueError("R2V2_VIEW must be one of: dual, right, front")
train_steps = int(os.getenv("NUM_TRAIN_STEPS", "50000"))
peak_lr = float(os.getenv("PEAK_LR", "2.5e-5"))
checkpoint = os.getenv("BASE_CHECKPOINT", "gs://openpi-assets/checkpoints/pi05_base/params")
precision = os.getenv("PYTORCH_TRAINING_PRECISION", "bfloat16")
if precision not in {"bfloat16", "float32"}:
raise ValueError("PYTORCH_TRAINING_PRECISION must be bfloat16 or float32")
return TrainConfig(
name=name,
project_name=os.getenv("WANDB_PROJECT", "openpi"),
exp_name=os.getenv("EXP_NAME", name),
model=pi0_config.Pi0Config(
pi05=_env_bool("MODEL_PI05", True),
action_dim=int(os.getenv("ACTION_DIM", "32")),
output_action_dim=int(os.getenv("OUTPUT_ACTION_DIM", "16")),
action_horizon=int(os.getenv("ACTION_HORIZON", "50")),
discrete_state_input=_env_bool("DISCRETE_STATE_INPUT", False),
),
data=data_factory(
repo_id=repo_id,
base_config=DataConfig(prompt_from_task=_env_bool("PROMPT_FROM_TASK", True)),
),
weight_loader=weight_loaders.CheckpointWeightLoader(checkpoint),
lr_schedule=_optimizer.CosineDecaySchedule(
warmup_steps=int(os.getenv("WARMUP_STEPS", "1000")),
peak_lr=peak_lr,
decay_steps=int(os.getenv("LR_DECAY_STEPS", str(train_steps))),
decay_lr=float(os.getenv("DECAY_LR", str(peak_lr / 10))),
),
optimizer=_optimizer.AdamW(
weight_decay=float(os.getenv("WEIGHT_DECAY", "1e-10")),
clip_gradient_norm=float(os.getenv("CLIP_GRADIENT_NORM", "1.0")),
),
assets_base_dir=os.getenv("ASSETS_BASE_DIR", "/workspace/openpi/assets"),
checkpoint_base_dir=os.getenv("CHECKPOINT_BASE_DIR", "/workspace/openpi/checkpoints"),
seed=int(os.getenv("SEED", "42")),
batch_size=int(os.getenv("BATCH_SIZE", "32")),
num_workers=int(os.getenv("NUM_WORKERS", "2")),
num_train_steps=train_steps,
log_interval=int(os.getenv("LOG_INTERVAL", "100")),
save_interval=int(os.getenv("SAVE_INTERVAL", "1000")),
keep_period=int(os.getenv("KEEP_PERIOD", "5000")),
overwrite=_env_bool("OVERWRITE", False),
resume=_env_bool("RESUME", False),
wandb_enabled=_env_bool("WANDB_ENABLED", True),
fsdp_devices=int(os.getenv("FSDP_DEVICES", "1")),
pytorch_training_precision=precision, # type: ignore[arg-type]
)
_PIPELINE_CONFIG = _env_config()
if _PIPELINE_CONFIG is not None:
_CONFIGS.append(_PIPELINE_CONFIG)
训练脚本
自动归一化数据+训练 在容器构建完成之后：
#!/usr/bin/env bash
set -euo pipefail
for name in TRAIN_CONFIG_NAME DATASET_REPO_ID; do
if [[ -z "${!name:-}" ]]; then
echo "ERROR: ${name} is required" >&2
exit 2
fi
done
echo "Config: ${TRAIN_CONFIG_NAME}"
echo "Dataset: ${HF_LEROBOT_HOME:-/data}/${DATASET_REPO_ID}"
echo "Experiment: ${EXP_NAME:-${TRAIN_CONFIG_NAME}}"
if [[ "${SKIP_NORM_STATS:-0}" != "1" ]]; then
echo "[1/2] Computing normalization statistics"
norm_args=(--config-name "${TRAIN_CONFIG_NAME}")
if [[ -n "${NORM_MAX_FRAMES:-}" ]]; then
norm_args+=(--max-frames "${NORM_MAX_FRAMES}")
fi
uv run python scripts/compute_norm_stats.py "${norm_args[@]}"
else
echo "[1/2] Skipping normalization statistics"
fi
echo "[2/2] Starting ${TRAIN_BACKEND:-jax} training"
case "${TRAIN_BACKEND:-jax}" in
jax)
uv run python scripts/train.py "${TRAIN_CONFIG_NAME}"
;;
pytorch)
gpu_count="${PYTORCH_GPU_COUNT:-1}"
if (( gpu_count > 1 )); then
uv run torchrun --standalone --nnodes=1 --nproc-per-node="${gpu_count}" \
scripts/train_pytorch.py "${TRAIN_CONFIG_NAME}"
else
uv run python scripts/train_pytorch.py "${TRAIN_CONFIG_NAME}"
fi
;;
*)
echo "ERROR: TRAIN_BACKEND must be jax or pytorch" >&2
exit 2
;;
esac
构建流程
最终文件树如图所示
- 构建镜像
docker compose -f docker/compose.train.yml --env-file .env build
- 查看构建的镜像
docker images | grep openpi-pi05
- 后台运行镜像自动运行entrypoint
(base) eft@eft-ThinkStation-P3-Tower:~/桌面/openpi$ docker compose -f docker/compose.train.yml --env-file .env up -d train
[+] Running 1/1
✔ Container openpi-pi05-train-1 Started