Hugging Face's logo

LeRobot documentation

多任务 DiT 策略

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

多任务 DiT 策略

多任务扩散 Transformer（DiT）策略是原始扩散策略架构的演进版本，它利用大型 DiT 并结合文本和视觉条件，用于多任务机器人学习。该实现支持扩散和流匹配目标以生成动作，使机器人能够根据语言指令执行多种多样的操作任务。

模型概览

该模型使用：

  • CLIP 视觉编码器：处理来自多个相机视角的 RGB 图像
  • CLIP 文本编码器：编码语言任务指令（冻结权重，带有可学习投影）
  • Diffusion Transformer：根据观测和语言预测动作序列
  • 两个目标：支持扩散（DDPM/DDIM）和流匹配用于动作生成

该模型令人兴奋，因为仅使用约 450M 参数和显著更少的训练，就能实现极高的灵巧性，可与数十亿参数的 VLA 相媲美。

安装要求

Multitask DiT Policy 有额外的依赖项。请使用以下命令安装：

pip install 'lerobot[multi_task_dit]'

这将安装所有必要的依赖项，包括用于 CLIP 模型的 HuggingFace Transformers 库。

用法

要在您的 LeRobot 配置中使用 Multitask DiT，请将策略类型指定为：

policy.type=multi_task_dit

训练

基础训练命令

以下是用于在您的数据集上训练 Multitask DiT 的完整训练命令：

lerobot-train \
  --dataset.repo_id=YOUR_DATASET \
  --output_dir=./outputs/multitask_dit_training \
  --batch_size=32 \
  --steps=5000 \
  --save_freq=500 \
  --log_freq=100 \
  --policy.type=multi_task_dit \
  --policy.device=cuda \
  --policy.repo_id="HF_USER/multitask-dit-your-robot" \
  --trackio.enable=true

推荐的超参数和数据集详情（30Hz 控制频率）

为了获得可靠的性能，请从以下建议的默认超参数开始：

lerobot-train \
  --dataset.repo_id=YOUR_DATASET \
  --output_dir=./outputs/mutitask_dit_training \
  --batch_size=320 \
  --steps=30000 \
  --policy.type=multi_task_dit \
  --policy.device=cuda \
  --policy.horizon=32 \
  --policy.n_action_steps=24 \
  --policy.objective=diffusion \
  --policy.noise_scheduler_type=DDPM \
  --policy.num_train_timesteps=100 \
  --policy.repo_id="HF_USER/multitask-dit-your-robot" \
  --trackio.enable=true

关键参数：

  • 批大小：192-320 - 如果你拥有能够支持此配置的 GPU，你将获得最佳的训练动态
  • Horizon：32 - 要预测的动作步数，在 30Hz 下约为 1.0 秒
  • n_action_steps：24 - 30Hz 下约 0.8 秒
  • 目标：diffusion - 从扩散开始，如果生成质量较差，则尝试流匹配
  • 训练步数：针对单一任务，建议超过 30k 步

训练配置参数

目标选择

在扩散和流匹配之间进行选择：

# Diffusion objective (default)
--policy.objective=diffusion \
--policy.noise_scheduler_type=DDPM \
--policy.num_train_timesteps=100 \
--policy.num_inference_steps=10 \
--policy.beta_schedule=squaredcos_cap_v2 \
--policy.prediction_type=epsilon \
--policy.clip_sample=true \
--policy.clip_sample_range=1.0

# Flow matching objective
--policy.objective=flow_matching \
--policy.timestep_sampling_strategy=beta \
--policy.num_integration_steps=100 \
--policy.integration_method=euler \
--policy.sigma_min=0.0
标志备注
--policy.noise_scheduler_typeDDPM 或 DDIM
--policy.num_inference_steps步骤更少，推理速度更快
--policy.beta_schedule噪声调度类型
--policy.prediction_typeepsilon（预测噪声）或 sample（预测干净样本）
--policy.clip_sample在去噪过程中裁剪样本
--policy.clip_sample_range裁剪范围 [-x, x]
--policy.timestep_sampling_strategybeta 或 uniform；在实践中，beta 策略的效果似乎要好得多
--policy.integration_methodeuler 或 rk4
--policy.sigma_min流插值路径中的最小噪声

Transformer 架构

根据数据集大小调整模型容量：

# Small datasets (< 100 examples)
--policy.num_layers=4 \
--policy.hidden_dim=512 \
--policy.num_heads=8  # should ideally be hidden_dim // 64

# Medium datasets (100-5k examples) - default
--policy.num_layers=6 \
--policy.hidden_dim=512 \
--policy.num_heads=8  # should ideally be hidden_dim // 64

