DreamJEPA

DreamJEPA is a 1-week-style research prototype for offline robot control under limited compute. It combines an ACT-lite behavior cloning policy with a JEPA-style latent world model. The world model predicts future latent representations conditioned on candidate action chunks, then uses a progress head to rerank those candidates.

The default environment is a CPU-friendly 2D PushT-lite simulator. This keeps the full pipeline runnable without external robot simulators while preserving the important contracts: state, action chunk, target future state, progress, and reranking.

What Is Implemented

• Offline demonstration generation.
• ACT-lite / behavior cloning policy.
• State encoder and EMA target encoder.
• Action-conditioned latent predictor.
• Progress head for reranking candidates.
• JEPA loss with stop-gradient target latent.
• Baseline policy evaluation and JEPA-reranked evaluation.
• GIF rollout export.
• Tests for the core model and reranker.

Quickstart

python -m venv .venv
.venv\Scripts\activate
pip install -e .[dev]

Run the full CPU smoke pipeline:

python src/train_policy.py --quick --output experiments/policy.pt
python src/train_jepa.py --quick --output experiments/jepa.pt
python src/eval_policy.py --policy experiments/policy.pt --episodes 20
python src/eval_rerank.py --policy experiments/policy.pt --jepa experiments/jepa.pt --episodes 20 --make-gif
python -m pytest

Results are written to results/tables/. GIFs are written to results/videos/ when --make-gif is used.

Method

The JEPA-style world model uses:

z_t = encoder(state_t)
z_target = target_encoder(state_t+H)
z_hat = predictor(z_t, action_chunk)
progress_hat = progress_head(z_hat)

The loss is:

L = mse(z_hat, stopgrad(z_target)) + alpha * mse(progress_hat, progress_true)

At inference, the policy proposes a base action chunk. DreamJEPA samples noisy candidate chunks, predicts each candidate's future latent/progress, and executes the first action from the highest-scoring chunk.

Honest Limitations

DreamJEPA does not claim to implement V-JEPA or a general-purpose VLA. The default setup uses state observations and templated task framing; language conditioning is represented as a VLA-lite interface, not as language generalization evidence.