Upload grp_model.py with huggingface_hub

grp_model.py

@@ -0,0 +1,356 @@
+import numpy as np
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+def get_patches_fast(images, cfg):
+    from einops import rearrange
+    batch_size, height, width, channels = images.shape
+    patch_size = cfg.patch_size ## n_patches = 8
+
+    patches = rearrange(images[:,:,:,:3], 'b (h p1) (w p2) c -> b (h w) (p1 p2 c)', p1 = patch_size, p2 = patch_size)
+    if channels > 3:
+        ## History stacking in the channel dimension for observations only, not goal images.
+        patches = rearrange(images, 'b (h p1) (w p2) (c hs) -> b (h w hs) (p1 p2 c)', p1 = patch_size, p2 = patch_size, hs=cfg.policy.obs_stacking) ## Stack the history in the channel dimension
+    return patches
+
+
+def calc_positional_embeddings(sequence_length, d):
+    result = torch.ones(sequence_length, d)
+    for i in range(sequence_length):
+        for j in range(d):
+            result[i][j] = np.sin(i / (10000 ** (j / d))) if j % 2 == 0 else np.cos(i / (10000 ** ((j - 1) / d)))
+    return result
+
+
+class Head(nn.Module):
+    """ one head of self-attention """
+
+    def __init__(self, head_size, n_embd, dropout):
+        super().__init__()
+        self.key = nn.Linear(n_embd, head_size, bias=False)
+        self.query = nn.Linear(n_embd, head_size, bias=False)
+        self.value = nn.Linear(n_embd, head_size, bias=False)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x, mask=None):
+        B,T,C = x.shape
+        # TODO: 
+        ## Provide the block masking logic for the attention head
+        k = self.key(x)
+        q = self.query(x)
+        wei = q @ k.transpose(-2,-1) * C**-0.5
+        #wei = wei.masked_fill(mask == 0, float('-inf'))
+        if mask is not None:
+            wei = wei.masked_fill(mask == 0, float('-inf'))
+        wei = F.softmax(wei, dim=-1)
+        wei = self.dropout(wei)
+        v = self.value(x)
+        out = wei @ v
+        return out
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, num_heads, head_size, n_embd, dropout):
+        super().__init__()
+        self.heads = nn.ModuleList([Head(head_size, n_embd=n_embd, dropout=dropout) for _ in range(num_heads)])
+        self.proj = nn.Linear(n_embd, n_embd)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x, mask=None):
+        with torch.profiler.record_function("Self-Attention"):
+            out = torch.cat([h(x, mask) for h in self.heads], dim=-1)
+            out = self.dropout(self.proj(out))
+        return out
+
+
+class FeedFoward(nn.Module):
+    def __init__(self, n_embd, dropout):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(n_embd, 4 * n_embd),
+            nn.ReLU(),
+            nn.Linear(4 * n_embd, n_embd),
+            nn.Dropout(dropout),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class Block(nn.Module):
+    def __init__(self, n_embd, n_head, dropout):
+        super().__init__()
+        head_size = n_embd // n_head
+        self.sa = MultiHeadAttention(n_head, head_size, n_embd=n_embd, dropout=dropout)
+        self.ffwd = FeedFoward(n_embd, dropout)
+        self.ln1 = nn.LayerNorm(n_embd)
+        self.ln2 = nn.LayerNorm(n_embd)
+
+    def forward(self, x, mask=None):
+        x = x + self.sa(self.ln1(x), mask)
+        x = x + self.ffwd(self.ln2(x))
+        return x
+
+
+class GRP(nn.Module):
+    def __init__(self, cfg, mlp_ratio=4):
+        super(GRP, self).__init__()
+        self._cfg = cfg
+        chars = cfg.dataset.chars_list
+        cfg.vocab_size = len(chars)
+        # TODO: 
+        ## Provide the logic for the GRP network
+
+        # 4) Transformer encoder blocks
+
+        # 5) Classification MLPk
+
+        # 1) Embeddings
+        # Calculate patch dimension: patch_size * patch_size * 3 (RGB)
+        patch_dim = cfg.patch_size * cfg.patch_size * 3
+        self.patch_embedding = nn.Linear(patch_dim, cfg.n_embd)
+        
+        # Token embedding for text goals (if not using T5)
+        # Check if dataset config exists and has encode_with_t5, else assume False or handle safely
+        use_t5 = False
+        if hasattr(cfg, 'dataset') and hasattr(cfg.dataset, 'encode_with_t5'):
+            use_t5 = cfg.dataset.encode_with_t5
+            
+        if not use_t5:
+            self.token_embedding_table = nn.Embedding(cfg.vocab_size, cfg.n_embd)
+
+        # Learnable positional embeddings or fixed buffer
+        # Calculate maximum sequence length: CLS + text tokens + separator + goal img patches + obs patches
+        num_patches_per_image = cfg.n_patches * cfg.n_patches
+        max_seq_len = 1 + cfg.max_block_size + 1 + num_patches_per_image + num_patches_per_image * cfg.policy.obs_stacking
+        # Using a fixed buffer as per the helper function provided
+        pos_emb = calc_positional_embeddings(max_seq_len, cfg.n_embd)
+        self.register_buffer('pos_embedding', pos_emb)
+
+        # Special Tokens
+        self.cls_token = nn.Parameter(torch.randn(1, 1, cfg.n_embd))
+        self.goal_token = nn.Parameter(torch.randn(1, 1, cfg.n_embd))
+        
+        # Dropout
+        self.dropout = nn.Dropout(cfg.dropout)
+
+        # 2) Transformer encoder blocks
+        # Corrected: using cfg.n_blocks instead of cfg.n_layer
+        self.blocks = nn.Sequential(*[
+            Block(cfg.n_embd, n_head=cfg.n_head, dropout=cfg.dropout) 
+            for _ in range(cfg.n_blocks)
+        ])
+        
+        self.ln_f = nn.LayerNorm(cfg.n_embd) # Final layer norm
+
+        # 3) Classification MLP / Action Head
+        # Corrected: using cfg.action_dim (root) instead of cfg.env.action_dim
+        output_dim = cfg.action_dim * cfg.policy.action_stacking
+        self.lm_head = nn.Linear(cfg.n_embd, output_dim)
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+    def forward(self, images, goals_txt, goal_imgs, targets=None, pose=None, mask_=False):
+        #device = images.device
+        n, c, h, w = images.shape
+        obs_patches = get_patches_fast(images, self._cfg)
+        patches_g = get_patches_fast(goal_imgs, self._cfg)
+        if self._cfg.dataset.encode_with_t5:
+            goals_e = goals_txt
+            B, T, E = goals_txt.shape
+        else:
+            goals_e = self.token_embedding_table(goals_txt)
+            B, E = goals_txt.shape
+            T = self._cfg.max_block_size
+
+        # TODO: 
+        ## Provide the logic to produce the output and loss for the GRP
+        
+        # Map the vector corresponding to each patch to the hidden size dimension
+
+        # Adding classification and goal_img tokens to the tokens
+
+        # Adding positional embedding
+
+        # Compute blocked masks
+
+        # Transformer Blocks
+
+        # Getting the classification token only
+
+        # Compute output and loss
+        obs_emb = self.patch_embedding(obs_patches)
+        goal_img_emb = self.patch_embedding(patches_g)
+        cls_tokens = self.cls_token.expand(B, -1, -1)
+        sep_tokens = self.goal_token.expand(B, -1, -1)
+
+        # Concatenate everything into one sequence:
+        # [CLS] + [Goal Text] + [Goal Separator] + [Goal Img] + [Observation Patches]
+        x = torch.cat((cls_tokens, goals_e, sep_tokens, goal_img_emb, obs_emb), dim=1)
+        
+        # Adding positional embedding
+        # We slice the pre-calculated buffer to the current sequence length
+        seq_len = x.shape[1]
+        x = x + self.pos_embedding[:seq_len, :].to(self._cfg.device)
+        x = self.dropout(x)
+
+        # Compute blocked masks
+        mask = None
+        if mask_:
+            # Create a causal mask (lower triangular)
+            mask = torch.tril(torch.ones(seq_len, seq_len, device=device))
+
+        # Transformer Blocks
+        for block in self.blocks:
+            x = block(x, mask)
+        
+        x = self.ln_f(x)
+
+        # Getting the classification token only (index 0)
+        # This token aggregates information from the entire sequence
+        cls_out = x[:, 0, :]
+
+        # Compute output
+        logits = self.lm_head(cls_out) # Shape: (B, action_dim * action_stacking)
+
+        # Compute loss
+        loss = None
+        if targets is not None:
+            # Typically MSE loss for continuous action regression
+            loss = F.mse_loss(logits, targets)
+
+        return (logits, loss)
+    
+    def resize_image(self, image):
+        """
+        Docstring for resize_image
+        
+        :param self: Description
+        :param image: Description
+        self._resize_state = lambda sf:   cv2.resize(np.array(sf, dtype=np.float32), (cfg.image_shape[0], cfg.image_shape[1]))  # resize state
+        """
+        import cv2
+        import numpy as _np
+        img = _np.array(image, dtype=_np.float32)
+        img = cv2.resize(img, (self._cfg.image_shape[0], self._cfg.image_shape[1]))
+        return img
+
+    def normalize_state(self, image):
+        """
+        Docstring for preprocess_state
+        
+        :param self: Description
+        :param image: Description
+        self._encode_state = lambda af:   ((af/(255.0)*2.0)-1.0) # encoder: take a float, output an integer
+        self._resize_state = lambda sf:   cv2.resize(np.array(sf, dtype=np.float32), (cfg.image_shape[0], cfg.image_shape[1]))  # resize state
+        """
+        # img = _np.array(image, dtype=_np.float32)
+        # img = cv2.resize(img, (self._cfg.image_shape[0], self._cfg.image_shape[1]))
+        enc = ((image / 255.0) * 2.0) - 1.0
+        # t = _torch.tensor(enc, dtype=_torch.float32, device=self._cfg.device)
+        return enc
+    
+    def preprocess_state(self, image):
+        img = self.resize_image(image)
+        img = self.normalize_state(img)
+        return img
+
+    def preprocess_goal_image(self, image):
+        return self.preprocess_state(image)
+
+    def encode_text_goal(self, goal, tokenizer=None, text_model=None):
+        import numpy as _np
+        import torch as _torch
+        if self._cfg.dataset.encode_with_t5:
+            if tokenizer is None or text_model is None:
+                raise ValueError("tokenizer and text_model must be provided when using T5 encoding")
+            # TODO:    
+            ## Provide the logic converting text goal to T5 embedding tensor
+            with _torch.no_grad():
+                # Tokenize the goal text
+                inputs = tokenizer(goal, return_tensors="pt", padding=True, truncation=True, max_length=self._cfg.max_block_size)
+                # Move inputs to the correct device
+                inputs = {k: v.to(self._cfg.device) for k, v in inputs.items()}
+                # Get embeddings from the text model (Encoder)
+                outputs = text_model(**inputs)
+                # Use last hidden state: (Batch, Seq_Len, Hidden_Dim)
+                embeddings = outputs.last_hidden_state
+                
+            return embeddings
+        else:
+            pad = " " * self._cfg.max_block_size
+            goal_ = goal[:self._cfg.max_block_size] + pad[len(goal):self._cfg.max_block_size]
+            try:
+                stoi = {c: i for i, c in enumerate(self._cfg.dataset.chars_list)}
+                ids = [stoi.get(c, 0) for c in goal_]
+            except Exception:
+                ids = [0] * self._cfg.max_block_size
+            return _torch.tensor(_np.expand_dims(_np.array(ids, dtype=_np.int64), axis=0), dtype=_torch.long, device=self._cfg.device)
+
+    def process_text_embedding_for_buffer(self, goal, tokenizer=None, text_model=None):
+        """
+        Process text goal embedding for storing in the circular buffer.
+        Returns a numpy array of shape (max_block_size, n_embd) without batch dimension.
+        """
+        import numpy as _np
+        if tokenizer is None or text_model is None:
+            raise ValueError("tokenizer and text_model must be provided when using T5 encoding")
+        
+        goal_ = _np.zeros((self._cfg.max_block_size, self._cfg.n_embd), dtype=_np.float32)
+        input_ids = tokenizer(goal, return_tensors="pt").input_ids
+        goal_t = text_model.encoder(input_ids).last_hidden_state.detach().cpu().numpy()
+        goal_[:len(goal_t[0]), :] = goal_t[0][:self._cfg.max_block_size]
+        return goal_
+
+    def decode_action(self, action_tensor):
+        
+        """
+        Docstring for decode_action
+        
+        :param self: Description
+        :param action_tensor: Description
+        self._decode_action = lambda binN: (binN * action_std) + action_mean  # Undo mapping to [-1, 1]
+        """
+        import torch as _torch
+        ## The action tensor is of shape (batch_size, action_dim * action_stacking) so we need to repeat the mean and std per action stacking
+        action_mean = _torch.tensor(np.repeat(self._cfg.dataset.action_mean, self._cfg.policy.action_stacking), dtype=action_tensor.dtype, device=action_tensor.device)
+        action_std = _torch.tensor(np.repeat(self._cfg.dataset.action_std, self._cfg.policy.action_stacking), dtype=action_tensor.dtype, device=action_tensor.device)
+        return (action_tensor * action_std) + action_mean
+    
+    def encode_action(self, action_float):
+        """
+        Docstring for encode_action
+        
+        :param self: Description
+        :param action_float: Description
+        self._encode_action = lambda af:   (af - action_mean)/(action_std) # encoder: take a float, output an integer
+        """
+        import torch as _torch
+        action_mean = _torch.tensor(self._cfg.dataset.action_mean, dtype=action_float.dtype, device=action_float.device)
+        action_std = _torch.tensor(self._cfg.dataset.action_std, dtype=action_float.dtype, device=action_float.device)
+        return (action_float - action_mean) / action_std
+
+
+@torch.no_grad()
+def estimate_loss(model, dataset):
+    out = {}
+    model.eval()
+    for split in ['train', 'val']:
+        losses = torch.zeros(model._cfg.eval_iters)
+        for k in range(model._cfg.eval_iters):
+            X, x_pose, x_goal, x_goal_img, Y = dataset.get_batch_grp(split, model._cfg, model._cfg.batch_size)
+            logits, loss = model(X, x_goal, x_goal_img, Y, pose=x_pose)
+            losses[k] = loss.item()
+        out[split] = losses.mean()
+    model.train()
+    return out