# Large datasets (> 5k examples)
--policy.num_layers=8 \
--policy.hidden_dim=512 \
--policy.num_heads=8   # should ideally be hidden_dim // 64

位置编码选项：

该模型支持两种用于动作序列的位置编码方法：

# Rotary Position Embedding (RoPE) - default, recommended
--policy.use_rope=true \
--policy.rope_base=10000.0  # Base frequency for RoPE

# Absolute positional encoding
--policy.use_positional_encoding=true  # Disables RoPE when true

其他 Transformer 参数：

--policy.dropout=0.1  # Dropout rate for DiT blocks (0.0-1.0)
--policy.timestep_embed_dim=256  # Timestep embedding dimension

视觉编码器配置

# Use different CLIP model for more expressivity at the cost of inference time
# experiment with larger or smaller models depending on the complexity of your tasks and size of dataset
--policy.vision_encoder_name=openai/clip-vit-large-patch14

# Use separate vision encoder per camera
# This may be useful when cameras have significantly different characteristics, but
# be wary of increased VRAM footprint.
--policy.use_separate_rgb_encoder_per_camera=true

# Image preprocessing (you may need to resize your images for inference speed ups)
--policy.image_resize_shape=[XXX,YYY] \
--policy.image_crop_shape=[224,224] \
--policy.image_crop_is_random=true  # Random during training, center at inference

文本编码器配置

# Use different CLIP text encoder model
# same as vision: experiment with larger or smaller models depending on the
# complexity of your tasks and size of dataset
--policy.text_encoder_name=openai/clip-vit-large-patch14

学习率配置

视觉编码器使用独立的学习率乘数，建议以 1/10 作为理想的起始点：

--policy.optimizer_lr=2e-5 \
--policy.vision_encoder_lr_multiplier=0.1  # Vision encoder LR = 0.1 * optimizer_lr

训练调优指南

1. 使用 Beta 采样的流匹配

此处的原始扩散实现基于TRI 的 LBM 论文中描述的工作。

此外，我们实现了一种流匹配目标，其高层描述见 Boston Dynamics 博客文章。

考虑测试 flow-matching 目标并评估您任务中的性能差异：

--policy.objective=flow_matching \
--policy.timestep_sampling_strategy=beta \
--policy.timestep_sampling_alpha=1.5 \
--policy.timestep_sampling_beta=1.0 \
--policy.timestep_sampling_s=0.999

这并非在所有用户场景下都被证明是万能解决方案，但偶尔能带来更流畅、更一致的操作体验。

2. Transformer 层数

将模型容量与数据集规模相匹配：

  • 小型数据集（< 100 个样本）：减少到 4 层
  • 大型数据集（> 5k 个样本）：增加到 8 层

3. horizon 调优

模型可能会对您选择的时间范围敏感。请根据控制频率，从大约 1 秒的时间范围开始：

  • 30 Hz 频率：horizon=30
  • 10 Hz 频率：horizon=10

然后尝试在此基础上逐步增加。horizon 决定了模型预测动作的预测范围。

4. n_action_steps 敏感性

模型对 n_action_steps 也可能非常敏感。请根据控制频率将其初始值设为约 0.8 秒，然后在此基础上进行调整：

  • 较低的值：响应更灵敏，但在长时程任务中可能稳定性较差
  • 较高的值：更适合长时程执行，但开环故障的恢复能力有限

推理调优

为了获得更快的推理速度，请使用步数更少的 DDIM：

--policy.noise_scheduler_type=DDIM \
--policy.num_inference_steps=10

恢复训练

要从检查点恢复训练：

lerobot-train \
  --config_path=./outputs/mutitask_dit_training/checkpoints/last/pretrained_model/train_config.json \
  --resume=true

检查点目录应包含 model.safetensors 和 config.json 文件（训练期间会自动保存）。恢复训练时，配置将从检查点加载，因此您无需指定其他参数。

常见故障模式与调试

训练这些模型可能会比较棘手。以下是一些常见的故障模式及调试方法：

空闲 / 无运动

模型在推理过程中可能会“坍缩”，导致画面静止或没有运动。这种情况可能发生在以下场景：

  1. 训练数据不足：如果您只有 20-50 个样本，请尝试将数据集大小大致翻倍。如果样本数量超过 300 个后仍然出现此问题，则任务可能过于复杂。

  2. 多个相似任务：当您的数据集包含多个相似任务（例如，抓取 2 个不同的物体）时，模型可能会过度依赖语言条件，而语言条件可能不够丰富。

调试技巧：

  • 增加数据集大小（不断加倍，直到超过 300 个样本）
  • 训练更长时间，最多可达 100k 步，即使损失曲线趋于平稳
  • 检查模型是否接收到了正确的语言指令，或增加指令的多样性

执行了错误的任务

