|
|
import math |
|
|
from einops import rearrange |
|
|
import torch |
|
|
import torch.nn as nn |
|
|
|
|
|
from .models.convnet import FrameEncoder, FrameDecoder |
|
|
from .models.quantizer import Quantizer |
|
|
|
|
|
class Tokenizer(nn.Module): |
|
|
def __init__(self, config: dict) -> None: |
|
|
super().__init__() |
|
|
self.config = config |
|
|
|
|
|
self.latent_res = config["image_size"] // 2 ** sum(config["encoder_config"]["down"]) |
|
|
self.tokens_grid_res = int(math.sqrt(config["num_tokens"])) |
|
|
self.token_res = self.latent_res // self.tokens_grid_res |
|
|
|
|
|
self.encoder_act_emb = nn.Embedding(config["num_actions"], config["image_size"] ** 2) |
|
|
self.decoder_act_emb = nn.Embedding(config["num_actions"], config["decoder_act_channels"] * self.latent_res ** 2) |
|
|
|
|
|
self.quantizer = Quantizer( |
|
|
config["codebook_size"], config["codebook_dim"], |
|
|
input_dim=config["encoder_config"]["latent_dim"] * self.token_res ** 2, |
|
|
max_codebook_updates_with_revival=config["max_codebook_updates_with_revival"] |
|
|
) |
|
|
|
|
|
self.encoder = FrameEncoder(config["encoder_config"]) |
|
|
self.decoder = FrameDecoder(config["decoder_config"]) |
|
|
self.frame_cnn = FrameEncoder(config["frame_cnn_config"]) |
|
|
|
|
|
def __repr__(self) -> str: |
|
|
return "tokenizer" |
|
|
|
|
|
def forward(self, x1: torch.FloatTensor, a: torch.LongTensor, x2: torch.FloatTensor) -> dict: |
|
|
z = self.encode(x1, a, x2) |
|
|
z = rearrange(z, 'b t c (h k) (w l) -> b t (h w) (k l c)', h=self.tokens_grid_res, w=self.tokens_grid_res) |
|
|
|
|
|
return self.quantizer(z) |
|
|
|
|
|
def encode(self, x1: torch.FloatTensor, a: torch.LongTensor, x2: torch.FloatTensor) -> torch.FloatTensor: |
|
|
a_emb = rearrange(self.encoder_act_emb(a), 'b t (h w) -> b t 1 h w', h=x1.size(3)) |
|
|
encoder_input = torch.cat((x1, a_emb, x2), dim=2) |
|
|
z = self.encoder(encoder_input) |
|
|
|
|
|
return z |
|
|
|
|
|
def decode(self, x1: torch.FloatTensor, a: torch.LongTensor, q2: torch.FloatTensor, should_clamp: bool = False) -> torch.FloatTensor: |
|
|
x1_emb = self.frame_cnn(x1) |
|
|
a_emb = rearrange(self.decoder_act_emb(a), 'b t (c h w) -> b t c h w', c=self.config["decoder_act_channels"], h=x1_emb.size(3)) |
|
|
|
|
|
decoder_input = torch.cat((x1_emb, a_emb, q2), dim=2) |
|
|
|
|
|
r = self.decoder(decoder_input) |
|
|
r = torch.clamp(r, 0, 1).mul(255).round().div(255) if should_clamp else r |
|
|
|
|
|
return r |
|
|
|
|
|
@torch.no_grad() |
|
|
def encode_decode(self, x1: torch.FloatTensor, a: torch.LongTensor, x2: torch.FloatTensor) -> torch.Tensor: |
|
|
z = self.encode(x1, a, x2) |
|
|
z = rearrange(z, 'b t c (h k) (w l) -> b t (h w) (k l c)', k=self.token_res, l=self.token_res) |
|
|
q = rearrange(self.quantizer(z).q, 'b t (h w) (k l e) -> b t e (h k) (w l)', h=self.tokens_grid_res, k=self.token_res, l=self.token_res) |
|
|
r = self.decode(x1, a, q, should_clamp=True) |
|
|
|
|
|
return r |
|
|
|
|
|
def embed_tokens(self, tokens): |
|
|
q = self.quantizer.embed_tokens(tokens) |
|
|
b, t, hw, kle = q.shape |
|
|
|
|
|
h = self.tokens_grid_res |
|
|
w = self.tokens_grid_res |
|
|
k = self.token_res |
|
|
l = self.token_res |
|
|
e = kle // (k * l) |
|
|
|
|
|
q = q.reshape(b, t, h, w, k, l, e) |
|
|
q = q.permute(0, 1, 6, 2, 4, 3, 5) |
|
|
q = q.reshape(b, t, e, h * k, w * l) |
|
|
|
|
|
return q |
|
|
|
|
|
@torch.no_grad() |
|
|
def burn_in(self, obs: torch.FloatTensor, act: torch.LongTensor) -> torch.LongTensor: |
|
|
assert obs.size(1) == act.size(1) + 1 |
|
|
quantizer_output = self(obs[:, :-1], act, obs[:, 1:]) |
|
|
|
|
|
return quantizer_output.tokens |
|
|
|
