Дообучение модели ACT (Action Chunking Transformer)

Обзор

ACT (Action Chunking Transformer) — это модель сквозного имитационного обучения (end-to-end imitation learning), разработанная для высокоточных манипуляционных задач. Предсказывая наборы действий (action chunks), модель преодолевает проблему накопления ошибок (compounding errors) в традиционном имитационном обучении, обеспечивая высокий уровень успеха робототехнических операций на недорогом оборудовании.

Основные характеристики

• Предсказание наборов действий (Action Chunking): предсказывает сразу несколько последовательных действий, уменьшая накопление ошибок.
• Архитектура Transformer: использует механизмы внимания (attention) для обработки последовательных данных.
• Сквозное обучение (End-to-End): предсказывает действия напрямую из необработанных данных наблюдений.
• Высокий показатель успеха: демонстрирует отличные результаты в задачах точного манипулирования.
• Дружелюбность к оборудованию: может работать на потребительском «железе».

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

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

• ОС: Linux (рекомендуется Ubuntu 20.04+) или macOS.
• Версия Python: 3.8+.
• GPU: NVIDIA GPU (рекомендуется RTX 3070 или выше) с минимум 6 ГБ видеопамяти.
• ОЗУ: минимум 16 ГБ.
• Место на диске: минимум 30 ГБ свободного пространства.

Подготовка окружения

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

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

# Создание виртуального окружения (рекомендуется venv или conda)
python -m venv .venv
source .venv/bin/activate
python -m pip install -U pip

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

2. Установка специфических зависимостей для ACT

# Установка дополнительных пакетов
pip install einops
pip install timm
pip install wandb

# Авторизация в Weights & Biases (опционально)
wandb login

Архитектура модели ACT

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

1. Визуальный энкодер: обрабатывает изображения с нескольких камер.
2. Энкодер состояния: обрабатывает информацию о состоянии робота.
3. Декодер Transformer: генерирует последовательность действий.
4. Голова действий (Action Head): выдает финальное предсказание действий.

Ключевые параметры

• Размер набора (Chunk Size): количество действий, предсказываемых за один раз (обычно 50-100).
• Длина контекста (Context Length): длина истории наблюдений.
• Скрытая размерность (Hidden Dimension): размерность скрытых слоев Transformer.
• Количество голов (Number of Heads): количество голов внимания.
• Количество слоев (Number of Layers): количество слоев в Transformer.

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

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

ACT требует наборы данных в формате LeRobot со следующей структурой:

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

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

• Минимум 50 эпизодов для базового обучения.
• Рекомендуется более 200 эпизодов для достижения оптимальных результатов.
• Каждый эпизод должен содержать полное выполнение задачи.
• Изображения с нескольких ракурсов (минимум 2 камеры).
• Высококачественная разметка действий.

Обучение (Fine-tuning)

Ограничение по параметрам

Параметр n_action_steps для модели ACT должен быть ≤ chunk_size. Рекомендуется устанавливать оба значения равными (например, 100).

• chunk_size: длина последовательности действий, предсказываемой моделью за один раз.
• n_action_steps: количество шагов действий, фактически выполняемых роботом.

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

# Установка переменных окружения
export HF_USER="your-huggingface-username"
export CUDA_VISIBLE_DEVICES=0

# Запуск обучения ACT
lerobot-train \
  --policy.type act \
  --dataset.repo_id ${HF_USER}/your_dataset \
  --batch_size 8 \
  --steps 50000 \
  --output_dir outputs/train/act_finetuned \
  --job_name act_finetuning \
  --policy.device cuda \
  --policy.chunk_size 100 \
  --policy.n_action_steps 100 \
  --policy.n_obs_steps 1 \
  --policy.optimizer_lr 1e-5 \
  --policy.optimizer_weight_decay 1e-4 \
  --policy.push_to_hub false \
  --save_checkpoint true \
  --save_freq 10000 \
  --wandb.enable true

