diff --git a/.gitmodules b/.gitmodules index e69de29..67465c0 100644 --- a/.gitmodules +++ b/.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 diff --git a/goat_bench/config.py b/goat_bench/config.py index cf0a528..9b59f96 100644 --- a/goat_bench/config.py +++ b/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", diff --git a/goat_bench/models/clip_policy.py b/goat_bench/models/clip_policy.py index 2ec5bf5..56c8492 100644 --- a/goat_bench/models/clip_policy.py +++ b/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 diff --git a/goat_bench/models/encoders/croco2_vit_decoder.py b/goat_bench/models/encoders/croco2_vit_decoder.py new file mode 100644 index 0000000..ba579c6 --- /dev/null +++ b/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 diff --git a/goat_bench/models/encoders/croco2_vit_encoder.py b/goat_bench/models/encoders/croco2_vit_encoder.py new file mode 100644 index 0000000..569b2e2 --- /dev/null +++ b/goat_bench/models/encoders/croco2_vit_encoder.py @@ -0,0 +1,137 @@ +# 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 Block, PatchEmbed +from .croco.models.pos_embed import get_2d_sincos_pos_embed, RoPE2D + +class Croco2ViTEncoder(nn.Module): + """ + Croco2ViTEncoder: An encoder 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(Croco2ViTEncoder, 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 + enc_pos_embed = get_2d_sincos_pos_embed(enc_embed_dim, int(self.patch_embed.num_patches**0.5), n_cls_token=0) + self.register_buffer('enc_pos_embed', torch.from_numpy(enc_pos_embed).float()) + self.rope = None # Nothing for cosine + elif pos_embed.startswith('RoPE'): # e.g., RoPE100 + self.enc_pos_embed = None # Nothing to add in the encoder 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) + + # Transformer for the encoder + self.enc_depth = enc_depth + self.enc_embed_dim = enc_embed_dim + self.enc_blocks = nn.ModuleList([ + Block(enc_embed_dim, enc_num_heads, mlp_ratio, qkv_bias=True, norm_layer=norm_layer, rope=self.rope, + adapter=adapter, adapter_bottleneck=adapter_bottleneck, adapter_scalar=adapter_scalar, adapter_style=adapter_style) + for i in range(enc_depth)]) + self.enc_norm = norm_layer(enc_embed_dim) + + # Initialize weights + self.initialize_weights() + + 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 _encode_image(self, image, return_all_blocks=False): + """ + Encode the image into patches and apply masking if needed. + + Args: + image (Tensor): Input image with shape B x 3 x img_size x img_size. + return_all_blocks (bool): If True, return features at the end of every block instead of the last block. + + Returns: + Encoded features, positions, and masks (if masking is performed). + """ + # Embed the image into patches (x has size B x Npatches x C) + # and get the position of each returned patch (pos has size B x Npatches x 2) + x, pos = self.patch_embed(image) + # Add positional embedding without cls token + if self.enc_pos_embed is not None: + x = x + self.enc_pos_embed[None, ...] + posvis = pos + # Apply the transformer encoder and normalization + if return_all_blocks: + out = [] + for blk in self.enc_blocks: + x = blk(x, posvis) + out.append(x) + out[-1] = self.enc_norm(out[-1]) + return out, pos + else: + for blk in self.enc_blocks: + x = blk(x, posvis) + x = self.enc_norm(x) + return x, pos + + def forward(self, img1): + """ + Forward pass for Croco2ViTEncoder. + + Args: + img1 (Tensor): Input image with shape B x 3 x img_size x img_size. + + Returns: + Encoded features and positions. + """ + feat1, pos1 = self._encode_image(img1) + return feat1, pos1 diff --git a/goat_bench/models/encoders/croco_binocular_encoder.py b/goat_bench/models/encoders/croco_binocular_encoder.py new file mode 100644 index 0000000..d907c80 --- /dev/null +++ b/goat_bench/models/encoders/croco_binocular_encoder.py @@ -0,0 +1,81 @@ +import torch +from torch import nn as nn +from goat_bench.models.encoders.croco2_vit_encoder import Croco2ViTEncoder +from goat_bench.models.encoders.croco2_vit_decoder import Croco2ViTDecoder +from torchvision import transforms as T +from gym import spaces + + +class CrocoBinocularEncoder(nn.Module): + def __init__( + self, + observation_space: spaces.Dict, + checkpoint: str, + adapter: bool, + hidden_size: int, + ): + super().__init__() + + self.goal_imagenav = "image_goal_rotation" in observation_space.spaces + self.goal_instance_imagenav = "goat_instance_imagegoal" or "instance_imagegoal" in observation_space.spaces + self.goal_features = "cache_croco_goal_feat" in observation_space.spaces and "cache_croco_goal_pos" in observation_space.spaces + imagenet_mean = [0.485, 0.456, 0.406] + imagenet_std = [0.229, 0.224, 0.225] + self.preprocess = T.Compose([ + T.Resize(112, interpolation=T.InterpolationMode.BICUBIC), + T.CenterCrop(112), + T.ConvertImageDtype(torch.float), + T.Normalize(mean=imagenet_mean, std=imagenet_std) + ]) + ckpt = torch.load(checkpoint) + kwargs = ckpt.get('croco_kwargs', {}) + if 'img_size' in kwargs: + del kwargs['img_size'] + self.encoder = Croco2ViTEncoder(adapter=adapter, img_size=112, **kwargs) + self.decoder = Croco2ViTDecoder(adapter=adapter, img_size=112, **kwargs) + encoder_weights = {k:v for k,v in ckpt['model'].items() if k.startswith('enc') or k.startswith('patch')} + msg = self.encoder.load_state_dict(encoder_weights, strict=False) + print(f"Loading Croco encoder weights from {checkpoint}: {msg}") + decoder_weights = {k:v for k,v in ckpt['model'].items() if k.startswith('dec') or k.startswith('patch')} + msg = self.decoder.load_state_dict(decoder_weights, strict=False) + print(f"Loading Croco decoder weights from {checkpoint}: {msg}") + + print("Freezing Croco encoder parameters") + for name, param in self.encoder.named_parameters(): + if "adaptmlp" not in name: + param.requires_grad = False + self.encoder.eval() + print("Freezing Croco decoder parameters") + for name, param in self.decoder.named_parameters(): + if "adaptmlp" not in name: + param.requires_grad = False + self.decoder.eval() + + dec_embed_dim = self.decoder.dec_embed_dim + self.fc = nn.Sequential( + nn.Linear( + dec_embed_dim, + hidden_size, + ), + nn.Flatten() + ) + + def forward(self, observations) -> torch.Tensor: # type: ignore + rgb = observations["rgb"] + rgb = rgb.permute(0, 3, 1, 2) + rgb = self.preprocess(rgb) + rgb_feat, rgb_pos = self.encoder(rgb) + + if self.goal_instance_imagenav: + instance_imagegoal = observations["instance_imagegoal"] if "instance_imagegoal" in observations else observations["goat_instance_imagegoal"] + instance_imagegoal = instance_imagegoal.permute(0, 3, 1, 2) + instance_imagegoal = self.preprocess(instance_imagegoal) + goal_feat, goal_pos = self.encoder(instance_imagegoal) + + else: + raise NotImplementedError("Required observations not found.") + + x = self.decoder(rgb_feat, rgb_pos, goal_feat, goal_pos) + + x = self.fc(x) + return x diff --git a/goat_bench/models/high_level_policy.py b/goat_bench/models/high_level_policy.py index 0cd3421..d983183 100644 --- a/goat_bench/models/high_level_policy.py +++ b/goat_bench/models/high_level_policy.py @@ -7,6 +7,7 @@ from gym import spaces from gym.spaces import Dict as SpaceDict from habitat.sims.habitat_simulator.actions import HabitatSimActions +from habitat.tasks.nav.instance_image_nav_task import InstanceImageGoalSensor 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 @@ -23,7 +24,7 @@ from goat_bench.task.sensors import (CacheImageGoalSensor, ClipGoalSelectorSensor, ClipImageGoalSensor, ClipObjectGoalSensor, - LanguageGoalSensor) + LanguageGoalSensor, GoatInstanceImageGoalSensor) @baseline_registry.register_policy @@ -117,6 +118,8 @@ def __init__( not in [ LanguageGoalSensor.cls_uuid, CacheImageGoalSensor.cls_uuid, + InstanceImageGoalSensor.cls_uuid, + GoatInstanceImageGoalSensor.cls_uuid, ] } ) @@ -140,6 +143,8 @@ def __init__( not in [ ClipObjectGoalSensor.cls_uuid, CacheImageGoalSensor.cls_uuid, + InstanceImageGoalSensor.cls_uuid, + GoatInstanceImageGoalSensor.cls_uuid, ] } ) @@ -238,6 +243,8 @@ def forward( not in [ CacheImageGoalSensor.cls_uuid, LanguageGoalSensor.cls_uuid, + InstanceImageGoalSensor.cls_uuid, + GoatInstanceImageGoalSensor.cls_uuid, ] }, rnn_hidden_states, @@ -266,6 +273,8 @@ def forward( not in [ CacheImageGoalSensor.cls_uuid, ClipObjectGoalSensor.cls_uuid, + InstanceImageGoalSensor.cls_uuid, + GoatInstanceImageGoalSensor.cls_uuid, ] }, rnn_hidden_states, diff --git a/goat_bench/task/sensors.py b/goat_bench/task/sensors.py index aaa4d22..0203dff 100644 --- a/goat_bench/task/sensors.py +++ b/goat_bench/task/sensors.py @@ -4,12 +4,17 @@ from typing import TYPE_CHECKING, Any, Optional import numpy as np -from gym import spaces +from gym import spaces, Space from habitat.core.embodied_task import EmbodiedTask from habitat.core.registry import registry -from habitat.core.simulator import RGBSensor, Sensor, SensorTypes, Simulator +from habitat.core.simulator import RGBSensor, Sensor, SensorTypes, Simulator, VisualObservation from habitat.core.utils import try_cv2_import +from habitat.tasks.nav.instance_image_nav_task import InstanceImageParameters from habitat.tasks.nav.nav import NavigationEpisode +from habitat.utils.geometry_utils import quaternion_from_coeff +import habitat_sim +from habitat_sim import bindings as hsim +from habitat_sim.agent.agent import AgentState, SixDOFPose from goat_bench.task.goat_task import GoatEpisode @@ -820,3 +825,152 @@ def get_observation( else: raise NotImplementedError return output_embedding + + +@registry.register_sensor +class GoatInstanceImageGoalSensor(RGBSensor): + """A sensor for instance-based image goal specification used by the + InstanceImageGoal Navigation task. Image goals are rendered according to + camera parameters (resolution, HFOV, extrinsics) specified by the dataset. + + Args: + sim: a reference to the simulator for rendering instance image goals. + config: a config for the InstanceImageGoalSensor sensor. + dataset: a Instance Image Goal navigation dataset that contains a + dictionary mapping goal IDs to instance image goals. + """ + + cls_uuid: str = "goat_instance_imagegoal" + _current_image_goal: Optional[VisualObservation] + _current_episode_id: Optional[str] + + def __init__( + self, + sim, + config: "DictConfig", + dataset: Any, + *args: Any, + **kwargs: Any, + ): + self._dataset = dataset + self._sim = sim + super().__init__(config=config) + self._current_episode_id = None + self._current_image_goal = None + self.add_noise = config.add_noise + + def _get_uuid(self, *args: Any, **kwargs: Any) -> str: + return self.cls_uuid + + def _get_observation_space(self, *args: Any, **kwargs: Any) -> Space: + # goals = next(iter(self._dataset.goals.values())) + # logger.info(goals) + # H, W = ( + # goals.image_goals[0] + # .image_dimensions + # ) + return spaces.Box(low=0, high=255, shape=(512, 512, 3), dtype=np.uint8) + + def _add_sensor( + self, img_params: InstanceImageParameters, sensor_uuid: str + ) -> None: + spec = habitat_sim.CameraSensorSpec() + spec.uuid = sensor_uuid + spec.sensor_type = habitat_sim.SensorType.COLOR + spec.resolution = img_params.image_dimensions + spec.hfov = img_params.hfov + spec.sensor_subtype = habitat_sim.SensorSubType.PINHOLE + self._sim.add_sensor(spec) + + agent = self._sim.get_agent(0) + agent_state = agent.get_state() + agent.set_state( + AgentState( + position=agent_state.position, + rotation=agent_state.rotation, + sensor_states={ + **agent_state.sensor_states, + sensor_uuid: SixDOFPose( + position=np.array(img_params.position), + rotation=quaternion_from_coeff(img_params.rotation), + ), + }, + ), + infer_sensor_states=False, + ) + + def _remove_sensor(self, sensor_uuid: str) -> None: + agent = self._sim.get_agent(0) + del self._sim._sensors[sensor_uuid] + hsim.SensorFactory.delete_subtree_sensor(agent.scene_node, sensor_uuid) + del agent._sensors[sensor_uuid] + agent.agent_config.sensor_specifications = [ + s + for s in agent.agent_config.sensor_specifications + if s.uuid != sensor_uuid + ] + + def _get_instance_image_goal( + self, img_params: InstanceImageParameters + ) -> VisualObservation: + sensor_uuid = f"{self.cls_uuid}_sensor" + self._add_sensor(img_params, sensor_uuid) + + self._sim._sensors[sensor_uuid].draw_observation() + img = self._sim._sensors[sensor_uuid].get_observation()[:, :, :3] + + self._remove_sensor(sensor_uuid) + return img + + def apply_noise(self, image): + mean = 0 + std = random.uniform(0.1, 2.0) + image = image + np.random.normal( + loc=mean, scale=std, size=image.shape + ).astype(np.float32) + return image.astype(np.uint8) + + def get_observation( + self, + *args: Any, + episode: Any, + task: Any, + **kwargs: Any, + ) -> Optional[VisualObservation]: + + dummy_image = np.zeros((512, 512, 3), dtype=np.uint8) + if isinstance(episode, GoatEpisode): + if task.active_subtask_idx < len(episode.tasks): + if episode.tasks[task.active_subtask_idx][1] == "image": + current_task = episode.tasks[task.active_subtask_idx] + instance_id = current_task[2] + goal_image_id = current_task[-1] + goal = [ + g for g in episode.goals[task.active_subtask_idx] + if g["object_id"] == instance_id + ] + scene_id = episode.scene_id.split('/')[-1].split(".")[0] + + goal = [ + g + for g in episode.goals[task.active_subtask_idx] + if g["object_id"] == instance_id + ] + img_params = InstanceImageParameters(**goal[0]['image_goals'][goal_image_id]) + image_goal = self._get_instance_image_goal(img_params) + # plt.imsave(f'{scene_id}_{instance_id}_{goal_image_id}_orig.png', image_goal) + if self.add_noise: + # logger.info("Applying Noise") + image_goal = self.apply_noise(image_goal) + # plt.imsave(f'{scene_id}_{instance_id}_{goal_image_id}_noise.png', image_goal) + self._current_image_goal = image_goal + else: + self._current_image_goal = dummy_image + else: + self._current_image_goal = dummy_image + else: + + img_params = episode.goals[0].image_goals[episode.goal_image_id] + self._current_image_goal = self._get_instance_image_goal(img_params) + + return self._current_image_goal diff --git a/scripts/eval/1-goat-skill-chaining-croco.sh b/scripts/eval/1-goat-skill-chaining-croco.sh new file mode 100644 index 0000000..08341ec --- /dev/null +++ b/scripts/eval/1-goat-skill-chaining-croco.sh @@ -0,0 +1,51 @@ +#!/bin/bash +#SBATCH --job-name=goat +#SBATCH --output=slurm_logs/eval/goat-%j.out +#SBATCH --error=slurm_logs/eval/goat-%j.err +#SBATCH --gpus a40:1 +#SBATCH --nodes 1 +#SBATCH --cpus-per-task 10 +#SBATCH --ntasks-per-node 1 +#SBATCH --exclude=xaea-12,nestor,shakey,dave,voltron,deebot +#SBATCH --signal=USR1@100 +#SBATCH --requeue +#SBATCH --partition=cvmlp-lab +#SBATCH --qos=short + +export GLOG_minloglevel=2 +export HABITAT_SIM_LOG=quiet +export MAGNUM_LOG=quiet + +MAIN_ADDR=$(scontrol show hostnames "${SLURM_JOB_NODELIST}" | head -n 1) +export MAIN_ADDR + +DATA_PATH="data/datasets/goat_bench/hm3d/v1/" +eval_ckpt_path_dir="data/new_checkpoints/goat_bench/ver/skill_chain_croco/" +tensorboard_dir="tb/goat_bench/ver/skill_chain_croco/val_seen/" +split="val_seen" + +srun python -um goat_bench.run \ + --run-type eval \ + --exp-config config/experiments/ver_goat_skill_chain.yaml \ + habitat_baselines.num_environments=1 \ + habitat_baselines.trainer_name="goat_ppo" \ + habitat_baselines.rl.policy.name=GoatHighLevelPolicy \ + habitat_baselines.tensorboard_dir=$tensorboard_dir \ + habitat_baselines.eval_ckpt_path_dir=$eval_ckpt_path_dir \ + habitat_baselines.checkpoint_folder=$eval_ckpt_path_dir \ + habitat.dataset.data_path="${DATA_PATH}/${split}/${split}.json.gz" \ + +habitat/task/lab_sensors@habitat.task.lab_sensors.clip_objectgoal_sensor=clip_objectgoal_sensor \ + +habitat/task/lab_sensors@habitat.task.lab_sensors.language_goal_sensor=language_goal_sensor \ + +habitat/task/lab_sensors@habitat.task.lab_sensors.goat_instance_imagegoal_sensor=goat_instance_imagegoal_sensor \ + ~habitat.task.lab_sensors.goat_goal_sensor \ + habitat.task.lab_sensors.language_goal_sensor.cache=data/goat-assets/goal_cache/language_nav/${split}_bert_embedding.pkl \ + habitat_baselines.load_resume_state_config=False \ + habitat_baselines.eval.use_ckpt_config=False \ + habitat_baselines.eval.split=$split \ + habitat_baselines.eval.should_load_ckpt=False \ + habitat_baselines.should_load_agent_state=False \ + habitat_baselines.rl.policy.add_instance_linear_projection=True \ + habitat_baselines.rl.policy.croco_adapter=True \ + habitat_baselines.rl.policy.use_croco=True + +touch $checkpoint_counter diff --git a/scripts/train/1-croco-instance-imagenav-ver.sh b/scripts/train/1-croco-instance-imagenav-ver.sh new file mode 100644 index 0000000..3ec0b25 --- /dev/null +++ b/scripts/train/1-croco-instance-imagenav-ver.sh @@ -0,0 +1,41 @@ +#!/bin/bash +#SBATCH --job-name=goat +#SBATCH --output=slurm_logs/goat-croco-ver-%j.out +#SBATCH --error=slurm_logs/goat-croco-ver-%j.err +#SBATCH --gpus a40:4 +#SBATCH --nodes 1 +#SBATCH --cpus-per-task 10 +#SBATCH --ntasks-per-node 4 +#SBATCH --exclude=xaea-12,nestor,shakey,dave,voltron,deebot +#SBATCH --signal=USR1@100 +#SBATCH --requeue +#SBATCH --partition=cvmlp-lab +#SBATCH --qos=short + +export GLOG_minloglevel=2 +export HABITAT_SIM_LOG=quiet +export MAGNUM_LOG=quiet + +MAIN_ADDR=$(scontrol show hostnames "${SLURM_JOB_NODELIST}" | head -n 1) +export MAIN_ADDR + +TENSORBOARD_DIR="tb/iin/ver/resnetclip_rgb_croco_image/seed_1" +CHECKPOINT_DIR="data/new_checkpoints/iin/ver/resnetclip_rgb_croco_image/seed_1" +DATA_PATH="data/datasets/iin/hm3d/v2" + +srun python -um goat_bench.run \ + --run-type train \ + --exp-config config/experiments/ver_instance_imagenav.yaml \ + habitat_baselines.trainer_name="ver" \ + habitat_baselines.num_environments=16 \ + habitat_baselines.rl.policy.name=GOATPolicy \ + habitat_baselines.rl.ddppo.train_encoder=False \ + habitat_baselines.rl.ddppo.backbone=resnet50_clip_avgpool \ + habitat_baselines.tensorboard_dir=${TENSORBOARD_DIR} \ + habitat_baselines.checkpoint_folder=${CHECKPOINT_DIR} \ + habitat.dataset.data_path=${DATA_PATH}/train/train.json.gz \ + habitat.task.measurements.success.success_distance=0.25 \ + habitat.simulator.type="GOATSim-v0" \ + habitat_baselines.rl.policy.add_instance_linear_projection=True \ + habitat_baselines.rl.policy.croco_adapter=True \ + habitat_baselines.rl.policy.use_croco=True