egogym 0.1.0__py3-none-any.whl

baselines/rum/models/bet/vqvae/vqvae.py

@@ -0,0 +1,313 @@
+import torch
+import torch.nn as nn
+import einops
+from baselines.rum.models.bet.vqvae.residual_vq import ResidualVQ
+from baselines.rum.models.bet.vqvae.vqvae_utils import get_tensor, weights_init_encoder
+
+
+class EncoderMLP(nn.Module):
+    def __init__(
+        self,
+        input_dim,
+        output_dim=16,
+        hidden_dim=128,
+        layer_num=1,
+        last_activation=None,
+    ):
+        super(EncoderMLP, self).__init__()
+        layers = []
+
+        layers.append(nn.Linear(input_dim, hidden_dim))
+        layers.append(nn.ReLU())
+        for _ in range(layer_num):
+            layers.append(nn.Linear(hidden_dim, hidden_dim))
+            layers.append(nn.ReLU())
+
+        self.encoder = nn.Sequential(*layers)
+        self.fc = nn.Linear(hidden_dim, output_dim)
+
+        if last_activation is not None:
+            self.last_layer = last_activation
+        else:
+            self.last_layer = None
+        self.apply(weights_init_encoder)
+
+    def forward(self, x):
+        h = self.encoder(x)
+        state = self.fc(h)
+        if self.last_layer:
+            state = self.last_layer(state)
+        return state
+
+
+class CondiitonalEncoderMLP(nn.Module):
+    def __init__(
+        self,
+        input_dim,
+        output_dim=16,
+        hidden_dim=512,
+        layer_num=2,
+        last_activation=None,
+        obs_dim=None,
+    ):
+        super(CondiitonalEncoderMLP, self).__init__()
+        layers = []
+
+        layers.append(nn.Linear(input_dim + obs_dim, hidden_dim))
+        layers.append(nn.ReLU())
+        for _ in range(layer_num):
+            layers.append(nn.Linear(hidden_dim, hidden_dim))
+            layers.append(nn.ReLU())
+
+        self.encoder = nn.Sequential(*layers)
+        self.fc = nn.Linear(hidden_dim, output_dim)
+
+        if last_activation is not None:
+            self.last_layer = last_activation
+        else:
+            self.last_layer = None
+        self.apply(weights_init_encoder)
+
+    def forward(self, x, obs):
+        x = torch.cat((x, obs), dim=1)
+        h = self.encoder(x)
+        state = self.fc(h)
+        if self.last_layer:
+            state = self.last_layer(state)
+        return state
+
+
+class VqVae(nn.Module):
+    def __init__(
+        self,
+        obs_dim=60,
+        input_dim_h=10,  # length of action chunk
+        input_dim_w=9,  # action dim
+        n_latent_dims=512,
+        vqvae_n_embed=32,
+        vqvae_groups=4,
+        eval=True,
+        device="cuda",
+        load_dir=None,
+        enc_loss_type="skip_vqlayer",
+        obs_cond=False,
+        encoder_loss_multiplier=1.0,
+        act_scale=1.0,
+    ):
+        super(VqVae, self).__init__()
+        self.n_latent_dims = n_latent_dims  # 64
+        self.input_dim_h = input_dim_h
+        self.input_dim_w = input_dim_w
+        self.rep_dim = self.n_latent_dims
+        self.vqvae_n_embed = vqvae_n_embed  # 120
+        self.vqvae_lr = 1e-3
+        self.vqvae_groups = vqvae_groups
+        self.device = device
+        self.enc_loss_type = enc_loss_type
+        self.obs_cond = obs_cond
+        self.encoder_loss_multiplier = encoder_loss_multiplier
+        self.act_scale = act_scale
+
+        discrete_cfg = {"groups": self.vqvae_groups, "n_embed": self.vqvae_n_embed}
+        self.vq_layer = ResidualVQ(
+            dim=self.n_latent_dims,
+            num_quantizers=discrete_cfg["groups"],
+            codebook_size=self.vqvae_n_embed,
+            eval=eval,
+        ).to(self.device)
+        self.embedding_dim = self.n_latent_dims
+        self.vq_layer.device = device
+
+        if self.input_dim_h == 1:
+            if self.obs_cond:
+                self.encoder = CondiitonalEncoderMLP(
+                    input_dim=input_dim_w, output_dim=n_latent_dims, obs_dim=obs_dim
+                ).to(self.device)
+                self.decoder = CondiitonalEncoderMLP(
+                    input_dim=n_latent_dims, output_dim=input_dim_w, obs_dim=obs_dim
+                ).to(self.device)
+            else:
+                self.encoder = EncoderMLP(
+                    input_dim=input_dim_w, output_dim=n_latent_dims
+                ).to(self.device)
+                self.decoder = EncoderMLP(
+                    input_dim=n_latent_dims, output_dim=input_dim_w
+                ).to(self.device)
+        else:
+            if self.obs_cond:
+                self.encoder = CondiitonalEncoderMLP(
+                    input_dim=input_dim_w * self.input_dim_h,
+                    output_dim=n_latent_dims,
+                    obs_dim=obs_dim,
+                ).to(self.device)
+                self.decoder = CondiitonalEncoderMLP(
+                    input_dim=n_latent_dims,
+                    output_dim=input_dim_w * self.input_dim_h,
+                    obs_dim=obs_dim,
+                ).to(self.device)
+            else:
+                self.encoder = EncoderMLP(
+                    input_dim=input_dim_w * self.input_dim_h, output_dim=n_latent_dims
+                ).to(self.device)
+                self.decoder = EncoderMLP(
+                    input_dim=n_latent_dims, output_dim=input_dim_w * self.input_dim_h
+                ).to(self.device)
+
+        # params = list(self.encoder.parameters()) + list(self.decoder.parameters()) + list(self.vq_layer.parameters())
+        # self.vqvae_optimizer = torch.optim.Adam(params, lr=self.vqvae_lr, weight_decay=0.0001)
+
+        if load_dir is not None:
+            try:
+                state_dict = torch.load(load_dir)
+            except RuntimeError:
+                state_dict = torch.load(load_dir, map_location=torch.device("cpu"))
+
+            new_dict = {}
+
+            prefix_to_remove = "_rvq."
+
+            for key, value in state_dict["loss_fn"].items():
+                if key.startswith(prefix_to_remove):
+                    new_key = key[len(prefix_to_remove) :]  # Remove the prefix
+                    new_dict[new_key] = value
+                else:
+                    new_dict[key] = value
+
+            self.load_state_dict(new_dict, strict=False)
+
+        if eval:
+            self.vq_layer.eval()
+        else:
+            self.vq_layer.train()
+
+    def draw_logits_forward(self, encoding_logits):
+        z_embed = self.vq_layer.draw_logits_forward(encoding_logits)
+        return z_embed
+
+    def draw_code_forward(self, encoding_indices):
+        with torch.no_grad():
+            z_embed = self.vq_layer.get_codes_from_indices(encoding_indices)
+            z_embed = z_embed.sum(dim=0)
+        return z_embed
+
+    def get_action_from_latent(self, latent, obs=None):
+        if self.obs_cond:
+            output = self.decoder(latent, obs[:, -1]) * self.act_scale
+        else:
+            output = self.decoder(latent) * self.act_scale
+        if self.input_dim_h == 1:
+            return einops.rearrange(output, "N (T A) -> N T A", A=self.input_dim_w)
+        else:
+            return einops.rearrange(output, "N (T A) -> N T A", A=self.input_dim_w)
+
+    def preprocess(self, state):
+        if not torch.is_tensor(state):
+            state = get_tensor(state, self.device)
+        if self.input_dim_h == 1:
+            state = state.squeeze(-2)  # state.squeeze(-1)
+        else:
+            state = einops.rearrange(state, "N T A -> N (T A)")
+        return state.to(self.device)
+
+    def get_code(self, state, obs=None, required_recon=False):
+        state = state / self.act_scale
+        state = self.preprocess(state)
+        with torch.no_grad():
+            if self.obs_cond:
+                state_rep = self.encoder(state, obs[:, -1])
+            else:
+                state_rep = self.encoder(state)
+            state_rep_shape = state_rep.shape[:-1]
+            state_rep_flat = state_rep.view(state_rep.size(0), -1, state_rep.size(1))
+            state_rep_flat, vq_code, vq_loss_state = self.vq_layer(state_rep_flat)
+            state_vq = state_rep_flat.view(*state_rep_shape, -1)
+            vq_code = vq_code.view(*state_rep_shape, -1)
+            vq_loss_state = torch.sum(vq_loss_state)
+            if required_recon:
+                if self.obs_cond:
+                    recon_state = self.decoder(state_vq, obs[:, -1]) * self.act_scale
+                    recon_state_ae = (
+                        self.decoder(state_rep, obs[:, -1]) * self.act_scale
+                    )
+                else:
+                    recon_state = self.decoder(state_vq) * self.act_scale
+                    recon_state_ae = self.decoder(state_rep) * self.act_scale
+                if self.input_dim_h == 1:
+                    return state_vq, vq_code, recon_state, recon_state_ae
+                else:
+                    return (
+                        state_vq,
+                        vq_code,
+                        torch.swapaxes(recon_state, -2, -1),
+                        torch.swapaxes(recon_state_ae, -2, -1),
+                    )
+            else:
+                return state_vq, vq_code
+
+    def vqvae_update(self, state, obs=None):
+        state = state / self.act_scale
+        state = self.preprocess(state)
+        if self.obs_cond:
+            state_rep = self.encoder(state, obs[:, -1])
+        else:
+            state_rep = self.encoder(state)
+        state_rep_shape = state_rep.shape[:-1]
+        state_rep_flat = state_rep.view(state_rep.size(0), -1, state_rep.size(1))
+        state_rep_flat, vq_code, vq_loss_state = self.vq_layer(state_rep_flat)
+        state_vq = state_rep_flat.view(*state_rep_shape, -1)
+        vq_code = vq_code.view(*state_rep_shape, -1)
+        vq_loss_state = torch.sum(vq_loss_state)
+
+        if self.obs_cond:
+            dec_out = self.decoder(state_vq, obs[:, -1])
+        else:
+            dec_out = self.decoder(state_vq)
+        encoder_loss = (state - dec_out).abs().mean()
+
+        rep_loss = encoder_loss * self.encoder_loss_multiplier + (vq_loss_state * 5)
+
+        # self.vqvae_optimizer.zero_grad()
+        # rep_loss.backward()
+        # self.vqvae_optimizer.step()
+        vqvae_recon_loss = torch.nn.MSELoss()(state, dec_out)
+        loss_dict = {
+            "rep_loss": rep_loss.detach().cpu().item(),
+            "vq_loss_state": vq_loss_state.detach().cpu().item(),
+            "vqvae_recon_loss_l1": encoder_loss.detach().cpu().item(),
+            "vqvae_recon_loss_l2": vqvae_recon_loss.detach().cpu().item(),
+            "n_different_codes": len(torch.unique(vq_code)),
+            "n_different_combinations": len(torch.unique(vq_code, dim=0)),
+        }
+        return rep_loss, loss_dict
+
+    def configure_optimizers(self, weight_decay, learning_rate, betas):
+        params = (
+            list(self.encoder.parameters())
+            + list(self.decoder.parameters())
+            + list(self.vq_layer.parameters())
+        )
+        optimizer = torch.optim.AdamW(
+            params,
+            weight_decay=weight_decay,
+            learning_rate=learning_rate,
+            betas=betas,
+        )
+        return optimizer
+
+    def _begin_epoch(self, optimizer, **kwargs):
+        # log codebook usage rate for debugging
+        # lr_0 = optimizer.param_groups[0]["lr"]
+        # lr_neg1 = optimizer.param_groups[-1]["lr"]
+        # return {"lr_0": lr_0, "lr_neg1": lr_neg1}
+        return None
+
+    # def state_dict(self):
+    #     return {'encoder': self.encoder.state_dict(),
+    #             'decoder': self.decoder.state_dict(),
+    #             'vq_embedding': self.vq_layer.state_dict()}
+
+    # def load_state_dict(self, state_dict):
+    #     self.encoder.load_state_dict(state_dict['encoder'])
+    #     self.decoder.load_state_dict(state_dict['decoder'])
+    #     self.vq_layer.load_state_dict(state_dict['vq_embedding'])
+    #     self.vq_layer.eval()