Расширенные настройки обучения

Конфигурация с несколькими камерами

# Обучение ACT для конфигурации с несколькими камерами
lerobot-train \
  --policy.type act \
  --dataset.repo_id ${HF_USER}/your_dataset \
  --batch_size 4 \
  --steps 100000 \
  --output_dir outputs/train/act_multicam \
  --job_name act_multicam_training \
  --policy.device cuda \
  --policy.chunk_size 100 \
  --policy.n_action_steps 100 \
  --policy.n_obs_steps 2 \
  --policy.vision_backbone resnet18 \
  --policy.dim_model 512 \
  --policy.dim_feedforward 3200 \
  --policy.n_encoder_layers 4 \
  --policy.n_decoder_layers 1 \
  --policy.n_heads 8 \
  --policy.optimizer_lr 1e-5 \
  --policy.optimizer_weight_decay 1e-4 \
  --policy.push_to_hub false \
  --save_checkpoint true \
  --wandb.enable true

Настройка для оптимизации памяти

# Для GPU с небольшим объемом видеопамяти
lerobot-train \
  --policy.type act \
  --dataset.repo_id ${HF_USER}/your_dataset \
  --batch_size 2 \
  --steps 75000 \
  --output_dir outputs/train/act_memory_opt \
  --job_name act_memory_optimized \
  --policy.device cuda \
  --policy.chunk_size 100 \
  --policy.n_action_steps 100 \
  --policy.n_obs_steps 1 \
  --policy.vision_backbone resnet18 \
  --policy.dim_model 256 \
  --policy.optimizer_lr 1e-5 \
  --policy.use_amp true \
  --num_workers 2 \
  --policy.push_to_hub false \
  --save_checkpoint true \
  --wandb.enable true

Детальное описание параметров

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

Параметр	Значение	Рекомендуемое	Описание
--policy.type	Тип политики	act	Тип модели ACT
--policy.pretrained_path	Путь к предобученной модели	lerobot/act	Официальная модель ACT от LeRobot (опц.)
--dataset.repo_id	ID репозитория данных	${HF_USER}/your_dataset	Ваш набор данных на HuggingFace
--batch_size	Размер батча	8	Регулируется под видеопамять (4-8 для RTX 3070)
--steps	Количество шагов	50000	50-100 тыс. шагов для точных задач
--output_dir	Директория вывода	outputs/train/act_finetuned	Путь сохранения модели
--job_name	Имя задачи	act_finetuning	Для логов и отслеживания экспериментов

Специфические параметры ACT

Параметр	Значение	Рекомендуемое	Описание
--policy.chunk_size	Размер набора действий	100	Кол-во предсказываемых действий за раз
--policy.n_action_steps	Шаги выполнения	100	Кол-во фактически выполняемых действий
--policy.n_obs_steps	Шаги наблюдения	1	Количество исторических наблюдений
--policy.vision_backbone	Визуальный энкодер	resnet18	Сеть извлечения признаков изображений
--policy.dim_model	Размерность модели	512	Основная размерность Transformer
--policy.dim_feedforward	Размерность FF	3200	Размерность слоя FF в Transformer
--policy.n_encoder_layers	Слои энкодера	4	Кол-во слоев энкодера Transformer
--policy.n_decoder_layers	Слои декодера	1	Кол-во слоев декодера Transformer
--policy.n_heads	Кол-во голов	8	Кол-во голов в multi-head attention
--policy.use_vae	Использование VAE	true	Вариационная оптимизация

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

Параметр	Значение	Рекомендуемое	Описание
--policy.optimizer_lr	Скорость обучения	1e-5	ACT требует низких значений LR
--policy.optimizer_weight_decay	Weight decay	0.0	Параметр регуляризации
--policy.optimizer_lr_backbone	LR энкодера	1e-5	Скорость обучения визуального энкодера
--policy.use_amp	Смешанная точность	true	Экономия видеопамяти
--num_workers	Потоки загрузки	4	Зависит от количества ядер CPU
--policy.push_to_hub	Публикация в Hub	false	Загрузка на HuggingFace
--save_checkpoint	Сохранение чекпоинтов	true	Сохранение промежуточных результатов
--save_freq	Частота сохранения	10000	Интервал сохранения чекпоинтов

