Views
No views yet
┌─────────────────────────────────────────────────────────────────────┐
│ ARC-AGI-3 Agent │
├─────────────────────────────────────────────────────────────────────┤
│ Observation Grid (64x64, 16 colors) │
│ ↓ │
│ ┌─────────────┐ │
│ │ Grid-JEPA │ ← I-JEPA adapted for discrete grid worlds │
│ │ Encoder │ 1×1 patches, latent-space prediction │
│ └─────────────┘ │
│ ↓ Latent Representation │
│ ┌─────────────┐ │
│ │ RSSM │ ← Recurrent State-Space Model (DreamerV3-style) │
│ │ World Model │ GRU dynamics + discrete latents │
│ └─────────────┘ │
│ ↓ Hidden State (PERSISTS across levels!) │
│ ┌─────────────┐ ┌─────────────┐ │
│ │ Planning │ ←→ │ Exploration │ │
│ │ (Imagination│ │ (Novelty) │ │
│ │ Rollouts) │ │ │ │
│ └─────────────┘ └─────────────┘ │
│ ↓ │
│ ┌─────────────┐ │
│ │Goal Inference│ ← Discovers objectives from terminal states │
│ └─────────────┘ │
│ ↓ │
│ Action (key, position) → Environment │
└─────────────────────────────────────────────────────────────────────┘arc-jepa/
├── src/
│ ├── models/
│ │ ├── encoder.py # GridPatchEmbed + ViT encoders + EMA
│ │ ├── predictor.py # Action-conditioned predictor
│ │ ├── grid_jepa.py # Complete Grid-JEPA system
│ │ ├── rssm.py # Recurrent State-Space Model
│ │ ├── agent.py # Full ARC agent (JEPA + RSSM + planning)
│ │ └── ttt_adapter.py # LoRA TTT adapter
│ ├── data/ # Dataset loaders + augmentations
│ ├── training/ # Training scripts
│ └── utils/ # Utilities
├── tests/ # Unit tests
└── README.md # This fileagent.py — Complete Agent (Central Module)ARCAgent: Full agent loop encoding the core insight of this projectGoalInferenceModule: Discovers objectives from terminal/done statesExplorationPolicy: Novelty-seeking with undo loop avoidancePlanningModule: Imagination-based action selection via RSSM rolloutsUncertaintyTracker: Hypothesis revision when predictions fail consistentlyencoder.py — Grid-JEPA EncoderGridPatchEmbed: 1×1 patch embeddings for color gridsViTEncoder: Multi-head attention transformer blocksEMATargetEncoder: EMA-updated target encoder (prevents collapse)predictor.py — Action-Conditioned PredictorDiscreteActionEmbed: Embeds (action_key, cell_position) pairsActionConditionedPredictor: Predicts target patches from context + actionGridWorldPredictor: Full predictor + decoder to color logitsrssm.py — Recurrent State-Space Modelobserve(): Update state with new observation (posterior)imagine(): Predict next state given action (prior)rollout(): Imagine future trajectories for planningttt_adapter.py — Test-Time TrainingLoRALayer: Low-rank adaptation (W' = W + BA)PredictorLoRAAdapter: Per-task LoRA on JEPA predictorTTTTrainer: Fine-tunes on demos with augmentation voting| Decision | Rationale |
|---|---|
| 1×1 patches | Grid cells are semantically meaningful, unlike image pixels |
| L2 latent loss | Reconstruction forces modeling irrelevant visual details |
| EMA target encoder | Prevents representation collapse in self-supervised learning |
| Feature conditioning | Outperforms concatenation for action conditioning |
| Straight-through latents | Enables gradient flow through discrete RSSM states |
| State persistence | ARC-AGI-3 levels build on each other |
| Uncertainty tracking | Prevents getting stuck on wrong hypotheses |
| LoRA TTT | Efficient per-task adaptation without catastrophic forgetting |