V-JEPA for Gray-Scott dynamics

Video Joint Embedding Predictive Architecture (V-JEPA) for spatiotemporal PDE dynamics — the Gray-Scott reaction-diffusion system.

This repository is a focused extraction of the work our team JEPAdormi produced during the 24-hour Hack the World(s) hackathon (we won first prize! 🏆). Rather than ship our full fork, we kept only the Gray-Scott track we worked on, plus the minimal upstream library code required for it to run. It also serves as a starting point for further work on V-JEPA for PDE prediction.

Presentation

Our hackathon presentation is included as HTW_JEPAdormi.pdf. View the live version here to see the GIFs in motion.

How it works

We train a temporal JEPA to predict the latent dynamics of the Gray-Scott field, then decode latents back to (u, v) frames and score multi-step rollouts with VRMSE against persistence and against The Well's neural baselines (ResUNet, FNO/TFNO, U-Net, ConvNeXt-U-Net). See gray_scott/README.md for the model design, the data, and the full script suite, and gray_scott/DESIGN.md for architecture notes.

No pretrained checkpoints are included. During the hackathon the models were trained on a GPU cluster and we did not export the weights, so there are no .pth.tar checkpoints in this repo — the eval/visualization scripts need a checkpoint you produce yourself. Anyone is welcome to run the training scripts on their own GPU(s); see Running below.

Setup

# Python 3.12 (see .python-version)
pip install -e .               # core: everything needed to train + run the basic VRMSE eval
pip install -e ".[baselines]"  # optional: The Well (neural baselines + the dataset-download CLI)

Data. The Gray-Scott trajectories come from The Well. Download them with The Well's own CLI, then point the loader at the result (no code edit needed). Our loader reads the on-disk HDF5 files directly — layout <ROOT>/data/{train,valid,test}/*.hdf5:

pip install the_well
the-well-download --base-path data --dataset gray_scott_reaction_diffusion
export GRAY_SCOTT_DATA_ROOT=data/datasets/gray_scott_reaction_diffusion

GRAY_SCOTT_DATA_ROOT is the explicit override; otherwise the loader looks under $EBJEPA_DSETS/the_well/gray_scott_reaction_diffusion (env.sh sets EBJEPA_DSETS), then ./data/the_well/gray_scott_reaction_diffusion. (The Well can also stream from HuggingFace via its WellDataset API, but this repo's loader expects local files.) See the track README's Data section for details.

Note: eb_jepa imports scikit-learn, which must match your installed NumPy. On a NumPy 1.x/2.x ABI error from inside sklearn, pin a matching pair (e.g. numpy<2 with an older sklearn, or an sklearn wheel built for NumPy 2).

Running

# Train
python -m gray_scott.main --fname gray_scott/cfgs/train.yaml

# Evaluate a checkpoint (VRMSE per horizon vs persistence)
python -m gray_scott.eval --ckpt <run>/latest.pth.tar --H 30

# Train + score the neural baselines
python -m gray_scott.baselines --split test --H 30

On a cluster (SLURM)

The train_*.sh, slurm_*.sh and run_*.sh files in scripts/ are the SLURM launchers we used on the cluster. Submit them from the repo root, e.g. sbatch scripts/train_gs_vjepa.sh. See scripts/example_experiment.sh for an annotated, end-to-end template (train → evaluate) explaining the #SBATCH directives and the env.sh setup. The full catalogue of analysis and visualization scripts is documented in gray_scott/README.md.

More

• CONTRIBUTIONS.md — what is ours vs. vendored from upstream, the fork point, and the commit trace.
• gray_scott/README.md — the track in depth: data, model, scripts, and the VRMSE metric variants.

Team

Team JEPAdormi — Hack the World(s) 2026:

• Jules Dupont (@Jules02)
• Adnan Ben Mansour (@Adnan-Ben-Mansour)
• Alexandre Duplessis (@alexandreduplessis)
• Arthur Gilles (@ArthurGilles)

Credits

Built on the hackathon organizer's eb_jepa, itself built on FAIR's eb_jepa.

License

Apache License 2.0 — the same license as the upstream eb_jepa we build on. The vendored eb_jepa/ code remains under its original Apache 2.0 license; our additions are also released under Apache 2.0. Attribution details are in NOTICE.