text: """Baseline models for use in the object navigation task.
  
  Object navigation is currently available as a Task in AI2-THOR and
  Facebook's Habitat.
  """
  from typing import Tuple, Dict, Optional, cast
  
  import gym
  import torch
  import torch.nn as nn
  from gym.spaces.dict import Dict as SpaceDict
  
  from allenact.algorithms.onpolicy_sync.policy import (
      ActorCriticModel,
      LinearCriticHead,
      LinearActorHead,
      DistributionType,
      Memory,
      ObservationType,
  )
  from allenact.base_abstractions.distributions import CategoricalDistr
  from allenact.base_abstractions.misc import ActorCriticOutput
  from allenact.embodiedai.models.basic_models import SimpleCNN, RNNStateEncoder
  
  
  class ObjectNavBaselineActorCritic(ActorCriticModel[CategoricalDistr]):
      """Baseline recurrent actor critic model for object-navigation.
  
      # Attributes
      action_space : The space of actions available to the agent. Currently only discrete
          actions are allowed (so this space will always be of type `gym.spaces.Discrete`).
      observation_space : The observation space expected by the agent. This observation space
          should include (optionally) 'rgb' images and 'depth' images and is required to
          have a component corresponding to the goal `goal_sensor_uuid`.
      goal_sensor_uuid : The uuid of the sensor of the goal object. See `GoalObjectTypeThorSensor`
          as an example of such a sensor.
      hidden_size : The hidden size of the GRU RNN.
      object_type_embedding_dim: The dimensionality of the embedding corresponding to the goal
          object type.
      """
  
      def __init__(
          self,
          action_space: gym.spaces.Discrete,
          observation_space: SpaceDict,
          goal_sensor_uuid: str,
          rgb_uuid: Optional[str],
          depth_uuid: Optional[str],
          hidden_size=512,
          object_type_embedding_dim=8,
          trainable_masked_hidden_state: bool = False,
          num_rnn_layers=1,
          rnn_type="GRU",
      ):
          """Initializer.
  
          See class documentation for parameter definitions.
          """
          super().__init__(action_space=action_space, observation_space=observation_space)
  
          self.goal_sensor_uuid = goal_sensor_uuid
          self._n_object_types = self.observation_space.spaces[self.goal_sensor_uuid].n
          self._hidden_size = hidden_size
          self.object_type_embedding_size = object_type_embedding_dim
  
          self.visual_encoder = SimpleCNN(
              observation_space=self.observation_space,
              output_size=self._hidden_size,
              rgb_uuid=rgb_uuid,
              depth_uuid=depth_uuid,
          )
  
          self.state_encoder = RNNStateEncoder(
              (0 if self.is_blind else self._hidden_size) + object_type_embedding_dim,
              self._hidden_size,
              trainable_masked_hidden_state=trainable_masked_hidden_state,
              num_layers=num_rnn_layers,
              rnn_type=rnn_type,
          )
  
          self.actor = LinearActorHead(self._hidden_size, action_space.n)
          self.critic = LinearCriticHead(self._hidden_size)
  
          self.object_type_embedding = nn.Embedding(
              num_embeddings=self._n_object_types,
              embedding_dim=object_type_embedding_dim,
          )
  
          self.train()
  
      @property
      def recurrent_hidden_state_size(self) -> int:
          """The recurrent hidden state size of the model."""
          return self._hidden_size
  
      @property
      def is_blind(self) -> bool:
          """True if the model is blind (e.g. neither 'depth' or 'rgb' is an
          input observation type)."""
          return self.visual_encoder.is_blind
  
      @property
      def num_recurrent_layers(self) -> int:
          """Number of recurrent hidden layers."""
          return self.state_encoder.num_recurrent_layers
  
      def _recurrent_memory_specification(self):
          return dict(
              rnn=(
                  (
                      ("layer", self.num_recurrent_layers),
                      ("sampler", None),
                      ("hidden", self.recurrent_hidden_state_size),
                  ),
                  torch.float32,
              )
          )
  
      def get_object_type_encoding(
          self, observations: Dict[str, torch.FloatTensor]
      ) -> torch.FloatTensor:
          """Get the object type encoding from input batched observations."""
          # noinspection PyTypeChecker
          return self.object_type_embedding(  # type:ignore
              observations[self.goal_sensor_uuid].to(torch.int64)
          )
  
      def forward(  # type:ignore
          self,
          observations: ObservationType,
          memory: Memory,
          prev_actions: torch.Tensor,
          masks: torch.FloatTensor,
      ) -> Tuple[ActorCriticOutput[DistributionType], Optional[Memory]]:
          """Processes input batched observations to produce new actor and critic
          values. Processes input batched observations (along with prior hidden
          states, previous actions, and masks denoting which recurrent hidden
          states should be masked) and returns an `ActorCriticOutput` object
          containing the model's policy (distribution over actions) and
          evaluation of the current state (value).
  
          # Parameters
          observations : Batched input observations.
          memory : `Memory` containing the hidden states from initial timepoints.
          prev_actions : Tensor of previous actions taken.
          masks : Masks applied to hidden states. See `RNNStateEncoder`.
          # Returns
          Tuple of the `ActorCriticOutput` and recurrent hidden state.
          """
          target_encoding = self.get_object_type_encoding(
              cast(Dict[str, torch.FloatTensor], observations)
          )
          x = [target_encoding]
  
          if not self.is_blind:
              perception_embed = self.visual_encoder(observations)
              x = [perception_embed] + x
  
          x_cat = torch.cat(x, dim=-1)  # type: ignore
          x_out, rnn_hidden_states = self.state_encoder(
              x_cat, memory.tensor("rnn"), masks
          )
  
          return (
              ActorCriticOutput(
                  distributions=self.actor(x_out), values=self.critic(x_out), extras={}
              ),
              memory.set_tensor("rnn", rnn_hidden_states),
          )
  
  
  class ResnetTensorObjectNavActorCritic(ActorCriticModel[CategoricalDistr]):
      def __init__(
          self,
          action_space: gym.spaces.Discrete,
          observation_space: SpaceDict,
          goal_sensor_uuid: str,
          rgb_resnet_preprocessor_uuid: Optional[str],
          depth_resnet_preprocessor_uuid: Optional[str] = None,
          hidden_size: int = 512,
          goal_dims: int = 32,
          resnet_compressor_hidden_out_dims: Tuple[int, int] = (128, 32),
          combiner_hidden_out_dims: Tuple[int, int] = (128, 32),
          include_auxiliary_head: bool = False,
      ):
  
          super().__init__(
              action_space=action_space, observation_space=observation_space,
          )
  
          self._hidden_size = hidden_size
          self.include_auxiliary_head = include_auxiliary_head
          if (
              rgb_resnet_preprocessor_uuid is None
              or depth_resnet_preprocessor_uuid is None
          ):
              resnet_preprocessor_uuid = (
                  rgb_resnet_preprocessor_uuid
                  if rgb_resnet_preprocessor_uuid is not None
                  else depth_resnet_preprocessor_uuid
              )
              self.goal_visual_encoder = ResnetTensorGoalEncoder(
                  self.observation_space,
                  goal_sensor_uuid,
                  resnet_preprocessor_uuid,
                  goal_dims,
                  resnet_compressor_hidden_out_dims,
                  combiner_hidden_out_dims,
              )
          else:
              self.goal_visual_encoder = ResnetDualTensorGoalEncoder(  # type:ignore
                  self.observation_space,
                  goal_sensor_uuid,
                  rgb_resnet_preprocessor_uuid,
                  depth_resnet_preprocessor_uuid,
                  goal_dims,
                  resnet_compressor_hidden_out_dims,
                  combiner_hidden_out_dims,
              )
          self.state_encoder = RNNStateEncoder(
              self.goal_visual_encoder.output_dims, self._hidden_size,
          )
          self.actor = LinearActorHead(self._hidden_size, action_space.n)
          self.critic = LinearCriticHead(self._hidden_size)
          if self.include_auxiliary_head:
              self.auxiliary_actor = LinearActorHead(self._hidden_size, action_space.n)
  
          self.train()
  
      @property
      def recurrent_hidden_state_size(self) -> int:
          """The recurrent hidden state size of the model."""
          return self._hidden_size
  
      @property
      def is_blind(self) -> bool:
          """True if the model is blind (e.g. neither 'depth' or 'rgb' is an
          input observation type)."""
          return self.goal_visual_encoder.is_blind
  
      @property
      def num_recurrent_layers(self) -> int:
          """Number of recurrent hidden layers."""
          return self.state_encoder.num_recurrent_layers
  
      def _recurrent_memory_specification(self):
          return dict(
              rnn=(
                  (
                      ("layer", self.num_recurrent_layers),
                      ("sampler", None),
                      ("hidden", self.recurrent_hidden_state_size),
                  ),
                  torch.float32,
              )
          )
  
      def get_object_type_encoding(
          self, observations: Dict[str, torch.FloatTensor]
      ) -> torch.FloatTensor:
          """Get the object type encoding from input batched observations."""
          return self.goal_visual_encoder.get_object_type_encoding(observations)
  
      def forward(  # type:ignore
          self,
          observations: ObservationType,
          memory: Memory,
          prev_actions: torch.Tensor,
          masks: torch.FloatTensor,
      ) -> Tuple[ActorCriticOutput[DistributionType], Optional[Memory]]:
          x = self.goal_visual_encoder(observations)
          x, rnn_hidden_states = self.state_encoder(x, memory.tensor("rnn"), masks)
          return (
              ActorCriticOutput(
                  distributions=self.actor(x),
                  values=self.critic(x),
                  extras={"auxiliary_distributions": self.auxiliary_actor(x)}
                  if self.include_auxiliary_head
                  else {},
              ),
              memory.set_tensor("rnn", rnn_hidden_states),
          )
  
  
  class ResnetTensorGoalEncoder(nn.Module):
      def __init__(
          self,
          observation_spaces: SpaceDict,
          goal_sensor_uuid: str,
          resnet_preprocessor_uuid: str,
          class_dims: int = 32,
          resnet_compressor_hidden_out_dims: Tuple[int, int] = (128, 32),
          combiner_hidden_out_dims: Tuple[int, int] = (128, 32),
      ) -> None:
          super().__init__()
          self.goal_uuid = goal_sensor_uuid
          self.resnet_uuid = resnet_preprocessor_uuid
          self.class_dims = class_dims
          self.resnet_hid_out_dims = resnet_compressor_hidden_out_dims
          self.combine_hid_out_dims = combiner_hidden_out_dims
          self.embed_class = nn.Embedding(
              num_embeddings=observation_spaces.spaces[self.goal_uuid].n,
              embedding_dim=self.class_dims,
          )
          self.blind = self.resnet_uuid not in observation_spaces.spaces
          if not self.blind:
              self.resnet_tensor_shape = observation_spaces.spaces[self.resnet_uuid].shape
              self.resnet_compressor = nn.Sequential(
                  nn.Conv2d(self.resnet_tensor_shape[0], self.resnet_hid_out_dims[0], 1),
                  nn.ReLU(),
                  nn.Conv2d(*self.resnet_hid_out_dims[0:2], 1),
                  nn.ReLU(),
              )
              self.target_obs_combiner = nn.Sequential(
                  nn.Conv2d(
                      self.resnet_hid_out_dims[1] + self.class_dims,
                      self.combine_hid_out_dims[0],
                      1,
                  ),
                  nn.ReLU(),
                  nn.Conv2d(*self.combine_hid_out_dims[0:2], 1),
              )
  
      @property
      def is_blind(self):
          return self.blind
  
      @property
      def output_dims(self):
          if self.blind:
              return self.class_dims
          else:
              return (
                  self.combine_hid_out_dims[-1]
                  * self.resnet_tensor_shape[1]
                  * self.resnet_tensor_shape[2]
              )
  
      def get_object_type_encoding(
          self, observations: Dict[str, torch.FloatTensor]
      ) -> torch.FloatTensor:
          """Get the object type encoding from input batched observations."""
          return cast(
              torch.FloatTensor,
              self.embed_class(observations[self.goal_uuid].to(torch.int64)),
          )
  
      def compress_resnet(self, observations):
          return self.resnet_compressor(observations[self.resnet_uuid])
  
      def distribute_target(self, observations):
          target_emb = self.embed_class(observations[self.goal_uuid])
          return target_emb.view(-1, self.class_dims, 1, 1).expand(
              -1, -1, self.resnet_tensor_shape[-2], self.resnet_tensor_shape[-1]
          )
  
      def adapt_input(self, observations):
          resnet = observations[self.resnet_uuid]
  
          use_agent = False
          nagent = 1
  
          if len(resnet.shape) == 6:
              use_agent = True
              nstep, nsampler, nagent = resnet.shape[:3]
          else:
              nstep, nsampler = resnet.shape[:2]
  
          observations[self.resnet_uuid] = resnet.view(-1, *resnet.shape[-3:])
          observations[self.goal_uuid] = observations[self.goal_uuid].view(-1, 1)
  
          return observations, use_agent, nstep, nsampler, nagent
  
      @staticmethod
      def adapt_output(x, use_agent, nstep, nsampler, nagent):
          if use_agent:
              return x.view(nstep, nsampler, nagent, -1)
          return x.view(nstep, nsampler * nagent, -1)
  
      def forward(self, observations):
          observations, use_agent, nstep, nsampler, nagent = self.adapt_input(
              observations
          )
  
          if self.blind:
              return self.embed_class(observations[self.goal_uuid])
          embs = [
              self.compress_resnet(observations),
              self.distribute_target(observations),
          ]
          x = self.target_obs_combiner(torch.cat(embs, dim=1,))
          x = x.reshape(x.size(0), -1)  # flatten
  
          return self.adapt_output(x, use_agent, nstep, nsampler, nagent)
  
  
  class ResnetDualTensorGoalEncoder(nn.Module):
      def __init__(
          self,
          observation_spaces: SpaceDict,
          goal_sensor_uuid: str,
          rgb_resnet_preprocessor_uuid: str,
          depth_resnet_preprocessor_uuid: str,
          class_dims: int = 32,
          resnet_compressor_hidden_out_dims: Tuple[int, int] = (128, 32),
          combiner_hidden_out_dims: Tuple[int, int] = (128, 32),
      ) -> None:
          super().__init__()
          self.goal_uuid = goal_sensor_uuid
          self.rgb_resnet_uuid = rgb_resnet_preprocessor_uuid
          self.depth_resnet_uuid = depth_resnet_preprocessor_uuid
          self.class_dims = class_dims
          self.resnet_hid_out_dims = resnet_compressor_hidden_out_dims
          self.combine_hid_out_dims = combiner_hidden_out_dims
          self.embed_class = nn.Embedding(
              num_embeddings=observation_spaces.spaces[self.goal_uuid].n,
              embedding_dim=self.class_dims,
          )
          self.blind = (
              self.rgb_resnet_uuid not in observation_spaces.spaces
              or self.depth_resnet_uuid not in observation_spaces.spaces
          )
          if not self.blind:
              self.resnet_tensor_shape = observation_spaces.spaces[
                  self.rgb_resnet_uuid
              ].shape
              self.rgb_resnet_compressor = nn.Sequential(
                  nn.Conv2d(self.resnet_tensor_shape[0], self.resnet_hid_out_dims[0], 1),
                  nn.ReLU(),
                  nn.Conv2d(*self.resnet_hid_out_dims[0:2], 1),
                  nn.ReLU(),
              )
              self.depth_resnet_compressor = nn.Sequential(
                  nn.Conv2d(self.resnet_tensor_shape[0], self.resnet_hid_out_dims[0], 1),
                  nn.ReLU(),
                  nn.Conv2d(*self.resnet_hid_out_dims[0:2], 1),
                  nn.ReLU(),
              )
              self.rgb_target_obs_combiner = nn.Sequential(
                  nn.Conv2d(
                      self.resnet_hid_out_dims[1] + self.class_dims,
                      self.combine_hid_out_dims[0],
                      1,
                  ),
                  nn.ReLU(),
                  nn.Conv2d(*self.combine_hid_out_dims[0:2], 1),
              )
              self.depth_target_obs_combiner = nn.Sequential(
                  nn.Conv2d(
                      self.resnet_hid_out_dims[1] + self.class_dims,
                      self.combine_hid_out_dims[0],
                      1,
                  ),
                  nn.ReLU(),
                  nn.Conv2d(*self.combine_hid_out_dims[0:2], 1),
              )
  
      @property
      def is_blind(self):
          return self.blind
  
      @property
      def output_dims(self):
          if self.blind:
              return self.class_dims
          else:
              return (
                  2
                  * self.combine_hid_out_dims[-1]
                  * self.resnet_tensor_shape[1]
                  * self.resnet_tensor_shape[2]
              )
  
      def get_object_type_encoding(
          self, observations: Dict[str, torch.FloatTensor]
      ) -> torch.FloatTensor:
          """Get the object type encoding from input batched observations."""
          return cast(
              torch.FloatTensor,
              self.embed_class(observations[self.goal_uuid].to(torch.int64)),
          )
  
      def compress_rgb_resnet(self, observations):
          return self.rgb_resnet_compressor(observations[self.rgb_resnet_uuid])
  
      def compress_depth_resnet(self, observations):
          return self.depth_resnet_compressor(observations[self.depth_resnet_uuid])
  
      def distribute_target(self, observations):
          target_emb = self.embed_class(observations[self.goal_uuid])
          return target_emb.view(-1, self.class_dims, 1, 1).expand(
              -1, -1, self.resnet_tensor_shape[-2], self.resnet_tensor_shape[-1]
          )
  
      def adapt_input(self, observations):
          rgb = observations[self.rgb_resnet_uuid]
          depth = observations[self.depth_resnet_uuid]
  
          use_agent = False
          nagent = 1
  
          if len(rgb.shape) == 6:
              use_agent = True
              nstep, nsampler, nagent = rgb.shape[:3]
          else:
              nstep, nsampler = rgb.shape[:2]
  
          observations[self.rgb_resnet_uuid] = rgb.view(-1, *rgb.shape[-3:])
          observations[self.depth_resnet_uuid] = depth.view(-1, *depth.shape[-3:])
          observations[self.goal_uuid] = observations[self.goal_uuid].view(-1, 1)
  
          return observations, use_agent, nstep, nsampler, nagent
  
      @staticmethod
      def adapt_output(x, use_agent, nstep, nsampler, nagent):
          if use_agent:
              return x.view(nstep, nsampler, nagent, -1)
          return x.view(nstep, nsampler * nagent, -1)
  
      def forward(self, observations):
          observations, use_agent, nstep, nsampler, nagent = self.adapt_input(
              observations
          )
  
          if self.blind:
              return self.embed_class(observations[self.goal_uuid])
          rgb_embs = [
              self.compress_rgb_resnet(observations),
              self.distribute_target(observations),
          ]
          rgb_x = self.rgb_target_obs_combiner(torch.cat(rgb_embs, dim=1,))
          depth_embs = [
              self.compress_depth_resnet(observations),
              self.distribute_target(observations),
          ]
          depth_x = self.depth_target_obs_combiner(torch.cat(depth_embs, dim=1,))
          x = torch.cat([rgb_x, depth_x], dim=1)
          x = x.reshape(x.size(0), -1)  # flatten
  
          return self.adapt_output(x, use_agent, nstep, nsampler, nagent)
  
  class ResnetDualTensorGoalGestureEncoder(nn.Module):
      def __init__(
          self,
          observation_spaces: SpaceDict,
          goal_sensor_uuid: str,
          # rel_position_uuid: str,
          gesture_sensor_uuid: str,
          human_pose_uuid: str,
          rgb_resnet_preprocessor_uuid: str,
          depth_resnet_preprocessor_uuid: str,
          class_dims: int = 512,
          gesture_compressor_hidden_out_dim: int = 512,
          human_pose_hidden_out_dim: int = 512, 
          resnet_compressor_hidden_out_dims: Tuple[int, int] = (128, 32),
          combiner_hidden_out_dims: Tuple[int, int] = (128, 32),
      ) -> None:
          super().__init__()
          self.goal_uuid = goal_sensor_uuid
          # self.rel_position_uuid = rel_position_uuid
          self.gesture_uuid = gesture_sensor_uuid
          self.human_pose_uuid = human_pose_uuid
          self.rgb_resnet_uuid = rgb_resnet_preprocessor_uuid
          self.depth_resnet_uuid = depth_resnet_preprocessor_uuid
          self.class_dims = class_dims
          self.gesture_hid_out_dim = gesture_compressor_hidden_out_dim
          self.human_pose_hid_out_dim = human_pose_hidden_out_dim
          self.resnet_hid_out_dims = resnet_compressor_hidden_out_dims
          self.combine_hid_out_dims = combiner_hidden_out_dims
          self.embed_class = nn.Embedding(
              num_embeddings=observation_spaces.spaces[self.goal_uuid].n,
              embedding_dim=self.class_dims,
          )
          self.blind = (
              self.rgb_resnet_uuid not in observation_spaces.spaces
              or self.depth_resnet_uuid not in observation_spaces.spaces
          )
          if not self.blind:
              self.resnet_tensor_shape = observation_spaces.spaces[
                  self.rgb_resnet_uuid
              ].shape
              self.rgb_resnet_compressor = nn.Sequential(
                  nn.Conv2d(self.resnet_tensor_shape[0], self.resnet_hid_out_dims[0], 1),
                  nn.ReLU(),
                  nn.Conv2d(*self.resnet_hid_out_dims[0:2], 1),
                  nn.ReLU(),
                  nn.Flatten(),
                  nn.Linear(self.resnet_tensor_shape[1]*self.resnet_tensor_shape[2]*self.resnet_hid_out_dims[-1], 512),
                  nn.ReLU(),
              )
              self.depth_resnet_compressor = nn.Sequential(
                  nn.Conv2d(self.resnet_tensor_shape[0], self.resnet_hid_out_dims[0], 1),
                  nn.ReLU(),
                  nn.Conv2d(*self.resnet_hid_out_dims[0:2], 1),
                  nn.ReLU(),
                  nn.Flatten(),
                  nn.Linear(self.resnet_tensor_shape[1]*self.resnet_tensor_shape[2]*self.resnet_hid_out_dims[-1], 512),
                  nn.ReLU(),
              )
              self.rgb_target_obs_combiner = nn.Sequential(
                  nn.Conv2d(
                      self.resnet_hid_out_dims[1] + self.class_dims,
                      self.combine_hid_out_dims[0],
                      1,
                  ),
                  nn.ReLU(),
                  nn.Conv2d(*self.combine_hid_out_dims[0:2], 1),
              )
              self.depth_target_obs_combiner = nn.Sequential(
                  nn.Conv2d(
                      self.resnet_hid_out_dims[1] + self.class_dims,
                      self.combine_hid_out_dims[0],
                      1,
                  ),
                  nn.ReLU(),
                  nn.Conv2d(*self.combine_hid_out_dims[0:2], 1),
              )
              
              self.rgb_target_gesture_obs_combiner = nn.Sequential(
                  nn.Conv2d(
                      self.resnet_hid_out_dims[1] + self.class_dims + self.gesture_hid_out_dim + self.human_pose_hid_out_dim,
                      self.combine_hid_out_dims[0],
                      1,
                  ),
                  nn.ReLU(),
                  nn.Conv2d(*self.combine_hid_out_dims[0:2], 1),
              )
              self.depth_target_gesture_obs_combiner = nn.Sequential(
                  nn.Conv2d(
                      self.resnet_hid_out_dims[1] + self.class_dims + self.gesture_hid_out_dim + self.human_pose_hid_out_dim,
                      self.combine_hid_out_dims[0],
                      1,
                  ),
                  nn.ReLU(),
                  nn.Conv2d(*self.combine_hid_out_dims[0:2], 1),
              )
              
              # self.vision_target_obs_combiner_compressor = nn.Sequential(
              #     nn.Flatten(),
              #     nn.Linear(
              #         in_features=self.resnet_tensor_shape[1]*self.resnet_tensor_shape[2]*self.resnet_hid_out_dims[1],
              #         out_features=self.gesture_hid_out_dim,
              #     ),
              #     # nn.ReLU(),
              #     # nn.Linear(
              #     #     in_features=self.gesture_hid_out_dim*2,
              #     #     out_features=self.gesture_hid_out_dim,
              #     # )
              # )
              
          self.gesture_tensor_shape = observation_spaces.spaces[self.gesture_uuid].shape
          # self.gesture_compressor = nn.LSTM(
          #     input_size=self.gesture_tensor_shape[1], 
          #     hidden_size=self.gesture_hid_out_dim,
          #     num_layers=1,
          #     batch_first=True,
          # )
          self.gesture_compressor = nn.Sequential(
              nn.Flatten(),
              nn.Linear(
                  in_features=self.gesture_tensor_shape[0]*self.gesture_tensor_shape[1],
                  out_features=self.gesture_hid_out_dim,
              ), 
              nn.ReLU(),
              # nn.Linear(
              #     in_features=self.gesture_hid_out_dim*4,
              #     out_features=self.gesture_hid_out_dim,
              # ),          
          )
          self.human_pose_shape = observation_spaces.spaces[self.human_pose_uuid].shape
          self.human_pose_compressor = nn.Sequential(
              nn.Linear(
              in_features=self.human_pose_shape[0],
              out_features=self.human_pose_hid_out_dim,
              ),
              nn.ReLU(),
          )
          # self.rel_position_shape = observation_spaces.spaces[self.rel_position_uuid].shape
  
      @property
      def is_blind(self):
          return self.blind
  
      @property
      def output_dims(self):
          if self.blind:
              return self.class_dims
          else:
              return (
                  # 2
                  # * self.combine_hid_out_dims[-1]
                  # * self.resnet_tensor_shape[1]
                  # * self.resnet_tensor_shape[2]
                  2*512 + self.class_dims + self.gesture_hid_out_dim + self.human_pose_hid_out_dim
              )
  
      def get_object_type_encoding(
          self, observations: Dict[str, torch.FloatTensor]
      ) -> torch.FloatTensor:
          """Get the object type encoding from input batched observations."""
          return cast(
              torch.FloatTensor,
              self.embed_class(observations[self.goal_uuid].to(torch.int64)),
          )
          
      def get_gesture_encoding(
          self, observations: Dict[str, torch.FloatTensor]
      ) -> torch.FloatTensor:
          """Get the gesture encoding from input batched observations. """
  
          # _, (hidden_n, _) = self.gesture_compressor(observations[self.gesture_uuid].float())
          # return cast(
          #     torch.FloatTensor,
          #     hidden_n[0],
          # )
  
          return cast(
              torch.FloatTensor,
              self.gesture_compressor(observations[self.gesture_uuid].float()),
          )
  
      def compress_rgb_resnet(self, observations):
          return self.rgb_resnet_compressor(observations[self.rgb_resnet_uuid])
  
      def compress_depth_resnet(self, observations):
          return self.depth_resnet_compressor(observations[self.depth_resnet_uuid])
      
      def compress_gesture(self, observations):
          return self.get_gesture_encoding(observations)
      
      def compress_human_pose(self, observations):
          return self.human_pose_compressor(observations[self.human_pose_uuid].float())
      
      def compress_rel_position(self, observations):
          return self.human_pose_compressor(observations[self.rel_position_uuid].float())
  
      def distribute_target(self, observations):
          target_emb = self.embed_class(observations[self.goal_uuid])
          return target_emb.view(-1, self.class_dims, 1, 1).expand(
              -1, -1, self.resnet_tensor_shape[-2], self.resnet_tensor_shape[-1]
          )
          
      def distribute_gesture(self, observations):
          gesture_emb = self.gesture_compressor(observations[self.gesture_uuid].float())
          return gesture_emb.view(-1, self.gesture_hid_out_dim, 1, 1).expand(
              -1, -1, self.resnet_tensor_shape[-2], self.resnet_tensor_shape[-1]
          )
          
      def distribute_human_pose(self, observations):
          human_pose_emb = self.human_pose_compressor(observations[self.human_pose_uuid].float())
          return human_pose_emb.view(-1, self.human_pose_hid_out_dim, 1, 1).expand(
              -1, -1, self.resnet_tensor_shape[-2], self.resnet_tensor_shape[-1]
          )
  
      def distribute_rel_position(self, observations):
          ref_position_emb = self.human_pose_compressor(observations[self.rel_position_uuid].float())
          return ref_position_emb.view(-1, self.human_pose_hid_out_dim, 1, 1).expand(
              -1, -1, self.resnet_tensor_shape[-2], self.resnet_tensor_shape[-1]
          )
  
      def adapt_input(self, observations):
          rgb = observations[self.rgb_resnet_uuid]
          depth = observations[self.depth_resnet_uuid]
  
          use_agent = False
          nagent = 1
  
          if len(rgb.shape) == 6:
              use_agent = True
              nstep, nsampler, nagent = rgb.shape[:3]
          else:
              nstep, nsampler = rgb.shape[:2]
  
          observations[self.rgb_resnet_uuid] = rgb.view(-1, *rgb.shape[-3:])
          observations[self.depth_resnet_uuid] = depth.view(-1, *depth.shape[-3:])
          observations[self.goal_uuid] = observations[self.goal_uuid].view(-1, 1)
          observations[self.gesture_uuid] = observations[self.gesture_uuid].view(-1, self.gesture_tensor_shape[0], self.gesture_tensor_shape[1])
          observations[self.human_pose_uuid] = observations[self.human_pose_uuid].view(-1, self.human_pose_shape[0])
          # observations[self.rel_position_uuid] = observations[self.rel_position_uuid].view(-1, self.rel_position_shape[0])
  
          return observations, use_agent, nstep, nsampler, nagent
  
      @staticmethod
      def adapt_output(x, use_agent, nstep, nsampler, nagent):
          if use_agent:
              return x.view(nstep, nsampler, nagent, -1)
          return x.view(nstep, nsampler * nagent, -1)
  
      def forward(self, observations):
          observations, use_agent, nstep, nsampler, nagent = self.adapt_input(
              observations
          )
  
          if self.blind:
              return self.embed_class(observations[self.goal_uuid])
          #rgb_embs = [
          #    self.compress_rgb_resnet(observations),
          #    self.distribute_target(observations),
          #    # self.distribute_rel_position(observations),
          #    self.distribute_gesture(observations),
          #    self.distribute_human_pose(observations),
          #]
          # rgb_x = self.rgb_target_gesture_obs_combiner(torch.cat(rgb_embs, dim=1,))
          rgb_x = self.compress_rgb_resnet(observations)
          #depth_embs = [
          #    self.compress_depth_resnet(observations),
          #    self.distribute_target(observations),
          #    # self.distribute_rel_position(observations),
          #    self.distribute_gesture(observations),
          #    self.distribute_human_pose(observations),
          #]
          #depth_x = self.depth_target_gesture_obs_combiner(torch.cat(depth_embs, dim=1,))
          depth_x = self.compress_depth_resnet(observations)
          
          target_x = torch.squeeze(self.get_object_type_encoding(observations), 1)
  
          gesture_motion_x = self.compress_gesture(observations)
  
          gesture_pose_x = self.compress_human_pose(observations)
  
          x = torch.cat([rgb_x, depth_x, target_x, gesture_motion_x, gesture_pose_x], dim=1)
          x = x.reshape(x.size(0), -1)  # flatten 
  
          return self.adapt_output(x, use_agent, nstep, nsampler, nagent)
      
  class ResnetTensorObjectGestureNavActorCritic(ActorCriticModel[CategoricalDistr]):
      def __init__(
          self,
          action_space: gym.spaces.Discrete,
          observation_space: SpaceDict,
          goal_sensor_uuid: str,
          # rel_position_uuid: str,
          gesture_sensor_uuid: Optional[str],
          human_pose_uuid: Optional[str],
          rgb_resnet_preprocessor_uuid: Optional[str],
          depth_resnet_preprocessor_uuid: Optional[str] = None,
          hidden_size: int = 512,
          goal_dims: int = 32,
          gesture_compressor_hidden_out_dim: int = 32,
          human_pose_hidden_out_dim: int = 32,
          resnet_compressor_hidden_out_dims: Tuple[int, int] = (128, 32),
          combiner_hidden_out_dims: Tuple[int, int] = (128, 32),
          include_auxiliary_head: bool = False,
      ):
  
          super().__init__(
              action_space=action_space, observation_space=observation_space,
          )
  
          self._hidden_size = hidden_size
          self.include_auxiliary_head = include_auxiliary_head
          if (
              rgb_resnet_preprocessor_uuid is None
              or depth_resnet_preprocessor_uuid is None
          ):
              resnet_preprocessor_uuid = (
                  rgb_resnet_preprocessor_uuid
                  if rgb_resnet_preprocessor_uuid is not None
                  else depth_resnet_preprocessor_uuid
              )
              self.goal_visual_encoder = ResnetTensorGoalEncoder(
                  self.observation_space,
                  goal_sensor_uuid,
                  resnet_preprocessor_uuid,
                  goal_dims,
                  resnet_compressor_hidden_out_dims,
                  combiner_hidden_out_dims,
              )
          else:
              self.goal_visual_encoder = ResnetDualTensorGoalGestureEncoder(  # type:ignore
                  self.observation_space,
                  goal_sensor_uuid,
                  # rel_position_uuid,
                  gesture_sensor_uuid,
                  human_pose_uuid,
                  rgb_resnet_preprocessor_uuid,
                  depth_resnet_preprocessor_uuid,
                  goal_dims,
                  gesture_compressor_hidden_out_dim,
                  human_pose_hidden_out_dim,
                  resnet_compressor_hidden_out_dims,
                  combiner_hidden_out_dims,
              )
          self.state_encoder = RNNStateEncoder(
              self.goal_visual_encoder.output_dims, self._hidden_size,
          )
          self.actor = LinearActorHead(self._hidden_size, action_space.n)
          self.critic = LinearCriticHead(self._hidden_size)
          if self.include_auxiliary_head:
              self.auxiliary_actor = LinearActorHead(self._hidden_size, action_space.n)
  
          self.train()
          
      @property
      def recurrent_hidden_state_size(self) -> int:
          """The recurrent hidden state size of the model."""
          return self._hidden_size
  
      @property
      def is_blind(self) -> bool:
          """True if the model is blind (e.g. neither 'depth' or 'rgb' is an
          input observation type)."""
          return self.goal_visual_encoder.is_blind
  
      @property
      def num_recurrent_layers(self) -> int:
          """Number of recurrent hidden layers."""
          return self.state_encoder.num_recurrent_layers
  
      def _recurrent_memory_specification(self):
          return dict(
              rnn=(
                  (
                      ("layer", self.num_recurrent_layers),
                      ("sampler", None),
                      ("hidden", self.recurrent_hidden_state_size),
                  ),
                  torch.float32,
              )
          )
  
      def get_object_type_encoding(
          self, observations: Dict[str, torch.FloatTensor]
      ) -> torch.FloatTensor:
          """Get the object type encoding from input batched observations."""
          return self.goal_visual_encoder.get_object_type_encoding(observations)
  
      def forward(  # type:ignore
          self,
          observations: ObservationType,
          memory: Memory,
          prev_actions: torch.Tensor,
          masks: torch.FloatTensor,
      ) -> Tuple[ActorCriticOutput[DistributionType], Optional[Memory]]:
          x = self.goal_visual_encoder(observations)
          x, rnn_hidden_states = self.state_encoder(x, memory.tensor("rnn"), masks)
          return (
              ActorCriticOutput(
                  distributions=self.actor(x),
                  values=self.critic(x),
                  extras={"auxiliary_distributions": self.auxiliary_actor(x)}
                  if self.include_auxiliary_head
                  else {},
              ),
              memory.set_tensor("rnn", rnn_hidden_states),
          )