initial commit

.gitattributes

@@ -0,0 +1,36 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/overview.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,113 @@
+---
+license: cc0-1.0
+tags:
+  - world-model
+  - arc-agi-3
+  - test-time-training
+  - next-frame-prediction
+  - vision-transformer
+  - pytorch
+---
+
+# WARC
+
+[![Code: GitHub](https://img.shields.io/badge/Code-GitHub-black?logo=github)](https://github.com/third-intelligence/WARC)
+[![License: CC0 1.0](https://img.shields.io/badge/License-CC0%201.0-lightgrey)](https://creativecommons.org/publicdomain/zero/1.0/)
+
+WARC (World model for ARC-AGI-3) is a 35M-parameter action-conditioned vision transformer. It
+predicts the next frame of a game, cell by cell, from the last three frames and
+the action, and it keeps learning while an agent plays: after every observed
+transition it updates all of its parameters by gradient descent. Pretrained on
+human recordings of other games, it learns to predict four in five mechanics
+of an unseen game with at least 80% accuracy on the cells they change and keeps
+that accuracy to the end of the recordings.
+
+## Model details
+
+| | |
+|---|---|
+| Input | The last three 64x64 frames of 16 colours, the action after the last frame, and the two actions between the frames |
+| Output | A distribution over the 16 colours of every cell of the next frame |
+| Architecture | 2x2 patches, 10 pre-norm transformer blocks of width 512 with rotary position embeddings, modulated by the action through adaptive layer normalization; a learned vector added to the token of the clicked cell |
+| Parameters | 35M |
+| Pretraining | Next-frame cross-entropy on human recordings of the public games; Muon and AdamW, 30 epochs |
+| Test-time training | 20 gradient steps on the ten newest and twenty earlier transitions after every observed transition; all parameters are updated |
+| License | CC0 1.0 |
+
+## Files
+
+| File | Pretrained on | Held-out games | Validation exact-frame accuracy |
+|---|---|---|---:|
+| `fold1/checkpoint_best.pt` | 16 games | bp35 lf52 lp85 r11l re86 sb26 sc25 su15 wa30 | 85.7% |
+| `fold2/checkpoint_best.pt` | 17 games | cd82 cn04 dc22 m0r0 sk48 tn36 tr87 tu93 | 79.2% |
+| `fold3/checkpoint_best.pt` | 17 games | ar25 ft09 g50t ka59 ls20 s5i5 sp80 vc33 | 78.5% |
+| `all25/checkpoint_best.pt` | all 25 games | none (one recording per game held out for validation) | 80.0% |
+
+Each checkpoint (146 MB) holds the weights, the model configuration, and the
+game split; `config.json` and `eval_history.jsonl` next to it record the
+training configuration and the per-epoch validation accuracy. `manifest.json`
+lists the SHA-256 of every checkpoint and the fold of every game.
+
+## Usage
+
+Install the code from the [GitHub repository](https://github.com/third-intelligence/WARC)
+and download the weights:
+
+```bash
+git clone https://github.com/third-intelligence/WARC
+cd WARC
+python3.12 -m venv .venv && source .venv/bin/activate
+pip install torch==2.7.0 torchvision==0.22.0 --index-url https://download.pytorch.org/whl/cu126
+pip install -r requirements.txt
+hf download third-intelligence/WARC --local-dir models    # the four checkpoints (pip install huggingface_hub)
+```
+
+Predict the next frame:
+
+```python
+import sys
+import torch
+
+sys.path[:0] = ["agent", "agent/VARC"]
+from warc_model.model import build_from_checkpoint
+
+device = torch.device("cuda")
+model, meta = build_from_checkpoint("models/all25/checkpoint_best.pt", device=device)
+
+# The last three frames, (1, 3, 64, 64) colours 0-15. Actions: 1-5 are the keys ACTION1-ACTION5,
+# 6 is a click at column x and row y; x and y are -1 for a key.
+frames = torch.randint(0, 16, (1, 3, 64, 64), device=device)     # replace with frames of a game
+action, x, y = (torch.tensor(v, device=device) for v in ([6], [10], [20]))
+hist_action, hist_x, hist_y = (torch.tensor(v, device=device) for v in ([[1, 2]], [[-1, -1]], [[-1, -1]]))
+
+with torch.no_grad():
+    logits = model(frames[:, -1], torch.zeros(1, dtype=torch.long, device=device), action, x, y,
+                   hist_obs=frames, hist_action_ids=hist_action, hist_action_xs=hist_x,
+                   hist_action_ys=hist_y)["ffp_logits"]          # (1, 16, 64, 64)
+next_frame = logits.argmax(1)[0]                                 # (64, 64) predicted colours
+```
+
+Test-time training inside a game loop (`FinalPredictor` in the repository), the
+agent, the Kaggle notebook, and the training commands are documented in the
+[README](https://github.com/third-intelligence/WARC#quick-start) and in
+[docs/reproducibility.md](https://github.com/third-intelligence/WARC/blob/main/docs/reproducibility.md).
+
+## Training data
+
+The 340 human gameplay recordings of the 25 public ARC-AGI-3 games published by
+the ARC Prize Foundation ([ARC-AGI-3 human dataset](https://arcprize.org/blog/arc-agi-3-human-dataset)).
+The fold models use the recordings of their training games only; the all-25
+model uses one recording per game for validation and the rest for training.
+
+## Acknowledgements
+
+WARC follows VARC ("ARC Is a Vision Problem!", Hu et al., 2025) and
+reuses parts of the [VARC code](https://github.com/lillian039/VARC) (MIT). The
+games and the recordings belong to the [ARC Prize Foundation](https://arcprize.org).
+The repository lists every third-party component and its licence.
+
+## License
+
+The weights and the code are released under
+[CC0 1.0](https://creativecommons.org/publicdomain/zero/1.0/). Third-party
+components keep their own licences.

all25/checkpoint_best.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:348faf916bb842fe27ccef24efa86f79f13f3175c78d1a1e6408fdddf0502922
+size 146368647

all25/config.json

@@ -0,0 +1,63 @@
+{
+ "fold": 0,
+ "split_seed": 0,
+ "patch_size": 2,
+ "embed_dim": null,
+ "depth": null,
+ "num_heads": null,
+ "mlp_dim": null,
+ "hist_n": 3,
+ "epochs": 30,
+ "batch_size": 86,
+ "global_batch_size": 516,
+ "learning_rate": 0.0003,
+ "weight_decay": 0.0,
+ "max_grad_norm": 1.0,
+ "lr_scheduler": "cosine",
+ "seed": 42,
+ "num_workers": 8,
+ "optimizer": "muon",
+ "muon_lr": 0.0,
+ "muon_momentum": 0.95,
+ "muon_adjust_lr": "match_rms_adamw",
+ "gamma": 0.95,
+ "no_amp": false,
+ "amp_dtype": "bf16",
+ "no_compile": false,
+ "compile_mode": "reduce-overhead",
+ "distributed": true,
+ "disable_augment": false,
+ "num_color_permute": 9,
+ "eval_batch_size": null,
+ "eval_every": 1,
+ "max_train_batches": 0,
+ "arch": "unified_hist_cm",
+ "model_config": {
+  "embed_dim": 512,
+  "depth": 10,
+  "num_heads": 8,
+  "mlp_dim": 512,
+  "dropout": 0.1,
+  "num_actions": 8,
+  "arch": "unified_hist_cm",
+  "backbone": "unified_hist",
+  "click_embed": true,
+  "hist_n": 3,
+  "patch_size": 2,
+  "num_input_colors": 17,
+  "conditioning": "adaln_zero_all_layers + learned vector on the clicked patch token (zero-init)",
+  "norm": "pre",
+  "positional": "rope_only",
+  "task_token": false,
+  "attention": "sdpa"
+ },
+ "world_size": 6,
+ "loss": "ffp_only",
+ "folds": {
+  "num_folds": 1,
+  "split_seed": 0,
+  "test_games_by_fold": {
+   "0": []
+  }
+ }
+}

fold1/checkpoint_best.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e002d2197aa9518cb9075562411d35f01c4cecd9abd02c81483cc4d597fb498e
+size 146368583

fold1/config.json

@@ -0,0 +1,100 @@
+{
+  "fold": 1,
+  "split_seed": 0,
+  "data_root": "ARC-AGI-3",
+  "save_dir": "runs/fold1",
+  "patch_size": 2,
+  "embed_dim": null,
+  "depth": null,
+  "num_heads": null,
+  "mlp_dim": null,
+  "hist_n": 3,
+  "epochs": 30,
+  "batch_size": 86,
+  "global_batch_size": 516,
+  "learning_rate": 0.0003,
+  "weight_decay": 0.0,
+  "max_grad_norm": 1.0,
+  "lr_scheduler": "cosine",
+  "seed": 42,
+  "num_workers": 8,
+  "optimizer": "muon",
+  "muon_lr": 0.0,
+  "muon_momentum": 0.95,
+  "muon_adjust_lr": "match_rms_adamw",
+  "gamma": 0.95,
+  "no_amp": false,
+  "amp_dtype": "bf16",
+  "no_compile": false,
+  "compile_mode": "reduce-overhead",
+  "distributed": true,
+  "disable_augment": false,
+  "num_color_permute": 9,
+  "eval_batch_size": null,
+  "eval_every": 1,
+  "max_train_batches": 0,
+  "resume_checkpoint": null,
+  "use_wandb": false,
+  "wandb_project": "warc",
+  "wandb_run_name": "",
+  "wandb_resume_id": "",
+  "arch": "unified_hist_cm",
+  "model_config": {
+    "embed_dim": 512,
+    "depth": 10,
+    "num_heads": 8,
+    "mlp_dim": 512,
+    "dropout": 0.1,
+    "num_actions": 8,
+    "arch": "unified_hist_cm",
+    "backbone": "unified_hist",
+    "click_marker": true,
+    "hist_n": 3,
+    "patch_size": 2,
+    "num_input_colors": 17,
+    "conditioning": "adaln_zero_all_layers + click marker on the clicked patch token (zero-init)",
+    "norm": "pre",
+    "positional": "rope_only",
+    "task_token": false,
+    "attention": "sdpa"
+  },
+  "world_size": 6,
+  "loss": "ffp_only",
+  "folds": {
+    "num_folds": 3,
+    "split_seed": 0,
+    "test_games_by_fold": {
+      "1": [
+        "bp35",
+        "lf52",
+        "lp85",
+        "r11l",
+        "re86",
+        "sb26",
+        "sc25",
+        "su15",
+        "wa30"
+      ],
+      "2": [
+        "cd82",
+        "cn04",
+        "dc22",
+        "m0r0",
+        "sk48",
+        "tn36",
+        "tr87",
+        "tu93"
+      ],
+      "3": [
+        "ar25",
+        "ft09",
+        "g50t",
+        "ka59",
+        "ls20",
+        "s5i5",
+        "sp80",
+        "vc33"
+      ]
+    }
+  }
+}

fold2/checkpoint_best.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7b9fee3f076d99ba05fd9ee1165a1ac4e4f11cac4bdfacf09ce729e3256c9982
+size 146368583

fold2/config.json

@@ -0,0 +1,100 @@
+{
+  "fold": 2,
+  "split_seed": 0,
+  "data_root": "ARC-AGI-3",
+  "save_dir": "runs/fold2",
+  "patch_size": 2,
+  "embed_dim": null,
+  "depth": null,
+  "num_heads": null,
+  "mlp_dim": null,
+  "hist_n": 3,
+  "epochs": 30,
+  "batch_size": 86,
+  "global_batch_size": 516,
+  "learning_rate": 0.0003,
+  "weight_decay": 0.0,
+  "max_grad_norm": 1.0,
+  "lr_scheduler": "cosine",
+  "seed": 42,
+  "num_workers": 8,
+  "optimizer": "muon",
+  "muon_lr": 0.0,
+  "muon_momentum": 0.95,
+  "muon_adjust_lr": "match_rms_adamw",
+  "gamma": 0.95,
+  "no_amp": false,
+  "amp_dtype": "bf16",
+  "no_compile": false,
+  "compile_mode": "reduce-overhead",
+  "distributed": true,
+  "disable_augment": false,
+  "num_color_permute": 9,
+  "eval_batch_size": null,
+  "eval_every": 1,
+  "max_train_batches": 0,
+  "resume_checkpoint": null,
+  "use_wandb": false,
+  "wandb_project": "warc",
+  "wandb_run_name": "",
+  "wandb_resume_id": "",
+  "arch": "unified_hist_cm",
+  "model_config": {
+    "embed_dim": 512,
+    "depth": 10,
+    "num_heads": 8,
+    "mlp_dim": 512,
+    "dropout": 0.1,
+    "num_actions": 8,
+    "arch": "unified_hist_cm",
+    "backbone": "unified_hist",
+    "click_marker": true,
+    "hist_n": 3,
+    "patch_size": 2,
+    "num_input_colors": 17,
+    "conditioning": "adaln_zero_all_layers + click marker on the clicked patch token (zero-init)",
+    "norm": "pre",
+    "positional": "rope_only",
+    "task_token": false,
+    "attention": "sdpa"
+  },
+  "world_size": 6,
+  "loss": "ffp_only",
+  "folds": {
+    "num_folds": 3,
+    "split_seed": 0,
+    "test_games_by_fold": {
+      "1": [
+        "bp35",
+        "lf52",
+        "lp85",
+        "r11l",
+        "re86",
+        "sb26",
+        "sc25",
+        "su15",
+        "wa30"
+      ],
+      "2": [
+        "cd82",
+        "cn04",
+        "dc22",
+        "m0r0",
+        "sk48",
+        "tn36",
+        "tr87",
+        "tu93"
+      ],
+      "3": [
+        "ar25",
+        "ft09",
+        "g50t",
+        "ka59",
+        "ls20",
+        "s5i5",
+        "sp80",
+        "vc33"
+      ]
+    }
+  }
+}

fold3/checkpoint_best.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:651ca2676710cc6b40d1f1eb4d4d4b838359dadc92038f20ae5c2d5f21c88ed7
+size 146368583

fold3/config.json

@@ -0,0 +1,100 @@
+{
+  "fold": 3,
+  "split_seed": 0,
+  "data_root": "ARC-AGI-3",
+  "save_dir": "runs/fold3",
+  "patch_size": 2,
+  "embed_dim": null,
+  "depth": null,
+  "num_heads": null,
+  "mlp_dim": null,
+  "hist_n": 3,
+  "epochs": 30,
+  "batch_size": 86,
+  "global_batch_size": 516,
+  "learning_rate": 0.0003,
+  "weight_decay": 0.0,
+  "max_grad_norm": 1.0,
+  "lr_scheduler": "cosine",
+  "seed": 42,
+  "num_workers": 8,
+  "optimizer": "muon",
+  "muon_lr": 0.0,
+  "muon_momentum": 0.95,
+  "muon_adjust_lr": "match_rms_adamw",
+  "gamma": 0.95,
+  "no_amp": false,
+  "amp_dtype": "bf16",
+  "no_compile": false,
+  "compile_mode": "reduce-overhead",
+  "distributed": true,
+  "disable_augment": false,
+  "num_color_permute": 9,
+  "eval_batch_size": null,
+  "eval_every": 1,
+  "max_train_batches": 0,
+  "resume_checkpoint": null,
+  "use_wandb": false,
+  "wandb_project": "warc",
+  "wandb_run_name": "",
+  "wandb_resume_id": "",
+  "arch": "unified_hist_cm",
+  "model_config": {
+    "embed_dim": 512,
+    "depth": 10,
+    "num_heads": 8,
+    "mlp_dim": 512,
+    "dropout": 0.1,
+    "num_actions": 8,
+    "arch": "unified_hist_cm",
+    "backbone": "unified_hist",
+    "click_marker": true,
+    "hist_n": 3,
+    "patch_size": 2,
+    "num_input_colors": 17,
+    "conditioning": "adaln_zero_all_layers + click marker on the clicked patch token (zero-init)",
+    "norm": "pre",
+    "positional": "rope_only",
+    "task_token": false,
+    "attention": "sdpa"
+  },
+  "world_size": 6,
+  "loss": "ffp_only",
+  "folds": {
+    "num_folds": 3,
+    "split_seed": 0,
+    "test_games_by_fold": {
+      "1": [
+        "bp35",
+        "lf52",
+        "lp85",
+        "r11l",
+        "re86",
+        "sb26",
+        "sc25",
+        "su15",
+        "wa30"
+      ],
+      "2": [
+        "cd82",
+        "cn04",
+        "dc22",
+        "m0r0",
+        "sk48",
+        "tn36",
+        "tr87",
+        "tu93"
+      ],
+      "3": [
+        "ar25",
+        "ft09",
+        "g50t",
+        "ka59",
+        "ls20",
+        "s5i5",
+        "sp80",
+        "vc33"
+      ]
+    }
+  }
+}

manifest.json

@@ -0,0 +1,81 @@
+{
+ "arch": "unified_hist_cm",
+ "model_config_note": "see each checkpoint config.model_config",
+ "game_to_fold": {
+  "bp35": 1,
+  "lf52": 1,
+  "lp85": 1,
+  "r11l": 1,
+  "re86": 1,
+  "sb26": 1,
+  "sc25": 1,
+  "su15": 1,
+  "wa30": 1,
+  "cd82": 2,
+  "cn04": 2,
+  "dc22": 2,
+  "m0r0": 2,
+  "sk48": 2,
+  "tn36": 2,
+  "tr87": 2,
+  "tu93": 2,
+  "ar25": 3,
+  "ft09": 3,
+  "g50t": 3,
+  "ka59": 3,
+  "ls20": 3,
+  "s5i5": 3,
+  "sp80": 3,
+  "vc33": 3
+ },
+ "folds": {
+  "1": {
+   "test_games": [
+    "bp35",
+    "lf52",
+    "lp85",
+    "r11l",
+    "re86",
+    "sb26",
+    "sc25",
+    "su15",
+    "wa30"
+   ],
+   "file": "fold1/checkpoint_best.pt",
+   "sha256": "e002d2197aa9518cb9075562411d35f01c4cecd9abd02c81483cc4d597fb498e"
+  },
+  "2": {
+   "test_games": [
+    "cd82",
+    "cn04",
+    "dc22",
+    "m0r0",
+    "sk48",
+    "tn36",
+    "tr87",
+    "tu93"
+   ],
+   "file": "fold2/checkpoint_best.pt",
+   "sha256": "7b9fee3f076d99ba05fd9ee1165a1ac4e4f11cac4bdfacf09ce729e3256c9982"
+  },
+  "3": {
+   "test_games": [
+    "ar25",
+    "ft09",
+    "g50t",
+    "ka59",
+    "ls20",
+    "s5i5",
+    "sp80",
+    "vc33"
+   ],
+   "file": "fold3/checkpoint_best.pt",
+   "sha256": "651ca2676710cc6b40d1f1eb4d4d4b838359dadc92038f20ae5c2d5f21c88ed7"
+  }
+ },
+ "all25": {
+  "file": "all25/checkpoint_best.pt",
+  "sha256": "348faf916bb842fe27ccef24efa86f79f13f3175c78d1a1e6408fdddf0502922",
+  "note": "the submitted model, pretrained on all 25 public games; not for evaluation on public games"
+ }
+}