TRON 2 DOCS
产品文档/TRON 2/VLA · 项目开发
TRON 2 · 使用与开发

案例 · 模型训练

适用条件见本页各节

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仓库密码不需要填写
4EntryPoint`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
❗
补充规范说明:
  1. 自定义任务配置文件(如`my_task.yaml`)已配置在`.gitignore` 中,不会提交至代码仓库;仓库仅保留通用模型`configs/train/tron2_tasks/example.yaml`;
  2. 模版 YAML 支持配置项:数据集`repo_id`、任务文本`prompt`、数据列定义、 动作预测时长`action_horizon`、状态维度`state_dim`、预训练基准权重`base checkpoint`、训练输出路径,以及可选的 `prompt_from_task` 和 `rtc_training_simulated_delay`。
继续阅读案例 · 部署脚本与配置