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 データセットでファインチューニング済み)のファインチューニングモデルも同梱しています。 ここでは、flux-3-action-so101モデルをベースにした LoRA を用いたタスク適応が中心的な例です。

ベースモデルからの新しいロボットエンボディメントの学習や、フルファインチューニングには、スタンドアロンのflux-actionリポジトリを使用してください。

インストール

Video VAEには、NATTEN 0.21.6以降が必要です。これは、お使いのtorchとCUDAのバージョンに正確に適合したwheelとして、https://whl.natten.org. から入手できます。NATTENはtorch 2.10から2.13用のwheelを公開していますが、LeRobotがデフォルトでインストールするtorch 2.15用はまだ公開されていません。そのため、まずサポートされているtorchをインストールしてください。例えば、LeRobotのチェックアウトから、CUDA 13.0とtorch 2.13、およびトレーニング、ポリシー、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 ウィールを使用してください。例えば、CUDA 13.0 付きの torch 2.11 には natten==0.21.6+torch2110cu130 を使用します。

SO-101のタスク適応

30 HzのLeRobot v3データセットを使用します。6つの絶対コマンドアクションチャネルと6つの測定状態チャネルを、対応する関節の順序と単位で用意してください。グリッパーは最後に配置します。生の値、一貫性のあるロボットのキャリブレーション、および指示をtaskに提供してください。

カメラ

カメラの名前、数、レイアウトは設定可能です。SO-101パッケージでは、2つのカメラにそれぞれが映す内容に基づいた名前が付けられています。

--rename_map を使用して、カメラ名を期待されるキーにマッピングできます。例:

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

ロールアウトには同じマッピングを使用します。一致するキーにはマッピングは不要です。リネームマップを参照してください。

読み込みと学習

black-forest-labs/flux-3-action-so101 から SO-101 ポリシーを読み込みます。 保存された設定により、black-forest-labs/flux-3-action-base の固定されたリビジョンから共有 VAE とテキストエンコーダーが自動的に読み込まれます。 アクセスが制限されている場合は、アクセス権限を持つアカウントで 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 カメラキー、8 つの観測タイムステップ、独立にエンコードされる 2 つの視覚スナップショット、および過去のアクション/状態の条件付けを使用します。42 個のアクションを予測し、最初の 32 個を 30 Hz で実行した後に再計画を行います。保存されたプロセッサには、ベースモデルと一致する正規化統計情報が含まれています。これらのチェックポイント設定は、タスク適応中も保持されます。 共有の examples/flux3/lora.json は、トレーニング/ランタイム設定とトレーニング予算のみを提供します。同じファイルは、各 Hub パッケージに lora.json としてバンドルされています。カメラレイアウト、条件付け、アクション表現、および損失設定は、チェックポイントから取得されます。これらのコマンドは、保存された実行長さ 32、guidance 3/3、および eager 実行を維持します。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 を公開しています。これらを変更するには、互換性のあるモデル/プロセッサの設定とファインチューニングが必要です。通常のタスク適応では、チェックポイントに保存された値を維持してください。

プリセットは、rank/alpha 32 の LoRA と、減衰 0.999 の EMA を有効にします。何も自動的にアップロードされることはありません。新しいアダプターは、保存されたアダプターを識別する --policy.use_peft=true ではなく、トップレベルの peft 設定で有効にしてください。

実験では、EMAはDROIDにおいて大幅な利点をもたらしましたが、SO-101の結果は現時点ではまちまちです。その利点はデータや学習設定に依存します。デプロイするモデルを選択する前に、EMAを有効にしたまま、生のチェックポイントとEMAチェックポイントの両方を同じロールアウトプロトコルで評価することを推奨します。定期的なオフライン検証損失は生の重みを使用しており、両方のバリアントを自動的に比較するわけではありません。

高速なイテレーションのために、まずはシングル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

