Перейти к основному содержимому

Тонкая настройка модели Diffusion Policy

Обзор​

Diffusion Policy - это метод обучения визуомоторной политики на основе диффузионных моделей, применяющий генеративные возможности диффузионных моделей в области управления роботами. Этот метод генерирует разнообразные и высококачественные последовательности действий роботов путем изучения процесса диффузии распределения действий, демонстрируя отличные результаты в сложных задачах манипуляции роботов.

Основные особенности​

  • Диффузионная генерация: Использует диффузионные модели для генерации непрерывных последовательностей действий
  • Мультимодальные действия: Может обрабатывать задачи с несколькими решениями
  • Высококачественный вывод: Генерирует плавные и естественные действия роботов
  • Сильная робастность: Хорошая робастность к шуму и возмущениям
  • Сильная выразительность: Может изучать сложные распределения действий

Предварительные требования​

Системные требования​

  • Операционная система: Linux (рекомендуется Ubuntu 20.04+) или macOS
  • Версия Python: 3.8+
  • GPU: NVIDIA GPU (рекомендуется RTX 3080 или выше), минимум 10GB видеопамяти
  • Память: Минимум 32GB RAM
  • Дисковое пространство: Минимум 50GB свободного места

Подготовка среды​

1. Установка LeRobot​

# Клонировать репозиторий LeRobot
git clone https://github.com/huggingface/lerobot.git
cd lerobot

# Создать виртуальную среду
conda create -n lerobot python=3.10
conda activate lerobot

# Установить зависимости
pip install -e .

2. Установка специфичных зависимостей Diffusion Policy​

# Установить зависимости, связанные с диффузионными моделями
pip install diffusers
pip install accelerate
pip install transformers
pip install einops
pip install wandb

# Установить библиотеки численных вычислений
pip install scipy
pip install scikit-learn

# Войти в Weights & Biases (опционально)
wandb login

Архитектура Diffusion Policy​

Основные компоненты​

  1. Кодировщик изображений: Извлекает признаки изображений
  2. Кодировщик состояний: Обрабатывает информацию о состоянии робота
  3. Условный кодировщик: Объединяет визуальную информацию и информацию о состоянии
  4. Диффузионная сеть: Изучает процесс диффузии распределения действий
  5. Планировщик шума: Управляет уровнями шума в процессе диффузии

Процесс диффузии​

  1. Прямой процесс: Постепенно добавляет шум к последовательностям действий
  2. Обратный процесс: Постепенно восстанавливает последовательности действий из шума
  3. Условная генерация: Генерирует действия на основе условий наблюдения
  4. Стратегия выборки: Использует выборку DDPM или DDIM

Подготовка данных​

Данные формата LeRobot​

Diffusion Policy требует использования наборов данных формата LeRobot:

your_dataset/
├── data/
│   ├── chunk-001/
│   │   ├── observation.images.cam_high.png
│   │   ├── observation.images.cam_low.png
│   │   ├── observation.state.npy
│   │   ├── action.npy
│   │   └── ...
│   └── chunk-002/
│       └── ...
├── meta.json
├── stats.safetensors
└── videos/
    ├── episode_000000.mp4
    └── ...

Требования к качеству данных​

  • Минимум 100 эпизодов для базового обучения
  • Рекомендуется 500+ эпизодов для оптимальных результатов
  • Последовательности действий должны быть плавными и непрерывными
  • Включите разнообразные сценарии задач
  • Высококачественные данные визуальных наблюдений

Тонкая настройка обучения​

Базовая команда обучения​

# Установить переменные среды
export HF_USER="your-huggingface-username"
export CUDA_VISIBLE_DEVICES=0

# Запустить обучение Diffusion Policy
lerobot-train \
  --policy.type diffusion \
  --policy.pretrained_path lerobot/diffusion_policy \
  --dataset.repo_id ${HF_USER}/your_dataset \
  --batch_size 64 \
  --steps 100000 \
  --output_dir outputs/train/diffusion_policy_finetuned \
  --job_name diffusion_policy_finetuning \
  --policy.device cuda \
  --policy.horizon 16 \
  --policy.n_action_steps 8 \
  --policy.n_obs_steps 2 \
  --policy.num_inference_steps 100 \
  --policy.optimizer_lr 1e-4 \
  --policy.optimizer_weight_decay 1e-6 \
  --policy.push_to_hub false \
  --save_checkpoint true \
  --save_freq 10000 \
  --wandb.enable true

