FLUX 3 Action

FLUX 3 Action prédit les actions du robot à partir des observations de la caméra, de l'état et d'une instruction de tâche. Un Video VAE et un encodeur de texte figés alimentent un transformateur de diffusion qui prédit les latents vidéo et les blocs d'actions ; seules les actions sont renvoyées lors de l'inférence. LeRobot prend en charge l'inférence et le PEFT : LoRA adapte le tronc, tandis que les têtes d'action et de conditionnement du robot s'entraînent en entier.

Nous fournissons flux-3-action-base en tant que modèle d’action de base, accompagné de finetunings pour les robots SO-101 flux-3-action-so101 et Franka Panda flux-3-action-droid (finetuné sur le jeu de données DROID). L’exemple central ici est une adaptation de tâche via un LoRA basé sur le modèle flux-3-action-so101.

Pour entraîner une nouvelle incarnation robotique à partir du modèle de base ou pour effectuer un affinement complet, utilisez notre dépôt autonome flux-action.

Installation

Depuis le dépôt LeRobot, installez les dépendances d’entraînement, de politique, PEFT et EMA :

pip install -e ".[training,flux3,peft,diffusion]"

Le VAE vidéo nécessite également une roue NATTEN compatible torch/CUDA, disponible sur https://whl.natten.org. Pour torch 2.11 avec CUDA 12.8 :

uv pip install natten==0.21.6+torch2110cu128 --find-links https://whl.natten.org

Utilisez natten==0.21.6+torch2100cu128 pour torch 2.10, ou compilez-le pour votre environnement.

Adaptation de la tâche SO-101

Utilisez un jeu de données LeRobot v3 à 30 Hz, avec six canaux d’actions commandées absolues et six canaux d’état mesuré, dans un ordre et des unités de joints correspondants. La pince est en dernier. Fournissez les valeurs brutes, une calibration robot cohérente et des instructions dans task.

Caméras

Les noms, le nombre et la disposition des caméras sont configurables. Le package SO-101 nomme ses deux caméras en fonction de ce qu’elles montrent :

Vous pouvez utiliser --rename_map pour associer les noms de caméras aux clés attendues, par exemple :

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

Utilisez la même correspondance pour le déroulement ; les clés correspondantes n’ont pas besoin de correspondance. Voir correspondances de renommage.

Charger et entraîner

Chargez la politique SO-101 depuis black-forest-labs/flux-3-action-so101. La configuration enregistrée charge automatiquement le VAE et l'encodeur de texte partagés depuis une révision épinglée de black-forest-labs/flux-3-action-base. Si l'accès est restreint, authentifiez-vous avec hf auth login en utilisant un compte disposant de l'accès. Exécutez depuis le checkout de 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

Le cache du Hub conserve le modèle téléchargé et les encodeurs partagés pour la reprise et l'inférence d'adaptateurs.

Le point de contrôle SO-101-finetuned épinglé utilise les clés caméra scene / wrist, huit pas de temps d’observation, deux instantanés visuels encodés indépendamment et un conditionnement sur les actions/états passés. Il prédit 42 actions, exécute les 32 premières à 30 Hz, puis relance la planification. Les processeurs enregistrés portent les statistiques de normalisation correspondantes de la base. Ces paramètres du point de contrôle sont conservés lors de l’adaptation à la tâche. Le fichier partagé examples/flux3/lora.json fournit uniquement les paramètres d’entraînement/temps d’exécution et le budget d’entraînement ; le même fichier est inclus sous le nom lora.json dans chaque paquet du Hub. La disposition des caméras, le conditionnement, la représentation des actions et les paramètres de perte proviennent du point de contrôle. Ces commandes conservent la longueur d’exécution enregistrée de 32, le guidage 3/3 et l’exécution en mode eager. Pour un guidage 4/1 et la compilation, ajoutez --policy.guidance_scale=4 --policy.guidance_scale_action=1 --policy.compile_model=true.

Expériences de conditionnement : LeRobot expose n_obs_steps, history_snapshots et condition_on_past_actions. Les modifier nécessite des paramètres de modèle/processeur compatibles et un fine-tuning ; conservez les valeurs enregistrées dans le checkpoint pour l'adaptation ordinaire à une tâche.

Le préréglage active LoRA avec un rang/alpha de 32 et EMA avec un facteur d'atténuation de 0,999 ; rien n'est téléversé automatiquement. Activez les nouveaux adaptateurs via les paramètres peft de niveau supérieur, et non via --policy.use_peft=true, qui identifie un adaptateur enregistré.

Dans nos expériences, EMA a apporté des avantages substantiels sur DROID, tandis que les résultats sur SO-101 ont été mitigés jusqu’à présent. Son bénéfice dépend des données et de la configuration d’entraînement. Nous recommandons de conserver EMA activé et d’évaluer à la fois les points de contrôle bruts et ceux avec EMA selon le même protocole de déploiement avant de choisir lequel déployer. La perte de validation hors ligne périodique utilise les poids bruts ; elle ne compare pas automatiquement les deux variantes.