Мониторинг и отладка

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

# Детальная настройка W&B
lerobot-train \
  --policy.type act \
  --dataset.repo_id your-name/your-dataset \
  --batch_size 8 \
  --steps 50000 \
  --policy.push_to_hub false \
  --wandb.enable true \
  --wandb.project act_experiments \
  --wandb.notes "Обучение ACT с 4 камерами" \
  # ... другие параметры

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

На какие метрики стоит обратить внимание:

• Total Loss: общая ошибка, должна стабильно снижаться.
• Action Loss: ошибка предсказания действий (L1/L2 loss).
• Learning Rate: кривая изменения скорости обучения.
• Gradient Norm: норма градиента (мониторинг взрыва градиентов).
• GPU Memory: использование видеопамяти.
• Training Speed: количество обрабатываемых примеров в секунду.

Анализ логов обучения

# log_analysis.py
import wandb
import matplotlib.pyplot as plt

def analyze_training_logs(project_name, run_name):
    api = wandb.Api()
    run = api.run(f"{project_name}/{run_name}")
    
    # Получение истории метрик
    history = run.history()
    
    # Построение графиков
    plt.figure(figsize=(12, 4))
    
    plt.subplot(1, 3, 1)
    plt.plot(history['step'], history['train/total_loss'])
    plt.title('Total Loss')
    plt.xlabel('Step')
    plt.ylabel('Loss')
    
    plt.subplot(1, 3, 2)
    plt.plot(history['step'], history['train/action_loss'])
    plt.title('Action Loss')
    plt.xlabel('Step')
    plt.ylabel('Loss')
    
    plt.subplot(1, 3, 3)
    plt.plot(history['step'], history['train/learning_rate'])
    plt.title('Learning Rate')
    plt.xlabel('Step')
    plt.ylabel('LR')
    
    plt.tight_layout()
    plt.savefig('training_analysis.png')
    plt.show()

if __name__ == "__main__":
    analyze_training_logs("act_experiments", "act_v1_multicam")

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

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

# offline_evaluation.py
import torch
import numpy as np
from lerobot.policies.act.modeling_act import ACTPolicy
from lerobot.datasets.lerobot_dataset import LeRobotDataset

def evaluate_act_model(model_path, dataset_path):
    # Загрузка модели
    policy = ACTPolicy.from_pretrained(model_path, device="cuda")
    policy.eval()
    
    # Загрузка тестового набора данных
    dataset = LeRobotDataset(dataset_path, split="test")
    
    total_l1_loss = 0
    total_l2_loss = 0
    num_samples = 0
    
    with torch.no_grad():
        for batch in dataset:
            # Предсказание модели
            prediction = policy(batch)
            
            # Расчет ошибки
            target_actions = batch['action']
            predicted_actions = prediction['action']
            
            l1_loss = torch.mean(torch.abs(predicted_actions - target_actions))
            l2_loss = torch.mean((predicted_actions - target_actions) ** 2)
            
            total_l1_loss += l1_loss.item()
            total_l2_loss += l2_loss.item()
            num_samples += 1
    
    avg_l1_loss = total_l1_loss / num_samples
    avg_l2_loss = total_l2_loss / num_samples
    
    print(f"Средняя ошибка L1: {avg_l1_loss:.4f}")
    print(f"Средняя ошибка L2: {avg_l2_loss:.4f}")
    
    return avg_l1_loss, avg_l2_loss

if __name__ == "__main__":
    model_path = "outputs/train/act_finetuned/checkpoints/last"
    dataset_path = "path/to/your/test/dataset"
    evaluate_act_model(model_path, dataset_path)

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

