Zunzhi/ngc-checkpoints

NGC: LeWM checkpoints

Checkpoints accompanying Latent World Models Should Plan with Distances Along the Latent Manifold, by Zunzhi You, Linwei Tao, Siyu Xu, Yunke Wang, and Chang Xu.

NGC scores candidate plans with distances along a graph of frozen world-model latents. These are LeWM world models trained for the experiments, not models trained with an additional NGC loss. Planning code, evaluation protocol, and fixed pair sets are in ngc-latent-manifold. The code repository may remain private while the code release is prepared.

Contents

All models use training seed 3072 (main results), 3073, or 3074 (training-seed study), epoch 10 of a 10-epoch run, and a 90% training split. Each weight file has its matching config.json in the same directory.

Use	Task	Directory
Main, seed 3072	Cube	checkpoints/project/cube/
Main, seed 3072	PushT	checkpoints/project/pusht/
Main, seed 3072	Reacher	checkpoints/project/reacher/
Main, seed 3072	TwoRoom	checkpoints/training/checkpoints/tworoom-seed3072/
Replicates, seeds 3073 and 3074	Each of the four tasks	checkpoints/training/checkpoints/<task>-seed<seed>/

The TwoRoom main checkpoint is the retrained 10-epoch model used in the paper. The older TwoRoom checkpoint is not included. Official LeWM checkpoints and all datasets remain upstream dependencies and are not redistributed here.

Download and use

From the code repository root, download the checkpoint directories only:

hf download Zunzhi/ngc-checkpoints --include "checkpoints/**" --local-dir .

For reproducible runs, pass --revision <commit SHA> to pin the download. The directory layout matches experiments/example.yaml in the code release. To load one model with the pinned stable-worldmodel dependency:

import stable_worldmodel as swm
model = swm.wm.utils.load_pretrained(
    'checkpoints/project/cube/weights_epoch_10.pt'
).cuda().eval()

Use scripts/evaluate.py in the code release to apply the paper's image preprocessing, action normalization, rest-start protocol, and executed-action history. Data must preserve the original episode and row ordering for the fixed pair sets. Consult the code's reproduction guide for official LeWM dataset downloads. This repository contains no datasets.

Verification and scope

checkpoint-inventory.json gives the size, SHA-256, task, and seed of each file. CHECKPOINT_SHA256SUMS verifies the relative checkpoint layout:

sha256sum -c CHECKPOINT_SHA256SUMS
# macOS: shasum -a 256 -c CHECKPOINT_SHA256SUMS

The 24 files (12 weights and 12 model configurations) total about 0.81 GiB. Two full reads produced identical hashes before upload. Their source paths were matched to archived evaluation commands, and recorded training settings were checked. Current hashes establish artifact identity, not independently its byte identity at the historical evaluation time. GPU kernels and architecture can affect latent rounding and planning trajectories.

These models are intended for the four simulated benchmark tasks. Performance and robustness outside those tasks have not been established. This model release alone does not validate a fresh end-to-end reproduction of all results.

License and citation

The released checkpoints are provided under MIT; see LICENSE. Third-party code and data retain their own terms. See CITATION.cff for the paper authors and title. A paper URL will be added when available.