baselines/rum/models/bet/vqvae/vqvae_utils.py

@@ -0,0 +1,30 @@
+import torch
+import numpy as np
+import torch.nn as nn
+import os.path as osp
+
+
+def weights_init_encoder(m):
+    if isinstance(m, nn.Linear):
+        nn.init.orthogonal_(m.weight.data)
+        m.bias.data.fill_(0.0)
+    elif isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
+        assert m.weight.size(2) == m.weight.size(3)
+        m.weight.data.fill_(0.0)
+        m.bias.data.fill_(0.0)
+        mid = m.weight.size(2) // 2
+        gain = nn.init.calculate_gain("relu")
+        nn.init.orthogonal_(m.weight.data[:, :, mid, mid], gain)
+
+
+def get_tensor(z, device):
+    if z is None:
+        return None
+    if z[0].dtype == np.dtype("O"):
+        return None
+    if len(z.shape) == 1:
+        return torch.FloatTensor(z.copy()).to(device).unsqueeze(0)
+        # return torch.from_numpy(z.copy()).float().to(device).unsqueeze(0)
+    else:
+        return torch.FloatTensor(z.copy()).to(device)
+        # return torch.from_numpy(z.copy()).float().to(device)

baselines/rum/models/custom.py

@@ -0,0 +1,33 @@
+import torch
+import torch.nn as nn
+
+class CustomModel(nn.Module):
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+        # TODO: set model
+        self.model = None
+
+    def to(self, device):
+        # TODO: move model to device
+        if self.model is not None:
+            self.model.to(device)
+        return self
+
+    def forward(self, x):
+        # TODO: forward pass of model
+        pass
+    
+    def step(self, data, *args, **kwargs):
+        images, actions = data
+        # images is a tensor of shape (1, image_buffer_size, 3, 256, 256), where the images are in chronological order
+        # actions is a tensor of shape (1, image_buffer_size, 7), where the actions are in chronological order, and the final action is padding
+        
+        # TODO: implement pass that takes in observations as described above out outputs a 7-dimensional action
+        action = torch.zeros(7)
+        logs = {}
+        
+        return action, logs
+
+    def reset(self):
+        # TODO: optional; this method is called once the robot is homed
+        pass