Расширенная конфигурация обучения​

Конфигурация многошагового предсказания​

# Конфигурация для предсказания длинных последовательностей
lerobot-train \
  --policy.type diffusion \
  --policy.pretrained_path lerobot/diffusion_policy \
  --dataset.repo_id ${HF_USER}/your_dataset \
  --batch_size 32 \
  --steps 150000 \
  --output_dir outputs/train/diffusion_policy_long_horizon \
  --job_name diffusion_policy_long_horizon \
  --policy.device cuda \
  --policy.horizon 32 \
  --policy.n_action_steps 16 \
  --policy.n_obs_steps 4 \
  --policy.num_inference_steps 100 \
  --policy.beta_schedule squaredcos_cap_v2 \
  --policy.clip_sample true \
  --policy.prediction_type epsilon \
  --policy.optimizer_lr 1e-4 \
  --policy.scheduler_name cosine \
  --policy.scheduler_warmup_steps 5000 \
  --policy.push_to_hub false \
  --save_checkpoint true \
  --wandb.enable true

Конфигурация оптимизации памяти​

# Для GPU с меньшей видеопамятью
lerobot-train \
  --policy.type diffusion \
  --policy.pretrained_path lerobot/diffusion_policy \
  --dataset.repo_id ${HF_USER}/your_dataset \
  --batch_size 16 \
  --steps 200000 \
  --output_dir outputs/train/diffusion_policy_memory_opt \
  --job_name diffusion_policy_memory_optimized \
  --policy.device cuda \
  --policy.horizon 16 \
  --policy.n_action_steps 8 \
  --policy.num_inference_steps 50 \
  --policy.optimizer_lr 5e-5 \
  --policy.use_amp true \
  --num_workers 2 \
  --policy.push_to_hub false \
  --save_checkpoint true \
  --wandb.enable true

Подробное описание параметров​

Основные параметры​

ПараметрЗначениеРекомендуемое значениеОписание
--policy.typeТип политикиdiffusionТип модели Diffusion Policy
--policy.pretrained_pathПуть к предобученной моделиlerobot/diffusion_policyОфициальная модель LeRobot (опционально)
--dataset.repo_idID репозитория набора данных${HF_USER}/datasetВаш набор данных HuggingFace
--batch_sizeРазмер пакета64Настроить в соответствии с видеопамятью, RTX 3080 рекомендуется 32-64
--stepsШаги обучения100000Диффузионные модели обычно требуют больше шагов обучения
--output_dirВыходной каталогoutputs/train/diffusion_policy_finetunedПуть сохранения модели
--job_nameИмя задачиdiffusion_policy_finetuningДля логирования и отслеживания экспериментов (опционально)

Специфичные для Diffusion Policy параметры​

ПараметрЗначениеРекомендуемое значениеОписание
--policy.horizonГоризонт предсказания16Длина предсказываемой последовательности действий
--policy.n_action_stepsШаги выполнения действий8Количество выполняемых действий каждый раз
--policy.n_obs_stepsШаги наблюдения2Количество исторических наблюдений
--policy.num_inference_stepsШаги вывода100Количество шагов диффузионной выборки (не действует во время обучения)
--policy.beta_scheduleРасписание шумаsquaredcos_cap_v2Стратегия планирования добавления шума
--policy.clip_sampleОбрезка образцовtrueОбрезать ли сгенерированные образцы
--policy.clip_sample_rangeДиапазон обрезки1.0Диапазон для обрезки образцов
--policy.prediction_typeТип предсказанияepsilonПредсказывать шум или образец
--policy.num_train_timestepsВременные шаги обучения100Количество шагов прямой диффузии

Параметры архитектуры сети​

