使用私有数据集训练与部署#
概述#
本教程将指导你如何使用自己采集并生成好的数据集,配置 FluxVLA 完成从训练到真机部署的完整流程。
在阅读本教程前,请确保你已经:
本教程以 configs/gr00t/gr00t_eagle_3b_aloha_full_finetune.py 为基础模板,演示如何为你的私有数据集创建一份完整的训练与部署配置。
整体流程#
使用私有数据集完成训练和部署的完整流程如下:
准备数据 — 确保数据已转换为 LeRobot v2.1 格式
创建配置文件 — 基于已有模板修改关键参数
启动训练 — 使用
scripts/train.sh执行分布式训练真机部署 — 使用训练好的 checkpoint 进行推理
其中,创建配置文件是最关键的一步。一份配置文件包含四大模块,你需要根据自己的数据集特点逐一进行调整:
配置模块 |
说明 |
是否必须修改 |
|---|---|---|
|
模型架构与预训练权重 |
视情况而定 |
|
训练数据路径、相机配置、维度参数 |
必须 |
|
训练超参数(epoch、学习率等) |
建议调整 |
|
真机部署的任务描述与硬件配置 |
必须 |
第一步:确认数据集结构#
在开始配置前,请确认你的数据集已经是 LeRobot v2.1 格式,目录结构如下:
your_dataset_path/
├── data/
│ ├── train/
│ │ ├── episode_0.parquet
│ │ ├── episode_1.parquet
│ │ └── ...
│ └── video/
│ ├── episode_0/
│ │ ├── observation.images.cam_high.mp4
│ │ ├── observation.images.cam_left_wrist.mp4
│ │ └── ...
│ └── ...
├── info.json
└── meta.json
你需要明确以下关键信息,后续配置中会反复用到:
需确认的信息 |
示例值(Aloha 双臂) |
你的值 |
|---|---|---|
数据集路径 |
|
|
相机名称列表 |
|
|
相机数量 |
3 |
|
机器人动作原始维度 ( |
14(双臂各 7 维) |
|
机器人状态维度 |
14 |
|
任务描述(自然语言) |
|
第二步:配置模型(model)#
模型配置定义了 VLA 的网络结构。如果你使用和 Aloha 相同的 GR00T-Eagle 架构进行 full finetune,通常只需要关注 vla_head 中与维度相关的参数。
model = dict(
type='LlavaVLA',
pretrained_name_or_path='./checkpoints/GR00T-N1.5-3B',
vlm_backbone=dict(
type='EagleBackbone',
vlm_path='fluxvla/models/third_party_models/eagle2_hg_model'),
vla_head=dict(
type='FlowMatchingHead',
state_dim=64, # 模型内部状态维度(padding 后)
hidden_size=1024,
input_embedding_dim=1536,
num_layers=1,
num_heads=4,
num_inference_timesteps=4,
num_steps=32,
traj_length=10, # 预测轨迹长度
action_dim=32, # 模型内部动作维度(padding 后)
ori_action_dim=14), # ⬅️ 修改:你的机器人原始动作维度
freeze_vlm_backbone=False,
name_mapping={
'vlm_backbone.vlm': 'backbone.eagle_model',
'vla_head': 'action_head'
},
freeze_projector=False)
需要修改的参数#
参数 |
说明 |
修改建议 |
|---|---|---|
|
机器人的原始动作维度 |
根据你的机器人修改。Aloha 双臂 = 14,UR3 单臂 = 7 |
|
padding 后的内部动作维度 |
一般保持 32 即可 |
|
padding 后的内部状态维度 |
一般保持 64 即可 |
|
预测的未来轨迹步数 |
默认 10,可根据任务复杂度调整 |
提示:
action_dim和state_dim是模型内部经过 padding 后的维度,不需要等于机器人的原始维度。ori_action_dim才是你的机器人的实际动作空间维度。
inference_model(可选)#
如果推理时需要不同的模型配置(例如不同的预训练路径),可以额外定义 inference_model,格式与 model 完全一致:
inference_model = dict(
type='LlavaVLA',
pretrained_name_or_path='./checkpoints/GR00T-N1.5-3B',
vlm_backbone=dict(
type='EagleBackbone',
dtype=None,
vlm_path='fluxvla/models/third_party_models/eagle2_hg_model'),
vla_head=dict(
type='FlowMatchingHead',
state_dim=64,
hidden_size=1024,
input_embedding_dim=1536,
num_layers=1,
num_heads=4,
num_steps=32,
num_inference_timesteps=4,
traj_length=10,
ori_action_dim=14, # ⬅️ 与 model 保持一致
action_dim=32),
freeze_vlm_backbone=False,
name_mapping={
'vlm_backbone.vlm': 'backbone.eagle_model',
'vla_head': 'action_head'
},
freeze_projector=False)
第三步:配置训练数据(train_dataloader)#
这是最需要仔细修改的部分,你需要根据自己的数据集特点调整数据路径、相机配置和维度参数。
train_dataloader = dict(
per_device_batch_size=8, # 每个 GPU 的 batch size
per_device_num_workers=4, # 每个 GPU 的数据加载进程数
dataset=dict(
type='DistributedRepeatingDataset',
name_mappings={'observation.state': ['proprio', 'action']},
statistic_keys=[
'observation.state', 'observation.eepose', 'timestamp'
],
datasets=[
dict(
type='ParquetDataset',
data_root_path=[
# ⬅️ 修改:替换为你的数据集路径
'./datasets/your_dataset_path/batch_01',
'./datasets/your_dataset_path/batch_02',
],
transforms=[...], # 见下方详细说明
action_window_size=32 # ⬅️ 修改:动作窗口大小
)
]))
3.1 数据集路径#
data_root_path 支持列表格式,你可以同时指定多个数据批次的路径:
data_root_path=[
'./datasets/RealRobot_AgileX_aloha_lerobot_v2/20250601_20250615_02_4090',
],
3.2 数据变换流水线(transforms)#
变换流水线是数据预处理的核心,需要按照以下顺序配置 5 个变换步骤:
Step 1:ProcessParquetInputs — 解析 Parquet 数据#
dict(
type='ProcessParquetInputs',
embodiment_id=0, # ⬅️ 机器人类型 ID(Aloha=0, UR3=2)
parquet_keys=[ # 需要从 parquet 中读取的字段
'observation.state',
'observation.eepose',
'timestamp',
'actions',
'info',
'stats',
'action_masks'
],
video_keys=[ # ⬅️ 修改:你的相机视频键名
'observation.images.cam_high',
'observation.images.cam_left_wrist',
'observation.images.cam_right_wrist'
],
name_mappings={'observation.state': ['states']}),
⚠️ 注意:
video_keys中的相机名称必须与数据集中video/目录下的文件名一致。例如,如果你的视频文件是observation.images.head_cam.mp4,则对应的 key 为observation.images.head_cam。
Step 2:ParquetPrompter — 生成文本提示#
dict(type='ParquetPrompter'),
此步骤会从 parquet 数据中提取任务描述(task 字段)作为文本提示词,无需额外配置。
Step 3:ProcessPromptsWithImage — 文本分词与图像关联#
dict(
type='ProcessPromptsWithImage',
max_len=900,
num_images=3, # ⬅️ 修改:相机数量
tokenizer=dict(
type='PretrainedTokenizer',
model_path='path/to/eagle2_hg_model',
)),
参数 |
说明 |
|---|---|
|
相机数量,必须与 |
|
文本序列最大长度 |
Step 4:图像预处理#
dict(type='ResizeImages', height=224, width=224),
dict(
type='NormalizeImages',
means=[[123.515625, 116.04492188, 103.59375], # 每个相机的 RGB 均值
[123.515625, 116.04492188, 103.59375],
[123.515625, 116.04492188, 103.59375]], # ⬅️ 列表长度 = 相机数量
stds=[[58.27148438, 57.02636719, 57.27539062],
[58.27148438, 57.02636719, 57.27539062],
[58.27148438, 57.02636719, 57.27539062]],
),
提示:
means和stds列表的长度必须等于相机数量。使用 Eagle 模型时,上述默认值即可。
Step 5:NormalizeStatesAndActions — 状态与动作归一化#
dict(
type='NormalizeStatesAndActions',
state_dim=64, # 与 model.vla_head.state_dim 一致
action_dim=32, # 与 model.vla_head.action_dim 一致
state_key='proprio',
action_key='action',
norm_type='mean_std') # 归一化方式:'mean_std' 或 'min_max'
3.3 核心参数对照表#
以下是数据配置中需要根据你的场景修改的所有参数汇总:
参数 |
位置 |
说明 |
Aloha 示例值 |
|---|---|---|---|
|
|
数据集路径 |
|
|
|
机器人类型 ID |
|
|
|
相机视频键名列表 |
3 个相机 |
|
|
相机数量 |
|
|
|
图像归一化参数(列表长度=相机数) |
3 组 |
|
|
内部状态维度(与 model 一致) |
|
|
|
内部动作维度(与 model 一致) |
|
|
|
归一化方式 |
|
|
|
动作窗口大小 |
|
第四步:配置训练参数(runner)#
训练参数通常不需要大幅修改,根据你的数据量和 GPU 资源做适当调整即可:
runner = dict(
type='FSDPTrainRunner',
max_epochs=6, # ⬅️ 训练轮数,数据量少可适当增大
learning_rate=2e-5, # ⬅️ 学习率
weight_decay=0.0,
max_grad_norm=1.0,
sampler=None,
tokenizer=dict(
type='PretrainedTokenizer',
model_path='path/to/eagle2_hg_model',
),
collator=dict(
type='DictCollator',
keys=[
'states', 'observation.eepose', 'timestamp', 'images',
'img_masks', 'lang_tokens', 'lang_masks', 'actions',
'action_masks', 'embodiment_ids'
],
meta_keys=['task_description', 'prompt', 'info', 'stats']),
metric=dict(
type='VLAMetric',
active_trackers=('jsonl', 'wandb'),
run_dir='work_dirs',
grad_accumulation_steps=1,
window_size=1),
lr_scheduler_type='constant', # 学习率调度策略
warmup_ratio=0.0,
enable_gradient_checkpointing=False,
enable_mixed_precision_training=True,
mixed_precision_dtype='bf16',
change_key_name=False)
常用调整建议#
参数 |
调整建议 |
|---|---|
|
数据量较少(< 100 episodes)时建议 10-24;较多时 3-6 即可 |
|
full finetune 建议 1e-5 ~ 5e-5 |
|
可选 |
|
使用余弦衰减时建议设为 0.03 |
|
显存不足时设为 |
|
显存不足时减小,如 4 或 2 |
第五步:配置推理部署(inference)#
推理配置用于真机部署,需要指定任务描述、数据处理流水线和机器人硬件接口。
inference = dict(
type='AlohaInferenceRunner', # ⬅️ 推理器类型,根据机器人选择
seed=7,
task_descriptions={ # ⬅️ 修改:你的任务描述
'1': 'pick up the green bowl with right arm',
'2': 'place it on the blue plate with left arm',
},
dataset=dict(
type='PrivateInferenceDataset',
embodiment_id=0, # ⬅️ 与训练一致
img_keys=[ # ⬅️ 修改:相机名称(不含前缀)
'cam_high',
'cam_left_wrist',
'cam_right_wrist'
],
transforms=[
dict(
type='ProcessPromptsWithImage',
max_len=900,
num_images=3, # ⬅️ 与训练一致
tokenizer=dict(type='PretrainedTokenizer')),
dict(type='ResizeImages', height=224, width=224),
dict(
type='NormalizeImages',
means=[[123.515625, 116.04492188, 103.59375],
[123.515625, 116.04492188, 103.59375],
[123.515625, 116.04492188, 103.59375]],
stds=[[58.27148438, 57.02636719, 57.27539062],
[58.27148438, 57.02636719, 57.27539062],
[58.27148438, 57.02636719, 57.27539062]],
),
dict(
type='NormalizeStatesAndActions',
state_dim=64, # ⬅️ 与训练一致
state_key='proprio',
action_key='action',
norm_type='mean_std') # ⬅️ 与训练一致
]),
denormalize_action=dict(
type='DenormalizePrivateAction',
norm_type='mean_std', # ⬅️ 与训练的 norm_type 一致
action_dim=14), # ⬅️ 机器人的原始动作维度
operator=dict( # ⬅️ 修改:你的机器人 ROS topic 配置
type='AlohaOperator',
img_front_topic='/camera_f/color/image_raw',
img_left_topic='/camera_l/color/image_raw',
img_right_topic='/camera_r/color/image_raw',
img_front_depth_topic='/camera_f/depth/image_raw',
img_left_depth_topic='/camera_l/depth/image_raw',
img_right_depth_topic='/camera_r/depth/image_raw',
puppet_arm_left_cmd_topic='/master/joint_left',
puppet_arm_right_cmd_topic='/master/joint_right',
puppet_arm_left_topic='/puppet/joint_left',
puppet_arm_right_topic='/puppet/joint_right',
robot_base_topic='/odom_raw',
robot_base_cmd_topic='/cmd_vel',
))
推理配置要点#
type:根据你的机器人平台选择推理器。Aloha 双臂使用AlohaInferenceRunner,UR 系列使用URInferenceRunner。task_descriptions:定义可用的任务指令。键为任务 ID(字符串),值为自然语言任务描述。推理时通过键盘按键选择任务。dataset.transforms:推理时的数据预处理流水线必须与训练时保持一致(归一化方式、图像尺寸等)。denormalize_action:将模型输出的归一化动作转换为真实机器人动作。norm_type和action_dim必须与训练一致。operator:定义机器人的 ROS topic 通信接口,需要根据你的实际硬件配置修改。
第六步:启动训练#
配置文件准备好后,即可启动训练。将配置文件保存到 configs/ 目录下(例如 configs/gr00t/my_custom_finetune.py),然后执行:
# 设置环境变量(单节点 8 GPU 示例)
export MLP_WORKER_GPU=8
export MLP_WORKER_NUM=1
export MLP_ROLE_INDEX=0
export MLP_WORKER_0_HOST=localhost
export MLP_WORKER_0_PORT=29500
# 启动训练
bash scripts/train.sh \
configs/gr00t/my_custom_finetune.py \
work_dirs/my_custom_finetune
训练完成后,在 work_dirs/my_custom_finetune/ 中会生成:
checkpoint_*.pt— 模型检查点dataset_statistics.json— 数据集统计信息(推理时使用)config.yaml/config.json— 配置备份
第七步:真机部署#
使用训练好的 checkpoint 进行真机推理:
python scripts/inference_real_robot.py \
--config configs/gr00t/my_custom_finetune.py \
--ckpt-path work_dirs/my_custom_finetune/checkpoint_epoch_6.pt
⚠️ 注意:推理时使用的配置文件应与训练时一致,确保数据归一化方式、维度参数等完全匹配。
完整示例:Aloha 双臂机器人#
以下是 configs/gr00t/gr00t_eagle_3b_aloha_full_finetune.py 的完整配置,作为创建你自己配置文件的参考模板。
# ============================================================
# 模型配置
# ============================================================
model = dict(
type='LlavaVLA',
pretrained_name_or_path='./checkpoints/GR00T-N1.5-3B',
vlm_backbone=dict(
type='EagleBackbone',
vlm_path='fluxvla/models/third_party_models/eagle2_hg_model'),
vla_head=dict(
type='FlowMatchingHead',
state_dim=64,
hidden_size=1024,
input_embedding_dim=1536,
num_layers=1,
num_heads=4,
num_inference_timesteps=4,
num_steps=32,
traj_length=10,
action_dim=32,
ori_action_dim=14), # Aloha 双臂:左臂 7 + 右臂 7 = 14
freeze_vlm_backbone=False,
name_mapping={
'vlm_backbone.vlm': 'backbone.eagle_model',
'vla_head': 'action_head'
},
freeze_projector=False)
inference_model = dict(
type='LlavaVLA',
pretrained_name_or_path='./checkpoints/GR00T-N1.5-3B',
vlm_backbone=dict(
type='EagleBackbone',
dtype=None,
vlm_path='fluxvla/models/third_party_models/eagle2_hg_model'),
vla_head=dict(
type='FlowMatchingHead',
state_dim=64,
hidden_size=1024,
input_embedding_dim=1536,
num_layers=1,
num_heads=4,
num_steps=32,
num_inference_timesteps=4,
traj_length=10,
ori_action_dim=14,
action_dim=32),
freeze_vlm_backbone=False,
name_mapping={
'vlm_backbone.vlm': 'backbone.eagle_model',
'vla_head': 'action_head'
},
freeze_projector=False)
# ============================================================
# 数据配置
# ============================================================
train_dataloader = dict(
per_device_batch_size=8,
per_device_num_workers=4,
dataset=dict(
type='DistributedRepeatingDataset',
name_mappings={'observation.state': ['proprio', 'action']},
statistic_keys=[
'observation.state', 'observation.eepose', 'timestamp'
],
datasets=[
dict(
type='ParquetDataset',
data_root_path=[
'./datasets/RealRobot_AgileX_aloha_lerobot_v2/20250601_20250615_02_4090',
],
transforms=[
dict(
type='ProcessParquetInputs',
embodiment_id=0, # Aloha
parquet_keys=[
'observation.state', 'observation.eepose',
'timestamp', 'actions', 'info', 'stats',
'action_masks'
],
video_keys=[ # 3 个相机
'observation.images.cam_high',
'observation.images.cam_left_wrist',
'observation.images.cam_right_wrist'
],
name_mappings={'observation.state': ['states']}),
dict(type='ParquetPrompter'),
dict(
type='ProcessPromptsWithImage',
max_len=900,
num_images=3, # 相机数量
tokenizer=dict(
type='PretrainedTokenizer',
model_path='path/to/eagle2_hg_model',
)),
dict(type='ResizeImages', height=224, width=224),
dict(
type='NormalizeImages',
means=[[123.515625, 116.04492188, 103.59375],
[123.515625, 116.04492188, 103.59375],
[123.515625, 116.04492188, 103.59375]],
stds=[[58.27148438, 57.02636719, 57.27539062],
[58.27148438, 57.02636719, 57.27539062],
[58.27148438, 57.02636719, 57.27539062]],
),
dict(
type='NormalizeStatesAndActions',
state_dim=64,
action_dim=32,
state_key='proprio',
action_key='action',
norm_type='mean_std')
],
action_window_size=32)
]))
# ============================================================
# 训练配置
# ============================================================
runner = dict(
type='FSDPTrainRunner',
max_epochs=6,
learning_rate=2e-5,
weight_decay=0.0,
max_grad_norm=1.0,
sampler=None,
tokenizer=dict(
type='PretrainedTokenizer',
model_path='path/to/eagle2_hg_model',
),
collator=dict(
type='DictCollator',
keys=[
'states', 'observation.eepose', 'timestamp', 'images',
'img_masks', 'lang_tokens', 'lang_masks', 'actions',
'action_masks', 'embodiment_ids'
],
meta_keys=['task_description', 'prompt', 'info', 'stats']),
metric=dict(
type='VLAMetric',
active_trackers=('jsonl', 'wandb'),
run_dir='work_dirs',
grad_accumulation_steps=1,
window_size=1),
lr_scheduler_type='constant',
warmup_ratio=0.0,
enable_gradient_checkpointing=False,
enable_mixed_precision_training=True,
mixed_precision_dtype='bf16',
change_key_name=False)
# ============================================================
# 推理配置
# ============================================================
inference = dict(
type='AlohaInferenceRunner',
seed=7,
task_descriptions={
'1': 'pick up the green bowl with right arm',
'2': 'place it on the blue plate with left arm',
# ... 根据你的任务添加更多
},
dataset=dict(
type='PrivateInferenceDataset',
embodiment_id=0,
img_keys=['cam_high', 'cam_left_wrist', 'cam_right_wrist'],
transforms=[
dict(
type='ProcessPromptsWithImage',
max_len=900,
num_images=3,
tokenizer=dict(type='PretrainedTokenizer')),
dict(type='ResizeImages', height=224, width=224),
dict(
type='NormalizeImages',
means=[[123.515625, 116.04492188, 103.59375],
[123.515625, 116.04492188, 103.59375],
[123.515625, 116.04492188, 103.59375]],
stds=[[58.27148438, 57.02636719, 57.27539062],
[58.27148438, 57.02636719, 57.27539062],
[58.27148438, 57.02636719, 57.27539062]],
),
dict(
type='NormalizeStatesAndActions',
state_dim=64,
state_key='proprio',
action_key='action',
norm_type='mean_std')
]),
denormalize_action=dict(
type='DenormalizePrivateAction',
norm_type='mean_std',
action_dim=14), # 原始动作维度
operator=dict(
type='AlohaOperator',
img_front_topic='/camera_f/color/image_raw',
img_left_topic='/camera_l/color/image_raw',
img_right_topic='/camera_r/color/image_raw',
img_front_depth_topic='/camera_f/depth/image_raw',
img_left_depth_topic='/camera_l/depth/image_raw',
img_right_depth_topic='/camera_r/depth/image_raw',
puppet_arm_left_cmd_topic='/master/joint_left',
puppet_arm_right_cmd_topic='/master/joint_right',
puppet_arm_left_topic='/puppet/joint_left',
puppet_arm_right_topic='/puppet/joint_right',
robot_base_topic='/odom_raw',
robot_base_cmd_topic='/cmd_vel',
))
常见问题#
Q:如何适配不同数量的相机?#
需要同步修改以下位置:
ProcessParquetInputs.video_keys— 列出所有相机的视频键ProcessPromptsWithImage.num_images— 设置相机数量NormalizeImages.means/stds— 每个相机一组归一化参数inference.dataset.img_keys— 列出推理时的相机名称
例如,使用 2 个相机的配置:
video_keys=[
'observation.images.cam_high',
'observation.images.cam_wrist'
],
# ...
num_images=2,
# ...
means=[[123.515625, 116.04492188, 103.59375],
[123.515625, 116.04492188, 103.59375]],
stds=[[58.27148438, 57.02636719, 57.27539062],
[58.27148438, 57.02636719, 57.27539062]],
Q:训练时和推理时的归一化方式不一致会怎样?#
推理效果会严重下降。请务必确保以下参数在训练和推理配置中完全一致:
norm_type(mean_std/min_max)state_dim、action_dim图像归一化的
means/stds图像尺寸
height/width
Q:如何使用多个数据集混合训练?#
在 datasets 列表中添加多个 ParquetDataset 配置即可:
datasets=[
dict(
type='ParquetDataset',
data_root_path=['./datasets/task_A/batch_01'],
transforms=[...],
action_window_size=32),
dict(
type='ParquetDataset',
data_root_path=['./datasets/task_B/batch_01'],
transforms=[...],
action_window_size=32),
]
Q:如何选择合适的 action_window_size?#
action_window_size 决定了每个训练样本中包含的动作序列长度。建议与 model.vla_head.action_dim 保持一致(如 32),这样模型在一次预测中生成一整个窗口的动作。