Add croco encoder and related scripts

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "goat_bench/models/encoders/croco"]
+	path = goat_bench/models/encoders/croco
+	url = git@github.com:gchhablani/croco-with-adapter.git

goat_bench/config.py

@@ -31,6 +31,12 @@ class ClipObjectGoalSensorConfig(LabSensorConfig):
     )
 
 
+@dataclass
+class GoatInstanceImageGoalSensorConfig(LabSensorConfig):
+    type: str = "GoatInstanceImageGoalSensor"
+    add_noise: bool = False
+
+
 @dataclass
 class GoatGoalSensorConfig(LabSensorConfig):
     type: str = "GoatGoalSensor"
@@ -204,6 +210,11 @@ class GOATPolicyConfig(PolicyConfig):
 
     clip_model: str = "RN50"
     add_clip_linear_projection: bool = False
+    add_instance_linear_projection: bool = False
+    croco_adapter: bool = False
+    use_croco: bool = False
+    croco_ckpt: str = "goat_bench/models/encoders/croco/pretrained_models/CroCo_V2_ViTBase_SmallDecoder.pth"
+    use_hfov: bool = False
     depth_ckpt: str = ""
     late_fusion: bool = False
 
@@ -257,6 +268,13 @@ class OVONHabitatConfig(HabitatConfig):
     node=ClipObjectGoalSensorConfig,
 )
 
+cs.store(
+    package=f"habitat.task.lab_sensors.goat_instance_imagegoal_sensor",
+    group="habitat/task/lab_sensors",
+    name="goat_instance_imagegoal_sensor",
+    node=GoatInstanceImageGoalSensorConfig,
+)
+
 cs.store(
     package=f"habitat.task.lab_sensors.goat_goal_sensor",
     group="habitat/task/lab_sensors",

goat_bench/models/clip_policy.py

@@ -7,6 +7,7 @@
 from gym import spaces
 from gym.spaces import Dict as SpaceDict
 from habitat import logger
+from habitat.tasks.nav.instance_image_nav_task import InstanceImageGoalSensor, InstanceImageGoalHFOVSensor
 from habitat.tasks.nav.nav import EpisodicCompassSensor, EpisodicGPSSensor
 from habitat.tasks.nav.object_nav_task import ObjectGoalSensor
 from habitat_baselines.common.baseline_registry import baseline_registry
@@ -27,9 +28,11 @@
     ClipImageGoalSensor,
     ClipObjectGoalSensor,
     GoatGoalSensor,
+    GoatInstanceImageGoalSensor,
     GoatMultiGoalSensor,
     LanguageGoalSensor,
 )
+from goat_bench.models.encoders.croco_binocular_encoder import CrocoBinocularEncoder
 
 
 @baseline_registry.register_policy(name="GOATPolicy")
