TRON 2 · 使用与开发
案例 · 模型训练
6. 模型训练
训练模型目前有两种方式,您可以根据情况自由选择:
| 训练方式 | 适用情况 | 优势 |
|---|---|---|
| 基于具身数据平台训练 | 暂无算力资源,或已有算力但希望快速跑通开发流程的初学者 | - 无代码训练: 通过可视化界面完成全训练流程,无需编程经验 - 实时监控: 提供训练指标、模型输出、系统日志的实时监控 - 自动化管理: 支持检查点管理、断点续训、参数调整等功能 - 自定义镜像:提供完整镜像,可直接复现官方开发流程 |
| 基于自有算力训练 | 已有算力(本地服务器或云资源)且熟悉云端开发的资深开发者 | - 复用已有算力资源 - 已开源训练与推理源码,开放性更强,可自行迭代优化(不可商用) |
6.1 基于具身平台训练
具身数据平台提供完整的模型训练功能,当前支持 OpenPI、ACT、自定义镜像 这几种机器人学习模型的训练,具体的训练方式如下述:
6.1.1 选择基座模型训练
具体使用流程请查阅原厂 具身数据平台使用手册 模型功能模块 - 训练工作流程。 |
|
6.1.2 选择自定义镜像训练
| 当用户需要借助具身平台的算力、基于我们的开源代码进行模型训练时,需通过自定义镜像的方式实现。具体填写内容如下: |
|
| 序号 | 参数名称 | 填写内容 |
|---|---|---|
| 1 | 训练镜像 | |
| 2 | 仓库用户名 | 不需要填写 |
| 3 | 仓库密码 | 不需要填写 |
| 4 | EntryPoint | `bash /root/tron2_openpi/scripts/cloud_train_entrypoint_portable.sh` |
| 5 | 启动参数 | 以下内容作为命令行参数传递给上述启动脚本,您可根据需要灵活调整steps、batch_size等参数: |
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`。
继续阅读案例 · 部署脚本与配置