Fine-tune FLUX 3 Action

Use the standalone trainer to adapt the action-pretrained base to new controls and tasks. This guide uses the game dataset adapter as a concrete example. For a task LoRA on the prepared SO-101 checkpoint, use the separate LeRobot guide.

Set up and download the base

Complete installation. Run commands from the repository root. Full fine-tuning stores model weights, gradients, optimizer state, and EMA profiles. The inference memory estimate does not describe training requirements. The launch below uses eight GPUs; the short validation run measures your actual memory use before you commit to a full job. Download the action-pretrained base and shared encoders:

uv run hf download black-forest-labs/flux-3-action-base \
  --revision 62878e2925e59b7a89ec14463ce89932624c490d \
  --include 'flux-3-action-base.safetensors' --include 'video_vae.safetensors' \
  --include 'text_encoder/*' --local-dir outputs/weights

Define the data contract

Fix action order, units, camera layout, and control rate before recording. The adapter below uses 15 Hz recordings. SO-101 uses a separate 30 Hz recipe; 32 actions cover different durations at those rates.

Field	Game adapter contract
GRUNT actions	`[move, strafe, turn, fire]`, four values in `[-1, 1]`
VECTOR actions	`[steer, throttle, nitro]`, three values in `[-1, 1]`
State	Last executed action; zero at the first frame of an episode
Camera	One RGB frame, recorded at 256 × 256 and resized to the 512 × 512 model canvas
Instruction	Fixed per game; the drone adapter reads an instruction per episode

For another robot or environment, define its measured state and action representation explicitly. Matching vector lengths does not make another checkpoint’s controls or normalization interchangeable. Record normal starts, varied starts, disturbances, and recovery behavior. Keep evaluation episodes and seeds separate from training. The published game experiments used 800 episodes of 16 seconds per game; the drone experiment used 800 episodes of 20 seconds. These are different datasets.

Episode files

The game simulator, bot, and recorder are maintained outside flux-action. Supply recordings in this format; the model repository includes the dataset adapter but not those environments or the drone recordings.

outputs/data/grunt/
  index.json
  ep_0000.npz
  ep_0001.npz

A minimal index.json for two 16-second GRUNT episodes at 15 Hz:

{
  "length": 240,
  "action_dim": 4,
  "episodes": [
    {"file": "ep_0000.npz"},
    {"file": "ep_0001.npz"}
  ]
}

This illustrates the schema; two episodes cannot fill the distributed training configuration below. Each listed NPZ contains:

Array	Shape and dtype	Meaning
`frames`	`(T, H, W, 3)`, uint8	RGB image before each action
`action`	`(T, D)`, float32	Command issued at that image

T must match length, which is shared by episodes in an index. Use separate roots for recordings of different lengths. At least 33 frames are needed per episode. For VECTOR, set action_dim to 3; the adapter pads to 4 and masks the last channel. For rotor, add "task" to every episode entry and use 4 actions. The adapter adds the fly the drone: prefix itself.

Training windows

Each window starts at frame s and stays within one episode:

Image s and state at s condition the prediction.
Actions s through s+31 are the 32 target commands.
Images s+1 through s+32 are the 32 future image targets.

For the game adapter, state at s is action s-1, or zero when s=0. Check this alignment before training. A shifted action label teaches the model to respond at the wrong time even if the loss decreases. For another data format, implement a dataset module using the game adapter as a reference. It defines the window fields, distributed sampling, and resume behavior.

Configure training

Create configs/games/local.json from the pinned game config, changing only local paths:

import json
from pathlib import Path

config = json.loads(Path("configs/games/train.json").read_text())
config["source_root"] = "grunt=outputs/data/grunt,vector=outputs/data/vector"
config["output_dir"] = "outputs/games"
config["policy"].update(
    trunk_weights="outputs/weights/flux-3-action-base.safetensors",
    video_vae_id="outputs/weights/video_vae.safetensors",
    text_encoder_id="outputs/weights/text_encoder",
)
Path("configs/games/local.json").write_text(json.dumps(config, indent=2))

index_dir: "unused" is intentional. TrainConfig requires the field, but examples.games.dataset:build reads each root’s index.json and ignores it. The default indexed DROID/LeRobot loader requires a real index directory.

Setting	What to preserve or change
`action_modality`, `action_dim`	Name and width of your controls; new embodiment heads are initialized for training
`camera_keys`, `camera_layout`, `canvas_hw`	Match the recordings; `canvas_hw` is `[height, width]`, each a multiple of 32
`fps`	Match the recorded action/frame rate; changing metadata does not resample data
`gripper_flip_dims`	`[]` for these games; DROID’s gripper inversion would corrupt the last game action
`single_frame_encode`	Keep `true` for this game config; it changes VAE encoding, not just runtime
`caption_dropout`	The game config uses `0.0`; choose and record a value explicitly for instruction-varying tasks
`n_action_steps`	Saved inference horizon: 8 in the config, 2 in the shooter playback experiment
`windows_per_rank`, `grad_accumulation`	Control global batch together with GPU count
`shard_size`	GPUs per HSDP shard group; 8 for the single-node launch below

Schedule and batch size

The game config runs for 3,000 optimizer updates. The trunk stays frozen through update 200, warms up over the next 600, and reaches its full learning rate at update 800. Cooldown starts at 2,600. See the scheduler implementation when adapting this schedule to another run length.

global batch = GPU count × windows_per_rank × grad_accumulation

