> ## Documentation Index
> Fetch the complete documentation index at: https://docs.bfl.ml/llms.txt
> Use this file to discover all available pages before exploring further.

# Fine-tune FLUX 3 Action

> Prepare demonstrations, train a custom action policy, and export it for evaluation.

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](/flux_3/flux3_action_so101).

## Set up and download the base

Complete [installation](/flux_3/flux3_action_inference#requirements-and-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:

```sh theme={null}
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.

```text theme={null}
outputs/data/grunt/
  index.json
  ep_0000.npz
  ep_0001.npz
```

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

```json theme={null}
{
  "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](https://github.com/black-forest-labs/flux-action/blob/main/examples/games/dataset.py) 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](https://github.com/black-forest-labs/flux-action/blob/main/configs/games/train.json), changing only local paths:

```python theme={null}
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](https://github.com/black-forest-labs/flux-action/blob/main/src/flux_action/training/schedule.py)
when adapting this schedule to another run length.

```text theme={null}
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:

```python theme={null}
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:

```sh theme={null}
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:

```sh theme={null}
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](https://github.com/black-forest-labs/flux-action/blob/main/src/flux_action/training/trainer.py)
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:

```sh theme={null}
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:

```python theme={null}
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](/flux_3/flux3_action_games#play-it) 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.