Commencez par le budget mono-GPU pour une itération rapide. Pour un budget d'entraînement plus important, utilisez la même configuration avec ces substitutions :

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

Le lot effectif est le nombre de GPU × le lot par GPU × l’accumulation. LeRobot compte les steps, save_freq et eval_steps en micro-étapes ; divisez par l’accumulation pour obtenir les mises à jour de l’optimiseur. Le préréglage enregistre et évalue toutes les 500 micro-étapes (125 mises à jour). Gardez les trois paramètres divisibles par l’accumulation lors de sa modification. Les longueurs de gel, d’échauffement, de décroissance et de refroidissement du planificateur utilisent également les micro-étapes. Multipliez le nombre de mises à jour de l’optimiseur souhaité par le facteur d’accumulation lors de la définition de ces longueurs de phase. Les deux budgets utilisent les mêmes paramètres LoRA, optimiseur et EMA, et conservent le conditionnement, les dimensions d’action et la normalisation de la base chargée. Les points de contrôle incluent les adaptateurs bruts et EMA. Choisissez le budget et le point de contrôle à l’aide d’une évaluation sur données retenues et de déroulements ; un lot plus grand ne garantit pas de meilleurs résultats.

Le préréglé réserve les 20 % derniers épisodes de chaque tâche pour l'évaluation. Ajustez la taille du lot et l'accumulation en fonction de votre matériel ; consultez l'entraînement multi-GPU et l'entraînement PEFT.

Adaptation de la tâche DROID

Expérimental : examples/flux3/lora.json réutilise les paramètres de l'optimiseur PEFT SO-101, de LoRA, de dropout, d'EMA et du budget d'entraînement. Le chargement de la configuration a été vérifié, mais l'entraînement PEFT DROID de bout en bout et les performances sur les tâches n'ont pas encore été validés.

Chargez la politique DROID directement par identifiant de dépôt ; ses encodeurs partagés sont résolus automatiquement :

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

Cela utilise un lot effectif de 8 et 2 500 mises à jour de l’optimiseur, avec l’EMA activé et des points de contrôle bruts/EMA enregistrés toutes les 125 mises à jour. Les téléversements automatiques sur le Hub sont désactivés. Pour le budget plus élevé, lancez avec quatre GPU comme dans l’exemple SO-101 et ajoutez --batch_size=8 --steps=14000.

Le préréglage charge la disposition des caméras, le conditionnement, la représentation des actions, le taux de contrôle et les processeurs du point de contrôle DROID, en ne remplaçant que les paramètres d’entraînement/temps d’exécution listés. Votre jeu de données doit correspondre au contrat du point de contrôle : le package actuel utilise trois caméras, huit canaux d’état/action et des commandes de position articulaire absolue à 15 Hz. Mappez les noms des caméras avec --rename_map si nécessaire. Les deux robots utilisent le même exemple d’entraînement ; aucune recette spécifique au robot n’est nécessaire pour le chargement ou l’inférence.

Valeurs par défaut et normalisation

Les nouvelles configurations utilisent ces valeurs par défaut de PEFT. Le chargement de --policy.path utilise les paramètres enregistrés dans le point de contrôle ; les configurations ou les substitutions CLI explicites ont la priorité. Les points de contrôle DROID conservent leur propre configuration.

Les nouvelles configurations utilisent n_obs_steps=1, history_snapshots=1 et condition_on_past_actions=false. Le chargement d'un point de contrôle conserve ses valeurs explicites : le paquet SO-101 épinglé utilise huit pas de temps d'historique, deux instantanés visuels et un conditionnement sur les actions/états passés. conditioning="history" sélectionne l'encodage indépendant des instantanés ; il ne détermine pas à lui seul le nombre de pas de temps. L'entraînement lit également les images et commandes futures comme cibles, ainsi que la commande précédente pour calculer les deltas d'action.

Les noms et l’ordre des caméras, les dimensions des états/actions et les canaux d’action absolus restent spécifiques à chaque robot. Le point de contrôle SO-101 conserve le dernier canal en absolu (delta_absolute_dims=[-1]) et utilise les poids de canaux [1, 1, 1, 1, 1, 2]. Les valeurs par défaut de la classe n’utilisent aucune exception pour les canaux absolus et des poids égaux.

Les processeurs enregistrés de la base possèdent les statistiques de normalisation. Fournissez des états/actions bruts et conservez les mêmes statistiques pour l’entraînement, la reprise et l’inférence. Les quantiles du package SO-101 supposent son ordre des articulations, ses unités et sa calibration ; ils ne sont pas des valeurs de calibration universelles pour SO-101. Pour utiliser des statistiques différentes, préparez une base séparée avec des processeurs correspondants avant l’entraînement. Les statistiques des ensembles de données d’actions absolues ne peuvent pas se substituer aux statistiques des deltas de commandes.

