### Install the package

Source: https://github.com/lucidrains/improving-transformers-world-model-for-rl/blob/main/README.md

Install the library via pip.

```bash
$ pip install improving-transformers-world-model
```

--------------------------------

### Initialize and Process with Impala Backbone

Source: https://context7.com/lucidrains/improving-transformers-world-model-for-rl/llms.txt

Initialize the Impala backbone and process single images or video sequences. Supports stateful processing with GRU carry and option to get separate CNN and RNN outputs.

```python
impala = Impala(
    dims=(64, 64, 128),         # CNN channel dimensions
    image_size=63,
    channels=3,
    init_conv_kernel=7,         # Initial convolution kernel size
    dim_rnn=32,                 # GRU input dimension
    dim_rnn_hidden=32           # GRU hidden dimension
)

# Process single image
image = torch.randn(4, 3, 63, 63)  # batch of 4 images
features, gru_hidden = impala(image)
# features: (4, 2080) - concatenated CNN + RNN output
# gru_hidden: GRU hidden state for next step

# Process video sequence
video = torch.randn(4, 3, 10, 63, 63)  # 4 batches, 10 frames
features, gru_hidden = impala(video)
# features: (4, 10, 2080)

# Stateful processing with GRU carry
frame1 = torch.randn(4, 3, 63, 63)
features1, hidden = impala(frame1)

frame2 = torch.randn(4, 3, 63, 63)
features2, hidden = impala(frame2, gru_hidden=hidden.detach())

# Get separate CNN and RNN outputs
cnn_out, rnn_out, hidden = impala(
    image,
    concat_cnn_rnn_outputs=False
)
# cnn_out: (4, 2048), rnn_out: (4, 32)
```

--------------------------------

### Initialize and Use Actor Network

Source: https://context7.com/lucidrains/improving-transformers-world-model-for-rl/llms.txt

Initialize the Actor network for policy output. Supports optional conditioning on world model embeddings via FiLM layers. Use `sample_action=True` to sample actions and get log probabilities.

```python
import torch
from improving_transformers_world_model.agent import Actor

# Initialize actor with world model conditioning
actor = Actor(
    dim_input=2080,             # IMPALA output dimension
    dim=256,
    num_actions=5,
    num_layers=3,
    expansion_factor=2.,
    dim_world_model_embed=128   # Enable FiLM conditioning
)

# Get action logits
state_features = torch.randn(4, 2080)
action_logits = actor(state_features)
# Output: (4, 5) - logits for 5 actions

# Sample action with log probability
world_embed = torch.randn(4, 128)
actions, log_probs = actor(
    state_features,
    world_model_embed=world_embed,
    sample_action=True
)
# actions: (4,) - sampled action indices
# log_probs: (4,) - log probabilities of sampled actions
```

--------------------------------

### Initialize and Use Critic Network

Source: https://context7.com/lucidrains/improving-transformers-world-model-for-rl/llms.txt

Initialize the Critic network for state value estimation. Uses HL-Gauss loss for distributional value prediction when `use_regression=False`. Set `return_value_and_logits=True` to get both expected values and distributional logits.

```python
import torch
from improving_transformers_world_model.agent import Critic

# Initialize critic with HL-Gauss value head
critic = Critic(
    dim_input=2080,             # IMPALA output dimension
    dim=256,
    num_layers=4,
    expansion_factor=2.,
    use_regression=False,       # Use distributional HL-Gauss
    hl_gauss_loss_kwargs=dict(
        min_value=0.,
        max_value=10.,
        num_bins=32,
        sigma=0.5
    )
)

# Get value estimates
state_features = torch.randn(4, 2080)
values = critic(state_features)
# Output: (4,) - scalar value estimates

# Get value with logits for loss computation
values, logits = critic(state_features, return_value_and_logits=True)
# values: (4,) - expected values
# logits: (4, 32) - distributional logits
```

--------------------------------

### Initialize and Train PPO Agent

Source: https://context7.com/lucidrains/improving-transformers-world-model-for-rl/llms.txt

Sets up an agent with IMPALA backbone and actor-critic networks, supporting training on both real and imagined trajectories.