# robot_evaluation.py
import torch
import numpy as np
from lerobot.policies.act.modeling_act import ACTPolicy

class ACTRobotController:
    def __init__(self, model_path, camera_names):
        self.policy = ACTPolicy.from_pretrained(model_path, device="cuda")
        self.policy.eval()
        self.camera_names = camera_names
        self.action_queue = []
        
    def get_action(self, observations):
        # Если очередь действий пуста, предсказываем новый набор (chunk)
        if len(self.action_queue) == 0:
            with torch.no_grad():
                # Подготовка входных данных
                batch = self.prepare_observation(observations)
                
                # Предсказание набора действий
                prediction = self.policy(batch)
                actions = prediction['action'].cpu().numpy()[0]  # [chunk_size, action_dim]
                
                # Добавление действий в очередь
                self.action_queue = list(actions)
        
        # Возвращаем следующее действие из очереди
        return self.action_queue.pop(0)
    
    def prepare_observation(self, observations):
        batch = {}
        
        # Обработка изображений
        for cam_name in self.camera_names:
            image_key = f"observation.images.{cam_name}"
            if image_key in observations:
                image = observations[image_key]
                # Предобработка (нормализация, ресайз и т.д.)
                image_tensor = self.preprocess_image(image)
                batch[image_key] = image_tensor.unsqueeze(0)
        
        # Обработка состояния
        if "observation.state" in observations:
            state = torch.tensor(observations["observation.state"], dtype=torch.float32)
            batch["observation.state"] = state.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 = ACTRobotController(
        model_path="outputs/train/act_finetuned/checkpoints/last",
        camera_names=["cam_high", "cam_low", "cam_left_wrist", "cam_right_wrist"]
    )
    
    # Цикл управления роботом
    for step in range(100):
        # Получение текущих наблюдений (здесь — от реального робота)
        observations = {
            "observation.images.cam_high": np.random.randint(0, 255, (480, 640, 3)),
            "observation.images.cam_low": np.random.randint(0, 255, (480, 640, 3)),
            "observation.state": np.random.randn(7)
        }
        
        # Получение действия
        action = controller.get_action(observations)
        
        # Выполнение действия (отправка роботу)
        print(f"Шаг {step}: Действие = {action}")
        
        # robot.execute_action(action)

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

Квантование модели

# quantization.py
import torch
from lerobot.policies.act.modeling_act import ACTPolicy

def quantize_act_model(model_path, output_path):
    # Загрузка модели
    policy = ACTPolicy.from_pretrained(model_path, device="cpu")
    policy.eval()
    
    # Динамическое квантование
    quantized_policy = torch.quantization.quantize_dynamic(
        policy, 
        {torch.nn.Linear}, 
        dtype=torch.qint8
    )
    
    # Сохранение квантованной модели
    torch.save(quantized_policy.state_dict(), output_path)
    print(f"Квантованная модель сохранена в {output_path}")
    
    return quantized_policy

if __name__ == "__main__":
    quantize_act_model(
        "outputs/train/act_finetuned/checkpoints/last",
        "outputs/act_quantized.pth"
    )

Оптимизация инференса

# optimized_inference.py
import torch
import torch.jit
from lerobot.policies.act.modeling_act import ACTPolicy

class OptimizedACTInference:
    def __init__(self, model_path, use_jit=True):
        self.policy = ACTPolicy.from_pretrained(model_path, device="cuda")
        self.policy.eval()
        
        if use_jit:
            # Оптимизация через TorchScript
            self.policy = torch.jit.script(self.policy)
        
        # Прогрев модели
        self.warmup()
    
    def warmup(self):
        dummy_batch = {
            "observation.images.cam_high": torch.randn(1, 3, 224, 224, device="cuda"),
            "observation.images.cam_low": torch.randn(1, 3, 224, 224, device="cuda"),
            "observation.state": torch.randn(1, 7, device="cuda")
        }
        
        with torch.no_grad():
            for _ in range(10):
                _ = self.policy(dummy_batch)
    
    @torch.no_grad()
    def predict(self, observations):
        # Быстрый инференс
        prediction = self.policy(observations)
        return prediction['action'].cpu().numpy()

