File size: 2,197 Bytes
14e8d4e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | ---
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}
}
``` |