多任务 DiT policy

多任务扩散 Transformer（DiT）policy 是原始 Diffusion Policy 架构的演进，它利用一个带文本和视觉条件的大型 DiT 进行多任务机器人学习。该实现同时支持用于 action 生成的扩散和流匹配目标，使机器人能够根据语言指令执行多种操作任务。

模型概览

该模型使用：

CLIP 视觉编码器：处理来自多个相机视角的 RGB 图像
CLIP 文本编码器：编码语言任务指令（冻结权重，带有可学习的投影）
扩散 Transformer：以 observation 和语言为条件预测 action sequence
两种目标：支持用于 action 生成的扩散（DDPM/DDIM）和流匹配

这个模型令人兴奋之处在于，你只需约 450M 参数且训练量显著更少，就能实现极高的灵巧性，与数十亿参数的 VLA 模型相媲美。

安装要求

多任务 DiT policy 有额外的依赖。使用以下命令安装：

pip install lerobot[multi_task_dit]

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

用法

要在你的 LeRobot 配置中使用多任务 DiT，请将 policy 类型指定为：

policy.type=multi_task_dit

训练

基本训练命令

下面是在你的 dataset 上训练多任务 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" \
  --wandb.enable=true

训练配置参数

目标选择

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

# Diffusion objective (default)
--policy.objective=diffusion \
--policy.noise_scheduler_type=DDPM \  # or "DDIM"
--policy.num_train_timesteps=100 \
--policy.num_inference_steps=10 \  # For faster inference
--policy.beta_schedule=squaredcos_cap_v2 \  # Noise schedule type
--policy.prediction_type=epsilon \  # "epsilon" (predict noise) or "sample" (predict clean)
--policy.clip_sample=true \  # Clip samples during denoising
--policy.clip_sample_range=1.0  # Clipping range [-x, x]

# Flow matching objective
--policy.objective=flow_matching \
--policy.timestep_sampling_strategy=beta \  # or "uniform" | the beta sampling strategy performance appears much better in practice
--policy.num_integration_steps=100 \
--policy.integration_method=euler \  # or "rk4"
--policy.sigma_min=0.0  # Minimum noise in flow interpolation path

Transformer 架构

根据 dataset 大小调整模型容量：

# 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

位置编码选项：

该模型支持两种用于 action sequence 的位置编码方法：

# 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
--policy.image_resize_shape=[XXX,YYY] \ # you may need to resize your images for inference speed ups
--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 博客文章中有高层次介绍。

可以考虑测试流匹配目标，并评估你的任务上的性能差异：

--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

这并未被证明在每种使用场景下都是银弹，但它偶尔能产生更平滑、更一致的 action。

2. Transformer 层数

使模型容量匹配你的 dataset 大小：

小 dataset（< 100 个样本）：减少到 4 层
大 dataset（> 5k 个样本）：增加到 8 层

3. horizon 调优

模型可能对你选择的预测长度（horizon）很敏感。根据你的控制频率，从约 1 秒的预测长度开始：

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

然后在此基础上尝试增加。预测长度决定了模型向未来预测 action 多远。

4. n_action_steps 敏感性

该模型对 n_action_steps 也可能非常敏感。根据你的控制频率，先将其设为约 0.8 秒，然后在此基础上调优：

较低的值：反应更快，但对于长时任务可能稳定性较差
较高的值：更适合长时执行，但开环失败时的恢复能力有限

inference 调优

为了更快的 inference，使用更少采样步数的 DDIM：

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

恢复训练

要从 checkpoint 恢复训练：

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

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

常见失败模式与调试

训练这些模型可能会比较棘手。以下是常见的失败模式和调试方法：

静止 / 无 action

模型可能在 inference 期间“坍缩”，导致静止或没有 action。这可能发生在以下情况：

训练数据不足：如果你只有 20-50 个样本，试着将 dataset 规模大约翻倍。一旦你有 300 个以上的样本，如果仍然出现这种情况，任务可能太复杂了。
多个相似任务：当你的 dataset 包含多个相似任务（例如拾取 2 个不同的物体）时，模型可能过度依赖语言条件，而语言条件可能不够丰富。

调试提示：

增加 dataset 规模（不断翻倍，直到超过 300 个样本）
训练更长时间，最多可达 100k 步，即使损失趋于平稳也要继续
检查模型是否接收到正确的语言指令，或增加指令的多样性

执行错误的任务

有时机器人会完全忽略你的指令，转而执行其他任务。这通常只在你训练了多个任务时才会发生。

可能的原因：

语言指令有歧义
特定任务的训练数据不足
模型混淆了多任务 dataset 中相似的任务