Trained Loop Depth

Phase 3 · full fine-tune of the backbone

model

FLUX 3 Action 7B transformer, video-only configuration: the video, video-condition and text streams (5.57 B parameters) are trained in full, no LoRA. The robot action and state streams are dropped. The video VAE and the Qwen3-VL text encoder stay frozen.

data

AgiBot World Beta head-camera clips, 15 fps, resized to 384×512. Train: 7960 episodes over 212 tasks, 159,983 windows; a window is k past frames plus a caption (the segment's action text) and the 33 frames that follow.

input

The past frames are VAE-encoded and enter the video-condition stream as clean tokens at timestep 0; the future frames are encoded, noised and enter the video stream (9 latent frames × 12 × 16 patches = 1728 tokens). Context length k cycles through 1, 5 and 9 per micro-batch, with k = 0 (caption only) 10% of the time. Captions are dropped to the empty string 10% of the time so guidance works at inference.

objective

Rectified-flow velocity MSE on the future tokens, timesteps drawn as sigmoid(N(1.08, 1.0)) as in the base model's recipe.

optimizer

AdamW, learning rate 1e-5, betas (0.9, 0.999), no weight decay, gradient clip 1.0, fp32 master weights with bf16 compute, EMA 0.999 updated every 10 steps (the EMA is what is evaluated). 32 windows per step: 2 GPUs × 4 windows × 4 accumulation steps.

schedule

6000 steps (1.2 epochs): 500-step warmup, constant, linear cooldown to zero over steps 5000–6000.

infrastructure

2 × A100 80 GB, FSDP2 sharding with activation checkpointing and compiled blocks, 16 s per step. Steps 0–3000 ran with a reduce-order bug that mixed gradients across ranks; it was fixed and steps 3000–6000 are clean.

Phase 5 · Looped-DiT retrofit

The 33 transformer blocks of the fine-tuned model (5 dual-stream mode blocks, 28 joint single blocks) are split into three stages. The retrofit adds no blocks: the core is 9 of the existing single blocks, run N times with shared weights.

initialisation

All trunk weights are the Phase-3 EMA at step 6000. The only new parameters are head-wise sigmoid attention gates in the 9 core blocks (one linear map for text tokens, one for video tokens, zero weights, bias +4), so at initialisation each gate passes 98% of the attention output and the model computes almost exactly what the fine-tune did at N = 1.

loss

Deep supervision: after each pass through the core the state is decoded through C and scored with the same velocity loss. Weights 1/3, 1/3, 1/3, 1 over exits 1–4, rescaled to sum 2, so the final exit counts most but every depth is trained.

schedule

Same data stream, seed, batch size, learning rate and optimizer as Phase 3; 200-step warmup. The 2500-step run cools down over steps 2000–2500 (half an epoch). The 4500-step run restarts from the step-2000 checkpoint and keeps the learning rate constant to step 4500 (0.9 epoch), where it was stopped without a cooldown; its EMA (decay 0.999) is the evaluated model.

cost

12 + 4 × 9 + 4 × 12 = 96 block passes per training forward against 33 for the plain model; 31 s per step on the same two GPUs, so the 2500-step run roughly matches Phase 3's compute.

1. The k past frames are encoded by the video VAE (4× in time, 32× in space, 96 channels) into clean conditioning tokens at timestep 0. With k = 0 there is no conditioning video.
2. The caption is encoded by the frozen Qwen3-VL text encoder (80 text tokens); the empty caption is encoded too, as the unconditional branch.
3. Gaussian noise is drawn for the 9 future latent frames from a fixed seed (seed 0 for every evaluation window), so the columns of a strip share their noise.
4. 30 Euler steps run on the time-shifted schedule (shift 5, t from 1 to 0). At each step the transformer predicts the velocity twice, with and without the caption, and the two are combined with classifier-free guidance scale 4.
5. The loop depth is a call-time argument: the same weights are sampled with the core run 1, 2, 4 or 8 times; the pre- and post-loop stages and the head are unchanged.
6. The final latents are decoded by the VAE to 33 frames at 384×512 and compared with the real continuation: PSNR, SSIM and LPIPS per window, FVD over 300 windows, and a VLM judge that asks whether the clip matches its caption, scored against a caption from another task.

Cost per depth

Seconds per 33-frame window at k = 1, including the guidance pass (two forwards per step) and the VAE decode; batch 4, one A100 for the 2500 model and one H100 for the 4500 model, so the two columns are not comparable in absolute terms. The plain Phase-3 model takes 7.9 s under the same protocol.