ПараметрЗначениеРекомендуемое значениеОписание
--policy.vision_backboneВизуальный остовresnet18Сеть извлечения признаков изображений
--policy.crop_shapeРазмер обрезки изображения84 84Размер обрезки для входных изображений
--policy.crop_is_randomСлучайная обрезкаtrueСлучайная обрезка во время обучения
--policy.use_group_normИспользовать групповую нормализациюtrueЗаменить пакетную нормализацию
--policy.spatial_softmax_num_keypointsКоличество ключевых точек пространственного softmax32Количество ключевых точек в слое пространственного softmax
--policy.down_dimsРазмерности понижающей дискретизации512 1024 2048Размерности пути понижающей дискретизации U-Net
--policy.kernel_sizeРазмер ядра свертки5Размер ядра для 1D свертки
--policy.n_groupsКоличество групп нормализации8Количество групп в GroupNorm
--policy.diffusion_step_embed_dimРазмерность встраивания шага128Размерность встраивания для шагов диффузии

Параметры обучения​

ПараметрЗначениеРекомендуемое значениеОписание
--policy.optimizer_lrСкорость обучения1e-4Рекомендуемая скорость обучения для диффузионных моделей
--policy.optimizer_weight_decayЗатухание весов1e-6Параметр регуляризации
--policy.optimizer_betasБета оптимизатора Adam0.95 0.999Параметры beta для оптимизатора Adam
--policy.optimizer_epsЭпсилон Adam1e-8Параметр численной стабильности
--policy.scheduler_nameПланировщик скорости обученияcosineРасписание косинусного затухания
--policy.scheduler_warmup_stepsШаги прогрева500Прогрев скорости обучения
--policy.use_ampСмешанная точностьtrueЭкономит видеопамять
--num_workersПотоки загрузки данных4Настроить в соответствии с количеством ядер CPU
--policy.push_to_hubОтправить в HubfalseЗагружать ли модель в HuggingFace (требуется repo_id)
--save_checkpointСохранить контрольные точкиtrueСохранять ли контрольные точки обучения
--save_freqЧастота сохранения10000Интервал сохранения контрольных точек

Мониторинг и отладка обучения​

Интеграция с Weights & Biases​

# Подробная конфигурация W&B
lerobot-train \
  --policy.type diffusion \
  --dataset.repo_id your-name/your-dataset \
  --batch_size 64 \
  --steps 100000 \
  --policy.push_to_hub false \
  --wandb.enable true \
  --wandb.project diffusion_policy_experiments \
  --wandb.notes "Diffusion Policy training with long horizon" \
  # ... другие параметры

Мониторинг ключевых метрик​

Метрики для отслеживания во время обучения:

  • Diffusion Loss: Общая потеря диффузионной модели
  • MSE Loss: Потеря среднеквадратичной ошибки
  • Learning Rate: Кривая изменения скорости обучения
  • Gradient Norm: Норма градиента
  • Inference Time: Время вывода
  • Sample Quality: Качество сгенерированных образцов

Визуализация обучения​

# visualization.py
import torch
import matplotlib.pyplot as plt
import numpy as np
from lerobot.policies.diffusion.modeling_diffusion import DiffusionPolicy