Effective batch は GPU 数 × 1 GPU あたりの batch × accumulation で計算されます。LeRobot は steps、save_freq、eval_steps をマイクロステップで数えます。オプティマイザーの更新回数を求めるには、accumulation で割ってください。プリセットでは、500 マイクロステップごと（125 回の更新ごと）に保存と評価を行います。accumulation を変更する際は、この 3 つの設定値すべてが accumulation で割り切れるようにしてください。 スケジューラーの freeze、warmup、decay、cooldown の長さもマイクロステップで指定します。これらのフェーズの長さを設定する際は、希望するオプティマイザーの更新回数を accumulation 係数で掛けてください。 両方の予算設定は、同じ LoRA、オプティマイザー、EMA 設定を使用し、読み込まれたベースモデルの conditioning、アクション次元、正規化を保持します。チェックポイントには、生アダプターと EMA アダプターの両方が含まれます。保持評価セットとロールアウトを用いて、予算とチェックポイントを選択してください。より大きな batch が必ずしもより良い結果を保証するわけではありません。

プリセットでは、各タスクのエピソードの最後の20%を評価用に保持します。ハードウェアに合わせてバッチサイズとアキュムレーションを調整してください。マルチGPUトレーニングおよびPEFTトレーニングを参照してください。

DROIDタスク適応

実験的: examples/flux3/lora.json は、SO-101 PEFT オプティマイザー、LoRA、dropout、EMA、および学習予算の設定を再利用しています。設定の読み込みは確認済みですが、エンドツーエンドの DROID PEFT 学習およびタスク性能は、まだ検証されていません。

リポジトリIDを指定して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と2,500回のオプティマイザー更新を使用し、EMAを有効化して、生データ/EMAのチェックポイントを125回ごとに保存します。Hubへの自動アップロードは無効化されています。より大きな予算の場合、SO-101の例と同様に4つのGPUで起動し、--batch_size=8 --steps=14000を追加してください。

プリセットは、DROIDチェックポイントのカメラレイアウト、条件付け、アクション表現、制御レート、プロセッサを読み込み、列挙されたトレーニング/ランタイム設定のみを上書きします。データセットはチェックポイントの契約に一致する必要があります。現在のパッケージでは、3つのカメラ、8つの状態/アクションチャネル、15 Hzの絶対関節位置コマンドを使用しています。必要に応じて、--rename_mapでカメラ名をマッピングしてください。両方のロボットは同じトレーニング例を使用しており、読み込みや推論にはロボット固有のレシピは不要です。

デフォルトと正規化

新しい設定では、これらの PEFT のデフォルト値が使用されます。--policy.path を読み込むと、チェックポイントに保存された設定が使用されます。明示的な設定または CLI による上書きが優先されます。DROID のチェックポイントは、独自の設定を保持します。

新しい設定では、n_obs_steps=1、history_snapshots=1、condition_on_past_actions=falseが使用されます。 チェックポイントを読み込むと、その明示的な値が保持されます。固定されているSO-101パッケージでは、8つの履歴タイムステップ、2つのビジュアルスナップショット、および過去のアクション/状態の条件付けを使用しています。conditioning="history"は独立したスナップショットエンコーディングを選択しますが、それ自体ではタイムステップの数を決定しません。トレーニングでは、ターゲットとして将来のフレームとコマンドを読み取り、アクションのデルタを計算するために直前のコマンドも読み取ります。

カメラの名前や順序、状態/アクションの次元数、絶対アクションチャネルはロボット固有です。 SO-101のチェックポイントは最後のチャネルを絶対値として保持し(delta_absolute_dims=[-1])、チャネル重みとして[1, 1, 1, 1, 1, 2]を使用します。クラスのデフォルトでは、絶対チャネルの例外は使用せず、重みは等しくなります。

ベースに保存されたプロセッサが正規化統計を保持しています。生の状態/アクションを入力し、学習、再開、推論のいずれにおいても同じ統計を使用してください。SO-101 パッケージの分位数は、その関節順序、単位、キャリブレーションを前提としています。これらは汎用的な SO-101 のキャリブレーション値ではありません。異なる統計を使用する場合は、学習前に、対応するプロセッサを持つ別のベースを準備してください。絶対アクションのデータセット統計は、コマンドデルタの統計の代わりにはなりません。

