Example · model training | 机器人技术学习笔记

6. 模型训练

训练模型目前有两种方式，您可以根据情况自由选择：

训练方式	适用情况	优势
基于具身数据平台训练	暂无算力资源，或已有算力但希望快速跑通开发流程的初学者	- 无代码训练: 通过可视化界面完成全训练流程，无需编程经验 - 实时监控: 提供训练指标、模型输出、系统日志的实时监控 - 自动化管理: 支持检查点管理、断点续训、参数调整等功能 - 自定义镜像：提供完整镜像，可直接复现官方开发流程
基于自有算力训练	已有算力（本地服务器或云资源）且熟悉云端开发的资深开发者	- 复用已有算力资源 - 已开源训练与推理源码，开放性更强，可自行迭代优化（不可商用）

6.1 基于具身平台训练

具身数据平台提供完整的模型训练功能，当前支持 OpenPI、ACT、自定义镜像这几种机器人学习模型的训练，具体的训练方式如下述：

6.1.1 选择基座模型训练

具体使用流程请查阅原厂具身数据平台使用手册 模型功能模块 - 训练工作流程。

6.1.2 选择自定义镜像训练

当用户需要借助具身平台的算力、基于我们的开源代码进行模型训练时，需通过自定义镜像的方式实现。具体填写内容如下：

序号	参数名称	填写内容
1	训练镜像	`limx-acr-wlcb-registry.cn-wulanchabu.cr.aliyuncs.com/limx_emb/hum:pytorch210-gpu-py312-cu128-ubuntu22-tron2-openpi-3`
2	仓库用户名	不需要填写
3	仓库密码	不需要填写
4	EntryPoint	`bash /root/tron2_openpi/scripts/cloud_train_entrypoint_portable.sh`
5	启动参数	以下内容作为命令行参数传递给上述启动脚本，您可根据需要灵活调整steps、batch_size等参数： `--data-dir /data --repo-id input --weight /models/openpi-assets/checkpoints/pi_base/params --output-dir /data/output --exp test --steps 30000 --save-interval 10000 --batch-size 64 --rtc-delay 10 --action-horizon 30 --prompt-from-task --overwrite`

6.2 基于自有算力训练

⚠️

注意：

使用自有算力训练模型前，请先参考本文开头「训练最低配置」部分。

💡

关于训练框架

OpenPI 开源项目中的底层技术栈包含 JAX 和 PyTorch 两种框架，由于训练其核心模型 Pi05。由于 OpenPI 原生训练框架为JAX，且核心模型、训练流程及基础 Checkpoint 均基于 JAX/Flax 构建，因此本项目默认使用 Jax 框架进行训练。
JAX 是一种用于模型训练的计算工具，定位类似 PyTorch — 它不是模型，也不是算法，只是“用来跑模型的框架”。
JAX 的优势在于计算效率高，尤其适合大模型训练场景。本项目选用 JAX 正是因为其在 Pi0.5 这类大模型训练中更为高效。

如需在自有算力环境中，使用开源数据集或自定义数据集训练模型，可参考下文流程并结合开源代码实现。当前开源代码已完成 OpenPI 源码与 TRON2 的适配，本文以 Pi05 为例进行说明：

6.2.1 配置训练环境【同推理】

git clone https://github.com/limx-tron2/tron2_openpi.git
git clone https://github.com/limx-tron2/tron2_env.git

cd tron2_openpi

sudo apt install pkg-config
sudo apt install software-properties-common -y
sudo add-apt-repository ppa:ubuntuhandbook1/ffmpeg7 -y
sudo apt install ffmpeg libavformat-dev libavcodec-dev libavdevice-dev libavutil-dev libswscale-dev libswresample-dev libavfilter-dev -y

uv sync
uv pip install -e .

source .venv/bin/activate 
uv pip install -e "../tron2_env[bridge,openpi]" # 激活当前环境 

# 验证是否正常安装成功，应输出路径：../tron2_env/src/tron2_env/__init__.py
PYTHONPATH="$(cd ../tron2_env/src && pwd)" uv run python -c "import tron2_env; print(tron2_env.__file__)"

6.2.2 训练流程

1.训练统一推荐独立 YAML 任务配置文件管理参数，不建议直接修改 src/openpi/training/config.py，便于配置隔离、版本管理与复用：

cp configs/train/tron2_tasks/example.yaml configs/train/tron2_tasks/my_task.yaml

2.打开配置文件： configs/train/tron2_tasks/my_task.yaml，修改数据集地址参数：repo_id: my_dataset；系统通过环境变量 HF_LEROBOT_HOME 拼接完整数据集路径，本地目录必须满足固定结构：

1.`$HF_LEROBOT_HOME/my_dataset/data/`：存放图像、轨迹、动作等原始数据

2.`$HF_LEROBOT_HOME/my_dataset/meta/`：存放数据集统计、任务描述、episode信息

export HF_LEROBOT_HOME=/path/to/datasets

3.数据集归一化参数是训练收敛的必要前置条件，首次训练前运行脚本 normalization statistics 计算统计文件：

uv run scripts/compute_norm_stats.py \
  --task-config configs/train/tron2_tasks/my_task.yaml

4.统计量计算完成后，执行训练脚本启动训练：

uv run scripts/train_tron2_task.py \
  --task-config configs/train/tron2_tasks/my_task.yaml

❗

补充：

我们提供了一体化入口脚本 `scripts/cloud_train_entrypoint_portable.sh`，无需分步执行：

不传 `--skip-norm` 参数时，脚本会自动先计算归一化统计量，再启动训练；
如需跳过统计量计算，增加入参 `--skip-norm`。

6.2.2.1 本地自定义路径模式：

自有算力无统一挂载目录时，手动传入数据集、权重、输出完整路径：

scripts/cloud_train_entrypoint_portable.sh \
  --data-dir /path/to/datasets \
  --repo-id my_dataset \
  --weight /path/to/checkpoint/params \
  --output-dir /path/to/output \
  --exp my_task \
  --prompt "perform the configured task" \
  --max-frames 100000

6.2.2.2 绑定自定义 YAML 配置文件模式：

直接加载已编辑完成的my_task.yaml，文件内统一管理repo_id、加载权重weight_loader、资源根目录assets_base_dir和保存路径checkpoint_base_dir等参数；仅需通过--data-dir 指定 HF_LEROBOT_HOME：

scripts/cloud_train_entrypoint_portable.sh \
  --task-config configs/train/tron2_tasks/my_task.yaml \
  --exp my_task \
  --data-dir "$HF_LEROBOT_HOME" \
  --max-frames 100000

❗

补充规范说明：

自定义任务配置文件（如`my_task.yaml`）已配置在`.gitignore` 中，不会提交至代码仓库；仓库仅保留通用模型`configs/train/tron2_tasks/example.yaml`；
模版 YAML 支持配置项：数据集`repo_id`、任务文本`prompt`、数据列定义、动作预测时长`action_horizon`、状态维度`state_dim`、预训练基准权重`base checkpoint`、训练输出路径，以及可选的 `prompt_from_task` 和 `rtc_training_simulated_delay`。