The config above gives 32 windows per update on 8 GPUs. To use global batch 128 on those same GPUs, set grad_accumulation=4. Keep enough episodes for every rank and worker to supply a batch; the adapter raises an error when an epoch cannot supply a complete update.

Inspect data before training

After setting up your recordings and local config, run this from the repo root:

import json
from pathlib import Path

import torch
from examples.games.dataset import build
from flux_action.training.trainer import TrainConfig

config = TrainConfig(**json.loads(Path("configs/games/local.json").read_text()))
dataset = build(config, seed=42, epoch=0, rank=0, world_size=1,
                num_workers=1, windows_per_rank=1, skip_batches=0,
                grad_accumulation=1)
sample = next(iter(dataset))
assert sample["images.game"].shape == (33, 3, 256, 256)
assert sample["images.game"].dtype == torch.uint8
assert sample["state"].shape == (4,)
assert sample["action"].shape == (32, 4)
assert torch.isfinite(sample["action"]).all()
assert sample["action"].abs().max() <= 1
print(sample["task"], sample["start"], sample["action_mask"])

Also inspect representative windows from each environment visually, including frame zero, its matching action, and padding. The assertions above check the example’s format, not whether the demonstrations accomplish the task.

Train

First run four optimizer updates on your intended hardware. Use a separate output directory and checkpoint so this test cannot resume into the real run:

uv run torchrun --nproc_per_node 8 -m flux_action.cli train \
  --config configs/games/local.json \
  --override output_dir=outputs/games-check --override resume=none \
  --override steps=4 --override checkpoint_every=4 --override log_every=1 \
  --override frozen_steps=0 --override trunk_warmup_steps=1 \
  --override heads_warmup_steps=1 --override cooldown_start=null

Check for finite losses, nonzero trunk/head learning rates, and a completed checkpoint. metrics.jsonl records action/video MSE, learning rates, gradient norm, update time, and peak_mem_gb. Then launch the full run:

uv run torchrun --nproc_per_node 8 -m flux_action.cli train \
  --config configs/games/local.json

Reported training costs

These are distinct reported configurations, not interchangeable estimates for the command above. The original experiment checkpoints and complete logs are not included in this guide.

Reported run	Hardware	Global batch	Updates	Reported wall time
Game example results	7 H200s	126: 7 × 9 × 2	2,000	About 5 hours at 9 to 10 s/update
Distributed game recipe timing	32 H200s across 4 nodes	128: 32 × 4 × 1	3,000	About 1 h 45 at 2 s/update

The eight-GPU command above uses a different configuration; measure its runtime and memory on your hardware.

Resume and select a checkpoint

With resume: "auto", rerun the training command to resume its last complete checkpoint. Keep the weights, encoders, data, and config available. See the trainer for exact-position resume and distributed loader settings. Select a checkpoint using task performance on held-out seeds as well as offline error. The adapter does not supply an environment runner or evaluation split.

Export and run

Export one checkpoint to a new directory:

uv run flux-action export-checkpoint --checkpoint outputs/games/step-3000 \
  --output outputs/games-export --profile ema_0p10 --dtype bfloat16

Profiles are model, ema_0p10, and ema_0p05. Compare them under the same evaluation protocol. BF16 export converts the FP32 training weights; it does not include the optimizer, and the VAE/text encoder remain external references. Keep those encoder paths available after moving an export. For the game config above, load one raw uint8 RGB frame and the last executed four-channel action:

import numpy as np
import torch
from flux_action.policy import FluxActionPolicy

policy = FluxActionPolicy.from_pretrained("outputs/games-export", device="cuda")
frame = np.load("game-frame.npy")  # (256, 256, 3), uint8 RGB from your environment
last_action = torch.zeros(1, 4, device="cuda")  # only at the start of an episode
observation = {
    "images.game": torch.from_numpy(frame).permute(2, 0, 1)[None].to("cuda").float() / 255,
    "state": last_action,
    "task": ["play the shooter: hunt the grunts, dodge the plasma, stay alive"],
}
with torch.inference_mode():
    plan = policy.predict_action_chunk(observation)
print(plan.shape)  # torch.Size([1, 32, 4])

See game playback for action decoding and execution. Its reported 79 ms plan time and the earlier 76 ms export measurement are separate observations. Compare timings only with matching hardware, precision, sampling settings, compilation, and warmup.

Several environments in one checkpoint

The game adapter pads VECTOR’s three actions and previous-action state to four channels. Its action_mask is [1, 1, 1, 0], while GRUNT uses [1, 1, 1, 1]. The loss ignores padding; playback sends only the real channels to each game. The task caption identifies the environment. This example trains one shared four-channel head; another action representation needs its own data contract.

Get Started

Account Management

FLUX 3

FLUX.2

FLUX Tools

API Integration

Training

Legacy Models

Set up and download the base

Define the data contract

Episode files

Training windows

Configure training

Schedule and batch size

Inspect data before training

Train

Reported training costs

Resume and select a checkpoint

Export and run

Several environments in one checkpoint

​Set up and download the base

​Define the data contract

​Episode files

​Training windows

​Configure training

​Schedule and batch size

​Inspect data before training

​Train

​Reported training costs

​Resume and select a checkpoint

​Export and run

​Several environments in one checkpoint

Set up and download the base

Define the data contract

Episode files

Training windows

Configure training

Schedule and batch size

Inspect data before training

Train

Reported training costs

Resume and select a checkpoint

Export and run

Several environments in one checkpoint