AnyGoal model checkpoints

Paper coming soon.

Model checkpoints for AnyGoal, a compositional generalization suite for world-model research. The repository provides checkpoints used by AnyGoal's rollout and planning integrations while preserving each model's native checkpoint format.

Quick start

Install AnyGoal with evaluation dependencies, then download every checkpoint into the directory used by the packaged configs:

pip install 'anygoal[eval]'

hf download HazelNam/anygoal --repo-type model \
  --local-dir "$(anygoal config path checkpoints)"

To download one checkpoint directly, keep its filename unchanged:

hf download HazelNam/anygoal prejepa_pusht_object.ckpt \
  --repo-type model \
  --local-dir "$(anygoal config path checkpoints)"

AnyGoal can download the published artifacts required for an evaluation, download all checkpoints registered by the installed package, or fetch missing artifacts immediately before evaluation:

anygoal artifacts download --model driftworld --dataset cubehop
anygoal artifacts download --checkpoints --all

anygoal eval planning \
  --model driftworld \
  --dataset cubehop \
  --download

Evaluation datasets use the checkpoint trained on their base family. For example, cubehop and cubediagonal use the Cube checkpoint, while avoidt, stackt, and pusht-antigoal use the PushT checkpoint. Existing local files are reused.

Supported model integrations

Model ID	Model	Checkpoint families in this repository	Format
`lewm`	LeWM	Insert Peg	`.pt`
`sensorimotor`	Sensorimotor World Model	PushT, Cube, TwoRoom, Insert Peg	`.ckpt`
`prejepa`	PreJEPA / DINO-WM	PushT, Cube, TwoRoom, Insert Peg	`.ckpt` / `.pt`
`driftworld`	DriftWorld	PushT, Cube, TwoRoom, Insert Peg	`.pt`
`irasim`	IRASim	PushT, Cube, TwoRoom, Insert Peg	`.pt`
`toy-adaln`	AdaLN predictor	PushT, Cube, TwoRoom, Insert Peg	`.pt`
`toy-action-duplicate`	Action-duplicate predictor	PushT, Cube, TwoRoom, Insert Peg	`.pt`
`toy-driftworld`	Toy DriftWorld predictor	PushT, Cube, TwoRoom, Insert Peg	`.pt`

AnyGoal also accepts dino-wm and dinowm as aliases for prejepa, and sensorimotor_wm as an alias for sensorimotor.

PreJEPA / DINO-WM checkpoints

The following base-task checkpoints are the PreJEPA checkpoints used by Stable World-Model:

Family	Filename
PushT	`prejepa_pusht_object.ckpt`
Cube	`prejepa_cube_object.ckpt`
TwoRoom	`prejepa_tworoom_object.ckpt`
Insert Peg	`prejepa_insertpeg.pt`

PreJEPA is AnyGoal's integration name for the DINO-WM-style baseline. The PushT, Cube, and TwoRoom checkpoints are redistributed here with permission from the original author. Their filenames are intentionally preserved because released AnyGoal configurations resolve these exact names.

Official LeWM checkpoints

The official PushT, Cube, and TwoRoom LeWM checkpoints are maintained by the LeWM authors and are not duplicated here. Download them from the official LeWM Hugging Face collection:

This repository provides the AnyGoal-trained Insert Peg checkpoint as lewm_insertpeg.pt.

Use a checkpoint

Run a dry run first to inspect the resolved checkpoint, dataset, configuration, and command:

anygoal eval rollout --model prejepa --dataset avoidt --seed 0 --dry-run
anygoal eval planning --model prejepa --dataset avoidt --seed 0 --dry-run

Remove --dry-run to execute the evaluation. A checkpoint can also be supplied explicitly:

anygoal eval rollout \
  --model prejepa \
  --dataset avoidt \
  --checkpoint "$(anygoal config path checkpoints)/prejepa_pusht_object.ckpt"

From Python:

import anygoal

result = anygoal.evaluate_rollout(
    "prejepa",
    data="/path/to/avoidt_test.h5",
    checkpoint="/path/to/prejepa_pusht_object.ckpt",
    seed=[0, 1, 2],
)

print(result[0].five_step_mse)

See the AnyGoal documentation for the complete model matrix, model-source installation, dataset catalog, training, rollout, and planning guides.

Metadata and attribution

AnyGoal-produced checkpoints are accompanied by same-named JSON sidecars containing provenance and recoverable training metadata. Third-party model architectures and imported checkpoints remain subject to their original attribution and licensing requirements.