# Next: Teaching an Agent to Imagine Code With JEPA
> RIAI Lab Notes #4: why exact-value prediction hit a wall, how a JEPA plans code in vector space, and first results: 10 of 12 when grounded.
- Author: Tai Pham Phu
- Published: 24 Oct 2026
- Tags: RIAI, AI Research, Self-Improving AI, JEPA, World Models
- URL: https://phamphutai.com/2026/10/24/riai-jepa-for-code/
```
WITH model: PASS after 184 real runs | WITHOUT: PASS after 46
```
That's one line from a test log in my RIAI project. The task was "sum of evens": find a small program that turns a list of numbers into the sum of its even ones. Both methods found it. The one **with** a learned model of code needed 184 real runs to get there. The one **without** any model, which simply tried the shortest programs first, needed 46.
The model was supposed to make search smarter. On that task, search with the model needed four times as many real runs.
This is the fourth of my RIAI Lab Notes. [The third post](https://phamphutai.com/2026/10/23/riai-writing-its-own-code/) ended with a coder that builds programs by search, with no language model writing the code, and solved 13 of 13 small tasks. It also ended with a gap. That coder's world model only knows coarse facts: *will this line crash?* and *what kind of value comes out?* The attempt to predict the **exact** value of each line was just starting, at about 2% of single steps right.
This post is about what happened when that exact-value model got much better, why it still didn't help search the way I hoped, and the direction I'm taking next: **JEPA**, applied to Python code.
---
## The vision, again
It hasn't changed: **the terminal does everything; it self-develops; it writes its own Python code to develop skills; and I expect it to keep developing, more and more.**
My next research direction, in one line: **apply JEPA in Python code to develop new skills.**
To see why, it helps to look at what the last day of experiments showed.
---
## What a better exact-value model bought, and what it didn't
### The model got much better at single steps
The step world model reads a line of code and the values going into it, and predicts exactly what happens next: the new value, which line runs next, what's returned, or which error is raised. In post 3, its first checkpoints got about 2% of single steps exactly right on held-out programs.
A few changes later (a code-aware tokenizer that keeps keywords whole and writes numbers digit by digit, and much longer training) the latest log shows, on held-out programs after 31,060 training steps:
| Steps predicted exactly right | Share |
|---|---|
| 1 step | 57% |
| 3 steps in a row | 24% |
| 5 steps in a row | 18% |
That's a real improvement. But look at the shape: every extra step multiplies the chance of a mistake. Imagining a whole program means getting every step right, and every character of every value right.
### Using it alone to guide search made search worse
The test that matters is search. Take a pool of candidate programs, "run" each one in the model's imagination on the examples, rank them by how close the imagined output is to the wanted output, and run them for real in that order.
On 16 tasks, an earlier version of the model (1-step accuracy around 50% on unseen functions) was put up against plain "shortest first":
- **With the model:** 16 of 16 solved, an average of **24.6** real runs per task.
- **Without:** 16 of 16 solved, an average of **14.9** real runs.
In 15 of the 16 tasks, the model imagined **zero** of the candidates exactly right. One wrong digit early on and the imagined output is off. Ranking by "close in characters" then sends the search the wrong way, as in the "sum of evens" line at the top.
### A hand-written calculator rescued it
So I split the job. The neural model decides the **structure**: which line runs next, which branch is taken, when a loop ends. A small **safe calculator** I wrote computes exact **values** for the lines it understands. It never runs code; it walks the expression tree and knows only a whitelist of operations. In that test it handled 48% of all lines, and the model decided the rest.
The hybrid worked well:
- **Straight-line tasks:** 14 of 15 solved, an average of **1.0** real run, against **15.4** without a model.
- **Programs with loops and branches:** 8 of 9 solved, an average of **19.6** real runs, against **61.0** without.
- **Text tasks** (most common word, count long words and so on): 11 of 12 solved, **22.6** real runs on average against **99.1**.
But be honest about where the values came from. The part that got the values right is a calculator *I* wrote. That's not the agent developing; that's me developing it.
### Two more honest results
**The flywheel didn't turn.** The agent's verified solutions were traced on new texts, added to the training data, and the model fine-tuned for 3 minutes. It imagined 2 of 11 of its own solutions exactly right on new examples before, and still **2 of 11** after.
**Writing code directly is still weak.** Taught to write the program from examples, on 12 exercises like "sort the array", "median" and "most common word", the same model's first guess was right **0 of 12** times, and one of 5 guesses **2 of 12**. A version shown the real values after each line it writes also got **2 of 12** within 5.
### The gap, in one sentence
**Predicting exactly what code produces is hard, and search doesn't need it.** Search needs to know which partial program is *heading toward* the answer. That's a question about meaning, not about characters.
That's the gap JEPA is built for.
---
## JEPA in plain words
I covered the debate around world models in [an earlier post](https://phamphutai.com/2026/10/20/world-models/), so here's only what matters for code.
In 2022, Yann LeCun published a position paper, *A Path Towards Autonomous Machine Intelligence*. Its central building block is the **Joint-Embedding Predictive Architecture**, or JEPA. The idea:
1. Take two related pieces of data, call them **x** and **y**. For video, x could be the frames so far and y the frames that come next.
2. Pass each through an **encoder** that turns it into a compact vector, a **representation** (also called an embedding).
3. Train a **predictor** to guess the representation of y from the representation of x.
The key move: **the prediction is made, and checked, in representation space, not in pixels or tokens.** A generative model has to predict every pixel of the next frame, including the exact shape of every leaf moving in the wind, which can't be predicted. A JEPA's encoder is free to leave out details that can't be predicted and keep what can.
### The catch: collapse
There's an easy way to cheat. If the encoder outputs **the same vector for everything**, the predictor is always perfectly right. The loss is zero and nothing has been learned. This is called **collapse**, and every JEPA needs a way to block it. LeCun's paper discusses two families: **contrastive** methods, which push representations of unrelated things apart, and **regularised** methods, which keep the representations informative directly. One regulariser it discusses is **VICReg** (Bardes, Ponce and LeCun, 2022). VICReg keeps the spread (variance) of each dimension of the vectors above a threshold, and pushes different dimensions to carry different information (low covariance).
The JEPA models Meta built since then lean on a different guard:
- **I-JEPA** (Assran and others, CVPR 2023) hides blocks of an image and predicts their representations from one visible block. The targets come from a **target encoder** whose weights are an **exponential moving average (EMA)** of the main encoder: a slowly-updated copy. The paper names collapse as a concern and uses this asymmetry between the two encoders to avoid it.
- **V-JEPA** (Bardes and others, 2024) does the same for video, trained on about 2 million videos with no text, no pixel reconstruction and no negative examples.
- **V-JEPA 2** (Assran and others, June 2025) scales to over a million hours of video, with an EMA target encoder and a **stop-gradient** (no learning signal flows back through the targets) to help prevent collapse. Closest to my problem is **V-JEPA 2-AC**, an *action-conditioned* predictor trained on under 62 hours of robot video: given the current representation and an action, it predicts the next one. To plan, it searches for actions whose predicted future is **closest to a goal image's representation**. Robot arms in two labs picked and placed objects that way, zero-shot.
### JEPA and code, so far
JEPA has been tried with language and code too. **LLM-JEPA** (Huang, LeCun and Balestriero, September 2025) treats paired data such as a natural-language description and the matching regular expression or SQL query as **two views of the same thing**. On top of the usual next-token loss, it adds a term that asks the model to predict the code's embedding from the text's embedding. On one regex benchmark, fine-tuning a 1-billion-parameter Llama model, accuracy went from 57.29% to 71.46%. The paper also found the usual next-token loss can't simply be dropped.
That's a different use from mine. LLM-JEPA pairs a description with its code. I want to pair a program **state** with the state **after** a line runs.
---
## Why code fits JEPA
Think of a program running on a few example inputs. At every moment it has a **state**: the current values of its variables. Each line, or each block of lines, is an **action** that turns one state into the next.
That's the V-JEPA 2-AC setup almost exactly: a state, an action, a next state, and a goal.
- **The state before and after a step are the two views.** The predictor sees the first plus the step's code, and guesses a summary of the second.
- **The goal is a state too.** A task given as examples already says what final state you want. Encode it, and you have a vector to plan toward.
- **Meaning matters more than spelling.** `[1, 2, 5, 9]` and `[1, 2, 5, 10]` differ by one character and mean nearly the same; `15` and `[15]` differ by two and mean something quite different. A learned representation can, in principle, place them by what they *are*. Whether mine does is something to measure.
- **Data is free.** Every program run gives (state, action, next state). Post 3 already collected 61,611 traced programs, though the first JEPA doesn't use them yet.
---
## The design, as the code does it today
The new file is `jepa.py`. Its first line describes itself as a *"LeCun-style JEPA world model for code, and planning in representation space"*. Here's what it does, piece by piece.
The JEPA for code: predict the summary of the next state, not its exact characters.
**What a state is.** The current value on every example, written as text and joined with a bar, like `[1, 2, 5, 9] | [4, 17, 30] | [3, 8]`. Each value is cut to at most 40 characters. One state covers all examples at once, so the vector describes what's happening on all of them together.
**What an action is.** One of **57 building blocks**, each a line or a few lines of Python with a known input type and output type. 25 work on lists and numbers, like `sorted(p0)`, `[v for v in p0 if v > 10]`, `sum(p0)` or `p0[len(p0) // 2]`. 32 work on text, like `.split()`, counting words into a dictionary, `max(words, key=len)` or a small loop that counts words longer than 3 letters.
**The encoder.** A small transformer: 2 layers, 4 attention heads, vectors of 128 numbers, reading at most 96 tokens with the same code tokenizer as the step world model. It averages over the tokens and maps the result to the final 128-number vector **z**.
**The predictor.** It takes z and a learned 64-number vector for the action, and passes both through a small network that outputs a **change** to z. The prediction is z plus that change:
```python
return z + self.mlp(torch.cat([z, self.act(a)], -1))
```
That's a natural fit for code: most steps change some things and leave the rest alone.
**The target encoder.** A copy of the encoder that is never trained directly. After every training step, its weights move 1% of the way toward the encoder's weights (an EMA with rate 0.99), and no gradient flows through it. It encodes the state **after** the step.
**The loss.** The mean squared distance between the predicted z and the target z, plus the two VICReg terms on the encoder's z: one that keeps each dimension's spread above 1 (weight 1.0) and one that pushes dimensions to be uncorrelated (weight 0.04). So it uses **both** guards against collapse: the EMA target encoder from I-JEPA and V-JEPA, and a VICReg-style regulariser from the 2022 paper.
**The training data.** Not the 61,611 traces, at least not yet. The script makes its own: 240,003 transitions, each built by taking 3 random inputs (lists of numbers or bits of text), applying 1 to 4 random building blocks whose types fit, and recording (state before, block, state after). The values come from the **safe calculator**, so this needs neither the VM nor the neural step model. 3,000 more transitions, made with a different random seed, are kept aside for checking.
**How progress is checked during training.** For 100 held-out transitions at a time, predict the next z, then ask: among the 100 real next states, is the true one the **closest** to the prediction? Guessing at random would be right 1% of the time.
### How it plans
This is where the coarse world model, the exact-value model and the JEPA differ the most.
Planning toward the goal's vector: only 5 programs reach reality.
Given a task as examples:
1. Encode the example **inputs** into a start vector, and the wanted **outputs** into a **goal vector**.
2. Grow programs block by block, in a **beam search**: try every block whose input type fits the current value (and don't repeat the same block twice in a row), and let the predictor move each candidate's vector. **No code runs.**
3. Keep the **64 candidates** whose predicted vector is closest to the goal, and repeat, up to **4 blocks**.
4. Whenever a candidate ends with the right **kind** of output (a number, a list or text, matching the examples), it becomes a finished program. Sort the finished ones by distance to the goal.
5. Send only the **best 5** to the offline VM, where they're checked on the 3 examples and a **fresh input** none of them saw.
This is V-JEPA 2-AC's idea brought down to toy size: plan by distance to a goal's representation. Mine uses a beam search over discrete blocks instead of that paper's sampling method over continuous robot actions.
There's also a **closed-loop** variant. After each step, instead of trusting the predicted vector, it computes the real next state with the safe calculator and encodes that. It's a check on how much error the predictor adds when it imagines several steps in a row.
Compare the three world models side by side:
| World model | Imagines | Ranks candidates by |
|---|---|---|
| Coarse (post 3) | crash or not, kind of value | doesn't rank; only rules out |
| Exact values (stepwm) | the exact output, character by character | how close the imagined output is in characters or size |
| JEPA (new) | a 128-number summary of the state | distance to the goal's summary |
---
## Results so far
Labels matter here, so each line says what it is.
- **Training data (measured):** 240,003 transitions from 57 building blocks, made with the safe calculator.
- **First checkpoint (measured):** after 794 training steps, the predicted vector picked the true next state out of 100 candidates **59%** of the time. Chance would be 1%.
Then the first test finished, the same afternoon. The model had trained for 2,126 steps.
### Test 1: the 12 exercises
The best 5 plans for each exercise were checked in the offline VM on the examples and a fresh input.
| Planning | Right on first plan | Right within 5 |
|---|---|---|
| **Open loop** (imagination only, no real states) | **1 of 12** | 4 of 12 |
| **Closed loop** (real state re-encoded after each step) | **10 of 12** | 10 of 12 |
| Best earlier model that writes code token by token | 1 of 12 | 2 of 12 |
The closed-loop planner found programs for sorting, averages, most common word, number of different words, counting long words and more. It failed on **median** and **range**, which can't be written as one chain of these blocks. The log also notes that one of the 10 was **lucky**: for "count numbers above 10", the program it found drops the first number and then adds one, which happened to pass every check but isn't the right logic.
### Test 2: held-out random programs
167 random programs made of the same blocks, never seen as whole programs during training: could the planner find a program that matches their behaviour?
| Planning | First plan right | Within 5 |
|---|---|---|
| Open loop | 39% | 56% |
| Closed loop | **91%** | **96%** |
### JEPA v2: choosing statements, not blocks
Within the hour I tried a second design, `jepa2.py`. The changes:
- The **state** is now every variable of the program so far, on every example.
- An **action** is one Python statement, read as code text by an action encoder, picked from a **grammar** of statements (one-input operations, filters with constants, small loops, two-input steps), so median, range and average are expressible.
- Besides the world model, it learns a **next-step scorer**: given the state, the goal and a candidate statement, does this statement lead toward the goal? It's trained by contrast: the real next statement of a solution against other statements that were possible at that point.
- Writing is closed loop by design. The scorer picks statements, the safe calculator runs them on the examples, and it stops as soon as a variable equals the goal on every example. The finished program is verified in the VM.
After 1,468 training steps: **8 of 12 right on the first try**, each in about half a second to a second and a half, with short, clean programs, for example:
```python
x1 = 0
for v in p0:
if v % 2 == 0:
x1 = x1 + v
return x1
```
It missed **median, range, average and number of different words**: for those, nothing it built reached the goal.
---
## The research plan
These are hypotheses, each with what would count as a yes. The test code for the first two is already written in `jepa.py`.
### 1. Does planning by vectors find right programs?
**Test:** the same 12 exercises as the coder tests above, with 3 examples chosen to rule out programs that fit by luck, plus a fresh input. Then up to 400 random programs made of the same blocks, with a different random seed from training, as many as fit in the time limit.
**A yes:** clearly more than the 2 of 12 the line-by-line writers reached within 5 guesses. By my reading of the block list, "range" (largest minus smallest) can't be written as one chain of these blocks, so 12 of 12 isn't possible. That's fine: the question is how many of the reachable ones it finds.
### 2. Does imagination drift?
**Test:** open loop (pure prediction for up to 4 steps) against closed loop (real state re-encoded after each step).
**A yes for the JEPA idea:** the two are close. A big gap would mean the predictor's errors pile up over steps, the same problem the exact-value model had, just in a different space.
### 3. Fewer real runs than the alternatives
**Test:** on the same task pools, count real runs per solved task for JEPA ranking, the exact-value model alone, the hybrid with my calculator, and plain shortest first.
**A yes:** JEPA beats the exact-value model alone (24.6 on the 16-task test) and approaches the hybrid **without** needing my hand-written calculator at planning time. If it needs the calculator to work, that's worth knowing too.
### 4. Does it generalise?
**Test:** task types and input kinds not in the training data, and states longer than anything seen.
**A yes:** distance to the goal still sorts right programs above wrong ones, measured as how high the correct program ranks.
### 5. Does experience compound?
This is the one that connects to **"more and more"**. When a task is solved and checked, the program can become a **new block**, a new skill. Then:
- the new block gets new transitions, and the JEPA is trained on them;
- the 61,611 traced programs from post 3, and the agent's own tools, become more training data;
- the next task can use the new block as a single step.
**A yes:** tasks that were out of reach become reachable after earlier skills are added, with the vectors trained on the agent's own experience, and checked on fresh inputs. The earlier flywheel test, which stayed at 2 of 11, is the bar to beat.
None of this is plugged into the terminal agent yet. Today `jepa.py` and `jepa2.py` are standalone experiments, like `wmcoder` was in post 3. The terminal agent still gets its new tools written by Claude Haiku.
---
## Temporary conclusion
This is where things stand on 24 October 2026, after one afternoon of JEPA tests. It's a temporary conclusion: small tests, early models, and it may change with the next run.
**1. Planning in representation space works, when it stays grounded.** With the real state re-encoded after every step, the JEPA planner solved 10 of 12 exercises on its first plan and matched 91% of held-out programs. The best models that write code token by token reached 1 or 2 of 12. Judging "is this step moving toward the goal?" in vector space was far more useful than predicting exact output.
**2. Pure imagination still drifts.** Open loop, the same model solved 1 of 12 and matched 39% of held-out programs. So hypothesis 2 has a first answer: the predictor's errors pile up over steps. Predicting summaries instead of exact characters didn't remove that problem; it moved it.
**3. The calculator is still doing important work.** Both the best JEPA v1 result and JEPA v2 rely on the safe calculator to compute real states between steps, and it also makes all the training data. That's a tool I wrote, not something the agent learned. The fair claim today is: **the learned part decides *which step to take*; a hand-written part tells it *where it is*.**
**4. Statement-level writing is promising.** JEPA v2 picked whole Python statements from a grammar and wrote short, clean programs for 8 of 12 exercises in about a second each, with no lucky passes. Its misses (median, range, average, number of different words) show what the grammar and the scorer can't yet reach.
**5. Compounding is still untested.** Nothing here yet shows "more and more": solved tasks becoming skills that make the next tasks easier. That's the test that matters most for the vision, and it hasn't been run.
So the direction holds up well enough to keep going: **JEPA-style scoring plus grounding beats token-by-token writing on these tasks.** The next steps are to shrink the hand-written part (learn the state from real VM traces instead of the calculator), close the gap on the four missed exercises, and run the compounding test.
### Risks I'm watching
- **Collapse, or near-collapse.** Both guards are on, but a space can avoid full collapse and still squash the differences that matter. The training log tracks the spread term for this reason.
- **Vectors that don't separate right from wrong.** The goal vector is an encoding of the wanted outputs. If "almost right" and "right" land close together, the ranking is no better than shortest first. Hypothesis 4 measures this directly.
- **The calculator is still in the loop.** It makes all the training data. Anything it can't compute (things outside its whitelist) the JEPA never sees. Moving to traces from the real VM is part of the plan.
- **Small test sets.** 12 exercises and a few hundred random programs. And I chose the blocks and the exercises, which helps any method. Fresh-input checks keep it honest about correctness, not about task choice.
- **These are first runs.** One model per design, 12 exercises and 167 held-out programs. The numbers may move either way.
---
## Why this direction
Every earlier step pushed toward imagining code more precisely. The tests say precision is expensive and fragile: one wrong digit and the imagined output misleads the search. JEPA asks a different question: not *"what exactly will this print?"* but *"is this step moving toward the goal?"*
If that works, it gives the agent something it doesn't have yet: a way to judge partial programs that it learned itself, from its own experiments, that gets better as it builds more skills. That's what I'll test next. Whatever comes out gets published, including a no.
---
## Key ideas
- **The exact-value world model got much better:** 57% of single steps exactly right on held-out programs, up from about 2%. But only 18% of 5-step runs.
- **Exact imagination alone made search worse:** 24.6 real runs per task with the model against 14.9 without, on 16 tasks.
- **A hand-written calculator fixed the values,** for example 1.0 real runs against 15.4, but that's my work, not the agent's.
- **JEPA predicts a summary of the next state, not its exact characters,** and plans by distance to the goal's summary, as V-JEPA 2-AC does for robot arms.
- **`jepa.py` uses both anti-collapse guards:** an EMA target encoder with no gradient, and VICReg-style spread and decorrelation terms.
- **First results (temporary):** closed-loop JEPA planning solved 10 of 12 exercises (one lucky) and 91% of held-out programs; open loop only 1 of 12 and 39%. JEPA v2 solved 8 of 12 cleanly, against 1 to 2 of 12 for token-by-token writers.
- **The safe calculator still grounds every step,** so the learned part chooses steps while a hand-written part tracks the state.
- **The real test is compounding:** solved tasks becoming blocks, and the next tasks getting easier.
*Sources: [LeCun, "A Path Towards Autonomous Machine Intelligence", 2022](https://openreview.net/forum?id=BZ5a1r-kVsf); [Bardes, Ponce and LeCun, "VICReg", ICLR 2022](https://arxiv.org/abs/2105.04906); [Assran et al., "Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture" (I-JEPA), CVPR 2023](https://openaccess.thecvf.com/content/CVPR2023/html/Assran_Self-Supervised_Learning_From_Images_With_a_Joint-Embedding_Predictive_Architecture_CVPR_2023_paper.html); [Bardes et al., "Revisiting Feature Prediction for Learning Visual Representations from Video" (V-JEPA), 2024](https://arxiv.org/abs/2404.08471); [Assran et al., "V-JEPA 2", June 2025](https://arxiv.org/abs/2506.09985); [Huang, LeCun and Balestriero, "LLM-JEPA", September 2025](https://arxiv.org/abs/2509.14252). Checked October 2026.*
*Related: [Writing Its Own Code: From a Teacher Model to Code Built by Search](https://phamphutai.com/2026/10/23/riai-writing-its-own-code/), and [World Models: Can AI Simulate Reality?](https://phamphutai.com/2026/10/20/world-models/).*