SAGE

Official checkpoints for SAGE: Subgoal-Conditioned Action Generation for Latent World Model Planning.

Paper | Code | Datasets

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:

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

Action normalization statistics accompany the native bundles: action_stats.json for LIBERO and robotwin_a2b/action_stats.npz for RoboTwin A2B. The RoboTwin file contains the paper's 14-dimensional mean and std arrays and count=1188823. Use the released statistics without recomputing them; the GitHub model installer includes and verifies this file.

Usage

Download the checkpoints with the Hugging Face Hub:

from huggingface_hub import snapshot_download

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

See the code repository for installation and evaluation commands, and native benchmark instructions for LIBERO and RoboTwin.

LIBERO Quick Start

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

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:

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:

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, then set LIBERO_ROOT to that checkout. Run the released full-episode queries:

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.

Citation

@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}
}