reflex-2048

a 14.5M-parameter transformer that plays 2048 with one forward pass per move and no search. it's distilled from the macroxue/2048-ai expectimax + lookup-table teacher on one A100, for $8.72 total. it averages 164k, about a quarter of its teacher, so treat it as a small, cheap data point rather than a record.

code, training pipeline and paper: github.com/a3ro-dev/reflex-2048.

results

10,000 fresh games, greedy argmax policy, end-of-schedule checkpoint (not selected by eval score):

mean	median	max	2048	4096	8192	16384	32768	65536
164,395	131,064	835,640	96.6%	86.4%	63.3%	30.0%	3.9%	0%

a second evaluation on 2,048 games from other seeds gave 166,767 mean. per-game scores weren't logged, so there are no standard errors.

agent	params	search	mean	32768
macroxue expectimax, large tables (the teacher)	tables	depth 4	625,898	70.7%
2048former v0.1 (self-reported)	53.2M	none	605,491	66.1%
Guei et al. 2021, n-tuple	268.4M	none	412,785	30.2%
reflex-2048	14.5M	none	164,395	3.9%
reflex-2048 + 1-ply over its value head (128 games)	14.5M	1-ply	4,189	0%

a move takes 9.1 ms median on a laptop CPU (Intel, 8 threads, batch 1, features included).

calibration against teacher moves (calibration_report.json):

boards	agreement	mean confidence	ECE
held-out depth-4 teacher games	93.6%	0.931	0.006
held-out depth-3 teacher games	91.1%	0.921	0.010
boards the model reaches itself (fresh seeds, never trained on)	61.0%	0.881	0.271

known failures

DAgger collapsed play. mixing teacher relabels of the student's own states into 30% of each batch dropped greedy play from 56,652 to 4,188 mean, while held-out agreement stayed at ~87%. the labels were verified correct and the cause is not isolated. this checkpoint is trained without DAgger data.

the value head can't drive search. it's an auxiliary head trained on log return-to-go, and it tracks the return across a game (0.88 log-space correlation on held-out teacher boards), but it can't rank the four moves of one turn. 1-ply search on it scores 4,189.

it's overconfident off its teacher's distribution. on boards it reaches by itself, it reports 88% confidence and agrees with the teacher 61% of the time.

use

pip install torch numpy huggingface_hub

from huggingface_hub import snapshot_download
import sys
path = snapshot_download("a3ro-dev/reflex-2048")
sys.path.insert(0, path)
from reflex.model import load

m = load(f"{path}/reflex.pt")          # cpu by default
# board as tile exponents (0 = empty, k = tile 2^k), row-major
m.decide([[0, 1, 0, 0], [0, 0, 0, 0], [0, 0, 1, 0], [0, 0, 0, 0]])
# {'UP': ..., 'DOWN': ..., 'LEFT': ..., 'RIGHT': ...}

to watch a full game in the terminal, run python -m reflex.play --ckpt reflex.pt from the snapshot folder.

the checkpoint is a plain torch.save dict (cfg, model state dict, step, samples). loading it needs weights_only=False, so only load files you trust.

model

21 tokens: 16 tiles (tile-exponent + row + column embeddings), 4 moves built from each move's afterstate (legality, reward, empties, max tile and corner, monotonicity, smoothness, merges, follow-up rewards), and 1 value token. they go through 8 pre-norm bidirectional transformer blocks (d=384, 12 heads, MLP ratio 4). each move token gives one logit, with illegal moves masked. action order is 0=UP, 1=DOWN, 2=LEFT, 3=RIGHT.

training

the loss is cross-entropy to the teacher's move plus 0.25× a Huber loss on log return-to-go, with a random D4 symmetry per sample. the data is 360M depth-3 and 100M depth-4 teacher self-play decisions.

the model starts from a first run's step-6000 checkpoint, taken before any DAgger data reached the trainer. it then trains 212M more samples on depth-3/depth-4 data at 50/50 (261M in total) with AdamW, peak lr 5e-4 and a 1.9 h cosine schedule, in bf16 on one A100.

the teacher is GPL-3.0 code, so the repo ships a patch script for it instead of the code.

thanks

all the compute was paid for with RunPod credits from Pranav Ramakrishnan. thank you, Pranav.