FLUX 3 Action

FLUX 3 Action предсказывает действия робота на основе наблюдений с камер, состояния и инструкции к задаче. Замороженные Video VAE и текстовый энкодер подают данные диффузионному трансформеру, который предсказывает латенты видео и блоки действий; на этапе инференса возвращаются только действия. LeRobot поддерживает инференс и PEFT: LoRA адаптирует ствол модели, в то время как головы действий и условий робота обучаются полностью.

Мы предоставляем flux-3-action-base в качестве базовой модели действий, а также дообученные версии для роботов SO-101 flux-3-action-so101 и Franka Panda flux-3-action-droid (дообученные на датасете DROID). Центральный пример здесь — адаптация задачи с помощью LoRA на основе модели flux-3-action-so101.

Для обучения новой роботизированной конфигурации на основе базовой модели или полного дообучения используйте наш отдельный репозиторий flux-action.

Установка

Video VAE требует NATTEN версии 0.21.6 или новее, в виде wheel-пакета, собранного для ваших точных версий torch и CUDA, с сайта https://whl.natten.org. NATTEN публикует wheel-пакеты для torch 2.10–2.13, но пока не для torch 2.15, который LeRobot устанавливает по умолчанию, поэтому сначала установите поддерживаемую версию torch. Например, из репозитория LeRobot установите torch 2.13 с CUDA 13.0 вместе с зависимостями для обучения, policy, PEFT и EMA. Как и в руководстве по установке, TorchCodec устанавливается из CPU-индекса:

pip install --index-url https://download.pytorch.org/whl/cu130 \
    'torch==2.13.0+cu130' 'torchvision==0.28.0+cu130' 'numpy>=2.0,<2.3'
pip install --index-url https://download.pytorch.org/whl/cpu 'torchcodec==0.13.0+cpu'
pip install -e ".[training,flux3,peft,diffusion]" \
    'torch==2.13.0+cu130' 'torchvision==0.28.0+cu130' 'torchcodec==0.13.0+cpu'
pip install 'natten==0.21.7+torch2130cu130' --find-links https://whl.natten.org

Для драйвера NVIDIA старше 580.65 используйте сборку CUDA 12.6: индекс cu126, теги +cu126 для torch и torchvision, а также natten==0.21.7+torch2130cu126. Для другой версии torch используйте колесо NATTEN с соответствующим тегом, например natten==0.21.6+torch2110cu130 для torch 2.11 с CUDA 13.0.

Адаптация задачи SO-101

Используйте датасет LeRobot v3 с частотой 30 Гц, содержащий шесть каналов абсолютных командных действий и шесть каналов измеренного состояния в соответствующем порядке и единицах измерения для суставов. Захватчик (gripper) указывается последним. Предоставляйте необработанные значения, согласованную калиброботу и инструкции в task.

Камеры

Имена, количество и расположение камер настраиваются. В пакете SO-101 две камеры названы по тому, что они показывают:

Вы можете использовать --rename_map, чтобы сопоставить имена камер с ожидаемыми ключами, например:

--rename_map='{"observation.images.top":"observation.images.scene","observation.images.gripper_cam":"observation.images.wrist"}'

Используйте то же самое сопоставление для rollout; для совпадающих ключей карта не требуется. См. карты переименования.

Загрузка и обучение

Загрузите политику SO-101 из black-forest-labs/flux-3-action-so101. Сохранённая конфигурация автоматически загружает общие VAE и текстовый энкодер из закреплённой реvisions black-forest-labs/flux-3-action-base. Если доступ ограничен, выполните аутентификацию с помощью hf auth login, используя учётную запись с доступом. Запустите из каталога LeRobot:

lerobot-train \
  --config_path=examples/flux3/lora.json \
  --policy.path=black-forest-labs/flux-3-action-so101 --policy.device=cuda \
  --dataset.repo_id=YOUR_ORG/YOUR_SO101_DATASET \
  --output_dir=outputs/so101_lora

Кэш Hub сохраняет загруженную модель и общие энкодеры для возобновления и адаптерного вывода.

