koushikz1/arc-agi-2-jepa-episodes

ARC-AGI-2 masked-demonstration episodes (JEPA-ready)

A task-level reasoning dataset built from the official ARC-AGI-2 tasks. Every row is an episode: a set of demonstration pairs (context), one held-out test_input, and the target_output the model must produce by inferring the rule shared by the demonstrations. It is not an input grid -> output grid dataset; the unit of learning is the task rule.

Built 2026-09-25 with scripts/build_arc2_episodes.py (seed 20260925). Everything below that is a number was computed from the files in this repository.

What is in it

Row counts per split are in analysis.md and stats.json.

Split rule (task-level, never example-level)

1000 training tasks  --seeded shuffle-->  train 800 | val 100 | test 100
 120 public evaluation tasks           -->  eval_public (benchmark only; never train on it)

No episode of a task appears in more than one split. arcgen_fresh, counterfactual and rule_programs inherit the split of their task, so arcgen_fresh/val and arcgen_fresh/test are held-out rule families with fresh grids — the validation signal you actually want (rule generalisation), rather than a random example split that leaks the rule. The full id lists are in splits.json.

Why an extra 800/100/100 split of the training tasks is not enough on its own: the training tasks are measurably easier than the evaluation tasks (table below), and public checkpoints such as the NVARC qwen3_4b_grids15_sft139 have already been trained on all 1,000 of them. For a model trained from scratch on this dataset, eval_public is a clean 120-task benchmark; for a model initialised from a public ARC checkpoint, only the Kaggle leaderboard is a clean readout.

Measured analysis: training vs public evaluation tasks

Reading: the evaluation set is a different regime, not a random sample of the training set. Grids are 3.4× larger by median area, one task in three uses the maximum 30×30, multi-test tasks are 5× more common, and rule-exact generators exist for none of them. ARC Prize's own report calls the training set "not difficulty-calibrated" and the evaluation sets calibrated against human testing. Any local score on the training distribution therefore overstates evaluation performance; that is the reason the arcgen_fresh held-out families are provided as a relative gate rather than an absolute one.

Schemas

episodes / episodes_aug / arcgen_fresh / sdg_hard

{"episode_id": "0934a4d8_loo_2", "task_id": "0934a4d8", "split": "train", "kind": "loo",
 "target_index": 2, "n_context": 3, "target_shape": [9, 3],
 "context": [{"input": [[...]], "output": [[...]]}, ...],
 "test_input": [[...]], "target_output": [[...]]}

episodes_aug adds base_episode_id, aug_id, dihedral (one of identity, rot90, rot180, rot270, fliph, flipv, transpose, antitranspose) and `colorperm (list of 10 ints, perm[0] == 0). arcgenfresh` / `sdghard carry source instead of kind/target_index`.

tasks: task_id, source_split, split, train, test, n_train, n_test, input_shapes, output_shapes, max_side, max_area, n_colors_in, n_colors_out, shape_relation, any_symmetric_input, output_shape_consistent, arcgen_generator, arcgen_fresh_pairs (test pairs include their outputs).

counterfactual: episode_id, task_id, split, target_shape, negatives: [{kind, detail, grid}].

rule_programs: task_id, split, language, source, license, generator_source.

Grids are lists of lists of ints 0–9, height and width 1–30.

Loading

from datasets import load_dataset
ep   = load_dataset("koushikz1/arc-agi-2-jepa-episodes", "episodes")        # train/val/test/eval_public
aug  = load_dataset("koushikz1/arc-agi-2-jepa-episodes", "episodes_aug")
gen  = load_dataset("koushikz1/arc-agi-2-jepa-episodes", "arcgen_fresh")    # held-out families in val/test
neg  = load_dataset("koushikz1/arc-agi-2-jepa-episodes", "counterfactual")

Using it with a JEPA-style model: what the evidence supports

No JEPA-family model (I-JEPA, V-JEPA 2, DSeq-JEPA) has a published ARC-AGI-1 or ARC-AGI-2 score, and the best verified hidden-set score of any small latent / recursive / object-centric ARC-AGI-2 system is 27.08 (a ~200 M LongT5 with test-time LoRA; HRM 2 %, TRM ~6–11 %, TOPAS 12–19 %). Every system above 50 % on the 2026 Kaggle leaderboard is an LLM-class model trained for weeks on private data. This dataset therefore cannot promise any accuracy; what it does is make the one JEPA variant with a defensible role trainable and testable without leaking the evaluation tasks:

ARC-JEPA-S (task-latent JEPA as a selector and hardness scorer, not as the solver).

• —Context encoder E over the demonstration pairs (+ test input) → task latent z; looped predictor P (weight-tied, K ∈ {1, 4}) → ẑ; EMA target encoder G over (test_input, target_output).
• —Loss L = 1 − cos(ẑ, sg[G(x_t, y_t)]) + InfoNCE(ẑ; G(x_t, y_t) vs G(x_t, y_neg)) with the negatives from counterfactual (wrong orientation, wrong colouring, near-miss corruption, other-task output) plus in-batch targets; optional VICReg term to prevent collapse.
• —Data: episodes_aug + arcgen_fresh/train + sdg_hard for training; arcgen_fresh/val for early stopping; arcgen_fresh/test and episodes/test for the held-out family readout; eval_public untouched.
• —Test time: a few gradient steps on z against the leave-one-out demonstration loss (LPN-style latent search), then score every candidate output y_c of an external proposer by cos(ẑ, G(x_t, y_c)).
• —Keep bar before it touches a submission: it must select at least 3 more correct outputs than the incumbent log-probability selector on saved candidate pools, with no losses on unanimous pools.

Do not use it as a stand-alone solver: the 22-op DSL executors and per-task latent models in the proposals this dataset was built for measure 0–10 % on ARC-AGI-2 in every published or locally measured attempt, and the benchmark was explicitly designed against brute-force program search.

Sources and licences

• —ARC-AGI-2 tasks: https://github.com/arcprize/ARC-AGI-2 (Apache-2.0). The 1,000 training tasks are byte-identical to the Kaggle competition files (verified).
• —ARC-GEN generators and the fresh pairs sampled from them: https://github.com/google/ARC-GEN (Copyright 2025 Google LLC, Apache-2.0). rule_programs redistributes the generator source verbatim with its licence header.
• —sdg_hard: synthetic puzzles written and sandbox-verified by the dataset author's pipeline (CC0).
• —Human-difficulty labels (arcprize human testing) are not included; the referenced repository was not reachable at build time.

Provenance

scripts/build_arc2_episodes.py reproduces every file from the sources above with --seed 20260925 (--aug 8 --arcgen-episodes 12 --sdg-episodes 8). Integrity checks run at build time: split id sets are disjoint; no episode contains its own target pair in the context; 300/300 sampled augmentations invert exactly to their base episode; no counterfactual negative equals its truth.

Citation

@misc{arc_agi_2_jepa_episodes_2026,
  title  = {ARC-AGI-2 masked-demonstration episodes (JEPA-ready)},
  author = {koushikz1},
  year   = {2026},
  url    = {https://huggingface.co/datasets/koushikz1/arc-agi-2-jepa-episodes}
}

Please also cite ARC-AGI-2 (Chollet et al., 2025, arXiv:2505.11831) and ARC-GEN (Google, 2025, arXiv:2511.00162).