Does intuitive physics emerge in a video-JEPA?

Hack the World(s) 24h Hackathon - Track 10: Does intuitive physics emerge?

This README presents our hackathon work. The original EB-JEPA library documentation is further down.

The JEPA's prediction energy spikes the moment physics breaks - across all three violation types.

teleport  ·  velocity reversal  ·  wall pass-through

The project in one sentence

We train a video-JEPA only on physically plausible Moving-MNIST clips, then show it impossible clips (teleport, velocity reversal, pass-through a wall) and measure whether its latent prediction energy (predcost) rises on the impossible ones - a violation-of-expectation signal, like an infant who is surprised by a "magic" event.

Key result: the JEPA detects the impossible with an AUROC of 0.997, whereas models that predict in pixel space fail (SimVP 0.50 = chance, ConvLSTM 0.74). Surprise in latent space beats surprise in pixel space.

What we built

All of the work lives in examples/intuitive_physics/.

Axis	Content	Files
1. Energy probe	Video-JEPA training + clip_energy (per-clip energy) + AUROC per violation type	main.py, eval.py, stimuli.py
2. Ablations / collapse	We show the VICReg anti-collapse terms (std/cov) are necessary: without them the energy becomes meaningless	cfgs/ablation/, collapse_metrics.py, make_ablation_figures.py → figures_ablation/
3. 3-model comparison	EB-JEPA (latent) vs SimVP (pixel CNN) vs ConvLSTM (pixel RNN) on the same VoE task	compare.py, simvp.py, convlstm.py, train_*.py
4. Latent visualization	t-SNE/UMAP maps, decoded position, rollout, latent-vs-pixel to understand what the JEPA encodes	visualize_latent.py → evidence/viz/

How to use the repo

The code runs on the GPU cluster (GB200, aarch64). See Installation. Trainings are always launched via SLURM (launch_sbatch), from the repo root after source env.sh.

source env.sh
source $UV_PROJECT_ENVIRONMENT/bin/activate

1) Train the video-JEPA (the main model)

python -m examples.launch_sbatch --example intuitive_physics --single \
    --fname examples/intuitive_physics/cfgs/train.yaml

2) Measure the energy on plausible/impossible pairs

python -m examples.intuitive_physics.eval \
    --ckpt $EBJEPA_CKPTS/intuitive_physics/<sweep>/exp_seed1/latest.pth.tar

3) Run the ablations (collapse evidence) — each config disables part of the regularization:

# one ablation, e.g. without the VICReg variance term
python -m examples.launch_sbatch --example intuitive_physics --single \
    --fname examples/intuitive_physics/cfgs/ablation/no_std.yaml
# (same with full / no_cov / no_reg / weak_std), then the figures:
python -m examples.intuitive_physics.make_ablation_figures

4) Compare EB-JEPA vs SimVP vs ConvLSTM — trains all 3 models then compares automatically:

bash examples/intuitive_physics/launch_comparison_pipeline.sh
# on SLURM: 3 trainings in parallel + a dependent comparison job
# (pixel models train via --example intuitive_physics_simvp / _convlstm)

5) Visualizations (inference only, no training — all read a trained checkpoint):

# the stimuli themselves: matched plausible/impossible pairs (no checkpoint needed)
python -m examples.intuitive_physics.visualize_stimuli

# latent space: t-SNE/UMAP maps, decoded position, rollout, latent-vs-pixel
sbatch viz_inference.sh <CKPT> evidence/viz          # on SLURM (GPU)
python -m examples.intuitive_physics.visualize_latent --ckpt <CKPT> --out evidence/viz

# energy-gap figures (AUROC, distributions) from a trained run
python -m examples.intuitive_physics.make_figures --ckpt_dir <CKPT_DIR>

# animated histograms of the plausible vs impossible energy gap
python -m examples.intuitive_physics.make_gif_histograms --ckpt_dir <CKPT_DIR>

---

⚡ EB-JEPA

Energy-Based Joint-Embedding Predictive Architectures

Meta AI Research, FAIR

Basile Terver, Randall Balestriero, Megi Dervishi, David Fan, Quentin Garrido, Tushar Nagarajan,
Koustuv Sinha, Wancong Zhang, Mike Rabbat, Yann LeCun, Amir Bar

An open source library and tutorial for learning representations for
prediction and planning using joint embedding predictive architectures.

Each example is (almost) self-contained and training takes up to a few hours on a single GPU card.

---

📚 Examples

Image JEPA

Self-supervised representations from unlabeled images on CIFAR-10, evaluated on classification.

Video JEPA

Predict next image representation in a sequence.

AC Video JEPA

JEPA for world modeling + planning in Two Rooms environment.

Planning Episode	Task Definition
Successful planning episode	From init to goal state

---

🚀 Installation

HTW cluster — quick start (hackathon only)

Skip this section unless you are on the HTW hackathon cluster — the generic install below is all you need locally.

Please follow the setup instructions before starting the project.

---

Local / generic (start here)

We use uv for package management.

# Install dependencies
uv sync
# Option 1: Activate virtual environment
source .venv/bin/activate
python -m examples.image_jepa.main
# Option 2: Run directly with uv
uv run python -m examples.image_jepa.main

If you need conda-specific packages, you can use Conda + uv

# Create conda environment with Python 3.12
conda create -n eb_jepa python=3.12 -y
conda activate eb_jepa
# Install package in editable mode with dev dependencies (pytest, black, isort, autoflake)
uv pip install -e . --group dev

Add these to your ~/.bashrc for persistent configuration.

