File size: 4,001 Bytes
0d47f33 f896212 929aeba f896212 929aeba f896212 0d47f33 f896212 929aeba f896212 929aeba f896212 929aeba f896212 929aeba f896212 929aeba 8b6c2eb 929aeba 8b6c2eb 929aeba 8b6c2eb 929aeba f896212 8b6c2eb f896212 b89d820 f896212 | 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 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | ---
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}
}
```
|