Политика Multitask DiT

Политика Multitask Diffusion Transformer (DiT) представляет собой развитие исходной архитектуры Diffusion Policy, использующей крупную модель DiT с текстовым и визуальным контекстом для многозадачного обучения роботов. Данная реализация поддерживает как цели диффузии, так и flow matching для генерации действий, что позволяет роботам выполнять разнообразные задачи манипулирования объектами на основе языковых инструкций.

Обзор модели

Модель использует:

• CLIP Vision Encoder: обрабатывает RGB-изображения с нескольких камер
• CLIP Text Encoder: кодирует инструкции языковых задач (замороженные веса с обучаемым проецированием)
• Diffusion Transformer: предсказывает последовательности действий с учётом наблюдений и языка
• Две цели: поддерживает как диффузию (DDPM/DDIM), так и сопоставление потоков для генерации действий

Эта модель интересна тем, что позволяет достичь исключительно высокой ловкости, сопоставимой с показателями VLAs с несколькими миллиардами параметров, используя всего ~450M параметров и значительно меньше данных для обучения.

Требования к установке

Политика Multitask DiT имеет дополнительные зависимости. Установите её с помощью:

pip install 'lerobot[multi_task_dit]'

Это установит все необходимые зависимости, включая библиотеку HuggingFace Transformers для моделей CLIP.

Использование

Чтобы использовать Multitask DiT в вашей конфигурации LeRobot, укажите тип политики как:

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

Рекомендуемые гиперпараметры и детали набора данных (частота управления 30 Гц)

Для стабильной работы начните со следующих рекомендуемых гиперпараметров по умолчанию:

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 — количество шагов действий для предсказания, ~1,0 сек при 30 Гц
• n_action_steps: 24 - ~0,8 секунды при 30 Гц
• Цель: diffusion — начать с диффузии и при низком качестве генерации экспериментировать с flow matching
• Шаги обучения: для одной задачи рекомендуется более 30 тыс. шагов

Параметры конфигурации обучения

Выбор цели

Выберите между диффузией и сопоставлением потоков:

# 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

Архитектура трансформера

Подбирайте ёмкость модели в зависимости от размера набора данных:

# 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. Flow Matching с выборкой по бета-распределению

Исходная реализация диффузии здесь основана на работе, описанной в статье TRI об LBM

Кроме того, мы реализовали цель flow-matching, которая описана на высоком уровне в статье в блоге 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 Гц: horizon=30
• Частота 10 Гц: 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 (сохраняются автоматически во время обучения). При возобновлении конфигурация загружается из контрольной точки, поэтому указывать другие параметры не нужно.

Частые причины сбоев и отладка

Обучение этих моделей может быть капризным. Ниже приведены типичные причины сбоев и подходы к отладке:

Ожидание / Отсутствие движения

Модель может «схлопнуться» во время инференса, что приводит к статичному изображению или отсутствию движения. Это может произойти, когда:

1. Недостаточно обучающих данных: если у вас есть только 20–50 примеров, попробуйте примерно удвоить размер набора данных. Если после достижения 300 примеров проблема сохраняется, задача, возможно, слишком сложна.

2. Несколько похожих задач: если ваш набор данных содержит несколько похожих задач (например, подбор двух разных объектов), модель может чрезмерно полагаться на языковое условие, которое может оказаться недостаточно информативным.

Советы по отладке:

• Увеличьте размер набора данных (удваивайте, пока не получите более 300 примеров)
• Тренируйте дольше, вплоть до 100k шагов, даже если кривая потерь выходит на плато
• Проверьте, получает ли модель корректные языковые инструкции, или увеличьте разнообразие инструкций

Выполнение неправильной задачи

Иногда робот полностью игнорирует вашу инструкцию и выполняет другую задачу. Как правило, это происходит только в том случае, если вы обучали модель на нескольких задачах.

Возможные причины:

• Неоднозначность языковых инструкций
• Недостаточное количество данных для обучения, специфичных для задачи
• Смешение моделей между похожими задачами в многозадачном наборе данных

Советы по отладке:

• Проверьте специфичность языковых инструкций, особенно если описания нескольких задач похожи
• Проверьте распределение задач в наборе данных для обучения и добавьте взвешивание для неудачных/игнорируемых задач
• Рассмотрите дообучение под конкретную задачу

Нестабильность обучения

Если потери обучения нестабильны или расходятся:

• Попробуйте подобрать скорость обучения в диапазоне от 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 \
  --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

Результаты:

Ссылки

Подробнее о технической реализации и архитектуре см. в:

• Внимательное исследование больших поведенческих моделей для многозадачного ловкого манипулирования
• Large Behavior Models и Atlas находят новое применение
• Разбор и публикация исходного кода многозадачного диффузионного трансформерного полиси