有时机器人会完全忽略您的指令，转而执行其他任务。这种情况通常只会在您使用多个任务进行训练时发生。

可能的原因：

  • 语言指令的歧义
  • 任务特定训练数据不足
  • 多任务数据集中相似任务之间的模型混淆

调试技巧：

  • 验证语言指令的具体性，尤其是在多个任务的描述相似时
  • 检查训练数据集中的任务分布，并为失败/被忽略的任务添加权重
  • 考虑针对特定任务的微调

训练不稳定

如果训练损失不稳定或发散：

  • 尝试将学习率调整到 1e-5 和 3e-4 之间
  • 如果可能，增加批处理大小
  • 检查您的数据集归一化是否正确
  • 验证图像预处理是否正常工作

性能考量

GPU 要求

  • 推理：建议至少使用 RTX 5070 Ti（或同等性能的 GPU）以获得合理的速度表现
  • 训练：理想情况下，应使用显存足以加载大于 64 的批处理大小的 GPU，具体需求取决于图像观测数量等因素。

批量大小建议

  • 最小值：64（低于此值可能导致训练不稳定）
  • 推荐：256-320（性能最佳，需要更大的 GPU）

示例：在自定义数据集上训练

以下是在自定义数据集上进行训练的完整示例：

lerobot-train \
  --dataset.repo_id=YOUR_DATASET \
  --output_dir=./outputs/mutitask_dit_training \
  --batch_size=320 \
  --steps=30000 \
  --save_freq=1000 \
  --log_freq=100 \
  --env_eval_freq=1000 \
  --policy.type=multi_task_dit \
  --policy.device=cuda \
  --policy.horizon=32 \
  --policy.n_action_steps=24 \
  --policy.objective=diffusion \
  --policy.noise_scheduler_type=DDPM \
  --policy.num_layers=6 \
  --policy.hidden_dim=512 \
  --policy.vision_encoder_name=openai/clip-vit-base-patch16 \
  --policy.image_resize_shape=[320,240] \
  --policy.image_crop_shape=[224,224] \
  --policy.repo_id="HF_USER/multitask-dit-your-robot" \
  --trackio.enable=true \
  --trackio.project=multitask_dit

Libero 结果

lerobot-train \
  --dataset.repo_id=HuggingFaceVLA/libero \
  --policy.type=multi_task_dit \
  --policy.push_to_hub=false \
  --output_dir="./outputs/multitask_dit_libero" \
  --job_name="multitask-dit-libero" \
  --trackio.enable=true \
  --trackio.project=multitask_dit_libero \
  --dataset.image_transforms.enable=true \
  --dataset.image_transforms.max_num_transforms=4 \
  --dataset.image_transforms.tfs='{"brightness":{"type":"ColorJitter","kwargs":{"brightness":[0.75,1.25]}},"contrast":{"type":"ColorJitter","kwargs":{"contrast":[0.6,1.4]}},"saturation":{"type":"ColorJitter","kwargs":{"saturation":[0.8,1.2]}},"hue":{"type":"ColorJitter","kwargs":{"hue":[-0.05,0.05]}},"sharpness":{"type":"SharpnessJitter","kwargs":{"sharpness":[0.6,1.4]}},"rotation":{"type":"RandomRotation","kwargs":{"degrees":[-5,5]}},"translation":{"type":"RandomAffine","kwargs":{"degrees":0,"translate":[0.1,0.1]}}}' \
  --dataset.video_backend=torchcodec \
  --policy.use_amp=true \
  --policy.horizon=48 \
  --policy.n_obs_steps=2 \
  --policy.use_rope=true \
  --policy.use_positional_encoding=false \
  --policy.hidden_dim=768 \
  --policy.num_layers=8 \
  --policy.num_heads=12 \
  --policy.dropout=0.1 \
  --policy.timestep_embed_dim=256 \
  --policy.objective=diffusion \
  --policy.optimizer_lr=3e-4 \
  --policy.optimizer_weight_decay=0 \
  --policy.scheduler_warmup_steps=0 \
  --policy.vision_encoder_name=openai/clip-vit-base-patch16 \
  --policy.image_resize_shape=[256,256] \
  --policy.image_crop_is_random=true \
  --policy.text_encoder_name=openai/clip-vit-base-patch16 \
  --policy.vision_encoder_lr_multiplier=0.1 \
  --policy.device=cuda \
  --num_workers=8 \
  --save_freq=4000 \
  --log_freq=100 \
  --steps=100000 \
  --batch_size=320

结果：

LIBERO 空间LIBERO 对象LIBERO 目标LIBERO 10平均
87.098.293.883.290.6

参考

有关技术实现和架构的更多详细信息，请参阅：

Update on GitHub

←X-VLA WALL-OSS→