Upload README.md with huggingface_hub

README.md

@@ -12,7 +12,30 @@ datasets:
 - facebook/jepa-wms
 ---
 
-# JEPA-WMs: Pretrained World Models
+<h1 align="center">
+    <p>🤖 <b>JEPA-WMs Pretrained Models</b></p>
+</h1>
+
+<h2 align="center">
+    <p><i>World models for robot planning with latent imagination 🧠</i></p>
+</h2>
+
+<div align="center" style="line-height: 1;">
+  <a href="https://github.com/facebookresearch/jepa-wms" target="_blank" style="margin: 2px;"><img alt="Github" src="https://img.shields.io/badge/Github-facebookresearch/jepa--wms-black?logo=github" style="display: inline-block; vertical-align: middle;"/></a>
+  <a href="https://huggingface.co/facebook/jepa-wms" target="_blank" style="margin: 2px;"><img alt="HuggingFace" src="https://img.shields.io/badge/🤗%20HuggingFace-facebook/jepa-wms-ffc107" style="display: inline-block; vertical-align: middle;"/></a>
+  <a href="https://arxiv.org/abs/2512.24497" target="_blank" style="margin: 2px;"><img alt="ArXiv" src="https://img.shields.io/badge/arXiv-2512.24497-b5212f?logo=arxiv" style="display: inline-block; vertical-align: middle;"/></a>
+</div>
+
+<br>
+
+<p align="center">
+  <b><a href="https://ai.facebook.com/research/">Meta AI Research, FAIR</a></b>
+</p>
+
+<p align="center">
+  This 🤗 HuggingFace repository hosts pretrained <b>JEPA-WM</b> world models.<br>
+  👉 See the <a href="https://github.com/facebookresearch/jepa-wms">main repository</a> for training code and datasets.
+</p>
 
 This repository contains pretrained world model checkpoints from the paper
 ["What Drives Success in Physical Planning with Joint-Embedding Predictive World Models?"](https://arxiv.org/abs/2512.24497)
@@ -85,10 +108,15 @@ checkpoint_path = hf_hub_download(
     filename="jepa_wm_droid.pth.tar"
 )
 
-# Load with PyTorch
+# Load checkpoint (contains 'encoder', 'predictor', and 'heads' state dicts)
 checkpoint = torch.load(checkpoint_path, map_location="cpu")
+print(checkpoint.keys())  # dict_keys(['encoder', 'predictor', 'heads', 'opt', 'scaler', 'epoch', 'batch_size', 'lr', 'amp'])
 ```
 
+> **Note**: This only downloads the weights. To instantiate the full model with the correct
+> architecture and load the weights, we recommend using PyTorch Hub (see above) or cloning the
+> [jepa-wms repository](https://github.com/facebookresearch/jepa-wms) and using the training/eval scripts.
+
 ## Citation
 
 ```bibtex
@@ -111,4 +139,5 @@ These models are licensed under [CC-BY-NC 4.0](https://creativecommons.org/licen
 
 - 📄 [Paper](https://arxiv.org/abs/2512.24497)
 - 💻 [GitHub Repository](https://github.com/facebookresearch/jepa-wms)
-- 🤗 [Dataset](https://huggingface.co/datasets/facebook/jepa-wms)
+- 🤗 [Datasets](https://huggingface.co/datasets/facebook/jepa-wms)
+- 🤗 [Models](https://huggingface.co/facebook/jepa-wms)