def visualize_diffusion_process(model_path, observation):
    # Загрузить модель
    policy = DiffusionPolicy.from_pretrained(model_path, device="cuda")
    policy.eval()
    
    # Генерировать процесс диффузии последовательности действий
    with torch.no_grad():
        # Начальный шум
        noise = torch.randn(1, policy.horizon, policy.action_dim, device="cuda")
        
        # Процесс диффузионной выборки
        actions_sequence = []
        for t in range(policy.num_inference_steps):
            # Предсказать шум
            noise_pred = policy.unet(noise, t, observation)
            
            # Обновить образец
            noise = policy.scheduler.step(noise_pred, t, noise).prev_sample
            
            # Сохранить промежуточные результаты
            if t % 10 == 0:
                actions_sequence.append(noise.cpu().numpy())
        
        final_actions = noise.cpu().numpy()
    
    # Визуализировать процесс диффузии
    fig, axes = plt.subplots(2, 3, figsize=(15, 10))
    
    for i, actions in enumerate(actions_sequence[:6]):
        ax = axes[i//3, i%3]
        ax.plot(actions[0, :, 0], label='Action Dim 0')
        ax.plot(actions[0, :, 1], label='Action Dim 1')
        ax.set_title(f'Diffusion Step {i*10}')
        ax.legend()
    
    plt.tight_layout()
    plt.savefig('diffusion_process.png')
    plt.show()
    
    return final_actions

if __name__ == "__main__":
    model_path = "outputs/train/diffusion_policy_finetuned/checkpoints/last"
    
    # Симулировать наблюдение
    observation = {
        "observation.images.cam_high": torch.randn(1, 3, 224, 224, device="cuda"),
        "observation.state": torch.randn(1, 7, device="cuda")
    }
    
    actions = visualize_diffusion_process(model_path, observation)
    print(f"Generated actions shape: {actions.shape}")

Оценка модели​

Офлайн оценка​

# offline_evaluation.py
import torch
import numpy as np
from lerobot.policies.diffusion.modeling_diffusion import DiffusionPolicy
from lerobot.datasets.lerobot_dataset import LeRobotDataset

def evaluate_diffusion_policy(model_path, dataset_path):
    # Загрузить модель
    policy = DiffusionPolicy.from_pretrained(model_path, device="cuda")
    policy.eval()
    
    # Загрузить тестовый набор данных
    dataset = LeRobotDataset(dataset_path, split="test")
    
    total_mse_loss = 0
    total_mae_loss = 0
    num_samples = 0
    
    with torch.no_grad():
        for batch in dataset:
            # Предсказание модели
            prediction = policy(batch)
            
            # Вычислить потерю
            target_actions = batch['action']
            predicted_actions = prediction['action']
            
            mse_loss = torch.mean((predicted_actions - target_actions) ** 2)
            mae_loss = torch.mean(torch.abs(predicted_actions - target_actions))
            
            total_mse_loss += mse_loss.item()
            total_mae_loss += mae_loss.item()
            num_samples += 1
    
    avg_mse_loss = total_mse_loss / num_samples
    avg_mae_loss = total_mae_loss / num_samples
    
    print(f"Average MSE Loss: {avg_mse_loss:.4f}")
    print(f"Average MAE Loss: {avg_mae_loss:.4f}")
    
    return avg_mse_loss, avg_mae_loss

def evaluate_action_diversity(model_path, observation, num_samples=10):
    # Оценить разнообразие действий
    policy = DiffusionPolicy.from_pretrained(model_path, device="cuda")
    policy.eval()
    
    actions_list = []
    
    with torch.no_grad():
        for _ in range(num_samples):
            prediction = policy(observation)
            actions_list.append(prediction['action'].cpu().numpy())
    
    actions_array = np.array(actions_list)  # [num_samples, horizon, action_dim]
    
    # Вычислить метрику разнообразия действий
    action_std = np.std(actions_array, axis=0)  # [horizon, action_dim]
    avg_std = np.mean(action_std)
    
    print(f"Average action standard deviation: {avg_std:.4f}")
    
    return avg_std, actions_array

if __name__ == "__main__":
    model_path = "outputs/train/diffusion_policy_finetuned/checkpoints/last"
    dataset_path = "path/to/your/test/dataset"
    
    # Офлайн оценка
    evaluate_diffusion_policy(model_path, dataset_path)
    
    # Оценка разнообразия
    observation = {
        "observation.images.cam_high": torch.randn(1, 3, 224, 224, device="cuda"),
        "observation.state": torch.randn(1, 7, device="cuda")
    }
    
    evaluate_action_diversity(model_path, observation)

Онлайн оценка (среда робота)​

# robot_evaluation.py
import torch
import numpy as np
from lerobot.policies.diffusion.modeling_diffusion import DiffusionPolicy

class DiffusionPolicyController:
    def __init__(self, model_path, num_inference_steps=50):
        self.policy = DiffusionPolicy.from_pretrained(model_path, device="cuda")
        self.policy.eval()
        self.num_inference_steps = num_inference_steps
        self.action_queue = []
        self.current_obs_history = []
        
    def get_action(self, observations):
        # Обновить историю наблюдений
        self.current_obs_history.append(observations)
        if len(self.current_obs_history) > self.policy.n_obs_steps:
            self.current_obs_history.pop(0)
        
        # Если очередь действий пуста или требуется переplanирование, сгенерировать новую последовательность действий
        if len(self.action_queue) == 0 or self.should_replan():
            with torch.no_grad():
                # Построить вход
                batch = self.prepare_observation_batch()
                
                # Установить шаги вывода
                self.policy.scheduler.set_timesteps(self.num_inference_steps)
                
                # Сгенерировать последовательность действий
                prediction = self.policy(batch)
                actions = prediction['action'].cpu().numpy()[0]  # [horizon, action_dim]
                
                # Обновить очередь действий
                self.action_queue = list(actions[:self.policy.n_action_steps])
        
        # Вернуть следующее действие
        return self.action_queue.pop(0)
    
    def should_replan(self):
        # Простая стратегия переplanирования: переplanировать, когда остается менее половины очереди действий
        return len(self.action_queue) < self.policy.n_action_steps // 2
    
    def prepare_observation_batch(self):
        batch = {}
        
        # Обработать наблюдения изображений
        if "observation.images.cam_high" in self.current_obs_history[-1]:
            images = []
            for obs in self.current_obs_history:
                image = obs["observation.images.cam_high"]
                image_tensor = self.preprocess_image(image)
                images.append(image_tensor)
            
            # Если истории недостаточно, повторить последнее наблюдение
            while len(images) < self.policy.n_obs_steps:
                images.insert(0, images[0])
            
            batch["observation.images.cam_high"] = torch.stack(images).unsqueeze(0)
        
        # Обработать наблюдения состояний
        if "observation.state" in self.current_obs_history[-1]:
            states = []
            for obs in self.current_obs_history:
                state = torch.tensor(obs["observation.state"], dtype=torch.float32)
                states.append(state)
            
            # Если истории недостаточно, повторить последнее состояние
            while len(states) < self.policy.n_obs_steps:
                states.insert(0, states[0])
            
            batch["observation.state"] = torch.stack(states).unsqueeze(0)
        
        return batch
    
    def preprocess_image(self, image):
        # Логика предобработки изображения
        image_tensor = torch.tensor(image).permute(2, 0, 1).float() / 255.0
        return image_tensor

# Пример использования
if __name__ == "__main__":
    controller = DiffusionPolicyController(
        model_path="outputs/train/diffusion_policy_finetuned/checkpoints/last",
        num_inference_steps=50
    )
    
    # Симуляция цикла управления роботом
    for step in range(100):
        # Получить текущее наблюдение
        observations = {
            "observation.images.cam_high": np.random.randint(0, 255, (224, 224, 3)),
            "observation.state": np.random.randn(7)
        }
        
        # Получить действие
        action = controller.get_action(observations)
        
        # Выполнить действие
        print(f"Step {step}: Action = {action}")
        
        # Это должно отправить действие реальному роботу
        # robot.execute_action(action)

Развертывание и оптимизация​

Ускорение вывода​

# fast_inference.py
import torch
from lerobot.policies.diffusion.modeling_diffusion import DiffusionPolicy
from diffusers import DDIMScheduler

class FastDiffusionInference:
    def __init__(self, model_path, num_inference_steps=10):
        self.policy = DiffusionPolicy.from_pretrained(model_path, device="cuda")
        self.policy.eval()
        
        # Использовать планировщик DDIM для быстрой выборки
        self.policy.scheduler = DDIMScheduler.from_config(self.policy.scheduler.config)
        self.num_inference_steps = num_inference_steps
        
        # Прогреть модель
        self.warmup()
    
    def warmup(self):
        # Прогреть модель с фиктивными данными
        dummy_batch = {
            "observation.images.cam_high": torch.randn(1, 2, 3, 224, 224, device="cuda"),
            "observation.state": torch.randn(1, 2, 7, device="cuda")
        }
        
        with torch.no_grad():
            for _ in range(5):
                _ = self.predict(dummy_batch)
    
    @torch.no_grad()
    def predict(self, observations):
        # Установить шаги вывода
        self.policy.scheduler.set_timesteps(self.num_inference_steps)
        
        # Быстрый вывод
        prediction = self.policy(observations)
        return prediction['action'].cpu().numpy()

if __name__ == "__main__":
    fast_inference = FastDiffusionInference(
        "outputs/train/diffusion_policy_finetuned/checkpoints/last",
        num_inference_steps=10
    )
    
    # Тестировать скорость вывода
    import time
    
    observations = {
        "observation.images.cam_high": torch.randn(1, 2, 3, 224, 224, device="cuda"),
        "observation.state": torch.randn(1, 2, 7, device="cuda")
    }
    
    start_time = time.time()
    for _ in range(100):
        action = fast_inference.predict(observations)
    end_time = time.time()
    
    avg_inference_time = (end_time - start_time) / 100
    print(f"Average inference time: {avg_inference_time:.4f} seconds")
    print(f"Inference frequency: {1/avg_inference_time:.2f} Hz")

Лучшие практики​

Рекомендации по сбору данных​

  1. Плавные действия: Обеспечьте плавность и непрерывность последовательностей действий в демонстрационных данных
  2. Разнообразные сценарии: Собирайте данные с различными начальными состояниями и целями
  3. Высококачественные аннотации: Обеспечьте точность аннотаций действий
  4. Достаточный объем данных: Диффузионные модели обычно требуют больше данных

Рекомендации по оптимизации обучения​

  1. Планирование шума: Выберите соответствующую стратегию планирования шума
  2. Шаги вывода: Сбалансируйте качество и скорость, выберите соответствующие шаги вывода
  3. Планирование скорости обучения: Используйте косинусное затухание или поэтапное снижение
  4. Регуляризация: Соответствующим образом используйте затухание весов

Рекомендации по оптимизации развертывания​

  1. Быстрая выборка: Используйте DDIM или другие методы быстрой выборки
  2. Сжатие модели: Используйте методы дистилляции знаний или квантования
  3. Параллельный вывод: Используйте возможности параллелизма GPU
  4. Оптимизация кэша: Кэшируйте промежуточные результаты вычислений

Часто задаваемые вопросы (FAQ)​

Q: Какие преимущества имеет Diffusion Policy по сравнению с другими методами обучения политики?​

A: Основные преимущества Diffusion Policy включают:

  • Мультимодальная генерация: Может обрабатывать задачи с несколькими решениями
  • Высококачественный вывод: Генерирует плавные и естественные последовательности действий
  • Сильная робастность: Хорошая робастность к шуму и возмущениям
  • Сильная выразительность: Может изучать сложные распределения действий

Q: Как выбрать соответствующее количество шагов вывода?​

A: Выбор шагов вывода требует баланса между качеством и скоростью:

  • Высокое качество: 100-1000 шагов, подходит для офлайн оценки
  • Приложения реального времени: 10-50 шагов, подходит для онлайн управления
  • Быстрое прототипирование: 5-10 шагов, подходит для быстрого тестирования

Q: Сколько времени занимает обучение?​

A: Время обучения зависит от нескольких факторов:

  • Размер набора данных: 500 эпизодов занимают примерно 12-24 часа (RTX 3080)
  • Сложность модели: Более крупные модели требуют больше времени
  • Шаги вывода: Больше шагов увеличивает время обучения
  • Требования к сходимости: Обычно требуется 100000-200000 шагов

Q: Как улучшить качество сгенерированных действий?​

A: Методы улучшения качества действий:

  • Увеличить шаги вывода: Больше шагов обычно производят лучшие результаты
  • Оптимизировать планирование шума: Выбрать соответствующую стратегию добавления шума
  • Качество данных: Обеспечить высокое качество обучающих данных
  • Архитектура модели: Использовать более крупные или глубокие сети
  • Техники регуляризации: Соответствующая регуляризация предотвращает переобучение

Q: Как обрабатывать требования реального времени?​

A: Методы удовлетворения требований реального времени:

  • Быстрая выборка: Использовать DDIM или DPM-Solver
  • Уменьшить шаги вывода: Найти баланс между качеством и скоростью
  • Дистилляция модели: Обучить меньшие студенческие модели
  • Параллельный вывод: Использовать несколько GPU или пакетную обработку
  • Предварительные вычисления: Предварительно вычислить частичные результаты

Связанные ресурсы​

Журнал изменений​

  • 2024-01: Релиз начальной версии
  • 2024-02: Добавлена поддержка быстрой выборки
  • 2024-03: Оптимизировано использование памяти и эффективность обучения
  • 2024-04: Добавлены оценка разнообразия и оптимизация развертывания