Закрепленный чекпоинт SO-101-finetuned использует ключи камер scene / wrist, восемь временных шагов наблюдений, два независимо закодированных визуальных снимка и обусловливание на основе прошлых действий/состояний. Он предсказывает 42 действия, выполняет первые 32 с частотой 30 Гц, а затем пересчитывает план. Сохранённые процессоры содержат соответствующие статистики нормализации базовой модели. Эти настройки чекпоинта сохраняются при адаптации под задачу. Общий файл examples/flux3/lora.json содержит только настройки обучения/выполнения и бюджет обучения; тот же файл поставляется как lora.json в каждом пакете Hub. Раскладка камер, обусловливание, представление действий и настройки функции потерь берутся из чекпоинта. Эти команды сохраняют сохранённую длину выполнения 32, guidance 3/3 и немедленное выполнение. Для guidance 4/1 и компиляции добавьте --policy.guidance_scale=4 --policy.guidance_scale_action=1 --policy.compile_model=true.

Эксперименты с условием: LeRobot предоставляет n_obs_steps, history_snapshots и condition_on_past_actions. Изменение этих параметров требует совместимых настроек модели/процессора и дообучения; для обычной адаптации под задачу сохраняйте значения, сохранённые в контрольной точке.

Пресет включает LoRA с rank/alpha 32 и EMA с коэффициентом затухания 0.999; автоматическая загрузка не выполняется. Новые адаптеры включаются через настройки верхнего уровня peft, а не через --policy.use_peft=true, который указывает на сохранённый адаптер.

В наших экспериментах EMA принесла существенную пользу на DROID, тогда как результаты для SO-101 пока неоднозначны. Её эффективность зависит от данных и конфигурации обучения. Мы рекомендуем оставлять EMA включённой и оценивать как исходные, так и EMA-чекпоинты по одному и тому же протоколу rollout перед выбором варианта для развёртывания. Периодическая офлайн-валидация по функции потерь использует исходные веса; она не сравнивает оба варианта автоматически.

Начните с бюджета для одного GPU, чтобы быстро итерировать. Для более крупного бюджета обучения используйте тот же конфиг с этими переопределениями:

accelerate launch --num_processes=4 --multi_gpu -m lerobot.scripts.lerobot_train \
  --config_path=examples/flux3/lora.json \
  --policy.path=black-forest-labs/flux-3-action-so101 --policy.device=cuda \
  --dataset.repo_id=YOUR_ORG/YOUR_SO101_DATASET \
  --batch_size=8 --steps=14000 \
  --output_dir=outputs/so101_lora_large

Эффективный размер батча равен числу GPU × батч на GPU × накопление. LeRobot считает steps, save_freq и eval_steps в микроступах; для получения числа обновлений оптимизатора разделите на коэффициент накопления. Пресет сохраняет и оценивает модель каждые 500 микроступов (125 обновлений). При изменении коэффициента накопления убедитесь, что все три параметра делятся на него без остатка. Длины фаз заморозки, разгона, затухания и охлаждения планировщика также задаются в микроступах. При установке длительности этих фаз умножьте желаемое число обновлений оптимизатора на коэффициент накопления. Оба бюджета используют одни и те же настройки LoRA, оптимизатора и EMA, а также сохраняют условия, размерность действий и нормализацию загруженной базовой модели. Чекпоинты содержат как исходные, так и EMA-адаптеры. Выбор бюджета и чекпоинта следует осуществлять на основе оценки на выделенной тестовой выборке и роллаутов; более крупный батч не гарантирует лучших результатов.

Пресет выделяет последние 20% эпизодов для каждой задачи на оценку. Настройте размер батча и накопление для вашего оборудования; см. обучение на нескольких GPU и обучение PEFT.

Адаптация задачи DROID

Экспериментально: examples/flux3/lora.json переиспользует настройки оптимизатора PEFT SO-101, LoRA, dropout, EMA и бюджета обучения. Загрузка конфигурации проверена, однако обучение DROID PEFT от начала до конца и качество выполнения задач пока не валидированы.

