Add model card for Sub-JEPA

README.md

@@ -0,0 +1,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}
+}
+```