def __init__(
self,
action_space: gym.spaces.Discrete,
observation_space: SpaceDict,
rgb_uuid: Optional[str] = "rgb",
hidden_size=512,
num_rnn_layers=1,
rnn_type="GRU",
):
super().__init__(action_space=action_space,
observation_space=observation_space)
self._hidden_size = hidden_size
self.rgb_uuid = rgb_uuid
self.visual_encoder = SimpleCNN(
observation_space=observation_space,
output_size=hidden_size,
rgb_uuid=self.rgb_uuid,
depth_uuid=None,
)
self.state_encoder = RNNStateEncoder(
0 if self.is_blind else self.recurrent_hidden_state_size,
self._hidden_size,
num_layers=num_rnn_layers,
rnn_type=rnn_type,
)
self.actor_critic = LinearActorCriticHead(self._hidden_size,
action_space.n)
self.train()
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()
def __init__(
self,
action_space: gym.spaces.Discrete,
observation_space: SpaceDict,
rgb_uuid: Optional[str],
depth_uuid: Optional[str],
goal_sensor_uuid: str,
hidden_size=512,
embed_coordinates=False,
coordinate_embedding_dim=8,
coordinate_dims=2,
num_rnn_layers=1,
rnn_type="GRU",
):
super().__init__(action_space=action_space, observation_space=observation_space)
self.goal_sensor_uuid = goal_sensor_uuid
self._hidden_size = hidden_size
self.embed_coordinates = embed_coordinates
if self.embed_coordinates:
self.coorinate_embedding_size = coordinate_embedding_dim
else:
self.coorinate_embedding_size = coordinate_dims
self.sensor_fusion = False
if "rgb" in observation_space.spaces and "depth" in observation_space.spaces:
self.sensor_fuser = nn.Linear(hidden_size * 2, hidden_size)
self.sensor_fusion = True
self.visual_encoder = SimpleCNN(
observation_space=observation_space,
output_size=hidden_size,
rgb_uuid=rgb_uuid,
depth_uuid=depth_uuid,
)
self.state_encoder = RNNStateEncoder(
(0 if self.is_blind else self.recurrent_hidden_state_size)
+ self.coorinate_embedding_size,
self._hidden_size,
num_layers=num_rnn_layers,
rnn_type=rnn_type,
)
self.actor = LinearActorHead(self._hidden_size, action_space.n)
self.critic = LinearCriticHead(self._hidden_size)
if self.embed_coordinates:
self.coordinate_embedding = nn.Linear(
coordinate_dims, coordinate_embedding_dim
)
self.train()
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()
def __init__(
self,
action_space: gym.spaces.Discrete,
observation_space: SpaceDict,
goal_sensor_uuid: str,
resnet_preprocessor_uuid: str,
detector_preprocessor_uuid: str,
rnn_hidden_size=512,
goal_dims: int = 32,
max_dets: int = 3,
resnet_compressor_hidden_out_dims: Tuple[int, int] = (128, 32),
box_embedder_hidden_out_dims: Tuple[int, int] = (128, 32),
class_embedder_hidden_out_dims: Tuple[int, int] = (128, 32),
combiner_hidden_out_dims: Tuple[int, int] = (128, 32),
):
super().__init__(
action_space=action_space,
observation_space=observation_space,
)
self.hidden_size = rnn_hidden_size
self.goal_visual_encoder = ResnetFasterRCNNTensorsGoalEncoder(
self.observation_space,
goal_sensor_uuid,
resnet_preprocessor_uuid,
detector_preprocessor_uuid,
goal_dims,
max_dets,
resnet_compressor_hidden_out_dims,
box_embedder_hidden_out_dims,
class_embedder_hidden_out_dims,
combiner_hidden_out_dims,
)
self.state_encoder = RNNStateEncoder(
self.goal_visual_encoder.output_dims,
rnn_hidden_size,
)
self.actor_critic = LinearActorCriticHead(self.hidden_size,
action_space.n)
self.train()