Skip to content

Repository files navigation

D-JEPA

D-JEPA: A Decision-Aligned Latent World Model

Learning decision-relevant structure from predicted futures.

Project website Model repository Decision-supervision dataset Code repository License: Apache-2.0 Paper on arXiv (link forthcoming) Video on YouTube (link forthcoming)

Paper and video links are placeholders; final links are forthcoming.

Overview  ·  Quick Start  ·  Results  ·  Documentation

Overview

D-JEPA aligns predictive geometry with action selection. It learns decision-relevant relations among predicted futures, combines evidence across predictive geometries, and expresses decision structure in latent future representations—all built on pretrained predictive models.

This repository includes scientific modules, checkpoint loading, cached-feature inference and training, tests, and reproducibility tools. Task/module checkpoints are distributed through the model repository; decision supervision, candidate inputs, fixed identities and result authorities are packaged in one supervision ZIP.

The offline robotics package is independently installed and maintained under real_robot/: recorded-data processing, V-JEPA 2-AC prediction, relational training, candidate scoring and offline planning. Simulation and robotics do not share environment-specific scripts or dependencies. Device-specific robot controllers and calibration are supplied by the user.

Task interfaces are organized by application: robotic manipulation, autonomous driving, and the separate physical-robot package. Each guide identifies data inputs, training/calibration commands and evaluation outputs.

Method at a glance

Module Role
Relational alignment Learn set-wise, bounded corrections from future–goal descriptors and ordinal coordinates.
Predictive plasticity Adapt the final predictor with preservation losses.
Multi-geometry alignment Combine two dense descriptors and four ordinal geometries.
Exact representation realization Encode terminal ordering in goal-distance geometry while preserving earlier futures.
Temporal transport Learn bounded, same-action updates to five-step future representations.
Ordinal / spatial adapters Support task-local inputs and sparse supervision.
Module composition and checkpoint dependencies

These are configurations of D-JEPA, not separate competing methods. Legacy identifiers in raw result authorities and provenance retain the identities of the original experiments.

The predictor boundary defines restricted adaptation; the transport objective accompanies the temporal module. See checkpoint profiles for full-model dependencies and loading requirements.

Quick Start

Python 3.10+ PyTorch 2.3+

The lightweight package supports CPU execution. From the repository root:

pip install -r requirements.txt

For verified Hub downloads, install pip install -e '.[download]'. See the reproduction guide for all configurations and populations.

Replay a released checkpoint

Download and verify the pusht-relational/ profile plus supervision:

python scripts/download_artifacts.py --profile pusht-relational --dataset
python -m zipfile -e data/D-JEPA-supervision-v1.zip data
python scripts/doctor.py --checkpoint checkpoints/pusht-relational
python scripts/verify_dataset.py --data-root data/D-JEPA-supervision-v1 --checksums
bash scripts/reproduce_paper.sh --only pusht-independent

This replays the relational checkpoint on the independent PushT population. Use --resume to verify and skip a completed, unchanged run. The release matrix also includes Granular formal decisions and the two task-local calibration replays, with distinct run identities. See the project website for reported results.

Evaluation protocol and file layout

The checkpoint argument is the downloaded profile directory, containing model.pt and config.json, not its .pt file alone. Choices are computed from frozen inputs first; labels are opened only afterward for evaluation. Output files are never silently overwritten. This replays cached candidate decisions, not a new simulator run.

See protocols and schema for candidate identities, split definitions and task-specific label semantics.

Train a task-local alignment module

bash scripts/train.sh --config configs/training/pushobj.yaml --epochs 120

YAML configures the actual run; CLI options take precedence. Each run records its resolved settings. Paths are relative to the working directory. Existing djepa-train / djepa-evaluate commands and Python import paths remain supported.

Training objective and reproducibility

This portable CPU trainer uses the released three-coordinate ordinal features, full-set success-mass objective, local ordering, preservation and calibration gate. It rejects overlapping train/calibration base identities. The historical paper checkpoints are released directly; this command is not a promise of bitwise-identical retraining across software/hardware versions.

For relational alignment and sparse spatial/multiview training, use the module-training recipes. The native execution guide covers feature preparation from upstream predictors and fixed-horizon physics rollouts, separately from cached replay. Run settings, inputs and output artifacts are recorded locally.

Results

Visit the project website for numerical comparisons and qualitative demonstrations. Machine-readable reference outcomes are distributed in the Hugging Face supervision archive, not as a separate results directory in this code repository.

Documentation

Guide Contents
Reproduction guide Installation, verified downloads, YAML runs and summary scripts
Module training Relational and sparse-supervision training recipes
Native execution Raw-input feature extraction, model factories and physics rollouts
Offline robotics Separate installation, recorded observations, relational training and offline planning
Robotic manipulation RoboTwin data, scene-conditioned candidates, preservation gates and matched execution
Autonomous driving Drive-JEPA features, relation/risk training, label-free selection and evaluation
Release scope Available workflows and future paper, video and experiment updates
Protocols and data schema Evaluation populations, identities and supervision semantics
Checkpoint profiles Loading, composition and upstream dependencies
Validation Tested release scope and cached-decision replay
Source provenance Source/export hashes for extracted modules
Website maintenance Local preview, theme and author-artwork slots

Repository map

src/djepa/      models · objectives · data · evaluation · native · robotics · driving · cli
real_robot/    independent offline robotics package · scripts · configs · tests
configs/       training · evaluation · reproduction · native · robotics · driving
scripts/       download · preflight · train · evaluate · reproduce · summarize · preview
examples/      label-free inference and sparse metrics
tests/         scientific behavior and workflow tests
docs/          project website and technical documentation
assets/        original brand assets

Legacy top-level Python modules forward to the organized implementations; scientific source hashes are mapped in PACKAGE_LAYOUT.json.

Run tests

bash scripts/smoke_test.sh

License

D-JEPA code, including the independent offline robotics package, is released under the Apache License 2.0. Third-party components retain their original licenses; see third-party notices. Model checkpoints and datasets are governed by the terms provided with their respective releases.

Acknowledgments

We thank the LeWorldModel (LeWM) and JEPA research teams for their foundational work and open-source resources.

Additional implementation dependencies, recorded revisions and release terms are documented in third-party notices.

About

D-JEPA: A Decision-Aligned Latent World Model

Topics

Resources

Stars

1 star

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages