JEPA World Model

A Joint Embedding Predictive Architecture (JEPA) that learns a self-supervised world model from agent trajectories in a 2D environment with walls, then predicts future latent states conditioned on actions — without ever decoding back to pixels.

Overview

Given a short sequence of observed states and the actions taken between them, the model learns to predict the latent representation of future states rather than reconstructing raw pixels. This follows the JEPA family of architectures (energy-based, non-generative self-supervised learning):

• Online encoder — a CNN that maps a (2, 64, 64) observation (agent + wall channels) to a latent state.
• Target encoder — an EMA copy of the online encoder (weights updated as target = ema_rate * target + (1 - ema_rate) * online), used to produce stable prediction targets and avoid representational collapse.
• Recurrent predictor — a convolutional recurrent module that rolls the latent state forward one action at a time: z_{t+1} = predictor(z_t, a_t).

Training combines a weighted latent MSE loss (predicted vs. target-encoder states, weighted more heavily at later rollout steps), VICReg-style variance and covariance regularization to prevent representational collapse, and a contrastive loss between predicted and target latents.

Because the encoder is never trained to decode pixels, the only way to check whether it's learned anything useful is indirectly: freeze the encoder, train a small probing head (Prober, in models.py) on top of the frozen latents to regress the agent's (x, y) location, and measure that probe's error across several held-out evaluation sets (probe_normal, probe_wall, probe_wall_other, probe_expert) that vary in wall layout and how the trajectories were generated.

Repo structure

File	Purpose
best_model.py	Single-device training loop, model definitions (Encoder, RecurrentPredictor, JEPA), and loss functions
best_model_DDP.py	Same model, adapted for multi-GPU training with PyTorch DistributedDataParallel
main.py	Loads a trained checkpoint and runs the downstream probing evaluation
models.py	Prober (the probing head) and MockModel (a random baseline for testing the eval pipeline)
dataset.py	WallDataset / create_wall_dataloader for the probing datasets
evaluator.py	ProbingEvaluator — trains the prober and reports per-task location loss
normalizer.py, schedulers.py, configs.py	Location normalization, LR scheduling, and config parsing helpers

Setup

pip install -r requirements.txt

Data

Each dataset split is a directory of states.npy, actions.npy, and (for probing splits) locations.npy. The full pretraining trajectories live at data/train/, and the probing splits live at data/probe_{normal,wall,wall_other,expert}/{train,val}/. data/ is gitignored (it's tens of GB) — point the scripts at wherever your copy lives:

export DL_DATA_PATH=/path/to/data   # defaults to ./data if unset

Training

Single device:

python best_model.py

Multi-GPU (DistributedDataParallel):

torchrun --nproc_per_node=<num_gpus> best_model_DDP.py

Both scripts save the trained encoder/predictor checkpoint to encoder_outputs/ and a training-loss plot to plots/ by default; see --help on best_model.py for overriding paths.

Evaluation

Runs the frozen encoder through the probing pipeline on all four evaluation sets and reports location-prediction loss for each:

python main.py --model-path encoder_outputs/<checkpoint>.pth

Defaults to DL_DATA_PATH / DL_MODEL_PATH (or ./data and encoder_outputs/trained_recurrent_jepa.pth) if the flags are omitted.

Results

Swept 28 configs across encoder (CNN/ViT), predictor design, data augmentation, and loss formulation, plus additional hyperparameter tuning on the best-performing config — full sweep and per-run metrics in this spreadsheet. The final model beat the standing benchmark and placed in the top 25th percentile of the class-wide end-of-semester competition.

License

MIT — see LICENSE.