Training は、完全なベースへの参照を保持した状態で、アダプターと学習済みヘッドを保存します。再読み込み時には、そのベースと凍結されたエンコーダーファイルにアクセス可能にしてください。アダプターは自動的にベースにマージされません。

再開

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

pretrained_model の生チェックポイントから再開します。これにより、オプティマイザーと EMA の状態が復元されます。推論では、EMA が有効な場合、その pretrained_model_ema の兄弟を選択できます。各アダプターは、元のベースモデルとプロセッサーを保持してください。新しくダウンロードしたベースモデルは、古いアダプターで使用されていたものの代替にはなりません。

推論

LeRobotランナーの--policy.pathを学習済みのアダプターチェックポイントに設定してください。これにより、参照されているベースモデルと保存済みのプロセッサが読み込まれます。完全なベースを直接呼び出す場合は、プロセッサをベースモデルと同時に読み込みます:

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に1つ前のコマンドを指定してください。これにより、モデルの条件付けにコマンドを追加せずに出力の積分をアンカーできます。複数の観測タイムステップに明示的に設定されたチェックポイントは、policy.config.n_obs_steps枚の画像/状態とコマンド履歴を必要とします。入力をプリプロセッサに通し、predict_action_chunkを呼び出し、その後で後処理を実行します。--policy.compile_model=trueを指定すると、凍結されたエンコーダーと推論用DiTの遅延コンパイルが有効になります。ただし、学習用DiTはコンパイルされません。

推論時のVRAM使用量を低減するには、読み込み後および最終的なpolicy.to(...)呼び出しで、QwenをCPUに移動します。

policy.frozen.text_encoder.cpu()

これにより、約 8.3 GiB の重みがホスト RAM に移動します。キャッシュされていない指示はエンコードが遅くなりますが、埋め込みは引き続きポリシーデバイスに転送され、キャッシュされます。読み込み時のピークメモリは削減されず、後続の policy.to(...) によって Qwen がポリシーとともに移動します。

その他の実装のトレーニング

まだファインチューニング済みチェックポイントを提供していない実装については、flux-actionのベースモデルから開始してください。このワークフローでは、ロボットのアクションヘッドと条件付けヘッド、カメラのレイアウト、アクションの表現方法と正規化を定義する必要があります。次元が等しいだけでは、別のロボットのヘッドを再利用することはできません。

その後のLeRobot推論やタスク適応のために、学習済みモデルを、設定および対応するプロセッサを含む完全なLeRobotポリシーとしてパッケージ化してください。生の重みファイルだけでは、読み込み可能な--policy.pathパッケージにはなりません。本ガイドでは、既存のSO-101およびDROIDポリシーを新しいタスクに適応させることに焦点を当てています。

すべてのポリシーは、videoおよびvideo_condのコンテンツストリームと、ロボットのアクションストリームのみを構築します。 使用されない画像/音声の重みは、読み込み前にフィルタリングされ、アクション予測を変更せずにDiTパラメータの約28%を節約します。 保存されたストリームリストは、画像/音声ブランチを有効にしません。完全なベースを読み込み、保存すると、削減されたレイアウトが記録されます。 削除されたストリームを対象とするアダプターには、元のベースと古いコードリビジョンが必要です。

完全なポリシーはデフォルトで厳格に読み込まれます。既知の未使用ストリームをフィルタリングした後、欠落、予期しない、または形状が誤っている重みはエラーを発生させます。明示的にstrict=Falseを指定すると、欠落しているエンボディメントヘッドの初期化や、形状が不一致のエンボディメントヘッドのリサイズのみが許可されます。トランクの重みは引き続き必須であり、形状チェックが行われます。

検証

CPU テストは、パッキング、プロセッサ、PEFT の更新、およびチェックポイント / 再開をカバーしています。現在のプロセッサのリファクタリングに先立ち、実際の重みを用いた H200 のチェックが実施されました。この正確なリビジョンは、まだ GPU およびロボットによる検証が必要です。既存のチェックは、4 GPU 環境での容量や、新しいエンボディメントにおける PEFT の品質を保証するものではありません。