使用私有数据集训练与部署#

概述#

本教程将指导你如何使用自己采集并生成好的数据集,配置 FluxVLA 完成从训练到真机部署的完整流程。

在阅读本教程前,请确保你已经:

  1. 按照 安装指南 完成环境搭建

  2. 按照 真机数据准备 将原始数据转换为 LeRobot v2.1 格式

  3. 下载了预训练模型权重(如 GR00T-N1.5-3B

本教程以 configs/gr00t/gr00t_eagle_3b_aloha_full_finetune.py 为基础模板,演示如何为你的私有数据集创建一份完整的训练与部署配置。


整体流程#

使用私有数据集完成训练和部署的完整流程如下:

  1. 准备数据 — 确保数据已转换为 LeRobot v2.1 格式

  2. 创建配置文件 — 基于已有模板修改关键参数

  3. 启动训练 — 使用 scripts/train.sh 执行分布式训练

  4. 真机部署 — 使用训练好的 checkpoint 进行推理

其中,创建配置文件是最关键的一步。一份配置文件包含四大模块,你需要根据自己的数据集特点逐一进行调整:

配置模块

说明

是否必须修改

model

模型架构与预训练权重

视情况而定

train_dataloader

训练数据路径、相机配置、维度参数

必须

runner

训练超参数(epoch、学习率等)

建议调整

inference

真机部署的任务描述与硬件配置

必须


第一步:确认数据集结构#

在开始配置前,请确认你的数据集已经是 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 双臂)

你的值

数据集路径

./datasets/RealRobot_AgileX_aloha_lerobot_v2/20250601_20250615_02_4090

相机名称列表

cam_high, cam_left_wrist, cam_right_wrist

相机数量

3

机器人动作原始维度 (ori_action_dim)

14(双臂各 7 维)

机器人状态维度

14

任务描述(自然语言)

"pick up the green bowl with right arm"


第二步:配置模型(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)

需要修改的参数#

参数

说明

修改建议

ori_action_dim

机器人的原始动作维度

根据你的机器人修改。Aloha 双臂 = 14,UR3 单臂 = 7

action_dim

padding 后的内部动作维度

一般保持 32 即可

state_dim

padding 后的内部状态维度

一般保持 64 即可

traj_length

预测的未来轨迹步数

默认 10,可根据任务复杂度调整

提示action_dimstate_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',
    )),

参数

说明

num_images

相机数量,必须与 video_keys 的数量一致

max_len

文本序列最大长度

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]],
),

提示meansstds 列表的长度必须等于相机数量。使用 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 示例值

data_root_path

ParquetDataset

数据集路径

./datasets/RealRobot_AgileX_aloha_lerobot_v2/...

embodiment_id

ProcessParquetInputs

机器人类型 ID

0

video_keys

ProcessParquetInputs

相机视频键名列表

3 个相机

num_images

ProcessPromptsWithImage

相机数量

3

means / stds

NormalizeImages

图像归一化参数(列表长度=相机数)

3 组

state_dim

NormalizeStatesAndActions

内部状态维度(与 model 一致)

64

action_dim

NormalizeStatesAndActions

内部动作维度(与 model 一致)

32

norm_type

NormalizeStatesAndActions

归一化方式

mean_std

action_window_size

ParquetDataset

动作窗口大小

32


第四步:配置训练参数(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)

常用调整建议#

参数

调整建议

max_epochs

数据量较少(< 100 episodes)时建议 10-24;较多时 3-6 即可

learning_rate

full finetune 建议 1e-5 ~ 5e-5

lr_scheduler_type

可选 constant(恒定)或 linear-warmup+cosine-decay(预热+余弦衰减)

warmup_ratio

使用余弦衰减时建议设为 0.03

enable_gradient_checkpointing

显存不足时设为 True,以时间换空间

per_device_batch_size(在 train_dataloader 中)

显存不足时减小,如 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',
    ))

推理配置要点#

  1. type:根据你的机器人平台选择推理器。Aloha 双臂使用 AlohaInferenceRunner,UR 系列使用 URInferenceRunner

  2. task_descriptions:定义可用的任务指令。键为任务 ID(字符串),值为自然语言任务描述。推理时通过键盘按键选择任务。

  3. dataset.transforms:推理时的数据预处理流水线必须与训练时保持一致(归一化方式、图像尺寸等)。

  4. denormalize_action:将模型输出的归一化动作转换为真实机器人动作。norm_typeaction_dim 必须与训练一致。

  5. 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:如何适配不同数量的相机?#

需要同步修改以下位置:

  1. ProcessParquetInputs.video_keys — 列出所有相机的视频键

  2. ProcessPromptsWithImage.num_images — 设置相机数量

  3. NormalizeImages.means / stds — 每个相机一组归一化参数

  4. 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_typemean_std / min_max

  • state_dimaction_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),这样模型在一次预测中生成一整个窗口的动作。