Загрузите DROID policy напрямую по идентификатору репозитория; его общие энкодеры будут разрешены автоматически:

lerobot-train \
  --config_path=examples/flux3/lora.json \
  --policy.path=black-forest-labs/flux-3-action-droid \
  --dataset.repo_id=YOUR_ORG/YOUR_DROID_TASK_DATASET \
  --output_dir=outputs/droid_lora

Здесь используется эффективный размер батча 8 и 2500 обновлений оптимизатора, с включённым EMA и сохранением контрольных точек raw/EMA каждые 125 обновлений. Автоматическая загрузка на Hub отключена. Для более крупного бюджета запустите обучение на четырёх GPU, как в примере SO-101, и добавьте --batch_size=8 --steps=14000.

Пресет загружает раскладку камер, условия, представление действий, частоту управления и процессоры из контрольной точки DROID, переопределяя только перечисленные настройки обучения/выполнения. Ваш набор данных должен соответствовать контракту контрольной точки: текущий пакет использует три камеры, восемь каналов состояния/действия и абсолютные команды позиционирования суставов с частотой 15 Гц. При необходимости сопоставьте имена камер с помощью --rename_map. Оба робота используют один и тот же пример обучения; для загрузки или инференса не требуется отдельный рецепт для конкретного робота.

Значения по умолчанию и нормализация

Новые конфигурации используют эти значения PEFT по умолчанию. Загрузка --policy.path применяет настройки, сохранённые в контрольной точке; явные переопределения в конфигурации или CLI имеют приоритет. Контрольные точки DROID сохраняют собственную конфигурацию.

Новые конфигурации используют n_obs_steps=1, history_snapshots=1 и condition_on_past_actions=false. При загрузке чекпоинта сохраняются его явные значения: закреплённый пакет SO-101 использует восемь временных шагов истории, два визуальных снимка и обусловливание по прошлым действиям/состоянию. conditioning="history" выбирает независимое кодирование снимков; само по себе оно не определяет количество временных шагов. Обучение также читает будущие кадры и команды как целевые значения, а также предшествующую команду для вычисления дельт действий.

Названия и порядок камер, размерности состояний/действий и абсолютные каналы действий остаются специфичными для конкретного робота. Чекпоинт SO-101 сохраняет последний канал абсолютным (delta_absolute_dims=[-1]) и использует веса каналов [1, 1, 1, 1, 1, 2]. Значения по умолчанию в классе не используют исключений для абсолютных каналов и применяют равные веса.

Сохранённые процессоры базовой модели содержат статистику нормализации. Передавайте сырые состояния/действия и поддерживайте одинаковую статистику для обучения, возобновления и инференса. Квантили пакета SO-101 предполагают определённый порядок суставов, единицы измерения и калибровку; они не являются универсальными значениями калибровки для SO-101. Чтобы использовать другую статистику, подготовьте отдельную базовую модель с соответствующими процессорами перед обучением. Статистика набора данных абсолютных действий не может заменить статистику командных дельт.

При обучении сохраняются адаптеры и обученные слои с ссылкой на полную базовую модель. При повторной загрузке убедитесь, что эта базовая модель и файлы её замороженного энкодера доступны; адаптеры не объединяются с ней автоматически.

Продолжить

lerobot-train \
  --config_path=outputs/so101_lora/checkpoints/last/pretrained_model/train_config.json \
  --resume=true

Возобновите обучение из исходного контрольного пункта pretrained_model, который восстанавливает состояние оптимизатора и EMA. Для инференса вы можете выбрать его «сестринский» контрольный пункт pretrained_model_ema, если включена EMA. Сохраняйте исходную базовую модель и процессоры для каждого адаптера; скачанная заново базовая модель не является заменой той, что использовалась старым адаптером.

Вывод

Укажите в параметре --policy.path исполнителя LeRobot путь к обученному контрольному пункту адаптера; он загрузит указанный базовый модельный контрольный пункт и сохранённые процессоры. Для прямых вызовов с полным базовым контрольным пунктом загрузите процессоры вместе с ним:

