def __init__(self, specs: BehaviorSpec, settings: CuriositySettings) -> None: super().__init__() self._policy_specs = specs state_encoder_settings = NetworkSettings( normalize=False, hidden_units=settings.encoding_size, num_layers=2, vis_encode_type=EncoderType.SIMPLE, memory=None, ) self._state_encoder = NetworkBody(specs.observation_shapes, state_encoder_settings) self._action_flattener = ModelUtils.ActionFlattener(specs) self.inverse_model_action_prediction = torch.nn.Sequential( LinearEncoder(2 * settings.encoding_size, 1, 256), linear_layer(256, self._action_flattener.flattened_size), ) self.forward_model_next_state_prediction = torch.nn.Sequential( LinearEncoder( settings.encoding_size + self._action_flattener.flattened_size, 1, 256), linear_layer(256, settings.encoding_size), )
def __init__( self, x_self_size: int, entities_sizes: List[int], embedding_size: int, output_size: Optional[int] = None, ): super().__init__() self.self_size = x_self_size self.entities_sizes = entities_sizes self.entities_num_max_elements: Optional[List[int]] = None self.ent_encoders = torch.nn.ModuleList([ LinearEncoder(self.self_size + ent_size, 2, embedding_size) for ent_size in self.entities_sizes ]) self.attention = MultiHeadAttention( query_size=embedding_size, key_size=embedding_size, value_size=embedding_size, output_size=embedding_size, num_heads=4, embedding_size=embedding_size, ) self.residual_layer = LinearEncoder(embedding_size, 1, embedding_size) if output_size is None: output_size = embedding_size self.x_self_residual_layer = LinearEncoder( embedding_size + x_self_size, 1, output_size)
def __init__(self, specs: BehaviorSpec, settings: CuriositySettings) -> None: super().__init__() self._action_spec = specs.action_spec state_encoder_settings = NetworkSettings( normalize=False, hidden_units=settings.encoding_size, num_layers=2, vis_encode_type=EncoderType.SIMPLE, memory=None, ) self._state_encoder = NetworkBody(specs.observation_specs, state_encoder_settings) self._action_flattener = ActionFlattener(self._action_spec) self.inverse_model_action_encoding = torch.nn.Sequential( LinearEncoder(2 * settings.encoding_size, 1, 256)) if self._action_spec.continuous_size > 0: self.continuous_action_prediction = linear_layer( 256, self._action_spec.continuous_size) if self._action_spec.discrete_size > 0: self.discrete_action_prediction = linear_layer( 256, sum(self._action_spec.discrete_branches)) self.forward_model_next_state_prediction = torch.nn.Sequential( LinearEncoder( settings.encoding_size + self._action_flattener.flattened_size, 1, 256), linear_layer(256, settings.encoding_size), )
def __init__( self, observation_specs: List[ObservationSpec], network_settings: NetworkSettings, encoded_act_size: int = 0, ): super().__init__() self.normalize = network_settings.normalize self.use_lstm = network_settings.memory is not None self.h_size = network_settings.hidden_units self.m_size = (network_settings.memory.memory_size if network_settings.memory is not None else 0) self.processors, self.embedding_sizes = ModelUtils.create_input_processors( observation_specs, self.h_size, network_settings.vis_encode_type, normalize=self.normalize, ) entity_num_max: int = 0 var_processors = [ p for p in self.processors if isinstance(p, EntityEmbedding) ] for processor in var_processors: entity_max: int = processor.entity_num_max_elements # Only adds entity max if it was known at construction if entity_max > 0: entity_num_max += entity_max if len(var_processors) > 0: if sum(self.embedding_sizes): self.x_self_encoder = LinearEncoder( sum(self.embedding_sizes), 1, self.h_size, kernel_init=Initialization.Normal, kernel_gain=(0.125 / self.h_size)**0.5, ) self.rsa = ResidualSelfAttention(self.h_size, entity_num_max) total_enc_size = sum(self.embedding_sizes) + self.h_size else: total_enc_size = sum(self.embedding_sizes) total_enc_size += encoded_act_size self.linear_encoder = LinearEncoder(total_enc_size, network_settings.num_layers, self.h_size) if self.use_lstm: self.lstm = LSTM(self.h_size, self.m_size) else: self.lstm = None # type: ignore
def __init__( self, observation_shapes: List[Tuple[int, ...]], network_settings: NetworkSettings, encoded_act_size: int = 0, ): super().__init__() self.normalize = network_settings.normalize self.use_lstm = network_settings.memory is not None self.h_size = network_settings.hidden_units self.m_size = (network_settings.memory.memory_size if network_settings.memory is not None else 0) self.visual_processors, self.vector_processors, encoder_input_size = ModelUtils.create_input_processors( observation_shapes, self.h_size, network_settings.vis_encode_type, normalize=self.normalize, ) total_enc_size = encoder_input_size + encoded_act_size self.linear_encoder = LinearEncoder(total_enc_size, network_settings.num_layers, self.h_size) if self.use_lstm: self.lstm = LSTM(self.h_size, self.m_size) else: self.lstm = None # type: ignore
def __init__( self, entity_size: int, entity_num_max_elements: Optional[int], embedding_size: int, ): """ Constructs an EntityEmbedding module. :param x_self_size: Size of "self" entity. :param entity_size: Size of other entities. :param entity_num_max_elements: Maximum elements for a given entity, None for unrestricted. Needs to be assigned in order for model to be exportable to ONNX and Barracuda. :param embedding_size: Embedding size for the entity encoder. :param concat_self: Whether to concatenate x_self to entities. Set True for ego-centric self-attention. """ super().__init__() self.self_size: int = 0 self.entity_size: int = entity_size self.entity_num_max_elements: int = -1 if entity_num_max_elements is not None: self.entity_num_max_elements = entity_num_max_elements self.embedding_size = embedding_size # Initialization scheme from http://www.cs.toronto.edu/~mvolkovs/ICML2020_tfixup.pdf self.self_ent_encoder = LinearEncoder( self.entity_size, 1, self.embedding_size, kernel_init=Initialization.Normal, kernel_gain=(0.125 / self.embedding_size)**0.5, )
def __init__( self, sensor_specs: List[SensorSpec], network_settings: NetworkSettings, encoded_act_size: int = 0, ): super().__init__() self.normalize = network_settings.normalize self.use_lstm = network_settings.memory is not None self.h_size = network_settings.hidden_units self.m_size = (network_settings.memory.memory_size if network_settings.memory is not None else 0) self.processors, self.embedding_sizes = ModelUtils.create_input_processors( sensor_specs, self.h_size, network_settings.vis_encode_type, normalize=self.normalize, ) total_enc_size = sum(self.embedding_sizes) + encoded_act_size self.linear_encoder = LinearEncoder(total_enc_size, network_settings.num_layers, self.h_size) if self.use_lstm: self.lstm = LSTM(self.h_size, self.m_size) else: self.lstm = None # type: ignore
def __init__( self, observation_specs: List[ObservationSpec], network_settings: NetworkSettings, encoded_act_size: int = 0, ): super().__init__() self.normalize = network_settings.normalize self.use_lstm = network_settings.memory is not None self.h_size = network_settings.hidden_units self.m_size = (network_settings.memory.memory_size if network_settings.memory is not None else 0) self.observation_encoder = ObservationEncoder( observation_specs, self.h_size, network_settings.vis_encode_type, self.normalize, ) self.processors = self.observation_encoder.processors total_enc_size = self.observation_encoder.total_enc_size total_enc_size += encoded_act_size self.linear_encoder = LinearEncoder(total_enc_size, network_settings.num_layers, self.h_size) if self.use_lstm: self.lstm = LSTM(self.h_size, self.m_size) else: self.lstm = None # type: ignore
def create_residual_self_attention( input_processors: nn.ModuleList, embedding_sizes: List[int], hidden_size: int ) -> Tuple[Optional[ResidualSelfAttention], Optional[LinearEncoder]]: """ Creates an RSA if there are variable length observations found in the input processors. :param input_processors: A ModuleList of input processors as returned by the function create_input_processors(). :param embedding sizes: A List of embedding sizes as returned by create_input_processors(). :param hidden_size: The hidden size to use for the RSA. :returns: A Tuple of the RSA itself, a self encoder, and the embedding size after the RSA. Returns None for the RSA and encoder if no var len inputs are detected. """ rsa, x_self_encoder = None, None entity_num_max: int = 0 var_processors = [ p for p in input_processors if isinstance(p, EntityEmbedding) ] for processor in var_processors: entity_max: int = processor.entity_num_max_elements # Only adds entity max if it was known at construction if entity_max > 0: entity_num_max += entity_max if len(var_processors) > 0: if sum(embedding_sizes): x_self_encoder = LinearEncoder( sum(embedding_sizes), 1, hidden_size, kernel_init=Initialization.Normal, kernel_gain=(0.125 / hidden_size)**0.5, ) rsa = ResidualSelfAttention(hidden_size, entity_num_max) return rsa, x_self_encoder
def test_predict_minimum_training(): # of 5 numbers, predict index of min np.random.seed(1336) torch.manual_seed(1336) n_k = 5 size = n_k + 1 embedding_size = 64 entity_embeddings = EntityEmbeddings(size, [size], embedding_size, [n_k], concat_self=False) transformer = ResidualSelfAttention(embedding_size) l_layer = LinearEncoder(embedding_size, 2, n_k) loss = torch.nn.CrossEntropyLoss() optimizer = torch.optim.Adam( list(entity_embeddings.parameters()) + list(transformer.parameters()) + list(l_layer.parameters()), lr=0.001, weight_decay=1e-6, ) batch_size = 200 onehots = ModelUtils.actions_to_onehot( torch.range(0, n_k - 1).unsqueeze(1), [n_k])[0] onehots = onehots.expand((batch_size, -1, -1)) losses = [] for _ in range(400): num = np.random.randint(0, n_k) inp = torch.rand((batch_size, num + 1, 1)) with torch.no_grad(): # create the target : The minimum argmin = torch.argmin(inp, dim=1) argmin = argmin.squeeze() argmin = argmin.detach() sliced_oh = onehots[:, :num + 1] inp = torch.cat([inp, sliced_oh], dim=2) embeddings = entity_embeddings(inp, [inp]) masks = EntityEmbeddings.get_masks([inp]) prediction = transformer(embeddings, masks) prediction = l_layer(prediction) ce = loss(prediction, argmin) losses.append(ce.item()) print(ce.item()) optimizer.zero_grad() ce.backward() optimizer.step() assert np.array(losses[-20:]).mean() < 0.1
def add_self_embedding(self, size: int) -> None: self.self_size = size self.self_ent_encoder = LinearEncoder( self.self_size + self.entity_size, 1, self.embedding_size, kernel_init=Initialization.Normal, kernel_gain=(0.125 / self.embedding_size)**0.5, )
def __init__(self, specs: BehaviorSpec, settings: CuriositySettings) -> None: super().__init__() self._action_spec = specs.action_spec state_encoder_settings = settings.network_settings if state_encoder_settings.memory is not None: state_encoder_settings.memory = None logger.warning( "memory was specified in network_settings but is not supported by Curiosity. It is being ignored." ) self._state_encoder = NetworkBody(specs.observation_specs, state_encoder_settings) self._action_flattener = ActionFlattener(self._action_spec) self.inverse_model_action_encoding = torch.nn.Sequential( LinearEncoder(2 * state_encoder_settings.hidden_units, 1, 256)) if self._action_spec.continuous_size > 0: self.continuous_action_prediction = linear_layer( 256, self._action_spec.continuous_size) if self._action_spec.discrete_size > 0: self.discrete_action_prediction = linear_layer( 256, sum(self._action_spec.discrete_branches)) self.forward_model_next_state_prediction = torch.nn.Sequential( LinearEncoder( state_encoder_settings.hidden_units + self._action_flattener.flattened_size, 1, 256, ), linear_layer(256, state_encoder_settings.hidden_units), )
def __init__( self, observation_specs: List[ObservationSpec], network_settings: NetworkSettings, action_spec: ActionSpec, ): super().__init__() self.normalize = network_settings.normalize self.use_lstm = network_settings.memory is not None self.h_size = network_settings.hidden_units self.m_size = (network_settings.memory.memory_size if network_settings.memory is not None else 0) self.action_spec = action_spec self.observation_encoder = ObservationEncoder( observation_specs, self.h_size, network_settings.vis_encode_type, self.normalize, ) self.processors = self.observation_encoder.processors # Modules for multi-agent self-attention obs_only_ent_size = self.observation_encoder.total_enc_size q_ent_size = (obs_only_ent_size + sum(self.action_spec.discrete_branches) + self.action_spec.continuous_size) attention_embeding_size = self.h_size self.obs_encoder = EntityEmbedding(obs_only_ent_size, None, attention_embeding_size) self.obs_action_encoder = EntityEmbedding(q_ent_size, None, attention_embeding_size) self.self_attn = ResidualSelfAttention(attention_embeding_size) self.linear_encoder = LinearEncoder( attention_embeding_size, network_settings.num_layers, self.h_size, kernel_gain=(0.125 / self.h_size)**0.5, ) if self.use_lstm: self.lstm = LSTM(self.h_size, self.m_size) else: self.lstm = None # type: ignore self._current_max_agents = torch.nn.Parameter(torch.as_tensor(1), requires_grad=False)
def __init__( self, x_self_size: int, entity_sizes: List[int], entity_num_max_elements: List[int], embedding_size: int, concat_self: bool = True, ): super().__init__() self.self_size: int = x_self_size self.entity_sizes: List[int] = entity_sizes self.entity_num_max_elements: List[int] = entity_num_max_elements self.concat_self: bool = concat_self # If not concatenating self, input to encoder is just entity size if not concat_self: self.self_size = 0 self.ent_encoders = torch.nn.ModuleList([ LinearEncoder(self.self_size + ent_size, 1, embedding_size) for ent_size in self.entity_sizes ])
def __init__( self, x_self_size: int, entity_sizes: List[int], embedding_size: int, entity_num_max_elements: Optional[List[int]] = None, concat_self: bool = True, ): """ Constructs an EntityEmbeddings module. :param x_self_size: Size of "self" entity. :param entity_sizes: List of sizes for other entities. Should be of length equivalent to the number of entities. :param embedding_size: Embedding size for entity encoders. :param entity_num_max_elements: Maximum elements in an entity, None for unrestricted. Needs to be assigned in order for model to be exportable to ONNX and Barracuda. :param concat_self: Whether to concatenate x_self to entites. Set True for ego-centric self-attention. """ super().__init__() self.self_size: int = x_self_size self.entity_sizes: List[int] = entity_sizes self.entity_num_max_elements: List[int] = [-1] * len(entity_sizes) if entity_num_max_elements is not None: self.entity_num_max_elements = entity_num_max_elements self.concat_self: bool = concat_self # If not concatenating self, input to encoder is just entity size if not concat_self: self.self_size = 0 # Initialization scheme from http://www.cs.toronto.edu/~mvolkovs/ICML2020_tfixup.pdf self.ent_encoders = torch.nn.ModuleList([ LinearEncoder( self.self_size + ent_size, 1, embedding_size, kernel_init=Initialization.Normal, kernel_gain=(0.125 / embedding_size)**0.5, ) for ent_size in self.entity_sizes ]) self.embedding_norm = LayerNorm()