```python
import torch
from improving_transformers_world_model import Agent, WorldModel
from improving_transformers_world_model.mock_env import Env

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Initialize world model
world_model = WorldModel(
    image_size=63,
    patch_size=7,
    channels=3,
    num_actions=5,
    reward_num_bins=10,
    transformer=dict(dim=128, depth=2, block_size=81),
    tokenizer=dict(dim=7 * 7 * 3, distance_threshold=0.5)
).to(device)

# Initialize agent
agent = Agent(
    impala=dict(
        image_size=63,
        channels=3,
        dims=(64, 64, 128),     # IMPALA CNN dimensions
        dim_rnn=32,              # GRU hidden size
    ),
    actor=dict(
        dim=128,
        num_actions=5,
        num_layers=3,
        dim_world_model_embed=128  # Enable FiLM conditioning
    ),
    critic=dict(
        dim=128,
        num_layers=4,
        use_regression=False,    # Use HL-Gauss for value prediction
        hl_gauss_loss_kwargs=dict(
            min_value=0.,
            max_value=5.,
            num_bins=32
        )
    ),
    actor_eps_clip=0.2,         # PPO clip parameter
    actor_beta_s=0.01,          # Entropy bonus weight
    actor_lr=1e-4,
    critic_lr=1e-4
).to(device)

# Collect imagined trajectories (dreaming)
initial_state = torch.randn(3, 63, 63, device=device)
dream_memories = agent(
    world_model,
    initial_state,
    max_steps=20,
    use_world_model_embed=True  # Condition actor on world model
)

# Collect real environment trajectories
env = Env((3, 63, 63))
real_memories = agent.interact_with_env(
    env,
    world_model=world_model,    # Optional: for world model conditioning
    max_steps=50
)

# Train agent on mixed real and imagined experiences
agent.learn(
    [dream_memories, real_memories],
    batch_size=32,
    epochs=4,
    gamma=0.99,
    lam=0.95
)
```

--------------------------------

### Initialize and Train WorldModel

Source: https://context7.com/lucidrains/improving-transformers-world-model-for-rl/llms.txt

Configures the WorldModel architecture and demonstrates a forward pass for computing losses during training.

```python
import torch
from improving_transformers_world_model import WorldModel

# Initialize world model for Craftax-style environment (63x63 RGB images)
world_model = WorldModel(
    image_size=63,
    patch_size=7,               # Patches will be 7x7 pixels
    channels=3,                 # RGB channels
    num_actions=5,              # Number of discrete actions
    reward_min_value=0.,        # Minimum reward value for HL-Gauss
    reward_max_value=1.,        # Maximum reward value for HL-Gauss
    reward_num_bins=10,         # Reward prediction bins
    transformer=dict(
        dim=512,                # Transformer hidden dimension
        depth=4,                # Number of transformer layers
        block_size=81           # 9x9 patches per frame (63/7)^2
    ),
    tokenizer=dict(
        dim=7 * 7 * 3,          # Patch dimension (147 for 7x7 RGB)
        distance_threshold=0.5   # Threshold for new codebook entries
    )
)

# Training forward pass with video sequence
# Input shape: (batch, channels, time, height, width)
state = torch.randn(2, 3, 20, 63, 63)  # 2 batches, 20 frames
rewards = torch.randint(0, 10, (2, 20)).float()
actions = torch.randint(0, 5, (2, 20, 1))
is_terminal = torch.randint(0, 2, (2, 20)).bool()

# Compute loss (state prediction + reward + terminal)
loss = world_model(
    state,
    actions=actions,
    rewards=rewards,
    is_terminal=is_terminal
)
loss.backward()

# Get loss breakdown
loss, (state_loss, reward_loss, terminal_loss) = world_model(
    state,
    actions=actions,
    rewards=rewards,
    is_terminal=is_terminal,
    return_loss_breakdown=True
)
```

--------------------------------

### Initialize and train the WorldModel

Source: https://github.com/lucidrains/improving-transformers-world-model-for-rl/blob/main/README.md

Configure the WorldModel architecture and perform a forward pass with backpropagation.

```python
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
```

--------------------------------

### Initialize Impala Feature Extractor

Source: https://context7.com/lucidrains/improving-transformers-world-model-for-rl/llms.txt

Imports the IMPALA-style CNN feature extractor for use in agent backbones.

```python
import torch
from improving_transformers_world_model import Impala
```

--------------------------------

### Generate Trajectories with WorldModel.sample

Source: https://context7.com/lucidrains/improving-transformers-world-model-for-rl/llms.txt

Uses the trained model to autoregressively generate future visual states, optionally conditioned on actions and rewards.