from lerobot.policies.factory import make_pre_post_processors
from lerobot.policies.flux3 import Flux3Policy

checkpoint = "black-forest-labs/flux-3-action-so101"
policy = Flux3Policy.from_pretrained(checkpoint).eval()
preprocessor, postprocessor = make_pre_post_processors(policy.config, pretrained_path=checkpoint)

# Run on every control tick, including while actions are queued:
prepared = preprocessor(observation)
normalized_action = policy.select_action(prepared)
absolute_command = postprocessor(normalized_action)

Предобработчик буферизует историю наблюдений и команд. Постобработчик отменяет нормализацию, интегрирует дельты и обновляет общую историю команд на каждом тике. Первая команда привязывается к измеренному состоянию; последующие команды интегрируются от предыдущей. Прямые выходные данные политики нормализуются и должны проходить постобработку. Сбрасывайте политику и оба процессора на границах эпизодов или после вмешательства; синхронный раннер LeRobot выполняет это автоматически. RTC/асинхронный вывод не поддерживается.

Для политики по умолчанию с одним временным шагом достаточно текущих изображений и измеренного состояния. Для офлайн-предсказания с якорем явного дельта-сдвига передайте observation.command_history с одной предыдущей командой; это фиксирует интеграцию вывода, не добавляя команду в условия модели. Чекпоинты, явно настроенные для нескольких временных шагов наблюдения, требуют policy.config.n_obs_steps изображений/состояний и истории команд. Пропустите входные данные через препроцессор, вызовите predict_action_chunk, затем выполните постобработку. --policy.compile_model=true включает ленивую компиляцию замороженных энкодеров и DiT для вывода; DiT для обучения не компилируется.

Для снижения VRAM при инференсе переместите Qwen на CPU после загрузки и финального вызова policy.to(...):

policy.frozen.text_encoder.cpu()

Это перемещает около 8,3 ГиБ весов в оперативную память хоста; некэшированные инструкции кодируются медленнее, при этом эмбеддинги по-прежнему передаются на устройство политики и кэшируются. Это не снижает пиковое потребление памяти при загрузке, а последующий вызов policy.to(...) перемещает Qwen обратно вместе с политикой.

Обучение других воплощений

Для вариантов, в которых мы пока не предлагаем дообученную контрольную точку, начните с базовой модели в flux-action. Этот рабочий процесс требует определения голов действий и условий робота, раскладки камер, представления и нормализации действий. Одинаковые размеры сами по себе не делают головы другого робота переиспользуемыми.

Для последующего вывода или адаптации задач в LeRobot упакуйте обученную модель в полный LeRobot policy вместе с её конфигурацией и соответствующими процессорами. Сам по себе файл с весами не является загружаемым пакетом --policy.path. В этом руководстве рассматривается адаптация существующих policy для SO-101 и DROID к новым задачам.

Все политики формируют только потоки контента video и video_cond, а также потоки действий робота. Неиспользуемые веса изображений/аудио отфильтровываются перед загрузкой, что экономит около 28% параметров DiT без изменения предсказаний действий. Сохранённые списки потоков не активируют ветви изображений/аудио; при загрузке и сохранении полной базовой модели фиксируется сокращённая структура. Адаптеры, нацеленные на удалённые потоки, требуют исходной базовой модели и более старой ревизии кода.

Политики загружаются строго по умолчанию. После отсечения известных неиспользуемых потоков отсутствующие, неожиданные или веса с неверной формой вызывают ошибку. Явное указание strict=False разрешает инициализировать отсутствующие или изменить размер несовпадающих головок embodiment; веса стержня остаются обязательными и проверяются по форме.

Валидация

Тесты на CPU охватывают упаковку, процессоры, обновления PEFT и восстановление из контрольных точек. Ранее проверки с реальными весами на H200 предшествовали текущему рефакторингу процессора. Для данной точной ревизии всё ещё требуется валидация на GPU и робототехнической платформе; существующие проверки не подтверждают работоспособность с четырьмя GPU или качество PEFT для нового воплощения.