tiny-jepa

World Labs for your laptop. A tiny, hackable, action-conditioned JEPA (Joint Embedding Predictive Architecture) world model — in pure numpy, zero required dependencies. The whole thing is ~1,500 lines you can read in an afternoon. Inspired by Yann LeCun's predictive world-model program and the spatial-intelligence spirit of Fei-Fei Li's World Labs.

The shipped checkpoint, hunting pellets with learned vision — no hand-coded game AI, just a world model that imagines, perceives, and plans.

Install

pip install git+https://github.com/OpenSourceRulesz/tiny-jepa.git
# PyPI (pip install tiny-jepa) coming soon
# extras: pip install "tiny-jepa[onnx]"  (ONNX export)
#         pip install "tiny-jepa[gym]"   (Gymnasium wrapper)

5-line quickstart

from tiny_jepa import TinyJEPA, OrbyWorld
model = TinyJEPA.load("tiny_jepa/checkpoints/orby_v1.npz")  # pretrained
env = OrbyWorld(seed=0)
for _ in range(64):
    env.step(model.plan(env.render(), env.agent))  # see, imagine, act

That's it — a world model hunting pellets from raw pixels. Now go crazy: model.imagine(z, actions) rolls out latent futures, model.dream(frame, actions) decodes them to images, train_on_env / train_on_data train your own worlds.

Try it with zero installs: — or drive it yourself in the playable web demo.

What's inside

tiny-spatial — does your world model understand space?

from tiny_jepa.bench import run_benchmark
run_benchmark("tiny_jepa/checkpoints/orby_v1.npz").print()

Five probes — object permanence under occlusion, hallucination (does it collapse to the dataset prior when blind?), surprise at expectation violations, spatial memory across rooms, memory updates. The shipped checkpoint scores spatial_iq 0.501: it notices surprises and limps along under occlusion, but half-hallucinates when blind and can't do rooms. Beat it.

Hierarchical JEPA — two timescales

tiny_jepa/hjepa.py adds a slow level on top of the frozen fast model: every 8 steps it picks a subgoal (pellet when visible, novel territory or last-seen position when not) using a learned command-conditioned forward model, then steers to it greedily with event-driven replans.

from tiny_jepa import HierarchicalJEPA
hier = HierarchicalJEPA.load("tiny_jepa/checkpoints/orby_v1.npz",
                             "tiny_jepa/checkpoints/orby_hier_v1_slow.npz")
hier.plan(frame, agent_xy)   # drop-in for TinyJEPA

The tradeoff is principled: hierarchy costs a little reactive hunting skill and buys spatial reasoning — rooms, exploration, last-seen memory. Same story as H-JEPA vs flat JEPA in the literature.

Challenges — make your mark

challenges/ has concrete quests with success metrics, from "beat the greedy planner" to "teach it object permanence." The leaderboard is empty. Your move.

Honest limits

32-dim latent, ~200k params, toy 2D world — not a general world model. Dreams blur with horizon; the planner is greedy, not full MPC. The architecture is the real JEPA idea end to end, and it's yours to extend: hierarchical timescales and CEM planning are the obvious next frontiers (see challenges/).

Citation

@software{tiny_jepa_2026,
  title  = {tiny-jepa: a tiny, hackable, action-conditioned JEPA world model},
  author = {OpenSourceRulesz},
  year   = {2026},
  url    = {https://github.com/OpenSourceRulesz/tiny-jepa}
}

MIT licensed.