baselines/rum/models/encoders/abstract_base_encoder.py

@@ -0,0 +1,70 @@
+import pathlib
+import warnings
+from abc import ABC, abstractmethod
+from typing import Union
+
+import torch
+import torch.nn as nn
+
+ENCODER_LOADING_ERROR_MSG = (
+    "Could not load encoder weights: defaulting to pretrained weights"
+)
+
+
+class AbstractEncoder(nn.Module, ABC):
+    def __init__(self):
+        super().__init__()
+
+    @property
+    @abstractmethod
+    def feature_dim(self):
+        pass
+
+    def transform(self, x):
+        return x
+
+    def to(self, device):
+        self.device = device
+        return super().to(device)
+
+    @abstractmethod
+    def forward(self, x):
+        pass
+
+    def load_weights(
+        self, weight_path: Union[pathlib.Path, str], strict: bool = True
+    ) -> None:
+        try:
+            checkpoint = torch.load(weight_path, map_location="cpu")
+            weight_dict = checkpoint.get("model", checkpoint)
+
+            self.model.load_state_dict(weight_dict, strict=strict)
+            return
+
+        except RuntimeError:
+            warnings.warn("Couldn't load full model weights. Trying fallback options...")
+
+        fallback_prefixes = [
+            "model.",
+            "encoder.",
+            "module.base_encoder._orig_mod.",
+        ]
+
+        for prefix in fallback_prefixes:
+            try:
+                filtered_weights = {
+                    k.replace(prefix, ""): v
+                    for k, v in weight_dict.items()
+                    if k.startswith(prefix)
+                }
+
+                filtered_weights.pop("fc.weight", None)
+                filtered_weights.pop("fc.bias", None)
+
+                self.model.load_state_dict(filtered_weights, strict=strict)
+                return
+
+            except RuntimeError:
+                warnings.warn(f"Failed loading with prefix: '{prefix}'")
+
+        warnings.warn(ENCODER_LOADING_ERROR_MSG)

