Improving Transformers World Model - Pytorch (wip)
Implementation of the new SOTA for model based RL, from the paper Improving Transformer World Models for Data-Efficient RL, in Pytorch.
They significantly outperformed DreamerV3 (as well as human experts) with a transformer world model and a less complicated setup, on Craftax (simplified Minecraft environment)
Install
$ pip install improving-transformers-world-model
Usage
import torch
from improving_transformers_world_model import (
WorldModel
)
world_model = WorldModel(
image_size = 63,
patch_size = 7,
channels = 3,
transformer = dict(
dim = 512,
depth = 4,
block_size = 81
),
tokenizer = dict(
dim = 7 * 7 * 3,
distance_threshold = 0.5
)
)
state = torch.randn(2, 3, 20, 63, 63) # batch, channels, time, height, width - craftax is 3 channels 63x63, and they used rollout of 20 frames. block size is presumably each image
loss = world_model(state)
loss.backward()
# dream up a trajectory to be mixed with real for training PPO
prompts = state[:, :, :2] # prompt frames
imagined_trajectories = world_model.sample(prompts, time_steps = 20)
assert imagined_trajectories.shape == state.shape
Citations
@inproceedings{Dedieu2025ImprovingTW,
title = {Improving Transformer World Models for Data-Efficient RL},
author = {Antoine Dedieu and Joseph Ortiz and Xinghua Lou and Carter Wendelken and Wolfgang Lehrach and J. Swaroop Guntupalli and Miguel L{\'a}zaro-Gredilla and Kevin Patrick Murphy},
year = {2025},
url = {https://api.semanticscholar.org/CorpusID:276107865}
}
1 maintainer
Metadata
1 maintainer
Release files for improving-transformers-world-model 0.0.59
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Source distribution (sdist)
| File | Size | Uploaded | |
|---|---|---|---|
| improving_transformers_world_model-0.0.59.tar.gz | 607.5 kB | Details |
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| improving_transformers_world_model-0.0.59-py3-none-any.whl | Python 3 | none | any | Details |
Total release size: 627.4 kB
Release files / improving_transformers_world_model-0.0.59.tar.gz
| Download URL | improving_transformers_world_model-0.0.59.tar.gz |
|---|---|
| Size | 607.5 kB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
402e09533e5cdab814f26ce3e0be59ce599b4093e0f9b53af63f251a1fe01521
|
|
BLAKE2b-256 checksum How to use checksums |
4ab79dcd51a6117a42b9bec068cb7f869e5da37a3327f907e8604ccf8469a05c
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/6.1.0 CPython/3.9.22
|
Release files / improving_transformers_world_model-0.0.59-py3-none-any.whl
| Download URL | improving_transformers_world_model-0.0.59-py3-none-any.whl |
|---|---|
| Size | 19.9 kB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
d53c61bd22fd8042cfdf092ef0637ab8a9a1a366a1aa1f381d0cec287816acfa
|
|
BLAKE2b-256 checksum How to use checksums |
0b72a8ca34002ef4a7b50eccab2ba644fc66e725776c4597e12722ee655d8853
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/6.1.0 CPython/3.9.22
|
