docs/METHOD.md

# FACT architecture

## Token layout

Each autoregressive block covers 0.8 s and packs **83 tokens**:
`[obs 50][state 1][action 16][future-state 16]`, preceded by one learnable task token
(and, with `use_text`, up to 64 umT5-XXL text tokens as a clean prefix, padding masked).
Training runs K = 4 blocks (3.2 s) in one forward pass with a block-causal mask: block *k*
attends to blocks `[k-window, k]` (window = 4) and the prefix; online K = 1.

| token | source | noise |
|---|---|---|
| obs (50) | current 3-view frame tiled into a 2×2 grid (head \| right ; left \| black) at 160×320, Wan2.2 VAE → 48×10×20 latent, 2×2 patches, `Linear(192→d)` | clean |
| state (1) | 62-D EEF62 state (zero-padded to 64), MLP → d | clean, τ = 0 |
| action (16) | 16 × 62 delta-base actions (T-Rex reference frame and q01/q99 normalisation), `Linear(64→d)` | noisy, τ ~ Beta(1.5, 1) per block |
| future state (16) | EEF62 states at +1..+16 steps (20 Hz) after the block anchor, delta-base w.r.t. the block reference state, normalised with the action statistics | noisy, same τ as action |

Positional information: slot embedding (83 slots) + block embedding + token-type embedding
(no RoPE).  Each token carries its own τ → sinusoidal(256) → MLP → per-token adaLN-zero
(6 modulation vectors, so different modalities can sit at different noise levels);
qk-RMSNorm; final adaLN followed by two linear heads (action flow, future-state flow).

## Two-stage flow

* **Coarse stage** (`FactCoarseTransformer`): joint rectified flow on `[action, future state]`
  from τ = 1 to τ = 0.4 (6 Euler steps at inference).
* **Coarse memory**: last-layer hidden states visible to block *k* (84 tokens online).
* **Force stage** (`ForceOnlyTransformer`, 768-d, 6 layers, re-used unchanged from T-Rex):
  16 action tokens + 10 current F6 force tokens + 10 tactile-VQ history tokens + memory;
  τ 0.4 → 0 in 4 steps, suffix-only refinement at offsets 0/4/8/12 (5 Hz tactile refresh),
  tactile dropout 0.1.  The tactile VQ (64 codes × 256-d) is trained jointly with the
  standard VQ + 0.25 · commitment loss.

Loss: `action_flow_mse + state_flow_mse + force_flow_mse + vq + 0.25·commitment` (all weights 1).

Sampling: 6 coarse steps (τ 1 → 0.4) → forward at τ = 0.4 to obtain memory and clean future
states → 4 force steps.  `coarse_to_zero` lets the coarse model run to τ = 0 instead (used
when the force stage is disabled or for analysis).

## Configurations

| config | d / layers / heads / FFN | params | notes |
|---|---|---|---|
| `fact_base.json` | 1024 / 12 / 16 / 4096 | 209 M | finalised model |
| `fact_base_text.json` | same + `use_text` (4096-d umT5 tokens, `LayerNorm + Linear(4096→d)`) | 213 M | multi-task training |
| `fact_small.json` | 768 / 8 / 12 / 3072 | 109 M | resource-constrained variant |

All configs keep the force stage (768 / 6 / 12 / 3072), window 4, 6 coarse + 4 force steps,
τ_split 0.4.

## Full-data training recipe

Global batch 16 × 8 = 128 samples (4 AR blocks each), AdamW, lr 3e-4 cosine to 10 % with
2000 warm-up steps, weight decay 0.01, grad clip 1.0, bf16 autocast (force stage in fp32),
force-stage lr × 0.3, optimizer step skipped when the pre-clip grad norm exceeds 20.
100 k steps ≈ 15 h on 32 accelerators at 0.55–1.1 s/it.

Every `--video-steps` the rank-0 process samples 4 consecutive chunks of one episode and logs
`eval/openloop_cascade_mse`, `eval/openloop_coarse_only_mse`, `eval/future_state_mse` and a
monitoring video (left: head camera at 20 Hz; right: predicted vs ground-truth action and
future-state curves — black GT / red cascade / blue coarse τ = 0.4 / green future state).