Trying 512 futures in its head before pressing a button

Expert track

The problem and its formalization

At the end of Chapters 1-3, we have fθ : x → s ∈ R^D (encoder, frozen since Phase 1) and g : (st, at) → ŝ{t+1} (world model, Phase 2). A policy is missing. Mine-JEPA chooses not to learn a parametric policy (no policy gradient, no value function network) — planning is done directly via search in latent space, at every time step, reusing the already trained modules without any additional gradients.

Goal embedding construction

good_frames = frames[food >= 7]          # ~16,000 frames out of 32,000
goal = encoder(good_frames).mean(dim=0)  # [D] — centroid

The goal is a post-hoc centroid in latent space, constructed from a subset of the collection dataset filtered on a survival proxy (food ≥ 7), with no task labels provided during the training of either the encoder or the world model.

Random-shooting MPC

For each time step:
  1. Encode current state  →  s_t
  2. Draw N=512 random action sequences of length H=12
  3. For each sequence: unroll world model for H steps → ŝ_{t+H}
  4. Score: score_i = -MSE(ŝ_{t+H,i}, s_goal)
  5. Execute the first action of the sequence with the best score
  6. Repeat (receding horizon)

class LatentMPCPlanner:
    @torch.no_grad()
    def plan(self, s_current, s_goal):
        s = s_current.expand(self.n_candidates, -1).clone()          # [N, D]
        actions = torch.randint(0, self.n_actions,
                                 (self.n_candidates, self.horizon))    # [N, H]
        for h in range(self.horizon):
            s = self.predictor(s, actions[:, h])                      # [N, D]
        scores = -(s - s_goal).pow(2).mean(dim=1)                     # [N]
        return actions[scores.argmax(), 0].item()

The entire planner fits in about ten lines of code — a direct corollary of pushing all the heavy lifting (scene understanding, dynamics) into the encoder and predictor during previous phases; the planner itself is merely a search algorithm.

Compute budget justification

Sequence space: 17^12 ≈ 6×10^11. The candidate budget (N=512) covers only a tiny fraction of this, but suffices because many actions have similar short-term effects (directional redundancy in Crafter's action space) — i.i.d. sampling adequately covers the broad useful directions without needing exhaustive search. The batched rollout of 512×12 steps runs in < 5 ms on GPU, since each step is just a forward pass of the small predictor (~140K parameters) on a batch of 512 — everything happens in latent space, without ever touching the rendering engine or the real environment.

Receding horizon

Re-planning at every step from the actually observed state, rather than blindly executing the 12-step plan, corrects the compounding error accumulated by the world model on long rollouts (cf. Chapter 03, section on k-step error accumulation). This is a structural choice: robustness comes from the re-planning frequency, not from the accuracy of the world model alone.

Measured results — two distinct evaluations

docs/05_planning.md reports an evaluation over 50 Crafter episodes:

Metric	JEPA-MPC Agent	Random Baseline
Mean reward	~2.1	~1.5
Achievements/episode	~3.0	~2.4
Success rate (≥1 achievement)	~100%	~98%
FPS	~150	—

CLAUDE.md (the official Phase 3 gate) reports distinct numbers on what appears to be a different batch of episodes: 100% success rate, 2.56 achievements/episode against 2.38 for the random baseline (+7.5%), and +14% reward. Both measurements stem from the same Phase 3 and point in the same direction (clear improvement across all three metrics, near-total success), but do not match digit-for-digit — consistent with being two distinct evaluation runs rather than the same measurement reported twice. Both are reported here as-is rather than forced to agree. The very first episode strings together 3 achievements (wakeup, collectwood, place_table) — a sequence requiring locating a tree, approaching it, hitting it, and placing an object, with zero hand-coded task logic.

Why it works

The action-conditioned predictor learned the transition function ŝ{t+1} = g(st, at). When the goal is a centroid of "good" states, the score -MSE(ŝ{t+H}, s_goal) naturally favors action sequences that make the latent trajectory converge toward that region — without any explicit task reward ever being injected into the encoder or predictor training. Goal-oriented behavior emerges entirely from the combination of (world model + search), not from a reinforcement learning objective.

Identified limits and project avenues

Main limitation: the goal embedding is a centroid — it encodes what the goal looks like, not how to reach it. Two visually distinct states can be close in latent space without being functionally equivalent (geometric ambiguity of the latent space); this foreshadows a broader issue — the latent geometry upon which planning relies is not guaranteed by prediction quality alone — which will resurface in a later phase of the project, on a different dataset.

Future extension noted in Phase 3 plan: CEM (Cross-Entropy Method) — iteratively refining the action sampling distribution over several iterations, rather than a single shot of 512 i.i.d. sequences — already implemented for continuous actions in minejepa/ebjepa/planning.py at Phase 3 time, but not yet plugged into Crafter's discrete planner. A discrete version of this idea returns in a later project phase.

The full pipeline at this stage

64×64 Frames  →  [Phase 1 Encoder]  →  s_t  [D=128]
                                          │
                                    s_t + a_t
                                          │
                                [Phase 2 World Model]  →  ŝ_{t+1}
                                          │
                            512 imagined sequences
                                          │
                              [Score vs s_goal]
                                          │
                                     best_action
                                          │
                              Crafter.step(action)
                                          │
                                obs_{t+1}  ←─ (loop)

Everything happens in latent space, except at the two endpoints: the pixel input and the final action.

References (verified, from docs/references/index.md)

• Maes, Le Lidec, Scieur, LeCun, Balestriero, LeWorldModel, arXiv:2603.19312 (2026) — the ratio = valpred/valcopy convention from Chapter 03, the foundation making the latent rollout exploited here reliable.
• Meta FAIR, ebjepa (github.com/facebookresearch/ebjepa) — planning logic (minejepa/ebjepa/planning.py) reused for the CEM extension mentioned at chapter's end.