Improving Transformers World Model for RL

The BlockCausalTransformer is designed for efficient temporal modeling using several advanced techniques:

1. Hyper Connections: Uses get_init_and_expand_reduce_stream_functions to manage multiple residual streams, allowing for complex information flow.
2. Block Causal Attention: Implements attention that is causal at a block level (defined by block_size), which can be optimized using PyTorch's flex_attention if available.
3. Value Residual Learning: In non-first layers, the attention mechanism can incorporate a value_residual to improve training stability and information retention.
4. Rotary Embeddings: Uses RotaryEmbedding for relative positional encoding.
5. SwiGLU FeedForward: Uses the SwiGLU activation variant for the feedforward layers.

The agent follows a hierarchical architecture:

1. Impala: Acts as the visual encoder and temporal memory. It processes images through a CNN and a GRU to produce a latent state representation.
2. Actor: Takes the latent state from Impala and maps it to action logits. If dim_world_model_embed was provided during initialization, it uses a FiLM (Feature-wise Linear Modulation) layer to condition the policy on world model embeddings.
3. Critic: Takes the latent state from Impala and predicts the state value. It uses an HLGaussLayer for value estimation, which supports both regression and histogram-based loss.

The calc_target_and_gae function computes the returns and the GAE advantages used for policy updates. It uses AssocScan for an accelerated implementation of the backward scan required for GAE.

Arguments:

• rewards: Tensor of rewards [... n].
• values: Tensor of values [... n+1] (including the value of the terminal state).
• masks: Boolean tensor indicating if the next state is terminal [... n].
• gamma: Discount factor.
• lam: GAE lambda parameter.
• use_accelerated: Whether to use the accelerated scan (defaults to True if on CUDA).

returns, gae = calc_target_and_gae(
    rewards=rewards_tensor,
    values=values_tensor,
    masks=masks_tensor,
    gamma=0.99,
    lam=0.95
)

The min_p_filter is a sampling technique used to prune low-probability tokens from the logits distribution. It keeps only tokens whose probability is at least min_p times the probability of the most likely token.

This is useful during the WorldModel.sample() process to improve sample quality by removing the 'long tail' of unlikely tokens.

from improving_transformers_world_model.world_model import min_p_filter

# logits shape: [batch, ..., vocab_size]
filtered_logits = min_p_filter(logits, min_p=0.1)

The Agent and Impala classes provide the interface for interacting with the world model in a Reinforcement Learning context. Impala is a specific implementation of an agent designed to work with these transformer-based world models.

from improving_transformers_world_model import Agent, Impala

# Example usage of an Impala agent
agent = Impala(
    # ... configuration parameters
)

The learn method performs a PPO-style update on the agent's parameters using collected memories. It calculates Generalized Advantage Estimation (GAE), computes policy and value losses, and updates the actor and critic optimizers.

Arguments:

• memories: A MemoriesWithNextState object or a list of them.
• lam: GAE lambda parameter (default 0.95).
• gamma: Discount factor (default 0.99).
• batch_size: Size of training batches.
• epochs: Number of training epochs.

# Assuming 'memories_with_next' is a MemoriesWithNextState object
loss_tuple = agent.learn(
    memories=memories_with_next,
    lam=0.95,
    gamma=0.99,
    batch_size=32,
    epochs=4
)

The Agent.forward method allows for 'dreaming' or generating trajectories using a WorldModel instead of a real environment. This is used for training on simulated data.

Arguments:

• world_model: The WorldModel to sample from.
• init_state: The starting frame.
• max_steps: Maximum number of steps to simulate.
• use_world_model_embed: If True, the agent uses world model embeddings to condition the actor.

# Generating trajectories from a world model
memories_with_next = agent.forward(
    world_model=my_world_model,
    init_state=initial_frame,
    max_steps=100,
    use_world_model_embed=True
)

The BlockCausalAttention module supports high-performance attention via PyTorch's flex_attention.

To use it, ensure torch.cuda.is_available() and that your environment supports torch.nn.attention.flex_attention. When use_flex_attn=True is passed to the BlockCausalTransformer, it will attempt to use the optimized kernel with a block_mask created via create_block_causal_mask.

Use interact_with_env to run the agent in a real environment. This method collects transitions (state, action, reward, etc.) and returns a MemoriesWithNextState object containing the episode's data.

If a world_model is provided, the agent will use it to generate embeddings that condition the actor's policy during interaction.

# Interaction with a real environment
memories_with_next = agent.interact_with_env(
    env=my_env,
    world_model=my_world_model
)

The WorldModel class is the primary entrypoint for the world model component. It is used to model environment dynamics, typically involving a tokenizer for observations and a transformer-based architecture for predicting future states or actions.

from improving_transformers_world_model import WorldModel

# Example initialization (exact parameters depend on implementation details)
world_model = WorldModel(
    # ... configuration parameters
)