README.md

---
license: mit
library_name: pytorch
pipeline_tag: robotics
tags:
  - robotics
  - world-model
  - planning
  - pusht
  - ogbench
  - libero
  - robotwin
datasets:
  - CLTRAY/SAGE-data
arxiv: 2607.17973
---

# SAGE

Official checkpoints for **[SAGE: Subgoal-Conditioned Action Generation for Latent World Model Planning](https://arxiv.org/abs/2607.17973)**.

[Paper](https://arxiv.org/abs/2607.17973) | [Code](https://github.com/PKU-ML/SAGE) | [Datasets](https://huggingface.co/datasets/CLTRAY/SAGE-data)

SAGE combines latent subgoal generation and goal-conditioned action proposals with a frozen world model for planning.

## Checkpoints

| Environment | Subgoal generator | Action prior | Far-goal action prior |
|:---|:---|:---|:---|
| PushT | `pusht_generator.pt` | `pusht_action_prior.pt` | `pusht_far_action_prior.pt` |
| OGBench Cube | `cube_generator.pt` | `cube_action_prior.pt` | `cube_far_action_prior.pt` |

| Environment | Complete SAGE bundle |
|:---|:---|
| LIBERO Scene2 | `libero_scene2/` |
| LIBERO Caddy | `libero_caddy/` |
| RoboTwin A2B | `robotwin_a2b/` |

Each native bundle contains three matching components:

```text
<suite>/
  world_model/config.json
  world_model/weights.pt
  generator.pt
  prior.pt
```

`world_model/weights.pt` contains the LeWM encoder and dynamics; `generator.pt`
and `prior.pt` contain the subgoal generator and GMM action prior separately.
Keep all components from the same suite together.

## Usage

Download the checkpoints with the Hugging Face Hub:

```python
from huggingface_hub import snapshot_download

snapshot_download(repo_id="CLTRAY/SAGE", local_dir="checkpoints")
```

See the [code repository](https://github.com/PKU-ML/SAGE) for installation and evaluation commands, and [native benchmark instructions](https://github.com/PKU-ML/SAGE/blob/main/NATIVE.md) for LIBERO and RoboTwin.

### LIBERO Quick Start

Clone the code and install its pinned LIBERO environment on Linux:

```bash
git clone https://github.com/PKU-ML/SAGE.git
cd SAGE
conda env create -f runtime/libero/environment.yml
conda activate sage-libero
pip install --no-deps -e .
```

Download a complete Scene2 bundle and its evaluation data:

```python
from huggingface_hub import snapshot_download

snapshot_download("CLTRAY/SAGE", allow_patterns=["libero_scene2/*"],
                  local_dir="checkpoints")
snapshot_download("CLTRAY/SAGE-data", repo_type="dataset",
                  allow_patterns=["evaluation/libero_scene2_evaluation.tar"],
                  local_dir="sage-data")
```

Verify the component loading and install the evaluation archive:

```bash
python -m sage.assets --suite libero_scene2 --out-dir checkpoints --verify-only
python -m sage.native_models checkpoints/libero_scene2
python scripts/install_native_dataset.py --suite libero_scene2 --split evaluation \
  --archive sage-data/evaluation/libero_scene2_evaluation.tar \
  --out datasets/libero_scene2
```

Install the LIBERO checkout and simulator assets specified in the
[native benchmark guide](https://github.com/PKU-ML/SAGE/blob/main/NATIVE.md),
then set `LIBERO_ROOT` to that checkout. Run the released full-episode queries:

```bash
python -m sage.reproduce_native --suite libero_scene2 --methods sage \
  --libero-root "$LIBERO_ROOT" --data-root datasets/libero_scene2 \
  --checkpoints checkpoints --full-episode --out results/libero_scene2_full
```

Replace `libero_scene2` with `libero_caddy` for Caddy. Query manifests and
evaluation settings are included in the code repository. Model-loading checks
do not run the simulator; online evaluation also requires the matching runtime
and assets. For training and validation archives, see the
[dataset download instructions](https://huggingface.co/datasets/CLTRAY/SAGE-data).

## Citation

```bibtex
@article{cheng2026sage,
  title={SAGE: Subgoal-Conditioned Action Generation for Latent World Model Planning},
  author={Cheng, Letian and Zhang, Qi and Wang, Yisen},
  journal={arXiv preprint arXiv:2607.17973},
  year={2026}
}
```