egogym 0.1.0__py3-none-any.whl

baselines/rum/models/bet/vqbet.py

@@ -0,0 +1,410 @@
+import logging
+from enum import Enum
+from typing import Dict, Optional, Tuple
+
+import einops
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import tqdm
+
+from baselines.rum.models.bet.gpt import GPT
+from baselines.rum.models.bet.utils import MLP
+from baselines.rum.models.bet.vqvae.vqvae import VqVae
+
+GENERATOR_SEED_FIXED = 123456789
+
+
+class VQBehaviorTransformer(nn.Module):
+    GOAL_SPEC = Enum("GOAL_SPEC", "concat stack unconditional")
+
+    def __init__(
+        self,
+        obs_dim: int,
+        act_dim: int,
+        goal_dim: int,
+        gpt_model: GPT,
+        vqvae_model: VqVae,
+        offset_loss_multiplier: float = 1.0e2,
+        secondary_code_multiplier: float = 0.5,
+        gamma: float = 2.0,
+        obs_window_size=10,
+        act_window_size=10,
+        sequentially_select=False,
+        use_og_bet_loss=False,
+        use_half_and_half_loss=False,
+        temperature=1.0,
+        device="cuda",
+    ):
+        super().__init__()
+        self._obs_dim = obs_dim
+        self._act_dim = act_dim
+        self._goal_dim = goal_dim
+        self.obs_window_size = obs_window_size
+        self.act_window_size = act_window_size
+        self.sequentially_select = sequentially_select
+        self._use_og_bet_loss = use_og_bet_loss
+        self._use_half_and_half_loss = use_half_and_half_loss
+        self.temperature = temperature
+        self.device = device
+
+        if goal_dim <= 0:
+            self._cbet_method = self.GOAL_SPEC.unconditional
+        # elif obs_dim == goal_dim:
+        #     self._cbet_method = self.GOAL_SPEC.concat
+        # TODO (haritheja): this is a temporary fix, we should be able to handle different types of goals like concat or stack
+        else:
+            self._cbet_method = self.GOAL_SPEC.stack
+
+        self._gpt_model = gpt_model
+        self._vqvae_model = vqvae_model
+        self._G = self._vqvae_model.vqvae_groups  # G(number of groups)
+        self._C = self._vqvae_model.vqvae_n_embed  # C(number of code integers)
+        self._D = self._vqvae_model.embedding_dim  # D(embedding dims)
+        self._current_steps = 0
+        if self.sequentially_select:
+            self._map_to_cbet_preds_bin1 = MLP(
+                in_channels=gpt_model.config.output_dim,
+                hidden_channels=[512, 512, self._C],
+            )
+            self._map_to_cbet_preds_bin2 = MLP(
+                in_channels=gpt_model.config.output_dim + self._C,
+                hidden_channels=[512, self._C],
+            )
+        else:
+            self._map_to_cbet_preds_bin = MLP(
+                in_channels=gpt_model.config.output_dim,
+                hidden_channels=[1024, 1024, self._G * self._C],
+            )
+        self._map_to_cbet_preds_offset = MLP(
+            in_channels=gpt_model.config.output_dim,
+            hidden_channels=[
+                1024,
+                1024,
+                self._G * self._C * (act_dim * self.act_window_size),
+            ],
+        )
+        self._offset_loss_multiplier = offset_loss_multiplier
+        self._secondary_code_multiplier = secondary_code_multiplier
+        self._criterion = FocalLoss(gamma=gamma)
+
+    def forward(
+        self,
+        obs_seq: torch.Tensor,
+        goal_seq: Optional[torch.Tensor],
+        action_seq: Optional[torch.Tensor],
+        second_half: bool = False,
+    ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:
+        # VQ-BeT doesn't use "padding_seq" and "predict_with_offset" input
+        return self._predict(obs_seq, goal_seq, action_seq, second_half)
+
+    def _predict(
+        self,
+        obs_seq: torch.Tensor,
+        goal_seq: Optional[torch.Tensor],
+        action_seq: Optional[torch.Tensor],
+        second_half: bool,
+    ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Dict[str, float]]:
+        """
+        Assume dimensions are N T D for N sequences of T timesteps with dimension D.
+
+        obs_seq : (Batch size) X (obs window) X (obs dim)
+        goal_seq : (Batch size) X (goal window) X (obs dim)
+        action_seq : (Batch size) X (total window) X (action dim)
+
+        assume that
+        total_w (total window) of action sequence is
+        total_w = obs_w + act_w - 1
+        e.g. is observation window is 5, and antion pred window is 3,
+        the obs seq (o_{t-4}, o_{t-3}, o_{t-2}, o_{t-1}, o_{t}) will predict (a_{t}, a_{t+1}, a_{t+2})
+
+        However, we not only predict (a_{t}, a_{t+1}, a_{t+2}), but also all the action seq for all the actions in the sequence
+
+        o_{t-4} -> |   | -> (a_{t-4}, a_{t-3}, a_{t-2})
+        o_{t-3} -> | G | -> (a_{t-3}, a_{t-2}, a_{t-1})
+        o_{t-2} -> | P | -> (a_{t-2}, a_{t-1}, a_{t})
+        o_{t-1} -> | T | -> (a_{t-1}, a_{t}, a_{t+1})
+        o_{t}   -> |   | -> (a_{t}, a_{t+1}, a_{t+2})
+        """
+        if obs_seq.shape[1] < self.obs_window_size:
+            # if input size is smaller than obs_window size (e.g. the initial steps of env eval episodes,
+            # VQ-BeT copy the obs and tile it to match obs_window_size
+            obs_seq = torch.cat(
+                (
+                    torch.tile(
+                        obs_seq[:, 0, :],
+                        (1, self.obs_window_size - obs_seq.shape[1], 1),
+                    ),
+                    obs_seq,
+                ),
+                dim=-2,
+            )
+            if self._cbet_method == self.GOAL_SPEC.stack:
+                goal_seq = torch.cat(
+                    (
+                        torch.tile(
+                            goal_seq[:, 0, :],
+                            (1, self.obs_window_size - goal_seq.shape[1], 1),
+                        ),
+                        goal_seq,
+                    ),
+                    dim=-2,
+                )
+        if self._cbet_method == self.GOAL_SPEC.unconditional:
+            gpt_input = obs_seq
+        elif self._cbet_method == self.GOAL_SPEC.concat:
+            gpt_input = torch.cat([goal_seq, obs_seq], dim=1)
+        elif self._cbet_method == self.GOAL_SPEC.stack:
+            gpt_input = torch.cat([goal_seq, obs_seq], dim=-1)
+        else:
+            raise NotImplementedError
+
+        gpt_output = self._gpt_model(gpt_input)
+        if self._cbet_method == self.GOAL_SPEC.concat:
+            # Chop off the goal encodings.
+            gpt_output = gpt_output[:, goal_seq.size(1) :, :]
+        gpt_output = einops.rearrange(gpt_output, "N T (G C) -> (N T) (G C)", G=self._G)
+        obs = einops.rearrange(obs_seq, "N T O -> (N T) O")
+        obs = obs.unsqueeze(dim=1)
+        # note that output of offset network is G C WA,
+        # where G is number of 'layers' of Residual VQ-VAE
+        # C is number of words in each layer's dictionary
+        # and W, A is predicted action window, and predicted action dims
+        if self.sequentially_select:
+            cbet_logits1 = self._map_to_cbet_preds_bin1(gpt_output)
+            cbet_offsets = self._map_to_cbet_preds_offset(gpt_output)
+            cbet_offsets = einops.rearrange(
+                cbet_offsets, "(NT) (G C WA) -> (NT) G C WA", G=self._G, C=self._C
+            )
+            cbet_probs1 = torch.softmax(cbet_logits1 / self.temperature, dim=-1)
+            NT, choices = cbet_probs1.shape
+            G = self._G
+            sampled_centers1 = einops.rearrange(
+                torch.multinomial(cbet_probs1.view(-1, choices), num_samples=1),
+                "(NT) 1 -> NT",
+                NT=NT,
+            )
+            cbet_logits2 = self._map_to_cbet_preds_bin2(
+                torch.cat(
+                    (gpt_output, F.one_hot(sampled_centers1, num_classes=self._C)),
+                    axis=1,
+                )
+            )
+            cbet_probs2 = torch.softmax(cbet_logits2 / self.temperature, dim=-1)
+            sampled_centers2 = einops.rearrange(
+                torch.multinomial(cbet_probs2.view(-1, choices), num_samples=1),
+                "(NT) 1 -> NT",
+                NT=NT,
+            )
+            sampled_centers = torch.stack(
+                (sampled_centers1, sampled_centers2), axis=1
+            )  # NT, G
+        else:
+            cbet_logits = self._map_to_cbet_preds_bin(gpt_output)
+            cbet_offsets = self._map_to_cbet_preds_offset(gpt_output)
+            cbet_logits = einops.rearrange(
+                cbet_logits, "(NT) (G C) -> (NT) G C", G=self._G
+            )
+            cbet_offsets = einops.rearrange(
+                cbet_offsets, "(NT) (G C WA) -> (NT) G C WA", G=self._G, C=self._C
+            )
+            cbet_probs = torch.softmax(cbet_logits / self.temperature, dim=-1)
+            NT, G, choices = cbet_probs.shape
+            sampled_centers = einops.rearrange(
+                torch.multinomial(cbet_probs.view(-1, choices), num_samples=1),
+                "(NT G) 1 -> NT G",
+                NT=NT,
+            )
+        if action_seq is not None:
+            n, total_w, act_dim = action_seq.shape
+            act_w = self._vqvae_model.input_dim_h
+            obs_w = total_w + 1 - act_w
+            output_shape = (n, obs_w, act_w, act_dim)
+            output = torch.empty(output_shape).to(action_seq.device)
+            for i in range(obs_w):
+                output[:, i, :, :] = action_seq[:, i : i + act_w, :]
+            action_seq = einops.rearrange(output, "N T W A -> (N T) W A")
+            _, action_bins = self._vqvae_model.get_code(
+                action_seq, obs
+            )  # action_bins: NT, G
+
+        with torch.no_grad():
+            centers = self._vqvae_model.draw_code_forward(sampled_centers).view(
+                NT, -1, self._D
+            )
+            return_decoder_input = einops.rearrange(
+                centers.clone().detach(), "NT G D -> NT (G D)"
+            )
+            decoded_action = (
+                self._vqvae_model.get_action_from_latent(return_decoder_input, obs)
+                .clone()
+                .detach()
+            )  # NT, A
+
+        def get_offset_from_centers(centers):
+            indices = (
+                torch.arange(NT, device=self.device).unsqueeze(1),
+                torch.arange(self._G, device=self.device).unsqueeze(0),
+                centers,
+            )
+            # Use advanced indexing to sample the values
+            sampled_offsets = cbet_offsets[indices]  # NT, G, WA or NT, G, A
+            sampled_offsets = sampled_offsets.sum(dim=1)
+            sampled_offsets = einops.rearrange(
+                sampled_offsets, "NT (W A) -> NT W A", W=self._vqvae_model.input_dim_h
+            )
+            return sampled_offsets
+
+        if self._use_og_bet_loss and action_seq is not None:
+            sampled_offsets = get_offset_from_centers(action_bins)
+        else:
+            sampled_offsets = get_offset_from_centers(sampled_centers)
+
+        a_hat = decoded_action + sampled_offsets
+
+        if action_seq is None:
+            return a_hat, None, {}
+        # Figure out the loss for the actions.
+        # First, we need to find the GT VQ codes for each action.
+
+        # Now we can compute the loss.
+        if action_seq.ndim == 2:
+            action_seq = action_seq.unsqueeze(0)
+
+        offset_target = action_seq - decoded_action
+        if self._use_half_and_half_loss and action_seq is not None:
+            offset_gt = get_offset_from_centers(action_bins)
+            offset_sampled = get_offset_from_centers(sampled_centers)
+            offset_loss = 0.5 * torch.nn.L1Loss()(
+                offset_target, offset_gt
+            ) + 0.5 * torch.nn.L1Loss()(offset_target, offset_sampled)
+        else:
+            offset_loss = torch.nn.L1Loss()(offset_target, sampled_offsets)
+
+        if self.sequentially_select:
+            cbet_loss1 = self._criterion(  # F.cross_entropy
+                cbet_logits1[:, :],
+                action_bins[:, 0],
+            )
+            cbet_logits2 = self._map_to_cbet_preds_bin2(
+                torch.cat(
+                    (gpt_output, F.one_hot(action_bins[:, 0], num_classes=self._C)),
+                    axis=1,
+                )
+            )
+            cbet_loss2 = self._criterion(  # F.cross_entropy
+                cbet_logits2[:, :],
+                action_bins[:, 1],
+            )
+        else:
+            cbet_loss1 = self._criterion(  # F.cross_entropy
+                cbet_logits[:, 0, :],
+                action_bins[:, 0],
+            )
+            cbet_loss2 = self._criterion(  # F.cross_entropy
+                cbet_logits[:, 1, :],
+                action_bins[:, 1],
+            )
+        cbet_loss = cbet_loss1 * 5 + cbet_loss2 * self._secondary_code_multiplier
+
+        equal_total_code_rate = (
+            torch.sum(
+                (torch.sum((action_bins == sampled_centers).int(), axis=1) == G).int()
+            )
+            / NT
+        )
+        equal_primary_code_rate = torch.sum(
+            (action_bins[:, 0] == sampled_centers[:, 0]).int()
+        ) / (NT)
+        equal_secondary_code_rate = torch.sum(
+            (action_bins[:, 1] == sampled_centers[:, 1]).int()
+        ) / (NT)
+        # if second_half:
+        #     cbet_loss = cbet_loss * 0
+        loss = cbet_loss + self._offset_loss_multiplier * offset_loss
+        action_mse = F.mse_loss(a_hat, action_seq, reduction="none")
+        action_l1 = F.l1_loss(a_hat, action_seq, reduction="none")
+        norm = torch.norm(action_seq, p=2, dim=-1, keepdim=True) + 1e-9
+        normalized_mse = (action_mse / norm).mean()
+
+        translation_loss = F.mse_loss(a_hat[:, :, :3], action_seq[:, :, :3]).detach()
+        rotation_loss = F.mse_loss(a_hat[:, :, 3:6], action_seq[:, :, 3:6]).detach()
+        gripper_loss = F.mse_loss(a_hat[:, :, 6:], action_seq[:, :, 6:]).detach()
+
+        loss_dict = {
+            "classification_loss": cbet_loss.detach().cpu().item(),
+            "offset_loss": offset_loss.detach().cpu().item(),
+            "loss": loss.detach().cpu().item(),
+            "equal_total_code_rate": equal_total_code_rate,
+            "equal_primary_code_rate": equal_primary_code_rate,
+            "equal_secondary_code_rate": equal_secondary_code_rate,
+            "L2_loss": action_mse.mean().detach().cpu().item(),
+            "L2_loss_normalized": normalized_mse.mean().detach().cpu().item(),
+            "L1_loss": action_l1.mean().detach().cpu().item(),
+            "translation_loss": translation_loss,
+            "rotation_loss": rotation_loss,
+            "gripper_loss": gripper_loss,
+        }
+        return a_hat, loss, loss_dict
+
+    # def configure_optimizers(self, weight_decay, learning_rate, betas):
+
+    #     optimizer1 = self._gpt_model.configure_optimizers(
+    #         weight_decay=weight_decay,
+    #         learning_rate=learning_rate,
+    #         betas=betas,
+    #     )
+    #     if self.sequentially_select:
+    #         optimizer1.add_param_group({"params": self._map_to_cbet_preds_bin1.parameters()})
+    #         optimizer1.add_param_group({"params": self._map_to_cbet_preds_bin2.parameters()})
+    #     else:
+    #         optimizer1.add_param_group({"params": self._map_to_cbet_preds_bin.parameters()})
+    #     optimizer2 = torch.optim.AdamW(self._map_to_cbet_preds_offset.parameters(), lr=learning_rate, weight_decay=weight_decay, betas=betas)
+    #     return {"optimizer1": optimizer1, "optimizer2": optimizer2}
+
+    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 _load_from_state_dict(self, *args, **kwargs):
+        # Don't fit kmeans if we are loading from a state dict.
+        # if (path / "cbet_model.pt").exists():
+        #     self.load_state_dict(torch.load(path / "cbet_model.pt"))
+        # elif (path / "gpt_model.pt").exists():
+        #     self._gpt_model.load_state_dict(torch.load(path / "gpt_model.pt"))
+        # else:
+        #     logging.warning("No model found at %s", path)
+        return super()._load_from_state_dict(*args, **kwargs)
+
+    # def load_model(self, path):
+    #     if (path / "cbet_model.pt").exists():
+    #         self.load_state_dict(torch.load(path / "cbet_model.pt"))
+    #     elif (path / "gpt_model.pt").exists():
+    #         self._gpt_model.load_state_dict(torch.load(path / "gpt_model.pt"))
+    #     else:
+    #         logging.warning("No model found at %s", path)
+
+
+class FocalLoss(nn.Module):
+    def __init__(self, gamma: float = 0, reduction: str = "mean"):
+        super(FocalLoss, self).__init__()
+        self.gamma = gamma
+        if reduction not in ("mean", "sum", "none"):
+            raise NotImplementedError
+        self.reduction = reduction
+
+    def forward(self, input, target):
+        logpt = F.log_softmax(input, dim=-1)
+        logpt = logpt.gather(1, target.view(-1, 1)).view(-1)
+        pt = logpt.exp()
+
+        loss = -1 * (1 - pt) ** self.gamma * logpt
+        if self.reduction == "mean":
+            return loss.mean()
+        elif self.reduction == "sum":
+            return loss.sum()
+        else:
+            return loss

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

@@ -0,0 +1,3 @@
+from baselines.rum.models.bet.vqvae.vector_quantize_pytorch import VectorQuantize
+from baselines.rum.models.bet.vqvae.residual_vq import ResidualVQ
+from baselines.rum.models.bet.vqvae.vqvae import VqVae