LeRobot 文件

多任務 DiT 策略

LeRobot

加入 Hugging Face 社群

並獲得增強的文件體驗

在模型、資料集和 Spaces 上進行協作

透過加速推論實現更快的範例

切換文件主題

開始使用

多工 DiT 策略

多工擴散變換器 (DiT) 策略是原始擴散策略架構的演進，它利用大型 DiT 並結合文本和視覺條件，實現多工機器人學習。此實作同時支援擴散（diffusion）和流匹配（flow matching）目標，用於動作生成，使機器人能夠根據語言指令執行各種操作任務。

模型概覽

此模型使用

CLIP 視覺編碼器：處理來自多個攝影機視角的 RGB 圖像
CLIP 文字編碼器：編碼語言任務指令（凍結權重與可學習投影）
擴散變換器：根據觀察和語言預測動作序列
兩種目標：支援擴散 (DDPM/DDIM) 和流匹配兩種動作生成目標

這個模型令人興奮，因為它僅需約 4.5 億個參數和顯著減少的訓練時間，就能達到與數十億參數的 VLA 競爭的極高靈巧度。

安裝需求

多工 DiT 策略有額外的依賴項。請使用以下指令安裝：

pip install lerobot[multi_task_dit]

這將安裝所有必要的依賴項，包括用於 CLIP 模型之 HuggingFace Transformers 函式庫。

使用方式

若要在您的 LeRobot 配置中使用多工 DiT，請將策略類型指定為

policy.type=multi_task_dit

訓練

基本訓練指令

以下是在您的資料集上訓練多工 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

建議的超參數和資料集詳細資訊 (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" \
  --wandb.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 \  # 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

變換器架構

根據資料集大小調整模型容量

# 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

其他變換器參數

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

這尚未被證明適用於所有使用者案例的萬靈丹，但它偶爾會帶來更流暢、更一致的動作。

2. 變換器層數

根據您的資料集大小匹配模型容量

小型資料集 (< 100 個範例)：減少到 4 層
大型資料集 (> 5k 個範例)：增加到 8 層

3. 時間範圍 (Horizon) 調校

模型對您選擇的時間範圍可能很敏感。根據您的控制頻率，從大約 1 秒的時間範圍開始：

30 Hz 頻率：horizon=30
10 Hz 頻率：horizon=10

然後從那裡開始嘗試增加。時間範圍決定了模型預測未來動作的距離。

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 檔案（在訓練期間自動儲存）。恢復時，配置會從檢查點載入，因此您無需指定其他參數。

常見故障模式與除錯

訓練這些模型可能很棘手。以下是常見的故障模式和除錯方法：

閒置 / 無動作

模型在推論期間可能會「崩潰」，導致靜止或無動作。這可能發生在以下情況：

訓練資料不足：如果您只有 20-50 個範例，請嘗試將資料集大小約翻倍。一旦超過 300 個範例仍出現此情況，則任務可能過於複雜。
多個相似任務：當您的資料集包含多個相似任務（例如，拾取 2 個不同的物體）時，模型可能會過度依賴語言條件，而這些條件可能不夠豐富。

除錯技巧

增加資料集大小（翻倍直到超過 300 個範例）
訓練更長時間，最多 10 萬步，即使損失已趨平穩。
檢查模型是否收到正確的語言指令，或增加指令的多樣性。

執行錯誤的任務

有時機器人會完全忽略您的指令並執行其他任務。這通常只發生在您已訓練多個任務的情況下。

潛在原因

語言指令模糊
任務特定訓練資料不足
模型在多工資料集中混淆相似任務

除錯技巧

驗證語言指令的具體性，特別是當多個任務的描述相似時。
檢查訓練資料集中的任務分佈，並對失敗/被忽略的任務增加權重。
考慮任務特定的微調

訓練不穩定

如果訓練損失不穩定或發散

嘗試在 1e-5 到 3e-4 之間調整學習率
如果可能，增加批次大小
檢查您的資料集正規化是否正確
驗證圖像預處理是否正常運作

性能考量

GPU 要求

推論：建議至少使用 RTX 5070 Ti（或同等 GPU）以獲得合理的運行速度效能。
訓練：理想情況下，GPU 應具備足夠的 VRAM 來載入批次大小 >64 的資料，這將根據圖像觀察數量等因素而異。

批次大小建議

最低：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 \
  --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" \
  --wandb.enable=true \
  --wandb.project=multitask_dit

Libero 結果

python -m lerobot.scripts.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" \
  --wandb.enable=true \
  --wandb.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 Spatial (空間)	LIBERO Object (物件)	LIBERO Goal (目標)	LIBERO 10	平均
87.0	98.2	93.8	83.2	90.6

參考資料

有關技術實作和架構的更多詳細資訊，請參閱

在 GitHub 上更新

←X-VLA WALL-OSS→