# Where datasets are stored / looked up
export EBJEPA_DSETS=/path/to/eb_jepa/datasets
# Optional: Directory for checkpoints and logs
export EBJEPA_CKPTS=/path/to/checkpoints

Verify the install with uv run pytest tests/.

🏋️ Training

Quick Start

# Local training
python -m examples.{image_jepa,video_jepa,ac_video_jepa}.main

Our default configs are tuned for H100 GPUs. With older GPUs (e.g., A100, V100), you may need to reduce batch size to fit in memory.

📂 Folder Structure

All experiments use a unified folder structure:

checkpoints/
└── {example_name}/
    ├── dev_2026-01-16_00-10/                 # Single/local runs (dev_ prefix)
    │   └── {exp_name}_seed1/
    │
    ├── sweep_2026-01-16_00-10/         # Auto-named 3-seed sweep
    │   ├── {exp_name}_seed1/
    │   ├── {exp_name}_seed1000/
    │   └── {exp_name}_seed10000/
    │
    └── sweep_my_experiment/            # Custom-named sweep
        └── ...

{exp_name} encodes key hyperparameters to avoid folder collisions, e.g.:

• image_jepa: resnet_vicreg_proj_bs256_ep300_ph2048_po2048_std1.0_cov80.0
• video_jepa: resnet_bs64_lr0.001_std10.0_cov100.0
• ac_video_jepa: impala_cov8_std16_simt12_idm1

🖥️ SLURM Launcher (optional)

Command	Description
--example {name}	Choose: image_jepa, video_jepa, ac_video_jepa, maze, fintime, ltsf, eeg, audio, pointcloud, gray_scott, intuitive_physics, factors_of_variation
--fname {path}	Run the sweep specified in the config at {path}
--single	Launch single job (dev mode)
--sweep {name}	Custom sweep name
--array-parallelism {N}	Limits the maximum number of concurrent jobs to N
--full-sweep	Full hyperparameter sweep from config
--use-wandb-sweep	Enable wandb sweep UI

# 3 seeds with wandb averaging (recommended)
python -m examples.launch_sbatch --example image_jepa --fname examples/image_jepa/cfgs/default.yaml

# Custom sweep name
python -m examples.launch_sbatch --example image_jepa --fname examples/image_jepa/cfgs/default.yaml --sweep my_experiment

# Single job
python -m examples.launch_sbatch --example image_jepa --fname examples/image_jepa/cfgs/default.yaml --single

# Full hyperparameter sweep
python -m examples.launch_sbatch --example image_jepa --fname examples/image_jepa/cfgs/default.yaml --full-sweep

# With wandb sweep UI for hyperparameter analysis
python -m examples.launch_sbatch --example image_jepa --fname examples/image_jepa/cfgs/default.yaml --use-wandb-sweep

Replace image_jepa with ac_video_jepa, video_jepa, or maze for other examples.

Full Sweep Configuration: The --full-sweep flag reads the sweep.param_grid section from the example's YAML config file (e.g., examples/image_jepa/cfgs/default.yaml). Without this flag, only a 3-seed sweep is launched. To customize sweep parameters, edit the sweep section in the config:

# Example: examples/image_jepa/cfgs/default.yaml
sweep:
  param_grid:
    loss.cov_coeff: [0.1, 1.0, 10.0, 100.0]
    loss.std_coeff: [1.0, 10.0]
    meta.seed: [1, 1000, 10000]

Wandb Seed Averaging

Runs with the same hyperparameters but different seeds share the same wandb run name, enabling automatic averaging:

1. Go to wandb web UI → Runs table
2. Click "Group by" → select "Name" → Groups runs with identical hyperparameters (different seeds) together

To filter runs from a specific sweep: 3. Click "Filter" → "Group" → select your sweep name

For detailed wandb sweep analysis (parallel coordinates, hyperparameter importance):

1. Use --use-wandb-sweep flag when launching
2. Go to wandb web UI → left pane → "Sweeps" → click your sweep name

SLURM Configuration: SLURM parameters default to the HTW cluster and are read from EBJEPA_SLURM_* env vars (set by env.sh, which also auto-detects your account/QOS per user). Override per launch with the CLI flags --partition/--account/--cpus-per-task/--time-min/--gpus-per-node, or export the matching EBJEPA_SLURM_* var. The SLURM_DEFAULTS dictionary at the top of examples/launch_sbatch.py holds the fallbacks.

🧪 Running test cases

Libraries added to eb_jepa must have their own test cases. To run the tests:

# With uv sync installation
uv run pytest tests/
# With conda + uv installation (no .venv created)
pytest tests/

👩‍💻 Development

Before contributing, please format your code with the following tools:

# Remove unused imports
autoflake --remove-all-unused-imports -r --in-place .
# Sort imports
python -m isort eb_jepa examples tests
# Format code
python -m black eb_jepa examples tests

📚 Citing EB-JEPA

If you find this repository useful, please consider giving a ⭐ and citing:

@misc{terver2026lightweightlibraryenergybasedjointembedding,
      title={A Lightweight Library for Energy-Based Joint-Embedding Predictive Architectures},
      author={Basile Terver and Randall Balestriero and Megi Dervishi and David Fan and Quentin Garrido and Tushar Nagarajan and Koustuv Sinha and Wancong Zhang and Mike Rabbat and Yann LeCun and Amir Bar},
      year={2026},
      eprint={2602.03604},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2602.03604},
}

📄 License

EB-JEPA is Apache licensed. See LICENSE.