L'entraînement enregistre les adaptateurs et les têtes entraînées avec une référence vers la base complète. Conservez cette base et ses fichiers d'encodeur gelés accessibles lors du rechargement ; les adaptateurs ne sont pas fusionnés automatiquement avec elle.

Reprendre

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

Reprendre depuis le point de contrôle brut pretrained_model, ce qui restaure l’état de l’optimiseur et de l’EMA. Pour l’inférence, vous pouvez sélectionner son homologue pretrained_model_ema lorsque l’EMA est activé. Conservez la base et les processeurs d’origine de chaque adaptateur ; une base téléchargée à neuf ne remplace pas celle utilisée par un adaptateur plus ancien.

Inférence

Indiquez au runner LeRobot le checkpoint de l’adaptateur entraîné via --policy.path ; il charge la base référencée ainsi que les processeurs enregistrés. Pour les appels directs avec une base complète, chargez les processeurs en même temps :

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)

Le préprocesseur met en mémoire tampon l’historique des observations et des commandes. Le postprocesseur annule la normalisation, intègre les deltas et met à jour l’historique partagé des commandes à chaque tick. La première commande est ancrée à l’état mesuré ; les commandes suivantes s’intègrent à partir de la commande précédente. Les sorties directes de la politique sont normalisées et doivent être post-traitées. Réinitialisez la politique et les deux processeurs aux limites d’épisode ou après une intervention ; le runner synchrone de LeRobot le fait. L’inférence RTC/asynchrone n’est pas prise en charge.

Pour la politique par défaut à pas de temps unique, les images actuelles et l’état mesuré suffisent. Pour une prédiction hors ligne avec une ancre delta explicite, fournissez observation.command_history avec une commande précédente ; cela ancre l’intégration de la sortie sans ajouter de commande au conditionnement du modèle. Les points de contrôle explicitement configurés pour plusieurs pas de temps d’observation nécessitent policy.config.n_obs_steps images/états et l’historique des commandes. Passez les entrées par le préprocesseur, appelez predict_action_chunk, puis post-traitez. --policy.compile_model=true active la compilation paresseuse des encodeurs figés et du DiT d’inférence ; il ne compile pas le DiT d’entraînement.

Pour réduire la VRAM d’inférence, déplacez Qwen sur le CPU après le chargement et l’appel final policy.to(...) :

policy.frozen.text_encoder.cpu()

Cela déplace environ 8,3 Go de poids vers la RAM hôte ; les instructions non mises en cache sont encodées plus lentement, tandis que les embeddings sont toujours transférés vers l’appareil de la politique et mis en cache. Cela ne réduit pas la mémoire de pointe lors du chargement, et un policy.to(...) ultérieur déplace Qwen à nouveau avec la politique.

Entraîner d'autres incarnations

Pour les cas d’usage pour lesquels nous ne proposons pas encore de point de contrôle affiné, partez du modèle de base de flux-action. Ce flux de travail nécessite de définir les têtes d’action et de conditionnement du robot, la disposition des caméras, ainsi que la représentation et la normalisation des actions. Le simple fait que les dimensions soient identiques ne suffit pas à rendre les têtes d’un autre robot réutilisables.

Pour les inférences ou l’adaptation de tâches ultérieures avec LeRobot, empaquetez le modèle entraîné en tant que politique LeRobot complète, avec sa configuration et ses processeurs correspondants. Un simple fichier de poids bruts ne constitue pas un paquet --policy.path chargeable. Ce guide se concentre sur l’adaptation des politiques SO-101 et DROID existantes à de nouvelles tâches.

Toutes les politiques ne construisent que les flux de contenu video et video_cond, ainsi que les flux d’actions du robot. Les poids d’images/audio inutilisés sont filtrés avant le chargement, ce qui permet d’économiser environ 28 % des paramètres de DiT sans modifier les prédictions d’actions. Les listes de flux enregistrées n’activent pas les branches image/audio ; le chargement et l’enregistrement d’une base complète enregistrent la disposition réduite. Les adaptateurs ciblant des flux supprimés nécessitent leur base d’origine et une révision de code antérieure.

Les politiques complètes sont chargées de manière stricte par défaut. Après le filtrage des flux connus non utilisés, les poids manquants, inattendus ou de forme incorrecte déclenchent une erreur. L’argument explicite strict=False autorise uniquement l’initialisation des têtes d’incarnation manquantes ou le redimensionnement des têtes d’incarnation dont la forme ne correspond pas ; les poids du tronc restent obligatoires et leur forme est vérifiée.

Validation

Les tests CPU couvrent le packing, les processeurs, les mises à jour PEFT et la reprise de checkpoint. Des vérifications H200 avec des poids réels ont précédé la refonte actuelle du processeur. Cette révision exacte nécessite encore une validation GPU et robotique ; les vérifications existantes n'établissent pas la capacité à quatre GPU ni la qualité PEFT sur une nouvelle incarnation.