```python
import torch
from improving_transformers_world_model import WorldModel

world_model = WorldModel(
    image_size=63,
    patch_size=7,
    channels=3,
    num_actions=5,
    reward_num_bins=10,
    transformer=dict(dim=256, depth=2, block_size=81),
    tokenizer=dict(dim=7 * 7 * 3, distance_threshold=0.5)
)

# Create prompt frames (initial context)
prompt = torch.randn(2, 3, 2, 63, 63)  # 2 batches, 2 prompt frames

# Generate imagined trajectories
imagined_trajectories = world_model.sample(
    prompt,
    time_steps=20,          # Total trajectory length (including prompt)
    temperature=1.5,        # Sampling temperature
    filter_kwargs=dict(min_p=0.1)  # Min-p filtering parameters
)

assert imagined_trajectories.shape == (2, 3, 20, 63, 63)

# Sample with rewards and terminal states returned
trajectories, rewards, is_terminals = world_model.sample(
    prompt,
    time_steps=20,
    return_rewards_and_done=True
)

# Sample with action conditioning
actions = torch.randint(0, 5, (2, 20, 1))
rewards_input = torch.zeros(2, 2)

trajectories = world_model.sample(
    prompt,
    time_steps=20,
    actions=actions,
    rewards=rewards_input
)
```

--------------------------------

### World Model Training with KV-Caching (Burn-in)

Source: https://context7.com/lucidrains/improving-transformers-world-model-for-rl/llms.txt

Demonstrates efficient training using KV-caching for burn-in sequences. The `burn_in_cache` is generated without gradients and then used for training on subsequent frames. `detach_cache=True` is crucial for preventing gradients from flowing into the burn-in phase.

```python
import torch
from improving_transformers_world_model import WorldModel

world_model = WorldModel(
    image_size=63,
    patch_size=7,
    channels=3,
    num_actions=5,
    reward_num_bins=10,
    transformer=dict(dim=128, depth=2, block_size=81),
    tokenizer=dict(dim=7 * 7 * 3, distance_threshold=0.5)
)

# Full sequence data
state = torch.randn(2, 3, 40, 63, 63)  # 40 frames
rewards = torch.randint(0, 10, (2, 40)).float()
actions = torch.randint(0, 5, (2, 40, 1))
is_terminal = torch.randint(0, 2, (2, 40)).bool()

# Burn-in: process first 20 frames without gradients
burn_in_state = state[:, :, :20]
burn_in_rewards = rewards[:, :20]
burn_in_actions = actions[:, :20]
burn_in_terminal = is_terminal[:, :20]

with torch.no_grad():
    _, burn_in_cache = world_model(
        burn_in_state,
        actions=burn_in_actions,
        rewards=burn_in_rewards,
        is_terminal=burn_in_terminal,
        return_cache=True,
        return_loss=False,
        detach_cache=True
    )

# Train on remaining frames with burned-in context
train_state = state[:, :, 20:]
train_rewards = rewards[:, 20:]
train_actions = actions[:, 20:]
train_terminal = is_terminal[:, 20:]

loss = world_model(
    train_state,
    actions=train_actions,
    rewards=train_rewards,
    is_terminal=train_terminal,
    cache=burn_in_cache,
    remove_cache_len_from_time=False
)
loss.backward()
```

--------------------------------

### Initialize and Use NearestNeighborTokenizer

Source: https://context7.com/lucidrains/improving-transformers-world-model-for-rl/llms.txt

Manages a codebook of patch vectors using Euclidean distance thresholds. Use train() to update the codebook and eval() for inference.

```python
import torch
from improving_transformers_world_model import NearestNeighborTokenizer

# Initialize tokenizer
tokenizer = NearestNeighborTokenizer(
    dim=147,                    # Patch dimension (7*7*3 for RGB)
    distance_threshold=0.5,     # Euclidean distance threshold
    max_codes=100_000           # Maximum codebook size
)

# Tokenize patches (training mode - will add new codes)
tokenizer.train()
patches = torch.randn(2, 10, 9, 9, 147)  # batch, time, height, width patches
token_ids = tokenizer(patches)
# Output shape: (2, 10, 9, 9) - integer indices

print(f"Codebook size: {tokenizer.num_codes.item()}")

# Inference mode - finds nearest codes without adding new ones
tokenizer.eval()
token_ids = tokenizer(patches)

# Retrieve patch vectors from token IDs
reconstructed_patches = tokenizer.codes_from_indices(token_ids)
# Output shape: (2, 10, 9, 9, 147)

# Check if codebook is full
if tokenizer.is_at_max_codes:
    print("Codebook has reached maximum capacity")
```

