README.md

---
pipeline_tag: reinforcement-learning
---

# Sub-JEPA: Subspace Gaussian Regularization for Stable End-to-End World Models

Official implementation of the paper: [Sub-JEPA: Subspace Gaussian Regularization for Stable End-to-End World Models](https://arxiv.org/abs/2605.09241).

## Overview

Joint-Embedding Predictive Architectures (JEPAs) provide a simple framework for learning world models by predicting future latent states. However, JEPA training can be subject to collapse without sufficient structural constraints. **Sub-JEPA** relaxes global constraints used in previous methods (like LeWM) by applying Gaussian regularization across multiple random subspaces rather than the original high-dimensional embedding space. This leads to a better balance between training stability and representation quality in continuous-control environments.

## Resources

- **GitHub:** [intcomp/Sub-JEPA](https://github.com/intcomp/Sub-JEPA)
- **Paper:** [arXiv:2605.09241](https://arxiv.org/abs/2605.09241)

## Installation

To set up the environment, clone the repository and apply the Sub-JEPA patch to the underlying LeWM codebase:

```bash
git clone --recursive https://github.com/intcomp/Sub-JEPA.git
cd Sub-JEPA

# Apply the Sub-JEPA patch to LeWM
git -C le-wm apply ../lewm_subjepa.patch
```

Please refer to the [official repository](https://github.com/intcomp/Sub-JEPA) for additional environment and data setup instructions.

## Usage

### Training

Training is configured with Hydra. To train on the `tworoom` environment:

```bash
PYTHONPATH=. python le-wm/train.py data=tworoom
```

### Evaluation

Evaluation configurations are located under `le-wm/config/eval/`:

```bash
python le-wm/eval.py --config-name=tworoom.yaml policy=tworoom/subjepa
```

## Citation

```bibtex
@misc{zhao2026subjepa,
  title        = {Sub-JEPA: Subspace Gaussian Regularization for Stable End-to-End World Models},
  author       = {Zhao, Kai and Nie, Dongliang and Lin, Yuchen and Luo, Zhehan and Gu, Yixiao and Fan, Deng-Ping and Zeng, Dan},
  year         = {2026},
  eprint       = {2605.09241},
  archivePrefix = {arXiv},
  primaryClass = {cs.LG},
  url          = {https://arxiv.org/abs/2605.09241}
}
```