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 | World model directory |
|:---|:---|
| LIBERO Scene2 | `libero_scene2/world_model/` |
| LIBERO Caddy | `libero_caddy/world_model/` |
| RoboTwin A2B | `robotwin_a2b/world_model/` |

Each world model directory contains `weights.pt` and `config.json`.

## 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.

## 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}
}
```