--------------------------------

### WorldModel.sample

Source: https://context7.com/lucidrains/improving-transformers-world-model-for-rl/llms.txt

Autoregressively generates future state trajectories from prompt frames using cached key-values.

```APIDOC
## WorldModel.sample

### Description
Generates future trajectories based on initial prompt frames. Supports temperature-based sampling, min-p filtering, and optional action conditioning.

### Parameters
- **prompt** (torch.Tensor) - Required - Initial context frames.
- **time_steps** (int) - Required - Total trajectory length to generate.
- **temperature** (float) - Optional - Sampling temperature.
- **filter_kwargs** (dict) - Optional - Filtering parameters like min_p.
- **actions** (torch.Tensor) - Optional - Action sequence for conditioning.

### Response
- **imagined_trajectories** (torch.Tensor) - Generated visual sequence.
- **rewards** (torch.Tensor) - Predicted rewards if return_rewards_and_done is True.
- **is_terminals** (torch.Tensor) - Predicted terminal states if return_rewards_and_done is True.
```

--------------------------------

### Initialize and Use BlockCausalTransformer

Source: https://context7.com/lucidrains/improving-transformers-world-model-for-rl/llms.txt

Configures a causal transformer for sequence processing and demonstrates autoregressive generation using a KV cache.

```python
transformer = BlockCausalTransformer(
    dim=512,                    # Model dimension
    depth=6,                    # Number of layers
    block_size=81,              # Tokens per block (patches per image)
    dim_head=64,                # Attention head dimension
    heads=8,                    # Number of attention heads
    ff_expand_factor=4.,        # Feedforward expansion
    num_residual_streams=4,     # Hyper-connections streams
    use_flex_attn=False,        # Use PyTorch FlexAttention (CUDA only)
    dropout_attn=0.1,
    dropout_ff=0.1
)

# Process sequence of patch tokens
# 10 frames * 81 patches = 810 tokens
tokens = torch.randn(2, 810, 512)

# Forward pass
embeddings = transformer(tokens)
# Output: (2, 810, 512)

# Forward with KV cache for autoregressive generation
embeddings, cache = transformer(tokens, return_cache=True)

# Continue generation with cache
new_tokens = torch.randn(2, 81, 512)  # One new frame
new_embeddings, updated_cache = transformer(
    new_tokens,
    cache=cache,
    return_cache=True
)
```

--------------------------------

### Initialize BlockCausalTransformer

Source: https://context7.com/lucidrains/improving-transformers-world-model-for-rl/llms.txt

Imports the block-wise causal attention transformer component.

```python
import torch
from improving_transformers_world_model import BlockCausalTransformer
```

--------------------------------

### Cite the paper

Source: https://github.com/lucidrains/improving-transformers-world-model-for-rl/blob/main/README.md

BibTeX citation for the Improving Transformer World Models for Data-Efficient RL paper.

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

--------------------------------

### WorldModel Class

Source: https://context7.com/lucidrains/improving-transformers-world-model-for-rl/llms.txt

The main class for initializing the world model and computing training losses for state, reward, and terminal predictions.

```APIDOC
## WorldModel

### Description
The WorldModel class combines a nearest neighbor tokenizer with a block causal transformer to learn visual dynamics. It supports training via forward passes that compute losses for state prediction, reward prediction (using HL-Gauss), and terminal state prediction.

### Parameters
#### Constructor Parameters
- **image_size** (int) - Required - Size of the input image (e.g., 63 for 63x63).
- **patch_size** (int) - Required - Size of the patches for tokenization.
- **channels** (int) - Required - Number of image channels (e.g., 3 for RGB).
- **num_actions** (int) - Required - Number of discrete actions.
- **reward_num_bins** (int) - Required - Number of bins for HL-Gauss reward prediction.
- **transformer** (dict) - Required - Configuration for the transformer (dim, depth, block_size).
- **tokenizer** (dict) - Required - Configuration for the tokenizer (dim, distance_threshold).

### Request Example
```python
world_model = WorldModel(image_size=63, patch_size=7, channels=3, num_actions=5, reward_num_bins=10, transformer=dict(dim=512, depth=4, block_size=81), tokenizer=dict(dim=147, distance_threshold=0.5))
loss = world_model(state, actions=actions, rewards=rewards, is_terminal=is_terminal)
```
```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.