# 多任务 DiT 策略

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

## 模型概览

该模型使用：

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

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

## 安装要求

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

```bash
pip install 'lerobot[multi_task_dit]'
```

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

## 用法

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

```python
policy.type=multi_task_dit
```

## 训练

### 基础训练命令

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

```bash
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 控制频率）

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

```bash
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 步

### 训练配置参数

#### 目标选择

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

```bash
# 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_type`       | `DDPM` 或 `DDIM`                                                                 |
| `--policy.num_inference_steps`        | 步骤更少，推理速度更快                                                |
| `--policy.beta_schedule`              | 噪声调度类型                                                              |
| `--policy.prediction_type`            | `epsilon`（预测噪声）或 `sample`（预测干净样本）                            |
| `--policy.clip_sample`                | 在去噪过程中裁剪样本                                                    |
| `--policy.clip_sample_range`          | 裁剪范围 `[-x, x]`                                                         |
| `--policy.timestep_sampling_strategy` | `beta` 或 `uniform`；在实践中，`beta` 策略的效果似乎要好得多 |
| `--policy.integration_method`         | `euler` 或 `rk4`                                                                 |
| `--policy.sigma_min`                  | 流插值路径中的最小噪声                                     |

#### Transformer 架构

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

```bash
# 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
```

**位置编码选项：**

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

```bash
# 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 参数：**

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

#### 视觉编码器配置

```bash
# 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
```

#### 文本编码器配置

```bash
# 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 作为理想的起始点：

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

### 训练调优指南

#### 1. 使用 Beta 采样的流匹配

此处的原始扩散实现基于[TRI 的 LBM 论文](https://arxiv.org/abs/2507.05331)中描述的工作。

此外，我们实现了一种流匹配目标，其高层描述见 [Boston Dynamics 博客文章](https://bostondynamics.com/blog/large-behavior-models-atlas-find-new-footing/)。

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

```bash
--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：

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

### 恢复训练

要从检查点恢复训练：

```bash
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）

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

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

```bash
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.0           | 98.2          | 93.8        | 83.2      | 90.6    |

## 参考

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

- [对多任务灵巧操作大型行为模型的细致考察](https://arxiv.org/abs/2507.05331)
- [大型行为模型和 Atlas 找到了新的立足点](https://bostondynamics.com/blog/large-behavior-models-atlas-find-new-footing/)
- [剖析并开源多任务扩散 Transformer 策略](https://brysonkjones.substack.com/p/dissecting-and-open-sourcing-multitask-diffusion-transformer-policy)