|
|
from typing import Any, Optional, Tuple |
|
|
|
|
|
from einops import rearrange |
|
|
import torch |
|
|
import torch.nn as nn |
|
|
import torch.nn.functional as F |
|
|
|
|
|
from .models.kv_caching import KeysValues |
|
|
from .models.slicer import Embedder, Head |
|
|
from .models.transformer import Transformer |
|
|
|
|
|
class WorldModel(nn.Module): |
|
|
def __init__(self, config: dict) -> None: |
|
|
super().__init__() |
|
|
self.obs_vocab_size, self.act_vocab_size = config["vocab_size"], config["act_vocab_size"] |
|
|
self.config = config |
|
|
self.transformer = Transformer(config) |
|
|
|
|
|
all_but_last_obs_tokens_pattern = torch.ones(config["tokens_per_block"]) |
|
|
all_but_last_obs_tokens_pattern[-2] = 0 |
|
|
act_tokens_pattern = torch.zeros(self.config["tokens_per_block"]) |
|
|
act_tokens_pattern[-1] = 1 |
|
|
obs_tokens_pattern = 1 - act_tokens_pattern |
|
|
|
|
|
self.pos_emb = nn.Embedding(config["max_tokens"], config["embed_dim"]) |
|
|
|
|
|
self.embedder = Embedder( |
|
|
max_blocks=config["max_blocks"], |
|
|
block_masks=[act_tokens_pattern, obs_tokens_pattern], |
|
|
embedding_tables=nn.ModuleList([nn.Embedding(self.act_vocab_size, config["embed_dim"]), nn.Embedding(self.obs_vocab_size, config["embed_dim"])]) |
|
|
) |
|
|
|
|
|
self.head_observations = Head( |
|
|
max_blocks=config["max_blocks"], |
|
|
block_mask=all_but_last_obs_tokens_pattern, |
|
|
head_module=nn.Sequential( |
|
|
nn.Linear(config["embed_dim"], config["embed_dim"]), |
|
|
nn.ReLU(), |
|
|
nn.Linear(config["embed_dim"], self.obs_vocab_size) |
|
|
) |
|
|
) |
|
|
|
|
|
self.head_rewards = Head( |
|
|
max_blocks=config["max_blocks"], |
|
|
block_mask=act_tokens_pattern, |
|
|
head_module=nn.Sequential( |
|
|
nn.Linear(config["embed_dim"], config["embed_dim"]), |
|
|
nn.ReLU(), |
|
|
nn.Linear(config["embed_dim"], 3) |
|
|
) |
|
|
) |
|
|
|
|
|
self.head_ends = Head( |
|
|
max_blocks=config["max_blocks"], |
|
|
block_mask=act_tokens_pattern, |
|
|
head_module=nn.Sequential( |
|
|
nn.Linear(config["embed_dim"], config["embed_dim"]), |
|
|
nn.ReLU(), |
|
|
nn.Linear(config["embed_dim"], 2) |
|
|
) |
|
|
) |
|
|
|
|
|
def __repr__(self) -> str: |
|
|
return "world_model" |
|
|
|
|
|
def forward(self, tokens: torch.LongTensor, past_keys_values: Optional[KeysValues] = None) -> dict: |
|
|
|
|
|
num_steps = tokens.size(1) |
|
|
assert num_steps <= self.config["max_tokens"] |
|
|
prev_steps = 0 if past_keys_values is None else past_keys_values.size |
|
|
|
|
|
sequences = self.embedder(tokens, num_steps, prev_steps) + self.pos_emb(prev_steps + torch.arange(num_steps, device=tokens.device)) |
|
|
|
|
|
x = self.transformer(sequences, past_keys_values) |
|
|
|
|
|
logits_observations = self.head_observations(x, num_steps=num_steps, prev_steps=prev_steps) |
|
|
logits_rewards = self.head_rewards(x, num_steps=num_steps, prev_steps=prev_steps) |
|
|
logits_ends = self.head_ends(x, num_steps=num_steps, prev_steps=prev_steps) |
|
|
return { |
|
|
"output_sequence": x, |
|
|
"logits_observations": logits_observations, |
|
|
"logits_rewards": logits_rewards, |
|
|
"logits_ends": logits_ends |
|
|
|
|
|
} |
|
|
|
|
|
def generate_empty_keys_values(self, n= 1): |
|
|
|
|
|
values = self.transformer.generate_empty_keys_values(n=n, max_tokens= self.config["max_tokens"]) |
|
|
return values |
|
|
|
|
|
def compute_labels_world_model(self, obs_tokens: torch.Tensor, rewards: torch.Tensor, ends: torch.Tensor, mask_padding: torch.BoolTensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
|
|
assert torch.all(ends.sum(dim=1) <= 1) |
|
|
mask_fill = torch.logical_not(mask_padding) |
|
|
labels_observations = rearrange(obs_tokens.masked_fill(mask_fill.unsqueeze(-1).expand_as(obs_tokens), -100), 'b t k -> b (t k)')[:, 1:] |
|
|
labels_rewards = (rewards.sign() + 1).masked_fill(mask_fill, -100).long() |
|
|
labels_ends = ends.masked_fill(mask_fill, -100) |
|
|
return labels_observations.reshape(-1), labels_rewards.reshape(-1), labels_ends.reshape(-1) |
|
|
|
