liujiting/openpi-wu-new

OpenPI-WU Spatial 训练包

这个目录包含训练代码、Pi0 PyTorch base 权重、LIBERO Spatial 数据、归一化统计以及 uv/Python 运行时。当前入口只训练 LIBERO Spatial，默认启动 Stage A：冻结 Pi0 base 和 VLM，只训练 rank-16、scale 1.0 的 Action-WU。

目录结构

最新openpi/
├── openpi-wu-2/   # 代码与训练脚本
├── runtime/       # Pi0 base、LIBERO 数据和 norm stats
├── uv/            # uv 及 Python 3.11
└── checkpoints/   # 训练输出

要求 Linux x86_64、NVIDIA GPU 和可用的 CUDA 驱动。建议使用 A100/A800 80GB；代码已在单张 A800 80GB、batch 32 下跑通。

1. 配置环境

进入解压后的目录：

cd openpi-wu-2

export UV_PYTHON_INSTALL_DIR="$(cd ../uv/python && pwd)"
PYTHON_BIN="$UV_PYTHON_INSTALL_DIR/cpython-3.11.15-linux-x86_64-gnu/bin/python3.11"
chmod +x ../uv/bin/uv "$PYTHON_BIN"
../uv/bin/uv venv --clear --python "$PYTHON_BIN"
../uv/bin/uv sync --frozen

.venv 必须在目标服务器重新创建，不要直接使用打包机器留下的虚拟环境。首次 uv sync 需要能够访问 Python/Git 依赖源。

2. 八卡正式训练

下面的命令使用 8 张 GPU、每卡 batch 32、global batch 256、梯度累积 1，训练 50,000 个 optimizer steps；学习率为 1e-5，warmup 1,000 steps，checkpoint 每 2,500 steps 保存一次，SwanLab 默认离线记录。

PACKAGE_DIR="$(cd .. && pwd)"

NUM_GPUS=8 \
PER_GPU_BATCH_SIZE=32 \
NUM_WORKERS=2 \
RUNTIME_DIR="$PACKAGE_DIR/runtime" \
CHECKPOINT_DIR="$PACKAGE_DIR/checkpoints" \
SWANLAB_MODE=local \
bash scripts/run_pi0_two_wu_spatial.sh action

训练输出位于：

checkpoints/pi0_libero_spatial/action-first-spatial-action-r16-gb256-50000/

3. 单卡检查

单卡只用于确认环境、数据和训练链路可用。NUM_WORKERS=0 可以避免首次启动时多进程 DataLoader 初始化过慢。

PACKAGE_DIR="$(cd .. && pwd)"

NUM_GPUS=1 \
PER_GPU_BATCH_SIZE=32 \
NUM_WORKERS=0 \
RUNTIME_DIR="$PACKAGE_DIR/runtime" \
CHECKPOINT_DIR="$PACKAGE_DIR/checkpoints" \
SWANLAB_MODE=local \
ACTION_EXP_NAME="action-first-spatial-action-r16-gb32-50k-single" \
bash scripts/run_pi0_two_wu_spatial.sh action

启动后脚本会先加载 Pi0 base、解析 52,970 条 Spatial 样本并准备首个 batch。使用 NUM_WORKERS>0 时，进度可能暂时停在 0/50000；看到 GPU 显存升至约 75GB 且利用率开始波动，即表示训练已经进入计算。

常用参数

参数	含义	正式训练默认值
NUM_GPUS	使用的 GPU 数量	8
PER_GPU_BATCH_SIZE	每张 GPU 的 batch	32
NUM_WORKERS	每个训练进程的数据加载 worker	2
ACTION_TRAIN_STEPS	Stage A optimizer steps	50000
ACTION_PEAK_LR	Stage A 峰值学习率	1e-5
SAVE_INTERVAL	checkpoint 保存间隔	2500
SWANLAB_MODE	local 离线或 cloud 在线	local

如需在线 SwanLab，设置 SWANLAB_MODE=cloud 和 SWANLAB_API_KEY。不要修改 runtime 内的 norm stats；它们与当前 Spatial 数据和 Pi0 policy 配置配套。