# Политика 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 имеет дополнительные зависимости. Установите её с помощью: ```bash pip install 'lerobot[multi_task_dit]' ``` Это установит все необходимые зависимости, включая библиотеку HuggingFace Transformers для моделей CLIP. ## Использование Чтобы использовать Multitask DiT в вашей конфигурации LeRobot, укажите тип политики как: ```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 ``` ### Рекомендуемые гиперпараметры и детали набора данных (частота управления 30 Гц) Для стабильной работы начните со следующих рекомендуемых гиперпараметров по умолчанию: ```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 — количество шагов действий для предсказания, ~1,0 сек при 30 Гц - **n_action_steps**: 24 - ~0,8 секунды при 30 Гц - **Цель**: `diffusion` — начать с диффузии и при низком качестве генерации экспериментировать с flow matching - **Шаги обучения**: для одной задачи рекомендуется более 30 тыс. шагов ### Параметры конфигурации обучения #### Выбор цели Выберите между диффузией и сопоставлением потоков: ```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` | Минимальный шум на пути интерполяции потока | #### Архитектура трансформера Подбирайте ёмкость модели в зависимости от размера набора данных: ```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. Flow Matching с выборкой по бета-распределению Исходная реализация диффузии здесь основана на работе, описанной в [статье TRI об LBM](https://arxiv.org/abs/2507.05331) Кроме того, мы реализовали цель flow-matching, которая описана на высоком уровне в [статье в блоге 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 Гц**: `horizon=30` - **Частота 10 Гц**: `horizon=10` Затем экспериментируйте с увеличением этого значения. Горизонт определяет, насколько далеко в будущее модель предсказывает действия. #### 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. **Несколько похожих задач**: если ваш набор данных содержит несколько похожих задач (например, подбор двух разных объектов), модель может чрезмерно полагаться на языковое условие, которое может оказаться недостаточно информативным. **Советы по отладке:** - Увеличьте размер набора данных (удваивайте, пока не получите более 300 примеров) - Тренируйте дольше, вплоть до 100k шагов, даже если кривая потерь выходит на плато - Проверьте, получает ли модель корректные языковые инструкции, или увеличьте разнообразие инструкций ### Выполнение неправильной задачи Иногда робот полностью игнорирует вашу инструкцию и выполняет другую задачу. Как правило, это происходит только в том случае, если вы обучали модель на нескольких задачах. **Возможные причины:** - Неоднозначность языковых инструкций - Недостаточное количество данных для обучения, специфичных для задачи - Смешение моделей между похожими задачами в многозадачном наборе данных **Советы по отладке:** - Проверьте специфичность языковых инструкций, особенно если описания нескольких задач похожи - Проверьте распределение задач в наборе данных для обучения и добавьте взвешивание для неудачных/игнорируемых задач - Рассмотрите дообучение под конкретную задачу ### Нестабильность обучения Если потери обучения нестабильны или расходятся: - Попробуйте подобрать скорость обучения в диапазоне от `1e-5` до `3e-4` - Увеличьте размер пакета, если это возможно - Проверьте, что нормализация вашего набора данных выполнена правильно - Проверьте, что предобработка изображений работает корректно ## Вопросы производительности ### Требования к GPU - **Вывод**: Для приемлемой скорости рекомендуется как минимум RTX 5070 Ti (или эквивалентный GPU) - **Обучение**: идеальным вариантом является GPU с достаточным объёмом VRAM для загрузки размеров батчей >64, что зависит от количества изображений-наблюдений и т. д. ### Рекомендации по размеру батча - **Минимум**: 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 Spatial | LIBERO Object | LIBERO Цель | LIBERO 10 | Среднее | | -------------- | ------------- | ----------- | --------- | ------- | | 87.0 | 98.2 | 93.8 | 83.2 | 90.6 | ## Ссылки Подробнее о технической реализации и архитектуре см. в: - [Внимательное исследование больших поведенческих моделей для многозадачного ловкого манипулирования](https://arxiv.org/abs/2507.05331) - [Large Behavior Models и Atlas находят новое применение](https://bostondynamics.com/blog/large-behavior-models-atlas-find-new-footing/) - [Разбор и публикация исходного кода многозадачного диффузионного трансформерного полиси](https://brysonkjones.substack.com/p/dissecting-and-open-sourcing-multitask-diffusion-transformer-policy)