4RyBbR8wrV40oYDWyIAtcQ.mp4

micro-jepa

micro-jepa is a compact, from-scratch implementation of a JEPA-style architecture for a 2D navigation task. It is designed as an educational toy: the agent collects experience, learns a latent world model, and plans actions in latent space without relying on reward labels or text.

The project is intentionally lightweight and CPU-friendly. It demonstrates the core idea behind Joint Embedding Predictive Architecture: predict the next latent state rather than reconstructing raw observations, then use that predictive abstraction to plan.

---

Why this project exists

Modern autoregressive systems such as LLMs are trained to predict the next token in a high-dimensional symbolic space. That makes them extremely good at pattern matching, but it also creates a structural tendency to hallucinate when asked to reason about the physical world.

JEPA-style systems try to avoid this by changing the objective:

• encode observations into a latent embedding;
• predict future latent states instead of future pixels or tokens;
• plan by comparing imagined future states against a task objective in latent space.

This repository makes that loop concrete with a small 2D environment, a simple MLP encoder, a residual world model, and a planning actor.

---

What is implemented

The project follows the six-module structure commonly associated with LeCun's autonomous machine intelligence view:

#	Module	File	Role
1	Configurator	micro_jepa/configurator.py	Converts a task into a structured goal and its associated cost weights
2	Perception / Encoder	micro_jepa/perception.py	Maps raw state vectors into a latent representation
3	World Model	micro_jepa/world_model.py	Predicts the next latent state from the current latent state and action
4	Cost / Energy	micro_jepa/cost.py	Scores imagined future states in latent space
5	Actor	micro_jepa/actor.py	Uses MPC-style rollouts to choose the next action
6	Memory	micro_jepa/memory.py	Stores recent transitions for self-supervised training

The environment itself lives in micro_jepa/environment.py. It is outside the JEPA core, just as physics or sensors are outside a real agent's internal model.

---

Training pipeline

The project runs in three phases:

1. Phase 1: self-supervised exploration

   • the agent acts randomly in the environment;
   • transitions are stored in the memory buffer.

2. Phase 2: joint training of encoder and world model

   • the encoder maps raw states to latents;
   • the world model predicts the next latent from the current latent and an action;
   • the loss is a latent-space prediction loss plus a lightweight geometric anchor regularization on the first two coordinates of the state.

3. Phase 3: planning and inference

   • the encoder and world model are frozen;
   • the actor performs MPC-style rollouts in the learned latent world;
   • the visualization shows the agent, moving obstacles, and the goal marker.

The core training objective is:

z_t = encoder(x_t)
z_next_target = encoder(x_next).detach()
z_next_pred = world_model(z_t, a_t)
loss = mse(z_next_pred, z_next_target) + geometric_anchor_loss

The stop-gradient on the target is important. Without it, the system could collapse to a trivial constant solution.

---

Installation

Requirements:

• Python 3.9+
• torch
• pydantic
• matplotlib

Install them with:

pip install torch pydantic matplotlib

---

Quick start

git clone <your-repo-url>
cd micro-jepa
python train.py

The script will:

• collect random transitions,
• train the encoder and world model,
• launch the planning visualization.

You can tune the main hyperparameters at the top of train.py:

• EXPLORATION_STEPS: more exploration data usually improves the learned world model;
• TRAINING_EPOCHS: longer training can improve latent consistency;
• MPC_CANDIDATES and MPC_HORIZON: larger values give more search depth but cost more compute;
• the goal weights in the configurator control how strongly the agent prefers reaching the target versus avoiding obstacles.

---

Interactive visualization

During Phase 3, the goal marker is rendered as a draggable green star.

Features:

• click the goal marker to start dragging it;
• move the mouse to reposition the target;
• the configurator and cost module update immediately so the actor can replan toward the new goal;
• the agent continues to follow the moved goal in real time.

This makes the simulation interactive and helps you inspect how the planning loop reacts to changes in the task.

---

Project layout

micro-jepa/
├── micro_jepa/
│   ├── actor.py
│   ├── configurator.py
│   ├── cost.py
│   ├── environment.py
│   ├── memory.py
│   ├── perception.py
│   ├── world_model.py
│   └── __init__.py
├── tests/
├── train.py
└── README.md

---

Practical notes

This repository is intentionally simple, so it is useful as a didactic reference rather than a production-grade world-model agent.

A few important simplifications are worth keeping in mind:

• the encoder and world model are small MLPs;
• the environment is low-dimensional and fully observable;
• planning uses random-shooting MPC rather than a learned policy;
• the representation is regularized only lightly, so this is not a full-scale JEPA system.

If you want to extend the project further, good next steps are:

• add a target encoder with EMA-style updates;
• add stronger regularization such as VICReg-style variance/covariance losses;
• move from 2D coordinates to image-based observations;
• replace the simple MPC actor with a learned planner or policy.

---

License

MIT. Built for learning, experimentation, and extension.