@@ -45,6 +48,11 @@ def __init__(
         policy_config: "DictConfig" = None,
         aux_loss_config: Optional["DictConfig"] = None,
         add_clip_linear_projection: bool = False,
+        add_instance_linear_projection: bool = False,
+        croco_adapter: bool = False,
+        use_croco: bool = False,
+        croco_ckpt: str = None,
+        use_hfov: bool = False,
         depth_ckpt: str = "",
         late_fusion: bool = False,
         **kwargs,
@@ -70,6 +78,11 @@ def __init__(
                 backbone=backbone,
                 discrete_actions=discrete_actions,
                 add_clip_linear_projection=add_clip_linear_projection,
+                add_instance_linear_projection=add_instance_linear_projection,
+                croco_adapter=croco_adapter,
+                use_croco=use_croco,
+                croco_ckpt=croco_ckpt,
+                use_hfov=use_hfov,
                 depth_ckpt=depth_ckpt,
                 late_fusion=late_fusion,
             ),
@@ -173,6 +186,10 @@ def __init__(
         add_clip_linear_projection: bool = False,
         add_language_linear_projection: bool = False,
         add_instance_linear_projection: bool = False,
+        croco_adapter: bool = False,
+        use_croco: bool = False,
+        croco_ckpt: str = None,
+        use_hfov: bool = False,
         late_fusion: bool = False,
     ):
         super().__init__()
@@ -181,6 +198,10 @@ def __init__(
         self.add_clip_linear_projection = add_clip_linear_projection
         self.add_language_linear_projection = add_language_linear_projection
         self.add_instance_linear_projection = add_instance_linear_projection
+        self.croco_adapter = croco_adapter
+        self.use_croco = use_croco
+        self.croco_ckpt = croco_ckpt
+        self.use_hfov = use_hfov
         self.late_fusion = late_fusion
         self._n_prev_action = 32
         if discrete_actions:
@@ -323,6 +344,43 @@ def __init__(
             rnn_input_size += instance_goal_size
             rnn_input_size_info["instance_goal"] = instance_goal_size
 
+        if self.use_croco and (
+            InstanceImageGoalSensor.cls_uuid in observation_space.spaces or
+            GoatInstanceImageGoalSensor.cls_uuid in observation_space.spaces
+        ):
+            self.croco_binocular_encoder = CrocoBinocularEncoder(
+                observation_space=observation_space,
+                checkpoint=self.croco_ckpt,
+                adapter=self.croco_adapter,
+                hidden_size=64, # NOTE: Total will be 49 * 64 per goal
+            )
+            self.croco_binocular_encoder.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
+            embedding_dim = 3136 # NOTE: Assuming the base decoder variant for now and FC is inside binocular encoder
+            print(
+                f"Binocular encoder embedding: {embedding_dim}, "
+                f"Add Instance linear: {add_instance_linear_projection}"
+            )
+            if self.add_instance_linear_projection:
+                self.instance_embedding = nn.Linear(embedding_dim, 256)
+                instance_goal_size = 256
+            else:
+                instance_goal_size = embedding_dim
+
+            rnn_input_size += instance_goal_size
+            rnn_input_size_info["instance_goal"] = instance_goal_size
+
+        if self.use_hfov and InstanceImageGoalHFOVSensor.cls_uuid in observation_space.spaces:
+            hfov_input_dim = int(
+                observation_space.spaces[InstanceImageGoalHFOVSensor.cls_uuid].high[0]
+            ) + 1
+            hfov_embedding_size = 32
+            self.hfov_embedding = nn.Embedding(
+                hfov_input_dim, hfov_embedding_size
+            )
+
+            rnn_input_size += hfov_embedding_size
+            rnn_input_size_info["hfov_embedding"] = hfov_embedding_size
+
         if EpisodicGPSSensor.cls_uuid in observation_space.spaces:
             input_gps_dim = observation_space.spaces[
                 EpisodicGPSSensor.cls_uuid
@@ -455,6 +513,19 @@ def forward(
                 instance_goal = self.instance_embedding(instance_goal)
             x.append(instance_goal)
 
+        if self.use_croco and (
+            InstanceImageGoalSensor.cls_uuid in observations or
+            GoatInstanceImageGoalSensor.cls_uuid in observations
+        ):
+            instance_goal = self.croco_binocular_encoder(observations)
+            if self.add_instance_linear_projection:
+                instance_goal = self.instance_embedding(instance_goal)
+            x.append(instance_goal)
+
+        if self.use_hfov and InstanceImageGoalHFOVSensor.cls_uuid in observations:
+            instance_goal_hfov = observations[InstanceImageGoalHFOVSensor.cls_uuid].long()
+            x.append(self.hfov_embedding(instance_goal_hfov).squeeze(dim=1))
+
         if (
             ClipImageGoalSensor.cls_uuid in observations
             and not self.late_fusion

goat_bench/models/encoders/croco2_vit_decoder.py

@@ -0,0 +1,155 @@
+# Copyright (C) 2022-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+
+import torch
+import torch.nn as nn
+torch.backends.cuda.matmul.allow_tf32 = True  # for GPU >= Ampere and PyTorch >= 1.12
+from functools import partial
+
+from .croco.models.blocks import DecoderBlock, PatchEmbed
+from .croco.models.pos_embed import get_2d_sincos_pos_embed, RoPE2D
+
+class Croco2ViTDecoder(nn.Module):
+    """
+    Croco2ViTDecoder: A decoder module for the CroCo model.
+
+    Args:
+        img_size (int): Input image size.
+        patch_size (int): Patch size.
+        mask_ratio (float): Ratio of masked tokens.
+        enc_embed_dim (int): Encoder feature dimension.
+        enc_depth (int): Encoder depth.
+        enc_num_heads (int): Number of heads in the encoder's transformer block.
+        dec_embed_dim (int): Decoder feature dimension.
+        dec_depth (int): Decoder depth.
+        dec_num_heads (int): Number of heads in the decoder's transformer block.
+        mlp_ratio (int): Ratio for MLP layers.
+        norm_layer (callable): Normalization layer constructor.
+        norm_im2_in_dec (bool): Whether to apply normalization to the 'memory' (second image) in the decoder.
+        pos_embed (str): Type of positional embedding (either 'cosine' or 'RoPE100').
+    """
+
+    def __init__(self, img_size=224, patch_size=16, mask_ratio=0.9, enc_embed_dim=768,
+                 enc_depth=12, enc_num_heads=12, dec_embed_dim=512, dec_depth=8,
+                 dec_num_heads=16, mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6),
+                 norm_im2_in_dec=True, pos_embed='cosine',
+                 adapter=False, adapter_bottleneck=64, adapter_scalar='0.1', adapter_style='parallel'):
+
+        super(Croco2ViTDecoder, self).__init__()
+
+        # Patch embeddings (with initialization done as in MAE)
+        self.patch_embed = PatchEmbed(img_size, patch_size, 3, enc_embed_dim)
+
+        self.pos_embed = pos_embed
+        if pos_embed == 'cosine':
+            # Positional embedding of the encoder
+            dec_pos_embed = get_2d_sincos_pos_embed(dec_embed_dim, int(self.patch_embed.num_patches**.5), n_cls_token=0)
+            self.register_buffer('dec_pos_embed', torch.from_numpy(dec_pos_embed).float())
+            # Position embedding in each block
+            self.rope = None  # Nothing for cosine
+        elif pos_embed.startswith('RoPE'):  # e.g., RoPE100
+            self.dec_pos_embed = None  # Nothing to add in the decoder with RoPE
+            if RoPE2D is None:
+                raise ImportError("Cannot find cuRoPE2D, please install it following the README instructions")
+            freq = float(pos_embed[len('RoPE'):])
+            self.rope = RoPE2D(freq=freq)
+        else:
+            raise NotImplementedError('Unknown pos_embed ' + pos_embed)
+
+        # Decoder
+        self._set_decoder(enc_embed_dim, dec_embed_dim, dec_num_heads, dec_depth, mlp_ratio, norm_layer, norm_im2_in_dec, adapter, adapter_bottleneck, adapter_scalar, adapter_style)
+
+        # Initialize weights
+        self.initialize_weights()
+
+    def _set_decoder(self, enc_embed_dim, dec_embed_dim, dec_num_heads, dec_depth, mlp_ratio, norm_layer, norm_im2_in_dec, adapter, adapter_bottleneck, adapter_scalar, adapter_style):
+        self.dec_depth = dec_depth
+        self.dec_embed_dim = dec_embed_dim
+        # Transfer from encoder to decoder
+        self.decoder_embed = nn.Linear(enc_embed_dim, dec_embed_dim, bias=True)
+        # Transformer for the decoder
+        self.dec_blocks = nn.ModuleList([
+            DecoderBlock(dec_embed_dim, dec_num_heads, mlp_ratio=mlp_ratio, qkv_bias=True, norm_layer=norm_layer, norm_mem=norm_im2_in_dec, rope=self.rope,
+                         adapter=adapter, adapter_bottleneck=adapter_bottleneck, adapter_scalar=adapter_scalar, adapter_style=adapter_style)
+            for i in range(dec_depth)])
+        # Final norm layer
+        self.dec_norm = norm_layer(dec_embed_dim)
+
+    def initialize_weights(self):
+        # patch embed 
+        self.patch_embed._init_weights()
+        # linears and layer norms
+        # self.apply(self._init_weights)
+        self.recursive_apply(self)
+
+    def recursive_apply(self, m):
+        for name, m2 in m.named_children():
+            if "adaptmlp" in name:
+                return
+            else:
+                self.recursive_apply(m2)
+        self._init_weights(m)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            # we use xavier_uniform following official JAX ViT:
+            torch.nn.init.xavier_uniform_(m.weight)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def _decoder(self, feat1, pos1, feat2, pos2, return_all_blocks=False):
+        """
+        Decoder function for Croco2ViTDecoder.
+
+        Args:
+            feat1: Features from the first image.
+            pos1: Position information for the first image.
+            feat2: Features from the second image.
+            pos2: Position information for the second image.
+            return_all_blocks: If True, return the features at the end of every block instead of just the features from the last block.
+
+        Returns:
+            Output features from the decoder.
+        """
+        # Encoder to decoder layer
+        visf1 = self.decoder_embed(feat1)
+        f2 = self.decoder_embed(feat2)
+        f1_ = visf1
+        # Add positional embedding
+        if self.dec_pos_embed is not None:
+            f1_ = f1_ + self.dec_pos_embed
+            f2 = f2 + self.dec_pos_embed
+        # Apply Transformer blocks
+        out = f1_
+        out2 = f2
+        if return_all_blocks:
+            _out, out = out, []
+            for blk in self.dec_blocks:
+                _out, out2 = blk(_out, out2, pos1, pos2)
+                out.append(_out)
+            out[-1] = self.dec_norm(out[-1])
+        else:
+            for blk in self.dec_blocks:
+                out, out2 = blk(out, out2, pos1, pos2)
+            out = self.dec_norm(out)
+        return out
+
+    def forward(self, feat1, pos1, feat2, pos2):
+        """
+        Forward pass for Croco2ViTDecoder.
+
+        Args:
+            feat1: Encoder features from the first image.
+            pos1: Position information for the first image.
+            feat2: Features from the second image.
+            pos2: Position information for the second image.
+
+        Returns:
+            Decoded features.
+        """
+        # Decoder
+        decfeat = self._decoder(feat1, pos1, feat2, pos2)
+        return decfeat