README.md

---
license: mit
language:
- en
tags:
- world-model
- robotics
- pusht
- jax
- latent-diffusion
datasets:
- smbml/future-sight-pusht
inference: false
---

<p align="center"><img src="assets/wordmark.svg" alt="Future Sight" width="420"></p>

# PushT world models

Action-conditioned world models from S. Bateman et al., **Training Controllable
World Models via Action Influence Maximization**, CoRL 2026.

PushT provides a controlled wind tunnel for studying how data scale and
composition affect world-model prediction and generalization. These models
compare large passive collections with smaller, task-specific active exploration
collections using a shared architecture. The paired
[Future Sight PushT dataset](https://huggingface.co/datasets/smbml/future-sight-pusht)
provides the training and evaluation data.

<video controls autoplay loop muted playsinline width="768" src="https://huggingface.co/smbml/future-sight-pusht-models/resolve/main/assets/rollouts.mp4"></video>

**Ground truth (left) and model predictions (right)**, using identical recorded
actions. Thirteen hand-selected examples show **Recovery1 + transferred AIM500**,
**PPO1 + AIM500 (fixed reward estimate)**, and **Recovery100** on BC, BC-failure,
and PushAway evaluation cases. Model names appear in the video; each clip shows
33 predicted transitions at 256 × 256 resolution. These examples illustrate
model behavior; they are not an aggregate benchmark or a ranking of models.

## Models at a glance

**16 checkpoints · 256 × 256 RGB · approximately 580 MB each · 9.28 GB total**

- **Inputs:** RGB observations and continuous 6D end-effector delta actions.
- **Outputs:** predicted RGB frames at 256 × 256 resolution.
- **Format:** native Orbax inference bundles, including the frozen VAE.
- **Weights:** exponential moving average (EMA); no optimizer or training-resume state.

## Choose a model

`PPO1` and `Recovery1` use the first 1% of the corresponding training metadata;
`PPO100` and `Recovery100` use the full corpus. The `500` suffix denotes the
additional collection budget. AIM500 uses 100 trajectories from each of five
collection iterations, with a 2:1 PPO-to-AIM sampling mixture.

| Model | Training step | Download |
| --- | ---: | ---: |
| [PPO1 + AIM500](https://huggingface.co/smbml/future-sight-pusht-models/tree/main/models/ppo1-aim500) | 36,000 | 580.1 MB |
| [PPO1](https://huggingface.co/smbml/future-sight-pusht-models/tree/main/models/ppo1) | 26,000 | 580.0 MB |
| [PPO1 + AIM500 (fixed reward estimate)](https://huggingface.co/smbml/future-sight-pusht-models/tree/main/models/ppo1-aim500-fixed-reward-estimate) | 42,000 | 580.0 MB |
| [PPO1 + AIM500 (MI only)](https://huggingface.co/smbml/future-sight-pusht-models/tree/main/models/ppo1-aim500-mi-only) | 39,000 | 580.0 MB |
| [PPO1 + AIM500 (novelty only)](https://huggingface.co/smbml/future-sight-pusht-models/tree/main/models/ppo1-aim500-novelty-only) | 42,000 | 580.1 MB |
| [PPO1 + DADS500](https://huggingface.co/smbml/future-sight-pusht-models/tree/main/models/ppo1-dads500) | 29,000 | 580.0 MB |
| [PPO1 + DIAYN500](https://huggingface.co/smbml/future-sight-pusht-models/tree/main/models/ppo1-diayn500) | 34,000 | 580.1 MB |
| [PPO1 + ICM500](https://huggingface.co/smbml/future-sight-pusht-models/tree/main/models/ppo1-icm500) | 29,000 | 579.9 MB |
| [PPO1 + METRA500](https://huggingface.co/smbml/future-sight-pusht-models/tree/main/models/ppo1-metra500) | 26,000 | 580.0 MB |
| [PPO1 + random walk500](https://huggingface.co/smbml/future-sight-pusht-models/tree/main/models/ppo1-random-walk500) | 28,000 | 580.0 MB |
| [PPO1 + Recovery500](https://huggingface.co/smbml/future-sight-pusht-models/tree/main/models/ppo1-recovery500) | 27,000 | 580.1 MB |
| [PPO1 + uniform random500](https://huggingface.co/smbml/future-sight-pusht-models/tree/main/models/ppo1-uniform-random500) | 33,000 | 580.0 MB |
| [PPO100](https://huggingface.co/smbml/future-sight-pusht-models/tree/main/models/ppo100) | 44,000 | 580.1 MB |
| [Recovery1](https://huggingface.co/smbml/future-sight-pusht-models/tree/main/models/recovery1) | 37,000 | 580.0 MB |
| [Recovery1 + transferred AIM500](https://huggingface.co/smbml/future-sight-pusht-models/tree/main/models/recovery1-transferred-aim500) | 42,000 | 580.1 MB |
| [Recovery100](https://huggingface.co/smbml/future-sight-pusht-models/tree/main/models/recovery100) | 40,000 | 580.1 MB |

The **fixed reward estimate** ablation uses AIM data without reward
re-estimation. Novelty-only and MI-only models ablate the reward components.
Transferred AIM adds exploration collected from the PPO-based setting to a
Recovery1 training mixture. See the experiment-specific recipes in Future Sight
for exact subset and mixture definitions; the budget suffix alone does not
specify sampling probabilities.

## Download and load one model

Run from a Future Sight checkout with its Pixi environment installed and one
visible GPU, for example `CUDA_VISIBLE_DEVICES=0 pixi run python example.py`
after saving the following snippet as `example.py`. The example
uses existing `huggingface_hub` and Future Sight APIs. No Transformers model
conversion or `trust_remote_code` is required.

```python
import os
from pathlib import Path

import jax
from huggingface_hub import snapshot_download

from future_sight.models.components.sharding import create_device_mesh
from future_sight.models.loading import load_world_model

cache = Path(os.environ.get("FUTURE_SIGHT_CACHE_DIR", "cache")).expanduser()
snapshot = Path(snapshot_download(
    repo_id="smbml/future-sight-pusht-models",
    revision="main",  # Use a full HF commit hash to pin an experiment.
    allow_patterns=["models.json", "models/ppo1-aim500/*"],
    cache_dir=cache / "models",
))

jax.devices()  # Initialize JAX before any TensorFlow-backed data loader.
mesh = create_device_mesh(data=1, model=1)
model, manifest = load_world_model(str(snapshot / "models/ppo1-aim500"), mesh)
print(manifest["weights"]["source_training_step"])
```

The example explicitly honors `FUTURE_SIGHT_CACHE_DIR`; the HF library itself
does not interpret that Future Sight variable. Change the folder in
`allow_patterns` and the load path to select another model. HF uses your existing
login when available; these public files are also downloadable anonymously.
Completed files are reused from the HF cache, and Orbax reads them in place.

The restored model can be used with Future Sight's Python inference session,
standalone PushT evaluator, or local interactive application. Model-specific HF
CLI integration and a hosted interactive Space are not provided by this repo.

## Architecture and provenance

All models use the latent diffusion forcing B/2 architecture: 12 transformer
layers, width 768, 12 attention heads, and 2 × 2 latent patches. The VAE maps
256 × 256 images to 32 × 32 × 4 latents. The published manifests specify the
exact action bounds, x-prediction / v-space training objective, BF16 runtime,
128 inference steps, four-frame sequence window and three cached predecessors.

Each bundle includes a versioned manifest with its training step, source commit,
and hashes of every restored tensor. [models.json](models.json) additionally
records downloadable file sizes and SHA-256 digests. Public metadata removes
private storage paths and uses the `future_sight` import namespace; checkpoint
payloads and tensor identities are unchanged. Historical source repository names
remain in provenance. Private training-contract files are not distributed.

Our base VAE is
[**`stabilityai/sd-vae-ft-mse`**](https://huggingface.co/stabilityai/sd-vae-ft-mse),
used through [pcuenq's Flax conversion](https://huggingface.co/pcuenq/sd-vae-ft-mse-flax)
and included as a frozen StableVAE in every inference bundle. Both upstream
model cards specify MIT.
The original upstream HF revision was not recorded; the bundle records the
converted VAE's content identity. The trained Future Sight weights are also
released under the [MIT license](LICENSE). See [upstream attribution](NOTICE.md)
for the bundled VAE.

## Validation and limitations

All 16 source checkpoints passed strict restoration and finite-output checks on two selected evaluation cases each. They also ran in the local interactive model switcher. Published model files are listed above. The PPO1 + AIM500 bundle additionally passed an authenticated Hugging Face upload/download round trip, file-hash verification, offline cache reuse, and exact raw-frame agreement with its local source on BC and PushAway when independently restored states use the same compiled inference graph. Separately compiled GPU runs were not pixel-identical; this check does not establish bitwise reproducibility across runs. Other models have local inference checks and remote file verification, not separate HF round-trip inference runs.

The local checks used JAX/JAXLIB 0.10.2, Flax 0.12.7 and Orbax 0.12.1 on an
NVIDIA RTX PRO 6000 Blackwell GPU. They do not establish compatibility with every
hardware or dependency combination. Small local checks and videos do not reproduce
the complete paper evaluation. Samples may drift or produce implausible dynamics,
particularly over longer horizons or unfamiliar actions.

## Citation

If you use these models or the paired dataset, please cite:

```bibtex
@inproceedings{bateman2026aim,
  title = {Training Controllable World Models via Action Influence Maximization},
  author = {Bateman, Samuel M. and Yin, Tenny and Zheng, Chongyi and Huang, Lei
            and Wang, Brian and Eysenbach, Ben and Fisac, Jaime Fern\'andez and Shah, Dhruv},
  booktitle = {Conference on Robot Learning},
  year = {2026}
}
```

Paper link forthcoming.