class BatchApply(Component): """ Takes an input with batch and time ranks, then folds the time rank into the batch rank, calls a certain API of some arbitrary child component, and unfolds the time rank again. """ def __init__(self, sub_component, api_method_name, scope="batch-apply", **kwargs): """ Args: sub_component (Component): The sub-Component to apply the batch to. api_method_name (str): The name of the API-method to call on the sub-component. """ super(BatchApply, self).__init__(scope=scope, **kwargs) self.sub_component = sub_component self.api_method_name = api_method_name # Create the necessary reshape components. self.folder = ReShape(fold_time_rank=True, scope="folder") self.unfolder = ReShape(unfold_time_rank=True, scope="unfolder") self.add_components(self.sub_component, self.folder, self.unfolder) @rlgraph_api def apply(self, input_): folded = self._graph_fn_fold(input_) applied = self._graph_fn_apply(folded) unfolded = self._graph_fn_unfold(applied, input_) return unfolded @graph_fn(flatten_ops=True, split_ops=True) def _graph_fn_fold(self, input_): if get_backend() == "tf": # Fold the time rank. input_folded = self.folder.apply(input_) return input_folded @graph_fn def _graph_fn_apply(self, input_folded): if get_backend() == "tf": # Send the folded input through the sub-component. sub_component_out = getattr(self.sub_component, self.api_method_name)(input_folded) return sub_component_out @graph_fn(flatten_ops=True, split_ops=True) def _graph_fn_unfold(self, sub_component_out, orig_input): if get_backend() == "tf": # Un-fold the time rank again. output = self.unfolder.apply( sub_component_out, input_before_time_rank_folding=orig_input) return output
class IMPALANetwork(NeuralNetwork): """ The base class for both "large and small architecture" versions of the networks used in [1]. [1] IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures - Espeholt, Soyer, Munos et al. - 2018 (https://arxiv.org/abs/1802.01561) """ def __init__(self, worker_sample_size=100, scope="impala-network", **kwargs): """ Args: worker_sample_size (int): How many time-steps an IMPALA actor will have performed in one rollout. """ super(IMPALANetwork, self).__init__(scope=scope, **kwargs) self.worker_sample_size = worker_sample_size # Create all needed sub-components. # ContainerSplitter for the Env signal (dict of 4 keys: for env image, env text, previous action and reward). self.splitter = ContainerSplitter("RGB_INTERLEAVED", "INSTR", "previous_action", "previous_reward", scope="input-splitter") # Fold the time rank into the batch rank. self.time_rank_fold_before_lstm = ReShape( fold_time_rank=True, scope="time-rank-fold-before-lstm") self.time_rank_unfold_before_lstm = ReShape( unfold_time_rank=True, time_major=True, scope="time-rank-unfold-before-lstm") # The Image Processing Stack (left side of "Large Architecture" Figure 3 in [1]). # Conv2D column + ReLU + fc(256) + ReLU. self.image_processing_stack = self.build_image_processing_stack() # The text processing pipeline: Takes a batch of string tensors as input, creates a hash-bucket thereof, # and passes the output of the hash bucket through an embedding-lookup(20) layer. The output of the embedding # lookup is then passed through an LSTM(64). self.text_processing_stack = self.build_text_processing_stack() #self.debug_slicer = Slice(scope="internal-states-slicer", squeeze=True) # The concatenation layer (concatenates outputs from image/text processing stacks, previous action/reward). self.concat_layer = ConcatLayer() # The main LSTM (going into the ActionAdapter (next in the Policy Component that uses this NN Component)). # Use time-major as it's faster (say tf docs). self.main_lstm = LSTMLayer(units=256, scope="lstm-256", time_major=True, static_loop=self.worker_sample_size) # Add all sub-components to this one. self.add_components( self.splitter, self.image_processing_stack, self.text_processing_stack, self.concat_layer, self.main_lstm, self.time_rank_fold_before_lstm, self.time_rank_unfold_before_lstm, #self.debug_slicer ) @staticmethod def build_image_processing_stack(): """ Builds the image processing pipeline for IMPALA and returns it. """ raise NotImplementedError @staticmethod def build_text_processing_stack(): """ Helper function to build the text processing pipeline for both the large and small architectures, consisting of: - ReShape preprocessor to fold the incoming time rank into the batch rank. - StringToHashBucket Layer taking a batch of sentences and converting them to an indices-table of dimensions: cols=length of longest sentences in input rows=number of items in the batch The cols dimension could be interpreted as the time rank into a consecutive LSTM. The StringToHashBucket Component returns the sequence length of each batch item for exactly that purpose. - Embedding Lookup Layer of embedding size 20 and number of rows == num_hash_buckets (see previous layer). - LSTM processing the batched sequences of words coming from the embedding layer as batches of rows. """ num_hash_buckets = 1000 # Create a hash bucket from the sentences and use that bucket to do an embedding lookup (instead of # a vocabulary). string_to_hash_bucket = StringToHashBucket( num_hash_buckets=num_hash_buckets) embedding = EmbeddingLookup(embed_dim=20, vocab_size=num_hash_buckets, pad_empty=True) # The time rank for the LSTM is now the sequence of words in a sentence, NOT the original env time rank. # We will only use the last output of the LSTM-64 for further processing as that is the output after having # seen all words in the sentence. # The original env stepping time rank is currently folded into the batch rank and must be unfolded again before # passing it into the main LSTM. lstm64 = LSTMLayer(units=64, scope="lstm-64", time_major=False) tuple_splitter = ContainerSplitter(tuple_length=2, scope="tuple-splitter") def custom_apply(self, inputs): hash_bucket, lengths = self.sub_components[ "string-to-hash-bucket"].apply(inputs) embedding_output = self.sub_components["embedding-lookup"].apply( hash_bucket) # Return only the last output (sentence of words, where we are not interested in intermediate results # where the LSTM has not seen the entire sentence yet). # Last output is the final internal h-state (slot 1 in the returned LSTM tuple; slot 0 is final c-state). lstm_output = self.sub_components["lstm-64"].apply( embedding_output, sequence_length=lengths) lstm_final_internals = lstm_output["last_internal_states"] # Need to split once more because the LSTM state is always a tuple of final c- and h-states. _, lstm_final_h_state = self.sub_components[ "tuple-splitter"].split(lstm_final_internals) return lstm_final_h_state text_processing_stack = Stack(string_to_hash_bucket, embedding, lstm64, tuple_splitter, api_methods={("apply", custom_apply)}, scope="text-stack") return text_processing_stack @rlgraph_api def apply(self, input_dict, internal_states=None): # Split the input dict coming directly from the Env. _, _, _, orig_previous_reward = self.splitter.split(input_dict) folded_input = self.time_rank_fold_before_lstm.apply(input_dict) image, text, previous_action, previous_reward = self.splitter.split( folded_input) # Get the left-stack (image) and right-stack (text) output (see [1] for details). text_processing_output = self.text_processing_stack.apply(text) image_processing_output = self.image_processing_stack.apply(image) # Concat everything together. concatenated_data = self.concat_layer.apply(image_processing_output, text_processing_output, previous_action, previous_reward) unfolded_concatenated_data = self.time_rank_unfold_before_lstm.apply( concatenated_data, orig_previous_reward) # Feed concat'd input into main LSTM(256). lstm_output = self.main_lstm.apply(unfolded_concatenated_data, internal_states) return lstm_output
class ActionAdapter(Component): """ A Component that cleans up a neural network's flat output and gets it ready for parameterizing a Distribution Component. Processing steps include: - Sending the raw, flattened NN output through a Dense layer whose number of units matches the flattened action space. - Reshaping (according to the action Space). - Translating the reshaped outputs (logits) into probabilities (by softmaxing) and log-probabilities (log). """ def __init__(self, action_space, add_units=0, units=None, weights_spec=None, biases_spec=None, activation=None, scope="action-adapter", **kwargs): """ Args: action_space (Space): The action Space within which this Component will create actions. add_units (Optional[int]): An optional number of units to add to the auto-calculated number of action- layer nodes. Can be negative to subtract units from the auto-calculated value. NOTE: Only one of either `add_units` or `units` must be provided. units (Optional[int]): An optional number of units to use for the action-layer. If None, will calculate the number of units automatically from the given action_space. NOTE: Only one of either `add_units` or `units` must be provided. weights_spec (Optional[any]): An optional RLGraph Initializer spec that will be used to initialize the weights of `self.action layer`. Default: None (use default initializer). biases_spec (Optional[any]): An optional RLGraph Initializer spec that will be used to initialize the biases of `self.action layer`. Default: None (use default initializer, which is usually 0.0). activation (Optional[str]): The activation function to use for `self.action_layer`. Default: None (=linear). """ super(ActionAdapter, self).__init__(scope=scope, **kwargs) self.action_space = action_space.with_batch_rank() self.weights_spec = weights_spec self.biases_spec = biases_spec self.activation = activation # Our (dense) action layer representing the flattened action space. self.action_layer = None # Calculate the number of nodes in the action layer (DenseLayer object) depending on our action Space # or using a given fixed number (`units`). # Also generate the ReShape sub-Component and give it the new_shape. if isinstance(self.action_space, IntBox): if units is None: units = add_units + self.action_space.flat_dim_with_categories self.reshape = ReShape( new_shape=self.action_space.get_shape(with_category_rank=True), flatten_categories=False) else: if units is None: units = add_units + 2 * self.action_space.flat_dim # Those two dimensions are the mean and log sd # Manually add moments after batch/time ranks. new_shape = tuple([2] + list(self.action_space.shape)) self.reshape = ReShape(new_shape=new_shape) assert units > 0, "ERROR: Number of nodes for action-layer calculated as {}! Must be larger 0.".format( units) # Create the action-layer and add it to this component. self.action_layer = DenseLayer(units=units, activation=self.activation, weights_spec=self.weights_spec, biases_spec=self.biases_spec, scope="action-layer") self.add_components(self.action_layer, self.reshape) def check_input_spaces(self, input_spaces, action_space=None): # Check the input Space. last_nn_layer_space = input_spaces["nn_output"] # type: Space sanity_check_space(last_nn_layer_space, non_allowed_types=[ContainerSpace]) # Check the action Space. sanity_check_space(self.action_space, must_have_batch_rank=True) if isinstance(self.action_space, IntBox): sanity_check_space(self.action_space, must_have_categories=True) else: # Fixme: Are there other restraints on continuous action spaces? E.g. no dueling layers? pass @rlgraph_api def get_action_layer_output(self, nn_output): """ Returns the raw, non-reshaped output of the action-layer (DenseLayer) after passing through it the raw nn_output (coming from the previous Component). Args: nn_output (DataOpRecord): The NN output of the preceding neural network. Returns: DataOpRecord: The output of the action layer (a DenseLayer) after passing `nn_output` through it. """ out = self.action_layer.apply(nn_output) return dict(output=out) @rlgraph_api def get_logits(self, nn_output): """ Args: nn_output (DataOpRecord): The NN output of the preceding neural network. Returns: SingleDataOp: The logits (raw nn_output, BUT reshaped). """ aa_output = self.get_action_layer_output(nn_output) logits = self.reshape.apply(aa_output["output"]) return logits @rlgraph_api def get_logits_probabilities_log_probs(self, nn_output): """ Args: nn_output (DataOpRecord): The NN output of the preceding neural network. Returns: Tuple[SingleDataOp]: - logits (raw nn_output, BUT reshaped) - probabilities (softmaxed(logits)) - log(probabilities) """ logits = self.get_logits(nn_output) probabilities, log_probs = self._graph_fn_get_probabilities_log_probs( logits) return dict(logits=logits, probabilities=probabilities, log_probs=log_probs) # TODO: Use a SoftMax Component instead (uses the same code as the one below). @graph_fn def _graph_fn_get_probabilities_log_probs(self, logits): """ Creates properties/parameters and log-probs from some reshaped output. Args: logits (SingleDataOp): The output of some layer that is already reshaped according to our action Space. Returns: tuple (2x SingleDataOp): parameters (DataOp): The parameters, ready to be passed to a Distribution object's get_distribution API-method (usually some probabilities or loc/scale pairs). log_probs (DataOp): Simply the log(parameters). """ if get_backend() == "tf": if isinstance(self.action_space, IntBox): # Discrete actions. parameters = tf.maximum(x=tf.nn.softmax(logits=logits, axis=-1), y=SMALL_NUMBER) # Log probs. log_probs = tf.log(x=parameters) elif isinstance(self.action_space, FloatBox): # Continuous actions. mean, log_sd = tf.split(value=logits, num_or_size_splits=2, axis=1) # Remove moments rank. mean = tf.squeeze(input=mean, axis=1) log_sd = tf.squeeze(input=log_sd, axis=1) # Clip log_sd. log(SMALL_NUMBER) is negative. log_sd = tf.clip_by_value(t=log_sd, clip_value_min=log(SMALL_NUMBER), clip_value_max=-log(SMALL_NUMBER)) # Turn log sd into sd. sd = tf.exp(x=log_sd) parameters = DataOpTuple(mean, sd) log_probs = DataOpTuple(tf.log(x=mean), log_sd) else: raise NotImplementedError return parameters, log_probs elif get_backend() == "pytorch": if isinstance(self.action_space, IntBox): # Discrete actions. softmax_logits = torch.softmax(logits, dim=-1) parameters = torch.max(softmax_logits, SMALL_NUMBER_TORCH) # Log probs. log_probs = torch.log(parameters) elif isinstance(self.action_space, FloatBox): # Continuous actions. mean, log_sd = torch.split(logits, split_size_or_sections=2, dim=1) # Remove moments rank. mean = torch.squeeze(mean, dim=1) log_sd = torch.squeeze(log_sd, dim=1) # Clip log_sd. log(SMALL_NUMBER) is negative. log_sd = torch.clamp(log_sd, min=LOG_SMALL_NUMBER, max=-LOG_SMALL_NUMBER) # Turn log sd into sd. sd = torch.exp(log_sd) parameters = DataOpTuple(mean, sd) log_probs = DataOpTuple(torch.log(mean), log_sd) else: raise NotImplementedError return parameters, log_probs
class Policy(Component): """ A Policy is a wrapper Component that contains a NeuralNetwork, an ActionAdapter and a Distribution Component. """ def __init__(self, network_spec, action_space=None, action_adapter_spec=None, max_likelihood=True, scope="policy", **kwargs): """ Args: network_spec (Union[NeuralNetwork,dict]): The NeuralNetwork Component or a specification dict to build one. action_space (Space): The action Space within which this Component will create actions. action_adapter_spec (Optional[dict]): A spec-dict to create an ActionAdapter. Use None for the default ActionAdapter object. max_likelihood (bool): Whether to pick actions according to the max-likelihood value or via sampling. Default: True. """ super(Policy, self).__init__(scope=scope, **kwargs) self.neural_network = NeuralNetwork.from_spec(network_spec) if action_space is None: self.action_adapter = ActionAdapter.from_spec(action_adapter_spec) action_space = self.action_adapter.action_space else: self.action_adapter = ActionAdapter.from_spec( action_adapter_spec, action_space=action_space) self.action_space = action_space self.max_likelihood = max_likelihood # TODO: Hacky trick to implement IMPALA post-LSTM256 time-rank folding and unfolding. # TODO: Replace entirely via sonnet-like BatchApply Component. is_impala = "IMPALANetwork" in type(self.neural_network).__name__ # Add API-method to get baseline output (if we use an extra value function baseline node). if isinstance(self.action_adapter, BaselineActionAdapter): # TODO: IMPALA attempt to speed up final pass after LSTM. if is_impala: self.time_rank_folder = ReShape(fold_time_rank=True, scope="time-rank-fold") self.time_rank_unfolder_v = ReShape(unfold_time_rank=True, time_major=True, scope="time-rank-unfold-v") self.time_rank_unfolder_a_probs = ReShape( unfold_time_rank=True, time_major=True, scope="time-rank-unfold-a-probs") self.time_rank_unfolder_logits = ReShape( unfold_time_rank=True, time_major=True, scope="time-rank-unfold-logits") self.time_rank_unfolder_log_probs = ReShape( unfold_time_rank=True, time_major=True, scope="time-rank-unfold-log-probs") self.add_components(self.time_rank_folder, self.time_rank_unfolder_v, self.time_rank_unfolder_a_probs, self.time_rank_unfolder_log_probs, self.time_rank_unfolder_logits) @rlgraph_api(component=self) def get_state_values_logits_probabilities_log_probs( self, nn_input, internal_states=None): nn_output = self.neural_network.apply(nn_input, internal_states) last_internal_states = nn_output.get("last_internal_states") nn_output = nn_output["output"] # TODO: IMPALA attempt to speed up final pass after LSTM. if is_impala: nn_output = self.time_rank_folder.apply(nn_output) out = self.action_adapter.get_logits_probabilities_log_probs( nn_output) # TODO: IMPALA attempt to speed up final pass after LSTM. if is_impala: state_values = self.time_rank_unfolder_v.apply( out["state_values"], nn_output) logits = self.time_rank_unfolder_logits.apply( out["logits"], nn_output) probs = self.time_rank_unfolder_a_probs.apply( out["probabilities"], nn_output) log_probs = self.time_rank_unfolder_log_probs.apply( out["log_probs"], nn_output) else: state_values = out["state_values"] logits = out["logits"] probs = out["probabilities"] log_probs = out["log_probs"] return dict(state_values=state_values, logits=logits, probabilities=probs, log_probs=log_probs, last_internal_states=last_internal_states) # Figure out our Distribution. if isinstance(action_space, IntBox): self.distribution = Categorical() # Continuous action space -> Normal distribution (each action needs mean and variance from network). elif isinstance(action_space, FloatBox): self.distribution = Normal() else: raise RLGraphError( "ERROR: `action_space` is of type {} and not allowed in {} Component!" .format(type(action_space).__name__, self.name)) self.add_components(self.neural_network, self.action_adapter, self.distribution) if is_impala: self.add_components(self.time_rank_folder, self.time_rank_unfolder_v, self.time_rank_unfolder_a_probs, self.time_rank_unfolder_log_probs, self.time_rank_unfolder_logits) # Define our interface. @rlgraph_api def get_nn_output(self, nn_input, internal_states=None): """ Args: nn_input (any): The input to our neural network. internal_states (Optional[any]): The initial internal states going into an RNN-based neural network. Returns: any: The raw output of the neural network (before it's cleaned-up and passed through the ActionAdapter). """ out = self.neural_network.apply(nn_input, internal_states) return dict(output=out["output"], last_internal_states=out.get("last_internal_states")) @rlgraph_api def get_action(self, nn_input, internal_states=None, max_likelihood=None): """ Returns an action based on NN output, action adapter output and distribution sampling. Args: nn_input (any): The input to our neural network. internal_states (Optional[any]): The initial internal states going into an RNN-based neural network. max_likelihood (Optional[bool]): If not None, use this to determine whether actions should be drawn from the distribution in max-likelihood or stochastic fashion. Returns: any: The drawn action. """ max_likelihood = self.max_likelihood if max_likelihood is None else max_likelihood nn_output = self.get_nn_output(nn_input, internal_states) # Skip our distribution, iff discrete action-space and max-likelihood acting (greedy). # In that case, one does not need to create a distribution in the graph each act (only to get the argmax # over the logits, which is the same as the argmax over the probabilities (or log-probabilities)). if max_likelihood is True and isinstance(self.action_space, IntBox): out = self.action_adapter.get_logits_probabilities_log_probs( nn_output["output"]) action = self._graph_fn_get_max_likelihood_action_wo_distribution( out["logits"]) else: out = self.action_adapter.get_logits_probabilities_log_probs( nn_output["output"]) action = self.distribution.draw(out["probabilities"], max_likelihood) return dict(action=action, last_internal_states=nn_output["last_internal_states"]) @rlgraph_api def get_max_likelihood_action(self, nn_input, internal_states=None): """ Args: nn_input (any): The input to our neural network. internal_states (Optional[any]): The initial internal states going into an RNN-based neural network. Returns: any: See `get_action`, but with max_likelihood force set to True. """ out = self.get_logits_probabilities_log_probs(nn_input, internal_states) if isinstance(self.action_space, IntBox): action = self._graph_fn_get_max_likelihood_action_wo_distribution( out["logits"]) else: action = self.distribution.sample_deterministic( out["probabilities"]) return dict(action=action, last_internal_states=out["last_internal_states"]) @rlgraph_api def get_stochastic_action(self, nn_input, internal_states=None): """ Args: nn_input (any): The input to our neural network. internal_states (Optional[any]): The initial internal states going into an RNN-based neural network. Returns: any: See `get_action`, but with max_likelihood force set to False. """ out = self.get_logits_probabilities_log_probs(nn_input, internal_states) action = self.distribution.sample_stochastic(out["probabilities"]) return dict(action=action, last_internal_states=out["last_internal_states"]) @rlgraph_api def get_action_layer_output(self, nn_input, internal_states=None): """ Args: nn_input (any): The input to our neural network. internal_states (Optional[any]): The initial internal states going into an RNN-based neural network. Returns: any: The raw output of the action layer of the ActionAdapter (including possibly the last internal states of a RNN-based NN). """ nn_output = self.get_nn_output(nn_input, internal_states) action_layer_output = self.action_adapter.get_action_layer_output( nn_output["output"]) # Add last internal states to return value. return dict(output=action_layer_output["output"], last_internal_states=nn_output["last_internal_states"]) @rlgraph_api def get_logits_probabilities_log_probs(self, nn_input, internal_states=None): """ Args: nn_input (any): The input to our neural network. internal_states (Optional[any]): The initial internal states going into an RNN-based neural network. Returns: Dict: logits: The (reshaped) logits from the ActionAdapter. probabilities: The probabilities gained from the softmaxed logits. log_probs: The log(probabilities) values. """ nn_output = self.get_nn_output(nn_input, internal_states) aa_output = self.action_adapter.get_logits_probabilities_log_probs( nn_output["output"]) return dict(logits=aa_output["logits"], probabilities=aa_output["probabilities"], log_probs=aa_output["log_probs"], last_internal_states=nn_output["last_internal_states"]) @rlgraph_api def get_entropy(self, nn_input, internal_states=None): """ Args: nn_input (any): The input to our neural network. internal_states (Optional[any]): The initial internal states going into an RNN-based neural network. Returns: any: See Distribution component. """ out = self.get_logits_probabilities_log_probs(nn_input, internal_states) entropy = self.distribution.entropy(out["probabilities"]) return dict(entropy=entropy, last_internal_states=out["last_internal_states"]) @graph_fn def _graph_fn_get_max_likelihood_action_wo_distribution(self, logits): """ Use this function only for discrete action spaces to circumvent using a full-blown backend-specific distribution object (e.g. tf.distribution.Multinomial). Args: logits (SingleDataOp): Logits over which to pick the argmax (greedy action). Returns: SingleDataOp: The argmax over the last rank of the input logits. """ if get_backend() == "tf": return tf.argmax(logits, axis=-1, output_type=tf.int32) elif get_backend() == "pytorch": return torch.argmax(logits, dim=-1).int()