if __name__ == "__main__":
    inference = OptimizedACTInference(
        "outputs/train/act_finetuned/checkpoints/last"
    )
    
    # Тест скорости инференса
    import time
    
    observations = {
        "observation.images.cam_high": torch.randn(1, 3, 224, 224, device="cuda"),
        "observation.images.cam_low": torch.randn(1, 3, 224, 224, device="cuda"),
        "observation.state": torch.randn(1, 7, device="cuda")
    }
    
    start_time = time.time()
    for _ in range(100):
        action = inference.predict(observations)
    end_time = time.time()
    
    avg_inference_time = (end_time - start_time) / 100
    print(f"Среднее время инференса: {avg_inference_time:.4f} сек")
    print(f"Частота инференса: {1/avg_inference_time:.2f} Гц")

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

Советы по сбору данных

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

Советы по оптимизации обучения

1. Размер набора действий: регулируйте chunk_size в зависимости от сложности задачи.
2. Планировщик LR: используйте косинусный отжиг или ступенчатое снижение LR.
3. Регуляризация: используйте Weight Decay и Dropout соответствующим образом.
4. Аугментация данных: применяйте подходящие методы аугментации изображений.

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

1. Сжатие модели: используйте квантование или прунинг для уменьшения размера модели.
2. Ускорение инференса: используйте TensorRT или ONNX для оптимизации работы модели.
3. Управление памятью: грамотно управляйте очередью действий и буферами наблюдений.
4. Гарантия реального времени: убедитесь, что частота инференса соответствует требованиям управления.

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

В: В чем преимущества ACT перед другими методами имитационного обучения?

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

• Снижение накопления ошибок: предсказание наборов действий уменьшает накопление погрешностей.
• Более высокий показатель успеха: отличные результаты в высокоточных манипуляциях.
• Сквозное обучение: отсутствие необходимости в ручном проектировании признаков.
• Мультимодальное слияние: эффективное объединение визуальных данных и данных о состоянии.

В: Как выбрать подходящий chunk_size?

О: Выбор chunk_size зависит от характеристик задачи:

• Быстрые задачи: chunk_size = 10-30.
• Задачи средней скорости: chunk_size = 50-100.
• Медленные задачи: chunk_size = 100-200.
• Обычно рекомендуется начинать с 50.

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

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

• Размер набора данных: около 4-8 часов для 100 эпизодов (на RTX 3070).
• Сложность модели: чем больше модель, тем дольше обучение.
• Конфигурация оборудования: мощные GPU значительно ускоряют процесс.
• Требования к сходимости: обычно требуется 50-100 тыс. шагов.

В: Как обрабатывать данные с нескольких камер?

О: Советы по работе с несколькими камерами:

• Выбор ракурсов: выбирайте ракурсы, которые дополняют друг друга.
• Слияние признаков: выполняйте слияние на уровне извлеченных признаков.
• Механизм внимания: позвольте модели научиться фокусироваться на важных ракурсах.
• Ресурсы: помните, что большее количество камер требует больше вычислительных ресурсов.

В: Как улучшить способность модели к обобщению?

О: Способы улучшения обобщения:

• Разнообразие данных: собирайте данные в различных условиях.
• Аугментация: используйте техники аугментации изображений и действий.
• Регуляризация: подходящие Weight Decay и Dropout.
• Рандомизация домена: используйте техники Domain Randomization в симуляциях.
• Мультизадачное обучение: совместное обучение на нескольких связанных задачах.

Дополнительные ресурсы

• Оригинальная статья ACT
• Реализация ACT в LeRobot
• Официальный код ACT
• Курс по робототехническому обучению
• Документация LeRobot

Журнал обновлений

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