baselines/rum/models/encoders/identity.py

@@ -0,0 +1,45 @@
+"""
+A generic model wrapper for dummy identity encoders.
+"""
+
+import pathlib
+from typing import Union
+from torch import nn
+import torch
+from .abstract_base_encoder import AbstractEncoder
+
+
+class IdentityEncoder(AbstractEncoder):
+    def __init__(
+        self,
+        model_name,
+        pretrained: bool = True,
+        weight_path: Union[None, str, pathlib.Path] = None,
+    ):
+        super().__init__()
+        self._model_name = model_name
+        self.device = None
+        self.model = nn.Identity()
+
+        # Use a placeholder parameter if the model has no parameters
+        if len(list(self.model.parameters())) == 0:
+            self.placeholder_param = nn.Parameter(torch.zeros(1, requires_grad=True))
+        else:
+            self.placeholder_param = None
+
+    def transform(self, x):
+        return x
+
+    @property
+    def feature_dim(self):
+        return 0
+
+    def to(self, device):
+        self.device = device
+        self.model.to(device)
+        if self.placeholder_param is not None:
+            self.placeholder_param.to(device)
+        return self
+
+    def forward(self, x):
+        return self.transform(x)

baselines/rum/models/encoders/timm_encoders.py

@@ -0,0 +1,82 @@
+"""
+A generic model wrapper for timm encoders.
+"""
+
+import pathlib
+from typing import Iterable, Optional, Union
+import warnings
+
+import einops
+import timm
+import torch
+import torch.nn as nn
+
+from baselines.rum.models.encoders.abstract_base_encoder import AbstractEncoder
+from baselines.rum.utils.decord_transforms import create_transform
+
+
+class TimmModel(AbstractEncoder):
+    def __init__(
+        self,
+        model_name: str = "hf-hub:notmahi/dobb-e",
+        pretrained: bool = True,
+        weight_path: Union[None, str, pathlib.Path] = None,
+    ):
+        super().__init__()
+        self._model_name = model_name
+
+        self.model = timm.create_model(model_name, pretrained=pretrained, num_classes=0)
+        if weight_path:
+            self.load_weights(weight_path, strict=False)
+
+    def transform(self, x):
+        return x
+
+    @property
+    def feature_dim(self):
+        return self.model.num_features
+
+    def to(self, device):
+        self.model.to(device)
+        return self
+
+    def forward(self, x):
+        return self.model(self.transform(x))
+
+
+class TimmSSL(TimmModel):
+    def __init__(
+        self,
+        model_name: str = "hf-hub:notmahi/dobb-e",
+        pretrained=True,
+        override_aug_kwargs=dict(
+            hflip=0.0, vflip=0.0, scale=(1.0, 1.0), crop_pct=0.875
+        ),
+        weight_path: Union[None, str, pathlib.Path] = None,
+    ):
+        super().__init__(model_name, pretrained=pretrained)
+        data_cfg = timm.data.resolve_data_config(self.model.pretrained_cfg)
+        # Now define the transforms.
+        data_cfg.update(override_aug_kwargs)
+        data_cfg["is_training"] = True
+        self._train_transform = create_transform(**data_cfg)
+        data_cfg["is_training"] = False
+        self._test_transform = create_transform(**data_cfg)
+        if weight_path is not None:
+            self.load_weights(weight_path, strict=True)
+
+    def transform(self, x):
+        return (
+            self._train_transform(x) if self.model.training else self._test_transform(x)
+        )
+
+    def forward(self, x):
+        # Split the input into frames and labels.
+        images, *labels = x
+        # Flatten the frames into a single batch.
+        flattened_images = einops.rearrange(images, "b t c h w -> (b t) c h w")
+        # Transform and pass through the model.
+        result = self.model(self.transform(flattened_images))
+        # Unflatten the result.
+        result = einops.rearrange(result, "(b t) c -> b t c", b=images.shape[0])
+        return result